Data processing backend for DataLens
from __future__ import annotations

import dl_version


app_version = f"dl_api:{dl_version.__version__}"

app/dl_control_api/dl_control_api/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

import flask

from dl_api_commons.sentry_config import (
    SentryConfig,
    hook_configure_configure_sentry_for_flask,
)
from dl_api_lib.app_settings import (
    ControlApiAppSettings,
    ControlApiAppTestingsSettings,
)
from dl_api_lib.loader import (
    ApiLibraryConfig,
    load_api_lib,
    preload_api_lib,
)
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.env_var_definitions import (
    jaeger_service_name_env_aware,
    use_jaeger_tracer,
)
from dl_configs.settings_loaders.loader_env import (
    load_connectors_settings_from_env_with_fallback,
    load_settings_from_env_with_fallback,
)
from dl_constants.enums import ConnectionType
from dl_control_api import app_version
from dl_control_api.app_factory import StandaloneControlApiAppFactory
from dl_core.connectors.settings.registry import (
    CONNECTORS_SETTINGS_CLASSES,
    CONNECTORS_SETTINGS_FALLBACKS,
)
from dl_core.loader import CoreLibraryConfig
from dl_core.logging_config import hook_configure_logging


def create_app(
    app_settings: ControlApiAppSettings,
    connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
    testing_app_settings: Optional[ControlApiAppTestingsSettings] = None,
    close_loop_after_request: bool = True,
) -> flask.Flask:
    mng_app_factory = StandaloneControlApiAppFactory(settings=app_settings)
    return mng_app_factory.create_app(
        connectors_settings=connectors_settings,
        testing_app_settings=testing_app_settings,
        close_loop_after_request=close_loop_after_request,
    )


def create_uwsgi_app() -> flask.Flask:
    preload_api_lib()
    settings = load_settings_from_env_with_fallback(ControlApiAppSettings)
    load_api_lib(
        ApiLibraryConfig(
            api_connector_ep_names=settings.BI_API_CONNECTOR_WHITELIST,
            core_lib_config=CoreLibraryConfig(core_connector_ep_names=settings.CORE_CONNECTOR_WHITELIST),
        )
    )
    connectors_settings = load_connectors_settings_from_env_with_fallback(
        settings_registry=CONNECTORS_SETTINGS_CLASSES,
        fallbacks=CONNECTORS_SETTINGS_FALLBACKS,
    )
    uwsgi_app = create_app(settings, connectors_settings)

    actual_sentry_dsn: Optional[str] = settings.SENTRY_DSN if settings.SENTRY_ENABLED else None

    hook_configure_logging(
        uwsgi_app,
        app_name="dl_api_app",
        app_prefix="a",
        use_jaeger_tracer=use_jaeger_tracer(),
        jaeger_service_name=jaeger_service_name_env_aware("dl-api-app"),
    )
    if actual_sentry_dsn is not None:
        hook_configure_configure_sentry_for_flask(
            uwsgi_app,
            SentryConfig(dsn=actual_sentry_dsn, release=app_version),
        )

    return uwsgi_app


app = create_uwsgi_app()

app/dl_control_api/dl_control_api/app.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

import flask

from dl_api_lib.app.control_api.app import (
    ControlApiAppFactory,
    EnvSetupResult,
)
from dl_api_lib.app_common import (
    SRFactoryBuilder,
    StandaloneServiceRegistryFactory,
)
from dl_api_lib.app_settings import (
    ControlApiAppSettings,
    ControlApiAppTestingsSettings,
)
from dl_api_lib.connector_availability.base import ConnectorAvailabilityConfig
from dl_configs.enums import RequiredService
from dl_constants.enums import USAuthMode
from dl_core.data_processing.cache.primitives import CacheTTLConfig
from dl_core.services_registry.entity_checker import EntityUsageChecker
from dl_core.services_registry.env_manager_factory import InsecureEnvManagerFactory
from dl_core.services_registry.env_manager_factory_base import EnvManagerFactory
from dl_core.services_registry.rqe_caches import RQECachesSetting


class StandaloneControlApiSRFactoryBuilder(SRFactoryBuilder[ControlApiAppSettings]):
    def _get_required_services(self, settings: ControlApiAppSettings) -> set[RequiredService]:
        return set()

    def _get_env_manager_factory(self, settings: ControlApiAppSettings) -> EnvManagerFactory:
        return InsecureEnvManagerFactory()

    def _get_inst_specific_sr_factory(
        self,
        settings: ControlApiAppSettings,
    ) -> StandaloneServiceRegistryFactory:
        return StandaloneServiceRegistryFactory()

    def _get_entity_usage_checker(self, settings: ControlApiAppSettings) -> Optional[EntityUsageChecker]:
        return None

    def _get_bleeding_edge_users(self, settings: ControlApiAppSettings) -> tuple[str, ...]:
        return tuple()

    def _get_rqe_caches_settings(self, settings: ControlApiAppSettings) -> Optional[RQECachesSetting]:
        return None

    def _get_default_cache_ttl_settings(self, settings: ControlApiAppSettings) -> Optional[CacheTTLConfig]:
        return None

    def _get_connector_availability(self, settings: ControlApiAppSettings) -> Optional[ConnectorAvailabilityConfig]:
        return settings.CONNECTOR_AVAILABILITY


class StandaloneControlApiAppFactory(ControlApiAppFactory[ControlApiAppSettings], StandaloneControlApiSRFactoryBuilder):
    def set_up_environment(
        self,
        app: flask.Flask,
        testing_app_settings: Optional[ControlApiAppTestingsSettings] = None,
    ) -> EnvSetupResult:
        us_auth_mode: USAuthMode
        from dl_core.flask_utils.trust_auth import TrustAuthService

        TrustAuthService(
            fake_user_id="_user_id_",
            fake_user_name="_user_name_",
            fake_tenant=None if testing_app_settings is None else testing_app_settings.fake_tenant,
        ).set_up(app)

        us_auth_mode_override = None if testing_app_settings is None else testing_app_settings.us_auth_mode_override
        us_auth_mode = USAuthMode.master if us_auth_mode_override is None else us_auth_mode_override

        result = EnvSetupResult(us_auth_mode=us_auth_mode)
        return result

app/dl_control_api/dl_control_api/app_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import dl_version


app_version = f"dl_api:{dl_version.__version__}"

app/dl_data_api/dl_data_api/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import os

from aiohttp import web

from dl_api_lib.app_settings import DataApiAppSettings
from dl_api_lib.loader import (
    ApiLibraryConfig,
    load_api_lib,
    preload_api_lib,
)
from dl_app_tools.aio_latency_tracking import LatencyTracker
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.env_var_definitions import (
    jaeger_service_name_env_aware,
    use_jaeger_tracer,
)
from dl_configs.settings_loaders.loader_env import (
    load_connectors_settings_from_env_with_fallback,
    load_settings_from_env_with_fallback,
)
from dl_constants.enums import ConnectionType
from dl_core.connectors.settings.registry import (
    CONNECTORS_SETTINGS_CLASSES,
    CONNECTORS_SETTINGS_FALLBACKS,
)
from dl_core.loader import CoreLibraryConfig
from dl_core.logging_config import configure_logging
from dl_data_api.app_factory import StandaloneDataApiAppFactory


LOGGER = logging.getLogger(__name__)


def create_app(
    setting: DataApiAppSettings,
    connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
) -> web.Application:
    data_api_app_factory = StandaloneDataApiAppFactory(settings=setting)
    return data_api_app_factory.create_app(
        connectors_settings=connectors_settings,
    )


async def create_gunicorn_app(start_selfcheck: bool = True) -> web.Application:
    preload_api_lib()
    settings = load_settings_from_env_with_fallback(DataApiAppSettings)
    load_api_lib(
        ApiLibraryConfig(
            api_connector_ep_names=settings.BI_API_CONNECTOR_WHITELIST,
            core_lib_config=CoreLibraryConfig(core_connector_ep_names=settings.CORE_CONNECTOR_WHITELIST),
        )
    )
    connectors_settings = load_connectors_settings_from_env_with_fallback(
        settings_registry=CONNECTORS_SETTINGS_CLASSES,
        fallbacks=CONNECTORS_SETTINGS_FALLBACKS,
    )
    configure_logging(
        app_name=settings.app_name,
        app_prefix=settings.app_prefix,
        use_jaeger_tracer=use_jaeger_tracer(),
        jaeger_service_name=jaeger_service_name_env_aware(settings.jaeger_service_name),
    )
    try:
        LOGGER.info("Creating application instance...")
        app = create_app(setting=settings, connectors_settings=connectors_settings)
        if start_selfcheck:
            LOGGER.info("Starting selfcheck aio task...")
            await LatencyTracker().run_task()
        LOGGER.info("Application instance was created")
        return app
    except Exception:
        LOGGER.exception("Exception during app creation")
        raise


def main() -> None:
    host = os.environ["APP_HOST"]
    port = int(os.environ["APP_PORT"])
    app_coro = create_gunicorn_app()
    web.run_app(app_coro, host=host, port=port)


if __name__ == "__main__":
    main()

app/dl_data_api/dl_data_api/app.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

from dl_api_commons.aio.middlewares.auth_trust_middleware import auth_trust_middleware
from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_api_lib.app.data_api.app import (
    DataApiAppFactory,
    EnvSetupResult,
)
from dl_api_lib.app_common import SRFactoryBuilder
from dl_api_lib.app_common_settings import ConnOptionsMutatorsFactory
from dl_api_lib.app_settings import (
    AppSettings,
    DataApiAppSettings,
)
from dl_api_lib.connector_availability.base import ConnectorAvailabilityConfig
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.enums import RequiredService
from dl_constants.enums import ConnectionType
from dl_core.aio.middlewares.services_registry import services_registry_middleware
from dl_core.aio.middlewares.us_manager import service_us_manager_middleware
from dl_core.data_processing.cache.primitives import CacheTTLConfig
from dl_core.services_registry.entity_checker import EntityUsageChecker
from dl_core.services_registry.env_manager_factory import InsecureEnvManagerFactory
from dl_core.services_registry.env_manager_factory_base import EnvManagerFactory
from dl_core.services_registry.inst_specific_sr import InstallationSpecificServiceRegistryFactory
from dl_core.services_registry.rqe_caches import RQECachesSetting
from dl_data_api import app_version


class StandaloneDataApiSRFactoryBuilder(SRFactoryBuilder[AppSettings]):
    @property
    def _is_async_env(self) -> bool:
        return True

    def _get_required_services(self, settings: AppSettings) -> set[RequiredService]:
        return set()

    def _get_env_manager_factory(self, settings: AppSettings) -> EnvManagerFactory:
        return InsecureEnvManagerFactory()

    def _get_inst_specific_sr_factory(
        self,
        settings: AppSettings,
    ) -> Optional[InstallationSpecificServiceRegistryFactory]:
        return None

    def _get_entity_usage_checker(self, settings: AppSettings) -> Optional[EntityUsageChecker]:
        return None

    def _get_bleeding_edge_users(self, settings: AppSettings) -> tuple[str, ...]:
        return tuple()

    def _get_rqe_caches_settings(self, settings: DataApiAppSettings) -> Optional[RQECachesSetting]:  # type: ignore[override]
        return None

    def _get_default_cache_ttl_settings(self, settings: DataApiAppSettings) -> Optional[CacheTTLConfig]:  # type: ignore[override]
        return None

    def _get_connector_availability(self, settings: AppSettings) -> Optional[ConnectorAvailabilityConfig]:
        return None


class StandaloneDataApiAppFactory(DataApiAppFactory[DataApiAppSettings], StandaloneDataApiSRFactoryBuilder):
    @property
    def _is_public(self) -> bool:
        return False

    def get_app_version(self) -> str:
        return app_version

    def set_up_environment(
        self,
        connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
    ) -> EnvSetupResult:
        auth_mw_list: list[AIOHTTPMiddleware]
        sr_middleware_list: list[AIOHTTPMiddleware]
        usm_middleware_list: list[AIOHTTPMiddleware]

        conn_opts_factory = ConnOptionsMutatorsFactory()
        sr_factory = self.get_sr_factory(
            settings=self._settings, conn_opts_factory=conn_opts_factory, connectors_settings=connectors_settings
        )

        # Auth middlewares
        auth_mw_list = [
            auth_trust_middleware(
                fake_user_id="_user_id_",
                fake_user_name="_user_name_",
            )
        ]

        # SR middlewares
        sr_middleware_list = [
            services_registry_middleware(
                services_registry_factory=sr_factory,
                use_query_cache=self._settings.CACHES_ON,
                use_mutation_cache=self._settings.MUTATIONS_CACHES_ON,
                mutation_cache_default_ttl=self._settings.MUTATIONS_CACHES_DEFAULT_TTL,
            ),
        ]

        # US manager middlewares
        common_us_kw = dict(
            us_base_url=self._settings.US_BASE_URL,
            crypto_keys_config=self._settings.CRYPTO_KEYS_CONFIG,
        )
        usm_middleware_list = [
            service_us_manager_middleware(us_master_token=self._settings.US_MASTER_TOKEN, **common_us_kw),  # type: ignore
            service_us_manager_middleware(us_master_token=self._settings.US_MASTER_TOKEN, as_user_usm=True, **common_us_kw),  # type: ignore
        ]

        result = EnvSetupResult(
            auth_mw_list=auth_mw_list,
            sr_middleware_list=sr_middleware_list,
            usm_middleware_list=usm_middleware_list,
        )

        return result

app/dl_data_api/dl_data_api/app_factory.py
datalens-tech/datalens-backend
# higher than the supposed default nginx default of 75 seconds
keepalive = 80
timeout = 90

# # Configuration in commandline:
# import os
# log_config = os.environ['GUNICORN_LOG_CONFIG']
# bind = "[::]:" + os.environ.get('GUNICORN_LISTEN_PORT', '8080')
# max_requests = int(os.environ.get('GUNICORN_MAX_RQ', '5000'))
# max_requests_jitter = int(os.environ.get('GUNICORN_MAX_RQ_JITTER', '500'))
# worker_class = 'aiohttp.GunicornWebWorker'
# workers = int(os.environ.get('GUNICORN_WORKERS_COUNT', '1'))

# # Other notable configuration:
# worker_connections = 1000
# threads = 2
# proc_name = 'app01'
# backlog = 2048
# accesslog = 'access.log'
# errorlog = 'error.log'
# user = 'ubuntu'
# group = 'ubuntu'

app/dl_data_api/docker/gunicorn_config.py
datalens-tech/datalens-backend

from .ext.declarative import get_declarative_base
from .orm.session import make_session
from .sql import Table, select


__version__ = '0.1.3.14'
VERSION = tuple(int(piece) for piece in __version__.split('.'))


__all__ = (
    'get_declarative_base',
    'make_session',
    'Table', 'select'
)

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/__init__.py
datalens-tech/datalens-backend
import re
import itertools
import enum

import sqlalchemy as sa
from sqlalchemy import schema, types as sqltypes, exc, util as sa_util
from sqlalchemy.engine import default, reflection
from sqlalchemy.sql import (
    compiler, expression, type_api, literal_column, elements, COLLECT_CARTESIAN_PRODUCTS
)
from sqlalchemy.sql.compiler import FromLinter, WARN_LINTING
from sqlalchemy.sql.ddl import CreateColumn
from sqlalchemy.types import DATE, DATETIME, FLOAT  # noqa
from sqlalchemy.util.compat import inspect_getfullargspec
from sqlalchemy.util import (
    warn,
    to_list,
    text_type,
)

from clickhouse_sqlalchemy import Table
from .. import types
from ..util import compat
from ..quoting import QUOTER

# Column specifications
colspecs = {}


# Type converters
ischema_names = {
    'Int64': types.Int64,
    'Int32': types.Int32,
    'Int16': types.Int16,
    'Int8': types.Int8,
    'UInt64': types.UInt64,
    'UInt32': types.UInt32,
    'UInt16': types.UInt16,
    'UInt8': types.UInt8,
    'Bool': types.Bool,
    'Date': types.Date,
    'Date32': types.Date,
    'DateTime': types.DateTime,
    'Float64': types.Float64,
    'Float32': types.Float32,
    'Decimal': types.Decimal,
    'String': types.String,
    'UUID': types.UUID,
    'IPv4': types.IPv4,
    'IPv6': types.IPv6,
    'FixedString': types.String,
    'Enum8': types.Enum8,
    'Enum16': types.Enum16,
    '_array': types.Array,
    '_nullable': types.Nullable,
    '_lowcardinality': types.LowCardinality,
    '_datetimewithtz': types.DateTimeWithTZ,
    '_datetime64': types.DateTime64,
    '_datetime64withtz': types.DateTime64WithTZ,
}


class ClickHouseIdentifierPreparer(compiler.IdentifierPreparer):

    reserved_words = compiler.IdentifierPreparer.reserved_words | set((
        'index',  # reserved in the 'create table' syntax, at least.
    ))
    # Alternatively, use `_requires_quotes = lambda self, value: True`

    def _escape_identifier(self, value):
        value = value.replace(self.escape_quote, self.escape_to_quote)
        # same semantics in CH regarding '\\', '\n', ...
        # (the `'` -> `\'` is not strictly necessary,
        # but semantically equal anyway)
        value = QUOTER.quote_str_inner(value)
        value = value.replace('%', '%%')
        return value

    # copypaste for sqlalchemy 1.2 compat
    def _requires_quotes_illegal_chars(self, value):
        return not self.legal_characters.match(text_type(value))


class ClickHouseInspector(reflection.Inspector):
    def reflect_table(
            self, table,
            include_columns=None, exclude_columns=(),
            resolve_fks=True, _extend_on=None,
            **kwargs,
    ):
        if not isinstance(table, Table):
            table.metadata.remove(table)
            ch_table = Table._make_from_standard(
                table, _extend_on=_extend_on,
            )
        else:
            ch_table = table
        super().reflect_table(
            ch_table,
            include_columns, exclude_columns,
            resolve_fks, _extend_on,
            **kwargs,
        )


class ClickHouseCompiler(compiler.SQLCompiler):

    _param_replace_re = r'[(){}:%"\'-]'

    def __init__(self, *args, **kwargs):
        self._param_replace_counter = itertools.count(1)
        self._param_replace_map = {}
        # Actually does the compilation:
        super(ClickHouseCompiler, self).__init__(*args, **kwargs)

    def bindparam_string(self, name, **kwargs):
        name_orig = name
        name = re.sub(self._param_replace_re, '_', name)
        if name != name_orig:
            name = '{}_{}'.format(name, next(self._param_replace_counter))
            self._param_replace_map[name_orig] = name
        return super(ClickHouseCompiler, self).bindparam_string(name, **kwargs)

    def visit_mod_binary(self, binary, operator, **kw):
        return self.process(binary.left, **kw) + ' %% ' + \
            self.process(binary.right, **kw)

    def post_process_text(self, text):
        return text.replace('%', '%%')

    def visit_count_func(self, fn, **kw):
        # count accepts zero arguments.
        return 'count%s' % self.process(fn.clause_expr, **kw)

    def visit_case(self, clause, **kwargs):
        text = 'CASE '
        if clause.value is not None:
            text += clause.value._compiler_dispatch(self, **kwargs) + ' '
        for cond, result in clause.whens:
            text += 'WHEN ' + cond._compiler_dispatch(
                self, **kwargs
            ) + ' THEN ' + result._compiler_dispatch(
                self, **kwargs) + " "
        if clause.else_ is None:
            raise exc.CompileError('ELSE clause is required in CASE')

        text += 'ELSE ' + clause.else_._compiler_dispatch(
            self, **kwargs
        ) + ' END'
        return text

    def visit_if__func(self, func, **kw):
        return "(%s) ? (%s) : (%s)" % (
            self.process(func.clauses.clauses[0], **kw),
            self.process(func.clauses.clauses[1], **kw),
            self.process(func.clauses.clauses[2], **kw)
        )

    def visit_column(self, column, include_table=True, **kwargs):
        # Columns prefixed with table name are not supported
        return super(ClickHouseCompiler, self).visit_column(
            column, include_table=False, **kwargs
        )

    def limit_clause(self, select, **kw):
        text = ''
        if select._limit_clause is not None:
            text += ' \n LIMIT '
            if select._offset_clause is not None:
                text += self.process(select._offset_clause, **kw) + ', '
            text += self.process(select._limit_clause, **kw)
        else:
            if select._offset_clause is not None:
                raise exc.CompileError('OFFSET without LIMIT is not supported')

        return text

    def visit_lambda(self, lambda_, **kw):
        func = lambda_.func
        spec = inspect_getfullargspec(func)

        if spec.varargs:
            raise exc.CompileError('Lambdas with *args are not supported')

        try:
            # ArgSpec in SA>=1.3.0b2
            keywords = spec.keywords
        except AttributeError:
            # FullArgSpec in SA>=1.3.0b2
            keywords = spec.varkw

        if keywords:
            raise exc.CompileError('Lambdas with **kwargs are not supported')

        text = ', '.join(spec.args) + ' -> '

        args = [literal_column(arg) for arg in spec.args]
        text += self.process(func(*args), **kw)

        return text

    def visit_parametrized_function(self, func, **kwargs):
        name = func.name
        name = (
            self.preparer.quote(name)
            if self.preparer._requires_quotes_illegal_chars(name)
            or isinstance(name, elements.quoted_name)
            else name
        )

        params = func.params_expr._compiler_dispatch(self, **kwargs)
        args = self.function_argspec(func, **kwargs)
        return "%(name)s%(params)s%(args)s" % dict(
            name=name, params=params, args=args)

    def visit_extract(self, extract, **kw):
        field = self.extract_map.get(extract.field, extract.field)
        column = self.process(extract.expr, **kw)
        if field == 'year':
            return 'toYear(%s)' % column
        elif field == 'month':
            return 'toMonth(%s)' % column
        elif field == 'day':
            return 'toDayOfMonth(%s)' % column
        else:
            return column

    def visit_join(self, join, asfrom=False, **kwargs):
        text = join.left._compiler_dispatch(self, asfrom=asfrom, **kwargs)
        # need to make a variable to prevent leaks in some debuggers
        join_type = getattr(join, 'type', None)
        if join_type is None:
            if join.isouter:
                join_type = 'LEFT OUTER'
            else:
                join_type = 'INNER'
        elif join_type is not None:
            join_type = join_type.upper()
            if join.isouter and 'INNER' in join_type:
                raise exc.CompileError(
                    "can't compile join with specified "
                    "INNER type and isouter=True"
                )
            # isouter=False by default, disable that checking
            # elif not join.isouter and 'OUTER' in join.type:
            #     raise exc.CompileError(
            #         "can't compile join with specified "
            #         "OUTER type and isouter=False"
            #     )
        if join.full and 'FULL' not in join_type:
            join_type = 'FULL ' + join_type

        if getattr(join, 'strictness', None):
            join_type = join.strictness.upper() + ' ' + join_type

        if getattr(join, 'distribution', None):
            join_type = join.distribution.upper() + ' ' + join_type

        if join_type is not None:
            text += ' ' + join_type.upper() + ' JOIN '

        onclause = join.onclause

        text += join.right._compiler_dispatch(self, asfrom=asfrom, **kwargs)
        if isinstance(onclause, elements.Tuple):
            text += ' USING ' + onclause._compiler_dispatch(self, **kwargs)
        else:
            text += ' ON ' + onclause._compiler_dispatch(self, **kwargs)
        return text

    def visit_array_join(self, array_join, **kwargs):
        kwargs['within_columns_clause'] = True
        return ' \nARRAY JOIN {columns}'.format(
            columns=', '.join(
                col._compiler_dispatch(self,
                                       within_label_clause=False,
                                       **kwargs)
                for col in array_join.clauses

            )
        )

    # TODO: integrational tests + send to upstream.
    def visit_notbetween_op_binary(self, binary, operator, **kw):
        """
        Overwrite default SA behavior that uses the NOT BETWEEN operator,
        which is not supported by ClickHouse
        """
        btw_text = self._generate_generic_binary(
            binary,
            " BETWEEN ",
            **kw
        )
        return "NOT ({})".format(btw_text)

    def visit_label(self,
                    label,
                    from_labeled_label=False,
                    **kw):
        if from_labeled_label:
            return super(ClickHouseCompiler, self).visit_label(
                label,
                render_label_as_label=label
            )
        else:
            return super(ClickHouseCompiler, self).visit_label(
                label,
                **kw
            )

    def _compose_select_body(
        self,
        text,
        select,
        compile_state,
        inner_columns,
        froms,
        byfrom,
        toplevel,
        kwargs,
    ):
        text += ", ".join(inner_columns)

        # TODO: confirm that we need this part of code in overridden method
        if self.linting & COLLECT_CARTESIAN_PRODUCTS:
            from_linter = FromLinter({}, set())
            warn_linting = self.linting & WARN_LINTING
            if toplevel:
                self.from_linter = from_linter
        else:
            from_linter = None
            warn_linting = False

        if froms:
            text += " \nFROM "

            if select._hints:
                text += ", ".join(
                    [
                        f._compiler_dispatch(
                            self,
                            asfrom=True,
                            fromhints=byfrom,
                            from_linter=from_linter,
                            **kwargs
                        )
                        for f in froms
                    ]
                )
            else:
                text += ", ".join(
                    [
                        f._compiler_dispatch(
                            self,
                            asfrom=True,
                            from_linter=from_linter,
                            **kwargs
                        )
                        for f in froms
                    ]
                )
        else:
            text += self.default_from()

        if getattr(select, '_array_join', None) is not None:
            text += select._array_join._compiler_dispatch(self, **kwargs)

        sample_clause = getattr(select, '_sample_clause', None)

        if sample_clause is not None:
            text += self.sample_clause(select, **kwargs)

        final_clause = getattr(select, '_final_clause', None)

        if final_clause is not None:
            text += self.final_clause()

        if select._where_criteria:
            t = self._generate_delimited_and_list(
                select._where_criteria, from_linter=from_linter, **kwargs
            )
            if t:
                text += " \nWHERE " + t

        if warn_linting:
            from_linter.warn()

        if select._group_by_clauses:
            text += self.group_by_clause(select, **kwargs)

        if select._having_criteria:
            t = self._generate_delimited_and_list(select._having_criteria, **kwargs)
            if t:
                text += " \nHAVING " + t

        if select._order_by_clauses:
            text += self.order_by_clause(select, **kwargs)

        if select._has_row_limiting_clause:
            text += self._row_limit_clause(select, **kwargs)

        if select._for_update_arg is not None:
            text += self.for_update_clause(select, **kwargs)

        return text

    def sample_clause(self, select, **kw):
        return " \nSAMPLE " + self.process(select._sample_clause, **kw)

    def final_clause(self):
        return " \nFINAL"

    def group_by_clause(self, select, **kw):
        text = ""

        group_by = select._group_by_clause._compiler_dispatch(
            self, **kw)

        if group_by:
            text = " GROUP BY " + group_by

            if getattr(select, '_with_totals', False):
                text += " WITH TOTALS"

        return text

    def render_literal_value(self, value, type_):
        """
        Primarily used for compiling with `literal_binds=True`.
        """
        return self.post_process_text(QUOTER.quote(value))


class ClickHouseDDLCompiler(compiler.DDLCompiler):

    def visit_create_column(self, create, **kw):
        column = create.element
        nullable = column.nullable

        # All columns including synthetic PKs must be 'nullable'
        column.nullable = True

        rv = super(ClickHouseDDLCompiler, self).visit_create_column(
            create, **kw
        )
        column.nullable = nullable

        return rv

    def visit_primary_key_constraint(self, constraint):
        # Do not render PKs.
        return ''

    def _compile_param(self, expr):
        compiler = self.sql_compiler
        if isinstance(expr, (list, tuple)):
            return '(' + ', '.join(
                self._compile_param(el) for el in expr
            ) + ')'
        if not isinstance(expr, expression.ColumnClause):
            if not hasattr(expr, 'self_group'):
                # assuming base type (int, string, etc.)
                return compat.text_type(expr)
            else:
                expr = expr.self_group()
        return compiler.process(
            expr, include_table=False, literal_binds=True
        )

    def visit_merge_tree(self, engine):
        param = engine.get_parameters()
        if param is None:
            param = '()'
        else:
            param = self._compile_param(to_list(param))

        text = '{0}{1}\n'.format(engine.name, param)
        if engine.partition_by:
            text += ' PARTITION BY {0}\n'.format(
                self._compile_param(
                    engine.partition_by.get_expressions_or_columns()[0]
                )
            )
        if engine.order_by:
            text += ' ORDER BY {0}\n'.format(
                self._compile_param(
                    engine.order_by.get_expressions_or_columns()
                )
            )
        if engine.primary_key:
            text += ' PRIMARY KEY {0}\n'.format(
                self._compile_param(
                    engine.primary_key.get_expressions_or_columns()
                )
            )
        if engine.sample:
            text += ' SAMPLE BY {0}\n'.format(
                self._compile_param(
                    engine.sample.get_expressions_or_columns()[0]
                )
            )
        if engine.settings:
            text += ' SETTINGS ' + ', '.join(
                '{key}={value}'.format(
                    key=key,
                    value=value
                )
                for key, value in sorted(engine.settings.items())
            )
        return text

    def visit_engine(self, engine):
        engine_params = engine.get_parameters()
        text = engine.name
        if not engine_params:
            return text

        text += '('

        compiled_params = []
        for param in engine_params:
            if isinstance(param, tuple):
                compiled = (
                    '('
                    + ', '.join(self._compile_param(p) for p in param)
                    + ')'
                )
            else:
                compiled = self._compile_param(param)

            compiled_params.append(compiled)

        text += ', '.join(compiled_params)

        return text + ')'

    def create_table_suffix(self, table):
        cluster = table.dialect_options['clickhouse']['cluster']
        if cluster:
            return 'ON CLUSTER ' + self.preparer.quote(cluster)

        return ''

    def post_create_table(self, table):
        engine = getattr(table, 'engine', None)

        if not engine:
            raise exc.CompileError("No engine for table '%s'" % table.name)

        return ' ENGINE = ' + self.process(engine)

    def visit_drop_table(self, drop):
        text = '\nDROP TABLE '

        if drop.if_exists:
            text += 'IF EXISTS '

        rv = text + self.preparer.format_table(drop.element)

        if drop.on_cluster:
            rv += (
                ' ON CLUSTER ' +
                self.preparer.quote(drop.on_cluster)
            )

        return rv


class ClickHouseTypeCompiler(compiler.GenericTypeCompiler):
    def visit_string(self, type_, **kw):
        if type_.length is None:
            return 'String'
        else:
            return 'FixedString(%s)' % type_.length

    def visit_array(self, type_, **kw):
        item_type = type_api.to_instance(type_.item_type)
        return 'Array(%s)' % self.process(item_type, **kw)

    def visit_nullable(self, type_, **kw):
        nested_type = type_api.to_instance(type_.nested_type)
        return 'Nullable(%s)' % self.process(nested_type, **kw)

    def visit_lowcardinality(self, type_, **kw):
        nested_type = type_api.to_instance(type_.nested_type)
        return "LowCardinality(%s)" % self.process(nested_type, **kw)

    def visit_int8(self, type_, **kw):
        return 'Int8'

    def visit_uint8(self, type_, **kw):
        return 'UInt8'

    def visit_int16(self, type_, **kw):
        return 'Int16'

    def visit_uint16(self, type_, **kw):
        return 'UInt16'

    def visit_int32(self, type_, **kw):
        return 'Int32'

    def visit_uint32(self, type_, **kw):
        return 'UInt32'

    def visit_int64(self, type_, **kw):
        return 'Int64'

    def visit_uint64(self, type_, **kw):
        return 'UInt64'

    def visit_bool(self, type_, **kw):
        return 'Bool'

    def visit_date(self, type_, **kw):
        return 'Date'

    def visit_date32(self, type_, **kw):
        return 'Date32'

    def visit_datetime(self, type_, **kw):
        return 'DateTime'

    def visit_datetime64(self, type_, **kw):
        return 'DateTime64({})'.format(type_.precision)

    def visit_datetimewithtz(self, type_, **kw):
        return "DateTime('{}')".format(type_.tz)

    def visit_datetime64withtz(self, type_, **kw):
        return "DateTime64({}, '{}')".format(type_.precision, type_.tz)

    def visit_float32(self, type_, **kw):
        return 'Float32'

    def visit_float64(self, type_, **kw):
        return 'Float64'

    def visit_numeric(self, type_, **kw):
        return 'Decimal(%s, %s)' % (type_.precision, type_.scale)

    def visit_nested(self, nested, **kwargs):
        ddl_compiler = self.dialect.ddl_compiler(self.dialect, None)
        cols_create = [
            ddl_compiler.visit_create_column(CreateColumn(nested_child))
            for nested_child in nested.columns
        ]
        return 'Nested(%s)' % ', '.join(cols_create)

    def _render_enum(self, db_type, type_, **kw):
        choices = (
            "'%s' = %d" %
            (x.name.replace("'", r"\'"), x.value) for x in type_.enum_class
        )
        return '%s(%s)' % (db_type, ', '.join(choices))

    def visit_enum(self, type_, **kw):
        return self._render_enum('Enum', type_, **kw)

    def visit_enum8(self, type_, **kw):
        return self._render_enum('Enum8', type_, **kw)

    def visit_enum16(self, type_, **kw):
        return self._render_enum('Enum16', type_, **kw)

    def visit_uuid(self, type_, **kw):
        return 'UUID'

    def visit_ipv4(self, type_, **kw):
        return 'IPv4'

    def visit_ipv6(self, type_, **kw):
        return 'IPv6'


class ClickHouseExecutionContextBase(default.DefaultExecutionContext):
    @sa_util.memoized_property
    def should_autocommit(self):
        return False  # No DML supported, never autocommit

    def _mangle_parameters(self):
        mapping = getattr(self.compiled, '_param_replace_map', None)
        if not mapping:
            return
        if not self.parameters:
            return
        multiparam = False
        parameters = self.parameters
        if isinstance(parameters, tuple):
            multiparam = True
        else:
            parameters = (parameters,)
        parameters = tuple(
            {mapping.get(key, key): val
             for key, val in params_item.items()}
            for params_item in parameters)
        if not multiparam:
            assert len(parameters) == 1
            parameters = parameters[0]
        self.parameters = parameters

    def pre_exec(self):
        self._mangle_parameters()


class ClickHouseDialect(default.DefaultDialect):
    name = 'clickhouse'
    supports_cast = True
    supports_unicode_statements = True
    supports_unicode_binds = True
    supports_sane_rowcount = False
    supports_sane_multi_rowcount = False
    supports_native_decimal = True
    supports_native_boolean = False
    supports_alter = True
    supports_sequences = False
    supports_native_enum = True  # Do not render check constraints on enums.
    supports_multivalues_insert = True

    max_identifier_length = 127
    default_paramstyle = 'pyformat'
    colspecs = colspecs
    ischema_names = ischema_names
    convert_unicode = True
    returns_unicode_strings = sqltypes.String.RETURNS_UNICODE
    description_encoding = None
    postfetch_lastrowid = False

    forced_server_version_string = None

    preparer = ClickHouseIdentifierPreparer
    type_compiler = ClickHouseTypeCompiler
    statement_compiler = ClickHouseCompiler
    ddl_compiler = ClickHouseDDLCompiler
    inspector = ClickHouseInspector

    construct_arguments = [
        (schema.Table, {
            'data': [],
            'cluster': None,
        })
    ]

    # Internal parameter:
    # Python enums will save the numeric value,
    # but then don't allow empty-string value, while clickhouse does.
    use_python_enum = False

    def _execute(self, connection, sql):
        raise NotImplementedError

    @reflection.cache
    def get_view_names(self, connection, schema=None, **kw):
        return self.get_table_names(connection, schema, **kw)

    def has_table(self, connection, table_name, schema=None):
        query = 'EXISTS TABLE {}'.format(self._quote_table_name(table_name))
        for r in self._execute(connection, query):
            if r.result == 1:
                return True
        return False

    def reflecttable(self, *args, **kwargs):
        "See reflect_table. This method name is deprecated"
        return self.reflect_table(*args, **kwargs)

    def reflect_table(
            self, connection, table, include_columns, exclude_columns,
            # Temporary(-ish): `*args` because there's a `resolve_fks` argument
            # added in sqlalchemy 1.3 but not present in sqlalchemy 1.2
            *args,
            **opts):
        """
        Hack to ensure the autoloaded table class is
        `clickhouse_sqlalchemy.Table`
        (to support CH-specific features e.g. joins).
        """
        # This check is necessary to support direct instantiation of
        # `clickhouse_sqlalchemy.Table` and then reflection of it.
        if not isinstance(table, Table):
            table.metadata.remove(table)
            ch_table = Table._make_from_standard(
                table, _extend_on=opts.get('_extend_on')
            )
        else:
            ch_table = table
        return super(ClickHouseDialect, self).reflect_table(
            connection, ch_table, include_columns, exclude_columns,
            *args, **opts)

    def _quote_table_name(self, table_name):
        # Use case: `describe table (select ...)`, over a TextClause.
        if isinstance(table_name, elements.TextClause):
            return str(table_name)
        return self.identifier_preparer.quote_identifier(table_name)

    # @reflection.cache  # XXXX: not clear why it has to be disabled.
    def get_columns(self, connection, table_name, schema=None, **kw):
        query = 'DESCRIBE TABLE {}'.format(self._quote_table_name(table_name))
        rows = self._execute(connection, query)

        return [self._get_column_info(row.name, row.type) for row in rows]

    def _get_column_info(self, name, format_type):
        col_type = self._get_column_type(name, format_type)
        result = {
            'name': name,
            'type': col_type,
            'default': None,
            'nullable': False,
            'lowcardinality': False,
        }

        # Hack to unwrap some wrap-types into flags;
        # mostly done because in most other databases `nullable` is a flag
        # rather than a wrapper.
        while True:  # assumes `.nested_type` chain is not a loop
            if isinstance(col_type, types.Nullable):
                result['nullable'] = True
                col_type = col_type.nested_type
            elif isinstance(col_type, types.LowCardinality):
                result['lowcardinality'] = True
                col_type = col_type.nested_type
            else:
                break

        result['type'] = col_type

        return result

    def _get_column_type(self, name, spec):
        # Try to parse a parametrized/wrapped type.

        if spec.startswith('Array'):
            inner = spec[6:-1]
            coltype = self.ischema_names['_array']
            return coltype(self._get_column_type(name, inner))

        elif spec.startswith('FixedString'):
            length = int(spec[12:-1])
            return self.ischema_names['FixedString'](length)

        elif spec.startswith('Nullable'):
            inner = spec[9:-1]
            coltype = self.ischema_names['_nullable']
            return coltype(self._get_column_type(name, inner))

        elif spec.startswith('LowCardinality'):
            inner = spec[15:-1]
            coltype = self.ischema_names['_lowcardinality']
            return coltype(self._get_column_type(name, inner))

        elif spec.startswith('Enum'):
            pos = spec.find('(')
            type = spec[:pos]
            coltype = self.ischema_names[type]

            options = dict()
            if pos >= 0:
                options = self._parse_options(
                    spec[pos + 1: spec.rfind(')')]
                )
            if not options:
                return sqltypes.NullType

            if not self.use_python_enum:
                # python enums can't accept empty string as enum item name.
                return lambda: coltype(*options.keys(), name=f'{name}_enum')
            type_enum = enum.Enum('%s_enum' % name, options)
            return lambda: coltype(type_enum)

        elif spec.lower().startswith('decimal'):
            coltype = self.ischema_names['Decimal']
            return coltype(*self._parse_decimal_params(spec))

        elif spec.lower().startswith('datetime('):
            coltype = self.ischema_names['_datetimewithtz']
            tz = spec[spec.find('(') + 1:spec.rfind(')')]
            tz = tz.strip().strip("'")
            return coltype(tz)
        elif spec.lower().startswith('datetime64('):
            # DateTime64(precision [, timezone])
            dt64_arg_str = spec[spec.find('(') + 1:spec.rfind(')')]
            if ',' in dt64_arg_str:
                raw_precision, raw_tz = dt64_arg_str.split(',')
                tz = raw_tz.strip().strip("'")
                coltype = self.ischema_names['_datetime64withtz']
                return coltype(int(raw_precision), tz)
            else:
                coltype = self.ischema_names['_datetime64']
                return coltype(int(dt64_arg_str))

        else:
            try:
                return self.ischema_names[spec]
            except KeyError:
                warn("Did not recognize type '%s' of column '%s'" %
                     (spec, name))
                return sqltypes.NullType

    @staticmethod
    def _parse_decimal_params(spec):
        ints = spec.split('(')[-1].split(')')[0]  # get all data in brackets
        precision, scale = ints.split(',')
        return int(precision.strip()), int(scale.strip())

    @staticmethod
    def _parse_options(option_string):
        options = dict()
        after_name = False
        escaped = False
        quote_character = None
        name = ''
        value = ''

        for ch in option_string:
            if escaped:
                name += ch
                escaped = False  # Accepting escaped character

            elif after_name:
                if ch in (' ', '='):
                    pass
                elif ch == ',':
                    options[name] = int(value)
                    after_name = False
                    name = ''
                    value = ''  # Reset before collecting new option
                else:
                    value += ch

            elif quote_character:
                if ch == '\\':
                    escaped = True
                elif ch == quote_character:
                    quote_character = None
                    after_name = True  # Start collecting option value
                else:
                    name += ch

            else:
                if ch == "'":
                    quote_character = ch

        if after_name:
            options.setdefault(name, int(value))  # Word after last comma

        return options

    @reflection.cache
    def get_schema_names(self, connection, **kw):
        rows = self._execute(connection, 'SHOW DATABASES')
        return [row.name for row in rows]

    @reflection.cache
    def get_foreign_keys(self, connection, table_name, schema=None, **kw):
        # No support for foreign keys.
        return []

    @reflection.cache
    def get_pk_constraint(self, connection, table_name, schema=None, **kw):
        # No support for primary keys.
        return []

    @reflection.cache
    def get_indexes(self, connection, table_name, schema=None, **kw):
        # No support for indexes.
        return []

    @reflection.cache
    def get_table_names(self, connection, schema=None, **kw):
        rows = self._execute(connection, 'SHOW TABLES')
        return [row.name for row in rows]

    def do_rollback(self, dbapi_connection):
        # No support for transactions.
        pass

    def do_executemany(self, cursor, statement, parameters, context=None):
        cursor.executemany(statement, parameters, context=context)

    def do_execute(self, cursor, statement, parameters, context=None):
        cursor.execute(statement, parameters, context=context)

    def _check_unicode_returns(self, connection, additional_tests=None):
        return True

    def _check_unicode_description(self, connection):
        return True

    def _get_server_version_info(self, connection):
        version = self.forced_server_version_string

        if version is None:
            version = self._query_server_version_string(connection)
            assert version

        return tuple(int(part) for part in version.split('.'))

    def _query_server_version_string(self, connection):
        raise NotImplementedError

    def connect(self, *cargs, **cparams):
        self.forced_server_version_string = cparams.pop(
            'server_version', self.forced_server_version_string)
        return super(ClickHouseDialect, self).connect(*cargs, **cparams)

    def format_string(self, connection):
        # `connection` is, usually, a `sqlalchemy.engine.base.Connection`
        # object;
        # `connection.connection` is e.g. a
        # `clickhouse_sqlalchemy.drivers.http.connector.Connection` object
        fmt_info = None
        if isinstance(connection, sa.engine.base.Connection):
            fmt_info = getattr(connection.connection, 'format_info', None)
        if fmt_info is not None:
            return fmt_info.format_string
        return 'TabSeparatedWithNamesAndTypes'

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/drivers/base.py
datalens-tech/datalens-backend

import sqlalchemy as sa
from ...util.compat import string_types
from ..base import ClickHouseDialect, ClickHouseExecutionContextBase
from . import connector


# Export connector version
VERSION = (0, 0, 2, None)


class ClickHouseExecutionContext(ClickHouseExecutionContextBase):

    def _add_statement_format(self):
        # TODO: refactor
        if not self.isinsert and not self.isddl:
            self.statement = '{} FORMAT {}'.format(
                self.statement,
                self.dialect.format_string(self.root_connection))

    def pre_exec(self):
        self._add_statement_format()
        super(ClickHouseExecutionContext, self).pre_exec()


class ClickHouseDialect_http(ClickHouseDialect):
    driver = 'http'
    execution_ctx_cls = ClickHouseExecutionContext

    @classmethod
    def dbapi(cls):
        return connector

    def create_connect_args(self, url):
        kwargs = {}
        protocol = url.query.get('protocol', 'http')
        port = url.port or 8123
        db_name = url.database or 'default'
        endpoint = url.query.get('endpoint', '')

        kwargs.update(url.query)

        db_url = '%s://%s:%d/%s' % (protocol, url.host, port, endpoint)

        return (db_url, db_name, url.username, url.password), kwargs

    def _execute(self, connection, sql):
        if isinstance(sql, string_types):
            # Makes sure the query will go through the
            # `ClickHouseExecutionContext` logic.
            sql = sa.sql.elements.TextClause(sql.replace(':', r'\:'))
        return connection.execute(sql)

    def _query_server_version_string(self, connection):
        query = 'select version()'
        return connection.scalar(query)


dialect = ClickHouseDialect_http

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/drivers/http/base.py
datalens-tech/datalens-backend
from uuid import uuid1

from ...parsers import FORMATS, Parts

from .escaper import Escaper
from .transport import RequestsTransport

# PEP 249 module globals
apilevel = '2.0'
# Threads may share the module and connections.
threadsafety = 2
# Python extended format codes, e.g. ...WHERE name=%(name)s
paramstyle = 'pyformat'


class Error(Exception):
    """
    Exception that is the base class of all other error exceptions.
    You can use this to catch all errors with one single except statement.
    """
    pass


def connect(*args, **kwargs):
    """
    Make new connection.
    """
    return Connection(*args, **kwargs)


class Connection(object):
    transport_cls = RequestsTransport

    def __init__(self, *args, **kwargs):
        format_name = kwargs.pop('format', 'default')
        self.format_info = FORMATS[format_name]()
        kwargs.update(format_info=self.format_info)

        stream = bool(kwargs.pop('stream', None))

        self._prefetch = not stream

        # TODO: support `stream` argument in the transport.
        self.transport = self.transport_cls(*args, **kwargs)
        super(Connection, self).__init__()

    def close(self):
        pass

    def commit(self):
        pass

    def rollback(self):
        pass

    def cursor(self):
        return Cursor(self)


class Cursor(object):
    """
    These objects represent a database cursor, which is used to manage
    the context of a fetch operation.

    Cursors are not isolated, i.e., any changes done to the database
    by a cursor are immediately visible by other cursors or connections.
    """
    class States(object):
        (
            NONE,
            RUNNING,
            FINISHED
        ) = range(3)

    _states = States()

    _params_escaper = Escaper()

    _columns = None
    _types = None
    # Result data iterable:
    _response = None
    # Result data prefetch cache:
    _rows = None

    _schema = None
    _totals = None
    _stats = None

    def __init__(self, connection):
        self._connection = connection
        self._reset_state()
        self._arraysize = 1
        super(Cursor, self).__init__()

    @property
    def rowcount(self):
        # TODO: initial len(self._rows)?
        return -1

    @property
    def description(self):
        columns = self._columns or []
        types = self._types or []

        return [
            (name, type_code, None, None, None, None, True)
            for name, type_code in zip(columns, types)
        ]

    def close(self):
        pass

    def execute(self, operation, parameters=None, context=None):
        raw_sql = operation

        if parameters is not None:
            raw_sql = raw_sql % self._params_escaper.escape(parameters)

        self._reset_state()
        self._begin_query()

        transport = self._connection.transport
        params = {'query_id': self._query_id}

        execute_ext = getattr(transport, 'execute_ext', None)
        if execute_ext is not None:
            response_gen = execute_ext(raw_sql, params=params)
            self._process_response_ext(response_gen)
        else:
            response_gen = transport.execute(raw_sql, params=params)
            self._process_response(response_gen)

        self._end_query()

    def executemany(self, operation, seq_of_parameters, context=None):
        index = operation.index('VALUES') + 7
        values_tpl = operation[index:]
        params = ', '.join(
            values_tpl % self._params_escaper.escape(params)
            for params in seq_of_parameters
        )
        self.execute(operation[:index] + params)

    def check_query_started(self):
        if self._state == self._states.NONE:
            raise RuntimeError("No query yet")

    def fetchone(self):
        self.check_query_started()

        # `self._prefetch` case:
        if self._rows is not None:
            if not self._rows:
                return None
            return self._rows.pop(0)

        return next(self._response, None)

    def fetchmany(self, size=1):
        self.check_query_started()

        # `self._prefetch` case:
        if self._rows is not None:
            rv = self._rows[:size]
            self._rows = self._rows[size:]
            return rv

        rows = []
        for _ in range(size):
            row = self.fetchone()
            if row is None:
                break
            rows.append(row)
        return rows

    def fetchall(self):
        self.check_query_started()

        # `self._prefetch` case:
        if self._rows is not None:
            rv = self._rows
            self._rows = []
            return rv

        rows = []
        while True:
            row = self.fetchone()
            if row is None:
                break
            rows.append(row)
        return rows

    @property
    def arraysize(self):
        return self._arraysize

    @arraysize.setter
    def arraysize(self, value):
        self._arraysize = value

    def setinputsizes(self, sizes):
        pass

    def setoutputsize(self, size, column=None):
        pass

    # Private and non-standard methods.
    def cancel(self):
        """
        Cancels query. Not in PEP 249 standard.
        """
        if self._state == self._states.NONE or self._query_id is None:
            raise RuntimeError("No query yet")

        # Try to cancel query by sending query with the same query_id.
        transport = self._connection.transport
        params = {'query_id': self._query_id}
        transport.execute('SELECT 1', params=params)

        self._end_query()
        self._query_id = None
        self._rows = None

    @property
    def ch_schema(self):
        return self._schema

    @property
    def ch_totals(self):
        return self._totals

    @property
    def ch_stats(self):
        return self._stats

    @property
    def _prefetch(self):
        return self._connection._prefetch

    def _process_response_ext(self, response):
        try:
            part, item = next(response)
        except StopIteration:
            raise Exception("Not a well-behaved extended response: empty")
        assert part == Parts.META
        self._schema = item
        self._columns = [col['name'] for col in item]
        self._types = [col['type'] for col in item]

        self._response = self._wrap_response_ext(response)
        if self._prefetch:
            self._rows = list(self._response)

    def _wrap_response_ext(self, response):
        for part, item in response:
            if part == Parts.DATACHUNK:
                for row in item:
                    yield row
            elif part == Parts.DATA:
                yield item
            elif part == Parts.TOTALS:
                self._totals = item
            elif part == Parts.STATS:
                self._stats = item
                self._totals = self._totals or item.get('totals')
            elif part == Parts.META:
                raise Exception((
                    "Not a well-behaved extended response: "
                    "unexpected meta piece"))

    def _process_response(self, response):
        response = iter(response)

        self._columns = next(response, None)
        self._types = next(response, None)
        self._response = response

        if self._prefetch:
            self._rows = list(self._response)

    def _reset_state(self):
        """
        Resets query state and get ready for another query.
        """
        self._query_id = None
        self._state = self._states.NONE

        self._columns = None
        self._types = None
        self._rows = None
        self._response = None

        self._schema = None
        self._totals = None
        self._stats = None

    def _begin_query(self):
        self._state = self._states.RUNNING
        self._query_id = uuid1()

    def _end_query(self):
        self._state = self._states.FINISHED

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/drivers/http/connector.py
datalens-tech/datalens-backend
"""
...

Note: this should probably be entirely replaced with `...quoting.QUOTER`.
"""
from datetime import date, datetime
from decimal import Decimal
import enum

from ...util import compat


class Escaper(object):

    number_types = compat.integer_types
    string_types = compat.string_types

    escape_chars = {
        "\b": "\\b",
        "\f": "\\f",
        "\r": "\\r",
        "\n": "\\n",
        "\t": "\\t",
        "\0": "\\0",
        "\\": "\\\\",
        "'": "\\'"
    }

    def escape_string(self, value):
        value = ''.join(self.escape_chars.get(c, c) for c in value)
        return "'" + value + "'"

    def escape(self, parameters):
        if isinstance(parameters, dict):
            return {k: self.escape_item(v) for k, v in parameters.items()}
        elif isinstance(parameters, (list, tuple)):
            return "[" + ",".join(
                [str(self.escape_item(x)) for x in parameters]) + "]"
        else:
            raise Exception("Unsupported param format: {}".format(parameters))

    def escape_number(self, item):
        return int(item)  # XXX: unclear why `int()` is necessary.

    def escape_date(self, item):
        # XXX: shouldn't this be `toDate(...)`?
        return self.escape_string(item.strftime('%Y-%m-%d'))

    def escape_datetime(self, item):
        # XXX: shouldn't this be `toDateTime(...)`?
        return self.escape_string(item.strftime('%Y-%m-%d %H:%M:%S'))

    def escape_decimal(self, item):
        return float(item)

    def escape_item(self, item):
        if item is None:
            return 'NULL'
        elif isinstance(item, self.number_types):
            return self.escape_number(item)
        elif isinstance(item, datetime):
            return self.escape_datetime(item)
        elif isinstance(item, date):
            return self.escape_date(item)
        elif isinstance(item, (Decimal, float)):
            return self.escape_decimal(item)
        elif isinstance(item, self.string_types):
            return self.escape_string(item)
        elif isinstance(item, (list, tuple)):
            return "[" + ", ".join(
                [str(self.escape_item(x)) for x in item]
            ) + "]"
        elif isinstance(item, enum.Enum):
            return self.escape_string(item.name)
        else:
            raise Exception("Unsupported object {}".format(item))

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/drivers/http/escaper.py
datalens-tech/datalens-backend

import itertools

import requests

from clickhouse_sqlalchemy.types import get_http_postprocess
from ...exceptions import DatabaseException
from .exceptions import HTTPException


DEFAULT_DDL_TIMEOUT = None
DEBUG_STREAMS = False


# TODO: commonize this with the schema-parsing logic.
def _get_type(type_str):
    from .base import ClickHouseDialect_http
    dialect = ClickHouseDialect_http()
    coltype = dialect._get_column_type('_', type_str)
    return get_http_postprocess(coltype)


class RequestsTransport(object):

    def __init__(
            self,
            db_url, db_name, username, password,
            timeout=None, ch_settings=None, verify=None,
            format_info=None,
            **kwargs):

        self.db_url = db_url
        self.db_name = db_name
        self.auth = None if username is None and password is None else (username, password)
        if format_info is None:
            raise TypeError('format_info is required here')
        if (format_info.generator_from_lines is None and
                format_info.generator is None):
            raise TypeError(
                ('formats here must have `generator_from_lines` '
                 'or `generator`'))
        self.format_info = format_info

        self.timeout = None
        connect_timeout = kwargs.get('connect_timeout')
        if timeout is not None:
            if connect_timeout is not None:
                self.timeout = (float(connect_timeout), float(timeout))
            else:
                self.timeout = float(timeout)

        self.verify = verify

        self.headers = {
            key[8:]: value
            for key, value in kwargs.items()
            if key.startswith('header__')
        }

        self.statuses_to_retry = (
            kwargs.get('statuses_to_retry', '502,504')
            .split(','))

        ch_settings = dict(ch_settings or {})
        self.ch_settings = ch_settings

        ddl_timeout = kwargs.pop('ddl_timeout', DEFAULT_DDL_TIMEOUT)
        if ddl_timeout is not None:
            self.ch_settings['distributed_ddl_task_timeout'] = int(ddl_timeout)

        super(RequestsTransport, self).__init__()

    def execute_ext(
            self, query, params=None,
            allow_chunks=False, strict=False):
        """
        Query is returning rows and these rows should be parsed or
        there is nothing to return.
        """
        resp = self._send(query, params=params, stream=True)
        fmt = self.format_info

        if fmt.generator_from_lines:
            chunks = resp.iter_lines()
            parser = fmt.generator_from_lines
        else:
            chunks = resp.iter_content(fmt.chunk_size)
            parser = fmt.generator

        if DEBUG_STREAMS:
            def _dbg_chunks(chunks):
                for chunk in chunks:
                    print(' >>> output >>> {!r}'.format(chunk))
                    yield chunk

            chunks = _dbg_chunks(chunks)

        try:
            chunk = next(chunks)
        except StopIteration:
            # Empty result; e.g. a DDL request.
            if not strict:
                yield fmt.parts.META, []
                return
            raise Exception("A not well-behaving parser", fmt)

        # Put it back:
        chunks = itertools.chain([chunk], chunks)
        results = parser(chunks)

        try:
            part, meta = next(results)
        except StopIteration:
            # Warning: this is more dubious;
            # sample case: 'drop table ... on cluster ...',
            # without 'FORMAT ...',
            # with exactly one host in the cluster.
            # chunks example: [b'127.0.0.1\t9000\t0\t\t0\t0']
            if not strict:
                yield fmt.parts.META, []
                return
            raise Exception("A not well-behaving parser", fmt)
        assert part == fmt.parts.META

        # TODO: move this to a different layer.
        convs = [_get_type(field['type']) for field in meta]

        def postprocess(row):
            assert len(row) == len(convs)
            return [
                converter(value) if converter is not None else value
                for value, converter in zip(row, convs)]

        yield fmt.parts.META, meta

        for part, item in results:
            if part == fmt.parts.DATACHUNK:
                item = [postprocess(row) for row in item]
                if allow_chunks:
                    yield part, item
                else:
                    for row in item:
                        yield fmt.parts.DATA, row
            elif part in (fmt.parts.DATA, fmt.parts.TOTALS):
                item = postprocess(item)
                yield part, item
            else:
                yield part, item

    def execute(self, query, params=None, **kwargs):
        fmt = self.format_info
        results = self.execute_ext(query, params=params, **kwargs)
        try:
            part, meta = next(results)
        except StopIteration:
            raise Exception("A not well-behaving `execute_ext`")
        assert part == fmt.parts.META

        yield [field['name'] for field in meta]
        yield [field['type'] for field in meta]
        for part, item in results:
            if part == fmt.parts.DATA:
                yield item

    def raw(self, query, params=None, stream=False):
        """
        Performs raw query to database. Returns its output
        :param query: Query to execute
        :param params: Additional params should be passed during query.
        :param stream: If flag is true, Http response from ClickHouse will be
            streamed.
        :return: Query execution result
        """
        return self._send(query, params=params, stream=stream).text

    def _send(self, data, params=None, stream=False):

        if DEBUG_STREAMS:
            print(' >>> input >>> {!r}'.format(data))

        query = data
        data = data.encode('utf-8')
        params = params or {}

        params['database'] = self.db_name
        params.update(self.ch_settings)

        session = requests.Session()

        headers = {
            key: (
                # On-before-send dynamic header values for e.g. tracing / timing.
                # `None` values should be ignored by `requests`.
                value(query=query, params=params, stream=stream)
                if callable(value)
                else value)
            for key, value in self.headers.items()}
        try:
            resp = session.post(
                self.db_url, auth=self.auth, params=params, data=data,
                stream=stream, timeout=self.timeout, headers=headers,
                allow_redirects=False, verify=self.verify,
            )
        except requests.exceptions.RequestException as exc:
            raise DatabaseException(exc)

        if resp.status_code != 200:
            orig = HTTPException(resp.text)
            orig.code = resp.status_code
            raise DatabaseException(orig)

        return resp

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/drivers/http/transport.py
datalens-tech/datalens-backend

from . import connector
from ..base import (
    ClickHouseDialect, ClickHouseExecutionContextBase, ClickHouseCompiler,
)

# Export connector version
VERSION = (0, 0, 2, None)


class ClickHouseExecutionContext(ClickHouseExecutionContextBase):
    def pre_exec(self):
        # Always do executemany on INSERT with VALUES clause.
        if self.isinsert and self.compiled.statement.select is None:
            self.executemany = True


class ClickHouseNativeCompiler(ClickHouseCompiler):

    def visit_insert(self, insert_stmt, asfrom=False, **kw):
        rv = super(ClickHouseNativeCompiler, self).visit_insert(
            insert_stmt, asfrom=asfrom, **kw)

        pos = rv.rfind('VALUES (')
        # Remove (%s)-templates from VALUES clause if exists.
        # ClickHouse server since version 19.3.3 parse query after VALUES and
        # allows inplace parameters.
        # Example: INSERT INTO test (x) VALUES (1), (2).
        if pos != -1:
            rv = rv[:pos + 6]
        return rv


class ClickHouseDialect_native(ClickHouseDialect):

    driver = 'native'
    execution_ctx_cls = ClickHouseExecutionContext
    statement_compiler = ClickHouseNativeCompiler

    @classmethod
    def dbapi(cls):
        return connector

    def create_connect_args(self, url):
        url = url.set(drivername='clickhouse')

        return (str(url), ), {}

    def _execute(self, connection, sql):
        return connection.execute(sql)

    def _query_server_version_string(self, connection):
        return connection.scalar('select version()')


dialect = ClickHouseDialect_native

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/drivers/native/base.py
datalens-tech/datalens-backend

import datetime

import pytz


DATE_NULL = '0000-00-00'
DATETIME_NULL = '0000-00-00 00:00:00'

# timezones missing in `pytz`
TZ_PATCHES = {
    # Technically, this is non-daylight-savings timezone only,
    # with 'MSD' being daylight-savings timezone,
    # but in practice, it is used interchangeably.
    'MSK': 'Europe/Moscow',
}


def parse_timezone(tz):
    tz = TZ_PATCHES.get(tz, tz)
    try:
        return pytz.timezone(tz)
    except pytz.UnknownTimeZoneError:
        raise Exception('Unknown timezone from CH', tz)


def parse_date(value):
    if value is None or value == DATE_NULL:
        return None
    return datetime.datetime.strptime(value, '%Y-%m-%d').date()


def parse_datetime(value, tzinfo=None):
    if value is None or value == DATETIME_NULL:
        return None
    result = datetime.datetime.strptime(value, '%Y-%m-%d %H:%M:%S')
    if tzinfo is not None:
        result = tzinfo.localize(result)
    return result


def parse_datetime64(value, tzinfo=None):
    if value is None or value == DATETIME_NULL:
        return None

    # CH datetime64 contains up to 9 digits after dot, so let's cut it to microseconds
    # len('2020-01-01 12:34:56.123456') == 26
    # maybe TODO: use pandas datetime64 type for nanoseconds handling
    result = datetime.datetime.fromisoformat(value[:26])

    if tzinfo is not None:
        result = tzinfo.localize(result)
    return result

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/dt_utils.py
datalens-tech/datalens-backend
from sqlalchemy import util, literal
from sqlalchemy.sql import (
    ClauseElement, coercions, roles
)
from sqlalchemy.sql.base import SchemaEventTarget
from sqlalchemy.sql.schema import ColumnCollectionMixin, SchemaItem
from sqlalchemy.sql.visitors import Visitable
from sqlalchemy.util import to_list


class Engine(SchemaEventTarget, Visitable):
    __visit_name__ = 'engine'

    def get_parameters(self):
        return []

    @property
    def name(self):
        return self.__class__.__name__

    def _set_parent(self, table, **kwargs):
        self.table = table
        self.table.engine = self


class TableCol(ColumnCollectionMixin, SchemaItem):
    def __init__(self, column, **kwargs):
        super(TableCol, self).__init__(*[column], **kwargs)

    def get_column(self):
        return list(self.columns)[0]


class KeysExpressionOrColumn(ColumnCollectionMixin, SchemaItem):
    def __init__(self, *expressions, **kwargs):
        columns = []
        self.expressions = []
        for expr, column, strname, add_element \
                in coercions.expect_col_expression_collection(roles.DDLConstraintColumnRole, expressions):
            if add_element is not None:
                columns.append(add_element)
            self.expressions.append(expr)

        super(KeysExpressionOrColumn, self).__init__(*columns, **kwargs)

    def _set_parent(self, table, **kwargs):
        super(KeysExpressionOrColumn, self)._set_parent(table)

    def get_expressions_or_columns(self):
        expr_columns = util.zip_longest(self.expressions, self.columns)
        return [
            (expr if isinstance(expr, ClauseElement) else colexpr)
            for expr, colexpr in expr_columns
        ]


class MergeTree(Engine):

    __visit_name__ = 'merge_tree'

    def __init__(
            self,
            partition_by=None,
            order_by=None,
            primary_key=None,
            sample=None,
            **settings
    ):
        self.partition_by = None
        if partition_by is not None:
            self.partition_by = KeysExpressionOrColumn(partition_by)

        self.order_by = None
        if order_by is not None:
            self.order_by = KeysExpressionOrColumn(*to_list(order_by))

        self.primary_key = None
        if primary_key is not None:
            self.primary_key = KeysExpressionOrColumn(*to_list(primary_key))

        self.sample = None
        if sample is not None:
            self.sample = KeysExpressionOrColumn(sample)
        self.settings = settings
        super(MergeTree, self).__init__()

    def _set_parent(self, table, **kwargs):
        super(MergeTree, self)._set_parent(table, **kwargs)
        if self.partition_by is not None:
            self.partition_by._set_parent(table)
        if self.order_by is not None:
            self.order_by._set_parent(table)
        if self.primary_key is not None:
            self.primary_key._set_parent(table)
        if self.sample is not None:
            self.sample._set_parent(table)

    def get_parameters(self):
        return []


class AggregatingMergeTree(MergeTree):
    pass


class GraphiteMergeTree(MergeTree):

    def __init__(self, config_name, *args, **kwargs):
        super(GraphiteMergeTree, self).__init__(
            *args,
            **kwargs
        )
        self.config_name = config_name

    def get_parameters(self):
        return literal(self.config_name)


class CollapsingMergeTree(MergeTree):
    def __init__(self, sign_col, *args, **kwargs):
        super(CollapsingMergeTree, self).__init__(
            *args, **kwargs
        )
        self.sign_col = TableCol(sign_col)

    def get_parameters(self):
        return self.sign_col.get_column()

    def _set_parent(self, table, **kwargs):
        super(CollapsingMergeTree, self)._set_parent(table, **kwargs)

        self.sign_col._set_parent(table)


class SummingMergeTree(MergeTree):
    def __init__(self, *args, **kwargs):
        summing_cols = kwargs.pop('columns', None)
        super(SummingMergeTree, self).__init__(
            *args, **kwargs
        )

        self.summing_cols = None
        if summing_cols is not None:
            self.summing_cols = KeysExpressionOrColumn(*summing_cols)

    def _set_parent(self, table, **kwargs):
        super(SummingMergeTree, self)._set_parent(table, **kwargs)

        if self.summing_cols is not None:
            self.summing_cols._set_parent(table)

    def get_parameters(self):
        if self.summing_cols is not None:
            return self.summing_cols.get_expressions_or_columns()


class ReplacingMergeTree(MergeTree):
    def __init__(self, *args, **kwargs):
        version_col = kwargs.pop('ver', None)
        super(ReplacingMergeTree, self).__init__(
            *args, **kwargs
        )

        self.version_col = None
        if version_col is not None:
            self.version_col = TableCol(version_col)

    def _set_parent(self, table, **kwargs):
        super(ReplacingMergeTree, self)._set_parent(table, **kwargs)

        if self.version_col is not None:
            self.version_col._set_parent(table)

    def get_parameters(self):
        if self.version_col is not None:
            return self.version_col.get_column()


class Distributed(Engine):
    def __init__(self, logs, default, hits, sharding_key=None):
        self.logs = logs
        self.default = default
        self.hits = hits
        self.sharding_key = sharding_key
        super(Distributed, self).__init__()

    def get_parameters(self):
        params = [
            self.logs,
            self.default,
            self.hits
        ]

        if self.sharding_key is not None:
            params.append(self.sharding_key)

        return params


class ReplicatedEngineMixin(object):
    def __init__(self, table_path, replica_name):
        self.table_path = literal(table_path)
        self.replica_name = literal(replica_name)

    def get_parameters(self):
        return [
            self.table_path,
            self.replica_name
        ]


class ReplicatedMergeTree(ReplicatedEngineMixin, MergeTree):
    def __init__(self, table_path, replica_name,
                 *args, **kwargs):
        ReplicatedEngineMixin.__init__(self, table_path, replica_name)
        MergeTree.__init__(self, *args, **kwargs)

    def get_parameters(self):
        return ReplicatedEngineMixin.get_parameters(self) + \
            MergeTree.get_parameters(self)


class ReplicatedCollapsingMergeTree(ReplicatedEngineMixin,
                                    CollapsingMergeTree):
    def __init__(self, table_path, replica_name,
                 *args, **kwargs):
        ReplicatedEngineMixin.__init__(self, table_path, replica_name)
        CollapsingMergeTree.__init__(self, *args, **kwargs)

    def get_parameters(self):
        return ReplicatedEngineMixin.get_parameters(self) + \
            CollapsingMergeTree.get_parameters(self)


class ReplicatedReplacingMergeTree(ReplicatedEngineMixin,
                                   ReplacingMergeTree):
    def __init__(self, table_path, replica_name,
                 *args, **kwargs):
        ReplicatedEngineMixin.__init__(self, table_path, replica_name)
        ReplacingMergeTree.__init__(self, *args, **kwargs)

    def get_parameters(self):
        rv = ReplicatedEngineMixin.get_parameters(self)
        replacing = ReplacingMergeTree.get_parameters(self)
        if replacing is not None:
            rv.append(replacing)
        return rv


class ReplicatedAggregatingMergeTree(ReplicatedEngineMixin,
                                     AggregatingMergeTree):
    def __init__(self, table_path, replica_name,
                 *args, **kwargs):
        ReplicatedEngineMixin.__init__(self, table_path, replica_name)
        AggregatingMergeTree.__init__(self, *args, **kwargs)

    def get_parameters(self):
        return ReplicatedEngineMixin.get_parameters(self) + \
            AggregatingMergeTree.get_parameters(self)


class ReplicatedSummingMergeTree(ReplicatedEngineMixin, SummingMergeTree):
    def __init__(self,
                 table_path,
                 replica_name,
                 *args,
                 **kwargs):
        ReplicatedEngineMixin.__init__(self, table_path, replica_name)
        SummingMergeTree.__init__(self, *args, **kwargs)

    def get_parameters(self):
        rv = ReplicatedEngineMixin.get_parameters(self)
        summing = SummingMergeTree.get_parameters(self)
        if summing is not None:
            rv.append(summing)
        return rv


class Buffer(Engine):
    def __init__(self, database, table, num_layers=16,
                 min_time=10, max_time=100, min_rows=10000, max_rows=1000000,
                 min_bytes=10000000, max_bytes=100000000):
        self.database = database
        self.table_name = table
        self.num_layers = num_layers
        self.min_time = min_time
        self.max_time = max_time
        self.min_rows = min_rows
        self.max_rows = max_rows
        self.min_bytes = min_bytes
        self.max_bytes = max_bytes
        super(Buffer, self).__init__()

    def get_parameters(self):
        return [
            self.database,
            self.table_name,
            self.num_layers,
            self.min_time,
            self.max_time,
            self.min_rows,
            self.max_rows,
            self.min_bytes,
            self.max_bytes
        ]


class _NoParamsEngine(Engine):
    def get_parameters(self):
        return []


class TinyLog(_NoParamsEngine):
    pass


class Log(_NoParamsEngine):
    pass


class Memory(_NoParamsEngine):
    pass


class Null(_NoParamsEngine):
    pass


class File(Engine):
    supported_data_formats = {
        'tabseparated': 'TabSeparated',
        'tabseparatedwithnames': 'TabSeparatedWithNames',
        'tabseparatedwithnamesandtypes': 'TabSeparatedWithNamesAndTypes',
        'template': 'Template',
        'csv': 'CSV',
        'csvwithnames': 'CSVWithNames',
        'customseparated': 'CustomSeparated',
        'values': 'Values',
        'jsoneachrow': 'JSONEachRow',
        'tskv': 'TSKV',
        'protobuf': 'Protobuf',
        'parquet': 'Parquet',
        'rowbinary': 'RowBinary',
        'rowbinarywithnamesandtypes': 'RowBinaryWithNamesAndTypes',
        'native': 'Native'
    }

    def __init__(self, data_format):
        fmt = self.supported_data_formats.get(data_format.lower())
        if not fmt:
            raise ValueError(
                'Format {0} not supproted for File engine'.format(
                    data_format
                )
            )
        self.data_format = fmt
        super(File, self).__init__()

    def get_parameters(self):
        return (self.data_format, )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/engines.py
datalens-tech/datalens-backend

import re

from .util import compat


class DatabaseException(Exception):

    code = None
    exc_type_name = None
    exc_message = None
    exc_message_extended = None
    server_version = None

    def __init__(self, orig):
        self.orig = orig
        self.maybe_parse_orig(orig)
        super(DatabaseException, self).__init__(orig)

    def maybe_parse_orig(self, orig):
        # Wrapped in case `str(orig)` is broken, primarily.
        # More generally, the parsing is best-effort here,
        # and shouldn't break the use-cases that don't rely on it.
        try:
            self.parse_orig(orig)
        except Exception:
            pass

    exc_text_re = (
        # re.IGNORECASE | re.DOTALL; *not* re.MULTILINE
        r'(?is)'
        r'^'
        r'(std::exception\.\s+)?'
        r'Code:\s+(?P<code>\d+)'
        # Everything but code is optional:
        r'(?:, type: (?P<exc_type_name>[^\s]+))?'
        r'(?:'  # exc_message + exc_message_extended
        r', e\.[^\s]+ = (?:DB::Exception: )?(?P<exc_message>[^\n]+)'
        r'(?:\n(?P<exc_message_extended>.*?))?'
        r')?'
    )
    exc_version_regexes = (
        (r'(?P<base>.*) \(version '
         r'(?P<version>[0-9][^\s]+)'
         r'(?: \(official build\))?'
         r'\)\s*$'),
        r'(?P<base>.*), version = (?P<version>[^\s]+)\s*$',
    )

    def _pick_out_version(self, text):
        if not text:
            return text, None
        for rex in self.exc_version_regexes:
            match = re.search(rex, text)
            if match:
                data = match.groupdict()
                return data['base'], data['version']
        return text, None

    def parse_orig(self, orig):
        text = str(orig)
        match = re.search(self.exc_text_re, text)
        if match:
            data = match.groupdict()
            self.code = data['code']
            self.exc_type_name = data['exc_type_name']
            exc_message, version_a = self._pick_out_version(
                data['exc_message'])
            exc_message_extended, version_b = self._pick_out_version(
                data['exc_message_extended'])

            self.exc_message = exc_message
            self.exc_message_extended = exc_message_extended
            self.server_version = version_a or version_b

    def __str__(self):
        text = 'Orig exception: {}'.format(self.orig)

        if compat.PY3:
            return compat.text_type(text)
        else:
            return compat.text_type(text).encode('utf-8')


if __name__ == '__main__':
    pass

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/exceptions.py
datalens-tech/datalens-backend
import dill

from sqlalchemy import util, exc
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql import type_api, functions
from sqlalchemy.sql.elements import (
    BindParameter,
    ColumnElement,
    ClauseList,
    ColumnClause,
    Label,
)
from sqlalchemy.sql.visitors import Visitable


class SampleParam(BindParameter):
    pass


def sample_clause(element):
    """Convert the given value to an "sample" clause.

    This handles incoming element to an expression; if
    an expression is already given, it is passed through.

    """
    if element is None:
        return None
    elif hasattr(element, '__clause_element__'):
        return element.__clause_element__()
    elif isinstance(element, Visitable):
        return element
    else:
        return SampleParam(None, element, unique=True)


class Lambda(ColumnElement):
    """Represent a lambda function, ``Lambda(lambda x: 2 * x)``."""

    __visit_name__ = 'lambda'

    def __init__(self, func):
        if not util.callable(func):
            raise exc.ArgumentError('func must be callable')

        self.type = type_api.NULLTYPE
        self.func = func

    def __getstate__(self):
        d = super().__getstate__()
        d['func'] = dill.dumps(self.func)
        return d

    def __setstate__(self, state):
        func = dill.loads(state.pop('func'))
        self.__dict__.update(state)
        self.func = func


class NestedColumn(ColumnClause):
    def __init__(self, parent, sub_column):
        self.parent = parent
        self.sub_column = sub_column
        if isinstance(self.parent, Label):
            table = self.parent.element.table
        else:
            table = self.parent.table
        super(NestedColumn, self).__init__(
            sub_column.name,
            sub_column.type,
            _selectable=table
        )


@compiles(NestedColumn)
def _comp(element, compiler, **kw):
    from_labeled_label = False
    if isinstance(element.parent, Label):
        from_labeled_label = True
    return "%s.%s" % (
        compiler.process(element.parent,
                         from_labeled_label=from_labeled_label,
                         within_label_clause=False,
                         within_columns_clause=True),
        compiler.visit_column(element, include_table=False),
    )


class ArrayJoin(ClauseList):
    __visit_name__ = 'array_join'


class ParametrizedFunction(functions.Function):
    """
    `funcname(funcparam)(funcarg)`,
    e.g. `quantile(0.5)(col1)`
    """

    __visit_name__ = 'parametrized_function'

    def __init__(self, name, params, *args, **kwargs):
        super(ParametrizedFunction, self).__init__(
            name, *args, **kwargs)
        self._func_name = name
        self._func_params = params
        self.params_expr = ClauseList(
            operator=functions.operators.comma_op,
            group_contents=True,
            *params
        ).self_group()

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/ext/clauses.py
datalens-tech/datalens-backend
from contextlib import contextmanager

from sqlalchemy import exc
import sqlalchemy.orm.query as query_module
import sqlalchemy.orm.context as context_module
from sqlalchemy.orm.query import Query as BaseQuery
from sqlalchemy.orm.util import _ORMJoin as _StandardORMJoin

from ..ext.clauses import (
    sample_clause,
    ArrayJoin,
)


class Query(BaseQuery):
    _with_totals = False
    _final = None
    _sample = None
    _array_join = None

    def _compile_context(self, *args, **kwargs):

        @contextmanager
        def replace_join():
            original = context_module._ORMJoin
            try:
                context_module._ORMJoin = _ORMJoin
                yield
            finally:
                context_module._ORMJoin = original

        with replace_join():
            context = super(Query, self)._compile_context(*args, **kwargs)
        statement = context.compile_state.statement

        statement._with_totals = self._with_totals
        statement._final_clause = self._final
        statement._sample_clause = sample_clause(self._sample)
        statement._array_join = self._array_join

        return context

    def with_totals(self):
        if not self._group_by_clauses:
            raise exc.InvalidRequestError(
                "Query.with_totals() can be used only with specified "
                "GROUP BY, call group_by()"
            )

        self._with_totals = True

        return self

    def array_join(self, *columns):
        self._array_join = ArrayJoin(*columns)
        return self

    def final(self):
        self._final = True

        return self

    def sample(self, sample):
        self._sample = sample

        return self

    def join(self, *props, **kwargs):
        type = kwargs.pop('type', None)
        strictness = kwargs.pop('strictness', None)
        distribution = kwargs.pop('distribution', None)
        rv = super(Query, self).join(*props, **kwargs)
        joineds = list(set(rv._from_obj) - set(self._from_obj))
        if not joineds:
            # Bad hack for sa1.4:
            rv._legacy_setup_joins[-1][-1].update(
                full=(rv._legacy_setup_joins[-1][-1]["full"], type, strictness, distribution),
            )
            return rv
        new = _ORMJoin._from_standard(joineds[0],
                                      type=type,
                                      strictness=strictness,
                                      distribution=distribution)

        @contextmanager
        def replace_join():
            original = query_module.orm_join
            try:
                query_module.orm_join = new
                yield
            finally:
                query_module.orm_join = original

        with replace_join():
            return super(Query, self).join(*props, **kwargs)

    def outerjoin(self, *props, **kwargs):
        kwargs['type'] = kwargs.get('type') or 'LEFT OUTER'
        return self.join(*props, **kwargs)


class _ORMJoin(_StandardORMJoin):
    @classmethod
    def _from_standard(cls, standard_join, type, strictness, distribution):
        return cls(
            standard_join.left,
            standard_join.right,
            standard_join.onclause,
            type=type,
            strictness=strictness,
            distribution=distribution
        )

    def __init__(
            self, left, right, onclause=None,
            isouter=False, full=False,
            type=None, strictness=None, distribution=None,
            **kwargs,
    ):
        super(_ORMJoin, self).__init__(left, right, onclause)

        # Bad hack for sa1.4:
        if isinstance(full, tuple):
            assert type is None
            assert strictness is None
            assert distribution is None
            full, type, strictness, distribution = full

        self.distribution = distribution
        self.strictness = str
        self.type = type
        self.strictness = None
        if strictness:
            self.strictness = strictness
        self.distribution = distribution
        self.type = type

    def __call__(self, *args, **kwargs):
        return self

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/orm/query.py
datalens-tech/datalens-backend

# from typing import Dict, Callable, Optional
from collections import namedtuple
from . import base
from .base import Parts, ParseError

__all__ = base.__all__ + (
    'FormatInfo', 'FORMATS',
    'Parts', 'ParseError',
)

# imported for namespace copy:
assert Parts
assert ParseError

_FormatInfo = namedtuple(
    '_FormatInfo',
    [
        'format_string', 'generator', 'generator_from_lines', 'sans_io',
        'chunk_size', 'parts',
    ])


class FormatInfo(_FormatInfo):

    def __new__(
            cls, format_string,
            generator=None, generator_from_lines=None, sans_io=None,
            chunk_size=64 * 1024, parts=base.Parts):
        """

        format_string: str: 'FORMAT {format_string}' part of the query text.

        generator:
        Callable[[Iterable[str]], Iterable[List[Any]]]
        (chunks to rows)

        generator_from_lines:
        Callable[[Iterable[str]], Iterable[List[Any]]]
        (lines to rows)

        sans_io: SansIOBase:
        feed-and-read class
        (`receive_data(data)`, `next_event()`)

        chunk_size: int: recommended chunk size for sans_io / generator.
        """
        return super(FormatInfo, cls).__new__(
            cls,
            format_string=format_string,
            generator=generator,
            generator_from_lines=generator_from_lines,
            sans_io=sans_io,
            chunk_size=chunk_size,
            parts=parts,
        )


# Wrapped into functions for lazy imports


def info_tsv():
    from .tsv import (
        parse_TabSeparatedWithNamesAndTypes,
        parse_TabSeparatedWithNamesAndTypes_lines,
    )
    return FormatInfo(
        format_string='TabSeparatedWithNamesAndTypes',
        generator_from_lines=parse_TabSeparatedWithNamesAndTypes_lines,
        generator=parse_TabSeparatedWithNamesAndTypes,
    )


def info_tsv_with_totals():
    from .tsv import parse_TabSeparatedWithNamesAndTypes_with_totals
    return FormatInfo(
        # implicitly expected: 'group by ... with totals' in the query text.
        format_string='TabSeparatedWithNamesAndTypes',
        generator=parse_TabSeparatedWithNamesAndTypes_with_totals,
    )


def info_jsoncompact():
    from .jsoncompact import JSONCompactChunksParser
    return FormatInfo(
        format_string='JSONCompact',
        generator=JSONCompactChunksParser.as_generator,
        sans_io=JSONCompactChunksParser,
    )


FORMATS = {  # Dict[Optional[str], Callable[[], FormatInfo]
    'default': info_tsv,
    'TabSeparatedWithNamesAndTypes': info_tsv,
    'TabSeparatedWithNamesAndTypes with totals': info_tsv_with_totals,
    'JSONCompact': info_jsoncompact,
}

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/parsers/__init__.py
datalens-tech/datalens-backend
# coding: utf8
"""
...
"""

from __future__ import unicode_literals


__all__ = (
    'Parts',
    'ParseError',
    'SansIOBase',
)

from ..util.compat import to_str


class Parts(object):

    # # Note: also testing the most basic structures
    # # to make it easier to cythonize if need be.
    # META, DATA, DATACHUNK, TOTALS, STATS, NEED_DATA, FINISHED = range(7)
    # # Debug-oriented:
    META = 'PART_META'
    DATA = 'PART_DATA'
    DATACHUNK = 'PART_DATACHUNK'
    TOTALS = 'PART_TOTALS'
    STATS = 'PART_STATS'
    NEED_DATA = 'NEED_DATA'
    FINISHED = 'FINISHED'


class ParseError(ValueError):

    def __init__(self, error, state='', line='', etcetera=None):
        self.error = to_str(error)
        self.state = to_str(state)
        self.line = line
        self.etcetera = etcetera
        message = 'state: %s; error: %s; line: %r' % (
            self.state, self.error, line)
        if etcetera:
            message += ' (%r)' % (etcetera,)
        super().__init__(message)


class SansIOBase(object):
    """
    Parser interface for async integration.
    See also: `h11.Connection`.
    """

    def receive_data(self, data):
        raise NotImplementedError

    def next_event(self):
        raise NotImplementedError

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/parsers/base.py
datalens-tech/datalens-backend
# coding: utf8
"""
...
"""

from __future__ import unicode_literals

import enum
import json
import logging
from typing import Optional

from .base import (
    Parts,
    ParseError,
    SansIOBase,
)
from ..exceptions import DatabaseException

logger = logging.getLogger(__name__)


@enum.unique
class ParserState(enum.Enum):
    START = "STATE_START"
    NEXT_BLOCK = "STATE_NEXT_BLOCK"
    META = "STATE_META"
    DATA = "STATE_DATA"
    FOOTER = "STATE_FOOTER"
    TOTALS = "STATE_TOTALS"  # deprecated
    STATS = "STATE_STATS"  # deprecated
    FINISHED = "STATE_FINISHED"
    ERROR = "STATE_ERROR"


class JSONCompactChunksParser(SansIOBase):
    """
    A parser class for JSONCompact chunks.

    Assumes a very particular positioning of newlines and whitespaces (and the
    data structure). Use a clickhouse version whitelist for reliability.

    Structured similarly to `h11.Connection`.
    """

    # minimal chunk size to ensure all the whitespaces for parsing actually fit
    # in and there's at least something in it.
    min_chunk_size = 1024
    # A bunch of things are expected to fit in this size:
    # metadata, each data line, footer.
    # On the other hand, garbage data will be read up to this size (+ chunk
    # size) into memory until raising an error.
    max_chunk_size = 8 * 1024 * 1024

    parts = Parts

    STATE_START = ParserState.START.value
    STATE_NEXT_BLOCK = ParserState.NEXT_BLOCK.value
    STATE_META = ParserState.META.value
    STATE_DATA = ParserState.DATA.value
    STATE_FOOTER = ParserState.FOOTER.value
    STATE_TOTALS = ParserState.TOTALS.value  # deprecated
    STATE_STATS = ParserState.STATS.value  # deprecated
    STATE_FINISHED = ParserState.FINISHED.value
    STATE_ERROR = ParserState.ERROR.value

    class RowTooLarge(ParseError):
        """ ... """

    def __init__(self, autoparse=True):
        self._unread = []
        self._results = []
        self.state = self.STATE_START
        self._autoparse = autoparse

    def receive_data(self, data: Optional[bytes]):
        """ Add data to the internal receive buffer """
        if not data:
            return  # allow and ignore empty chunks, just in case.
        self._unread.append(data)

    @staticmethod
    def _splitpair(chunk, splitters, start=0):
        assert not isinstance(splitters, (bytes, str))
        for idx, (splitter, rightsplit) in enumerate(splitters):
            if rightsplit:
                func = chunk.rindex
            else:
                func = chunk.index
            try:
                pos = func(splitter, start)
            except ValueError:  # `ValueError: subsection not found`
                continue

            return chunk[:pos], chunk[pos + len(splitter):], idx

        return chunk, None, None

    def _read_more(self, current_chunk):
        if not self._unread:
            return False, current_chunk
        extra_data = self._unread.pop(0)
        if current_chunk is not None:
            # max_chunk_size check
            l1 = len(current_chunk)
            l2 = len(extra_data)
            mcs = self.max_chunk_size
            if l1 <= mcs and l2 <= mcs and l1 + l2 > mcs:
                # Sample case: looking for metadata end, reading until it is
                # found, but it is not found and there's a lot of data.
                raise self.RowTooLarge(
                    "Tried to read more data than `max_chunk_size`",
                    dict(l1=l1, l2=l2, mcs=mcs))

            # Same as below, but only done when it is likely to help
            # (i.e. on large rows).
            if (l1 > self.min_chunk_size * 5 and not isinstance(current_chunk, bytearray)):
                current_chunk = bytearray(current_chunk)

            current_chunk += extra_data

            return True, current_chunk

        # # This should allow append (above) to happen inplace,
        # # but adds an extra copy.
        # extra_data = bytearray(extra_data)
        return True, extra_data

    def _read_until(self, current_chunk, end_markers):
        assert not isinstance(end_markers, (bytes, str))
        # Semantics note:
        # `right is None` means `not found`,
        # `right == b""` means `endswith`
        left, right, idx = self._splitpair(current_chunk, end_markers)
        max_marker_len = max(len(marker[0]) for marker in end_markers)
        while right is None:
            chunk_prev_len = len(current_chunk)
            # raises ValueError:
            more, current_chunk = self._read_more(current_chunk)
            if not more:
                break
            # Try again.
            start_pos = chunk_prev_len - max_marker_len
            if start_pos < 0:
                start_pos = 0
            left, right, idx = self._splitpair(
                current_chunk,
                end_markers,
                start=start_pos,
            )
        if right is not None:
            return left, right, idx
        return None, current_chunk, None

    def _put_back(self, chunk):
        if not chunk:
            return
        self._unread.insert(0, chunk)

    def _add_result(self, part, item):
        if self._autoparse:
            item = json.loads(item)
        self._results.append((part, item))

    def next_event(self):
        """ ... """
        if self.state == self.STATE_FINISHED:  # terminating state
            return self.parts.FINISHED, None
        if self.state == self.STATE_ERROR:  # terminating state
            raise ParseError(
                "next event requested in error-state",
                "STATE_ERROR", b"")

        if self._results:
            return self._results.pop(0)

        more, chunk = self._read_more(None)
        if not chunk:
            return self.parts.NEED_DATA, None

        while len(chunk) < self.min_chunk_size:
            more, chunk = self._read_more(chunk)
            if not more:
                break

        if self.state == self.STATE_START:
            prefix = b'{\n\t"meta":\n\t[\n'
            chunk = chunk.lstrip()  # allow extra whitespaces at start.

            if len(chunk) < len(prefix):
                self._put_back(chunk)
                return self.parts.NEED_DATA, None

            if not chunk.startswith(prefix):
                self.state = self.STATE_ERROR
                raise ParseError("unexpected data start", "STATE_START", chunk)

            chunk = chunk[len(prefix):]
            self.state = self.STATE_META

        if self.state == self.STATE_META:
            # (implicit "[") + data + (implicit "]")
            # maybe: `meta_end = b"\n\t\t}\n\t],\n\n"`
            meta_end = b"\n\t],\n"
            meta_chunk, chunk, _ = self._read_until(
                chunk,
                (
                    (meta_end, False),
                ))
            if meta_chunk is not None:
                self.state = self.STATE_NEXT_BLOCK
                self._add_result(self.parts.META, b"[" + meta_chunk + b"]")

        if self.state == self.STATE_NEXT_BLOCK:
            next_blocks = (
                (b'\n\t"data":\n\t[\n', self.STATE_DATA),  # data_start
                # Footer examples:
                # '\n\n\t"totals": ['
                # '\n\n\t"extremes":\n\t{'
                # '\n\n\t"rows": '
                (b'\n\t"', self.STATE_FOOTER),
            )
            for prefix, next_state in next_blocks:
                left, right, _ = self._splitpair(chunk, ((prefix, False),))
                if right is not None:
                    if left.strip():
                        self.state = self.STATE_ERROR
                        raise ParseError(
                            ("unexpected data between blocks "
                             "(to state {})").format(next_state),
                            "STATE_NEXT_BLOCK",
                            chunk)
                    chunk = right
                    self.state = next_state
                    break
            # `else:`
            if self.state == self.STATE_NEXT_BLOCK:
                if len(chunk) > self.min_chunk_size:
                    self.state = self.STATE_ERROR
                    raise ParseError(
                        "could not find the next block",
                        "STATE_NEXT_BLOCK", chunk)
                # Otherwise perhaps the chunk was too small.

        if self.state == self.STATE_DATA:
            # Up to multiple b"...\n\t\t[...],\n..." chunks
            data_end = b"\n\t],\n"
            data_line_end = b"],\n"
            data_chunk, chunk, match_idx = self._read_until(
                chunk,
                # Tricky ordering: most of the time this will result in going
                # over the data twice,
                # but it has to be done because `data_end` can be found after
                # the data block.
                (
                    (data_end, False),
                    (data_line_end, True),
                ))
            if data_chunk is not None:
                # maybe: `and b"[" in data_chunk`
                # maybe: `and data_chunk != b"\n\t"`
                self._add_result(
                    self.parts.DATACHUNK,
                    (b"[" +
                     data_chunk +
                     (b"]]" if match_idx == 1 else b"]")),
                )
            if match_idx == 0:  # `data_end`
                self.state = self.STATE_NEXT_BLOCK

        if self.state == self.STATE_FOOTER:
            suffix = b"\t}\n}"
            stats_chunk, chunk, _ = self._read_until(chunk, ((suffix, False),))
            if stats_chunk is not None:
                self.state = self.STATE_FINISHED
                self._add_result(
                    self.parts.STATS,
                    b'{"' + stats_chunk + suffix)
                chunk = chunk.strip()
                if chunk:
                    # Data in this state is not expected. Check if it's a timeout error and log the error.
                    # Raise ParseError otherwise.
                    db_exc = DatabaseException(chunk.decode())
                    if db_exc.code == '159':  # Timeout exceeded
                        logger.error("Got timeout error message after finished response processing.")
                    else:
                        self.state = self.STATE_ERROR
                        raise ParseError(
                            "Extra data after the end",
                            "STATE_STATS", chunk)

        if chunk:
            self._put_back(chunk)

        if self._results:
            return self._results.pop(0)
        return self.parts.NEED_DATA, None

    @classmethod
    def as_generator(cls, data_chunks, **kwargs):
        parser = cls(**kwargs)
        event = None
        chunk = None
        for chunk in data_chunks:
            parser.receive_data(chunk)
            event, data = parser.next_event()
            while (event != parser.parts.NEED_DATA and
                   event != parser.parts.FINISHED):
                yield event, data
                event, data = parser.next_event()
        if event is not None and event != parser.parts.FINISHED:
            raise ValueError(
                "Unexpected last event",
                dict(event=event, data=data, chunk=chunk))

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/parsers/jsoncompact.py
datalens-tech/datalens-backend

import codecs

from .base import Parts, ParseError, to_str


__all__ = [
    'parse_tsv_line',
    'parse_TabSeparatedWithNamesAndTypes_general',
    'parse_TabSeparatedWithNamesAndTypes',
    'parse_TabSeparatedWithNamesAndTypes_with_totals',
]


def unescape(value):
    return codecs.escape_decode(value)[0].decode('utf-8', errors='replace')


def parse_tsv_line(line):
    if line and line[-1] == b'\n':
        line = line[:-1]
    return [
        unescape(value) if value != b'\\N' else None
        for value in line.split(b'\t')]


def chunks_to_lines(chunks, sep=b'\n'):
    r"""

    >>> data = b'value\nString\na\nb\nc\nd\n'
    >>> data.split(b'\n')
    [b'value', b'String', b'a', b'b', b'c', b'd', b'']
    >>> list(chunks_to_lines([data])) == data.split(b'\n')
    True
    """
    unflushed = None
    for chunk in chunks:
        pieces = chunk.split(sep)
        for piece in pieces[:-1]:  # all-but-last
            if unflushed is not None:
                piece = unflushed + piece
                unflushed = None
            yield piece
        if unflushed is None:
            unflushed = pieces[-1]
        else:
            unflushed += pieces[-1]
    if unflushed is not None:
        yield unflushed


def chunks_with_terminating_newline_to_lines(chunks, sep=b'\n'):
    r"""
    Wrapper over `chunks_to_lines` that requires and strips the last newline.

    >>> data = b'value\nString\na\nb\nc\nd\n'
    >>> list(chunks_with_terminating_newline_to_lines([data]))
    [b'value', b'String', b'a', b'b', b'c', b'd', b'']
    >>> list(chunks_with_terminating_newline_to_lines([data.rstrip()]))
    Traceback (most recent call last):
      ...
    ValueError: ('Missing the terminating newline', b'd')
    >>> list(chunks_with_terminating_newline_to_lines([b'xx']))
    Traceback (most recent call last):
      ...
    ValueError: ('Missing the terminating newline', b'xx')
    """
    items = iter(chunks_to_lines(chunks, sep=sep))
    try:
        prev_item = next(items)
    except StopIteration:
        return
    item = prev_item
    for item in items:
        yield prev_item
        prev_item = item
    if item != b'':
        raise ValueError('Missing the terminating newline', item)


def parse_tsv_chunks(chunks):
    for line in chunks_with_terminating_newline_to_lines(chunks):
        yield parse_tsv_line(line)


UNSPECIFIED = object()


def parse_TabSeparatedWithNamesAndTypes_general(
        chunks, with_totals=UNSPECIFIED, empty_string=''):
    """
    General / parametric case:
    required totals / no totals / ambiguous maybe-there-will-be-totals.
    """
    items = iter(parse_tsv_chunks(chunks))
    with_totals_ambiguous = with_totals is UNSPECIFIED
    with_totals = bool(with_totals)

    try:
        names = next(items)
        types = next(items)
    except StopIteration:
        raise ParseError('Unexpected stream end', 'STATE_START')

    if len(names) != len(types):
        raise ParseError(
            'Mismatch in amount of column names and column types',
            'STATE_META',
            etcetera=dict(names=names, types=types))

    metadata = [
        {'name': to_str(name), 'type': to_str(type_name)}
        for name, type_name in zip(names, types)]
    yield Parts.META, metadata

    def get_next_item():
        if lookahead:
            return lookahead.pop(0)
        return next(items)  # raises StopIteration

    def data_result(some_item):
        if len(some_item) != len(metadata):
            raise ParseError('Unexpected row width', 'STATE_DATA',
                             etcetera=dict(item=some_item, metadata=metadata))
        return Parts.DATA, some_item

    def totals_result(some_item):
        if len(some_item) != len(metadata):
            raise ParseError('Unexpected row width', 'STATE_TOTALS',
                             etcetera=dict(item=some_item, metadata=metadata))
        return Parts.TOTALS, some_item

    lookahead = []
    state = 'STATE_DATA'
    while True:
        try:
            item = get_next_item()
        except StopIteration:
            break

        if item == [empty_string]:
            # Cases:
            #
            #  * A single string column which happens to have an empty value,
            #  * An empty line before totals.

            while len(lookahead) < 2:
                try:
                    lookahead.append(next(items))
                except StopIteration:
                    break

            if lookahead == []:  # '…\n' – end of the stream.
                break
            assert lookahead
            if len(lookahead) == 1:  # end of the stream
                if with_totals:
                    totals_item = lookahead[0]
                    state = 'FINISHED'
                    yield totals_result(totals_item)
                    break
            yield data_result(item)
            continue
        else:
            yield data_result(item)

    if with_totals and not with_totals_ambiguous and state == 'STATE_DATA':
        raise ParseError(
            'Stream ended before totals', state,
            etcetera=dict(item=item))


def parse_TabSeparatedWithNamesAndTypes_lines(lines, **kwargs):
    items = (parse_tsv_line(line) for line in lines)

    try:
        names = next(items)
        types = next(items)
    except StopIteration:
        # Have to support parsing a completely empty stream
        # (for e.g. DDL requests).
        return

    if len(names) != len(types):
        raise ParseError(
            'Mismatch in amount of column names and column types',
            'STATE_META',
            etcetera=dict(names=names, types=types))

    metadata = [
        {'name': to_str(name), 'type': to_str(type_name)}
        for name, type_name in zip(names, types)]
    yield Parts.META, metadata

    for item in items:
        if len(item) != len(metadata):
            raise ParseError('Unexpected row width', 'STATE_DATA',
                             etcetera=dict(item=item, metadata=metadata))
        yield Parts.DATA, item


def parse_TabSeparatedWithNamesAndTypes(chunks, **kwargs):
    lines = chunks_with_terminating_newline_to_lines(chunks)
    return parse_TabSeparatedWithNamesAndTypes_lines(lines, **kwargs)


def parse_TabSeparatedWithNamesAndTypes_with_totals(chunks, **kwargs):
    return parse_TabSeparatedWithNamesAndTypes_general(
        chunks, with_totals=True, **kwargs)

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/parsers/tsv.py
datalens-tech/datalens-backend
"""
...

Note: this is a reworked version of `.drivers.http.escaper`, currently used for
literal_binds, might be used for everything.

Unlike escaper, the quoting output is always a string.

Moreover, it should technically be a bytestring,
but currently that does not seem feasible.
"""

from datetime import date, datetime
from decimal import Decimal
import enum

from sqlalchemy.sql.elements import Null

from .util import compat


class Quoter(object):
    """
    Value quoting and serialization for queries.

    Note: the return value is always a string.

    See also: `.drivers.http.escaper`.

    See also: python-clickhouse-client's  `clickhouse.tools.quote_value`.
    """

    number_types = compat.integer_types
    string_types = compat.string_types

    escape_chars = {
        "\b": r"\b",
        "\f": r"\f",
        "\r": r"\r",
        "\n": r"\n",
        "\t": r"\t",
        "\0": r"\0",
        "\\": "\\\\",
        "'": r"\'",
    }

    def quote_str_inner(self, value):
        # TODO: compare the performance with:
        #     import re
        #     escape_chars_re = re.compile(
        #         '[{}]'.format(
        #             ''.join(
        #                 re.escape(char)
        #                 for char in escape_chars)))
        #     value = self.escape_chars_re.sub(
        #         lambda match, escape_chars=escape_chars: (
        #             escape_chars[match.group(0)]),
        #         value)

        value = ''.join(self.escape_chars.get(char, char) for char in value)
        return value

    def quote_str(self, value):
        value = self.quote_str_inner(value)
        return "'" + value + "'"

    def quote_number(self, item):
        return '%d' % (int(item),)

    def quote_date(self, item):
        return 'toDate(%s)' % (
            self.quote_str(item.strftime('%Y-%m-%d')),)

    def quote_datetime(self, item):
        return 'toDateTime(%s)' % (
            self.quote_str(item.strftime('%Y-%m-%d %H:%M:%S')),)

    def quote_decimal(self, item):
        return str(float(item))

    def quote(self, value):
        if value is None or isinstance(value, Null):
            return 'NULL'

        if isinstance(value, self.number_types):
            return self.quote_number(value)

        if isinstance(value, datetime):
            return self.quote_datetime(value)
        if isinstance(value, date):
            return self.quote_date(value)

        if isinstance(value, (Decimal, float)):
            return self.quote_decimal(value)

        if isinstance(value, self.string_types):
            return self.quote_str(value)

        # XXX: shouldn't `()` be used for tuples?
        if isinstance(value, (list, tuple)):
            return "[" + ", ".join(
                self.quote(subitem)
                for subitem in value
            ) + "]"

        if isinstance(value, enum.Enum):
            return self.quote_str(value.name)

        raise Exception("Unsupported object type: {}".format(type(value)))


QUOTER = Quoter()

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/quoting.py
datalens-tech/datalens-backend
from sqlalchemy.sql.ddl import (
    SchemaDropper as SchemaDropperBase, DropTable as DropTableBase
)


class DropTable(DropTableBase):
    def __init__(self, element, **kw):
        self.on_cluster = element.dialect_options['clickhouse']['cluster']
        super(DropTable, self).__init__(element, **kw)


class SchemaDropper(SchemaDropperBase):
    def __init__(self, dialect, connection, if_exists=False, **kwargs):
        self.if_exists = if_exists
        super(SchemaDropper, self).__init__(dialect, connection, **kwargs)

    def visit_table(self, table, drop_ok=False, _is_metadata_operation=False, _ignore_sequences=None):
        self.connection.execute(DropTable(table, if_exists=self.if_exists))

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/sql/ddl.py
datalens-tech/datalens-backend
from sqlalchemy import Table as TableBase
from sqlalchemy.sql.base import (
    _bind_or_error,
)

from clickhouse_sqlalchemy.sql.selectable import (
    Join,
    Select,
)
from . import ddl


class Table(TableBase):

    def drop(self, bind=None, checkfirst=False, if_exists=False):
        if bind is None:
            bind = _bind_or_error(self)
        bind._run_ddl_visitor(ddl.SchemaDropper, self, checkfirst=checkfirst, if_exists=if_exists)

    def join(self, right, onclause=None, isouter=False, full=False,
             type=None, strictness=None, distribution=None):
        return Join(self, right,
                    onclause=onclause, type=type,
                    isouter=isouter, full=full,
                    strictness=strictness, distribution=distribution)

    def select(self, whereclause=None, **kwargs):
        if whereclause is not None:
            kwargs["whereclause"] = whereclause
        return Select._create_select_from_fromclause(self, [self], **kwargs)

    @classmethod
    def _make_from_standard(cls, std_table, _extend_on=None):
        ch_table = cls(std_table.name, std_table.metadata)
        ch_table.schema = std_table.schema
        ch_table.fullname = std_table.fullname
        ch_table.implicit_returning = std_table.implicit_returning
        ch_table.comment = std_table.comment
        ch_table.info = std_table.info
        ch_table._prefixes = std_table._prefixes
        ch_table.dialect_options = std_table.dialect_options

        if _extend_on is None:
            ch_table._columns = std_table._columns

        return ch_table

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/sql/schema.py
datalens-tech/datalens-backend
from sqlalchemy.sql.selectable import (
    Select as StandardSelect,
    Join as StandardJoin,
)

from clickhouse_sqlalchemy.ext.clauses import ArrayJoin
from ..ext.clauses import sample_clause


__all__ = ('Select', 'select')


class Join(StandardJoin):

    def __init__(self, left, right,
                 onclause=None, isouter=False, full=False,
                 type=None, strictness=None, distribution=None):
        if type is not None:
            type = type.upper()
        super(Join, self).__init__(left, right, onclause,
                                   isouter=isouter, full=full)
        self.strictness = None
        if strictness:
            self.strictness = strictness
        self.distribution = distribution
        self.type = type


class Select(StandardSelect):
    _with_totals = False
    _final_clause = None
    _sample_clause = None
    _array_join = None

    @classmethod
    def _create_select_from_fromclause(cls, target, entities, *arg, **kw):
        if arg or kw:
            return cls.create_legacy_select(entities, *arg, **kw)
        else:
            return cls._create_select(*entities)

    def with_totals(self):
        self._with_totals = True
        return self

    def final(self):
        self._final_clause = True
        return self

    def sample(self, sample):
        self._sample_clause = sample_clause(sample)
        return self

    def array_join(self, *columns):
        self._array_join = ArrayJoin(*columns)
        return self

    def join(self, right, onclause=None, isouter=False, full=False, type=None,
             strictness=None, distribution=None):
        return Join(self, right,
                    onclause=onclause, type=type,
                    isouter=isouter, full=full,
                    strictness=strictness, distribution=distribution)


select = Select
join = Join

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/sql/selectable.py
datalens-tech/datalens-backend

import decimal
from ipaddress import (
    IPv4Address, IPv6Address,
    IPv4Network, IPv6Network,
)

import sqlalchemy as sa
from sqlalchemy.sql.type_api import to_instance, UserDefinedType
from sqlalchemy import types, func, and_, or_

from .dt_utils import (
    parse_date, parse_datetime, parse_datetime64, parse_timezone,
)
from .ext.clauses import NestedColumn
from .util.compat import text_type


def _noop(value):
    return value


def get_http_postprocess(coltype):
    """
    A hack to account for the fact that non-parametrized types can often happen
    to be non-instantiated.
    """
    if isinstance(coltype, type):
        coltype = coltype()  # should be parameterless in this case
    return getattr(coltype, '_chsa_http_postprocess', _noop)


class CHTypeMixin(object):
    """ Common base-class marker for the CHSA types """

    def _chsa_http_postprocess(self, value):
        return value


class String(CHTypeMixin, types.String):
    pass


class Int(CHTypeMixin, types.Integer):

    def _chsa_http_postprocess(self, value):
        return int(value)


class Float(CHTypeMixin, types.Float):

    def _chsa_http_postprocess(self, value):
        if value is None:
            # sample non-nullable case: `select 1 / 0 format JSON`
            return None
        return float(value)


class Array(CHTypeMixin, types.TypeEngine):
    __visit_name__ = 'array'

    def __init__(self, item_type):
        self.item_type = item_type
        self.item_type_impl = to_instance(item_type)
        super(Array, self).__init__()

    def literal_processor(self, dialect):
        item_processor = self.item_type_impl.literal_processor(dialect)

        def process(value):
            processed_value = []
            for x in value:
                if item_processor:
                    x = item_processor(x)
                processed_value.append(x)
            return '[' + ', '.join(processed_value) + ']'
        return process


class Nullable(CHTypeMixin, types.TypeEngine):
    __visit_name__ = 'nullable'

    def __init__(self, nested_type):
        self.nested_type = nested_type
        self._nested_http_postprocess = get_http_postprocess(self.nested_type)
        super(Nullable, self).__init__()

    def _chsa_http_postprocess(self, value):
        if value is None:
            return None
        return self._nested_http_postprocess(value)


class UUID(String):
    __visit_name__ = 'uuid'


class LowCardinality(CHTypeMixin, types.TypeEngine):
    __visit_name__ = 'lowcardinality'

    def __init__(self, nested_type):
        self.nested_type = nested_type
        super(LowCardinality, self).__init__()


class Int8(Int):
    __visit_name__ = 'int8'


class UInt8(Int):
    __visit_name__ = 'uint8'


class Int16(Int):
    __visit_name__ = 'int16'


class UInt16(Int):
    __visit_name__ = 'uint16'


class Int32(Int):
    __visit_name__ = 'int32'


class UInt32(Int):
    __visit_name__ = 'uint32'


class Int64(Int):
    __visit_name__ = 'int64'


class UInt64(Int):
    __visit_name__ = 'uint64'


class Float32(Float):
    __visit_name__ = 'float32'


class Float64(Float):
    __visit_name__ = 'float64'


class Bool(CHTypeMixin, types.Boolean):
    __visit_name__ = 'bool'

    def _chsa_http_postprocess(self, value):
        return bool(value)


class Date(CHTypeMixin, types.Date):
    __visit_name__ = 'date'

    def _chsa_http_postprocess(self, value):
        return parse_date(value)


class Date32(CHTypeMixin, types.Date):
    __visit_name__ = 'date32'

    def _chsa_http_postprocess(self, value):
        return parse_date(value)


class DateTime(CHTypeMixin, types.DateTime):
    __visit_name__ = 'datetime'

    def _chsa_http_postprocess(self, value):
        return parse_datetime(value)


class DateTime64(CHTypeMixin, types.DateTime):
    __visit_name__ = 'datetime64'

    def __init__(self, precision):
        super(DateTime64, self).__init__()
        self.precision = precision

    def _chsa_http_postprocess(self, value):
        return parse_datetime64(value)


class DateTimeWithTZ(DateTime):
    """
    'Aware' datetime, a datetime with timezone.
    """

    __visit_name__ = 'datetimewithtz'

    def __init__(self, tz):
        super(DateTimeWithTZ, self).__init__()
        self.tz = tz
        self.tzinfo = parse_timezone(tz)

    def _chsa_http_postprocess(self, value):
        return parse_datetime(value, tzinfo=self.tzinfo)


class DateTime64WithTZ(DateTime64):
    """
    'Aware' 64-bit datetime, a datetime with timezone.
    """

    __visit_name__ = 'datetime64withtz'

    def __init__(self, precision, tz):
        super(DateTime64WithTZ, self).__init__(precision)
        self.tz = tz
        self.tzinfo = parse_timezone(tz)

    def _chsa_http_postprocess(self, value):
        return parse_datetime64(value, tzinfo=self.tzinfo)


class Enum(CHTypeMixin, types.Enum):
    __visit_name__ = 'enum'

    def __init__(self, *enums, **kw):
        if not enums:
            enums = kw.get('_enums', ())  # passed as keyword

        super(Enum, self).__init__(*enums, **kw)


class Enum8(Enum):
    __visit_name__ = 'enum8'


class Enum16(Enum):
    __visit_name__ = 'enum16'


class Decimal(CHTypeMixin, types.Numeric):
    __visit_name__ = 'numeric'

    def _chsa_http_postprocess(self, value):
        if value is None:
            return None
        return decimal.Decimal(value)


class Nested(CHTypeMixin, types.TypeEngine):
    __visit_name__ = 'nested'

    def __init__(self, *columns):
        if not columns:
            raise ValueError('columns must be specified for nested type')
        self.columns = columns
        self._columns_dict = {col.name: col for col in columns}
        super(Nested, self).__init__()

    class Comparator(UserDefinedType.Comparator):
        def __getattr__(self, key):
            str_key = key.rstrip("_")
            try:
                sub = self.type._columns_dict[str_key]
            except KeyError:
                raise AttributeError(key)
            else:
                original_type = sub.type
                try:
                    sub.type = Array(sub.type)
                    expr = NestedColumn(self.expr, sub)
                    return expr
                finally:
                    sub.type = original_type

    comparator_factory = Comparator


class IPv4(CHTypeMixin, types.UserDefinedType):
    __visit_name__ = "ipv4"

    def bind_processor(self, dialect):
        def process(value):
            return text_type(value)

        return process

    def bind_expression(self, bindvalue):
        bindvalue = bindvalue._clone(maintain_key=True)
        bindvalue.expanding = False
        return func.toIPv4(bindvalue)

    class comparator_factory(types.UserDefinedType.Comparator):
        def in_(self, other):
            if isinstance(other, sa.sql.elements.BindParameter):
                other = other.value
            if not isinstance(other, IPv4Network):
                other = IPv4Network(other)

            return and_(
                self >= other[0],
                self <= other[-1]
            )

        def not_in(self, other):
            return self.notin_(other)

        def notin_(self, other):
            if isinstance(other, sa.sql.elements.BindParameter):
                other = other.value
            if not isinstance(other, IPv4Network):
                other = IPv4Network(other)

            return or_(
                self < other[0],
                self > other[-1]
            )

    def _chsa_http_postprocess(self, value):
        return IPv4Address(value)


class IPv6(CHTypeMixin, types.UserDefinedType):
    __visit_name__ = "ipv6"

    def bind_processor(self, dialect):
        def process(value):
            return text_type(value)

        return process

    def bind_expression(self, bindvalue):
        return func.toIPv6(bindvalue)

    class comparator_factory(types.UserDefinedType.Comparator):
        def in_(self, other):
            if isinstance(other, sa.sql.elements.BindParameter):
                other = other.value
            if not isinstance(other, IPv6Network):
                other = IPv6Network(other)

            return and_(
                self >= other[0],
                self <= other[-1]
            )

        def not_in(self, other):
            return self.notin_(other)

        def notin_(self, other):
            if isinstance(other, sa.sql.elements.BindParameter):
                other = other.value
            if not isinstance(other, IPv6Network):
                other = IPv6Network(other)

            return or_(
                self < other[0],
                self > other[-1]
            )

    def _chsa_http_postprocess(self, value):
        return IPv6Address(value)

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/types.py
datalens-tech/datalens-backend
import types
import sys


PY2 = sys.version_info[0] == 2
PY3 = sys.version_info[0] == 3


if PY3:
    string_types = str,
    integer_types = int,
    class_types = type,
    text_type = str
    binary_type = bytes

else:
    string_types = basestring,    # noqa
    integer_types = (int, long)  # noqa
    class_types = (type, types.ClassType)
    text_type = unicode  # noqa
    binary_type = str


def to_bytes(value):
    if isinstance(value, text_type):
        return value.encode('utf-8', errors='replace')
    return value


def to_str(value):
    if isinstance(value, binary_type):
        return value.decode('utf-8', errors='replace')
    return value

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy/util/compat.py
datalens-tech/datalens-backend
import pkgutil
from sqlalchemy.dialects import registry

from clickhouse_sqlalchemy.util import compat

from . import log

if compat.PY3:
    import configparser
else:
    import ConfigParser as configparser


registry.register(
    "clickhouse", "clickhouse_sqlalchemy.drivers.http.base", "dialect"
)
registry.register(
    "clickhouse.http", "clickhouse_sqlalchemy.drivers.http.base", "dialect"
)
registry.register(
    "clickhouse.native", "clickhouse_sqlalchemy.drivers.native.base", "dialect"
)

file_config = configparser.ConfigParser()
# file_config.read(['setup.cfg'])
setup_cfg_data = pkgutil.get_data(
    __package__ or __name__.rsplit('.', 1)[0],
    'setup.cfg')
file_config.read_string(setup_cfg_data.decode('utf-8'))

log.configure(file_config.get('log', 'level'))

host = file_config.get('db', 'host')
port = file_config.getint('db', 'port')
http_port = file_config.getint('db', 'http_port')
database = file_config.get('db', 'database')
user = file_config.get('db', 'user')
password = file_config.get('db', 'password')

uri_template = '{schema}://{user}:{password}@{host}:{port}/{database}'

http_uri = uri_template.format(
    schema='clickhouse+http', user=user, password=password, host=host,
    port=http_port, database=database,
)
http_tsv_uri = http_uri
http_jc_uri = http_uri + '?format=JSONCompact'

uri = http_uri

native_uri = uri_template.format(
    schema='clickhouse+native', user=user, password=password, host=host,
    port=port, database=database)

system_http_uri = uri_template.format(
    schema='clickhouse+http', user=user, password=password, host=host,
    port=http_port, database='system')
system_native_uri = uri_template.format(
    schema='clickhouse+native', user=user, password=password, host=host,
    port=port, database='system')

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/config.py
datalens-tech/datalens-backend
from clickhouse_sqlalchemy_tests.testcase import HttpSessionTestCase


class CursorTestCase(HttpSessionTestCase):
    def test_check_iter_cursor(self):
        rv = self.session.execute('SELECT number FROM system.numbers LIMIT 5')
        self.assertListEqual(list(rv), [(x,) for x in range(5)])

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/http/test_cursor.py
datalens-tech/datalens-backend
from decimal import Decimal
from datetime import datetime, date

from sqlalchemy import Column, literal

from clickhouse_sqlalchemy import types, engines, Table
from clickhouse_sqlalchemy.drivers.http.escaper import Escaper
from clickhouse_sqlalchemy_tests.testcase import HttpSessionTestCase


class EscapingTestCase(HttpSessionTestCase):
    def escaped_compile(self, clause, **kwargs):
        return str(self._compile(clause, **kwargs))

    def test_select_escaping(self):
        query = self.session.query(literal('\t'))
        self.assertEqual(
            self.escaped_compile(query, literal_binds=True),
            r"SELECT '\t' AS anon_1"
        )

    def test_escaper(self):
        e = Escaper()
        self.assertEqual(e.escape([None]), '[NULL]')
        self.assertEqual(e.escape([[None]]), '[[NULL]]')
        self.assertEqual(e.escape([[123]]), '[[123]]')
        self.assertEqual(e.escape({'x': 'str'}), {'x': "'str'"})
        self.assertEqual(e.escape([Decimal('10')]), '[10.0]')
        self.assertEqual(e.escape([10.0]), '[10.0]')
        self.assertEqual(e.escape([date(2017, 1, 2)]), "['2017-01-02']")

        with self.assertRaises(Exception) as ex:
            e.escape([object()])

        self.assertIn('Unsupported object', str(ex.exception))

        with self.assertRaises(Exception) as ex:
            e.escape('str')

        self.assertIn('Unsupported param format', str(ex.exception))

    def test_escape_binary_mod(self):
        query = self.session.query(literal(1) % literal(2))
        self.assertEqual(
            self.compile(query, literal_binds=True),
            'SELECT 1 %% 2 AS anon_1'
        )

        table = Table(
            't', self.metadata(),
            Column('x', types.Int32, primary_key=True),
            engines.Memory()
        )

        query = self.session.query(table.c.x % table.c.x)
        self.assertEqual(
            self.compile(query, literal_binds=True),
            'SELECT x %% x AS anon_1 FROM t'
        )

    def test_literal_binds(self):
        null = None
        query = self.session.query(
            literal(null).label('null'),
            literal([null]).label('nulla'),
            literal([[null]]).label('nullaa'),
            literal([(123,)]).label('intaa'),
            literal('zxc').label('str'),
            literal('фыва⋯').label('stru'),
            literal(123).label('int'),
            literal(
                340282366920938463463374607431768211456
            ).label('longint'),
            literal(date(2018, 2, 28)).label('date'),
            literal(
                datetime(2018, 2, 28, 23, 59, 59, 999999)
            ).label('datetime'),
            literal(Decimal(10.0000000010)).label('decimal'),
        )

        query_text = self.compile(query, literal_binds=True)
        expected_query_text = 'SELECT ' + ', '.join((
            'NULL AS "null"',
            '[NULL] AS nulla',
            '[[NULL]] AS nullaa',
            '[[123]] AS intaa',
            '\'zxc\' AS str',
            '\'фыва⋯\' AS stru',
            '123 AS int',
            '340282366920938463463374607431768211456 AS longint',
            'toDate(\'2018-02-28\') AS date',
            'toDateTime(\'2018-02-28 23:59:59\') AS datetime',
            '10.000000001 AS decimal'
        ))
        self.assertEqual(query_text, expected_query_text)


class InsertEscapingTestCase(HttpSessionTestCase):
    columns = [')', '(', '%', '{']  # might as well use full `string.ascii`?
    table = Table(
        'test',
        HttpSessionTestCase.metadata(),
        *(
            [Column(name, types.String, primary_key=True)
             for name in columns] +
            [engines.Memory()]))

    def test_insert_data(self):
        data = [{col: '1' for col in self.columns}]

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), data)

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/http/test_escaping.py
datalens-tech/datalens-backend
from sqlalchemy import Column

from clickhouse_sqlalchemy import types, Table
from clickhouse_sqlalchemy_tests.testcase import HttpSessionTestCase


class FormatSectionTestCase(HttpSessionTestCase):

    @property
    def execution_ctx_cls(self):
        return self.session.bind.dialect.execution_ctx_cls

    def _compile(self, clause, bind=None, **kwargs):
        if bind is None:
            bind = self.session.bind
        statement = super(FormatSectionTestCase, self)._compile(
            clause, bind=bind, **kwargs
        )

        context = self.execution_ctx_cls._init_compiled(
            dialect=bind.dialect, connection=bind, dbapi_connection=bind,
            execution_options={}, compiled=statement, parameters=(),
            invoked_statement=statement, extracted_parameters=(),
        )
        context.pre_exec()

        return context.statement

    def test_select_format_clause(self):
        metadata = self.metadata()

        bind = self.session.bind
        bind.cursor = lambda: None

        table = Table(
            't1', metadata,
            Column('x', types.Int32, primary_key=True)
        )

        statement = self.compile(self.session.query(table.c.x), bind=bind)
        self.assertEqual(
            statement,
            'SELECT x AS t1_x FROM t1 FORMAT TabSeparatedWithNamesAndTypes'
        )

    def test_insert_from_select_no_format_clause(self):
        metadata = self.metadata()

        bind = self.session.bind
        bind.cursor = lambda: None

        t1 = Table(
            't1', metadata,
            Column('x', types.Int32, primary_key=True)
        )

        t2 = Table(
            't2', metadata,
            Column('x', types.Int32, primary_key=True)
        )

        query = t2.insert() \
            .from_select(['x'], self.session.query(t1.c.x).subquery())
        statement = self.compile(query, bind=bind)
        self.assertIn(statement, [
            'INSERT INTO t2 (x) SELECT x FROM t1',
            'INSERT INTO t2 (x) SELECT x FROM (SELECT x AS x FROM t1) AS anon_1',
        ])

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/http/test_select.py
datalens-tech/datalens-backend
from clickhouse_sqlalchemy_tests.testcase import HttpSessionTestCase
from clickhouse_sqlalchemy_tests.session import http_stream_session


class StreamingHttpTestCase(HttpSessionTestCase):

    session = http_stream_session
    power = 4

    def make_query(self, power=None):
        if power is None:
            power = self.power
        # `10**power` rows.
        # Timing order of magnitude:
        # est. 1 second on 10**5 lines for TSV,
        # 2.5 seconds on uncythonized native protocol.
        query = 'select ' + ', '.join(
            ("arrayJoin(["
             "'{0}', '2', '3', '4', '5', "
             "'6', '7', '8', '9', '10']) as a{0}").format(num)
            for num in range(power))
        return query

    def test_streaming(self):
        power = self.power
        query = self.make_query(power=power)
        res = self.session.execute(query)
        count = sum(1 for _ in res)
        self.assertEqual(count, 10 ** power)

    def test_fetchmany(self):
        power = self.power - 1
        query = self.make_query(power=power)
        res = self.session.execute(query)

        count = 0
        while True:
            block = res.fetchmany(1000)
            if not block:
                break
            count += sum(1 for _ in block)

        self.assertEqual(count, 10 ** power)

    def test_fetchall(self):
        power = self.power - 2
        if power < 1:
            raise Exception(
                "Misconfigured test case:"
                " `power` should be at least 3.")
        query = self.make_query(power=power)
        res = self.session.execute(query)

        count = 0

        block = res.fetchmany(7)
        count += sum(1 for _ in block)

        block = res.fetchall()
        count += sum(1 for _ in block)

        self.assertEqual(count, 10 ** power)

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/http/test_stream.py
datalens-tech/datalens-backend
from datetime import date, datetime, timezone

import pytz
import sqlalchemy as sa
from responses import mock
from mock import patch
from sqlalchemy import Column, func, create_engine

from clickhouse_sqlalchemy import types, Table, make_session
from clickhouse_sqlalchemy.drivers.http.base import ClickHouseDialect_http
from clickhouse_sqlalchemy_tests.config import http_uri
from clickhouse_sqlalchemy_tests.session import http_tsv_session
from clickhouse_sqlalchemy_tests.testcase import HttpSessionTestCase


class TransportCase(HttpSessionTestCase):

    # Mocked data is in TSV format.
    session = http_tsv_session

    @property
    def url(self):
        return 'http://{host}:{port}'.format(host=self.host, port=self.port)

    @mock.activate
    def test_parse_func_count(self):
        mock.add(
            mock.POST, self.url, status=200,
            body='count_1\nUInt64\n42\n'
        )

        table = Table(
            't1', self.metadata(),
            Column('x', types.Int32, primary_key=True)
        )

        rv = self.session.query(func.count()).select_from(table).scalar()
        self.assertEqual(rv, 42)

    @mock.activate
    def test_parse_int_types(self):
        types_ = [
            'Int8', 'UInt8', 'Int16', 'UInt16', 'Int32', 'UInt32', 'Int64',
            'UInt64'
        ]
        columns = [chr(i + ord('a')) for i in range(len(types_))]

        mock.add(
            mock.POST, self.url, status=200,
            body=(
                '\t'.join(columns) + '\n' +
                '\t'.join(types_) + '\n' +
                '\t'.join(['42'] * len(types_)) + '\n'
            )
        )

        table = Table(
            't1', self.metadata(),
            *[Column(col, types.Int) for col in columns]
        )

        rv = self.session.query(*table.c).first()
        self.assertEqual(rv, tuple([42] * len(columns)))

    @mock.activate
    def test_parse_float_types(self):
        types_ = ['Float32', 'Float64']
        columns = ['a', 'b']

        mock.add(
            mock.POST, self.url, status=200,
            body=(
                '\t'.join(columns) + '\n' +
                '\t'.join(types_) + '\n' +
                '\t'.join(['42'] * len(types_)) + '\n'
            )
        )

        table = Table(
            't1', self.metadata(),
            *[Column(col, types.Float) for col in columns]
        )

        rv = self.session.query(*table.c).first()
        self.assertEqual(rv, tuple([42.0] * len(columns)))

    # do not call that method
    @patch.object(ClickHouseDialect_http, '_get_server_version_info')
    @mock.activate
    def test_parse_date_types(self, patched_server_info):
        mock.add(
            mock.POST, self.url, status=200,
            body=(
                'a\n' +
                'Date\n' +
                '2012-10-25\n'
            )
        )

        table = Table(
            't1', self.metadata(),
            Column('a', types.Date)
        )

        rv = self.session.query(*table.c).first()
        self.assertEqual(rv, (date(2012, 10, 25), ))

    @mock.activate
    def test_parse_datetime_type(self):
        mock.add(
            mock.POST, 'http://localhost:8123', status=200,
            body=(
                'a\n' +
                'DateTime\n' +
                '2012-10-25 01:02:03\n'
            )
        )

        table = Table(
            't1', self.metadata(),
            Column('a', types.DateTime)
        )

        rv = (
            self.session.query(*table.c)
            .filter(table.c.a > datetime(2000, 12, 23, 13, 14, 15, 123))
            .first())
        self.assertEqual(rv, (datetime(2012, 10, 25, 1, 2, 3), ))

    @mock.activate
    def test_parse_datetime_utc_type(self):
        mock.add(
            mock.POST, 'http://localhost:8123', status=200,
            body=(
                'a\n' +
                'DateTime(\'UTC\')\n' +
                '2012-10-25 01:02:03\n'
            )
        )

        table = Table(
            't1', self.metadata(),
            Column('a', types.DateTime)
        )

        rv = self.session.query(*table.c).first()
        expected_value = (
            datetime(2012, 10, 25, 1, 2, 3)
            .replace(tzinfo=timezone.utc))
        self.assertEqual(rv, (expected_value,))

    @mock.activate
    def test_parse_datetime_any_tz_type(self):
        mock.add(
            mock.POST, 'http://localhost:8123', status=200,
            body=(
                'a\n' +
                'DateTime(\'MSK\')\n' +
                '2012-10-25 01:02:03\n'
            )
        )

        table = Table(
            't1', self.metadata(),
            Column('a', types.DateTime)
        )

        rv = self.session.query(*table.c).first()
        expected_value = pytz.timezone('Europe/Moscow').localize(
            datetime(2012, 10, 25, 1, 2, 3))
        self.assertEqual(rv, (expected_value,))

    @mock.activate
    def test_parse_nullable_datetime_any_tz_type(self):
        mock.add(
            mock.POST, 'http://localhost:8123', status=200,
            body=(
                'a\n' +
                'Nullable(DateTime(\'EST\'))\n' +
                '2012-10-25 01:02:03\n'
            )
        )

        table = Table(
            't1', self.metadata(),
            Column('a', types.DateTime)
        )

        rv = self.session.query(*table.c).first()
        expected_value = (
            datetime(2012, 10, 25, 1, 2, 3)
            .replace(tzinfo=pytz.timezone('EST')))
        self.assertEqual(rv, (expected_value,))

    @mock.activate
    def test_parse_nullable_type(self):
        mock.add(
            mock.POST, self.url, status=200,
            body=(
                'a\n' +
                'String\n' +
                '\\N\n' +
                '\\\\N\n' +
                '\n'
            )
        )

        table = Table(
            't1', self.metadata(),
            Column('a', types.String)
        )

        rv = self.session.query(*table.c).all()
        self.assertEqual(rv, [(None, ), ('\\N', ), ('', )])


class TransportHeadersCase(TransportCase):

    @mock.activate
    def test_dynamic_headers(self):
        header_value = 'h1'
        engine = create_engine(
            http_uri,
            connect_args={
                'server_version': '20.3.0.0',  # otherwise will break under `mock`.
                'header__X-Test-Header': lambda *_, **__: header_value,
            },
        )
        session = make_session(engine)

        mock.add(
            mock.POST, self.url, status=200,
            body='a\nUInt8\n1\n'
        )

        session.query(sa.text('select 1 as a')).scalar()
        assert mock.calls[-1].request.headers['X-Test-Header'] == header_value

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/http/test_transport.py
datalens-tech/datalens-backend
from clickhouse_sqlalchemy.drivers.http.utils import unescape, parse_tsv
from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


class HttpUtilsTestCase(BaseTestCase):
    def test_unescape(self):
        test_values = [b'', b'a', b'\xff']
        actual = [unescape(t) for t in test_values]
        self.assertListEqual(actual, [u'', u'a', u'\ufffd'])

    def test_parse_tsv(self):
        test_values = [b'', b'a\tb\tc', b'a\tb\t\xff']
        try:
            for value in test_values:
                parse_tsv(value)
        except IndexError:
            self.fail('"parse_tsv" raised IndexError exception!')
        except TypeError:
            self.fail('"parse_tsv" raised TypeError exception!')

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/http/test_utils.py
datalens-tech/datalens-backend
from sqlalchemy.engine.url import URL

from clickhouse_sqlalchemy.drivers.native.base import ClickHouseDialect_native
from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


class TestConnectArgs(BaseTestCase):
    def setUp(self):
        self.dialect = ClickHouseDialect_native()

    def test_simple_url(self):
        url = URL(
            drivername='clickhouse+native',
            host='localhost',
            database='default',
        )
        connect_args = self.dialect.create_connect_args(url)
        self.assertEqual(
            str(connect_args[0][0]), 'clickhouse://localhost/default'
        )

    def test_secure_false(self):
        url = URL(
            drivername='clickhouse+native',
            username='default',
            password='default',
            host='localhost',
            port='9001',
            database='default',
            query={'secure': 'False'}
        )
        connect_args = self.dialect.create_connect_args(url)
        self.assertEqual(
            str(connect_args[0][0]),
            'clickhouse://default:default@localhost:9001/default?secure=False'
        )

    def test_no_auth(self):
        url = URL(
            drivername='clickhouse+native',
            host='localhost',
            port='9001',
            database='default',
        )
        connect_args = self.dialect.create_connect_args(url)
        self.assertEqual(
            str(connect_args[0][0]), 'clickhouse://localhost:9001/default'
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/native/test_base.py
datalens-tech/datalens-backend
from clickhouse_sqlalchemy_tests.testcase import NativeSessionTestCase


class CursorTestCase(NativeSessionTestCase):
    def test_execute_without_context(self):
        raw = self.session.bind.raw_connection()
        cur = raw.cursor()

        cur.execute("SELECT * FROM system.numbers LIMIT 1")
        rv = cur.fetchall()
        self.assertEqual(len(rv), 1)

    def test_execute_with_context(self):
        rv = self.session.execute("SELECT * FROM system.numbers LIMIT 1")

        self.assertEqual(len(rv.fetchall()), 1)

    def test_check_iter_cursor(self):
        rv = self.session.execute('SELECT number FROM system.numbers LIMIT 5')
        self.assertListEqual(list(rv), [(x,) for x in range(5)])

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/native/test_cursor.py
datalens-tech/datalens-backend
from sqlalchemy import Column, func

from clickhouse_sqlalchemy import engines, types, Table
from clickhouse_sqlalchemy.exceptions import DatabaseException
from clickhouse_sqlalchemy_tests.session import native_session
from clickhouse_sqlalchemy_tests.testcase import NativeSessionTestCase


class NativeInsertTestCase(NativeSessionTestCase):
    session = native_session

    def test_rowcount_return(self):
        table = Table(
            'test', self.metadata(),
            Column('x', types.Int32, primary_key=True),
            engines.Memory()
        )
        table.drop(if_exists=True)
        table.create()

        rv = self.session.execute(table.insert(), [{'x': x} for x in range(5)])
        self.assertEqual(rv.rowcount, 5)
        self.assertEqual(
            self.session.query(func.count()).select_from(table).scalar(), 5
        )

        rv = self.session.execute(
            'INSERT INTO test SELECT * FROM system.numbers LIMIT 5'
        )
        self.assertEqual(rv.rowcount, -1)

    def test_types_check(self):
        table = Table(
            'test', self.metadata(),
            Column('x', types.UInt32, primary_key=True),
            engines.Memory()
        )
        table.drop(if_exists=True)
        table.create()

        with self.assertRaises(DatabaseException) as ex:
            self.session.execute(
                table.insert().execution_options(types_check=True),
                [{'x': -1}]
            )
        self.assertIn('-1 for column "x"', str(ex.exception.orig))

        with self.assertRaises(DatabaseException) as ex:
            self.session.execute(table.insert(), [{'x': -1}])
        self.assertIn(
            'Repeat query with types_check=True for detailed info',
            str(ex.exception.orig)
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/native/test_insert.py
datalens-tech/datalens-backend
from sqlalchemy import Column

from clickhouse_sqlalchemy import engines, types, Table
from clickhouse_sqlalchemy_tests.session import native_session, mocked_engine
from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


class SanityTestCase(BaseTestCase):
    session = native_session

    def test_sanity(self):
        with mocked_engine(self.session) as engine:
            table = Table(
                't1', self.metadata(session=engine.session),
                Column('x', types.Int32, primary_key=True),
                engines.Memory()
            )
            table.drop(if_exists=True)
            self.assertEqual(engine.history, ['DROP TABLE IF EXISTS t1'])

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/native/test_select.py
datalens-tech/datalens-backend
from sqlalchemy import Column, create_engine

from clickhouse_sqlalchemy import make_session, engines, types, Table
from clickhouse_sqlalchemy_tests.testcase import (
    BaseAbstractTestCase, BaseTestCase,
    HttpSessionTestCase, NativeSessionTestCase,
)
from clickhouse_sqlalchemy_tests.config import (
    system_http_uri, system_native_uri,
    database as test_database,
)


class ClickHouseDialectTestCase(BaseAbstractTestCase):

    @property
    def dialect(self):
        return self.session.bind.dialect

    def setUp(self):
        super(ClickHouseDialectTestCase, self).setUp()
        self.table = Table(
            'test_exists_table',
            self.metadata(),
            Column('x', types.Int32, primary_key=True),
            engines.Memory()
        )
        self.table.drop(self.session.bind, if_exists=True)

    def tearDown(self):
        self.table.drop(self.session.bind, if_exists=True)
        super(ClickHouseDialectTestCase, self).tearDown()

    def test_has_table(self):
        self.assertFalse(
            self.dialect.has_table(self.session,
                                   self.table.name)
        )

        self.table.create(self.session.bind)
        self.assertTrue(
            self.dialect.has_table(self.session,
                                   self.table.name)
        )

    def test_get_table_names(self):
        self.table.create(self.session.bind)
        db_tables = self.dialect.get_table_names(self.session)
        self.assertIn(self.table.name, db_tables)

    def test_get_schema_names(self):
        schemas = self.dialect.get_schema_names(self.session)
        self.assertIn(test_database, schemas)


class ClickHouseDialectHttpTestCase(
        ClickHouseDialectTestCase, HttpSessionTestCase):
    """ HTTP version """


class ClickHouseDialectNativeTestCase(
        ClickHouseDialectTestCase, NativeSessionTestCase):
    """ Native version """


class CachedServerVersionTestCase(BaseTestCase):

    def _test_server_version_any(self, uri):
        engine_session = make_session(create_engine(
            system_http_uri,
            connect_args=dict(server_version='123.45.67.89')))
        # Connection and version get initialized on first query:
        engine_session.scalar('select 1')
        ver = engine_session.get_bind().dialect.server_version_info
        self.assertEqual(ver, (123, 45, 67, 89))

    def test_server_version_http(self):
        return self._test_server_version_any(system_http_uri)

    def test_server_version_native(self):
        return self._test_server_version_any(system_native_uri)

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/drivers/test_clickhouse_dialect.py
datalens-tech/datalens-backend
from sqlalchemy import Column
from sqlalchemy.sql.ddl import CreateTable

from clickhouse_sqlalchemy import types, engines, get_declarative_base
from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


class DeclarativeTestCase(BaseTestCase):
    def test_create_table(self):
        base = get_declarative_base()

        class TestTable(base):
            x = Column(types.Int32, primary_key=True)
            y = Column(types.String)

            __table_args__ = (
                engines.Memory(),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            'CREATE TABLE test_table (x Int32, y String) ENGINE = Memory'
        )

    def test_create_table_custom_name(self):
        base = get_declarative_base()

        class TestTable(base):
            __tablename__ = 'testtable'

            x = Column(types.Int32, primary_key=True)
            y = Column(types.String)

            __table_args__ = (
                engines.Memory(),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            'CREATE TABLE testtable (x Int32, y String) ENGINE = Memory'
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/ext/test_declative.py
datalens-tech/datalens-backend
from sqlalchemy import Column, func

from clickhouse_sqlalchemy import types, Table
from clickhouse_sqlalchemy_tests.testcase import (
    BaseAbstractTestCase, HttpSessionTestCase, NativeSessionTestCase,
)


class CountTestCaseBase(BaseAbstractTestCase):
    def create_table(self):
        metadata = self.metadata()

        return Table(
            't1', metadata,
            Column('x', types.Int32, primary_key=True)
        )

    def test_count(self):
        table = self.create_table()

        self.assertEqual(
            self.compile(self.session.query(func.count(table.c.x))),
            'SELECT count(x) AS count_1 FROM t1'
        )

    def test_count_distinct(self):
        table = self.create_table()
        query = self.session.query(func.count(func.distinct(table.c.x)))
        self.assertEqual(
            self.compile(query),
            'SELECT count(distinct(x)) AS count_1 FROM t1'
        )

    def test_count_no_column_specified(self):
        table = self.create_table()
        query = self.session.query(func.count()).select_from(table)
        self.assertEqual(
            self.compile(query),
            'SELECT count(*) AS count_1 FROM t1'
        )


class CountHttpTestCase(CountTestCaseBase, HttpSessionTestCase):
    """ ... """


class CountNativeTestCase(CountTestCaseBase, NativeSessionTestCase):
    """ ... """

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/functions/test_count.py
datalens-tech/datalens-backend
"""
Test EXTRACT
"""
from sqlalchemy import Column, extract

from clickhouse_sqlalchemy import types
from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


def get_date_column(name):
    """
    Return types.Date column
    :param name: Column name
    :return: sqlalchemy.Column
    """
    return Column(name, types.Date)


class ExtractTestCase(BaseTestCase):
    def test_extract_year(self):
        self.assertEqual(
            self.compile(extract('year', get_date_column('x'))),
            'toYear(x)'
        )

    def test_extract_month(self):
        self.assertEqual(
            self.compile(extract('month', get_date_column('x'))),
            'toMonth(x)'
        )

    def test_extract_day(self):
        self.assertEqual(
            self.compile(extract('day', get_date_column('x'))),
            'toDayOfMonth(x)'
        )

    def test_extract_unknown(self):
        self.assertEqual(
            self.compile(extract('test', get_date_column('x'))),
            'x'
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/functions/test_extract.py
datalens-tech/datalens-backend
from sqlalchemy import literal, func, text

from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


class IfTestCase(BaseTestCase):
    def test_if(self):
        expression = func.if_(
            literal(1) > literal(2),
            text('a'),
            text('b'),
        )

        self.assertEqual(
            self.compile(expression, literal_binds=True),
            'if(1 > 2, a, b)'
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/functions/test_if.py
datalens-tech/datalens-backend
from sqlalchemy import (
    Column,
    exc,
    func,
    literal,
    select,
)
from sqlalchemy import text
from sqlalchemy import literal_column
from sqlalchemy import tuple_

from clickhouse_sqlalchemy import types, Table
from clickhouse_sqlalchemy.ext.clauses import Lambda, ParametrizedFunction
from clickhouse_sqlalchemy_tests.testcase import (
    BaseAbstractTestCase, HttpSessionTestCase, NativeSessionTestCase,
)


class SelectTestCase(BaseAbstractTestCase):

    def create_table(self, *columns):
        metadata = self.metadata()

        return Table(
            't1', metadata,
            Column('x', types.Int32, primary_key=True),
            *columns
        )

    def test_select(self):
        table = self.create_table()

        query = self.session.query(table.c.x)\
            .filter(table.c.x.in_([1, 2]))\
            .having(func.count() > 0)\
            .order_by(table.c.x.desc())
        self.assertEqual(
            self.compile(query),
            'SELECT x AS t1_x '
            'FROM t1 '
            'WHERE x IN (__[POSTCOMPILE_x_1]) '
            'HAVING count(*) > %(count_1)s '
            'ORDER BY x DESC'
        )

    def test_very_simple_select(self):
        """
        A non-CH specific select statement should work too.
        """
        query = select([literal(1).label('col1')])
        self.assertEqual(
            self.compile(query),
            'SELECT %(param_1)s AS col1'
        )

    def test_group_by_query(self):
        table = self.create_table()

        query = self.session.query(table.c.x).group_by(table.c.x)
        self.assertEqual(
            self.compile(query),
            'SELECT x AS t1_x FROM t1 GROUP BY x'
        )

        query = self.session.query(table.c.x).group_by(table.c.x).with_totals()
        self.assertEqual(
            self.compile(query),
            'SELECT x AS t1_x FROM t1 GROUP BY x WITH TOTALS'
        )

        with self.assertRaises(exc.InvalidRequestError) as ex:
            self.session.query(table.c.x).with_totals()

        self.assertIn('with_totals', str(ex.exception))

    def test_array_join(self):
        table = self.create_table(
            Column('nested.array_column', types.Array(types.Int8)),
            Column('nested.another_array_column', types.Array(types.Int8))
        )
        first_label = table.c['nested.array_column'].label('from_array')
        second_not_label = table.c['nested.another_array_column']
        query = self.session.query(first_label, second_not_label)\
            .array_join(first_label, second_not_label)
        self.assertEqual(
            self.compile(query),
            'SELECT '
            '"nested.array_column" AS from_array, '
            '"nested.another_array_column" '
            'AS "t1_nested.another_array_column" '
            'FROM t1 '
            'ARRAY JOIN "nested.array_column" AS from_array, '
            '"nested.another_array_column"'
        )

    def test_sample(self):
        table = self.create_table()

        query = self.session.query(table.c.x).sample(0.1).group_by(table.c.x)
        self.assertEqual(
            self.compile(query),
            'SELECT x AS t1_x FROM t1 SAMPLE %(param_1)s GROUP BY x'
        )
        self.assertEqual(
            self.compile(query, literal_binds=True),
            'SELECT x AS t1_x FROM t1 SAMPLE 0.1 GROUP BY x'
        )

    def test_label_read(self):
        query_text = (
            r'select '
            r'0 as "a z", '
            r'2 as "a \\z", '
            r'4 as "a \\\\z", '
            r'6 as "a \\\\\\z"')
        result = self.session.execute(query_text)

        result_keys = result.keys()
        expected = [r'a z', r'a \z', r'a \\z', r'a \\\z']
        self.assertEqual(result_keys, expected)
        result_columns = [item[0] for item in result.cursor.description]
        self.assertEqual(result_columns, expected)
        self.assertEqual(len(result.fetchall()), 1)

    def test_label_define(self):
        columns = [r'a z', r'a \z', r'a \\z', r'a \\\z', 'a \'\n"\\z']
        query = self.session.query(*[
            literal(1).label(col)
            for col in columns])
        query_text = self.compile(query, literal_binds=True)
        query_text = query_text % ()  # dbapi interpolate level
        expected_query_text = (
            r'SELECT '
            r'1 AS "a z", '
            r'1 AS "a \\z", '
            r'1 AS "a \\\\z", '
            r'1 AS "a \\\\\\z", '
            '1 AS "a \\\'\\n""\\\\z"'
        )
        self.assertEqual(query_text, expected_query_text)

        result = self.session.execute(query)

        result_keys = result.keys()
        self.assertEqual(result_keys, columns)

        result_columns = [item[0] for item in result.raw.cursor.description]
        self.assertEqual(result_columns, columns)

        self.assertEqual(len(result.fetchmany(3)), 1)

    def test_like_simple(self):
        query = self.session.query(
            literal('Hello World%').like('%World%').label('value'),
        )
        self.assertEqual(
            self.compile(query, literal_binds=True),
            # Note: '%%' is for the python string interpolation
            # (dbapi layer, where parameters are applied);
            # the database should receive undoubled percentages.
            "SELECT 'Hello World%%' LIKE '%%World%%' AS value"
        )

    def test_like_extended(self):
        col = literal('Hello World%').label('value')
        query = self.session.query(col)
        query = query.filter(col.like('%World%'))
        result = self.session.execute(query)
        result = list(result)
        assert result == [('Hello World%',)]

    def test_like_full(self):
        values = [
            r'a b c % d e',
            r'a b c %% d e',
            r'a b c \% d e',
            r'a b c \\% d e',
        ]
        matches = [
            # r' c % d ' in value
            r'% c \% d %',
            # ' c \\' in value, ' d ' in value
            # (semantically significant '%')
            r'% c \\% d %',
            # r' c \% d ' in value
            r'% c \\\% d %',
            # r' b % d ' in value
            r'% b \% d %',
            # The rest should match all of the values
            # (semantically significant '%'):
            r'% c % d %', r'% c %% d %',
            r'% b % d %', r'% b %% d %',
        ]

        col = func.arrayJoin(values).label('value')
        col_ref = literal_column('value')
        query = self.session.query(
            col,
            *[
                col_ref.like(ptn).label('m{}'.format(idx))
                for idx, ptn in enumerate(matches)])

        query_text = self.compile(query, literal_binds=True)
        query_text = query_text % ()  # dbapi interpolate level
        expected_query_text = (
            r"SELECT arrayJoin(["
            r"'a b c % d e', 'a b c %% d e', "
            r"'a b c \\% d e', 'a b c \\\\% d e'"
            r"]) AS value, "
            r"value LIKE '% c \\% d %' AS m0, "
            r"value LIKE '% c \\\\% d %' AS m1, "
            r"value LIKE '% c \\\\\\% d %' AS m2, "
            r"value LIKE '% b \\% d %' AS m3, "
            r"value LIKE '% c % d %' AS m4, "
            r"value LIKE '% c %% d %' AS m5, "
            r"value LIKE '% b % d %' AS m6, "
            r"value LIKE '% b %% d %' AS m7"
        )
        self.assertEqual(query_text, expected_query_text)

        result = self.session.execute(query)

        result = list(result)
        result.sort()
        self.assertEqual(
            result,
            # value,
            # ' c % d ', 'c \\' … ' d', r'c \% d', 'b % d' (never),
            # ... (always)
            [(r'a b c % d e',
              True, False, False, False,
              True, True, True, True),
             (r'a b c %% d e',
              False, False, False, False,
              True, True, True, True),
             (r'a b c \% d e',
              False, True, True, False,
              True, True, True, True),
             (r'a b c \\% d e',
              False, True, False, False,
              True, True, True, True)])

    def test_parametrized_function(self):
        values = [1, 2, 3, 4, 5]
        table = select([func.arrayJoin(values).label('value')]).alias('sq')
        col_ref = literal_column('value')
        expr = ParametrizedFunction(
            'quantileExact',
            [literal(0.75)],
            col_ref)
        query = self.session.query(expr.label('result')).select_from(table)
        self.assertEqual(
            self.compile(query),
            'SELECT quantileExact(%(param_1)s)(value) AS result '
            'FROM (SELECT arrayJoin(%(arrayJoin_1)s) AS value) AS sq'
        )
        result = self.session.scalar(query)
        self.assertEqual(result, 4)

    def test_final(self):
        table = self.create_table()

        query = self.session.query(table.c.x).final().group_by(table.c.x)
        self.assertEqual(
            self.compile(query),
            'SELECT x AS t1_x FROM t1 FINAL GROUP BY x'
        )

    def test_lambda_functions(self):
        query = self.session.query(
            func.arrayFilter(
                Lambda(lambda x: x.like('%World%')),
                literal(['Hello', 'abc World'], types.Array(types.String))
            ).label('test')
        )
        self.assertEqual(
            self.compile(query, literal_binds=True),
            "SELECT arrayFilter("
            # See `def test_like` about the double percentages.
            "x -> x LIKE '%%World%%', ['Hello', 'abc World']"
            ") AS test"
        )


class SelectHttpTestCase(SelectTestCase, HttpSessionTestCase):
    """ ... """


class SelectNativeTestCase(SelectTestCase, NativeSessionTestCase):
    """ ... """


class JoinTestCase(BaseAbstractTestCase):
    def create_tables(self, num):
        metadata = self.metadata()

        return [Table(
            't{}'.format(i), metadata,
            Column('x', types.Int32, primary_key=True),
            Column('y', types.Int32, primary_key=True),
        ) for i in range(num)]

    def test_joins(self):
        t1, t2 = self.create_tables(2)

        query = self.session.query(t1.c.x, t2.c.x) \
            .join(
            t2,
            t1.c.x == t1.c.y,
            strictness='any')

        self.assertEqual(
            self.compile(query),
            "SELECT x AS t0_x, x AS t1_x FROM t0 "
            "ANY INNER JOIN t1 ON x = y"
        )

        query = self.session.query(t1.c.x, t2.c.x) \
            .join(
            t2,
            t1.c.x == t1.c.y,
            type='inner',
            strictness='any')

        self.assertEqual(
            self.compile(query),
            "SELECT x AS t0_x, x AS t1_x FROM t0 "
            "ANY INNER JOIN t1 ON x = y"
        )

        query = self.session.query(t1.c.x, t2.c.x) \
            .join(
            t2,
            tuple_(t1.c.x, t1.c.y),
            type='inner',
            strictness='all'
        )

        self.assertEqual(
            self.compile(query),
            "SELECT x AS t0_x, x AS t1_x FROM t0 "
            "ALL INNER JOIN t1 USING (x, y)"
        )

        query = self.session.query(t1.c.x, t2.c.x) \
            .join(t2,
                  tuple_(t1.c.x, t1.c.y),
                  type='inner',
                  strictness='all',
                  distribution='global')

        self.assertEqual(
            self.compile(query),
            "SELECT x AS t0_x, x AS t1_x FROM t0 "
            "GLOBAL ALL INNER JOIN t1 USING (x, y)"
        )

        query = self.session.query(t1.c.x, t2.c.x) \
            .outerjoin(
            t2,
            tuple_(t1.c.x, t1.c.y),
            type='left outer',
            strictness='all',
            distribution='global'
        )

        self.assertEqual(
            self.compile(query),
            "SELECT x AS t0_x, x AS t1_x FROM t0 "
            "GLOBAL ALL LEFT OUTER JOIN t1 USING (x, y)"
        )

        query = self.session.query(t1.c.x, t2.c.x) \
            .outerjoin(
            t2,
            tuple_(t1.c.x, t1.c.y),
            type='LEFT OUTER',
            strictness='ALL',
            distribution='GLOBAL'
        )

        self.assertEqual(
            self.compile(query),
            "SELECT x AS t0_x, x AS t1_x FROM t0 "
            "GLOBAL ALL LEFT OUTER JOIN t1 USING (x, y)"
        )

        query = self.session.query(t1.c.x, t2.c.x) \
            .outerjoin(t2,
                       tuple_(t1.c.x, t1.c.y),
                       strictness='ALL',
                       type='FULL OUTER')

        self.assertEqual(
            self.compile(query),
            "SELECT x AS t0_x, x AS t1_x FROM t0 "
            "ALL FULL OUTER JOIN t1 USING (x, y)"
        )


class JoinHttpTestCase(JoinTestCase, HttpSessionTestCase):
    """ Join over a HTTP session """


class JoinNativeTestCase(JoinTestCase, NativeSessionTestCase):
    """ Join over a native protocol session """


class YieldTest(NativeSessionTestCase):
    def test_yield_per_and_execution_options(self):
        numbers = Table(
            'numbers', self.metadata(),
            Column('number', types.UInt64, primary_key=True),
        )

        query = self.session.query(numbers.c.number).limit(100).yield_per(15)
        query = query.execution_options(foo='bar')
        self.assertIsNotNone(query.load_options._yield_per)
        self.assertEqual(
            query._execution_options,
            {
                # 'stream_results': True,
                'foo': 'bar',
                # 'max_row_buffer': 15,
            }
        )

    def test_basic(self):
        numbers = Table(
            'numbers', self.metadata(),
            Column('number', types.UInt64, primary_key=True),
        )

        q = iter(
            self.session.query(numbers.c.number)
            .yield_per(1)
            .from_statement(text('SELECT * FROM system.numbers LIMIT 3'))
        )

        ret = []
        ret.append(next(q))
        ret.append(next(q))
        ret.append(next(q))
        try:
            next(q)
            self.assertTrue(False)
        except StopIteration:
            pass

        self.assertEqual(ret, [(0, ), (1, ), (2, )])

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/orm/test_select.py
datalens-tech/datalens-backend
import re

from sqlalchemy import create_engine

from clickhouse_sqlalchemy import make_session

from .config import (
    http_uri, native_uri, system_native_uri, system_http_uri,
)

http_session = make_session(create_engine(http_uri))
http_stream_session = make_session(create_engine(http_uri + '?stream=1'))
http_tsv_session = make_session(create_engine(
    http_uri + '?format=TabSeparatedWithNamesAndTypes'))
http_jc_session = make_session(create_engine(
    http_uri + '?format=JSONCompact&stream=1'))

native_session = make_session(create_engine(native_uri))

system_native_session = make_session(create_engine(system_native_uri))
system_http_session = make_session(create_engine(system_http_uri))


class MockedEngine(object):

    prev_do_execute = None
    prev_do_executemany = None
    prev_query_server_version_string = None

    def __init__(self, session=None):
        self._buffer = []

        if session is None:
            session = make_session(create_engine(http_uri))

        self.session = session
        self.dialect_cls = session.bind.dialect.__class__

    @property
    def history(self):
        return [re.sub(r'[\n\t]', '', str(ssql)) for ssql in self._buffer]

    def __enter__(self):
        self.prev_do_execute = self.dialect_cls.do_execute
        self.prev_do_executemany = self.dialect_cls.do_executemany
        self.prev_query_server_version_string = \
            self.dialect_cls._query_server_version_string

        def do_executemany(
                instance, cursor, statement, parameters, context=None):
            self._buffer.append(statement)

        def do_execute(instance, cursor, statement, parameters, context=None):
            self._buffer.append(statement)

        def query_server_version_string(*args, **kwargs):
            return '19.16.2.2'

        self.dialect_cls.do_execute = do_execute
        self.dialect_cls.do_executemany = do_executemany
        self.dialect_cls._query_server_version_string = \
            query_server_version_string

        return self

    def __exit__(self, *exc_info):
        self.dialect_cls.do_execute = self.prev_do_execute
        self.dialect_cls.do_executemany = self.prev_do_executemany
        self.dialect_cls._query_server_version_string = \
            self.prev_query_server_version_string


mocked_engine = MockedEngine

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/session.py
datalens-tech/datalens-backend
from sqlalchemy import literal, not_

from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


class NotBetweenTestCase(BaseTestCase):
    def test_notbetween(self):
        self.assertEqual(
            self.compile(not_(literal(2).between(1, 3)), literal_binds=True),
            '2 NOT BETWEEN 1 AND 3'
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/sql/test_basic_operators.py
datalens-tech/datalens-backend
from sqlalchemy import literal, exc, case

from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


class CaseTestCase(BaseTestCase):
    def test_else_required(self):
        with self.assertRaises(exc.CompileError) as ex:
            self.compile(case([(literal(1), 0)]))

        self.assertEqual(
            str(ex.exception),
            'ELSE clause is required in CASE'
        )

    def test_case(self):
        expression = case([(literal(1), 0)], else_=1)
        self.assertEqual(
            self.compile(expression, literal_binds=True),
            'CASE WHEN 1 THEN 0 ELSE 1 END'
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/sql/test_case.py
datalens-tech/datalens-backend
from sqlalchemy import exc

from clickhouse_sqlalchemy.ext.clauses import Lambda
from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


class LabmdaTestCase(BaseTestCase):
    def test_lambda_is_callable(self):
        with self.assertRaises(exc.ArgumentError) as ex:
            self.compile(Lambda(1))

        self.assertEqual(str(ex.exception), 'func must be callable')

    def test_lambda_no_arg_kwargs(self):
        with self.assertRaises(exc.CompileError) as ex:
            self.compile(Lambda(lambda x, *args: x * 2))

        self.assertEqual(
            str(ex.exception),
            'Lambdas with *args are not supported'
        )

        with self.assertRaises(exc.CompileError) as ex:
            self.compile(Lambda(lambda x, **kwargs: x * 2))

        self.assertEqual(
            str(ex.exception),
            'Lambdas with **kwargs are not supported'
        )

    def test_lambda_ok(self):
        self.assertEqual(
            self.compile(Lambda(lambda x: x * 2), literal_binds=True),
            'x -> x * 2'
        )

        def mul(x):
            return x * 2

        self.assertEqual(
            self.compile(Lambda(mul), literal_binds=True),
            'x -> x * 2'
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/sql/test_lambda.py
datalens-tech/datalens-backend
from sqlalchemy import Column, exc

from clickhouse_sqlalchemy import types, Table
from clickhouse_sqlalchemy_tests.testcase import (
    BaseAbstractTestCase, HttpSessionTestCase, NativeSessionTestCase,
)


class LimitTestCase(BaseAbstractTestCase):
    def create_table(self):
        metadata = self.metadata()

        return Table(
            't1', metadata,
            Column('x', types.Int32, primary_key=True)
        )

    def test_limit(self):
        table = self.create_table()

        query = self.session.query(table.c.x).limit(10)
        self.assertEqual(
            self.compile(query, literal_binds=True),
            'SELECT x AS t1_x FROM t1  LIMIT 10'
        )

    def test_limit_with_offset(self):
        table = self.create_table()

        query = self.session.query(table.c.x).limit(10).offset(5)
        self.assertEqual(
            self.compile(query, literal_binds=True),
            'SELECT x AS t1_x FROM t1  LIMIT 5, 10'
        )

        query = self.session.query(table.c.x).offset(5).limit(10)
        self.assertEqual(
            self.compile(query, literal_binds=True),
            'SELECT x AS t1_x FROM t1  LIMIT 5, 10'
        )

    def test_offset_without_limit(self):
        table = self.create_table()

        with self.assertRaises(exc.CompileError) as ex:
            query = self.session.query(table.c.x).offset(5)
            self.compile(query, literal_binds=True)

        self.assertEqual(
            str(ex.exception),
            'OFFSET without LIMIT is not supported'
        )


class LimitHttpTestCase(LimitTestCase, HttpSessionTestCase):
    """ Limit over a HTTP session """


class LimitNativeTestCase(LimitTestCase, NativeSessionTestCase):
    """ Limit over a native protocol session """

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/sql/test_limit.py
datalens-tech/datalens-backend

from sqlalchemy import Column, text, Table, inspect
from sqlalchemy.sql.elements import TextClause

from clickhouse_sqlalchemy import types, Table as CHTable, engines
from clickhouse_sqlalchemy_tests.testcase import (
    BaseAbstractTestCase, HttpSessionTestCase, NativeSessionTestCase,
)


class SchemaTestCase(BaseAbstractTestCase):

    def test_reflect(self):
        """
        checking, that after call metadata.reflect()
        we have a clickhouse-specific table, which has overridden join methods
        """
        session = self.session
        metadata = self.metadata()

        # ### Maybe: ###
        # # session = native_session  # database=test
        # session = self.__class__.session
        # same as in `self.metadata()`  # database=default
        # unbound_metadata = MetaData(bind=session.bind)

        table = CHTable(
            'test_reflect',
            metadata,
            Column('x', types.Int32),
            engines.Log()
        )

        table.drop(session.bind, if_exists=True)
        table.create(session.bind)

        # Sub-test: ensure the `metadata.reflect` makes a CHTable
        metadata.clear()  # reflect from clean state
        self.assertFalse(metadata.tables)
        metadata.reflect(only=[table.name])
        table2 = metadata.tables.get(table.name)
        self.assertIsNotNone(table2)
        self.assertListEqual([c.name for c in table2.columns], ['x'])
        self.assertTrue(isinstance(table2, CHTable))

        # Sub-test: ensure `CHTable(..., autoload=True)` works too
        metadata.clear()
        table3 = CHTable('test_reflect', metadata, autoload=True)
        self.assertListEqual([c.name for c in table3.columns], ['x'])

        # Sub-test: check that they all reflected the same.
        for table_x in [table, table2, table3]:
            query = table_x.select().join(
                text('another_table'),
                table.c.x == 'xxx',
                type='INNER',
                strictness='ALL',
                distribution='GLOBAL'
            )
            self.assertEqual(
                self.compile(query),
                "SELECT x FROM test_reflect "
                "GLOBAL ALL INNER JOIN another_table "
                "ON x = %(x_1)s"
            )

    def test_reflect_generic_table(self):
        """
        checking, that generic table columns are reflected properly
        """
        metadata = self.metadata()

        table = Table(
            'test_reflect',
            metadata,
            Column('x', types.Int32),
            engines.Log()
        )

        self.session.execute('DROP TABLE IF EXISTS test_reflect')
        table.create(self.session.bind)

        # Sub-test: ensure the `metadata.reflect` makes a CHTable
        metadata.clear()  # reflect from clean state
        self.assertFalse(metadata.tables)

        table = Table('test_reflect', metadata, autoload=True)
        self.assertListEqual([c.name for c in table.columns], ['x'])

    def test_reflect_subquery(self):
        table_node_sql = (
            '(select arrayJoin([1, 2]) as a, arrayJoin([3, 4]) as b)')
        table_node = TextClause(table_node_sql)

        metadata = self.metadata()
        # Cannot use `Table` as it only works with a simple string.
        columns = inspect(metadata.bind).get_columns(table_node)
        self.assertListEqual(
            sorted([col['name'] for col in columns]),
            ['a', 'b'])

    def test_get_schema_names(self):
        insp = inspect(self.session.bind)
        schema_names = insp.get_schema_names()
        self.assertGreater(len(schema_names), 0)

    def test_get_table_names(self):
        table = Table(
            'test_reflect',
            self.metadata(),
            Column('x', types.Int32),
            engines.Log()
        )

        self.session.execute('DROP TABLE IF EXISTS test_reflect')
        table.create(self.session.bind)

        insp = inspect(self.session.bind)
        self.assertListEqual(insp.get_table_names(), ['test_reflect'])


class SchemaHttpTestCase(SchemaTestCase, HttpSessionTestCase):
    """ Schema over a HTTP session """


class SchemaNativeTestCase(SchemaTestCase, NativeSessionTestCase):
    """ Schema over a native protocol session """

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/sql/test_schema.py
datalens-tech/datalens-backend
from sqlalchemy import Column, and_
from sqlalchemy.exc import CompileError
from clickhouse_sqlalchemy import types, select, Table
from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


class SelectTestCase(BaseTestCase):
    def create_table(self, table_name=None, *columns):
        metadata = self.metadata()
        columns = columns or [
            Column('x', types.Int32, primary_key=True)
        ]

        return Table(
            table_name or 't1',
            metadata,
            *columns
        )

    def test_group_by_with_totals(self):
        table = self.create_table()

        query = select.create_legacy_select([table.c.x]).group_by(table.c.x)
        self.assertEqual(
            self.compile(query),
            'SELECT x FROM t1 GROUP BY x'
        )

        query = select.create_legacy_select([table.c.x]).group_by(table.c.x).with_totals()
        self.assertEqual(
            self.compile(query),
            'SELECT x FROM t1 GROUP BY x WITH TOTALS'
        )

    def test_sample(self):
        table = self.create_table()

        query = select.create_legacy_select([table.c.x]).sample(0.1).group_by(table.c.x)
        self.assertEqual(
            self.compile(query),
            'SELECT x FROM t1 SAMPLE %(param_1)s GROUP BY x'
        )

        self.assertEqual(
            self.compile(query, literal_binds=True),
            'SELECT x FROM t1 SAMPLE 0.1 GROUP BY x'
        )

    def test_final(self):
        table = self.create_table()

        query = select.create_legacy_select([table.c.x]).final().group_by(table.c.x)
        self.assertEqual(
            self.compile(query),
            'SELECT x FROM t1 FINAL GROUP BY x'
        )

    def test_nested_type(self):
        table = self.create_table(
            't1',
            Column('x', types.Int32, primary_key=True),
            Column('parent', types.Nested(
                Column('child1', types.Int32),
                Column('child2', types.String),
            ))
        )

        query = select.create_legacy_select(
            [table.c.parent.child1]
        ).where(
            and_(
                table.c.parent.child1 == [1, 2],
                table.c.parent.child2 == ['foo', 'bar']
            )
        )
        self.assertEqual(
            self.compile(query, literal_binds=True),
            'SELECT parent.child1 FROM t1 '
            'WHERE parent.child1 = [1, 2] '
            'AND parent.child2 = [\'foo\', \'bar\']'
        )

    def test_nested_array_join(self):
        table = self.create_table(
            't1',
            Column('x', types.Int32, primary_key=True),
            Column('parent', types.Nested(
                Column('child1', types.Int32),
                Column('child2', types.String),
            ))
        )
        query = select.create_legacy_select(
            [
                table.c.parent.child1,
                table.c.parent.child2,
            ]
        ).array_join(
            table.c.parent
        )
        self.assertEqual(
            self.compile(query),
            'SELECT parent.child1, parent.child2 '
            'FROM t1 '
            'ARRAY JOIN parent'
        )

        query = select.create_legacy_select(
            [
                table.c.parent.child1,
                table.c.parent.child2,
            ]
        ).array_join(
            table.c.parent.child1,
        )
        self.assertEqual(
            self.compile(query),
            'SELECT parent.child1, parent.child2 '
            'FROM t1 '
            'ARRAY JOIN parent.child1'
        )

        query = select.create_legacy_select(
            [
                table.c.parent.child1,
                table.c.parent.child2,
            ]
        ).array_join(
            table.c.parent.child1,
            table.c.parent.child2,
        )
        self.assertEqual(
            self.compile(query),
            'SELECT parent.child1, parent.child2 '
            'FROM t1 '
            'ARRAY JOIN parent.child1, parent.child2'
        )

        parent_labeled = table.c.parent.label('p')

        query = select.create_legacy_select(
            [
                table.c.parent.child1,
                table.c.parent.child2,
            ]
        ).array_join(
            parent_labeled
        )
        self.assertEqual(
            self.compile(query),
            'SELECT parent.child1, parent.child2 '
            'FROM t1 '
            'ARRAY JOIN parent AS p'
        )

        query = select.create_legacy_select(
            [
                parent_labeled.child1,
                parent_labeled.child2,
                table.c.parent.child1,
            ]
        ).array_join(
            parent_labeled
        )
        self.assertEqual(
            self.compile(query),
            'SELECT p.child1, p.child2, parent.child1 AS child1_1 '
            'FROM t1 '
            'ARRAY JOIN parent AS p'
        )

    def test_join(self):
        table_1 = self.create_table(
            'table_1', Column('x', types.UInt32, primary_key=True)
        )
        table_2 = self.create_table(
            'table_2', Column('y', types.UInt32, primary_key=True)
        )

        def make_statement(type=None,
                           strictness=None,
                           distribution=None,
                           full=False,
                           isouter=False):
            join = table_1.join(
                table_2,
                table_2.c.y == table_1.c.x,
                isouter=isouter,
                full=full,
                type=type,
                strictness=strictness,
                distribution=distribution
            )
            return select.create_legacy_select([table_1.c.x]).select_from(join)

        self.assertEqual(
            self.compile(make_statement(
                type='INNER'
            )),
            'SELECT x FROM table_1 '
            'INNER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                type='INNER',
                strictness='all'
            )),
            'SELECT x FROM table_1 '
            'ALL INNER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                type='INNER',
                strictness='any'
            )),
            'SELECT x FROM table_1 '
            'ANY INNER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                type='INNER',
                distribution='global'
            )),
            'SELECT x FROM table_1 '
            'GLOBAL INNER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                type='INNER',
                distribution='global',
                strictness='any'
            )),
            'SELECT x FROM table_1 '
            'GLOBAL ANY INNER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                type='INNER',
                distribution='global',
                strictness='all'
            )),
            'SELECT x FROM table_1 '
            'GLOBAL ALL INNER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                type='LEFT OUTER',
                distribution='global',
                strictness='all')),
            'SELECT x FROM table_1 '
            'GLOBAL ALL LEFT OUTER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                type='RIGHT OUTER',
                distribution='global',
                strictness='all')),
            'SELECT x FROM table_1 '
            'GLOBAL ALL RIGHT OUTER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                type='CROSS',
                distribution='global',
                strictness='all')),
            'SELECT x FROM table_1 '
            'GLOBAL ALL CROSS JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(type='FULL OUTER')),
            'SELECT x FROM table_1 '
            'FULL OUTER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(isouter=False,
                                        full=False)),
            'SELECT x FROM table_1 '
            'INNER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(isouter=True,
                                        full=False)),
            'SELECT x FROM table_1 '
            'LEFT OUTER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(isouter=True,
                                        full=True)),
            'SELECT x FROM table_1 '
            'FULL LEFT OUTER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(isouter=False,
                                        full=True)),
            'SELECT x FROM table_1 '
            'FULL INNER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                isouter=False, full=False,
                type='INNER'
            )),
            'SELECT x FROM table_1 '
            'INNER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                isouter=False, full=False,
                type='LEFT OUTER'
            )),
            'SELECT x FROM table_1 '
            'LEFT OUTER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                isouter=False, full=False,
                type='LEFT', strictness='ALL'
            )),
            'SELECT x FROM table_1'
            ' ALL LEFT JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                isouter=False, full=False,
                type='LEFT', strictness='ANY'
            )),
            'SELECT x FROM table_1 '
            'ANY LEFT JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                isouter=False, full=False,
                type='LEFT', strictness='ANY', distribution='GLOBAL'
            )),
            'SELECT x FROM table_1 '
            'GLOBAL ANY LEFT JOIN table_2 ON y = x'
        )

        self.assertRaises(
            CompileError,
            lambda: self.compile(make_statement(
                isouter=True, full=False,
                type='INNER', strictness='ANY', distribution='GLOBAL'
            )),
        )

        self.assertEqual(
            self.compile(make_statement(
                isouter=True, full=True,
                type='LEFT OUTER', strictness='ANY', distribution='GLOBAL'
            )),
            'SELECT x FROM table_1 GLOBAL '
            'ANY FULL LEFT OUTER JOIN table_2 ON y = x'
        )

        self.assertEqual(
            self.compile(make_statement(
                isouter=True, full=False,
                type='CROSS',
            )),
            'SELECT x FROM table_1 CROSS JOIN table_2 ON y = x'
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/sql/test_selectable.py
datalens-tech/datalens-backend
import enum
from sqlalchemy import sql, Column

from clickhouse_sqlalchemy import types, Table, engines
from clickhouse_sqlalchemy.util import compat
from clickhouse_sqlalchemy_tests.testcase import (
    BaseAbstractTestCase, HttpSessionTestCase, NativeSessionTestCase,
)


class VisitTestCase(BaseAbstractTestCase):
    def test_true_false(self):
        self.assertEqual(self.compile(sql.false()), '0')
        self.assertEqual(self.compile(sql.true()), '1')

    def test_array(self):
        self.assertEqual(
            self.compile(types.Array(types.Int32())),
            'Array(Int32)'
        )
        self.assertEqual(
            self.compile(types.Array(types.Array(types.Int32()))),
            'Array(Array(Int32))'
        )

    def test_enum(self):
        class MyEnum(enum.Enum):
            __order__ = 'foo bar'
            foo = 100
            bar = 500

        self.assertEqual(
            self.compile(types.Enum(MyEnum)),
            "Enum('foo' = 100, 'bar' = 500)"
        )

        self.assertEqual(
            self.compile(types.Enum16(MyEnum)),
            "Enum16('foo' = 100, 'bar' = 500)"
        )

        if compat.PY3:
            data = [" ' t = ", "test"]
        else:
            from collections import OrderedDict
            data = OrderedDict([(" ' t = ", 1), ("test", 2)])

        MyEnum = enum.Enum('MyEnum', data)

        self.assertEqual(
            self.compile(types.Enum8(MyEnum)),
            "Enum8(' \\' t = ' = 1, 'test' = 2)"
        )


class VisitHttpTestCase(VisitTestCase, HttpSessionTestCase):
    """ Visit over a HTTP session """


class VisitNativeTestCase(VisitTestCase, NativeSessionTestCase):
    """ Visit over a native protocol session """

    def test_insert_no_templates_after_value(self):
        """ Optimized non-templating insert test (native protocol only) """
        table = Table(
            't1', self.metadata(),
            Column('x', types.Int32),
            engines.Memory()
        )
        self.assertEqual(
            self.compile(table.insert()),
            'INSERT INTO t1 (x) VALUES'
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/test_compiler.py
datalens-tech/datalens-backend
from sqlalchemy import Column
from sqlalchemy.sql.ddl import CreateTable

from clickhouse_sqlalchemy import types, engines, Table
from clickhouse_sqlalchemy.sql.ddl import DropTable
from clickhouse_sqlalchemy_tests.testcase import BaseTestCase
from clickhouse_sqlalchemy_tests.session import mocked_engine


class DDLTestCase(BaseTestCase):
    def test_create_table(self):
        table = Table(
            't1', self.metadata(),
            Column('x', types.Int32, primary_key=True),
            Column('y', types.String),
            Column('z', types.String(10)),
            # Must be quoted:
            Column('index', types.String),
            engines.Memory()
        )

        # No NOT NULL. And any PKS.
        self.assertEqual(
            self.compile(CreateTable(table)),
            'CREATE TABLE t1 ('
            'x Int32, y String, z FixedString(10), "index" String) '
            'ENGINE = Memory'
        )

    def test_create_table_nested_types(self):
        table = Table(
            't1', self.metadata(),
            Column('x', types.Int32, primary_key=True),
            Column('y', types.Array(types.String)),
            engines.Memory()
        )

        self.assertEqual(
            self.compile(CreateTable(table)),
            'CREATE TABLE t1 '
            '(x Int32, y Array(String)) '
            'ENGINE = Memory'
        )

        table = Table(
            't1', self.metadata(),
            Column('x', types.Int32, primary_key=True),
            Column('y', types.Array(types.Array(types.String))),
            engines.Memory()
        )

        self.assertEqual(
            self.compile(CreateTable(table)),
            'CREATE TABLE t1 '
            '(x Int32, y Array(Array(String))) '
            'ENGINE = Memory'
        )

        table = Table(
            't1', self.metadata(),
            Column('x', types.Int32, primary_key=True),
            Column('y', types.Array(types.Array(types.String))),
            engines.Memory()
        )

        self.assertEqual(
            self.compile(CreateTable(table)),
            'CREATE TABLE t1 '
            '(x Int32, y Array(Array(String))) '
            'ENGINE = Memory'
        )

    def test_create_table_nullable(self):
        table = Table(
            't1', self.metadata(),
            Column('x', types.Int32, primary_key=True),
            Column('y', types.Nullable(types.String)),
            Column('z', types.Nullable(types.String(10))),
            engines.Memory()
        )

        self.assertEqual(
            self.compile(CreateTable(table)),
            'CREATE TABLE t1 '
            '(x Int32, y Nullable(String), z Nullable(FixedString(10))) '
            'ENGINE = Memory'
        )

    def test_create_table_nested(self):
        table = Table(
            't1',
            self.metadata(),
            Column('x', types.Int32, primary_key=True),
            Column('parent', types.Nested(
                Column('child1', types.Int32),
                Column('child2', types.String),
            )),
            engines.Memory()
        )
        self.assertEqual(
            self.compile(CreateTable(table)),
            'CREATE TABLE t1 ('
            'x Int32, '
            'parent Nested('
            'child1 Int32, '
            "child2 String"
            ')'
            ') ENGINE = Memory'
        )

    def test_create_table_nested_nullable(self):
        table = Table(
            't1', self.metadata(),
            Column('x', types.Int32, primary_key=True),
            Column('y', types.Array(types.Nullable(types.String))),
            engines.Memory()
        )

        self.assertEqual(
            self.compile(CreateTable(table)),
            'CREATE TABLE t1 '
            '(x Int32, y Array(Nullable(String))) '
            'ENGINE = Memory'
        )

    def test_create_table_nullable_nested_nullable(self):
        table = Table(
            't1', self.metadata(),
            Column('x', types.Int32, primary_key=True),
            Column('y', types.Nullable(
                types.Array(types.Nullable(types.String)))
            ),
            engines.Memory()
        )

        self.assertEqual(
            self.compile(CreateTable(table)),
            'CREATE TABLE t1 '
            '(x Int32, y Nullable(Array(Nullable(String)))) '
            'ENGINE = Memory'
        )

    def test_table_create_on_cluster(self):
        create_sql = (
            'CREATE TABLE t1 ON CLUSTER test_cluster '
            '(x Int32) ENGINE = Memory'
        )

        with mocked_engine() as engine:
            table = Table(
                't1', self.metadata(session=engine.session),
                Column('x', types.Int32, primary_key=True),
                engines.Memory(),
                clickhouse_cluster='test_cluster'
            )

            table.create()
            self.assertEqual(engine.history, [create_sql])

        self.assertEqual(
            self.compile(CreateTable(table)),
            create_sql
        )

    def test_drop_table_clause(self):
        table = Table(
            't1', self.metadata(),
            Column('x', types.Int32, primary_key=True)
        )

        self.assertEqual(
            self.compile(DropTable(table)),
            'DROP TABLE t1'
        )
        self.assertEqual(
            self.compile(DropTable(table, if_exists=True)),
            'DROP TABLE IF EXISTS t1'
        )

    def test_table_drop(self):
        with mocked_engine() as engine:
            table = Table(
                't1', self.metadata(session=engine.session),
                Column('x', types.Int32, primary_key=True)
            )
            table.drop(if_exists=True)
            self.assertEqual(engine.history, ['DROP TABLE IF EXISTS t1'])

    def test_table_drop_on_cluster(self):
        drop_sql = 'DROP TABLE IF EXISTS t1 ON CLUSTER test_cluster'

        with mocked_engine() as engine:
            table = Table(
                't1', self.metadata(session=engine.session),
                Column('x', types.Int32, primary_key=True),
                clickhouse_cluster='test_cluster'
            )
            table.drop(if_exists=True)
            self.assertEqual(engine.history, [drop_sql])

        self.assertEqual(
            self.compile(DropTable(table, if_exists=True)),
            drop_sql
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/test_ddl.py
datalens-tech/datalens-backend
from sqlalchemy import Column, func, exc
from sqlalchemy.sql.ddl import CreateTable

from clickhouse_sqlalchemy import types, engines, get_declarative_base, Table
from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


class EnginesDeclarativeTestCase(BaseTestCase):
    @property
    def base(self):
        return get_declarative_base()

    def test_text_engine_columns_declarative(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)

            __table_args__ = (
                engines.MergeTree('date', ('date', 'x')),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            'CREATE TABLE test_table (date Date, x Int32, y String) '
            'ENGINE = MergeTree() '
            'PARTITION BY date '
            'ORDER BY (date, x)'
        )

    def test_text_engine_columns(self):
        table = Table(
            't1', self.metadata(),
            Column('date', types.Date, primary_key=True),
            Column('x', types.Int32),
            Column('y', types.String),
            engines.MergeTree('date', ('date', 'x')),
        )

        self.assertEqual(
            self.compile(CreateTable(table)),
            'CREATE TABLE t1 (date Date, x Int32, y String) '
            'ENGINE = MergeTree() '
            'PARTITION BY date '
            'ORDER BY (date, x)'
        )

    def test_func_engine_columns(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)

            __table_args__ = (
                engines.MergeTree('date', ('date', func.intHash32(x)),
                                  sample=func.intHash32(x)),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            'CREATE TABLE test_table (date Date, x Int32, y String) '
            'ENGINE = MergeTree() PARTITION BY date '
            'ORDER BY (date, intHash32(x)) '
            'SAMPLE BY intHash32(x)'
        )

    def test_merge_tree_all_settings(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)

            __table_args__ = (
                engines.MergeTree(
                    partition_by=date,
                    order_by=(date, x),
                    primary_key=(x, y),
                    sample=func.hashFunc(x),
                    setting1=2,
                    setting2=5
                ),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            'CREATE TABLE test_table '
            '(date Date, x Int32, y String) '
            'ENGINE = MergeTree() '
            'PARTITION BY date '
            'ORDER BY (date, x) '
            'PRIMARY KEY (x, y) '
            'SAMPLE BY hashFunc(x) '
            'SETTINGS setting1=2, setting2=5'
        )

    def test_create_table_without_engine(self):
        no_engine_table = Table(
            't1', self.metadata(),
            Column('x', types.Int32, primary_key=True),
            Column('y', types.String)
        )

        with self.assertRaises(exc.CompileError) as ex:
            self.compile(CreateTable(no_engine_table))

        self.assertEqual(str(ex.exception), "No engine for table 't1'")

    def test_collapsing_merge_tree(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)
            sign = Column(types.Int8)

            __table_args__ = (
                engines.CollapsingMergeTree(
                    sign,
                    date,
                    (date, x),
                    (x, y),
                    func.random(),
                    key='value'),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            'CREATE TABLE test_table '
            '(date Date, x Int32, y String, sign Int8) '
            'ENGINE = CollapsingMergeTree(sign) '
            'PARTITION BY date '
            'ORDER BY (date, x) '
            'PRIMARY KEY (x, y) '
            'SAMPLE BY random() '
            'SETTINGS key=value'
        )

    def test_summing_merge_tree(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.Int32)

            __table_args__ = (
                engines.SummingMergeTree(date, (date, x), columns=(y, )),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            'CREATE TABLE test_table (date Date, x Int32, y Int32) '
            'ENGINE = SummingMergeTree(y) '
            'PARTITION BY date '
            'ORDER BY (date, x)'
        )

    def test_buffer(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)

            __table_args__ = (
                engines.Buffer('db', 'table'),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            "CREATE TABLE test_table (date Date, x Int32, y String) "
            "ENGINE = Buffer("
            "db, table, 16, 10, 100, 10000, 1000000, 10000000, 100000000"
            ")"
        )

    def test_distributed(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)

            __table_args__ = (
                engines.Distributed(
                    'cluster', 'test', 'merge_distributed1', 'rand()'
                ),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            "CREATE TABLE test_table (date Date, x Int32) "
            "ENGINE = Distributed(cluster, test, merge_distributed1, rand())"
        )

    def test_graphite_merge_tree(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)

            __table_args__ = (
                engines.GraphiteMergeTree('config_section', date, (date, x)),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            "CREATE TABLE test_table (date Date, x Int32) "
            "ENGINE = GraphiteMergeTree('config_section') "
            "PARTITION BY date "
            "ORDER BY (date, x)"
        )

    def test_memory(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)

            __table_args__ = (
                engines.Memory(),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            "CREATE TABLE test_table (date Date, x Int32) "
            "ENGINE = Memory"
        )

    def test_replicated_merge_tree(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)

            __table_args__ = (
                engines.ReplicatedMergeTree(
                    '/table/path', 'name', 'date', ('date', 'x')
                ),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            "CREATE TABLE test_table (date Date, x Int32, y String) "
            "ENGINE = ReplicatedMergeTree('/table/path', 'name') "
            "PARTITION BY date "
            "ORDER BY (date, x)"
        )

    def test_replacing_merge_tree(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)
            version = Column(types.Int32)

            __table_args__ = (
                engines.ReplacingMergeTree(
                    'date',
                    ('date', 'x'),
                    ver='version',
                ),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            "CREATE TABLE test_table ("
            "date Date, x Int32, y String, version Int32"
            ") "
            "ENGINE = ReplacingMergeTree(version) "
            "PARTITION BY date "
            "ORDER BY (date, x)"
        )

    def test_replacing_merge_tree_no_version(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)
            version = Column(types.Int32)

            __table_args__ = (
                engines.ReplacingMergeTree(
                    'date',
                    ('date', 'x'),
                ),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            "CREATE TABLE test_table ("
            "date Date, x Int32, y String, version Int32"
            ") "
            "ENGINE = ReplacingMergeTree() "
            "PARTITION BY date "
            "ORDER BY (date, x)"
        )

    def test_replicated_replacing_merge_tree(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)
            version = Column(types.Int32)

            __table_args__ = (
                engines.ReplicatedReplacingMergeTree(
                    '/table/path', 'name',
                    'date', ('date', 'x'),
                    ver='version',
                ),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            "CREATE TABLE test_table ("
            "date Date, x Int32, y String, version Int32"
            ") "
            "ENGINE = ReplicatedReplacingMergeTree("
            "'/table/path', 'name', version"
            ") "
            "PARTITION BY date "
            "ORDER BY (date, x)"
        )

    def test_partition_by_func(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)
            version = Column(types.Int32)

            __table_args__ = (
                engines.ReplacingMergeTree(
                    ver=version,
                    partition_by=func.toYYYYMM(date),
                    order_by=(date, x),
                ),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            "CREATE TABLE test_table ("
            "date Date, x Int32, y String, version Int32"
            ") "
            "ENGINE = ReplacingMergeTree(version) "
            "PARTITION BY toYYYYMM(date) "
            "ORDER BY (date, x)"
        )

    def test_file(self):
        class TestTable(self.base):
            date = Column(types.Date, primary_key=True)
            x = Column(types.Int32)
            y = Column(types.String)

            __table_args__ = (
                engines.File(
                    'JSONEachRow'
                ),
            )

        self.assertEqual(
            self.compile(CreateTable(TestTable.__table__)),
            "CREATE TABLE test_table (date Date, x Int32, y String)"
            " ENGINE = File(JSONEachRow)"
        )

    def test_file_raises(self):
        with self.assertRaises(Exception):
            class TestTable(self.base):
                date = Column(types.Date, primary_key=True)
                x = Column(types.Int32)
                y = Column(types.String)

                __table_args__ = (
                    engines.File(
                        'unsupported_format'
                    ),
                )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/test_engines.py
datalens-tech/datalens-backend

from clickhouse_sqlalchemy.exceptions import DatabaseException
from clickhouse_sqlalchemy_tests.testcase import BaseTestCase


CASES = (
    (r"Code: 60, e.displayText() = DB::Exception: Table default.abc "
     r"doesn't exist. (version 19.16.2.2 (official build))"),
    (r"Code: 53. Type mismatch in VALUES section. Repeat query with "
     r"types_check=True for detailed info. Column x: argument out of range"),
    (r"std::exception. Code: 1001, type: ...::TErrorException, "
     r"e.what() = Error resolving path ..."),
)


class ExceptionParsingTestCase(BaseTestCase):

    def test_exception_parsing(self):
        for case in CASES:
            exc = DatabaseException(Exception(case))
            self.assertIsNotNone(exc.code)

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/test_exceptions.py
datalens-tech/datalens-backend
# coding: utf8
"""
...
"""

from __future__ import unicode_literals

import json

from clickhouse_sqlalchemy.parsers.jsoncompact import JSONCompactChunksParser

try:
    from io import BytesIO
except ImportError:
    from cStringIO import StringIO as BytesIO

from clickhouse_sqlalchemy_tests.testcase import BaseTestCase

from clickhouse_sqlalchemy.parsers import (
    FORMATS,
    Parts,
)


class TSVParserTestCase(BaseTestCase):

    def test_the_edge_case(self):
        """
        SQL

            select '' as value group by value with totals
            format TabSeparatedWithNamesAndTypes

        """
        chunks = [b'value\nString\n\n\n\n']
        fmt_info = FORMATS['TabSeparatedWithNamesAndTypes with totals']()
        result = fmt_info.generator(chunks)
        result = list(result)
        expected = [
            (Parts.META, [{'name': 'value', 'type': 'String'}]),
            (Parts.DATA, ['']),
            (Parts.TOTALS, ['']),
        ]
        self.assertEqual(result, expected)

    def test_the_edgiest_case(self):
        """
        SQL:

            select arrayJoin(['', 'a', 'b']) as value group by value with
            totals order by value desc format TabSeparatedWithNamesAndTypes

        """
        chunks = [b'value\nString\nb\na\n\n\n\n']
        fmt_info = FORMATS['TabSeparatedWithNamesAndTypes with totals']()
        result = fmt_info.generator(chunks)
        result = list(result)
        expected = [
            (Parts.META, [{'name': 'value', 'type': 'String'}]),
            (Parts.DATA, ['b']),
            (Parts.DATA, ['a']),
            (Parts.DATA, ['']),
            (Parts.TOTALS, ['']),
        ]
        self.assertEqual(result, expected)

    def test_the_normal_case(self):
        """
        SQL:

            select arrayJoin(['', '', '', '']) as value
            format TabSeparatedWithNamesAndTypes

        """
        chunks = [b'value\nString\n\n\n\n\n']
        fmt_info = FORMATS['TabSeparatedWithNamesAndTypes']()
        result = fmt_info.generator(chunks)
        result = list(result)
        expected = [
            (Parts.META, [{'name': 'value', 'type': 'String'}]),
            (Parts.DATA, ['']),
            (Parts.DATA, ['']),
            (Parts.DATA, ['']),
            (Parts.DATA, ['']),
        ]
        self.assertEqual(result, expected)


JSONCOMPACT_SAMPLE_DATA_SIMPLE = '''
{
\t"meta":
\t[
\t\t{
\t\t\t"name": "fielddate",
\t\t\t"type": "Nullable(String)"
\t\t},
\t\t{
\t\t\t"name": "lock_wait_times",
\t\t\t"type": "Nullable(String)"
\t\t}
\t],

\t"data":
\t[
\t\t["2019-04-01", "1"],
\t\t["2019-04-01", "2"],
\t\t["2019-04-01", "3"]
\t],

\t"totals": [null,null],

\t"extremes":
\t{
\t\t"min": ["2019-04-01","1"],
\t\t"max": ["2019-04-01","3"]
\t},

\t"rows": 3,

\t"rows_before_limit_at_least": 3,

\t"statistics":
\t{
\t\t"elapsed": 0.481989638,
\t\t"rows_read": 0,
\t\t"bytes_read": 0
\t}
}
'''.encode('utf-8')

JSONCOMPACT_SAMPLE_DATA_EMPTY = '''{
\t"meta":
\t[
\t\t{
\t\t\t"name": "test_x",
\t\t\t"type": "DateTime"
\t\t}
\t],

\t"data":
\t[

\t],

\t"rows": 0,

\t"statistics":
\t{
\t\t"elapsed": 0.000019615,
\t\t"rows_read": 0,
\t\t"bytes_read": 0
\t}
}
'''.encode('utf-8')

JSONCOMPACT_SAMPLE_WITH_TIMEOUT = '''
{
\t"meta":
\t[
\t\t{
\t\t\t"name": "res_0",
\t\t\t"type": "Nullable(String)"
\t\t}
\t],

\t"data":
\t[
\t\t["EKB"],
\t\t["RST"],
\t\t["SAM"],
\t\t["SOF"],
\t\t["SPB"]
\t],

\t"rows": 5,

\t"rows_before_limit_at_least": 105,

\t"statistics":
\t{
\t\t"elapsed": 1.193702828,
\t\t"rows_read": 96642607,
\t\t"bytes_read": 6624351318
\t}
}
Code: 159. DB::Exception: Timeout exceeded: elapsed 32.001073883 seconds, maximum: 32. (TIMEOUT_EXCEEDED) (version 22.8.16.1)
'''.encode('utf-8')

EMPTY_DATA_CHUNKS = [
    b'{\n\t"meta":\n\t[\n\t\t{\n\t\t\t"name": "res_o",\n\t\t\t"type": "String"\n\t\t}\n\t],\n\n\t"data":\n\t[\n\n\t]',
    b',\n\n\t"rows": 0,\n\n\t"statistics":\n\t{\n\t\t"elapsed": 0.000345268,\n\t\t"rows_read": 0,\n\t\t"bytes_read": 0\n\t}\n}\n',
]

SPLIT_AFTER_DATA_CHUNKS = [
    b'{\n\t"meta":\n\t[\n\t\t{\n\t\t\t"name": "res_0",\n\t\t\t"type": "Nullable(Int64)"\n\t\t}\n\t],\n\n\t"data":\n\t[\n\t\t["1"],\n\t\t["2"],\n\t\t["3"],\n\t\t["4"],\n\t\t'
    b'["5"],\n\t\t["6"],\n\t\t["7"],\n\t\t["8"],\n\t\t["9"],\n\t\t["1' + 1200 * b'0' + b'"]\n\t]',
    b',\n\n\t"rows": 10,\n\n\t"rows_before_limit_at_least": 10,\n\n\t"statistics":\n\t{\n\t\t"elapsed": 0.00487046,\n\t\t"rows_read": 10,\n\t\t"bytes_read": 13375\n\t}\n}\n',
]

DATA_WITH_TIMEOUT_CHUNKS = [
    b'{\n\t"meta":\n\t[\n\t\t{\n\t\t\t"name": "res_0",\n\t\t\t"type": "Nullable(String)"\n\t\t}\n\t],\n\n\t"data":\n\t[\n\t\t["EKB"],\n\t\t["RST"],\n\t\t["SAM"],\n\t\t["SOF"],',
    b'\n\t\t["SPB"]\n\t],\n\n\t"rows": 5,\n\n\t"rows_before_limit_at_least": 105,\n\n\t"statistics":\n\t{\n\t\t"elapsed": 1.193702828,\n\t\t"rows_read": 96642607,\n\t\t',
    b'"bytes_read": 6624351318\n\t}\n}\nCode: 159. DB::Exception: Timeout exceeded: elapsed 32.001073883 seconds, maximum: 32. (TIMEOUT_EXCEEDED) (version 22.8.16.1)\n'
]


def build_mocked_jsoncompact_result(events):
    """
    Given an iterable of events, build an entire JSONCompact data result.
    """
    mocked_result = {'data': []}
    for part, item in events:
        if isinstance(item, (bytes, str)):
            item = json.loads(item)
        if part == Parts.META:
            mocked_result['meta'] = item
        elif part == Parts.DATA:
            mocked_result['data'].append(item)
        elif part == Parts.DATACHUNK:
            mocked_result['data'].extend(item)
        elif part == Parts.TOTALS:
            mocked_result['totals'] = item
        elif part == Parts.STATS:
            mocked_result.update(item)
        else:
            raise Exception("Unknown part", part)
    return mocked_result


class JSONCompactParserTestCase(BaseTestCase):

    def _process(self, data_source_bytes, chunk_size, check_json=True):

        data_source = BytesIO(data_source_bytes)

        results = []
        chunk_inputs = []
        chunk_results = []

        def gather_results():
            results.append((
                list(chunk_inputs),
                list(chunk_results)))
            chunk_inputs[:] = []
            chunk_results[:] = []

        fmt_info = FORMATS['JSONCompact']()
        parser = fmt_info.sans_io()
        while True:
            event, data = parser.next_event()
            if event == Parts.NEED_DATA:
                gather_results()
                assert data is None

                chunk = data_source.read(chunk_size)
                if not chunk:
                    break
                parser.receive_data(chunk)
                chunk_inputs.append(chunk)

            elif event == Parts.FINISHED:
                gather_results()
                assert data is None

                chunk = data_source.read()
                if chunk:
                    chunk_inputs.append(chunk)
                    gather_results()

                break

            else:
                chunk_results.append((event, data))

        gather_results()

        data_results = [
            (part, json.loads(item) if isinstance(item, bytes) else item)
            for chunk_inputs, chunk_results in results
            for part, item in chunk_results]

        result = build_mocked_jsoncompact_result(data_results)
        if check_json:
            expected = json.loads(data_source_bytes)
            self.assertEqual(result, expected)
        else:
            expected = None
        return data_results, result, expected

    def test_jsoncompact_chunks_parser(self):
        self._process(
            data_source_bytes=JSONCOMPACT_SAMPLE_DATA_SIMPLE,
            chunk_size=123)

    def test_jsoncompact_empty_data(self):
        self._process(
            data_source_bytes=JSONCOMPACT_SAMPLE_DATA_EMPTY,
            chunk_size=128000)

    def test_jsoncompact_with_timeout(self):
        self._process(
            data_source_bytes=JSONCOMPACT_SAMPLE_WITH_TIMEOUT,
            chunk_size=128000,
            check_json=False)

    def _parse(self, chunks):
        parser = JSONCompactChunksParser()
        event = None
        chunk = None
        data = None
        for chunk in chunks:
            parser.receive_data(chunk)
            event, data = parser.next_event()
            while (event != parser.parts.NEED_DATA and
                   event != parser.parts.FINISHED):
                yield event, data
                event, data = parser.next_event()
        if event is not None and event != parser.parts.FINISHED:
            raise ValueError(
                "Unexpected last event",
                dict(event=event, data=data, chunk=chunk),
            )

    def test_empty_data_chunks(self):
        resp = list(self._parse(EMPTY_DATA_CHUNKS))
        data = resp[1]
        assert data[0] == 'PART_DATACHUNK'
        assert data[1] == []

    def test_split_after_data_chunks(self):
        resp = list(self._parse(SPLIT_AFTER_DATA_CHUNKS))

        data_parts = [part[1] for part in resp if part[0] == 'PART_DATACHUNK']
        rows = sum(len(part) for part in data_parts)
        assert rows == 10

        parts = [part[0] for part in resp]
        assert parts == ['PART_META', 'PART_DATACHUNK', 'PART_DATACHUNK', 'PART_STATS']

    def test_data_with_timeout(self):
        resp = list(self._parse(DATA_WITH_TIMEOUT_CHUNKS))
        assert resp[1] == ('PART_DATACHUNK', [['EKB'], ['RST'], ['SAM']])
        assert resp[2] == ('PART_DATACHUNK', [['SOF'], ['SPB']])
        assert resp[3] == ('PART_STATS', {'rows': 5, 'rows_before_limit_at_least': 105,
                                          'statistics': {'bytes_read': 6624351318, 'elapsed': 1.193702828,
                                                         'rows_read': 96642607}})

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/test_parsers.py
datalens-tech/datalens-backend
import enum

from sqlalchemy import Column, inspect

from clickhouse_sqlalchemy import types, engines, Table
from clickhouse_sqlalchemy_tests.testcase import (
    TypesTestCase, HttpSessionTestCase, NativeSessionTestCase,
)
from clickhouse_sqlalchemy_tests.util import require_server_version


class ReflectionTestCase(TypesTestCase):
    def _type_round_trip(self, *types):
        metadata = self.metadata()

        args = (
            [Column('t%d' % i, type_) for i, type_ in enumerate(types)] +
            [engines.Memory()]
        )

        table = Table('t', metadata, *args)
        with self.create_table(table):
            return [
                c['type'] for c in
                inspect(metadata.bind).get_columns('t')
            ]

    def test_array(self):
        coltype = self._type_round_trip(types.Array(types.Int32))[0]

        self.assertIsInstance(coltype, types.Array)
        self.assertEqual(coltype.item_type, types.Int32)

    def test_array_of_array(self):
        coltype = self._type_round_trip(
            types.Array(types.Array(types.Int32))
        )[0]

        self.assertIsInstance(coltype, types.Array)
        self.assertIsInstance(coltype.item_type, types.Array)
        self.assertEqual(coltype.item_type.item_type, types.Int32)

    def test_sting_length(self):
        coltype = self._type_round_trip(types.String(10))[0]

        self.assertIsInstance(coltype, types.String)
        self.assertEqual(coltype.length, 10)

    def test_nullable(self):
        coltype = self._type_round_trip(types.Nullable(types.Int32))[0]

        # self.assertIsInstance(coltype, types.Nullable)
        # self.assertEqual(coltype.nested_type, types.Int32)
        self.assertIsInstance(coltype, types.Int32)

    @require_server_version(19, 3, 3)
    def test_low_cardinality(self):
        coltype = self._type_round_trip(types.LowCardinality(types.String))[0]

        # self.assertIsInstance(coltype, types.LowCardinality)
        # self.assertEqual(coltype.nested_type, types.String)
        self.assertIsInstance(coltype, types.String)

    def test_enum8(self):
        enum_options = {'three': 3, "quoted' ": 9, 'comma, ': 14}
        coltype = self._type_round_trip(
            types.Enum8(enum.Enum('any8_enum', enum_options))
        )[0]

        self.assertIsInstance(coltype, types.Enum8)

        self.assertEqual(sorted(coltype.enums), sorted(enum_options))
        if coltype.enum_class is not None:
            self.assertEqual(
                {o.name: o.value for o in coltype.enum_class}, enum_options,
            )

    def test_enum16(self):
        enum_options = {'first': 1024, 'second': 2048}
        coltype = self._type_round_trip(
            types.Enum16(enum.Enum('any16_enum', enum_options))
        )[0]

        self.assertIsInstance(coltype, types.Enum16)

        self.assertEqual(sorted(coltype.enums), sorted(enum_options))
        if coltype.enum_class is not None:
            self.assertEqual(
                {o.name: o.value for o in coltype.enum_class}, enum_options,
            )

    def test_decimal(self):
        coltype = self._type_round_trip(types.Decimal(38, 38))[0]

        self.assertIsInstance(coltype, types.Decimal)
        self.assertEqual(coltype.precision, 38)
        self.assertEqual(coltype.scale, 38)

    def test_datetimewithtz(self):
        coltype = self._type_round_trip(types.DateTimeWithTZ('Etc/UTC'))[0]
        self.assertIsInstance(coltype, types.DateTimeWithTZ)
        self.assertEqual(coltype.tz, 'Etc/UTC')


class ReflectionHttpTestCase(ReflectionTestCase, HttpSessionTestCase):
    """ Reflection over a HTTP session """


class ReflectionNativeTestCase(ReflectionTestCase, NativeSessionTestCase):
    """ Reflection over a native protocol session """

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/test_reflection.py
datalens-tech/datalens-backend
from datetime import datetime

from clickhouse_sqlalchemy.dt_utils import parse_datetime64

from clickhouse_sqlalchemy_tests.testcase import BaseAbstractTestCase


class Datetime64ParserTestCase(BaseAbstractTestCase):
    def test_7_digits_after_dot(self):
        self.assertEqual(
            parse_datetime64('2020-01-01 12:34:56.1234567'),
            datetime(year=2020, month=1, day=1, hour=12, minute=34, second=56, microsecond=123456)
        )

    def test_6_digits_after_dot(self):
        self.assertEqual(
            parse_datetime64('2020-01-01 12:34:56.123456'),
            datetime(year=2020, month=1, day=1, hour=12, minute=34, second=56, microsecond=123456)
        )

    def test_3_digits_after_dot(self):
        self.assertEqual(
            parse_datetime64('2020-01-01 12:34:56.123'),
            datetime(year=2020, month=1, day=1, hour=12, minute=34, second=56, microsecond=123000)
        )

    def test_no_digits_after_dot(self):
        self.assertEqual(
            parse_datetime64('2020-01-01 12:34:56'),
            datetime(year=2020, month=1, day=1, hour=12, minute=34, second=56, microsecond=0)
        )

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/test_stuff.py
datalens-tech/datalens-backend
# coding: utf8
"""
...
"""

from __future__ import unicode_literals

import sqlalchemy as sa

from clickhouse_sqlalchemy_tests.testcase import BaseTestCase
from clickhouse_sqlalchemy_tests.session import http_jc_session


class HttpTotalsTestCase(BaseTestCase):

    session = http_jc_session

    def test_select_totals_simple(self):
        query = "select '' as value group by value with totals"
        res = self.session.execute(query)
        cursor = res.cursor
        results = list(res)
        self.assertEqual(results, [('',)])
        self.assertTrue(isinstance(results[0], sa.engine.row.RowProxy))
        self.assertEqual(
            cursor.ch_schema,
            [{'name': 'value', 'type': 'String'}])
        self.assertTrue(cursor.ch_stats)
        try:
            elapsed = cursor.ch_stats['statistics']['elapsed']
        except Exception:
            elapsed = None
        expected = {
            'rows': 1,
            'totals': [''],
            'statistics': {
                'elapsed': elapsed,
                'rows_read': 1,
                'bytes_read': 1}}
        self.assertEqual(cursor.ch_stats, expected)
        # Note: one row, postprocessed, but not currently wrapped in
        # `sqlalchemy.engine.result.RowProxy`»
        self.assertEqual(cursor.ch_totals, [''])

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/test_totals.py
datalens-tech/datalens-backend
from __future__ import unicode_literals

from datetime import datetime
from decimal import Decimal
from ipaddress import IPv4Address, IPv4Network, IPv6Address, IPv6Network
from unittest import skip

from sqlalchemy import Column, Numeric, and_, literal, cast
from sqlalchemy.sql.ddl import CreateTable

from clickhouse_sqlalchemy import types, engines, Table
from clickhouse_sqlalchemy.exceptions import DatabaseException

from clickhouse_sqlalchemy_tests.testcase import (
    TypesTestCase, HttpSessionTestCase, NativeSessionTestCase,
)
from clickhouse_sqlalchemy_tests.util import require_server_version


class DateTimeTypeTestCase(TypesTestCase):
    table = Table(
        'test', TypesTestCase.metadata(),
        Column('x', types.DateTime, primary_key=True),
        engines.Memory()
    )

    def test_select_insert(self):
        dt = datetime(2018, 1, 1, 15, 20)

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': dt}])
            self.assertEqual(self.session.query(self.table.c.x).scalar(), dt)

    def test_create_table(self):
        self.assertEqual(
            self.compile(CreateTable(self.table)),
            'CREATE TABLE test (x DateTime) ENGINE = Memory'
        )


class DateTimeTypeHttpTestCase(DateTimeTypeTestCase, HttpSessionTestCase):
    """ DateTimeType over a HTTP session """


class DateTimeTypeNativeTestCase(DateTimeTypeTestCase, NativeSessionTestCase):
    """ DateTimeType over a native protocol session """


class TypeCastTestCase(TypesTestCase):

    def test_cast_to_nullable(self):
        query = cast(literal(127), types.Nullable(types.Int64))
        self.assertEqual(self.session.query(query).scalar(), 127)
        query = cast(literal(None), types.Nullable(types.Int64))
        self.assertIsNone(self.session.query(query).scalar())


class TypeCastHttpTestCase(TypeCastTestCase, HttpSessionTestCase):
    """ TypeCast over a HTTP session """


class TypeCastNativeTestCase(TypeCastTestCase, NativeSessionTestCase):
    """ TypeCast over a native protocol session """


class NumericTypeTestCase(TypesTestCase):
    table = Table(
        'test', TypesTestCase.metadata(),
        Column('x', Numeric(10, 2)),
        engines.Memory()
    )

    def test_create_table(self):
        self.assertEqual(
            self.compile(CreateTable(self.table)),
            'CREATE TABLE test (x Decimal(10, 2)) ENGINE = Memory'
        )

    def test_select_insert(self):
        x = Decimal('12345678.12')

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': x}])
            self.assertEqual(self.session.query(self.table.c.x).scalar(), x)

    def test_insert_truncate(self):
        value = Decimal('123.129999')
        expected = Decimal('123.12')

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': value}])
            self.assertEqual(self.session.query(self.table.c.x).scalar(),
                             expected)

    def test_insert_overflow(self):
        value = Decimal('12345678901234567890.1234567890')

        with self.create_table(self.table):
            with self.assertRaises(DatabaseException) as ex:
                self.session.execute(self.table.insert(), [{'x': value}])

            # 'Too many digits' is written in the CH response;
            # 'out of range' is raised from `struct` within
            # `clickhouse_driver`,
            # before the query is sent to CH.
            message = str(ex.exception)
            self.assertTrue(
                'out of range' in message or
                'Too many digits' in message or
                'Decimal value is too big' in message)

    def test_create_table_decimal_symlink(self):
        table = Table(
            'test', TypesTestCase.metadata(),
            Column('x', types.Decimal(10, 2)),
            engines.Memory()
        )

        self.assertEqual(
            self.compile(CreateTable(table)),
            'CREATE TABLE test (x Decimal(10, 2)) ENGINE = Memory'
        )


class NumericTypeHttpTestCase(NumericTypeTestCase, HttpSessionTestCase):
    """ NumericType over a HTTP session """

    @skip("doesn't happen in CH 20.5 anymore")
    def test_insert_truncate(self):
        value = Decimal('123.129999')
        with self.create_table(self.table):
            with self.assertRaises(DatabaseException) as ex:
                self.session.execute(self.table.insert(), [{'x': value}])
            self.assertIn('value is too small', str(ex.exception))


class NumericTypeNativeTestCase(NumericTypeTestCase, NativeSessionTestCase):
    """ NumericType over a native protocol session """


class IPv4TestCase(TypesTestCase):
    table = Table(
        'test', TypesTestCase.metadata(),
        Column('x', types.IPv4),
        engines.Memory()
    )

    @require_server_version(19, 3, 3)
    def test_create_table(self):
        self.assertEqual(
            self.compile(CreateTable(self.table)),
            'CREATE TABLE test (x IPv4) ENGINE = Memory'
        )

    @require_server_version(19, 3, 3)
    def test_select_insert(self):
        a = IPv4Address('10.0.0.1')

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': a}])
            self.assertEqual(self.session.query(self.table.c.x).scalar(), a)

    @require_server_version(19, 3, 3)
    def test_select_insert_string(self):
        a = '10.0.0.1'

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': a}])
            self.assertEqual(self.session.query(self.table.c.x).scalar(),
                             IPv4Address('10.0.0.1'))

    @require_server_version(19, 3, 3)
    def test_select_where_address(self):
        a = IPv4Address('10.0.0.1')

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': a}])

            self.assertEqual(self.session.query(self.table.c.x).filter(
                self.table.c.x == IPv4Address('10.0.0.1')).scalar(), a)
            self.assertEqual(self.session.query(self.table.c.x).filter(
                and_(IPv4Address('10.0.0.0') < self.table.c.x,
                     self.table.c.x < IPv4Address('10.0.0.2'))).scalar(), a)

    @require_server_version(19, 3, 3)
    def test_select_where_string(self):
        a = IPv4Address('10.0.0.1')

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': a}])

            self.assertEqual(self.session.query(self.table.c.x).filter(
                and_('10.0.0.0' < self.table.c.x,
                     self.table.c.x < '10.0.0.2')).scalar(), a)

    @require_server_version(19, 3, 3)
    def test_select_in_network(self):
        ips = [
            IPv4Address('10.0.0.1'),
            IPv4Address('10.0.0.2'),
            IPv4Address('10.0.0.3'),
            IPv4Address('192.168.0.1')
        ]

        with self.create_table(self.table):
            self.session.execute(self.table.insert(),
                                 [{'x': ip} for ip in ips])

            col = self.table.c.x
            values = (
                self.session.query(self.table.c.x)
                .filter(col.in_(IPv4Network('10.0.0.0/8')))
                .order_by(col)
                .all())
            self.assertEqual(
                values,
                [
                    (IPv4Address('10.0.0.1'),),
                    (IPv4Address('10.0.0.2'),),
                    (IPv4Address('10.0.0.3'),)
                ])

    @require_server_version(19, 3, 3)
    def test_select_in_string(self):
        ips = [
            IPv4Address('10.0.0.1'),
            IPv4Address('10.0.0.2'),
            IPv4Address('10.0.0.3'),
            IPv4Address('192.168.0.1')
        ]

        with self.create_table(self.table):
            self.session.execute(self.table.insert(),
                                 [{'x': ip} for ip in ips])

            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    self.table.c.x.in_('10.0.0.0/8')).all(), [
                    (IPv4Address('10.0.0.1'),),
                    (IPv4Address('10.0.0.2'),),
                    (IPv4Address('10.0.0.3'),)
                ])

    @require_server_version(19, 3, 3)
    def test_select_not_in_network(self):
        ips = [
            IPv4Address('10.0.0.1'),
            IPv4Address('10.0.0.2'),
            IPv4Address('10.0.0.3'),
            IPv4Address('192.168.0.1')
        ]

        with self.create_table(self.table):
            self.session.execute(self.table.insert(),
                                 [{'x': ip} for ip in ips])

            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    self.table.c.x.notin_(IPv4Network('10.0.0.0/8'))).all(), [
                    (IPv4Address('192.168.0.1'),)
                ])
            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    ~self.table.c.x.in_(IPv4Network('10.0.0.0/8'))).all(), [
                    (IPv4Address('192.168.0.1'),)
                ])

    @require_server_version(19, 3, 3)
    def test_select_not_in_string(self):
        ips = [
            IPv4Address('10.0.0.1'),
            IPv4Address('10.0.0.2'),
            IPv4Address('10.0.0.3'),
            IPv4Address('192.168.0.1')
        ]

        with self.create_table(self.table):
            self.session.execute(self.table.insert(),
                                 [{'x': ip} for ip in ips])

            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    self.table.c.x.notin_('10.0.0.0/8')).all(), [
                    (IPv4Address('192.168.0.1'),)
                ])
            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    ~self.table.c.x.in_('10.0.0.0/8')).all(), [
                    (IPv4Address('192.168.0.1'),)
                ])


class IPv4HttpTestCase(IPv4TestCase, HttpSessionTestCase):
    """ IPv4 over a HTTP session """


class IPv4NativeTestCase(IPv4TestCase, NativeSessionTestCase):
    """ IPv4 over a native protocol session """


class IPv6TestCase(TypesTestCase):
    table = Table(
        'test', TypesTestCase.metadata(),
        Column('x', types.IPv6),
        engines.Memory()
    )

    @require_server_version(19, 3, 3)
    def test_select_insert(self):
        a = IPv6Address('79f4:e698:45de:a59b:2765:28e3:8d3a:35ae')

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': a}])
            self.assertEqual(self.session.query(self.table.c.x).scalar(), a)

    @require_server_version(19, 3, 3)
    def test_select_insert_string(self):
        a = '79f4:e698:45de:a59b:2765:28e3:8d3a:35ae'

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': a}])
            self.assertEqual(self.session.query(self.table.c.x).scalar(),
                             IPv6Address(a))

    @require_server_version(19, 3, 3)
    def test_create_table(self):
        self.assertEqual(
            self.compile(CreateTable(self.table)),
            'CREATE TABLE test (x IPv6) ENGINE = Memory'
        )

    @require_server_version(19, 3, 3)
    def test_select_where_address(self):
        a = IPv6Address('42e::2')

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': a}])

            self.assertEqual(self.session.query(self.table.c.x).filter(
                self.table.c.x == IPv6Address('42e::2')).scalar(), a)

            self.assertEqual(self.session.query(self.table.c.x).filter(
                and_(IPv6Address('42e::1') < self.table.c.x,
                     self.table.c.x < IPv6Address('42e::3'))).scalar(), a)

    @require_server_version(19, 3, 3)
    def test_select_where_string(self):
        a = IPv6Address('42e::2')

        with self.create_table(self.table):
            self.session.execute(self.table.insert(), [{'x': a}])

            self.assertEqual(self.session.query(self.table.c.x).filter(
                and_('42e::1' < self.table.c.x,
                     self.table.c.x < '42e::3')).scalar(), a)

    @require_server_version(19, 3, 3)
    def test_select_in_network(self):
        ips = [
            IPv6Address('42e::1'),
            IPv6Address('42e::2'),
            IPv6Address('42e::3'),
            IPv6Address('f42e::ffff')
        ]

        with self.create_table(self.table):
            self.session.execute(self.table.insert(),
                                 [{'x': ip} for ip in ips])

            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    self.table.c.x.in_(IPv6Network('42e::/64'))).all(), [
                    (IPv6Address('42e::1'),),
                    (IPv6Address('42e::2'),),
                    (IPv6Address('42e::3'),)
                ])

    @require_server_version(19, 3, 3)
    def test_select_in_string(self):
        ips = [
            IPv6Address('42e::1'),
            IPv6Address('42e::2'),
            IPv6Address('42e::3'),
            IPv6Address('f42e::ffff')
        ]

        with self.create_table(self.table):
            self.session.execute(self.table.insert(),
                                 [{'x': ip} for ip in ips])

            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    self.table.c.x.in_('42e::/64')).all(), [
                    (IPv6Address('42e::1'),),
                    (IPv6Address('42e::2'),),
                    (IPv6Address('42e::3'),)
                ])

    @require_server_version(19, 3, 3)
    def test_select_not_in_network(self):
        ips = [
            IPv6Address('42e::1'),
            IPv6Address('42e::2'),
            IPv6Address('42e::3'),
            IPv6Address('f42e::ffff')
        ]

        with self.create_table(self.table):
            self.session.execute(self.table.insert(),
                                 [{'x': ip} for ip in ips])

            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    self.table.c.x.notin_(IPv6Network('42e::/64'))).all(), [
                    (IPv6Address('f42e::ffff'),)
                ])
            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    ~self.table.c.x.in_(IPv6Network('42e::/64'))).all(), [
                    (IPv6Address('f42e::ffff'),)
                ])

    @require_server_version(19, 3, 3)
    def test_select_not_in_string(self):
        ips = [
            IPv6Address('42e::1'),
            IPv6Address('42e::2'),
            IPv6Address('42e::3'),
            IPv6Address('f42e::ffff')
        ]

        with self.create_table(self.table):
            self.session.execute(self.table.insert(),
                                 [{'x': ip} for ip in ips])

            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    self.table.c.x.notin_('42e::/64')).all(), [
                    (IPv6Address('f42e::ffff'),)
                ])
            self.assertEqual(
                self.session.query(self.table.c.x).filter(
                    ~self.table.c.x.in_('42e::/64')).all(), [
                    (IPv6Address('f42e::ffff'),)
                ])


class IPv6HttpTestCase(IPv6TestCase, HttpSessionTestCase):
    """ IPv6 over a HTTP session """


class IPv6NativeTestCase(IPv6TestCase, NativeSessionTestCase):
    """ IPv6 over a native protocol session """

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/test_types.py
datalens-tech/datalens-backend
import re
from contextlib import contextmanager
from unittest import TestCase

from sqlalchemy import MetaData
from sqlalchemy.orm import Query

from .config import database, host, port, http_port, user, password
from .session import (
    http_session, native_session,
    system_http_session, system_native_session,
    http_tsv_session, http_jc_session,
)


class BaseAbstractTestCase(object):
    """ Supporting code for tests """

    host = host
    port = http_port
    database = database
    user = user
    password = password

    strip_spaces = re.compile(r'[\n\t]')
    session = http_session

    @classmethod
    def metadata(cls, session=None):
        if session is None:
            session = cls.session
        return MetaData(bind=session.bind)

    def _compile(self, clause, bind=None, literal_binds=False):
        if bind is None:
            bind = self.session.bind
        if isinstance(clause, Query):
            context = clause._compile_context()
            context.compile_state.statement.use_labels = True
            clause = context.compile_state.statement

        kw = {}
        compile_kwargs = {}
        if literal_binds:
            compile_kwargs['literal_binds'] = True

        if compile_kwargs:
            kw['compile_kwargs'] = compile_kwargs

        return clause.compile(dialect=bind.dialect, **kw)

    def compile(self, clause, **kwargs):
        return self.strip_spaces.sub(
            '', str(self._compile(clause, **kwargs))
        )

    @contextmanager
    def create_table(self, table):
        table.drop(bind=self.session.bind, if_exists=True)
        table.create(bind=self.session.bind)
        try:
            yield
        except Exception:
            raise
        finally:
            table.drop(bind=self.session.bind)


class BaseTestCase(BaseAbstractTestCase, TestCase):
    """ Actually tests """


class BaseHttpTestCase(BaseTestCase):

    @classmethod
    def setUpClass(cls):
        system_http_session.execute(
            'DROP DATABASE IF EXISTS {}'.format(cls.database))
        system_http_session.execute(
            'CREATE DATABASE {}'.format(cls.database))

        super(BaseHttpTestCase, cls).setUpClass()


class HttpSessionTestCase(BaseHttpTestCase):
    """ Explicitly HTTP-based session Test Case """

    session = http_tsv_session
    # session = http_jc_session


class HttpJCSessionTestCase(BaseTestCase):

    session = http_jc_session


class NativeSessionTestCase(BaseTestCase):
    """ Explicitly Native-protocol-based session Test Case """

    port = port

    session = native_session

    @classmethod
    def setUpClass(cls):
        # System database is always present.
        system_native_session.execute(
            'DROP DATABASE IF EXISTS {}'.format(cls.database)
        )
        system_native_session.execute(
            'CREATE DATABASE {}'.format(cls.database)
        )

        super(NativeSessionTestCase, cls).setUpClass()


class TypesTestCase(BaseAbstractTestCase):
    """ ... """

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/testcase.py
datalens-tech/datalens-backend
from functools import wraps


def require_server_version(*version_required):

    def check(f):

        @wraps(f)
        def wrapper(*args, **kwargs):
            self = args[0]
            conn = self.session.bind.raw_connection()

            dialect = self.session.bind.dialect.name
            if dialect == 'clickhouse+native':
                i = conn.transport.connection.server_info
                current = (i.version_major, i.version_minor, i.version_patch)
            else:
                cursor = conn.cursor()
                cursor.execute(
                    'SELECT version() FORMAT TabSeparatedWithNamesAndTypes'
                )
                version = cursor.fetchall()[0][0].split('.')
                current = tuple(int(x) for x in version[:3])

            conn.close()
            if current is not None and version_required <= current:
                return f(*args, **kwargs)
            else:
                self.skipTest(
                    'Mininum revision required: {}'.format(
                        '.'.join(str(x) for x in version_required)
                    )
                )

        return wrapper

    return check

lib/clickhouse-sqlalchemy/clickhouse_sqlalchemy_tests/util.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
    List,
    Optional,
)

import attr


@attr.s(frozen=True)
class ApiResponse:
    pass


@attr.s(frozen=True)
class HttpApiResponse(ApiResponse):
    """
    Client-independent BI HTTP API response object
    """

    _status_code: int = attr.ib(kw_only=True)
    _json: Dict[str, Any] = attr.ib(kw_only=True)

    @property
    def status_code(self) -> int:
        return self._status_code

    @property
    def bi_status_code(self) -> Optional[str]:
        return self.json.get("code")

    @property
    def json(self) -> dict:
        return self._json

    @property
    def response_errors(self) -> List[str]:
        return self.extract_response_errors(self._json)

    @classmethod
    def extract_response_errors(cls, response_json: Dict[str, Any]) -> List[str]:
        return []


class ApiBase:
    pass

lib/dl_api_client/dl_api_client/dsmaker/api/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from http import HTTPStatus
from typing import (
    Any,
    ClassVar,
    Dict,
    List,
    Optional,
    Union,
)

import attr

from dl_api_client.dsmaker.api.base import HttpApiResponse
from dl_api_client.dsmaker.api.http_async_base import AsyncHttpApiV1Base
from dl_api_client.dsmaker.api.http_sync_base import (
    ClientResponse,
    SyncHttpApiV1Base,
)
from dl_api_client.dsmaker.api.schemas.data import (
    FilterClauseSchema,
    PivotItemSchema,
    PivotPaginationSchema,
    PivotResponseSchema,
    PivotTotalsSchema,
    QueryBlockSchema,
    QueryFieldsItemSchema,
    ResultResponseSchema,
)
from dl_api_client.dsmaker.api.schemas.dataset import WhereClauseSchema
from dl_api_client.dsmaker.api.serialization_base import BaseApiV1SerializationAdapter
from dl_api_client.dsmaker.primitives import (
    BlockSpec,
    Dataset,
    OrderedField,
    ParameterFieldValue,
    PivotPagination,
    PivotTotals,
    RequestLegendItem,
    RequestPivotItem,
    ResultField,
    UpdateAction,
    WhereClause,
)
from dl_constants.enums import (
    FieldRole,
    OrderDirection,
    RangeType,
)


@attr.s(frozen=True)
class HttpDataApiResponse(HttpApiResponse):
    api_version: str = attr.ib(kw_only=True)
    _data: Optional[dict] = attr.ib(kw_only=True, default=None)

    @property
    def data(self) -> dict:
        assert self._data is not None
        return self._data

    @classmethod
    def extract_response_errors(cls, response_json: Dict[str, Any]) -> List[str]:
        errors = super().extract_response_errors(response_json)
        if "message" in response_json:
            errors.append(response_json["message"])
        return errors


@attr.s
class DataApiV1SerializationAdapter(BaseApiV1SerializationAdapter):
    def make_req_data_get_preview(
        self,
        *,
        dataset: Dataset,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
    ) -> Dict[str, Any]:
        data = self.dump_dataset(dataset)
        updates = list(updates or ()) + self.generate_implicit_updates(dataset)  # type: ignore  # TODO: fix
        data["updates"] = self.dump_updates(updates)

        if limit is not None:
            data["limit"] = limit

        return data

    def make_req_data_get_result(
        self,
        *,
        # TODO FIX: Remove after check that frontend team does not use it
        dataset: Optional[Dataset] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        fields: Optional[List[ResultField]] = None,
        group_by: Optional[List[ResultField]] = None,
        order_by: Optional[List[Union[ResultField, OrderedField]]] = None,
        where: Optional[List[WhereClause]] = None,
        disable_group_by: Optional[bool] = None,
        with_totals: Optional[bool] = None,
        updates: List[Union[UpdateAction, dict]] = None,
    ) -> Dict[str, Any]:
        if dataset is not None:
            data = self.dump_dataset(dataset)
            updates = list(updates or ()) + self.generate_implicit_updates(dataset)  # type: ignore  # TODO: fix
        else:
            data = {}

        if limit is not None:
            data["limit"] = limit
        if offset is not None:
            data["offset"] = offset
        if fields is not None:
            data["columns"] = [f.id for f in fields]
        if group_by is not None:
            data["group_by"] = [f.id for f in group_by]
        if order_by is not None:
            data["order_by"] = []
            direction_map = {OrderDirection.asc: "ASC", OrderDirection.desc: "DESC"}
            for ob in order_by:
                if isinstance(ob, ResultField):
                    ob = ob.asc
                assert isinstance(ob, OrderedField)
                data["order_by"].append({"column": ob.field_id, "direction": direction_map[ob.direction]})
        if where is not None:
            data["where"] = []
            where_schema = WhereClauseSchema()
            for condition in where:
                data["where"].append(where_schema.dump(condition))
        if updates:
            data["updates"] = self.dump_updates(updates)
        if disable_group_by is not None:
            data["disable_group_by"] = disable_group_by
        if with_totals is not None:
            data["with_totals"] = with_totals

        return data

    def make_req_data_get_value_range(
        self,
        *,
        dataset: Optional[Dataset] = None,
        field: ResultField,
        filters: Optional[List[WhereClause]] = None,
        updates: List[Union[UpdateAction, dict]] = None,
    ) -> Dict[str, Any]:
        if dataset is not None:
            data = self.dump_dataset(dataset)
            updates = list(updates or ()) + self.generate_implicit_updates(dataset)  # type: ignore  # TODO: fix
        else:
            data = {}

        data["field_guid"] = field.id
        if filters is not None:
            data["where"] = []
            where_schema = WhereClauseSchema()
            for condition in filters:
                data["where"].append(where_schema.dump(condition))
        if updates:
            data["updates"] = self.dump_updates(updates)

        return data

    def make_req_data_get_distinct(
        self,
        *,
        dataset: Optional[Dataset] = None,
        limit: Optional[int] = None,
        field: ResultField,
        filters: Optional[List[WhereClause]] = None,
        updates: List[Union[UpdateAction, dict]] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
    ) -> Dict[str, Any]:
        if dataset is not None:
            data = self.dump_dataset(dataset)
            updates = list(updates or ()) + self.generate_implicit_updates(dataset)  # type: ignore  # TODO: fix
        else:
            data = {}

        if limit is not None:
            data["limit"] = limit
        data["field_guid"] = field.id
        if filters is not None:
            data["where"] = []
            where_schema = WhereClauseSchema()
            for condition in filters:
                data["where"].append(where_schema.dump(condition))
        if updates:
            data["updates"] = self.dump_updates(updates)
        if ignore_nonexistent_filters is not None:
            data["ignore_nonexistent_filters"] = ignore_nonexistent_filters

        return data


@attr.s
class DataApiBaseMixin:
    api_v: ClassVar[str]

    def make_response_obj(self, client_response: ClientResponse, data: dict) -> HttpDataApiResponse:
        return HttpDataApiResponse(
            json=client_response.json,
            status_code=client_response.status_code,
            data=data,
            api_version=self.api_v,
        )


@attr.s
class SyncHttpDataApiBase(SyncHttpApiV1Base, DataApiBaseMixin):
    pass


@attr.s
class AsyncHttpDataApiBase(AsyncHttpApiV1Base, DataApiBaseMixin):
    pass


@attr.s
class SyncHttpDataApiV1(SyncHttpDataApiBase):
    api_v = "v1"
    # Override serial_adapter with subclass
    serial_adapter: DataApiV1SerializationAdapter = attr.ib(init=False, factory=DataApiV1SerializationAdapter)

    def get_response_for_dataset_preview(
        self,
        *,
        dataset_id: Optional[str] = None,
        raw_body: Optional[dict] = None,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        if dataset_id is None:
            url = f"/api/data/{self.api_v}/datasets/data/preview"
        else:
            url = f"/api/data/{self.api_v}/datasets/{dataset_id}/versions/draft/preview"

        return self._request(url, method="post", data=raw_body or {}, headers=headers)

    def get_response_for_dataset_result(
        self, dataset_id: str, raw_body: dict, headers: Optional[dict] = None
    ) -> ClientResponse:
        """:returns: Flask test client like response"""
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/versions/draft/result"
        return self._request(url, method="post", data=raw_body, headers=headers)

    def get_response_for_dataset_value_range(
        self,
        dataset_id: str,
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        """:returns: Flask test client like response"""
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/versions/draft/values/range"
        return self._request(url, method="post", data=raw_body, headers=headers)

    def get_response_for_dataset_value_distinct(
        self,
        dataset_id: str,
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        """:returns: Flask test client like response"""
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/versions/draft/values/distinct"
        return self._request(url, method="post", data=raw_body, headers=headers)

    def get_preview(
        self,
        *,
        dataset: Dataset,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_preview(
            dataset=dataset,
            updates=updates,
            limit=limit,
        )
        response = self.get_response_for_dataset_preview(
            dataset_id=(dataset.id if dataset.created_ else None),
            raw_body=data,
        )

        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_data = response.json["result"]["data"]
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    def get_result(
        self,
        *,
        dataset: Optional[Dataset] = None,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
        offset: int = None,
        fields: Optional[List[ResultField]] = None,
        group_by: Optional[List[ResultField]] = None,
        order_by: Optional[List[Union[ResultField, OrderedField]]] = None,
        where: Optional[List[WhereClause]] = None,  # TODO: Rename -> filters
        filters: Optional[List[WhereClause]] = None,
        disable_group_by: Optional[bool] = None,
        with_totals: Optional[bool] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        filters = filters or where  # For compatibility
        data = self.serial_adapter.make_req_data_get_result(
            dataset=dataset,
            limit=limit,
            offset=offset,
            fields=fields,
            group_by=group_by,
            order_by=order_by,
            where=filters,
            disable_group_by=disable_group_by,
            with_totals=with_totals,
            updates=updates,
        )
        response = self.get_response_for_dataset_result(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_data = response.json["result"]["data"]
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    def get_value_range(
        self,
        *,
        dataset: Optional[Dataset],
        field: ResultField,
        updates: List[Union[UpdateAction, dict]] = None,
        where: Optional[List[WhereClause]] = None,  # TODO: Rename -> filters
        filters: Optional[List[WhereClause]] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        filters = filters or where
        data = self.serial_adapter.make_req_data_get_value_range(
            field=field,
            dataset=dataset,
            filters=filters,
            updates=updates,
        )
        response = self.get_response_for_dataset_value_range(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_data = response.json["result"]["data"]
            except KeyError:
                raise ValueError(response.data)

        return HttpDataApiResponse(
            api_version=self.api_v,
            json=response.json,
            status_code=response.status_code,
            data=resp_data,
        )

    def get_distinct(
        self,
        *,
        dataset: Optional[Dataset] = None,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
        field: ResultField,
        where: Optional[List[WhereClause]] = None,  # TODO: Rename -> filters
        filters: Optional[List[WhereClause]] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        filters = filters or where
        data = self.serial_adapter.make_req_data_get_distinct(
            dataset=dataset,
            limit=limit,
            field=field,
            filters=filters,
            updates=updates,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
        )
        response = self.get_response_for_dataset_value_distinct(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_data = response.json["result"]["data"]
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    def get_fields(self, dataset_id: str, fail_ok: bool = False) -> HttpDataApiResponse:
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/fields"
        response = self._request(url, method="get")
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json

        return self.make_response_obj(client_response=response, data=resp_data)


@attr.s
class DataApiV2SerializationAdapter(BaseApiV1SerializationAdapter):
    def make_req_data_common(
        self,
        *,
        dataset: Optional[Dataset] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        fields: Optional[List[Union[ResultField, RequestLegendItem]]] = None,
        order_by: Optional[List[Union[ResultField, OrderedField]]] = None,
        filters: Optional[List[WhereClause]] = None,
        parameters: Optional[List[Union[ResultField, ParameterFieldValue]]] = None,
        blocks: Optional[List[BlockSpec]] = None,
        autofill_legend: bool = False,
        ignore_nonexistent_filters: Optional[bool] = None,
        row_count_hard_limit: Optional[int] = None,
        updates: List[Union[UpdateAction, dict]] = None,
    ) -> Dict[str, Any]:
        if dataset is not None:
            data = self.dump_dataset(dataset)
            updates = list(updates or ()) + self.generate_implicit_updates(dataset)  # type: ignore  # TODO: fix
        else:
            data = {}

        if limit is not None:
            data["limit"] = limit

        if offset is not None:
            data["offset"] = offset

        if fields is not None:
            fields_item_schema = QueryFieldsItemSchema()
            legend_data: List[dict] = []
            for field in fields:
                if isinstance(field, ResultField):
                    field = field.as_req_legend_item(role=FieldRole.row)
                assert isinstance(field, RequestLegendItem)
                legend_data.append(fields_item_schema.dump(field))
            data["fields"] = legend_data

        if order_by is not None:
            order_by_data: List[dict] = []
            for ob_field in order_by:
                if isinstance(ob_field, ResultField):
                    ob_field = ob_field.asc
                assert isinstance(ob_field, OrderedField)
                order_by_data.append(
                    {
                        "ref": {"type": "id", "id": ob_field.field_id},
                        "direction": ob_field.direction.name,
                        "block_id": ob_field.block_id,
                    }
                )
            data["order_by"] = order_by_data

        if filters is not None:
            filter_data: List[dict] = []
            filter_schema = FilterClauseSchema()
            for condition in filters:
                filter_data.append(filter_schema.dump(condition))
            data["filters"] = filter_data

        if parameters is not None:
            parameters_data: List[dict] = []
            for param_field in parameters:
                if isinstance(param_field, ResultField):
                    param_field = param_field.parameter_value()
                assert isinstance(param_field, ParameterFieldValue)
                parameters_data.append(
                    {
                        "ref": {
                            "type": "id",
                            "id": param_field.field_id,
                        },
                        "value": param_field.value,
                        "block_id": param_field.block_id,
                    }
                )
            data["parameter_values"] = parameters_data

        if blocks is not None:
            block_data: List[dict] = []
            block_schema = QueryBlockSchema()
            for block in blocks:
                block_data.append(block_schema.dump(block))
            data["blocks"] = block_data

        if updates:
            data["updates"] = self.dump_updates(updates)

        data["autofill_legend"] = autofill_legend

        if ignore_nonexistent_filters is not None:
            data["ignore_nonexistent_filters"] = ignore_nonexistent_filters

        if row_count_hard_limit is not None:
            data["row_count_hard_limit"] = row_count_hard_limit

        return data

    def make_req_data_get_preview(
        self,
        *,
        dataset: Dataset,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
    ) -> Dict[str, Any]:
        data = self.dump_dataset(dataset)
        updates = list(updates or ()) + self.generate_implicit_updates(dataset)
        data["updates"] = self.dump_updates(updates)

        if limit is not None:
            data["limit"] = limit

        return data

    def make_req_data_get_result(
        self,
        *,
        dataset: Optional[Dataset] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        fields: Optional[List[Union[ResultField, RequestLegendItem]]] = None,
        order_by: Optional[List[Union[ResultField, OrderedField]]] = None,
        filters: Optional[List[WhereClause]] = None,
        parameters: Optional[List[Union[ResultField, ParameterFieldValue]]] = None,
        blocks: Optional[List[BlockSpec]] = None,
        disable_group_by: Optional[bool] = None,
        with_totals: Optional[bool] = None,
        result_with_totals: Optional[bool] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        row_count_hard_limit: Optional[int] = None,
        updates: List[Union[UpdateAction, dict]] = None,
    ) -> Dict[str, Any]:
        data = self.make_req_data_common(
            dataset=dataset,
            limit=limit,
            offset=offset,
            fields=fields,
            order_by=order_by,
            filters=filters,
            parameters=parameters,
            blocks=blocks,
            updates=updates,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
            row_count_hard_limit=row_count_hard_limit,
        )
        if disable_group_by is not None:
            data["disable_group_by"] = disable_group_by
        if with_totals is not None:
            data["with_totals"] = with_totals
        if result_with_totals is not None:
            data["result"] = {}
            data["result"]["with_totals"] = with_totals

        return data

    def make_req_data_get_value_range(
        self,
        *,
        dataset: Optional[Dataset] = None,
        field: ResultField,
        filters: Optional[List[WhereClause]] = None,
        parameters: Optional[List[Union[ResultField, ParameterFieldValue]]] = None,
        blocks: Optional[List[BlockSpec]] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        updates: List[Union[UpdateAction, dict]] = None,
    ) -> Dict[str, Any]:
        fields = [
            field.as_req_legend_item(role=FieldRole.range, range_type=RangeType.min),
            field.as_req_legend_item(role=FieldRole.range, range_type=RangeType.max),
        ]
        data = self.make_req_data_common(
            dataset=dataset,
            fields=fields,
            filters=filters,
            parameters=parameters,
            blocks=blocks,
            updates=updates,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
        )

        return data

    def make_req_data_get_distinct(
        self,
        *,
        dataset: Optional[Dataset] = None,
        field: Union[ResultField, RequestLegendItem],
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        order_by: Optional[List[Union[ResultField, OrderedField]]] = None,
        filters: Optional[List[WhereClause]] = None,
        parameters: Optional[List[Union[ResultField, ParameterFieldValue]]] = None,
        blocks: Optional[List[BlockSpec]] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        updates: List[Union[UpdateAction, dict]] = None,
    ) -> Dict[str, Any]:
        field_li: RequestLegendItem
        if isinstance(field, RequestLegendItem):
            assert field.role_spec.role == FieldRole.distinct
            field_li = field
        else:
            field_li = field.as_req_legend_item(role=FieldRole.distinct)
        fields = [field_li]
        data = self.make_req_data_common(
            dataset=dataset,
            limit=limit,
            offset=offset,
            fields=fields,
            order_by=order_by,
            filters=filters,
            parameters=parameters,
            blocks=blocks,
            updates=updates,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
        )

        return data

    def make_req_data_get_pivot(
        self,
        *,
        dataset: Optional[Dataset] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        fields: Optional[List[Union[ResultField, RequestLegendItem]]] = None,
        order_by: Optional[List[OrderedField]] = None,
        filters: Optional[List[WhereClause]] = None,
        blocks: Optional[List[BlockSpec]] = None,
        autofill_legend: bool = False,
        ignore_nonexistent_filters: Optional[bool] = None,
        pivot_pagination: Optional[PivotPagination] = None,
        pivot_structure: Optional[list[RequestPivotItem]] = None,
        pivot_totals: Optional[PivotTotals] = None,
        with_totals: Optional[bool] = None,
        updates: List[Union[UpdateAction, dict]] = None,
    ) -> Dict[str, Any]:
        data = self.make_req_data_common(
            dataset=dataset,
            limit=limit,
            offset=offset,
            fields=fields,
            order_by=order_by,
            filters=filters,
            blocks=blocks,
            autofill_legend=autofill_legend,
            updates=updates,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
        )
        data["pivot"] = {}
        if pivot_structure is not None:
            data["pivot"]["structure"] = [PivotItemSchema().dump(item) for item in pivot_structure]
        if pivot_pagination is not None:
            data["pivot"]["pagination"] = PivotPaginationSchema().dump(pivot_pagination)
        if with_totals is not None:
            data["pivot"]["with_totals"] = with_totals
        if pivot_totals is not None:
            data["pivot"]["totals"] = PivotTotalsSchema().dump(pivot_totals)

        return data


@attr.s
class SyncHttpDataApiV2(SyncHttpDataApiBase):
    api_v = "v2"
    # Override serial_adapter with subclass
    serial_adapter: DataApiV2SerializationAdapter = attr.ib(init=False, factory=DataApiV2SerializationAdapter)

    def _make_resp_data(self, full_response_data: dict) -> dict:
        return full_response_data

    def get_response_for_dataset_preview(
        self,
        dataset_id: Optional[str],
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        if dataset_id is None:
            url = f"/api/data/{self.api_v}/datasets/data/preview"
        else:
            url = f"/api/data/{self.api_v}/datasets/{dataset_id}/preview"
        return self._request(url, method="post", data=raw_body, headers=headers)

    def get_response_for_dataset_result(
        self,
        dataset_id: str,
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/result"
        return self._request(url, method="post", data=raw_body, headers=headers)

    def get_response_for_dataset_value_range(
        self,
        dataset_id: str,
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/values/range"
        return self._request(url, method="post", data=raw_body, headers=headers)

    def get_response_for_dataset_value_distinct(
        self,
        dataset_id: str,
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/values/distinct"
        return self._request(url, method="post", data=raw_body, headers=headers)

    def get_response_for_dataset_pivot(
        self,
        dataset_id: str,
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/pivot"
        return self._request(url, method="post", data=raw_body, headers=headers)

    def get_preview(
        self,
        *,
        dataset: Dataset,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_preview(dataset=dataset, limit=limit, updates=updates)
        response = self.get_response_for_dataset_preview(
            dataset_id=(dataset.id if dataset.created_ else None),
            raw_body=data,
        )
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_schema = ResultResponseSchema()
                resp_data = resp_schema.load(response.json)
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    def get_result(
        self,
        *,
        dataset: Optional[Dataset] = None,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        fields: Optional[List[Union[ResultField, RequestLegendItem]]] = None,
        order_by: Optional[List[Union[ResultField, OrderedField]]] = None,
        filters: Optional[List[WhereClause]] = None,
        parameters: Optional[List[Union[ResultField, ParameterFieldValue]]] = None,
        blocks: Optional[List[BlockSpec]] = None,
        disable_group_by: Optional[bool] = None,
        with_totals: Optional[bool] = None,
        result_with_totals: Optional[bool] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        row_count_hard_limit: Optional[int] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_result(
            dataset=dataset,
            limit=limit,
            offset=offset,
            fields=fields,
            order_by=order_by,
            filters=filters,
            parameters=parameters,
            blocks=blocks,
            disable_group_by=disable_group_by,
            with_totals=with_totals,
            result_with_totals=result_with_totals,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
            row_count_hard_limit=row_count_hard_limit,
            updates=updates,
        )
        response = self.get_response_for_dataset_result(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_schema = ResultResponseSchema()
                resp_data = resp_schema.load(response.json)
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    def get_value_range(
        self,
        *,
        dataset: Optional[Dataset],
        field: ResultField,
        updates: List[Union[UpdateAction, dict]] = None,
        filters: Optional[List[WhereClause]] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_value_range(
            field=field,
            dataset=dataset,
            filters=filters,
            updates=updates,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
        )
        response = self.get_response_for_dataset_value_range(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_schema = ResultResponseSchema()
                resp_data = resp_schema.load(response.json)
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    def get_distinct(
        self,
        *,
        dataset: Optional[Dataset] = None,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
        field: ResultField,
        filters: Optional[List[WhereClause]] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_distinct(
            dataset=dataset,
            limit=limit,
            field=field,
            filters=filters,
            updates=updates,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
        )
        response = self.get_response_for_dataset_value_distinct(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_schema = ResultResponseSchema()
                resp_data = resp_schema.load(response.json)
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    def get_pivot(
        self,
        *,
        dataset: Optional[Dataset] = None,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
        offset: Optional[int] = None,
        fields: Optional[List[Union[ResultField, RequestLegendItem]]] = None,
        order_by: Optional[List[OrderedField]] = None,
        filters: Optional[List[WhereClause]] = None,
        blocks: Optional[List[BlockSpec]] = None,
        autofill_legend: bool = False,
        pivot_structure: Optional[list[RequestPivotItem]] = None,
        pivot_pagination: Optional[PivotPagination] = None,
        pivot_totals: Optional[PivotTotals] = None,
        with_totals: Optional[bool] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_pivot(
            dataset=dataset,
            limit=limit,
            offset=offset,
            fields=fields,
            order_by=order_by,
            filters=filters,
            blocks=blocks,
            updates=updates,
            autofill_legend=autofill_legend,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
            pivot_structure=pivot_structure,
            pivot_pagination=pivot_pagination,
            pivot_totals=pivot_totals,
            with_totals=with_totals,
        )
        response = self.get_response_for_dataset_pivot(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_schema = PivotResponseSchema()
                resp_data = resp_schema.load(response.json)
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)


@attr.s
class AsyncHttpDataApiV2(AsyncHttpDataApiBase):
    api_v = "v2"
    # Override serial_adapter with subclass
    serial_adapter: DataApiV2SerializationAdapter = attr.ib(init=False, factory=DataApiV2SerializationAdapter)

    def _make_resp_data(self, full_response_data: dict) -> dict:
        return full_response_data

    async def get_response_for_dataset_preview(
        self,
        dataset_id: Optional[str],
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        if dataset_id is None:
            url = f"/api/data/{self.api_v}/datasets/data/preview"
        else:
            url = f"/api/data/{self.api_v}/datasets/{dataset_id}/preview"
        return await self._request(url, method="post", data=raw_body, headers=headers)

    async def get_response_for_dataset_result(
        self,
        dataset_id: str,
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/result"
        return await self._request(url, method="post", data=raw_body, headers=headers)

    async def get_response_for_dataset_value_range(
        self,
        dataset_id: str,
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/values/range"
        return await self._request(url, method="post", data=raw_body, headers=headers)

    async def get_response_for_dataset_value_distinct(
        self,
        dataset_id: str,
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/values/distinct"
        return await self._request(url, method="post", data=raw_body, headers=headers)

    async def get_response_for_dataset_pivot(
        self,
        dataset_id: str,
        raw_body: dict,
        headers: Optional[dict] = None,
    ) -> ClientResponse:
        url = f"/api/data/{self.api_v}/datasets/{dataset_id}/pivot"
        return await self._request(url, method="post", data=raw_body, headers=headers)

    async def get_preview(
        self,
        *,
        dataset: Dataset,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_preview(dataset=dataset, limit=limit, updates=updates)
        response = await self.get_response_for_dataset_preview(
            dataset_id=(dataset.id if dataset.created_ else None),
            raw_body=data,
        )
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_data = self._make_resp_data(response.json)
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    async def get_result(
        self,
        *,
        dataset: Optional[Dataset] = None,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        fields: Optional[List[Union[ResultField, RequestLegendItem]]] = None,
        order_by: Optional[List[Union[ResultField, OrderedField]]] = None,
        filters: Optional[List[WhereClause]] = None,
        parameters: Optional[List[Union[ResultField, ParameterFieldValue]]] = None,
        blocks: Optional[List[BlockSpec]] = None,
        disable_group_by: Optional[bool] = None,
        with_totals: Optional[bool] = None,
        result_with_totals: Optional[bool] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        row_count_hard_limit: Optional[int] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_result(
            dataset=dataset,
            limit=limit,
            offset=offset,
            fields=fields,
            order_by=order_by,
            filters=filters,
            parameters=parameters,
            blocks=blocks,
            disable_group_by=disable_group_by,
            with_totals=with_totals,
            result_with_totals=result_with_totals,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
            row_count_hard_limit=row_count_hard_limit,
            updates=updates,
        )
        response = await self.get_response_for_dataset_result(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_data = self._make_resp_data(response.json)
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    async def get_value_range(
        self,
        *,
        dataset: Optional[Dataset],
        field: ResultField,
        updates: List[Union[UpdateAction, dict]] = None,
        filters: Optional[List[WhereClause]] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_value_range(
            field=field,
            dataset=dataset,
            filters=filters,
            updates=updates,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
        )
        response = await self.get_response_for_dataset_value_range(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_data = self._make_resp_data(response.json)
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    async def get_distinct(
        self,
        *,
        dataset: Optional[Dataset] = None,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
        field: ResultField,
        filters: Optional[List[WhereClause]] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_distinct(
            dataset=dataset,
            limit=limit,
            field=field,
            filters=filters,
            updates=updates,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
        )
        response = await self.get_response_for_dataset_value_distinct(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_data = self._make_resp_data(response.json)
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)

    async def get_pivot(
        self,
        *,
        dataset: Optional[Dataset] = None,
        updates: List[Union[UpdateAction, dict]] = None,
        limit: int = None,
        offset: Optional[int] = None,
        fields: Optional[List[Union[ResultField, RequestLegendItem]]] = None,
        order_by: Optional[List[OrderedField]] = None,
        filters: Optional[List[WhereClause]] = None,
        blocks: Optional[List[BlockSpec]] = None,
        autofill_legend: bool = False,
        pivot_structure: Optional[list[RequestPivotItem]] = None,
        pivot_pagination: Optional[PivotPagination] = None,
        pivot_totals: Optional[PivotTotals] = None,
        with_totals: Optional[bool] = None,
        ignore_nonexistent_filters: Optional[bool] = None,
        fail_ok: bool = False,
    ) -> HttpDataApiResponse:
        data = self.serial_adapter.make_req_data_get_pivot(
            dataset=dataset,
            limit=limit,
            offset=offset,
            fields=fields,
            order_by=order_by,
            filters=filters,
            blocks=blocks,
            updates=updates,
            autofill_legend=autofill_legend,
            ignore_nonexistent_filters=ignore_nonexistent_filters,
            pivot_structure=pivot_structure,
            pivot_pagination=pivot_pagination,
            pivot_totals=pivot_totals,
            with_totals=with_totals,
        )
        response = await self.get_response_for_dataset_pivot(dataset_id=dataset.id, raw_body=data)
        resp_data: Optional[dict] = None
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        if response.status_code == HTTPStatus.OK:
            try:
                resp_schema = PivotResponseSchema()
                resp_data = resp_schema.load(response.json)
            except KeyError:
                raise ValueError(response.data)

        return self.make_response_obj(client_response=response, data=resp_data)


@attr.s
class SyncHttpDataApiV1_5(SyncHttpDataApiV2):
    api_v = "v1.5"

    def _make_resp_data(self, full_response_data: dict) -> dict:
        return full_response_data["result"]["data"]

lib/dl_api_client/dl_api_client/dsmaker/api/data_api.py
datalens-tech/datalens-backend
from __future__ import annotations

from http import HTTPStatus
from itertools import chain
from typing import (
    Any,
    Iterable,
    Optional,
    Union,
)

import attr

from dl_api_client.dsmaker.api.base import HttpApiResponse
from dl_api_client.dsmaker.api.http_sync_base import SyncHttpApiV1Base
from dl_api_client.dsmaker.api.schemas.dataset import (
    AvatarRelationSchema,
    ComponentErrorListSchema,
    DataSourceSchema,
    ObligatoryFilterSchema,
    ResultFieldSchema,
    ResultSchemaAuxSchema,
    SourceAvatarSchema,
)
from dl_api_client.dsmaker.api.serialization_base import BaseApiV1SerializationAdapter
from dl_api_client.dsmaker.primitives import (
    Dataset,
    ResultField,
    UpdateAction,
)


@attr.s(frozen=True)
class HttpDatasetApiResponse(HttpApiResponse):
    _dataset: Optional[Dataset] = attr.ib(kw_only=True, default=None)

    @property
    def dataset(self) -> Dataset:
        assert self._dataset is not None
        return self._dataset

    @classmethod
    def extract_response_errors(cls, response_json: dict[str, Any]) -> list[str]:
        errors = super().extract_response_errors(response_json)
        for key, value in chain.from_iterable((response_json.items(), response_json.get("dataset", {}).items())):
            if key.endswith("errors") and value:
                errors.append(f"{key}: {value}")

        return errors


@attr.s
class DatasetApiV1SerializationAdapter(BaseApiV1SerializationAdapter):
    @staticmethod
    def load_dataset_from_response_body(dataset: Dataset, body: dict) -> Dataset:
        """
        Load dataset contents from response

        :param dataset: old dataset object, used for loading aliases of nested items and other properties
        :param body: API response body
        """
        new_id = body.get("id") or dataset.id
        new_name = body.get("name", dataset.name)
        body = body["dataset"]
        new_dataset = Dataset(
            created_=dataset.created_,
            name=new_name,
            id=new_id,
            revision_id=body.get("revision_id"),
            rls=body.get("rls") or {},
            component_errors=ComponentErrorListSchema().load(body.get("component_errors") or {}),
            result_schema_aux=ResultSchemaAuxSchema().load(body.get("result_schema_aux") or {}),
            obligatory_filters=[
                ObligatoryFilterSchema().load(filter_info) for filter_info in body["obligatory_filters"]
            ],
        )

        for dsrc_data in body["sources"]:
            source_id = dsrc_data["id"]
            source_alias = dataset.sources.get_alias(source_id) or source_id
            new_dataset.sources[source_alias] = DataSourceSchema().load(dsrc_data)

        for avatar_data in body["source_avatars"]:
            avatar_id = avatar_data["id"]
            avatar_alias = dataset.source_avatars.get_alias(avatar_id) or avatar_id
            new_dataset.source_avatars[avatar_alias] = SourceAvatarSchema().load(avatar_data)

        for relation_data in body["avatar_relations"]:
            relation_id = relation_data["id"]
            relation_alias = dataset.avatar_relations.get_alias(relation_id) or relation_id
            new_dataset.avatar_relations[relation_alias] = AvatarRelationSchema().load(relation_data)

        for field_data in body["result_schema"]:
            field_id = field_data["guid"]
            field_alias = dataset.result_schema.get_alias(field_id) or field_id
            new_dataset.result_schema[field_alias] = ResultFieldSchema().load(field_data)

        return new_dataset


@attr.s
class SyncHttpDatasetApiV1(SyncHttpApiV1Base):
    _created_dataset_id_list: list[str] = attr.ib(init=False, factory=list)
    # Override serial_adapter with subclass
    serial_adapter: DatasetApiV1SerializationAdapter = attr.ib(init=False, factory=DatasetApiV1SerializationAdapter)

    def save_dataset(
        self,
        dataset: Dataset,
        preview: bool = True,
        fail_ok: bool = False,
        dir_path: str = None,
        workbook_id: str = None,
        created_via: str = None,
        lock_timeout: int = 3,
    ) -> HttpDatasetApiResponse:
        dataset.prepare()

        if self.serial_adapter.generate_implicit_updates(dataset=dataset):
            raise RuntimeError("Dataset has pending updates that have not been applied yet!")

        data = dict(self.dump_dataset_to_request_body(dataset), preview=preview)
        if not dataset.created_:
            if dir_path is not None and workbook_id is not None:
                raise RuntimeError("workbook_id and dir_path can not be used simultaneously")
            if workbook_id is not None:
                data["workbook_id"] = workbook_id
            if dir_path is not None:
                data["dir_path"] = dir_path
            if created_via is not None:
                data["created_via"] = created_via
            response = self._request(
                "/api/v1/datasets",
                method="post",
                data=dict(data, name=dataset.name),
                lock_timeout=lock_timeout,
            )

        else:
            response = self._request(
                f"/api/v1/datasets/{dataset.id}/versions/draft",
                method="put",
                data=data,
                lock_timeout=lock_timeout,
            )

        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        else:
            # `fail_ok` should not mean `allow 5xx`
            assert response.status_code < 500, response.json

        new_dataset: Optional[Dataset] = None
        if response.status_code == HTTPStatus.OK:
            new_dataset = self.serial_adapter.load_dataset_from_response_body(dataset=dataset, body=response.json)
            if not dataset.created_:
                new_dataset_id = response.json["id"]
                new_dataset.created_ = True
                new_dataset.id = new_dataset_id
                self._created_dataset_id_list.append(new_dataset_id)

        return HttpDatasetApiResponse(
            json=response.json,
            status_code=response.status_code,
            dataset=new_dataset,
        )

    def delete_dataset(self, dataset_id: str, fail_ok: bool = False) -> None:
        response = self._request(
            f"/api/v1/datasets/{dataset_id}",
            method="delete",
        )
        if not fail_ok:
            assert response.status_code == HTTPStatus.OK, response.json
        else:
            # `fail_ok` should not mean `allow 5xx`
            assert response.status_code < 500, response.json

    def cleanup_created_resources(self):  # type: ignore  # TODO: fix
        for dataset_id in self._created_dataset_id_list:
            try:
                self.delete_dataset(dataset_id, fail_ok=True)
            except Exception:  # noqa
                pass

    def copy_dataset(
        self,
        dataset: Dataset,
        new_key: str,
    ) -> HttpDatasetApiResponse:
        data = {"new_key": new_key}

        response = self._request(f"/api/v1/datasets/{dataset.id}/copy", method="post", data=data)
        assert response.status_code == 200, f"Dataset copy fail: {response.data}"
        dataset = self.serial_adapter.load_dataset_from_response_body(dataset=dataset, body=response.json)
        return HttpDatasetApiResponse(
            json=response.json,
            status_code=response.status_code,
            dataset=dataset,
        )

    def load_dataset(self, dataset: Dataset) -> HttpDatasetApiResponse:
        response = self._request(f"/api/v1/datasets/{dataset.id}/versions/draft", method="get")
        dataset = self.serial_adapter.load_dataset_from_response_body(dataset=dataset, body=response.json)
        dataset.created_ = True
        return HttpDatasetApiResponse(
            json=response.json,
            status_code=response.status_code,
            dataset=dataset,
        )

    def apply_updates(  # type: ignore  # TODO: fix
        self,
        dataset: Dataset,
        updates: list[Union[UpdateAction, dict]] = None,
        fail_ok: bool = False,
    ) -> HttpDatasetApiResponse:
        if dataset.created_:
            url = f"/api/v1/datasets/{dataset.id}/versions/draft/validators/schema"
        else:
            url = "/api/v1/datasets/validators/dataset"

        data = self.dump_dataset_to_request_body(dataset)
        updates = list(updates or ()) + self.serial_adapter.generate_implicit_updates(dataset)
        data["updates"] = self.serial_adapter.dump_updates(updates)
        response = self._request(url, method="post", data=data)

        if fail_ok:
            # `fail_ok` should not mean `allow 5xx`
            assert response.status_code < 500, response.json
        else:
            try:
                assert response.status_code == HTTPStatus.OK, response.json
            except AssertionError:
                print("\n".join(HttpApiResponse.extract_response_errors(response)))
                raise

        if "dataset" in response.json:
            dataset = self.serial_adapter.load_dataset_from_response_body(dataset=dataset, body=response.json)
        return HttpDatasetApiResponse(
            json=response.json,
            status_code=response.status_code,
            dataset=dataset,
        )

    def validate_field(self, dataset: Dataset, field: ResultField) -> HttpDatasetApiResponse:
        if dataset.created_:
            url = f"/api/v1/datasets/{dataset.id}/versions/draft/validators/field"
        else:
            url = "/api/v1/datasets/validators/field"

        if not field.id:
            field.prepare()

        data = self.dump_dataset_to_request_body(dataset)
        data.update({"field": {"guid": field.id, "title": field.title, "formula": field.formula}})
        response = self._request(url, method="post", data=data)
        return HttpDatasetApiResponse(
            json=response.json,
            status_code=response.status_code,
            dataset=None,
        )

    def refresh_dataset_sources(
        self,
        dataset: Dataset,
        source_ids: Iterable[str],
        fail_ok: bool = False,
    ) -> HttpDatasetApiResponse:
        refresh_resp = self.apply_updates(
            dataset=dataset,
            fail_ok=fail_ok,
            updates=[
                {
                    "action": "refresh_source",
                    "source": {
                        "id": source_id,
                    },
                }
                for source_id in source_ids
            ],
        )
        return refresh_resp

lib/dl_api_client/dl_api_client/dsmaker/api/dataset_api.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
from http import HTTPStatus
import json
import time
from typing import (
    Mapping,
    Optional,
)

import attr

from dl_api_client.dsmaker.api.base import ApiBase
from dl_api_client.dsmaker.api.http_sync_base import ClientResponse
from dl_api_client.dsmaker.api.serialization_base import BaseApiV1SerializationAdapter
from dl_api_client.dsmaker.primitives import Dataset


class AsyncHttpClientBase(abc.ABC):
    @abc.abstractmethod
    async def open(
        self,
        url: str,
        method: str,
        headers: dict,
        data: Optional[str] = None,
        content_type: Optional[str] = None,
    ) -> ClientResponse:
        raise NotImplementedError


@attr.s
class AsyncHttpApiBase(ApiBase):
    client: AsyncHttpClientBase = attr.ib()
    headers: Mapping[str, str] = attr.ib(factory=dict)

    def __attrs_post_init__(self) -> None:
        self.headers = dict(self.headers)

    async def _request(
        self,
        url: str,
        method: str,
        data: Optional[dict] = None,
        headers: Optional[dict] = None,
        lock_timeout: int = None,
    ) -> ClientResponse:
        data_str: Optional[str] = None
        content_type: Optional[str] = None
        if method not in ("get", "delete") and data is not None:
            content_type = "application/json"
            data_str = json.dumps(data)

        headers = headers or {}
        assert isinstance(headers, dict)
        headers = {**self.headers, **headers}

        started = time.monotonic()

        async def send_request() -> ClientResponse:
            return await self.client.open(
                url=url,
                method=method,
                headers=headers,
                data=data_str,
                content_type=content_type,
            )

        response = await send_request()

        while response.status_code == HTTPStatus.LOCKED and lock_timeout and time.monotonic() - started < lock_timeout:
            await asyncio.sleep(0.3)
            response = send_request()

        return response


@attr.s
class AsyncHttpApiV1Base(AsyncHttpApiBase):
    serial_adapter: BaseApiV1SerializationAdapter = attr.ib(init=False, factory=BaseApiV1SerializationAdapter)

    def dump_dataset_to_request_body(self, dataset: Dataset) -> dict:
        return self.serial_adapter.dump_dataset(dataset)

lib/dl_api_client/dl_api_client/dsmaker/api/http_async_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from http import HTTPStatus
import json
import time
from typing import (
    Mapping,
    Optional,
)

import attr

from dl_api_client.dsmaker.api.base import ApiBase
from dl_api_client.dsmaker.api.serialization_base import BaseApiV1SerializationAdapter
from dl_api_client.dsmaker.primitives import Dataset


@attr.s(frozen=True)
class ClientResponse:
    status_code: int = attr.ib(kw_only=True)
    data: bytes = attr.ib(kw_only=True)
    _json: Optional[dict] = attr.ib(kw_only=True)

    @property
    def json(self) -> dict:
        assert self._json is not None
        return self._json


class SyncHttpClientBase(abc.ABC):
    @abc.abstractmethod
    def open(
        self,
        url: str,
        method: str,
        headers: dict,
        data: Optional[str] = None,
        content_type: Optional[str] = None,
    ) -> ClientResponse:
        raise NotImplementedError

    def _strtict_dict(self, d: Optional[dict]) -> dict:
        if d is None:
            d = {}
        assert d is not None
        return d

    def post(
        self,
        url: str,
        headers: Optional[dict] = None,
        data: Optional[str] = None,
        content_type: Optional[str] = None,
    ) -> ClientResponse:
        return self.open(
            url=url,
            method="post",
            headers=self._strtict_dict(headers),
            data=data,
            content_type=content_type,
        )

    def put(
        self,
        url: str,
        headers: Optional[dict] = None,
        data: Optional[str] = None,
        content_type: Optional[str] = None,
    ) -> ClientResponse:
        return self.open(
            url=url,
            method="put",
            headers=self._strtict_dict(headers),
            data=data,
            content_type=content_type,
        )

    def get(self, url: str, headers: Optional[dict] = None) -> ClientResponse:
        return self.open(url=url, method="get", headers=self._strtict_dict(headers))

    def delete(self, url: str, headers: Optional[dict] = None) -> ClientResponse:
        return self.open(url=url, method="delete", headers=self._strtict_dict(headers))


@attr.s
class SyncHttpApiBase(ApiBase):
    client: SyncHttpClientBase = attr.ib()
    headers: Mapping[str, str] = attr.ib(factory=dict)

    def __attrs_post_init__(self) -> None:
        self.headers = dict(self.headers)

    def _request(
        self,
        url: str,
        method: str,
        data: Optional[dict] = None,
        headers: Optional[dict] = None,
        lock_timeout: int = None,
    ) -> ClientResponse:
        data_str: Optional[str] = None
        content_type: Optional[str] = None
        if method not in ("get", "delete") and data is not None:
            content_type = "application/json"
            data_str = json.dumps(data)

        headers = headers or {}
        assert isinstance(headers, dict)
        headers = {**self.headers, **headers}

        started = time.monotonic()

        def send_request() -> ClientResponse:
            return self.client.open(
                url=url,
                method=method,
                headers=headers,
                data=data_str,
                content_type=content_type,
            )

        response = send_request()

        while response.status_code == HTTPStatus.LOCKED and lock_timeout and time.monotonic() - started < lock_timeout:
            time.sleep(0.3)
            response = send_request()

        return response


@attr.s
class SyncHttpApiV1Base(SyncHttpApiBase):
    serial_adapter: BaseApiV1SerializationAdapter = attr.ib(init=False, factory=BaseApiV1SerializationAdapter)

    def dump_dataset_to_request_body(self, dataset: Dataset) -> dict:
        return self.serial_adapter.dump_dataset(dataset)

lib/dl_api_client/dl_api_client/dsmaker/api/http_sync_base.py
datalens-tech/datalens-backend
from typing import (
    ClassVar,
    Generic,
    Type,
    TypeVar,
)

from marshmallow import (
    EXCLUDE,
    Schema,
    post_load,
)


class BaseSchema(Schema):
    class Meta:
        unknown = EXCLUDE


_TARGET_OBJECT_TV = TypeVar("_TARGET_OBJECT_TV")


class DefaultSchema(BaseSchema, Generic[_TARGET_OBJECT_TV]):
    TARGET_CLS: ClassVar[Type[_TARGET_OBJECT_TV]]

    @classmethod
    def get_target_cls(cls) -> Type[_TARGET_OBJECT_TV]:
        return cls.TARGET_CLS

    def to_object(self, data: dict) -> _TARGET_OBJECT_TV:
        return self.get_target_cls()(**data)  # type: ignore

    @post_load(pass_many=False)
    def post_load(self, data, **_) -> _TARGET_OBJECT_TV:  # type: ignore
        obj = self.to_object(data)
        return obj

lib/dl_api_client/dl_api_client/dsmaker/api/schemas/base.py
datalens-tech/datalens-backend
from typing import (
    Any,
    TypeVar,
)

from marshmallow import EXCLUDE
from marshmallow import fields as ma_fields
from marshmallow_oneofschema import OneOfSchema

from dl_api_client.dsmaker.api.schemas.base import (
    BaseSchema,
    DefaultSchema,
)
from dl_api_client.dsmaker.primitives import (
    AnnotationPivotRoleSpec,
    BlockPlacement,
    DimensionPivotRoleSpec,
    LegendItem,
    OrderByRoleSpec,
    PivotHeaderInfo,
    PivotHeaderRoleSpec,
    PivotHeaderValue,
    PivotItem,
    PivotMeasureRoleSpec,
    PivotMeasureSorting,
    PivotMeasureSortingSettings,
    PivotPagination,
    PivotRoleSpec,
    RangeRoleSpec,
    RequestLegendItemRef,
    RoleSpec,
    RowRoleSpec,
    TemplateRoleSpec,
    TreeRoleSpec,
)
from dl_constants.enums import (
    FieldRole,
    FieldType,
    FieldVisibility,
    LegendItemType,
    NotificationLevel,
    OrderDirection,
    PivotHeaderRole,
    PivotRole,
    QueryBlockPlacementType,
    QueryItemRefType,
    RangeType,
    UserDataType,
    WhereClauseOperation,
)


_ROLE_SPEC_TV = TypeVar("_ROLE_SPEC_TV", bound=RoleSpec)


class DimensionValueSpecSchema(BaseSchema):
    legend_item_id = ma_fields.Integer()
    value = ma_fields.Raw()


class RoleSpecSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "role"

    class RoleSpecSchemaVariantBase(DefaultSchema[_ROLE_SPEC_TV]):
        role = ma_fields.Enum(FieldRole)

    class RoleSpecSchemaVariant(RoleSpecSchemaVariantBase[RoleSpec]):
        TARGET_CLS = RoleSpec

    class RowRoleSpecSchemaVariant(RoleSpecSchemaVariantBase[RoleSpec]):
        TARGET_CLS = RowRoleSpec

        visibility = ma_fields.Enum(FieldVisibility)

    class OrderByRoleSpecSchemaVariant(RoleSpecSchemaVariantBase[RoleSpec]):
        TARGET_CLS = OrderByRoleSpec

        direction = ma_fields.Enum(OrderDirection)

    class RangeRoleSpecSchemaVariant(RoleSpecSchemaVariantBase[RoleSpec]):
        TARGET_CLS = RangeRoleSpec

        range_type = ma_fields.Enum(RangeType)

    class TemplateRoleSpecSchemaVariant(RoleSpecSchemaVariantBase[RoleSpec]):
        TARGET_CLS = TemplateRoleSpec

        template = ma_fields.String()
        visibility = ma_fields.Enum(FieldVisibility)

    class TreeRoleSpecSchemaVariant(RoleSpecSchemaVariantBase[RoleSpec]):
        TARGET_CLS = TreeRoleSpec

        level = ma_fields.Integer()
        prefix_req = ma_fields.String(dump_only=True, attribute="prefix", data_key="prefix")  # sent as a string
        prefix_resp = ma_fields.Raw(load_only=True, attribute="prefix", data_key="prefix")  # but returned as an array
        dimension_values = ma_fields.Nested(DimensionValueSpecSchema, many=True, allow_none=True)
        visibility = ma_fields.Enum(FieldVisibility)

    type_schemas = {
        FieldRole.row.name: RowRoleSpecSchemaVariant,
        FieldRole.measure.name: RoleSpecSchemaVariant,
        FieldRole.info.name: RoleSpecSchemaVariant,
        FieldRole.order_by.name: OrderByRoleSpecSchemaVariant,
        FieldRole.distinct.name: RoleSpecSchemaVariant,
        FieldRole.range.name: RangeRoleSpecSchemaVariant,
        FieldRole.total.name: RoleSpecSchemaVariant,
        FieldRole.template.name: TemplateRoleSpecSchemaVariant,
        FieldRole.tree.name: TreeRoleSpecSchemaVariant,
        FieldRole.filter.name: RoleSpecSchemaVariant,
        FieldRole.parameter.name: RoleSpecSchemaVariant,
    }

    def get_obj_type(self, obj: Any) -> str:
        assert isinstance(obj, RoleSpec)
        return obj.role.name


class ItemRefSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    class RefSchema(BaseSchema):
        type = ma_fields.Enum(QueryItemRefType)

    class IdRefSchema(RefSchema):
        id = ma_fields.String(required=True)

    class TitleRefSchema(RefSchema):
        title = ma_fields.String(required=True)

    class EmptyRefSchema(RefSchema):
        pass

    class PlaceholderRefSchema(RefSchema):
        pass

    type_field_remove = False
    type_field = "type"

    type_schemas = {
        QueryItemRefType.id.name: IdRefSchema,
        QueryItemRefType.title.name: TitleRefSchema,
        QueryItemRefType.measure_name.name: EmptyRefSchema,
        QueryItemRefType.dimension_name.name: EmptyRefSchema,
        QueryItemRefType.placeholder.name: PlaceholderRefSchema,
    }

    def get_obj_type(self, obj: Any) -> str:
        assert isinstance(obj, RequestLegendItemRef)
        return obj.type.name


class QueryFieldsItemSchema(BaseSchema):
    legend_item_id = ma_fields.Integer(allow_none=True)
    block_id = ma_fields.Integer(allow_none=True)
    role_spec = ma_fields.Nested(RoleSpecSchema)
    ref = ma_fields.Nested(ItemRefSchema, allow_none=False)


class FilterClauseSchema(BaseSchema):
    ref = ma_fields.Nested(ItemRefSchema)
    title = ma_fields.String(allow_none=True)
    operation = ma_fields.Enum(WhereClauseOperation)
    values = ma_fields.List(ma_fields.String())
    block_id = ma_fields.Integer(dump_default=None)


class PivotPaginationSchema(DefaultSchema[PivotPagination]):
    TARGET_CLS = PivotPagination

    offset_rows = ma_fields.Integer(allow_none=True)
    limit_rows = ma_fields.Integer(allow_none=True)


class LegendItemSchema(DefaultSchema[LegendItem]):
    TARGET_CLS = LegendItem

    legend_item_id = ma_fields.Integer()
    id = ma_fields.String()
    title = ma_fields.String()
    role_spec = ma_fields.Nested(RoleSpecSchema)
    data_type = ma_fields.Enum(UserDataType)
    field_type = ma_fields.Enum(FieldType)
    item_type = ma_fields.Enum(LegendItemType)


class ResultSpecSchema(BaseSchema):
    with_totals = ma_fields.Boolean()


class PivotHeaderRoleSpecSchema(DefaultSchema[PivotHeaderRoleSpec]):
    TARGET_CLS = PivotHeaderRoleSpec

    role = ma_fields.Enum(PivotHeaderRole, required=True)


class PivotHeaderValueSchema(DefaultSchema[PivotHeaderValue]):
    TARGET_CLS = PivotHeaderValue

    value = ma_fields.String(allow_none=False, required=True)


class PivotMeasureSortingSettingsSchema(DefaultSchema[PivotMeasureSortingSettings]):
    TARGET_CLS = PivotMeasureSortingSettings

    header_values = ma_fields.List(ma_fields.Nested(PivotHeaderValueSchema), allow_none=False, required=True)
    direction = ma_fields.Enum(OrderDirection, load_default=OrderDirection.asc)
    role_spec = ma_fields.Nested(PivotHeaderRoleSpecSchema, allow_none=False, required=True)


class PivotMeasureSortingSchema(DefaultSchema[PivotMeasureSorting]):
    TARGET_CLS = PivotMeasureSorting

    column = ma_fields.Nested(PivotMeasureSortingSettingsSchema, allow_none=True, load_default=None)
    row = ma_fields.Nested(PivotMeasureSortingSettingsSchema, allow_none=True, load_default=None)


class PivotRoleSpecSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "role"

    class PivotRoleSpecSchemaVariant(DefaultSchema[PivotRoleSpec]):
        TARGET_CLS = PivotRoleSpec

        role = ma_fields.Enum(PivotRole)

    class DimensionPivotRoleSpecSchemaVariant(DefaultSchema[DimensionPivotRoleSpec]):
        TARGET_CLS = DimensionPivotRoleSpec

        role = ma_fields.Enum(PivotRole)
        direction = ma_fields.Enum(OrderDirection)

    class AnnotationPivotRoleSpecSchemaVariant(DefaultSchema[AnnotationPivotRoleSpec]):
        TARGET_CLS = AnnotationPivotRoleSpec

        role = ma_fields.Enum(PivotRole)
        annotation_type = ma_fields.String()
        target_legend_item_ids = ma_fields.List(ma_fields.Integer(), allow_none=True)

    class MeasurePivotRoleSpecSchemaVariant(DefaultSchema[PivotMeasureRoleSpec]):
        TARGET_CLS = PivotMeasureRoleSpec

        role = ma_fields.Enum(PivotRole, required=True)
        sorting = ma_fields.Nested(PivotMeasureSortingSchema, allow_none=True, load_default=None)

    type_schemas = {
        PivotRole.pivot_row.name: DimensionPivotRoleSpecSchemaVariant,
        PivotRole.pivot_column.name: DimensionPivotRoleSpecSchemaVariant,
        PivotRole.pivot_measure.name: MeasurePivotRoleSpecSchemaVariant,
        PivotRole.pivot_annotation.name: AnnotationPivotRoleSpecSchemaVariant,
        PivotRole.pivot_info.name: PivotRoleSpecSchemaVariant,
    }

    def get_obj_type(self, obj: Any) -> str:
        assert isinstance(obj, PivotRoleSpec)
        return obj.role.name


class PivotItemSchema(DefaultSchema[PivotItem]):
    TARGET_CLS = PivotItem

    pivot_item_id = ma_fields.Integer()
    legend_item_ids = ma_fields.List(ma_fields.Integer())
    role_spec = ma_fields.Nested(PivotRoleSpecSchema)
    title = ma_fields.String()


class PivotTotalItemSchema(BaseSchema):
    level = ma_fields.Integer()


class PivotTotalsSchema(BaseSchema):
    rows = ma_fields.Nested(PivotTotalItemSchema, many=True)
    columns = ma_fields.Nested(PivotTotalItemSchema, many=True)


class PivotSpecSchema(BaseSchema):
    with_totals = ma_fields.Boolean()
    structure = ma_fields.Nested(PivotItemSchema, many=True)
    totals = ma_fields.Nested(PivotTotalsSchema)


class PivotHeaderInfoSchema(DefaultSchema[PivotHeaderInfo]):
    TARGET_CLS = PivotHeaderInfo

    sorting_direction = ma_fields.Enum(OrderDirection, allow_none=True)
    role_spec = ma_fields.Nested(PivotHeaderRoleSpecSchema, allow_none=False, required=True)


class PivotDimensionWithInfoSchema(BaseSchema):
    cells = ma_fields.List(
        ma_fields.List(  # Vector of values (can be None)
            ma_fields.List(ma_fields.Raw()),
            allow_none=True,
        ),
    )
    header_info = ma_fields.Nested(PivotHeaderInfoSchema)


class PivotRowSchema(BaseSchema):
    header = ma_fields.List(  # List of dimensions in row
        ma_fields.List(  # Vector of values (can be None)
            ma_fields.List(ma_fields.Raw()),  # Individual value is a <value, legend_item_id> pair
            allow_none=True,
        ),
    )
    header_with_info = ma_fields.Nested(PivotDimensionWithInfoSchema)
    values = ma_fields.List(  # List of cells in row (per column)
        ma_fields.List(  # Vector of values (can be None)
            ma_fields.List(ma_fields.Raw()),  # Individual value is a <value, legend_item_id> pair
            allow_none=True,
        ),
    )


class PivotDataResponseSchema(BaseSchema):
    columns = ma_fields.List(  # List of columns
        ma_fields.List(  # List of dimensions in column
            ma_fields.List(  # Vector of values (can be None)
                ma_fields.List(ma_fields.Raw()),  # Individual value is a <value, legend_item_id> pair
                allow_none=True,
            ),
        ),
    )
    columns_with_info = ma_fields.List(ma_fields.Nested(PivotDimensionWithInfoSchema))
    row_dimension_headers = ma_fields.List(  # List of row dimension headers
        ma_fields.List(  # Vector of values (can be None)
            ma_fields.List(ma_fields.Raw()),  # Individual value is a <value, legend_item_id> pair
            allow_none=True,
        ),
    )
    rows = ma_fields.Nested(PivotRowSchema, many=True)


class NotificationSchema(BaseSchema):
    title = ma_fields.String()
    message = ma_fields.String()
    level = ma_fields.Enum(NotificationLevel)
    locator = ma_fields.String()


class BlockPlacementSchema(BaseSchema):
    type = ma_fields.Enum(QueryBlockPlacementType)


class RootBlockPlacementSchema(BlockPlacementSchema):
    pass


class AfterBlockPlacementSchema(BlockPlacementSchema):
    dimension_values = ma_fields.Nested(DimensionValueSpecSchema, many=True, allow_none=True)


class BlockPlacementSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = True
    type_field = "type"

    type_schemas = {
        QueryBlockPlacementType.root.name: RootBlockPlacementSchema,
        QueryBlockPlacementType.after.name: AfterBlockPlacementSchema,
    }

    def get_obj_type(self, obj: Any) -> str:
        assert isinstance(obj, BlockPlacement)
        return obj.type.name


class QueryBlockSchema(DefaultSchema):
    block_id = ma_fields.Integer(allow_none=False)
    parent_block_id = ma_fields.Integer(allow_none=True, load_default=None)
    placement = ma_fields.Nested(BlockPlacementSchema, allow_none=False)


class BaseDataApiV2ResponseSchema(BaseSchema):
    """
    Base class for Data API v2 responses.
    """

    sel_fields = ma_fields.Nested(LegendItemSchema, data_key="fields", attribute="fields", many=True)
    blocks = ma_fields.Raw()
    notifications = ma_fields.Nested(NotificationSchema, many=True)


class ResultResponseSchema(BaseDataApiV2ResponseSchema):
    """
    Generic v2 response.
    Mainly for the /result endpoint.
    """

    result_data = ma_fields.Raw()


class PivotResponseSchema(BaseDataApiV2ResponseSchema):
    """
    /pivot-specific schema
    """

    pivot_data = ma_fields.Nested(PivotDataResponseSchema)
    pivot = ma_fields.Nested(PivotSpecSchema)

lib/dl_api_client/dl_api_client/dsmaker/api/schemas/data.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Mapping,
)

from marshmallow import (
    post_dump,
    pre_load,
)
from marshmallow import EXCLUDE
from marshmallow import fields as ma_fields
from marshmallow_oneofschema import OneOfSchema

from dl_api_client.dsmaker.api.schemas.base import DefaultSchema
from dl_api_client.dsmaker.primitives import (
    ArrayFloatParameterValue,
    ArrayIntParameterValue,
    ArrayStrParameterValue,
    AvatarRelation,
    BooleanParameterValue,
    Column,
    ComponentError,
    ComponentErrorPack,
    ComponentErrorRegistry,
    Dataset,
    DataSource,
    DateParameterValue,
    DateTimeParameterValue,
    DateTimeTZParameterValue,
    DirectJoinPart,
    FloatParameterValue,
    FormulaJoinPart,
    GenericDateTimeParameterValue,
    GeoPointParameterValue,
    GeoPolygonParameterValue,
    IntegerParameterValue,
    JoinCondition,
    MarkupParameterValue,
    ObligatoryFilter,
    ParameterValue,
    ParameterValueConstraint,
    RangeParameterValueConstraint,
    ResultField,
    ResultFieldJoinPart,
    ResultSchemaAux,
    SetParameterValueConstraint,
    SourceAvatar,
    StringParameterValue,
    TreeStrParameterValue,
    UuidParameterValue,
    WhereClause,
)
from dl_constants.enums import (
    AggregationFunction,
    BinaryJoinOperator,
    CalcMode,
    ComponentErrorLevel,
    ComponentType,
    ConditionPartCalcMode,
    DataSourceType,
    FieldType,
    JoinConditionType,
    JoinType,
    ManagedBy,
    ParameterValueConstraintType,
    UserDataType,
    WhereClauseOperation,
)
from dl_model_tools.schema.dynamic_enum_field import DynamicEnumField
from dl_utils.schemas import OneOfSchemaWithDumpLoadHooks


if TYPE_CHECKING:
    from dl_api_client.dsmaker.primitives import JoinPart


class StringValueSchema(DefaultSchema[StringParameterValue]):
    TARGET_CLS = StringParameterValue
    value = ma_fields.String()


class IntegerValueSchema(DefaultSchema[IntegerParameterValue]):
    TARGET_CLS = IntegerParameterValue
    value = ma_fields.Integer()


class FloatValueSchema(DefaultSchema[FloatParameterValue]):
    TARGET_CLS = FloatParameterValue
    value = ma_fields.Float()


class DateValueSchema(DefaultSchema[DateParameterValue]):
    TARGET_CLS = DateParameterValue
    value = ma_fields.Date()


class DateTimeValueSchema(DefaultSchema[DateTimeParameterValue]):
    TARGET_CLS = DateTimeParameterValue
    value = ma_fields.DateTime()


class DateTimeTZValueSchema(DefaultSchema[DateTimeTZParameterValue]):
    TARGET_CLS = DateTimeTZParameterValue
    value = ma_fields.DateTime()


class GenericDateTimeValueSchema(DefaultSchema[GenericDateTimeParameterValue]):
    TARGET_CLS = GenericDateTimeParameterValue
    value = ma_fields.DateTime()


class BooleanValueSchema(DefaultSchema[BooleanParameterValue]):
    TARGET_CLS = BooleanParameterValue
    value = ma_fields.Boolean()


class GeoPointValueSchema(DefaultSchema[GeoPointParameterValue]):
    TARGET_CLS = GeoPointParameterValue
    value = ma_fields.List(ma_fields.Float())


class GeoPolygonValueSchema(DefaultSchema[GeoPolygonParameterValue]):
    TARGET_CLS = GeoPolygonParameterValue
    value = ma_fields.List(ma_fields.List(ma_fields.List(ma_fields.Float())))


class UuidValueSchema(DefaultSchema[UuidParameterValue]):
    TARGET_CLS = UuidParameterValue
    value = ma_fields.String()


class MarkupValueSchema(DefaultSchema[MarkupParameterValue]):
    TARGET_CLS = MarkupParameterValue
    value = ma_fields.String()


class ArrayStrValueSchema(DefaultSchema[ArrayStrParameterValue]):
    TARGET_CLS = ArrayStrParameterValue
    value = ma_fields.List(ma_fields.String())


class ArrayIntValueSchema(DefaultSchema[ArrayIntParameterValue]):
    TARGET_CLS = ArrayIntParameterValue
    value = ma_fields.List(ma_fields.Integer())


class ArrayFloatValueSchema(DefaultSchema[ArrayFloatParameterValue]):
    TARGET_CLS = ArrayFloatParameterValue
    value = ma_fields.List(ma_fields.Float())


class ValueSchema(OneOfSchemaWithDumpLoadHooks):
    CONTEXT_KEY = "bi_value_type"
    type_field = "type"
    type_schemas = {
        UserDataType.string.name: StringValueSchema,
        UserDataType.integer.name: IntegerValueSchema,
        UserDataType.float.name: FloatValueSchema,
        UserDataType.date.name: DateValueSchema,
        UserDataType.datetime.name: DateTimeValueSchema,
        UserDataType.datetimetz.name: DateTimeTZValueSchema,
        UserDataType.genericdatetime.name: GenericDateTimeValueSchema,
        UserDataType.boolean.name: BooleanValueSchema,
        UserDataType.geopoint.name: GeoPointValueSchema,
        UserDataType.geopolygon.name: GeoPolygonValueSchema,
        UserDataType.uuid.name: UuidValueSchema,
        UserDataType.markup.name: MarkupValueSchema,
        UserDataType.array_str.name: ArrayStrValueSchema,
        UserDataType.array_int.name: ArrayIntValueSchema,
        UserDataType.array_float.name: ArrayFloatValueSchema,
        UserDataType.tree_str.name: TreeStrParameterValue,
    }

    @pre_load(pass_many=False)
    def wrap_value_with_type(self, data: Any, **_: Any) -> Dict[str, Any]:
        type = getattr(self, "context", {}).get(self.CONTEXT_KEY)
        return {"type": type, "value": data}

    @post_dump(pass_many=False)
    def extract_value(self, data: Dict[str, Any], **_: Any) -> Any:
        return data["value"]

    def get_obj_type(self, obj: ParameterValue) -> str:
        return getattr(obj, self.type_field).name


class AllParameterValueConstraintSchema(DefaultSchema[ParameterValueConstraint]):
    TARGET_CLS = ParameterValueConstraint


class RangeParameterValueConstraintSchema(DefaultSchema[RangeParameterValueConstraint]):
    TARGET_CLS = RangeParameterValueConstraint

    min = ma_fields.Nested(ValueSchema, allow_none=True)
    max = ma_fields.Nested(ValueSchema, allow_none=True)


class SetParameterValueConstraintSchema(DefaultSchema[SetParameterValueConstraint]):
    TARGET_CLS = SetParameterValueConstraint

    values = ma_fields.List(ma_fields.Nested(ValueSchema))


class ParameterValueConstraintSchema(OneOfSchema):
    type_field = "type"
    type_schemas = {
        ParameterValueConstraintType.all.name: AllParameterValueConstraintSchema,
        ParameterValueConstraintType.range.name: RangeParameterValueConstraintSchema,
        ParameterValueConstraintType.set.name: SetParameterValueConstraintSchema,
    }

    def get_obj_type(self, obj: ParameterValueConstraint) -> str:
        return getattr(obj, self.type_field).name


class ResultFieldSchema(DefaultSchema[ResultField]):
    TARGET_CLS = ResultField

    title = ma_fields.String(required=True)
    source = ma_fields.String()
    guid = ma_fields.String(attribute="id")
    calc_mode = ma_fields.Enum(CalcMode, required=True)
    hidden = ma_fields.Boolean(load_default=False)
    description = ma_fields.String()
    formula = ma_fields.String(load_default="")
    initial_data_type = ma_fields.Enum(UserDataType, allow_none=True)
    cast = ma_fields.Enum(UserDataType)
    type = ma_fields.Enum(FieldType, readonly=True)
    data_type = ma_fields.Enum(UserDataType, allow_none=True)
    valid = ma_fields.Boolean(allow_none=True)
    avatar_id = ma_fields.String(allow_none=True)
    aggregation = ma_fields.Enum(AggregationFunction, load_default=AggregationFunction.none.name)
    has_auto_aggregation = ma_fields.Boolean(allow_none=True)
    lock_aggregation = ma_fields.Boolean(allow_none=True)
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user)
    default_value = ma_fields.Nested(ValueSchema, allow_none=True)
    value_constraint = ma_fields.Nested(ParameterValueConstraintSchema, allow_none=True)

    # Only locally used
    created_ = ma_fields.Boolean(load_default=True, load_only=True, attribute="created_")  # Always True

    @pre_load(pass_many=False)
    def store_type_in_context(self, data: Mapping[str, Any], **_: Any) -> Mapping[str, Any]:
        if "cast" in data:
            self.context[ValueSchema.CONTEXT_KEY] = data["cast"]
        return data


class ResultSchemaAuxSchema(DefaultSchema[ResultSchemaAux]):
    TARGET_CLS = ResultSchemaAux

    inter_dependencies = ma_fields.Dict(allow_none=False)


class WhereClauseSchema(DefaultSchema[WhereClause]):
    TARGET_CLS = WhereClause

    column = ma_fields.String()
    operation = ma_fields.Enum(WhereClauseOperation)
    values = ma_fields.List(ma_fields.String())


class ObligatoryFilterSchema(DefaultSchema[ObligatoryFilter]):
    TARGET_CLS = ObligatoryFilter

    id = ma_fields.String()
    field_guid = ma_fields.String()
    default_filters = ma_fields.List(ma_fields.Nested(WhereClauseSchema))
    managed_by = ma_fields.Enum(ManagedBy)
    valid = ma_fields.Boolean()


class ColumnSchema(DefaultSchema[Column]):
    TARGET_CLS = Column

    name = ma_fields.String()
    title = ma_fields.String()
    native_type = ma_fields.Dict(allow_none=True)
    user_type = ma_fields.Enum(UserDataType)
    description = ma_fields.String(dump_default="", allow_none=True)
    has_auto_aggregation = ma_fields.Boolean(dump_default=False, allow_none=True)
    lock_aggregation = ma_fields.Boolean(dump_default=False, allow_none=True)
    nullable = ma_fields.Boolean(dump_default=None, allow_none=True)


class DataSourceSchema(DefaultSchema[DataSource]):
    TARGET_CLS = DataSource

    id = ma_fields.String()
    title = ma_fields.String()
    connection_id = ma_fields.String(allow_none=True)
    source_type = DynamicEnumField(DataSourceType)
    raw_schema = ma_fields.Nested(ColumnSchema, allow_none=True, required=False, many=True)
    index_info_set = ma_fields.List(ma_fields.Dict, allow_none=True)
    parameters = ma_fields.Dict()
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user)
    valid = ma_fields.Boolean()

    # Only locally used
    created_ = ma_fields.Boolean(load_default=True, load_only=True, attribute="created_")  # Always True


class SourceAvatarSchema(DefaultSchema[SourceAvatar]):
    TARGET_CLS = SourceAvatar

    id = ma_fields.String(required=True)
    source_id = ma_fields.String()
    title = ma_fields.String()
    is_root = ma_fields.Boolean()
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user)
    valid = ma_fields.Boolean()

    # Only locally used
    created_ = ma_fields.Boolean(load_default=True, load_only=True, attribute="created_")  # Always True


class ConditionPartDirectSchema(DefaultSchema[DirectJoinPart]):
    TARGET_CLS = DirectJoinPart

    source = ma_fields.String(required=True)


class ConditionPartFormulaSchema(DefaultSchema[FormulaJoinPart]):
    TARGET_CLS = FormulaJoinPart

    formula = ma_fields.String(required=True)


class ConditionPartResultFieldSchema(DefaultSchema[ResultFieldJoinPart]):
    TARGET_CLS = ResultFieldJoinPart

    field_id = ma_fields.String(required=True)


class ConditionPartGenericSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "calc_mode"
    type_schemas = {
        ConditionPartCalcMode.direct.name: ConditionPartDirectSchema,
        ConditionPartCalcMode.formula.name: ConditionPartFormulaSchema,
        ConditionPartCalcMode.result_field.name: ConditionPartResultFieldSchema,
    }

    def get_obj_type(self, obj: JoinPart) -> str:
        return getattr(obj, self.type_field).name


class AvatarRelationSchema(DefaultSchema[AvatarRelation]):
    TARGET_CLS = AvatarRelation

    class JoinConditionSchema(DefaultSchema[JoinCondition]):
        TARGET_CLS = JoinCondition

        type = ma_fields.Enum(JoinConditionType, dump_default=JoinConditionType.binary, dump_only=True)
        operator = ma_fields.Enum(BinaryJoinOperator, required=True)
        left = ma_fields.Nested(ConditionPartGenericSchema, attribute="left_part", required=True)
        right = ma_fields.Nested(ConditionPartGenericSchema, attribute="right_part", required=True)

    id = ma_fields.String(required=True)
    left_avatar_id = ma_fields.String()
    right_avatar_id = ma_fields.String()
    conditions = ma_fields.Nested(JoinConditionSchema, many=True)
    join_type = ma_fields.Enum(JoinType)
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user, load_default=ManagedBy.user)

    # Only locally used
    created_ = ma_fields.Boolean(load_default=True, load_only=True, attribute="created_")  # Always True


class ComponentErrorListSchema(DefaultSchema[ComponentErrorRegistry]):
    TARGET_CLS = ComponentErrorRegistry

    class ComponentErrorPackSchema(DefaultSchema[ComponentErrorPack]):
        TARGET_CLS = ComponentErrorPack

        class ComponentErrorSchema(DefaultSchema[ComponentError]):
            TARGET_CLS = ComponentError

            message = ma_fields.String()
            level = ma_fields.Enum(ComponentErrorLevel)
            code = ma_fields.String()
            details = ma_fields.Dict()

        id = ma_fields.String()
        type = ma_fields.Enum(ComponentType)
        errors = ma_fields.List(ma_fields.Nested(ComponentErrorSchema))

    items = ma_fields.List(ma_fields.Nested(ComponentErrorPackSchema))


class DatasetContentInternalSchema(DefaultSchema[Dataset]):
    """
    A base class for schemas that need to contain the full dataset description
    """

    TARGET_CLS = Dataset

    sources = ma_fields.Nested(DataSourceSchema, many=True, required=False)
    source_avatars = ma_fields.Nested(SourceAvatarSchema, many=True, required=False)
    avatar_relations = ma_fields.Nested(AvatarRelationSchema, many=True, required=False)
    result_schema = ma_fields.Nested(ResultFieldSchema, many=True, required=False)
    result_schema_aux = ma_fields.Nested(ResultSchemaAuxSchema, allow_none=False)
    rls = ma_fields.Dict(load_default=dict, required=False)
    component_errors = ma_fields.Nested(ComponentErrorListSchema, required=False)
    obligatory_filters = ma_fields.Nested(ObligatoryFilterSchema, many=True, load_default=list)
    revision_id = ma_fields.String(allow_none=True, dump_default=None, load_default=None)

lib/dl_api_client/dl_api_client/dsmaker/api/schemas/dataset.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import singledispatchmethod
from typing import (
    List,
    Union,
)

import attr
from marshmallow import fields as ma_fields

from dl_api_client.dsmaker.api.schemas.base import DefaultSchema
from dl_api_client.dsmaker.api.schemas.dataset import (
    DatasetContentInternalSchema,
    ObligatoryFilterSchema,
    ParameterValueConstraintSchema,
    WhereClauseSchema,
)
from dl_api_client.dsmaker.primitives import (
    Action,
    ApiProxyObject,
    AvatarRelation,
    Container,
    Dataset,
    DataSource,
    DirectJoinPart,
    FormulaJoinPart,
    JoinPart,
    ObligatoryFilter,
    ResultField,
    ResultFieldJoinPart,
    SourceAvatar,
    UpdateAction,
)
from dl_constants.enums import JoinConditionType


class ObligatoryFilterUpdateSchema(DefaultSchema[ObligatoryFilter]):
    TARGET_CLS = ObligatoryFilter

    id = ma_fields.String()
    default_filters = ma_fields.List(ma_fields.Nested(WhereClauseSchema))
    valid = ma_fields.Boolean()


@attr.s
class BaseApiV1SerializationAdapter:
    # TODO: switch to marshmallow

    def dump_item(self, item: ApiProxyObject, action: Action = None) -> dict:
        """
        Dump item data for API request. Return dict

        :param item: the object to dump
        :param action: dump object data for given update action
        :return:
        """
        return self._dump_item(item, action=action)

    @singledispatchmethod
    def _dump_item(self, item: ApiProxyObject, action: Action) -> dict:
        raise TypeError(f"Unknown item type: {type(item)}")

    @_dump_item.register(DataSource)
    def dump_data_source(self, item: DataSource, action: Action) -> dict:
        if action == Action.delete:
            return dict(id=item.id)
        else:
            return dict(
                id=item.id,
                connection_id=item.connection_id,
                source_type=item.source_type.name,
                title=item.title,
                raw_schema=[
                    dict(
                        name=col.name,
                        title=col.title,
                        user_type=col.user_type.name,
                        native_type=col.native_type,
                        description=col.description,
                        nullable=col.nullable,
                        has_auto_aggregation=col.has_auto_aggregation,
                        lock_aggregation=col.lock_aggregation,
                    )
                    for col in item.raw_schema
                ]
                if item.raw_schema is not None
                else None,
                index_info_set=None if item.index_info_set is None else list(item.index_info_set),
                parameters=item.parameters,
                managed_by=item.managed_by.name,
                valid=item.valid,
            )

    @_dump_item.register(SourceAvatar)
    def dump_source_avatar(self, item: SourceAvatar, action: Action) -> dict:
        if action == Action.delete:
            return dict(id=item.id)
        else:
            return dict(
                id=item.id,
                source_id=item.source_id,
                title=item.title,
                is_root=item.is_root,
                managed_by=item.managed_by.name,
            )

    @_dump_item.register(AvatarRelation)
    def dump_avatar_relation(self, item: AvatarRelation, action: Action) -> dict:
        if action == Action.delete:
            return dict(id=item.id)
        else:

            def dump_condition_part(part: JoinPart) -> dict:
                data = dict(calc_mode=part.calc_mode.name)
                if isinstance(part, DirectJoinPart):
                    data["source"] = part.source
                elif isinstance(part, FormulaJoinPart):
                    data["formula"] = part.formula
                elif isinstance(part, ResultFieldJoinPart):
                    data["field_id"] = part.field_id
                return data

            return dict(
                id=item.id,
                left_avatar_id=item.left_avatar_id,
                right_avatar_id=item.right_avatar_id,
                conditions=[
                    dict(
                        type=JoinConditionType.binary.name,
                        left=dump_condition_part(condition.left_part),
                        right=dump_condition_part(condition.right_part),
                        operator=condition.operator.name,
                    )
                    for condition in item.conditions
                ],
                join_type=item.join_type.name,
                managed_by=item.managed_by.name,
            )

    @_dump_item.register(ResultField)
    def dump_field(self, item: ResultField, action: Action) -> dict:
        if action == Action.delete:
            return dict(guid=item.id)
        if action == Action.add:
            return dict(
                guid=item.id,
                title=item.title,
                calc_mode=item.calc_mode.name if item.calc_mode is not None else None,
                aggregation=item.aggregation.name if item.aggregation is not None else None,
                type=item.type.name if item.type is not None else None,
                source=item.source,
                hidden=item.hidden,
                description=item.description,
                formula=item.formula,
                initial_data_type=item.initial_data_type.name if item.initial_data_type is not None else None,
                cast=item.cast.name if item.cast is not None else None,
                data_type=item.data_type.name if item.data_type is not None else None,
                valid=item.valid,
                has_auto_aggregation=item.has_auto_aggregation,
                lock_aggregation=item.lock_aggregation,
                avatar_id=item.avatar_id,
                managed_by=item.managed_by.name,
                default_value=item.default_value.value if item.default_value is not None else None,
                value_constraint=ParameterValueConstraintSchema().dump(item.value_constraint)
                if item.value_constraint is not None
                else None,
            )
        else:
            return dict(guid=item.id)

    @_dump_item.register(ObligatoryFilter)
    def dump_obligatory_filter(self, item: ObligatoryFilter, action: Action = None) -> dict:
        if action == Action.delete:
            return dict(id=item.id)
        if action == Action.add:
            return ObligatoryFilterSchema().dump(item)
        else:
            return ObligatoryFilterUpdateSchema().dump(item)

    def _strip_implicit_updates_from_dataset(self, dataset: Dataset) -> Dataset:
        return dataset.clone(
            sources=Container({name: item for name, item in dataset.sources.items() if item.created_}),
            source_avatars=Container({name: item for name, item in dataset.source_avatars.items() if item.created_}),
            avatar_relations=Container(
                {name: item for name, item in dataset.avatar_relations.items() if item.created_}
            ),
            result_schema=Container({name: item for name, item in dataset.result_schema.items() if item.created_}),
        )

    def dump_dataset(self, item: Dataset) -> dict:
        stripped_dataset = self._strip_implicit_updates_from_dataset(dataset=item)
        dataset_data = DatasetContentInternalSchema().dump(stripped_dataset)
        return {"dataset": dataset_data}

    def generate_implicit_updates(self, dataset: Dataset) -> List[UpdateAction]:
        updates = []
        for dsrc in dataset.sources:
            if not dsrc.created_:
                dsrc.prepare()
                updates.append(dsrc.add())
        for avatar in dataset.source_avatars:
            if not avatar.created_:
                avatar.prepare()
                updates.append(avatar.add())
        for relation in dataset.avatar_relations:
            if not relation.created_:
                relation.prepare()
                updates.append(relation.add())
        for field in dataset.result_schema:
            if not field.created_:
                field.prepare()
                updates.append(field.add())
        return updates

    @staticmethod
    def _get_action_postfix(item: ApiProxyObject) -> str:
        return {
            DataSource: "source",
            SourceAvatar: "source_avatar",
            AvatarRelation: "avatar_relation",
            ResultField: "field",
            ObligatoryFilter: "obligatory_filter",
        }[type(item)]

    def dump_updates(self, updates: List[Union[UpdateAction, dict]] = None) -> List[dict]:
        result = []
        for update in updates or ():
            if isinstance(update, UpdateAction):
                action_postfix = self._get_action_postfix(update.item)
                update.item.prepare()
                item_data = dict(self.dump_item(update.item, action=update.action), **(update.custom_data or {}))
                item_data = {k: v for k, v in item_data.items() if v is not None}
                result.append(
                    {
                        "action": "{}_{}".format(update.action.name, action_postfix),
                        "order_index": update.order_index,
                        action_postfix: item_data,
                    }
                )
            elif isinstance(update, dict):
                # as raw dict; pass it on directly
                result.append(update)
            else:
                raise TypeError(type(update))

        return result

lib/dl_api_client/dl_api_client/dsmaker/api/serialization_base.py
datalens-tech/datalens-backend
from typing import Iterator

import attr

from dl_api_client.dsmaker.primitives import LegendItem


@attr.s(frozen=True)
class FieldLegend:
    fields: tuple[LegendItem, ...] = attr.ib(converter=tuple)
    _liid_to_item_map: dict[int, LegendItem] = attr.ib(init=False)

    @_liid_to_item_map.default
    def _make_liid_to_item_map(self) -> dict[int, LegendItem]:
        return {field.legend_item_id: field for field in self.fields}

    def get_item_by_legend_item_id(self, legend_item_id: int) -> LegendItem:
        return self._liid_to_item_map[legend_item_id]

    def get_title_by_legend_item_id(self, legend_item_id: int) -> str:
        return self._liid_to_item_map[legend_item_id].title

    def __iter__(self) -> Iterator[LegendItem]:
        return iter(self.fields)

lib/dl_api_client/dl_api_client/dsmaker/data_abstraction/legend.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from itertools import chain
from typing import (
    AbstractSet,
    Callable,
    Iterable,
)

import attr

from dl_api_client.dsmaker.data_abstraction.primitives import (
    DataCellTuple,
    DataItem,
    DataItemTag,
)


class DataCellMapper1D(abc.ABC):
    """
    Mapping-like interface for data where a single `DataCellTuple`
    defines a single measure cell.
    """

    @abc.abstractmethod
    def items(self) -> Iterable[tuple[DataCellTuple, DataItem]]:
        raise NotImplementedError

    def as_dict(self) -> dict[DataCellTuple, DataItem]:
        return {key: value for key, value in self.items()}

    @abc.abstractmethod
    def get_value_count(self) -> int:
        raise NotImplementedError

    def apply_tagger(
        self,
        tagger: Callable[[DataCellTuple, DataItem], AbstractSet[DataItemTag]],
    ) -> DataCellMapper1D:
        return TaggerProxyDataCellMapper1D(nested=self, tagger=tagger)

    def apply_filter(
        self,
        require_tag_combos: AbstractSet[frozenset[DataItemTag]] = frozenset(),
        exclude_tag_combos: AbstractSet[frozenset[DataItemTag]] = frozenset(),
    ) -> DataCellMapper1D:
        return FilterProxyDataCellMapper1D(
            nested=self,
            require_tag_combos=require_tag_combos,
            exclude_tag_combos=exclude_tag_combos,
        )


@attr.s(frozen=True)
class SimpleDataCellMapper1D(DataCellMapper1D):
    """
    A simple dict-based implementation of `DataCellMapper1D`
    """

    _cells: dict[DataCellTuple, DataItem] = attr.ib()

    def items(self) -> Iterable[tuple[DataCellTuple, DataItem]]:
        return ((key, value) for key, value in self._cells.items())

    def as_dict(self) -> dict[DataCellTuple, DataItem]:
        return self._cells

    def get_value_count(self) -> int:
        return len(self._cells)


@attr.s(frozen=True)
class TaggerProxyDataCellMapper1D(DataCellMapper1D):
    """
    A proxy mapper that adds tags to items of the nested mapper.
    """

    nested: DataCellMapper1D = attr.ib(kw_only=True)
    tagger: Callable[[DataCellTuple, DataItem], AbstractSet[DataItemTag]] = attr.ib(kw_only=True)

    def items(self) -> Iterable[tuple[DataCellTuple, DataItem]]:
        for dims, item in self.nested.items():
            new_tags = self.tagger(dims, item)
            if new_tags != item.meta.tags:
                item = item.clone(meta=item.meta.clone(tags=frozenset(new_tags)))
            item = item
            yield dims, item

    def get_value_count(self) -> int:
        return self.nested.get_value_count()


@attr.s(frozen=True)
class FilterProxyDataCellMapper1D(DataCellMapper1D):
    """
    A proxy mapper that filters the contents of a nested mapper.
    """

    nested: DataCellMapper1D = attr.ib(kw_only=True)
    require_tag_combos: AbstractSet[frozenset[DataItemTag]] = attr.ib(kw_only=True, factory=frozenset)
    exclude_tag_combos: AbstractSet[frozenset[DataItemTag]] = attr.ib(kw_only=True, factory=frozenset)

    def __attrs_post_init__(self) -> None:
        assert not self.require_tag_combos and self.exclude_tag_combos, "Cannot specify both requirements and excludes"

    def items(self) -> Iterable[tuple[DataCellTuple, DataItem]]:
        for dims, item in self.nested.items():
            # if there are requirements, apply them
            if self.require_tag_combos and not any(
                item.meta.tags.issuperset(combo) for combo in self.require_tag_combos
            ):
                continue

            # if there are excludes, apply them
            if any(item.meta.tags.issuperset(combo) for combo in self.exclude_tag_combos):
                continue

            yield dims, item

    def get_value_count(self) -> int:
        return len(list(self.items()))


class DataCellMapper2D(abc.ABC):
    """
    Mapping-like interface for data where a pair of `DataCellTuple` instances
    defines a single measure cell.
    Can be converted to a `DataCellMapper1D` via the `as_1d_mapper` method.
    """

    @abc.abstractmethod
    def items(self) -> Iterable[tuple[tuple[DataCellTuple, DataCellTuple], DataItem]]:
        raise NotImplementedError

    def as_dict(self) -> dict[tuple[DataCellTuple, DataCellTuple], DataItem]:
        return {key: value for key, value in self.items()}

    def _make_1d_coords(self, orig_dim_coords: tuple[DataCellTuple, DataCellTuple]) -> DataCellTuple:
        """
        Turn a pair of `DataCellTuple`s into a single `DataCellTuple`.
        Sort the data cells by cell title.
        """
        return DataCellTuple(
            tuple(
                sorted(
                    (dim_cell for dim_cell in chain(orig_dim_coords[0].cells, orig_dim_coords[1].cells)),
                    key=lambda _cell: _cell.title,
                )
            )
        )

    def as_1d_mapper(self) -> DataCellMapper1D:
        """
        Convert into a `DataCellMapper1D` by flattening the dimension coord pair into a single `DataCellTuple`
        """

        cells_1d = {self._make_1d_coords(orig_dim_coords): value for orig_dim_coords, value in self.items()}
        return SimpleDataCellMapper1D(cells_1d)

lib/dl_api_client/dl_api_client/dsmaker/data_abstraction/mapping_base.py
datalens-tech/datalens-backend
import abc
from itertools import chain

from dl_api_client.dsmaker.data_abstraction.mapping_base import DataCellMapper1D


class Mapper1DComparator(abc.ABC):
    @abc.abstractmethod
    def assert_equal(self, left: DataCellMapper1D, right: DataCellMapper1D) -> None:
        raise NotImplementedError


class SimpleMapper1DComparator(Mapper1DComparator):
    def assert_equal(self, left: DataCellMapper1D, right: DataCellMapper1D) -> None:
        left_dict = left.as_dict()
        right_dict = right.as_dict()

        for key in set(chain(left_dict.keys(), right_dict.keys())):
            left_item = left_dict.get(key)
            right_item = right_dict.get(key)
            left_cell = left_item.cell if left_item is not None else None
            right_cell = right_item.cell if right_item is not None else None
            assert left_cell == right_cell, f"Value mismatch for key {key}: {left_cell} vs. {right_cell}"

lib/dl_api_client/dl_api_client/dsmaker/data_abstraction/mapping_comparator.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from itertools import chain
from typing import (
    Any,
    Generator,
    Iterable,
)

import attr

from dl_api_client.dsmaker.api.data_api import HttpDataApiResponse
from dl_api_client.dsmaker.data_abstraction.legend import FieldLegend
from dl_api_client.dsmaker.data_abstraction.mapping_base import (
    DataCellMapper1D,
    SimpleDataCellMapper1D,
)
from dl_api_client.dsmaker.data_abstraction.primitives import (
    DataCell,
    DataCellTuple,
    DataItem,
    DataItemMeta,
    DataItemTag,
)
from dl_constants.enums import FieldRole
from dl_constants.internal_constants import MEASURE_NAME_TITLE


@attr.s(frozen=True)
class PivotDataSliceProxy(abc.ABC):
    resp_data: dict = attr.ib(kw_only=True)

    @abc.abstractmethod
    def get_compound_header(self) -> tuple[Any, ...]:
        raise ModuleNotFoundError

    @abc.abstractmethod
    def get_flat_header(self) -> Any:
        raise ModuleNotFoundError


@attr.s(frozen=True)
class PivotDataRowProxy(PivotDataSliceProxy):
    row_idx: int = attr.ib(kw_only=True)

    def get_compound_header(self) -> tuple[Any, ...]:
        row_data = self.resp_data["pivot_data"]["rows"][self.row_idx]["header"]
        return tuple(raw_vect[0][0] for raw_vect in row_data)

    def get_flat_header(self) -> Any:
        row_data = self.resp_data["pivot_data"]["rows"][self.row_idx]["header"]
        assert len(row_data) == 1
        return row_data[0][0][0]


@attr.s(frozen=True)
class PivotDataColumnProxy(PivotDataSliceProxy):
    col_idx: int = attr.ib(kw_only=True)

    def get_compound_header(self) -> tuple[Any, ...]:
        col_data = self.resp_data["pivot_data"]["columns"][self.col_idx]
        return tuple(raw_vect[0][0] for raw_vect in col_data)

    def get_flat_header(self) -> Any:
        col_data = self.resp_data["pivot_data"]["columns"][self.col_idx]
        assert len(col_data) == 1
        return col_data[0][0][0]


@attr.s
class PivotDataAbstraction:
    """
    A data abstraction interface for pivot table responses.
    """

    resp_data: dict = attr.ib(kw_only=True)

    def get_field_legend(self) -> FieldLegend:
        return FieldLegend(self.resp_data["fields"])

    def _iter_all_cells(self) -> Generator[tuple[DataCellTuple, DataItem], None, None]:
        """
        Iterate over all (row x column) dimension value combinations
        and yield `dimensions: measure` pairs.
        Key points:
        - `Measure Name` should be present as a dimension in each such pair.
        - annotation measures are yielded as well as the "main" measure values,
          and annotations have their original `Measure Name` values
        """

        legend = self.get_field_legend()
        columns = self.resp_data["pivot_data"]["columns"]

        def make_cell_from_raw_cell(raw_cell: list[Any]) -> DataCell:
            """Convert raw value pair into a `DataCell` object."""
            return DataCell(
                value=raw_cell[0],
                title=legend.get_title_by_legend_item_id(raw_cell[1]),
            )

        def remove_measure_name_cells(cells: Iterable[DataCell]) -> tuple[DataCell, ...]:
            """Strip dimension sequence of `Measure Name` cells."""
            return tuple(cell for cell in cells if cell.title != MEASURE_NAME_TITLE)

        # Since `Measure Name` might be present once, twice or not at all,
        # we will first remove it, and the add it again manually
        # to normalized the `dimensions -> measure` mapping

        rows = self.resp_data["pivot_data"]["rows"]
        column_value_sets = [
            remove_measure_name_cells(make_cell_from_raw_cell(dim_cell_vector[0]) for dim_cell_vector in column)
            for column in columns
        ]

        for row in rows:
            row_value_set = set(
                remove_measure_name_cells(
                    make_cell_from_raw_cell(dim_cell_vector[0]) for dim_cell_vector in row["header"]
                )
            )

            for col_idx, col_value_set in enumerate(column_value_sets):
                measure_raw_cell_vector = row["values"][col_idx]
                if measure_raw_cell_vector is None:
                    # no data to yield here, so skip
                    continue

                # Iterate over "layers":
                # - 0th is the main value
                # - 1st (is present) is the annotation
                for measure_raw_cell in measure_raw_cell_vector:
                    measure_cell = make_cell_from_raw_cell(measure_raw_cell)

                    measure_liid = measure_raw_cell[1]
                    tags: set[DataItemTag] = set()
                    role = legend.get_item_by_legend_item_id(measure_liid).role_spec.role
                    if role == FieldRole.total:
                        tags.add(DataItemTag.total)

                    # Create measure name cell and add it to dimensions
                    mname_cell = DataCell(value=measure_cell.title, title=MEASURE_NAME_TITLE)

                    dim_tuple = DataCellTuple(
                        cells=tuple(
                            sorted(chain(row_value_set, col_value_set, [mname_cell]), key=lambda cell: cell.title)
                        ),
                    )

                    yield dim_tuple, DataItem(cell=measure_cell, meta=DataItemMeta(tags=frozenset(tags)))

    def get_1d_mapper(self) -> DataCellMapper1D:
        return SimpleDataCellMapper1D(
            cells={dim_tuple: cell for dim_tuple, cell in self._iter_all_cells()},
        )

    def get_row(self, idx: int) -> PivotDataRowProxy:
        return PivotDataRowProxy(resp_data=self.resp_data, row_idx=idx)

    def get_column(self, idx: int) -> PivotDataColumnProxy:
        return PivotDataColumnProxy(resp_data=self.resp_data, col_idx=idx)

    def iter_rows(self) -> Generator[PivotDataRowProxy, None, None]:
        for row_idx in range(len(self.resp_data["pivot_data"]["rows"])):
            yield self.get_row(row_idx)

    def iter_columns(self) -> Generator[PivotDataColumnProxy, None, None]:
        for col_idx in range(len(self.resp_data["pivot_data"]["columns"])):
            yield self.get_column(col_idx)

    def get_flat_row_headers(self) -> list[Any]:
        return [row.get_flat_header() for row in self.iter_rows()]

    def get_flat_column_headers(self) -> list[Any]:
        return [col.get_flat_header() for col in self.iter_columns()]

    def get_compound_row_headers(self) -> list[tuple[Any, ...]]:
        return [row.get_compound_header() for row in self.iter_rows()]

    def get_compound_column_headers(self) -> list[tuple[Any, ...]]:
        return [col.get_compound_header() for col in self.iter_columns()]

    @classmethod
    def from_response(cls, response: HttpDataApiResponse) -> PivotDataAbstraction:
        assert response.api_version == "v2"
        return cls(resp_data=response.data)

lib/dl_api_client/dl_api_client/dsmaker/data_abstraction/pivot.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import Enum
from typing import Any

import attr


@attr.s(frozen=True)
class DataCell:
    value: Any = attr.ib(kw_only=True)
    title: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class DataCellTuple:
    cells: tuple[DataCell, ...] = attr.ib()

    @property
    def titles(self) -> frozenset[str]:
        return frozenset(cell.title for cell in self.cells)


class DataItemTag(Enum):
    total = "total"
    annotation = "annotation"


@attr.s(frozen=True)
class DataItemMeta:
    tags: frozenset[DataItemTag] = attr.ib(kw_only=True, factory=frozenset)

    def clone(self, **kwargs: Any) -> DataItemMeta:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True)
class DataItem:
    cell: DataCell = attr.ib(kw_only=True)
    meta: DataItemMeta = attr.ib(kw_only=True, factory=DataItemMeta)

    def clone(self, **kwargs: Any) -> DataItem:
        return attr.evolve(self, **kwargs)

lib/dl_api_client/dl_api_client/dsmaker/data_abstraction/primitives.py
datalens-tech/datalens-backend
from __future__ import annotations

from itertools import chain
from typing import (
    AbstractSet,
    Any,
    Generator,
    Optional,
)

import attr

from dl_api_client.dsmaker.api.data_api import HttpDataApiResponse
from dl_api_client.dsmaker.data_abstraction.legend import FieldLegend
from dl_api_client.dsmaker.data_abstraction.mapping_base import (
    DataCellMapper1D,
    SimpleDataCellMapper1D,
)
from dl_api_client.dsmaker.data_abstraction.primitives import (
    DataCell,
    DataCellTuple,
    DataItem,
    DataItemMeta,
    DataItemTag,
)
from dl_constants.enums import (
    FieldRole,
    FieldType,
)
from dl_constants.internal_constants import MEASURE_NAME_TITLE


@attr.s(frozen=True)
class ResultRawDataRow:
    _raw_row_data: dict[str, Any] = attr.ib()

    @property
    def data(self) -> list[Any]:
        return self._raw_row_data["data"]

    @property
    def legend(self) -> list[int]:
        return self._raw_row_data["legend"]


@attr.s(frozen=True)
class ResultRowSplitter:
    _field_legend: FieldLegend = attr.ib(kw_only=True)
    _dimension_liids: AbstractSet[int] = attr.ib(kw_only=True)
    _measure_liids: AbstractSet[int] = attr.ib(kw_only=True)

    @_measure_liids.default
    def _make_measure_liids(self) -> set[int]:
        return {field.legend_item_id for field in self._field_legend.fields if field.field_type == FieldType.MEASURE}

    @_dimension_liids.default
    def _make_dimension_liids(self) -> set[int]:
        return {field.legend_item_id for field in self._field_legend.fields if field.field_type == FieldType.DIMENSION}

    def _make_and_filter_cells(
        self,
        raw_row: ResultRawDataRow,
        legend_item_ids: AbstractSet[int],
    ) -> Generator[tuple[DataCell, int], None, None]:
        for liid, value in zip(raw_row.legend, raw_row.data):
            if liid in legend_item_ids:
                yield DataCell(value=value, title=self._field_legend.get_title_by_legend_item_id(liid)), liid

    def split_row_by_measures(self, raw_row: ResultRawDataRow) -> Generator[tuple[DataCellTuple, DataCell], None, None]:
        """
        Split raw data row into an iterable (generator) of pairs:
        1. a `DataCellTuple` instance containing all of the dimension values, including `Measure Name`,
        2. a `DataCell` containing the measure value.
        """

        if not self._measure_liids:
            return

        dimension_cells = tuple(cell for cell, liid in self._make_and_filter_cells(raw_row, self._dimension_liids))
        measure_cells_and_liids = self._make_and_filter_cells(raw_row, self._measure_liids)

        for measure_cell, measure_liid in measure_cells_and_liids:
            tags: set[DataItemTag] = set()
            role = self._field_legend.get_item_by_legend_item_id(measure_liid).role_spec.role
            if role == FieldRole.total:
                tags.add(DataItemTag.total)

            mname_cell = DataCell(value=measure_cell.title, title=MEASURE_NAME_TITLE)
            dim_tuple = DataCellTuple(
                tuple(sorted(chain(dimension_cells + (mname_cell,)), key=lambda cell: cell.title))
            )
            yield dim_tuple, DataItem(cell=measure_cell, meta=DataItemMeta(tags=frozenset(tags)))


@attr.s
class ResultDataAbstraction:
    """
    A data abstraction interface for regular data API v2 responses.
    """

    resp_data: dict = attr.ib(kw_only=True)

    def get_field_legend(self) -> FieldLegend:
        return FieldLegend(self.resp_data["fields"])

    @property
    def _raw_rows(self) -> list[dict]:
        assert len(self.resp_data["result_data"]) == 1, "Only responses with one data array are supported"
        return self.resp_data["result_data"][0]["rows"]

    def iter_raw_rows(self) -> Generator[ResultRawDataRow, None, None]:
        for raw_row_data in self._raw_rows:
            yield ResultRawDataRow(raw_row_data)

    def _iter_split_rows_by_measure(
        self,
        dimension_liids: Optional[frozenset[int]] = None,
        measure_liids: Optional[frozenset[int]] = None,
    ) -> Generator[tuple[DataCellTuple, DataCell], None, None]:
        liid_sets = {}
        if dimension_liids is not None:
            liid_sets["dimension_liids"] = dimension_liids
        if measure_liids is not None:
            liid_sets["measure_liids"] = measure_liids

        splitter = ResultRowSplitter(field_legend=self.get_field_legend(), **liid_sets)
        for raw_row in self.iter_raw_rows():
            yield from splitter.split_row_by_measures(raw_row)

    def get_1d_mapper(
        self,
        dimension_liids: Optional[frozenset[int]] = None,
        measure_liids: Optional[frozenset[int]] = None,
    ) -> DataCellMapper1D:
        return SimpleDataCellMapper1D(
            cells={
                dim_tuple: cell
                for dim_tuple, cell in self._iter_split_rows_by_measure(
                    dimension_liids=dimension_liids,
                    measure_liids=measure_liids,
                )
            },
        )

    @classmethod
    def from_response(cls, response: HttpDataApiResponse) -> ResultDataAbstraction:
        assert response.api_version == "v2"
        return cls(resp_data=response.data)

lib/dl_api_client/dl_api_client/dsmaker/data_abstraction/result.py
datalens-tech/datalens-backend
from collections import defaultdict
from http import HTTPStatus
from itertools import chain
from typing import (
    AbstractSet,
    Optional,
    Tuple,
)

from dl_api_client.dsmaker.api.data_api import (
    HttpDataApiResponse,
    SyncHttpDataApiV2,
)
from dl_api_client.dsmaker.data_abstraction.mapping_base import DataCellMapper1D
from dl_api_client.dsmaker.data_abstraction.mapping_comparator import SimpleMapper1DComparator
from dl_api_client.dsmaker.data_abstraction.pivot import PivotDataAbstraction
from dl_api_client.dsmaker.data_abstraction.primitives import (
    DataCellTuple,
    DataItem,
    DataItemTag,
)
from dl_api_client.dsmaker.data_abstraction.result import ResultDataAbstraction
from dl_api_client.dsmaker.primitives import (
    Dataset,
    OrderedField,
    PivotMeasureSorting,
    PivotPagination,
    PivotTotals,
    RequestLegendItem,
    RequestPivotItem,
    WhereClause,
)
from dl_constants.enums import (
    FieldRole,
    FieldType,
    NotificationType,
    OrderDirection,
    PivotRole,
    QueryItemRefType,
)
from dl_constants.internal_constants import (
    DIMENSION_NAME_TITLE,
    MEASURE_NAME_TITLE,
)


SPECIAL_TITLES = {
    MEASURE_NAME_TITLE,
    DIMENSION_NAME_TITLE,
}


def _get_regular_legend_for_pivot(
    dataset: Dataset,
    columns: list[str],
    rows: list[str],
    measures: list[str],
    annotations: Optional[list[str]] = None,
    totals: Optional[list[tuple[str, ...]]] = None,
) -> tuple[list[RequestLegendItem], dict[str, list[int]]]:
    totals = totals or []
    assert totals is not None

    # Make regular legend
    main_block_id = 0
    legend_item_id = 0
    legend: list[RequestLegendItem] = []
    liid_map: dict[str, list[int]] = defaultdict(list)
    all_titles = chain(rows, columns, measures, annotations or ())
    item: RequestLegendItem
    for title in all_titles:
        assert title != DIMENSION_NAME_TITLE
        if liid_map[title]:  # Don't duplicate fields
            continue
        if title == MEASURE_NAME_TITLE:
            item = dataset.measure_name_as_req_legend_item(
                role=FieldRole.row, legend_item_id=legend_item_id, block_id=main_block_id
            )
        else:
            item = dataset.find_field(title=title).as_req_legend_item(
                role=FieldRole.row, legend_item_id=legend_item_id, block_id=main_block_id
            )
        legend.append(item)
        liid_map[title].append(legend_item_id)
        legend_item_id += 1

        if totals and title != MEASURE_NAME_TITLE:
            field = dataset.find_field(title=title)
            for totals_idx, totals_set in enumerate(totals):
                block_id = main_block_id + 1 + totals_idx  # the first block for totals should be 1
                if field.type == FieldType.DIMENSION:
                    if title in totals_set:
                        item = field.as_req_legend_item(legend_item_id=legend_item_id, block_id=block_id)
                    else:
                        item = dataset.placeholder_as_req_legend_item(
                            legend_item_id=legend_item_id,
                            block_id=block_id,
                            role=FieldRole.template,
                            template="",
                        )
                else:
                    item = field.as_req_legend_item(
                        legend_item_id=legend_item_id,
                        block_id=block_id,
                        role=FieldRole.total,
                    )

                legend.append(item)
                liid_map[title].append(legend_item_id)
                legend_item_id += 1

    return legend, liid_map


def get_pivot_response(
    dataset: Dataset,
    data_api: SyncHttpDataApiV2,
    columns: list[str],
    rows: list[str],
    measures: list[str],
    measures_sorting_settings: Optional[list[Optional[PivotMeasureSorting]]] = None,
    annotations: Optional[list[str]] = None,
    title_mapping: Optional[dict[str, str]] = None,
    totals: Optional[list[tuple[str, ...]]] = None,
    simple_totals: Optional[PivotTotals] = None,
    with_totals: Optional[bool] = None,
    order_fields: Optional[dict[str, Optional[OrderDirection]]] = None,
    filters: Optional[list[WhereClause]] = None,
    pivot_pagination: Optional[PivotPagination] = None,
    autofill_legend: bool = False,
) -> HttpDataApiResponse:
    title_mapping = title_mapping or {}
    assert title_mapping is not None

    order_fields = order_fields or {}
    assert order_fields is not None

    totals = totals or []
    assert totals is not None

    if measures_sorting_settings is None:
        measures_sorting_settings = [None] * len(measures)
    assert len(measures_sorting_settings) == len(measures)

    # Make regular legend
    legend, liid_map = _get_regular_legend_for_pivot(
        dataset=dataset,
        columns=columns,
        rows=rows,
        measures=measures,
        annotations=annotations,
        totals=totals,
    )

    def make_pivot_item(
        title: str, role: PivotRole, measure_sorting_settings: Optional[PivotMeasureSorting] = None
    ) -> RequestPivotItem:
        legend_item_ids = liid_map[title]
        remapped_title = title_mapping.get(title, title)
        if title == MEASURE_NAME_TITLE:
            return dataset.make_req_pivot_item(role=role, legend_item_ids=legend_item_ids, title=remapped_title)
        if role == PivotRole.pivot_annotation:
            return dataset.make_req_pivot_item(
                role=role, legend_item_ids=legend_item_ids, title=remapped_title, annotation_type="color"
            )
        if role in (PivotRole.pivot_row, PivotRole.pivot_column):
            return dataset.make_req_pivot_item(
                role=role,
                legend_item_ids=legend_item_ids,
                title=remapped_title,
                direction=order_fields.get(title, OrderDirection.asc),
            )
        if role == PivotRole.pivot_measure:
            return dataset.make_req_pivot_item(
                role=role,
                legend_item_ids=legend_item_ids,
                title=remapped_title,
                measure_sorting_settings=measure_sorting_settings,
            )
        return dataset.make_req_pivot_item(role=role, legend_item_ids=legend_item_ids, title=remapped_title)

    def order_by_item(title: str, direction: OrderDirection) -> OrderedField:
        field = dataset.find_field(title=title)
        if direction == OrderDirection.desc:
            return field.desc
        return field.asc

    pivot_resp = data_api.get_pivot(
        dataset=dataset,
        fields=legend,
        pivot_structure=[
            *[make_pivot_item(title=title, role=PivotRole.pivot_column) for title in columns],
            *[make_pivot_item(title=title, role=PivotRole.pivot_row) for title in rows],
            *[
                make_pivot_item(
                    title=title, role=PivotRole.pivot_measure, measure_sorting_settings=measure_sorting_settings
                )
                for title, measure_sorting_settings in zip(measures, measures_sorting_settings)
            ],
            *[make_pivot_item(title=title, role=PivotRole.pivot_annotation) for title in annotations or ()],
        ],
        order_by=[order_by_item(title=title, direction=direction) for title, direction in (order_fields or {}).items()],
        filters=filters,
        pivot_pagination=pivot_pagination,
        pivot_totals=simple_totals,
        with_totals=with_totals,
        autofill_legend=autofill_legend,
        fail_ok=True,
    )
    return pivot_resp


def check_pivot_response(
    dataset: Dataset,
    data_api: SyncHttpDataApiV2,
    columns: list[str],
    rows: list[str],
    measures: list[str],
    measures_sorting_settings: Optional[list[Optional[PivotMeasureSorting]]] = None,
    annotations: Optional[list[str]] = None,
    title_mapping: Optional[dict[str, str]] = None,
    totals: Optional[list[tuple[str, ...]]] = None,
    simple_totals: Optional[PivotTotals] = None,
    with_totals: Optional[bool] = None,
    order_fields: Optional[dict[str, Optional[OrderDirection]]] = None,
    check_totals: Optional[list[tuple[str, ...]]] = None,
    filters: Optional[list[WhereClause]] = None,
    autofill_legend: bool = False,
    custom_pivot_legend_check: Optional[list[Tuple[str, PivotRole]]] = None,
    min_col_cnt: Optional[int] = None,
    max_col_cnt: Optional[int] = None,
    min_row_cnt: Optional[int] = None,
    max_row_cnt: Optional[int] = None,
    min_value_cnt: Optional[int] = None,
    max_value_cnt: Optional[int] = None,
    expected_notifications: Optional[list[NotificationType]] = None,
) -> PivotDataAbstraction:
    title_mapping = title_mapping or {}
    assert title_mapping is not None

    def norm_title(title: str) -> str:
        return title_mapping.get(title, title)

    assert simple_totals is None or totals is None
    if check_totals is None:
        if with_totals:
            check_totals = [
                *([tuple(columns)] if rows else []),
                *([tuple(rows)] if columns else []),
                *([()] if rows and columns else []),
            ]
        if totals is not None:
            check_totals = totals
        if simple_totals is not None:
            if not measures:
                check_totals = None
            else:
                check_totals = (
                    (
                        [tuple(rows[: row_tot_item.level]) + tuple(columns) for row_tot_item in simple_totals.rows]
                        if rows
                        else []
                    )
                    + (
                        [tuple(columns[: col_tot_item.level]) + tuple(rows) for col_tot_item in simple_totals.columns]
                        if columns
                        else []
                    )
                    + (
                        [
                            tuple(rows[: row_tot_item.level]) + tuple(columns[: col_tot_item.level])
                            for row_tot_item in simple_totals.rows
                            for col_tot_item in simple_totals.columns
                        ]
                        if rows and columns
                        else []
                    )
                )

    pivot_resp = get_pivot_response(
        dataset=dataset,
        data_api=data_api,
        columns=columns,
        rows=rows,
        measures=measures,
        measures_sorting_settings=measures_sorting_settings,
        annotations=annotations,
        title_mapping=title_mapping,
        totals=totals,
        simple_totals=simple_totals,
        with_totals=with_totals,
        order_fields=order_fields,
        filters=filters,
        autofill_legend=autofill_legend,
    )

    assert pivot_resp.status_code == HTTPStatus.OK, f"Expected HTTP status 200, got {pivot_resp.status_code}"

    expected_pivot_legend_tuples: list[Tuple[str, PivotRole]]
    if custom_pivot_legend_check is not None:
        # make a copy that can be appended to
        expected_pivot_legend_tuples = [item for item in custom_pivot_legend_check]
    else:
        expected_pivot_legend_tuples = [
            *[(norm_title(title), PivotRole.pivot_column) for title in columns],
            *[(norm_title(title), PivotRole.pivot_row) for title in rows],
            *[(norm_title(title), PivotRole.pivot_measure) for title in measures],
            *[(norm_title(title), PivotRole.pivot_annotation) for title in (annotations or ())],
            (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
        ]

    actual_pivot_legend_tuples = [(item.title, item.role_spec.role) for item in pivot_resp.data["pivot"]["structure"]]
    assert actual_pivot_legend_tuples == expected_pivot_legend_tuples

    legend, liid_map = _get_regular_legend_for_pivot(
        dataset=dataset,
        columns=columns,
        rows=rows,
        measures=measures,
        annotations=annotations,
        totals=check_totals,
    )
    legend_for_result = [item for item in legend if item.ref.type != QueryItemRefType.measure_name]
    result_resp = data_api.get_result(dataset=dataset, fields=legend_for_result, filters=filters, fail_ok=True)
    assert result_resp.status_code == HTTPStatus.OK, result_resp.json

    result_mapper_measure_liids = frozenset(
        chain.from_iterable(liid_map[title] for title in chain(measures, annotations or ()))
    )

    result_data_abstraction = ResultDataAbstraction.from_response(result_resp)
    result_data_mapper: DataCellMapper1D = result_data_abstraction.get_1d_mapper(
        measure_liids=result_mapper_measure_liids
    )
    pivot_data_abstraction = PivotDataAbstraction.from_response(pivot_resp)
    pivot_data_mapper: DataCellMapper1D = pivot_data_abstraction.get_1d_mapper()

    def _is_total(dimension: Optional[list], totals_liids: set[int]):
        return dimension is not None and any(value[0][1] in totals_liids for value in dimension)

    totals_liids = {
        item.legend_item_id for item in pivot_resp.data["fields"] if item.role_spec.role.value == "template"
    }
    for column in pivot_resp.data["pivot_data"]["columns_with_info"]:
        role = column["header_info"].role_spec.role.value
        expected_role = "total" if _is_total(column["cells"], totals_liids) else "data"
        assert role == expected_role, f'Column {column["cells"]} is {role}, but it expected to be {expected_role}'
    for row in pivot_resp.data["pivot_data"]["rows"]:
        role = row["header_with_info"]["header_info"].role_spec.role.value
        expected_role = "total" if _is_total(row["header"], totals_liids) else "data"
        assert role == expected_role, f'Row {row["header"]} is {role}, but it expected to be {expected_role}'

    # workaround for missing annotations for totals
    # remove them from the result mapper
    if annotations:

        def apply_anno_tag(dims: DataCellTuple, data_item: DataItem) -> AbstractSet[DataItemTag]:
            if data_item.cell.title in anno_only:
                return data_item.meta.tags | {DataItemTag.annotation}
            return data_item.meta.tags

        anno_only = set(annotations) - set(measures)  # annotations that are not present as regular measures
        if anno_only:
            result_data_mapper = result_data_mapper.apply_tagger(apply_anno_tag)
            result_data_mapper = result_data_mapper.apply_filter(
                exclude_tag_combos={frozenset({DataItemTag.total, DataItemTag.annotation})},
            )

    SimpleMapper1DComparator().assert_equal(pivot_data_mapper, result_data_mapper)

    if expected_notifications is not None:
        actual_notifications = [
            NotificationType(notification["locator"]) for notification in pivot_resp.data["notifications"]
        ]
        assert actual_notifications == expected_notifications

    pivot_data = pivot_resp.data["pivot_data"]
    pivot_rows = pivot_data["rows"]

    # Check columns
    if min_col_cnt is not None:
        assert len(pivot_data["columns"]) >= min_col_cnt
    if max_col_cnt is not None:
        assert len(pivot_data["columns"]) <= max_col_cnt

    # Check row dimension headers
    actual_row_dim_headers = [header[0][0] for header in pivot_data["row_dimension_headers"]]
    expected_row_dim_headers = [
        item.title for item in pivot_resp.data["pivot"]["structure"] if item.role_spec.role == PivotRole.pivot_row
    ]
    assert actual_row_dim_headers == expected_row_dim_headers

    # Check row headers
    if min_row_cnt is not None:
        assert len(pivot_rows) >= min_row_cnt
    if max_row_cnt is not None:
        assert len(pivot_rows) <= max_row_cnt

    # Check measure values
    pivot_value_cnt = pivot_data_mapper.get_value_count()
    if min_value_cnt is not None:
        assert pivot_value_cnt >= min_value_cnt
    if max_value_cnt is not None:
        assert pivot_value_cnt <= max_value_cnt

    block_meta_list = pivot_resp.json["blocks"]
    if not check_totals:
        assert len(block_meta_list) == 1
    block_meta = block_meta_list[0]
    assert block_meta["block_id"] == 0
    assert block_meta["query"] is not None

    return pivot_data_abstraction


def get_all_measure_cells(pivot_rows: list[dict]) -> list[list[list]]:
    return [cell for row in pivot_rows for cell in row["values"] if cell is not None]  # Each cell is a list of pairs

lib/dl_api_client/dl_api_client/dsmaker/pivot_utils.py
datalens-tech/datalens-backend
from __future__ import annotations

from datetime import (
    date,
    datetime,
)
from enum import Enum
import json
from typing import (
    Any,
    ClassVar,
    Collection,
    Dict,
    Generator,
    Generic,
    List,
    Optional,
    Tuple,
    TypeVar,
    Union,
)
import uuid

import attr

from dl_constants.enums import (
    AggregationFunction,
    BinaryJoinOperator,
    CalcMode,
    ComponentErrorLevel,
    ComponentType,
    ConditionPartCalcMode,
    DataSourceType,
    FieldRole,
    FieldType,
    FieldVisibility,
    JoinType,
    LegendItemType,
    ManagedBy,
    OrderDirection,
    ParameterValueConstraintType,
    PivotHeaderRole,
    PivotRole,
    QueryBlockPlacementType,
    QueryItemRefType,
    RangeType,
    UserDataType,
    WhereClauseOperation,
)


class Action(Enum):
    add = "add"
    update = "update"
    delete = "delete"
    refresh = "refresh"


_API_OBJ_TV = TypeVar("_API_OBJ_TV", bound="ApiProxyObject")


@attr.s
class ApiProxyObject:
    id: str = attr.ib(kw_only=True, factory=lambda: str(uuid.uuid4()))
    # special attributes, not part of object data, postfixed with `_`
    created_: bool = attr.ib(kw_only=True, default=False)  # True means that the item has already been added via API

    def clone(self: _API_OBJ_TV, **kwargs: Any) -> _API_OBJ_TV:
        return attr.evolve(self, **kwargs)

    def prepare(self) -> None:
        """Custom initialization of properties"""

    def _make_action(self, action: Action, data: Optional[dict] = None) -> UpdateAction:
        return UpdateAction(
            action=action,
            item=self,
            custom_data=data,
        )

    def add(self) -> UpdateAction:
        """Generate ``add_*`` action"""
        return self._make_action(Action.add)

    def update(self, **data) -> UpdateAction:
        """Generate ``update_*`` action"""
        return self._make_action(Action.update, data=data)

    def delete(self) -> UpdateAction:
        """Generate ``delete_*`` action"""
        return self._make_action(Action.delete)

    def set_name(self, name: str) -> None:
        """
        Called when ``self`` is added to a container under alias ``name``.
        Can be used to define item's name /title.
        """


@attr.s(auto_attribs=True)
class UpdateAction:
    action: Action
    item: ApiProxyObject
    custom_data: Optional[dict]
    order_index: int = 0


_ITEM_TV = TypeVar("_ITEM_TV", bound=ApiProxyObject)


class Container(Generic[_ITEM_TV]):
    """
    A partially dict-like container for nested items where each item has its own string alias.
    Items can be fetched both by integer (by index) and string (by alias) keys.
    """

    def __init__(self, data: Union[list, tuple, dict, "Container"] = None):
        self._item_ids: List[str] = []  # order list of object IDs
        self._items: Dict[str, _ITEM_TV] = {}  # objects by ID
        self._id_by_alias: Dict[str, str] = {}  # alias -> ID

        if data:
            self.__iadd__(data)

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}({self._items})"

    def __str__(self) -> str:
        return self.__repr__()

    def __iter__(self) -> Generator[_ITEM_TV, None, None]:
        """
        Iterate over items in container.
        Note that this is different from ``dict`` in that it yields values, not keys.
        """
        for id in self._item_ids:
            yield self._items[id]

    def __setitem__(self, alias: str, item: _ITEM_TV):  # type: ignore  # TODO: fix
        """
        Add item to container under given alias.
        Item order is preserved.
        """
        if not isinstance(alias, str):
            raise TypeError(f"Invalid key type for item assignment in container: {type(alias)}")

        self._items[item.id] = item
        self._item_ids.append(item.id)
        self._id_by_alias[alias] = item.id
        if hasattr(item, "set_name"):
            item.set_name(alias)

    def __getitem__(self, key: Union[str, int]) -> _ITEM_TV:
        """Return item either by index (if ``key`` is ``int``) or alias (``key`` is ``str``)"""
        if isinstance(key, str):
            # key is an alias
            return self._items[self._id_by_alias[key]]
        if isinstance(key, int):
            # key is an index
            return self._items[self._item_ids[key]]

        raise TypeError(f"Invalid key type for container: {type(key)}")

    def items(self) -> Generator[Tuple[str, _ITEM_TV], None, None]:
        """Just like ``dict.items()`` - iterate over key-value tuples."""
        alias_by_id = {id: alias for alias, id in self._id_by_alias.items()}
        for id in self._item_ids:
            yield alias_by_id[id], self._items[id]

    def __iadd__(self, other: Union[list, tuple, dict, Container]) -> Container:
        """
        Extend by adding items from another collection.
        If items don't have aliases, then use IDs as aliases
        """
        if isinstance(other, (list, tuple)):
            other = {item.id: item for item in other}
        if isinstance(other, (dict, Container)):
            for alias, item in other.items():
                self[alias] = item
        else:
            raise KeyError(type(other))
        return self

    def __add__(self, other: Union[list, tuple, dict, Container]) -> Container:
        """
        Return new container that consists of items from ``self`` and from ``other``.
        The new container is constructed using the logic in ``__iadd__``.
        """
        new_cont = self.__class__()
        new_cont += other
        return new_cont

    def __len__(self) -> int:
        """Return the number of items in the container."""
        return len(self._items)

    def get_alias(self, id: str) -> Optional[str]:
        """Find item by its ID and return its alias."""
        alias_by_id = {id: alias for alias, id in self._id_by_alias.items()}
        return alias_by_id.get(id)

    def __bool__(self) -> bool:
        return bool(self._items)

    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, Container):
            return False
        return (
            other._items == self._items
            and other._item_ids == self._item_ids
            and other._id_by_alias == self._id_by_alias
        )


@attr.s
class DataSource(ApiProxyObject):
    connection_id: str = attr.ib(default=None)
    source_type: Optional[DataSourceType] = attr.ib(default=None, converter=DataSourceType.normalize)
    parameters: dict = attr.ib(default=None)
    raw_schema: list = attr.ib(factory=list)
    index_info_set: Optional[list] = attr.ib(default=None)
    title: str = attr.ib(default=None)
    managed_by: ManagedBy = attr.ib(default=ManagedBy.user)
    valid: bool = attr.ib(default=True)

    def set_name(self, name: str) -> None:
        if self.title is None:
            self.title = name

    def prepare(self) -> None:
        super().prepare()
        if self.title is None:
            self.title = f"Source {self.id}"

    def avatar(self, **kwargs: Any) -> SourceAvatar:
        return SourceAvatar(source_id=self.id, **kwargs)

    def refresh(self, force_update_fields: bool = False) -> "UpdateAction":
        """Generate ``refresh_source`` action"""
        return self._make_action(Action.refresh, data={"force_update_fields": force_update_fields})


@attr.s
class SourceAvatar(ApiProxyObject):
    source_id: str = attr.ib(default=None)
    title: str = attr.ib(default=None)
    is_root: bool = attr.ib(default=False)
    managed_by: ManagedBy = attr.ib(default=ManagedBy.user)
    valid: bool = attr.ib(default=True)

    def prepare(self) -> None:
        super().prepare()
        if self.title is None:
            self.title = f"Avatar {self.id}"

    def set_name(self, name: str) -> None:
        if self.title is None:
            self.title = name

    def join(
        self,
        other: "SourceAvatar",
        conditions: List["JoinCondition"] = None,
        join_type: JoinType = JoinType.inner,
    ) -> "AvatarRelation":
        return AvatarRelation(
            left_avatar_id=self.id,
            right_avatar_id=other.id,
            conditions=conditions or [],
            join_type=join_type,
        )

    def field(self, **kwargs) -> "ResultField":  # type: ignore  # TODO: fix
        return ResultField(avatar_id=self.id, **kwargs)


class JoinPart:
    calc_mode: ConditionPartCalcMode


@attr.s
class DirectJoinPart(JoinPart):
    calc_mode: ConditionPartCalcMode = ConditionPartCalcMode.direct
    source: str = attr.ib(default=None)


@attr.s
class ResultFieldJoinPart(JoinPart):
    calc_mode: ConditionPartCalcMode = ConditionPartCalcMode.result_field
    field_id: str = attr.ib(default=None)


@attr.s
class FormulaJoinPart(JoinPart):
    calc_mode: ConditionPartCalcMode = ConditionPartCalcMode.formula
    formula: str = attr.ib(default=None)


class ConditionMakerMixin:
    """Mixin that enables construction of join conditions using Python's comparison operators"""

    def _simple_condition(self, other: "ConditionMakerMixin", operator: BinaryJoinOperator) -> "JoinCondition":
        return JoinCondition(
            left_part=self.get_cpart_from_self(), right_part=other.get_cpart_from_self(), operator=operator
        )

    def __eq__(self, other: "ConditionMakerMixin"):  # type: ignore  # TODO: fix
        return self._simple_condition(other, BinaryJoinOperator.eq)

    def __ne__(self, other: "ConditionMakerMixin"):  # type: ignore  # TODO: fix
        return self._simple_condition(other, BinaryJoinOperator.ne)

    def __lt__(self, other: "ConditionMakerMixin"):  # type: ignore  # TODO: fix
        return self._simple_condition(other, BinaryJoinOperator.lt)

    def __le__(self, other: "ConditionMakerMixin"):  # type: ignore  # TODO: fix
        return self._simple_condition(other, BinaryJoinOperator.lte)

    def __gt__(self, other: "ConditionMakerMixin"):  # type: ignore  # TODO: fix
        return self._simple_condition(other, BinaryJoinOperator.gt)

    def __ge__(self, other: "ConditionMakerMixin"):  # type: ignore  # TODO: fix
        return self._simple_condition(other, BinaryJoinOperator.gte)

    def get_cpart_from_self(self) -> JoinPart:
        raise NotImplementedError


@attr.s
class _Column:
    # use attr.s on superclass of the "real" class so that comparison methods from ConditionMakerMixin are used
    title: str = attr.ib(default=None)
    name: str = attr.ib(default=None)
    user_type: Optional[UserDataType] = attr.ib(default=None, converter=UserDataType.normalize)
    native_type: Optional[dict] = attr.ib(default=None)
    nullable: bool = attr.ib(default=True)
    description: str = attr.ib(default="")
    has_auto_aggregation: bool = attr.ib(default=False)
    lock_aggregation: bool = attr.ib(default=False)


class Column(ConditionMakerMixin, _Column):  # type: ignore  # TODO: fix
    def get_cpart_from_self(self) -> JoinPart:
        return DirectJoinPart(source=self.name)


@attr.s(auto_attribs=True)
class JoinCondition:
    operator: BinaryJoinOperator
    left_part: JoinPart
    right_part: JoinPart


@attr.s
class AvatarRelation(ApiProxyObject):
    left_avatar_id: str = attr.ib(default=None)
    right_avatar_id: str = attr.ib(default=None)
    conditions: List[JoinCondition] = attr.ib(factory=list)
    join_type: JoinType = attr.ib(default=JoinType.inner)
    managed_by: ManagedBy = attr.ib(default=ManagedBy.user)

    def on(self, *conditions: List[JoinCondition]):  # type: ignore  # TODO: fix
        self.conditions += conditions  # type: ignore  # TODO: fix
        return self


_INNER_TYPE = TypeVar("_INNER_TYPE")


@attr.s
class ParameterValue(Generic[_INNER_TYPE]):
    type: UserDataType
    value: _INNER_TYPE = attr.ib()


@attr.s
class StringParameterValue(ParameterValue[str]):
    type: UserDataType = UserDataType.string


@attr.s
class IntegerParameterValue(ParameterValue[int]):
    type: UserDataType = UserDataType.integer


@attr.s
class FloatParameterValue(ParameterValue[float]):
    type: UserDataType = UserDataType.float


@attr.s
class DateParameterValue(ParameterValue[date]):
    type: UserDataType = UserDataType.date


@attr.s
class DateTimeParameterValue(ParameterValue[datetime]):
    type: UserDataType = UserDataType.datetime


@attr.s
class DateTimeTZParameterValue(ParameterValue[datetime]):
    type: UserDataType = UserDataType.datetimetz


@attr.s
class GenericDateTimeParameterValue(ParameterValue[datetime]):
    type: UserDataType = UserDataType.genericdatetime


@attr.s
class BooleanParameterValue(ParameterValue[bool]):
    type: UserDataType = UserDataType.boolean


@attr.s
class GeoPointParameterValue(ParameterValue[List[Union[int, float]]]):
    type: UserDataType = UserDataType.geopoint


@attr.s
class GeoPolygonParameterValue(ParameterValue[List[List[List[Union[int, float]]]]]):
    type: UserDataType = UserDataType.geopolygon


@attr.s
class UuidParameterValue(ParameterValue[str]):
    type: UserDataType = UserDataType.uuid


@attr.s
class MarkupParameterValue(ParameterValue[str]):
    type: UserDataType = UserDataType.markup


@attr.s
class ArrayStrParameterValue(ParameterValue[List[str]]):
    type: UserDataType = UserDataType.array_str


@attr.s
class ArrayIntParameterValue(ParameterValue[List[int]]):
    type: UserDataType = UserDataType.array_int


@attr.s
class ArrayFloatParameterValue(ParameterValue[List[float]]):
    type: UserDataType = UserDataType.array_float


@attr.s
class TreeStrParameterValue(ParameterValue[List[str]]):
    type: UserDataType = UserDataType.tree_str


@attr.s
class ParameterValueConstraint:
    type: ParameterValueConstraintType = ParameterValueConstraintType.all


@attr.s
class RangeParameterValueConstraint(ParameterValueConstraint):
    type: ParameterValueConstraintType = ParameterValueConstraintType.range
    min: Optional[ParameterValue] = attr.ib(default=None)
    max: Optional[ParameterValue] = attr.ib(default=None)


@attr.s
class SetParameterValueConstraint(ParameterValueConstraint):
    type: ParameterValueConstraintType = ParameterValueConstraintType.set
    values: List[ParameterValue] = attr.ib(factory=list)


def _make_pivot_role_spec(
    role: PivotRole,
    annotation_type: Optional[str] = None,
    target_legend_item_ids: Optional[list[int]] = None,
    direction: OrderDirection = OrderDirection.asc,
    measure_sorting_settings: Optional[PivotMeasureSorting] = None,
    allow_roles: Optional[Collection[PivotRole]] = None,
) -> PivotRoleSpec:
    if allow_roles is not None:
        assert role in allow_roles

    pivot_role_spec: PivotRoleSpec
    if role == PivotRole.pivot_annotation:
        pivot_role_spec = AnnotationPivotRoleSpec(
            role=role,
            annotation_type=annotation_type,
            target_legend_item_ids=target_legend_item_ids,
        )
    elif role == PivotRole.pivot_measure:
        pivot_role_spec = PivotMeasureRoleSpec(
            role=role,
            sorting=measure_sorting_settings,
        )
    elif role in (PivotRole.pivot_row, PivotRole.pivot_column):
        pivot_role_spec = DimensionPivotRoleSpec(
            role=role,
            direction=direction,
        )
    else:
        raise ValueError(role)

    return pivot_role_spec


@attr.s
class _ResultField(ApiProxyObject):
    title: Optional[str] = attr.ib(default=None)
    calc_mode: CalcMode = attr.ib(default=CalcMode.direct, converter=CalcMode.normalize)
    aggregation: AggregationFunction = attr.ib(
        default=AggregationFunction.none, converter=AggregationFunction.normalize
    )
    type: FieldType = attr.ib(default=FieldType.DIMENSION, converter=FieldType.normalize)
    source: Optional[str] = attr.ib(default=None)
    hidden: bool = attr.ib(default=False)
    description: str = attr.ib(default="")
    formula: str = attr.ib(default="")
    initial_data_type: Optional[UserDataType] = attr.ib(default=None, converter=UserDataType.normalize)
    cast: Optional[UserDataType] = attr.ib(default=None, converter=UserDataType.normalize)
    data_type: Optional[UserDataType] = attr.ib(default=None, converter=UserDataType.normalize)
    valid: bool = attr.ib(default=True)
    has_auto_aggregation: bool = attr.ib(default=False)
    lock_aggregation: bool = attr.ib(default=False)
    avatar_id: Optional[str] = attr.ib(default=None)
    managed_by: ManagedBy = attr.ib(default=ManagedBy.user, converter=ManagedBy.normalize)
    default_value: Optional[ParameterValue] = attr.ib(default=None)
    value_constraint: Optional[ParameterValueConstraint] = attr.ib(default=None)

    def set_name(self, name: str) -> None:
        if self.title is None:
            self.title = name

    def prepare(self) -> None:
        super().prepare()
        if self.source and not self.formula and self.default_value is None:
            self.calc_mode = CalcMode.direct
            self.formula = ""
            self.default_value = None
            self.value_constraint = None
        elif self.formula and not self.source and self.default_value is None:
            self.calc_mode = CalcMode.formula
            self.source = ""
            self.default_value = None
            self.value_constraint = None
        elif self.default_value is not None and not self.source and not self.formula:
            self.calc_mode = CalcMode.parameter
            self.source = ""
            self.formula = ""

        if self.title is None:
            if self.calc_mode == CalcMode.direct:
                self.title = self.source
            else:
                self.title = f"Formula {self.id}"

        self.avatar_id = self.avatar_id if self.calc_mode == CalcMode.direct else None

    def filter(self, op: Union[str, WhereClauseOperation], values: list) -> WhereClause:
        if isinstance(op, str):
            op = WhereClauseOperation[op]
        return WhereClause(column=self.id if self.id is not None else self.title, operation=op, values=values)

    @property
    def asc(self) -> OrderedField:
        return OrderedField(field_id=self.id, direction=OrderDirection.asc)

    @property
    def desc(self) -> OrderedField:
        return OrderedField(field_id=self.id, direction=OrderDirection.desc)

    def parameter_value(self, value: Any = None) -> ParameterFieldValue:
        return ParameterFieldValue(field_id=self.id, value=value if value is not None else self.default_value.value)

    def as_req_legend_item(
        self,
        role: FieldRole = FieldRole.row,
        range_type: Optional[RangeType] = None,
        dimension_values: Optional[Dict[int, Any]] = None,
        tree_prefix: Optional[list] = None,
        tree_level: Optional[int] = None,
        legend_item_id: Optional[int] = None,
        block_id: Optional[int] = None,
    ) -> RequestLegendItem:
        role_spec: RoleSpec  # TODO: Move role_spec creation to a separate func
        if role == FieldRole.range:
            assert range_type is not None
            role_spec = RangeRoleSpec(
                role=role,
                range_type=range_type,
            )
        elif role == FieldRole.tree:
            assert dimension_values is not None
            assert tree_prefix is not None
            if tree_level is None:
                tree_level = len(tree_prefix) + 1
            assert tree_level is not None
            role_spec = TreeRoleSpec(
                role=role,
                level=tree_level,
                prefix=json.dumps(tree_prefix),
                dimension_values=[
                    DimensionValueSpec(legend_item_id=dim_liid, value=dim_value)
                    for dim_liid, dim_value in dimension_values.items()
                ],
            )
        elif role == FieldRole.row:
            role_spec = RowRoleSpec(role=role)
        else:
            role_spec = RoleSpec(role=role)
        assert isinstance(self, ResultField)
        return RequestLegendItem(
            ref=IdRequestLegendItemRef(
                type=QueryItemRefType.id,
                id=self.id,
            ),
            role_spec=role_spec,
            legend_item_id=legend_item_id,
            block_id=block_id,
        )


@attr.s(frozen=True)
class WhereClause:
    operation: WhereClauseOperation = attr.ib(kw_only=True)
    column: str = attr.ib(kw_only=True)
    values: list = attr.ib(kw_only=True)
    block_id: Optional[int] = attr.ib(kw_only=True, default=None)
    ref: RequestLegendItemRef = attr.ib(kw_only=True)

    @ref.default
    def _make_ref(self) -> RequestLegendItemRef:
        return IdRequestLegendItemRef(id=self.column)

    def for_block(self, block_id: int) -> WhereClause:
        return attr.evolve(self, block_id=block_id)


@attr.s
class ObligatoryFilter(ApiProxyObject):
    field_guid: str = attr.ib(default="")
    default_filters: List[WhereClause] = attr.ib(factory=list)
    managed_by: ManagedBy = attr.ib(default=ManagedBy.user)
    valid: bool = attr.ib(default=True)


class ResultField(ConditionMakerMixin, _ResultField):  # type: ignore  # TODO: fix
    def get_cpart_from_self(self) -> JoinPart:
        return ResultFieldJoinPart(field_id=self.id)


@attr.s
class RoleSpec:  # noqa
    role: FieldRole = attr.ib(kw_only=True)


@attr.s
class DimensionRoleSpec(RoleSpec):
    visibility: FieldVisibility = attr.ib(kw_only=True, default=FieldVisibility.visible)


@attr.s
class RowRoleSpec(DimensionRoleSpec):
    pass


@attr.s
class TemplateRoleSpec(DimensionRoleSpec):  # noqa
    template: str = attr.ib(kw_only=True)


@attr.s
class TreeRoleSpec(DimensionRoleSpec):
    level: int = attr.ib(kw_only=True)
    prefix: str = attr.ib(kw_only=True)
    dimension_values: List[DimensionValueSpec] = attr.ib(kw_only=True, factory=list)


@attr.s
class RangeRoleSpec(RoleSpec):
    range_type: RangeType = attr.ib(kw_only=True)


@attr.s
class OrderByRoleSpec(RoleSpec):
    direction: OrderDirection = attr.ib(kw_only=True)


@attr.s
class LegendItemBase:  # noqa
    role_spec: RoleSpec = attr.ib(kw_only=True)
    legend_item_id: Optional[int] = attr.ib(kw_only=True, default=None)
    block_id: Optional[int] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class RequestLegendItemRef:
    type: QueryItemRefType = attr.ib(kw_only=True, default=QueryItemRefType.id)


@attr.s
class IdRequestLegendItemRef(RequestLegendItemRef):
    # Set defaults
    type: QueryItemRefType = attr.ib(kw_only=True, default=QueryItemRefType.id)
    id: str = attr.ib(kw_only=True)


@attr.s
class MeasureNameRequestLegendItemRef(RequestLegendItemRef):
    # Set defaults
    type: QueryItemRefType = attr.ib(kw_only=True, default=QueryItemRefType.measure_name)


@attr.s
class PlaceholderRequestLegendItemRef(RequestLegendItemRef):
    # Set defaults
    type: QueryItemRefType = attr.ib(kw_only=True, default=QueryItemRefType.placeholder)


@attr.s
class RequestLegendItem(LegendItemBase):
    ref: RequestLegendItemRef = attr.ib(kw_only=True)


@attr.s
class LegendItem(LegendItemBase):  # noqa
    legend_item_id: int = attr.ib(kw_only=True)  # redefine as strictly not None
    id: str = attr.ib(kw_only=True)
    title: str = attr.ib(kw_only=True)
    data_type: UserDataType = attr.ib(kw_only=True)
    field_type: FieldType = attr.ib(kw_only=True)
    item_type: LegendItemType = attr.ib(kw_only=True)


@attr.s(frozen=True)
class BlockPlacement:
    type: ClassVar[QueryBlockPlacementType]


@attr.s(frozen=True)
class RootBlockPlacement(BlockPlacement):
    type = QueryBlockPlacementType.root


@attr.s(frozen=True)
class DimensionValueSpec:
    legend_item_id: int = attr.ib(kw_only=True)
    value: Any = attr.ib(kw_only=True)


@attr.s(frozen=True)
class AfterBlockPlacement(BlockPlacement):
    type = QueryBlockPlacementType.after

    dimension_values: Optional[List[DimensionValueSpec]] = attr.ib(default=None)


@attr.s
class BlockSpec:
    block_id: int = attr.ib(kw_only=True)
    parent_block_id: Optional[int] = attr.ib(kw_only=True, default=None)
    placement: BlockPlacement = attr.ib(kw_only=True)


@attr.s(frozen=True)
class PivotPagination:
    offset_rows: Optional[int] = attr.ib(kw_only=True, default=None)
    limit_rows: Optional[int] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class PivotTotalItem:
    level: int = attr.ib(kw_only=True)


@attr.s(frozen=True)
class PivotTotals:
    rows: list[PivotTotalItem] = attr.ib(kw_only=True, factory=list)
    columns: list[PivotTotalItem] = attr.ib(kw_only=True, factory=list)

    @staticmethod
    def item(level: int) -> PivotTotalItem:
        return PivotTotalItem(level=level)


@attr.s
class PivotRoleSpec:
    role: PivotRole = attr.ib(kw_only=True)


@attr.s
class DimensionPivotRoleSpec(PivotRoleSpec):
    direction: OrderDirection = attr.ib(kw_only=True)


@attr.s
class AnnotationPivotRoleSpec(PivotRoleSpec):
    annotation_type: str = attr.ib(kw_only=True)
    target_legend_item_ids: Optional[list[int]] = attr.ib(kw_only=True)


@attr.s
class PivotMeasureRoleSpec(PivotRoleSpec):
    sorting: Optional[PivotMeasureSorting] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class PivotItemBase:
    legend_item_ids: list[int] = attr.ib(kw_only=True)
    role_spec: PivotRoleSpec = attr.ib(kw_only=True)


@attr.s(frozen=True)
class RequestPivotItem(PivotItemBase):
    title: Optional[str] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class PivotItem(PivotItemBase):
    pivot_item_id: int = attr.ib(kw_only=True)
    title: str = attr.ib(kw_only=True)


@attr.s
class PivotHeaderRoleSpec:
    role: PivotHeaderRole = attr.ib(kw_only=True, default=PivotHeaderRole.data)


@attr.s(frozen=True)
class PivotHeaderValue:
    value: str = attr.ib(kw_only=True)


@attr.s(slots=True)
class PivotHeaderInfo:
    sorting_direction: Optional[OrderDirection] = attr.ib(kw_only=True, default=None)
    role_spec: PivotHeaderRoleSpec = attr.ib(kw_only=True, factory=PivotHeaderRoleSpec)


@attr.s(frozen=True)
class PivotMeasureSortingSettings:
    header_values: list[PivotHeaderValue] = attr.ib(kw_only=True)
    direction: OrderDirection = attr.ib(kw_only=True, default=OrderDirection.asc)
    role_spec: PivotHeaderRoleSpec = attr.ib(kw_only=True, factory=PivotHeaderRoleSpec)


@attr.s(frozen=True)
class PivotMeasureSorting:
    column: Optional[PivotMeasureSortingSettings] = attr.ib(kw_only=True, default=None)
    row: Optional[PivotMeasureSortingSettings] = attr.ib(kw_only=True, default=None)


@attr.s
class OrderedField:
    field_id: str = attr.ib(kw_only=True)
    direction: OrderDirection = attr.ib(kw_only=True)
    block_id: Optional[int] = attr.ib(kw_only=True, default=None)

    def for_block(self, block_id: int) -> OrderedField:
        return attr.evolve(self, block_id=block_id)


@attr.s
class ParameterFieldValue:
    field_id: str = attr.ib(kw_only=True)
    value: Any = attr.ib(kw_only=True)
    block_id: Optional[int] = attr.ib(kw_only=True, default=None)


@attr.s
class ComponentError:
    level: ComponentErrorLevel = attr.ib()
    message: str = attr.ib()
    code: str = attr.ib(default="")
    details: dict = attr.ib(factory=dict)


@attr.s
class ComponentErrorPack:
    id: str = attr.ib()
    type: ComponentType = attr.ib()
    errors: List[ComponentError] = attr.ib(factory=list)


@attr.s
class ComponentErrorRegistry:
    items: List[ComponentErrorPack] = attr.ib(factory=list)

    def get_pack(self, id: str) -> Optional[ComponentErrorPack]:  # type: ignore  # TODO: fix
        for item in self.items:
            if item.id == id:
                return item


@attr.s
class ResultSchemaAux:
    inter_dependencies: dict = attr.ib(factory=dict)  # No need to go into details


@attr.s
class Dataset(ApiProxyObject):
    name: str = attr.ib(default=None)
    revision_id: Optional[str] = attr.ib(default=None)
    sources: Container[DataSource] = attr.ib(factory=Container, converter=Container)
    source_avatars: Container[SourceAvatar] = attr.ib(factory=Container, converter=Container)
    avatar_relations: Container[AvatarRelation] = attr.ib(factory=Container, converter=Container)
    result_schema: Container[ResultField] = attr.ib(factory=Container, converter=Container)
    result_schema_aux: ResultSchemaAux = attr.ib(factory=ResultSchemaAux)
    rls: dict = attr.ib(factory=dict)
    component_errors: ComponentErrorRegistry = attr.ib(factory=ComponentErrorRegistry)
    obligatory_filters: List[ObligatoryFilter] = attr.ib(default=attr.Factory(list))

    def prepare(self) -> None:
        super().prepare()
        if self.name is None:
            self.name = f"My Dataset {str(uuid.uuid4())}"

    @staticmethod
    def source(**kwargs: Any) -> DataSource:
        return DataSource(**kwargs)

    @staticmethod
    def col(name: str, **kwargs: Any) -> Column:
        if "title" not in kwargs:
            kwargs["title"] = name
        return Column(name=name, **kwargs)

    def field(self, avatar: SourceAvatar = None, **kwargs) -> "ResultField":  # type: ignore  # TODO: fix
        avatar_id = avatar.id if avatar is not None else kwargs.pop("avatar_id", None)
        return ResultField(avatar_id=avatar_id, **kwargs)

    def find_field(self, title: Optional[str] = None, id: Optional[str] = None) -> ResultField:
        for field in self.result_schema:
            if id is not None and field.id == id:
                return field
            if title is not None and field.title == title:
                return field
        raise KeyError("Field not found")

    def measure_name_as_req_legend_item(
        self,
        role: FieldRole = FieldRole.row,
        legend_item_id: Optional[int] = None,
        block_id: Optional[int] = None,
    ) -> RequestLegendItem:
        assert role in (FieldRole.row, FieldRole.info)
        if role == FieldRole.row:
            role_spec = RowRoleSpec(role=role)
        else:
            role_spec = RoleSpec(role=role)
        return RequestLegendItem(
            ref=MeasureNameRequestLegendItemRef(),
            role_spec=role_spec,
            legend_item_id=legend_item_id,
            block_id=block_id,
        )

    def placeholder_as_req_legend_item(
        self,
        role: FieldRole = FieldRole.row,
        template: Optional[str] = None,
        legend_item_id: Optional[int] = None,
        block_id: Optional[int] = None,
    ) -> RequestLegendItem:
        role_spec: RoleSpec
        if role == FieldRole.template:
            assert template is not None
            role_spec = TemplateRoleSpec(
                role=FieldRole.template,
                template=template,
            )
        elif role == FieldRole.row:
            role_spec = RowRoleSpec(role=role)
        else:
            role_spec = RoleSpec(role=role)
        return RequestLegendItem(
            ref=PlaceholderRequestLegendItemRef(),
            role_spec=role_spec,
            legend_item_id=legend_item_id,
            block_id=block_id,
        )

    def make_req_pivot_item(
        self,
        legend_item_ids: list[int],
        role: PivotRole = PivotRole.pivot_row,
        title: Optional[str] = None,
        annotation_type: Optional[str] = None,
        target_legend_item_ids: Optional[list[int]] = None,
        direction: OrderDirection = OrderDirection.asc,
        measure_sorting_settings: Optional[PivotMeasureSorting] = None,
    ) -> RequestPivotItem:
        role_spec = _make_pivot_role_spec(
            role=role,
            annotation_type=annotation_type,
            target_legend_item_ids=target_legend_item_ids,
            direction=direction,
            measure_sorting_settings=measure_sorting_settings,
        )
        return RequestPivotItem(
            role_spec=role_spec,
            legend_item_ids=legend_item_ids,
            title=title,
        )

lib/dl_api_client/dl_api_client/dsmaker/primitives.py
datalens-tech/datalens-backend
from http import HTTPStatus
from typing import (
    Any,
    Callable,
    Dict,
    Optional,
    Tuple,
)

from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.primitives import (
    Dataset,
    ParameterValue,
    ParameterValueConstraint,
)


def _add_anything_to_dataset(
    *,
    control_api: SyncHttpDatasetApiV1,
    dataset: Optional[Dataset] = None,
    dataset_id: Optional[str] = None,
    updater: Callable[[Dataset], Dataset],
    exp_status: int = HTTPStatus.OK,
    save: bool = True,
) -> Dataset:
    if dataset is None:
        assert dataset_id is not None
        ds = control_api.load_dataset(dataset=Dataset(id=dataset_id)).dataset
    else:
        ds = dataset

    ds = updater(ds)

    ds_resp = control_api.apply_updates(dataset=ds, fail_ok=True)
    assert ds_resp.status_code == exp_status, ds_resp.response_errors
    ds = ds_resp.dataset

    if save:
        ds = control_api.save_dataset(ds).dataset

    return ds


def add_formulas_to_dataset(
    *,
    api_v1: SyncHttpDatasetApiV1,  # FIXME: Rename to control_api
    dataset: Optional[Dataset] = None,
    dataset_id: Optional[str] = None,
    formulas: Dict[str, str],
    exp_status: int = HTTPStatus.OK,
    save: bool = True,
) -> Dataset:
    def _add_formulas(ds: Dataset) -> Dataset:
        for field_name, formula in formulas.items():
            ds.result_schema[field_name] = ds.field(formula=formula)
        return ds

    return _add_anything_to_dataset(
        control_api=api_v1,
        dataset=dataset,
        dataset_id=dataset_id,
        updater=_add_formulas,
        exp_status=exp_status,
        save=save,
    )


def add_parameters_to_dataset(
    *,
    api_v1: SyncHttpDatasetApiV1,  # FIXME: Rename to control_api
    dataset: Optional[Dataset] = None,
    dataset_id: Optional[str] = None,
    parameters: Dict[str, Tuple[ParameterValue, Optional[ParameterValueConstraint]]],
    exp_status: int = HTTPStatus.OK,
    save: bool = True,
) -> Dataset:
    def _add_parameters(ds: Dataset) -> Dataset:
        for field_name, parameter in parameters.items():
            ds.result_schema[field_name] = ds.field(
                cast=parameter[0].type, default_value=parameter[0], value_constraint=parameter[1]
            )
        return ds

    return _add_anything_to_dataset(
        control_api=api_v1,
        dataset=dataset,
        dataset_id=dataset_id,
        updater=_add_parameters,
        exp_status=exp_status,
        save=save,
    )


def create_basic_dataset(
    *,
    api_v1: SyncHttpDatasetApiV1,  # FIXME: Rename to control_api
    connection_id: str,
    data_source_settings: Dict[str, Any],
    formulas: Optional[Dict[str, str]] = None,
    save: bool = True,
) -> Dataset:
    ds = Dataset()
    ds.sources["source_1"] = ds.source(connection_id=connection_id, **data_source_settings)
    ds.source_avatars["avatar_1"] = ds.sources["source_1"].avatar()
    ds_resp = api_v1.apply_updates(dataset=ds, fail_ok=True)
    assert ds_resp.status_code == HTTPStatus.OK, ds_resp.response_errors
    ds = ds_resp.dataset

    if formulas is not None:
        ds = add_formulas_to_dataset(api_v1=api_v1, dataset=ds, formulas=formulas, save=False)

    if save:
        ds = api_v1.save_dataset(ds).dataset

    return ds

lib/dl_api_client/dl_api_client/dsmaker/shortcuts/dataset.py
datalens-tech/datalens-backend
from typing import Any

from dl_api_client.dsmaker.api.data_api import HttpDataApiResponse
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows


def get_range_values(response: HttpDataApiResponse) -> tuple[Any, Any]:
    data_rows = get_data_rows(response)
    if response.api_version in ("v1", "v1.5"):
        values = (data_rows[0][0], data_rows[1][0])
    elif response.api_version == "v2":
        values = (data_rows[0][0], data_rows[0][1])
    else:
        raise ValueError(response.api_version)
    return values

lib/dl_api_client/dl_api_client/dsmaker/shortcuts/range_data.py
datalens-tech/datalens-backend
from http import HTTPStatus
from typing import (
    Any,
    Union,
)

from dl_api_client.dsmaker.api.data_api import (
    HttpDataApiResponse,
    SyncHttpDataApiV1,
    SyncHttpDataApiV2,
)
from dl_api_client.dsmaker.primitives import Dataset


def get_data_rows(response: HttpDataApiResponse) -> list[list[Any]]:
    if response.api_version == "v2":
        raw_data_rows = response.json["result_data"][0]["rows"]
        data_rows = [row["data"] for row in raw_data_rows]
    else:
        data_rows = response.json["result"]["data"]["Data"]
    return data_rows


def get_regular_result_data(
    ds: Dataset,
    data_api: Union[SyncHttpDataApiV1, SyncHttpDataApiV2],
    field_names: list[str],
) -> dict[str, list[str]]:
    field_names = sorted(set(field_names))  # Deduplicate fields
    result_resp = data_api.get_result(
        dataset=ds,
        fields=[ds.find_field(title=field_name) for field_name in field_names],
        fail_ok=True,
    )
    assert result_resp.status_code == HTTPStatus.OK, result_resp.json
    data_rows = get_data_rows(result_resp)

    return {field_name: [row[field_idx] for row in data_rows] for field_idx, field_name in enumerate(field_names)}

lib/dl_api_client/dl_api_client/dsmaker/shortcuts/result_data.py
datalens-tech/datalens-backend
from __future__ import annotations

import json

from dl_api_client.dsmaker.primitives import (
    Dataset,
    RequestLegendItem,
)
from dl_constants.enums import FieldRole


def make_request_legend_items_for_tree_branches(
    tree_title: str,
    branches: list[tuple[int, list[str]]],
    dataset: Dataset,
) -> list[RequestLegendItem]:
    """
    Generate request legend items for tree branches specified like this:

    branches=[
        (1, []),
        (2, ['Company']),
        (3, ['Company', 'R&D']),
    ]
    """

    dept_tree = dataset.find_field(title=tree_title)
    branch_items = []
    for idx, (level, prefix) in enumerate(branches):
        dimension_values: dict[int, str]
        if level <= 1:
            dimension_values = {}
        else:
            # The only dimension is the tree itself,
            # so add the prefix to the condition
            assert idx > 0
            dimension_values = {idx - 1: json.dumps(prefix)}

        branch_items.append(
            dept_tree.as_req_legend_item(
                role=FieldRole.tree,
                tree_level=level,
                tree_prefix=prefix,
                dimension_values=dimension_values,
                legend_item_id=idx,
                block_id=idx,
            )
        )

    return branch_items

lib/dl_api_client/dl_api_client/dsmaker/shortcuts/tree.py
datalens-tech/datalens-backend
from dl_api_client.dsmaker.primitives import (
    JoinCondition,
    ResultField,
    ResultFieldJoinPart,
)
from dl_constants.enums import (
    BinaryJoinOperator,
    CalcMode,
    UserDataType,
)


# FIXME: This is practically a dummy. Add tests please


def test_field_condition_operators():
    field_1 = ResultField(
        title="field_1",
        id="12345",
        data_type=UserDataType.string,
        initial_data_type=UserDataType.string,
        cast=UserDataType.string,
        calc_mode=CalcMode.direct,
        source="column_1",
    )
    field_2 = ResultField(
        title="field_2",
        id="67890",
        data_type=UserDataType.string,
        initial_data_type=UserDataType.string,
        cast=UserDataType.string,
        calc_mode=CalcMode.direct,
        source="column_2",
    )
    condition = field_1 == field_2
    assert condition == JoinCondition(
        left_part=ResultFieldJoinPart(field_id=field_1.id),
        right_part=ResultFieldJoinPart(field_id=field_2.id),
        operator=BinaryJoinOperator.eq,
    )

lib/dl_api_client/dl_api_client_tests/unit/test_primitives.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_commons.logging import (
    clean_secret_data_in_headers,
    log_request_end,
    log_request_start,
)
from dl_api_commons.request_id import (
    make_uuid_from_parts,
    request_id_generator,
)


__all__ = (
    "make_uuid_from_parts",
    "request_id_generator",
    "log_request_start",
    "log_request_end",
    "clean_secret_data_in_headers",
)

lib/dl_api_commons/dl_api_commons/__init__.py
datalens-tech/datalens-backend
import enum
import itertools
from typing import (
    Iterable,
    Optional,
    Tuple,
    overload,
)


class AuthFailureError(Exception):
    internal_message: str
    user_message: Optional[str]
    response_code: Optional[int]

    def __init__(
        self, internal_message: str, *, user_message: Optional[str] = None, response_code: Optional[int] = None
    ) -> None:
        self.internal_message = internal_message
        self.user_message = user_message
        self.response_code = response_code
        super().__init__(internal_message, user_message, response_code)


class BadHeaderPrefixError(AuthFailureError):
    response_code: Optional[int] = 401


_default_skip_auth_path_list_raw: Tuple[str, ...] = (
    "/ping",
    "/unistat",
    "/metrics",
)

default_skip_auth_path_list: Tuple[str, ...] = tuple(
    itertools.chain(
        (path.rstrip("/") for path in _default_skip_auth_path_list_raw),
        (path.rstrip("/") + "/" for path in _default_skip_auth_path_list_raw),
    )
)


def match_path_prefix(prefix_list: Iterable[str], path: str) -> bool:
    for skip_path_template in prefix_list:
        if skip_path_template.endswith("*"):
            return path.startswith(skip_path_template[:-1])
        if path == skip_path_template:
            return True
    return False


class AuthTokenType(enum.Enum):
    bearer = "Bearer"
    oauth = "OAuth"


@overload
def get_token_from_authorization_header(secret_header_value: str, token_type: AuthTokenType) -> str:
    pass


@overload
def get_token_from_authorization_header(secret_header_value: None, token_type: AuthTokenType) -> None:
    pass


def get_token_from_authorization_header(secret_header_value, token_type):
    if secret_header_value is None:
        return None
    assert isinstance(secret_header_value, str)
    prefix = f"{token_type.value} "

    if secret_header_value.startswith(prefix):
        return secret_header_value.removeprefix(prefix)

    raise BadHeaderPrefixError(
        user_message="Invalid format of Authorization header",
        internal_message=f"Authorization header does not start with {prefix!r}",
    )

lib/dl_api_commons/dl_api_commons/access_control_common.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
from asyncio import AbstractEventLoop
from concurrent.futures.thread import ThreadPoolExecutor
import enum
import logging
import threading
from typing import (
    Generator,
    Generic,
    Optional,
    TypeVar,
    Union,
)

import attr


try:
    from asyncio.exceptions import TimeoutError
except ImportError:
    from asyncio.futures import TimeoutError  # type: ignore  # TODO: fix


class JobState(enum.Enum):
    worker_not_started = enum.auto()  # Initial state
    worker_startup_confirmed = enum.auto()  # Control sets this to confirm that worker should perform initialization
    worker_ready = enum.auto()  # Worker sets this state to indicate that generator is instantiated and
    ready_for_next_chunk = enum.auto()
    chunk_in_buffer = enum.auto()
    end_of_stream = enum.auto()
    pending_close = enum.auto()
    closed = enum.auto()
    error = enum.auto()


class _ErrorInWorkerThread(Exception):
    """
    Exception is thrown by sync control methods to indicate that exception fires in worker.
    Desired action in control coroutines: await worker thread future to re-raise exception.
    """


class _WaitForStateTimeout(Exception):
    """Indicates that target state was not reached within timeout"""


class AWFSGExeption(Exception):
    """Base exception"""


class InitializationFailed(AWFSGExeption):
    """
    Fires in case of errors during worker initialization phase:
      sync generator creation fails
      worker thread was not acquired within timeout
    """


class AlreadyStarted(AWFSGExeption):
    pass


class WorkerIsClosed(AWFSGExeption):
    pass


class WorkerDoesNotThrowExceptionInErrorState(AWFSGExeption):
    pass


class EndOfStream(AWFSGExeption):
    pass


class AWFSGRuntimeError(AWFSGExeption):
    """Something strange and totally unexpected happens"""

    pass


_STATE_ITEM_TV = TypeVar("_STATE_ITEM_TV")


class _NoSet:
    instance: "_NoSet" = None  # type: ignore  # TODO: fix


_NoSet.instance = _NoSet()


@attr.s
class SynchronizedJobState(Generic[_STATE_ITEM_TV]):
    _log: logging.LoggerAdapter = attr.ib()
    _monitor: threading.Condition = attr.ib(init=False, factory=threading.Condition)
    _buffer: Union[_NoSet, _STATE_ITEM_TV] = attr.ib(init=False, default=_NoSet.instance)
    _state: JobState = attr.ib(init=False, default=JobState.worker_not_started)

    def _ensure_monitor(self):  # type: ignore  # TODO: fix
        if not self._monitor._is_owned():  # type: ignore  # TODO: fix  # noqa
            msg = "Attempt to use synchronized method without monitor acquiring"
            self._log.error(msg)
            raise AWFSGRuntimeError(msg)

    def fetch_buffer(self) -> _STATE_ITEM_TV:
        self._ensure_monitor()
        if self._buffer is _NoSet.instance:
            msg = "Trying to fetch from empty buffer"
            self._log.error(msg)
            raise AWFSGRuntimeError(msg)
        self._buffer, ret = _NoSet.instance, self._buffer
        return ret  # type: ignore  # TODO: fix

    def set_buffer(self, val: _STATE_ITEM_TV):  # type: ignore  # TODO: fix
        self._ensure_monitor()
        if self._buffer is not _NoSet.instance:
            msg = "Trying to reset buffer"
            self._log.error(msg)
            raise AWFSGRuntimeError(msg)
        self._buffer = val

    def set_state(self, state: JobState):  # type: ignore  # TODO: fix
        self._ensure_monitor()
        self._state = state
        self._monitor.notify_all()

    @property
    def state(self) -> JobState:
        return self._state

    def wait_for_state_change(self, state: Optional[JobState] = None, timeout: Optional[float] = None) -> None:
        self._ensure_monitor()
        target_state = state
        state_before = self._state
        if state is None:
            self._log.debug("Waiting for state monitor event")
            self._monitor.wait(timeout=timeout)
        else:
            self._log.debug("Waiting for particular state %s for %s", target_state, timeout)
            result = self._monitor.wait_for(lambda: self._state == target_state, timeout=timeout)
            if not result:
                self._log.debug("Requested state was not achieved within timeout %s", timeout)
                raise _WaitForStateTimeout()
        state_after = self._state
        self._log.debug("State monitor event received. Transition: %s -> %s", state_before, state_after)

    def __enter__(self):  # type: ignore  # TODO: fix
        return self._monitor.__enter__()

    def __exit__(self, *args, **kwargs):  # type: ignore  # TODO: fix
        return self._monitor.__exit__(*args, **kwargs)  # type: ignore  # TODO: fix


_JOB_ITEM_TV = TypeVar("_JOB_ITEM_TV")


# TODO Switch level to debug
@attr.s(cmp=False, hash=False)
class Job(Generic[_JOB_ITEM_TV], metaclass=abc.ABCMeta):
    _service_tpe: ThreadPoolExecutor = attr.ib()
    _workers_tpe: ThreadPoolExecutor = attr.ib()
    _worker_thread_start_timeout: float = attr.ib(default=0.1)
    _worker_thread_start_confirmation_timeout: float = attr.ib(default=0.1)

    _loop: AbstractEventLoop = attr.ib(init=False, factory=asyncio.get_event_loop)
    _worker_done_fut: Optional[asyncio.Future] = attr.ib(init=False, default=None)

    _ss: SynchronizedJobState = attr.ib(init=False, default=None)
    _startup_lock: asyncio.Lock = attr.ib(factory=asyncio.Lock)
    _worker_thread_started_event = attr.ib(init=False, factory=asyncio.Event)
    _log: logging.LoggerAdapter = attr.ib(init=False, default=None)

    def __attrs_post_init__(self):  # type: ignore  # TODO: fix
        self._log = logging.LoggerAdapter(
            logging.getLogger(__name__),
            extra=dict(
                job_sys_id=id(self),
            ),
        )
        self._ss = SynchronizedJobState(log=self._log)

    @property
    def state(self) -> JobState:
        return self._ss.state

    @abc.abstractmethod
    def make_generator(self) -> Generator[_JOB_ITEM_TV, None, None]:
        pass

    def _worker_do_step(self, generator: Generator) -> bool:
        """:returns True if there is work to handle in next step"""
        if self.state == JobState.ready_for_next_chunk:
            try:
                fetched_chunk = next(generator)

            except StopIteration:
                self._ss.set_state(JobState.end_of_stream)
                return False

            except Exception:
                self._ss.set_state(JobState.error)
                raise

            else:
                self._ss.set_buffer(fetched_chunk)
                self._ss.set_state(JobState.chunk_in_buffer)

        elif self.state == JobState.pending_close:
            # Notifying generator that we won't receive any more data
            try:
                generator.close()
            except Exception:  # noqa
                self._log.exception("Exception during generator closing")

            self._ss.set_state(JobState.closed)
            return False

        else:
            self._log.error("Unknown state value, stopping processing: %s", self.state)
            self._ss.set_state(JobState.error)
            return False

        self._log.debug("Waiting for job state monitor")
        self._ss.wait_for_state_change()
        return True

    def _worker_thread_target(self) -> None:
        self._loop.call_soon_threadsafe(self._worker_thread_started_event.set)
        with self._ss:
            self._log.debug("Waiting for 'start_confirmed' state")
            try:
                self._ss.wait_for_state_change(
                    state=JobState.worker_startup_confirmed, timeout=self._worker_thread_start_confirmation_timeout
                )
            except _WaitForStateTimeout:
                self._log.error(
                    "Worker thread did not receive start confirmation within timeout %s",
                    self._worker_thread_start_confirmation_timeout,
                )
                self._ss.set_state(JobState.error)
                return

            self._log.debug("Start was confirmed. Trying to create generator")

            try:
                generator = self.make_generator()
            except Exception:
                self._log.debug("Generator creation was failed. Setting state 'error'")
                self._ss.set_state(JobState.error)
                raise

            self._ss.set_state(JobState.worker_ready)
            # Here worker thread is totally ready for processing items in
            self._ss.wait_for_state_change()

            try:
                should_continue = True
                while should_continue:
                    should_continue = self._worker_do_step(generator)
            finally:
                if generator.gi_frame is not None:
                    self._log.error("Worker main cycle exits with running generator")
                    try:
                        generator.close()
                    except Exception:  # noqa
                        self._log.exception("Exception during generator closing")

    # Initialization
    def _sync_confirm_worker_start_set_ready_for_data(self) -> None:
        with self._ss:
            self._ss.set_state(JobState.worker_startup_confirmed)
            self._ss.wait_for_state_change()

            if self.state == JobState.worker_ready:
                self._ss.set_state(JobState.ready_for_next_chunk)
            elif self.state == JobState.error:
                self._ss.set_state(JobState.closed)
                raise _ErrorInWorkerThread
            else:
                self._log.error("Unexpected state during startup procedure: %s", self.state)
                self._ss.set_state(JobState.pending_close)

    async def run(self) -> None:
        """Start worker in configured TPE and ensure it will start in configured timeout."""
        self._log.info("Starting async wrapper for sync generator")
        async with self._startup_lock:
            if self.state == JobState.worker_not_started:
                self._log.debug("Scheduling worker thread")
                self._worker_done_fut = asyncio.ensure_future(
                    self._loop.run_in_executor(self._workers_tpe, self._worker_thread_target)
                )
                self._log.debug("Waiting for 'worker_thread_started_event' from worker thread")
                try:
                    await asyncio.wait_for(
                        self._worker_thread_started_event.wait(),
                        timeout=self._worker_thread_start_timeout,
                    )
                # We don't need to close worker thread in case of timeout:
                #  Worker will wait for 'start_confirmed' state with timeout
                #  If timeout fires - it will terminates
                except TimeoutError:
                    raise InitializationFailed(f"Worker thread did not start in {self._worker_thread_start_timeout}")
                except Exception as err:
                    raise InitializationFailed() from err

                self._log.debug("Event 'worker_thread_started_event' received")
                try:
                    await self._loop.run_in_executor(
                        self._service_tpe, self._sync_confirm_worker_start_set_ready_for_data
                    )
                except _ErrorInWorkerThread:
                    try:
                        self._log.info("Worker state was switched to error. Awaiting worker thread...")
                        await self._worker_done_fut
                    except Exception:
                        self._log.info("Exception from worker thread was caught", exc_info=True)
                        raise
                    else:
                        raise WorkerDoesNotThrowExceptionInErrorState()
            else:
                raise AWFSGRuntimeError("Job already started")

    # Fetching
    def _sync_fetch_chunk_from_buffer_schedule_next(self) -> _JOB_ITEM_TV:
        self._log.debug("Waiting for state monitor to proceed chunk fetch")
        with self._ss:
            self._log.debug("State monitor acquired to proceed chunk fetch")
            while True:
                if self.state == JobState.end_of_stream:
                    self._ss.set_state(JobState.closed)
                    raise EndOfStream()

                elif self.state == JobState.chunk_in_buffer:
                    ret = self._ss.fetch_buffer()
                    self._ss.set_state(JobState.ready_for_next_chunk)
                    return ret

                elif self.state == JobState.ready_for_next_chunk:
                    # Case: worker thread did not catch monitor after previous fetch procedure launch
                    #  So we should give it another chance to
                    self._ss.wait_for_state_change()
                    continue

                elif self.state == JobState.error:
                    self._ss.set_state(JobState.closed)
                    raise _ErrorInWorkerThread()

                elif self.state in (JobState.closed, JobState.pending_close):
                    raise WorkerIsClosed()

                else:
                    self._log.error(
                        "Unexpected state encountered during fetching next, scheduling close: %s", self.state
                    )
                    self._ss.set_state(JobState.pending_close)
                    raise AWFSGRuntimeError(f"Unexpected state encountered during fetching next: {self.state}")

    # TODO FIX: Use async lock to prevent concurrent calls
    # TODO FIX: Assert that job is running
    async def get_next(self) -> _JOB_ITEM_TV:
        try:
            next_chunk = await self._loop.run_in_executor(
                self._service_tpe, self._sync_fetch_chunk_from_buffer_schedule_next
            )
            return next_chunk
        except _ErrorInWorkerThread:
            try:
                self._log.info("Worker state was switched to error. Awaiting worker thread...")
                await self._worker_done_fut  # type: ignore  # TODO: fix
            except Exception:
                self._log.info("Exception from worker thread was caught", exc_info=True)
                raise
            else:
                raise WorkerDoesNotThrowExceptionInErrorState()

    # Cancellation
    def _sync_close_wait_closed(self) -> None:
        self._log.info("Awaiting monitor to make state transition to closing")
        with self._ss:
            if self.state in (JobState.ready_for_next_chunk, JobState.chunk_in_buffer, JobState.pending_close):
                if self.state != JobState.pending_close:
                    self._log.info("Settings state to closing")
                    self._ss.set_state(JobState.pending_close)
                else:
                    self._log.info("State already is 'pending_close'")
                # TODO FIX: May be timeout
                while True:
                    self._log.info("Waiting for worker to stop")
                    self._ss.wait_for_state_change()
                    if self.state.closed:
                        self._log.info("State 'closed' was reached")
                        return

            elif self.state == JobState.closed:
                self._log.info("Already closed")
                return

            elif self.state == JobState.end_of_stream:
                self._ss.set_state(JobState.closed)
                return

            elif self.state == JobState.error:
                self._ss.set_state(JobState.closed)
                raise _ErrorInWorkerThread()

            elif self.state == JobState.worker_not_started:
                return

            else:
                raise AWFSGRuntimeError(f"Unexpected state during close: {self.state}")

    # TODO FIX: Use async lock to prevent concurrent calls
    # TODO FIX: Assert that job is running
    async def cancel(self) -> None:
        self._log.info("Launching close procedure in service TPE")
        try:
            await self._loop.run_in_executor(self._service_tpe, self._sync_close_wait_closed)
        except _ErrorInWorkerThread:
            await self._worker_done_fut  # type: ignore  # TODO: fix
        self._log.info("Generator was successfully cancelled")

lib/dl_api_commons/dl_api_commons/aio/async_wrapper_for_sync_generator.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Optional,
)

from aiohttp import web

from dl_api_commons.aiohttp import aiohttp_wrappers
from dl_api_commons.base_models import (
    NoAuthData,
    TenantCommon,
    TenantDef,
)


if TYPE_CHECKING:
    from aiohttp.typedefs import Handler

    from dl_api_commons.aio.typing import AIOHTTPMiddleware


def auth_trust_middleware(
    fake_user_id: Optional[str] = None,
    fake_user_name: Optional[str] = None,
    fake_tenant: Optional[TenantDef] = None,
) -> AIOHTTPMiddleware:
    @web.middleware
    @aiohttp_wrappers.DLRequestBase.use_dl_request
    async def actual_auth_trust_middleware(
        app_request: aiohttp_wrappers.DLRequestBase, handler: Handler
    ) -> web.StreamResponse:
        if aiohttp_wrappers.RequiredResourceCommon.SKIP_AUTH in app_request.required_resources:
            # LOGGER.info("Auth was skipped due to SKIP_AUTH flag in target view")
            pass
        elif app_request.request.headers.get("x-unauthorized"):
            raise web.HTTPUnauthorized()
        else:
            updated_rci = app_request.temp_rci.clone(
                auth_data=NoAuthData(), tenant=TenantCommon() if fake_tenant is None else fake_tenant
            )
            if fake_user_id is not None:
                updated_rci = updated_rci.clone(user_id=fake_user_id)
            if fake_user_name is not None:
                updated_rci = updated_rci.clone(user_name=fake_user_name)
            app_request.replace_temp_rci(updated_rci)

        return await handler(app_request.request)

    return actual_auth_trust_middleware

lib/dl_api_commons/dl_api_commons/aio/middlewares/auth_trust_middleware.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    Sequence,
)

from aiohttp import web
from aiohttp.typedefs import Handler
from multidict import CIMultiDict

from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_api_commons.aiohttp import aiohttp_wrappers
from dl_api_commons.headers import (
    DEFAULT_RCI_PLAIN_HEADERS,
    DEFAULT_RCI_SECRET_HEADERS,
    append_extra_headers_and_normalize,
)


def commit_rci_middleware(
    rci_extra_plain_headers: Optional[Sequence[str]] = None,
    rci_extra_secret_headers: Optional[Sequence[str]] = None,
) -> AIOHTTPMiddleware:
    plain_headers_to_rci = append_extra_headers_and_normalize(DEFAULT_RCI_PLAIN_HEADERS, rci_extra_plain_headers or ())
    secret_headers_to_rci = append_extra_headers_and_normalize(
        DEFAULT_RCI_SECRET_HEADERS, rci_extra_secret_headers or ()
    )

    @web.middleware
    @aiohttp_wrappers.DLRequestBase.use_dl_request
    async def actual_commit_rci_middleware(
        dl_request: aiohttp_wrappers.DLRequestBase, handler: Handler
    ) -> web.StreamResponse:
        headers = dl_request.request.headers
        dl_request.replace_temp_rci(
            dl_request.temp_rci.clone(
                plain_headers=CIMultiDict(
                    ((key, val) for key, val in headers.items() if key.lower() in plain_headers_to_rci)
                ),
                secret_headers=CIMultiDict(
                    ((key, val) for key, val in headers.items() if key.lower() in secret_headers_to_rci)
                ),
            )
        )
        dl_request.commit_rci()
        return await handler(dl_request.request)

    return actual_commit_rci_middleware

lib/dl_api_commons/dl_api_commons/aio/middlewares/commit_rci.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

from aiohttp import web

from dl_api_commons.aiohttp.aiohttp_wrappers import DLRequestBase
from dl_api_commons.logging import (
    CompositeLoggingContextController,
    RequestLoggingContextController,
)


def get_root_logging_context_controller(request: web.Request) -> CompositeLoggingContextController:
    key = DLRequestBase.KEY_LOG_CTX_CONTROLLER

    if key in request:
        composite_logging_controller = request[key]
        assert isinstance(composite_logging_controller, CompositeLoggingContextController)
    else:
        composite_logging_controller = CompositeLoggingContextController()
        request[key] = composite_logging_controller

    return composite_logging_controller


def add_logging_ctx_controller(request: web.Request, ctrl: RequestLoggingContextController) -> None:
    """
    Adds to current AIOHTTP request new request logging context controller.
    If composite controller is not created for current request it will be created automatically.
    """
    composite_logging_controller = get_root_logging_context_controller(request)
    composite_logging_controller.add_sub_controller(ctrl)


_endpoint_code_attr_name = "endpoint_code"


def get_endpoint_code(request: web.Request) -> Optional[str]:
    final_handler = request.match_info.handler

    if hasattr(final_handler, _endpoint_code_attr_name):
        maybe_endpoint_code = getattr(final_handler, _endpoint_code_attr_name)
        if isinstance(maybe_endpoint_code, str):
            return maybe_endpoint_code

    return None

lib/dl_api_commons/dl_api_commons/aio/middlewares/commons.py
datalens-tech/datalens-backend
from typing import Optional

from aiohttp import (
    hdrs,
    web,
)
from aiohttp.typedefs import Handler

from dl_api_commons.aio.typing import AIOHTTPMiddleware


ALL_METHODS = ("DELETE", "GET", "OPTIONS", "PATCH", "POST", "PUT")


def cors_middleware(
    allow_origins: tuple[str, ...],
    expose_headers: Optional[tuple[str, ...]] = None,
    allow_headers: Optional[tuple[str, ...]] = None,
    allow_methods: tuple[str, ...] = ALL_METHODS,
    allow_credentials: bool = False,
    max_age: Optional[int] = None,
) -> AIOHTTPMiddleware:
    @web.middleware
    async def middleware(request: web.Request, handler: Handler) -> web.StreamResponse:
        is_options_request = request.method == "OPTIONS"
        is_preflight_request = is_options_request and hdrs.ACCESS_CONTROL_REQUEST_METHOD in request.headers
        origin = request.headers.get(hdrs.ORIGIN)
        allow_all_origins = "*" in allow_origins

        if is_preflight_request:
            response = web.StreamResponse()
        else:
            response = await handler(request)

        if not origin:
            return response

        if not (origin in allow_origins or allow_all_origins):
            return response

        if allow_all_origins and not allow_credentials:
            cors_origin = "*"
        else:
            cors_origin = origin
        response.headers[hdrs.ACCESS_CONTROL_ALLOW_ORIGIN] = cors_origin

        if allow_credentials:
            response.headers[hdrs.ACCESS_CONTROL_ALLOW_CREDENTIALS] = "true"

        if expose_headers:
            response.headers[hdrs.ACCESS_CONTROL_EXPOSE_HEADERS] = ",".join(expose_headers)

        if is_options_request:
            if allow_headers:
                response.headers[hdrs.ACCESS_CONTROL_ALLOW_HEADERS] = ",".join(allow_headers)
            if allow_methods:
                response.headers[hdrs.ACCESS_CONTROL_ALLOW_METHODS] = ",".join(allow_methods)
            if max_age is not None:
                response.headers[hdrs.ACCESS_CONTROL_MAX_AGE] = str(max_age)

        if is_preflight_request:
            # skip other middlewares
            raise web.HTTPOk(text="", headers=response.headers)

        return response

    return middleware

lib/dl_api_commons/dl_api_commons/aio/middlewares/cors.py
datalens-tech/datalens-backend
import hmac
import logging
import time
from typing import (
    ClassVar,
    Optional,
)

from aiohttp import web
from aiohttp.typedefs import Handler
import attr

from dl_api_commons.aiohttp import aiohttp_wrappers


LOGGER = logging.getLogger(__name__)


def generate_csrf_token(user_id: str, timestamp: int, csrf_secret: str) -> str:
    msg = bytes("{}:{}".format(user_id, timestamp), encoding="utf-8")
    secret = bytes(csrf_secret, encoding="utf-8")
    h = hmac.new(key=secret, msg=msg, digestmod="sha1")
    return h.hexdigest()


@attr.s(frozen=True)
class CSRFMiddleware:
    USER_ID_COOKIES: ClassVar[tuple[str, ...]] = ()

    csrf_header_name: str = attr.ib()
    csrf_time_limit: int = attr.ib()
    csrf_secret: str = attr.ib()
    csrf_methods: tuple[str, ...] = attr.ib(default=("POST", "PUT", "DELETE"))

    def validate_csrf_token(self, token_header_value: Optional[str], user_token: str) -> bool:
        if not token_header_value:
            return False

        try:
            token, timestamp_str = token_header_value.split(":")
            timestamp = int(timestamp_str)
        except ValueError:
            return False

        if not token or timestamp <= 0:
            return False

        ts_now = int(time.time())
        if ts_now - timestamp > self.csrf_time_limit:
            return False

        if not hmac.compare_digest(generate_csrf_token(user_token, timestamp, self.csrf_secret), token):
            return False

        return True

    @web.middleware
    @aiohttp_wrappers.DLRequestBase.use_dl_request_on_method
    async def middleware(self, dl_request: aiohttp_wrappers.DLRequestBase, handler: Handler) -> web.StreamResponse:
        async def _continue_request() -> web.StreamResponse:
            return await handler(dl_request.request)

        if aiohttp_wrappers.RequiredResourceCommon.SKIP_CSRF in dl_request.required_resources:
            LOGGER.info("CSRF check was skipped due to SKIP_CSRF flag in target view")
            return await _continue_request()

        if dl_request.request.method not in self.csrf_methods or not dl_request.request.cookies:
            return await _continue_request()

        rci = dl_request.last_resort_rci
        user_tokens = []
        if rci is not None and rci.user_id:
            user_tokens.append(rci.user_id)
        for user_id_cookie in self.USER_ID_COOKIES:
            user_token = dl_request.request.cookies.get(user_id_cookie)
            if user_token is not None:
                user_tokens.append(user_token)

        if not user_tokens:
            return await _continue_request()

        csrf_token = dl_request.request.headers.get(self.csrf_header_name)

        LOGGER.info("Checking CSRF token for user tokens: %s", user_tokens)

        token_is_valid = None
        for user_token in user_tokens:
            token_is_valid = self.validate_csrf_token(csrf_token, user_token)
            if token_is_valid:
                LOGGER.info("CSRF token is valid for user token %s", user_token)
                break

        if not token_is_valid:
            LOGGER.info("CSRF validation failed.")
            return web.Response(body="CSRF validation failed", content_type="text/html", status=400)

        return await handler(dl_request.request)

lib/dl_api_commons/dl_api_commons/aio/middlewares/csrf.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from asyncio import CancelledError
from contextlib import ExitStack
import enum
import logging
import os
import sys
from typing import (
    Any,
    Dict,
    Optional,
    Tuple,
)

from aiohttp import web
import attr
import sentry_sdk
from typing_extensions import final

from dl_api_commons.logging import RequestLoggingContextController


LOGGER = logging.getLogger(__name__)


@final
class ErrorLevel(enum.Enum):
    info = enum.auto()
    warning = enum.auto()
    error = enum.auto()


@attr.s(auto_attribs=True)
class ErrorData:
    status_code: int
    response_body: Dict[str, Any]
    level: ErrorLevel = attr.ib(validator=[attr.validators.instance_of(ErrorLevel)])
    http_reason: Optional[str] = attr.ib(default=None)


@attr.s
class AIOHTTPErrorHandler(metaclass=abc.ABCMeta):
    sentry_app_name_tag: Optional[str] = attr.ib()
    use_sentry: bool = attr.ib(default=False)

    @sentry_app_name_tag.validator
    def _validate_sentry_app_name_tag(self, attribute, value):  # type: ignore  # TODO: fix
        if value is not None and not isinstance(value, str):
            raise ValueError(f"Unexpected type of sentry_app_name_tag: '{value}'")

    # noinspection PyUnusedLocal,PyMethodMayBeStatic
    def make_response(self, err_data: ErrorData, err: Exception, request: web.Request) -> web.Response:
        return web.json_response(err_data.response_body, status=err_data.status_code, reason=err_data.http_reason)

    @abc.abstractmethod
    def _classify_error(self, err: Exception, request: web.Request) -> ErrorData:
        pass

    def handle_error(
        self, err: Exception, request: web.Request, req_logging_ctx_ctrl: RequestLoggingContextController
    ) -> Tuple[web.Response, ErrorData]:
        if isinstance(err, CancelledError):
            LOGGER.warning("Client request was cancelled", exc_info=True)
            raise  # noqa
        elif isinstance(err, web.HTTPSuccessful):
            raise  # noqa
        else:
            self.maybe_debug_catch(err)
            try:
                # TODO CONSIDER: Validate that response is serializable
                err_data = self._classify_error(err, request)
                self.log_error_http_response(err, err_data, req_logging_ctx_ctrl)
                return self.make_response(err_data, err, request), err_data

            except Exception as on_error_error:  # noqa
                req_logging_ctx_ctrl.put_to_context("is_error", True)
                LOGGER.error(
                    "Error handler raised an error during handling this exception",
                    exc_info=(type(err), err, err.__traceback__),
                )
                LOGGER.critical("Error handler raised an error during creating error response", exc_info=True)
                err_data = ErrorData(
                    500,
                    http_reason="Internal server error",
                    response_body=dict(message="Internal server error"),
                    level=ErrorLevel.error,
                )
                return (
                    web.json_response(
                        dict(message="Internal Server Error"),
                        status=500,
                    ),
                    err_data,
                )

    def maybe_debug_catch(self, err):  # type: ignore  # TODO: fix
        if os.environ.get("BI_ERR_PDB"):
            sys.last_traceback = err.__traceback__
            import traceback

            traceback.print_exc()
            import ipdb

            ipdb.pm()

    def log_error_http_response(  # type: ignore  # TODO: fix
        self,
        err: Exception,
        err_data: ErrorData,
        request_logging_context_ctrl: RequestLoggingContextController,
    ):
        with ExitStack() as err_logging_stack:
            # skip all 4xx codes
            use_sentry = self.use_sentry and not (400 <= err_data.status_code < 500)
            if use_sentry:
                reg_exc_scope = err_logging_stack.enter_context(sentry_sdk.push_scope())
                # Tagging events as captured in this middleware
                reg_exc_scope.set_tag("event_source", "outer_error_handling_middleware")
                reg_exc_scope.set_tag("http_response_code", err_data.status_code)
                reg_exc_scope.set_extra("http_response_body", err_data.response_body)
                reg_exc_scope.set_level(err_data.level.name)

            if err_data.level == ErrorLevel.info:
                # Assumed that events will bi captured for warnings and higher, so we capture manually only for info
                if use_sentry:
                    sentry_sdk.capture_exception(err)
                LOGGER.info("Regular exception fired", exc_info=True)
            elif err_data.level == ErrorLevel.warning:
                LOGGER.warning("Warning exception fired", exc_info=True)
            else:
                request_logging_context_ctrl.put_to_context("is_error", True)
                LOGGER.exception("Caught an exception in request handler")

lib/dl_api_commons/dl_api_commons/aio/middlewares/error_handling_outer.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import contextlib
from http import HTTPStatus
import json
import logging
import time
from typing import (
    Any,
    Iterable,
    Optional,
)

from aiohttp import web
from aiohttp.typedefs import Handler
import attr
import sentry_sdk

from dl_api_commons.aio.middlewares.commons import (
    add_logging_ctx_controller,
    get_root_logging_context_controller,
)
from dl_api_commons.aio.middlewares.error_handling_outer import AIOHTTPErrorHandler
from dl_api_commons.aio.middlewares.request_id import RequestId
from dl_api_commons.aiohttp.aiohttp_wrappers import DLRequestBase
from dl_api_commons.exc import RequestTimeoutError
from dl_api_commons.logging import RequestLoggingContextController
from dl_api_commons.reporting.records import RequestResultReportingRecord
from dl_utils.aio import timeout


LOGGER = logging.getLogger(__name__)


@attr.s
class SentryRequestLoggingContextController(RequestLoggingContextController):
    _scope: sentry_sdk.Scope = attr.ib()
    _user_info: dict = attr.ib(init=False, factory=dict)

    allowed_tags = frozenset(
        (
            "request_id",
            "parent_request_id",
            "folder_id",
            "dataset_id",
            "conn_id",
            "conn_exec_cls",
            "conn_type",
        )
    )

    user_info_remap = {"user_name": "username", "user_id": "id"}

    def put_to_context(self, key: str, value: Any) -> None:
        val_to_write: Any
        try:
            json.dumps(value)
            val_to_write = value
        except Exception:  # noqa
            LOGGER.exception("Attempt to put to context non-jsonable value")
            try:
                val_to_write = str(value)
            except Exception:
                LOGGER.exception("Attempt to put to context non-stringifiable value")
                raise

        if key in self.user_info_remap:
            self._user_info[self.user_info_remap[key]] = val_to_write
            self._scope.set_user({**self._user_info})
        elif key in self.allowed_tags:
            self._scope.set_tag(key, val_to_write)
        else:
            self._scope.set_extra(key, val_to_write)


@contextlib.contextmanager  # type: ignore  # TODO: fix
def last_chance_log() -> Iterable[None]:
    try:
        yield
    except Exception:
        LOGGER.exception("Exception fired in outer_error_handling_middleware")
        raise


@attr.s
class RequestBootstrap:
    req_id_service: RequestId = attr.ib()
    error_handler: Optional[AIOHTTPErrorHandler] = attr.ib(default=None)
    timeout_sec: Optional[float] = attr.ib(default=None)

    @web.middleware
    async def middleware(self, request: web.Request, handler: Handler) -> web.StreamResponse:
        dl_request: Optional[DLRequestBase] = None
        result = None
        err_code: Optional[str] = None
        try:
            with contextlib.ExitStack() as top_level_stack:
                top_level_stack.enter_context(last_chance_log())
                if self.error_handler is not None and self.error_handler.use_sentry:
                    scope = top_level_stack.enter_context(sentry_sdk.configure_scope())
                    if self.error_handler.sentry_app_name_tag:
                        scope.set_tag("app_name", self.error_handler.sentry_app_name_tag)
                    add_logging_ctx_controller(request, SentryRequestLoggingContextController(scope))

                req_logging_ctx_ctrl = get_root_logging_context_controller(request)

                dl_request = self.req_id_service.dl_request_init(request)

                if self.timeout_sec is not None:
                    try:
                        async with timeout(self.timeout_sec):
                            result = await handler(request)
                    except asyncio.TimeoutError:
                        raise RequestTimeoutError() from None
                else:
                    result = await handler(request)
        except Exception as err:
            if self.error_handler is not None:
                result, err_data = self.error_handler.handle_error(err, request, req_logging_ctx_ctrl)
                err_code = err_data.response_body.get("code", None)
            else:
                raise
        finally:
            response_status_code = (
                result.status if isinstance(result, web.Response) else HTTPStatus.INTERNAL_SERVER_ERROR
            )
            if dl_request is not None:
                dl_request.reporting_registry.save_reporting_record(
                    RequestResultReportingRecord(
                        timestamp=time.time(),
                        response_status_code=response_status_code,
                        err_code=err_code,
                    )
                )
                self.req_id_service.dl_request_finilize(dl_request)
            self.req_id_service.on_request_end(request.method, request.path_qs, response_status_code)
        assert isinstance(result, web.StreamResponse)
        return result

lib/dl_api_commons/dl_api_commons/aio/middlewares/request_bootstrap.py
datalens-tech/datalens-backend
from __future__ import annotations

import json
import logging
from typing import (
    Optional,
    Type,
)

from aiohttp import web
import attr

from dl_api_commons import (
    log_request_start,
    make_uuid_from_parts,
    request_id_generator,
)
from dl_api_commons.aio.middlewares.commons import get_endpoint_code
from dl_api_commons.aiohttp import aiohttp_wrappers
from dl_api_commons.aiohttp.aiohttp_wrappers import DLRequestBase
from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.headers import (
    HEADER_DEBUG_MODE_ENABLED,
    HEADER_LOGGING_CONTEXT,
    get_x_dl_context,
)
from dl_api_commons.logging import (
    NON_TRANSITIVE_LOGGING_CTX_KEYS,
    log_request_end,
)
from dl_api_commons.reporting.profiler import DefaultReportingProfiler
from dl_api_commons.reporting.registry import DefaultReportingRegistry
from dl_app_tools import log
from dl_constants.api_constants import DLHeadersCommon


LOGGER = logging.getLogger(__name__)


@attr.s
class RequestId:
    header_name: str = attr.ib(default="X-Request-ID")
    dl_request_cls: Type[aiohttp_wrappers.DLRequestBase] = attr.ib(default=aiohttp_wrappers.DLRequestBase)
    append_own_req_id: bool = attr.ib(default=False)
    app_prefix: Optional[str] = attr.ib(default=None)
    accept_logging_ctx: bool = attr.ib(default=False)
    logging_ctx_header_name: Optional[str] = attr.ib(default=None)
    # Is used only for reporting registry to determine if request is anonymous
    #  Remove after adding mechanisms to detect if request in anonymous via RCI (e.g. based on user ID/AuthData)
    is_public_env: bool = attr.ib(default=False)

    def __attrs_post_init__(self) -> None:
        if self.append_own_req_id and self.accept_logging_ctx:
            raise ValueError("Flags accept_logging_ctx and append_own_req_id can not be used together")

        if self.logging_ctx_header_name is None:
            self.logging_ctx_header_name = HEADER_LOGGING_CONTEXT

    async def on_response_prepare(self, request: web.Request, response: web.StreamResponse) -> None:
        dl_req = self.dl_request_cls.get_for_request(request)
        if dl_req is not None and dl_req.last_resort_rci is not None and dl_req.last_resort_rci.request_id is not None:
            response.headers.add(self.header_name, dl_req.last_resort_rci.request_id)

    def dl_request_init(self, request: web.Request) -> DLRequestBase:
        dl_request = self.dl_request_cls.init_and_bind_to_aiohttp_request(request)

        parent_request_id = dl_request.request.headers.get(self.header_name)

        if self.append_own_req_id:
            request_id = make_uuid_from_parts(
                current=request_id_generator(self.app_prefix),
                parent=parent_request_id,
            )
        else:
            request_id = parent_request_id or request_id_generator(self.app_prefix)

        endpoint_code = get_endpoint_code(request)

        if dl_request.log_ctx_controller:
            dl_request.log_ctx_controller.put_to_context("request_id", request_id)
            dl_request.log_ctx_controller.put_to_context("parent_request_id", parent_request_id)
            dl_request.log_ctx_controller.put_to_context("endpoint_code", endpoint_code)

        log.context.reset_context()
        initial_log_context = log.context.get_log_context()
        assert len(initial_log_context) == 0, (
            f"Initial log context for request {request_id}" f" is not empty: {initial_log_context}"
        )

        log.context.put_to_context("request_id", request_id)
        log.context.put_to_context("parent_request_id", parent_request_id)
        log.context.put_to_context("endpoint_code", endpoint_code)

        if self.accept_logging_ctx and self.logging_ctx_header_name in request.headers:
            assert self.logging_ctx_header_name is not None
            try:
                logging_ctx_from_header = json.loads(request.headers.get(self.logging_ctx_header_name) or "")
                for ctx_key in NON_TRANSITIVE_LOGGING_CTX_KEYS:
                    logging_ctx_from_header.pop(ctx_key, None)
            except Exception:  # noqa
                LOGGER.exception("Can not parse logging context: %s", request.headers.get(self.logging_ctx_header_name))
            else:
                for ctx_key, ctx_val in logging_ctx_from_header.items():
                    if ctx_key in initial_log_context:
                        LOGGER.warning(
                            "Logging context key was in initial logging context: '%s' -> '%s'."
                            " Ignoring value from header",
                            ctx_key,
                            initial_log_context[ctx_key],
                        )
                    else:
                        log.context.put_to_context(ctx_key, ctx_val)

        log_request_start(
            logger=LOGGER, method=request.method, full_path=request.path_qs, headers=request.headers.items()
        )

        dl_request.init_temp_rci(
            RequestContextInfo.create(
                request_id=request_id,
                x_dl_debug_mode=bool(int(request.headers.get(HEADER_DEBUG_MODE_ENABLED, "0"))),
                endpoint_code=endpoint_code,
                x_dl_context=get_x_dl_context(dl_request.request.headers.get(DLHeadersCommon.DL_CONTEXT, "{}")),
                # This props will be filled in commit_rci middleware
                plain_headers=None,
                secret_headers=None,
                # This props will be filled in auth middleware
                tenant=None,
                user_id=None,
                user_name=None,
                auth_data=None,
            )
        )

        reporting_registry = DefaultReportingRegistry(rci=dl_request.temp_rci)
        dl_request.reporting_registry = reporting_registry
        reporting_profiler = DefaultReportingProfiler(
            reporting_registry=reporting_registry,
            is_public_env=self.is_public_env,
        )
        dl_request.reporting_profiler = reporting_profiler

        return dl_request

    def dl_request_finilize(self, dl_request: DLRequestBase) -> None:
        dl_request.reporting_profiler.on_request_end()

    def on_request_end(self, method: str, full_path: str, status_code: int) -> None:
        log_request_end(LOGGER, method, full_path, status_code)

lib/dl_api_commons/dl_api_commons/aio/middlewares/request_id.py
datalens-tech/datalens-backend
from __future__ import annotations

import contextlib
import logging
from typing import (
    Any,
    Generator,
    Optional,
)

from aiohttp import web
from aiohttp.typedefs import Handler
import opentracing
import opentracing.tags

from dl_api_commons.aio.middlewares.commons import (
    add_logging_ctx_controller,
    get_endpoint_code,
)
from dl_api_commons.logging_tracing import OpenTracingLoggingContextController


LOGGER = logging.getLogger(__name__)


class TracingService:
    """
    This middleware creates a root span for request to component.
    It provides interface for setting tags/other data to span via logging context controller.
    It is injected into AIOHTTP request and can be accessed via DLRequestBase.
    Assumed that it will be top-level middleware.
    """

    @contextlib.contextmanager
    def prepared_span_cm(self, request: web.Request, operation_name: str) -> Generator[opentracing.Span, None, None]:
        tracer = opentracing.global_tracer()
        span_context = tracer.extract(format=opentracing.Format.HTTP_HEADERS, carrier=request.headers)

        with tracer.start_active_span(operation_name, child_of=span_context) as scope:
            span = scope.span
            span.set_tag(opentracing.tags.SPAN_KIND, opentracing.tags.SPAN_KIND_RPC_SERVER)
            # TODO FIX: consider move to request ID middleware
            span.set_tag(opentracing.tags.HTTP_METHOD, request.method)
            span.set_tag(opentracing.tags.HTTP_URL, str(request.url))

            logging_context_controller = OpenTracingLoggingContextController(scope.span)
            add_logging_ctx_controller(request, logging_context_controller)
            yield span

    @staticmethod
    def set_span_tag(span: Optional[opentracing.Span], tag_name: str, tag_value: Any) -> None:
        if span is None:
            return
        span.set_tag(tag_name, tag_value)

    @web.middleware
    async def middleware(self, request: web.Request, handler: Handler) -> web.StreamResponse:  # type: ignore  # TODO: fix
        root_span: Optional[opentracing.Span] = None
        operation_name = get_endpoint_code(request)

        with contextlib.ExitStack() as exit_stack:
            if operation_name is not None:
                if isinstance(operation_name, str):
                    root_span = exit_stack.enter_context(self.prepared_span_cm(request, operation_name))
                else:
                    LOGGER.warning("Got not-string value for root span operation name : %s", operation_name)

            try:
                resp = await handler(request)
            except Exception:  # noqa
                self.set_span_tag(root_span, opentracing.tags.ERROR, True)
                raise
            else:
                self.set_span_tag(root_span, opentracing.tags.HTTP_STATUS_CODE, resp.status)
                return resp

lib/dl_api_commons/dl_api_commons/aio/middlewares/tracing.py
datalens-tech/datalens-backend
from __future__ import annotations

from aiohttp import web
import attr


@attr.s(frozen=True)
class ServerHeader:
    _server_header: str = attr.ib()

    @_server_header.validator
    def validate_server_header(self, _, value):  # type: ignore  # TODO: fix
        if not isinstance(value, str):
            raise TypeError(f"Server header must be a 'str', not '{type(value).__qualname__}'")
        if len(value) < 1:
            raise ValueError("Server header must have non-zero length")

    def add_signal_handlers(self, app: web.Application):  # type: ignore  # TODO: fix
        app.on_response_prepare.append(self.on_response_prepare)

    async def on_response_prepare(self, _: web.Request, response: web.Response):  # type: ignore  # TODO: fix
        response.headers["Server"] = self._server_header

lib/dl_api_commons/dl_api_commons/aio/server_header.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Awaitable,
    Callable,
)

from aiohttp import web
from aiohttp.typedefs import Handler


# TODO: replace AIOHTTPMiddleware with aiohttp.typedefs.Middleware
AIOHTTPMiddleware = Callable[[web.Request, Handler], Awaitable[web.StreamResponse]]
AIOHTTPMethodMiddleware = Callable[[Any, web.Request, Handler], Awaitable[web.StreamResponse]]

lib/dl_api_commons/dl_api_commons/aio/typing.py
datalens-tech/datalens-backend
from __future__ import annotations

from abc import (
    ABC,
    abstractmethod,
)
import asyncio
from contextlib import asynccontextmanager
import logging
import ssl
from types import TracebackType
from typing import (
    Any,
    AsyncGenerator,
    Awaitable,
    Callable,
    Optional,
    Type,
)

import aiohttp
import aiohttp.web
import attr

from dl_configs.utils import get_root_certificates_path


LOGGER = logging.getLogger(__name__)


THeaders = dict[str, str]
TCookies = dict[str, str]
TParams = dict[str, str]


class BaseRetrier(ABC):
    @abstractmethod
    async def retry_request(
        self, req_func: Callable[..., Awaitable[aiohttp.ClientResponse]], method: str, *args: Any, **kwargs: Any
    ) -> aiohttp.ClientResponse:
        """make some retry magic"""


class NoRetriesRetrier(BaseRetrier):
    async def retry_request(
        self, req_func: Callable[..., Awaitable[aiohttp.ClientResponse]], method: str, *args: Any, **kwargs: Any
    ) -> aiohttp.ClientResponse:
        return await req_func(method, *args, **kwargs)


@attr.s(kw_only=True, frozen=True)
class PredefinedIntervalsRetrier(BaseRetrier):
    retry_intervals: tuple[float, ...] = attr.ib(default=(0.5, 1.0, 2.0))
    retry_codes: set[int] = attr.ib(default={408, 429, 500, 502, 503, 504})
    retry_methods: set[str] = attr.ib(default={"GET"})
    retry_exception_classes: tuple[Type[Exception]] = attr.ib(default=(aiohttp.ClientError,))

    @retry_methods.validator
    def check_all_chars_upper(self, attribute: str, value: set[str]) -> None:
        for method in value:
            if not all("A" <= char <= "Z" for char in method):
                raise ValueError("method names should be in uppercase")

    def should_retry_exc(self, exc: Exception) -> bool:
        return any(isinstance(exc, retry_exc_cls) for retry_exc_cls in self.retry_exception_classes)

    async def retry_request(
        self, req_func: Callable[..., Awaitable[aiohttp.ClientResponse]], method: str, *args: Any, **kwargs: Any
    ) -> aiohttp.ClientResponse:
        if method.upper() not in self.retry_methods:
            return await req_func(method, *args, **kwargs)

        for idx, ret_interval in enumerate(self.retry_intervals):
            last = idx == len(self.retry_intervals) - 1

            try:
                resp = await req_func(method, *args, **kwargs)
            except Exception as err:
                if not self.should_retry_exc(err):
                    raise
                LOGGER.warning("aiohttp client error: %r", err)
                if last:
                    raise
            else:
                if resp.status in self.retry_codes:
                    LOGGER.warning("HTTP error: %r, ", resp.status)

                if resp.status not in self.retry_codes or last:
                    return resp

            await asyncio.sleep(ret_interval)

        raise Exception("You should not be here")


def default_ssl_context() -> ssl.SSLContext:
    return ssl.create_default_context(cafile=get_root_certificates_path())


@attr.s(kw_only=True)
class BIAioHTTPClient:
    base_url: str = attr.ib()

    cookies: TCookies = attr.ib(factory=dict)
    headers: THeaders = attr.ib(factory=dict)

    conn_timeout_sec: float = attr.ib(default=0.7)
    read_timeout_sec: float = attr.ib(default=30)
    # TODO: total_timeout_sec (retries!)

    raise_for_status: bool = attr.ib(default=True)

    retrier: BaseRetrier = attr.ib(factory=NoRetriesRetrier)

    ssl_context: Optional[ssl.SSLContext] = attr.ib(factory=default_ssl_context)
    _session: Optional[aiohttp.ClientSession] = attr.ib()
    close_session_on_exit: Optional[bool] = attr.ib(default=None)

    def make_default_session(self) -> aiohttp.ClientSession:
        return aiohttp.ClientSession(
            cookies=self.cookies, headers=self.headers, connector=aiohttp.TCPConnector(ssl_context=self.ssl_context)
        )

    @property
    def session(self) -> aiohttp.ClientSession:
        if self._session is None:
            self._session = self.make_default_session()
            if self.close_session_on_exit is None:
                self.close_session_on_exit = True
        return self._session

    async def close(self) -> None:
        if self.close_session_on_exit:
            await self.session.close()

    async def __aenter__(self) -> BIAioHTTPClient:
        return self

    async def __aexit__(
        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
    ) -> None:
        await self.close()

    def url(self, path: str) -> str:
        return "/".join(map(lambda s: s.strip("/"), (self.base_url, path)))

    @asynccontextmanager
    async def request(self, method: str, *args: Any, **kwargs: Any) -> AsyncGenerator[aiohttp.ClientResponse, None]:
        response = await self.retrier.retry_request(self._request, method, *args, **kwargs)  # type: ignore  # TODO: fix
        if self.raise_for_status:
            response.raise_for_status()
        yield response
        response.close()

    async def _request(
        self,
        method: str,
        path: str = "",
        params: Optional[TParams] = None,
        data: Optional[Any] = None,
        json_data: Optional[Any] = None,
        headers: Optional[THeaders] = None,
        cookies: Optional[TCookies] = None,
        conn_timeout_sec: Optional[float] = None,
        read_timeout_sec: Optional[float] = None,
    ) -> Optional[Any]:
        timeout = aiohttp.ClientTimeout(
            sock_connect=conn_timeout_sec or self.conn_timeout_sec,
            sock_read=read_timeout_sec or self.read_timeout_sec,
        )
        return await self.session.request(
            method=method,
            url=self.url(path),
            params=params,
            data=data,
            json=json_data,
            headers=self.headers | (headers or {}),
            cookies=self.cookies | (cookies or {}),
            timeout=timeout,
        )

lib/dl_api_commons/dl_api_commons/aiohttp/aiohttp_client.py
datalens-tech/datalens-backend
from __future__ import annotations

import enum
import functools
import inspect
import json
from typing import (
    Any,
    Awaitable,
    Callable,
    FrozenSet,
    Generic,
    Literal,
    Optional,
    Type,
    TypeVar,
    Union,
    overload,
)

from aiohttp import web
from aiohttp.typedefs import Handler

from dl_api_commons.aio.typing import (
    AIOHTTPMethodMiddleware,
    AIOHTTPMiddleware,
)
from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.exc import InvalidHeaderException
from dl_api_commons.logging import RequestLoggingContextController
from dl_api_commons.reporting.profiler import ReportingProfiler
from dl_api_commons.reporting.registry import ReportingRegistry
from dl_constants.api_constants import DLHeaders


class RequiredResource(enum.Enum):
    pass


class RequiredResourceCommon(RequiredResource):
    US_MANAGER = enum.auto()
    SERVICE_US_MANAGER = enum.auto()
    MASTER_KEY = enum.auto()
    SKIP_AUTH = enum.auto()
    SKIP_CSRF = enum.auto()


class RCINotSet(Exception):
    pass


class DLRequestBase:
    __slots__ = ("request",)
    request: web.Request

    KEY_DL_REQUEST = "dl_request"
    KEY_RCI_TEMP = "bi_request_context_info_temp"
    KEY_RCI = "bi_request_context_info"
    KEY_LOG_CTX_CONTROLLER = "bi_log_ctx_controller"
    KEY_REPORTING_REGISTRY = "reporting_registry"
    KEY_REPORTING_PROFILER = "reporting_profiler"

    HANDLER_ATTR_NAME_REQUIRED_RESOURCES = "REQUIRED_RESOURCES"
    NO_RESOURCES = frozenset()  # type: ignore  # TODO: fix

    # flag that forces `.url` to use `https` scheme
    enforce_https_on_self_url: bool = True

    def __init__(self, request: web.Request):
        self.request = request

    @classmethod
    def init_and_bind_to_aiohttp_request(cls, request: web.Request) -> DLRequestBase:
        dl_request = cls(request)
        request[cls.KEY_DL_REQUEST] = dl_request
        return dl_request

    @classmethod
    def get_for_request(cls, request: web.Request) -> Optional[DLRequestBase]:
        if cls.KEY_DL_REQUEST in request:
            return request[cls.KEY_DL_REQUEST]
        return None

    def _set_attr_once(self, name, value) -> None:
        if name in self.request:
            raise ValueError(f"Request key '{name}' already set")
        self.request[name] = value

    @property
    def url(self) -> str:
        if self.enforce_https_on_self_url and self.request.scheme == "http":
            # 'correct' but possibly problematic: `return self.request.clone(scheme='https').url`
            result = str(self.request.url)
            pfx = "http://"
            if result.startswith(pfx):
                return "https://" + result[len(pfx) :]
        return self.request.url  # type: ignore  # TODO: fix

    # TODO FIX: Check that is not used and remove
    @property
    def request_id(self) -> Optional[str]:
        return self.rci.request_id

    # TODO FIX: Check that is not used and remove
    @property
    def user_id(self) -> str:
        return self.rci.user_id  # type: ignore  # TODO: fix

    # TODO FIX: Check that is not used and remove
    @property
    def user_name(self):  # type: ignore  # TODO: fix
        return self.rci.user_name

    @property
    def temp_rci(self) -> RequestContextInfo:
        """
        Returns not-committed RCI. Should be used only during RCI building.
        """
        if self.KEY_RCI_TEMP not in self.request:
            raise RCINotSet("Temp RCI was not initiated")
        return self.request[self.KEY_RCI_TEMP]

    def init_temp_rci(self, rci: RequestContextInfo):  # type: ignore  # TODO: fix
        """This method should be called in request_id middleware (or another top-level middleware)"""
        self._set_attr_once(self.KEY_RCI_TEMP, rci)

    def replace_temp_rci(self, rci: RequestContextInfo) -> None:
        if self.is_rci_committed():
            raise ValueError("Attempt to modify temp RequestContextInfo after commit")
        self.request[self.KEY_RCI_TEMP] = rci
        if self.reporting_registry is not None:
            self.reporting_registry.rci = rci

    def update_temp_rci(self, **kwargs: Any) -> None:
        return self.replace_temp_rci(self.temp_rci.clone(**kwargs))

    def commit_rci(self) -> None:
        """
        Mark RCI as fully build. Should be called after all RCI-forming middleware done their job.
        """
        self._set_attr_once(self.KEY_RCI, self.temp_rci)
        if self.reporting_registry is not None:
            self.reporting_registry.rci = self.rci

    def is_rci_committed(self) -> bool:
        return self.KEY_RCI in self.request

    @property
    def rci(self) -> RequestContextInfo:
        """
        Returns RCI if committed. If not committed - exception will be thrown.
        """
        if self.is_rci_committed():
            return self.request[self.KEY_RCI]
        raise RCINotSet("RequestContextInfo is not committed for this request")

    @property
    def last_resort_rci(self) -> Optional[RequestContextInfo]:
        """
        :return: Returns committed RCI if exists
        """
        try:
            return self.rci
        except RCINotSet:
            try:
                return self.temp_rci
            except RCINotSet:
                return None

    @property
    def log_ctx_controller(self) -> Optional[RequestLoggingContextController]:
        return self.request.get(self.KEY_LOG_CTX_CONTROLLER)

    @property
    def reporting_registry(self) -> ReportingRegistry:
        return self.request.get(self.KEY_REPORTING_REGISTRY)  # type: ignore  # TODO: fix

    @reporting_registry.setter
    def reporting_registry(self, value: ReportingRegistry):  # type: ignore  # TODO: fix
        value.rci = self.temp_rci
        self._set_attr_once(self.KEY_REPORTING_REGISTRY, value)

    @property
    def reporting_profiler(self) -> ReportingProfiler:
        return self.request.get(self.KEY_REPORTING_PROFILER)  # type: ignore  # TODO: fix

    @reporting_profiler.setter
    def reporting_profiler(self, value: ReportingProfiler) -> None:
        self._set_attr_once(self.KEY_REPORTING_PROFILER, value)

    @overload
    def get_single_header(self, header: Union[DLHeaders, str]) -> Optional[str]:
        pass

    @overload  # noqa
    def get_single_header(self, header: Union[DLHeaders, str], required: Literal[False]) -> Optional[str]:
        pass

    @overload  # noqa
    def get_single_header(self, header: Union[DLHeaders, str], required: Literal[True]) -> str:
        pass

    def get_single_header(self, header, required=False):  # type: ignore  # TODO: fix  # noqa
        header_name = header.value if isinstance(header, DLHeaders) else header
        header_value_list = self.request.headers.getall(header_name, ())

        if len(header_value_list) == 0:
            if required:
                raise InvalidHeaderException("Header required, but missing", header_name=header_name)
            return None
        elif len(header_value_list) > 1:
            raise InvalidHeaderException("Expecting single header but multiple received", header_name=header_name)

        return header_value_list[0]

    def get_single_json_header(self, header: DLHeaders) -> Union[bool, int, float, list, dict, None]:
        raw_header = self.request.headers.get(header.value)
        if raw_header is None:
            return None
        try:
            return json.loads(raw_header)
        except json.JSONDecodeError:
            raise InvalidHeaderException(
                "Invalid JSON in header content",
                header_name=header.value,
            )

    @property
    def required_resources(self) -> FrozenSet[RequiredResource]:
        handler = self.request.match_info.handler

        if hasattr(handler, self.HANDLER_ATTR_NAME_REQUIRED_RESOURCES):
            return getattr(handler, self.HANDLER_ATTR_NAME_REQUIRED_RESOURCES)

        return self.NO_RESOURCES

    _SELF_TYPE = TypeVar("_SELF_TYPE", bound="DLRequestBase")

    @classmethod
    def use_dl_request(
        cls: Type[_SELF_TYPE], coro: Callable[[_SELF_TYPE, Handler], Awaitable[web.StreamResponse]]
    ) -> AIOHTTPMiddleware:
        if not inspect.iscoroutinefunction(coro):
            raise ValueError("This decorator may only be applied to a coroutine")

        @functools.wraps(coro)
        async def wrapper(request: web.Request, handler: Handler) -> web.StreamResponse:
            dl_request = request[cls.KEY_DL_REQUEST]
            if not isinstance(dl_request, cls):
                raise TypeError(
                    f"Error in .use_dl_request decorator on {coro}."
                    f" Actual class of DLRequest {type(dl_request).__name__}"
                    f" is not subclass of request: {cls.__name__}."
                )
            return await coro(dl_request, handler)

        return wrapper

    @classmethod
    def use_dl_request_on_method(
        cls, coro: Callable[[Any, _SELF_TYPE, Handler], Awaitable[web.StreamResponse]]
    ) -> AIOHTTPMethodMiddleware:
        if not inspect.iscoroutinefunction(coro):
            raise ValueError("This decorator may only be applied to a coroutine")

        @functools.wraps(coro)
        async def wrapper(self: Any, request: web.Request, handler: Handler) -> web.StreamResponse:
            dl_request = request[cls.KEY_DL_REQUEST]
            return await coro(self, dl_request, handler)

        return wrapper


_DL_REQUEST_TV = TypeVar("_DL_REQUEST_TV", bound=DLRequestBase)


class DLRequestView(web.View, Generic[_DL_REQUEST_TV]):
    dl_request_cls: Type[_DL_REQUEST_TV] = DLRequestBase  # type: ignore

    @property
    def dl_request(self) -> _DL_REQUEST_TV:
        dl_request = self.request[DLRequestBase.KEY_DL_REQUEST]
        if not isinstance(dl_request, self.dl_request_cls):
            raise TypeError(
                f"Actual class of DLRequest {type(dl_request).__name__}"
                f" is not subclass of request: {self.dl_request_cls.__name__}"
            )
        return dl_request

lib/dl_api_commons/dl_api_commons/aiohttp/aiohttp_wrappers.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    Any,
    Optional,
    Type,
    TypeVar,
    Union,
)

import attr
from multidict import (
    CIMultiDict,
    CIMultiDictProxy,
)

from dl_constants.api_constants import (
    DLCookies,
    DLHeaders,
    DLHeadersCommon,
)


class TenantDef(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_outbound_tenancy_headers(self) -> dict[DLHeaders, str]:
        raise NotImplementedError()

    @abc.abstractmethod
    def get_tenant_id(self) -> str:
        raise NotImplementedError()


class AuthData(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_headers(self) -> dict[DLHeaders, str]:
        raise NotImplementedError()

    @abc.abstractmethod
    def get_cookies(self) -> dict[DLCookies, str]:
        raise NotImplementedError()


@attr.s(frozen=True)
class RequestContextInfo:
    request_id: Optional[str] = attr.ib(default=None)
    tenant: Optional[TenantDef] = attr.ib(default=None)
    user_id: Optional[str] = attr.ib(default=None)
    user_name: Optional[str] = attr.ib(default=None)
    x_dl_debug_mode: Optional[bool] = attr.ib(default=None)
    endpoint_code: Optional[str] = attr.ib(default=None)

    _x_dl_context: dict[str, str] = attr.ib(factory=dict)
    _plain_headers: CIMultiDict = attr.ib(factory=CIMultiDict)
    _secret_headers: CIMultiDict = attr.ib(repr=False, factory=CIMultiDict)
    auth_data: Optional[AuthData] = attr.ib(repr=False, default=None)

    @property
    def x_dl_context(self) -> dict[str, str]:
        return self._x_dl_context.copy()

    @property
    def plain_headers(self) -> CIMultiDictProxy:
        return CIMultiDictProxy(self._plain_headers)

    @property
    def secret_headers(self) -> CIMultiDictProxy:
        return CIMultiDictProxy(self._secret_headers)

    @property
    def forwarder_for(self) -> Optional[str]:
        return self.plain_headers.get(DLHeadersCommon.FORWARDED_FOR.value)

    @property
    def client_ip(self) -> Optional[str]:
        if self.forwarder_for is not None:
            ip_list = [ip.strip() for ip in self.forwarder_for.split(",")]
            if len(ip_list) > 1:
                return ip_list[-2]
            else:
                # unlikely to happen
                return ip_list[0]
        return None

    @property
    def host(self) -> Optional[str]:
        return self.plain_headers.get("Host")

    @property
    def locale(self) -> Optional[str]:
        return self.plain_headers.get(DLHeadersCommon.ACCEPT_LANGUAGE.value)

    @staticmethod
    def normalize_headers_dict(headers: Union[CIMultiDict, dict, None]) -> CIMultiDict:
        if headers is None:
            return CIMultiDict()
        if isinstance(headers, dict):
            return CIMultiDict(headers)
        if isinstance(headers, CIMultiDict):
            return CIMultiDict(headers.items())
        raise TypeError(f"Unexpected type of headers: {type(headers)}")

    @classmethod
    def create_empty(cls: Type[_REQUEST_CONTEXT_INFO_TV]) -> _REQUEST_CONTEXT_INFO_TV:
        return cls()

    @classmethod
    def create(
        cls: Type[_REQUEST_CONTEXT_INFO_TV],
        request_id: Optional[str],
        tenant: Optional[TenantDef],
        user_id: Optional[str],
        user_name: Optional[str],
        x_dl_debug_mode: Optional[bool],
        endpoint_code: Optional[str],
        x_dl_context: Optional[dict],
        plain_headers: Union[CIMultiDict, dict, None],
        secret_headers: Union[CIMultiDict, dict, None],
        auth_data: Optional[AuthData] = None,
    ) -> _REQUEST_CONTEXT_INFO_TV:
        return cls(
            request_id=request_id,
            tenant=tenant,
            user_id=user_id,
            user_name=user_name,
            x_dl_debug_mode=x_dl_debug_mode,
            endpoint_code=endpoint_code,
            x_dl_context=x_dl_context or {},
            plain_headers=cls.normalize_headers_dict(plain_headers),
            secret_headers=cls.normalize_headers_dict(secret_headers),
            auth_data=auth_data,
        )

    def clone(self: _REQUEST_CONTEXT_INFO_TV, **kwargs: Any) -> _REQUEST_CONTEXT_INFO_TV:
        return attr.evolve(self, **kwargs)


_REQUEST_CONTEXT_INFO_TV = TypeVar("_REQUEST_CONTEXT_INFO_TV", bound="RequestContextInfo")


@attr.s(frozen=True)
class TenantCommon(TenantDef):
    def get_outbound_tenancy_headers(self) -> dict[DLHeaders, str]:
        return {}

    def get_tenant_id(self) -> str:
        return "common"


@attr.s()
class NoAuthData(AuthData):
    def get_headers(self) -> dict[DLHeaders, str]:
        return {}

    def get_cookies(self) -> dict[DLCookies, str]:
        return {}

lib/dl_api_commons/dl_api_commons/base_models.py
datalens-tech/datalens-backend
import logging
import ssl
from typing import (
    Any,
    AsyncIterable,
    Mapping,
    Optional,
    Union,
)
import uuid

import aiohttp
import attr

from dl_api_commons.base_models import (
    AuthData,
    TenantDef,
)
from dl_api_commons.tracing import get_current_tracing_headers
from dl_configs.utils import get_root_certificates_path
from dl_constants.api_constants import (
    DLHeaders,
    DLHeadersCommon,
)


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class Req:
    method: str = attr.ib()
    url: str = attr.ib()
    params: dict[str, Union[str, int]] = attr.ib(default=None)
    data_json: Optional[dict[str, Any]] = attr.ib(default=None)
    data: Union[bytes, str, AsyncIterable[bytes], aiohttp.MultipartWriter, None] = attr.ib(default=None, kw_only=True)

    require_ok: bool = attr.ib(default=True, kw_only=True)
    extra_headers: Optional[dict[DLHeaders, str]] = attr.ib(default=None, kw_only=True)
    add_common_headers: bool = attr.ib(default=True, kw_only=True)


@attr.s(frozen=True, auto_attribs=True)
class Resp:
    status: int
    content: bytes
    content_type: Optional[str]
    json: dict
    req_id: str
    headers: Mapping


class RequestExecutionException(Exception):
    def __init__(self, msg: str, content: bytes, status: int, request_id: str):
        super().__init__(msg, content)
        self.status = status
        self.request_id = request_id


def get_default_aiohttp_session() -> aiohttp.ClientSession:
    return aiohttp.ClientSession(
        connector=aiohttp.TCPConnector(ssl_context=ssl.create_default_context(cafile=get_root_certificates_path())),
    )


@attr.s(auto_attribs=True)
class DLCommonAPIClient:
    _base_url: str = attr.ib()
    _tenant: TenantDef = attr.ib()
    _auth_data: AuthData = attr.ib()
    _req_id: Optional[str] = attr.ib(default=None)

    _extra_headers: Optional[dict[DLHeaders, str]] = attr.ib(default=None)
    _session: aiohttp.ClientSession = attr.ib(factory=get_default_aiohttp_session)

    @staticmethod
    def update_dl_headers(acc: dict[DLHeaders, str], update: dict[DLHeaders, str], stage: str) -> None:
        headers_intersection = acc.keys() & update.keys()
        if headers_intersection:
            LOGGER.warning(f"Got headers intersection on update on stage {stage!r}: {headers_intersection}")
        acc.update(update)

    # TODO FIX: Check if no overrides (take in account that headers are CI)
    @staticmethod
    def dl_headers_to_plain(dl_headers: dict[DLHeaders, str], extra_plain_headers: dict[str, str]) -> dict[str, str]:
        result: dict[str, str] = {dl_header.value: value for dl_header, value in dl_headers.items()}
        result.update(extra_plain_headers)
        return result

    def get_common_headers(self, req_id: str) -> dict[DLHeaders, str]:
        result: dict[DLHeaders, str] = {}

        self.update_dl_headers(result, {DLHeadersCommon.REQUEST_ID: req_id}, "req_id")
        self.update_dl_headers(result, self._tenant.get_outbound_tenancy_headers(), "tenancy")
        self.update_dl_headers(result, self._auth_data.get_headers(), "auth")

        return result

    def get_effective_headers_for_request(self, rq: Req, req_id: str) -> dict[str, str]:
        tracing_headers = get_current_tracing_headers()

        result: dict[DLHeaders, str] = {}

        if rq.add_common_headers:
            self.update_dl_headers(result, self.get_common_headers(req_id), "common_headers")

        if rq.extra_headers:
            self.update_dl_headers(result, rq.extra_headers, "request_extra_headers")

        if self._extra_headers:
            self.update_dl_headers(result, self._extra_headers, "common_extra_headers")

        return self.dl_headers_to_plain(result, tracing_headers)

    # TODO FIX: Sanitize path
    def make_full_url(self, path: str) -> str:
        return f"{self._base_url.rstrip('/')}/{path.lstrip('/')}"

    async def make_request(self, rq: Req) -> Resp:
        if rq.data is not None and rq.data_json is not None:
            raise ValueError("data_json and data can not be provided simultaneously")

        req_id = uuid.uuid4().hex if self._req_id is None else self._req_id
        resp = await self._session.request(
            rq.method,
            self.make_full_url(rq.url),
            headers=self.get_effective_headers_for_request(rq, req_id),
            params=rq.params,
            json=rq.data_json,
            data=rq.data,
        )
        content = await resp.read()

        if resp.content_type == "application/json":
            resp_json = await resp.json()
        else:
            resp_json = None

        if rq.require_ok:
            if resp.status >= 400:
                raise RequestExecutionException("Non-ok response", content, status=resp.status, request_id=req_id)

        return Resp(
            req_id=req_id,
            status=resp.status,
            json=resp_json,
            content=content,
            content_type=resp.content_type,
            headers=resp.headers,
        )

    async def close(self) -> None:
        await self._session.close()

lib/dl_api_commons/dl_api_commons/client/base.py
datalens-tech/datalens-backend
import contextlib
import json
import logging
from typing import (
    Any,
    Iterator,
    NoReturn,
    Optional,
)

import attr
import marshmallow.exceptions

from dl_api_commons.client.base import (
    DLCommonAPIClient,
    Req,
    Resp,
)
import dl_api_commons.exc


LOGGER = logging.getLogger(__name__)


@attr.s()
class CommonInternalAPIClient(DLCommonAPIClient):
    _use_workbooks_api: bool = attr.ib(default=False)
    _read_only: bool = attr.ib(default=False)

    async def make_request(self, rq: Req) -> Resp:
        if self._read_only:
            if rq.method.lower() != "get":
                raise AssertionError(
                    f"Can not execute {rq.method!r} {rq.url!r} due to {type(self).__name__!r} is read-only"
                )

        return await super().make_request(rq)

    @contextlib.contextmanager
    def deserialization_err_handler(self, resp: Resp, op_code: str) -> Iterator[None]:
        try:
            yield None
        except marshmallow.exceptions.ValidationError as ma_exc:
            LOGGER.info(f"Deserialization fail in internal API client: {type(self).__name__}", exc_info=True)
            effective_messages: dict[str, Any] = ma_exc.messages  # type: ignore

            raise dl_api_commons.exc.MalformedAPIResponseErr(
                self.create_exc_data(
                    resp,
                    operation=op_code,
                    validation_errors=effective_messages,
                )
            ) from ma_exc

    def create_exc_data(
        self, resp: Resp, operation: str, validation_errors: Optional[dict[str, Any]] = None
    ) -> dl_api_commons.exc.APIResponseData:
        body: Any
        try:
            body = resp.json
        except Exception:  # noqa
            body = resp.content

        return dl_api_commons.exc.APIResponseData(
            status_code=resp.status,
            response_body=body,
            operation=operation,
            response_body_validation_errors=validation_errors,
        )

    def raise_from_resp(self, resp: Resp, *, op_code: str) -> NoReturn:
        common_data = self.create_exc_data(resp, op_code)
        resp_data: str
        try:
            resp_data = json.dumps(common_data.response_body)
        except Exception:  # noqa
            resp_data = repr(common_data.response_body)

        LOGGER.info(
            f"Handling unexpected response in internal API client: {type(self).__name__}",
            extra=dict(
                status_code=resp.status,
                resp_data=resp_data,
            ),
        )

        if resp.status in (401, 403):
            raise dl_api_commons.exc.AccessDeniedErr(common_data)
        if resp.status == 404:
            raise dl_api_commons.exc.NotFoundErr(common_data)

        raise dl_api_commons.exc.InternalAPIBadResponseErr(common_data)

lib/dl_api_commons/dl_api_commons/client/common.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Generic,
    Optional,
    TypeVar,
)

import attr

from dl_constants.exc import DLBaseException


_EXC_DATA_TV = TypeVar("_EXC_DATA_TV")


def format_response_body(obj: Any) -> str:
    str_limit = 500
    tail_placeholder = "...truncated..."

    if isinstance(obj, str):
        if len(obj) > str_limit:
            return repr(obj[: str_limit - len(tail_placeholder)] + tail_placeholder)
        else:
            return repr(obj)
    else:
        orig_repr = repr(obj)
        return format_response_body(orig_repr)


class ExceptionWithData(Generic[_EXC_DATA_TV], Exception):
    _data: _EXC_DATA_TV

    def __init__(self, data: _EXC_DATA_TV):
        super().__init__(data)
        self._data = data

    @property
    def data(self) -> _EXC_DATA_TV:
        return self._data


@attr.s()
class APIResponseData:
    operation: str = attr.ib()
    status_code: Optional[int] = attr.ib()
    response_body: Optional[Any] = attr.ib(repr=format_response_body)
    response_body_validation_errors: Optional[dict[str, Any]] = attr.ib(default=None)


class InternalAPIBadResponseErr(ExceptionWithData[APIResponseData]):
    CODE = "INTERNAL_API_BAD_RESPONSE"


class AccessDeniedErr(InternalAPIBadResponseErr):
    CODE = "ACCESS_DENIED"


class NotFoundErr(InternalAPIBadResponseErr):
    CODE = "NOT_FOUND"


class MalformedAPIResponseErr(InternalAPIBadResponseErr):
    CODE = "MALFORMED_API_RESPONSE"


class InvalidHeaderException(Exception):
    schema_validation_messages: Optional[dict]
    header_name: str

    def __init__(self, *args: Any, header_name: str, schema_validation_messages: Optional[dict] = None):
        self.header_name = header_name
        self.schema_validation_messages = schema_validation_messages
        super().__init__(*args)


class FlaskRCINotSet(Exception):
    pass


class RequestTimeoutError(DLBaseException):
    err_code = DLBaseException.err_code + ["REQUEST_TIMEOUT"]
    default_message = "Backend app request timeout exceeded"

lib/dl_api_commons/dl_api_commons/exc.py
datalens-tech/datalens-backend
"""
Middleware to reset asyncio event loop on each request
"""
from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

import attr

from dl_api_commons.flask.middlewares.wsgi_middleware import FlaskWSGIMiddleware


if TYPE_CHECKING:
    from dl_api_commons.flask.types import (
        WSGIEnviron,
        WSGIReturn,
        WSGIStartResponse,
    )


@attr.s
class AIOEventLoopMiddleware(FlaskWSGIMiddleware):
    _APP_FLAG_ATTR_NAME = "_bi_middleware_aio_event_loop"

    debug: bool = attr.ib(default=False)  # If `True` - debug mode for event loop will be turned on

    def wsgi_app(self, environ: WSGIEnviron, start_response: WSGIStartResponse) -> WSGIReturn:
        loop = asyncio.new_event_loop()
        loop.set_debug(self.debug)
        asyncio.set_event_loop(loop)

        try:
            return self.original_wsgi_app(environ, start_response)
        finally:
            loop.close()
            asyncio.set_event_loop(None)

lib/dl_api_commons/dl_api_commons/flask/middlewares/aio_event_loop_middleware.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    Tuple,
)

import attr
import flask
from multidict import CIMultiDict

from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.exc import FlaskRCINotSet
from dl_api_commons.headers import (
    DEFAULT_RCI_PLAIN_HEADERS,
    DEFAULT_RCI_SECRET_HEADERS,
    append_extra_headers_and_normalize,
)


@attr.s(frozen=True)
class ReqCtxInfoMiddleware:
    """
    This middleware gathers context info from products of another middleware and request headers.
    It should be setup after request id & auth middleware.
    """

    plain_headers: Tuple[str, ...] = attr.ib(
        default=(),
        converter=lambda extra: append_extra_headers_and_normalize(  # type: ignore
            default=DEFAULT_RCI_PLAIN_HEADERS,
            extra=extra,
        ),
    )
    secret_headers: Tuple[str, ...] = attr.ib(
        default=(),
        converter=lambda extra: append_extra_headers_and_normalize(  # type: ignore
            default=DEFAULT_RCI_SECRET_HEADERS,
            extra=extra,
        ),
    )

    _G_ATTR_NAME_COMMITTED_RCI = "_bi_request_context_info"
    _G_ATTR_NAME_TEMP_RCI = "_bi_temp_request_context_info"

    def _commit_rci(self) -> None:
        # TODO CONSIDER: raise if something is missing
        temp_rci = self.get_temp_rci()

        final_rci = temp_rci.clone(
            # TODO CONSIDER: Add default user agent if not passed in header
            plain_headers=CIMultiDict(
                (
                    (
                        header,
                        flask.request.headers.get(header),
                    )
                    for header in self.plain_headers
                )
            ),
            secret_headers=CIMultiDict(
                (
                    (
                        header,
                        flask.request.headers.get(header),
                    )
                    for header in self.secret_headers
                )
            ),
        )
        setattr(flask.g, self._G_ATTR_NAME_COMMITTED_RCI, final_rci)

    def set_up(self, app: flask.Flask) -> None:
        app.before_request(self._commit_rci)

    @classmethod
    def get_last_resort_rci(cls) -> Optional[RequestContextInfo]:
        try:
            return cls.get_request_context_info()
        except FlaskRCINotSet:
            try:
                return cls.get_temp_rci()
            except FlaskRCINotSet:
                return None

    @classmethod
    def get_temp_rci(cls) -> RequestContextInfo:
        temp_rci = getattr(flask.g, cls._G_ATTR_NAME_TEMP_RCI, None)
        if temp_rci is None:
            raise FlaskRCINotSet("Temp RCI was not set for current request context")
        assert isinstance(temp_rci, RequestContextInfo)
        return temp_rci

    @classmethod
    def replace_temp_rci(cls, rci: RequestContextInfo) -> None:
        assert isinstance(rci, RequestContextInfo)
        setattr(flask.g, cls._G_ATTR_NAME_TEMP_RCI, rci)

    @classmethod
    def is_rci_committed(cls) -> bool:
        return hasattr(flask.g, cls._G_ATTR_NAME_COMMITTED_RCI)

    @classmethod
    def get_request_context_info(cls) -> RequestContextInfo:
        if hasattr(flask.g, cls._G_ATTR_NAME_COMMITTED_RCI):
            return getattr(flask.g, cls._G_ATTR_NAME_COMMITTED_RCI)

        raise FlaskRCINotSet("Request context info was not flushed")

lib/dl_api_commons/dl_api_commons/flask/middlewares/commit_rci_middleware.py
datalens-tech/datalens-backend
"""
Middleware to reset context on each request.
"""
from __future__ import annotations

import contextvars
from typing import TYPE_CHECKING

from dl_api_commons.flask.middlewares.wsgi_middleware import FlaskWSGIMiddleware


if TYPE_CHECKING:
    from dl_api_commons.flask.types import (
        WSGIEnviron,
        WSGIReturn,
        WSGIStartResponse,
    )


class ContextVarMiddleware(FlaskWSGIMiddleware):
    _APP_FLAG_ATTR_NAME = "_bi_middleware_flag_context_var"

    def wsgi_app(self, environ: WSGIEnviron, start_response: WSGIStartResponse) -> WSGIReturn:
        ctx = contextvars.Context()
        return ctx.run(self.original_wsgi_app, environ, start_response)

lib/dl_api_commons/dl_api_commons/flask/middlewares/context_var_middleware.py
datalens-tech/datalens-backend
from __future__ import annotations

import functools
from typing import (
    Any,
    Callable,
    TypeVar,
)

import flask

from dl_api_commons.flask.middlewares.context_var_middleware import ContextVarMiddleware
from dl_api_commons.logging import (
    CompositeLoggingContextController,
    LogRequestLoggingContextController,
)


class RequestLoggingContextControllerMiddleWare:
    _G_ATTR_NAME = "_bi_logging_ctx_controller"
    _APP_INIT_FLAG = "_bi_middleware_flag_req_log_ctx_ctrl"

    def set_up(self, app: flask.Flask) -> None:
        assert ContextVarMiddleware.is_wrapping_app(
            app
        ), "ContextVarMiddleware should be initialized before RequestLoggingContextControllerMiddleWare"
        app.before_request(self._bind_logging_request_controller_to_request)
        setattr(app, self._APP_INIT_FLAG, True)

    def _bind_logging_request_controller_to_request(self) -> None:
        setattr(
            flask.g,
            self._G_ATTR_NAME,
            CompositeLoggingContextController(
                [
                    LogRequestLoggingContextController(),
                ]
            ),
        )

    @classmethod
    def is_initialized_for_app(cls, app: flask.Flask) -> bool:
        return hasattr(app, cls._APP_INIT_FLAG)

    @classmethod
    def get_for_request(cls) -> CompositeLoggingContextController:
        if hasattr(flask.g, cls._G_ATTR_NAME):
            return getattr(flask.g, cls._G_ATTR_NAME)

        raise ValueError("Request logging context controller info was not created")


_RESULT_TV = TypeVar("_RESULT_TV")


def put_to_request_context(**ctx: Any) -> Callable[[Callable[..., _RESULT_TV]], Callable[..., _RESULT_TV]]:
    def real_decorator(wrapped: Callable[..., _RESULT_TV]) -> Callable[..., _RESULT_TV]:
        @functools.wraps(wrapped)
        def wrapper(*args: Any, **kwargs: Any) -> _RESULT_TV:
            for ctx_key, ctx_value in ctx.items():
                RequestLoggingContextControllerMiddleWare.get_for_request().put_to_context(ctx_key, ctx_value)
            return wrapped(*args, **kwargs)

        return wrapper

    return real_decorator

lib/dl_api_commons/dl_api_commons/flask/middlewares/logging_context.py
datalens-tech/datalens-backend
from __future__ import annotations

import json
import logging
import time
from typing import Optional

import attr
import flask
from flask import request

from dl_api_commons import (
    log_request_start,
    make_uuid_from_parts,
    request_id_generator,
)
from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.flask.middlewares.commit_rci_middleware import ReqCtxInfoMiddleware
from dl_api_commons.flask.middlewares.logging_context import RequestLoggingContextControllerMiddleWare
from dl_api_commons.headers import (
    HEADER_DEBUG_MODE_ENABLED,
    HEADER_LOGGING_CONTEXT,
    get_x_dl_context,
)
from dl_api_commons.logging import (
    NON_TRANSITIVE_LOGGING_CTX_KEYS,
    log_request_end_extended,
)
from dl_app_tools.log.context import put_to_context
from dl_constants.api_constants import DLHeadersCommon


LOGGER = logging.getLogger(__name__)


@attr.s
class RequestIDService:
    _request_id_app_prefix: Optional[str] = attr.ib()
    _append_local_req_id: bool = attr.ib(default=True)
    _accept_logging_context: bool = attr.ib(default=False)
    _logging_ctx_header_name: str = attr.ib(default=HEADER_LOGGING_CONTEXT)

    def _before_request(self) -> None:
        flask.g.request_start_time = time.monotonic()

        incoming_req_id = flask.request.headers.get(DLHeadersCommon.REQUEST_ID.value)
        current_req_id: str

        if self._append_local_req_id:
            current_req_id = make_uuid_from_parts(
                current=request_id_generator(self._request_id_app_prefix),
                parent=incoming_req_id,
            )
        else:
            current_req_id = incoming_req_id or request_id_generator(self._request_id_app_prefix)

        req_log_ctx_ctrl = RequestLoggingContextControllerMiddleWare.get_for_request()
        req_log_ctx_ctrl.put_to_context("request_id", current_req_id)
        req_log_ctx_ctrl.put_to_context("parent_request_id", incoming_req_id)

        # Updating logging context with outer values
        logging_ctx_header = request.headers.get(self._logging_ctx_header_name)
        if self._accept_logging_context and logging_ctx_header:
            try:
                logging_ctx_from_header = json.loads(logging_ctx_header)
                for ctx_key in NON_TRANSITIVE_LOGGING_CTX_KEYS:
                    logging_ctx_from_header.pop(ctx_key, None)

                for ctx_key, ctx_val in logging_ctx_from_header.items():
                    put_to_context(ctx_key, ctx_val)

            except Exception:  # noqa
                LOGGER.exception("Can not parse logging context: %s", logging_ctx_header)

        log_request_start(
            logger=LOGGER,
            method=request.method,
            full_path=request.full_path,
            headers=request.headers.items(),
        )

        # Initially setup RCI for request
        ReqCtxInfoMiddleware.replace_temp_rci(
            RequestContextInfo.create(
                request_id=current_req_id,
                x_dl_debug_mode=bool(int(request.headers.get(HEADER_DEBUG_MODE_ENABLED, "0"))),
                endpoint_code=None,
                x_dl_context=get_x_dl_context(request.headers.get(DLHeadersCommon.DL_CONTEXT, "{}")),
                # This props will be filled in commit_rci middleware
                plain_headers=None,
                secret_headers=None,
                # This props will be filled in auth middleware
                tenant=None,
                user_id=None,
                user_name=None,
                auth_data=None,
            )
        )

    def set_up(self, app: flask.Flask) -> None:
        assert RequestLoggingContextControllerMiddleWare.is_initialized_for_app(
            app
        ), "RequestLoggingContextControllerMiddleWare should be setup before request_id"
        app.before_request(self._before_request)
        app.after_request(_set_request_id_response_header)
        app.after_request(_post_log_request_id)


def _post_log_request_id(response: flask.Response) -> flask.Response:
    if response:
        response_timing = None
        request_start_time = getattr(flask.g, "request_start_time", None)
        if request_start_time is not None:
            response_timing = time.monotonic() - request_start_time

        rci = ReqCtxInfoMiddleware.get_last_resort_rci()
        user_name: Optional[str] = None
        user_id: Optional[str] = None
        if rci is not None:
            user_name = rci.user_name
            user_id = rci.user_id

        log_request_end_extended(
            logger=LOGGER,
            request_method=request.method,
            request_path=request.full_path,
            request_headers=dict(request.headers),
            response_status=response.status_code,
            response_headers=dict(response.headers),
            response_timing=response_timing,
            user_id=user_id,
            username=user_name,
        )
        # LOGGER.info('Response content: %r', response.get_json(silent=True))
    return response


def _set_request_id_response_header(response: flask.Response) -> flask.Response:
    if response:
        rci = ReqCtxInfoMiddleware.get_last_resort_rci()
        if rci is not None and rci.request_id is not None:
            response.headers.set("x-request-id", rci.request_id)
    return response

lib/dl_api_commons/dl_api_commons/flask/middlewares/request_id.py
datalens-tech/datalens-backend
"""
Middleware to extract OpenTracing context from HTTP headers and creating root request span.
Should be placed after ContextVarMiddleware
"""
from __future__ import annotations

import contextlib
import functools
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Sequence,
    Tuple,
)

import attr
import flask
import opentracing
import opentracing.tags

from dl_api_commons.flask.middlewares.logging_context import RequestLoggingContextControllerMiddleWare
from dl_api_commons.flask.middlewares.wsgi_middleware import FlaskWSGIMiddleware
from dl_api_commons.logging_tracing import OpenTracingLoggingContextController


if TYPE_CHECKING:
    from dl_api_commons.flask.types import (
        WSGIEnviron,
        WSGIReturn,
        WSGIStartResponse,
    )

LOGGER = logging.getLogger(__name__)


@attr.s
class TracingMiddleware(FlaskWSGIMiddleware):
    _APP_FLAG_ATTR_NAME = "_bi_middleware_flag_open_tracing"

    url_prefix_exclude: Sequence[str] = attr.ib(default=())

    @staticmethod
    def normalize_headers(wsgi_environ: Dict[str, str]) -> Dict[str, str]:
        prefix = "HTTP_"
        headers = {
            key[len(prefix) :].replace("_", "-").lower(): val
            for (key, val) in wsgi_environ.items()
            if key.startswith(prefix)
        }
        return headers

    def should_create_root_span(self, http_path: str) -> bool:
        for prefix in self.url_prefix_exclude:
            if http_path.startswith(prefix):
                return False
        return True

    def wsgi_app(self, environ: WSGIEnviron, start_response: WSGIStartResponse) -> WSGIReturn:  # type: ignore
        http_path = environ["PATH_INFO"]
        http_method = environ["REQUEST_METHOD"]
        tracer = opentracing.global_tracer()

        span_context = tracer.extract(opentracing.Format.HTTP_HEADERS, self.normalize_headers(environ))
        default_operation_name = "inbound_request"

        effective_start_response: WSGIStartResponse = start_response

        with contextlib.ExitStack() as cm_stack:
            if self.should_create_root_span(http_path=http_path):
                scope = cm_stack.enter_context(tracer.start_active_span(default_operation_name, child_of=span_context))
                scope.span.set_tag(opentracing.tags.SPAN_KIND, opentracing.tags.SPAN_KIND_RPC_SERVER)
                scope.span.set_tag(opentracing.tags.HTTP_URL, http_path)
                scope.span.set_tag(opentracing.tags.HTTP_METHOD, http_method)

                @functools.wraps(start_response)
                def start_response_set_span_tags(
                    status: str, response_headers: List[Tuple[str, str]], exc_info: Any = None
                ) -> Any:
                    try:
                        status_code = int(status.split(" ")[0])
                        scope.span.set_tag(opentracing.tags.HTTP_STATUS_CODE, status_code)
                        if status_code >= 500:
                            scope.span.set_tag(opentracing.tags.ERROR, True)
                    except Exception:  # noqa
                        LOGGER.exception("Exception during extracting tags from WSGI start_response args")

                    return start_response(status, response_headers, exc_info)

                effective_start_response = start_response_set_span_tags

            return self.original_wsgi_app(environ, effective_start_response)


class TracingContextMiddleware:
    def set_up(self, app: flask.Flask) -> None:
        app.before_request(self._bind_logging_request_controller_to_request)

    @classmethod
    def _bind_logging_request_controller_to_request(cls) -> None:
        tracer = opentracing.global_tracer()

        if tracer.active_span is not None:
            composite_ctrl = RequestLoggingContextControllerMiddleWare.get_for_request()
            composite_ctrl.add_sub_controller(OpenTracingLoggingContextController(tracer.active_span))

lib/dl_api_commons/dl_api_commons/flask/middlewares/tracing.py
datalens-tech/datalens-backend
"""
Usage example:

.. code-block:: python
    app = Flask(__name__)
    FlaskWSGIMiddleware().wrap_flask_app(app)
"""
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Optional,
)

import attr
import flask


if TYPE_CHECKING:
    from dl_api_commons.flask.types import (
        WSGICallable,
        WSGIEnviron,
        WSGIReturn,
        WSGIStartResponse,
    )


@attr.s
class FlaskWSGIMiddleware(metaclass=abc.ABCMeta):
    _APP_FLAG_ATTR_NAME: ClassVar[str]

    _app: Optional[flask.Flask] = attr.ib(init=False, default=None)
    _original_wsgi_app: Optional[WSGICallable] = attr.ib(init=False, default=None)

    @property
    def original_wsgi_app(self) -> WSGICallable:
        if self._original_wsgi_app is None:
            raise ValueError(f"Middleware {self} was not bind to application")

        return self._original_wsgi_app

    @abc.abstractmethod
    def wsgi_app(self, environ: WSGIEnviron, start_response: WSGIStartResponse) -> WSGIReturn:
        return self.original_wsgi_app(environ, start_response)

    def wrap_flask_app(self, app: flask.Flask) -> None:
        if self._app is not None:
            raise ValueError("Middleware already wrapping Flask app")

        self._app = app
        self._original_wsgi_app = app.wsgi_app
        app.wsgi_app = self.wsgi_app  # type: ignore
        setattr(app, self._APP_FLAG_ATTR_NAME, True)

    @classmethod
    def is_wrapping_app(cls, app: flask.Flask) -> bool:
        return hasattr(app, cls._APP_FLAG_ATTR_NAME)

lib/dl_api_commons/dl_api_commons/flask/middlewares/wsgi_middleware.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Callable,
    Dict,
    Iterable,
    List,
    Tuple,
)


WSGIEnviron = Dict[str, Any]
WSGIStartResponse = Callable[[str, List[Tuple[str, str]], Any], Any]
WSGIReturn = Iterable[bytes]

WSGICallable = Callable[[WSGIEnviron, WSGIStartResponse], WSGIReturn]

lib/dl_api_commons/dl_api_commons/flask/types.py
datalens-tech/datalens-backend
from __future__ import annotations

import json
import logging
from typing import Sequence

from dl_constants.api_constants import DLHeadersCommon


LOGGER = logging.getLogger(__name__)


HEADER_LOGGING_CONTEXT = "X-DL-LoggingContext"
HEADER_DEBUG_MODE_ENABLED = "X-DL-Debug-Mode"
INTERNAL_HEADER_PROFILING_STACK = "X-DL-Profiling-Stack"

# Headers that will be stored to request context info
DEFAULT_RCI_PLAIN_HEADERS = (
    DLHeadersCommon.FORWARDED_FOR.value,
    "User-Agent",
    "X-DL-Folder-ID",
    DLHeadersCommon.TENANT_ID.value,
    "Referer",
    "X-Chart-Id",
    "X-Dash-Id",
    "uber-trace-id",
    "Host",
    DLHeadersCommon.ACCEPT_LANGUAGE.value,
    # TODO FIX: Move to RCI attributes instead of headers passing
    DLHeadersCommon.SUDO.value,
    DLHeadersCommon.ALLOW_SUPERUSER.value,
)
DEFAULT_RCI_SECRET_HEADERS = (
    "Authorization",
    "Cookie",
    # TODO: BI-4918 move to local injection and reuse bi_api_commons_ya_cloud.constants.DLHeadersYC
    DLHeadersCommon.IAM_TOKEN.value,  # TODO add injection of extra secret headers to app factories
)


def normalize_header_name(hdr: str) -> str:
    return hdr.lower().replace("_", "-")


# TODO FIX: Add headers normalization function and use it here and in RCI handling middleware
def append_extra_headers_and_normalize(default: Sequence[str], extra: Sequence[str]) -> tuple[str, ...]:
    """Chains default and extra with CI deduplication. Translating result to lower case."""
    default_norm = tuple(normalize_header_name(hdr) for hdr in default)
    extra_norm = tuple(normalize_header_name(hdr) for hdr in extra)
    default_set = set(default_norm)
    return default_norm + tuple(hdr for hdr in extra_norm if hdr not in default_set)


def get_x_dl_context(header_value: str) -> dict:
    try:
        x_dl_context = json.loads(header_value)
    except (json.decoder.JSONDecodeError, TypeError):
        x_dl_context = {}
        LOGGER.warning(f"Got a malformed {DLHeadersCommon.DL_CONTEXT} header: {header_value!r}")

    x_dl_context_filtered = {k: v for k, v in x_dl_context.items() if isinstance(v, str)}

    if x_dl_context_filtered != x_dl_context:
        LOGGER.warning(f"Got unexpected data structure in {DLHeadersCommon.DL_CONTEXT} header: {header_value!r}")

    return x_dl_context_filtered

lib/dl_api_commons/dl_api_commons/headers.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from http.cookies import SimpleCookie
import logging
import os
import re
from typing import (
    Any,
    ClassVar,
    Iterable,
    Optional,
    Pattern,
    Sequence,
)

import attr

from dl_api_commons.headers import normalize_header_name
from dl_app_tools.log import context


LOGGER = logging.getLogger(__name__)

NON_TRANSITIVE_LOGGING_CTX_KEYS = frozenset(
    {
        "request_id",
        "pid",
    }
)

SECRET_HEADERS: frozenset[str] = frozenset(
    h.lower()
    for h in (
        "Authorization",
        "X-Us-Master-Token",
        "Master-Token",
        "X-DL-API-Key",
    )
)

SECRET_HEADERS_PATTERNS: Sequence[Pattern] = (re.compile(r".*token.*", re.IGNORECASE),)

SECRET_COOKIES: frozenset[str] = frozenset(
    c.lower()
    for c in (
        "Session_id",
        "sessionid2",
        "yc_session",
        "iam_cookie",
    )
)


def _is_secret_header(header_name: str) -> bool:
    if header_name.lower() in SECRET_HEADERS:
        return True

    for pattern in SECRET_HEADERS_PATTERNS:
        if pattern.match(header_name):
            return True

    return False


def _obfuscate_value(secret_value: str) -> str:
    repl_str = "<hidden>"
    if len(secret_value) > 8:
        return secret_value[:3] + repl_str + secret_value[-3:]
    else:
        return repl_str


def _obfuscate_cookie_header_value(cookie_string: str) -> str:
    cookie: SimpleCookie = SimpleCookie(cookie_string)
    for cookie_name in cookie:
        if cookie_name.lower() in SECRET_COOKIES:
            repl_str = "<hidden>"
            cookie[cookie_name].set(cookie_name, repl_str, repl_str)

    return cookie.output(header="", sep=";").strip()


def clean_secret_data_in_headers(headers: Iterable[tuple[str, str]]) -> Iterable[tuple[str, str]]:
    return tuple(
        (
            name,
            _obfuscate_value(val)
            if _is_secret_header(name)
            else _obfuscate_cookie_header_value(val)
            if name.lower() == "cookie"
            else val,
        )
        for name, val in headers
    )


def log_request_start(logger: logging.Logger, method: str, full_path: str, headers: Iterable[tuple[str, str]]) -> None:
    clean_headers = clean_secret_data_in_headers(headers)

    logger.info(
        "Received request. method: %s, path: %s, headers: %s, pid: %s",
        method.upper(),
        full_path,
        # Complexity to be compatible with previous version of logger
        "{{{}}}".format(", ".join([f"{k!r}: {v!r}" for k, v in clean_headers])),
        os.getpid(),
    )


def log_request_end(logger: logging.Logger, method: str, full_path: str, status_code: int) -> None:
    logger.info(
        "Response. method: %s, path: %s, status: %d",
        method.upper(),
        full_path,
        status_code,
        extra=dict(
            event_code="http_response",
            request_method=method,
            request_path=full_path,
            response_status=status_code,
        ),
    )


# TODO FIX: Make more strict typing for headers
def _normalize_headers(headers: Any) -> Optional[dict[str, str]]:
    if headers is None:
        return None
    if hasattr(headers, "items"):
        headers = headers.items()
    headers = sorted(headers)
    headers = {normalize_header_name(key): value for key, value in headers}

    headers = dict(clean_secret_data_in_headers(headers.items()))

    return headers


# TODO CONSIDER: Create custom type for headers
def log_request_end_extended(
    logger: logging.Logger,
    request_method: str,
    request_path: str,
    request_headers: Optional[dict],
    response_status: int,
    response_headers: Optional[dict],
    response_timing: Optional[float],
    # ...
    user_id: Optional[str] = None,
    username: Optional[str] = None,
    # extra extra for when they're not in the context.
    # TODO: tenant_id
    request_id: Optional[str] = None,
    endpoint_code: Optional[str] = None,
) -> None:
    """
    Response pre-return detailed (extended) logging.

    :param logger: logger to send the log to.

    :param request_method: HTTP method; uppercase, e.g. 'GET'.

    :param request_path: HTTP path; should, generally, include the query string,
        i.e. '/a/b?c=d&e=f'.

    :param request_headers: HTTP headers, unordered, normalized (lowercase
        dash names, i.e. 'user-agent'), unduplicated, scrubbed.
        Headers like 'remote-addr' should be authoritative.

    :param user_id: primary identifier of the *authorized* initial user (authoritative);
        e.g. '1120000000092758'.

    :param username: readable identifier of the user, i.e. login, e.g. 'robot-datalens'.

    :param response_status: HTTP status of the response; e.g. '502'.

    :param response_headers: HTTP headers of the response, normalized.

    :param response_timing: duration of the request handling, in seconds.

    :param request_id: unique request identifier; should be specified when not
        available in context.

    :param endpoint_code: name of the handler; should be specified when not
        available in context; use an empty string if it is not applicable.
    """
    request_method = request_method.upper()
    request_headers = _normalize_headers(request_headers)
    response_headers = _normalize_headers(response_headers)
    if response_timing is not None:
        response_timing = round(response_timing, 4)

    extra = dict(
        event_code="http_response",
        request_method=request_method,
        request_path=request_path,
        request_headers=request_headers,
        user_id=user_id,
        username=username,
        response_status=response_status,
        response_headers=response_headers,
        response_timing=response_timing,
        # Other possibilities:
        # response_body_info=dict(body_piece=body[:max_size], body_size=body_size, ...),
        # response_details=dict(...),
    )
    if request_id is not None:
        extra["request_id"] = request_id
    if endpoint_code is not None:
        extra["endpoint_code"] = endpoint_code

    logger.info(
        "Response. method: %s, path: %s, status: %d", request_method, request_path, response_status, extra=extra
    )


class RequestLoggingContextController(metaclass=abc.ABCMeta):
    """
    Just an interface to mutate logging context (python logging, sentry, etc...)
    At this moment there is only one implementation: for sentry context.
    In near future it should completely replace log.context.log_context(...)
    """

    @abc.abstractmethod
    def put_to_context(self, key: str, value: Any) -> None:
        """
        For flexibility - any keys for this moment. Later key set will be restricted
        """


@attr.s
class CompositeLoggingContextController(RequestLoggingContextController):
    _sub_controllers: list[RequestLoggingContextController] = attr.ib(factory=list)

    def put_to_context(self, key: str, value: Any) -> None:
        for ctrl in self._sub_controllers:
            ctrl.put_to_context(key, value)

    def add_sub_controller(self, ctrl: RequestLoggingContextController) -> None:
        self._sub_controllers.append(ctrl)


def extra_with_evt_code(event_code: str, extra: dict[str, Any]) -> dict[str, Any]:
    return dict(extra, event_code=event_code)


def _get_map_key_label(*args: str, **kwargs: str) -> dict[str, str]:
    map_field_name_label: dict[str, str] = {}
    for field_name in args:
        map_field_name_label[field_name] = field_name
    for label, field_name in kwargs.items():
        map_field_name_label[field_name] = label

    return map_field_name_label


def format_dict(extra: dict[str, Any], separator: str = " ", *args: str, **kwargs: str) -> str:
    map_extra_name_label = _get_map_key_label(*args, **kwargs)
    parts = []

    for extra_name, label in map_extra_name_label.items():
        if extra_name in extra:
            parts.append(f"{label}={repr(extra[extra_name])}")
        else:
            parts.append(f"{label}=N/A")
            LOGGER.warning(
                "Can not found extra key during message formatting: %s",
                extra_name,
                extra=extra_with_evt_code("logging_missing_extra_in_formatting", dict(extra_name=extra_name)),
            )

    return separator.join(parts)


def mask_sensitive_fields_by_name_in_json_recursive(
    source: Optional[dict[str, Any]], extra_sensitive_key_names: Iterable[str] = ()
) -> Optional[dict[str, Any]]:
    if source is None:
        return None

    all_sensitive_key_names: set[str] = {
        "password",
        "token",
        "secret",
        "private_key",
        "cypher_text",
    }
    all_sensitive_key_names.update(extra_sensitive_key_names)

    def process_value(key_name: Optional[str], value: Any) -> Any:
        if value is None:
            return None

        if isinstance(value, dict):
            return {nested_key: process_value(nested_key, nested_value) for nested_key, nested_value in value.items()}
        if isinstance(
            value,
            (
                list,
                tuple,
            ),
        ):
            return [process_value(key_name, nested_value) for nested_value in value]
        if isinstance(value, str):
            if key_name in all_sensitive_key_names:
                return _obfuscate_value(value)
            return value

        if isinstance(
            value,
            (
                bool,
                float,
                int,
            ),
        ):
            if key_name in all_sensitive_key_names:
                LOGGER.error("Non-string type for sensitive field '%s': %s", key_name, type(value))
                return _obfuscate_value(str(value))
            return value

        raise TypeError(f"Unexpected value type: {type}")

    return process_value(None, source)


class LogRequestLoggingContextController(RequestLoggingContextController):
    allowed_keys: ClassVar[tuple[str, ...]] = (
        "request_id",
        "parent_request_id",
        "endpoint_code",
        "user_id",
        "folder_id",
        "project_id",
        "org_id",
    )

    def put_to_context(self, key: str, value: Any) -> None:
        if key in self.allowed_keys:
            # Each request assumed to be executed in individual ContextVars context so we don't need to pop it back
            # see `bi_core.flask_utils.context_var_middleware`
            context.put_to_context(key, value)

lib/dl_api_commons/dl_api_commons/logging.py
datalens-tech/datalens-backend
from logging import LogRecord

import dl_app_tools.log.context as log_context


def add_log_context(record: LogRecord) -> None:
    context = log_context.get_log_context()
    for key, value in context.items():
        setattr(record, key, value)

lib/dl_api_commons/dl_api_commons/logging_config.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import re
from typing import Any

import attr

from dl_api_commons import clean_secret_data_in_headers


log = logging.getLogger()

SECRET_VAR_RE = re.compile(
    "(?i)^.*("
    "token"
    "|"
    "fernet"
    "|"
    "password"
    "|"
    "pwd"
    "|"
    "secret"
    "|"
    "cypher"
    "|"
    "master_key"
    "|"
    "api_key"
    ").*$"
)

SECRET_VAR_CONTENT_RE = re.compile(
    "(?i).*("
    "BEGIN PRIVATE KEY"
    "|"
    "yc_session"
    "|"
    "CheckSessionRequest"  # GRPC request class that contain in repr secret cookies
    ")"
)

S3_TBL_FUNC_RE = re.compile("s3\([^,]*['\"]http([^,]+,){3}")  # noqa


def is_secret_var(var_name: str) -> bool:
    return bool(SECRET_VAR_RE.match(var_name))


def get_content_repr(content: Any) -> str:
    if isinstance(content, str):
        return content
    return repr(content)


def is_secret_content(content: Any) -> bool:
    if content is None:
        return False
    content_repr = get_content_repr(content)
    return bool(SECRET_VAR_CONTENT_RE.match(content_repr))


def contains_s3_tbl_func(content: Any) -> bool:
    if content is None:
        return False
    content_repr = get_content_repr(content)
    return bool(S3_TBL_FUNC_RE.search(content_repr))


def mask_s3_tbl_func(content: Any) -> str:
    content_repr = get_content_repr(content)
    while match := S3_TBL_FUNC_RE.search(content_repr):
        content_repr = content_repr[: match.start() + 3] + "<hidden>" + content_repr[match.end() :]
    return content_repr


def cleanup_local_vars(local_vars: dict) -> None:
    """Mutates local vars dict inplace"""
    for name, val in local_vars.items():
        if is_secret_var(name) or is_secret_content(val):
            local_vars[name] = "[hidden]"
        elif contains_s3_tbl_func(val):
            local_vars[name] = mask_s3_tbl_func(val)


def cleanup_event_headers(original_headers: dict[str, str]) -> dict[str, str]:
    return {
        name: value
        for name, value in clean_secret_data_in_headers(
            (
                original_name,
                original_value,
            )
            for original_name, original_value in original_headers.items()
        )
    }


def cleanup_event_request_section(req_section: dict[str, str | dict[str, str]]) -> dict[str, str | dict[str, str]]:
    clean_req_section = dict(req_section)

    secret_original_headers = req_section.get("headers")

    if secret_original_headers is None:
        pass
    elif isinstance(secret_original_headers, dict):
        clean_req_section["headers"] = cleanup_event_headers(secret_original_headers)
    else:
        log.error(f"Unexpected type of request headers section in outgoing Sentry event: {type(req_section)}")

    return clean_req_section


def cleanup_common_secret_data(
    event: dict,
    hint: dict,  # noqa
) -> dict:
    for exc_data in event.get("exception", {}).get("values"):
        for frame in exc_data.get("stacktrace", {}).get("frames", ()):
            local_vars = frame.get("vars", {})
            cleanup_local_vars(local_vars)

    secret_original_req_section = event.get("request")

    if secret_original_req_section is None:
        pass
    elif isinstance(secret_original_req_section, dict):
        event["request"] = cleanup_event_request_section(secret_original_req_section)
    else:
        log.error(f"Unexpected type of request section in outgoing Sentry event: {type(secret_original_req_section)}")

    return event

lib/dl_api_commons/dl_api_commons/logging_sentry.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Any,
    ClassVar,
    Sequence,
)

import attr
import opentracing

from dl_api_commons.logging import RequestLoggingContextController


LOGGER = logging.getLogger(__name__)


@attr.s
class OpenTracingLoggingContextController(RequestLoggingContextController):
    root_span: opentracing.Span = attr.ib()

    root_span_tags: ClassVar[Sequence[str]] = (
        "dataset_id",
        "conn_id",
        "conn_type",
        "request_id",
        "parent_request_id",
    )
    span_tag_prefix: ClassVar[str] = "bi"

    def put_to_context(self, key: str, value: Any) -> None:
        if key in self.root_span_tags:
            self.root_span.set_tag(f"{self.span_tag_prefix}.{key}", value)
        elif key == "endpoint_code":
            if isinstance(value, str):
                self.root_span.set_operation_name(value)
            else:
                LOGGER.warning("Got not-string value for root span operation name : %s", value)
        elif key == "is_error":
            self.root_span.set_tag(opentracing.span.tags.ERROR, bool(value))

lib/dl_api_commons/dl_api_commons/logging_tracing.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    Tuple,
)

import attr

from dl_api_commons.reporting.records import ReportingRecord
from dl_constants.enums import (
    ConnectionType,
    NotificationLevel,
    ReportingQueryType,
)


@attr.s(frozen=True, auto_attribs=True)
class QueryExecutionReportingRecord(ReportingRecord):
    query_id: str


@attr.s(frozen=True, auto_attribs=True)
class QueryExecutionStartReportingRecord(QueryExecutionReportingRecord):
    dataset_id: Optional[str]
    query_type: Optional[ReportingQueryType]
    connection_type: ConnectionType
    conn_reporting_data: dict
    query: str  # SQL query
    workbook_id: Optional[str]


@attr.s(frozen=True, auto_attribs=True)
class QueryExecutionCacheInfoReportingRecord(QueryExecutionReportingRecord):
    # Note: using it as ternary: hit / miss / undefined (error/disabled/...)
    # TODO?: write a bi_core.data_processing.cache_processing.CacheSituation name too?
    cache_full_hit: Optional[bool] = None


@attr.s(frozen=True, auto_attribs=True)
class QueryExecutionEndReportingRecord(QueryExecutionReportingRecord):
    exception: Optional[BaseException]


@attr.s(frozen=True, auto_attribs=True)
class DataProcessingReportingRecord(ReportingRecord):
    processing_id: str


@attr.s(frozen=True, auto_attribs=True)
class DataProcessingStartReportingRecord(DataProcessingReportingRecord):
    source_query_ids: Tuple[str, ...]


@attr.s(frozen=True, auto_attribs=True)
class DataProcessingEndReportingRecord(DataProcessingReportingRecord):
    exception: Optional[Exception]


@attr.s(frozen=True, auto_attribs=True)
class DataProcessingCacheInfoReportingRecord(DataProcessingReportingRecord):
    cache_full_hit: Optional[bool] = None


@attr.s(frozen=True, auto_attribs=True)
class NotificationReportingRecord(ReportingRecord):
    title: str
    message: str
    level: NotificationLevel
    locator: str

lib/dl_api_commons/dl_api_commons/reporting/models.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
import time
import traceback
from typing import (
    TYPE_CHECKING,
    Dict,
    Optional,
    Tuple,
    Type,
    TypeVar,
    Union,
)

import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.reporting.models import (
    DataProcessingCacheInfoReportingRecord,
    DataProcessingEndReportingRecord,
    DataProcessingReportingRecord,
    DataProcessingStartReportingRecord,
    QueryExecutionCacheInfoReportingRecord,
    QueryExecutionEndReportingRecord,
    QueryExecutionReportingRecord,
    QueryExecutionStartReportingRecord,
)
from dl_api_commons.reporting.records import ReportingRecord
from dl_constants.api_constants import DLContextKey


if TYPE_CHECKING:
    from dl_api_commons.reporting.registry import ReportingRegistry


PROFILING_LOG_NAME = "bi_core.profiling_db"
QUERY_PROFILING_ENTRY = "dump request profile"


_RECORD_TV = TypeVar("_RECORD_TV", bound=ReportingRecord)


class ReportingProfiler(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def on_request_end(self) -> None:
        pass


@attr.s
class DefaultReportingProfiler(ReportingProfiler):
    _reporting_registry: ReportingRegistry = attr.ib(default=None)
    _is_public_env: bool = attr.ib(default=False)

    @property
    def reporting_registry(self) -> ReportingRegistry:
        return self._reporting_registry

    @property
    def rci(self) -> RequestContextInfo:
        return self._reporting_registry.rci

    def get_profiling_logger(self) -> logging.Logger:
        return logging.getLogger(PROFILING_LOG_NAME)

    def _find_record_of_type(
        self,
        records: Tuple[ReportingRecord, ...],
        rtype: Type[_RECORD_TV],
    ) -> Optional[_RECORD_TV]:
        return next((r for r in records if isinstance(r, rtype)), None)

    def get_selector_records_for_query(self, query_id: str) -> Tuple[QueryExecutionReportingRecord, ...]:
        records = self.reporting_registry.get_reporting_records()
        return tuple(r for r in records if isinstance(r, QueryExecutionReportingRecord) and r.query_id == query_id)

    def get_data_proc_records_for_query(self, query_id: str) -> Tuple[DataProcessingReportingRecord, ...]:
        records = self.reporting_registry.get_reporting_records()
        proc_ids = {
            r.processing_id
            for r in records
            if isinstance(r, DataProcessingStartReportingRecord) and query_id in r.source_query_ids
        }
        return tuple(r for r in records if isinstance(r, DataProcessingReportingRecord) and r.processing_id in proc_ids)

    def flush_query_report(self, query_id: str) -> None:
        request_result_records = self.reporting_registry.get_request_result_reporting_records()
        query_records = self.get_selector_records_for_query(query_id=query_id)
        data_proc_records = self.get_data_proc_records_for_query(query_id=query_id)

        start_record = self._find_record_of_type(query_records, rtype=QueryExecutionStartReportingRecord)
        assert start_record is not None, "need a matching `QueryExecutionStartReportingRecord`"
        end_selector_record = self._find_record_of_type(query_records, rtype=QueryExecutionEndReportingRecord)
        end_data_proc_record = self._find_record_of_type(query_records, rtype=DataProcessingEndReportingRecord)
        cache_selector_record = self._find_record_of_type(query_records, rtype=QueryExecutionCacheInfoReportingRecord)
        cache_data_proc_record = self._find_record_of_type(
            data_proc_records, rtype=DataProcessingCacheInfoReportingRecord
        )
        response_status_code, err_code = None, None
        if request_result_records:
            response_status_code = request_result_records[0].response_status_code
            err_code = request_result_records[0].err_code

        cache_used = cache_selector_record is not None or cache_data_proc_record is not None
        cache_full_hit = any(
            cache_record.cache_full_hit
            for cache_record in (cache_selector_record, cache_data_proc_record)
            if cache_record is not None
        )

        error: Optional[BaseException] = None
        end_record = end_data_proc_record if end_data_proc_record is not None else end_selector_record
        if end_record is not None:
            error = end_record.exception
            end_timestamp = end_record.timestamp
        else:
            end_timestamp = time.time()

        error_text: Optional[str]
        if error is None:
            error_text = None
        else:
            error_text = traceback.format_exception_only(type(error), error)[-1]

        x_dl_context = self.rci.x_dl_context
        extra: Dict[str, Union[str, int, float, None]] = dict(
            event_code="profile_db_request",
            dataset_id=start_record.dataset_id,
            # auto: request_id
            user_id=self.rci.user_id,
            billing_folder_id=self.rci.tenant.get_tenant_id() if self.rci.tenant is not None else None,
            connection_type=start_record.connection_type.name,
            source=self.rci.endpoint_code,
            username=self.rci.user_name,
            execution_time=int(round((end_timestamp - start_record.timestamp) * 1000)),
            query=start_record.query,
            status="success" if error is None else "error",
            error=error_text,
            err_code=err_code,
            cache_used=cache_used,
            cache_full_hit=cache_full_hit,
            # not in action fields:
            query_type=(start_record.query_type.name if start_record.query_type is not None else None),
            # TODO FIX: Use value from RCI when become done
            is_public=int(self._is_public_env),
            dash_id=x_dl_context.get(DLContextKey.DASH_ID),
            dash_tab_id=x_dl_context.get(DLContextKey.DASH_TAB_ID),
            chart_id=x_dl_context.get(DLContextKey.CHART_ID),
            chart_kind=x_dl_context.get(DLContextKey.CHART_KIND),
            response_status_code=response_status_code,
            workbook_id=start_record.workbook_id,
            **start_record.conn_reporting_data,
        )
        # TODO FIX: Change logger
        log = self.get_profiling_logger()
        log.info(QUERY_PROFILING_ENTRY, extra=extra)

    def flush_all_query_reports(self) -> None:
        all_query_ids = sorted(
            {
                r.query_id
                for r in self.reporting_registry.get_reporting_records()
                if isinstance(r, QueryExecutionReportingRecord)
            }
        )
        for query_id in all_query_ids:
            self.flush_query_report(query_id=query_id)

    def on_request_end(self) -> None:
        self.flush_all_query_reports()

lib/dl_api_commons/dl_api_commons/reporting/profiler.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Optional,
    TypeVar,
)

import attr


_RR_TV = TypeVar("_RR_TV", bound="ReportingRecord")


@attr.s(frozen=True, auto_attribs=True)
class ReportingRecord:
    timestamp: float

    def clone(self: _RR_TV, **kwargs: Any) -> _RR_TV:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True, auto_attribs=True)
class RequestResultReportingRecord(ReportingRecord):
    response_status_code: Optional[int] = 200
    err_code: Optional[str] = None

lib/dl_api_commons/dl_api_commons/reporting/records.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    List,
    Optional,
    Tuple,
    Type,
    TypeVar,
)

import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.reporting.records import (
    ReportingRecord,
    RequestResultReportingRecord,
)


_RECORD_TV = TypeVar("_RECORD_TV", bound=ReportingRecord)


@attr.s
class ReportingRegistry(metaclass=abc.ABCMeta):
    _rci: RequestContextInfo = attr.ib()

    @property
    def rci(self) -> RequestContextInfo:
        return self._rci

    @rci.setter
    def rci(self, rci: RequestContextInfo) -> None:
        self._rci = rci

    @abc.abstractmethod
    def clear_records(self) -> None:
        pass

    @abc.abstractmethod
    def save_reporting_record(self, report: Optional[ReportingRecord]) -> None:
        pass

    @abc.abstractmethod
    def get_reporting_records(self) -> Tuple[ReportingRecord, ...]:
        pass

    @abc.abstractmethod
    def get_request_result_reporting_records(self) -> Tuple[RequestResultReportingRecord, ...]:
        pass

    @abc.abstractmethod
    def get_records_of_type(self, clz: Type[_RECORD_TV]) -> Tuple[_RECORD_TV, ...]:
        pass


@attr.s
class DefaultReportingRegistry(ReportingRegistry):
    _reporting_records: List[ReportingRecord] = attr.ib(init=False, factory=list)

    def clear_records(self) -> None:
        self._reporting_records.clear()

    def save_reporting_record(self, report: Optional[ReportingRecord]) -> None:
        if report is not None:
            self._reporting_records.append(report)

    def get_reporting_records(self) -> Tuple[ReportingRecord, ...]:
        return tuple(self._reporting_records)

    def get_request_result_reporting_records(self) -> Tuple[RequestResultReportingRecord, ...]:
        return tuple(r for r in self._reporting_records if isinstance(r, RequestResultReportingRecord))

    def get_records_of_type(self, clz: Type[_RECORD_TV]) -> Tuple[_RECORD_TV, ...]:
        return tuple(r for r in self._reporting_records if isinstance(r, clz))

lib/dl_api_commons/dl_api_commons/reporting/registry.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional
import uuid


def make_uuid_from_parts(current: str, parent: Optional[str]) -> str:
    assert current

    if not parent:
        return current

    uuid_maxlen = 120
    uuid_sep = "--"  # need to be a non-word character, to get successfully separated by elasticsearch
    cutted_replace = "..."

    if len(parent) + len(current) > uuid_maxlen:
        cutted_half_len = int((uuid_maxlen - len(current) - len(cutted_replace)) / 2)
        cutted_parent = parent[:cutted_half_len] + cutted_replace + parent[-cutted_half_len:]
        parent = cutted_parent

    result = parent + uuid_sep + current
    return result


def request_id_generator(prefix: Optional[str] = None) -> str:
    result = uuid.uuid4().hex
    if prefix is not None:
        result = prefix + "." + result
    return result

lib/dl_api_commons/dl_api_commons/request_id.py
datalens-tech/datalens-backend
import logging
from typing import (
    Any,
    Optional,
    Sequence,
)

import attr
import flask

from dl_api_commons.logging_sentry import cleanup_common_secret_data


@attr.s(frozen=True)
class SentryConfig:
    dsn: str = attr.ib(repr=False)
    release: Optional[str] = attr.ib(default=None)


def configure_sentry(cfg: SentryConfig, extra_integrations: Sequence[Any] = ()) -> None:
    import sentry_sdk
    from sentry_sdk.integrations.argv import ArgvIntegration
    from sentry_sdk.integrations.atexit import AtexitIntegration
    from sentry_sdk.integrations.excepthook import ExcepthookIntegration
    from sentry_sdk.integrations.logging import LoggingIntegration
    from sentry_sdk.integrations.modules import ModulesIntegration
    from sentry_sdk.integrations.stdlib import StdlibIntegration
    from sentry_sdk.integrations.threading import ThreadingIntegration

    # from sentry_sdk.integrations.logging import LoggingIntegration
    sentry_sdk.init(
        dsn=cfg.dsn,
        release=cfg.release,
        default_integrations=False,
        before_send=cleanup_common_secret_data,
        integrations=[
            # # Default
            AtexitIntegration(),
            ExcepthookIntegration(),
            StdlibIntegration(),
            ModulesIntegration(),
            ArgvIntegration(),
            LoggingIntegration(event_level=logging.WARNING),
            ThreadingIntegration(),
            #  # Custom
            *extra_integrations,
        ],
    )


def configure_sentry_for_aiohttp(cfg: SentryConfig) -> None:
    from sentry_sdk.integrations.aiohttp import AioHttpIntegration

    configure_sentry(cfg, extra_integrations=(AioHttpIntegration(),))


def configure_sentry_for_flask(cfg: SentryConfig) -> None:
    from sentry_sdk.integrations.flask import FlaskIntegration

    configure_sentry(cfg, extra_integrations=(FlaskIntegration(),))


def hook_configure_configure_sentry_for_flask(app: flask.Flask, cfg: SentryConfig) -> None:
    """
    Try to configure sentry in uwsgi `postfork` if possible,
    but ensure it is configured in `before_first_request` (flask app).
    """
    try:
        import uwsgidecorators  # type: ignore
    except Exception:  # noqa
        pass
    else:

        @uwsgidecorators.postfork
        def _init_logging_in_uwsgi_postfork():  # type: ignore
            configure_sentry_for_flask(cfg)

    @app.before_first_request
    def _init_logging_in_before_first_request():  # type: ignore
        configure_sentry_for_flask(cfg)

lib/dl_api_commons/dl_api_commons/sentry_config.py
datalens-tech/datalens-backend
import abc
from typing import Optional

from dl_api_commons.base_models import (
    TenantCommon,
    TenantDef,
)


class TenantResolver(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def resolve_tenant_def_by_tenant_id(self, tenant_id_str: Optional[str]) -> TenantDef:
        raise NotImplementedError()


class CommonTenantResolver(TenantResolver):
    def resolve_tenant_def_by_tenant_id(self, tenant_id_str: Optional[str]) -> TenantDef:
        return TenantCommon()

lib/dl_api_commons/dl_api_commons/tenant_resolver.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Dict,
    Optional,
)

import opentracing


def get_current_tracing_headers(tracer: Optional[opentracing.Tracer] = None) -> Dict[str, str]:
    actual_tracer = opentracing.global_tracer() if tracer is None else tracer
    tracing_headers: Dict[str, str] = {}
    active_span = actual_tracer.active_span

    if active_span is not None:
        actual_tracer.inject(active_span.context, opentracing.Format.HTTP_HEADERS, tracing_headers)

    return tracing_headers

lib/dl_api_commons/dl_api_commons/tracing.py
datalens-tech/datalens-backend
from dl_constants.api_constants import (
    DLCookies,
    DLHeaders,
)


def stringify_dl_headers(headers: dict[DLHeaders, str]) -> dict[str, str]:
    return {name.value.lower(): val for name, val in headers.items()}


def stringify_dl_cookies(headers: dict[DLCookies, str]) -> dict[str, str]:
    return {name.value: val for name, val in headers.items()}

lib/dl_api_commons/dl_api_commons/utils.py
datalens-tech/datalens-backend
from __future__ import annotations

from aiohttp import web
from multidict import CIMultiDict
import pytest

from dl_api_commons.aio.middlewares.commit_rci import commit_rci_middleware
from dl_api_commons.aio.middlewares.error_handling_outer import (
    AIOHTTPErrorHandler,
    ErrorData,
    ErrorLevel,
)
from dl_api_commons.aio.middlewares.request_bootstrap import RequestBootstrap
from dl_api_commons.aio.middlewares.request_id import RequestId
from dl_api_commons.aio.server_header import ServerHeader
from dl_api_commons.aiohttp.aiohttp_wrappers import DLRequestView


@pytest.mark.asyncio
async def test_service_header_normal_case(aiohttp_client):
    server_header = "Ololo"

    class SimpleView(web.View):
        async def get(self):
            return web.json_response({})

    class FailView(web.View):
        async def get(self):
            raise ValueError()

    class StreamedView(web.View):
        async def get(self):
            streamed_resp = web.StreamResponse()
            await streamed_resp.prepare(self.request)
            await streamed_resp.write(b"ololo")
            return streamed_resp

    app = web.Application()
    app.router.add_route("*", "/simple_view", SimpleView)
    app.router.add_route("*", "/fail_view", FailView)
    app.router.add_route("*", "/streamed_view", StreamedView)

    ServerHeader(server_header).add_signal_handlers(app)

    client = await aiohttp_client(app)

    resp = await client.get("/not_found")
    assert 404 == resp.status and [server_header] == resp.headers.getall("Server")

    resp = await client.get("/simple_view")
    assert 200 == resp.status and [server_header] == resp.headers.getall("Server")

    resp = await client.get("/fail_view")
    assert 500 == resp.status and [server_header] == resp.headers.getall("Server")

    resp = await client.get("/streamed_view")
    assert 200 == resp.status and [server_header] == resp.headers.getall("Server")


def test_service_header_validation_fail():
    with pytest.raises(ValueError):
        ServerHeader("")

    with pytest.raises(TypeError):
        ServerHeader(123)  # noqa


@pytest.mark.asyncio
async def test_commit_rci_middleware(caplog, aiohttp_client):
    caplog.set_level("INFO")

    app = web.Application(
        middlewares=[
            RequestBootstrap(
                req_id_service=RequestId(),
            ).middleware,
            commit_rci_middleware(
                rci_extra_plain_headers=("X-Add-Me", "x-add-me-2"), rci_extra_secret_headers=("X-Add-Me-Secret",)
            ),
        ]
    )

    class TestView(DLRequestView):
        async def get(self):
            rci = self.dl_request.rci
            return web.json_response(
                {
                    "secret_headers": [[k, v] for k, v in rci.secret_headers.items()],
                    "plain_headers": [[k, v] for k, v in rci.plain_headers.items()],
                    "req_id": rci.request_id,
                }
            )

    app.router.add_route("*", "/simple_view", TestView)
    client = await aiohttp_client(app)

    async def get_headers(inbound_headers):
        resp = await client.get("/simple_view", headers=inbound_headers, skip_auto_headers=("User-Agent",))
        try:
            assert resp.status == 200
            return await resp.json()
        finally:
            resp.close()

    actual_headers = await get_headers(
        CIMultiDict(
            (
                ("X-Add-Me", "plain_val"),
                ("X-Ignore-Me", "some"),
                ("X-Add-Me-Secret", "secret_val"),
                (
                    "X-Request-Id",
                    "reqid1",
                ),
                ("X-Chart-Id", "some_chart_id"),
                ("Host", "127.0.0.1"),
                ("Accept-Language", "en"),
            )
        )
    )

    assert actual_headers == {
        "secret_headers": [["X-Add-Me-Secret", "secret_val"]],
        "plain_headers": [
            ["Host", "127.0.0.1"],
            ["X-Add-Me", "plain_val"],
            ["X-Chart-Id", "some_chart_id"],
            ["Accept-Language", "en"],
        ],
        "req_id": "reqid1",
    }

    actual_headers = await get_headers(
        CIMultiDict(
            (
                ("x-add-me", "plain_val"),
                ("x-ignore-me", "some"),
                ("x-add-me-secret", "secret_val"),
                (
                    "x-request-id",
                    "reqid2",
                ),
                ("x-chart-id", "some_chart_id"),
                ("host", "127.0.0.1"),
                ("Accept-Language", "en"),
            )
        )
    )

    assert actual_headers == {
        "secret_headers": [["x-add-me-secret", "secret_val"]],
        "plain_headers": [
            ["Host", "127.0.0.1"],
            ["x-add-me", "plain_val"],
            ["x-chart-id", "some_chart_id"],
            ["Accept-Language", "en"],
        ],
        "req_id": "reqid2",
    }


@pytest.mark.asyncio
async def test_error_handling_middleware(caplog, aiohttp_client):
    caplog.set_level("INFO")
    log_name = "dl_api_commons.aio.middlewares.error_handling_outer"

    class MyException(Exception):
        def __init__(self, msg: str, level: str, status_code: int):
            super().__init__(msg)
            self.msg = msg
            self.level = level
            self.status_code = status_code

    class ErrorHandler(AIOHTTPErrorHandler):
        def _classify_error(self, err: Exception, request: web.Request) -> ErrorData:
            if isinstance(err, MyException):
                return ErrorData(
                    status_code=err.status_code,
                    level=ErrorLevel[err.level],
                    response_body=dict(msg=err.msg),
                )
            if isinstance(err, web.HTTPException):
                return ErrorData(
                    status_code=err.status_code,
                    level=ErrorLevel.info,
                    response_body={},
                    http_reason=err.reason,
                )
            else:
                raise ValueError("Can not format exception")

    async def info(request):
        raise MyException("info_exc", "info", 400)

    async def warn(request):
        raise MyException("warn_exc", "warning", 400)

    async def error(request):
        raise MyException("err_exc", "error", 500)

    async def fmt_error(request):
        raise ValueError()

    async def ok(request):
        return web.json_response({})

    app = web.Application(
        middlewares=[
            RequestBootstrap(
                req_id_service=RequestId(),
                error_handler=ErrorHandler(sentry_app_name_tag=None),
            ).middleware,
        ]
    )
    app.router.add_get("/ok", ok)
    app.router.add_get("/info", info)
    app.router.add_get("/warn", warn)
    app.router.add_get("/error", error)
    app.router.add_get("/fmt_error", fmt_error)

    client = await aiohttp_client(app)

    async def get_status_and_body(url) -> tuple[int, dict]:
        resp = await client.get(url)
        return resp.status, await resp.json()

    assert await get_status_and_body("/ok") == (200, {})

    caplog.clear()
    assert await get_status_and_body("/info") == (400, {"msg": "info_exc"})
    log_rec = [r for r in caplog.records if r.message == "Regular exception fired" and r.name == log_name][0]
    assert log_rec.levelname == "INFO"
    assert log_rec.exc_info

    caplog.clear()
    assert await get_status_and_body("/warn") == (400, {"msg": "warn_exc"})
    log_rec = [r for r in caplog.records if r.levelname == "WARNING" and r.name == log_name][0]
    assert log_rec.message == "Warning exception fired"
    assert log_rec.exc_info

    caplog.clear()
    assert await get_status_and_body("/error") == (500, {"msg": "err_exc"})
    log_rec = [r for r in caplog.records if r.levelname == "ERROR" and r.name == log_name][0]
    assert log_rec.message == "Caught an exception in request handler"
    assert log_rec.exc_info

    caplog.clear()
    assert await get_status_and_body("/fmt_error") == (500, {"message": "Internal Server Error"})
    fmt_log_rec = [r for r in caplog.records if r.levelname == "CRITICAL" and r.name == log_name][0]
    cause_log_rec = [r for r in caplog.records if r.levelname == "ERROR" and r.name == log_name][0]

    assert fmt_log_rec.message == "Error handler raised an error during creating error response"
    assert fmt_log_rec.exc_info

    assert cause_log_rec.message == "Error handler raised an error during handling this exception"
    assert cause_log_rec.exc_info
    assert cause_log_rec.exc_info != fmt_log_rec.exc_info

    caplog.clear()
    resp = await client.get("/not_found")
    resp_json = await resp.json()
    assert resp_json == {}
    assert resp.reason == "Not Found"
    log_rec = [r for r in caplog.records if r.message == "Regular exception fired" and r.name == log_name][0]
    assert log_rec.levelname == "INFO"
    assert log_rec.exc_info

lib/dl_api_commons/dl_api_commons_tests/unit/aio/test_aiohttp_services.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
from concurrent.futures.thread import ThreadPoolExecutor
import inspect
import logging
import random
import threading
import time
from typing import (
    Any,
    Callable,
    Generator,
    Union,
)

import attr
import pytest

from dl_api_commons.aio.async_wrapper_for_sync_generator import (
    EndOfStream,
    InitializationFailed,
    Job,
    JobState,
)


pytestmark = pytest.mark.asyncio  # TODO: should work through conftest?


@pytest.fixture()
def service_tpe() -> ThreadPoolExecutor:
    tpe = ThreadPoolExecutor(thread_name_prefix="SERVICE_TPE_")
    yield tpe
    tpe.shutdown()


@pytest.fixture()
def worker_tpe() -> ThreadPoolExecutor:
    tpe = ThreadPoolExecutor(thread_name_prefix="WORKER_TPE_")
    yield tpe
    tpe.shutdown()


@pytest.fixture()
def wrapper_factory(worker_tpe, service_tpe) -> Callable[[Generator], Job]:
    @attr.s
    class MyGenerator(Job):
        _generator: Union[Generator, Callable[[], Generator]] = attr.ib(default=None)

        def make_generator(self) -> Generator[Any, bool, None]:
            if inspect.isgenerator(self._generator):
                return self._generator
            else:
                return self._generator()

    return lambda gen: MyGenerator(
        service_tpe=service_tpe,
        workers_tpe=worker_tpe,
        generator=gen,
    )


async def async_gen_adapter(wr: Job):
    await wr.run()

    while True:
        try:
            yield await wr.get_next()
        except EndOfStream:
            return


async def test_simple_finite_generator(wrapper_factory, caplog):
    caplog.set_level("DEBUG")
    count = int(1e3)

    def test_generator():
        for i in range(count):
            yield i

    job = wrapper_factory(test_generator())

    result = []
    async for item in async_gen_adapter(job):
        result.append(item)

    assert result == list(range(count))
    assert job.state == JobState.closed
    await job.cancel()


async def test_error_in_start_generator(wrapper_factory, caplog):
    caplog.set_level("DEBUG")

    def test_generator():
        raise ValueError("Error in generator")
        # noqa
        yield

    job = wrapper_factory(test_generator())
    await job.run()

    with pytest.raises(ValueError, match=r"^Error in generator$"):
        await job.get_next()

    assert job.state == JobState.closed


async def test_error_in_generator_creation(wrapper_factory, caplog):
    caplog.set_level("DEBUG")

    class MarkerException(Exception):
        pass

    def broken_generator_factory():
        raise MarkerException

    job = wrapper_factory(lambda: broken_generator_factory())

    with pytest.raises(MarkerException):
        await job.run()


async def test_start_confirmation_timeout(caplog):
    caplog.set_level("DEBUG")

    local_worker_tpe = ThreadPoolExecutor(max_workers=1)
    local_service_tpe = ThreadPoolExecutor(max_workers=10)

    class LocalJob(Job):
        def make_generator(self) -> Generator[Any, None, None]:
            def g():
                yield None

            return g()

    loop = asyncio.get_running_loop()
    garbage_started = asyncio.Event()
    stop_garbage = threading.Event()

    def garbage():
        logging.info("Garbage thread was started")

        def cb():
            logging.info("Triggering event 'garbage_started'")
            garbage_started.set()

        loop.call_soon_threadsafe(cb)
        got_stop_event = stop_garbage.wait(timeout=15)
        if got_stop_event:
            logging.info("Stopping garbage thread due to stop event")
        else:
            logging.info("Stopping garbage thread due to timeout")
        return

    try:
        asyncio.ensure_future(asyncio.get_running_loop().run_in_executor(local_worker_tpe, garbage))
        await asyncio.wait_for(garbage_started.wait(), timeout=1)

        job = LocalJob(
            service_tpe=local_service_tpe,
            workers_tpe=local_worker_tpe,
        )

        with pytest.raises(InitializationFailed):
            await job.run()

        await job.cancel()

    finally:
        logging.info("Executing test clean-up")
        # Terminating garbage thread
        stop_garbage.set()
        local_worker_tpe.shutdown(wait=True)
        local_service_tpe.shutdown(wait=True)
        logging.info("Test clean-up done")


async def test_closing_before_full_consuming(wrapper_factory, caplog):
    caplog.set_level("INFO")

    def test_generator():
        try:
            while True:
                time.sleep(0.5)
                yield random.randint(0, int(1e6))
        finally:
            logging.info("Generator was correctly closed")

    job = wrapper_factory(test_generator())
    await job.run()

    await job.get_next()
    await job.get_next()

    await job.cancel()

    assert job.state == JobState.closed
    assert "Generator was correctly closed" in caplog.messages

lib/dl_api_commons/dl_api_commons_tests/unit/aio/test_async_wrapper_for_sync_generator.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import os
from typing import Any

from aiohttp import web
import flask
from multidict import CIMultiDict
import pytest

from dl_api_commons import (
    clean_secret_data_in_headers,
    log_request_end,
    log_request_start,
)
from dl_api_commons.aio.middlewares import request_id as aio_request_id
from dl_api_commons.aio.middlewares.request_bootstrap import RequestBootstrap
from dl_api_commons.flask.middlewares.context_var_middleware import ContextVarMiddleware
from dl_api_commons.flask.middlewares.logging_context import RequestLoggingContextControllerMiddleWare
from dl_api_commons.flask.middlewares.request_id import RequestIDService
from dl_api_commons.logging import mask_sensitive_fields_by_name_in_json_recursive


LOGGER = logging.getLogger(__name__)

_EXPECTED_MASKED_HEADERS = (
    ("SomeToken", "<hidden>"),
    ("some-token", "<hidden>"),
    ("some-tokeN", "<hidden>"),
    ("SomeToken1", "<hidden>"),
    ("SomeToken-1", "<hidden>"),
    ("x-us-master-token", "<hidden>"),
    ("X-US-Master-Token", "<hidden>"),
    ("x-us-master-token", "lon<hidden>890"),
    ("NonSecretHeader", "Ololo"),
    ("NonSecretHeader", "Azaza"),
    ("X-DL-API-Key", "123<hidden>789"),
    ("Cookie", "Session_id=<hidden>; _ym_isad=2; sessionid2=<hidden>"),
    ("Cookie", "_ym_isad=2; session_id=<hidden>; sessionid2=<hidden>"),
)

_INPUT_HEADERS = (
    ("SomeToken", "token"),
    ("some-token", "token"),
    ("some-tokeN", "token"),
    ("SomeToken1", "token"),
    ("SomeToken-1", "token"),
    ("x-us-master-token", "token"),
    ("X-US-Master-Token", "token"),
    ("x-us-master-token", "long67890"),
    ("NonSecretHeader", "Ololo"),
    ("NonSecretHeader", "Azaza"),
    ("X-DL-API-Key", "123456789"),
    ("Cookie", "Session_id=sec_val; _ym_isad=2; sessionid2=sec_val2"),
    ("Cookie", "_ym_isad=2; session_id=sec_valci; sessionid2=sec_val2"),
)


def test_headers_cleaning():
    assert _EXPECTED_MASKED_HEADERS == clean_secret_data_in_headers(_INPUT_HEADERS)


def test_common_logging_request_start(caplog):
    caplog.set_level("INFO")
    caplog.clear()
    log_request_start(
        logger=LOGGER,
        method="get",
        full_path="/api/ololo/azaza?param=value",
        headers=(
            ("Some-Token", "Some_secret_value"),
            ("User-Agent", "Python-urllib/2.7"),
        ),
    )
    assert len(caplog.records) == 1

    rec = caplog.records[0]
    assert rec.name == LOGGER.name
    expected_message = (
        "Received request."
        " method: GET, path: /api/ololo/azaza?param=value,"
        " headers: {'Some-Token': 'Som<hidden>lue', 'User-Agent': 'Python-urllib/2.7'},"
        f" pid: {os.getpid()}"
    )
    assert rec.message == expected_message


def test_common_logging_request_end(caplog):
    caplog.set_level("INFO")
    caplog.clear()
    log_request_end(
        logger=LOGGER,
        method="get",
        full_path="/unistat/?",
        status_code=200,
    )
    assert len(caplog.records) == 1

    rec = caplog.records[0]
    assert rec.name == LOGGER.name
    assert "Response. method: GET, path: /unistat/?, status: 200" == rec.message


request_path = "/?a=b"
request_cookies = {
    "Session_id": "sec_val",
    "sessionid2": "sec_val2",
    "non_secret_cookie": "non_secret_val",
}
request_headers = (
    ("X-YaCloud-SubjectToken", "1234567890"),
    ("Plain-Header", "Plain value"),
    ("User-Agent", "Test client"),
    ("X-Request-Id", "parentreqid1234"),
)


# TODO FIX: Request ID tests
def test_common_logging_flask(caplog):
    caplog.set_level("INFO")
    caplog.clear()

    app = flask.Flask(__name__)

    ContextVarMiddleware().wrap_flask_app(app)

    RequestLoggingContextControllerMiddleWare().set_up(app)
    RequestIDService(request_id_app_prefix=None).set_up(app)

    @app.route("/")
    def home():
        return flask.jsonify({"msg": "Hello!"})

    client = app.test_client()
    for c_name, c_val in request_cookies.items():
        client.set_cookie("localhost", c_name, c_val)

    resp = client.get(
        request_path,
        headers={k: v for k, v in request_headers},
    )
    assert resp.status_code == 200
    records = caplog.records

    # FIXME: fix the original stochastic 'Unclosed client session' errors.
    other_records = [rec for rec in records if rec.name == "asyncio"]
    records = [rec for rec in records if rec.name != "asyncio"]
    if other_records:
        LOGGER.error("Extraneous asyncio logrecords: %r", other_records)

    assert len(records) == 2

    expected_start_msg = (
        "Received request. method: GET, path: /?a=b,"
        " headers:"
        " {'User-Agent': 'Test client',"
        " 'Host': 'localhost',"
        " 'X-Yacloud-Subjecttoken': '123<hidden>890',"
        " 'Plain-Header': 'Plain value',"
        " 'X-Request-Id': 'parentreqid1234',"
        " 'Cookie': 'Session_id=<hidden>; non_secret_cookie=non_secret_val; sessionid2=<hidden>'},"
        f" pid: {os.getpid()}"
    )
    expected_end_msg = "Response. method: GET, path: /?a=b, status: 200"

    assert expected_start_msg == records[0].message
    assert expected_end_msg == records[1].message

    # Check that own request id was appended
    internal_request_id = records[0].request_id
    assert internal_request_id and internal_request_id.startswith("parentreqid1234--")

    parent_request_id = records[0].parent_request_id
    assert parent_request_id and parent_request_id == "parentreqid1234"


@pytest.mark.asyncio
async def test_common_logging_aiohttp(caplog, aiohttp_client):
    caplog.set_level("INFO")
    caplog.clear()

    app = web.Application(
        middlewares=[
            RequestBootstrap(
                req_id_service=aio_request_id.RequestId(append_own_req_id=True),
            ).middleware,
        ]
    )

    async def handle(_: web.Request):
        return web.json_response({})

    app.router.add_get("/", handle)

    client = await aiohttp_client(app, cookies=request_cookies)

    resp = await client.get(request_path, headers=CIMultiDict(request_headers))
    assert resp.status == 200

    expected_start_msg = (
        "Received request. method: GET, path: /?a=b,"
        " headers:"
        f" {{'Host': '{client.host}:{client.port}',"
        " 'X-YaCloud-SubjectToken': '123<hidden>890',"
        " 'Plain-Header': 'Plain value',"
        " 'User-Agent': 'Test client',"
        " 'X-Request-Id': 'parentreqid1234',"
        " 'Accept': '*/*',"
        " 'Accept-Encoding': 'gzip, deflate',"
        " 'Cookie': 'Session_id=<hidden>; non_secret_cookie=non_secret_val; sessionid2=<hidden>'},"
        f" pid: {os.getpid()}"
    )
    expected_end_msg = "Response. method: GET, path: /?a=b, status: 200"

    # noqa
    req_id_records = [rec for rec in caplog.records if rec.name == aio_request_id.LOGGER.name]
    assert len(req_id_records) == 2
    assert expected_start_msg == req_id_records[0].message
    assert expected_end_msg == req_id_records[1].message

    # Check that own request id was appended
    internal_request_id = req_id_records[0].request_id
    assert internal_request_id and internal_request_id.startswith("parentreqid1234--")


@pytest.mark.parametrize(
    "source, expected_masked",
    (
        (None, None),
        ({}, {}),
        (
            dict(password="1", a=2),
            dict(password="<hidden>", a=2),
        ),
        (
            dict(password=1, token=False, cypher_text=5.0, a=2),
            dict(password="<hidden>", token="<hidden>", cypher_text="<hidden>", a=2),
        ),
        (
            dict(a=dict(a=dict(token="", password="", cypher_text="<hidden>", sample="asdf"))),
            dict(
                a=dict(
                    a=dict(
                        token="<hidden>",
                        password="<hidden>",
                        cypher_text="<hidden>",
                        sample="asdf",
                    )
                )
            ),
        ),
        (
            dict(items=[dict(password="asdf"), dict(password="asdf")]),
            dict(items=[dict(password="<hidden>"), dict(password="<hidden>")]),
        ),
        (
            dict(password=["asdf", "fdsa"]),
            dict(password=["<hidden>", "<hidden>"]),
        ),
    ),
)
def test_mask_sensitive_fields_by_name_recursive(source: dict[str, Any], expected_masked: dict[str, Any]):
    actual_masked = mask_sensitive_fields_by_name_in_json_recursive(source)
    assert actual_masked == expected_masked

lib/dl_api_commons/dl_api_commons_tests/unit/aio/test_common_logging.py
datalens-tech/datalens-backend
from __future__ import annotations

import time
from typing import (
    Awaitable,
    Callable,
    ClassVar,
    Optional,
)

from aiohttp import web
from aiohttp.test_utils import TestClient
import pytest

from dl_api_commons.aio.middlewares.auth_trust_middleware import auth_trust_middleware
from dl_api_commons.aio.middlewares.csrf import (
    CSRFMiddleware,
    generate_csrf_token,
)
from dl_api_commons.aio.middlewares.request_bootstrap import RequestBootstrap
from dl_api_commons.aio.middlewares.request_id import RequestId
from dl_api_commons.aiohttp.aiohttp_wrappers import (
    DLRequestView,
    RequiredResource,
    RequiredResourceCommon,
)


_SAMPLE_USER_ID = "123"
_SAMPLE_TIMESTAMP = 1
_VALID_CSRF_SECRET = "valid_secret"
_INVALID_CSRF_SECRET = "invalid_secret"
_CSRF_TIME_LIMIT = 3600 * 12

_AppFactory = Callable[[bool], Awaitable[TestClient]]


def ts_now():
    return int(time.time())


class TestingCSRFMiddleware(CSRFMiddleware):
    USER_ID_COOKIES = ("user_id_cookie",)


@pytest.fixture(scope="function")
async def csrf_app_factory(aiohttp_client) -> _AppFactory:
    async def f(authorized: Optional[bool]) -> TestClient:
        async def non_csrf_handler(request: web.Request):
            return web.json_response(dict(ok="ok"), status=200)

        async def csrf_handler(request: web.Request):
            return web.json_response(dict(ok="ok"), status=200)

        class SkipCSRFView(DLRequestView):
            REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset({RequiredResourceCommon.SKIP_CSRF})

            async def put(self) -> web.StreamResponse:
                return web.json_response(dict(ok="ok"), status=200)

        req_id_service = RequestId(
            append_own_req_id=True,
            app_prefix="f",
        )

        app = web.Application(
            middlewares=[
                RequestBootstrap(
                    req_id_service=req_id_service,
                ).middleware,
                auth_trust_middleware(fake_user_id=_SAMPLE_USER_ID if authorized else None),
                TestingCSRFMiddleware(
                    csrf_header_name="x-csrf-token",
                    csrf_time_limit=_CSRF_TIME_LIMIT,
                    csrf_secret=_VALID_CSRF_SECRET,
                    csrf_methods=("POST", "PUT", "DELETE"),
                ).middleware,
            ]
        )

        app.router.add_get("/", non_csrf_handler)
        app.router.add_post("/", csrf_handler)
        app.router.add_put("/", SkipCSRFView)

        return await aiohttp_client(app)

    return f


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "case_name, method, authorized, headers, cookies",
    [
        (
            "OK_just_cookie",
            "POST",
            False,
            {
                "x-csrf-token": "{}:{}".format(
                    generate_csrf_token(_SAMPLE_USER_ID, ts_now(), _VALID_CSRF_SECRET), ts_now()
                )
            },
            {"user_id_cookie": _SAMPLE_USER_ID},
        ),
        (
            "OK_just_user_id",
            "POST",
            True,
            {
                "x-csrf-token": "{}:{}".format(
                    generate_csrf_token(_SAMPLE_USER_ID, ts_now(), _VALID_CSRF_SECRET), ts_now()
                )
            },
            {"does_not": "matter"},
        ),
        (
            "OK_cookie_and_user_id",
            "POST",
            True,
            {
                "x-csrf-token": "{}:{}".format(
                    generate_csrf_token(_SAMPLE_USER_ID, ts_now(), _VALID_CSRF_SECRET), ts_now()
                )
            },
            {"user_id_cookie": _SAMPLE_USER_ID},
        ),
        (
            "SKIP_non_csrf_method",
            "GET",
            False,
            {
                "x-csrf-token": "{}:{}".format(
                    generate_csrf_token(_SAMPLE_USER_ID, ts_now(), _VALID_CSRF_SECRET), ts_now()
                )
            },
            {"user_id_cookie": _SAMPLE_USER_ID},
        ),
        (
            "SKIP_no_cookies",
            "POST",
            True,
            {
                "x-csrf-token": "{}:{}".format(
                    generate_csrf_token(_SAMPLE_USER_ID, ts_now(), _VALID_CSRF_SECRET), ts_now()
                )
            },
            None,
        ),
        (
            "SKIP_no_user_tokens",
            "POST",
            False,
            {
                "x-csrf-token": "{}:{}".format(
                    generate_csrf_token(_SAMPLE_USER_ID, ts_now(), _VALID_CSRF_SECRET), ts_now()
                )
            },
            {"does_not": "matter"},
        ),
        (
            "SKIP_view_marker",
            "PUT",
            False,
            {"x-csrf-token": "some_token"},
            {"user_id_cookie": _SAMPLE_USER_ID},
        ),
    ],
)
async def test_csrf_ok(
    case_name: str,
    method: str,
    authorized: bool,
    headers: dict[str, str],
    cookies: dict[str, str],
    csrf_app_factory: _AppFactory,
):
    client = await csrf_app_factory(authorized)
    resp = await client.request(
        method=method,
        path="/",
        headers=headers,
        cookies=cookies,
    )
    assert resp.status == 200
    resp_json = await resp.json()
    assert resp_json == dict(ok="ok")


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "case_name, method, authorized, headers, cookies",
    [
        (
            "INVALID_no_csrf_token_provided",
            "POST",
            False,
            {},
            {"user_id_cookie": _SAMPLE_USER_ID},
        ),
        (
            "INVALID_malformed_csrf_token_1",
            "POST",
            False,
            {"x-csrf-token": "asdf:1234:5678"},
            {"user_id_cookie": _SAMPLE_USER_ID},
        ),
        (
            "INVALID_malformed_csrf_token_2",
            "POST",
            False,
            {"x-csrf-token": "asdf:qwer"},
            {"user_id_cookie": _SAMPLE_USER_ID},
        ),
        (
            "INVALID_empty_token",
            "POST",
            False,
            {"x-csrf-token": ":1234"},
            {"user_id_cookie": _SAMPLE_USER_ID},
        ),
        (
            "INVALID_token_expired",
            "POST",
            False,
            {
                "x-csrf-token": "{}:{}".format(
                    generate_csrf_token(_SAMPLE_USER_ID, ts_now() - 2 * _CSRF_TIME_LIMIT, _VALID_CSRF_SECRET),
                    ts_now() - 2 * _CSRF_TIME_LIMIT,
                )
            },
            {"user_id_cookie": _SAMPLE_USER_ID},
        ),
        (
            "INVALID_bad_secret",
            "POST",
            False,
            {
                "x-csrf-token": "{}:{}".format(
                    generate_csrf_token(_SAMPLE_USER_ID, ts_now(), _INVALID_CSRF_SECRET), ts_now()
                )
            },
            {"user_id_cookie": _SAMPLE_USER_ID},
        ),
    ],
)
async def test_csrf_invalid(
    case_name: str,
    method: str,
    authorized: bool,
    headers: dict[str, str],
    cookies: dict[str, str],
    csrf_app_factory: _AppFactory,
):
    validation_failed_text = "CSRF validation failed"
    validation_failed_status = 400

    client = await csrf_app_factory(authorized)
    resp = await client.request(
        method=method,
        path="/",
        headers=headers,
        cookies=cookies,
    )
    assert resp.status == validation_failed_status
    resp_text = await resp.text()
    assert resp_text == validation_failed_text

lib/dl_api_commons/dl_api_commons_tests/unit/aio/test_csrf.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Awaitable,
    Callable,
    Dict,
    Sequence,
)

from aiohttp import web
from aiohttp.test_utils import TestClient
from aiohttp.typedefs import Handler
import attr
import pytest
import shortuuid

from dl_api_commons.aio.middlewares.commit_rci import commit_rci_middleware
from dl_api_commons.aio.middlewares.request_bootstrap import RequestBootstrap
from dl_api_commons.aio.middlewares.request_id import RequestId
from dl_api_commons.aiohttp.aiohttp_wrappers import DLRequestBase


pytestmark = [pytest.mark.asyncio]


@attr.s
class _AppConfig:
    request_id: RequestId = attr.ib()
    extra_mw: Sequence = attr.ib(default=())
    extra_handlers: Dict[str, Handler] = attr.ib(factory=dict)


_AppFactory = Callable[[_AppConfig], Awaitable[TestClient]]


@pytest.fixture(scope="function")
async def app_factory(aiohttp_client) -> _AppFactory:
    async def f(config: _AppConfig) -> TestClient:
        async def not_streamed(_):
            return web.json_response({})

        async def streamed(request):
            resp = web.StreamResponse()
            await resp.prepare(request)
            await resp.write_eof(b"ololo")
            return resp

        app = web.Application(
            middlewares=[
                RequestBootstrap(
                    config.request_id,
                ).middleware,
                *config.extra_mw,
            ]
        )
        app.on_response_prepare.append(config.request_id.on_response_prepare)
        app.router.add_get("/not_streamed", not_streamed)
        app.router.add_get("/streamed", streamed)
        for route, handler in config.extra_handlers.items():
            app.router.add_route("*", route, handler)
        return await aiohttp_client(app)

    return f


async def test_no_req_id(app_factory: _AppFactory):
    app = await app_factory(_AppConfig(RequestId()))

    async with app.get("/not_streamed") as resp:
        assert resp.status == 200
        await resp.read()
        assert resp.headers["X-Request-ID"]

    async with app.get("/streamed") as resp:
        assert resp.status == 200
        body = await resp.read()
        assert body == b"ololo"
        assert resp.headers["X-Request-ID"]

    async with app.get("/not_found_1234125") as resp:
        assert resp.status == 404
        assert resp.headers["X-Request-ID"]


async def test_req_id_from_client_no_append(app_factory: _AppFactory):
    app = await app_factory(_AppConfig(RequestId()))
    req_id = shortuuid.uuid()

    async with app.get("/streamed", headers={"x-request-id": req_id}) as resp:
        assert resp.status == 200
        assert resp.headers["X-Request-ID"] == req_id


async def test_req_id_from_client_append(app_factory: _AppFactory):
    app = await app_factory(_AppConfig(RequestId(append_own_req_id=True, app_prefix="sp")))
    req_id = shortuuid.uuid()

    async with app.get("/streamed", headers={"x-request-id": req_id}) as resp:
        assert resp.status == 200
        assert resp.headers["X-Request-ID"].startswith(f"{req_id}--sp.")


async def test_committed_rci(app_factory: _AppFactory):
    async def ensure_rci_committed_handler(request: web.Request) -> web.Response:
        dl_request = DLRequestBase.get_for_request(request)
        return web.json_response({"request_id": dl_request.rci.request_id})

    app = await app_factory(
        _AppConfig(
            RequestId(),
            extra_mw=(commit_rci_middleware(),),
            extra_handlers={"/ensure_rci": ensure_rci_committed_handler},
        )
    )

    req_id = shortuuid.uuid()

    async with app.get("/ensure_rci", headers={"x-request-id": req_id}) as resp:
        assert resp.status == 200
        assert resp.headers["X-Request-ID"] == req_id
        resp_json = await resp.json()
        assert resp_json == {"request_id": req_id}


async def test_uncommitted_committed_rci(app_factory: _AppFactory):
    async def ensure_rci_is_not_committed(request: web.Request) -> web.Response:
        dl_request = DLRequestBase.get_for_request(request)
        if dl_request.is_rci_committed():
            return web.json_response({}, status=500, reason="RCI is committed but it was not expected")
        else:
            return web.json_response({"request_id": dl_request.temp_rci.request_id})

    app = await app_factory(
        _AppConfig(
            RequestId(),
            extra_handlers={"/ensure_rci": ensure_rci_is_not_committed},
        )
    )

    req_id = shortuuid.uuid()

    async with app.get("/ensure_rci", headers={"x-request-id": req_id}) as resp:
        assert resp.status == 200, resp.reason
        assert resp.headers["X-Request-ID"] == req_id
        resp_json = await resp.json()
        assert resp_json == {"request_id": req_id}


async def test_endpoint_code(app_factory: _AppFactory, caplog):
    caplog.set_level("INFO")

    msg = "hello from dummy view"

    class SomeView(web.View):
        endpoint_code = "some_view_epc"

        async def get(self):
            logging.getLogger().info(msg)
            return web.Response()

    class NoEpc(web.View):
        async def get(self):
            logging.getLogger().info(msg)
            return web.Response()

    app = await app_factory(
        _AppConfig(
            RequestId(),
            extra_handlers={
                "/some_view": SomeView,
                "/no_epc": NoEpc,
            },
        )
    )

    # Case 'has endpoint code'
    async with app.get("/some_view") as resp:
        assert resp.status == 200, resp.reason

    log_rec = next(r for r in caplog.records if r.msg == msg)
    assert log_rec.endpoint_code == SomeView.endpoint_code

    # Case 'no endpoint code (not our view)'
    caplog.clear()
    async with app.get("/no_epc") as resp:
        assert resp.status == 200, resp.reason

    log_rec = next(r for r in caplog.records if r.msg == msg)
    assert log_rec.endpoint_code is None


async def test_request_id(app_factory: _AppFactory, caplog):
    caplog.set_level("INFO")

    msg = "hello from dummy view"
    err_msg = "hello from evil view"

    class SomeView(web.View):
        endpoint_code = "some_view_epc"

        async def get(self):
            logging.getLogger().info(msg)
            return web.Response()

    class ErrorView(web.View):
        endpoint_code = "error_view_epc"

        async def get(self):
            logging.getLogger().info(err_msg)
            return web.Response(status=400)

        async def post(self):
            logging.getLogger().info(err_msg)
            raise ValueError("Throwing exception on purpose cause I am evil")

    app = await app_factory(
        _AppConfig(
            RequestId(append_own_req_id=True, app_prefix="sp"),
            extra_handlers={
                "/some_view": SomeView,
                "/error_view": ErrorView,
            },
        )
    )

    async with app.get("/some_view", headers={"x-request-id": "some_req_id"}) as resp:
        assert resp.status == 200, resp.reason
    log_rec = next(r for r in caplog.records if r.msg == msg)
    assert log_rec.request_id.startswith("some_req_id--sp.")
    assert log_rec.parent_request_id == "some_req_id"
    log_rec = next(r for r in caplog.records if r.msg.startswith("Response"))
    assert log_rec.message == "Response. method: GET, path: /some_view, status: 200"

    caplog.clear()
    async with app.get("/error_view", headers={"x-request-id": "some_req_id"}) as resp:
        assert resp.status == 400, resp.reason
    log_rec = next(r for r in caplog.records if r.msg == err_msg)
    assert log_rec.request_id.startswith("some_req_id--sp.")
    assert log_rec.parent_request_id == "some_req_id"
    log_rec = next(r for r in caplog.records if r.msg.startswith("Response"))
    assert log_rec.message == "Response. method: GET, path: /error_view, status: 400"

    caplog.clear()
    async with app.post("/error_view", headers={"x-request-id": "some_req_id"}) as resp:
        assert resp.status == 500, resp.reason
    log_rec = next(r for r in caplog.records if r.msg == err_msg)
    assert log_rec.request_id.startswith("some_req_id--sp.")
    assert log_rec.parent_request_id == "some_req_id"
    log_rec = next(r for r in caplog.records if r.msg.startswith("Response"))
    assert log_rec.message == "Response. method: POST, path: /error_view, status: 500"

lib/dl_api_commons/dl_api_commons_tests/unit/aio/test_request_id.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Callable,
    Optional,
)

from aiohttp import web
import pytest

from dl_api_commons.aio.middlewares.error_handling_outer import (
    AIOHTTPErrorHandler,
    ErrorData,
    ErrorLevel,
)
from dl_api_commons.aio.middlewares.request_bootstrap import RequestBootstrap
from dl_api_commons.aio.middlewares.request_id import RequestId
from dl_api_commons.aio.middlewares.tracing import TracingService


class MockError(Exception):
    body: dict
    status: int
    level: ErrorLevel

    def __init__(self, status: int, level: ErrorLevel, body: dict):
        super().__init__(status, level)
        self.status = status
        self.level = level
        self.body = body


class MockErrorHandlingMW(AIOHTTPErrorHandler):
    def _classify_error(self, err: Exception, request: web.Request) -> ErrorData:
        if isinstance(err, MockError):
            return ErrorData(err.status, err.body, err.level)
        return ErrorData(500, {}, ErrorLevel.error)


def get_view(the_endpoint_code: Optional[str], action: Callable[[], dict]):
    if the_endpoint_code is not None:

        class EPView(web.View):
            endpoint_code = the_endpoint_code

            async def get(self):
                return web.json_response(action())

        return EPView
    else:

        class NonEPView(web.View):
            async def get(self):
                return web.json_response(action())

        return NonEPView


def register_view(app: web.Application, path: str, ep_code: Optional[str]):
    def deco(func: Callable[[], dict]):
        app.router.add_route("*", path, get_view(ep_code, func))
        return func

    return deco


@pytest.mark.asyncio
async def test_tracing_scenarios(aiohttp_client):
    """
    At this moment only checks that tracing/error handling middleware doesn't breaks something.
    TODO: Add InMemory reporter to tests tracer to be able to check that spans got correct tags
    """
    app = web.Application(
        middlewares=[
            TracingService().middleware,
            RequestBootstrap(
                req_id_service=RequestId(),
                error_handler=MockErrorHandlingMW(use_sentry=False, sentry_app_name_tag="tests"),
            ).middleware,
        ]
    )

    @register_view(app, "/ok", "ok")
    def ok():
        return dict(ok="ok")

    @register_view(app, "/err_err", "err")
    def err_err():
        raise MockError(200, ErrorLevel.error, dict(err_err=None))

    @register_view(app, "/err_info", "err_info")
    def err_info():
        raise MockError(400, ErrorLevel.info, dict(err_info=None))

    @register_view(app, "/no_ep_ok", None)
    def no_ep_ok():
        return dict(no_ep_ok=None)

    @register_view(app, "/no_ep_err", None)
    def no_ep_err():
        raise MockError(401, ErrorLevel.error, dict(no_ep_err=None))

    client = await aiohttp_client(app)

    resp = await client.get("/ok")
    resp_json = await resp.json()
    assert resp.status == 200 and resp_json == dict(ok="ok")

    resp = await client.get("/err_err")
    resp_json = await resp.json()
    assert resp.status == 200 and resp_json == dict(err_err=None)

    resp = await client.get("/err_info")
    resp_json = await resp.json()
    assert resp.status == 400 and resp_json == dict(err_info=None)

    resp = await client.get("/no_ep_ok")
    resp_json = await resp.json()
    assert resp.status == 200 and resp_json == dict(no_ep_ok=None)

    resp = await client.get("/no_ep_err")
    resp_json = await resp.json()
    assert resp.status == 401 and resp_json == dict(no_ep_err=None)

lib/dl_api_commons/dl_api_commons_tests/unit/aio/test_tracing.py
datalens-tech/datalens-backend
import asyncio

import aiohttp.pytest_plugin
import aiohttp.test_utils
import pytest
from statcommons.logs import LOGMUTATORS

from dl_api_commons.logging_config import add_log_context


pytest_plugins = ("aiohttp.pytest_plugin",)

try:
    del aiohttp.pytest_plugin.loop
except AttributeError:
    pass


def pytest_configure(config):  # noqa
    LOGMUTATORS.apply(require=False)
    LOGMUTATORS.add_mutator("log_context", add_log_context)


@pytest.fixture
def loop(event_loop):
    asyncio.set_event_loop(event_loop)
    yield event_loop
    # Attempt to cover an old version of pytest-asyncio:
    # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
    asyncio.set_event_loop_policy(None)

lib/dl_api_commons/dl_api_commons_tests/unit/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import os
import threading

import flask

from dl_api_commons.flask.middlewares.aio_event_loop_middleware import AIOEventLoopMiddleware


def test_simulation(loop):
    pre_test_count = 1000
    pre_test_ids = set()
    marker_attr_n = "_tainted_n"  # marker 'known new loop'
    marker_attr_c = "_tainted_c"  # marker 'known new loop used as current'

    strict_id = True  # whether to consider id() as unique. Was supposed to be mostly so for small numbers.
    loop_keeper = []  # just to prevent GCing

    for idx in range(pre_test_count):
        new_loop = asyncio.new_event_loop()
        loop_keeper.append(new_loop)
        if strict_id:
            assert id(new_loop) not in pre_test_ids, dict(case="new_loop_is_not_new", step=idx)
        new_loop.set_debug(False)
        assert not getattr(new_loop, marker_attr_n, None), "new loop should be new"
        setattr(new_loop, marker_attr_n, True)
        asyncio.set_event_loop(new_loop)
        try:
            current_loop = asyncio.get_event_loop()
            assert current_loop is new_loop, dict(
                case="wrong_current_loop", step=idx, new_loop=new_loop, current_loop=current_loop
            )
            assert getattr(current_loop, marker_attr_n, None), "current loop should have been marked as new"
            assert not getattr(current_loop, marker_attr_c, None), "current loop should not have been used"
            setattr(current_loop, marker_attr_c, True)
            current_loop.run_until_complete(asyncio.sleep(0.0001))
            loop_id = id(current_loop)
        finally:
            new_loop.close()
            asyncio.set_event_loop(None)
        if strict_id:
            assert loop_id not in pre_test_ids, dict(step=idx, case="current_loop_is_not_new")
        pre_test_ids.add(loop_id)

    if strict_id:
        assert len(pre_test_ids) == pre_test_count, "did not get enough unique loops"


def test_app(caplog, loop):
    caplog.set_level("DEBUG")
    app = flask.Flask(__name__)
    AIOEventLoopMiddleware(debug=False).wrap_flask_app(app)

    marker_attr = "_test_tainted"
    strict_id = False  # whether to consider id() as unique. Was supposed to be mostly so for small numbers.

    @app.route("/")
    def home():
        async def make_some_sleep():
            await asyncio.sleep(0.0001)

        loop = asyncio.get_event_loop()
        assert not getattr(loop, marker_attr, None), "loop should be new"
        setattr(loop, marker_attr, True)

        loop.run_until_complete(make_some_sleep())
        return flask.jsonify(
            dict(
                event_loop_id=id(loop),
                pid=os.getpid(),
                thread=repr(threading.current_thread()),
                thread_id=threading.get_ident(),
            )
        )

    client = app.test_client()

    event_loop_id_set = set()
    rq_count = 300
    debug_data = []

    for _ in range(rq_count):
        resp = client.get("/")
        assert resp.status_code == 200
        resp_data = resp.json
        debug_data.append(resp_data)
        event_loop_id_set.add(resp_data["event_loop_id"])

    if strict_id:
        assert len(event_loop_id_set) == rq_count, debug_data

lib/dl_api_commons/dl_api_commons_tests/unit/flask/test_aio_event_loop_middleware.py
datalens-tech/datalens-backend
from __future__ import annotations

import flask

from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.flask.middlewares.commit_rci_middleware import ReqCtxInfoMiddleware
from dl_api_commons.flask.middlewares.context_var_middleware import ContextVarMiddleware
from dl_api_commons.flask.middlewares.logging_context import RequestLoggingContextControllerMiddleWare
from dl_api_commons.flask.middlewares.request_id import RequestIDService


def test_integration(caplog):
    caplog.set_level("DEBUG")
    caplog.clear()
    app = flask.Flask(__name__)

    ContextVarMiddleware().wrap_flask_app(app)

    RequestLoggingContextControllerMiddleWare().set_up(app)
    RequestIDService(
        request_id_app_prefix=None,
        append_local_req_id=False,
    ).set_up(app)
    ReqCtxInfoMiddleware(
        plain_headers=("include-me",),
    ).set_up(app)

    rci_lst = []

    @app.route("/test")
    def test_connection():
        rci_lst.append(ReqCtxInfoMiddleware.get_request_context_info())
        return flask.jsonify({})

    client = app.test_client()

    resp = client.get(
        "/test",
        headers={
            "x-request-id": "asdf1234",
            "include-me": "include_me",
            "exclude-me": "exclude_me",
            "referer": "referer",
            "x-chart-id": "x-chart-id",
        },
    )
    assert resp.status_code == 200
    assert len(rci_lst) == 1
    rci: RequestContextInfo = rci_lst[0]

    assert rci.request_id == "asdf1234"
    assert rci.plain_headers.get("include-me") == "include_me"
    assert rci.plain_headers.get("exclude-me") is None
    # TODO FIX: Check headers

lib/dl_api_commons/dl_api_commons_tests/unit/flask/test_commit_rci_middleware.py
datalens-tech/datalens-backend
from __future__ import annotations

import flask

from dl_api_commons.flask.middlewares.context_var_middleware import ContextVarMiddleware
from dl_api_commons.flask.middlewares.tracing import TracingMiddleware


def test_app(caplog, loop):
    """Just to check that middleware doesn't break something"""
    caplog.set_level("DEBUG")

    app = flask.Flask(__name__)
    TracingMiddleware(
        url_prefix_exclude=(),
    ).wrap_flask_app(app)
    ContextVarMiddleware().wrap_flask_app(app)

    @app.route("/ok")
    def ok():
        return flask.jsonify({})

    @app.route("/err")
    def err():
        return 500, flask.jsonify({})

    client = app.test_client()

    resp = client.get("/ok")
    assert resp.status_code == 200

    resp = client.get("/err")
    assert resp.status_code == 500

lib/dl_api_commons/dl_api_commons_tests/unit/flask/test_tracing_middleware.py
datalens-tech/datalens-backend
from __future__ import annotations

import uuid

import pytest

from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.reporting.models import (
    QueryExecutionCacheInfoReportingRecord,
    QueryExecutionEndReportingRecord,
    QueryExecutionStartReportingRecord,
)
from dl_api_commons.reporting.profiler import DefaultReportingProfiler
from dl_api_commons.reporting.records import RequestResultReportingRecord
from dl_api_commons.reporting.registry import DefaultReportingRegistry
from dl_constants.api_constants import DLContextKey
from dl_constants.enums import (
    ConnectionType,
    ReportingQueryType,
)


CONNECTION_TYPE_CHYT_TEST = ConnectionType.declare("chyt_test")
CONNECTION_TYPE_TEST = ConnectionType.declare("test")


_QID = "some_qid_1234"

_DEFAULT_START_RECORD_TS_0 = QueryExecutionStartReportingRecord(
    query_id=_QID,
    timestamp=0,
    dataset_id="ds_123",
    connection_type=CONNECTION_TYPE_TEST,
    conn_reporting_data={
        "connection_id": "conn_123",
        "host": "8.8.8.8",
    },
    query_type=ReportingQueryType.external,
    query="SELECT 1",
    workbook_id="wb_123",
)

_DEFAULT_REPORT_FIELDS_FROM_START = dict(
    dataset_id="ds_123",
    connection_id="conn_123",
    connection_type=CONNECTION_TYPE_TEST.name,
    host="8.8.8.8",
    query_type=ReportingQueryType.external.name,
    query="SELECT 1",
    workbook_id="wb_123",
)

_CHYT_START_RECORD_TS_0 = QueryExecutionStartReportingRecord(
    query_id=_QID,
    timestamp=0,
    dataset_id="ds_123",
    connection_type=CONNECTION_TYPE_CHYT_TEST,
    conn_reporting_data={
        "connection_id": "conn_123",
        "cluster": "my_cluster",
        "clique_alias": "*ch_my_clique",
    },
    query_type=ReportingQueryType.external,
    query="SELECT 1",
    workbook_id=None,
)

_CHYT_REPORT_FIELDS_FROM_START = dict(
    dataset_id="ds_123",
    connection_id="conn_123",
    connection_type=CONNECTION_TYPE_CHYT_TEST.name,
    cluster="my_cluster",
    clique_alias="*ch_my_clique",
    query_type=ReportingQueryType.external.name,
    query="SELECT 1",
    workbook_id=None,
)


@pytest.mark.parametrize(
    "rci",
    [
        RequestContextInfo.create_empty(),
        RequestContextInfo.create(
            endpoint_code="ololo",
            tenant=None,
            user_name="user_name_123",
            user_id="uid_113",
            x_dl_context={k.value: str(uuid.uuid4()) for k in DLContextKey},
            plain_headers=None,
            secret_headers=None,
            auth_data=None,
            request_id="req999",
            x_dl_debug_mode=False,
        ),
    ],
)
@pytest.mark.parametrize(
    "case_name, records_seq, expected_query_data",
    [
        (
            "OK_no_cache_used",
            (
                _DEFAULT_START_RECORD_TS_0,
                QueryExecutionEndReportingRecord(query_id=_QID, timestamp=1, exception=None),
                RequestResultReportingRecord(
                    timestamp=1.0,
                    response_status_code=200,
                ),
            ),
            dict(
                _DEFAULT_REPORT_FIELDS_FROM_START,
                cache_used=False,
                cache_full_hit=False,
                status="success",
                execution_time=1000,
                error=None,
                err_code=None,
                response_status_code=200,
            ),
        ),
        (
            "OK_cache_hit",
            (
                _DEFAULT_START_RECORD_TS_0,
                QueryExecutionCacheInfoReportingRecord(query_id=_QID, timestamp=0.5, cache_full_hit=True),
                QueryExecutionEndReportingRecord(query_id=_QID, timestamp=1, exception=None),
                RequestResultReportingRecord(timestamp=1.0, response_status_code=200),
            ),
            dict(
                _DEFAULT_REPORT_FIELDS_FROM_START,
                cache_used=True,
                cache_full_hit=True,
                status="success",
                execution_time=1000,
                error=None,
                err_code=None,
                response_status_code=200,
            ),
        ),
        (
            "OK_no_cache_hit",
            (
                _DEFAULT_START_RECORD_TS_0,
                QueryExecutionCacheInfoReportingRecord(query_id=_QID, timestamp=0.5, cache_full_hit=False),
                QueryExecutionEndReportingRecord(query_id=_QID, timestamp=1, exception=None),
                RequestResultReportingRecord(
                    timestamp=1.0,
                    response_status_code=200,
                ),
            ),
            dict(
                _DEFAULT_REPORT_FIELDS_FROM_START,
                cache_used=True,
                cache_full_hit=False,
                status="success",
                execution_time=1000,
                error=None,
                err_code=None,
                response_status_code=200,
            ),
        ),
        (
            "ERR_no_cache_used",
            (
                _DEFAULT_START_RECORD_TS_0,
                QueryExecutionEndReportingRecord(query_id=_QID, timestamp=1, exception=ValueError("val err")),
                RequestResultReportingRecord(
                    timestamp=1.0,
                    response_status_code=500,
                    err_code="ERR.VAL_ERR",
                ),
            ),
            dict(
                _DEFAULT_REPORT_FIELDS_FROM_START,
                cache_used=False,
                cache_full_hit=False,
                status="error",
                execution_time=1000,
                error="ValueError: val err\n",
                err_code="ERR.VAL_ERR",
                response_status_code=500,
            ),
        ),
        (
            "OK_chyt",
            (
                _CHYT_START_RECORD_TS_0,
                QueryExecutionEndReportingRecord(query_id=_QID, timestamp=1, exception=None),
                RequestResultReportingRecord(
                    timestamp=1.0,
                    response_status_code=200,
                ),
            ),
            dict(
                _CHYT_REPORT_FIELDS_FROM_START,
                cache_used=False,
                cache_full_hit=False,
                status="success",
                execution_time=1000,
                error=None,
                err_code=None,
                response_status_code=200,
            ),
        ),
    ],
)
def test_db_query_report_generation(case_name, records_seq, expected_query_data, rci, caplog):
    required_extras = (
        "event_code",
        "dataset_id",
        "user_id",
        "billing_folder_id",
        "connection_id",
        "connection_type",
        "source",
        "username",
        "execution_time",
        "query",
        "status",
        "error",
        "err_code",
        "cache_used",
        "cache_full_hit",
        "query_type",
        "is_public",
        "dash_id",
        "dash_tab_id",
        "chart_id",
        "chart_kind",
        "response_status_code",
        "workbook_id",
    )
    if "chyt" in case_name:
        required_extras += (
            "cluster",
            "clique_alias",
        )
    else:
        required_extras += ("host",)

    expected_extras = dict(
        event_code="profile_db_request",
        source=rci.endpoint_code,
        billing_folder_id=rci.tenant.get_tenant_id() if rci.tenant is not None else None,
        user_id=rci.user_id,
        username=rci.user_name,
        is_public=0,
        dash_id=rci.x_dl_context.get(DLContextKey.DASH_ID),
        dash_tab_id=rci.x_dl_context.get(DLContextKey.DASH_TAB_ID),
        chart_id=rci.x_dl_context.get(DLContextKey.CHART_ID),
        chart_kind=rci.x_dl_context.get(DLContextKey.CHART_KIND),
        **expected_query_data,
    )

    assert set(expected_extras.keys()) == set(required_extras)

    rr = DefaultReportingRegistry(rci=rci)
    rep_profiler = DefaultReportingProfiler(reporting_registry=rr)
    for evt in records_seq:
        rr.save_reporting_record(evt)

    caplog.set_level("INFO")
    caplog.clear()

    rep_profiler.flush_query_report(next(iter(records_seq)).query_id)

    report_log_records = [r for r in caplog.records if getattr(r, "event_code", None) == "profile_db_request"]

    assert len(report_log_records) == 1, f"Got non-single query report in logs: {report_log_records}"
    log_record = report_log_records[0]

    actual_extras = {
        extra_name: getattr(log_record, extra_name) for extra_name in required_extras if hasattr(log_record, extra_name)
    }

    assert actual_extras == expected_extras

lib/dl_api_commons/dl_api_commons_tests/unit/test_reporting.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_constants.enums import (
    ComponentErrorLevel,
    ComponentType,
)
from dl_core import component_errors
from dl_core.marshmallow import ErrorCodeField
from dl_model_tools.schema.base import DefaultSchema


class ComponentErrorListSchema(DefaultSchema[component_errors.ComponentErrorRegistry]):
    TARGET_CLS = component_errors.ComponentErrorRegistry

    class ComponentErrorPackSchema(DefaultSchema[component_errors.ComponentErrorPack]):
        TARGET_CLS = component_errors.ComponentErrorPack

        class ComponentErrorSchema(DefaultSchema[component_errors.ComponentError]):
            TARGET_CLS = component_errors.ComponentError

            message = ma_fields.String()
            level = ma_fields.Enum(ComponentErrorLevel)
            code = ErrorCodeField()
            details = ma_fields.Dict()

        id = ma_fields.String()
        type = ma_fields.Enum(ComponentType)
        errors = ma_fields.List(ma_fields.Nested(ComponentErrorSchema))

    items = ma_fields.List(ma_fields.Nested(ComponentErrorPackSchema))

lib/dl_api_connector/dl_api_connector/api_schema/component_errors.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
)

from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.top_level import USEntryBaseSchema
from dl_constants.enums import ConnectionState
from dl_constants.enums import ConnectionType as CT
from dl_core.us_connection_base import ConnectionBase
from dl_model_tools.schema.dynamic_enum_field import DynamicEnumField


class ConnectionSchema(USEntryBaseSchema):
    TARGET_CLS = ConnectionBase
    CONN_TYPE_CTX_KEY: ClassVar[str] = "conn_type"

    # From ConnectionBase.as_dict()
    db_type = DynamicEnumField(CT, attribute="conn_type", dump_only=True)
    # TODO FIX: Change for datetime field after created/updated _at become datetime instead of string
    created_at = ma_fields.String(dump_only=True)
    updated_at = ma_fields.String(dump_only=True)

    # Create only fields
    permissions_mode = ma_fields.Raw(load_default=None, load_only=True)
    initial_permissions = ma_fields.Raw(load_default=None, load_only=True)

    def create_data_model_constructor_kwargs(self, data_attributes: dict[str, Any]) -> dict[str, Any]:
        """
        Point for adding extra kwargs for DataModel constructor:
         - Fields that is not defined in schema
         - Context-dependant defaults
         - Required fields defaults for fields which was turned into partials
        """
        return dict(data_attributes)

    def create_data_model(self, data_attributes: dict[str, Any]) -> Any:
        return self.TARGET_CLS.DataModel(**self.create_data_model_constructor_kwargs(data_attributes))

    def default_create_from_dict_kwargs(self, data: dict[str, Any]) -> dict[str, Any]:
        try:
            ct = CT(self.context[self.CONN_TYPE_CTX_KEY])
        except ValueError:
            ct = CT.unknown
        ret = dict(
            super().default_create_from_dict_kwargs(data),
            type_=ct.name,
            permissions_mode=data["permissions_mode"],
            initial_permissions=data["initial_permissions"],
            meta=dict(
                # TODO CONSIDER: May be do it in handler?
                state=ConnectionState.saved.name
            ),
        )

        return ret


class ConnectionMetaMixin(ConnectionSchema):
    """
    This is a base class for Schematic-based connections
    To be removed after migration to attrs
    """

    meta = ma_fields.Dict(dump_only=True)  # In ConnectionBase.as_dict() meta was included

lib/dl_api_connector/dl_api_connector/api_schema/connection_base.py
datalens-tech/datalens-backend
from typing import Optional

from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.extras import FieldExtra


def cache_ttl_field(
    attribute: str,
    required: bool = False,
    allow_none: bool = True,
    missing: Optional[int] = None,
    default: Optional[int] = None,
    bi_extra: FieldExtra = FieldExtra(editable=True),
) -> ma_fields.Integer:
    return ma_fields.Integer(
        attribute=attribute,
        required=required,
        allow_none=allow_none,
        load_default=missing,
        dump_default=default,
        bi_extra=bi_extra,
    )


def secret_string_field(
    attribute: str,
    required: bool = True,
    allow_none: bool = False,
    default: Optional[str] = None,
    bi_extra: FieldExtra = FieldExtra(editable=True),
) -> ma_fields.String:
    return ma_fields.String(
        attribute=attribute,
        load_only=True,  # do not expose never
        required=required,
        allow_none=allow_none,
        dump_default=default,
        bi_extra=bi_extra,
    )

lib/dl_api_connector/dl_api_connector/api_schema/connection_base_fields.py
datalens-tech/datalens-backend
from __future__ import annotations

import marshmallow
from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.extras import FieldExtra
from dl_constants.enums import RawSQLLevel
import dl_core.marshmallow as core_ma_fields


class RawSQLLevelMixin(marshmallow.Schema):
    raw_sql_level = ma_fields.Enum(
        RawSQLLevel,
        attribute="data.raw_sql_level",
        required=False,
        allow_none=False,
        load_default=RawSQLLevel.off,
        dump_default=RawSQLLevel.off,
        bi_extra=FieldExtra(editable=True),
    )


class DataExportForbiddenMixin(marshmallow.Schema):
    data_export_forbidden = core_ma_fields.OnOffField(
        attribute="data.data_export_forbidden",
        required=False,
        load_default=False,
        dump_default="off",
        bi_extra=FieldExtra(editable=True),
    )

lib/dl_api_connector/dl_api_connector/api_schema/connection_mixins.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Mapping,
    Optional,
)

import marshmallow
from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.connection_base import ConnectionSchema
from dl_api_connector.api_schema.connection_base_fields import (
    cache_ttl_field,
    secret_string_field,
)
from dl_api_connector.api_schema.extras import FieldExtra
from dl_core.utils import (
    parse_comma_separated_hosts,
    validate_hostname_or_ip_address,
)


class DBHostField(ma_fields.String):
    """
    Field that validates comma-separated list of hosts
    If parent schema is instance of ClassicSQLConnectionSchema,
     schema attribute `ALLOW_MULTI_HOST` will be checked to determine if multiple hosts is allowed.
    Otherwise only single host is allowed.
    """

    @staticmethod
    def _validate_host_str(user_hosts_str: str, allow_multi_host: bool) -> None:
        if allow_multi_host:
            host_list = [host for host in parse_comma_separated_hosts(user_hosts_str) if host]
        else:
            host_list = [user_hosts_str]

        if host_list:
            for user_host in host_list:
                try:
                    validate_hostname_or_ip_address(user_host)
                except ValueError as err:
                    raise marshmallow.ValidationError("Unable to parse host: {}".format(user_host)) from err
        else:
            raise marshmallow.ValidationError("Unable to parse host: {}".format(user_hosts_str))

    def _deserialize(self, value: Any, attr: Optional[str], data: Optional[Mapping[str, Any]], **kwargs: Any) -> Any:
        user_host_str = super()._deserialize(value, attr, data, **kwargs)

        schema = self.parent
        # May be assert?
        if isinstance(schema, ClassicSQLConnectionSchema):
            # FIXME: Remove reference to schema. Replace with mixin or flag
            allow_multi_host = schema.ALLOW_MULTI_HOST
        else:
            allow_multi_host = False

        self._validate_host_str(user_host_str, allow_multi_host=allow_multi_host)

        return user_host_str


class ClassicSQLConnectionSchema(ConnectionSchema):
    ALLOW_MULTI_HOST = False

    host = DBHostField(attribute="data.host", required=True, bi_extra=FieldExtra(editable=True))
    port = ma_fields.Integer(attribute="data.port", required=True, bi_extra=FieldExtra(editable=True))
    username = ma_fields.String(attribute="data.username", required=True, bi_extra=FieldExtra(editable=True))
    password = secret_string_field(attribute="data.password", bi_extra=FieldExtra(editable=True))
    db_name = ma_fields.String(attribute="data.db_name", allow_none=True, bi_extra=FieldExtra(editable=True))
    cache_ttl_sec = cache_ttl_field(attribute="data.cache_ttl_sec")

lib/dl_api_connector/dl_api_connector/api_schema/connection_sql.py
datalens-tech/datalens-backend
from __future__ import annotations

import enum
from typing import (
    Optional,
    Sequence,
    TypedDict,
    Union,
)

import attr


class OperationsMode(enum.Enum):
    pass


class CreateMode(OperationsMode):
    create = enum.auto()
    test = enum.auto()


class EditMode(OperationsMode):
    edit = enum.auto()
    test = enum.auto()


class SchemaKWArgs(TypedDict):
    only: Optional[Sequence[str]]
    partial: Union[Sequence[str], bool]
    exclude: Sequence[str]
    load_only: Sequence[str]
    dump_only: Sequence[str]


@attr.s(frozen=True, auto_attribs=True)
class FieldExtra:
    partial_in: Sequence[OperationsMode] = ()
    exclude_in: Sequence[OperationsMode] = ()
    editable: Union[bool, Sequence[OperationsMode]] = ()

lib/dl_api_connector/dl_api_connector/api_schema/extras.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    AbstractSet,
    Any,
    Iterable,
    Mapping,
    Optional,
    Sequence,
    Type,
    Union,
)

from marshmallow import (
    EXCLUDE,
    ValidationError,
)
from marshmallow_oneofschema import OneOfSchema

from dl_api_connector.api_schema.source_base import (
    DataSourceBaseSchema,
    DataSourceTemplateBaseSchema,
)
from dl_constants.enums import DataSourceType


class DataSourceSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "source_type"
    type_schemas: dict[str, Type[DataSourceBaseSchema]] = {}

    def get_obj_type(self, obj: dict[str, Any]) -> str:
        return obj[self.type_field].name


class DataSourceStrictSchema(DataSourceSchema):
    def load(
        self,
        data: Union[Mapping[str, Any], Iterable[Mapping[str, Any]]],
        *,
        many: Optional[bool] = None,
        partial: Optional[Union[bool, Sequence[str], AbstractSet[str]]] = None,
        unknown: Optional[str] = None,
        **kwargs: Any,
    ) -> Any:
        assert isinstance(data, (dict, list))
        many = many if many is not None else isinstance(data, list)
        if many and any("managed_by" not in source for source in data) or not many and "managed_by" not in data:
            raise ValidationError("Missing `managed_by`")
        return super().load(data, many=many, partial=partial, unknown=unknown, **kwargs)


class DataSourceTemplateResponseSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "source_type"
    type_schemas: dict[str, Type[DataSourceTemplateBaseSchema]] = {}

    def get_obj_type(self, obj: dict[str, Any]) -> str:
        return obj[self.type_field].name


def register_source_api_schema(source_type: DataSourceType, schema_cls: Type[DataSourceBaseSchema]) -> None:
    DataSourceSchema.type_schemas[source_type.name] = schema_cls


def register_source_template_api_schema(
    source_type: DataSourceType, schema_cls: Type[DataSourceTemplateBaseSchema]
) -> None:
    DataSourceTemplateResponseSchema.type_schemas[source_type.name] = schema_cls

lib/dl_api_connector/dl_api_connector/api_schema/source.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Mapping,
    Optional,
    Union,
)

from marshmallow import EXCLUDE
from marshmallow import fields as ma_fields
from marshmallow import post_load
from marshmallow_oneofschema import OneOfSchema

from dl_constants.enums import (
    BinaryJoinOperator,
    ConditionPartCalcMode,
    DataSourceType,
    IndexKind,
    JoinConditionType,
    JoinType,
    ManagedBy,
    UserDataType,
)
from dl_core.db import (
    IndexInfo,
    SchemaColumn,
)
from dl_core.db.native_type_schema import OneOfNativeTypeSchema
from dl_core.marshmallow import FrozenSetField
from dl_core.multisource import (
    BinaryCondition,
    ConditionPartDirect,
    ConditionPartFormula,
    ConditionPartResultField,
)
from dl_model_tools.schema.base import BaseSchema
from dl_model_tools.schema.dynamic_enum_field import DynamicEnumField


if TYPE_CHECKING:
    from dl_core.multisource import ConditionPart


class RawSchemaColumnSchema(BaseSchema):
    name = ma_fields.String()
    title = ma_fields.String()

    native_type = ma_fields.Nested(OneOfNativeTypeSchema, allow_none=True)

    user_type = ma_fields.Enum(UserDataType)
    description = ma_fields.String(dump_default="", allow_none=True)
    has_auto_aggregation = ma_fields.Boolean(dump_default=False, allow_none=True)
    lock_aggregation = ma_fields.Boolean(dump_default=False, allow_none=True)
    nullable = ma_fields.Boolean(dump_default=None, allow_none=True)

    @post_load
    def make_column(self, data: dict, **kwargs: Any) -> SchemaColumn:
        return SchemaColumn(
            name=data["name"],
            title=data["title"],
            user_type=data["user_type"],
            native_type=data["native_type"],
            description=data.get("description", ""),
            has_auto_aggregation=data.get("has_auto_aggregation", False),
            lock_aggregation=data.get("lock_aggregation", False),
            nullable=data["nullable"],
        )


class IndexInfoSchema(BaseSchema):
    columns = ma_fields.List(ma_fields.String, required=True)
    kind = ma_fields.Enum(IndexKind, allow_none=True)

    @post_load
    def make_index_info(self, data: dict, **kwargs: Any) -> IndexInfo:
        columns = tuple(data.pop("columns"))
        return IndexInfo(columns=columns, **data)


class SimpleParametersSchema(BaseSchema):
    pass


class SQLParametersSchema(BaseSchema):
    db_name = ma_fields.String(allow_none=True)
    table_name = ma_fields.String(allow_none=True)
    db_version = ma_fields.String(allow_none=True)


class SchematizedParametersSchema(SQLParametersSchema):
    schema_name = ma_fields.String(allow_none=True)


class SubselectParametersSchema(SimpleParametersSchema):
    subsql = ma_fields.String()


class DataSourceCommonSchema(BaseSchema):
    title = ma_fields.String(required=True)
    connection_id = ma_fields.String(allow_none=True)
    source_type = DynamicEnumField(DataSourceType)
    raw_schema = ma_fields.Nested(RawSchemaColumnSchema, many=True, allow_none=True)
    index_info_set = FrozenSetField(
        ma_fields.Nested(IndexInfoSchema),
        sort_output=True,
        allow_none=True,
        load_default=None,
        dump_default=None,
    )
    parameters = ma_fields.Nested(SimpleParametersSchema)  # redefined in subclasses
    parameter_hash = ma_fields.String(dump_only=True)


class DataSourceTemplateResponseField(ma_fields.Field):
    """
    Optimized version of the `DataSourceTemplateResponseSchema` that does a bare minimum.
    """

    _allowed_keys = frozenset(
        (
            "title",
            "tab_title",
            "connection_id",
            "source_type",
            # In `DataSourceCommonSchema` but not intended for dsrc templates: 'raw_schema'
            "index_info_set",
            "parameters",
            "parameter_hash",
            "id",
            "managed_by",
            "virtual",
            "valid",
            "group",
            "form",
            "disabled",
        )
    )

    def _serialize(self, value: Optional[dict], attr: Optional[str], obj: Any, **kwargs: Any) -> Optional[dict]:
        if value is None:
            return None
        assert isinstance(value, dict)
        allowed_keys = self._allowed_keys
        value = {key: val for key, val in value.items() if key in allowed_keys}
        st = value.get("source_type", None)
        if st is not None:
            value["source_type"] = st.name
        return value

    def _deserialize(self, value: Any, attr: Optional[str], data: Optional[Mapping[str, Any]], **kwargs: Any) -> Any:
        raise Exception("Not Applicable")


class VirtualFlagField(ma_fields.Field):
    def _serialize(self, value: Union[str, ManagedBy, None], attr: Optional[str], obj: Any, **kwargs: Any) -> bool:
        if isinstance(value, str):
            value = ManagedBy[value]
        if value is None:
            return False
        return value != ManagedBy.user


class DataSourceBaseSchema(DataSourceCommonSchema):
    id = ma_fields.String(required=True)
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user)
    virtual = VirtualFlagField(attribute="managed_by", dump_only=True)
    valid = ma_fields.Boolean(load_default=True)


class DataSourceTemplateBaseSchema(DataSourceCommonSchema):
    group = ma_fields.List(ma_fields.String(), required=True)
    form = ma_fields.List(ma_fields.Dict(), required=False)
    tab_title = ma_fields.String(load_default=None, required=False)
    disabled = ma_fields.Boolean(load_default=False, required=False)


class SimpleDataSourceSchema(DataSourceBaseSchema):
    pass


class SimpleDataSourceTemplateSchema(DataSourceTemplateBaseSchema):
    pass


class SQLDataSourceSchema(SimpleDataSourceSchema):
    parameters = ma_fields.Nested(SQLParametersSchema)


class SQLDataSourceTemplateSchema(DataSourceTemplateBaseSchema):
    parameters = ma_fields.Nested(SQLParametersSchema)


class SchematizedSQLDataSourceSchema(SQLDataSourceSchema):
    parameters = ma_fields.Nested(SchematizedParametersSchema)


class SchematizedSQLDataSourceTemplateSchema(DataSourceTemplateBaseSchema):
    parameters = ma_fields.Nested(SchematizedParametersSchema)


class SubselectDataSourceSchema(SimpleDataSourceSchema):
    parameters = ma_fields.Nested(SubselectParametersSchema)


class SubselectDataSourceTemplateSchema(DataSourceTemplateBaseSchema):
    parameters = ma_fields.Nested(SubselectParametersSchema)


class SourceAvatarSchema(BaseSchema):
    id = ma_fields.String(required=True)
    source_id = ma_fields.String()
    title = ma_fields.String()
    is_root = ma_fields.Boolean()
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user)
    virtual = VirtualFlagField(attribute="managed_by", dump_only=True)
    valid = ma_fields.Boolean(dump_only=True)


class SourceAvatarStrictSchema(SourceAvatarSchema):
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user, required=True)


class ConditionPartSchema(BaseSchema):
    calc_mode = ma_fields.Enum(ConditionPartCalcMode, required=True, dump_default=ConditionPartCalcMode.direct)


class ConditionPartDirectSchema(ConditionPartSchema):
    source = ma_fields.String(required=True)

    @post_load
    def make_condition_part(self, data: dict, **kwargs: Any) -> ConditionPartDirect:
        return ConditionPartDirect(source=data["source"])


class ConditionPartFormulaSchema(ConditionPartSchema):
    formula = ma_fields.String(required=True)

    @post_load
    def make_condition_part(self, data: dict, **kwargs: Any) -> ConditionPartFormula:
        return ConditionPartFormula(formula=data["formula"])


class ConditionPartResultFieldSchema(ConditionPartSchema):
    field_id = ma_fields.String(required=True)

    @post_load
    def make_condition_part(self, data: dict, **kwargs: Any) -> ConditionPartResultField:
        return ConditionPartResultField(field_id=data["field_id"])


class ConditionPartGenericSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "calc_mode"
    type_schemas = {
        ConditionPartCalcMode.direct.name: ConditionPartDirectSchema,
        ConditionPartCalcMode.formula.name: ConditionPartFormulaSchema,
        ConditionPartCalcMode.result_field.name: ConditionPartResultFieldSchema,
    }

    def get_obj_type(self, obj: ConditionPart) -> str:
        return getattr(obj, self.type_field).name


class AvatarRelationSchema(BaseSchema):
    class JoinConditionSchema(BaseSchema):
        type = ma_fields.Enum(JoinConditionType, attribute="condition_type", required=True)
        operator = ma_fields.Enum(BinaryJoinOperator, required=True)
        left = ma_fields.Nested(ConditionPartGenericSchema, attribute="left_part", required=True)
        right = ma_fields.Nested(ConditionPartGenericSchema, attribute="right_part", required=True)

        @post_load
        def make_condition(self, data: dict, **kwargs: Any) -> BinaryCondition:
            return BinaryCondition(
                operator=data["operator"],
                left_part=data["left_part"],
                right_part=data["right_part"],
            )

    id = ma_fields.String(required=True)
    left_avatar_id = ma_fields.String()
    right_avatar_id = ma_fields.String()
    conditions = ma_fields.Nested(JoinConditionSchema, many=True)
    join_type = ma_fields.Enum(JoinType)
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user, load_default=ManagedBy.user)
    virtual = VirtualFlagField(attribute="managed_by", dump_only=True)

lib/dl_api_connector/dl_api_connector/api_schema/source_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import copy
import itertools
import logging
import os
from typing import (
    Any,
    ClassVar,
    Generic,
    Iterable,
    Optional,
    Sequence,
    TypeVar,
    Union,
    final,
)

import marshmallow
from marshmallow import (
    missing,
    post_load,
    pre_load,
)
from marshmallow import Schema
from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.extras import (
    CreateMode,
    EditMode,
    FieldExtra,
    OperationsMode,
    SchemaKWArgs,
)
from dl_core import exc as bi_core_exc
from dl_core.base_models import (
    EntryLocation,
    PathEntryLocation,
    WorkbookEntryLocation,
)
from dl_core.us_entry import USEntry
from dl_core.us_manager.us_manager import USManagerBase


LOGGER = logging.getLogger(__name__)

_TARGET_OBJECT_TV = TypeVar("_TARGET_OBJECT_TV")


class BaseTopLevelSchema(Schema, Generic[_TARGET_OBJECT_TV]):
    CTX_KEY_EDITABLE_OBJECT: ClassVar[str] = "editable_object"
    CTX_KEY_OPERATIONS_MODE: ClassVar[str] = "operations_mode"

    TARGET_CLS: ClassVar[type]

    def __init__(
        self,
        *,
        only: Optional[Sequence[str]] = None,
        exclude: Sequence[str] = (),
        many: bool = False,
        context: Optional[dict] = None,
        load_only: Sequence[str] = (),
        dump_only: Sequence[str] = (),
        partial: Union[Sequence[str], bool] = False,
        unknown: Optional[str] = None,
    ) -> None:
        refined_kwargs = self._refine_init_kwargs(
            dict(
                only=only,
                partial=partial,
                exclude=exclude,
                load_only=load_only,
                dump_only=dump_only,
            ),
            # Context is not bind to self here
            None if context is None else context.get(self.CTX_KEY_OPERATIONS_MODE),
        )

        super().__init__(many=many, context=context, unknown=unknown, **refined_kwargs)

    @staticmethod
    def get_field_extra(f: ma_fields.Field) -> Optional[FieldExtra]:
        return f.metadata.get("bi_extra", None)

    @property
    def operations_mode(self) -> Optional[CreateMode]:
        return self.context.get(self.CTX_KEY_OPERATIONS_MODE)

    @classmethod
    def all_fields_dict(cls) -> dict[str, ma_fields.Field]:
        return cls._declared_fields

    @classmethod
    def all_fields_with_extra_info(cls) -> Iterable[tuple[str, ma_fields.Field, FieldExtra]]:
        for field_name, field in cls.all_fields_dict().items():
            extra = cls.get_field_extra(field)
            if extra is not None:
                yield field_name, field, extra

    def _refine_init_kwargs(self, kw_args: SchemaKWArgs, operations_mode: Optional[OperationsMode]) -> SchemaKWArgs:
        if operations_mode is None:
            return kw_args

        ret = copy.deepcopy(kw_args)

        if isinstance(operations_mode, EditMode) and kw_args["only"] is None:
            editable_fields = []
            for field_name, _field, field_extra in self.all_fields_with_extra_info():
                if (isinstance(field_extra.editable, bool) and field_extra.editable is True) or (
                    isinstance(field_extra.editable, Sequence) and operations_mode in field_extra.editable
                ):
                    editable_fields.append(field_name)

            # To even not accept non-editable fields
            ret["only"] = editable_fields
            # To reset required flags & prevent fallback to `missing=`
            ret["partial"] = editable_fields
        else:
            # Handling extra partials only in non-edit mode
            extra_partial = []
            for field_name, _field, field_extra in self.all_fields_with_extra_info():
                if operations_mode in field_extra.partial_in:
                    extra_partial.append(field_name)
            # Do not override partials if was passed via schema constructor
            # TODO CONSIDER: May be append to constructor partials if present?
            if kw_args["partial"] is False and extra_partial:
                ret["partial"] = tuple(extra_partial)

        # Handling extra excludes in every case
        extra_exclude = []
        for field_name, _field, field_extra in self.all_fields_with_extra_info():
            if operations_mode in field_extra.exclude_in:
                extra_exclude.append(field_name)

        if extra_exclude:
            ret["exclude"] = tuple(
                set(
                    itertools.chain(
                        kw_args["exclude"],
                        extra_exclude,
                    )
                )
            )

        return ret

    def create_object(self, data: dict[str, Any]) -> _TARGET_OBJECT_TV:
        raise NotImplementedError()

    def update_object(self, obj: _TARGET_OBJECT_TV, data: dict[str, Any]) -> _TARGET_OBJECT_TV:
        raise NotImplementedError()

    @post_load(pass_many=False)
    def post_load(self, data: dict[str, Any], **_: Any) -> Union[_TARGET_OBJECT_TV, dict]:
        if isinstance(self.operations_mode, EditMode):
            editable_object = self.context.get(self.CTX_KEY_EDITABLE_OBJECT)
            if editable_object is None:
                return data
            assert isinstance(editable_object, self.TARGET_CLS)
            return self.update_object(editable_object, data)
        if isinstance(self.operations_mode, CreateMode):
            return self.create_object(data)
        raise ValueError(f"Can not perform load. Unknown operations mode: {self.operations_mode!r}")

    def get_allowed_unknown_fields(self) -> set[str]:
        """
        Should return set of fields that does not exist at schema, but it is ok.
        This fields will be removed from data in pre_load hook.
        Main purpose: create a transition period for frontend to adopt to schema changes.
        """
        return set()

    @final
    def handle_unknown_fields(self, data: dict[str, Any]) -> dict[str, Any]:
        refined_data = {}
        allowed_unknown_data = {}
        disallowed_unknown_data = {}

        allowed_unknown_fields = self.get_allowed_unknown_fields()

        for field_name, field_value in data.items():
            if field_name in self.fields:
                refined_data[field_name] = field_value
            else:
                if field_name in allowed_unknown_fields:
                    allowed_unknown_data[field_name] = field_value
                else:
                    disallowed_unknown_data[field_name] = field_value

        # Passing disallowed unknown fields to schema to get correct error messages
        refined_data.update(disallowed_unknown_data)

        if allowed_unknown_data or disallowed_unknown_data:
            sorted_all_unknown_keys = list(
                sorted(
                    itertools.chain(
                        allowed_unknown_data.keys(),
                        disallowed_unknown_data.keys(),
                    )
                )
            )
            sorted_allowed_unknown_keys = list(sorted(allowed_unknown_data.keys()))
            sorted_disallowed_unknown_keys = list(sorted(disallowed_unknown_data.keys()))

            LOGGER.info(
                "Got unknown fields for schema %s/%s. Allowed: %s. Disallowed: %s",
                type(self).__qualname__,
                self.operations_mode,
                ",".join(sorted_allowed_unknown_keys),
                ",".join(sorted_disallowed_unknown_keys),
                extra=dict(
                    schema_unk_fields=sorted_all_unknown_keys,
                    schema_unk_fields_allowed=sorted_allowed_unknown_keys,
                    schema_unk_fields_disallowed=sorted_disallowed_unknown_keys,
                ),
            )

        return refined_data

    @pre_load(pass_many=False)
    def pre_load(self, data: dict[str, Any], **_: Any) -> dict[str, Any]:
        all_data_keys = list(sorted(data.keys()))
        LOGGER.debug(
            "Got fields for schema %s/%s: %s",
            type(self).__qualname__,
            self.operations_mode,
            ",".join(all_data_keys),
            extra=dict(
                schema_input_keys=all_data_keys,
            ),
        )
        return self.handle_unknown_fields(data)


_US_ENTRY_TV = TypeVar("_US_ENTRY_TV", bound=USEntry)


class USEntryBaseSchema(BaseTopLevelSchema[_US_ENTRY_TV]):
    CTX_KEY_USM: ClassVar[str] = "us_manager"

    id = ma_fields.String(dump_only=True, attribute="uuid")
    key = ma_fields.String(dump_only=True, attribute="raw_us_key")

    dir_path = ma_fields.String(
        allow_none=False,
        load_only=True,
        load_default="Connection",
        bi_extra=FieldExtra(exclude_in=[CreateMode.test]),
        attribute="entry_key.dir_path",
    )
    workbook_id = ma_fields.String(
        required=False,
        allow_none=True,
        load_default=None,
        bi_extra=FieldExtra(exclude_in=[CreateMode.test]),
        attribute="entry_key.workbook_id",
    )
    name = ma_fields.String(
        required=True,
        allow_none=False,
        bi_extra=FieldExtra(exclude_in=[CreateMode.test]),
        attribute="entry_key.entry_name",
    )

    @property
    def us_manager(self) -> USManagerBase:
        return self.context[self.CTX_KEY_USM]

    def default_create_from_dict_kwargs(self, data: dict[str, Any]) -> dict[str, Any]:
        entry_name = data["entry_key"]["entry_name"]
        entry_dir_path = data["entry_key"]["dir_path"]
        entry_wb_id = data["entry_key"]["workbook_id"]

        if entry_wb_id is None and entry_dir_path is None:
            raise marshmallow.ValidationError("dir_path can not be null if workbook_id is null", field_name="dir_path")

        return dict(
            ds_key=resolve_entry_loc_from_api_req_body(
                name=entry_name,
                dir_path=entry_dir_path,
                workbook_id=entry_wb_id,
            ),
            us_manager=self.us_manager,
        )

    # TODO FIX: Find a way to specify return type
    def create_data_model(self, data_attributes: dict[str, Any]) -> Any:
        raise NotImplementedError()

    @final
    def validate_object(self, obj: _US_ENTRY_TV) -> None:
        try:
            obj.validate()
        except bi_core_exc.EntryValidationError as err:
            if err.model_field:
                cause_attribute = f"data.{err.model_field}"
                cause_field_name = (
                    next(
                        map(
                            lambda field: field.name,
                            filter(lambda field: field.attribute == cause_attribute, self.fields.values()),
                        ),
                        None,
                    )
                    or marshmallow.exceptions.SCHEMA
                )
            else:
                cause_field_name = marshmallow.exceptions.SCHEMA

            raise marshmallow.ValidationError(err.message, field_name=cause_field_name) from err

    @final
    def create_object(self, data: dict[str, Any]) -> _US_ENTRY_TV:
        data_attributes = data.get("data", {})

        obj = self.TARGET_CLS.create_from_dict(  # type: ignore  # TODO: fix
            data_dict=self.create_data_model(data_attributes),
            **self.default_create_from_dict_kwargs(data),
        )
        self.validate_object(obj)
        return obj

    @final
    def update_object(self, obj: _US_ENTRY_TV, data: dict[str, Any]) -> _US_ENTRY_TV:
        # Assumed that only data of USEntry can be modified with schema
        assert not (data.keys() - {"data"})

        for data_field_name, data_field_value in data.get("data", {}).items():
            if data_field_value is missing:
                continue

            if isinstance(obj.data, dict):
                obj.data[data_field_name] = obj.data[data_field_value]
            else:
                setattr(obj.data, data_field_name, data_field_value)

        self.validate_object(obj)
        return obj


def resolve_entry_loc_from_api_req_body(
    *,
    name: str,
    dir_path: Optional[str],
    workbook_id: Optional[str],
) -> EntryLocation:
    assert name is not None, "name can not be None"

    if workbook_id is None:
        assert dir_path is not None, "dir_path can not be None"
        return PathEntryLocation(os.path.join(dir_path, name))
    else:
        return WorkbookEntryLocation(
            workbook_id=workbook_id,
            entry_name=name,
        )

lib/dl_api_connector/dl_api_connector/api_schema/top_level.py
datalens-tech/datalens-backend
from typing import (
    ClassVar,
    Optional,
    final,
)

from dl_i18n.localizer_base import (
    Localizer,
    Translatable,
)


class ConnectionInfoProvider:
    title_translatable: ClassVar[Translatable]
    alias: ClassVar[Optional[str]] = None

    @classmethod
    @final
    def get_title(cls, localizer: Localizer) -> str:
        return localizer.translate(cls.title_translatable)

lib/dl_api_connector/dl_api_connector/connection_info.py
datalens-tech/datalens-backend
import abc
from typing import (
    ClassVar,
    Optional,
    Tuple,
    Type,
)

from dl_api_connector.api_schema.connection_base import ConnectionSchema
from dl_api_connector.api_schema.source_base import (
    DataSourceBaseSchema,
    DataSourceTemplateBaseSchema,
)
from dl_api_connector.connection_info import ConnectionInfoProvider
from dl_api_connector.dashsql import (
    DashSQLParamLiteralizer,
    DefaultDashSQLParamLiteralizer,
)
from dl_api_connector.form_config.models.base import ConnectionFormFactory
from dl_api_lib.query.registry import MQMFactorySettingItem
from dl_constants.enums import QueryProcessingMode
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
    CoreSourceDefinition,
)
from dl_formula.core.dialect import (
    DialectCombo,
    DialectName,
)
from dl_i18n.localizer_base import TranslationConfig
from dl_query_processing.compilation.filter_compiler import (
    FilterFormulaCompiler,
    MainFilterFormulaCompiler,
)
from dl_query_processing.multi_query.factory import (
    DefaultMultiQueryMutatorFactory,
    NoCompengMultiQueryMutatorFactory,
)


class ApiSourceDefinition(abc.ABC):
    core_source_def_cls: ClassVar[Type[CoreSourceDefinition]]
    api_schema_cls: ClassVar[Type[DataSourceBaseSchema]] = DataSourceBaseSchema
    template_api_schema_cls: ClassVar[Type[DataSourceTemplateBaseSchema]] = DataSourceTemplateBaseSchema


class ApiConnectionDefinition(abc.ABC):
    core_conn_def_cls: ClassVar[Type[CoreConnectionDefinition]]
    api_generic_schema_cls: ClassVar[Type[ConnectionSchema]]
    alias: ClassVar[Optional[str]] = None  # TODO remove in favor of info provider
    info_provider_cls: ClassVar[Type[ConnectionInfoProvider]]
    form_factory_cls: ClassVar[Optional[Type[ConnectionFormFactory]]] = None


class ApiConnector(abc.ABC):
    # backend_type-bound properties - TODO: move to a separate entity
    formula_dialect_name: ClassVar[DialectName] = DialectName.DUMMY
    multi_query_mutation_factories: tuple[MQMFactorySettingItem, ...] = (
        MQMFactorySettingItem(
            query_proc_mode=QueryProcessingMode.basic,
            factory_cls=DefaultMultiQueryMutatorFactory,
        ),
        MQMFactorySettingItem(
            query_proc_mode=QueryProcessingMode.no_compeng,
            factory_cls=NoCompengMultiQueryMutatorFactory,
        ),
    )
    is_forkable: ClassVar[bool] = True
    is_compeng_executable: ClassVar[bool] = False
    filter_formula_compiler_cls: ClassVar[Type[FilterFormulaCompiler]] = MainFilterFormulaCompiler
    dashsql_literalizer_cls: ClassVar[Type[DashSQLParamLiteralizer]] = DefaultDashSQLParamLiteralizer
    # others
    core_connector_cls: ClassVar[Type[CoreConnector]]
    connection_definitions: ClassVar[Tuple[Type[ApiConnectionDefinition], ...]] = ()
    source_definitions: ClassVar[Tuple[Type[ApiSourceDefinition], ...]] = ()
    translation_configs: ClassVar[frozenset[TranslationConfig]] = frozenset()
    compeng_dialect: Optional[DialectCombo] = None

lib/dl_api_connector/dl_api_connector/connector.py
datalens-tech/datalens-backend
import abc

import sqlalchemy as sa
from sqlalchemy.types import TypeEngine

from dl_constants.enums import UserDataType
from dl_core.exc import DashSQLError


TValueBase = str | list[str] | tuple[str, ...]

BI_TYPE_TO_SA_TYPE: dict[UserDataType, TypeEngine] = {
    UserDataType.string: sa.TEXT(),
    UserDataType.integer: sa.BIGINT(),
    UserDataType.float: sa.FLOAT(),
    UserDataType.date: sa.DATE(),
    UserDataType.datetime: sa.DATETIME(),
    UserDataType.boolean: sa.BOOLEAN(),
    UserDataType.datetimetz: sa.DATETIME(timezone=True),
    UserDataType.genericdatetime: sa.DATETIME(),
}


class DashSQLParamLiteralizer(abc.ABC):
    @abc.abstractmethod
    def get_sa_type(self, bi_type: UserDataType, value_base: TValueBase) -> TypeEngine:
        raise NotImplementedError


class DefaultDashSQLParamLiteralizer(DashSQLParamLiteralizer):
    def get_sa_type(self, bi_type: UserDataType, value_base: TValueBase) -> TypeEngine:
        try:
            sa_type = BI_TYPE_TO_SA_TYPE[bi_type]
            return sa_type
        except (KeyError, ValueError):
            raise DashSQLError(f"Unsupported type {bi_type.name!r}")

lib/dl_api_connector/dl_api_connector/dashsql.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import Enum
from typing import (
    Any,
    Literal,
    Optional,
)

import attr

from dl_api_connector.form_config.models.common import (
    SerializableConfig,
    TFieldName,
    remap_skip_if_null,
    skip_if_null,
)


@attr.s(kw_only=True, frozen=True)
class FormFieldSelector(SerializableConfig):
    """Can be used to refer to a specific field of the form"""

    name: TFieldName = attr.ib()


class FormFieldApiAction(Enum):
    """What to do with the field when sending a request to the API (validate and include it in the request)"""

    include = "include"
    skip = "skip"


@attr.s(kw_only=True, frozen=True)
class FormFieldApiSchema(SerializableConfig):
    name: TFieldName = attr.ib()
    type: Optional[Literal["string", "boolean"]] = attr.ib(default=None, metadata=skip_if_null())
    required: bool = attr.ib(default=False)
    length: Optional[int] = attr.ib(default=None, metadata=skip_if_null())
    nullable: Optional[bool] = attr.ib(default=None, metadata=skip_if_null())
    default_action: FormFieldApiAction = attr.ib(
        default=FormFieldApiAction.include, metadata=remap_skip_if_null("defaultAction")
    )


@attr.s(kw_only=True, frozen=True)
class FormFieldConditionalApiAction(SerializableConfig):
    selector: FormFieldSelector = attr.ib()
    action: FormFieldApiAction = attr.ib()


@attr.s(kw_only=True, frozen=True)
class FormFieldApiActionCondition(SerializableConfig):
    when: FormFieldSelector = attr.ib()
    equals: Any = attr.ib()
    then: list[FormFieldConditionalApiAction] = attr.ib()


@attr.s(kw_only=True, frozen=True)
class FormActionApiSchema(SerializableConfig):
    items: list[FormFieldApiSchema] = attr.ib()
    conditions: list[FormFieldApiActionCondition] = attr.ib(factory=list)


@attr.s(kw_only=True, frozen=True)
class FormApiSchema(SerializableConfig):
    create: Optional[FormActionApiSchema] = attr.ib(default=None, metadata=skip_if_null())
    edit: Optional[FormActionApiSchema] = attr.ib(default=None, metadata=skip_if_null())
    check: Optional[FormActionApiSchema] = attr.ib(default=None, metadata=skip_if_null())

lib/dl_api_connector/dl_api_connector/form_config/models/api_schema.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from enum import Enum
from typing import (
    Any,
    Optional,
)

import attr

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.common import (
    SerializableConfig,
    TFieldName,
    TopLevelFieldName,
    remap_skip_if_null,
)
from dl_api_connector.form_config.models.rows import CustomizableRow
from dl_api_connector.form_config.models.rows.base import (
    DisplayConditionsMixin,
    FormFieldMixin,
    FormRow,
)
from dl_api_connector.form_config.models.rows.prepared.base import PreparedRow
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_i18n.localizer_base import Localizer


@attr.s(kw_only=True, frozen=True)
class FormUIOverride(SerializableConfig):
    show_create_dataset_btn: Optional[bool] = attr.ib(
        default=None, metadata=remap_skip_if_null("showCreateDatasetButton")
    )
    show_create_ql_chart_btn: Optional[bool] = attr.ib(
        default=None, metadata=remap_skip_if_null("showCreateQlChartButton")
    )


TOP_LEVEL_NON_CONFIG_FIELDS: set[TopLevelFieldName] = {field_name for field_name in TopLevelFieldName}


@attr.s(kw_only=True, frozen=True)
class ConnectionForm(SerializableConfig):
    """
    title               shown at the top of the form
    rows                form body (labels, fields, etc. => without action buttons)
    api_schema          specifies form fields' validation and action request body schemas
    template_name       name of the template that stored in the US; for connectors with templates (service, partner)
    form_ui_override    override form elements' visibility

    Keep in mind that the naming convention here is camelCase,
    so either use it in field names or remap keys using SerializableConfig features where it is needed
    """

    title: str = attr.ib()
    rows: list[FormRow] = attr.ib()
    api_schema: Optional[FormApiSchema] = attr.ib(default=None, metadata=remap_skip_if_null("apiSchema"))
    template_name: Optional[str] = attr.ib(default=None, metadata=remap_skip_if_null("templateName"))
    form_ui_override: Optional[FormUIOverride] = attr.ib(default=None, metadata=remap_skip_if_null("uiSchema"))

    def validate_conditional_fields(self) -> None:
        form_field_names = self._get_form_field_names()
        fields_in_conditions = self._get_conditional_fields()

        rogue_fields = fields_in_conditions - form_field_names - TOP_LEVEL_NON_CONFIG_FIELDS

        assert (
            not rogue_fields
        ), f"Fields {rogue_fields} are used in conditions, but not defined in the form or component inner fields"

    def validate_api_schema_fields(self) -> None:
        form_field_names = self._get_form_field_names()
        api_schema_fields = self._get_api_schema_fields()

        rogue_fields = api_schema_fields - form_field_names - TOP_LEVEL_NON_CONFIG_FIELDS

        assert not rogue_fields, (
            f"Fields {rogue_fields} are referred to in api schema conditions,"
            f" but not defined in the form or component inner fields"
        )

    def _get_form_field_names(self) -> set[TFieldName]:
        form_field_names: set[TFieldName] = set()

        for row in self.rows:
            if isinstance(row, CustomizableRow):
                for row_item in row.items:
                    if isinstance(row_item, FormFieldMixin):
                        form_field_names.add(row_item.name)
            if isinstance(row, PreparedRow):
                for inner_field in row.Inner:
                    form_field_names.add(inner_field)
                if isinstance(row, FormFieldMixin):
                    form_field_names.add(row.name)

        return form_field_names

    def _get_conditional_fields(self) -> set[TFieldName]:
        fields_in_conditions: set[TFieldName] = set()

        for row in self.rows:
            if isinstance(row, CustomizableRow):
                for row_item in row.items:
                    if isinstance(row_item, DisplayConditionsMixin) and row_item.display_conditions is not None:
                        fields_in_conditions |= set(row_item.display_conditions.keys())
            elif isinstance(row, PreparedRow):
                if isinstance(row, DisplayConditionsMixin) and row.display_conditions is not None:
                    fields_in_conditions |= set(row.display_conditions.keys())

        return fields_in_conditions

    def _get_api_schema_fields(self) -> set[TFieldName]:
        api_schema_fields: set[TFieldName] = set()

        defined_api_schemas = (
            [sch for sch in [self.api_schema.create, self.api_schema.edit, self.api_schema.check] if sch is not None]
            if self.api_schema is not None
            else []
        )
        for api_schema in defined_api_schemas:
            api_schema_fields |= set(form_field_api_schema.name for form_field_api_schema in api_schema.items)
            for condition in api_schema.conditions or []:
                api_schema_fields.add(condition.when.name)
                api_schema_fields |= set(conditional_action.selector.name for conditional_action in condition.then)

        return api_schema_fields


class ConnectionFormMode(Enum):
    create = "create"
    edit = "edit"


class ConnectionFormFactory:
    def __init__(self, mode: ConnectionFormMode, localizer: Localizer):
        self.mode = mode
        self._localizer = localizer

    def _filter_nulls(self, coll: list[Optional[Any]]) -> list[Any]:
        return [item for item in coll if item is not None]

    def _get_top_level_create_api_schema_items(self) -> list[FormFieldApiSchema]:
        return [
            FormFieldApiSchema(name=TopLevelFieldName.type_),
        ]

    def _get_top_level_check_api_schema_items(self) -> list[FormFieldApiSchema]:
        return (
            [
                FormFieldApiSchema(name=TopLevelFieldName.type_),
            ]
            if self.mode == ConnectionFormMode.create
            else []
        )

    @abc.abstractmethod
    def get_form_config(
        self, connector_settings: Optional[ConnectorSettingsBase], tenant: Optional[TenantDef]
    ) -> ConnectionForm:
        """Returns a form config built according to the specified settings"""

        raise NotImplementedError

lib/dl_api_connector/dl_api_connector/form_config/models/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import (
    Enum,
    unique,
)
from typing import (
    Any,
    Literal,
    Optional,
    Union,
    final,
)

import attr
from dynamic_enum import DynamicEnum


Align = Literal["start", "center", "end"]
WidthVariant = Literal["s", "m", "l", "auto"]
Width = Union[WidthVariant, int, float]


class MarkdownString(str):
    ...


MarkdownStr = Union[MarkdownString, str]  # marker type to show if a field supports markdown or not


@unique
class FormFieldName(Enum):
    """
    Field key in the API schema
    Enum to eliminate typos
    Left empty, so it can be extended with connector specific fields
    """


@unique
class CommonFieldName(FormFieldName):
    host = "host"
    port = "port"
    username = "username"
    password = "password"
    db_name = "db_name"
    raw_sql_level = "raw_sql_level"
    cache_ttl_sec = "cache_ttl_sec"
    secure = "secure"
    access_token = "access_token"
    token = "token"
    is_auto_create_dashboard = "is_auto_create_dashboard"
    advanced_settings = "advanced_settings"
    ssl_enable = "ssl_enable"
    data_export_forbidden = "data_export_forbidden"
    ssl_ca = "ssl_ca"


@unique
class BooleanField(Enum):
    on = "on"
    off = "off"


@unique
class TopLevelFieldName(Enum):
    """Fields that are sent into the API but do not represent data filled by the user in the main form"""

    type_ = "type"


@unique
class InnerFieldName(Enum):
    """
    Prepared rows/components can have inner fields that can' be controlled by the config,
    but can be used in conditions

    Each row/component can have its own set of inner fields,
    but their names can't intersect with each other or `FormFieldName`s
    """


TFieldName = Union[FormFieldName, TopLevelFieldName, InnerFieldName]


@unique
class OAuthApplication(Enum):
    """Identifier of the OAuth application that should be used by the frontend"""


@attr.s(kw_only=True, frozen=True)
class CFGMeta:
    METADATA_KEY = "_SERIALIZABLE_CONFIG_META_KEY_"

    inner: bool = attr.ib(
        default=False
    )  # whether it is a service field involved only in inner logic and which needs to be skipped
    skip_if_null: bool = attr.ib(default=False)  # sometimes it is more convenient for the UI to receive undefined
    key: Optional[str] = attr.ib(default=None)  # remap key

    def attr_meta(self) -> dict[str, CFGMeta]:
        return {self.METADATA_KEY: self}


@attr.s
class SerializableConfig:
    """
    Base class for attrs decorated configs, that prepares values for serialization.
    Include ClassVars and other non-attr fields by overriding as_dict method.
    Set attr field's metadata to `CFGMeta().attr_meta` to change the way it is processed.
    Has to be a parent class for all top-level config classes.
    """

    _SKIP_SENTINEL = object
    _remap_keys_buffer: dict[str | TFieldName, str] = attr.ib(
        init=False, factory=dict, metadata=CFGMeta(inner=True).attr_meta()
    )

    @staticmethod
    @final
    def prepare_value(value: Any) -> Any:
        if isinstance(value, (Enum, DynamicEnum)):
            return value.value
        if isinstance(value, SerializableConfig):
            return value.as_dict()
        if attr.has(value.__class__):
            return attr.asdict(value, value_serializer=SerializableConfig._attr_value_serializer, recurse=False)  # type: ignore
        if isinstance(value, (list, tuple)):
            return [SerializableConfig.prepare_value(item) for item in value]
        if isinstance(value, dict):
            return {SerializableConfig.prepare_value(k): SerializableConfig.prepare_value(v) for k, v in value.items()}
        return value

    @staticmethod
    @final
    def _attr_value_serializer(inst: type, field: attr.Attribute, value: Any) -> Any:
        cls = SerializableConfig
        if CFGMeta.METADATA_KEY in field.metadata:
            if not isinstance(inst, cls):
                raise ValueError(
                    f"Cannot use serialization meta for an object of type {type(inst).__name__},"
                    f" add {cls.__name__} as a parent class to enable this."
                )
            meta = field.metadata[CFGMeta.METADATA_KEY]
            assert isinstance(meta, CFGMeta), f"Unexpected meta class, {CFGMeta.__name__} expected"
            if meta.inner or meta.skip_if_null and value is None:
                return cls._SKIP_SENTINEL
            if meta.key is not None:
                inst._remap_keys_buffer[field.name] = meta.key

        return cls.prepare_value(value)

    def as_dict(self) -> dict[str, Any]:
        result = attr.asdict(self, value_serializer=self._attr_value_serializer, recurse=False)  # type: ignore
        return {
            k.value if isinstance(k, (Enum, DynamicEnum)) else self._remap_keys_buffer.get(k, k): v
            for k, v in result.items()
            if v is not self._SKIP_SENTINEL
        }


# frequently used meta shortcuts
def skip_if_null() -> dict[str, CFGMeta]:
    return CFGMeta(skip_if_null=True).attr_meta()


def remap(key: str) -> dict[str, CFGMeta]:
    return CFGMeta(key=key).attr_meta()


def remap_skip_if_null(key: str) -> dict[str, CFGMeta]:
    return CFGMeta(skip_if_null=True, key=key).attr_meta()


def inner() -> dict[str, CFGMeta]:
    return CFGMeta(inner=True).attr_meta()

lib/dl_api_connector/dl_api_connector/form_config/models/common.py
datalens-tech/datalens-backend
from .customizable import *  # noqa
from .prepared import *  # noqa

lib/dl_api_connector/dl_api_connector/form_config/models/rows/__init__.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Optional,
    Union,
)

import attr

from dl_api_connector.form_config.models.common import (
    SerializableConfig,
    TFieldName,
    remap_skip_if_null,
    skip_if_null,
)


@attr.s(kw_only=True, frozen=True)
class FormFieldMixin:
    """Form field that can be sent to the API"""

    name: TFieldName = attr.ib()


TDisplayConditions = dict[Union[TFieldName], Any]


@attr.s(kw_only=True, frozen=True)
class DisplayConditionsMixin(SerializableConfig):
    """Allow to control item visibility based on form field values"""

    display_conditions: Optional[TDisplayConditions] = attr.ib(
        default=None, metadata=remap_skip_if_null("displayConditions")
    )


@attr.s(kw_only=True, frozen=True)
class InnerFieldMixin(SerializableConfig):
    """Inner fields are not send to the API, but can affect other fields"""

    inner: Optional[bool] = attr.ib(default=None, metadata=skip_if_null())  # false if undefined


class FormRow(SerializableConfig):
    pass

lib/dl_api_connector/dl_api_connector/form_config/models/rows/base.py
datalens-tech/datalens-backend
from .base import CustomizableRow
from .components import (
    CheckboxRowItem,
    DatepickerRowItem,
    DescriptionRowItem,
    FileInputRowItem,
    HiddenRowItem,
    InputRowItem,
    LabelRowItem,
    PlainTextRowItem,
    RadioButtonRowItem,
    RadioGroupRowItem,
    RadioGroupRowItemOption,
    SelectableOption,
    SelectOption,
    SelectRowItem,
)


__all__ = (
    "CustomizableRow",
    "HiddenRowItem",
    "LabelRowItem",
    "SelectableOption",
    "InputRowItem",
    "SelectOption",
    "SelectRowItem",
    "RadioButtonRowItem",
    "RadioGroupRowItemOption",
    "RadioGroupRowItem",
    "CheckboxRowItem",
    "DatepickerRowItem",
    "PlainTextRowItem",
    "DescriptionRowItem",
    "FileInputRowItem",
)

lib/dl_api_connector/dl_api_connector/form_config/models/rows/customizable/__init__.py
datalens-tech/datalens-backend
from typing import (
    Any,
    ClassVar,
    Optional,
)

import attr

from dl_api_connector.form_config.models.common import (
    MarkdownStr,
    SerializableConfig,
    Width,
    remap_skip_if_null,
    skip_if_null,
)
from dl_api_connector.form_config.models.rows.base import (
    DisplayConditionsMixin,
    FormFieldMixin,
    FormRow,
    InnerFieldMixin,
)


@attr.s(kw_only=True, frozen=True)
class RowItem(SerializableConfig):
    component_id: ClassVar[str]

    def as_dict(self) -> dict[str, Any]:
        return dict(
            id=self.component_id,
            **super().as_dict(),
        )


@attr.s(kw_only=True, frozen=True)
class PlaceholderMixin(SerializableConfig):
    placeholder: Optional[str] = attr.ib(default=None, metadata=skip_if_null())


@attr.s(kw_only=True, frozen=True)
class DefaultValueMixin(SerializableConfig):
    # comma separated if multiple=true
    default_value: Optional[str | bool] = attr.ib(default=None, metadata=remap_skip_if_null("defaultValue"))


@attr.s(kw_only=True, frozen=True)
class ControlRowItem(RowItem, FormFieldMixin, DisplayConditionsMixin, InnerFieldMixin, DefaultValueMixin):
    fake_value: Optional[str] = attr.ib(default=None, metadata=remap_skip_if_null("fakeValue"))
    hint_text: Optional[MarkdownStr] = attr.ib(default=None, metadata=remap_skip_if_null("hintText"))
    width: Optional[Width] = attr.ib(default=None, metadata=skip_if_null())


@attr.s(kw_only=True, frozen=True)
class CustomizableRow(FormRow):
    items: list[RowItem] = attr.ib()

lib/dl_api_connector/dl_api_connector/form_config/models/rows/customizable/base.py
datalens-tech/datalens-backend
from typing import (
    Literal,
    Optional,
)

import attr

from dl_api_connector.form_config.models.common import (
    Align,
    MarkdownStr,
    SerializableConfig,
    Width,
    inner,
    remap_skip_if_null,
    skip_if_null,
)
from dl_api_connector.form_config.models.rows.base import (
    DisplayConditionsMixin,
    FormFieldMixin,
    InnerFieldMixin,
)
from dl_api_connector.form_config.models.rows.customizable.base import (
    ControlRowItem,
    DefaultValueMixin,
    PlaceholderMixin,
    RowItem,
)


@attr.s(kw_only=True, frozen=True)
class HiddenRowItem(RowItem, FormFieldMixin, InnerFieldMixin, DisplayConditionsMixin, DefaultValueMixin):
    """
    Can be used to send values that do not have a corresponding control item
    (e.g. calculated token expiration time, that is received from an external API)
    """

    component_id = "hidden"


@attr.s(kw_only=True, frozen=True)
class LabelRowItem(RowItem, DisplayConditionsMixin):
    component_id = "label"
    text: str = attr.ib()
    align: Optional[Align] = attr.ib(default=None, metadata=skip_if_null())
    help_text: Optional[MarkdownStr] = attr.ib(default=None, metadata=remap_skip_if_null("helpText"))


@attr.s(kw_only=True, frozen=True)
class SelectableOption:  # TODO rename to RadioButtonOption
    text: str = attr.ib()
    value: str = attr.ib()


@attr.s(kw_only=True, frozen=True)
class InputRowItem(ControlRowItem, PlaceholderMixin):
    @attr.s(kw_only=True, frozen=True)
    class Props(SerializableConfig):
        multiline: Optional[bool] = attr.ib(default=None, metadata=skip_if_null())  # false if undefined
        type: Literal["text", "password", "number"] = attr.ib(default="text")

    component_id = "input"

    control_props: Optional[Props] = attr.ib(default=None, metadata=remap_skip_if_null("controlProps"))


@attr.s(kw_only=True, frozen=True)
class SelectOption(SerializableConfig):
    class Data(SerializableConfig):
        description: Optional[str] = attr.ib(default=None, metadata=skip_if_null())

    content: str = attr.ib()
    value: str = attr.ib()
    data: Optional[Data] = attr.ib(default=None, metadata=skip_if_null())


@attr.s(kw_only=True, frozen=True)
class SelectRowItem(ControlRowItem, PlaceholderMixin):
    @attr.s(kw_only=True, frozen=True)
    class Props(SerializableConfig):
        show_search: Optional[bool] = attr.ib(default=None, metadata=remap_skip_if_null("showSearch"))

    component_id = "select"

    available_values: Optional[list[SelectOption]] = attr.ib(
        default=None, metadata=remap_skip_if_null("availableValues")
    )
    control_props: Optional[Props] = attr.ib(default=None, metadata=remap_skip_if_null("controlProps"))


@attr.s(kw_only=True, frozen=True)
class RadioButtonRowItem(ControlRowItem):
    component_id = "radio_button"

    options: list[SelectableOption] = attr.ib()


@attr.s(kw_only=True, frozen=True)
class RadioGroupRowItemOption(SerializableConfig):
    @attr.s(kw_only=True, frozen=True)
    class ValueContent(SerializableConfig):
        text: str = attr.ib()
        hint_text: Optional[MarkdownStr] = attr.ib(default=None, metadata=remap_skip_if_null("hintText"))

    content: ValueContent = attr.ib()
    value: str = attr.ib()


@attr.s(kw_only=True, frozen=True)
class RadioGroupRowItem(RowItem, DisplayConditionsMixin, FormFieldMixin, InnerFieldMixin):
    @attr.s(kw_only=True, frozen=True)
    class Props(SerializableConfig):
        disabled: Optional[bool] = attr.ib(default=None, metadata=skip_if_null())

    component_id = "radio_group"

    options: list[RadioGroupRowItemOption] = attr.ib()
    default_value: Optional[str] = attr.ib(default=None, metadata=remap_skip_if_null("defaultValue"))
    control_props: Optional[Props] = attr.ib(default=None, metadata=remap_skip_if_null("controlProps"))

    def __attrs_post_init__(self) -> None:
        if self.default_value is not None and self.default_value not in (
            possible_values := [option.value for option in self.options]
        ):
            raise ValueError(f"Invalid default value for radio group: {self.default_value} is not in {possible_values}")


@attr.s(kw_only=True, frozen=True)
class CheckboxRowItem(RowItem, DisplayConditionsMixin, FormFieldMixin, InnerFieldMixin, DefaultValueMixin):
    @attr.s(kw_only=True, frozen=True)
    class Props(SerializableConfig):
        size: Optional[Literal["m", "l"]] = attr.ib(default=None, metadata=skip_if_null())
        qa: Optional[str] = attr.ib(
            default=None, metadata=skip_if_null()
        )  # UI-specific testing stuff, should be removed at some point

    component_id = "checkbox"

    text: str = attr.ib()
    control_props: Optional[Props] = attr.ib(default=None, metadata=remap_skip_if_null("controlProps"))


@attr.s(kw_only=True, frozen=True)
class DatepickerRowItem(ControlRowItem):
    @attr.s(kw_only=True, frozen=True)
    class Props(SerializableConfig):
        size: Optional[Literal["s", "m", "l", "xl"]] = attr.ib(default=None, metadata=skip_if_null())

    component_id = "datepicker"

    width: Optional[Width] = attr.ib(default=None, init=False, metadata=inner())  # not supported by Datepicker
    control_props: Optional[Props] = attr.ib(default=None, metadata=remap_skip_if_null("controlProps"))


@attr.s(kw_only=True, frozen=True)
class PlainTextRowItem(RowItem, DisplayConditionsMixin):
    component_id = "plain_text"

    text: str = attr.ib()
    hint_text: Optional[MarkdownStr] = attr.ib(default=None, metadata=remap_skip_if_null("hintText"))


@attr.s(kw_only=True, frozen=True)
class DescriptionRowItem(RowItem, DisplayConditionsMixin):
    component_id = "description"

    text: MarkdownStr = attr.ib()


@attr.s(kw_only=True, frozen=True)
class FileInputRowItem(RowItem, FormFieldMixin, InnerFieldMixin, DisplayConditionsMixin):
    component_id = "file-input"

lib/dl_api_connector/dl_api_connector/form_config/models/rows/customizable/components.py
datalens-tech/datalens-backend
from .components import (
    CacheTTLRow,
    CollapseRow,
    OAuthTokenRow,
)


__all__ = (
    "OAuthTokenRow",
    "CacheTTLRow",
    "CollapseRow",
)

lib/dl_api_connector/dl_api_connector/form_config/models/rows/prepared/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
    Optional,
)

import attr

from dl_api_connector.form_config.models.common import (
    InnerFieldName,
    SerializableConfig,
    skip_if_null,
)
from dl_api_connector.form_config.models.rows.base import FormRow


@attr.s(kw_only=True, frozen=True)
class PreparedRow(FormRow):
    type: ClassVar[str]

    class Inner(InnerFieldName):
        pass

    def as_dict(self) -> dict[str, Any]:
        return dict(
            type=self.type,
            **super().as_dict(),
        )


@attr.s(kw_only=True, frozen=True)
class DisabledMixin(SerializableConfig):
    disabled: Optional[bool] = attr.ib(default=None, metadata=skip_if_null())

lib/dl_api_connector/dl_api_connector/form_config/models/rows/prepared/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

import attr

from dl_api_connector.form_config.models.common import (
    OAuthApplication,
    SerializableConfig,
    remap_skip_if_null,
)
from dl_api_connector.form_config.models.rows.base import (
    DisplayConditionsMixin,
    FormFieldMixin,
    InnerFieldMixin,
)
from dl_api_connector.form_config.models.rows.prepared.base import (
    DisabledMixin,
    PreparedRow,
)


@attr.s(kw_only=True, frozen=True)
class OAuthTokenRow(PreparedRow, DisplayConditionsMixin, FormFieldMixin, DisabledMixin):
    type = "oauth"

    fake_value: Optional[str] = attr.ib(default=None, metadata=remap_skip_if_null("fakeValue"))

    application: OAuthApplication = attr.ib()
    label_text: Optional[str] = attr.ib(default=None, metadata=remap_skip_if_null("labelText"))
    button_text: Optional[str] = attr.ib(default=None, metadata=remap_skip_if_null("buttonText"))


@attr.s(kw_only=True, frozen=True)
class CacheTTLRow(PreparedRow, DisplayConditionsMixin, FormFieldMixin, DisabledMixin):
    type = "cache_ttl_sec"

    class Inner(PreparedRow.Inner):
        cache_ttl_mode = "cache_ttl_mode"

    label_text: Optional[str] = attr.ib(default=None, metadata=remap_skip_if_null("labelText"))


@attr.s(kw_only=True, frozen=True)
class CollapseRow(PreparedRow, DisplayConditionsMixin, FormFieldMixin, InnerFieldMixin):
    type = "collapse"

    @attr.s(kw_only=True, frozen=True)
    class Props(SerializableConfig):
        default_expanded: Optional[bool] = attr.ib(default=None, metadata=remap_skip_if_null("defaultIsExpand"))

    text: str = attr.ib()
    component_props: Optional[Props] = attr.ib(default=None, metadata=remap_skip_if_null("componentProps"))

lib/dl_api_connector/dl_api_connector/form_config/models/rows/prepared/components.py
datalens-tech/datalens-backend
import typing
from typing import Optional

import attr

from dl_api_connector.form_config.models import rows as C
from dl_api_connector.form_config.models.base import ConnectionFormMode
from dl_api_connector.form_config.models.common import (
    BooleanField,
    CommonFieldName,
    MarkdownStr,
)
from dl_api_connector.form_config.models.rows.base import TDisplayConditions
from dl_api_connector.i18n.localizer import Translatable
from dl_constants.enums import RawSQLLevel
from dl_i18n.localizer_base import Localizer
from dl_i18n.localizer_base import Translatable as BaseTranslatable


@attr.s
class RowConstructor:
    _localizer: Localizer = attr.ib()

    def host_row(
        self,
        default_value: Optional[str] = None,
        display_conditions: Optional[TDisplayConditions] = None,
        label_help_text: Optional[str] = None,
    ) -> C.CustomizableRow:
        text = self._localizer.translate(Translatable("field_host"))
        return C.CustomizableRow(
            items=[
                C.LabelRowItem(text=text, display_conditions=display_conditions, help_text=label_help_text),
                C.InputRowItem(
                    name=CommonFieldName.host,
                    width="l",
                    default_value=default_value,
                    display_conditions=display_conditions,
                ),
            ]
        )

    def port_row(
        self,
        label_text: BaseTranslatable = Translatable("field_port"),
        default_value: Optional[str] = None,
    ) -> C.CustomizableRow:
        text = self._localizer.translate(label_text)
        return C.CustomizableRow(
            items=[
                C.LabelRowItem(text=text),
                C.InputRowItem(
                    name=CommonFieldName.port,
                    width="s",
                    control_props=C.InputRowItem.Props(type="number"),
                    default_value=default_value,
                ),
            ]
        )

    def username_row(
        self,
        label_text: BaseTranslatable = Translatable("field_username"),
        default_value: Optional[str] = None,
        display_conditions: Optional[TDisplayConditions] = None,
    ) -> C.CustomizableRow:
        return C.CustomizableRow(
            items=[
                C.LabelRowItem(text=self._localizer.translate(label_text), display_conditions=display_conditions),
                C.InputRowItem(
                    name=CommonFieldName.username,
                    width="m",
                    default_value=default_value,
                    display_conditions=display_conditions,
                ),
            ]
        )

    def password_row(self, mode: ConnectionFormMode) -> C.CustomizableRow:
        label_text = self._localizer.translate(Translatable("field_password"))
        return C.CustomizableRow(
            items=[
                C.LabelRowItem(text=label_text),
                C.InputRowItem(
                    name=CommonFieldName.password,
                    width="m",
                    default_value="" if mode == ConnectionFormMode.create else None,
                    fake_value="******" if mode == ConnectionFormMode.edit else None,
                    control_props=C.InputRowItem.Props(type="password"),
                ),
            ]
        )

    def db_name_row(
        self,
        label_text: BaseTranslatable = Translatable("field_db-name"),
        default_value: Optional[str] = None,
        display_conditions: Optional[TDisplayConditions] = None,
    ) -> C.CustomizableRow:
        return C.CustomizableRow(
            items=[
                C.LabelRowItem(text=self._localizer.translate(label_text), display_conditions=display_conditions),
                C.InputRowItem(
                    name=CommonFieldName.db_name,
                    width="l",
                    default_value=default_value,
                    display_conditions=display_conditions,
                ),
            ]
        )

    def raw_sql_level_to_radio_group_option(
        self,
        raw_sql_level: RawSQLLevel,
    ) -> C.RadioGroupRowItemOption:
        return {
            RawSQLLevel.off: C.RadioGroupRowItemOption(
                value=raw_sql_level.value,
                content=C.RadioGroupRowItemOption.ValueContent(
                    text=self._localizer.translate(Translatable("value_raw-sql-level-off")),
                ),
            ),
            RawSQLLevel.subselect: C.RadioGroupRowItemOption(
                value=raw_sql_level.value,
                content=C.RadioGroupRowItemOption.ValueContent(
                    text=self._localizer.translate(Translatable("value_raw-sql-level-subselect")),
                    hint_text=self._localizer.translate(Translatable("value_raw-sql-level-subselect-note")),
                ),
            ),
            RawSQLLevel.dashsql: C.RadioGroupRowItemOption(
                value=raw_sql_level.value,
                content=C.RadioGroupRowItemOption.ValueContent(
                    text=self._localizer.translate(Translatable("value_raw-sql-level-dashsql")),
                    hint_text=self._localizer.translate(Translatable("value_raw-sql-level-dashsql-note")),
                ),
            ),
        }[raw_sql_level]

    def _raw_sql_level_default_options(self) -> list[C.RadioGroupRowItemOption]:
        return [
            self.raw_sql_level_to_radio_group_option(raw_sql_level)
            for raw_sql_level in (RawSQLLevel.off, RawSQLLevel.subselect, RawSQLLevel.dashsql)
        ]

    def raw_sql_level_row(
        self,
        default_value: RawSQLLevel = RawSQLLevel.off,
        options: Optional[list[C.RadioGroupRowItemOption]] = None,
        disabled: Optional[bool] = None,
    ) -> C.CustomizableRow:
        label_allow_subselect_hint_2 = self._localizer.translate(Translatable("label_allow-subselect-hint-2"))
        label_allow_subselect_hint_3 = self._localizer.translate(Translatable("label_allow-subselect-hint-3"))
        return C.CustomizableRow(
            items=[
                C.LabelRowItem(
                    align="start",
                    text=self._localizer.translate(Translatable("field_raw-sql-level")),
                    help_text=f"- {label_allow_subselect_hint_2}\n- {label_allow_subselect_hint_3}",
                ),
                C.RadioGroupRowItem(
                    name=CommonFieldName.raw_sql_level,
                    options=options if options is not None else self._raw_sql_level_default_options(),
                    default_value=self.raw_sql_level_to_radio_group_option(default_value).value,
                    control_props=C.RadioGroupRowItem.Props(disabled=disabled),
                ),
            ]
        )

    def access_token_input_row(
        self,
        mode: ConnectionFormMode,
        label_help_text: Optional[MarkdownStr] = None,
    ) -> C.CustomizableRow:
        return C.CustomizableRow(
            items=[
                C.LabelRowItem(
                    text=self._localizer.translate(Translatable("field_access-token")),
                    help_text=label_help_text,
                ),
                C.InputRowItem(
                    name=CommonFieldName.access_token,
                    width="l",
                    control_props=C.InputRowItem.Props(multiline=True),
                    fake_value="******" if mode == ConnectionFormMode.edit else None,
                ),
            ]
        )

    def auto_create_dash_row(self, default_value: Optional[bool] = True) -> C.CustomizableRow:
        return C.CustomizableRow(
            items=[
                C.CheckboxRowItem(
                    name=CommonFieldName.is_auto_create_dashboard,
                    text=self._localizer.translate(Translatable("field_auto-create-dashboard")),
                    inner=True,
                    default_value=default_value,
                    control_props=C.CheckboxRowItem.Props(qa="conn-auto-create-dash-checkbox"),
                ),
            ]
        )

    def collapse_advanced_settings_row(self) -> C.CollapseRow:
        return C.CollapseRow(
            inner=True,
            name=CommonFieldName.advanced_settings,
            text=self._localizer.translate(Translatable("label_advanced-settings")),
        )

    def ssl_rows(
        self,
        enabled_name: CommonFieldName,
        enabled_help_text: str,
        enabled_default_value: bool,
    ) -> typing.Sequence[C.CustomizableRow]:
        return [
            C.CustomizableRow(
                items=[
                    C.LabelRowItem(
                        text=self._localizer.translate(Translatable("label_ssl-enabled")),
                        display_conditions={CommonFieldName.advanced_settings: "opened"},
                        help_text=enabled_help_text,
                    ),
                    C.RadioButtonRowItem(
                        name=enabled_name,
                        options=[
                            C.SelectableOption(
                                text=self._localizer.translate(Translatable("value_ssl-enabled-off")),
                                value=BooleanField.off.value,
                            ),
                            C.SelectableOption(
                                text=self._localizer.translate(Translatable("value_ssl-enabled-on")),
                                value=BooleanField.on.value,
                            ),
                        ],
                        default_value=BooleanField.on.value if enabled_default_value else BooleanField.off.value,
                        display_conditions={CommonFieldName.advanced_settings: "opened"},
                    ),
                ]
            ),
            C.CustomizableRow(
                items=[
                    C.LabelRowItem(
                        text=self._localizer.translate(Translatable("label_ssl-ca")),
                        display_conditions={CommonFieldName.advanced_settings: "opened"},
                    ),
                    C.FileInputRowItem(
                        name=CommonFieldName.ssl_ca,
                        display_conditions={CommonFieldName.advanced_settings: "opened"},
                    ),
                ]
            ),
        ]

    def data_export_forbidden_row(self) -> C.CustomizableRow:
        return C.CustomizableRow(
            items=[
                C.LabelRowItem(
                    text=self._localizer.translate(Translatable("label_data-export-forbidden")),
                    display_conditions={CommonFieldName.advanced_settings: "opened"},
                    help_text=self._localizer.translate(Translatable("help_text_data-export-forbidden")),
                ),
                C.RadioButtonRowItem(
                    name=CommonFieldName.data_export_forbidden,
                    options=[
                        C.SelectableOption(
                            text=self._localizer.translate(Translatable("value_data-export-forbidden-off")),
                            value=BooleanField.off.value,
                        ),
                        C.SelectableOption(
                            text=self._localizer.translate(Translatable("value_data-export-forbidden-on")),
                            value=BooleanField.on.value,
                        ),
                    ],
                    default_value=BooleanField.off.value,
                    display_conditions={CommonFieldName.advanced_settings: "opened"},
                ),
            ]
        )

lib/dl_api_connector/dl_api_connector/form_config/models/shortcuts/rows.py
datalens-tech/datalens-backend
import os

import attr

import dl_api_connector as package
from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_api_connector/dl_api_connector/i18n/localizer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
    Set,
)

from marshmallow import (
    ValidationError,
    fields,
)
import pytest

from dl_api_connector.api_schema.extras import (
    CreateMode,
    EditMode,
    FieldExtra,
)
from dl_api_connector.api_schema.top_level import BaseTopLevelSchema
from dl_testing.utils import get_log_record


def test_allowed_unknown_fields(caplog):
    caplog.set_level("INFO")

    def get_unk_fields_record():
        return get_log_record(
            caplog,
            predicate=lambda r: r.msg.startswith("Got unknown fields for schema"),
            single=True,
        )

    class SampleSchema(BaseTopLevelSchema[Dict[str, Any]]):
        TARGET_CLS = dict

        a = fields.String(required=True)
        b = fields.String(required=True, bi_extra=FieldExtra(editable=True))

        def create_object(self, data: Dict[str, Any]) -> dict:
            return dict(data)

        def update_object(self, obj: dict, data: Dict[str, Any]) -> dict:
            obj.update(data)
            return obj

        def get_allowed_unknown_fields(self) -> Set[str]:
            super_fields = super().get_allowed_unknown_fields()
            super_fields.update(
                [
                    "allowed_unknown_field",
                    "a",
                ]
            )
            return super_fields

    # Allowed unknown field create
    caplog.clear()

    loaded_data = SampleSchema(
        context={SampleSchema.CTX_KEY_OPERATIONS_MODE: CreateMode.create},
    ).load(dict(a="a", b="b", allowed_unknown_field="val"))
    assert loaded_data == dict(a="a", b="b")

    unk_record = get_unk_fields_record()
    assert unk_record.schema_unk_fields == ["allowed_unknown_field"]
    assert unk_record.schema_unk_fields_allowed == ["allowed_unknown_field"]
    assert unk_record.schema_unk_fields_disallowed == []

    # Disallowed unknown fields create
    caplog.clear()

    with pytest.raises(ValidationError) as validation_err:
        SampleSchema(
            context={SampleSchema.CTX_KEY_OPERATIONS_MODE: CreateMode.create},
        ).load(dict(a="", b="", len=1))

    assert validation_err.value.messages == dict(
        len=["Unknown field."],
    )

    unk_record = get_unk_fields_record()
    assert unk_record.schema_unk_fields == ["len"]
    assert unk_record.schema_unk_fields_allowed == []
    assert unk_record.schema_unk_fields_disallowed == ["len"]

    # Allowed unknown field update
    caplog.clear()

    edit_target = dict(a="old_val_a", b="old_val_b")
    SampleSchema(
        context={
            SampleSchema.CTX_KEY_OPERATIONS_MODE: EditMode.edit,
            SampleSchema.CTX_KEY_EDITABLE_OBJECT: edit_target,
        },
    ).load(
        dict(a="new_val_a", b="new_val_b", allowed_unknown_field="val")  # this field should be ignored
    )
    assert edit_target == dict(a="old_val_a", b="new_val_b")

    unk_record = get_unk_fields_record()
    assert unk_record.schema_unk_fields == ["a", "allowed_unknown_field"]
    assert unk_record.schema_unk_fields_allowed == ["a", "allowed_unknown_field"]
    assert unk_record.schema_unk_fields_disallowed == []

    # Disallowed unknown fields update
    caplog.clear()

    with pytest.raises(ValidationError) as validation_err:
        edit_target = dict(a="old_val_a", b="old_val_b")
        SampleSchema(
            context={
                SampleSchema.CTX_KEY_OPERATIONS_MODE: EditMode.edit,
                SampleSchema.CTX_KEY_EDITABLE_OBJECT: edit_target,
            },
        ).load(dict(b="new_val_b", some_unknown_field="val", a="new_val_a"))

    assert validation_err.value.messages == dict(
        some_unknown_field=["Unknown field."],
    )

    unk_record = get_unk_fields_record()
    assert unk_record.schema_unk_fields == ["a", "some_unknown_field"]
    assert unk_record.schema_unk_fields_allowed == ["a"]
    assert unk_record.schema_unk_fields_disallowed == ["some_unknown_field"]

lib/dl_api_connector/dl_api_connector_tests/unit/test_schemas_base.py
datalens-tech/datalens-backend
from enum import Enum
from typing import Optional

import attr

from dl_api_connector.form_config.models.common import (
    SerializableConfig,
    inner,
    remap,
    skip_if_null,
)


class MyEnum(Enum):
    one = "one"


@attr.s(kw_only=True, frozen=True)
class SomeItemConfig(SerializableConfig):
    my_field: MyEnum = attr.ib()


@attr.s(kw_only=True, frozen=True)
class SomeConfig(SerializableConfig):
    basic_field: str = attr.ib()
    enum_field: MyEnum = attr.ib()
    undefined: Optional[int] = attr.ib(default=None, metadata=skip_if_null())
    remapped: str = attr.ib(default="orig_value", metadata=remap("remapped_2"))
    list_field: list[SomeItemConfig] = attr.ib()
    dict_field: dict[MyEnum, list[SomeItemConfig]] = attr.ib()
    magic_field: int = attr.ib(default=637, metadata=inner())


def test_config_as_dict():
    my_config = SomeConfig(
        basic_field="asdf",
        enum_field=MyEnum.one,
        list_field=[
            SomeItemConfig(my_field=MyEnum.one),
        ],
        dict_field={
            MyEnum.one: [
                SomeItemConfig(my_field=MyEnum.one),
                SomeItemConfig(my_field=MyEnum.one),
            ]
        },
    )

    result = my_config.as_dict()
    assert "undefined" not in result
    assert "magic_field" not in result
    assert result == {
        "basic_field": "asdf",
        "enum_field": "one",
        "remapped_2": "orig_value",
        "list_field": [
            {"my_field": "one"},
        ],
        "dict_field": {
            "one": [
                {"my_field": "one"},
                {"my_field": "one"},
            ],
        },
    }

lib/dl_api_connector/dl_api_connector_tests/unit/test_serializable_config.py
datalens-tech/datalens-backend
from __future__ import annotations

import typing
from typing import (
    Any,
    Dict,
)

from aiohttp import web
import attr

from dl_api_commons.aiohttp.aiohttp_wrappers import RequiredResource
from dl_core.aio.aiohttp_wrappers_data_core import DLRequestDataCore


@attr.s(frozen=True)
class AppWrapper:
    allow_query_cache_usage: bool = attr.ib()
    allow_notifications: bool = attr.ib()

    _APP_KEY = "DL_APP_WRAPPER"

    @classmethod
    def from_app(cls, app: web.Application) -> "AppWrapper":
        return app[cls._APP_KEY]

    def bind(self, app: web.Application) -> None:
        app[self._APP_KEY] = self


class DSAPIRequest(DLRequestDataCore):
    @property
    def app_wrapper(self) -> AppWrapper:
        return AppWrapper.from_app(self.request.app)

    # TODO FIX: Remove this property after all middleware will be switched to grab resources from view object
    @property
    def required_resources(self) -> typing.FrozenSet[RequiredResource]:
        from dl_api_lib.app.data_api.resources.base import BaseView

        view_cls = self.request.match_info.handler
        if isinstance(view_cls, type) and issubclass(view_cls, BaseView):
            return view_cls.get_required_resources(self.request.method)
        return frozenset()

    # TODO FIX: Move to common
    KEY_BODY_JSON = "body_json"

    # TODO FIX: Move to common
    def store_parsed_json_body(self, val: Dict[str, Any]) -> None:
        self._set_attr_once(self.KEY_BODY_JSON, val)

    # TODO FIX: Move to common
    @property
    def json(self) -> Dict[str, Any]:
        if self.KEY_BODY_JSON in self.request:
            return self.request[self.KEY_BODY_JSON]
        raise ValueError("JSON was not parsed for this request")

lib/dl_api_lib/dl_api_lib/aio/aiohttp_wrappers.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

from aiohttp import web
import attr

from dl_api_commons.aio.middlewares.error_handling_outer import (
    AIOHTTPErrorHandler,
    ErrorData,
    ErrorLevel,
)
from dl_api_lib.error_handling import (
    BIError,
    PublicAPIErrorSchema,
    RegularAPIErrorSchema,
)


def bi_error_to_error_data(bi_error: BIError, public_mode: bool, reason: Optional[str] = None) -> ErrorData:
    bi_error_schema = PublicAPIErrorSchema() if public_mode else RegularAPIErrorSchema()

    http_response_code = 500 if bi_error.http_code is None else bi_error.http_code
    return ErrorData(
        status_code=http_response_code,
        response_body=bi_error_schema.dump(bi_error),
        level=ErrorLevel.info if http_response_code < 500 else ErrorLevel.error,
    )


# noinspection PyDataclass
@attr.s
class DatasetAPIErrorHandler(AIOHTTPErrorHandler):
    public_mode: bool = attr.ib(kw_only=True)

    def _classify_error(self, err: Exception, request: web.Request) -> ErrorData:
        bi_error: BIError
        reason: Optional[str]

        if isinstance(err, web.HTTPException):
            bi_error = BIError(
                http_code=err.status,
                application_code_stack=(),
                message=err.reason,
                details={},
                debug={},
            )
            reason = err.reason
        else:
            bi_error = BIError.from_exception(err)
            reason = None
        return bi_error_to_error_data(bi_error, public_mode=self.public_mode, reason=reason)

lib/dl_api_lib/dl_api_lib/aio/middlewares/error_handling_outer.py
datalens-tech/datalens-backend
# TODO FIX: Move to common

from __future__ import annotations

import json
import logging

from aiohttp import web
from aiohttp.typedefs import Handler

from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_api_commons.logging import mask_sensitive_fields_by_name_in_json_recursive
from dl_api_lib.aio.aiohttp_wrappers import DSAPIRequest
from dl_api_lib.app.data_api.resources.base import (
    BaseView,
    RequiredResourceDSAPI,
)


LOGGER = logging.getLogger(__name__)


def json_body_middleware() -> AIOHTTPMiddleware:
    @web.middleware
    @DSAPIRequest.use_dl_request
    async def actual_body_log_middleware(dl_request: DSAPIRequest, handler: Handler) -> web.StreamResponse:
        request = dl_request.request
        view_cls = request.match_info.handler

        if (
            RequiredResourceDSAPI.JSON_REQUEST in dl_request.required_resources
            # TODO FIX: ensure methods
            and request.method.lower() in ("post", "put", "patch", "delete")
        ):
            # TODO CONSIDER: Maybe check content type first
            try:
                body = await request.json()
            # TODO CONSIDER: Maybe catch all exceptions
            except json.JSONDecodeError as err:
                raise web.HTTPBadRequest(reason="Invalid JSON in body") from err

            dl_request.store_parsed_json_body(body)

            dbg_body_data = mask_sensitive_fields_by_name_in_json_recursive(body)
            dbg_body = json.dumps(dbg_body_data)
            url = dl_request.url  # url with http->https override
            extra = dict(
                request_path=url,
                request_body=dbg_body,
            )
            LOGGER.debug("Body for %s: %s...", url, dbg_body[:100], extra=extra)

        if isinstance(view_cls, type) and issubclass(view_cls, BaseView):

            async def json_is_ready():
                raise Exception("Direct usage of request.json() is prohibited. Use DLRequest.json_body")

            request.json = json_is_ready

        return await handler(request)

    return actual_body_log_middleware

lib/dl_api_lib/dl_api_lib/aio/middlewares/json_body_middleware.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging

from aiohttp import web
from aiohttp.typedefs import Handler

from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_api_commons.aiohttp.aiohttp_wrappers import RequiredResourceCommon
from dl_api_commons.base_models import TenantCommon
from dl_api_lib.aio.aiohttp_wrappers import DSAPIRequest
from dl_constants.api_constants import DLHeadersCommon


LOGGER = logging.getLogger(__name__)


def public_api_key_middleware(api_key: str) -> AIOHTTPMiddleware:
    if not isinstance(api_key, str):
        raise TypeError(f"API key must be a string, not '{type(api_key)}'")

    @web.middleware
    @DSAPIRequest.use_dl_request
    async def actual_public_api_key_middleware(dl_request: DSAPIRequest, handler: Handler) -> web.StreamResponse:
        if RequiredResourceCommon.SKIP_AUTH in dl_request.required_resources:
            LOGGER.info("Auth was skipped due to SKIP_AUTH flag in target view")
            return await handler(dl_request.request)

        inbound_api_key = dl_request.get_single_header(DLHeadersCommon.PUBLIC_API_KEY, required=False)

        if inbound_api_key is None:
            raise web.HTTPForbidden(reason="public api key required")

        if inbound_api_key != api_key:
            LOGGER.info("Invalid API key, rejecting request...")
            raise web.HTTPForbidden(reason="invalid public api key")

        LOGGER.info("API key check passed")

        dl_request.replace_temp_rci(
            dl_request.temp_rci.clone(
                user_id="__ANONYMOUS_USER_OF_PUBLIC_DATALENS__",
                user_name="Anonymous (Public) User",
                tenant=TenantCommon(),
            )
        )

        return await handler(dl_request.request)

    return actual_public_api_key_middleware

lib/dl_api_lib/dl_api_lib/aio/middlewares/public_api_key_middleware.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Optional,
    Sequence,
)

from dl_api_commons.reporting.models import NotificationReportingRecord
from dl_api_commons.reporting.registry import ReportingRegistry
from dl_api_lib.api_common.data_types import bi_to_yql
from dl_api_lib.schemas.data import (
    DataApiV2ResponseSchema,
    DatasetFieldsResponseSchema,
    NotificationSchema,
)
from dl_api_lib.schemas.legend import LegendSchema
from dl_api_lib.schemas.pivot import (
    PivotHeaderInfoSchema,
    PivotItemSchema,
)
from dl_constants.enums import ManagedBy
from dl_core.us_dataset import Dataset
from dl_query_processing.enums import QueryType
from dl_query_processing.legend.field_legend import LegendItem
from dl_query_processing.merging.primitives import MergedQueryDataStream
from dl_utils.utils import enum_not_none


if TYPE_CHECKING:
    from dl_api_lib.pivot.primitives import DataCellVector
    from dl_api_lib.pivot.table import PivotTable
    from dl_constants.types import TJSONLike


LOGGER = logging.getLogger(__name__)


class DataRequestResponseSerializerV2Mixin:
    # TODO: Split this into subclasses

    @classmethod
    def make_data_response_v2_block_meta(cls, merged_stream: MergedQueryDataStream) -> List[Dict[str, Any]]:
        return [
            {
                "block_id": block.block_id,
                "query": block.debug_query,
            }
            for block in merged_stream.meta.blocks
        ]


class DataRequestResponseSerializer(DataRequestResponseSerializerV2Mixin):
    """Serializes data request responses"""

    _SHOW_QUERY_FOR_QUERY_TYPES = frozenset(
        (
            QueryType.result,
            QueryType.pivot,
        )
    )

    @staticmethod
    def get_yql_schema(legend_items: List[LegendItem]) -> list:
        """
        Create YQL schema for given list of fields (`self._query_spec.select_specs`).
        Is needed in the API response to be used by frontend to render results
        """
        return [
            "ListType",
            [
                "StructType",
                [
                    [
                        item.title,  # No longer used by frontend
                        ["OptionalType", ["DataType", bi_to_yql(enum_not_none(item.data_type))]],
                    ]
                    for item in legend_items  # FIXME
                ],
            ],
        ]

    @classmethod
    def make_data_response_v1(
        cls,
        merged_stream: MergedQueryDataStream,
        totals: Optional[Sequence] = None,
        totals_query: Optional[str] = None,
        data_export_forbidden: Optional[bool] = None,
        fields_data: Optional[List[Dict[str, Any]]] = None,
    ) -> Dict[str, Any]:
        legend_item_ids = merged_stream.legend_item_ids
        assert legend_item_ids is not None  # in v1 there is only one stream
        query_type = merged_stream.meta.blocks[0].query_type
        if query_type == QueryType.value_range:
            legend_item_ids = [legend_item_ids[0]]
        legend_items = [merged_stream.legend.get_item(liid) for liid in legend_item_ids]
        data: Dict[str, Any] = {
            "result": {
                "data": {
                    "Data": [row.data for row in merged_stream.rows],
                    "Type": cls.get_yql_schema(legend_items),
                },
            },
        }

        if merged_stream.meta.blocks[0].debug_query and query_type in cls._SHOW_QUERY_FOR_QUERY_TYPES:
            data["result"]["query"] = merged_stream.meta.blocks[0].debug_query  # type: ignore  # TODO: fix
        if totals is not None or totals_query is not None:
            # Note: `totals is None` also happens when everything was filtered out.
            data["result"]["totals"] = totals  # type: ignore  # TODO: fix
            data["result"]["totals_query"] = totals_query  # type: ignore  # TODO: fix
        if data_export_forbidden is not None:
            data["result"]["data_export_forbidden"] = data_export_forbidden  # type: ignore  # TODO: fix
        if fields_data is not None:
            data["result"]["fields"] = fields_data  # type: ignore  # TODO: fix

        return data

    @classmethod
    def make_data_response_v2(
        cls,
        merged_stream: MergedQueryDataStream,
        reporting_registry: Optional[ReportingRegistry] = None,
        data_export_forbidden: Optional[bool] = None,
    ) -> Dict[str, Any]:
        data: Dict[str, Any] = {}

        if data_export_forbidden is not None:
            data["data_export_forbidden"] = data_export_forbidden

        data["fields"] = merged_stream.legend.items

        data_rows: List[dict] = [{"data": row.data, "legend": row.legend_item_ids} for row in merged_stream.rows]
        data["result_data"] = [{"rows": data_rows}]

        data["blocks"] = cls.make_data_response_v2_block_meta(merged_stream=merged_stream)
        if reporting_registry is not None:
            data["notifications"] = reporting_registry.get_records_of_type(NotificationReportingRecord)

        data = DataApiV2ResponseSchema().dump(data)
        return data


def get_fields_data_raw(dataset: Dataset, for_result: bool = False) -> List[Dict[str, Any]]:
    return [
        {
            "title": fld.title,
            "guid": fld.guid,
            "data_type": fld.data_type,
            "calc_mode": fld.calc_mode,
            **(
                {
                    "hidden": fld.hidden,
                    "type": fld.type,
                }
                if not for_result
                else {}
            ),
        }
        for fld in dataset.result_schema
        if fld.managed_by == ManagedBy.user
    ]


def get_fields_data_serializable(dataset: Dataset, for_result: bool = False) -> List[Dict[str, TJSONLike]]:
    fields = get_fields_data_raw(dataset, for_result=for_result)
    return DatasetFieldsResponseSchema().dump(dict(fields=fields))["fields"]


class PivotDataRequestResponseSerializer(DataRequestResponseSerializerV2Mixin):
    @classmethod
    def make_pivot_response(
        cls,
        merged_stream: MergedQueryDataStream,
        pivot_table: PivotTable,
        reporting_registry: Optional[ReportingRegistry] = None,
    ) -> dict:
        legend_data = LegendSchema().dump(merged_stream.legend)
        pivot_spec_data = {
            "structure": [PivotItemSchema().dump(item) for item in pivot_table.pivot_legend.items],
        }
        # TODO: Dump pivot_table.pivot_legend
        block_meta_list = cls.make_data_response_v2_block_meta(merged_stream=merged_stream)

        def _serialize_cell_vector(value_vector: Optional[DataCellVector]) -> Optional[list[list[Any]]]:
            if value_vector is None:
                return None

            return [
                # [value, legend_item_id, pivot_item_id]
                [pivot_table.unpack_value(cell[0]), cell[1], cell[2]]
                for cell in value_vector.cells
            ]

        table_data = {
            "columns": [
                [_serialize_cell_vector(value_vector) for value_vector in col.values]
                for col in pivot_table.get_columns()
            ],
            "columns_with_info": [
                {
                    "cells": [_serialize_cell_vector(value_vector) for value_vector in col.values],
                    "header_info": PivotHeaderInfoSchema().dump(col.info),
                }
                for col in pivot_table.get_columns()
            ],
            "row_dimension_headers": [
                _serialize_cell_vector(flat_header) for flat_header in pivot_table.get_row_dim_headers()
            ],
            "rows": [
                {
                    "header": [_serialize_cell_vector(value_vector) for value_vector in row.header.values],
                    "header_with_info": {
                        "cells": [_serialize_cell_vector(value_vector) for value_vector in row.header.values],
                        "header_info": PivotHeaderInfoSchema().dump(row.header.info),
                    },
                    "values": [_serialize_cell_vector(value_vector) for value_vector in row.values],
                }
                for row in pivot_table.get_rows()
            ],
        }

        response_data = {
            **legend_data,
            "pivot": pivot_spec_data,
            "pivot_data": table_data,
            "blocks": block_meta_list,
        }

        if reporting_registry is not None:
            response_data["notifications"] = [
                NotificationSchema().dump(notification)
                for notification in reporting_registry.get_records_of_type(NotificationReportingRecord)
            ]
        return response_data

lib/dl_api_lib/dl_api_lib/api_common/data_serialization.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.enums import UserDataType


_BI_TO_YQL = {
    UserDataType.string: "String",
    UserDataType.integer: "Int64",
    UserDataType.float: "Double",
    UserDataType.date: "Date",
    UserDataType.datetime: "Datetime",
    UserDataType.datetimetz: "DatetimeTZ",
    UserDataType.genericdatetime: "GenericDatetime",
    UserDataType.boolean: "Bool",
    UserDataType.geopoint: "GeoPoint",
    UserDataType.geopolygon: "GeoPolygon",
    UserDataType.uuid: "UUID",
    UserDataType.markup: "Markup",
    UserDataType.array_str: "ArrayStr",
    UserDataType.array_int: "ArrayInt",
    UserDataType.array_float: "ArrayFloat",
    UserDataType.tree_str: "TreeStr",
    # Should not ever be in the output: `UserDataType.unsupported`
}


# TODO: Legacy stuff. Should be removed with data-api-v1
def bi_to_yql(bi_type: UserDataType) -> str:
    return _BI_TO_YQL[bi_type]

lib/dl_api_lib/dl_api_lib/api_common/data_types.py
datalens-tech/datalens-backend
from __future__ import annotations

from copy import deepcopy
import logging
from typing import (
    NamedTuple,
    Optional,
)

import attr

from dl_api_lib import exc
from dl_api_lib.dataset.utils import allow_rls_for_dataset
from dl_api_lib.service_registry.service_registry import ApiServiceRegistry
from dl_api_lib.utils.rls import FieldRLSSerializer
from dl_app_tools.profiling_base import generic_profiler
from dl_constants.exc import (
    DEFAULT_ERR_CODE_API_PREFIX,
    GLOBAL_ERR_PREFIX,
)
from dl_core.base_models import (
    DefaultConnectionRef,
    DefaultWhereClause,
)
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.editor import DatasetComponentEditor
from dl_core.data_source import get_parameters_hash
from dl_core.data_source.collection import DataSourceCollectionBase
from dl_core.db import are_raw_schemas_same
import dl_core.exc as core_exc
from dl_core.us_dataset import (
    Dataset,
    DataSourceRole,
)
from dl_core.us_manager.local_cache import USEntryBuffer
from dl_core.us_manager.us_manager import USManagerBase
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_utils.aio import await_sync


LOGGER = logging.getLogger(__name__)


class DatasetUpdateInfo(NamedTuple):
    """A simple container that contains info about what has changed in the dataset"""

    added_own_source_ids: list[str]
    updated_own_source_ids: list[str]


EMPTY_DS_UPDATE_INFO = DatasetUpdateInfo(
    added_own_source_ids=[],
    updated_own_source_ids=[],
)


@attr.s
class DatasetApiLoader:
    _service_registry: ApiServiceRegistry = attr.ib(kw_only=True)

    def update_dataset_from_body(
        self,
        dataset: Dataset,
        us_manager: USManagerBase,
        dataset_data: Optional[dict],
        allow_rls_change: bool = True,
    ) -> DatasetUpdateInfo:
        update_info = EMPTY_DS_UPDATE_INFO

        if dataset_data:
            update_info = self.populate_dataset_from_body(
                dataset=dataset,
                us_manager=us_manager,
                body=dataset_data,
                allow_rls_change=allow_rls_change,
            )
        return update_info

    def _ensure_additional_connections_are_preloaded(
        self,
        dataset_data: Optional[dict],
        us_manager: USManagerBase,
    ) -> None:
        """
        Ensure that the connections used by the dataset are pre-loaded.
        This can be one of 2 cases:
        1. The connection is already a part of the dataset.
           In this case it should have already been loaded with the dataset,
           and nothing will be done.
        2. The connection is not present in the dataset, but is being added
           to it via the patch `dataset_data`.

        Both of these cases will work with sync US manager,
        but the second one will not with an async one.
        On the other hand, this should not be a problem because
        we don't allow modification of data sources in the async server (data-api).
        """

        if dataset_data is None:
            return

        for source_data in dataset_data.get("sources", []):
            if isinstance(us_manager, SyncUSManager):
                # Force pre-loading of the used connection
                connection_id = source_data["connection_id"]
                connection_ref = DefaultConnectionRef(conn_id=connection_id)
                try:
                    us_manager.ensure_entry_preloaded(connection_ref)
                except core_exc.ReferencedUSEntryNotFound:
                    # Ignore deleted connections here - an error will be raised
                    # when it is fetched from the buffer
                    pass
            else:
                pass  # TODO: Maybe validate that the patch doesn't contain any new connections?

    def _update_dataset_sources_from_body(
        self,
        dataset: Dataset,
        body: dict,
        us_entry_buffer: USEntryBuffer,
    ) -> tuple[list, list, set]:
        ds_accessor = DatasetComponentAccessor(dataset=dataset)
        ds_editor = DatasetComponentEditor(dataset=dataset)
        dsrc_coll_factory = self._service_registry.get_data_source_collection_factory(us_entry_buffer=us_entry_buffer)

        # sources
        updated_own_source_ids = []
        added_own_source_ids = []
        handled_source_ids = set()
        old_src_coll: Optional[DataSourceCollectionBase]
        for source_data in body.get("sources", []):
            source_id = source_data["id"]
            title = source_data.get("title") or None
            connection_id = source_data["connection_id"]
            source_type = source_data["source_type"]
            valid = source_data["valid"]
            handled_source_ids.add(source_id)

            # check whether it exists already
            old_dsrc_coll_spec = ds_accessor.get_data_source_coll_spec_opt(source_id=source_id)

            if old_dsrc_coll_spec is not None:
                old_src_coll = dsrc_coll_factory.get_data_source_collection(spec=old_dsrc_coll_spec)
                LOGGER.info("Source %s already exists, checking if it needs to be updated.", source_id)
                old_hash = old_src_coll.get_param_hash()
                new_hash = get_parameters_hash(
                    connection_id=connection_id,
                    source_type=source_type,
                    **source_data["parameters"],
                )  # not that this does not include title and raw_schema updates
                old_raw_schema = old_src_coll.get_cached_raw_schema(role=DataSourceRole.origin)
                schema_updated = not are_raw_schemas_same(old_raw_schema, source_data["raw_schema"])  # type: ignore  # TODO: fix

                old_index_info_set = old_src_coll.get_strict(role=DataSourceRole.origin).saved_index_info_set
                new_index_info_set = source_data["index_info_set"]
                indexes_updated = old_index_info_set != new_index_info_set

                should_update = new_hash != old_hash or schema_updated or indexes_updated
                # To do not mark source as changed if only indexes was changed
                should_update_and_notify = new_hash != old_hash or schema_updated

                assert not (should_update_and_notify and not should_update)

                if should_update:
                    # source really was updated
                    LOGGER.info("Source %s is an in-dataset source. It will be updated.", source_id)
                    ds_editor.update_data_source(
                        source_id=source_id,
                        role=DataSourceRole.origin,
                        created_from=source_type,
                        connection_id=connection_id,
                        raw_schema=source_data["raw_schema"],
                        index_info_set=new_index_info_set,
                        **source_data["parameters"],
                    )
                    if should_update_and_notify:
                        updated_own_source_ids.append(source_id)

                if old_src_coll is not None and (title != old_src_coll.title or valid != old_src_coll.valid):
                    LOGGER.info("Collection data %s has changed. Updating.", source_id)
                    ds_editor.update_data_source_collection(source_id=source_id, title=title, valid=valid)

            if old_dsrc_coll_spec is None:
                # data source doesn't exist - it has to be (re)created
                LOGGER.info("Creating new plain source %s for connection %s.", source_id, connection_id)
                ds_editor.add_data_source_collection(
                    source_id=source_id,
                    title=title,
                    managed_by=source_data["managed_by"],
                    valid=valid,
                )
                ds_editor.add_data_source(
                    source_id=source_id,
                    role=DataSourceRole.origin,
                    created_from=source_data["source_type"],
                    connection_id=connection_id,
                    raw_schema=source_data["raw_schema"],
                    index_info_set=source_data["index_info_set"],
                    parameters=source_data["parameters"],
                )
                added_own_source_ids.append(source_id)

        return updated_own_source_ids, added_own_source_ids, handled_source_ids

    @classmethod
    def _update_dataset_source_avatars_from_body(cls, dataset: Dataset, body: dict) -> tuple[str, set]:
        ds_accessor = DatasetComponentAccessor(dataset=dataset)
        ds_editor = DatasetComponentEditor(dataset=dataset)
        # source avatars
        root_avatar_id = None
        handled_source_avatar_ids = set()
        for avatar_data in body.get("source_avatars", []):
            handled_source_avatar_ids.add(avatar_data["id"])
            if ds_accessor.has_avatar(avatar_id=avatar_data["id"]):
                ds_editor.update_avatar(
                    avatar_id=avatar_data["id"],
                    source_id=avatar_data["source_id"],
                    title=avatar_data["title"],
                )
            else:
                ds_editor.add_avatar(
                    avatar_id=avatar_data["id"],
                    source_id=avatar_data["source_id"],
                    title=avatar_data["title"],
                    managed_by=avatar_data["managed_by"],
                )
            if avatar_data["is_root"]:
                root_avatar_id = avatar_data["id"]
        return root_avatar_id, handled_source_avatar_ids  # type: ignore  # TODO: fix

    @classmethod
    def _update_dataset_source_avatar_relations_from_body(cls, dataset: Dataset, body: dict) -> set:
        ds_accessor = DatasetComponentAccessor(dataset=dataset)
        ds_editor = DatasetComponentEditor(dataset=dataset)
        # avatar relations
        handled_avatar_relation_ids = set()
        for relation_data in body.get("avatar_relations", []):
            handled_avatar_relation_ids.add(relation_data["id"])
            if ds_accessor.has_avatar_relation(relation_id=relation_data["id"]):
                ds_editor.update_avatar_relation(
                    relation_id=relation_data["id"],
                    conditions=relation_data["conditions"],
                    join_type=relation_data["join_type"],
                )
            else:
                ds_editor.add_avatar_relation(
                    relation_id=relation_data["id"],
                    left_avatar_id=relation_data["left_avatar_id"],
                    right_avatar_id=relation_data["right_avatar_id"],
                    conditions=relation_data["conditions"],
                    join_type=relation_data["join_type"],
                    managed_by=relation_data["managed_by"],
                )
        return handled_avatar_relation_ids

    @staticmethod
    def _rls_list_to_set(rls_list):  # type: ignore  # TODO: fix
        return set(
            (
                rlse.field_guid,
                rlse.allowed_value,
                rlse.subject.subject_name,
                rlse.pattern_type,
            )
            for rlse in rls_list
        )

    def _update_dataset_rls_from_body(self, dataset: Dataset, body: dict, allow_rls_change: bool = True) -> None:
        if not allow_rls_for_dataset(dataset):
            return

        subject_resolver = None
        if allow_rls_change:
            # E.g. dataset editing in sync api
            subject_resolver = await_sync(self._service_registry.get_subject_resolver())

        for field in dataset.result_schema:
            rls_text_config = body["rls"].get(field.guid, "")
            # TODO: split into "pre-parse all fields", "fetch all names", "make field results",
            # to combine name-to-info fetching for all fields.
            if allow_rls_change:
                rls_entries = FieldRLSSerializer.from_text_config(
                    rls_text_config,
                    field.guid,
                    subject_resolver=subject_resolver,
                )
                dataset.rls.items = [rlse for rlse in dataset.rls.items if rlse.field_guid != field.guid]
                dataset.rls.items.extend(rls_entries)

            else:
                # e.g. preview request in async api
                rls_entries_pre = FieldRLSSerializer.from_text_config(
                    rls_text_config, field.guid, subject_resolver=None
                )
                saved_field_rls = [
                    rlse.ensure_removed_failed_subject_prefix()
                    for rlse in dataset.rls.items
                    if rlse.field_guid == field.guid
                ]
                if self._rls_list_to_set(saved_field_rls) != self._rls_list_to_set(rls_entries_pre):
                    raise exc.RLSConfigParsingError(
                        "For this feature to work, save dataset after editing the RLS config.", details=dict()
                    )
                # otherwise no effective config changes (that are worth checking in preview)

    @classmethod
    def _update_dataset_obligatory_filters_from_body(cls, dataset: Dataset, body: dict) -> str:  # type: ignore  # TODO: fix
        # obligatory_filters
        ds_accessor = DatasetComponentAccessor(dataset=dataset)
        ds_editor = DatasetComponentEditor(dataset=dataset)
        handled_oblig_filter_ids = set()
        for filter_info in body.get("obligatory_filters") or ():
            handled_oblig_filter_ids.add(filter_info.id)
            default_where_clauses = [
                DefaultWhereClause(operation=default_filter.operation, values=default_filter.values)
                for default_filter in filter_info.default_filters
            ]

            filter_by_field_guid = ds_accessor.get_obligatory_filter_opt(field_guid=filter_info.field_guid)
            if filter_by_field_guid is not None and filter_by_field_guid.id != filter_info.id:
                ds_editor.remove_obligatory_filter(obfilter_id=filter_by_field_guid.id)

            if ds_accessor.has_obligatory_filter(obfilter_id=filter_info.id):
                ds_editor.update_obligatory_filter(
                    obfilter_id=filter_info.id, default_filters=default_where_clauses, valid=filter_info.valid
                )
            else:
                ds_editor.add_obligatory_filter(
                    obfilter_id=filter_info.id,
                    field_guid=filter_info.field_guid,
                    default_filters=default_where_clauses,
                    managed_by=filter_info.managed_by,
                    valid=filter_info.valid,
                )
        for oblig_filter in ds_accessor.get_obligatory_filter_list():
            if oblig_filter.id not in handled_oblig_filter_ids:
                ds_editor.remove_obligatory_filter(obfilter_id=oblig_filter.id)

    @generic_profiler("populate-dataset-from-body")
    def populate_dataset_from_body(
        self,
        dataset: Dataset,
        us_manager: USManagerBase,
        body: dict,
        allow_rls_change: bool = True,
    ) -> DatasetUpdateInfo:
        """
        Synchronize dataset object with configuration received in ``body``.
        """

        self._ensure_additional_connections_are_preloaded(dataset_data=body, us_manager=us_manager)

        ds_accessor = DatasetComponentAccessor(dataset=dataset)
        ds_editor = DatasetComponentEditor(dataset=dataset)

        if dataset.revision_id != body["revision_id"]:
            raise exc.DatasetRevisionMismatch()

        # fields (result_schema)
        ds_editor.set_result_schema(body.get("result_schema", []))
        if body.get("result_schema_aux"):
            dataset.data.result_schema_aux = body["result_schema_aux"]

        updated_own_source_ids, added_own_source_ids, handled_source_ids = self._update_dataset_sources_from_body(
            dataset=dataset, body=body, us_entry_buffer=us_manager.get_entry_buffer()
        )
        root_avatar_id, handled_source_avatar_ids = self._update_dataset_source_avatars_from_body(
            dataset=dataset, body=body
        )
        handled_avatar_relation_ids = self._update_dataset_source_avatar_relations_from_body(dataset=dataset, body=body)

        # cleanup (done in reverse order because of dependencies)
        for relation in ds_accessor.get_avatar_relation_list():
            if relation.id not in handled_avatar_relation_ids:
                ds_editor.remove_avatar_relation(relation_id=relation.id)
        for avatar in ds_accessor.get_avatar_list():
            if avatar.id not in handled_source_avatar_ids:
                ds_editor.remove_avatar(avatar_id=avatar.id)
        for source_id in set(ds_accessor.get_data_source_id_list()) - handled_source_ids:
            ds_editor.remove_data_source_collection(source_id=source_id)

        # additional tweaks
        if root_avatar_id is not None:
            if not ds_accessor.get_avatar_strict(root_avatar_id).is_root:
                ds_editor.set_root_avatar(avatar_id=root_avatar_id, rebuild_relations=False)

        # rls
        self._update_dataset_rls_from_body(dataset=dataset, body=body, allow_rls_change=allow_rls_change)

        # errors
        if "component_errors" in body:
            dataset.error_registry.items = deepcopy(body["component_errors"].items)
            for item in dataset.error_registry.items:
                for error in item.errors:
                    if error.code[:2] == [GLOBAL_ERR_PREFIX, DEFAULT_ERR_CODE_API_PREFIX]:
                        error.code = error.code[2:]

        self._update_dataset_obligatory_filters_from_body(dataset=dataset, body=body)

        return DatasetUpdateInfo(
            added_own_source_ids=added_own_source_ids,
            updated_own_source_ids=updated_own_source_ids,
        )

lib/dl_api_lib/dl_api_lib/api_common/dataset_loader.py
datalens-tech/datalens-backend
from __future__ import annotations

import hashlib
import json
from typing import (
    Any,
    List,
)

import attr

from dl_api_lib.request_model.data import FieldAction
from dl_core.serialization import RedisDatalensDataJSONEncoder
from dl_core.us_manager.mutation_cache.mutation_key_base import MutationKey


class MutationKeySerializationError(ValueError):
    pass


@attr.s(auto_attribs=True, kw_only=True, frozen=True)
class UpdateDatasetMutationKey(MutationKey):
    _dumped: str
    _hash: str

    def get_collision_tier_breaker(self) -> Any:
        return self._dumped

    def get_hash(self) -> str:
        return self._hash

    @classmethod
    def create(cls, dataset_revision_id: str, updates: List[FieldAction]) -> UpdateDatasetMutationKey:
        try:
            serialized = [upd.serialized for upd in updates]
        except Exception:
            raise MutationKeySerializationError()
        serialized.sort(key=lambda x: json.dumps(x, indent=None, sort_keys=True, cls=RedisDatalensDataJSONEncoder))
        dumped = json.dumps(
            dict(ds_rev=dataset_revision_id, mutation=serialized),
            sort_keys=True,
            indent=None,
            separators=(",", ":"),
            cls=RedisDatalensDataJSONEncoder,
        )
        hashed = hashlib.sha256(dumped.encode()).hexdigest()
        return UpdateDatasetMutationKey(dumped=dumped, hash=hashed)

lib/dl_api_lib/dl_api_lib/api_common/update_dataset_mutation_key.py
datalens-tech/datalens-backend
""" ... """

from __future__ import annotations

from functools import wraps
import json
import logging
import sys
from typing import (
    Callable,
    Optional,
)

import flask
from flask import current_app as capp
from flask import request
from flask_restx import Namespace
from marshmallow import Schema
from marshmallow import ValidationError as MValidationError
from werkzeug.exceptions import HTTPException

from dl_api_commons.logging import mask_sensitive_fields_by_name_in_json_recursive
from dl_api_lib import utils
from dl_api_lib.error_handling import (
    BIError,
    RegularAPIErrorSchema,
    status,
)
from dl_api_lib.schemas.main import get_api_model
from dl_utils.utils import maybe_postmortem


LOGGER = logging.getLogger(__name__)


def abort_request(code, message=None, response_data: Optional[dict] = None) -> None:
    """
    Copy-paste of `flask_restx.abort`, but without passing additional data through **kwargs.
    So we can to pass field "code" to response json.

    :param int code: The associated HTTP status code
    :param str message: An optional details message
    :param response_data: Any additional data to pass to the error payload
    :raise HTTPException:
    """
    if response_data is None:
        response_data = {}
    try:
        flask.abort(code)
    except HTTPException as e:
        if message:
            response_data["message"] = str(message)
        if response_data:
            e.data = response_data  # type: ignore  # TODO: fix
        raise


def schematic_request(  # type: ignore  # TODO: fix
    ns: Namespace,
    body: Optional[Schema] = None,
    query: Optional[Schema] = None,
    responses: Optional[dict[int, tuple[str, Schema]]] = None,
    dump: bool = True,
):
    """
    Decorator for REST API handlers.
    Allows the specification of marshmallow schemas for:
    - request body
    - request query
    - responses

    Translates exception classes into HTTP status codes
    """
    responses = responses or {}
    api_model_responses = {
        code: (desc, get_api_model(model, ns) if model is not None else None)
        for code, (desc, model) in responses.items()
    }
    body_schema = body
    query_schema = query

    def decorator(f):  # type: ignore  # TODO: fix
        if responses:
            for code, (description, model) in api_model_responses.items():
                f = ns.response(code, description, model)(f)

        @wraps(f)
        def wrapper(*args, **kwargs):  # type: ignore  # TODO: fix
            body = request.get_json()

            if LOGGER.isEnabledFor(logging.INFO):
                dbg_body_data = mask_sensitive_fields_by_name_in_json_recursive(body)
                dbg_body = json.dumps(dbg_body_data)
                url = request.url
                pfx = "http://"
                if url.startswith(pfx):
                    url = "https://" + url[len(pfx) :]
                LOGGER.info("Body (piece) for %s: %s...", url, dbg_body[:1000])
                extra = dict(
                    request_path=request.full_path,
                    # Tricky point:
                    # Can either use the `dbg_body` as a string, pre-dumped
                    # into json, in which case it will be handled as a
                    # continuous homogenous string.
                    # or give the `dbg_body_data` to the logger, which will
                    # dump it into json, then, in the logfeller, it will be
                    # loaded (full structure) into memory and then dumped into
                    # YSON.
                    request_body=dbg_body,
                )
                LOGGER.debug("Body for %s: %s...", url, dbg_body[:100], extra=extra)
                del dbg_body_data
                del dbg_body

            if query_schema is not None:
                try:
                    kwargs["query"] = query_schema.load(request.args)
                except MValidationError as err:
                    abort_request(status.BAD_REQUEST, err.messages)

            if body_schema is not None:
                try:
                    kwargs["body"] = body_schema.load(body)
                except MValidationError as err:
                    LOGGER.exception("Validation errors")
                    abort_request(status.BAD_REQUEST, err.messages)

            try:
                result = f(*args, **kwargs)
                if isinstance(result, tuple) and len(result) == 2:
                    result, code = result
                else:
                    code = status.OK
                code = int(code)
                if dump and code in responses:  # type: ignore  # TODO: fix
                    response_schema = responses[code][1]  # type: ignore  # TODO: fix
                    if response_schema:
                        result = response_schema.dump(result)
            except HTTPException:
                # which is the exception raised by flask's `abort`
                raise
            except Exception as err:
                ei = sys.exc_info()

                exc_api_prefix = capp.config.get("ERR_CODE_API_PREFIX")
                error_schema = RegularAPIErrorSchema(context=dict(api_prefix=exc_api_prefix))

                # TODO FIX: Make fallback for exceptions during creating error response
                bi_error = BIError.from_exception(err)
                error_info = error_schema.dump(bi_error)
                error_code = bi_error.http_code

                if error_code is None:
                    error_code = status.INTERNAL_SERVER_ERROR
                    LOGGER.exception("Caught an exception in request handler")
                    maybe_postmortem(ei=ei)
                else:
                    LOGGER.info("Regular exception fired", exc_info=True)

                abort_request(error_code, response_data=error_info)

            return result, code

        if body_schema is not None:
            # decorate wrapper
            wrapper = ns.doc(body=get_api_model(body_schema, ns), responses=api_model_responses)(wrapper)

        return wrapper

    return decorator


def with_profiler_stats(stats_dir: str, condition_check: Callable = None):  # type: ignore  # TODO: fix
    """Run function with cProfiler and save stats to file"""

    def decorator(func):  # type: ignore  # TODO: fix
        @wraps(func)
        def wrapper(*args, **kwargs):  # type: ignore  # TODO: fix
            if condition_check is None or condition_check(*args, **kwargs):
                with utils.profile_stats(stats_dir):
                    return func(*args, **kwargs)
            else:
                return func(*args, **kwargs)

        return wrapper

    return decorator

lib/dl_api_lib/dl_api_lib/api_decorators.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Generic,
    Optional,
    TypeVar,
)

import attr
import flask
from flask import Flask
from flask_marshmallow import Marshmallow
from statcommons.unistat.flask import (
    register_metrics,
    register_unistat_hax,
)

from dl_api_commons.flask.middlewares.aio_event_loop_middleware import AIOEventLoopMiddleware
from dl_api_commons.flask.middlewares.commit_rci_middleware import ReqCtxInfoMiddleware
from dl_api_commons.flask.middlewares.context_var_middleware import ContextVarMiddleware
from dl_api_commons.flask.middlewares.logging_context import RequestLoggingContextControllerMiddleWare
from dl_api_commons.flask.middlewares.request_id import RequestIDService
from dl_api_commons.flask.middlewares.tracing import (
    TracingContextMiddleware,
    TracingMiddleware,
)
from dl_api_lib.app.control_api.resources import init_apis
from dl_api_lib.app_common import SRFactoryBuilder
from dl_api_lib.app_common_settings import ConnOptionsMutatorsFactory
from dl_api_lib.app_settings import (
    ControlApiAppSettings,
    ControlApiAppTestingsSettings,
)
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_constants.enums import (
    ConnectionType,
    USAuthMode,
)
from dl_core import profiling_middleware
from dl_core.flask_utils.services_registry_middleware import ServicesRegistryMiddleware
from dl_core.flask_utils.us_manager_middleware import USManagerFlaskMiddleware
from dl_core.ping import register_ping_handler_hax


if TYPE_CHECKING:
    from dl_core.connection_models import ConnectOptions
    from dl_core.us_connection_base import ExecutorBasedMixin


@attr.s(frozen=True)
class EnvSetupResult:
    us_auth_mode: USAuthMode = attr.ib(kw_only=True)


TControlApiAppSettings = TypeVar("TControlApiAppSettings", bound=ControlApiAppSettings)


@attr.s(kw_only=True)
class ControlApiAppFactory(SRFactoryBuilder, Generic[TControlApiAppSettings], abc.ABC):
    _settings: TControlApiAppSettings = attr.ib()

    @abc.abstractmethod
    def set_up_environment(
        self,
        app: flask.Flask,
        testing_app_settings: Optional[ControlApiAppTestingsSettings] = None,
    ) -> EnvSetupResult:
        raise NotImplementedError()

    @property
    def _is_async_env(self) -> bool:
        return False

    def _get_conn_opts_mutators_factory(self) -> ConnOptionsMutatorsFactory:
        conn_opts_mutators_factory = ConnOptionsMutatorsFactory()

        def enable_index_fetching_mutator(
            conn_opts: ConnectOptions, conn: ExecutorBasedMixin
        ) -> Optional[ConnectOptions]:
            return conn_opts.clone(fetch_table_indexes=True)

        if self._settings.DO_DSRC_IDX_FETCH:
            conn_opts_mutators_factory.add_mutator(enable_index_fetching_mutator)

        return conn_opts_mutators_factory

    def create_app(
        self,
        connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
        testing_app_settings: Optional[ControlApiAppTestingsSettings] = None,
        close_loop_after_request: bool = True,
    ) -> flask.Flask:
        app = Flask(__name__)

        TracingMiddleware(
            url_prefix_exclude=(
                "/ping",
                "/unistat",
                "/metrics",
            ),
        ).wrap_flask_app(app)
        ContextVarMiddleware().wrap_flask_app(app)

        if close_loop_after_request:
            AIOEventLoopMiddleware().wrap_flask_app(app)

        app.url_map.strict_slashes = False

        # To avoid error messages replacement by flask_restx
        app.config["ERROR_404_HELP"] = False
        app.config["SWAGGER_UI_DOC_EXPANSION"] = "list"

        RequestLoggingContextControllerMiddleWare().set_up(app)
        TracingContextMiddleware().set_up(app)
        RequestIDService(
            request_id_app_prefix=self._settings.app_prefix,
        ).set_up(app)
        profiling_middleware.set_up(app, accept_outer_stages=False)

        # HAX to avoid bumping into inappropriate authorization added in 'before_request'.
        # TODO: move the inappropriate authorization to common handler wrap.
        register_ping_handler_hax(app)
        register_unistat_hax(app)
        register_metrics(app)

        env_setup_result = self.set_up_environment(app=app, testing_app_settings=testing_app_settings)

        ReqCtxInfoMiddleware().set_up(app)

        conn_opts_mutators_factory = self._get_conn_opts_mutators_factory()

        sr_factory = self.get_sr_factory(
            settings=self._settings,
            conn_opts_factory=conn_opts_mutators_factory,
            connectors_settings=connectors_settings,
        )

        ServicesRegistryMiddleware(
            services_registry_factory=sr_factory,
        ).set_up(app)

        USManagerFlaskMiddleware(
            crypto_keys_config=self._settings.CRYPTO_KEYS_CONFIG,
            us_base_url=self._settings.US_BASE_URL,
            us_master_token=self._settings.US_MASTER_TOKEN,
            us_auth_mode=env_setup_result.us_auth_mode,
        ).set_up(app)

        app.logger

        ma = Marshmallow()
        ma.init_app(app)

        init_apis(app)

        return app

lib/dl_api_lib/dl_api_lib/app/control_api/app.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import Optional

import flask
from flask import request
from flask_restx import Api
from werkzeug.exceptions import BadRequest

from dl_api_commons.access_control_common import AuthFailureError
from dl_api_commons.flask.middlewares.commit_rci_middleware import ReqCtxInfoMiddleware
from dl_core.exc import USReqException


API = Api(prefix="/api/v1")

LOGGER = logging.getLogger(__name__)


def init_apis(app: flask.Flask) -> None:
    from dl_api_lib.app.control_api.resources import connections  # noqa
    from dl_api_lib.app.control_api.resources import dataset  # noqa
    from dl_api_lib.app.control_api.resources import info  # noqa

    API.init_app(app)


@API.errorhandler(USReqException)
def handle_us_error(error):  # type: ignore
    resp = error.orig_exc.response
    try:
        text = resp.text
    except Exception:
        text = resp.content.decode("utf-8", errors="replace")
    LOGGER.warning(
        "Handling US error: %r (%s)",
        text,
        resp.status_code,
    )
    return {"message": text}, resp.status_code


@API.errorhandler(BadRequest)
def handle_bad_request(error):  # type: ignore
    rci = ReqCtxInfoMiddleware.get_last_resort_rci()
    req_id: Optional[str] = rci.request_id if rci is not None else None

    LOGGER.warning(
        "Bad request on %s: %s, req_id: %s",
        request.url,
        error.data,
        req_id,
    )
    return error.data, 400


@API.errorhandler(AuthFailureError)
def handle_auth_error(error):
    if error.response_code in (401, 403):
        if error.user_message is None:
            LOGGER.warning("No user message for AuthFailureError with defined response code", exc_info=True)
            user_message = "Auth failure"
        else:
            user_message = error.user_message
        http_response_code = error.response_code
    else:
        user_message = "Internal server error"
        http_response_code = 500

    return {"message": user_message}, http_response_code

lib/dl_api_lib/dl_api_lib/app/control_api/resources/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import os
import re
from typing import (
    Any,
    Optional,
)

import flask
from flask_restx import Resource

from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.flask.middlewares.commit_rci_middleware import ReqCtxInfoMiddleware
from dl_api_connector.api_schema.extras import OperationsMode
from dl_api_lib import api_decorators
from dl_api_lib.schemas.tools import prepare_schema_context
from dl_api_lib.service_registry.service_registry import ApiServiceRegistry
from dl_core.flask_utils.us_manager_middleware import USManagerFlaskMiddleware
from dl_core.us_manager.us_manager_sync import SyncUSManager


PROFILE_REQUESTS = os.environ.get("PROFILE_REQUESTS", "")
PROFILE_REQ_CLASSES = {v for v in set(os.environ.get("PROFILE_REQ_CLASSES", "").split(",")) if v}
PROFILE_REQ_METHODS = {v for v in set(os.environ.get("PROFILE_REQ_METHODS", "").lower().split(",")) if v}
PROFILE_REQ_PATH_RE = os.environ.get("PROFILE_REQ_PATH_RE")

# When turning on the profiler try to be as specific as possible. Example:
# PROFILE_REQUESTS="/tmp/cprofiler/"
# PROFILE_REQ_CLASSES="DatasetVersionResult"
# PROFILE_REQ_METHODS="POST"
# PROFILE_REQ_PATH_RE="/api/v1/datasets/njJYHhj8oikj/versions/draft/result"


def _profile_request_check(*args, **kwargs) -> bool:  # type: ignore  # TODO: fix
    """
    Check whether request handler should be profiled with current settings:
    - ``PROFILE_PATH_RE``: limit profiling only to requests with path matching the given regex
    """
    if PROFILE_REQ_PATH_RE:
        if not re.match(PROFILE_REQ_PATH_RE, flask.request.path):
            return False

    return True


COMMON_INFO_HEADERS = ("Referer", "X-Chart-Id")  # TODO: bi-common


class BIResourceMeta(type(Resource)):  # type: ignore  # TODO: fix
    """
    Metaclass for applying decorators to all BI handler classes and their methods
    (e.g. profiling, logging, etc.)
    """

    @staticmethod
    def _get_future_class_attr(bases, attrs, attr_name):  # type: ignore  # TODO: fix
        """Get attribute of future class described by ``bases`` and ``attrs``"""
        if attr_name in attrs:
            # attribute is redefined in new class
            return attrs[attr_name]
        else:
            # attribute is inherited from one of the base classes
            # (i.e. the first one in `bases` with this attribute)
            for base in bases:
                if hasattr(base, attr_name):
                    return getattr(base, attr_name)

    def __new__(mcs, name, bases, attrs):  # type: ignore  # TODO: fix
        profile_methods = mcs._get_future_class_attr(bases, attrs, "profile_methods")
        if PROFILE_REQ_METHODS:
            profile_methods = set(profile_methods) & PROFILE_REQ_METHODS

        if PROFILE_REQUESTS and profile_methods and (not PROFILE_REQ_CLASSES or (name in PROFILE_REQ_CLASSES)):
            # enable request profiler
            stats_dir = os.path.join(PROFILE_REQUESTS, name)
            attrs = {
                # decorate handler methods with profiler
                meth_name: (
                    api_decorators.with_profiler_stats(
                        os.path.join(stats_dir, meth_name),
                        # turn on checker only if there is something to check
                        condition_check=_profile_request_check if PROFILE_REQ_PATH_RE else None,
                    )(method)
                    if meth_name in profile_methods
                    else method
                )
                for meth_name, method in attrs.items()
            }

        return super().__new__(mcs, name, bases, attrs)


class BIResource(Resource, metaclass=BIResourceMeta):
    """Base class for all BI handler classes"""

    @classmethod
    def get_current_rci(cls) -> RequestContextInfo:
        return ReqCtxInfoMiddleware.get_request_context_info()

    @classmethod
    def get_schema_ctx(
        cls, schema_operations_mode: Optional[OperationsMode] = None, editable_object: Any = None
    ) -> dict:
        return prepare_schema_context(
            usm=cls.get_us_manager(),
            op_mode=schema_operations_mode,
            editable_object=editable_object,
        )

    @classmethod
    def get_us_manager(cls) -> SyncUSManager:
        return USManagerFlaskMiddleware.get_request_us_manager()

    @classmethod
    def get_service_us_manager(cls) -> SyncUSManager:
        return USManagerFlaskMiddleware.get_request_service_us_manager()

    @classmethod
    def get_service_registry(cls) -> ApiServiceRegistry:
        service_registry = cls.get_us_manager().get_services_registry()
        assert isinstance(service_registry, ApiServiceRegistry)
        return service_registry

    profile_methods = ("get", "post", "delete", "put", "patch", "options", "head")

lib/dl_api_lib/dl_api_lib/app/control_api/resources/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from flask import request
from flask_restx import abort
from marshmallow import ValidationError as MValidationError

from dl_api_commons.flask.middlewares.logging_context import put_to_request_context
from dl_api_connector.api_schema.extras import (
    CreateMode,
    EditMode,
)
from dl_api_lib import exc
from dl_api_lib.api_decorators import schematic_request
from dl_api_lib.app.control_api.resources import API
from dl_api_lib.app.control_api.resources.base import BIResource
from dl_api_lib.enums import USPermissionKind
from dl_api_lib.schemas.connection import (
    ConnectionInfoSourceSchemaQuerySchema,
    ConnectionInfoSourceSchemaResponseSchema,
    ConnectionSourcesQuerySchema,
    ConnectionSourceTemplatesResponseSchema,
    GenericConnectionSchema,
)
from dl_api_lib.utils import need_permission_on_entry
from dl_constants.enums import ConnectionType as ct
from dl_constants.exc import DLBaseException
from dl_core.data_source.type_mapping import get_data_source_class
from dl_core.data_source_merge_tools import make_spec_from_dict
from dl_core.exc import (
    DatabaseUnavailable,
    USPermissionRequired,
)
from dl_core.us_connection_base import ConnectionBase


if TYPE_CHECKING:
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase


LOGGER = logging.getLogger(__name__)

ns = API.namespace("Connections", path="/connections")


def _handle_conn_test_exc(exception: Exception):  # type: ignore  # TODO: fix
    if isinstance(exception, DLBaseException):
        raise exception
    else:
        LOGGER.exception("Got unhandled exception")
        raise DatabaseUnavailable() from exception


@ns.route("/test_connection_params")
class ConnectionParamsTester(BIResource):
    @schematic_request(ns=ns)
    def post(self):  # type: ignore  # TODO: fix
        service_registry = self.get_service_registry()
        schema = GenericConnectionSchema(context=self.get_schema_ctx(schema_operations_mode=CreateMode.test))
        body_json = dict(request.json)
        body_json["name"] = "mocked_name"  # FIXME: BI-4639
        try:
            conn: ConnectionBase = schema.load(body_json)
            conn_type = body_json.get("type")
        except MValidationError as err:
            return err.messages, 400

        conn.validate_new_data_sync(
            services_registry=service_registry,
        )

        try:
            conn.test(conn_executor_factory=service_registry.get_conn_executor_factory().get_sync_conn_executor)
        except NotImplementedError as err:
            raise exc.UnsupportedForEntityType(f"Connector {conn_type} does not support testing") from err
        except Exception as err:
            _handle_conn_test_exc(err)


@ns.route("/test_connection/<connection_id>")
class ConnectionTester(BIResource):
    @schematic_request(ns=ns)
    def post(self, connection_id):  # type: ignore  # TODO: fix
        usm = self.get_us_manager()
        service_registry = self.get_service_registry()
        conn = usm.get_by_id(connection_id, expected_type=ConnectionBase)
        conn_orig = usm.clone_entry_instance(conn)
        assert isinstance(conn_orig, ConnectionBase)  # for typing
        need_permission_on_entry(conn, USPermissionKind.read)

        body_json = dict(request.json)

        has_changes = len(body_json.keys()) > 0
        if has_changes:
            need_permission_on_entry(conn, USPermissionKind.edit)
        schema_ctx = self.get_schema_ctx(schema_operations_mode=EditMode.test, editable_object=None)

        schema_cls = GenericConnectionSchema().get_edit_schema_cls(conn)
        schema = schema_cls(context=schema_ctx)

        try:
            changes = schema.load(body_json)
            schema.update_object(conn, changes)
        except MValidationError as err:
            return err.messages, 400

        conn.validate_new_data_sync(services_registry=service_registry, changes=changes, original_version=conn_orig)

        try:
            conn.test(conn_executor_factory=service_registry.get_conn_executor_factory().get_sync_conn_executor)
        except NotImplementedError as err:
            raise exc.UnsupportedForEntityType(f"Connector {conn.conn_type.name} does not support testing") from err
        except Exception as e:
            _handle_conn_test_exc(e)


@ns.route("/")
class ConnectionsList(BIResource):
    @put_to_request_context(endpoint_code="ConnectionCreate")
    @schematic_request(ns=ns)
    def post(self):  # type: ignore  # TODO: fix
        us_manager = self.get_us_manager()

        conn_availability = self.get_service_registry().get_connector_availability()
        conn_type_is_available = conn_availability.check_connector_is_available(ct[request.json.get("type")])
        if not conn_type_is_available:
            # TODO: remove `abort` after migration to schematic_request decorator with common error handling
            abort(400, "This connection type is not editable")
            raise exc.UnsupportedForEntityType("This connection type is not editable")

        schema = GenericConnectionSchema(context=self.get_schema_ctx(schema_operations_mode=CreateMode.create))
        try:
            conn: ConnectionBase = schema.load(request.json)
        except MValidationError as e:
            return e.messages, 400

        conn.validate_new_data_sync(services_registry=self.get_service_registry())

        us_manager.save(conn)

        return {"id": conn.uuid}


@ns.route("/<connection_id>")
class ConnectionItem(BIResource):
    @put_to_request_context(endpoint_code="ConnectionGet")
    @schematic_request(
        ns=ns,
        responses={
            # 200: ('Success', GetConnectionResponseSchema()),
        },
    )
    def get(self, connection_id):  # type: ignore  # TODO: fix
        conn = self.get_us_manager().get_by_id(connection_id, expected_type=ConnectionBase)
        need_permission_on_entry(conn, USPermissionKind.read)

        result = GenericConnectionSchema(context=self.get_schema_ctx(EditMode.edit)).dump(conn)
        return result

    @put_to_request_context(endpoint_code="ConnectionDelete")
    @schematic_request(ns=ns)
    def delete(self, connection_id: str) -> None:
        us_manager = self.get_us_manager()
        conn = us_manager.get_by_id(connection_id, expected_type=ConnectionBase)
        need_permission_on_entry(conn, USPermissionKind.admin)

        us_manager.delete(conn)

    @put_to_request_context(endpoint_code="ConnectionUpdate")
    @schematic_request(ns=ns)
    def put(self, connection_id):  # type: ignore  # TODO: fix
        us_manager = self.get_us_manager()

        with us_manager.get_locked_entry_cm(ConnectionBase, connection_id) as conn:  # type: ignore  # TODO: fix
            need_permission_on_entry(conn, USPermissionKind.edit)
            conn_orig = us_manager.clone_entry_instance(conn)
            assert isinstance(conn_orig, ConnectionBase)  # for typing
            schema_ctx = self.get_schema_ctx(schema_operations_mode=EditMode.edit, editable_object=None)

            schema_cls = GenericConnectionSchema(context=schema_ctx).get_edit_schema_cls(conn)
            schema = schema_cls(context=schema_ctx)

            try:
                changes = schema.load(request.json)
                schema.update_object(conn, changes)
            except MValidationError as e:
                return e.messages, 400
            conn.validate_new_data_sync(
                services_registry=self.get_service_registry(),
                changes=changes,
                original_version=conn_orig,
            )
            us_manager.save(conn)


def _dump_source_templates(tpls) -> dict:  # type: ignore  # TODO: fix
    if tpls is None:
        return None  # type: ignore  # TODO: fix
    return [dict(tpl._asdict(), parameter_hash=tpl.get_param_hash()) for tpl in tpls]  # type: ignore  # TODO: fix


@ns.route("/<connection_id>/info/metadata_sources")
class ConnectionInfoMetadataSources(BIResource):
    @schematic_request(ns=ns, responses={200: ("Success", ConnectionSourceTemplatesResponseSchema())})
    def get(self, connection_id):  # type: ignore  # TODO: fix
        connection = self.get_us_manager().get_by_id(connection_id, expected_type=ConnectionBase)

        localizer = self.get_service_registry().get_localizer()
        source_template_templates = connection.get_data_source_template_templates(localizer=localizer)

        source_templates = []
        try:
            need_permission_on_entry(connection, USPermissionKind.read)
        except USPermissionRequired:
            pass
        else:
            source_templates = connection.get_data_source_local_templates()

        return {
            "sources": _dump_source_templates(source_templates),
            "freeform_sources": _dump_source_templates(source_template_templates),
        }


@ns.route("/<connection_id>/info/sources")
class ConnectionInfoSources(BIResource):
    @schematic_request(
        ns=ns,
        query=ConnectionSourcesQuerySchema(),
        responses={200: ("Success", ConnectionSourceTemplatesResponseSchema())},
    )
    def get(self, connection_id, query):  # type: ignore  # TODO: fix
        connection = self.get_us_manager().get_by_id(connection_id, expected_type=ConnectionBase)

        service_registry = self.get_service_registry()
        localizer = service_registry.get_localizer()
        source_template_templates = connection.get_data_source_template_templates(localizer=localizer)

        source_templates = []
        try:
            need_permission_on_entry(connection, USPermissionKind.read)
        except USPermissionRequired:
            pass
        else:
            source_templates = connection.get_data_source_templates(
                conn_executor_factory=service_registry.get_conn_executor_factory().get_sync_conn_executor,
            )

        search_text = query.get("search_text")
        if search_text is not None:
            # TODO: for some connections, do the filtering in the database.
            search_text = search_text.lower()
            source_templates = [item for item in source_templates if search_text in item.title.lower()]
        limit = query.get("limit")
        if limit is not None:  # Note that `load_default=1000` in the schema, so it should always be defined.
            source_templates = source_templates[:limit]
        return {
            "sources": _dump_source_templates(source_templates),
            "freeform_sources": _dump_source_templates(source_template_templates),
        }


@ns.route("/<connection_id>/info/source/schema")
class ConnectionInfoSourceSchema(BIResource):
    @schematic_request(
        ns=ns,
        body=ConnectionInfoSourceSchemaQuerySchema(),
        responses={200: ("Success", ConnectionInfoSourceSchemaResponseSchema())},
    )
    def post(self, connection_id: str, body: dict):  # type: ignore  # TODO: fix
        connection = self.get_us_manager().get_by_id(connection_id, expected_type=ConnectionBase)
        sr = self.get_service_registry()

        def conn_executor_factory_func() -> SyncConnExecutorBase:
            conn_executor = sr.get_conn_executor_factory().get_sync_conn_executor(conn=connection)
            return conn_executor

        need_permission_on_entry(connection, USPermissionKind.read)
        if not connection.is_dashsql_allowed:
            abort(400, "Not supported for connector without dashsql allowed")

        src = body["source"]
        dsrc_spec = make_spec_from_dict(source_type=src["source_type"], data=src["parameters"])
        dsrc_cls = get_data_source_class(src["source_type"])
        dsrc = dsrc_cls(spec=dsrc_spec, connection=connection)  # type: ignore

        schema_info = dsrc.get_schema_info(conn_executor_factory=conn_executor_factory_func)

        return {
            "raw_schema": schema_info.schema,
        }

lib/dl_api_lib/dl_api_lib/app/control_api/resources/connections.py
datalens-tech/datalens-backend
from __future__ import annotations

from http import HTTPStatus
import logging
from typing import (
    Any,
    Dict,
    Sequence,
)
import uuid

from dl_api_commons.flask.middlewares.logging_context import put_to_request_context
from dl_api_connector.api_schema.top_level import resolve_entry_loc_from_api_req_body
from dl_api_lib import (
    exc,
    utils,
)
from dl_api_lib.api_decorators import schematic_request
from dl_api_lib.app.control_api.resources import API
from dl_api_lib.app.control_api.resources.base import BIResource
from dl_api_lib.app.control_api.resources.dataset_base import DatasetResource
from dl_api_lib.const import DEFAULT_DATASET_LOCK_WAIT_TIMEOUT
from dl_api_lib.dataset.utils import (
    check_permissions_for_origin_sources,
    invalidate_sample_sources,
    log_dataset_field_stats,
)
from dl_api_lib.enums import USPermissionKind
import dl_api_lib.schemas.data
import dl_api_lib.schemas.dataset_base
import dl_api_lib.schemas.validation
from dl_constants.enums import ManagedBy
from dl_constants.exc import (
    CODE_OK,
    DEFAULT_ERR_CODE_API_PREFIX,
    GLOBAL_ERR_PREFIX,
)
from dl_core.base_models import (
    EntryLocation,
    PathEntryLocation,
    WorkbookEntryLocation,
)
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.editor import DatasetComponentEditor
from dl_core.constants import DatasetConstraints
from dl_core.us_dataset import Dataset
from dl_core.utils import generate_revision_id
import dl_query_processing.exc


LOGGER = logging.getLogger(__name__)

ns = API.namespace("Datasets", path="/datasets")

VALIDATION_OK_MESSAGE = "Validation was successful"


def _make_api_err_code(raw_code: Sequence[str]) -> str:
    return ".".join(
        (
            GLOBAL_ERR_PREFIX,
            DEFAULT_ERR_CODE_API_PREFIX,
            *raw_code,
        )
    )


@ns.route("/")
class DatasetCollection(DatasetResource):
    @classmethod
    def generate_dataset_location(cls, body: dict) -> EntryLocation:
        name = body.get("name", "Dataset {}".format(str(uuid.uuid4())))
        return resolve_entry_loc_from_api_req_body(
            name=name,
            workbook_id=body.get("workbook_id"),
            dir_path=body.get("dir_path", "datasets"),
        )

    @put_to_request_context(endpoint_code="DatasetCreate")
    @schematic_request(
        ns=ns,
        body=dl_api_lib.schemas.main.CreateDatasetSchema(),
        responses={
            200: ("Success", dl_api_lib.schemas.main.CreateDatasetResponseSchema()),
        },
    )
    def post(self, body):  # type: ignore  # TODO: fix
        """Create dataset"""
        us_manager = self.get_us_manager()
        dataset = Dataset.create_from_dict(
            Dataset.DataModel(name=""),  # TODO: Remove name - it's not used, but is required
            ds_key=self.generate_dataset_location(body),
            us_manager=us_manager,
        )
        ds_editor = DatasetComponentEditor(dataset=dataset)

        if "created_via" in body:
            ds_editor.set_created_via(created_via=body["created_via"])

        result_schema = body["dataset"].get("result_schema", [])
        if len(result_schema) > DatasetConstraints.FIELD_COUNT_LIMIT_SOFT:
            raise dl_query_processing.exc.DatasetTooManyFieldsFatal()

        loader = self.create_dataset_api_loader()
        loader.populate_dataset_from_body(dataset=dataset, body=body["dataset"], us_manager=us_manager)

        us_manager.save(dataset)

        LOGGER.info("New dataset was saved with ID %s", dataset.uuid)

        return self.make_dataset_response_data(dataset=dataset, us_entry_buffer=us_manager.get_entry_buffer())


@ns.route("/<string:dataset_id>")
class DatasetItem(BIResource):
    @schematic_request(
        ns=ns,
        responses={
            200: ("Success", None),  # type: ignore  # TODO: fix
            404: ("Not found", None),  # type: ignore  # TODO: fix
        },
    )
    def delete(self, dataset_id):  # type: ignore  # TODO: fix
        """Delete dataset"""
        us_manager = self.get_us_manager()
        ds, _ = DatasetResource.get_dataset(dataset_id=dataset_id, body={})
        utils.need_permission_on_entry(ds, USPermissionKind.admin)

        us_manager.delete(ds)


@ns.route("/<dataset_id>/fields")
class DatasetItemFields(BIResource):
    # TODO FIX: Move serialization logic to schema
    # TODO WARNING: Keep synced with async version
    @schematic_request(
        ns=ns,
        responses={
            200: ("Success", dl_api_lib.schemas.data.DatasetFieldsResponseSchema()),
        },
    )
    def get(self, dataset_id):  # type: ignore  # TODO: fix
        ds, _ = DatasetResource.get_dataset(dataset_id=dataset_id, body={})
        fields = [
            {
                "title": f.title,
                "guid": f.guid,
                "data_type": f.data_type,
                "hidden": f.hidden,
                "type": f.type,
                "calc_mode": f.calc_mode,
            }
            for f in ds.result_schema
            if f.managed_by == ManagedBy.user
        ]
        return {"fields": fields}


@ns.route("/<dataset_id>/copy")
class DatasetCopy(DatasetResource):
    @put_to_request_context(endpoint_code="DatasetCopy")
    @schematic_request(
        ns=ns,
        body=dl_api_lib.schemas.main.DatasetCopyRequestSchema(),
        responses={
            200: ("Success", dl_api_lib.schemas.main.DatasetCopyResponseSchema()),
            400: ("Failed", dl_api_lib.schemas.main.BadRequestResponseSchema()),
        },
    )
    def post(self, dataset_id, body):  # type: ignore  # TODO: fix
        copy_us_key = body["new_key"]
        us_manager = self.get_us_manager()
        ds, _ = self.get_dataset(dataset_id=dataset_id, body={})
        us_manager.load_dependencies(ds)
        orig_ds_loc = ds.entry_key
        copy_ds_loc: PathEntryLocation

        if isinstance(orig_ds_loc, PathEntryLocation):
            copy_ds_loc = PathEntryLocation(copy_us_key)
        else:
            raise exc.FeatureNotAvailable(message="Dataset copy in workbooks is not supported yet")

        utils.need_permission_on_entry(ds, USPermissionKind.edit)

        LOGGER.info("Going to copy dataset %s with new key %r", dataset_id, copy_us_key)
        ds_copy = us_manager.copy_entry(ds, key=copy_ds_loc)
        us_manager.save(ds_copy)
        LOGGER.info("Dataset copy was saved with ID %s", ds_copy.uuid)

        return self.make_dataset_response_data(dataset=ds_copy, us_entry_buffer=us_manager.get_entry_buffer())


@ns.route("/<dataset_id>/versions/<version>")
class DatasetVersionItem(DatasetResource):
    @put_to_request_context(endpoint_code="DatasetGet")
    @schematic_request(
        ns=ns,
        responses={
            200: ("Success", dl_api_lib.schemas.main.GetDatasetVersionResponseSchema()),
            400: ("Failed", dl_api_lib.schemas.main.BadRequestResponseSchema()),
        },
    )
    def get(self, dataset_id, version):  # type: ignore  # TODO: fix
        """Get dataset version"""
        us_manager = self.get_us_manager()
        ds, _ = self.get_dataset(dataset_id=dataset_id, body={})
        utils.need_permission_on_entry(ds, USPermissionKind.read)
        ds_dict = ds.as_dict()  # FIXME
        # TODO FIX: determine desired behaviour in case of workbooks
        ds_dict["key"] = ds.raw_us_key

        dl_loc = ds.entry_key
        if isinstance(dl_loc, WorkbookEntryLocation):
            ds_dict["workbook_id"] = dl_loc.workbook_id

        ds_dict["is_favorite"] = ds.is_favorite

        us_manager.load_dependencies(ds)
        ds_dict.update(self.make_dataset_response_data(dataset=ds, us_entry_buffer=us_manager.get_entry_buffer()))
        return ds_dict

    @put_to_request_context(endpoint_code="DatasetUpdate")
    @schematic_request(
        ns=ns,
        body=dl_api_lib.schemas.main.DatasetUpdateSchema(),
        responses={
            200: ("Success", dl_api_lib.schemas.dataset_base.DatasetContentSchema()),
            400: ("Failed", dl_api_lib.schemas.main.BadRequestResponseSchema()),
        },
    )
    def put(self, dataset_id: str, version: str, body: Dict[str, Any]) -> dict:
        """Update dataset version"""
        us_manager = self.get_us_manager()
        with us_manager.get_locked_entry_cm(
            Dataset, dataset_id, wait_timeout=DEFAULT_DATASET_LOCK_WAIT_TIMEOUT
        ) as ds:  # type: Dataset
            utils.need_permission_on_entry(ds, USPermissionKind.edit)
            us_manager.load_dependencies(ds)

            result_schema = body["dataset"].get("result_schema", [])
            if len(result_schema) > DatasetConstraints.FIELD_COUNT_LIMIT_SOFT:
                raise dl_query_processing.exc.DatasetTooManyFieldsFatal()

            ds_accessor = DatasetComponentAccessor(dataset=ds)

            old_sources = ds_accessor.get_data_source_id_list()
            loader = self.create_dataset_api_loader()
            update_info = loader.populate_dataset_from_body(dataset=ds, body=body["dataset"], us_manager=us_manager)
            new_sources = update_info.added_own_source_ids + update_info.updated_own_source_ids
            invalidate_sample_sources(
                dataset=ds,
                source_ids=new_sources,
                us_manager=us_manager,
            )

            # checks that all dataset sources have read rights
            sources_to_check = set(new_sources) - set(old_sources)
            check_permissions_for_origin_sources(
                dataset=ds,
                source_ids=sources_to_check,
                permission_kind=USPermissionKind.read,
                us_entry_buffer=us_manager.get_entry_buffer(),
            )

            ds_editor = DatasetComponentEditor(dataset=ds)
            ds_editor.set_revision_id(revision_id=generate_revision_id())
            us_manager.save(ds)

            return self.make_dataset_response_data(dataset=ds, us_entry_buffer=us_manager.get_entry_buffer())


@ns.route("/validators/dataset")
@ns.route("/<dataset_id>/versions/<version>/validators/schema")
class DatasetVersionValidator(DatasetResource):
    @put_to_request_context(endpoint_code="DatasetValidate")
    @schematic_request(
        ns=ns,
        body=dl_api_lib.schemas.validation.DatasetValidationSchema(),
        responses={
            200: ("Success", dl_api_lib.schemas.validation.DatasetValidationResponseSchema()),
            400: ("Failed", dl_api_lib.schemas.validation.DatasetValidationResponseSchema()),
        },
    )
    def post(self, dataset_id: str = None, version: str = None, body: dict = None):  # type: ignore  # TODO: fix
        """Validate dataset version schema"""
        us_manager = self.get_us_manager()
        dataset, _ = self.get_dataset(dataset_id=dataset_id, body=body)  # type: ignore  # TODO: fix
        dataset_validator_factory = self.get_service_registry().get_dataset_validator_factory()
        ds_validator = dataset_validator_factory.get_dataset_validator(ds=dataset, us_manager=us_manager)
        data = {}

        # apply updates
        try:
            ds_validator.apply_batch(action_batch=body.get("updates", ()))  # type: ignore  # TODO: fix
        except exc.DLValidationFatal as err:
            any_errors = True
            code = _make_api_err_code(exc.DLValidationFatal.err_code)
            message = err.message
        else:
            # collect connections not found in us
            ds_validator.collect_nonexistent_connection_errors()
            ds_validator.find_and_remove_phantom_error_refs()

            # check for errors
            any_errors = bool(dataset.error_registry.items)
            code = _make_api_err_code(exc.DLValidationError.err_code) if any_errors else CODE_OK
            message = exc.DLValidationError.default_message if any_errors else VALIDATION_OK_MESSAGE
            data.update(self.make_dataset_response_data(dataset=dataset, us_entry_buffer=us_manager.get_entry_buffer()))

        data.update(
            {
                "code": code,
                "message": message,
                "dataset_errors": [],  # TODO: Remove
            }
        )
        status = HTTPStatus.OK if not any_errors else HTTPStatus.BAD_REQUEST

        log_dataset_field_stats(dataset=dataset)

        return data, status


@ns.route("/validators/field")
@ns.route("/<dataset_id>/versions/<version>/validators/field")
class DatasetVersionFieldValidator(DatasetResource):
    @put_to_request_context(endpoint_code="DatasetFieldValidate")
    @schematic_request(
        ns=ns,
        body=dl_api_lib.schemas.validation.FieldValidationSchema(),
        responses={
            200: ("Success", dl_api_lib.schemas.validation.FieldValidationResponseSchema()),
            400: ("Failed", dl_api_lib.schemas.validation.FieldValidationResponseSchema()),
        },
    )
    def post(self, *, dataset_id: str = None, version: str = None, body):  # type: ignore  # TODO: fix
        """Validate formula field of dataset version"""
        us_manager = self.get_us_manager()
        dataset, _ = self.get_dataset(dataset_id=dataset_id, body=body)
        dataset_validator_factory = self.get_service_registry().get_dataset_validator_factory()
        ds_validator = dataset_validator_factory.get_dataset_validator(ds=dataset, us_manager=us_manager)
        formula = body["field"]["formula"]
        # TODO full validation (with aggregation check for this field), not just formula
        field_errors = ds_validator.get_single_formula_errors(formula)
        code = _make_api_err_code(exc.DLValidationError.err_code) if field_errors else CODE_OK
        message = exc.DLValidationError.default_message if field_errors else VALIDATION_OK_MESSAGE
        error_data = {
            "field_errors": [
                {
                    "title": body["field"].get("title"),
                    "guid": body["field"].get("guid"),
                    "errors": field_errors,
                }
            ]
            if field_errors
            else [],
            "code": code,
            "message": message,
        }
        status = HTTPStatus.OK if not field_errors else HTTPStatus.BAD_REQUEST
        return error_data, status

lib/dl_api_lib/dl_api_lib/app/control_api/resources/dataset.py
datalens-tech/datalens-backend
from __future__ import annotations

from copy import deepcopy
import logging
from typing import (
    Any,
    List,
    Optional,
    Set,
    Tuple,
)

from dl_api_lib.api_common.dataset_loader import (
    DatasetApiLoader,
    DatasetUpdateInfo,
)
from dl_api_lib.app.control_api.resources.base import BIResource
from dl_api_lib.dataset.dialect import resolve_dialect_name
from dl_api_lib.dataset.utils import (
    allow_rls_for_dataset,
    get_dataset_conn_types,
)
from dl_api_lib.enums import (
    BI_TYPE_AGGREGATIONS,
    CASTS_BY_TYPE,
)
from dl_api_lib.query.registry import (
    get_compeng_dialect,
    is_compeng_executable,
)
from dl_api_lib.service_registry.service_registry import ApiServiceRegistry
from dl_api_lib.utils.rls import FieldRLSSerializer
from dl_constants.enums import (
    AggregationFunction,
    BinaryJoinOperator,
    ConnectionType,
    ManagedBy,
    UserDataType,
)
from dl_constants.exc import (
    DEFAULT_ERR_CODE_API_PREFIX,
    GLOBAL_ERR_PREFIX,
)
from dl_core.backend_types import get_backend_type
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.data_source.base import DbInfo
from dl_core.dataset_capabilities import DatasetCapabilities
from dl_core.exc import (
    ReferencedUSEntryNotFound,
    UnexpectedUSEntryType,
    USObjectNotFoundException,
)
from dl_core.services_registry.top_level import ServicesRegistry
from dl_core.us_dataset import (
    Dataset,
    DataSourceRole,
)
from dl_core.us_manager.local_cache import USEntryBuffer
from dl_formula.core.dialect import (
    DialectCombo,
    StandardDialect,
    from_name_and_version,
)


LOGGER = logging.getLogger(__name__)

DEFAULT_DATABASE_DIALECT = StandardDialect.DUMMY
DEFAULT_MAX_AVATARS = 32
DEFAULT_MAX_CONNECTIONS = 20
UNAVAILABLE = object()


class DatasetResource(BIResource):
    """
    Base class for handlers that need to serialize and deserialize datasets
    """

    @classmethod
    def create_dataset_api_loader(cls) -> DatasetApiLoader:
        return DatasetApiLoader(
            service_registry=cls.get_service_registry(),
        )

    @classmethod
    def get_dataset(cls, dataset_id: Optional[str], body: dict) -> Tuple[Dataset, DatasetUpdateInfo]:
        us_manager = cls.get_us_manager()
        if dataset_id:
            try:
                dataset = us_manager.get_by_id(dataset_id, expected_type=Dataset)
            except UnexpectedUSEntryType:
                raise USObjectNotFoundException("Dataset with id {} does not exist".format(dataset_id))
        else:
            dataset = Dataset.create_from_dict(
                Dataset.DataModel(name=""),  # TODO: Remove name - it's not used, but is required
                us_manager=us_manager,
            )

        loader = cls.create_dataset_api_loader()
        update_info = loader.update_dataset_from_body(
            dataset=dataset,
            us_manager=us_manager,
            dataset_data=body.get("dataset"),
        )
        return dataset, update_info

    @classmethod
    def dump_dataset_data(
        cls,
        dataset: Dataset,
        service_registry: ServicesRegistry,
        us_entry_buffer: USEntryBuffer,
    ) -> dict:
        ds_accessor = DatasetComponentAccessor(dataset=dataset)
        dsrc_coll_factory = service_registry.get_data_source_collection_factory(us_entry_buffer=us_entry_buffer)

        data: dict[str, Any] = {}
        # generate info about sources
        sources = []
        preview_enabled_result = True
        for source_id in ds_accessor.get_data_source_id_list():
            dsrc_coll_spec = ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
            dsrc_coll = dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)

            origin_dsrc = dsrc_coll.get_strict(role=DataSourceRole.origin)
            connection_id = dsrc_coll.get_connection_id(DataSourceRole.origin)
            sources.append(
                {
                    "id": source_id,
                    "title": dsrc_coll.title or origin_dsrc.default_title,
                    "connection_id": connection_id,
                    "managed_by": dsrc_coll.managed_by,
                    "valid": dsrc_coll.valid,
                    "source_type": origin_dsrc.spec.source_type,
                    "raw_schema": origin_dsrc.saved_raw_schema,
                    "index_info_set": origin_dsrc.saved_index_info_set,
                    "parameters": origin_dsrc.get_parameters(),
                    "parameter_hash": dsrc_coll.get_param_hash(),
                }
            )
            if not origin_dsrc.preview_enabled:
                preview_enabled_result = origin_dsrc.preview_enabled
        data["sources"] = sources

        data["preview_enabled"] = preview_enabled_result

        # generate info about source avatars
        data["source_avatars"] = ds_accessor.get_avatar_list()

        # generate info about source relations
        data["avatar_relations"] = ds_accessor.get_avatar_relation_list()

        # fields (result_schema)
        result_schema = []
        for field in dataset.result_schema:
            field_data = field._asdict()
            field_data["autoaggregated"] = field.autoaggregated
            field_data["aggregation_locked"] = field.aggregation_locked
            result_schema.append(field_data)
        data["result_schema"] = result_schema
        data["result_schema_aux"] = dataset.data.result_schema_aux

        # rls
        rls = {}
        if allow_rls_for_dataset(dataset):
            for field in dataset.result_schema:
                field_rls = [e for e in dataset.rls.items if e.field_guid == field.guid]
                rls[field.guid] = FieldRLSSerializer.to_text_config(field_rls)
        data["rls"] = rls

        data["obligatory_filters"] = ds_accessor.get_obligatory_filter_list()

        data["component_errors"] = deepcopy(dataset.error_registry)
        for item in data["component_errors"].items:
            for error in item.errors:
                if error.code[:2] != [GLOBAL_ERR_PREFIX, DEFAULT_ERR_CODE_API_PREFIX]:
                    error.code = [GLOBAL_ERR_PREFIX, DEFAULT_ERR_CODE_API_PREFIX] + error.code

        data["created_via"] = dataset.created_via

        data["revision_id"] = dataset.revision_id

        return {"dataset": data}

    @classmethod
    def dump_option_data(
        cls,
        dataset: Dataset,
        service_registry: ApiServiceRegistry,
        us_entry_buffer: USEntryBuffer,
    ) -> dict:
        ds_accessor = DatasetComponentAccessor(dataset=dataset)
        dsrc_coll_factory = service_registry.get_data_source_collection_factory(us_entry_buffer=us_entry_buffer)
        sfm = service_registry.get_supported_functions_manager()
        capabilities = DatasetCapabilities(dataset=dataset, dsrc_coll_factory=dsrc_coll_factory)

        opt_data: dict[str, Any] = {}
        opt_data["preview"] = dict(
            enabled=capabilities.supports_preview(),
        )
        opt_data["join"] = dict(
            types=capabilities.get_supported_join_types(),
            operators=[BinaryJoinOperator.eq],
        )
        role = capabilities.resolve_source_role()  # the "main" role, not for preview

        # Determine dialect
        dialect: Optional[DialectCombo] = DEFAULT_DATABASE_DIALECT
        one_source_id = dataset.get_single_data_source_id()
        db_info: Optional[DbInfo]

        if one_source_id is not None:
            try:
                dsrc_coll_spec = ds_accessor.get_data_source_coll_spec_strict(source_id=one_source_id)
                dsrc_coll = dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
                dsrc = dsrc_coll.get_strict(role=role)
                db_info = dsrc.get_cached_db_info()
            except ReferencedUSEntryNotFound:
                db_info = None

            if db_info is not None:
                backend_type = capabilities.get_backend_type(role=role)
                dialect_name = resolve_dialect_name(backend_type=backend_type)
                dialect = from_name_and_version(
                    dialect_name=dialect_name,
                    dialect_version=db_info.version,
                )
                if not dialect:
                    LOGGER.error(
                        "Failed to get dialect from SA dialect %r and DB version %r", dialect_name.name, db_info.version
                    )
                    dialect = DEFAULT_DATABASE_DIALECT

        assert dialect is not None
        funcs_dialect = dialect
        conn_types = set(get_dataset_conn_types(dataset=dataset, us_entry_buffer=us_entry_buffer))
        is_compeng_executable_set = {is_compeng_executable(get_backend_type(conn_type)) for conn_type in conn_types}
        if len(is_compeng_executable_set) == 1 and next(iter(is_compeng_executable_set)):
            funcs_dialect = get_compeng_dialect()

        opt_data["data_types"] = dict(
            items=[
                dict(
                    type=user_type,
                    aggregations=sfm.get_supported_aggregations(dialect=funcs_dialect, user_type=user_type),
                    casts=CASTS_BY_TYPE.get(user_type, []),
                    filter_operations=sfm.get_supported_filters(dialect=funcs_dialect, user_type=user_type),
                )
                for user_type in UserDataType
            ],
        )

        opt_data["sources"] = dict(items=[])
        connection_ids = set()
        connection_types: Set[Optional[ConnectionType]] = set()
        for source_id in ds_accessor.get_data_source_id_list():
            dsrc_coll_spec = ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
            dsrc_coll = dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
            if dsrc_coll.managed_by != ManagedBy.user:
                continue
            dsrc = dsrc_coll.get_strict(role=DataSourceRole.origin)
            opt_data["sources"]["items"].append(
                dict(
                    id=source_id,
                    schema_update_enabled=dsrc.supports_schema_update,
                )
            )
            connection_id = dsrc_coll.get_connection_id(role=DataSourceRole.origin)
            if connection_id is not None:
                connection_ids.add(connection_id)
                connection_types.add(dsrc.conn_type)

        opt_data["sources"]["compatible_types"] = [
            dict(source_type=source_type) for source_type in capabilities.get_compatible_source_types()
        ]

        only_one_connection = len(connection_ids) <= 1

        opt_data["connections"] = dict(items=[])
        for conn_id in sorted(connection_ids):
            replacement_types = []
            repl_conn_types = capabilities.get_compatible_connection_types(ignore_connection_ids=[conn_id])
            for conn_type in repl_conn_types:
                # it seems we can do this check once
                if not only_one_connection:
                    continue
                item = dict(conn_type=conn_type)
                if item not in replacement_types:
                    replacement_types.append(item)

            opt_data["connections"]["items"].append(
                dict(
                    id=conn_id,
                    replacement_types=replacement_types,
                )
            )

        compatible_conn_types: List[dict] = []
        for conn_type in capabilities.get_compatible_connection_types():
            if connection_ids:  # There already are connections in the dataset
                continue
            compatible_conn_types.append(dict(conn_type=conn_type))
        opt_data["connections"]["compatible_types"] = compatible_conn_types

        opt_data["connections"]["max"] = DEFAULT_MAX_CONNECTIONS if compatible_conn_types else 1

        opt_data["schema_update_enabled"] = any(
            [dsrc_data["schema_update_enabled"] for dsrc_data in opt_data["sources"]["items"]]
        )

        opt_data["source_avatars"] = dict(
            # if JOINs are supported, then new tables can be added
            max=DEFAULT_MAX_AVATARS if capabilities.get_supported_join_types() else 1,
            items=[],
        )
        for avatar in ds_accessor.get_avatar_list():
            dsrc_coll_spec = ds_accessor.get_data_source_coll_spec_strict(source_id=avatar.source_id)
            dsrc_coll = dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
            dsrc = dsrc_coll.get_strict(role=DataSourceRole.origin)
            opt_data["source_avatars"]["items"].append(
                dict(
                    id=avatar.id,
                    schema_update_enabled=dsrc.supports_schema_update,
                )
            )

        opt_data["supports_offset"] = capabilities.supports_offset(role)

        opt_data["supported_functions"] = sfm.get_supported_function_names(dialect=dialect)

        opt_data["fields"] = dict(
            items=[],
        )
        for field in dataset.result_schema:
            if field.has_auto_aggregation or field.lock_aggregation:
                # Aggregation is disabled
                aggregations = [AggregationFunction.none]
            else:
                aggregations = BI_TYPE_AGGREGATIONS.get(field.cast, [])

            opt_data["fields"]["items"].append(
                dict(
                    guid=field.guid,
                    casts=CASTS_BY_TYPE.get(field.initial_data_type, ()),
                    aggregations=aggregations,
                )
            )

        return {"options": opt_data}

    def make_dataset_response_data(self, dataset: Dataset, us_entry_buffer: USEntryBuffer) -> dict:
        service_registry = self.get_service_registry()
        ds_dict = self.dump_dataset_data(
            dataset=dataset,
            us_entry_buffer=us_entry_buffer,
            service_registry=service_registry,
        )
        ds_dict.update(
            self.dump_option_data(
                dataset=dataset,
                us_entry_buffer=us_entry_buffer,
                service_registry=service_registry,
            )
        )
        ds_dict["id"] = dataset.uuid
        if dataset.permissions is not None:
            ds_dict["permissions"] = dataset.permissions

        return ds_dict

lib/dl_api_lib/dl_api_lib/app/control_api/resources/dataset_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import Schema
from marshmallow import fields as ma_fields

from dl_api_connector.form_config.models.base import ConnectionFormMode
from dl_api_lib.api_decorators import schematic_request
from dl_api_lib.app.control_api.resources import API
from dl_api_lib.app.control_api.resources.base import BIResource
from dl_api_lib.connection_forms.registry import (
    NoForm,
    get_connection_form_factory_cls,
)
from dl_api_lib.enums import BI_TYPE_AGGREGATIONS
from dl_api_lib.exc import UnsupportedForEntityType
from dl_api_lib.public.entity_usage_checker import PublicEnvEntityUsageChecker
from dl_api_lib.schemas.main import BadRequestResponseSchema
from dl_constants.enums import (
    ConnectionType,
    UserDataType,
)
from dl_core.exc import EntityUsageNotAllowed
from dl_core.us_dataset import Dataset


class FieldTypeInfoSchema(Schema):
    name = ma_fields.String()
    aggregations = ma_fields.List(ma_fields.String())


class GetFieldTypeCollectionResponseSchema(Schema):
    types = ma_fields.Nested(FieldTypeInfoSchema, many=True)


class DatasetsPublicityCheckerRequestSchema(Schema):
    datasets = ma_fields.List(ma_fields.String())


class DatasetsPublicityCheckerResponseSchema(Schema):
    class DatasetResponseSchema(Schema):
        dataset_id = ma_fields.String()
        allowed = ma_fields.Boolean()
        reason = ma_fields.String()

    result = ma_fields.Nested(DatasetResponseSchema, many=True)


ns = API.namespace("Info", path="/info")


@ns.route("/field_types")
class FieldTypeCollection(BIResource):
    @schematic_request(ns=ns, responses={200: ("Success", GetFieldTypeCollectionResponseSchema())})
    def get(self):  # type: ignore  # TODO: fix
        return {
            "types": [
                {"name": k.name, "aggregations": [x.name for x in v]}
                for k, v in BI_TYPE_AGGREGATIONS.items()
                if k
                not in (
                    UserDataType.uuid,
                    UserDataType.markup,
                    UserDataType.datetimetz,
                    UserDataType.datetime,
                    UserDataType.unsupported,
                )
            ]
        }


@ns.route("/datasets_publicity_checker")
class DatasetsPublicityChecker(BIResource):
    @schematic_request(
        ns=ns,
        body=DatasetsPublicityCheckerRequestSchema(),
        responses={200: ("Success", DatasetsPublicityCheckerResponseSchema())},
    )
    def post(self, body):  # type: ignore  # TODO: fix
        ds_ids = body["datasets"]
        responses = []
        us_manager = self.get_us_manager()

        public_usage_checker = PublicEnvEntityUsageChecker()

        for ds in self.get_us_manager().get_collection(
            Dataset, raise_on_broken_entry=True, include_data=True, ids=ds_ids
        ):
            try:
                us_manager.load_dependencies(ds)
                public_usage_checker.ensure_dataset_can_be_used(
                    rci=self.get_current_rci(),
                    dataset=ds,
                    us_manager=us_manager,
                )
            except EntityUsageNotAllowed as exc:
                allowed = False
                reason = exc.message
            else:
                allowed = True
                reason = None  # type: ignore  # TODO: fix

            responses.append(
                {
                    "dataset_id": ds.uuid,
                    "allowed": allowed,
                    "reason": reason,
                }
            )

        return {"result": responses}


@ns.route("/connectors")
class AvailableConnectorsCollection(BIResource):
    def get(self) -> dict:
        conn_availability = self.get_service_registry().get_connector_availability()
        localizer = self.get_service_registry().get_localizer()
        return conn_availability.as_dict(localizer=localizer)


@ns.route("/connectors/forms/<string:conn_type>/<string:form_mode>")
class ConnectorForm(BIResource):
    @schematic_request(
        ns=ns,
        responses={
            400: ("Failed", BadRequestResponseSchema()),
        },
    )
    def get(self, conn_type: str, form_mode: str) -> dict:
        try:
            ct = ConnectionType(conn_type)
        except ValueError:
            raise UnsupportedForEntityType(f"Unknown connector type: {conn_type}")

        try:
            mode = ConnectionFormMode(form_mode)
        except ValueError:
            raise UnsupportedForEntityType(f"Unknown form mode: {form_mode}")

        try:
            form_factory_cls = get_connection_form_factory_cls(ct)
        except NoForm:
            return {"form": None}

        localizer = self.get_service_registry().get_localizer()

        form_factory = form_factory_cls(mode=mode, localizer=localizer)
        form_config = form_factory.get_form_config(
            connector_settings=self.get_service_registry().get_connectors_settings(ct),
            tenant=self.get_current_rci().tenant,
        )

        return {
            "form": form_config.as_dict(),
        }


@ns.route("/internal/pseudo_workbook/<path:us_path>")
class WorkbookInfo(BIResource):
    @schematic_request(ns=ns)
    def get(self, us_path: str) -> dict:
        usm = self.get_us_manager()
        all_entries = usm.load_get_entries_at_path(us_path)

        return dict(
            connections={
                conn_dict["key"].split("/")[-1]: dict(
                    id=conn_dict["entryId"],
                    type=conn_dict["type"],
                )
                for conn_dict in all_entries
                if conn_dict["scope"] == "connection"
            },
            datasets={
                ds_dict["key"].split("/")[-1]: dict(
                    id=ds_dict["entryId"],
                )
                for ds_dict in all_entries
                if ds_dict["scope"] == "dataset"
            },
            charts={
                chart_dict["key"].split("/")[-1]: dict(
                    id=chart_dict["entryId"],
                )
                for chart_dict in all_entries
                if chart_dict["scope"] == "widget"
            },
            dashboards={
                dash["key"].split("/")[-1]: dict(
                    id=dash["entryId"],
                )
                for dash in all_entries
                if dash["scope"] == "dash"
            },
        )

lib/dl_api_lib/dl_api_lib/app/control_api/resources/info.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import functools
import logging.config
from typing import (
    Generic,
    TypeVar,
)

from aiohttp import web
import attr

from dl_api_commons.aio.middlewares.commit_rci import commit_rci_middleware
from dl_api_commons.aio.middlewares.request_bootstrap import RequestBootstrap
from dl_api_commons.aio.middlewares.request_id import RequestId
from dl_api_commons.aio.middlewares.tracing import TracingService
from dl_api_commons.aio.server_header import ServerHeader
from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_api_commons.sentry_config import (
    SentryConfig,
    configure_sentry_for_aiohttp,
)
from dl_api_lib.aio.aiohttp_wrappers import (
    AppWrapper,
    DSAPIRequest,
)
from dl_api_lib.aio.middlewares.error_handling_outer import DatasetAPIErrorHandler
from dl_api_lib.aio.middlewares.json_body_middleware import json_body_middleware
from dl_api_lib.app.data_api.resources.dashsql import DashSQLView
from dl_api_lib.app.data_api.resources.dataset.distinct import (
    DatasetDistinctViewV1,
    DatasetDistinctViewV1_5,
    DatasetDistinctViewV2,
)
from dl_api_lib.app.data_api.resources.dataset.fields import DatasetFieldsView
from dl_api_lib.app.data_api.resources.dataset.pivot import DatasetPivotView
from dl_api_lib.app.data_api.resources.dataset.preview import (
    DatasetPreviewViewV1,
    DatasetPreviewViewV1_5,
    DatasetPreviewViewV2,
)
from dl_api_lib.app.data_api.resources.dataset.range import (
    DatasetRangeViewV1,
    DatasetRangeViewV1_5,
    DatasetRangeViewV2,
)
from dl_api_lib.app.data_api.resources.dataset.result import (
    DatasetResultViewV1,
    DatasetResultViewV1_5,
    DatasetResultViewV2,
)
from dl_api_lib.app.data_api.resources.metrics import DSDataApiMetricsView
from dl_api_lib.app.data_api.resources.ping import (
    PingReadyView,
    PingView,
)
from dl_api_lib.app.data_api.resources.unistat import UnistatView
from dl_api_lib.app_common import SRFactoryBuilder
from dl_api_lib.app_settings import DataApiAppSettings
from dl_compeng_pg.compeng_pg_base.data_processor_service_pg import CompEngPgConfig
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.enums import RedisMode
from dl_constants.enums import (
    ConnectionType,
    ProcessorType,
    RedisInstanceKind,
)
from dl_core.aio.web_app_services.data_processing.factory import make_compeng_service
from dl_core.aio.web_app_services.redis import (
    RedisSentinelService,
    SingleHostSimpleRedisService,
)
from dl_core.utils import make_url


LOGGER = logging.getLogger(__name__)


def _log_exc(coro):  # type: ignore
    @functools.wraps(coro)
    async def wrapper(*args, **kwargs):  # type: ignore
        try:
            return await coro(*args, **kwargs)
        except Exception:
            LOGGER.exception("Unhandled exception in application signal listener", exc_info=True)
            raise

    return wrapper


async def add_connection_close(request: web.Request, response: web.StreamResponse) -> None:
    """
    Fix a keepalive race condition (HTTP server closing the connection when
    HTTP client is sending a next request) by forcing
    single-request-per-connection (telling the client to close the
    connection after the first request).
    """
    response.headers.add("Connection", "close")


@attr.s(frozen=True)
class EnvSetupResult:
    auth_mw_list: list[AIOHTTPMiddleware] = attr.ib(kw_only=True)
    sr_middleware_list: list[AIOHTTPMiddleware] = attr.ib(kw_only=True)
    usm_middleware_list: list[AIOHTTPMiddleware] = attr.ib(kw_only=True)


TDataApiSettings = TypeVar("TDataApiSettings", bound=DataApiAppSettings)


@attr.s(kw_only=True)
class DataApiAppFactory(SRFactoryBuilder, Generic[TDataApiSettings], abc.ABC):
    _settings: TDataApiSettings = attr.ib()

    @abc.abstractmethod
    def get_app_version(self) -> str:
        raise NotImplementedError()

    @abc.abstractmethod
    def set_up_environment(
        self,
        connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
    ) -> EnvSetupResult:
        raise NotImplementedError()

    @property
    @abc.abstractmethod
    def _is_public(self) -> bool:
        raise NotImplementedError()

    @property
    def _is_async_env(self) -> bool:
        return True

    def set_up_routes(self, app: web.Application) -> None:
        app.router.add_route("get", "/ping", PingView)
        app.router.add_route("get", "/ping_ready", PingReadyView)
        app.router.add_route("get", "/unistat/", UnistatView)
        app.router.add_route("get", "/metrics", DSDataApiMetricsView)

        # TODO: at some point the non-'/data/' routes should probably be removed.
        app.router.add_route("post", "/api/v1/connections/{conn_id}/dashsql", DashSQLView)  # FIXME: Remove
        app.router.add_route("post", "/api/data/v1/connections/{conn_id}/dashsql", DashSQLView)
        app.router.add_route(
            "post", "/api/v1/datasets/{ds_id}/versions/draft/result", DatasetResultViewV1
        )  # FIXME: Remove
        app.router.add_route("post", "/api/data/v1/datasets/{ds_id}/versions/draft/result", DatasetResultViewV1)
        app.router.add_route("post", "/api/data/v1.5/datasets/{ds_id}/result", DatasetResultViewV1_5)
        app.router.add_route("post", "/api/data/v2/datasets/{ds_id}/result", DatasetResultViewV2)
        app.router.add_route(
            "post", "/api/v1/datasets/{ds_id}/versions/draft/values/distinct", DatasetDistinctViewV1
        )  # FIXME: Remove
        app.router.add_route(
            "post", "/api/data/v1/datasets/{ds_id}/versions/draft/values/distinct", DatasetDistinctViewV1
        )
        app.router.add_route("post", "/api/data/v1.5/datasets/{ds_id}/values/distinct", DatasetDistinctViewV1_5)
        app.router.add_route("post", "/api/data/v2/datasets/{ds_id}/values/distinct", DatasetDistinctViewV2)
        app.router.add_route(
            "post", "/api/v1/datasets/{ds_id}/versions/draft/values/range", DatasetRangeViewV1
        )  # FIXME: Remove
        app.router.add_route("post", "/api/data/v1/datasets/{ds_id}/versions/draft/values/range", DatasetRangeViewV1)
        app.router.add_route("post", "/api/data/v1.5/datasets/{ds_id}/values/range", DatasetRangeViewV1_5)
        app.router.add_route("post", "/api/data/v2/datasets/{ds_id}/values/range", DatasetRangeViewV2)
        app.router.add_route(
            "post", "/api/data/v1/datasets/{ds_id}/versions/draft/pivot", DatasetPivotView
        )  # FIXME: Remove
        app.router.add_route("post", "/api/data/v2/datasets/{ds_id}/pivot", DatasetPivotView)
        app.router.add_route("get", "/api/v1/datasets/{ds_id}/fields", DatasetFieldsView)  # FIXME: Remove
        app.router.add_route("get", "/api/data/v1/datasets/{ds_id}/fields", DatasetFieldsView)
        app.router.add_route("get", "/api/data/v2/datasets/{ds_id}/fields", DatasetFieldsView)

        if not self._is_public:
            app.router.add_route("post", "/api/v1/datasets/data/preview", DatasetPreviewViewV1)  # FIXME: Remove
            app.router.add_route("post", "/api/data/v1/datasets/data/preview", DatasetPreviewViewV1)
            app.router.add_route(
                "post", "/api/v1/datasets/{ds_id}/versions/draft/preview", DatasetPreviewViewV1
            )  # FIXME: Remove
            app.router.add_route("post", "/api/data/v1/datasets/{ds_id}/versions/draft/preview", DatasetPreviewViewV1)
            app.router.add_route("post", "/api/data/v1.5/datasets/data/preview", DatasetPreviewViewV1_5)
            app.router.add_route("post", "/api/data/v1.5/datasets/{ds_id}/preview", DatasetPreviewViewV1_5)
            app.router.add_route("post", "/api/data/v2/datasets/data/preview", DatasetPreviewViewV2)
            app.router.add_route("post", "/api/data/v2/datasets/{ds_id}/preview", DatasetPreviewViewV2)

    def set_up_sentry(self) -> None:
        configure_sentry_for_aiohttp(
            SentryConfig(
                dsn=self._settings.SENTRY_DSN,
                release=self.get_app_version(),
            )
        )

    def create_app(
        self,
        connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
    ) -> web.Application:
        if self._settings.SENTRY_ENABLED:
            self.set_up_sentry()

        env_setup_result = self.set_up_environment(
            connectors_settings=connectors_settings,
        )

        req_id_service = RequestId(
            dl_request_cls=DSAPIRequest,
            append_own_req_id=True,
            app_prefix=self._settings.app_prefix,
            is_public_env=self._is_public,
        )

        error_handler = DatasetAPIErrorHandler(
            public_mode=self._is_public,
            use_sentry=self._settings.SENTRY_ENABLED,
            sentry_app_name_tag=self._settings.app_name,
        )

        middleware_list = [
            TracingService().middleware,
            RequestBootstrap(
                req_id_service=req_id_service,
                error_handler=error_handler,
                timeout_sec=self._settings.COMMON_TIMEOUT_SEC,
            ).middleware,
            *env_setup_result.auth_mw_list,
            commit_rci_middleware(),
            *env_setup_result.sr_middleware_list,
            *env_setup_result.usm_middleware_list,
            json_body_middleware(),
        ]

        app = web.Application(
            middlewares=middleware_list,
        )

        wrapper = AppWrapper(
            allow_query_cache_usage=self._settings.CACHES_ON,
            allow_notifications=not self._is_public,
        )
        wrapper.bind(app)

        ServerHeader("DataLensAPI").add_signal_handlers(app)
        app.on_response_prepare.append(req_id_service.on_response_prepare)

        if self._settings.BI_ASYNC_APP_DISABLE_KEEPALIVE:
            app.on_response_prepare.append(add_connection_close)

        if self._settings.CACHES_ON and self._settings.CACHES_REDIS:
            if self._settings.CACHES_REDIS.MODE == RedisMode.single_host:
                redis_server_single_host = SingleHostSimpleRedisService(
                    instance_kind=RedisInstanceKind.caches,
                    url=make_url(
                        protocol="rediss" if self._settings.CACHES_REDIS.SSL else "redis",
                        host=self._settings.CACHES_REDIS.HOSTS[0],
                        port=self._settings.CACHES_REDIS.PORT,
                        path=str(self._settings.CACHES_REDIS.DB),
                    ),
                    password=self._settings.CACHES_REDIS.PASSWORD,
                    ssl=self._settings.CACHES_REDIS.SSL,
                )
                app.on_startup.append(_log_exc(redis_server_single_host.init_hook))
                app.on_cleanup.append(_log_exc(redis_server_single_host.tear_down_hook))
            elif self._settings.CACHES_REDIS.MODE == RedisMode.sentinel:
                redis_server_sentinel = RedisSentinelService(
                    instance_kind=RedisInstanceKind.caches,
                    namespace=self._settings.CACHES_REDIS.CLUSTER_NAME,
                    sentinel_hosts=self._settings.CACHES_REDIS.HOSTS,
                    sentinel_port=self._settings.CACHES_REDIS.PORT,
                    db=self._settings.CACHES_REDIS.DB,
                    password=self._settings.CACHES_REDIS.PASSWORD,
                    ssl=self._settings.CACHES_REDIS.SSL,
                )
                app.on_startup.append(_log_exc(redis_server_sentinel.init_hook))
                app.on_cleanup.append(_log_exc(redis_server_sentinel.tear_down_hook))
            else:
                raise ValueError(f"Unknown redis mode {self._settings.CACHES_REDIS.MODE}")

        if self._settings.MUTATIONS_CACHES_ON and self._settings.MUTATIONS_REDIS:
            if self._settings.MUTATIONS_REDIS.MODE == RedisMode.single_host:
                mutations_redis_server_single_host = SingleHostSimpleRedisService(
                    instance_kind=RedisInstanceKind.mutations,
                    url=make_url(
                        protocol="rediss" if self._settings.MUTATIONS_REDIS.SSL else "redis",
                        host=self._settings.MUTATIONS_REDIS.HOSTS[0],
                        port=self._settings.MUTATIONS_REDIS.PORT,
                        path=str(self._settings.MUTATIONS_REDIS.DB),
                    ),
                    password=self._settings.MUTATIONS_REDIS.PASSWORD,
                    ssl=self._settings.MUTATIONS_REDIS.SSL,
                )
                app.on_startup.append(_log_exc(mutations_redis_server_single_host.init_hook))
                app.on_cleanup.append(_log_exc(mutations_redis_server_single_host.tear_down_hook))
            else:
                mutations_redis_server_sentinel: SingleHostSimpleRedisService | RedisSentinelService = (
                    RedisSentinelService(
                        instance_kind=RedisInstanceKind.mutations,
                        namespace=self._settings.MUTATIONS_REDIS.CLUSTER_NAME,
                        sentinel_hosts=self._settings.MUTATIONS_REDIS.HOSTS,
                        sentinel_port=self._settings.MUTATIONS_REDIS.PORT,
                        db=self._settings.MUTATIONS_REDIS.DB,
                        password=self._settings.MUTATIONS_REDIS.PASSWORD,
                        ssl=self._settings.MUTATIONS_REDIS.SSL,
                    )
                )
                app.on_startup.append(_log_exc(mutations_redis_server_sentinel.init_hook))
                app.on_cleanup.append(_log_exc(mutations_redis_server_sentinel.tear_down_hook))

        if self._settings.BI_COMPENG_PG_ON and self._settings.BI_COMPENG_PG_URL is not None:
            compeng_service = make_compeng_service(
                processor_type=ProcessorType.ASYNCPG,
                config=CompEngPgConfig(
                    url=self._settings.BI_COMPENG_PG_URL,
                ),
            )
            app.on_startup.append(_log_exc(compeng_service.init_hook))
            app.on_cleanup.append(_log_exc(compeng_service.tear_down_hook))

        # TODO: don't use it again!
        # special hack for gettings bleeding edge users in dashsql
        app["BLEEDING_EDGE_USERS"] = self._settings.BLEEDING_EDGE_USERS

        self.set_up_routes(app=app)

        return app

lib/dl_api_lib/dl_api_lib/app/data_api/app.py
datalens-tech/datalens-backend
# TODO FIX: Move to common

from __future__ import annotations

import enum
import functools
import inspect
import itertools
import typing

from dl_api_commons.aiohttp.aiohttp_wrappers import (
    DLRequestView,
    RequiredResource,
)
from dl_api_lib.aio import aiohttp_wrappers


_METHOD_REQ_RESOURCES_ATTR_NAME = "__dl_required_resources__"


class RequiredResourceDSAPI(RequiredResource):
    JSON_REQUEST = enum.auto()


class BaseView(DLRequestView[aiohttp_wrappers.DSAPIRequest]):
    dl_request_cls = aiohttp_wrappers.DSAPIRequest

    _COMMON_REQUIRED_RESOURCES: typing.ClassVar[typing.FrozenSet[RequiredResource]] = frozenset()

    @classmethod
    def get_required_resources(cls, method_name: str) -> typing.FrozenSet[RequiredResource]:
        method_name = method_name.lower()
        if hasattr(cls, method_name):
            method = getattr(cls, method_name)
            if hasattr(method, _METHOD_REQ_RESOURCES_ATTR_NAME):
                return frozenset(
                    itertools.chain(cls._COMMON_REQUIRED_RESOURCES, getattr(method, _METHOD_REQ_RESOURCES_ATTR_NAME))
                )
            else:
                return cls._COMMON_REQUIRED_RESOURCES

        # If method is not implemented - we will not use any resource
        return frozenset()

    @staticmethod
    def with_resolved_entities(coro):  # type: ignore  # TODO: fix
        @functools.wraps(coro)
        async def wrapper(self, *args, **kwargs):  # type: ignore  # TODO: fix
            await self.resolve_entities()
            return await coro(self, *args, **kwargs)

        return wrapper

    async def resolve_entities(self):  # type: ignore  # TODO: fix
        pass


# TODO FIX: add ability to exclude resources for particular handles
def requires(*resources: RequiredResource, skip_parent_resources: bool = False):  # type: ignore  # TODO: fix
    def real_deco(func_or_class):  # type: ignore  # TODO: fix
        if isinstance(func_or_class, type) and issubclass(func_or_class, BaseView):
            clz = func_or_class
            if skip_parent_resources:
                all_resources = resources
            else:
                all_resources = itertools.chain(  # type: ignore  # TODO: fix
                    clz._COMMON_REQUIRED_RESOURCES,
                    resources,
                )
            clz._COMMON_REQUIRED_RESOURCES = frozenset(all_resources)

        elif inspect.iscoroutinefunction(func_or_class):
            func = func_or_class
            setattr(func, _METHOD_REQ_RESOURCES_ATTR_NAME, frozenset(resources))

        else:
            raise ValueError(
                "Only BaseView or coroutine functions inside of them can be used with requires() decorator."
                f" Not '{func_or_class}'"
            )

        return func_or_class

    return real_deco

lib/dl_api_lib/dl_api_lib/app/data_api/resources/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import decimal
import json
import math
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Optional,
    Type,
)
import uuid

from aiohttp import web

from dl_api_commons.aiohttp.aiohttp_wrappers import RequiredResourceCommon
from dl_api_connector.dashsql import TValueBase
from dl_api_lib.app.data_api.resources.base import (
    BaseView,
    RequiredResourceDSAPI,
    requires,
)
from dl_api_lib.dashsql.registry import get_dash_sql_param_literalizer
from dl_api_lib.enums import USPermissionKind
import dl_api_lib.schemas.main
from dl_api_lib.utils.base import need_permission_on_entry
from dl_app_tools.profiling_base import generic_profiler_async
from dl_constants.enums import (
    ConnectionType,
    UserDataType,
)
from dl_core.backend_types import get_backend_type
from dl_core.data_processing.dashsql import (
    BindParamInfo,
    DashSQLCachedSelector,
    DashSQLEvent,
    TRow,
)
from dl_core.exc import (
    DashSQLError,
    UnexpectedUSEntryType,
)
from dl_core.us_connection_base import (
    ConnectionBase,
    ExecutorBasedMixin,
    SubselectMixin,
)
from dl_query_processing.utils.datetime import parse_datetime


if TYPE_CHECKING:
    from dl_constants.types import TJSONLike
    from dl_core.data_processing.dashsql import (
        DashSQLSelector,
        TResultEvents,
    )


TRowProcessor = Callable[[TRow], TRow]


def parse_value(value: Optional[str], bi_type: UserDataType) -> Any:
    if value is None:
        return None
    if bi_type == UserDataType.string:
        return value
    if bi_type == UserDataType.integer:
        return int(value)
    if bi_type == UserDataType.float:
        return float(value)
    if bi_type == UserDataType.date:
        return datetime.datetime.strptime(value, "%Y-%m-%d").date()
    if bi_type == UserDataType.datetime:
        return parse_datetime(value)
    if bi_type == UserDataType.boolean:
        if value == "true":
            return True
        if value == "false":
            return False
        raise ValueError(f"Not a valid boolean: {value!r}")
    raise DashSQLError(f"Unsupported type: {bi_type.name!r}")


def make_param_obj(name: str, param: dict, conn_type: ConnectionType) -> BindParamInfo:
    type_name: str = param["type_name"]
    value_base: TValueBase = param["value"]

    try:
        bi_type = UserDataType[type_name]
    except KeyError:
        raise DashSQLError(f"Unknown type name {type_name!r}")

    backend_type = get_backend_type(conn_type=conn_type)
    literalizer_cls = get_dash_sql_param_literalizer(backend_type=backend_type)
    sa_type = literalizer_cls.get_sa_type(bi_type=bi_type, value_base=value_base)

    try:
        if isinstance(value_base, (list, tuple)):
            value = [parse_value(sub_value, bi_type) for sub_value in value_base]
        else:
            value = parse_value(value_base, bi_type)
    except ValueError:
        raise DashSQLError(f"Unsupported value for type {type_name!r}: {value_base!r}")

    return BindParamInfo(
        name=name,
        type_name=type_name,
        sa_type=sa_type,
        value=value,
        expanding=isinstance(value, (list, tuple)),
    )


@requires(RequiredResourceCommon.US_MANAGER)
class DashSQLView(BaseView):
    """
    Connection + SQL query -> special-format data result.

    Partially related to `DatasetPreviewView` and its bases, but much more stripped down.
    """

    # TODO?: cache support

    endpoint_code = "DashSQL"
    profiler_prefix = "dashsql_result"
    dashsql_selector_cls: Type[DashSQLSelector] = DashSQLCachedSelector

    @property
    def conn_id(self) -> Optional[str]:
        return self.request.match_info.get("conn_id")

    def enrich_logging_context(self, conn: ConnectionBase) -> None:
        """
        See also:
        `dl_api_lib.app.data_api.resources.dataset.base.DatasetDataBaseView.default_query_execution_cm_stack`
        """
        if not self.dl_request.log_ctx_controller:
            # TODO?: warn
            return
        self.dl_request.log_ctx_controller.put_to_context("conn_type", conn.conn_type.name)
        sr = self.dl_request.services_registry
        if isinstance(conn, ExecutorBasedMixin):
            ce_cls_str = sr.get_conn_executor_factory().get_async_conn_executor_cls(conn).__qualname__
            self.dl_request.log_ctx_controller.put_to_context("conn_exec_cls", ce_cls_str)

    @staticmethod
    def _json_default(value: Any) -> TJSONLike:
        if isinstance(value, bytes):
            return value.decode("utf-8", errors="replace")
        if isinstance(value, (datetime.datetime, datetime.date, datetime.time)):
            return value.isoformat()
        if isinstance(value, datetime.timedelta):
            return value.total_seconds()
        if isinstance(value, decimal.Decimal):
            return str(value)
        if isinstance(value, uuid.UUID):
            return str(value)
        raise DashSQLError("Unexpected value type in dashsql serialization: {!r}".format(type(value)))

    @classmethod
    def _postprocess_key(cls, value: Any) -> str:
        return str(value)

    @classmethod
    def _postprocess_any(cls, value: Any) -> Any:
        if isinstance(value, (list, tuple)):
            return tuple(cls._postprocess_any(sub_value) for sub_value in value)
        if isinstance(value, dict):
            return {
                cls._postprocess_key(sub_key): cls._postprocess_any(sub_value) for sub_key, sub_value in value.items()
            }
        if isinstance(value, float) and (math.isnan(value) or math.isinf(value)):
            return str(value)
        return value

    @classmethod
    def _make_postprocess_row(cls, bi_type_names: tuple[str, ...]) -> TRowProcessor:
        # Nothing type-specific at the moment
        funcs = tuple(cls._postprocess_any for _ in bi_type_names)

        def _postprocess_row(row: TRow, funcs: tuple[Callable, ...] = funcs) -> TRow:
            return tuple(func(value) for func, value in zip(funcs, row))

        return _postprocess_row

    async def _postprocess_events(self, result_events: TResultEvents) -> TResultEvents:
        postprocess_row: Optional[TRowProcessor] = None
        async for event_name, event_data in result_events:
            if postprocess_row is None and event_name == DashSQLEvent.metadata.value:
                assert isinstance(event_data, dict)
                postprocess_row = self._make_postprocess_row(event_data["bi_types"])  # type: ignore  # TODO: fix
            elif event_name == DashSQLEvent.rowchunk.value and postprocess_row is not None:
                event_data = tuple(postprocess_row(row) for row in event_data)  # type: ignore  # TODO: fix
            elif event_name == DashSQLEvent.row.value and postprocess_row is not None:  # should be obsolete
                event_data = postprocess_row(event_data)  # type: ignore  # TODO: fix
            yield event_name, event_data

    @staticmethod
    async def _flatten_chunk_events(result_events: TResultEvents) -> TResultEvents:
        """Unwrap `rowchunk` events into `row` events, to avoid changing the protocol (for now)"""
        async for event_name, event_data in result_events:
            if event_name == DashSQLEvent.rowchunk.value:
                for row in event_data:
                    assert isinstance(row, (tuple, list))
                    yield DashSQLEvent.row.value, row  # type: ignore  # TODO: fix
            else:
                yield event_name, event_data

    async def collect_result_events_into_response(self, result_events: TResultEvents) -> web.Response:
        events: list = []
        async for event_name, event_data in result_events:
            events.append(dict(event=event_name, data=event_data))

        data = json.dumps(events, default=self._json_default)
        response = web.Response(
            body=data.encode(),
            content_type="application/json",
        )
        return response

    @generic_profiler_async("dashsql-result")
    @requires(RequiredResourceDSAPI.JSON_REQUEST)
    async def post(self) -> web.Response:
        conn_id = self.conn_id
        assert conn_id

        if self.dl_request.log_ctx_controller:
            self.dl_request.log_ctx_controller.put_to_context("conn_id", conn_id)

        schema = dl_api_lib.schemas.main.DashSQLRequestSchema()
        body = schema.load(self.dl_request.json)
        sql_query = body["sql_query"]
        params = body.get("params")
        db_params = body.get("db_params")
        connector_specific_params = body.get("connector_specific_params")

        conn = await self.dl_request.us_manager.get_by_id(conn_id, ConnectionBase)
        assert isinstance(conn, ConnectionBase)

        self.enrich_logging_context(conn)

        if not isinstance(conn, SubselectMixin) or not isinstance(conn, ExecutorBasedMixin):
            raise UnexpectedUSEntryType("Expecting a subselect-compatible connection")

        # A slightly more explicit check (which should have been done in `get_by_id` anyway):
        need_permission_on_entry(conn, USPermissionKind.execute)

        # TODO: instead of this, use something like:
        #     formula_dialect = dl_formula.core.dialect.from_name_and_version(conn.get_dialect().name)
        #     bindparam = dl_formula.definitions.literals.literal(parsed_value, formula_dialect)
        # (but account for `expanding`)
        conn_type = conn.conn_type
        param_objs = None
        if params is not None:
            param_objs = [make_param_obj(name, param, conn_type=conn_type) for name, param in params.items()]

        # TODO: move dashsql selector's construction to factory
        bleeding_edge_users = self.request.app.get("BLEEDING_EDGE_USERS", ())
        is_bleeding_edge_user = self.dl_request.rci.user_name in bleeding_edge_users

        dashsql_selector = self.dashsql_selector_cls(
            conn=conn,
            sql_query=sql_query,
            # Note: `None` means `no substitution` (legacy),
            # `[]` means `attempt to substitute`.
            params=param_objs,
            db_params=db_params or {},
            connector_specific_params=connector_specific_params,
            service_registry=self.dl_request.services_registry,
            is_bleeding_edge_user=is_bleeding_edge_user,
        )
        result_events = await dashsql_selector.execute()
        result_events = self._postprocess_events(result_events)

        # TODO: modify API consumers to allow `rowchunk` events;
        # but for now, unroll them into `row` events.
        result_events = self._flatten_chunk_events(result_events)

        return await self.collect_result_events_into_response(result_events)

lib/dl_api_lib/dl_api_lib/app/data_api/resources/dashsql.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import (
    AsyncExitStack,
    asynccontextmanager,
)
import logging
import os
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncIterable,
    ClassVar,
    Collection,
    Dict,
    List,
    Optional,
    Type,
)

from aiohttp import web

from dl_api_commons.aiohttp.aiohttp_wrappers import RequiredResourceCommon
from dl_api_lib import utils
from dl_api_lib.api_common.data_serialization import (
    DataRequestResponseSerializer,
    get_fields_data_serializable,
)
from dl_api_lib.api_common.dataset_loader import (
    DatasetApiLoader,
    DatasetUpdateInfo,
)
from dl_api_lib.api_common.update_dataset_mutation_key import UpdateDatasetMutationKey
from dl_api_lib.app.data_api.resources.base import (
    BaseView,
    requires,
)
from dl_api_lib.dataset.view import DatasetView
from dl_api_lib.query.formalization.block_formalizer import BlockFormalizer
from dl_api_lib.query.formalization.legend_formalizer import (
    DistinctLegendFormalizer,
    LegendFormalizer,
    PivotLegendFormalizer,
    PreviewLegendFormalizer,
    RangeLegendFormalizer,
    ResultLegendFormalizer,
)
from dl_api_lib.request_model.data import (
    Action,
    DataRequestModel,
    FieldAction,
)
from dl_api_lib.service_registry.service_registry import ApiServiceRegistry
from dl_app_tools.profiling_base import GenericProfiler
from dl_constants.enums import DataSourceRole
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.data_source.base import DataSource
from dl_core.data_source.collection import DataSourceCollectionFactory
from dl_core.dataset_capabilities import DatasetCapabilities
from dl_core.exc import USObjectNotFoundException
from dl_core.us_connection_base import ExecutorBasedMixin
from dl_core.us_dataset import Dataset
from dl_core.us_manager.mutation_cache.engine_factory import CacheInitializationError
from dl_core.us_manager.mutation_cache.mutation_key_base import MutationKey
from dl_core.us_manager.mutation_cache.usentry_mutation_cache import (
    RedisCacheEngine,
    USEntryMutationCache,
)
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_query_processing.enums import QueryType
from dl_query_processing.execution.exec_info import QueryExecutionInfo
from dl_query_processing.legend.block_legend import BlockSpec
from dl_query_processing.merging.merger import DataStreamMerger
from dl_query_processing.merging.primitives import MergedQueryDataStream
from dl_query_processing.pagination.paginator import QueryPaginator
from dl_query_processing.postprocessing.postprocessor import DataPostprocessor
from dl_query_processing.postprocessing.primitives import (
    PostprocessedQuery,
    PostprocessedQueryBlock,
    PostprocessedQueryUnion,
    PostprocessedRow,
)
from dl_utils.task_runner import ConcurrentTaskRunner


if TYPE_CHECKING:
    from dl_api_lib.query.formalization.raw_specs import RawQuerySpecUnion


LOGGER = logging.getLogger(__name__)


@requires(RequiredResourceCommon.US_MANAGER)
class DatasetDataBaseView(BaseView):
    STORED_DATASET_REQUIRED: ClassVar[bool] = True

    profiler_prefix: ClassVar[str]

    dataset: Dataset
    ds_accessor: DatasetComponentAccessor

    @property
    def dataset_id(self) -> Optional[str]:
        return self.request.match_info.get("ds_id")

    @property
    def allow_query_cache_usage(self) -> bool:
        """
        Indicates if query cache may be used for request.
         Now only config flag but in future - may be special header.
        """
        return self.dl_request.app_wrapper.allow_query_cache_usage

    @property
    def allow_notifications(self) -> bool:
        """
        Indicates if may be returned in response.
        Now only config flag but in future - may be more complex logic
        """
        return self.dl_request.app_wrapper.allow_notifications

    @asynccontextmanager  # type: ignore  # TODO: fix
    async def default_query_execution_cm_stack(
        self,
        exec_info: QueryExecutionInfo,
        body: dict,
        profiling_code: str = "query-execute",
    ) -> AsyncIterable[AsyncExitStack]:
        """
        See also:
        `dl_api_lib.app.data_api.resources.dashsql.DashSQLView.enrich_logging_context`
        """
        async with AsyncExitStack() as stack:
            if self.dl_request.log_ctx_controller:
                try:
                    target_connections = exec_info.target_connections
                    assert len(target_connections) == 1
                    target_conn = target_connections[0]
                    self.dl_request.log_ctx_controller.put_to_context("conn_id", target_conn.uuid)
                    self.dl_request.log_ctx_controller.put_to_context("conn_type", target_conn.conn_type.name)

                    sr = self.dl_request.services_registry
                    try:
                        if isinstance(target_conn, ExecutorBasedMixin):
                            ce_cls_str = (
                                sr.get_conn_executor_factory().get_async_conn_executor_cls(target_conn).__qualname__
                            )
                        else:
                            ce_cls_str = "N/D"
                        self.dl_request.log_ctx_controller.put_to_context("conn_exec_cls", ce_cls_str)
                    except Exception:  # noqa
                        LOGGER.exception("Can not get CE class for connection %s", target_conn.uuid)

                except Exception:  # noqa
                    LOGGER.exception("Can not save connection info to logging context")

            stack.enter_context(GenericProfiler(f"{self.profiler_prefix}-{profiling_code}"))  # type: ignore  # TODO: fix
            stack.enter_context(utils.query_execution_context(dataset_id=self.dataset.uuid, version="draft", body=body))
            yield stack

    def resolve_dataset_source_role(self, dataset: Dataset, log_reasons: bool = False) -> DataSourceRole:
        dsrc_coll_factory = DataSourceCollectionFactory(us_entry_buffer=self.dl_request.us_manager.get_entry_buffer())
        capabilities = DatasetCapabilities(dataset=dataset, dsrc_coll_factory=dsrc_coll_factory)
        return capabilities.resolve_source_role(log_reasons=log_reasons)

    async def resolve_entities(self) -> None:
        us_manager = self.dl_request.us_manager
        if self.dl_request.log_ctx_controller:
            self.dl_request.log_ctx_controller.put_to_context("dataset_id", self.dataset_id)

        if self.dataset_id is None:
            if self.STORED_DATASET_REQUIRED:
                raise ValueError(f"View {self} requires stored dataset, but no ID found in match info")

            dataset = Dataset.create_from_dict(
                Dataset.DataModel(name=""),  # TODO: Remove name - it's not used, but is required
                us_manager=us_manager,  # type: ignore  # TODO: fix # WTF??? sync or async??
            )
        else:
            try:
                dataset = await us_manager.get_by_id(self.dataset_id, Dataset)
            except USObjectNotFoundException:
                raise web.HTTPNotFound(reason="Entity not found")

            await us_manager.load_dependencies(dataset)

        self.dataset = dataset
        self.ds_accessor = DatasetComponentAccessor(dataset=dataset)

    @staticmethod
    def _updates_only_fields(updates: List[Action]) -> bool:
        # Checks if updates has only field updates
        return all([isinstance(upd, FieldAction) for upd in updates])

    def try_get_mutation_key(self, updates: List[Action]) -> Optional[MutationKey]:
        if self.dataset_id is not None and self.dataset.revision_id is not None:
            if self._updates_only_fields(updates):
                return UpdateDatasetMutationKey.create(self.dataset.revision_id, updates)  # type: ignore
        return None

    def try_get_cache(self) -> Optional[USEntryMutationCache]:
        try:
            mc_factory = self.dl_request.services_registry.get_mutation_cache_factory()
            if mc_factory is None:
                LOGGER.debug("Mutation cache is disabled")
                return None
            mce_factory = self.dl_request.services_registry.get_mutation_cache_engine_factory(RedisCacheEngine)
            cache_engine = mce_factory.get_cache_engine()
            mutation_cache = mc_factory.get_mutation_cache(
                usm=self.dl_request.us_manager,
                engine=cache_engine,
            )
            return mutation_cache
        except CacheInitializationError:  # Error creating factory with redis cache engine or something
            LOGGER.error("Mutation cache error", exc_info=True)
            return None

    async def try_get_dataset_from_cache(
        self,
        mutation_cache: Optional[USEntryMutationCache],
        mutation_key: Optional[MutationKey],
    ) -> Optional[Dataset]:
        if mutation_key is None or mutation_cache is None:
            return None
        if self.dataset_id is None or self.dataset.revision_id is None:
            return None

        cached_dataset = await mutation_cache.get_mutated_entry_from_cache(
            Dataset,
            self.dataset_id,
            self.dataset.revision_id,
            mutation_key,
        )
        if cached_dataset is None:
            return None
        LOGGER.info("Found dataset in mutation cache")
        cached_dataset.permissions_mode = self.dataset.permissions_mode
        cached_dataset.permissions = self.dataset.permissions
        assert isinstance(cached_dataset, Dataset)
        return cached_dataset

    async def try_save_dataset_to_cache(
        self,
        mutation_cache: Optional[USEntryMutationCache],
        mutation_key: Optional[MutationKey],
        dataset: Dataset,
    ) -> None:
        if mutation_key is None:
            return None
        if mutation_cache is None:
            return None
        await mutation_cache.save_mutation_cache(dataset, mutation_key)

    async def prepare_dataset_for_request(
        self,
        req_model: DataRequestModel,
        allow_rls_change: bool = False,
        enable_mutation_caching: bool = False,
    ) -> DatasetUpdateInfo:
        us_manager = self.dl_request.us_manager
        services_registry = self.dl_request.services_registry
        assert isinstance(services_registry, ApiServiceRegistry)
        loader = DatasetApiLoader(service_registry=services_registry)

        with GenericProfiler("dataset-prepare"):
            if enable_mutation_caching:
                mutation_cache = self.try_get_cache()
                mutation_key = self.try_get_mutation_key(req_model.updates)
                cached_dataset = await self.try_get_dataset_from_cache(mutation_cache, mutation_key)
                if cached_dataset:
                    self.dataset = cached_dataset
                    return loader.update_dataset_from_body(
                        dataset=self.dataset,
                        us_manager=us_manager,
                        dataset_data=req_model.dataset,
                        allow_rls_change=allow_rls_change,
                    )

            update_info = loader.update_dataset_from_body(
                dataset=self.dataset,
                us_manager=us_manager,
                dataset_data=req_model.dataset,
                allow_rls_change=allow_rls_change,
            )
            await us_manager.load_dependencies(self.dataset)

            services_registry = self.dl_request.services_registry
            assert isinstance(services_registry, ApiServiceRegistry)

            await self.check_for_notifications(services_registry, us_manager)

            dataset_validator_factory = services_registry.get_dataset_validator_factory()
            ds_validator = dataset_validator_factory.get_dataset_validator(
                ds=self.dataset, us_manager=us_manager, is_data_api=True
            )
            executor = services_registry.get_compute_executor()
            await executor.execute(lambda: ds_validator.apply_batch(action_batch=req_model.updates))
            if enable_mutation_caching:
                await self.try_save_dataset_to_cache(mutation_cache, mutation_key, self.dataset)  # noqa

        return update_info

    async def check_for_notifications(self, services_registry: ApiServiceRegistry, us_manager: AsyncUSManager) -> None:
        ds_lc_manager = us_manager.get_lifecycle_manager(self.dataset)
        for conn_id in ds_lc_manager.collect_links().values():
            try:
                conn = us_manager.get_loaded_us_connection(conn_id)
            except Exception:
                LOGGER.info("Failed to get loaded us connection %s", conn_id, exc_info=True)
            else:
                conn_notifications = conn.check_for_notifications()
                if not conn_notifications:
                    continue

                reporting_registry = services_registry.get_reporting_registry()
                for notification_record in conn_notifications:
                    reporting_registry.save_reporting_record(notification_record)

    def make_legend_formalizer(self, query_type: QueryType, autofill_legend: bool = False) -> LegendFormalizer:
        legend_formalizer_cls: Type[LegendFormalizer]
        if query_type == QueryType.pivot:
            legend_formalizer_cls = PivotLegendFormalizer
        elif query_type == QueryType.result:
            legend_formalizer_cls = ResultLegendFormalizer
        elif query_type == QueryType.preview:
            legend_formalizer_cls = PreviewLegendFormalizer
        elif query_type == QueryType.distinct:
            legend_formalizer_cls = DistinctLegendFormalizer
        elif query_type == QueryType.value_range:
            legend_formalizer_cls = RangeLegendFormalizer
        else:
            raise ValueError(f"Legend formalization is not supported for {query_type.name} query type")
        return legend_formalizer_cls(
            dataset=self.dataset,
            autofill_legend=autofill_legend,
        )

    def make_block_formalizer(self) -> BlockFormalizer:
        return BlockFormalizer(
            dataset=self.dataset,
            reporting_registry=self.dl_request.services_registry.get_reporting_registry(),
        )

    async def _call_post_exec_async_hook(self, target_connection_ids: set[str]) -> None:
        for connection_id in target_connection_ids:
            connection = self.dl_request.us_manager.get_loaded_us_connection(connection_id)
            lifecycle_manager = self.dl_request.us_manager.get_lifecycle_manager(
                entry=connection,
                service_registry=self.dl_request.services_registry,
            )
            await lifecycle_manager.post_exec_async_hook()

    async def execute_query(
        self,
        block_spec: BlockSpec,
        possible_data_lengths: Optional[Collection] = None,
        profiling_postfix: str = "",
    ) -> PostprocessedQuery:
        # TODO: Move to a separate class

        us_manager = self.dl_request.us_manager

        ds_view = DatasetView(
            ds=self.dataset,
            us_manager=us_manager,
            block_spec=block_spec,
            rci=self.dl_request.rci,
        )

        with GenericProfiler(f"{self.profiler_prefix}-query-build{profiling_postfix}"):
            exec_info = ds_view.build_exec_info()

        async with self.default_query_execution_cm_stack(exec_info, body=self.dl_request.json):
            executed_query = await ds_view.get_data_async(
                exec_info=exec_info,
                allow_cache_usage=self.allow_query_cache_usage,
            )
            if possible_data_lengths is not None:
                assert len(executed_query.rows) in possible_data_lengths

        postprocessor = DataPostprocessor(profiler_prefix=self.profiler_prefix)
        postprocessed_query = postprocessor.get_postprocessed_data(
            executed_query=executed_query,
            block_spec=block_spec,
        )

        return postprocessed_query

    async def execute_all_queries(
        self,
        raw_query_spec_union: RawQuerySpecUnion,
        autofill_legend: bool,
        call_post_exec_async_hook: bool = False,
    ) -> MergedQueryDataStream:
        # TODO: Move to a separate class

        legend_formalizer = self.make_legend_formalizer(
            query_type=raw_query_spec_union.meta.query_type, autofill_legend=autofill_legend
        )
        legend = legend_formalizer.make_legend(raw_query_spec_union=raw_query_spec_union)

        block_legend = self.make_block_formalizer().make_block_legend(
            raw_query_spec_union=raw_query_spec_union,
            legend=legend,
        )
        paginator = QueryPaginator()
        pre_paginator = paginator.get_pre_paginator()
        post_paginator = paginator.get_post_paginator()
        block_legend = pre_paginator.pre_paginate(block_legend=block_legend)

        runner = ConcurrentTaskRunner()
        for block_spec in block_legend.blocks:
            await runner.schedule(self.execute_query(block_spec=block_spec))
        executed_queries = await runner.finalize()
        postprocessed_query_blocks = [
            PostprocessedQueryBlock.from_block_spec(block_spec, postprocessed_query=postprocessed_query)
            for block_spec, postprocessed_query in zip(block_legend.blocks, executed_queries)
        ]

        postprocessed_query_union = PostprocessedQueryUnion(
            blocks=postprocessed_query_blocks,
            legend=legend,
            limit=block_legend.meta.limit,
            offset=block_legend.meta.offset,
        )
        merged_stream = DataStreamMerger().merge(postprocessed_query_union=postprocessed_query_union)
        merged_stream = post_paginator.post_paginate(merged_stream=merged_stream)

        if call_post_exec_async_hook:
            await self._call_post_exec_async_hook(merged_stream.meta.target_connection_ids)

        return merged_stream

    def _make_response_v1(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
        totals_query: Optional[str] = None,
        totals: Optional[PostprocessedRow] = None,
    ) -> Dict[str, Any]:
        add_fields_data = req_model.add_fields_data
        fields_data: Optional[List[Dict[str, Any]]] = None
        if add_fields_data:
            fields_data = get_fields_data_serializable(self.dataset, for_result=True)
            LOGGER.info("Field schema data", extra=dict(fields=fields_data))

        response_json = DataRequestResponseSerializer.make_data_response_v1(
            merged_stream=merged_stream,
            totals=totals,
            totals_query=totals_query,
            data_export_forbidden=self.get_data_export_forbidden_flag(),
            fields_data=fields_data,
        )
        return response_json

    def _make_response_v2(self, merged_stream: MergedQueryDataStream) -> Dict[str, Any]:
        result = DataRequestResponseSerializer.make_data_response_v2(
            merged_stream=merged_stream,
            reporting_registry=self.dl_request.services_registry.get_reporting_registry()
            if self.allow_notifications
            else None,
            data_export_forbidden=self.get_data_export_forbidden_flag(),
        )
        return result

    def get_data_export_forbidden_flag(self) -> bool:
        role = self.resolve_dataset_source_role(dataset=self.dataset)
        avatar_ids = [avatar.id for avatar in self.ds_accessor.get_avatar_list()]
        dsrc_coll_factory = DataSourceCollectionFactory(us_entry_buffer=self.dl_request.us_manager.get_entry_buffer())
        for avatar_id in avatar_ids:
            avatar = self.ds_accessor.get_avatar_strict(avatar_id=avatar_id)
            dsrc_coll_spec = self.ds_accessor.get_data_source_coll_spec_strict(source_id=avatar.source_id)
            dsrc_coll = dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)

            dsrc: DataSource
            if DataSourceRole.origin in dsrc_coll:
                dsrc = dsrc_coll.get_strict(role=DataSourceRole.origin)
            else:
                dsrc = dsrc_coll.get_strict(role)

            if dsrc.data_export_forbidden:
                return True

        return False

    def check_dataset_revision_id(self, req_model: DataRequestModel) -> None:
        dataset_revision_id = req_model.dataset_revision_id
        if dataset_revision_id is not None and dataset_revision_id != self.dataset.revision_id:
            LOGGER.warning(
                f"Dataset revision id mismatch: got {dataset_revision_id} from the request, "
                f"but found {self.dataset.revision_id} in the current dataset"
            )

lib/dl_api_lib/dl_api_lib/app/data_api/resources/dataset/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
)

from aiohttp import web

from dl_api_lib.app.data_api.resources.base import (
    RequiredResourceDSAPI,
    requires,
)
from dl_api_lib.app.data_api.resources.dataset.base import DatasetDataBaseView
import dl_api_lib.schemas.data
import dl_api_lib.schemas.main
from dl_app_tools.profiling_base import generic_profiler_async
from dl_query_processing.merging.primitives import MergedQueryDataStream


if TYPE_CHECKING:
    from aiohttp.web_response import Response

    from dl_api_lib.request_model.data import DataRequestModel

LOGGER = logging.getLogger(__name__)


class DatasetDistinctView(DatasetDataBaseView, abc.ABC):
    endpoint_code = "DatasetVersionValuesDistinct"
    # TODO FIX: Move to constants
    profiler_prefix = "distinct"

    # TODO FIX: Move request deserialization logic to schema
    # TODO FIX: Move response serialization logic to schema
    # TODO FIX: Add docs/schemas decorator
    # @schematic_request(
    #     ns=ns, body=dl_api_lib.schemas.data.DatasetVersionValuesDistinctPostSchema(),
    #     responses={200: ('Success', dl_api_lib.schemas.data.DatasetVersionValuesDistinctResponseSchema())}
    # )
    @generic_profiler_async("ds-values-distinct-full")
    @DatasetDataBaseView.with_resolved_entities
    @requires(RequiredResourceDSAPI.JSON_REQUEST)
    async def post(self) -> Response:
        req_model: DataRequestModel = self.load_req_model()
        self.check_dataset_revision_id(req_model)

        enable_mutation_caching = self.dl_request.services_registry.get_mutation_cache_factory() is not None
        await self.prepare_dataset_for_request(req_model=req_model, enable_mutation_caching=enable_mutation_caching)

        merged_stream = await self.execute_all_queries(
            raw_query_spec_union=req_model.raw_query_spec_union,
            autofill_legend=req_model.autofill_legend,
            call_post_exec_async_hook=False,
        )

        response_json = self.make_response(req_model=req_model, merged_stream=merged_stream)
        return web.json_response(response_json)

    @abc.abstractmethod
    def load_req_model(self) -> DataRequestModel:
        raise NotImplementedError()

    @abc.abstractmethod
    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        raise NotImplementedError()


class DatasetDistinctViewV1(DatasetDistinctView):
    """
    Old API v1 format (input and output)
    """

    def load_req_model(self) -> DataRequestModel:
        schema = dl_api_lib.schemas.data.DatasetVersionValuesDistinctPostSchema()
        req_model: DataRequestModel = schema.load(self.dl_request.json)
        return req_model

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        return self._make_response_v1(req_model=req_model, merged_stream=merged_stream)


class DatasetDistinctViewV1_5(DatasetDistinctView):
    """
    Same as v1, for full v1.5 coverage
    """

    def load_req_model(self) -> DataRequestModel:
        schema = dl_api_lib.schemas.data.DistinctDataRequestV2Schema()
        req_model: DataRequestModel = schema.load(self.dl_request.json)
        return req_model

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        return self._make_response_v1(req_model=req_model, merged_stream=merged_stream)


class DatasetDistinctViewV2(DatasetDistinctView):
    """
    New API v2 format (input and output)
    """

    def load_req_model(self) -> DataRequestModel:
        schema = dl_api_lib.schemas.data.DistinctDataRequestV2Schema()
        req_model: DataRequestModel = schema.load(self.dl_request.json)
        return req_model

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        return self._make_response_v2(merged_stream=merged_stream)

lib/dl_api_lib/dl_api_lib/app/data_api/resources/dataset/distinct.py
datalens-tech/datalens-backend
from __future__ import annotations

from aiohttp import web

from dl_api_lib.api_common.data_serialization import get_fields_data_serializable
from dl_api_lib.app.data_api.resources.dataset.base import DatasetDataBaseView


# TODO FIX: Generalize with sync version
class DatasetFieldsView(DatasetDataBaseView):
    endpoint_code = "DatasetFieldsGet"

    @DatasetDataBaseView.with_resolved_entities
    async def get(self) -> web.Response:
        fields = get_fields_data_serializable(self.dataset)
        return web.json_response(dict(fields=fields, revision_id=self.dataset.revision_id))

lib/dl_api_lib/dl_api_lib/app/data_api/resources/dataset/fields.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from aiohttp import web

from dl_api_lib.api_common.data_serialization import PivotDataRequestResponseSerializer
from dl_api_lib.app.data_api.resources.base import (
    RequiredResourceDSAPI,
    requires,
)
from dl_api_lib.app.data_api.resources.dataset.base import DatasetDataBaseView
from dl_api_lib.pivot.pandas.transformer import PdPivotTransformer
from dl_api_lib.query.formalization.pivot_formalizer import PivotFormalizer
from dl_api_lib.query.formalization.pivot_legend import PivotLegend
from dl_api_lib.request_model.normalization.drm_normalizer_pivot import PivotRequestModelNormalizer
import dl_api_lib.schemas.data
import dl_api_lib.schemas.main
from dl_app_tools.profiling_base import (
    GenericProfiler,
    generic_profiler_async,
)
from dl_query_processing.merging.primitives import MergedQueryDataStream
from dl_utils.utils import exc_catch_awrap


if TYPE_CHECKING:
    from aiohttp.web_response import Response

    from dl_api_lib.pivot.table import PivotTable
    from dl_api_lib.query.formalization.raw_pivot_specs import PivotPaginationSpec
    from dl_api_lib.request_model.data import PivotDataRequestModel
    from dl_query_processing.legend.field_legend import Legend


LOGGER = logging.getLogger(__name__)


class DatasetPivotView(DatasetDataBaseView):
    endpoint_code = "DatasetVersionPivot"
    # TODO FIX: Move to constants
    profiler_prefix = "pivot"

    # TODO FIX: Move request deserialization logic to schema
    # TODO FIX: Move response serialization logic to schema
    # TODO FIX: Add docs/schemas decorator
    # @schematic_request(
    #     ns=ns, body=dl_api_lib.schemas.data.PivotDataRequestBaseSchema(),
    #     responses={200: ('Success', dl_api_lib.schemas.data.DatasetVersionResultResponseSchema())}
    # )
    @generic_profiler_async("ds-pivot-full")
    @exc_catch_awrap
    @DatasetDataBaseView.with_resolved_entities
    @requires(RequiredResourceDSAPI.JSON_REQUEST)
    async def post(self) -> Response:
        schema = dl_api_lib.schemas.data.PivotDataRequestBaseSchema()
        req_model: PivotDataRequestModel = schema.load(self.dl_request.json)

        req_normalizer = PivotRequestModelNormalizer()
        req_model = req_normalizer.normalize_drm(req_model)

        enable_mutation_caching = self.dl_request.services_registry.get_mutation_cache_factory() is not None
        await self.prepare_dataset_for_request(req_model=req_model, enable_mutation_caching=enable_mutation_caching)

        merged_stream = await self.execute_all_queries(
            raw_query_spec_union=req_model.raw_query_spec_union,
            autofill_legend=req_model.autofill_legend,
            call_post_exec_async_hook=True,
        )

        pivot_legend: PivotLegend
        pivot_formalizer = PivotFormalizer(dataset=self.dataset, legend=merged_stream.legend)
        pivot_legend = pivot_formalizer.make_pivot_legend(raw_pivot_spec=req_model.pivot)

        pivot_table = await self.pivot_data(
            merged_stream=merged_stream,
            legend=merged_stream.legend,
            pagination=req_model.pivot.pagination,
            pivot_legend=pivot_legend,
        )

        # Serialize to response
        response_json = self.make_pivot_response(merged_stream=merged_stream, pivot_table=pivot_table)
        return web.json_response(response_json)

    async def pivot_data(
        self,
        merged_stream: MergedQueryDataStream,
        legend: Legend,
        pagination: PivotPaginationSpec,
        pivot_legend: PivotLegend,
    ) -> PivotTable:
        """
        Async wrapper for CPU-intensive pivot table creation
        """
        executor = self.dl_request.services_registry.get_compute_executor()
        return await executor.execute(
            lambda: self._pivot_data_sync(
                merged_stream=merged_stream,
                legend=legend,
                pagination=pagination,
                pivot_legend=pivot_legend,
            )
        )

    def _pivot_data_sync(
        self,
        merged_stream: MergedQueryDataStream,
        legend: Legend,
        pagination: PivotPaginationSpec,
        pivot_legend: PivotLegend,
    ) -> PivotTable:
        pivot_op_name = "pivot"
        # Create pivot table
        with GenericProfiler(f"{self.profiler_prefix}-{pivot_op_name}-transform-full"):
            legend_data = dl_api_lib.schemas.legend.LegendSchema().dump(legend)
            LOGGER.info("Pivot legend", extra=dict(legend=legend_data))
            transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
            with GenericProfiler(f"{self.profiler_prefix}-{pivot_op_name}-transform"):
                pivot_table = transformer.pivot(rows=merged_stream.rows)

            # Sort the table
            with GenericProfiler(f"{self.profiler_prefix}-{pivot_op_name}-sort"):
                pivot_table.facade.sort()
            col_cnt, row_cnt = pivot_table.get_column_count(), pivot_table.get_row_count()
            LOGGER.info(
                "Pivot table has %s columns and %s rows",
                col_cnt,
                row_cnt,
                extra=dict(pivot_table_size=dict(columns=col_cnt, rows=row_cnt)),
            )

            # Paginate
            if pagination.non_default:
                with GenericProfiler(f"{self.profiler_prefix}-{pivot_op_name}-paginate"):
                    pivot_table.facade.paginate(
                        limit_rows=pagination.limit_rows,
                        offset_rows=pagination.offset_rows,
                    )

        return pivot_table

    def make_pivot_response(self, merged_stream: MergedQueryDataStream, pivot_table: PivotTable) -> dict:
        reporting_registry = (
            self.dl_request.services_registry.get_reporting_registry() if self.allow_notifications else None
        )
        return PivotDataRequestResponseSerializer.make_pivot_response(
            merged_stream=merged_stream,
            pivot_table=pivot_table,
            reporting_registry=reporting_registry,
        )

lib/dl_api_lib/dl_api_lib/app/data_api/resources/dataset/pivot.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Collection,
    Dict,
    Optional,
)

from aiohttp import web

from dl_api_lib.app.data_api.resources.base import (
    RequiredResourceDSAPI,
    requires,
)
from dl_api_lib.app.data_api.resources.dataset.base import DatasetDataBaseView
from dl_api_lib.dataset.utils import invalidate_sample_sources
import dl_api_lib.schemas.data
import dl_api_lib.schemas.main
from dl_app_tools.profiling_base import generic_profiler_async
from dl_constants.enums import DataSourceRole
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.data_source.collection import DataSourceCollectionFactory
from dl_core.dataset_capabilities import DatasetCapabilities
import dl_core.exc as common_exc
from dl_query_processing.legend.block_legend import BlockSpec
from dl_query_processing.merging.primitives import MergedQueryDataStream
from dl_query_processing.postprocessing.primitives import PostprocessedQuery


if TYPE_CHECKING:
    from aiohttp.web_response import Response

    from dl_api_lib.request_model.data import DataRequestModel
    from dl_core.us_dataset import Dataset


LOGGER = logging.getLogger(__name__)


class DatasetPreviewView(DatasetDataBaseView, abc.ABC):
    endpoint_code = "DatasetVersionPreview"
    # TODO FIX: Move to constants
    profiler_prefix = "preview"

    STORED_DATASET_REQUIRED = False

    def resolve_dataset_source_role(self, dataset: Dataset, log_reasons: bool = False) -> DataSourceRole:
        dsrc_coll_factory = DataSourceCollectionFactory(us_entry_buffer=self.dl_request.us_manager.get_entry_buffer())
        capabilities = DatasetCapabilities(dataset=dataset, dsrc_coll_factory=dsrc_coll_factory)
        try:
            return capabilities.resolve_source_role(for_preview=True, log_reasons=log_reasons)
        except common_exc.NoCommonRoleError:
            raise common_exc.TableNameNotConfiguredError(
                "Dataset's sources are not configured correctly. Direct access is not possible"
            )

    # TODO FIX: Add docs/schemas decorator
    # @schematic_request(
    #     ns=ns, body=dl_api_lib.schemas.data.DatasetPreviewRequestSchema(),
    #     responses={200: ('Success', dl_api_lib.schemas.data.DatasetVersionResultResponseSchema())}
    # )
    @generic_profiler_async("ds-preview-full")
    @DatasetDataBaseView.with_resolved_entities
    @requires(RequiredResourceDSAPI.JSON_REQUEST)
    async def post(self) -> Response:
        req_model = self.load_req_model()

        update_info = await self.prepare_dataset_for_request(
            req_model=req_model,
            allow_rls_change=False,
        )

        new_sources = update_info.added_own_source_ids + update_info.updated_own_source_ids
        invalidate_sample_sources(
            dataset=self.dataset,
            source_ids=new_sources,
            us_manager=self.dl_request.us_manager,
        )

        merged_stream = await self.execute_all_queries(
            raw_query_spec_union=req_model.raw_query_spec_union,
            autofill_legend=req_model.autofill_legend,
            call_post_exec_async_hook=False,
        )

        response_json = self.make_response(req_model=req_model, merged_stream=merged_stream)
        return web.json_response(response_json)

    async def execute_query(
        self,
        block_spec: BlockSpec,
        possible_data_lengths: Optional[Collection] = None,
        profiling_postfix: str = "",
    ) -> PostprocessedQuery:
        ds_accessor = DatasetComponentAccessor(dataset=self.dataset)
        if not ds_accessor.get_data_source_id_list():
            return PostprocessedQuery()

        return await super().execute_query(
            block_spec=block_spec,
            profiling_postfix=profiling_postfix,
        )

    def load_req_model(self) -> DataRequestModel:
        schema = dl_api_lib.schemas.data.DatasetPreviewRequestSchema()
        req_model: DataRequestModel = schema.load(self.dl_request.json)
        return req_model

    @abc.abstractmethod
    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        raise NotImplementedError()


class DatasetPreviewViewV1(DatasetPreviewView):
    """
    Old API v1 format (input and output)
    """

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        return self._make_response_v1(req_model=req_model, merged_stream=merged_stream)


class DatasetPreviewViewV1_5(DatasetPreviewViewV1):
    """
    Same as v1, for full v1.5 coverage
    """


class DatasetPreviewViewV2(DatasetPreviewView):
    """
    New API v2 format (input and output)
    """

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        return self._make_response_v2(merged_stream=merged_stream)

lib/dl_api_lib/dl_api_lib/app/data_api/resources/dataset/preview.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Collection,
    Dict,
    Optional,
)

from aiohttp import web

from dl_api_lib.app.data_api.resources.base import (
    RequiredResourceDSAPI,
    requires,
)
from dl_api_lib.app.data_api.resources.dataset.base import DatasetDataBaseView
from dl_api_lib.dataset.view import DatasetView
import dl_api_lib.schemas.data
import dl_api_lib.schemas.main
from dl_app_tools.profiling_base import generic_profiler_async
from dl_query_processing.execution.primitives import (
    ExecutedQuery,
    ExecutedQueryMetaInfo,
)
from dl_query_processing.legend.block_legend import BlockSpec
from dl_query_processing.merging.primitives import MergedQueryDataStream
from dl_query_processing.postprocessing.postprocessor import DataPostprocessor
from dl_query_processing.postprocessing.primitives import PostprocessedQuery
from dl_query_processing.translation.primitives import DetailedType
from dl_utils.utils import enum_not_none


if TYPE_CHECKING:
    from aiohttp.web_response import Response

    from dl_api_lib.request_model.data import DataRequestModel


# TODO FIX: Generalize with sync version
class DatasetRangeView(DatasetDataBaseView, abc.ABC):
    endpoint_code = "DatasetVersionValuesRange"
    # TODO FIX: Move to constants
    profiler_prefix = "range"
    transpose_data: ClassVar[bool]

    # TODO FIX: Move request deserialization logic to schema
    # TODO FIX: Move response serialization logic to schema
    # TODO FIX: Add docs/schemas decorator
    # @schematic_request(
    #     ns=ns, body=dl_api_lib.schemas.data.DatasetVersionValuesRangePostSchema(),
    #     responses={200: ('Success', dl_api_lib.schemas.data.DatasetVersionValuesRangeResponseSchema())}
    # )
    @generic_profiler_async("ds-values-range-full")
    @DatasetDataBaseView.with_resolved_entities
    @requires(RequiredResourceDSAPI.JSON_REQUEST)
    async def post(self) -> Response:
        req_model = self.load_req_model()
        self.check_dataset_revision_id(req_model)

        enable_mutation_caching = self.dl_request.services_registry.get_mutation_cache_factory() is not None
        await self.prepare_dataset_for_request(req_model=req_model, enable_mutation_caching=enable_mutation_caching)

        merged_stream = await self.execute_all_queries(
            raw_query_spec_union=req_model.raw_query_spec_union,
            autofill_legend=req_model.autofill_legend,
            call_post_exec_async_hook=False,
        )

        response_json = self.make_response(req_model=req_model, merged_stream=merged_stream)
        return web.json_response(response_json)

    @abc.abstractmethod
    def load_req_model(self) -> DataRequestModel:
        raise NotImplementedError()

    @abc.abstractmethod
    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        raise NotImplementedError()

    async def execute_query(
        self,
        block_spec: BlockSpec,
        possible_data_lengths: Optional[Collection] = None,
        profiling_postfix: str = "",
    ) -> PostprocessedQuery:
        us_manager = self.dl_request.us_manager

        ds_view = DatasetView(
            ds=self.dataset,
            us_manager=us_manager,
            block_spec=block_spec,
            rci=self.dl_request.rci,
        )

        # try getting cached values from source (MetricaAPI)
        field, min_value, max_value = ds_view.fast_get_expression_value_range()

        if min_value is not None and max_value is not None:
            field_id = field.guid
            meta_for_range = ExecutedQueryMetaInfo(
                phantom_select_ids=[],
                field_order=[(0, field_id), (1, field_id)],
                detailed_types=[
                    DetailedType(
                        field_id=field_id,
                        data_type=enum_not_none(field.data_type),
                        formula_data_type=None,
                        formula_data_type_params=None,
                    )
                    for _ in range(2)
                ],
            )
            executed_query = ExecutedQuery(rows=[[min_value, max_value]], meta=meta_for_range)
            postprocessor = DataPostprocessor(profiler_prefix=self.profiler_prefix)
            postprocessed_query = postprocessor.get_postprocessed_data(
                executed_query=executed_query,
                block_spec=block_spec,
            )

        else:
            postprocessed_query = await super().execute_query(
                block_spec=block_spec,
                profiling_postfix=profiling_postfix,
            )

        if self.transpose_data:
            # In v1 and v1.5 min and max are returned as two rows of the same column,
            # so the response has to be transposed
            postprocessed_query = self._transpose_range_data(postprocessed_query)
        return postprocessed_query

    def _transpose_range_data(self, postprocessed_query: PostprocessedQuery) -> PostprocessedQuery:
        min_value, max_value = next(iter(postprocessed_query.postprocessed_data))

        field_order = postprocessed_query.meta.field_order
        assert field_order is not None
        detailed_types = postprocessed_query.meta.detailed_types
        assert detailed_types is not None

        postprocessed_query = postprocessed_query.clone(
            postprocessed_data=[[min_value], [max_value]],
            meta=postprocessed_query.meta.clone(
                field_order=field_order[:1],
                detailed_types=detailed_types[:1],
            ),
        )
        return postprocessed_query


class DatasetRangeViewV1(DatasetRangeView):
    """
    Old API v1 format (input and output)
    """

    transpose_data = True

    def load_req_model(self) -> DataRequestModel:
        schema = dl_api_lib.schemas.data.DatasetVersionValuesRangePostSchema()
        req_model: DataRequestModel = schema.load(self.dl_request.json)
        return req_model

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        return self._make_response_v1(req_model=req_model, merged_stream=merged_stream)


class DatasetRangeViewV1_5(DatasetRangeView):
    """
    Same as v1, for full v1.5 coverage
    """

    transpose_data = True

    def load_req_model(self) -> DataRequestModel:
        schema = dl_api_lib.schemas.data.RangeDataRequestV2Schema()
        req_model: DataRequestModel = schema.load(self.dl_request.json)
        return req_model

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        return self._make_response_v1(req_model=req_model, merged_stream=merged_stream)


class DatasetRangeViewV2(DatasetRangeView):
    """
    New API v2 format (input and output)
    """

    transpose_data = False

    def load_req_model(self) -> DataRequestModel:
        schema = dl_api_lib.schemas.data.RangeDataRequestV2Schema()
        req_model: DataRequestModel = schema.load(self.dl_request.json)
        return req_model

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
    ) -> Dict[str, Any]:
        return self._make_response_v2(merged_stream=merged_stream)

lib/dl_api_lib/dl_api_lib/app/data_api/resources/dataset/range.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Optional,
    Tuple,
)

from aiohttp import web

from dl_api_lib.app.data_api.resources.base import (
    RequiredResourceDSAPI,
    requires,
)
from dl_api_lib.app.data_api.resources.dataset.base import DatasetDataBaseView
from dl_api_lib.query.formalization.raw_specs import RawQuerySpecUnion
import dl_api_lib.schemas.data
import dl_api_lib.schemas.main
from dl_app_tools.profiling_base import generic_profiler_async
from dl_constants.enums import (
    FieldRole,
    FieldType,
    UserDataType,
)
from dl_query_processing.enums import (
    EmptyQueryMode,
    GroupByPolicy,
    QueryType,
)
from dl_query_processing.legend.block_legend import BlockSpec
from dl_query_processing.legend.field_legend import (
    FieldObjSpec,
    Legend,
    LegendItem,
    PlaceholderObjSpec,
    RoleSpec,
    TemplateRoleSpec,
)
from dl_query_processing.merging.primitives import MergedQueryDataStream
from dl_query_processing.postprocessing.primitives import PostprocessedRow
from dl_utils.utils import exc_catch_awrap


if TYPE_CHECKING:
    from aiohttp.web_response import Response

    from dl_api_lib.request_model.data import DataRequestModel


LOGGER = logging.getLogger(__name__)


class DatasetResultView(DatasetDataBaseView, abc.ABC):
    endpoint_code = "DatasetVersionResult"
    # TODO FIX: Move to constants
    profiler_prefix = "result"

    # TODO FIX: Move request deserialization logic to schema
    # TODO FIX: Move response serialization logic to schema
    # TODO FIX: Add docs/schemas decorator
    # @schematic_request(
    #     ns=ns, body=dl_api_lib.schemas.data.DatasetVersionResultRequestSchema(),
    #     responses={200: ('Success', dl_api_lib.schemas.data.DatasetVersionResultResponseSchema())}
    # )
    @generic_profiler_async("ds-result-full")
    @exc_catch_awrap
    @DatasetDataBaseView.with_resolved_entities
    @requires(RequiredResourceDSAPI.JSON_REQUEST)
    async def post(self) -> Response:
        req_model = self.load_req_model()
        self.check_dataset_revision_id(req_model)

        enable_mutation_caching = self.dl_request.services_registry.get_mutation_cache_factory() is not None
        await self.prepare_dataset_for_request(req_model=req_model, enable_mutation_caching=enable_mutation_caching)

        merged_stream = await self.execute_all_queries(
            raw_query_spec_union=req_model.raw_query_spec_union,
            autofill_legend=req_model.autofill_legend,
            call_post_exec_async_hook=True,
        )

        totals, totals_query = None, None
        if req_model.with_totals:
            totals_query, totals = await self._make_totals(
                raw_query_spec_union=req_model.raw_query_spec_union,
                legend=merged_stream.legend,
            )

        response_json = self.make_response(
            req_model=req_model,
            merged_stream=merged_stream,
            totals=totals,
            totals_query=totals_query,
        )
        return web.json_response(response_json)

    @abc.abstractmethod
    def load_req_model(self) -> DataRequestModel:
        raise NotImplementedError

    @abc.abstractmethod
    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
        totals_query: Optional[str],
        totals: Optional[PostprocessedRow],
    ) -> Dict[str, Any]:
        raise NotImplementedError()

    async def _make_totals(
        self,
        *,
        raw_query_spec_union: RawQuerySpecUnion,
        legend: Legend,
    ) -> Tuple[Optional[str], Optional[PostprocessedRow]]:
        block_id = 0
        assert all(item.block_id is None or item.block_id == block_id for item in legend.items)
        streamable_items = legend.list_streamable_items()

        # Patch legend by replacing all dimensions with placeholders
        updated_items: List[LegendItem] = []
        for item in streamable_items:
            if item.field_type == FieldType.DIMENSION:
                updated_item = item.clone(
                    obj=PlaceholderObjSpec(),
                    role_spec=TemplateRoleSpec(role=FieldRole.template, template=None),
                    data_type=UserDataType.string,
                )
            else:
                assert isinstance(item.obj, FieldObjSpec)
                updated_item = item.clone(role_spec=RoleSpec(role=FieldRole.total))
            updated_items.append(updated_item)

        # Add filters
        updated_items += legend.list_for_role(FieldRole.filter)

        legend_for_block = Legend(items=updated_items)

        block_spec = BlockSpec(  # Fake block spec just for use in the  postprocessing interface
            block_id=block_id,
            parent_block_id=None,
            legend_item_ids=[item.legend_item_id for item in streamable_items],
            legend=legend_for_block,
            query_type=QueryType.totals,
            group_by_policy=GroupByPolicy.force,
            ignore_nonexistent_filters=raw_query_spec_union.ignore_nonexistent_filters,
            disable_rls=raw_query_spec_union.disable_rls,
            allow_measure_fields=raw_query_spec_union.allow_measure_fields,
            empty_query_mode=EmptyQueryMode.empty_row,
        )
        postprocessed_query = await self.execute_query(
            block_spec=block_spec,
            possible_data_lengths=(0, 1),
        )

        first_item: Optional[PostprocessedRow] = None
        if postprocessed_query.postprocessed_data:
            first_item = postprocessed_query.postprocessed_data[0]

        return (
            postprocessed_query.meta.debug_query,
            first_item,
        )


class DatasetResultViewV1(DatasetResultView):
    """
    Old API v1 format (input and output)
    """

    def load_req_model(self) -> DataRequestModel:
        schema = dl_api_lib.schemas.data.DatasetVersionResultRequestSchema()
        req_model: DataRequestModel = schema.load(self.dl_request.json)
        return req_model

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
        totals_query: Optional[str],
        totals: Optional[PostprocessedRow],
    ) -> Dict[str, Any]:
        return self._make_response_v1(
            req_model=req_model,
            merged_stream=merged_stream,
            totals_query=totals_query,
            totals=totals,
        )


class DatasetResultViewV1_5(DatasetResultView):
    """
    Accepts requests in v2 format, but responds in old v1 format (with YQL schema).
    To make the transition smoother for the frontend.
    """

    def load_req_model(self) -> DataRequestModel:
        schema = dl_api_lib.schemas.data.ResultDataRequestV1_5Schema()
        req_model: DataRequestModel = schema.load(self.dl_request.json)
        return req_model

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
        totals_query: Optional[str],
        totals: Optional[PostprocessedRow],
    ) -> Dict[str, Any]:
        return self._make_response_v1(
            req_model=req_model,
            merged_stream=merged_stream,
            totals_query=totals_query,
            totals=totals,
        )


class DatasetResultViewV2(DatasetResultView):
    """
    New API v2 format (input and output)
    """

    def load_req_model(self) -> DataRequestModel:
        schema = dl_api_lib.schemas.data.ResultDataRequestV2Schema()
        req_model: DataRequestModel = schema.load(self.dl_request.json)
        return req_model

    def make_response(
        self,
        req_model: DataRequestModel,
        merged_stream: MergedQueryDataStream,
        totals_query: Optional[str],
        totals: Optional[PostprocessedRow],
    ) -> Dict[str, Any]:
        return self._make_response_v2(merged_stream=merged_stream)

lib/dl_api_lib/dl_api_lib/app/data_api/resources/dataset/result.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_core.aio.metrics_view import MetricsView

from .base import BaseView


class DSDataApiMetricsView(MetricsView, BaseView):
    """View for solomon-agent, based on dl_core class"""

lib/dl_api_lib/dl_api_lib/app/data_api/resources/metrics.py
datalens-tech/datalens-backend
# TODO FIX: Move to bi-common

from __future__ import annotations

import logging
from typing import (
    ClassVar,
    Optional,
)

import aiohttp
from aiohttp import (
    ClientTimeout,
    TCPConnector,
    web,
)

from dl_api_commons.aiohttp.aiohttp_wrappers import RequiredResourceCommon
from dl_compeng_pg.compeng_pg_base.data_processor_service_pg import CompEngPgService
from dl_configs.enums import (
    RQE_SERVICES,
    RequiredService,
)
from dl_configs.rqe import RQEBaseURL
from dl_core.connection_executors.models.constants import HEADER_REQUEST_ID
from dl_core.services_registry.conn_executor_factory import DefaultConnExecutorFactory
from dl_utils.aio import timeout

from .base import (
    BaseView,
    requires,
)


LOGGER = logging.getLogger(__name__)


@requires(RequiredResourceCommon.SKIP_AUTH)
class PingView(BaseView):
    async def get(self) -> web.Response:
        return web.json_response(
            {
                "request_id": self.dl_request.rci.request_id,
            }
        )


@requires(RequiredResourceCommon.SKIP_AUTH)
class PingReadyView(BaseView):
    default_timeout: ClassVar[float] = 10.0

    async def is_pg_ready(self) -> bool:
        compeng: CompEngPgService = CompEngPgService.get_app_instance(self.request.app)
        processor = compeng.get_data_processor()

        try:
            async with timeout(self.default_timeout):
                await processor.start()
                result = await processor.ping()
                assert result == 1, result  # TODO: encapsulate this logic
        except Exception as e:
            LOGGER.info("Got an exception trying to check PG readiness: %s", e)
            return False
        finally:
            await processor.end()

        return True

    async def is_rqe_ready(self, base_url: RQEBaseURL) -> Optional[int]:
        session = aiohttp.ClientSession(
            connector=TCPConnector(force_close=True),
            timeout=ClientTimeout(
                total=self.default_timeout * 2,
                connect=5,
                sock_read=self.default_timeout * 2,
                sock_connect=5,
            ),
        )

        async with session as rqe_session:
            try:
                resp = await rqe_session.get(
                    url=f"{base_url}/ping",
                    headers={
                        HEADER_REQUEST_ID: self.dl_request.rci.request_id,
                    },
                )
            except Exception as e:
                LOGGER.info("Got an exception trying to check RQE readiness: %s", e)
                return None

        return resp.status

    async def get(self) -> web.Response:
        response_details: dict[str, Optional[int]] = dict()
        required_services = self.dl_request.services_registry.get_required_services()

        if RequiredService.POSTGRES in required_services:
            LOGGER.info("Going to check PG readiness")
            response_details[RequiredService.POSTGRES.name] = await self.is_pg_ready()

        if required_rqe_services := required_services & RQE_SERVICES:
            conn_exec_factory = self.dl_request.services_registry.get_conn_executor_factory()
            assert isinstance(conn_exec_factory, DefaultConnExecutorFactory)
            rqe_config = conn_exec_factory.rqe_config
            assert rqe_config is not None, "RQE readiness check requested, but no rqe_config provided"

            rqe_base_url_map: dict[RequiredService, RQEBaseURL] = {
                RequiredService.RQE_INT_SYNC: rqe_config.int_sync_rqe,
                RequiredService.RQE_EXT_SYNC: rqe_config.ext_sync_rqe,
                RequiredService.RQE_EXT_ASYNC: rqe_config.ext_async_rqe,
            }

            for rqe_service in required_rqe_services:
                rqe_base_url = rqe_base_url_map[rqe_service]
                LOGGER.info(f"Going to check RQE readiness ({rqe_service.name} at {rqe_base_url})")
                response_details[rqe_service.name] = await self.is_rqe_ready(rqe_base_url)

        status = (
            200 if all(ping_result == 200 or ping_result is True for ping_result in response_details.values()) else 503
        )

        return web.json_response(
            status=status,
            data={
                "request_id": self.dl_request.rci.request_id,
                "details": {
                    **response_details,
                },
            },
        )

lib/dl_api_lib/dl_api_lib/app/data_api/resources/ping.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import json

from aiohttp import web
from statcommons.unistat.uwsgi import process_uwsgi_data_for_unistat

from dl_api_commons.aiohttp.aiohttp_wrappers import RequiredResourceCommon

from .base import (
    BaseView,
    requires,
)


# TODO FIX: Move to bi-common
async def read_usocket(path):  # type: ignore  # TODO: fix
    """
    Simple "connect to a unix socket at `path` and read the stream".
    """
    reader, writer = await asyncio.open_unix_connection(path=path)
    try:
        return await reader.read()
    finally:
        writer.close()
        await writer.wait_closed()


async def uwsgi_unistat(path):  # type: ignore  # TODO: fix
    data = await read_usocket(path)
    data = json.loads(data)
    return process_uwsgi_data_for_unistat(data)


@requires(RequiredResourceCommon.SKIP_AUTH)
class UnistatView(BaseView):
    async def get(self) -> web.Response:
        result = await uwsgi_unistat("/tmp/uwsgi_stats_rqe_int_sync.sock")
        result = list(result.items())
        return web.json_response(result)

lib/dl_api_lib/dl_api_lib/app/data_api/resources/unistat.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging.config
from typing import (
    TYPE_CHECKING,
    Generic,
    Optional,
    TypeVar,
    final,
)

import attr

from dl_api_lib.app_common_settings import ConnOptionsMutatorsFactory
from dl_api_lib.app_settings import AppSettings
from dl_api_lib.connector_availability.base import ConnectorAvailabilityConfig
from dl_api_lib.i18n.registry import (
    LOCALIZATION_CONFIGS,
    register_translation_configs,
)
from dl_api_lib.service_registry.sr_factory import DefaultApiSRFactory
from dl_api_lib.service_registry.supported_functions_manager import SupportedFunctionsManager
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.enums import RequiredService
from dl_constants.enums import (
    ConnectionType,
    RLSSubjectType,
)
from dl_core.data_processing.cache.primitives import CacheTTLConfig
from dl_core.rls import (
    RLS_FAILED_USER_NAME_PREFIX,
    BaseSubjectResolver,
    RLSSubject,
)
from dl_core.services_registry.entity_checker import EntityUsageChecker
from dl_core.services_registry.file_uploader_client_factory import FileUploaderSettings
from dl_core.services_registry.inst_specific_sr import (
    InstallationSpecificServiceRegistry,
    InstallationSpecificServiceRegistryFactory,
)
from dl_core.services_registry.rqe_caches import RQECachesSetting
from dl_core.services_registry.top_level import ServicesRegistry
from dl_core.utils import FutureRef
from dl_i18n.localizer_base import (
    LocalizerLoader,
    TranslationConfig,
)
from dl_task_processor.arq_wrapper import create_arq_redis_settings


if TYPE_CHECKING:
    from dl_core.services_registry.env_manager_factory_base import EnvManagerFactory


LOGGER = logging.getLogger(__name__)


@attr.s
class NotFoundSubjectResolver(BaseSubjectResolver):
    def get_subjects_by_names(self, names: list[str]) -> list[RLSSubject]:
        return [
            RLSSubject(
                subject_id="",
                subject_type=RLSSubjectType.notfound,
                subject_name=RLS_FAILED_USER_NAME_PREFIX + name,
            )
            for name in names
        ]


@attr.s
class StandaloneServiceRegistry(InstallationSpecificServiceRegistry):
    async def get_subject_resolver(self) -> BaseSubjectResolver:
        return NotFoundSubjectResolver()


@attr.s
class StandaloneServiceRegistryFactory(InstallationSpecificServiceRegistryFactory):
    def get_inst_specific_sr(self, sr_ref: FutureRef[ServicesRegistry]) -> StandaloneServiceRegistry:
        return StandaloneServiceRegistry(service_registry_ref=sr_ref)


TSettings = TypeVar("TSettings", bound=AppSettings)


class SRFactoryBuilder(Generic[TSettings], abc.ABC):
    @property
    @abc.abstractmethod
    def _is_async_env(self) -> bool:
        raise NotImplementedError()

    @abc.abstractmethod
    def _get_required_services(self, settings: TSettings) -> set[RequiredService]:
        raise NotImplementedError

    @abc.abstractmethod
    def _get_env_manager_factory(self, settings: TSettings) -> EnvManagerFactory:
        raise NotImplementedError

    @abc.abstractmethod
    def _get_inst_specific_sr_factory(
        self,
        settings: TSettings,
    ) -> Optional[InstallationSpecificServiceRegistryFactory]:
        raise NotImplementedError

    @abc.abstractmethod
    def _get_entity_usage_checker(self, settings: TSettings) -> Optional[EntityUsageChecker]:
        raise NotImplementedError

    @abc.abstractmethod
    def _get_bleeding_edge_users(self, settings: TSettings) -> tuple[str, ...]:
        raise NotImplementedError

    @abc.abstractmethod
    def _get_rqe_caches_settings(self, settings: TSettings) -> Optional[RQECachesSetting]:
        raise NotImplementedError

    @abc.abstractmethod
    def _get_default_cache_ttl_settings(self, settings: TSettings) -> Optional[CacheTTLConfig]:
        raise NotImplementedError

    @abc.abstractmethod
    def _get_connector_availability(self, settings: TSettings) -> Optional[ConnectorAvailabilityConfig]:
        raise NotImplementedError

    @property
    def _extra_translation_configs(self) -> set[TranslationConfig]:
        return set()

    @final
    def get_sr_factory(
        self,
        settings: TSettings,
        conn_opts_factory: ConnOptionsMutatorsFactory,
        connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
    ) -> DefaultApiSRFactory:
        supported_functions_manager = SupportedFunctionsManager(supported_tags=settings.FORMULA_SUPPORTED_FUNC_TAGS)

        file_uploader_settings = (
            FileUploaderSettings(
                base_url=settings.FILE_UPLOADER_BASE_URL,
                master_token=settings.FILE_UPLOADER_MASTER_TOKEN,
            )
            if settings.FILE_UPLOADER_BASE_URL and settings.FILE_UPLOADER_MASTER_TOKEN
            else None
        )

        required_services: set[RequiredService] = set(self._get_required_services(settings))
        if settings.BI_COMPENG_PG_ON:
            required_services.add(RequiredService.POSTGRES)

        register_translation_configs(self._extra_translation_configs)
        localization_loader = LocalizerLoader(configs=LOCALIZATION_CONFIGS)
        localization_factory = localization_loader.load()
        localizer_fallback = (
            localization_factory.get_for_locale(locale=settings.DEFAULT_LOCALE) if settings.DEFAULT_LOCALE else None
        )

        sr_factory = DefaultApiSRFactory(
            async_env=self._is_async_env,
            rqe_config=settings.RQE_CONFIG,
            default_cache_ttl_config=self._get_default_cache_ttl_settings(settings),
            bleeding_edge_users=self._get_bleeding_edge_users(settings),
            connect_options_factory=conn_opts_factory,
            env_manager_factory=self._get_env_manager_factory(settings),
            default_formula_parser_type=settings.FORMULA_PARSER_TYPE,
            connectors_settings=connectors_settings,
            entity_usage_checker=self._get_entity_usage_checker(settings),
            field_id_generator_type=settings.FIELD_ID_GENERATOR_TYPE,
            file_uploader_settings=file_uploader_settings,
            redis_pool_settings=create_arq_redis_settings(settings.REDIS_ARQ) if settings.REDIS_ARQ else None,
            rqe_caches_settings=self._get_rqe_caches_settings(settings),
            supported_functions_manager=supported_functions_manager,
            required_services=required_services,
            localizer_factory=localization_factory,
            localizer_fallback=localizer_fallback,
            connector_availability=self._get_connector_availability(settings),
            inst_specific_sr_factory=self._get_inst_specific_sr_factory(settings),
            force_non_rqe_mode=settings.RQE_FORCE_OFF,
            query_proc_mode=settings.QUERY_PROCESSING_MODE,
        )
        return sr_factory

lib/dl_api_lib/dl_api_lib/app_common.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Callable,
    List,
    Optional,
)


if TYPE_CHECKING:
    from dl_core.connection_models import ConnectOptions
    from dl_core.us_connection_base import ExecutorBasedMixin


_T_CONN_OPTIONS_MUTATOR = Callable[["ConnectOptions", "ExecutorBasedMixin"], Optional["ConnectOptions"]]


class ConnOptionsMutatorsFactory:
    def __init__(self) -> None:
        self.mutators: List[_T_CONN_OPTIONS_MUTATOR] = []

    def add_mutator(self, func: _T_CONN_OPTIONS_MUTATOR) -> None:
        self.mutators.append(func)

    def __call__(self, conn: ExecutorBasedMixin) -> Optional[ConnectOptions]:
        any_mutated = False
        conn_opts = conn.get_conn_options()
        for func in self.mutators:
            mutated = func(conn_opts, conn)
            if mutated is not None:
                any_mutated = True
                conn_opts = mutated
        # TODO?: is there any sense in returning `None`, though?
        return conn_opts if any_mutated else None

lib/dl_api_lib/dl_api_lib/app_common_settings.py
datalens-tech/datalens-backend
from typing import (
    Collection,
    Optional,
    Type,
)

import attr

from dl_api_connector.connector import ApiConnector
import dl_api_lib as package
from dl_api_lib.connector_registrator import CONN_REG_BI_API
from dl_utils.entrypoints import EntrypointClassManager


def _register_connector(connector_cls: Type[ApiConnector]) -> None:
    CONN_REG_BI_API.register_connector(connector_cls)


_CONNECTOR_EP_GROUP = f"{package.__name__}.connectors"


@attr.s
class ApiConnectorEntrypointManager(EntrypointClassManager[ApiConnector]):
    entrypoint_group_name = attr.ib(init=False, default=_CONNECTOR_EP_GROUP)


def register_all_connectors(connector_ep_names: Optional[Collection[str]] = None) -> None:
    connectors = ApiConnectorEntrypointManager().get_all_ep_classes()
    for ep_name, connector_cls in sorted(connectors.items()):
        if connector_ep_names is not None and ep_name not in connector_ep_names:
            continue
        CONN_REG_BI_API.register_connector(connector_cls)

lib/dl_api_lib/dl_api_lib/app_connectors.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
    Optional,
)

import attr

from dl_api_commons.base_models import TenantDef
from dl_api_lib.connector_availability.base import ConnectorAvailabilityConfig
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_configs.enums import RedisMode
from dl_configs.environments import is_setting_applicable
from dl_configs.rqe import RQEConfig
from dl_configs.settings_loaders.meta_definition import (
    required,
    s_attrib,
)
from dl_configs.settings_loaders.settings_obj_base import SettingsBase
from dl_configs.settings_submodels import RedisSettings
from dl_configs.utils import split_by_comma
from dl_constants.enums import (
    QueryProcessingMode,
    USAuthMode,
)
from dl_core.components.ids import FieldIdGeneratorType
from dl_formula.parser.factory import ParserType


@attr.s(frozen=True)
class CachesTTLSettings(SettingsBase):
    MATERIALIZED: Optional[int] = s_attrib("SEC_MATERIALIZED_DATASET")  # type: ignore
    OTHER: Optional[int] = s_attrib("SEC_OTHER")  # type: ignore


def _list_to_tuple(value: Any) -> Any:
    if isinstance(value, list):
        return tuple(value)
    return value


@attr.s(frozen=True)
class AppSettings:
    BLEEDING_EDGE_USERS: tuple[str, ...] = s_attrib(  # type: ignore
        "DL_BLEEDING_EDGE_USERS",
        env_var_converter=split_by_comma,
        missing=(),
    )

    SENTRY_ENABLED: bool = s_attrib("DL_SENTRY_ENABLED", missing=False)  # type: ignore
    SENTRY_DSN: Optional[str] = s_attrib("DL_SENTRY_DSN", missing=None)  # type: ignore

    CRYPTO_KEYS_CONFIG: CryptoKeysConfig = s_attrib(  # type: ignore
        "DL_CRY",
        json_converter=CryptoKeysConfig.from_json,
        sensitive=True,
    )
    US_BASE_URL: str = s_attrib("US_HOST", fallback_cfg_key="US_BASE_URL")  # type: ignore
    US_MASTER_TOKEN: Optional[str] = s_attrib("US_MASTER_TOKEN", sensitive=True, missing=None)  # type: ignore

    RQE_FORCE_OFF: bool = s_attrib("RQE_FORCE_OFF", missing=False)  # type: ignore
    RQE_CONFIG: RQEConfig = s_attrib("RQE", fallback_factory=RQEConfig.get_default)  # type: ignore
    RQE_CACHES_ON: bool = s_attrib("RQE_CACHES_ON", missing=False)  # type: ignore
    RQE_CACHES_TTL: int = s_attrib("RQE_CACHES_TTL", missing=60 * 10)  # type: ignore
    RQE_CACHES_REDIS: Optional[RedisSettings] = s_attrib(  # type: ignore
        "RQE_CACHES_REDIS",
        fallback_factory=(
            lambda cfg: RedisSettings(  # type: ignore
                MODE=RedisMode.single_host,
                CLUSTER_NAME=cfg.REDIS_RQE_CACHES_CLUSTER_NAME,
                HOSTS=_list_to_tuple(cfg.REDIS_RQE_CACHES_HOSTS),
                PORT=cfg.REDIS_RQE_CACHES_PORT,
                SSL=cfg.REDIS_RQE_CACHES_SSL,
                DB=cfg.REDIS_RQE_CACHES_DB,
                PASSWORD=required(str),
            )
            if is_setting_applicable(cfg, "REDIS_RQE_CACHES_DB")
            else None
        ),
        missing=None,
    )

    SAMPLES_CH_HOSTS: tuple[str, ...] = s_attrib(  # type: ignore
        "SAMPLES_CH_HOST", env_var_converter=split_by_comma, missing_factory=list
    )

    BI_COMPENG_PG_ON: bool = s_attrib("BI_COMPENG_PG_ON", missing=True)  # type: ignore
    BI_COMPENG_PG_URL: Optional[str] = s_attrib("BI_COMPENG_PG_URL", missing=None)  # type: ignore

    FORMULA_PARSER_TYPE: Optional[ParserType] = s_attrib(  # type: ignore
        "BI_FORMULA_PARSER_TYPE",
        env_var_converter=lambda s: ParserType[s.lower()],
        missing=ParserType.antlr_py,
    )
    FORMULA_SUPPORTED_FUNC_TAGS: tuple[str] = s_attrib(
        "FORMULA_SUPPORTED_FUNC_TAGS",
        env_var_converter=split_by_comma,
        missing=("stable",),
    )

    BI_API_CONNECTOR_WHITELIST: Optional[list[str]] = s_attrib(  # type: ignore
        "BI_API_CONNECTOR_WHITELIST",
        env_var_converter=lambda s: list(split_by_comma(s)),
        fallback_cfg_key="BI_API_CONNECTOR_WHITELIST",
        missing=None,
    )
    CORE_CONNECTOR_WHITELIST: Optional[list[str]] = s_attrib(  # type: ignore
        "CORE_CONNECTOR_WHITELIST",
        env_var_converter=lambda s: list(split_by_comma(s)),
        fallback_cfg_key="CORE_CONNECTOR_WHITELIST",
        missing=None,
    )

    FIELD_ID_GENERATOR_TYPE: FieldIdGeneratorType = s_attrib(  # type: ignore
        "FIELD_ID_GENERATOR_TYPE",
        env_var_converter=lambda s: FieldIdGeneratorType[s.lower()],
        missing=FieldIdGeneratorType.readable,
    )

    REDIS_ARQ: Optional[RedisSettings] = None

    FILE_UPLOADER_BASE_URL: Optional[str] = s_attrib(  # type: ignore
        "FILE_UPLOADER_BASE_URL",
        fallback_cfg_key="DATALENS_API_LB_UPLOADS_BASE_URL",
        missing=None,
    )
    FILE_UPLOADER_MASTER_TOKEN: Optional[str] = s_attrib(  # type: ignore
        "FILE_UPLOADER_MASTER_TOKEN",
        sensitive=True,
        missing=None,
    )

    DEFAULT_LOCALE: Optional[str] = "en"

    QUERY_PROCESSING_MODE: QueryProcessingMode = s_attrib(  # type: ignore
        "QUERY_PROCESSING_MODE",
        env_var_converter=lambda s: QueryProcessingMode[s.lower()],
        missing=QueryProcessingMode.basic,
    )


@attr.s(frozen=True)
class ControlApiAppSettings(AppSettings):
    DO_DSRC_IDX_FETCH: bool = s_attrib("DL_DO_DS_IDX_FETCH", missing=False)  # type: ignore

    CONNECTOR_AVAILABILITY: ConnectorAvailabilityConfig = s_attrib(  # type: ignore
        "CONNECTOR_AVAILABILITY",
        fallback_factory=lambda cfg: ConnectorAvailabilityConfig.from_settings(cfg.CONNECTOR_AVAILABILITY),
    )

    REDIS_ARQ: Optional[RedisSettings] = s_attrib(  # type: ignore
        # TODO: move this values to a separate key
        "REDIS_ARQ",
        fallback_factory=(
            lambda cfg: RedisSettings(  # type: ignore
                MODE=RedisMode(cfg.REDIS_PERSISTENT_MODE),
                CLUSTER_NAME=cfg.REDIS_PERSISTENT_CLUSTER_NAME,
                HOSTS=_list_to_tuple(cfg.REDIS_PERSISTENT_HOSTS),
                PORT=cfg.REDIS_PERSISTENT_PORT,
                SSL=cfg.REDIS_PERSISTENT_SSL,
                DB=cfg.REDIS_FILE_UPLOADER_TASKS_DB,
                PASSWORD=required(str),
            )
            if is_setting_applicable(cfg, "REDIS_PERSISTENT_CLUSTER_NAME")
            else None
        ),
        missing=None,
    )

    app_prefix: ClassVar[str] = "a"


@attr.s(frozen=True)
class DataApiAppSettings(AppSettings):
    COMMON_TIMEOUT_SEC: int = s_attrib("COMMON_TIMEOUT_SEC", missing=90)  # type: ignore

    CACHES_ON: bool = s_attrib("CACHES_ON", missing=True)  # type: ignore
    CACHES_REDIS: Optional[RedisSettings] = s_attrib(  # type: ignore
        # TODO: move this values to a separate key
        "CACHES_REDIS",
        fallback_factory=(
            lambda cfg: RedisSettings(  # type: ignore
                MODE=RedisMode.sentinel,
                CLUSTER_NAME=cfg.REDIS_CACHES_CLUSTER_NAME,
                HOSTS=_list_to_tuple(cfg.REDIS_CACHES_HOSTS),
                PORT=cfg.REDIS_CACHES_PORT,
                SSL=cfg.REDIS_CACHES_SSL,
                DB=cfg.REDIS_CACHES_DB,
                PASSWORD=required(str),
            )
            if is_setting_applicable(cfg, "REDIS_CACHES_DB")
            else None
        ),
        missing=None,
    )
    MUTATIONS_CACHES_ON: bool = s_attrib("MUTATIONS_CACHES_ON", missing=False)  # type: ignore
    MUTATIONS_CACHES_DEFAULT_TTL: float = s_attrib("MUTATIONS_CACHES_DEFAULT_TTL", missing=3 * 60 * 60)  # type: ignore
    MUTATIONS_REDIS: Optional[RedisSettings] = s_attrib(
        # TODO: move this values to a separate key
        "MUTATIONS_REDIS",
        fallback_factory=(
            lambda cfg: RedisSettings(  # type: ignore
                MODE=RedisMode.sentinel,
                CLUSTER_NAME=cfg.REDIS_CACHES_CLUSTER_NAME,
                HOSTS=_list_to_tuple(cfg.REDIS_CACHES_HOSTS),
                PORT=cfg.REDIS_CACHES_PORT,
                SSL=cfg.REDIS_CACHES_SSL,
                DB=cfg.REDIS_MUTATIONS_CACHES_DB,
                PASSWORD=required(str),
            )
            if is_setting_applicable(cfg, "REDIS_MUTATIONS_CACHES_DB")
            else None
        ),
        missing=None,
    )
    CACHES_TTL_SETTINGS: Optional[CachesTTLSettings] = s_attrib(  # type: ignore
        "CACHES_TTL",
        fallback_factory=lambda: CachesTTLSettings(  # type: ignore
            MATERIALIZED=3600,
            OTHER=300,
        ),
        missing=None,
    )

    BI_ASYNC_APP_DISABLE_KEEPALIVE: bool = s_attrib("BI_ASYNC_APP_DISABLE_KEEPALIVE", missing=False)  # type: ignore

    @property
    def app_name(self) -> str:
        return "bi_data_api"

    @property
    def jaeger_service_name(self) -> str:
        return "bi-data-api"

    app_prefix: ClassVar[str] = "y"


@attr.s(frozen=True, kw_only=True)
class ControlApiAppTestingsSettings:
    us_auth_mode_override: Optional[USAuthMode] = attr.ib(default=None)
    fake_tenant: Optional[TenantDef] = attr.ib(default=None)

lib/dl_api_lib/dl_api_lib/app_settings.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Type

from dl_api_connector.form_config.models.base import ConnectionFormFactory
from dl_constants.enums import ConnectionType


class NoForm(Exception):
    pass


CONN_FORM_FACTORY_BY_TYPE: dict[ConnectionType, Type[ConnectionFormFactory]] = {}


def register_connection_form_factory_cls(conn_type: ConnectionType, factory_cls: Type[ConnectionFormFactory]) -> None:
    CONN_FORM_FACTORY_BY_TYPE[conn_type] = factory_cls


def get_connection_form_factory_cls(conn_type: ConnectionType) -> Type[ConnectionFormFactory]:
    if (conn_form := CONN_FORM_FACTORY_BY_TYPE.get(conn_type)) is not None:
        return conn_form
    raise NoForm(conn_type)

lib/dl_api_lib/dl_api_lib/connection_forms/registry.py
datalens-tech/datalens-backend
from typing import Type

from dl_api_connector.connection_info import ConnectionInfoProvider
from dl_constants.enums import ConnectionType


class InfoProviderNotFound(Exception):
    pass


CONNECTOR_INFO_PROVIDER_CLS_BY_TYPE: dict[ConnectionType, Type[ConnectionInfoProvider]] = {}


def register_connector_info_provider_cls(
    conn_type: ConnectionType,
    info_provider_cls: Type[ConnectionInfoProvider],
) -> None:
    CONNECTOR_INFO_PROVIDER_CLS_BY_TYPE[conn_type] = info_provider_cls


def get_connector_info_provider_cls(conn_type: ConnectionType) -> Type[ConnectionInfoProvider]:
    if (conn_info_provider := CONNECTOR_INFO_PROVIDER_CLS_BY_TYPE.get(conn_type)) is not None:
        return conn_info_provider
    raise InfoProviderNotFound(conn_type)

lib/dl_api_lib/dl_api_lib/connection_info.py
datalens-tech/datalens-backend
from typing import Optional

from dl_constants.enums import ConnectionType


CONNECTOR_ALIAS_BY_TYPE: dict[ConnectionType, Optional[str]] = {}


def register_connector_alias(conn_type: ConnectionType, alias: Optional[str]) -> None:
    CONNECTOR_ALIAS_BY_TYPE[conn_type] = alias


def get_connector_alias(conn_type: ConnectionType) -> Optional[str]:
    return CONNECTOR_ALIAS_BY_TYPE.get(conn_type, None)

lib/dl_api_lib/dl_api_lib/connector_alias.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import itertools
import logging
from typing import (
    Any,
    Iterator,
    Optional,
)

import attr

from dl_api_lib.connection_forms.registry import CONN_FORM_FACTORY_BY_TYPE
from dl_api_lib.connection_info import get_connector_info_provider_cls
from dl_api_lib.i18n.localizer import Translatable
from dl_configs.connector_availability import (
    ConnectorAvailabilityConfigSettings,
    ConnectorBaseSettings,
    ConnectorContainerSettings,
    ConnectorSettings,
    SectionSettings,
)
from dl_configs.settings_loaders.fallback_cfg_resolver import ObjectLikeConfig
from dl_configs.settings_loaders.meta_definition import s_attrib
from dl_configs.settings_loaders.settings_obj_base import SettingsBase
from dl_configs.utils import conn_type_set_env_var_converter
from dl_constants.enums import (
    ConnectionType,
    ConnectorAvailability,
)
from dl_i18n.localizer_base import Localizer


LOGGER = logging.getLogger(__name__)


class LocalizedSerializable(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def as_dict(self, localizer: Localizer) -> dict[str, Any]:
        raise NotImplementedError


def _make_translatable(text: str, domain: Optional[str]) -> Translatable:
    translatable = Translatable(text)
    if domain is not None:
        translatable.domain = domain
    return translatable


@attr.s(kw_only=True)
class Section(LocalizedSerializable):
    title_translatable: Translatable = attr.ib()
    connectors: list[ConnectorBase] = attr.ib(validator=attr.validators.min_len(1))  # type: ignore

    @classmethod
    def from_settings(cls, settings: SectionSettings) -> Section:
        def _connector_from_settings(settings: ConnectorBaseSettings | ObjectLikeConfig) -> ConnectorBase:
            if hasattr(settings, "conn_type"):
                return Connector.from_settings(settings)  # type: ignore
            elif hasattr(settings, "includes"):
                return ConnectorContainer.from_settings(settings)  # typeL ignore
            raise ValueError('Can\'t create a connector, neither "conn_type" nor "includes" found among settings')

        return cls(
            title_translatable=_make_translatable(settings.title_translatable.text, settings.title_translatable.domain),
            connectors=[_connector_from_settings(item) for item in settings.connectors],
        )

    def as_dict(self, localizer: Localizer) -> dict[str, Any]:
        return dict(
            title=localizer.translate(self.title_translatable),
            connectors=[connector.as_dict(localizer) for connector in self.connectors],
        )


@attr.s(kw_only=True)
class ConnectorBase(LocalizedSerializable, metaclass=abc.ABCMeta):
    @property
    @abc.abstractmethod
    def hidden(self) -> bool:
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def alias(self) -> str:
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def conn_type_str(self) -> str:
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def backend_driven_form(self) -> bool:
        raise NotImplementedError

    @abc.abstractmethod
    def get_title(self, localizer: Localizer) -> str:
        raise NotImplementedError

    def as_dict(self, localizer: Localizer) -> dict[str, Any]:
        return dict(
            hidden=self.hidden,
            alias=self.alias,
            conn_type=self.conn_type_str,
            title=self.get_title(localizer),
            backend_driven_form=self.backend_driven_form,
        )


@attr.s(kw_only=True)
class Connector(ConnectorBase):
    """Represents an actual connection type"""

    _hidden: bool = attr.ib(init=False, default=False)
    conn_type: ConnectionType = attr.ib(
        converter=lambda v: ConnectionType(v) if not isinstance(v, ConnectionType) else v
    )
    availability: ConnectorAvailability = attr.ib(default=ConnectorAvailability.free, converter=ConnectorAvailability)

    @classmethod
    def from_settings(cls, settings: ConnectorSettings) -> Connector:
        return cls(
            conn_type=ConnectionType(settings.conn_type),
            availability=settings.availability,
        )

    @property
    def hidden(self) -> bool:
        return self._hidden

    @property
    def alias(self) -> str:
        return get_connector_info_provider_cls(self.conn_type).alias or self.conn_type_str

    @property
    def conn_type_str(self) -> str:
        return self.conn_type.name

    @property
    def backend_driven_form(self) -> bool:
        return self.conn_type in CONN_FORM_FACTORY_BY_TYPE

    def get_title(self, localizer: Localizer) -> str:
        return localizer.translate(get_connector_info_provider_cls(self.conn_type).title_translatable)


@attr.s(kw_only=True)
class ConnectorContainer(ConnectorBase):
    """
    Acts like a meta-connection, the only purpose of which is to show that it combines other connections
    It is up to the UI to decide what to do with it
    """

    conn_type_str = "__meta__"
    _alias: str = attr.ib()
    includes: list[Connector] = attr.ib(validator=attr.validators.min_len(1))  # type: ignore
    title_translatable: Translatable = attr.ib()

    @property
    def hidden(self) -> bool:
        return all(connector.hidden for connector in self.includes)

    @property
    def alias(self) -> str:
        return self._alias

    @property
    def backend_driven_form(self) -> bool:
        return False

    def get_title(self, localizer: Localizer) -> str:
        return localizer.translate(self.title_translatable)

    def iter_connectors(self) -> Iterator[Connector]:
        return itertools.chain(self.includes)

    @classmethod
    def from_settings(cls, settings: ConnectorContainerSettings) -> ConnectorContainer:
        return cls(
            alias=settings.alias,
            title_translatable=_make_translatable(settings.title_translatable.text, settings.title_translatable.domain),
            includes=[Connector.from_settings(item) for item in settings.includes],
        )

    def as_dict(self, localizer: Localizer) -> dict[str, Any]:
        return dict(
            **super().as_dict(localizer),
            includes=[connector.as_dict(localizer) for connector in self.includes],
        )


@attr.s(kw_only=True)
class ConnectorAvailabilityConfig(SettingsBase):
    uncategorized: list[Connector] = attr.ib(factory=list)
    sections: list[Section] = attr.ib(factory=list)

    visible_connectors: set[ConnectionType] = s_attrib(  # type: ignore
        "VISIBLE",
        env_var_converter=conn_type_set_env_var_converter,
        missing_factory=set,
    )

    def __attrs_post_init__(self) -> None:
        for connector in self._iter_connectors():
            connector._hidden = connector.conn_type not in self.visible_connectors

    @classmethod
    def from_settings(
        cls, settings: ConnectorAvailabilityConfigSettings | ObjectLikeConfig
    ) -> ConnectorAvailabilityConfig:
        visible_connectors: set[ConnectionType] = {ConnectionType(item) for item in settings.visible_connectors}

        return cls(
            uncategorized=[Connector.from_settings(item) for item in settings.uncategorized],
            sections=[Section.from_settings(item) for item in settings.sections],
            visible_connectors=visible_connectors,
        )

    def _iter_connectors(self) -> Iterator[Connector]:
        return itertools.chain(
            self.uncategorized,
            itertools.chain.from_iterable(
                connector.includes if isinstance(connector, ConnectorContainer) else (connector,)  # type: ignore
                for connector in itertools.chain.from_iterable(section.connectors for section in self.sections)
            ),
        )

    def as_dict(self, localizer: Localizer) -> dict[str, Any]:
        return dict(
            uncategorized=[connector.as_dict(localizer) for connector in self.uncategorized],
            sections=[section.as_dict(localizer) for section in self.sections],
        )

    def check_connector_is_available(self, conn_type: ConnectionType) -> bool:
        conn_options = next((conn for conn in self._iter_connectors() if conn.conn_type == conn_type), None)
        if conn_options is None:
            LOGGER.warning("Connector %s is not available in current env", conn_type.name)
            return False

        return conn_options.availability == ConnectorAvailability.free

lib/dl_api_lib/dl_api_lib/connector_availability/base.py
datalens-tech/datalens-backend
from typing import Type

from dl_api_connector.api_schema.source import (
    register_source_api_schema,
    register_source_template_api_schema,
)
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)
from dl_api_lib.connection_forms.registry import register_connection_form_factory_cls
from dl_api_lib.connection_info import register_connector_info_provider_cls
from dl_api_lib.connector_alias import register_connector_alias
from dl_api_lib.dashsql.registry import register_dash_sql_param_literalizer_cls
from dl_api_lib.dataset.dialect import register_dialect_name
from dl_api_lib.i18n.registry import register_translation_configs
from dl_api_lib.query.registry import (
    register_compeng_dialect,
    register_filter_formula_compiler_cls,
    register_is_compeng_executable,
    register_is_forkable_source,
    register_multi_query_mutator_factory_cls,
)
from dl_api_lib.schemas.connection import register_sub_schema_class


class ApiConnectorRegistrator:
    @classmethod
    def register_source_definition(cls, source_def: Type[ApiSourceDefinition]) -> None:
        register_source_api_schema(
            source_type=source_def.core_source_def_cls.source_type,
            schema_cls=source_def.api_schema_cls,
        )
        register_source_template_api_schema(
            source_type=source_def.core_source_def_cls.source_type,
            schema_cls=source_def.template_api_schema_cls,
        )

    @classmethod
    def register_connection_definition(cls, conn_def: Type[ApiConnectionDefinition]) -> None:
        conn_type = conn_def.core_conn_def_cls.conn_type
        register_connector_info_provider_cls(conn_type=conn_type, info_provider_cls=conn_def.info_provider_cls)
        register_connector_alias(conn_type=conn_type, alias=conn_def.alias)
        register_sub_schema_class(conn_type=conn_type, schema_cls=conn_def.api_generic_schema_cls)
        if conn_def.form_factory_cls is not None:
            register_connection_form_factory_cls(conn_type=conn_type, factory_cls=conn_def.form_factory_cls)

    @classmethod
    def register_connector(cls, connector: Type[ApiConnector]) -> None:
        # backend_type-related stuff - TODO: Move to a separate entity
        backend_type = connector.core_connector_cls.backend_type
        register_dialect_name(backend_type=backend_type, dialect_name=connector.formula_dialect_name)
        for mqm_setting_item in connector.multi_query_mutation_factories:
            register_multi_query_mutator_factory_cls(
                query_proc_mode=mqm_setting_item.query_proc_mode,
                backend_type=backend_type,
                dialects=mqm_setting_item.dialects,
                factory_cls=mqm_setting_item.factory_cls,
            )
        register_is_forkable_source(backend_type=backend_type, is_forkable=connector.is_forkable)
        register_is_compeng_executable(backend_type=backend_type, is_compeng_executable=connector.is_compeng_executable)
        register_filter_formula_compiler_cls(
            backend_type=backend_type,
            filter_compiler_cls=connector.filter_formula_compiler_cls,
        )
        register_dash_sql_param_literalizer_cls(
            backend_type=backend_type,
            literalizer_cls=connector.dashsql_literalizer_cls,
        )

        # everything else
        for source_def in connector.source_definitions:
            cls.register_source_definition(source_def=source_def)
        for conn_def in connector.connection_definitions:
            cls.register_connection_definition(conn_def=conn_def)
        register_translation_configs(connector.translation_configs)
        if connector.compeng_dialect is not None:
            register_compeng_dialect(connector.compeng_dialect)


CONN_REG_BI_API = ApiConnectorRegistrator

lib/dl_api_lib/dl_api_lib/connector_registrator.py
datalens-tech/datalens-backend
DEFAULT_DATASET_LOCK_WAIT_TIMEOUT = 30  # seconds

lib/dl_api_lib/dl_api_lib/const.py
datalens-tech/datalens-backend
from typing import Type

from dl_api_connector.dashsql import DashSQLParamLiteralizer
from dl_constants.enums import SourceBackendType


_LITERALIZER_CLASSES: dict[SourceBackendType, Type[DashSQLParamLiteralizer]] = {}


def get_dash_sql_param_literalizer(backend_type: SourceBackendType) -> DashSQLParamLiteralizer:
    literalizer_cls = _LITERALIZER_CLASSES[backend_type]
    return literalizer_cls()


def register_dash_sql_param_literalizer_cls(
    backend_type: SourceBackendType,
    literalizer_cls: Type[DashSQLParamLiteralizer],
) -> None:
    if (registered_literalizer_cls := _LITERALIZER_CLASSES.get(backend_type)) is not None:
        assert registered_literalizer_cls is literalizer_cls
    else:
        _LITERALIZER_CLASSES[backend_type] = literalizer_cls

lib/dl_api_lib/dl_api_lib/dashsql/registry.py
datalens-tech/datalens-backend
from __future__ import annotations

from itertools import chain
import logging
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Optional,
    Sequence,
)

from dl_api_lib.dataset.dialect import resolve_dialect_name
from dl_api_lib.query.formalization.query_formalizer import SimpleQuerySpecFormalizer
from dl_api_lib.query.formalization.query_formalizer_base import QuerySpecFormalizerBase
from dl_api_lib.query.registry import (
    get_compeng_dialect,
    is_compeng_enabled,
)
from dl_api_lib.service_registry.service_registry import ApiServiceRegistry
from dl_constants.enums import (
    DataSourceRole,
    SourceBackendType,
)
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.dependencies.factory import ComponentDependencyManagerFactory
from dl_core.components.dependencies.factory_base import ComponentDependencyManagerFactoryBase
from dl_core.components.editor import DatasetComponentEditor
from dl_core.components.ids import (
    AvatarId,
    FieldId,
    FieldIdValidator,
)
from dl_core.data_source.base import DataSource
from dl_core.data_source.collection import (
    DataSourceCollectionBase,
    DataSourceCollectionFactory,
)
from dl_core.dataset_capabilities import DatasetCapabilities
from dl_core.exc import (
    FieldNotFound,
    ReferencedUSEntryNotFound,
)
from dl_core.us_dataset import Dataset
from dl_core.us_manager.us_manager import USManagerBase
from dl_core.utils import shorten_uuid
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.core.dialect import from_name_and_version
import dl_formula.core.exc as formula_exc
from dl_formula.definitions.scope import Scope
from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.parser.base import FormulaParser
from dl_query_processing.column_registry import ColumnRegistry
from dl_query_processing.compilation.base import RawQueryCompilerBase
from dl_query_processing.compilation.formula_compiler import FormulaCompiler
from dl_query_processing.compilation.primitives import (
    CompiledMultiQuery,
    CompiledMultiQueryBase,
    CompiledQuery,
)
from dl_query_processing.compilation.query_compiler import DefaultQueryCompiler
from dl_query_processing.compilation.query_mutator import (
    ExtendedAggregationQueryMutator,
    OptimizingQueryMutator,
    QueryMutator,
)
from dl_query_processing.compilation.specs import QuerySpec
from dl_query_processing.enums import (
    QueryType,
    SelectValueType,
)
import dl_query_processing.exc
from dl_query_processing.legend.block_legend import BlockSpec
from dl_query_processing.legend.field_legend import Legend
from dl_query_processing.multi_query.mutators.base import MultiQueryMutatorBase
from dl_query_processing.translation.multi_level_translator import MultiLevelQueryTranslator
from dl_query_processing.translation.primitives import TranslatedMultiQueryBase


if TYPE_CHECKING:
    from dl_core.components.ids import FieldIdGenerator
    from dl_core.db.elements import SchemaColumn
    from dl_core.fields import BIField
    from dl_query_processing.compilation.filter_compiler import FilterFormulaCompiler


LOGGER = logging.getLogger(__name__)


class AvatarAliasMapper:
    def __init__(self) -> None:
        self.next_idx = 1
        self.avatar_map: dict[AvatarId, str] = {}

    def __call__(self, avatar_id: AvatarId) -> str:
        try:
            shorten_uuid(avatar_id)
        except ValueError:
            # FIXME: the subquery aliases, such as `qq_i0_0_0`, aren't mapped consistently;
            # so have to discern those somehow, and skip them here.
            # For testing, just try
            # return 's' + avatar_id
            return avatar_id

        result = self.avatar_map.get(avatar_id)
        if result is not None:
            return result
        avatar_idx = self.next_idx
        self.next_idx += 1
        result = f"t{avatar_idx}"
        LOGGER.debug(f"AvatarAliasMapper: mapping {avatar_id!r} -> {result!r}")
        self.avatar_map[avatar_id] = result
        return result


INTERNALLY_APPLIED_WRAPPERS = {
    SelectValueType.array_prefix,
}


class DatasetBaseWrapper:
    _validation_mode: ClassVar[bool] = False
    _verbose_logging: ClassVar[bool] = False

    def __init__(
        self,
        ds: Dataset,
        *,
        us_manager: USManagerBase,
        block_spec: Optional[BlockSpec] = None,
        function_scopes: int = Scope.EXPLICIT_USAGE,
        debug_mode: bool = False,
    ):
        self._ds = ds
        self._us_manager = us_manager
        service_registry = self._us_manager.get_services_registry()
        assert isinstance(service_registry, ApiServiceRegistry)
        self._service_registry: ApiServiceRegistry = service_registry
        self._ds_accessor = DatasetComponentAccessor(dataset=ds)
        self._ds_editor = DatasetComponentEditor(dataset=ds)
        self._dsrc_coll_factory = DataSourceCollectionFactory(us_entry_buffer=self._us_manager.get_entry_buffer())
        self._function_scopes = function_scopes
        self._capabilities = DatasetCapabilities(
            dataset=self._ds,
            dsrc_coll_factory=self._dsrc_coll_factory,
        )
        self._ds_dep_mgr_factory = self.make_dep_mgr_factory()
        if block_spec is None:
            block_spec = BlockSpec(block_id=0, legend=Legend(items=[]), parent_block_id=None)
        assert block_spec is not None
        self._formalizer = self.make_formalizer(query_type=block_spec.query_type)
        self._avatar_alias_mapper = AvatarAliasMapper()

        self._has_sources = False
        self.dialect: DialectCombo = D.DUMMY

        # field-dependent stuff (helper mappings)
        self.inspect_env: Optional[InspectionEnvironment] = None
        self._column_reg: Optional[ColumnRegistry] = None
        self._formula_compiler: Optional[FormulaCompiler] = None
        self._query_spec: Optional[QuerySpec] = None

        self._reload_sources()
        self._reload_formalized_specs(block_spec=block_spec)
        self._id_generator: FieldIdGenerator = self.make_id_generator()
        self._id_validator = FieldIdValidator()

        self._debug_mode = debug_mode

    @property
    def query_spec(self) -> QuerySpec:
        assert self._query_spec is not None
        return self._query_spec

    @property
    def formula_compiler(self) -> FormulaCompiler:
        assert self._formula_compiler is not None
        return self._formula_compiler

    def make_dep_mgr_factory(self) -> ComponentDependencyManagerFactoryBase:
        return ComponentDependencyManagerFactory(dataset=self._ds)

    def make_formalizer(self, query_type: QueryType) -> QuerySpecFormalizerBase:
        return SimpleQuerySpecFormalizer(
            dataset=self._ds,
            us_entry_buffer=self._us_manager.get_entry_buffer(),
            verbose_logging=self._verbose_logging,
        )

    def _get_data_source_coll_strict(self, source_id: str) -> DataSourceCollectionBase:
        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
        dsrc_coll = self._dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
        return dsrc_coll

    def _get_data_source_strict(self, source_id: str, role: Optional[DataSourceRole] = None) -> DataSource:
        if role is None:
            role = self.resolve_role()
        assert role is not None
        dsrc_coll = self._get_data_source_coll_strict(source_id=source_id)
        dsrc = dsrc_coll.get_strict(role=role)
        return dsrc

    def make_id_generator(self) -> FieldIdGenerator:
        id_generator_factory = self._service_registry.get_field_id_generator_factory()
        return id_generator_factory.get_field_id_generator(ds=self._ds)

    def make_formula_parser(self) -> FormulaParser:
        parser_factory = self._service_registry.get_formula_parser_factory()
        return parser_factory.get_formula_parser()

    def make_filter_compiler(self) -> FilterFormulaCompiler:
        raise NotImplementedError

    def make_query_compiler(self) -> RawQueryCompilerBase:
        assert self._column_reg is not None
        return DefaultQueryCompiler(
            dataset=self._ds,
            column_reg=self._column_reg,
            formula_compiler=self.formula_compiler,
            filter_compiler=self.make_filter_compiler(),
        )

    def make_query_mutators(self) -> Sequence[QueryMutator]:
        return [
            OptimizingQueryMutator(
                dialect=self.dialect,
                disable_optimizations=self._validation_mode,
            ),
            ExtendedAggregationQueryMutator(
                allow_empty_dimensions_for_forks=self._validation_mode,
                allow_arbitrary_toplevel_lod_dimensions=self._validation_mode,
                new_subquery_mode=True,
            ),
        ]

    def make_multi_query_mutators(self) -> Sequence[MultiQueryMutatorBase]:
        backend_type = self.get_backend_type()
        mqm_factory_factory = self._service_registry.get_multi_query_mutator_factory_factory()
        return mqm_factory_factory.get_multi_query_mutators(
            backend_type=backend_type,
            dataset=self._ds,
            dialect=self.dialect,
        )

    def make_multi_query_translator(self) -> MultiLevelQueryTranslator:
        assert self.inspect_env is not None
        assert self._column_reg is not None
        return MultiLevelQueryTranslator(
            inspect_env=self.inspect_env,
            function_scopes=self._function_scopes,
            verbose_logging=self._verbose_logging,
            source_db_columns=self._column_reg,
            avatar_alias_mapper=self._avatar_alias_mapper,
            collect_stats=not self._validation_mode,
            dialect=self.dialect,
            compeng_dialect=get_compeng_dialect() if is_compeng_enabled() else None,
        )

    def resolve_role(self) -> DataSourceRole:
        return self._capabilities.resolve_source_role()

    def get_backend_type(self) -> SourceBackendType:
        return self._capabilities.get_backend_type(role=self.resolve_role())

    def _reload_sources(self) -> None:
        self._has_sources = bool(self._ds.get_single_data_source_id())
        # resolve database characteristics
        # never go to database from here --> only_cache=True
        if self._has_sources:
            role = self.resolve_role()
            try:
                backend_type = self._capabilities.get_backend_type(role=role)
                dialect_name = resolve_dialect_name(backend_type=backend_type)
                source_id = self._ds.get_single_data_source_id()
                dsrc = self._get_data_source_strict(source_id=source_id, role=role)
                db_info = dsrc.get_cached_db_info()
                db_version = db_info.version
                self.dialect = from_name_and_version(
                    dialect_name=dialect_name,
                    dialect_version=db_version,
                )
            except ReferencedUSEntryNotFound:
                self.dialect = D.DUMMY

        if self._query_spec is not None:
            self.load_exbuilders()

    def _reload_formalized_specs(self, block_spec: Optional[BlockSpec] = None) -> None:
        assert block_spec is not None, "block_spec must not be None in this implementation"
        self._query_spec = self._formalizer.make_query_spec(block_spec=block_spec)
        if self._formula_compiler is None:
            self.load_exbuilders()
        if self._formula_compiler is not None:
            self._formula_compiler.update_environments(
                group_by_ids=[spec.field_id for spec in self.query_spec.group_by_specs],
                order_by_specs=self.query_spec.order_by_specs,
            )

    def allow_nested_window_functions(self) -> bool:
        return True

    def load_exbuilders(self):  # type: ignore  # TODO: fix
        self.inspect_env = InspectionEnvironment()
        self._column_reg = ColumnRegistry(
            db_columns=self._generate_raw_column_list(),
            avatar_source_map={avatar.id: avatar.source_id for avatar in self._ds_accessor.get_avatar_list()},
        )

        self._formula_compiler = FormulaCompiler(
            columns=self._column_reg,
            inspect_env=self.inspect_env,
            formula_parser=self.make_formula_parser(),
            all_fields=list(self._ds.result_schema),
            group_by_ids=[spec.field_id for spec in self.query_spec.group_by_specs],
            filter_ids=[spec.field_id for spec in self.query_spec.filter_specs],
            field_wrappers={
                spec.field_id: spec.wrapper
                for spec in self.query_spec.select_specs
                if spec.wrapper.type in INTERNALLY_APPLIED_WRAPPERS
            },
            order_by_specs=self.query_spec.order_by_specs,
            mock_among_dimensions=self._validation_mode,
            allow_nested_window_functions=self.allow_nested_window_functions(),
            parameter_value_specs=self.query_spec.parameter_value_specs,
            validate_aggregations=self.query_spec.meta.query_type != QueryType.totals,
        )

    def _generate_raw_column_list(self) -> list[SchemaColumn]:
        """Generate info about columns of the data table"""

        role = self.resolve_role()
        return list(
            chain.from_iterable(
                (self._get_data_source_coll_strict(source_id=source_id).get_cached_raw_schema(role=role) or ())
                for source_id in self._ds_accessor.get_data_source_id_list()
            )
        )

    def get_field_by_id(self, field_id: FieldId) -> BIField:
        try:
            return self._ds.result_schema.by_guid(field_id)
        except KeyError:
            raise FieldNotFound(f"Field {field_id} not found in dataset")

    def process_compiled_query(self, compiled_query: CompiledQuery) -> CompiledMultiQueryBase:
        try:
            # Apply whole-query mutations
            for mutator in self.make_query_mutators():
                compiled_query = mutator.mutate_query(compiled_query=compiled_query)

            # Transform flat query into multi-query object
            compiled_multi_query: CompiledMultiQueryBase = CompiledMultiQuery(queries=[compiled_query])

            # Apply multi-query mutations (extended aggregations, lookups and window functions)
            multi_query_mutators = self.make_multi_query_mutators()
            for multi_mutator in multi_query_mutators:
                compiled_multi_query = multi_mutator.mutate_multi_query(compiled_multi_query)

        except formula_exc.FormulaError as err:
            raise dl_query_processing.exc.FormulaHandlingError(*err.errors)

        return compiled_multi_query

    def compile_and_translate_query(self, query_spec: QuerySpec) -> TranslatedMultiQueryBase:
        try:
            # Compile query from spec
            query_compiler = self.make_query_compiler()
            compiled_query = query_compiler.make_compiled_query(query_spec=query_spec)

            # Plan, slice and separate query
            compiled_multi_query = self.process_compiled_query(compiled_query=compiled_query)

            # Translate query.
            query_translator = self.make_multi_query_translator()
            translated_multi_query = query_translator.translate_multi_query(compiled_multi_query=compiled_multi_query)
        except formula_exc.FormulaError as err:
            raise dl_query_processing.exc.FormulaHandlingError(*err.errors)

        return translated_multi_query

lib/dl_api_lib/dl_api_lib/dataset/base_wrapper.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    ClassVar,
    Generator,
    NamedTuple,
    Optional,
    Union,
)

import attr

from dl_constants.enums import (
    ComponentType,
    ManagedBy,
)
from dl_core.base_models import ObligatoryFilter
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.editor import DatasetComponentEditor
from dl_core.data_source import DataSourceCollectionBase
from dl_core.data_source.collection import DataSourceCollectionFactory
import dl_core.exc as common_exc
from dl_core.fields import (
    BIField,
    ResultSchema,
)
from dl_core.multisource import (
    AvatarRelation,
    SourceAvatar,
)
from dl_core.us_dataset import Dataset
from dl_core.us_manager.local_cache import USEntryBuffer


DatasetComponent = Union[
    DataSourceCollectionBase,
    SourceAvatar,
    AvatarRelation,
    BIField,
    ObligatoryFilter,
    ResultSchema,
]


class DatasetComponentRef(NamedTuple):
    component_type: ComponentType
    component_id: str


@attr.s(frozen=True)
class DatasetComponentAbstraction:
    _dataset: Dataset = attr.ib(kw_only=True)
    _us_entry_buffer: USEntryBuffer = attr.ib(kw_only=True)
    _ds_accessor: DatasetComponentAccessor = attr.ib(init=False)
    _ds_editor: DatasetComponentEditor = attr.ib(init=False)
    _dsrc_coll_factory: DataSourceCollectionFactory = attr.ib(init=False)

    _can_check_valid: ClassVar[frozenset[ComponentType]] = frozenset(
        {
            ComponentType.data_source,
            ComponentType.source_avatar,
            ComponentType.avatar_relation,
            ComponentType.field,
            ComponentType.obligatory_filter,
        }
    )

    @_ds_accessor.default
    def _make_ds_accessor(self) -> DatasetComponentAccessor:
        return DatasetComponentAccessor(dataset=self._dataset)

    @_ds_editor.default
    def _make_ds_editor(self) -> DatasetComponentEditor:
        return DatasetComponentEditor(dataset=self._dataset)

    @_dsrc_coll_factory.default
    def _make_dsrc_coll_factory(self) -> DataSourceCollectionFactory:
        return DataSourceCollectionFactory(us_entry_buffer=self._us_entry_buffer)

    @property
    def _ds(self) -> Dataset:  # FIXME
        return self._dataset

    def _get_data_source_coll_strict(self, source_id: str) -> DataSourceCollectionBase:
        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
        dsrc_coll = self._dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
        return dsrc_coll

    def _get_data_source_coll_opt(self, source_id: str) -> Optional[DataSourceCollectionBase]:
        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_opt(source_id=source_id)
        if not dsrc_coll_spec:
            return None
        dsrc_coll = self._dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
        return dsrc_coll

    def get_component(self, component_ref: DatasetComponentRef) -> Optional[DatasetComponent]:
        """Get dataset component by its ID and type"""
        if component_ref.component_type == ComponentType.data_source:
            return self._get_data_source_coll_opt(source_id=component_ref.component_id)
        if component_ref.component_type == ComponentType.source_avatar:
            return self._ds_accessor.get_avatar_opt(avatar_id=component_ref.component_id)
        if component_ref.component_type == ComponentType.avatar_relation:
            return self._ds_accessor.get_avatar_relation_opt(relation_id=component_ref.component_id)
        if component_ref.component_type == ComponentType.field:
            for field in self._ds.result_schema:
                if field.guid == component_ref.component_id:
                    return field
            return None
        if component_ref.component_type == ComponentType.obligatory_filter:
            return self._ds_accessor.get_obligatory_filter_opt(obfilter_id=component_ref.component_id)
        if component_ref.component_type == ComponentType.result_schema:
            if self._ds.result_schema.id == component_ref.component_id:
                return self._ds.result_schema
            return None

        raise ValueError(f"Unsupported component_type {component_ref.component_type.name}")

    def update_component_validity(self, component_ref: DatasetComponentRef, valid: bool) -> None:
        """
        Update ``valid`` attribute of given dataset component.
        If component doesn't exist, return peacefully.
        """
        component = self.get_component(component_ref=component_ref)
        if (
            component_ref.component_type not in self._can_check_valid or component is None
        ):  # or valid == component.valid:
            return

        if valid != component.valid:
            if component_ref.component_type == ComponentType.data_source:
                self._ds_editor.update_data_source_collection(source_id=component_ref.component_id, valid=valid)
            elif component_ref.component_type == ComponentType.obligatory_filter:
                self._ds_editor.update_obligatory_filter(obfilter_id=component_ref.component_id, valid=valid)
            elif component_ref.component_type == ComponentType.source_avatar:
                self._ds_editor.update_avatar(avatar_id=component_ref.component_id, valid=valid)
            elif component_ref.component_type == ComponentType.avatar_relation:
                self._ds_editor.update_avatar_relation(relation_id=component_ref.component_id, valid=valid)
            elif component_ref.component_type == ComponentType.field:
                for field_idx, field in enumerate(self._ds.result_schema):
                    if field.guid == component_ref.component_id:
                        self._ds.result_schema.update_field(idx=field_idx, field=field.clone(valid=valid))
            else:
                raise ValueError(f"Unsupported component_type {component_ref.component_type.name}")

    def iter_dataset_components_by_type(self, component_type: ComponentType) -> Generator[DatasetComponent, None, None]:
        """Iterate over dataset components of given type"""
        if component_type == ComponentType.data_source:
            for other_component_id in self._ds_accessor.get_data_source_id_list():
                yield self._get_data_source_coll_strict(source_id=other_component_id)
        elif component_type == ComponentType.source_avatar:
            yield from self._ds_accessor.get_avatar_list()
        elif component_type == ComponentType.avatar_relation:
            yield from self._ds_accessor.get_avatar_relation_list()
        elif component_type == ComponentType.obligatory_filter:
            yield from self._ds_accessor.get_obligatory_filter_list()
        elif component_type == ComponentType.field:
            yield from self._ds.result_schema
        elif component_type == ComponentType.result_schema:
            yield self._ds.result_schema
        else:
            raise ValueError(f"Unsupported component_type {component_type.name}")

    def validate_component_can_be_managed(self, component_ref: DatasetComponentRef, by: Optional[ManagedBy]) -> None:
        if by is not None:
            component = self.get_component(component_ref=component_ref)
            if component.managed_by != by:  # type: ignore  # TODO: fix
                raise common_exc.DatasetConfigurationError(
                    f"Component {component_ref.component_type.name} {component_ref.component_id} "
                    f"cannot be managed by {by.name}"
                )

lib/dl_api_lib/dl_api_lib/dataset/component_abstraction.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Dict

from dl_constants.enums import SourceBackendType
from dl_formula.core.dialect import DialectName


_DIALECT_NAMES_FROM_SA: Dict[SourceBackendType, DialectName] = {}


def register_dialect_name(backend_type: SourceBackendType, dialect_name: DialectName) -> None:
    if backend_type in _DIALECT_NAMES_FROM_SA:
        assert _DIALECT_NAMES_FROM_SA[backend_type] == dialect_name
    _DIALECT_NAMES_FROM_SA[backend_type] = dialect_name


def resolve_dialect_name(backend_type: SourceBackendType) -> DialectName:
    return _DIALECT_NAMES_FROM_SA.get(backend_type, DialectName.DUMMY)

lib/dl_api_lib/dl_api_lib/dataset/dialect.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
import logging
from typing import (
    Generator,
    Iterable,
    List,
    Optional,
    Tuple,
)

from dl_api_lib import utils as bi_utils
from dl_api_lib.enums import USPermissionKind
from dl_constants.enums import (
    ConnectionType,
    DataSourceRole,
)
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.editor import DatasetComponentEditor
from dl_core.data_source.collection import (
    DataSourceCollectionBase,
    DataSourceCollectionFactory,
)
from dl_core.us_dataset import Dataset
from dl_core.us_manager.local_cache import USEntryBuffer
from dl_core.us_manager.us_manager import USManagerBase


LOGGER = logging.getLogger(__name__)


def _iter_data_source_collections(
    dataset: Dataset,
    us_entry_buffer: USEntryBuffer,
    source_ids: Optional[Iterable[str]] = None,
) -> Generator[DataSourceCollectionBase, None, None]:
    ds_accessor = DatasetComponentAccessor(dataset=dataset)

    if source_ids is None:
        source_ids = ds_accessor.get_data_source_id_list()
    assert source_ids is not None

    dsrc_coll_factory = DataSourceCollectionFactory(us_entry_buffer=us_entry_buffer)
    for source_id in source_ids:
        dsrc_coll_spec = ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
        dsrc_coll = dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
        yield dsrc_coll


def check_permissions_for_origin_sources(
    dataset: Dataset,
    source_ids: Iterable[str],
    permission_kind: USPermissionKind,
    us_entry_buffer: USEntryBuffer,
) -> None:
    """Check whether data source has read rights or not."""
    for dsrc_coll in _iter_data_source_collections(
        dataset=dataset,
        us_entry_buffer=us_entry_buffer,
        source_ids=source_ids,
    ):
        data_source = dsrc_coll.get_opt(role=DataSourceRole.origin)
        if data_source is not None:
            bi_utils.need_permission_on_entry(data_source.connection, permission_kind)


def log_dataset_field_stats(dataset: Dataset) -> None:
    stats = defaultdict(int)
    stats["total"] = len(dataset.result_schema)
    for field in dataset.result_schema:
        # We expect no name collisions here
        stats[field.calc_mode.name] += 1
        stats[field.type.name.lower()] += 1
    LOGGER.info("Dataset field stats", extra=dict(dataset_field_stats=dict(stats)))


def allow_rls_for_dataset(dataset: Dataset) -> bool:
    return (
        dataset.permissions is None
        or dataset.permissions[  # dataset is not "real" (saved to US), used for validation
            "edit"
        ]  # dataset is US-bound and should have the required permissions
    )


def get_dataset_conn_types(
    dataset: Dataset,
    us_entry_buffer: USEntryBuffer,
    roles: Tuple[DataSourceRole] = (DataSourceRole.origin,),
) -> List[ConnectionType]:
    colls = list(_iter_data_source_collections(dataset=dataset, us_entry_buffer=us_entry_buffer))
    if not colls:
        return []
    result = []
    for dsrc_coll in colls:
        for role in roles:
            dsrc = dsrc_coll.get_opt(role)
            if dsrc is None:
                continue
            assert dsrc.conn_type is not None
            result.append(dsrc.conn_type)
    return result


def invalidate_sample_sources(dataset: Dataset, source_ids: List[str], us_manager: USManagerBase) -> None:
    """
    Remove sample data sources for given source IDs.
    """
    ds_editor = DatasetComponentEditor(dataset=dataset)
    for dsrc_coll in _iter_data_source_collections(
        dataset=dataset,
        us_entry_buffer=us_manager.get_entry_buffer(),
        source_ids=source_ids,
    ):
        source_id = dsrc_coll.id
        if dsrc_coll.exists(DataSourceRole.sample):
            ds_editor.remove_data_source(source_id=source_id, role=DataSourceRole.sample)

lib/dl_api_lib/dl_api_lib/dataset/utils.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import contextmanager
import logging
from typing import (
    TYPE_CHECKING,
    Collection,
    Generator,
    Optional,
    Sequence,
)
import uuid

import attr

from dl_api_lib import (
    exc,
    utils,
)
from dl_api_lib.dataset.base_wrapper import DatasetBaseWrapper
from dl_api_lib.dataset.component_abstraction import (
    DatasetComponentAbstraction,
    DatasetComponentRef,
)
from dl_api_lib.dataset.utils import check_permissions_for_origin_sources
from dl_api_lib.enums import (
    FILTERS_BY_TYPE,
    DatasetAction,
    USPermissionKind,
)
from dl_api_lib.query.formalization.block_formalizer import BlockFormalizer
from dl_api_lib.query.formalization.legend_formalizer import ValidationLegendFormalizer
from dl_api_lib.query.formalization.raw_specs import RawQuerySpecUnion
from dl_api_lib.request_model.data import (
    Action,
    AddField,
    AddUpdateObligatoryFilter,
    AvatarActionBase,
    DeleteField,
    DeleteObligatoryFilter,
    FieldAction,
    ObligatoryFilterActionBase,
    ObligatoryFilterBase,
    RelationActionBase,
    ReplaceConnection,
    ReplaceConnectionAction,
    SourceActionBase,
    UpdateField,
)
from dl_app_tools.profiling_base import generic_profiler
from dl_constants.enums import (
    AggregationFunction,
    BinaryJoinOperator,
    CalcMode,
    ComponentErrorLevel,
    ComponentType,
    ConnectionType,
    DataSourceRole,
    DataSourceType,
    ManagedBy,
    TopLevelComponentId,
)
from dl_core.base_models import (
    DefaultConnectionRef,
    DefaultWhereClause,
)
from dl_core.connectors.base.data_source_migration import get_data_source_migrator
from dl_core.constants import DatasetConstraints
from dl_core.data_source.base import DataSource
from dl_core.data_source.collection import DataSourceCollectionBase
from dl_core.data_source.sql import BaseSQLDataSource
from dl_core.data_source.type_mapping import get_data_source_class
from dl_core.db import (
    SchemaColumn,
    SchemaInfo,
)
import dl_core.exc as common_exc
from dl_core.fields import (
    BIField,
    DirectCalculationSpec,
    FormulaCalculationSpec,
    create_calc_spec_from,
    del_calc_spec_kwargs_from,
)
from dl_core.multisource import (
    AvatarRelation,
    BinaryCondition,
    ConditionPartDirect,
    SourceAvatar,
)
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_dataset import Dataset
from dl_core.us_manager.us_manager import USManagerBase
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_formula.core import exc as formula_exc
from dl_formula.core.message_ctx import (
    FormulaErrorCtx,
    MessageLevel,
)
from dl_query_processing.compilation.helpers import single_formula_comp_query_for_validation
import dl_query_processing.exc
from dl_query_processing.legend.block_legend import BlockSpec
from dl_utils.utils import enum_not_none


if TYPE_CHECKING:
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class ErrorInfo:
    code: Sequence[str] = attr.ib(kw_only=True)
    message: str = attr.ib(kw_only=True)
    details: dict = attr.ib(kw_only=True, factory=dict)
    level: ComponentErrorLevel = attr.ib(kw_only=True, default=ComponentErrorLevel.error)


@attr.s(frozen=True)
class FormulaErrorInfo(ErrorInfo):  # noqa
    row: Optional[int] = attr.ib(kw_only=True)
    column: Optional[int] = attr.ib(kw_only=True)
    token: Optional[str] = attr.ib(kw_only=True)


_FORMULA_TO_CORE_ERR_LEVEL_MAP = {
    MessageLevel.ERROR: ComponentErrorLevel.error,
    MessageLevel.WARNING: ComponentErrorLevel.warning,
}


def comp_err_level_from_formula_err_level(formula_level: MessageLevel) -> ComponentErrorLevel:
    return _FORMULA_TO_CORE_ERR_LEVEL_MAP[formula_level]


def field_not_none(field: Optional[BIField]) -> BIField:
    assert field is not None
    return field


class DatasetValidator(DatasetBaseWrapper):
    """
    An extension of DatasetView that allows "smooth" updates of the DS configuration
    and provides validation methods.
    """

    # TODO: refactor the class's structure (separate into several classes: one per action)
    _validation_mode = True
    _is_bleeding_edge_user: bool = False

    def __init__(
        self,
        ds: Dataset,
        us_manager: USManagerBase,
        is_data_api: bool = False,
        is_bleeding_edge_user: bool = False,
    ):
        super().__init__(ds=ds, us_manager=us_manager)
        self._is_data_api = is_data_api
        self._is_bleeding_edge_user = is_bleeding_edge_user

        self._ds_ca = DatasetComponentAbstraction(dataset=self._ds, us_entry_buffer=self._us_manager.get_entry_buffer())
        self._remapped_source_ids: dict[str, str] = {}
        self._affected_components: set[DatasetComponentRef] = set()

    def _reload_formalized_specs(self, block_spec: Optional[BlockSpec] = None) -> None:
        # Ignore incoming spec.
        raw_query_spec_union = RawQuerySpecUnion()
        legend = ValidationLegendFormalizer(dataset=self._ds).make_legend(raw_query_spec_union=raw_query_spec_union)
        block_legend = BlockFormalizer(dataset=self._ds).make_block_legend(
            raw_query_spec_union=raw_query_spec_union, legend=legend
        )
        block_spec = block_legend.blocks[0]
        super()._reload_formalized_specs(block_spec=block_spec)

    @property
    def _sync_us_manager(self) -> SyncUSManager:
        us_manager = self._us_manager
        if isinstance(us_manager, SyncUSManager):
            return us_manager
        raise TypeError(
            f"Unsupported type of USManager ({type(us_manager).__qualname__}). " f"Only SyncUSManager supported."
        )

    @generic_profiler("validator-apply-action-batch")
    def apply_batch(self, action_batch: Sequence[Action]) -> None:
        self.update_unpatched_fields()  # TODO: remove after all fields have been patched
        try:
            for action_patch in action_batch:
                self.apply_action(
                    item_data=action_patch,
                )
        except common_exc.DatasetConfigurationError as err:
            LOGGER.info("Error occurred during applying action: %r", err)
            raise exc.DLValidationFatal(f"Invalid action: {err}")

        self._check_for_orphaned_component_errors()

    def apply_action(self, item_data: Action, by: Optional[ManagedBy] = ManagedBy.user) -> None:
        """Apply update to the dataset configuration"""

        action = DatasetAction.remap_legacy(item_data.action)
        if self._is_data_api and action not in DatasetAction.get_actions_whitelist_for_data_api():
            raise exc.DatasetActionNotAllowedError(f"Action {action.value} is not allowed.")

        extra = dict(
            dataset_action_name=action.name,
            dataset_action_data=item_data,
        )
        LOGGER.info(f"Going to apply action {action.name}", extra=extra)

        # TODO: ManagedBy is not present in any schema. Why is this check needed here?
        if by is not None and (item_data.managed_by or ManagedBy.user) != by:
            raise exc.DLValidationFatal(f"Item cannot be managed by {by.name}")

        if isinstance(item_data, FieldAction):
            self.apply_field_action(
                action=action, field_data=item_data.field, order_index=item_data.order_index or 0, by=by
            )
        elif isinstance(item_data, SourceActionBase):
            self.apply_source_action(action=action, source_data=item_data.source, by=by)
        elif isinstance(item_data, AvatarActionBase):
            self.apply_source_avatar_action(
                action=action,
                source_avatar_data=item_data.source_avatar,
                by=by,
                disable_fields_update=item_data.disable_fields_update,
            )
        elif isinstance(item_data, RelationActionBase):
            self.apply_avatar_relation_action(action=action, avatar_relation_data=item_data.avatar_relation, by=by)
        elif isinstance(item_data, ReplaceConnectionAction):
            self.apply_connection_action(action=action, connection_data=item_data.connection, by=by)
        elif isinstance(item_data, ObligatoryFilterActionBase):
            self.apply_obligatory_filter_action(action=action, filter_data=item_data.obligatory_filter, by=by)

        self.update_validity_of_affected_components()

    def _check_for_orphaned_component_errors(self) -> None:
        missing_components: list[DatasetComponentRef] = []
        for err_pack in self._ds.error_registry.items:
            component_ref = DatasetComponentRef(component_type=err_pack.type, component_id=err_pack.id)
            component = self._ds_ca.get_component(component_ref=component_ref)
            if component is None:
                missing_components.append(component_ref)
        if missing_components:
            LOGGER.warning(
                "Nonexistent components are referenced in error registry",
                extra={
                    "dataset_id": self._ds.uuid,
                    "components": [
                        {"type": component_ref.component_type, "id": component_ref.component_id}
                        for component_ref in missing_components
                    ],
                },
            )

    def validate_result_schema_length(self, field_cnt_diff: int = 0) -> None:
        if self._is_data_api:
            return

        # only one limit related record in error registry + remove potentially invalid errors
        self._ds.error_registry.remove_errors(id=self._ds.result_schema.id, code=exc.TooManyFieldsError.err_code)

        new_result_schema_len = len(self._ds.result_schema) + field_cnt_diff
        if new_result_schema_len > DatasetConstraints.FIELD_COUNT_LIMIT_HARD and field_cnt_diff > 0:
            raise dl_query_processing.exc.DatasetTooManyFieldsFatal()
        if new_result_schema_len > DatasetConstraints.FIELD_COUNT_LIMIT_SOFT:
            self._ds.error_registry.add_error(
                id=self._ds.result_schema.id,
                type=ComponentType.result_schema,
                message="Too many fields in the result schema",
                code=exc.TooManyFieldsError.err_code,
            )

    def get_dependent_fields(self, field: Optional[BIField]) -> list[BIField]:
        """Return list of fields directly dependent on the given one"""

        if field is None:
            return []

        result = []
        for other_field in self._ds.result_schema:
            if other_field.depends_on(field):
                result.append(other_field)

        return result

    def _get_field_errors(self, field: BIField) -> list[FormulaErrorCtx]:
        errors = self.formula_compiler.get_field_errors(field=field)
        try:
            formula_info = self.formula_compiler.compile_field_formula(field=field, collect_errors=True)
            assert self._column_reg is not None
            formula_comp_query = single_formula_comp_query_for_validation(
                formula=formula_info, ds_accessor=self._ds_accessor, column_reg=self._column_reg
            )
            formula_comp_multi_query = self.process_compiled_query(compiled_query=formula_comp_query)
            multi_translator = self.make_multi_query_translator()
            errors += multi_translator.collect_errors(compiled_multi_query=formula_comp_multi_query)  # type: ignore  # TODO: fix
        except formula_exc.FormulaError:
            pass
        return errors  # type: ignore  # TODO: fix

    def _get_relation_errors(self, relation: AvatarRelation) -> list[FormulaErrorCtx]:
        # TODO: switch to ComponentError items
        errors: list[FormulaErrorCtx] = []

        # Validate expressions
        try:
            formula_info = self.formula_compiler.compile_relation_formula(relation=relation, collect_errors=True)
            assert self._column_reg is not None
            formula_comp_query = single_formula_comp_query_for_validation(
                formula=formula_info, ds_accessor=self._ds_accessor, column_reg=self._column_reg
            )
            formula_comp_multi_query = self.process_compiled_query(compiled_query=formula_comp_query)
            multi_translator = self.make_multi_query_translator()
            errors = multi_translator.collect_errors(compiled_multi_query=formula_comp_multi_query)
        except formula_exc.FormulaError as err:
            errors += err.errors

        # Validate join_type
        left_coll = self._get_data_source_coll_strict(
            source_id=self._ds_accessor.get_avatar_strict(avatar_id=relation.left_avatar_id).source_id
        )
        right_coll = self._get_data_source_coll_strict(
            source_id=self._ds_accessor.get_avatar_strict(avatar_id=relation.right_avatar_id).source_id
        )

        if not left_coll.supports_join_type(relation.join_type) or not right_coll.supports_join_type(
            relation.join_type
        ):
            errors.append(
                FormulaErrorCtx(
                    message=common_exc.AvatarRelationJoinTypeError.default_message,
                    code=common_exc.AvatarRelationJoinTypeError.err_code,  # type: ignore  # TODO: fix
                    level=MessageLevel.ERROR,
                )
            )

        return errors

    def _get_formula_errors(self, formula: str, feature_errors: bool = True) -> list[FormulaErrorCtx]:
        errors = self.formula_compiler.get_formula_errors(formula=formula)
        try:
            formula_info = self.formula_compiler.compile_text_formula(formula=formula, collect_errors=True)
            assert self._column_reg is not None
            formula_comp_query = single_formula_comp_query_for_validation(
                formula=formula_info, ds_accessor=self._ds_accessor, column_reg=self._column_reg
            )
            formula_comp_multi_query = self.process_compiled_query(compiled_query=formula_comp_query)
            multi_translator = self.make_multi_query_translator()
            errors += multi_translator.collect_errors(  # type: ignore  # TODO: fix
                compiled_multi_query=formula_comp_multi_query, feature_errors=feature_errors
            )
        except formula_exc.FormulaError:
            pass
        return errors

    def _automanage_calc_mode_attrs(self, field: BIField) -> BIField:
        updated_field = field
        if updated_field.calc_mode == CalcMode.direct:
            if updated_field.avatar_id is None:
                root_avatar = self._ds_accessor.get_root_avatar_strict()
                updated_field = updated_field.clone(avatar_id=root_avatar.id)
            col = self._get_raw_schema_column(field=updated_field)
            if col is not None:
                updated_field = updated_field.clone(
                    has_auto_aggregation=col.has_auto_aggregation,
                    lock_aggregation=col.lock_aggregation,
                )
            else:
                updated_field = updated_field.clone(
                    has_auto_aggregation=False,
                    lock_aggregation=False,
                )
        elif updated_field.calc_mode == CalcMode.formula:
            if updated_field.formula:
                updated_field = updated_field.clone(
                    guid_formula=BIField.rename_in_formula(
                        formula=updated_field.formula, key_map=self._ds.result_schema.titles_to_guids
                    )
                )
            elif updated_field.guid_formula:
                updated_field = updated_field.clone(
                    formula=BIField.rename_in_formula(
                        formula=updated_field.guid_formula, key_map=self._ds.result_schema.guids_to_titles
                    )
                )

        return updated_field

    def _get_autoupdated_field(self, field: BIField, explicitly_updated: bool = False) -> BIField:
        updated_field = field
        updated_field = self._automanage_calc_mode_attrs(updated_field)

        self.formula_compiler.update_field(updated_field)
        updated_field = updated_field.clone(
            data_type=self.formula_compiler.get_field_final_data_type(updated_field),
            initial_data_type=self.formula_compiler.get_field_initial_data_type(updated_field),
            has_auto_aggregation=self.formula_compiler.field_has_auto_aggregation(updated_field),
            type=self.formula_compiler.get_field_type(updated_field),
        )
        field_is_new = field.initial_data_type is None

        # Auto-manage cast
        if field_is_new and updated_field.cast is None:
            LOGGER.info(
                f"Automatically initializing cast in {field.title} "
                f"as {enum_not_none(updated_field.initial_data_type).name}"
            )
            updated_field = updated_field.clone(cast=updated_field.initial_data_type)
        elif not field_is_new and updated_field.initial_data_type != field.initial_data_type:
            # autofix data type if derived type has changed
            if updated_field.cast != updated_field.initial_data_type:
                LOGGER.info(
                    f"Automatically setting cast in {field.title} "
                    f"to {enum_not_none(updated_field.initial_data_type).name}"
                )
                updated_field = updated_field.clone(cast=updated_field.initial_data_type)
                # the change of cast might result in new data_type
                updated_field = self._get_autoupdated_field(updated_field)

        # Auto-manage has_auto_aggregation
        if (
            explicitly_updated
            and updated_field.has_auto_aggregation
            and updated_field.aggregation != AggregationFunction.none
        ):
            # Field is autoaggregated and also has an explicit aggregation.
            # We can drop the additional aggregation if this is the field
            # being updated by the user (`explicitly_updated`) explicitly.
            updated_field = updated_field.clone(
                aggregation=AggregationFunction.none,
            )

        return updated_field

    def _update_dependent_fields(
        self,
        old_field: Optional[BIField],
        new_field: Optional[BIField],
        visited_guids: frozenset[str],
    ) -> None:
        """Automatically update all fields dependent on the given one"""
        dep_field_ids = set()
        dep_fields = []
        for dep_field in self.get_dependent_fields(old_field) + self.get_dependent_fields(new_field):
            if dep_field.guid not in dep_field_ids:
                dep_fields.append(dep_field)
                dep_field_ids.add(dep_field.guid)

        # Handle renaming and guid changing
        if new_field is not None and old_field is not None:
            if new_field.title != old_field.title:
                for i, dep_field in enumerate(dep_fields):
                    assert isinstance(dep_field.calc_spec, FormulaCalculationSpec)
                    dep_fields[i] = dep_field.clone(
                        formula=BIField.rename_in_formula(dep_field.formula, key_map={old_field.title: new_field.title})
                    )
            if new_field.guid != old_field.guid:
                for i, dep_field in enumerate(dep_fields):
                    assert isinstance(dep_field.calc_spec, FormulaCalculationSpec)
                    dep_fields[i] = dep_field.clone(
                        guid_formula=BIField.rename_in_formula(
                            dep_field.guid_formula, key_map={old_field.guid: new_field.guid}
                        )
                    )

        # Update the fields
        for dep_field in dep_fields:
            self._update_field(
                old_field=self.get_field_by_id(dep_field.guid),
                new_field=dep_field,
                recursive=True,
                visited_guids=visited_guids,
            )

    def _update_field(
        self,
        old_field: Optional[BIField],
        new_field: Optional[BIField],
        order_index: Optional[int] = None,
        recursive: bool = False,
        visited_guids: frozenset[str] = frozenset(),
        explicitly_updated: bool = False,
    ) -> None:
        """Update/delete/add field and apply autofixes."""

        latest_field = field_not_none(new_field or old_field)
        field_id = latest_field.guid
        if field_id in visited_guids:
            # recursion detected
            return
        visited_guids |= frozenset([field_id])

        component_ref = DatasetComponentRef(component_type=ComponentType.field, component_id=field_id)

        # Clear old errors
        self._ds.error_registry.remove_errors(id=field_id)

        # Remove expression and dependency caches for field
        self.formula_compiler.uncache_field(latest_field)

        if new_field is None and old_field is not None:
            # Field is being deleted
            self.formula_compiler.unregister_field(old_field)
            self._ds.result_schema.remove(field_id=old_field.guid)
        elif old_field is None and new_field is not None:
            # Field is being created
            if order_index is None:
                order_index = len(self._ds.result_schema)
            self.formula_compiler.register_field(new_field)
            new_field = self._get_autoupdated_field(new_field, explicitly_updated=explicitly_updated)
            self.formula_compiler.update_field(new_field)
            self._ds.result_schema.add(idx=order_index, field=new_field)
        else:
            assert new_field is not None
            # Field is being updated
            new_field = self._get_autoupdated_field(new_field, explicitly_updated=explicitly_updated)
            if new_field != old_field:
                self.formula_compiler.update_field(new_field)
                field_idx = self._ds.result_schema.index(field_id=field_id)
                self._ds.result_schema.update_field(idx=field_idx, field=new_field)

        self._reload_formalized_specs()
        if recursive and new_field != old_field:
            self._update_dependent_fields(
                old_field,
                new_field,
                visited_guids=visited_guids,
            )

        if new_field is not None and old_field is not None and new_field.cast != old_field.cast:
            oblig_filter = self._ds_accessor.get_obligatory_filter_opt(field_guid=old_field.guid)
            if oblig_filter is not None:
                allowed_filters = FILTERS_BY_TYPE[enum_not_none(new_field.cast)]
                for default_filter in oblig_filter.default_filters:
                    if default_filter.operation not in allowed_filters:
                        self._ds.error_registry.add_error(
                            # FIXME: Filter errors will be accumulated here
                            #  (new ones being added without the old ones being removed).
                            #  This kind of ruins the idempotency of this method
                            id=oblig_filter.id,
                            type=ComponentType.obligatory_filter,
                            message=f"Filtration type {default_filter.operation.name} "
                            f"is not allowed for field type {enum_not_none(new_field.cast).name}",
                            code=common_exc.DatasetConfigurationError.err_code,  # TODO: add specific err_code
                            details=dict(
                                filter_type=default_filter.operation.name,
                                field_type=enum_not_none(new_field.cast).name,
                            ),
                        )
                        of_component_ref = DatasetComponentRef(
                            component_type=ComponentType.obligatory_filter, component_id=oblig_filter.id
                        )
                        self.add_affected_component(of_component_ref)

        if new_field is not None:
            for error in self._get_field_errors(field=new_field):
                self._ds.error_registry.add_error(
                    id=field_id,
                    type=ComponentType.field,
                    message=error.message,
                    code=error.code,
                    details={"token": error.token, "row": error.coords[0], "column": error.coords[1]},
                )
        self.add_affected_component(component_ref)
        old_title = old_field.title if old_field is not None else None
        new_title = new_field.title if new_field is not None else None
        self.perform_component_title_validation(component_ref=component_ref, old_title=old_title, new_title=new_title)

        # Update field dependencies in dataset
        field_dep_mgr = self._ds_dep_mgr_factory.get_field_shallow_inter_dependency_manager()
        if new_field is None:
            field_dep_mgr.clear_field_direct_references(dep_field_id=field_id)  # type: ignore  # TODO: fix
        else:
            field_dep_mgr.set_field_direct_references(
                dep_field_id=field_id, ref_field_ids=self.formula_compiler.get_referenced_fields(new_field)
            )

    def _get_unpatched_field_ids(self) -> list[str]:
        """Return fields that required patching of properties (e.g. cast, initial_data_type, etc.)"""
        return [
            f.guid
            for f in self._ds.result_schema
            if (f.type is None or f.initial_data_type is None or f.cast is None or f.data_type is None)
        ]

    def update_unpatched_fields(self) -> None:
        """Patch all fields that require patching (non-recursive)"""
        for field_id in self._get_unpatched_field_ids():
            field = self.get_field_by_id(field_id)
            self._update_field(field, field, recursive=False)

    def _get_raw_schema_column(self, field: BIField) -> Optional[SchemaColumn]:
        avatar_id = field.avatar_id
        if avatar_id is None:
            LOGGER.error("Empty avatar_id for field %s, cannot check auto aggregation", field.guid)
            return None
        avatar = self._ds_accessor.get_avatar_opt(avatar_id=avatar_id)
        if avatar is None:
            LOGGER.warning(
                "Avatar not found for id %s from field %s, cannot check auto aggregation", avatar_id, field.guid
            )
            return None
        dsrc_coll = self._get_data_source_coll_strict(source_id=avatar.source_id)
        raw_schema = dsrc_coll.get_cached_raw_schema()
        try:
            return [c for c in raw_schema if c.name == field.source][0]  # type: ignore  # TODO: fix
        except IndexError:
            LOGGER.warning("No field with source %s found in raw_schema", field.source)
            return None

    @generic_profiler("validator-apply-field-action")
    def apply_field_action(
        self,
        action: DatasetAction,
        field_data: UpdateField,
        order_index: Optional[int] = None,
        by: Optional[ManagedBy] = ManagedBy.user,
    ) -> None:
        """Apply update to the result schema configuration"""

        field_data_dict = field_data.as_dict()
        field_data_dict = {k: v for k, v in field_data_dict.items() if v is not None}
        field_id = field_data_dict.pop("guid", None) or self._id_generator.make_field_id(field_data_dict.get("title"))
        new_field_id = field_data_dict.pop("new_id", None)
        order_index = order_index or 0  # added to the beginning by default for usability reasons
        component_ref = DatasetComponentRef(component_type=ComponentType.field, component_id=field_id)
        strict = field_data_dict.pop("strict", False)

        update_field_id = new_field_id is not None and action == DatasetAction.update_field
        if update_field_id:
            if not self._id_validator.is_valid(new_field_id):
                raise exc.DLValidationFatal("Field ID must be [a-z0-9_\\-]{1,36}")
            new_component_ref = DatasetComponentRef(component_type=ComponentType.field, component_id=new_field_id)
            self.perform_component_id_validation(component_ref=new_component_ref)
            # non strict option is used only for old charts. not needed for id update
            strict = True

        old_field: Optional[BIField] = None
        new_field: Optional[BIField] = None

        if action == DatasetAction.add_field:
            self.perform_component_id_validation(component_ref=component_ref)
            self.validate_result_schema_length(field_cnt_diff=1)

        if action in (DatasetAction.add_field, DatasetAction.update_field):
            if "formula" in field_data_dict and "guid_formula" not in field_data_dict:
                field_data_dict["guid_formula"] = ""
            elif "guid_formula" in field_data_dict and "formula" not in field_data_dict:
                field_data_dict["formula"] = ""

        if action in (DatasetAction.update_field, DatasetAction.delete_field):
            try:
                old_field = self.get_field_by_id(field_id)
            except common_exc.FieldNotFound:
                err_msg = f"Field with ID {field_id} doesn't exist"
                LOGGER.error(err_msg)
                if strict:
                    raise exc.DLValidationFatal(err_msg)
                return

            self._ds_ca.validate_component_can_be_managed(component_ref=component_ref, by=by)
            LOGGER.info(f"Field {field_id} old data", extra=old_field._asdict())

        if action == DatasetAction.add_field or update_field_id:
            checked_field_id = field_id if not update_field_id else new_field_id
            try:
                self.get_field_by_id(checked_field_id)
            except common_exc.FieldNotFound:
                pass
            else:
                err_msg = "Field with ID {} already exists".format(checked_field_id)
                LOGGER.error(err_msg)
                if strict:
                    raise exc.DLValidationFatal(err_msg)
                return

        if action == DatasetAction.update_field:
            # Update field id
            assert old_field is not None
            if update_field_id:
                index = self._ds.result_schema.index(field_id)
                self.apply_field_action(
                    action=DatasetAction.delete_field,
                    field_data=DeleteField(guid=field_id),
                    by=by,
                )
                add_field_data = dict(
                    guid=new_field_id,
                    title=old_field.title,
                    type=old_field.type,
                    aggregation=old_field.aggregation,
                    cast=old_field.cast,
                    has_auto_aggregation=old_field.has_auto_aggregation,
                    lock_aggregation=old_field.lock_aggregation,
                    hidden=old_field.hidden,
                    calc_spec=old_field.calc_spec,
                )
                self.apply_field_action(
                    action=DatasetAction.add_field,
                    field_data=AddField(**add_field_data),
                    by=by,
                    order_index=index,
                )
                self._ds.rename_field_id_usages(old_id=field_id, new_id=new_field_id)
                return
            # Update existing field
            new_field = old_field.clone(**field_data_dict)
            if new_field == old_field:
                # Nothing changed, so just quit
                return

        elif action == DatasetAction.add_field:
            # add new field
            # TODO: use create_result_schema_field from dataset instead
            field_data_dict["hidden"] = field_data_dict.get("hidden", False)
            field_data_dict["aggregation"] = field_data_dict.get("aggregation", AggregationFunction.none)
            if "calc_spec" not in field_data_dict:
                default_calc_mode = CalcMode.formula if field_data_dict.get("formula") else CalcMode.direct
                field_data_dict["calc_mode"] = field_data_dict.get("calc_mode", default_calc_mode)
                field_data_dict["calc_spec"] = field_data_dict.get("calc_spec", create_calc_spec_from(field_data_dict))
            field_data_dict = del_calc_spec_kwargs_from(field_data_dict)

            if (
                isinstance(field_data_dict["calc_spec"], DirectCalculationSpec)
                and not field_data_dict["calc_spec"].avatar_id
            ):
                # TODO: remove this after old charts are somehow updated
                root_avatar = self._ds_accessor.get_root_avatar_strict()
                field_data_dict["calc_spec"] = field_data_dict["calc_spec"].clone(avatar_id=root_avatar.id)

            new_field = BIField.make(guid=field_id, **field_data_dict)

        elif action == DatasetAction.clone_field:
            # clone existing field under a new title
            from_field_id = field_data_dict["from_guid"]
            from_field = self.get_field_by_id(from_field_id)
            if by != ManagedBy.user:
                raise common_exc.DatasetConfigurationError("Non-user-managed fields cannot be cloned")

            cast = field_data_dict.get("cast")
            cast = cast if cast is not None else from_field.cast
            aggregation = field_data_dict.get("aggregation")
            aggregation = aggregation if aggregation is not None else from_field.aggregation

            # It has the same type as the original field if both have aggregations or both don't
            same_type_as_original = (
                from_field.aggregation != AggregationFunction.none
                and aggregation != AggregationFunction.none
                or from_field.aggregation == AggregationFunction.none
                and aggregation == AggregationFunction.none
            )

            add_field_data = dict(
                guid=field_id,
                title=field_data_dict["title"],
                type=from_field.type if same_type_as_original else None,
                aggregation=aggregation,
                cast=cast,
                has_auto_aggregation=from_field.has_auto_aggregation,
                lock_aggregation=from_field.lock_aggregation,
                hidden=from_field.hidden,
                calc_spec=from_field.calc_spec,
            )
            self.apply_field_action(
                action=DatasetAction.add_field,
                field_data=AddField(**add_field_data),
                by=ManagedBy.user,
                order_index=order_index,
            )
            return

        elif action == DatasetAction.delete_field:
            # delete field

            # Check if field obligatory filter is set up
            if self._ds_accessor.get_obligatory_filter_opt(field_guid=field_id) is not None:
                raise common_exc.DatasetConfigurationError(
                    f'Obligatory filter present for field "{field_id}". It should be deleted first.'
                )

        self._update_field(
            old_field=old_field,
            new_field=new_field,
            order_index=order_index,
            recursive=True,
            explicitly_updated=True,
        )

        # errors can be cleared or fixed by now, so revalidate the result schema length
        if action in (DatasetAction.add_field, DatasetAction.delete_field):
            self.validate_result_schema_length()

    def _update_direct_fields_for_updated_raw_schema(
        self,
        old_raw_schema: Sequence[SchemaColumn],
        new_raw_schema: Optional[Sequence[SchemaColumn]],
        avatar_ids: list[str],
        do_delete_fields: bool = False,
        force_update_fields: bool = False,
    ) -> None:
        old_title_by_name = {col.name: col.title for col in (old_raw_schema or ())}
        new_name_by_title = {col.title: col.name for col in (new_raw_schema or ())}

        old_column_by_name = {col.name: col for col in (old_raw_schema or ())}
        new_column_by_name = {col.name: col for col in (new_raw_schema or ())}

        def apply_field_changes() -> None:
            # Before updating all of them, uncache all affected fields to avoid loading incorrect dependencies
            for old_field, new_field in updated_fields:
                self.formula_compiler.uncache_field(field_not_none(old_field or new_field))
            # Now move on to updating
            updated_ids: list[str] = []
            by_id: dict[str, tuple[Optional[BIField], Optional[BIField]]] = {}
            for old_field, new_field in updated_fields:
                self._update_field(old_field, new_field, recursive=False)
                field_id = field_not_none(old_field or new_field).guid
                if field_id not in by_id:
                    by_id[field_id] = (old_field, new_field)
                    updated_ids.append(field_id)

            # Now update dependent fields
            for field_id in updated_ids:
                new_field, old_field = by_id[field_id]
                self._update_dependent_fields(
                    old_field=old_field,
                    new_field=new_field,
                    visited_guids=frozenset(),
                )

            updated_fields.clear()

        field_cnt_diff = 0
        fields_by_avatar_and_source = set()
        updated_fields: list[tuple[Optional[BIField], Optional[BIField]]] = []
        for old_field in self._ds.result_schema:
            if old_field.calc_mode != CalcMode.direct or old_field.avatar_id not in avatar_ids:
                continue

            # try to match them to the new raw schema by title from the original raw schema
            original_title = old_title_by_name.get(old_field.source)
            new_field = old_field
            if original_title in new_name_by_title:
                if new_name_by_title[original_title] != old_field.source:
                    LOGGER.info(
                        f'Replacing source column "{original_title}" '
                        f'with "{new_name_by_title[original_title]}" in field {old_field.guid}'
                    )
                    # match by title
                    new_field = new_field.clone(source=new_name_by_title[original_title])

            else:
                # Column is no longer present
                if do_delete_fields:  # Source was deleted
                    # In this case delete the fields
                    new_field = None
                    field_cnt_diff -= 1

            # re-validate ALL fields that belong to the specified avatars
            # and that have changed in some way
            latest_field = field_not_none(new_field or old_field)
            old_column = old_column_by_name.get(latest_field.source)
            new_column = new_column_by_name.get(latest_field.source)
            if force_update_fields or (new_field != old_field) or (new_column != old_column):
                updated_fields.append((old_field, new_field))

            fields_by_avatar_and_source.add((latest_field.avatar_id, latest_field.source))

        apply_field_changes()

        # generate new fields for new columns
        if new_raw_schema:
            for avatar_id in avatar_ids:
                for new_col in new_raw_schema:
                    if (avatar_id, new_col.name) not in fields_by_avatar_and_source:
                        field_cnt_diff += 1
                        new_field = BIField.make(
                            **self._ds.create_result_schema_field(
                                column=new_col,
                                field_id_generator=self._id_generator,
                                avatar_id=avatar_id,
                            )
                        )
                        updated_fields.append((None, new_field))
        self.validate_result_schema_length(field_cnt_diff)

        apply_field_changes()

    def add_affected_component(self, component_ref: DatasetComponentRef) -> None:
        """
        Register component as affected by current update action.
        This registry is used for updating validity status of components
        """
        self._affected_components.add(component_ref)

    def update_validity_of_affected_components(self, clear: bool = True) -> None:
        """
        Iterate over components registered in ``self._affected_components``
        and update their ``valid`` attribute based on the presence or absence of errors
        in the dataset's error registry.
        """
        for component_ref in self._affected_components:
            self._ds_ca.update_component_validity(
                component_ref=component_ref, valid=not self._ds.error_registry.get_pack(id=component_ref.component_id)
            )
        if clear:
            self._affected_components.clear()

    def perform_component_id_validation(self, component_ref: DatasetComponentRef) -> None:
        if component_ref.component_id in (toplevel_id.value for toplevel_id in TopLevelComponentId):
            raise exc.DLValidationFatal(f"Cannot use a reserved ID - {component_ref.component_id}")

        # TODO validate component id itself: length, allowed symbols, etc.

        # TODO bring this back when external-api defaults are patched
        # other_types = {component_type for component_type in ComponentType} - {component_ref.component_type}
        #
        # for other_type in other_types:
        #     for other_component in self._ds_ca.iter_dataset_components_by_type(other_type):
        #         other_component_id = other_component.guid if isinstance(other_component, BIField) else other_component.id
        #
        #         if component_ref.component_id == other_component_id:
        #             raise exc.DLValidationFatal(
        #                 f'Component IDs are not unique: '
        #                 f'found {component_ref.component_type.name} and {other_type.name} '
        #                 f'with ID {other_component_id}'
        #             )

    def perform_component_title_validation(
        self, component_ref: DatasetComponentRef, old_title: Optional[str] = None, new_title: Optional[str] = None
    ) -> None:
        """
        Validate title change for the given component.
        If a conflict for this title is found among some of the components,
        then register errors for all of them (if not registered yet).
        If a conflict was present for the old title, but is not there for the new one,
        then remove errors.

        Return list of affected component IDs

        Supports all component types with titles: data_source (collection), source_avatar, field
        """
        old_title_conflict_ids = set()
        new_title_conflict_ids = set()
        exc_cls = {
            ComponentType.data_source: common_exc.DataSourceTitleConflict,
            ComponentType.source_avatar: common_exc.SourceAvatarTitleConflict,
            ComponentType.field: common_exc.FieldTitleConflict,
        }[component_ref.component_type]

        for other_component in self._ds_ca.iter_dataset_components_by_type(component_type=component_ref.component_type):
            if not isinstance(other_component, (DataSourceCollectionBase, SourceAvatar, BIField)):
                raise TypeError(f"Component type {type(other_component)} is not supported here")

            other_component_id = other_component.guid if isinstance(other_component, BIField) else other_component.id
            if other_component_id == component_ref.component_id:
                continue
            if old_title is not None and other_component.title == old_title:
                old_title_conflict_ids.add(other_component_id)
            if new_title is not None and other_component.title == new_title:
                new_title_conflict_ids.add(other_component_id)

        if old_title_conflict_ids:
            old_title_conflict_ids.add(component_ref.component_id)
        if new_title_conflict_ids:
            new_title_conflict_ids.add(component_ref.component_id)

        for other_component_id in old_title_conflict_ids | new_title_conflict_ids:
            other_component_ref = DatasetComponentRef(
                component_type=component_ref.component_type,
                component_id=other_component_id,
            )
            if other_component_id in new_title_conflict_ids:
                # new conflict
                pack = self._ds.error_registry.get_pack(id=other_component_id)
                if pack is None or not pack.get_errors(code=exc_cls.err_code):
                    self._ds.error_registry.add_error(
                        id=other_component_id,
                        type=component_ref.component_type,
                        message=exc_cls.default_message,
                        code=exc_cls.err_code,
                        details={"title": new_title, "other_id": other_component_id},
                    )
            else:
                # old conflict, now resolved, so remove error
                self._ds.error_registry.remove_errors(id=other_component_id, code_prefix=exc_cls.err_code)

            self.add_affected_component(other_component_ref)

    def _are_schemas_identical(
        self,
        old_raw_schema: Optional[Sequence[SchemaColumn]],
        new_raw_schema: Optional[Sequence[SchemaColumn]],
    ) -> bool:
        if old_raw_schema is None and new_raw_schema is None:
            return True
        if old_raw_schema is None or new_raw_schema is None:
            return False
        assert old_raw_schema is not None and new_raw_schema is not None

        if type(old_raw_schema) is not type(new_raw_schema):
            old_raw_schema = tuple(old_raw_schema)
            new_raw_schema = tuple(new_raw_schema)
        return new_raw_schema == old_raw_schema

    def refresh_data_source(
        self,
        source_id: str,
        old_raw_schema: Optional[Sequence[SchemaColumn]],
        force_update_fields: bool = False,
    ) -> None:
        """
        Patch data source with parameters from the database.
        If any errors are encountered, then register them in the dataset.
        """

        sr = self._service_registry

        @contextmanager
        def source_error_ctx() -> Generator[None, None, None]:
            try:
                yield
            except common_exc.DLBaseException as err:
                LOGGER.exception(f"Failed to check source {source_id}. Assuming it doesn't exist")
                self._ds.error_registry.add_error(
                    id=source_id,
                    type=ComponentType.data_source,
                    message=err.message,
                    code=err.err_code,
                    details=err.details,
                )

        dsrc_coll = self._get_data_source_coll_strict(source_id=source_id)
        origin_dsrc = dsrc_coll.get_strict(role=DataSourceRole.origin)

        def conn_executor_factory_func() -> SyncConnExecutorBase:
            conn_executor = sr.get_conn_executor_factory().get_sync_conn_executor(conn=origin_dsrc.connection)
            return conn_executor

        @generic_profiler("validator-get-source-exists")
        def get_source_exists() -> bool:
            exists = False  # noqa
            with source_error_ctx():
                exists = origin_dsrc.source_exists(conn_executor_factory=conn_executor_factory_func)
            return exists

        @generic_profiler("validator-get-db-info")
        def get_db_version(exists: bool) -> Optional[str]:
            db_version = None
            if exists:
                db_version = origin_dsrc.get_db_info(conn_executor_factory=conn_executor_factory_func).version
                LOGGER.info(f"Got version {db_version} for data source {source_id}")
            return db_version

        @generic_profiler("validator-get-schema")
        def get_schema_info(exists: bool) -> Optional[SchemaInfo]:
            with source_error_ctx():
                if exists:
                    schema_info = origin_dsrc.get_schema_info(conn_executor_factory=conn_executor_factory_func)
                    # TODO FIX: BI-2355 Ensure that dsrc.get_schema_info() do not suppress exceptions and remove
                    if not schema_info.schema:  # TODO: remove when get_schema starts raising errors
                        raise common_exc.DLBaseException(message="Failed to load table schema")
                    return schema_info
                else:
                    params = dict()
                    if isinstance(origin_dsrc, BaseSQLDataSource):
                        table_def = origin_dsrc.get_table_definition()
                        params["table_definition"] = str(table_def)

                    raise common_exc.SourceDoesNotExist(
                        db_message="",
                        query="",
                        params=params,
                    )
            return None

        exists = get_source_exists()
        LOGGER.info(f"Data source {source_id} exists: {exists}")

        new_raw_schema_info = get_schema_info(exists=exists)

        new_raw_schema = new_raw_schema_info.schema if new_raw_schema_info is not None else []
        new_idx_info_set = new_raw_schema_info.indexes if new_raw_schema_info is not None else None

        LOGGER.info(f"Got raw_schema with {len(new_raw_schema)} columns for data source {source_id}")

        db_version = get_db_version(exists=exists)
        self._ds_editor.update_data_source(
            source_id=source_id,
            role=DataSourceRole.origin,
            raw_schema=new_raw_schema,
            db_version=db_version,
            index_info_set=new_idx_info_set,
        )

        self._reload_sources()

        if not self._are_schemas_identical(new_raw_schema, old_raw_schema) or force_update_fields:
            # try to match old and new schemas against each other
            # and update result_schema fields accordingly
            avatar_ids = [avatar.id for avatar in self._ds_accessor.get_avatar_list(source_id=source_id)]
            self._update_direct_fields_for_updated_raw_schema(
                old_raw_schema=old_raw_schema,  # type: ignore  # TODO: fix
                new_raw_schema=new_raw_schema,
                avatar_ids=avatar_ids,
                do_delete_fields=False,
                force_update_fields=force_update_fields,
            )

    def validate_source_already_exists(self, source_id: str, source_data: dict) -> bool:
        """
        Check if an identical source already exists.
        If it does, then register the ID as remapped, return ``True``.
        Otherwise return ``False``.
        """

        existing_id = self._ds.find_data_source_configuration(
            connection_id=source_data.get("connection_id"),
            created_from=source_data["source_type"],
            parameters=source_data.get("parameters"),
            # Currently ignoring: `title`
        )
        if existing_id:
            # such source already exists, don't add its copy to the dataset
            # remap all following actions to use the existing source's ID
            LOGGER.info("Such a source already exists, so remapping ID %s to %s", source_id, existing_id)
            self._remapped_source_ids[source_id] = existing_id
            return True

        return False

    @generic_profiler("validator-apply-source-action")
    def apply_source_action(
        self,
        action: DatasetAction,
        source_data: dict,
        by: Optional[ManagedBy] = ManagedBy.user,
        ignore_source_ids: Optional[Collection[str]] = None,
    ) -> None:
        """Apply update to the data source configuration"""

        source_data = source_data.copy()
        source_id = source_data.pop("id") or str(uuid.uuid4())
        component_ref = DatasetComponentRef(component_type=ComponentType.data_source, component_id=source_id)
        old_title: Optional[str] = None
        new_title: Optional[str] = None

        if source_id in self._remapped_source_ids:
            LOGGER.warning("Source %s has not been added, so action %s will be ignored", source_id, action.name)
            return

        if action == DatasetAction.add_source:
            if self.validate_source_already_exists(source_id=source_id, source_data=source_data):
                return
            self.perform_component_id_validation(component_ref)

        def add_source(title: str) -> None:
            connection_id = source_data.get("connection_id")
            if connection_id is None:
                raise exc.DLValidationFatal("source connection_id cannot be None")

            connection_ref = DefaultConnectionRef(conn_id=connection_id)
            self._sync_us_manager.ensure_entry_preloaded(connection_ref)
            self._ds_editor.add_data_source(
                source_id=source_id,
                role=DataSourceRole.origin,
                connection_id=connection_id,
                created_from=source_data["source_type"],
                title=title,
                parameters=source_data["parameters"],
            )

            check_permissions_for_origin_sources(
                dataset=self._ds,
                source_ids=[source_id],
                permission_kind=USPermissionKind.read,
                us_entry_buffer=self._us_manager.get_entry_buffer(),
            )

        if action in (DatasetAction.update_source, DatasetAction.delete_source):
            dsrc_coll = self._get_data_source_coll_strict(source_id=source_id)
            old_title = dsrc_coll.title
            self._ds_ca.validate_component_can_be_managed(component_ref=component_ref, by=by)

        if action in (DatasetAction.update_source, DatasetAction.add_source):
            existing_source_id = source_id if action == DatasetAction.update_source else None
            if not self._capabilities.source_can_be_added(
                connection_id=source_data.get("connection_id"),
                created_from=source_data["source_type"],
                ignore_source_ids=ignore_source_ids or [existing_source_id],  # type: ignore  # TODO: fix
            ):
                raise exc.DLValidationFatal("Source cannot be added to dataset")

        # we have verified that the action is valid, so we can proceed to modifying the dataset

        # clear old errors
        self._ds.error_registry.remove_errors(id=source_id)

        if action == DatasetAction.add_source:
            add_source(title=source_data.get("title") or None)  # type: ignore  # TODO: fix
            self.refresh_data_source(source_id=source_id, old_raw_schema=None)
            new_title = source_data["title"]
            self.add_affected_component(component_ref)

        elif action == DatasetAction.update_source:
            self.add_affected_component(component_ref)
            parameters = source_data.get("parameters") or {}

            dsrc_coll = self._get_data_source_coll_strict(source_id=source_id)
            origin_dsrc = dsrc_coll.get_strict(role=DataSourceRole.origin)
            old_raw_schema = origin_dsrc.saved_raw_schema
            new_title = source_data.get("title") or dsrc_coll.title
            self._ds_editor.update_data_source(
                source_id=source_id,
                created_from=source_data.get("source_type"),
                role=DataSourceRole.origin,
                connection_id=source_data.get("connection_id"),
                **parameters,
            )
            if new_title != old_title:
                self._ds_editor.update_data_source_collection(source_id=source_id, title=new_title)

            if (set(source_data) - {"title", "id", "source_type", "parameters"}) or parameters:
                # something besides the title was updated
                self.refresh_data_source(source_id=source_id, old_raw_schema=old_raw_schema)  # type: ignore  # TODO: fix

        elif action == DatasetAction.delete_source:
            self._ds_editor.remove_data_source_collection(source_id=source_id)
            self._reload_sources()

        elif action == DatasetAction.refresh_source:
            dsrc_coll = self._get_data_source_coll_strict(source_id=source_id)
            origin_dsrc = dsrc_coll.get_strict(role=DataSourceRole.origin)
            new_title = dsrc_coll.title
            assert origin_dsrc is not None
            self.refresh_data_source(
                source_id=source_id,
                old_raw_schema=origin_dsrc.saved_raw_schema,
                # Just in case some of the fields were erroneously marked as invalid,
                # force-update them if `force_update_fields` is true
                force_update_fields=source_data.get("force_update_fields", False),
            )
            self.add_affected_component(component_ref)

        # check title against conflicts
        self.perform_component_title_validation(component_ref=component_ref, old_title=old_title, new_title=new_title)

    @generic_profiler("validator-apply-avatar-action")
    def apply_source_avatar_action(
        self,
        action: DatasetAction,
        source_avatar_data: dict,
        by: Optional[ManagedBy] = ManagedBy.user,
        disable_fields_update: bool = False,
    ) -> None:
        """Apply update to the data source configuration"""

        source_avatar_data = source_avatar_data.copy()
        avatar_id = source_avatar_data.pop("id")
        source_id = source_avatar_data.get("source_id")
        component_ref = DatasetComponentRef(component_type=ComponentType.source_avatar, component_id=avatar_id)
        old_avatar: Optional[SourceAvatar] = None
        old_raw_schema = None
        new_raw_schema = None
        old_title: Optional[str] = None
        new_title: Optional[str] = None
        do_delete_fields = False
        if source_id is not None:
            source_id = self._remapped_source_ids.get(source_id, source_id)

        update_fields = False

        if action in (DatasetAction.update_source_avatar, DatasetAction.delete_source_avatar):
            old_avatar = self._ds_accessor.get_avatar_strict(avatar_id=avatar_id)
            old_dsrc_coll = self._get_data_source_coll_strict(source_id=old_avatar.source_id)
            old_raw_schema = old_dsrc_coll.get_cached_raw_schema(role=DataSourceRole.origin)
            old_title = old_avatar.title
            self._ds_ca.validate_component_can_be_managed(component_ref=component_ref, by=by)

        # clear old errors
        self._ds.error_registry.remove_errors(id=avatar_id)

        if action == DatasetAction.add_source_avatar:
            self.perform_component_id_validation(component_ref)
            new_title = source_avatar_data["title"]
            self._ds_editor.add_avatar(
                avatar_id=avatar_id,
                source_id=source_id,
                title=new_title,  # type: ignore  # TODO: fix
            )
            new_dsrc_coll = self._get_data_source_coll_strict(source_id=source_id)
            new_raw_schema = new_dsrc_coll.get_cached_raw_schema(role=DataSourceRole.origin)
            self.formula_compiler.register_avatar(avatar_id=avatar_id, source_id=source_id)
            update_fields = True

        if action == DatasetAction.update_source_avatar:
            new_title = source_avatar_data.get("title") or old_avatar.title  # type: ignore  # TODO: fix
            if new_title != old_title:
                self._ds_editor.update_avatar(avatar_id=avatar_id, title=new_title)
            if source_id is not None and source_id != old_avatar.source_id:  # type: ignore  # TODO: fix
                new_dsrc_coll = self._get_data_source_coll_strict(source_id=source_id)
                self._ds_editor.update_avatar(avatar_id=avatar_id, source_id=source_id)
                new_raw_schema = new_dsrc_coll.get_cached_raw_schema(role=DataSourceRole.origin)
                self.formula_compiler.unregister_avatar(avatar_id=avatar_id)
                self.formula_compiler.register_avatar(avatar_id=avatar_id, source_id=source_id)
                update_fields = True

        elif action == DatasetAction.delete_source_avatar:
            # clean up relations
            left_relations = self._ds_accessor.get_avatar_relation_list(right_avatar_id=avatar_id)
            for relation in left_relations:
                self._ds_editor.remove_avatar_relation(relation_id=relation.id)

            # remove avatar
            self._ds_editor.remove_avatar(avatar_id=avatar_id)
            self.formula_compiler.unregister_avatar(avatar_id=avatar_id)
            update_fields = True
            do_delete_fields = True

        # update result schema fields
        if update_fields and not disable_fields_update:
            self._update_direct_fields_for_updated_raw_schema(
                old_raw_schema=old_raw_schema,  # type: ignore  # TODO: fix
                new_raw_schema=new_raw_schema,
                avatar_ids=[avatar_id],
                do_delete_fields=do_delete_fields,
            )

        # check title against conflicts
        self.perform_component_title_validation(component_ref=component_ref, old_title=old_title, new_title=new_title)

        self.load_exbuilders()

    def _try_generate_conditions_from_columns(
        self,
        left_avatar_id: str,
        right_avatar_id: str,
    ) -> list[BinaryCondition]:
        """Automatically generate condition for first pair of matching field names"""

        left_columns = self._get_data_source_strict(
            source_id=self._ds_accessor.get_avatar_strict(avatar_id=left_avatar_id).source_id,
            role=DataSourceRole.origin,
        ).saved_raw_schema
        right_columns = self._get_data_source_strict(
            source_id=self._ds_accessor.get_avatar_strict(avatar_id=right_avatar_id).source_id,
            role=DataSourceRole.origin,
        ).saved_raw_schema

        if not left_columns or not right_columns:
            if right_columns and not left_columns:
                which = "left"
            elif left_columns and not right_columns:
                which = "right"
            else:
                which = "both"
            LOGGER.warning(f"Can't generate condition: schema for {which} is missing")
            return []

        left_by_name = {col.name: col for col in left_columns}
        right_by_name = {col.name: col for col in right_columns}
        right_by_lower_name = {col.name.lower(): col for col in right_columns}
        column_match = None
        for left_col in sorted(left_by_name):
            # first try to match by exact name,
            # and then - ignore case
            if left_col in right_by_name:
                column_match = (left_by_name[left_col], right_by_name[left_col])
                break
            if left_col.lower() in right_by_lower_name:
                column_match = (left_by_name[left_col], right_by_lower_name[left_col.lower()])
                break
        if column_match:
            return [
                BinaryCondition(
                    left_part=ConditionPartDirect(source=column_match[0].name),
                    right_part=ConditionPartDirect(source=column_match[1].name),
                    operator=BinaryJoinOperator.eq,
                )
            ]
        else:
            return []

    @generic_profiler("validator-apply-relation-action")
    def apply_avatar_relation_action(
        self, action: DatasetAction, avatar_relation_data: dict, by: Optional[ManagedBy] = ManagedBy.user
    ) -> None:
        """Apply update to the data source configuration"""

        avatar_relation_data = avatar_relation_data.copy()
        relation_id = avatar_relation_data.pop("id")
        component_ref = DatasetComponentRef(component_type=ComponentType.avatar_relation, component_id=relation_id)

        if by is not None and action in (DatasetAction.update_avatar_relation, DatasetAction.delete_avatar_relation):
            self._ds_ca.validate_component_can_be_managed(component_ref=component_ref, by=by)

        # clear old errors
        self._ds.error_registry.remove_errors(id=relation_id)

        if action == DatasetAction.add_avatar_relation:
            self.perform_component_id_validation(component_ref)
            conditions = avatar_relation_data.get("conditions")
            left_avatar_id = avatar_relation_data["left_avatar_id"]
            right_avatar_id = avatar_relation_data["right_avatar_id"]
            if not conditions:
                conditions = self._try_generate_conditions_from_columns(
                    left_avatar_id=left_avatar_id, right_avatar_id=right_avatar_id
                )
            self._ds_editor.add_avatar_relation(
                relation_id=relation_id,
                left_avatar_id=left_avatar_id,
                right_avatar_id=right_avatar_id,
                conditions=conditions,
                join_type=avatar_relation_data.get("join_type"),
            )

        if action == DatasetAction.update_avatar_relation:
            self._ds_editor.update_avatar_relation(
                relation_id=relation_id,
                conditions=avatar_relation_data.get("conditions"),
                join_type=avatar_relation_data.get("join_type"),
            )

        if action == DatasetAction.delete_avatar_relation:
            self._ds_editor.remove_avatar_relation(relation_id=relation_id)

        if action in (DatasetAction.add_avatar_relation, DatasetAction.update_avatar_relation):
            errors = self._get_relation_errors(
                relation=self._ds_accessor.get_avatar_relation_strict(relation_id=relation_id),
            )
            for error in errors:
                self._ds.error_registry.add_error(
                    id=relation_id,
                    type=ComponentType.avatar_relation,
                    message=error.message,
                    code=error.code,
                    details={},
                )
            self.add_affected_component(component_ref)

        self.load_exbuilders()

    def _migrate_source_parameters(
        self,
        old_connection: Optional[ConnectionBase],
        new_connection: ConnectionBase,
        dsrc: DataSource,
    ) -> tuple[dict, DataSourceType]:
        old_conn_type: ConnectionType
        if old_connection is not None:
            old_conn_type = old_connection.conn_type
        else:
            old_conn_type = dsrc.conn_type

        src_migrator = get_data_source_migrator(conn_type=old_conn_type)
        dst_migrator = get_data_source_migrator(conn_type=new_connection.conn_type)

        migration_dtos = src_migrator.export_migration_dtos(
            data_source_spec=dsrc.spec,
        )
        try:
            new_dsrc_spec = dst_migrator.import_migration_dtos(
                migration_dtos=migration_dtos,
                connection_ref=new_connection.conn_ref,
            )
        except common_exc.DataSourceMigrationImpossible:
            # Failed to migrate anything.
            # Let's just go with an empty default data source
            return {}, new_connection.source_type

        new_dsrc_cls = get_data_source_class(new_dsrc_spec.source_type)
        new_dsrc_dummy = new_dsrc_cls(
            id="",
            spec=new_dsrc_spec,
            us_entry_buffer=self._us_manager.get_entry_buffer(),
        )
        new_dsrc_parameters = new_dsrc_dummy.get_parameters()
        return new_dsrc_parameters, new_dsrc_spec.source_type

    @generic_profiler("validator-apply-connection-action")
    def apply_connection_action(
        self, action: DatasetAction, connection_data: ReplaceConnection, by: Optional[ManagedBy] = ManagedBy.user
    ) -> None:
        assert action == DatasetAction.replace_connection
        old_connection: Optional[ConnectionBase]
        try:
            old_connection_ref = DefaultConnectionRef(conn_id=connection_data.id)
            self._sync_us_manager.ensure_entry_preloaded(old_connection_ref)
            old_connection = self._sync_us_manager.get_loaded_us_connection(old_connection_ref)
        except common_exc.ReferencedUSEntryNotFound:
            old_connection = None

        new_connection_ref = DefaultConnectionRef(conn_id=connection_data.new_id)
        self._sync_us_manager.ensure_entry_preloaded(new_connection_ref)
        new_connection = self._sync_us_manager.get_loaded_us_connection(new_connection_ref)
        utils.need_permission_on_entry(new_connection, USPermissionKind.read)

        updates = []
        ignore_source_ids = []  # exclude all the sources that are going to be updated from compatibility checks
        for source_id in self._ds_accessor.get_data_source_id_list():
            dsrc_coll = self._get_data_source_coll_strict(source_id=source_id)
            dsrc = dsrc_coll.get_strict(role=DataSourceRole.origin)
            connection_ref = dsrc.connection_ref
            if not isinstance(connection_ref, DefaultConnectionRef):
                raise TypeError(f"Unexpected connection_ref type: {type(connection_ref)}")
            if connection_ref.conn_id == connection_data.id:
                parameters, new_source_type = self._migrate_source_parameters(
                    old_connection=old_connection,
                    new_connection=new_connection,
                    dsrc=dsrc,  # type: ignore  # TODO: fix
                )

                # migrate parameters for new source type
                updates.append(
                    dict(
                        id=dsrc_coll.id,
                        connection_id=connection_data.new_id,
                        source_type=new_source_type,
                        parameters=parameters,
                    )
                )
                ignore_source_ids.append(dsrc_coll.id)

        LOGGER.info(f"Found {len(updates)} sources to update")
        for dsrc_data in updates:
            LOGGER.info(
                f'Going to update source {dsrc_data["id"]} as {dsrc_data["source_type"].name} '
                f"for connection {connection_data.new_id}"
            )
            self.apply_source_action(
                action=DatasetAction.update_source,
                source_data=dsrc_data,
                by=by,
                ignore_source_ids=ignore_source_ids,
            )

    @generic_profiler("validator-apply-obligatory-filter-action")
    def apply_obligatory_filter_action(
        self, action: DatasetAction, filter_data: ObligatoryFilterBase, by: Optional[ManagedBy] = ManagedBy.user
    ) -> None:
        assert isinstance(filter_data, ObligatoryFilterBase)
        component_ref = DatasetComponentRef(component_type=ComponentType.obligatory_filter, component_id=filter_data.id)

        if action in (DatasetAction.add_obligatory_filter, DatasetAction.update_obligatory_filter):
            assert isinstance(filter_data, AddUpdateObligatoryFilter)
            if action == DatasetAction.update_obligatory_filter:
                filter_object = self._ds_accessor.get_obligatory_filter_strict(obfilter_id=filter_data.id)
                field_guid = filter_object.field_guid
            else:
                assert isinstance(filter_data.field_guid, str)
                self.perform_component_id_validation(component_ref)
                field_guid = filter_data.field_guid

            field = self.get_field_by_id(field_guid)  # Validating field guid
            default_where_clauses = []
            allowed_filter_types = FILTERS_BY_TYPE[enum_not_none(field.cast)]
            for default_filter in filter_data.default_filters:
                clause = DefaultWhereClause(operation=default_filter.operation, values=default_filter.values)
                if clause.operation not in allowed_filter_types:
                    raise common_exc.DatasetConfigurationError(
                        f"Filtration type {clause.operation.name} is not allowed "
                        f"for field type {enum_not_none(field.cast).name}"
                    )
                default_where_clauses.append(clause)

            if action == DatasetAction.add_obligatory_filter:
                valid = filter_data.valid if filter_data.valid is not None else True
                assert valid is not None
                self._ds_editor.add_obligatory_filter(
                    obfilter_id=filter_data.id,
                    field_guid=field_guid,
                    default_filters=default_where_clauses,
                    managed_by=by,
                    valid=valid,
                )
            else:
                self._ds_editor.update_obligatory_filter(
                    obfilter_id=filter_data.id,
                    default_filters=default_where_clauses,
                    valid=filter_data.valid,
                )

        if action == DatasetAction.delete_obligatory_filter:
            assert isinstance(filter_data, DeleteObligatoryFilter)
            self._ds_editor.remove_obligatory_filter(obfilter_id=filter_data.id)

    @generic_profiler("validator-get-single-formula-errors")
    def get_single_formula_errors(self, formula: str, feature_errors: bool = False) -> list[ErrorInfo]:
        if not self._has_sources:
            errors = [FormulaErrorCtx(message="No data sources have been defined", level=MessageLevel.ERROR)]

        else:
            formula_errors = self._get_formula_errors(formula, feature_errors=feature_errors)
            if len(formula_errors) == 1:
                # Hide error for empty formulas
                if formula_errors[0].is_sub_error(code=formula_exc.ParseEmptyFormulaError.default_code):
                    formula_errors = []

            errors = [
                FormulaErrorInfo(  # type: ignore  # TODO: fix
                    code=err.code,
                    message=err.message,
                    level=comp_err_level_from_formula_err_level(err.level),
                    row=err.coords[0],
                    column=err.coords[1],
                    token=err.token,
                )
                for err in formula_errors
            ]

        return errors  # type: ignore  # TODO: fix

    def collect_nonexistent_connection_errors(self) -> None:
        for source_id in self._ds_accessor.get_data_source_id_list():
            dsrc_coll = self._get_data_source_coll_strict(source_id=source_id)
            dsrc = dsrc_coll.get_strict(role=DataSourceRole.origin)
            try:
                self._sync_us_manager.ensure_entry_preloaded(dsrc.connection_ref)
            except common_exc.ReferencedUSEntryNotFound:
                self._ds.error_registry.add_error(
                    id=dsrc_coll.id,
                    type=ComponentType.data_source,
                    message=f"Connection for {dsrc_coll.title} not found",
                    code=common_exc.ReferencedUSEntryNotFound.err_code,
                    details={},
                )

    def _find_phantom_error_refs(self) -> list[DatasetComponentRef]:
        all_error_refs = [
            DatasetComponentRef(component_type=item.type, component_id=item.id)
            for item in self._ds.error_registry.items
        ]
        phantom_refs: list[DatasetComponentRef] = []
        for component_ref in all_error_refs:
            if self._ds_ca.get_component(component_ref=component_ref) is None:
                phantom_refs.append(component_ref)

        return phantom_refs

    def _remove_phantom_errors_for_refs(self, phantom_refs: list[DatasetComponentRef]) -> None:
        for ref in phantom_refs:
            self._ds.error_registry.remove_errors(id=ref.component_id)

    def find_and_remove_phantom_error_refs(self) -> None:
        phantom_error_refs = self._find_phantom_error_refs()
        if phantom_error_refs:
            LOGGER.warning(
                f"Detected phantom errors during dataset validation: {phantom_error_refs},"
                " removing from component errors"
            )

        self._remove_phantom_errors_for_refs(phantom_error_refs)

lib/dl_api_lib/dl_api_lib/dataset/validator.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import logging
from typing import (
    TYPE_CHECKING,
    Type,
)

from dl_api_commons.base_models import RequestContextInfo
from dl_api_lib.dataset.base_wrapper import DatasetBaseWrapper
from dl_api_lib.query.formalization.query_formalizer import (
    DataQuerySpecFormalizer,
    TotalsSpecFormalizer,
    ValueDistinctSpecFormalizer,
    ValueRangeSpecFormalizer,
)
from dl_api_lib.query.registry import get_filter_formula_compiler_cls
from dl_constants.enums import (
    CalcMode,
    DataSourceRole,
    ProcessorType,
    SelectorType,
)
from dl_core.fields import BIField
from dl_core.us_connection_base import ClassicConnectionSQL
from dl_core.us_dataset import Dataset
from dl_core.us_manager.us_manager import USManagerBase
from dl_query_processing.enums import QueryType
from dl_query_processing.execution.exec_info import QueryExecutionInfo
from dl_query_processing.execution.executor import QueryExecutor
from dl_query_processing.execution.executor_base import QueryExecutorBase
from dl_query_processing.execution.primitives import ExecutedQuery
from dl_query_processing.legend.block_legend import BlockSpec


if TYPE_CHECKING:
    from dl_constants.types import TBIDataValue
    from dl_query_processing.compilation.filter_compiler import FilterFormulaCompiler


LOGGER = logging.getLogger(__name__)


class DatasetView(DatasetBaseWrapper):
    """
    A ``Dataset`` wrapper for performing data queries.

    The whole process can be broken down into the following steps:

    1. Formalization.
        Performed by subclasses of ``QuerySpecFormalizerBase``.

        Various lists defining the future query such as
        ``select_ids``, ``group_by_ids``, ``order_by_specs``, etc.
        are normalized (see more detailed info about that in formalizer descriptions)
        and combined into a single object, namely ``QuerySpec``.

        Result: ``QuerySpec`` instance

    2. Compilation.
        Performed by subclasses of ``RawQueryCompilerBase``;
        ``FormulaCompiler`` and ``FilterFormulaCompiler`` subclasses
        are used internally.

        Query compilation starts with a ``QuerySpec`` instance as input.
        The query spec tells the compiler which fields are to be used
        in which parts of the query. The compiler generates formula objects
        for each item in the spec with all the respective wrappers
        (custom operations for filters, ASC/DESC for ordering)
        and combines them in a ``CompiledQuery`` object,
        which somewhat resembles the structure of the original ``QuerySpec`` object,
        as will the other query representations in the subsequent steps.

        Result: ``CompiledQuery`` instance

    3. Planning.
        Performed by subclasses of ``ExecutionPlanner``.

        All parts of the compiled query are analyzed,
        as a result of which the following happens:
        - The total number of execution levels (nested sub-queries) is determined.
        - For each expression its top execution level is determined.
          It may be different for each formula (e.g. some filters are applied at source_db level,
          whole others - at various compeng levels).
        - For each expression a slicer configuration is gnerated.
          It will determine how it is cut up into parts for these execution levels.

        An example is described in detail in ``MainExecutionPlanner``

        Result: ``ExecutionPlan`` instance

    4. Slicing.
        Performed by ``DefaultQuerySlicer``.

        Every part of the query is cut into several formulas for each execution level
        using the configurations from the previous step (Planning).

        The resulting object still resembles the query classes from all previous steps
        with formulas organized into sections corresponding to the clauses
        of an SQL query. But it is worth noting that not all of these sections
        will be applied to the top-level query. Each "formula" in its every section
        (select, group_by, where, etc.) is really a collection of:
        - the top-level formula that will be applied to its corresponding clause
        - aliased lower level formulas that will go to the SELECT clause of querues
          executed at lower levels

        Result: ``SlicedQuery`` instance

    5. Separation.
        Performed by ``QuerySeparator``.

        The sliced query is separated into what will actually represent
        the real queries to be performed at each execution level.
        Multiple queries per execution level;
        multiple execution levels in the resulting object
        (as many as was determined at the Planning stage).

        Result: ``CompiledMultiLevelQuery`` instance

    6. Translation.
        TODO

    7. Execution.
        TODO

    8. Postprocessing.
        TODO
    """

    _SOURCE_DB_PROCESSOR_TYPE = ProcessorType.SOURCE_DB
    _COMPENG_PROCESSOR_TYPE = ProcessorType.ASYNCPG
    _SELECTOR_TYPE = SelectorType.CACHED_LAZY

    _verbose_logging = True

    def __init__(
        self,
        ds: Dataset,
        *,
        us_manager: USManagerBase,
        block_spec: BlockSpec,
        rci: RequestContextInfo,
    ):
        self._rci = rci
        self._query_type = block_spec.query_type  # FIXME: Remove
        self._compeng_semaphore = asyncio.Semaphore()
        super().__init__(
            ds=ds, block_spec=block_spec, us_manager=us_manager, debug_mode=self._rci.x_dl_debug_mode or False
        )

    def make_formalizer(self, query_type: QueryType) -> DataQuerySpecFormalizer:
        query_form_cls: Type[DataQuerySpecFormalizer]
        if query_type == QueryType.value_range:
            query_form_cls = ValueRangeSpecFormalizer
        elif query_type == QueryType.distinct:
            query_form_cls = ValueDistinctSpecFormalizer
        elif query_type == QueryType.totals:
            query_form_cls = TotalsSpecFormalizer
        else:
            query_form_cls = DataQuerySpecFormalizer
        return query_form_cls(
            verbose_logging=self._verbose_logging,
            dataset=self._ds,
            rci=self._rci,
            role=self.resolve_role(),
            dep_mgr_factory=self._ds_dep_mgr_factory,
            us_entry_buffer=self._us_manager.get_entry_buffer(),
            service_registry=self._service_registry,
        )

    def make_filter_compiler(self) -> FilterFormulaCompiler:
        backend_type = self.get_backend_type()
        filter_cmp_cls = get_filter_formula_compiler_cls(backend_type=backend_type)
        return filter_cmp_cls(formula_compiler=self.formula_compiler)

    def make_query_executor(self, allow_cache_usage: bool) -> QueryExecutorBase:
        return QueryExecutor(
            dataset=self._ds,
            avatar_alias_mapper=self._avatar_alias_mapper,
            compeng_processor_type=self._COMPENG_PROCESSOR_TYPE,
            source_db_processor_type=self._SOURCE_DB_PROCESSOR_TYPE,
            selector_type=self._SELECTOR_TYPE,
            allow_cache_usage=allow_cache_usage,
            us_manager=self._us_manager,
            compeng_semaphore=self._compeng_semaphore,
        )

    @property
    def is_preview(self) -> bool:
        return self._query_type == QueryType.preview

    def resolve_role(self) -> DataSourceRole:
        return self._capabilities.resolve_source_role(for_preview=self.is_preview)

    def build_exec_info(self) -> QueryExecutionInfo:
        LOGGER.info(f"Select field IDs: {[spec.field_id for spec in self.query_spec.select_specs]}")

        assert (
            self._formula_compiler is not None and self.inspect_env is not None
        ), "perhaps the sources were not reloaded properly"

        translated_multi_query = self.compile_and_translate_query(query_spec=self.query_spec)

        role = self.resolve_role()

        target_connections: list[ClassicConnectionSQL] = []
        for avatar_id in translated_multi_query.get_base_root_from_ids():
            avatar = self._ds_accessor.get_avatar_strict(avatar_id=avatar_id)
            dsrc = self._get_data_source_strict(source_id=avatar.source_id, role=role)
            connection = dsrc.connection
            if connection not in target_connections:
                target_connections.append(connection)

        return QueryExecutionInfo(
            role=role,
            translated_multi_query=translated_multi_query,
            target_connections=target_connections,
        )

    def fast_get_expression_value_range(self) -> tuple[BIField, TBIDataValue, TBIDataValue]:
        """Try to get fast (cached or pre-defined) range from data source"""

        assert len(self.query_spec.select_specs) == 2, "Fast value range is supported only for single field (MIN+MAX)"
        select_field_ids = {item_spec.field_id for item_spec in self.query_spec.select_specs}
        assert len(select_field_ids) == 1, "Got more than one field ID for range"
        field_id = next(iter(select_field_ids))
        field = self._ds.result_schema.by_guid(field_id)
        if field.calc_mode != CalcMode.direct:
            return field, None, None

        avatar_id = field.avatar_id
        if avatar_id is None:
            return field, None, None

        source_ids = self._ds_accessor.get_data_source_id_list()
        if len(source_ids) != 1:
            return field, None, None

        avatar = self._ds_accessor.get_avatar_strict(avatar_id=avatar_id)
        dsrc = self._get_data_source_strict(source_id=avatar.source_id, role=self.resolve_role())

        min_value, max_value = dsrc.get_expression_value_range(field.source)

        return field, min_value, max_value

    async def get_data_async(
        self,
        exec_info: QueryExecutionInfo,
        allow_cache_usage: bool = True,
    ) -> ExecutedQuery:
        executor = self.make_query_executor(
            allow_cache_usage=allow_cache_usage,
        )
        return await executor.execute_query(exec_info)

lib/dl_api_lib/dl_api_lib/dataset/view.py
datalens-tech/datalens-backend
"""
Only locally used non-API enums
"""

from __future__ import annotations

from enum import (
    Enum,
    unique,
)
from typing import Set

from dl_constants.enums import (
    UserDataType,
    WhereClauseOperation,
)
from dl_constants.enums import AggregationFunction as ag


W = WhereClauseOperation
_FILT_MINIMAL = frozenset((W.EQ, W.NE, W.IN, W.NIN))
_FILT_EQNEQ = frozenset((W.EQ, W.NE))
_FILT_COMP = _FILT_EQNEQ | frozenset((W.GT, W.GTE, W.LT, W.LTE))
_FILT_NULL = frozenset((W.ISNULL, W.ISNOTNULL))
_FILT_COMMON = _FILT_COMP | frozenset((W.IN, W.NIN, W.BETWEEN))
_FILT_STR = frozenset(
    (
        W.STARTSWITH,
        W.ISTARTSWITH,
        W.ENDSWITH,
        W.IENDSWITH,
        W.CONTAINS,
        W.ICONTAINS,
        W.NOTCONTAINS,
        W.NOTICONTAINS,
    )
)
_FILT_ARRAY = (
    _FILT_NULL
    | _FILT_EQNEQ
    | frozenset(
        (
            W.STARTSWITH,
            W.CONTAINS,
            W.NOTCONTAINS,
            W.LENEQ,
            W.LENNE,
            W.LENGT,
            W.LENGTE,
            W.LENLT,
            W.LENLTE,
        )
    )
)

FILTERS_BY_TYPE = {
    UserDataType.datetime: _FILT_NULL | _FILT_COMMON,
    UserDataType.datetimetz: _FILT_NULL | _FILT_COMMON,
    UserDataType.genericdatetime: _FILT_NULL | _FILT_COMMON,
    UserDataType.date: _FILT_NULL | _FILT_COMMON,
    UserDataType.boolean: _FILT_NULL | _FILT_MINIMAL,
    UserDataType.integer: _FILT_NULL | _FILT_COMMON | _FILT_STR,
    UserDataType.float: _FILT_NULL | _FILT_COMMON,
    UserDataType.string: _FILT_NULL | _FILT_COMMON | _FILT_STR,
    UserDataType.geopoint: _FILT_NULL | _FILT_MINIMAL,
    UserDataType.geopolygon: _FILT_NULL | _FILT_MINIMAL,
    UserDataType.uuid: _FILT_NULL | _FILT_COMMON | _FILT_STR,
    # intentionally not supporting comparison between markups:
    # it has too much dependency on internal optimizations.
    UserDataType.markup: _FILT_NULL,
    UserDataType.unsupported: _FILT_NULL,
    UserDataType.array_float: _FILT_ARRAY,
    UserDataType.array_int: _FILT_ARRAY,
    UserDataType.array_str: _FILT_ARRAY,
    UserDataType.tree_str: _FILT_ARRAY,
}

CASTS_BY_TYPE = {
    UserDataType.datetime: [
        UserDataType.genericdatetime,
        UserDataType.datetime,
        UserDataType.date,
        UserDataType.string,
        UserDataType.integer,
        UserDataType.float,
        UserDataType.boolean,
    ],
    UserDataType.datetimetz: [
        UserDataType.datetimetz,
        UserDataType.genericdatetime,
        UserDataType.date,
        UserDataType.string,
        UserDataType.integer,
        UserDataType.float,
        UserDataType.boolean,
    ],
    UserDataType.genericdatetime: [
        UserDataType.genericdatetime,
        UserDataType.date,
        UserDataType.string,
        UserDataType.integer,
        UserDataType.float,
        UserDataType.boolean,
    ],
    UserDataType.date: [
        UserDataType.date,
        UserDataType.genericdatetime,
        UserDataType.string,
        UserDataType.integer,
        UserDataType.float,
        UserDataType.boolean,
    ],
    UserDataType.boolean: [UserDataType.boolean, UserDataType.integer, UserDataType.float, UserDataType.string],
    UserDataType.integer: [
        UserDataType.integer,
        UserDataType.float,
        UserDataType.string,
        UserDataType.boolean,
        UserDataType.date,
        UserDataType.genericdatetime,
    ],
    UserDataType.float: [
        UserDataType.float,
        UserDataType.integer,
        UserDataType.string,
        UserDataType.boolean,
        UserDataType.date,
        UserDataType.genericdatetime,
    ],
    UserDataType.string: [
        UserDataType.string,
        UserDataType.date,
        UserDataType.genericdatetime,
        UserDataType.boolean,
        UserDataType.integer,
        UserDataType.float,
        UserDataType.geopoint,
        UserDataType.geopolygon,
    ],
    UserDataType.geopoint: [UserDataType.geopoint, UserDataType.string],
    UserDataType.geopolygon: [UserDataType.geopolygon, UserDataType.string],
    UserDataType.uuid: [UserDataType.uuid, UserDataType.string],
    UserDataType.markup: [UserDataType.markup],
    UserDataType.unsupported: [UserDataType.unsupported],  # `, UserDataType.string` would be too implicit.
    UserDataType.array_float: [UserDataType.array_float, UserDataType.string],
    UserDataType.array_int: [UserDataType.array_int, UserDataType.string],
    UserDataType.array_str: [UserDataType.array_str, UserDataType.string],
    UserDataType.tree_str: [UserDataType.tree_str],
}

_AGG_BASIC = [ag.none, ag.count]
BI_TYPE_AGGREGATIONS = {
    UserDataType.string: _AGG_BASIC + [ag.countunique],
    UserDataType.integer: _AGG_BASIC + [ag.sum, ag.avg, ag.min, ag.max, ag.countunique],
    UserDataType.float: _AGG_BASIC + [ag.sum, ag.avg, ag.min, ag.max, ag.countunique],
    UserDataType.date: _AGG_BASIC + [ag.min, ag.max, ag.countunique, ag.avg],
    UserDataType.datetime: _AGG_BASIC + [ag.min, ag.max, ag.countunique, ag.avg],
    UserDataType.datetimetz: _AGG_BASIC + [ag.min, ag.max, ag.countunique, ag.avg],
    UserDataType.genericdatetime: _AGG_BASIC + [ag.min, ag.max, ag.countunique],  # TODO: 'avg'?
    UserDataType.boolean: _AGG_BASIC + [],
    UserDataType.geopoint: _AGG_BASIC + [],
    UserDataType.geopolygon: _AGG_BASIC + [],
    UserDataType.uuid: _AGG_BASIC + [ag.countunique],
    UserDataType.markup: _AGG_BASIC + [],  # TODO: 'any'
    UserDataType.unsupported: [ag.none],  # only explicit formula-based processing is allowed
    UserDataType.array_float: _AGG_BASIC + [ag.countunique],
    UserDataType.array_int: _AGG_BASIC + [ag.countunique],
    UserDataType.array_str: _AGG_BASIC + [ag.countunique],
    UserDataType.tree_str: _AGG_BASIC + [ag.countunique],
}


@unique
class DatasetAction(Enum):
    # field
    update_field = "update_field"
    add_field = "add_field"
    delete_field = "delete_field"
    clone_field = "clone_field"

    # source
    update_source = "update_source"
    add_source = "add_source"
    delete_source = "delete_source"
    refresh_source = "refresh_source"

    # source avatar
    update_source_avatar = "update_source_avatar"
    add_source_avatar = "add_source_avatar"
    delete_source_avatar = "delete_source_avatar"

    # source relation
    update_avatar_relation = "update_avatar_relation"
    add_avatar_relation = "add_avatar_relation"
    delete_avatar_relation = "delete_avatar_relation"

    # connection actions
    replace_connection = "replace_connection"

    # filters
    update_obligatory_filter = "update_obligatory_filter"
    add_obligatory_filter = "add_obligatory_filter"
    delete_obligatory_filter = "delete_obligatory_filter"

    # TODO: remove legacy:
    update = "update"
    add = "add"
    delete = "delete"

    @staticmethod
    def remap_legacy(action: "DatasetAction") -> "DatasetAction":
        return _LEGACY_ACTIONS.get(action, action)

    @staticmethod
    def get_actions_whitelist_for_data_api() -> Set["DatasetAction"]:
        return {DatasetAction.add_field, DatasetAction.update_field, DatasetAction.delete_field}


_LEGACY_ACTIONS = {
    DatasetAction.add: DatasetAction.add_field,
    DatasetAction.update: DatasetAction.update_field,
    DatasetAction.delete: DatasetAction.delete_field,
}


@unique
class USPermissionKind(Enum):
    execute = "execute"
    read = "read"
    edit = "edit"
    admin = "admin"

lib/dl_api_lib/dl_api_lib/enums.py
datalens-tech/datalens-backend
from __future__ import annotations

import http
import inspect
import itertools
import logging
from typing import (
    Any,
    ClassVar,
    Optional,
    Type,
)

import attr
from marshmallow import Schema
from marshmallow import ValidationError as MValidationError
from marshmallow import fields

from dl_api_commons.exc import RequestTimeoutError
from dl_api_lib import exc
from dl_constants.exc import (
    DEFAULT_ERR_CODE_API_PREFIX,
    GLOBAL_ERR_PREFIX,
)
from dl_core import exc as common_exc
from dl_formula.core import exc as formula_exc
import dl_query_processing.exc


LOGGER = logging.getLogger(__name__)

status = http.HTTPStatus
EXCEPTION_CODES = {
    dl_query_processing.exc.LogicError: status.BAD_REQUEST,
    dl_query_processing.exc.DatasetError: status.BAD_REQUEST,
    formula_exc.FormulaError: status.BAD_REQUEST,
    common_exc.DatabaseQueryError: status.BAD_REQUEST,
    common_exc.SourceAccessDenied: status.FORBIDDEN,
    common_exc.DataSourceErrorFromComponentError: status.BAD_REQUEST,
    common_exc.DatabaseUnavailable: status.BAD_REQUEST,
    common_exc.DatasetConfigurationError: status.BAD_REQUEST,
    common_exc.ConnectionConfigurationError: status.BAD_REQUEST,
    common_exc.USAccessDeniedException: status.FORBIDDEN,
    common_exc.USObjectNotFoundException: status.NOT_FOUND,
    common_exc.USPermissionCheckError: 530,
    common_exc.USLockUnacquiredException: status.LOCKED,
    common_exc.USBadRequestException: status.BAD_REQUEST,
    common_exc.USIncorrectTenantIdException: status.CONFLICT,
    common_exc.QueryConstructorError: status.BAD_REQUEST,
    common_exc.ResultRowCountLimitExceeded: status.BAD_REQUEST,
    common_exc.TableNameNotConfiguredError: status.BAD_REQUEST,
    common_exc.NotAvailableError: status.BAD_REQUEST,
    common_exc.InvalidFieldError: status.BAD_REQUEST,
    common_exc.FieldNotFound: status.BAD_REQUEST,
    common_exc.USPermissionRequired: status.FORBIDDEN,
    exc.RLSConfigParsingError: status.BAD_REQUEST,
    common_exc.RLSSubjectNotFound: status.BAD_REQUEST,
    exc.FeatureNotAvailable: status.BAD_REQUEST,
    dl_query_processing.exc.FilterError: status.BAD_REQUEST,
    exc.UnsupportedForEntityType: status.BAD_REQUEST,
    common_exc.SourceAvatarNotFound: status.BAD_REQUEST,
    common_exc.ReferencedUSEntryNotFound: status.BAD_REQUEST,
    common_exc.ReferencedUSEntryAccessDenied: status.FORBIDDEN,
    common_exc.DataSourceTitleConflict: status.BAD_REQUEST,
    common_exc.DataSourcesInconsistent: status.BAD_REQUEST,
    dl_query_processing.exc.EmptyQuery: status.BAD_REQUEST,
    dl_query_processing.exc.InvalidQueryStructure: status.BAD_REQUEST,
    common_exc.PlatformPermissionRequired: status.FORBIDDEN,
    exc.DatasetRevisionMismatch: status.BAD_REQUEST,
    common_exc.EntityUsageNotAllowed: status.BAD_REQUEST,
    exc.DatasetActionNotAllowedError: status.BAD_REQUEST,
    dl_query_processing.exc.LegendError: status.BAD_REQUEST,
    dl_query_processing.exc.PivotError: status.BAD_REQUEST,
    common_exc.TypeCastFailed: status.BAD_REQUEST,
    dl_query_processing.exc.BlockSpecError: status.BAD_REQUEST,
    dl_query_processing.exc.TreeError: status.BAD_REQUEST,
    dl_query_processing.exc.ParameterError: status.BAD_REQUEST,
    common_exc.DashSQLError: status.BAD_REQUEST,
    dl_query_processing.exc.GenericInvalidRequestError: status.BAD_REQUEST,
    dl_query_processing.exc.InvalidGroupByConfiguration: status.BAD_REQUEST,
    common_exc.WrongQueryParameterization: status.BAD_REQUEST,
    RequestTimeoutError: status.FAILED_DEPENDENCY,
}


# TODO CONSIDER: Add private fields to store in logs for further quantity analytics
@attr.s(frozen=True, auto_attribs=True)
class BIError:
    """
    Class that represents errors occurred during BI requests handling.
     Is used for metrics collection and forming API error responses.
    """

    http_code: Optional[int]
    application_code_stack: tuple[str, ...]

    message: str  # treated as safe to display to final user
    details: dict
    debug: dict

    DEFAULT_ERROR_MESSAGE = "Internal Server Error"

    @staticmethod
    def get_default_error_code(
        err: Exception, exc_code_mapping: Optional[dict[Type[Exception], int]] = None
    ) -> Optional[int]:
        """
        :param err: Exception to map to HTTP status code
        :param exc_code_mapping: Override for default `EXCEPTION_CODES` exception to HTTP status code mapping
        :return: Mapped HTTP status code
        """

        if exc_code_mapping is None:
            exc_code_mapping = EXCEPTION_CODES

        for err_cls in inspect.getmro(type(err)):
            if err_cls in exc_code_mapping:
                return exc_code_mapping[err_cls]

        return None

    @classmethod
    def from_exception(
        cls,
        ex: Exception,
        default_message: Optional[str] = None,
        exc_code_mapping: Optional[dict[Type[Exception], int]] = None,
    ) -> "BIError":
        """
        Creates BIError from exception
        :param ex: Exception to create BIError from
        :param default_message: Message for exceptions that are not in whitelist
        :param exc_code_mapping: Override for default exception to HTTP status code mapping
        :return:
        """
        if default_message is None:
            default_message = cls.DEFAULT_ERROR_MESSAGE

        message: str = default_message
        details: dict[str, Any] = {}
        debug: dict[str, Any] = {}
        application_code_stack: tuple[str, ...] = ()
        http_code: Optional[int] = cls.get_default_error_code(ex, exc_code_mapping)

        if isinstance(ex, common_exc.DLBaseException):
            message = ex.message
            details = ex.details
            debug = ex.debug_info
            application_code_stack = tuple(ex.err_code)

        elif isinstance(ex, formula_exc.FormulaError):
            if not isinstance(ex, dl_query_processing.exc.FormulaHandlingError):
                # FIXME: Temporary measure. Later we'll switch to catching this error instead of FormulaError
                LOGGER.exception("Caught a non-wrapped FormulaError")

            if len(ex.errors) == 1:
                # There is only one error, the pass the code directly
                application_code_stack = ex.errors[0].code or tuple(dl_query_processing.exc.DLFormulaError.err_code)
            else:
                application_code_stack = tuple(dl_query_processing.exc.DLFormulaError.err_code)
            message = str(ex)

        elif isinstance(ex, MValidationError):
            message = str(ex)
            http_code = status.BAD_REQUEST

        else:
            # Any unknown exceptions leaves default values
            pass

        return BIError(
            http_code=http_code,
            application_code_stack=application_code_stack,
            message=message,
            details=details,
            debug=debug,
        )


class RegularAPIErrorSchema(Schema):
    code = fields.Method(serialize="serialize_error_code")
    message = fields.String()

    details = fields.Dict()  # In future will be replaced with schemas for each exception
    debug = fields.Dict()

    def get_api_prefix(self) -> str:
        # TODO CONSIDER: Warning in case of no api_prefix in context
        return self.context.get("api_prefix") or DEFAULT_ERR_CODE_API_PREFIX

    def serialize_error_code(self, data: BIError) -> str:
        return ".".join(
            itertools.chain(
                (
                    GLOBAL_ERR_PREFIX,
                    self.get_api_prefix(),
                ),
                data.application_code_stack,
            )
        )


class PublicAPIErrorSchema(RegularAPIErrorSchema):
    class Meta(RegularAPIErrorSchema.Meta):
        PUBLIC_FORWARDED_ERROR_CODES: ClassVar[frozenset[tuple[str, ...]]] = frozenset(
            (tuple(common_exc.MaterializationNotFinished.err_code),)
        )
        PUBLIC_DEFAULT_MESSAGE = "Something went wrong"
        PUBLIC_DEFAULT_ERR_CODE = "ERR.UNKNOWN"

    message = fields.Method(serialize="serialize_message")  # type: ignore  # TODO: fix

    debug = fields.Constant({})  # type: ignore  # TODO: fix
    details = fields.Constant({})  # type: ignore  # TODO: fix

    def serialize_error_code(self, data: BIError) -> str:
        if tuple(data.application_code_stack) in self.Meta.PUBLIC_FORWARDED_ERROR_CODES:
            return super().serialize_error_code(data)

        return self.Meta.PUBLIC_DEFAULT_ERR_CODE

    def serialize_message(self, data: BIError) -> str:
        if tuple(data.application_code_stack) in self.Meta.PUBLIC_FORWARDED_ERROR_CODES:
            return data.message

        else:
            return self.Meta.PUBLIC_DEFAULT_MESSAGE

lib/dl_api_lib/dl_api_lib/error_handling.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.exc import DLBaseException


class FeatureNotAvailable(DLBaseException):
    err_code = DLBaseException.err_code + ["FEATURE_NOT_AVAILABLE"]


class RLSError(DLBaseException):
    err_code = DLBaseException.err_code + ["RLS"]


class RLSConfigParsingError(RLSError):
    err_code = RLSError.err_code + ["PARSE"]


class DatasetActionNotAllowedError(DLBaseException):
    err_code = DLBaseException.err_code + ["ACTION_NOT_ALLOWED"]


class UnsupportedForEntityType(DLBaseException):
    err_code = DLBaseException.err_code + ["UNSUPPORTED"]
    default_message = "This entity type does not support this operation"


class DatasetRevisionMismatch(DLBaseException):
    err_code = DLBaseException.err_code + ["DATASET_REVISION_MISMATCH"]
    default_message = "Dataset version mismatch. Refresh the page to continue."


class _DLValidationResult(DLBaseException):
    err_code = DLBaseException.err_code + ["VALIDATION"]
    default_message = ""


class DLValidationError(_DLValidationResult):
    err_code = _DLValidationResult.err_code + ["ERROR"]
    default_message = "Validation finished with errors."


class TooManyFieldsError(DLValidationError):
    err_code = DLValidationError.err_code + ["TOO_MANY_FIELDS"]


class DLValidationFatal(_DLValidationResult):
    err_code = _DLValidationResult.err_code + ["FATAL"]
    default_message = "Validation encountered a fatal error."

lib/dl_api_lib/dl_api_lib/exc.py
datalens-tech/datalens-backend
import os

import attr

import dl_api_lib as package
from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_api_lib/dl_api_lib/i18n/localizer.py
datalens-tech/datalens-backend
from typing import Iterable

from dl_api_connector.i18n.localizer import CONFIGS as BI_API_CONNECTOR_CONFIG
from dl_api_lib.i18n.localizer import CONFIGS as BI_API_LIB_CONFIGS
from dl_core.i18n.localizer import CONFIGS as BI_CORE_CONFIGS
from dl_i18n.localizer_base import TranslationConfig


LOCALIZATION_CONFIGS: set[TranslationConfig] = set()


def register_translation_configs(config: Iterable[TranslationConfig]) -> None:
    LOCALIZATION_CONFIGS.update(config)


register_translation_configs(BI_API_LIB_CONFIGS)
register_translation_configs(BI_CORE_CONFIGS)
register_translation_configs(BI_API_CONNECTOR_CONFIG)

lib/dl_api_lib/dl_api_lib/i18n/registry.py
datalens-tech/datalens-backend
from typing import (
    Collection,
    Optional,
)

import attr

from dl_api_lib.app_connectors import register_all_connectors
from dl_core.loader import (
    CoreLibraryConfig,
    load_core_lib,
    preload_bi_core,
)
from dl_formula.loader import (
    FormulaLibraryConfig,
    load_formula_lib,
    preload_bi_formula,
)


@attr.s(frozen=True)
class ApiLibraryConfig:
    formula_lib_config: FormulaLibraryConfig = attr.ib(kw_only=True, default=FormulaLibraryConfig())
    core_lib_config: CoreLibraryConfig = attr.ib(kw_only=True, default=CoreLibraryConfig())
    api_connector_ep_names: Optional[Collection[str]] = attr.ib(kw_only=True, default=None)


def preload_api_lib() -> None:
    """Loads all entrypoint connectors without registering them, e.g. to trigger constants declaration"""

    preload_bi_core()
    preload_bi_formula()


def load_api_lib(api_lib_config: ApiLibraryConfig = ApiLibraryConfig()) -> None:
    load_formula_lib(formula_lib_config=api_lib_config.formula_lib_config)
    load_core_lib(core_lib_config=api_lib_config.core_lib_config)
    register_all_connectors(connector_ep_names=api_lib_config.api_connector_ep_names)

lib/dl_api_lib/dl_api_lib/loader.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from collections import defaultdict
from copy import deepcopy
from itertools import chain
from operator import itemgetter
from typing import (
    TYPE_CHECKING,
    Any,
    Generator,
    TypeVar,
)

import attr

from dl_api_lib.pivot.primitives import PivotHeaderInfo
from dl_constants.enums import (
    FieldRole,
    PivotHeaderRole,
)


if TYPE_CHECKING:
    from dl_api_lib.pivot.primitives import (
        MeasureValues,
        PivotHeader,
    )
    from dl_query_processing.legend.field_legend import Legend


_PIVOT_DATA_FRAME_TV = TypeVar("_PIVOT_DATA_FRAME_TV", bound="PivotDataFrame")


@attr.s
class PivotDataFrame(abc.ABC):
    headers_info: dict[tuple, PivotHeaderInfo] = attr.ib(default=attr.Factory(lambda: defaultdict(PivotHeaderInfo)))

    @abc.abstractmethod
    def iter_columns(self) -> Generator[PivotHeader, None, None]:
        raise NotImplementedError

    @abc.abstractmethod
    def iter_row_headers(self) -> Generator[PivotHeader, None, None]:
        raise NotImplementedError

    @abc.abstractmethod
    def iter_rows(self) -> Generator[tuple[PivotHeader, MeasureValues], None, None]:
        raise NotImplementedError

    def clone(self: _PIVOT_DATA_FRAME_TV, **kwargs: Any) -> _PIVOT_DATA_FRAME_TV:
        kwargs["headers_info"] = kwargs.pop("headers_info", deepcopy(self.headers_info))
        return attr.evolve(self, **kwargs)

    @abc.abstractmethod
    def get_column_count(self) -> int:
        raise NotImplementedError

    @abc.abstractmethod
    def get_row_count(self) -> int:
        raise NotImplementedError

    def populate_headers_info(self, legend: Legend) -> None:
        total_liids = {item.legend_item_id for item in legend.list_for_role(FieldRole.template)}
        for header in chain(self.iter_columns(), map(itemgetter(0), self.iter_rows())):
            if any(value.main_legend_item_id in total_liids for value in header.values):
                header.info.role_spec.role = PivotHeaderRole.total

lib/dl_api_lib/dl_api_lib/pivot/base/data_frame.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Any,
    Generator,
    Optional,
    TypeVar,
)

import attr

from dl_api_lib.pivot.primitives import (
    DataCell,
    DataCellVector,
)
from dl_constants.enums import PivotRole


if TYPE_CHECKING:
    from dl_api_lib.pivot.base.data_frame import PivotDataFrame
    from dl_api_lib.pivot.base.paginator import PivotPaginator
    from dl_api_lib.pivot.base.sorter import PivotSorter
    from dl_api_lib.pivot.primitives import (
        MeasureValues,
        PivotHeader,
    )
    from dl_api_lib.query.formalization.pivot_legend import PivotLegend
    from dl_query_processing.legend.field_legend import Legend


_FACADE_TV = TypeVar("_FACADE_TV", bound="TableDataFacade")


@attr.s
class TableDataFacade(abc.ABC):
    _legend: Legend = attr.ib(kw_only=True)
    _pivot_legend: PivotLegend = attr.ib(kw_only=True)
    _pivot_dframe: PivotDataFrame = attr.ib(init=False)
    _sorter: PivotSorter = attr.ib(init=False)
    _paginator: PivotPaginator = attr.ib(init=False)

    def __attrs_post_init__(self) -> None:
        self._pivot_dframe = self._make_pivot_dframe()
        self._pivot_dframe.populate_headers_info(self._legend)
        self._sorter = self._make_sorter()
        self._paginator = self._make_paginator()

    @property
    def pivot_dframe(self) -> PivotDataFrame:
        return self._pivot_dframe

    @abc.abstractmethod
    def _make_pivot_dframe(self) -> PivotDataFrame:
        raise NotImplementedError

    @abc.abstractmethod
    def _make_sorter(self) -> PivotSorter:
        raise NotImplementedError

    @abc.abstractmethod
    def _make_paginator(self) -> PivotPaginator:
        raise NotImplementedError

    def iter_columns(self) -> Generator[PivotHeader, None, None]:
        return self._pivot_dframe.iter_columns()

    def iter_row_dim_headers(self) -> Generator[DataCellVector, None, None]:
        dname_legend_item_id = next(iter(self._legend.get_dimension_name_legend_item_ids()))
        dname_pivot_item_id = next(iter(self._pivot_legend.get_dimension_name_pivot_item_ids()))
        for row_dim_item in self._pivot_legend.list_for_role(PivotRole.pivot_row):
            yield DataCellVector(
                cells=(
                    DataCell(
                        value=row_dim_item.title,
                        legend_item_id=dname_legend_item_id,
                        pivot_item_id=dname_pivot_item_id,
                    ),
                )
            )

    def iter_rows(self) -> Generator[tuple[PivotHeader, MeasureValues], None, None]:
        return self._pivot_dframe.iter_rows()

    def get_column_count(self) -> int:
        return self._pivot_dframe.get_column_count()

    def get_row_count(self) -> int:
        return self._pivot_dframe.get_row_count()

    def clone(self: _FACADE_TV, **kwargs: Any) -> _FACADE_TV:
        return attr.evolve(self, **kwargs)

    def sort(self) -> None:
        self._sorter.sort()

    def paginate(self, offset_rows: Optional[int], limit_rows: Optional[int]) -> None:
        self._pivot_dframe = self._paginator.paginate(
            pivot_dframe=self._pivot_dframe,
            offset_rows=offset_rows,
            limit_rows=limit_rows,
        )
        self._sorter = self._make_sorter()

lib/dl_api_lib/dl_api_lib/pivot/base/facade.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Optional,
)


if TYPE_CHECKING:
    from dl_api_lib.pivot.base.data_frame import PivotDataFrame


class PivotPaginator(abc.ABC):
    @abc.abstractmethod
    def paginate(
        self,
        pivot_dframe: PivotDataFrame,
        limit_rows: Optional[int] = None,
        offset_rows: Optional[int] = None,
    ) -> PivotDataFrame:
        raise NotImplementedError

lib/dl_api_lib/dl_api_lib/pivot/base/paginator.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import TYPE_CHECKING

import attr

from dl_api_lib.pivot.sort_strategy import (
    DimensionSortValueStrategy,
    MeasureSortStrategy,
)


if TYPE_CHECKING:
    from dl_api_lib.pivot.base.data_frame import PivotDataFrame
    from dl_api_lib.pivot.sort_strategy import SortValueStrategy
    from dl_api_lib.query.formalization.pivot_legend import PivotLegend
    from dl_query_processing.legend.field_legend import Legend


@attr.s
class PivotSorter(abc.ABC):
    _pivot_dframe: PivotDataFrame = attr.ib(kw_only=True)
    _legend: Legend = attr.ib(kw_only=True)
    _pivot_legend: PivotLegend = attr.ib(kw_only=True)
    _dimension_sort_strategy: SortValueStrategy = attr.ib(kw_only=True)
    _measure_sort_strategy: SortValueStrategy = attr.ib(kw_only=True)

    @_dimension_sort_strategy.default
    def _make_dimension_sort_strategy(self) -> SortValueStrategy:
        return DimensionSortValueStrategy(legend=self._legend, pivot_legend=self._pivot_legend)

    @_measure_sort_strategy.default
    def _make_measure_sort_strategy(self) -> SortValueStrategy:
        return MeasureSortStrategy(legend=self._legend, pivot_legend=self._pivot_legend)

    @abc.abstractmethod
    def sort(self) -> None:
        """
        Order by row and column dimensions. Specs of other fields will be ignored
        Extract ASC/DESC directions from the `OrderByFieldSpec` items.
        `Measure Name` pseudo-dimension is not ordered,
        measure order from the legend is preserved.
        """
        raise NotImplementedError

lib/dl_api_lib/dl_api_lib/pivot/base/sorter.py
datalens-tech/datalens-backend
import abc
from typing import Iterable

import attr

from dl_api_lib.pivot.hashable_packing import (
    FastJsonHashableValuePacker,
    HashableValuePackerBase,
)
from dl_api_lib.pivot.table import PivotTable
from dl_api_lib.query.formalization.pivot_legend import PivotLegend
from dl_query_processing.legend.field_legend import Legend
from dl_query_processing.merging.primitives import MergedQueryDataRow


@attr.s
class PivotTransformer(abc.ABC):
    _legend: Legend = attr.ib(kw_only=True)
    _pivot_legend: PivotLegend = attr.ib(kw_only=True)
    _cell_packer: HashableValuePackerBase = attr.ib(kw_only=True, factory=FastJsonHashableValuePacker)

    @abc.abstractmethod
    def pivot(self, rows: Iterable[MergedQueryDataRow]) -> PivotTable:
        """
        Turns raw data stream into a pivot table (``PivotTable``) instance
        composed of ``DataCell`` items. Each cell is a 2-tuple containing
        the actual value and a legend reference index that can be used
        to figure out, what properties this value possesses (data type,
        field id, etc.).
        """

        raise NotImplementedError

lib/dl_api_lib/dl_api_lib/pivot/base/transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Generator,
    TypeVar,
)

import attr

from dl_api_lib.pivot.base.data_frame import PivotDataFrame


if TYPE_CHECKING:
    from dl_api_lib.pivot.primitives import (
        MeasureValues,
        PivotHeader,
    )


_PIVOT_DATA_FRAME_TV = TypeVar("_PIVOT_DATA_FRAME_TV", bound="EmptyPivotDataFrame")


@attr.s
class EmptyPivotDataFrame(PivotDataFrame):
    def iter_columns(self) -> Generator[PivotHeader, None, None]:
        yield from ()

    def iter_row_headers(self) -> Generator[PivotHeader, None, None]:
        yield from ()

    def iter_rows(self) -> Generator[tuple[PivotHeader, MeasureValues], None, None]:
        yield from ()

    def clone(self: _PIVOT_DATA_FRAME_TV, **kwargs: Any) -> _PIVOT_DATA_FRAME_TV:
        return attr.evolve(self, **kwargs)

    def get_column_count(self) -> int:
        return 0

    def get_row_count(self) -> int:
        return 0

lib/dl_api_lib/dl_api_lib/pivot/empty/data_frame.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_api_lib.pivot.base.facade import TableDataFacade
from dl_api_lib.pivot.empty.data_frame import EmptyPivotDataFrame
from dl_api_lib.pivot.empty.paginator import EmptyPivotPaginator
from dl_api_lib.pivot.empty.sorter import EmptyPivotSorter


@attr.s
class EmptyDataFrameFacade(TableDataFacade):
    def _make_pivot_dframe(self) -> EmptyPivotDataFrame:
        return EmptyPivotDataFrame()

    def _make_sorter(self) -> EmptyPivotSorter:
        return EmptyPivotSorter(
            legend=self._legend,
            pivot_legend=self._pivot_legend,
            pivot_dframe=self._pivot_dframe,
        )

    def _make_paginator(self) -> EmptyPivotPaginator:
        return EmptyPivotPaginator()

lib/dl_api_lib/dl_api_lib/pivot/empty/facade.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Optional,
)

from dl_api_lib.pivot.base.paginator import PivotPaginator
from dl_api_lib.pivot.empty.data_frame import EmptyPivotDataFrame


if TYPE_CHECKING:
    from dl_api_lib.pivot.base.data_frame import PivotDataFrame


class EmptyPivotPaginator(PivotPaginator):
    def paginate(
        self,
        pivot_dframe: PivotDataFrame,
        limit_rows: Optional[int] = None,
        offset_rows: Optional[int] = None,
    ) -> PivotDataFrame:
        assert isinstance(pivot_dframe, EmptyPivotDataFrame)
        return pivot_dframe

lib/dl_api_lib/dl_api_lib/pivot/empty/paginator.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_api_lib.pivot.base.sorter import PivotSorter


@attr.s
class EmptyPivotSorter(PivotSorter):
    def sort(self) -> None:
        pass

lib/dl_api_lib/dl_api_lib/pivot/empty/sorter.py
datalens-tech/datalens-backend
import abc
import json
from typing import (
    Any,
    Hashable,
    NamedTuple,
    get_args,
)

from dl_query_processing.postprocessing.primitives import PostprocessedValue


class HashableValuePackerBase(abc.ABC):
    """
    Base class for value packers.
    Used for packing unhashable `PostprocessedValue` values into `DataCell` objects.
    """

    @abc.abstractmethod
    def pack(self, value: PostprocessedValue) -> Hashable:
        """Turn any value into a hashable one"""

        raise NotImplementedError

    @abc.abstractmethod
    def unpack(self, hashable_value: Hashable) -> PostprocessedValue:
        """Restore original value from custom hashable"""

        raise NotImplementedError


class JsonWrapper(NamedTuple):
    """
    Simple wrapper for JSON-packed values
    to make them distinguishable from regular strings.
    """

    json: str


class JsonHashableValuePacker(HashableValuePackerBase):
    """Packs `dict`s into JSON string wrappers. Unpacks by loading JSON"""

    def pack(self, value: PostprocessedValue) -> Hashable:
        if isinstance(value, (dict, list)):
            return JsonWrapper(json=json.dumps(value, sort_keys=True))

        return value

    def unpack(self, hashable_value: Hashable) -> PostprocessedValue:
        value = hashable_value
        if isinstance(hashable_value, JsonWrapper):
            value = json.loads(hashable_value.json)

        assert isinstance(value, get_args(PostprocessedValue))
        return value  # type: ignore  # (see assert above)


class HashableWrapper:
    """
    Wrapper that makes unhashable objects hashable.
    Contains:
    - `value` - original value;
    - `hashable` - a hashable lossless representation.

    Neither `NamedTuple` nor `attr.s` objects work well with customized hashing,
    so using just a plain slotted object.
    """

    __slots__ = ("value", "hashable")

    def __init__(self, value: Any, hashable: Hashable):
        self.value = value
        self.hashable = hashable

    def __hash__(self) -> int:
        return hash(self.hashable)

    def _norm_other_hashable(self, other: Any) -> Any:
        if isinstance(other, HashableWrapper):
            return other.hashable
        return other

    # Ordering operations for use by the pivoting logic

    def __gt__(self, other: Any) -> bool:
        return self.hashable > self._norm_other_hashable(other)

    def __ge__(self, other: Any) -> bool:
        return self.hashable >= self._norm_other_hashable(other)

    def __lt__(self, other: Any) -> bool:
        return self.hashable < self._norm_other_hashable(other)

    def __le__(self, other: Any) -> bool:
        return self.hashable <= self._norm_other_hashable(other)

    def __eq__(self, other: Any) -> bool:
        return self.hashable == self._norm_other_hashable(other)

    def __ne__(self, other: Any) -> bool:
        return self.hashable != self._norm_other_hashable(other)


class FastJsonHashableValuePacker(HashableValuePackerBase):
    """
    Packs `dict`s into wrappers containing a JSON representation
    and also the original value (for fast unpacking).
    Unpacks by returning the saved original value.
    """

    def pack(self, value: PostprocessedValue) -> Hashable:
        if isinstance(value, (dict, list)):
            return HashableWrapper(value=value, hashable=json.dumps(value, sort_keys=True))

        return value

    def unpack(self, hashable_value: Hashable) -> PostprocessedValue:
        value = hashable_value
        if isinstance(hashable_value, HashableWrapper):
            value = hashable_value.value

        assert isinstance(value, get_args(PostprocessedValue))
        return value  # type: ignore  # (see assert above)

lib/dl_api_lib/dl_api_lib/pivot/hashable_packing.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Generator,
    Optional,
)

import attr
import pandas as pd

from dl_api_lib.pivot.base.data_frame import PivotDataFrame
from dl_api_lib.pivot.primitives import (
    DataCellVector,
    MeasureValues,
    PivotHeader,
)


@attr.s
class PdPivotDataFrame(PivotDataFrame):
    """
    A simple wrapper for `pandas.DataFrame`
    that simplifies iteration over columns and rows a little.
    """

    _pd_df: pd.DataFrame = attr.ib(kw_only=True)

    @property
    def pd_df(self) -> pd.DataFrame:
        return self._pd_df

    def iter_columns(self) -> Generator[PivotHeader, None, None]:
        for values in self._pd_df.columns:
            if isinstance(values, DataCellVector):
                # Single column dimension, so normalize
                values = (values,)

            yield PivotHeader(values=values, info=self.headers_info[values])

    def iter_row_headers(self) -> Generator[PivotHeader, None, None]:
        for values in self._pd_df.index:
            if isinstance(values, DataCellVector):
                # Single row dimension, so normalize
                values = (values,)

            yield PivotHeader(values=values, info=self.headers_info[values])

    def iter_rows(self) -> Generator[tuple[PivotHeader, MeasureValues], None, None]:
        headers = self.iter_row_headers()
        for row in self._pd_df.itertuples(index=False):
            values: tuple[Optional[DataCellVector], ...] = tuple(
                (cell if isinstance(cell, DataCellVector) else None) for cell in row
            )
            yield next(headers), values

    def get_column_count(self) -> int:
        return len(self._pd_df.columns)

    def get_row_count(self) -> int:
        return len(self._pd_df.index)


@attr.s
class PdSeriesPivotDataFrameBase(PivotDataFrame):
    """
    A base for DataFrame wrappers for `pandas.Series`
    """

    _pd_series: pd.Series = attr.ib(kw_only=True)

    @property
    def pd_series(self) -> pd.Series:
        return self._pd_series

    def _iter_headers(self) -> Generator[PivotHeader, None, None]:
        for values in self._pd_series.index:
            if isinstance(values, DataCellVector):
                # Single column dimension, so normalize
                values = (values,)

            yield PivotHeader(values=values, info=self.headers_info[values])

    def _get_values(self) -> MeasureValues:
        values: tuple[Optional[DataCellVector], ...] = tuple(
            (value if isinstance(value, DataCellVector) else None) for value in self._pd_series.values
        )
        return values


@attr.s
class PdHSeriesPivotDataFrame(PdSeriesPivotDataFrameBase):
    """
    A horizontal `pandas.Series`
    """

    def iter_columns(self) -> Generator[PivotHeader, None, None]:
        yield from self._iter_headers()

    def iter_row_headers(self) -> Generator[PivotHeader, None, None]:
        yield PivotHeader()

    def iter_rows(self) -> Generator[tuple[PivotHeader, MeasureValues], None, None]:
        yield PivotHeader(), self._get_values()

    def get_column_count(self) -> int:
        return len(self._pd_series)

    def get_row_count(self) -> int:
        return 1


@attr.s
class PdVSeriesPivotDataFrame(PdSeriesPivotDataFrameBase):
    """
    A vertical `pandas.Series`
    """

    def iter_columns(self) -> Generator[PivotHeader, None, None]:
        yield PivotHeader()

    def iter_row_headers(self) -> Generator[PivotHeader, None, None]:
        yield from self._iter_headers()

    def iter_rows(self) -> Generator[tuple[PivotHeader, MeasureValues], None, None]:
        for header, value in zip(self._iter_headers(), self._get_values()):
            yield header, (value,)

    def get_column_count(self) -> int:
        return 1

    def get_row_count(self) -> int:
        return len(self._pd_series)

lib/dl_api_lib/dl_api_lib/pivot/pandas/data_frame.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr
import pandas as pd

from dl_api_lib.pivot.base.facade import TableDataFacade
from dl_api_lib.pivot.pandas.data_frame import (
    PdHSeriesPivotDataFrame,
    PdPivotDataFrame,
    PdVSeriesPivotDataFrame,
)
from dl_api_lib.pivot.pandas.paginator import (
    PdHSeriesPivotPaginator,
    PdPivotPaginator,
    PdVSeriesPivotPaginator,
)
from dl_api_lib.pivot.pandas.sorter import (
    PdHSeriesPivotSorter,
    PdPivotSorter,
    PdVSeriesPivotSorter,
)


@attr.s
class PdDataFrameFacade(TableDataFacade):
    _pd_df: pd.DataFrame = attr.ib(kw_only=True)

    def _make_pivot_dframe(self) -> PdPivotDataFrame:
        return PdPivotDataFrame(pd_df=self._pd_df)

    def _make_sorter(self) -> PdPivotSorter:
        return PdPivotSorter(
            legend=self._legend,
            pivot_legend=self._pivot_legend,
            pivot_dframe=self._pivot_dframe,
        )

    def _make_paginator(self) -> PdPivotPaginator:
        return PdPivotPaginator()


@attr.s
class PdSeriesDataFrameFacadeBase(TableDataFacade):
    _pd_series: pd.Series = attr.ib(kw_only=True)


class PdHSeriesDataFrameFacade(PdSeriesDataFrameFacadeBase):
    def _make_pivot_dframe(self) -> PdHSeriesPivotDataFrame:
        return PdHSeriesPivotDataFrame(pd_series=self._pd_series)

    def _make_sorter(self) -> PdHSeriesPivotSorter:
        return PdHSeriesPivotSorter(
            legend=self._legend,
            pivot_legend=self._pivot_legend,
            pivot_dframe=self._pivot_dframe,
        )

    def _make_paginator(self) -> PdHSeriesPivotPaginator:
        return PdHSeriesPivotPaginator()


@attr.s
class PdVSeriesDataFrameFacade(PdSeriesDataFrameFacadeBase):
    def _make_pivot_dframe(self) -> PdVSeriesPivotDataFrame:
        return PdVSeriesPivotDataFrame(pd_series=self._pd_series)

    def _make_sorter(self) -> PdVSeriesPivotSorter:
        return PdVSeriesPivotSorter(
            legend=self._legend,
            pivot_legend=self._pivot_legend,
            pivot_dframe=self._pivot_dframe,
        )

    def _make_paginator(self) -> PdVSeriesPivotPaginator:
        return PdVSeriesPivotPaginator()

lib/dl_api_lib/dl_api_lib/pivot/pandas/facade.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Optional,
)

from dl_api_lib.pivot.base.paginator import PivotPaginator
from dl_api_lib.pivot.pandas.data_frame import (
    PdHSeriesPivotDataFrame,
    PdPivotDataFrame,
    PdVSeriesPivotDataFrame,
)


if TYPE_CHECKING:
    from dl_api_lib.pivot.base.data_frame import PivotDataFrame


class PdPivotPaginator(PivotPaginator):
    def paginate(
        self,
        pivot_dframe: PivotDataFrame,
        limit_rows: Optional[int] = None,
        offset_rows: Optional[int] = None,
    ) -> PivotDataFrame:
        assert isinstance(pivot_dframe, PdPivotDataFrame)
        pd_df = pivot_dframe.pd_df
        total_rows = len(pd_df.index)

        if offset_rows is not None and offset_rows != 0:
            assert offset_rows > 0
            tail_size = total_rows - min(total_rows, offset_rows)
            pd_df = pd_df.tail(tail_size)
            total_rows = tail_size

        if limit_rows is not None:
            assert limit_rows > 0
            pd_df = pd_df.head(min(total_rows, limit_rows))

        # Create table copy with truncated data frame
        pivot_dframe = pivot_dframe.clone(pd_df=pd_df)
        return pivot_dframe


class PdHSeriesPivotPaginator(PivotPaginator):
    def paginate(
        self,
        pivot_dframe: PivotDataFrame,
        limit_rows: Optional[int] = None,
        offset_rows: Optional[int] = None,
    ) -> PivotDataFrame:
        assert isinstance(pivot_dframe, PdHSeriesPivotDataFrame)
        pd_series = pivot_dframe.pd_series
        if limit_rows is not None:
            assert limit_rows > 1
        if offset_rows is not None and offset_rows > 0:
            assert pivot_dframe.get_row_count() == 1
            pd_series = pd_series.transform(lambda x: None)  # type: ignore

        pivot_dframe = pivot_dframe.clone(pd_series=pd_series)
        return pivot_dframe


class PdVSeriesPivotPaginator(PivotPaginator):
    def paginate(
        self,
        pivot_dframe: PivotDataFrame,
        limit_rows: Optional[int] = None,
        offset_rows: Optional[int] = None,
    ) -> PivotDataFrame:
        assert isinstance(pivot_dframe, PdVSeriesPivotDataFrame)
        pd_series = pivot_dframe.pd_series
        total_rows = len(pd_series.index)

        if offset_rows is not None and offset_rows != 0:
            assert offset_rows > 0
            tail_size = total_rows - min(total_rows, offset_rows)
            pd_series = pd_series.tail(tail_size)
            total_rows = tail_size

        if limit_rows is not None:
            assert limit_rows > 0
            pd_series = pd_series.head(min(total_rows, limit_rows))

        # Create table copy with truncated data frame
        pivot_dframe = pivot_dframe.clone(pd_series=pd_series)
        return pivot_dframe

lib/dl_api_lib/dl_api_lib/pivot/pandas/paginator.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from enum import (
    Enum,
    auto,
    unique,
)
from itertools import count
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Generator,
    Optional,
    Sequence,
    Union,
    cast,
)

import pandas as pd

from dl_api_lib.pivot.base.sorter import PivotSorter
from dl_api_lib.pivot.pandas.data_frame import (
    PdHSeriesPivotDataFrame,
    PdPivotDataFrame,
    PdVSeriesPivotDataFrame,
)
from dl_api_lib.pivot.primitives import DataCellVector
from dl_api_lib.query.formalization.pivot_legend import (
    PivotDimensionRoleSpec,
    PivotMeasureRoleSpec,
)
from dl_constants.enums import (
    OrderDirection,
    PivotHeaderRole,
    PivotRole,
)
import dl_query_processing.exc as exc


if TYPE_CHECKING:
    from dl_api_lib.pivot.primitives import (
        PivotHeader,
        PivotMeasureSortingSettings,
    )


@unique
class SortAxis(Enum):
    columns = auto()
    rows = auto()


_PD_AXIS_MAP = {
    SortAxis.columns: 1,
    SortAxis.rows: 0,
}


class PdPivotSorterBase(PivotSorter):
    """
    pd.DataFrame-specific sorter.
    """

    @abc.abstractmethod
    def _get_pd_obj(self) -> Union[pd.DataFrame, pd.Series]:
        raise NotImplementedError

    def _resolve_axis_order(self) -> dict[SortAxis, list[OrderDirection]]:
        return {
            SortAxis.columns: [
                cast(PivotDimensionRoleSpec, f.role_spec).direction
                for f in self._pivot_legend.list_for_role(role=PivotRole.pivot_column)
            ],
            SortAxis.rows: [
                cast(PivotDimensionRoleSpec, f.role_spec).direction
                for f in self._pivot_legend.list_for_role(role=PivotRole.pivot_row)
            ],
        }

    @staticmethod
    def _complementary_axis(axis: SortAxis) -> SortAxis:
        return next(iter(set(SortAxis) - {axis}))

    def _get_pd_index(self, axis: SortAxis) -> Generator[PivotHeader, None, None]:
        if axis == SortAxis.columns:
            return self._pivot_dframe.iter_columns()
        return self._pivot_dframe.iter_row_headers()

    def _has_total(self, axis: SortAxis) -> bool:
        index = self._get_pd_index(axis)
        total_count = sum(header.info.role_spec.role == PivotHeaderRole.total for header in index)
        if total_count > 1:
            raise exc.PivotSortingWithSubtotalsIsNotAllowed()
        return bool(total_count)

    @abc.abstractmethod
    def _get_pd_axis(self, axis: SortAxis) -> int:
        raise NotImplementedError

    @abc.abstractmethod
    def _get_sorting_key(self, sorting_idx: int, axis: SortAxis) -> pd.Series:
        raise NotImplementedError

    @abc.abstractmethod
    def _apply_sorting_key(self, axis: SortAxis, sorting_key: pd.Series) -> None:
        raise NotImplementedError

    def _single_axis_sort(self, axis: SortAxis, directions: Sequence[OrderDirection]) -> None:
        ascending_list = [direction == OrderDirection.asc for direction in directions]
        ascending: Union[bool, list[bool]] = ascending_list  # Separate var to avoid typing errors
        if len(ascending_list) == 1:
            # If there is only one dimension, then it must be a scalar
            ascending = ascending_list[0]

        idx_iter = count()  # see the `dimension_idx` function below

        def dimension_idx() -> int:
            """
            Returns the index of the current dimension along the current axis.
            The first call returns 0, the second - 1, and so on.
            """
            return next(idx_iter)

        def sorting_key_func(idx: pd.Index) -> pd.Index:
            """Convert cell vectors in data frame index to simple orderable values"""
            idx_item = idx[0]
            assert isinstance(idx_item, DataCellVector)
            legend_item_id = idx_item.main_legend_item_id
            pivot_item_ids = self._pivot_legend.leg_item_id_to_pivot_item_id_list(legend_item_id=legend_item_id)
            # FIXME: Find a way to determine which of the several possible pivot items this really is
            pivot_item_id = next(iter(pivot_item_ids))
            dim_idx = dimension_idx()
            direction = directions[dim_idx]
            normalizer = self._dimension_sort_strategy.get_normalizer(
                pivot_item_id=pivot_item_id,
                direction=direction,
            )

            return pd.Index([normalizer.normalize_vector_value(vector) for vector in idx])

        self._get_pd_obj().sort_index(
            axis=self._get_pd_axis(axis),
            ascending=ascending,
            inplace=True,
            key=sorting_key_func,
        )

    def _sort_by_measure(self, axis: SortAxis, sorting_piid: int, settings: PivotMeasureSortingSettings) -> None:
        sorting_idx: Optional[int] = None

        for idx, header in enumerate(self._get_pd_index(axis)):
            if header.compare_sorting_settings(settings):
                if sorting_idx is None:
                    sorting_idx = idx
                    header.info.sorting_direction = settings.direction
                else:  # should never actually occur, as header_values + role_spec uniquely identify sorting_idx
                    raise exc.PivotSortingRowOrColumnIsAmbiguous()
        if sorting_idx is None:
            raise exc.PivotSortingRowOrColumnNotFound()

        # we can use sort_values to sort a DataFrame by row or column,
        # however tuples in indexes make it incredibly hard to apply,
        # so instead we find a row/column by index, get a sorting permutation
        # using argsort and apply it to the DataFrame
        sorting_key = self._get_sorting_key(sorting_idx, axis)
        normalizer = self._measure_sort_strategy.get_normalizer(
            pivot_item_id=sorting_piid, direction=settings.direction
        )
        sorting_key = sorting_key.map(normalizer.normalize_vector_value).argsort()
        if settings.direction == OrderDirection.desc:
            sorting_key = sorting_key[::-1]
        if self._has_total(self._complementary_axis(axis)):
            # manually put the total at the end
            key_len = len(sorting_key)
            total_pos = list(sorting_key).index(key_len - 1)
            order = list(range(key_len))
            order.append(order.pop(total_pos))
            sorting_key = sorting_key[order]
        self._apply_sorting_key(axis, sorting_key)

    def sort(self) -> None:
        # sort by measure; overrides provide additional sorting by dimensions
        for pivot_item in self._pivot_legend.list_for_role(role=PivotRole.pivot_measure):
            sorting_settings = cast(PivotMeasureRoleSpec, pivot_item.role_spec).sorting
            if sorting_settings is None:
                continue
            for axis, settings in zip(SortAxis, [sorting_settings.column, sorting_settings.row]):
                if settings is not None:
                    self._sort_by_measure(axis, pivot_item.pivot_item_id, settings)


class PdPivotSorter(PdPivotSorterBase):
    """
    pd.DataFrame-specific sorter.
    """

    def _get_pd_obj(self) -> pd.DataFrame:
        assert isinstance(self._pivot_dframe, PdPivotDataFrame)
        return self._pivot_dframe.pd_df

    def _get_pd_axis(self, axis: SortAxis) -> int:
        return _PD_AXIS_MAP[axis]

    def _get_sorting_key(self, sorting_idx: int, axis: SortAxis) -> pd.Series:
        if axis == SortAxis.columns:
            return self._get_pd_obj().iloc[:, sorting_idx]
        return self._get_pd_obj().iloc[sorting_idx]

    def _apply_sorting_key(self, axis: SortAxis, sorting_key: pd.Series) -> None:
        if axis == SortAxis.columns:
            new_df = self._get_pd_obj().iloc[sorting_key]  # ordering by column, so reorder rows
        else:
            new_df = self._get_pd_obj().iloc[:, sorting_key]  # ordering by row, so reorder columns
        self._pivot_dframe._pd_df = new_df  # type: ignore

    def sort(self) -> None:
        directions_by_axis = self._resolve_axis_order()
        for axis in (SortAxis.rows, SortAxis.columns):
            if not directions_by_axis[axis]:
                # Nothing to sort here
                continue
            self._single_axis_sort(axis=axis, directions=directions_by_axis[axis])
        super().sort()


class PdSeriesPivotSorterBase(PdPivotSorterBase):
    _AXIS: ClassVar[SortAxis]

    def _get_pd_obj(self) -> pd.Series:
        assert isinstance(self._pivot_dframe, (PdHSeriesPivotDataFrame, PdVSeriesPivotDataFrame))
        return self._pivot_dframe.pd_series  # type: ignore

    def _get_pd_axis(self, axis: SortAxis) -> int:
        return 0

    def _get_sorting_key(self, sorting_idx: int, axis: SortAxis) -> pd.Series:
        assert axis == self._AXIS
        # returning a sentinel value, as measure sorting in case of series is trivial
        return pd.Series()

    def _apply_sorting_key(self, axis: SortAxis, sorting_key: pd.Series) -> None:
        assert axis == self._AXIS
        # nothing to do here: sorting by one value is trivial

    def sort(self) -> None:
        directions_by_axis = self._resolve_axis_order()
        axis = self._AXIS
        assert not directions_by_axis[self._complementary_axis(axis)]
        self._single_axis_sort(axis=axis, directions=directions_by_axis[axis])
        super().sort()


class PdHSeriesPivotSorter(PdSeriesPivotSorterBase):
    _AXIS = SortAxis.columns


class PdVSeriesPivotSorter(PdSeriesPivotSorterBase):
    _AXIS = SortAxis.rows

lib/dl_api_lib/dl_api_lib/pivot/pandas/sorter.py
datalens-tech/datalens-backend
from collections import defaultdict
import logging
from typing import (
    TYPE_CHECKING,
    Dict,
    Iterable,
    List,
    Optional,
    Sequence,
    Union,
)

import attr
import pandas as pd

from dl_api_lib.pivot.base.transformer import PivotTransformer
from dl_api_lib.pivot.empty.facade import EmptyDataFrameFacade
from dl_api_lib.pivot.pandas.facade import (
    PdDataFrameFacade,
    PdHSeriesDataFrameFacade,
    PdVSeriesDataFrameFacade,
)
from dl_api_lib.pivot.primitives import (
    DataCell,
    DataCellVector,
)
from dl_api_lib.pivot.stream_modifiers import (
    DataCellConverter,
    MeasureDataTransposer,
)
from dl_api_lib.pivot.table import PivotTable
from dl_constants.enums import PivotRole
import dl_query_processing.exc
from dl_query_processing.merging.primitives import MergedQueryDataRow


if TYPE_CHECKING:
    from dl_api_lib.pivot.base.facade import TableDataFacade
    from dl_api_lib.pivot.stream_modifiers import TransposedDataRow


LOGGER = logging.getLogger(__name__)


@attr.s
class PdPivotTransformer(PivotTransformer):
    """Created pandas-based pivot table"""

    def pivot(self, rows: Iterable[MergedQueryDataRow]) -> PivotTable:
        # Prepare data stream
        raw_dcell_stream: Iterable[Sequence[DataCell]] = DataCellConverter(
            rows=rows,
            cell_packer=self._cell_packer,
            legend=self._legend,
            pivot_legend=self._pivot_legend,
        )
        # Transpose measures (turn multiple measure fields into a single pseudo-field).
        transp_stream: Iterable[TransposedDataRow] = MeasureDataTransposer(
            dcell_stream=raw_dcell_stream,
            pivot_legend=self._pivot_legend,
            cell_packer=self._cell_packer,
        )

        fake_measure_piid_str = str(self._pivot_legend.get_unused_pivot_item_id())

        data: Dict[str, List[Optional[DataCellVector]]] = defaultdict(list)
        for dim_vectors, value_vector in transp_stream:
            for dim_vector in dim_vectors:
                base_dim_cell = dim_vector.cells[0]
                data[str(base_dim_cell.pivot_item_id)].append(dim_vector)
            data[fake_measure_piid_str].append(value_vector)

        any_data = bool(data)

        data_lengths = {len(values) for values in data.values()}
        if len(data_lengths) > 1:
            raise dl_query_processing.exc.PivotUnevenDataColumnsError(
                "Expected all data columns to contain the same number of values. " f"Got: {sorted(data_lengths)}"
            )

        raw_pd_df = pd.DataFrame(data)
        column_ids = [str(item.pivot_item_id) for item in self._pivot_legend.list_for_role(PivotRole.pivot_column)]
        row_ids = [str(item.pivot_item_id) for item in self._pivot_legend.list_for_role(PivotRole.pivot_row)]

        pd_pivot_res: Union[pd.DataFrame, pd.Series]
        facade: TableDataFacade

        pd_series: pd.Series
        pd_df: pd.DataFrame

        def wrapped_pivot(columns: Sequence[str], index: Sequence[str]) -> Union[pd.Series, pd.DataFrame]:
            try:
                return raw_pd_df.pivot(columns=columns, index=index)
            except ValueError as err:
                if str(err).startswith("Index contains duplicate entries"):
                    raise dl_query_processing.exc.PivotDuplicateDimensionValue() from err

                raise

        if not any_data:
            facade = EmptyDataFrameFacade(legend=self._legend, pivot_legend=self._pivot_legend)

        elif column_ids and row_ids:
            pd_pivot_res = wrapped_pivot(columns=column_ids, index=row_ids)
            assert isinstance(pd_pivot_res, pd.DataFrame)
            pd_df = pd_pivot_res
            # Drop the fake dimension `fake_measure_liid_str` from the column index:
            pd_df = pd_df.set_axis(axis="columns", labels=pd_df.columns.droplevel(0), copy=False)
            facade = PdDataFrameFacade(
                pd_df=pd_df,
                legend=self._legend,
                pivot_legend=self._pivot_legend,
            )

        elif column_ids:
            pd_pivot_res = wrapped_pivot(columns=column_ids, index=row_ids)
            assert isinstance(pd_pivot_res, pd.Series)
            pd_series = pd_pivot_res
            # Drop `fake_measure_liid_str` from the index:
            pd_series = pd_series.set_axis(labels=pd_series.index.droplevel(0), copy=False)
            facade = PdHSeriesDataFrameFacade(
                pd_series=pd_series,
                legend=self._legend,
                pivot_legend=self._pivot_legend,
            )

        elif row_ids:
            # Fool pandas into thinking columns are rows and vice-versa.
            # Otherwise it creates a DataFrame with an empty index for columns,
            # which does not behave the way a "normal" DataFrame should
            pd_pivot_res = wrapped_pivot(columns=row_ids, index=column_ids)
            assert isinstance(pd_pivot_res, pd.Series)
            pd_series = pd_pivot_res
            # Drop `fake_measure_liid_str` from the index:
            pd_series = pd_series.set_axis(labels=pd_series.index.droplevel(0), copy=False)
            facade = PdVSeriesDataFrameFacade(
                pd_series=pd_series,
                legend=self._legend,
                pivot_legend=self._pivot_legend,
            )

        else:
            raise TypeError("Invalid pivot configuration")

        LOGGER.info(f"Using {type(facade)} for pivot table")

        table = PivotTable(
            facade=facade,
            pivot_legend=self._pivot_legend,
            cell_packer=self._cell_packer,
        )
        return table

lib/dl_api_lib/dl_api_lib/pivot/pandas/transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Hashable,
    NamedTuple,
    Optional,
)

import attr

from dl_constants.enums import (
    OrderDirection,
    PivotHeaderRole,
)


class DataCell(NamedTuple):
    value: Hashable
    legend_item_id: int
    pivot_item_id: int


class DataCellVector(NamedTuple):
    cells: tuple[DataCell, ...]

    @property
    def main_cell(self) -> DataCell:
        return self.cells[0]

    @property
    def main_value(self) -> Hashable:
        return self[0][0][0]  # self.cells[0].value

    @property
    def main_legend_item_id(self) -> int:
        return self[0][0][1]  # self.cells[0].legend_item_id

    @property
    def main_pivot_item_id(self) -> int:
        return self[0][0][2]  # self.cells[0].pivot_item_id


class MeasureNameValue(NamedTuple):
    value: Hashable  # The actual measure name
    measure_piid: int  # legend item ID of the referenced measure () for handling of duplicate measures in table


@attr.s
class PivotHeaderRoleSpec:
    role: PivotHeaderRole = attr.ib(kw_only=True, default=PivotHeaderRole.data)


@attr.s(frozen=True)
class PivotHeaderValue:
    value: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class PivotMeasureSortingSettings:
    header_values: list[PivotHeaderValue] = attr.ib(kw_only=True)
    direction: OrderDirection = attr.ib(kw_only=True, default=OrderDirection.asc)
    role_spec: PivotHeaderRoleSpec = attr.ib(kw_only=True, factory=PivotHeaderRoleSpec)


@attr.s(frozen=True)
class PivotMeasureSorting:
    column: Optional[PivotMeasureSortingSettings] = attr.ib(kw_only=True, default=None)
    row: Optional[PivotMeasureSortingSettings] = attr.ib(kw_only=True, default=None)


@attr.s(slots=True)
class PivotHeaderInfo:
    sorting_direction: Optional[OrderDirection] = attr.ib(kw_only=True, default=None)
    role_spec: PivotHeaderRoleSpec = attr.ib(kw_only=True, factory=PivotHeaderRoleSpec)


@attr.s(slots=True)
class PivotHeader:
    values: tuple[DataCellVector, ...] = attr.ib(kw_only=True, default=())
    info: PivotHeaderInfo = attr.ib(kw_only=True, factory=PivotHeaderInfo)

    def main_header_values(self) -> list[Hashable]:
        return [value.main_value for value in self.values]

    @staticmethod
    def _normalize_value(value: Hashable) -> Hashable:
        if isinstance(value, MeasureNameValue):
            return value.value
        return value

    def compare_sorting_settings(self, sorting_settings: PivotMeasureSortingSettings) -> bool:
        if self.info.role_spec != sorting_settings.role_spec:
            return False

        main_values = self.main_header_values()
        sorting_values: list[Hashable] = [header_value.value for header_value in sorting_settings.header_values]
        if not main_values:
            return not sorting_values

        last_value = main_values.pop()
        main_values = list(map(self._normalize_value, main_values))
        if isinstance(last_value, MeasureNameValue):
            # the last Measure name value may be hidden, so try to compare without it first
            if sorting_values == main_values:
                return True

            main_values.append(last_value.value)
            return sorting_values == main_values
        main_values.append(last_value)
        return sorting_values == main_values


MeasureValues = tuple[Optional[DataCellVector], ...]


class DataRow(NamedTuple):
    header: PivotHeader
    values: MeasureValues

lib/dl_api_lib/dl_api_lib/pivot/primitives.py
datalens-tech/datalens-backend
"""
Module defines common logic of sorting values in pivot tables.
This is independent from the specific pivot table implementations.
"""

from __future__ import annotations

import abc
from enum import (
    Enum,
    auto,
    unique,
)
from math import isnan
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Optional,
    Union,
)

import attr

from dl_api_lib.pivot.primitives import (
    DataCellVector,
    MeasureNameValue,
)
from dl_constants.enums import (
    FieldRole,
    OrderDirection,
    PivotRole,
    UserDataType,
)


if TYPE_CHECKING:
    from dl_api_lib.query.formalization.pivot_legend import PivotLegend
    from dl_query_processing.legend.field_legend import Legend


@unique
class SortAxis(Enum):
    columns = auto()
    rows = auto()


_PD_AXIS_MAP = {
    SortAxis.columns: 1,
    SortAxis.rows: 0,
}


class OrderableNullValueBase(abc.ABC):
    """
    An orderable, hashable and immutable value that can be used
    instead of ``None`` when sorting something.
    """

    __slots__ = ("weight",)

    weight: int

    def __init__(self, weight: int = 0):
        super().__setattr__("weight", weight)

    def __setattr__(self, key: str, value: Any) -> None:
        raise AttributeError(f"{type(self).__name__} is immutable")

    def __hash__(self) -> int:
        return 0

    def __eq__(self, other: Any) -> bool:
        return type(other) is type(self) and other.weight == self.weight

    def __ne__(self, other: Any) -> bool:
        return type(other) is not type(self) or other.weight != self.weight

    @abc.abstractmethod
    def __gt__(self, other: Any) -> bool:
        raise NotImplementedError()

    @abc.abstractmethod
    def __ge__(self, other: Any) -> bool:
        raise NotImplementedError()

    @abc.abstractmethod
    def __lt__(self, other: Any) -> bool:
        raise NotImplementedError()

    @abc.abstractmethod
    def __le__(self, other: Any) -> bool:
        raise NotImplementedError()


class LeastNullValue(OrderableNullValueBase):
    """A null/None placeholder that is less than everything else."""

    __slots__ = ()

    def __gt__(self, other: Any) -> bool:
        return isinstance(other, type(self)) and self.weight > other.weight

    def __ge__(self, other: Any) -> bool:
        return isinstance(other, type(self)) and self.weight >= other.weight

    def __lt__(self, other: Any) -> bool:
        return not isinstance(other, type(self)) or self.weight < other.weight

    def __le__(self, other: Any) -> bool:
        return not isinstance(other, type(self)) or self.weight <= other.weight


class GreatestNullValue(OrderableNullValueBase):
    """A null/None placeholder that is greater than everything else."""

    __slots__ = ()

    def __gt__(self, other: Any) -> bool:
        return not isinstance(other, type(self)) or self.weight > other.weight

    def __ge__(self, other: Any) -> bool:
        return not isinstance(other, type(self)) or self.weight >= other.weight

    def __lt__(self, other: Any) -> bool:
        return isinstance(other, type(self)) and self.weight < other.weight

    def __le__(self, other: Any) -> bool:
        return isinstance(other, type(self)) and self.weight <= other.weight


NORMAL_LEAST_NULL_VALUE = LeastNullValue()
TOTAL_GREATEST_NULL_VALUE = GreatestNullValue(1)
TOTAL_LEAST_NULL_VALUE = LeastNullValue(-1)


@attr.s
class SortValueNormalizer(abc.ABC):
    _legend: Legend = attr.ib(kw_only=True)
    _pivot_legend: PivotLegend = attr.ib(kw_only=True)
    _pivot_item_id: int = attr.ib(kw_only=True)  # To be replaced with pivot_item_id
    _direction: OrderDirection = attr.ib()

    @abc.abstractmethod
    def normalize_vector_value(self, vector: Union[DataCellVector, float]) -> Any:
        raise NotImplementedError


@attr.s
class BaseSortValueNormalizer(SortValueNormalizer):
    _legend_item_ids: set[int] = attr.ib(init=False)
    _value_converter: Optional[Callable[[Any], Any]] = attr.ib(init=False)

    @_legend_item_ids.default
    def _make_legend_item_ids(self) -> set[int]:
        liid_list = self._pivot_legend.get_item(self._pivot_item_id).legend_item_ids
        legend_items = [self._legend.get_item(legend_item_id) for legend_item_id in liid_list]
        legend_items = [item for item in legend_items if item.role_spec.role != FieldRole.template]
        return {item.legend_item_id for item in legend_items}

    @_value_converter.default
    def _make_value_converter(self) -> Optional[Callable[[Any], Any]]:
        data_types = {self._legend.get_item(legend_item_id).data_type for legend_item_id in self._legend_item_ids}
        assert len(data_types) == 1, "Only single data type is supported within a pivot dimension"
        data_type = next(iter(data_types))
        # Normalize numbers for correct sorting
        if data_type is UserDataType.integer:
            return int
        if data_type is UserDataType.float:
            return float

        return None

    def normalize_vector_value(self, vector: Union[DataCellVector, float]) -> Any:
        # the case of NaN values
        if isinstance(vector, float) and isnan(vector):
            return NORMAL_LEAST_NULL_VALUE

        assert isinstance(vector, DataCellVector)
        value = vector.main_value
        if value is None:
            # Control how `None`/`null` is sorted
            return NORMAL_LEAST_NULL_VALUE

        if self._value_converter is not None:
            return self._value_converter(value)

        if isinstance(value, str):
            # Casefold strings for case-insensitive sorting
            return value.casefold()

        return value


@attr.s
class MeasureSortValueNormalizer(BaseSortValueNormalizer):
    pass


@attr.s
class DimensionSortValueNormalizer(BaseSortValueNormalizer):
    def normalize_vector_value(self, vector: Union[DataCellVector, float]) -> Any:
        assert isinstance(vector, DataCellVector)
        legend_item_id = vector.main_legend_item_id
        if legend_item_id not in self._legend_item_ids:
            # Control the sorting of values that are not part of the main data
            # (e.g. template labels for totals).
            # The data type of these values may be different from the main data.
            if self._direction == OrderDirection.asc:
                return TOTAL_GREATEST_NULL_VALUE
            else:  # desc
                # Sorting is reversed, but totals still must be at the end,
                # even after regular NULLs.
                # So a `LeastNullValue` value with a smaller weight is used.
                return TOTAL_LEAST_NULL_VALUE

        return super().normalize_vector_value(vector)


@attr.s
class MeasureNameSortValueNormalizer(SortValueNormalizer):
    _measure_name_order_values: Dict[MeasureNameValue, int] = attr.ib(init=False)

    @_measure_name_order_values.default
    def _make_measure_name_order_values(self) -> Dict[MeasureNameValue, int]:
        return {
            MeasureNameValue(value=item.title, measure_piid=item.pivot_item_id): order_value
            for order_value, item in enumerate(self._pivot_legend.list_for_role(role=PivotRole.pivot_measure))
        }

    def normalize_vector_value(self, vector: Union[DataCellVector, float]) -> Any:
        assert isinstance(vector, DataCellVector)
        main_value = vector.main_value
        assert isinstance(main_value, MeasureNameValue)
        return self._measure_name_order_values[main_value]


@attr.s
class SortValueStrategy(abc.ABC):
    _legend: Legend = attr.ib(kw_only=True)
    _pivot_legend: PivotLegend = attr.ib(kw_only=True)

    @abc.abstractmethod
    def get_normalizer(self, pivot_item_id: int, direction: OrderDirection) -> SortValueNormalizer:
        """
        Return a `SortValueNormalizer` instance for the specific type of value.
        The normalizer will adapt the value for sorting taking into account its data type,
        legend role and pivot role.
        """
        raise NotImplementedError


@attr.s
class DimensionSortValueStrategy(SortValueStrategy):
    def get_normalizer(self, pivot_item_id: int, direction: OrderDirection) -> SortValueNormalizer:
        measure_name_pivot_item_ids = self._pivot_legend.get_measure_name_pivot_item_ids()
        if pivot_item_id in measure_name_pivot_item_ids:
            return MeasureNameSortValueNormalizer(
                legend=self._legend,
                pivot_legend=self._pivot_legend,
                pivot_item_id=pivot_item_id,
                direction=direction,
            )

        return DimensionSortValueNormalizer(
            legend=self._legend,
            pivot_legend=self._pivot_legend,
            pivot_item_id=pivot_item_id,
            direction=direction,
        )


@attr.s
class MeasureSortStrategy(SortValueStrategy):
    def get_normalizer(self, pivot_item_id: int, direction: OrderDirection) -> SortValueNormalizer:
        return MeasureSortValueNormalizer(
            legend=self._legend,
            pivot_legend=self._pivot_legend,
            pivot_item_id=pivot_item_id,
            direction=direction,
        )

lib/dl_api_lib/dl_api_lib/pivot/sort_strategy.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (  # It thinks NamedTuple isn't used, but it is  # noqa: F401
    TYPE_CHECKING,
    Dict,
    Generator,
    Iterable,
    Iterator,
    List,
    NamedTuple,
    Optional,
    Sequence,
    Set,
    Tuple,
)

import attr

from dl_api_lib.pivot.primitives import (
    DataCell,
    DataCellVector,
    MeasureNameValue,
)
from dl_api_lib.query.formalization.pivot_legend import (
    PivotAnnotationRoleSpec,
    PivotLegend,
)
from dl_constants.enums import (
    FieldType,
    PivotRole,
)
from dl_query_processing.legend.field_legend import (
    FieldObjSpec,
    TemplateRoleSpec,
)
from dl_query_processing.merging.primitives import MergedQueryDataRow


if TYPE_CHECKING:
    from dl_api_lib.pivot.hashable_packing import HashableValuePackerBase
    from dl_query_processing.legend.field_legend import Legend


LOGGER = logging.getLogger(__name__)


@attr.s
class DataCellConverter:
    """
    Converts raw data stream into a stream of ``DataCell`` lists.
    """

    _rows: Iterable[MergedQueryDataRow] = attr.ib()
    _cell_packer: HashableValuePackerBase = attr.ib(kw_only=True)
    _legend: Legend = attr.ib(kw_only=True)
    _pivot_legend: PivotLegend = attr.ib(kw_only=True)
    _legend_item_id_map: dict[int, int] = attr.ib(init=False)

    @_legend_item_id_map.default
    def _make_legend_item_id_map(self) -> dict[int, int]:
        main_field_items_by_field_id: dict[tuple[str, int], int] = {}
        main_template_items_by_template: dict[tuple[Optional[str], int], int] = {}

        legend_item_id_map: dict[int, int] = {}
        for item in self._legend.items:
            for pivot_item_id in self._pivot_legend.leg_item_id_to_pivot_item_id_list(item.legend_item_id):
                if item.field_type == FieldType.MEASURE:
                    continue

                if isinstance(item.obj, FieldObjSpec):
                    field_id = item.obj.id
                    field_item_key = (field_id, pivot_item_id)
                    remapped_legend_item_id = main_field_items_by_field_id.get(field_item_key)
                    if remapped_legend_item_id is not None:
                        legend_item_id_map[item.legend_item_id] = remapped_legend_item_id
                    else:
                        main_field_items_by_field_id[field_item_key] = item.legend_item_id

                elif isinstance(item.role_spec, TemplateRoleSpec):
                    template = item.role_spec.template
                    templ_item_key = (template, pivot_item_id)
                    remapped_legend_item_id = main_template_items_by_template.get(templ_item_key)
                    if remapped_legend_item_id is not None:
                        legend_item_id_map[item.legend_item_id] = remapped_legend_item_id
                    else:
                        main_template_items_by_template[templ_item_key] = item.legend_item_id

        return legend_item_id_map

    def _remap_legend_item_id(self, legend_item_id: int) -> int:
        """
        Re-map legend_item_id for dimension values so that
        same values of the same dimension from different blocks
        fall into the same group during pivoting.
        """
        return self._legend_item_id_map.get(legend_item_id, legend_item_id)

    def _iter_rows(self) -> Generator[Sequence[DataCell], None, None]:
        for row in self._rows:
            dcell_row: list[DataCell] = []
            for value, legend_item_id in zip(row.data, row.legend_item_ids):
                pivot_item_id_list = self._pivot_legend.leg_item_id_to_pivot_item_id_list(legend_item_id)
                for pivot_item_id in pivot_item_id_list:
                    cell = DataCell(
                        value=self._cell_packer.pack(value),
                        legend_item_id=self._remap_legend_item_id(legend_item_id),
                        pivot_item_id=pivot_item_id,
                    )
                    dcell_row.append(cell)

            yield dcell_row

    def __iter__(self) -> Iterator[Sequence[DataCell]]:
        return iter(self._iter_rows())


class TransposedDataRow(NamedTuple):
    dimensions: Tuple[DataCellVector, ...]
    value: Optional[DataCellVector]


@attr.s
class MeasureDataTransposer:
    """
    Transposes measure fields within the stream by
    turning multiple measure fields into two pseudo fields:
    1. measure name (string, referencing the MEASURE_NAME legend item);
    2. measure value (original value referencing the original legend item);
    and multiplying the number of rows by the number of measures.
    """

    _dcell_stream: Iterable[Sequence[DataCell]] = attr.ib()
    _pivot_legend: PivotLegend = attr.ib(kw_only=True)
    _cell_packer: HashableValuePackerBase = attr.ib(kw_only=True)

    # Internal
    _dimension_pivot_item_ids: Set[int] = attr.ib(init=False)
    _measure_name_pivot_item_ids: List[int] = attr.ib(init=False)
    _measure_piid_and_name_mask_transp: Sequence[Tuple[int, str]] = attr.ib(init=False)
    _anno_piid_list_by_measure_liid: Dict[int, List[int]] = attr.ib(init=False)
    _should_add_fake_measure: bool = attr.ib(init=False)
    _measure_name_legend_item_id: Optional[int] = attr.ib(init=False)

    def __attrs_post_init__(self) -> None:
        self._dimension_pivot_item_ids = set()
        for role in (PivotRole.pivot_row, PivotRole.pivot_column):
            for item in self._pivot_legend.list_for_role(role):
                self._dimension_pivot_item_ids.add(item.pivot_item_id)

        self._measure_name_pivot_item_ids = self._pivot_legend.get_measure_name_pivot_item_ids()
        measures = self._pivot_legend.list_for_role(PivotRole.pivot_measure)
        anno_items = self._pivot_legend.list_for_role(PivotRole.pivot_annotation)

        self._measure_piid_and_name_mask_transp = tuple(
            # Transposed mask. Each item will generate a new row
            (item.pivot_item_id, item.title)
            for item in self._pivot_legend.list_for_role(PivotRole.pivot_measure)
        )

        self._measure_name_legend_item_id = None
        if self._measure_name_pivot_item_ids:
            self._measure_name_legend_item_id = self._pivot_legend.get_item(
                self._measure_name_pivot_item_ids[0]
            ).legend_item_ids[0]

        # Direct mask for annotations (on a per-measure basis)
        self._anno_piid_list_by_measure_liid = {}
        _anno_mapping = self._anno_piid_list_by_measure_liid  # shortcut for readability
        for measure_item in measures:
            for measure_legend_item_id in measure_item.legend_item_ids:
                if measure_legend_item_id not in _anno_mapping:
                    _anno_mapping[measure_legend_item_id] = []
                for anno_item in anno_items:
                    anno_role_spec = anno_item.role_spec
                    assert isinstance(anno_role_spec, PivotAnnotationRoleSpec)
                    if (
                        anno_role_spec.target_legend_item_ids is None
                        or measure_legend_item_id in anno_role_spec.target_legend_item_ids
                    ):
                        _anno_mapping[measure_legend_item_id].append(anno_item.pivot_item_id)

        self._should_add_fake_measure = not self._measure_piid_and_name_mask_transp
        if self._should_add_fake_measure:
            LOGGER.info("No measures found in legend, so using a fake one")

    @property
    def measure_name_legend_item_id(self) -> int:
        assert self._measure_name_legend_item_id is not None
        return self._measure_name_legend_item_id

    def _expand_row(self, row: Sequence[DataCell]) -> Generator[TransposedDataRow, None, None]:
        # Apply direct dimension mask
        common_dim_part = tuple(
            DataCellVector(cells=(cell,)) for cell in row if cell[2] in self._dimension_pivot_item_ids
        )
        measure_value_dict: Dict[int, DataCell] = {
            cell[2]: cell for cell in row if cell[2] not in self._dimension_pivot_item_ids  # pivot_item_id: cell
        }

        # Expand transposed measure mask
        for measure_piid, measure_name in self._measure_piid_and_name_mask_transp:
            mnames_dim_part = tuple(
                # There can be multiple `measure_name` items
                DataCellVector(
                    cells=(
                        DataCell(  # the measure name cell
                            value=MeasureNameValue(
                                value=self._cell_packer.pack(measure_name),
                                measure_piid=measure_piid,
                            ),
                            legend_item_id=self.measure_name_legend_item_id,
                            pivot_item_id=measure_name_pivot_item_id,
                        ),
                    )
                )
                for measure_name_pivot_item_id in self._measure_name_pivot_item_ids
            )
            dim_part = common_dim_part + mnames_dim_part
            value_cell = measure_value_dict[measure_piid]
            measure_liid = value_cell.legend_item_id
            anno_cells = (
                measure_value_dict[anno_piid] for anno_piid in self._anno_piid_list_by_measure_liid[measure_liid]
            )  # generator!!

            value_vector = DataCellVector(cells=(value_cell, *anno_cells))
            yield TransposedDataRow(dimensions=dim_part, value=value_vector)

        if self._should_add_fake_measure:
            # There are no measures
            yield TransposedDataRow(dimensions=common_dim_part, value=None)

    def _iter_rows(self) -> Generator[TransposedDataRow, None, None]:
        for row in self._dcell_stream:
            yield from self._expand_row(row)

    def __iter__(self) -> Iterator[TransposedDataRow]:
        """Iterable object interface"""
        return iter(self._iter_rows())

lib/dl_api_lib/dl_api_lib/pivot/stream_modifiers.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Generator,
    Hashable,
)

import attr

from dl_api_lib.pivot.primitives import (
    DataRow,
    MeasureNameValue,
)
from dl_api_lib.query.formalization.pivot_legend import PivotLegend
from dl_query_processing.postprocessing.primitives import PostprocessedValue


if TYPE_CHECKING:
    from dl_api_lib.pivot.base.facade import TableDataFacade
    from dl_api_lib.pivot.hashable_packing import HashableValuePackerBase
    from dl_api_lib.pivot.primitives import (
        DataCellVector,
        PivotHeader,
    )


@attr.s
class PivotTable:
    _facade: TableDataFacade = attr.ib(kw_only=True)
    _pivot_legend: PivotLegend = attr.ib(kw_only=True)
    _cell_packer: HashableValuePackerBase = attr.ib(kw_only=True)

    @property
    def pivot_legend(self) -> PivotLegend:
        return self._pivot_legend

    @property
    def facade(self) -> TableDataFacade:
        return self._facade

    def get_columns(self) -> list[PivotHeader]:
        return [header for header in self.facade.iter_columns()]

    def get_row_dim_headers(self) -> list[DataCellVector]:
        return [header for header in self.facade.iter_row_dim_headers()]

    def get_rows(self) -> Generator[DataRow, None, None]:
        for header, values in self.facade.iter_rows():
            yield DataRow(header=header, values=values)

    def get_column_count(self) -> int:
        return self._facade.get_column_count()

    def get_row_count(self) -> int:
        return self._facade.get_row_count()

    def unpack_value(self, hashable_value: Hashable) -> PostprocessedValue:
        if isinstance(hashable_value, MeasureNameValue):
            hashable_value = hashable_value.value
        return self._cell_packer.unpack(hashable_value)

    def clone(self, **kwargs: Any) -> PivotTable:
        return attr.evolve(self, **kwargs)

lib/dl_api_lib/dl_api_lib/pivot/table.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import typing

import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_constants.enums import DataSourceRole
from dl_core import us_dataset
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.data_source.collection import DataSourceCollectionFactory
from dl_core.dataset_capabilities import DatasetCapabilities
from dl_core.exc import EntityUsageNotAllowed
from dl_core.services_registry.entity_checker import EntityUsageChecker
from dl_core.us_connection_base import ConnectionBase


if typing.TYPE_CHECKING:
    from dl_core.us_manager.us_manager import USManagerBase


LOGGER = logging.getLogger(__name__)


@attr.s
class PublicEnvEntityUsageChecker(EntityUsageChecker):
    def ensure_dataset_can_be_used(
        self,
        rci: RequestContextInfo,
        dataset: us_dataset.Dataset,
        us_manager: USManagerBase,
    ) -> None:
        ds_accessor = DatasetComponentAccessor(dataset=dataset)
        dsrc_coll_factory = DataSourceCollectionFactory(us_entry_buffer=us_manager.get_entry_buffer())
        capabilities = DatasetCapabilities(dataset=dataset, dsrc_coll_factory=dsrc_coll_factory)
        LOGGER.info("Checking if dataset %s can be used in public env", dataset.uuid)
        main_source_role = capabilities.resolve_source_role()
        LOGGER.info("Dataset source role is %s", main_source_role)

        if main_source_role == DataSourceRole.materialization:
            LOGGER.info("Dataset source role is %s. Usage is allowed", main_source_role)
            return
        elif main_source_role == DataSourceRole.origin:
            LOGGER.info("Dataset source role is %s. Underlying connections will be checked", main_source_role)
            dsrc_coll_spec_list = ds_accessor.get_data_source_coll_spec_list()
            dsrc_coll_list = [
                dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
                for dsrc_coll_spec in dsrc_coll_spec_list
            ]
            effective_data_source_list = [dsrc_coll.get_strict() for dsrc_coll in dsrc_coll_list]
            conn_set = {dsrc.connection for dsrc in effective_data_source_list}
            LOGGER.info(
                "Dataset has %d data source, %d connections used",
                len(effective_data_source_list),
                len(conn_set),
            )

            conn_validation_exc_map = {}

            for conn in conn_set:
                try:
                    self.ensure_data_connection_can_be_used(rci, conn)
                except EntityUsageNotAllowed as conn_not_allowed_exc:
                    conn_validation_exc_map[conn.uuid] = conn_not_allowed_exc

            if len(conn_validation_exc_map) == 0:
                LOGGER.info("All connections are allowed to be used in public env")
                return
            else:
                for conn_id, exc in conn_validation_exc_map.items():
                    LOGGER.info("Connection %s is not allowed to be used in public env: %s", conn_id, exc)

                # TODO FIX: Clarify error messages
                raise EntityUsageNotAllowed("Dataset data are not materialized")
        else:
            raise EntityUsageNotAllowed(f"Unexpected data source role resolved for dataset {dataset.uuid}")

    def ensure_data_connection_can_be_used(self, rci: RequestContextInfo, conn: ConnectionBase):  # type: ignore  # TODO: fix
        LOGGER.info("Checking if connection %s %s can be used in public env", conn.conn_type, conn.uuid)
        if conn.allow_public_usage:
            return
        raise EntityUsageNotAllowed(f"Connection type {conn.conn_type.value} can not be used in public dataset")

lib/dl_api_lib/dl_api_lib/public/entity_usage_checker.py
datalens-tech/datalens-backend
import logging
from typing import AbstractSet

from dl_constants.enums import ManagedBy
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.ids import AvatarId
from dl_core.us_dataset import Dataset
import dl_query_processing.exc


LOGGER = logging.getLogger(__name__)


def normalize_explicit_avatar_ids(
    dataset: Dataset, required_avatar_ids: AbstractSet[AvatarId]
) -> AbstractSet[AvatarId]:
    user_managed_required_avatar_ids = set()
    accessor = DatasetComponentAccessor(dataset=dataset)

    for avatar_id in required_avatar_ids:
        avatar = accessor.get_avatar_strict(avatar_id=avatar_id)
        if avatar.managed_by == ManagedBy.user:
            user_managed_required_avatar_ids.add(avatar_id)

    if not user_managed_required_avatar_ids:
        # In this case only one non-user-managed avatar can be handled
        if len(required_avatar_ids) > 1:
            raise dl_query_processing.exc.InvalidQueryStructure("Cannot build query with only feature-managed avatars")

        elif not required_avatar_ids:
            # Attempt to fall back to root avatar
            root_avatar = accessor.get_root_avatar_opt()
            if root_avatar is not None:
                LOGGER.warning(
                    f"Requested to build source without avatars. "
                    f"Should only happen when selecting constants. "
                    f"Falling back to root avatar {root_avatar.id}"
                )
                required_avatar_ids = required_avatar_ids | {root_avatar.id}

    return required_avatar_ids

lib/dl_api_lib/dl_api_lib/query/formalization/avatar_tools.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    List,
    Optional,
    Tuple,
    TypeVar,
)

import attr

from dl_api_commons.reporting.registry import ReportingRegistry
from dl_api_lib.query.formalization.field_resolver import FieldResolver
from dl_api_lib.query.formalization.id_gen import IdGenerator
from dl_api_lib.query.formalization.raw_specs import (
    RawAfterBlockPlacement,
    RawBlockSpec,
    RawQuerySpecUnion,
    RawRootBlockPlacement,
)
from dl_constants.enums import (
    FieldRole,
    FieldType,
    NotificationType,
)
from dl_core.reporting.notifications import get_notification_record
from dl_core.us_dataset import Dataset
from dl_query_processing.enums import (
    EmptyQueryMode,
    QueryType,
)
import dl_query_processing.exc
from dl_query_processing.legend.block_legend import (
    AfterBlockPlacement,
    BlockLegend,
    BlockLegendMeta,
    BlockPlacement,
    BlockSpec,
    DimensionSpec,
    DimensionValueSpec,
    DispersedAfterBlockPlacement,
    RootBlockPlacement,
)
from dl_query_processing.legend.field_legend import (
    Legend,
    TreeRoleSpec,
)


_VAL_TV = TypeVar("_VAL_TV")


def ifnull(value: _VAL_TV, null_value: _VAL_TV) -> _VAL_TV:
    if value is not None:
        return value
    return null_value


@attr.s
class BlockFormalizer(abc.ABC):
    _dataset: Dataset = attr.ib(kw_only=True)
    _reporting_registry: Optional[ReportingRegistry] = attr.ib(kw_only=True, default=None)
    _field_resolver: FieldResolver = attr.ib(init=False)

    def __attrs_post_init__(self) -> None:
        self._field_resolver = FieldResolver(dataset=self._dataset)

    def _gen_legend_and_ids_for_block(self, legend: Legend, block_id: int) -> Tuple[Legend, List[int]]:
        legend_for_block = legend.limit_to_block(block_id=block_id)
        legend_item_ids = [item.legend_item_id for item in legend_for_block.list_streamable_items()]
        return legend_for_block, legend_item_ids

    def _resolve_block_placement_from_raw_spec(self, raw_block_spec: RawBlockSpec) -> BlockPlacement:
        # Make placement
        placement: BlockPlacement
        if isinstance(raw_block_spec.placement, RawRootBlockPlacement):
            placement = RootBlockPlacement()
        elif isinstance(raw_block_spec.placement, RawAfterBlockPlacement):
            dimension_values: Optional[List[DimensionValueSpec]] = None
            if raw_block_spec.placement.dimension_values is not None:
                dimension_values = [
                    DimensionValueSpec(
                        legend_item_id=dim_val_spec.legend_item_id,
                        value=dim_val_spec.value,
                    )
                    for dim_val_spec in raw_block_spec.placement.dimension_values
                ]
            placement = AfterBlockPlacement(dimension_values=dimension_values)
        else:
            raise TypeError(f"Unsupported placement type: {type(raw_block_spec.placement).__name__}")

        return placement

    def _generate_default_block_spec(
        self,
        block_id: int,
        legend: Legend,
        raw_query_spec_union: RawQuerySpecUnion,
        raw_block_spec: Optional[RawBlockSpec] = None,
        main_block: Optional[BlockSpec] = None,
    ) -> BlockSpec:
        legend_for_block, legend_item_ids = self._gen_legend_and_ids_for_block(legend=legend, block_id=block_id)

        streamable_items = legend_for_block.list_streamable_items()

        empty_query_mode = EmptyQueryMode.error
        if streamable_items and all(item.role_spec.role == FieldRole.template for item in streamable_items):
            # Always return an empty row for template-only query
            empty_query_mode = EmptyQueryMode.empty_row

        query_type = raw_query_spec_union.meta.query_type
        if any(item.role_spec.role in (FieldRole.total, FieldRole.template) for item in streamable_items):
            query_type = QueryType.totals

        parent_block_id: Optional[int]
        limit: Optional[int]
        offset: Optional[int]
        row_count_hard_limit: Optional[int]
        placement: BlockPlacement

        if raw_block_spec is not None:
            parent_block_id = raw_block_spec.parent_block_id
            limit = raw_block_spec.limit
            offset = raw_block_spec.offset
            row_count_hard_limit = ifnull(
                raw_block_spec.row_count_hard_limit,
                raw_query_spec_union.meta.row_count_hard_limit,
            )
            # Make placement
            placement = self._resolve_block_placement_from_raw_spec(raw_block_spec)
        else:
            parent_block_id = None
            limit = None
            offset = None
            row_count_hard_limit = raw_query_spec_union.meta.row_count_hard_limit
            placement = self._resolve_block_placement_from_legend(
                legend_for_block,
                main_block=main_block,
                query_type=raw_query_spec_union.meta.query_type,
            )

        return BlockSpec(
            block_id=block_id,
            parent_block_id=parent_block_id,
            placement=placement,
            legend_item_ids=legend_item_ids,
            legend=legend_for_block,
            limit=limit,
            offset=offset,
            query_type=query_type,
            row_count_hard_limit=row_count_hard_limit,
            group_by_policy=raw_query_spec_union.group_by_policy,
            ignore_nonexistent_filters=raw_query_spec_union.ignore_nonexistent_filters,
            disable_rls=raw_query_spec_union.disable_rls,
            allow_measure_fields=raw_query_spec_union.allow_measure_fields,
            empty_query_mode=empty_query_mode,
        )

    def _resolve_block_placement_from_legend(
        self,
        legend_for_block: Legend,
        main_block: Optional[BlockSpec],
        query_type: QueryType,
    ) -> BlockPlacement:
        streamable_items = legend_for_block.list_streamable_items()
        if streamable_items and all(item.role_spec.role == FieldRole.template for item in streamable_items):
            # All items are templates
            # (Most likely a query for totals, but without any measures)
            # The only thing we can do is place it at the end
            return AfterBlockPlacement()

        total_items = legend_for_block.list_for_role(FieldRole.total)
        tree_items = legend_for_block.list_for_role(FieldRole.tree)

        if tree_items and total_items:
            raise dl_query_processing.exc.BlockItemCompatibilityError(
                "Got tree and totals in one block. They are incompatible"
            )

        if total_items and query_type == QueryType.result:
            assert main_block is not None
            # This is a block of totals
            main_legend_item_ids = main_block.legend_item_ids
            # There must be the same number of items as in  the main block
            if len(streamable_items) != len(main_legend_item_ids):
                raise dl_query_processing.exc.UnevenBlockColumnCountError(
                    "Blocks have different column counts. "
                    f"Main block: {len(main_legend_item_ids)}, secondary: {len(streamable_items)}"
                )
            parent_dim_specs: list[DimensionSpec] = []
            child_dim_specs: list[DimensionSpec] = []
            for item_idx, item in enumerate(streamable_items):
                if item.field_type == FieldType.DIMENSION and item.role_spec.role != FieldRole.template:
                    parent_dim_specs.append(DimensionSpec(legend_item_id=main_legend_item_ids[item_idx]))
                    child_dim_specs.append(DimensionSpec(legend_item_id=item.legend_item_id))

            return DispersedAfterBlockPlacement(
                parent_dimensions=parent_dim_specs,
                child_dimensions=child_dim_specs,
            )

        if total_items:  # pivot, maybe others
            # No point in doing complex data merging if it will all be re-mixed in pivot table.
            # So just stick totals at the end.
            return AfterBlockPlacement()

        if tree_items:
            if len(tree_items) > 1:
                raise dl_query_processing.exc.MultipleTreeError(
                    "Multiple trees detected in block. This is not supported"
                )
            # This is a tree block
            tree_item = tree_items[0]
            tree_spec = tree_item.role_spec
            assert isinstance(tree_spec, TreeRoleSpec)
            return AfterBlockPlacement(dimension_values=tree_spec.dimension_values)

        return RootBlockPlacement()

    def validate_block_legend(self, block_legend: BlockLegend) -> None:
        root_cnt = 0
        for block_spec in block_legend.blocks:
            is_root = isinstance(block_spec.placement, RootBlockPlacement)
            if is_root:
                root_cnt += 1

        if root_cnt > 1:
            raise dl_query_processing.exc.MultipleRootBlockError()
        if root_cnt < 1:
            raise dl_query_processing.exc.NoRootBlockError()

    def finalize_block_legend(self, block_legend: BlockLegend) -> BlockLegend:
        if len(block_legend.blocks) > 1:
            first_block, *other_blocks = block_legend.blocks
            if (  # FIXME: separate `QueryType` into `DataRequestType` and `BlockType`
                # Regular main block
                first_block.query_type in {QueryType.result, QueryType.pivot}
                # and all others are (sub)totals
                and all(block.query_type == QueryType.totals for block in other_blocks)
            ):
                filter_items = first_block.legend.list_for_role(FieldRole.filter)
                if any(item.field_type == FieldType.MEASURE for item in filter_items):
                    # And if there are MEASURE filters
                    # (which, for now, we can't support in totals correctly),
                    # then remove totals - that is all blocks except for `first_block`
                    non_total_blocks = [block for block in block_legend.blocks if block.query_type != QueryType.totals]
                    block_legend = block_legend.clone(blocks=non_total_blocks)
                    if self._reporting_registry is not None:
                        self._reporting_registry.save_reporting_record(
                            get_notification_record(NotificationType.totals_removed_due_to_measure_filter)
                        )

        return block_legend

    def make_block_legend(self, raw_query_spec_union: RawQuerySpecUnion, legend: Legend) -> BlockLegend:
        used_ids = set(raw_query_spec_union.get_unique_block_ids())
        id_gen = IdGenerator(used_ids=used_ids)
        blocks: List[BlockSpec] = []

        if not used_ids:
            # There is always at least one block.
            # So if none were explicitly specified, then generate one ID automatically
            id_gen.generate_id()  # It adds to used_ids

        # Create explicitly listed block specs
        for raw_block_spec in raw_query_spec_union.block_specs:
            block_spec = self._generate_default_block_spec(
                block_id=raw_block_spec.block_id,
                raw_block_spec=raw_block_spec,
                legend=legend,
                raw_query_spec_union=raw_query_spec_union,
                main_block=blocks[0] if blocks else None,
            )
            blocks.append(block_spec)

        # Generate block specs not listed explicitly, but referred to in the query legend
        orphaned_ids = used_ids - {block_spec.block_id for block_spec in blocks}
        for block_id in sorted(orphaned_ids):
            block_spec = self._generate_default_block_spec(
                block_id=block_id,
                legend=legend,
                raw_query_spec_union=raw_query_spec_union,
                main_block=blocks[0] if blocks else None,
            )
            blocks.append(block_spec)

        # The above logic ensures that there is at least one block,
        # but to make sure `blocks` is not empty let's
        assert blocks
        # Force the first block to be root
        blocks[0] = blocks[0].clone(placement=RootBlockPlacement())

        block_legend = BlockLegend(
            blocks=blocks,
            meta=BlockLegendMeta(
                limit=raw_query_spec_union.limit,
                offset=raw_query_spec_union.offset,
                row_count_hard_limit=raw_query_spec_union.meta.row_count_hard_limit,
            ),
        )
        block_legend = self.finalize_block_legend(block_legend=block_legend)
        self.validate_block_legend(block_legend=block_legend)
        return block_legend

lib/dl_api_lib/dl_api_lib/query/formalization/block_formalizer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Optional,
)

import attr

from dl_api_lib.query.formalization.raw_specs import (
    IdFieldRef,
    IdOrTitleFieldRef,
    TitleFieldRef,
)
import dl_core.exc


if TYPE_CHECKING:
    from dl_api_lib.query.formalization.raw_specs import FieldRef
    from dl_core.components.ids import FieldId
    from dl_core.fields import BIField
    from dl_core.us_dataset import Dataset


@attr.s
class FieldResolver:
    _dataset: Dataset = attr.ib(kw_only=True)

    def _find_field_by_id_or_title(self, id_or_title: str) -> BIField:
        try:
            field = self._dataset.result_schema.by_guid(id_or_title)
        except dl_core.exc.FieldNotFound:
            field = self._dataset.result_schema.by_title(id_or_title)
            # if it still raises `FieldNotFound`, well, then it raises
        return field

    def field_id_from_spec(self, field_ref: Optional[FieldRef]) -> FieldId:
        if isinstance(field_ref, IdFieldRef):
            # Run the search against `result_schema` to make sure the field exists
            # And that the correct exception is thrown if it doesn't
            field_id = self._dataset.result_schema.by_guid(field_ref.id).guid
        elif isinstance(field_ref, TitleFieldRef):
            field_id = self._dataset.result_schema.by_title(field_ref.title).guid
        elif isinstance(field_ref, IdOrTitleFieldRef):
            field_id = self._find_field_by_id_or_title(field_ref.id_or_title).guid
        else:
            raise TypeError(f"Invalid field spec: {field_ref}")
        return field_id

lib/dl_api_lib/dl_api_lib/query/formalization/field_resolver.py
datalens-tech/datalens-backend
from __future__ import annotations

import itertools
from typing import (
    Generator,
    Set,
)

import attr


@attr.s
class IdGenerator:
    _used_ids: Set[int] = attr.ib(factory=set)
    _item_counter: Generator[int, None, None] = attr.ib(init=False, factory=itertools.count)

    def add_id(self, item_id: int) -> None:
        self._used_ids.add(item_id)

    def generate_id(self) -> int:
        while True:
            item_id = next(self._item_counter)
            if item_id not in self._used_ids:
                self._used_ids.add(item_id)
                return item_id

lib/dl_api_lib/dl_api_lib/query/formalization/id_gen.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import json
from typing import (
    ClassVar,
    List,
    Optional,
    Set,
)

import attr

from dl_api_lib.query.formalization.field_resolver import FieldResolver
from dl_api_lib.query.formalization.id_gen import IdGenerator
from dl_api_lib.query.formalization.raw_specs import (
    MeasureNameRef,
    PlaceholderRef,
    RawFieldSpec,
    RawFilterFieldSpec,
    RawGroupByFieldSpec,
    RawOrderByFieldSpec,
    RawParameterValueSpec,
    RawQuerySpecUnion,
    RawRangeRoleSpec,
    RawRowRoleSpec,
    RawSelectFieldSpec,
    RawTemplateRoleSpec,
    RawTreeRoleSpec,
)
from dl_constants.enums import (
    CalcMode,
    FieldRole,
    FieldType,
    ManagedBy,
    UserDataType,
    WhereClauseOperation,
)
import dl_core.exc as core_exc
from dl_core.fields import BIField
from dl_core.us_dataset import Dataset
import dl_query_processing.exc
from dl_query_processing.legend.field_legend import (
    SELECTABLE_ROLES,
    DimensionNameObjSpec,
    FieldObjSpec,
    FilterRoleSpec,
    Legend,
    LegendItem,
    MeasureNameObjSpec,
    OrderByRoleSpec,
    ParameterRoleSpec,
    PlaceholderObjSpec,
    RangeRoleSpec,
    RoleSpec,
    RowRoleSpec,
    TemplateRoleSpec,
    TreeRoleSpec,
)
from dl_utils.utils import enum_not_none


DATA_TYPES_SUPPORTING_TREE = frozenset(
    {
        UserDataType.tree_str,
    }
)


@attr.s
class LegendFormalizer(abc.ABC):
    SUPPORTS_ROLES: ClassVar[Set[FieldRole]]
    SUPPORTS_MEASURE_NAME: ClassVar[bool] = False

    _dataset: Dataset = attr.ib(kw_only=True)
    _autofill_legend: bool = attr.ib(kw_only=True, default=False)
    _field_resolver: FieldResolver = attr.ib(init=False)

    def __attrs_post_init__(self) -> None:
        self._field_resolver = FieldResolver(dataset=self._dataset)

    def validate_legend_item(self, item: LegendItem) -> None:
        """Redefine this to impose restrictions on legend items"""
        if item.role_spec.role not in self.SUPPORTS_ROLES:
            raise dl_query_processing.exc.UnsopportedRoleInLegend(
                f"Legend role {item.role_spec.role.name} is not supported"
            )

        if not self.SUPPORTS_MEASURE_NAME and isinstance(item.obj, MeasureNameObjSpec):
            raise dl_query_processing.exc.MeasureNameUnsupported()

        if item.role_spec.role == FieldRole.tree:
            if item.data_type not in DATA_TYPES_SUPPORTING_TREE:
                raise dl_query_processing.exc.RoleDataTypeMismatch("Unsupported data type for tree role")

    def validate_legend(self, legend: Legend) -> None:
        for item in legend:
            self.validate_legend_item(item)

    def _get_field(self, field_id: str) -> BIField:
        return self._dataset.result_schema.by_guid(field_id)

    def _resolve_item_spec(
        self,
        legend_item_id: int,
        item_spec: RawFieldSpec,
        ignore_nonexistent_filters: bool,
    ) -> Optional[LegendItem]:
        # Make role spec
        role_spec: RoleSpec
        if isinstance(item_spec, RawSelectFieldSpec):
            raw_role_spec = item_spec.role_spec
            if raw_role_spec.role == FieldRole.template:
                assert isinstance(raw_role_spec, RawTemplateRoleSpec)
                role_spec = TemplateRoleSpec(
                    role=raw_role_spec.role,
                    template=raw_role_spec.template,
                    visibility=raw_role_spec.visibility,
                )
            elif raw_role_spec.role == FieldRole.tree:
                assert isinstance(raw_role_spec, RawTreeRoleSpec)
                try:
                    prefix = json.loads(raw_role_spec.prefix)
                    if not isinstance(prefix, list):
                        raise ValueError(prefix)
                except ValueError:
                    raise dl_query_processing.exc.GenericInvalidRequestError("Invalid value for tree prefix")
                role_spec = TreeRoleSpec(
                    role=raw_role_spec.role,
                    level=raw_role_spec.level,
                    prefix=prefix,
                    dimension_values=raw_role_spec.dimension_values,
                    visibility=raw_role_spec.visibility,
                )
            elif raw_role_spec.role == FieldRole.range:
                assert isinstance(raw_role_spec, RawRangeRoleSpec)
                role_spec = RangeRoleSpec(role=raw_role_spec.role, range_type=raw_role_spec.range_type)
            elif raw_role_spec.role == FieldRole.row:
                assert isinstance(raw_role_spec, RawRowRoleSpec)
                role_spec = RowRoleSpec(role=raw_role_spec.role, visibility=raw_role_spec.visibility)
            else:
                role_spec = RoleSpec(role=raw_role_spec.role)
        elif isinstance(item_spec, RawGroupByFieldSpec):
            role_spec = RowRoleSpec(role=FieldRole.row)
        elif isinstance(item_spec, RawOrderByFieldSpec):
            role_spec = OrderByRoleSpec(role=FieldRole.order_by, direction=item_spec.direction)
        elif isinstance(item_spec, RawFilterFieldSpec):
            role_spec = FilterRoleSpec(role=FieldRole.filter, operation=item_spec.operation, values=item_spec.values)
        elif isinstance(item_spec, RawParameterValueSpec):
            role_spec = ParameterRoleSpec(role=FieldRole.parameter, value=item_spec.value)
        else:
            raise TypeError(str(type(item_spec)))

        # Make the legend item
        legend_item: LegendItem
        if isinstance(item_spec.ref, MeasureNameRef):
            legend_item = LegendItem(
                legend_item_id=legend_item_id,
                obj=MeasureNameObjSpec(),
                block_id=item_spec.block_id,
                role_spec=role_spec,
                data_type=UserDataType.string,
                field_type=FieldType.DIMENSION,
            )
        elif isinstance(item_spec.ref, PlaceholderRef):
            legend_item = LegendItem(
                legend_item_id=legend_item_id,
                obj=PlaceholderObjSpec(),
                block_id=item_spec.block_id,
                role_spec=role_spec,
                data_type=UserDataType.string,
                field_type=FieldType.DIMENSION,
            )
        else:
            try:
                field_id = self._field_resolver.field_id_from_spec(item_spec.ref)
            except core_exc.FieldNotFound:
                if role_spec.role == FieldRole.filter and ignore_nonexistent_filters:
                    return None
                raise
            field = self._get_field(field_id)
            legend_item = LegendItem(
                legend_item_id=legend_item_id,
                block_id=item_spec.block_id,
                obj=FieldObjSpec(id=field_id, title=field.title),
                role_spec=role_spec,
                data_type=enum_not_none(field.data_type),
                field_type=field.type,
            )

        return legend_item

    def _validate_explicit_legend_item_ids(self, raw_query_spec_union: RawQuerySpecUnion) -> Set[int]:
        explicit_ids = [
            spec.legend_item_id for spec in raw_query_spec_union.iter_item_specs() if spec.legend_item_id is not None
        ]
        unique_explicit_ids = set(explicit_ids)
        if len(explicit_ids) != len(unique_explicit_ids):
            raise dl_query_processing.exc.NonUniqueLegendIdsError("Got non-unique legend item IDs")
        return unique_explicit_ids

    def _generate_info_items(self, legend: Legend, id_gen: IdGenerator) -> None:
        """
        Fill legend with all fields from result_schema that are not already there,
        use role info so that they are in no way used in the data request or pivot.
        """
        found_field_ids: Set[str] = {item.id for item in legend}
        for field in self._dataset.result_schema:
            if field.guid in found_field_ids:
                continue
            legend.add_item(
                LegendItem(
                    legend_item_id=id_gen.generate_id(),
                    obj=FieldObjSpec(id=field.guid, title=field.title),
                    data_type=field.data_type,
                    field_type=field.type,
                    role_spec=RoleSpec(role=FieldRole.info),
                )
            )

    def _generate_tree_filter_items(self, legend: Legend, id_gen: IdGenerator) -> None:
        """
        Generate additional filters for each item with role tree.
        Use the same block_id as in the original item.
        """

        for item in legend.list_for_role(role=FieldRole.tree):
            tree_spec = item.role_spec
            assert isinstance(tree_spec, TreeRoleSpec)

            # Generate filters for the tree field itself
            filter_role_specs: List[FilterRoleSpec] = []
            filter_role_specs.append(
                FilterRoleSpec(
                    role=FieldRole.filter,
                    operation=WhereClauseOperation.LENGTE,
                    values=[tree_spec.level],
                )
            )
            if len(tree_spec.prefix) > 0:
                filter_role_specs.append(
                    FilterRoleSpec(
                        role=FieldRole.filter,
                        operation=WhereClauseOperation.STARTSWITH,
                        values=[tree_spec.prefix],
                    )
                )
            for filter_role_spec in filter_role_specs:
                filter_item = LegendItem(
                    obj=item.obj,
                    data_type=item.data_type,
                    field_type=item.field_type,
                    role_spec=filter_role_spec,
                    legend_item_id=id_gen.generate_id(),
                    block_id=item.block_id,
                )
                legend.add_item(filter_item)

            # Generate filters for locked dimensions
            for dim_spec in tree_spec.dimension_values:
                locked_dim_item = legend.get_item(dim_spec.legend_item_id)
                if locked_dim_item.id == item.id:
                    # Ignore filter for the tree itself
                    continue
                filter_role_spec = FilterRoleSpec(
                    role=FieldRole.filter,
                    operation=WhereClauseOperation.EQ,
                    values=[dim_spec.value],
                )
                filter_item = LegendItem(
                    obj=locked_dim_item.obj,
                    data_type=locked_dim_item.data_type,
                    field_type=locked_dim_item.field_type,
                    role_spec=filter_role_spec,
                    legend_item_id=id_gen.generate_id(),
                    block_id=item.block_id,
                )
                legend.add_item(filter_item)

    def patch_legend(self, legend: Legend, id_gen: IdGenerator) -> None:
        """Patch the legend if something is missing"""

        self._generate_tree_filter_items(legend=legend, id_gen=id_gen)

        if self._autofill_legend:
            self._generate_info_items(legend=legend, id_gen=id_gen)

    def generate_legend_items(self, raw_query_spec_union: RawQuerySpecUnion, id_gen: IdGenerator) -> List[LegendItem]:
        items: List[LegendItem] = []
        already_used_field_ids: Set[str] = set()
        for item_spec in raw_query_spec_union.iter_item_specs():
            if isinstance(item_spec, RawGroupByFieldSpec):
                # Ignore group_by fields unless this field is not present in select
                if self._field_resolver.field_id_from_spec(item_spec.ref) in already_used_field_ids:
                    continue

            item = self._resolve_item_spec(
                legend_item_id=(
                    item_spec.legend_item_id if item_spec.legend_item_id is not None else id_gen.generate_id()
                ),
                item_spec=item_spec,
                ignore_nonexistent_filters=raw_query_spec_union.ignore_nonexistent_filters,
            )
            if item is None:
                # Skip it
                continue

            self.validate_legend_item(item)
            items.append(item)
            if item.role_spec.role in SELECTABLE_ROLES:
                already_used_field_ids.add(item.id)

        return items

    def make_legend(self, raw_query_spec_union: RawQuerySpecUnion) -> Legend:
        used_ids = self._validate_explicit_legend_item_ids(raw_query_spec_union)
        id_gen = IdGenerator(used_ids=used_ids)
        items = self.generate_legend_items(raw_query_spec_union=raw_query_spec_union, id_gen=id_gen)
        legend = Legend(items=items)
        self.patch_legend(legend=legend, id_gen=id_gen)
        self.validate_legend(legend)
        return legend


@attr.s
class ResultLegendFormalizer(LegendFormalizer):
    SUPPORTS_ROLES = {
        FieldRole.order_by,
        FieldRole.filter,
        FieldRole.parameter,
        FieldRole.info,
        FieldRole.row,
        FieldRole.total,
        FieldRole.template,
        FieldRole.tree,
        FieldRole.measure,
    }


@attr.s
class PreviewLegendFormalizer(LegendFormalizer):
    SUPPORTS_ROLES = {
        FieldRole.order_by,
        FieldRole.filter,
        FieldRole.parameter,
        FieldRole.info,
        FieldRole.row,
        FieldRole.measure,
    }

    def generate_legend_items(self, raw_query_spec_union: RawQuerySpecUnion, id_gen: IdGenerator) -> List[LegendItem]:
        assert not raw_query_spec_union.select_specs  # preview's select is always autogenerated # FIXME: DL exception
        items: List[LegendItem] = []

        for field in self._dataset.result_schema:
            if field.hidden:
                continue
            if field.managed_by != ManagedBy.user:
                continue
            if field.calc_mode == CalcMode.parameter:
                continue

            item = LegendItem(
                legend_item_id=id_gen.generate_id(),
                obj=FieldObjSpec(id=field.guid, title=field.title),
                data_type=field.data_type,
                field_type=field.type,
                role_spec=RowRoleSpec(role=FieldRole.row),
            )
            items.append(item)

        return items


@attr.s
class DistinctLegendFormalizer(LegendFormalizer):
    SUPPORTS_ROLES = {
        FieldRole.order_by,
        FieldRole.filter,
        FieldRole.parameter,
        FieldRole.info,
        FieldRole.distinct,
    }

    def validate_legend_item(self, item: LegendItem) -> None:
        super().validate_legend_item(item)
        if item.role_spec.role == FieldRole.filter and item.field_type == FieldType.MEASURE:
            raise dl_query_processing.exc.MeasureFilterUnsupportedError()

    def validate_legend(self, legend: Legend) -> None:
        super().validate_legend(legend=legend)
        if len(legend.list_for_role(FieldRole.distinct)) != 1:
            raise dl_query_processing.exc.LegendError("Legend must have one item")


@attr.s
class RangeLegendFormalizer(LegendFormalizer):
    SUPPORTS_ROLES = {
        FieldRole.filter,
        FieldRole.parameter,
        FieldRole.info,
        FieldRole.range,
    }

    def validate_legend_item(self, item: LegendItem) -> None:
        super().validate_legend_item(item)
        if item.role_spec.role == FieldRole.filter and item.field_type == FieldType.MEASURE:
            raise dl_query_processing.exc.MeasureFilterUnsupportedError()

    def validate_legend(self, legend: Legend) -> None:
        super().validate_legend(legend=legend)
        if len(legend.list_for_role(FieldRole.range)) > 2:
            raise dl_query_processing.exc.LegendError("Legend must have one or two range items")


@attr.s
class PivotLegendFormalizer(LegendFormalizer):
    SUPPORTS_ROLES = {
        FieldRole.order_by,
        FieldRole.filter,
        FieldRole.parameter,
        FieldRole.info,
        FieldRole.row,
        FieldRole.measure,
        FieldRole.template,
        FieldRole.total,
    }
    SUPPORTS_MEASURE_NAME = True

    def patch_legend(self, legend: Legend, id_gen: IdGenerator) -> None:
        super().patch_legend(legend=legend, id_gen=id_gen)

        # Always add `Measure Name` to the legend for pivot tables if it is not already there
        mname_legend_item_ids = legend.get_measure_name_legend_item_ids()
        if not mname_legend_item_ids:
            legend.add_item(
                LegendItem(
                    legend_item_id=id_gen.generate_id(),
                    obj=MeasureNameObjSpec(),
                    # using `row` here would corrupt the expected structure, so use `info`
                    role_spec=RoleSpec(role=FieldRole.info),
                    data_type=UserDataType.string,
                    field_type=FieldType.DIMENSION,
                )
            )

        # Same goes for `Dimension Name`
        dname_legend_item_ids = legend.get_dimension_name_legend_item_ids()
        if not dname_legend_item_ids:
            legend.add_item(
                LegendItem(
                    legend_item_id=id_gen.generate_id(),
                    obj=DimensionNameObjSpec(),
                    # using `row` here would corrupt the expected structure, so use `info`
                    role_spec=RoleSpec(role=FieldRole.info),
                    data_type=UserDataType.string,
                    field_type=FieldType.DIMENSION,
                )
            )


@attr.s
class ValidationLegendFormalizer(LegendFormalizer):
    SUPPORTS_ROLES = {
        FieldRole.row,
    }

    def generate_legend_items(self, raw_query_spec_union: RawQuerySpecUnion, id_gen: IdGenerator) -> List[LegendItem]:
        """Assume that all fields participate in the query"""

        items: List[LegendItem] = []

        for field in self._dataset.result_schema:
            if field.managed_by != ManagedBy.user:
                continue

            item = LegendItem(
                legend_item_id=id_gen.generate_id(),
                obj=FieldObjSpec(id=field.guid, title=field.title),
                data_type=field.data_type,
                field_type=field.type,
                role_spec=RowRoleSpec(role=FieldRole.row),
            )
            items.append(item)

        return items

lib/dl_api_lib/dl_api_lib/query/formalization/legend_formalizer.py
datalens-tech/datalens-backend
from typing import cast

import attr

from dl_api_lib.query.formalization.id_gen import IdGenerator
from dl_api_lib.query.formalization.pivot_legend import (
    PivotAnnotationRoleSpec,
    PivotDimensionRoleSpec,
    PivotLegend,
    PivotLegendItem,
    PivotMeasureRoleSpec,
    PivotRoleSpec,
)
from dl_api_lib.query.formalization.raw_pivot_specs import (
    RawAnnotationPivotRoleSpec,
    RawDimensionPivotRoleSpec,
    RawPivotLegendItem,
    RawPivotMeasureRoleSpec,
    RawPivotRoleSpec,
    RawPivotSpec,
)
from dl_constants.enums import (
    FieldRole,
    FieldType,
    PivotItemType,
    PivotRole,
)
from dl_constants.internal_constants import (
    DIMENSION_NAME_TITLE,
    MEASURE_NAME_TITLE,
)
from dl_core.us_dataset import Dataset
import dl_query_processing.exc
from dl_query_processing.legend.field_legend import (
    Legend,
    LegendItem,
    MeasureNameObjSpec,
)


DIMENSION_ROLES = {PivotRole.pivot_row, PivotRole.pivot_column}
MEASURE_ROLES = {PivotRole.pivot_measure}
ANNOTATION_ROLES = {PivotRole.pivot_annotation}


@attr.s
class PivotFormalizerBase:
    _legend: Legend = attr.ib(kw_only=True)

    def patch_pivot_legend(self, pivot_legend: PivotLegend, id_gen: IdGenerator) -> None:
        """
        Add pivot legend items if they are required
        to make the pivot table configuration valid.
        This mostly concerns `Measure Name` items.
        """

        measure_cnt = len(pivot_legend.list_for_role(PivotRole.pivot_measure))
        row_cnt = len(pivot_legend.list_for_role(PivotRole.pivot_row))
        column_cnt = len(pivot_legend.list_for_role(PivotRole.pivot_column))

        mname_legend_item_ids = self._legend.get_measure_name_legend_item_ids()
        assert mname_legend_item_ids
        mname_legend_item_id = next(iter(mname_legend_item_ids))

        mname_cnt = len([item for item in pivot_legend.items if item.item_type == PivotItemType.measure_name])

        if row_cnt == 0 or column_cnt == 0:
            # Patch one-sided cases
            if measure_cnt == 0:
                # Only dimensions on one side.
                pass  # Nothing to do here - a pseudo-dimension is added at transformer level
            elif measure_cnt == 1 or mname_cnt == 0:
                # One measure or mnames not yet used,
                # so it is safe to add mnames to the side that's empty
                if column_cnt == 0:  # Columns have priority if both are missing
                    pivot_legend.add_item(
                        PivotLegendItem(
                            pivot_item_id=id_gen.generate_id(),
                            legend_item_ids=[mname_legend_item_id],
                            role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                            item_type=PivotItemType.measure_name,
                            title=MEASURE_NAME_TITLE,
                        )
                    )
                elif row_cnt == 0:
                    pivot_legend.add_item(
                        PivotLegendItem(
                            pivot_item_id=id_gen.generate_id(),
                            legend_item_ids=[mname_legend_item_id],
                            role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                            item_type=PivotItemType.measure_name,
                            title=MEASURE_NAME_TITLE,
                        )
                    )
            elif measure_cnt > 1:  # and mname_cnt > 0
                # Can't add another mnames.
                # Must use a dummy dimension
                pass  # Nothing to do here - single-sided pivots are supported

        dname_legend_item_ids = self._legend.get_dimension_name_legend_item_ids()
        assert dname_legend_item_ids
        dname_legend_item_id = next(iter(dname_legend_item_ids))
        dname_cnt = len([item for item in pivot_legend.items if item.item_type == PivotItemType.dimension_name])
        if dname_cnt == 0:
            pivot_legend.add_item(
                PivotLegendItem(
                    pivot_item_id=id_gen.generate_id(),
                    legend_item_ids=[dname_legend_item_id],
                    role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_info),
                    item_type=PivotItemType.dimension_name,
                    title=DIMENSION_NAME_TITLE,
                )
            )

    def validate_pivot_legend(self, pivot_legend: PivotLegend) -> None:
        measures = pivot_legend.list_for_role(PivotRole.pivot_measure)
        measure_cnt = len(measures)
        mnames = [item for item in pivot_legend.items if item.item_type == PivotItemType.measure_name]
        mname_cnt = len(mnames)
        mname_role = mnames[0].role_spec.role if mnames else None

        measure_sorts_by_column, measure_sorts_by_row = 0, 0
        for item in measures:
            role_spec = cast(PivotMeasureRoleSpec, item.role_spec)
            if role_spec.sorting is None:
                continue
            measure_sorts_by_column += int(role_spec.sorting.column is not None)
            measure_sorts_by_row += int(role_spec.sorting.row is not None)

        if mname_cnt == 0 and measure_cnt > 1:
            # Measure names are not used and there are more than 1 dimensions
            raise dl_query_processing.exc.PivotMeasureNameRequired()
        if mname_cnt > 0 and measure_cnt == 0:
            # Measure names are used, but there are no dimensions
            raise dl_query_processing.exc.PivotMeasureNameForbidden()
        if mname_cnt > 1 and measure_cnt > 1:
            # Measure names are more than once, but there are multiple dimensions
            raise dl_query_processing.exc.PivotMeasureNameDuplicate()
        if measure_sorts_by_column > 1 or measure_sorts_by_row > 1:
            # Multiple sort along the same axis
            raise dl_query_processing.exc.PivotSortingMultipleColumnsOrRows()
        if measure_cnt > 1 and (
            (mname_role == PivotRole.pivot_row and measure_sorts_by_column > 0)
            or (mname_role == PivotRole.pivot_column and measure_sorts_by_row > 0)
        ):
            # Multiple measures, but some sorts are directed across them, which makes no sense
            raise dl_query_processing.exc.PivotSortingAgainstMultipleMeasures()


@attr.s
class PivotFormalizer(PivotFormalizerBase):
    """
    Generates PivotLegend form RawPivotSpec.
    """

    _dataset: Dataset = attr.ib(kw_only=True)
    _allow_dimension_annotations: bool = attr.ib(kw_only=True, default=True)

    def make_pivot_role_spec(self, raw_pivot_role_spec: RawPivotRoleSpec) -> PivotRoleSpec:
        role_spec: PivotRoleSpec
        if raw_pivot_role_spec.role in (PivotRole.pivot_row, PivotRole.pivot_column):
            assert isinstance(raw_pivot_role_spec, RawDimensionPivotRoleSpec)
            role_spec = PivotDimensionRoleSpec(
                role=raw_pivot_role_spec.role,
                direction=raw_pivot_role_spec.direction,
            )
        elif raw_pivot_role_spec.role == PivotRole.pivot_measure:
            assert isinstance(raw_pivot_role_spec, RawPivotMeasureRoleSpec)
            role_spec = PivotMeasureRoleSpec(
                role=raw_pivot_role_spec.role,
                sorting=raw_pivot_role_spec.sorting,
            )
        elif raw_pivot_role_spec.role == PivotRole.pivot_annotation:
            assert isinstance(raw_pivot_role_spec, RawAnnotationPivotRoleSpec)
            role_spec = PivotAnnotationRoleSpec(
                role=raw_pivot_role_spec.role,
                annotation_type=raw_pivot_role_spec.annotation_type,
                target_legend_item_ids=raw_pivot_role_spec.target_legend_item_ids,
            )
        else:
            raise ValueError(f"Unknown pivot role: {raw_pivot_role_spec.role}")

        return role_spec

    def make_item_type(self, legend_items: list[LegendItem], role: PivotRole) -> PivotItemType:
        is_mname_set = {isinstance(legend_item.obj, MeasureNameObjSpec) for legend_item in legend_items}
        assert len(is_mname_set) > 0
        if len(is_mname_set) > 1:
            raise dl_query_processing.exc.PivotItemsIncompatibleError(
                f"Pivot items have incompatible types in the same {role.name} role"
            )
        is_mname = next(iter(is_mname_set))
        if is_mname:
            return PivotItemType.measure_name
        return PivotItemType.stream_item

    def make_item_title(self, legend_items: list[LegendItem], role: PivotRole) -> str:
        title_set = {
            legend_item.title
            for legend_item in legend_items
            # ignore templates for totals
            if legend_item.role_spec.role != FieldRole.template
        }
        assert len(title_set) > 0
        if len(title_set) > 1:
            raise dl_query_processing.exc.PivotItemsIncompatibleError(
                f'Pivot items have incompatible titles: {",".join(sorted(title_set))} ' f"in the same {role.name} role"
            )

        return next(iter(title_set))

    def make_pivot_item(self, raw_pivot_item: RawPivotLegendItem, id_gen: IdGenerator) -> PivotLegendItem:
        role_spec = self.make_pivot_role_spec(raw_pivot_item.role_spec)
        target_legend_items = [
            self._legend.get_item(legend_item_id) for legend_item_id in raw_pivot_item.legend_item_ids
        ]

        has_measures = any(item.field_type == FieldType.MEASURE for item in target_legend_items)
        has_dimensions = any(item.field_type == FieldType.DIMENSION for item in target_legend_items)
        if has_measures and role_spec.role in DIMENSION_ROLES:
            raise dl_query_processing.exc.PivotInvalidRoleLegendError(
                "Measure field cannot be used as a pivot dimension"
            )
        if has_dimensions and role_spec.role in MEASURE_ROLES:
            raise dl_query_processing.exc.PivotInvalidRoleLegendError(
                "Dimension field cannot be used as a pivot measure"
            )
        if not self._allow_dimension_annotations and has_dimensions and role_spec.role in ANNOTATION_ROLES:
            raise dl_query_processing.exc.PivotInvalidRoleLegendError(
                "Dimension field cannot be used as a pivot annotation"
            )

        item_type = self.make_item_type(target_legend_items, role=raw_pivot_item.role_spec.role)

        title: str
        if raw_pivot_item.title is not None:
            title = raw_pivot_item.title
        else:
            title = self.make_item_title(target_legend_items, role=raw_pivot_item.role_spec.role)

        item = PivotLegendItem(
            pivot_item_id=id_gen.generate_id(),
            legend_item_ids=raw_pivot_item.legend_item_ids,
            role_spec=role_spec,
            item_type=item_type,
            title=title,
        )
        return item

    def make_pivot_legend(self, raw_pivot_spec: RawPivotSpec) -> PivotLegend:
        id_gen = IdGenerator()
        items = [self.make_pivot_item(raw_pivot_item, id_gen=id_gen) for raw_pivot_item in raw_pivot_spec.structure]
        pivot_legend = PivotLegend(items=items)

        self.patch_pivot_legend(pivot_legend, id_gen=id_gen)
        self.validate_pivot_legend(pivot_legend)

        return pivot_legend

lib/dl_api_lib/dl_api_lib/query/formalization/pivot_formalizer.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from typing import (
    Dict,
    List,
    Optional,
    Sequence,
)

import attr

from dl_api_lib.pivot.primitives import PivotMeasureSorting
from dl_constants.enums import (
    OrderDirection,
    PivotItemType,
    PivotRole,
)
import dl_query_processing.exc


@attr.s(frozen=True)
class PivotRoleSpec:
    role: PivotRole = attr.ib(kw_only=True)


@attr.s(frozen=True)
class PivotDimensionRoleSpec(PivotRoleSpec):
    direction: OrderDirection = attr.ib(kw_only=True, default=OrderDirection.asc)


@attr.s(frozen=True)
class PivotAnnotationRoleSpec(PivotRoleSpec):
    annotation_type: str = attr.ib(kw_only=True)
    target_legend_item_ids: Optional[list[int]] = attr.ib(
        kw_only=True, default=None
    )  # if None, then applies to all measures


@attr.s(frozen=True)
class PivotMeasureRoleSpec(PivotRoleSpec):
    sorting: Optional[PivotMeasureSorting] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class PivotLegendItem:
    pivot_item_id: int = attr.ib(kw_only=True)
    legend_item_ids: list[int] = attr.ib(kw_only=True)
    role_spec: PivotRoleSpec = attr.ib(kw_only=True)
    title: str = attr.ib(kw_only=True)
    item_type: PivotItemType = attr.ib(kw_only=True, default=PivotItemType.stream_item)


@attr.s
class PivotLegend:
    _items: List[PivotLegendItem] = attr.ib(kw_only=True)
    _piid_to_item: Dict[int, PivotLegendItem] = attr.ib(init=False)
    _liid_to_piid: Dict[int, list[int]] = attr.ib(init=False)

    def __attrs_post_init__(self) -> None:
        self._piid_to_item = {}
        self._liid_to_piid = defaultdict(list)
        self._repopulate_mappings()

    def _repopulate_mappings(self) -> None:
        self._piid_to_item.clear()
        self._liid_to_piid.clear()
        for item in self._items:
            self._piid_to_item[item.pivot_item_id] = item
            for legend_item_id in item.legend_item_ids:
                self._liid_to_piid[legend_item_id].append(item.pivot_item_id)

    @property
    def items(self) -> Sequence[PivotLegendItem]:
        return self._items

    def get_item(self, pivot_item_id: int) -> PivotLegendItem:
        try:
            return self._piid_to_item[pivot_item_id]
        except KeyError:
            raise dl_query_processing.exc.PivotLegendItemReferenceError(f"Unknown pivot legend item: {pivot_item_id}")

    def leg_item_id_to_pivot_item_id_list(self, legend_item_id: int) -> list[int]:
        return self._liid_to_piid[legend_item_id]

    def add_item(self, item: PivotLegendItem) -> None:
        self._items.append(item)
        self._repopulate_mappings()

    def list_for_role(self, role: PivotRole) -> Sequence[PivotLegendItem]:
        return tuple(item for item in self._items if item.role_spec.role == role)

    def get_measure_name_pivot_item_ids(self) -> List[int]:
        return [item.pivot_item_id for item in self.items if item.item_type == PivotItemType.measure_name]

    def get_dimension_name_pivot_item_ids(self) -> List[int]:
        return [item.pivot_item_id for item in self.items if item.item_type == PivotItemType.dimension_name]

    def get_unused_pivot_item_id(self) -> int:
        return max(item.pivot_item_id for item in self._items) + 1

lib/dl_api_lib/dl_api_lib/query/formalization/pivot_legend.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import os
from typing import (
    TYPE_CHECKING,
    AbstractSet,
    Collection,
    Dict,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
)

import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_api_lib.query.formalization.avatar_tools import normalize_explicit_avatar_ids
from dl_api_lib.query.formalization.field_resolver import FieldResolver
from dl_api_lib.query.formalization.query_formalizer_base import QuerySpecFormalizerBase
from dl_constants.enums import (
    CalcMode,
    DataSourceRole,
    FieldRole,
    FieldType,
    OrderDirection,
    RangeType,
    RLSSubjectType,
    UserDataType,
    WhereClauseOperation,
)
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.dependencies.factory_base import ComponentDependencyManagerFactoryBase
from dl_core.components.ids import (
    AvatarId,
    FieldId,
)
from dl_core.constants import DataAPILimits
from dl_core.data_source.base import DataSource
from dl_core.data_source.collection import DataSourceCollectionFactory
import dl_core.exc
from dl_core.fields import BIField
from dl_core.us_dataset import Dataset
from dl_core.us_manager.local_cache import USEntryBuffer
from dl_query_processing.compilation.specs import (
    ArrayPrefixSelectWrapperSpec,
    FilterFieldSpec,
    FilterSourceColumnSpec,
    GroupByFieldSpec,
    OrderByFieldSpec,
    ParameterValueSpec,
    RelationSpec,
    SelectFieldSpec,
    SelectWrapperSpec,
)
from dl_query_processing.enums import (
    GroupByPolicy,
    QueryType,
    SelectValueType,
)
import dl_query_processing.exc
from dl_query_processing.legend.field_legend import (
    FilterRoleSpec,
    OrderByRoleSpec,
    ParameterRoleSpec,
    RangeRoleSpec,
    RoleSpec,
    TreeRoleSpec,
)


if TYPE_CHECKING:
    from dl_core.services_registry.top_level import ServicesRegistry
    from dl_query_processing.compilation.query_meta import QueryMetaInfo
    from dl_query_processing.legend.block_legend import BlockSpec


LOGGER = logging.getLogger(__name__)


_RANGE_TYPE_TO_SELECT_TYPE = {
    RangeType.min: SelectWrapperSpec(type=SelectValueType.min),
    RangeType.max: SelectWrapperSpec(type=SelectValueType.max),
}


def need_permission_on_entry(us_entry, permission) -> None:  # type: ignore
    # TODO: DELETE ME after the check is moved up the stack
    assert us_entry.permissions is not None
    assert us_entry.uuid is not None
    if not us_entry.permissions[permission]:
        raise dl_core.exc.USPermissionRequired(us_entry.uuid, permission)


@attr.s
class SimpleQuerySpecFormalizer(QuerySpecFormalizerBase):  # noqa
    """
    Simple formalizer.
    Can be used for validation (no real data query is performed).
    """

    _dataset: Dataset = attr.ib(kw_only=True)
    _us_entry_buffer: USEntryBuffer = attr.ib(kw_only=True)

    _ds_accessor: DatasetComponentAccessor = attr.ib(init=False)
    _dsrc_coll_factory: DataSourceCollectionFactory = attr.ib(init=False)
    _field_resolver: FieldResolver = attr.ib(init=False)
    _avatar_resolution_cache: Dict[str, bool] = attr.ib(init=False, factory=dict)

    @_ds_accessor.default
    def _make_ds_accessor(self) -> DatasetComponentAccessor:
        return DatasetComponentAccessor(dataset=self._dataset)

    @_dsrc_coll_factory.default
    def _make_dsrc_coll_factory(self) -> DataSourceCollectionFactory:
        return DataSourceCollectionFactory(us_entry_buffer=self._us_entry_buffer)

    @_field_resolver.default
    def _make_field_resolver(self) -> FieldResolver:
        return FieldResolver(dataset=self._dataset)

    def _get_data_source_strict(self, source_id: str, role: DataSourceRole) -> DataSource:
        assert role is not None
        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
        dsrc_coll = self._dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
        dsrc = dsrc_coll.get_strict(role=role)
        return dsrc

    def _is_preview(self, block_spec: BlockSpec) -> bool:
        return block_spec.query_type == QueryType.preview

    def make_phantom_select_ids(
        self,
        block_spec: BlockSpec,
        order_by_specs: Sequence[OrderByFieldSpec],
    ) -> List[FieldId]:
        return []

    def _avatar_exists(self, avatar_id: str) -> bool:
        try:
            return self._avatar_resolution_cache[avatar_id]
        except KeyError:
            avatar_exists = self._ds_accessor.get_avatar_opt(avatar_id=avatar_id) is not None
            self._avatar_resolution_cache[avatar_id] = avatar_exists
            return avatar_exists

    def validate_select_field(self, block_spec: BlockSpec, field: BIField) -> None:
        pass

    def make_select_specs(
        self,
        block_spec: BlockSpec,
        phantom_select_ids: List[FieldId],
    ) -> List[SelectFieldSpec]:
        select_specs: List[SelectFieldSpec] = []
        select_spec_set: Set[SelectFieldSpec] = set()
        for legend_select_spec in block_spec.legend.list_selectable_items():
            select_role_spec = legend_select_spec.role_spec

            select_wrapper: SelectWrapperSpec
            if select_role_spec.role == FieldRole.range:
                assert isinstance(select_role_spec, RangeRoleSpec)
                select_wrapper = _RANGE_TYPE_TO_SELECT_TYPE[select_role_spec.range_type]
            elif select_role_spec.role == FieldRole.tree:
                assert isinstance(select_role_spec, TreeRoleSpec)
                select_wrapper = self._make_tree_wrapper(select_role_spec)
            else:
                select_wrapper = SelectWrapperSpec(type=SelectValueType.plain)

            field_id = legend_select_spec.id
            select_spec = SelectFieldSpec(field_id=field_id, wrapper=select_wrapper)
            if select_spec not in select_spec_set:
                select_specs.append(select_spec)
                select_spec_set.add(select_spec)

        for field_id in phantom_select_ids:
            select_spec = SelectFieldSpec(
                field_id=field_id,
                wrapper=SelectWrapperSpec(type=SelectValueType.plain),
            )
            if select_spec not in select_specs:
                select_specs.append(select_spec)
                select_spec_set.add(select_spec)

        for select_spec in select_specs:
            field = self._dataset.result_schema.by_guid(select_spec.field_id)
            self.validate_select_field(block_spec=block_spec, field=field)

        return select_specs

    def make_group_by_specs(
        self,
        block_spec: BlockSpec,
        select_specs: Sequence[SelectFieldSpec],
    ) -> List[GroupByFieldSpec]:
        group_by_specs: List[GroupByFieldSpec] = []
        group_by_spec_set: Set[GroupByFieldSpec] = set()
        for select_spec in select_specs:
            field = self._dataset.result_schema.by_guid(select_spec.field_id)
            if field.type == FieldType.DIMENSION:
                group_by_spec = GroupByFieldSpec(
                    field_id=select_spec.field_id,
                    wrapper=select_spec.wrapper,
                )
                if group_by_spec not in group_by_spec_set:
                    group_by_specs.append(group_by_spec)
                    group_by_spec_set.add(group_by_spec)

        return group_by_specs

    def make_filter_specs(
        self,
        block_spec: BlockSpec,
    ) -> List[FilterFieldSpec]:
        result: List[FilterFieldSpec] = []
        for legend_filter_spec in block_spec.legend.list_for_role(FieldRole.filter):
            filter_role_spec = legend_filter_spec.role_spec
            assert isinstance(filter_role_spec, FilterRoleSpec)
            try:
                field_id = legend_filter_spec.id
            except dl_core.exc.FieldNotFound:
                if block_spec.ignore_nonexistent_filters:
                    self._log_info("Skipping filter for unknown field %s", legend_filter_spec)
                    continue
                else:
                    raise

            filter_spec = FilterFieldSpec(
                field_id=field_id,
                operation=filter_role_spec.operation,
                values=filter_role_spec.values,
            )
            result.append(filter_spec)

        return result

    def _make_tree_wrapper(self, role_spec: RoleSpec) -> ArrayPrefixSelectWrapperSpec:
        assert isinstance(role_spec, TreeRoleSpec)
        return ArrayPrefixSelectWrapperSpec(length=role_spec.level)

    def make_order_by_specs(
        self,
        block_spec: BlockSpec,
    ) -> List[OrderByFieldSpec]:
        selectable_items = block_spec.legend.list_selectable_items()
        result: List[OrderByFieldSpec] = []
        for legend_order_by_spec in block_spec.legend.list_for_role(FieldRole.order_by):
            order_by_role_spec = legend_order_by_spec.role_spec
            assert isinstance(order_by_role_spec, OrderByRoleSpec)
            field_id = legend_order_by_spec.id
            selectable_matches = [item for item in selectable_items if item.id == field_id]
            select_wrapper: SelectWrapperSpec = SelectWrapperSpec(type=SelectValueType.plain)
            if selectable_matches:
                first_selectable_match = selectable_matches[0]
                if first_selectable_match.role_spec.role == FieldRole.tree:
                    # Special case: if we're selecting this field (and wrapping it) as a tree,
                    # then it must also be wrapped for ORDER BY
                    select_wrapper = self._make_tree_wrapper(first_selectable_match.role_spec)
            order_by_spec = OrderByFieldSpec(
                field_id=field_id,
                direction=order_by_role_spec.direction,
                wrapper=select_wrapper,
            )
            result.append(order_by_spec)

        return result

    def make_source_column_filter_specs(self) -> List[FilterSourceColumnSpec]:
        return []

    def make_relation_and_avatar_specs(
        self,
        used_field_ids: Collection[FieldId],
    ) -> Tuple[List[RelationSpec], AbstractSet[AvatarId], Optional[AvatarId]]:
        return [], set(), None

    def make_parameter_value_specs(
        self,
        block_spec: BlockSpec,
    ) -> List[ParameterValueSpec]:
        result: List[ParameterValueSpec] = []
        for legend_parameter_spec in block_spec.legend.list_for_role(FieldRole.parameter):
            parameter_role_spec = legend_parameter_spec.role_spec
            assert isinstance(parameter_role_spec, ParameterRoleSpec)
            try:
                field_id = legend_parameter_spec.id
            except dl_core.exc.FieldNotFound:
                if block_spec.ignore_nonexistent_filters:
                    self._log_info("Skipping parameter value for unknown field %s", legend_parameter_spec)
                    continue
                else:
                    raise

            parameter_value_spec = ParameterValueSpec(field_id=field_id, value=parameter_role_spec.value)
            result.append(parameter_value_spec)

        return result


@attr.s
class DataQuerySpecFormalizer(SimpleQuerySpecFormalizer):  # noqa
    """
    A more complex formalizer for real data queries:
    ``/result``, ``/preview``, etc.

    Since this is used for real data queries, more ad-hoc filters,
    additional select ID generators and validations kick in here.

    Ordering by fields omitted from SELECT.
        If some IDs are contained in ORDER BY, but not in SELECT,
        then add them as ``phantom_select_ids`` to enable this kind of ordering.

    RLS.
        Most data queries must be executed with RLS filters
        if such filters are defined. So additional filter specs are generated here.

    Source column filters.
        Some data sources may define filters that apply to their raw columns.
        If such filters exist, then they are enforced.

    Preview & hidden fields.
        Some of the dataset's fields may be marked as "hidden".
        Among other things this means that such fields should be ignored
        in preview queries. So exclude them from the SELECT list.

    Cache-based features & uncached fields.
        Selected fields (or, in general, fields used in the query somehow)
        may include function calls that require the usage of cache-based features
        such as Geocoding. Problems can arise if caches for these features
        have not been generated yet. The handling of fields
        that are affected by missing cache data is controlled
        by the query type - if it is preview,
        such fields are ignored, otherwise errors are raised if such fields are used.

    ``GROUP BY`` control
        In some cases we don't want to do ``GROUP BY`` at all - it saves a couple
        of glaciers from melting and preserves the data in its original form.
        The only problem is that it's is not possible to have aggregations in a
        query ``GROUP BY``-less query without collapsing all rows into a single
        one, so some validation is done against that scenario.

    """

    _rci: RequestContextInfo = attr.ib(kw_only=True)
    _role: DataSourceRole = attr.ib(kw_only=True)
    _dep_mgr_factory: ComponentDependencyManagerFactoryBase = attr.ib(kw_only=True)
    _service_registry: ServicesRegistry = attr.ib(kw_only=True)

    def _validate_phantom_select_ids(self, phantom_select_ids: Sequence[FieldId]) -> None:
        if any(
            self._dataset.result_schema.by_guid(phantom_select_id).type == FieldType.DIMENSION
            for phantom_select_id in phantom_select_ids
        ):
            # FIXME: Not really an InvalidFieldError... replace with some other exception
            raise dl_core.exc.InvalidFieldError("Order by dimension must be in select")

    def _ensure_not_measure(self, field: BIField) -> None:
        if field.type != FieldType.DIMENSION:
            raise dl_query_processing.exc.LogicError(
                f"Request does not support measure fields. Measure field: {field.guid}"
            )

    def _ensure_not_unsupported_type(self, field: BIField) -> None:
        if field.cast == UserDataType.unsupported:
            raise dl_query_processing.exc.LogicError(f"Cannot select fields of unsupported type: {field.title}")

    def validate_select_field(self, block_spec: BlockSpec, field: BIField) -> None:
        super().validate_select_field(block_spec=block_spec, field=field)

        field_id = field.guid

        if field.calc_mode == CalcMode.direct:
            avatar_id = field.avatar_id
            assert avatar_id is not None
            if not self._avatar_exists(avatar_id=avatar_id):
                raise dl_core.exc.UnknownReferencedAvatar(
                    f"Field {field.title!r} ({field_id}) references unknown source avatar " f"{field.avatar_id}."
                )

        if not field.valid:
            # FIXME: BI-2714 Investigate if/why this error is happening and return the raise
            # raise exc.InvalidFieldError(
            LOGGER.error(f"Field {field.title!r} ({field_id}) is invalid " f"and cannot be selected. Error ignored.")

        self._ensure_not_unsupported_type(field)
        if not block_spec.allow_measure_fields:
            self._ensure_not_measure(field)

    def _make_rls_filter_specs(self, subject_type: RLSSubjectType) -> List[FilterFieldSpec]:
        subject_id = self._rci.user_id
        if not subject_type or not subject_id:
            raise Exception("No subject to use in RLS")

        result: List[FilterFieldSpec] = []
        restrictions = self._dataset.rls.get_subject_restrictions(subject_type=subject_type, subject_id=subject_id)
        for field_guid, values in restrictions.items():
            result.append(
                FilterFieldSpec(
                    field_id=field_guid,
                    operation=WhereClauseOperation.IN,
                    values=values,  # type: ignore  # TODO: fix
                    anonymous=True,
                )
            )
        self._log_info("RLS filters for %s %s: %s", subject_type.name, subject_id, result)
        return result

    def make_phantom_select_ids(
        self,
        block_spec: BlockSpec,
        order_by_specs: Sequence[OrderByFieldSpec],
    ) -> List[FieldId]:
        order_by_id_set: Set[str] = {order_by_spec.field_id for order_by_spec in order_by_specs}
        select_id_set = {legend_select_spec.id for legend_select_spec in block_spec.legend.list_selectable_items()}
        phantom_select_ids = list(order_by_id_set - select_id_set)
        self._validate_phantom_select_ids(phantom_select_ids)
        return phantom_select_ids

    def make_filter_specs(self, block_spec: BlockSpec) -> List[FilterFieldSpec]:
        filter_specs = super().make_filter_specs(block_spec=block_spec)

        if block_spec.disable_rls:
            need_permission_on_entry(self._dataset, "admin")  # TODO: move it up the stack
            LOGGER.info(
                "Switching off RLS for request: " 'got "disable_rls" parameter and user has admin role for dataset'
            )

        if not block_spec.disable_rls:
            filter_specs += self._make_rls_filter_specs(subject_type=RLSSubjectType.user)

        return filter_specs

    def make_source_column_filter_specs(self) -> List[FilterSourceColumnSpec]:
        source_column_filter_specs = super().make_source_column_filter_specs()

        # This is not quite right, because we collect filters from all avatars,
        # even unused in current request. So we can break joins optimizations.
        # It's better to collect source filters after all other request parts have been processed -
        # only for really useful avatars.
        avatar_ids = [avatar.id for avatar in self._ds_accessor.get_avatar_list()]

        for avatar_id in sorted(avatar_ids):
            avatar = self._ds_accessor.get_avatar_strict(avatar_id=avatar_id)
            dsrc = self._get_data_source_strict(source_id=avatar.source_id, role=self._role)
            dsrc_filters = dsrc.get_filters(service_registry=self._service_registry)
            if dsrc_filters:
                self._log_info("Filters for source %s: %s", dsrc.id, dsrc_filters)

            for dsrc_filter_spec in dsrc_filters:
                source_column_filter_spec = FilterSourceColumnSpec(
                    avatar_id=avatar_id,
                    column_name=dsrc_filter_spec.name,
                    operation=dsrc_filter_spec.operation,
                    values=dsrc_filter_spec.values,
                )
                source_column_filter_specs.append(source_column_filter_spec)

        return source_column_filter_specs

    def make_group_by_specs(
        self,
        block_spec: BlockSpec,
        select_specs: Sequence[SelectFieldSpec],
    ) -> List[GroupByFieldSpec]:
        has_measures = any(
            self._dataset.result_schema.by_guid(spec.field_id).type == FieldType.MEASURE for spec in select_specs
        )

        if block_spec.group_by_policy == GroupByPolicy.if_measures:
            # If there are no measures, then don't GROUP BY
            if not has_measures:
                return []
            # Otherwise fall back to default grouping
        elif block_spec.group_by_policy == GroupByPolicy.disable:
            # GROUP BY is disabled. If there are measures in the query, it is not valid
            if has_measures:
                raise dl_query_processing.exc.InvalidGroupByConfiguration(
                    "Invalid parameter disable_group_by for dataset with measure fields"
                )
            return []

        # Default grouping
        return super().make_group_by_specs(block_spec=block_spec, select_specs=select_specs)

    def make_relation_and_avatar_specs(
        self,
        used_field_ids: Collection[FieldId],
    ) -> Tuple[List[RelationSpec], AbstractSet[AvatarId], Optional[AvatarId]]:
        # Resolve avatars explicitly specified in fields
        deep_dep_mgr = self._dep_mgr_factory.get_field_deep_inter_dependency_manager()
        field_ava_dep_mgr = self._dep_mgr_factory.get_field_avatar_dependency_manager()
        tree_resolver = self._dep_mgr_factory.get_avatar_tree_resolver()
        required_field_ids = set(used_field_ids) | {
            ref_field_id
            for field_id in used_field_ids
            for ref_field_id in deep_dep_mgr.get_field_deep_references(field_id)
        }
        explicitly_required_avatar_ids = {
            avatar_id
            for field_id in required_field_ids
            for avatar_id in field_ava_dep_mgr.get_field_avatar_references(field_id)
        }

        # Normalize avatars (fix them if there are no user-managed ones)
        explicitly_required_avatar_ids = normalize_explicit_avatar_ids(  # type: ignore  # TODO: fix
            dataset=self._dataset, required_avatar_ids=explicitly_required_avatar_ids
        )
        if not explicitly_required_avatar_ids:
            return [], set(), self._ds_accessor.get_root_avatar_strict().id

        # Resolve the "missing links" and get the required relation IDs
        root_avatar_id, required_avatar_ids, required_relation_ids = tree_resolver.expand_required_avatar_ids(
            required_avatar_ids=explicitly_required_avatar_ids
        )
        relation_specs = [RelationSpec(relation_id=relation_id) for relation_id in sorted(required_relation_ids)]
        return relation_specs, required_avatar_ids, root_avatar_id

    def make_limit_offset(self, block_spec: BlockSpec) -> Tuple[Optional[int], Optional[int]]:
        # Make defaults
        limit, offset = super().make_limit_offset(block_spec=block_spec)

        if self._is_preview(block_spec=block_spec) and self._role != DataSourceRole.sample:
            # Using direct/materialization mode, so we must to limit the number of entries
            limit = min(block_spec.limit or DataAPILimits.PREVIEW_ROW_LIMIT, DataAPILimits.PREVIEW_ROW_LIMIT)

        return limit, offset

    def make_query_meta(
        self,
        block_spec: BlockSpec,
        phantom_select_ids: List[FieldId],
        select_specs: List[SelectFieldSpec],
        root_avatar_id: Optional[AvatarId],
    ) -> QueryMetaInfo:
        query_meta = super().make_query_meta(
            block_spec=block_spec,
            phantom_select_ids=phantom_select_ids,
            select_specs=select_specs,
            root_avatar_id=root_avatar_id,
        )

        allow_subquery = os.environ.get("ALLOW_SUBQUERY_IN_PREVIEW", "0") != "0"  # FIXME: Do not get from env
        # No point in using subquery when fetching from sample - the table is already pretty small there,
        # so disable it by default
        if self._is_preview(block_spec=block_spec) and self._role != DataSourceRole.sample:
            # using direct/materialization mode, so we must to limit the number of entries
            from_subquery = bool(allow_subquery)

            if root_avatar_id is not None:
                avatar = self._ds_accessor.get_avatar_strict(avatar_id=root_avatar_id)
                dsrc = self._get_data_source_strict(source_id=avatar.source_id, role=self._role)
                from_subquery = from_subquery and dsrc.supports_preview_from_subquery

            if from_subquery != query_meta.from_subquery:
                query_meta = query_meta.clone(from_subquery=from_subquery)

        return query_meta


class NoGroupByQuerySpecFormalizerBase(DataQuerySpecFormalizer):
    def make_group_by_specs(
        self,
        block_spec: BlockSpec,
        select_specs: Sequence[SelectFieldSpec],
    ) -> List[GroupByFieldSpec]:
        return []


class ValueDistinctSpecFormalizer(NoGroupByQuerySpecFormalizerBase):
    def make_select_specs(
        self,
        block_spec: BlockSpec,
        phantom_select_ids: List[FieldId],
    ) -> List[SelectFieldSpec]:
        select_specs = super().make_select_specs(block_spec=block_spec, phantom_select_ids=phantom_select_ids)
        assert len(select_specs) == 1
        return select_specs

    def make_order_by_specs(
        self,
        block_spec: BlockSpec,
    ) -> List[OrderByFieldSpec]:
        legend_select_spec_list = block_spec.legend.list_selectable_items()
        assert len(legend_select_spec_list) == 1
        legend_select_spec = legend_select_spec_list[0]
        return [
            OrderByFieldSpec(
                field_id=legend_select_spec.id,
                direction=OrderDirection.asc,
                wrapper=SelectWrapperSpec(type=SelectValueType.plain),
            )
        ]


class SingleRowQuerySpecFormalizerBase(NoGroupByQuerySpecFormalizerBase):
    def make_order_by_specs(
        self,
        block_spec: BlockSpec,
    ) -> List[OrderByFieldSpec]:
        return []

    def make_limit_offset(self, block_spec: BlockSpec) -> Tuple[Optional[int], Optional[int]]:
        return (None, None)


class ValueRangeSpecFormalizer(SingleRowQuerySpecFormalizerBase):
    def validate_select_field(self, block_spec: BlockSpec, field: BIField) -> None:
        if field.type != FieldType.DIMENSION:
            raise dl_query_processing.exc.LogicError("Value range can only be fetched for dimensions")

    def make_select_specs(
        self,
        block_spec: BlockSpec,
        phantom_select_ids: List[FieldId],
    ) -> List[SelectFieldSpec]:
        select_specs = super().make_select_specs(block_spec=block_spec, phantom_select_ids=phantom_select_ids)
        assert len(select_specs) == 2
        assert all(
            select_spec.wrapper.type in (SelectValueType.min, SelectValueType.max) for select_spec in select_specs
        )
        return select_specs


class TotalsSpecFormalizer(DataQuerySpecFormalizer):
    pass

lib/dl_api_lib/dl_api_lib/query/formalization/query_formalizer.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    AbstractSet,
    Collection,
    List,
    Optional,
    Sequence,
    Tuple,
)

import attr

from dl_core.components.ids import (
    AvatarId,
    FieldId,
)
from dl_core.constants import DataAPILimits
from dl_query_processing.compilation.query_meta import QueryMetaInfo
from dl_query_processing.compilation.specs import (
    FilterFieldSpec,
    FilterSourceColumnSpec,
    GroupByFieldSpec,
    OrderByFieldSpec,
    ParameterValueSpec,
    QuerySpec,
    RelationSpec,
    SelectFieldSpec,
)
from dl_query_processing.legend.block_legend import BlockSpec


LOGGER = logging.getLogger(__name__)


@attr.s
class QuerySpecFormalizerBase(abc.ABC):
    """
    Normalizes lists of field IDs/titles or more complex specs (for filters and ORDER BY)
    for different parts of the future query and combines them into a single ``QuerySpec`` object.
    """

    _verbose_logging: bool = attr.ib(kw_only=True, default=False)

    def _log_info(self, *args, **kwargs) -> None:  # type: ignore  # TODO: fix
        if self._verbose_logging:
            LOGGER.info(*args, **kwargs)

    @abc.abstractmethod
    def make_phantom_select_ids(
        self,
        block_spec: BlockSpec,
        order_by_specs: Sequence[OrderByFieldSpec],
    ) -> List[FieldId]:
        raise NotImplementedError

    @abc.abstractmethod
    def make_select_specs(
        self,
        block_spec: BlockSpec,
        phantom_select_ids: List[FieldId],
    ) -> List[SelectFieldSpec]:
        raise NotImplementedError

    @abc.abstractmethod
    def make_group_by_specs(
        self,
        block_spec: BlockSpec,
        select_specs: Sequence[SelectFieldSpec],
    ) -> List[GroupByFieldSpec]:
        raise NotImplementedError

    @abc.abstractmethod
    def make_filter_specs(
        self,
        block_spec: BlockSpec,
    ) -> List[FilterFieldSpec]:
        raise NotImplementedError

    @abc.abstractmethod
    def make_order_by_specs(
        self,
        block_spec: BlockSpec,
    ) -> List[OrderByFieldSpec]:
        raise NotImplementedError

    @abc.abstractmethod
    def make_source_column_filter_specs(self) -> List[FilterSourceColumnSpec]:
        raise NotImplementedError

    @abc.abstractmethod
    def make_relation_and_avatar_specs(
        self,
        used_field_ids: Collection[FieldId],
    ) -> Tuple[List[RelationSpec], AbstractSet[AvatarId], Optional[AvatarId]]:
        raise NotImplementedError

    @abc.abstractmethod
    def make_parameter_value_specs(
        self,
        block_spec: BlockSpec,
    ) -> List[ParameterValueSpec]:
        raise NotImplementedError

    def make_limit_offset(self, block_spec: BlockSpec) -> Tuple[Optional[int], Optional[int]]:
        return block_spec.limit, block_spec.offset

    def make_query_meta(
        self,
        block_spec: BlockSpec,
        phantom_select_ids: List[FieldId],
        select_specs: List[SelectFieldSpec],
        root_avatar_id: Optional[AvatarId],
    ) -> QueryMetaInfo:
        row_count_hard_limit = block_spec.row_count_hard_limit
        if row_count_hard_limit is None:
            row_count_hard_limit = DataAPILimits.DEFAULT_SOURCE_DB_LIMIT
        assert row_count_hard_limit is not None
        row_count_hard_limit = min(row_count_hard_limit, DataAPILimits.DEFAULT_SOURCE_DB_LIMIT)

        query_meta = QueryMetaInfo(
            query_type=block_spec.query_type,
            phantom_select_ids=phantom_select_ids,
            field_order=[
                (idx, field_spec.field_id)
                for idx, field_spec in enumerate(select_specs)
                if field_spec.field_id not in phantom_select_ids
            ],
            row_count_hard_limit=row_count_hard_limit,
            empty_query_mode=block_spec.empty_query_mode,
        )
        return query_meta

    def make_query_spec(self, block_spec: BlockSpec) -> QuerySpec:
        order_by_specs = self.make_order_by_specs(block_spec=block_spec)
        phantom_select_ids = self.make_phantom_select_ids(block_spec=block_spec, order_by_specs=order_by_specs)
        select_specs = self.make_select_specs(block_spec=block_spec, phantom_select_ids=phantom_select_ids)
        group_by_specs = self.make_group_by_specs(block_spec=block_spec, select_specs=select_specs)
        filter_specs = self.make_filter_specs(block_spec=block_spec)
        used_field_ids = (
            {spec.field_id for spec in select_specs}
            | {fs.field_id for fs in filter_specs}
            | {obs.field_id for obs in order_by_specs}
        )
        relation_specs, required_avatar_ids, root_avatar_id = self.make_relation_and_avatar_specs(
            used_field_ids=used_field_ids
        )
        source_column_filter_specs = self.make_source_column_filter_specs()
        parameter_value_specs = self.make_parameter_value_specs(block_spec=block_spec)
        limit, offset = self.make_limit_offset(block_spec=block_spec)

        query_meta = self.make_query_meta(
            block_spec=block_spec,
            phantom_select_ids=phantom_select_ids,
            select_specs=select_specs,
            root_avatar_id=root_avatar_id,
        )

        return QuerySpec(
            select_specs=select_specs,
            group_by_specs=group_by_specs,
            filter_specs=filter_specs,
            source_column_filter_specs=source_column_filter_specs,
            order_by_specs=order_by_specs,
            relation_specs=relation_specs,
            parameter_value_specs=parameter_value_specs,
            limit=limit,
            offset=offset,
            root_avatar_id=root_avatar_id,
            required_avatar_ids=required_avatar_ids,
            meta=query_meta,
        )

lib/dl_api_lib/dl_api_lib/query/formalization/query_formalizer_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Optional,
    TypeVar,
)

import attr

from dl_api_lib.pivot.primitives import PivotMeasureSorting
from dl_constants.enums import (
    OrderDirection,
    PivotRole,
)


_PIVOT_ROLE_SPEC_TV = TypeVar("_PIVOT_ROLE_SPEC_TV", bound="RawPivotRoleSpec")


@attr.s(frozen=True)
class RawPivotRoleSpec:
    role: PivotRole = attr.ib(kw_only=True)

    def clone(self: _PIVOT_ROLE_SPEC_TV, **updates: Any) -> _PIVOT_ROLE_SPEC_TV:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class RawDimensionPivotRoleSpec(RawPivotRoleSpec):
    direction: OrderDirection = attr.ib(kw_only=True, default=OrderDirection.asc)


@attr.s(frozen=True)
class RawAnnotationPivotRoleSpec(RawPivotRoleSpec):
    target_legend_item_ids: Optional[list[int]] = attr.ib(kw_only=True, default=None)
    annotation_type: str = attr.ib(kw_only=True, default="")


@attr.s(frozen=True)
class RawPivotMeasureRoleSpec(RawPivotRoleSpec):
    sorting: Optional[PivotMeasureSorting] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class RawPivotLegendItem:
    pivot_item_id: Optional[int] = attr.ib(kw_only=True, default=None)
    legend_item_ids: list[int] = attr.ib(kw_only=True)
    role_spec: RawPivotRoleSpec = attr.ib(kw_only=True)
    title: Optional[str] = attr.ib(kw_only=True, default=None)

    def clone(self, **updates: Any) -> RawPivotLegendItem:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class RawPivotTotalsItemSpec:
    level: int = attr.ib(kw_only=True)


@attr.s(frozen=True)
class RawPivotTotalsSpec:
    rows: list[RawPivotTotalsItemSpec] = attr.ib(kw_only=True, factory=list)
    columns: list[RawPivotTotalsItemSpec] = attr.ib(kw_only=True, factory=list)

    def clone(self, **updates: Any) -> RawPivotTotalsSpec:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class PivotPaginationSpec:
    offset_rows: Optional[int] = attr.ib(kw_only=True, default=None)
    limit_rows: Optional[int] = attr.ib(kw_only=True, default=None)

    @property
    def non_default(self) -> bool:
        return self.offset_rows is not None or self.limit_rows is not None


@attr.s(frozen=True)
class RawPivotSpec:
    pagination: PivotPaginationSpec = attr.ib(kw_only=True, factory=PivotPaginationSpec)
    structure: list[RawPivotLegendItem] = attr.ib(kw_only=True, factory=list)
    with_totals: Optional[bool] = attr.ib(kw_only=True, default=False)
    totals: Optional[RawPivotTotalsSpec] = attr.ib(kw_only=True, factory=RawPivotTotalsSpec)

    def clone(self, **updates: Any) -> RawPivotSpec:
        return attr.evolve(self, **updates)

lib/dl_api_lib/dl_api_lib/query/formalization/raw_pivot_specs.py
datalens-tech/datalens-backend
from __future__ import annotations

from itertools import chain
from typing import (
    AbstractSet,
    Any,
    ClassVar,
    Iterable,
    List,
    Optional,
    TypeVar,
    Union,
)

import attr

from dl_constants.enums import (
    FieldRole,
    FieldVisibility,
    OrderDirection,
    QueryBlockPlacementType,
    RangeType,
    WhereClauseOperation,
)
from dl_constants.internal_constants import (
    DIMENSION_NAME_ID,
    MEASURE_NAME_ID,
    PLACEHOLDER_ID,
)
from dl_query_processing.base_specs.dimensions import DimensionValueSpec
from dl_query_processing.enums import (
    GroupByPolicy,
    QueryType,
)
from dl_query_processing.postprocessing.primitives import PostprocessedValue


FilterArgType = Union[str, int, float, None]


@attr.s(frozen=True)
class FieldRef:
    pass


@attr.s(frozen=True)
class IdFieldRef(FieldRef):
    id: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class TitleFieldRef(FieldRef):
    title: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class IdOrTitleFieldRef(FieldRef):
    id_or_title: str = attr.ib(kw_only=True)  # ID or title, depending on the context


@attr.s(frozen=True)
class MeasureNameRef(FieldRef):
    id: str = attr.ib(kw_only=True, init=False, default=MEASURE_NAME_ID)


@attr.s(frozen=True)
class DimensionNameRef(FieldRef):
    id: str = attr.ib(kw_only=True, init=False, default=DIMENSION_NAME_ID)


@attr.s(frozen=True)
class PlaceholderRef(FieldRef):
    id: str = attr.ib(kw_only=True, init=False, default=PLACEHOLDER_ID)


_FIELD_SPEC_TV = TypeVar("_FIELD_SPEC_TV", bound="RawFieldSpec")


@attr.s(frozen=True)
class RawFieldSpec:
    ref: FieldRef = attr.ib(kw_only=True, default=None)
    legend_item_id: Optional[int] = attr.ib(kw_only=True, default=None)
    block_id: Optional[int] = attr.ib(kw_only=True, default=None)

    def clone(self: _FIELD_SPEC_TV, **updates: Any) -> _FIELD_SPEC_TV:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class RawRoleSpec:
    role: FieldRole = attr.ib(kw_only=True)


@attr.s(frozen=True)
class RawOrderByRoleSpec(RawRoleSpec):
    direction: OrderDirection = attr.ib(kw_only=True)


@attr.s(frozen=True)
class RawRangeRoleSpec(RawRoleSpec):
    range_type: RangeType = attr.ib(kw_only=True)


@attr.s(frozen=True)
class RawDimensionRoleSpec(RawRoleSpec):
    visibility: FieldVisibility = attr.ib(kw_only=True, default=FieldVisibility.visible)


@attr.s(frozen=True)
class RawRowRoleSpec(RawDimensionRoleSpec):
    pass


@attr.s(frozen=True)
class RawTemplateRoleSpec(RawDimensionRoleSpec):
    template: str = attr.ib(kw_only=True, default="")


@attr.s(frozen=True)
class RawTreeRoleSpec(RawDimensionRoleSpec):
    level: int = attr.ib(kw_only=True)
    prefix: str = attr.ib(kw_only=True)
    dimension_values: List[DimensionValueSpec] = attr.ib(kw_only=True, factory=list)


@attr.s(frozen=True)
class RawSelectFieldSpec(RawFieldSpec):
    role_spec: RawRoleSpec = attr.ib(kw_only=True, default=RawRowRoleSpec(role=FieldRole.row))
    label: Optional[str] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class RawGroupByFieldSpec(RawFieldSpec):
    pass


@attr.s(frozen=True)
class RawOrderByFieldSpec(RawFieldSpec):  # noqa
    """
    Ambiguous ORDER BY field clause specifier that serves as input from the API,
    where either field ID or title can be specified.
    """

    direction: OrderDirection = attr.ib(kw_only=True)


@attr.s(frozen=True)
class RawFilterFieldSpec(RawFieldSpec):  # noqa
    """
    Ambiguous filter specifier
    """

    operation: WhereClauseOperation = attr.ib(kw_only=True)
    values: List[FilterArgType] = attr.ib(kw_only=True)
    block_id: Optional[int] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class RawParameterValueSpec(RawFieldSpec):
    value: Any = attr.ib(kw_only=True, default=None)


@attr.s
class RawQueryMetaInfo:
    query_type: QueryType = attr.ib(kw_only=True, default=QueryType.result)
    row_count_hard_limit: Optional[int] = attr.ib(kw_only=True, default=None)

    def clone(self, **updates: Any) -> RawQueryMetaInfo:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class RawQuerySpecUnion:
    # TODO: Combine with RawQuerySpecUnion
    limit: Optional[int] = attr.ib(kw_only=True, default=None)
    offset: Optional[int] = attr.ib(kw_only=True, default=None)
    group_by_policy: GroupByPolicy = attr.ib(kw_only=True, default=GroupByPolicy.force)
    disable_rls: bool = attr.ib(kw_only=True, default=False)
    ignore_nonexistent_filters: bool = attr.ib(kw_only=True, default=True)
    allow_measure_fields: bool = attr.ib(kw_only=True, default=True)
    meta: RawQueryMetaInfo = attr.ib(kw_only=True, factory=RawQueryMetaInfo)
    # item lists
    select_specs: List[RawSelectFieldSpec] = attr.ib(kw_only=True, factory=list)
    group_by_specs: List[RawGroupByFieldSpec] = attr.ib(kw_only=True, factory=list)
    order_by_specs: List[RawOrderByFieldSpec] = attr.ib(kw_only=True, factory=list)
    filter_specs: List[RawFilterFieldSpec] = attr.ib(kw_only=True, factory=list)
    parameter_value_specs: List[RawParameterValueSpec] = attr.ib(kw_only=True, factory=list)
    block_specs: List[RawBlockSpec] = attr.ib(kw_only=True, factory=list)

    def clone(self, **updates: Any) -> RawQuerySpecUnion:
        return attr.evolve(self, **updates)

    def iter_item_specs(self) -> Iterable[RawFieldSpec]:
        return chain(
            self.select_specs,
            self.group_by_specs,
            self.order_by_specs,
            self.filter_specs,
            self.parameter_value_specs,
        )

    def get_unique_block_ids(self) -> AbstractSet[int]:
        return {spec.block_id for spec in self.iter_item_specs() if spec.block_id is not None}


@attr.s(frozen=True)
class RawDimensionValueSpec:
    legend_item_id: int = attr.ib(kw_only=True)
    value: PostprocessedValue = attr.ib(kw_only=True)


@attr.s(frozen=True)
class RawBlockPlacement:
    type: ClassVar[QueryBlockPlacementType]


@attr.s(frozen=True)
class RawRootBlockPlacement(RawBlockPlacement):
    type = QueryBlockPlacementType.root


@attr.s(frozen=True)
class RawAfterBlockPlacement(RawBlockPlacement):
    type = QueryBlockPlacementType.after

    dimension_values: Optional[List[RawDimensionValueSpec]] = attr.ib(default=None)


@attr.s
class RawBlockSpec:
    block_id: int = attr.ib(kw_only=True)
    parent_block_id: Optional[int] = attr.ib(kw_only=True)
    placement: RawBlockPlacement = attr.ib(kw_only=True, factory=RawRootBlockPlacement)
    limit: Optional[int] = attr.ib(kw_only=True, default=None)
    offset: Optional[int] = attr.ib(kw_only=True, default=None)
    row_count_hard_limit: Optional[int] = attr.ib(kw_only=True, default=None)


def spec_is_field_name_pseudo_dimension(spec: RawFieldSpec) -> bool:
    return isinstance(spec.ref, (MeasureNameRef, DimensionNameRef))


@attr.s(frozen=True)
class RawResultSpec:
    with_totals: bool = attr.ib(kw_only=True, default=False)

lib/dl_api_lib/dl_api_lib/query/formalization/raw_specs.py
datalens-tech/datalens-backend
from typing import (
    Collection,
    Optional,
    Type,
)

import attr

from dl_constants.enums import (
    QueryProcessingMode,
    SourceBackendType,
)
from dl_core.fields import ResultSchema
from dl_formula.core.dialect import DialectCombo
from dl_query_processing.compilation.filter_compiler import (
    FilterFormulaCompiler,
    MainFilterFormulaCompiler,
)
from dl_query_processing.multi_query.factory import (
    DefaultMultiQueryMutatorFactory,
    MultiQueryMutatorFactoryBase,
)


_FILTER_FORMULA_COMPILER_BY_BACKEND: dict[SourceBackendType, Type[FilterFormulaCompiler]] = {}
_DEFAULT_FILTER_FORMULA_COMPILER = MainFilterFormulaCompiler


def get_filter_formula_compiler_cls(backend_type: SourceBackendType) -> Type[FilterFormulaCompiler]:
    return _FILTER_FORMULA_COMPILER_BY_BACKEND.get(backend_type, _DEFAULT_FILTER_FORMULA_COMPILER)


def register_filter_formula_compiler_cls(
    backend_type: SourceBackendType,
    filter_compiler_cls: Type[FilterFormulaCompiler],
) -> None:
    try:
        assert _FILTER_FORMULA_COMPILER_BY_BACKEND[backend_type] is filter_compiler_cls
    except KeyError:
        _FILTER_FORMULA_COMPILER_BY_BACKEND[backend_type] = filter_compiler_cls


_IS_FORKABLE_BACKEND_TYPE: dict[SourceBackendType, bool] = {}
_IS_FORKABLE_DEFAULT = True


def is_forkable_source(backend_type: SourceBackendType) -> bool:
    return _IS_FORKABLE_BACKEND_TYPE.get(backend_type, _IS_FORKABLE_DEFAULT)


def register_is_forkable_source(backend_type: SourceBackendType, is_forkable: bool) -> None:
    try:
        assert _IS_FORKABLE_BACKEND_TYPE[backend_type] is is_forkable
    except KeyError:
        _IS_FORKABLE_BACKEND_TYPE[backend_type] = is_forkable


_IS_COMPENG_EXECUTABLE_BACKEND_TYPE: dict[SourceBackendType, bool] = {}
_IS_COMPENG_EXECUTABLE_DEFAULT = False


def is_compeng_executable(backend_type: SourceBackendType) -> bool:
    return _IS_COMPENG_EXECUTABLE_BACKEND_TYPE.get(backend_type, _IS_COMPENG_EXECUTABLE_DEFAULT)


def register_is_compeng_executable(backend_type: SourceBackendType, is_compeng_executable: bool) -> None:
    try:
        assert _IS_COMPENG_EXECUTABLE_BACKEND_TYPE[backend_type] is is_compeng_executable
    except KeyError:
        _IS_COMPENG_EXECUTABLE_BACKEND_TYPE[backend_type] = is_compeng_executable


@attr.s(frozen=True, auto_attribs=True, kw_only=True)
class MQMFactoryKey:
    query_proc_mode: QueryProcessingMode
    backend_type: SourceBackendType
    dialect: Optional[DialectCombo]


@attr.s(frozen=True, auto_attribs=True, kw_only=True)
class MQMFactorySettingItem:
    query_proc_mode: QueryProcessingMode
    factory_cls: Type[MultiQueryMutatorFactoryBase]
    dialects: Collection[Optional[DialectCombo]] = attr.ib(default=(None,))


_MQM_FACTORY_REGISTRY: dict[MQMFactoryKey, Type[MultiQueryMutatorFactoryBase]] = {}


def get_multi_query_mutator_factory(
    query_proc_mode: QueryProcessingMode,
    backend_type: SourceBackendType,
    dialect: DialectCombo,
    result_schema: ResultSchema,
) -> Optional[MultiQueryMutatorFactoryBase]:
    factory_cls = _MQM_FACTORY_REGISTRY.get(
        # First try with exact dialect
        MQMFactoryKey(query_proc_mode=query_proc_mode, backend_type=backend_type, dialect=dialect),
        _MQM_FACTORY_REGISTRY.get(
            # Then try without the dialect, just the backend
            MQMFactoryKey(query_proc_mode=query_proc_mode, backend_type=backend_type, dialect=None),
            DefaultMultiQueryMutatorFactory,  # If still nothing, then use the default
        ),
    )

    if factory_cls is None:
        return None

    return factory_cls(result_schema=result_schema)


def register_multi_query_mutator_factory_cls(
    query_proc_mode: QueryProcessingMode,
    backend_type: SourceBackendType,
    dialects: Collection[Optional[DialectCombo]],
    factory_cls: Type[MultiQueryMutatorFactoryBase],
) -> None:
    for dialect in dialects:
        key = MQMFactoryKey(query_proc_mode=query_proc_mode, backend_type=backend_type, dialect=dialect)
        if key in _MQM_FACTORY_REGISTRY:
            assert _MQM_FACTORY_REGISTRY[key] is factory_cls
        else:
            _MQM_FACTORY_REGISTRY[key] = factory_cls


_COMPENG_DIALECT: set[DialectCombo] = set()


def get_compeng_dialect() -> DialectCombo:
    assert len(_COMPENG_DIALECT) == 1, "No compeng dialect has been found"
    return next(iter(_COMPENG_DIALECT))


def is_compeng_enabled() -> bool:
    return bool(_COMPENG_DIALECT)


def register_compeng_dialect(dialect: DialectCombo) -> None:
    if not _COMPENG_DIALECT:
        _COMPENG_DIALECT.add(dialect)
    else:
        assert next(iter(_COMPENG_DIALECT)) == dialect

lib/dl_api_lib/dl_api_lib/query/registry.py
datalens-tech/datalens-backend
from enum import Enum
from typing import (
    Any,
    Dict,
    List,
    Optional,
    TypeVar,
)

import attr

from dl_api_lib.enums import DatasetAction
from dl_api_lib.query.formalization.raw_pivot_specs import RawPivotSpec
from dl_api_lib.query.formalization.raw_specs import (
    RawFilterFieldSpec,
    RawQuerySpecUnion,
    RawResultSpec,
)
from dl_constants.enums import (
    AggregationFunction,
    CalcMode,
    FieldType,
    ManagedBy,
    UserDataType,
)
from dl_core.fields import CalculationSpec
from dl_core.values import BIValue


@attr.s(frozen=True, kw_only=True)
class Action:
    action: DatasetAction = attr.ib()
    order_index: int = attr.ib()
    managed_by: Optional[ManagedBy] = attr.ib(default=None)


@attr.s(frozen=True, kw_only=True)
class FieldBase:
    guid: Optional[str] = attr.ib(default=None)
    strict: Optional[bool] = attr.ib(default=None)

    def as_dict(self) -> Dict:
        return attr.asdict(self, recurse=False)


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class FieldAction(Action):
    field: FieldBase

    @property
    def serialized(self) -> Dict[str, Any]:
        serializer = lambda t, at, val: val.name if isinstance(val, Enum) else val  # noqa
        return attr.asdict(self, recurse=True, value_serializer=serializer)  # type: ignore


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class SourceActionBase(Action):
    source: dict


@attr.s(frozen=True, kw_only=True)
class AddUpdateSourceAction(SourceActionBase):
    pass


@attr.s(frozen=True, kw_only=True)
class DeleteField(FieldBase):
    pass


@attr.s(frozen=True, kw_only=True)
class UpdateField(FieldBase):
    title: Optional[str] = attr.ib(default=None)
    source: Optional[str] = attr.ib(default=None)
    calc_mode: Optional[CalcMode] = attr.ib(default=None)
    calc_spec: Optional[CalculationSpec] = attr.ib(default=None)
    hidden: Optional[bool] = attr.ib(default=None)
    description: Optional[str] = attr.ib(default=None)
    aggregation: Optional[AggregationFunction] = attr.ib(default=None)
    has_auto_aggregation: Optional[bool] = attr.ib(default=None)
    lock_aggregation: Optional[bool] = attr.ib(default=None)
    formula: Optional[str] = attr.ib(default=None)
    guid_formula: Optional[str] = attr.ib(default=None)
    cast: Optional[UserDataType] = attr.ib(default=None)
    avatar_id: Optional[str] = attr.ib(default=None)
    new_id: Optional[str] = attr.ib(default=None)
    default_value: Optional[BIValue] = attr.ib(default=None)
    value_constraint: dict = attr.ib(default=None)


@attr.s(frozen=True, kw_only=True)
class AddField(UpdateField):
    title: str = attr.ib()
    type: Optional[FieldType] = attr.ib(default=None)


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class CloneField(UpdateField):
    title: str
    from_guid: str


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class SourceBase:
    id: str


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class ObligatoryFilterBase:
    id: str


@attr.s(frozen=True, kw_only=True)
class AddUpdateObligatoryFilter(ObligatoryFilterBase):
    field_guid: Optional[str] = attr.ib(default=None)
    default_filters: List[RawFilterFieldSpec] = attr.ib(default=[])
    managed_by: Optional[ManagedBy] = attr.ib(default=None)
    valid: Optional[bool] = attr.ib(default=None)


@attr.s(frozen=True, kw_only=True)
class DeleteObligatoryFilter(ObligatoryFilterBase):
    pass


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class ObligatoryFilterActionBase(Action):
    obligatory_filter: ObligatoryFilterBase


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class AddUpdateObligatoryFilterAction(ObligatoryFilterActionBase):
    obligatory_filter: AddUpdateObligatoryFilter


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class DeleteObligatoryFilterAction(ObligatoryFilterActionBase):
    obligatory_filter: DeleteObligatoryFilter


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class ReplaceConnection:
    id: str
    new_id: str


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class ReplaceConnectionAction(Action):
    connection: ReplaceConnection


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class RelationActionBase(Action):
    avatar_relation: dict


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class SourceAvatarBase(Action):
    id: str


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class SourceAvatar(SourceAvatarBase):
    source_id: Optional[str] = attr.ib(default=None)
    title: Optional[str] = attr.ib(default=None)
    is_root: Optional[bool] = attr.ib(default=None)
    managed_by: Optional[ManagedBy] = attr.ib(default=None)


@attr.s(frozen=True, kw_only=True, auto_attribs=True)
class AvatarActionBase(Action):
    source_avatar: SourceAvatar
    disable_fields_update: bool


_DRM_TV = TypeVar("_DRM_TV", bound="DataRequestModel")


@attr.s()
class DataRequestModel:
    # Dataset state
    dataset: Optional[Dict[str, Any]] = attr.ib(kw_only=True, default=None)  # TODO: schematize
    # Updates to apply to the given dataset state
    updates: List[Action] = attr.ib(kw_only=True, factory=list)

    # Specifies what the data query will consist of
    raw_query_spec_union: RawQuerySpecUnion = attr.ib(kw_only=True, factory=RawQuerySpecUnion)

    # Other parameters controlling the request
    add_fields_data: bool = attr.ib(kw_only=True, default=False)
    with_totals: bool = attr.ib(kw_only=True, default=False)
    autofill_legend: bool = attr.ib(kw_only=True, default=False)
    dataset_revision_id: Optional[str] = attr.ib(kw_only=True, default=None)

    def clone(self: _DRM_TV, **updates: Any) -> _DRM_TV:
        return attr.evolve(self, **updates)


@attr.s()
class ResultDataRequestModel(DataRequestModel):
    result: Optional[RawResultSpec] = attr.ib(kw_only=True, default=None)


@attr.s()
class PivotDataRequestModel(DataRequestModel):
    pivot: RawPivotSpec = attr.ib(kw_only=True, factory=RawPivotSpec)

lib/dl_api_lib/dl_api_lib/request_model/data.py
datalens-tech/datalens-backend
import abc
from typing import (
    Generic,
    TypeVar,
)

from dl_api_lib.query.formalization.id_gen import IdGenerator
from dl_api_lib.query.formalization.raw_specs import RawQuerySpecUnion
from dl_api_lib.request_model.data import DataRequestModel


_SPEC_TV = TypeVar("_SPEC_TV")


class RequestPartSpecNormalizerBase(abc.ABC, Generic[_SPEC_TV]):
    def get_liid_generator(self, raw_query_spec_union: RawQuerySpecUnion) -> IdGenerator:
        used_ids = {
            spec.legend_item_id for spec in raw_query_spec_union.iter_item_specs() if spec.legend_item_id is not None
        }
        id_gen = IdGenerator(used_ids=used_ids)
        return id_gen

    def normalize_spec(
        self,
        spec: _SPEC_TV,
        raw_query_spec_union: RawQuerySpecUnion,
    ) -> tuple[_SPEC_TV, RawQuerySpecUnion]:
        return spec, raw_query_spec_union


_DRM_TV = TypeVar("_DRM_TV", bound=DataRequestModel)


class RequestModelNormalizerBase(abc.ABC, Generic[_DRM_TV]):
    """
    The idea behind this class is that it expands various API shortcuts
    into a full shortcut-less version of the request
    (e.g. converts the short pivot totals spec into additional
    legend items and pivot structure items).
    """

    def normalize_drm(self, drm: _DRM_TV) -> _DRM_TV:
        return drm

lib/dl_api_lib/dl_api_lib/request_model/normalization/drm_normalizer_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import itertools
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Optional,
)

import attr

from dl_api_lib.query.formalization.id_gen import IdGenerator
from dl_api_lib.query.formalization.raw_pivot_specs import (
    RawAnnotationPivotRoleSpec,
    RawPivotSpec,
    RawPivotTotalsItemSpec,
    RawPivotTotalsSpec,
)
from dl_api_lib.query.formalization.raw_specs import (
    PlaceholderRef,
    RawQuerySpecUnion,
    RawRoleSpec,
    RawRowRoleSpec,
    RawTemplateRoleSpec,
    spec_is_field_name_pseudo_dimension,
)
from dl_api_lib.request_model.data import PivotDataRequestModel
from dl_api_lib.request_model.normalization.drm_normalizer_base import (
    RequestModelNormalizerBase,
    RequestPartSpecNormalizerBase,
)
from dl_constants.enums import (
    FieldRole,
    PivotRole,
)
import dl_query_processing.exc


if TYPE_CHECKING:
    from dl_api_lib.query.formalization.raw_pivot_specs import RawPivotLegendItem
    from dl_api_lib.query.formalization.raw_specs import RawSelectFieldSpec


@attr.s
class PivotTotalsNormalizerHelper:
    _original_select_specs: list[RawSelectFieldSpec] = attr.ib(kw_only=True)
    _legend_item_id_gen: IdGenerator = attr.ib(kw_only=True)
    _block_id_gen: IdGenerator = attr.ib(kw_only=True)
    _base_block_id: int = attr.ib(kw_only=True)

    # Item types:
    T_DIMENSION: ClassVar[int] = 1
    T_TEMPLATE: ClassVar[int] = 2
    T_MEASURE: ClassVar[int] = 3

    def _get_legend_item_by_id(self, legend_item_id: int) -> RawSelectFieldSpec:
        for item in self._original_select_specs:
            if item.legend_item_id == legend_item_id:
                return item
        raise dl_query_processing.exc.LegendItemReferenceError(f"Unknown legend item: {legend_item_id}")

    def make_role_spec(self, item_type: int) -> RawRoleSpec:
        if item_type is self.T_TEMPLATE:
            return RawTemplateRoleSpec(role=FieldRole.template, template="")
        if item_type is self.T_MEASURE:
            return RawRoleSpec(role=FieldRole.total)
        return RawRowRoleSpec(role=FieldRole.row)

    def resolve_main_legend_item(self, pivot_item: RawPivotLegendItem) -> Optional[RawSelectFieldSpec]:
        # TODO: Some validation would be nice here
        for legend_item_id in pivot_item.legend_item_ids:
            item = self._get_legend_item_by_id(legend_item_id=legend_item_id)
            # We are looking for the item from the original base block
            if item.block_id == self._base_block_id:
                return item
        return None

    def make_legend_item(
        self,
        item_type: int,
        orig_pivot_item: RawPivotLegendItem,
        block_id: int,
    ) -> tuple[Optional[RawSelectFieldSpec], RawPivotLegendItem]:
        """
        Generate raw legend item and an updated version of the given pivot item.
        """
        original_item = self.resolve_main_legend_item(orig_pivot_item)
        if original_item is None:
            return None, orig_pivot_item
        new_role_spec = self.make_role_spec(item_type=item_type)
        legend_item_id = self._legend_item_id_gen.generate_id()
        ref = original_item.ref
        if item_type is self.T_TEMPLATE:
            ref = PlaceholderRef()
        new_item = original_item.clone(
            ref=ref,
            role_spec=new_role_spec,
            legend_item_id=legend_item_id,
            block_id=block_id,
        )
        updated_pivot_item = orig_pivot_item.clone(legend_item_ids=orig_pivot_item.legend_item_ids + [legend_item_id])
        return new_item, updated_pivot_item

    def gen_single_dir_totals(
        self,
        this_level: Optional[int],
        other_level: Optional[int],
        this_dir_dimensions: list[RawPivotLegendItem],
        other_dir_dimensions: list[RawPivotLegendItem],
        measures: list[RawPivotLegendItem],
    ) -> tuple[list[RawSelectFieldSpec], list[RawPivotLegendItem], list[RawPivotLegendItem], list[RawPivotLegendItem],]:
        """
        Generate new items for the legend (raw_query_spec_union)
        and pivot spec.
        """

        block_id = self._block_id_gen.generate_id()

        this_dir_dimensions = this_dir_dimensions[:]
        other_dir_dimensions = other_dir_dimensions[:]
        measures = measures[:]

        # Generate items
        new_block_select_items: list[RawSelectFieldSpec] = []
        for item_idx, orig_pivot_item in enumerate(this_dir_dimensions):
            item_type = self.T_DIMENSION if (this_level is None or item_idx < this_level) else self.T_TEMPLATE
            new_legend_item, updated_pivot_item = self.make_legend_item(
                item_type=item_type,
                orig_pivot_item=orig_pivot_item,
                block_id=block_id,
            )
            if new_legend_item is not None:
                new_block_select_items.append(new_legend_item)
                this_dir_dimensions[item_idx] = updated_pivot_item

        for item_idx, orig_pivot_item in enumerate(other_dir_dimensions):
            item_type = self.T_DIMENSION if (other_level is None or item_idx < other_level) else self.T_TEMPLATE
            new_legend_item, updated_pivot_item = self.make_legend_item(
                item_type=item_type,
                orig_pivot_item=orig_pivot_item,
                block_id=block_id,
            )
            if new_legend_item is not None:
                new_block_select_items.append(new_legend_item)
                other_dir_dimensions[item_idx] = updated_pivot_item

        for item_idx, orig_pivot_item in enumerate(measures):
            if orig_pivot_item.role_spec.role == PivotRole.pivot_annotation:
                # Totals are not annotated
                continue
            item_type = self.T_MEASURE
            new_legend_item, updated_pivot_item = self.make_legend_item(
                item_type=item_type,
                orig_pivot_item=orig_pivot_item,
                block_id=block_id,
            )
            if new_legend_item is not None:
                new_block_select_items.append(new_legend_item)
                measures[item_idx] = updated_pivot_item

        return (
            new_block_select_items,
            this_dir_dimensions,
            other_dir_dimensions,
            measures,
        )


class PivotSpecNormalizer(RequestPartSpecNormalizerBase[RawPivotSpec]):
    DEFAULT_BLOCK_ID: ClassVar[int] = 0

    def _normalize_simple_totals_from_flag(self, pivot_spec: RawPivotSpec) -> RawPivotSpec:
        if pivot_spec.with_totals:
            pivot_totals: Optional[RawPivotTotalsSpec]
            if pivot_spec.totals is None:
                pivot_totals = RawPivotTotalsSpec(
                    rows=[RawPivotTotalsItemSpec(level=0)],
                    columns=[RawPivotTotalsItemSpec(level=0)],
                )
            else:
                pivot_totals = pivot_spec.totals

            pivot_spec = pivot_spec.clone(
                totals=pivot_totals,
                with_totals=None,
            )

        return pivot_spec

    def _normalize_totals(
        self,
        pivot_spec: RawPivotSpec,
        raw_query_spec_union: RawQuerySpecUnion,
        liid_gen: IdGenerator,
    ) -> tuple[RawPivotSpec, RawQuerySpecUnion]:
        raw_pivot_items = pivot_spec.structure
        totals = pivot_spec.totals
        assert totals is not None

        # Verify that there is only one block
        block_ids = raw_query_spec_union.get_unique_block_ids()
        if len(block_ids) > 1:
            raise dl_query_processing.exc.MultipleBlocksUnsupportedError(
                "Incoming multiple blocks not supported in simplified API"
            )
        base_block_id = next(iter(block_ids)) if block_ids else self.DEFAULT_BLOCK_ID
        block_id_gen = IdGenerator(used_ids={base_block_id})

        # Patch original items with explicit block_id
        select_specs: list[RawSelectFieldSpec] = []
        new_item: RawSelectFieldSpec
        for item in raw_query_spec_union.select_specs:
            if spec_is_field_name_pseudo_dimension(item):
                # Measure Names and such - they must not be block-specific
                new_item = item.clone(block_id=None)
            else:
                new_item = item.clone(block_id=base_block_id)
            select_specs.append(new_item)
        order_by_specs = [
            # All global ORDER BY items must be localized to the base block
            # (i.e. can't sort totals)
            item.clone(block_id=base_block_id)
            for item in raw_query_spec_union.order_by_specs
        ]

        def filter_raw_items(role: PivotRole) -> list[RawPivotLegendItem]:
            return [item for item in raw_pivot_items if item.role_spec.role == role]

        def normalize_target_legend_item_ids(item: RawPivotLegendItem) -> RawPivotLegendItem:
            anno_role_spec = item.role_spec
            assert isinstance(anno_role_spec, RawAnnotationPivotRoleSpec)
            if anno_role_spec.target_legend_item_ids is None:
                item = item.clone(
                    role_spec=anno_role_spec.clone(
                        target_legend_item_ids=measure_legend_item_ids,
                    )
                )
            return item

        row_dimensions = filter_raw_items(PivotRole.pivot_row)
        column_dimensions = filter_raw_items(PivotRole.pivot_column)
        strict_measures = filter_raw_items(PivotRole.pivot_measure)
        measure_legend_item_ids = sorted(
            set(itertools.chain.from_iterable(item.legend_item_ids for item in strict_measures))
        )
        annotations = [normalize_target_legend_item_ids(item) for item in filter_raw_items(PivotRole.pivot_annotation)]
        measures = strict_measures + annotations
        other_pivot_items = [
            item
            for item in raw_pivot_items
            if item.role_spec.role
            not in {
                PivotRole.pivot_row,
                PivotRole.pivot_column,
                PivotRole.pivot_measure,
                PivotRole.pivot_annotation,
            }
        ]

        # Check whether either `rows` or `columns` contains only Measure Names
        mname_liids = {item.legend_item_id for item in select_specs if spec_is_field_name_pseudo_dimension(item)}
        non_mname_row_dims = any(
            legend_item_id not in mname_liids for item in row_dimensions for legend_item_id in item.legend_item_ids
        )
        non_mname_column_dims = any(
            legend_item_id not in mname_liids for item in column_dimensions for legend_item_id in item.legend_item_ids
        )

        if not row_dimensions or not non_mname_row_dims:
            totals = totals.clone(rows=[])
        if not column_dimensions or not non_mname_column_dims:
            totals = totals.clone(columns=[])

        if not measures:
            # Totals are not possible in this case
            return pivot_spec, raw_query_spec_union

        helper = PivotTotalsNormalizerHelper(
            original_select_specs=select_specs,
            legend_item_id_gen=liid_gen,
            block_id_gen=block_id_gen,
            base_block_id=base_block_id,
        )

        # 1. "Horizontal" totals
        for row_total_item in totals.rows:
            add_select_specs, row_dimensions, column_dimensions, measures = helper.gen_single_dir_totals(
                this_level=row_total_item.level,
                other_level=None,
                this_dir_dimensions=row_dimensions,
                other_dir_dimensions=column_dimensions,
                measures=measures,
            )
            select_specs += add_select_specs

        # 2. "Vertical" totals
        for column_total_item in totals.columns:
            add_select_specs, column_dimensions, row_dimensions, measures = helper.gen_single_dir_totals(
                this_level=column_total_item.level,
                other_level=None,
                this_dir_dimensions=column_dimensions,
                other_dir_dimensions=row_dimensions,
                measures=measures,
            )
            select_specs += add_select_specs

        # 3. The intersections
        for row_total_item in totals.rows:
            for column_total_item in totals.columns:
                add_select_specs, row_dimensions, column_dimensions, measures = helper.gen_single_dir_totals(
                    this_level=row_total_item.level,
                    other_level=column_total_item.level,
                    this_dir_dimensions=row_dimensions,
                    other_dir_dimensions=column_dimensions,
                    measures=measures,
                )
                select_specs += add_select_specs

        # Make updated clones
        raw_query_spec_union = raw_query_spec_union.clone(
            select_specs=select_specs,
            order_by_specs=order_by_specs,
        )
        pivot_spec = pivot_spec.clone(
            structure=list(
                itertools.chain(
                    column_dimensions,
                    row_dimensions,
                    measures,
                    other_pivot_items,
                )
            ),
            totals=None,  # We've used it all up, so remove it
        )

        return pivot_spec, raw_query_spec_union

    def normalize_spec(
        self,
        spec: RawPivotSpec,
        raw_query_spec_union: RawQuerySpecUnion,
    ) -> tuple[RawPivotSpec, RawQuerySpecUnion]:
        spec, raw_query_spec_union = super().normalize_spec(
            spec=spec,
            raw_query_spec_union=raw_query_spec_union,
        )
        pivot_spec = spec

        liid_gen = self.get_liid_generator(raw_query_spec_union=raw_query_spec_union)

        if pivot_spec.with_totals:
            pivot_spec = self._normalize_simple_totals_from_flag(pivot_spec=pivot_spec)

        if pivot_spec.totals is not None:
            pivot_spec, raw_query_spec_union = self._normalize_totals(
                pivot_spec=pivot_spec,
                raw_query_spec_union=raw_query_spec_union,
                liid_gen=liid_gen,
            )

        return pivot_spec, raw_query_spec_union


class PivotRequestModelNormalizer(RequestModelNormalizerBase[PivotDataRequestModel]):
    """
    Pivot-specific normalizer.

    Expands:
    - simplified totals API
    """

    def normalize_drm(self, drm: PivotDataRequestModel) -> PivotDataRequestModel:
        drm = super().normalize_drm(drm=drm)

        pivot_spec_normalizer = PivotSpecNormalizer()
        pivot_spec, raw_query_spec_union = pivot_spec_normalizer.normalize_spec(
            spec=drm.pivot,
            raw_query_spec_union=drm.raw_query_spec_union,
        )
        if pivot_spec is not drm.pivot or raw_query_spec_union is not drm.raw_query_spec_union:
            drm = drm.clone(
                raw_query_spec_union=raw_query_spec_union,
                pivot=pivot_spec,
            )

        return drm

lib/dl_api_lib/dl_api_lib/request_model/normalization/drm_normalizer_pivot.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Any,
    Dict,
)

from marshmallow import EXCLUDE
from marshmallow import fields as ma_fields
from marshmallow_oneofschema import OneOfSchema

from dl_api_connector.api_schema.source import DataSourceSchema
from dl_api_connector.api_schema.source_base import (
    AvatarRelationSchema,
    SourceAvatarSchema,
)
from dl_api_lib.enums import DatasetAction
from dl_api_lib.request_model.data import (
    AddField,
    AddUpdateObligatoryFilterAction,
    AddUpdateSourceAction,
    AvatarActionBase,
    CloneField,
    DeleteField,
    DeleteObligatoryFilter,
    DeleteObligatoryFilterAction,
    FieldAction,
    RelationActionBase,
    ReplaceConnection,
    ReplaceConnectionAction,
    SourceActionBase,
    UpdateField,
)
from dl_api_lib.schemas.filter import ObligatoryFilterSchema
from dl_api_lib.schemas.parameters import ParameterValueConstraintSchema
from dl_api_lib.schemas.values import (
    ValueSchema,
    WithNestedValueSchema,
)
from dl_constants.enums import (
    AggregationFunction,
    CalcMode,
    UserDataType,
)
from dl_model_tools.schema.base import (
    BaseSchema,
    DefaultValidateSchema,
)


LOGGER = logging.getLogger(__name__)


class ActionBaseSchema(BaseSchema):
    action = ma_fields.Enum(DatasetAction, required=True)
    order_index = ma_fields.Integer(load_default=0, required=False)


class FieldActionBaseSchema(ActionBaseSchema, DefaultValidateSchema[FieldAction]):
    class FieldBaseSchema(BaseSchema):
        guid = ma_fields.String()
        strict = ma_fields.Boolean(required=False, load_default=False)

    field = ma_fields.Nested(FieldBaseSchema, required=True)


class UpdateFieldActionSchema(FieldActionBaseSchema, DefaultValidateSchema[FieldAction]):
    TARGET_CLS = FieldAction

    class UpdateFieldBaseSchema(WithNestedValueSchema, FieldActionBaseSchema.FieldBaseSchema):
        TYPE_FIELD_NAME = "cast"
        source = ma_fields.String()
        calc_mode = ma_fields.Enum(CalcMode)
        hidden = ma_fields.Boolean()
        description = ma_fields.String()
        aggregation = ma_fields.Enum(AggregationFunction)
        formula = ma_fields.String()
        guid_formula = ma_fields.String()
        cast = ma_fields.Enum(UserDataType, allow_none=True)
        avatar_id = ma_fields.String(allow_none=True)
        new_id = ma_fields.String(allow_none=True)
        default_value = ma_fields.Nested(ValueSchema, allow_none=True)
        value_constraint = ma_fields.Nested(ParameterValueConstraintSchema, allow_none=True)

    class UpdateFieldSchema(UpdateFieldBaseSchema, DefaultValidateSchema[UpdateField]):
        TARGET_CLS = UpdateField
        title = ma_fields.String()

    field = ma_fields.Nested(UpdateFieldSchema, required=False)


class AddFieldActionSchema(UpdateFieldActionSchema):
    TARGET_CLS = FieldAction

    class AddFieldSchema(UpdateFieldActionSchema.UpdateFieldBaseSchema, DefaultValidateSchema[AddField]):
        TARGET_CLS = AddField
        title = ma_fields.String(required=True)

    field = ma_fields.Nested(AddFieldSchema, required=False)


class CloneFieldActionSchema(FieldActionBaseSchema, DefaultValidateSchema[FieldAction]):
    TARGET_CLS = FieldAction

    class CloneFieldSchema(FieldActionBaseSchema.FieldBaseSchema, DefaultValidateSchema[CloneField]):
        TARGET_CLS = CloneField

        title = ma_fields.String()
        from_guid = ma_fields.String()
        aggregation = ma_fields.Enum(AggregationFunction, allow_none=True)
        cast = ma_fields.Enum(UserDataType, allow_none=True)

    field = ma_fields.Nested(CloneFieldSchema, required=True)


class DeleteFieldActionSchema(FieldActionBaseSchema):
    TARGET_CLS = FieldAction

    class DeleteFieldSchema(FieldActionBaseSchema.FieldBaseSchema, DefaultValidateSchema[DeleteField]):
        TARGET_CLS = DeleteField

    field = ma_fields.Nested(DeleteFieldSchema, required=True)


class SourceActionBaseSchema(ActionBaseSchema, DefaultValidateSchema[SourceActionBase]):
    class SourceBaseSchema(BaseSchema):
        id = ma_fields.String()

    source = ma_fields.Nested(SourceBaseSchema, required=True)


class AddUpdateSourceActionSchema(SourceActionBaseSchema, DefaultValidateSchema[SourceActionBase]):
    TARGET_CLS = AddUpdateSourceAction

    source = ma_fields.Nested(DataSourceSchema, required=True)


class DeleteSourceActionSchema(SourceActionBaseSchema, DefaultValidateSchema[SourceActionBase]):
    TARGET_CLS = SourceActionBase

    pass


class RefreshSourceActionSchema(SourceActionBaseSchema, DefaultValidateSchema[SourceActionBase]):
    TARGET_CLS = SourceActionBase

    class RefreshSourceSchema(SourceActionBaseSchema.SourceBaseSchema):
        force_update_fields = ma_fields.Boolean(load_default=False)

    source = ma_fields.Nested(RefreshSourceSchema, required=True)


class AvatarActionBaseSchema(ActionBaseSchema):
    class AvatarBaseSchema(BaseSchema):
        id = ma_fields.String()

    source_avatar = ma_fields.Nested(AvatarBaseSchema, required=True)
    # Should be used by external API only
    disable_fields_update = ma_fields.Boolean(load_default=False, dump_default=False, required=False, allow_none=False)


class AddUpdateAvatarActionSchema(AvatarActionBaseSchema, DefaultValidateSchema[AvatarActionBase]):
    TARGET_CLS = AvatarActionBase
    source_avatar = ma_fields.Nested(SourceAvatarSchema, required=True)


class DeleteAvatarActionSchema(AvatarActionBaseSchema, DefaultValidateSchema[AvatarActionBase]):
    TARGET_CLS = AvatarActionBase
    pass


class RelationActionBaseSchema(ActionBaseSchema):
    class RelationBaseSchema(BaseSchema):
        id = ma_fields.String()

    avatar_relation = ma_fields.Nested(RelationBaseSchema, required=True)


class AddUpdateRelationActionSchema(RelationActionBaseSchema, DefaultValidateSchema[RelationActionBase]):
    TARGET_CLS = RelationActionBase
    avatar_relation = ma_fields.Nested(AvatarRelationSchema, required=True)


class DeleteRelationActionSchema(RelationActionBaseSchema, DefaultValidateSchema[RelationActionBase]):
    TARGET_CLS = RelationActionBase
    pass


class ReplaceConnectionActionSchema(ActionBaseSchema, DefaultValidateSchema[ReplaceConnectionAction]):
    TARGET_CLS = ReplaceConnectionAction

    class ReplaceConnectionSchema(DefaultValidateSchema[ReplaceConnection]):
        TARGET_CLS = ReplaceConnection
        id = ma_fields.String(allow_none=False, required=True)
        new_id = ma_fields.String(allow_none=False, required=True)

    connection = ma_fields.Nested(ReplaceConnectionSchema, required=False)


class ObligatoryFilterActionBaseSchema(ActionBaseSchema):
    class ObligatoryFilterBaseSchema(BaseSchema):
        id = ma_fields.String(allow_none=False, required=True)

    obligatory_filter = ma_fields.Nested(ObligatoryFilterBaseSchema, required=True)


class AddUpdateObligatoryFilterActionSchema(
    ObligatoryFilterActionBaseSchema, DefaultValidateSchema[AddUpdateObligatoryFilterAction]
):
    TARGET_CLS = AddUpdateObligatoryFilterAction
    obligatory_filter = ma_fields.Nested(ObligatoryFilterSchema, required=True)


class DeleteObligatoryFilterActionSchema(
    ObligatoryFilterActionBaseSchema, DefaultValidateSchema[DeleteObligatoryFilterAction]
):
    TARGET_CLS = DeleteObligatoryFilterAction

    class DeleteObligatoryFilterSchema(
        ObligatoryFilterActionBaseSchema.ObligatoryFilterBaseSchema, DefaultValidateSchema[DeleteObligatoryFilter]
    ):
        TARGET_CLS = DeleteObligatoryFilter
        id = ma_fields.String(allow_none=False, required=True)

    obligatory_filter = ma_fields.Nested(DeleteObligatoryFilterSchema, required=True)


class ActionSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "action"
    type_schemas = {
        # legacy  TODO: remove
        DatasetAction.add.name: AddFieldActionSchema,
        DatasetAction.update.name: UpdateFieldActionSchema,
        DatasetAction.delete.name: DeleteFieldActionSchema,
        # fields
        DatasetAction.add_field.name: AddFieldActionSchema,
        DatasetAction.update_field.name: UpdateFieldActionSchema,
        DatasetAction.delete_field.name: DeleteFieldActionSchema,
        DatasetAction.clone_field.name: CloneFieldActionSchema,
        # sources
        DatasetAction.add_source.name: AddUpdateSourceActionSchema,
        DatasetAction.update_source.name: AddUpdateSourceActionSchema,
        DatasetAction.delete_source.name: DeleteSourceActionSchema,
        DatasetAction.refresh_source.name: RefreshSourceActionSchema,
        # avatars
        DatasetAction.add_source_avatar.name: AddUpdateAvatarActionSchema,
        DatasetAction.update_source_avatar.name: AddUpdateAvatarActionSchema,
        DatasetAction.delete_source_avatar.name: DeleteAvatarActionSchema,
        # relations
        DatasetAction.add_avatar_relation.name: AddUpdateRelationActionSchema,
        DatasetAction.update_avatar_relation.name: AddUpdateRelationActionSchema,
        DatasetAction.delete_avatar_relation.name: DeleteRelationActionSchema,
        # connections
        DatasetAction.replace_connection.name: ReplaceConnectionActionSchema,
        # filters
        DatasetAction.add_obligatory_filter.name: AddUpdateObligatoryFilterActionSchema,
        DatasetAction.update_obligatory_filter.name: AddUpdateObligatoryFilterActionSchema,
        DatasetAction.delete_obligatory_filter.name: DeleteObligatoryFilterActionSchema,
    }

    def get_obj_type(self, obj: Dict[str, Any]) -> str:
        return obj[self.type_field].name

lib/dl_api_lib/dl_api_lib/schemas/action.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Any,
    ClassVar,
    Dict,
    Type,
)

from marshmallow import fields as ma_fields
from marshmallow_oneofschema import OneOfSchema

from dl_api_connector.api_schema.connection_base import (
    ConnectionMetaMixin,
    ConnectionSchema,
)
from dl_api_connector.api_schema.source import DataSourceSchema
from dl_api_connector.api_schema.source_base import (
    DataSourceTemplateResponseField,
    RawSchemaColumnSchema,
)
from dl_constants.enums import ConnectionType as CT
from dl_core.us_connection_base import (
    ConnectionBase,
    UnknownConnection,
)
from dl_model_tools.schema.base import BaseSchema


LOGGER = logging.getLogger(__name__)


class ConnectionSourcesQuerySchema(BaseSchema):
    search_text = ma_fields.String(required=False, load_default=None)
    limit = ma_fields.Integer(required=False, load_default=10_000)


class ConnectionInfoSourceSchemaQuerySchema(BaseSchema):
    source = ma_fields.Nested(DataSourceSchema, required=True)


class ConnectionSourceTemplatesResponseSchema(BaseSchema):
    # # Complete:
    # sources = ma_fields.Nested(DataSourceTemplateResponseSchema, many=True)
    # freeform_sources = ma_fields.Nested(DataSourceTemplateResponseSchema, many=True)
    # # Minimal processing:
    sources = ma_fields.List(DataSourceTemplateResponseField)
    freeform_sources = ma_fields.List(DataSourceTemplateResponseField)


class ConnectionInfoSourceSchemaResponseSchema(BaseSchema):
    raw_schema = ma_fields.Nested(RawSchemaColumnSchema, many=True, allow_none=True)


class GenericConnectionSchema(OneOfSchema):
    type_schemas: dict[str, Type[ConnectionSchema]] = {}
    supported_connections: ClassVar[set[Type[ConnectionBase]]] = set()

    def get_data_type(self, data):  # type: ignore
        data_type = super().get_data_type(data)
        self.context[ConnectionSchema.CONN_TYPE_CTX_KEY] = data_type
        return data_type

    @classmethod
    def is_supported_connection(cls, conn: ConnectionBase) -> bool:
        return type(conn) in cls.supported_connections

    @classmethod
    def is_supported_type_discriminator(cls, type_discriminator: str) -> bool:
        return type_discriminator in cls.type_schemas

    def get_obj_type(self, obj: ConnectionBase) -> str:
        return obj.conn_type.name

    def get_edit_schema_cls(self, obj: ConnectionBase) -> Type[ConnectionSchema]:
        """
        Returns schema for connection editing
        :param obj: connection to modify
        :return: Schema class that will modify connection on `load`
        """
        type_discriminator = self.get_obj_type(obj)
        return self.type_schemas[type_discriminator]

    # TODO FIX: remove after type discriminator key will be normalized now we use 'db_type' in GET and 'type' in POST
    def _dump(self, obj: ConnectionBase, *, update_fields: bool = True, **kwargs: Any) -> Dict[str, Any]:
        ret = super()._dump(obj, update_fields=update_fields, **kwargs)
        ret.pop(self.type_field)
        return ret


def register_sub_schema_class(conn_type: CT, schema_cls: Type[ConnectionSchema]) -> None:
    if conn_type.name in GenericConnectionSchema.type_schemas:
        registered_schema = GenericConnectionSchema.type_schemas[conn_type.name]
        assert (
            registered_schema == schema_cls
        ), f"Already have {registered_schema} serving {conn_type}, can not register {schema_cls} for it"

    GenericConnectionSchema.type_schemas[conn_type.name] = schema_cls
    GenericConnectionSchema.supported_connections.add(schema_cls.TARGET_CLS)


class UnknownConnectionSchema(ConnectionMetaMixin, ConnectionSchema):
    TARGET_CLS = UnknownConnection


register_sub_schema_class(CT.unknown, UnknownConnectionSchema)

lib/dl_api_lib/dl_api_lib/schemas/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import Enum
import json
from typing import (
    Any,
    ClassVar,
    Dict,
    List,
)

from marshmallow import (
    pre_dump,
    pre_load,
)
from marshmallow import EXCLUDE
from marshmallow import fields as ma_fields
from marshmallow_oneofschema import OneOfSchema

from dl_api_lib.query.formalization.raw_pivot_specs import RawPivotSpec
from dl_api_lib.query.formalization.raw_specs import (
    DimensionNameRef,
    IdFieldRef,
    IdOrTitleFieldRef,
    MeasureNameRef,
    PlaceholderRef,
    RawAfterBlockPlacement,
    RawBlockSpec,
    RawDimensionValueSpec,
    RawFilterFieldSpec,
    RawGroupByFieldSpec,
    RawOrderByFieldSpec,
    RawOrderByRoleSpec,
    RawParameterValueSpec,
    RawQueryMetaInfo,
    RawQuerySpecUnion,
    RawRangeRoleSpec,
    RawResultSpec,
    RawRoleSpec,
    RawRootBlockPlacement,
    RawRowRoleSpec,
    RawSelectFieldSpec,
    RawTemplateRoleSpec,
    RawTreeRoleSpec,
    TitleFieldRef,
)
from dl_api_lib.request_model.data import (
    DataRequestModel,
    PivotDataRequestModel,
    ResultDataRequestModel,
)
from dl_api_lib.schemas.action import ActionSchema
from dl_api_lib.schemas.dataset_base import DatasetContentSchema
from dl_api_lib.schemas.filter import WhereSchema
from dl_api_lib.schemas.legend import LegendItemSchema
from dl_api_lib.schemas.pivot import RequestPivotSpecSchema
from dl_constants.enums import (
    CalcMode,
    FieldRole,
    FieldType,
    FieldVisibility,
    NotificationLevel,
    OrderDirection,
    QueryBlockPlacementType,
    QueryItemRefType,
    RangeType,
    UserDataType,
    WhereClauseOperation,
)
from dl_core.constants import DataAPILimits
from dl_model_tools.schema.base import (
    BaseSchema,
    DefaultSchema,
)
from dl_query_processing.enums import (
    GroupByPolicy,
    QueryType,
)
from dl_query_processing.legend.block_legend import BlockPlacement


class DatasetDataRequestBaseSchema(DefaultSchema[DataRequestModel]):
    TARGET_CLS = DataRequestModel

    def to_object(self, data: dict) -> DataRequestModel:
        raw_query_spec_union = self._make_raw_query_spec_union(data)
        drm = self._make_drm(raw_query_spec_union=raw_query_spec_union, data=data)
        return drm

    def _make_raw_query_spec_union(self, data: Dict[str, Any]) -> RawQuerySpecUnion:
        raise NotImplementedError

    def _make_drm(self, raw_query_spec_union: RawQuerySpecUnion, data: Dict[str, Any]) -> DataRequestModel:
        raise NotImplementedError


class DatasetPreviewRequestSchema(DatasetDataRequestBaseSchema, DatasetContentSchema):
    limit = ma_fields.Integer(load_default=100)
    row_count_hard_limit = ma_fields.Integer(load_default=10**5)
    updates = ma_fields.Nested(ActionSchema, many=True, load_default=[])
    TARGET_CLS = DataRequestModel

    def _make_raw_query_spec_union(self, data: Dict[str, Any]) -> RawQuerySpecUnion:
        raw_query_spec_union = RawQuerySpecUnion(
            limit=data.get("limit"),
            group_by_policy=GroupByPolicy.if_measures,
            meta=RawQueryMetaInfo(
                query_type=QueryType.preview,
                row_count_hard_limit=data["row_count_hard_limit"],
            ),
        )
        return raw_query_spec_union

    def _make_drm(self, raw_query_spec_union: RawQuerySpecUnion, data: Dict[str, Any]) -> DataRequestModel:
        return DataRequestModel(
            raw_query_spec_union=raw_query_spec_union,
            dataset=data.get("dataset"),
            updates=data["updates"],
        )


class FieldsResponseFieldSchema(BaseSchema):
    title = ma_fields.String()
    guid = ma_fields.String()
    data_type = ma_fields.Enum(UserDataType)
    hidden = ma_fields.Boolean()
    type = ma_fields.Enum(FieldType)
    calc_mode = ma_fields.Enum(CalcMode)


class DatasetFieldsResponseSchema(BaseSchema):
    fields_ = ma_fields.Nested(FieldsResponseFieldSchema, attribute="fields", data_key="fields", many=True)


class YqlDataSchema(BaseSchema):
    Type = ma_fields.Raw()
    Data = ma_fields.List(ma_fields.List(ma_fields.Raw()))


class OrderBySchema(DefaultSchema[RawOrderByFieldSpec]):
    class ApiOrderingDirection(Enum):
        DESC = "desc"
        ASC = "asc"

    TARGET_CLS = RawOrderByFieldSpec

    column = ma_fields.String(required=True)
    direction = ma_fields.Enum(ApiOrderingDirection, required=True)

    def to_object(self, data: dict) -> RawOrderByFieldSpec:
        data["direction"] = {
            self.ApiOrderingDirection.ASC: OrderDirection.asc,
            self.ApiOrderingDirection.DESC: OrderDirection.desc,
        }[data["direction"]]
        data["ref"] = IdOrTitleFieldRef(id_or_title=data.pop("column"))
        return self.get_target_cls()(**data)


class DatasetVersionResultRequestSchema(DatasetDataRequestBaseSchema):
    columns = ma_fields.List(ma_fields.String(), required=True)
    where = ma_fields.Nested(WhereSchema, many=True, load_default=[])
    group_by = ma_fields.List(ma_fields.String(), load_default=[])
    order_by = ma_fields.Nested(OrderBySchema, many=True, load_default=[])
    limit = ma_fields.Integer(load_default=None)
    offset = ma_fields.Integer(load_default=None)
    updates = ma_fields.Nested(ActionSchema, many=True, load_default=[])
    row_count_hard_limit = ma_fields.Integer(load_default=DataAPILimits.DATA_API_DEFAULT_ROW_COUNT_HARD_LIMIT)
    with_totals = ma_fields.Boolean(load_default=False)
    disable_group_by = ma_fields.Boolean(load_default=False)
    add_fields_data = ma_fields.Boolean(load_default=True)
    ignore_nonexistent_filters = ma_fields.Boolean(load_default=False)
    revision_id = ma_fields.String(load_default=None)

    def _make_raw_query_spec_union(self, data: Dict[str, Any]) -> RawQuerySpecUnion:
        raw_query_spec_union = RawQuerySpecUnion(
            select_specs=[RawSelectFieldSpec(ref=IdOrTitleFieldRef(id_or_title=field)) for field in data["columns"]],
            group_by_specs=[
                RawGroupByFieldSpec(ref=IdOrTitleFieldRef(id_or_title=field)) for field in data["group_by"]
            ],
            filter_specs=[
                RawFilterFieldSpec(ref=filter_spec.ref, operation=filter_spec.operation, values=filter_spec.values)
                for filter_spec in data["where"]
            ],
            order_by_specs=data["order_by"],
            limit=data["limit"],
            offset=data["offset"],
            group_by_policy=GroupByPolicy.disable if data.get("disable_group_by") else GroupByPolicy.force,
            ignore_nonexistent_filters=data["ignore_nonexistent_filters"],
            meta=RawQueryMetaInfo(
                query_type=QueryType.result,
                row_count_hard_limit=data["row_count_hard_limit"],
            ),
        )
        return raw_query_spec_union

    def _make_drm(self, raw_query_spec_union: RawQuerySpecUnion, data: Dict[str, Any]) -> DataRequestModel:
        return DataRequestModel(
            raw_query_spec_union=raw_query_spec_union,
            updates=data["updates"],
            with_totals=data["with_totals"],
            add_fields_data=data["add_fields_data"],
            dataset_revision_id=data.get("revision_id"),
        )


class DatasetVersionResultResponseSchema(BaseSchema):
    class Result(BaseSchema):
        query = ma_fields.String(required=False)
        data = ma_fields.Nested(YqlDataSchema)

        # Rare flag, thus making it possible to not return it on every
        # unrelated request.
        not_aggregated = ma_fields.Boolean(required=False)

        totals_query = ma_fields.String(required=False)
        totals = ma_fields.List(ma_fields.String())

        data_export_forbidden = ma_fields.Boolean(required=False)
        fields_ = ma_fields.Nested(
            FieldsResponseFieldSchema, attribute="fields", data_key="fields", many=True, required=False
        )

    result = ma_fields.Nested(Result)


class DatasetVersionValuesBasePostSchema(DatasetDataRequestBaseSchema):
    QUERY_TYPE: ClassVar[QueryType]

    field_guid = ma_fields.String(required=True)
    where = ma_fields.Nested(WhereSchema, many=True, load_default=[])
    updates = ma_fields.Nested(ActionSchema, many=True, load_default=[])
    revision_id = ma_fields.String(load_default=None)

    def _make_select_specs(self, field_id: str) -> List[RawSelectFieldSpec]:
        return [RawSelectFieldSpec(ref=IdFieldRef(id=field_id))]

    def _make_raw_query_spec_union(self, data: Dict[str, Any]) -> RawQuerySpecUnion:
        raw_query_spec_union = RawQuerySpecUnion(
            select_specs=self._make_select_specs(field_id=data["field_guid"]),
            order_by_specs=[],
            filter_specs=[
                RawFilterFieldSpec(ref=filter_spec.ref, operation=filter_spec.operation, values=filter_spec.values)
                for filter_spec in data["where"]
            ],
            limit=data.get("limit"),
            offset=data.get("offset"),
            group_by_policy=GroupByPolicy.disable,
            disable_rls=data.get("disable_rls", False),
            ignore_nonexistent_filters=data.get("ignore_nonexistent_filters", False),
            allow_measure_fields=False,  # Cannot get distincts or ranges for measures
            meta=RawQueryMetaInfo(
                query_type=self.QUERY_TYPE,
            ),
        )
        return raw_query_spec_union


class DatasetVersionValuesDistinctPostSchema(DatasetVersionValuesBasePostSchema):
    QUERY_TYPE = QueryType.distinct

    row_count_hard_limit = ma_fields.Integer(load_default=DataAPILimits.DATA_API_DEFAULT_ROW_COUNT_HARD_LIMIT)
    limit = ma_fields.Integer(load_default=1000)
    offset = ma_fields.Integer(load_default=0)
    ignore_nonexistent_filters = ma_fields.Boolean(load_default=False)
    disable_rls = ma_fields.Boolean(load_default=False)

    def _make_select_specs(self, field_id: str) -> List[RawSelectFieldSpec]:
        # Right now DISTINCT is handled at query level,
        # so the expressions are compiled in the regular manner
        return [RawSelectFieldSpec(ref=IdFieldRef(id=field_id), role_spec=RawRoleSpec(role=FieldRole.distinct))]

    def _make_drm(self, raw_query_spec_union: RawQuerySpecUnion, data: Dict[str, Any]) -> DataRequestModel:
        return DataRequestModel(
            raw_query_spec_union=raw_query_spec_union,
            updates=data["updates"],
            dataset_revision_id=data.get("revision_id"),
        )


class DatasetVersionValuesDistinctResponseSchema(DatasetVersionResultResponseSchema):
    pass


class DatasetVersionValuesRangePostSchema(DatasetVersionValuesBasePostSchema):
    QUERY_TYPE = QueryType.value_range

    def _make_select_specs(self, field_id: str) -> List[RawSelectFieldSpec]:
        return [
            RawSelectFieldSpec(
                ref=IdFieldRef(id=field_id),
                role_spec=RawRangeRoleSpec(role=FieldRole.range, range_type=RangeType.min),
            ),
            RawSelectFieldSpec(
                ref=IdFieldRef(id=field_id),
                role_spec=RawRangeRoleSpec(role=FieldRole.range, range_type=RangeType.max),
            ),
        ]

    def _make_drm(self, raw_query_spec_union: RawQuerySpecUnion, data: Dict[str, Any]) -> DataRequestModel:
        return DataRequestModel(
            raw_query_spec_union=raw_query_spec_union,
            updates=data["updates"],
            dataset_revision_id=data.get("revision_id"),
        )


class DatasetVersionValuesRangeResponseSchema(DatasetVersionResultResponseSchema):
    pass


class IdRefSchema(DefaultSchema[IdFieldRef]):
    TARGET_CLS = IdFieldRef

    id = ma_fields.String(required=True)


class TitleRefSchema(DefaultSchema[TitleFieldRef]):
    TARGET_CLS = TitleFieldRef

    title = ma_fields.String(required=True)


class MeasureNameRefSchema(DefaultSchema[MeasureNameRef]):
    TARGET_CLS = MeasureNameRef


class DimensionNameRefSchema(DefaultSchema[DimensionNameRef]):
    TARGET_CLS = DimensionNameRef


class PlaceholderRefSchema(DefaultSchema[PlaceholderRef]):
    TARGET_CLS = PlaceholderRef


class ItemRefSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = True
    type_field = "type"

    type_schemas: Dict[str, Any] = {
        QueryItemRefType.id.name: IdRefSchema,
        QueryItemRefType.title.name: TitleRefSchema,
        QueryItemRefType.measure_name.name: MeasureNameRefSchema,
        QueryItemRefType.dimension_name.name: DimensionNameRefSchema,
        QueryItemRefType.placeholder.name: PlaceholderRefSchema,
    }


class DimensionValueSpecSchema(DefaultSchema[RawDimensionValueSpec]):
    TARGET_CLS = RawDimensionValueSpec

    legend_item_id = ma_fields.Integer(required=True)
    value = ma_fields.Raw(required=True)


class RoleSpecSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "role"

    class RoleSpecSchemaVariant(DefaultSchema[RawRoleSpec]):
        TARGET_CLS = RawRoleSpec

        role = ma_fields.Enum(FieldRole)

    class OrderByRoleSpecSchemaVariant(DefaultSchema[RawOrderByRoleSpec]):
        TARGET_CLS = RawOrderByRoleSpec

        role = ma_fields.Enum(FieldRole)
        direction = ma_fields.Enum(OrderDirection)

    class RangeRoleSpecSchemaVariant(DefaultSchema[RawRangeRoleSpec]):
        TARGET_CLS = RawRangeRoleSpec

        role = ma_fields.Enum(FieldRole)
        range_type = ma_fields.Enum(RangeType)

    class RowRoleSpecSchemaVariant(DefaultSchema[RawRowRoleSpec]):
        TARGET_CLS = RawRowRoleSpec

        role = ma_fields.Enum(FieldRole)
        visibility = ma_fields.Enum(FieldVisibility)

    class TemplateRoleSpecSchemaVariant(DefaultSchema[RawTemplateRoleSpec]):
        TARGET_CLS = RawTemplateRoleSpec

        role = ma_fields.Enum(FieldRole)
        template = ma_fields.String()
        visibility = ma_fields.Enum(FieldVisibility)

    class TreeRoleSpecSchemaVariant(DefaultSchema[RawTreeRoleSpec]):
        TARGET_CLS = RawTreeRoleSpec

        role = ma_fields.Enum(FieldRole)
        level = ma_fields.Integer(allow_none=False)
        prefix = ma_fields.String(allow_none=False)
        dimension_values = ma_fields.Nested(DimensionValueSpecSchema, many=True, allow_none=False)
        visibility = ma_fields.Enum(FieldVisibility)

        @pre_dump
        def dump_prefix(self, data: dict, *args: Any, **kwargs: Any) -> dict:
            data["prefix"] = json.dumps(data["prefix"])
            return data

    type_schemas = {
        FieldRole.row.name: RowRoleSpecSchemaVariant,
        FieldRole.measure.name: RoleSpecSchemaVariant,
        FieldRole.info.name: RoleSpecSchemaVariant,
        FieldRole.order_by.name: OrderByRoleSpecSchemaVariant,
        FieldRole.distinct.name: RoleSpecSchemaVariant,
        FieldRole.range.name: RangeRoleSpecSchemaVariant,
        FieldRole.total.name: RoleSpecSchemaVariant,
        FieldRole.template.name: TemplateRoleSpecSchemaVariant,
        FieldRole.tree.name: TreeRoleSpecSchemaVariant,
    }


class _FlattenRefMixin(BaseSchema):
    @pre_load(pass_many=False)
    def nest_ref_obj(self, data: Dict[str, Any], **_: Any) -> Dict[str, Any]:
        if "ref" not in data:
            ref_type = QueryItemRefType[data["ref_type"]]
            ref_data: Dict[str, Any] = {"type": data["ref_type"]}
            if ref_type == QueryItemRefType.id:
                ref_data["id"] = data.get("id")
            elif ref_type == QueryItemRefType.title:
                ref_data["title"] = data.get("title")
            elif ref_type == QueryItemRefType.measure_name:
                pass
            elif ref_type == QueryItemRefType.placeholder:
                pass
            data["ref"] = ref_data
        return data


class QueryFieldsItemSchema(DefaultSchema[RawSelectFieldSpec], _FlattenRefMixin):
    TARGET_CLS = RawSelectFieldSpec

    ref = ma_fields.Nested(ItemRefSchema, allow_none=False)
    legend_item_id = ma_fields.Integer(allow_none=True)
    role_spec = ma_fields.Nested(RoleSpecSchema, allow_none=False)
    label = ma_fields.String(allow_none=True)
    block_id = ma_fields.Integer(allow_none=True)


class QueryOrderByItemSchema(DefaultSchema[RawOrderByFieldSpec], _FlattenRefMixin):
    TARGET_CLS = RawOrderByFieldSpec

    ref = ma_fields.Nested(ItemRefSchema, allow_none=False)
    direction = ma_fields.Enum(OrderDirection, required=True)
    block_id = ma_fields.Integer(allow_none=True)


class QueryFiltersItemSchema(DefaultSchema[RawFilterFieldSpec], _FlattenRefMixin):
    TARGET_CLS = RawFilterFieldSpec

    ref = ma_fields.Nested(ItemRefSchema, allow_none=False)
    operation = ma_fields.Enum(WhereClauseOperation, required=True)
    values = ma_fields.List(ma_fields.Raw(allow_none=True), required=True, allow_none=True)
    block_id = ma_fields.Integer(allow_none=True)


class RootBlockPlacementSchema(DefaultSchema[RawRootBlockPlacement]):
    TARGET_CLS = RawRootBlockPlacement

    type = ma_fields.Enum(QueryBlockPlacementType)


class AfterBlockPlacementSchema(DefaultSchema[RawAfterBlockPlacement]):
    TARGET_CLS = RawAfterBlockPlacement

    type = ma_fields.Enum(QueryBlockPlacementType)
    dimension_values = ma_fields.Nested(DimensionValueSpecSchema, many=True, allow_none=True)


class BlockPlacementSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = True
    type_field = "type"

    type_schemas = {
        QueryBlockPlacementType.root.name: RootBlockPlacementSchema,
        QueryBlockPlacementType.after.name: AfterBlockPlacementSchema,
    }

    def get_obj_type(self, obj: Any) -> str:
        assert isinstance(obj, BlockPlacement)
        return obj.type.name


class QueryBlockSchema(DefaultSchema[RawBlockSpec]):
    TARGET_CLS = RawBlockSpec

    block_id = ma_fields.Integer(allow_none=False)
    parent_block_id = ma_fields.Integer(allow_none=True, load_default=None)
    placement = ma_fields.Nested(BlockPlacementSchema, allow_none=False)
    limit = ma_fields.Integer(load_default=None)
    offset = ma_fields.Integer(load_default=None)
    row_count_hard_limit = ma_fields.Integer(load_default=DataAPILimits.DATA_API_DEFAULT_ROW_COUNT_HARD_LIMIT)


class ParameterValueSchema(DefaultSchema[RawParameterValueSpec], _FlattenRefMixin):
    TARGET_CLS = RawParameterValueSpec

    ref = ma_fields.Nested(ItemRefSchema, allow_none=False)
    value = ma_fields.Raw(required=True, allow_none=False)
    block_id = ma_fields.Integer(allow_none=True)


class NewDatasetDataRequestBaseSchema(DatasetDataRequestBaseSchema):
    QUERY_TYPE: ClassVar[QueryType]

    # Query spec
    sel_fields = ma_fields.Nested(
        QueryFieldsItemSchema,
        data_key="fields",
        attribute="fields",
        many=True,
        required=True,
    )
    order_by = ma_fields.Nested(QueryOrderByItemSchema, many=True, load_default=[])
    filters = ma_fields.Nested(QueryFiltersItemSchema, many=True, load_default=[])
    limit = ma_fields.Integer(load_default=None)
    offset = ma_fields.Integer(load_default=None)
    disable_group_by = ma_fields.Boolean(load_default=False)
    row_count_hard_limit = ma_fields.Integer(load_default=DataAPILimits.DATA_API_DEFAULT_ROW_COUNT_HARD_LIMIT)
    ignore_nonexistent_filters = ma_fields.Boolean(load_default=False)
    blocks = ma_fields.Nested(QueryBlockSchema, many=True, load_default=[])
    # Updates
    updates = ma_fields.Nested(ActionSchema, many=True)
    # Misc
    autofill_legend = ma_fields.Boolean(load_default=False)
    revision_id = ma_fields.String(load_default=None)
    # Parameters
    parameter_values = ma_fields.Nested(ParameterValueSchema, many=True, load_default=[])

    def _make_select_specs(self, data: Dict[str, Any]) -> List[RawSelectFieldSpec]:
        return data["fields"]

    def _make_raw_query_spec_union(self, data: Dict[str, Any]) -> RawQuerySpecUnion:
        raw_query_spec_union = RawQuerySpecUnion(
            select_specs=self._make_select_specs(data),
            order_by_specs=data.get("order_by", []),
            filter_specs=data.get("filters", []),
            parameter_value_specs=data.get("parameter_values", []),
            block_specs=data.get("blocks", []),
            limit=data.get("limit"),
            offset=data.get("offset"),
            group_by_policy=GroupByPolicy.disable if data.get("disable_group_by") else GroupByPolicy.force,
            ignore_nonexistent_filters=data["ignore_nonexistent_filters"],
            meta=RawQueryMetaInfo(
                query_type=self.QUERY_TYPE,
                row_count_hard_limit=data["row_count_hard_limit"],
            ),
        )
        return raw_query_spec_union

    def _make_drm(self, raw_query_spec_union: RawQuerySpecUnion, data: Dict[str, Any]) -> DataRequestModel:
        return self.TARGET_CLS(
            raw_query_spec_union=raw_query_spec_union,
            updates=data.get("updates", []),
            autofill_legend=data["autofill_legend"],
            dataset_revision_id=data.get("revision_id"),
        )


class ResultDataRequestV1_5Schema(NewDatasetDataRequestBaseSchema):
    TARGET_CLS = DataRequestModel
    QUERY_TYPE = QueryType.result

    add_fields_data = ma_fields.Boolean(load_default=True)
    with_totals = ma_fields.Boolean(load_default=False)

    def _make_drm(self, raw_query_spec_union: RawQuerySpecUnion, data: Dict[str, Any]) -> DataRequestModel:
        return self.TARGET_CLS(
            raw_query_spec_union=raw_query_spec_union,
            updates=data.get("updates", []),
            autofill_legend=data["autofill_legend"],
            add_fields_data=data["add_fields_data"],
            with_totals=data["with_totals"],
            dataset_revision_id=data.get("revision_id"),
        )


class RequestResultSpecSchema(DefaultSchema[RawResultSpec]):
    TARGET_CLS = RawResultSpec

    with_totals = ma_fields.Boolean(allow_none=False, load_default=False)


class ResultDataRequestV2Schema(NewDatasetDataRequestBaseSchema):
    TARGET_CLS = ResultDataRequestModel
    QUERY_TYPE = QueryType.result

    result = ma_fields.Nested(RequestResultSpecSchema(), allow_none=True)

    def _make_drm(self, raw_query_spec_union: RawQuerySpecUnion, data: Dict[str, Any]) -> ResultDataRequestModel:
        return self.TARGET_CLS(
            raw_query_spec_union=raw_query_spec_union,
            updates=data.get("updates", []),
            result=data.get("result"),
            autofill_legend=data["autofill_legend"],
        )


class DistinctDataRequestV2Schema(NewDatasetDataRequestBaseSchema):
    TARGET_CLS = DataRequestModel
    QUERY_TYPE = QueryType.distinct

    limit = ma_fields.Integer(load_default=1000)  # Define missing value


class RangeDataRequestV2Schema(NewDatasetDataRequestBaseSchema):
    TARGET_CLS = DataRequestModel
    QUERY_TYPE = QueryType.value_range


class PivotDataRequestBaseSchema(NewDatasetDataRequestBaseSchema):
    TARGET_CLS = PivotDataRequestModel
    QUERY_TYPE = QueryType.pivot

    row_count_hard_limit = ma_fields.Integer(
        load_default=DataAPILimits.PIVOT_API_DEFAULT_ROW_COUNT_HARD_LIMIT
    )  # Override default
    pivot = ma_fields.Nested(RequestPivotSpecSchema(), allow_none=False, load_default=lambda: RawPivotSpec())

    def _make_drm(self, raw_query_spec_union: RawQuerySpecUnion, data: Dict[str, Any]) -> DataRequestModel:
        return self.TARGET_CLS(
            raw_query_spec_union=raw_query_spec_union,
            updates=data.get("updates", []),
            pivot=data["pivot"],
            autofill_legend=data["autofill_legend"],
        )


class ResponseBlockInfoSchema(BaseSchema):
    block_id = ma_fields.Integer()
    query = ma_fields.String(allow_none=True)


class V2DataStreamResponseSchema(BaseSchema):
    rows = ma_fields.Raw(allow_none=False)  # Applying schema to data would be very slow


class NotificationSchema(BaseSchema):
    title = ma_fields.String()
    message = ma_fields.String()
    level = ma_fields.Enum(NotificationLevel)
    locator = ma_fields.String()


class DataApiV2ResponseSchema(BaseSchema):
    fields_ = ma_fields.Nested(LegendItemSchema, attribute="fields", data_key="fields", many=True, allow_none=False)
    result_data = ma_fields.Nested(V2DataStreamResponseSchema, many=True, allow_none=False)
    data_export_forbidden = ma_fields.Boolean(allow_none=False)
    blocks = ma_fields.Nested(ResponseBlockInfoSchema, many=True)
    notifications = ma_fields.Nested(NotificationSchema, many=True)

lib/dl_api_lib/dl_api_lib/schemas/data.py
datalens-tech/datalens-backend
from __future__ import annotations

from copy import deepcopy
from typing import (
    Any,
    Dict,
)
import uuid

from marshmallow import (
    post_dump,
    pre_load,
    validates_schema,
)
from marshmallow import ValidationError
from marshmallow import fields as ma_fields
from marshmallow_oneofschema import OneOfSchema

from dl_api_connector.api_schema.component_errors import ComponentErrorListSchema
from dl_api_connector.api_schema.source import DataSourceStrictSchema
from dl_api_connector.api_schema.source_base import (
    AvatarRelationSchema,
    SourceAvatarStrictSchema,
    VirtualFlagField,
)
from dl_api_lib.schemas.fields import ResultSchemaAuxSchema
from dl_api_lib.schemas.filter import ObligatoryFilterSchema
from dl_api_lib.schemas.options import OptionsMixin
from dl_api_lib.schemas.parameters import ParameterValueConstraintSchema
from dl_api_lib.schemas.values import (
    ValueSchema,
    WithNestedValueSchema,
)
from dl_constants.enums import (
    AggregationFunction,
    CalcMode,
    FieldType,
    ManagedBy,
    UserDataType,
)
from dl_core.fields import (
    BIField,
    CalculationSpec,
    DirectCalculationSpec,
    FormulaCalculationSpec,
    ParameterCalculationSpec,
    del_calc_spec_kwargs_from,
    filter_calc_spec_kwargs,
)
from dl_model_tools.schema.base import (
    BaseSchema,
    DefaultSchema,
)


class DirectCalculationSpecSchema(DefaultSchema[DirectCalculationSpec]):
    TARGET_CLS = DirectCalculationSpec

    avatar_id = ma_fields.String(allow_none=True)
    source = ma_fields.String()


class FormulaCalculationSpecSchema(DefaultSchema[FormulaCalculationSpec]):
    TARGET_CLS = FormulaCalculationSpec

    formula = ma_fields.String(load_default="")
    guid_formula = ma_fields.String(load_default="")


class ParameterCalculationSpecSchema(DefaultSchema[ParameterCalculationSpec]):
    TARGET_CLS = ParameterCalculationSpec

    default_value = ma_fields.Nested(ValueSchema, allow_none=True)
    value_constraint = ma_fields.Nested(ParameterValueConstraintSchema, allow_none=True)


class ResultSchemaSchema(WithNestedValueSchema, DefaultSchema[BIField]):
    TYPE_FIELD_NAME = "cast"
    TARGET_CLS = BIField

    class CalculationSpecSchema(OneOfSchema):
        type_field = "mode"
        type_schemas = {
            CalcMode.direct.name: DirectCalculationSpecSchema,
            CalcMode.formula.name: FormulaCalculationSpecSchema,
            CalcMode.parameter.name: ParameterCalculationSpecSchema,
        }

        def get_obj_type(self, obj: CalculationSpec) -> str:
            return obj.mode.name

    title = ma_fields.String(required=True)
    guid = ma_fields.String()
    hidden = ma_fields.Boolean(load_default=False)
    description = ma_fields.String()
    initial_data_type = ma_fields.Enum(UserDataType, allow_none=True)
    cast = ma_fields.Enum(UserDataType)
    type = ma_fields.Enum(FieldType, readonly=True)
    data_type = ma_fields.Enum(UserDataType, allow_none=True)
    valid = ma_fields.Boolean(allow_none=True)

    # this will be flattened on dump and un-flattened before load
    # TODO: dump/load as is and update usage on front end respectively
    calc_spec = ma_fields.Nested(CalculationSpecSchema)

    aggregation = ma_fields.Enum(AggregationFunction, load_default=AggregationFunction.none.name)
    aggregation_locked = ma_fields.Boolean(readonly=True, allow_none=True, load_default=False, dump_only=True)
    autoaggregated = ma_fields.Boolean(readonly=True, allow_none=True, dump_only=True)
    has_auto_aggregation = ma_fields.Boolean(allow_none=True)
    lock_aggregation = ma_fields.Boolean(allow_none=True)

    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user)
    virtual = VirtualFlagField(attribute="managed_by", dump_only=True)

    @post_dump(pass_many=False)
    def add_calc_spec(self, data: Dict[str, Any], **_: Any) -> Dict[str, Any]:
        data = deepcopy(data)
        calc_spec_data = data.pop("calc_spec")
        calc_spec_data["calc_mode"] = calc_spec_data.pop("mode")
        data.update(calc_spec_data)
        # For backward compatibility use '' for formula and source; avatar_id must be present even if None
        for key in ("formula", "guid_formula", "source"):
            data.setdefault(key, "")
        for key in ("avatar_id", "default_value", "value_constraint"):
            data.setdefault(key, None)
        return data

    @pre_load(pass_many=False)
    def extract_calc_spec(self, data: Dict[str, Any], **_: Any) -> Dict[str, Any]:
        data = deepcopy(data)
        mode = data["calc_mode"]
        data["calc_spec"] = dict(filter_calc_spec_kwargs(mode, data), mode=mode)
        return del_calc_spec_kwargs_from(data)

    @validates_schema
    def check_source_is_set_for_direct_and_aggregation_types(
        self, data: Dict[str, Any], *args: Any, **kwargs: Any
    ) -> None:
        if data["calc_spec"].mode == CalcMode.direct.name:
            if not data["calc_spec"].source:
                raise ValidationError("source is required for {}".format(data["title"]))

    def to_object(self, data: dict) -> BIField:
        return BIField.make(**data)


class DatasetContentInternalSchema(BaseSchema):
    """
    A base class for schemas that need to contain the full dataset description
    """

    sources = ma_fields.Nested(DataSourceStrictSchema, many=True, required=False)
    source_avatars = ma_fields.Nested(SourceAvatarStrictSchema, many=True, required=False)
    avatar_relations = ma_fields.Nested(AvatarRelationSchema, many=True, required=False)
    result_schema = ma_fields.List(ma_fields.Nested(ResultSchemaSchema, required=False))
    result_schema_aux = ma_fields.Nested(ResultSchemaAuxSchema, allow_none=False)
    rls = ma_fields.Dict(load_default=dict, required=False)
    preview_enabled = ma_fields.Boolean(load_default=False)
    component_errors = ma_fields.Nested(ComponentErrorListSchema, required=False)
    obligatory_filters = ma_fields.Nested(ObligatoryFilterSchema, many=True, load_default=list)
    revision_id = ma_fields.String(allow_none=True, dump_default=None, load_default=None)

    @pre_load
    def prepare_guids(self, in_data: Dict[str, Any], *args: Any, **kwargs: Any) -> Dict[str, Any]:
        for item in in_data.get("result_schema", ()):
            if not item.get("guid"):
                item["guid"] = str(uuid.uuid4())

        return in_data


class DatasetContentSchema(OptionsMixin):
    dataset = ma_fields.Nested(DatasetContentInternalSchema, required=False)

lib/dl_api_lib/dl_api_lib/schemas/dataset_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import fields as ma_fields

from dl_core.components.dependencies.primitives import (
    FieldInterDependencyInfo,
    FieldInterDependencyItem,
)
from dl_core.us_dataset import Dataset
from dl_model_tools.schema.base import DefaultSchema


class ResultSchemaAuxSchema(DefaultSchema):
    TARGET_CLS = Dataset.DataModel.ResultSchemaAux

    class FieldInterDependencyInfoSchema(DefaultSchema):
        TARGET_CLS = FieldInterDependencyInfo

        class FieldInterDependencyItemSchema(DefaultSchema):
            TARGET_CLS = FieldInterDependencyItem

            dep_field_id = ma_fields.String(allow_none=False)
            ref_field_ids = ma_fields.List(ma_fields.String(allow_none=False), allow_none=False)

        deps = ma_fields.List(ma_fields.Nested(FieldInterDependencyItemSchema()), allow_none=False)

    inter_dependencies = ma_fields.Nested(FieldInterDependencyInfoSchema(), allow_none=False)

lib/dl_api_lib/dl_api_lib/schemas/fields.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
)

from marshmallow import fields as ma_fields
from marshmallow import post_dump

from dl_api_lib.enums import WhereClauseOperation
from dl_api_lib.query.formalization.raw_specs import (
    IdOrTitleFieldRef,
    RawFilterFieldSpec,
)
from dl_api_lib.request_model.data import AddUpdateObligatoryFilter
from dl_constants.enums import ManagedBy
from dl_model_tools.schema.base import (
    DefaultSchema,
    DefaultValidateSchema,
)


class WhereSchema(DefaultSchema[RawFilterFieldSpec]):
    TARGET_CLS = RawFilterFieldSpec

    column = ma_fields.String(required=True)
    operation = ma_fields.Enum(WhereClauseOperation, required=True)
    values = ma_fields.List(ma_fields.Raw(allow_none=True), required=True, allow_none=True)

    def to_object(self, data: dict) -> RawFilterFieldSpec:
        return RawFilterFieldSpec(
            ref=IdOrTitleFieldRef(id_or_title=data["column"]),
            operation=data["operation"],
            values=data["values"],
        )


class ObligatoryFilterSchema(DefaultValidateSchema[AddUpdateObligatoryFilter]):
    TARGET_CLS = AddUpdateObligatoryFilter

    id = ma_fields.String(required=True)
    field_guid = ma_fields.String()
    default_filters = ma_fields.Nested(WhereSchema, many=True, load_default=[])
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user, load_default=ManagedBy.user)
    valid = ma_fields.Boolean(load_default=True)

    @post_dump
    def set_where_column(self, data: Dict[str, Any], **_: Any) -> Dict[str, Any]:
        if data.get("field_guid"):
            for filter in data["default_filters"]:
                if not filter.get("column"):
                    filter["column"] = data["field_guid"]
        return data

lib/dl_api_lib/dl_api_lib/schemas/filter.py
datalens-tech/datalens-backend
from typing import Any

from marshmallow import EXCLUDE
from marshmallow import fields as ma_fields
from marshmallow_oneofschema import OneOfSchema

from dl_constants.enums import (
    FieldRole,
    FieldType,
    FieldVisibility,
    LegendItemType,
    OrderDirection,
    RangeType,
    UserDataType,
)
from dl_model_tools.schema.base import BaseSchema
from dl_query_processing.legend.field_legend import RoleSpec


class DimensionValueSpecSchema(BaseSchema):
    legend_item_id = ma_fields.Integer()
    value = ma_fields.Raw()


class RoleSpecSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "role"

    class RoleSpecSchemaVariant(BaseSchema):
        role = ma_fields.Enum(FieldRole)

    class OrderByRoleSpecSchemaVariant(RoleSpecSchemaVariant):
        direction = ma_fields.Enum(OrderDirection)

    class AnnotationRoleSpecSchemaVariant(RoleSpecSchemaVariant):
        annotation_type = ma_fields.String()
        target_legend_item_ids = ma_fields.List(ma_fields.Integer(), allow_none=True)

    class RangeRoleSpecSchemaVariant(RoleSpecSchemaVariant):
        range_type = ma_fields.Enum(RangeType)

    class DimensionRoleSpecSchemaVariant(RoleSpecSchemaVariant):
        visibility = ma_fields.Enum(FieldVisibility)

    class RowRoleSpecSchemaVariant(DimensionRoleSpecSchemaVariant):
        pass

    class TemplateRoleSpecSchemaVariant(DimensionRoleSpecSchemaVariant):
        template = ma_fields.String()

    class TreeRoleSpecSchemaVariant(DimensionRoleSpecSchemaVariant):
        level = ma_fields.Integer()
        prefix = ma_fields.Raw()
        dimension_values = ma_fields.Nested(DimensionValueSpecSchema, many=True, allow_none=True)

    type_schemas = {
        FieldRole.row.name: RowRoleSpecSchemaVariant,
        FieldRole.measure.name: RoleSpecSchemaVariant,
        FieldRole.info.name: RoleSpecSchemaVariant,
        FieldRole.order_by.name: OrderByRoleSpecSchemaVariant,
        FieldRole.filter.name: RoleSpecSchemaVariant,  # TODO: extend if needed
        FieldRole.parameter.name: RoleSpecSchemaVariant,  # TODO: extend if needed
        FieldRole.range.name: RangeRoleSpecSchemaVariant,
        FieldRole.distinct.name: RoleSpecSchemaVariant,
        FieldRole.template.name: TemplateRoleSpecSchemaVariant,
        FieldRole.total.name: RoleSpecSchemaVariant,
        FieldRole.tree.name: TreeRoleSpecSchemaVariant,
    }

    def get_obj_type(self, obj: Any) -> str:
        assert isinstance(obj, RoleSpec)
        return obj.role.name


class LegendItemSchema(BaseSchema):
    legend_item_id = ma_fields.Integer()
    id = ma_fields.String()
    title = ma_fields.String()
    role_spec = ma_fields.Nested(RoleSpecSchema)
    data_type = ma_fields.Enum(UserDataType)
    field_type = ma_fields.Enum(FieldType)
    item_type = ma_fields.Enum(LegendItemType)


class LegendSchema(BaseSchema):
    fields_ = ma_fields.Nested(LegendItemSchema, data_key="fields", many=True, attribute="items")

lib/dl_api_lib/dl_api_lib/schemas/legend.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from typing import (
    Any,
    Dict,
    Optional,
)

from flask_restx import (
    Namespace,
    fields,
)
from flask_restx.model import RawModel
from marshmallow import Schema
from marshmallow import fields as ma_fields

from dl_api_lib.schemas.dataset_base import DatasetContentSchema
from dl_constants.enums import DataSourceCreatedVia
from dl_model_tools.schema.base import BaseSchema
from dl_model_tools.schema.dynamic_enum_field import DynamicEnumField


def get_api_model(
    ma_schema: Optional[Schema],
    ns: Namespace,
    name: Optional[str] = None,
    schema_fields: Optional[Dict[str, ma_fields.Field]] = None,
) -> RawModel:
    """Generate a ``flask_restx`` schema that will be used for generating Swagger documentation."""

    name = name or ma_schema.__class__.__name__
    schema_fields = schema_fields or (ma_schema._declared_fields if ma_schema is not None else None)

    fields_map = {
        ma_fields.Integer: fields.Integer,
        ma_fields.String: fields.String,
        ma_fields.Boolean: fields.Boolean,
        ma_fields.Raw: fields.Raw,
    }

    def _translate_field(_field: ma_fields.Field) -> fields.Raw:
        """
        Convert marshmallow field instance into a ``flask_restx`` field.
        Do this recursively for complex fields like ``Nested`` and ``List``.
        """

        if isinstance(_field, ma_fields.Nested):
            nested = _field.nested
            nested = nested() if isinstance(nested, type) else nested
            assert isinstance(nested, Schema)
            return fields.Nested(
                get_api_model(nested, ns),
                as_list=_field.many,
                required=_field.required,
            )
        elif isinstance(_field, ma_fields.Enum):
            return fields.String(required=_field.required, enum=[x.name for x in list(_field.enum)])
        elif isinstance(_field, ma_fields.List):
            return fields.List(
                _translate_field(_field.inner),
                required=_field.required,
            )
        elif isinstance(_field, ma_fields.DateTime):
            return fields.DateTime(
                example=datetime.datetime(2018, 1, 1).strftime(_field.format or _field.DEFAULT_FORMAT),
            )
        elif isinstance(_field, ma_fields.Date):
            return fields.Date(
                example=datetime.date(2018, 1, 1).strftime(_field.format or _field.DEFAULT_FORMAT),
            )
        else:
            api_field_type = fields_map.get(type(_field), fields.String)
            return api_field_type(required=_field.required)

    api_model_dict = {}
    if schema_fields is not None:
        for field_name, field in schema_fields.items():
            api_model_dict[field_name] = _translate_field(field)

    return ns.model(name, api_model_dict)


class CreateDatasetSchema(DatasetContentSchema):
    name = ma_fields.String()
    dir_path = ma_fields.String()
    workbook_id = ma_fields.String()
    # TODO FIX: Ensure that not used and remove
    preview = ma_fields.Boolean(load_default=False, required=False)
    created_via = DynamicEnumField(DataSourceCreatedVia, load_default=DataSourceCreatedVia.user)


class CreateDatasetResponseSchema(DatasetContentSchema):
    id = ma_fields.String()


class NormalizedDateTime(ma_fields.DateTime):
    """A DateTime field that also accepts strings (formatted datetimes)"""

    def _serialize(
        self,
        value: str | datetime.datetime,
        attr: Optional[str],
        obj: Any,
        **kwargs: Any,
    ) -> Optional[str | float]:
        if isinstance(value, str):
            value = datetime.datetime.strptime(value.split(".")[0], "%Y-%m-%d %H:%M:%S")
        return super()._serialize(value, attr, obj, **kwargs)


class GetDatasetResponseSchema(DatasetContentSchema):
    id = ma_fields.String()
    name = ma_fields.String()
    pub_operation_id = ma_fields.String()
    row_count = ma_fields.Integer()
    ctime = NormalizedDateTime(format="%Y-%m-%d %H:%M:%S")
    mtime = NormalizedDateTime(format="%Y-%m-%d %H:%M:%S")
    is_favorite = ma_fields.Boolean()
    permissions = ma_fields.Dict(keys=ma_fields.String(), values=ma_fields.Boolean())


class UpdateDatasetSchema(BaseSchema):
    name = ma_fields.String()


class DatasetUpdateSchema(DatasetContentSchema):
    pass


class DatasetCopyRequestSchema(BaseSchema):
    new_key = ma_fields.String(required=True)


class DatasetCopyResponseSchema(DatasetContentSchema):
    id = ma_fields.String(required=True)


class GetDatasetVersionResponseSchema(GetDatasetResponseSchema):
    key = ma_fields.String()
    workbook_id = ma_fields.String()


class BadRequestResponseSchema(BaseSchema):
    message = ma_fields.String()


class DashSQLBindParam(BaseSchema):
    type_name = ma_fields.String(required=True)
    # value = ma_fields.String(required=True, allow_none=True)
    # ^ Can also be a list of strings, so:
    value = ma_fields.Raw(required=True, allow_none=True)


class DashSQLRequestSchema(BaseSchema):
    sql_query = ma_fields.String(required=True)
    params = ma_fields.Dict(ma_fields.String(), ma_fields.Nested(DashSQLBindParam()), required=False)
    db_params = ma_fields.Dict(ma_fields.String(), ma_fields.String(), required=False)
    connector_specific_params = ma_fields.Dict(ma_fields.String(), ma_fields.String(), required=False)

lib/dl_api_lib/dl_api_lib/schemas/main.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging

from marshmallow import fields as ma_fields

from dl_api_lib.enums import WhereClauseOperation
from dl_constants.enums import (
    AggregationFunction,
    BinaryJoinOperator,
    ConnectionType,
    DataSourceType,
    JoinType,
    UserDataType,
)
from dl_model_tools.schema.base import BaseSchema
from dl_model_tools.schema.dynamic_enum_field import DynamicEnumField


LOGGER = logging.getLogger(__name__)


class JoinSchema(BaseSchema):
    types = ma_fields.List(ma_fields.Enum(JoinType))
    operators = ma_fields.List(ma_fields.Enum(BinaryJoinOperator))


class PreviewSchema(BaseSchema):
    enabled = ma_fields.Boolean()


class DataTypesSchema(BaseSchema):
    class DataTypeListItemSchema(BaseSchema):
        type = ma_fields.Enum(UserDataType)
        aggregations = ma_fields.List(ma_fields.Enum(AggregationFunction))
        casts = ma_fields.List(ma_fields.Enum(UserDataType))
        filter_operations = ma_fields.List(ma_fields.Enum(WhereClauseOperation))

    items = ma_fields.List(ma_fields.Nested(DataTypeListItemSchema))


class FieldsSchema(BaseSchema):
    class FieldListItemSchema(BaseSchema):
        guid = ma_fields.String()
        casts = ma_fields.List(ma_fields.Enum(UserDataType))
        aggregations = ma_fields.List(ma_fields.Enum(AggregationFunction))

    items = ma_fields.List(ma_fields.Nested(FieldListItemSchema))


class CompatConnectionTypeListItemSchema(BaseSchema):
    conn_type = DynamicEnumField(ConnectionType)


class ConnectionsSchema(BaseSchema):
    class ConnectionListItemSchema(BaseSchema):
        id = ma_fields.String()
        replacement_types = ma_fields.List(ma_fields.Nested(CompatConnectionTypeListItemSchema))

    max = ma_fields.Integer()
    compatible_types = ma_fields.List(ma_fields.Nested(CompatConnectionTypeListItemSchema))
    items = ma_fields.List(ma_fields.Nested(ConnectionListItemSchema))


class CompatSourceTypeListItemSchema(BaseSchema):
    source_type = DynamicEnumField(DataSourceType)


class SourcesSchema(BaseSchema):
    class SourceListItemSchema(BaseSchema):
        id = ma_fields.String()
        schema_update_enabled = ma_fields.Boolean()

    max = ma_fields.Integer()
    compatible_types = ma_fields.List(ma_fields.Nested(CompatSourceTypeListItemSchema))
    items = ma_fields.List(ma_fields.Nested(SourceListItemSchema))


class AvatarsSchema(BaseSchema):
    class SourceListItemSchema(BaseSchema):
        id = ma_fields.String()
        schema_update_enabled = ma_fields.Boolean()

    max = ma_fields.Integer()
    items = ma_fields.List(ma_fields.Nested(SourceListItemSchema))


class OptionsSchema(BaseSchema):
    join = ma_fields.Nested(JoinSchema)
    schema_update_enabled = ma_fields.Boolean()
    data_types = ma_fields.Nested(DataTypesSchema)
    preview = ma_fields.Nested(PreviewSchema)
    fields_ = ma_fields.Nested(FieldsSchema, attribute="fields", data_key="fields")
    connections = ma_fields.Nested(ConnectionsSchema)
    sources = ma_fields.Nested(SourcesSchema)
    source_avatars = ma_fields.Nested(AvatarsSchema)
    supports_offset = ma_fields.Boolean()
    supported_functions = ma_fields.List(ma_fields.String())


class OptionsMixin(BaseSchema):
    options = ma_fields.Nested(OptionsSchema, dump_only=True)

lib/dl_api_lib/dl_api_lib/schemas/options.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Any

from marshmallow import EXCLUDE
from marshmallow import fields as ma_fields
from marshmallow_oneofschema import OneOfSchema

from dl_api_lib.schemas.values import ValueSchema
from dl_constants.enums import ParameterValueConstraintType
from dl_core.fields import (
    ParameterValueConstraint,
    RangeParameterValueConstraint,
    SetParameterValueConstraint,
)
from dl_model_tools.schema.base import DefaultSchema


class ParameterValueConstraintSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "type"

    class BaseParameterValueConstraintSchema(DefaultSchema):
        type = ma_fields.Enum(ParameterValueConstraintType)

    class AllParameterValueConstraintSchema(BaseParameterValueConstraintSchema):
        TARGET_CLS = ParameterValueConstraint

    class RangeParameterValueConstraintSchema(BaseParameterValueConstraintSchema):
        TARGET_CLS = RangeParameterValueConstraint

        min = ma_fields.Nested(ValueSchema, allow_none=True)
        max = ma_fields.Nested(ValueSchema, allow_none=True)

    class SetParameterValueConstraintSchema(BaseParameterValueConstraintSchema):
        TARGET_CLS = SetParameterValueConstraint

        values = ma_fields.List(ma_fields.Nested(ValueSchema))

    type_schemas = {
        ParameterValueConstraintType.all.name: AllParameterValueConstraintSchema,
        ParameterValueConstraintType.range.name: RangeParameterValueConstraintSchema,
        ParameterValueConstraintType.set.name: SetParameterValueConstraintSchema,
    }

    def get_obj_type(self, obj: Any) -> str:
        assert isinstance(obj, ParameterValueConstraint)
        return obj.type.name

lib/dl_api_lib/dl_api_lib/schemas/parameters.py
datalens-tech/datalens-backend
from typing import Any

from marshmallow import EXCLUDE
from marshmallow import fields as ma_fields
from marshmallow import validate as ma_validate
from marshmallow_oneofschema import OneOfSchema

from dl_api_lib.pivot.primitives import (
    PivotHeaderInfo,
    PivotHeaderRoleSpec,
    PivotHeaderValue,
    PivotMeasureSorting,
    PivotMeasureSortingSettings,
)
from dl_api_lib.query.formalization.pivot_legend import PivotRoleSpec
from dl_api_lib.query.formalization.raw_pivot_specs import (
    PivotPaginationSpec,
    RawAnnotationPivotRoleSpec,
    RawDimensionPivotRoleSpec,
    RawPivotLegendItem,
    RawPivotMeasureRoleSpec,
    RawPivotRoleSpec,
    RawPivotSpec,
    RawPivotTotalsItemSpec,
    RawPivotTotalsSpec,
)
from dl_constants.enums import (
    OrderDirection,
    PivotHeaderRole,
    PivotRole,
)
from dl_model_tools.schema.base import (
    BaseSchema,
    DefaultSchema,
)


class PivotHeaderRoleSpecSchema(DefaultSchema[PivotHeaderRoleSpec]):
    TARGET_CLS = PivotHeaderRoleSpec

    role = ma_fields.Enum(PivotHeaderRole, required=True)


class PivotHeaderValueSchema(DefaultSchema[PivotHeaderValue]):
    TARGET_CLS = PivotHeaderValue

    value = ma_fields.String(allow_none=False, required=True)


class PivotMeasureSortingSettingsSchema(DefaultSchema[PivotMeasureSortingSettings]):
    TARGET_CLS = PivotMeasureSortingSettings

    header_values = ma_fields.List(ma_fields.Nested(PivotHeaderValueSchema), allow_none=False, required=True)
    direction = ma_fields.Enum(OrderDirection, load_default=OrderDirection.asc)
    role_spec = ma_fields.Nested(PivotHeaderRoleSpecSchema, allow_none=False, required=True)


class PivotMeasureSortingSchema(DefaultSchema[PivotMeasureSorting]):
    TARGET_CLS = PivotMeasureSorting

    column = ma_fields.Nested(PivotMeasureSortingSettingsSchema, allow_none=True, load_default=None)
    row = ma_fields.Nested(PivotMeasureSortingSettingsSchema, allow_none=True, load_default=None)


class PivotRoleSpecSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "role"

    class PivotRoleSpecSchemaVariant(DefaultSchema[RawPivotRoleSpec]):
        TARGET_CLS = RawPivotRoleSpec

        role = ma_fields.Enum(PivotRole, required=True)

    class DimensionPivotRoleSpecSchemaVariant(DefaultSchema[RawDimensionPivotRoleSpec]):
        TARGET_CLS = RawDimensionPivotRoleSpec

        role = ma_fields.Enum(PivotRole, required=True)
        direction = ma_fields.Enum(OrderDirection, load_default=OrderDirection.asc)

    class AnnotationPivotRoleSpecSchemaVariant(DefaultSchema[RawAnnotationPivotRoleSpec]):
        TARGET_CLS = RawAnnotationPivotRoleSpec

        role = ma_fields.Enum(PivotRole, required=True)
        annotation_type = ma_fields.String(required=True)
        target_legend_item_ids = ma_fields.List(ma_fields.Integer(), allow_none=True)

    class MeasurePivotRoleSpecSchemaVariant(DefaultSchema[RawPivotMeasureRoleSpec]):
        TARGET_CLS = RawPivotMeasureRoleSpec

        role = ma_fields.Enum(PivotRole, required=True)
        sorting = ma_fields.Nested(PivotMeasureSortingSchema, allow_none=True, load_default=None)

    type_schemas = {
        PivotRole.pivot_row.name: DimensionPivotRoleSpecSchemaVariant,
        PivotRole.pivot_column.name: DimensionPivotRoleSpecSchemaVariant,
        PivotRole.pivot_measure.name: MeasurePivotRoleSpecSchemaVariant,
        PivotRole.pivot_annotation.name: AnnotationPivotRoleSpecSchemaVariant,
        PivotRole.pivot_info.name: PivotRoleSpecSchemaVariant,
    }

    def get_obj_type(self, obj: Any) -> str:
        assert isinstance(obj, PivotRoleSpec)
        return obj.role.name


class RequestPivotItemSchema(DefaultSchema[RawPivotLegendItem]):
    TARGET_CLS = RawPivotLegendItem

    pivot_item_id = ma_fields.Integer(allow_none=True, required=False)
    legend_item_ids = ma_fields.List(ma_fields.Integer(allow_none=False), allow_none=False, required=True)
    role_spec = ma_fields.Nested(PivotRoleSpecSchema, allow_none=False, required=True)
    title = ma_fields.String(allow_none=True, required=False)


class PivotItemSchema(BaseSchema):
    pivot_item_id = ma_fields.Integer(allow_none=True)
    legend_item_ids = ma_fields.List(ma_fields.Integer(allow_none=False), allow_none=False, required=True)
    role_spec = ma_fields.Nested(PivotRoleSpecSchema, allow_none=False, required=True)
    title = ma_fields.String()


class PivotPaginationSchema(DefaultSchema[PivotPaginationSpec]):
    TARGET_CLS = PivotPaginationSpec

    limit_rows = ma_fields.Integer(load_default=None, validate=ma_validate.Range(min=1))
    offset_rows = ma_fields.Integer(load_default=None, validate=ma_validate.Range(min=0))


class PivotTotalsItemSchema(DefaultSchema[RawPivotTotalsItemSpec]):
    TARGET_CLS = RawPivotTotalsItemSpec

    level = ma_fields.Integer(allow_none=True, required=True)


class PivotTotalsSchema(DefaultSchema[RawPivotTotalsSpec]):
    TARGET_CLS = RawPivotTotalsSpec

    rows = ma_fields.Nested(PivotTotalsItemSchema, many=True, required=True)
    columns = ma_fields.Nested(PivotTotalsItemSchema, many=True, required=True)


class RequestPivotSpecSchema(DefaultSchema[RawPivotSpec]):
    TARGET_CLS = RawPivotSpec

    structure = ma_fields.Nested(RequestPivotItemSchema, many=True, allow_none=False)
    pagination = ma_fields.Nested(PivotPaginationSchema, load_default=PivotPaginationSpec())
    with_totals = ma_fields.Boolean(allow_none=True, load_default=None)
    totals = ma_fields.Nested(PivotTotalsSchema, allow_none=True, load_default=None)


class PivotHeaderInfoSchema(DefaultSchema[PivotHeaderInfo]):
    TARGET_CLS = PivotHeaderInfo

    sorting_direction = ma_fields.Enum(OrderDirection, allow_none=True)
    role_spec = ma_fields.Nested(PivotHeaderRoleSpecSchema, allow_none=False, required=True)

lib/dl_api_lib/dl_api_lib/schemas/pivot.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
    Optional,
)

from dl_api_connector.api_schema.extras import OperationsMode
from dl_api_connector.api_schema.top_level import (
    BaseTopLevelSchema,
    USEntryBaseSchema,
)
from dl_core.us_manager.us_manager import USManagerBase


def prepare_schema_context(
    usm: Optional[USManagerBase] = None,
    op_mode: Optional[OperationsMode] = None,
    editable_object: Optional[Any] = None,
) -> Dict[Any, Any]:
    return {
        BaseTopLevelSchema.CTX_KEY_EDITABLE_OBJECT: editable_object,
        BaseTopLevelSchema.CTX_KEY_OPERATIONS_MODE: op_mode,
        USEntryBaseSchema.CTX_KEY_USM: usm,
    }

lib/dl_api_lib/dl_api_lib/schemas/tools.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import fields as ma_fields

from dl_api_lib.schemas.action import ActionSchema
from dl_api_lib.schemas.dataset_base import DatasetContentSchema
from dl_constants.enums import ComponentErrorLevel
from dl_core.marshmallow import ErrorCodeField
from dl_model_tools.schema.base import BaseSchema


class DatasetValidationSchema(DatasetContentSchema):
    """
    Full dataset configuration with a batch of updates that should be applied consecutively to the configuration.
    Validation errors are returned only for the final version (after all of the updates are applied)
    """

    updates = ma_fields.Nested(ActionSchema, many=True, required=False)


class FieldErrorListSchema(BaseSchema):
    class FieldErrorSchema(BaseSchema):
        message = ma_fields.String()
        code = ErrorCodeField()
        column = ma_fields.Integer()
        row = ma_fields.Integer()
        level = ma_fields.Enum(ComponentErrorLevel)
        token = ma_fields.String()

    title = ma_fields.String()
    guid = ma_fields.String()
    errors = ma_fields.Nested(FieldErrorSchema, many=True)


class DatasetValidationResponseSchema(DatasetContentSchema):
    message = ma_fields.String()
    code = ma_fields.String()
    dataset_errors = ma_fields.List(ma_fields.String())  # TODO: Remove


class FieldValidationSchema(DatasetContentSchema):
    class FieldFormulaSchema(BaseSchema):
        title = ma_fields.String(required=True)
        guid = ma_fields.String()
        formula = ma_fields.String(load_default="")

    field = ma_fields.Nested(FieldFormulaSchema)


class FieldValidationResponseSchema(BaseSchema):
    message = ma_fields.String()
    code = ma_fields.String()
    # errors
    field_errors = ma_fields.Nested(FieldErrorListSchema, many=True)

lib/dl_api_lib/dl_api_lib/schemas/validation.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
    Dict,
    Mapping,
    cast,
)

from marshmallow import (
    EXCLUDE,
    Schema,
)
from marshmallow import (
    post_dump,
    pre_load,
)
from marshmallow import fields as ma_fields

from dl_constants.enums import UserDataType
from dl_core.values import (
    ArrayFloatValue,
    ArrayIntValue,
    ArrayStrValue,
    BIValue,
    BooleanValue,
    DateTimeTZValue,
    DateTimeValue,
    DateValue,
    FloatValue,
    GenericDateTimeValue,
    GeoPointValue,
    GeoPolygonValue,
    IntegerValue,
    MarkupValue,
    StringValue,
    UuidValue,
)
from dl_model_tools.schema.base import DefaultSchema
from dl_utils.schemas import OneOfSchemaWithDumpLoadHooks


VALUE_TYPE_CONTEXT_KEY = "bi_value_type"


class WithNestedValueSchema:
    TYPE_FIELD_NAME: ClassVar[str]

    @pre_load(pass_many=False)
    def store_type_in_context(self, data: Mapping[str, Any], **_: Any) -> Mapping[str, Any]:
        if self.TYPE_FIELD_NAME in data:
            cast(Schema, self).context[VALUE_TYPE_CONTEXT_KEY] = data[self.TYPE_FIELD_NAME]
        return data


class ValueSchema(OneOfSchemaWithDumpLoadHooks):
    """
    This schema is designed to dump and load BIValue objects as value itself without any info about its type
    On load it relies on type info stored in context by outer schema object (e.g. see WithNestedValueSchema mixin)
    """

    class Meta:
        unknown = EXCLUDE

    @pre_load(pass_many=False)
    def wrap_value_with_type(self, data: Any, **_: Any) -> Dict[str, Any]:
        type = getattr(self, "context", {}).get(VALUE_TYPE_CONTEXT_KEY)
        return {"type": type, "value": data}

    @post_dump(pass_many=False)
    def extract_value(self, data: Dict[str, Any], **_: Any) -> Any:
        return data["value"]

    class BaseValueSchema(DefaultSchema):
        type = ma_fields.Enum(UserDataType)
        value = ma_fields.Field()

    class StringValueSchema(BaseValueSchema):
        TARGET_CLS = StringValue
        value = ma_fields.String()

    class IntegerValueSchema(BaseValueSchema):
        TARGET_CLS = IntegerValue
        value = ma_fields.Integer()

    class FloatValueSchema(BaseValueSchema):
        TARGET_CLS = FloatValue
        value = ma_fields.Float()

    class DateValueSchema(BaseValueSchema):
        TARGET_CLS = DateValue
        value = ma_fields.Date()

    class DateTimeValueSchema(BaseValueSchema):
        TARGET_CLS = DateTimeValue
        value = ma_fields.DateTime()

    class DateTimeTZValueSchema(BaseValueSchema):
        TARGET_CLS = DateTimeTZValue
        value = ma_fields.DateTime()

    class GenericDateTimeValueSchema(BaseValueSchema):
        TARGET_CLS = GenericDateTimeValue
        value = ma_fields.DateTime()

    class BooleanValueSchema(BaseValueSchema):
        TARGET_CLS = BooleanValue
        value = ma_fields.Boolean()

    class GeoPointValueSchema(BaseValueSchema):
        TARGET_CLS = GeoPointValue
        value = ma_fields.List(ma_fields.Float())

    class GeoPolygonValueSchema(BaseValueSchema):
        TARGET_CLS = GeoPolygonValue
        value = ma_fields.List(ma_fields.List(ma_fields.List(ma_fields.Float())))

    class UuidValueSchema(BaseValueSchema):
        TARGET_CLS = UuidValue
        value = ma_fields.String()

    class MarkupValueSchema(BaseValueSchema):
        TARGET_CLS = MarkupValue
        value = ma_fields.String()

    class ArrayStrValueSchema(BaseValueSchema):
        TARGET_CLS = ArrayStrValue
        value = ma_fields.List(ma_fields.String())

    class ArrayIntValueSchema(BaseValueSchema):
        TARGET_CLS = ArrayIntValue
        value = ma_fields.List(ma_fields.Integer())

    class ArrayFloatValueSchema(BaseValueSchema):
        TARGET_CLS = ArrayFloatValue
        value = ma_fields.List(ma_fields.Float())

    class TreeStrValueSchema(BaseValueSchema):
        TARGET_CLS = ArrayStrValue
        value = ma_fields.List(ma_fields.String())

    type_schemas = {
        UserDataType.string.name: StringValueSchema,
        UserDataType.integer.name: IntegerValueSchema,
        UserDataType.float.name: FloatValueSchema,
        UserDataType.date.name: DateValueSchema,
        UserDataType.datetime.name: DateTimeValueSchema,
        UserDataType.datetimetz.name: DateTimeTZValueSchema,
        UserDataType.genericdatetime.name: GenericDateTimeValueSchema,
        UserDataType.boolean.name: BooleanValueSchema,
        UserDataType.geopoint.name: GeoPointValueSchema,
        UserDataType.geopolygon.name: GeoPolygonValueSchema,
        UserDataType.uuid.name: UuidValueSchema,
        UserDataType.markup.name: MarkupValueSchema,
        UserDataType.array_str.name: ArrayStrValueSchema,
        UserDataType.array_int.name: ArrayIntValueSchema,
        UserDataType.array_float.name: ArrayFloatValueSchema,
        UserDataType.tree_str.name: TreeStrValueSchema,
    }

    def get_obj_type(self, obj: BIValue) -> str:
        return obj.type.name

lib/dl_api_lib/dl_api_lib/schemas/values.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc

import attr

from dl_api_lib.dataset.validator import DatasetValidator
from dl_core.us_dataset import Dataset
from dl_core.us_manager.us_manager import USManagerBase


@attr.s(frozen=True)
class DatasetValidatorFactory(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_dataset_validator(
        self, ds: Dataset, us_manager: USManagerBase, is_data_api: bool = False
    ) -> DatasetValidator:
        pass


@attr.s(frozen=True)
class DefaultDatasetValidatorFactory(DatasetValidatorFactory):
    _is_bleeding_edge_user: bool = attr.ib(default=False)

    def get_dataset_validator(
        self, ds: Dataset, us_manager: USManagerBase, is_data_api: bool = False
    ) -> DatasetValidator:
        return DatasetValidator(
            ds=ds,
            us_manager=us_manager,
            is_data_api=is_data_api,
            is_bleeding_edge_user=self._is_bleeding_edge_user,
        )

lib/dl_api_lib/dl_api_lib/service_registry/dataset_validator_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_core.components.ids import (
    DEFAULT_FIELD_ID_GENERATOR_TYPE,
    FIELD_ID_GENERATOR_MAP,
    FieldIdGenerator,
    FieldIdGeneratorType,
)
from dl_core.us_dataset import Dataset


@attr.s
class FieldIdGeneratorFactory:
    field_id_generator_type: FieldIdGeneratorType = attr.ib(kw_only=True, default=DEFAULT_FIELD_ID_GENERATOR_TYPE)

    def __attrs_post_init__(self) -> None:
        if self.field_id_generator_type not in FIELD_ID_GENERATOR_MAP:
            self.field_id_generator_type = DEFAULT_FIELD_ID_GENERATOR_TYPE

    def get_field_id_generator(
        self,
        ds: Dataset,
    ) -> FieldIdGenerator:
        FieldIdGeneratorClass = FIELD_ID_GENERATOR_MAP[self.field_id_generator_type]
        return FieldIdGeneratorClass(dataset=ds)

lib/dl_api_lib/dl_api_lib/service_registry/field_id_generator_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Dict,
    Optional,
)

import attr

from dl_formula.parser.base import FormulaParser
from dl_formula.parser.factory import (
    ParserType,
    get_parser,
)


_DEFAULT_PARSER_TYPE = ParserType.antlr_py

LOGGER = logging.getLogger(__name__)


@attr.s
class FormulaParserFactory:
    _default_formula_parser_type: Optional[ParserType] = attr.ib(kw_only=True, default=None)
    _saved_parsers: Dict[ParserType, FormulaParser] = attr.ib(init=False, factory=dict)

    def get_formula_parser(self, parser_type: Optional[ParserType] = None) -> FormulaParser:
        if parser_type is None:
            parser_type = self._default_formula_parser_type or _DEFAULT_PARSER_TYPE
        assert parser_type is not None

        if parser_type not in self._saved_parsers:
            self._saved_parsers[parser_type] = get_parser(parser_type=parser_type)

        return self._saved_parsers[parser_type]

    def log_parser_stats_for_type(self, parser_type: ParserType) -> None:
        parser = self._saved_parsers[parser_type]
        stats = parser.get_usage_stats()
        data = {
            "global_cache_hits": stats.global_cache_hits,
            "global_cache_misses": stats.global_cache_misses,
            "global_cache_maxsize": stats.global_cache_maxsize,
            "global_cache_currsize": stats.global_cache_currsize,
            "call_count": stats.call_count,
            "call_time": stats.call_time,
            "length_counts": stats.length_counts,
            "length_avg_times": stats.length_avg_times,
            "length_median_times": stats.length_median_times,
            "parser_type": parser_type.name,
        }
        LOGGER.info(
            f"Formula parser global statistics ({parser_type.name})", extra=dict(function_parser_statistics=data)
        )

    def log_parser_stats_for_all_used_parsers(self) -> None:
        for parser_type in self._saved_parsers:
            self.log_parser_stats_for_type(parser_type=parser_type)

    def close(self) -> None:
        self.log_parser_stats_for_all_used_parsers()

lib/dl_api_lib/dl_api_lib/service_registry/formula_parser_factory.py
datalens-tech/datalens-backend
import abc
from typing import Sequence

import attr

from dl_api_lib.query.registry import get_multi_query_mutator_factory
from dl_constants.enums import (
    QueryProcessingMode,
    SourceBackendType,
)
from dl_core.us_dataset import Dataset
from dl_formula.core.dialect import DialectCombo
from dl_query_processing.multi_query.factory import MultiQueryMutatorFactoryBase
from dl_query_processing.multi_query.mutators.base import MultiQueryMutatorBase


class SRMultiQueryMutatorFactory(abc.ABC):
    @abc.abstractmethod
    def get_mqm_factory(
        self,
        backend_type: SourceBackendType,
        dialect: DialectCombo,
        dataset: Dataset,
    ) -> MultiQueryMutatorFactoryBase:
        raise NotImplementedError

    def get_multi_query_mutators(
        self,
        backend_type: SourceBackendType,
        dialect: DialectCombo,
        dataset: Dataset,
    ) -> Sequence[MultiQueryMutatorBase]:
        mqm_factory = self.get_mqm_factory(backend_type=backend_type, dialect=dialect, dataset=dataset)
        return mqm_factory.get_mutators()


@attr.s
class DefaultSRMultiQueryMutatorFactory(SRMultiQueryMutatorFactory):
    _query_proc_mode: QueryProcessingMode = attr.ib(kw_only=True)

    def get_mqm_factory(
        self,
        backend_type: SourceBackendType,
        dialect: DialectCombo,
        dataset: Dataset,
    ) -> MultiQueryMutatorFactoryBase:
        # Try to get for the specified query mode
        factory = get_multi_query_mutator_factory(
            query_proc_mode=self._query_proc_mode,
            backend_type=backend_type,
            dialect=dialect,
            result_schema=dataset.result_schema,
        )
        if factory is None:
            # Try again for the basic mode
            factory = get_multi_query_mutator_factory(
                query_proc_mode=QueryProcessingMode.basic,
                backend_type=backend_type,
                dialect=dialect,
                result_schema=dataset.result_schema,
            )

        assert factory is not None
        return factory

lib/dl_api_lib/dl_api_lib/service_registry/multi_query_mutator_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Optional,
)

import attr

from dl_api_lib.connector_availability.base import ConnectorAvailabilityConfig
from dl_api_lib.service_registry.field_id_generator_factory import FieldIdGeneratorFactory
from dl_api_lib.service_registry.formula_parser_factory import FormulaParserFactory
from dl_api_lib.service_registry.multi_query_mutator_factory import (
    DefaultSRMultiQueryMutatorFactory,
    SRMultiQueryMutatorFactory,
)
from dl_api_lib.service_registry.supported_functions_manager import SupportedFunctionsManager
from dl_api_lib.utils.rls import BaseSubjectResolver
from dl_constants.enums import QueryProcessingMode
from dl_core.services_registry.top_level import (
    DefaultServicesRegistry,
    ServicesRegistry,
)
from dl_formula.parser.factory import ParserType
from dl_i18n.localizer_base import (
    Localizer,
    LocalizerFactory,
)


if TYPE_CHECKING:
    from dl_api_lib.service_registry.dataset_validator_factory import DatasetValidatorFactory


@attr.s
class ApiServiceRegistry(ServicesRegistry, metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_formula_parser_factory(self) -> FormulaParserFactory:
        raise NotImplementedError

    @abc.abstractmethod
    def get_dataset_validator_factory(self) -> "DatasetValidatorFactory":
        raise NotImplementedError

    @abc.abstractmethod
    def get_field_id_generator_factory(self) -> FieldIdGeneratorFactory:
        raise NotImplementedError

    @abc.abstractmethod
    def get_supported_functions_manager(self) -> SupportedFunctionsManager:
        raise NotImplementedError

    @abc.abstractmethod
    async def get_subject_resolver(self) -> BaseSubjectResolver:
        raise NotImplementedError

    @abc.abstractmethod
    def get_localizer(self) -> Localizer:
        raise NotImplementedError

    @abc.abstractmethod
    def get_connector_availability(self) -> ConnectorAvailabilityConfig:
        raise NotImplementedError

    @abc.abstractmethod
    def get_multi_query_mutator_factory_factory(self) -> SRMultiQueryMutatorFactory:
        raise NotImplementedError


@attr.s
class DefaultApiServiceRegistry(DefaultServicesRegistry, ApiServiceRegistry):  # noqa
    _default_formula_parser_type: Optional[ParserType] = attr.ib(kw_only=True, default=None)
    _formula_parser_factory: Optional[FormulaParserFactory] = attr.ib(kw_only=True)
    _dataset_validator_factory: Optional["DatasetValidatorFactory"] = attr.ib(kw_only=True)
    _field_id_generator_factory: Optional[FieldIdGeneratorFactory] = attr.ib(kw_only=True)
    _supported_functions_manager: Optional[SupportedFunctionsManager] = attr.ib(kw_only=True, default=None)
    _localizer_factory: Optional[LocalizerFactory] = attr.ib(kw_only=True, default=None)
    _localizer_fallback: Optional[Localizer] = attr.ib(kw_only=True, default=None)
    _connector_availability: Optional[ConnectorAvailabilityConfig] = attr.ib(kw_only=True, default=None)
    _query_proc_mode: QueryProcessingMode = attr.ib(kw_only=True, default=QueryProcessingMode.basic)

    @_formula_parser_factory.default  # noqa
    def _default_formula_parser_factory(self) -> FormulaParserFactory:
        return FormulaParserFactory(default_formula_parser_type=self._default_formula_parser_type)

    @_field_id_generator_factory.default  # noqa
    def _default_field_id_generator_factory(self) -> FieldIdGeneratorFactory:
        return FieldIdGeneratorFactory()

    def get_formula_parser_factory(self) -> FormulaParserFactory:
        assert self._formula_parser_factory is not None
        return self._formula_parser_factory

    def get_dataset_validator_factory(self) -> "DatasetValidatorFactory":
        assert self._dataset_validator_factory is not None
        return self._dataset_validator_factory

    def get_field_id_generator_factory(self) -> FieldIdGeneratorFactory:
        assert self._field_id_generator_factory is not None
        return self._field_id_generator_factory

    def get_supported_functions_manager(self) -> SupportedFunctionsManager:
        assert self._supported_functions_manager is not None
        return self._supported_functions_manager

    async def get_subject_resolver(self) -> BaseSubjectResolver:
        assert self._inst_specific_sr is not None
        return await self._inst_specific_sr.get_subject_resolver()

    def get_localizer(self) -> Localizer:
        assert self._localizer_factory is not None
        return self._localizer_factory.get_for_locale(
            self.rci.locale or "unknown",
            fallback=self._localizer_fallback,
        )

    def get_connector_availability(self) -> ConnectorAvailabilityConfig:
        assert self._connector_availability is not None
        return self._connector_availability

    def get_multi_query_mutator_factory_factory(self) -> SRMultiQueryMutatorFactory:
        return DefaultSRMultiQueryMutatorFactory(query_proc_mode=self._query_proc_mode)

    def close(self) -> None:
        if self._formula_parser_factory is not None:
            self._formula_parser_factory.close()

        super().close()

    async def close_async(self) -> None:
        if self._formula_parser_factory is not None:
            self._formula_parser_factory.close()

        await super().close_async()

lib/dl_api_lib/dl_api_lib/service_registry/service_registry.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
)

import attr

from dl_api_lib.connector_availability.base import ConnectorAvailabilityConfig
from dl_api_lib.service_registry.dataset_validator_factory import DefaultDatasetValidatorFactory
from dl_api_lib.service_registry.field_id_generator_factory import FieldIdGeneratorFactory
from dl_api_lib.service_registry.service_registry import DefaultApiServiceRegistry
from dl_api_lib.service_registry.supported_functions_manager import SupportedFunctionsManager
from dl_constants.enums import QueryProcessingMode
from dl_core.components.ids import FieldIdGeneratorType
from dl_core.services_registry.sr_factories import DefaultSRFactory
from dl_core.services_registry.top_level import ServicesRegistry
from dl_core.utils import FutureRef
from dl_formula.parser.factory import ParserType
from dl_i18n.localizer_base import (
    Localizer,
    LocalizerFactory,
)


if TYPE_CHECKING:
    from dl_api_commons.base_models import RequestContextInfo


@attr.s
class DefaultApiSRFactory(DefaultSRFactory[DefaultApiServiceRegistry]):
    service_registry_cls = DefaultApiServiceRegistry
    _supported_functions_manager: Optional[SupportedFunctionsManager] = attr.ib(kw_only=True, default=None)
    _default_formula_parser_type: Optional[ParserType] = attr.ib(kw_only=True, default=None)
    _field_id_generator_type: FieldIdGeneratorType = attr.ib(kw_only=True, default=FieldIdGeneratorType.readable)
    _localizer_factory: Optional[LocalizerFactory] = attr.ib(default=None)
    _localizer_fallback: Optional[Localizer] = attr.ib(default=None)
    _connector_availability: Optional[ConnectorAvailabilityConfig] = attr.ib(default=None)
    _query_proc_mode: QueryProcessingMode = attr.ib(kw_only=True, default=QueryProcessingMode.basic)

    def additional_sr_constructor_kwargs(
        self,
        request_context_info: RequestContextInfo,
        sr_ref: FutureRef[ServicesRegistry],
    ) -> dict[str, Any]:
        return dict(
            default_formula_parser_type=self._default_formula_parser_type,
            dataset_validator_factory=DefaultDatasetValidatorFactory(
                is_bleeding_edge_user=self.is_bleeding_edge_user(request_context_info),
            ),
            field_id_generator_factory=FieldIdGeneratorFactory(
                field_id_generator_type=self._field_id_generator_type,
            ),
            supported_functions_manager=self._supported_functions_manager,
            localizer_factory=self._localizer_factory,
            localizer_fallback=self._localizer_fallback,
            connector_availability=self._connector_availability,
            query_proc_mode=self._query_proc_mode,
        )

lib/dl_api_lib/dl_api_lib/service_registry/sr_factory.py
datalens-tech/datalens-backend
from functools import reduce
from itertools import chain
from operator import ior

import attr

from dl_api_lib.enums import (
    BI_TYPE_AGGREGATIONS,
    FILTERS_BY_TYPE,
)
from dl_api_lib.query.registry import (
    get_compeng_dialect,
    is_compeng_enabled,
)
from dl_constants.enums import (
    AggregationFunction,
    UserDataType,
    WhereClauseOperation,
)
from dl_formula.core.dialect import DialectCombo
import dl_formula.core.nodes as formula_nodes
from dl_formula.definitions.registry import (
    OPERATION_REGISTRY,
    FuncKey,
    OperationRegistry,
)
from dl_formula.definitions.scope import Scope
from dl_formula.inspect.expression import iter_operation_calls
from dl_query_processing.compilation.filter_compiler import FilterFormulaCompiler
from dl_query_processing.compilation.formula_compiler import FormulaCompiler
from dl_utils.func_tools import method_lru


_FUNCTION_TAG_TO_SCOPE_MAP: dict[str, Scope] = {
    "stable": Scope.STABLE,
}


@attr.s(frozen=True)
class SupportedFunctionsManager:
    """
    Service-helper.
    Translate env type to formula's format and provide some wrappers for get_supported_functions.
    """

    _supported_tags: tuple[str] = attr.ib()
    _operation_registry: OperationRegistry = attr.ib(default=OPERATION_REGISTRY)

    @method_lru(maxsize=1000)
    def get_supported_filters(self, dialect: DialectCombo, user_type: UserDataType) -> list[WhereClauseOperation]:
        return [
            op
            for op in self._get_supported_filters_for_dialect(dialect=dialect)
            if op in FILTERS_BY_TYPE.get(user_type, [])
        ]

    @method_lru(maxsize=1000)
    def get_supported_aggregations(self, dialect: DialectCombo, user_type: UserDataType) -> list[AggregationFunction]:
        supported_func_names = set(
            name
            for name, *_ in self._get_supported_functions(
                dialect=dialect,
                scope=Scope.EXPLICIT_USAGE,
            )
        )

        return [
            ag_type
            for ag_type in BI_TYPE_AGGREGATIONS[user_type]
            # TODO: make _agg_functions public?
            if FormulaCompiler._agg_functions.get(ag_type, ag_type.name) in supported_func_names
        ]

    def get_supported_function_names(self, dialect: DialectCombo) -> list[str]:
        native_supp_funcs = self._get_supported_functions(
            dialect=dialect,
            scope=Scope.SUGGESTED,
        )

        if is_compeng_enabled():
            compeng_supp_funcs = self._get_supported_functions(
                dialect=get_compeng_dialect(),
                scope=Scope.SUGGESTED,
            )
        else:
            compeng_supp_funcs = []

        supported_function_names = sorted(
            {
                func[0]
                for func in chain.from_iterable(
                    (
                        native_supp_funcs,
                        compeng_supp_funcs,
                    )
                )
            }
        )
        return supported_function_names

    @property
    def _supported_scopes(self) -> tuple[int, ...]:
        return tuple(_FUNCTION_TAG_TO_SCOPE_MAP[tag].value for tag in self._supported_tags)

    @method_lru(maxsize=1000)
    def _get_supported_functions(
        self,
        dialect: DialectCombo,
        scope: int = Scope.EXPLICIT_USAGE,
        only_functions: bool = True,
    ) -> list[FuncKey]:
        return self._operation_registry.get_supported_functions(
            require_dialects=dialect,
            only_functions=only_functions,
            function_scopes=scope | reduce(ior, self._supported_scopes),
        )

    @method_lru(maxsize=20)
    def _get_supported_filters_for_dialect(self, dialect: DialectCombo) -> list[WhereClauseOperation]:
        result = []
        supported_funcs = {
            (name, arg_cnt)
            for name, arg_cnt, *_ in self._get_supported_functions(
                dialect=dialect,
                only_functions=False,
                scope=Scope.EXPLICIT_USAGE,
            )
        }
        for filter_op, filter_def in FilterFormulaCompiler.FILTER_OPERATIONS.items():
            arg_cnt = filter_def.arg_cnt
            expr_factory = filter_def.callable
            if arg_cnt is None:
                # args are interpreted as a single list
                # and the field itself is the first argument
                arg_cnt = 2
            else:
                arg_cnt += 1  # The field itself is also an argument

            expr = expr_factory(*(formula_nodes.Null() for _ in range(arg_cnt)))
            for op_node in iter_operation_calls(expr):
                func_key = (op_node.name, len(op_node.args))
                if func_key in supported_funcs:
                    # FIXME: type check
                    result.append(filter_op)
        return result

lib/dl_api_lib/dl_api_lib/service_registry/supported_functions_manager.py
datalens-tech/datalens-backend
"""
Compatibility imports.

It is recommended to import the necessary modules directly,
i.e. `from dl_api_lib.utils.something import ...` instead of `from dl_api_lib.utils import ...`.
"""

from __future__ import annotations

from .base import (  # ...
    chunks,
    need_permission,
    need_permission_on_entry,
    profile_stats,
    query_execution_context,
)


__all__ = (
    "query_execution_context",
    "profile_stats",
    "need_permission",
    "need_permission_on_entry",
    "chunks",
)

lib/dl_api_lib/dl_api_lib/utils/__init__.py
datalens-tech/datalens-backend
""" ... """

from __future__ import annotations

import cProfile
from contextlib import contextmanager
import datetime
import logging
import os
from typing import (
    TYPE_CHECKING,
    Any,
    Iterable,
    Iterator,
    Optional,
)
import uuid

from dl_api_lib.enums import USPermissionKind
from dl_app_tools.profiling_base import GenericProfiler
import dl_core.exc as common_exc
from dl_core.flask_utils.us_manager_middleware import USManagerFlaskMiddleware


# noinspection PyUnresolvedReferences


if TYPE_CHECKING:
    from dl_core.us_entry import USEntry

LOGGER = logging.getLogger(__name__)


@contextmanager
def query_execution_context(
    log_error: bool = True,
    dataset_id: Optional[str] = None,
    version: Optional[str] = None,
    body: Optional[dict] = None,
) -> Iterator[None]:
    try:
        yield  # execute query here
    except common_exc.DatabaseQueryError as err:
        if log_error:
            LOGGER.info(
                "Failed SQL query for dataset %s-%s. Request: %s.  Query: %s. Error: %s",
                dataset_id,
                version,
                str(body)[:1000],
                err.query,
                err.db_message,
                exc_info=True,
            )
        raise


@contextmanager
def profile_stats(stats_dir: Optional[str] = None) -> Iterator[None]:
    """Save profiler stats to file"""
    stats_dir = stats_dir or "./cprofiler"
    if not os.path.exists(stats_dir):
        os.makedirs(stats_dir)
    pr = cProfile.Profile()
    try:
        pr.enable()
        yield
    finally:
        pr.disable()
        filename = os.path.join(
            stats_dir,
            "{}-{}.stats".format(datetime.datetime.utcnow().strftime("%Y-%m-%d-%H-%M-%S"), str(uuid.uuid4())[:4]),
        )
        pr.dump_stats(filename)


def need_permission(us_entry_id: str, permission: USPermissionKind) -> None:
    usm_user = USManagerFlaskMiddleware.get_request_us_manager()
    entry = usm_user.get_by_id(us_entry_id)
    need_permission_on_entry(entry, permission)


def need_permission_on_entry(us_entry: USEntry, permission: USPermissionKind) -> None:
    assert us_entry.permissions is not None
    assert us_entry.uuid is not None
    if not us_entry.permissions[permission.name]:
        raise common_exc.USPermissionRequired(us_entry.uuid, permission.name)


def chunks(lst: list[Any], size: int) -> Iterable[list[Any]]:
    """Yield successive chunks from lst. No padding."""
    for idx in range(0, len(lst), size):
        yield lst[idx : idx + size]


def split_by_quoted_quote(value: str, quote: str = "'") -> tuple[str, str]:
    """
    Parse out a quoted value at the beginning,
    where quotes are quoted by doubling (CSV-like).

    >>> split_by_quoted_quote("'abc'de")
    ('abc', 'de')
    >>> split_by_quoted_quote("'ab''c'''de")
    ("ab'c'", 'de')
    >>> split_by_quoted_quote("'ab''c'''")
    ("ab'c'", '')
    """
    ql = len(quote)
    if not value.startswith(quote):
        raise ValueError("Value does not start with quote")
    value = value[ql:]
    result = []
    while True:
        try:
            next_quote = value.index(quote)
        except ValueError:
            raise ValueError("Unclosed quote")
        value_piece = value[:next_quote]
        result.append(value_piece)
        value = value[next_quote + ql :]
        if value.startswith(quote):
            result.append(quote)
            value = value[ql:]
        else:  # some other text, or end-of-line.
            break

    return "".join(result), value


def quote_by_quote(value: str, quote: str = "'") -> str:
    """
    ...

    >>> quote_by_quote("a'b'")
    "'a''b'''"
    >>> split_by_quoted_quote(quote_by_quote("a'b'") + "and 'stuff'")
    ("a'b'", "and 'stuff'")
    """
    return "{}{}{}".format(quote, value.replace(quote, quote + quote), quote)

lib/dl_api_lib/dl_api_lib/utils/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import itertools
import logging
import re
from typing import (
    Any,
    ClassVar,
    Dict,
    Iterable,
    List,
    NamedTuple,
    Optional,
    Tuple,
)

import attr
import more_itertools

from dl_api_lib import exc
from dl_api_lib.utils.base import (
    chunks,
    quote_by_quote,
    split_by_quoted_quote,
)
from dl_constants.enums import RLSSubjectType
from dl_core.rls import (
    RLS_ALL_SUBJECT,
    RLS_ALL_SUBJECT_NAME,
    RLS_FAILED_USER_NAME_PREFIX,
    RLS_USERID_SUBJECT,
    RLS_USERID_SUBJECT_NAME,
    BaseSubjectResolver,
    RLSEntry,
    RLSPatternType,
    RLSSubject,
)


LOGGER = logging.getLogger(__name__)


class RLSConfigItem(NamedTuple):
    """Structural item of a minimally parsed RLS config"""

    idxes: List[int]  # row numbers
    names: List[str]  # user identifiers


class FieldRLSSerializer:
    allow_all_subject_name: ClassVar[str] = RLS_ALL_SUBJECT_NAME  # `'value': *` cases
    allow_all_subject: ClassVar[RLSSubject] = RLS_ALL_SUBJECT
    userid_subject_name: ClassVar[str] = RLS_USERID_SUBJECT_NAME
    userid_subject: ClassVar[RLSSubject] = RLS_USERID_SUBJECT
    userid_line: ClassVar[str] = "userid: userid"

    @classmethod
    def to_text_config(cls, data: Iterable[RLSEntry]) -> str:
        # Reminder: this only groups consecutive matching values.
        # Should be good enough for this purpose.
        rls_entries_by_value = itertools.groupby(
            data, key=lambda rls_entry: (rls_entry.pattern_type, rls_entry.allowed_value)
        )
        lines = []
        for (pattern_type, allowed_value), rls_entries in rls_entries_by_value:
            subjects_text = ", ".join(rls_entry.subject.subject_name for rls_entry in rls_entries)
            if pattern_type == RLSPatternType.all:
                line = "*: {}".format(subjects_text)
            elif pattern_type == RLSPatternType.userid:
                line = cls.userid_line
            elif pattern_type == RLSPatternType.value:
                assert allowed_value is not None
                line = "{}: {}".format(quote_by_quote(allowed_value), subjects_text)
            else:
                raise Exception("RLS pattern type not yet supported", type(pattern_type), pattern_type)
            lines.append(line)

        return "\n".join(lines)

    # Subject names list parsing regex.
    # Note that this includes allow_all `*` subjects.
    _subjects_re_s = (
        r"(?P<subjects>"  # name
        r"(?:"  # repeat:
        r"[^,]+"  # subject_name
        r"(?:, *)?"  # `'' or ',' or ', +'`
        r")+"
        r")$"
    )
    # allow_all line `*: user1, user2, …`
    _aa_line_re = re.compile(r"^\s*\*\s*:" + _subjects_re_s)
    _uid_line_re = re.compile(r"\s*".join(("^", "userid", ":", "userid", "$")))  # 'userid: userid'
    assert _uid_line_re.match(userid_line)  # self-check
    # 'value': user1, user2, …`. Note that the value might contain more quotes.
    _line_re = re.compile(r"^'.+': " + _subjects_re_s)

    @classmethod
    def _parse_single_line(cls, line: str) -> Tuple[RLSPatternType, Optional[str], List[str]]:
        """
        `'value: subjects'` line to `pattern_type, value, subject_names` tuple.
        """
        aa_match = cls._aa_line_re.match(line)

        if aa_match:  # `*: user1, user2, …`
            pattern_type = RLSPatternType.all
            value = None
            subjects_line = aa_match.group("subjects")

        elif cls._uid_line_re.match(line):
            pattern_type = RLSPatternType.userid
            value = None
            subjects_line = "userid"

        else:  # `'value ''with "quotes"''': user3, user4, …`
            pattern_type = RLSPatternType.value

            if not cls._line_re.match(line):
                raise ValueError("Wrong format")

            value, rest_of_the_line = split_by_quoted_quote(line)
            rest_of_the_line = rest_of_the_line.strip()
            if not rest_of_the_line.startswith(":"):
                raise ValueError(f"Separating ':' expected but not found" f" at {rest_of_the_line[:10]!r}")
            rest_of_the_line = rest_of_the_line[1:].strip()
            subjects_line = rest_of_the_line

        subject_names = [name.strip() for name in subjects_line.split(",")]

        # Some extra validation.
        if cls.allow_all_subject_name in subject_names:
            if pattern_type == RLSPatternType.all:
                raise ValueError(("Wildcard `*: *` is not allowed." " It would effectively disable RLS for the field."))
            if len(subject_names) != 1:
                # Note that this does not check for
                #     value1: *
                #     value1: user, …
                # lines.
                raise ValueError(("Wildcard user must be the only user in line" ", i.e. `…: *`."))

        return pattern_type, value, subject_names

    @classmethod
    def _try_parse_single_line(cls, line: str, idx: int) -> Tuple[RLSPatternType, Optional[str], List[str]]:
        try:
            return cls._parse_single_line(line)
        except ValueError as exc_value:
            raise exc.RLSConfigParsingError(
                f"RLS: Parsing failed at line {idx + 1}", details=dict(description=str(exc_value))
            )

    @classmethod
    def _resolve_subject_names(cls, subject_names: List[str], subject_resolver: BaseSubjectResolver) -> Dict[str, Any]:
        """
        Obtain the subject infos from a subject resolver.
        """
        name_to_subject = {}
        # The chunk size is up to tuning,
        # the primary point is to avoid creating too large JSONs,
        # as it gets deserialized in one async iteration.
        # Note that it makes the 'Logins do not exist' error potentially incomplete.
        names_by_account_type = cls._group_names_by_account_type(subject_names)
        for names_chunk in chunks(names_by_account_type.regular_subject_names, 1000):
            subjects = subject_resolver.get_subjects_by_names(names_chunk)
            for subject in subjects:
                # User name without prefix in dict
                name_to_subject[subject.subject_name.removeprefix(RLS_FAILED_USER_NAME_PREFIX)] = subject
        for service_account_str in names_by_account_type.service_account_strs:
            sa_subject: RLSSubject = cls._parse_sa_account_str(service_account_str)
            name_to_subject[sa_subject.subject_name] = sa_subject
        return name_to_subject

    @attr.s
    class AccountGroups:
        regular_subject_names: List[str] = attr.ib()
        service_account_strs: List[str] = attr.ib()  # sa strs have the following format. @sa:{sa_id}

    @classmethod
    def _group_names_by_account_type(cls, subject_names: List[str]) -> AccountGroups:
        regular_subject_names, service_account_strs = more_itertools.partition(cls._is_sa, subject_names)
        return cls.AccountGroups(list(regular_subject_names), list(service_account_strs))

    @classmethod
    def _is_sa(cls, subject_name: str) -> bool:
        return subject_name.startswith("@sa:")

    @classmethod
    def _parse_sa_account_str(cls, sa_str: str) -> RLSSubject:
        # This is a hack. sa names are passed like this -> @sa:{sa_id}
        sa_id = sa_str.removeprefix("@sa:")
        return RLSSubject(subject_type=RLSSubjectType.user, subject_id=sa_id, subject_name=sa_str)

    @classmethod
    def pre_parse_text_config(cls, config: str) -> Tuple[RLSConfigItem, RLSConfigItem, Dict[str, RLSConfigItem]]:
        """
        Parse the text config into a
        `(allow_all: RLSConfigItem, {field_value: RLSConfigItem, ...})`
        pair.

        `allow_all` here is `*: subject1, …` cases, not to be confused with
        wildcard subjects.
        """
        allow_all_item = RLSConfigItem(idxes=[], names=[])
        userid_item = RLSConfigItem(idxes=[], names=[])
        value_to_item: Dict[str, RLSConfigItem] = {}

        if not config:
            return allow_all_item, userid_item, value_to_item

        for idx, line in enumerate(config.strip().split("\n")):
            pattern_type, value, subject_names = cls._try_parse_single_line(line=line, idx=idx)

            if pattern_type == RLSPatternType.all:
                assert not value
                value_item = allow_all_item
            elif pattern_type == RLSPatternType.userid:
                assert not value
                assert subject_names == ["userid"]
                value_item = userid_item
            elif pattern_type == RLSPatternType.value:
                assert value is not None
                maybe_value_item = value_to_item.get(value)
                if maybe_value_item is not None:
                    value_item = maybe_value_item
                else:
                    value_item = RLSConfigItem(idxes=[], names=[])
                    value_to_item[value] = value_item
            else:
                raise Exception("RLS pattern type is not supported", type(pattern_type), pattern_type)

            value_item.idxes.append(idx)
            # remove failed user name prefix
            value_item.names.extend([name.removeprefix(RLS_FAILED_USER_NAME_PREFIX) for name in subject_names])

        return allow_all_item, userid_item, value_to_item

    @classmethod
    def from_text_config(
        cls, config: str, field_guid: str, subject_resolver: Optional[BaseSubjectResolver]
    ) -> List[RLSEntry]:
        if not config:
            return []

        allow_all_item, userid_item, value_to_item = cls.pre_parse_text_config(config)

        # TODO: check for
        #     value1: *
        #     value1: user1, …
        # lines.

        # TODO: check for
        #     *: user1
        #     value: user1
        # lines.

        all_names = set(
            name
            for value_item in itertools.chain(value_to_item.values(), [allow_all_item])  # wildcard value
            for name in value_item.names
        )
        all_names -= {cls.allow_all_subject_name}  # wildcard subject
        all_names_lst = sorted(all_names)

        if subject_resolver is not None:
            name_to_subject = cls._resolve_subject_names(all_names_lst, subject_resolver=subject_resolver)
        else:
            name_to_subject = {
                name: RLSSubject(
                    subject_type=RLSSubjectType.unknown,
                    subject_id="",
                    subject_name=name,
                )
                for name in all_names_lst
            }
        # `resolve_subjects`-independent hack
        name_to_subject[cls.allow_all_subject_name] = cls.allow_all_subject

        # Combine the results.
        rls_entries = []
        for value, value_info in value_to_item.items():
            names = sorted(set(value_info.names))
            for name in names:
                # TODO?: write down the source config line idx too.
                rls_entries.append(
                    RLSEntry(
                        field_guid=field_guid,
                        allowed_value=value,
                        subject=name_to_subject[name],
                        pattern_type=RLSPatternType.value,
                    )
                )
        for name in allow_all_item.names:
            rls_entries.append(
                RLSEntry(
                    field_guid=field_guid,
                    allowed_value=None,
                    subject=name_to_subject[name],
                    pattern_type=RLSPatternType.all,
                )
            )
        for name in userid_item.names:
            assert name == cls.userid_subject_name
            rls_entries.append(
                RLSEntry(
                    field_guid=field_guid,
                    allowed_value=None,
                    subject=cls.userid_subject,
                    pattern_type=RLSPatternType.userid,
                )
            )

        return rls_entries

lib/dl_api_lib/dl_api_lib/utils/rls.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest

from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_api_lib_testing.data_api_base import DataApiTestBase
from dl_api_lib_testing.dataset_base import DatasetTestBase
from dl_api_lib_tests.db.config import (
    API_TEST_CONFIG,
    DB_CORE_URL,
    CoreConnectionSettings,
)

from dl_connector_clickhouse.core.clickhouse.constants import SOURCE_TYPE_CH_TABLE
from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
from dl_connector_clickhouse.db_testing.engine_wrapper import ClickhouseDbEngineConfig


class DefaultApiTestBase(DataApiTestBase, DatasetTestBase, ConnectionTestBase):
    """The knowledge that this is a ClickHouse connector should not go beyond this class"""

    bi_compeng_pg_on = True
    conn_type = CONNECTION_TYPE_CLICKHOUSE

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return DB_CORE_URL

    @pytest.fixture(scope="class")
    def engine_config(self, db_url: str, engine_params: dict) -> ClickhouseDbEngineConfig:
        return ClickhouseDbEngineConfig(url=db_url, engine_params=engine_params)

    @pytest.fixture(scope="function")
    def environment_readiness(self, enable_all_connectors) -> None:
        pass

    @pytest.fixture(scope="class")
    def connection_params(self) -> dict:
        return dict(
            db_name=CoreConnectionSettings.DB_NAME,
            host=CoreConnectionSettings.HOST,
            port=CoreConnectionSettings.PORT,
            username=CoreConnectionSettings.USERNAME,
            password=CoreConnectionSettings.PASSWORD,
        )

    @pytest.fixture(scope="class")
    def dataset_params(self, sample_table) -> dict:
        return dict(
            source_type=SOURCE_TYPE_CH_TABLE.name,
            parameters=dict(
                db_name=sample_table.db.name,
                table_name=sample_table.name,
            ),
        )

    @pytest.fixture(scope="function")
    def dataset_id(self, saved_dataset: Dataset) -> str:
        return saved_dataset.id

lib/dl_api_lib/dl_api_lib_tests/db/base.py
datalens-tech/datalens-backend
from typing import ClassVar

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=52500).host,
    port_us_http=get_test_container_hostport("us", fallback_port=52500).port,
    host_us_pg=get_test_container_hostport("db-postgres", fallback_port=52509).host,
    port_us_pg_5432=get_test_container_hostport("db-postgres", fallback_port=52509).port,
    us_master_token="AC1ofiek8coB",
    core_connector_ep_names=["clickhouse", "postgresql"],
)


class CoreConnectionSettings:
    DB_NAME: ClassVar[str] = "test_data"
    HOST: ClassVar[str] = get_test_container_hostport("db-clickhouse", fallback_port=52510).host
    PORT: ClassVar[int] = get_test_container_hostport("db-clickhouse", fallback_port=52510).port
    USERNAME: ClassVar[str] = "datalens"
    PASSWORD: ClassVar[str] = "qwerty"


DB_URLS = {
    D.CLICKHOUSE_22_10: f"clickhouse://datalens:qwerty@"
    f'{get_test_container_hostport("db-clickhouse", fallback_port=52510).as_pair()}/test_data',
}
DB_CORE_URL = DB_URLS[D.CLICKHOUSE_22_10]

API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["clickhouse", "postgresql"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
    bi_compeng_pg_url=(
        f"postgresql://us:us@"
        f'{get_test_container_hostport("db-postgres", fallback_port=52509).as_pair()}/us-db-ci_purgeable'
    ),
)

lib/dl_api_lib/dl_api_lib_tests/db/config.py
datalens-tech/datalens-backend
from dl_api_lib_testing.initialization import initialize_api_lib_test
from dl_api_lib_tests.db.config import API_TEST_CONFIG


pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)

lib/dl_api_lib/dl_api_lib_tests/db/conftest.py
datalens-tech/datalens-backend
import shortuuid

from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_core.base_models import PathEntryLocation


class TestDataset(DefaultApiTestBase):
    def test_invalid_dataset_id(self, control_api, saved_connection_id, saved_dataset, sync_us_manager):
        usm = sync_us_manager
        us_client = usm._us_client
        path = PathEntryLocation(shortuuid.uuid())
        dash = us_client.create_entry(scope="dash", key=path)
        dash_id = dash["entryId"]

        resp = control_api.client.get("/api/v1/datasets/{}/versions/draft".format(saved_dataset.id))
        assert resp.status_code == 200

        resp = control_api.client.get("/api/v1/datasets/{}/versions/draft".format(saved_connection_id))
        assert resp.status_code == 404

        resp = control_api.client.get("/api/v1/datasets/{}/versions/draft".format(dash_id))
        assert resp.status_code == 404

    def test_create_entity_with_existing_name(self, control_api, saved_connection_id, saved_dataset, dataset_params):
        name = saved_dataset.name

        second_ds = Dataset(name=name)
        second_ds.sources["source_1"] = second_ds.source(
            connection_id=saved_connection_id,
            **dataset_params,
        )
        second_ds.source_avatars["avatar_1"] = second_ds.sources["source_1"].avatar()

        second_ds = control_api.apply_updates(dataset=second_ds).dataset
        resp = control_api.save_dataset(dataset=second_ds, fail_ok=True)
        assert resp.status_code == 400
        assert resp.json["message"] == "The entry already exists"
        assert resp.json["code"] == "ERR.DS_API.US.BAD_REQUEST.ALREADY_EXISTS"

lib/dl_api_lib/dl_api_lib_tests/db/control_api/test_dataset.py
datalens-tech/datalens-backend
from http import HTTPStatus

import pytest

from dl_api_client.dsmaker.pivot_utils import (
    check_pivot_response,
    get_all_measure_cells,
)
from dl_api_client.dsmaker.shortcuts.dataset import add_formulas_to_dataset
from dl_api_client.dsmaker.shortcuts.result_data import get_regular_result_data
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import PivotRole
from dl_constants.internal_constants import (
    DIMENSION_NAME_TITLE,
    MEASURE_NAME_TITLE,
)


class TestPivotWithAnnotations(DefaultApiTestBase):
    def test_pivot_multiple_measures_with_annotation(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([profit])",
                "order count": "COUNT([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", "sub_category"],
            rows=["order_date", "city", MEASURE_NAME_TITLE],
            measures=["sales sum", "profit sum"],
            annotations=["order count"],
            min_col_cnt=10,
            min_row_cnt=100,
            min_value_cnt=100,
        )

    def test_pivot_multiple_measures_with_targeted_annotation(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([profit])",
                "order count": "COUNT([sales])",
            },
        )

        data_by_field = get_regular_result_data(
            ds,
            data_api,
            field_names=[
                "category",
                "sub_category",
                "order_date",
                "city",
                "sales sum",
                "profit sum",
                "order count",
            ],
        )
        category_leg_id = 0
        subcategory_leg_id = 1
        orderdate_leg_id = 2
        city_leg_id = 3
        mname_leg_id = 4
        profit_leg_id = 5
        sales_leg_id = 6
        anno_leg_id = 7

        result_resp = data_api.get_pivot(
            dataset=ds,
            fields=[
                ds.find_field(title="category").as_req_legend_item(legend_item_id=category_leg_id),
                ds.find_field(title="sub_category").as_req_legend_item(legend_item_id=subcategory_leg_id),
                ds.find_field(title="order_date").as_req_legend_item(legend_item_id=orderdate_leg_id),
                ds.find_field(title="city").as_req_legend_item(legend_item_id=city_leg_id),
                ds.measure_name_as_req_legend_item(legend_item_id=mname_leg_id),
                ds.find_field(title="sales sum").as_req_legend_item(legend_item_id=sales_leg_id),
                ds.find_field(title="profit sum").as_req_legend_item(legend_item_id=profit_leg_id),
                ds.find_field(title="order count").as_req_legend_item(legend_item_id=anno_leg_id),
            ],
            pivot_structure=[
                ds.make_req_pivot_item(role=PivotRole.pivot_column, legend_item_ids=[category_leg_id]),
                ds.make_req_pivot_item(role=PivotRole.pivot_column, legend_item_ids=[subcategory_leg_id]),
                ds.make_req_pivot_item(role=PivotRole.pivot_row, legend_item_ids=[orderdate_leg_id]),
                ds.make_req_pivot_item(role=PivotRole.pivot_row, legend_item_ids=[city_leg_id]),
                ds.make_req_pivot_item(role=PivotRole.pivot_row, legend_item_ids=[mname_leg_id]),
                ds.make_req_pivot_item(role=PivotRole.pivot_measure, legend_item_ids=[sales_leg_id]),
                ds.make_req_pivot_item(role=PivotRole.pivot_measure, legend_item_ids=[profit_leg_id]),
                ds.make_req_pivot_item(
                    role=PivotRole.pivot_annotation,
                    legend_item_ids=[anno_leg_id],
                    annotation_type="color",
                    target_legend_item_ids=[profit_leg_id],
                ),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK
        result_data = result_resp.data

        # Check legend
        legend_tuples = [(item.title, item.role_spec.role) for item in result_data["pivot"]["structure"]]
        assert legend_tuples == [
            ("category", PivotRole.pivot_column),
            ("sub_category", PivotRole.pivot_column),
            ("order_date", PivotRole.pivot_row),
            ("city", PivotRole.pivot_row),
            (MEASURE_NAME_TITLE, PivotRole.pivot_row),
            ("sales sum", PivotRole.pivot_measure),
            ("profit sum", PivotRole.pivot_measure),
            ("order count", PivotRole.pivot_annotation),
            (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
        ]

        # Check measure values
        pivot_rows = result_data["pivot_data"]["rows"]
        # Check annotation values
        all_measure_cells = get_all_measure_cells(pivot_rows)
        sales_annotation_values = sorted(
            {float(cell[1][0]) for cell in all_measure_cells if cell[0][1] == sales_leg_id and len(cell) > 1}
        )
        profit_annotation_values = sorted(
            {float(cell[1][0]) for cell in all_measure_cells if cell[0][1] == profit_leg_id and len(cell) > 1}
        )
        original_annotation_values = sorted({float(val) for val in data_by_field["order count"]})
        assert sales_annotation_values == []
        assert profit_annotation_values == pytest.approx(original_annotation_values)

    def test_pivot_multi_measures_with_annotation_same_as_one_measure(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "Count": "COUNT()",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", "sub_category"],
            rows=["order_date", "city", MEASURE_NAME_TITLE],
            measures=["sales sum", "Count"],
            annotations=["sales sum"],
            min_col_cnt=10,
            min_row_cnt=100,
            min_value_cnt=100,
        )

    def test_pivot_with_multiple_annotations(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([profit])",
                "order count": "COUNT([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["order_date"],
            measures=["sales sum"],
            annotations=["profit sum", "order count"],
        )

    def test_pivot_with_annotation_by_dimension(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "order count": "COUNT([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["order_date"],
            measures=["sales sum"],
            annotations=["order_date"],
        )

lib/dl_api_lib/dl_api_lib_tests/db/data_api/pivot/test_annotations.py
datalens-tech/datalens-backend
from http import HTTPStatus
from typing import Optional

from dl_api_client.dsmaker.pivot_utils import (
    check_pivot_response,
    get_pivot_response,
)
from dl_api_client.dsmaker.primitives import PivotPagination
from dl_api_client.dsmaker.shortcuts.dataset import add_formulas_to_dataset
from dl_api_client.dsmaker.shortcuts.result_data import get_regular_result_data
from dl_api_lib.pivot.primitives import (
    PivotHeaderValue,
    PivotMeasureSorting,
    PivotMeasureSortingSettings,
)
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import OrderDirection
from dl_constants.internal_constants import MEASURE_NAME_TITLE


class TestBasicPivot(DefaultApiTestBase):
    def test_basic_pivot(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["order_date"],
            measures=["sales sum"],
            min_col_cnt=3,
            min_row_cnt=100,
            min_value_cnt=100,
        )

    def test_pivot_multiple_measures(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", "sub_category"],
            rows=["order_date", MEASURE_NAME_TITLE, "city"],
            measures=["sales sum", "profit sum"],
            measures_sorting_settings=[
                PivotMeasureSorting(
                    row=PivotMeasureSortingSettings(
                        header_values=[
                            PivotHeaderValue(value="2014-01-04"),
                            PivotHeaderValue(value="sales sum"),
                            PivotHeaderValue(value="Naperville"),
                        ],
                    )
                ),
                None,
            ],
            min_col_cnt=10,
            min_row_cnt=100,
            min_value_cnt=100,
        )

        sorting_row_idx = None
        for row in pivot_abs.iter_rows():
            if row.get_compound_header() == ("2014-01-04", "sales sum", "Naperville"):
                sorting_row_idx = row.row_idx
                break
        assert sorting_row_idx is not None

        def _get_value(value: Optional[tuple]) -> float:
            if value is None:
                return float("-inf")
            return float(value[0][0])

        row_values = list(map(_get_value, pivot_abs.resp_data["pivot_data"]["rows"][sorting_row_idx]["values"]))
        assert sorted(row_values) == row_values

    def test_pivot_with_order_by(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        data_by_field = get_regular_result_data(
            ds,
            data_api,
            field_names=[
                "category",
                "sub_category",
                "order_date",
                "city",
                "sales sum",
                "profit sum",
            ],
        )

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", "sub_category"],
            rows=["order_date", "city", MEASURE_NAME_TITLE],
            measures=["sales sum", "profit sum"],
            order_fields={"category": OrderDirection.desc, "order_date": OrderDirection.desc},
        )
        assert pivot_resp.status_code == HTTPStatus.OK
        result_data = pivot_resp.data

        # Check first column dimension
        category_values = []
        for col in result_data["pivot_data"]["columns"]:
            value = col[0][0][0]
            if value not in category_values:
                category_values.append(value)
        assert category_values == sorted(set(data_by_field["category"]), reverse=True)

        # Check first row dimension
        pivot_rows = result_data["pivot_data"]["rows"]
        date_values = []
        for row in pivot_rows:
            value = row["header"][0][0][0]
            if value not in date_values:
                date_values.append(value)
        assert date_values == sorted(set(data_by_field["order_date"]), reverse=True)

    def test_pivot_with_pagination(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        def get_pivot(pivot_pagination: Optional[PivotPagination] = None):
            pivot_resp = get_pivot_response(
                dataset=ds,
                data_api=data_api,
                columns=["category", "sub_category"],
                rows=["order_date", "city", MEASURE_NAME_TITLE],
                measures=["sales sum", "profit sum"],
                pivot_pagination=pivot_pagination,
            )
            assert pivot_resp.status_code == HTTPStatus.OK, pivot_resp.response_errors
            return pivot_resp.data

        # Save unpaginated table:
        result_data = get_pivot()
        initial_columns = result_data["pivot_data"]["columns"]
        initial_rows = result_data["pivot_data"]["rows"]

        # Get paginated table and compare
        result_data = get_pivot(pivot_pagination=PivotPagination(offset_rows=1, limit_rows=2))
        paginated_columns = result_data["pivot_data"]["columns"]
        paginated_rows = result_data["pivot_data"]["rows"]
        assert paginated_columns == initial_columns
        assert paginated_rows == initial_rows[1:3]

        # Pseudo-pagination
        result_data = get_pivot(pivot_pagination=PivotPagination(offset_rows=0, limit_rows=None))
        paginated_columns = result_data["pivot_data"]["columns"]
        paginated_rows = result_data["pivot_data"]["rows"]
        assert paginated_columns == initial_columns
        assert paginated_rows == initial_rows

    def test_pivot_with_remapped_titles(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["order_date"],
            measures=["sales sum"],
            title_mapping={
                "category": "My Dimension 1",
                "order_date": "My Dimension 2",
                "sales sum": "Measure",
            },
            min_col_cnt=3,
            min_row_cnt=100,
            min_value_cnt=100,
        )

lib/dl_api_lib/dl_api_lib_tests/db/data_api/pivot/test_basic.py
datalens-tech/datalens-backend
from http import HTTPStatus
import json
from typing import Any

from dl_api_client.dsmaker.data_abstraction.pivot import PivotDataAbstraction
from dl_api_client.dsmaker.pivot_utils import (
    check_pivot_response,
    get_all_measure_cells,
)
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_client.dsmaker.shortcuts.dataset import add_formulas_to_dataset
from dl_api_client.dsmaker.shortcuts.result_data import get_regular_result_data
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import (
    OrderDirection,
    PivotRole,
)
from dl_constants.internal_constants import (
    DIMENSION_NAME_TITLE,
    MEASURE_NAME_TITLE,
)


class TestPivotCornerCases(DefaultApiTestBase):
    def test_pivot_with_markup(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "category MU": "BOLD([category])",
                "order_date MU": "ITALIC(STR([order_date]))",
                "measure MU": 'BOLD(STR(SUM([sales]))) + " - " + ITALIC(STR(SUM([profit])))',
            },
        )

        data_by_field = get_regular_result_data(
            ds, data_api, field_names=["category MU", "order_date MU", "measure MU"]
        )

        category_liid = 0
        orderdate_liid = 1
        measure_liid = 2

        result_resp = data_api.get_pivot(
            dataset=ds,
            fields=[
                ds.find_field(title="category MU").as_req_legend_item(legend_item_id=category_liid),
                ds.find_field(title="order_date MU").as_req_legend_item(legend_item_id=orderdate_liid),
                ds.find_field(title="measure MU").as_req_legend_item(legend_item_id=measure_liid),
            ],
            pivot_structure=[
                ds.make_req_pivot_item(role=PivotRole.pivot_column, legend_item_ids=[category_liid]),
                ds.make_req_pivot_item(role=PivotRole.pivot_row, legend_item_ids=[orderdate_liid]),
                ds.make_req_pivot_item(role=PivotRole.pivot_measure, legend_item_ids=[measure_liid]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK
        result_data = result_resp.data

        # Check legend
        legend_tuples = [(item.title, item.role_spec.role) for item in result_data["pivot"]["structure"]]
        assert legend_tuples == [
            ("category MU", PivotRole.pivot_column),
            ("order_date MU", PivotRole.pivot_row),
            ("measure MU", PivotRole.pivot_measure),
            (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
        ]

        def dumps(value: Any) -> str:
            return json.dumps(value, sort_keys=True)

        # Check columns
        category_values = {dumps(col[0][0][0]) for col in result_data["pivot_data"]["columns"]}
        assert category_values == {
            dumps({"type": "bold", "content": {"type": "text", "content": "Office Supplies"}}),
            dumps({"type": "bold", "content": {"type": "text", "content": "Furniture"}}),
            dumps({"type": "bold", "content": {"type": "text", "content": "Technology"}}),
        }

        # Check row headers
        pivot_rows = result_data["pivot_data"]["rows"]
        date_values = {dumps(row["header"][0][0][0]) for row in pivot_rows}
        assert date_values.issuperset(
            {
                dumps({"type": "italics", "content": {"type": "text", "content": "2014-01-03"}}),
                dumps({"type": "italics", "content": {"type": "text", "content": "2014-01-06"}}),
                dumps({"type": "italics", "content": {"type": "text", "content": "2014-01-09"}}),
            }
        )

        # Check measure values
        assert len(pivot_rows) > 100
        all_measure_cells = get_all_measure_cells(pivot_rows)
        measure_values = sorted({dumps(cell[0][0]) for cell in all_measure_cells})
        original_measure_values = sorted({dumps(val) for val in data_by_field["measure MU"]})
        assert measure_values == original_measure_values

    def test_pivot_no_measures(self, control_api, data_api, dataset_id):
        ds = control_api.load_dataset(dataset=Dataset(id=dataset_id)).dataset

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", "sub_category"],
            rows=["order_date", "city"],
            measures=[],
            min_col_cnt=10,
            min_row_cnt=100,
            max_value_cnt=0,
        )

    def test_pivot_no_dimensions_multiple_measures(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=[],
            rows=[],
            measures=["sales sum", "profit sum"],
            min_col_cnt=2,
            max_col_cnt=2,
            min_row_cnt=1,
            max_row_cnt=1,
            min_value_cnt=2,
            max_value_cnt=2,
            custom_pivot_legend_check=[
                ("sales sum", PivotRole.pivot_measure),
                ("profit sum", PivotRole.pivot_measure),
                (MEASURE_NAME_TITLE, PivotRole.pivot_column),
                (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
            ],
        )

    def test_pivot_only_row_dimensions_one_measure(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=[],
            rows=["category", "order_date"],
            measures=["sales sum"],
            min_col_cnt=1,
            max_col_cnt=1,
            custom_pivot_legend_check=[
                ("category", PivotRole.pivot_row),
                ("order_date", PivotRole.pivot_row),
                ("sales sum", PivotRole.pivot_measure),
                (MEASURE_NAME_TITLE, PivotRole.pivot_column),
                (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
            ],
        )

    def test_pivot_only_column_dimensions_multiple_measures_no_mnames(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", "order_date"],
            rows=[],
            measures=["sales sum", "profit sum"],
            min_row_cnt=2,
            max_row_cnt=2,
            custom_pivot_legend_check=[
                ("category", PivotRole.pivot_column),
                ("order_date", PivotRole.pivot_column),
                ("sales sum", PivotRole.pivot_measure),
                ("profit sum", PivotRole.pivot_measure),
                (MEASURE_NAME_TITLE, PivotRole.pivot_row),
                (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
            ],
        )

    def test_pivot_only_column_dimensions_multiple_measures_with_mnames(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", "order_date", MEASURE_NAME_TITLE],
            rows=[],
            measures=["sales sum", "profit sum"],
            min_row_cnt=1,
            max_row_cnt=1,
        )

    def test_single_measure_with_duplicate_measure_name(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        category_liid = 0
        orderdate_liid = 1
        sales_liid = 2
        mname_liid = 3

        def get_pivot(duplicate_measure_name: bool) -> dict:
            result_resp = data_api.get_pivot(
                dataset=ds,
                fields=[
                    ds.find_field(title="category").as_req_legend_item(legend_item_id=category_liid),
                    ds.find_field(title="order_date").as_req_legend_item(legend_item_id=orderdate_liid),
                    ds.find_field(title="sales sum").as_req_legend_item(legend_item_id=sales_liid),
                    ds.measure_name_as_req_legend_item(legend_item_id=mname_liid),
                ],
                pivot_structure=[
                    ds.make_req_pivot_item(role=PivotRole.pivot_column, legend_item_ids=[category_liid]),
                    ds.make_req_pivot_item(role=PivotRole.pivot_column, legend_item_ids=[mname_liid]),
                    ds.make_req_pivot_item(role=PivotRole.pivot_row, legend_item_ids=[orderdate_liid]),
                    *(
                        (ds.make_req_pivot_item(role=PivotRole.pivot_row, legend_item_ids=[mname_liid]),)
                        if duplicate_measure_name
                        else ()
                    ),
                    ds.make_req_pivot_item(role=PivotRole.pivot_measure, legend_item_ids=[sales_liid]),
                ],
                fail_ok=True,
            )
            assert result_resp.status_code == HTTPStatus.OK
            return result_resp.data

        single_result_data = get_pivot(duplicate_measure_name=False)
        double_result_data = get_pivot(duplicate_measure_name=True)

        # Check legend
        legend_tuples = [(item.title, item.role_spec.role) for item in double_result_data["pivot"]["structure"]]
        assert legend_tuples == [
            ("category", PivotRole.pivot_column),
            ("Measure Names", PivotRole.pivot_column),
            ("order_date", PivotRole.pivot_row),
            ("Measure Names", PivotRole.pivot_row),
            ("sales sum", PivotRole.pivot_measure),
            (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
        ]

        # Check column headers
        for col in double_result_data["pivot_data"]["columns"]:
            assert len(col) == 2

        # Check row headers
        for row in double_result_data["pivot_data"]["rows"]:
            assert len(row["header"]) == 2

        assert len(double_result_data["pivot_data"]["rows"]) == len(single_result_data["pivot_data"]["rows"])
        for single_row, double_row in zip(
            single_result_data["pivot_data"]["rows"], double_result_data["pivot_data"]["rows"]
        ):
            assert len(single_row["values"]) == len(double_row["values"])
            for single_cell, double_cell in zip(single_row["values"], double_row["values"]):
                if single_cell is None or double_cell is None:
                    assert single_cell == double_cell
                else:
                    assert len(single_cell) == len(double_cell) == 1
                    assert single_cell[0][0:2] == double_cell[0][0:2]  # exclude pivot_item_id (idx=2) from comparison

    def test_pivot_only_row_dimensions_no_measures(self, control_api, data_api, dataset_id):
        ds = control_api.load_dataset(dataset=Dataset(id=dataset_id)).dataset

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=[],
            rows=["category", "order_date"],
            measures=[],
            max_col_cnt=1,  # There will be 1 column without any headers or values
            min_row_cnt=100,
            max_value_cnt=0,
        )

    def test_pivot_only_single_column_dimension_no_measures(self, control_api, data_api, dataset_id):
        ds = control_api.load_dataset(dataset=Dataset(id=dataset_id)).dataset

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["order_date"],
            rows=[],
            measures=[],
            max_row_cnt=1,  # There will be 1 row without any headers or values
            min_col_cnt=100,
            max_value_cnt=0,
        )

    def test_pivot_duplicate_measures(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", "sub_category"],
            rows=["order_date", "city", MEASURE_NAME_TITLE],
            measures=["sales sum", "profit sum", "sales sum"],
            min_col_cnt=10,
            min_row_cnt=100,
            min_value_cnt=100,
        )

    def test_pivot_empty_string_dimension_values(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "empty city": 'IF [city] = "New York" THEN "" ELSE [city] END',
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["empty city"],
            rows=["category"],
            measures=["sales sum"],
            order_fields={"empty city": OrderDirection.desc},
            with_totals=True,
        )

    def test_pivot_null_dimension_values(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "null city": 'IF [city] = "New York" THEN NULL ELSE [city] END',
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["null city"],
            rows=["category"],
            measures=["sales sum"],
            with_totals=True,
        )

    def test_pivot_sorting_with_totals(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        def get_pivot(direction: OrderDirection) -> PivotDataAbstraction:
            return check_pivot_response(
                dataset=ds,
                data_api=data_api,
                columns=["category"],
                rows=["region"],
                measures=["sales sum"],
                order_fields={"category": direction},
                with_totals=True,
            )

        pivot_abs = get_pivot(OrderDirection.asc)
        col_titles = pivot_abs.get_flat_column_headers()
        assert col_titles == ["Furniture", "Office Supplies", "Technology", ""]

        pivot_abs = get_pivot(OrderDirection.desc)
        col_titles = pivot_abs.get_flat_column_headers()
        assert col_titles == ["Technology", "Office Supplies", "Furniture", ""]

lib/dl_api_lib/dl_api_lib_tests/db/data_api/pivot/test_corner_cases.py
datalens-tech/datalens-backend
from http import HTTPStatus

from dl_api_client.dsmaker.pivot_utils import get_pivot_response
from dl_api_client.dsmaker.primitives import (
    Dataset,
    PivotPagination,
    PivotTotals,
)
from dl_api_client.dsmaker.shortcuts.dataset import add_formulas_to_dataset
from dl_api_lib.pivot.primitives import (
    PivotHeaderValue,
    PivotMeasureSorting,
    PivotMeasureSortingSettings,
)
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import (
    FieldRole,
    PivotRole,
)
from dl_constants.internal_constants import MEASURE_NAME_TITLE


class TestPivotErrors(DefaultApiTestBase):
    def test_multiple_measures_without_measure_name(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["order_date"],
            measures=["sales sum", "profit sum"],
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.MEASURE_NAME.REQUIRED"

    def test_no_measures_with_measure_name(self, control_api, data_api, dataset_id):
        ds = control_api.load_dataset(dataset=Dataset(id=dataset_id)).dataset

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["order_date", MEASURE_NAME_TITLE],
            measures=[],
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.MEASURE_NAME.FORBIDDEN"

    def test_multiple_measures_with_double_measure_name(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", MEASURE_NAME_TITLE],
            rows=["order_date", MEASURE_NAME_TITLE],
            measures=["sales sum", "profit sum"],
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.MEASURE_NAME.DUPLICATE"

    def test_pagination_range(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        def check_pagination_error(pivot_pagination: PivotPagination) -> None:
            pivot_resp = get_pivot_response(
                dataset=ds,
                data_api=data_api,
                columns=["category"],
                rows=["order_date"],
                measures=["sales sum"],
                pivot_pagination=pivot_pagination,
            )
            assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST

        check_pagination_error(PivotPagination(limit_rows=-9, offset_rows=8))
        check_pagination_error(PivotPagination(limit_rows=9, offset_rows=-8))
        check_pagination_error(PivotPagination(limit_rows=0, offset_rows=1))

    def test_measure_as_dimension(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([profit])",
            },
        )

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["sales sum"],
            measures=["profit sum"],
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.LEGEND.INVALID_ROLE"

    def test_dimension_as_measure(self, control_api, data_api, dataset_id):
        ds = control_api.load_dataset(dataset=Dataset(id=dataset_id)).dataset

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["city"],
            measures=["order_date"],
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.LEGEND.INVALID_ROLE"

    def test_uneven_data_columns(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        pivot_resp = data_api.get_pivot(
            dataset=ds,
            fields=[
                ds.find_field(title="category").as_req_legend_item(legend_item_id=0, block_id=0),
                ds.find_field(title="sales sum").as_req_legend_item(legend_item_id=1, block_id=0),
                ds.find_field(title="order_date").as_req_legend_item(legend_item_id=2, block_id=1),
                ds.find_field(title="sales sum").as_req_legend_item(legend_item_id=3, role=FieldRole.total, block_id=1),
            ],
            pivot_structure=[
                ds.make_req_pivot_item(role=PivotRole.pivot_row, legend_item_ids=[0]),
                ds.make_req_pivot_item(role=PivotRole.pivot_column, legend_item_ids=[2]),
                ds.make_req_pivot_item(role=PivotRole.pivot_measure, legend_item_ids=[1, 3]),
            ],
            fail_ok=True,
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.UNEVEN_DATA_COLUMNS"

    def test_wrong_column_in_measure_sorting(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["order_date"],
            measures=["sales sum"],
            measures_sorting_settings=[
                PivotMeasureSorting(
                    column=PivotMeasureSortingSettings(header_values=[PivotHeaderValue(value="Not found")])
                )
            ],
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.SORTING.ROW_OR_COLUMN_NOT_FOUND"

    def test_wrong_role_spec_in_measure_sorting(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["order_date"],
            measures=["sales sum"],
            totals=[()],
            measures_sorting_settings=[
                PivotMeasureSorting(
                    column=PivotMeasureSortingSettings(
                        header_values=[PivotHeaderValue(value="")]
                    )  # sort by data, not totals!
                )
            ],
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.SORTING.ROW_OR_COLUMN_NOT_FOUND"

    def test_measure_sorting_multiple_sort_on_the_same_axis(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )
        sorting_settings = PivotMeasureSorting(
            column=PivotMeasureSortingSettings(
                header_values=[PivotHeaderValue(value="Furniture"), PivotHeaderValue(value="sales sum")]
            )
        )

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", MEASURE_NAME_TITLE],
            rows=["order_date"],
            measures=["sales sum", "profit sum"],
            measures_sorting_settings=[sorting_settings, sorting_settings],
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.SORTING.MULTIPLE_COLUMNS_OR_ROWS"

    def test_measure_sorting_by_column_with_multiple_measures_in_rows(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )
        sorting_settings = PivotMeasureSorting(
            column=PivotMeasureSortingSettings(
                header_values=[PivotHeaderValue(value="Furniture"), PivotHeaderValue(value="sales sum")]
            )
        )

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["order_date", MEASURE_NAME_TITLE],
            measures=["sales sum", "profit sum"],
            measures_sorting_settings=[sorting_settings, None],
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.SORTING.AGAINST_MULTIPLE_MEASURES"

    def test_measure_sorting_with_subtotals(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        pivot_resp = get_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=[MEASURE_NAME_TITLE],
            rows=["category", "order_date"],
            measures=["sales sum", "profit sum"],
            measures_sorting_settings=[
                PivotMeasureSorting(
                    column=PivotMeasureSortingSettings(header_values=[PivotHeaderValue(value="sales sum")])
                ),
                None,
            ],
            simple_totals=PivotTotals(
                rows=[PivotTotals.item(level=1)],
            ),
        )
        assert pivot_resp.status_code == HTTPStatus.BAD_REQUEST
        assert pivot_resp.bi_status_code == "ERR.DS_API.PIVOT.SORTING.SUBTOTALS_ARE_NOT_ALLOWED"

lib/dl_api_lib/dl_api_lib_tests/db/data_api/pivot/test_errors.py
datalens-tech/datalens-backend
from dl_api_client.dsmaker.pivot_utils import check_pivot_response
from dl_api_client.dsmaker.primitives import PivotTotals
from dl_api_client.dsmaker.shortcuts.dataset import add_formulas_to_dataset
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import PivotRole
from dl_constants.internal_constants import (
    DIMENSION_NAME_TITLE,
    MEASURE_NAME_TITLE,
)


class TestPivotWithSimpleTotals(DefaultApiTestBase):
    def test_main_totals(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["region"],
            measures=["sales sum"],
            simple_totals=PivotTotals(
                rows=[PivotTotals.item(level=0)],
                columns=[PivotTotals.item(level=0)],
            ),
        )
        col_titles = pivot_abs.get_flat_column_headers()
        row_titles = pivot_abs.get_flat_row_headers()
        assert col_titles == ["Furniture", "Office Supplies", "Technology", ""]
        assert row_titles == ["Central", "East", "South", "West", ""]

    def test_with_totals_flag(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["region"],
            measures=["sales sum"],
            with_totals=True,
        )
        col_titles = pivot_abs.get_flat_column_headers()
        row_titles = pivot_abs.get_flat_row_headers()
        assert col_titles == ["Furniture", "Office Supplies", "Technology", ""]
        assert row_titles == ["Central", "East", "South", "West", ""]

    def test_corner_case_totals(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=[],
            measures=["sales sum"],
            simple_totals=PivotTotals(
                rows=[],
                columns=[PivotTotals.item(level=0)],
            ),
            custom_pivot_legend_check=[
                ("category", PivotRole.pivot_column),
                ("sales sum", PivotRole.pivot_measure),
                (MEASURE_NAME_TITLE, PivotRole.pivot_row),
                (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
            ],
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=[],
            rows=["category"],
            measures=["sales sum"],
            simple_totals=PivotTotals(
                rows=[PivotTotals.item(level=0)],
                columns=[],
            ),
            custom_pivot_legend_check=[
                ("category", PivotRole.pivot_row),
                ("sales sum", PivotRole.pivot_measure),
                (MEASURE_NAME_TITLE, PivotRole.pivot_column),
                (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
            ],
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["region"],
            measures=[],
            simple_totals=PivotTotals(
                rows=[PivotTotals.item(level=0)],
                columns=[PivotTotals.item(level=0)],
            ),
        )

    def test_multi_measure_corner_case_totals_flag(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([profit])",
            },
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=[],
            measures=["sales sum", "profit sum"],
            with_totals=True,
            custom_pivot_legend_check=[
                ("category", PivotRole.pivot_column),
                ("sales sum", PivotRole.pivot_measure),
                ("profit sum", PivotRole.pivot_measure),
                (MEASURE_NAME_TITLE, PivotRole.pivot_row),
                (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
            ],
        )

        check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=[],
            rows=["category"],
            measures=["sales sum", "profit sum"],
            simple_totals=PivotTotals(
                rows=[PivotTotals.item(level=0)],
                columns=[],
            ),
            custom_pivot_legend_check=[
                ("category", PivotRole.pivot_row),
                ("sales sum", PivotRole.pivot_measure),
                ("profit sum", PivotRole.pivot_measure),
                (MEASURE_NAME_TITLE, PivotRole.pivot_column),
                (DIMENSION_NAME_TITLE, PivotRole.pivot_info),
            ],
        )

    def test_main_totals_with_annotation(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["region"],
            measures=["sales sum"],
            annotations=["profit sum"],
            simple_totals=PivotTotals(
                rows=[PivotTotals.item(level=0)],
                columns=[PivotTotals.item(level=0)],
            ),
        )
        col_titles = pivot_abs.get_flat_column_headers()
        row_titles = pivot_abs.get_flat_row_headers()
        assert col_titles == ["Furniture", "Office Supplies", "Technology", ""]
        assert row_titles == ["Central", "East", "South", "West", ""]

    def test_main_totals_with_multiple_measures(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["region", MEASURE_NAME_TITLE],
            measures=["sales sum", "profit sum"],
            simple_totals=PivotTotals(
                rows=[PivotTotals.item(level=0)],
                columns=[PivotTotals.item(level=0)],
            ),
        )
        col_titles = pivot_abs.get_flat_column_headers()
        row_compound_titles = pivot_abs.get_compound_row_headers()
        assert col_titles == ["Furniture", "Office Supplies", "Technology", ""]
        assert row_compound_titles == [
            ("Central", "sales sum"),
            ("Central", "profit sum"),
            ("East", "sales sum"),
            ("East", "profit sum"),
            ("South", "sales sum"),
            ("South", "profit sum"),
            ("West", "sales sum"),
            ("West", "profit sum"),
            ("", "sales sum"),
            ("", "profit sum"),
        ]

    def test_subtotals_with_multiple_measures(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=[MEASURE_NAME_TITLE, "region"],
            measures=["sales sum", "profit sum"],
            simple_totals=PivotTotals(
                rows=[PivotTotals.item(level=1)],
                columns=[PivotTotals.item(level=0)],
            ),
        )
        col_titles = pivot_abs.get_flat_column_headers()
        row_compound_titles = pivot_abs.get_compound_row_headers()
        assert col_titles == ["Furniture", "Office Supplies", "Technology", ""]
        assert row_compound_titles == [
            ("sales sum", "Central"),
            ("sales sum", "East"),
            ("sales sum", "South"),
            ("sales sum", "West"),
            ("sales sum", ""),
            ("profit sum", "Central"),
            ("profit sum", "East"),
            ("profit sum", "South"),
            ("profit sum", "West"),
            ("profit sum", ""),
        ]

    def test_main_totals_with_only_mnames_one_one_side(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([sales])",
            },
        )

        # 1. Only Measure Names in rows
        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=[MEASURE_NAME_TITLE],
            measures=["sales sum", "profit sum"],
            with_totals=True,
            check_totals=[()],  # override the autogenerated value because this is a corner case
        )
        col_titles = pivot_abs.get_flat_column_headers()
        row_titles = pivot_abs.get_flat_row_headers()
        assert col_titles == ["Furniture", "Office Supplies", "Technology", ""]
        assert row_titles == ["sales sum", "profit sum"]

        # 2. Only Measure Names in columns
        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=[MEASURE_NAME_TITLE],
            rows=["category"],
            measures=["sales sum", "profit sum"],
            with_totals=True,
            check_totals=[()],  # override the autogenerated value because this is a corner case
        )
        col_titles = pivot_abs.get_flat_column_headers()
        row_titles = pivot_abs.get_flat_row_headers()
        assert col_titles == ["sales sum", "profit sum"]
        assert row_titles == ["Furniture", "Office Supplies", "Technology", ""]

lib/dl_api_lib/dl_api_lib_tests/db/data_api/pivot/test_simple_totals.py
datalens-tech/datalens-backend
from typing import Optional

import pytest

from dl_api_client.dsmaker.pivot_utils import check_pivot_response
from dl_api_client.dsmaker.shortcuts.dataset import add_formulas_to_dataset
from dl_api_lib.pivot.primitives import (
    PivotHeaderRoleSpec,
    PivotHeaderValue,
    PivotMeasureSorting,
    PivotMeasureSortingSettings,
)
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import (
    NotificationType,
    OrderDirection,
    PivotHeaderRole,
    WhereClauseOperation,
)
from dl_constants.internal_constants import MEASURE_NAME_TITLE


class TestPivotWithTotals(DefaultApiTestBase):
    def test_basic_pivot_with_grand_total(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["region"],
            measures=["sales sum"],
            totals=[()],
        )
        col_titles = pivot_abs.get_flat_column_headers()
        row_titles = pivot_abs.get_flat_row_headers()
        assert col_titles == ["Furniture", "Office Supplies", "Technology", ""]
        assert row_titles == ["Central", "East", "South", "West", ""]

    def test_pivot_empty_data_with_grand_totals(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "profit sum": "SUM([profit])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category", MEASURE_NAME_TITLE],
            rows=[],
            measures=["sales sum", "profit sum"],
            measures_sorting_settings=[
                PivotMeasureSorting(
                    column=PivotMeasureSortingSettings(
                        header_values=[PivotHeaderValue(value=""), PivotHeaderValue(value="sales sum")],
                        direction=OrderDirection.desc,
                        role_spec=PivotHeaderRoleSpec(role=PivotHeaderRole.total),
                    )
                ),
                None,
            ],
            totals=[()],
            filters=[ds.find_field("category").filter(WhereClauseOperation.STARTSWITH, ["lol"])],
        )

        assert pivot_abs.get_compound_column_headers() == [("", "sales sum"), ("", "profit sum")]
        assert pivot_abs.get_compound_row_headers() == [()]
        assert pivot_abs.get_1d_mapper().get_value_count() == 2  # no values except grand totals

    @pytest.mark.parametrize("role", [PivotHeaderRole.total, PivotHeaderRole.data])
    def test_pivot_measure_sorting_with_total_and_empty_name_column(self, control_api, data_api, dataset_id, role):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
                "category without Furniture": "IF([category] = 'Furniture', '', [category])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category without Furniture"],
            rows=["region"],
            measures=["sales sum"],
            measures_sorting_settings=[
                PivotMeasureSorting(
                    column=PivotMeasureSortingSettings(
                        header_values=[PivotHeaderValue(value="")],
                        direction=OrderDirection.desc,
                        role_spec=PivotHeaderRoleSpec(role=role),
                    )
                )
            ],
            totals=[()],
        )
        col_titles = pivot_abs.get_flat_column_headers()
        assert col_titles == ["", "Office Supplies", "Technology", ""]  # first '' is from formula, second is from total

        def _get_value(value: Optional[tuple]) -> float:
            if value is None:
                return float("-inf")
            return float(value[0][0])

        column_idx = 0 if role == PivotHeaderRole.data else -1
        total_values = [_get_value(row["values"][column_idx]) for row in pivot_abs.resp_data["pivot_data"]["rows"]]
        last_value = total_values.pop()
        # pivot is sorted by values of '' column with selected role,
        # except for the last value, which is from totals row
        assert sorted(total_values, reverse=True) == total_values
        if role == PivotHeaderRole.total:  # total column: the last value is a grand total => is not None:
            assert last_value != float("-inf")
        else:  # data column: the last value is a column total => is None
            assert last_value == float("-inf")

    def test_basic_pivot_with_grand_total_and_simple_subtotals(self, control_api, data_api, dataset_id):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["region"],
            measures=["sales sum"],
            totals=[
                ("category",),  # region total by category
                ("region",),  # category total by region
                (),  # Grand Total
            ],
        )
        col_titles = pivot_abs.get_flat_column_headers()
        row_titles = pivot_abs.get_flat_row_headers()
        assert col_titles == ["Furniture", "Office Supplies", "Technology", ""]
        assert row_titles == ["Central", "East", "South", "West", ""]

    def test_pivot_with_totals_and_measure_filters(self, control_api, data_api, dataset_id):
        """
        Check that if measure filters are added to the request,
        then totals are disabled.
        """

        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset_id=dataset_id,
            formulas={
                "sales sum": "SUM([sales])",
            },
        )

        pivot_abs = check_pivot_response(
            dataset=ds,
            data_api=data_api,
            columns=["category"],
            rows=["region"],
            measures=["sales sum"],
            totals=[("category",), ("region",), ()],
            filters=[ds.find_field("sales sum").filter(WhereClauseOperation.GT, [10.0])],
            expected_notifications=[NotificationType.totals_removed_due_to_measure_filter],
        )
        col_titles = pivot_abs.get_flat_column_headers()
        row_titles = pivot_abs.get_flat_row_headers()
        assert col_titles == ["Furniture", "Office Supplies", "Technology"]  # No totals
        assert row_titles == ["Central", "East", "South", "West"]  # No totals

lib/dl_api_lib/dl_api_lib_tests/db/data_api/pivot/test_totals.py
datalens-tech/datalens-backend
from __future__ import annotations

import random
from typing import (
    Any,
    Sequence,
)

import attr

from dl_constants.enums import WhereClauseOperation


@attr.s
class AutoGeneratorSettings:
    dimensions: Sequence[str] = attr.ib(kw_only=True)
    dates: frozenset[str] = attr.ib(kw_only=True, factory=frozenset)
    filters: dict[str, tuple[WhereClauseOperation, list[str]]] = attr.ib(kw_only=True)
    dimension_cnts: Sequence[int] = attr.ib(kw_only=True)
    aggregations: Sequence[str] = attr.ib(kw_only=True)
    formula_cnts: Sequence[int] = attr.ib(kw_only=True)
    measure_base_expressions: Sequence[str] = attr.ib(kw_only=True)
    filter_probability: float = attr.ib(kw_only=True)
    bfb_probability: float = attr.ib(kw_only=True)
    lookup_probability: float = attr.ib(kw_only=True)


@attr.s(frozen=True)
class FormulaTemplate:
    template_str: str = attr.ib()
    placeholder: str = attr.ib(kw_only=True)

    def wrap(self, nested_formula: str) -> str:
        return self.template_str.format(**{self.placeholder: nested_formula})


@attr.s(frozen=True)
class FormulaRecursionState:
    agg_level: int = attr.ib(kw_only=True)
    effective_dims: list[str] = attr.ib(kw_only=True)
    remaining_dims: list[str] = attr.ib(kw_only=True)
    available_filters: frozenset[str] = attr.ib(kw_only=True, factory=frozenset)
    top_level: bool = attr.ib(kw_only=True, default=False)

    def clone(self, **kwargs: Any) -> FormulaRecursionState:
        return attr.evolve(self, **kwargs)

    def pop_agg_level(self) -> FormulaRecursionState:
        assert self.agg_level > 0
        return self.clone(top_level=False, agg_level=self.agg_level - 1)

    def pop_dimension(self) -> tuple[FormulaRecursionState, str]:
        assert self.remaining_dims
        chosen_dim = random.choice(self.remaining_dims)
        new_effective_dims = self.effective_dims + [chosen_dim]
        new_remaining_dims = [dim for dim in self.remaining_dims if dim != chosen_dim]
        new_state = self.clone(
            remaining_dims=new_remaining_dims,
            effective_dims=new_effective_dims,
        )
        return new_state, chosen_dim

    def pop_bfb_field(self) -> tuple[FormulaRecursionState, str]:
        bfb_field = next(iter(self.available_filters))
        new_available_filters = self.available_filters - frozenset((bfb_field,))
        return self.clone(available_filters=new_available_filters), bfb_field


@attr.s
class TestSettings:
    base_dimensions: Sequence[str] = attr.ib(kw_only=True)
    measure_formulas: Sequence[str] = attr.ib(kw_only=True)
    filters: dict[str, tuple[WhereClauseOperation, list[str]]] = attr.ib(kw_only=True, factory=dict)

    def serialize(self) -> dict:
        return {
            "base_dimensions": list(self.base_dimensions),
            "measure_formulas": list(self.measure_formulas),
            "filters": {name: {"op": op.name, "values": values} for name, (op, values) in self.filters.items()},
        }

    @classmethod
    def deserialize(cls, data: dict) -> TestSettings:
        return cls(
            base_dimensions=data["base_dimensions"],
            measure_formulas=data["measure_formulas"],
            filters={
                name: (WhereClauseOperation[params["op"]], params["values"]) for name, params in data["filters"].items()
            },
        )


@attr.s
class LODTestAutoGenerator:
    settings: AutoGeneratorSettings = attr.ib(kw_only=True)

    def _match_probability(self, value: float) -> bool:
        return random.random() <= value

    def _generate_formula_iteration_generic_func(
        self,
        recursion_state: FormulaRecursionState,
        func_name: str,
        add_args: Sequence[str] = (),
        agg_level: bool = False,
        lod: bool = False,
        bfb: bool = False,
        igdim: bool = False,
    ) -> tuple[FormulaTemplate, FormulaRecursionState]:
        top_level = recursion_state.top_level

        if agg_level:
            recursion_state = recursion_state.pop_agg_level()

        lod_str: str = ""
        if lod and not top_level:
            recursion_state, chosen_dim = recursion_state.pop_dimension()
            lod_str = f" INCLUDE [{chosen_dim}]"

        bfb_str: str = ""
        if bfb and recursion_state.available_filters:
            if self._match_probability(self.settings.bfb_probability):
                recursion_state, bfb_field = recursion_state.pop_bfb_field()
                bfb_str = f" BEFORE FILTER BY [{bfb_field}]"

        add_args_str = (", " + ", ".join(add_args)) if add_args else ""
        formula_tmpl = FormulaTemplate(
            f"{func_name}({{nested_formula}}{add_args_str}{lod_str}{bfb_str})",
            placeholder="nested_formula",
        )
        return formula_tmpl, recursion_state

    def _generate_formula_iteration_agg(
        self,
        recursion_state: FormulaRecursionState,
    ) -> tuple[FormulaTemplate, FormulaRecursionState]:
        return self._generate_formula_iteration_generic_func(
            recursion_state=recursion_state,
            func_name=random.choice(self.settings.aggregations),
            agg_level=True,
            lod=True,
            bfb=True,
        )

    def _generate_formula_iteration_lookup(
        self,
        recursion_state: FormulaRecursionState,
    ) -> tuple[FormulaTemplate, FormulaRecursionState]:
        date_dims = list(set(recursion_state.effective_dims) & self.settings.dates)
        assert date_dims
        return self._generate_formula_iteration_generic_func(
            recursion_state=recursion_state,
            func_name="AGO",
            add_args=[f"[{random.choice(date_dims)}]"],
            bfb=True,
        )

    def _generate_formula_iteration(
        self,
        recursion_state: FormulaRecursionState,
    ) -> tuple[FormulaTemplate, FormulaRecursionState]:
        date_dims = set(recursion_state.effective_dims) & self.settings.dates
        if date_dims and recursion_state.agg_level > 0 and self._match_probability(self.settings.lookup_probability):
            return self._generate_formula_iteration_lookup(recursion_state=recursion_state)

        return self._generate_formula_iteration_agg(recursion_state=recursion_state)

    def generate_formula(self, recursion_state: FormulaRecursionState) -> str:
        if recursion_state.agg_level == 0:
            return random.choice(self.settings.measure_base_expressions)

        formula_tmpl, child_recursion_state = self._generate_formula_iteration(recursion_state=recursion_state)
        nested_formula = self.generate_formula(recursion_state=child_recursion_state)
        formula = formula_tmpl.wrap(nested_formula)
        return formula

    def generate_test_settings(self) -> TestSettings:
        dimension_cnt = random.choice(self.settings.dimension_cnts)
        base_dimensions = random.sample(self.settings.dimensions, k=dimension_cnt)
        effective_dims = [dim for dim in self.settings.dimensions if dim in base_dimensions]
        remaining_dims = [dim for dim in self.settings.dimensions if dim not in base_dimensions]
        formula_cnt = random.choice(self.settings.formula_cnts)

        # Generate filters
        chosen_filters: dict[str, tuple[WhereClauseOperation, list[str]]] = {}
        for name, params in self.settings.filters.items():
            if random.random() <= self.settings.filter_probability:
                chosen_filters[name] = params

        recursion_state = FormulaRecursionState(
            effective_dims=effective_dims,
            remaining_dims=remaining_dims,
            available_filters=frozenset(chosen_filters.keys()),
            top_level=True,
            agg_level=random.randrange(1, len(remaining_dims) + 1),
        )

        measure_formulas = [self.generate_formula(recursion_state=recursion_state) for i in range(formula_cnt)]

        test_settings = TestSettings(
            base_dimensions=base_dimensions,
            measure_formulas=measure_formulas,
            filters=chosen_filters,
        )
        return test_settings

    def generate_setting_list(self, test_cnt: int) -> list[TestSettings]:
        return [self.generate_test_settings() for i in range(test_cnt)]

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/complex_queries/generation/generator.py
datalens-tech/datalens-backend
from __future__ import annotations

from http import HTTPStatus
import json
from typing import Optional

import attr
import pytest

from dl_api_client.dsmaker.api.data_api import SyncHttpDataApiV2
from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_tests.db.data_api.result.complex_queries.generation.generator import TestSettings


class Error400(Exception):
    pass


@attr.s
class PreGeneratedLODTestRunner:
    dataset_id: str = attr.ib(kw_only=True)
    control_api: SyncHttpDatasetApiV1 = attr.ib(kw_only=True)
    data_api: SyncHttpDataApiV2 = attr.ib(kw_only=True)

    def get_measure_data(
        self,
        ds: Dataset,
        test_settings: TestSettings,
        measures: list[str],
    ) -> dict[str, list[float]]:
        result_resp = self.data_api.get_result(
            dataset=ds,
            fields=[
                *[ds.find_field(title=name) for name in test_settings.base_dimensions],
                *[ds.find_field(title=name) for name in measures],
            ],
            filters=[
                ds.find_field(title=name).filter(op, values) for name, (op, values) in test_settings.filters.items()
            ],
            order_by=[
                *[ds.find_field(title=name) for name in test_settings.base_dimensions],
            ],
            fail_ok=True,
        )

        if result_resp.status_code == HTTPStatus.BAD_REQUEST:
            raise Error400(result_resp.json)

        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert data_rows

        def float_or_none(value: Optional[str]) -> Optional[float]:
            if value is None:
                return None
            return float(value)

        result: dict[str, list[float]] = {}
        dim_cnt = len(test_settings.base_dimensions)
        for measure_idx, name in enumerate(measures):
            result[name] = [float_or_none(row[measure_idx + dim_cnt]) for row in data_rows]

        return result

    def run_test(self, test_settings: TestSettings, ignore_400_error: bool = False) -> None:
        ds = self.control_api.load_dataset(Dataset(id=self.dataset_id)).dataset
        measure_fields = [
            (f"Measure {measure_idx}", measure_formula)
            for measure_idx, measure_formula in enumerate(test_settings.measure_formulas)
        ]
        for measure_name, measure_formula in measure_fields:
            ds.result_schema[measure_name] = ds.field(formula=measure_formula)

        try:
            all_measures = self.get_measure_data(
                ds=ds, test_settings=test_settings, measures=[name for name, _ in measure_fields]
            )
            for measure_name, measure_formula in measure_fields:
                one_measure = self.get_measure_data(ds=ds, test_settings=test_settings, measures=[measure_name])
                assert pytest.approx(all_measures[measure_name]) == one_measure[measure_name]
        except Error400:
            if not ignore_400_error:
                raise
            print("Ignoring Error400")

    def run_test_list(self, test_list: list[TestSettings], ignore_400_error: bool) -> None:
        for test_idx, test_settings in enumerate(test_list):
            print(f"\nTest # {test_idx}\n{json.dumps(test_settings.serialize())}\n")
            self.run_test(test_settings, ignore_400_error=ignore_400_error)

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/complex_queries/generation/runner.py
datalens-tech/datalens-backend
from __future__ import annotations

from http import HTTPStatus

import pytest

from dl_api_client.dsmaker.shortcuts.dataset import add_formulas_to_dataset
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_tests.db.base import DefaultApiTestBase


class TestCompengCornerCases(DefaultApiTestBase):
    def test_zero_division_in_compeng(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Measure Reg Div": "1 / (SUM(SUM([sales]) AMONG) * 0)",
                "Measure Int Div": "DIV(1, SUM(COUNT([sales]) AMONG) * 0)",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="Measure Reg Div"),
                ds.find_field(title="Measure Int Div"),
            ],
            fail_ok=True,
        )

        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 1
        row = data_rows[0]
        assert row[0] is None
        assert row[1] is None

    def test_integer_division_in_compeng(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Measure": "COUNT([sales])",
                "Fraction": 'DB_CAST([Measure], "Int64") / DB_CAST(SUM([Measure] TOTAL), "integer")',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="Fraction"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) > 1

        fraction_sum = sum(float(row[1]) for row in data_rows)
        assert pytest.approx(fraction_sum) == 1.0

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/complex_queries/test_compeng_corner_cases.py
datalens-tech/datalens-backend
from collections import defaultdict
import datetime
from http import HTTPStatus
from typing import (
    Any,
    Iterable,
)

import pytest

from dl_api_client.dsmaker.primitives import Dataset
from dl_api_client.dsmaker.shortcuts.dataset import add_formulas_to_dataset
from dl_api_client.dsmaker.shortcuts.result_data import (
    get_data_rows,
    get_regular_result_data,
)
from dl_api_lib_testing.connector.complex_queries import DefaultBasicExtAggregationTestSuite
from dl_api_lib_testing.helpers.data_source import data_source_settings_from_table
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import (
    UserDataType,
    WhereClauseOperation,
)
from dl_core_testing.database import (
    C,
    make_table,
)
from dl_query_processing.compilation.primitives import CompiledMultiQueryBase
from dl_query_processing.translation.multi_level_translator import MultiLevelQueryTranslator


class TestBasicExtendedAggregations(DefaultApiTestBase, DefaultBasicExtAggregationTestSuite):
    # Add to the default set of tests

    def test_fixed_same_as_exclude(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "sales sum": "SUM([sales])",
                "sales sum fx city": "SUM([sales] FIXED [city])",
                "sales sum exc category": "SUM([sales] EXCLUDE [category])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="category"),
                ds.find_field(title="sales sum"),
                ds.find_field(title="sales sum fx city"),
                ds.find_field(title="sales sum exc category"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        for row in data_rows:
            assert float(row[3]) == float(row[4])

    def test_nested_lod(self, control_api, data_api, saved_dataset):
        """
        Check nested LODs.

        Calculate the per-city average per-category sum of sales
        (first the dimensions expand to allow `category` in, and then collapse
        to the request default of just `city`)
        """
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "sales sum": "SUM([sales])",
                "sales sum inc category": "SUM([sales] INCLUDE [category])",
                "nested LOD AVG": "AVG([sales sum inc category])",
                "nested LOD SUM": "SUM([sales sum inc category])",
            },
        )

        def get_expected_data() -> tuple[dict[str, float], dict[str, float]]:
            # 1. Get ungrouped data
            ungrouped_result_resp = data_api.get_result(
                dataset=ds,
                fields=[
                    ds.find_field(title="city"),
                    ds.find_field(title="category"),
                    ds.find_field(title="sales"),
                ],
                disable_group_by=True,
                fail_ok=True,
            )
            assert ungrouped_result_resp.status_code == HTTPStatus.OK, ungrouped_result_resp.json

            # 2. Group data by city and category and calculate the sum of sales per group
            data_rows = get_data_rows(ungrouped_result_resp)
            sum_by_city_and_cat: dict[tuple[str, str], float] = defaultdict(lambda: 0)
            for row in data_rows:
                sum_by_city_and_cat[(row[0], row[1])] += float(row[2])
            sum_list_by_city: dict[str, list[float]] = defaultdict(list)
            for (city, cat), sum_value in sum_by_city_and_cat.items():
                sum_list_by_city[city].append(sum_value)

            # 3. Group by city and calculate average of sales sums
            avg_sum_by_city = {city: sum(sum_list) / len(sum_list) for city, sum_list in sum_list_by_city.items()}
            sum_sum_by_city = {city: sum(sum_list) for city, sum_list in sum_list_by_city.items()}
            return avg_sum_by_city, sum_sum_by_city

        # Calculate expected results
        exp_avg_by_city, exp_sum_by_city = get_expected_data()
        # Get actual results
        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                # ds.find_field(title='sales sum'),
                ds.find_field(title="nested LOD AVG"),
                ds.find_field(title="nested LOD SUM"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json

        # Compare values for each `city`
        data_rows = get_data_rows(result_resp)
        for row in data_rows:
            city = row[0]
            city_sum_avg = float(row[1])
            # Use approx here because DB and Python float calculations do not match exactly
            assert city_sum_avg == pytest.approx(exp_avg_by_city[city])
            city_sum_sum = float(row[2])
            assert city_sum_sum == pytest.approx(exp_sum_by_city[city])

    def test_lod_nested_zero_dim_aggregation_with_filter(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "sales sum": "SUM([sales])",
                "sales sum total": "SUM(SUM([sales] INCLUDE [city]))",
            },
        )

        def get_total_value() -> float:
            total_sum_result = data_api.get_result(
                dataset=ds,
                fields=[ds.find_field(title="sales sum")],
                filters=[
                    ds.find_field(title="category").filter(
                        WhereClauseOperation.EQ,
                        values=["Office Supplies"],
                    ),
                ],
            )
            data_rows = get_data_rows(total_sum_result)
            return float(data_rows[0][0])

        expected_total_value = get_total_value()

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="sales sum total"),
            ],
            filters=[
                ds.find_field(title="category").filter(
                    WhereClauseOperation.EQ,
                    values=["Office Supplies"],
                ),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        for row_idx, row in enumerate(data_rows):
            assert float(row[0]) == expected_total_value, f"total sum doesn't match expected number in row {row_idx}"

    def test_double_agg_ratio(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "sales sum": "SUM([sales])",
                "simple ratio": "SUM([sales]) / SUM([sales] FIXED)",
                "double agg ratio": "SUM([sales]) / SUM(SUM([sales]) FIXED)",
            },
        )

        def get_total_value() -> float:
            total_sum_result = data_api.get_result(
                dataset=ds,
                fields=[ds.find_field(title="sales sum")],
            )
            data_rows = get_data_rows(total_sum_result)
            return float(data_rows[0][0])

        expected_total_value = get_total_value()

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="sales sum"),
                ds.find_field(title="simple ratio"),
                ds.find_field(title="double agg ratio"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        for row_idx, row in enumerate(data_rows):
            expected_value = float(row[1]) / expected_total_value
            assert (
                float(row[3]) == float(row[2]) == expected_value
            ), f"total sum doesn't match expected number in row {row_idx}"

    def test_nested_zero_dim_aggregation(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "sales sum": "SUM([sales])",
                "city avg sales sum": "AVG(SUM([sales] INCLUDE [city]))",
            },
        )

        def get_expected_value() -> float:
            # 1. Get per-city data
            result_resp = data_api.get_result(
                dataset=ds,
                fields=[
                    ds.find_field(title="city"),
                    ds.find_field(title="sales sum"),
                ],
                fail_ok=True,
            )
            assert result_resp.status_code == HTTPStatus.OK, result_resp.json
            data_rows = get_data_rows(result_resp)
            # Calculate the average
            return sum(float(row[1]) for row in data_rows) / len(data_rows)

        expected_avg_value = get_expected_value()

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city avg sales sum"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 1
        assert float(data_rows[0][0]) == pytest.approx(expected_avg_value)

    def test_lod_with_bfb(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "sales sum": "SUM([sales])",
                "sales sum fx Date": "SUM([sales] FIXED [order_date])",
                "sales sum fx Date BFB": "SUM([sales] FIXED [order_date] BEFORE FILTER BY [category])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="category"),
                ds.find_field(title="sales sum"),
                ds.find_field(title="sales sum fx Date"),
                ds.find_field(title="sales sum fx Date BFB"),
            ],
            filters=[
                ds.find_field(title="category").filter(
                    WhereClauseOperation.EQ,
                    values=["Office Supplies"],
                ),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        sales_by_date = defaultdict(lambda: 0)
        lod_sales_by_date = {}
        case_1_cnt = 0
        case_2_cnt = 0
        for row in data_rows:
            sales_by_date[row[0]] += float(row[2])
            lod_sales_by_date[row[0]] = float(row[4])
            # BFB sum is equal to non-BFB sum only when there is exactly one category value,
            # so it also equals the regular per-row sum
            if float(row[4]) == float(row[3]):
                assert float(row[4]) == float(row[3]) == float(row[2])
                case_1_cnt += 1
            else:  # Otherwise BFB is greater than non-BFB
                assert float(row[4]) > float(row[3])
                case_2_cnt += 1

        assert lod_sales_by_date == pytest.approx(lod_sales_by_date)
        assert case_1_cnt > 0
        assert case_2_cnt > 0

    def test_workaround_for_inconsistent_agg(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "My Field": "SUM(SUM([sales] INCLUDE [Order ID]) / SUM([sales] FIXED))",
            },
        )
        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="My Field"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 1
        assert pytest.approx(float(data_rows[0][0])) == 1

    def test_lod_nested_multiple_times(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                # Plain and simple row count
                "Agg 1": "COUNT([sales])",
                # All the others are pretty much the same thing, but with multiple aggregation levels
                "Agg 2": "SUM(COUNT([sales] INCLUDE [city]))",
                "Agg 3": "SUM(SUM(COUNT([sales] INCLUDE [city]) INCLUDE [category]))",
                "Agg 4": "SUM(SUM(SUM(COUNT([sales] INCLUDE [city]) INCLUDE [category]) INCLUDE [order_date]))",
                "Agg 5": (
                    "SUM(SUM(SUM(SUM(COUNT([sales] INCLUDE [city]) INCLUDE [category]) "
                    "INCLUDE [order_date]) INCLUDE [region]))"
                ),
            },
        )

        def get_single_row_data(field_names: Iterable[str]) -> tuple[int, ...]:
            result_resp = data_api.get_result(
                dataset=ds,
                fail_ok=True,
                fields=[ds.find_field(title=field_name) for field_name in field_names],
            )
            assert result_resp.status_code == HTTPStatus.OK, result_resp.json
            data_rows = get_data_rows(result_resp)
            assert len(data_rows) == 1, "There must be exactly 1 row of data"
            return tuple(int(item) for item in data_rows[0])

        value_1 = get_single_row_data(["Agg 1"])[0]
        value_2 = get_single_row_data(["Agg 2"])[0]
        assert value_2 == value_1
        value_3 = get_single_row_data(["Agg 3"])[0]
        assert value_3 == value_1
        value_4 = get_single_row_data(["Agg 4"])[0]
        assert value_4 == value_1
        value_5 = get_single_row_data(["Agg 4"])[0]
        assert value_5 == value_1

        # Check the two at a time with the order of aggregation increasing by 1
        simultaneous_values = get_single_row_data(["Agg 1", "Agg 2"])
        assert len(set(simultaneous_values)) == 1
        assert next(iter(simultaneous_values)) == value_1

        simultaneous_values = get_single_row_data(["Agg 2", "Agg 3"])
        assert len(set(simultaneous_values)) == 1
        assert next(iter(simultaneous_values)) == value_1

        simultaneous_values = get_single_row_data(["Agg 3", "Agg 4"])
        assert len(set(simultaneous_values)) == 1
        assert next(iter(simultaneous_values)) == value_1

        simultaneous_values = get_single_row_data(["Agg 4", "Agg 5"])
        assert len(set(simultaneous_values)) == 1
        assert next(iter(simultaneous_values)) == value_1

        # Check with diffs 2, 3 and 4
        simultaneous_values = get_single_row_data(["Agg 1", "Agg 3"])
        assert len(set(simultaneous_values)) == 1
        assert next(iter(simultaneous_values)) == value_1

        simultaneous_values = get_single_row_data(["Agg 1", "Agg 4"])
        assert len(set(simultaneous_values)) == 1
        assert next(iter(simultaneous_values)) == value_1

        simultaneous_values = get_single_row_data(["Agg 1", "Agg 5"])
        assert len(set(simultaneous_values)) == 1
        assert next(iter(simultaneous_values)) == value_1

        # The killer query - do them all at once!
        simultaneous_values = get_single_row_data(["Agg 1", "Agg 2", "Agg 3", "Agg 4", "Agg 5"])
        assert len(set(simultaneous_values)) == 1
        assert next(iter(simultaneous_values)) == value_1

    def test_two_same_dim_lods_in_different_subqueries(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "sales sum": "SUM([sales])",
                "same dim LOD": "SUM([sales] FIXED [city], [category])",
                "by city": "SUM(SUM([same dim LOD] FIXED [city]))",
                "by category": "SUM(SUM([same dim LOD] FIXED [category]))",
            },
        )

        def get_total_value() -> float:
            total_sum_result = data_api.get_result(
                dataset=ds,
                fields=[ds.find_field(title="sales sum")],
            )
            data_rows = get_data_rows(total_sum_result)
            return float(data_rows[0][0])

        total_value = get_total_value()

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="by city"),
                ds.find_field(title="by category"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 1
        by_city, by_category = float(data_rows[0][0]), float(data_rows[0][1])
        assert by_city == pytest.approx(by_category) == pytest.approx(total_value)

    def test_two_total_sums_with_different_nested_includes(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "By City": "SUM(SUM([sales] INCLUDE [city]))",
                "By Category": "SUM(SUM([sales] INCLUDE [category]))",
            },
        )

        def get_total_value() -> float:
            total_sum_result = data_api.get_result(
                dataset=ds,
                fields=[ds.find_field(title="Sales Sum")],
            )
            data_rows = get_data_rows(total_sum_result)
            return float(data_rows[0][0])

        total_value = get_total_value()

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="By City"),
                ds.find_field(title="By Category"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 1
        by_city, by_category = float(data_rows[0][0]), float(data_rows[0][1])
        assert by_city == pytest.approx(by_category) == pytest.approx(total_value)

    def test_agg_with_lod_over_window_function(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Lod Over Win 1": "MIN(RANK_UNIQUE(SUM([sales] INCLUDE [city]) TOTAL) FIXED)",
                "Lod Over Win 2": "MAX(RANK_UNIQUE(SUM([sales] INCLUDE [city]) TOTAL) FIXED)",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="Lod Over Win 1"),
                ds.find_field(title="Lod Over Win 2"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) > 10  # There should be a lot of them
        for row in data_rows:
            assert int(row[1]) == 1
            assert int(row[2]) == len(data_rows)

    def test_lod_compatibility_error(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            exp_status=HTTPStatus.BAD_REQUEST,
            formulas={
                "Invalid Field": "SUM(AVG([sales] INCLUDE [city]) - AVG([sales] INCLUDE [category]))",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Invalid Field"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST, result_resp.json
        assert result_resp.bi_status_code == "ERR.DS_API.FORMULA.VALIDATION.LOD.INCOMPATIBLE_DIMENSIONS"

    def test_dimension_with_single_and_double_agg(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Single Agg": "SUM([sales])",
                "Double Agg": "SUM(SUM([sales] INCLUDE [city]))",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Single Agg"),
            ],
        )
        data_rows = get_data_rows(result_resp)
        expected_by_date = {row[0]: float(row[1]) for row in data_rows}

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Single Agg"),
                ds.find_field(title="Double Agg"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == len(expected_by_date)
        for row in data_rows:
            date_str, single_agg_value, double_agg_value = row[0], float(row[1]), float(row[2])
            assert single_agg_value == pytest.approx(double_agg_value) == pytest.approx(expected_by_date[date_str])

    def test_lod_with_ago(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Lod With Ago": """
                AVG(
                    SUM(  /* Doesn't really aggregate anything
                           * because it has the same dims as the nested SUMs */
                        SUM([sales]) - ZN(AGO(SUM([sales]), [order_date], "day"))
                        INCLUDE [order_date]
                    )
                )
            """,
            },
        )

        def get_expected_value() -> float:
            result_resp = data_api.get_result(
                dataset=ds,
                fields=[
                    ds.find_field(title="order_date"),
                    ds.find_field(title="Sales Sum"),
                ],
            )
            data_rows = get_data_rows(result_resp)
            sales_by_date = {row[0]: float(row[1]) for row in data_rows}

            total_sum = 0.0
            for date_str in sorted(sales_by_date):
                ago_date_str = str(datetime.date.fromisoformat(date_str) - datetime.timedelta(days=1))
                sales_value = sales_by_date[date_str]
                ago_sales_value = sales_by_date.get(ago_date_str, 0)
                total_sum += sales_value - ago_sales_value

            return total_sum / len(sales_by_date)

        expected_value = get_expected_value()

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="Lod With Ago"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 1
        actual_value = float(data_rows[0][0])
        assert actual_value == pytest.approx(expected_value)

    def test_toplevel_lod_extra_dimension_error(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Extra Dim Field": "AVG([sales] INCLUDE [category])",
                "Missing Dim Field": "AVG([sales] FIXED)",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Extra Dim Field"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST, result_resp.json
        assert result_resp.bi_status_code == "ERR.DS_API.FORMULA.VALIDATION.LOD.INVALID_TOPLEVEL_DIMENSIONS"

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Missing Dim Field"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json

    def test_double_agg_no_lod(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Double Agg": "COUNT(COUNT([sales]))",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="Double Agg"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        assert len(data_rows) == 1
        assert int(data_rows[0][0]) == 1

    def test_bfb_no_lod(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Bfb No Lod": "SUM([sales] BEFORE FILTER BY [category])",
            },
        )

        result_resp = data_api.get_result(dataset=ds, fields=[ds.find_field(title="Sales Sum")])
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 1
        expected_total_value = float(data_rows[0][0])

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="Bfb No Lod"),
            ],
            filters=[
                ds.find_field(title="category").filter(op=WhereClauseOperation.EQ, values=["Office Supplies"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        assert len(data_rows) == 1
        assert float(data_rows[0][0]) == expected_total_value

    def test_regular_agg_with_bfb_agg(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Agg Bfb": "SUM([sales] EXCLUDE [category] BEFORE FILTER BY [category])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="region"),
                ds.find_field(title="category"),
                ds.find_field(title="Sales Sum"),
            ],
            filters=[
                ds.find_field(title="category").filter(WhereClauseOperation.EQ, ["Office Supplies"]),
            ],
        )
        data_rows = get_data_rows(result_resp)
        expected_by_region = {row[0]: float(row[2]) for row in data_rows}

        # Same request, but with `Agg Bfb`
        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="region"),
                ds.find_field(title="category"),
                ds.find_field(title="Sales Sum"),
                ds.find_field(title="Agg Bfb"),
            ],
            filters=[
                ds.find_field(title="category").filter(WhereClauseOperation.EQ, ["Office Supplies"]),
            ],
        )
        data_rows = get_data_rows(result_resp)
        actual_by_region = {row[0]: float(row[2]) for row in data_rows}

        assert actual_by_region == expected_by_region

    def test_double_aggregation_optimization(self, control_api, data_api, saved_dataset, monkeypatch):
        """
        Test double aggregation optimizations in formulas like `SUM(SUM([sales]))` (-> `SUM([sales])`)
        """
        data_container: dict[str, Any] = {}
        dataset_id = saved_dataset.id

        def _log_query_complexity_stats(self, compiled_multi_query: CompiledMultiQueryBase) -> None:
            data_container.update(
                {
                    "query_count": compiled_multi_query.query_count(),
                }
            )

        # Capture query complexity info
        monkeypatch.setattr(MultiLevelQueryTranslator, "_log_query_complexity_stats", _log_query_complexity_stats)

        def check_agg(first_agg_name: str, second_agg_name: str) -> None:
            ds = add_formulas_to_dataset(
                api_v1=control_api,
                dataset_id=dataset_id,
                formulas={
                    f"{first_agg_name} {second_agg_name} Opt": f"{second_agg_name}({first_agg_name}([sales]))",
                    f"{first_agg_name} {second_agg_name}": f"{second_agg_name}({first_agg_name}([sales])+0)",
                },
            )

            data = get_regular_result_data(
                ds,
                data_api,
                field_names=["category", f"{first_agg_name} {second_agg_name}"],
            )[f"{first_agg_name} {second_agg_name}"]
            assert data_container["query_count"] >= 3
            data_opt = get_regular_result_data(
                ds,
                data_api,
                field_names=["category", f"{first_agg_name} {second_agg_name} Opt"],
            )[f"{first_agg_name} {second_agg_name} Opt"]
            assert data_container["query_count"] == 1
            assert sorted([float(val) for val in data_opt]) == pytest.approx(
                sorted([float(val) for val in data])
            ), f"Got different values for {first_agg_name}({second_agg_name}(...))"

        for first_agg_name in ("avg", "count"):
            for second_agg_name in ("sum", "any", "max", "min", "count", "countd"):
                # no optimization for AVG
                check_agg(first_agg_name.upper(), second_agg_name.upper())

    def test_lod_fixed_markup(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Bold city": "BOLD([city])",
                "Sales Sum fx": "AVG(SUM([sales] FIXED))",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="Bold city"),
                ds.find_field(title="Sales Sum fx"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json

    def test_bug_bi_3425_deeply_nested_bfb(self, control_api, data_api, db, saved_connection_id):
        raw_data = [
            {"id": 10, "city": "New York", "category": "Office Supplies", "sales": 1},
            {"id": 11, "city": "New York", "category": "Office Supplies", "sales": 10},
            {"id": 12, "city": "New York", "category": "Furniture", "sales": 100},
            {"id": 13, "city": "New York", "category": "Furniture", "sales": 1},
            {"id": 14, "city": "New Rochelle", "category": "Office Supplies", "sales": 10000},
            {"id": 15, "city": "New Rochelle", "category": "Office Supplies", "sales": 100000},
            {"id": 16, "city": "New Rochelle", "category": "Furniture", "sales": 10000},
            {"id": 17, "city": "New Rochelle", "category": "Furniture", "sales": 10000000},
            {"id": 18, "city": "Detroit", "category": "Office Supplies", "sales": 1},
            {"id": 19, "city": "Detroit", "category": "Office Supplies", "sales": 100},
            {"id": 20, "city": "Detroit", "category": "Furniture", "sales": 10000},
            {"id": 21, "city": "Detroit", "category": "Furniture", "sales": 1000000},
        ]
        columns = [
            C("id", UserDataType.integer, vg=lambda rn, **kwargs: raw_data[rn]["id"]),
            C("city", UserDataType.string, vg=lambda rn, **kwargs: raw_data[rn]["city"]),
            C("category", UserDataType.string, vg=lambda rn, **kwargs: raw_data[rn]["category"]),
            C("sales", UserDataType.integer, vg=lambda rn, **kwargs: raw_data[rn]["sales"]),
        ]
        db_table = make_table(db, columns=columns, rows=len(raw_data))

        ds = Dataset()
        ds.sources["source_1"] = ds.source(
            connection_id=saved_connection_id, **data_source_settings_from_table(db_table)
        )
        ds.source_avatars["avatar_1"] = ds.sources["source_1"].avatar()

        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=ds,
            formulas={
                "city_1st": 'GET_ITEM(SPLIT([city], " "), 1)',
                "sales_sum_bfb": "SUM(SUM([sales]) EXCLUDE [city] BEFORE FILTER BY [city_1st])",
                "sales_sum_if": 'SUM_IF([sales], [city_1st] = "New")',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="sales_sum_if"),
                ds.find_field(title="sales_sum_bfb"),
            ],
            order_by=[
                ds.find_field(title="city"),
            ],
            filters=[
                ds.find_field(title="city_1st").filter(WhereClauseOperation.EQ, ["New"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows_1 = get_data_rows(result_resp)
        data_rows_1_stripped = [row[:2] for row in data_rows_1]  # Strip off `sales_sum_bfb` values

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="sales_sum_if"),
                # This time without `sales_sum_bfb`
            ],
            order_by=[
                ds.find_field(title="city"),
            ],
            filters=[
                ds.find_field(title="city_1st").filter(WhereClauseOperation.EQ, ["New"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows_2 = get_data_rows(result_resp)

        assert data_rows_1_stripped == data_rows_2

    def test_lod_only_in_filter(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Daily Profit Sum": "AVG(SUM([sales] INCLUDE [order_date]))",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="Sales Sum"),
            ],
            filters=[
                ds.find_field(title="Daily Profit Sum").filter(WhereClauseOperation.GT, ["10.0"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) > 0

    def test_lod_in_filter_and_select(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Daily Profit Sum": "AVG(SUM([sales] INCLUDE [order_date]))",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="Daily Profit Sum"),
            ],
            filters=[
                ds.find_field(title="Daily Profit Sum").filter(WhereClauseOperation.GT, ["10.0"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) > 0

    def test_replace_original_dim_with_another(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Measure": "AVG(SUM([sales] FIXED [city]))",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="Measure"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 3  # There are 3 categories

    def test_bi_4534_inconsistent_aggregation(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sum": "SUM([sales])",
                "Measure": "SUM([Sum]) - SUM(AGO([Sum], [order_date]))",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Measure"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK

    def test_bi_4652_measure_filter_with_total_in_select(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Measure": "SUM([sales])",
                "Total Measure": "SUM([sales] FIXED)",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="Total Measure"),
            ],
            filters=[
                ds.find_field(title="Measure").filter(WhereClauseOperation.GT, ["10.0"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        dim_values = [row[0] for row in data_rows]
        assert len(dim_values) == len(set(dim_values)), "Dimension values are not unique"

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/complex_queries/test_ext_agg_basic.py
datalens-tech/datalens-backend
from __future__ import annotations

from http import HTTPStatus

from dl_api_client.dsmaker.shortcuts.dataset import add_formulas_to_dataset
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_tests.db.base import DefaultApiTestBase


class TestExtendedAggregationCornerCases(DefaultApiTestBase):
    def test_lod_with_const_dim(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Const Dim": '"something"',
                "Measure": "SUM([sales] EXCLUDE [category])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="Const Dim"),
                ds.find_field(title="Measure"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        assert len(data_rows) == 3
        assert len(set(row[2] for row in data_rows)) == 1  # They should all be the same

    def test_lod_with_date_dimension(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Date Dim": 'DB_CAST([order_date], "Date") + 1',
                "Measure": "SUM([sales] EXCLUDE [category])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="Date Dim"),
                ds.find_field(title="Measure"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        assert len(data_rows) > 1

    def test_duplicate_main_dimensions(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Category 2": "[category]",
                "Measure": "AVG(SUM([sales] INCLUDE [city]))",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="Category 2"),
                ds.find_field(title="Measure"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 3  # There are 3 categories
        for row in data_rows:
            assert row[0] == row[1]

    def test_duplicate_lod_dimensions(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "City 2": "[city]",
                "Measure Double Dim": "AVG(SUM([sales] INCLUDE [city], [City 2]))",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="Measure Double Dim"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 3  # There are 3 categories

    def test_lod_include_measure(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Measure": "SUM(SUM([sales] INCLUDE SUM([profit])))",
            },
            exp_status=HTTPStatus.BAD_REQUEST,
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="Measure"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/complex_queries/test_ext_agg_corner_cases.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_api_lib_tests.db.data_api.result.complex_queries.generation.generator import (
    AutoGeneratorSettings,
    LODTestAutoGenerator,
    TestSettings,
)
from dl_api_lib_tests.db.data_api.result.complex_queries.generation.runner import PreGeneratedLODTestRunner
from dl_constants.enums import WhereClauseOperation


GENERATED_TESTS = [
    {
        "base_dimensions": ["category"],
        "measure_formulas": [
            "SUM(AVG(AVG(SUM(AVG(SUM(AVG([sales] INCLUDE [region]) INCLUDE [ship_date]) INCLUDE [city]) INCLUDE [sub_category]) INCLUDE [order_date]) INCLUDE [ship_mode] BEFORE FILTER BY [category]))",  # noqa
            "AVG(SUM(AVG(SUM(SUM([sales] INCLUDE [region]) INCLUDE [city]) INCLUDE [sub_category]) INCLUDE [ship_mode]))",
        ],
        "filters": {"category": {"op": "EQ", "values": ["Furniture"]}},
    },
    {
        "base_dimensions": ["sub_category", "city"],
        "measure_formulas": [
            "SUM(SUM([sales] INCLUDE [region]) BEFORE FILTER BY [profit])",
            "AVG(SUM(AVG(AVG([sales] INCLUDE [region]) INCLUDE [order_date]) INCLUDE [ship_mode]))",
            "SUM([sales])",
            "AVG(SUM(AVG([sales] INCLUDE [order_date]) INCLUDE [category]) BEFORE FILTER BY [profit])",
        ],
        "filters": {"profit": {"op": "GT", "values": ["1.0"]}},
    },
    {
        "base_dimensions": ["region"],
        "measure_formulas": [
            "SUM(SUM(SUM(SUM([sales] INCLUDE [category]) INCLUDE [city] BEFORE FILTER BY [profit]) INCLUDE [ship_date]))",
            "AVG([sales])",
        ],
        "filters": {"profit": {"op": "GT", "values": ["1.0"]}},
    },
    {
        "base_dimensions": ["region"],
        "measure_formulas": [
            "AVG(AVG(AVG(AVG([sales] INCLUDE [category]) INCLUDE [ship_mode]) INCLUDE [sub_category]))",
            "SUM(AVG(AVG(SUM(AVG(SUM(SUM([sales] INCLUDE [order_date]) INCLUDE [city]) INCLUDE [sub_category]) INCLUDE [ship_date]) INCLUDE [category]) INCLUDE [ship_mode] BEFORE FILTER BY [order_date]))",  # noqa
            "SUM(AVG([sales] INCLUDE [city]) BEFORE FILTER BY [order_date])",
        ],
        "filters": {"order_date": {"op": "GT", "values": ["2014-03-01"]}},
    },
    {
        "base_dimensions": ["category"],
        "measure_formulas": [
            "SUM(AVG([sales] INCLUDE [ship_mode] BEFORE FILTER BY [category]))",
            "SUM(SUM(SUM(SUM(SUM(AVG(SUM([sales] INCLUDE [ship_mode] BEFORE FILTER BY [category]) INCLUDE [region]) INCLUDE [order_date]) INCLUDE [city]) INCLUDE [ship_date]) INCLUDE [sub_category]))",  # noqa
            "AVG(AVG(SUM([sales] INCLUDE [ship_mode]) INCLUDE [region] BEFORE FILTER BY [category]))",
        ],
        "filters": {"category": {"op": "EQ", "values": ["Furniture"]}},
    },
    {
        "base_dimensions": ["region"],
        "measure_formulas": [
            "SUM(SUM(AVG(SUM(SUM(AVG(AVG([sales] INCLUDE [order_date]) INCLUDE [sub_category]) INCLUDE [city]) INCLUDE [ship_mode]) INCLUDE [ship_date]) INCLUDE [category]) BEFORE FILTER BY [category])",  # noqa
            "SUM(AVG(AVG(AVG(SUM(AVG([sales] INCLUDE [ship_mode]) INCLUDE [city]) INCLUDE [sub_category]) INCLUDE [ship_date]) INCLUDE [category]))",
            "SUM(AVG(AVG(SUM(SUM([sales] INCLUDE [ship_mode]) INCLUDE [city] BEFORE FILTER BY [category]) INCLUDE [ship_date]) INCLUDE [category]))",
            "AVG(AVG(SUM(AVG([sales] INCLUDE [ship_date]) INCLUDE [ship_mode] BEFORE FILTER BY [category]) INCLUDE [city]))",
        ],
        "filters": {"category": {"op": "EQ", "values": ["Furniture"]}},
    },
    {
        "base_dimensions": ["region"],
        "measure_formulas": [
            "AVG(SUM([sales] INCLUDE [city]))",
            "SUM(AVG(AVG(SUM(SUM(AVG([sales] INCLUDE [order_date]) INCLUDE [ship_date] BEFORE FILTER BY [category]) INCLUDE [city]) INCLUDE [category]) INCLUDE [sub_category]))",
            "AVG(AVG(SUM(AVG(SUM([sales] INCLUDE [order_date]) INCLUDE [ship_date]) INCLUDE [category]) INCLUDE [ship_mode]) BEFORE FILTER BY [category])",
            "SUM(SUM([sales] INCLUDE [order_date]))",
        ],
        "filters": {"category": {"op": "EQ", "values": ["Furniture"]}},
    },
]


class TestPreGeneratedComplexQueryTests(DefaultApiTestBase):
    @pytest.mark.parametrize("raw_test_settings", GENERATED_TESTS)
    def test_pre_generated(self, control_api, data_api, dataset_id, raw_test_settings):
        test_runner = PreGeneratedLODTestRunner(control_api=control_api, data_api=data_api, dataset_id=dataset_id)
        test_runner.run_test(
            test_settings=TestSettings.deserialize(raw_test_settings),
        )


class TestNewGeneratedComplexQueryTests(DefaultApiTestBase):
    @pytest.mark.skip  # Should only be used for generating new tests
    def test_new_auto_generated(self, control_api, data_api, dataset_id):
        autogen_settings = AutoGeneratorSettings(
            dimensions=(
                "category",
                "city",
                "order_date",
                "region",
                "ship_date",
                "ship_mode",
                "sub_category",
            ),
            dates=frozenset(("order_date", "ship_date")),
            filters={
                "city": (WhereClauseOperation.STARTSWITH, ["New"]),
                "order_date": (WhereClauseOperation.GT, ["2014-03-01"]),
                "profit": (WhereClauseOperation.GT, ["1.0"]),
                "category": (WhereClauseOperation.EQ, ["Furniture"]),
            },
            dimension_cnts=(1, 2),
            aggregations=("AVG", "SUM"),
            formula_cnts=(2, 3, 4),
            measure_base_expressions=("[sales]",),
            filter_probability=0.1,
            bfb_probability=0.2,
            lookup_probability=0.1,
        )
        auto_gen = LODTestAutoGenerator(settings=autogen_settings)
        setting_list = auto_gen.generate_setting_list(100)

        test_runner = PreGeneratedLODTestRunner(control_api=control_api, data_api=data_api, dataset_id=dataset_id)
        test_runner.run_test_list(setting_list, ignore_400_error=False)

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/complex_queries/test_generated.py
datalens-tech/datalens-backend
import datetime
from http import HTTPStatus
from typing import Optional

from dl_api_client.dsmaker.primitives import WhereClause
from dl_api_client.dsmaker.shortcuts.dataset import (
    add_formulas_to_dataset,
    create_basic_dataset,
)
from dl_api_client.dsmaker.shortcuts.range_data import get_range_values
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_testing.connector.complex_queries import DefaultBasicLookupFunctionTestSuite
from dl_api_lib_testing.helpers.data_source import data_source_settings_from_table
from dl_api_lib_testing.helpers.lookup_checkers import (
    check_ago_data,
    check_at_date_data,
    read_date,
)
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import WhereClauseOperation
from dl_core_testing.database import make_table


class TestBasicLookupFunctions(DefaultApiTestBase, DefaultBasicLookupFunctionTestSuite):
    def test_ago(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Sales Sum Yesterday": 'AGO([Sales Sum], [order_date], "day", 1)',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="order_date"),
                ds.find_field(title="Sales Sum"),
                ds.find_field(title="Sales Sum Yesterday"),
            ],
            order_by=[
                ds.find_field(title="category"),
                ds.find_field(title="order_date"),
            ],
            filters=[
                ds.find_field(title="category").filter(op=WhereClauseOperation.EQ, values=["Office Supplies"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        check_ago_data(data_rows=data_rows, date_idx=1, value_idx=2, ago_idx=3, day_offset=1)

    def test_ago_variants(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Ago 2": "AGO([Sales Sum], [order_date])",
                "Ago 3 unit": 'AGO([Sales Sum], [order_date], "day")',
                "Ago 3 number": "AGO([Sales Sum], [order_date], 1)",
                "Ago 4": 'AGO([Sales Sum], [order_date], "day", 1)',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Sales Sum"),
                ds.find_field(title="Ago 2"),
                ds.find_field(title="Ago 3 unit"),
                ds.find_field(title="Ago 3 number"),
                ds.find_field(title="Ago 4"),
            ],
            order_by=[
                ds.find_field(title="order_date"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=2, day_offset=1)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=3, day_offset=1)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=4, day_offset=1)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=5, day_offset=1)

    def test_ago_errors(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Ago Sum": 'AGO([Sales Sum], [order_date], "day", 1)',
            },
        )

        # Dimension in AGO doesn't match dimensions in the query
        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="Sales Sum"),
                ds.find_field(title="Ago Sum"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST, result_resp.json
        # FIXME: The more specific error for AGO is temporarily reverted to the generic inconsistent agg error
        assert result_resp.bi_status_code == "ERR.DS_API.FORMULA.VALIDATION.AGG.INCONSISTENT"
        # assert result_resp.bi_status_code == 'ERR.DS_API.FORMULA.VALIDATION.LOOKUP_FUNC.UNSELECTED_DIMENSION'

        # There are no dimensions in the query
        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="Ago Sum"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST, result_resp.json
        # FIXME: Same as above
        assert result_resp.bi_status_code == "ERR.DS_API.FORMULA.VALIDATION.AGG.INCONSISTENT"
        # assert result_resp.bi_status_code == 'ERR.DS_API.FORMULA.VALIDATION.LOOKUP_FUNC.UNSELECTED_DIMENSION'

    def test_ago_in_compeng(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales RSum": "RSUM(SUM([sales]))",
                "Ago RSum": 'AGO([Sales RSum], [order_date], "day", 1)',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Sales RSum"),
                ds.find_field(title="Ago RSum"),
            ],
            order_by=[
                ds.find_field(title="order_date"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=2, day_offset=1)

    def test_dimensions_in_ago_identical_to_dims_in_query(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "order_date Clone": "[order_date]",
                "Group Sales": "SUM([sales])",
                "Ago Along Clone": "AGO([Group Sales], [order_date Clone])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Group Sales"),
                ds.find_field(title="Ago Along Clone"),
            ],
            order_by=[
                ds.find_field(title="order_date"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=2, day_offset=1)

    def test_ago_with_non_ago_aggregation(self, control_api, data_api, saved_dataset):
        """
        Check that an expression containing ago and a simple aggregation
        at the same level is sliced correctly.
        """
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "Ago": "AGO([Group Sales], [order_date])",
                "Ago And not Ago Agg": "[Group Sales] - [Ago]",
                # In this formula the left part ([Group Sales]) has no AGO (QueryFork) nodes in it,
                # while the right part ([Ago]) does.
                # This means that fork slicing will be used, but it needs to slice above aggregations (not below)
                # in the parts of the expression where there are no QueryFork nodes
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Group Sales"),
                ds.find_field(title="Ago And not Ago Agg"),
            ],
            order_by=[
                ds.find_field(title="order_date"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json

    def test_ago_with_avatarless_measure(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Avatarless Measure": "COUNT()",
                "Ago": "AGO([Avatarless Measure], [order_date])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Avatarless Measure"),
                ds.find_field(title="Ago"),
            ],
            order_by=[
                ds.find_field(title="order_date"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=2, day_offset=1)

    def test_ago_with_ignore_dimensions(self, control_api, data_api, saved_connection_id, db):
        db_table = make_table(db)
        ds = create_basic_dataset(
            api_v1=control_api,
            connection_id=saved_connection_id,
            data_source_settings=data_source_settings_from_table(table=db_table),
            formulas={
                "sum": "SUM([int_value])",
                # Create a dimension that would cause regular AGO to always return NULL:
                "day": "DAY([date_value])",
                # Regular AGO:
                "ago": "AGO([sum], [date_value])",
                # AGO that ignores [day] in JOIN
                "ago_igdim": "AGO([sum], [date_value] IGNORE DIMENSIONS [day])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="date_value"),
                ds.find_field(title="sum"),
                ds.find_field(title="ago"),
                ds.find_field(title="ago_igdim"),
                # "Bad" dimension:
                ds.find_field(title="day"),
            ],
            order_by=[
                ds.find_field(title="date_value"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        # [ago] should have all NULL values because of the "bad" dimension [day]
        for row_idx, row in enumerate(data_rows):
            assert row[2] is None, f"Expected a None value in row {row_idx}, but got {row[2]}"

        # [ago_igdim] should act the same way as a regular AGO under regular circumstances
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=3, day_offset=1)

    def test_at_date(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Sales Sum Fixed": "AT_DATE([Sales Sum], [order_date], #2014-02-02#)",
                "Sales Sum Nullable": 'AT_DATE([Sales Sum], [order_date], DATE_PARSE("2014-02-02"))',
                "Sales Sum Trunc": 'AT_DATE([Sales Sum], [order_date], DATETRUNC([order_date], "month"))',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="order_date"),
                ds.find_field(title="Sales Sum"),
                ds.find_field(title="Sales Sum Fixed"),
                ds.find_field(title="Sales Sum Nullable"),
                ds.find_field(title="Sales Sum Trunc"),
            ],
            order_by=[
                ds.find_field(title="category"),
                ds.find_field(title="order_date"),
            ],
            filters=[
                ds.find_field(title="category").filter(op=WhereClauseOperation.EQ, values=["Office Supplies"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        check_at_date_data(
            data_rows=data_rows,
            date_idx=1,
            value_idx=2,
            ago_idx=3,
            ago_date_callable=lambda d: datetime.date(2014, 2, 2),
        )
        check_at_date_data(
            data_rows=data_rows,
            date_idx=1,
            value_idx=2,
            ago_idx=4,
            ago_date_callable=lambda d: datetime.date(2014, 2, 2),
        )
        check_at_date_data(
            data_rows=data_rows, date_idx=1, value_idx=2, ago_idx=5, ago_date_callable=lambda d: d.replace(day=1)
        )

    def test_at_date_with_measure_as_third_arg(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Profit Sum": "SUM([profit])",
                "Sum At Date": "AT_DATE([Sales Sum], [order_date], #2014-02-02# + [Profit Sum]*0)",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="order_date"),
                ds.find_field(title="Sales Sum"),
                ds.find_field(title="Sum At Date"),
            ],
            order_by=[
                ds.find_field(title="category"),
                ds.find_field(title="order_date"),
            ],
            filters=[
                ds.find_field(title="category").filter(op=WhereClauseOperation.EQ, values=["Office Supplies"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        check_at_date_data(
            data_rows=data_rows,
            date_idx=1,
            value_idx=2,
            ago_idx=3,
            ago_date_callable=lambda d: datetime.date(2014, 2, 2),
        )

    def test_ago_with_bfb(self, control_api, data_api, saved_connection_id, db):
        day_offset = 3

        db_table = make_table(db)

        ds = create_basic_dataset(
            api_v1=control_api,
            connection_id=saved_connection_id,
            data_source_settings=data_source_settings_from_table(table=db_table),
            formulas={
                "sum": "SUM([int_value])",
                "date_duplicate": "[date_value]",
                "ago": f'AGO([sum], [date_value], "day", {day_offset})',
                "ago_bfb": (f'AGO([sum], [date_value], "day", {day_offset} ' f"BEFORE FILTER BY [date_duplicate])"),
                "ago_bfb_nested": (
                    f'AGO(AGO([sum], [date_value], "day", 1), [date_value], "day", {day_offset - 1} '
                    "BEFORE FILTER BY [date_duplicate])"
                ),
            },
        )

        min_date_s, _ = get_range_values(data_api.get_value_range(dataset=ds, field=ds.find_field(title="date_value")))
        min_date = read_date(min_date_s)
        gte_date_s = (min_date + datetime.timedelta(days=day_offset)).isoformat()

        def get_data_rows_with_filter(
            filters: Optional[list[WhereClause]] = None,
        ) -> list[list]:
            result_resp = data_api.get_result(
                dataset=ds,
                fields=[
                    ds.find_field(title="date_value"),
                    ds.find_field(title="sum"),
                    ds.find_field(title="ago"),
                    ds.find_field(title="ago_bfb"),
                    ds.find_field(title="ago_bfb_nested"),
                ],
                order_by=[
                    ds.find_field(title="date_value"),
                ],
                filters=filters,
                fail_ok=True,
            )
            assert result_resp.status_code == HTTPStatus.OK, result_resp.json
            return get_data_rows(result_resp)

        # Check default BFB (main dimension) in AGO without explicit BFB
        data_rows = get_data_rows_with_filter(
            filters=[
                # Apply filter to main dimension - it should be BFBed by default
                ds.find_field(title="date_value").filter(
                    # > min_date + {day_offset} days
                    op=WhereClauseOperation.GTE,
                    values=[gte_date_s],
                ),
            ]
        )
        check_ago_data(
            data_rows=data_rows,
            date_idx=0,
            value_idx=1,
            ago_idx=2,
            day_offset=day_offset,
            allow_missing_date_values=True,
        )
        # Now make sure there really are non-NULL values
        for row_idx in range(1, day_offset):  # Skip 0th row because because it has no AGO value
            assert data_rows[row_idx][2] is not None, f"Expected a non-None value in row {row_idx}"

        # Now check the explicit BFB (with filter for non-main dimension)
        data_rows = get_data_rows_with_filter(
            filters=[
                # Filter has to be applied to a dimension other than the one in AGO
                ds.find_field(title="date_duplicate").filter(
                    # > min_date + {day_offset} days
                    op=WhereClauseOperation.GTE,
                    values=[gte_date_s],
                ),
            ]
        )
        check_ago_data(
            data_rows=data_rows,
            date_idx=0,
            value_idx=1,
            ago_idx=2,
            day_offset=day_offset,
        )
        # Omit the first 2 rows because their values are not None
        check_ago_data(
            data_rows=data_rows,
            date_idx=0,
            value_idx=1,
            ago_idx=3,
            day_offset=day_offset,
            allow_missing_date_values=True,
        )
        check_ago_data(
            data_rows=data_rows,
            date_idx=0,
            value_idx=1,
            ago_idx=4,
            day_offset=day_offset,
            allow_missing_date_values=True,
        )
        # Now make sure there really are non-NULL values
        for row_idx in range(1, day_offset):  # Skip 0th row because because it has no AGO value
            assert data_rows[row_idx][3] is not None, f"Expected a non-None value in row {row_idx}"
            assert data_rows[row_idx][4] is not None, f"Expected a non-None value in row {row_idx}"

    def test_ago_with_corner_case_dimensions(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Measure": "SUM([sales])",
                "Invalid Field": "[Whaaa?...]",
                "Invalid AGO": "AGO([Measure], [Invalid Field])",
            },
            exp_status=HTTPStatus.BAD_REQUEST,
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Invalid Field"),
                ds.find_field(title="Invalid AGO"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST

    def test_ago_with_different_measures(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Measure 1": "SUM([sales])",
                "Measure 2": "COUNTD([sales])",
                "Ago Measure 1": "AGO([Measure 1], [order_date])",
                "Ago Measure 2": "AGO([Measure 2], [order_date])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Measure 1"),
                ds.find_field(title="Measure 2"),
                ds.find_field(title="Ago Measure 1"),
                ds.find_field(title="Ago Measure 2"),
            ],
            order_by=[
                ds.find_field(title="order_date"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=3, day_offset=1)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=2, ago_idx=4, day_offset=1)

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/complex_queries/test_lookup_functions.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from http import HTTPStatus

import pytest

from dl_api_client.dsmaker.api.data_api import HttpDataApiResponse
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_client.dsmaker.shortcuts.dataset import add_formulas_to_dataset
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_testing.connector.complex_queries import DefaultBasicWindowFunctionTestSuite
from dl_api_lib_testing.helpers.data_source import data_source_settings_from_table
from dl_api_lib_testing.helpers.multi_query import (
    MultiQueryInterceptor,
    count_joins,
)
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import (
    OrderDirection,
    UserDataType,
    WhereClauseOperation,
)
from dl_core_testing.database import (
    C,
    make_table,
)
from dl_query_processing.compilation.primitives import CompiledMultiQueryBase
from dl_query_processing.enums import ExecutionLevel


class TestBasicLookupFunctions(DefaultApiTestBase, DefaultBasicWindowFunctionTestSuite):
    def test_winfunc_lod_combination(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Win Func": "SUM(SUM([sales]) TOTAL)",
                "Lod Func": "SUM(SUM([sales]) FIXED)",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="region"),
                ds.find_field(title="Sales Sum"),
                ds.find_field(title="Win Func"),
                ds.find_field(title="Lod Func"),
            ],
            fail_ok=True,
        )

        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        total_value = sum(float(row[1]) for row in data_rows)

        for row in data_rows:
            win_value = float(row[2])
            lod_value = float(row[3])
            assert pytest.approx(win_value) == total_value
            assert pytest.approx(lod_value) == total_value

    def test_window_function_filter(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "Rank of Sales": 'RANK_UNIQUE([Group Sales], "asc" TOTAL)',
            },
        )

        def get_result(filters: list, window_in_select: bool = True) -> HttpDataApiResponse:
            select_fields = [
                ds.find_field(title="Group Sales"),
                ds.find_field(title="Rank of Sales"),
            ]
            if not window_in_select:
                # for testing queries with window functions only in the filter section
                select_fields = select_fields[:-1]

            result_resp = data_api.get_result(
                dataset=ds,
                fields=select_fields
                + [
                    ds.find_field(title="order_date"),
                    ds.find_field(title="city"),
                ],
                filters=filters,
                fail_ok=True,
            )
            return result_resp

        result_resp = get_result(
            filters=[
                ds.find_field(title="Rank of Sales").filter(WhereClauseOperation.LTE, [5]),
            ]
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 5

        # Combine with filters before window functions
        result_resp = get_result(
            filters=[
                # This one goes after
                ds.find_field(title="Rank of Sales").filter(WhereClauseOperation.LTE, [5]),
                # This one goes before
                ds.find_field(title="city").filter(WhereClauseOperation.STARTSWITH, ["S"]),
            ]
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 5

        # Without window functions in SELECT
        result_resp = get_result(
            filters=[
                ds.find_field(title="Rank of Sales").filter(WhereClauseOperation.LTE, [5]),
            ],
            window_in_select=False,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 5

    def test_window_function_filter_with_before_filter_by(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "Rank of Sales BFB": 'RANK_UNIQUE([Group Sales], "asc" TOTAL BEFORE FILTER BY [city])',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="city"),
                ds.find_field(title="Group Sales"),
                ds.find_field(title="Rank of Sales BFB"),
            ],
            filters=[
                ds.find_field(title="city").filter(WhereClauseOperation.STARTSWITH, ["S"]),
            ],
            fail_ok=True,
        )

        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        max_rank = 0
        possible_ranks = list(range(10000))
        for row in data_rows:
            assert row[1].startswith("S")
            row_rank = int(row[3])
            max_rank = max(max_rank, row_rank)
            possible_ranks.remove(row_rank)
        # rank should have gaps because of filters, so the max rank will be much greater than the number of rows
        assert max_rank > len(data_rows)
        possible_ranks = [r for r in possible_ranks if r <= max_rank]
        # ranks that were skipped should have remained in the list
        assert possible_ranks  # <- gaps

    def test_window_function_filter_with_before_filter_by_in_ordered_window_function(
        self, control_api, data_api, saved_dataset
    ):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "RSUM of Sales BFB": "RSUM([Group Sales] BEFORE FILTER BY [order_date])",
                "MAVG of RSUM of Sales": "MAVG([RSUM of Sales BFB], 50)",
                "MAVG of MAVG of RSUM of Sales": "MAVG([MAVG of RSUM of Sales], 50)",
                "RSUM of Sales non-BFB": "RSUM([Group Sales])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="Group Sales"),
                ds.find_field(title="RSUM of Sales BFB"),
                ds.find_field(title="order_date"),
                ds.find_field(title="MAVG of MAVG of RSUM of Sales"),
                ds.find_field(title="RSUM of Sales non-BFB"),
            ],
            order_by=[
                ds.find_field(title="order_date"),
            ],
            filters=[
                ds.find_field(title="order_date").filter(WhereClauseOperation.GTE, ["2015-05-01"]),
            ],
            fail_ok=True,
        )

        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        first_row = data_rows[0]
        # `* 2` - to eliminate float errors
        assert float(first_row[1]) > float(first_row[0]) * 2  # [RSUM of Sales BFB] > [Group Sales] * 2
        # Make sure the non-BFB field was calculated correctly too
        assert pytest.approx(float(first_row[4])) == float(first_row[0])  # [RSUM of Sales non-BFB] == [Group Sales]

    def test_order_dependent_window_function(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "RSUM ASC": 'RSUM([Group Sales], "asc" TOTAL)',
                "RSUM DESC": 'RSUM([Group Sales], "desc" TOTAL)',
            },
        )

        def _get_data(direction: OrderDirection) -> list[list]:
            order_field = ds.find_field(title="city")
            result_resp = data_api.get_result(
                dataset=ds,
                fields=[
                    ds.find_field(title="city"),
                    ds.find_field(title="Group Sales"),
                    ds.find_field(title="RSUM ASC"),
                    ds.find_field(title="RSUM DESC"),
                ],
                order_by=[
                    order_field if direction == OrderDirection.asc else order_field.desc,
                ],
                fail_ok=True,
            )
            assert result_resp.status_code == HTTPStatus.OK, result_resp.json
            data_rows = get_data_rows(result_resp)
            return data_rows

        # ORDER BY City ASC
        data_1 = _get_data(OrderDirection.asc)
        city_values_1 = [row[0] for row in data_1]
        sum_values_1 = [float(row[1]) for row in data_1]
        rsum_asc_values_1 = [float(row[2]) for row in data_1]
        rsum_desc_values_1 = [float(row[3]) for row in data_1]
        assert city_values_1[0] == min(city_values_1)
        assert city_values_1[-1] == max(city_values_1)
        assert all(
            pytest.approx(rsum_asc_values_1[i]) == rsum_asc_values_1[i - 1] + sum_values_1[i]
            for i in range(1, len(data_1))
        )
        assert all(
            pytest.approx(rsum_desc_values_1[i]) == rsum_desc_values_1[i - 1] - sum_values_1[i - 1]
            for i in range(1, len(data_1))
        )

        # ORDER BY City DESC
        # (all of the patterns are the same except for the actual values of RSUM and the order of cities)
        data_2 = _get_data(OrderDirection.desc)
        city_values_2 = [row[0] for row in data_2]
        sum_values_2 = [float(row[1]) for row in data_2]
        rsum_asc_values_2 = [float(row[2]) for row in data_2]
        rsum_desc_values_2 = [float(row[3]) for row in data_2]
        assert city_values_2[0] == max(city_values_2)
        assert city_values_2[-1] == min(city_values_2)
        assert all(
            pytest.approx(rsum_asc_values_2[i]) == rsum_asc_values_2[i - 1] + sum_values_2[i]
            for i in range(1, len(data_2))
        )
        assert all(
            pytest.approx(rsum_desc_values_2[i]) == rsum_desc_values_2[i - 1] - sum_values_2[i - 1]
            for i in range(1, len(data_2))
        )

        assert pytest.approx(rsum_asc_values_2) == rsum_desc_values_1[::-1]
        assert pytest.approx(rsum_desc_values_2) == rsum_asc_values_1[::-1]

    def test_order_by_window_function(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "Sales Rank": 'RANK_UNIQUE([Group Sales], "asc")',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="Group Sales"),
                ds.find_field(title="Sales Rank"),
            ],
            order_by=[
                ds.find_field(title="Sales Rank"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        values = [int(row[2]) for row in data_rows]
        assert values == list(range(1, len(values) + 1))

    def test_nested_window_functions(self, control_api, data_api, saved_dataset):
        scary_formula_tmpl = 'RANK_UNIQUE(10 * {repl}, "asc")'
        scary_formula = scary_formula_tmpl
        for i in range(10):
            scary_formula = scary_formula_tmpl.replace("{repl}", scary_formula)
        scary_formula = scary_formula.replace("{repl}", "[Group Sales]")

        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "Rank of Sales": 'RANK_UNIQUE([Group Sales], "asc" TOTAL)',
                "Rank of Rank of Sales": 'RANK_UNIQUE([Rank of Sales], "desc" TOTAL)',
                "RSUM of Rank of Sales": "RSUM([Rank of Sales] TOTAL ORDER BY [Rank of Sales])",
                "10x Rank of Sales": scary_formula,
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Group Sales"),
                ds.find_field(title="Rank of Sales"),
                ds.find_field(title="Rank of Rank of Sales"),
                ds.find_field(title="RSUM of Rank of Sales"),
                ds.find_field(title="10x Rank of Sales"),
            ],
            order_by=[
                ds.find_field(title="Rank of Sales"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        total_cnt = len(data_rows)
        rsum = 0
        for i, row in enumerate(data_rows):
            num = i + 1
            rsum += num
            assert int(row[2]) == num
            assert int(row[3]) == total_cnt - i
            assert int(row[4]) == rsum
            assert int(row[5]) == num

    def test_window_function_bfb_unselected_dimension_error(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Rank of Sales BFB": 'RANK_UNIQUE(SUM([sales]), "asc" TOTAL BEFORE FILTER BY [city])',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Rank of Sales BFB"),
            ],
            filters=[
                ds.find_field(title="city").filter(WhereClauseOperation.STARTSWITH, ["S"]),
            ],
            fail_ok=True,
        )

        assert result_resp.status_code == HTTPStatus.BAD_REQUEST, result_resp.json
        assert result_resp.json["code"] == "ERR.DS_API.FORMULA.VALIDATION.WIN_FUNC.BFB_UNSELECTED_DIMENSION"
        assert "neither an aggregation nor a dimension in the query" in result_resp.json["message"]

    def test_order_by_field_that_depends_on_window_function(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "Rank": "RANK_UNIQUE([Group Sales])",
                "Stringified Rank": 'CONCAT("Rank ", [Rank])',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="Rank"),
            ],
            order_by=[
                ds.find_field(title="Stringified Rank"),
            ],
            fail_ok=True,
        )

        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert data_rows

    def test_dimensions_in_window_function_identical_to_dims_in_query(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "City Clone": "[city]",
                "order_date Clone": "[order_date]",
                "Group Sales": "SUM([sales])",
                "RSum Among Clone": "RSUM([Group Sales] AMONG [City Clone])",
                "RSum Within Clone": "RSUM([Group Sales] WITHIN [order_date Clone])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="order_date"),
                ds.find_field(title="RSum Among Clone"),
                ds.find_field(title="RSum Within Clone"),
            ],
            order_by=[
                ds.find_field(title="city"),
                ds.find_field(title="order_date"),
            ],
            fail_ok=True,
        )

        assert result_resp.status_code == HTTPStatus.OK, result_resp.json

    def test_order_by_multilevel_window_function(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "Multi Rank": 'RANK(RANK(RANK(RANK_UNIQUE([Group Sales], "asc"), "asc"), "asc"), "asc")',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="Group Sales"),
            ],
            order_by=[
                ds.find_field(title="Multi Rank"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json

    def test_extra_dimensions_in_within(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Value": "SUM(SUM([sales]) WITHIN [category], [city])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="category"),
                ds.find_field(title="Value"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert data_rows
        by_cat: dict[str, set[str]] = defaultdict(set)
        for row in data_rows:
            by_cat[row[0]].add(row[1])

        assert len(by_cat) == 3
        # Values should be equivalent to `SUM(SUM([sales]) WITHIN [Category])`,
        # so all values within a category should be the same
        for cat_values in by_cat.values():
            assert len(cat_values) == 1

    def test_measure_in_within(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "My Measure": "COUNT()",
                "Win Value": "SUM(SUM([sales]) WITHIN [My Measure])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="My Measure"),
                ds.find_field(title="Win Value"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert data_rows

        # Make sure that all `Win Value` values are the same for matching `My Measure` values
        by_my_measure = {row[2]: row[1] for row in data_rows}
        for row in data_rows:
            assert row[1] == by_my_measure[row[2]]

    def test_dim_complex_dimension_in_within(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Measure": "SUM([sales])",
                "Complex Dim": 'DATETRUNC([order_date], "week")',
                "Win Value": "SUM([Measure] WITHIN [Complex Dim])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Win Value"),
            ],
            order_by=[
                ds.find_field(title="order_date"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        date_values = {row[0] for row in data_rows}
        measure_values = {row[1] for row in data_rows}
        assert len(date_values) > len(measure_values) > 10  # the last is an arbitrary nonzero number

    def test_nested_among(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Nested Among": "SUM(RANK(SUM([sales]) AMONG [category]) AMONG [category])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="region"),
                ds.find_field(title="category"),
                ds.find_field(title="Nested Among"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert data_rows

        for row in data_rows:
            # 3 Categories in each Region => RANK = 1,2,3; SUM(RANK) = 6
            assert int(row[2]) == 6

    def test_null_dimensions(self, control_api, data_api, saved_connection_id, db):
        raw_data = [
            {"id": 1, "city": "New York", "category": "Office Supplies", "sales": 1},
            {"id": 2, "city": "New York", "category": "Furniture", "sales": 10},
            {"id": 3, "city": "New Rochelle", "category": "Office Supplies", "sales": 100},
            {"id": 4, "city": "New Rochelle", "category": "Furniture", "sales": 1000},
            {"id": 5, "city": None, "category": "Office Supplies", "sales": 10000},
            {"id": 6, "city": None, "category": "Furniture", "sales": 100000},
        ]
        columns = [
            C("id", UserDataType.integer, vg=lambda rn, **kwargs: raw_data[rn]["id"]),
            C("city", UserDataType.string, vg=lambda rn, **kwargs: raw_data[rn]["city"]),
            C("category", UserDataType.string, vg=lambda rn, **kwargs: raw_data[rn]["category"]),
            C("sales", UserDataType.integer, vg=lambda rn, **kwargs: raw_data[rn]["sales"]),
        ]
        db_table = make_table(db, columns=columns, rows=len(raw_data))

        ds = Dataset()
        ds.sources["source_1"] = ds.source(
            connection_id=saved_connection_id, **data_source_settings_from_table(db_table)
        )
        ds.source_avatars["avatar_1"] = ds.sources["source_1"].avatar()

        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=ds,
            formulas={
                "rank_within_cat": "RANK(SUM([sales]) WITHIN [category])",
                "rank_within_city": "RANK(SUM([sales]) WITHIN [city])",
                "max_id": "MAX([id])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="category"),
                ds.find_field(title="rank_within_cat"),
                ds.find_field(title="rank_within_city"),
            ],
            order_by=[
                ds.find_field(title="max_id"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert data_rows

        city = [row[0] for row in data_rows]
        rank_within_cat = [int(row[2]) for row in data_rows]
        rank_within_city = [int(row[3]) for row in data_rows]

        assert city == ["New York", "New York", "New Rochelle", "New Rochelle", None, None]
        assert rank_within_cat == [3, 3, 2, 2, 1, 1]
        assert rank_within_city == [2, 1, 2, 1, 2, 1]

    def test_compeng_part_has_no_joins(self, control_api, data_api, saved_dataset, monkeypatch):
        """
        Check that when selecting:
        1. dimension,
        2. aggregation and
        3. window function
        in one simple request, there are no JOINs in compeng
        """

        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales Sum": "SUM([sales])",
                "Rank of Sales Sum": 'RANK_UNIQUE(SUM([sales]), "asc" TOTAL)',
            },
        )

        def intercept_query(multi_query: CompiledMultiQueryBase) -> None:
            compeng_part = multi_query.for_level_type(ExecutionLevel.compeng)
            assert compeng_part.query_count() > 1
            assert count_joins(compeng_part) == 0

        interceptor = MultiQueryInterceptor(mpatch=monkeypatch, callback=intercept_query)

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="region"),
                ds.find_field(title="Sales Sum"),
                ds.find_field(title="Rank of Sales Sum"),
            ],
            fail_ok=True,
        )

        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        assert interceptor.intercepted

    def test_wf_and_non_wf_in_filter(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Rank": "RANK(SUM([sales]) TOTAL)",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
            ],
            filters=[
                ds.find_field(title="city").filter(WhereClauseOperation.GT, ["A"]),
                ds.find_field(title="Rank").filter(WhereClauseOperation.LT, ["7"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 6

    def test_round_2_in_compeng(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Measure Non-Win": "ROUND(SUM([sales]) / 1000.0, 2)",
                "Measure Win": "ROUND(MAX(SUM([sales]) TOTAL) / 1000.0, 2)",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="Measure Non-Win"),
                ds.find_field(title="Measure Win"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 1

        non_win = float(data_rows[0][0])
        win = float(data_rows[0][1])

        assert win == pytest.approx(non_win)

    def test_subquery_column_name_conflict(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "Sales RSUM BFB": ("RSUM(SUM([sales]) TOTAL BEFORE FILTER BY [order_date])"),
                "Sales RSUM BFB MAVG BFB": (
                    "MAVG(RSUM(SUM([sales]) TOTAL BEFORE FILTER BY [order_date]), 100 BEFORE FILTER BY [order_date])"
                ),
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="Sales RSUM BFB MAVG BFB"),
                ds.find_field(title="Group Sales"),
            ],
            order_by=[
                ds.find_field(title="order_date"),
            ],
            filters=[
                ds.find_field(title="order_date").filter(
                    op=WhereClauseOperation.EQ,
                    values=["2014-05-01"],
                ),
                ds.find_field(title="Sales RSUM BFB").filter(
                    op=WhereClauseOperation.LT,
                    values=["2000000"],
                ),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json

    def test_bfb_with_only_winfunc_measure(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Sales SUM": "SUM([sales])",
                "Sales RSUM BFB": "RSUM(SUM([sales]) TOTAL ORDER BY [order_date] BEFORE FILTER BY [order_date])",
            },
        )

        def get_data(measures: tuple[str, ...]) -> list[list[str]]:
            result_resp = data_api.get_result(
                dataset=ds,
                fields=[ds.find_field(title="order_date"), *[ds.find_field(title=measure) for measure in measures]],
                filters=[
                    ds.find_field(title="category").filter(WhereClauseOperation.EQ, ["Office Supplies"]),
                    ds.find_field(title="order_date").filter(WhereClauseOperation.GT, ["2014-04-01"]),
                ],
                fail_ok=True,
            )
            assert result_resp.status_code == HTTPStatus.OK, result_resp.json
            data_rows = get_data_rows(result_resp)
            assert len(data_rows) > 1
            return data_rows

        data_rows_no_sum = get_data(measures=("Sales RSUM BFB",))
        data_rows_with_sum = get_data(measures=("Sales RSUM BFB", "Sales SUM"))
        assert len(data_rows_no_sum) == len(data_rows_with_sum)
        for row_no_sum, row_with_sum in zip(data_rows_no_sum, data_rows_with_sum):
            assert row_no_sum[0] == row_with_sum[0]  # The dimension
            assert row_no_sum[1] == row_with_sum[1]  # The measure

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/complex_queries/test_window_functions.py
datalens-tech/datalens-backend
from dl_api_client.dsmaker.primitives import ResultField
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import (
    AggregationFunction,
    CalcMode,
)


class TestResultErrors(DefaultApiTestBase):
    def test_empty_query(self, saved_dataset, data_api):
        result_resp = data_api.get_result(dataset=saved_dataset, fields=[], fail_ok=True)
        assert result_resp.status_code == 400
        assert result_resp.bi_status_code == "ERR.DS_API.EMPTY_QUERY"

    def test_get_nonexisting_field(self, saved_dataset, data_api):
        result_resp = data_api.get_result(dataset=saved_dataset, fields=[ResultField(title="unknown")], fail_ok=True)
        assert result_resp.status_code == 400
        assert result_resp.bi_status_code == "ERR.DS_API.FIELD.NOT_FOUND"

    def test_calcmode_formula_without_formula_field(self, saved_dataset, data_api):
        ds = saved_dataset
        title = "Not a formula"
        ds.result_schema[title] = ds.field(title=title, calc_mode=CalcMode.formula)

        result_resp = data_api.get_result(dataset=ds, fields=[ds.find_field(title)], fail_ok=True)
        assert result_resp.status_code == 400
        assert result_resp.bi_status_code == "ERR.DS_API.FORMULA.PARSE.UNEXPECTED_EOF.EMPTY_FORMULA"

    def test_filter_errors(self, saved_dataset, data_api):
        ds = saved_dataset
        ds.result_schema["New date"] = ds.field(formula="DATE(NOW())")  # just a date field
        ds.result_schema["New bool"] = ds.field(formula='"P" = "NP"')  # just a bool field

        # invalid value for date
        result_resp = data_api.get_result(
            dataset=ds,
            fields=[ds.find_field(title="New date")],
            filters=[ds.find_field(title="New date").filter("EQ", values=["Office Supplies"])],
            fail_ok=True,
        )
        assert result_resp.status_code == 400
        assert result_resp.bi_status_code == "ERR.DS_API.FILTER.INVALID_VALUE"

        # a very special check for invalid bool values
        result_resp = data_api.get_result(
            dataset=ds,
            fields=[ds.find_field(title="New date")],
            filters=[ds.result_schema["New bool"].filter("EQ", values=[""])],
            fail_ok=True,
        )
        assert result_resp.status_code == 400
        assert result_resp.bi_status_code == "ERR.DS_API.FILTER.INVALID_VALUE"

        # unknown field
        result_resp = data_api.get_result(
            dataset=ds,
            fields=[ds.find_field(title="New date")],
            filters=[ds.field(title="Made-up field").filter("EQ", values=["whatever"])],
            fail_ok=True,
        )
        assert result_resp.status_code == 400
        assert result_resp.bi_status_code == "ERR.DS_API.FIELD.NOT_FOUND"

    def test_invalid_group_by_configuration(self, saved_dataset, data_api):
        ds = saved_dataset
        ds.result_schema["Measure"] = ds.field(title="Measure", aggregation=AggregationFunction.sum)

        result_resp = data_api.get_result(
            dataset=saved_dataset,
            fields=[ds.find_field(title="Measure")],
            disable_group_by=True,
            fail_ok=True,
        )
        assert result_resp.status_code == 400
        assert result_resp.bi_status_code == "ERR.DS_API.INVALID_GROUP_BY_CONFIGURATION"
        assert result_resp.json["message"] == "Invalid parameter disable_group_by for dataset with measure fields"

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/test_errors.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import OrderedDict
from http import HTTPStatus
import json
from typing import (
    Any,
    Callable,
)
from typing import OrderedDict as OrderedDictTyping
from typing import Sequence

from dl_api_client.dsmaker.api.data_api import SyncHttpDataApiV2
from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.primitives import (
    Dataset,
    OrderedField,
    ResultField,
)
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_client.dsmaker.shortcuts.tree import make_request_legend_items_for_tree_branches
from dl_api_lib_testing.helpers.data_source import data_source_settings_from_table
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import (
    FieldRole,
    UserDataType,
)
from dl_core_testing.database import (
    C,
    Db,
    make_table,
)


TREE_DATA = [
    # / Company
    {"id": 10, "dept": ["Company"], "salary": 0, "is_person": 0},
    # / Company / R&D
    {"id": 11, "dept": ["Company", "R&D"], "salary": 0, "is_person": 0},
    # / Company / R&D / Cloud
    {"id": 12, "dept": ["Company", "R&D", "Cloud"], "salary": 0, "is_person": 0},
    {"id": 13, "dept": ["Company", "R&D", "Cloud", "Dan"], "salary": 100, "is_person": 1},
    # / Company / R&D / Cloud / Rain
    {"id": 14, "dept": ["Company", "R&D", "Cloud", "Rain"], "salary": 0, "is_person": 0},
    {"id": 15, "dept": ["Company", "R&D", "Cloud", "Rain", "Jessica"], "salary": 100, "is_person": 1},
    {"id": 16, "dept": ["Company", "R&D", "Cloud", "Rain", "Claude"], "salary": 100, "is_person": 1},
    # / Company / R&D / Cloud / Snow
    {"id": 17, "dept": ["Company", "R&D", "Cloud", "Snow"], "salary": 0, "is_person": 0},
    {"id": 18, "dept": ["Company", "R&D", "Cloud", "Snow", "Ivan"], "salary": 100, "is_person": 1},
    {"id": 19, "dept": ["Company", "R&D", "Cloud", "Snow", "Ivonne"], "salary": 100, "is_person": 1},
    {"id": 20, "dept": ["Company", "R&D", "Cloud", "Snow", "Angela"], "salary": 100, "is_person": 1},
    # / Company / R&D / Cloud / Lightning
    {"id": 21, "dept": ["Company", "R&D", "Cloud", "Lightning"], "salary": 0, "is_person": 0},
    {"id": 22, "dept": ["Company", "R&D", "Cloud", "Lightning", "Gary"], "salary": 100, "is_person": 1},
    # / Company / R&D / Search
    {"id": 23, "dept": ["Company", "R&D", "Search"], "salary": 0, "is_person": 0},
    {"id": 24, "dept": ["Company", "R&D", "Search", "Agnes"], "salary": 100, "is_person": 1},
    # / Company / Sales
    {"id": 25, "dept": ["Company", "Sales"], "salary": 0, "is_person": 0},
    {"id": 26, "dept": ["Company", "Sales", "Alice"], "salary": 10000, "is_person": 1},
    # / Company / Sales / CN
    {"id": 27, "dept": ["Company", "Sales", "CN"], "salary": 0, "is_person": 0},
    # / Company / Sales / EU
    {"id": 28, "dept": ["Company", "Sales", "EU"], "salary": 0, "is_person": 0},
    {"id": 29, "dept": ["Company", "Sales", "EU", "Noname"], "salary": 50, "is_person": 1},
    # / Company / Sales / US
    {"id": 30, "dept": ["Company", "Sales", "US"], "salary": 0, "is_person": 0},
]


def make_tree_dataset(db: Db, connection_id: str, control_api: SyncHttpDatasetApiV1) -> Dataset:
    columns = [
        C("id", UserDataType.integer, vg=lambda rn, **kwargs: TREE_DATA[rn]["id"]),
        C("dept", UserDataType.array_str, vg=lambda rn, **kwargs: TREE_DATA[rn]["dept"]),
        C("salary", UserDataType.integer, vg=lambda rn, **kwargs: TREE_DATA[rn]["salary"]),
        C("is_person", UserDataType.integer, vg=lambda rn, **kwargs: TREE_DATA[rn]["is_person"]),
    ]
    db_table = make_table(db, columns=columns, rows=len(TREE_DATA))
    ds = Dataset()
    ds.sources["source_1"] = ds.source(connection_id=connection_id, **data_source_settings_from_table(db_table))
    ds.source_avatars["avatar_1"] = ds.sources["source_1"].avatar()
    ds.result_schema["dept_tree"] = ds.field(formula="TREE([dept])")
    ds.result_schema["person_count"] = ds.field(formula="COUNT_IF(BOOL([is_person]))")
    ds = control_api.apply_updates(dataset=ds).dataset
    ds = control_api.save_dataset(ds).dataset
    return ds


def get_tree(
    branches: list[tuple[int, list[str]]],
    dataset: Dataset,
    data_api: SyncHttpDataApiV2,
    order_by: Sequence[ResultField | OrderedField] = (),
    measures: Sequence[tuple[str, Callable]] = (("person_count", int),),
) -> OrderedDictTyping[tuple[str, ...], tuple[Any, ...]]:
    branch_items = make_request_legend_items_for_tree_branches(
        branches=branches,
        dataset=dataset,
        tree_title="dept_tree",
    )
    result_resp = data_api.get_result(
        dataset=dataset,
        fields=[
            *branch_items,
            *[dataset.find_field(title=name) for name, cast_func in measures],
        ],
        order_by=list(order_by),
        fail_ok=True,
    )
    assert result_resp.status_code == HTTPStatus.OK, result_resp.json
    data_rows = get_data_rows(result_resp)
    result = OrderedDict(
        (tuple(json.loads(row[0])), tuple([cast_func(row[idx + 1]) for idx, (name, cast_func) in enumerate(measures)]))
        for row in data_rows
    )

    block_metas = [block_meta for block_meta in result_resp.json["blocks"]]
    assert len(block_metas) == len(branch_items)
    for block_meta in block_metas:
        assert block_meta["query"] is not None

    return result


class TestResultWithTrees(DefaultApiTestBase):
    def test_simple_tree(self, control_api, data_api, db, saved_connection_id):
        connection_id = saved_connection_id
        ds = make_tree_dataset(db=db, connection_id=connection_id, control_api=control_api)

        tree = get_tree(
            [
                (1, []),
            ],
            dataset=ds,
            data_api=data_api,
        )
        assert tree == {
            ("Company",): (10,),
        }

        tree = get_tree(
            [
                (1, []),
                (2, ["Company"]),
            ],
            dataset=ds,
            data_api=data_api,
        )
        assert tree == {
            ("Company",): (10,),
            (
                "Company",
                "R&D",
            ): (8,),
            (
                "Company",
                "Sales",
            ): (2,),
        }

        tree = get_tree(
            [
                (1, []),
                (2, ["Company"]),
                (3, ["Company", "R&D"]),
            ],
            dataset=ds,
            data_api=data_api,
        )
        assert tree == {
            ("Company",): (10,),
            (
                "Company",
                "R&D",
            ): (8,),
            ("Company", "R&D", "Cloud"): (7,),
            ("Company", "R&D", "Search"): (1,),
            (
                "Company",
                "Sales",
            ): (2,),
        }

        tree = get_tree(
            [
                (1, []),
                (2, ["Company"]),
                (3, ["Company", "R&D"]),
                (4, ["Company", "R&D", "Cloud"]),
            ],
            dataset=ds,
            data_api=data_api,
        )
        assert tree == {
            ("Company",): (10,),
            (
                "Company",
                "R&D",
            ): (8,),
            ("Company", "R&D", "Cloud"): (7,),
            ("Company", "R&D", "Cloud", "Dan"): (1,),
            ("Company", "R&D", "Cloud", "Rain"): (2,),
            ("Company", "R&D", "Cloud", "Snow"): (3,),
            ("Company", "R&D", "Cloud", "Lightning"): (1,),
            ("Company", "R&D", "Search"): (1,),
            (
                "Company",
                "Sales",
            ): (2,),
        }

    def test_simple_tree_with_order_by(self, control_api, data_api, db, saved_connection_id):
        connection_id = saved_connection_id
        ds = make_tree_dataset(db=db, connection_id=connection_id, control_api=control_api)

        # Order by measure
        tree = get_tree(
            dataset=ds,
            data_api=data_api,
            branches=[(1, []), (2, ["Company"]), (3, ["Company", "R&D"])],
            order_by=[
                ds.find_field(title="person_count"),
            ],
        )
        assert tree == OrderedDict(
            (
                (("Company",), (10,)),
                (
                    (
                        "Company",
                        "Sales",
                    ),
                    (2,),
                ),
                (
                    (
                        "Company",
                        "R&D",
                    ),
                    (8,),
                ),
                (("Company", "R&D", "Search"), (1,)),
                (("Company", "R&D", "Cloud"), (7,)),
            )
        )

        # Order by branch
        tree = get_tree(
            dataset=ds,
            data_api=data_api,
            branches=[(1, []), (2, ["Company"]), (3, ["Company", "R&D"])],
            order_by=[
                ds.find_field(title="dept_tree"),
            ],
        )
        assert tree == OrderedDict(
            (
                (("Company",), (10,)),
                (
                    (
                        "Company",
                        "R&D",
                    ),
                    (8,),
                ),
                (("Company", "R&D", "Cloud"), (7,)),
                (("Company", "R&D", "Search"), (1,)),
                (
                    (
                        "Company",
                        "Sales",
                    ),
                    (2,),
                ),
            )
        )

    def test_simple_tree_with_order_by_and_rsum(
        self,
        control_api,
        data_api,
        db,
        saved_connection_id,
    ):
        connection_id = saved_connection_id
        ds = make_tree_dataset(db=db, connection_id=connection_id, control_api=control_api)
        ds.result_schema["rsum_measure"] = ds.field(formula="RSUM([person_count])")

        # Order by measure
        tree = get_tree(
            dataset=ds,
            data_api=data_api,
            branches=[(1, []), (2, ["Company"]), (3, ["Company", "R&D"])],
            order_by=[
                ds.find_field(title="person_count"),
            ],
            measures=[
                ("person_count", int),
                ("rsum_measure", int),
            ],
        )
        assert tree == OrderedDict(
            (
                (("Company",), (10, 10)),
                (
                    (
                        "Company",
                        "Sales",
                    ),
                    (2, 2),
                ),
                (
                    (
                        "Company",
                        "R&D",
                    ),
                    (8, 10),
                ),
                (("Company", "R&D", "Search"), (1, 1)),
                (("Company", "R&D", "Cloud"), (7, 8)),
            )
        )

        # Order by tree path
        tree = get_tree(
            dataset=ds,
            data_api=data_api,
            branches=[(1, []), (2, ["Company"]), (3, ["Company", "R&D"])],
            order_by=[
                ds.find_field(title="dept_tree"),
            ],
            measures=[
                ("person_count", int),
                ("rsum_measure", int),
            ],
        )
        assert tree == OrderedDict(
            (
                (("Company",), (10, 10)),
                (
                    (
                        "Company",
                        "R&D",
                    ),
                    (8, 8),
                ),
                (("Company", "R&D", "Cloud"), (7, 7)),
                (("Company", "R&D", "Search"), (1, 8)),
                (
                    (
                        "Company",
                        "Sales",
                    ),
                    (2, 10),
                ),
            )
        )

    def test_tree_and_totals_incompatibility_error(self, control_api, data_api, db, saved_connection_id):
        connection_id = saved_connection_id
        ds = make_tree_dataset(db=db, connection_id=connection_id, control_api=control_api)

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="dept_tree").as_req_legend_item(
                    role=FieldRole.tree,
                    tree_level=1,
                    tree_prefix=[],
                    dimension_values={},
                ),
                ds.find_field(title="person_count").as_req_legend_item(role=FieldRole.total),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST
        assert result_resp.bi_status_code == "ERR.DS_API.BLOCK.ITEM_COMPATIBILITY"

    def test_multiple_trees_error(self, control_api, data_api, db, saved_connection_id):
        connection_id = saved_connection_id
        ds = make_tree_dataset(db=db, connection_id=connection_id, control_api=control_api)

        ds.result_schema["dept_tree_2"] = ds.field(formula="TREE([dept])")

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="dept_tree").as_req_legend_item(
                    role=FieldRole.tree,
                    tree_level=1,
                    tree_prefix=[],
                    dimension_values={},
                ),
                ds.find_field(title="dept_tree_2").as_req_legend_item(
                    role=FieldRole.tree,
                    tree_level=1,
                    tree_prefix=[],
                    dimension_values={},
                ),
                ds.find_field(title="person_count"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST
        assert result_resp.bi_status_code == "ERR.DS_API.TREE.MULTIPLE"

    def test_invalid_data_type_for_tree_error(self, control_api, data_api, db, saved_connection_id):
        connection_id = saved_connection_id
        ds = make_tree_dataset(db=db, connection_id=connection_id, control_api=control_api)

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="id").as_req_legend_item(
                    role=FieldRole.tree,
                    tree_level=1,
                    tree_prefix=[],
                    dimension_values={},
                ),
                ds.find_field(title="person_count"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST
        assert result_resp.bi_status_code == "ERR.DS_API.LEGEND.ROLE_DATA_TYPE_MISMATCH"

    def test_invalid_legend_item_reference_error(self, control_api, data_api, db, saved_connection_id):
        connection_id = saved_connection_id
        ds = make_tree_dataset(db=db, connection_id=connection_id, control_api=control_api)

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="dept_tree").as_req_legend_item(
                    role=FieldRole.tree,
                    tree_level=1,
                    tree_prefix=[],
                    dimension_values={5: "8"},
                ),
                ds.find_field(title="person_count"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST
        assert result_resp.bi_status_code == "ERR.DS_API.LEGEND.ITEM_REFERENCE"

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/test_trees.py
datalens-tech/datalens-backend
from __future__ import annotations

import uuid

import pytest

from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_testing.data_api_base import DataApiTestParams
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import FieldType


class TestUpdates(DefaultApiTestBase):
    @pytest.fixture(scope="function")
    def data_api_test_params(self, sample_table) -> DataApiTestParams:
        # This default is defined for the sample table
        return DataApiTestParams(
            two_dims=("category", "city"),
            summable_field="sales",
            range_field="sales",
            distinct_field="city",
            date_field="order_date",
        )

    def test_result_with_updates(self, saved_dataset, data_api, data_api_test_params):
        ds = saved_dataset
        id_1, id_2, id_3 = (str(uuid.uuid4()) for _ in range(3))
        result_resp = data_api.get_result(
            dataset=ds,
            updates=[
                ds.field(
                    id=id_1,
                    title="First",
                    formula=f"SUM([{data_api_test_params.summable_field}]) / 100",
                    type=FieldType.MEASURE,
                ).add(),
                # use an invalid field type for the second one to make sure it fixes itself
                ds.field(
                    id=id_2,
                    title="Second",
                    formula=f"COUNTD([{data_api_test_params.distinct_field}])",
                    type=FieldType.DIMENSION,
                ).add(),
                ds.field(id=id_3, title="Third", formula="[First] / [Second]", type=FieldType.MEASURE).add(),
            ],
            fields=[
                ds.field(id=id_1),
                ds.field(id=id_2),
                ds.field(id=id_3),
            ],
        )
        assert result_resp.status_code == 200, result_resp.response_errors
        result_data = result_resp.data
        titles = [field.title for field in result_data["fields"]]
        assert titles == ["First", "Second", "Third"]

    def test_get_result_add_update_field_without_avatar(self, saved_dataset, data_api):
        ds = saved_dataset
        ds.result_schema["Some Field"] = ds.field(formula="NOW()")

        result_resp = data_api.get_result(
            dataset=ds,
            updates=[
                ds.result_schema["Some Field"].update(calc_mode="direct"),
            ],
            fields=[
                ds.find_field(title="Some Field"),
            ],
        )
        assert result_resp.status_code == 200, result_resp.response_errors
        data_rows = get_data_rows(result_resp)
        assert data_rows

lib/dl_api_lib/dl_api_lib_tests/db/data_api/result/test_updates.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib_testing.data_api_base import DataApiTestParams
from dl_api_lib_tests.db.base import DefaultApiTestBase
from dl_constants.enums import WhereClauseOperation


class TestDataApiErrors(DefaultApiTestBase):
    @pytest.fixture(scope="function")
    def data_api_test_params(self, sample_table) -> DataApiTestParams:
        # This default is defined for the sample table
        return DataApiTestParams(
            two_dims=("category", "city"),
            summable_field="sales",
            range_field="sales",
            distinct_field="city",
            date_field="order_date",
        )

    def test_distinct_measure_filter_error(self, saved_dataset, data_api, data_api_test_params):
        ds = saved_dataset
        ds.result_schema["Measure"] = ds.field(formula=f"SUM([{data_api_test_params.summable_field}])")

        distinct_resp = data_api.get_distinct(
            dataset=ds,
            field=ds.find_field(title=data_api_test_params.distinct_field),
            filters=[
                ds.find_field(title="Measure").filter(WhereClauseOperation.GT, [100]),
            ],
            fail_ok=True,
        )
        assert distinct_resp.status_code == 400
        assert distinct_resp.bi_status_code == "ERR.DS_API.FILTER.MEASURE_UNSUPPORTED"

    def test_range_measure_filter_error(self, saved_dataset, data_api, data_api_test_params):
        ds = saved_dataset
        ds.result_schema["Measure"] = ds.field(formula=f"SUM([{data_api_test_params.summable_field}])")

        range_resp = data_api.get_value_range(
            dataset=ds,
            field=ds.find_field(title=data_api_test_params.range_field),
            filters=[
                ds.find_field(title="Measure").filter(WhereClauseOperation.GT, [100]),
            ],
            fail_ok=True,
        )
        assert range_resp.status_code == 400
        assert range_resp.bi_status_code == "ERR.DS_API.FILTER.MEASURE_UNSUPPORTED"

lib/dl_api_lib/dl_api_lib_tests/db/data_api/test_errors.py
datalens-tech/datalens-backend
from dl_api_lib.query.formalization.raw_pivot_specs import (
    RawDimensionPivotRoleSpec,
    RawPivotLegendItem,
    RawPivotMeasureRoleSpec,
    RawPivotSpec,
)
from dl_api_lib.query.formalization.raw_specs import (
    PlaceholderRef,
    RawQuerySpecUnion,
    RawRoleSpec,
    RawSelectFieldSpec,
    RawTemplateRoleSpec,
    TitleFieldRef,
)
from dl_api_lib.request_model.data import PivotDataRequestModel
from dl_api_lib.request_model.normalization.drm_normalizer_pivot import PivotRequestModelNormalizer
from dl_constants.enums import (
    FieldRole,
    PivotRole,
)


def test_simple_totals():
    original_drm = PivotDataRequestModel(
        raw_query_spec_union=RawQuerySpecUnion(
            select_specs=[
                RawSelectFieldSpec(ref=TitleFieldRef(title="First Dim"), legend_item_id=0),
                RawSelectFieldSpec(ref=TitleFieldRef(title="Second Dim"), legend_item_id=1),
                RawSelectFieldSpec(ref=TitleFieldRef(title="First Measure"), legend_item_id=2),
            ],
        ),
        pivot=RawPivotSpec(
            structure=[
                RawPivotLegendItem(
                    legend_item_ids=[0],
                    role_spec=RawDimensionPivotRoleSpec(role=PivotRole.pivot_row),
                ),
                RawPivotLegendItem(
                    legend_item_ids=[1],
                    role_spec=RawDimensionPivotRoleSpec(role=PivotRole.pivot_column),
                ),
                RawPivotLegendItem(
                    legend_item_ids=[2],
                    role_spec=RawPivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                ),
            ],
            totals=None,
            with_totals=True,
        ),
    )
    normalizer = PivotRequestModelNormalizer()
    normalized_drm = normalizer.normalize_drm(original_drm)

    expected_drm = PivotDataRequestModel(
        raw_query_spec_union=RawQuerySpecUnion(
            select_specs=[
                RawSelectFieldSpec(ref=TitleFieldRef(title="First Dim"), legend_item_id=0, block_id=0),
                RawSelectFieldSpec(ref=TitleFieldRef(title="Second Dim"), legend_item_id=1, block_id=0),
                RawSelectFieldSpec(ref=TitleFieldRef(title="First Measure"), legend_item_id=2, block_id=0),
                RawSelectFieldSpec(
                    ref=PlaceholderRef(),
                    role_spec=RawTemplateRoleSpec(role=FieldRole.template, template=""),
                    legend_item_id=3,
                    block_id=1,
                ),
                RawSelectFieldSpec(ref=TitleFieldRef(title="Second Dim"), legend_item_id=4, block_id=1),
                RawSelectFieldSpec(
                    ref=TitleFieldRef(title="First Measure"),
                    role_spec=RawRoleSpec(role=FieldRole.total),
                    legend_item_id=5,
                    block_id=1,
                ),
                RawSelectFieldSpec(
                    ref=PlaceholderRef(),
                    role_spec=RawTemplateRoleSpec(role=FieldRole.template, template=""),
                    legend_item_id=6,
                    block_id=2,
                ),
                RawSelectFieldSpec(ref=TitleFieldRef(title="First Dim"), legend_item_id=7, block_id=2),
                RawSelectFieldSpec(
                    ref=TitleFieldRef(title="First Measure"),
                    role_spec=RawRoleSpec(role=FieldRole.total),
                    legend_item_id=8,
                    block_id=2,
                ),
                RawSelectFieldSpec(
                    ref=PlaceholderRef(),
                    role_spec=RawTemplateRoleSpec(role=FieldRole.template, template=""),
                    legend_item_id=9,
                    block_id=3,
                ),
                RawSelectFieldSpec(
                    ref=PlaceholderRef(),
                    role_spec=RawTemplateRoleSpec(role=FieldRole.template, template=""),
                    legend_item_id=10,
                    block_id=3,
                ),
                RawSelectFieldSpec(
                    ref=TitleFieldRef(title="First Measure"),
                    role_spec=RawRoleSpec(role=FieldRole.total),
                    legend_item_id=11,
                    block_id=3,
                ),
            ],
        ),
        pivot=RawPivotSpec(
            structure=[
                RawPivotLegendItem(
                    legend_item_ids=[1, 4, 6, 10],
                    role_spec=RawDimensionPivotRoleSpec(role=PivotRole.pivot_column),
                ),
                RawPivotLegendItem(
                    legend_item_ids=[0, 3, 7, 9],
                    role_spec=RawDimensionPivotRoleSpec(role=PivotRole.pivot_row),
                ),
                RawPivotLegendItem(
                    legend_item_ids=[2, 5, 8, 11],
                    role_spec=RawPivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                ),
            ],
            totals=None,
            with_totals=None,
        ),
    )

    assert normalized_drm == expected_drm

lib/dl_api_lib/dl_api_lib_tests/unit/drm_normalization/test_pivot.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib.pivot.pandas.transformer import PdPivotTransformer
from dl_api_lib.pivot.primitives import DataCell as DC
from dl_api_lib.pivot.primitives import DataCellVector as DV
from dl_api_lib.pivot.primitives import DataRow
from dl_api_lib.pivot.primitives import MeasureNameValue as MNV
from dl_api_lib.pivot.primitives import PivotHeader
from dl_api_lib.pivot.table import PivotTable
from dl_api_lib.query.formalization.pivot_legend import (
    PivotDimensionRoleSpec,
    PivotLegend,
    PivotLegendItem,
    PivotMeasureRoleSpec,
)
from dl_constants.enums import (
    FieldType,
    PivotItemType,
    PivotRole,
    UserDataType,
)
import dl_query_processing.exc
from dl_query_processing.legend.field_legend import (
    FieldObjSpec,
    Legend,
    LegendItem,
    MeasureNameObjSpec,
)
from dl_query_processing.merging.primitives import MergedQueryDataRow


def test_pivot_empty_data():
    fid_city, fid_ctgry, fid_sales = ("123", "456", "789")
    data = []

    # Legend item IDs
    liid_ctgry, liid_city, liid_sales = (0, 1, 2)
    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_city = 10
    piid_ctgry = 20
    piid_sales = 30
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    actual_pivot_table = transformer.pivot(data)

    expected_pivot_columns = []
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = []
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def test_pivot_duplicate_dimension_values():
    fid_city, fid_ctgry, fid_sales = ("123", "456", "789")
    liid_ctgry, liid_city, liid_sales = (0, 1, 2)
    legend_item_ids = (liid_ctgry, liid_city, liid_sales)

    data = [
        MergedQueryDataRow(data=("Detroit", "Furniture", 100), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Furniture", 200), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Office Supplies", 300), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", "Furniture", 400), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_city = 10
    piid_ctgry = 20
    piid_sales = 30
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    with pytest.raises(dl_query_processing.exc.PivotDuplicateDimensionValue):
        transformer.pivot(data)


def _check_pagination(pivot_table: PivotTable, offset_rows: int = 1, limit_rows: int = 3) -> None:
    initial_columns = pivot_table.get_columns()
    initial_rows = list(pivot_table.get_rows())

    pivot_table.facade.paginate(offset_rows=offset_rows, limit_rows=limit_rows)

    paginated_columns = list(pivot_table.get_columns())
    paginated_rows = list(pivot_table.get_rows())

    assert paginated_columns == initial_columns

    if pivot_table.get_row_count() == 1:
        # A Dummy row (filled with None) is returned in this case
        assert paginated_rows == [
            DataRow(header=PivotHeader(), values=tuple(None for _ in range(len(paginated_columns))))
        ]
    else:
        assert paginated_rows == initial_rows[offset_rows : offset_rows + limit_rows]


def test_pivot_only_row_dims_multiple_measures():
    fid_city, fid_sales, fid_profit = ("123", "456", "789")
    liid_city, liid_mnames, liid_sales, liid_profit = (0, 1, 2, 3)
    legend_item_ids = (liid_city, liid_sales, liid_profit)

    data = [
        MergedQueryDataRow(data=("Detroit", 100, 10), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", 200, 20), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", 300, 30), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_mnames,
                obj=MeasureNameObjSpec(),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
            LegendItem(
                legend_item_id=liid_profit,
                obj=FieldObjSpec(id=fid_profit, title="Profit"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_city = 10
    piid_mnames = 20
    piid_sales = 30
    piid_profit = 40
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
            PivotLegendItem(
                pivot_item_id=piid_profit,
                legend_item_ids=[liid_profit],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Profit",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    actual_pivot_table = transformer.pivot(data)

    # Sort it
    actual_pivot_table.facade.sort()

    expected_pivot_columns = [PivotHeader()]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(DV(cells=(DC(100, liid_sales, piid_sales),)),),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(DV(cells=(DC(10, liid_profit, piid_profit),)),),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(DV(cells=(DC(300, liid_sales, piid_sales),)),),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(DV(cells=(DC(30, liid_profit, piid_profit),)),),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(DV(cells=(DC(200, liid_sales, piid_sales),)),),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(DV(cells=(DC(20, liid_profit, piid_profit),)),),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows

    _check_pagination(actual_pivot_table)


def test_pivot_only_column_dims_multiple_measures():
    fid_city, fid_sales, fid_profit = ("123", "456", "789")
    liid_city, liid_mnames, liid_sales, liid_profit = (0, 1, 2, 3)
    legend_item_ids = (liid_city, liid_sales, liid_profit)

    data = [
        MergedQueryDataRow(data=("Detroit", 100, 10), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", 200, 20), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", 300, 30), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_mnames,
                obj=MeasureNameObjSpec(),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
            LegendItem(
                legend_item_id=liid_profit,
                obj=FieldObjSpec(id=fid_profit, title="Profit"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_city = 10
    piid_mnames = 20
    piid_sales = 30
    piid_profit = 40
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
            PivotLegendItem(
                pivot_item_id=piid_profit,
                legend_item_ids=[liid_profit],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Profit",
            ),
        ],
    )

    # Pivot
    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    actual_pivot_table = transformer.pivot(data)

    # Sort it
    actual_pivot_table.facade.sort()

    expected_pivot_columns = [
        PivotHeader(
            values=(
                DV(cells=(DC("Detroit", liid_city, piid_city),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            )
        ),
        PivotHeader(
            values=(
                DV(cells=(DC("Detroit", liid_city, piid_city),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            )
        ),
        PivotHeader(
            values=(
                DV(cells=(DC("Moscow", liid_city, piid_city),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            )
        ),
        PivotHeader(
            values=(
                DV(cells=(DC("Moscow", liid_city, piid_city),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            )
        ),
        PivotHeader(
            values=(
                DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            )
        ),
        PivotHeader(
            values=(
                DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            )
        ),
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(),
            values=(
                DV(cells=(DC(100, liid_sales, piid_sales),)),
                DV(cells=(DC(10, liid_profit, piid_profit),)),
                DV(cells=(DC(300, liid_sales, piid_sales),)),
                DV(cells=(DC(30, liid_profit, piid_profit),)),
                DV(cells=(DC(200, liid_sales, piid_sales),)),
                DV(cells=(DC(20, liid_profit, piid_profit),)),
            ),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows

    # TODO: Check pagination


def test_pivot_only_row_dims_no_measures():
    fid_city = "123"
    liid_city = 0
    legend_item_ids = (liid_city,)

    data = [
        MergedQueryDataRow(data=("Detroit",), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco",), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow",), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
        ]
    )

    # Pivot item IDs
    piid_city = 10
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    actual_pivot_table = transformer.pivot(data)

    # Sort it
    actual_pivot_table.facade.sort()

    expected_pivot_columns = [PivotHeader()]  # One pseudo-column with empty header (no dims)
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("Detroit", liid_city, piid_city),)),)),
            values=(None,),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("Moscow", liid_city, piid_city),)),)),
            values=(None,),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("San Francisco", liid_city, piid_city),)),)),
            values=(None,),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows

    _check_pagination(actual_pivot_table)


def test_pivot_only_column_dims_no_measures():
    fid_city = "123"
    liid_city = 0
    legend_item_ids = (liid_city,)

    data = [
        MergedQueryDataRow(data=("Detroit",), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco",), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow",), legend_item_ids=legend_item_ids),
    ]

    # Legend item IDs
    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
        ]
    )

    # Pivot item IDs
    piid_city = 10
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="City",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    actual_pivot_table = transformer.pivot(data)

    # Sort it
    actual_pivot_table.facade.sort()

    expected_pivot_columns = [
        PivotHeader(values=(DV(cells=(DC("Detroit", liid_city, piid_city),)),)),
        PivotHeader(values=(DV(cells=(DC("Moscow", liid_city, piid_city),)),)),
        PivotHeader(values=(DV(cells=(DC("San Francisco", liid_city, piid_city),)),)),
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(  # One pseudo-row
            header=PivotHeader(),  # with no row-level dimensions
            values=(None, None, None),  # and values for pseudo-measure
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows

    _check_pagination(actual_pivot_table)

lib/dl_api_lib/dl_api_lib_tests/unit/pivot/test_corner_cases.py
datalens-tech/datalens-backend
from dl_api_lib.pivot.pandas.transformer import PdPivotTransformer
from dl_api_lib.pivot.primitives import (
    PivotHeaderRoleSpec,
    PivotHeaderValue,
    PivotMeasureSorting,
    PivotMeasureSortingSettings,
)
from dl_api_lib.pivot.primitives import DataCell as DC
from dl_api_lib.pivot.primitives import DataCellVector as DV
from dl_api_lib.pivot.primitives import MeasureNameValue as MNV
from dl_api_lib.query.formalization.pivot_legend import (
    PivotDimensionRoleSpec,
    PivotLegend,
    PivotLegendItem,
    PivotMeasureRoleSpec,
)
from dl_constants.enums import (
    FieldType,
    OrderDirection,
    PivotHeaderRole,
    PivotItemType,
    PivotRole,
    UserDataType,
)
from dl_query_processing.legend.field_legend import (
    FieldObjSpec,
    Legend,
    LegendItem,
    MeasureNameObjSpec,
)
from dl_query_processing.merging.primitives import MergedQueryDataRow


def test_paginate():
    fid_city, fid_ctgry, fid_sales, fid_profit = ("123", "456", "789", "000")
    liid_ctgry, liid_city, liid_mnames, liid_sales, liid_profit = (0, 1, 2, 3, 4)
    legend_item_ids = (liid_city, liid_ctgry, liid_sales, liid_profit)

    data = [
        MergedQueryDataRow(data=("Detroit", "Furniture", 100, 10), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Technology", 200, 20), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Furniture", 300, 30), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Office Supplies", 400, 40), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Office Supplies", 500, 50), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", "Office Supplies", 600, 60), legend_item_ids=legend_item_ids),
    ]

    # Legend item IDs
    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_mnames,
                obj=MeasureNameObjSpec(),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
            LegendItem(
                legend_item_id=liid_profit,
                obj=FieldObjSpec(id=fid_profit, title="Profit"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_city = 10
    piid_ctgry = 20
    piid_mnames = 30
    piid_sales = 40
    piid_profit = 50
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(
                    role=PivotRole.pivot_measure,
                    sorting=PivotMeasureSorting(
                        column=PivotMeasureSortingSettings(
                            header_values=[PivotHeaderValue(value="Furniture")],
                            direction=OrderDirection.desc,
                            role_spec=PivotHeaderRoleSpec(role=PivotHeaderRole.data),
                        ),
                        row=PivotMeasureSortingSettings(
                            header_values=[PivotHeaderValue(value="Moscow"), PivotHeaderValue(value="Sales")],
                            direction=OrderDirection.asc,
                            role_spec=PivotHeaderRoleSpec(role=PivotHeaderRole.data),
                        ),
                    ),
                ),
                title="Sales",
            ),
            PivotLegendItem(
                pivot_item_id=piid_profit,
                legend_item_ids=[liid_profit],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Profit",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    pivot_table = transformer.pivot(data)
    pivot_table.facade.sort()

    initial_columns = list(pivot_table.get_columns())
    initial_rows = list(pivot_table.get_rows())

    offset_rows = 1
    limit_rows = 3
    pivot_table.facade.paginate(offset_rows=offset_rows, limit_rows=limit_rows)

    paginated_columns = list(pivot_table.get_columns())
    paginated_rows = list(pivot_table.get_rows())

    assert paginated_columns == initial_columns
    assert paginated_rows == initial_rows[offset_rows : offset_rows + limit_rows]

    # check that paginator preserves headers' info
    sorting_headers = [
        (DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),),
        (
            DV(cells=(DC("Moscow", liid_city, piid_city),)),
            DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
        ),
    ]
    # (Moscow, Sales) header is paginated, so the check is a bit hacky
    for header, direction in zip(sorting_headers, [OrderDirection.desc, OrderDirection.asc]):
        assert pivot_table.facade._pivot_dframe.headers_info[header].sorting_direction == direction

lib/dl_api_lib/dl_api_lib_tests/unit/pivot/test_paginator.py
datalens-tech/datalens-backend
from dl_api_lib.pivot.hashable_packing import JsonHashableValuePacker
from dl_api_lib.pivot.pandas.transformer import PdPivotTransformer
from dl_api_lib.pivot.primitives import (
    PivotHeader,
    PivotHeaderInfo,
    PivotHeaderRoleSpec,
    PivotHeaderValue,
    PivotMeasureSorting,
    PivotMeasureSortingSettings,
)
from dl_api_lib.pivot.primitives import DataCell as DC
from dl_api_lib.pivot.primitives import DataCellVector as DV
from dl_api_lib.pivot.primitives import DataRow
from dl_api_lib.pivot.primitives import MeasureNameValue as MNV
from dl_api_lib.query.formalization.pivot_legend import (
    PivotDimensionRoleSpec,
    PivotLegend,
    PivotLegendItem,
    PivotMeasureRoleSpec,
)
from dl_constants.enums import (
    FieldType,
    OrderDirection,
    PivotHeaderRole,
    PivotItemType,
    PivotRole,
    UserDataType,
)
from dl_query_processing.legend.field_legend import (
    FieldObjSpec,
    Legend,
    LegendItem,
    MeasureNameObjSpec,
)
from dl_query_processing.merging.primitives import MergedQueryDataRow


def test_measure_sort_basic():
    fid_city, fid_ctgry, fid_sales = "123", "456", "789"
    liid_ctgry = 0
    liid_city = 1
    liid_mnames = 2
    liid_sales = 3
    legend_item_ids = (liid_city, liid_ctgry, liid_sales)

    data = [
        MergedQueryDataRow(data=("Detroit", "Furniture", 100), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Technology", 200), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Furniture", 300), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Office Supplies", 400), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Office Supplies", 500), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", "Office Supplies", 600), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_mnames,
                obj=MeasureNameObjSpec(),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_mnames = 10
    piid_city = 20
    piid_ctgry = 30
    piid_sales = 40
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column, direction=OrderDirection.desc),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row, direction=OrderDirection.asc),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(
                    role=PivotRole.pivot_measure,
                    sorting=PivotMeasureSorting(
                        column=PivotMeasureSortingSettings(
                            header_values=[PivotHeaderValue(value="Furniture")],
                            direction=OrderDirection.desc,
                            role_spec=PivotHeaderRoleSpec(role=PivotHeaderRole.data),
                        ),
                        row=PivotMeasureSortingSettings(
                            header_values=[PivotHeaderValue(value="Moscow")],
                            direction=OrderDirection.asc,
                            role_spec=PivotHeaderRoleSpec(role=PivotHeaderRole.data),
                        ),
                    ),
                ),
                title="Sales",
            ),
        ],
    )

    transformer = PdPivotTransformer(
        legend=legend,
        pivot_legend=pivot_legend,
        cell_packer=JsonHashableValuePacker(),
    )
    actual_pivot_table = transformer.pivot(data)
    actual_pivot_table.facade.sort()

    expected_pivot_columns = [
        PivotHeader(
            values=(DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),),
            info=PivotHeaderInfo(sorting_direction=OrderDirection.desc),
        ),
        PivotHeader(values=(DV(cells=(DC("Technology", liid_ctgry, piid_ctgry),)),)),
        PivotHeader(values=(DV(cells=(DC("Office Supplies", liid_ctgry, piid_ctgry),)),)),
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(300, liid_sales, piid_sales),)),
                None,
                DV(cells=(DC(400, liid_sales, piid_sales),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(100, liid_sales, piid_sales),)),
                None,
                DV(cells=(DC(600, liid_sales, piid_sales),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                ),
                info=PivotHeaderInfo(sorting_direction=OrderDirection.asc),
            ),
            values=(
                None,
                DV(cells=(DC(200, liid_sales, piid_sales),)),
                DV(cells=(DC(500, liid_sales, piid_sales),)),
            ),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def test_measure_sort_no_rows_dimension():
    fid_ctgry, fid_sales = "123", "456"
    liid_ctgry = 0
    liid_mnames = 1
    liid_sales = 2
    legend_item_ids = (liid_ctgry, liid_sales)

    data = [
        MergedQueryDataRow(data=("Furniture", 200), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Office Supplies", 400), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Technology", 100), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_mnames,
                obj=MeasureNameObjSpec(),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_mnames = 10
    piid_ctgry = 20
    piid_sales = 30
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(
                    role=PivotRole.pivot_measure,
                    sorting=PivotMeasureSorting(
                        column=PivotMeasureSortingSettings(
                            header_values=[PivotHeaderValue(value="Office Supplies"), PivotHeaderValue(value="Sales")],
                            direction=OrderDirection.asc,
                            role_spec=PivotHeaderRoleSpec(role=PivotHeaderRole.data),
                        ),
                        row=None,
                    ),
                ),
                title="Sales",
            ),
        ],
    )

    transformer = PdPivotTransformer(
        legend=legend,
        pivot_legend=pivot_legend,
        cell_packer=JsonHashableValuePacker(),
    )
    actual_pivot_table = transformer.pivot(data)
    actual_pivot_table.facade.sort()

    # data should be left intact, as "sorting" was performed by column with only one value
    mnames_dv = DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),))
    expected_pivot_columns = [
        PivotHeader(values=(DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)), mnames_dv)),
        PivotHeader(
            values=(DV(cells=(DC("Office Supplies", liid_ctgry, piid_ctgry),)), mnames_dv),
            info=PivotHeaderInfo(sorting_direction=OrderDirection.asc),
        ),
        PivotHeader(values=(DV(cells=(DC("Technology", liid_ctgry, piid_ctgry),)), mnames_dv)),
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(),
            values=(
                DV(cells=(DC(200, liid_sales, piid_sales),)),
                DV(cells=(DC(400, liid_sales, piid_sales),)),
                DV(cells=(DC(100, liid_sales, piid_sales),)),
            ),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def test_measure_sort_with_multiple_measures():
    fid_city, fid_ctgry, fid_sales, fid_profit = ("123", "456", "789", "000")
    liid_ctgry = 0
    liid_city = 1
    liid_mnames = 2
    liid_sales = 3
    liid_profit = 4
    legend_item_ids = (liid_city, liid_ctgry, liid_sales, liid_profit)

    data = [
        MergedQueryDataRow(data=("Detroit", "Furniture", 100, 10), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Technology", 200, 20), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Furniture", 300, 30), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Office Supplies", 400, 40), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Office Supplies", 500, 50), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", "Office Supplies", 600, 60), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_mnames,
                obj=MeasureNameObjSpec(),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
            LegendItem(
                legend_item_id=liid_profit,
                obj=FieldObjSpec(id=fid_profit, title="Profit"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_mnames = 10
    piid_city = 20
    piid_ctgry = 30
    piid_sales = 40
    piid_profit = 50
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row, direction=OrderDirection.desc),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(
                    role=PivotRole.pivot_measure,
                    sorting=PivotMeasureSorting(
                        column=None,
                        row=PivotMeasureSortingSettings(
                            header_values=[PivotHeaderValue(value="Moscow"), PivotHeaderValue(value="Sales")],
                            direction=OrderDirection.desc,
                            role_spec=PivotHeaderRoleSpec(role=PivotHeaderRole.data),
                        ),
                    ),
                ),
                title="Sales",
            ),
            PivotLegendItem(
                pivot_item_id=piid_profit,
                legend_item_ids=[liid_profit],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Profit",
            ),
        ],
    )

    transformer = PdPivotTransformer(
        legend=legend,
        pivot_legend=pivot_legend,
        cell_packer=JsonHashableValuePacker(),
    )
    actual_pivot_table = transformer.pivot(data)
    actual_pivot_table.facade.sort()

    expected_pivot_columns = [
        PivotHeader(values=(DV(cells=(DC("Office Supplies", liid_ctgry, piid_ctgry),)),)),
        PivotHeader(values=(DV(cells=(DC("Technology", liid_ctgry, piid_ctgry),)),)),
        PivotHeader(values=(DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),)),
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(400, liid_sales, piid_sales),)),
                None,
                DV(cells=(DC(300, liid_sales, piid_sales),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(40, liid_profit, piid_profit),)),
                None,
                DV(cells=(DC(30, liid_profit, piid_profit),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                ),
                info=PivotHeaderInfo(sorting_direction=OrderDirection.desc),
            ),
            values=(
                DV(cells=(DC(500, liid_sales, piid_sales),)),
                DV(cells=(DC(200, liid_sales, piid_sales),)),
                None,
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(50, liid_profit, piid_profit),)),
                DV(cells=(DC(20, liid_profit, piid_profit),)),
                None,
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(600, liid_sales, piid_sales),)),
                None,
                DV(cells=(DC(100, liid_sales, piid_sales),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(60, liid_profit, piid_profit),)),
                None,
                DV(cells=(DC(10, liid_profit, piid_profit),)),
            ),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def test_dimension_sort_with_multiple_measures():
    fid_city, fid_ctgry, fid_sales, fid_profit = ("123", "456", "789", "000")
    liid_ctgry = 0
    liid_city = 1
    liid_mnames = 2
    liid_sales = 3
    liid_profit = 4
    legend_item_ids = (liid_city, liid_ctgry, liid_sales, liid_profit)

    data = [
        MergedQueryDataRow(data=("Detroit", "Furniture", 100, 10), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Technology", 200, 20), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Furniture", 300, 30), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Office Supplies", 400, 40), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Office Supplies", 500, 50), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", "Office Supplies", 600, 60), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_mnames,
                obj=MeasureNameObjSpec(),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
            LegendItem(
                legend_item_id=liid_profit,
                obj=FieldObjSpec(id=fid_profit, title="Profit"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_mnames = 10
    piid_city = 20
    piid_ctgry = 30
    piid_sales = 40
    piid_profit = 50
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column, direction=OrderDirection.desc),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row, direction=OrderDirection.asc),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
            PivotLegendItem(
                pivot_item_id=piid_profit,
                legend_item_ids=[liid_profit],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Profit",
            ),
        ],
    )

    transformer = PdPivotTransformer(
        legend=legend,
        pivot_legend=pivot_legend,
        cell_packer=JsonHashableValuePacker(),
    )
    actual_pivot_table = transformer.pivot(data)
    actual_pivot_table.facade.sort()

    expected_pivot_columns = [
        # DESC
        PivotHeader(values=(DV(cells=(DC("Technology", liid_ctgry, piid_ctgry),)),)),
        PivotHeader(values=(DV(cells=(DC("Office Supplies", liid_ctgry, piid_ctgry),)),)),
        PivotHeader(values=(DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),)),
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        # ASC by City;
        # Measure Name should have the same order as in the legend ('Sales', 'Profit')
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                None,
                DV(cells=(DC(600, liid_sales, piid_sales),)),
                DV(cells=(DC(100, liid_sales, piid_sales),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                None,
                DV(cells=(DC(60, liid_profit, piid_profit),)),
                DV(cells=(DC(10, liid_profit, piid_profit),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(200, liid_sales, piid_sales),)),
                DV(cells=(DC(500, liid_sales, piid_sales),)),
                None,
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(20, liid_profit, piid_profit),)),
                DV(cells=(DC(50, liid_profit, piid_profit),)),
                None,
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                None,
                DV(cells=(DC(400, liid_sales, piid_sales),)),
                DV(cells=(DC(300, liid_sales, piid_sales),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                None,
                DV(cells=(DC(40, liid_profit, piid_profit),)),
                DV(cells=(DC(30, liid_profit, piid_profit),)),
            ),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def test_dimension_sort_mixed_case_strings():
    fid_thing, fid_string, fid_measure = ("123", "456", "789")
    liid_thing, liid_string, liid_measure = (0, 1, 2)
    legend_item_ids = (liid_thing, liid_string, liid_measure)

    data = [
        MergedQueryDataRow(data=("This", "Abc", None), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("This", "bCd", None), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("This", "ab", None), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("This", "CDE", None), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("This", "cd", None), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("This", "BC", None), legend_item_ids=legend_item_ids),
    ]

    # Legend item IDs
    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_thing,
                obj=FieldObjSpec(id=fid_thing, title="Thing"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_string,
                obj=FieldObjSpec(id=fid_string, title="String"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_measure,
                obj=FieldObjSpec(id=fid_measure, title="Measure"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_thing = 10
    piid_string = 20
    piid_measure = 30
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_thing,
                legend_item_ids=[liid_thing],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Thing",
            ),
            PivotLegendItem(
                pivot_item_id=piid_string,
                legend_item_ids=[liid_string],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row, direction=OrderDirection.asc),
                title="String",
            ),
            PivotLegendItem(
                pivot_item_id=piid_measure,
                legend_item_ids=[liid_measure],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Measure",
            ),
        ],
    )

    transformer = PdPivotTransformer(
        legend=legend,
        pivot_legend=pivot_legend,
        cell_packer=JsonHashableValuePacker(),
    )
    actual_pivot_table = transformer.pivot(data)
    actual_pivot_table.facade.sort()

    assert len(actual_pivot_table.get_columns()) == 1

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("ab", liid_string, piid_string),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("Abc", liid_string, piid_string),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("BC", liid_string, piid_string),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("bCd", liid_string, piid_string),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("cd", liid_string, piid_string),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("CDE", liid_string, piid_string),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def test_dimension_sort_stringified_numbers():
    fid_thing, fid_number, fid_measure = ("123", "456", "789")
    liid_thing, liid_number, liid_measure = (0, 1, 2)
    legend_item_ids = (liid_thing, liid_number, liid_measure)
    data = [
        MergedQueryDataRow(data=("This", "1", None), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("This", "24", None), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("This", "2", None), legend_item_ids=legend_item_ids),
        # None is a valid value in data from DB
        MergedQueryDataRow(data=("This", None, None), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("This", "12", None), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("This", "3", None), legend_item_ids=legend_item_ids),
    ]
    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_thing,
                obj=FieldObjSpec(id=fid_thing, title="Thing"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_number,
                obj=FieldObjSpec(id=fid_number, title="Number"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.integer,
            ),
            LegendItem(
                legend_item_id=liid_measure,
                obj=FieldObjSpec(id=fid_measure, title="Measure"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_thing = 10
    piid_number = 20
    piid_measure = 30
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_thing,
                legend_item_ids=[liid_thing],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Thing",
            ),
            PivotLegendItem(
                pivot_item_id=piid_number,
                legend_item_ids=[liid_number],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row, direction=OrderDirection.asc),
                title="Number",
            ),
            PivotLegendItem(
                pivot_item_id=piid_measure,
                legend_item_ids=[liid_measure],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Measure",
            ),
        ],
    )

    transformer = PdPivotTransformer(
        legend=legend,
        pivot_legend=pivot_legend,
        cell_packer=JsonHashableValuePacker(),
    )
    actual_pivot_table = transformer.pivot(data)
    actual_pivot_table.facade.sort()

    assert len(actual_pivot_table.get_columns()) == 1

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC(None, liid_number, piid_number),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("1", liid_number, piid_number),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("2", liid_number, piid_number),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("3", liid_number, piid_number),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("12", liid_number, piid_number),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("24", liid_number, piid_number),)),)),
            values=(DV(cells=(DC(None, liid_measure, piid_measure),)),),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows

lib/dl_api_lib/dl_api_lib_tests/unit/pivot/test_sorter.py
datalens-tech/datalens-backend
from dl_api_lib.pivot.hashable_packing import JsonHashableValuePacker
from dl_api_lib.pivot.hashable_packing import JsonWrapper as JW
from dl_api_lib.pivot.primitives import DataCell as DC
from dl_api_lib.pivot.primitives import DataCellVector as DV
from dl_api_lib.pivot.primitives import MeasureNameValue as MNV
from dl_api_lib.pivot.stream_modifiers import (
    DataCellConverter,
    MeasureDataTransposer,
    TransposedDataRow,
)
from dl_api_lib.query.formalization.pivot_legend import (
    PivotAnnotationRoleSpec,
    PivotDimensionRoleSpec,
    PivotLegend,
    PivotLegendItem,
    PivotMeasureRoleSpec,
)
from dl_constants.enums import (
    FieldRole,
    FieldType,
    PivotItemType,
    PivotRole,
    UserDataType,
)
from dl_query_processing.legend.field_legend import (
    FieldObjSpec,
    Legend,
    LegendItem,
    RowRoleSpec,
)
from dl_query_processing.merging.primitives import MergedQueryDataRow


def test_data_cell_converter():
    liid_ctgry, liid_city, liid_sales, liid_profit = (0, 1, 2, 3)
    legend_item_ids = (liid_city, liid_ctgry, liid_sales, liid_profit)
    data = [
        MergedQueryDataRow(data=("Detroit", "Furniture", 100, 10), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Portland", "Furniture", 200, 20), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Plants", 300, 30), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", "Plants", 400, 40), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                role_spec=RowRoleSpec(role=FieldRole.row),
                obj=FieldObjSpec(id="1", title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                role_spec=RowRoleSpec(role=FieldRole.row),
                obj=FieldObjSpec(id="1", title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                role_spec=RowRoleSpec(role=FieldRole.row),
                obj=FieldObjSpec(id="1", title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.float,
            ),
            LegendItem(
                legend_item_id=liid_profit,
                role_spec=RowRoleSpec(role=FieldRole.row),
                obj=FieldObjSpec(id="1", title="Profit"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.float,
            ),
        ],
    )

    # Pivot item IDs
    piid_ctgry = 10
    piid_city = 20
    piid_sales = 30
    piid_profit = 40
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
            PivotLegendItem(
                pivot_item_id=piid_profit,
                legend_item_ids=[liid_profit],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Profit",
            ),
        ],
    )

    dcell_stream = DataCellConverter(
        rows=data,
        cell_packer=JsonHashableValuePacker(),
        legend=legend,
        pivot_legend=pivot_legend,
    )
    converted_data = list(dcell_stream)
    expected_data = [
        [
            DC("Detroit", liid_city, piid_city),
            DC("Furniture", liid_ctgry, piid_ctgry),
            DC(100, liid_sales, piid_sales),
            DC(10, liid_profit, piid_profit),
        ],
        [
            DC("Portland", liid_city, piid_city),
            DC("Furniture", liid_ctgry, piid_ctgry),
            DC(200, liid_sales, piid_sales),
            DC(20, liid_profit, piid_profit),
        ],
        [
            DC("Moscow", liid_city, piid_city),
            DC("Plants", liid_ctgry, piid_ctgry),
            DC(300, liid_sales, piid_sales),
            DC(30, liid_profit, piid_profit),
        ],
        [
            DC("Detroit", liid_city, piid_city),
            DC("Plants", liid_ctgry, piid_ctgry),
            DC(400, liid_sales, piid_sales),
            DC(40, liid_profit, piid_profit),
        ],
    ]
    assert converted_data == expected_data


def test_measure_data_transposer():
    # Legend item IDs
    liid_mnames = 0
    liid_city = 1
    liid_ctgry = 2
    liid_sales = 3
    liid_profit = 4

    # Pivot item IDs
    piid_mnames = 10
    piid_city = 20
    piid_ctgry = 30
    piid_sales = 40
    piid_profit = 50

    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
            PivotLegendItem(
                pivot_item_id=piid_profit,
                legend_item_ids=[liid_profit],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Profit",
            ),
        ],
    )

    dcell_stream = [
        [
            DC("Detroit", liid_city, piid_city),
            DC("Furniture", liid_ctgry, piid_ctgry),
            DC(100, liid_sales, piid_sales),
            DC(10, liid_profit, piid_profit),
        ],
        [
            DC("Portland", liid_city, piid_city),
            DC("Furniture", liid_ctgry, piid_ctgry),
            DC(200, liid_sales, piid_sales),
            DC(20, liid_profit, piid_profit),
        ],
        [
            DC("Moscow", liid_city, piid_city),
            DC("Plants", liid_ctgry, piid_ctgry),
            DC(300, liid_sales, piid_sales),
            DC(30, liid_profit, piid_profit),
        ],
        [
            DC("Detroit", liid_city, piid_city),
            DC("Plants", liid_ctgry, piid_ctgry),
            DC(400, liid_sales, piid_sales),
            DC(40, liid_profit, piid_profit),
        ],
    ]

    transposed_stream = MeasureDataTransposer(
        dcell_stream=dcell_stream,
        pivot_legend=pivot_legend,
        cell_packer=JsonHashableValuePacker(),
    )
    transposed_data = list(transposed_stream)
    expected_data = [
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Detroit", liid_city, piid_city),)),
                DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            ),
            value=DV(cells=(DC(100, liid_sales, piid_sales),)),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Detroit", liid_city, piid_city),)),
                DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            ),
            value=DV(cells=(DC(10, liid_profit, piid_profit),)),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Portland", liid_city, piid_city),)),
                DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            ),
            value=DV(cells=(DC(200, liid_sales, piid_sales),)),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Portland", liid_city, piid_city),)),
                DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            ),
            value=DV(cells=(DC(20, liid_profit, piid_profit),)),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Moscow", liid_city, piid_city),)),
                DV(cells=(DC("Plants", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            ),
            value=DV(cells=(DC(300, liid_sales, piid_sales),)),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Moscow", liid_city, piid_city),)),
                DV(cells=(DC("Plants", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            ),
            value=DV(cells=(DC(30, liid_profit, piid_profit),)),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Detroit", liid_city, piid_city),)),
                DV(cells=(DC("Plants", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            ),
            value=DV(cells=(DC(400, liid_sales, piid_sales),)),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Detroit", liid_city, piid_city),)),
                DV(cells=(DC("Plants", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            ),
            value=DV(cells=(DC(40, liid_profit, piid_profit),)),
        ),
    ]
    assert transposed_data == expected_data


def test_measure_data_transposer_with_annotations():
    # Legend item IDs:
    liid_mnames = 0
    liid_city = 1
    liid_ctgry = 2
    liid_sales = 3
    liid_profit = 4
    liid_customers = 5
    liid_orders = 6

    # Pivot item IDs
    piid_mnames = 10
    piid_city = 20
    piid_ctgry = 30
    piid_sales = 40
    piid_profit = 50
    piid_customers = 60
    piid_orders = 70
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
            PivotLegendItem(
                pivot_item_id=piid_profit,
                legend_item_ids=[liid_profit],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Profit",
            ),
            PivotLegendItem(
                pivot_item_id=piid_customers,
                legend_item_ids=[liid_customers],
                role_spec=PivotAnnotationRoleSpec(role=PivotRole.pivot_annotation, annotation_type="color"),
                title="Customers",
            ),
            PivotLegendItem(
                pivot_item_id=piid_orders,
                legend_item_ids=[liid_orders],
                role_spec=PivotAnnotationRoleSpec(
                    role=PivotRole.pivot_annotation,
                    annotation_type="color",
                    target_legend_item_ids=[liid_profit],
                ),
                title="Orders",
            ),
        ],
    )
    dcell_stream = [
        [
            DC("Detroit", liid_city, piid_city),
            DC("Furniture", liid_ctgry, piid_ctgry),
            DC(100, liid_sales, piid_sales),
            DC(10, liid_profit, piid_profit),
            DC(3, liid_customers, piid_customers),
            DC(4, liid_orders, piid_orders),
        ],
        [
            DC("Portland", liid_city, piid_city),
            DC("Furniture", liid_ctgry, piid_ctgry),
            DC(200, liid_sales, piid_sales),
            DC(20, liid_profit, piid_profit),
            DC(5, liid_customers, piid_customers),
            DC(6, liid_orders, piid_orders),
        ],
        [
            DC("Moscow", liid_city, piid_city),
            DC("Plants", liid_ctgry, piid_ctgry),
            DC(300, liid_sales, piid_sales),
            DC(30, liid_profit, piid_profit),
            DC(7, liid_customers, piid_customers),
            DC(10, liid_orders, piid_orders),
        ],
        [
            DC("Detroit", liid_city, piid_city),
            DC("Plants", liid_ctgry, piid_ctgry),
            DC(400, liid_sales, piid_sales),
            DC(40, liid_profit, piid_profit),
            DC(9, liid_customers, piid_customers),
            DC(23, liid_orders, piid_orders),
        ],
    ]
    transposed_stream = MeasureDataTransposer(
        dcell_stream=dcell_stream,
        pivot_legend=pivot_legend,
        cell_packer=JsonHashableValuePacker(),
    )
    transposed_data = list(transposed_stream)
    expected_data = [
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Detroit", liid_city, piid_city),)),
                DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            ),
            value=DV(
                cells=(
                    DC(100, liid_sales, piid_sales),
                    DC(3, liid_customers, piid_customers),
                )
            ),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Detroit", liid_city, piid_city),)),
                DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            ),
            value=DV(
                cells=(
                    DC(10, liid_profit, piid_profit),
                    DC(3, liid_customers, piid_customers),
                    DC(4, liid_orders, piid_orders),
                )
            ),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Portland", liid_city, piid_city),)),
                DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            ),
            value=DV(
                cells=(
                    DC(200, liid_sales, piid_sales),
                    DC(5, liid_customers, piid_customers),
                )
            ),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Portland", liid_city, piid_city),)),
                DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            ),
            value=DV(
                cells=(
                    DC(20, liid_profit, piid_profit),
                    DC(5, liid_customers, piid_customers),
                    DC(6, liid_orders, piid_orders),
                )
            ),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Moscow", liid_city, piid_city),)),
                DV(cells=(DC("Plants", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            ),
            value=DV(
                cells=(
                    DC(300, liid_sales, piid_sales),
                    DC(7, liid_customers, piid_customers),
                )
            ),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Moscow", liid_city, piid_city),)),
                DV(cells=(DC("Plants", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            ),
            value=DV(
                cells=(
                    DC(30, liid_profit, piid_profit),
                    DC(7, liid_customers, piid_customers),
                    DC(10, liid_orders, piid_orders),
                )
            ),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Detroit", liid_city, piid_city),)),
                DV(cells=(DC("Plants", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
            ),
            value=DV(
                cells=(
                    DC(400, liid_sales, piid_sales),
                    DC(9, liid_customers, piid_customers),
                )
            ),
        ),
        TransposedDataRow(
            dimensions=(
                DV(cells=(DC("Detroit", liid_city, piid_city),)),
                DV(cells=(DC("Plants", liid_ctgry, piid_ctgry),)),
                DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
            ),
            value=DV(
                cells=(
                    DC(40, liid_profit, piid_profit),
                    DC(9, liid_customers, piid_customers),
                    DC(23, liid_orders, piid_orders),
                ),
            ),
        ),
    ]
    assert transposed_data == expected_data


def test_data_cell_converter_with_dicts():
    liid_city = 0
    liid_value = 1
    legend_item_ids = (liid_city, liid_value)
    data = [
        MergedQueryDataRow(data=("Detroit", {"a": 1, "f": 6}), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Portland", {"c": 2, "s": 7}), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", {"z": 3, "q": {"some": "thing"}}), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", {"y": 4, "k": 100500}), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_city,
                role_spec=RowRoleSpec(role=FieldRole.row),
                obj=FieldObjSpec(id="1", title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
        ],
    )

    piid_city = 10
    piid_value = 20
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_value,
                legend_item_ids=[liid_value],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Value",
            ),
        ],
    )

    dcell_stream = DataCellConverter(
        rows=data,
        cell_packer=JsonHashableValuePacker(),
        legend=legend,
        pivot_legend=pivot_legend,
    )
    converted_data = list(dcell_stream)
    expected_data = [
        [
            DC("Detroit", liid_city, piid_city),
            DC(JW('{"a": 1, "f": 6}'), liid_value, piid_value),
        ],
        [
            DC("Portland", liid_city, piid_city),
            DC(JW('{"c": 2, "s": 7}'), liid_value, piid_value),
        ],
        [
            DC("Moscow", liid_city, piid_city),
            DC(JW('{"q": {"some": "thing"}, "z": 3}'), liid_value, piid_value),
        ],
        [
            DC("Detroit", liid_city, piid_city),
            DC(JW('{"k": 100500, "y": 4}'), liid_value, piid_value),
        ],
    ]
    assert converted_data == expected_data

lib/dl_api_lib/dl_api_lib_tests/unit/pivot/test_stream_modifers.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib.pivot.pandas.transformer import PdPivotTransformer
from dl_api_lib.pivot.primitives import DataCell as DC
from dl_api_lib.pivot.primitives import DataCellVector as DV
from dl_api_lib.pivot.primitives import DataRow
from dl_api_lib.pivot.primitives import MeasureNameValue as MNV
from dl_api_lib.pivot.primitives import PivotHeader
from dl_api_lib.query.formalization.pivot_legend import (
    PivotAnnotationRoleSpec,
    PivotDimensionRoleSpec,
    PivotLegend,
    PivotLegendItem,
    PivotMeasureRoleSpec,
)
from dl_constants.enums import (
    FieldType,
    PivotItemType,
    PivotRole,
    UserDataType,
)
from dl_query_processing.legend.field_legend import (
    FieldObjSpec,
    Legend,
    LegendItem,
    MeasureNameObjSpec,
)
from dl_query_processing.merging.primitives import MergedQueryDataRow


def test_pivot_single_measure():
    fid_city, fid_ctgry, fid_sales = ("123", "456", "789")
    liid_city, liid_ctgry, liid_sales = (0, 1, 2)
    legend_item_ids = (liid_city, liid_ctgry, liid_sales)

    data = [
        MergedQueryDataRow(data=("Detroit", "Furniture", 100), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Furniture", 200), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Office Supplies", 300), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", "Office Supplies", 400), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    piid_city, piid_ctgry, piid_sales = (10, 20, 30)
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    actual_pivot_table = transformer.pivot(data)

    expected_pivot_columns = [
        PivotHeader(values=(DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),)),
        PivotHeader(values=(DV(cells=(DC("Office Supplies", liid_ctgry, piid_ctgry),)),)),
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("Detroit", liid_city, piid_city),)),)),
            values=(DV(cells=(DC(100, liid_sales, piid_sales),)), DV(cells=(DC(400, liid_sales, piid_sales),))),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("Moscow", liid_city, piid_city),)),)),
            values=(None, DV(cells=(DC(300, liid_sales, piid_sales),))),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("San Francisco", liid_city, piid_city),)),)),
            values=(DV(cells=(DC(200, liid_sales, piid_sales),)), None),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def test_pivot_no_measures():
    fid_city, fid_ctgry = ("123", "456")
    liid_city, liid_ctgry = (0, 1)
    legend_item_ids = (liid_city, liid_ctgry)

    data = [
        MergedQueryDataRow(data=("Detroit", "Furniture"), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Furniture"), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Office Supplies"), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", "Office Supplies"), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
        ]
    )

    piid_city, piid_ctgry = (10, 20)
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    actual_pivot_table = transformer.pivot(data)

    expected_pivot_columns = [
        PivotHeader(values=(DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),)),
        PivotHeader(values=(DV(cells=(DC("Office Supplies", liid_ctgry, piid_ctgry),)),)),
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("Detroit", liid_city, piid_city),)),)),
            values=(None, None),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("Moscow", liid_city, piid_city),)),)),
            values=(None, None),
        ),
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC("San Francisco", liid_city, piid_city),)),)),
            values=(None, None),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def test_pivot_multiple_measures():
    fid_city, fid_ctgry, fid_sales, fid_profit = ("123", "456", "789", "000")
    liid_city, liid_ctgry, liid_mnames, liid_sales, liid_profit = (0, 1, 2, 3, 4)
    legend_item_ids = (liid_city, liid_ctgry, liid_sales, liid_profit)

    data = [
        MergedQueryDataRow(data=("Detroit", "Furniture", 100, 10), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Furniture", 200, 20), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Office Supplies", 300, 30), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", "Office Supplies", 400, 40), legend_item_ids=legend_item_ids),
    ]

    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_mnames,
                obj=MeasureNameObjSpec(),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
            LegendItem(
                legend_item_id=liid_profit,
                obj=FieldObjSpec(id=fid_profit, title="Profit"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_city = 10
    piid_ctgry = 20
    piid_mnames = 30
    piid_sales = 40
    piid_profit = 50
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
            PivotLegendItem(
                pivot_item_id=piid_profit,
                legend_item_ids=[liid_profit],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Profit",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    actual_pivot_table = transformer.pivot(data)

    expected_pivot_columns = [
        PivotHeader(values=(DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),)),
        PivotHeader(values=(DV(cells=(DC("Office Supplies", liid_ctgry, piid_ctgry),)),)),
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(10, liid_profit, piid_profit),)),
                DV(cells=(DC(40, liid_profit, piid_profit),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(100, liid_sales, piid_sales),)),
                DV(cells=(DC(400, liid_sales, piid_sales),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                None,
                DV(cells=(DC(30, liid_profit, piid_profit),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                None,
                DV(cells=(DC(300, liid_sales, piid_sales),)),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(20, liid_profit, piid_profit),)),
                None,
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(cells=(DC(200, liid_sales, piid_sales),)),
                None,
            ),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def test_pivot_multiple_measures_and_annotations():
    fid_city, fid_ctgry, fid_sales, fid_profit, fid_customers, fid_orders = (
        "123",
        "456",
        "789",
        "000",
        "222",
        "333",
    )
    # Legend item IDs:
    liid_mnames = 0
    liid_city = 1
    liid_ctgry = 2
    liid_sales = 3
    liid_profit = 4
    liid_customers = 5
    liid_orders = 6
    legend_item_ids = (liid_city, liid_ctgry, liid_sales, liid_profit, liid_customers, liid_orders)
    data = [
        MergedQueryDataRow(data=("Detroit", "Furniture", 100, 10, 3, 4), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("San Francisco", "Furniture", 200, 20, 5, 7), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Moscow", "Office Supplies", 300, 30, 4, 4), legend_item_ids=legend_item_ids),
        MergedQueryDataRow(data=("Detroit", "Office Supplies", 400, 40, 6, 10), legend_item_ids=legend_item_ids),
    ]
    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_mnames,
                obj=MeasureNameObjSpec(),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
            LegendItem(
                legend_item_id=liid_profit,
                obj=FieldObjSpec(id=fid_profit, title="Profit"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
            LegendItem(
                legend_item_id=liid_customers,
                obj=FieldObjSpec(id=fid_customers, title="Customers"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
            LegendItem(
                legend_item_id=liid_orders,
                obj=FieldObjSpec(id=fid_orders, title="Orders"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    # Pivot item IDs
    piid_mnames = 10
    piid_city = 20
    piid_ctgry = 30
    piid_sales = 40
    piid_profit = 50
    piid_customers = 60
    piid_orders = 70
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_mnames,
                item_type=PivotItemType.measure_name,
                legend_item_ids=[liid_mnames],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="Measure Name",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
            PivotLegendItem(
                pivot_item_id=piid_profit,
                legend_item_ids=[liid_profit],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Profit",
            ),
            PivotLegendItem(
                pivot_item_id=piid_customers,
                legend_item_ids=[liid_customers],
                role_spec=PivotAnnotationRoleSpec(role=PivotRole.pivot_annotation, annotation_type="color"),
                title="Customers",
            ),
            PivotLegendItem(
                pivot_item_id=piid_orders,
                legend_item_ids=[liid_orders],
                role_spec=PivotAnnotationRoleSpec(
                    role=PivotRole.pivot_annotation,
                    annotation_type="color",
                    target_legend_item_ids=[liid_profit],
                ),
                title="Orders",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    actual_pivot_table = transformer.pivot(data)

    expected_pivot_columns = [
        PivotHeader(values=(DV(cells=(DC("Furniture", liid_ctgry, piid_ctgry),)),)),
        PivotHeader(values=(DV(cells=(DC("Office Supplies", liid_ctgry, piid_ctgry),)),)),
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(
                    cells=(
                        DC(10, liid_profit, piid_profit),
                        DC(3, liid_customers, piid_customers),
                        DC(4, liid_orders, piid_orders),
                    )
                ),
                DV(
                    cells=(
                        DC(40, liid_profit, piid_profit),
                        DC(6, liid_customers, piid_customers),
                        DC(10, liid_orders, piid_orders),
                    )
                ),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Detroit", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(
                    cells=(
                        DC(100, liid_sales, piid_sales),
                        DC(3, liid_customers, piid_customers),
                    )
                ),
                DV(
                    cells=(
                        DC(400, liid_sales, piid_sales),
                        DC(6, liid_customers, piid_customers),
                    )
                ),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                None,
                DV(
                    cells=(
                        DC(30, liid_profit, piid_profit),
                        DC(4, liid_customers, piid_customers),
                        DC(4, liid_orders, piid_orders),
                    )
                ),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("Moscow", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                None,
                DV(
                    cells=(
                        DC(300, liid_sales, piid_sales),
                        DC(4, liid_customers, piid_customers),
                    )
                ),
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Profit", piid_profit), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(
                    cells=(
                        DC(20, liid_profit, piid_profit),
                        DC(5, liid_customers, piid_customers),
                        DC(7, liid_orders, piid_orders),
                    )
                ),
                None,
            ),
        ),
        DataRow(
            header=PivotHeader(
                values=(
                    DV(cells=(DC("San Francisco", liid_city, piid_city),)),
                    DV(cells=(DC(MNV("Sales", piid_sales), liid_mnames, piid_mnames),)),
                )
            ),
            values=(
                DV(
                    cells=(
                        DC(200, liid_sales, piid_sales),
                        DC(5, liid_customers, piid_customers),
                    )
                ),
                None,
            ),
        ),
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def check_diagonal_pivot_table(row_cnt: int):
    fid_city, fid_ctgry, fid_sales = ("123", "456", "789")
    liid_ctgry, liid_city, liid_sales = (0, 1, 2)
    legend_item_ids = (liid_city, liid_ctgry, liid_sales)

    def city_gen(num):
        return f"City {num:020}"

    def cat_gen(num):
        return f"Cat no. {num:020}"

    data = [
        MergedQueryDataRow(data=(city_gen(num), cat_gen(num), num), legend_item_ids=legend_item_ids)
        for num in range(row_cnt)
    ]
    legend = Legend(
        items=[
            LegendItem(
                legend_item_id=liid_ctgry,
                obj=FieldObjSpec(id=fid_ctgry, title="Category"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_city,
                obj=FieldObjSpec(id=fid_city, title="City"),
                field_type=FieldType.DIMENSION,
                data_type=UserDataType.string,
            ),
            LegendItem(
                legend_item_id=liid_sales,
                obj=FieldObjSpec(id=fid_sales, title="Sales"),
                field_type=FieldType.MEASURE,
                data_type=UserDataType.integer,
            ),
        ]
    )

    piid_city, piid_ctgry, piid_sales = (10, 20, 30)
    pivot_legend = PivotLegend(
        items=[
            PivotLegendItem(
                pivot_item_id=piid_ctgry,
                legend_item_ids=[liid_ctgry],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_column),
                title="Category",
            ),
            PivotLegendItem(
                pivot_item_id=piid_city,
                legend_item_ids=[liid_city],
                role_spec=PivotDimensionRoleSpec(role=PivotRole.pivot_row),
                title="City",
            ),
            PivotLegendItem(
                pivot_item_id=piid_sales,
                legend_item_ids=[liid_sales],
                role_spec=PivotMeasureRoleSpec(role=PivotRole.pivot_measure),
                title="Sales",
            ),
        ],
    )

    transformer = PdPivotTransformer(legend=legend, pivot_legend=pivot_legend)
    actual_pivot_table = transformer.pivot(data)

    expected_pivot_columns = [
        PivotHeader(values=(DV(cells=(DC(cat_gen(col_num), liid_ctgry, piid_ctgry),)),)) for col_num in range(row_cnt)
    ]
    actual_pivot_columns = actual_pivot_table.get_columns()
    assert actual_pivot_columns == expected_pivot_columns

    expected_pivot_rows = [
        DataRow(
            header=PivotHeader(values=(DV(cells=(DC(city_gen(row_num), liid_city, piid_city),)),)),
            values=(
                *[None for _ in range(row_num)],
                DV(cells=(DC(row_num, liid_sales, piid_sales),)),
                *[None for _ in range(row_num + 1, row_cnt)],
            ),
        )
        for row_num in range(row_cnt)
    ]
    actual_pivot_rows = list(actual_pivot_table.get_rows())
    assert actual_pivot_rows == expected_pivot_rows


def test_diagonal_pivot():
    check_diagonal_pivot_table(100)


@pytest.mark.skip  # too much memory in CI
def test_pivot_overflow():
    with pytest.raises(MemoryError):
        check_diagonal_pivot_table(50000)

lib/dl_api_lib/dl_api_lib_tests/unit/pivot/test_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.exc import (
    DEFAULT_ERR_CODE_API_PREFIX,
    GLOBAL_ERR_PREFIX,
)
import dl_core.exc
import dl_formula.core.exc
import dl_query_processing.exc


def exc_cls_to_code(exc_cls):
    exc_cls_code = (
        getattr(exc_cls, "err_code", None)
        # formula excs.
        # TODO: document how it is built or use a common function.
        or getattr(exc_cls, "default_code", [])[1:]
        or []
    )
    return ".".join((GLOBAL_ERR_PREFIX, DEFAULT_ERR_CODE_API_PREFIX) + tuple(exc_cls_code))


# Error codes referenced in docs:
# https://cloud.yandex.ru/docs/datalens/troubleshooting/errors/all
STABILIZED_CODES = (
    (dl_core.exc.CannotParseDateTime, "ERR.DS_API.DB.CANNOT_PARSE.DATETIME"),
    (dl_core.exc.CannotParseNumber, "ERR.DS_API.DB.CANNOT_PARSE.NUMBER"),
    (dl_core.exc.ColumnDoesNotExist, "ERR.DS_API.DB.COLUMN_DOES_NOT_EXIST"),
    (dl_core.exc.DatabaseQueryError, "ERR.DS_API.DB"),
    (dl_core.exc.DatabaseUnavailable, "ERR.DS_API.DATABASE_UNAVAILABLE"),
    (dl_core.exc.DbMemoryLimitExceeded, "ERR.DS_API.DB.MEMORY_LIMIT_EXCEEDED"),
    (dl_core.exc.InvalidQuery, "ERR.DS_API.DB.INVALID_QUERY"),
    (dl_core.exc.JoinColumnTypeMismatch, "ERR.DS_API.DB.JOIN_COLUMN_TYPE_MISMATCH"),
    (dl_core.exc.MaterializationNotFinished, "ERR.DS_API.DB.MATERIALIZATION_NOT_FINISHED"),
    (dl_core.exc.ResultRowCountLimitExceeded, "ERR.DS_API.ROW_COUNT_LIMIT"),
    (dl_core.exc.SourceDoesNotExist, "ERR.DS_API.DB.SOURCE_DOES_NOT_EXIST"),
    (dl_core.exc.UnexpectedInfOrNan, "ERR.DS_API.DB.UNEXPECTED_INF_OR_NAN"),
    (dl_core.exc.TableNameNotConfiguredError, "ERR.DS_API.SOURCE_CONFIG.TABLE_NOT_CONFIGURED"),
    (dl_core.exc.USBadRequestException, "ERR.DS_API.US.BAD_REQUEST"),
    (dl_core.exc.USAlreadyExistsException, "ERR.DS_API.US.BAD_REQUEST.ALREADY_EXISTS"),
    (dl_core.exc.USObjectNotFoundException, "ERR.DS_API.US.OBJ_NOT_FOUND"),
    (dl_query_processing.exc.DLFormulaError, "ERR.DS_API.FORMULA"),
    (dl_core.exc.FieldNotFound, "ERR.DS_API.FIELD.NOT_FOUND"),
    (dl_query_processing.exc.FilterValueError, "ERR.DS_API.FILTER.INVALID_VALUE"),
    (dl_formula.core.exc.DoubleAggregationError, "ERR.DS_API.VALIDATION.AGG.DOUBLE"),
    (dl_formula.core.exc.InconsistentAggregationError, "ERR.DS_API.VALIDATION.AGG.INCONSISTENT"),
    (dl_formula.core.exc.WindowFunctionWOAggregationError, "ERR.DS_API.VALIDATION.WIN_FUNC.NO_AGG"),
)


def test_exc_code_match():
    for exc_cls, expected_code in STABILIZED_CODES:
        actual_code = exc_cls_to_code(exc_cls)
        assert actual_code == expected_code, exc_cls

lib/dl_api_lib/dl_api_lib_tests/unit/test_exc_codes.py
datalens-tech/datalens-backend
from typing import Any

import attr
import pytest

from dl_api_lib.api_common.update_dataset_mutation_key import UpdateDatasetMutationKey
from dl_api_lib.request_model.data import FieldAction
from dl_api_lib.schemas.action import ActionSchema


TEST_DATASET_REVISION_ID = "123"
TEST_DATASET_ANOTHER_REVISION_ID = "321"


def test_empty_mutation_keys():
    mutation_key_empty = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, [])
    mutation_key_empty_copy = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, [])
    mutation_key_empty_different_dataset = UpdateDatasetMutationKey.create(TEST_DATASET_ANOTHER_REVISION_ID, [])
    assert mutation_key_empty.get_hash() == mutation_key_empty_copy.get_hash()
    assert mutation_key_empty.get_hash() != mutation_key_empty_different_dataset.get_hash()


@pytest.fixture
def original_mutation_list() -> list[dict[str, Any]]:
    return [
        {
            "action": "add_field",
            "field": {
                "formula": "sum([Discount])*2",
                "cast": "float",
                "description": "",
                "source": "",
                "hidden": False,
                "aggregation": "none",
                "title": "title-11e2dd40-6739-11ec-9e02-a5fd2b33324c",
                "guid": "test",
                "calc_mode": "formula",
            },
            "debug_info": "merged-update",
        },
        {
            "action": "add_field",
            "field": {
                "valid": True,
                "calc_mode": "parameter",
                "source": "",
                "description": "",
                "title": "datetime_param2",
                "type": "DIMENSION",
                "aggregation": "none",
                "hidden": False,
                "guid": "datetime_param2",
                "cast": "datetime",
                "default_value": "2022-05-18T12:33:21.000+03:00",
            },
        },
    ]


@pytest.fixture
def original_mutation(original_mutation_list) -> list[FieldAction]:
    return ActionSchema(many=True).load(original_mutation_list)


@pytest.fixture
def original_mutation_shuffle_dicts(original_mutation_list) -> list[FieldAction]:
    return ActionSchema(many=True).load(original_mutation_list[::-1])


@pytest.fixture
def original_mutation_shuffle_keys(original_mutation_list) -> list[FieldAction]:
    shuffled = []
    for update in original_mutation_list:
        update["field"] = {k: v for k, v in list(update["field"].items())[::-1]}
        shuffled.append({k: v for k, v in list(update.items())[::-1]})
    return ActionSchema(many=True).load(shuffled)


def test_reordered_mutation_keys(
    original_mutation,
    original_mutation_shuffle_dicts,
    original_mutation_shuffle_keys,
):
    mutation_key_original = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, original_mutation)
    mutation_key_shuffled_1 = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, original_mutation_shuffle_dicts)
    mutation_key_shuffled_2 = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, original_mutation_shuffle_keys)
    assert mutation_key_original.get_hash() == mutation_key_shuffled_1.get_hash()
    assert mutation_key_original.get_hash() == mutation_key_shuffled_2.get_hash()


@pytest.fixture
def mutation_add_one_dict(original_mutation_list) -> list[FieldAction]:
    return ActionSchema(many=True).load(
        original_mutation_list
        + [
            {
                "action": "add_field",
                "field": {
                    "calc_mode": "formula",
                    "aggregation": "none",
                    "title": "title-1234d4a0-9309-11ec-9b5b-91d1cff8b10b",
                    "guid": "test",
                    "cast": "integer",
                    "hidden": True,
                    "description": "",
                    "source": "",
                    "formula": "[Profit]*2",
                },
            }
        ]
    )


@pytest.fixture
def mutation_add_one_key(original_mutation_list) -> list[FieldAction]:
    original_mutation_list[1]["field"]["avatar_id"] = "132132"
    return ActionSchema(many=True).load(original_mutation_list)


@pytest.fixture
def mutation_different_value(original_mutation_list) -> list[FieldAction]:
    original_mutation_list[0]["field"]["description"] = "extra description"
    return ActionSchema(many=True).load(original_mutation_list)


def test_cached_updates_mutation_keys(
    original_mutation,
    mutation_add_one_dict,
    mutation_add_one_key,
    mutation_different_value,
):
    mutation_key_original = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, original_mutation)
    mutation_key_more_dicts = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, mutation_add_one_dict)
    mutation_key_more_keys = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, mutation_add_one_key)
    mutation_key_different_value = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, mutation_different_value)
    assert mutation_key_original.get_hash() != mutation_key_more_dicts.get_hash()
    assert mutation_key_original.get_hash() != mutation_key_more_keys.get_hash()
    assert mutation_key_original.get_hash() != mutation_key_different_value.get_hash()


def test_mutation_comparation_logic_and_collision(original_mutation):
    mutation_key_1 = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, original_mutation)
    mutation_key_2 = UpdateDatasetMutationKey.create(TEST_DATASET_REVISION_ID, original_mutation)
    mutation_key_diff = UpdateDatasetMutationKey.create(TEST_DATASET_ANOTHER_REVISION_ID, original_mutation)
    assert mutation_key_1 == mutation_key_2
    assert mutation_key_1 != mutation_key_diff
    assert mutation_key_1 != 123
    mutation_collision = attr.evolve(mutation_key_diff, hash=mutation_key_1.get_hash())
    assert mutation_key_1.get_hash() == mutation_collision.get_hash()
    assert mutation_key_1.get_collision_tier_breaker() != mutation_collision.get_collision_tier_breaker()
    assert mutation_key_1 != mutation_collision
    assert mutation_key_diff != mutation_collision

lib/dl_api_lib/dl_api_lib_tests/unit/test_mutation_cache.py
datalens-tech/datalens-backend
import contextlib
from typing import (
    Generator,
    Optional,
)

import attr
import flask

from dl_api_commons.aio.middlewares.auth_trust_middleware import auth_trust_middleware
from dl_api_lib.app.control_api.app import ControlApiAppFactory
from dl_api_lib.app.control_api.app import EnvSetupResult as ControlApiEnvSetupResult
from dl_api_lib.app.data_api.app import DataApiAppFactory
from dl_api_lib.app.data_api.app import EnvSetupResult as DataApiEnvSetupResult
from dl_api_lib.app_common import (
    SRFactoryBuilder,
    StandaloneServiceRegistryFactory,
)
from dl_api_lib.app_common_settings import ConnOptionsMutatorsFactory
from dl_api_lib.app_settings import (
    AppSettings,
    ControlApiAppSettings,
    ControlApiAppTestingsSettings,
    DataApiAppSettings,
)
from dl_api_lib.connector_availability.base import ConnectorAvailabilityConfig
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.enums import (
    RedisMode,
    RequiredService,
)
from dl_configs.rqe import (
    RQEBaseURL,
    RQEConfig,
)
from dl_configs.settings_submodels import RedisSettings
from dl_constants.enums import (
    ConnectionType,
    USAuthMode,
)
from dl_core.aio.middlewares.services_registry import services_registry_middleware
from dl_core.aio.middlewares.us_manager import service_us_manager_middleware
from dl_core.data_processing.cache.primitives import CacheTTLConfig
from dl_core.flask_utils.trust_auth import TrustAuthService
from dl_core.services_registry.entity_checker import EntityUsageChecker
from dl_core.services_registry.env_manager_factory_base import EnvManagerFactory
from dl_core.services_registry.inst_specific_sr import InstallationSpecificServiceRegistryFactory
from dl_core.services_registry.rqe_caches import RQECachesSetting
from dl_core_testing.app_test_workarounds import TestEnvManagerFactory
from dl_core_testing.fixture_server_runner import WSGIRunner


@attr.s
class RQEConfigurationMaker:
    bi_test_config: ApiTestEnvironmentConfiguration = attr.ib(kw_only=True)

    @contextlib.contextmanager
    def sync_rqe_netloc_subprocess_cm(self) -> Generator[RQEBaseURL, None, None]:
        whitelist = self.bi_test_config.core_test_config.get_core_library_config().core_connector_ep_names
        env = dict(
            EXT_QUERY_EXECUTER_SECRET_KEY=self.bi_test_config.ext_query_executer_secret_key,
            DEV_LOGGING="1",
        )
        if whitelist:
            env["CORE_CONNECTOR_WHITELIST"] = ",".join(whitelist)

        with WSGIRunner(
            module="dl_core.bin.query_executor_sync",
            callable="app",
            ping_path="/ping",
            env=env,
        ) as runner:
            yield RQEBaseURL(  # type: ignore  # TODO: fix compatibility of models using `s_attrib` with mypy
                host=runner.bind_addr,
                port=runner.bind_port,
            )

    @contextlib.contextmanager
    def async_rqe_netloc_subprocess_cm(self) -> Generator[RQEBaseURL, None, None]:
        yield RQEBaseURL(  # type: ignore  # TODO: fix compatibility of models using `s_attrib` with mypy
            host="127.0.0.1",
            port=65500,
        )

    @contextlib.contextmanager
    def rqe_config_subprocess_cm(self) -> Generator[RQEConfig, None, None]:
        with (
            self.sync_rqe_netloc_subprocess_cm() as sync_rqe_netloc,
            self.async_rqe_netloc_subprocess_cm() as async_rqe_netloc,
        ):
            yield RQEConfig(  # type: ignore  # TODO: fix compatibility of models using `s_attrib` with mypy
                hmac_key=self.bi_test_config.ext_query_executer_secret_key.encode(),
                ext_sync_rqe=sync_rqe_netloc,
                ext_async_rqe=async_rqe_netloc,
                int_sync_rqe=sync_rqe_netloc,
                int_async_rqe=async_rqe_netloc,
            )


@attr.s
class RedisSettingMaker:
    bi_test_config: ApiTestEnvironmentConfiguration = attr.ib(kw_only=True)

    def get_redis_settings(self, db: int) -> RedisSettings:
        return RedisSettings(  # type: ignore  # TODO: fix compatibility of models using `s_attrib` with mypy
            MODE=RedisMode.single_host,
            CLUSTER_NAME="",
            HOSTS=(self.bi_test_config.redis_host,),
            PORT=self.bi_test_config.redis_port,
            DB=db,
            PASSWORD=self.bi_test_config.redis_password,
        )

    def get_redis_settings_default(self) -> RedisSettings:
        return self.get_redis_settings(self.bi_test_config.redis_db_default)

    def get_redis_settings_cache(self) -> RedisSettings:
        return self.get_redis_settings(self.bi_test_config.redis_db_cache)

    def get_redis_settings_mutation(self) -> RedisSettings:
        return self.get_redis_settings(self.bi_test_config.redis_db_mutation)

    def get_redis_settings_arq(self) -> RedisSettings:
        return self.get_redis_settings(self.bi_test_config.redis_db_arq)


class TestingSRFactoryBuilder(SRFactoryBuilder[AppSettings]):
    def _get_required_services(self, settings: AppSettings) -> set[RequiredService]:
        return {RequiredService.RQE_INT_SYNC, RequiredService.RQE_EXT_SYNC}

    def _get_env_manager_factory(self, settings: AppSettings) -> EnvManagerFactory:
        return TestEnvManagerFactory()

    def _get_inst_specific_sr_factory(
        self,
        settings: AppSettings,
    ) -> Optional[InstallationSpecificServiceRegistryFactory]:
        return StandaloneServiceRegistryFactory()

    def _get_entity_usage_checker(self, settings: AppSettings) -> Optional[EntityUsageChecker]:
        return None

    def _get_bleeding_edge_users(self, settings: AppSettings) -> tuple[str, ...]:
        return ()

    def _get_rqe_caches_settings(self, settings: AppSettings) -> Optional[RQECachesSetting]:
        return None

    def _get_default_cache_ttl_settings(self, settings: AppSettings) -> Optional[CacheTTLConfig]:
        return None

    def _get_connector_availability(self, settings: AppSettings) -> Optional[ConnectorAvailabilityConfig]:
        return settings.CONNECTOR_AVAILABILITY if isinstance(settings, ControlApiAppSettings) else None


class TestingControlApiAppFactory(ControlApiAppFactory[ControlApiAppSettings], TestingSRFactoryBuilder):
    """Management API app factory for tests"""

    def set_up_environment(
        self,
        app: flask.Flask,
        testing_app_settings: Optional[ControlApiAppTestingsSettings] = None,
    ) -> ControlApiEnvSetupResult:
        us_auth_mode: USAuthMode
        TrustAuthService(
            fake_user_id="_the_tests_syncapp_user_id_",
            fake_user_name="_the_tests_syncapp_user_name_",
            fake_tenant=None if testing_app_settings is None else testing_app_settings.fake_tenant,
        ).set_up(app)

        us_auth_mode_override = None if testing_app_settings is None else testing_app_settings.us_auth_mode_override

        us_auth_mode = USAuthMode.master if us_auth_mode_override is None else us_auth_mode_override

        result = ControlApiEnvSetupResult(us_auth_mode=us_auth_mode)
        return result


class TestingDataApiAppFactory(DataApiAppFactory[DataApiAppSettings], TestingSRFactoryBuilder):
    """Data API app factory for tests"""

    @property
    def _is_public(self) -> bool:
        return False

    def get_app_version(self) -> str:
        return "tests"

    def set_up_environment(
        self,
        connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
    ) -> DataApiEnvSetupResult:
        conn_opts_factory = ConnOptionsMutatorsFactory()
        sr_factory = self.get_sr_factory(
            settings=self._settings, conn_opts_factory=conn_opts_factory, connectors_settings=connectors_settings
        )

        auth_mw_list = [
            auth_trust_middleware(
                fake_user_id="_the_tests_asyncapp_user_id_",
                fake_user_name="_the_tests_asyncapp_user_name_",
            )
        ]

        sr_middleware_list = [
            services_registry_middleware(
                services_registry_factory=sr_factory,
                use_query_cache=self._settings.CACHES_ON,
                use_mutation_cache=self._settings.MUTATIONS_CACHES_ON,
                mutation_cache_default_ttl=self._settings.MUTATIONS_CACHES_DEFAULT_TTL,
            ),
        ]

        common_us_kw = dict(
            us_base_url=self._settings.US_BASE_URL,
            crypto_keys_config=self._settings.CRYPTO_KEYS_CONFIG,
        )
        usm_middleware_list = [
            service_us_manager_middleware(us_master_token=self._settings.US_MASTER_TOKEN, **common_us_kw),  # type: ignore[arg-type]
            service_us_manager_middleware(us_master_token=self._settings.US_MASTER_TOKEN, as_user_usm=True, **common_us_kw),  # type: ignore[arg-type]
        ]

        result = DataApiEnvSetupResult(
            auth_mw_list=auth_mw_list,
            sr_middleware_list=sr_middleware_list,
            usm_middleware_list=usm_middleware_list,
        )

        return result

lib/dl_api_lib_testing/dl_api_lib_testing/app.py
datalens-tech/datalens-backend
import abc
from asyncio import AbstractEventLoop
from typing import (
    Any,
    ClassVar,
    Generator,
    Optional,
)

from flask.app import Flask
from flask.testing import FlaskClient
import pytest

from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase
from dl_api_commons.base_models import (
    RequestContextInfo,
    TenantCommon,
)
from dl_api_lib.app.control_api.app import ControlApiAppFactory
from dl_api_lib.app_common_settings import ConnOptionsMutatorsFactory
from dl_api_lib.app_settings import (
    ControlApiAppSettings,
    ControlApiAppTestingsSettings,
)
from dl_api_lib.connector_availability.base import ConnectorAvailabilityConfig
from dl_api_lib.loader import preload_api_lib
from dl_api_lib_testing.app import (
    RedisSettingMaker,
    RQEConfigurationMaker,
    TestingControlApiAppFactory,
)
from dl_api_lib_testing.client import FlaskSyncApiClient
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.rqe import RQEConfig
from dl_constants.enums import (
    ConnectionType,
    QueryProcessingMode,
)
from dl_core.components.ids import FieldIdGeneratorType
from dl_core.united_storage_client import USAuthContextMaster
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.configuration import CoreTestEnvironmentConfigurationBase
from dl_core_testing.flask_utils import (
    FlaskTestClient,
    FlaskTestResponse,
)


class ApiTestBase(abc.ABC):
    """
    Base class defining the basic fixtures of bi-api tests
    """

    bi_compeng_pg_on: ClassVar[bool] = True
    query_processing_mode: ClassVar[QueryProcessingMode] = QueryProcessingMode.basic

    @pytest.fixture(scope="function", autouse=True)
    def preload(self):
        preload_api_lib()

    @pytest.fixture(scope="class")
    @abc.abstractmethod
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        raise NotImplementedError

    @pytest.fixture(scope="class")
    def core_test_config(
        self,
        bi_test_config: ApiTestEnvironmentConfiguration,
    ) -> CoreTestEnvironmentConfigurationBase:
        return bi_test_config.core_test_config

    @pytest.fixture(scope="class")
    def connectors_settings(self) -> dict[ConnectionType, ConnectorSettingsBase]:
        return {}

    def is_connector_available(self, conn_type: ConnectionType) -> bool:
        return True

    @pytest.fixture(scope="function")
    def enable_all_connectors(self, monkeypatch: Any) -> None:
        monkeypatch.setattr(ConnectorAvailabilityConfig, "check_connector_is_available", self.is_connector_available)

    @pytest.fixture(scope="function")
    def environment_readiness(self, enable_all_connectors: Any) -> None:
        """Make sure the environment is ready for tests"""

    @pytest.fixture(scope="function")
    def rqe_config_subprocess(
        self,
        bi_test_config: ApiTestEnvironmentConfiguration,
    ) -> Generator[RQEConfig, None, None]:
        with RQEConfigurationMaker(bi_test_config=bi_test_config).rqe_config_subprocess_cm() as rqe_config:
            yield rqe_config

    @classmethod
    def create_control_api_settings(
        cls,
        bi_test_config: ApiTestEnvironmentConfiguration,
        rqe_config_subprocess: RQEConfig,
    ) -> ControlApiAppSettings:
        core_test_config = bi_test_config.core_test_config
        us_config = core_test_config.get_us_config()

        redis_setting_maker = RedisSettingMaker(bi_test_config=bi_test_config)

        settings = ControlApiAppSettings(
            CONNECTOR_AVAILABILITY=ConnectorAvailabilityConfig.from_settings(
                bi_test_config.connector_availability_settings,
            ),
            US_BASE_URL=us_config.us_host,
            US_MASTER_TOKEN=us_config.us_master_token,
            CRYPTO_KEYS_CONFIG=core_test_config.get_crypto_keys_config(),
            BI_API_CONNECTOR_WHITELIST=bi_test_config.get_api_library_config().api_connector_ep_names,
            CORE_CONNECTOR_WHITELIST=core_test_config.get_core_library_config().core_connector_ep_names,
            RQE_CONFIG=rqe_config_subprocess,
            BI_COMPENG_PG_ON=cls.bi_compeng_pg_on,
            BI_COMPENG_PG_URL=bi_test_config.bi_compeng_pg_url,
            DO_DSRC_IDX_FETCH=True,
            FIELD_ID_GENERATOR_TYPE=FieldIdGeneratorType.suffix,
            REDIS_ARQ=redis_setting_maker.get_redis_settings_arq(),
            FILE_UPLOADER_BASE_URL="http://127.0.0.1:9999",  # fake url
            FILE_UPLOADER_MASTER_TOKEN="qwerty",
            QUERY_PROCESSING_MODE=cls.query_processing_mode,
        )
        return settings

    @pytest.fixture(scope="function")
    def control_api_app_settings(
        self,
        bi_test_config: ApiTestEnvironmentConfiguration,
        rqe_config_subprocess: RQEConfig,
    ) -> ControlApiAppSettings:
        return self.create_control_api_settings(
            bi_test_config=bi_test_config,
            rqe_config_subprocess=rqe_config_subprocess,
        )

    @pytest.fixture(scope="function")
    def control_api_app_factory(self, control_api_app_settings: ControlApiAppSettings) -> ControlApiAppFactory:
        return TestingControlApiAppFactory(settings=control_api_app_settings)

    @pytest.fixture(scope="function")
    def control_api_app(
        self,
        environment_readiness: None,
        control_api_app_factory: ControlApiAppFactory,
        connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
    ) -> Generator[Flask, None, None]:
        """Session-wide test `Flask` application."""

        app = control_api_app_factory.create_app(
            connectors_settings=connectors_settings,
            testing_app_settings=ControlApiAppTestingsSettings(fake_tenant=TenantCommon()),
            close_loop_after_request=False,
        )

        app.config["WE_ARE_IN_TESTS"] = True

        # Establish an application context before running the tests.
        with app.app_context() as ctx:
            assert ctx
            yield app

    @pytest.fixture(scope="function")
    def client(
        self, control_api_app: Flask, loop: AbstractEventLoop
    ) -> FlaskClient:  # FIXME: Rename to control_app_client
        class TestClient(FlaskTestClient):
            def get_default_headers(self) -> dict[str, Optional[str]]:
                return {}

        control_api_app.test_client_class = TestClient
        control_api_app.response_class = FlaskTestResponse  # for the `json` property

        return control_api_app.test_client()

    @pytest.fixture(scope="function")
    def control_api_sync_client(self, client: FlaskClient) -> SyncHttpClientBase:
        return FlaskSyncApiClient(int_wclient=client)

    @pytest.fixture(scope="function")
    def control_api(self, control_api_sync_client: SyncHttpClientBase) -> SyncHttpDatasetApiV1:
        return SyncHttpDatasetApiV1(client=control_api_sync_client)

    @pytest.fixture(scope="function")
    def sync_us_manager(
        self,
        core_test_config: CoreTestEnvironmentConfigurationBase,
        control_api_app_factory: ControlApiAppFactory,
        connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
        control_api_app_settings: ControlApiAppSettings,
    ) -> SyncUSManager:
        bi_context = RequestContextInfo.create_empty()
        us_config = core_test_config.get_us_config()
        us_manager = SyncUSManager(
            bi_context=bi_context,
            services_registry=control_api_app_factory.get_sr_factory(
                conn_opts_factory=ConnOptionsMutatorsFactory(),
                connectors_settings=connectors_settings,
                settings=control_api_app_settings,
            ).make_service_registry(request_context_info=bi_context),
            us_base_url=us_config.us_host,
            us_auth_context=USAuthContextMaster(us_config.us_master_token),
            crypto_keys_config=core_test_config.get_crypto_keys_config(),
        )
        return us_manager

lib/dl_api_lib_testing/dl_api_lib_testing/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from asyncio import AbstractEventLoop
import json
from typing import (
    Any,
    Optional,
)

import aiohttp.test_utils
import attr
from werkzeug.test import Client as WClient

from dl_api_client.dsmaker.api.http_sync_base import (
    ClientResponse,
    SyncHttpClientBase,
)
from dl_utils.aio import await_sync


@attr.s
class TestClientConverterAiohttpToFlask:
    _loop: AbstractEventLoop = attr.ib()
    _aio_client: aiohttp.test_utils.TestClient = attr.ib()
    _extra_headers: Optional[dict[str, str]] = attr.ib(default=None)
    _cookies: Optional[dict[str, str]] = attr.ib(default=None)

    @attr.s(frozen=True, auto_attribs=True)
    class Resp:
        status_code: int
        data: bytes
        _json: Optional[dict]

        @property
        def json(self) -> dict:
            assert self._json is not None
            return self._json

    def set_cookie(self, server_name: str, key: str, value: str = "", **kwargs: Any) -> None:
        if self._cookies is None:
            self._cookies = {}
        self._cookies.update({key: value})

    def delete_cookie(self, server_name: str, key: str, **kwargs: Any) -> None:
        if self._cookies is not None:
            self._cookies.pop(key, None)

    def post(self, url: str, **kwargs: Any) -> Resp:
        return self.open(url, method="post", **kwargs)

    def get(self, url: str, **kwargs: Any) -> Resp:
        return self.open(url, method="get", **kwargs)

    def open(
        self,
        url: str,
        *,
        method: str,
        data: Optional[str] = None,
        content_type: Optional[str] = None,
        headers: Optional[dict] = None,
    ) -> Resp:
        headers = dict(headers) if headers else {}
        if self._extra_headers:
            headers.update(self._extra_headers)

        if content_type is not None:
            headers["Content-Type"] = content_type
        resp = await_sync(
            self._aio_client.request(
                method,
                url,
                data=data,
                headers=headers,
                cookies=self._cookies,
            ),
            loop=self._loop,
        )
        resp_data = await_sync(resp.read())

        js: Optional[dict] = None
        if resp_data is not None:
            try:
                js = json.loads(resp_data)
            except json.JSONDecodeError:
                pass

        return self.Resp(status_code=resp.status, data=resp_data, json=js)


@attr.s
class WrappedAioSyncApiClient(SyncHttpClientBase):
    _int_wrapped_client: TestClientConverterAiohttpToFlask = attr.ib()

    def __attrs_post_init__(self) -> None:
        assert isinstance(self._int_wrapped_client, TestClientConverterAiohttpToFlask)

    def open(
        self,
        url: str,
        method: str,
        headers: dict,
        data: Optional[str] = None,
        content_type: Optional[str] = None,
    ) -> ClientResponse:
        method = method.lower()
        int_response = self._int_wrapped_client.open(
            url=url,
            method=method,
            headers=headers,
            data=data,
            content_type=content_type,
        )
        return ClientResponse(
            status_code=int_response.status_code,
            data=int_response.data,
            json=int_response.json,
        )


@attr.s
class FlaskSyncApiClient(SyncHttpClientBase):
    _int_wclient: WClient = attr.ib()

    def __attrs_post_init__(self) -> None:
        assert isinstance(self._int_wclient, WClient)

    def open(
        self,
        url: str,
        method: str,
        headers: dict,
        data: Optional[str] = None,
        content_type: Optional[str] = None,
    ) -> ClientResponse:
        method = method.lower()
        int_response = self._int_wclient.open(
            url,
            method=method,
            headers=headers,
            data=data,
            content_type=content_type,
        )
        return ClientResponse(
            status_code=int_response.status_code,
            data=int_response.data,
            json=int_response.json,
        )

lib/dl_api_lib_testing/dl_api_lib_testing/client.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Collection,
    Optional,
)

import attr

from dl_api_lib.loader import ApiLibraryConfig
from dl_configs.connector_availability import ConnectorAvailabilityConfigSettings
from dl_core_testing.configuration import CoreTestEnvironmentConfigurationBase


@attr.s(kw_only=True)
class ApiTestEnvironmentConfiguration:
    core_test_config: CoreTestEnvironmentConfigurationBase = attr.ib()

    ext_query_executer_secret_key: str = attr.ib()

    mutation_cache_enabled: bool = attr.ib(default=True)

    bi_compeng_pg_url: str = attr.ib(default="")

    file_uploader_api_host: str = attr.ib(default="http://127.0.0.1")
    file_uploader_api_port: int = attr.ib(default=9999)

    redis_host: str = attr.ib(default="")
    redis_port: int = attr.ib(default=6379)
    redis_password: str = attr.ib(default="")
    redis_db_default: int = attr.ib(default=0)
    redis_db_cache: int = attr.ib(default=1)
    redis_db_mutation: int = attr.ib(default=2)
    redis_db_arq: int = attr.ib(default=11)

    connector_availability_settings: ConnectorAvailabilityConfigSettings = attr.ib(
        factory=ConnectorAvailabilityConfigSettings,
    )
    api_connector_ep_names: Optional[Collection[str]] = attr.ib(default=None)

    def get_api_library_config(self) -> ApiLibraryConfig:
        return ApiLibraryConfig(
            api_connector_ep_names=self.api_connector_ep_names,
            core_lib_config=self.core_test_config.get_core_library_config(),
        )

lib/dl_api_lib_testing/dl_api_lib_testing/configuration.py
datalens-tech/datalens-backend
import abc
import contextlib
import json
from typing import (
    ClassVar,
    Generator,
)
import uuid

import pytest

from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase
from dl_api_lib_testing.base import ApiTestBase
from dl_constants.enums import ConnectionType
from dl_core_testing.database import (
    Db,
    DbTable,
)
from dl_core_testing.fixtures.dispenser import DbCsvTableDispenser
from dl_core_testing.fixtures.sample_tables import TABLE_SPEC_SAMPLE_SUPERSTORE
from dl_core_testing.testcases.service_base import DbServiceFixtureTextClass


class ConnectionTestBase(ApiTestBase, DbServiceFixtureTextClass):
    conn_type: ClassVar[ConnectionType]

    db_table_dispenser = DbCsvTableDispenser()

    @abc.abstractmethod
    @pytest.fixture(scope="class")
    def connection_params(self) -> dict:
        raise NotImplementedError

    @pytest.fixture(scope="function")
    def saved_connection_id(
        self,
        control_api_sync_client: SyncHttpClientBase,
        connection_params: dict,
    ) -> Generator[str, None, None]:
        with self.create_connection(
            control_api_sync_client=control_api_sync_client,
            connection_params=connection_params,
        ) as conn_id:
            yield conn_id

    @contextlib.contextmanager
    def create_connection(
        self,
        control_api_sync_client: SyncHttpClientBase,
        connection_params: dict,
    ) -> Generator[str, None, None]:
        data = dict(
            name=f"{self.conn_type.name} conn {uuid.uuid4()}",
            type=self.conn_type.name,
            **connection_params,
        )
        resp = control_api_sync_client.post(
            "/api/v1/connections",
            content_type="application/json",
            data=json.dumps(data),
        )
        assert resp.status_code == 200, resp.json
        conn_id = resp.json["id"]

        yield conn_id

        control_api_sync_client.delete("/api/v1/connections/{}".format(conn_id))

    @pytest.fixture(scope="class")
    def sample_table(self, db: Db) -> DbTable:
        return self.db_table_dispenser.get_csv_table(db=db, spec=TABLE_SPEC_SAMPLE_SUPERSTORE)

lib/dl_api_lib_testing/dl_api_lib_testing/connection_base.py
datalens-tech/datalens-backend
import json
from typing import (
    ClassVar,
    Optional,
    Type,
    final,
)

import pytest

from dl_api_commons.base_models import (
    TenantCommon,
    TenantDef,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
)
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_i18n.localizer_base import (
    LocalizerLoader,
    TranslationConfig,
)


class ConnectionFormTestBase:
    CONN_FORM_FACTORY_CLS: ClassVar[Type[ConnectionFormFactory]]
    TRANSLATION_CONFIGS: ClassVar[list[TranslationConfig]]

    @pytest.fixture
    def connectors_settings(self) -> Optional[ConnectorSettingsBase]:
        """Parametrize if a form has extra settings"""

        return None

    @pytest.fixture
    def tenant(self) -> TenantDef:
        """Parametrize if a form depends on tenant type"""

        return TenantCommon()

    @final
    @pytest.fixture(
        params=[mode.name for mode in ConnectionFormMode],
    )
    def mode(self, request) -> ConnectionFormMode:  # type: ignore
        return ConnectionFormMode(request.param)

    @pytest.fixture
    def form_config(
        self, connectors_settings: Optional[ConnectorSettingsBase], tenant: TenantDef, mode: ConnectionFormMode
    ) -> ConnectionForm:
        loader = LocalizerLoader(
            configs=self.TRANSLATION_CONFIGS,
        )
        factory = loader.load()
        localizer = factory.get_for_locale("en")
        form_factory = self.CONN_FORM_FACTORY_CLS(mode=mode, localizer=localizer)
        form_config = form_factory.get_form_config(connectors_settings, tenant)
        return form_config

    def test_validate_conditional_fields(self, form_config: ConnectionForm) -> None:
        form_config.validate_conditional_fields()

    def test_validate_api_schema_fields(self, form_config: ConnectionForm) -> None:
        form_config.validate_api_schema_fields()

    def test_serialize(self, form_config: ConnectionForm) -> None:
        serializable = form_config.as_dict()
        json.dumps(serializable)

lib/dl_api_lib_testing/dl_api_lib_testing/connection_form_base.py
datalens-tech/datalens-backend
from collections import defaultdict
import functools
from http import HTTPStatus
from typing import Iterable

import pytest

from dl_api_client.dsmaker.primitives import (
    Dataset,
    ResultField,
)
from dl_api_client.dsmaker.shortcuts.dataset import (
    add_formulas_to_dataset,
    create_basic_dataset,
)
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_testing.data_api_base import DataApiTestBase
from dl_api_lib_testing.dataset_base import DatasetTestBase
from dl_api_lib_testing.helpers.data_source import data_source_settings_from_table
from dl_api_lib_testing.helpers.lookup_checkers import check_ago_data
from dl_constants.enums import (
    QueryProcessingMode,
    UserDataType,
    WhereClauseOperation,
)
from dl_core_testing.database import (
    C,
    make_table,
)
from dl_core_testing.testcases.service_base import DbServiceFixtureTextClass
from dl_testing.regulated_test import (
    Feature,
    for_features,
)


class DefaultBasicExtAggregationTestSuite(DataApiTestBase, DatasetTestBase, DbServiceFixtureTextClass):
    def test_lod_fixed_single_dim_in_two_dim_query(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "sales sum": "SUM([sales])",
                "sales sum fx city": "SUM([sales] FIXED [city])",
                "sales sum fx category": "SUM([sales] FIXED [category])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="category"),
                ds.find_field(title="sales sum"),
                ds.find_field(title="sales sum fx city"),
                ds.find_field(title="sales sum fx category"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        sum_by_city = defaultdict(lambda: 0)
        for row in data_rows:
            sum_by_city[row[0]] += float(row[2])
        sum_by_category = defaultdict(lambda: 0)
        for row in data_rows:
            sum_by_category[row[1]] += float(row[2])

        for row in data_rows:
            assert float(row[3]) == pytest.approx(sum_by_city[row[0]])
            assert float(row[4]) == pytest.approx(sum_by_category[row[1]])

    def test_null_dimensions(self, request, control_api, data_api, db, saved_connection_id):
        connection_id = saved_connection_id

        raw_data = [
            {"id": 1, "city": "New York", "category": "Office Supplies", "sales": 1},
            {"id": 2, "city": "New York", "category": "Furniture", "sales": 10},
            {"id": 3, "city": "New Rochelle", "category": "Office Supplies", "sales": 100},
            {"id": 4, "city": "New Rochelle", "category": "Furniture", "sales": 1000},
            {"id": 5, "city": None, "category": "Office Supplies", "sales": 10000},
            {"id": 6, "city": None, "category": "Furniture", "sales": 100000},
        ]
        columns = [
            C("id", UserDataType.integer, vg=lambda rn, **kwargs: raw_data[rn]["id"]),
            C("city", UserDataType.string, vg=lambda rn, **kwargs: raw_data[rn]["city"]),
            C("category", UserDataType.string, vg=lambda rn, **kwargs: raw_data[rn]["category"]),
            C("sales", UserDataType.integer, vg=lambda rn, **kwargs: raw_data[rn]["sales"]),
        ]
        db_table = make_table(db, columns=columns, rows=len(raw_data))
        request.addfinalizer(functools.partial(db.drop_table, db_table.table))

        ds = Dataset()
        ds.sources["source_1"] = ds.source(connection_id=connection_id, **data_source_settings_from_table(db_table))
        ds.source_avatars["avatar_1"] = ds.sources["source_1"].avatar()

        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=ds,
            formulas={
                "sum_lod": "SUM(SUM([sales] INCLUDE [category]))",
                "max_id": "MAX([id])",
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="sum_lod"),
            ],
            order_by=[
                ds.find_field(title="max_id"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert data_rows

        city = [row[0] for row in data_rows]
        sum_lod = [int(row[1]) for row in data_rows]

        assert city == ["New York", "New Rochelle", None]
        assert sum_lod == [11, 1100, 110000]

    def test_total_lod(self, control_api, data_api, saved_dataset):
        data_api = data_api
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "sales sum": "SUM([sales])",
                "sales sum total": "SUM([sales] FIXED)",
            },
        )

        def get_total_value() -> float:
            total_sum_result = data_api.get_result(
                dataset=ds,
                fields=[ds.find_field(title="sales sum")],
            )
            data_rows = get_data_rows(total_sum_result)
            return float(data_rows[0][0])

        expected_total_value = get_total_value()

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="city"),
                ds.find_field(title="sales sum"),
                ds.find_field(title="sales sum total"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        total_sum = sum(float(row[1]) for row in data_rows)

        for row_idx, row in enumerate(data_rows):
            assert (
                float(row[2]) == pytest.approx(total_sum) == pytest.approx(expected_total_value)
            ), f"total sum doesn't match expected number in row {row_idx}"

    def test_total_lod_2(
        self,
        request,
        control_api,
        data_api,
        saved_connection_id,
        db,
    ):
        db_table = make_table(db=db)
        request.addfinalizer(functools.partial(db.drop_table, db_table.table))
        ds = create_basic_dataset(
            api_v1=control_api,
            connection_id=saved_connection_id,
            data_source_settings=data_source_settings_from_table(table=db_table),
            formulas={
                "div 2": "DIV([int_value], 2)",
                "div 3": "DIV([int_value], 3)",
                "Agg 1": "SUM([int_value])",
                "Agg 2": "SUM(SUM([int_value] INCLUDE [div 2]))",
                "Agg 3": "SUM(SUM(SUM([int_value] INCLUDE [div 2]) INCLUDE [div 3]))",
            },
        )

        def get_single_row_data(field_names: Iterable[str]) -> tuple[int, ...]:
            result_resp = data_api.get_result(
                dataset=ds,
                fail_ok=True,
                fields=[ds.find_field(title=field_name) for field_name in field_names],
            )
            assert result_resp.status_code == HTTPStatus.OK, result_resp.json
            data_rows = get_data_rows(result_resp)
            assert len(data_rows) == 1, "There must be exactly 1 row of data"
            return tuple(int(item) for item in data_rows[0])

        value_1 = get_single_row_data(["Agg 1"])[0]
        value_2 = get_single_row_data(["Agg 2"])[0]
        assert value_2 == value_1
        value_3 = get_single_row_data(["Agg 3"])[0]
        assert value_3 == value_1

        def check_equality_of_totals(*field_names: str) -> None:
            simultaneous_values = get_single_row_data(field_names)
            assert len(set(simultaneous_values)) == 1
            assert next(iter(simultaneous_values)) == value_1

        check_equality_of_totals("Agg 1", "Agg 2")
        check_equality_of_totals("Agg 2", "Agg 3")
        check_equality_of_totals("Agg 1", "Agg 3")
        check_equality_of_totals("Agg 1", "Agg 2", "Agg 3")

    def test_lod_in_order_by(
        self,
        request,
        control_api,
        data_api,
        saved_connection_id,
        db,
    ):
        db_table = make_table(db=db)
        request.addfinalizer(functools.partial(db.drop_table, db_table.table))

        data_api = data_api
        ds = create_basic_dataset(
            api_v1=control_api,
            connection_id=saved_connection_id,
            data_source_settings=data_source_settings_from_table(table=db_table),
            formulas={
                "Dimension": "[int_value] % 3",
                "LOD Measure": "SUM([int_value]) / SUM([int_value] FIXED)",
            },
        )

        def get_data(order_by: list[ResultField]) -> list:
            result_resp = data_api.get_result(
                dataset=ds,
                fields=[
                    ds.find_field(title="Dimension"),
                    ds.find_field(title="LOD Measure"),
                ],
                order_by=order_by,
                fail_ok=True,
            )
            assert result_resp.status_code == HTTPStatus.OK, result_resp.json
            return get_data_rows(result_resp)

        data_rows = get_data(order_by=[])
        ordered_data_rows = get_data(
            order_by=[
                ds.find_field(title="LOD Measure"),
                ds.find_field(title="Dimension"),
            ]
        )

        data_rows.sort(key=lambda row: (float(row[1]), row[0]))  # (LOD Measure, City)

        assert ordered_data_rows == data_rows


class DefaultBasicLookupFunctionTestSuite(DataApiTestBase, DatasetTestBase, DbServiceFixtureTextClass):
    def test_ago_any_db(self, request, saved_connection_id, control_api, data_api, db):
        db_table = make_table(db=db)
        request.addfinalizer(functools.partial(db.drop_table, db_table.table))
        ds = create_basic_dataset(
            api_v1=control_api,
            connection_id=saved_connection_id,
            data_source_settings=data_source_settings_from_table(table=db_table),
            formulas={
                "sum": "SUM([int_value])",
                "ago": 'AGO([sum], [date_value], "day", 2)',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="date_value"),
                ds.find_field(title="sum"),
                ds.find_field(title="ago"),
            ],
            order_by=[
                ds.find_field(title="date_value"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=2, day_offset=2)

    def test_triple_ago_any_db(self, request, saved_connection_id, control_api, data_api, db):
        db_table = make_table(db)
        request.addfinalizer(functools.partial(db.drop_table, db_table.table))
        ds = create_basic_dataset(
            api_v1=control_api,
            connection_id=saved_connection_id,
            data_source_settings=data_source_settings_from_table(table=db_table),
            formulas={
                "sum": "SUM([int_value])",
                "ago_1": 'AGO([sum], [date_value], "day", 1)',
                "ago_2": 'AGO([sum], [date_value], "day", 2)',
                "ago_3": 'AGO([sum], [date_value], "day", 3)',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="date_value"),
                ds.find_field(title="sum"),
                ds.find_field(title="ago_1"),
                ds.find_field(title="ago_2"),
                ds.find_field(title="ago_3"),
            ],
            order_by=[
                ds.find_field(title="date_value"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=2, day_offset=1)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=3, day_offset=2)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=4, day_offset=3)

    def test_ago_any_db_multisource(self, request, saved_connection_id, control_api, data_api, db):
        connection_id = saved_connection_id
        table_1 = make_table(db)
        table_2 = make_table(db)

        def teardown(db, *tables):
            for table in tables:
                db.drop_table(table)

        request.addfinalizer(functools.partial(teardown, db, table_1.table, table_2.table))

        ds = Dataset()
        ds.sources["source_1"] = ds.source(
            connection_id=connection_id,
            **data_source_settings_from_table(table=table_1),
        )
        ds.sources["source_2"] = ds.source(
            connection_id=connection_id,
            **data_source_settings_from_table(table=table_2),
        )
        ds.source_avatars["avatar_1"] = ds.sources["source_1"].avatar()
        ds.source_avatars["avatar_2"] = ds.sources["source_2"].avatar()
        ds.avatar_relations["rel_1"] = (
            ds.source_avatars["avatar_1"]
            .join(ds.source_avatars["avatar_2"])
            .on(ds.col("string_value") == ds.col("string_value"))
        )

        ds.result_schema["date_1"] = ds.source_avatars["avatar_1"].field(source="date_value")
        ds.result_schema["int_2"] = ds.source_avatars["avatar_2"].field(source="int_value")
        ds.result_schema["sum"] = ds.field(formula="SUM([int_2])")
        ds.result_schema["ago"] = ds.field(formula='AGO([sum], [date_1], "day", 2)')
        ds_resp = control_api.apply_updates(dataset=ds, fail_ok=True)
        assert ds_resp.status_code == HTTPStatus.OK, ds_resp.response_errors
        ds = ds_resp.dataset
        ds = control_api.save_dataset(ds).dataset

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="date_1"),
                ds.find_field(title="sum"),
                ds.find_field(title="ago"),
            ],
            order_by=[
                ds.find_field(title="date_1"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=2, day_offset=2)

    def test_nested_ago(self, request, saved_connection_id, control_api, data_api, db):
        db_table = make_table(db)
        request.addfinalizer(functools.partial(db.drop_table, db_table.table))
        ds = create_basic_dataset(
            api_v1=control_api,
            connection_id=saved_connection_id,
            data_source_settings=data_source_settings_from_table(table=db_table),
            formulas={
                "sum": "SUM([int_value])",
                "ago_1": 'AGO([sum], [date_value], "day", 1)',
                "ago_2": 'AGO([ago_1], [date_value], "day", 1)',
                "ago_3": 'AGO([ago_2], [date_value], "day", 1)',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="date_value"),
                ds.find_field(title="sum"),
                ds.find_field(title="ago_1"),
                ds.find_field(title="ago_2"),
                ds.find_field(title="ago_3"),
            ],
            order_by=[
                ds.find_field(title="date_value"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=2, day_offset=1)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=3, day_offset=2)
        check_ago_data(data_rows=data_rows, date_idx=0, value_idx=1, ago_idx=4, day_offset=3)

    def test_month_ago_for_shorter_month(self, request, db, saved_connection_id, control_api, data_api):
        any_db_table_200 = make_table(db, rows=200)
        request.addfinalizer(functools.partial(db.drop_table, any_db_table_200.table))

        # FIXME
        # if any_db.conn_type == CONNECTION_TYPE_ORACLE:
        #     # Oracle cannot add a month to 2021-01-31 (2021-02-31 doesn't exist)
        #     pytest.skip()

        ds = create_basic_dataset(
            api_v1=control_api,
            connection_id=saved_connection_id,
            data_source_settings=data_source_settings_from_table(table=any_db_table_200),
            formulas={
                "new_date_value": "#2021-01-01# + [int_value]",
                "sum": "SUM([int_value])",
                "ago": 'AGO([sum], [new_date_value], "month", 1)',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="new_date_value"),
                ds.find_field(title="sum"),
                ds.find_field(title="ago"),
            ],
            filters=[
                ds.find_field(title="new_date_value").filter(op=WhereClauseOperation.EQ, values=["2021-02-28"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        data_rows = get_data_rows(result_resp)

        # Check whether rows are duplicated
        assert len(data_rows) == 1


class DefaultBasicWindowFunctionTestSuite(DataApiTestBase, DatasetTestBase, DbServiceFixtureTextClass):
    feature_window_functions = Feature("window_functions")

    @for_features(feature_window_functions)
    def test_window_functions(self, control_api, data_api, saved_dataset):
        ds = add_formulas_to_dataset(
            api_v1=control_api,
            dataset=saved_dataset,
            formulas={
                "Group Sales": "SUM([sales])",
                "Rank of Sales": 'RANK([Group Sales], "asc" TOTAL)',
                "Unique Rank of Sales": 'RANK_UNIQUE([Group Sales], "asc" TOTAL)',
                "Rank of City Sales for Date": 'RANK([Group Sales], "asc" AMONG [city])',
                "Total Sales": "SUM([Group Sales] TOTAL)",
                "Date Sales": "SUM([Group Sales] WITHIN [order_date])",
                "City Sales": "SUM([Group Sales] AMONG [order_date])",
                "Total RSUM": 'RSUM([Group Sales], "asc" TOTAL)',
            },
        )

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="order_date"),
                ds.find_field(title="city"),
                ds.find_field(title="Group Sales"),
                ds.find_field(title="Rank of Sales"),
                ds.find_field(title="Unique Rank of Sales"),
                ds.find_field(title="Rank of City Sales for Date"),
                ds.find_field(title="Total Sales"),
                ds.find_field(title="Date Sales"),
                ds.find_field(title="City Sales"),
                ds.find_field(title="Total RSUM"),
            ],
            order_by=[
                ds.find_field(title="order_date"),
                ds.find_field(title="city"),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json

        data_rows = get_data_rows(result_resp)
        cnt = len(data_rows)

        # TODO: More thorough tests

        # [Rank of Sales] values are a subset of the full range of row numbers
        assert {row[3] for row in data_rows}.issubset({str(i) for i in range(1, cnt + 1)})

        # There are as many [Unique Rank of Sales] values as there are rows
        assert {row[4] for row in data_rows} == ({str(i) for i in range(1, cnt + 1)})

        # [Rank of City Sales for Date] values are not greater than the number of [City] values
        assert len({row[5] for row in data_rows}) <= len({row[1] for row in data_rows})

        # all rows have the same [Total Sales] value
        assert len({row[6] for row in data_rows}) == 1

        # as many values of [Date Sales] as there are [order_date] values (or less - because there may be duplicates)
        assert len({row[7] for row in data_rows}) <= len({row[0] for row in data_rows})

        # as many values of [City Sales] as there are [City] values (or less - because there may be duplicates)
        assert len({row[8] for row in data_rows}) <= len({row[1] for row in data_rows})

        for i in range(1, len(data_rows)):
            # RSUM = previous RSUM value + value of current arg
            assert pytest.approx(float(data_rows[i][9])) == float(data_rows[i - 1][9]) + float(data_rows[i][2])


class DefaultBasicComplexQueryTestSuite(
    DefaultBasicExtAggregationTestSuite,
    DefaultBasicLookupFunctionTestSuite,
    DefaultBasicWindowFunctionTestSuite,
):
    """Put them all together"""

    query_processing_mode = QueryProcessingMode.native_wf

lib/dl_api_lib_testing/dl_api_lib_testing/connector/complex_queries.py
datalens-tech/datalens-backend
import json

from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_testing.regulated_test import RegulatedTestCase


class DefaultConnectorConnectionTestSuite(ConnectionTestBase, RegulatedTestCase):
    def test_create_connection(self, control_api_sync_client: SyncHttpClientBase, saved_connection_id: str) -> None:
        assert saved_connection_id
        resp = control_api_sync_client.get(
            url=f"/api/v1/connections/{saved_connection_id}",
        )
        assert resp.status_code == 200, resp.json

    def test_test_connection(self, control_api_sync_client: SyncHttpClientBase, saved_connection_id: str) -> None:
        resp = control_api_sync_client.post(
            f"/api/v1/connections/test_connection/{saved_connection_id}",
            content_type="application/json",
            data=json.dumps({}),
        )
        assert resp.status_code == 200, resp.json

    def test_cache_ttl_sec_override(
        self, control_api_sync_client: SyncHttpClientBase, saved_connection_id: str
    ) -> None:
        resp = control_api_sync_client.get(
            url=f"/api/v1/connections/{saved_connection_id}",
        )
        assert resp.status_code == 200, resp.json
        assert resp.json["cache_ttl_sec"] is None, resp.json

        cache_ttl_override = 100500
        resp = control_api_sync_client.put(
            url=f"/api/v1/connections/{saved_connection_id}",
            content_type="application/json",
            data=json.dumps({"cache_ttl_sec": cache_ttl_override}),
        )
        assert resp.status_code == 200, resp.json

        resp = control_api_sync_client.get(
            url=f"/api/v1/connections/{saved_connection_id}",
        )
        assert resp.status_code == 200, resp.json
        assert resp.json["cache_ttl_sec"] == cache_ttl_override, resp.json

lib/dl_api_lib_testing/dl_api_lib_testing/connector/connection_suite.py
datalens-tech/datalens-backend
import abc
from typing import Optional

from aiohttp.test_utils import TestClient
import pytest

from dl_api_lib_testing.dashsql_base import DashSQLTestBase
from dl_testing.regulated_test import RegulatedTestCase


class DefaultDashSQLTestSuite(DashSQLTestBase, RegulatedTestCase, metaclass=abc.ABCMeta):
    @pytest.fixture(scope="class")
    def dashsql_headers(self) -> Optional[dict[str, str]]:
        return None

    @pytest.fixture(scope="class")
    def dashsql_basic_query(self) -> str:
        return "select 1, 2, 3"

    @pytest.mark.asyncio
    async def test_basic_select(
        self,
        data_api_lowlevel_aiohttp_client: TestClient,
        saved_connection_id: str,
        dashsql_basic_query: str,
        dashsql_headers: Optional[dict[str, str]],
    ) -> None:
        resp = await self.get_dashsql_response(
            data_api_aio=data_api_lowlevel_aiohttp_client,
            conn_id=saved_connection_id,
            query=dashsql_basic_query,
            headers=dashsql_headers,
        )
        data = await resp.json()
        assert data[1]["data"] == [1, 2, 3]

lib/dl_api_lib_testing/dl_api_lib_testing/connector/dashsql_suite.py
datalens-tech/datalens-backend
import datetime
import functools
import json
from typing import Any

import pytest

from dl_api_client.dsmaker.api.data_api import SyncHttpDataApiV2
from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.primitives import (
    Dataset,
    WhereClause,
)
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_testing.data_api_base import (
    DataApiTestParams,
    StandardizedDataApiTestBase,
)
from dl_constants.enums import (
    UserDataType,
    WhereClauseOperation,
)
from dl_core_testing.database import (
    C,
    Db,
    make_table,
)
from dl_testing.regulated_test import (
    Feature,
    RegulatedTestCase,
    for_features,
)


class DefaultConnectorDataResultTestSuite(StandardizedDataApiTestBase, RegulatedTestCase):
    array_support = Feature("Connector supports arrays")

    def test_basic_result(
        self,
        saved_dataset: Dataset,
        data_api_test_params: DataApiTestParams,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        ds = saved_dataset
        ds.result_schema["Measure"] = ds.field(formula=f"SUM([{data_api_test_params.summable_field}])")

        result_resp_1 = self.get_result(ds, data_api, field_names=("Measure", data_api_test_params.two_dims[0]))
        result_resp_2 = self.get_result(ds, data_api, field_names=("Measure", *data_api_test_params.two_dims))

        rows_1 = get_data_rows(result_resp_1)
        rows_2 = get_data_rows(result_resp_2)

        min_row_cnt = 2  # just an arbitrary number
        assert len(rows_1) > min_row_cnt
        assert len(rows_2) > len(rows_1)

        total_sum_1 = sum(float(row[0]) for row in rows_1)
        total_sum_2 = sum(float(row[0]) for row in rows_2)

        assert pytest.approx(total_sum_1) == total_sum_2

    def _test_contains(
        self,
        request,
        db: Db,
        saved_connection_id: str,
        dataset_params: dict,
        control_api: SyncHttpDatasetApiV1,
        data_api: SyncHttpDataApiV2,
        filter_op: WhereClauseOperation,
    ) -> None:
        columns = [
            C("int_value", UserDataType.integer, vg=lambda rn, **kwargs: rn),
            C("array_int_value", UserDataType.array_int, vg=lambda rn, **kwargs: [i for i in reversed(range(rn))]),
            C(
                "array_str_value",
                UserDataType.array_str,
                vg=lambda rn, **kwargs: [str(i) if i != 5 else None for i in reversed(range(rn))],
            ),
            C(
                "array_float_value",
                UserDataType.array_float,
                vg=lambda rn, **kwargs: [i / 100.0 for i in reversed(range(rn))],
            ),
        ]
        db_table = make_table(db, columns=columns)
        request.addfinalizer(functools.partial(db.drop_table, db_table.table))

        params = self.get_dataset_params(dataset_params, db_table)
        ds = self.make_basic_dataset(control_api, connection_id=saved_connection_id, dataset_params=params)

        def check_filter(field_title: str, filter_value: Any) -> None:
            result_resp = data_api.get_result(
                dataset=ds,
                fields=[
                    ds.find_field(title="int_value"),
                    ds.find_field(title=field_title),
                ],
                filters=[
                    ds.find_field(title=field_title).filter(
                        op=filter_op,
                        values=[filter_value],
                    ),
                ],
            )

            data_rows = get_data_rows(result_resp)
            assert data_rows, repr(filter_value)
            for row in data_rows:
                result_array = json.loads(row[1])
                if filter_op == WhereClauseOperation.CONTAINS:
                    assert filter_value in result_array
                elif filter_op == WhereClauseOperation.NOTCONTAINS:
                    assert filter_value not in result_array
                else:
                    raise ValueError(f"Unknown operation {filter_op}")

        check_filter("array_int_value", 3)
        check_filter("array_str_value", "3")
        check_filter("array_float_value", 0.03)
        check_filter("array_str_value", None)

    @for_features(array_support)
    def test_array_contains_filter(
        self,
        db: Db,
        saved_connection_id: str,
        dataset_params: dict,
        control_api: SyncHttpDatasetApiV1,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        self._test_contains(
            db, saved_connection_id, dataset_params, control_api, data_api, WhereClauseOperation.CONTAINS
        )

    @for_features(array_support)
    def test_array_not_contains_filter(
        self,
        db: Db,
        saved_connection_id: str,
        dataset_params: dict,
        control_api: SyncHttpDatasetApiV1,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        self._test_contains(
            db, saved_connection_id, dataset_params, control_api, data_api, WhereClauseOperation.NOTCONTAINS
        )

    @pytest.mark.parametrize(
        "field_title,filter_field_title,is_numeric",
        (
            ("array_int_value", "int_value", True),
            ("array_str_value", "str_value", False),
            ("array_str_value", "concat_const", False),
            ("array_str_value", "concat_field", False),
            ("array_float_value", "float_value", True),
            ("array_float_value", "none_value", False),
        ),
    )
    @for_features(array_support)
    def test_array_contains_field(
        self,
        request,
        db: Db,
        saved_connection_id: str,
        dataset_params: dict,
        control_api: SyncHttpDatasetApiV1,
        data_api: SyncHttpDataApiV2,
        field_title: str,
        filter_field_title: str,
        is_numeric: bool,
    ) -> None:
        columns = [
            C("int_value", UserDataType.integer, vg=lambda rn, **kwargs: 3),
            C("str_value", UserDataType.string, vg=lambda rn, **kwargs: "3"),
            C("float_value", UserDataType.float, vg=lambda rn, **kwargs: 0.03),
            C("none_value", UserDataType.float, vg=lambda rn, **kwargs: None),
            C("array_int_value", UserDataType.array_int, vg=lambda rn, **kwargs: [i for i in reversed(range(rn))]),
            C("array_str_value", UserDataType.array_str, vg=lambda rn, **kwargs: [str(i) for i in reversed(range(rn))]),
            C(
                "array_float_value",
                UserDataType.array_float,
                vg=lambda rn, **kwargs: [i / 100.0 if i != 5 else None for i in reversed(range(rn))],
            ),
        ]
        db_table = make_table(db, columns=columns)
        request.addfinalizer(functools.partial(db.drop_table, db_table.table))
        params = self.get_dataset_params(dataset_params, db_table)
        ds = self.make_basic_dataset(control_api, connection_id=saved_connection_id, dataset_params=params)

        filter_name = filter_field_title + "_ff"
        ds.result_schema["concat_const"] = ds.field(formula="CONCAT('3', '')")
        ds.result_schema["concat_field"] = ds.field(formula="CONCAT([str_value], '')")
        ds.result_schema[filter_name] = ds.field(formula=f"CONTAINS([{field_title}], [{filter_field_title}])")

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title=filter_field_title),
                ds.find_field(title=field_title),
            ],
            filters=[
                ds.find_field(title=filter_name).filter(
                    op=WhereClauseOperation.EQ,
                    values=[True],
                ),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == 200, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert data_rows, filter_field_title
        for row in data_rows:
            result_value = json.loads(row[0]) if is_numeric else row[0]
            result_array = json.loads(row[1])
            assert result_value in result_array

    def test_dates(
        self, saved_dataset: Dataset, data_api_test_params: DataApiTestParams, data_api: SyncHttpDataApiV2
    ) -> None:
        ds = saved_dataset
        new_field_name = "New shiny field"
        ds.result_schema[new_field_name] = ds.field(
            formula=f"IF [{data_api_test_params.date_field}] > DATE('2020-01-01') THEN 1 ELSE 2 END"
        )
        ds.result_schema[new_field_name].cast = UserDataType.float
        result_resp = self.get_result(ds, data_api, field_names=(data_api_test_params.date_field, new_field_name))
        assert result_resp.status_code == 200, result_resp.json

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[ds.find_field(title=data_api_test_params.date_field)],
            filters=[
                ds.find_field(title=data_api_test_params.date_field).filter(
                    op=WhereClauseOperation.GT,
                    values=["1990-01-01"],
                )
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == 200, result_resp.json
        assert get_data_rows(result_resp)

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[ds.find_field(title=data_api_test_params.date_field)],
            filters=[
                ds.find_field(title=data_api_test_params.date_field).filter(
                    op=WhereClauseOperation.BETWEEN,
                    values=["1990-01-01", "2023-10-02"],
                )
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == 200, result_resp.json
        assert get_data_rows(result_resp)

    def test_get_result_with_formula_in_where(
        self, saved_dataset: Dataset, data_api_test_params: DataApiTestParams, data_api: SyncHttpDataApiV2
    ) -> None:
        ds = saved_dataset
        ds.result_schema["Summable Percentage"] = ds.field(formula=f"[{data_api_test_params.summable_field}] * 100")

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="Summable Percentage"),
            ],
            filters=[
                ds.find_field(title="Summable Percentage").filter(WhereClauseOperation.GT, [30]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == 200, result_resp.json
        data_rows = get_data_rows(result_resp)
        assert data_rows
        assert all(float(row[0]) > 30 for row in data_rows)

    def test_get_result_with_string_filter_operations_for_numbers(
        self, saved_dataset: Dataset, data_api_test_params: DataApiTestParams, data_api: SyncHttpDataApiV2
    ) -> None:
        ds = saved_dataset
        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title=data_api_test_params.summable_field),
            ],
            filters=[
                ds.find_field(title=data_api_test_params.summable_field).filter(WhereClauseOperation.ICONTAINS, ["2"]),
            ],
            fail_ok=True,
        )
        assert result_resp.status_code == 200, result_resp.json

        data_rows = get_data_rows(result_resp)
        values: set[str] = {row[0] for row in data_rows}
        assert len(values) > 1  # we just need to make sure there are several different values
        assert all("2" in value for value in values), values


class DefaultConnectorDataGroupByFormulaTestSuite(StandardizedDataApiTestBase, RegulatedTestCase):
    def test_ordered_result(
        self,
        saved_dataset: Dataset,
        data_api_test_params: DataApiTestParams,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        ds = saved_dataset
        grouped_resp = self.get_result_ordered(
            ds,
            data_api,
            field_names=(data_api_test_params.two_dims[0], data_api_test_params.distinct_field),
            order_by=(data_api_test_params.distinct_field,),
        )
        grouped_rows = get_data_rows(grouped_resp)

        min_row_cnt = 5  # just an arbitrary number
        assert len(grouped_rows) > min_row_cnt

    def test_complex_result(
        self,
        saved_dataset: Dataset,
        data_api_test_params: DataApiTestParams,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        ds = saved_dataset
        ds.result_schema["LengthField"] = ds.field(formula=f"LEN([{data_api_test_params.distinct_field}])")

        grouped_resp = self.get_result_ordered(
            ds, data_api, field_names=(data_api_test_params.two_dims[0], "LengthField"), order_by=("LengthField",)
        )
        grouped_rows = get_data_rows(grouped_resp)

        min_row_cnt = 5  # just an arbitrary number
        assert len(grouped_rows) > min_row_cnt


class DefaultConnectorDataRangeTestSuite(StandardizedDataApiTestBase, RegulatedTestCase):
    def test_basic_range(
        self,
        saved_dataset: Dataset,
        data_api_test_params: DataApiTestParams,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        ds = saved_dataset

        range_resp = self.get_range(ds, data_api, field_name=data_api_test_params.range_field)
        range_rows = get_data_rows(range_resp)
        min_val, max_val = float(range_rows[0][0]), float(range_rows[0][1])
        assert max_val > min_val


class DefaultConnectorDataDistinctTestSuite(StandardizedDataApiTestBase, RegulatedTestCase):
    def test_basic_distinct(
        self,
        saved_dataset: Dataset,
        data_api_test_params: DataApiTestParams,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        ds = saved_dataset

        distinct_resp = self.get_distinct(ds, data_api, field_name=data_api_test_params.distinct_field)
        distinct_rows = get_data_rows(distinct_resp)
        min_distinct_row_cnt = 5  # just an arbitrary number
        assert len(distinct_rows) > min_distinct_row_cnt
        values = [row[0] for row in distinct_rows]
        assert len(set(values)) == len(values), "Values are not unique"

    def test_distinct_with_nonexistent_filter(
        self,
        saved_dataset: Dataset,
        data_api_test_params: DataApiTestParams,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        ds = saved_dataset

        distinct_resp = data_api.get_distinct(
            dataset=ds,
            field=ds.find_field(title=data_api_test_params.distinct_field),
            filters=[WhereClause(column="idontexist", operation=WhereClauseOperation.EQ, values=[0])],
            ignore_nonexistent_filters=True,
        )
        assert distinct_resp.status_code == 200, distinct_resp.json

        distinct_rows = get_data_rows(distinct_resp)
        min_distinct_row_cnt = 5  # just an arbitrary number
        assert len(distinct_rows) > min_distinct_row_cnt
        values = [row[0] for row in distinct_rows]
        assert len(set(values)) == len(values), "Values are not unique"

    def test_date_filter_distinct(
        self,
        request,
        db: Db,
        saved_connection_id: str,
        dataset_params: dict,
        control_api: SyncHttpDatasetApiV1,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        columns = [
            C(name="date_val", user_type=UserDataType.date, nullable=True),
        ]
        data = [
            {"date_val": datetime.date(2002, 1, 2)},
            {"date_val": datetime.date(2023, 4, 2)},
        ]
        db_table = make_table(db, columns=columns, data=data)
        request.addfinalizer(functools.partial(db.drop_table, db_table.table))

        params = self.get_dataset_params(dataset_params, db_table)
        ds = self.make_basic_dataset(control_api, connection_id=saved_connection_id, dataset_params=params)

        distinct_resp = data_api.get_distinct(
            dataset=ds,
            field=ds.find_field(title="date_val"),
            filters=[
                ds.find_field(title="date_val").filter(
                    op=WhereClauseOperation.ICONTAINS,
                    values=["2023-04-02"],
                ),
            ],
        )
        assert distinct_resp.status_code == 200, distinct_resp.json
        data_rows = get_data_rows(distinct_resp)
        assert len(data_rows) == 1

        distinct_resp = data_api.get_distinct(
            dataset=ds,
            field=ds.find_field(title="date_val"),
            filters=[
                ds.find_field(title="date_val").filter(
                    op=WhereClauseOperation.ICONTAINS,
                    values=["202"],
                ),
            ],
        )
        assert distinct_resp.status_code == 200, distinct_resp.json
        data_rows = get_data_rows(distinct_resp)
        assert len(data_rows) == 1


class DefaultConnectorDataPreviewTestSuite(StandardizedDataApiTestBase, RegulatedTestCase):
    def test_basic_preview(
        self,
        saved_dataset: Dataset,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        ds = saved_dataset

        preview_resp = self.get_preview(ds, data_api)
        preview_rows = get_data_rows(preview_resp)
        min_preview_row_cnt = 10  # just an arbitrary number
        assert len(preview_rows) > min_preview_row_cnt

lib/dl_api_lib_testing/dl_api_lib_testing/connector/data_api_suites.py
datalens-tech/datalens-backend
import abc

from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib_testing.dataset_base import DatasetTestBase
from dl_testing.regulated_test import RegulatedTestCase


class DefaultConnectorDatasetTestSuite(DatasetTestBase, RegulatedTestCase, metaclass=abc.ABCMeta):
    def check_basic_dataset(self, ds: Dataset) -> None:
        """Additional dataset checks can be defined here"""
        assert ds.id
        assert len(ds.result_schema)

        field_names = {field.title for field in ds.result_schema}
        assert {"category", "city"}.issubset(field_names)

    def test_create_basic_dataset(
        self,
        saved_dataset: Dataset,
    ) -> None:
        self.check_basic_dataset(saved_dataset)

lib/dl_api_lib_testing/dl_api_lib_testing/connector/dataset_suite.py
datalens-tech/datalens-backend
import abc
from typing import (
    Any,
    Optional,
)

from aiohttp.test_utils import (
    ClientResponse,
    TestClient,
)

from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_api_lib_testing.data_api_base import DataApiTestBase


class DashSQLTestBase(ConnectionTestBase, DataApiTestBase, metaclass=abc.ABCMeta):
    async def get_dashsql_response(
        self,
        data_api_aio: TestClient,
        conn_id: str,
        query: str,
        fail_ok: bool = False,
        params: Optional[dict] = None,
        connector_specific_params: Optional[dict] = None,
        headers: Optional[dict[str, str]] = None,
    ) -> ClientResponse:
        request_body: dict[str, Any] = {"sql_query": query}
        if params:
            request_body["params"] = params
        if connector_specific_params:
            request_body["connector_specific_params"] = connector_specific_params

        resp = await data_api_aio.post(f"api/v1/connections/{conn_id}/dashsql", json=request_body, headers=headers)
        resp_data = await resp.json()
        if not fail_ok:
            assert resp.status == 200, resp_data
        return resp

lib/dl_api_lib_testing/dl_api_lib_testing/dashsql_base.py
datalens-tech/datalens-backend
import abc
import asyncio
from http import HTTPStatus
from typing import (
    ClassVar,
    Generator,
    Iterable,
    NamedTuple,
)

from aiohttp import web
from aiohttp.pytest_plugin import AiohttpClient
from aiohttp.test_utils import TestClient
import pytest

from dl_api_client.dsmaker.api.data_api import (
    HttpDataApiResponse,
    SyncHttpDataApiV2,
)
from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib.app.data_api.app import DataApiAppFactory
from dl_api_lib.app_settings import DataApiAppSettings
from dl_api_lib_testing.app import (
    RedisSettingMaker,
    TestingDataApiAppFactory,
)
from dl_api_lib_testing.base import ApiTestBase
from dl_api_lib_testing.client import (
    TestClientConverterAiohttpToFlask,
    WrappedAioSyncApiClient,
)
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.dataset_base import DatasetTestBase
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.rqe import RQEConfig
from dl_constants.enums import ConnectionType
from dl_core.components.ids import FieldIdGeneratorType
from dl_core_testing.database import DbTable


class DataApiTestParams(NamedTuple):
    two_dims: tuple[str, str]
    summable_field: str
    range_field: str
    distinct_field: str
    date_field: str


class DataApiTestBase(ApiTestBase, metaclass=abc.ABCMeta):
    mutation_caches_on: ClassVar[bool] = True
    data_caches_on: ClassVar[bool] = True

    @pytest.fixture
    def loop(self, event_loop: asyncio.AbstractEventLoop) -> Generator[asyncio.AbstractEventLoop, None, None]:
        asyncio.set_event_loop(event_loop)
        yield event_loop
        # Attempt to cover an old version of pytest-asyncio:
        # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
        asyncio.set_event_loop_policy(None)

    @classmethod
    def create_data_api_settings(
        cls,
        bi_test_config: ApiTestEnvironmentConfiguration,
        rqe_config_subprocess: RQEConfig,
    ) -> DataApiAppSettings:
        core_test_config = bi_test_config.core_test_config
        us_config = core_test_config.get_us_config()
        redis_setting_maker = RedisSettingMaker(bi_test_config=bi_test_config)

        return DataApiAppSettings(
            SENTRY_ENABLED=False,
            US_BASE_URL=us_config.us_host,
            US_MASTER_TOKEN=us_config.us_master_token,
            CRYPTO_KEYS_CONFIG=core_test_config.get_crypto_keys_config(),
            # TODO FIX: Configure caches
            CACHES_ON=cls.data_caches_on,
            SAMPLES_CH_HOSTS=(),
            RQE_CONFIG=rqe_config_subprocess,
            BI_API_CONNECTOR_WHITELIST=bi_test_config.get_api_library_config().api_connector_ep_names,
            CORE_CONNECTOR_WHITELIST=core_test_config.get_core_library_config().core_connector_ep_names,
            MUTATIONS_CACHES_ON=cls.mutation_caches_on,
            CACHES_REDIS=redis_setting_maker.get_redis_settings_cache(),
            BI_COMPENG_PG_ON=cls.bi_compeng_pg_on,
            BI_COMPENG_PG_URL=bi_test_config.bi_compeng_pg_url,
            FIELD_ID_GENERATOR_TYPE=FieldIdGeneratorType.suffix,
            FILE_UPLOADER_BASE_URL=f"{bi_test_config.file_uploader_api_host}:{bi_test_config.file_uploader_api_port}",
            FILE_UPLOADER_MASTER_TOKEN="qwerty",
            QUERY_PROCESSING_MODE=cls.query_processing_mode,
        )  # type: ignore

    @pytest.fixture(scope="function")
    def data_api_app_settings(
        self,
        bi_test_config: ApiTestEnvironmentConfiguration,
        rqe_config_subprocess: RQEConfig,
    ) -> DataApiAppSettings:
        return self.create_data_api_settings(
            bi_test_config=bi_test_config,
            rqe_config_subprocess=rqe_config_subprocess,
        )

    @pytest.fixture(scope="function")
    def data_api_app_factory(self, data_api_app_settings: DataApiAppSettings) -> DataApiAppFactory:
        return TestingDataApiAppFactory(settings=data_api_app_settings)

    @pytest.fixture(scope="function")
    def data_api_app(
        self,
        data_api_app_factory: DataApiAppFactory,
        connectors_settings: dict[ConnectionType, ConnectorSettingsBase],
    ) -> web.Application:
        return data_api_app_factory.create_app(
            connectors_settings=connectors_settings,
        )

    @pytest.fixture(scope="function")
    def data_api_lowlevel_aiohttp_client(
        self,
        loop: asyncio.AbstractEventLoop,
        aiohttp_client: AiohttpClient,
        data_api_app: web.Application,
    ) -> TestClient:
        return loop.run_until_complete(aiohttp_client(data_api_app))

    @pytest.fixture(scope="function")
    def data_api_sync_client(
        self,
        loop: asyncio.AbstractEventLoop,
        data_api_lowlevel_aiohttp_client: TestClient,
    ) -> WrappedAioSyncApiClient:
        return WrappedAioSyncApiClient(
            int_wrapped_client=TestClientConverterAiohttpToFlask(
                loop=loop,
                aio_client=data_api_lowlevel_aiohttp_client,
            ),
        )

    @pytest.fixture(scope="function")
    def data_api(self, data_api_sync_client: SyncHttpClientBase) -> SyncHttpDataApiV2:
        return SyncHttpDataApiV2(client=data_api_sync_client)


class StandardizedDataApiTestBase(DataApiTestBase, DatasetTestBase, metaclass=abc.ABCMeta):
    @pytest.fixture(scope="class")
    def data_api_test_params(self) -> DataApiTestParams:
        # This default is defined for the sample table
        return DataApiTestParams(
            two_dims=("category", "city"),
            summable_field="sales",
            range_field="sales",
            distinct_field="city",
            date_field="order_date",
        )

    def get_dataset_params(self, dataset_params: dict, db_table: DbTable) -> dict:
        return dataset_params | dict(
            parameters=dataset_params.get("parameters", {})
            | dict(
                schema_name=db_table.schema,
                table_name=db_table.name,
            )
        )

    def get_result(
        self,
        ds: Dataset,
        data_api: SyncHttpDataApiV2,
        field_names: Iterable[str],
    ) -> HttpDataApiResponse:
        data_resp = data_api.get_result(
            dataset=ds,
            fields=[ds.find_field(title=field_name) for field_name in field_names],
        )
        assert data_resp.status_code == HTTPStatus.OK, data_resp.response_errors
        return data_resp

    def get_range(
        self,
        ds: Dataset,
        data_api: SyncHttpDataApiV2,
        field_name: str,
    ) -> HttpDataApiResponse:
        data_resp = data_api.get_value_range(
            dataset=ds,
            field=ds.find_field(title=field_name),
        )
        assert data_resp.status_code == HTTPStatus.OK, data_resp.response_errors
        return data_resp

    def get_distinct(
        self,
        ds: Dataset,
        data_api: SyncHttpDataApiV2,
        field_name: str,
    ) -> HttpDataApiResponse:
        data_resp = data_api.get_distinct(
            dataset=ds,
            field=ds.find_field(title=field_name),
        )
        assert data_resp.status_code == HTTPStatus.OK, data_resp.response_errors
        return data_resp

    def get_result_ordered(
        self,
        ds: Dataset,
        data_api: SyncHttpDataApiV2,
        field_names: Iterable[str],
        order_by: Iterable[str],
    ) -> HttpDataApiResponse:
        data_resp = data_api.get_result(
            dataset=ds,
            fields=[ds.find_field(title=field_name) for field_name in field_names],
            order_by=[ds.find_field(title=field_name) for field_name in order_by],
        )
        assert data_resp.status_code == HTTPStatus.OK, data_resp.response_errors
        return data_resp

    def get_preview(self, ds: Dataset, data_api: SyncHttpDataApiV2) -> HttpDataApiResponse:
        data_resp = data_api.get_preview(dataset=ds)
        assert data_resp.status_code == HTTPStatus.OK, data_resp.response_errors
        return data_resp

lib/dl_api_lib_testing/dl_api_lib_testing/data_api_base.py
datalens-tech/datalens-backend
import abc
from typing import Generator

import pytest

from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib_testing.connection_base import ConnectionTestBase


class DatasetTestBase(ConnectionTestBase, metaclass=abc.ABCMeta):
    @abc.abstractmethod
    @pytest.fixture(scope="session")
    def dataset_params(self) -> dict:
        raise NotImplementedError

    def make_basic_dataset(
        self,
        control_api: SyncHttpDatasetApiV1,
        connection_id: str,
        dataset_params: dict,
    ) -> Dataset:
        ds = Dataset()
        ds.sources["source_1"] = ds.source(
            connection_id=connection_id,
            **dataset_params,
        )

        ds.source_avatars["avatar_1"] = ds.sources["source_1"].avatar()
        ds = control_api.apply_updates(dataset=ds, fail_ok=False).dataset
        ds = control_api.save_dataset(dataset=ds).dataset
        return ds

    @pytest.fixture(scope="function")
    def saved_dataset(
        self,
        control_api: SyncHttpDatasetApiV1,
        saved_connection_id: str,
        dataset_params: dict,
    ) -> Generator[Dataset, None, None]:
        ds = self.make_basic_dataset(
            control_api=control_api,
            connection_id=saved_connection_id,
            dataset_params=dataset_params,
        )
        yield ds
        control_api.delete_dataset(dataset_id=ds.id, fail_ok=False)

lib/dl_api_lib_testing/dl_api_lib_testing/dataset_base.py
datalens-tech/datalens-backend
from dl_core_testing.database import DbTable
from dl_core_testing.dataset import get_created_from


def data_source_settings_from_table(table: DbTable):
    source_type = get_created_from(db=table.db)
    data = {  # this still requires connection_id to be defined
        "source_type": source_type,
        "parameters": {
            "table_name": table.name,
            # FIXME: learn to include or exclude it by looking at the db type
            "db_name": table.db.name,
        },
    }

    if table.schema:
        data["parameters"]["schema_name"] = table.schema

    return data

lib/dl_api_lib_testing/dl_api_lib_testing/helpers/data_source.py
datalens-tech/datalens-backend
import datetime
from typing import (
    Any,
    Callable,
)


def read_date(s: str) -> datetime.date:
    try:
        return datetime.date.fromisoformat(s)
    except ValueError:
        return datetime.datetime.fromisoformat(s).date()


def check_at_date_data(
    data_rows: list[list[Any]],
    date_idx: int,
    value_idx: int,
    ago_idx: int,
    ago_date_callable: Callable[[datetime.date], datetime.date],
    allow_missing_date_values: bool = False,
) -> None:
    assert len(data_rows) > 0
    value_by_date = {read_date(row[date_idx]): row[value_idx] for row in data_rows}
    rows_checked = 0

    for row_idx, row in enumerate(data_rows):
        cur_date = read_date(row[date_idx])
        ago_date = ago_date_callable(cur_date)
        expected_ago_value = value_by_date.get(ago_date)
        actual_ago_value = row[ago_idx]

        if expected_ago_value is None:
            if allow_missing_date_values:
                pass  # Do not check in this case
            else:
                assert actual_ago_value is None
        else:
            assert actual_ago_value == expected_ago_value

        rows_checked += 1

    # Make sure that rows were checked
    assert rows_checked > 5


def check_ago_data(
    data_rows: list[list[Any]],
    date_idx: int,
    value_idx: int,
    ago_idx: int,
    day_offset: int,
    allow_missing_date_values: bool = False,
) -> None:
    def ago_date_callable(cur_date: datetime.date) -> datetime.date:  # noqa
        return cur_date - datetime.timedelta(days=day_offset)

    check_at_date_data(
        data_rows=data_rows,
        date_idx=date_idx,
        value_idx=value_idx,
        ago_idx=ago_idx,
        ago_date_callable=ago_date_callable,
        allow_missing_date_values=allow_missing_date_values,
    )

lib/dl_api_lib_testing/dl_api_lib_testing/helpers/lookup_checkers.py
datalens-tech/datalens-backend
from typing import Callable

import attr
from pytest import MonkeyPatch

from dl_query_processing.compilation.primitives import CompiledMultiQueryBase
from dl_query_processing.translation.multi_level_translator import MultiLevelQueryTranslator


@attr.s
class MultiQueryInterceptor:
    _mpatch: MonkeyPatch = attr.ib(kw_only=True)
    _callback: Callable[[CompiledMultiQueryBase], None] = attr.ib(kw_only=True)
    _intercepted: bool = attr.ib(init=False, default=False)

    @property
    def intercepted(self) -> bool:
        return self._intercepted

    def __attrs_post_init__(self) -> None:
        self._prepare()

    def _prepare(self) -> None:
        self._mpatch.setattr(MultiLevelQueryTranslator, "_log_query_complexity_stats", self._log_query_complexity_stats)

    def _log_query_complexity_stats(self, compiled_multi_query: CompiledMultiQueryBase) -> None:
        self._callback(compiled_multi_query)
        self._intercepted = True


def count_joins(multi_query: CompiledMultiQueryBase) -> int:
    result = 0
    for query in multi_query.iter_queries():
        result += len(query.join_on)
    return result

lib/dl_api_lib_testing/dl_api_lib_testing/helpers/multi_query.py
datalens-tech/datalens-backend
from pytest import Config

from dl_api_lib.loader import load_api_lib
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.initialization import initialize_core_test


def initialize_api_lib_test(pytest_config: Config, api_test_config: ApiTestEnvironmentConfiguration) -> None:
    load_api_lib(api_lib_config=api_test_config.get_api_library_config())
    initialize_core_test(pytest_config=pytest_config, core_test_config=api_test_config.core_test_config)

lib/dl_api_lib_testing/dl_api_lib_testing/initialization.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import logging
import math
import time
from typing import Optional


LOGGER = logging.getLogger(__name__)


class LatencyTracker:
    bins_base: float = 1.5  # logarithmic binning exponent
    sleep_interval_sec: float = 100 / 1000
    stats_log_interval_sec: float = 5.0
    min_warn_interval_sec: float = 2.5

    def __init__(self) -> None:
        # duration lower bound (milliseconds) -> count
        self._bins: dict[int, int] = {}
        self._last_log_time = time.monotonic()
        self._logger = LOGGER.getChild(self.__class__.__name__)
        self._task: Optional[asyncio.Task] = None

    def _to_bin(self, value: float) -> int:
        return int(self.bins_base ** math.floor(math.log(value) / math.log(self.bins_base)))

    def _log_duration(self, duration_sec: float) -> None:
        duration_msec = duration_sec * 1000
        duration_bin = self._to_bin(duration_msec)
        self._bins[duration_bin] = self._bins.setdefault(duration_bin, 0) + 1

    def _log_stats(self) -> None:
        bins = self._bins
        self._bins = {}
        bins_s = ", ".join(f">={dur}ms: {cnt}" for dur, cnt in bins.items())
        loop = asyncio.get_event_loop()
        loop_scheduled = len(getattr(loop, "_scheduled", None) or ())
        loop_ready = len(getattr(loop, "_ready", None) or ())
        self._logger.debug(
            "Latency stats: %s, len(loop._ready)=%s, len(loop._scheduled)=%s",
            bins_s,
            loop_ready,
            loop_scheduled,
            extra=dict(
                latency_stats=bins,
                loop_ready=loop_ready,
                loop_scheduled=loop_scheduled,
            ),
        )

    def _maybe_log_stats(self) -> None:
        now = time.monotonic()
        if now - self._last_log_time < self.stats_log_interval_sec:
            return
        self._last_log_time = now
        self._log_stats()

    def _log_warn(self, duration_sec: float) -> None:
        # The idea is to add a log line explicitly,
        # to look for logs before it,
        # to maybe find whatever caused the latency.
        self._logger.warning("High latency: %.3fs", duration_sec)

    def _handle_duration(self, duration_sec: float) -> None:
        self._log_duration(duration_sec)
        self._maybe_log_stats()
        if duration_sec >= self.min_warn_interval_sec:
            self._log_warn(duration_sec)

    async def _task_main(self) -> None:
        while True:
            ts0 = time.monotonic()
            await asyncio.sleep(self.sleep_interval_sec)
            ts1 = time.monotonic()
            self._handle_duration(ts1 - ts0 - self.sleep_interval_sec)

    async def run_task(self) -> None:
        # Does not have to be `async`, but it is more consistent this way.
        asyncio.create_task(self._task_main())


async def atest_latency_tracker() -> None:
    # Not a unit test because this is long and non-checking.
    logging.basicConfig(
        level=logging.DEBUG,
        format="%(asctime)s: %(levelname)-13s: %(name)s: %(message)s",
    )
    tracker = LatencyTracker()
    await tracker.run_task()

    part_one_duration_sec = 7.0
    LOGGER.debug("Sleeping in small steps for %rs", part_one_duration_sec)
    part_one_end = time.monotonic() + part_one_duration_sec
    while time.monotonic() < part_one_end:
        await asyncio.sleep(0.001)
        time.sleep(0.005)

    for idx in range(1, 3000, 50):
        await asyncio.sleep(0.001)
        LOGGER.debug("Sleeping for %r msec", idx)
        time.sleep(idx / 1000)


def test_latency_tracker() -> None:
    return asyncio.run(atest_latency_tracker())


if __name__ == "__main__":
    test_latency_tracker()

lib/dl_app_tools/dl_app_tools/aio_latency_tracking.py
datalens-tech/datalens-backend
from __future__ import annotations

from .format import JsonFormatter


__all__ = ("JsonFormatter",)

lib/dl_app_tools/dl_app_tools/log/__init__.py
datalens-tech/datalens-backend
# coding: utf-8
from __future__ import annotations

from collections import defaultdict
import contextvars
import sys

import six


_LOG_CONTEXT = contextvars.ContextVar("_LOG_CONTEXT", default=None)


def force_unicode(obj):  # type: ignore  # TODO: fix
    if not isinstance(obj, six.string_types):
        obj = repr(obj)
    return obj.decode(sys.getdefaultencoding())


class ContextFormatter:
    @classmethod
    def is_context_exist(cls):  # type: ignore  # TODO: fix
        return _LOG_CONTEXT.get() is not None

    @classmethod
    def get_or_create_logging_context(cls):  # type: ignore  # TODO: fix
        if not cls.is_context_exist():
            _LOG_CONTEXT.set(defaultdict(list))  # type: ignore  # TODO: fix
        return _LOG_CONTEXT.get()


def reset_context():  # type: ignore  # TODO: fix
    _LOG_CONTEXT.set(defaultdict(list))  # type: ignore  # TODO: fix


def put_to_context(key, value):  # type: ignore  # TODO: fix
    ctx = ContextFormatter.get_or_create_logging_context()
    ctx[key].append(value)


def pop_from_context(key):  # type: ignore  # TODO: fix
    if not ContextFormatter.is_context_exist():
        return
    ctx = ContextFormatter.get_or_create_logging_context()
    if key not in ctx:
        return
    value = ctx[key].pop()
    if not ctx[key]:
        del ctx[key]
    return value


class LogContext(object):
    context = None

    def __init__(self, **context):  # type: ignore  # TODO: fix
        self.context = context

    def __enter__(self):  # type: ignore  # TODO: fix
        for key, value in self.context.items():  # type: ignore  # TODO: fix
            put_to_context(key, value)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):  # type: ignore  # TODO: fix
        for key in self.context:  # type: ignore  # TODO: fix
            pop_from_context(key)
        return False


log_context = LogContext


def get_log_context():  # type: ignore  # TODO: fix
    if not ContextFormatter.is_context_exist():
        return {}
    ctx = ContextFormatter.get_or_create_logging_context()
    return {key: values[-1] for key, values in ctx.items()}

lib/dl_app_tools/dl_app_tools/log/context.py
datalens-tech/datalens-backend
# coding:utf-8

from __future__ import annotations

import json
import logging
import os

from .context import get_log_context


IS_DEPLOY = "DEPLOY_BOX_ID" in os.environ

# To catch the 'extra'-added attributes.
DEFAULT_RECORD_ATTRS = frozenset(
    (
        "name",
        "msg",
        "args",
        "levelname",
        "levelno",
        "pathname",
        "filename",
        "module",
        "exc_info",
        "exc_text",
        "stack_info",
        "lineno",
        "funcName",
        "created",
        "msecs",
        "relativeCreated",
        "thread",
        "threadName",
        "processName",
        "process",
        # Added around here:
        "message",
        # Around here and in BI:
        "log_context",
        "dlog_context",
    )
)


def get_record_extra(record):  # type: ignore  # TODO: fix
    return {key: value for key, value in vars(record).items() if key not in DEFAULT_RECORD_ATTRS}


class JsonFormatter(logging.Formatter):
    """
    Formatting log records as JSON.

    Message will be formatted as JSON with the following structure:
    {
        "msg": "message",
        "stackTrace": "your very long stacktrace \n with newlines \n and all the things",
        "level": "WARN",
        "@fields": {
            "std": {
                "funcName": "get_values",
                "lineno": 274,
                "name": "mymicroservice.dao.utils",
            },
            "context": {
                "foo": "qwerty",
                "bar": 42,
            }
        }
    }

    Dict @fields.context will contain all the fields from log_context.
    """

    LOG_RECORD_USEFUL_FIELDS = ("funcName", "lineno", "name")

    def format(self, record):  # type: ignore  # TODO: fix
        record.message = record.getMessage()

        log_data = {
            "message": record.message,
            "level": record.levelname,
        }
        if IS_DEPLOY:
            log_data["levelStr"] = record.levelname
            log_data["loggerName"] = record.name
            log_data["level"] = record.levelno

        if record.exc_info:
            exc = logging.Formatter.formatException(self, record.exc_info)
            log_data["stackTrace"] = exc

        fields = {}

        standard_fields = self._get_standard_fields(record)
        if standard_fields:
            fields["std"] = standard_fields

        context = {}
        log_context_fields = record.log_context if hasattr(record, "log_context") else get_log_context()
        context.update(log_context_fields)
        context.update(get_record_extra(record))
        if context:
            fields["context"] = context

        if fields:
            log_data["@fields"] = fields

        return json.dumps(log_data, default=repr)

    def _get_standard_fields(self, record):  # type: ignore  # TODO: fix
        return {field: getattr(record, field) for field in self.LOG_RECORD_USEFUL_FIELDS if hasattr(record, field)}

lib/dl_app_tools/dl_app_tools/log/format.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextvars import ContextVar
import functools
import inspect
import itertools as it
import logging
import time
import traceback
from typing import (
    Any,
    Awaitable,
    Callable,
    Dict,
    Iterable,
    List,
    Optional,
    Tuple,
    Type,
    TypeVar,
    Union,
    cast,
)

import attr
import opentracing

from dl_utils.utils import get_type_full_name


LOGGER = logging.getLogger(__name__)


TTraceback = Any  # ?builtins?.?traceback?
TExcInfo = Tuple[Type[Exception], Exception, TTraceback]
TLogData = Dict[str, Any]
TStageList = Union[List[str], Tuple[str, ...], str]


@attr.s
class ProfileResult:
    """
    A mutable object that gets filled on context manager exit.

    Usage example:

        with GenericProfiler("some-profiled-block") as prof_result:
           some_code
        timing = prof_result.exec_time_sec
    """

    exec_time_sec: Optional[float] = attr.ib(default=None)
    log_data: Optional[TLogData] = attr.ib(default=None)
    exc_info: Optional[TExcInfo] = attr.ib(default=None)
    exc_value: Optional[Exception] = attr.ib(default=None)


class GenericProfiler:
    EVENT_CODE = "profiling_result"

    CTX_OUTER_STAGES: ContextVar[Tuple[str, ...]] = ContextVar("CTX_OUTER_STAGES")
    CTX_LOCAL_PROFILER_STACK: ContextVar[List[GenericProfiler]] = ContextVar("CTX_LOCAL_PROFILER_STACK")
    USE_CTX_VARS = True

    @classmethod
    def get_outer_stages(cls) -> Tuple[str, ...]:
        outer_stages = cls.CTX_OUTER_STAGES.get(None)
        if outer_stages is None:
            outer_stages = ()
            cls.CTX_OUTER_STAGES.set(outer_stages)
        return outer_stages

    @classmethod
    def get_profilers_stack(cls) -> List["GenericProfiler"]:
        profilers_stack = cls.CTX_LOCAL_PROFILER_STACK.get(None)
        if profilers_stack is None:
            profilers_stack = []
            cls.CTX_LOCAL_PROFILER_STACK.set(profilers_stack)
        return profilers_stack

    @classmethod
    def reset_outer_stages(cls, outer_stages: Optional[TStageList] = None) -> None:
        if isinstance(outer_stages, str):
            outer_stages = cls.load_stage_stack(outer_stages)
        elif outer_stages is None:
            outer_stages = tuple()
        elif isinstance(outer_stages, (list, tuple)):
            outer_stages = tuple(outer_stages)
        else:
            raise TypeError(f"Unexpected type of outer stages: {type(outer_stages)}")

        if cls.get_profilers_stack():
            raise RuntimeError("Trying to reset outer stages of profiler while profilers stack is not empty")

        cls.CTX_OUTER_STAGES.set(tuple(outer_stages))

    @classmethod
    def reset_all(cls, outer_stages: Optional[TStageList] = None) -> None:
        if cls.get_current_stages_stack():
            current_stack = cls.get_current_stages_stack_str()
            tb_limit = 5
            LOGGER.warning(
                "Attempt to reset profiling stack %s. Traceback (last %d): %s",
                current_stack,
                tb_limit,
                traceback.format_stack(limit=tb_limit),
            )

        cls.CTX_LOCAL_PROFILER_STACK.set([])
        cls.reset_outer_stages(outer_stages)

    @classmethod
    def get_current_stages_stack(cls) -> Tuple[str, ...]:
        return tuple(it.chain(cls.get_outer_stages(), (p.stage for p in cls.get_profilers_stack())))

    @classmethod
    def get_current_stages_stack_str(cls) -> str:
        return cls.dump_stage_stack(cls.get_current_stages_stack())

    @staticmethod
    def dump_stage_stack(stage_stack: Iterable[str]) -> str:
        return "/".join(stage_stack)

    @staticmethod
    def load_stage_stack(stage_stack_str: str) -> Tuple[str, ...]:
        return tuple(stage_stack_str.split("/"))

    def __init__(
        self,
        stage: str,
        extra_data: Optional[dict] = None,
        logger: Optional[logging.Logger] = None,
    ):
        self.logger = LOGGER if logger is None else logger  # type: logging.Logger
        self.stage = stage
        self.extra_data = {} if extra_data is None else extra_data

        self.start_time: Optional[float] = None
        self.end_time: Optional[float] = None
        self.profile_result: Optional[ProfileResult] = None
        self._ot_span_context: Optional[opentracing.scope.Scope] = None  # Will be filled on __enter__

    def _cleanup(self) -> None:
        self.start_time = None
        self.end_time = None
        self.profile_result = None

    def __enter__(self) -> ProfileResult:
        tracer = opentracing.global_tracer()
        self._ot_span_context = tracer.start_active_span(self.stage, child_of=tracer.active_span)
        self.get_profilers_stack().append(self)

        self.start_time = time.monotonic()

        self.profile_result = ProfileResult()
        self.pre_profile()
        return self.profile_result

    def __exit__(self, exc_type: Type[Exception], exc_val: Exception, exc_tb: TTraceback) -> None:
        try:
            self.post_profile()

            self.end_time = time.monotonic()

            if exc_type is None:
                exc_info = None
            else:
                exc_info = (exc_type, exc_val, exc_tb)
                assert self._ot_span_context is not None
                span = self._ot_span_context.span
                span.set_tag(opentracing.span.tags.ERROR, True)
                span.set_tag("bi.exc_type", get_type_full_name(exc_type))

            log_data = self._save_profiling_log(exc_type, exc_info=exc_info)
            self._update_profile_result(log_data=log_data, exc_info=exc_info)
            self._cleanup()
        finally:
            # Restore profiler stack in any case
            assert self._ot_span_context is not None
            self._ot_span_context.close()
            try:
                self.get_profilers_stack().pop()
            except IndexError:
                LOGGER.warning("Attempt to restore empty profiling stack on %s", self.stage)

    def pre_profile(self) -> None:
        pass

    def post_profile(self) -> None:
        pass

    def _save_profiling_log(self, exc_type: Type[Exception], exc_info: Optional[TExcInfo] = None) -> TLogData:
        assert self.end_time is not None
        assert self.start_time is not None
        time_diff = self.end_time - self.start_time

        # HAX: strvalue in context.exc_info; to be removed later.
        exc_info_text = None
        if exc_info is not None:
            exc_type, exc_val, _ = exc_info
            exc_info_text = traceback.format_exception_only(exc_type, exc_val)[-1]

        extra = dict(
            self.extra_data,
            event_code=self.EVENT_CODE,
            stage=self.stage,
            stage_stack=list(self.get_current_stages_stack()),
            stage_stack_str=self.get_current_stages_stack_str(),
            execution_time=int(round(time_diff * 1e7)),  # TODO: drop this weird stuff
            success=exc_type is None,
            exc_type=exc_type.__name__ if exc_type else None,
            exc_info_text=exc_info_text,
            # not in action fields:
            exec_time_sec=round(time_diff, 4),
        )

        log_exc_info: Optional[TExcInfo] = None
        if exc_info is not None:
            if isinstance(exc_info, tuple) and len(exc_info) == 3:
                log_exc_info = exc_info
            else:
                LOGGER.error("Invalid exc_info passed to _save_profiling_log: %r", exc_info)

        self.logger.info("Profiling results: %s in %.4fs", self.stage, time_diff, extra=extra, exc_info=log_exc_info)

        return extra

    def _update_profile_result(self, exc_info: Optional[TExcInfo] = None, log_data: Optional[TLogData] = None) -> None:
        res = self.profile_result
        assert res is not None

        if exc_info is not None:
            res.exc_info = exc_info
            res.exc_value = exc_info[1]
        if log_data is not None:
            res.log_data = log_data
            res.exec_time_sec = log_data["exec_time_sec"]


_GP_FUNC_RET_TV = TypeVar("_GP_FUNC_RET_TV")
_GP_FUNC_T = Callable[..., _GP_FUNC_RET_TV]


def generic_profiler(
    stage: str,
    extra_data: dict = None,
    logger: logging.Logger = None,
) -> Callable[[_GP_FUNC_T], _GP_FUNC_T]:
    def generic_profiler_wrap(func: _GP_FUNC_T) -> _GP_FUNC_T:
        if inspect.iscoroutinefunction(func):
            raise ValueError("This decorator shouldn't be used on coroutines; use `generic_profiler_async` instead")

        @functools.wraps(func)
        def generic_profiler_wrapper(*args: Any, **kwargs: Any) -> _GP_FUNC_RET_TV:
            with GenericProfiler(stage, extra_data=extra_data, logger=logger):
                return func(*args, **kwargs)

        return cast(_GP_FUNC_T, generic_profiler_wrapper)

    return generic_profiler_wrap


_GPA_CORO_RET_TV = TypeVar("_GPA_CORO_RET_TV")
_GPA_CORO_TV = TypeVar("_GPA_CORO_TV", bound=Callable[..., Awaitable[_GPA_CORO_RET_TV]])


def generic_profiler_async(
    stage: str, extra_data: dict = None, logger: logging.Logger = None
) -> Callable[[_GPA_CORO_TV], _GPA_CORO_TV]:
    def generic_profiler_wrap_async(coro: _GPA_CORO_TV) -> _GPA_CORO_TV:
        if not inspect.iscoroutinefunction(coro):
            raise ValueError("This decorator may only be applied to a coroutine; use `generic_profiler` instead")

        @functools.wraps(coro)
        async def generic_profiler_wrapper_async(*args: Any, **kwargs: Any) -> _GPA_CORO_RET_TV:
            with GenericProfiler(stage, extra_data=extra_data, logger=logger):
                return await coro(*args, **kwargs)

        return cast(_GPA_CORO_TV, generic_profiler_wrapper_async)

    return generic_profiler_wrap_async

lib/dl_app_tools/dl_app_tools/profiling_base.py
datalens-tech/datalens-backend
from __future__ import annotations


def test_some():
    assert True

lib/dl_app_tools/dl_app_tools_tests/unit/test_stuff.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Type

import attr

from dl_compeng_pg.compeng_aiopg.pool_aiopg import AiopgPoolWrapper
from dl_compeng_pg.compeng_aiopg.processor_aiopg import AiopgOperationProcessor
from dl_compeng_pg.compeng_pg_base.data_processor_service_pg import CompEngPgService
from dl_compeng_pg.compeng_pg_base.pool_base import BasePgPoolWrapper
from dl_core.data_processing.processing.processor import OperationProcessorAsyncBase


@attr.s
class AiopgCompEngService(CompEngPgService[AiopgPoolWrapper]):
    def _get_pool_wrapper_cls(self) -> Type[BasePgPoolWrapper]:
        return AiopgPoolWrapper

    def get_data_processor(self) -> OperationProcessorAsyncBase:
        return AiopgOperationProcessor(pg_pool=self.pool)

lib/dl_compeng_pg/dl_compeng_pg/compeng_aiopg/data_processor_service_aiopg.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    AsyncGenerator,
    ClassVar,
    Optional,
    Sequence,
    Union,
)

import aiopg.sa
import attr
from sqlalchemy.sql.base import Executable
from sqlalchemy.sql.selectable import Select

from dl_compeng_pg.compeng_pg_base.exec_adapter_base import PostgreSQLExecAdapterAsync
from dl_constants.enums import UserDataType
from dl_core.data_processing.prepared_components.primitives import PreparedMultiFromInfo
from dl_core.data_processing.streaming import (
    AsyncChunked,
    AsyncChunkedBase,
)


if TYPE_CHECKING:
    from dl_constants.types import TBIDataValue


DEFAULT_CHUNK_SIZE = 1000

LOGGER = logging.getLogger(__name__)


@attr.s
class AiopgExecAdapter(PostgreSQLExecAdapterAsync[aiopg.sa.SAConnection]):  # noqa
    _conn: aiopg.sa.SAConnection = attr.ib()
    _log: ClassVar[logging.Logger] = LOGGER

    async def _execute(self, query: Union[str, Executable]) -> None:
        """Execute query without fetching data"""
        await self._conn.execute(query)

    async def _execute_ddl(self, query: Union[str, Executable]) -> None:
        """Execute a DDL statement"""
        await self._execute(query)

    async def _execute_and_fetch(
        self,
        *,
        query: Union[Select, str],
        user_types: Sequence[UserDataType],
        chunk_size: int,
        joint_dsrc_info: Optional[PreparedMultiFromInfo] = None,
        query_id: str,
    ) -> AsyncChunkedBase[Sequence[TBIDataValue]]:
        async def chunked_data_gen() -> AsyncGenerator[list[list[TBIDataValue]], None]:
            """Fetch data in chunks"""

            result = await self._conn.execute(query)
            while True:
                chunk = []
                for row in await result.fetchmany(chunk_size):
                    assert len(row) == len(user_types), "user_types is not the same length as the data row"
                    chunk.append(
                        [
                            self._tt.cast_for_output(value=value, user_t=user_t)
                            for value, user_t in zip(row.as_tuple(), user_types)
                        ]
                    )
                if not chunk:
                    break
                yield chunk

        return AsyncChunked(chunked_data=chunked_data_gen())

    async def insert_data_into_table(
        self,
        *,
        table_name: str,
        names: Sequence[str],
        user_types: Sequence[UserDataType],
        data: AsyncChunkedBase,
    ) -> None:
        """Insert data into a table."""

        table = self._make_sa_table(table_name=table_name, names=names, user_types=user_types)
        self._log.info(f"Inserting data into table {table_name}")

        async for raw_chunk in data.chunks:
            chunk = []
            for row in raw_chunk:
                chunk.append(
                    [self._tt.cast_for_input(value=value, user_t=user_t) for value, user_t in zip(row, user_types)]
                )
            query = table.insert(values=chunk)
            await self._execute(query=query)

lib/dl_compeng_pg/dl_compeng_pg/compeng_aiopg/exec_adapter_aiopg.py
datalens-tech/datalens-backend
from dl_compeng_pg.compeng_aiopg.data_processor_service_aiopg import AiopgCompEngService
from dl_constants.enums import ProcessorType
from dl_core.data_processors.base.plugin import DataProcessorPlugin


class AiopgCompengPlugin(DataProcessorPlugin):
    data_processor_type = ProcessorType.AIOPG
    data_processor_service_cls = AiopgCompEngService

lib/dl_compeng_pg/dl_compeng_pg/compeng_aiopg/plugin.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    Type,
)

import aiopg.sa
import attr

from dl_compeng_pg.compeng_pg_base.pool_base import (
    DEFAULT_OPERATION_TIMEOUT,
    DEFAULT_POOL_MAX_SIZE,
    DEFAULT_POOL_MIN_SIZE,
    BasePgPoolWrapper,
)


@attr.s
class AiopgPoolWrapper(BasePgPoolWrapper):
    _pool: Optional[aiopg.sa.Engine] = attr.ib()

    @classmethod
    async def connect(
        cls: Type["AiopgPoolWrapper"],
        url: str,
        pool_min_size: int = DEFAULT_POOL_MIN_SIZE,  # Initial pool size
        pool_max_size: int = DEFAULT_POOL_MAX_SIZE,  # Maximum pool size
        operation_timeout: float = DEFAULT_OPERATION_TIMEOUT,  # SQL operation timeout
    ) -> "AiopgPoolWrapper":
        pool = await aiopg.sa.create_engine(
            url,
            minsize=pool_min_size,
            maxsize=pool_max_size,
            timeout=operation_timeout,
        )
        return cls(pool=pool)

    @property
    def pool(self) -> aiopg.sa.Engine:
        if self._pool is None:
            raise RuntimeError("Aiopg pool is closed")
        return self._pool

    async def disconnect(self) -> None:
        pool = self._pool
        self._pool = None
        pool.terminate()  # type: ignore  # TODO: fix
        await pool.wait_closed()  # type: ignore  # TODO: fix

lib/dl_compeng_pg/dl_compeng_pg/compeng_aiopg/pool_aiopg.py
datalens-tech/datalens-backend
from __future__ import annotations

import aiopg.sa
import attr

from dl_compeng_pg.compeng_aiopg.exec_adapter_aiopg import AiopgExecAdapter
from dl_compeng_pg.compeng_aiopg.pool_aiopg import AiopgPoolWrapper
from dl_compeng_pg.compeng_pg_base.processor_base import PostgreSQLOperationProcessor


@attr.s
class AiopgOperationProcessor(PostgreSQLOperationProcessor[AiopgExecAdapter, AiopgPoolWrapper, aiopg.sa.SAConnection]):
    async def start(self) -> None:
        self._pg_conn = await self._pg_pool.pool.acquire()
        self._pgex_adapter = AiopgExecAdapter(conn=self._pg_conn)  # type: ignore  # TODO: fix

    async def end(self) -> None:
        self._pgex_adapter = None
        if self._pg_conn is not None:
            await self._pg_conn.close()
        self._pg_conn = None

lib/dl_compeng_pg/dl_compeng_pg/compeng_aiopg/processor_aiopg.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Type

import attr

from dl_compeng_pg.compeng_asyncpg.pool_asyncpg import AsyncpgPoolWrapper
from dl_compeng_pg.compeng_asyncpg.processor_asyncpg import AsyncpgOperationProcessor
from dl_compeng_pg.compeng_pg_base.data_processor_service_pg import CompEngPgService
from dl_compeng_pg.compeng_pg_base.pool_base import BasePgPoolWrapper
from dl_core.data_processing.processing.processor import OperationProcessorAsyncBase


@attr.s
class AsyncpgCompEngService(CompEngPgService[AsyncpgPoolWrapper]):
    def _get_pool_wrapper_cls(self) -> Type[BasePgPoolWrapper]:
        return AsyncpgPoolWrapper

    def get_data_processor(self) -> OperationProcessorAsyncBase:
        return AsyncpgOperationProcessor(pg_pool=self.pool)

lib/dl_compeng_pg/dl_compeng_pg/compeng_asyncpg/data_processor_service_asyncpg.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import contextmanager
import logging
from typing import (
    TYPE_CHECKING,
    AsyncGenerator,
    ClassVar,
    Generator,
    Optional,
    Sequence,
    Union,
)

import asyncpg
import attr
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import pypostgresql

from dl_compeng_pg.compeng_pg_base.exec_adapter_base import PostgreSQLExecAdapterAsync
from dl_constants.enums import UserDataType
from dl_core.connectors.base.error_transformer import DbErrorTransformer
from dl_core.data_processing.prepared_components.primitives import PreparedMultiFromInfo
from dl_core.data_processing.streaming import (
    AsyncChunked,
    AsyncChunkedBase,
)
from dl_sqlalchemy_postgres.asyncpg import DBAPIMock

from dl_connector_postgresql.core.postgresql_base.error_transformer import make_async_pg_error_transformer
from dl_connector_postgresql.core.postgresql_base.utils import compile_pg_query


if TYPE_CHECKING:
    from dl_constants.types import TBIDataValue


LOGGER = logging.getLogger(__name__)


@attr.s
class AsyncpgExecAdapter(PostgreSQLExecAdapterAsync[asyncpg.pool.PoolConnectionProxy]):  # noqa
    # `<PoolConnectionProxy <asyncpg.connection.Connection object at ...> ...>`
    _conn: asyncpg.pool.PoolConnectionProxy = attr.ib()
    _error_transformer: DbErrorTransformer = attr.ib(init=False, factory=make_async_pg_error_transformer)

    _log: ClassVar[logging.Logger] = LOGGER

    @property
    def dialect(self) -> sa.engine.default.DefaultDialect:
        # we should replace it with sqlalchemy.dialects.postgresql.asyncpg
        # now it produces wrong types
        dialect = pypostgresql.dialect(paramstyle="pyformat", dbapi=DBAPIMock())

        dialect.implicit_returning = True
        dialect.supports_native_enum = True
        dialect.supports_smallserial = True  # 9.2+
        dialect._backslash_escapes = False
        dialect.supports_sane_multi_rowcount = True  # psycopg 2.0.9+
        dialect._has_native_hstore = True  # type: ignore  # TODO: fix

        return dialect

    def _compile_query(self, query):  # type: ignore  # TODO: fix
        # exclude {AsyncAdapt_asyncpg_dbapi.ENUM}
        # https://github.com/sqlalchemy/sqlalchemy/blob/rel_1_4/lib/sqlalchemy/dialects/postgresql/asyncpg.py#L345
        return compile_pg_query(query, self.dialect, exclude_types={DBAPIMock.ENUM}, add_types=False)

    async def _execute(self, query: Union[str, sa.sql.base.Executable]) -> None:
        """Execute query without fetching data"""
        query_text, params = self._compile_query(query)
        await self._conn.execute(query_text, *params)

    async def _execute_ddl(self, query: Union[str, sa.sql.base.Executable]) -> None:
        """Execute a DDL statement"""
        await self._execute(query)

    @contextmanager
    def handle_db_errors(self, query_text: str) -> Generator[None, None, None]:
        try:
            yield
        except Exception as wrapper_exc:
            trans_exc = self._error_transformer.make_bi_error(wrapper_exc=wrapper_exc, debug_compiled_query=query_text)
            raise trans_exc from wrapper_exc

    async def _execute_and_fetch(  # type: ignore  # TODO: fix
        self,
        *,
        query: Union[str, sa.sql.selectable.Select],
        user_types: Sequence[UserDataType],
        chunk_size: int,
        joint_dsrc_info: Optional[PreparedMultiFromInfo] = None,
        query_id: str,
    ) -> AsyncChunked[list[TBIDataValue]]:
        query_text, params = self._compile_query(query)

        async def chunked_data_gen() -> AsyncGenerator[list[list[TBIDataValue]], None]:
            """Fetch data in chunks"""

            with self.handle_db_errors(query_text=query_text):
                cur = await self._conn.cursor(query_text, *params)
                while True:
                    chunk = []
                    for row in await cur.fetch(chunk_size):
                        assert len(row) == len(user_types), "user_types is not the same length as the data row"
                        chunk.append(
                            [
                                self._tt.cast_for_output(value=value, user_t=user_t)
                                for value, user_t in zip(row, user_types)
                            ]
                        )
                    if not chunk:
                        break
                    yield chunk

        return AsyncChunked(chunked_data=chunked_data_gen())

    async def insert_data_into_table(
        self,
        *,
        table_name: str,
        names: Sequence[str],
        user_types: Sequence[UserDataType],
        data: AsyncChunkedBase,
    ) -> None:
        """Insert data into a table."""
        self._log.info(f"Inserting data into table {table_name}")

        async for raw_chunk in data.chunks:
            chunk = []
            for row in raw_chunk:
                chunk.append(
                    [self._tt.cast_for_input(value=value, user_t=user_t) for value, user_t in zip(row, user_types)]
                )
            await self._conn.copy_records_to_table(table_name=table_name, columns=names, records=chunk)

lib/dl_compeng_pg/dl_compeng_pg/compeng_asyncpg/exec_adapter_asyncpg.py
datalens-tech/datalens-backend
from dl_compeng_pg.compeng_asyncpg.data_processor_service_asyncpg import AsyncpgCompEngService
from dl_constants.enums import ProcessorType
from dl_core.data_processors.base.plugin import DataProcessorPlugin


class AsyncpgCompengPlugin(DataProcessorPlugin):
    data_processor_type = ProcessorType.ASYNCPG
    data_processor_service_cls = AsyncpgCompEngService

lib/dl_compeng_pg/dl_compeng_pg/compeng_asyncpg/plugin.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
from typing import (
    Optional,
    Type,
)

import asyncpg
import attr

from dl_compeng_pg.compeng_pg_base.pool_base import (
    DEFAULT_OPERATION_TIMEOUT,
    DEFAULT_POOL_MAX_SIZE,
    DEFAULT_POOL_MIN_SIZE,
    BasePgPoolWrapper,
)


@attr.s
class AsyncpgPoolWrapper(BasePgPoolWrapper):
    _pool: Optional[asyncpg.pool.Pool] = attr.ib()

    @classmethod
    async def connect(
        cls: Type["AsyncpgPoolWrapper"],
        url: str,
        pool_min_size: int = DEFAULT_POOL_MIN_SIZE,  # Initial pool size
        pool_max_size: int = DEFAULT_POOL_MAX_SIZE,  # Maximum pool size
        operation_timeout: float = DEFAULT_OPERATION_TIMEOUT,  # SQL operation timeout
    ) -> "AsyncpgPoolWrapper":
        pool = await asyncpg.create_pool(
            url,
            min_size=pool_min_size,
            max_size=pool_max_size,
            command_timeout=operation_timeout,
            # Should disable the prepared statements and reduce the problems
            # with constants.
            statement_cache_size=0,
        )
        return cls(pool=pool)

    @property
    def pool(self) -> asyncpg.pool.Pool:
        if self._pool is None:
            raise RuntimeError("Asyncpg pool is closed")
        return self._pool

    async def disconnect(self) -> None:
        pool = self._pool
        self._pool = None
        await asyncio.wait_for(pool.close(), timeout=2.5)  # type: ignore  # TODO: fix

lib/dl_compeng_pg/dl_compeng_pg/compeng_asyncpg/pool_asyncpg.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import AsyncExitStack
from typing import Optional

import asyncpg
import attr

from dl_compeng_pg.compeng_asyncpg.exec_adapter_asyncpg import AsyncpgExecAdapter
from dl_compeng_pg.compeng_asyncpg.pool_asyncpg import AsyncpgPoolWrapper
from dl_compeng_pg.compeng_pg_base.processor_base import PostgreSQLOperationProcessor


@attr.s
class AsyncpgOperationProcessor(
    PostgreSQLOperationProcessor[AsyncpgExecAdapter, AsyncpgPoolWrapper, asyncpg.pool.PoolConnectionProxy]
):
    _cmstack: Optional[AsyncExitStack] = attr.ib(init=False, default=None)
    _timeout = 1.5

    async def start(self) -> None:
        assert self._cmstack is None, "should not double-enter"
        cmstack = AsyncExitStack()
        self._cmstack = cmstack
        pg_conn = await cmstack.enter_async_context(self._pg_pool.pool.acquire(timeout=self._timeout))
        self._pg_conn = pg_conn
        await cmstack.enter_async_context(pg_conn.transaction())
        self._pgex_adapter = AsyncpgExecAdapter(conn=pg_conn)

    async def end(self) -> None:
        self._pgex_adapter = None
        await self._cmstack.aclose()  # type: ignore  # TODO: fix
        self._pg_conn = None
        self._cmstack = None

lib/dl_compeng_pg/dl_compeng_pg/compeng_asyncpg/processor_asyncpg.py
datalens-tech/datalens-backend
# TODO: Move to data_processing
from __future__ import annotations

import abc
import logging
from typing import (
    ClassVar,
    Generic,
    Optional,
    Type,
    TypeVar,
)

import attr

from dl_compeng_pg.compeng_pg_base.pool_base import (
    DEFAULT_OPERATION_TIMEOUT,
    DEFAULT_POOL_MAX_SIZE,
    DEFAULT_POOL_MIN_SIZE,
    BasePgPoolWrapper,
)
from dl_core.aio.web_app_services.data_processing.data_processor import (
    DataProcessorConfig,
    DataProcessorService,
)


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class CompEngPgConfig(DataProcessorConfig):
    url: str = attr.ib(kw_only=True)


_POOL_TV = TypeVar("_POOL_TV", bound=BasePgPoolWrapper)
_COMPENG_PR_SRV_TV = TypeVar("_COMPENG_PR_SRV_TV", bound="CompEngPgService")


@attr.s
class CompEngPgService(DataProcessorService, Generic[_POOL_TV], metaclass=abc.ABCMeta):
    APP_KEY: ClassVar[str] = "compeng_service"

    _pg_url: str = attr.ib()  # DSN
    _pool_min_size: int = attr.ib(default=DEFAULT_POOL_MIN_SIZE)
    _pool_max_size: int = attr.ib(default=DEFAULT_POOL_MAX_SIZE)
    _operation_timeout: float = attr.ib(default=DEFAULT_OPERATION_TIMEOUT)
    _pool: Optional[_POOL_TV] = attr.ib(init=False, default=None)

    @property
    def pool(self) -> _POOL_TV:
        if self._pool is None:
            raise ValueError("Pool was not created")
        return self._pool

    async def initialize(self) -> None:
        await self._init_pool()

    @abc.abstractmethod
    def _get_pool_wrapper_cls(self) -> Type[BasePgPoolWrapper]:
        raise NotImplementedError

    async def _init_pool(self) -> None:
        self._pool = await self._get_pool_wrapper_cls().connect(  # type: ignore  # TODO: fix
            self._pg_url,
            pool_min_size=self._pool_min_size,
            pool_max_size=self._pool_max_size,
            operation_timeout=self._operation_timeout,
        )

    async def finalize(self) -> None:
        if self._pool is not None:
            await self._deinit_pool()

    async def _deinit_pool(self) -> None:
        """..."""
        LOGGER.info("Tear down compeng pg pool %r...", self)
        await self._pool.disconnect()  # type: ignore  # TODO: fix
        self._pool = None
        LOGGER.info("Tear down compeng pg pool %r: done.", self)

    @classmethod
    def from_config(cls: Type[_COMPENG_PR_SRV_TV], config: DataProcessorConfig) -> _COMPENG_PR_SRV_TV:
        assert isinstance(config, CompEngPgConfig)
        return cls(
            pg_url=config.url,
        )

lib/dl_compeng_pg/dl_compeng_pg/compeng_pg_base/data_processor_service_pg.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    ClassVar,
    Generic,
    Sequence,
    TypeVar,
    Union,
)

import attr
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
from sqlalchemy.engine.default import DefaultDialect
from sqlalchemy.sql.base import Executable

from dl_constants.enums import UserDataType
from dl_core.data_processing.processing.db_base.exec_adapter_base import ProcessorDbExecAdapterBase
from dl_core.data_processing.streaming import AsyncChunkedBase
from dl_core.db.sa_types import make_sa_type

from dl_connector_postgresql.core.postgresql_base.type_transformer import PostgreSQLTypeTransformer


LOGGER = logging.getLogger(__name__)

_CONN_TV = TypeVar("_CONN_TV")


@attr.s
class PostgreSQLExecAdapterAsync(Generic[_CONN_TV], ProcessorDbExecAdapterBase, metaclass=abc.ABCMeta):  # noqa
    """
    PG-CompEng-specific adapter.
    Adds DDL functionality and PostgreSQL specificity to the base DB adapter.
    """

    _conn: _CONN_TV = attr.ib(kw_only=True)
    _tt: PostgreSQLTypeTransformer = attr.ib(factory=PostgreSQLTypeTransformer, init=False)

    _log: ClassVar[logging.Logger] = LOGGER.getChild("PostgreSQLExecAdapterAsync")  # type: ignore  # TODO: fix

    @property
    def dialect(self) -> DefaultDialect:
        # Note: not necessarily psycopg2, but should be close enough
        # (especially for debug-compile).
        return PGDialect_psycopg2()

    @abc.abstractmethod
    async def _execute(self, query: Union[str, Executable]) -> None:
        """Execute query without (necessarily) fetching data"""

    async def _execute_ddl(self, query: Union[str, Executable]) -> None:
        """Execute a DDL statement"""
        await self._execute(query)

    def _make_sa_table(self, table_name: str, names: Sequence[str], user_types: Sequence[UserDataType]) -> sa.Table:
        assert len(names) == len(user_types)
        columns = [
            sa.Column(name=name, type_=make_sa_type(native_type=self._tt.type_user_to_native(user_t=user_t)))
            for name, user_t in zip(names, user_types)
        ]
        return sa.Table(table_name, sa.MetaData(), *columns, prefixes=["TEMPORARY"])

    async def create_table(
        self,
        *,
        table_name: str,
        names: Sequence[str],
        user_types: Sequence[UserDataType],
    ) -> sa.sql.selectable.TableClause:
        """Create table in database"""

        table = self._make_sa_table(table_name=table_name, names=names, user_types=user_types)
        self._log.info(f"Creating PG processor table {table_name}: {table}")
        await self._execute_ddl(sa.schema.CreateTable(table))
        return table

    async def _drop_table(self, table_name: str) -> None:
        await self._execute_ddl(sa.schema.DropTable(sa.table(table_name)))  # type: ignore

    async def drop_table(self, table_name: str) -> None:
        """Drop table in database"""

        self._log.info(f"Dropping PG processor table {table_name}")
        await self._drop_table(table_name=table_name)

    @abc.abstractmethod
    async def insert_data_into_table(
        self,
        *,
        table_name: str,
        names: Sequence[str],
        user_types: Sequence[UserDataType],
        data: AsyncChunkedBase,
    ) -> None:
        """,,,"""

lib/dl_compeng_pg/dl_compeng_pg/compeng_pg_base/exec_adapter_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr
import shortuuid
import sqlalchemy as sa

from dl_compeng_pg.compeng_pg_base.exec_adapter_base import PostgreSQLExecAdapterAsync
from dl_constants.enums import JoinType
from dl_core.connectors.base.query_compiler import QueryCompiler
from dl_core.data_processing.prepared_components.primitives import PreparedSingleFromInfo
from dl_core.data_processing.processing.db_base.op_executors import (
    OpExecutorAsync,
    log_op,
)
from dl_core.data_processing.processing.operation import (
    BaseOp,
    UploadOp,
)
from dl_core.data_processing.stream_base import (
    DataSourceVS,
    DataStreamAsync,
)


COMPENG_SUPPORTED_JOIN_TYPES = frozenset(
    {
        JoinType.inner,
        JoinType.left,
        JoinType.right,
        JoinType.full,
    }
)


@attr.s
class PgOpExecutorAsync(OpExecutorAsync):
    @property
    def pgex_adapter(self) -> PostgreSQLExecAdapterAsync:
        assert isinstance(self.db_ex_adapter, PostgreSQLExecAdapterAsync)
        return self.db_ex_adapter


class UploadOpExecutorAsync(PgOpExecutorAsync):
    """Dumps incoming stream to a database table"""

    @log_op  # type: ignore  # TODO: fix
    async def execute(self, op: BaseOp) -> DataSourceVS:  # type: ignore  # TODO: fix
        assert isinstance(op, UploadOp)

        source_stream = self.ctx.get_stream(op.source_stream_id)
        assert isinstance(source_stream, DataStreamAsync)

        table_name = shortuuid.uuid()
        await self.pgex_adapter.create_table(
            table_name=table_name,
            names=source_stream.names,
            user_types=source_stream.user_types,
        )
        await self.pgex_adapter.insert_data_into_table(
            table_name=table_name,
            names=source_stream.names,
            user_types=source_stream.user_types,
            data=source_stream.data,
        )

        alias = op.alias
        sql_source = sa.alias(sa.table(table_name), alias)

        prep_src_info = PreparedSingleFromInfo(
            id=op.result_id,
            alias=alias,
            query_compiler=QueryCompiler(dialect=self.pgex_adapter.dialect),
            sql_source=sql_source,
            col_names=source_stream.names,
            user_types=source_stream.user_types,
            data_source_list=None,
            supported_join_types=COMPENG_SUPPORTED_JOIN_TYPES,
            db_name=None,
            connect_args={},
            pass_db_query_to_user=False,
            target_connection_ref=None,
        )

        return DataSourceVS(
            id=op.dest_stream_id,
            result_id=prep_src_info.id,
            prep_src_info=prep_src_info,
            names=prep_src_info.col_names,
            user_types=prep_src_info.user_types,
            alias=op.alias,
        )

lib/dl_compeng_pg/dl_compeng_pg/compeng_pg_base/op_executors.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from contextlib import asynccontextmanager
from typing import (
    AsyncGenerator,
    Type,
    TypeVar,
)


_POOL_WRAPPER_TV = TypeVar("_POOL_WRAPPER_TV", bound="BasePgPoolWrapper")

DEFAULT_POOL_MIN_SIZE = 10
DEFAULT_POOL_MAX_SIZE = 20
DEFAULT_OPERATION_TIMEOUT = 60.0


class BasePgPoolWrapper(metaclass=abc.ABCMeta):
    @classmethod
    async def connect(
        cls: Type[_POOL_WRAPPER_TV],
        url: str,
        pool_min_size: int = DEFAULT_POOL_MIN_SIZE,  # Initial pool size
        pool_max_size: int = DEFAULT_POOL_MAX_SIZE,  # Maximum pool size
        operation_timeout: float = DEFAULT_OPERATION_TIMEOUT,  # SQL operation timeout
    ) -> _POOL_WRAPPER_TV:
        raise NotImplementedError

    async def disconnect(self) -> None:
        raise NotImplementedError

    @classmethod
    @asynccontextmanager
    async def context(
        cls: Type[_POOL_WRAPPER_TV],
        url: str,
        pool_min_size: int = DEFAULT_POOL_MIN_SIZE,
        pool_max_size: int = DEFAULT_POOL_MAX_SIZE,
        operation_timeout: float = DEFAULT_OPERATION_TIMEOUT,
    ) -> AsyncGenerator[_POOL_WRAPPER_TV, None]:
        pool_wrapper = await cls.connect(
            url=url, pool_min_size=pool_min_size, pool_max_size=pool_max_size, operation_timeout=operation_timeout
        )
        try:
            yield pool_wrapper
        finally:
            await pool_wrapper.disconnect()

lib/dl_compeng_pg/dl_compeng_pg/compeng_pg_base/pool_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    ClassVar,
    Dict,
    Generic,
    Optional,
    Type,
    TypeVar,
)

import attr

from dl_compeng_pg.compeng_pg_base.exec_adapter_base import PostgreSQLExecAdapterAsync
from dl_compeng_pg.compeng_pg_base.op_executors import UploadOpExecutorAsync
from dl_compeng_pg.compeng_pg_base.pool_base import BasePgPoolWrapper
from dl_constants.enums import UserDataType
from dl_core.data_processing.processing.context import OpExecutionContext
from dl_core.data_processing.processing.db_base.op_executors import (
    CalcOpExecutorAsync,
    DownloadOpExecutorAsync,
    JoinOpExecutorAsync,
    OpExecutorAsync,
)
from dl_core.data_processing.processing.operation import (
    BaseOp,
    CalcOp,
    DownloadOp,
    JoinOp,
    UploadOp,
)
from dl_core.data_processing.processing.processor import OperationProcessorAsyncBase
from dl_core.data_processing.stream_base import AbstractStream


_ADAPTER_TV = TypeVar("_ADAPTER_TV", bound=PostgreSQLExecAdapterAsync)
_POOL_TV = TypeVar("_POOL_TV", bound=BasePgPoolWrapper)
_CONN_TV = TypeVar("_CONN_TV")


@attr.s
class PostgreSQLOperationProcessor(
    OperationProcessorAsyncBase, Generic[_ADAPTER_TV, _POOL_TV, _CONN_TV], metaclass=abc.ABCMeta
):
    _pg_pool: _POOL_TV = attr.ib()
    _task_timeout: Optional[int] = attr.ib(default=None)
    _pgex_adapter: Optional[_ADAPTER_TV] = attr.ib(init=False, default=None)
    _pg_conn: Optional[_CONN_TV] = attr.ib(init=False, default=None)

    @abc.abstractmethod
    async def start(self) -> None:
        """Prepare for work."""

    @abc.abstractmethod
    async def end(self) -> None:
        """Cleanup."""

    async def ping(self) -> Optional[int]:
        assert self._pgex_adapter is not None
        result = await self._pgex_adapter.scalar("select 1", user_type=UserDataType.integer)
        assert result is None or isinstance(result, int)
        return result

    _executors: ClassVar[Dict[Type[BaseOp], Type[OpExecutorAsync]]] = {
        UploadOp: UploadOpExecutorAsync,
        DownloadOp: DownloadOpExecutorAsync,
        CalcOp: CalcOpExecutorAsync,
        JoinOp: JoinOpExecutorAsync,
    }

    async def execute_operation(self, op: BaseOp, ctx: OpExecutionContext) -> AbstractStream:
        opex_cls: Type[OpExecutorAsync] = self._executors[type(op)]
        assert self._pgex_adapter is not None
        opex = opex_cls(db_ex_adapter=self._pgex_adapter, ctx=ctx)  # type: ignore  # TODO: fix
        return await opex.execute(op)

lib/dl_compeng_pg/dl_compeng_pg/compeng_pg_base/processor_base.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_configs.settings_loaders.settings_obj_base import SettingsBase
from dl_constants.enums import ConnectorAvailability


@attr.s(kw_only=True)
class TranslatableSettings(SettingsBase):
    text: str = attr.ib()
    domain: Optional[str] = attr.ib(default=None)


@attr.s(kw_only=True)
class ConnectorBaseSettings(SettingsBase):
    pass


@attr.s(kw_only=True)
class ConnectorSettings(ConnectorBaseSettings):
    conn_type: str = attr.ib()
    availability: ConnectorAvailability = attr.ib(default=ConnectorAvailability.free)


@attr.s(kw_only=True)
class ConnectorContainerSettings(ConnectorBaseSettings):
    title_translatable: TranslatableSettings = attr.ib()
    alias: str = attr.ib()
    includes: list[ConnectorSettings] = attr.ib()


@attr.s(kw_only=True)
class SectionSettings(SettingsBase):
    title_translatable: TranslatableSettings = attr.ib()
    connectors: list[ConnectorBaseSettings] = attr.ib(validator=attr.validators.min_len(1))  # type: ignore


@attr.s(kw_only=True)
class ConnectorAvailabilityConfigSettings(SettingsBase):
    uncategorized: list[ConnectorSettings] = attr.ib(factory=list)
    sections: list[SectionSettings] = attr.ib(factory=list)

    visible_connectors: set[str] = attr.ib(factory=set)

lib/dl_configs/dl_configs/connector_availability.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc


class ConnectorsDataBase(metaclass=abc.ABCMeta):
    @classmethod
    @abc.abstractmethod
    def connector_name(cls) -> str:
        """
        TODO: BI-4359 just remove it after migration to new scheme
        """
        raise NotImplementedError

lib/dl_configs/dl_configs/connectors_data.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Callable

import attr

from dl_configs.environments import LegacyDefaults
from dl_configs.settings_loaders.fallback_cfg_resolver import ObjectLikeConfig
from dl_configs.settings_loaders.settings_obj_base import SettingsBase


@attr.s(frozen=True)
class ConnectorSettingsBase(SettingsBase):
    """"""


ConnectorsConfigType = ObjectLikeConfig | LegacyDefaults
SettingsFallbackType = Callable[[ConnectorsConfigType], dict[str, ConnectorSettingsBase]]

lib/dl_configs/dl_configs/connectors_settings.py
datalens-tech/datalens-backend
from __future__ import annotations

from os import environ
from typing import (
    Any,
    Dict,
    Mapping,
    Optional,
)

import attr
from cryptography import fernet

from dl_configs.settings_loaders.loader_env import EnvSettingsLoader
from dl_configs.settings_loaders.meta_definition import s_attrib
from dl_configs.settings_loaders.settings_obj_base import SettingsBase


@attr.s(frozen=True)
class CryptoKeysConfig(SettingsBase):
    map_id_key: Dict[str, str] = s_attrib("KEY_VAL_ID", sensitive=True)
    actual_key_id: str = s_attrib("ACTUAL_KEY_ID")

    @actual_key_id.default
    def _default_actual_key_id(self):
        if len(self.map_id_key) == 1:
            return next(iter(self.map_id_key.keys()))
        # TODO FIX: Integrate with env loader exceptions handling
        raise ValueError("Missing actual_key_id (acceptable only in single key case)")

    def __attrs_post_init__(self):
        if self.actual_key_id not in self.map_id_key:
            raise ValueError()

    @classmethod
    def from_json(cls, json_value: Dict[str, Any]) -> CryptoKeysConfig:
        try:
            actual_key_id = json_value["actual_key_id"]
        except KeyError as exc:
            # TODO FIX: Integrate with env loader exceptions handling
            raise ValueError(f"Crypto keys config missing key {exc!r}")
        try:
            map_id_key = json_value["keys"]
        except KeyError as exc:
            # TODO FIX: Integrate with env loader exceptions handling
            raise ValueError(f"Crypto keys config missing key {exc!r}")
        return CryptoKeysConfig(
            actual_key_id=actual_key_id,
            map_id_key=map_id_key,
        )


def get_single_key_crypto_keys_config(*, key_id: str, key_value: str) -> CryptoKeysConfig:
    return CryptoKeysConfig(
        map_id_key={key_id: key_value},
        actual_key_id=key_id,
    )


def get_dummy_crypto_keys_config() -> CryptoKeysConfig:
    return CryptoKeysConfig(
        map_id_key={"dummy_key_id_qwerty123": fernet.Fernet.generate_key().decode()},
        actual_key_id="dummy_key_id_qwerty123",
    )


@attr.s()
class _CryptoKeysConfigOnlySettings(SettingsBase):
    """
    We can not load CryptoKeysConfig if it encoded in env as single JSON string.
    Loader can not handle this case due to root settings class should be env dict.
    So we make this wrapper.
    """

    cry: CryptoKeysConfig = s_attrib("DL_CRY", json_converter=CryptoKeysConfig.from_json, sensitive=True)


def get_crypto_keys_config_from_env(config_source: Optional[Mapping[str, str]] = None) -> CryptoKeysConfig:
    if config_source is None:
        config_source = environ

    loader = EnvSettingsLoader(config_source)
    return loader.load_settings(
        settings_type=_CryptoKeysConfigOnlySettings,
    ).cry

lib/dl_configs/dl_configs/crypto_keys.py
datalens-tech/datalens-backend
import enum


class FeatureEnablingMode(enum.Enum):
    enabled = "enabled"
    bleeding_edge = "bleeding_edge"
    disabled = "disabled"


class RequiredService(enum.Enum):
    POSTGRES = enum.auto()
    RQE_EXT_ASYNC = enum.auto()
    RQE_EXT_SYNC = enum.auto()
    RQE_INT_SYNC = enum.auto()


RQE_SERVICES: set[RequiredService] = {
    RequiredService.RQE_INT_SYNC,
    RequiredService.RQE_EXT_SYNC,
    RequiredService.RQE_EXT_ASYNC,
}


class RedisMode(enum.Enum):
    sentinel = enum.auto()
    single_host = enum.auto()

lib/dl_configs/dl_configs/enums.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_configs.env_var_reader import get_from_env


def use_jaeger_tracer() -> bool:
    return get_from_env("DL_USE_JAEGER_TRACER", lambda s: bool(int(s)), default=False)


def jaeger_service_name_env_aware(service_name: str) -> str:
    return service_name + get_from_env("DL_JAEGER_TRACER_SERVICE_NAME_SUFFIX", str, "")

lib/dl_configs/dl_configs/env_var_definitions.py
datalens-tech/datalens-backend
from __future__ import annotations

import os
from typing import (
    Callable,
    Generic,
    Literal,
    NamedTuple,
    Optional,
    TypeVar,
    Union,
    overload,
)


_RET_TV = TypeVar("_RET_TV")
_FACTORY_RET_TV = TypeVar("_FACTORY_RET_TV")


class SecretStr(str):
    def __repr__(self) -> str:
        return "sec_str<***>"

    def __str__(self) -> str:
        return "sec_str<***>"


class Required(NamedTuple):
    message: Optional[str] = None


class Factory(Generic[_FACTORY_RET_TV]):
    def __init__(self, factory: Callable[[], _FACTORY_RET_TV]):
        self._factory = factory

    def build(self) -> _FACTORY_RET_TV:
        return self._factory()


class EnvError(Exception):
    pass


@overload
def get_from_env(env_key: str, converter: Callable[[str], _RET_TV], default: Literal[None]) -> Optional[_RET_TV]:
    pass


@overload  # noqa
def get_from_env(env_key: str, converter: Callable[[str], _RET_TV], default: Required) -> _RET_TV:
    pass


@overload  # noqa
def get_from_env(
    env_key: str, converter: Callable[[str], _RET_TV], default: Union[_RET_TV, Factory[_RET_TV]]
) -> _RET_TV:
    pass


@overload  # noqa
def get_from_env(env_key: str, converter: Callable[[str], _RET_TV]) -> _RET_TV:
    pass


def get_from_env(  # noqa
    env_key: str,
    converter: Callable[[str], _RET_TV],
    default: Union[_RET_TV, Factory[_RET_TV], Required, None] = Required(),
) -> Optional[_RET_TV]:
    if env_key in os.environ:
        try:
            return converter(os.environ[env_key])
        except Exception as err:
            raise EnvError(f"Could not convert env var {env_key}") from err

    if isinstance(default, Required):
        err_msg = f"Env missing '{env_key}' but required"
        if default.message is not None:
            err_msg = f"{err_msg}: {default.message}"
        raise EnvError(err_msg, env_key)

    if isinstance(default, Factory):
        return default.build()

    return default

lib/dl_configs/dl_configs/env_var_reader.py
datalens-tech/datalens-backend
"""
Information about the outside context for all environments.

All the classes (instead of dicts) for static checking.
"""

from __future__ import annotations


def is_setting_applicable(cfg: object, key: str):
    # TODO: remove it
    # it's a crutch-function
    # we will remove this after migrating values to dedicated keys
    return bool(getattr(cfg, key, False))


class LegacyDefaults:
    pass


class BaseInstallationsMap:
    pass


class LegacyEnvAliasesMap:
    pass

lib/dl_configs/dl_configs/environments.py
datalens-tech/datalens-backend
from __future__ import annotations

import os
from typing import Optional

import attr

from dl_configs.settings_loaders.common import SDict
from dl_configs.settings_loaders.loader_env import EnvSettingsLoader
from dl_configs.settings_loaders.meta_definition import (
    required,
    s_attrib,
)
from dl_configs.settings_loaders.settings_obj_base import SettingsBase
from dl_configs.utils import validate_one_of


@attr.s(frozen=True)
class RQEBaseURL(SettingsBase):
    scheme: str = s_attrib("SCHEME", env_var_converter=validate_one_of({"http", "https"}), missing="http")
    host: str = s_attrib("HOST")
    port: int = s_attrib("PORT")

    def __str__(self) -> str:
        return f"{self.scheme}://{self.host}:{self.port}"


@attr.s(frozen=True)
class RQEConfig(SettingsBase):
    ext_sync_rqe: RQEBaseURL = s_attrib("EXT_SYNC")
    ext_async_rqe: RQEBaseURL = s_attrib("EXT_ASYNC")
    int_sync_rqe: RQEBaseURL = s_attrib("INT_SYNC")
    int_async_rqe: RQEBaseURL = s_attrib("INT_ASYNC")
    hmac_key: bytes = s_attrib(
        "SECRET_KEY",
        sensitive=True,
        env_var_converter=lambda s: s.encode("ascii"),
    )

    @classmethod
    def get_default(cls) -> RQEConfig:
        return RQEConfig(
            hmac_key=required(bytes),
            int_sync_rqe=RQEBaseURL(scheme="http", host="[::1]", port=9874),
            int_async_rqe=RQEBaseURL(scheme="http", host="[::1]", port=9875),
            ext_sync_rqe=RQEBaseURL(scheme="http", host="[::1]", port=9876),
            ext_async_rqe=RQEBaseURL(scheme="http", host="[::1]", port=9877),
        )

    def clone(self, **kwargs):
        return attr.evolve(self, **kwargs)


def rqe_config_from_env(env: Optional[SDict] = None, add_dummy_hmac_if_missing: bool = False) -> RQEConfig:
    """
    :return: RQE config (default with overrides from environment variables)
    """
    hmac_env_key = "RQE_SECRET_KEY"
    actual_env: dict[str, str]

    if env is None:
        actual_env = dict(os.environ)
        # TODO FIX: BI-2497 remove this kludge when configs loading in tests import time will be eliminated (mostly uWSGI apps instantiation)
        if add_dummy_hmac_if_missing and hmac_env_key not in actual_env:
            actual_env[hmac_env_key] = ""
    else:
        actual_env = env

    loader = EnvSettingsLoader(actual_env)
    return loader.load_settings(
        key_prefix="RQE",
        settings_type=RQEConfig,
        default_value=RQEConfig.get_default(),
    )

lib/dl_configs/dl_configs/rqe.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Callable,
    Mapping,
    Union,
)


SDict = Mapping[str, str]
FallbackFactory = Union[
    Callable[[Any, Any], Any],  # fallback_cfg & app_cfg_type
    Callable[[Any], Any],  # fallback_cfg only
]

lib/dl_configs/dl_configs/settings_loaders/common.py
datalens-tech/datalens-backend
from typing import (
    Collection,
    Optional,
    Type,
)

import attr
from attr import make_class

from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.settings_loaders.meta_definition import s_attrib
from dl_configs.settings_loaders.settings_obj_base import SettingsBase
from dl_constants.enums import ConnectionType


def generate_connectors_settings_class(
    settings_registry: dict[ConnectionType, Type[ConnectorSettingsBase]],
    whitelist: Optional[Collection[ConnectionType]] = None,
) -> Type[attr.AttrsInstance]:
    attrs: dict[str, attr.Attribute] = {}
    for conn_type, settings_class in settings_registry.items():
        if whitelist is not None and conn_type not in whitelist:
            continue
        field_name = conn_type.value.upper()
        attrs[field_name] = s_attrib(field_name, type_class=Optional[settings_class], missing=None)

    return make_class("ConnectorSettingsByTypeGenerated", attrs=attrs, bases=(SettingsBase,), frozen=True)

lib/dl_configs/dl_configs/settings_loaders/connectors_settings.py
datalens-tech/datalens-backend
import logging
from typing import List

from dl_configs.settings_loaders.common import SDict


LOGGER = logging.getLogger(__name__)

REMAP = {
    "EXT_QUERY_EXECUTER_SECRET_KEY": "RQE_SECRET_KEY",
    # TODO FIX: BI-1753 remove after sensitive info will leave versioned data of US entry
    "FERNET_KEY": "DL_CRY_KEY_VAL_ID_0",
}

_LEGACY_AND_ACTUAL_KEYS_INTERSECTION = set(REMAP.keys()) & set(REMAP.values())

assert not _LEGACY_AND_ACTUAL_KEYS_INTERSECTION, (
    f"Got non-empty intersection of legacy & actual keys:" f" {_LEGACY_AND_ACTUAL_KEYS_INTERSECTION}"
)


def remap_env(s_dict: SDict) -> SDict:
    ret = dict(s_dict)

    legacy_keys_to_remap_candidates = s_dict.keys() & REMAP.keys()
    legacy_keys_to_remap: List[str] = []

    for legacy_key in sorted(legacy_keys_to_remap_candidates):
        actual_key = REMAP[legacy_key]

        if actual_key in s_dict:
            LOGGER.warning(
                "Environment contains both legacy (%r) and actual (%r) keys. Actual key taking precedence.",
                legacy_key,
                actual_key,
            )
            continue

        legacy_keys_to_remap.append(legacy_key)

    for legacy_key in legacy_keys_to_remap:
        actual_key = REMAP[legacy_key]
        LOGGER.info("Remapping environment key: %s -> %s", legacy_key, actual_key)
        ret[actual_key] = ret.pop(legacy_key)

    return ret

lib/dl_configs/dl_configs/settings_loaders/env_remap.py
datalens-tech/datalens-backend
class SettingsLoadingException(Exception):
    pass

lib/dl_configs/dl_configs/settings_loaders/exc.py
datalens-tech/datalens-backend
from collections.abc import Mapping
import logging
from typing import (
    Any,
    ClassVar,
    Iterator,
    Optional,
)

import attr
import yaml

from dl_configs.settings_loaders.common import SDict


LOGGER = logging.getLogger(__name__)


@attr.s
class FallbackConfigResolver:
    def resolve(self, s_dict: SDict) -> Any:
        raise NotImplementedError()


@attr.s
class ObjectLikeConfig(Mapping):
    _data: dict = attr.ib()
    _path: list = attr.ib(default=[])

    def __len__(self) -> int:
        return len(self._data)

    def __iter__(self) -> Iterator:
        return iter(self._data)

    def _get_key(self, key: Any):
        if key not in self._data:
            path = ".".join(self._path + [key])
            raise AttributeError(f'There is no record in config by path: "{path}"')
        return self._data[key]

    def __getattr__(self, key: Any) -> Any:
        return self._get_key(key)

    def get(self, key: Any):
        return self._data.get(key)

    def __getitem__(self, key: Any) -> Any:
        return self._get_key(key)

    @classmethod
    def from_dict(cls, data: dict, path: Optional[list] = None) -> "ObjectLikeConfig":
        if path is None:
            path = []

        def _get_value_for_cfg(k: Any, v: Any) -> Any:
            if isinstance(v, dict):
                return cls.from_dict(v, path + [k])
            if isinstance(v, (list, tuple)):
                return [
                    cls.from_dict(item, path + [k] + [str(idx)]) if isinstance(item, dict) else item
                    for idx, item in enumerate(v)
                ]
            return v

        ret = cls(
            data={k: _get_value_for_cfg(k, v) for k, v in data.items()},
            path=path,
        )
        return ret

    def to_dict(self) -> dict[str, Any]:
        def _get_value_for_dict(v) -> Any:
            if isinstance(v, ObjectLikeConfig):
                return v.to_dict()
            if isinstance(v, (list, tuple)):
                return [item.to_dict if isinstance(item, ObjectLikeConfig) else item for item in v]
            return v

        return {k: _get_value_for_dict(v) for k, v in self._data.items()}


@attr.s
class YamlFileConfigResolver(FallbackConfigResolver):
    config_path_key: ClassVar[str] = "CONFIG_PATH"

    @classmethod
    def _get_config_path(cls, s_dict: SDict) -> Optional[str]:
        path = s_dict.get(cls.config_path_key)
        LOGGER.info('Config path is "%s"', path)
        return path

    @classmethod
    def is_config_enabled(cls, s_dict: SDict):
        result = cls._get_config_path(s_dict) is not None
        LOGGER.info('Check path key in s_dict: "%s"', result)
        return result

    def resolve(self, s_dict: SDict) -> ObjectLikeConfig:
        LOGGER.info("Resolve by YamlFileConfigResolver")
        path = self._get_config_path(s_dict)
        with open(path) as f:
            config = yaml.safe_load(f.read())
        return ObjectLikeConfig.from_dict(config)


@attr.s
class ConstantFallbackConfigResolver(FallbackConfigResolver):
    config: Any = attr.ib()

    def resolve(self, s_dict: SDict) -> Any:
        return self.config

lib/dl_configs/dl_configs/settings_loaders/fallback_cfg_resolver.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import reduce
import inspect
import json
import logging
import os
from types import MappingProxyType
import typing
from typing import (
    Any,
    Callable,
    ClassVar,
    Collection,
    Optional,
    Type,
    TypeVar,
    Union,
    final,
    no_type_check,
)

import attr
import typeguard

from dl_configs.connectors_settings import (
    ConnectorsConfigType,
    ConnectorSettingsBase,
    SettingsFallbackType,
)
from dl_configs.settings_loaders.common import (
    FallbackFactory,
    SDict,
)
from dl_configs.settings_loaders.connectors_settings import generate_connectors_settings_class
from dl_configs.settings_loaders.env_remap import remap_env
from dl_configs.settings_loaders.exc import SettingsLoadingException
from dl_configs.settings_loaders.fallback_cfg_resolver import (
    FallbackConfigResolver,
    YamlFileConfigResolver,
)
from dl_configs.settings_loaders.meta_definition import (
    RequiredValue,
    ReservedKeys,
    SMeta,
    s_attrib,
)
from dl_configs.settings_loaders.settings_obj_base import SettingsBase
from dl_constants.enums import ConnectionType
from dl_utils.utils import get_type_full_name


_SETTINGS_TV = TypeVar("_SETTINGS_TV")

SEP = "_"


LOGGER = logging.getLogger(__name__)


class InvalidConfigValueException(SettingsLoadingException):
    pass


class ConfigFieldMissing(InvalidConfigValueException):
    def __init__(self, field_set: set[str]):
        super().__init__(field_set)
        self._field_set = field_set

    @property
    def field_set(self) -> set[str]:
        return self._field_set


def get_sub_keys(prefix: str, env: SDict):
    effective_prefix = prefix if prefix.endswith(SEP) else prefix + SEP
    return {key[len(effective_prefix) :]: value for key, value in env.items() if key.startswith(effective_prefix)}


NOT_SET = object()


@attr.s
class SDictExtractor:
    expected_type: Type = attr.ib()
    key: Optional[str] = attr.ib()
    # TODO FIX: Validate default on creation time
    default: Any = attr.ib()

    def _get_scoped_s_dict(self, s_dict: SDict) -> SDict:
        base_key = self.key
        return s_dict if base_key is None else get_sub_keys(base_key.upper(), s_dict)

    def has_field(self, s_dict: SDict) -> bool:
        raise NotImplementedError()

    def _extract(self, s_dict: SDict) -> Any:
        raise NotImplementedError()

    @final
    def extract(self, s_dict: SDict) -> Any:
        ret = self._extract(s_dict)
        typeguard.check_type(self.key or "root", ret, self.expected_type)
        return ret


@attr.s
class ScalarExtractor(SDictExtractor):
    converter: Optional[Callable[[str], Any]] = attr.ib()

    def has_field(self, s_dict: SDict) -> bool:
        return self.key in s_dict

    def _extract(self, s_dict: SDict) -> Any:
        if self.key not in s_dict:
            if self.default is not NOT_SET:
                return self.default

            return None

        value = s_dict[self.key]

        if self.converter is None:
            return value

        try:
            return self.converter(value)
        # TODO FIX: take in account potential sensivity of value
        except Exception as exc:
            raise InvalidConfigValueException(f"Could not convert value {self.key}: {self.expected_type}") from exc


@attr.s
class DefaultOnlyExtractor(SDictExtractor):
    def __attrs_post_init__(self):
        assert self.default is not NOT_SET

    def has_field(self, s_dict: SDict) -> bool:
        return False

    def _extract(self, s_dict: SDict) -> Any:
        return self.default


@attr.s
class DictExtractor(SDictExtractor):
    value_converter: Optional[Callable[[str], Any]] = attr.ib()

    def has_field(self, s_dict: SDict) -> bool:
        return len(self._get_scoped_s_dict(s_dict)) > 0

    def _extract(self, s_dict: SDict) -> Any:
        value_converter = self.value_converter
        return {
            key: value_converter(value) if value_converter is not None else value
            for key, value in self._get_scoped_s_dict(s_dict).items()
        }


@attr.s
class CompositeExtractor(SDictExtractor):
    map_field_name_subextractor: dict[str, SDictExtractor] = attr.ib()
    map_field_name_field_override: dict[str, Any] = attr.ib()
    composition_factory: Callable = attr.ib()
    json_value_converter: Optional[Callable[[Any], Any]] = attr.ib(default=None)
    is_root: bool = attr.ib(default=False)
    field_enables_flag_extractor: Optional[ScalarExtractor] = attr.ib(default=None)

    def should_ignore(self, s_dict) -> bool:
        scoped_s_dict = self._get_scoped_s_dict(s_dict)
        field_enables_flag_extractor = self.field_enables_flag_extractor

        if field_enables_flag_extractor is None:
            return False

        if self.field_enables_flag_extractor.has_field(scoped_s_dict):
            should_read = self.field_enables_flag_extractor.extract(scoped_s_dict)
            return not bool(should_read)

        return False

    def has_field(self, s_dict: SDict) -> bool:
        if len(self.map_field_name_field_override):
            return True

        scoped_s_dict = self._get_scoped_s_dict(s_dict)
        if ReservedKeys.JSON_VALUE in scoped_s_dict:
            return True
        all_extractors = list(self.map_field_name_subextractor.values())
        if self.field_enables_flag_extractor is not None:
            all_extractors.append(self.field_enables_flag_extractor)

        return any(extractor.has_field(scoped_s_dict) for extractor in all_extractors)

    def run_composition_factory_with_validation(self, fields: dict[str, Any]) -> Any:
        sig = inspect.signature(self.composition_factory)
        required_params = {
            p_name for p_name, p_obj in sig.parameters.items() if p_obj.default is inspect.Parameter.empty
        }
        missing_params = required_params - set(fields.keys())
        if missing_params:
            raise ConfigFieldMissing(missing_params)

        return self.composition_factory(**fields)

    @staticmethod
    def _ensure_no_missing_fields(
        obj: Any = None, map_field_name_missing_exc: Optional[dict[str, ConfigFieldMissing]] = None
    ) -> None:
        missing_required_fields: list[attr.Attribute] = []
        missing_sub_field_code_set: set[str] = set()

        def handle_field_sub_exc(faulty_field_name: str, exc: ConfigFieldMissing):
            missing_sub_field_code_set.update(f"{faulty_field_name}.{sub_field}" for sub_field in exc.field_set)

        if obj is None:
            if map_field_name_missing_exc is not None:
                for field_name, sub_exc in map_field_name_missing_exc.items():
                    handle_field_sub_exc(field_name, sub_exc)
        else:
            for field in attr.fields(type(obj)):
                field_value = getattr(obj, field.name)
                if isinstance(field_value, RequiredValue):
                    missing_required_fields.append(field)

                if map_field_name_missing_exc is not None and field.name in map_field_name_missing_exc:
                    handle_field_sub_exc(field.name, map_field_name_missing_exc[field.name])

        if missing_required_fields or missing_sub_field_code_set:
            raise ConfigFieldMissing(
                field_set=set(field.name for field in missing_required_fields) | missing_sub_field_code_set
            )

    def _extract_from_json(self, s_dict: SDict) -> Any:
        converter = self.json_value_converter
        if converter is None:
            raise InvalidConfigValueException(f"JSON converter is not defined for field {self.key!r}")

        json_str = s_dict[ReservedKeys.JSON_VALUE]

        illegal_fields = s_dict.keys() - {ReservedKeys.JSON_VALUE}

        # It is ok if field enabler in env
        if self.field_enables_flag_extractor is not None:
            illegal_fields = illegal_fields - {self.field_enables_flag_extractor.key}

        if illegal_fields:
            raise InvalidConfigValueException(
                f"{ReservedKeys.JSON_VALUE} can not be combined with any other keys: {' '.join(sorted(illegal_fields))}"
            )

        try:
            json_value = json.loads(json_str)
        except json.JSONDecodeError:
            # TODO FIX: take in account potential sensivity of value
            raise InvalidConfigValueException(f"Malformed JSON for {self.key!r}")

        try:
            return converter(json_value)
        except Exception as exc:  # noqa
            # TODO FIX: take in account potential sensivity of value
            raise InvalidConfigValueException(f"Can not convert JSON value to target object {self.key!r}")

    def _extract(self, s_dict: SDict) -> Any:
        default_value = self.default
        scoped_s_dict = self._get_scoped_s_dict(s_dict)

        if not self.is_root:
            if self.should_ignore(s_dict):
                return None

            if ReservedKeys.JSON_VALUE in scoped_s_dict:
                return self._extract_from_json(scoped_s_dict)

        fields = {}
        map_field_name_missing_fields: dict[str, ConfigFieldMissing] = {}

        for field_name, sub_extractor in self.map_field_name_subextractor.items():
            if field_name in self.map_field_name_field_override:
                fields[field_name] = self.map_field_name_field_override[field_name]
            elif sub_extractor.has_field(scoped_s_dict) or sub_extractor.default is not NOT_SET:
                try:
                    fields[field_name] = sub_extractor.extract(scoped_s_dict)
                except ConfigFieldMissing as child_field_missing:
                    map_field_name_missing_fields[field_name] = child_field_missing

        if default_value is None and not self.has_field(s_dict):
            if map_field_name_missing_fields:
                self._ensure_no_missing_fields(map_field_name_missing_exc=map_field_name_missing_fields)
            return None

        elif default_value is NOT_SET or default_value is None:
            if map_field_name_missing_fields:
                self._ensure_no_missing_fields(map_field_name_missing_exc=map_field_name_missing_fields)
            return self.run_composition_factory_with_validation(fields)

        else:
            updated_default_value = attr.evolve(self.default, **fields)
            self._ensure_no_missing_fields(
                obj=updated_default_value, map_field_name_missing_exc=map_field_name_missing_fields
            )
            return updated_default_value


@attr.s
class EnvSettingsLoader:
    file_remap_prefix: ClassVar[str] = "BIE_FILE_MAP"

    _s_dict: Optional[SDict] = attr.ib(repr=False)

    def __attrs_post_init__(self):
        if self._s_dict is None:
            self._s_dict = MappingProxyType(dict(os.environ))
        else:
            self._s_dict = MappingProxyType(dict(self._s_dict))

    @property
    def s_dict(self) -> SDict:
        assert self._s_dict is not None
        return self._s_dict

    def load_file_mapped_keys(self, s_dict: SDict) -> SDict:
        map_key_file_path = get_sub_keys(self.file_remap_prefix, s_dict)
        loaded_keys: dict[str, str] = {}

        for key, file_path in map_key_file_path.items():
            with open(file_path, "r", encoding="ascii") as file:
                loaded_keys[key] = file.read()

        return loaded_keys

    @classmethod
    def build_default_extractor(
        cls,
        name: str,
        field_type: Type,
        meta: SMeta,
        default: Any,
    ) -> Union[ScalarExtractor, DictExtractor]:
        env_var_name = meta.name.upper()
        unwrapped_type_set = cls.unwrap_union(field_type, ignore_none=True)
        assert len(unwrapped_type_set) == 1
        effective_type = next(iter(unwrapped_type_set))

        simple_type_converters: dict[Type, Callable[[str], Any]] = {
            int: int,
            float: float,
            str: None,
            bool: lambda x: {"1": True, "0": False, "true": True, "false": False}[x.lower()],
        }
        if meta.env_var_converter is not None:
            return ScalarExtractor(
                expected_type=field_type,
                key=env_var_name,
                converter=meta.env_var_converter,
                default=default,
            )
        elif effective_type in simple_type_converters:
            return ScalarExtractor(
                expected_type=field_type,
                key=env_var_name,
                converter=simple_type_converters[effective_type],
                default=default,
            )
        elif typing.get_origin(effective_type) == dict:
            key_type, value_type = typing.get_args(effective_type)
            assert key_type == str
            value_converter: Optional[Callable[[str], Any]]

            if value_type in simple_type_converters:
                value_converter = simple_type_converters[value_type]
            else:
                raise TypeError(f"Unsupprted map value type: {field_type!r}")

            return DictExtractor(
                expected_type=field_type,
                key=env_var_name,
                value_converter=value_converter,
                default=default,
            )
        else:
            raise TypeError(f"Unsupported field type: {field_type!r} while extracting {env_var_name}")

    @staticmethod
    def unwrap_union(the_type: Type, ignore_none: bool) -> frozenset[Type]:
        if typing.get_origin(the_type) == Union:
            unwrapped_types = typing.get_args(the_type)
            if ignore_none:
                # TODO FIX: try to not use noqa: E721
                # (to consider: `not issubclass(arg_type, type(None))`
                unwrapped_types = [argtype for argtype in unwrapped_types if argtype is not type(None)]  # noqa: E721
            return frozenset(unwrapped_types)
        return frozenset([the_type])

    @classmethod
    def is_sub_settings_field(cls, attrib: attr.Attribute) -> bool:
        the_type = attrib.type
        possible_types = cls.unwrap_union(the_type, ignore_none=True)
        map_type_is_sub_setting: dict[Type, bool] = {
            t: isinstance(t, type) and issubclass(t, SettingsBase) for t in possible_types
        }
        if any(map_type_is_sub_setting.values()):
            assert all(
                map_type_is_sub_setting.values()
            ), f"Divergence in {attrib!r} type classification: {map_type_is_sub_setting!r}"
            return True
        return False

    @classmethod
    def build_composite_extractor(
        cls,
        settings_type: Type[SettingsBase],
        key_prefix: tuple[str, ...],
        default: Any = NOT_SET,
        fallback_cfg: Any = None,
        app_cfg_type: Any = None,
        is_root: bool = False,
        field_enables_flag_extractor: Optional[ScalarExtractor] = None,
        json_converter: Optional[Callable[[Any], Any]] = None,
        field_overrides: Optional[dict[str, Any]] = None,
    ) -> CompositeExtractor:
        if is_root:
            assert field_enables_flag_extractor is None
        else:
            assert app_cfg_type is None, (
                "App type passing is not allowed for non-root composite extractor for type"
                f" ({get_type_full_name(settings_type)})"
            )
            assert field_overrides is None, "Field overrides are supported only for root extractors"

        effective_fields_overrides: dict[str, Any] = {}
        if field_overrides is not None:
            effective_fields_overrides.update(**field_overrides)

        possible_types_set = cls.unwrap_union(settings_type, ignore_none=True)
        assert len(possible_types_set) == 1, "Polymorphism is not supported at this moment"

        for possible_type in possible_types_set:
            if attr.has(possible_type):
                attr.resolve_types(possible_type)

        field_actual_type = next(iter(possible_types_set))

        map_field_name_field = attr.fields_dict(field_actual_type)
        extractors_map: dict[str, SDictExtractor] = {}

        for field_name, field in map_field_name_field.items():
            field_s_meta = SMeta.from_attrib(field)
            if field_s_meta is None:
                continue

            # Adding app type to fields override if was provided in args
            # Extractor will be also created (but not used during extraction)
            if field_s_meta.is_app_cfg_type and app_cfg_type is not None:
                effective_fields_overrides[field_name] = app_cfg_type

            # Determining field default
            field_default: Any

            if field_s_meta.fallback_cfg_key is not None:
                assert field_s_meta.fallback_factory is None
                field_default = getattr(fallback_cfg, field_s_meta.fallback_cfg_key, NOT_SET)

            elif field_s_meta.fallback_factory is not None:
                assert field_s_meta.fallback_cfg_key is None
                # It's a fuse to prevent unexpected behaviour in case of
                #  fallbacks defined in both nested class and root class.
                # If become critical to have fallback factory in nested set
                #  try to rewrite this check to prevent overlapping fallback factories
                assert is_root, f"Fallback factories may be applied only on top-level configs, not on {settings_type}"

                fallback_factory: FallbackFactory

                if isinstance(field_s_meta.fallback_factory, staticmethod):
                    fallback_factory = field_s_meta.fallback_factory.__func__
                else:
                    fallback_factory = field_s_meta.fallback_factory

                assert inspect.isfunction(fallback_factory) or inspect.ismethod(fallback_factory)
                factory_signature = inspect.signature(fallback_factory)

                if len(factory_signature.parameters) == 0:
                    field_default = fallback_factory()
                elif len(factory_signature.parameters) == 1:
                    field_default = fallback_factory(fallback_cfg)
                elif len(factory_signature.parameters) == 2:
                    field_default = fallback_factory(fallback_cfg, app_cfg_type)
                else:
                    raise SettingsLoadingException(f"Unexpected signature of fallback factory for {field}")

            elif cls.is_sub_settings_field(field) and default is not None and default is not NOT_SET:
                field_default = getattr(default, field_name)

            else:
                field_default = NOT_SET

            child_extractor: SDictExtractor

            if field_s_meta.name is None:
                child_extractor = DefaultOnlyExtractor(
                    key=None,
                    default=field_default,
                    expected_type=field.type,
                )
            elif cls.is_sub_settings_field(field):
                assert field_s_meta.fallback_cfg_key is None, "At this moment nested objects can not have defaults"
                enabled_key_name = field_s_meta.enabled_key_name

                child_extractor = cls.build_composite_extractor(
                    settings_type=field.type,
                    key_prefix=key_prefix + (field_s_meta.name,),
                    default=field_default,
                    json_converter=field_s_meta.json_converter,
                    field_enables_flag_extractor=None
                    if enabled_key_name is None
                    else cls.build_default_extractor(
                        name="enabled",
                        field_type=bool,
                        meta=SMeta(enabled_key_name),
                        default=NOT_SET,
                    ),
                )
            else:
                assert field_s_meta.enabled_key_name is None, (
                    f"enabled_key_name is supported only for nested configs:"
                    f" {get_type_full_name(settings_type)}.{field_name}"
                )
                child_extractor = cls.build_default_extractor(
                    name=field_name,
                    field_type=field.type,
                    meta=field_s_meta,
                    default=field_default,
                )

            extractors_map[field_name] = child_extractor

        return CompositeExtractor(
            key=key_prefix[-1] if key_prefix else None,
            expected_type=settings_type,
            composition_factory=field_actual_type,
            map_field_name_subextractor=extractors_map,
            default=default,
            is_root=is_root,
            json_value_converter=json_converter,
            field_enables_flag_extractor=field_enables_flag_extractor,
            map_field_name_field_override=effective_fields_overrides,
        )

    @classmethod
    def get_app_cfg_type_field(cls, settings_type: Type[SettingsBase]) -> Optional[attr.Attribute]:
        candidates = [field for field in attr.fields(settings_type) if SMeta.from_attrib(field).is_app_cfg_type]
        if len(candidates) > 1:
            raise SettingsLoadingException(
                f"Settings class {get_type_full_name(settings_type)}"
                f" has more than one app_type attributes: {candidates}"
            )
        else:
            return next(iter(candidates)) if candidates else None

    @classmethod
    def get_app_cfg_type_value_from_env(cls, settings_type: Type[SettingsBase], s_dict: SDict) -> Any:
        app_cfg_type_field = cls.get_app_cfg_type_field(settings_type)
        if app_cfg_type_field is None:
            return None

        extractor = cls.build_default_extractor(
            name=app_cfg_type_field.name,
            field_type=app_cfg_type_field.type,
            meta=SMeta.from_attrib(app_cfg_type_field),
            default=NOT_SET,
        )
        return extractor.extract(s_dict)

    @classmethod
    def build_top_level_extractor(
        cls,
        settings_type: Type[_SETTINGS_TV],
        key_prefix: Optional[str] = None,
        fallback_cfg: Any = None,
        app_cfg_type: Any = None,
        default_value: Optional[_SETTINGS_TV] = None,
        field_overrides: Optional[dict[str, Any]] = None,
    ) -> CompositeExtractor:
        return cls.build_composite_extractor(
            settings_type=settings_type,
            key_prefix=() if key_prefix is None else (key_prefix,),
            fallback_cfg=fallback_cfg,
            app_cfg_type=app_cfg_type,
            is_root=True,
            default=NOT_SET if default_value is None else default_value,
            field_overrides=field_overrides,
        )

    def load_settings(
        self,
        settings_type: Type[_SETTINGS_TV],
        key_prefix: Optional[str] = None,
        fallback_cfg_resolver: Optional[FallbackConfigResolver] = None,
        default_value: Optional[_SETTINGS_TV] = None,
    ) -> _SETTINGS_TV:
        effective_s_dict = self.s_dict
        effective_s_dict = remap_env(effective_s_dict)

        # TODO FIX: Error handling
        keys_from_files = self.load_file_mapped_keys(self.s_dict)

        keys_from_files_intersection = effective_s_dict.keys() & keys_from_files.keys()

        if keys_from_files_intersection:
            raise InvalidConfigValueException(
                f"Got intersection of explicit env vars & file mapped vars:"
                f" {' '.join(sorted(keys_from_files_intersection))}"
            )

        effective_s_dict = {**effective_s_dict, **keys_from_files}

        fallback_cfg: Any = None

        if fallback_cfg_resolver is not None:
            fallback_cfg = fallback_cfg_resolver.resolve(effective_s_dict)

        app_cfg_type = self.get_app_cfg_type_value_from_env(settings_type, effective_s_dict)

        extractor = self.build_top_level_extractor(
            settings_type=settings_type,
            key_prefix=key_prefix,
            fallback_cfg=fallback_cfg,
            app_cfg_type=app_cfg_type,
            default_value=default_value,
        )

        return extractor.extract(effective_s_dict)


def load_settings_from_env_with_fallback_legacy(
    settings_type: Type[_SETTINGS_TV],
    fallback_cfg_resolver: FallbackConfigResolver,
    env: Optional[SDict] = None,
) -> _SETTINGS_TV:
    effective_env = os.environ if env is None else env

    return EnvSettingsLoader(effective_env).load_settings(
        settings_type,
        fallback_cfg_resolver=fallback_cfg_resolver,
    )


def load_settings_from_env_with_fallback(
    settings_type: Type[_SETTINGS_TV],
    env: Optional[SDict] = None,
    default_fallback_cfg_resolver: Optional[FallbackConfigResolver] = None,
) -> _SETTINGS_TV:
    effective_env = os.environ if env is None else env

    if YamlFileConfigResolver.is_config_enabled(s_dict=effective_env):
        default_fallback_cfg_resolver = YamlFileConfigResolver()

    if default_fallback_cfg_resolver is None:
        LOGGER.warning("No fallback resolver")

    return EnvSettingsLoader(effective_env).load_settings(
        settings_type,
        fallback_cfg_resolver=default_fallback_cfg_resolver,
    )


@no_type_check  # mypy is barely working with dynamic attrs classes
def load_connectors_settings_from_env_with_fallback(
    settings_registry: dict[ConnectionType, Type[ConnectorSettingsBase]],
    fallbacks: dict[ConnectionType, SettingsFallbackType],
    whitelist: Optional[Collection[ConnectionType]] = None,
    env: Optional[SDict] = None,
    fallback_cfg_resolver: Optional[FallbackConfigResolver] = None,
) -> dict[ConnectionType, ConnectorSettingsBase]:
    settings_class = generate_connectors_settings_class(settings_registry, whitelist)

    def connectors_fallback(full_cfg: ConnectorsConfigType):
        full_settings = reduce(lambda settings, fallback: settings | fallback(full_cfg), fallbacks.values(), {})
        return settings_class(**full_settings)

    connectors = attr.make_class(
        "Connectors",
        {
            "CONNECTORS": s_attrib(
                "CONNECTORS",
                type_class=Optional[settings_class],
                fallback_factory=connectors_fallback,
            )
        },
        frozen=True,
    )

    loaded_settings = load_settings_from_env_with_fallback(connectors, env, fallback_cfg_resolver)
    connectors_settings: dict[ConnectionType, ConnectorSettingsBase] = {}
    if loaded_settings.CONNECTORS is None:
        return connectors_settings

    for name in attr.fields_dict(settings_class):
        conn_type = ConnectionType(name.lower())
        connector_settings = getattr(loaded_settings.CONNECTORS, name)
        if connector_settings is not None:
            assert isinstance(connector_settings, ConnectorSettingsBase)
            connectors_settings[conn_type] = connector_settings
    return connectors_settings

lib/dl_configs/dl_configs/settings_loaders/loader_env.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Callable,
    ClassVar,
    Optional,
    Type,
    TypeVar,
    cast,
)

import attr

from dl_configs.settings_loaders.common import FallbackFactory


class ReservedKeys:
    JSON_VALUE: ClassVar[str] = "JSON_VALUE"


@attr.s(frozen=True)
class SMeta:
    attrs_meta_key: ClassVar[str] = "s_meta"

    name: Optional[str] = attr.ib()
    fallback_cfg_key: Optional[str] = attr.ib(default=None)
    fallback_factory: Optional[Callable[[Any], Any]] = attr.ib(default=None)
    env_var_converter: Optional[Callable[[str], Any]] = attr.ib(default=None)
    json_converter: Optional[Callable[[Any], Any]] = attr.ib(default=None)
    enabled_key_name: Optional[str] = attr.ib(default=None)
    is_app_cfg_type: bool = attr.ib(default=False)

    def __attrs_post_init__(self):
        assert not (
            self.fallback_cfg_key is not None and self.fallback_factory is not None
        ), "fallback_factory and fallback_cfg_key can not be defined together"

    def to_dict(self):
        return {self.attrs_meta_key: self}

    @classmethod
    def from_attrib(cls, attrib: Any) -> Optional[SMeta]:
        return attrib.metadata.get(cls.attrs_meta_key)


def s_attrib(
    key_name: Optional[str],
    type_class: object = None,
    fallback_cfg_key: Optional[str] = None,
    fallback_factory: Optional[FallbackFactory] = None,
    env_var_converter: Optional[Callable[[str], Any]] = None,
    json_converter: Optional[Callable[[Any], Any]] = None,
    missing: Any = attr.NOTHING,
    missing_factory: Any = None,
    sensitive: bool = False,
    enabled_key_name: Optional[str] = None,
    is_app_cfg_type: bool = False,
) -> attr.Attribute:
    if is_app_cfg_type:
        assert fallback_cfg_key is None
        assert fallback_factory is None
        assert missing_factory is None
        assert missing is attr.NOTHING

    s_meta = SMeta(
        name=key_name,
        fallback_cfg_key=fallback_cfg_key,
        fallback_factory=fallback_factory,
        env_var_converter=env_var_converter,
        json_converter=json_converter,
        enabled_key_name=enabled_key_name,
        is_app_cfg_type=is_app_cfg_type,
    )

    return attr.ib(
        type=type_class,
        default=missing,
        factory=missing_factory,
        repr=not sensitive,
        kw_only=True,
        metadata=s_meta.to_dict(),
    )


_REQUIRED_TV = TypeVar("_REQUIRED_TV")


@attr.s(frozen=True)
class RequiredValue:
    message: Optional[str] = attr.ib(default=None)


def required(t: Type[_REQUIRED_TV]) -> _REQUIRED_TV:
    """
    Usage: set constructor args of attrs objects in fallback_config_factory.
    Fields with `REQUIRED_VALUE` value must be replaced by settings loader or exception must be thrown
    Typing and cast are only to trick MyPy
    """
    return cast(t, RequiredValue())

lib/dl_configs/dl_configs/settings_loaders/meta_definition.py
datalens-tech/datalens-backend
from __future__ import annotations


class SettingsBase:
    """
    Just a marker class
    """

lib/dl_configs/dl_configs/settings_loaders/settings_obj_base.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_configs.enums import RedisMode
from dl_configs.settings_loaders.meta_definition import s_attrib
from dl_configs.settings_loaders.settings_obj_base import SettingsBase
from dl_configs.utils import split_by_comma


def redis_mode_env_var_converter(env_value: str) -> RedisMode:
    return {
        "sentinel": RedisMode.sentinel,
        "single_host": RedisMode.single_host,
    }[env_value.lower()]


@attr.s(frozen=True)
class RedisSettings(SettingsBase):
    MODE: RedisMode = s_attrib("MODE", env_var_converter=redis_mode_env_var_converter)
    CLUSTER_NAME: str = s_attrib("NAME", missing="")
    HOSTS: tuple[str, ...] = s_attrib("HOSTS", env_var_converter=split_by_comma)
    PORT: int = s_attrib("PORT")
    DB: int = s_attrib("DB")
    PASSWORD: str = s_attrib("PASSWORD", sensitive=True, missing=None)
    SSL: Optional[bool] = s_attrib("SSL", missing=None)


@attr.s(frozen=True)
class CorsSettings(SettingsBase):
    ALLOWED_ORIGINS: tuple[str, ...] = s_attrib("ALLOWED_ORIGINS", env_var_converter=split_by_comma)
    ALLOWED_HEADERS: tuple[str, ...] = s_attrib("ALLOWED_HEADERS", env_var_converter=split_by_comma)
    EXPOSE_HEADERS: tuple[str, ...] = s_attrib("EXPOSE_HEADERS", env_var_converter=split_by_comma)


@attr.s(frozen=True)
class CsrfSettings(SettingsBase):
    METHODS: tuple[str, ...] = s_attrib("METHODS", env_var_converter=split_by_comma)
    HEADER_NAME: str = s_attrib("HEADER_NAME")
    TIME_LIMIT: int = s_attrib("TIME_LIMIT")
    SECRET: str = s_attrib("SECRET", sensitive=True, missing=None)


@attr.s(frozen=True)
class S3Settings(SettingsBase):
    ACCESS_KEY_ID: str = s_attrib("ACCESS_KEY_ID", sensitive=True, missing=None)
    SECRET_ACCESS_KEY: str = s_attrib("SECRET_ACCESS_KEY", sensitive=True, missing=None)
    ENDPOINT_URL: str = s_attrib("ENDPOINT_URL")


@attr.s(frozen=True)
class GoogleAppSettings(SettingsBase):
    API_KEY: str = s_attrib("API_KEY", sensitive=True)
    CLIENT_ID: str = s_attrib("CLIENT_ID", sensitive=True)
    CLIENT_SECRET: str = s_attrib("CLIENT_SECRET", sensitive=True)

lib/dl_configs/dl_configs/settings_submodels.py
datalens-tech/datalens-backend
import os
from typing import (
    Callable,
    Container,
    TypeVar,
)

from dl_constants.enums import ConnectionType


ROOT_CERTIFICATES_FILENAME = "/etc/ssl/certs/ca-certificates.crt"
TEMP_ROOT_CERTIFICATES_FOLDER_PATH = "/tmp/ssl/certs/"


_T = TypeVar("_T")


def validate_one_of(valid_values: Container[_T]) -> Callable[[_T], _T]:
    def validator(value: _T) -> _T:
        if value not in valid_values:
            raise ValueError(f"Expected one of {valid_values}, but got '{value}'")
        return value

    return validator


def get_root_certificates() -> bytes:
    with open(get_root_certificates_path(), "rb") as fobj:
        return fobj.read()


def get_root_certificates_path() -> str:
    return ROOT_CERTIFICATES_FILENAME


def get_temp_root_certificates_folder_path() -> str:
    os.makedirs(TEMP_ROOT_CERTIFICATES_FOLDER_PATH, exist_ok=True)

    return TEMP_ROOT_CERTIFICATES_FOLDER_PATH


def split_by_comma(s: str) -> tuple[str, ...]:
    return tuple(entry.strip() for entry in s.split(",") if entry)


def conn_type_set_env_var_converter(s: str) -> set[ConnectionType]:
    return set(ConnectionType(entry.strip()) for entry in s.split(",") if entry)

lib/dl_configs/dl_configs/utils.py
datalens-tech/datalens-backend
from typing import Optional

import attr
from attr import fields

from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.settings_loaders.connectors_settings import generate_connectors_settings_class
from dl_configs.settings_loaders.meta_definition import s_attrib
from dl_configs.settings_loaders.settings_obj_base import SettingsBase
from dl_constants.enums import ConnectionType


CONNECTION_TYPE_APPALLING = ConnectionType.declare("appalling")
CONNECTION_TYPE_EMPTY_YET_COOL = ConnectionType.declare("empty_yet_cool")
CONNECTION_TYPE_FANTASTIC = ConnectionType.declare("fantastic")
CONNECTION_TYPE_WONDERFUL = ConnectionType.declare("wonderful")


@attr.s(frozen=True)
class WonderfulConnectorSettings(ConnectorSettingsBase):
    WONDERFUL_STRING_SETTING: str = s_attrib("WONDERFUL_STRING_SETTING", missing=None)
    WONDERFUL_LIST_OF_STRINGS_SETTING: list[str] = s_attrib("WONDERFUL_LIST_OF_STRINGS_SETTING", missing=None)


@attr.s(frozen=True)
class FantasticConnectorSettings(ConnectorSettingsBase):
    FANTASTIC_INTEGER_SETTING: int = s_attrib("FANTASTIC_INTEGER_SETTING", missing=None)
    FANTASTIC_FLOAT_SETTING: float = s_attrib("FANTASTIC_FLOAT_SETTING", missing=None)


@attr.s(frozen=True)
class EmptyYetCoolConnectorSettings(ConnectorSettingsBase):
    pass


@attr.s(frozen=True)
class AppallingConnectorSettings(ConnectorSettingsBase):
    APPALING_SET_SETTING: set[int] = s_attrib("APPALING_SET_SETTING", missing=None)


@attr.s(frozen=True)
class ExpectedConnectorsSettings(SettingsBase):
    EMPTY_YET_COOL: Optional[EmptyYetCoolConnectorSettings] = s_attrib("EMPTY_YET_COOL", missing=None)
    FANTASTIC: Optional[FantasticConnectorSettings] = s_attrib("FANTASTIC", missing=None)
    WONDERFUL: Optional[WonderfulConnectorSettings] = s_attrib("WONDERFUL", missing=None)


def test_connectors_settings_class_generator():
    settings_registry = {
        CONNECTION_TYPE_EMPTY_YET_COOL: EmptyYetCoolConnectorSettings,
        CONNECTION_TYPE_FANTASTIC: FantasticConnectorSettings,
        CONNECTION_TYPE_WONDERFUL: WonderfulConnectorSettings,
    }
    generated_class = generate_connectors_settings_class(settings_registry=settings_registry)
    assert issubclass(generated_class, SettingsBase)
    assert fields(generated_class) == fields(ExpectedConnectorsSettings)


def test_connectors_settings_class_generation_with_whitelist():
    settings_registry = {
        CONNECTION_TYPE_APPALLING: AppallingConnectorSettings,
        CONNECTION_TYPE_EMPTY_YET_COOL: EmptyYetCoolConnectorSettings,
        CONNECTION_TYPE_FANTASTIC: FantasticConnectorSettings,
        CONNECTION_TYPE_WONDERFUL: WonderfulConnectorSettings,
    }
    whitelist = {CONNECTION_TYPE_EMPTY_YET_COOL, CONNECTION_TYPE_FANTASTIC, CONNECTION_TYPE_WONDERFUL}
    generated_class = generate_connectors_settings_class(settings_registry=settings_registry, whitelist=whitelist)
    assert issubclass(generated_class, SettingsBase)
    assert fields(generated_class) == fields(ExpectedConnectorsSettings)

lib/dl_configs/dl_configs_tests/unit/test_connectors_settings_generation.py
datalens-tech/datalens-backend
from __future__ import annotations

import json

from cryptography import fernet

from dl_configs.crypto_keys import get_crypto_keys_config_from_env


def test_minimal():
    env = dict(FERNET_KEY=fernet.Fernet.generate_key().decode("ascii"))
    cc = get_crypto_keys_config_from_env(env)
    assert cc.map_id_key == {"0": env["FERNET_KEY"]}
    assert cc.actual_key_id == "0"


def test_single_zero_id_key():
    env = dict(DL_CRY_KEY_VAL_ID_0=fernet.Fernet.generate_key().decode("ascii"))
    cc = get_crypto_keys_config_from_env(env)
    assert cc.map_id_key == {"0": env["DL_CRY_KEY_VAL_ID_0"]}
    assert cc.actual_key_id == "0"


def test_multiple_keys():
    env = dict(
        DL_CRY_KEY_VAL_ID_0=fernet.Fernet.generate_key().decode("ascii"),
        DL_CRY_KEY_VAL_ID_1=fernet.Fernet.generate_key().decode("ascii"),
        DL_CRY_ACTUAL_KEY_ID="1",
    )
    cc = get_crypto_keys_config_from_env(env)
    assert cc.map_id_key == {
        "0": env["DL_CRY_KEY_VAL_ID_0"],
        "1": env["DL_CRY_KEY_VAL_ID_1"],
    }
    assert cc.actual_key_id == "1"


def test_multiple_keys_complex_names():
    env = dict(
        DL_CRY_KEY_VAL_ID_0=fernet.Fernet.generate_key().decode("ascii"),
        DL_CRY_KEY_VAL_ID_CL_PRD_1=fernet.Fernet.generate_key().decode("ascii"),
        DL_CRY_ACTUAL_KEY_ID="CL_PRD_1",
    )
    cc = get_crypto_keys_config_from_env(env)
    assert cc.map_id_key == {
        "0": env["DL_CRY_KEY_VAL_ID_0"],
        "CL_PRD_1": env["DL_CRY_KEY_VAL_ID_CL_PRD_1"],
    }
    assert cc.actual_key_id == "CL_PRD_1"


def test_json():
    actual_key_id = "some_k_id_0"
    cry_dict = dict(
        keys={
            actual_key_id: fernet.Fernet.generate_key().decode("ascii"),
            "outdated_0": fernet.Fernet.generate_key().decode("ascii"),
        },
        actual_key_id=actual_key_id,
    )
    env = dict(
        DL_CRY_JSON_VALUE=json.dumps(cry_dict),
    )
    cc = get_crypto_keys_config_from_env(env)
    assert cc.map_id_key == cry_dict["keys"]
    assert cc.actual_key_id == actual_key_id

lib/dl_configs/dl_configs_tests/unit/test_crypto_controller_from_env.py
datalens-tech/datalens-backend
# Commented due to attr.resolve_types() does not work with classes declared in non-global context
# from __future__ import annotations
import enum
import json
from os import path
import tempfile
from typing import (
    Any,
    Callable,
    Dict,
    Optional,
    Type,
)

import attr
import pytest
import shortuuid

from dl_configs.settings_loaders.common import SDict
from dl_configs.settings_loaders.fallback_cfg_resolver import ConstantFallbackConfigResolver
from dl_configs.settings_loaders.loader_env import (
    ConfigFieldMissing,
    EnvSettingsLoader,
)
from dl_configs.settings_loaders.meta_definition import (
    required,
    s_attrib,
)
from dl_configs.settings_loaders.settings_obj_base import SettingsBase


def load_settings(
    env: SDict,
    settings_type: Type[SettingsBase],
    key_prefix: Optional[str] = None,
    fallback_env: Any = None,
) -> Any:
    fallback_cfg_resolver = ConstantFallbackConfigResolver(fallback_env)
    loader = EnvSettingsLoader(env)
    return loader.load_settings(
        settings_type,
        key_prefix=key_prefix,
        fallback_cfg_resolver=fallback_cfg_resolver,
    )


def perform_loader_env_check(
    env: SDict,
    expected_settings: Any,
    key_prefix: Optional[str] = None,
    fallback_env: Any = None,
) -> None:
    actual_settings = load_settings(
        env=env,
        settings_type=type(expected_settings),
        key_prefix=key_prefix,
        fallback_env=fallback_env,
    )

    assert actual_settings == expected_settings


def test_simple_scalars_loading():
    @attr.s
    class SomeSettings(SettingsBase):
        a: int = s_attrib("a")
        b: str = s_attrib("b")
        c: float = s_attrib("c")
        bool_0: bool = s_attrib("bool_0")
        bool_py_true: bool = s_attrib("bool_py_true")
        bool_py_false: bool = s_attrib("bool_py_false")
        bool_true: bool = s_attrib("bool_true")
        bool_false: bool = s_attrib("bool_false")
        bool_1: bool = s_attrib("bool_1")
        z: Dict[str, str] = s_attrib("z", env_var_converter=json.loads)

    perform_loader_env_check(
        key_prefix="DL",
        env=dict(
            DL_A="123",
            DL_B="abc",
            DL_C="1.2",
            DL_BOOL_0="0",
            DL_BOOL_1="1",
            DL_BOOL_PY_TRUE="True",
            DL_BOOL_PY_FALSE="False",
            DL_BOOL_TRUE="true",
            DL_BOOL_FALSE="false",
            DL_Z="""{"a": "b"}""",
        ),
        expected_settings=SomeSettings(
            a=123,
            b="abc",
            c=1.2,
            bool_0=False,
            bool_1=True,
            bool_py_true=True,
            bool_py_false=False,
            bool_true=True,
            bool_false=False,
            z=dict(a="b"),
        ),
    )


def test_map_loading_simple():
    @attr.s
    class SettingsMapStrStr(SettingsBase):
        m: Dict[str, str] = s_attrib("M_MAP")

    perform_loader_env_check(
        dict(M_MAP_k1="val11", M_MAP_k2="val22"), expected_settings=SettingsMapStrStr(m=dict(k1="val11", k2="val22"))
    )


def test_map_loading_nested():
    @attr.s
    class Nested(SettingsBase):
        m: Dict[str, str] = s_attrib("M_MAP")

    @attr.s
    class SettingsMapStrStr(SettingsBase):
        nested: Nested = s_attrib("NESTED")

    perform_loader_env_check(
        env=dict(
            NESTED_M_MAP_k1="val11",
            NESTED_M_MAP_k2="val22",
        ),
        expected_settings=SettingsMapStrStr(nested=Nested(m=dict(k1="val11", k2="val22"))),
    )


def test_map_loading_int():
    @attr.s
    class SettingsMapStrStr(SettingsBase):
        m: Dict[str, int] = s_attrib("M_MAP")

    perform_loader_env_check(
        env=dict(
            M_MAP_k1="11",
            M_MAP_k2="22",
        ),
        expected_settings=SettingsMapStrStr(m=dict(k1=11, k2=22)),
    )


def test_map_loading_empty():
    @attr.s
    class SettingsMapStrStr(SettingsBase):
        m: Dict[str, int] = s_attrib("M_MAP", missing_factory=dict)

    perform_loader_env_check(env=dict(), expected_settings=SettingsMapStrStr(m={}))


def test_composite_loading():
    @attr.s
    class SuperNestedSettings(SettingsBase):
        param_1: bool = s_attrib("param_1")

    @attr.s
    class NestedSettings(SettingsBase):
        n1: int = s_attrib("n1")
        n2: int = s_attrib("n2")
        nested: SuperNestedSettings = s_attrib("nested")

    @attr.s
    class SomeSettings(SettingsBase):
        a: int = s_attrib("a")
        nested: NestedSettings = s_attrib("nested")
        nullable_nested: Optional[NestedSettings] = s_attrib("nullable_nested", missing=None)

    perform_loader_env_check(
        key_prefix=None,
        env=dict(A="1", NESTED_N1="11", NESTED_N2="12", NESTED_NESTED_PARAM_1="1"),
        expected_settings=SomeSettings(
            a=1,
            nested=NestedSettings(n1=11, n2=12, nested=SuperNestedSettings(param_1=True)),
            nullable_nested=None,
        ),
    )


def test_scalar_defaults():
    class Fallback:
        THE_A = 75
        THE_B = 88

    @attr.s
    class SomeSettings(SettingsBase):
        a: int = s_attrib("a", fallback_cfg_key="THE_A")
        b: int = s_attrib("b", fallback_cfg_key="THE_B")
        c: Optional[int] = s_attrib("c", fallback_cfg_key="THE_NON_EXISTING_KEY", missing=None)

    perform_loader_env_check(
        key_prefix=None,
        env=dict(
            B="99",
        ),
        expected_settings=SomeSettings(
            a=75,
            b=99,
            c=None,
        ),
        fallback_env=Fallback,
    )


def test_composite_default():
    class Fallback:
        HOST = "example.com"
        PORT = 443

    @attr.s
    class BackendGroup(SettingsBase):
        host: str = s_attrib("host")
        port: int = s_attrib("port")
        proto: str = s_attrib("proto")

    @attr.s
    class SomeSettings(SettingsBase):
        bg: BackendGroup = s_attrib(
            "bg", fallback_factory=lambda env: BackendGroup(host=env.HOST, port=env.PORT, proto="http")
        )

    perform_loader_env_check(
        key_prefix=None,
        env=dict(BG_PROTO="https", BG_PORT=8443),
        expected_settings=SomeSettings(
            bg=BackendGroup(host=Fallback.HOST, port=8443, proto="https")  # From env  # From env
        ),
        fallback_env=Fallback,
    )


def test_scalar_fallback_order():
    class FallbackEnv:
        THE_B = "THE B FROM FALLBACK"

    @attr.s
    class SomeSettings:
        a: str = s_attrib("dl_a", fallback_cfg_key="THE_A", missing="attrs default a")
        b: str = s_attrib("dl_b", fallback_cfg_key="THE_B", missing="attrs default b")
        c: str = s_attrib("dl_c", fallback_cfg_key="THE_C", missing="attrs default c")

    perform_loader_env_check(
        env=dict(DL_A="a from proc env"),
        expected_settings=SomeSettings(
            # A was declared in user env so we expect value from it
            a="a from proc env",
            # There is no B in user env but there is value in fallback env. So we expect value from fallback
            b=FallbackEnv.THE_B,
            # There is no C neither in user env nor in fallback one. So we expect default from attrs
            c="attrs default c",
        ),
        fallback_env=FallbackEnv,
    )


def test_ignore():
    class Fallback:
        EP_HOST: str = "host-from-fb"
        EP_PORT: int = 0

    @attr.s(frozen=True)
    class Nested(SettingsBase):
        host: str = s_attrib("HOST")
        port: int = s_attrib("PORT")

    @attr.s(frozen=True)
    class DefaultedSettings(SettingsBase):
        endpoint: Optional[Nested] = s_attrib(
            "EP",
            enabled_key_name="ENABLED",
            fallback_factory=lambda fb: Nested(host=fb.EP_HOST, port=fb.EP_PORT),
            missing=Nested(host="localhost", port=321),
        )

    @attr.s(frozen=True)
    class NonDefaultedSettings(SettingsBase):
        endpoint: Optional[Nested] = s_attrib("EP", enabled_key_name="ENABLED")

    # Check that endpoint is not despite of fallback
    perform_loader_env_check(
        env=dict(EP_ENABLED="0"),
        expected_settings=DefaultedSettings(endpoint=None),
        fallback_env=Fallback,
    )

    # Check that we got a fallback if no keys in env and ignore key is not true
    perform_loader_env_check(
        env=dict(
            EP_ENABLED="1",
            EP_PORT=9999,
        ),
        expected_settings=DefaultedSettings(endpoint=Nested(host=Fallback.EP_HOST, port=9999)),
        fallback_env=Fallback,
    )

    # Check that we will not fail with ignored non-defaulted key
    perform_loader_env_check(
        env=dict(
            EP_ENABLED="0",
        ),
        expected_settings=NonDefaultedSettings(endpoint=None),
    )


def test_sensitive_simple():
    # TODO FIX: add loading report check when it's ready
    @attr.s
    class Nested(SettingsBase):
        a: int = s_attrib("a", sensitive=True)

    @attr.s
    class SomeSettings(SettingsBase):
        nested: Nested = s_attrib("n")
        nested_sensitive: Nested = s_attrib("n_sens", sensitive=True)
        sensitive: str = s_attrib("sens", sensitive=True)

    assert attr.fields_dict(Nested)["a"].repr is False

    assert attr.fields_dict(SomeSettings)["nested"].repr is True
    assert attr.fields_dict(SomeSettings)["nested_sensitive"].repr is False
    assert attr.fields_dict(SomeSettings)["sensitive"].repr is False


def test_required_in_fallback_factory():
    class Fallback:
        DB_HOST = "127.0.0.1"

    @attr.s
    class DBConfig(SettingsBase):
        host: str = s_attrib("HOST")
        user: str = s_attrib("USER")
        password: str = s_attrib("PASSWORD", sensitive=True)

    @attr.s
    class SomeSettings(SettingsBase):
        db: DBConfig = s_attrib(
            "DB",
            fallback_factory=lambda fb: DBConfig(
                host=fb.DB_HOST,
                user=required(str),
                password=required(str),
            ),
        )

    perform_loader_env_check(
        env=dict(
            DB_USER="user_from_env",
            DB_PASSWORD="password_from_env",
        ),
        expected_settings=SomeSettings(
            db=DBConfig(host=Fallback.DB_HOST, user="user_from_env", password="password_from_env")
        ),
        fallback_env=Fallback,
    )

    with pytest.raises(ConfigFieldMissing) as exc:
        load_settings(env=dict(), settings_type=SomeSettings, fallback_env=Fallback)

    assert exc.value.field_set == {"db.password", "db.user"}

    with pytest.raises(ConfigFieldMissing) as exc:
        load_settings(env=dict(DB_USER="user_from_env"), settings_type=SomeSettings, fallback_env=Fallback)

    assert exc.value.field_set == {"db.password"}


def test_fallback_factory_static_method():
    class Fallback1:
        fb_1_only = "fb_1_only"

    class Fallback2:
        fb_2_only = "fb_2_only"

    @attr.s
    class Nested(SettingsBase):
        @staticmethod
        def a_default(env):
            if isinstance(env, Fallback1):
                return env.fb_1_only
            elif isinstance(env, Fallback2):
                return env.fb_2_only
            else:
                raise ValueError("Unexpected fallback type")

        a: str = s_attrib("a", fallback_factory=a_default)

    perform_loader_env_check(env={}, expected_settings=Nested(a=Fallback1.fb_1_only), fallback_env=Fallback1())

    perform_loader_env_check(env={}, expected_settings=Nested(a=Fallback2.fb_2_only), fallback_env=Fallback2())


def test_app_type():
    class Fallback:
        SOME_KEY = "whatever"

    class Mode(enum.Enum):
        mode_a = enum.auto()
        mode_b = enum.auto()

    @attr.s
    class TypedSettings(SettingsBase):
        mode: Mode = s_attrib("APP_MODE", is_app_cfg_type=True, env_var_converter=lambda s: Mode[s])

        @staticmethod
        def default_depends_on_mode(fallback_config: Fallback, mode: Mode):
            return f"{mode.name} + {fallback_config.SOME_KEY}"

        depends_on_mode: str = s_attrib("DEPENDS_ON_MODE", fallback_factory=default_depends_on_mode)

        @staticmethod
        def default_not_depends_on_mode(fallback_config: Fallback):
            return fallback_config.SOME_KEY

        not_depends_on_mode: str = s_attrib("NOT_DEPENDS_ON_MODE", fallback_factory=default_not_depends_on_mode)

    perform_loader_env_check(
        dict(APP_MODE="mode_a"),
        TypedSettings(
            mode=Mode.mode_a,
            depends_on_mode=f"{Mode.mode_a.name} + {Fallback.SOME_KEY}",
            not_depends_on_mode=Fallback.SOME_KEY,
        ),
        fallback_env=Fallback(),
    )


def test_app_type_field_override():
    class Fallback:
        SOME_KEY = "whatever"

    class Mode(enum.Enum):
        mode_a = enum.auto()
        mode_b = enum.auto()

    @attr.s
    class TypedSettings(SettingsBase):
        mode: Mode = s_attrib("APP_MODE", is_app_cfg_type=True, env_var_converter=lambda s: Mode[s])
        override_me: str = s_attrib("OVERRIDE_ME")

        @staticmethod
        def default_depends_on_mode(fallback_config: Fallback, mode: Mode):
            return f"{mode.name} + {fallback_config.SOME_KEY}"

        depends_on_mode: str = s_attrib("DEPENDS_ON_MODE", fallback_factory=default_depends_on_mode)

    env = dict(OVERRIDE_ME="yes")

    extractor = EnvSettingsLoader(env).build_top_level_extractor(
        TypedSettings,
        app_cfg_type=Mode.mode_a,
        fallback_cfg=Fallback,
    )
    actual_settings = extractor.extract(env)

    assert actual_settings == TypedSettings(
        mode=Mode.mode_a, depends_on_mode=f"{Mode.mode_a.name} + {Fallback.SOME_KEY}", override_me="yes"
    )

    env = dict(OVERRIDE_ME="yes")
    # TODO FIX: Make dedicated test when tests will migrate on classes
    extractor = EnvSettingsLoader(env).build_top_level_extractor(
        TypedSettings,
        app_cfg_type=Mode.mode_b,
        fallback_cfg=Fallback,
        field_overrides=dict(override_me="totally_overridden"),
    )
    actual_settings = extractor.extract(env)

    assert actual_settings == TypedSettings(
        mode=Mode.mode_b, depends_on_mode=f"{Mode.mode_b.name} + {Fallback.SOME_KEY}", override_me="totally_overridden"
    )


def test_fallback_only():
    class Fallback:
        SOME_KEY = "whatever"

    @attr.s
    class Nested(SettingsBase):
        a: int = s_attrib("A", sensitive=True)

    @attr.s
    class TypedSettings(SettingsBase):
        no_env_1: Optional[str] = s_attrib(None, fallback_cfg_key="SOME_KEY", missing=None)
        via_env: Optional[str] = s_attrib("VIA_ENV", fallback_cfg_key="SOME_KEY")
        no_env_nested: Nested = s_attrib(None, fallback_factory=lambda: Nested(a=1))

    perform_loader_env_check(
        dict(VIA_ENV="mode_a"),
        TypedSettings(
            no_env_1=Fallback.SOME_KEY,
            via_env="mode_a",
            no_env_nested=Nested(a=1),
        ),
        fallback_env=Fallback(),
    )


def test_fallback_value_propagation_in_nested_settings():
    @attr.s()
    class ConnectorA(SettingsBase):
        host: str = s_attrib("HOST")
        password: str = s_attrib("PASSWORD")

    @attr.s()
    class ConnBKey(SettingsBase):
        id: Optional[str] = s_attrib("ID")
        secret: str = s_attrib("SECRET")

    @attr.s()
    class ConnectorB(SettingsBase):
        url: str = s_attrib("URL")
        key: ConnBKey = s_attrib("KEY")

    @attr.s()
    class ConnectorsSettings(SettingsBase):
        conn_a: ConnectorA = s_attrib("A")
        conn_b: ConnectorB = s_attrib("B")

    @attr.s()
    class Settings(SettingsBase):
        conns: ConnectorsSettings = s_attrib(
            "CS",
            fallback_factory=lambda: ConnectorsSettings(
                conn_a=ConnectorA(
                    host="127.0.0.1",
                    password=required(str),
                ),
                conn_b=ConnectorB(
                    url="https://127.0.0.1:8080",
                    key=ConnBKey(id=None, secret=required(str)),
                ),
            ),
        )

    # Check required only
    perform_loader_env_check(
        dict(
            CS_A_PASSWORD="pass",
            CS_B_KEY_SECRET="key",
        ),
        Settings(
            conns=ConnectorsSettings(
                conn_a=ConnectorA(
                    host="127.0.0.1",
                    password="pass",
                ),
                conn_b=ConnectorB(
                    url="https://127.0.0.1:8080",
                    key=ConnBKey(
                        id=None,
                        secret="key",
                    ),
                ),
            )
        ),
    )

    # Check that deep overrides are applied
    perform_loader_env_check(
        dict(
            CS_A_PASSWORD="pass",
            CS_B_KEY_SECRET="key_secret",
            CS_B_KEY_ID="key_id_from_env",
        ),
        Settings(
            conns=ConnectorsSettings(
                conn_a=ConnectorA(
                    host="127.0.0.1",
                    password="pass",
                ),
                conn_b=ConnectorB(
                    url="https://127.0.0.1:8080",
                    key=ConnBKey(
                        id="key_id_from_env",
                        secret="key_secret",
                    ),
                ),
            )
        ),
    )

    # Check that exception fires on missing keys
    with pytest.raises(ConfigFieldMissing) as exc:
        load_settings(
            env=dict(),
            settings_type=Settings,
        )

    assert exc.value.field_set == {"conns.conn_b.key.secret", "conns.conn_a.password"}


def test_complex_nested_with_required_dft_ffactory_no_env():
    @attr.s()
    class Nested(SettingsBase):
        val: str = s_attrib("VAL")

    @attr.s()
    class RootWithFallbackFactoryNone(SettingsBase):
        nested_with_ff: Optional[Nested] = s_attrib("NWFF", fallback_factory=lambda: None)

    perform_loader_env_check(
        dict(),
        RootWithFallbackFactoryNone(nested_with_ff=None),
    )


def test_complex_nested_with_required_dft_missing_no_env():
    @attr.s()
    class Nested(SettingsBase):
        val: str = s_attrib("VAL")

    @attr.s()
    class RootWithMissingNone(SettingsBase):
        nested_with_ff: Optional[Nested] = s_attrib("NWFF", missing=None)

    perform_loader_env_check(
        dict(),
        RootWithMissingNone(nested_with_ff=None),
    )


def test_complex_nested_with_required_dft_not_defined_no_env():
    @attr.s()
    class Nested(SettingsBase):
        val: str = s_attrib("VAL")

    @attr.s()
    class RootWithMissingNone(SettingsBase):
        nested_with_ff: Optional[Nested] = s_attrib("NWFF")

    with pytest.raises(ConfigFieldMissing) as exc:
        load_settings(
            env=dict(),
            settings_type=RootWithMissingNone,
        )

    # There is keys for any of `nested_with_ff` attributes so it's treated as missing object, not particular fields
    assert exc.value.field_set == {"nested_with_ff"}


def test_complex_double_nested_with_required_dft_not_defined_partial_env():
    @attr.s()
    class Pair(SettingsBase):
        left: str = s_attrib("L")
        right: str = s_attrib("R")
        comment: str = s_attrib("C")

    @attr.s()
    class Nested(SettingsBase):
        val: Pair = s_attrib("PAIR")

    @attr.s()
    class RootWithMissingNone(SettingsBase):
        nested_with_ff: Optional[Nested] = s_attrib("NWFF")

    with pytest.raises(ConfigFieldMissing) as exc:
        load_settings(
            env=dict(NWFF_PAIR_C="some comment"),
            settings_type=RootWithMissingNone,
        )

    assert exc.value.field_set == {"nested_with_ff.val.right", "nested_with_ff.val.left"}


def test_complex_double_nested_dft_fallback_factory_nested_none():
    @attr.s()
    class Pair(SettingsBase):
        left: str = s_attrib("L")
        right: str = s_attrib("R")
        comment: str = s_attrib("C")

    @attr.s()
    class Nested(SettingsBase):
        val: Optional[Pair] = s_attrib("PAIR")

    @attr.s()
    class RootWithMissingNone(SettingsBase):
        nested_with_ff: Optional[Nested] = s_attrib("NWFF", fallback_factory=lambda: Nested(val=None))

    # Check regular loading
    perform_loader_env_check(
        dict(
            NWFF_PAIR_L="L",
            NWFF_PAIR_R="R",
            NWFF_PAIR_C="1",
        ),
        RootWithMissingNone(
            nested_with_ff=Nested(
                val=Pair(
                    left="L",
                    right="R",
                    comment="1",
                )
            )
        ),
    )

    # Check that exception is correct in case of missing vars
    with pytest.raises(ConfigFieldMissing) as exc:
        load_settings(
            env=dict(
                NWFF_PAIR_C="1",
            ),
            settings_type=RootWithMissingNone,
        )

    assert exc.value.field_set == {"nested_with_ff.val.right", "nested_with_ff.val.left"}


def test_json_value():
    @attr.s
    class Nested(SettingsBase):
        dft_key: str = s_attrib("DFT_KEY")
        map_key_id_value: Dict[str, str] = s_attrib("MAP")

    @attr.s
    class TypedSettings(SettingsBase):
        nested: Optional[Nested] = s_attrib(
            "N",
            json_converter=lambda js: Nested(dft_key=js["dft_key"], map_key_id_value=js["keys"]),
            enabled_key_name="ENABLED",
            missing=None,
        )

    expected_dft_key = "asf"
    expected_key_map = dict(a="a", b="b")
    expected_settings = TypedSettings(
        nested=Nested(dft_key=expected_dft_key, map_key_id_value=dict(expected_key_map)),
    )

    json_encoded_nested_value = json.dumps(dict(dft_key=expected_dft_key, keys=expected_key_map))

    # Case None
    perform_loader_env_check(
        dict(),
        TypedSettings(nested=None),
    )

    # Case of JSON value
    perform_loader_env_check(
        dict(
            N_JSON_VALUE=json_encoded_nested_value,
        ),
        expected_settings,
    )

    # Case disabled
    perform_loader_env_check(
        dict(
            N_JSON_VALUE=json_encoded_nested_value,
            N_ENABLED="false",
        ),
        TypedSettings(nested=None),
    )

    # Case disabled
    perform_loader_env_check(
        dict(
            N_JSON_VALUE=json_encoded_nested_value,
            N_ENABLED="true",
        ),
        expected_settings,
    )

    # Case of plain vars
    perform_loader_env_check(
        dict(
            N_DFT_KEY="asf",
            **{f"N_MAP_{key}": value for key, value in expected_key_map.items()},
        ),
        expected_settings,
    )


@pytest.fixture(scope="function")
def temp_file_factory() -> Callable[[str], str]:
    with tempfile.TemporaryDirectory() as tmp_dir_name:

        def create_file_with_content(content: str) -> str:
            file_name = path.join(tmp_dir_name, shortuuid.uuid())
            with open(file_name, "wb") as file:
                file.write(content.encode("ascii"))
            return file_name

        yield create_file_with_content


def test_file_mapping(temp_file_factory):
    @attr.s
    class TypedSettings(SettingsBase):
        a: str = s_attrib("A")
        b: int = s_attrib("B")

    a_value = "the_a_value"
    b_value = 434

    a_file_name = temp_file_factory(a_value)
    b_file_name = temp_file_factory(str(b_value))

    perform_loader_env_check(
        dict(
            BIE_FILE_MAP_A=a_file_name,
            BIE_FILE_MAP_B=b_file_name,
        ),
        TypedSettings(
            a=a_value,
            b=b_value,
        ),
    )

lib/dl_configs/dl_configs_tests/unit/test_settings_loader_env.py
datalens-tech/datalens-backend
"""
Here we test old autonomous loaders of sub-configs (which actually use new declarative framework).
In future all loaders which are tested here must be removed.
"""
from dl_configs.rqe import (
    RQEConfig,
    rqe_config_from_env,
)


def test_rqe_config():
    hmac_key = "sfgasdfdamjowerre"

    actual_config = rqe_config_from_env(
        {
            # Passing legacy env key to also check env remap
            "EXT_QUERY_EXECUTER_SECRET_KEY": hmac_key
        }
    )

    assert actual_config == RQEConfig.get_default().clone(hmac_key=hmac_key.encode())

lib/dl_configs/dl_configs_tests/unit/test_settings_loader_env_legacy.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.connection_base import ConnectionSchema
from dl_api_connector.api_schema.connection_base_fields import (
    cache_ttl_field,
    secret_string_field,
)
from dl_api_connector.api_schema.connection_mixins import RawSQLLevelMixin
from dl_api_connector.api_schema.extras import FieldExtra

from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery


class BigQueryConnectionSchema(ConnectionSchema, RawSQLLevelMixin):
    TARGET_CLS = ConnectionSQLBigQuery

    project_id = ma_fields.String(attribute="data.project_id", required=True, bi_extra=FieldExtra(editable=True))
    credentials = secret_string_field(attribute="data.credentials")
    cache_ttl_sec = cache_ttl_field(attribute="data.cache_ttl_sec")

lib/dl_connector_bigquery/dl_connector_bigquery/api/api_schema/connection.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.source_base import (
    SimpleDataSourceSchema,
    SimpleDataSourceTemplateSchema,
)
from dl_model_tools.schema.base import BaseSchema


class BigQueryTableParametersSchema(BaseSchema):
    dataset_name = ma_fields.String(allow_none=True)
    table_name = ma_fields.String(allow_none=True)
    db_version = ma_fields.String(allow_none=True)  # FIXME


class BigQueryTableDataSourceSchema(SimpleDataSourceSchema):
    parameters = ma_fields.Nested(BigQueryTableParametersSchema)


class BigQueryTableDataSourceTemplateSchema(SimpleDataSourceTemplateSchema):
    parameters = ma_fields.Nested(BigQueryTableParametersSchema)

lib/dl_connector_bigquery/dl_connector_bigquery/api/api_schema/source.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import unique
from typing import Optional

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormActionApiSchema,
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
)
from dl_api_connector.form_config.models.common import (
    CommonFieldName,
    FormFieldName,
)
import dl_api_connector.form_config.models.rows as C
from dl_api_connector.form_config.models.shortcuts.rows import RowConstructor
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_constants.enums import RawSQLLevel

from dl_connector_bigquery.api.connection_info import BigQueryConnectionInfoProvider
from dl_connector_bigquery.api.i18n.localizer import Translatable


@unique
class BigQueryFieldName(FormFieldName):
    project_id = "project_id"
    credentials = "credentials"


class BigQueryConnectionFormFactory(ConnectionFormFactory):
    def get_form_config(
        self,
        connector_settings: Optional[ConnectorSettingsBase],
        tenant: Optional[TenantDef],
    ) -> ConnectionForm:
        rc = RowConstructor(localizer=self._localizer)

        common_api_schema_items: list[FormFieldApiSchema] = [
            FormFieldApiSchema(name=BigQueryFieldName.project_id, required=True),
            FormFieldApiSchema(name=BigQueryFieldName.credentials, required=self.mode == ConnectionFormMode.create),
        ]

        edit_api_schema = FormActionApiSchema(
            items=[
                *common_api_schema_items,
                FormFieldApiSchema(name=CommonFieldName.cache_ttl_sec, nullable=True),
                FormFieldApiSchema(name=CommonFieldName.raw_sql_level),
            ]
        )

        create_api_schema = FormActionApiSchema(
            items=[
                *edit_api_schema.items,
                *self._get_top_level_create_api_schema_items(),
            ]
        )

        check_api_schema = FormActionApiSchema(
            items=[
                *common_api_schema_items,
                *self._get_top_level_check_api_schema_items(),
            ]
        )

        return ConnectionForm(
            title=BigQueryConnectionInfoProvider.get_title(self._localizer),
            rows=[
                C.CustomizableRow(
                    items=[
                        C.LabelRowItem(text=self._localizer.translate(Translatable("label_project-id"))),
                        C.InputRowItem(name=BigQueryFieldName.project_id, width="m"),
                    ]
                ),
                C.CustomizableRow(
                    items=[
                        C.LabelRowItem(text=self._localizer.translate(Translatable("label_service-acc-key-file"))),
                        C.FileInputRowItem(name=BigQueryFieldName.credentials),
                    ]
                ),
                rc.raw_sql_level_row(
                    options=[
                        rc.raw_sql_level_to_radio_group_option(RawSQLLevel.off),
                        rc.raw_sql_level_to_radio_group_option(RawSQLLevel.subselect),
                    ],
                ),
                C.CacheTTLRow(name=CommonFieldName.cache_ttl_sec),
            ],
            api_schema=FormApiSchema(
                create=create_api_schema if self.mode == ConnectionFormMode.create else None,
                edit=edit_api_schema if self.mode == ConnectionFormMode.edit else None,
                check=check_api_schema,
            ),
        )

lib/dl_connector_bigquery/dl_connector_bigquery/api/connection_form/form_config.py
datalens-tech/datalens-backend
from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_bigquery.api.i18n.localizer import Translatable


class BigQueryConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-bigquery")

lib/dl_connector_bigquery/dl_connector_bigquery/api/connection_info.py
datalens-tech/datalens-backend
from dl_api_connector.api_schema.source_base import (
    SubselectDataSourceSchema,
    SubselectDataSourceTemplateSchema,
)
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)

from dl_connector_bigquery.api.api_schema.connection import BigQueryConnectionSchema
from dl_connector_bigquery.api.api_schema.source import (
    BigQueryTableDataSourceSchema,
    BigQueryTableDataSourceTemplateSchema,
)
from dl_connector_bigquery.api.connection_form.form_config import BigQueryConnectionFormFactory
from dl_connector_bigquery.api.connection_info import BigQueryConnectionInfoProvider
from dl_connector_bigquery.api.i18n.localizer import CONFIGS
from dl_connector_bigquery.core.connector import (
    BigQueryCoreConnectionDefinition,
    BigQueryCoreConnector,
    BigQueryCoreSubselectSourceDefinition,
    BigQueryCoreTableSourceDefinition,
)
from dl_connector_bigquery.formula.constants import DIALECT_NAME_BIGQUERY


class BigQueryApiTableSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = BigQueryCoreTableSourceDefinition
    api_schema_cls = BigQueryTableDataSourceSchema
    template_api_schema_cls = BigQueryTableDataSourceTemplateSchema


class BigQueryApiSubselectSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = BigQueryCoreSubselectSourceDefinition
    api_schema_cls = SubselectDataSourceSchema
    template_api_schema_cls = SubselectDataSourceTemplateSchema


class BigQueryApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = BigQueryCoreConnectionDefinition
    api_generic_schema_cls = BigQueryConnectionSchema
    info_provider_cls = BigQueryConnectionInfoProvider
    form_factory_cls = BigQueryConnectionFormFactory


class BigQueryApiConnector(ApiConnector):
    core_connector_cls = BigQueryCoreConnector
    source_definitions = (
        BigQueryApiTableSourceDefinition,
        BigQueryApiSubselectSourceDefinition,
    )
    connection_definitions = (BigQueryApiConnectionDefinition,)
    formula_dialect_name = DIALECT_NAME_BIGQUERY
    translation_configs = frozenset(CONFIGS)

lib/dl_connector_bigquery/dl_connector_bigquery/api/connector.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_bigquery as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_bigquery/dl_connector_bigquery/api/i18n/localizer.py
datalens-tech/datalens-backend
from __future__ import annotations

import base64
import json
from typing import (
    Optional,
    Tuple,
)

from google.auth.credentials import Credentials as BQCredentials
from google.cloud.bigquery import Client as BQClient
import google.oauth2.service_account as g_service_account
import sqlalchemy as sa
import sqlalchemy_bigquery._types as bq_types

from dl_core.connection_executors.adapters.adapters_base_sa_classic import (
    BaseClassicAdapter,
    BaseConnLineConstructor,
)
from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
from dl_core.connection_models.common_models import (
    DBIdent,
    SchemaIdent,
    TableIdent,
)

from dl_connector_bigquery.core.constants import CONNECTION_TYPE_BIGQUERY
from dl_connector_bigquery.core.error_transformer import big_query_db_error_transformer
from dl_connector_bigquery.core.target_dto import BigQueryConnTargetDTO


class BigQueryConnLineConstructor(BaseConnLineConstructor[BigQueryConnTargetDTO]):
    def _get_dsn_params(
        self,
        safe_db_symbols: Tuple[str, ...] = (),
        db_name: Optional[str] = None,
        standard_auth: Optional[bool] = True,
    ) -> dict:
        return {
            "dialect": self._dialect_name,
            "project_id": self._target_dto.project_id,
        }


class BigQueryDefaultAdapter(BaseClassicAdapter[BigQueryConnTargetDTO]):
    conn_type = CONNECTION_TYPE_BIGQUERY
    dsn_template = "{dialect}://{project_id}"
    conn_line_constructor_type = BigQueryConnLineConstructor
    _error_transformer = big_query_db_error_transformer

    _type_code_to_sa = {
        "STRING": bq_types.STRING,
        "DATE": bq_types.DATE,
        "DATETIME": bq_types.DATETIME,
        "BOOL": bq_types.BOOL,
        "BOOLEAN": bq_types.BOOLEAN,
        "INTEGER": bq_types.INTEGER,
        "INT64": bq_types.INT64,
        "FLOAT": bq_types.FLOAT,
        "FLOAT64": bq_types.FLOAT64,
        "NUMERIC": bq_types.NUMERIC,
        "BIGNUMERIC": bq_types.BIGNUMERIC,
        # TODO: ARRAY
    }

    def get_default_db_name(self) -> Optional[str]:
        return self._target_dto.project_id

    def get_engine_kwargs(self) -> dict:
        return {
            "credentials_base64": self._target_dto.credentials,
        }

    def _get_bq_credentials(self) -> BQCredentials:
        credentials_info = json.loads(base64.b64decode(self._target_dto.credentials))
        credentials = g_service_account.Credentials.from_service_account_info(credentials_info)
        return credentials

    def get_bq_client(self) -> BQClient:
        client = BQClient(credentials=self._get_bq_credentials())
        return client

    def _get_tables(self, schema_ident: SchemaIdent) -> list[TableIdent]:
        client = self.get_bq_client()
        bq_datasets = list(client.list_datasets())
        project_id = self._target_dto.project_id
        db_engine = self.get_db_engine(db_name=project_id)
        quoter = db_engine.dialect.identifier_preparer.quote

        subqueries = [
            sa.select(
                [
                    sa.literal_column("project_id"),
                    sa.literal_column("dataset_id"),
                    sa.literal_column("table_id"),
                ]
            ).select_from(
                sa.text(
                    quoter(
                        "{project_id}.{dataset_id}.__TABLES__".format(
                            project_id=project_id, dataset_id=dataset_info.dataset_id
                        )
                    )
                )
            )
            for dataset_info in bq_datasets
        ]
        query = DBAdapterQuery(query=sa.union_all(*subqueries))
        result = self.execute(query)
        return [
            TableIdent(
                db_name=_project_id,
                schema_name=dataset_id,
                table_name=table_name,
            )
            for _project_id, dataset_id, table_name in result.get_all()
        ]

    def _get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        return None

lib/dl_connector_bigquery/dl_connector_bigquery/core/adapters.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import List

import attr

from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor

from dl_connector_bigquery.core.adapters import BigQueryDefaultAdapter
from dl_connector_bigquery.core.dto import BigQueryConnDTO
from dl_connector_bigquery.core.target_dto import BigQueryConnTargetDTO


@attr.s(cmp=False, hash=False)
class BigQueryAsyncConnExecutor(DefaultSqlAlchemyConnExecutor[BigQueryDefaultAdapter]):
    TARGET_ADAPTER_CLS = BigQueryDefaultAdapter

    _conn_dto: BigQueryConnDTO = attr.ib()

    async def _make_target_conn_dto_pool(self) -> List[BigQueryConnTargetDTO]:
        dto_pool: list[BigQueryConnTargetDTO] = []
        dto_pool.append(
            BigQueryConnTargetDTO(
                conn_id=self._conn_dto.conn_id,
                pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                project_id=self._conn_dto.project_id,
                credentials=self._conn_dto.credentials,
            )
        )
        return dto_pool

lib/dl_connector_bigquery/dl_connector_bigquery/core/connection_executors.py
datalens-tech/datalens-backend
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
    CoreSourceDefinition,
)

from dl_connector_bigquery.core.adapters import BigQueryDefaultAdapter
from dl_connector_bigquery.core.connection_executors import BigQueryAsyncConnExecutor
from dl_connector_bigquery.core.constants import (
    BACKEND_TYPE_BIGQUERY,
    CONNECTION_TYPE_BIGQUERY,
    SOURCE_TYPE_BIGQUERY_SUBSELECT,
    SOURCE_TYPE_BIGQUERY_TABLE,
)
from dl_connector_bigquery.core.data_source import (
    BigQuerySubselectDataSource,
    BigQueryTableDataSource,
)
from dl_connector_bigquery.core.data_source_spec import (
    BigQuerySubselectDataSourceSpec,
    BigQueryTableDataSourceSpec,
)
from dl_connector_bigquery.core.query_compiler import BigQueryQueryCompiler
from dl_connector_bigquery.core.sa_types import SQLALCHEMY_BIGQUERY_TYPES
from dl_connector_bigquery.core.storage_schemas.connection import BigQueryConnectionDataStorageSchema
from dl_connector_bigquery.core.storage_schemas.data_source_spec import (
    BigQuerySubselectDataSourceSpecStorageSchema,
    BigQueryTableDataSourceSpecStorageSchema,
)
from dl_connector_bigquery.core.type_transformer import BigQueryTypeTransformer
from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery


class BigQueryCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_BIGQUERY
    connection_cls = ConnectionSQLBigQuery
    us_storage_schema_cls = BigQueryConnectionDataStorageSchema
    type_transformer_cls = BigQueryTypeTransformer
    sync_conn_executor_cls = BigQueryAsyncConnExecutor
    async_conn_executor_cls = BigQueryAsyncConnExecutor
    dialect_string = "bigquery"


class BigQueryCoreTableSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_BIGQUERY_TABLE
    source_cls = BigQueryTableDataSource
    source_spec_cls = BigQueryTableDataSourceSpec
    us_storage_schema_cls = BigQueryTableDataSourceSpecStorageSchema


class BigQueryCoreSubselectSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_BIGQUERY_SUBSELECT
    source_cls = BigQuerySubselectDataSource
    source_spec_cls = BigQuerySubselectDataSourceSpec
    us_storage_schema_cls = BigQuerySubselectDataSourceSpecStorageSchema


class BigQueryCoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_BIGQUERY
    compiler_cls = BigQueryQueryCompiler
    connection_definitions = (BigQueryCoreConnectionDefinition,)
    source_definitions = (
        BigQueryCoreTableSourceDefinition,
        BigQueryCoreSubselectSourceDefinition,
    )
    rqe_adapter_classes = frozenset({BigQueryDefaultAdapter})
    sa_types = SQLALCHEMY_BIGQUERY_TYPES

lib/dl_connector_bigquery/dl_connector_bigquery/core/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
    SourceBackendType,
)


BACKEND_TYPE_BIGQUERY = SourceBackendType.declare("BIGQUERY")

CONNECTION_TYPE_BIGQUERY = ConnectionType.declare("bigquery")

SOURCE_TYPE_BIGQUERY_TABLE = DataSourceType.declare("BIGQUERY_TABLE")
SOURCE_TYPE_BIGQUERY_SUBSELECT = DataSourceType.declare("BIGQUERY_SUBSELECT")

lib/dl_connector_bigquery/dl_connector_bigquery/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Optional,
)

from dl_constants.enums import DataSourceType
from dl_core.connection_models import (
    TableDefinition,
    TableIdent,
)
from dl_core.data_source.sql import (
    BaseSQLDataSource,
    SQLDataSource,
    SubselectDataSource,
    TableSQLDataSourceMixin,
    require_table_name,
)
from dl_core.utils import sa_plain_text

from dl_connector_bigquery.core.constants import (
    CONNECTION_TYPE_BIGQUERY,
    SOURCE_TYPE_BIGQUERY_SUBSELECT,
    SOURCE_TYPE_BIGQUERY_TABLE,
)
from dl_connector_bigquery.core.data_source_spec import (
    BigQuerySubselectDataSourceSpec,
    BigQueryTableDataSourceSpec,
)
from dl_connector_bigquery.core.query_compiler import BigQueryQueryCompiler
from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery


class BigQueryDataSourceMixin(BaseSQLDataSource):
    conn_type = CONNECTION_TYPE_BIGQUERY

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type in (SOURCE_TYPE_BIGQUERY_TABLE, SOURCE_TYPE_BIGQUERY_SUBSELECT)


class BigQueryTableDataSource(BigQueryDataSourceMixin, TableSQLDataSourceMixin, SQLDataSource):
    """
    BigQuery table
    """

    compiler_cls = BigQueryQueryCompiler

    @property
    def spec(self) -> BigQueryTableDataSourceSpec:
        assert isinstance(self._spec, BigQueryTableDataSourceSpec)
        return self._spec

    @property
    def connection(self) -> ConnectionSQLBigQuery:
        connection = super().connection
        assert isinstance(connection, ConnectionSQLBigQuery)
        return connection

    @property
    def db_name(self) -> Optional[str]:
        return self.connection.project_id

    def get_table_definition(self) -> TableDefinition:
        assert self.table_name is not None
        return TableIdent(
            db_name=self.db_name,
            schema_name=self.spec.dataset_name,
            table_name=self.table_name,
        )

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            dataset_name=self.spec.dataset_name,
            table_name=self.spec.table_name,
        )

    @require_table_name
    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        q = self.quote
        alias_str = "" if alias is None else f" AS {q(alias)}"
        return sa_plain_text(f"{q(self.db_name)}" f".{q(self.spec.dataset_name)}" f".{q(self.table_name)}{alias_str}")


class BigQuerySubselectDataSource(BigQueryDataSourceMixin, SubselectDataSource):
    """
    BigQuery subselect
    """

    @property
    def spec(self) -> BigQuerySubselectDataSourceSpec:
        assert isinstance(self._spec, BigQuerySubselectDataSourceSpec)
        return self._spec

lib/dl_connector_bigquery/dl_connector_bigquery/core/data_source.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

import attr

from dl_core.data_source_spec.sql import (
    SQLDataSourceSpecBase,
    SubselectDataSourceSpec,
    TableSQLDataSourceSpec,
)


@attr.s
class BigQueryDataSourceSpecMixin(SQLDataSourceSpecBase):
    pass


@attr.s
class BigQueryTableDataSourceSpec(BigQueryDataSourceSpecMixin, TableSQLDataSourceSpec):
    dataset_name: str = attr.ib(kw_only=True)


@attr.s
class BigQuerySubselectDataSourceSpec(BigQueryDataSourceSpecMixin, SubselectDataSourceSpec):
    subsql: Optional[str] = attr.ib(kw_only=True, default=None)

lib/dl_connector_bigquery/dl_connector_bigquery/core/data_source_spec.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_core.connection_models.dto_defs import ConnDTO
from dl_core.utils import secrepr

from dl_connector_bigquery.core.constants import CONNECTION_TYPE_BIGQUERY


@attr.s(frozen=True)
class BigQueryConnDTO(ConnDTO):
    conn_type = CONNECTION_TYPE_BIGQUERY

    credentials: str = attr.ib(repr=secrepr)  # corresponds to `credentials_base64` in BQ SQLA
    project_id: str = attr.ib()

lib/dl_connector_bigquery/dl_connector_bigquery/core/dto.py
datalens-tech/datalens-backend
import google.api_core.exceptions as ga_exc

from dl_core.connectors.base.error_transformer import (
    ChainedDbErrorTransformer,
    DbErrorTransformer,
)
from dl_core.connectors.base.error_transformer import ErrorTransformerRule as Rule
from dl_core.connectors.base.error_transformer import wrapper_exc_is
import dl_core.exc as exc


big_query_db_error_transformer: DbErrorTransformer = ChainedDbErrorTransformer(
    [Rule(when=wrapper_exc_is(wrapper_exc_cls=ga_exc.Forbidden), then_raise=exc.UserQueryAccessDenied)]
)

lib/dl_connector_bigquery/dl_connector_bigquery/core/error_transformer.py
datalens-tech/datalens-backend
from typing import ClassVar

from dl_core.connectors.base.query_compiler import (
    QueryCompiler,
    SectionAliasMode,
)


class BigQueryQueryCompiler(QueryCompiler):
    groupby_alias_mode: ClassVar[SectionAliasMode] = SectionAliasMode.by_alias_in_section

lib/dl_connector_bigquery/dl_connector_bigquery/core/query_compiler.py
datalens-tech/datalens-backend
import sqlalchemy_bigquery._types as bq_types

from dl_core.db.sa_types_base import (
    make_native_type,
    simple_instantiator,
)

from dl_connector_bigquery.core.constants import CONNECTION_TYPE_BIGQUERY as CT_BQ


SQLALCHEMY_BIGQUERY_TYPES = {
    make_native_type(CT_BQ, bq_types.DATE): simple_instantiator(bq_types.DATE),
    make_native_type(CT_BQ, bq_types.DATETIME): simple_instantiator(bq_types.DATE),
    make_native_type(CT_BQ, bq_types.STRING): simple_instantiator(bq_types.STRING),
    make_native_type(CT_BQ, bq_types.BOOLEAN): simple_instantiator(bq_types.BOOLEAN),
    make_native_type(CT_BQ, bq_types.INTEGER): simple_instantiator(bq_types.INTEGER),
    make_native_type(CT_BQ, bq_types.FLOAT): simple_instantiator(bq_types.FLOAT),
    make_native_type(CT_BQ, bq_types.NUMERIC): simple_instantiator(bq_types.NUMERIC),
    make_native_type(CT_BQ, bq_types.ARRAY(bq_types.INTEGER)): simple_instantiator(bq_types.ARRAY(bq_types.INTEGER)),
    make_native_type(CT_BQ, bq_types.ARRAY(bq_types.STRING)): simple_instantiator(bq_types.ARRAY(bq_types.STRING)),
    make_native_type(CT_BQ, bq_types.ARRAY(bq_types.FLOAT)): simple_instantiator(bq_types.ARRAY(bq_types.FLOAT)),
}

lib/dl_connector_bigquery/dl_connector_bigquery/core/sa_types.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_core.us_manager.storage_schemas.connection import (
    BaseConnectionDataStorageSchema,
    CacheableConnectionDataSchemaMixin,
    SubselectConnectionDataSchemaMixin,
)

from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery


class BigQueryConnectionDataStorageSchema(
    BaseConnectionDataStorageSchema[ConnectionSQLBigQuery.DataModel],
    CacheableConnectionDataSchemaMixin,
    SubselectConnectionDataSchemaMixin,
):
    TARGET_CLS = ConnectionSQLBigQuery.DataModel

    project_id = ma_fields.String(allow_none=False, required=True)
    credentials = ma_fields.String(allow_none=False, required=True)

lib/dl_connector_bigquery/dl_connector_bigquery/core/storage_schemas/connection.py
datalens-tech/datalens-backend
import marshmallow.fields as ma_fields

from dl_core.us_manager.storage_schemas.data_source_spec_base import (
    BaseSQLDataSourceSpecStorageSchema,
    SubselectSQLDataSourceSpecStorageSchemaMixin,
    TableSQLDataSourceSpecStorageSchemaMixin,
)

from dl_connector_bigquery.core.data_source_spec import (
    BigQuerySubselectDataSourceSpec,
    BigQueryTableDataSourceSpec,
)


class DatasetSQLDataSourceSpecStorageSchemaMixin(BaseSQLDataSourceSpecStorageSchema):  # noqa
    dataset_name = ma_fields.String(required=True)


class BigQueryTableDataSourceSpecStorageSchema(
    TableSQLDataSourceSpecStorageSchemaMixin,
    DatasetSQLDataSourceSpecStorageSchemaMixin,
    BaseSQLDataSourceSpecStorageSchema,
):
    TARGET_CLS = BigQueryTableDataSourceSpec  # type: ignore


class BigQuerySubselectDataSourceSpecStorageSchema(
    SubselectSQLDataSourceSpecStorageSchemaMixin,
    BaseSQLDataSourceSpecStorageSchema,
):
    TARGET_CLS = BigQuerySubselectDataSourceSpec  # type: ignore

lib/dl_connector_bigquery/dl_connector_bigquery/core/storage_schemas/data_source_spec.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO
from dl_core.utils import secrepr


@attr.s(frozen=True, kw_only=True)
class BigQueryConnTargetDTO(ConnTargetDTO):
    credentials: str = attr.ib(repr=secrepr)  # corresponds to `credentials_base64` in BQ SQLA
    project_id: str = attr.ib()

    def get_effective_host(self) -> Optional[str]:
        return None

lib/dl_connector_bigquery/dl_connector_bigquery/core/target_dto.py
datalens-tech/datalens-backend
from typing import Any
import uuid

from dl_constants.enums import RawSQLLevel
from dl_core.us_manager.us_manager_sync import SyncUSManager

from dl_connector_bigquery.core.constants import CONNECTION_TYPE_BIGQUERY
from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery


def make_bigquery_saved_connection(
    sync_usm: SyncUSManager,
    project_id: str,
    credentials: str,
    raw_sql_level: RawSQLLevel = RawSQLLevel.off,
    **kwargs: Any,
) -> ConnectionSQLBigQuery:
    conn_name = "bigquery test conn {}".format(uuid.uuid4())
    conn = ConnectionSQLBigQuery.create_from_dict(
        data_dict=ConnectionSQLBigQuery.DataModel(
            project_id=project_id,
            credentials=credentials,
            raw_sql_level=raw_sql_level,
        ),
        ds_key=conn_name,
        type_=CONNECTION_TYPE_BIGQUERY.name,
        us_manager=sync_usm,
        **kwargs,
    )
    sync_usm.save(conn)
    return conn

lib/dl_connector_bigquery/dl_connector_bigquery/core/testing/connection.py
datalens-tech/datalens-backend
import sqlalchemy_bigquery._types as bq_types

from dl_constants.enums import UserDataType
from dl_core.db.conversion_base import (
    TypeTransformer,
    make_native_type,
)

from dl_connector_bigquery.core.constants import CONNECTION_TYPE_BIGQUERY


class BigQueryTypeTransformer(TypeTransformer):
    conn_type = CONNECTION_TYPE_BIGQUERY
    native_to_user_map = {
        make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.DATE): UserDataType.date,
        make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.DATETIME): UserDataType.genericdatetime,
        make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.STRING): UserDataType.string,
        make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.BOOLEAN): UserDataType.boolean,
        make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.INTEGER): UserDataType.integer,
        make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.FLOAT): UserDataType.float,
        make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.NUMERIC): UserDataType.float,
        make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.ARRAY(bq_types.INTEGER)): UserDataType.array_int,
        make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.ARRAY(bq_types.STRING)): UserDataType.array_str,
        make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.ARRAY(bq_types.FLOAT)): UserDataType.array_float,
    }
    user_to_native_map = {
        UserDataType.date: make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.DATE),
        UserDataType.genericdatetime: make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.DATETIME),
        UserDataType.string: make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.STRING),
        UserDataType.boolean: make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.BOOLEAN),
        UserDataType.integer: make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.INTEGER),
        UserDataType.float: make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.FLOAT),
        UserDataType.array_int: make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.ARRAY(bq_types.INTEGER)),
        UserDataType.array_str: make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.ARRAY(bq_types.STRING)),
        UserDataType.array_float: make_native_type(CONNECTION_TYPE_BIGQUERY, bq_types.ARRAY(bq_types.FLOAT)),
    }

lib/dl_connector_bigquery/dl_connector_bigquery/core/type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Callable,
    ClassVar,
)

import attr

from dl_core.base_models import (
    ConnCacheableDataModelMixin,
    ConnectionDataModelBase,
    ConnSubselectDataModelMixin,
)
from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.us_connection_base import (
    ConnectionBase,
    ConnectionSQL,
    DataSourceTemplate,
)
from dl_i18n.localizer_base import Localizer
from dl_utils.utils import DataKey

from dl_connector_bigquery.core.constants import (
    CONNECTION_TYPE_BIGQUERY,
    SOURCE_TYPE_BIGQUERY_SUBSELECT,
    SOURCE_TYPE_BIGQUERY_TABLE,
)
from dl_connector_bigquery.core.dto import BigQueryConnDTO


class ConnectionSQLBigQuery(ConnectionSQL):
    conn_type = CONNECTION_TYPE_BIGQUERY

    has_schema: ClassVar[bool] = True
    default_schema_name = None
    source_type = SOURCE_TYPE_BIGQUERY_TABLE
    allowed_source_types = frozenset((SOURCE_TYPE_BIGQUERY_TABLE, SOURCE_TYPE_BIGQUERY_SUBSELECT))
    allow_dashsql: ClassVar[bool] = False
    allow_cache: ClassVar[bool] = True
    is_always_user_source: ClassVar[bool] = True

    @attr.s
    class DataModel(ConnCacheableDataModelMixin, ConnSubselectDataModelMixin, ConnectionDataModelBase):
        credentials: str = attr.ib(kw_only=True)
        project_id: str = attr.ib(kw_only=True)
        data_export_forbidden: bool = attr.ib(default=False)

        @classmethod
        def get_secret_keys(cls) -> set[DataKey]:
            return {DataKey(parts=("credentials",))}

    @property
    def project_id(self) -> str:
        return self.data.project_id

    def get_conn_dto(self) -> BigQueryConnDTO:
        return BigQueryConnDTO(
            conn_id=self.uuid,
            project_id=self.data.project_id,
            credentials=self.data.credentials,
        )

    def get_data_source_template_group(self, parameters: dict) -> list[str]:
        return [val for val in (parameters.get("project_id"), parameters.get("dataset_name")) if val is not None]

    def get_parameter_combinations(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[dict]:
        tables = self.get_tables(conn_executor_factory)
        return [
            dict(
                project_id=table_ident.db_name,
                dataset_name=table_ident.schema_name,
                table_name=table_ident.table_name,
            )
            for table_ident in tables
        ]

    def get_data_source_template_templates(self, localizer: Localizer) -> list[DataSourceTemplate]:
        return self._make_subselect_templates(source_type=SOURCE_TYPE_BIGQUERY_SUBSELECT, localizer=localizer)

    @property
    def allow_public_usage(self) -> bool:
        return False

lib/dl_connector_bigquery/dl_connector_bigquery/core/us_connection.py
datalens-tech/datalens-backend
from dl_db_testing.connectors.base.connector import DbTestingConnector

from dl_connector_bigquery.db_testing.engine_wrapper import BigQueryEngineWrapper


class BigQueryDbTestingConnector(DbTestingConnector):
    engine_wrapper_classes = (BigQueryEngineWrapper,)

lib/dl_connector_bigquery/dl_connector_bigquery/db_testing/connector.py
datalens-tech/datalens-backend
from typing import (
    Optional,
    Sequence,
)

import attr
import sqlalchemy as sa

from dl_db_testing.database.engine_wrapper import (
    DbEngineConfig,
    EngineWrapperBase,
)


@attr.s(frozen=True)
class BigQueryDbEngineConfig(DbEngineConfig):
    default_dataset_name: str = attr.ib(kw_only=True)


class BigQueryEngineWrapper(EngineWrapperBase):
    URL_PREFIX = "bigquery"
    TABLE_AVAILABILITY_TIMEOUT = 30.0

    @property
    def config(self) -> BigQueryDbEngineConfig:
        assert isinstance(self._config, BigQueryDbEngineConfig)
        return self._config

    def get_conn_credentials(self, full: bool = False) -> dict:
        creds = dict(
            db_name=self.url.database,
        )
        return creds

    def table_from_columns(
        self,
        columns: Sequence[sa.Column],
        *,
        schema: Optional[str] = None,
        table_name: Optional[str] = None,
    ) -> sa.Table:
        table = super().table_from_columns(
            columns=columns,
            schema=schema or self.config.default_dataset_name,
            table_name=table_name,
        )
        return table

lib/dl_connector_bigquery/dl_connector_bigquery/db_testing/engine_wrapper.py
datalens-tech/datalens-backend
import sqlalchemy.sql.functions as sa_funcs
from sqlalchemy_bigquery.base import BigQueryDialect as SABigQueryDialect

from dl_formula.connectors.base.connector import FormulaConnector

from dl_connector_bigquery.formula.constants import BigQueryDialect as BigQueryDialectNS
from dl_connector_bigquery.formula.definitions.all import DEFINITIONS


class BigQueryFormulaConnector(FormulaConnector):
    dialect_ns_cls = BigQueryDialectNS
    dialects = BigQueryDialectNS.BIGQUERY
    default_dialect = BigQueryDialectNS.BIGQUERY
    op_definitions = DEFINITIONS
    sa_dialect = SABigQueryDialect()

    @classmethod
    def registration_hook(cls) -> None:
        # Unregister BigQuery's custom implementation of `unnest` because it breaks other connectors
        # https://github.com/googleapis/python-bigquery-sqlalchemy/issues/882
        del sa_funcs._registry["_default"]["unnest"]  # noqa

lib/dl_connector_bigquery/dl_connector_bigquery/formula/connector.py
datalens-tech/datalens-backend
from dl_formula.core.dialect import (
    DialectName,
    DialectNamespace,
    simple_combo,
)


DIALECT_NAME_BIGQUERY = DialectName.declare("BIGQUERY")


class BigQueryDialect(DialectNamespace):
    BIGQUERY = simple_combo(name=DIALECT_NAME_BIGQUERY)

lib/dl_connector_bigquery/dl_connector_bigquery/formula/constants.py
datalens-tech/datalens-backend
from dl_connector_bigquery.formula.definitions.conditional_blocks import DEFINITIONS_COND_BLOCKS
from dl_connector_bigquery.formula.definitions.functions_aggregation import DEFINITIONS_AGG
from dl_connector_bigquery.formula.definitions.functions_datetime import DEFINITIONS_DATETIME
from dl_connector_bigquery.formula.definitions.functions_logical import DEFINITIONS_LOGICAL
from dl_connector_bigquery.formula.definitions.functions_markup import DEFINITIONS_MARKUP
from dl_connector_bigquery.formula.definitions.functions_math import DEFINITIONS_MATH
from dl_connector_bigquery.formula.definitions.functions_string import DEFINITIONS_STRING
from dl_connector_bigquery.formula.definitions.functions_type import DEFINITIONS_TYPE
from dl_connector_bigquery.formula.definitions.operators_binary import DEFINITIONS_BINARY
from dl_connector_bigquery.formula.definitions.operators_ternary import DEFINITIONS_TERNARY
from dl_connector_bigquery.formula.definitions.operators_unary import DEFINITIONS_UNARY


DEFINITIONS = [
    *DEFINITIONS_COND_BLOCKS,
    *DEFINITIONS_AGG,
    *DEFINITIONS_DATETIME,
    *DEFINITIONS_LOGICAL,
    *DEFINITIONS_MARKUP,
    *DEFINITIONS_MATH,
    *DEFINITIONS_STRING,
    *DEFINITIONS_TYPE,
    *DEFINITIONS_UNARY,
    *DEFINITIONS_BINARY,
    *DEFINITIONS_TERNARY,
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/all.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.conditional_blocks as base

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


V = TranslationVariant.make


DEFINITIONS_COND_BLOCKS = [
    # _case_block_
    base.CaseBlock.for_dialect(D.BIGQUERY),
    # _if_block_
    base.IfBlock3(
        variants=[
            V(D.BIGQUERY, sa.func.IF),
        ]
    ),
    base.IfBlockMulti.for_dialect(D.BIGQUERY),
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/conditional_blocks.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_aggregation as base

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


V = TranslationVariant.make


DEFINITIONS_AGG = [
    # any
    base.AggAny(
        variants=[
            V(D.BIGQUERY, sa.func.ANY_VALUE),
        ]
    ),
    # avg
    base.AggAvgFromNumber.for_dialect(D.BIGQUERY),
    base.AggAvgFromDate.for_dialect(D.BIGQUERY),
    base.AggAvgFromDatetime.for_dialect(D.BIGQUERY),
    base.AggAvgFromDatetimeTZ.for_dialect(D.BIGQUERY),
    # avg_if
    base.AggAvgIf.for_dialect(D.BIGQUERY),
    # count
    base.AggCount0.for_dialect(D.BIGQUERY),
    base.AggCount1.for_dialect(D.BIGQUERY),
    # count_if
    base.AggCountIf(variants=[V(D.BIGQUERY, sa.func.COUNTIF)]),
    # countd
    base.AggCountd.for_dialect(D.BIGQUERY),
    # countd_approx
    base.AggCountdApprox(
        variants=[
            V(D.BIGQUERY, sa.func.APPROX_COUNT_DISTINCT),
        ]
    ),
    # countd_if
    base.AggCountdIf.for_dialect(D.BIGQUERY),
    # max
    base.AggMax.for_dialect(D.BIGQUERY),
    # min
    base.AggMin.for_dialect(D.BIGQUERY),
    # stdev
    base.AggStdev.for_dialect(D.BIGQUERY),
    # stdevp
    base.AggStdevp.for_dialect(D.BIGQUERY),
    # sum
    base.AggSum.for_dialect(D.BIGQUERY),
    # sum_if
    base.AggSumIf.for_dialect(D.BIGQUERY),
    # var
    base.AggVar.for_dialect(D.BIGQUERY),
    # varp
    base.AggVarp.for_dialect(D.BIGQUERY),
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/functions_aggregation.py
datalens-tech/datalens-backend
import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

import dl_formula.core.exc as exc
from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common import raw_sql
from dl_formula.definitions.common_datetime import datetime_interval
import dl_formula.definitions.functions_datetime as base
from dl_formula.definitions.literals import un_literal

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


V = TranslationVariant.make


def _bq_dynamic_interval(what: str, num: int) -> ClauseElement:
    what = what.lower()
    parts = [0, 0, 0, 0, 0, 0]
    part_names = ["year", "month", "day", "hour", "minute", "second"]
    if what not in part_names:
        raise exc.TranslationError(f"Invalid interval name: {what}")
    parts[part_names.index(what)] = num
    return sa.func.MAKE_INTERVAL(*parts)


DEFINITIONS_DATETIME = [
    # dateadd
    base.FuncDateadd1.for_dialect(D.BIGQUERY),
    base.FuncDateadd2Unit.for_dialect(D.BIGQUERY),
    base.FuncDateadd2Number.for_dialect(D.BIGQUERY),
    base.FuncDateadd3DateConstNum(
        variants=[
            V(
                D.BIGQUERY,
                lambda date, what, num: sa.func.DATE_ADD(
                    date, datetime_interval(un_literal(what), mult=un_literal(num))
                ),
            ),
        ]
    ),
    base.FuncDateadd3DateNonConstNum(
        variants=[
            V(D.BIGQUERY, lambda date, what, num: sa.func.DATE_ADD(date, _bq_dynamic_interval(un_literal(what), num))),
        ]
    ),
    base.FuncDateadd3DatetimeConstNum(
        variants=[
            V(
                D.BIGQUERY,
                lambda date, what, num: sa.func.DATE_ADD(
                    date, datetime_interval(un_literal(what), mult=un_literal(num))
                ),
            ),
        ]
    ),
    base.FuncDateadd3DatetimeNonConstNum(
        variants=[
            V(D.BIGQUERY, lambda date, what, num: sa.func.DATE_ADD(date, _bq_dynamic_interval(un_literal(what), num))),
        ]
    ),
    base.FuncDateadd3GenericDatetimeNonConstNum(
        variants=[
            V(D.BIGQUERY, lambda date, what, num: sa.func.DATE_ADD(date, _bq_dynamic_interval(un_literal(what), num))),
        ]
    ),
    # datepart
    base.FuncDatepart2.for_dialect(D.BIGQUERY),
    base.FuncDatepart3Const.for_dialect(D.BIGQUERY),
    base.FuncDatepart3NonConst.for_dialect(D.BIGQUERY),
    # datetrunc
    base.FuncDatetrunc2Date(
        variants=[
            V(
                D.BIGQUERY,
                lambda date, unit: (
                    sa.func.DATE_TRUNC(date, raw_sql(base.norm_datetrunc_unit(unit).replace("week", "week(monday)")))
                    if base.norm_datetrunc_unit(unit) in {"year", "quarter", "month", "week", "day"}
                    else date
                ),
            ),
        ]
    ),
    base.FuncDatetrunc2Datetime(
        variants=[
            V(
                D.BIGQUERY,
                lambda date, unit: sa.func.DATETIME_TRUNC(
                    date, raw_sql(base.norm_datetrunc_unit(unit).replace("week", "week(monday)"))
                ),
            ),
        ]
    ),
    # day
    base.FuncDay(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.EXTRACT("DAY", value)),
        ]
    ),
    # dayofweek
    base.FuncDayofweek1.for_dialect(D.BIGQUERY),
    base.FuncDayofweek2(
        variants=[
            V(
                D.BIGQUERY,
                lambda date, firstday: sa.func.MOD(sa.func.EXTRACT("DAYOFWEEK", date) - 1 - base.norm_fd(firstday), 7)
                + 1,
            ),
        ]
    ),
    # genericnow
    base.FuncGenericNow(
        variants=[
            V(D.BIGQUERY, sa.func.CURRENT_DATETIME),
        ]
    ),
    # hour
    base.FuncHourDate.for_dialect(D.BIGQUERY),
    base.FuncHourDatetime(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.EXTRACT("HOUR", value)),
        ]
    ),
    # minute
    base.FuncMinuteDate.for_dialect(D.BIGQUERY),
    base.FuncMinuteDatetime(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.EXTRACT("MINUTE", value)),
        ]
    ),
    # month
    base.FuncMonth(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.EXTRACT("MONTH", value)),
        ]
    ),
    # now
    base.FuncNow(
        variants=[
            V(D.BIGQUERY, sa.func.CURRENT_DATETIME),
        ]
    ),
    # quarter
    base.FuncQuarter(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.EXTRACT("QUARTER", value)),
        ]
    ),
    # second
    base.FuncSecondDate.for_dialect(D.BIGQUERY),
    base.FuncSecondDatetime(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.EXTRACT("SECOND", value)),
        ]
    ),
    # today
    base.FuncToday(
        variants=[
            V(D.BIGQUERY, sa.func.CURRENT_DATE),
        ]
    ),
    # week
    base.FuncWeek(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.EXTRACT("ISOWEEK", value)),
        ]
    ),
    # year
    base.FuncYear(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.EXTRACT("YEAR", value)),
        ]
    ),
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/functions_datetime.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_logical as base

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


V = TranslationVariant.make


DEFINITIONS_LOGICAL = [
    # case
    base.FuncCase.for_dialect(D.BIGQUERY),
    # if
    base.FuncIf.for_dialect(D.BIGQUERY),
    # ifnull
    base.FuncIfnull(
        variants=[
            V(D.BIGQUERY, sa.func.IFNULL),
        ]
    ),
    # isnan
    base.FuncIsnan(
        variants=[
            V(D.BIGQUERY, sa.func.IS_NAN),
        ]
    ),
    # isnull
    base.FuncIsnull.for_dialect(D.BIGQUERY),
    # zn
    base.FuncZn(
        variants=[
            V(D.BIGQUERY, lambda x: sa.func.IFNULL(x, 0)),
        ]
    ),
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/functions_logical.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_markup as base

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


DEFINITIONS_MARKUP = [
    # +
    base.BinaryPlusMarkup.for_dialect(D.BIGQUERY),
    # __str
    base.FuncInternalStrConst.for_dialect(D.BIGQUERY),
    base.FuncInternalStr.for_dialect(D.BIGQUERY),
    # bold
    base.FuncBold.for_dialect(D.BIGQUERY),
    # italic
    base.FuncItalics.for_dialect(D.BIGQUERY),
    # markup
    base.ConcatMultiMarkup.for_dialect(D.BIGQUERY),
    # url
    base.FuncUrl.for_dialect(D.BIGQUERY),
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/functions_markup.py
datalens-tech/datalens-backend
import math

import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_math as base

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


V = TranslationVariant.make

DEFINITIONS_MATH = [
    # abs
    base.FuncAbs.for_dialect(D.BIGQUERY),
    # acos
    base.FuncAcos.for_dialect(D.BIGQUERY),
    # asin
    base.FuncAsin.for_dialect(D.BIGQUERY),
    # atan
    base.FuncAtan.for_dialect(D.BIGQUERY),
    # atan2
    base.FuncAtan2.for_dialect(D.BIGQUERY),
    # ceiling
    base.FuncCeiling.for_dialect(D.BIGQUERY),
    # cos
    base.FuncCos.for_dialect(D.BIGQUERY),
    # cot
    base.FuncCot.for_dialect(D.BIGQUERY),
    # degrees
    base.FuncDegrees(
        variants=[
            V(D.BIGQUERY, lambda x: x / math.pi * 180),
        ]
    ),
    # div
    base.FuncDivBasic.for_dialect(D.BIGQUERY),
    # div_safe
    base.FuncDivSafe2.for_dialect(D.BIGQUERY),
    base.FuncDivSafe3.for_dialect(D.BIGQUERY),
    # exp
    base.FuncExp.for_dialect(D.BIGQUERY),
    # floor
    base.FuncFloor.for_dialect(D.BIGQUERY),
    # greatest
    base.FuncGreatest1.for_dialect(D.BIGQUERY),
    base.FuncGreatestMain.for_dialect(D.BIGQUERY),
    # base.GreatestMulti.for_dialect(D.BIGQUERY),  # FIXME
    base.GreatestMulti(
        variants=[
            V(D.BIGQUERY, sa.func.GREATEST),
        ]
    ),
    # least
    base.FuncLeast1.for_dialect(D.BIGQUERY),
    base.FuncLeastMain.for_dialect(D.BIGQUERY),
    # base.LeastMulti.for_dialect(D.BIGQUERY),  # FIXME
    base.LeastMulti(
        variants=[
            V(D.BIGQUERY, sa.func.LEAST),
        ]
    ),
    # ln
    base.FuncLn.for_dialect(D.BIGQUERY),
    # log
    base.FuncLog(
        variants=[
            V(D.BIGQUERY, lambda x, base: sa.func.LOG(x, base)),
        ]
    ),
    # log10
    base.FuncLog10.for_dialect(D.BIGQUERY),
    # pi
    base.FuncPi(
        variants=[
            V(D.BIGQUERY, lambda: sa.literal(math.pi)),
        ]
    ),
    # power
    base.FuncPower(
        variants=[
            V(D.BIGQUERY, sa.func.POWER),
        ]
    ),
    # radians
    base.FuncRadians(
        variants=[
            V(D.BIGQUERY, lambda x: x * math.pi / 180),
        ]
    ),
    # round
    base.FuncRound1.for_dialect(D.BIGQUERY),
    base.FuncRound2.for_dialect(D.BIGQUERY),
    # sign
    base.FuncSign.for_dialect(D.BIGQUERY),
    # sin
    base.FuncSin.for_dialect(D.BIGQUERY),
    # sqrt
    base.FuncSqrt.for_dialect(D.BIGQUERY),
    # square
    base.FuncSquare(
        variants=[
            V(D.BIGQUERY, lambda x: sa.func.POW(x, 2)),
        ]
    ),
    # tan
    base.FuncTan.for_dialect(D.BIGQUERY),
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/functions_math.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_string as base
from dl_formula.shortcuts import n

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


V = TranslationVariant.make


DEFINITIONS_STRING = [
    # ascii
    base.FuncAscii.for_dialect(D.BIGQUERY),
    # char
    base.FuncChar(
        variants=[
            V(D.BIGQUERY, sa.func.CHR),
        ]
    ),
    # concat
    base.Concat1.for_dialect(D.BIGQUERY),
    base.ConcatMultiStrConst.for_dialect(D.BIGQUERY),
    base.ConcatMultiStr.for_dialect(D.BIGQUERY),
    base.ConcatMultiAny.for_dialect(D.BIGQUERY),
    # contains
    base.FuncContainsNonConst(
        variants=[
            V(D.BIGQUERY, lambda x, y: sa.func.STRPOS(x, y) != 0),
        ]
    ),
    base.FuncContainsNonString.for_dialect(D.BIGQUERY),
    # notcontains
    base.FuncNotContainsNonConst.for_dialect(D.BIGQUERY),
    base.FuncNotContainsNonString.for_dialect(D.BIGQUERY),
    # endswith
    base.FuncEndswithNonConst(variants=[V(D.BIGQUERY, sa.func.ENDS_WITH)]),
    base.FuncEndswithNonString.for_dialect(D.BIGQUERY),
    # find
    base.FuncFind2(variants=[V(D.BIGQUERY, sa.func.STRPOS)]),
    base.FuncFind3(
        variants=[
            V(
                D.BIGQUERY,
                lambda x, y, z: n.if_(sa.func.STRPOS(sa.func.SUBSTRING(x, z), y) > 0)  # type: ignore
                .then(sa.func.STRPOS(sa.func.SUBSTRING(x, z), y) + z - 1)
                .else_(0),
            ),
        ]
    ),
    # icontains
    base.FuncIContainsNonConst.for_dialect(D.BIGQUERY),
    base.FuncIContainsNonString.for_dialect(D.BIGQUERY),
    # iendswith
    base.FuncIEndswithNonConst.for_dialect(D.BIGQUERY),
    base.FuncIEndswithNonString.for_dialect(D.BIGQUERY),
    # istartswith
    base.FuncIStartswithNonConst.for_dialect(D.BIGQUERY),
    base.FuncIStartswithNonString.for_dialect(D.BIGQUERY),
    # left
    base.FuncLeft.for_dialect(D.BIGQUERY),
    # len
    base.FuncLenString(variants=[V(D.BIGQUERY, sa.func.LENGTH)]),
    # lower
    base.FuncLowerConst.for_dialect(D.BIGQUERY),
    base.FuncLowerNonConst.for_dialect(D.BIGQUERY),
    # ltrim
    base.FuncLtrim(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.LTRIM(value, " ")),
        ]
    ),
    # regexp_extract
    base.FuncRegexpExtract(variants=[V(D.BIGQUERY, sa.func.REGEXP_EXTRACT)]),
    # regexp_extract_nth
    base.FuncRegexpExtractNth(
        variants=[
            V(
                D.BIGQUERY,
                lambda text, pattern, ind: sa.type_coerce(
                    sa.func.REGEXP_EXTRACT_ALL(text, pattern), sa.ARRAY(sa.String)
                )[ind - 1],
            ),
        ]
    ),
    # regexp_match
    base.FuncRegexpMatch(variants=[V(D.BIGQUERY, sa.func.REGEXP_CONTAINS)]),
    # regexp_replace
    base.FuncRegexpReplace(variants=[V(D.BIGQUERY, sa.func.REGEXP_REPLACE)]),
    # replace
    base.FuncReplace.for_dialect(D.BIGQUERY),
    # right
    base.FuncRight.for_dialect(D.BIGQUERY),
    # rtrim
    base.FuncRtrim(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.RTRIM(value, " ")),
        ]
    ),
    # space
    base.FuncSpaceConst.for_dialect(D.BIGQUERY),
    base.FuncSpaceNonConst(variants=[V(D.BIGQUERY, lambda value: sa.func.REPEAT(" ", value))]),
    # split
    # base.FuncSplit1,  # FIXME: Add array support
    # base.FuncSplit2,  # FIXME: Add array support
    base.FuncSplit3(
        variants=[
            V(
                D.BIGQUERY,
                lambda text, delim, ind: (sa.type_coerce(sa.func.SPLIT(text, delim), sa.ARRAY(sa.String))[ind - 1]),
            ),
        ]
    ),
    # startswith
    base.FuncStartswithNonConst(variants=[V(D.BIGQUERY, sa.func.STARTS_WITH)]),
    base.FuncStartswithNonString.for_dialect(D.BIGQUERY),
    # substr
    base.FuncSubstr2.for_dialect(D.BIGQUERY),
    base.FuncSubstr3.for_dialect(D.BIGQUERY),
    # trim
    base.FuncTrim(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.TRIM(value, " ")),
        ]
    ),
    # upper
    base.FuncUpperConst.for_dialect(D.BIGQUERY),
    base.FuncUpperNonConst.for_dialect(D.BIGQUERY),
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/functions_string.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_type as base

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


V = TranslationVariant.make


DEFINITIONS_TYPE = [
    # bool
    base.FuncBoolFromNull.for_dialect(D.BIGQUERY),
    base.FuncBoolFromNumber.for_dialect(D.BIGQUERY),
    base.FuncBoolFromBool.for_dialect(D.BIGQUERY),
    base.FuncBoolFromStrGeo.for_dialect(D.BIGQUERY),
    base.FuncBoolFromDateDatetime.for_dialect(D.BIGQUERY),
    # date
    base.FuncDate1FromNull.for_dialect(D.BIGQUERY),
    base.FuncDate1FromDatetime.for_dialect(D.BIGQUERY),
    base.FuncDate1FromString.for_dialect(D.BIGQUERY),
    base.FuncDate1FromNumber(
        variants=[V(D.BIGQUERY, lambda expr: sa.cast(sa.func.TIMESTAMP_SECONDS(sa.cast(expr, sa.INTEGER)), sa.DATE))]
    ),
    # datetime
    base.FuncDatetime1FromNull.for_dialect(D.BIGQUERY),
    base.FuncDatetime1FromDatetime.for_dialect(D.BIGQUERY),
    base.FuncDatetime1FromDate.for_dialect(D.BIGQUERY),
    base.FuncDatetime1FromNumber(
        variants=[
            V(D.BIGQUERY, lambda expr: sa.cast(sa.func.TIMESTAMP_SECONDS(sa.cast(expr, sa.INTEGER)), sa.DATETIME))
        ]
    ),
    base.FuncDatetime1FromString.for_dialect(D.BIGQUERY),
    # datetimetz
    base.FuncDatetimeTZConst.for_dialect(D.BIGQUERY),
    # float
    base.FuncFloatNumber(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.FLOAT)),
        ]
    ),
    base.FuncFloatString(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.FLOAT)),
        ]
    ),
    base.FuncFloatFromBool(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.IF(value, 1.0, 0.0)),
        ]
    ),
    base.FuncFloatFromDate(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(sa.func.UNIX_SECONDS(sa.cast(value, sa.TIMESTAMP)), sa.FLOAT)),
        ]
    ),
    base.FuncFloatFromDatetime(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(sa.func.UNIX_SECONDS(sa.cast(value, sa.TIMESTAMP)), sa.FLOAT)),
        ]
    ),
    base.FuncFloatFromGenericDatetime(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(sa.func.UNIX_SECONDS(sa.cast(value, sa.TIMESTAMP)), sa.FLOAT)),
        ]
    ),
    # genericdatetime
    base.FuncGenericDatetime1FromNull.for_dialect(D.BIGQUERY),
    base.FuncGenericDatetime1FromDatetime.for_dialect(D.BIGQUERY),
    base.FuncGenericDatetime1FromDate.for_dialect(D.BIGQUERY),
    base.FuncGenericDatetime1FromNumber(
        variants=[
            V(D.BIGQUERY, lambda expr: sa.cast(sa.func.TIMESTAMP_SECONDS(sa.cast(expr, sa.INTEGER)), sa.DATETIME))
        ]
    ),
    base.FuncGenericDatetime1FromString.for_dialect(D.BIGQUERY),
    # geopoint
    base.FuncGeopointFromStr.for_dialect(D.BIGQUERY),
    base.FuncGeopointFromCoords.for_dialect(D.BIGQUERY),
    # geopolygon
    base.FuncGeopolygon.for_dialect(D.BIGQUERY),
    # int
    base.FuncIntFromNull(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.INTEGER)),
        ]
    ),
    base.FuncIntFromInt(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.INTEGER)),
        ]
    ),
    base.FuncIntFromFloat(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(sa.func.FLOOR(value), sa.INTEGER)),
        ]
    ),
    base.FuncIntFromBool(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.IF(value, 1, 0)),
        ]
    ),
    base.FuncIntFromStr(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.INTEGER)),
        ]
    ),
    base.FuncIntFromDate(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.UNIX_SECONDS(sa.cast(value, sa.TIMESTAMP))),
        ]
    ),
    base.FuncIntFromDatetime(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.UNIX_SECONDS(sa.cast(value, sa.TIMESTAMP))),
        ]
    ),
    base.FuncIntFromGenericDatetime(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.UNIX_SECONDS(sa.cast(value, sa.TIMESTAMP))),
        ]
    ),
    base.FuncIntFromDatetimeTZ(
        variants=[
            V(D.BIGQUERY, lambda value: sa.func.UNIX_SECONDS(sa.cast(value, sa.TIMESTAMP))),
        ]
    ),
    # str
    base.FuncStrFromNull(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(sa.null(), sa.String())),
        ]
    ),
    base.FuncStrFromUnsupported(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    base.FuncStrFromInteger(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    base.FuncStrFromFloat(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    base.FuncStrFromBool(
        variants=[
            V(D.BIGQUERY, lambda value: sa.case(whens=[(value.is_(None), sa.null()), (value, "True")], else_="False")),
        ]
    ),
    base.FuncStrFromStrGeo.for_dialect(D.BIGQUERY),
    base.FuncStrFromDate(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    base.FuncStrFromDatetime(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    base.FuncStrFromDatetimeTZ(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    base.FuncStrFromString.for_dialect(D.BIGQUERY),
    base.FuncStrFromUUID(
        variants=[
            V(D.BIGQUERY, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    # base.FuncStrFromArray,  # FIXME
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/functions_type.py
datalens-tech/datalens-backend
import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common import raw_sql
from dl_formula.definitions.common_datetime import DAY_SEC
import dl_formula.definitions.operators_binary as base

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


V = TranslationVariant.make


def _bq_sec_interval(sec: ClauseElement) -> ClauseElement:
    return sa.func.MAKE_INTERVAL(0, 0, 0, 0, 0, sec)


DEFINITIONS_BINARY = [
    # !=
    base.BinaryNotEqual.for_dialect(D.BIGQUERY),
    # %
    base.BinaryModInteger(
        variants=[
            V(D.BIGQUERY, sa.func.MOD),
        ]
    ),
    base.BinaryModFloat(
        variants=[
            V(D.BIGQUERY, lambda left, right: left - sa.func.FLOOR(left / right) * right),
        ]
    ),
    # *
    base.BinaryMultNumbers.for_dialect(D.BIGQUERY),
    base.BinaryMultStringConst.for_dialect(D.BIGQUERY),
    base.BinaryMultStringNonConst(
        variants=[
            V(D.BIGQUERY, sa.func.REPEAT),
        ]
    ),
    # +
    base.BinaryPlusNumbers.for_dialect(D.BIGQUERY),
    base.BinaryPlusStrings.for_dialect(D.BIGQUERY),
    base.BinaryPlusDateInt(
        variants=[
            V(D.BIGQUERY, lambda date, days: date + days),
        ]
    ),
    base.BinaryPlusDateFloat(
        variants=[
            V(D.BIGQUERY, lambda date, days: date + sa.cast(days, sa.INTEGER)),
        ]
    ),
    base.BinaryPlusDatetimeNumber(
        variants=[
            V(D.BIGQUERY, lambda date, days: date + _bq_sec_interval(sa.cast(days * DAY_SEC, sa.INTEGER))),
        ]
    ),
    base.BinaryPlusGenericDatetimeNumber(
        variants=[
            V(D.BIGQUERY, lambda date, days: date + _bq_sec_interval(sa.cast(days * DAY_SEC, sa.INTEGER))),
        ]
    ),
    # -
    base.BinaryMinusNumbers.for_dialect(D.BIGQUERY),
    base.BinaryMinusDateInt(
        variants=[
            V(D.BIGQUERY, lambda date, days: date - days),
        ]
    ),
    base.BinaryMinusDateFloat(
        variants=[
            V(
                D.BIGQUERY,
                lambda date, days: sa.cast(
                    sa.cast(date, sa.DATETIME) - _bq_sec_interval(sa.cast(days * DAY_SEC, sa.INTEGER)), sa.DATE
                ),
            ),
        ]
    ),
    base.BinaryMinusDatetimeNumber(
        variants=[
            V(
                D.BIGQUERY,
                lambda date, days: sa.cast(date, sa.DATETIME) - _bq_sec_interval(sa.cast(days * DAY_SEC, sa.INTEGER)),
            ),
        ]
    ),
    base.BinaryMinusGenericDatetimeNumber(
        variants=[
            V(D.BIGQUERY, lambda date, days: date - _bq_sec_interval(sa.cast(days * DAY_SEC, sa.INTEGER))),
        ]
    ),
    base.BinaryMinusDates(
        variants=[
            V(D.BIGQUERY, lambda left, right: sa.func.EXTRACT(raw_sql("DAY"), left - right)),
        ]
    ),
    base.BinaryMinusDatetimes(
        variants=[
            V(
                D.BIGQUERY,
                lambda left, right: (
                    sa.func.UNIX_SECONDS(sa.cast(left, sa.TIMESTAMP))
                    - sa.func.UNIX_SECONDS(sa.cast(right, sa.TIMESTAMP))
                )
                / DAY_SEC,
            ),
        ]
    ),
    base.BinaryMinusGenericDatetimes(
        variants=[
            V(
                D.BIGQUERY,
                lambda left, right: (
                    sa.func.UNIX_SECONDS(sa.cast(left, sa.TIMESTAMP))
                    - sa.func.UNIX_SECONDS(sa.cast(right, sa.TIMESTAMP))
                )
                / DAY_SEC,
            ),
        ]
    ),
    # /
    base.BinaryDivInt.for_dialect(D.BIGQUERY),
    base.BinaryDivFloat.for_dialect(D.BIGQUERY),
    # <
    base.BinaryLessThan.for_dialect(D.BIGQUERY),
    # <=
    base.BinaryLessThanOrEqual.for_dialect(D.BIGQUERY),
    # ==
    base.BinaryEqual.for_dialect(D.BIGQUERY),
    # >
    base.BinaryGreaterThan.for_dialect(D.BIGQUERY),
    # >=
    base.BinaryGreaterThanOrEqual.for_dialect(D.BIGQUERY),
    # ^
    base.BinaryPower.for_dialect(D.BIGQUERY),
    # _!=
    base.BinaryNotEqualInternal.for_dialect(D.BIGQUERY),
    # _==
    base.BinaryEqualInternal.for_dialect(D.BIGQUERY),
    # _dneq
    base.BinaryEqualDenullified.for_dialect(D.BIGQUERY),
    # and
    base.BinaryAnd.for_dialect(D.BIGQUERY),
    # in
    base.BinaryIn.for_dialect(D.BIGQUERY),
    # like
    base.BinaryLike.for_dialect(D.BIGQUERY),
    # notin
    base.BinaryNotIn.for_dialect(D.BIGQUERY),
    # notlike
    base.BinaryNotLike.for_dialect(D.BIGQUERY),
    # or
    base.BinaryOr.for_dialect(D.BIGQUERY),
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/operators_binary.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_ternary as base

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


DEFINITIONS_TERNARY = [
    # between
    base.TernaryBetween.for_dialect(D.BIGQUERY),
    # notbetween
    base.TernaryNotBetween.for_dialect(D.BIGQUERY),
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/operators_ternary.py
datalens-tech/datalens-backend
from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.operators_unary as base

from dl_connector_bigquery.formula.constants import BigQueryDialect as D


V = TranslationVariant.make


DEFINITIONS_UNARY = [
    # isfalse
    base.UnaryIsFalseStringGeo.for_dialect(D.BIGQUERY),
    base.UnaryIsFalseNumbers.for_dialect(D.BIGQUERY),
    base.UnaryIsFalseDateTime.for_dialect(D.BIGQUERY),
    base.UnaryIsFalseBoolean.for_dialect(D.BIGQUERY),
    # istrue
    base.UnaryIsTrueStringGeo.for_dialect(D.BIGQUERY),
    base.UnaryIsTrueNumbers.for_dialect(D.BIGQUERY),
    base.UnaryIsTrueDateTime.for_dialect(D.BIGQUERY),
    base.UnaryIsTrueBoolean.for_dialect(D.BIGQUERY),
    # neg
    base.UnaryNegate.for_dialect(D.BIGQUERY),
    # not
    base.UnaryNotBool.for_dialect(D.BIGQUERY),
    base.UnaryNotNumbers.for_dialect(D.BIGQUERY),
    base.UnaryNotStringGeo.for_dialect(D.BIGQUERY),
    base.UnaryNotDateDatetime.for_dialect(D.BIGQUERY),
]

lib/dl_connector_bigquery/dl_connector_bigquery/formula/definitions/operators_unary.py
datalens-tech/datalens-backend
import abc
from typing import ClassVar

import attr

from dl_testing.env_params.getter import EnvParamGetter


class BigQuerySecretReaderBase(abc.ABC):
    KEY_CONFIG: ClassVar[str] = "BIGQUERY_CONFIG"
    KEY_CREDS: ClassVar[str] = "BIGQUERY_CREDS"

    @property
    @abc.abstractmethod
    def project_config(self) -> dict:
        raise NotImplementedError

    @abc.abstractmethod
    def get_creds(self) -> str:
        raise NotImplementedError

    def get_project_id(self) -> str:
        return self.project_config["project_id"]

    def get_dataset_name(self) -> str:
        return self.project_config["dataset_name"]

    def get_table_name(self) -> str:
        return self.project_config["table_name"]


@attr.s
class BigQuerySecretReader(BigQuerySecretReaderBase):
    _env_param_getter: EnvParamGetter = attr.ib()
    _project_config: dict = attr.ib(init=False)

    @_project_config.default
    def _make_project_config(self) -> dict:
        return self._env_param_getter.get_json_value(self.KEY_CONFIG)

    @property
    def project_config(self) -> dict:
        return self._project_config

    def get_creds(self) -> str:
        return self._env_param_getter.get_str_value(self.KEY_CREDS)

lib/dl_connector_bigquery/dl_connector_bigquery/testing/secrets.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_api_lib_testing.data_api_base import StandardizedDataApiTestBase
from dl_api_lib_testing.dataset_base import DatasetTestBase

from dl_connector_bigquery.core.constants import (
    CONNECTION_TYPE_BIGQUERY,
    SOURCE_TYPE_BIGQUERY_TABLE,
)
from dl_connector_bigquery_tests.ext.config import API_TEST_CONFIG
from dl_connector_bigquery_tests.ext.core.base import BaseBigQueryTestClass


class BigQueryConnectionTestBase(BaseBigQueryTestClass, ConnectionTestBase):
    conn_type = CONNECTION_TYPE_BIGQUERY
    bi_compeng_pg_on = False

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def connection_params(self, bq_secrets) -> dict:
        return dict(
            project_id=bq_secrets.get_project_id(),
            credentials=bq_secrets.get_creds(),
        )


class BigQueryDatasetTestBase(BigQueryConnectionTestBase, DatasetTestBase):
    @pytest.fixture(scope="class")
    def dataset_params(self, sample_table) -> dict:
        return dict(
            source_type=SOURCE_TYPE_BIGQUERY_TABLE.name,
            parameters=dict(
                dataset_name=sample_table.schema,
                table_name=sample_table.name,
            ),
        )


class BigQueryDataApiTestBase(BigQueryDatasetTestBase, StandardizedDataApiTestBase):
    mutation_caches_on = False

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/base.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.complex_queries import DefaultBasicComplexQueryTestSuite
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_bigquery_tests.ext.api.base import BigQueryDataApiTestBase


class TestBigQueryBasicComplexQueries(BigQueryDataApiTestBase, DefaultBasicComplexQueryTestSuite):
    test_params = RegulatedTestParams(
        mark_features_skipped={
            DefaultBasicComplexQueryTestSuite.feature_window_functions: "Native window functions are not implemented"
        }
    )

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/test_complex_queries.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite

from dl_connector_bigquery_tests.ext.api.base import BigQueryConnectionTestBase


class TestBigQueryConnection(BigQueryConnectionTestBase, DefaultConnectorConnectionTestSuite):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/test_connection.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.data_api_suites import (
    DefaultConnectorDataDistinctTestSuite,
    DefaultConnectorDataGroupByFormulaTestSuite,
    DefaultConnectorDataPreviewTestSuite,
    DefaultConnectorDataRangeTestSuite,
    DefaultConnectorDataResultTestSuite,
)
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_bigquery_tests.ext.api.base import BigQueryDataApiTestBase


class TestBigQueryDataResult(BigQueryDataApiTestBase, DefaultConnectorDataResultTestSuite):
    test_params = RegulatedTestParams(
        mark_features_skipped={
            DefaultConnectorDataResultTestSuite.array_support: "BigQuery doesn't support arrays",
        }
    )


class TestBigQueryDataGroupBy(BigQueryDataApiTestBase, DefaultConnectorDataGroupByFormulaTestSuite):
    pass


class TestBigQueryDataRange(BigQueryDataApiTestBase, DefaultConnectorDataRangeTestSuite):
    pass


class TestBigQueryDataDistinct(BigQueryDataApiTestBase, DefaultConnectorDataDistinctTestSuite):
    pass


class TestBigQueryDataPreview(BigQueryDataApiTestBase, DefaultConnectorDataPreviewTestSuite):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/test_data.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite

from dl_connector_bigquery_tests.ext.api.base import BigQueryDatasetTestBase


class TestBigQueryDataset(BigQueryDatasetTestBase, DefaultConnectorDatasetTestSuite):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/api/test_dataset.py
datalens-tech/datalens-backend
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport


# Infra settings
CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=51600).host,
    port_us_http=get_test_container_hostport("us", fallback_port=51600).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=51609).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=51609).port,
    us_master_token="AC1ofiek8coB",
    core_connector_ep_names=["bigquery"],
)


API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["bigquery"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
)

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/config.py
datalens-tech/datalens-backend
import os

import pytest

from dl_api_lib_testing.initialization import initialize_api_lib_test
from dl_formula_testing.forced_literal import forced_literal_use
from dl_testing.env_params.generic import GenericEnvParamGetter

from dl_connector_bigquery.testing.secrets import (
    BigQuerySecretReader,
    BigQuerySecretReaderBase,
)
from dl_connector_bigquery_tests.ext.config import API_TEST_CONFIG


pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)


@pytest.fixture(scope="session")
def env_param_getter() -> GenericEnvParamGetter:
    filepath = os.path.join(os.path.dirname(__file__), "params.yml")
    return GenericEnvParamGetter.from_yaml_file(filepath)


@pytest.fixture(scope="session")
def bq_secrets(env_param_getter) -> BigQuerySecretReaderBase:
    return BigQuerySecretReader(env_param_getter=env_param_getter)


__all__ = (
    # auto-use fixtures:
    "forced_literal_use",
)

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/conftest.py
datalens-tech/datalens-backend
import asyncio
from typing import Generator

import pytest

from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import (
    C,
    Db,
    DbTable,
    make_table,
)
from dl_core_testing.fixtures.sample_tables import TABLE_SPEC_SAMPLE_SUPERSTORE
from dl_core_testing.testcases.connection import BaseConnectionTestClass

from dl_connector_bigquery.core.constants import CONNECTION_TYPE_BIGQUERY
from dl_connector_bigquery.core.testing.connection import make_bigquery_saved_connection
from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery
from dl_connector_bigquery.db_testing.engine_wrapper import BigQueryDbEngineConfig
import dl_connector_bigquery_tests.ext.config as test_config


class BaseBigQueryTestClass(BaseConnectionTestClass[ConnectionSQLBigQuery]):
    conn_type = CONNECTION_TYPE_BIGQUERY
    core_test_config = test_config.CORE_TEST_CONFIG

    @pytest.fixture(autouse=True)
    # FIXME: This fixture is a temporary solution for failing core tests when they are run together with api tests
    def loop(self, event_loop: asyncio.AbstractEventLoop) -> Generator[asyncio.AbstractEventLoop, None, None]:
        asyncio.set_event_loop(event_loop)
        yield event_loop
        # Attempt to cover an old version of pytest-asyncio:
        # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
        asyncio.set_event_loop_policy(None)

    @pytest.fixture(scope="class")
    def db_url(self, bq_secrets) -> str:
        return f"bigquery://{bq_secrets.get_project_id()}"

    @pytest.fixture(scope="class")
    def engine_params(self, bq_secrets) -> dict:
        return dict(
            credentials_base64=bq_secrets.get_creds(),
        )

    @pytest.fixture(scope="class")
    def engine_config(self, db_url: str, engine_params: dict, bq_secrets) -> BigQueryDbEngineConfig:
        return BigQueryDbEngineConfig(
            url=db_url,
            engine_params=engine_params,
            default_dataset_name=bq_secrets.get_dataset_name(),
        )

    @pytest.fixture(scope="function")
    def connection_creation_params(self, bq_secrets) -> dict:
        return dict(
            project_id=bq_secrets.get_project_id(),
            credentials=bq_secrets.get_creds(),
            **(dict(raw_sql_level=self.raw_sql_level) if self.raw_sql_level is not None else {}),
        )

    @pytest.fixture(scope="function")
    def saved_connection(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params: dict,
    ) -> ConnectionSQLBigQuery:
        conn = make_bigquery_saved_connection(sync_usm=sync_us_manager, **connection_creation_params)
        return conn

    @pytest.fixture(scope="class")
    def sample_table(self, db: Db, bq_secrets) -> DbTable:
        # FIXME: Do this via a "smart" dispenser
        return make_table(
            db=db,
            name=bq_secrets.get_table_name(),
            schema=bq_secrets.get_dataset_name(),
            columns=[
                C(name=name, user_type=user_type) for name, user_type in TABLE_SPEC_SAMPLE_SUPERSTORE.table_schema
            ],
            create_in_db=False,
        )

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/core/base.py
datalens-tech/datalens-backend
from dl_core.us_connection_base import DataSourceTemplate
from dl_core_testing.testcases.connection import DefaultConnectionTestClass

from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery
from dl_connector_bigquery_tests.ext.core.base import BaseBigQueryTestClass


class TestBigQueryConnection(
    BaseBigQueryTestClass,
    DefaultConnectionTestClass[ConnectionSQLBigQuery],
):
    def check_saved_connection(self, conn: ConnectionSQLBigQuery, params: dict) -> None:
        assert conn.uuid is not None
        assert conn.data.project_id == params["project_id"]
        assert conn.data.credentials == params["credentials"]

    def check_data_source_templates(
        self,
        conn: ConnectionSQLBigQuery,
        dsrc_templates: list[DataSourceTemplate],
    ) -> None:
        for dsrc_tmpl in dsrc_templates:
            assert conn.project_id in dsrc_tmpl.group
            assert dsrc_tmpl.title

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/core/test_connection.py
datalens-tech/datalens-backend
import pytest

from dl_constants.enums import UserDataType
from dl_core.connection_models.common_models import DBIdent
from dl_core_testing.database import (
    C,
    Db,
)
from dl_core_testing.testcases.connection_executor import (
    DefaultAsyncConnectionExecutorTestSuite,
    DefaultSyncAsyncConnectionExecutorCheckBase,
    DefaultSyncConnectionExecutorTestSuite,
)
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery
from dl_connector_bigquery_tests.ext.core.base import BaseBigQueryTestClass


class BigQuerySyncAsyncConnectionExecutorCheckBase(
    BaseBigQueryTestClass,
    DefaultSyncAsyncConnectionExecutorCheckBase[ConnectionSQLBigQuery],
):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultAsyncConnectionExecutorTestSuite.test_error_on_select_from_nonexistent_source: "Need to learn to detect this error",  # TODO: FIXME
            DefaultAsyncConnectionExecutorTestSuite.test_closing_sql_sessions: "Not implemented",
        },
    )

    @pytest.fixture(scope="class")
    def db_table_columns(self, db: Db) -> list[C]:
        return [
            col_spec
            for col_spec in C.full_house()
            if col_spec.user_type
            not in (
                UserDataType.uuid,  # UUID is not supported
                UserDataType.datetime,
                UserDataType.genericdatetime,  # datetimes with fractional seconds are not supported  # FIXME
            )
        ]

    @pytest.fixture(scope="function")
    def db_ident(self) -> DBIdent:
        return DBIdent(db_name=None)


class TestBigQuerySyncConnectionExecutor(
    BigQuerySyncAsyncConnectionExecutorCheckBase,
    DefaultSyncConnectionExecutorTestSuite[ConnectionSQLBigQuery],
):
    pass


class TestBigQueryAsyncConnectionExecutor(
    BigQuerySyncAsyncConnectionExecutorCheckBase,
    DefaultAsyncConnectionExecutorTestSuite[ConnectionSQLBigQuery],
):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/core/test_connection_executor.py
datalens-tech/datalens-backend
import pytest

from dl_constants.enums import (
    RawSQLLevel,
    UserDataType,
)
from dl_core_testing.fixtures.sample_tables import TABLE_SPEC_SAMPLE_SUPERSTORE
from dl_core_testing.testcases.data_source import DefaultDataSourceTestClass

from dl_connector_bigquery.core.constants import (
    SOURCE_TYPE_BIGQUERY_SUBSELECT,
    SOURCE_TYPE_BIGQUERY_TABLE,
)
from dl_connector_bigquery.core.data_source import (
    BigQuerySubselectDataSource,
    BigQueryTableDataSource,
)
from dl_connector_bigquery.core.data_source_spec import (
    BigQuerySubselectDataSourceSpec,
    BigQueryTableDataSourceSpec,
)
from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery
from dl_connector_bigquery_tests.ext.core.base import BaseBigQueryTestClass


class TestBigQueryTableDataSource(
    BaseBigQueryTestClass,
    DefaultDataSourceTestClass[
        ConnectionSQLBigQuery,
        BigQueryTableDataSourceSpec,
        BigQueryTableDataSource,
    ],
):
    DSRC_CLS = BigQueryTableDataSource

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self, sample_table) -> BigQueryTableDataSourceSpec:
        dsrc_spec = BigQueryTableDataSourceSpec(
            source_type=SOURCE_TYPE_BIGQUERY_TABLE,
            dataset_name=sample_table.schema,
            table_name=sample_table.name,
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return list(TABLE_SPEC_SAMPLE_SUPERSTORE.table_schema)


class TestBigQuerySubselectDataSource(
    BaseBigQueryTestClass,
    DefaultDataSourceTestClass[
        ConnectionSQLBigQuery,
        BigQuerySubselectDataSourceSpec,
        BigQuerySubselectDataSource,
    ],
):
    DSRC_CLS = BigQuerySubselectDataSource

    raw_sql_level = RawSQLLevel.subselect

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self, sample_table) -> BigQuerySubselectDataSourceSpec:
        dsrc_spec = BigQuerySubselectDataSourceSpec(
            source_type=SOURCE_TYPE_BIGQUERY_SUBSELECT,
            subsql=f"SELECT * FROM {sample_table.schema}.{sample_table.name}",
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return list(TABLE_SPEC_SAMPLE_SUPERSTORE.table_schema)

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/core/test_data_source.py
datalens-tech/datalens-backend
import pytest

from dl_core_testing.database import DbTable
from dl_core_testing.testcases.dataset import DefaultDatasetTestSuite
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_bigquery.core.constants import SOURCE_TYPE_BIGQUERY_TABLE
from dl_connector_bigquery.core.us_connection import ConnectionSQLBigQuery
from dl_connector_bigquery_tests.ext.core.base import BaseBigQueryTestClass


class TestBigQueryDataset(BaseBigQueryTestClass, DefaultDatasetTestSuite[ConnectionSQLBigQuery]):
    source_type = SOURCE_TYPE_BIGQUERY_TABLE

    test_params = RegulatedTestParams(
        mark_tests_failed={
            DefaultDatasetTestSuite.test_get_param_hash: "",  # TODO: FIXME
        },
    )

    @pytest.fixture(scope="function")
    def dsrc_params(self, dataset_table: DbTable) -> dict:
        return dict(
            project_id=dataset_table.db.name,
            dataset_name=dataset_table.schema,
            table_name=dataset_table.name,
        )

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/core/test_dataset.py
datalens-tech/datalens-backend
import re

import pytest
import sqlalchemy.exc as sa_exc

from dl_formula_testing.testcases.base import FormulaConnectorTestBase

from dl_connector_bigquery.db_testing.engine_wrapper import BigQueryDbEngineConfig
from dl_connector_bigquery.formula.constants import BigQueryDialect as D


class BigQueryTestBase(FormulaConnectorTestBase):
    dialect = D.BIGQUERY
    eval_attempts = 5
    retry_on_exceptions = ((sa_exc.DatabaseError, re.compile("404 Not found: Table")),)
    engine_config_cls = BigQueryDbEngineConfig

    supports_arrays = False
    supports_uuid = False
    null_casts_to_false = True

    @pytest.fixture(scope="class")
    def engine_params(self, bq_secrets) -> dict:
        return dict(
            credentials_base64=bq_secrets.get_creds(),
        )

    @pytest.fixture(scope="class")
    def engine_config(self, db_url: str, engine_params: dict, bq_secrets) -> BigQueryDbEngineConfig:
        default_dataset_name = bq_secrets.get_dataset_name()
        return BigQueryDbEngineConfig(
            url=db_url,
            engine_params=engine_params,
            default_dataset_name=default_dataset_name,
        )

    @pytest.fixture(scope="class")
    def db_url(self, bq_secrets) -> str:
        return f"bigquery://{bq_secrets.get_project_id()}"

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/formula/base.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.conditional_blocks import DefaultConditionalBlockFormulaConnectorTestSuite

from dl_connector_bigquery_tests.ext.formula.base import BigQueryTestBase


class TestConditionalBlockBigQuery(BigQueryTestBase, DefaultConditionalBlockFormulaConnectorTestSuite):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/formula/test_conditional_blocks.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_aggregation import DefaultMainAggFunctionFormulaConnectorTestSuite

from dl_connector_bigquery_tests.ext.formula.base import BigQueryTestBase


class TestMainAggFunctionBigQuery(BigQueryTestBase, DefaultMainAggFunctionFormulaConnectorTestSuite):
    supports_countd_approx = True

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/formula/test_functions_aggregation.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_datetime import DefaultDateTimeFunctionFormulaConnectorTestSuite

from dl_connector_bigquery_tests.ext.formula.base import BigQueryTestBase


class TestDateTimeFunctionBigQuery(BigQueryTestBase, DefaultDateTimeFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/formula/test_functions_datetime.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_logical import DefaultLogicalFunctionFormulaConnectorTestSuite

from dl_connector_bigquery_tests.ext.formula.base import BigQueryTestBase


class TestLogicalFunctionBigQuery(BigQueryTestBase, DefaultLogicalFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/formula/test_functions_logical.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_markup import DefaultMarkupFunctionFormulaConnectorTestSuite

from dl_connector_bigquery_tests.ext.formula.base import BigQueryTestBase


class TestMarkupFunctionBigQuery(BigQueryTestBase, DefaultMarkupFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/formula/test_functions_markup.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_math import DefaultMathFunctionFormulaConnectorTestSuite

from dl_connector_bigquery_tests.ext.formula.base import BigQueryTestBase


class TestMathFunctionBigQuery(BigQueryTestBase, DefaultMathFunctionFormulaConnectorTestSuite):
    supports_float_div = False

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/formula/test_functions_math.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_string import DefaultStringFunctionFormulaConnectorTestSuite

from dl_connector_bigquery_tests.ext.formula.base import BigQueryTestBase


class TestStringFunctionBigQuery(BigQueryTestBase, DefaultStringFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/formula/test_functions_string.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_type_conversion import (
    DefaultBoolTypeFunctionFormulaConnectorTestSuite,
    DefaultDateTypeFunctionFormulaConnectorTestSuite,
    DefaultFloatTypeFunctionFormulaConnectorTestSuite,
    DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite,
    DefaultGeopointTypeFunctionFormulaConnectorTestSuite,
    DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite,
    DefaultIntTypeFunctionFormulaConnectorTestSuite,
    DefaultStrTypeFunctionFormulaConnectorTestSuite,
)

from dl_connector_bigquery_tests.ext.formula.base import BigQueryTestBase


# STR


class TestStrTypeFunctionBigQuery(BigQueryTestBase, DefaultStrTypeFunctionFormulaConnectorTestSuite):
    skip_custom_tz = True


# FLOAT


class TestFloatTypeFunctionBigQuery(BigQueryTestBase, DefaultFloatTypeFunctionFormulaConnectorTestSuite):
    pass


# BOOL


class TestBoolTypeFunctionBigQuery(BigQueryTestBase, DefaultBoolTypeFunctionFormulaConnectorTestSuite):
    pass


# INT


class TestIntTypeFunctionBigQuery(BigQueryTestBase, DefaultIntTypeFunctionFormulaConnectorTestSuite):
    pass


# DATE


class TestDateTypeFunctionBigQuery(BigQueryTestBase, DefaultDateTypeFunctionFormulaConnectorTestSuite):
    pass


# GENERICDATETIME (& DATETIME)


class TestGenericDatetimeTypeFunctionBigQuery(
    BigQueryTestBase,
    DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite,
):
    pass


# GEOPOINT


class TestGeopointTypeFunctionBigQuery(
    BigQueryTestBase,
    DefaultGeopointTypeFunctionFormulaConnectorTestSuite,
):
    pass


# GEOPOLYGON


class TestGeopolygonTypeFunctionBigQuery(
    BigQueryTestBase,
    DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite,
):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/formula/test_functions_type_conversion.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.operators import DefaultOperatorFormulaConnectorTestSuite

from dl_connector_bigquery_tests.ext.formula.base import BigQueryTestBase


class TestOperatorBigQuery(BigQueryTestBase, DefaultOperatorFormulaConnectorTestSuite):
    pass

lib/dl_connector_bigquery/dl_connector_bigquery_tests/ext/formula/test_operators.py
datalens-tech/datalens-backend
from dl_api_connector.i18n.localizer import CONFIGS as BI_API_CONNECTOR_CONFIGS
from dl_api_lib_testing.connection_form_base import ConnectionFormTestBase

from dl_connector_bigquery.api.connection_form.form_config import BigQueryConnectionFormFactory
from dl_connector_bigquery.api.i18n.localizer import CONFIGS as BI_CONNECTOR_BIGQUERY_CONFIGS


class TestBigQueryConnectionForm(ConnectionFormTestBase):
    CONN_FORM_FACTORY_CLS = BigQueryConnectionFormFactory
    TRANSLATION_CONFIGS = BI_CONNECTOR_BIGQUERY_CONFIGS + BI_API_CONNECTOR_CONFIGS

lib/dl_connector_bigquery/dl_connector_bigquery_tests/unit/test_connection_form.py
datalens-tech/datalens-backend
from __future__ import annotations

from urllib.parse import urlparse

from marshmallow import fields as ma_fields
from marshmallow import validate as ma_validate

from dl_api_connector.api_schema.connection_base import (
    ConnectionMetaMixin,
    ConnectionSchema,
)
from dl_api_connector.api_schema.connection_base_fields import (
    cache_ttl_field,
    secret_string_field,
)
from dl_api_connector.api_schema.extras import FieldExtra

from dl_connector_bitrix_gds.core.us_connection import BitrixGDSConnection


class BitrixPortalValidator(ma_validate.Validator):
    error = "Not a valid portal name"

    def __call__(self, portal: str) -> str:
        parsed_host = urlparse("//{}".format(portal)).hostname
        if portal.lower() != parsed_host:
            raise ma_validate.ValidationError(self.error)
        return portal


class BitrixGDSConnectionSchema(ConnectionMetaMixin, ConnectionSchema):
    TARGET_CLS = BitrixGDSConnection

    portal = ma_fields.String(
        attribute="data.portal",
        required=True,
        allow_none=False,
        bi_extra=FieldExtra(editable=True),
        validate=BitrixPortalValidator(),
    )
    token = secret_string_field(attribute="data.token")
    cache_ttl_sec = cache_ttl_field(attribute="data.cache_ttl_sec")

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/api/api_schema/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import unique
from typing import Optional

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormActionApiSchema,
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
)
from dl_api_connector.form_config.models.common import (
    CommonFieldName,
    FormFieldName,
)
import dl_api_connector.form_config.models.rows as C
from dl_api_connector.form_config.models.rows.base import FormRow
from dl_api_connector.form_config.models.shortcuts.rows import RowConstructor
from dl_configs.connectors_settings import ConnectorSettingsBase

from dl_connector_bitrix_gds.api.connection_info import BitrixGDSConnectionInfoProvider
from dl_connector_bitrix_gds.api.i18n.localizer import Translatable


@unique
class BitrixGDSFieldName(FormFieldName):
    portal = "portal"


class BitrixGDSConnectionFormFactory(ConnectionFormFactory):
    def get_form_config(
        self,
        connector_settings: Optional[ConnectorSettingsBase],
        tenant: Optional[TenantDef],
    ) -> ConnectionForm:
        rc = RowConstructor(localizer=self._localizer)

        token_row = C.CustomizableRow(
            items=[
                C.LabelRowItem(text=self._localizer.translate(Translatable("label_token"))),
                C.InputRowItem(
                    name=CommonFieldName.token,
                    width="l",
                    control_props=C.InputRowItem.Props(type="password"),
                    default_value="" if self.mode == ConnectionFormMode.create else None,
                    fake_value=None if self.mode == ConnectionFormMode.create else "******",
                ),
            ]
        )

        portal_row = C.CustomizableRow(
            items=[
                C.LabelRowItem(text=self._localizer.translate(Translatable("label_portal"))),
                C.InputRowItem(name=BitrixGDSFieldName.portal),
            ]
        )

        rows: list[FormRow] = [
            portal_row,
            token_row,
        ]
        if self.mode == ConnectionFormMode.create:
            rows.append(rc.auto_create_dash_row())

        edit_api_schema = FormActionApiSchema(
            items=[
                FormFieldApiSchema(name=BitrixGDSFieldName.portal, required=True),
                FormFieldApiSchema(name=CommonFieldName.token, required=self.mode == ConnectionFormMode.create),
            ]
        )

        create_api_schema = FormActionApiSchema(
            items=[
                *edit_api_schema.items,
                *self._get_top_level_create_api_schema_items(),
            ]
        )

        check_api_schema = FormActionApiSchema(
            items=[
                *edit_api_schema.items,
                *self._get_top_level_check_api_schema_items(),
            ]
        )

        return ConnectionForm(
            title=BitrixGDSConnectionInfoProvider.get_title(self._localizer),
            template_name="bitrix24",
            rows=rows,
            api_schema=FormApiSchema(
                create=create_api_schema if self.mode == ConnectionFormMode.create else None,
                edit=edit_api_schema if self.mode == ConnectionFormMode.edit else None,
                check=check_api_schema,
            ),
        )

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/api/connection_form/form_config.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_bitrix_gds.api.i18n.localizer import Translatable


class BitrixGDSConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-bitrix")

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/api/connection_info.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_connector.api_schema.source_base import (
    SQLDataSourceSchema,
    SQLDataSourceTemplateSchema,
)
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)
from dl_api_lib.query.registry import MQMFactorySettingItem
from dl_constants.enums import QueryProcessingMode

from dl_connector_bitrix_gds.api.api_schema.connection import BitrixGDSConnectionSchema
from dl_connector_bitrix_gds.api.connection_form.form_config import BitrixGDSConnectionFormFactory
from dl_connector_bitrix_gds.api.connection_info import BitrixGDSConnectionInfoProvider
from dl_connector_bitrix_gds.api.filter_compiler import BitrixGDSFilterFormulaCompiler
from dl_connector_bitrix_gds.api.i18n.localizer import CONFIGS
from dl_connector_bitrix_gds.api.multi_query import BitrixGDSMultiQueryMutatorFactory
from dl_connector_bitrix_gds.core.connector import (
    BitrixGDSCoreConnectionDefinition,
    BitrixGDSCoreConnector,
    BitrixGDSCoreSourceDefinition,
)
from dl_connector_bitrix_gds.formula.constants import DIALECT_NAME_BITRIX


class BitrixGDSApiSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = BitrixGDSCoreSourceDefinition
    api_schema_cls = SQLDataSourceSchema
    template_api_schema_cls = SQLDataSourceTemplateSchema


class BitrixGDSApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = BitrixGDSCoreConnectionDefinition
    api_generic_schema_cls = BitrixGDSConnectionSchema
    form_factory_cls = BitrixGDSConnectionFormFactory
    info_provider_cls = BitrixGDSConnectionInfoProvider


class BitrixGDSApiConnector(ApiConnector):
    core_connector_cls = BitrixGDSCoreConnector
    formula_dialect_name = DIALECT_NAME_BITRIX
    multi_query_mutation_factories = (
        MQMFactorySettingItem(
            query_proc_mode=QueryProcessingMode.basic,
            factory_cls=BitrixGDSMultiQueryMutatorFactory,
        ),
    )
    connection_definitions = (BitrixGDSApiConnectionDefinition,)
    source_definitions = (BitrixGDSApiSourceDefinition,)
    is_forkable = False
    is_compeng_executable = True
    filter_formula_compiler_cls = BitrixGDSFilterFormulaCompiler
    translation_configs = frozenset(CONFIGS)

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/api/connector.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import TYPE_CHECKING

from dl_query_processing.compilation.filter_compiler import MainFilterFormulaCompiler


if TYPE_CHECKING:
    from dl_query_processing.compilation.filter_compiler import FilterParams


class BitrixGDSFilterFormulaCompiler(MainFilterFormulaCompiler):
    """connector-specific customizations point"""

    def _mangle_date_filter(self, filter_params: FilterParams) -> FilterParams:
        return filter_params  # Disable the datetime mangling

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/api/filter_compiler.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_bitrix_gds as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/api/i18n/localizer.py
datalens-tech/datalens-backend
from typing import ClassVar

import attr

from dl_constants.enums import UserDataType
from dl_core.fields import ResultSchema
import dl_formula.core.nodes as formula_nodes
from dl_query_processing.compilation.primitives import CompiledFormulaInfo
from dl_query_processing.enums import ExecutionLevel
from dl_query_processing.multi_query.factory import MultiQueryMutatorFactoryBase
from dl_query_processing.multi_query.mutators.base import MultiQueryMutatorBase
from dl_query_processing.multi_query.mutators.splitter_based import SplitterMultiQueryMutator
from dl_query_processing.multi_query.splitters.prefiltered import PrefilteredFieldMultiQuerySplitter
from dl_query_processing.multi_query.splitters.query_fork import QueryForkQuerySplitter


@attr.s
class BitrixGDSMultiQuerySplitter(PrefilteredFieldMultiQuerySplitter):
    expr_names: ClassVar[set[str]] = {"between", ">", ">=", "<", "<=", "=="}
    data_types: ClassVar[set[UserDataType]] = {UserDataType.datetime, UserDataType.date, UserDataType.genericdatetime}

    result_schema: ResultSchema = attr.ib(kw_only=True)

    def is_pre_filter(self, formula: CompiledFormulaInfo) -> bool:
        assert formula.original_field_id is not None
        expr = formula.formula_obj.expr
        if not isinstance(expr, formula_nodes.OperationCall):
            return False
        field = self.result_schema.by_guid(formula.original_field_id)
        if field.data_type in self.data_types and expr.name in self.expr_names:
            # FIXME: Refactor this
            return True
        return False


class BitrixGDSMultiQueryMutatorFactory(MultiQueryMutatorFactoryBase):
    def get_mutators(self) -> list[MultiQueryMutatorBase]:
        return [
            SplitterMultiQueryMutator(
                splitters=[
                    BitrixGDSMultiQuerySplitter(
                        crop_to_level_type=ExecutionLevel.compeng,
                        result_schema=self.result_schema,
                    ),
                    QueryForkQuerySplitter(),
                ],
            )
        ]

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/api/multi_query.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import json
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Optional,
    Union,
)
import uuid

import attr
from redis_cache_lock.types import TClientACM
from redis_cache_lock.utils import wrap_generate_func
import sqlalchemy as sa
from sqlalchemy.sql.elements import TypeCoerce

from dl_app_tools.profiling_base import generic_profiler_async
from dl_constants.enums import ConnectionType
from dl_core.aio.web_app_services.redis import RedisConnParams
from dl_core.connection_executors.adapters.async_adapters_aiohttp import AiohttpDBAdapter
from dl_core.connection_executors.adapters.async_adapters_base import AsyncRawExecutionResult
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    RawColumnInfo,
    RawSchemaInfo,
)
from dl_core.connection_models import (
    DBIdent,
    SchemaIdent,
    TableDefinition,
    TableIdent,
)
from dl_core.connectors.base.error_handling import ETBasedExceptionMaker
from dl_core.data_processing.cache.engine import RedisCacheLockWrapped
from dl_core.db.native_type import CommonNativeType
from dl_core.exc import DatabaseQueryError

from dl_connector_bitrix_gds.core.caches import (
    bitrix_cache_deserializer,
    bitrix_cache_serializer,
    build_local_key_rep,
    get_redis_cli_acm_from_params,
)
from dl_connector_bitrix_gds.core.constants import (
    CONNECTION_TYPE_BITRIX24,
    DEFAULT_DB,
)
from dl_connector_bitrix_gds.core.error_transformer import bitrix_error_transformer
from dl_connector_bitrix_gds.core.tables import (
    BITRIX_TABLES_MAP,
    CRM_DYNAMIC_ITEMS_TABLE,
    SMART_PROCESS_TABLE_PREFIX,
)


if TYPE_CHECKING:
    from sqlalchemy.sql.elements import Label

    from dl_connector_bitrix_gds.core.tables import BitrixGDSTable
    from dl_connector_bitrix_gds.core.target_dto import BitrixGDSConnTargetDTO

LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True, kw_only=True)
class BitrixRequestPayload:
    portal: str = attr.ib()
    table: str = attr.ib()
    json_body: dict = attr.ib()
    flatten_body: dict = attr.ib()


def extract_select_column_name(column: Label) -> str:
    element = column.element
    if isinstance(element, TypeCoerce):
        element = element.typed_expression
    return element.name


class BitrixGDSDefaultAdapter(AiohttpDBAdapter, ETBasedExceptionMaker):
    conn_type: ClassVar[ConnectionType] = CONNECTION_TYPE_BITRIX24
    _target_dto: BitrixGDSConnTargetDTO = attr.ib()
    _redis_cli_acm: Optional[TClientACM] = attr.ib(init=False)

    table: Optional[BitrixGDSTable] = None

    _error_transformer = bitrix_error_transformer

    EXTRA_EXC_CLS = (json.JSONDecodeError,)

    def __attrs_post_init__(self) -> None:
        super().__attrs_post_init__()
        redis_conn_params: Optional[RedisConnParams]
        if self._target_dto.redis_conn_params is not None:
            redis_conn_params = RedisConnParams(**self._target_dto.redis_conn_params)
        else:
            redis_conn_params = None
        self._redis_cli_acm = get_redis_cli_acm_from_params(redis_conn_params)

    @generic_profiler_async("db-query-cached")  # type: ignore  # TODO: fix
    async def _run_query_cached(self, dba_query: DBAdapterQuery) -> Any:
        async def wrap_run_query() -> Any:
            result = await self._run_query(dba_query)
            return result

        payload = self._build_request_payload(dba_query)
        local_key_rep = build_local_key_rep(payload.portal, payload.table, payload.flatten_body)

        assert self._redis_cli_acm is not None
        rcl = RedisCacheLockWrapped(
            key=local_key_rep.key_parts_hash,
            client_acm=self._redis_cli_acm,
            resource_tag="bic_conn_bitrix_query_cache",
            lock_ttl_sec=60,
            data_ttl_sec=self._target_dto.redis_caches_ttl or 600,
        )

        result_b, result = await rcl.generate_with_lock(
            generate_func=wrap_generate_func(
                func=wrap_run_query,
                serialize=bitrix_cache_serializer,
            ),
        )
        if result is None:
            LOGGER.info("Result found in cache: %s", local_key_rep.key_parts_hash)
            result = bitrix_cache_deserializer(result_b)

        return result

    @generic_profiler_async("db-query")  # type: ignore  # TODO: fix
    async def _run_query(self, dba_query: DBAdapterQuery) -> Any:
        query_text = self.compile_query_for_execution(dba_query.query)
        payload = self._build_request_payload(dba_query)

        api_url = f"https://{self._target_dto.portal}/bitrix/tools/biconnector/pbi.php"
        request_id = self._req_ctx_info.request_id or str(uuid.uuid4())
        LOGGER.info(
            "Sending query to Bitrix:\nrequest_id: %s\nurl: %s\nquery: %s\nparams: %s",
            request_id,
            api_url,
            query_text,
            json.dumps({k: (v if k != "key" else "...") for k, v in payload.json_body.items()}),
        )

        with self.handle_execution_error(query_text):
            resp = await self._session.post(
                url=api_url,
                params={
                    "table": payload.table,
                    "consumer": "datalens",
                    "request_id": request_id,
                },
                json=payload.json_body,
            )

            if resp.status != 200:
                body = await resp.text()
                raise DatabaseQueryError(db_message=body, query=query_text)

            resp_body = await resp.json()

        return resp_body

    def _parse_response_body_data(self, body: list, selected_columns: Optional[list] = None) -> dict:
        if not len(body):
            raise ValueError("empty response")
        cols = body[0]
        rows = body[1:]

        assert self.table is not None
        columns_type = self.table.get_columns_type()
        if selected_columns is None:
            selected_columns = cols

        try:
            normalized_data = dict(
                cols=[dict(id=col, label=col, type=columns_type.get(col, "string")) for col in selected_columns],
                rows=[[dict(zip(cols, row))[col] for col in selected_columns] for row in rows],
            )
        except (KeyError, TypeError, ValueError):
            raise ValueError("unexpected data structure")

        return normalized_data

    def _table_schema(self, table: str) -> BitrixGDSTable:
        if table.startswith(SMART_PROCESS_TABLE_PREFIX):
            return CRM_DYNAMIC_ITEMS_TABLE
        else:
            return BITRIX_TABLES_MAP[table]

    def _build_request_payload(self, dba_query: DBAdapterQuery) -> BitrixRequestPayload:
        table = self._extract_table_name(dba_query.query)
        if self.table is None:
            self.table = self._table_schema(table)
        json_body, flatten_body = self.generate_body(dba_query)
        payload = BitrixRequestPayload(
            table=table,
            portal=self._target_dto.portal,
            json_body=json_body,
            flatten_body=flatten_body,
        )
        return payload

    def _extract_table_name(self, query: Union[sa.sql.Select, str]) -> str:
        assert isinstance(query, sa.sql.Select)
        froms = query.froms[0]
        if isinstance(froms, sa.sql.Subquery) and hasattr(froms, "element"):
            froms = froms.element.froms[0]
        assert isinstance(froms, sa.sql.TableClause)
        return froms.name

    def _parse_response_body(self, body: Any, dba_query: DBAdapterQuery) -> dict:
        assert isinstance(dba_query.query, sa.sql.Select)
        selected_columns_values = dba_query.query.selected_columns.values()
        selected_columns: Optional[list[str]] = None
        if "*" not in set(column.name for column in selected_columns_values):
            selected_columns = [extract_select_column_name(column) for column in selected_columns_values]
            # 'table."COLUMN_NAME"' -> 'COLUMN_NAME'
            selected_columns = [col.split(".")[-1].replace('"', "").replace("`", "") for col in selected_columns]

        try:
            if not isinstance(body, list):
                raise TypeError("Unexpected response format")
            return self._parse_response_body_data(body, selected_columns=selected_columns)
        except (ValueError, TypeError) as err:
            LOGGER.debug("Unexpected API response")
            raise DatabaseQueryError(
                message=f"Unexpected API response body: {err.args[0]}",
                db_message="",
                query=dba_query.debug_compiled_query,
                orig=None,
                details={},
            )

    @generic_profiler_async("db-full")  # type: ignore  # TODO: fix
    async def execute(self, query: DBAdapterQuery) -> AsyncRawExecutionResult:
        with self.wrap_execute_excs(query=query, stage="request"):
            if self._redis_cli_acm is not None:
                resp_body = await self._run_query_cached(query)
            else:
                resp_body = await self._run_query(query)

        rd = self._parse_response_body(resp_body, query)

        async def chunk_gen(chunk_size=query.chunk_size or self._default_chunk_size):  # type: ignore  # TODO: fix
            data = rd["rows"]
            while data:
                chunk = data[:chunk_size]
                data = data[chunk_size:]
                yield chunk

        return AsyncRawExecutionResult(
            raw_cursor_info=dict(cols=rd["cols"]),
            raw_chunk_generator=chunk_gen(),
        )

    async def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        return None  # Not Applicable

    async def get_schema_names(self, db_ident: DBIdent) -> list[str]:
        raise NotImplementedError()

    async def get_tables(self, schema_ident: SchemaIdent) -> list[TableIdent]:
        known_general_tables = BITRIX_TABLES_MAP.keys()
        user_tables = await self._get_user_tables(schema_ident)
        tables = [
            table_name
            for table_name in user_tables
            if table_name in known_general_tables or table_name.startswith(SMART_PROCESS_TABLE_PREFIX)
        ]
        return [
            TableIdent(
                db_name=schema_ident.db_name,
                schema_name=schema_ident.schema_name,
                table_name=table_name,
            )
            for table_name in tables
        ]

    async def _get_user_tables(self, schema_ident: SchemaIdent) -> list[str]:
        body: dict[str, Any] = {
            "key": self._target_dto.token,
        }
        api_url: str = f"https://{self._target_dto.portal}/bitrix/tools/biconnector/gds.php?show_tables"
        resp = await self._session.post(
            url=api_url,
            json=body,
        )
        tables: list[str] = [table[0] for table in await resp.json()]
        return tables

    @generic_profiler_async("db-table-info")  # type: ignore  # TODO: fix
    async def get_table_info(
        self, table_def: Optional[TableDefinition] = None, fetch_idx_info: bool = False
    ) -> RawSchemaInfo:
        assert isinstance(table_def, TableIdent)
        table_name = table_def.table_name

        assert table_name in [table.table_name for table in await self.get_tables(SchemaIdent(DEFAULT_DB, None))]
        bitrix_table = self._table_schema(table_name)
        columns_type = bitrix_table.get_columns_type()
        query = sa.select(["*"]).select_from(sa.table(table_name))
        if bitrix_table.daterange_col_name is not None:
            query = query.where(
                sa.column(bitrix_table.daterange_col_name) == "2000-01-01",
            )
        query_obj = DBAdapterQuery(
            query=query,
        )
        res = await self.execute(query_obj)
        res_cols = res.raw_cursor_info["cols"]

        return RawSchemaInfo(
            columns=tuple(
                RawColumnInfo(
                    name=col["id"],
                    title=col["id"],
                    nullable=True,
                    native_type=CommonNativeType(
                        conn_type=self.conn_type,
                        name=columns_type.get(col["id"], "string"),
                        nullable=True,
                    ),
                )
                for col in res_cols
            )
        )

    async def is_table_exists(self, table_ident: TableIdent) -> bool:
        return table_ident in await self.get_tables(SchemaIdent(DEFAULT_DB, None))
        # db_name and schema_name from SchemaIdent is used in method get_tables.
        # Later method get_tables is used in get_parameter_combination, which sets
        # db_name as 'default' and doesn't actually use schema_name

    async def test(self) -> None:
        tables = await self.get_tables(SchemaIdent(DEFAULT_DB, None))
        table_name = tables[0].table_name
        table = self._table_schema(table_name)
        query_obj = DBAdapterQuery(
            query=sa.select(["*"])
            .select_from(sa.table(table_name))
            .where(
                sa.column(table.daterange_col_name) == "2000-01-01",
            ),
        )
        await self.execute(query_obj)

    def generate_body(self, dba_query: DBAdapterQuery) -> tuple[dict[str, Any], dict[str, str]]:
        assert self.table is not None
        body: dict[str, Any] = {
            "key": self._target_dto.token,
            "dateRange": {},
        }

        def date_converter(value: Any) -> Any:
            if isinstance(value, (datetime.datetime, datetime.date)):
                return value.isoformat()
            return value

        def build_date_range(body: dict[str, Any], clause: sa.sql.expression.BinaryExpression) -> dict[str, Any]:
            label: str = ""
            if isinstance(clause.left, sa.sql.elements.ColumnClause):
                label = clause.left.name
            elif isinstance(clause.left, sa.sql.elements.Cast):
                label = str(clause.left.anon_label)
            col_name = label.split(".")[-1].replace('"', "").replace("`", "")
            body["configParams"] = {"timeFilterColumn": col_name}
            op = clause.operator.__name__
            if op == "eq":
                body["dateRange"]["startDate"] = date_converter(clause.right.effective_value)
                body["dateRange"]["endDate"] = date_converter(clause.right.effective_value)
            elif op == "between_op":
                body["dateRange"]["startDate"] = date_converter(clause.right.clauses[0].effective_value)
                body["dateRange"]["endDate"] = date_converter(clause.right.clauses[1].effective_value)
            # a bold assumption
            elif op in ("gt", "ge"):
                body["dateRange"]["startDate"] = date_converter(clause.right.effective_value)
            elif op in ("lt", "le"):
                body["dateRange"]["endDate"] = date_converter(clause.right.effective_value)
            return body

        assert isinstance(dba_query.query, sa.sql.Select)
        sa_whereclause = dba_query.query.whereclause
        if sa_whereclause is not None:
            if isinstance(sa_whereclause, sa.sql.expression.BooleanClauseList):
                for clause in sa_whereclause.clauses:
                    body = build_date_range(body, clause)
                    # getting just first datetime filtration
                    break
            elif isinstance(sa_whereclause, sa.sql.expression.BinaryExpression):
                clause = sa_whereclause
                body = build_date_range(body, clause)

        flatten_body = {
            "key": body["key"],
            "startDate": body["dateRange"].get("startDate"),
            "endDate": body["dateRange"].get("endDate"),
            "timeFilterColumn": body.get("configParams", dict()).get("timeFilterColumn"),
        }

        return body, flatten_body

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/adapter.py
datalens-tech/datalens-backend
from contextlib import asynccontextmanager
import gzip
import logging
from typing import (
    Any,
    AsyncGenerator,
    Optional,
)

import attr
import redis.asyncio
from redis_cache_lock.types import TClientACM

from dl_app_tools.profiling_base import GenericProfiler
from dl_constants.types import TJSONExt
from dl_core.aio.web_app_services.redis import RedisConnParams
from dl_core.data_processing.cache.primitives import LocalKeyRepresentation
from dl_core.serialization import (
    common_dumps,
    common_loads,
)


LOGGER = logging.getLogger(__name__)


def build_local_key_rep(portal: str, table: str, body: dict) -> LocalKeyRepresentation:
    local_key_rep = LocalKeyRepresentation()
    local_key_rep = local_key_rep.extend(part_type="portal", part_content=portal)
    local_key_rep = local_key_rep.extend(part_type="table", part_content=table)
    local_key_rep = local_key_rep.extend(part_type="body", part_content=frozenset(body.items()))

    return local_key_rep


def make_simple_cli_acm(conn_params: RedisConnParams) -> TClientACM:
    @asynccontextmanager
    async def cli_acm(**_: Any) -> AsyncGenerator[redis.asyncio.Redis, None]:
        rcli: redis.asyncio.Redis = redis.asyncio.Redis(**attr.asdict(conn_params))
        try:
            yield rcli
        finally:
            await rcli.connection_pool.disconnect()

    return cli_acm


def get_redis_cli_acm_from_params(redis_conn_params: Optional[RedisConnParams]) -> Optional[TClientACM]:
    if redis_conn_params is None:
        return None
    return make_simple_cli_acm(conn_params=redis_conn_params)


def bitrix_cache_serializer(data: TJSONExt) -> bytes:
    with GenericProfiler("qcache-serialize"):
        serialized_result_data = common_dumps(data)
    with GenericProfiler("qcache-compress"):
        result_data = gzip.compress(serialized_result_data)
    return result_data


def bitrix_cache_deserializer(data_repr: bytes) -> TJSONExt:
    with GenericProfiler("qcache-decompress"):
        encoded_data = gzip.decompress(data_repr)
    with GenericProfiler("qcache-deserialize"):
        data = common_loads(encoded_data)
    return data

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/caches.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Optional,
    Sequence,
)

import attr

from dl_core.aio.web_app_services.redis import RedisConnParams
from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor

from dl_connector_bitrix_gds.core.adapter import BitrixGDSDefaultAdapter
from dl_connector_bitrix_gds.core.target_dto import BitrixGDSConnTargetDTO


if TYPE_CHECKING:
    from dl_connector_bitrix_gds.core.dto import BitrixGDSConnDTO


@attr.s(cmp=False, hash=False)
class BitrixGDSAsyncAdapterConnExecutor(DefaultSqlAlchemyConnExecutor[BitrixGDSDefaultAdapter]):
    TARGET_ADAPTER_CLS = BitrixGDSDefaultAdapter

    _conn_dto: BitrixGDSConnDTO = attr.ib()

    async def _make_target_conn_dto_pool(self) -> Sequence[BitrixGDSConnTargetDTO]:
        redis_conn_params: Optional[RedisConnParams] = None
        caches_ttl: Optional[int] = None

        assert self._services_registry is not None
        rqe_caches_setting = self._services_registry.get_rqe_caches_settings()
        if rqe_caches_setting is not None:
            assert rqe_caches_setting.redis_settings is not None
            redis_conn_params = RedisConnParams(
                host=rqe_caches_setting.redis_settings.HOSTS[0],
                port=rqe_caches_setting.redis_settings.PORT,
                db=rqe_caches_setting.redis_settings.DB,
                password=rqe_caches_setting.redis_settings.PASSWORD,
                ssl=rqe_caches_setting.redis_settings.SSL,
            )
            caches_ttl = rqe_caches_setting.caches_ttl

        conn_params: Optional[dict]
        if isinstance(redis_conn_params, RedisConnParams):
            conn_params = attr.asdict(redis_conn_params)
        else:
            conn_params = None
        return [
            BitrixGDSConnTargetDTO(
                conn_id=self._conn_dto.conn_id,
                pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                portal=self._conn_dto.portal,
                token=self._conn_dto.token,
                max_execution_time=self._conn_options.max_execution_time,  # type: ignore  # TODO: fix
                total_timeout=self._conn_options.total_timeout,  # type: ignore  # TODO: fix
                connect_timeout=self._conn_options.connect_timeout,  # type: ignore  # TODO: fix
                redis_conn_params=conn_params,
                redis_caches_ttl=caches_ttl,
            )
        ]

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/connection_executors.py
datalens-tech/datalens-backend
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
    CoreSourceDefinition,
)
from dl_core.data_source_spec.sql import StandardSQLDataSourceSpec
from dl_core.us_manager.storage_schemas.data_source_spec_base import SQLDataSourceSpecStorageSchema

from dl_connector_bitrix_gds.core.adapter import BitrixGDSDefaultAdapter
from dl_connector_bitrix_gds.core.connection_executors import BitrixGDSAsyncAdapterConnExecutor
from dl_connector_bitrix_gds.core.constants import (
    BACKEND_TYPE_BITRIX_GDS,
    CONNECTION_TYPE_BITRIX24,
    SOURCE_TYPE_BITRIX_GDS,
)
from dl_connector_bitrix_gds.core.data_source import BitrixGDSDataSource
from dl_connector_bitrix_gds.core.storage_schemas.connection import BitrixGDSConnectionDataStorageSchema
from dl_connector_bitrix_gds.core.type_transformer import BitrixGDSTypeTransformer
from dl_connector_bitrix_gds.core.us_connection import BitrixGDSConnection


class BitrixGDSCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_BITRIX24
    connection_cls = BitrixGDSConnection
    us_storage_schema_cls = BitrixGDSConnectionDataStorageSchema
    type_transformer_cls = BitrixGDSTypeTransformer
    sync_conn_executor_cls = BitrixGDSAsyncAdapterConnExecutor
    async_conn_executor_cls = BitrixGDSAsyncAdapterConnExecutor
    dialect_string = "bi_bitrix"


class BitrixGDSCoreSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_BITRIX_GDS
    source_cls = BitrixGDSDataSource
    source_spec_cls = StandardSQLDataSourceSpec
    us_storage_schema_cls = SQLDataSourceSpecStorageSchema


class BitrixGDSCoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_BITRIX_GDS
    connection_definitions = (BitrixGDSCoreConnectionDefinition,)
    source_definitions = (BitrixGDSCoreSourceDefinition,)
    rqe_adapter_classes = frozenset({BitrixGDSDefaultAdapter})

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/connector.py
datalens-tech/datalens-backend
from enum import (
    Enum,
    unique,
)

from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
    SourceBackendType,
)


BACKEND_TYPE_BITRIX_GDS = SourceBackendType.declare("BITRIX_GDS")
CONNECTION_TYPE_BITRIX24 = ConnectionType.declare("bitrix24")
SOURCE_TYPE_BITRIX_GDS = DataSourceType.declare("BITRIX_GDS")

DEFAULT_DB = "default"


@unique
class BitrixGDSTableType(Enum):
    crm_deal = "crm_deal"
    crm_lead = "crm_lead"
    crm_company = "crm_company"
    crm_contact = "crm_contact"
    crm_deal_stage_history = "crm_deal_stage_history"
    crm_lead_status_history = "crm_lead_status_history"
    socialnetwork_group = "socialnetwork_group"
    telephony_call = "telephony_call"
    crm_activity = "crm_activity"
    crm_lead_uf = "crm_lead_uf"
    crm_deal_uf = "crm_deal_uf"
    crm_lead_product_row = "crm_lead_product_row"
    crm_deal_product_row = "crm_deal_product_row"
    crm_dynamic_items = "crm_dynamic_items"
    user = "user"
    crm_company_uf = "crm_company_uf"
    crm_contact_uf = "crm_contact_uf"

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    ClassVar,
    Optional,
)

from dl_constants.enums import (
    DataSourceType,
    JoinType,
)
from dl_core.connection_models import (
    TableDefinition,
    TableIdent,
)
from dl_core.data_source.sql import PseudoSQLDataSource

from dl_connector_bitrix_gds.core.constants import (
    CONNECTION_TYPE_BITRIX24,
    SOURCE_TYPE_BITRIX_GDS,
)


class BitrixGDSDataSource(PseudoSQLDataSource):
    supported_join_types: ClassVar[frozenset[JoinType]] = frozenset()
    conn_type = CONNECTION_TYPE_BITRIX24

    @property
    def db_version(self) -> Optional[str]:
        return None

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type == SOURCE_TYPE_BITRIX_GDS

    def get_table_definition(self) -> TableDefinition:
        assert self.table_name is not None
        return TableIdent(db_name=self.db_name, schema_name=None, table_name=self.table_name)

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/data_source.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_core.connection_models.dto_defs import ConnDTO
from dl_core.utils import secrepr

from dl_connector_bitrix_gds.core.constants import CONNECTION_TYPE_BITRIX24


@attr.s(frozen=True)
class BitrixGDSConnDTO(ConnDTO):
    conn_type = CONNECTION_TYPE_BITRIX24

    portal: str = attr.ib(kw_only=True)
    token: str = attr.ib(kw_only=True, repr=secrepr)

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/dto.py
datalens-tech/datalens-backend
from __future__ import annotations

import json

import dl_core.connectors.base.error_transformer as error_transformer
from dl_core.connectors.base.error_transformer import DbErrorTransformer
from dl_core.connectors.base.error_transformer import ErrorTransformerRule as Rule
from dl_core.connectors.base.error_transformer import orig_exc_is
import dl_core.exc as exc


bitrix_error_transformer: DbErrorTransformer = error_transformer.make_default_transformer_with_custom_rules(
    Rule(
        when=orig_exc_is(orig_exc_cls=json.JSONDecodeError),
        then_raise=exc.SourceResponseError,
    )
)

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/error_transformer.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_core.us_manager.storage_schemas.connection import (
    CacheableConnectionDataSchemaMixin,
    ConnectionBaseDataStorageSchema,
)

from dl_connector_bitrix_gds.core.us_connection import BitrixGDSConnection


class BitrixGDSConnectionDataStorageSchema(
    CacheableConnectionDataSchemaMixin,
    ConnectionBaseDataStorageSchema[BitrixGDSConnection.DataModel],
):
    TARGET_CLS = BitrixGDSConnection.DataModel

    portal = ma_fields.String(required=True, allow_none=False)
    token = ma_fields.String(required=True, allow_none=False)

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/storage_schemas/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Iterable,
    Optional,
)

import attr

from dl_connector_bitrix_gds.core.constants import BitrixGDSTableType


CRM_DEAL = [
    {"name": "ID", "type": "integer"},
    {"name": "DATE_CREATE", "type": "datetime"},
    {"name": "DATE_MODIFY", "type": "datetime"},
    {"name": "CREATED_BY_ID", "type": "integer"},
    {"name": "CREATED_BY_NAME", "type": "string"},
    {"name": "CREATED_BY", "type": "string"},
    {"name": "MODIFY_BY_ID", "type": "integer"},
    {"name": "MODIFIED_BY_NAME", "type": "string"},
    {"name": "MODIFIED_BY", "type": "string"},
    {"name": "ASSIGNED_BY_ID", "type": "integer"},
    {"name": "ASSIGNED_BY_NAME", "type": "string"},
    {"name": "ASSIGNED_BY", "type": "string"},
    {"name": "ASSIGNED_BY_DEPARTMENT", "type": "string"},
    {"name": "OPENED", "type": "string"},
    {"name": "LEAD_ID", "type": "integer"},
    {"name": "COMPANY_ID", "type": "integer"},
    {"name": "COMPANY_NAME", "type": "string"},
    {"name": "COMPANY", "type": "string"},
    {"name": "CONTACT_ID", "type": "string"},
    {"name": "CONTACT_NAME", "type": "string"},
    {"name": "CONTACT", "type": "string"},
    {"name": "TITLE", "type": "string"},
    {"name": "CRM_PRODUCT", "type": "string"},
    {"name": "CRM_PRODUCT_COUNT", "type": "string"},
    {"name": "CATEGORY_ID", "type": "string"},
    {"name": "CATEGORY_NAME", "type": "string"},
    {"name": "CATEGORY", "type": "string"},
    {"name": "STAGE_ID", "type": "string"},
    {"name": "STAGE_NAME", "type": "string"},
    {"name": "STAGE", "type": "string"},
    {"name": "STAGE_SEMANTIC_ID", "type": "string"},
    {"name": "STAGE_SEMANTIC", "type": "string"},
    {"name": "IS_NEW", "type": "string"},
    {"name": "IS_RECURRING", "type": "string"},
    {"name": "IS_RETURN_CUSTOMER", "type": "string"},
    {"name": "CLOSED", "type": "string"},
    {"name": "TYPE_ID", "type": "string"},
    {"name": "OPPORTUNITY", "type": "string"},
    {"name": "IS_MANUAL_OPPORTUNITY", "type": "string"},
    {"name": "TAX_VALUE", "type": "string"},
    {"name": "CURRENCY_ID", "type": "string"},
    {"name": "OPPORTUNITY_ACCOUNT", "type": "string"},
    {"name": "TAX_VALUE_ACCOUNT", "type": "string"},
    {"name": "ACCOUNT_CURRENCY_ID", "type": "string"},
    {"name": "PROBABILITY", "type": "string"},
    {"name": "COMMENTS", "type": "string"},
    {"name": "BEGINDATE", "type": "datetime"},
    {"name": "CLOSEDATE", "type": "datetime"},
    {"name": "LOCATION_ID", "type": "string"},
    {"name": "SOURCE_ID", "type": "string"},
    {"name": "SOURCE_NAME", "type": "string"},
    {"name": "SOURCE", "type": "string"},
    {"name": "SOURCE_DESCRIPTION", "type": "string"},
    {"name": "ORIGINATOR_ID", "type": "string"},
    {"name": "ORIGIN_ID", "type": "string"},
    {"name": "ADDITIONAL_INFO", "type": "string"},
    {"name": "UTM_SOURCE", "type": "string"},
    {"name": "UTM_MEDIUM", "type": "string"},
    {"name": "UTM_CAMPAIGN", "type": "string"},
    {"name": "UTM_CONTENT", "type": "string"},
    {"name": "UTM_TERM", "type": "string"},
    {"name": "BANK_DETAIL_ID", "type": "string"},
]


CRM_LEAD = [
    {"name": "ID", "type": "integer"},
    {"name": "DATE_MODIFY", "type": "datetime"},
    {"name": "DATE_CREATE", "type": "datetime"},
    {"name": "CREATED_BY_ID", "type": "integer"},
    {"name": "CREATED_BY", "type": "string"},
    {"name": "MODIFY_BY_ID", "type": "integer"},
    {"name": "MODIFIED_BY_NAME", "type": "string"},
    {"name": "MODIFIED_BY", "type": "string"},
    {"name": "ASSIGNED_BY_ID", "type": "integer"},
    {"name": "ASSIGNED_BY_NAME", "type": "string"},
    {"name": "ASSIGNED_BY", "type": "string"},
    {"name": "ASSIGNED_BY_DEPARTMENT", "type": "string"},
    {"name": "OPENED", "type": "string"},
    {"name": "COMPANY_ID", "type": "integer"},
    {"name": "COMPANY_NAME", "type": "string"},
    {"name": "COMPANY", "type": "string"},
    {"name": "CONTACT_ID", "type": "string"},
    {"name": "CONTACT_NAME", "type": "string"},
    {"name": "CONTACT", "type": "string"},
    {"name": "STATUS_ID", "type": "string"},
    {"name": "STATUS_NAME", "type": "string"},
    {"name": "STATUS", "type": "string"},
    {"name": "STATUS_DESCRIPTION", "type": "string"},
    {"name": "STATUS_SEMANTIC_ID", "type": "string"},
    {"name": "STATUS_SEMANTIC", "type": "string"},
    {"name": "CRM_PRODUCT_ID", "type": "string"},
    {"name": "CRM_PRODUCT_NAME", "type": "string"},
    {"name": "CRM_PRODUCT", "type": "string"},
    {"name": "OPPORTUNITY", "type": "string"},
    {"name": "CURRENCY_ID", "type": "string"},
    {"name": "SOURCE_ID", "type": "string"},
    {"name": "SOURCE_NAME", "type": "string"},
    {"name": "SOURCE", "type": "string"},
    {"name": "SOURCE_DESCRIPTION", "type": "string"},
    {"name": "TITLE", "type": "string"},
    {"name": "FULL_NAME", "type": "string"},
    {"name": "NAME", "type": "string"},
    {"name": "LAST_NAME", "type": "string"},
    {"name": "SECOND_NAME", "type": "string"},
    {"name": "COMPANY_TITLE", "type": "string"},
    {"name": "POST", "type": "string"},
    {"name": "ADDRESS_1", "type": "string"},
    {"name": "ADDRESS_2", "type": "string"},
    {"name": "ADDRESS_CITY", "type": "string"},
    {"name": "ADDRESS_POSTAL_CODE", "type": "string"},
    {"name": "ADDRESS_REGION", "type": "string"},
    {"name": "ADDRESS_PROVINCE", "type": "string"},
    {"name": "ADDRESS_COUNTRY", "type": "string"},
    {"name": "ADDRESS_COUNTRY_CODE", "type": "string"},
    {"name": "COMMENTS", "type": "string"},
    {"name": "ORIGINATOR_ID", "type": "string"},
    {"name": "ORIGIN_ID", "type": "string"},
    {"name": "DATE_CLOSED", "type": "datetime"},
    {"name": "BIRTHDATE", "type": "datetime"},
    {"name": "HONORIFIC", "type": "string"},
    {"name": "UTM_SOURCE", "type": "string"},
    {"name": "UTM_MEDIUM", "type": "string"},
    {"name": "UTM_CAMPAIGN", "type": "string"},
    {"name": "UTM_CONTENT", "type": "string"},
    {"name": "UTM_TERM", "type": "string"},
    {"name": "PHONE", "type": "string"},
    {"name": "WEB", "type": "string"},
    {"name": "EMAIL", "type": "string"},
    {"name": "IM", "type": "string"},
]


CRM_COMPANY = [
    {"name": "ID", "type": "integer"},
    {"name": "DATE_CREATE", "type": "datetime"},
    {"name": "DATE_MODIFY", "type": "datetime"},
    {"name": "CREATED_BY_ID", "type": "integer"},
    {"name": "CREATED_BY_NAME", "type": "string"},
    {"name": "CREATED_BY", "type": "string"},
    {"name": "MODIFY_BY_ID", "type": "integer"},
    {"name": "MODIFIED_BY_NAME", "type": "string"},
    {"name": "MODIFIED_BY", "type": "string"},
    {"name": "ASSIGNED_BY_ID", "type": "integer"},
    {"name": "ASSIGNED_BY_NAME", "type": "string"},
    {"name": "ASSIGNED_BY", "type": "string"},
    {"name": "OPENED", "type": "string"},
    {"name": "TITLE", "type": "string"},
    {"name": "ADDRESS_1", "type": "string"},
    {"name": "ADDRESS_2", "type": "string"},
    {"name": "ADDRESS_CITY", "type": "string"},
    {"name": "ADDRESS_POSTAL_CODE", "type": "string"},
    {"name": "ADDRESS_REGION", "type": "string"},
    {"name": "ADDRESS_PROVINCE", "type": "string"},
    {"name": "ADDRESS_COUNTRY", "type": "string"},
    {"name": "ADDRESS_COUNTRY_CODE", "type": "string"},
    {"name": "BANKING_DETAILS", "type": "string"},
    {"name": "COMMENTS", "type": "string"},
    {"name": "COMPANY_TYPE_ID", "type": "string"},
    {"name": "COMPANY_TYPE_NAME", "type": "string"},
    {"name": "COMPANY_TYPE", "type": "string"},
    {"name": "INDUSTRY_ID", "type": "string"},
    {"name": "INDUSTRY_NAME", "type": "string"},
    {"name": "INDUSTRY", "type": "string"},
    {"name": "REVENUE", "type": "float"},
    {"name": "CURRENCY_ID", "type": "string"},
    {"name": "EMPLOYEES", "type": "string"},
    {"name": "LEAD_ID", "type": "integer"},
    {"name": "ORIGINATOR_ID", "type": "string"},
    {"name": "ORIGIN_ID", "type": "string"},
    {"name": "ORIGIN_VERSION", "type": "string"},
    {"name": "IS_MY_COMPANY", "type": "string"},
    {"name": "UTM_SOURCE", "type": "string"},
    {"name": "UTM_MEDIUM", "type": "string"},
    {"name": "UTM_CAMPAIGN", "type": "string"},
    {"name": "UTM_CONTENT", "type": "string"},
    {"name": "UTM_TERM", "type": "string"},
    {"name": "PHONE", "type": "string"},
    {"name": "WEB", "type": "string"},
    {"name": "EMAIL", "type": "string"},
    {"name": "IM", "type": "string"},
]


CRM_CONTACT = [
    {"name": "ID", "type": "integer"},
    {"name": "DATE_CREATE", "type": "datetime"},
    {"name": "DATE_MODIFY", "type": "datetime"},
    {"name": "CREATED_BY_ID", "type": "integer"},
    {"name": "CREATED_BY_NAME", "type": "string"},
    {"name": "CREATED_BY", "type": "string"},
    {"name": "MODIFY_BY_ID", "type": "integer"},
    {"name": "MODIFIED_BY_NAME", "type": "string"},
    {"name": "MODIFIED_BY", "type": "string"},
    {"name": "ASSIGNED_BY_ID", "type": "integer"},
    {"name": "ASSIGNED_BY_NAME", "type": "string"},
    {"name": "ASSIGNED_BY", "type": "string"},
    {"name": "OPENED", "type": "string"},
    {"name": "COMPANY_ID", "type": "integer"},
    {"name": "SOURCE_ID", "type": "string"},
    {"name": "SOURCE_DESCRIPTION", "type": "string"},
    {"name": "NAME", "type": "string"},
    {"name": "LAST_NAME", "type": "string"},
    {"name": "SECOND_NAME", "type": "string"},
    {"name": "POST", "type": "string"},
    {"name": "ADDRESS_1", "type": "string"},
    {"name": "ADDRESS_2", "type": "string"},
    {"name": "ADDRESS_CITY", "type": "string"},
    {"name": "ADDRESS_POSTAL_CODE", "type": "string"},
    {"name": "ADDRESS_REGION", "type": "string"},
    {"name": "ADDRESS_PROVINCE", "type": "string"},
    {"name": "ADDRESS_COUNTRY", "type": "string"},
    {"name": "ADDRESS_COUNTRY_CODE", "type": "string"},
    {"name": "COMMENTS", "type": "string"},
    {"name": "LEAD_ID", "type": "integer"},
    {"name": "EXPORT", "type": "string"},
    {"name": "TYPE_ID", "type": "string"},
    {"name": "ORIGINATOR_ID", "type": "string"},
    {"name": "ORIGIN_ID", "type": "string"},
    {"name": "ORIGIN_VERSION", "type": "string"},
    {"name": "BIRTHDATE", "type": "date"},
    {"name": "HONORIFIC", "type": "string"},
    {"name": "FACE_ID", "type": "string"},
    {"name": "PHONE", "type": "string"},
    {"name": "WEB", "type": "string"},
    {"name": "EMAIL", "type": "string"},
    {"name": "IM", "type": "string"},
    {"name": "UTM_SOURCE", "type": "string"},
    {"name": "UTM_MEDIUM", "type": "string"},
    {"name": "UTM_CAMPAIGN", "type": "string"},
    {"name": "UTM_CONTENT", "type": "string"},
    {"name": "UTM_TERM", "type": "string"},
]


CRM_DEAL_STAGE_HISTORY = [
    {"name": "ID", "type": "integer"},
    {"name": "TYPE_ID", "type": "integer"},
    {"name": "DEAL_ID", "type": "integer"},
    {"name": "DATE_CREATE", "type": "datetime"},
    {"name": "START_DATE", "type": "string"},
    {"name": "END_DATE", "type": "string"},
    {"name": "ASSIGNED_BY_ID", "type": "integer"},
    {"name": "ASSIGNED_BY_NAME", "type": "string"},
    {"name": "ASSIGNED_BY", "type": "string"},
    {"name": "ASSIGNED_BY_DEPARTMENT", "type": "string"},
    {"name": "STAGE_SEMANTIC_ID", "type": "string"},
    {"name": "STAGE_SEMANTIC", "type": "string"},
    {"name": "STAGE_ID", "type": "string"},
    {"name": "STAGE_NAME", "type": "string"},
    {"name": "STAGE", "type": "string"},
]


CRM_LEAD_STATUS_HISTORY = [
    {"name": "ID", "type": "integer"},
    {"name": "TYPE_ID", "type": "integer"},
    {"name": "LEAD_ID", "type": "integer"},
    {"name": "DATE_CREATE", "type": "datetime"},
    {"name": "ASSIGNED_BY_ID", "type": "integer"},
    {"name": "ASSIGNED_BY_NAME", "type": "string"},
    {"name": "ASSIGNED_BY", "type": "string"},
    {"name": "ASSIGNED_BY_DEPARTMENT", "type": "string"},
    {"name": "STATUS_SEMANTIC_ID", "type": "string"},
    {"name": "STATUS_SEMANTIC", "type": "string"},
    {"name": "STATUS_ID", "type": "string"},
    {"name": "STATUS_NAME", "type": "string"},
    {"name": "STATUS", "type": "string"},
]


SOCIALNETWORK_GROUP = [
    {"name": "ID", "type": "integer"},
    {"name": "SITE_ID", "type": "string"},
    {"name": "NAME", "type": "string"},
    {"name": "DESCRIPTION", "type": "string"},
    {"name": "DATE_CREATE", "type": "datetime"},
    {"name": "DATE_MODIFY", "type": "datetime"},
    {"name": "OPENED", "type": "string"},
    {"name": "SUBJECT", "type": "string"},
    {"name": "OWNER_ID", "type": "integer"},
    {"name": "OWNER_NAME", "type": "string"},
    {"name": "OWNER", "type": "string"},
    {"name": "KEYWORDS", "type": "string"},
    {"name": "NUMBER_OF_MEMBERS", "type": "integer"},
    {"name": "DATE_ACTIVITY", "type": "datetime"},
    {"name": "CLOSED", "type": "string"},
    {"name": "PROJECT", "type": "string"},
]


TELEPHONY_CALL = [
    {"name": "CALL_ID", "type": "string"},
    {"name": "PORTAL_USER_ID", "type": "integer"},
    {"name": "PORTAL_USER", "type": "string"},
    {"name": "PORTAL_USER_DEPARTMENT", "type": "string"},
    {"name": "PORTAL_NUMBER", "type": "string"},
    {"name": "PHONE_NUMBER", "type": "string"},
    {"name": "CALL_TYPE", "type": "integer"},
    {"name": "CALL_DURATION", "type": "integer"},
    {"name": "CALL_START_TIME", "type": "datetime"},
    {"name": "CALL_STATUS_CODE", "type": "string"},
    {"name": "CALL_STATUS_REASON", "type": "string"},
    {"name": "RECORD_FILE_ID", "type": "integer"},
    {"name": "CALL_VOTE", "type": "integer"},
    {"name": "COST", "type": "float"},
    {"name": "COST_CURRENCY", "type": "string"},
    {"name": "CRM_ENTITY_TYPE", "type": "string"},
    {"name": "CRM_ENTITY_ID", "type": "integer"},
    {"name": "CRM_ACTIVITY_ID", "type": "integer"},
    {"name": "REST_APP_ID", "type": "integer"},
    {"name": "REST_APP_NAME", "type": "string"},
    {"name": "TRANSCRIPT_PENDING", "type": "string"},
    {"name": "TRANSCRIPT_ID", "type": "integer"},
    {"name": "REDIAL_ATTEMPT", "type": "integer"},
    {"name": "COMMENT", "type": "string"},
]


CRM_ACTIVITY = [
    {"name": "ID", "type": "integer"},
    {"name": "TYPE_ID", "type": "integer"},
    {"name": "TYPE_NAME", "type": "string"},
    {"name": "OWNER_ID", "type": "integer"},
    {"name": "OWNER_TYPE_ID", "type": "integer"},
    {"name": "OWNER_TYPE_NAME", "type": "string"},
    {"name": "ASSOCIATED_ENTITY_ID", "type": "integer"},
    {"name": "SUBJECT", "type": "string"},
    {"name": "COMPLETED", "type": "string"},
    {"name": "STATUS_ID", "type": "integer"},
    {"name": "STATUS_NAME", "type": "string"},
    {"name": "RESPONSIBLE_ID", "type": "integer"},
    {"name": "RESPONSIBLE_NAME", "type": "string"},
    {"name": "RESPONSIBLE", "type": "string"},
    {"name": "PRIORITY_ID", "type": "integer"},
    {"name": "PRIORITY_NAME", "type": "string"},
    {"name": "DESCRIPTION", "type": "string"},
    {"name": "DESCRIPTION_TYPE_ID", "type": "integer"},
    {"name": "DESCRIPTION_TYPE_NAME", "type": "string"},
    {"name": "DIRECTION_ID", "type": "integer"},
    {"name": "DIRECTION_NAME", "type": "string"},
    {"name": "DATE_CREATE", "type": "datetime"},
    {"name": "DATE_MODIFY", "type": "datetime"},
    {"name": "START_TIME", "type": "datetime"},
    {"name": "END_TIME", "type": "datetime"},
    {"name": "DEADLINE", "type": "datetime"},
    {"name": "ORIGINATOR_ID", "type": "string"},
    {"name": "ORIGIN_ID", "type": "string"},
    {"name": "AUTHOR_ID", "type": "integer"},
    {"name": "AUTHOR_NAME", "type": "string"},
    {"name": "AUTHOR", "type": "string"},
    {"name": "EDITOR_ID", "type": "integer"},
    {"name": "EDITOR_NAME", "type": "string"},
    {"name": "EDITOR", "type": "string"},
]


CRM_LEAD_UF = [
    {"name": "LEAD_ID", "type": "integer"},
    {"name": "DATE_MODIFY", "type": "datetime"},
    {"name": "DATE_CREATE", "type": "datetime"},
    {"name": "DATE_CLOSED", "type": "datetime"},
]


CRM_DEAL_UF = [
    {"name": "DEAL_ID", "type": "integer"},
    {"name": "DATE_CREATE", "type": "datetime"},
    {"name": "CLOSEDATE", "type": "datetime"},
]


CRM_LEAD_PRODUCT_ROW = [
    {"name": "ID", "type": "integer"},
    {"name": "LEAD_ID", "type": "integer"},
    {"name": "LEAD_DATE_MODIFY", "type": "datetime"},
    {"name": "LEAD_DATE_CREATE", "type": "datetime"},
    {"name": "LEAD_DATE_CLOSED", "type": "datetime"},
    {"name": "PRODUCT", "type": "string"},
    {"name": "PRODUCT_ID", "type": "integer"},
    {"name": "PRODUCT_NAME", "type": "string"},
    {"name": "PRICE", "type": "string"},
    {"name": "PRICE_EXCLUSIVE", "type": "string"},
    {"name": "PRICE_NETTO", "type": "string"},
    {"name": "PRICE_BRUTTO", "type": "string"},
    {"name": "QUANTITY", "type": "string"},
    {"name": "DISCOUNT_TYPE", "type": "string"},
    {"name": "DISCOUNT_TYPE_ID", "type": "integer"},
    {"name": "DISCOUNT_TYPE_NAME", "type": "string"},
    {"name": "DISCOUNT_RATE", "type": "string"},
    {"name": "DISCOUNT_SUM", "type": "string"},
    {"name": "TAX_RATE", "type": "string"},
    {"name": "TAX_INCLUDED", "type": "string"},
    {"name": "CUSTOMIZED", "type": "string"},
    {"name": "MEASURE", "type": "string"},
    {"name": "MEASURE_CODE", "type": "integer"},
    {"name": "MEASURE_NAME", "type": "string"},
    {"name": "SORT", "type": "integer"},
]


CRM_DEAL_PRODUCT_ROW = [
    {"name": "ID", "type": "integer"},
    {"name": "DEAL_ID", "type": "integer"},
    {"name": "DEAL_DATE_CREATE", "type": "datetime"},
    {"name": "DEAL_CLOSEDATE", "type": "datetime"},
    {"name": "PRODUCT", "type": "string"},
    {"name": "PRODUCT_ID", "type": "integer"},
    {"name": "PRODUCT_NAME", "type": "string"},
    {"name": "PRICE", "type": "string"},
    {"name": "PRICE_EXCLUSIVE", "type": "string"},
    {"name": "PRICE_NETTO", "type": "string"},
    {"name": "PRICE_BRUTTO", "type": "string"},
    {"name": "QUANTITY", "type": "string"},
    {"name": "DISCOUNT_TYPE", "type": "string"},
    {"name": "DISCOUNT_TYPE_ID", "type": "integer"},
    {"name": "DISCOUNT_TYPE_NAME", "type": "string"},
    {"name": "DISCOUNT_RATE", "type": "string"},
    {"name": "DISCOUNT_SUM", "type": "string"},
    {"name": "TAX_RATE", "type": "string"},
    {"name": "TAX_INCLUDED", "type": "string"},
    {"name": "CUSTOMIZED", "type": "string"},
    {"name": "MEASURE", "type": "string"},
    {"name": "MEASURE_CODE", "type": "integer"},
    {"name": "MEASURE_NAME", "type": "string"},
    {"name": "SORT", "type": "integer"},
]


CRM_DYNAMIC_ITEMS = [
    {"name": "ID", "type": "integer"},
    {"name": "XML_ID", "type": "string"},
    {"name": "TITLE", "type": "string"},
    {"name": "CREATED_BY", "type": "string"},
    {"name": "CREATED_BY_ID", "type": "integer"},
    {"name": "CREATED_BY_NAME", "type": "string"},
    {"name": "UPDATED_BY", "type": "string"},
    {"name": "UPDATED_BY_ID", "type": "integer"},
    {"name": "UPDATED_BY_NAME", "type": "string"},
    {"name": "MOVED_BY", "type": "string"},
    {"name": "MOVED_BY_ID", "type": "integer"},
    {"name": "MOVED_BY_NAME", "type": "string"},
    {"name": "CREATED_TIME", "type": "datetime"},
    {"name": "UPDATED_TIME", "type": "datetime"},
    {"name": "MOVED_TIME", "type": "datetime"},
    {"name": "CATEGORY", "type": "string"},
    {"name": "CATEGORY_ID", "type": "integer"},
    {"name": "CATEGORY_NAME", "type": "string"},
    {"name": "OPENED", "type": "string"},
    {"name": "STAGE", "type": "string"},
    {"name": "STAGE_ID", "type": "string"},
    {"name": "STAGE_NAME", "type": "string"},
    {"name": "PREVIOUS_STAGE_ID", "type": "string"},
    {"name": "BEGINDATE", "type": "datetime"},
    {"name": "CLOSEDATE", "type": "datetime"},
    {"name": "COMPANY", "type": "string"},
    {"name": "COMPANY_ID", "type": "integer"},
    {"name": "COMPANY_NAME", "type": "string"},
    {"name": "CONTACT", "type": "string"},
    {"name": "CONTACT_ID", "type": "integer"},
    {"name": "CONTACT_NAME", "type": "string"},
    {"name": "OPPORTUNITY", "type": "string"},
    {"name": "IS_MANUAL_OPPORTUNITY", "type": "string"},
    {"name": "TAX_VALUE", "type": "string"},
    {"name": "CURRENCY_ID", "type": "string"},
    {"name": "OPPORTUNITY_ACCOUNT", "type": "string"},
    {"name": "TAX_VALUE_ACCOUNT", "type": "string"},
    {"name": "ACCOUNT_CURRENCY_ID", "type": "string"},
    {"name": "MYCOMPANY", "type": "string"},
    {"name": "MYCOMPANY_ID", "type": "integer"},
    {"name": "MYCOMPANY_NAME", "type": "string"},
    {"name": "SOURCE", "type": "string"},
    {"name": "SOURCE_ID", "type": "string"},
    {"name": "SOURCE_NAME", "type": "string"},
    {"name": "SOURCE_DESCRIPTION", "type": "string"},
    {"name": "ASSIGNED_BY", "type": "string"},
    {"name": "ASSIGNED_BY_ID", "type": "integer"},
    {"name": "ASSIGNED_BY_NAME", "type": "string"},
    {"name": "WEBFORM_ID", "type": "integer"},
]

USER = [
    {"name": "ID", "type": "integer"},
    {"name": "ACTIVE", "type": "string"},
    {"name": "NAME", "type": "string"},
    {"name": "DEPARTMENT", "type": "string"},
]

CRM_COMPANY_UF = [
    {"name": "COMPANY_ID", "type": "integer"},
    {"name": "DATE_CREATE", "type": "datetime"},
]

CRM_CONTACT_UF = [
    {"name": "CONTACT_ID", "type": "integer"},
    {"name": "DATE_CREATE", "type": "datetime"},
]


@attr.s
class BitrixGDSColumn:
    name: str = attr.ib()
    type: str = attr.ib()


@attr.s
class BitrixGDSTable:
    type: BitrixGDSTableType = attr.ib()
    schema: Iterable[BitrixGDSColumn] = attr.ib()
    daterange_col_name: Optional[str] = attr.ib(default=None)

    def get_columns(self) -> list[str]:
        return [column.name for column in self.schema]

    def get_columns_type(self) -> dict[str, str]:
        return {column.name: column.type for column in self.schema}


BITRIX_TABLES_MAP = {
    "crm_deal": BitrixGDSTable(
        type=BitrixGDSTableType.crm_deal,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_DEAL],
    ),
    "crm_lead": BitrixGDSTable(
        type=BitrixGDSTableType.crm_lead,
        daterange_col_name="DATE_MODIFY",
        schema=[BitrixGDSColumn(**column) for column in CRM_LEAD],
    ),
    "crm_company": BitrixGDSTable(
        type=BitrixGDSTableType.crm_company,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_COMPANY],
    ),
    "crm_contact": BitrixGDSTable(
        type=BitrixGDSTableType.crm_contact,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_CONTACT],
    ),
    "crm_deal_stage_history": BitrixGDSTable(
        type=BitrixGDSTableType.crm_deal_stage_history,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_DEAL_STAGE_HISTORY],
    ),
    "crm_lead_status_history": BitrixGDSTable(
        type=BitrixGDSTableType.crm_lead_status_history,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_LEAD_STATUS_HISTORY],
    ),
    "socialnetwork_group": BitrixGDSTable(
        type=BitrixGDSTableType.socialnetwork_group,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in SOCIALNETWORK_GROUP],
    ),
    "telephony_call": BitrixGDSTable(
        type=BitrixGDSTableType.telephony_call,
        daterange_col_name="CALL_START_TIME",
        schema=[BitrixGDSColumn(**column) for column in TELEPHONY_CALL],
    ),
    "crm_activity": BitrixGDSTable(
        type=BitrixGDSTableType.crm_activity,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_ACTIVITY],
    ),
    "crm_lead_uf": BitrixGDSTable(
        type=BitrixGDSTableType.crm_lead_uf,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_LEAD_UF],
    ),
    "crm_deal_uf": BitrixGDSTable(
        type=BitrixGDSTableType.crm_deal_uf,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_DEAL_UF],
    ),
    "crm_lead_product_row": BitrixGDSTable(
        type=BitrixGDSTableType.crm_lead_product_row,
        daterange_col_name="LEAD_DATE_MODIFY",
        schema=[BitrixGDSColumn(**column) for column in CRM_LEAD_PRODUCT_ROW],
    ),
    "crm_deal_product_row": BitrixGDSTable(
        type=BitrixGDSTableType.crm_deal_product_row,
        daterange_col_name="DEAL_DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_DEAL_PRODUCT_ROW],
    ),
    "user": BitrixGDSTable(
        type=BitrixGDSTableType.user,
        schema=[BitrixGDSColumn(**column) for column in USER],
    ),
    "crm_company_uf": BitrixGDSTable(
        type=BitrixGDSTableType.crm_company_uf,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_COMPANY_UF],
    ),
    "crm_contact_uf": BitrixGDSTable(
        type=BitrixGDSTableType.crm_contact_uf,
        daterange_col_name="DATE_CREATE",
        schema=[BitrixGDSColumn(**column) for column in CRM_CONTACT_UF],
    ),
}

CRM_DYNAMIC_ITEMS_TABLE = BitrixGDSTable(
    type=BitrixGDSTableType.crm_dynamic_items,
    daterange_col_name="CREATED_TIME",
    schema=[BitrixGDSColumn(**column) for column in CRM_DYNAMIC_ITEMS],
)

SMART_PROCESS_TABLE_PREFIX = "crm_dynamic_items_"

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/tables.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO
from dl_core.utils import secrepr


def hide_pass(value: Optional[dict]) -> str:
    if value is None:
        return repr(value)
    if not value:
        return repr(value)
    return repr({k: v for k, v in value.items() if k != "password"})


@attr.s(frozen=True)
class BitrixGDSConnTargetDTO(ConnTargetDTO):
    portal: str = attr.ib(kw_only=True)
    token: str = attr.ib(kw_only=True, repr=secrepr)

    max_execution_time: Optional[int] = attr.ib()
    connect_timeout: Optional[int] = attr.ib()
    total_timeout: Optional[int] = attr.ib()

    redis_conn_params: Optional[dict] = attr.ib(repr=hide_pass)
    redis_caches_ttl: Optional[int] = attr.ib()

    def get_effective_host(self) -> Optional[str]:
        return None  # Not Applicable

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/target_dto.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.enums import UserDataType
from dl_core.db.conversion_base import (
    TypeTransformer,
    make_native_type,
)

from dl_connector_bitrix_gds.core.constants import CONNECTION_TYPE_BITRIX24


class BitrixGDSTypeTransformer(TypeTransformer):
    conn_type = CONNECTION_TYPE_BITRIX24
    native_to_user_map = {
        make_native_type(CONNECTION_TYPE_BITRIX24, "integer"): UserDataType.integer,
        make_native_type(CONNECTION_TYPE_BITRIX24, "float"): UserDataType.float,
        make_native_type(CONNECTION_TYPE_BITRIX24, "string"): UserDataType.string,
        make_native_type(CONNECTION_TYPE_BITRIX24, "date"): UserDataType.date,
        make_native_type(CONNECTION_TYPE_BITRIX24, "datetime"): UserDataType.genericdatetime,
    }
    user_to_native_map = dict(
        [(bi_type, native_type) for native_type, bi_type in native_to_user_map.items()]
        + [
            (UserDataType.datetime, make_native_type(CONNECTION_TYPE_BITRIX24, "datetime")),
        ]
    )

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Callable,
    ClassVar,
    Optional,
)

import attr

from dl_core.base_models import ConnCacheableDataModelMixin
from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.connection_models.conn_options import ConnectOptions
from dl_core.us_connection_base import (
    ConnectionBase,
    DataSourceTemplate,
    ExecutorBasedMixin,
)
from dl_core.utils import secrepr
from dl_utils.utils import DataKey

from dl_connector_bitrix_gds.core.constants import (
    DEFAULT_DB,
    SOURCE_TYPE_BITRIX_GDS,
)
from dl_connector_bitrix_gds.core.dto import BitrixGDSConnDTO


@attr.s(frozen=True, hash=True)
class BitrixGDSConnectOptions(ConnectOptions):
    max_execution_time: Optional[int] = attr.ib(default=None)
    connect_timeout: Optional[int] = attr.ib(default=None)
    total_timeout: Optional[int] = attr.ib(default=None)


class BitrixGDSConnection(ExecutorBasedMixin, ConnectionBase):
    allow_cache: ClassVar[bool] = True

    @attr.s(kw_only=True)
    class DataModel(ConnCacheableDataModelMixin, ConnectionBase.DataModel):
        portal: str = attr.ib()
        token: str = attr.ib(repr=secrepr)

        @classmethod
        def get_secret_keys(cls) -> set[DataKey]:
            return {
                *super().get_secret_keys(),
                DataKey(parts=("token",)),
            }

    @property
    def cache_ttl_sec_override(self) -> Optional[int]:
        return self.data.cache_ttl_sec

    def get_conn_options(self) -> BitrixGDSConnectOptions:
        return super().get_conn_options().to_subclass(BitrixGDSConnectOptions)

    def get_conn_dto(self) -> BitrixGDSConnDTO:
        return BitrixGDSConnDTO(
            conn_id=self.uuid,
            portal=self.data.portal,
            token=self.data.token,
        )

    def get_parameter_combinations(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[dict]:
        return [
            dict(db_name=DEFAULT_DB, table_name=item.table_name)
            for item in self.get_tables(
                conn_executor_factory=conn_executor_factory,
                db_name=DEFAULT_DB,
                schema_name=None,
            )
        ]

    def get_data_source_templates(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[DataSourceTemplate]:
        return [
            DataSourceTemplate(
                title=parameters["table_name"],
                group=[],
                source_type=SOURCE_TYPE_BITRIX_GDS,
                connection_id=self.uuid,  # type: ignore  # TODO: fix
                parameters=parameters,
            )
            for parameters in self.get_parameter_combinations(conn_executor_factory=conn_executor_factory)
        ]

    @property
    def allow_public_usage(self) -> bool:
        return True

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/core/us_connection.py
datalens-tech/datalens-backend
from dl_formula.connectors.base.connector import FormulaConnector
from dl_sqlalchemy_bitrix.base import BitrixDialect as SABitrixDialect

from dl_connector_bitrix_gds.formula.constants import BitrixDialect as BitrixDialectNS
from dl_connector_bitrix_gds.formula.definitions.all import DEFINITIONS


class BitrixGDSFormulaConnector(FormulaConnector):
    dialect_ns_cls = BitrixDialectNS
    dialects = BitrixDialectNS.BITRIX
    default_dialect = BitrixDialectNS.BITRIX
    op_definitions = DEFINITIONS
    sa_dialect = SABitrixDialect()

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/formula/connector.py
datalens-tech/datalens-backend
from dl_formula.core.dialect import (
    DialectName,
    DialectNamespace,
    simple_combo,
)


DIALECT_NAME_BITRIX = DialectName.declare("BITRIX")


class BitrixDialect(DialectNamespace):
    BITRIX = simple_combo(name=DIALECT_NAME_BITRIX)

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/formula/constants.py
datalens-tech/datalens-backend
from dl_connector_bitrix_gds.formula.definitions.functions_datetime import DEFINITIONS_DATETIME
from dl_connector_bitrix_gds.formula.definitions.functions_markup import DEFINITIONS_MARKUP
from dl_connector_bitrix_gds.formula.definitions.functions_string import DEFINITIONS_STRING
from dl_connector_bitrix_gds.formula.definitions.functions_type import DEFINITIONS_TYPE
from dl_connector_bitrix_gds.formula.definitions.operators_binary import DEFINITIONS_BINARY
from dl_connector_bitrix_gds.formula.definitions.operators_ternary import DEFINITIONS_TERNARY


DEFINITIONS = [
    *DEFINITIONS_DATETIME,
    *DEFINITIONS_MARKUP,
    *DEFINITIONS_STRING,
    *DEFINITIONS_TYPE,
    *DEFINITIONS_BINARY,
    *DEFINITIONS_TERNARY,
]

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/formula/definitions/all.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_datetime as base

from dl_connector_bitrix_gds.formula.constants import BitrixDialect as D


DEFINITIONS_DATETIME = [
    # dateadd
    base.FuncDateadd1.for_dialect(D.BITRIX),
    base.FuncDateadd2Unit.for_dialect(D.BITRIX),
    base.FuncDateadd2Number.for_dialect(D.BITRIX),
]

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/formula/definitions/functions_datetime.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_markup as base

from dl_connector_bitrix_gds.formula.constants import BitrixDialect as D


DEFINITIONS_MARKUP = [
    # __str
    base.FuncInternalStrConst.for_dialect(D.BITRIX),
    base.FuncInternalStr.for_dialect(D.BITRIX),
]

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/formula/definitions/functions_markup.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_string as base

from dl_connector_bitrix_gds.formula.constants import BitrixDialect as D


DEFINITIONS_STRING = [
    # concat
    base.ConcatMultiStrConst.for_dialect(D.BITRIX),
]

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/formula/definitions/functions_string.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_type as base

from dl_connector_bitrix_gds.formula.constants import BitrixDialect as D


DEFINITIONS_TYPE = [
    # date
    base.FuncDate1FromDatetime.for_dialect(D.BITRIX),
    base.FuncDate1FromString.for_dialect(D.BITRIX),
    # datetime
    base.FuncDatetime1FromDatetime.for_dialect(D.BITRIX),
    base.FuncDatetime1FromString.for_dialect(D.BITRIX),
    # datetimetz
    base.FuncDatetimeTZConst.for_dialect(D.BITRIX),
    # genericdatetime
    base.FuncGenericDatetime1FromDatetime.for_dialect(D.BITRIX),
    base.FuncGenericDatetime1FromString.for_dialect(D.BITRIX),
]

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/formula/definitions/functions_type.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_binary as base

from dl_connector_bitrix_gds.formula.constants import BitrixDialect as D


DEFINITIONS_BINARY = [
    # !=
    base.BinaryNotEqual.for_dialect(D.BITRIX),
    # <
    base.BinaryLessThan.for_dialect(D.BITRIX),
    # <=
    base.BinaryLessThanOrEqual.for_dialect(D.BITRIX),
    # ==
    base.BinaryEqual.for_dialect(D.BITRIX),
    # >
    base.BinaryGreaterThan.for_dialect(D.BITRIX),
    # >=
    base.BinaryGreaterThanOrEqual.for_dialect(D.BITRIX),
    # _!=
    base.BinaryNotEqualInternal.for_dialect(D.BITRIX),
    # _==
    base.BinaryEqualInternal.for_dialect(D.BITRIX),
    # _dneq
    base.BinaryEqualDenullified.for_dialect(D.BITRIX),
]

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/formula/definitions/operators_binary.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_ternary as base

from dl_connector_bitrix_gds.formula.constants import BitrixDialect as D


DEFINITIONS_TERNARY = [
    # between
    base.TernaryBetween.for_dialect(D.BITRIX),
]

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds/formula/definitions/operators_ternary.py
datalens-tech/datalens-backend
pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/conftest.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_api_lib_testing.data_api_base import (
    DataApiTestParams,
    StandardizedDataApiTestBase,
)
from dl_api_lib_testing.dataset_base import DatasetTestBase
from dl_core_testing.database import (
    CoreDbConfig,
    Db,
)
from dl_core_testing.engine_wrapper import TestingEngineWrapper

from dl_connector_bitrix_gds.core.constants import (
    CONNECTION_TYPE_BITRIX24,
    SOURCE_TYPE_BITRIX_GDS,
)
from dl_connector_bitrix_gds_tests.ext.config import (
    API_TEST_CONFIG,
    BITRIX_PORTALS,
    DB_NAME,
    SMART_TABLE_NAME,
    TABLE_NAME,
)


class BitrixConnectionTestBase(ConnectionTestBase):
    conn_type = CONNECTION_TYPE_BITRIX24

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return ""

    @pytest.fixture(scope="class")
    def db(self, db_config: CoreDbConfig) -> Db:
        engine_wrapper = TestingEngineWrapper(config=db_config.engine_config)
        return Db(config=db_config, engine_wrapper=engine_wrapper)

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def connection_params(self, bitrix_token: str) -> dict:
        return dict(
            portal=BITRIX_PORTALS["default"],
            token=bitrix_token,
        )


class BitrixDatalensConnectionTestBase(BitrixConnectionTestBase):
    @pytest.fixture(scope="class")
    def connection_params(self, bitrix_datalens_token: str) -> dict:
        return dict(
            portal=BITRIX_PORTALS["datalens"],
            token=bitrix_datalens_token,
        )


class BitrixInvalidConnectionTestBase(BitrixConnectionTestBase):
    @pytest.fixture(scope="class")
    def connection_params(self, bitrix_token: str) -> dict:
        return dict(
            portal=BITRIX_PORTALS["invalid"],
            token=bitrix_token,
        )


class BitrixDatasetTestBase(BitrixConnectionTestBase, DatasetTestBase):
    @pytest.fixture(scope="class")
    def dataset_params(self) -> dict:
        return dict(
            source_type=SOURCE_TYPE_BITRIX_GDS.name,
            title=TABLE_NAME,
            parameters=dict(
                db_name=DB_NAME,
                table_name=TABLE_NAME,
            ),
        )


class BitrixSmartTablesDatasetTestBase(BitrixDatalensConnectionTestBase, DatasetTestBase):
    @pytest.fixture(scope="class")
    def dataset_params(self) -> dict:
        return dict(
            source_type=SOURCE_TYPE_BITRIX_GDS.name,
            title=SMART_TABLE_NAME,
            parameters=dict(
                db_name=DB_NAME,
                table_name=SMART_TABLE_NAME,
            ),
        )


class BitrixDataApiTestBase(BitrixDatasetTestBase, StandardizedDataApiTestBase):
    mutation_caches_on = False

    @pytest.fixture(scope="class")
    def data_api_test_params(self) -> DataApiTestParams:
        return DataApiTestParams(
            two_dims=("ASSIGNED_BY_NAME", "ID"),
            summable_field="ID",
            range_field="ID",
            distinct_field="ASSIGNED_BY_NAME",
            date_field="DATE_CREATE",
        )


class BitrixSmartTablesDataApiTestBase(BitrixSmartTablesDatasetTestBase, StandardizedDataApiTestBase):
    mutation_caches_on = False

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/ext/api/base.py
datalens-tech/datalens-backend
import json

from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase
from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite

from dl_connector_bitrix_gds_tests.ext.api.base import (
    BitrixConnectionTestBase,
    BitrixDatalensConnectionTestBase,
    BitrixInvalidConnectionTestBase,
)
from dl_connector_bitrix_gds_tests.ext.config import BITRIX_PORTALS


class TestBitrixConnection(BitrixConnectionTestBase, DefaultConnectorConnectionTestSuite):
    def test_portal_override(self, control_api_sync_client: SyncHttpClientBase, saved_connection_id: str) -> None:
        resp = control_api_sync_client.get(
            url=f"/api/v1/connections/{saved_connection_id}",
        )
        assert resp.status_code == 200, resp.json
        assert resp.json["portal"] == BITRIX_PORTALS["default"], resp.json

        new_portal = BITRIX_PORTALS["datalens"]
        resp = control_api_sync_client.put(
            url=f"/api/v1/connections/{saved_connection_id}",
            content_type="application/json",
            data=json.dumps({"portal": new_portal}),
        )
        assert resp.status_code == 200, resp.json

        resp = control_api_sync_client.get(
            url=f"/api/v1/connections/{saved_connection_id}",
        )
        assert resp.status_code == 200, resp.json
        assert resp.json["portal"] == BITRIX_PORTALS["datalens"], resp.json

        resp = control_api_sync_client.put(
            url=f"/api/v1/connections/{saved_connection_id}",
            content_type="application/json",
            data=json.dumps({"portal": BITRIX_PORTALS["default"]}),
        )
        assert resp.status_code == 200, resp.json


class TestBitrixDatalensConnection(BitrixDatalensConnectionTestBase, DefaultConnectorConnectionTestSuite):
    pass


class TestBitrixInvalidConnection(BitrixInvalidConnectionTestBase, DefaultConnectorConnectionTestSuite):
    def test_test_connection(self, control_api_sync_client: SyncHttpClientBase, saved_connection_id: str) -> None:
        resp = control_api_sync_client.post(
            f"/api/v1/connections/test_connection/{saved_connection_id}",
            content_type="application/json",
            data=json.dumps({}),
        )
        assert resp.status_code == 400, resp.json

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/ext/api/test_connection.py
datalens-tech/datalens-backend
from dl_api_client.dsmaker.api.data_api import SyncHttpDataApiV2
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_testing.connector.data_api_suites import (
    DefaultConnectorDataDistinctTestSuite,
    DefaultConnectorDataGroupByFormulaTestSuite,
    DefaultConnectorDataPreviewTestSuite,
    DefaultConnectorDataRangeTestSuite,
    DefaultConnectorDataResultTestSuite,
)
from dl_constants.enums import WhereClauseOperation
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_bitrix_gds_tests.ext.api.base import (
    BitrixDataApiTestBase,
    BitrixSmartTablesDataApiTestBase,
)


class TestBitrixDataResult(BitrixDataApiTestBase, DefaultConnectorDataResultTestSuite):
    test_params = RegulatedTestParams(
        mark_features_skipped={
            DefaultConnectorDataResultTestSuite.array_support: "Bitrix doesn't support arrays",
        }
    )


class TestBitrixDataGroupBy(BitrixDataApiTestBase, DefaultConnectorDataGroupByFormulaTestSuite):
    pass


class TestBitrixDataRange(BitrixDataApiTestBase, DefaultConnectorDataRangeTestSuite):
    pass


class TestBitrixDataDistinct(BitrixDataApiTestBase, DefaultConnectorDataDistinctTestSuite):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultConnectorDataDistinctTestSuite.test_date_filter_distinct: "Can't create a new table in bitrix tests",
        }
    )


class TestBitrixDataPreview(BitrixDataApiTestBase, DefaultConnectorDataPreviewTestSuite):
    pass


class TestBitrixSmartTablesData(BitrixSmartTablesDataApiTestBase):
    def test_bitrix_string_to_date(
        self,
        saved_dataset: Dataset,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        ds = saved_dataset
        ds.result_schema["Date from string user_field"] = ds.field(formula=f"DATE([UF_CRM_5_1694020695771])")
        self.get_preview(ds, data_api)

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[ds.find_field(title="Date from string user_field")],
            filters=[
                ds.find_field(title="Date from string user_field").filter(
                    op=WhereClauseOperation.BETWEEN,
                    values=["2023-09-05", "2023-09-06"],
                )
            ],
        )
        assert result_resp.status_code == 200, result_resp.json
        assert get_data_rows(result_resp)

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/ext/api/test_data.py
datalens-tech/datalens-backend
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite

from dl_connector_bitrix_gds_tests.ext.api.base import (
    BitrixDatasetTestBase,
    BitrixSmartTablesDatasetTestBase,
)


class TestBitrixDataset(BitrixDatasetTestBase, DefaultConnectorDatasetTestSuite):
    def check_basic_dataset(self, ds: Dataset) -> None:
        assert ds.id
        assert len(ds.result_schema)

        field_names = {field.title for field in ds.result_schema}
        assert {"ID", "DATE_CREATE", "DATE_MODIFY", "ASSIGNED_BY_NAME"}.issubset(field_names)


class TestBitrixSmartTablesDataset(BitrixSmartTablesDatasetTestBase, DefaultConnectorDatasetTestSuite):
    def check_basic_dataset(self, ds: Dataset) -> None:
        assert ds.id
        assert len(ds.result_schema)

        field_names = {field.title for field in ds.result_schema}
        assert {"ID", "UF_CRM_5_1694020695771", "ASSIGNED_BY_NAME"}.issubset(field_names)

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/ext/api/test_dataset.py
datalens-tech/datalens-backend
import abc

from dl_api_client.dsmaker.api.data_api import SyncHttpDataApiV2
from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase
from dl_api_lib_testing.data_api_base import DataApiTestBase
from dl_api_lib_testing.dataset_base import DatasetTestBase

from dl_connector_bitrix_gds.core.constants import SOURCE_TYPE_BITRIX_GDS
from dl_connector_bitrix_gds_tests.ext.api.base import (
    BitrixDatasetTestBase,
    BitrixSmartTablesDatasetTestBase,
)
from dl_connector_bitrix_gds_tests.ext.config import DB_NAME


class BitrixSourcesTestBase(DatasetTestBase, DataApiTestBase, metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def do_check_source(self, source: dict) -> bool:
        """Should return true if and only if we want to test this source in this class"""
        pass

    def test_sources(
        self,
        control_api_sync_client: SyncHttpClientBase,
        control_api: SyncHttpDatasetApiV1,
        saved_connection_id: str,
        data_api: SyncHttpDataApiV2,
    ):
        conn_id = saved_connection_id

        sources_resp = control_api_sync_client.get(f"/api/v1/connections/{conn_id}/info/sources")
        assert sources_resp.status_code == 200, sources_resp.json

        sources_checked = 0
        for source in sources_resp.json["sources"]:
            if not self.do_check_source(source):
                continue

            dataset_params = dict(
                source_type=SOURCE_TYPE_BITRIX_GDS.name,
                title=source["title"],
                parameters=dict(
                    db_name=DB_NAME,
                    table_name=source["title"],
                ),
            )
            ds = self.make_basic_dataset(
                control_api=control_api,
                connection_id=saved_connection_id,
                dataset_params=dataset_params,
            )
            preview_resp = data_api.get_preview(dataset=ds)
            assert preview_resp.status_code == 200, preview_resp.response_errors

            sources_checked += 1

        assert sources_checked > 1


class TestBitrixSources(BitrixDatasetTestBase, BitrixSourcesTestBase):
    def do_check_source(self, source: dict) -> bool:
        title = source["title"]
        return title not in [
            "telephony_call",
            "crm_lead_uf",
            "crm_deal_uf",
            "crm_company_uf",
            "crm_contact_uf",
        ] and not title.startswith("crm_dynamic_items_")


class TestBitrixSmartTablesSources(BitrixSmartTablesDatasetTestBase, BitrixSourcesTestBase):
    def do_check_source(self, source: dict) -> bool:
        return source["title"].startswith("crm_dynamic_items_")

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/ext/api/test_sources.py
datalens-tech/datalens-backend
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport


# Infra settings
CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=51911).host,
    port_us_http=get_test_container_hostport("us", fallback_port=51911).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=51910).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=51910).port,
    us_master_token="AC1ofiek8coB",
    core_connector_ep_names=["bitrix_gds", "postgresql"],
)

COMPENG_URL = f'postgresql://datalens:qwerty@{get_test_container_hostport("db-postgres-13", fallback_port=52301).as_pair()}/test_data'
API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["bitrix_gds", "postgresql"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
    bi_compeng_pg_url=COMPENG_URL,
)

BITRIX_PORTALS = dict(
    default="gds.office.bitrix.ru",
    datalens="datalens.bitrix24.ru",
    invalid="some_portal",
)
DB_NAME = "default"
TABLE_NAME = "crm_deal"
SMART_TABLE_NAME = "crm_dynamic_items_180"

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/ext/config.py
datalens-tech/datalens-backend
import os

import pytest

from dl_api_lib_testing.initialization import initialize_api_lib_test
from dl_testing.env_params.generic import GenericEnvParamGetter

from dl_connector_bitrix_gds_tests.ext.config import API_TEST_CONFIG


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)


@pytest.fixture(scope="session")
def env_param_getter() -> GenericEnvParamGetter:
    filepath = os.path.join(os.path.dirname(__file__), "params.yml")
    return GenericEnvParamGetter.from_yaml_file(filepath)


@pytest.fixture(scope="session")
def bitrix_token(env_param_getter):
    return env_param_getter.get_str_value("BITRIX_TOKEN")


@pytest.fixture(scope="session")
def bitrix_datalens_token(env_param_getter):
    return env_param_getter.get_str_value("BITRIX_DATALENS_TOKEN")

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/ext/conftest.py
datalens-tech/datalens-backend
from dl_api_connector.i18n.localizer import CONFIGS as BI_API_CONNECTOR_CONFIGS
from dl_api_lib_testing.connection_form_base import ConnectionFormTestBase

from dl_connector_bitrix_gds.api.connection_form.form_config import BitrixGDSConnectionFormFactory
from dl_connector_bitrix_gds.api.i18n.localizer import CONFIGS as BI_CONNECTOR_BITRIX_GDS_CONFIGS


class TestBitrixGDSConnectionForm(ConnectionFormTestBase):
    CONN_FORM_FACTORY_CLS = BitrixGDSConnectionFormFactory
    TRANSLATION_CONFIGS = BI_API_CONNECTOR_CONFIGS + BI_CONNECTOR_BITRIX_GDS_CONFIGS

lib/dl_connector_bitrix_gds/dl_connector_bitrix_gds_tests/unit/test_connection_form.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Any

from marshmallow import (
    fields,
    post_dump,
    pre_dump,
    validate,
)

from dl_api_commons.flask.middlewares.commit_rci_middleware import ReqCtxInfoMiddleware
from dl_api_connector.api_schema.component_errors import ComponentErrorListSchema
from dl_api_connector.api_schema.connection_base import ConnectionSchema
from dl_api_connector.api_schema.extras import FieldExtra
from dl_constants.exc import (
    DEFAULT_ERR_CODE_API_PREFIX,
    GLOBAL_ERR_PREFIX,
)
from dl_core import exc
from dl_core.flask_utils.us_manager_middleware import USManagerFlaskMiddleware
from dl_core.utils import (
    make_user_auth_cookies,
    make_user_auth_headers,
)

from dl_connector_bundle_chs3.chs3_base.api.api_schema.source import (
    BaseFileSourceSchema,
    ReplaceFileSourceSchema,
)
from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection


class BaseFileS3ConnectionSchema(ConnectionSchema):
    TARGET_CLS = BaseFileS3Connection

    sources = fields.Nested(
        BaseFileSourceSchema,
        many=True,
        attribute="data.sources",
        bi_extra=FieldExtra(editable=True),
        validate=validate.Length(min=1, max=10),
    )
    replace_sources = fields.Nested(
        ReplaceFileSourceSchema,
        many=True,
        load_default=[],
        attribute="data.replace_sources",
        bi_extra=FieldExtra(editable=True),
        load_only=True,
    )

    component_errors = fields.Nested(ComponentErrorListSchema, attribute="data.component_errors", required=False)

    def create_data_model_constructor_kwargs(self, data_attributes: dict[str, Any]) -> dict[str, Any]:
        base_kwargs = super().create_data_model_constructor_kwargs(data_attributes)
        base_kwargs.pop("replace_sources")  # TODO
        return base_kwargs

    @post_dump(pass_original=True)
    def load_preview(self, data: dict[str, Any], conn: BaseFileS3Connection, **kwargs: Any) -> dict[str, Any]:
        # TODO: read preview data from s3 if status == FileProcessingStatus.ready ???

        rci = ReqCtxInfoMiddleware.get_request_context_info()
        headers = make_user_auth_headers(rci=rci)
        cookies = make_user_auth_cookies(rci=rci)

        usm = USManagerFlaskMiddleware.get_request_service_us_manager()
        service_registry = usm.get_services_registry()
        fu_client_factory = service_registry.get_file_uploader_client_factory()

        file_sources = [src.get_desc() for src in conn.data.sources]
        with fu_client_factory.get_client(headers=headers, cookies=cookies) as fu_client:
            sources_preview = fu_client.get_preview_batch_sync(file_sources)
        preview_by_source_id = {sp.source_id: sp.preview for sp in sources_preview}

        for source in data["sources"]:
            source["preview"] = preview_by_source_id[source["id"]]

        return data

    @pre_dump
    def expand_source_errors(self, conn: BaseFileS3Connection, **kwargs: Any) -> BaseFileS3Connection:
        expected_prefix = [GLOBAL_ERR_PREFIX, DEFAULT_ERR_CODE_API_PREFIX] + exc.DataSourceError.err_code
        for err_pack in conn.data.component_errors.items:
            for error in err_pack.errors:
                if len(error.code) < len(expected_prefix) or error.code[: len(expected_prefix)] != expected_prefix:
                    error.code = expected_prefix + error.code
        return conn

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/api/api_schema/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Type,
)

import marshmallow as ma
from marshmallow import (
    RAISE,
    Schema,
    fields,
    post_load,
)

from dl_api_connector.api_schema.extras import FieldExtra
from dl_api_connector.api_schema.source_base import (
    SQLDataSourceSchema,
    SQLDataSourceTemplateSchema,
)
from dl_constants.enums import (
    FileProcessingStatus,
    UserDataType,
)
from dl_model_tools.schema.base import BaseSchema

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection


class RawSchemaColumnSchema(Schema):
    name = fields.String()
    title = fields.String()
    user_type = fields.Enum(UserDataType)


class BaseFileSourceSchema(Schema):
    class Meta:
        unknown = RAISE

        target: Type[BaseFileS3Connection.FileDataSource]

    @post_load(pass_many=False)
    def to_object(self, data: dict[str, Any], **kwargs: Any) -> BaseFileS3Connection.FileDataSource:
        return self.Meta.target(**data)

    id = fields.String()
    file_id = fields.String(load_default=None)
    title = fields.String(bi_extra=FieldExtra(editable=True))
    status = fields.Enum(FileProcessingStatus, dump_only=True)
    raw_schema = fields.Nested(
        RawSchemaColumnSchema,
        many=True,
        dump_only=True,
    )
    preview = fields.List(fields.List(fields.Raw), dump_only=True, dump_default=None)
    # meta_info = {"row_num": 100500, "size": 256, "file_name": 'file_name.csv', }


class ReplaceFileSourceSchema(Schema):
    class Meta:
        unknown = RAISE

    old_source_id = fields.String()
    new_source_id = fields.String()


class BaseFileS3DataSourceParametersSchema(BaseSchema):
    origin_source_id = ma.fields.String(allow_none=True, load_default=None)


class BaseFileS3DataSourceSchema(SQLDataSourceSchema):
    parameters = ma.fields.Nested(BaseFileS3DataSourceParametersSchema)


class BaseFileS3DataSourceTemplateSchema(SQLDataSourceTemplateSchema):
    parameters = ma.fields.Nested(BaseFileS3DataSourceParametersSchema)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/api/api_schema/source.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)

from dl_connector_bundle_chs3.chs3_base.api.api_schema.connection import BaseFileS3ConnectionSchema
from dl_connector_bundle_chs3.chs3_base.api.api_schema.source import (
    BaseFileS3DataSourceSchema,
    BaseFileS3DataSourceTemplateSchema,
)
from dl_connector_bundle_chs3.chs3_base.api.i18n.localizer import CONFIGS
from dl_connector_bundle_chs3.chs3_base.core.connector import (
    BaseFileS3CoreConnectionDefinition,
    BaseFileS3CoreConnector,
    BaseFileS3TableCoreSourceDefinition,
)
from dl_connector_clickhouse.formula.constants import DIALECT_NAME_CLICKHOUSE


class BaseFileS3TableApiSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = BaseFileS3TableCoreSourceDefinition
    api_schema_cls = BaseFileS3DataSourceSchema
    template_api_schema_cls = BaseFileS3DataSourceTemplateSchema


class BaseFileS3ApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = BaseFileS3CoreConnectionDefinition
    api_generic_schema_cls = BaseFileS3ConnectionSchema


class BaseFileS3ApiConnector(ApiConnector):
    core_connector_cls = BaseFileS3CoreConnector
    formula_dialect_name = DIALECT_NAME_CLICKHOUSE
    connection_definitions = (BaseFileS3ApiConnectionDefinition,)
    source_definitions = (BaseFileS3TableApiSourceDefinition,)
    translation_configs = frozenset(CONFIGS)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/api/connector.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_bundle_chs3 as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/api/i18n/localizer.py
datalens-tech/datalens-backend
from functools import reduce
import logging
from typing import (
    Dict,
    Optional,
)

from aiohttp import ClientResponse
import attr

from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
from dl_core.connection_models import (
    DBIdent,
    TableIdent,
)

from dl_connector_bundle_chs3.chs3_base.core.target_dto import BaseFileS3ConnTargetDTO
from dl_connector_clickhouse.core.clickhouse_base.adapters import BaseAsyncClickHouseAdapter
from dl_connector_clickhouse.core.clickhouse_base.ch_commons import (
    ClickHouseBaseUtils,
    get_ch_settings,
)


class FileS3Utils(ClickHouseBaseUtils):
    pass


LOGGER = logging.getLogger(__name__)


@attr.s(kw_only=True)
class BaseAsyncFileS3Adapter(BaseAsyncClickHouseAdapter):
    ch_utils = FileS3Utils
    _target_dto: BaseFileS3ConnTargetDTO = attr.ib()  # type: ignore  # TODO: FIX

    async def is_table_exists(self, table_ident: TableIdent) -> bool:  # type: ignore
        return True

    async def get_db_version(self, db_ident: DBIdent) -> Optional[str]:  # type: ignore
        return ""

    def get_request_params(self, dba_q: DBAdapterQuery) -> Dict[str, str]:
        return dict(
            # TODO FIX: Move to utils
            database=dba_q.db_name or self._target_dto.db_name or "system",
            **get_ch_settings(
                read_only_level=2,
                max_execution_time=self._target_dto.max_execution_time,
            ),
        )

    async def _make_query(self, dba_q: DBAdapterQuery, mirroring_mode: bool = False) -> ClientResponse:
        if not isinstance(dba_q.query, str):
            dialect = self.get_dialect()
            query_str_raw = dba_q.query.compile(dialect=dialect, compile_kwargs={"literal_binds": True}).string
        else:
            query_str_raw = dba_q.query

        if query_str_raw.startswith("select version()"):  # skip formatting for special queries
            return await super()._make_query(dba_q, mirroring_mode)

        replace_secret = self._target_dto.replace_secret
        secrets = (
            (f"key_id_{replace_secret}", self._target_dto.access_key_id),
            (f"secret_key_{replace_secret}", self._target_dto.secret_access_key),
        )
        secrets_hidden = (
            (f"key_id_{replace_secret}", "<hidden>"),
            (f"secret_key_{replace_secret}", "<hidden>"),
        )

        debug_query = dba_q.debug_compiled_query if dba_q.debug_compiled_query is not None else query_str_raw

        dba_q = dba_q.clone(
            query=reduce(lambda a, kv: a.replace(*kv), secrets, query_str_raw) % (),
            debug_compiled_query=reduce(lambda a, kv: a.replace(*kv), secrets_hidden, debug_query) % (),
        )
        return await super()._make_query(dba_q, mirroring_mode)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/adapter.py
datalens-tech/datalens-backend
import attr

from dl_connector_bundle_chs3.chs3_base.core.dto import BaseFileS3ConnDTO
from dl_connector_bundle_chs3.chs3_base.core.target_dto import BaseFileS3ConnTargetDTO
from dl_connector_clickhouse.core.clickhouse_base.conn_options import CHConnectOptions
from dl_connector_clickhouse.core.clickhouse_base.connection_executors import AsyncClickHouseConnExecutor


@attr.s(cmp=False, hash=False)
class BaseFileS3AsyncAdapterConnExecutor(AsyncClickHouseConnExecutor):
    _conn_dto: BaseFileS3ConnDTO = attr.ib()
    _conn_options: CHConnectOptions = attr.ib()

    async def _make_target_conn_dto_pool(self) -> list[BaseFileS3ConnTargetDTO]:  # type: ignore
        dto_pool = []
        for host in self._conn_hosts_pool:
            dto_pool.append(
                BaseFileS3ConnTargetDTO(
                    conn_id=self._conn_dto.conn_id,
                    pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                    pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                    protocol=self._conn_dto.protocol,
                    host=host,
                    port=self._conn_dto.port,
                    db_name=None,
                    username=self._conn_dto.username,
                    password=self._conn_dto.password,
                    max_execution_time=self._conn_options.max_execution_time,
                    total_timeout=self._conn_options.total_timeout,
                    connect_timeout=self._conn_options.connect_timeout,
                    disable_value_processing=self._conn_options.disable_value_processing,
                    s3_endpoint=self._conn_dto.s3_endpoint,
                    bucket=self._conn_dto.bucket,
                    access_key_id=self._conn_dto.access_key_id,
                    secret_access_key=self._conn_dto.secret_access_key,
                    replace_secret=self._conn_dto.replace_secret,
                )
            )
        return dto_pool

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/connection_executors.py
datalens-tech/datalens-backend
from clickhouse_sqlalchemy.orm.query import Query as CHQuery

from dl_core.connections_security.base import (
    ConnSecuritySettings,
    NonUserInputConnectionSafetyChecker,
)
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
    CoreSourceDefinition,
)

from dl_connector_bundle_chs3.chs3_base.core.constants import BACKEND_TYPE_CHS3
from dl_connector_bundle_chs3.chs3_base.core.dto import BaseFileS3ConnDTO
from dl_connector_bundle_chs3.chs3_base.core.type_transformer import FileTypeTransformer


class BaseFileS3CoreConnectionDefinition(CoreConnectionDefinition):
    type_transformer_cls = FileTypeTransformer
    dialect_string = "bi_clickhouse"


class BaseFileS3TableCoreSourceDefinition(CoreSourceDefinition):
    pass


class BaseFileS3CoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_CHS3
    conn_security = frozenset(
        {ConnSecuritySettings(NonUserInputConnectionSafetyChecker, frozenset({BaseFileS3ConnDTO}))}
    )
    query_cls = CHQuery

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/connector.py
datalens-tech/datalens-backend
from dl_connector_clickhouse.core.clickhouse_base.constants import BACKEND_TYPE_CLICKHOUSE


BACKEND_TYPE_CHS3 = BACKEND_TYPE_CLICKHOUSE

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import wraps
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Optional,
)

from clickhouse_sqlalchemy.quoting import Quoter

from dl_constants.enums import (
    DataSourceType,
    FileProcessingStatus,
)
from dl_core import exc
from dl_core.db import SchemaInfo
from dl_core.utils import sa_plain_text

from dl_connector_bundle_chs3.chs3_base.core.data_source_spec import BaseFileS3DataSourceSpec
from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3.file.core.adapter import AsyncFileS3Adapter
from dl_connector_clickhouse.core.clickhouse_base.ch_commons import create_column_sql
from dl_connector_clickhouse.core.clickhouse_base.data_source import ClickHouseDataSourceBase


if TYPE_CHECKING:
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase


def require_file_configured(func: Callable) -> Callable:
    @wraps(func)
    def wrapper(self: BaseFileS3DataSource, *args: Any, **kwargs: Any) -> Any:
        if self._get_origin_src().status != FileProcessingStatus.ready:
            raise exc.TableNameNotConfiguredError

        return func(self, *args, **kwargs)

    return wrapper


class BaseFileS3DataSource(ClickHouseDataSourceBase):
    preview_enabled: ClassVar[bool] = True
    store_raw_schema = True
    _quoter: Optional[Quoter] = None

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        raise NotImplementedError

    @property
    def origin_source_id(self) -> Optional[str]:
        return self.spec.origin_source_id

    def _get_origin_src(self) -> BaseFileS3Connection.FileDataSource:
        return self.connection.get_file_source_by_id(self.origin_source_id)

    @property
    def connection(self) -> BaseFileS3Connection:
        connection = super().connection
        assert isinstance(connection, BaseFileS3Connection)
        return connection

    @require_file_configured
    def get_schema_info(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> SchemaInfo:
        origin_src = self._get_origin_src()
        assert origin_src.raw_schema is not None
        source_schema = origin_src.raw_schema
        return SchemaInfo.from_schema(source_schema)

    @require_file_configured
    def source_exists(
        self,
        conn_executor_factory: Callable[[], SyncConnExecutorBase],
        force_refresh: bool = False,
    ) -> bool:
        try:
            self.connection.get_file_source_by_id(self.origin_source_id)
        except exc.SourceDoesNotExist:
            return False
        return True

    @property
    def s3_endpoint(self) -> Optional[str]:
        return self.spec.s3_endpoint

    @property
    def bucket(self) -> Optional[str]:
        return self.spec.bucket

    @property
    def status(self) -> FileProcessingStatus:
        return self.spec.status

    @property
    def spec(self) -> BaseFileS3DataSourceSpec:
        assert isinstance(self._spec, BaseFileS3DataSourceSpec)
        return self._spec

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            origin_source_id=self.origin_source_id,
        )

    def quote_str(self, value: str) -> str:
        if self._quoter is None:
            self._quoter = Quoter()
        return self._quoter.quote_str(value)

    def _handle_component_errors(self) -> None:
        pass

    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        origin_src = self._get_origin_src()
        status = origin_src.status
        raw_schema = self.spec.raw_schema
        s3_filename = origin_src.s3_filename

        self._handle_component_errors()

        if status != FileProcessingStatus.ready or raw_schema is None or s3_filename is None:
            raise exc.MaterializationNotFinished

        replace_secret = self.connection.get_replace_secret()

        s3_endpoint = self.connection.s3_endpoint
        bucket = self.connection.bucket
        assert s3_endpoint is not None and bucket is not None and raw_schema is not None

        s3_path = self.quote_str("{}/{}/{}".format(s3_endpoint.strip("/"), bucket.strip("/"), s3_filename))
        key_id_placeholder = self.quote_str(f"key_id_{replace_secret}")
        secret_key_placeholder = self.quote_str(f"secret_key_{replace_secret}")
        file_fmt = self.quote_str("Native")
        dialect = AsyncFileS3Adapter.get_dialect()
        schema_str = self.quote_str(
            ", ".join(create_column_sql(dialect, col, self.type_transformer) for col in raw_schema)
        )
        alias_str = "" if alias is None else f" AS {self.quote(alias)}"

        return sa_plain_text(
            f"s3({s3_path}, {key_id_placeholder}, {secret_key_placeholder}, {file_fmt}, {schema_str}){alias_str}"
        )

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/data_source.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_constants.enums import FileProcessingStatus
from dl_core.data_source_spec.sql import StandardSQLDataSourceSpec


@attr.s
class BaseFileS3DataSourceSpec(StandardSQLDataSourceSpec):
    s3_endpoint: Optional[str] = attr.ib(kw_only=True, default=None)
    bucket: Optional[str] = attr.ib(kw_only=True, default=None)
    origin_source_id: Optional[str] = attr.ib(kw_only=True, default=None)
    status: FileProcessingStatus = attr.ib(kw_only=True, default=FileProcessingStatus.in_progress)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/data_source_spec.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_connector_clickhouse.core.clickhouse_base.dto import ClickHouseConnDTO


@attr.s(frozen=True)
class BaseFileS3ConnDTO(ClickHouseConnDTO):
    s3_endpoint: str = attr.ib(kw_only=True)
    access_key_id: str = attr.ib(kw_only=True, repr=False)
    secret_access_key: str = attr.ib(kw_only=True, repr=False)
    bucket: str = attr.ib(kw_only=True)
    replace_secret: str = attr.ib(kw_only=True)

    # Set default=None
    db_name: Optional[str] = attr.ib(kw_only=True, default=None)  # type: ignore
    cluster_name: Optional[str] = attr.ib(kw_only=True, default=None)  # type: ignore
    endpoint: Optional[str] = attr.ib(kw_only=True, default=None)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/dto.py
datalens-tech/datalens-backend
import logging
from typing import (
    Iterable,
    Optional,
)

import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_core.connectors.base.lifecycle import ConnectionLifecycleManager
from dl_file_uploader_task_interface.tasks import (
    DeleteFileTask,
    SaveSourceTask,
)
from dl_task_processor.processor import TaskProcessor
from dl_utils.aio import await_sync

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection


LOGGER = logging.getLogger(__name__)


@attr.s
class FileConnTaskScheduler:
    _task_processor: TaskProcessor = attr.ib(kw_only=True)
    _rci: RequestContextInfo = attr.ib(kw_only=True)

    def schedule_sources_delete(
        self, conn: BaseFileS3Connection, source_to_del: Optional[Iterable[str]] = None
    ) -> None:
        """Removes all _saved_sources if `sources_to_del` is not specified"""

        if source_to_del is None:
            source_to_del = (src.id for src in conn._saved_sources or [])
        for src_id in source_to_del:
            source = conn.get_saved_source_by_id(src_id)
            if source.s3_filename is None:
                LOGGER.warning(f"Cannot schedule file deletion for source_id {source.id} - s3_filename not set")
                continue
            task = DeleteFileTask(
                s3_filename=source.s3_filename,
                preview_id=source.preview_id,
            )
            task_instance = await_sync(self._task_processor.schedule(task))
            LOGGER.info(
                f"Scheduled task DeleteFileTask for source_id {source.id}, filename {source.s3_filename}. "
                f"instance_id: {task_instance.instance_id}"
            )

    def schedule_sources_update(self, conn: BaseFileS3Connection) -> None:
        saved_sources = set(src.id for src in (conn._saved_sources or []))
        current_sources = set(src.id for src in conn.data.sources)
        sources_to_add = current_sources - saved_sources
        sources_to_del = saved_sources - current_sources
        sources_to_replace = {
            upd_src["old_source_id"]: upd_src["new_source_id"] for upd_src in conn.data.replace_sources
        }
        assert self._rci.tenant is not None

        for src_id in sources_to_add:
            source = conn.get_file_source_by_id(src_id)
            src_source_id = sources_to_replace.get(source.id, source.id)
            task = SaveSourceTask(
                tenant_id=self._rci.tenant.get_tenant_id(),
                file_id=source.file_id,
                src_source_id=src_source_id,
                dst_source_id=source.id,
                connection_id=conn.uuid,  # type: ignore
            )
            task_instance = await_sync(self._task_processor.schedule(task))
            LOGGER.info(
                f"Scheduled task SaveSourceTask for source_id {source.id}. instance_id: {task_instance.instance_id}"
            )
        self.schedule_sources_delete(conn, sources_to_del)


class BaseFileS3ConnectionLifecycleManager(ConnectionLifecycleManager[BaseFileS3Connection]):
    ENTRY_CLS = BaseFileS3Connection

    def get_task_processor(self, req_id: Optional[str] = None) -> TaskProcessor:
        task_processor_factory = self._service_registry.get_task_processor_factory()
        task_processor = task_processor_factory.make(req_id)
        return task_processor

    def post_save_hook(self) -> None:
        super().post_save_hook()

        rci = self._us_manager.bi_context
        task_processor = self.get_task_processor(rci.request_id)
        scheduler = FileConnTaskScheduler(task_processor=task_processor, rci=rci)
        scheduler.schedule_sources_update(self.entry)

    def post_delete_hook(self) -> None:
        super().post_delete_hook()

        rci = self._us_manager.bi_context
        task_processor = self.get_task_processor(rci.request_id)
        scheduler = FileConnTaskScheduler(task_processor=task_processor, rci=rci)
        scheduler.schedule_sources_delete(self.entry)

    async def post_init_async_hook(self) -> None:
        await super().post_init_async_hook()

        self.entry._saved_sources = self.entry.data.sources

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/lifecycle.py
datalens-tech/datalens-backend
from typing import (
    ClassVar,
    Optional,
)
import uuid

import attr

from dl_configs.connectors_data import ConnectorsDataBase
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.settings_loaders.meta_definition import s_attrib


@attr.s(frozen=True)
class FileS3ConnectorSettings(ConnectorSettingsBase):
    SECURE: bool = s_attrib("SECURE", missing=True)  # type: ignore
    HOST: str = s_attrib("HOST")  # type: ignore
    PORT: int = s_attrib("PORT", missing=8443)  # type: ignore
    USERNAME: str = s_attrib("USERNAME")  # type: ignore
    PASSWORD: str = s_attrib("PASSWORD", sensitive=True, missing=None)  # type: ignore

    ACCESS_KEY_ID: str = s_attrib("ACCESS_KEY_ID", sensitive=True, missing=None)  # type: ignore
    SECRET_ACCESS_KEY: str = s_attrib("SECRET_ACCESS_KEY", sensitive=True, missing=None)  # type: ignore
    S3_ENDPOINT: str = s_attrib("S3_ENDPOINT")  # type: ignore
    BUCKET: str = s_attrib("BUCKET")  # type: ignore

    REPLACE_SECRET_SALT: str = s_attrib(  # type: ignore
        "REPLACE_SECRET_SALT", sensitive=True, missing_factory=lambda: str(uuid.uuid4())
    )
    # ^ Note that this is used in a query, which, in turn, is used in a cache key at the moment
    #   This means that the value must be set explicitly to preserve caches between restarts and instances


class ConnectorsDataFileBase(ConnectorsDataBase):
    CONN_FILE_CH_HOST: ClassVar[Optional[str]] = None
    CONN_FILE_CH_PORT: ClassVar[int] = 8443
    CONN_FILE_CH_USERNAME: ClassVar[Optional[str]] = None

    @classmethod
    def connector_name(cls) -> str:
        return "FILE"

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/settings.py
datalens-tech/datalens-backend
from marshmallow import (
    Schema,
    fields,
)

from dl_constants.enums import FileProcessingStatus
from dl_core.us_manager.storage_schemas.base import DefaultStorageSchema
from dl_core.us_manager.storage_schemas.base_types import SchemaColumnStorageSchema
from dl_core.us_manager.storage_schemas.error_registry import ComponentErrorListSchema


class BaseFileConnectionSourceStorageSchema(DefaultStorageSchema):
    id = fields.String()
    file_id = fields.String(allow_none=True, load_default=None)
    preview_id = fields.String(allow_none=True, load_default=None)
    title = fields.String()
    s3_filename = fields.String(allow_none=True, load_default=None)
    raw_schema = fields.Nested(SchemaColumnStorageSchema, many=True, allow_none=True, load_default=None)
    status = fields.Enum(FileProcessingStatus)


class BaseFileConnectionDataStorageSchema(Schema):
    sources = fields.Nested(BaseFileConnectionSourceStorageSchema, many=True)
    component_errors = fields.Nested(ComponentErrorListSchema)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/storage_schemas/connection.py
datalens-tech/datalens-backend
from marshmallow import fields

from dl_constants.enums import FileProcessingStatus
from dl_core.us_manager.storage_schemas.data_source_spec_base import SQLDataSourceSpecStorageSchema

from dl_connector_bundle_chs3.chs3_base.core.data_source_spec import BaseFileS3DataSourceSpec


class BaseFileS3DataSourceSpecStorageSchema(SQLDataSourceSpecStorageSchema):
    TARGET_CLS = BaseFileS3DataSourceSpec

    s3_endpoint = fields.String(required=False, allow_none=True, load_default=None)
    bucket = fields.String(required=False, allow_none=True, load_default=None)
    status = fields.Enum(FileProcessingStatus, required=False, allow_none=True, load_default=None)
    origin_source_id = fields.String(required=False, allow_none=True, load_default=None)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/storage_schemas/data_source_spec.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_core.connection_executors.models.connection_target_dto_base import BaseSQLConnTargetDTO
from dl_core.utils import secrepr


@attr.s
class BaseFileS3ConnTargetDTO(BaseSQLConnTargetDTO):
    protocol: str = attr.ib(kw_only=True)
    disable_value_processing: bool = attr.ib(kw_only=True)

    endpoint: Optional[str] = attr.ib(kw_only=True, default=None)
    max_execution_time: Optional[int] = attr.ib(kw_only=True, default=None)
    connect_timeout: Optional[int] = attr.ib(kw_only=True, default=None)
    total_timeout: Optional[int] = attr.ib(kw_only=True, default=None)

    s3_endpoint: str = attr.ib(kw_only=True)
    bucket: str = attr.ib(kw_only=True)
    access_key_id: str = attr.ib(kw_only=True, repr=secrepr)
    secret_access_key: str = attr.ib(kw_only=True, repr=secrepr)
    replace_secret: str = attr.ib(kw_only=True)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/target_dto.py
datalens-tech/datalens-backend
from dl_configs.settings_submodels import S3Settings
from dl_core_testing.database import DbTable
from dl_testing.s3_utils import s3_tbl_func_maker

from dl_connector_bundle_chs3.chs3_base.core.dto import BaseFileS3ConnDTO


def create_s3_native_from_ch_table(
    filename: str,
    s3_bucket: str,
    s3_settings: S3Settings,
    clickhouse_table: DbTable,
    tbl_schema: str,
) -> None:
    tbl = clickhouse_table
    db = clickhouse_table.db
    s3_tbl_func = s3_tbl_func_maker(s3_settings)
    s3_tbl_func_for_db = s3_tbl_func(
        for_="db",
        conn_dto=BaseFileS3ConnDTO(
            conn_id=None,
            protocol="http",
            host=db.url.host,
            port=db.url.port,
            username=db.url.username,
            password=db.url.password,
            multihosts=db.get_conn_hosts(),  # type: ignore
            s3_endpoint="http://s3-storage:8000",
            access_key_id=s3_settings.ACCESS_KEY_ID,
            secret_access_key=s3_settings.SECRET_ACCESS_KEY,
            bucket=s3_bucket,
            replace_secret="WE_ONLY_NEED_DB_MODE_HERE_SO_NO_SECRET",
        ),
        filename=filename,
        file_fmt="Native",
        schema_line=tbl_schema,
    )
    insert_stmt = f"INSERT INTO FUNCTION {s3_tbl_func_for_db} SELECT * FROM {tbl.db.quote(tbl.name)}"
    db.execute(insert_stmt)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/testing/utils.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import re
from typing import (
    Any,
    Optional,
)

from clickhouse_sqlalchemy import types as ch_types

from dl_constants.enums import UserDataType
from dl_core.db.conversion_base import (
    BooleanTypeCaster,
    DatetimeTypeCaster,
    DatetimeTZTypeCaster,
    DateTypeCaster,
    GenericDatetimeTypeCaster,
    TypeCaster,
    make_int,
    make_native_type,
)

from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
from dl_connector_clickhouse.core.clickhouse_base.type_transformer import ClickHouseTypeTransformer


def make_int_cleanup_spaces(value: Any) -> Optional[int]:
    if isinstance(value, str):
        if " " in value:
            value = value.replace(" ", "")
    return make_int(value)


_spaces_in_numbers_re = re.compile(r"(?<=\d) (?=\d)")


def make_float_cleanup_spaces(value: Any) -> Optional[float]:
    if isinstance(value, str):
        if "," in value:
            value = value.replace(",", ".")
        if " " in value:
            value = _spaces_in_numbers_re.sub("", value)
    return float(value)


def make_boolean(value: Any) -> Optional[bool]:
    if value is None:
        return None
    if isinstance(value, str):
        value = value.lower()
        if value in ("true", "1"):
            return True
        elif value in ("false", "0"):
            return False
    return bool(value)


class IntegerFileTypeCaster(TypeCaster):
    def _cast_for_input(self, value: Any) -> Any:
        return make_int_cleanup_spaces(value)


class FloatFileTypeCaster(TypeCaster):
    def _cast_for_input(self, value: Any) -> Any:
        return make_float_cleanup_spaces(value)


class DateFileTypeCaster(DateTypeCaster):
    def _cast_for_input(self, value: Any) -> Any:
        return DateTypeCaster.cast_func(value).isoformat()  # type: ignore


class DatetimeFileTypeCaster(DatetimeTypeCaster):
    def _cast_for_input(self, value: Any) -> Any:
        dt: Optional[datetime.datetime] = DatetimeTypeCaster.cast_func(value)
        if dt is not None:
            if dt.tzinfo is not None and dt.utcoffset() is not None:
                dt = dt.replace(tzinfo=None) - dt.utcoffset()  # type: ignore
            return dt.isoformat()
        return None


class DatetimeTZFileTypeCaster(DatetimeTZTypeCaster):
    def _cast_for_input(self, value: Any) -> Any:
        return DatetimeTZTypeCaster.cast_func(value).isoformat()  # type: ignore


class GenericDatetimeFileTypeCaster(GenericDatetimeTypeCaster):
    def _cast_for_input(self, value: Any) -> Any:
        dt: Optional[datetime.datetime] = GenericDatetimeTypeCaster.cast_func(value)
        if dt is not None:
            if dt.tzinfo is not None and dt.utcoffset() is not None:
                dt = dt.replace(tzinfo=None) - dt.utcoffset()  # type: ignore
            return dt.isoformat()
        return None


class BooleanFileTypeCaster(BooleanTypeCaster):
    cast_func = make_boolean


class FileTypeTransformer(ClickHouseTypeTransformer):
    casters = {
        **ClickHouseTypeTransformer.casters,
        UserDataType.integer: IntegerFileTypeCaster(),
        UserDataType.float: FloatFileTypeCaster(),
        UserDataType.date: DateFileTypeCaster(),
        UserDataType.datetime: DatetimeFileTypeCaster(),
        UserDataType.datetimetz: DatetimeTZFileTypeCaster(),
        UserDataType.genericdatetime: GenericDatetimeFileTypeCaster(),
        UserDataType.boolean: BooleanFileTypeCaster(),
    }

    user_to_native_map = {
        **ClickHouseTypeTransformer.user_to_native_map,
        UserDataType.datetime: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.DateTime64),
        UserDataType.genericdatetime: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.DateTime64),
        UserDataType.date: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Date32),
    }

    @classmethod
    def cast_for_input(cls, value: Any, user_t: UserDataType) -> Any:
        """Prepare value for insertion into the database"""
        if value == "" or value is None:
            return None
        return cls.casters[user_t].cast_for_input(value=value)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Optional,
)

import attr
import xxhash

from dl_constants.enums import (
    DataSourceRole,
    FileProcessingStatus,
)
from dl_core import (
    connection_models,
    exc,
)
from dl_core.base_models import ConnectionDataModelBase
from dl_core.component_errors import ComponentErrorRegistry
from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.db.elements import SchemaColumn
from dl_core.services_registry.file_uploader_client_factory import FileSourceDesc
from dl_core.us_connection_base import (
    ConnectionBase,
    ConnectionHardcodedDataMixin,
    DataSourceTemplate,
)
from dl_core.utils import (
    make_user_auth_cookies,
    make_user_auth_headers,
    parse_comma_separated_hosts,
)

from dl_connector_bundle_chs3.chs3_base.core.dto import BaseFileS3ConnDTO
from dl_connector_bundle_chs3.chs3_base.core.settings import FileS3ConnectorSettings
from dl_connector_clickhouse.core.clickhouse_base.conn_options import CHConnectOptions
from dl_connector_clickhouse.core.clickhouse_base.us_connection import ConnectionClickhouseBase


if TYPE_CHECKING:
    from dl_core.services_registry.top_level import ServicesRegistry


LOGGER = logging.getLogger(__name__)


class BaseFileS3Connection(ConnectionHardcodedDataMixin[FileS3ConnectorSettings], ConnectionClickhouseBase):  # type: ignore
    is_always_internal_source: ClassVar[bool] = True
    allow_cache: ClassVar[bool] = True
    settings_type = FileS3ConnectorSettings

    editable_data_source_parameters: ClassVar[tuple[str, ...]] = (
        "file_id",
        "title",
        "s3_filename",
        "status",
        "preview_id",
    )

    @attr.s(kw_only=True)
    class FileDataSource:
        id: str = attr.ib()
        file_id: str = attr.ib()
        title: str = attr.ib()
        preview_id: Optional[str] = attr.ib(default=None)
        status: FileProcessingStatus = attr.ib(default=FileProcessingStatus.in_progress)
        s3_filename: Optional[str] = attr.ib(default=None)
        raw_schema: Optional[list[SchemaColumn]] = attr.ib(factory=list[SchemaColumn])

        def str_for_hash(self) -> str:
            return ",".join(
                [
                    self.id,
                    self.file_id,
                    self.title,
                    str(self.s3_filename),
                    self.status.name,
                ]
            )

        def get_desc(self) -> FileSourceDesc:
            raise NotImplementedError()

    @attr.s(eq=False, kw_only=True)
    class DataModel(ConnectionDataModelBase):
        sources: list["BaseFileS3Connection.FileDataSource"] = attr.ib()
        replace_sources: list[dict] = attr.ib(factory=list)  # should not be saved in US
        data_export_forbidden: bool = attr.ib(default=False)

        component_errors: ComponentErrorRegistry = attr.ib(factory=ComponentErrorRegistry)

        def str_for_hash(self) -> str:
            return "|".join(src.str_for_hash() for src in self.sources)

    data: DataModel

    def get_replace_secret(self) -> str:
        return xxhash.xxh64(self.data.str_for_hash() + self._connector_settings.REPLACE_SECRET_SALT, seed=0).hexdigest()

    @property
    def bucket(self) -> str:
        return self._connector_settings.BUCKET

    @property
    def s3_endpoint(self) -> str:
        return self._connector_settings.S3_ENDPOINT

    @property
    def s3_access_key_id(self) -> str:
        return self._connector_settings.ACCESS_KEY_ID

    @property
    def s3_secret_access_key(self) -> str:
        return self._connector_settings.SECRET_ACCESS_KEY

    def get_conn_dto(self) -> BaseFileS3ConnDTO:  # type: ignore
        cs = self._connector_settings
        conn_dto = BaseFileS3ConnDTO(
            conn_id=self.uuid,
            s3_endpoint=cs.S3_ENDPOINT,
            access_key_id=cs.ACCESS_KEY_ID,
            secret_access_key=cs.SECRET_ACCESS_KEY,
            bucket=cs.BUCKET,
            replace_secret=self.get_replace_secret(),
            protocol="https" if cs.SECURE else "http",
            host=cs.HOST,
            multihosts=parse_comma_separated_hosts(cs.HOST),  # type: ignore
            port=cs.PORT,
            username=cs.USERNAME,
            password=cs.PASSWORD,
        )
        return conn_dto

    def get_conn_options(self) -> CHConnectOptions:
        return self.get_effective_conn_options(
            base_conn_opts=connection_models.ConnectOptions(pass_db_messages_to_user=False),
            max_allowed_max_execution_time=self.MAX_ALLOWED_MAX_EXECUTION_TIME,
            user_max_execution_time=None,
        )

    @property
    def cache_ttl_sec_override(self) -> Optional[int]:
        return 600

    def get_data_source_templates(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[DataSourceTemplate]:
        assert self.source_type is not None
        return [
            DataSourceTemplate(
                title=source.title,
                group=[],
                source_type=self.source_type,
                connection_id=self.uuid,  # type: ignore
                parameters={
                    "origin_source_id": source.id,
                },
            )
            for source in self.data.sources
        ]

    def get_file_source_by_id(self, id: Optional[str]) -> FileDataSource:
        file_source = next(iter(src for src in self.data.sources if src.id == id), None)
        if file_source is None:
            raise exc.SourceDoesNotExist(f"DataSource id={id} not found in connection id={self.uuid}")
        return file_source

    def remove_source_by_id(self, id: str) -> None:
        idx = next(iter(i for i, src in enumerate(self.data.sources) if src.id == id), None)
        if idx is not None:
            del self.data.sources[idx]
        else:
            LOGGER.warning(f"DataSource id={id} not found in connection id={self.uuid} when trying to remove it")

    def _remove_saved_source_by_id(self, id: str) -> None:
        assert self._saved_sources is not None
        idx = next(iter(i for i, src in enumerate(self._saved_sources) if src.id == id), None)
        if idx is not None:
            del self._saved_sources[idx]
        else:
            LOGGER.warning(
                f"DataSource id={id} not found in connection id={self.uuid} saved sources when trying to remove it"
            )

    def get_saved_source_by_id(self, id: Optional[str]) -> FileDataSource:
        file_source = next(iter(src for src in self._saved_sources or [] if src.id == id), None)
        if file_source is None:
            raise exc.SourceDoesNotExist(f"DataSource id={id} not found in saved sources of connection id={self.uuid}")
        return file_source

    def update_data_source(
        self,
        id: str,
        role: Optional[DataSourceRole] = None,
        raw_schema: Optional[list] = None,
        remove_raw_schema: bool = False,
        **parameters: Any,
    ) -> None:
        if role != DataSourceRole.origin:
            raise ValueError(f"Unsupported role for {self.__class__.__name__} datasource.")

        source = self.get_file_source_by_id(id)

        sentinel = object()
        filtered_parameters = {
            param_name: param_value
            for param_name in self.editable_data_source_parameters
            if (param_value := parameters.pop(param_name, sentinel)) is not sentinel
        }
        if parameters:
            raise ValueError(f"Unknown update_data_source parameter: {parameters}")

        LOGGER.info(f"Updating data source {id} with role {role}. Parameters: {filtered_parameters}")

        if raw_schema is not None:
            source.raw_schema = raw_schema
            LOGGER.info("Updated data source raw_schema")
        if remove_raw_schema:
            source.raw_schema = None
            LOGGER.info("Removed data source raw_schema")

        for param_name, value in filtered_parameters.items():
            setattr(source, param_name, value)

    def replace_sources(self) -> None:
        assert self._saved_sources is not None
        for upd_source in self.data.replace_sources:
            old_src = self.get_saved_source_by_id(upd_source["old_source_id"])
            new_src = self.get_file_source_by_id(upd_source["new_source_id"])
            new_src.id, old_src.id = old_src.id, new_src.id

    def restore_source_params_from_orig(self, src_id: str, original_version: BaseFileS3Connection) -> None:
        raise NotImplementedError()

    def test(self, conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase]) -> None:
        """
        Don't execute `select 1` on our service databases - it's useless because user can't
        manage it anyway.
        """
        pass

    def validate(self) -> None:
        try:
            pass  # TODO
        except exc.EntryValidationError as err:
            raise err

    _saved_sources: Optional[list[FileDataSource]] = None

    async def validate_new_data(
        self,
        services_registry: ServicesRegistry,
        changes: Optional[dict] = None,
        original_version: Optional[ConnectionBase] = None,
    ) -> None:
        assert isinstance(original_version, (type(self), type(None)))
        if original_version is None:
            saved_sources = set()
        else:
            saved_sources = set(src.id for src in original_version.data.sources)

        if changes is None:
            sources_to_update = set()
        else:
            # save source if we receive file_id regardless of whether it was saved or not
            sources_to_update = set(src.id for src in changes["data"]["sources"] if src.file_id is not None)

        current_sources = set(src.id for src in self.data.sources)
        sources_to_add = (current_sources - saved_sources) | sources_to_update
        unaltered_sources = current_sources & saved_sources - sources_to_update

        if len(current_sources) != len(set(src.title for src in self.data.sources)):
            raise exc.DataSourceTitleConflict

        err_details = defaultdict(list)
        for src in self.data.sources:
            if src.file_id is None and src.id not in saved_sources:
                err_details["not_configured_not_saved"].append(src.id)

        for replace_src in self.data.replace_sources:
            if replace_src["old_source_id"] not in saved_sources:
                err_details["replaced_not_saved"].append(replace_src["old_source_id"])

        if err_details:
            raise exc.DataSourcesInconsistent(details=err_details)

        if original_version is not None:
            for src_id in unaltered_sources:
                # Restore original internal source properties
                self.restore_source_params_from_orig(src_id, original_version)

        rci = self.us_manager.bi_context
        headers = make_user_auth_headers(rci=rci)
        cookies = make_user_auth_cookies(rci=rci)
        fu_client_factory = self.us_manager.get_services_registry().get_file_uploader_client_factory()
        fu_client = fu_client_factory.get_client(headers=headers, cookies=cookies)

        sources_desc = []
        for src_id in sources_to_add:
            src = self.get_file_source_by_id(src_id)
            sources_desc.append(src.get_desc())

        if sources_desc:
            async with fu_client:
                internal_params = await fu_client.get_internal_params_batch(sources_desc)

            for src_order, src_id in enumerate(sources_to_add):
                self.update_data_source(
                    src_id,
                    role=DataSourceRole.origin,
                    raw_schema=internal_params[src_order].raw_schema,
                    preview_id=internal_params[src_order].preview_id,
                )

        if self._saved_sources is not None:
            saved_source_ids = set(src.id for src in self._saved_sources)
            for src_id in sources_to_update & saved_source_ids:  # trigger update by removing source from _saved_sources
                self._remove_saved_source_by_id(src_id)

        # clear errors for updated sources
        for src_id in sources_to_add:
            self.data.component_errors.remove_errors(id=src_id)

        if self.data.replace_sources:
            self.replace_sources()

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/us_connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import (
    fields,
    validate,
)

from dl_api_connector.api_schema.extras import FieldExtra

from dl_connector_bundle_chs3.chs3_base.api.api_schema.connection import BaseFileS3ConnectionSchema
from dl_connector_bundle_chs3.chs3_gsheets.api.api_schema.source import GSheetsFileSourceSchema
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection


class GSheetsFileS3ConnectionSchema(BaseFileS3ConnectionSchema):
    TARGET_CLS = GSheetsFileS3Connection

    sources = fields.Nested(
        GSheetsFileSourceSchema,
        many=True,
        attribute="data.sources",
        bi_extra=FieldExtra(editable=True),
        validate=validate.Length(min=1, max=10),
    )

    refresh_token = fields.String(
        attribute="data.refresh_token",
        load_default=None,
        allow_none=True,
        load_only=True,
        bi_extra=FieldExtra(editable=True),
    )
    refresh_enabled = fields.Boolean(attribute="data.refresh_enabled", bi_extra=FieldExtra(editable=True))
    authorized = fields.Boolean(dump_only=True)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/api/api_schema/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import fields

from dl_connector_bundle_chs3.chs3_base.api.api_schema.source import BaseFileSourceSchema
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection


class GSheetsFileSourceSchema(BaseFileSourceSchema):
    class Meta(BaseFileSourceSchema.Meta):
        target = GSheetsFileS3Connection.FileDataSource

    spreadsheet_id = fields.String(dump_only=True)
    sheet_id = fields.Integer(dump_only=True)
    first_line_is_header = fields.Boolean(dump_only=True)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/api/api_schema/source.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_bundle_chs3.chs3_base.api.i18n.localizer import Translatable


class GSheetsFileS3ConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-gsheets_v2")

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/api/connection_info.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_connector_bundle_chs3.chs3_base.api.connector import (
    BaseFileS3ApiConnectionDefinition,
    BaseFileS3ApiConnector,
    BaseFileS3TableApiSourceDefinition,
)
from dl_connector_bundle_chs3.chs3_gsheets.api.api_schema.connection import GSheetsFileS3ConnectionSchema
from dl_connector_bundle_chs3.chs3_gsheets.api.connection_info import GSheetsFileS3ConnectionInfoProvider
from dl_connector_bundle_chs3.chs3_gsheets.core.connector import (
    GSheetsFileS3CoreConnectionDefinition,
    GSheetsFileS3CoreConnector,
    GSheetsFileS3TableCoreSourceDefinition,
)


class GSheetsFileS3TableApiSourceDefinition(BaseFileS3TableApiSourceDefinition):
    core_source_def_cls = GSheetsFileS3TableCoreSourceDefinition


class GSheetsFileS3ApiConnectionDefinition(BaseFileS3ApiConnectionDefinition):
    core_conn_def_cls = GSheetsFileS3CoreConnectionDefinition
    api_generic_schema_cls = GSheetsFileS3ConnectionSchema
    info_provider_cls = GSheetsFileS3ConnectionInfoProvider


class GSheetsFileS3ApiConnector(BaseFileS3ApiConnector):
    core_connector_cls = GSheetsFileS3CoreConnector
    connection_definitions = (GSheetsFileS3ApiConnectionDefinition,)
    source_definitions = (GSheetsFileS3TableApiSourceDefinition,)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/api/connector.py
datalens-tech/datalens-backend
import logging

from dl_connector_bundle_chs3.chs3_base.core.adapter import BaseAsyncFileS3Adapter
from dl_connector_bundle_chs3.chs3_gsheets.core.constants import CONNECTION_TYPE_GSHEETS_V2


LOGGER = logging.getLogger(__name__)


class AsyncGSheetsFileS3Adapter(BaseAsyncFileS3Adapter):
    conn_type = CONNECTION_TYPE_GSHEETS_V2

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/adapter.py
datalens-tech/datalens-backend
import attr

from dl_connector_bundle_chs3.chs3_base.core.connection_executors import BaseFileS3AsyncAdapterConnExecutor
from dl_connector_bundle_chs3.chs3_gsheets.core.adapter import AsyncGSheetsFileS3Adapter


@attr.s(cmp=False, hash=False)
class GSheetsFileS3AsyncAdapterConnExecutor(BaseFileS3AsyncAdapterConnExecutor):
    TARGET_ADAPTER_CLS = AsyncGSheetsFileS3Adapter  # type: ignore  # TODO: FIX

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/connection_executors.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.chs3_base.core.connector import (
    BaseFileS3CoreConnectionDefinition,
    BaseFileS3CoreConnector,
    BaseFileS3TableCoreSourceDefinition,
)
from dl_connector_bundle_chs3.chs3_gsheets.core.adapter import AsyncGSheetsFileS3Adapter
from dl_connector_bundle_chs3.chs3_gsheets.core.connection_executors import GSheetsFileS3AsyncAdapterConnExecutor
from dl_connector_bundle_chs3.chs3_gsheets.core.constants import (
    CONNECTION_TYPE_GSHEETS_V2,
    SOURCE_TYPE_GSHEETS_V2,
)
from dl_connector_bundle_chs3.chs3_gsheets.core.data_source import GSheetsFileS3DataSource
from dl_connector_bundle_chs3.chs3_gsheets.core.data_source_spec import GSheetsFileS3DataSourceSpec
from dl_connector_bundle_chs3.chs3_gsheets.core.lifecycle import GSheetsFileS3ConnectionLifecycleManager
from dl_connector_bundle_chs3.chs3_gsheets.core.notifications import (
    DataUpdateFailureNotification,
    StaleDataNotification,
)
from dl_connector_bundle_chs3.chs3_gsheets.core.settings import GSheetsFileS3SettingDefinition
from dl_connector_bundle_chs3.chs3_gsheets.core.storage_schemas.connection import GSheetsFileConnectionDataStorageSchema
from dl_connector_bundle_chs3.chs3_gsheets.core.storage_schemas.data_source_spec import (
    GSheetsFileS3DataSourceSpecStorageSchema,
)
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection
from dl_connector_clickhouse.core.clickhouse_base.adapters import ClickHouseAdapter


class GSheetsFileS3CoreConnectionDefinition(BaseFileS3CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_GSHEETS_V2
    connection_cls = GSheetsFileS3Connection
    us_storage_schema_cls = GSheetsFileConnectionDataStorageSchema
    sync_conn_executor_cls = GSheetsFileS3AsyncAdapterConnExecutor
    async_conn_executor_cls = GSheetsFileS3AsyncAdapterConnExecutor
    lifecycle_manager_cls = GSheetsFileS3ConnectionLifecycleManager
    dialect_string = "bi_clickhouse"
    settings_definition = GSheetsFileS3SettingDefinition


class GSheetsFileS3TableCoreSourceDefinition(BaseFileS3TableCoreSourceDefinition):
    source_type = SOURCE_TYPE_GSHEETS_V2
    source_cls = GSheetsFileS3DataSource
    source_spec_cls = GSheetsFileS3DataSourceSpec
    us_storage_schema_cls = GSheetsFileS3DataSourceSpecStorageSchema


class GSheetsFileS3CoreConnector(BaseFileS3CoreConnector):
    connection_definitions = (GSheetsFileS3CoreConnectionDefinition,)
    source_definitions = (GSheetsFileS3TableCoreSourceDefinition,)
    rqe_adapter_classes = frozenset({ClickHouseAdapter, AsyncGSheetsFileS3Adapter})
    notification_classes = (
        StaleDataNotification,
        DataUpdateFailureNotification,
    )

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
    NotificationType,
)


CONNECTION_TYPE_GSHEETS_V2 = ConnectionType.declare("gsheets_v2")
SOURCE_TYPE_GSHEETS_V2 = DataSourceType.declare("GSHEETS_V2")

NOTIF_TYPE_GSHEETS_V2_STALE_DATA = NotificationType.declare("stale_data")
NOTIF_TYPE_GSHEETS_V2_DATA_UPDATE_FAILURE = NotificationType.declare("data_update_failure")

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.enums import (
    ComponentErrorLevel,
    DataSourceType,
)
from dl_core import exc
from dl_core.reporting.notifications import get_notification_record

from dl_connector_bundle_chs3.chs3_base.core.data_source import BaseFileS3DataSource
from dl_connector_bundle_chs3.chs3_gsheets.core.constants import (
    CONNECTION_TYPE_GSHEETS_V2,
    NOTIF_TYPE_GSHEETS_V2_DATA_UPDATE_FAILURE,
    SOURCE_TYPE_GSHEETS_V2,
)


class GSheetsFileS3DataSource(BaseFileS3DataSource):
    conn_type = CONNECTION_TYPE_GSHEETS_V2

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type in {
            SOURCE_TYPE_GSHEETS_V2,
        }

    def _handle_component_errors(self) -> None:
        conn_src_id = self.origin_source_id
        if conn_src_id is not None and (error_pack := self.connection.data.component_errors.get_pack(conn_src_id)):
            single_error = error_pack.errors[0]

            if single_error.level == ComponentErrorLevel.error:

                class ThisDataSourceError(exc.DataSourceErrorFromComponentError):
                    err_code = exc.DataSourceErrorFromComponentError.err_code + single_error.code
                    default_message = single_error.message

                raise ThisDataSourceError(
                    details=single_error.details,
                )
            else:
                reporting_registry = self._get_connection().us_manager.get_services_registry().get_reporting_registry()
                # this is the only case of connection component errors at the moment
                # may be generalized in the future
                reporting_registry.save_reporting_record(
                    get_notification_record(
                        NOTIF_TYPE_GSHEETS_V2_DATA_UPDATE_FAILURE,
                        err_code=".".join(single_error.code),
                        request_id=single_error.details.get("request-id"),
                    )
                )

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/data_source.py
datalens-tech/datalens-backend
import attr

from dl_connector_bundle_chs3.chs3_base.core.data_source_spec import BaseFileS3DataSourceSpec


@attr.s
class GSheetsFileS3DataSourceSpec(BaseFileS3DataSourceSpec):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/data_source_spec.py
datalens-tech/datalens-backend
import datetime
from typing import ClassVar

from dl_constants.enums import FileProcessingStatus
from dl_core.connectors.base.lifecycle import ConnectionLifecycleManager
from dl_core.reporting.notifications import get_notification_record
from dl_core.utils import (
    make_user_auth_cookies,
    make_user_auth_headers,
)

from dl_connector_bundle_chs3.chs3_base.core.lifecycle import BaseFileS3ConnectionLifecycleManager
from dl_connector_bundle_chs3.chs3_gsheets.core.constants import NOTIF_TYPE_GSHEETS_V2_STALE_DATA
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection


class GSheetsFileS3ConnectionLifecycleManager(
    BaseFileS3ConnectionLifecycleManager,
    ConnectionLifecycleManager[GSheetsFileS3Connection],
):
    ENTRY_CLS = GSheetsFileS3Connection

    STALE_THRESHOLD_SECONDS: ClassVar[int] = 30 * 60

    async def post_exec_async_hook(self) -> None:
        await super().post_exec_async_hook()

        assert isinstance(self.entry, GSheetsFileS3Connection)
        assert isinstance(self.entry.data, GSheetsFileS3Connection.DataModel)
        data = self.entry.data

        if not data.refresh_enabled:
            return

        dt_now = datetime.datetime.now(datetime.timezone.utc)

        data_updated_at_all = data.oldest_data_update_time()
        if (
            data_updated_at_all is not None
            and (dt_now - data_updated_at_all).total_seconds() >= self.STALE_THRESHOLD_SECONDS
        ):
            reporting_registry = self._service_registry.get_reporting_registry()
            reporting_registry.save_reporting_record(get_notification_record(NOTIF_TYPE_GSHEETS_V2_STALE_DATA))

        data_updated_at = data.oldest_data_update_time(exclude_statuses={FileProcessingStatus.in_progress})
        if data_updated_at is None or (dt_now - data_updated_at).total_seconds() < self.STALE_THRESHOLD_SECONDS:
            return

        fu_client_factory = self._service_registry.get_file_uploader_client_factory()
        rci = self._us_manager.bi_context
        headers = make_user_auth_headers(rci=rci)
        cookies = make_user_auth_cookies(rci=rci)
        fu_client = fu_client_factory.get_client(headers=headers, cookies=cookies)

        sources = [src.get_desc() for src in data.sources]
        assert self.entry.uuid is not None
        await fu_client.update_connection_data_internal(
            conn_id=self.entry.uuid,
            sources=sources,
            authorized=self.entry.authorized,
            tenant_id=rci.tenant.get_tenant_id() if rci.tenant is not None else None,
        )

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/lifecycle.py
datalens-tech/datalens-backend
from typing import Optional

from dl_constants.enums import NotificationLevel
from dl_core.reporting.notifications import BaseNotification

from dl_connector_bundle_chs3.chs3_gsheets.core.constants import (
    NOTIF_TYPE_GSHEETS_V2_DATA_UPDATE_FAILURE,
    NOTIF_TYPE_GSHEETS_V2_STALE_DATA,
)


class StaleDataNotification(BaseNotification):
    type = NOTIF_TYPE_GSHEETS_V2_STALE_DATA
    _title = "Stale data"
    _message = "The data has not been updated for more than 30 minutes, a background update is in progress"
    _level = NotificationLevel.info


class DataUpdateFailureNotification(BaseNotification):
    def __init__(self, err_code: str, request_id: Optional[str]) -> None:
        super().__init__()
        self.err_code = err_code
        self.request_id = request_id or "unknown"

    type = NOTIF_TYPE_GSHEETS_V2_DATA_UPDATE_FAILURE
    _title = "Data update failed"
    _message = (
        "The displayed data may be outdated due to the failure of the last update.\n"
        "Reason: {err_code}, Request-ID: {request_id}."
    )
    _level = NotificationLevel.warning

    @property
    def message(self) -> str:
        return self._message.format(err_code=self.err_code, request_id=self.request_id)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/notifications.py
datalens-tech/datalens-backend
from dl_configs.connectors_settings import (
    ConnectorsConfigType,
    ConnectorSettingsBase,
)
from dl_configs.settings_loaders.meta_definition import required
from dl_core.connectors.settings.primitives import (
    ConnectorSettingsDefinition,
    get_connectors_settings_config,
)

from dl_connector_bundle_chs3.chs3_base.core.settings import (
    ConnectorsDataFileBase,
    FileS3ConnectorSettings,
)


def gsheets_file_s3_settings_fallback(full_cfg: ConnectorsConfigType) -> dict[str, ConnectorSettingsBase]:
    cfg = get_connectors_settings_config(
        full_cfg,
        object_like_config_key="FILE",
        connector_data_class=ConnectorsDataFileBase,
    )
    if cfg is None:
        return {}
    return dict(
        GSHEETS_V2=FileS3ConnectorSettings(  # type: ignore
            HOST=cfg.CONN_FILE_CH_HOST,
            PORT=cfg.CONN_FILE_CH_PORT,
            USERNAME=cfg.CONN_FILE_CH_USERNAME,
            PASSWORD=required(str),
            ACCESS_KEY_ID=required(str),
            SECRET_ACCESS_KEY=required(str),
            S3_ENDPOINT=full_cfg.S3_ENDPOINT_URL,
            BUCKET=full_cfg.FILE_UPLOADER_S3_PERSISTENT_BUCKET_NAME,
        )
    )


class GSheetsFileS3SettingDefinition(ConnectorSettingsDefinition):
    settings_class = FileS3ConnectorSettings
    fallback = gsheets_file_s3_settings_fallback

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/settings.py
datalens-tech/datalens-backend
from marshmallow import fields

from dl_core.us_manager.storage_schemas.connection import BaseConnectionDataStorageSchema

from dl_connector_bundle_chs3.chs3_base.core.storage_schemas.connection import (
    BaseFileConnectionDataStorageSchema,
    BaseFileConnectionSourceStorageSchema,
)
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection


class GSheetsFileConnectionSourceStorageSchema(BaseFileConnectionSourceStorageSchema):
    TARGET_CLS = GSheetsFileS3Connection.FileDataSource

    spreadsheet_id = fields.String(allow_none=True, load_default=None)
    sheet_id = fields.Integer(allow_none=True, load_default=None)
    first_line_is_header = fields.Boolean(allow_none=True, load_default=None)
    data_updated_at = fields.DateTime()


class GSheetsFileConnectionDataStorageSchema(
    BaseFileConnectionDataStorageSchema,
    BaseConnectionDataStorageSchema[GSheetsFileS3Connection.DataModel],
):
    TARGET_CLS = GSheetsFileS3Connection.DataModel

    sources = fields.Nested(GSheetsFileConnectionSourceStorageSchema, many=True)
    refresh_token = fields.String(load_default=None)
    refresh_enabled = fields.Boolean()

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/storage_schemas/connection.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.chs3_base.core.storage_schemas.data_source_spec import (
    BaseFileS3DataSourceSpecStorageSchema,
)
from dl_connector_bundle_chs3.chs3_gsheets.core.data_source_spec import GSheetsFileS3DataSourceSpec


class GSheetsFileS3DataSourceSpecStorageSchema(BaseFileS3DataSourceSpecStorageSchema):
    TARGET_CLS = GSheetsFileS3DataSourceSpec

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/storage_schemas/data_source_spec.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Any
import uuid

from dl_core.us_manager.us_manager_sync import SyncUSManager

from dl_connector_bundle_chs3.chs3_gsheets.core.constants import CONNECTION_TYPE_GSHEETS_V2
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection


def make_saved_gsheets_v2_connection(
    sync_usm: SyncUSManager,
    sources: list[GSheetsFileS3Connection.FileDataSource],
    **kwargs: Any,
) -> GSheetsFileS3Connection:
    conn_type = CONNECTION_TYPE_GSHEETS_V2

    conn_name = "{} test conn {}".format(conn_type.name, uuid.uuid4())
    conn = GSheetsFileS3Connection.create_from_dict(
        data_dict=GSheetsFileS3Connection.DataModel(
            sources=sources,
        ),
        ds_key=conn_name,
        type_=CONNECTION_TYPE_GSHEETS_V2.name,
        us_manager=sync_usm,
        **kwargs,
    )
    sync_usm.save(conn)
    return conn

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/testing/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import logging
from typing import (
    ClassVar,
    Optional,
)

import attr

from dl_constants.enums import (
    DataSourceRole,
    FileProcessingStatus,
)
from dl_core.services_registry.file_uploader_client_factory import GSheetsFileSourceDesc
from dl_utils.utils import DataKey

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3.chs3_gsheets.core.constants import SOURCE_TYPE_GSHEETS_V2


LOGGER = logging.getLogger(__name__)


class GSheetsFileS3Connection(BaseFileS3Connection):
    source_type = SOURCE_TYPE_GSHEETS_V2
    allowed_source_types = frozenset((SOURCE_TYPE_GSHEETS_V2,))

    editable_data_source_parameters: ClassVar[
        tuple[str, ...]
    ] = BaseFileS3Connection.editable_data_source_parameters + (
        "spreadsheet_id",
        "sheet_id",
        "first_line_is_header",
        "data_updated_at",
    )

    @attr.s(eq=False, kw_only=True)
    class FileDataSource(BaseFileS3Connection.FileDataSource):
        spreadsheet_id: Optional[str] = attr.ib(default=None)
        sheet_id: Optional[int] = attr.ib(default=None)
        first_line_is_header: Optional[bool] = attr.ib(default=None)
        data_updated_at: datetime.datetime = attr.ib(factory=lambda: datetime.datetime.now(datetime.timezone.utc))

        def __hash__(self) -> int:
            raw_schema = tuple(self.raw_schema) if self.raw_schema is not None else tuple()
            return hash(
                (
                    self.id,
                    self.file_id,
                    self.title,
                    self.s3_filename,
                    raw_schema,
                    self.status,
                    self.sheet_id,
                    self.spreadsheet_id,
                )
            )

        def str_for_hash(self) -> str:
            return ",".join(
                [
                    super().str_for_hash(),
                    str(self.spreadsheet_id),
                    str(self.sheet_id),
                ]
            )

        def get_desc(self) -> GSheetsFileSourceDesc:
            return GSheetsFileSourceDesc(
                file_id=self.file_id,
                source_id=self.id,
                title=self.title,
                raw_schema=self.raw_schema,
                preview_id=self.preview_id,
                spreadsheet_id=self.spreadsheet_id,
                sheet_id=self.sheet_id,
                first_line_is_header=self.first_line_is_header,
            )

    @attr.s(eq=False, kw_only=True)
    class DataModel(BaseFileS3Connection.DataModel):
        sources: list["GSheetsFileS3Connection.FileDataSource"] = attr.ib()  # type: ignore

        refresh_token: Optional[str] = attr.ib(default=None, repr=False)
        refresh_enabled: bool = attr.ib(default=False)

        def oldest_data_update_time(
            self, exclude_statuses: Optional[set[FileProcessingStatus]] = None
        ) -> Optional[datetime.datetime]:
            if exclude_statuses is None:
                exclude_statuses = set()

            data_updated_list = [src.data_updated_at for src in self.sources if src.status not in exclude_statuses]
            if not data_updated_list:
                return None

            return min(data_updated_list)

        @classmethod
        def get_secret_keys(cls) -> set[DataKey]:
            return {DataKey(parts=("refresh_token",))}

    data: DataModel

    @property
    def authorized(self) -> bool:
        return self.data.refresh_token is not None

    def restore_source_params_from_orig(self, src_id: str, original_version: BaseFileS3Connection) -> None:
        orig_src = original_version.get_file_source_by_id(src_id)
        assert isinstance(orig_src, GSheetsFileS3Connection.FileDataSource)
        self.update_data_source(
            src_id,
            role=DataSourceRole.origin,
            raw_schema=orig_src.raw_schema,
            file_id=orig_src.file_id,
            s3_filename=orig_src.s3_filename,
            status=orig_src.status,
            preview_id=orig_src.preview_id,
            first_line_is_header=orig_src.first_line_is_header,
            spreadsheet_id=orig_src.spreadsheet_id,
            sheet_id=orig_src.sheet_id,
        )

    @property
    def allow_public_usage(self) -> bool:
        return True

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_gsheets/core/us_connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import (
    fields,
    validate,
)

from dl_api_connector.api_schema.extras import FieldExtra

from dl_connector_bundle_chs3.chs3_base.api.api_schema.connection import BaseFileS3ConnectionSchema
from dl_connector_bundle_chs3.file.api.api_schema.source import FileSourceSchema
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection


class FileS3ConnectionSchema(BaseFileS3ConnectionSchema):
    TARGET_CLS = FileS3Connection

    sources = fields.Nested(
        FileSourceSchema,
        many=True,
        attribute="data.sources",
        bi_extra=FieldExtra(editable=True),
        validate=validate.Length(min=1, max=10),
    )

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/api/api_schema/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import (
    RAISE,
    Schema,
    fields,
)

from dl_constants.enums import UserDataType

from dl_connector_bundle_chs3.chs3_base.api.api_schema.connection import BaseFileSourceSchema
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection


class FileSourceColumnTypeSchema(Schema):
    class Meta:
        unknown = RAISE

    name = fields.String()
    user_type = fields.Enum(UserDataType)


class FileSourceSchema(BaseFileSourceSchema):
    class Meta(BaseFileSourceSchema.Meta):
        target = FileS3Connection.FileDataSource

    column_types = fields.Nested(FileSourceColumnTypeSchema, many=True, load_only=True)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/api/api_schema/source.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_bundle_chs3.chs3_base.api.i18n.localizer import Translatable


class FileS3ConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-file")

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/api/connection_info.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_connector_bundle_chs3.chs3_base.api.connector import (
    BaseFileS3ApiConnectionDefinition,
    BaseFileS3ApiConnector,
    BaseFileS3TableApiSourceDefinition,
)
from dl_connector_bundle_chs3.file.api.api_schema.connection import FileS3ConnectionSchema
from dl_connector_bundle_chs3.file.api.connection_info import FileS3ConnectionInfoProvider
from dl_connector_bundle_chs3.file.core.connector import (
    FileS3CoreConnectionDefinition,
    FileS3CoreConnector,
    FileS3TableCoreSourceDefinition,
)


class FileS3TableApiSourceDefinition(BaseFileS3TableApiSourceDefinition):
    core_source_def_cls = FileS3TableCoreSourceDefinition


class FileS3ApiConnectionDefinition(BaseFileS3ApiConnectionDefinition):
    core_conn_def_cls = FileS3CoreConnectionDefinition
    api_generic_schema_cls = FileS3ConnectionSchema
    info_provider_cls = FileS3ConnectionInfoProvider


class FileS3ApiConnector(BaseFileS3ApiConnector):
    core_connector_cls = FileS3CoreConnector
    connection_definitions = (FileS3ApiConnectionDefinition,)
    source_definitions = (FileS3TableApiSourceDefinition,)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/api/connector.py
datalens-tech/datalens-backend
import logging

from dl_connector_bundle_chs3.chs3_base.core.adapter import BaseAsyncFileS3Adapter
from dl_connector_bundle_chs3.file.core.constants import CONNECTION_TYPE_FILE


LOGGER = logging.getLogger(__name__)


class AsyncFileS3Adapter(BaseAsyncFileS3Adapter):
    conn_type = CONNECTION_TYPE_FILE

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/adapter.py
datalens-tech/datalens-backend
import attr

from dl_connector_bundle_chs3.chs3_base.core.connection_executors import BaseFileS3AsyncAdapterConnExecutor
from dl_connector_bundle_chs3.file.core.adapter import AsyncFileS3Adapter


@attr.s(cmp=False, hash=False)
class FileS3AsyncAdapterConnExecutor(BaseFileS3AsyncAdapterConnExecutor):
    TARGET_ADAPTER_CLS = AsyncFileS3Adapter  # type: ignore  # TODO: FIX

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/connection_executors.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.chs3_base.core.connector import (
    BaseFileS3CoreConnectionDefinition,
    BaseFileS3CoreConnector,
    BaseFileS3TableCoreSourceDefinition,
)
from dl_connector_bundle_chs3.file.core.adapter import AsyncFileS3Adapter
from dl_connector_bundle_chs3.file.core.connection_executors import FileS3AsyncAdapterConnExecutor
from dl_connector_bundle_chs3.file.core.constants import (
    CONNECTION_TYPE_FILE,
    SOURCE_TYPE_FILE_S3_TABLE,
)
from dl_connector_bundle_chs3.file.core.data_source import FileS3DataSource
from dl_connector_bundle_chs3.file.core.data_source_spec import FileS3DataSourceSpec
from dl_connector_bundle_chs3.file.core.lifecycle import FileS3ConnectionLifecycleManager
from dl_connector_bundle_chs3.file.core.settings import FileS3SettingDefinition
from dl_connector_bundle_chs3.file.core.storage_schemas.connection import FileConnectionDataStorageSchema
from dl_connector_bundle_chs3.file.core.storage_schemas.data_source_spec import FileS3DataSourceSpecStorageSchema
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_clickhouse.core.clickhouse_base.adapters import ClickHouseAdapter


class FileS3CoreConnectionDefinition(BaseFileS3CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_FILE
    connection_cls = FileS3Connection
    us_storage_schema_cls = FileConnectionDataStorageSchema
    sync_conn_executor_cls = FileS3AsyncAdapterConnExecutor
    async_conn_executor_cls = FileS3AsyncAdapterConnExecutor
    lifecycle_manager_cls = FileS3ConnectionLifecycleManager
    dialect_string = "bi_clickhouse"
    settings_definition = FileS3SettingDefinition


class FileS3TableCoreSourceDefinition(BaseFileS3TableCoreSourceDefinition):
    source_type = SOURCE_TYPE_FILE_S3_TABLE
    source_cls = FileS3DataSource
    source_spec_cls = FileS3DataSourceSpec
    us_storage_schema_cls = FileS3DataSourceSpecStorageSchema


class FileS3CoreConnector(BaseFileS3CoreConnector):
    connection_definitions = (FileS3CoreConnectionDefinition,)
    source_definitions = (FileS3TableCoreSourceDefinition,)
    rqe_adapter_classes = frozenset({ClickHouseAdapter, AsyncFileS3Adapter})

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
)


CONNECTION_TYPE_FILE = ConnectionType.declare("file")
SOURCE_TYPE_FILE_S3_TABLE = DataSourceType.declare("FILE_S3_TABLE")

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.enums import DataSourceType

from dl_connector_bundle_chs3.chs3_base.core.data_source import BaseFileS3DataSource
from dl_connector_bundle_chs3.file.core.constants import (
    CONNECTION_TYPE_FILE,
    SOURCE_TYPE_FILE_S3_TABLE,
)


class FileS3DataSource(BaseFileS3DataSource):
    conn_type = CONNECTION_TYPE_FILE

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type in {
            SOURCE_TYPE_FILE_S3_TABLE,
        }

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/data_source.py
datalens-tech/datalens-backend
import attr

from dl_connector_bundle_chs3.chs3_base.core.data_source_spec import BaseFileS3DataSourceSpec


@attr.s
class FileS3DataSourceSpec(BaseFileS3DataSourceSpec):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/data_source_spec.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.chs3_base.core.lifecycle import BaseFileS3ConnectionLifecycleManager
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection


class FileS3ConnectionLifecycleManager(BaseFileS3ConnectionLifecycleManager):
    ENTRY_CLS = FileS3Connection

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/lifecycle.py
datalens-tech/datalens-backend
from dl_configs.connectors_settings import (
    ConnectorsConfigType,
    ConnectorSettingsBase,
)
from dl_configs.settings_loaders.meta_definition import required
from dl_core.connectors.settings.primitives import (
    ConnectorSettingsDefinition,
    get_connectors_settings_config,
)

from dl_connector_bundle_chs3.chs3_base.core.settings import (
    ConnectorsDataFileBase,
    FileS3ConnectorSettings,
)


def file_s3_settings_fallback(full_cfg: ConnectorsConfigType) -> dict[str, ConnectorSettingsBase]:
    cfg = get_connectors_settings_config(
        full_cfg,
        object_like_config_key="FILE",
        connector_data_class=ConnectorsDataFileBase,
    )
    if cfg is None:
        return {}
    return dict(
        FILE=FileS3ConnectorSettings(  # type: ignore
            HOST=cfg.CONN_FILE_CH_HOST,
            PORT=cfg.CONN_FILE_CH_PORT,
            USERNAME=cfg.CONN_FILE_CH_USERNAME,
            PASSWORD=required(str),
            ACCESS_KEY_ID=required(str),
            SECRET_ACCESS_KEY=required(str),
            S3_ENDPOINT=full_cfg.S3_ENDPOINT_URL,
            BUCKET=full_cfg.FILE_UPLOADER_S3_PERSISTENT_BUCKET_NAME,
        )
    )


class FileS3SettingDefinition(ConnectorSettingsDefinition):
    settings_class = FileS3ConnectorSettings
    fallback = file_s3_settings_fallback

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/settings.py
datalens-tech/datalens-backend
from marshmallow import fields

from dl_core.us_manager.storage_schemas.connection import BaseConnectionDataStorageSchema

from dl_connector_bundle_chs3.chs3_base.core.storage_schemas.connection import (
    BaseFileConnectionDataStorageSchema,
    BaseFileConnectionSourceStorageSchema,
)
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection


class FileConnectionSourceStorageSchema(BaseFileConnectionSourceStorageSchema):
    TARGET_CLS = FileS3Connection.FileDataSource


class FileConnectionDataStorageSchema(
    BaseFileConnectionDataStorageSchema,
    BaseConnectionDataStorageSchema[FileS3Connection.DataModel],
):
    TARGET_CLS = FileS3Connection.DataModel

    sources = fields.Nested(FileConnectionSourceStorageSchema, many=True)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/storage_schemas/connection.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.chs3_base.core.storage_schemas.data_source_spec import (
    BaseFileS3DataSourceSpecStorageSchema,
)
from dl_connector_bundle_chs3.file.core.data_source_spec import FileS3DataSourceSpec


class FileS3DataSourceSpecStorageSchema(BaseFileS3DataSourceSpecStorageSchema):
    TARGET_CLS = FileS3DataSourceSpec

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/storage_schemas/data_source_spec.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Any
import uuid

from dl_core.us_manager.us_manager_sync import SyncUSManager

from dl_connector_bundle_chs3.file.core.constants import CONNECTION_TYPE_FILE
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection


def make_saved_file_connection(
    sync_usm: SyncUSManager,
    sources: list[FileS3Connection.FileDataSource],
    **kwargs: Any,
) -> FileS3Connection:
    conn_type = CONNECTION_TYPE_FILE

    conn_name = "{} test conn {}".format(conn_type.name, uuid.uuid4())
    conn = FileS3Connection.create_from_dict(
        data_dict=FileS3Connection.DataModel(
            sources=sources,
        ),
        ds_key=conn_name,
        type_=CONNECTION_TYPE_FILE.name,
        us_manager=sync_usm,
        **kwargs,
    )
    sync_usm.save(conn)
    return conn

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/testing/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging

import attr

from dl_constants.enums import DataSourceRole
from dl_core.services_registry.file_uploader_client_factory import FileSourceDesc

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3.file.core.constants import SOURCE_TYPE_FILE_S3_TABLE


LOGGER = logging.getLogger(__name__)


class FileS3Connection(BaseFileS3Connection):
    source_type = SOURCE_TYPE_FILE_S3_TABLE
    allowed_source_types = frozenset((SOURCE_TYPE_FILE_S3_TABLE,))

    @attr.s(eq=False, kw_only=True)
    class FileDataSource(BaseFileS3Connection.FileDataSource):
        column_types: list[dict] = attr.ib(factory=list)  # should not be saved in US

        def get_desc(self) -> FileSourceDesc:
            return FileSourceDesc(
                file_id=self.file_id,
                source_id=self.id,
                title=self.title,
                raw_schema=self.column_types or self.raw_schema,
                preview_id=self.preview_id,
            )

    def restore_source_params_from_orig(self, src_id: str, original_version: BaseFileS3Connection) -> None:
        orig_src = original_version.get_file_source_by_id(src_id)
        assert isinstance(orig_src, FileS3Connection.FileDataSource)
        self.update_data_source(
            src_id,
            role=DataSourceRole.origin,
            raw_schema=orig_src.raw_schema,
            file_id=orig_src.file_id,
            s3_filename=orig_src.s3_filename,
            status=orig_src.status,
            preview_id=orig_src.preview_id,
        )

    @property
    def allow_public_usage(self) -> bool:
        return True

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/file/core/us_connection.py
datalens-tech/datalens-backend
import abc
import datetime
import math
from typing import Generator

import attr
import pytest

from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_constants.enums import (
    ConnectionType,
    UserDataType,
)
from dl_core.services_registry.file_uploader_client_factory import (
    FileSourceDesc,
    FileUploaderClient,
    FileUploaderClientFactory,
    SourceInternalParams,
    SourcePreview,
)
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import (
    C,
    make_sample_data,
)

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.base import (
    FILE_CONN_TV,
    BaseCHS3TestClass,
)
from dl_connector_bundle_chs3_tests.db.config import API_TEST_CONFIG


@attr.s
class FileUploaderClientMockup(FileUploaderClient):
    # In reality, we already receive json-serialized values for the preview from the source,
    # so tune value generators here as well
    _VALUE_GENERATORS = {
        UserDataType.date: lambda rn, ts, **kwargs: (ts.date() + datetime.timedelta(days=rn)).isoformat(),
        UserDataType.datetime: lambda rn, ts, **kwargs: (ts + datetime.timedelta(days=rn / math.pi)).isoformat(),
        UserDataType.genericdatetime: lambda rn, ts, **kwargs: (ts + datetime.timedelta(days=rn / math.pi)).isoformat(),
    }

    async def get_preview(self, src: FileSourceDesc) -> SourcePreview:
        if src.raw_schema is None:
            return SourcePreview(source_id=src.source_id, preview=[])
        cols = [
            C(sch_col.name, sch_col.user_type, sch_col.nullable, vg=self._VALUE_GENERATORS.get(sch_col.user_type))
            for sch_col in src.raw_schema
        ]
        preview_dicts = make_sample_data(cols, rows=20)
        preview = [[row[c.name] for c in cols] for row in preview_dicts]
        return SourcePreview(source_id=src.source_id, preview=preview)

    async def get_internal_params(self, src: FileSourceDesc) -> SourceInternalParams:
        """Should normally return actual connection params from US, but this will do for tests"""

        return SourceInternalParams(
            preview_id=src.preview_id,
            raw_schema=src.raw_schema,
        )


class CHS3ConnectionApiTestBase(BaseCHS3TestClass[FILE_CONN_TV], ConnectionTestBase, metaclass=abc.ABCMeta):
    bi_compeng_pg_on = False

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def monkeyclass(self) -> Generator[pytest.MonkeyPatch, None, None]:
        with pytest.MonkeyPatch.context() as mp:
            yield mp

    @pytest.fixture(scope="class", autouse=True)
    def patch_file_uploader_client(self, monkeyclass: pytest.MonkeyPatch) -> None:
        monkeyclass.setattr(FileUploaderClientFactory, "_file_uploader_client_cls", FileUploaderClientMockup)

    @pytest.fixture(scope="class")
    def connectors_settings(self) -> dict[ConnectionType, ConnectorSettingsBase]:
        return {self.conn_type: self.connection_settings}

    @pytest.fixture(scope="function")
    def saved_connection_id(
        self,
        control_api_sync_client: SyncHttpClientBase,
        connection_params: dict,
        sync_us_manager: SyncUSManager,
        sample_file_data_source: BaseFileS3Connection.FileDataSource,
    ) -> Generator[str, None, None]:
        """
        Normally connections are updated by file uploader worker,
        but that would require the whole file-uploader pipeline to be setup in tests
        """

        with super().create_connection(
            control_api_sync_client=control_api_sync_client,
            connection_params=connection_params,
        ) as conn_id:
            conn = sync_us_manager.get_by_id(conn_id, BaseFileS3Connection)
            for src in conn.data.sources:
                src.status = sample_file_data_source.status
                src.raw_schema = sample_file_data_source.raw_schema
                src.s3_filename = sample_file_data_source.s3_filename
            sync_us_manager.save(conn)
            yield conn_id

    @abc.abstractmethod
    @pytest.fixture(scope="function")
    def connection_params(self, sample_file_data_source: BaseFileS3Connection.FileDataSource) -> dict:
        raise NotImplementedError()

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/base/api/base.py
datalens-tech/datalens-backend
import abc
from http import HTTPStatus
import uuid

from flask.testing import FlaskClient
import pytest

from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite
from dl_core.exc import DataSourceTitleConflict
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3_tests.db.base.api.base import CHS3ConnectionApiTestBase
from dl_connector_bundle_chs3_tests.db.base.core.base import FILE_CONN_TV


class CHS3ConnectionTestSuite(
    CHS3ConnectionApiTestBase[FILE_CONN_TV],
    DefaultConnectorConnectionTestSuite,
    metaclass=abc.ABCMeta,
):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultConnectorConnectionTestSuite.test_test_connection: "Not implemented",
            DefaultConnectorConnectionTestSuite.test_cache_ttl_sec_override: "Unavailable for CHS3 connectors",
        },
    )

    @abc.abstractmethod
    @pytest.fixture(scope="function")
    def single_new_conn_source_params(self) -> dict:
        raise NotImplementedError()

    def test_add_and_drop_connection_source(
        self,
        client: FlaskClient,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
        single_new_conn_source_params: dict,
    ) -> None:
        conn_id = saved_connection_id
        orig_conn = sync_us_manager.get_by_id(conn_id, BaseFileS3Connection)

        new_source = single_new_conn_source_params
        add_resp = client.put(
            "/api/v1/connections/{}".format(conn_id),
            json={
                "sources": [
                    {"id": orig_conn.data.sources[0].id, "title": orig_conn.data.sources[0].title},
                    new_source,
                ],
            },
        )
        assert add_resp.status_code == HTTPStatus.OK, add_resp.json

        conn = sync_us_manager.get_by_id(conn_id)
        assert len(conn.data.sources) == len(orig_conn.data.sources) + 1

        drop_resp = client.put(
            "/api/v1/connections/{}".format(conn_id),
            json={
                "sources": [
                    {"id": orig_conn.data.sources[0].id, "title": orig_conn.data.sources[0].title},
                ],
            },
        )
        assert drop_resp.status_code == HTTPStatus.OK, drop_resp.json

        conn = sync_us_manager.get_by_id(conn_id)
        assert len(conn.data.sources) == len(orig_conn.data.sources)

    def test_rename_connection_source(
        self,
        client: FlaskClient,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
    ) -> None:
        conn_id = saved_connection_id
        orig_conn = sync_us_manager.get_by_id(conn_id, BaseFileS3Connection)

        resp = client.put(
            "/api/v1/connections/{}".format(conn_id),
            json={
                "sources": [
                    {"id": orig_conn.data.sources[0].id, "title": "renamed source"},
                ],
            },
        )
        assert resp.status_code == HTTPStatus.OK, resp.json

        conn = sync_us_manager.get_by_id(conn_id)
        assert conn.data.sources[0].title == "renamed source"

    def test_replace_connection_source(
        self,
        client: FlaskClient,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
        single_new_conn_source_params: dict,
    ) -> None:
        conn_id = saved_connection_id
        orig_conn = sync_us_manager.get_by_id(conn_id, BaseFileS3Connection)

        source_to_replace = orig_conn.data.sources[0]
        new_source = single_new_conn_source_params
        resp = client.put(
            "/api/v1/connections/{}".format(conn_id),
            json={
                "sources": [new_source],
                "replace_sources": [
                    {
                        "old_source_id": source_to_replace.id,
                        "new_source_id": new_source["id"],
                    },
                ],
            },
        )
        assert resp.status_code == HTTPStatus.OK, resp.json
        conn = sync_us_manager.get_by_id(conn_id, BaseFileS3Connection)

        old_source_ids = set(src.id for src in orig_conn.data.sources)
        new_source_ids = set(src.id for src in conn.data.sources)
        assert old_source_ids == new_source_ids
        new_replaced_source = conn.get_file_source_by_id(source_to_replace.id)
        assert new_replaced_source.file_id != source_to_replace.file_id

    def test_consistency_checks_pass_not_configured_source(
        self,
        client: FlaskClient,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
        single_new_conn_source_params: dict,
    ) -> None:
        conn = sync_us_manager.get_by_id(saved_connection_id, BaseFileS3Connection)

        new_source = single_new_conn_source_params
        new_source.pop("file_id")  # as if this source already exists
        resp = client.put(
            "/api/v1/connections/{}".format(saved_connection_id),
            json={
                "sources": [
                    {"id": conn.data.sources[0].id, "title": conn.data.sources[0].title},
                    new_source,
                ],
            },
        )
        assert resp.status_code == HTTPStatus.BAD_REQUEST, resp.json
        details: dict[str, list[str]] = resp.json["details"]
        assert details == {
            "not_configured_not_saved": [new_source["id"]],
        }

    def test_consistency_checks_replace_non_existent_source(
        self,
        client: FlaskClient,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
        single_new_conn_source_params: dict,
    ) -> None:
        conn = sync_us_manager.get_by_id(saved_connection_id, BaseFileS3Connection)

        replaced_source_id = str(uuid.uuid4())  # replacing a non-existent source
        new_source = single_new_conn_source_params
        resp = client.put(
            "/api/v1/connections/{}".format(saved_connection_id),
            json={
                "sources": [
                    {"id": conn.data.sources[0].id, "title": conn.data.sources[0].title},
                    new_source,
                ],
                "replace_sources": [
                    {
                        "old_source_id": replaced_source_id,
                        "new_source_id": new_source["id"],
                    },
                ],
            },
        )
        assert resp.status_code == HTTPStatus.BAD_REQUEST, resp.json
        details: dict[str, list[str]] = resp.json["details"]
        assert details == {
            "replaced_not_saved": [replaced_source_id],
        }

    def test_consistency_checks_non_unique_titles(
        self,
        client: FlaskClient,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
        single_new_conn_source_params: dict,
    ) -> None:
        conn = sync_us_manager.get_by_id(saved_connection_id, BaseFileS3Connection)

        new_source = single_new_conn_source_params
        new_source["title"] = conn.data.sources[0].title
        resp = client.put(
            "/api/v1/connections/{}".format(saved_connection_id),
            json={
                "sources": [
                    {"id": conn.data.sources[0].id, "title": conn.data.sources[0].title},
                    new_source,
                ],
            },
        )
        assert resp.status_code == HTTPStatus.BAD_REQUEST, resp.json
        assert resp.json["message"] == DataSourceTitleConflict().message

    def test_table_name_spoofing(
        self,
        client: FlaskClient,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
    ) -> None:
        usm = sync_us_manager
        orig_conn = usm.get_by_id(saved_connection_id, BaseFileS3Connection)
        orig_filename = orig_conn.data.sources[0].s3_filename

        fake_filename = "hack_me.native"
        resp = client.put(
            "/api/v1/connections/{}".format(saved_connection_id),
            json={
                "sources": [
                    {
                        "id": orig_conn.data.sources[0].id,
                        "title": orig_conn.data.sources[0].title,
                        "s3_filename": fake_filename,
                    }
                ],
            },
        )
        assert resp.status_code == HTTPStatus.BAD_REQUEST, resp.json
        assert resp.json["sources"]["0"]["s3_filename"] == ["Unknown field."]
        conn = usm.get_by_id(saved_connection_id, BaseFileS3Connection)
        assert conn.data.sources[0].s3_filename == orig_filename

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/base/api/connection.py
datalens-tech/datalens-backend
import abc
from http import HTTPStatus

from flask.testing import FlaskClient
import pytest

from dl_api_client.dsmaker.api.data_api import SyncHttpDataApiV2
from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib_testing.connector.data_api_suites import DefaultConnectorDataResultTestSuite
from dl_api_lib_testing.data_api_base import (
    DataApiTestParams,
    StandardizedDataApiTestBase,
)
from dl_constants.enums import (
    DataSourceRole,
    FileProcessingStatus,
)
from dl_core.us_manager.us_manager_sync import SyncUSManager

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3_tests.db.base.api.dataset import CHS3DatasetTestBase
from dl_connector_bundle_chs3_tests.db.base.core.base import FILE_CONN_TV


class CHS3DataApiTestBase(CHS3DatasetTestBase[FILE_CONN_TV], StandardizedDataApiTestBase, metaclass=abc.ABCMeta):
    mutation_caches_on = False


class CHS3DataResultTestSuite(
    CHS3DataApiTestBase[FILE_CONN_TV],
    DefaultConnectorDataResultTestSuite,
    metaclass=abc.ABCMeta,
):
    @pytest.mark.xfail(reason="TODO")
    def test_date32(self) -> None:  # TODO implement
        assert 0

    @pytest.mark.parametrize(
        "case, expected_status",
        (
            ("w_status_in_progress", HTTPStatus.BAD_REQUEST),
            ("w_raw_schema_removed", HTTPStatus.OK),  # ok, because raw schema is stored in the dataset
        ),
    )
    def test_file_not_ready_result(
        self,
        case: str,
        expected_status: HTTPStatus,
        sync_us_manager: SyncUSManager,
        data_api: SyncHttpDataApiV2,
        saved_dataset: Dataset,
        saved_connection_id: str,
        sample_file_data_source: BaseFileS3Connection.FileDataSource,
        data_api_test_params: DataApiTestParams,
    ) -> None:
        ds = saved_dataset
        conn = sync_us_manager.get_by_id(saved_connection_id, BaseFileS3Connection)
        conn.update_data_source(
            id=sample_file_data_source.id,
            role=DataSourceRole.origin,
            status=FileProcessingStatus.in_progress if case == "w_status_in_progress" else FileProcessingStatus.ready,
            remove_raw_schema=True if case == "w_raw_schema_removed" else False,
        )
        sync_us_manager.save(conn)

        result_resp = data_api.get_result(
            dataset=ds, fields=[ds.find_field(title=data_api_test_params.date_field)], fail_ok=True
        )
        assert result_resp.status_code == expected_status, result_resp.json

    def test_dataset_with_removed_file(
        self,
        client: FlaskClient,
        control_api: SyncHttpDatasetApiV1,
        data_api: SyncHttpDataApiV2,
        saved_dataset: Dataset,
        saved_connection_id: str,
        data_api_test_params: DataApiTestParams,
    ) -> None:
        ds = saved_dataset
        ds.result_schema["Measure"] = ds.field(formula=f"SUM([{data_api_test_params.summable_field}])")

        ds_resp = control_api.apply_updates(dataset=ds, fail_ok=True)
        assert ds_resp.status_code == HTTPStatus.OK, ds_resp.response_errors
        ds = ds_resp.dataset
        ds = control_api.save_dataset(ds).dataset

        preview_resp = data_api.get_preview(dataset=ds)
        assert preview_resp.status_code == HTTPStatus.OK, preview_resp.json

        # remove source from the connection
        update_resp = client.put(
            "/api/v1/connections/{}".format(saved_connection_id),
            json={"sources": [{"id": "dummy", "title": "dummy", "file_id": "dummy"}]},
        )
        assert update_resp.status_code == HTTPStatus.OK, update_resp.json

        get_ds_resp = client.get(f"/api/v1/datasets/{ds.id}/versions/draft")
        assert get_ds_resp.status_code == HTTPStatus.OK, get_ds_resp.json

        refresh_resp = control_api.refresh_dataset_sources(ds, [ds.sources[0].id], fail_ok=True)
        assert refresh_resp.status_code == HTTPStatus.BAD_REQUEST
        assert refresh_resp.json["code"] == "ERR.DS_API.VALIDATION.ERROR"

        preview_resp = data_api.get_preview(dataset=ds, fail_ok=True)
        assert preview_resp.status_code == HTTPStatus.BAD_REQUEST, preview_resp.json
        assert preview_resp.json["code"] == "ERR.DS_API.DB.SOURCE_DOES_NOT_EXIST"

        result_resp = data_api.get_result(dataset=ds, fields=[ds.result_schema["Measure"]], fail_ok=True)
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST, preview_resp.json
        assert result_resp.json["code"] == "ERR.DS_API.DB.SOURCE_DOES_NOT_EXIST"

    def test_table_name_spoofing(
        self,
        data_api: SyncHttpDataApiV2,
        saved_dataset: Dataset,
        data_api_test_params: DataApiTestParams,
    ) -> None:
        fake_filename = "hack_me.native"
        fake_parameters = {
            "db_name": "fake db_name",
            "db_version": "fake db_version",
            "table_name": fake_filename,
            "origin_source_id": "fake_source_id",
        }

        ds = saved_dataset
        ds.sources[0].parameters = fake_parameters
        result_resp = data_api.get_result(
            dataset=ds, fields=[ds.find_field(title=data_api_test_params.date_field)], fail_ok=True
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json

        preview_resp = data_api.get_preview(dataset=ds, fail_ok=True)
        assert preview_resp.status_code == HTTPStatus.BAD_REQUEST, preview_resp.json
        # ^ because of fake origin_source_id in parameters

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/base/api/data.py
datalens-tech/datalens-backend
import abc
from http import HTTPStatus

import pytest

from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite
from dl_api_lib_testing.dataset_base import DatasetTestBase
from dl_constants.enums import (
    DataSourceRole,
    FileProcessingStatus,
)
from dl_core.us_manager.us_manager_sync import SyncUSManager

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3_tests.db.base.api.base import CHS3ConnectionApiTestBase
from dl_connector_bundle_chs3_tests.db.base.core.base import FILE_CONN_TV


class CHS3DatasetTestBase(CHS3ConnectionApiTestBase[FILE_CONN_TV], DatasetTestBase, metaclass=abc.ABCMeta):
    @pytest.fixture(scope="function")
    def dataset_params(self, sample_file_data_source) -> dict:
        return dict(
            source_type=self.source_type.name,
            parameters=dict(
                origin_source_id=sample_file_data_source.id,
            ),
        )


class CHS3DatasetTestSuite(CHS3DatasetTestBase, DefaultConnectorDatasetTestSuite, metaclass=abc.ABCMeta):
    @pytest.mark.parametrize(
        "source_status, expected_status",
        (
            (FileProcessingStatus.ready, HTTPStatus.OK),
            (FileProcessingStatus.in_progress, HTTPStatus.BAD_REQUEST),
            (FileProcessingStatus.failed, HTTPStatus.BAD_REQUEST),
        ),
    )
    def test_add_dataset_source(
        self,
        sync_us_manager: SyncUSManager,
        saved_connection: FILE_CONN_TV,
        sample_file_data_source: BaseFileS3Connection.FileDataSource,
        control_api: SyncHttpDatasetApiV1,
        source_status: FileProcessingStatus,
        expected_status: int,
    ) -> None:
        saved_connection.update_data_source(
            id=sample_file_data_source.id,
            role=DataSourceRole.origin,
            status=source_status,
        )
        sync_us_manager.save(saved_connection)

        ds = Dataset()
        ds.sources["source_1"] = ds.source(
            source_type=self.source_type,
            connection_id=saved_connection.uuid,
            parameters=dict(
                origin_source_id=sample_file_data_source.id,
            ),
        )
        ds.source_avatars["avatar_1"] = ds.sources["source_1"].avatar()
        ds_resp = control_api.apply_updates(dataset=ds, fail_ok=True)
        assert ds_resp.status_code == expected_status, ds_resp.json

        ds = ds_resp.dataset
        ds_resp = control_api.save_dataset(dataset=ds)
        assert ds_resp.status_code == HTTPStatus.OK, ds_resp.response_errors

    def test_table_name_spoofing(
        self,
        control_api: SyncHttpDatasetApiV1,
        saved_connection_id: str,
    ) -> None:
        fake_parameters = {
            "db_name": "fake db_name",
            "db_version": "fake db_version",
            "table_name": "hack_me.native",
            "origin_source_id": "fake_source_id",
        }

        ds = Dataset()
        ds.sources["source_1"] = ds.source(
            source_type=self.source_type,
            connection_id=saved_connection_id,
            parameters=fake_parameters,
        )
        ds.source_avatars["avatar_1"] = ds.sources["source_1"].avatar()
        ds_resp = control_api.apply_updates(dataset=ds, fail_ok=True)
        assert ds_resp.status_code == HTTPStatus.BAD_REQUEST, ds_resp.json

    def test_update_connection_source(
        self,
        sync_us_manager: SyncUSManager,
        control_api: SyncHttpDatasetApiV1,
        saved_dataset: Dataset,
        saved_connection_id: str,
        sample_file_data_source: BaseFileS3Connection.FileDataSource,
    ) -> None:
        orig_ds = saved_dataset

        # --- change connection source: add one column ---
        conn = sync_us_manager.get_by_id(saved_connection_id, BaseFileS3Connection)
        new_field = sample_file_data_source.raw_schema[-1].clone(name="new_field")
        conn.update_data_source(
            id=sample_file_data_source.id,
            role=DataSourceRole.origin,
            raw_schema=sample_file_data_source.raw_schema + [new_field],
        )
        sync_us_manager.save(conn)
        # --- ---------------------------------------- ---

        refresh_resp = control_api.refresh_dataset_sources(orig_ds, [orig_ds.sources[0].id])
        assert refresh_resp.status_code == HTTPStatus.OK, refresh_resp.json
        ds = refresh_resp.dataset

        old_raw_schema, new_raw_schema = orig_ds.sources[0].raw_schema, ds.sources[0].raw_schema
        assert len(new_raw_schema) == len(old_raw_schema) + 1

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/base/api/dataset.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
from typing import (
    TYPE_CHECKING,
    AsyncGenerator,
    ClassVar,
    Generator,
    TypeVar,
)
import uuid

import attr
import pytest

from dl_api_commons.base_models import (
    RequestContextInfo,
    TenantCommon,
)
from dl_api_lib_testing.app import RedisSettingMaker
from dl_configs.settings_submodels import S3Settings
from dl_constants.enums import DataSourceType
from dl_core.db import (
    SchemaColumn,
    get_type_transformer,
)
from dl_core.services_registry import ServicesRegistry
from dl_core_testing.database import DbTable
from dl_core_testing.fixtures.primitives import FixtureTableSpec
from dl_core_testing.fixtures.sample_tables import TABLE_SPEC_SAMPLE_SUPERSTORE
from dl_core_testing.testcases.connection import BaseConnectionTestClass
from dl_file_uploader_worker_lib.utils.parsing_utils import get_field_id_generator
from dl_task_processor.processor import (
    DummyTaskProcessorFactory,
    TaskProcessorFactory,
)
from dl_testing.s3_utils import (
    create_s3_bucket,
    create_s3_client,
)

from dl_connector_bundle_chs3.chs3_base.core.testing.utils import create_s3_native_from_ch_table
from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3_tests.db import config as test_config
from dl_connector_clickhouse.db_testing.engine_wrapper import ClickhouseDbEngineConfig


if TYPE_CHECKING:
    from types_aiobotocore_s3 import S3Client as AsyncS3Client


FILE_CONN_TV = TypeVar("FILE_CONN_TV", bound=BaseFileS3Connection)


class BaseCHS3TestClass(BaseConnectionTestClass[FILE_CONN_TV], metaclass=abc.ABCMeta):
    core_test_config = test_config.CORE_TEST_CONFIG
    connection_settings = test_config.SR_CONNECTION_SETTINGS

    source_type: ClassVar[DataSourceType]

    @pytest.fixture(scope="session")
    def event_loop(self):
        """Avoid spontaneous event loop closes between tests"""
        loop = asyncio.get_event_loop()
        yield loop
        loop.close()

    @pytest.fixture(scope="session", autouse=True)
    # FIXME: This fixture is a temporary solution for failing core tests when they are run together with api tests
    def loop(self, event_loop: asyncio.AbstractEventLoop) -> Generator[asyncio.AbstractEventLoop, None, None]:
        asyncio.set_event_loop(event_loop)
        yield event_loop
        # Attempt to cover an old version of pytest-asyncio:
        # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
        asyncio.set_event_loop_policy(None)

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return test_config.DB_CH_URL

    @pytest.fixture(scope="class")
    def engine_config(self, db_url: str, engine_params: dict) -> ClickhouseDbEngineConfig:
        return ClickhouseDbEngineConfig(url=db_url, engine_params=engine_params)

    @pytest.fixture(scope="session")
    def conn_bi_context(self) -> RequestContextInfo:
        return RequestContextInfo(tenant=TenantCommon())

    @pytest.fixture(scope="session")
    def redis_setting_maker(self) -> RedisSettingMaker:
        bi_test_config = test_config.API_TEST_CONFIG
        return RedisSettingMaker(bi_test_config=bi_test_config)

    @pytest.fixture(scope="session")
    def s3_settings(self) -> S3Settings:
        return S3Settings(
            ENDPOINT_URL=test_config.S3_ENDPOINT_URL,
            ACCESS_KEY_ID=self.connection_settings.ACCESS_KEY_ID,
            SECRET_ACCESS_KEY=self.connection_settings.SECRET_ACCESS_KEY,
        )

    @pytest.fixture(scope="session")
    def task_processor_factory(self) -> TaskProcessorFactory:
        return DummyTaskProcessorFactory()

    @pytest.fixture(scope="session")
    def conn_sync_service_registry(
        self,
        conn_bi_context: RequestContextInfo,
        task_processor_factory: TaskProcessorFactory,
    ) -> ServicesRegistry:
        return self.service_registry_factory(
            conn_exec_factory_async_env=False,
            conn_bi_context=conn_bi_context,
            task_processor_factory=task_processor_factory,
        )

    @pytest.fixture(scope="session")
    def conn_async_service_registry(
        self,
        conn_bi_context: RequestContextInfo,
        task_processor_factory: TaskProcessorFactory,
    ) -> ServicesRegistry:
        return self.service_registry_factory(
            conn_exec_factory_async_env=True,
            conn_bi_context=conn_bi_context,
            task_processor_factory=task_processor_factory,
        )

    @pytest.fixture(scope="function")
    async def s3_client(self, s3_settings: S3Settings) -> AsyncS3Client:
        async with create_s3_client(s3_settings) as client:
            yield client

    @pytest.fixture(scope="function")
    async def s3_bucket(self, s3_client: AsyncS3Client) -> str:
        bucket_name = self.connection_settings.BUCKET
        await create_s3_bucket(s3_client, bucket_name)
        return bucket_name

    @pytest.fixture(scope="class")
    def sample_table_spec(self) -> FixtureTableSpec:
        return attr.evolve(TABLE_SPEC_SAMPLE_SUPERSTORE, nullable=True)

    def _get_s3_func_schema_for_table(self, table: DbTable) -> str:
        field_id_gen = get_field_id_generator(self.conn_type)
        tbl_schema = ", ".join(
            "{} {}".format(field_id_gen.make_field_id(dict(title=col.name, index=idx)), col.type.compile())
            for idx, col in enumerate(table.table.columns)
        )
        tbl_schema = tbl_schema.replace("()", "")  # String() -> String: type arguments are not needed here
        return tbl_schema

    @pytest.fixture(scope="function")
    async def sample_s3_file(
        self,
        s3_client: AsyncS3Client,
        s3_bucket: str,
        s3_settings: S3Settings,
        sample_table: DbTable,
    ) -> AsyncGenerator[str, None]:
        filename = f"my_file_{uuid.uuid4()}.native"

        tbl_schema = self._get_s3_func_schema_for_table(sample_table)
        create_s3_native_from_ch_table(filename, s3_bucket, s3_settings, sample_table, tbl_schema)

        yield filename

        await s3_client.delete_object(Bucket=s3_bucket, Key=filename)

    def _get_raw_schema_for_ch_table(self, table_spec: FixtureTableSpec) -> list[SchemaColumn]:
        field_id_gen = get_field_id_generator(self.conn_type)
        type_transformer = get_type_transformer(self.conn_type)
        raw_schema = [
            SchemaColumn(
                name=field_id_gen.make_field_id(dict(title=col[0], index=idx)),
                title=col[0],
                user_type=(user_type := table_spec.get_user_type_for_col(col[0])),
                native_type=type_transformer.type_user_to_native(user_type),
            )
            for idx, col in enumerate(table_spec.table_schema)
        ]
        return raw_schema

    @abc.abstractmethod
    @pytest.fixture(scope="function")
    def sample_file_data_source(self, sample_s3_file: str) -> BaseFileS3Connection.FileDataSource:
        raise NotImplementedError()

    @pytest.fixture(scope="function")
    def connection_creation_params(self, sample_file_data_source: BaseFileS3Connection.FileDataSource) -> dict:
        return dict(sources=[sample_file_data_source])

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/base/core/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc

from dl_core.us_connection_base import DataSourceTemplate
from dl_core_testing.testcases.connection import DefaultConnectionTestClass

from dl_connector_bundle_chs3_tests.db.base.core.base import (
    FILE_CONN_TV,
    BaseCHS3TestClass,
)


class CHS3ConnectionTestBase(
    BaseCHS3TestClass,
    DefaultConnectionTestClass[FILE_CONN_TV],
    metaclass=abc.ABCMeta,
):
    do_check_data_export_flag = False

    def check_saved_connection(self, conn: FILE_CONN_TV, params: dict) -> None:
        assert set(src.title for src in conn.data.sources) == set(src.title for src in params["sources"])

    def check_data_source_templates(
        self,
        conn: FILE_CONN_TV,
        dsrc_templates: list[DataSourceTemplate],
    ) -> None:
        assert dsrc_templates
        for dsrc_tmpl in dsrc_templates:
            conn_src = conn.get_file_source_by_id(dsrc_tmpl.parameters["origin_source_id"])
            assert conn_src.title == dsrc_tmpl.title

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/base/core/connection.py
datalens-tech/datalens-backend
import abc

import pytest
import sqlalchemy as sa

from dl_configs.settings_submodels import S3Settings
from dl_constants.enums import UserDataType
from dl_core.connection_executors import (
    AsyncConnExecutorBase,
    ConnExecutorQuery,
)
from dl_core.connection_models import DBIdent
from dl_core_testing.database import DbTable
from dl_core_testing.testcases.connection_executor import (
    DefaultAsyncConnectionExecutorTestSuite,
    DefaultSyncAsyncConnectionExecutorCheckBase,
    DefaultSyncConnectionExecutorTestSuite,
)
from dl_testing.regulated_test import RegulatedTestParams
from dl_testing.s3_utils import s3_tbl_func_maker

from dl_connector_bundle_chs3_tests.db.base.core.base import (
    FILE_CONN_TV,
    BaseCHS3TestClass,
)


class CHS3SyncAsyncConnectionExecutorTestBase(
    BaseCHS3TestClass,
    DefaultSyncAsyncConnectionExecutorCheckBase[FILE_CONN_TV],
    metaclass=abc.ABCMeta,
):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultAsyncConnectionExecutorTestSuite.test_get_db_version: "Not implemented",
            DefaultAsyncConnectionExecutorTestSuite.test_table_not_exists: "Assumes it always exists",
            DefaultAsyncConnectionExecutorTestSuite.test_get_table_schema_info: "Not implemented",
            DefaultAsyncConnectionExecutorTestSuite.test_get_table_schema_info_for_nonexistent_table: "Not implemented",
        },
    )

    @pytest.fixture(scope="function")
    def db_ident(self) -> DBIdent:
        pass


class CHS3SyncConnectionExecutorTestBase(
    CHS3SyncAsyncConnectionExecutorTestBase[FILE_CONN_TV],
    DefaultSyncConnectionExecutorTestSuite[FILE_CONN_TV],
    metaclass=abc.ABCMeta,
):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultSyncConnectionExecutorTestSuite.test_type_recognition: "Not implemented",
            DefaultSyncConnectionExecutorTestSuite.test_get_table_names: "Not implemented",
        },
    )


class CHS3AsyncConnectionExecutorTestBase(
    CHS3SyncAsyncConnectionExecutorTestBase[FILE_CONN_TV],
    DefaultAsyncConnectionExecutorTestSuite[FILE_CONN_TV],
    metaclass=abc.ABCMeta,
):
    async def test_select_data(
        self,
        sample_table: DbTable,
        saved_connection: FILE_CONN_TV,
        async_connection_executor: AsyncConnExecutorBase,
        s3_settings: S3Settings,
        sample_s3_file: str,
    ) -> None:
        schema_line = self._get_s3_func_schema_for_table(sample_table)
        s3_tbl_func = s3_tbl_func_maker(s3_settings)(
            for_="dba",
            conn_dto=saved_connection.get_conn_dto(),
            filename=sample_s3_file,
            file_fmt="Native",
            schema_line=schema_line,
        )
        file_columns = [sa.column(col_desc.split(" ")[0]) for col_desc in schema_line.split(", ")]
        # ^ "col1 String, col2 Int32, col3 Date32" -> [col1, col2, col3]

        n_rows = 3
        result = await async_connection_executor.execute(
            ConnExecutorQuery(
                query=sa.select(columns=file_columns)
                .select_from(sa.text(s3_tbl_func))
                .order_by(file_columns[0])
                .limit(n_rows),
                chunk_size=6,
            )
        )
        rows = await result.get_all()
        assert len(rows) == n_rows

    @pytest.mark.asyncio
    async def test_cast_row_to_output(
        self,
        sample_table: DbTable,
        saved_connection: FILE_CONN_TV,
        async_connection_executor: AsyncConnExecutorBase,
        s3_settings: S3Settings,
        sample_s3_file: str,
    ) -> None:
        schema_line = self._get_s3_func_schema_for_table(sample_table)
        s3_tbl_func = s3_tbl_func_maker(s3_settings)(
            for_="dba",
            conn_dto=saved_connection.get_conn_dto(),
            filename=sample_s3_file,
            file_fmt="Native",
            schema_line=schema_line,
        )

        result = await async_connection_executor.execute(
            ConnExecutorQuery(
                sa.select(columns=[sa.literal(1), sa.literal(2), sa.literal(3)])
                .select_from(sa.text(s3_tbl_func))
                .limit(1),
                user_types=[
                    UserDataType.boolean,
                    UserDataType.float,
                    UserDataType.integer,
                ],
            )
        )
        rows = await result.get_all()
        assert rows == [(True, 2.0, 3)], rows

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/base/core/connection_executor.py
datalens-tech/datalens-backend
import abc
from typing import TypeVar

import pytest

from dl_constants.enums import UserDataType
from dl_core.db import SchemaColumn
from dl_core.db.native_type import GenericNativeType
from dl_core_testing.fixtures.sample_tables import TABLE_SPEC_SAMPLE_SUPERSTORE
from dl_core_testing.testcases.data_source import DefaultDataSourceTestClass

from dl_connector_bundle_chs3.chs3_base.core.data_source import BaseFileS3DataSource
from dl_connector_bundle_chs3.chs3_base.core.data_source_spec import BaseFileS3DataSourceSpec
from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3.file.core.constants import SOURCE_TYPE_FILE_S3_TABLE
from dl_connector_bundle_chs3_tests.db.base.core.base import (
    FILE_CONN_TV,
    BaseCHS3TestClass,
)


FILE_DSRC_SPEC_TV = TypeVar("FILE_DSRC_SPEC_TV", bound=BaseFileS3DataSourceSpec)
FILE_DSRC_TV = TypeVar("FILE_DSRC_TV", bound=BaseFileS3DataSource)


class CHS3TableDataSourceTestBase(
    BaseCHS3TestClass,
    DefaultDataSourceTestClass[FILE_CONN_TV, FILE_DSRC_SPEC_TV, FILE_DSRC_TV],
    metaclass=abc.ABCMeta,
):
    @pytest.fixture(scope="function")
    def initial_data_source_spec(
        self,
        sample_file_data_source: BaseFileS3Connection.FileDataSource,
    ) -> FILE_DSRC_SPEC_TV:
        dsrc_spec = BaseFileS3DataSourceSpec(
            source_type=SOURCE_TYPE_FILE_S3_TABLE,
            raw_schema=sample_file_data_source.raw_schema,
            s3_endpoint=self.connection_settings.S3_ENDPOINT,
            bucket=self.connection_settings.BUCKET,
            origin_source_id=sample_file_data_source.id,
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return list(TABLE_SPEC_SAMPLE_SUPERSTORE.table_schema)

    def test_build_from_clause(
        self,
        data_source: FILE_DSRC_TV,
        sample_file_data_source: BaseFileS3Connection.FileDataSource,
        saved_connection: FILE_CONN_TV,
    ) -> None:
        data_source._spec.raw_schema = [  # leaving one column to simplify the test
            SchemaColumn(
                name="c1",
                native_type=GenericNativeType(conn_type=self.conn_type, name="Int64"),
                user_type=UserDataType.integer,
            ),
        ]
        query_from = data_source.get_sql_source().compile(compile_kwargs={"literal_binds": True}).string

        replace_secret = saved_connection.get_conn_dto().replace_secret

        expected = (
            f"s3("
            f"'{self.connection_settings.S3_ENDPOINT}/{self.connection_settings.BUCKET}/"
            f"{sample_file_data_source.s3_filename}', "
            f"'key_id_{replace_secret}', 'secret_key_{replace_secret}', 'Native', "
            f"'c1 Nullable(Int64)')"
        )
        assert query_from == expected, query_from

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/base/core/data_source.py
datalens-tech/datalens-backend
import abc
from typing import Generic

import pytest

from dl_core.services_registry import ServicesRegistry
from dl_core.us_dataset import Dataset
from dl_core_testing.database import DbTable
from dl_core_testing.dataset_wrappers import DatasetTestWrapper
from dl_core_testing.testcases.dataset import DefaultDatasetTestSuite

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.base import (
    FILE_CONN_TV,
    BaseCHS3TestClass,
)


class CHS3DatasetTestBase(BaseCHS3TestClass, DefaultDatasetTestSuite, Generic[FILE_CONN_TV], abc.ABC):
    @pytest.fixture(scope="function")
    def dsrc_params(self, dataset_table: DbTable, sample_file_data_source: BaseFileS3Connection.FileDataSource) -> dict:
        return dict(
            origin_source_id=sample_file_data_source.id,
        )

    def test_get_param_hash(
        self,
        sample_table: DbTable,
        saved_connection: FileS3Connection,
        saved_dataset: Dataset,
        conn_default_service_registry: ServicesRegistry,
        dataset_wrapper: DatasetTestWrapper,
        sample_file_data_source: BaseFileS3Connection.FileDataSource,
    ) -> None:
        dataset = saved_dataset
        service_registry = conn_default_service_registry
        source_id = dataset.get_single_data_source_id()
        dsrc_coll = dataset_wrapper.get_data_source_coll_strict(source_id=source_id)
        hash_from_dataset = dsrc_coll.get_param_hash()

        templates = saved_connection.get_data_source_templates(
            conn_executor_factory=service_registry.get_conn_executor_factory().get_sync_conn_executor,
        )
        found_template = False
        for template in templates:
            if template.parameters["origin_source_id"] == sample_file_data_source.id:
                found_template = True
                hash_from_template = template.get_param_hash()
                assert hash_from_dataset == hash_from_template

        assert found_template

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/base/core/dataset.py
datalens-tech/datalens-backend
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport

from dl_connector_bundle_chs3.chs3_base.core.settings import FileS3ConnectorSettings


CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=52611).host,
    port_us_http=get_test_container_hostport("us", fallback_port=52611).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=52610).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=52610).port,
    us_master_token="AC1ofiek8coB",
    core_connector_ep_names=["clickhouse", "file", "gsheets_v2"],
)

SR_CONNECTION_SETTINGS = FileS3ConnectorSettings(
    SECURE=False,
    HOST=get_test_container_hostport("db-clickhouse", original_port=8123).host,
    PORT=get_test_container_hostport("db-clickhouse", original_port=8123).port,
    USERNAME="datalens",
    PASSWORD="qwerty",
    ACCESS_KEY_ID="accessKey1",
    SECRET_ACCESS_KEY="verySecretKey1",
    BUCKET="dl-file-uploader",
    S3_ENDPOINT="http://s3-storage:8000",  # compose svc name, because this is a container interaction (ch <-> s3)
)

DB_CH_URL = (
    f"clickhouse://datalens:qwerty@"
    f"{get_test_container_hostport('db-clickhouse', fallback_port=52604).as_pair()}/test_data"
)

S3_ENDPOINT_URL = f"http://{get_test_container_hostport('s3-storage', fallback_port=52620).as_pair()}"

API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["clickhouse", "file", "gsheets_v2"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
    redis_host=get_test_container_hostport("redis", fallback_port=52604).host,
    redis_port=get_test_container_hostport("redis", fallback_port=52604).port,
    redis_password="AwockEuvavDyinmeakmiRiopanbesBepsensUrdIz5",
)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/config.py
datalens-tech/datalens-backend
from dl_api_lib_testing.initialization import initialize_api_lib_test

from dl_connector_bundle_chs3_tests.db.config import API_TEST_CONFIG


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/conftest.py
datalens-tech/datalens-backend
import pytest

from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_bundle_chs3_tests.db.base.api.base import CHS3ConnectionApiTestBase
from dl_connector_bundle_chs3_tests.db.base.api.data import CHS3DataApiTestBase
from dl_connector_bundle_chs3_tests.db.base.api.dataset import CHS3DatasetTestBase
from dl_connector_bundle_chs3_tests.db.file.core.base import BaseFileS3TestClass


class FileS3ApiConnectionTestBase(
    BaseFileS3TestClass,
    CHS3ConnectionApiTestBase[FileS3Connection],
):
    @pytest.fixture(scope="function")
    def connection_params(
        self,
        sample_file_data_source: FileS3Connection.FileDataSource,
    ) -> dict:
        return dict(
            sources=[
                dict(
                    file_id=sample_file_data_source.file_id,
                    id=sample_file_data_source.id,
                    title=sample_file_data_source.title,
                    column_types=sample_file_data_source.column_types,
                ),
            ],
        )


class FileS3DatasetTestBase(FileS3ApiConnectionTestBase, CHS3DatasetTestBase[FileS3Connection]):
    pass


class FileS3DataApiTestBase(FileS3DatasetTestBase, CHS3DataApiTestBase[FileS3Connection]):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/file/api/base.py
datalens-tech/datalens-backend
import uuid

import pytest

from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_bundle_chs3_tests.db.base.api.connection import CHS3ConnectionTestSuite
from dl_connector_bundle_chs3_tests.db.file.api.base import FileS3ApiConnectionTestBase


class TestFileS3Connection(FileS3ApiConnectionTestBase, CHS3ConnectionTestSuite[FileS3Connection]):
    @pytest.fixture(scope="function")
    def single_new_conn_source_params(self) -> dict:
        return {
            "id": str(uuid.uuid4()),
            "file_id": str(uuid.uuid4()),
            "title": f"New File {str(uuid.uuid4())}",
            "column_types": [
                {"name": "new_field", "user_type": "string"},
            ],
        }

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/file/api/test_connection.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.data_api_suites import (
    DefaultConnectorDataDistinctTestSuite,
    DefaultConnectorDataGroupByFormulaTestSuite,
    DefaultConnectorDataPreviewTestSuite,
    DefaultConnectorDataRangeTestSuite,
)
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_bundle_chs3_tests.db.base.api.data import CHS3DataResultTestSuite
from dl_connector_bundle_chs3_tests.db.file.api.base import FileS3DataApiTestBase


class TestFileS3DataResult(FileS3DataApiTestBase, CHS3DataResultTestSuite):
    test_params = RegulatedTestParams(
        mark_features_skipped={
            CHS3DataResultTestSuite.array_support: "File connector doesn't support arrays",
        }
    )


class TestFileS3DataGroupBy(FileS3DataApiTestBase, DefaultConnectorDataGroupByFormulaTestSuite):
    pass


class TestFileS3DataRange(FileS3DataApiTestBase, DefaultConnectorDataRangeTestSuite):
    pass


class TestFileDataDistinct(FileS3DataApiTestBase, DefaultConnectorDataDistinctTestSuite):
    test_params = RegulatedTestParams(
        mark_tests_failed={
            DefaultConnectorDataDistinctTestSuite.test_date_filter_distinct: "FIXME",
        }
    )


class TestFileS3DataPreview(FileS3DataApiTestBase, DefaultConnectorDataPreviewTestSuite):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/file/api/test_data.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3_tests.db.base.api.dataset import CHS3DatasetTestSuite
from dl_connector_bundle_chs3_tests.db.file.api.base import FileS3DatasetTestBase


class TestFileS3Dataset(FileS3DatasetTestBase, CHS3DatasetTestSuite):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/file/api/test_dataset.py
datalens-tech/datalens-backend
import uuid

import pytest

from dl_constants.enums import FileProcessingStatus
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.fixtures.primitives import FixtureTableSpec

from dl_connector_bundle_chs3.file.core.constants import (
    CONNECTION_TYPE_FILE,
    SOURCE_TYPE_FILE_S3_TABLE,
)
from dl_connector_bundle_chs3.file.core.testing.connection import make_saved_file_connection
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.base import BaseCHS3TestClass


class BaseFileS3TestClass(BaseCHS3TestClass[FileS3Connection]):
    conn_type = CONNECTION_TYPE_FILE
    source_type = SOURCE_TYPE_FILE_S3_TABLE

    @pytest.fixture(scope="function")
    def sample_file_data_source(
        self,
        sample_table_spec: FixtureTableSpec,
        sample_s3_file: str,
    ) -> FileS3Connection.FileDataSource:
        raw_schema = self._get_raw_schema_for_ch_table(sample_table_spec)
        return FileS3Connection.FileDataSource(
            id=str(uuid.uuid4()),
            file_id=str(uuid.uuid4()),
            title=sample_s3_file,
            s3_filename=sample_s3_file,
            raw_schema=raw_schema,
            status=FileProcessingStatus.ready,
            column_types=[{"name": col[0], "user_type": col[1].name} for col in sample_table_spec.table_schema],
        )

    @pytest.fixture(scope="function")
    def saved_connection(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params: dict,
    ) -> FileS3Connection:
        conn = make_saved_file_connection(
            sync_usm=sync_us_manager,
            **connection_creation_params,
        )
        return conn

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/file/core/base.py
datalens-tech/datalens-backend
from dl_core_testing.testcases.adapter import BaseAsyncAdapterTestClass

from dl_connector_bundle_chs3.chs3_base.core.target_dto import BaseFileS3ConnTargetDTO
from dl_connector_bundle_chs3.file.core.adapter import AsyncFileS3Adapter
from dl_connector_bundle_chs3_tests.db.file.core.base import BaseFileS3TestClass


class TestAsyncFileS3Adapter(
    BaseFileS3TestClass,
    BaseAsyncAdapterTestClass[BaseFileS3ConnTargetDTO],
):
    ASYNC_ADAPTER_CLS = AsyncFileS3Adapter

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/file/core/test_adapter.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.connection import CHS3ConnectionTestBase
from dl_connector_bundle_chs3_tests.db.file.core.base import BaseFileS3TestClass


class TestFileS3Connection(BaseFileS3TestClass, CHS3ConnectionTestBase[FileS3Connection]):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/file/core/test_connection.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.connection_executor import (
    CHS3AsyncConnectionExecutorTestBase,
    CHS3SyncConnectionExecutorTestBase,
)
from dl_connector_bundle_chs3_tests.db.file.core.base import BaseFileS3TestClass


class TestFileS3SyncConnectionExecutor(BaseFileS3TestClass, CHS3SyncConnectionExecutorTestBase[FileS3Connection]):
    pass


class TestFileS3AsyncConnectionExecutor(BaseFileS3TestClass, CHS3AsyncConnectionExecutorTestBase[FileS3Connection]):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/file/core/test_connection_executor.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.file.core.data_source import FileS3DataSource
from dl_connector_bundle_chs3.file.core.data_source_spec import FileS3DataSourceSpec
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.data_source import CHS3TableDataSourceTestBase
from dl_connector_bundle_chs3_tests.db.file.core.base import BaseFileS3TestClass


class TestFileS3TableDataSource(
    BaseFileS3TestClass,
    CHS3TableDataSourceTestBase[FileS3Connection, FileS3DataSourceSpec, FileS3DataSource],
):
    DSRC_CLS = FileS3DataSource

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/file/core/test_data_source.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.dataset import CHS3DatasetTestBase
from dl_connector_bundle_chs3_tests.db.file.core.base import BaseFileS3TestClass


class TestFileS3Dataset(BaseFileS3TestClass, CHS3DatasetTestBase[FileS3Connection]):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/file/core/test_dataset.py
datalens-tech/datalens-backend
import datetime
import logging

from aiohttp import web
import pytest

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core.us_manager.us_manager_async import AsyncUSManager

from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection
from dl_connector_bundle_chs3_tests.db.base.api.base import CHS3ConnectionApiTestBase
from dl_connector_bundle_chs3_tests.db.base.api.data import CHS3DataApiTestBase
from dl_connector_bundle_chs3_tests.db.base.api.dataset import CHS3DatasetTestBase
from dl_connector_bundle_chs3_tests.db.gsheets_v2.core.base import BaseGSheetsFileS3TestClass


LOGGER = logging.getLogger(__name__)


class GSheetsFileS3ApiConnectionTestBase(
    BaseGSheetsFileS3TestClass,
    CHS3ConnectionApiTestBase[GSheetsFileS3Connection],
):
    @pytest.fixture(scope="function")
    async def mock_file_uploader_api(
        self,
        aiohttp_server,
        bi_test_config: ApiTestEnvironmentConfiguration,
        async_us_manager: AsyncUSManager,
    ) -> None:
        async def mocked_update_connection_data_internal(request: web.Request) -> web.Response:
            req_data = await request.json()
            conn_id: str = req_data["connection_id"]
            sources_to_update = [src["id"] for src in req_data["sources"]]

            conn = await async_us_manager.get_by_id(conn_id, GSheetsFileS3Connection)
            for src_id in sources_to_update:
                src: GSheetsFileS3Connection.FileDataSource = conn.get_file_source_by_id(src_id)
                src.data_updated_at = datetime.datetime.now(datetime.timezone.utc)
                LOGGER.info(f"Successfully updated source id {src_id}")
            await async_us_manager.save(conn)

            return web.HTTPOk()

        app = web.Application()
        app.router.add_route("POST", "/api/v2/update_connection_data_internal", mocked_update_connection_data_internal)

        server = await aiohttp_server(app, port=bi_test_config.file_uploader_api_port)

        yield

        await server.close()

    @pytest.fixture(scope="function")
    def connection_params(
        self,
        sample_file_data_source: GSheetsFileS3Connection.FileDataSource,
    ) -> dict:
        return dict(
            refresh_enabled=True,
            sources=[
                dict(
                    file_id=sample_file_data_source.file_id,
                    id=sample_file_data_source.id,
                    title=sample_file_data_source.title,
                ),
            ],
        )


class GSheetsFileS3DatasetTestBase(GSheetsFileS3ApiConnectionTestBase, CHS3DatasetTestBase[GSheetsFileS3Connection]):
    pass


class GSheetsFileS3DataApiTestBase(GSheetsFileS3DatasetTestBase, CHS3DataApiTestBase[GSheetsFileS3Connection]):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/gsheets_v2/api/base.py
datalens-tech/datalens-backend
from http import HTTPStatus
import logging
import uuid

from flask.testing import FlaskClient
import pytest

from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_testing.utils import get_log_record

from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection
from dl_connector_bundle_chs3_tests.db.base.api.connection import CHS3ConnectionTestSuite
from dl_connector_bundle_chs3_tests.db.gsheets_v2.api.base import GSheetsFileS3ApiConnectionTestBase


class TestGSheetsFileS3Connection(GSheetsFileS3ApiConnectionTestBase, CHS3ConnectionTestSuite[GSheetsFileS3Connection]):
    @pytest.fixture(scope="function")
    def single_new_conn_source_params(self) -> dict:
        return {
            "id": str(uuid.uuid4()),
            "file_id": str(uuid.uuid4()),
            "title": f"New File {str(uuid.uuid4())}",
        }

    def test_authorization_field(
        self,
        client: FlaskClient,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
    ) -> None:
        conn_id = saved_connection_id
        conn = sync_us_manager.get_by_id(conn_id, GSheetsFileS3Connection)
        base_update_data = {
            "refresh_enabled": True,
            "sources": [{"id": src.id, "title": src.title} for src in conn.data.sources],
        }

        # no token => not authorized
        assert conn.authorized is False
        conn_resp = client.get(f"/api/v1/connections/{conn_id}")
        assert conn_resp.status_code == HTTPStatus.OK, conn_resp.json
        assert conn_resp.json["authorized"] is False
        assert "refresh_token" not in conn_resp.json

        # add token into connection
        resp = client.put(
            "/api/v1/connections/{}".format(conn_id),
            json={
                **base_update_data,
                "refresh_token": "some_token",
            },
        )
        assert resp.status_code == HTTPStatus.OK, resp.json
        conn = sync_us_manager.get_by_id(conn_id, GSheetsFileS3Connection)
        assert conn.authorized is True
        conn_resp = client.get(f"/api/v1/connections/{conn_id}")
        assert conn_resp.status_code == HTTPStatus.OK, conn_resp.json
        assert conn_resp.json["authorized"] is True
        assert "refresh_token" not in conn_resp.json

        # remove token from the connection
        resp = client.put(
            "/api/v1/connections/{}".format(conn_id),
            json={
                **base_update_data,
                "refresh_token": None,
            },
        )
        assert resp.status_code == HTTPStatus.OK, resp.json
        conn: GSheetsFileS3Connection = sync_us_manager.get_by_id(conn_id, GSheetsFileS3Connection)
        assert conn.authorized is False
        conn_resp = client.get(f"/api/v1/connections/{conn_id}")
        assert conn_resp.status_code == HTTPStatus.OK
        assert conn_resp.json["authorized"] is False
        assert "refresh_token" not in conn_resp.json

    def test_force_update_with_file_id(
        self,
        caplog,
        client: FlaskClient,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
    ) -> None:
        """Passed file_id to an existing source means that it has been updated => this should trigger data update"""

        caplog.set_level(logging.INFO)

        conn_id = saved_connection_id
        usm = sync_us_manager
        conn = usm.get_by_id(conn_id, GSheetsFileS3Connection)

        resp = client.put(
            "/api/v1/connections/{}".format(conn_id),
            json={
                "sources": [
                    {
                        "file_id": str(uuid.uuid4()),  # force source update by passing file_id
                        "id": conn.data.sources[0].id,
                        "title": conn.data.sources[0].title,
                    },
                ],
            },
        )
        assert resp.status_code == HTTPStatus.OK, resp.json

        schedule_save_src_log_record = get_log_record(
            caplog,
            predicate=lambda r: r.message.startswith("Scheduled task SaveSourceTask for source_id"),
            single=True,
        )
        assert conn.data.sources[0].id in schedule_save_src_log_record.message

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/gsheets_v2/api/test_connection.py
datalens-tech/datalens-backend
import datetime
from http import HTTPStatus

from flask.testing import FlaskClient
import pytest

from dl_api_client.dsmaker.api.data_api import SyncHttpDataApiV2
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib_testing.connector.data_api_suites import (
    DefaultConnectorDataDistinctTestSuite,
    DefaultConnectorDataGroupByFormulaTestSuite,
    DefaultConnectorDataPreviewTestSuite,
    DefaultConnectorDataRangeTestSuite,
    DefaultConnectorDataResultTestSuite,
)
from dl_api_lib_testing.data_api_base import DataApiTestParams
from dl_constants.enums import (
    ComponentErrorLevel,
    ComponentType,
    DataSourceRole,
    FileProcessingStatus,
)
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_bundle_chs3.chs3_gsheets.core.constants import NOTIF_TYPE_GSHEETS_V2_STALE_DATA
from dl_connector_bundle_chs3.chs3_gsheets.core.lifecycle import GSheetsFileS3ConnectionLifecycleManager
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection
from dl_connector_bundle_chs3_tests.db.base.api.data import CHS3DataResultTestSuite
from dl_connector_bundle_chs3_tests.db.gsheets_v2.api.base import GSheetsFileS3DataApiTestBase


class TestGSheetsFileS3DataResult(GSheetsFileS3DataApiTestBase, CHS3DataResultTestSuite):
    test_params = RegulatedTestParams(
        mark_features_skipped={
            DefaultConnectorDataResultTestSuite.array_support: "GSheets V2 connector doesn't support arrays",
        },
    )

    @pytest.mark.asyncio
    def test_update_data(
        self,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
        data_api: SyncHttpDataApiV2,
        saved_dataset: Dataset,
        data_api_test_params: DataApiTestParams,
        mock_file_uploader_api,
    ) -> None:
        ds = saved_dataset

        # prepare connection sources: set updated time to the current moment
        conn = sync_us_manager.get_by_id(saved_connection_id, GSheetsFileS3Connection)
        dt_now = datetime.datetime.now(datetime.timezone.utc)
        data_updated_at_orig = dt_now
        for src in conn.data.sources:
            src.data_updated_at = dt_now
        sync_us_manager.save(conn)

        def get_notifications_from_result_resp() -> list[dict]:
            result_resp = data_api.get_result(
                dataset=ds, fields=[ds.find_field(title=data_api_test_params.date_field)], fail_ok=True
            )
            return result_resp.json.get("notifications", [])

        # it is not time to update data yet, so we expect no data updates or corresponding notifications
        notifications = get_notifications_from_result_resp()
        assert all(
            notification["locator"] != NOTIF_TYPE_GSHEETS_V2_STALE_DATA.value for notification in notifications
        ), notifications
        conn = sync_us_manager.get_by_id(saved_connection_id, GSheetsFileS3Connection)
        assert conn.data.sources[0].data_updated_at == data_updated_at_orig

        # trigger data update by setting the data update time in the connection to N minutes ago
        data_updated_at = conn.data.oldest_data_update_time() - datetime.timedelta(
            seconds=GSheetsFileS3ConnectionLifecycleManager.STALE_THRESHOLD_SECONDS + 60,  # just in case
        )
        for src in conn.data.sources:
            src.data_updated_at = data_updated_at
        sync_us_manager.save(conn)

        # now notifications should be there, as well as connection sources should be updated
        notifications = get_notifications_from_result_resp()
        assert any(
            notification["locator"] == NOTIF_TYPE_GSHEETS_V2_STALE_DATA.value for notification in notifications
        ), notifications
        conn = sync_us_manager.get_by_id(saved_connection_id, GSheetsFileS3Connection)
        assert conn.data.sources[0].data_updated_at != data_updated_at

    def test_component_error(
        self,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
        saved_dataset: Dataset,
        client: FlaskClient,
        data_api: SyncHttpDataApiV2,
        data_api_test_params: DataApiTestParams,
    ) -> None:
        conn = sync_us_manager.get_by_id(saved_connection_id, GSheetsFileS3Connection)
        err_details = {"error": "details", "request-id": "637"}
        conn.data.component_errors.add_error(
            id=conn.data.sources[0].id,
            type=ComponentType.data_source,
            message="Custom error message",
            code=["FILE", "CUSTOM_FILE_ERROR"],
            details=err_details,
        )
        conn.update_data_source(
            conn.data.sources[0].id,
            role=DataSourceRole.origin,
            s3_filename=None,
            status=FileProcessingStatus.failed,
            preview_id=None,
            data_updated_at=datetime.datetime.now(datetime.timezone.utc),
        )
        sync_us_manager.save(conn)

        ds = saved_dataset
        result_resp = data_api.get_result(
            dataset=ds, fields=[ds.find_field(title=data_api_test_params.distinct_field)], fail_ok=True
        )
        assert result_resp.status_code == HTTPStatus.BAD_REQUEST, result_resp.json
        assert result_resp.json["details"] == err_details
        assert result_resp.json["message"] == "Custom error message"
        assert result_resp.json["code"] == "ERR.DS_API.SOURCE.FILE.CUSTOM_FILE_ERROR"

        conn_resp = client.get(f"/api/v1/connections/{saved_connection_id}")
        assert conn_resp.status_code == HTTPStatus.OK, conn_resp.json
        assert conn_resp.json["component_errors"], conn_resp.json
        actual_errors = conn_resp.json["component_errors"]["items"][0]["errors"]
        assert len(actual_errors) == 1, actual_errors
        assert actual_errors[0]["code"] == "ERR.DS_API.SOURCE.FILE.CUSTOM_FILE_ERROR"

    @pytest.mark.asyncio
    def test_component_error_warning(
        self,
        sync_us_manager: SyncUSManager,
        saved_connection_id: str,
        saved_dataset: Dataset,
        data_api: SyncHttpDataApiV2,
        data_api_test_params: DataApiTestParams,
        mock_file_uploader_api,
    ) -> None:
        conn = sync_us_manager.get_by_id(saved_connection_id, GSheetsFileS3Connection)

        long_long_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(
            seconds=GSheetsFileS3ConnectionLifecycleManager.STALE_THRESHOLD_SECONDS + 60,  # just in case
        )
        err_details = {"error": "details", "request-id": "637"}
        conn.data.component_errors.add_error(
            id=conn.data.sources[0].id,
            type=ComponentType.data_source,
            message="Custom error message",
            code=["FILE", "CUSTOM_FILE_ERROR"],
            details=err_details,
            level=ComponentErrorLevel.warning,
        )
        conn.update_data_source(
            conn.data.sources[0].id,
            role=DataSourceRole.origin,
            data_updated_at=long_long_ago,
        )
        sync_us_manager.save(conn)

        ds = saved_dataset
        result_resp = data_api.get_result(
            dataset=ds, fields=[ds.find_field(title=data_api_test_params.distinct_field)], fail_ok=True
        )
        assert result_resp.status_code == HTTPStatus.OK, result_resp.json
        assert len(result_resp.json["notifications"]) == 2
        assert "Reason: FILE.CUSTOM_FILE_ERROR, Request-ID: 637" in result_resp.json["notifications"][0]["message"]
        conn = sync_us_manager.get_by_id(saved_connection_id, GSheetsFileS3Connection)
        assert conn.data.sources[0].data_updated_at > long_long_ago  # data update was triggered


class TestGSheetsFileS3DataGroupBy(GSheetsFileS3DataApiTestBase, DefaultConnectorDataGroupByFormulaTestSuite):
    pass


class TestGSheetsFileS3DataRange(GSheetsFileS3DataApiTestBase, DefaultConnectorDataRangeTestSuite):
    pass


class TestGSheetsFileDataDistinct(GSheetsFileS3DataApiTestBase, DefaultConnectorDataDistinctTestSuite):
    test_params = RegulatedTestParams(
        mark_tests_failed={
            DefaultConnectorDataDistinctTestSuite.test_date_filter_distinct: "FIXME",
        }
    )


class TestGSheetsFileS3DataPreview(GSheetsFileS3DataApiTestBase, DefaultConnectorDataPreviewTestSuite):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/gsheets_v2/api/test_data.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3_tests.db.base.api.dataset import CHS3DatasetTestSuite
from dl_connector_bundle_chs3_tests.db.gsheets_v2.api.base import GSheetsFileS3DatasetTestBase


class TestGSheetsFileS3Dataset(GSheetsFileS3DatasetTestBase, CHS3DatasetTestSuite):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/gsheets_v2/api/test_dataset.py
datalens-tech/datalens-backend
import uuid

import pytest

from dl_constants.enums import FileProcessingStatus
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.fixtures.primitives import FixtureTableSpec

from dl_connector_bundle_chs3.chs3_gsheets.core.constants import (
    CONNECTION_TYPE_GSHEETS_V2,
    SOURCE_TYPE_GSHEETS_V2,
)
from dl_connector_bundle_chs3.chs3_gsheets.core.testing.connection import make_saved_gsheets_v2_connection
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.base import BaseCHS3TestClass


class BaseGSheetsFileS3TestClass(BaseCHS3TestClass[GSheetsFileS3Connection]):
    conn_type = CONNECTION_TYPE_GSHEETS_V2
    source_type = SOURCE_TYPE_GSHEETS_V2

    @pytest.fixture(scope="function")
    def sample_file_data_source(
        self,
        sample_table_spec: FixtureTableSpec,
        sample_s3_file: str,
    ) -> GSheetsFileS3Connection.FileDataSource:
        raw_schema = self._get_raw_schema_for_ch_table(sample_table_spec)
        return GSheetsFileS3Connection.FileDataSource(
            id=str(uuid.uuid4()),
            file_id=str(uuid.uuid4()),
            title=sample_s3_file,
            s3_filename=sample_s3_file,
            raw_schema=raw_schema,
            status=FileProcessingStatus.ready,
        )

    @pytest.fixture(scope="function")
    def saved_connection(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params: dict,
    ) -> GSheetsFileS3Connection:
        conn = make_saved_gsheets_v2_connection(
            sync_usm=sync_us_manager,
            **connection_creation_params,
        )
        return conn

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/gsheets_v2/core/base.py
datalens-tech/datalens-backend
from dl_core_testing.testcases.adapter import BaseAsyncAdapterTestClass

from dl_connector_bundle_chs3.chs3_base.core.target_dto import BaseFileS3ConnTargetDTO
from dl_connector_bundle_chs3.chs3_gsheets.core.adapter import AsyncGSheetsFileS3Adapter
from dl_connector_bundle_chs3_tests.db.gsheets_v2.core.base import BaseGSheetsFileS3TestClass


class TestAsyncGSheetsFileS3Adapter(
    BaseGSheetsFileS3TestClass,
    BaseAsyncAdapterTestClass[BaseFileS3ConnTargetDTO],
):
    ASYNC_ADAPTER_CLS = AsyncGSheetsFileS3Adapter

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/gsheets_v2/core/test_adapter.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.connection import CHS3ConnectionTestBase
from dl_connector_bundle_chs3_tests.db.gsheets_v2.core.base import BaseGSheetsFileS3TestClass


class TestGSheetsFileS3Connection(BaseGSheetsFileS3TestClass, CHS3ConnectionTestBase[FileS3Connection]):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/gsheets_v2/core/test_connection.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.connection_executor import (
    CHS3AsyncConnectionExecutorTestBase,
    CHS3SyncConnectionExecutorTestBase,
)
from dl_connector_bundle_chs3_tests.db.gsheets_v2.core.base import BaseGSheetsFileS3TestClass


class TestGSheetsFileS3SyncConnectionExecutor(
    BaseGSheetsFileS3TestClass,
    CHS3SyncConnectionExecutorTestBase[FileS3Connection],
):
    pass


class TestGSheetsFileS3AsyncConnectionExecutor(
    BaseGSheetsFileS3TestClass,
    CHS3AsyncConnectionExecutorTestBase[FileS3Connection],
):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/gsheets_v2/core/test_connection_executor.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.chs3_gsheets.core.data_source import GSheetsFileS3DataSource
from dl_connector_bundle_chs3.chs3_gsheets.core.data_source_spec import GSheetsFileS3DataSourceSpec
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.data_source import CHS3TableDataSourceTestBase
from dl_connector_bundle_chs3_tests.db.gsheets_v2.core.base import BaseGSheetsFileS3TestClass


class TestGSheetsFileS3TableDataSource(
    BaseGSheetsFileS3TestClass,
    CHS3TableDataSourceTestBase[GSheetsFileS3Connection, GSheetsFileS3DataSourceSpec, GSheetsFileS3DataSource],
):
    DSRC_CLS = GSheetsFileS3DataSource

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/gsheets_v2/core/test_data_source.py
datalens-tech/datalens-backend
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection
from dl_connector_bundle_chs3_tests.db.base.core.dataset import CHS3DatasetTestBase
from dl_connector_bundle_chs3_tests.db.gsheets_v2.core.base import BaseGSheetsFileS3TestClass


class TestGSheetsFileS3Dataset(BaseGSheetsFileS3TestClass, CHS3DatasetTestBase[GSheetsFileS3Connection]):
    pass

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3_tests/db/gsheets_v2/core/test_dataset.py
datalens-tech/datalens-backend
from __future__ import annotations

import marshmallow as ma

from dl_api_connector.api_schema.connection_base import (
    ConnectionMetaMixin,
    ConnectionSchema,
)
from dl_api_connector.api_schema.connection_base_fields import (
    cache_ttl_field,
    secret_string_field,
)
from dl_api_connector.api_schema.connection_mixins import (
    DataExportForbiddenMixin,
    RawSQLLevelMixin,
)
from dl_api_connector.api_schema.connection_sql import DBHostField
from dl_api_connector.api_schema.extras import FieldExtra

from dl_connector_chyt.core.us_connection import ConnectionCHYTToken


class CHYTConnectionSchema(ConnectionMetaMixin, RawSQLLevelMixin, DataExportForbiddenMixin, ConnectionSchema):
    TARGET_CLS = ConnectionCHYTToken

    ALLOW_MULTIHOST = False

    host = DBHostField(attribute="data.host", required=True, bi_extra=FieldExtra(editable=True))
    port = ma.fields.Integer(attribute="data.port", required=True, bi_extra=FieldExtra(editable=True))
    token = secret_string_field(attribute="data.token", required=True, bi_extra=FieldExtra(editable=True))
    alias = ma.fields.String(attribute="data.alias", required=True, bi_extra=FieldExtra(editable=True))
    secure = ma.fields.Boolean(attribute="data.secure", bi_extra=FieldExtra(editable=True))
    cache_ttl_sec = cache_ttl_field(attribute="data.cache_ttl_sec")

lib/dl_connector_chyt/dl_connector_chyt/api/api_schema/connection.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.source_base import (
    DataSourceTemplateBaseSchema,
    SimpleDataSourceSchema,
    SimpleParametersSchema,
)


# Effectively `CHYTTableSourceSchema = SQLDataSourceSchema` because it is in `SQL_SOURCE_TYPES`.


class CHYTTableListParametersSchema(SimpleParametersSchema):
    table_names = ma_fields.String()  # newline-separated tables


class CHYTTableListDataSourceSchema(SimpleDataSourceSchema):
    parameters = ma_fields.Nested(CHYTTableListParametersSchema)


class CHYTTableListDataSourceTemplateSchema(DataSourceTemplateBaseSchema):
    parameters = ma_fields.Nested(CHYTTableListParametersSchema)


class CHYTTableRangeParametersSchema(SimpleParametersSchema):
    directory_path = ma_fields.String()
    range_from = ma_fields.String()
    range_to = ma_fields.String()


class CHYTTableRangeDataSourceSchema(SimpleDataSourceSchema):
    parameters = ma_fields.Nested(CHYTTableRangeParametersSchema)


class CHYTTableRangeDataSourceTemplateSchema(DataSourceTemplateBaseSchema):
    parameters = ma_fields.Nested(CHYTTableRangeParametersSchema)

lib/dl_connector_chyt/dl_connector_chyt/api/api_schema/source.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import unique
from typing import Optional

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormActionApiSchema,
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
)
from dl_api_connector.form_config.models.common import (
    CommonFieldName,
    FormFieldName,
)
import dl_api_connector.form_config.models.rows as C
from dl_api_connector.form_config.models.shortcuts.rows import RowConstructor
from dl_configs.connectors_settings import ConnectorSettingsBase

from dl_connector_chyt.api.connection_info import CHYTConnectionInfoProvider
from dl_connector_chyt.api.i18n.localizer import Translatable
from dl_connector_chyt.core.settings import CHYTConnectorSettings


@unique
class CHYTFieldName(FormFieldName):
    alias = "alias"


class CHYTConnectionFormFactory(ConnectionFormFactory):
    def get_form_config(
        self, connector_settings: Optional[ConnectorSettingsBase], tenant: Optional[TenantDef]
    ) -> ConnectionForm:
        assert connector_settings is not None and isinstance(connector_settings, CHYTConnectorSettings)

        rc = RowConstructor(localizer=self._localizer)

        clique_alias_row = C.CustomizableRow(
            items=[
                C.LabelRowItem(text=self._localizer.translate(Translatable("field_clique-alias"))),
                C.InputRowItem(name=CHYTFieldName.alias, width="l", default_value=connector_settings.DEFAULT_CLIQUE),
            ]
        )

        token_row = C.CustomizableRow(
            items=[
                C.LabelRowItem(text=self._localizer.translate(Translatable("field_chyt-token"))),
                C.InputRowItem(
                    name=CommonFieldName.token,
                    width="m",
                    default_value="" if self.mode == ConnectionFormMode.create else None,
                    fake_value="******" if self.mode == ConnectionFormMode.edit else None,
                    control_props=C.InputRowItem.Props(type="password"),
                ),
            ]
        )

        secure_row = C.CustomizableRow(
            items=[
                C.CheckboxRowItem(name=CommonFieldName.secure, text="HTTPS", default_value=False),
            ]
        )

        common_api_schema_items: list[FormFieldApiSchema] = [
            FormFieldApiSchema(name=CommonFieldName.host, required=True),
            FormFieldApiSchema(name=CommonFieldName.port, required=True),
            FormFieldApiSchema(name=CHYTFieldName.alias, required=True),
            FormFieldApiSchema(name=CommonFieldName.token, required=self.mode == ConnectionFormMode.create),
            FormFieldApiSchema(name=CommonFieldName.secure, type="boolean"),
            FormFieldApiSchema(name=CommonFieldName.data_export_forbidden),
        ]

        edit_api_schema = FormActionApiSchema(
            items=[
                *common_api_schema_items,
                FormFieldApiSchema(name=CommonFieldName.cache_ttl_sec, nullable=True),
                FormFieldApiSchema(name=CommonFieldName.raw_sql_level),
            ]
        )

        create_api_schema = FormActionApiSchema(
            items=[
                *edit_api_schema.items,
                *self._get_top_level_create_api_schema_items(),
            ]
        )

        check_api_schema = FormActionApiSchema(
            items=[
                *common_api_schema_items,
                *self._get_top_level_check_api_schema_items(),
            ]
        )

        return ConnectionForm(
            title=CHYTConnectionInfoProvider.get_title(self._localizer),
            rows=[
                rc.host_row(),
                rc.port_row(),
                clique_alias_row,
                token_row,
                C.CacheTTLRow(name=CommonFieldName.cache_ttl_sec),
                rc.raw_sql_level_row(),
                secure_row,
                rc.collapse_advanced_settings_row(),
                rc.data_export_forbidden_row(),
            ],
            api_schema=FormApiSchema(
                create=create_api_schema if self.mode == ConnectionFormMode.create else None,
                edit=edit_api_schema if self.mode == ConnectionFormMode.edit else None,
                check=check_api_schema,
            ),
        )

lib/dl_connector_chyt/dl_connector_chyt/api/connection_form/form_config.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_chyt.api.i18n.localizer import Translatable


class CHYTConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-chyt")

lib/dl_connector_chyt/dl_connector_chyt/api/connection_info.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_connector.api_schema.source_base import (
    SQLDataSourceSchema,
    SQLDataSourceTemplateSchema,
    SubselectDataSourceSchema,
    SubselectDataSourceTemplateSchema,
)
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)

from dl_connector_chyt.api.api_schema.connection import CHYTConnectionSchema
from dl_connector_chyt.api.api_schema.source import (
    CHYTTableListDataSourceSchema,
    CHYTTableListDataSourceTemplateSchema,
    CHYTTableRangeDataSourceSchema,
    CHYTTableRangeDataSourceTemplateSchema,
)
from dl_connector_chyt.api.connection_form.form_config import CHYTConnectionFormFactory
from dl_connector_chyt.api.connection_info import CHYTConnectionInfoProvider
from dl_connector_chyt.api.i18n.localizer import CONFIGS
from dl_connector_chyt.core.connector import (
    CHYTCoreConnectionDefinition,
    CHYTCoreConnector,
    CHYTTableCoreSourceDefinition,
    CHYTTableListCoreSourceDefinition,
    CHYTTableRangeCoreSourceDefinition,
    CHYTTableSubselectCoreSourceDefinition,
)
from dl_connector_clickhouse.formula.constants import DIALECT_NAME_CLICKHOUSE


class CHYTApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = CHYTCoreConnectionDefinition
    api_generic_schema_cls = CHYTConnectionSchema
    info_provider_cls = CHYTConnectionInfoProvider
    form_factory_cls = CHYTConnectionFormFactory


class CHYTTableApiSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = CHYTTableCoreSourceDefinition
    api_schema_cls = SQLDataSourceSchema
    template_api_schema_cls = SQLDataSourceTemplateSchema


class CHYTTableListApiSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = CHYTTableListCoreSourceDefinition
    api_schema_cls = CHYTTableListDataSourceSchema
    template_api_schema_cls = CHYTTableListDataSourceTemplateSchema


class CHYTTableRangeApiSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = CHYTTableRangeCoreSourceDefinition
    api_schema_cls = CHYTTableRangeDataSourceSchema
    template_api_schema_cls = CHYTTableRangeDataSourceTemplateSchema


class CHYTSubselectApiSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = CHYTTableSubselectCoreSourceDefinition
    api_schema_cls = SubselectDataSourceSchema
    template_api_schema_cls = SubselectDataSourceTemplateSchema


class CHYTApiConnector(ApiConnector):
    core_connector_cls = CHYTCoreConnector
    formula_dialect_name = DIALECT_NAME_CLICKHOUSE
    connection_definitions = (CHYTApiConnectionDefinition,)
    source_definitions = (
        CHYTTableApiSourceDefinition,
        CHYTTableListApiSourceDefinition,
        CHYTTableRangeApiSourceDefinition,
        CHYTSubselectApiSourceDefinition,
    )
    translation_configs = frozenset(CONFIGS)

lib/dl_connector_chyt/dl_connector_chyt/api/connector.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_chyt as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_chyt/dl_connector_chyt/api/i18n/localizer.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    Any,
    ClassVar,
    Optional,
    Type,
)

from aiochclient.http_clients import aiohttp
import attr

from dl_constants.enums import IndexKind
from dl_core import exc
from dl_core.connection_executors.models.db_adapter_data import RawIndexInfo
from dl_core.connection_models import TableIdent
from dl_core.connectors.base.error_transformer import DBExcKWArgs
from dl_core.utils import (
    get_current_w3c_tracing_headers,
    make_url,
)
from dl_utils.aio import await_sync

from dl_connector_chyt.core.constants import CONNECTION_TYPE_CHYT
from dl_connector_chyt.core.target_dto import (
    BaseCHYTConnTargetDTO,
    CHYTConnTargetDTO,
)
from dl_connector_chyt.core.utils import CHYTUtils
from dl_connector_clickhouse.core.clickhouse_base.adapters import (
    BaseClickHouseAdapter,
    BaseClickHouseConnLineConstructor,
)
from dl_connector_clickhouse.core.clickhouse_base.ch_commons import (
    ensure_db_message,
    get_ch_settings,
)


class CHYTConnLineConstructor(BaseClickHouseConnLineConstructor):
    def _get_dsn_params(
        self,
        safe_db_symbols: tuple[str, ...] = (),
        db_name: Optional[str] = None,
        standard_auth: Optional[bool] = True,
    ) -> dict:
        new_safe_symbols = safe_db_symbols + ("*",)
        return super()._get_dsn_params(safe_db_symbols=new_safe_symbols, db_name=db_name, standard_auth=standard_auth)


class BaseCHYTAdapter(BaseClickHouseAdapter, abc.ABC):
    ch_utils = CHYTUtils
    conn_line_constructor_type: ClassVar[Type[CHYTConnLineConstructor]] = CHYTConnLineConstructor

    _target_dto: BaseCHYTConnTargetDTO = attr.ib()

    def _get_dsn_params_from_headers(self) -> dict[str, str]:
        params = {
            "statuses_to_retry": "500,502,504",
            **super()._get_dsn_params_from_headers(),
        }
        params.update(
            **self._convert_headers_to_dsn_params(
                get_current_w3c_tracing_headers(
                    override_sampled=self.ch_utils.get_tracing_sample_flag_override(self._req_ctx_info),
                    req_id=self._req_ctx_info.request_id,
                ),
            )
        )
        return params

    def get_ch_settings(self) -> dict:
        return get_ch_settings(
            max_execution_time=self._target_dto.max_execution_time,
            read_only_level=2,
            output_format_json_quote_denormals=1,
        )

    # Fat boilerplate of an override.
    @classmethod
    def make_exc(  # TODO:  Move to ErrorTransformer
        cls,
        wrapper_exc: Exception,
        orig_exc: Optional[Exception],
        debug_compiled_query: Optional[str],
        **kwargs: Any,
    ) -> tuple[Type[exc.DatabaseQueryError], DBExcKWArgs]:
        exc_cls, kw = super().make_exc(
            wrapper_exc=wrapper_exc,
            orig_exc=orig_exc,
            debug_compiled_query=debug_compiled_query,
            **kwargs,  # type: ignore
        )
        return ensure_db_message(exc_cls, kw)

    # noinspection PyMethodMayBeStatic
    async def _fetch_yt_sorting_columns(
        self,
        host: str,
        table_path: str,
        secret_auth_headers: dict[str, str],
    ) -> Optional[RawIndexInfo]:
        async with aiohttp.ClientSession() as session:
            async with session.get(
                make_url(
                    protocol="https",
                    host=host,
                    port=443,
                    path="/api/v3/get",
                ),
                params=dict(path=f"{table_path}/@schema"),
                headers={
                    **secret_auth_headers,
                    "accept": "application/json",
                },
            ) as resp:
                resp.raise_for_status()
                resp_js = await resp.json()

        sorting_columns = tuple(
            col_js["name"] for col_js in resp_js["$value"] if isinstance(col_js.get("sort_order"), str)
        )

        if sorting_columns:
            return RawIndexInfo(
                columns=sorting_columns,
                unique=resp_js["$attributes"]["unique_keys"],
                kind=IndexKind.table_sorting,
            )
        else:
            return None

    @abc.abstractmethod
    async def _get_yt_table_index_info(
        self,
        table_path: str,
        secret_auth_headers: dict[str, str],  # Strange name to prevent leaking to Sentry
    ) -> Optional[RawIndexInfo]:
        raise NotImplementedError()

    @abc.abstractmethod
    def _get_yt_proxy_auth_headers(self) -> dict[str, str]:
        raise NotImplementedError()

    def _get_table_indexes(self, table_ident: TableIdent) -> tuple[RawIndexInfo, ...]:
        """Even does not try to use SA to fetch indexes"""

        may_be_sorting_idx = await_sync(
            self._get_yt_table_index_info(
                table_path=table_ident.table_name, secret_auth_headers=self._get_yt_proxy_auth_headers()
            )
        )
        return (may_be_sorting_idx,) if may_be_sorting_idx is not None else ()


class CHYTAdapter(BaseCHYTAdapter):
    conn_type = CONNECTION_TYPE_CHYT

    _target_dto: CHYTConnTargetDTO = attr.ib()

    def _get_yt_proxy_auth_headers(self) -> dict[str, str]:
        return {"Authorization": f"OAuth {self._target_dto.password}"}

    async def _get_yt_table_index_info(
        self,
        table_path: str,
        secret_auth_headers: dict[str, str],
    ) -> Optional[RawIndexInfo]:
        assert isinstance(table_path, str) and table_path.startswith("//"), "Incorrect YT table path"
        return await self._fetch_yt_sorting_columns(self._target_dto.host, table_path, secret_auth_headers)

lib/dl_connector_chyt/dl_connector_chyt/core/adapters.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging

import attr
from clickhouse_sqlalchemy.drivers.http.transport import _get_type  # noqa

from dl_core.utils import get_current_w3c_tracing_headers

from dl_connector_chyt.core.constants import CONNECTION_TYPE_CHYT
from dl_connector_chyt.core.target_dto import BaseCHYTConnTargetDTO
from dl_connector_chyt.core.utils import CHYTUtils
from dl_connector_clickhouse.core.clickhouse_base.adapters import BaseAsyncClickHouseAdapter


LOGGER = logging.getLogger(__name__)


@attr.s(kw_only=True)
class AsyncCHYTAdapter(BaseAsyncClickHouseAdapter):
    ch_utils = CHYTUtils
    conn_type = CONNECTION_TYPE_CHYT

    _target_dto: BaseCHYTConnTargetDTO = attr.ib()

    def _get_current_tracing_headers(self) -> dict[str, str]:
        return get_current_w3c_tracing_headers(
            override_sampled=self.ch_utils.get_tracing_sample_flag_override(self._req_ctx_info),
            req_id=self._req_ctx_info.request_id,
        )

lib/dl_connector_chyt/dl_connector_chyt/core/async_adapters.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_connector_clickhouse.core.clickhouse_base.conn_options import CHConnectOptions


@attr.s(frozen=True, hash=True)
class CHYTConnectOptions(CHConnectOptions):
    pass

lib/dl_connector_chyt/dl_connector_chyt/core/conn_options.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging

import attr

from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor

from dl_connector_chyt.core.adapters import (
    BaseCHYTAdapter,
    CHYTAdapter,
)
from dl_connector_chyt.core.async_adapters import AsyncCHYTAdapter
from dl_connector_chyt.core.conn_options import CHYTConnectOptions
from dl_connector_chyt.core.dto import (
    CHYTDTO,
    BaseCHYTDTO,
)
from dl_connector_chyt.core.target_dto import (
    BaseCHYTConnTargetDTO,
    CHYTConnTargetDTO,
)


LOGGER = logging.getLogger(__name__)


@attr.s(cmp=False, hash=False)
class BaseCHYTAdapterConnExecutor(DefaultSqlAlchemyConnExecutor[BaseCHYTAdapter], abc.ABC):
    _conn_dto: BaseCHYTDTO = attr.ib()
    _conn_options: CHYTConnectOptions = attr.ib()

    @abc.abstractmethod
    async def _get_target_conn_dto(self) -> BaseCHYTConnTargetDTO:
        raise NotImplementedError()

    async def _make_target_conn_dto_pool(self) -> list[BaseCHYTConnTargetDTO]:
        target_dto = await self._get_target_conn_dto()
        return [target_dto]


@attr.s(cmp=False, hash=False)
class CHYTAdapterConnExecutor(BaseCHYTAdapterConnExecutor):
    _conn_dto: CHYTDTO = attr.ib()

    async def _get_target_conn_dto(self) -> CHYTConnTargetDTO:
        return CHYTConnTargetDTO(
            conn_id=self._conn_dto.conn_id,
            pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
            pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
            host=self._conn_dto.host,
            port=self._conn_dto.port,
            username="default",
            password=self._conn_dto.token,
            db_name=self._conn_dto.clique_alias,
            cluster_name=None,
            endpoint="query",
            protocol=self._conn_dto.protocol,
            max_execution_time=self._conn_options.max_execution_time,
            total_timeout=self._conn_options.total_timeout,
            connect_timeout=self._conn_options.connect_timeout,
            insert_quorum=self._conn_options.insert_quorum,
            insert_quorum_timeout=self._conn_options.insert_quorum_timeout,
            disable_value_processing=self._conn_options.disable_value_processing,
        )


@attr.s(cmp=False, hash=False)
class CHYTSyncAdapterConnExecutor(CHYTAdapterConnExecutor):
    TARGET_ADAPTER_CLS = CHYTAdapter


@attr.s(cmp=False, hash=False)
class CHYTAsyncAdapterConnExecutor(CHYTAdapterConnExecutor):
    TARGET_ADAPTER_CLS = AsyncCHYTAdapter  # type: ignore

lib/dl_connector_chyt/dl_connector_chyt/core/connection_executors.py
datalens-tech/datalens-backend
from clickhouse_sqlalchemy.orm.query import Query as CHQuery

from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
    CoreSourceDefinition,
)

from dl_connector_chyt.core.adapters import CHYTAdapter
from dl_connector_chyt.core.async_adapters import AsyncCHYTAdapter
from dl_connector_chyt.core.connection_executors import (
    CHYTAsyncAdapterConnExecutor,
    CHYTSyncAdapterConnExecutor,
)
from dl_connector_chyt.core.constants import (
    BACKEND_TYPE_CHYT,
    CONNECTION_TYPE_CHYT,
    SOURCE_TYPE_CHYT_YTSAURUS_SUBSELECT,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE_LIST,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE_RANGE,
)
from dl_connector_chyt.core.data_source import (
    CHYTTableDataSource,
    CHYTTableListDataSource,
    CHYTTableRangeDataSource,
    CHYTTableSubselectDataSource,
)
from dl_connector_chyt.core.data_source_migration import CHYTDataSourceMigrator
from dl_connector_chyt.core.data_source_spec import (
    CHYTSubselectDataSourceSpec,
    CHYTTableDataSourceSpec,
    CHYTTableListDataSourceSpec,
    CHYTTableRangeDataSourceSpec,
)
from dl_connector_chyt.core.settings import CHYTSettingDefinition
from dl_connector_chyt.core.storage_schemas.connection import ConnectionCHYTDataStorageSchema
from dl_connector_chyt.core.storage_schemas.data_source_spec import (
    CHYTSubselectDataSourceSpecStorageSchema,
    CHYTTableDataSourceSpecStorageSchema,
    CHYTTableListDataSourceSpecStorageSchema,
    CHYTTableRangeDataSourceSpecStorageSchema,
)
from dl_connector_chyt.core.type_transformer import CHYTTypeTransformer
from dl_connector_chyt.core.us_connection import ConnectionCHYTToken


class CHYTCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_CHYT
    connection_cls = ConnectionCHYTToken
    us_storage_schema_cls = ConnectionCHYTDataStorageSchema
    type_transformer_cls = CHYTTypeTransformer
    sync_conn_executor_cls = CHYTSyncAdapterConnExecutor
    async_conn_executor_cls = CHYTAsyncAdapterConnExecutor
    dialect_string = "bi_chyt"
    settings_definition = CHYTSettingDefinition
    data_source_migrator_cls = CHYTDataSourceMigrator


class CHYTTableCoreSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_CHYT_YTSAURUS_TABLE
    source_cls = CHYTTableDataSource
    source_spec_cls = CHYTTableDataSourceSpec
    us_storage_schema_cls = CHYTTableDataSourceSpecStorageSchema


class CHYTTableListCoreSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_CHYT_YTSAURUS_TABLE_LIST
    source_cls = CHYTTableListDataSource
    source_spec_cls = CHYTTableListDataSourceSpec
    us_storage_schema_cls = CHYTTableListDataSourceSpecStorageSchema


class CHYTTableRangeCoreSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_CHYT_YTSAURUS_TABLE_RANGE
    source_cls = CHYTTableRangeDataSource
    source_spec_cls = CHYTTableRangeDataSourceSpec
    us_storage_schema_cls = CHYTTableRangeDataSourceSpecStorageSchema


class CHYTTableSubselectCoreSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_CHYT_YTSAURUS_SUBSELECT
    source_cls = CHYTTableSubselectDataSource
    source_spec_cls = CHYTSubselectDataSourceSpec
    us_storage_schema_cls = CHYTSubselectDataSourceSpecStorageSchema


class CHYTCoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_CHYT
    connection_definitions = (CHYTCoreConnectionDefinition,)
    source_definitions = (
        CHYTTableCoreSourceDefinition,
        CHYTTableListCoreSourceDefinition,
        CHYTTableRangeCoreSourceDefinition,
        CHYTTableSubselectCoreSourceDefinition,
    )
    rqe_adapter_classes = frozenset(
        {
            CHYTAdapter,
            AsyncCHYTAdapter,
        }
    )
    query_cls = CHQuery

lib/dl_connector_chyt/dl_connector_chyt/core/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
)

from dl_connector_clickhouse.core.clickhouse_base.constants import BACKEND_TYPE_CLICKHOUSE


BACKEND_TYPE_CHYT = BACKEND_TYPE_CLICKHOUSE
CONNECTION_TYPE_CHYT = ConnectionType.declare("chyt")

SOURCE_TYPE_CHYT_YTSAURUS_TABLE = DataSourceType.declare("CHYT_YTSAURUS_TABLE")
SOURCE_TYPE_CHYT_YTSAURUS_SUBSELECT = DataSourceType.declare("CHYT_YTSAURUS_SUBSELECT")
SOURCE_TYPE_CHYT_YTSAURUS_TABLE_LIST = DataSourceType.declare("CHYT_YTSAURUS_TABLE_LIST")
SOURCE_TYPE_CHYT_YTSAURUS_TABLE_RANGE = DataSourceType.declare("CHYT_YTSAURUS_TABLE_RANGE")

lib/dl_connector_chyt/dl_connector_chyt/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import collections
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Optional,
)
import urllib.parse

import attr
import sqlalchemy as sa

from dl_constants.enums import (
    DataSourceType,
    UserDataType,
)
from dl_core import exc
from dl_core.connection_models.common_models import (
    SATextTableDefinition,
    TableDefinition,
    TableIdent,
)
from dl_core.data_source.sql import (
    BaseSQLDataSource,
    TableSQLDataSourceMixin,
    require_table_name,
)
from dl_core.db import (
    SchemaColumn,
    SchemaInfo,
)
from dl_core.db.native_type import ClickHouseNativeType
from dl_sqlalchemy_chyt import (
    CHYTTablesConcat,
    CHYTTablesRange,
    CHYTTableSubselect,
)

from dl_connector_chyt.core.constants import (
    CONNECTION_TYPE_CHYT,
    SOURCE_TYPE_CHYT_YTSAURUS_SUBSELECT,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE_LIST,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE_RANGE,
)
from dl_connector_chyt.core.data_source_spec import (
    CHYTTableDataSourceSpec,
    CHYTTableListDataSourceSpec,
    CHYTTableRangeDataSourceSpec,
)
from dl_connector_clickhouse.core.clickhouse_base.data_source import (
    ClickHouseBaseMixin,
    CommonClickHouseSubselectDataSource,
)


if TYPE_CHECKING:
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase


LOGGER = logging.getLogger(__name__)


class CHYTDataSourceBaseMixin(ClickHouseBaseMixin):
    @staticmethod
    def normalize_path(value: str) -> str:
        """
        ...


        >>> CHYTDataSourceBaseMixin.normalize_path('\\n //a / b/c ')
        '//a / b/c '
        >>> CHYTDataSourceBaseMixin.normalize_path('https://host.net/cluster/navigation?path=//a%20/ b/c+&offsetMode=row')
        '//a / b/c '
        """
        raw_value = value.lstrip()
        if raw_value.startswith("//"):
            return raw_value.rstrip("/")
        parsed_url = urllib.parse.urlparse(raw_value)
        parsed_query = urllib.parse.parse_qs(parsed_url.query)
        res = parsed_query.get("path", [None])[0]  # type: ignore
        if res:
            return res  # No extra normalization
        raise exc.TableNameInvalidError("Invalid CHYT table path")


class BaseCHYTTableDataSource(CHYTDataSourceBaseMixin, TableSQLDataSourceMixin, abc.ABC):
    @property
    def spec(self) -> CHYTTableDataSourceSpec:
        assert isinstance(self._spec, CHYTTableDataSourceSpec)
        return self._spec

    @property
    def default_title(self) -> str:
        return self.spec.table_name.split("/")[-1]  # type: ignore  # TODO: fix

    @require_table_name
    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        if alias:
            return sa.alias(self.get_sql_source(), name=alias)
        path = self.spec.table_name
        path = self.normalize_path(path)  # type: ignore  # TODO: fix
        return sa.table(path)

    @require_table_name
    def get_table_definition(self) -> TableDefinition:
        table_name = self.table_name
        assert table_name is not None
        return TableIdent(db_name=None, schema_name=None, table_name=table_name)


class BaseCHYTSpecialDataSource(CHYTDataSourceBaseMixin, BaseSQLDataSource, abc.ABC):
    """A common base class for CHYT custom sources (like table funcs or subselects)"""

    @property
    def default_title(self) -> str:
        raise NotImplementedError

    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        raise NotImplementedError

    def get_table_definition(self) -> TableDefinition:
        return SATextTableDefinition(self.get_sql_source())


class BaseCHYTTableFuncDataSource(BaseCHYTSpecialDataSource, abc.ABC):
    """A common base class for CHYT Table-Function-based data source (concat / concat range)."""

    def get_schema_info(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> SchemaInfo:
        assert self.conn_type is not None
        schema_info = super().get_schema_info(conn_executor_factory=conn_executor_factory)
        return schema_info.clone(
            schema=tuple(
                [
                    *schema_info.schema,
                    *[
                        SchemaColumn(
                            name=key,
                            title=key,
                            user_type=UserDataType.string,
                            native_type=ClickHouseNativeType.normalize_name_and_create(
                                conn_type=self.conn_type, name="string"
                            ),
                            nullable=False,
                        )
                        for key in ("$table_path", "$table_name", "$table_index")
                    ],
                ]
            )
        )


@attr.s
class BaseCHYTTableListDataSource(BaseCHYTTableFuncDataSource, abc.ABC):
    table_names: str = attr.ib(kw_only=True, default=None)  # stored as single `\n` separated string

    @property
    def spec(self) -> CHYTTableListDataSourceSpec:
        assert isinstance(self._spec, CHYTTableListDataSourceSpec)
        return self._spec

    @classmethod
    def normalize_tables_paths(cls, value: str) -> list[str]:
        table_paths = value.splitlines()
        table_paths = [cls.normalize_path(path) for path in table_paths if path.lstrip()]
        if not table_paths:
            raise exc.TableNameInvalidError("Table list is empty")
        if len(table_paths) != len(set(table_paths)):
            duplicates = [item for item, cnt in collections.Counter(table_paths).items() if cnt > 1]
            assert duplicates, ("should match the pre-check", table_paths)
            duplicates_s = "\n".join(duplicates)
            raise exc.TableNameInvalidError(f"Table list contains duplicates:\n{duplicates_s}")
        return table_paths

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            table_names=self.spec.table_names,
        )

    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        if not self.spec.table_names:
            raise exc.TableNameNotConfiguredError
        table_names = self.normalize_tables_paths(self.spec.table_names)
        return CHYTTablesConcat(*table_names, alias=alias)

    @property
    def default_title(self) -> str:
        assert self.spec.table_names is not None
        return self.normalize_tables_paths(self.spec.table_names)[0].split("/")[-1]


def _value_or_none(value: Optional[str]) -> Optional[str]:
    return value or None


@attr.s
class BaseCHYTTableRangeDataSource(BaseCHYTTableFuncDataSource, abc.ABC):
    _directory_path: Optional[str] = attr.ib(default=None)
    _range_from: Optional[str] = attr.ib(default=None, converter=_value_or_none)
    _range_to: Optional[str] = attr.ib(default=None, converter=_value_or_none)

    @property
    def spec(self) -> CHYTTableRangeDataSourceSpec:
        assert isinstance(self._spec, CHYTTableRangeDataSourceSpec)
        return self._spec

    @property
    def directory_path(self) -> Optional[str]:
        return self.spec.directory_path

    @property
    def range_from(self) -> Optional[str]:
        return self.spec.range_from

    @property
    def range_to(self) -> Optional[str]:
        return self.spec.range_to

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            directory_path=self.directory_path,
            range_from=self.range_from or "",
            range_to=self.range_to or "",
        )

    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        if not self.directory_path:
            raise exc.TableNameNotConfiguredError
        path = self.directory_path
        path = self.normalize_path(path)
        return CHYTTablesRange(directory=path, start=self.range_from, end=self.range_to, alias=alias)

    @property
    def default_title(self) -> str:
        assert self.directory_path is not None
        return self.normalize_path(self.directory_path).split("/")[-1]


# Expected MRO:
# CHYT-side: BaseCHYTTableSubselectDataSource, BaseCHYTTableFuncDataSource, CHYTBaseMixin, ClickHouseBaseMixin,
# Subselect-side: CommonClickHouseSubselectDataSource, SubselectDataSource,
# Base-side: BaseSQLDataSource, DataSource
# The Subselect-side isn't particularly necessary, but it marks the class as relevant.
class BaseCHYTTableSubselectDataSource(BaseCHYTSpecialDataSource, CommonClickHouseSubselectDataSource, abc.ABC):
    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            subsql=self.subsql,
        )

    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        if not self.connection.is_subselect_allowed:
            raise exc.SubselectNotAllowed()

        subsql = self.subsql
        if not subsql:
            raise exc.TableNameNotConfiguredError
        return CHYTTableSubselect(subsql, alias=alias)

    @property
    def default_title(self) -> str:
        return "(select ...)"


class CHYTTokenAuthDataSourceMixin:
    conn_type = CONNECTION_TYPE_CHYT

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type in {
            SOURCE_TYPE_CHYT_YTSAURUS_TABLE,
            SOURCE_TYPE_CHYT_YTSAURUS_SUBSELECT,
            SOURCE_TYPE_CHYT_YTSAURUS_TABLE_LIST,
            SOURCE_TYPE_CHYT_YTSAURUS_TABLE_RANGE,
        }


class CHYTTableDataSource(CHYTTokenAuthDataSourceMixin, BaseCHYTTableDataSource):
    pass


class CHYTTableListDataSource(CHYTTokenAuthDataSourceMixin, BaseCHYTTableListDataSource):
    pass


class CHYTTableRangeDataSource(CHYTTokenAuthDataSourceMixin, BaseCHYTTableRangeDataSource):
    pass


class CHYTTableSubselectDataSource(CHYTTokenAuthDataSourceMixin, BaseCHYTTableSubselectDataSource):
    pass

lib/dl_connector_chyt/dl_connector_chyt/core/data_source.py
datalens-tech/datalens-backend
from typing import (
    ClassVar,
    Optional,
    Type,
)

import attr

from dl_constants.enums import DataSourceType
from dl_core.connectors.base.data_source_migration import (
    DataSourceMigrationInterface,
    MigrationKeyMappingItem,
    MigrationSpec,
)
from dl_core.connectors.sql_base.data_source_migration import DefaultSQLDataSourceMigrator
from dl_core.data_source_spec.base import DataSourceSpec

from dl_connector_chyt.core.constants import (
    SOURCE_TYPE_CHYT_YTSAURUS_SUBSELECT,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE_LIST,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE_RANGE,
)
from dl_connector_chyt.core.data_source_spec import (
    CHYTSubselectDataSourceSpec,
    CHYTTableDataSourceSpec,
    CHYTTableListDataSourceSpec,
    CHYTTableRangeDataSourceSpec,
)


@attr.s(frozen=True)
class CHYTTableListDSMI(DataSourceMigrationInterface):
    table_names: Optional[str] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class CHYTTableRangeDSMI(DataSourceMigrationInterface):
    directory_path: Optional[str] = attr.ib(kw_only=True, default=None)
    range_from: Optional[str] = attr.ib(kw_only=True, default=None)
    range_to: Optional[str] = attr.ib(kw_only=True, default=None)


class BaseCHYTDataSourceMigrator(DefaultSQLDataSourceMigrator):
    table_list_source_type: ClassVar[Optional[DataSourceType]]
    table_list_dsrc_spec_cls: ClassVar[Optional[Type[DataSourceSpec]]]

    table_range_source_type: ClassVar[Optional[DataSourceType]]
    table_range_dsrc_spec_cls: ClassVar[Optional[Type[DataSourceSpec]]]

    def get_migration_specs(self) -> list[MigrationSpec]:
        result = super().get_migration_specs()

        if self.table_list_source_type is not None:
            assert self.table_list_dsrc_spec_cls is not None
            result.append(
                MigrationSpec(
                    source_type=self.table_list_source_type,
                    dto_cls=CHYTTableListDSMI,
                    dsrc_spec_cls=self.table_list_dsrc_spec_cls,
                    migration_mapping_items=(
                        MigrationKeyMappingItem(migration_dto_key="table_names", source_spec_key="table_names"),
                    ),
                )
            )

        if self.table_range_source_type is not None:
            assert self.table_range_dsrc_spec_cls is not None
            result.append(
                MigrationSpec(
                    source_type=self.table_range_source_type,
                    dto_cls=CHYTTableRangeDSMI,
                    dsrc_spec_cls=self.table_range_dsrc_spec_cls,
                    migration_mapping_items=(
                        MigrationKeyMappingItem(migration_dto_key="directory_path", source_spec_key="directory_path"),
                        MigrationKeyMappingItem(migration_dto_key="range_from", source_spec_key="range_from"),
                        MigrationKeyMappingItem(migration_dto_key="range_to", source_spec_key="range_to"),
                    ),
                )
            )

        return result


class CHYTDataSourceMigrator(BaseCHYTDataSourceMigrator):
    table_source_type = SOURCE_TYPE_CHYT_YTSAURUS_TABLE
    table_dsrc_spec_cls = CHYTTableDataSourceSpec

    subselect_source_type = SOURCE_TYPE_CHYT_YTSAURUS_SUBSELECT
    subselect_dsrc_spec_cls = CHYTSubselectDataSourceSpec

    table_list_source_type = SOURCE_TYPE_CHYT_YTSAURUS_TABLE_LIST
    table_list_dsrc_spec_cls = CHYTTableListDataSourceSpec

    table_range_source_type = SOURCE_TYPE_CHYT_YTSAURUS_TABLE_RANGE
    table_range_dsrc_spec_cls = CHYTTableRangeDataSourceSpec

lib/dl_connector_chyt/dl_connector_chyt/core/data_source_migration.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    TypeVar,
)

import attr

from dl_core.data_source_spec.sql import (
    SQLDataSourceSpecBase,
    SubselectDataSourceSpec,
    TableSQLDataSourceSpec,
)


class CHYTDataSourceSpecMixin(SQLDataSourceSpecBase):
    pass


class CHYTTableDataSourceSpec(CHYTDataSourceSpecMixin, TableSQLDataSourceSpec):
    pass


class CHYTSubselectDataSourceSpec(CHYTDataSourceSpecMixin, SubselectDataSourceSpec):
    pass


@attr.s
class CHYTTableListDataSourceSpec(CHYTDataSourceSpecMixin, SQLDataSourceSpecBase):
    # stored as a single `\n` separated string
    table_names: Optional[str] = attr.ib(kw_only=True, default=None)


_VAL_TV = TypeVar("_VAL_TV")


def _zero_to_none(value: _VAL_TV) -> Optional[_VAL_TV]:  # type: ignore
    return value or None


@attr.s
class CHYTTableRangeDataSourceSpec(CHYTDataSourceSpecMixin, SQLDataSourceSpecBase):
    directory_path: Optional[str] = attr.ib(kw_only=True, default=None)
    range_from: Optional[str] = attr.ib(kw_only=True, default=None, converter=_zero_to_none)  # type: ignore
    range_to: Optional[str] = attr.ib(kw_only=True, default=None, converter=_zero_to_none)  # type: ignore

lib/dl_connector_chyt/dl_connector_chyt/core/data_source_spec.py
datalens-tech/datalens-backend
import attr

from dl_connector_chyt.core.constants import CONNECTION_TYPE_CHYT
from dl_connector_clickhouse.core.clickhouse_base.dto import ClickHouseBaseDTO


@attr.s(frozen=True, kw_only=True)
class BaseCHYTDTO(ClickHouseBaseDTO):
    clique_alias: str = attr.ib(kw_only=True)


@attr.s(frozen=True, kw_only=True)
class CHYTDTO(BaseCHYTDTO):
    conn_type = CONNECTION_TYPE_CHYT

    port: int = attr.ib()
    host: str = attr.ib()
    protocol: str = attr.ib()
    token: str = attr.ib(repr=False)

lib/dl_connector_chyt/dl_connector_chyt/core/dto.py
datalens-tech/datalens-backend
from dl_core import exc


class CHYTQueryError(exc.DatabaseQueryError):
    err_code = exc.DatabaseQueryError.err_code + ["CHYT"]
    default_message = "CHYT Error"


class CHYTAuthError(CHYTQueryError):
    err_code = CHYTQueryError.err_code + ["AUTH_FAILED"]
    default_message = "Authentication failed"


class CHYTCliqueError(exc.DatabaseQueryError):
    err_code = CHYTQueryError.err_code + ["CLIQUE"]
    default_message = "Clique Error"


class CHYTCliqueIsNotRunning(CHYTCliqueError):
    err_code = CHYTCliqueError.err_code + ["NOT_RUNNING"]
    default_message = "Clique is not running."
    formatting_messages = {frozenset({"clique"}): "Clique {clique} is not running."}


class CHYTCliqueIsSuspended(CHYTCliqueError):
    err_code = CHYTCliqueError.err_code + ["SUSPENDED"]
    default_message = "Clique is suspended."
    formatting_messages = {frozenset({"clique"}): "Clique {clique} is suspended."}


class CHYTCliqueNotExists(CHYTCliqueError):
    err_code = CHYTCliqueError.err_code + ["INVALID_SPECIFICATION"]
    default_message = "Invalid clique specification. Probably, clique does not exists."


class CHYTCliqueAccessDenied(CHYTCliqueError):
    err_code = CHYTCliqueError.err_code + ["ACCESS_DENIED"]
    default_message = "Access to clique was denied."
    formatting_messages = {frozenset({"clique", "user"}): "Access to clique {clique} for user {user} was denied."}


class CHYTTableHasNoSchema(CHYTQueryError):
    err_code = CHYTQueryError.err_code + ["TABLE_HAS_NO_SCHEMA"]
    default_message = "The table has no schema. Only schematized tables are supported."


class CHYTTableAccessDenied(CHYTQueryError):
    err_code = CHYTQueryError.err_code + ["TABLE_ACCESS_DENIED"]
    default_message = "Access to table was denied."


class CHYTInvalidSortedJoin(CHYTQueryError):
    err_code = CHYTQueryError.err_code + ["INVALID_SORTED_JOIN"]
    default_message = "Invalid sorted JOIN."


class CHYTInvalidSortedJoinNotAKeyColumn(CHYTInvalidSortedJoin):
    err_code = CHYTInvalidSortedJoin.err_code + ["NOT_A_KEY_COLUMN"]
    default_message = "Column used in join expression is not a key column."
    formatting_messages = {frozenset({"col"}): "Column {col} used in join expression is not a key column."}


class CHYTInvalidSortedJoinNotKeyPrefixColumn(CHYTInvalidSortedJoin):
    err_code = CHYTInvalidSortedJoin.err_code + ["NOT_KEY_PREFIX_COLUMN"]
    default_message = "Joined columns should form prefix of joined table key columns."


class CHYTInvalidSortedJoinMoreThanOneTable(CHYTInvalidSortedJoin):
    err_code = CHYTInvalidSortedJoin.err_code + ["MORE_THAN_ONE_TABLE"]
    default_message = (
        "Cannot join a concatenation of tables with another table. " "Wrap concatenation into an SQL query."
    )


class CHYTInvalidSortedJoinTableNotSorted(CHYTInvalidSortedJoin):
    err_code = CHYTInvalidSortedJoin.err_code + ["TABLE_NOT_SORTED"]
    default_message = "Tables should be sorted"
    formatting_messages = {frozenset({"table"}): "Table {table} should be sorted."}


class CHYTInvalidSortedJoinCompoundExpressionsNotSupported(CHYTInvalidSortedJoin):
    err_code = CHYTInvalidSortedJoin.err_code + ["COMPOUND_EXPR_NOT_SUPPORTED"]
    default_message = "CHYT does not support compound expressions in ON/USING clause"


class CHYTInvalidSortedJoinKeyIsEmpty(CHYTInvalidSortedJoin):
    err_code = CHYTInvalidSortedJoin.err_code + ["KEY_IS_EMPTY"]
    default_message = "Cannot join: key is empty"


class CHYTInvalidSortedJoinConcatNotSupported(CHYTInvalidSortedJoin):
    err_code = CHYTInvalidSortedJoin.err_code + ["CONCAT_NOT_SUPPORTED"]
    default_message = "Joining concatenation of multiple tables is not supported"


class CHYTInvalidSortedJoinNotSameKeyPosition(CHYTInvalidSortedJoin):
    err_code = CHYTInvalidSortedJoin.err_code + ["NOT_SAME_KEY_POSITION"]
    default_message = "Joined columns do not occupy same positions in key columns of joined tables"
    formatting_messages = {
        frozenset({"col1", "col2"}): "Joined columns {col1} and {col2}"
        " do not occupy same positions in key columns of joined tables",
    }


class CHYTMultipleDynamicTablesNotSupported(CHYTQueryError):
    err_code = CHYTQueryError.err_code + ["MULTI_DYN_NOT_SUPPORTED"]
    default_message = (
        "Reading multiple dynamic tables " "or dynamic table together with static table is not supported in CHYT"
    )


class CHYTSubqueryWeightLimitExceeded(CHYTQueryError):
    err_code = CHYTQueryError.err_code + ["SUBQ_WEIGHT_LIMIT_EXCEEDED"]
    default_message = "Subquery exceeds data weight limit"

lib/dl_connector_chyt/dl_connector_chyt/core/exc.py
datalens-tech/datalens-backend
from typing import (
    ClassVar,
    Optional,
)

import attr

from dl_configs.connectors_data import ConnectorsDataBase
from dl_configs.connectors_settings import (
    ConnectorsConfigType,
    ConnectorSettingsBase,
)
from dl_configs.settings_loaders.meta_definition import s_attrib
from dl_configs.utils import split_by_comma
from dl_core.connectors.settings.primitives import (
    ConnectorSettingsDefinition,
    get_connectors_settings_config,
)


@attr.s(frozen=True)
class CHYTConnectorSettings(ConnectorSettingsBase):
    """
    PUBLIC_CLIQUES:     cliques which usage is discouraged due to their high load by other users
    FORBIDDEN_CLIQUES:  cliques that cannot be used at all
    DEFAULT_CLIQUE:     clique that is set by default in the connection form
    """

    PUBLIC_CLIQUES: tuple[str] = s_attrib("PUBLIC_CLIQUES", missing_factory=tuple, env_var_converter=split_by_comma)  # type: ignore
    FORBIDDEN_CLIQUES: tuple[str] = s_attrib("FORBIDDEN_CLIQUES", missing_factory=tuple, env_var_converter=split_by_comma)  # type: ignore
    DEFAULT_CLIQUE: Optional[str] = s_attrib("DEFAULT_CLIQUE", missing=None)  # type: ignore


class ConnectorsDataCHYTBase(ConnectorsDataBase):
    CONN_CHYT_PUBLIC_CLIQUES: ClassVar[Optional[tuple[str]]] = None
    CONN_CHYT_FORBIDDEN_CLIQUES: ClassVar[Optional[tuple[str]]] = None
    CONN_CHYT_DEFAULT_CLIQUE: ClassVar[Optional[str]] = None

    @classmethod
    def connector_name(cls) -> str:
        return "CHYT"


def chyt_settings_fallback(full_cfg: ConnectorsConfigType) -> dict[str, ConnectorSettingsBase]:
    cfg = get_connectors_settings_config(
        full_cfg,
        object_like_config_key="CHYT",
        connector_data_class=ConnectorsDataCHYTBase,
    )
    if cfg is None:
        return {}
    return dict(
        CHYT=CHYTConnectorSettings(  # type: ignore
            PUBLIC_CLIQUES=tuple(cfg.CONN_CHYT_PUBLIC_CLIQUES),
            FORBIDDEN_CLIQUES=tuple(cfg.CONN_CHYT_FORBIDDEN_CLIQUES),
            DEFAULT_CLIQUE=cfg.CONN_CHYT_DEFAULT_CLIQUE,
        )
    )


class CHYTSettingDefinition(ConnectorSettingsDefinition):
    settings_class = CHYTConnectorSettings
    fallback = chyt_settings_fallback

lib/dl_connector_chyt/dl_connector_chyt/core/settings.py
datalens-tech/datalens-backend
from typing import (
    Generic,
    TypeVar,
)

from marshmallow import fields as ma_fields

from dl_core.us_manager.storage_schemas.connection import (
    CacheableConnectionDataSchemaMixin,
    ConnectionBaseDataStorageSchema,
    SubselectConnectionDataSchemaMixin,
)

from dl_connector_chyt.core.us_connection import (
    BaseConnectionCHYT,
    ConnectionCHYTToken,
)


_CHYT_CONN_DATA_TV = TypeVar("_CHYT_CONN_DATA_TV", bound=BaseConnectionCHYT.DataModel)


class ConnectionCHYTBaseDataStorageSchema(
    CacheableConnectionDataSchemaMixin,
    SubselectConnectionDataSchemaMixin,
    ConnectionBaseDataStorageSchema[_CHYT_CONN_DATA_TV],
    Generic[_CHYT_CONN_DATA_TV],
):
    alias = ma_fields.String(required=True, allow_none=False)
    max_execution_time = ma_fields.Integer(required=False, allow_none=True, load_default=None, dump_default=None)


class ConnectionCHYTDataStorageSchema(ConnectionCHYTBaseDataStorageSchema[ConnectionCHYTToken.DataModel]):
    TARGET_CLS = ConnectionCHYTToken.DataModel

    token = ma_fields.String(allow_none=False, required=True)
    host = ma_fields.String(allow_none=False, required=True)
    port = ma_fields.Integer(allow_none=False, required=True)
    secure = ma_fields.Boolean(allow_none=False, required=True)

lib/dl_connector_chyt/dl_connector_chyt/core/storage_schemas/connection.py
datalens-tech/datalens-backend
import marshmallow.fields as ma_fields

from dl_core.us_manager.storage_schemas.data_source_spec_base import (
    BaseSQLDataSourceSpecStorageSchema,
    SubselectDataSourceSpecStorageSchema,
    TableSQLDataSourceSpecStorageSchemaMixin,
)

from dl_connector_chyt.core.data_source_spec import (
    CHYTSubselectDataSourceSpec,
    CHYTTableDataSourceSpec,
    CHYTTableListDataSourceSpec,
    CHYTTableRangeDataSourceSpec,
)


class CHYTTableDataSourceSpecStorageSchema(TableSQLDataSourceSpecStorageSchemaMixin):
    TARGET_CLS = CHYTTableDataSourceSpec


class CHYTSubselectDataSourceSpecStorageSchema(SubselectDataSourceSpecStorageSchema):
    TARGET_CLS = CHYTSubselectDataSourceSpec


class CHYTTableListDataSourceSpecStorageSchema(BaseSQLDataSourceSpecStorageSchema):
    TARGET_CLS = CHYTTableListDataSourceSpec

    table_names = ma_fields.String(required=False, allow_none=True)


class CHYTTableRangeDataSourceSpecStorageSchema(BaseSQLDataSourceSpecStorageSchema):
    TARGET_CLS = CHYTTableRangeDataSourceSpec

    directory_path = ma_fields.String(required=False, allow_none=True, load_default=None)
    range_from = ma_fields.String(required=False, allow_none=True, load_default=None)
    range_to = ma_fields.String(required=False, allow_none=True, load_default=None)

lib/dl_connector_chyt/dl_connector_chyt/core/storage_schemas/data_source_spec.py
datalens-tech/datalens-backend
import attr

from dl_connector_clickhouse.core.clickhouse_base.target_dto import BaseClickHouseConnTargetDTO


@attr.s(frozen=True)
class BaseCHYTConnTargetDTO(BaseClickHouseConnTargetDTO):
    pass


@attr.s(frozen=True)
class CHYTConnTargetDTO(BaseCHYTConnTargetDTO):
    pass

lib/dl_connector_chyt/dl_connector_chyt/core/target_dto.py
datalens-tech/datalens-backend
from typing import ClassVar

from dl_constants.enums import ConnectionType as CT
from dl_constants.enums import UserDataType
from dl_core.db.conversion_base import (
    TypeCaster,
    YTBooleanTypeCaster,
    make_native_type,
)
from dl_core.db.native_type import GenericNativeType
from dl_sqlalchemy_chyt import types as chyt_types

from dl_connector_chyt.core.constants import CONNECTION_TYPE_CHYT
from dl_connector_clickhouse.core.clickhouse_base.type_transformer import ClickHouseTypeTransformer


def make_chyt_native_to_user_map(conn_type: CT) -> dict[GenericNativeType, UserDataType]:
    return {
        **{
            make_native_type(conn_type, native_type.name): bi_type
            for native_type, bi_type in ClickHouseTypeTransformer.native_to_user_map.items()
        },
        make_native_type(conn_type, chyt_types.YtBoolean): UserDataType.boolean,
    }


def make_chyt_user_to_native_map(conn_type: CT) -> dict[UserDataType, GenericNativeType]:
    return {
        **{
            bi_type: make_native_type(conn_type, native_type.name)
            for bi_type, native_type in ClickHouseTypeTransformer.user_to_native_map.items()
        },
        UserDataType.boolean: make_native_type(conn_type, chyt_types.YtBoolean),
    }


class CHYTTypeTransformer(ClickHouseTypeTransformer):
    conn_type = CONNECTION_TYPE_CHYT

    native_to_user_map = make_chyt_native_to_user_map(CONNECTION_TYPE_CHYT)
    user_to_native_map = make_chyt_user_to_native_map(CONNECTION_TYPE_CHYT)

    casters: ClassVar[dict[UserDataType, TypeCaster]] = {
        **ClickHouseTypeTransformer.casters,
        UserDataType.boolean: YTBooleanTypeCaster(),
    }

lib/dl_connector_chyt/dl_connector_chyt/core/type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Optional,
)

import attr
import marshmallow as ma

from dl_constants.enums import DataSourceType
from dl_core.base_models import (
    ConnCacheableDataModelMixin,
    ConnSubselectDataModelMixin,
)
from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.us_connection_base import (
    ConnectionBase,
    ConnectionHardcodedDataMixin,
    DataSourceTemplate,
    ExecutorBasedMixin,
    SubselectMixin,
)
from dl_core.utils import secrepr
from dl_i18n.localizer_base import Localizer
from dl_utils.utils import DataKey

from dl_connector_chyt.api.i18n.localizer import Translatable
from dl_connector_chyt.core.conn_options import CHYTConnectOptions
from dl_connector_chyt.core.constants import (
    SOURCE_TYPE_CHYT_YTSAURUS_SUBSELECT,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE_LIST,
    SOURCE_TYPE_CHYT_YTSAURUS_TABLE_RANGE,
)
from dl_connector_chyt.core.dto import CHYTDTO
from dl_connector_chyt.core.settings import CHYTConnectorSettings
from dl_connector_clickhouse.core.clickhouse_base.us_connection import ConnectionClickhouseBase


if TYPE_CHECKING:
    from dl_core.services_registry.top_level import ServicesRegistry


class BaseConnectionCHYT(
    SubselectMixin,
    ExecutorBasedMixin,
    ConnectionBase,
    ConnectionHardcodedDataMixin[CHYTConnectorSettings],
    abc.ABC,
):
    allow_dashsql: ClassVar[bool] = True
    is_always_user_source: ClassVar[bool] = True
    settings_type = CHYTConnectorSettings

    chyt_table_source_type: ClassVar[DataSourceType]
    chyt_table_list_source_type: ClassVar[DataSourceType]
    chyt_table_range_source_type: ClassVar[DataSourceType]
    chyt_subselect_source_type: ClassVar[DataSourceType]

    @attr.s(kw_only=True)
    class DataModel(ConnCacheableDataModelMixin, ConnSubselectDataModelMixin, ConnectionBase.DataModel):
        alias: str = attr.ib()
        max_execution_time: Optional[int] = attr.ib(default=None)

    async def validate_new_data(
        self,
        services_registry: ServicesRegistry,
        changes: Optional[dict] = None,
        original_version: Optional[ConnectionBase] = None,
    ) -> None:
        chyt_settings = self._connector_settings
        if self.data.alias in chyt_settings.FORBIDDEN_CLIQUES:
            err_msg = f'It is forbidden to use the following cliques: {", ".join(chyt_settings.FORBIDDEN_CLIQUES)}.'
            if chyt_settings.PUBLIC_CLIQUES:
                err_msg += f' Valid public cliques are: {", ".join(chyt_settings.PUBLIC_CLIQUES)}.'
            raise ma.ValidationError(message=err_msg)

    @property
    def cache_ttl_sec_override(self) -> Optional[int]:
        return self.data.cache_ttl_sec

    def get_conn_options(self) -> CHYTConnectOptions:
        base = super().get_conn_options()
        base_ch_connect_options = ConnectionClickhouseBase.get_effective_conn_options(
            base_conn_opts=base,
            user_max_execution_time=self.data.max_execution_time,
            max_allowed_max_execution_time=ConnectionClickhouseBase.MAX_ALLOWED_MAX_EXECUTION_TIME,
        )
        return base_ch_connect_options.to_subclass(CHYTConnectOptions)

    @abc.abstractmethod
    def get_data_source_template_templates(self, localizer: Localizer) -> list[DataSourceTemplate]:
        raise NotImplementedError()

    def get_data_source_templates(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[DataSourceTemplate]:
        return []


class ConnectionCHYTToken(BaseConnectionCHYT):
    allow_cache: ClassVar[bool] = True

    source_type = SOURCE_TYPE_CHYT_YTSAURUS_TABLE
    allowed_source_types = frozenset(
        (
            SOURCE_TYPE_CHYT_YTSAURUS_TABLE,
            SOURCE_TYPE_CHYT_YTSAURUS_SUBSELECT,
            SOURCE_TYPE_CHYT_YTSAURUS_TABLE_LIST,
            SOURCE_TYPE_CHYT_YTSAURUS_TABLE_RANGE,
        )
    )

    chyt_table_source_type = SOURCE_TYPE_CHYT_YTSAURUS_TABLE
    chyt_table_list_source_type = SOURCE_TYPE_CHYT_YTSAURUS_TABLE_LIST
    chyt_table_range_source_type = SOURCE_TYPE_CHYT_YTSAURUS_TABLE_RANGE
    chyt_subselect_source_type = SOURCE_TYPE_CHYT_YTSAURUS_SUBSELECT

    @attr.s(kw_only=True)
    class DataModel(BaseConnectionCHYT.DataModel):
        host: str = attr.ib()
        port: int = attr.ib()
        token: str = attr.ib(repr=secrepr)
        secure: bool = attr.ib(default=False)

        @classmethod
        def get_secret_keys(cls) -> set[DataKey]:
            return {
                *super().get_secret_keys(),
                DataKey(parts=("token",)),
            }

    @property
    def token(self) -> str:
        return self.data.token

    def get_data_source_template_templates(self, localizer: Localizer) -> list[DataSourceTemplate]:
        common: dict[str, Any] = dict(
            group=[],
            connection_id=self.uuid,
            parameters={},
        )
        return [
            DataSourceTemplate(
                title="YTsaurus table via CHYT",
                tab_title=localizer.translate(Translatable("source_templates-tab_title-table")),
                source_type=self.chyt_table_source_type,
                form=[
                    {
                        "name": "table_name",
                        "input_type": "text",
                        "default": "",
                        "required": True,
                        "title": localizer.translate(Translatable("source_templates-label-ytsaurus_table")),
                        "field_doc_key": "YTsaurus/CHYT_TABLE/table_name",
                    },
                ],
                **common,
            ),
            DataSourceTemplate(
                title="List of YTsaurus tables via CHYT",
                tab_title=localizer.translate(Translatable("source_templates-tab_title-concat")),
                source_type=self.chyt_table_list_source_type,
                form=[
                    {
                        "name": "table_names",
                        "input_type": "textarea",
                        "default": "",
                        "required": True,
                        "title": localizer.translate(Translatable("source_templates-label-ytasurus_table_list")),
                        "field_doc_key": "YTsaurus/CHYT_TABLE_LIST/table_names",
                    },
                ],
                **common,
            ),
            DataSourceTemplate(
                title="Range of YTsaurus tables via CHYT",
                tab_title=localizer.translate(Translatable("source_templates-tab_title-range")),
                source_type=self.chyt_table_range_source_type,
                form=[
                    {
                        "name": "directory_path",
                        "input_type": "text",
                        "default": "",
                        "required": True,
                        "title": localizer.translate(Translatable("source_templates-label-ytsaurus_dir")),
                        "field_doc_key": "YTsaurus/CHYT_TABLE_RANGE/directory_path",
                    },
                    {
                        "name": "range_from",
                        "input_type": "text",
                        "default": "",
                        "required": False,
                        "title": localizer.translate(Translatable("source_templates-label-range_from")),
                    },
                    {
                        "name": "range_to",
                        "input_type": "text",
                        "default": "",
                        "required": False,
                        "title": localizer.translate(Translatable("source_templates-label-range_to")),
                    },
                ],
                **common,
            ),
        ] + self._make_subselect_templates(
            title="SQL query via CHYT",
            source_type=self.chyt_subselect_source_type,
            field_doc_key="YTsaurus/CHYT_SUBSELECT/subsql",
            localizer=localizer,
        )

    def get_conn_dto(self) -> CHYTDTO:
        return CHYTDTO(
            conn_id=self.uuid,
            clique_alias=self.data.alias,
            token=self.token,
            host=self.data.host,
            port=self.data.port,
            protocol="https" if self.data.secure else "http",
        )

lib/dl_connector_chyt/dl_connector_chyt/core/us_connection.py
datalens-tech/datalens-backend
import logging
import re
from typing import (
    Optional,
    Type,
)

from dl_core import exc
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI

from dl_connector_chyt.core import exc as chyt_exc
from dl_connector_clickhouse.core.clickhouse_base.ch_commons import (
    ClickHouseBaseUtils,
    ParsedErrorMsg,
)


LOGGER = logging.getLogger(__name__)


class CHYTUtils(ClickHouseBaseUtils):
    add_real_user_header = True

    chyt_expr_exc = {
        r"Invalid sorted JOIN: (?P<col>.*) is not a key column": chyt_exc.CHYTInvalidSortedJoinNotAKeyColumn,
        r"Invalid sorted JOIN: joined columns should form prefix of joined table key columns": chyt_exc.CHYTInvalidSortedJoinNotKeyPrefixColumn,
        r"Invalid sorted JOIN: only single table may currently be joined": chyt_exc.CHYTInvalidSortedJoinMoreThanOneTable,
        r"Invalid sorted JOIN: table (?P<table>.*) is not sorted": chyt_exc.CHYTInvalidSortedJoinTableNotSorted,
        r"Invalid sorted JOIN: CHYT does not support compound expressions in ON/USING clause": chyt_exc.CHYTInvalidSortedJoinCompoundExpressionsNotSupported,
        r"Invalid sorted JOIN: key is empty": chyt_exc.CHYTInvalidSortedJoinKeyIsEmpty,
        r"Invalid sorted JOIN: joining concatenation of multiple tables is not supported": chyt_exc.CHYTInvalidSortedJoinConcatNotSupported,
        r"Invalid sorted JOIN: joined columns (?P<col1>.*) and (?P<col2>.*)"
        r" do not occupy same positions in key columns of joined tables": chyt_exc.CHYTInvalidSortedJoinNotSameKeyPosition,
        r"Invalid sorted JOIN": chyt_exc.CHYTInvalidSortedJoin,
        r"Access denied": chyt_exc.CHYTTableAccessDenied,
        r"Error validating permissions for user": chyt_exc.CHYTTableAccessDenied,
        r"CHYT does not support tables without schema": chyt_exc.CHYTTableHasNoSchema,
        r"NYT::TErrorException: Memory limit \(total\) exceeded": exc.DbMemoryLimitExceeded,
        r"Error resolving path": exc.SourceDoesNotExist,
        r"No tables to read from": exc.SourceDoesNotExist,
        r"Reading multiple dynamic tables or dynamic table together with static table is not supported": chyt_exc.CHYTMultipleDynamicTablesNotSupported,
        r"Subquery exceeds data weight limit": chyt_exc.CHYTSubqueryWeightLimitExceeded,
    }
    chyt_fallback_exc_cls = chyt_exc.CHYTQueryError
    clique_expr_exc = {
        r'User "(?P<user>[-0-9a-zA-Z]+)" has no access to clique (?P<clique>\*\S+)': chyt_exc.CHYTCliqueAccessDenied,
        r"Clique (?P<clique>\*\S+) is not running": chyt_exc.CHYTCliqueIsNotRunning,
        r"Clique (?P<clique>\*\S+) is suspended": chyt_exc.CHYTCliqueIsSuspended,
        r"Invalid clique specification": chyt_exc.CHYTCliqueNotExists,
        r"Authentication failed": chyt_exc.CHYTAuthError,
    }

    @classmethod
    def parse_clique_message(cls, err_msg: str) -> Optional[tuple[Type[exc.DatabaseQueryError], dict[str, str]]]:
        for err_re, chyt_exc_cls in cls.clique_expr_exc.items():
            match = re.search(err_re, err_msg)
            if match:
                LOGGER.info("Recognized as CHYT error without code")
                return chyt_exc_cls, match.groupdict()
        return None

    @classmethod
    def get_exc_class_by_parsed_message(
        cls, msg: ParsedErrorMsg
    ) -> Optional[tuple[Type[exc.DatabaseQueryError], dict[str, str]]]:
        if msg.code == 1001:
            LOGGER.info("Recognized as CHYT error code")
            for err_re, chyt_exc_cls in cls.chyt_expr_exc.items():
                match = re.search(err_re, msg.full_msg)
                if match:
                    return chyt_exc_cls, match.groupdict()
            return cls.chyt_fallback_exc_cls, {}

        return super().get_exc_class_by_parsed_message(msg)

    @classmethod
    def get_exc_class(cls, err_msg: str) -> Optional[tuple[Type[exc.DatabaseQueryError], dict[str, str]]]:
        # Clique exception
        not_ch_exc = cls.parse_clique_message(err_msg)
        if not_ch_exc:
            return not_ch_exc
        return super().get_exc_class(err_msg)

    @classmethod
    def get_tracing_sample_flag_override(cls, rci: DBAdapterScopedRCI) -> Optional[bool]:
        # We should set sample flag only if x_dl_debug_mode is presented and is True
        return rci.x_dl_debug_mode is True

lib/dl_connector_chyt/dl_connector_chyt/core/utils.py
datalens-tech/datalens-backend
from dl_db_testing.connectors.base.connector import DbTestingConnector

from dl_connector_chyt.db_testing.engine_wrapper import CHYTEngineWrapperBase


class CHYTDbTestingConnector(DbTestingConnector):
    engine_wrapper_classes = (CHYTEngineWrapperBase,)

lib/dl_connector_chyt/dl_connector_chyt/db_testing/connector.py
datalens-tech/datalens-backend
from typing import Sequence

import sqlalchemy as sa

from dl_db_testing.database.engine_wrapper import EngineWrapperBase


class CHYTEngineWrapperBase(EngineWrapperBase):
    URL_PREFIX = "bi_chyt"

    def get_conn_credentials(self, full: bool = False) -> dict:
        return {}  # FIXME

    def drop_table(self, db_name: str, table: sa.Table) -> None:
        raise NotImplementedError

    def create_table(self, table: sa.Table) -> None:
        raise NotImplementedError

    def insert_into_table(self, table: sa.Table, data: Sequence[dict]) -> None:
        raise NotImplementedError

    def test(self) -> bool:
        return True

lib/dl_connector_chyt/dl_connector_chyt/db_testing/engine_wrapper.py
datalens-tech/datalens-backend
from dl_constants.exc import (
    DEFAULT_ERR_CODE_API_PREFIX,
    GLOBAL_ERR_PREFIX,
)

from dl_connector_chyt.core import exc


def exc_cls_to_code(exc_cls: exc.CHYTQueryError | exc.CHYTCliqueError) -> str:
    return ".".join([GLOBAL_ERR_PREFIX, DEFAULT_ERR_CODE_API_PREFIX] + exc_cls.err_code)


STABILIZED_CODES = (
    (exc.CHYTInvalidSortedJoin, "ERR.DS_API.DB.CHYT.INVALID_SORTED_JOIN"),
    (exc.CHYTInvalidSortedJoinMoreThanOneTable, "ERR.DS_API.DB.CHYT.INVALID_SORTED_JOIN.MORE_THAN_ONE_TABLE"),
    (exc.CHYTInvalidSortedJoinNotAKeyColumn, "ERR.DS_API.DB.CHYT.INVALID_SORTED_JOIN.NOT_A_KEY_COLUMN"),
    (exc.CHYTInvalidSortedJoinNotKeyPrefixColumn, "ERR.DS_API.DB.CHYT.INVALID_SORTED_JOIN.NOT_KEY_PREFIX_COLUMN"),
    (exc.CHYTMultipleDynamicTablesNotSupported, "ERR.DS_API.DB.CHYT.MULTI_DYN_NOT_SUPPORTED"),
    (exc.CHYTQueryError, "ERR.DS_API.DB.CHYT"),
    (exc.CHYTSubqueryWeightLimitExceeded, "ERR.DS_API.DB.CHYT.SUBQ_WEIGHT_LIMIT_EXCEEDED"),
    (exc.CHYTTableAccessDenied, "ERR.DS_API.DB.CHYT.TABLE_ACCESS_DENIED"),
    (exc.CHYTCliqueIsNotRunning, "ERR.DS_API.DB.CHYT.CLIQUE.NOT_RUNNING"),
    (exc.CHYTCliqueIsSuspended, "ERR.DS_API.DB.CHYT.CLIQUE.SUSPENDED"),
    (exc.CHYTCliqueNotExists, "ERR.DS_API.DB.CHYT.CLIQUE.INVALID_SPECIFICATION"),
    (exc.CHYTCliqueAccessDenied, "ERR.DS_API.DB.CHYT.CLIQUE.ACCESS_DENIED"),
    (exc.CHYTTableHasNoSchema, "ERR.DS_API.DB.CHYT.TABLE_HAS_NO_SCHEMA"),
    (exc.CHYTAuthError, "ERR.DS_API.DB.CHYT.AUTH_FAILED"),
)


def test_exc_code_match():
    for exc_cls, expected_code in STABILIZED_CODES:
        actual_code = exc_cls_to_code(exc_cls)
        assert actual_code == expected_code, exc_cls

lib/dl_connector_chyt/dl_connector_chyt_tests/unit/test_exc_codes.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_connector.api_schema.connection_base import ConnectionMetaMixin
from dl_api_connector.api_schema.connection_mixins import (
    DataExportForbiddenMixin,
    RawSQLLevelMixin,
)
from dl_api_connector.api_schema.connection_sql import ClassicSQLConnectionSchema
from dl_api_connector.api_schema.extras import FieldExtra
import dl_core.marshmallow as core_ma_fields

from dl_connector_clickhouse.core.clickhouse.us_connection import ConnectionClickhouse


class ClickHouseConnectionSchema(
    ConnectionMetaMixin, RawSQLLevelMixin, DataExportForbiddenMixin, ClassicSQLConnectionSchema
):
    TARGET_CLS = ConnectionClickhouse
    ALLOW_MULTI_HOST = True

    secure = core_ma_fields.OnOffField(attribute="data.secure", bi_extra=FieldExtra(editable=True))
    ssl_ca = core_ma_fields.Base64StringField(
        attribute="data.ssl_ca",
        bi_extra=FieldExtra(editable=True),
        required=False,
        allow_none=True,
        load_default=None,
        load_only=True,
    )

lib/dl_connector_clickhouse/dl_connector_clickhouse/api/api_schema/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    Sequence,
)

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormActionApiSchema,
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
)
from dl_api_connector.form_config.models.common import CommonFieldName
import dl_api_connector.form_config.models.rows as C
from dl_api_connector.form_config.models.rows.base import FormRow
from dl_api_connector.form_config.models.shortcuts.rows import RowConstructor
from dl_configs.connectors_settings import ConnectorSettingsBase

from dl_connector_clickhouse.api.connection_info import ClickHouseConnectionInfoProvider
from dl_connector_clickhouse.api.i18n.localizer import Translatable


class ClickHouseConnectionFormFactory(ConnectionFormFactory):
    @staticmethod
    def get_common_api_schema_items() -> list[FormFieldApiSchema]:
        return [
            FormFieldApiSchema(name=CommonFieldName.host, required=True),
            FormFieldApiSchema(name=CommonFieldName.port, required=True),
            FormFieldApiSchema(name=CommonFieldName.username),
            FormFieldApiSchema(name=CommonFieldName.password),
            FormFieldApiSchema(name=CommonFieldName.secure),
            FormFieldApiSchema(name=CommonFieldName.ssl_ca),
            FormFieldApiSchema(name=CommonFieldName.data_export_forbidden),
        ]

    def _get_base_edit_api_schema(self) -> FormActionApiSchema:
        return FormActionApiSchema(
            items=[
                *self.get_common_api_schema_items(),
                FormFieldApiSchema(name=CommonFieldName.cache_ttl_sec, nullable=True),
                FormFieldApiSchema(name=CommonFieldName.raw_sql_level),
            ],
        )

    def _get_base_create_api_schema(self, edit_api_schema: FormActionApiSchema) -> FormActionApiSchema:
        return FormActionApiSchema(
            items=[
                *edit_api_schema.items,
                *self._get_top_level_create_api_schema_items(),
            ],
            conditions=edit_api_schema.conditions.copy(),
        )

    def _get_base_check_api_schema(self) -> FormActionApiSchema:
        return FormActionApiSchema(
            items=[
                *self.get_common_api_schema_items(),
                *self._get_top_level_check_api_schema_items(),
            ]
        )

    def _get_base_form_config(
        self,
        host_section: Sequence[FormRow],
        username_section: Sequence[FormRow],
        create_api_schema: FormActionApiSchema,
        edit_api_schema: FormActionApiSchema,
        check_api_schema: FormActionApiSchema,
        rc: RowConstructor,
    ) -> ConnectionForm:
        return ConnectionForm(
            title=ClickHouseConnectionInfoProvider.get_title(self._localizer),
            rows=self._filter_nulls(
                [
                    *host_section,
                    rc.port_row(label_text=Translatable("field_click-house-port"), default_value="8443"),
                    *username_section,
                    rc.password_row(mode=self.mode),
                    C.CacheTTLRow(name=CommonFieldName.cache_ttl_sec),
                    rc.raw_sql_level_row(),
                    rc.collapse_advanced_settings_row(),
                    *rc.ssl_rows(
                        enabled_name=CommonFieldName.secure,
                        enabled_help_text=self._localizer.translate(
                            Translatable("label_clickhouse-ssl-enabled-tooltip")
                        ),
                        enabled_default_value=True,
                    ),
                    rc.data_export_forbidden_row(),
                ]
            ),
            api_schema=FormApiSchema(
                create=create_api_schema if self.mode == ConnectionFormMode.create else None,
                edit=edit_api_schema if self.mode == ConnectionFormMode.edit else None,
                check=check_api_schema,
            ),
        )

    def get_form_config(
        self,
        connector_settings: Optional[ConnectorSettingsBase],
        tenant: Optional[TenantDef],
    ) -> ConnectionForm:
        rc = RowConstructor(localizer=self._localizer)

        host_section: list[FormRow] = [rc.host_row()]
        username_section: list[FormRow] = [rc.username_row()]

        edit_api_schema = self._get_base_edit_api_schema()
        create_api_schema = self._get_base_create_api_schema(edit_api_schema)
        check_api_schema = self._get_base_check_api_schema()

        return self._get_base_form_config(
            host_section=host_section,
            username_section=username_section,
            create_api_schema=create_api_schema,
            edit_api_schema=edit_api_schema,
            check_api_schema=check_api_schema,
            rc=rc,
        )

lib/dl_connector_clickhouse/dl_connector_clickhouse/api/connection_form/form_config.py
datalens-tech/datalens-backend
from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_clickhouse.api.i18n.localizer import Translatable


class ClickHouseConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-clickhouse")

lib/dl_connector_clickhouse/dl_connector_clickhouse/api/connection_info.py
datalens-tech/datalens-backend
from dl_api_connector.api_schema.source_base import (
    SQLDataSourceSchema,
    SQLDataSourceTemplateSchema,
    SubselectDataSourceSchema,
    SubselectDataSourceTemplateSchema,
)
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)
from dl_api_lib.query.registry import MQMFactorySettingItem
from dl_constants.enums import QueryProcessingMode
from dl_query_processing.multi_query.factory import NoCompengMultiQueryMutatorFactory

from dl_connector_clickhouse.api.api_schema.connection import ClickHouseConnectionSchema
from dl_connector_clickhouse.api.connection_form.form_config import ClickHouseConnectionFormFactory
from dl_connector_clickhouse.api.connection_info import ClickHouseConnectionInfoProvider
from dl_connector_clickhouse.api.i18n.localizer import CONFIGS
from dl_connector_clickhouse.core.clickhouse.connector import (
    ClickHouseCoreConnectionDefinition,
    ClickHouseCoreConnector,
    ClickHouseSubselectCoreSourceDefinition,
    ClickHouseTableCoreSourceDefinition,
)
from dl_connector_clickhouse.formula.constants import (
    DIALECT_NAME_CLICKHOUSE,
    ClickHouseDialect,
)


class ClickHouseApiTableSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = ClickHouseTableCoreSourceDefinition
    api_schema_cls = SQLDataSourceSchema
    template_api_schema_cls = SQLDataSourceTemplateSchema


class ClickHouseApiSubselectSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = ClickHouseSubselectCoreSourceDefinition
    api_schema_cls = SubselectDataSourceSchema
    template_api_schema_cls = SubselectDataSourceTemplateSchema


class ClickHouseApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = ClickHouseCoreConnectionDefinition
    api_generic_schema_cls = ClickHouseConnectionSchema
    info_provider_cls = ClickHouseConnectionInfoProvider
    form_factory_cls = ClickHouseConnectionFormFactory


class ClickHouseApiConnector(ApiConnector):
    core_connector_cls = ClickHouseCoreConnector
    connection_definitions = (ClickHouseApiConnectionDefinition,)
    source_definitions = (
        ClickHouseApiTableSourceDefinition,
        ClickHouseApiSubselectSourceDefinition,
    )
    formula_dialect_name = DIALECT_NAME_CLICKHOUSE
    translation_configs = frozenset(CONFIGS)
    multi_query_mutation_factories = ApiConnector.multi_query_mutation_factories + (
        MQMFactorySettingItem(
            query_proc_mode=QueryProcessingMode.native_wf,
            dialects=ClickHouseDialect.and_above(ClickHouseDialect.CLICKHOUSE_22_10).to_list(),
            factory_cls=NoCompengMultiQueryMutatorFactory,
        ),
    )

lib/dl_connector_clickhouse/dl_connector_clickhouse/api/connector.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_clickhouse as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_clickhouse/dl_connector_clickhouse/api/i18n/localizer.py
datalens-tech/datalens-backend
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreSourceDefinition,
)
from dl_core.connectors.sql_base.connector import SQLSubselectCoreSourceDefinitionBase
from dl_core.data_source_spec.sql import StandardSQLDataSourceSpec
from dl_core.us_manager.storage_schemas.data_source_spec_base import SQLDataSourceSpecStorageSchema

from dl_connector_clickhouse.core.clickhouse.constants import (
    SOURCE_TYPE_CH_SUBSELECT,
    SOURCE_TYPE_CH_TABLE,
)
from dl_connector_clickhouse.core.clickhouse.data_source import (
    ClickHouseDataSource,
    ClickHouseSubselectDataSource,
)
from dl_connector_clickhouse.core.clickhouse.data_source_migration import ClickHouseDataSourceMigrator
from dl_connector_clickhouse.core.clickhouse.storage_schemas.connection import ConnectionClickhouseDataStorageSchema
from dl_connector_clickhouse.core.clickhouse.us_connection import ConnectionClickhouse
from dl_connector_clickhouse.core.clickhouse_base.connection_executors import (
    AsyncClickHouseConnExecutor,
    ClickHouseConnExecutor,
)
from dl_connector_clickhouse.core.clickhouse_base.connector import ClickHouseCoreConnectorBase
from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
from dl_connector_clickhouse.core.clickhouse_base.type_transformer import ClickHouseTypeTransformer


class ClickHouseCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_CLICKHOUSE
    connection_cls = ConnectionClickhouse
    us_storage_schema_cls = ConnectionClickhouseDataStorageSchema
    type_transformer_cls = ClickHouseTypeTransformer
    sync_conn_executor_cls = ClickHouseConnExecutor
    async_conn_executor_cls = AsyncClickHouseConnExecutor
    dialect_string = "bi_clickhouse"
    data_source_migrator_cls = ClickHouseDataSourceMigrator


class ClickHouseTableCoreSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_CH_TABLE
    source_cls = ClickHouseDataSource
    source_spec_cls = StandardSQLDataSourceSpec
    us_storage_schema_cls = SQLDataSourceSpecStorageSchema


class ClickHouseSubselectCoreSourceDefinition(SQLSubselectCoreSourceDefinitionBase):
    source_type = SOURCE_TYPE_CH_SUBSELECT
    source_cls = ClickHouseSubselectDataSource


class ClickHouseCoreConnector(ClickHouseCoreConnectorBase):
    connection_definitions = (ClickHouseCoreConnectionDefinition,)
    source_definitions = (
        ClickHouseTableCoreSourceDefinition,
        ClickHouseSubselectCoreSourceDefinition,
    )

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import DataSourceType


SOURCE_TYPE_CH_TABLE = DataSourceType.declare("CH_TABLE")
SOURCE_TYPE_CH_SUBSELECT = DataSourceType.declare("CH_SUBSELECT")

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.enums import DataSourceType

from dl_connector_clickhouse.core.clickhouse.constants import (
    SOURCE_TYPE_CH_SUBSELECT,
    SOURCE_TYPE_CH_TABLE,
)
from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
from dl_connector_clickhouse.core.clickhouse_base.data_source import (
    ActualClickHouseBaseMixin,
    ClickHouseDataSourceBase,
    CommonClickHouseSubselectDataSource,
)


class ClickHouseDataSource(ClickHouseDataSourceBase):
    conn_type = CONNECTION_TYPE_CLICKHOUSE

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type in {
            SOURCE_TYPE_CH_TABLE,
            SOURCE_TYPE_CH_SUBSELECT,
        }


class ClickHouseSubselectDataSource(ActualClickHouseBaseMixin, CommonClickHouseSubselectDataSource):  # type: ignore  # TODO: fix
    """Clickhouse subselect"""

    conn_type = CONNECTION_TYPE_CLICKHOUSE

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type in {
            SOURCE_TYPE_CH_TABLE,
            SOURCE_TYPE_CH_SUBSELECT,
        }

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/data_source.py
datalens-tech/datalens-backend
from dl_core.connectors.sql_base.data_source_migration import DefaultSQLDataSourceMigrator

from dl_connector_clickhouse.core.clickhouse.constants import (
    SOURCE_TYPE_CH_SUBSELECT,
    SOURCE_TYPE_CH_TABLE,
)


class ClickHouseDataSourceMigrator(DefaultSQLDataSourceMigrator):
    table_source_type = SOURCE_TYPE_CH_TABLE
    subselect_source_type = SOURCE_TYPE_CH_SUBSELECT
    with_db_name = True

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/data_source_migration.py
datalens-tech/datalens-backend
from dl_connector_clickhouse.core.clickhouse.us_connection import ConnectionClickhouse
from dl_connector_clickhouse.core.clickhouse_base.storage_schemas.connection import (
    ConnectionClickHouseBaseDataStorageSchema,
)


class ConnectionClickhouseDataStorageSchema(ConnectionClickHouseBaseDataStorageSchema[ConnectionClickhouse.DataModel]):
    TARGET_CLS = ConnectionClickhouse.DataModel

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/storage_schemas/connection.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Optional,
)
import uuid

from dl_constants.enums import RawSQLLevel
from dl_core.us_manager.us_manager_sync import SyncUSManager

from dl_connector_clickhouse.core.clickhouse.us_connection import ConnectionClickhouse
from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE


def make_clickhouse_saved_connection(
    sync_usm: SyncUSManager,
    db_name: Optional[str],
    host: str,
    port: Optional[int],
    username: Optional[str],
    password: Optional[str],
    raw_sql_level: RawSQLLevel = RawSQLLevel.off,
    secure: bool = False,
    ssl_ca: Optional[str] = None,
    **kwargs: Any,
) -> ConnectionClickhouse:
    conn_name = "ch test conn {}".format(uuid.uuid4())
    conn = ConnectionClickhouse.create_from_dict(
        data_dict=ConnectionClickhouse.DataModel(
            db_name=db_name,
            host=host,
            port=port,
            username=username,
            password=password,
            raw_sql_level=raw_sql_level,
            secure=secure,
            ssl_ca=ssl_ca,
        ),
        ds_key=conn_name,
        type_=CONNECTION_TYPE_CLICKHOUSE.name,
        us_manager=sync_usm,
        **kwargs,
    )
    sync_usm.save(conn)
    return conn

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/testing/connection.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Type,
)

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.models.connection_target_dto_base import BaseSQLConnTargetDTO
from dl_core_testing.executors import ExecutorFactoryBase

from dl_connector_clickhouse.core.clickhouse_base.adapters import ClickHouseAdapter
from dl_connector_clickhouse.core.clickhouse_base.target_dto import ClickHouseConnTargetDTO


class ClickHouseExecutorFactory(ExecutorFactoryBase):
    def get_dto_class(self) -> Type[BaseSQLConnTargetDTO]:
        return ClickHouseConnTargetDTO

    def get_dba_class(self) -> Type[CommonBaseDirectAdapter]:
        return ClickHouseAdapter

    def get_dto_kwargs(self) -> dict[str, Any]:
        return dict(
            **super().get_dto_kwargs(),
            protocol="http",
            cluster_name=None,
            endpoint=None,
            connect_timeout=None,
            max_execution_time=None,
            total_timeout=None,
            insert_quorum=None,
            insert_quorum_timeout=None,
            disable_value_processing=False,
        )

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/testing/exec_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import ClassVar

from dl_core.us_connection_base import DataSourceTemplate
from dl_i18n.localizer_base import Localizer

from dl_connector_clickhouse.core.clickhouse.constants import (
    SOURCE_TYPE_CH_SUBSELECT,
    SOURCE_TYPE_CH_TABLE,
)
from dl_connector_clickhouse.core.clickhouse_base.us_connection import ConnectionClickhouseBase


class ConnectionClickhouse(ConnectionClickhouseBase):
    """
    User's ClickHouse database.
    Should not be used for internal clickhouses.
    """

    source_type = SOURCE_TYPE_CH_TABLE
    allowed_source_types = frozenset((SOURCE_TYPE_CH_TABLE, SOURCE_TYPE_CH_SUBSELECT))
    allow_dashsql: ClassVar[bool] = True
    allow_cache: ClassVar[bool] = True
    is_always_user_source: ClassVar[bool] = False  # TODO: should be `True`, but need some cleanup for that.

    def get_data_source_template_templates(self, localizer: Localizer) -> list[DataSourceTemplate]:
        return self._make_subselect_templates(source_type=SOURCE_TYPE_CH_SUBSELECT, localizer=localizer)

    @property
    def allow_public_usage(self) -> bool:
        return True

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse/us_connection.py
datalens-tech/datalens-backend
from __future__ import annotations

import contextlib
import json
import logging
import ssl
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncGenerator,
    Callable,
    ClassVar,
    ContextManager,
    Generator,
    Optional,
    Sequence,
    Type,
    TypeVar,
)

from aiohttp import BasicAuth
from aiohttp.client import (
    ClientResponse,
    ClientTimeout,
)
import attr
from clickhouse_sqlalchemy import exceptions as ch_exc
from clickhouse_sqlalchemy import types as ch_types
from clickhouse_sqlalchemy.drivers.http.transport import _get_type  # noqa
from clickhouse_sqlalchemy.parsers.jsoncompact import JSONCompactChunksParser
import requests
import sqlalchemy as sa
from sqlalchemy.sql.type_api import TypeEngine

from dl_app_tools.profiling_base import generic_profiler_async
from dl_constants.enums import IndexKind
from dl_core import exc
from dl_core.connection_executors.adapters.adapters_base_sa_classic import (
    BaseClassicAdapter,
    ClassicSQLConnLineConstructor,
)
from dl_core.connection_executors.adapters.async_adapters_aiohttp import AiohttpDBAdapter
from dl_core.connection_executors.adapters.async_adapters_base import AsyncRawExecutionResult
from dl_core.connection_executors.adapters.common_base import get_dialect_string
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    RawColumnInfo,
    RawIndexInfo,
    RawSchemaInfo,
)
from dl_core.connection_models import (
    SATextTableDefinition,
    TableDefinition,
    TableIdent,
)
from dl_core.connectors.base.error_transformer import DBExcKWArgs
from dl_core.connectors.ssl_common.adapter import BaseSSLCertAdapter
from dl_core.db.native_type import (
    ClickHouseDateTime64NativeType,
    ClickHouseDateTime64WithTZNativeType,
    ClickHouseDateTimeWithTZNativeType,
    ClickHouseNativeType,
    GenericNativeType,
    norm_native_type,
)
from dl_core.utils import make_url

from dl_connector_clickhouse.core.clickhouse_base.ch_commons import (
    ClickHouseBaseUtils,
    ClickHouseUtils,
    get_ch_settings,
)
from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
from dl_connector_clickhouse.core.clickhouse_base.exc import CHRowTooLarge


if TYPE_CHECKING:
    from dl_constants.types import TBIChunksGen
    from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
    from dl_core.connection_models.common_models import (
        DBIdent,
        SchemaIdent,
    )

    from dl_connector_clickhouse.core.clickhouse_base.target_dto import (  # noqa: F401
        BaseClickHouseConnTargetDTO,
        ClickHouseConnTargetDTO,
    )

LOGGER = logging.getLogger(__name__)

_DBA_TV = TypeVar("_DBA_TV", bound="BaseAsyncClickHouseAdapter")
_TARGET_DTO_TV = TypeVar("_TARGET_DTO_TV", bound="BaseClickHouseConnTargetDTO")


@attr.s()
class BaseClickHouseConnLineConstructor(ClassicSQLConnLineConstructor[_TARGET_DTO_TV]):
    _converted_headers: dict[str, str] = attr.ib()

    def _get_dsn_query_params(self) -> dict:
        # TODO FIX: Move to utils
        connect_timeout = self._target_dto.connect_timeout if self._target_dto.connect_timeout is not None else 1
        # TODO FIX: Move to utils
        # TODO CONSIDER: May be another constant?
        total_timeout = self._target_dto.total_timeout if self._target_dto.total_timeout is not None else 290

        return {
            **dict(
                protocol=self._target_dto.protocol,
                endpoint=self._target_dto.endpoint,
                timeout=total_timeout,
                connect_timeout=connect_timeout,
                stream=True,
            ),
            **self._converted_headers,
        }


class BaseClickHouseAdapter(BaseClassicAdapter["BaseClickHouseConnTargetDTO"], BaseSSLCertAdapter):
    allow_sa_text_as_columns_source = True
    ch_utils: ClassVar[Type[ClickHouseBaseUtils]] = ClickHouseBaseUtils
    conn_line_constructor_type: ClassVar[Type[BaseClickHouseConnLineConstructor]] = BaseClickHouseConnLineConstructor

    _dt_with_system_tz = True

    def _get_dsn_params_from_headers(self) -> dict[str, str]:
        return self._convert_headers_to_dsn_params(self.ch_utils.get_context_headers(self._req_ctx_info))

    def get_conn_line(self, db_name: Optional[str] = None, params: Optional[dict[str, Any]] = None) -> str:
        return self.conn_line_constructor_type(
            dsn_template=self.dsn_template,
            dialect_name=get_dialect_string(self.conn_type),
            target_dto=self._target_dto,
            converted_headers=self._get_dsn_params_from_headers(),
        ).make_conn_line(db_name=db_name, params=params)

    # TODO FIX: Move to utils
    def get_ch_settings(self) -> dict:
        return get_ch_settings(
            max_execution_time=self._target_dto.max_execution_time,
            read_only_level=None,
            insert_quorum=self._target_dto.insert_quorum,
            insert_quorum_timeout=self._target_dto.insert_quorum_timeout,
            output_format_json_quote_denormals=1,
        )

    def get_connect_args(self) -> dict:
        args = {
            **super().get_connect_args(),
            # Hard-code server version
            #  to prevent querying in on each engine instantiation against CH server (default dialect behaviour).
            #  Dialect at this moment does not depends on CH version, so we can hard code it.
            "server_version": "19.16.2.2",
            "ch_settings": self.get_ch_settings(),
            "format": "JSONCompact",
        }

        if self._target_dto.secure:
            args["verify"] = self.get_ssl_cert_path(self._target_dto.ssl_ca)

        return args

    @contextlib.contextmanager
    def execution_context(self) -> Generator[None, None, None]:
        contexts: list[ContextManager[None]] = [super().execution_context()]

        if self._target_dto.ssl_ca:
            contexts.append(self.ssl_cert_context(self._target_dto.ssl_ca))

        with contextlib.ExitStack() as stack:
            for context in contexts:
                stack.enter_context(context)
            try:
                yield
            finally:
                stack.close()

    @staticmethod
    def _convert_headers_to_dsn_params(headers: dict[str, str]) -> dict[str, str]:
        return {f"header__{h_name}": h_val for h_name, h_val in headers.items()}

    EXTRA_EXC_CLS = (ch_exc.DatabaseException,)

    @classmethod
    def make_exc(  # TODO:  Move to ErrorTransformer
        cls, wrapper_exc: Exception, orig_exc: Optional[Exception], debug_compiled_query: Optional[str]
    ) -> tuple[Type[exc.DatabaseQueryError], DBExcKWArgs]:
        exc_cls, kw = super().make_exc(wrapper_exc, orig_exc, debug_compiled_query)

        ch_exc_cls: Optional[Type[exc.DatabaseQueryError]] = None
        if isinstance(wrapper_exc, ch_exc.DatabaseException):
            # Special case for ClickHouse
            # try to differentiate errors by error code
            db_msg = kw["db_message"]
            if db_msg:
                try:
                    # TODO: Temporary, will be fixed for sync adapter later
                    ch_exc_cls_with_params = cls.ch_utils.get_exc_class(db_msg)
                    if ch_exc_cls_with_params:
                        ch_exc_cls, _ = ch_exc_cls_with_params
                        ch_exc_cls = ch_exc_cls or exc_cls
                except Exception:  # noqa
                    LOGGER.info("Can not parse ClickHouse error message")

        if ch_exc_cls:
            exc_cls = ch_exc_cls

        elif isinstance(orig_exc, requests.exceptions.ReadTimeout):
            LOGGER.info("ClickHouse timed out")
            exc_cls = exc.SourceTimeout

        elif isinstance(orig_exc, requests.exceptions.ConnectionError):
            exc_cls = exc.SourceConnectError

        return exc_cls, kw

    # TODO CONSIDER: Do we really need to overwrite native type???
    def normalize_sa_col_type(self, sa_col_type: TypeEngine) -> TypeEngine:
        if isinstance(sa_col_type, ch_types.Decimal):
            return sa.Float()
        elif isinstance(sa_col_type, (ch_types.Enum8, ch_types.Enum16)):
            return sa.String()

        return super().normalize_sa_col_type(sa_col_type)

    def _test(self) -> None:
        self.execute(DBAdapterQuery("select 1"))

    def _get_subselect_table_info(self, subquery: SATextTableDefinition) -> RawSchemaInfo:
        # `describe table {table_func}/{subselect}` also works in CH.
        return RawSchemaInfo(columns=self._get_raw_columns_info(subquery))

    def _get_raw_columns_info(self, table_def: TableDefinition) -> tuple[RawColumnInfo, ...]:
        table_columns = self._get_sa_table_columns(table_def)
        db_name = table_def.db_name if isinstance(table_def, TableIdent) else None
        db_engine = self.get_db_engine(db_name)
        if self._dt_with_system_tz:
            system_tz = db_engine.scalar(sa.text("select timezone()"))

        def make_native_type(col):  # type: ignore  # TODO: fix
            col_type = col["type"]
            col_type = self.normalize_sa_col_type(col_type)

            nullable = col.get("nullable", True)
            lowcardinality = col.get("lowcardinality", False)
            is_array = False

            if isinstance(col_type, ch_types.Array):
                is_array = True
                col_type = col_type.item_type
            if isinstance(col_type, ch_types.LowCardinality):
                lowcardinality = True
                col_type = col_type.nested_type
            if isinstance(col_type, ch_types.Nullable):
                nullable = True
                col_type = col_type.nested_type

            explicit_timezone = True
            if self._dt_with_system_tz:
                if isinstance(col_type, ch_types.DateTime) and not isinstance(col_type, ch_types.DateTimeWithTZ):
                    explicit_timezone = False
                    col_type = ch_types.DateTimeWithTZ(system_tz)
                elif isinstance(col_type, ch_types.DateTime64) and not isinstance(col_type, ch_types.DateTime64WithTZ):
                    explicit_timezone = False
                    col_type = ch_types.DateTime64WithTZ(col_type.precision, system_tz)

            conn_type = self.conn_type
            if is_array:
                name = norm_native_type(ch_types.Array(col_type))
            else:
                name = norm_native_type(col_type)

            if isinstance(col_type, ch_types.DateTimeWithTZ):
                return ClickHouseDateTimeWithTZNativeType(
                    conn_type=conn_type,
                    name=name,  # type: ignore  # TODO: fix
                    nullable=nullable,
                    lowcardinality=lowcardinality,
                    timezone_name=col_type.tz,
                    explicit_timezone=explicit_timezone,
                )
            if isinstance(col_type, ch_types.DateTime64WithTZ):
                return ClickHouseDateTime64WithTZNativeType(
                    conn_type=conn_type,
                    name=name,  # type: ignore  # TODO: fix
                    nullable=nullable,
                    lowcardinality=lowcardinality,
                    precision=col_type.precision,
                    timezone_name=col_type.tz,
                    explicit_timezone=explicit_timezone,
                )
            if isinstance(col_type, ch_types.DateTime64):
                return ClickHouseDateTime64NativeType(
                    conn_type=conn_type,
                    name=name,  # type: ignore  # TODO: fix
                    nullable=nullable,
                    lowcardinality=lowcardinality,
                    precision=col_type.precision,
                )
            return ClickHouseNativeType(
                conn_type=conn_type,
                name=name,  # type: ignore  # TODO: fix
                nullable=nullable,
                lowcardinality=lowcardinality,
            )

        return tuple(
            RawColumnInfo(
                name=column["name"],
                title=column.get("title"),
                nullable=True,
                # column.get('nullable'),  # forced `True` for table-creation.  # to be deprecated entirely.
                native_type=make_native_type(column),
            )
            for column in table_columns
        )


class ClickHouseAdapter(BaseClickHouseAdapter):
    conn_type = CONNECTION_TYPE_CLICKHOUSE
    warn_on_default_db_name_override = False
    ch_utils = ClickHouseUtils

    def get_default_db_name(self) -> Optional[str]:
        return self._target_dto.db_name or "system"

    def _get_table_indexes(self, table_ident: TableIdent) -> tuple[RawIndexInfo, ...]:
        common_indexes = super()._get_table_indexes(table_ident)
        result = self.execute(
            DBAdapterQuery(
                query=sa.select(
                    [sa.column("sorting_key")],
                )
                .select_from(sa.table("tables"))
                .where(
                    sa.and_(
                        sa.column("name") == table_ident.table_name,
                        sa.column("database") == self.get_db_name_for_query(table_ident.db_name),
                    )
                ),
                db_name="system",
            )
        ).get_all()
        assert len(result) == 1, "Number of row for table name "

        ch_specific_idx_list = []
        sorting_key_str = result[0][0]

        if sorting_key_str and sorting_key_str.strip():
            sorting_key = tuple(col_name.strip() for col_name in sorting_key_str.split(","))
            ch_specific_idx_list.append(
                RawIndexInfo(
                    columns=sorting_key,
                    unique=False,
                    kind=IndexKind.table_sorting,
                )
            )

        return common_indexes + tuple(ch_specific_idx_list)


@attr.s(kw_only=True)
class BaseAsyncClickHouseAdapter(AiohttpDBAdapter):
    ch_utils: ClassVar[Type[ClickHouseBaseUtils]] = ClickHouseBaseUtils

    _target_dto: BaseClickHouseConnTargetDTO = attr.ib()

    _url: str = attr.ib(init=False, default=None)

    def __attrs_post_init__(self) -> None:
        super().__attrs_post_init__()
        self._url = make_url(
            protocol=self._target_dto.protocol,
            host=self._target_dto.host,
            port=self._target_dto.port,
            path=self._target_dto.endpoint,
        )

    def get_session_timeout(self) -> Optional[ClientTimeout]:  # type: ignore  # TODO: fix
        return ClientTimeout(connect=self._target_dto.connect_timeout, total=self._target_dto.total_timeout)

    def get_session_auth(self) -> Optional[BasicAuth]:
        return BasicAuth(
            login=self._target_dto.username,
            password=self._target_dto.password,
            encoding="utf-8",
        )

    def get_session_headers(self) -> dict[str, str]:
        return {
            **super().get_session_headers(),
            **self.ch_utils.get_context_headers(self._req_ctx_info),
        }

    def get_request_params(self, dba_q: DBAdapterQuery) -> dict[str, str]:
        read_only_level = None if dba_q.trusted_query else 2
        return dict(
            # TODO FIX: Move to utils
            database=dba_q.db_name or self._target_dto.db_name or "system",
            **get_ch_settings(
                read_only_level=read_only_level,
                max_execution_time=self._target_dto.max_execution_time,
                # doesn't matter until materializer uses async ch adapter
                insert_quorum=self._target_dto.insert_quorum,
                insert_quorum_timeout=self._target_dto.insert_quorum_timeout,
            ),
        )

    def _get_current_tracing_headers(self) -> dict[str, str]:
        return {}

    def _get_ssl_context(self) -> Optional[ssl.SSLContext]:
        return None

    # TODO FIX: Add logging from dl_core.connection_executors.adapters.sa_utils.CursorLogger
    async def _make_query(self, dba_q: DBAdapterQuery, mirroring_mode: bool = False) -> ClientResponse:
        query_str: str
        if isinstance(dba_q.query, str):
            _query_for_check = dba_q.query.lower()
            if not dba_q.trusted_query:
                assert _query_for_check.startswith(("select ", "select\n", "explain ", "explain\n"))
            query_str = dba_q.query
        else:
            query_str = dba_q.query.compile(dialect=self.get_dialect(), compile_kwargs={"literal_binds": True}).string
            # SA generates query for DBAPI, so mod is represented as `%%` so next hack is needed
            query_str = query_str % ()

        final_query = query_str + "\n" + "FORMAT JSONCompact"
        ch_params = self.get_request_params(dba_q)

        query_for_log = (
            dba_q.debug_compiled_query + "\n" + "FORMAT JSONCompact"
            if isinstance(dba_q.debug_compiled_query, str)
            else final_query
        )

        max_log_len = 1024 * 4
        LOGGER.info(
            "Going to send query to CH%s: %s%s",
            " (mirroring_mode=True)" if mirroring_mode else "",
            query_for_log[:max_log_len],
            "..." if len(query_for_log) > max_log_len else "",
            extra={
                "ch_query": query_for_log,
                "ch_url": self._url,
                "ch_params": json.dumps(ch_params),
                "event_code": "async_dba_ch_query_ready",
            },
        )

        tracing_headers = self._get_current_tracing_headers()

        resp = await self._session.post(
            url=self._url,
            params=ch_params,
            data=final_query.encode(),
            allow_redirects=False,
            headers={**tracing_headers},
            ssl_context=self._get_ssl_context(),
        )

        # CHYT rewrites incoming query_id so we log returned one instead of constructing own id
        LOGGER.info("Query has been sent to CH with query_id %s", resp.headers.get("X-ClickHouse-Query-Id"))

        return resp

    async def _parse_response_body(self, resp: ClientResponse) -> AsyncGenerator[tuple[str, Any], None]:
        parser = JSONCompactChunksParser()

        event = None
        data = None
        bytes_chunk: Optional[bytes] = None
        finish_sent = False

        async for bytes_chunk in resp.content.iter_chunked(self._http_read_chunk_size):
            parser.receive_data(bytes_chunk)

            while True:
                try:
                    event, data = parser.next_event()
                except parser.RowTooLarge:
                    raise CHRowTooLarge()
                if event == parser.parts.NEED_DATA:
                    break
                elif event == parser.parts.FINISHED:
                    if not finish_sent:
                        finish_sent = True
                        yield event, data
                    else:
                        # TODO FIX: Remove when JSONCompactChunksParser become able
                        #  to properly finalize remaining data in buffer after FINISH event
                        assert bytes_chunk is not None
                        LOGGER.info(
                            "Got duplicated FINISHED event from stream parser, chunk len=%s, truncated chunk: %r",
                            len(bytes_chunk),
                            bytes_chunk[:1024],
                        )

                    break
                else:
                    yield event, data

        else:
            if event != parser.parts.FINISHED:
                # TODO: return a sensible user-facing error in this case
                # (see a test in `bi-api/tests/db/result/test_error_handling.py`)
                decoded_chunk: Optional[str] = bytes_chunk.decode(errors="replace") if bytes_chunk is not None else None
                raise exc.SourceProtocolError(
                    debug_info=dict(msg="Unexpected last event", event=event, data=data, chunk=decoded_chunk),
                    # TODO: provide the actual query
                    query="",
                    # TODO?: show the chunk here for the user databases?
                    db_message="",
                )

        try:
            LOGGER.info(
                "Dumping clickhouse summary from headers",
                extra={"ch_summary": json.loads(resp.headers["X-ClickHouse-Summary"])},
            )
        except KeyError:
            LOGGER.warning("No clickhouse summary in headers")
        except Exception:
            LOGGER.warning("Failed to dump clickhouse summary", exc_info=True)

    @generic_profiler_async("db-full")  # type: ignore  # TODO: fix
    async def execute(self, query: DBAdapterQuery) -> AsyncRawExecutionResult:
        et = JSONCompactChunksParser.parts

        with self.wrap_execute_excs(query=query, stage="request"):
            # TODO FIX: Statuses for retry
            resp = await self._make_query(query)

        if resp.status != 200:
            bytes_body = await resp.read()
            body = bytes_body.decode(errors="replace")
            db_exc = self.make_exc(resp.status, body, debug_compiled_query=query.debug_compiled_query)
            raise db_exc

        # Primarily for DashSQL
        ch_resp_headers = {key: val for key, val in resp.headers.items() if key.lower().startswith("x-clickhouse-")}

        if query.is_ddl_dml_query:

            async def empty_chunk_gen() -> TBIChunksGen:
                return
                yield  # noqa

            return AsyncRawExecutionResult(
                raw_cursor_info=dict(clickhouse_headers=ch_resp_headers),
                raw_chunk_generator=empty_chunk_gen(),
            )

        def _safe_col_converter(col_conv: Callable[[str], Any], val: str) -> Any:
            try:
                return col_conv(val)
            except ValueError as err:
                raise exc.DataParseError(f"Cannot convert {val!r}", query=query.debug_compiled_query) from err

        events_generator = self._parse_response_body(resp)
        with self.wrap_execute_excs(query=query, stage="meta"):
            first_evt_type, first_evt_data = await events_generator.__anext__()
        assert first_evt_type == et.META
        if self._target_dto.disable_value_processing:
            row_converters = tuple(None for _ in first_evt_data)
        else:
            row_converters = tuple(_get_type(field["type"]) for field in first_evt_data)

        # TODO FIX: Here we ignore requested size of chunk
        # TODO FIX: Handle to large row exception
        async def chunk_gen() -> TBIChunksGen:
            expected_types: Sequence[Any] = (et.DATACHUNK, et.STATS)
            evt_type, evt_data = None, None

            with self.wrap_execute_excs(query=query, stage="rows"):
                async for evt_tuple in events_generator:
                    evt_type, evt_data = evt_tuple
                    if evt_type not in expected_types:
                        raise exc.SourceProtocolError()  # type: ignore  # TODO: fix

                    if evt_type == et.DATACHUNK:
                        yield tuple(
                            tuple(
                                _safe_col_converter(col_converter, val)  # type: ignore  # TODO: fix
                                if col_converter is not None
                                else val
                                for val, col_converter in zip(raw_row, row_converters)
                            )
                            for raw_row in evt_data
                        )
                    elif evt_type == et.STATS:
                        LOGGER.info(
                            "Dumping clickhouse statistics from response",
                            extra=evt_data,  # statistics (elapsed, rows_read, bytes_read) + rows (num)
                        )
                        expected_types = (et.FINISHED,)
                    elif evt_type == et.FINISHED:
                        expected_types = ()

            if evt_type != et.FINISHED:
                # TODO: For user databases this should probably be passed to the user.
                raise exc.SourceProtocolError(  # type: ignore  # TODO: fix
                    debug_info=dict(msg="Unexpected last event", event=evt_type, data=evt_data)
                )

        return AsyncRawExecutionResult(
            # Reminder: this `raw_cursor_info` gets (sometimes) serialized over RQE (async-ext).
            # Primarily useful for DashSQL
            raw_cursor_info=dict(
                # Deprecating this in favor of `names` + `db_types`:
                columns=[{"name": col["name"], "clickhouse_type_name": col["type"]} for col in first_evt_data],
                # ...
                names=[col["name"] for col in first_evt_data],
                driver_types=[col["type"] for col in first_evt_data],
                db_types=[self._ch_type_name_to_native_type(col["type"]) for col in first_evt_data],
                clickhouse_headers=ch_resp_headers,
            ),
            # Data
            raw_chunk_generator=chunk_gen(),
        )

    def _ch_type_name_to_native_type(self, type_name: str) -> GenericNativeType:
        """A simplified CH type parser for some commonly useful cases"""
        type_pieces = type_name.lower().strip(")").split("(")
        if type_pieces and type_pieces[0] == "lowcardinality":
            type_pieces = type_pieces[1:]
        if type_pieces and type_pieces[0] == "nullable":
            type_pieces = type_pieces[1:]
        return GenericNativeType.normalize_name_and_create(
            conn_type=self.conn_type,
            name=type_pieces[0] if type_pieces else "",
        )

    def make_exc(  # TODO:  Move to ErrorTransformer
        self,
        status_code: int,  # noqa
        err_body: str,
        debug_compiled_query: Optional[str] = None,
    ) -> exc.DatabaseQueryError:
        exc_cls: Type[exc.DatabaseQueryError]
        try:
            exc_cls, exc_params = self.ch_utils.get_exc_class(err_body) or (exc.DatabaseQueryError, None)
        except Exception:  # noqa
            exc_cls, exc_params = exc.DatabaseQueryError, None

        return exc_cls(
            db_message=err_body,
            query=debug_compiled_query,
            orig=None,
            details={},
            params=exc_params,
        )

    async def test(self) -> None:
        await self.execute(DBAdapterQuery("select 1"))

    async def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        raise NotImplementedError()

    async def get_schema_names(self, db_ident: DBIdent) -> list[str]:
        raise NotImplementedError()

    async def get_tables(self, schema_ident: SchemaIdent) -> list[TableIdent]:
        raise NotImplementedError()

    async def get_table_info(self, table_def: TableDefinition, fetch_idx_info: bool) -> RawSchemaInfo:
        raise NotImplementedError()

    async def is_table_exists(self, table_ident: TableIdent) -> bool:
        raise NotImplementedError()

    @classmethod
    def create(  # type: ignore  # TODO: fix
        cls: Type[_DBA_TV],
        target_dto: _TARGET_DTO_TV,
        req_ctx_info: DBAdapterScopedRCI,
        default_chunk_size: int,
    ) -> _DBA_TV:
        return cls(
            target_dto=target_dto,
            req_ctx_info=req_ctx_info,
            default_chunk_size=default_chunk_size,
        )


class AsyncClickHouseAdapter(BaseAsyncClickHouseAdapter):
    conn_type = CONNECTION_TYPE_CLICKHOUSE
    ch_utils = ClickHouseUtils

    def _get_ssl_context(self) -> Optional[ssl.SSLContext]:
        if not self._target_dto.secure or self._target_dto.ssl_ca is None:
            return None

        return ssl.create_default_context(cadata=self._target_dto.ssl_ca)

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/adapters.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import re
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Optional,
    Pattern,
    Type,
)

import attr
from sqlalchemy.engine.default import DefaultDialect
from sqlalchemy.sql import ddl as sa_ddl
from sqlalchemy.sql import schema as sa_schema

from dl_core import exc
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
from dl_core.db import (
    SchemaColumn,
    make_sa_type,
)

from dl_connector_clickhouse.core.clickhouse_base.exc import (
    CannotInsertNullInOrdinaryColumn,
    CHIncorrectData,
    CHReadonlyUser,
    ClickHouseSourceDoesNotExistError,
    EstimatedExecutionTooLong,
    InvalidSplitSeparator,
    TooManyColumns,
)


if TYPE_CHECKING:
    from dl_core.db.conversion_base import TypeTransformer

LOGGER = logging.getLogger(__name__)


def get_ch_settings(
    read_only_level: Optional[int] = None,
    max_execution_time: Optional[int] = None,
    insert_quorum: Optional[int] = None,
    insert_quorum_timeout: Optional[int] = None,
    output_format_json_quote_denormals: Optional[int] = None,
) -> dict:
    settings = {
        # https://clickhouse.com/docs/en/operations/settings/settings#settings-join_use_nulls
        # 1 — JOIN behaves the same way as in standard SQL.
        # The type of the corresponding field is converted to Nullable, and empty cells are filled with NULL.
        "join_use_nulls": 1,
        # Limits the time to wait for a response from the servers in the cluster.
        # If a ddl request has not been performed on all hosts, a response will contain
        # a timeout error and a request will be executed in an async mode.
        "distributed_ddl_task_timeout": 280,
        # https://clickhouse.com/docs/en/operations/settings/query-complexity#max-execution-time
        # Maximum query execution time in seconds.
        # By default, specify a large value to ensure there are no
        # forever-running queries (which is also known to break old-version CH
        # hosts at around 100_000 second long queries).
        # Note that in CH the value is rounded down to integer, and 0 seems to mean 'no limit'.
        "max_execution_time": max_execution_time if max_execution_time is None else 3600 * 4,
        "readonly": read_only_level,
        # https://clickhouse.com/docs/en/operations/settings/settings#settings-insert_quorum
        # INSERT succeeds only when ClickHouse manages to correctly write data to the insert_quorum
        # of replicas during the insert_quorum_timeout
        "insert_quorum": insert_quorum,
        "insert_quorum_timeout": insert_quorum_timeout,
        # request clickhouse stat in response headers
        # otherwise clickhouse sends nulls in X-ClickHouse-Summary
        "send_progress_in_http_headers": 0,
        # https://clickhouse.com/docs/en/operations/settings/formats#output_format_json_quote_denormals
        # Enables +nan, -nan, +inf, -inf outputs in JSON output format.
        # After support from frontend should be enabled by default in all cases.
        "output_format_json_quote_denormals": output_format_json_quote_denormals,
    }

    return {k: v for k, v in settings.items() if v is not None}


@attr.s(frozen=True, auto_attribs=True)
class ParsedErrorMsg:
    full_msg: str
    code: int


class ClickHouseBaseUtils:
    add_real_user_header: ClassVar[bool] = False

    ch_err_msg_re: ClassVar[Pattern] = re.compile(r"^(std::exception\.\s+)?Code:\s+(?P<code>\d+)")

    # ClickHouse error codes list: https://github.com/ClickHouse/ClickHouse/blob/master/src/Common/ErrorCodes.cpp
    # TODO: verify codes meaning and pick some other useful codes from CH list.
    map_err_code_exc_cls: ClassVar[dict[int, Type[exc.DatabaseQueryError]]] = {
        10: exc.ColumnDoesNotExist,
        36: InvalidSplitSeparator,
        41: exc.CannotParseDateTime,
        43: exc.InvalidArgumentType,
        46: exc.UnknownFunction,
        47: exc.ColumnDoesNotExist,
        53: exc.JoinColumnTypeMismatch,
        60: ClickHouseSourceDoesNotExistError,
        62: exc.InvalidQuery,
        70: exc.UnexpectedInfOrNan,
        72: exc.CannotParseNumber,
        81: exc.DatabaseDoesNotExist,
        117: CHIncorrectData,
        153: exc.DivisionByZero,
        159: exc.SourceTimeout,
        160: EstimatedExecutionTooLong,
        161: TooManyColumns,
        164: CHReadonlyUser,
        241: exc.DbMemoryLimitExceeded,
        277: exc.DBIndexNotUsed,
        349: CannotInsertNullInOrdinaryColumn,
        516: exc.DbAuthenticationFailed,
        1000: exc.NoSpaceLeft,
    }

    @classmethod
    def parse_message(cls, err_msg: str) -> Optional[ParsedErrorMsg]:
        match = cls.ch_err_msg_re.match(err_msg)

        if match is None:
            return None

        code = int(match.group("code"))
        LOGGER.info("Got CH error code %s", code)

        return ParsedErrorMsg(
            code=code,
            full_msg=err_msg,
        )

    @classmethod
    def get_exc_class_by_parsed_message(
        cls, msg: ParsedErrorMsg
    ) -> Optional[tuple[Type[exc.DatabaseQueryError], dict[str, str]]]:
        if msg.code in cls.map_err_code_exc_cls:
            return cls.map_err_code_exc_cls[msg.code], {}
        return None

    @classmethod
    def get_exc_class(cls, err_msg: str) -> Optional[tuple[Type[exc.DatabaseQueryError], dict[str, str]]]:
        parse_msg = cls.parse_message(err_msg)
        if parse_msg:
            return cls.get_exc_class_by_parsed_message(parse_msg)
        raise ValueError("Can not parse error message", err_msg)

    @classmethod
    def get_context_headers(cls, rci: DBAdapterScopedRCI) -> dict[str, str]:
        headers: dict[str, str] = {
            # TODO: bi_constants / bi_configs.constants
            "user-agent": "DataLens",
        }

        if rci.user_name and cls.add_real_user_header:
            headers["X-DataLens-Real-User"] = rci.user_name

        if rci.request_id:
            headers["x-request-id"] = rci.request_id

        if rci.client_ip:
            headers["X-Forwarded-For-Y"] = rci.client_ip

        return headers

    @classmethod
    def get_tracing_sample_flag_override(cls, rci: DBAdapterScopedRCI) -> Optional[bool]:
        """
        Return overridden sample flag value. If value should not be overridden `None` will be returned.
        """
        return None


class ClickHouseUtils(ClickHouseBaseUtils):
    pass


def create_column_sql(
    sa_dialect: DefaultDialect,
    col: SchemaColumn,
    tt: TypeTransformer,
    partition_fields: Optional[list[str]] = None,
) -> str:
    native_type = tt.type_user_to_native(user_t=col.user_type, native_t=col.native_type)

    if partition_fields and col.name in partition_fields:
        # Partition column cannot be nullable. Enforcing it in here.
        nullable = False
    else:
        nullable = None  # type: ignore  # TODO: fix

    if nullable is not None:
        native_type = native_type.as_common().clone(nullable=nullable)

    sa_type = make_sa_type(native_type)

    sa_column_obj = sa_schema.Column(col.name, sa_type)
    sa_cc_obj = sa_ddl.CreateColumn(sa_column_obj)
    sa_ddl_obj = sa_cc_obj.compile(dialect=sa_dialect)
    return sa_ddl_obj.string


def ensure_db_message(exc_cls, kw):  # type: ignore  # TODO: fix
    db_message = kw.get("db_message")
    details = kw.get("details")
    if db_message and details is not None and not details.get("db_message"):
        details.update(db_message=db_message)
    return exc_cls, kw

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/ch_commons.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

import attr

from dl_core.connection_models.conn_options import ConnectOptions


@attr.s(frozen=True, hash=True)
class CHConnectOptions(ConnectOptions):
    max_execution_time: Optional[int] = attr.ib(default=None)
    connect_timeout: Optional[int] = attr.ib(default=None)
    total_timeout: Optional[int] = attr.ib(default=None)
    insert_quorum: Optional[int] = attr.ib(default=None)
    insert_quorum_timeout: Optional[int] = attr.ib(default=None)
    disable_value_processing: bool = attr.ib(default=False)

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/conn_options.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    TypeVar,
)

import attr

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor

from dl_connector_clickhouse.core.clickhouse_base.adapters import (
    AsyncClickHouseAdapter,
    ClickHouseAdapter,
)
from dl_connector_clickhouse.core.clickhouse_base.conn_options import CHConnectOptions
from dl_connector_clickhouse.core.clickhouse_base.dto import ClickHouseConnDTO
from dl_connector_clickhouse.core.clickhouse_base.target_dto import ClickHouseConnTargetDTO


_BASE_CLICKHOUSE_ADAPTER_TV = TypeVar("_BASE_CLICKHOUSE_ADAPTER_TV", bound=CommonBaseDirectAdapter)


@attr.s(cmp=False, hash=False)
class _BaseClickHouseConnExecutor(DefaultSqlAlchemyConnExecutor[_BASE_CLICKHOUSE_ADAPTER_TV]):
    _conn_dto: ClickHouseConnDTO = attr.ib()
    _conn_options: CHConnectOptions = attr.ib()

    async def _make_target_conn_dto_pool(self) -> list[ClickHouseConnTargetDTO]:  # type: ignore  # TODO: fix
        dto_pool = []
        for host in self._conn_hosts_pool:
            dto_pool.append(
                ClickHouseConnTargetDTO(
                    conn_id=self._conn_dto.conn_id,
                    pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                    pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                    protocol=self._conn_dto.protocol,
                    host=host,
                    port=self._conn_dto.port,
                    db_name=self._conn_dto.db_name,
                    cluster_name=self._conn_dto.cluster_name,
                    endpoint=self._conn_dto.endpoint,
                    username=self._conn_dto.username,
                    password=self._conn_dto.password,
                    max_execution_time=self._conn_options.max_execution_time,
                    total_timeout=self._conn_options.total_timeout,
                    connect_timeout=self._conn_options.connect_timeout,
                    insert_quorum=self._conn_options.insert_quorum,
                    insert_quorum_timeout=self._conn_options.insert_quorum_timeout,
                    disable_value_processing=self._conn_options.disable_value_processing,
                    secure=self._conn_dto.secure,
                    ssl_ca=self._conn_dto.ssl_ca,
                )
            )
        return dto_pool

    def mutate_for_dashsql(self, db_params: Optional[dict[str, str]] = None):  # type: ignore  # TODO: fix
        if db_params:  # TODO.
            raise Exception("`db_params` are not supported at the moment")
        return self.clone(
            conn_options=self._conn_options.clone(
                disable_value_processing=True,
            ),
        )


@attr.s(cmp=False, hash=False)
class ClickHouseConnExecutor(_BaseClickHouseConnExecutor[ClickHouseAdapter]):
    TARGET_ADAPTER_CLS = ClickHouseAdapter


@attr.s(cmp=False, hash=False)
class AsyncClickHouseConnExecutor(_BaseClickHouseConnExecutor[AsyncClickHouseAdapter]):
    TARGET_ADAPTER_CLS = AsyncClickHouseAdapter

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/connection_executors.py
datalens-tech/datalens-backend
from clickhouse_sqlalchemy.orm.query import Query as CHQuery

from dl_core.connectors.base.connector import CoreConnector

from dl_connector_clickhouse.core.clickhouse_base.adapters import (
    AsyncClickHouseAdapter,
    ClickHouseAdapter,
)
from dl_connector_clickhouse.core.clickhouse_base.constants import BACKEND_TYPE_CLICKHOUSE
from dl_connector_clickhouse.core.clickhouse_base.query_compiler import ClickHouseQueryCompiler
from dl_connector_clickhouse.core.clickhouse_base.sa_types import SQLALCHEMY_CLICKHOUSE_TYPES


class ClickHouseCoreConnectorBase(CoreConnector):
    backend_type = BACKEND_TYPE_CLICKHOUSE
    compiler_cls = ClickHouseQueryCompiler
    query_cls = CHQuery
    rqe_adapter_classes = frozenset({ClickHouseAdapter, AsyncClickHouseAdapter})
    sa_types = SQLALCHEMY_CLICKHOUSE_TYPES

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    SourceBackendType,
)


BACKEND_TYPE_CLICKHOUSE = SourceBackendType.declare("CLICKHOUSE")
CONNECTION_TYPE_CLICKHOUSE = ConnectionType.declare("clickhouse")

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    ClassVar,
)

from dl_constants.enums import JoinType
from dl_core.data_source.sql import (
    BaseSQLDataSource,
    StandardSQLDataSource,
    SubselectDataSource,
)

from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
from dl_connector_clickhouse.core.clickhouse_base.query_compiler import ClickHouseQueryCompiler


if TYPE_CHECKING:
    from dl_core import us_connection  # noqa
    from dl_core.db import (  # noqa
        SchemaColumn,
        SchemaInfo,
    )


LOGGER = logging.getLogger(__name__)


class CommonClickHouseSubselectDataSource(SubselectDataSource):
    """Common subselect base for ClickHouse-based databases"""

    pass


class ClickHouseBaseMixin(BaseSQLDataSource):
    default_server_version = None
    supported_join_types: ClassVar[frozenset[JoinType]] = frozenset(
        {
            JoinType.inner,
            JoinType.left,
            JoinType.full,
            JoinType.right,
        }
    )
    compiler_cls = ClickHouseQueryCompiler

    def get_connect_args(self) -> dict:
        return dict(super().get_connect_args(), server_version=self.db_version or self.default_server_version)  # type: ignore  # TODO: fix  # noqa


class ActualClickHouseBaseMixin(ClickHouseBaseMixin):
    """
    Data source base class / mixin for actual ClickHouse (i.e. not chyt).
    """

    conn_type = CONNECTION_TYPE_CLICKHOUSE


class ClickHouseDataSourceBase(ActualClickHouseBaseMixin, StandardSQLDataSource, metaclass=abc.ABCMeta):  # type: ignore  # TODO: fix
    """ClickHouse table. Might not work correctly for views."""

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/data_source.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

import attr

from dl_core.connection_models.dto_defs import (
    ConnDTO,
    DefaultSQLDTO,
)

from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE


@attr.s(frozen=True)
class ClickHouseBaseDTO(ConnDTO):
    pass


@attr.s(frozen=True)
class ClickHouseConnDTO(ClickHouseBaseDTO, DefaultSQLDTO):  # noqa
    conn_type = CONNECTION_TYPE_CLICKHOUSE
    secure: bool = attr.ib(kw_only=True, default=False)
    ssl_ca: Optional[str] = attr.ib(kw_only=True, default=None)

    protocol: str = attr.ib(kw_only=True)
    # TODO CONSIDER: Is really optional
    endpoint: Optional[str] = attr.ib(kw_only=True)
    cluster_name: str = attr.ib(kw_only=True)

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/dto.py
datalens-tech/datalens-backend
from __future__ import annotations

import re
from typing import (
    Any,
    Optional,
)

from dl_core import exc


class ClickHouseSourceDoesNotExistError(exc.SourceDoesNotExist):
    ERR_RE = re.compile(r".*Table\s(?P<table>.*)\sdoesn't exist.*")

    def __init__(
        self,
        db_message: Optional[str] = None,
        query: Optional[str] = None,
        message: Optional[str] = None,
        details: Optional[dict[str, Any]] = None,
        orig: Optional[Exception] = None,
        debug_info: Optional[dict[str, Any]] = None,
        params: Optional[dict[str, Any]] = None,
    ):
        super(ClickHouseSourceDoesNotExistError, self).__init__(
            db_message=db_message,
            query=query,
            message=message,
            details=details,
            orig=orig,
            debug_info=debug_info,
            params=params,
        )
        if self.db_message and (match := self.ERR_RE.match(self.db_message)):
            if table := match.group("table"):
                self.params["table_definition"] = table


class CannotInsertNullInOrdinaryColumn(exc.DatabaseQueryError):
    err_code = exc.DatabaseQueryError.err_code + ["CANNOT_INSERT_NULL_IN_ORDINARY_COLUMN"]
    default_message = "Cannot convert NULL value to non-Nullable type."


class CHQueryError(exc.DatabaseQueryError):
    err_code = exc.DatabaseQueryError.err_code + ["CH"]
    default_message = "CH Error"


class CHIncorrectData(CHQueryError):
    err_code = CHQueryError.err_code + ["INCORRECT_DATA"]
    default_message = "Clickhouse could not parse the data in the specified source"


class CHReadonlyUser(CHQueryError):
    err_code = CHQueryError.err_code + ["READONLY_USER"]
    default_message = "Clickhouse user should have parameter readonly set to 0 or 2"


class EstimatedExecutionTooLong(exc.DatabaseQueryError):
    err_code = exc.DatabaseQueryError.err_code + ["EST_EXEC_TOO_LONG"]
    default_message = "Estimated query execution time is too long. Maximum: 14400."


class TooManyColumns(exc.DatabaseQueryError):
    err_code = exc.DatabaseQueryError.err_code + ["TOO_MANY_COLUMNS"]
    default_message = "Too many columns - limit exceeded."


class InvalidSplitSeparator(exc.DatabaseQueryError):
    err_code = exc.DatabaseQueryError.err_code + ["INVALID_SPLIT_SEPARATOR"]
    default_message = "Invalid separator for splitting."


class CHRowTooLarge(exc.SourceProtocolError):
    err_code = exc.SourceProtocolError.err_code + ["TOO_LARGE_ROW"]
    default_message = "Data source failed to respond correctly (too large row)."

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/exc.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import ClassVar

from dl_core.connectors.base.query_compiler import (
    QueryCompiler,
    SectionAliasMode,
)


class ClickHouseQueryCompiler(QueryCompiler):
    groupby_alias_mode: ClassVar[SectionAliasMode] = SectionAliasMode.by_alias_in_section
    force_nulls_lower_than_values = True

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/query_compiler.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import partial
from typing import (
    Callable,
    Collection,
    Type,
)

from clickhouse_sqlalchemy import types as ch_types
import sqlalchemy as sa
from sqlalchemy.types import TypeEngine

from dl_constants.enums import ConnectionType
from dl_core.db.native_type import (
    ClickHouseDateTime64NativeType,
    ClickHouseDateTime64WithTZNativeType,
    ClickHouseDateTimeWithTZNativeType,
    ClickHouseNativeType,
    CommonNativeType,
    GenericNativeType,
)
from dl_core.db.sa_types_base import make_native_type

from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
from dl_connector_clickhouse.core.clickhouse_base.type_transformer import (
    CH_TYPES_DATE,
    CH_TYPES_FLOAT,
    CH_TYPES_INT,
)


def _make_ch_type(nt: GenericNativeType, typeobj: TypeEngine) -> TypeEngine:
    nullable = True
    lowcardinality = False
    if isinstance(nt, CommonNativeType):
        nullable = nt.nullable
    if isinstance(nt, ClickHouseNativeType):
        lowcardinality = nt.lowcardinality
    result = typeobj
    if nullable:
        result = ch_types.Nullable(result)
    if lowcardinality:
        result = ch_types.LowCardinality(result)
    return result


def ch_instantiator(typecls: Type[TypeEngine]) -> Callable:
    def type_gen(nt: GenericNativeType) -> TypeEngine:
        return _make_ch_type(nt=nt, typeobj=typecls())

    return type_gen


def ch_fallback_type_gen(*args, **kwargs):  # type: ignore  # TODO: fix
    return ch_types.Nullable(ch_types.String())


def _make_ch_dtwtz(nt: GenericNativeType, typecls: Type[TypeEngine] = ch_types.DateTimeWithTZ) -> TypeEngine:
    if isinstance(nt, ClickHouseDateTimeWithTZNativeType):
        tz = nt.timezone_name
    else:
        tz = "UTC"
    typeobj = typecls(tz)
    return _make_ch_type(nt=nt, typeobj=typeobj)


DEFAULT_DT64_PRECISION = 9


def _make_ch_dt64(nt: GenericNativeType, typecls: Type[TypeEngine] = ch_types.DateTime64) -> TypeEngine:
    if isinstance(nt, ClickHouseDateTime64NativeType):
        precision = nt.precision
    else:
        precision = DEFAULT_DT64_PRECISION
    typeobj = typecls(precision)
    return _make_ch_type(nt=nt, typeobj=typeobj)


def _make_ch_dt64wtz(nt: GenericNativeType, typecls: Type[TypeEngine] = ch_types.DateTime64WithTZ) -> TypeEngine:
    if isinstance(nt, ClickHouseDateTime64WithTZNativeType):
        tz = nt.timezone_name
        precision = nt.precision
    else:
        tz = "UTC"
        precision = DEFAULT_DT64_PRECISION
    typeobj = typecls(precision, tz)
    return _make_ch_type(nt=nt, typeobj=typeobj)


def _make_ch_array(nt: GenericNativeType, inner_typecls: Type[TypeEngine]) -> TypeEngine:
    ch_type_inner = _make_ch_type(nt=nt, typeobj=inner_typecls())
    return ch_types.Array(ch_type_inner)


SQLALCHEMY_CLICKHOUSE_BASE_TYPES = (
    *CH_TYPES_INT,
    *CH_TYPES_FLOAT,
    ch_types.String,  # TODO: FixedString
    # # TODO: corresponding ClickHouseNativeType subclasses.
    # ch_types.Enum8, ch_types.Enum16, ch_types.Decimal,
    *CH_TYPES_DATE,
    ch_types.DateTime,
    ch_types.UUID,
    ch_types.Bool,
)


def _generate_complex_ch_types(
    base_ch_types: Collection[Type[TypeEngine]], conn_type: ConnectionType
) -> dict[GenericNativeType, Callable[..., TypeEngine]]:
    return {
        **{make_native_type(conn_type, typecls): ch_instantiator(typecls) for typecls in base_ch_types},
        make_native_type(conn_type, ch_types.DateTimeWithTZ): _make_ch_dtwtz,
        make_native_type(conn_type, ch_types.DateTime64): _make_ch_dt64,
        make_native_type(conn_type, ch_types.DateTime64WithTZ): _make_ch_dt64wtz,
        **{
            make_native_type(conn_type, ch_types.Array(typecls)): partial(_make_ch_array, inner_typecls=typecls)
            for typecls in CH_TYPES_INT
        },
        **{
            make_native_type(conn_type, ch_types.Array(typecls)): partial(_make_ch_array, inner_typecls=typecls)
            for typecls in CH_TYPES_FLOAT
        },
        make_native_type(conn_type, ch_types.Array(ch_types.String)): partial(
            _make_ch_array, inner_typecls=ch_types.String
        ),
        # For the `UserDataType.unsupported`; should only be filled with `NULL`s in materialization.
        # See also: `dl_core.data_source.sql.BaseSQLDataSource._make_raw_column_select`
        make_native_type(conn_type, sa.sql.sqltypes.NullType): ch_fallback_type_gen,
    }


SQLALCHEMY_CLICKHOUSE_TYPES = _generate_complex_ch_types(
    SQLALCHEMY_CLICKHOUSE_BASE_TYPES, conn_type=CONNECTION_TYPE_CLICKHOUSE
)

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/sa_types.py
datalens-tech/datalens-backend
from typing import (
    Generic,
    TypeVar,
)

from marshmallow import fields as ma_fields

from dl_core.us_manager.storage_schemas.connection import ConnectionSQLDataStorageSchema

from dl_connector_clickhouse.core.clickhouse_base.us_connection import ConnectionClickhouseBase


_CH_CONN_DATA_TV = TypeVar("_CH_CONN_DATA_TV", bound=ConnectionClickhouseBase.DataModel)


class ConnectionClickHouseBaseDataStorageSchema(
    ConnectionSQLDataStorageSchema[_CH_CONN_DATA_TV],
    Generic[_CH_CONN_DATA_TV],
):
    secure = ma_fields.Boolean(allow_none=False, required=False, load_default=False, dump_default=False)
    ssl_ca = ma_fields.String(required=False, allow_none=True, load_default=None, dump_default=None)
    endpoint = ma_fields.String(required=False, allow_none=True, load_default=None, dump_default=None)
    cluster_name = ma_fields.String(required=False, allow_none=True, load_default=None, dump_default=None)
    max_execution_time = ma_fields.Integer(required=False, allow_none=True, load_default=None, dump_default=None)

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/storage_schemas/connection.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_core.connection_executors.models.connection_target_dto_base import BaseSQLConnTargetDTO


@attr.s(frozen=True)
class BaseClickHouseConnTargetDTO(BaseSQLConnTargetDTO):
    protocol: str = attr.ib()
    # TODO CONSIDER: Is really optional?
    endpoint: Optional[str] = attr.ib()
    cluster_name: Optional[str] = attr.ib()
    max_execution_time: Optional[int] = attr.ib()
    connect_timeout: Optional[int] = attr.ib()
    total_timeout: Optional[int] = attr.ib()
    insert_quorum: Optional[int] = attr.ib()
    insert_quorum_timeout: Optional[int] = attr.ib()
    disable_value_processing: bool = attr.ib()
    secure: bool = attr.ib(kw_only=True, default=False)
    ssl_ca: Optional[str] = attr.ib(kw_only=True, default=None)


class ClickHouseConnTargetDTO(BaseClickHouseConnTargetDTO):
    """..."""

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/target_dto.py
datalens-tech/datalens-backend
from __future__ import annotations

from clickhouse_sqlalchemy import types as ch_types
import sqlalchemy as sa

from dl_constants.enums import UserDataType
from dl_core.backend_types import get_backend_type
from dl_core.db.conversion_base import (
    TypeTransformer,
    make_native_type,
)
from dl_core.db.elements import GenericNativeType

from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE


CH_TYPES_INT = frozenset(
    (
        ch_types.Int,
        ch_types.Int8,
        ch_types.Int16,
        ch_types.Int32,
        ch_types.Int64,
        ch_types.UInt8,
        ch_types.UInt16,
        ch_types.UInt32,
        ch_types.UInt64,
    )
)
CH_TYPES_FLOAT = frozenset((ch_types.Float, ch_types.Float32, ch_types.Float64, ch_types.Decimal))
CH_TYPES_DATE = frozenset((ch_types.Date, ch_types.Date32))


class ClickHouseTypeTransformer(TypeTransformer):
    conn_type = CONNECTION_TYPE_CLICKHOUSE
    native_to_user_map = {
        **{make_native_type(CONNECTION_TYPE_CLICKHOUSE, typecls): UserDataType.integer for typecls in CH_TYPES_INT},  # type: ignore  # TODO: fix
        **{
            make_native_type(CONNECTION_TYPE_CLICKHOUSE, typecls): UserDataType.string
            for typecls in (ch_types.String,)  # TODO: FixedString
        },
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Enum8): UserDataType.string,
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Enum16): UserDataType.string,
        **{make_native_type(CONNECTION_TYPE_CLICKHOUSE, typecls): UserDataType.float for typecls in CH_TYPES_FLOAT},
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Date): UserDataType.date,
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Date32): UserDataType.date,
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Bool): UserDataType.boolean,
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.DateTime): UserDataType.genericdatetime,
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.DateTime64): UserDataType.genericdatetime,
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.DateTimeWithTZ): UserDataType.genericdatetime,
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.DateTime64WithTZ): UserDataType.genericdatetime,
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.UUID): UserDataType.uuid,
        **{
            make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Array(typecls)): UserDataType.array_int
            for typecls in CH_TYPES_INT
        },
        **{
            make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Array(typecls)): UserDataType.array_float
            for typecls in CH_TYPES_FLOAT
        },
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Array(ch_types.String())): UserDataType.array_str,
        make_native_type(CONNECTION_TYPE_CLICKHOUSE, sa.sql.sqltypes.NullType): UserDataType.unsupported,
    }
    user_to_native_map = {
        UserDataType.integer: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Int64),
        UserDataType.float: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Float64),
        UserDataType.boolean: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Bool),
        UserDataType.string: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.String),
        UserDataType.date: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Date),
        UserDataType.datetime: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.DateTime),  # TODO: DateTime64
        UserDataType.genericdatetime: make_native_type(
            CONNECTION_TYPE_CLICKHOUSE, ch_types.DateTime
        ),  # TODO: DateTime64
        # WARNING: underparametrized
        UserDataType.datetimetz: make_native_type(
            CONNECTION_TYPE_CLICKHOUSE, ch_types.DateTimeWithTZ
        ),  # TODO: DateTime64WithTZ
        UserDataType.geopoint: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.String),
        UserDataType.geopolygon: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.String),
        UserDataType.uuid: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.UUID),
        UserDataType.markup: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.String),
        UserDataType.array_int: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Array(ch_types.Int64)),
        UserDataType.array_float: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Array(ch_types.Float64)),
        UserDataType.array_str: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Array(ch_types.String)),
        UserDataType.tree_str: make_native_type(CONNECTION_TYPE_CLICKHOUSE, ch_types.Array(ch_types.String)),
        UserDataType.unsupported: make_native_type(CONNECTION_TYPE_CLICKHOUSE, sa.sql.sqltypes.NullType),
    }

    def make_foreign_native_type_conversion(self, native_t: GenericNativeType) -> GenericNativeType:
        # All CH conn_types are supposed to have equivalent types.
        nt_backend_type = get_backend_type(conn_type=native_t.conn_type)
        own_backend_type = get_backend_type(self.conn_type)

        if nt_backend_type == own_backend_type:
            return native_t.clone(conn_type=self.conn_type)
        return super().make_foreign_native_type_conversion(native_t)

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import Enum
from typing import (
    Any,
    Callable,
    Optional,
)

import attr

from dl_core.connection_executors.common_base import ConnExecutorQuery
from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.connection_models import ConnectOptions
from dl_core.us_connection_base import (
    ClassicConnectionSQL,
    ConnectionBase,
)

from dl_connector_clickhouse.core.clickhouse_base.conn_options import CHConnectOptions
from dl_connector_clickhouse.core.clickhouse_base.dto import ClickHouseConnDTO


@attr.s
class SubselectTemplate:
    title: str = attr.ib()
    sql_query: str = attr.ib()


class SubselectParameterType(Enum):  # TODO: maybe redundant?
    single_value = "single_value"
    list_of_values = "list_of_values"


@attr.s
class SubselectParameter:
    name: str = attr.ib()
    ss_type: SubselectParameterType = attr.ib()
    values: Any = attr.ib()


class ConnectionClickhouseBase(ClassicConnectionSQL):
    MAX_ALLOWED_MAX_EXECUTION_TIME = 280  # 290 sec - requests http timeout, 300 sec - uwsgi harakiri

    @attr.s(kw_only=True)
    class DataModel(ClassicConnectionSQL.DataModel):
        secure: bool = attr.ib(default=False)
        endpoint: str = attr.ib(default="")
        cluster_name: Optional[str] = attr.ib(default=None)
        max_execution_time: Optional[int] = attr.ib(default=None)
        ssl_ca: Optional[str] = attr.ib(kw_only=True, default=None)

    def get_conn_dto(self) -> ClickHouseConnDTO:
        return ClickHouseConnDTO(
            conn_id=self.uuid,
            protocol="https" if self.data.secure else "http",
            host=self.data.host,
            multihosts=self.parse_multihosts(),  # type: ignore  # TODO: fix
            port=self.data.port,
            endpoint=self.data.endpoint,
            cluster_name=self.data.cluster_name,
            db_name=self.data.db_name,
            username=self.data.username,
            password=self.password,
            secure=self.data.secure,
            ssl_ca=self.data.ssl_ca,
        )

    @staticmethod
    def get_effective_conn_options(
        base_conn_opts: ConnectOptions,
        user_max_execution_time: Optional[int],
        max_allowed_max_execution_time: int,
    ) -> CHConnectOptions:
        if user_max_execution_time is None:
            max_execution_time = None
            total_timeout = max_allowed_max_execution_time + 10
        else:
            if 0 < user_max_execution_time <= max_allowed_max_execution_time:
                total_timeout = user_max_execution_time + 10
                max_execution_time = user_max_execution_time
            else:
                total_timeout = max_allowed_max_execution_time + 10
                max_execution_time = max_allowed_max_execution_time

        return base_conn_opts.to_subclass(
            CHConnectOptions,
            total_timeout=total_timeout,
            connect_timeout=1,
            max_execution_time=max_execution_time,
        )

    def get_conn_options(self) -> CHConnectOptions:
        base = super().get_conn_options()
        return self.get_effective_conn_options(
            base_conn_opts=base,
            user_max_execution_time=self.data.max_execution_time,
            max_allowed_max_execution_time=self.MAX_ALLOWED_MAX_EXECUTION_TIME,
        )

    def get_parameter_combinations(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[dict]:
        ch_system_dbs = (
            "system",
            "mdb_system",
            "_system",
            "information_schema",
        )
        conn_executor = conn_executor_factory(self)
        query = ConnExecutorQuery(query="SELECT `database`, `name` from `system`.`tables`", db_name="system")
        return [
            dict(db_name=db_name, table_name=table_name)
            for db_name, table_name in conn_executor.execute(query=query).get_all()
            if db_name.lower() not in ch_system_dbs
        ]

lib/dl_connector_clickhouse/dl_connector_clickhouse/core/clickhouse_base/us_connection.py
datalens-tech/datalens-backend
from dl_db_testing.connectors.base.connector import DbTestingConnector

from dl_connector_clickhouse.db_testing.engine_wrapper import (
    BiClickHouseEngineWrapper,
    ClickHouseEngineWrapper,
)


class ClickHouseDbTestingConnector(DbTestingConnector):
    engine_wrapper_classes = (
        ClickHouseEngineWrapper,
        BiClickHouseEngineWrapper,
    )

lib/dl_connector_clickhouse/dl_connector_clickhouse/db_testing/connector.py
datalens-tech/datalens-backend
from typing import (
    Optional,
    Sequence,
)

import attr
from clickhouse_sqlalchemy import Table as CHTable
from clickhouse_sqlalchemy.engines import Log
import shortuuid
import sqlalchemy as sa

from dl_db_testing.database.engine_wrapper import (
    DbEngineConfig,
    EngineWrapperBase,
)


@attr.s(frozen=True)
class ClickhouseDbEngineConfig(DbEngineConfig):
    cluster: Optional[str] = attr.ib(kw_only=True, default=None)


class ClickHouseEngineWrapperBase(EngineWrapperBase):
    CONFIG_CLS = ClickhouseDbEngineConfig

    @property
    def config(self) -> ClickhouseDbEngineConfig:
        assert isinstance(self._config, ClickhouseDbEngineConfig)
        return self._config

    def get_conn_credentials(self, full: bool = False) -> dict:
        creds = super().get_conn_credentials(full)
        if not full:
            creds.pop("db_name")
        return creds

    def load_table(self, table_name: str, schema: Optional[str] = None) -> sa.Table:
        return CHTable(table_name, sa.MetaData(bind=self.engine), autoload=True)

    def drop_table(self, db_name: str, table: sa.Table) -> None:
        assert isinstance(table, CHTable)
        if self.config.cluster:
            self.execute(
                sa.text(f"DROP TABLE IF EXISTS `{db_name}`.`{table.name}` " f"ON CLUSTER `{self.config.cluster}`")
            )
        else:
            table.drop(bind=self.engine, if_exists=True)

    def table_from_columns(
        self,
        columns: Sequence[sa.Column],
        *,
        schema: Optional[str] = None,
        table_name: Optional[str] = None,
    ) -> sa.Table:
        table_name = table_name or f"test_table_{shortuuid.uuid()[:10]}"
        table = CHTable(table_name, sa.MetaData(), *columns)
        table.engine = Log()
        return table


class ClickHouseEngineWrapper(ClickHouseEngineWrapperBase):
    URL_PREFIX = "clickhouse"


class BiClickHouseEngineWrapper(ClickHouseEngineWrapperBase):
    URL_PREFIX = "bi_clickhouse"

    def _db_connect_params(self) -> dict:
        return {"format": "JSONCompact"}

lib/dl_connector_clickhouse/dl_connector_clickhouse/db_testing/engine_wrapper.py
datalens-tech/datalens-backend
from clickhouse_sqlalchemy.drivers.base import ClickHouseDialect as SAClickHouseDialect

from dl_formula.connectors.base.connector import FormulaConnector

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as ClickHouseDialectNS
from dl_connector_clickhouse.formula.definitions.all import DEFINITIONS
from dl_connector_clickhouse.formula.literal import ClickHouseLiteralizer
from dl_connector_clickhouse.formula.type_constructor import ClickHouseTypeConstructor


class ClickHouseFormulaConnector(FormulaConnector):
    dialect_ns_cls = ClickHouseDialectNS
    dialects = ClickHouseDialectNS.CLICKHOUSE
    default_dialect = ClickHouseDialectNS.CLICKHOUSE_21_8
    op_definitions = DEFINITIONS
    literalizer_cls = ClickHouseLiteralizer
    type_constructor_cls = ClickHouseTypeConstructor
    sa_dialect = SAClickHouseDialect()

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/connector.py
datalens-tech/datalens-backend
from dl_formula.core.dialect import (
    DialectName,
    DialectNamespace,
    simple_combo,
)


DIALECT_NAME_CLICKHOUSE = DialectName.declare("CLICKHOUSE")


class ClickHouseDialect(DialectNamespace):
    CLICKHOUSE_21_8 = simple_combo(name=DIALECT_NAME_CLICKHOUSE, version=(21, 8))
    CLICKHOUSE_22_10 = simple_combo(name=DIALECT_NAME_CLICKHOUSE, version=(22, 10))
    CLICKHOUSE = CLICKHOUSE_21_8 | CLICKHOUSE_22_10

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/constants.py
datalens-tech/datalens-backend
from dl_connector_clickhouse.formula.definitions.conditional_blocks import DEFINITIONS_COND_BLOCKS
from dl_connector_clickhouse.formula.definitions.functions_aggregation import DEFINITIONS_AGG
from dl_connector_clickhouse.formula.definitions.functions_array import DEFINITIONS_ARRAY
from dl_connector_clickhouse.formula.definitions.functions_datetime import DEFINITIONS_DATETIME
from dl_connector_clickhouse.formula.definitions.functions_logical import DEFINITIONS_LOGICAL
from dl_connector_clickhouse.formula.definitions.functions_markup import DEFINITIONS_MARKUP
from dl_connector_clickhouse.formula.definitions.functions_math import DEFINITIONS_MATH
from dl_connector_clickhouse.formula.definitions.functions_special import DEFINITIONS_SPECIAL
from dl_connector_clickhouse.formula.definitions.functions_string import DEFINITIONS_STRING
from dl_connector_clickhouse.formula.definitions.functions_type import DEFINITIONS_TYPE
from dl_connector_clickhouse.formula.definitions.functions_window import DEFINITIONS_WINDOW
from dl_connector_clickhouse.formula.definitions.operators_binary import DEFINITIONS_BINARY
from dl_connector_clickhouse.formula.definitions.operators_ternary import DEFINITIONS_TERNARY
from dl_connector_clickhouse.formula.definitions.operators_unary import DEFINITIONS_UNARY


DEFINITIONS = [
    *DEFINITIONS_COND_BLOCKS,
    *DEFINITIONS_AGG,
    *DEFINITIONS_DATETIME,
    *DEFINITIONS_LOGICAL,
    *DEFINITIONS_MARKUP,
    *DEFINITIONS_MATH,
    *DEFINITIONS_SPECIAL,
    *DEFINITIONS_STRING,
    *DEFINITIONS_TYPE,
    *DEFINITIONS_WINDOW,
    *DEFINITIONS_UNARY,
    *DEFINITIONS_BINARY,
    *DEFINITIONS_TERNARY,
    *DEFINITIONS_ARRAY,
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/all.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import (
    TranslationVariant,
    TranslationVariantWrapped,
)
import dl_formula.definitions.conditional_blocks as base

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


DEFINITIONS_COND_BLOCKS = [
    # _case_block_
    base.CaseBlock(
        variants=[
            VW(D.CLICKHOUSE, lambda *args: base.CaseBlock.translation(*args, as_multiif=True, replace_nulls=True)),
        ]
    ),
    # _if_block_
    base.IfBlock3(
        variants=[
            V(D.CLICKHOUSE, getattr(sa.func, "if")),
        ]
    ),
    base.IfBlockMulti.for_dialect(D.CLICKHOUSE),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/conditional_blocks.py
datalens-tech/datalens-backend
import clickhouse_sqlalchemy.ext.clauses as chsq_clauses
import sqlalchemy as sa
from sqlalchemy.sql import ClauseElement

from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common import quantile_value
import dl_formula.definitions.functions_aggregation as base
from dl_formula.definitions.literals import (
    literal,
    un_literal,
)
from dl_formula.shortcuts import n

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make


def _explicit_ch_tostring(value: ClauseElement) -> ClauseElement:
    """
    A hax-fix function that
    1. has explicit arguments, for CHSA Lambda
    2. is pickleable (not a lambda)
    """
    return sa.func.toString(value)


def _all_concat_ch(expr: ClauseElement, sep: str = ", ") -> ClauseElement:
    res = expr
    # Any --agg-> Array(Any)
    res = sa.func.groupUniqArray(res)
    # Reduce the stochasticity.
    res = sa.func.arraySort(res)
    # Array(Any) -> Array(String)
    res = sa.func.arrayMap(chsq_clauses.Lambda(_explicit_ch_tostring), res)
    # Array(String) -> String
    res = sa.func.arrayStringConcat(res, sep)
    return res


def _top_concat_ch(expr: ClauseElement, amount: ClauseElement, sep: str = ", ") -> ClauseElement:
    # TODO?: validate `amount > 0`?
    res = expr
    # Any --agg-> Array(Any)
    res = chsq_clauses.ParametrizedFunction("topK", [amount], expr)
    # Result is frequency-ordered, so not sorting it, even though it can be
    # stochastic for equal frequencies.
    # Array(Any) -> Array(String)
    res = sa.func.arrayMap(chsq_clauses.Lambda(lambda x: sa.func.toString(x)), res)
    # Array(String) -> String
    res = sa.func.arrayStringConcat(res, sep)
    return res


DEFINITIONS_AGG = [
    # all_concat
    base.AggAllConcat1(
        variants=[
            V(D.CLICKHOUSE, _all_concat_ch),
        ]
    ),
    base.AggAllConcat2(
        variants=[
            V(D.CLICKHOUSE, _all_concat_ch),
        ]
    ),
    # any
    base.AggAny(
        variants=[
            V(D.CLICKHOUSE, sa.func.any),
        ]
    ),
    # arg_max
    base.AggArgMax(
        variants=[
            V(D.CLICKHOUSE, sa.func.argMax),
        ]
    ),
    # arg_min
    base.AggArgMin(
        variants=[
            V(D.CLICKHOUSE, sa.func.argMin),
        ]
    ),
    # avg
    base.AggAvgFromNumber.for_dialect(D.CLICKHOUSE),
    base.AggAvgFromDate.for_dialect(D.CLICKHOUSE),
    base.AggAvgFromDatetime.for_dialect(D.CLICKHOUSE),
    base.AggAvgFromDatetimeTZ.for_dialect(D.CLICKHOUSE),
    # avg_if
    base.AggAvgIf(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda expr, cond: getattr(sa.func, "if")(
                    sa.func.isNaN(sa.func.avgIf(expr, cond)),
                    sa.null(),
                    sa.func.avgIf(expr, cond),
                ),
            ),
        ]
    ),
    # count
    base.AggCount0.for_dialect(D.CLICKHOUSE),
    base.AggCount1.for_dialect(D.CLICKHOUSE),
    # count_if
    base.AggCountIf(
        variants=[
            V(D.CLICKHOUSE, lambda cond: n.func.IFNULL(sa.func.countIf(cond), 0)),
        ]
    ),
    # countd
    base.AggCountd(
        variants=[
            V(D.CLICKHOUSE, sa.func.uniqExact),
        ]
    ),
    # countd_approx
    base.AggCountdApprox(
        variants=[
            V(D.CLICKHOUSE, sa.func.uniq),
        ]
    ),
    # countd_if
    base.AggCountdIf(
        variants=[
            V(D.CLICKHOUSE, sa.func.uniqExactIf),
        ]
    ),
    # max
    base.AggMax.for_dialect(D.CLICKHOUSE),
    # median
    base.AggMedian(
        variants=[
            V(D.CLICKHOUSE, sa.func.medianExact),
        ]
    ),
    # min
    base.AggMin.for_dialect(D.CLICKHOUSE),
    # quantile
    base.AggQuantile(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda expr, quant: chsq_clauses.ParametrizedFunction(
                    "quantileExact", [literal(quantile_value(un_literal(quant)), d=D.CLICKHOUSE)], expr
                ),
            ),
        ]
    ),
    # quantile_approx
    base.AggQuantileApproximate(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda expr, quant: chsq_clauses.ParametrizedFunction(
                    "quantile", [literal(quantile_value(un_literal(quant)), d=D.CLICKHOUSE)], expr
                ),
            ),
        ]
    ),
    # stdev
    base.AggStdev(
        variants=[
            V(D.CLICKHOUSE, sa.func.stddevSamp),
        ]
    ),
    # stdevp
    base.AggStdevp(
        variants=[
            V(D.CLICKHOUSE, sa.func.stddevPop),
        ]
    ),
    # sum
    base.AggSum.for_dialect(D.CLICKHOUSE),
    # sum_if
    base.AggSumIf(
        variants=[
            V(D.CLICKHOUSE, sa.func.sumIf),
        ]
    ),
    # top_concat
    base.AggTopConcat1(
        variants=[
            V(D.CLICKHOUSE, _top_concat_ch),
        ]
    ),
    base.AggTopConcat2(
        variants=[
            V(D.CLICKHOUSE, _top_concat_ch),
        ]
    ),
    # var
    base.AggVar(
        variants=[
            V(D.CLICKHOUSE, sa.func.varSamp),
        ]
    ),
    # varp
    base.AggVarp(
        variants=[
            V(D.CLICKHOUSE, sa.func.varPop),
        ]
    ),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/functions_aggregation.py
datalens-tech/datalens-backend
from clickhouse_sqlalchemy import types as ch_types
from clickhouse_sqlalchemy.ext.clauses import Lambda
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_array as base
from dl_formula.definitions.literals import un_literal

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make


DEFINITIONS_ARRAY = [
    # arr_avg
    base.FuncArrayAvg(
        variants=[
            V(D.and_above(D.CLICKHOUSE_21_8), sa.func.arrayAvg),
        ]
    ),
    # arr_max
    base.FuncArrayMaxFloat(
        variants=[
            V(D.and_above(D.CLICKHOUSE_21_8), sa.func.arrayMax),
        ]
    ),
    base.FuncArrayMaxInt(
        variants=[
            V(D.and_above(D.CLICKHOUSE_21_8), sa.func.arrayMax),
        ]
    ),
    # arr_min
    base.FuncArrayMinFloat(
        variants=[
            V(D.and_above(D.CLICKHOUSE_21_8), sa.func.arrayMin),
        ]
    ),
    base.FuncArrayMinInt(
        variants=[
            V(D.and_above(D.CLICKHOUSE_21_8), sa.func.arrayMin),
        ]
    ),
    # arr_product
    base.FuncArrayProduct(
        variants=[
            V(D.and_above(D.CLICKHOUSE_21_8), sa.func.arrayProduct),
        ]
    ),
    # arr_remove
    base.FuncArrayRemoveLiteralNull(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda array, value: sa.func.arrayMap(
                    Lambda(lambda x: sa.func.assumeNotNull(x)),
                    sa.func.arrayFilter(Lambda(lambda x: sa.func.isNotNull(x)), array),
                ),
            ),
        ]
    ),
    base.FuncArrayRemoveDefault(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda array, value: sa.func.arrayFilter(Lambda(lambda x: sa.func.ifNull(x != value, 1)), array),
            ),
        ]
    ),
    # arr_str
    base.FuncArrStr1(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda array: sa.func.arrayStringConcat(
                    sa.func.arrayMap(
                        Lambda(lambda x: sa.cast(x, ch_types.String)),
                        sa.func.arrayFilter(Lambda(lambda x: sa.func.isNotNull(x)), array),
                    ),
                    ",",
                ),
            ),
        ]
    ),
    base.FuncArrStr2(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda array, sep: sa.func.arrayStringConcat(
                    sa.func.arrayMap(
                        Lambda(lambda x: sa.cast(x, ch_types.String)),
                        sa.func.arrayFilter(Lambda(lambda x: sa.func.isNotNull(x)), array),
                    ),
                    sep,
                ),
            ),
        ]
    ),
    base.FuncArrStr3(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda array, sep, null_str: sa.func.arrayStringConcat(
                    sa.func.arrayMap(
                        Lambda(lambda x: sa.cast(sa.func.ifNull(sa.func.toString(x), null_str), ch_types.String)), array
                    ),
                    sep,
                ),
            ),
        ]
    ),
    # arr_sum
    base.FuncArraySumFloat(
        variants=[
            V(D.and_above(D.CLICKHOUSE_21_8), sa.func.arraySum),
        ]
    ),
    base.FuncArraySumInt(
        variants=[
            V(D.and_above(D.CLICKHOUSE_21_8), sa.func.arraySum),
        ]
    ),
    # array
    base.FuncConstArrayFloat.for_dialect(D.CLICKHOUSE),
    base.FuncConstArrayInt.for_dialect(D.CLICKHOUSE),
    base.FuncConstArrayStr.for_dialect(D.CLICKHOUSE),
    base.FuncNonConstArrayInt(
        variants=[
            V(D.CLICKHOUSE, lambda *args: sa.func.array(*args)),
        ]
    ),
    base.FuncNonConstArrayFloat(
        variants=[
            V(D.CLICKHOUSE, lambda *args: sa.func.array(*args)),
        ]
    ),
    base.FuncNonConstArrayStr(
        variants=[
            V(D.CLICKHOUSE, lambda *args: sa.func.array(*args)),
        ]
    ),
    # cast_arr_float
    base.FuncFloatArrayFromIntArray(
        variants=[
            V(D.CLICKHOUSE, lambda array: sa.func.arrayMap(Lambda(lambda x: sa.func.toFloat64(x)), array)),
        ]
    ),
    base.FuncFloatArrayFromFloatArray.for_dialect(D.CLICKHOUSE),
    base.FuncFloatArrayFromStringArray(
        variants=[
            V(D.CLICKHOUSE, lambda array: sa.func.arrayMap(Lambda(lambda x: sa.func.toFloat64OrNull(x)), array)),
        ]
    ),
    # cast_arr_int
    base.FuncIntArrayFromIntArray.for_dialect(D.CLICKHOUSE),
    base.FuncIntArrayFromFloatArray(
        variants=[
            V(D.CLICKHOUSE, lambda array: sa.func.arrayMap(Lambda(lambda x: sa.func.toInt64(sa.func.floor(x))), array)),
        ]
    ),
    base.FuncIntArrayFromStringArray(
        variants=[
            V(D.CLICKHOUSE, lambda array: sa.func.arrayMap(Lambda(lambda x: sa.func.toInt64OrNull(x)), array)),
        ]
    ),
    # cast_arr_str
    base.FuncStringArrayFromIntArray(
        variants=[
            V(D.CLICKHOUSE, lambda array: sa.func.arrayMap(Lambda(lambda x: sa.func.toString(x)), array)),
        ]
    ),
    base.FuncStringArrayFromFloatArray(
        variants=[
            V(D.CLICKHOUSE, lambda array: sa.func.arrayMap(Lambda(lambda x: sa.func.toString(x)), array)),
        ]
    ),
    base.FuncStringArrayFromStringArray.for_dialect(D.CLICKHOUSE),
    # contains
    base.FuncArrayContains(
        variants=[
            V(D.CLICKHOUSE, sa.func.has),
        ]
    ),
    # notcontains
    base.FuncArrayNotContains.for_dialect(D.CLICKHOUSE),
    # contains_all
    base.FuncArrayContainsAll(
        variants=[
            V(D.CLICKHOUSE, sa.func.hasAll),
        ]
    ),
    # contains_any
    base.FuncArrayContainsAny(
        variants=[
            V(D.CLICKHOUSE, sa.func.hasAny),
        ]
    ),
    # contains_subsequence
    base.FuncArrayContainsSubsequence(
        variants=[
            V(D.and_above(D.CLICKHOUSE_21_8), sa.func.hasSubstr),
        ]
    ),
    # count_item
    base.FuncArrayCountItemInt(
        variants=[
            V(D.CLICKHOUSE, sa.func.countEqual),
        ]
    ),
    base.FuncArrayCountItemFloat(
        variants=[
            V(D.CLICKHOUSE, sa.func.countEqual),
        ]
    ),
    base.FuncArrayCountItemStr(
        variants=[
            V(D.CLICKHOUSE, sa.func.countEqual),
        ]
    ),
    # get_item
    base.FuncGetArrayItemFloat(
        variants=[
            V(D.CLICKHOUSE, sa.func.arrayElement),
        ]
    ),
    base.FuncGetArrayItemInt(
        variants=[
            V(D.CLICKHOUSE, sa.func.arrayElement),
        ]
    ),
    base.FuncGetArrayItemStr(
        variants=[
            V(D.CLICKHOUSE, sa.func.arrayElement),
        ]
    ),
    # len
    base.FuncLenArray(
        variants=[
            V(D.CLICKHOUSE, sa.func.length),
        ]
    ),
    # replace
    base.FuncReplaceArrayLiteralNull(
        variants=[
            V(D.CLICKHOUSE, lambda array, old, new: sa.func.arrayMap(Lambda(lambda x: sa.func.ifNull(x, new)), array)),
        ]
    ),
    base.FuncReplaceArrayDefault(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda array, old, new: sa.func.arrayMap(Lambda(lambda x: sa.func.if_(x == old, new, x)), array),
            ),
        ]
    ),
    # slice
    base.FuncArraySlice(
        variants=[
            V(D.CLICKHOUSE, sa.func.arraySlice),
        ]
    ),
    # startswith
    base.FuncStartswithArrayConst(
        variants=[
            V(D.CLICKHOUSE, lambda x, y, _env: sa.func.arraySlice(x, 1, len(un_literal(y))) == y),
        ]
    ),
    base.FuncStartswithArrayNonConst(
        variants=[
            V(D.CLICKHOUSE, lambda x, y, _env: sa.func.arraySlice(x, 1, sa.func.length(y)) == y),
        ]
    ),
    # unnest
    base.FuncUnnestArrayFloat(
        variants=[
            V(D.CLICKHOUSE, sa.func.arrayJoin),
        ]
    ),
    base.FuncUnnestArrayInt(
        variants=[
            V(D.CLICKHOUSE, sa.func.arrayJoin),
        ]
    ),
    base.FuncUnnestArrayStr(
        variants=[
            V(D.CLICKHOUSE, sa.func.arrayJoin),
        ]
    ),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/functions_array.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Union,
)

import clickhouse_sqlalchemy.types as ch_types
import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

from dl_formula.definitions.base import (
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.common_datetime import (
    ch_date_with_tz,
    datetime_interval,
    make_ch_tz_args,
)
import dl_formula.definitions.functions_datetime as base
from dl_formula.definitions.literals import un_literal

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


if TYPE_CHECKING:
    from dl_formula.translation.context import TranslationCtx


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


def _datetrunc2_ch_impl_19_3_3(
    date_ctx: TranslationCtx,
    unit_ctx: TranslationCtx,
) -> ClauseElement:
    return _datetrunc3_ch_impl_19_3_3(date_ctx, unit_ctx, num_ctx=1)


def _datetrunc3_ch_impl_19_3_3(
    date_ctx: TranslationCtx,
    unit_ctx: TranslationCtx,
    num_ctx: Union[TranslationCtx, int],
) -> ClauseElement:
    date_expr = date_ctx.expression
    unit = base.norm_datetrunc_unit(unit_ctx.expression)
    # for plugging the datetrunc2 here:
    num = num_ctx if isinstance(num_ctx, int) else un_literal(num_ctx.expression)
    tz_args, output_tz_args = make_ch_tz_args(date_ctx)
    if unit in {"year", "quarter", "month", "week"}:
        return sa.func.toDateTime(
            sa.func.toStartOfInterval(date_expr, datetime_interval(unit, num, caps=False), *tz_args), *output_tz_args
        )
    if unit in {"day", "hour", "minute", "second"}:
        return sa.func.toStartOfInterval(date_expr, datetime_interval(unit, num, caps=False), *tz_args)

    # This normally should not happen
    raise NotImplementedError(f"Unsupported unit {unit}")


DEFINITIONS_DATETIME = [
    # dateadd
    base.FuncDateadd1.for_dialect(D.CLICKHOUSE),
    base.FuncDateadd2Unit.for_dialect(D.CLICKHOUSE),
    base.FuncDateadd2Number.for_dialect(D.CLICKHOUSE),
    base.FuncDateadd3Legacy.for_dialect(D.CLICKHOUSE),
    base.FuncDateadd3DateNonConstNum(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda date, what, num: sa.cast(
                    date + datetime_interval(what.value, sa.func.ifNull(num, 0), ch_func=True),
                    ch_types.Nullable(ch_types.Date),
                ),
            ),
        ]
    ),
    base.FuncDateadd3DatetimeNonConstNum(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda dt, what, num: (
                    sa.func.toDateTime(dt + datetime_interval(what.value, sa.func.ifNull(num, 0), ch_func=True), "UTC")
                ),
            ),
        ]
    ),
    base.FuncDateadd3GenericDatetimeNonConstNum(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda dt, what, num: (dt + datetime_interval(what.value, sa.func.ifNull(num, 0), ch_func=True)),
            ),
        ]
    ),
    base.FuncDateadd3DatetimeTZNonConstNum(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda dt, what, num: (dt + datetime_interval(what.value, sa.func.ifNull(num, 0), ch_func=True)),
            ),
        ]
    ),
    # datepart
    base.FuncDatepart2Legacy.for_dialect(D.CLICKHOUSE),
    base.FuncDatepart2.for_dialect(D.CLICKHOUSE),
    base.FuncDatepart3Const.for_dialect(D.CLICKHOUSE),
    base.FuncDatepart3NonConst.for_dialect(D.CLICKHOUSE),
    # datetrunc
    base.FuncDatetrunc2Date(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda date, unit: (
                    sa.func.toStartOfInterval(sa.func.toDate(date), datetime_interval(unit.value, 1, caps=False))
                    if base.norm_datetrunc_unit(unit) in {"year", "quarter", "month", "week"}
                    else sa.func.toDate(date)
                ),
            ),
        ]
    ),
    base.FuncDatetrunc2Datetime(
        variants=[
            VW(D.CLICKHOUSE, _datetrunc2_ch_impl_19_3_3),
        ]
    ),
    base.FuncDatetrunc2DatetimeTZ(
        variants=[
            VW(D.CLICKHOUSE, _datetrunc2_ch_impl_19_3_3),
        ]
    ),
    base.FuncDatetrunc3Date(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda date, unit, num: (
                    sa.func.toStartOfInterval(
                        sa.func.toDate(date), datetime_interval(unit.value, num.value, caps=False)
                    )
                    if base.norm_datetrunc_unit(unit) in {"year", "quarter", "month", "week"}
                    else sa.func.toDate(
                        sa.func.toStartOfInterval(
                            sa.func.toDate(date), datetime_interval(unit.value, num.value, caps=False)
                        )
                    )
                    if base.norm_datetrunc_unit(unit) == "day"
                    else sa.func.toDate(date)
                ),
            ),
        ]
    ),
    base.FuncDatetrunc3Datetime(
        variants=[
            VW(D.CLICKHOUSE, _datetrunc3_ch_impl_19_3_3),
        ]
    ),
    # day
    base.FuncDay(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toDayOfMonth(*ch_date_with_tz(date_ctx))),
        ]
    ),
    base.FuncDayDatetimeTZ(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toDayOfMonth(*ch_date_with_tz(date_ctx))),
        ]
    ),
    # dayofweek
    base.FuncDayofweek1.for_dialect(D.CLICKHOUSE),
    base.FuncDayofweek2(
        variants=[
            VW(
                D.CLICKHOUSE,
                lambda date_ctx, firstday_ctx: base.dow_firstday_shift(
                    sa.func.toDayOfWeek(*ch_date_with_tz(date_ctx)), firstday_ctx.expression
                ),
            ),
        ]
    ),
    base.FuncDayofweek2TZ(
        variants=[
            VW(
                D.CLICKHOUSE,
                lambda date_ctx, firstday_ctx: base.dow_firstday_shift(
                    sa.func.toDayOfWeek(*ch_date_with_tz(date_ctx)), firstday_ctx.expression
                ),
            ),
        ]
    ),
    # genericnow
    base.FuncGenericNow(
        variants=[
            V(D.CLICKHOUSE, lambda: sa.func.toDateTime(sa.func.now())),
        ]
    ),
    # hour
    base.FuncHourDate.for_dialect(D.CLICKHOUSE),
    base.FuncHourDatetime(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toHour(*ch_date_with_tz(date_ctx))),
        ]
    ),
    base.FuncHourDatetimeTZ(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toHour(*ch_date_with_tz(date_ctx))),
        ]
    ),
    # minute
    base.FuncMinuteDate.for_dialect(D.CLICKHOUSE),
    base.FuncMinuteDatetime(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toMinute(*ch_date_with_tz(date_ctx))),
        ]
    ),
    base.FuncMinuteDatetimeTZ(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toMinute(*ch_date_with_tz(date_ctx))),
        ]
    ),
    # month
    base.FuncMonth(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toMonth(*ch_date_with_tz(date_ctx))),
        ]
    ),
    base.FuncMonthDatetimeTZ(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toMonth(*ch_date_with_tz(date_ctx))),
        ]
    ),
    # now
    base.FuncNow(
        variants=[
            V(D.CLICKHOUSE, lambda: sa.func.toDateTime(sa.func.now())),
        ]
    ),
    # quarter
    base.FuncQuarter(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toQuarter(*ch_date_with_tz(date_ctx))),
        ]
    ),
    base.FuncQuarterDatetimeTZ(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toQuarter(*ch_date_with_tz(date_ctx))),
        ]
    ),
    # second
    base.FuncSecondDate.for_dialect(D.CLICKHOUSE),
    base.FuncSecondDatetime(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toSecond(*ch_date_with_tz(date_ctx))),
        ]
    ),
    base.FuncSecondDatetimeTZ(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toSecond(*ch_date_with_tz(date_ctx))),
        ]
    ),
    # today
    base.FuncToday(
        variants=[
            V(D.CLICKHOUSE, lambda: sa.func.toDate(sa.func.now())),
        ]
    ),
    # week
    base.FuncWeek(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toISOWeek(*ch_date_with_tz(date_ctx))),
        ]
    ),
    base.FuncWeekTZ(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toISOWeek(*ch_date_with_tz(date_ctx))),
        ]
    ),
    # year
    base.FuncYear(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toYear(*ch_date_with_tz(date_ctx))),
        ]
    ),
    base.FuncYearDatetimeTZ(
        variants=[
            VW(D.CLICKHOUSE, lambda date_ctx: sa.func.toYear(*ch_date_with_tz(date_ctx))),
        ]
    ),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/functions_datetime.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import (
    TranslationVariant,
    TranslationVariantWrapped,
)
import dl_formula.definitions.functions_logical as base
from dl_formula.shortcuts import n

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


DEFINITIONS_LOGICAL = [
    # case
    base.FuncCase.for_dialect(D.CLICKHOUSE),
    # if
    base.FuncIf.for_dialect(D.CLICKHOUSE),
    # ifnan
    base.FuncIfnan(
        variants=[
            VW(
                D.CLICKHOUSE,
                lambda check_value, alt_value: n.func.IF(n.func.ISNAN(check_value), alt_value, check_value),
            ),
        ]
    ),
    # ifnull
    base.FuncIfnull(
        variants=[
            V(D.CLICKHOUSE, sa.func.ifNull),
        ]
    ),
    # iif
    base.FuncIif3Legacy.for_dialect(D.CLICKHOUSE),
    # isnan
    base.FuncIsnan(
        variants=[
            V(D.CLICKHOUSE, sa.func.isNaN),
        ]
    ),
    # isnull
    base.FuncIsnull(
        variants=[
            V(D.CLICKHOUSE, sa.func.isNull),
        ]
    ),
    # zn
    base.FuncZn(
        variants=[
            V(D.CLICKHOUSE, lambda x: sa.func.ifNull(x, 0)),
        ]
    ),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/functions_logical.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_markup as base

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


DEFINITIONS_MARKUP = [
    # +
    base.BinaryPlusMarkup.for_dialect(D.CLICKHOUSE),
    # __str
    base.FuncInternalStrConst.for_dialect(D.CLICKHOUSE),
    base.FuncInternalStr.for_dialect(D.CLICKHOUSE),
    # bold
    base.FuncBold.for_dialect(D.CLICKHOUSE),
    # italic
    base.FuncItalics.for_dialect(D.CLICKHOUSE),
    # markup
    base.ConcatMultiMarkup.for_dialect(D.CLICKHOUSE),
    # url
    base.FuncUrl.for_dialect(D.CLICKHOUSE),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/functions_markup.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    TranslationVariant,
    TranslationVariantWrapped,
)
import dl_formula.definitions.functions_math as base
from dl_formula.shortcuts import n

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class FuncGreatestCH(base.FuncGreatestBase):
    variants = [V(D.CLICKHOUSE, lambda x, y: getattr(sa.func, "if")(x >= y, x, y))]
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME, DataType.DATETIME]),
        ArgTypeSequence([DataType.DATETIMETZ, DataType.DATETIMETZ]),
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.GENERICDATETIME]),
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
        ArgTypeSequence([DataType.BOOLEAN, DataType.BOOLEAN]),
    ]


class FuncGreatestNumbersCH(base.FuncGreatestBase):
    variants = [
        V(D.CLICKHOUSE, sa.func.greatest),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]


class FuncLeastCH(base.FuncLeastBase):
    variants = [
        V(D.CLICKHOUSE, sa.func.least),
    ]
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME, DataType.DATETIME]),
        ArgTypeSequence([DataType.DATETIMETZ, DataType.DATETIMETZ]),
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.GENERICDATETIME]),
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
        ArgTypeSequence([DataType.BOOLEAN, DataType.BOOLEAN]),
    ]


class FuncLeastNumbersCH(base.FuncLeastBase):
    variants = [
        V(D.CLICKHOUSE, sa.func.least),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]


DEFINITIONS_MATH = [
    # abs
    base.FuncAbs.for_dialect(D.CLICKHOUSE),
    # acos
    base.FuncAcos.for_dialect(D.CLICKHOUSE),
    # asin
    base.FuncAsin.for_dialect(D.CLICKHOUSE),
    # atan
    base.FuncAtan.for_dialect(D.CLICKHOUSE),
    # atan2
    base.FuncAtan2(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda x, y: n.if_(
                    n.if_(sa.and_(x == 0, y == 0)).then(0),  # type: ignore  # TODO: fix
                    n.if_(sa.and_(y == 0, x < 0)).then(-sa.func.pi() / 2),  # type: ignore  # TODO: fix
                    n.if_(sa.and_(y == 0, x > 0)).then(sa.func.pi() / 2),  # type: ignore  # TODO: fix
                    n.if_(sa.and_(x == 0, y < 0)).then(sa.func.pi()),  # type: ignore  # TODO: fix
                    n.if_(sa.and_(x == 0, y > 0)).then(0),  # type: ignore  # TODO: fix
                    n.if_(y > 0).then(sa.func.atan(x / y)),  # type: ignore  # TODO: fix
                    n.if_(x > 0).then(sa.func.atan(x / y) + sa.func.pi()),  # type: ignore  # TODO: fix
                ).else_(
                    sa.func.atan(x / y) - sa.func.pi()  # for x < 0
                ),
            ),
        ]
    ),
    # ceiling
    base.FuncCeiling.for_dialect(D.CLICKHOUSE),
    # compare
    base.FuncCompare(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda left, right, epsilon: sa.func.multiIf(
                    sa.func.lessOrEquals(sa.func.abs(left - right), epsilon),
                    0,
                    sa.func.less(left, right),
                    -1,
                    1,
                ),
            ),
        ]
    ),
    # cos
    base.FuncCos.for_dialect(D.CLICKHOUSE),
    # cot
    base.FuncCot(
        variants=[
            V(D.CLICKHOUSE, lambda x: sa.func.cos(x) / sa.func.sin(x)),
        ]
    ),
    # degrees
    base.FuncDegrees(
        variants=[
            V(D.CLICKHOUSE, lambda x: x / sa.func.pi() * 180),
        ]
    ),
    # div
    base.FuncDivBasic(
        variants=[
            V(D.CLICKHOUSE, sa.func.intDiv),
        ]
    ),
    # div_safe
    base.FuncDivSafe2(
        variants=[
            V(D.CLICKHOUSE, lambda x, y: sa.func.IF(y != 0, sa.func.intDivOrZero(x, y), None)),
        ]
    ),
    base.FuncDivSafe3(
        variants=[
            V(D.CLICKHOUSE, lambda x, y, default: sa.func.IF(y != 0, sa.func.intDivOrZero(x, y), default)),
        ]
    ),
    # exp
    base.FuncExp.for_dialect(D.CLICKHOUSE),
    # fdiv_safe
    base.FuncFDivSafe2.for_dialect(D.CLICKHOUSE),
    base.FuncFDivSafe3.for_dialect(D.CLICKHOUSE),
    # floor
    base.FuncFloor.for_dialect(D.CLICKHOUSE),
    # greatest
    base.FuncGreatest1.for_dialect(D.CLICKHOUSE),
    FuncGreatestCH(),
    FuncGreatestNumbersCH(),
    base.GreatestMulti.for_dialect(D.CLICKHOUSE),
    # least
    base.FuncLeast1.for_dialect(D.CLICKHOUSE),
    FuncLeastCH(),
    FuncLeastNumbersCH(),
    base.LeastMulti.for_dialect(D.CLICKHOUSE),
    # ln
    base.FuncLn(
        variants=[
            V(D.CLICKHOUSE, sa.func.log),
        ]
    ),
    # log
    base.FuncLog(
        variants=[
            V(D.CLICKHOUSE, lambda x, y: sa.func.log(x) / sa.func.log(y)),
        ]
    ),
    # log10
    base.FuncLog10.for_dialect(D.CLICKHOUSE),
    # pi
    base.FuncPi.for_dialect(D.CLICKHOUSE),
    # power
    base.FuncPower(
        variants=[
            V(D.CLICKHOUSE, sa.func.pow),
        ]
    ),
    # radians
    base.FuncRadians(
        variants=[
            V(D.CLICKHOUSE, lambda x: x * sa.func.pi() / 180),
        ]
    ),
    # round
    base.FuncRound1.for_dialect(D.CLICKHOUSE),
    base.FuncRound2.for_dialect(D.CLICKHOUSE),
    # sign
    base.FuncSign(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda x: n.if_(
                    n.if_(x < 0).then(-1),  # type: ignore  # TODO: fix
                    n.if_(x > 0).then(1),  # type: ignore  # TODO: fix
                ).else_(0),
            ),
        ]
    ),
    # sin
    base.FuncSin.for_dialect(D.CLICKHOUSE),
    # sqrt
    base.FuncSqrt.for_dialect(D.CLICKHOUSE),
    # square
    base.FuncSquare(
        variants=[
            V(D.CLICKHOUSE, lambda x: sa.func.pow(x, 2)),
        ]
    ),
    # tan
    base.FuncTan.for_dialect(D.CLICKHOUSE),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/functions_math.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_special as base

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make


DEFINITIONS_SPECIAL = [
    # _make_nan
    base.FuncMakeNan(
        variants=[
            V(D.CLICKHOUSE, lambda: sa.column("cast('NaN' AS Float64)", is_literal=True)),
        ]
    ),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/functions_special.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_string as base
from dl_formula.shortcuts import n

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make


DEFINITIONS_STRING = [
    # ascii
    base.FuncAscii(
        variants=[
            V(D.CLICKHOUSE, sa.func.reinterpretAsUInt8),
        ]
    ),
    # char
    base.FuncChar(
        variants=[
            V(D.CLICKHOUSE, sa.func.reinterpretAsString),
        ]
    ),
    # concat
    base.Concat1.for_dialect((D.CLICKHOUSE)),
    base.ConcatMultiStrConst.for_dialect(D.CLICKHOUSE),
    base.ConcatMultiStr(
        variants=[
            V(D.CLICKHOUSE, lambda *args: sa.func.concat(*args)),
        ]
    ),
    base.ConcatMultiAny.for_dialect(D.CLICKHOUSE),
    # contains
    base.FuncContainsConst(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda x, y: x.like(
                    "%{}%".format(base.quote_like(y.value)),
                ),
            ),
        ]
    ),
    base.FuncContainsNonConst(
        variants=[
            V(D.CLICKHOUSE, lambda x, y: sa.func.position(x, y) > 0),
        ]
    ),
    base.FuncContainsNonString.for_dialect(D.CLICKHOUSE),
    # notcontains
    base.FuncNotContainsConst.for_dialect(D.CLICKHOUSE),
    base.FuncNotContainsNonConst.for_dialect(D.CLICKHOUSE),
    base.FuncNotContainsNonString.for_dialect(D.CLICKHOUSE),
    # endswith
    base.FuncEndswithNonConst(
        variants=[
            V(D.CLICKHOUSE, lambda x, y: sa.func.endsWith(x, y)),
        ]
    ),
    base.FuncEndswithNonString.for_dialect(D.CLICKHOUSE),
    # find
    base.FuncFind2(
        variants=[
            V(D.CLICKHOUSE, sa.func.positionUTF8),
        ]
    ),
    base.FuncFind3(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda x, y, z: n.if_(  # type: ignore  # TODO: fix
                    sa.func.positionUTF8(sa.func.substringUTF8(x, z, sa.func.lengthUTF8(x)), y) > 0
                )
                .then(sa.func.positionUTF8(sa.func.substringUTF8(x, z, sa.func.lengthUTF8(x)), y) + z - 1)
                .else_(0),
            ),
        ]
    ),
    # icontains
    base.FuncIContainsConst(
        variants=[
            V(D.CLICKHOUSE, lambda x, y: sa.func.positionCaseInsensitiveUTF8(x, y) != 0),
        ]
    ),
    base.FuncIContainsNonConst(
        variants=[
            V(D.CLICKHOUSE, lambda x, y: sa.func.positionCaseInsensitiveUTF8(x, y) != 0),
        ]
    ),
    base.FuncIContainsNonString.for_dialect(D.CLICKHOUSE),
    # iendswith
    base.FuncIEndswithNonConst.for_dialect(D.CLICKHOUSE),
    base.FuncIEndswithNonString.for_dialect(D.CLICKHOUSE),
    # istartswith
    base.FuncIStartswithNonConst.for_dialect(D.CLICKHOUSE),
    base.FuncIStartswithNonString.for_dialect(D.CLICKHOUSE),
    # left
    base.FuncLeft(
        variants=[
            V(D.CLICKHOUSE, lambda x, y: sa.func.substringUTF8(x, 1, y)),
        ]
    ),
    # len
    base.FuncLenString(
        variants=[
            V(D.CLICKHOUSE, sa.func.lengthUTF8),
        ]
    ),
    # lower
    base.FuncLowerConst.for_dialect(D.CLICKHOUSE),
    base.FuncLowerNonConst(
        variants=[
            V(D.CLICKHOUSE, sa.func.lowerUTF8),
        ]
    ),
    # ltrim
    base.FuncLtrim(
        variants=[
            V(D.CLICKHOUSE, sa.func.trimLeft),
        ]
    ),
    # regexp_extract
    base.FuncRegexpExtract(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda text, pattern: sa.func.arrayElement(sa.func.extractAll(sa.func.assumeNotNull(text), pattern), 1),
            ),
        ]
    ),
    # regexp_extract_nth
    base.FuncRegexpExtractNth(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda text, pattern, ind: sa.func.arrayElement(
                    sa.func.extractAll(sa.func.assumeNotNull(text), pattern), ind
                ),
            ),
        ]
    ),
    # regexp_match
    base.FuncRegexpMatch(
        variants=[
            V(D.CLICKHOUSE, sa.func.match),
        ]
    ),
    # regexp_replace
    base.FuncRegexpReplace(
        variants=[
            V(D.CLICKHOUSE, sa.func.replaceRegexpAll),
        ]
    ),
    # replace
    base.FuncReplace(
        variants=[
            V(D.CLICKHOUSE, sa.func.replaceAll),
        ]
    ),
    # right
    base.FuncRight(
        variants=[
            V(D.CLICKHOUSE, lambda x, y: sa.func.substringUTF8(x, sa.func.lengthUTF8(x) - y + 1, y)),
        ]
    ),
    # rtrim
    base.FuncRtrim(
        variants=[
            V(D.CLICKHOUSE, sa.func.trimRight),
        ]
    ),
    # space
    base.FuncSpaceConst.for_dialect(D.CLICKHOUSE),
    base.FuncSpaceNonConst(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda size: sa.func.arrayStringConcat(sa.func.arrayResize(sa.func.emptyArrayString(), size, " ")),
            ),
        ]
    ),
    # split
    base.FuncSplit1(
        variants=[
            V(D.CLICKHOUSE, lambda text: sa.func.splitByString(",", sa.func.assumeNotNull(text))),
        ]
    ),
    base.FuncSplit2(
        variants=[
            V(D.CLICKHOUSE, lambda text, delim: sa.func.splitByString(delim, sa.func.assumeNotNull(text))),
        ]
    ),
    base.FuncSplit3(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda text, delim, ind: sa.func.arrayElement(
                    sa.func.splitByString(delim, sa.func.assumeNotNull(text)), ind
                ),
            ),
        ]
    ),
    # startswith
    base.FuncStartswithNonConst(
        variants=[
            V(D.CLICKHOUSE, lambda x, y: sa.func.startsWith(x, y)),
        ]
    ),
    base.FuncStartswithNonString.for_dialect(D.CLICKHOUSE),
    # substr
    base.FuncSubstr2(
        variants=[
            V(D.CLICKHOUSE, lambda text, start: sa.func.substringUTF8(text, start, sa.func.lengthUTF8(text))),
        ]
    ),
    base.FuncSubstr3(
        variants=[
            V(D.CLICKHOUSE, lambda text, start, length: sa.func.substringUTF8(text, start, length)),
        ]
    ),
    # trim
    base.FuncTrim(
        variants=[
            V(D.CLICKHOUSE, sa.func.trimBoth),
        ]
    ),
    # upper
    base.FuncUpperConst.for_dialect(D.CLICKHOUSE),
    base.FuncUpperNonConst(
        variants=[
            V(D.CLICKHOUSE, sa.func.upperUTF8),
        ]
    ),
    # utf8
    base.FuncUtf8(
        variants=[
            V(D.CLICKHOUSE, lambda text, charset: sa.func.convertCharset(text, charset, "UTF-8")),
        ]
    ),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/functions_string.py
datalens-tech/datalens-backend
from typing import Sequence

import clickhouse_sqlalchemy.types as ch_types
import sqlalchemy as sa
from sqlalchemy.types import TypeEngine

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    Function,
    SingleVariantTranslationBase,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.common_datetime import (
    ch_date_with_tz,
    ensure_naive_first_arg,
)
import dl_formula.definitions.functions_type as base
from dl_formula.definitions.scope import Scope
from dl_formula.shortcuts import n
from dl_formula.translation.context import TranslationCtx

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class FuncDate2FromNumber(base.FuncDate2):
    variants = [
        V(D.CLICKHOUSE, lambda expr, tz: sa.func.toDate(expr, tz)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.CONST_STRING]),
        ArgTypeSequence([DataType.INTEGER, DataType.CONST_STRING]),
    ]


class FuncDatetimeTZCH(SingleVariantTranslationBase, base.FuncDatetimeTZ):
    dialects = D.CLICKHOUSE
    argument_types = [
        ArgTypeSequence(
            [
                {
                    DataType.DATETIME,
                    DataType.GENERICDATETIME,
                    DataType.DATETIMETZ,
                    # Not even tried: `DataType.DATE`.
                    DataType.INTEGER,
                    DataType.FLOAT,
                    DataType.STRING,
                },
                DataType.CONST_STRING,
            ]
        ),
    ]

    @classmethod
    def _translate_main(cls, value_ctx, tz_ctx):  # type: ignore
        # NOTE: Tricky point:
        # in CH, converting all aware datetimes to UTC,
        # primarily for the correct output.
        # In particular datetimetz-taking functions such as `DATE()`, *MUST*
        # convert to the `data_type_params.timezone` in CH before applying
        # the CH function.
        # (similarly, in PG, there's only UTC tz-aware datetimes)
        # Reminder: tz-naive datetimes in CH should also be, effectively, in
        # UTC (as there's no proper tz-naive datetimes in CH).
        expr = value_ctx.expression
        if value_ctx.data_type == DataType.DATETIME:
            # Have to pass the naive-datetime values through a string,
            # otherwise CH uses the internal TZ knowledge to shift the value,
            # rather than re-interpret it at the specified timezone.
            expr = sa.func.formatDateTime(expr, "%Y-%m-%d %H:%M:%S")
        expr = sa.func.toDateTime(expr, tz_ctx.expression)  # 'interpret in', effectively
        expr = sa.func.toDateTime(expr, "UTC")
        return expr


class FuncDatetimeTZToNaiveCH(base.FuncDatetimeTZToNaive):
    dialects = D.CLICKHOUSE

    @classmethod
    def _translate_main(cls, value_ctx):  # type: ignore
        expr = value_ctx.expression
        tz = value_ctx.data_type_params.timezone
        assert tz
        # use cases: datepart, grouping (tz transitions should actually group into one value)
        expr = sa.func.toDateTime(expr, tz)
        # TODO: find a better way to tz-shift the value in CH: https://github.com/ClickHouse/ClickHouse/issues/19768
        expr = sa.func.toDateTime(sa.func.formatDateTime(expr, "%Y-%m-%d %H:%M:%S"), "UTC")
        return expr


# Note: `SingleVariantTranslationBase` here essentially acts as a mixin, providing
# custom `get_variants` implementation that plugs into `cls.dialects` and `cls._translate_main`.
class FuncTypeGenericDatetime2CHImpl(SingleVariantTranslationBase, base.FuncTypeGenericDatetime2Impl):
    dialects = D.CLICKHOUSE
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.DATETIME, DataType.GENERICDATETIME, DataType.INTEGER, DataType.FLOAT, DataType.STRING},
                DataType.CONST_STRING,
            ]
        ),
    ]

    @classmethod
    def _translate_main(cls, expr, tz):  # type: ignore
        return sa.func.toDateTime(expr.expression, tz.expression)


class FuncDatetime2CH(FuncTypeGenericDatetime2CHImpl):
    name = "datetime"


class FuncGenericDatetime2CH(FuncTypeGenericDatetime2CHImpl):
    name = "genericdatetime"
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncDbCastClickHouseBase(base.FuncDbCastBase):
    WHITELISTS = {
        D.CLICKHOUSE: {
            DataType.INTEGER: [
                base.WhitelistTypeSpec(name="Int8", sa_type=ch_types.Int8),
                base.WhitelistTypeSpec(name="Int16", sa_type=ch_types.Int16),
                base.WhitelistTypeSpec(name="Int32", sa_type=ch_types.Int32),
                base.WhitelistTypeSpec(name="Int64", sa_type=ch_types.Int64),
                base.WhitelistTypeSpec(name="UInt8", sa_type=ch_types.UInt8),
                base.WhitelistTypeSpec(name="UInt16", sa_type=ch_types.UInt16),
                base.WhitelistTypeSpec(name="UInt32", sa_type=ch_types.UInt32),
                base.WhitelistTypeSpec(name="UInt64", sa_type=ch_types.UInt64),
            ],
            DataType.FLOAT: [
                base.WhitelistTypeSpec(name="Float32", sa_type=ch_types.Float32),
                base.WhitelistTypeSpec(name="Float64", sa_type=ch_types.Float64),
                base.WhitelistTypeSpec(name="Decimal", sa_type=ch_types.Decimal, arg_types=base.DECIMAL_CAST_ARG_T),
            ],
            DataType.STRING: [
                base.WhitelistTypeSpec(name="String", sa_type=ch_types.String),
                # TODO: FixedString
            ],
            DataType.DATE: [
                base.WhitelistTypeSpec(name="Date", sa_type=ch_types.Date),
                base.WhitelistTypeSpec(name="Date32", sa_type=ch_types.Date32),
            ],
        }
    }

    @classmethod
    def generate_cast_type(
        cls, dialect: DialectCombo, wr_name: TranslationCtx, value: TranslationCtx, type_args: Sequence[TranslationCtx]
    ) -> TypeEngine:
        # Add Nullable wrapper
        type_ = super().generate_cast_type(dialect=dialect, wr_name=wr_name, value=value, type_args=type_args)
        type_ = ch_types.Nullable(type_)
        return type_


class FuncDbCastClickHouse2(FuncDbCastClickHouseBase, base.FuncDbCast2):
    pass


class FuncDbCastClickHouse3(FuncDbCastClickHouseBase, base.FuncDbCast3):
    pass


class FuncDbCastClickHouse4(FuncDbCastClickHouseBase, base.FuncDbCast4):
    pass


DEFINITIONS_TYPE = [
    # bool
    base.FuncBoolFromNull(
        variants=[
            V(D.CLICKHOUSE, lambda _: sa.func.toInt8(sa.null())),
        ]
    ),
    base.FuncBoolFromNumber.for_dialect(D.CLICKHOUSE),
    base.FuncBoolFromBool.for_dialect(D.CLICKHOUSE),
    base.FuncBoolFromStrGeo.for_dialect(D.CLICKHOUSE),
    base.FuncBoolFromDateDatetime.for_dialect(D.CLICKHOUSE),
    # date
    base.FuncDate1FromNull(
        variants=[
            V(D.CLICKHOUSE, lambda _: sa.func.toDate(sa.null())),
        ]
    ),
    base.FuncDate1FromDatetime(
        variants=[
            VW(D.CLICKHOUSE, lambda value: sa.func.toDate(value.expression)),
        ]
    ),
    base.FuncDate1FromDatetimeTZ(
        variants=[
            VW(D.CLICKHOUSE, lambda value_ctx: sa.func.toDate(*ch_date_with_tz(value_ctx))),
        ]
    ),
    base.FuncDate1FromString(
        variants=[
            V(D.CLICKHOUSE, lambda expr: sa.cast(expr, ch_types.Nullable(ch_types.Date))),
            # V(D.CLICKHOUSE, lambda expr: sa.func.toDateOrNull(expr)),  FIXME: depends on ClickHouse/issues/9584
        ]
    ),
    base.FuncDate1FromNumber(
        variants=[
            V(D.CLICKHOUSE, lambda expr: sa.func.toDate(expr)),
        ]
    ),
    FuncDate2FromNumber(),
    # date_parse
    base.FuncDateParse(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.toDate(sa.func.parseDateTimeBestEffortOrNull(value))),
        ]
    ),
    # datetime
    base.FuncDatetime1FromNull(
        variants=[
            V(D.CLICKHOUSE, lambda _: sa.func.toDateTime(sa.null())),
        ]
    ),
    base.FuncDatetime1FromDatetime.for_dialect(D.CLICKHOUSE),
    base.FuncDatetime1FromDate(
        variants=[
            V(D.CLICKHOUSE, lambda expr: sa.func.toDateTime(expr)),
        ]
    ),
    base.FuncDatetime1FromNumber(
        variants=[
            V(D.CLICKHOUSE, lambda expr: sa.func.toDateTime(expr)),
        ]
    ),
    base.FuncDatetime1FromString(
        variants=[
            V(D.CLICKHOUSE, lambda expr: sa.func.toDateTimeOrNull(expr)),
        ]
    ),
    FuncDatetime2CH(),
    # datetime_parse
    base.FuncDatetimeParse(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.parseDateTimeBestEffortOrNull(value)),
        ]
    ),
    # datetimetz
    base.FuncDatetimeTZConst.for_dialect(D.CLICKHOUSE),
    FuncDatetimeTZCH(),
    # datetimetz_to_naive
    FuncDatetimeTZToNaiveCH(),
    # db_cast
    FuncDbCastClickHouse2(),
    FuncDbCastClickHouse3(),
    FuncDbCastClickHouse4(),
    # float
    base.FuncFloatNumber(
        variants=[
            V(D.CLICKHOUSE, sa.func.toFloat64),
        ]
    ),
    base.FuncFloatString(
        variants=[
            V(D.CLICKHOUSE, sa.func.toFloat64OrNull),
        ]
    ),
    base.FuncFloatFromBool(
        variants=[
            V(D.CLICKHOUSE, sa.func.toFloat64),
        ]
    ),
    base.FuncFloatFromDate(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda value: sa.func.toFloat64(sa.func.toUnixTimestamp(sa.func.toDateTime(value, "UTC"), "UTC")),
            ),
        ]
    ),
    base.FuncFloatFromDatetime(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.toFloat64(sa.func.toUnixTimestamp(value, "UTC"))),
        ]
    ),
    base.FuncFloatFromGenericDatetime(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.toFloat64(sa.func.toUnixTimestamp(value))),
        ]
    ),
    # genericdatetime
    base.FuncGenericDatetime1FromNull(
        variants=[
            V(D.CLICKHOUSE, lambda _: sa.func.toDateTime(sa.null())),
        ]
    ),
    base.FuncGenericDatetime1FromDatetime.for_dialect(D.CLICKHOUSE),
    base.FuncGenericDatetime1FromDate(
        variants=[
            V(D.CLICKHOUSE, lambda expr: sa.func.toDateTime(expr)),
        ]
    ),
    base.FuncGenericDatetime1FromNumber(
        variants=[
            V(D.CLICKHOUSE, lambda expr: sa.func.toDateTime(expr)),
        ]
    ),
    base.FuncGenericDatetime1FromString(
        variants=[
            V(D.CLICKHOUSE, lambda expr: sa.func.toDateTimeOrNull(expr)),
        ]
    ),
    FuncGenericDatetime2CH(),
    # genericdatetime_parse
    base.FuncGenericDatetimeParse(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.parseDateTimeBestEffortOrNull(value)),
        ]
    ),
    # geopoint
    base.FuncGeopointFromStr.for_dialect(D.CLICKHOUSE),
    base.FuncGeopointFromCoords.for_dialect(D.CLICKHOUSE),
    # geopolygon
    base.FuncGeopolygon.for_dialect(D.CLICKHOUSE),
    # int
    base.FuncIntFromNull(
        variants=[
            V(D.CLICKHOUSE, lambda _: sa.func.toInt64(sa.null())),
        ]
    ),
    base.FuncIntFromInt.for_dialect(D.CLICKHOUSE),
    base.FuncIntFromFloat(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.toInt64(sa.func.floor(value))),
        ]
    ),
    base.FuncIntFromBool(
        variants=[
            V(D.CLICKHOUSE, lambda value: value),
        ]
    ),
    base.FuncIntFromStr(
        variants=[
            V(D.CLICKHOUSE, sa.func.toInt64OrNull),
        ]
    ),
    base.FuncIntFromDate(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.toUnixTimestamp(sa.func.toDateTime(value, "UTC"), "UTC")),
        ]
    ),
    base.FuncIntFromDatetime(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.toUnixTimestamp(value, "UTC")),
        ]
    ),
    base.FuncIntFromGenericDatetime(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.toUnixTimestamp(value)),
        ]
    ),
    base.FuncIntFromDatetimeTZ(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.toUnixTimestamp(value, "UTC")),
        ]
    ),
    # str
    base.FuncStrFromUnsupported(
        variants=[
            V(D.CLICKHOUSE, sa.func.toString),
        ]
    ),
    base.FuncStrFromInteger(
        variants=[
            V(D.CLICKHOUSE, sa.func.toString),
        ]
    ),
    base.FuncStrFromFloat(
        variants=[
            V(D.CLICKHOUSE, sa.func.toString),
        ]
    ),
    base.FuncStrFromBool(
        variants=[
            V(D.CLICKHOUSE, lambda value: getattr(sa.func, "if")(value, "True", "False")),
        ]
    ),
    base.FuncStrFromStrGeo.for_dialect(D.CLICKHOUSE),
    base.FuncStrFromDate(
        variants=[
            V(D.CLICKHOUSE, sa.func.toString),
        ]
    ),
    base.FuncStrFromDatetime(
        variants=[
            V(D.CLICKHOUSE, lambda value: sa.func.toString(value)),
        ]
    ),
    base.FuncStrFromDatetimeTZ(
        variants=[
            VW(D.CLICKHOUSE, ensure_naive_first_arg(n.func.STR)),
        ]
    ),
    base.FuncStrFromString.for_dialect(D.CLICKHOUSE),
    base.FuncStrFromUUID(
        variants=[
            V(D.CLICKHOUSE, sa.func.toString),
        ]
    ),
    base.FuncStrFromArray(
        variants=[
            V(D.CLICKHOUSE, sa.func.toString),
        ]
    ),
    # tree
    base.FuncTreeStr.for_dialect(D.CLICKHOUSE),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/functions_type.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Optional,
    cast,
)

import sqlalchemy as sa
from sqlalchemy.sql.elements import (
    ClauseElement,
    ClauseList,
)
from sqlalchemy.sql.functions import Function as SAFunction

from dl_formula.definitions.base import (
    FuncTranslationImplementationBase,
    TranslateCallback,
    TranslationVariant,
)
import dl_formula.definitions.functions_window as base
from dl_formula.translation.context import TranslationCtx
from dl_formula.translation.env import TranslationEnvironment

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make

SUPPORTED_DIALECTS = D.and_above(D.CLICKHOUSE_22_10)


class WinLagClickHouseBase(base.WinLagBase):
    variants = [
        V(
            SUPPORTED_DIALECTS,
            translation=lambda x, offset=sa.literal(1), default=sa.null(), *_: (
                base.lag_implementation(
                    x,
                    offset=offset,
                    default=default,
                    lag_name="lagInFrame",
                    lead_name="leadInFrame",
                )
            ),
            translation_rows=lambda x, offset=sa.literal(1), *_: (None, None),
            as_winfunc=True,
        ),
    ]


class WinLag1ClickHouse(WinLagClickHouseBase):
    arg_cnt = 1


class WinLag2ClickHouse(WinLagClickHouseBase):
    arg_cnt = 2


class WinLag3ClickHouse(WinLagClickHouseBase):
    arg_cnt = 3


class RankPercentileTranslationImplementation(FuncTranslationImplementationBase):
    def translate(
        self,
        *raw_args: TranslationCtx,
        translator_cb: TranslateCallback,
        partition_by: Optional[ClauseList] = None,
        default_order_by: Optional[ClauseList] = None,
        translation_ctx: Optional[TranslationCtx] = None,
        translation_env: Optional[TranslationEnvironment] = None,
    ) -> ClauseElement:
        def translation_rank(value: Any, *args: Any) -> SAFunction:
            return sa.func.RANK(value)

        args, kwargs = self._handle_args(
            raw_args,
            translation_func=translation_rank,
            translation_ctx=translation_ctx,
            translation_env=translation_env,
        )
        order_by_part = base._order_by_from_args(*args)  # type: ignore

        wf_rank = translation_rank(*args).over(partition_by=partition_by, order_by=order_by_part)
        wf_total_partition_rows = sa.func.COUNT(1).over(partition_by=partition_by)  # ORDER BY is unnecessary

        result = (wf_rank - 1) / (wf_total_partition_rows - 1)
        return cast(ClauseElement, result)


class WinRankPercentileClickHouseBase(base.WinRankPercentileBase):
    variants = [
        TranslationVariant(
            dialects=SUPPORTED_DIALECTS,
            value=RankPercentileTranslationImplementation(),
        ),
    ]


class WinRankPercentile1ClickHouse(WinRankPercentileClickHouseBase):
    arg_cnt = 1


class WinRankPercentile2ClickHouse(WinRankPercentileClickHouseBase):
    arg_cnt = 2


DEFINITIONS_WINDOW = [
    # avg
    base.WinAvg.for_dialect(SUPPORTED_DIALECTS),
    # avg_if
    base.WinAvgIf.for_dialect(SUPPORTED_DIALECTS),
    # count
    base.WinCount0.for_dialect(SUPPORTED_DIALECTS),
    base.WinCount1.for_dialect(SUPPORTED_DIALECTS),
    # count_if
    base.WinCountIf.for_dialect(SUPPORTED_DIALECTS),
    # first
    base.WinFirst.for_dialect(SUPPORTED_DIALECTS),
    # lag
    WinLag1ClickHouse(),
    WinLag2ClickHouse(),
    WinLag3ClickHouse(),
    # last
    base.WinLast.for_dialect(SUPPORTED_DIALECTS),
    # mavg
    base.WinMAvg2.for_dialect(SUPPORTED_DIALECTS),
    base.WinMAvg3.for_dialect(SUPPORTED_DIALECTS),
    # max
    base.WinMax.for_dialect(SUPPORTED_DIALECTS),
    # mcount
    base.WinMCount2.for_dialect(SUPPORTED_DIALECTS),
    base.WinMCount3.for_dialect(SUPPORTED_DIALECTS),
    # min
    base.WinMin.for_dialect(SUPPORTED_DIALECTS),
    # mmax
    base.WinMMax2.for_dialect(SUPPORTED_DIALECTS),
    base.WinMMax3.for_dialect(SUPPORTED_DIALECTS),
    # mmin
    base.WinMMin2.for_dialect(SUPPORTED_DIALECTS),
    base.WinMMin3.for_dialect(SUPPORTED_DIALECTS),
    # msum
    base.WinMSum2.for_dialect(SUPPORTED_DIALECTS),
    base.WinMSum3.for_dialect(SUPPORTED_DIALECTS),
    # rank
    base.WinRank1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRank2.for_dialect(SUPPORTED_DIALECTS),
    # rank_dense
    base.WinRankDense1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRankDense2.for_dialect(SUPPORTED_DIALECTS),
    # rank_percentile
    WinRankPercentile1ClickHouse(),
    WinRankPercentile2ClickHouse(),
    # rank_unique
    base.WinRankUnique1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRankUnique2.for_dialect(SUPPORTED_DIALECTS),
    # ravg
    base.WinRAvg1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRAvg2.for_dialect(SUPPORTED_DIALECTS),
    # rcount
    base.WinRCount1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRCount2.for_dialect(SUPPORTED_DIALECTS),
    # rmax
    base.WinRMax1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRMax2.for_dialect(SUPPORTED_DIALECTS),
    # rmin
    base.WinRMin1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRMin2.for_dialect(SUPPORTED_DIALECTS),
    # rsum
    base.WinRSum1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRSum2.for_dialect(SUPPORTED_DIALECTS),
    # sum
    base.WinSum.for_dialect(SUPPORTED_DIALECTS),
    # sum_if
    base.WinSumIf.for_dialect(SUPPORTED_DIALECTS),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/functions_window.py
datalens-tech/datalens-backend
import functools
from typing import Any

import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.common import raw_sql
from dl_formula.definitions.common_datetime import DAY_SEC
from dl_formula.definitions.literals import (
    Literal,
    literal,
)
import dl_formula.definitions.operators_binary as base
from dl_formula.translation.context import TranslationCtx
from dl_formula.translation.env import TranslationEnvironment

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class CHBinaryEqNEqOrderedBaseArgTransformer(
    base.BinaryEqNEqOrderedBase.BinaryEqNEqOrderedBaseArgTransformer,
):
    @staticmethod
    def _to_datetime(value: TranslationCtx) -> TranslationCtx:
        value = value.copy()
        value.set_expression(sa.func.toDateTime(value.expression, "UTC"))
        value.set_type(DataType.DATETIME)
        return value

    @staticmethod
    def _to_genericdatetime(value: TranslationCtx) -> TranslationCtx:
        value = value.copy()
        value.set_expression(sa.func.toDateTime(value.expression))
        value.set_type(DataType.GENERICDATETIME)
        return value

    def transform_args(
        self, env: TranslationEnvironment, args: list[TranslationCtx], arg_types: list[DataType]
    ) -> list[TranslationCtx]:
        args = super().transform_args(env=env, args=args, arg_types=arg_types)

        if env.dialect & D.CLICKHOUSE:  # FIXME: Connectorize
            # if `args` is a mix of DATE and DATETIME/GENERICDATETIME, then DATE must be cast to DATETIME/GENERICDATETIME
            # TODO: DATETIMETZ
            if arg_types[0] in {DataType.DATE, DataType.CONST_DATE} and arg_types[1] in {
                DataType.DATETIME,
                DataType.CONST_DATETIME,
            }:
                args = [self._to_datetime(args[0]), args[1]]
            elif arg_types[0] in {DataType.DATE, DataType.CONST_DATE} and arg_types[1] in {
                DataType.GENERICDATETIME,
                DataType.CONST_GENERICDATETIME,
            }:
                args = [self._to_genericdatetime(args[0]), args[1]]
            elif arg_types[0] in {DataType.DATETIME, DataType.CONST_DATETIME} and arg_types[1] in {
                DataType.DATE,
                DataType.CONST_DATE,
            }:
                args = [args[0], self._to_datetime(args[1])]
            elif arg_types[0] in {DataType.GENERICDATETIME, DataType.CONST_GENERICDATETIME} and arg_types[1] in {
                DataType.DATE,
                DataType.CONST_DATE,
            }:
                args = [args[0], self._to_genericdatetime(args[1])]

        return args


def _get_null_substitute_value(data_type: DataType, dialect: DialectCombo) -> tuple[Literal, bool]:
    def lit(v: Any) -> Literal:
        return literal(v, d=dialect)

    stringify_value = False

    null_value: Literal
    if data_type in {
        DataType.STRING,
        DataType.CONST_STRING,
        DataType.GEOPOINT,
        DataType.CONST_GEOPOINT,
        DataType.GEOPOLYGON,
        DataType.CONST_GEOPOLYGON,
        DataType.MARKUP,
        DataType.CONST_MARKUP,
    }:
        null_value = lit("-")
    elif data_type in {
        DataType.DATE,
        DataType.CONST_DATE,
        DataType.DATETIME,
        DataType.CONST_DATETIME,
        DataType.GENERICDATETIME,
        DataType.CONST_GENERICDATETIME,
    }:
        null_value = lit("-")
        # TODO: Revert to using date values once dialects prior to 22.10 are no longer supported
        stringify_value = True  # Convert original value to string
    elif data_type in {
        DataType.INTEGER,
        DataType.CONST_INTEGER,
        DataType.FLOAT,
        DataType.CONST_FLOAT,
        DataType.BOOLEAN,
        DataType.CONST_BOOLEAN,
    }:
        null_value = lit(0)
    elif data_type in {DataType.UUID, DataType.CONST_UUID}:
        null_value = sa.func.toUUID(lit("00000000-0000-0000-0000-000000000000"))
    elif data_type in {
        DataType.ARRAY_INT,
        DataType.CONST_ARRAY_INT,
        DataType.ARRAY_FLOAT,
        DataType.CONST_ARRAY_FLOAT,
    }:
        null_value = lit([0])
    elif data_type in {
        DataType.ARRAY_STR,
        DataType.CONST_ARRAY_STR,
        DataType.TREE_STR,
        DataType.CONST_TREE_STR,
    }:
        null_value = lit(["__null_value__"])
    else:
        raise ValueError

    return null_value, stringify_value


def _denullified_eq(
    left_ctx: TranslationCtx,
    right_ctx: TranslationCtx,
    translation_env: TranslationEnvironment,
) -> ClauseElement:
    data_type = left_ctx.data_type
    assert data_type is not None

    right_type = right_ctx.data_type
    assert right_type is not None
    assert right_type.non_const_type == data_type.non_const_type

    if data_type == DataType.UNSUPPORTED:
        # We don't know the real type,
        # so the best we can do is build a regular equality expression
        return left_ctx.expression == right_ctx.expression  # type: ignore

    dialect = translation_env.dialect
    null_value, stringify_value = _get_null_substitute_value(data_type, dialect=dialect)

    left = left_ctx.expression
    right = right_ctx.expression
    if stringify_value:
        left = sa.func.toString(left)
        right = sa.func.toString(right)

    expr = sa.and_(
        sa.func.ifNull(left, null_value) == sa.func.ifNull(right, null_value),
        sa.func.isNull(left) == sa.func.isNull(right),
    )
    return expr  # type: ignore


class BinaryEqualDenullified(base.BinaryEqualDenullified):
    variants = [
        VW(D.CLICKHOUSE, lambda left, right, _env: _denullified_eq(left, right, _env)),
    ]


class BinaryInDate(base.BinaryIn):
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.DATE]),
    ]
    variants = [
        V(
            D.CLICKHOUSE_22_10,
            functools.partial(base._in_fix_null, stringify_values=True),  # noqa
        ),
    ]


class BinaryNotInDate(base.BinaryNotIn):
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.DATE]),
    ]
    variants = [
        V(
            D.CLICKHOUSE_22_10,
            functools.partial(base._not_in_fix_null, stringify_values=True),
        ),
    ]


DEFINITIONS_BINARY = [
    # !=
    base.BinaryNotEqual.for_dialect(D.CLICKHOUSE, arg_transformer=CHBinaryEqNEqOrderedBaseArgTransformer()),
    # %
    base.BinaryModInteger.for_dialect(D.CLICKHOUSE),
    base.BinaryModFloat(
        variants=[
            V(D.CLICKHOUSE, lambda left, right: left - sa.func.floor(left / right) * right),
        ]
    ),
    # *
    base.BinaryMultNumbers.for_dialect(D.CLICKHOUSE),
    base.BinaryMultStringConst.for_dialect(D.CLICKHOUSE),
    base.BinaryMultStringNonConst(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda text, size: sa.func.arrayStringConcat(
                    sa.func.arrayResize(sa.func.emptyArrayString(), size, text)
                ),
            ),
        ]
    ),
    # +
    base.BinaryPlusNumbers.for_dialect(D.CLICKHOUSE),
    base.BinaryPlusStrings.for_dialect(D.CLICKHOUSE),
    base.BinaryPlusDateInt(
        variants=[
            V(D.CLICKHOUSE, lambda date, days: sa.func.date_add(raw_sql("day"), days, date)),
        ]
    ),
    base.BinaryPlusArray(
        variants=[
            V(D.CLICKHOUSE, sa.func.arrayConcat),
        ]
    ),
    base.BinaryPlusDateFloat(
        variants=[
            V(D.CLICKHOUSE, lambda date, days: sa.func.date_add(raw_sql("day"), base.as_bigint(days), date)),
        ]
    ),
    base.BinaryPlusDatetimeNumber(
        variants=[
            V(D.CLICKHOUSE, lambda dt, days: sa.func.toDateTime(dt + base.as_bigint(days * DAY_SEC), "UTC")),
        ]
    ),
    base.BinaryPlusGenericDatetimeNumber(
        variants=[
            V(D.CLICKHOUSE, lambda dt, days: sa.func.date_add(raw_sql("second"), days * DAY_SEC, dt)),
        ]
    ),
    # -
    base.BinaryMinusNumbers.for_dialect(D.CLICKHOUSE),
    base.BinaryMinusDateInt(
        variants=[
            V(D.CLICKHOUSE, lambda date, days: sa.func.date_add(raw_sql("day"), -days, date)),
        ]
    ),
    base.BinaryMinusDateFloat(
        variants=[
            V(
                D.CLICKHOUSE,
                lambda date, days: sa.func.date_add(raw_sql("day"), base.as_bigint(-sa.func.ceil(days)), date),
            ),
        ]
    ),
    base.BinaryMinusDatetimeNumber(
        variants=[
            V(D.CLICKHOUSE, lambda dt, days: sa.func.toDateTime(dt - base.as_bigint(days * DAY_SEC), "UTC")),
        ]
    ),
    base.BinaryMinusGenericDatetimeNumber(
        variants=[
            V(D.CLICKHOUSE, lambda dt, days: sa.func.date_sub(raw_sql("second"), days * DAY_SEC, dt)),
        ]
    ),
    base.BinaryMinusDates(
        variants=[
            V(D.CLICKHOUSE, lambda left, right: sa.func.DATEDIFF("day", right, left)),
        ]
    ),
    base.BinaryMinusDatetimes(
        variants=[
            V(D.CLICKHOUSE, lambda left, right: (sa.func.toUInt32(left) - sa.func.toUInt32(right)) / DAY_SEC),
        ]
    ),
    base.BinaryMinusGenericDatetimes(
        variants=[
            V(D.CLICKHOUSE, lambda left, right: sa.func.DATEDIFF("second", right, left) / DAY_SEC),
        ]
    ),
    # /
    base.BinaryDivInt.for_dialect(D.CLICKHOUSE),
    base.BinaryDivFloat.for_dialect(D.CLICKHOUSE),
    # <
    base.BinaryLessThan.for_dialect(D.CLICKHOUSE, arg_transformer=CHBinaryEqNEqOrderedBaseArgTransformer()),
    # <=
    base.BinaryLessThanOrEqual.for_dialect(D.CLICKHOUSE, arg_transformer=CHBinaryEqNEqOrderedBaseArgTransformer()),
    # ==
    base.BinaryEqual.for_dialect(D.CLICKHOUSE, arg_transformer=CHBinaryEqNEqOrderedBaseArgTransformer()),
    # >
    base.BinaryGreaterThan.for_dialect(D.CLICKHOUSE, arg_transformer=CHBinaryEqNEqOrderedBaseArgTransformer()),
    # >=
    base.BinaryGreaterThanOrEqual.for_dialect(D.CLICKHOUSE, arg_transformer=CHBinaryEqNEqOrderedBaseArgTransformer()),
    # ^
    base.BinaryPower.for_dialect(D.CLICKHOUSE),
    # _!=
    base.BinaryNotEqualInternal.for_dialect(D.CLICKHOUSE, arg_transformer=CHBinaryEqNEqOrderedBaseArgTransformer()),
    # _==
    base.BinaryEqualInternal.for_dialect(D.CLICKHOUSE, arg_transformer=CHBinaryEqNEqOrderedBaseArgTransformer()),
    # _dneq
    BinaryEqualDenullified(arg_transformer=CHBinaryEqNEqOrderedBaseArgTransformer()),
    # and
    base.BinaryAnd.for_dialect(D.CLICKHOUSE),
    # in
    BinaryInDate(),
    base.BinaryIn.for_dialect(D.CLICKHOUSE),
    # like
    base.BinaryLike.for_dialect(D.CLICKHOUSE),
    # notin
    BinaryNotInDate(),
    base.BinaryNotIn.for_dialect(D.CLICKHOUSE),
    # notlike
    base.BinaryNotLike.for_dialect(D.CLICKHOUSE),
    # or
    base.BinaryOr.for_dialect(D.CLICKHOUSE),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/operators_binary.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_ternary as base

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


DEFINITIONS_TERNARY = [
    # between
    base.TernaryBetween.for_dialect(D.CLICKHOUSE),
    # notbetween
    base.TernaryNotBetween.for_dialect(D.CLICKHOUSE),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/operators_ternary.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.operators_unary as base

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


V = TranslationVariant.make


DEFINITIONS_UNARY = [
    # isfalse
    base.UnaryIsFalseStringGeo.for_dialect(D.CLICKHOUSE),
    base.UnaryIsFalseNumbers.for_dialect(D.CLICKHOUSE),
    base.UnaryIsFalseDateTime(
        variants=[
            V(D.CLICKHOUSE, lambda x: sa.literal(0)),
        ]
    ),
    base.UnaryIsFalseBoolean(
        variants=[
            V(D.CLICKHOUSE, lambda x: x == 0),
        ]
    ),
    # istrue
    base.UnaryIsTrueStringGeo.for_dialect(D.CLICKHOUSE),
    base.UnaryIsTrueNumbers.for_dialect(D.CLICKHOUSE),
    base.UnaryIsTrueDateTime(
        variants=[
            V(D.CLICKHOUSE, lambda x: sa.literal(1)),
        ]
    ),
    base.UnaryIsTrueBoolean(
        variants=[
            V(D.CLICKHOUSE, lambda x: x != 0),
        ]
    ),
    # neg
    base.UnaryNegate.for_dialect(D.CLICKHOUSE),
    # not
    base.UnaryNotBool.for_dialect(D.CLICKHOUSE),
    base.UnaryNotNumbers.for_dialect(D.CLICKHOUSE),
    base.UnaryNotStringGeo.for_dialect(D.CLICKHOUSE),
    base.UnaryNotDateDatetime(
        variants=[
            V(D.CLICKHOUSE, lambda x: sa.literal(0)),
        ]
    ),
]

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/definitions/operators_unary.py
datalens-tech/datalens-backend
import datetime
from typing import Union

import sqlalchemy as sa

from dl_formula.connectors.base.literal import (
    Literal,
    Literalizer,
)
from dl_formula.core.dialect import DialectCombo

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


class ClickHouseLiteralizer(Literalizer):
    __slots__ = ()

    def literal_datetime(self, value: datetime.datetime, dialect: DialectCombo) -> Literal:
        # TODO: DateTime64 support for fresh CHs.
        tzinfo = value.tzinfo
        value = value.replace(microsecond=0, tzinfo=None)
        if tzinfo:
            tzname = getattr(tzinfo, "zone", tzinfo.tzname(value))
            return sa.func.toDateTime(value.isoformat(), tzname)
        else:
            return sa.func.toDateTime(value.isoformat())

    def literal_date(self, value: datetime.date, dialect: DialectCombo) -> Literal:
        if dialect & D.and_above(D.CLICKHOUSE_22_10):
            return sa.func.toDate32(value.isoformat())
        return sa.func.toDate(value.isoformat())

    def literal_bool(self, value: bool, dialect: DialectCombo) -> Literal:
        return sa.literal(int(value))

    def literal_array(self, value: Union[tuple, list], dialect: DialectCombo) -> Literal:
        return sa.func.array(*value)

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/literal.py
datalens-tech/datalens-backend
import clickhouse_sqlalchemy.types as ch_types
from sqlalchemy.types import TypeEngine

from dl_formula.connectors.base.type_constructor import DefaultSATypeConstructor
from dl_formula.core.datatype import DataType


class ClickHouseTypeConstructor(DefaultSATypeConstructor):
    def get_sa_type(self, data_type: DataType) -> TypeEngine:
        type_map: dict[DataType, TypeEngine] = {
            **{
                str_type: ch_types.Nullable(ch_types.String)
                for str_type in (
                    DataType.STRING,
                    DataType.GEOPOINT,
                    DataType.GEOPOLYGON,
                )
            },
            DataType.INTEGER: ch_types.Nullable(ch_types.Int64),
            DataType.FLOAT: ch_types.Nullable(ch_types.Float64),
            DataType.BOOLEAN: ch_types.Nullable(ch_types.UInt8),
            DataType.DATE: ch_types.Nullable(ch_types.Date),
            DataType.DATETIME: ch_types.Nullable(ch_types.DateTime),
            DataType.UUID: ch_types.Nullable(ch_types.UUID),
            DataType.ARRAY_INT: ch_types.Array(ch_types.Nullable(ch_types.Int64)),
            DataType.ARRAY_FLOAT: ch_types.Array(ch_types.Nullable(ch_types.Float64)),
            DataType.ARRAY_STR: ch_types.Array(ch_types.Nullable(ch_types.String)),
        }
        if (type_eng := type_map.get(data_type)) is not None:
            return type_eng
        else:
            return super().get_sa_type(data_type)

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula/type_constructor.py
datalens-tech/datalens-backend
from dl_formula_ref.texts import StyledDialect

from dl_connector_clickhouse.formula.constants import ClickHouseDialect
from dl_connector_clickhouse.formula_ref.i18n import Translatable


HUMAN_DIALECTS = {
    ClickHouseDialect.CLICKHOUSE: StyledDialect(
        "`ClickHouse`",
        "`ClickHouse`",
        "`ClickHouse`",
    ),
    ClickHouseDialect.CLICKHOUSE_21_8: StyledDialect(
        "`ClickHouse 21.8`",
        "`ClickHouse`<br/>`21.8`",
        Translatable("`ClickHouse` version `21.8`"),
    ),
    ClickHouseDialect.CLICKHOUSE_22_10: StyledDialect(
        "`ClickHouse 22.10`",
        "`ClickHouse`<br/>`22.10`",
        Translatable("`ClickHouse` version `22.10`"),
    ),
}

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula_ref/human_dialects.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_clickhouse as package


DOMAIN = f"dl_formula_ref_{package.__name__}"

_LOCALE_DIR = os.path.join(os.path.dirname(__file__), "..", "locales")

CONFIGS = [
    TranslationConfig(path=_LOCALE_DIR, domain=DOMAIN, locale="en"),
    TranslationConfig(path=_LOCALE_DIR, domain=DOMAIN, locale="ru"),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula_ref/i18n.py
datalens-tech/datalens-backend
from dl_formula_ref.functions.type_conversion import DbCastExtension
from dl_formula_ref.plugins.base.plugin import FormulaRefPlugin

from dl_connector_clickhouse.formula.constants import ClickHouseDialect
from dl_connector_clickhouse.formula.definitions.functions_type import FuncDbCastClickHouseBase
from dl_connector_clickhouse.formula_ref.human_dialects import HUMAN_DIALECTS
from dl_connector_clickhouse.formula_ref.i18n import CONFIGS


class ClickHouseFormulaRefPlugin(FormulaRefPlugin):
    any_dialects = frozenset((*ClickHouseDialect.CLICKHOUSE.to_list(),))
    compeng_support_dialects = frozenset((*ClickHouseDialect.CLICKHOUSE.to_list(),))
    human_dialects = HUMAN_DIALECTS
    translation_configs = frozenset(CONFIGS)
    db_cast_extension = DbCastExtension(
        type_whitelists=FuncDbCastClickHouseBase.WHITELISTS,
    )

lib/dl_connector_clickhouse/dl_connector_clickhouse/formula_ref/plugin.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_api_lib_testing.data_api_base import StandardizedDataApiTestBase
from dl_api_lib_testing.dataset_base import DatasetTestBase

from dl_connector_clickhouse.core.clickhouse.constants import SOURCE_TYPE_CH_TABLE
from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
from dl_connector_clickhouse_tests.db.config import (
    API_TEST_CONFIG,
    CoreConnectionSettings,
)
from dl_connector_clickhouse_tests.db.core.base import BaseClickHouseTestClass


class ClickHouseConnectionTestBase(BaseClickHouseTestClass, ConnectionTestBase):
    conn_type = CONNECTION_TYPE_CLICKHOUSE
    bi_compeng_pg_on = False

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def connection_params(self) -> dict:
        return dict(
            db_name=CoreConnectionSettings.DB_NAME,
            host=CoreConnectionSettings.HOST,
            port=CoreConnectionSettings.PORT,
            username=CoreConnectionSettings.USERNAME,
            password=CoreConnectionSettings.PASSWORD,
        )


class ClickHouseDatasetTestBase(ClickHouseConnectionTestBase, DatasetTestBase):
    @pytest.fixture(scope="class")
    def dataset_params(self, sample_table) -> dict:
        return dict(
            source_type=SOURCE_TYPE_CH_TABLE.name,
            parameters=dict(
                db_name=sample_table.db.name,
                table_name=sample_table.name,
            ),
        )


class ClickHouseDataApiTestBase(ClickHouseDatasetTestBase, StandardizedDataApiTestBase):
    mutation_caches_on = False

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/api/base.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.complex_queries import DefaultBasicComplexQueryTestSuite

from dl_connector_clickhouse_tests.db.api.base import ClickHouseDataApiTestBase


class TestClickHouseBasicComplexQueries(ClickHouseDataApiTestBase, DefaultBasicComplexQueryTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/api/test_complex_queries.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite

from dl_connector_clickhouse_tests.db.api.base import ClickHouseConnectionTestBase


class TestClickHouseConnection(ClickHouseConnectionTestBase, DefaultConnectorConnectionTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/api/test_connection.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.data_api_suites import (
    DefaultConnectorDataDistinctTestSuite,
    DefaultConnectorDataGroupByFormulaTestSuite,
    DefaultConnectorDataPreviewTestSuite,
    DefaultConnectorDataRangeTestSuite,
    DefaultConnectorDataResultTestSuite,
)

from dl_connector_clickhouse_tests.db.api.base import ClickHouseDataApiTestBase


class TestClickHouseDataResult(ClickHouseDataApiTestBase, DefaultConnectorDataResultTestSuite):
    pass


class TestClickHouseDataGroupBy(ClickHouseDataApiTestBase, DefaultConnectorDataGroupByFormulaTestSuite):
    pass


class TestClickHouseDataRange(ClickHouseDataApiTestBase, DefaultConnectorDataRangeTestSuite):
    pass


class TestClickHouseDataDistinct(ClickHouseDataApiTestBase, DefaultConnectorDataDistinctTestSuite):
    pass


class TestClickHouseDataPreview(ClickHouseDataApiTestBase, DefaultConnectorDataPreviewTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/api/test_data.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite

from dl_connector_clickhouse_tests.db.api.base import ClickHouseDatasetTestBase


class TestClickHouseDataset(ClickHouseDatasetTestBase, DefaultConnectorDatasetTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/api/test_dataset.py
datalens-tech/datalens-backend
import os
from typing import ClassVar

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D


# Infra settings
SSL_CA_FILENAME = "marsnet_ca.crt"


def get_clickhouse_ssl_ca_path() -> str:
    return os.path.join(os.path.dirname(__file__), "../../docker-compose/db-clickhouse/ssl", SSL_CA_FILENAME)


def get_clickhouse_ssl_ca() -> str:
    path = get_clickhouse_ssl_ca_path()

    with open(path) as f:
        return f.read()


CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=52211).host,
    port_us_http=get_test_container_hostport("us", fallback_port=52211).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=52210).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=52210).port,
    us_master_token="AC1ofiek8coB",
    core_connector_ep_names=["clickhouse"],
)


class CoreConnectionSettings:
    DB_NAME: ClassVar[str] = "test_data"
    HOST: ClassVar[str] = get_test_container_hostport("db-clickhouse-22-10", fallback_port=52204).host
    PORT: ClassVar[int] = get_test_container_hostport("db-clickhouse-22-10", fallback_port=52204).port
    USERNAME: ClassVar[str] = "datalens"
    PASSWORD: ClassVar[str] = "qwerty"


class CoreSslConnectionSettings:
    DB_NAME: ClassVar[str] = "test_data"
    # don't use get_test_container_hostport, localhost in needed due to IP usage deprecation in CN
    HOST: ClassVar[str] = "localhost"
    PORT: ClassVar[int] = 52206
    USERNAME: ClassVar[str] = "datalens"
    PASSWORD: ClassVar[str] = "qwerty"


DB_URLS = {
    D.CLICKHOUSE_21_8: f"clickhouse://datalens:qwerty@"
    f'{get_test_container_hostport("db-clickhouse-21-8", fallback_port=52202).as_pair()}/test_data',
    D.CLICKHOUSE_22_10: f"clickhouse://datalens:qwerty@"
    f'{get_test_container_hostport("db-clickhouse-22-10", fallback_port=52204).as_pair()}/test_data',
}
DB_CORE_URL = DB_URLS[D.CLICKHOUSE_22_10]

API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["clickhouse"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
)

# don't use get_test_container_hostport, localhost in needed due to IP usage deprecation in CN
DB_CORE_SSL_URL = "clickhouse://datalens:qwerty@localhost:52206/test_data?protocol=https"

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/config.py
datalens-tech/datalens-backend
from dl_api_lib_testing.initialization import initialize_api_lib_test
from dl_formula_testing.forced_literal import forced_literal_use

from dl_connector_clickhouse_tests.db.config import API_TEST_CONFIG


pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)


__all__ = (
    # auto-use fixtures:
    "forced_literal_use",
)

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/conftest.py
datalens-tech/datalens-backend
import asyncio
import os
from typing import Generator

from frozendict import frozendict
import pytest

import dl_configs.utils as bi_configs_utils
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.testcases.connection import BaseConnectionTestClass

from dl_connector_clickhouse.core.clickhouse.testing.connection import make_clickhouse_saved_connection
from dl_connector_clickhouse.core.clickhouse.us_connection import ConnectionClickhouse
from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
from dl_connector_clickhouse.db_testing.engine_wrapper import ClickhouseDbEngineConfig
import dl_connector_clickhouse_tests.db.config as test_config


class BaseClickHouseTestClass(BaseConnectionTestClass[ConnectionClickhouse]):
    conn_type = CONNECTION_TYPE_CLICKHOUSE
    core_test_config = test_config.CORE_TEST_CONFIG

    @pytest.fixture(autouse=True)
    # FIXME: This fixture is a temporary solution for failing core tests when they are run together with api tests
    def loop(self, event_loop: asyncio.AbstractEventLoop) -> Generator[asyncio.AbstractEventLoop, None, None]:
        asyncio.set_event_loop(event_loop)
        yield event_loop
        # Attempt to cover an old version of pytest-asyncio:
        # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
        asyncio.set_event_loop_policy(None)

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return test_config.DB_CORE_URL

    @pytest.fixture(scope="class")
    def engine_config(self, db_url: str, engine_params: dict) -> ClickhouseDbEngineConfig:
        return ClickhouseDbEngineConfig(url=db_url, engine_params=engine_params)

    @pytest.fixture(scope="function")
    def connection_creation_params(self) -> dict:
        return dict(
            db_name=test_config.CoreConnectionSettings.DB_NAME,
            host=test_config.CoreConnectionSettings.HOST,
            port=test_config.CoreConnectionSettings.PORT,
            username=test_config.CoreConnectionSettings.USERNAME,
            password=test_config.CoreConnectionSettings.PASSWORD,
            **(dict(raw_sql_level=self.raw_sql_level) if self.raw_sql_level is not None else {}),
        )

    @pytest.fixture(scope="function")
    def saved_connection(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params: dict,
    ) -> ConnectionClickhouse:
        conn = make_clickhouse_saved_connection(sync_usm=sync_us_manager, **connection_creation_params)
        return conn


class BaseSslClickHouseTestClass(BaseClickHouseTestClass):
    @pytest.fixture(autouse=True)
    def clear_ssl_folder(self):
        folder_path = bi_configs_utils.get_temp_root_certificates_folder_path()

        for filename in os.listdir(folder_path):
            file_path = os.path.join(folder_path, filename)
            os.unlink(file_path)

    def check_ssl_directory_is_empty(self) -> None:
        assert not os.listdir(bi_configs_utils.get_temp_root_certificates_folder_path())

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return test_config.DB_CORE_SSL_URL

    @pytest.fixture(scope="class")
    def engine_params(self) -> dict:
        return {
            "connect_args": frozendict(
                verify=test_config.get_clickhouse_ssl_ca_path(),
            )
        }

    @pytest.fixture(scope="function")
    def connection_creation_params(self) -> dict:
        return dict(
            db_name=test_config.CoreSslConnectionSettings.DB_NAME,
            host=test_config.CoreSslConnectionSettings.HOST,
            port=test_config.CoreSslConnectionSettings.PORT,
            username=test_config.CoreSslConnectionSettings.USERNAME,
            password=test_config.CoreSslConnectionSettings.PASSWORD,
            **(dict(raw_sql_level=self.raw_sql_level) if self.raw_sql_level is not None else {}),
            secure=True,
            ssl_ca=test_config.get_clickhouse_ssl_ca(),
        )

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/core/base.py
datalens-tech/datalens-backend
from dl_core_testing.testcases.adapter import BaseAsyncAdapterTestClass

from dl_connector_clickhouse.core.clickhouse_base.adapters import AsyncClickHouseAdapter
from dl_connector_clickhouse.core.clickhouse_base.target_dto import ClickHouseConnTargetDTO
from dl_connector_clickhouse_tests.db.core.base import BaseClickHouseTestClass


class TestAsyncClickHouseAdapter(
    BaseClickHouseTestClass,
    BaseAsyncAdapterTestClass[ClickHouseConnTargetDTO],
):
    ASYNC_ADAPTER_CLS = AsyncClickHouseAdapter

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/core/test_adapter.py
datalens-tech/datalens-backend
from __future__ import annotations

import os

import pytest

from dl_core.us_connection_base import DataSourceTemplate
from dl_core_testing.testcases.connection import DefaultConnectionTestClass

from dl_connector_clickhouse.core.clickhouse.us_connection import ConnectionClickhouse
from dl_connector_clickhouse_tests.db.core.base import (
    BaseClickHouseTestClass,
    BaseSslClickHouseTestClass,
)


class TestClickHouseConnection(
    BaseClickHouseTestClass,
    DefaultConnectionTestClass[ConnectionClickhouse],
):
    do_check_data_export_flag = True

    def check_saved_connection(self, conn: ConnectionClickhouse, params: dict) -> None:
        assert conn.uuid is not None
        assert conn.data.db_name == params["db_name"]
        assert conn.data.username == params["username"]
        assert conn.data.secure is False
        assert conn.data.ssl_ca is None

    def check_data_source_templates(
        self,
        conn: ConnectionClickhouse,
        dsrc_templates: list[DataSourceTemplate],
    ) -> None:
        assert dsrc_templates
        for dsrc_tmpl in dsrc_templates:
            assert conn.db_name in dsrc_tmpl.group
            assert dsrc_tmpl.title

        # Make sure
        tmpl_db_names = {dsrc_tmpl.group[0] for dsrc_tmpl in dsrc_templates}
        assert "system" not in tmpl_db_names


@pytest.mark.skipif(os.environ.get("WE_ARE_IN_CI"), reason="can't use localhost")
class TestSslClickHouseConnection(
    BaseSslClickHouseTestClass,
    TestClickHouseConnection,
    DefaultConnectionTestClass[ConnectionClickhouse],
):
    def check_saved_connection(self, conn: ConnectionClickhouse, params: dict) -> None:
        assert conn.uuid is not None
        assert conn.data.db_name == params["db_name"]
        assert conn.data.username == params["username"]
        assert conn.data.secure is True
        assert conn.data.ssl_ca is params["ssl_ca"]

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/core/test_connection.py
datalens-tech/datalens-backend
import asyncio
import enum
import os
from typing import (
    Optional,
    Sequence,
)

import attr
from clickhouse_sqlalchemy import types as ch_types
import pytest
import sqlalchemy as sa

from dl_constants.enums import (
    ConnectionType,
    UserDataType,
)
from dl_core.connection_executors import (
    AsyncConnExecutorBase,
    ConnExecutorQuery,
    SyncConnExecutorBase,
)
from dl_core.connection_models.common_models import (
    DBIdent,
    SchemaIdent,
)
from dl_core.db.native_type import (
    ClickHouseDateTime64WithTZNativeType,
    ClickHouseDateTimeWithTZNativeType,
    ClickHouseNativeType,
    GenericNativeType,
    norm_native_type,
)
from dl_core_testing.database import (
    Db,
    DbTable,
)
from dl_core_testing.testcases.connection_executor import (
    DefaultAsyncConnectionExecutorTestSuite,
    DefaultSyncAsyncConnectionExecutorCheckBase,
    DefaultSyncConnectionExecutorTestSuite,
)
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_clickhouse.core.clickhouse.us_connection import ConnectionClickhouse
from dl_connector_clickhouse_tests.db.config import CoreConnectionSettings
from dl_connector_clickhouse_tests.db.core.base import (
    BaseClickHouseTestClass,
    BaseSslClickHouseTestClass,
)


class ClickHouseSyncAsyncConnectionExecutorCheckBase(
    BaseClickHouseTestClass,
    DefaultSyncAsyncConnectionExecutorCheckBase[ConnectionClickhouse],
):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultAsyncConnectionExecutorTestSuite.test_closing_sql_sessions: "Not implemented",
        },
    )

    @pytest.fixture(scope="function")
    def db_ident(self) -> DBIdent:
        return DBIdent(db_name=CoreConnectionSettings.DB_NAME)

    def check_db_version(self, db_version: Optional[str]) -> None:
        assert db_version is not None
        assert "." in db_version


class TestClickHouseSyncConnectionExecutor(
    ClickHouseSyncAsyncConnectionExecutorCheckBase,
    DefaultSyncConnectionExecutorTestSuite[ConnectionClickhouse],
):
    @attr.s(frozen=True)
    class CD(DefaultSyncConnectionExecutorTestSuite.CD):
        def get_expected_native_type(self, conn_type: ConnectionType) -> GenericNativeType:
            actual_type = (
                self.sa_type.nested_type
                if isinstance(self.sa_type, (ch_types.Nullable, ch_types.LowCardinality))
                else self.sa_type
            )
            return self.nt or ClickHouseNativeType(
                conn_type=conn_type,
                name=norm_native_type(self.nt_name if self.nt_name is not None else actual_type),
                nullable=isinstance(self.sa_type, ch_types.Nullable),  # note: self.nullable is not taken into account
                lowcardinality=False,
            )

    def get_schemas_for_type_recognition(self) -> dict[str, Sequence[DefaultSyncConnectionExecutorTestSuite.CD]]:
        enum_values = ["allchars '\"\t,= etc", "test", "value1"]
        tst_enum8 = enum.Enum("TstEnum8", enum_values)
        tst_enum16 = enum.Enum("TstEnum16", enum_values)

        return {
            "ch_types_numbers": [
                self.CD(ch_types.Nullable(ch_types.Int8()), UserDataType.integer),
                self.CD(ch_types.Nullable(ch_types.Int16()), UserDataType.integer),
                self.CD(ch_types.Nullable(ch_types.Int32()), UserDataType.integer),
                self.CD(ch_types.Nullable(ch_types.Int64()), UserDataType.integer),
                self.CD(ch_types.Nullable(ch_types.UInt8()), UserDataType.integer),
                self.CD(ch_types.Nullable(ch_types.UInt16()), UserDataType.integer),
                self.CD(ch_types.Nullable(ch_types.UInt32()), UserDataType.integer),
                self.CD(ch_types.Nullable(ch_types.UInt64()), UserDataType.integer),
                self.CD(ch_types.Nullable(ch_types.Float32()), UserDataType.float),
                self.CD(ch_types.Nullable(ch_types.Float64()), UserDataType.float),
                self.CD(ch_types.Decimal(8, 4), UserDataType.float, nt_name="float"),
            ],
            "ch_types_string": [
                self.CD(sa.String(length=256), UserDataType.string),
                self.CD(ch_types.Nullable(ch_types.String()), UserDataType.string),
                # Note: `Nullable(LowCardinality(String))` is actually not allowed:
                # "Nested type LowCardinality(String) cannot be inside Nullable type"
                self.CD(
                    ch_types.LowCardinality(ch_types.Nullable(ch_types.String())),
                    UserDataType.string,
                    nt=ClickHouseNativeType(
                        conn_type=self.conn_type,
                        name="string",
                        nullable=True,
                        lowcardinality=True,
                    ),
                ),
                self.CD(
                    ch_types.LowCardinality(ch_types.String()),
                    UserDataType.string,
                    nt=ClickHouseNativeType(
                        conn_type=self.conn_type,
                        name="string",
                        nullable=False,
                        lowcardinality=True,
                    ),
                ),
            ],
            "ch_types_date": [
                # not nullable so we can check 0000-00-00
                self.CD(ch_types.Date(), UserDataType.date),
                self.CD(ch_types.Date32(), UserDataType.date, nt_name="date"),
                # not nullable so we can check 0000-00-00 00:00:00
                self.CD(
                    ch_types.DateTime(),
                    UserDataType.genericdatetime,
                    nt=ClickHouseDateTimeWithTZNativeType(
                        conn_type=self.conn_type,
                        name="datetimewithtz",
                        nullable=False,
                        lowcardinality=False,
                        timezone_name="UTC",  # the CH system timezone
                        explicit_timezone=False,
                    ),
                ),
                self.CD(
                    ch_types.DateTimeWithTZ("Europe/Moscow"),
                    UserDataType.genericdatetime,
                    nt=ClickHouseDateTimeWithTZNativeType(
                        conn_type=self.conn_type,
                        name="datetimewithtz",
                        nullable=False,
                        lowcardinality=False,
                        timezone_name="Europe/Moscow",
                        explicit_timezone=True,
                    ),
                ),
                self.CD(
                    ch_types.DateTime64(6),
                    UserDataType.genericdatetime,
                    nt=ClickHouseDateTime64WithTZNativeType(
                        conn_type=self.conn_type,
                        name="datetime64withtz",
                        nullable=False,
                        lowcardinality=False,
                        precision=6,
                        timezone_name="UTC",  # the CH system timezone
                        explicit_timezone=False,
                    ),
                ),
                self.CD(
                    ch_types.DateTime64WithTZ(6, "Europe/Moscow"),
                    UserDataType.genericdatetime,
                    nt=ClickHouseDateTime64WithTZNativeType(
                        conn_type=self.conn_type,
                        name="datetime64withtz",
                        nullable=False,
                        lowcardinality=False,
                        precision=6,
                        timezone_name="Europe/Moscow",
                        explicit_timezone=True,
                    ),
                ),
            ],
            "ch_types_other": [
                self.CD(ch_types.Enum8(tst_enum8), UserDataType.string, nt_name="string"),
                self.CD(ch_types.Enum16(tst_enum16), UserDataType.string, nt_name="string"),
                self.CD(ch_types.Nullable(ch_types.Bool()), UserDataType.boolean),
            ],
        }


class TestClickHouseAsyncConnectionExecutor(
    ClickHouseSyncAsyncConnectionExecutorCheckBase,
    DefaultAsyncConnectionExecutorTestSuite[ConnectionClickhouse],
):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultAsyncConnectionExecutorTestSuite.test_get_db_version: "Not implemented",
            DefaultAsyncConnectionExecutorTestSuite.test_test: "Not implemented",
            DefaultAsyncConnectionExecutorTestSuite.test_table_exists: "Not implemented",
            DefaultAsyncConnectionExecutorTestSuite.test_table_not_exists: "Not implemented",
            DefaultAsyncConnectionExecutorTestSuite.test_get_table_schema_info: "Not implemented",
            DefaultAsyncConnectionExecutorTestSuite.test_get_table_schema_info_for_nonexistent_table: "Not implemented",
        },
    )

    @pytest.mark.asyncio
    async def test_sa_mod(self, async_connection_executor: AsyncConnExecutorBase) -> None:
        result = await async_connection_executor.execute(ConnExecutorQuery(sa.select([sa.literal(3) % sa.literal(2)])))
        rows = await result.get_all()
        assert rows == [(1,)]

    @pytest.mark.asyncio
    async def test_inf(self, async_connection_executor: AsyncConnExecutorBase) -> None:
        result = await async_connection_executor.execute(ConnExecutorQuery("select 1 / 0"))
        rows = await result.get_all()
        assert rows == [(None,)]


@pytest.mark.skipif(os.environ.get("WE_ARE_IN_CI"), reason="can't use localhost")
class TestSslClickHouseSyncConnectionExecutor(
    BaseSslClickHouseTestClass,
    TestClickHouseSyncConnectionExecutor,
):
    def test_test(self, sync_connection_executor: SyncConnExecutorBase) -> None:
        super().test_test(sync_connection_executor)
        self.check_ssl_directory_is_empty()


@pytest.mark.skipif(os.environ.get("WE_ARE_IN_CI"), reason="can't use localhost")
class TestSslClickHouseAsyncConnectionExecutor(
    BaseSslClickHouseTestClass,
    TestClickHouseAsyncConnectionExecutor,
):
    async def test_test(self, async_connection_executor: AsyncConnExecutorBase) -> None:
        await super().test_test(async_connection_executor)
        await asyncio.sleep(0.1)
        self.check_ssl_directory_is_empty()

    async def test_multiple_connection_test(self, async_connection_executor: AsyncConnExecutorBase) -> None:
        await super().test_multiple_connection_test(async_connection_executor)
        await asyncio.sleep(0.1)
        self.check_ssl_directory_is_empty()

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/core/test_connection_executor.py
datalens-tech/datalens-backend
import pytest

from dl_constants.enums import (
    RawSQLLevel,
    UserDataType,
)
from dl_core.data_source_spec.sql import (
    StandardSQLDataSourceSpec,
    SubselectDataSourceSpec,
)
from dl_core_testing.fixtures.sample_tables import TABLE_SPEC_SAMPLE_SUPERSTORE
from dl_core_testing.testcases.data_source import DefaultDataSourceTestClass

from dl_connector_clickhouse.core.clickhouse.constants import (
    SOURCE_TYPE_CH_SUBSELECT,
    SOURCE_TYPE_CH_TABLE,
)
from dl_connector_clickhouse.core.clickhouse.data_source import (
    ClickHouseDataSource,
    ClickHouseSubselectDataSource,
)
from dl_connector_clickhouse.core.clickhouse.us_connection import ConnectionClickhouse
from dl_connector_clickhouse_tests.db.core.base import BaseClickHouseTestClass


class TestClickHouseTableDataSource(
    BaseClickHouseTestClass,
    DefaultDataSourceTestClass[
        ConnectionClickhouse,
        StandardSQLDataSourceSpec,
        ClickHouseDataSource,
    ],
):
    DSRC_CLS = ClickHouseDataSource

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self, sample_table) -> StandardSQLDataSourceSpec:
        dsrc_spec = StandardSQLDataSourceSpec(
            source_type=SOURCE_TYPE_CH_TABLE,
            db_name=sample_table.db.name,
            table_name=sample_table.name,
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return list(TABLE_SPEC_SAMPLE_SUPERSTORE.table_schema)


class TestClickHouseSubselectDataSource(
    BaseClickHouseTestClass,
    DefaultDataSourceTestClass[
        ConnectionClickhouse,
        SubselectDataSourceSpec,
        ClickHouseSubselectDataSource,
    ],
):
    DSRC_CLS = ClickHouseSubselectDataSource

    raw_sql_level = RawSQLLevel.subselect

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self, sample_table) -> SubselectDataSourceSpec:
        dsrc_spec = SubselectDataSourceSpec(
            source_type=SOURCE_TYPE_CH_SUBSELECT,
            subsql=f'SELECT * FROM "{sample_table.name}"',
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return list(TABLE_SPEC_SAMPLE_SUPERSTORE.table_schema)

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/core/test_data_source.py
datalens-tech/datalens-backend
from dl_core_testing.testcases.dataset import DefaultDatasetTestSuite

from dl_connector_clickhouse.core.clickhouse.constants import SOURCE_TYPE_CH_TABLE
from dl_connector_clickhouse.core.clickhouse.us_connection import ConnectionClickhouse
from dl_connector_clickhouse_tests.db.core.base import BaseClickHouseTestClass


class TestClickHouseDataset(BaseClickHouseTestClass, DefaultDatasetTestSuite[ConnectionClickhouse]):
    source_type = SOURCE_TYPE_CH_TABLE

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/core/test_dataset.py
datalens-tech/datalens-backend
import datetime

import pytest
import pytz

from dl_formula_testing.testcases.base import FormulaConnectorTestBase

from dl_connector_clickhouse.db_testing.engine_wrapper import ClickhouseDbEngineConfig
from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D
from dl_connector_clickhouse_tests.db.config import DB_URLS


class ClickHouseTestBase(FormulaConnectorTestBase):
    engine_config_cls = ClickhouseDbEngineConfig
    supports_arrays = True
    supports_uuid = True

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return DB_URLS[self.dialect]


class ClickHouse_21_8TestBase(ClickHouseTestBase):
    dialect = D.CLICKHOUSE_21_8

    @pytest.fixture(scope="class")
    def tzinfo(self) -> datetime.tzinfo:
        # The database is configured to use a custom timezone
        return pytz.timezone("America/New_York")


class ClickHouse_22_10TestBase(ClickHouseTestBase):
    dialect = D.CLICKHOUSE_22_10

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/base.py
datalens-tech/datalens-backend
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.conditional_blocks import DefaultConditionalBlockFormulaConnectorTestSuite
from dl_formula_testing.util import to_str

from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_21_8TestBase


class ConditionalBlockClickHouseTestSuite(DefaultConditionalBlockFormulaConnectorTestSuite):
    def test_case_block_returning_null(self, dbe: DbEvaluator):
        # https://github.com/ClickHouse/ClickHouse/issues/7237
        assert to_str(dbe.eval('CASE 3 WHEN 1 THEN "1st" WHEN 2 THEN "2nd" END')) == ""
        assert dbe.eval("CASE 3 WHEN 1 THEN 1 WHEN 2 THEN 2 END") == 0
        assert dbe.eval("CASE 3 WHEN 1 THEN 1.1 WHEN 2 THEN 2.2 END") == 0.0
        assert dbe.eval("CASE 3 WHEN 1 THEN TRUE WHEN 2 THEN TRUE END") is False
        # NULL in THEN
        assert to_str(dbe.eval('CASE 2 WHEN 1 THEN NULL WHEN 2 THEN "2nd" ELSE "3rd" END')) == "2nd"
        assert to_str(dbe.eval('CASE 1 WHEN 1 THEN NULL WHEN 2 THEN "2nd" ELSE "3rd" END')) == ""


class TestMainAggFunctionClickHouse_21_8(ClickHouse_21_8TestBase, ConditionalBlockClickHouseTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_conditional_blocks.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.functions_aggregation import DefaultMainAggFunctionFormulaConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_21_8TestBase


class MainAggFunctionClickHouseTestSuite(DefaultMainAggFunctionFormulaConnectorTestSuite):
    supports_countd_approx = True
    supports_quantile = True
    supports_median = True
    supports_arg_min_max = True
    supports_any = True
    supports_all_concat = True
    supports_top_concat = True

    def test_quantile(self, dbe: DbEvaluator, data_table: sa.Table) -> None:  # additional checks for CH
        value = dbe.eval("QUANTILE([int_value], 0.9)", from_=data_table)
        assert 80 <= value <= 90
        assert abs(value - 90) < 0.5  # allow for float approximation.

        value = dbe.eval("QUANTILE_APPROX([int_value], 0.9)", from_=data_table)
        assert 80 < value < 91  # can be either 81 or 90, apparently.


class TestMainAggFunctionClickHouse_21_8(ClickHouse_21_8TestBase, MainAggFunctionClickHouseTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_functions_aggregation.py
datalens-tech/datalens-backend
from functools import reduce

import pytest
import sqlalchemy as sa

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.functions_array import DefaultArrayFunctionFormulaConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_21_8TestBase


class ArrayFunctionClickHouseTestSuite(DefaultArrayFunctionFormulaConnectorTestSuite):
    def test_startswith_string_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("STARTSWITH([arr_str_value], [arr_str_value])", from_=data_table)
        assert not dbe.eval('STARTSWITH([arr_str_value], ARRAY("", "cde", NULL))', from_=data_table)


class TestArrayFunctionClickHouse_21_8(ClickHouse_21_8TestBase, ArrayFunctionClickHouseTestSuite):
    def test_array_contains_subsequence(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("CONTAINS_SUBSEQUENCE(ARRAY(1, 2, 3), ARRAY(1, 2))", from_=data_table)
        assert dbe.eval("CONTAINS_SUBSEQUENCE(ARRAY(1.1, 2.2, 3.3), ARRAY(1.1, 2.2))", from_=data_table)
        assert dbe.eval('CONTAINS_SUBSEQUENCE(ARRAY("a", "b", "c"), ARRAY("a", "b"))', from_=data_table)

        assert not dbe.eval("CONTAINS_SUBSEQUENCE(ARRAY(1, 2, 3), ARRAY(1, 3))", from_=data_table)
        assert not dbe.eval("CONTAINS_SUBSEQUENCE(ARRAY(1.1, 2.2, 3.3), ARRAY(2.2, 1.1))", from_=data_table)
        assert not dbe.eval('CONTAINS_SUBSEQUENCE(ARRAY("a", "b", "c"), ARRAY("a", "a"))', from_=data_table)
        assert not dbe.eval('CONTAINS_SUBSEQUENCE(ARRAY("a", "b", "c"), ARRAY("a", "e"))', from_=data_table)

        assert dbe.eval('CONTAINS_SUBSEQUENCE(ARRAY("a", NULL, "c"), ARRAY("a", NULL))', from_=data_table)
        assert not dbe.eval('CONTAINS_SUBSEQUENCE(ARRAY("a", NULL, "c"), ARRAY("a", "c"))', from_=data_table)
        assert not dbe.eval("CONTAINS_SUBSEQUENCE(ARRAY(1.1, 2.2, 3.3), ARRAY(1.1, NULL))", from_=data_table)

        assert dbe.eval("CONTAINS_SUBSEQUENCE([arr_int_value], ARRAY(23, 456))", from_=data_table)
        assert not dbe.eval("CONTAINS_SUBSEQUENCE([arr_int_value], ARRAY(24, 456))", from_=data_table)

        assert dbe.eval("CONTAINS_SUBSEQUENCE([arr_int_value], SLICE([arr_int_value], 2, 1))", from_=data_table)
        assert dbe.eval("CONTAINS_SUBSEQUENCE(ARRAY(0, 23, 456, NULL, 123), [arr_int_value])", from_=data_table)

        assert dbe.eval('CONTAINS_SUBSEQUENCE([arr_str_value], ARRAY("cde"))', from_=data_table)
        assert dbe.eval('CONTAINS_SUBSEQUENCE([arr_str_value], ARRAY("cde", NULL))', from_=data_table)
        assert not dbe.eval('CONTAINS_SUBSEQUENCE(ARRAY("cde"), [arr_str_value])', from_=data_table)

    @pytest.mark.parametrize(
        "bi_func, eval_func",
        [
            ("ARR_MIN", min),
            ("ARR_MAX", max),
            ("ARR_SUM", sum),
            ("ARR_AVG", lambda a: sum(a) / len(a)),
            ("ARR_PRODUCT", lambda a: reduce(lambda x, y: x * y, a, 1)),
        ],
    )
    def test_array_aggregation_functions(self, dbe: DbEvaluator, data_table: sa.Table, bi_func, eval_func):
        """
        Array aggregation functions (arrayMin/Max/Sum/Avg) were added in CH 21.1
        (arrayProduct in 21.6).
        """

        inp_int = (1, 2, 3, -1)
        inp_float = (1.2, 12, 0.1, 12.0)

        bi_inp_int = ", ".join((str(item) for item in inp_int))
        bi_inp_float = ", ".join((str(item) for item in inp_float))

        assert dbe.eval(f"{bi_func}(ARRAY({bi_inp_int}))", from_=data_table) == eval_func(inp_int)
        assert dbe.eval(f"{bi_func}(REPLACE([arr_int_value], NULL, 1))", from_=data_table) == eval_func((0, 23, 456, 1))
        assert dbe.eval(f"{bi_func}(ARRAY({bi_inp_float}))", from_=data_table) == eval_func(inp_float)
        assert dbe.eval(f"{bi_func}(REPLACE([arr_float_value], NULL, 1))", from_=data_table) == eval_func(
            (0, 45, 0.123, 1)
        )

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_functions_array.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_datetime import DefaultDateTimeFunctionFormulaConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import (
    ClickHouse_21_8TestBase,
    ClickHouse_22_10TestBase,
)


class DateTimeFunctionClickHouseTestSuite(DefaultDateTimeFunctionFormulaConnectorTestSuite):
    supports_dateadd_non_const_unit_num = True
    supports_deprecated_dateadd = True
    supports_deprecated_datepart_2 = True
    supports_datetrunc_3 = True
    supports_datetimetz = True


class TestDateTimeFunctionClickHouse_21_8(ClickHouse_21_8TestBase, DateTimeFunctionClickHouseTestSuite):
    pass


class TestDateTimeFunctionClickHouse_22_10(ClickHouse_22_10TestBase, DateTimeFunctionClickHouseTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_functions_datetime.py
datalens-tech/datalens-backend
import pytest

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.functions_logical import DefaultLogicalFunctionFormulaConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_21_8TestBase


class LogicalFunctionClickHouseTestSuite(DefaultLogicalFunctionFormulaConnectorTestSuite):
    supports_nan_funcs = True
    supports_iif = True

    @pytest.mark.skip("Not fixed in ClickHouse yet")
    def test_bi_1052(self, dbe: DbEvaluator) -> None:
        assert dbe.eval('IF MONTH(DATE("1989-03-17")) = 8 THEN "first" ELSE "second" END')


class TestMainAggFunctionClickHouse_21_8(ClickHouse_21_8TestBase, LogicalFunctionClickHouseTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_functions_logical.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_markup import DefaultMarkupFunctionFormulaConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_21_8TestBase


class TestMarkupFunctionClickHouse_21_8(ClickHouse_21_8TestBase, DefaultMarkupFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_functions_markup.py
datalens-tech/datalens-backend
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.functions_math import DefaultMathFunctionFormulaConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_21_8TestBase


class MathFunctionClickHouseTestSuite(DefaultMathFunctionFormulaConnectorTestSuite):
    def test_compare(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("COMPARE(0, 0, 0)") == 0
        assert dbe.eval("COMPARE(0.123, 0.123, 0)") == 0
        assert dbe.eval("COMPARE(0.123, 0.123, 0.1)") == 0
        assert dbe.eval("COMPARE(0.123, 0.12, 0.1)") == 0
        assert dbe.eval("COMPARE(0.123, 0.12, 0.001)") == 1
        assert dbe.eval("COMPARE(0.12, 0.123, 0.001)") == -1


class TestMathFunctionClickHouse_21_8(ClickHouse_21_8TestBase, MathFunctionClickHouseTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_functions_math.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.core import nodes
from dl_formula.core.datatype import DataType
from dl_formula.translation import ext_nodes
from dl_formula.translation.context import TranslationCtx
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.functions_string import DefaultStringFunctionFormulaConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_21_8TestBase


class StringFunctionClickHouseTestSuite(DefaultStringFunctionFormulaConnectorTestSuite):
    def test_utf8(self, dbe: DbEvaluator) -> None:
        # UTF8() is only available for CLICKHOUSE
        binary_str = TranslationCtx(expression=sa.text(r"'\xcf\xf0\xe8\xe2\xe5\xf2'"), data_type=DataType.STRING)
        assert (
            dbe.eval(
                nodes.Formula.make(
                    nodes.FuncCall.make(
                        name="UTF8",
                        args=[ext_nodes.CompiledExpression.make(binary_str), nodes.LiteralString.make("CP-1251")],
                    )
                )
            )
            == "Привет"
        )


class TestStringFunctionClickHouse_21_8(ClickHouse_21_8TestBase, StringFunctionClickHouseTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_functions_string.py
datalens-tech/datalens-backend
import datetime

import clickhouse_sqlalchemy.exceptions as ch_exc
import pytest
import pytz
import sqlalchemy as sa

import dl_formula.core.exc as exc
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.testcases.functions_type_conversion import (
    DefaultBoolTypeFunctionFormulaConnectorTestSuite,
    DefaultDateTypeFunctionFormulaConnectorTestSuite,
    DefaultDbCastTypeFunctionFormulaConnectorTestSuite,
    DefaultFloatTypeFunctionFormulaConnectorTestSuite,
    DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite,
    DefaultGeopointTypeFunctionFormulaConnectorTestSuite,
    DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite,
    DefaultIntTypeFunctionFormulaConnectorTestSuite,
    DefaultStrTypeFunctionFormulaConnectorTestSuite,
    DefaultTreeTypeFunctionFormulaConnectorTestSuite,
)
from dl_formula_testing.util import (
    as_tz,
    dt_strip,
    to_str,
    utc_ts,
)

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D
from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_21_8TestBase


# STR


class StrTypeFunctionClickHouseTestSuite(DefaultStrTypeFunctionFormulaConnectorTestSuite):
    bool_of_null_is_false = True

    def test_str_from_array(self, dbe: DbEvaluator, data_table) -> None:
        assert to_str(dbe.eval("STR([arr_int_value])", from_=data_table)) == "[0,23,456,NULL]"
        assert to_str(dbe.eval("STR([arr_float_value])", from_=data_table)) == "[0,45,0.123,NULL]"
        assert to_str(dbe.eval("STR([arr_str_value])", from_=data_table)) == "['','','cde',NULL]"


class TestStrTypeFunctionClickHouse_21_8(ClickHouse_21_8TestBase, StrTypeFunctionClickHouseTestSuite):
    pass


# FLOAT


class TestFloatTypeFunctionClickHouse_21_8(ClickHouse_21_8TestBase, DefaultFloatTypeFunctionFormulaConnectorTestSuite):
    pass


# BOOL


class TestBoolTypeFunctionClickHouse_21_8(ClickHouse_21_8TestBase, DefaultBoolTypeFunctionFormulaConnectorTestSuite):
    pass


# INT


class TestIntTypeFunctionClickHouse_21_8(ClickHouse_21_8TestBase, DefaultIntTypeFunctionFormulaConnectorTestSuite):
    pass


# DATE


class DateTypeFunctionClickHouseTestSuite(DefaultDateTypeFunctionFormulaConnectorTestSuite):
    def test_date2(self, dbe: DbEvaluator) -> None:
        ts = utc_ts(2019, 1, 2, 23, 4, 5)
        tz = "Europe/Moscow"
        expected = datetime.date(2019, 1, 3)
        assert dbe.eval('DATE({}, "{}")'.format(int(ts), tz)) == expected
        assert dbe.eval('DATE({}, "{}")'.format(float(ts), tz)) == expected


class TestDateTypeFunctionClickHouse_21_8(ClickHouse_21_8TestBase, DateTypeFunctionClickHouseTestSuite):
    pass


# DATE_PARSE


class TestDateParseTypeFunctionClickHouse_21_8(ClickHouse_21_8TestBase, FormulaConnectorTestBase):
    pass


# GENERICDATETIME (& DATETIME)


class GenericDatetimeTypeFunctionClickHouseTestSuite(DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite):
    @pytest.mark.parametrize("func_name", ("GENERICDATETIME", "DATETIME"))
    def test_genericdatetime2_ch(self, dbe: DbEvaluator, func_name: str) -> None:
        mos_tz = "Europe/Moscow"
        dt_naive = datetime.datetime(2019, 1, 3, 2, 4, 5)
        dt_naive_iso_str = dt_naive.isoformat()
        dt_as_mos = as_tz(dt_naive, tzinfo=pytz.timezone(mos_tz))
        dt_as_mos_ts = dt_as_mos.timestamp()

        dt_as_dbtz = as_tz(dt_naive, tzinfo=dbe.db.tzinfo)
        dt_as_dbtz_to_mos = dt_as_dbtz.astimezone(pytz.timezone(mos_tz))
        dt_as_dbtz_to_mos_stripped = dt_strip(dt_as_dbtz_to_mos)

        assert dt_strip(dbe.eval(f'{func_name}({int(dt_as_mos_ts)}, "{mos_tz}")')) == dt_naive
        assert dt_strip(dbe.eval(f'{func_name}({float(dt_as_mos_ts)}, "{mos_tz}")')) == dt_naive

        # Notably, the string and genericdatetime are interpreted at the specified timezone
        # (with undefined resolution of ambiguities),
        # the number / datetime is interpreted as UTC

        # from string
        assert dt_strip(dbe.eval(f"{func_name}('{dt_naive_iso_str}', '{mos_tz}')")) == dt_naive
        # from literal
        assert dt_strip(dbe.eval(f"{func_name}(#{dt_naive_iso_str}#, '{mos_tz}')")) == dt_as_dbtz_to_mos_stripped

        # Double-wrap (from string)
        assert dt_strip(dbe.eval(f"{func_name}({func_name}('{dt_naive_iso_str}', '{mos_tz}'), '{mos_tz}')")) == dt_naive


class TestGenericDatetimeTypeFunctionClickHouse_21_8(
    ClickHouse_21_8TestBase,
    GenericDatetimeTypeFunctionClickHouseTestSuite,
):
    pass


# GENERICDATETIME_PARSE (& DATETIME_PARSE)


class GenericDatetimeParseTypeFunctionClickHouseTestSuite(FormulaConnectorTestBase):
    @pytest.mark.parametrize("func_name", ("GENERICDATETIME_PARSE", "DATETIME_PARSE"))
    def test_new_datetime_parse(self, dbe: DbEvaluator, func_name: str) -> None:
        if not dbe.dialect & D.and_above(D.CLICKHOUSE_21_8):
            pytest.skip()

        assert dt_strip(dbe.eval(f'{func_name}("2019-01-02 03:04:05")')) == (datetime.datetime(2019, 1, 2, 3, 4, 5))
        assert dt_strip(dbe.eval(f'{func_name}("2019-01-02 03:04:05+02")')) == (
            datetime.datetime(2019, 1, 1, 20, 4, 5)
        )  # moved to NYC timezone
        assert dt_strip(dbe.eval(f'{func_name}("20190102030405")')) == (datetime.datetime(2019, 1, 2, 3, 4, 5))
        assert dt_strip(dbe.eval(f'{func_name}("20190102 030405")')) == (datetime.datetime(2019, 1, 2, 3, 4, 5))
        assert dt_strip(dbe.eval(f'{func_name}("2019.01.02 03:04:05")')) == (datetime.datetime(2019, 1, 2, 3, 4, 5))
        assert dt_strip(dbe.eval(f'{func_name}("02/01/2019 03:04:05")')) == (datetime.datetime(2019, 1, 2, 3, 4, 5))
        assert dt_strip(dbe.eval(f'{func_name}("2019-01-02 03:04")')) == (datetime.datetime(2019, 1, 2, 3, 4, 0))
        assert dt_strip(dbe.eval(f'{func_name}("2019-01-02 030405")')) == (datetime.datetime(2019, 1, 2, 3, 4, 5))
        assert dt_strip(dbe.eval(f'{func_name}("2019 Jan 02 03:04:05")')) == (datetime.datetime(2019, 1, 2, 3, 4, 5))


class TestGenericDatetimeParseTypeFunctionClickHouse_21_8(
    ClickHouse_21_8TestBase,
    GenericDatetimeParseTypeFunctionClickHouseTestSuite,
):
    pass


# GEOPOINT


class TestGeopointTypeFunctionClickHouse_21_8(
    ClickHouse_21_8TestBase,
    DefaultGeopointTypeFunctionFormulaConnectorTestSuite,
):
    pass


# GEOPOLYGON


class TestGeopolygonTypeFunctionClickHouse_21_8(
    ClickHouse_21_8TestBase,
    DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite,
):
    pass


# DB_CAST


class DbCastTypeFunctionClickHouseTestSuite(DefaultDbCastTypeFunctionFormulaConnectorTestSuite):
    def test_db_cast_clickhouse(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        value = dbe.eval("[int_value]", from_=data_table)
        assert dbe.eval('DB_CAST(FLOAT([int_value]), "Float32")', from_=data_table) == pytest.approx(float(value))
        assert dbe.eval('DB_CAST(FLOAT([int_value]), "Float64")', from_=data_table) == pytest.approx(float(value))
        assert dbe.eval('DB_CAST(FLOAT([int_value]), "Decimal", 5, 0)', from_=data_table) == value

        assert dbe.eval('DB_CAST(FLOAT([int_value]), "Float64") / 0', from_=data_table)  # Doesn't raise an error
        with pytest.raises(ch_exc.DatabaseException):
            dbe.eval('DB_CAST(FLOAT([int_value]), "Decimal", 5, 0) / 0', from_=data_table)  # Not supported by CH

        with pytest.raises(exc.TranslationError):
            dbe.eval('DB_CAST(FLOAT([int_value]), "Decimal", 1)', from_=data_table)


class TestDbCastTypeFunctionClickHouse_21_8(
    ClickHouse_21_8TestBase,
    DbCastTypeFunctionClickHouseTestSuite,
):
    pass


# TREE


class TestTreeTypeFunctionClickHouse_21_8(
    ClickHouse_21_8TestBase,
    DefaultTreeTypeFunctionFormulaConnectorTestSuite,
):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_functions_type_conversion.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_window import DefaultWindowFunctionFormulaConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_22_10TestBase


class TestWindowFunctionClickHouse_22_10(ClickHouse_22_10TestBase, DefaultWindowFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_functions_window.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.literals import DefaultLiteralFormulaConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_21_8TestBase


class LiteralFunctionClickHouseTestSuite(DefaultLiteralFormulaConnectorTestSuite):
    supports_microseconds = False
    supports_utc = True
    supports_custom_tz = True
    default_tz = None


class TestLiteralFunctionClickHouse_21_8(ClickHouse_21_8TestBase, LiteralFunctionClickHouseTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_literals.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.misc_funcs import DefaultMiscFunctionalityConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import ClickHouse_21_8TestBase


class TestMiscFunctionalityClickHouse_21_8(ClickHouse_21_8TestBase, DefaultMiscFunctionalityConnectorTestSuite):
    pass

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_misc_funcs.py
datalens-tech/datalens-backend
import datetime

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.operators import DefaultOperatorFormulaConnectorTestSuite

from dl_connector_clickhouse_tests.db.formula.base import (
    ClickHouse_21_8TestBase,
    ClickHouse_22_10TestBase,
)


class TestOperatorClickHouse_21_8(ClickHouse_21_8TestBase, DefaultOperatorFormulaConnectorTestSuite):
    pass


class TestOperatorClickHouse_22_10(ClickHouse_22_10TestBase, DefaultOperatorFormulaConnectorTestSuite):
    def test_date_before_1970_add(self, dbe: DbEvaluator):
        assert dbe.eval("#1931-01-01# + 1") == datetime.date(1931, 1, 2)

    def test_date_before_1970_subtract(self, dbe: DbEvaluator):
        assert dbe.eval("#1931-01-03# - 1") == datetime.date(1931, 1, 2)

    def test_date_before_1970_diff(self, dbe: DbEvaluator):
        assert dbe.eval("#1931-01-03# - #1931-01-01#") == 2

    def test_date_before_and_after_1970_diff(self, dbe: DbEvaluator):
        assert dbe.eval("#1971-01-03# - #1931-01-03#") == 14610

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/db/formula/test_operators.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.dialect import DefaultDialectFormulaConnectorTestSuite

from dl_connector_clickhouse.formula.constants import ClickHouseDialect as D
from dl_connector_clickhouse.formula.constants import DIALECT_NAME_CLICKHOUSE


class DialectClickHouseTestSuite(DefaultDialectFormulaConnectorTestSuite):
    dialect_name = DIALECT_NAME_CLICKHOUSE
    default_dialect = D.CLICKHOUSE_21_8
    dialect_matches = (
        ("21.9.10", D.CLICKHOUSE_21_8),
        ("23.1.3", D.CLICKHOUSE_22_10),
    )

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/unit/formula/test_dialect.py
datalens-tech/datalens-backend
from dl_api_connector.i18n.localizer import CONFIGS as BI_API_CONNECTOR_CONFIGS
from dl_api_lib_testing.connection_form_base import ConnectionFormTestBase

from dl_connector_clickhouse.api.connection_form.form_config import ClickHouseConnectionFormFactory
from dl_connector_clickhouse.api.i18n.localizer import CONFIGS as BI_API_LIB_CONFIGS


class TestClickhouseConnectionForm(ConnectionFormTestBase):
    CONN_FORM_FACTORY_CLS = ClickHouseConnectionFormFactory
    TRANSLATION_CONFIGS = BI_API_CONNECTOR_CONFIGS + BI_API_LIB_CONFIGS

lib/dl_connector_clickhouse/dl_connector_clickhouse_tests/unit/test_connection_form.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_connector_greenplum.core.us_connection import GreenplumConnection
from dl_connector_postgresql.api.api_schema.connection import PostgreSQLConnectionSchema


class GreenplumConnectionSchema(PostgreSQLConnectionSchema):
    TARGET_CLS = GreenplumConnection  # type: ignore

lib/dl_connector_greenplum/dl_connector_greenplum/api/api_schema/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    Sequence,
)

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormActionApiSchema,
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
)
from dl_api_connector.form_config.models.common import CommonFieldName
import dl_api_connector.form_config.models.rows as C
from dl_api_connector.form_config.models.rows.base import FormRow
from dl_api_connector.form_config.models.shortcuts.rows import RowConstructor
from dl_configs.connectors_settings import ConnectorSettingsBase

from dl_connector_greenplum.api.connection_info import GreenplumConnectionInfoProvider


class GreenplumConnectionFormFactory(ConnectionFormFactory):
    def _get_base_edit_api_schema(self) -> FormActionApiSchema:
        return FormActionApiSchema(
            items=[
                FormFieldApiSchema(name=CommonFieldName.host, required=True),
                FormFieldApiSchema(name=CommonFieldName.port, required=True),
                FormFieldApiSchema(name=CommonFieldName.username, required=True),
                FormFieldApiSchema(name=CommonFieldName.db_name, required=True),
                FormFieldApiSchema(
                    name=CommonFieldName.password,
                    required=self.mode == ConnectionFormMode.create,
                ),
                FormFieldApiSchema(name=CommonFieldName.cache_ttl_sec, nullable=True),
                FormFieldApiSchema(name=CommonFieldName.raw_sql_level),
                FormFieldApiSchema(name=CommonFieldName.data_export_forbidden),
            ],
        )

    def _get_base_create_api_schema(self, edit_api_schema: FormActionApiSchema) -> FormActionApiSchema:
        return FormActionApiSchema(
            items=[
                *edit_api_schema.items,
                *self._get_top_level_create_api_schema_items(),
            ],
            conditions=edit_api_schema.conditions.copy(),
        )

    def _get_base_check_api_schema(self) -> FormActionApiSchema:
        return FormActionApiSchema(
            items=[
                FormFieldApiSchema(name=CommonFieldName.host, required=True),
                FormFieldApiSchema(name=CommonFieldName.port, required=True),
                FormFieldApiSchema(name=CommonFieldName.db_name, required=True),
                FormFieldApiSchema(name=CommonFieldName.username, required=True),
                FormFieldApiSchema(
                    name=CommonFieldName.password,
                    required=self.mode == ConnectionFormMode.create,
                ),
                *self._get_top_level_check_api_schema_items(),
            ]
        )

    def _get_base_form_config(
        self,
        host_section: Sequence[FormRow],
        username_section: Sequence[FormRow],
        db_name_section: Sequence[FormRow],
        create_api_schema: FormActionApiSchema,
        edit_api_schema: FormActionApiSchema,
        check_api_schema: FormActionApiSchema,
        rc: RowConstructor,
    ) -> ConnectionForm:
        return ConnectionForm(
            title=GreenplumConnectionInfoProvider.get_title(self._localizer),
            rows=self._filter_nulls(
                [
                    *host_section,
                    rc.port_row(default_value="5432"),
                    *db_name_section,
                    *username_section,
                    rc.password_row(mode=self.mode),
                    C.CacheTTLRow(name=CommonFieldName.cache_ttl_sec),
                    rc.raw_sql_level_row(),
                    rc.collapse_advanced_settings_row(),
                    rc.data_export_forbidden_row(),
                ]
            ),
            api_schema=FormApiSchema(
                create=create_api_schema if self.mode == ConnectionFormMode.create else None,
                edit=edit_api_schema if self.mode == ConnectionFormMode.edit else None,
                check=check_api_schema,
            ),
        )

    def get_form_config(
        self,
        connector_settings: Optional[ConnectorSettingsBase],
        tenant: Optional[TenantDef],
    ) -> ConnectionForm:
        rc = RowConstructor(localizer=self._localizer)

        host_section: list[FormRow] = [rc.host_row()]
        username_section: list[FormRow] = [rc.username_row()]
        db_name_section: list[FormRow] = [rc.db_name_row()]

        edit_api_schema = self._get_base_edit_api_schema()
        create_api_schema = self._get_base_create_api_schema(edit_api_schema)
        check_api_schema = self._get_base_check_api_schema()

        return self._get_base_form_config(
            host_section=host_section,
            username_section=username_section,
            db_name_section=db_name_section,
            create_api_schema=create_api_schema,
            edit_api_schema=edit_api_schema,
            check_api_schema=check_api_schema,
            rc=rc,
        )

lib/dl_connector_greenplum/dl_connector_greenplum/api/connection_form/form_config.py
datalens-tech/datalens-backend
from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_greenplum.api.i18n.localizer import Translatable


class GreenplumConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-greenplum")

lib/dl_connector_greenplum/dl_connector_greenplum/api/connection_info.py
datalens-tech/datalens-backend
from dl_api_connector.api_schema.source_base import (
    SchematizedSQLDataSourceSchema,
    SchematizedSQLDataSourceTemplateSchema,
    SubselectDataSourceSchema,
    SubselectDataSourceTemplateSchema,
)
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)

from dl_connector_greenplum.api.api_schema.connection import GreenplumConnectionSchema
from dl_connector_greenplum.api.connection_form.form_config import GreenplumConnectionFormFactory
from dl_connector_greenplum.api.connection_info import GreenplumConnectionInfoProvider
from dl_connector_greenplum.api.i18n.localizer import CONFIGS
from dl_connector_greenplum.core.connector import (
    GreenplumCoreConnectionDefinition,
    GreenplumCoreConnector,
    GreenplumSubselectCoreSourceDefinition,
    GreenplumTableCoreSourceDefinition,
)
from dl_connector_postgresql.formula.constants import DIALECT_NAME_POSTGRESQL


class GreenplumApiTableSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = GreenplumTableCoreSourceDefinition
    api_schema_cls = SchematizedSQLDataSourceSchema
    template_api_schema_cls = SchematizedSQLDataSourceTemplateSchema


class GreenplumApiSubselectSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = GreenplumSubselectCoreSourceDefinition
    api_schema_cls = SubselectDataSourceSchema
    template_api_schema_cls = SubselectDataSourceTemplateSchema


class GreenplumApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = GreenplumCoreConnectionDefinition
    api_generic_schema_cls = GreenplumConnectionSchema
    info_provider_cls = GreenplumConnectionInfoProvider
    form_factory_cls = GreenplumConnectionFormFactory


class GreenplumApiConnector(ApiConnector):
    core_connector_cls = GreenplumCoreConnector
    connection_definitions = (GreenplumApiConnectionDefinition,)
    source_definitions = (
        GreenplumApiTableSourceDefinition,
        GreenplumApiSubselectSourceDefinition,
    )
    formula_dialect_name = DIALECT_NAME_POSTGRESQL
    translation_configs = frozenset(CONFIGS)

lib/dl_connector_greenplum/dl_connector_greenplum/api/connector.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_greenplum as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_greenplum/dl_connector_greenplum/api/i18n/localizer.py
datalens-tech/datalens-backend
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
)
from dl_core.connectors.sql_base.connector import (
    SQLSubselectCoreSourceDefinitionBase,
    SQLTableCoreSourceDefinitionBase,
)

from dl_connector_greenplum.core.constants import (
    BACKEND_TYPE_GREENPLUM,
    CONNECTION_TYPE_GREENPLUM,
    SOURCE_TYPE_GP_SUBSELECT,
    SOURCE_TYPE_GP_TABLE,
)
from dl_connector_greenplum.core.data_source import (
    GreenplumSubselectDataSource,
    GreenplumTableDataSource,
)
from dl_connector_greenplum.core.data_source_migration import GreenPlumDataSourceMigrator
from dl_connector_greenplum.core.storage_schemas.connection import GreenplumConnectionDataStorageSchema
from dl_connector_greenplum.core.us_connection import GreenplumConnection
from dl_connector_postgresql.core.postgresql_base.adapters_postgres import PostgresAdapter
from dl_connector_postgresql.core.postgresql_base.async_adapters_postgres import AsyncPostgresAdapter
from dl_connector_postgresql.core.postgresql_base.connection_executors import (
    AsyncPostgresConnExecutor,
    PostgresConnExecutor,
)
from dl_connector_postgresql.core.postgresql_base.sa_types import SQLALCHEMY_POSTGRES_TYPES
from dl_connector_postgresql.core.postgresql_base.type_transformer import PostgreSQLTypeTransformer


class GreenplumCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_GREENPLUM
    connection_cls = GreenplumConnection
    us_storage_schema_cls = GreenplumConnectionDataStorageSchema
    type_transformer_cls = PostgreSQLTypeTransformer
    sync_conn_executor_cls = PostgresConnExecutor
    async_conn_executor_cls = AsyncPostgresConnExecutor
    dialect_string = "bi_postgresql"
    data_source_migrator_cls = GreenPlumDataSourceMigrator


class GreenplumTableCoreSourceDefinition(SQLTableCoreSourceDefinitionBase):
    source_type = SOURCE_TYPE_GP_TABLE
    source_cls = GreenplumTableDataSource


class GreenplumSubselectCoreSourceDefinition(SQLSubselectCoreSourceDefinitionBase):
    source_type = SOURCE_TYPE_GP_SUBSELECT
    source_cls = GreenplumSubselectDataSource


class GreenplumCoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_GREENPLUM
    connection_definitions = (GreenplumCoreConnectionDefinition,)
    source_definitions = (
        GreenplumTableCoreSourceDefinition,
        GreenplumSubselectCoreSourceDefinition,
    )
    rqe_adapter_classes = frozenset({PostgresAdapter, AsyncPostgresAdapter})
    sa_types = SQLALCHEMY_POSTGRES_TYPES

lib/dl_connector_greenplum/dl_connector_greenplum/core/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
    SourceBackendType,
)


BACKEND_TYPE_GREENPLUM = SourceBackendType.declare("GREENPLUM")
CONNECTION_TYPE_GREENPLUM = ConnectionType.declare("greenplum")
SOURCE_TYPE_GP_TABLE = DataSourceType.declare("GP_TABLE")
SOURCE_TYPE_GP_SUBSELECT = DataSourceType.declare("GP_SUBSELECT")

lib/dl_connector_greenplum/dl_connector_greenplum/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.enums import DataSourceType
from dl_core.data_source.sql import (
    BaseSQLDataSource,
    StandardSchemaSQLDataSource,
    SubselectDataSource,
)

from dl_connector_greenplum.core.constants import (
    CONNECTION_TYPE_GREENPLUM,
    SOURCE_TYPE_GP_SUBSELECT,
    SOURCE_TYPE_GP_TABLE,
)
from dl_connector_postgresql.core.postgresql_base.query_compiler import PostgreSQLQueryCompiler


class GreenplumDataSourceMixin(BaseSQLDataSource):
    compiler_cls = PostgreSQLQueryCompiler

    conn_type = CONNECTION_TYPE_GREENPLUM

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type in (SOURCE_TYPE_GP_TABLE, SOURCE_TYPE_GP_SUBSELECT)


class GreenplumTableDataSource(GreenplumDataSourceMixin, StandardSchemaSQLDataSource):
    """Greenplum table"""


class GreenplumSubselectDataSource(GreenplumDataSourceMixin, SubselectDataSource):
    """Greenplum subselect"""

lib/dl_connector_greenplum/dl_connector_greenplum/core/data_source.py
datalens-tech/datalens-backend
from dl_core.connectors.sql_base.data_source_migration import DefaultSQLDataSourceMigrator

from dl_connector_greenplum.core.constants import (
    SOURCE_TYPE_GP_SUBSELECT,
    SOURCE_TYPE_GP_TABLE,
)


class GreenPlumDataSourceMigrator(DefaultSQLDataSourceMigrator):
    table_source_type = SOURCE_TYPE_GP_TABLE
    subselect_source_type = SOURCE_TYPE_GP_SUBSELECT
    default_schema_name = "public"

lib/dl_connector_greenplum/dl_connector_greenplum/core/data_source_migration.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_connector_greenplum.core.constants import CONNECTION_TYPE_GREENPLUM
from dl_connector_postgresql.core.postgresql_base.dto import PostgresConnDTOBase


@attr.s(frozen=True)
class GreenplumConnDTO(PostgresConnDTOBase):
    conn_type = CONNECTION_TYPE_GREENPLUM

lib/dl_connector_greenplum/dl_connector_greenplum/core/dto.py
datalens-tech/datalens-backend
from dl_connector_greenplum.core.us_connection import GreenplumConnection
from dl_connector_postgresql.core.postgresql_base.storage_schemas.connection import (
    ConnectionPostgreSQLBaseDataStorageSchema,
)


class GreenplumConnectionDataStorageSchema(ConnectionPostgreSQLBaseDataStorageSchema):
    TARGET_CLS = GreenplumConnection.DataModel

lib/dl_connector_greenplum/dl_connector_greenplum/core/storage_schemas/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import ClassVar

from dl_core.us_connection_base import DataSourceTemplate
from dl_i18n.localizer_base import Localizer

from dl_connector_greenplum.core.constants import (
    SOURCE_TYPE_GP_SUBSELECT,
    SOURCE_TYPE_GP_TABLE,
)
from dl_connector_greenplum.core.dto import GreenplumConnDTO
from dl_connector_postgresql.core.postgresql_base.us_connection import ConnectionPostgreSQLBase


class GreenplumConnection(ConnectionPostgreSQLBase):
    source_type = SOURCE_TYPE_GP_TABLE
    allowed_source_types = frozenset((SOURCE_TYPE_GP_TABLE, SOURCE_TYPE_GP_SUBSELECT))
    allow_dashsql: ClassVar[bool] = True
    allow_cache: ClassVar[bool] = True
    is_always_user_source: ClassVar[bool] = True

    def get_conn_dto(self) -> GreenplumConnDTO:
        return GreenplumConnDTO(
            conn_id=self.uuid,
            host=self.data.host,
            multihosts=self.parse_multihosts(),  # type: ignore  # TODO: fix
            port=self.data.port,
            db_name=self.data.db_name,
            username=self.data.username,
            password=self.password,
            enforce_collate=self.data.enforce_collate,
        )

    def get_data_source_template_templates(self, localizer: Localizer) -> list[DataSourceTemplate]:
        return self._make_subselect_templates(
            source_type=SOURCE_TYPE_GP_SUBSELECT,
            field_doc_key="PG_SUBSELECT/subsql",  # shared, currently.
            localizer=localizer,
        )

    @property
    def allow_public_usage(self) -> bool:
        return True

lib/dl_connector_greenplum/dl_connector_greenplum/core/us_connection.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_api_lib_testing.data_api_base import StandardizedDataApiTestBase
from dl_api_lib_testing.dataset_base import DatasetTestBase
from dl_constants.enums import RawSQLLevel
from dl_core_testing.database import DbTable
from dl_core_testing.testcases.service_base import ServiceFixtureTextClass

from dl_connector_greenplum.core.constants import (
    CONNECTION_TYPE_GREENPLUM,
    SOURCE_TYPE_GP_TABLE,
)
from dl_connector_greenplum_tests.db.config import (
    API_TEST_CONFIG,
    CONNECTION_PARAMS,
    CORE_TEST_CONFIG,
    DB_CORE_URL,
)


class GreenplumConnectionTestBase(ConnectionTestBase, ServiceFixtureTextClass):
    conn_type = CONNECTION_TYPE_GREENPLUM
    core_test_config = CORE_TEST_CONFIG
    bi_compeng_pg_on = False

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return DB_CORE_URL

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def connection_params(self) -> dict:
        return CONNECTION_PARAMS


class GreenplumDashSQLConnectionTest(GreenplumConnectionTestBase):
    @pytest.fixture(scope="class")
    def connection_params(self) -> dict:
        return CONNECTION_PARAMS | dict(raw_sql_level=RawSQLLevel.dashsql.value)


class GreenplumDatasetTestBase(GreenplumConnectionTestBase, DatasetTestBase):
    @pytest.fixture(scope="class")
    def dataset_params(self, sample_table: DbTable) -> dict:
        return dict(
            source_type=SOURCE_TYPE_GP_TABLE.name,
            parameters=dict(
                db_name=sample_table.db.name,
                table_name=sample_table.name,
            ),
        )


class GreenplumDataApiTestBase(GreenplumDatasetTestBase, StandardizedDataApiTestBase):
    mutation_caches_on = False

lib/dl_connector_greenplum/dl_connector_greenplum_tests/db/api/base.py
datalens-tech/datalens-backend
from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase
from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite

from dl_connector_greenplum_tests.db.api.base import GreenplumConnectionTestBase


class TestGreenplumSQLConnection(GreenplumConnectionTestBase, DefaultConnectorConnectionTestSuite):
    def test_create_connection(self, control_api_sync_client: SyncHttpClientBase, saved_connection_id: str) -> None:
        assert saved_connection_id
        response = control_api_sync_client.get(
            url=f"/api/v1/connections/{saved_connection_id}",
        )
        assert response.status_code == 200
        assert response.json["db_type"] == "greenplum"
        assert response.json["cache_ttl_sec"] is None

lib/dl_connector_greenplum/dl_connector_greenplum_tests/db/api/test_connection.py
datalens-tech/datalens-backend
from aiohttp.test_utils import TestClient
import pytest

from dl_api_lib_testing.connector.dashsql_suite import DefaultDashSQLTestSuite

from dl_connector_greenplum_tests.db.api.base import GreenplumDashSQLConnectionTest
from dl_connector_greenplum_tests.db.config import DASHSQL_QUERY


class TestGreenplumDashSQL(GreenplumDashSQLConnectionTest, DefaultDashSQLTestSuite):
    @pytest.mark.asyncio
    async def test_result(self, data_api_lowlevel_aiohttp_client: TestClient, saved_connection_id: str):
        await self.get_dashsql_response(
            data_api_aio=data_api_lowlevel_aiohttp_client,
            conn_id=saved_connection_id,
            query=DASHSQL_QUERY,
        )

lib/dl_connector_greenplum/dl_connector_greenplum_tests/db/api/test_dashsql.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.data_api_suites import (
    DefaultConnectorDataDistinctTestSuite,
    DefaultConnectorDataGroupByFormulaTestSuite,
    DefaultConnectorDataPreviewTestSuite,
    DefaultConnectorDataRangeTestSuite,
    DefaultConnectorDataResultTestSuite,
)
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_greenplum_tests.db.api.base import GreenplumDataApiTestBase


class TestGreenplumDataResult(GreenplumDataApiTestBase, DefaultConnectorDataResultTestSuite):
    pass


class TestGreenplumDataGroupBy(GreenplumDataApiTestBase, DefaultConnectorDataGroupByFormulaTestSuite):
    pass


class TestGreenplumDataRange(GreenplumDataApiTestBase, DefaultConnectorDataRangeTestSuite):
    pass


class TestGreenplumDataDistinct(GreenplumDataApiTestBase, DefaultConnectorDataDistinctTestSuite):
    pass


class TestGreenplumDataPreview(GreenplumDataApiTestBase, DefaultConnectorDataPreviewTestSuite):
    pass

lib/dl_connector_greenplum/dl_connector_greenplum_tests/db/api/test_data.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite

from dl_connector_greenplum_tests.db.api.base import GreenplumDatasetTestBase


class TestGreenplumDataset(GreenplumDatasetTestBase, DefaultConnectorDatasetTestSuite):
    pass

lib/dl_connector_greenplum/dl_connector_greenplum_tests/db/api/test_dataset.py
datalens-tech/datalens-backend
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport


# Infra settings
CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=51911).host,
    port_us_http=get_test_container_hostport("us", fallback_port=51911).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=51910).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=51910).port,
    us_master_token="AC1ofiek8coB",
    core_connector_ep_names=["greenplum"],
)

DB_CORE_URL = f'bi_postgresql://datalens:qwerty@{get_test_container_hostport("db-postgres-13", fallback_port=52301).as_pair()}/test_data'

API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["greenplum"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
)

DB_NAME = "test_data"
TABLE_NAME = "sample"
CONNECTION_PARAMS = dict(
    db_name=DB_NAME,
    host=get_test_container_hostport("db-postgres-13", fallback_port=52301).host,
    port=get_test_container_hostport("db-postgres-13", fallback_port=52301).port,
    username="datalens",
    password="qwerty",
)


DASHSQL_QUERY = r"""
with base as (
    select generate_series(0, 6) as number
)
select
    number,
    'test' || number::text as str,
    number::bool as num_bool,
    number::text::bytea as num_bytea,
    number::char as num_char,
    -- number::name as num_name,
    number::int8 as num_int8,
    number::int2 as num_int2,
    number::int4 as num_int4,
    number::text as num_text,
    number::oid as num_oid,
    number::text::json as num_json,
    -- number::text::jsonb as num_jsonb, -- tests-pg too old
    number::float4 as num_float4,
    number::float8 as num_float8,
    number::numeric as num_numeric,
    (number || ' second')::interval as num_interval,
    number::varchar(12) as num_varchar,
    ('2020-01-0' || (number + 1))::date as num_date,
    ('00:01:0' || number)::time as num_time,
    ('2020-01-01T00:00:0' || number)::timestamp as num_timestamp,
    ('2020-01-01T00:00:0' || number)::timestamptz as num_timestamptz,
    ARRAY[number, 2, 3] as num_array,
    'nan'::double precision + number as some_nan
from base
limit 10
"""

lib/dl_connector_greenplum/dl_connector_greenplum_tests/db/config.py
datalens-tech/datalens-backend
from dl_api_lib_testing.initialization import initialize_api_lib_test

from dl_connector_greenplum_tests.db.config import API_TEST_CONFIG


pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)

lib/dl_connector_greenplum/dl_connector_greenplum_tests/db/conftest.py
datalens-tech/datalens-backend
from dl_api_connector.i18n.localizer import CONFIGS as BI_API_CONNECTOR_CONFIGS
from dl_api_lib_testing.connection_form_base import ConnectionFormTestBase

from dl_connector_greenplum.api.connection_form.form_config import GreenplumConnectionFormFactory
from dl_connector_greenplum.api.i18n.localizer import CONFIGS as BI_CONNECTOR_GREENPLUM_CONFIGS


class TestGreenplumConnectionForm(ConnectionFormTestBase):
    CONN_FORM_FACTORY_CLS = GreenplumConnectionFormFactory
    TRANSLATION_CONFIGS = BI_API_CONNECTOR_CONFIGS + BI_CONNECTOR_GREENPLUM_CONFIGS

lib/dl_connector_greenplum/dl_connector_greenplum_tests/unit/test_connection_form.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import fields as ma_fields

from dl_api_lib import exc
from dl_api_lib.api_decorators import schematic_request
from dl_api_lib.app.control_api.resources.base import BIResource
from dl_api_lib.app.control_api.resources.connections import ns
from dl_api_lib.enums import USPermissionKind
from dl_api_lib.utils import need_permission_on_entry
from dl_app_tools.profiling_base import generic_profiler
from dl_core.exc import ConnectionConfigurationError
from dl_core.us_connection_base import ConnectionBase
from dl_model_tools.schema.base import BaseSchema
from dl_sqlalchemy_metrica_api.exceptions import MetrikaApiAccessDeniedException

from dl_connector_metrica.core.constants import (
    CONNECTION_TYPE_APPMETRICA_API,
    CONNECTION_TYPE_METRICA_API,
)
from dl_connector_metrica.core.us_connection import MetrikaBaseMixin


class MetricaConnectionAvailableCountersSchema(BaseSchema):
    class CounterInfo(BaseSchema):
        id = ma_fields.String()
        name = ma_fields.String()

    counters = ma_fields.Nested(CounterInfo, many=True)


@ns.route("/<connection_id>/metrica_available_counters")
class MetricaConnectionAvailableCounters(BIResource):
    @schematic_request(ns=ns, responses={200: ("Success", MetricaConnectionAvailableCountersSchema())})
    @generic_profiler("metrica-available-counters")
    def get(self, connection_id: str) -> dict:
        conn = self.get_us_manager().get_by_id(connection_id, expected_type=ConnectionBase)
        need_permission_on_entry(conn, USPermissionKind.edit)
        if conn.conn_type not in (CONNECTION_TYPE_METRICA_API, CONNECTION_TYPE_APPMETRICA_API):
            raise exc.UnsupportedForEntityType("Unsupported connection type")
        assert isinstance(conn, MetrikaBaseMixin)
        try:
            counters = conn.get_available_counters()
        except MetrikaApiAccessDeniedException as ex:
            raise ConnectionConfigurationError(
                "Unable to load available counters. Possibly caused by invalid OAuth token."
            ) from ex
        else:
            return dict(counters=counters)

lib/dl_connector_metrica/dl_connector_metrica/api/api_handler.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Optional,
)

from marshmallow import ValidationError
from marshmallow import fields as ma_fields
from marshmallow import validates_schema

from dl_api_connector.api_schema.connection_base import (
    ConnectionMetaMixin,
    ConnectionSchema,
)
from dl_api_connector.api_schema.connection_base_fields import secret_string_field
from dl_api_connector.api_schema.extras import FieldExtra

from dl_connector_metrica.core.us_connection import (
    AppMetricaApiConnection,
    MetrikaApiConnection,
    parse_metrica_ids,
)


class ConnectionMetrikaLikeAPI(ConnectionMetaMixin, ConnectionSchema):
    token = secret_string_field(attribute="data.token")
    counter_id = ma_fields.String(attribute="data.counter_id", required=True, bi_extra=FieldExtra(editable=True))
    accuracy = ma_fields.Float(
        attribute="data.accuracy",
        allow_none=True,
        dump_default=None,
        load_default=None,
        bi_extra=FieldExtra(editable=True),
    )

    @validates_schema
    def validate_counter_id(self, data: Optional[dict[str, Any]], *args: Any, **kwargs: Any) -> None:
        if data is None or "data" not in data or "counter_id" not in data["data"]:
            return

        ids_orig = data["data"]["counter_id"]
        ids = list(filter(lambda t: t, parse_metrica_ids(ids_orig)))
        if ids:
            for id_str in ids:
                try:
                    id_value = int(id_str)
                except Exception as ex:
                    raise ValidationError(f"Unable to parse id: {id_str!r}") from ex
                if id_value <= 0:
                    raise ValidationError(f"Value should be positive: {id_str!r}")
        else:
            raise ValidationError(f"Unable to parse id: {ids_orig!r}")


class ConnectionMetrikaAPISchema(ConnectionMetrikaLikeAPI):
    TARGET_CLS = MetrikaApiConnection


class ConnectionAppMetricaAPISchema(ConnectionMetrikaLikeAPI):
    TARGET_CLS = AppMetricaApiConnection

lib/dl_connector_metrica/dl_connector_metrica/api/api_schema/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from enum import unique
from typing import (
    ClassVar,
    Optional,
)

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormActionApiSchema,
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
)
from dl_api_connector.form_config.models.common import (
    CommonFieldName,
    FormFieldName,
    OAuthApplication,
)
import dl_api_connector.form_config.models.rows as C
from dl_api_connector.form_config.models.rows.base import FormRow
from dl_api_connector.form_config.models.shortcuts.rows import RowConstructor
from dl_configs.connectors_settings import ConnectorSettingsBase

from dl_connector_metrica.api.connection_form.rows import (
    AccuracyRow,
    AppMetricaCounterRowItem,
    CounterRow,
    MetricaCounterRowItem,
)
from dl_connector_metrica.api.connection_info import (
    AppMetricaConnectionInfoProvider,
    MetricaConnectionInfoProvider,
)
from dl_connector_metrica.api.i18n.localizer import Translatable
from dl_connector_metrica.core.settings import (
    AppmetricaConnectorSettings,
    MetricaConnectorSettings,
)


class MetricaOAuthApplication(OAuthApplication):
    metrika_api = "metrika_api"
    appmetrica_api = "appmetrica_api"


@unique
class MetricaFieldName(FormFieldName):
    counter_id = "counter_id"
    accuracy = "accuracy"


class MetricaLikeBaseFormFactory(ConnectionFormFactory, metaclass=abc.ABCMeta):
    template_name: ClassVar[str]
    oauth_application: ClassVar[OAuthApplication]

    @abc.abstractmethod
    def _title(self) -> str:
        raise NotImplementedError

    @abc.abstractmethod
    def _counter_row(self, manual_input: bool) -> CounterRow:
        raise NotImplementedError

    @abc.abstractmethod
    def _allow_manual_counter_input(self, connector_settings: ConnectorSettingsBase) -> bool:
        raise NotImplementedError

    @abc.abstractmethod
    def _allow_auto_dash_creation(self, connector_settings: ConnectorSettingsBase) -> bool:
        raise NotImplementedError

    def get_form_config(
        self,
        connector_settings: Optional[ConnectorSettingsBase],
        tenant: Optional[TenantDef],
    ) -> ConnectionForm:
        assert connector_settings is not None
        rc = RowConstructor(localizer=self._localizer)

        rows: list[FormRow] = [
            C.OAuthTokenRow(
                name=CommonFieldName.token,
                fake_value="******" if self.mode == ConnectionFormMode.edit else None,
                application=self.oauth_application,
                label_text=self._localizer.translate(Translatable("field_oauth-token")),
                button_text=self._localizer.translate(Translatable("button_get-token")),
            ),
            self._counter_row(manual_input=self._allow_manual_counter_input(connector_settings)),
            AccuracyRow(name=MetricaFieldName.accuracy),
        ]

        edit_api_schema = (
            FormActionApiSchema(
                items=[
                    FormFieldApiSchema(name=MetricaFieldName.counter_id, required=True),
                    FormFieldApiSchema(name=CommonFieldName.token),
                    FormFieldApiSchema(name=MetricaFieldName.accuracy, nullable=True),
                ]
            )
            if self.mode == ConnectionFormMode.edit
            else None
        )

        create_api_schema = (
            FormActionApiSchema(
                items=[
                    FormFieldApiSchema(name=MetricaFieldName.counter_id, required=True),
                    FormFieldApiSchema(name=CommonFieldName.token, required=True),
                    FormFieldApiSchema(name=MetricaFieldName.accuracy, nullable=True),
                    *self._get_top_level_create_api_schema_items(),
                ]
            )
            if self.mode == ConnectionFormMode.create
            else None
        )

        if self.mode == ConnectionFormMode.create and self._allow_auto_dash_creation(connector_settings):
            rows.append(rc.auto_create_dash_row())

        return ConnectionForm(
            title=self._title(),
            template_name=self.template_name,
            rows=rows,
            api_schema=FormApiSchema(create=create_api_schema, edit=edit_api_schema),
        )


class MetricaAPIConnectionFormFactory(MetricaLikeBaseFormFactory):
    template_name = "metrica_api"
    oauth_application = MetricaOAuthApplication.metrika_api

    def _title(self) -> str:
        return MetricaConnectionInfoProvider.get_title(self._localizer)

    def _counter_row(self, manual_input: bool) -> MetricaCounterRowItem:
        return MetricaCounterRowItem(
            name=MetricaFieldName.counter_id,
            allow_manual_input=manual_input,
        )

    def _allow_manual_counter_input(self, connector_settings: ConnectorSettingsBase) -> bool:
        assert isinstance(connector_settings, MetricaConnectorSettings)
        return connector_settings.COUNTER_ALLOW_MANUAL_INPUT

    def _allow_auto_dash_creation(self, connector_settings: ConnectorSettingsBase) -> bool:
        assert isinstance(connector_settings, MetricaConnectorSettings)
        return connector_settings.ALLOW_AUTO_DASH_CREATION


class AppMetricaAPIConnectionFormFactory(MetricaLikeBaseFormFactory):
    template_name = "appmetrica_api"
    oauth_application = MetricaOAuthApplication.appmetrica_api

    def _title(self) -> str:
        return AppMetricaConnectionInfoProvider.get_title(self._localizer)

    def _counter_row(self, manual_input: bool) -> AppMetricaCounterRowItem:
        return AppMetricaCounterRowItem(
            name=MetricaFieldName.counter_id,
            allow_manual_input=manual_input,
        )

    def _allow_manual_counter_input(self, connector_settings: ConnectorSettingsBase) -> bool:
        assert isinstance(connector_settings, AppmetricaConnectorSettings)
        return connector_settings.COUNTER_ALLOW_MANUAL_INPUT

    def _allow_auto_dash_creation(self, connector_settings: ConnectorSettingsBase) -> bool:
        assert isinstance(connector_settings, AppmetricaConnectorSettings)
        return connector_settings.ALLOW_AUTO_DASH_CREATION

lib/dl_connector_metrica/dl_connector_metrica/api/connection_form/form_config.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_api_connector.form_config.models.common import remap_skip_if_null
from dl_api_connector.form_config.models.rows.base import FormFieldMixin
from dl_api_connector.form_config.models.rows.prepared.base import (
    DisabledMixin,
    PreparedRow,
)


@attr.s(kw_only=True, frozen=True)
class CounterRow(PreparedRow, FormFieldMixin, DisabledMixin):
    label_text: Optional[str] = attr.ib(default=None, metadata=remap_skip_if_null("labelText"))
    allow_manual_input: Optional[bool] = attr.ib(default=None, metadata=remap_skip_if_null("allowManualInput"))

    class Inner(PreparedRow.Inner):
        counter_input_method = "counter_input_method"


@attr.s(kw_only=True, frozen=True)
class MetricaCounterRowItem(CounterRow):
    type = "metrica_counter"


@attr.s(kw_only=True, frozen=True)
class AppMetricaCounterRowItem(CounterRow):
    type = "appmetrica_counter"


@attr.s(kw_only=True, frozen=True)
class AccuracyRow(PreparedRow, FormFieldMixin, DisabledMixin):
    type = "metrica_accuracy"

lib/dl_connector_metrica/dl_connector_metrica/api/connection_form/rows.py
datalens-tech/datalens-backend
from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_metrica.api.i18n.localizer import Translatable


class MetricaConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-metrica")
    alias = "metrica"


class AppMetricaConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-appmetrica")
    alias = "appmetrica"

lib/dl_connector_metrica/dl_connector_metrica/api/connection_info.py
datalens-tech/datalens-backend
from dl_api_connector.api_schema.source_base import (
    SQLDataSourceSchema,
    SQLDataSourceTemplateSchema,
)
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)

from dl_connector_metrica.api.api_schema.connection import (
    ConnectionAppMetricaAPISchema,
    ConnectionMetrikaAPISchema,
)
from dl_connector_metrica.api.connection_form.form_config import (
    AppMetricaAPIConnectionFormFactory,
    MetricaAPIConnectionFormFactory,
)
from dl_connector_metrica.api.connection_info import (
    AppMetricaConnectionInfoProvider,
    MetricaConnectionInfoProvider,
)
from dl_connector_metrica.api.filter_compiler import MetricaApiFilterFormulaCompiler
from dl_connector_metrica.api.i18n.localizer import CONFIGS
from dl_connector_metrica.core.connector import (
    AppMetricaApiCoreConnectionDefinition,
    AppMetricaApiCoreConnector,
    AppMetricaApiCoreSourceDefinition,
    MetricaApiCoreConnectionDefinition,
    MetricaApiCoreConnector,
    MetricaApiCoreSourceDefinition,
)
from dl_connector_metrica.formula.constants import DIALECT_NAME_METRICAAPI


class MetricaApiFilteredApiTableSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = MetricaApiCoreSourceDefinition
    api_schema_cls = SQLDataSourceSchema
    template_api_schema_cls = SQLDataSourceTemplateSchema


class MetricaApiApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = MetricaApiCoreConnectionDefinition
    api_generic_schema_cls = ConnectionMetrikaAPISchema
    alias = "metrica"
    info_provider_cls = MetricaConnectionInfoProvider
    form_factory_cls = MetricaAPIConnectionFormFactory


class MetricaApiApiConnector(ApiConnector):
    core_connector_cls = MetricaApiCoreConnector
    connection_definitions = (MetricaApiApiConnectionDefinition,)
    source_definitions = (MetricaApiFilteredApiTableSourceDefinition,)
    filter_formula_compiler_cls = MetricaApiFilterFormulaCompiler
    formula_dialect_name = DIALECT_NAME_METRICAAPI
    translation_configs = frozenset(CONFIGS)


class AppMetricaApiFilteredApiTableSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = AppMetricaApiCoreSourceDefinition
    api_schema_cls = SQLDataSourceSchema
    template_api_schema_cls = SQLDataSourceTemplateSchema


class AppMetricaApiApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = AppMetricaApiCoreConnectionDefinition
    api_generic_schema_cls = ConnectionAppMetricaAPISchema
    alias = "appmetrica"
    info_provider_cls = AppMetricaConnectionInfoProvider
    form_factory_cls = AppMetricaAPIConnectionFormFactory


class AppMetricaApiApiConnector(ApiConnector):
    core_connector_cls = AppMetricaApiCoreConnector
    connection_definitions = (AppMetricaApiApiConnectionDefinition,)
    source_definitions = (AppMetricaApiFilteredApiTableSourceDefinition,)
    filter_formula_compiler_cls = MetricaApiFilterFormulaCompiler
    formula_dialect_name = DIALECT_NAME_METRICAAPI
    translation_configs = frozenset(CONFIGS)

lib/dl_connector_metrica/dl_connector_metrica/api/connector.py
datalens-tech/datalens-backend
from dl_query_processing.compilation.filter_compiler import (
    FilterParams,
    MainFilterFormulaCompiler,
)


class MetricaApiFilterFormulaCompiler(MainFilterFormulaCompiler):
    """Does not support datetime casting used for most other sources"""

    def _mangle_date_filter(self, filter_params: FilterParams) -> FilterParams:
        return filter_params

lib/dl_connector_metrica/dl_connector_metrica/api/filter_compiler.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_metrica as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_metrica/dl_connector_metrica/api/i18n/localizer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
    Optional,
    Tuple,
    Type,
    TypeVar,
)
from urllib.parse import (
    quote_plus,
    urlencode,
)

import sqlalchemy as sa
from sqlalchemy.engine import Engine

from dl_core import exc
from dl_core.connection_executors.adapters.adapters_base_sa import BaseSAAdapter
from dl_core.connection_models import DBIdent
from dl_core.connectors.base.error_transformer import DBExcKWArgs
from dl_sqlalchemy_metrica_api import exceptions as sqla_metrika_exc  # type: ignore

from dl_connector_metrica.core.constants import (
    CONNECTION_TYPE_APPMETRICA_API,
    CONNECTION_TYPE_METRICA_API,
)
from dl_connector_metrica.core.exc import MetricaAPIDatabaseQueryError
from dl_connector_metrica.core.target_dto import (
    AppMetricaAPIConnTargetDTO,
    MetricaAPIConnTargetDTO,
)


_M_CONN_T_DTO_TV = TypeVar("_M_CONN_T_DTO_TV", bound=MetricaAPIConnTargetDTO)


class MetricaAPIDefaultAdapter(BaseSAAdapter[_M_CONN_T_DTO_TV]):
    conn_type = CONNECTION_TYPE_METRICA_API

    def _get_db_engine(self, db_name: str, disable_streaming: bool = False) -> Engine:
        if disable_streaming:
            raise Exception("`disable_streaming` is not applicable here")
        dsn = "{dialect}://:{token}@/{db_name}".format(
            dialect=self.get_dialect_str(),
            token=quote_plus(self._target_dto.token),
            db_name=db_name,
        )
        dsn_params: Dict[str, Any] = {}
        if self._target_dto.accuracy is not None:
            dsn_params.update(accuracy=self._target_dto.accuracy)

        if dsn_params:
            dsn += "?" + urlencode(dsn_params)

        return sa.create_engine(dsn).execution_options(compiled_cache=None)

    @classmethod
    def make_exc(  # TODO:  Move to ErrorTransformer
        cls, wrapper_exc: Exception, orig_exc: Optional[Exception], debug_compiled_query: Optional[str]
    ) -> Tuple[Type[exc.DatabaseQueryError], DBExcKWArgs]:
        exc_cls, kw = super().make_exc(wrapper_exc, orig_exc, debug_compiled_query)

        if isinstance(
            orig_exc,
            (
                sqla_metrika_exc.MetrikaHttpApiException,
                sqla_metrika_exc.NotSupportedError,
                sqla_metrika_exc.ProgrammingError,
            ),
        ):
            exc_cls = MetricaAPIDatabaseQueryError

        return exc_cls, kw

    def get_default_db_name(self) -> Optional[str]:
        return None

    def _get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        return ""


class AppMetricaAPIDefaultAdapter(MetricaAPIDefaultAdapter[AppMetricaAPIConnTargetDTO]):  # type: ignore
    conn_type = CONNECTION_TYPE_APPMETRICA_API

lib/dl_connector_metrica/dl_connector_metrica/core/adapters_metrica_x.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor

from dl_connector_metrica.core.adapters_metrica_x import (
    AppMetricaAPIDefaultAdapter,
    MetricaAPIDefaultAdapter,
)
from dl_connector_metrica.core.dto import (
    AppMetricaAPIConnDTO,
    MetricaAPIConnDTO,
)
from dl_connector_metrica.core.target_dto import (
    AppMetricaAPIConnTargetDTO,
    MetricaAPIConnTargetDTO,
)


@attr.s(cmp=False, hash=False)
class MetricaAPIConnExecutor(DefaultSqlAlchemyConnExecutor[MetricaAPIDefaultAdapter]):
    TARGET_ADAPTER_CLS = MetricaAPIDefaultAdapter

    _conn_dto: MetricaAPIConnDTO = attr.ib()

    async def _make_target_conn_dto_pool(self) -> list[MetricaAPIConnTargetDTO]:
        return [
            MetricaAPIConnTargetDTO(
                conn_id=self._conn_dto.conn_id,
                pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                accuracy=self._conn_dto.accuracy,
                token=self._conn_dto.token,
            )
        ]


@attr.s(cmp=False, hash=False)
class AppMetricaAPIConnExecutor(DefaultSqlAlchemyConnExecutor[AppMetricaAPIDefaultAdapter]):
    TARGET_ADAPTER_CLS = AppMetricaAPIDefaultAdapter

    _conn_dto: AppMetricaAPIConnDTO = attr.ib()

    async def _make_target_conn_dto_pool(self) -> list[AppMetricaAPIConnTargetDTO]:
        return [
            AppMetricaAPIConnTargetDTO(
                conn_id=self._conn_dto.conn_id,
                pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                accuracy=self._conn_dto.accuracy,
                token=self._conn_dto.token,
            )
        ]

lib/dl_connector_metrica/dl_connector_metrica/core/connection_executors.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_core.connections_security.base import (
    ConnSecuritySettings,
    NonUserInputConnectionSafetyChecker,
)
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
    CoreSourceDefinition,
)
from dl_core.data_source_spec.sql import StandardSQLDataSourceSpec
from dl_core.us_manager.storage_schemas.data_source_spec_base import SQLDataSourceSpecStorageSchema

from dl_connector_metrica.core.adapters_metrica_x import (
    AppMetricaAPIDefaultAdapter,
    MetricaAPIDefaultAdapter,
)
from dl_connector_metrica.core.connection_executors import (
    AppMetricaAPIConnExecutor,
    MetricaAPIConnExecutor,
)
from dl_connector_metrica.core.constants import (
    BACKEND_TYPE_APPMETRICA_API,
    BACKEND_TYPE_METRICA_API,
    CONNECTION_TYPE_APPMETRICA_API,
    CONNECTION_TYPE_METRICA_API,
    SOURCE_TYPE_APPMETRICA_API,
    SOURCE_TYPE_METRICA_API,
)
from dl_connector_metrica.core.data_source import (
    AppMetrikaApiDataSource,
    MetrikaApiDataSource,
)
from dl_connector_metrica.core.data_source_migration import (
    AppMetricaApiDataSourceMigrator,
    MetricaApiDataSourceMigrator,
)
from dl_connector_metrica.core.dto import (
    AppMetricaAPIConnDTO,
    MetricaAPIConnDTO,
)
from dl_connector_metrica.core.lifecycle import MetricaConnectionLifecycleManager
from dl_connector_metrica.core.settings import (
    AppMetricaSettingDefinition,
    MetricaSettingDefinition,
)
from dl_connector_metrica.core.storage_schemas.connection import (
    ConnectionAppMetricaApiDataStorageSchema,
    ConnectionMetrikaApiDataStorageSchema,
)
from dl_connector_metrica.core.type_transformer import MetrikaApiTypeTransformer
from dl_connector_metrica.core.us_connection import (
    AppMetricaApiConnection,
    MetrikaApiConnection,
)


class MetricaApiCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_METRICA_API
    connection_cls = MetrikaApiConnection
    us_storage_schema_cls = ConnectionMetrikaApiDataStorageSchema
    type_transformer_cls = MetrikaApiTypeTransformer
    sync_conn_executor_cls = MetricaAPIConnExecutor
    async_conn_executor_cls = MetricaAPIConnExecutor
    lifecycle_manager_cls = MetricaConnectionLifecycleManager
    dialect_string = "metrika_api"
    settings_definition = MetricaSettingDefinition
    data_source_migrator_cls = MetricaApiDataSourceMigrator


class MetricaApiCoreSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_METRICA_API
    source_cls = MetrikaApiDataSource
    source_spec_cls = StandardSQLDataSourceSpec
    us_storage_schema_cls = SQLDataSourceSpecStorageSchema


class MetricaApiCoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_METRICA_API
    connection_definitions = (MetricaApiCoreConnectionDefinition,)
    source_definitions = (MetricaApiCoreSourceDefinition,)
    rqe_adapter_classes = frozenset({MetricaAPIDefaultAdapter})
    conn_security = frozenset(
        {
            ConnSecuritySettings(NonUserInputConnectionSafetyChecker, frozenset({MetricaAPIConnDTO})),
        }
    )


class AppMetricaApiCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_APPMETRICA_API
    connection_cls = AppMetricaApiConnection
    us_storage_schema_cls = ConnectionAppMetricaApiDataStorageSchema
    type_transformer_cls = MetrikaApiTypeTransformer
    sync_conn_executor_cls = AppMetricaAPIConnExecutor
    async_conn_executor_cls = AppMetricaAPIConnExecutor
    lifecycle_manager_cls = MetricaConnectionLifecycleManager
    dialect_string = "appmetrica_api"
    settings_definition = AppMetricaSettingDefinition
    data_source_migrator_cls = AppMetricaApiDataSourceMigrator


class AppMetricaApiCoreSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_APPMETRICA_API
    source_cls = AppMetrikaApiDataSource
    source_spec_cls = StandardSQLDataSourceSpec
    us_storage_schema_cls = SQLDataSourceSpecStorageSchema


class AppMetricaApiCoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_APPMETRICA_API
    connection_definitions = (AppMetricaApiCoreConnectionDefinition,)
    source_definitions = (AppMetricaApiCoreSourceDefinition,)
    rqe_adapter_classes = frozenset({AppMetricaAPIDefaultAdapter})
    conn_security = frozenset(
        {
            ConnSecuritySettings(NonUserInputConnectionSafetyChecker, frozenset({AppMetricaAPIConnDTO})),
        }
    )

lib/dl_connector_metrica/dl_connector_metrica/core/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
    SourceBackendType,
)


BACKEND_TYPE_METRICA_API = SourceBackendType.declare("METRICA_API")
CONNECTION_TYPE_METRICA_API = ConnectionType.declare("metrika_api")  # Note the K in the value
SOURCE_TYPE_METRICA_API = DataSourceType.declare("METRIKA_API")  # Note the K in the value

BACKEND_TYPE_APPMETRICA_API = SourceBackendType.declare("APPMETRICA_API")
CONNECTION_TYPE_APPMETRICA_API = ConnectionType.declare("appmetrica_api")
SOURCE_TYPE_APPMETRICA_API = DataSourceType.declare("APPMETRICA_API")

lib/dl_connector_metrica/dl_connector_metrica/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Optional,
    Tuple,
)

from dl_constants.enums import UserDataType
from dl_core import exc
from dl_core.data_source.sql import PseudoSQLDataSource
from dl_core.db import (
    SchemaColumn,
    SchemaInfo,
)

from dl_connector_metrica.core.constants import (
    CONNECTION_TYPE_APPMETRICA_API,
    CONNECTION_TYPE_METRICA_API,
)
from dl_connector_metrica.core.us_connection import MetrikaApiConnection


if TYPE_CHECKING:
    from dl_core.connection_executors.async_base import AsyncConnExecutorBase
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase


LOGGER = logging.getLogger(__name__)


class MetrikaApiDataSource(PseudoSQLDataSource):
    store_raw_schema: ClassVar[bool] = False
    preview_enabled: ClassVar[bool] = False
    supports_schema_update: ClassVar[bool] = False

    conn_type = CONNECTION_TYPE_METRICA_API

    def get_schema_info(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> SchemaInfo:
        assert self.saved_raw_schema is not None
        return SchemaInfo.from_schema(self.saved_raw_schema)

    def _check_existence(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> bool:
        return True

    async def _check_existence_async(self, conn_executor_factory: Callable[[], AsyncConnExecutorBase]) -> bool:
        return True

    @property
    def saved_raw_schema(self) -> Optional[list[SchemaColumn]]:
        assert self.conn_type is not None
        db_name = self.db_name
        assert db_name is not None
        conn_cls = self.get_connection_cls()
        assert issubclass(conn_cls, MetrikaApiConnection)
        return [
            sch_column.clone(source_id=self.id)
            for sch_column in conn_cls.get_raw_schema(
                metrica_namespace=db_name,
                actual_conn_type=self.conn_type,
            )
        ]

    def get_expression_value_range(self, col_name: str) -> Tuple[Any, Any]:
        """Date/datetime column value ranges are defined as ``(<counter creation>, <now>)``"""
        try:
            assert self.saved_raw_schema is not None
            column = next(col for col in self.saved_raw_schema if col.name == col_name)
        except StopIteration:
            raise exc.InvalidColumnError("Invalid field name")

        if column.user_type not in (UserDataType.date, UserDataType.datetime, UserDataType.genericdatetime):
            raise exc.InvalidColumnError("Invalid field for value range")

        creation_date = self.connection.data.counter_creation_date
        now = datetime.datetime.utcnow()
        if column.user_type in (UserDataType.datetime, UserDataType.genericdatetime):
            min_value = datetime.datetime(creation_date.year, creation_date.month, creation_date.day)
            max_value = now
        else:
            min_value = creation_date
            max_value = now.date()  # type: ignore  # TODO: fix

        return min_value, max_value


class AppMetrikaApiDataSource(MetrikaApiDataSource):
    conn_type = CONNECTION_TYPE_APPMETRICA_API

lib/dl_connector_metrica/dl_connector_metrica/core/data_source.py
datalens-tech/datalens-backend
from dl_core.connectors.sql_base.data_source_migration import DefaultSQLDataSourceMigrator

from dl_connector_metrica.core.constants import (
    SOURCE_TYPE_APPMETRICA_API,
    SOURCE_TYPE_METRICA_API,
)


class MetricaApiDataSourceMigrator(DefaultSQLDataSourceMigrator):
    table_source_type = SOURCE_TYPE_METRICA_API


class AppMetricaApiDataSourceMigrator(DefaultSQLDataSourceMigrator):
    table_source_type = SOURCE_TYPE_APPMETRICA_API
    with_db_name = True

lib/dl_connector_metrica/dl_connector_metrica/core/data_source_migration.py
datalens-tech/datalens-backend
import attr

from dl_core.connection_models.dto_defs import ConnDTO

from dl_connector_metrica.core.constants import (
    CONNECTION_TYPE_APPMETRICA_API,
    CONNECTION_TYPE_METRICA_API,
)


@attr.s(frozen=True)
class MetricaAPIConnDTO(ConnDTO):  # noqa
    conn_type = CONNECTION_TYPE_METRICA_API

    token: str = attr.ib(repr=False, kw_only=True)
    accuracy: float = attr.ib(kw_only=True)


@attr.s(frozen=True)
class AppMetricaAPIConnDTO(MetricaAPIConnDTO):
    conn_type = CONNECTION_TYPE_APPMETRICA_API

lib/dl_connector_metrica/dl_connector_metrica/core/dto.py
datalens-tech/datalens-backend
from dl_core.exc import DatabaseQueryError


class MetricaAPIDatabaseQueryError(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["METRICA"]
    default_message = "Metrica API error."

lib/dl_connector_metrica/dl_connector_metrica/core/exc.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging

from dl_core import exc
from dl_core.connectors.base.lifecycle import ConnectionLifecycleManager
from dl_sqlalchemy_metrica_api.exceptions import MetrikaApiAccessDeniedException

from dl_connector_metrica.core.us_connection import MetrikaBaseMixin


LOGGER = logging.getLogger(__name__)


class MetricaConnectionLifecycleManager(ConnectionLifecycleManager[MetrikaBaseMixin]):
    ENTRY_CLS = MetrikaBaseMixin

    # TODO FIX: split into sync and async hooks
    def pre_save_hook(self) -> None:
        super().pre_save_hook()

        if self.entry.counter_id != self.entry._initial_counter_id or not self.entry.data.counter_creation_date:  # type: ignore  # TODO: fix
            LOGGER.info(
                "initial counter_id = %s, current counter_id = %s. "
                "Retrieving current counter creation date from Metrika API",
                self.entry._initial_counter_id,
                self.entry.counter_id,  # noqa
            )
            try:
                self.entry.data.counter_creation_date = self.entry.get_counter_creation_date()  # type: ignore  # TODO: fix
            except MetrikaApiAccessDeniedException as ex:
                raise exc.ConnectionConfigurationError("No access to counter info. Check your OAuth token.") from ex

lib/dl_connector_metrica/dl_connector_metrica/core/lifecycle.py
datalens-tech/datalens-backend
import attr

from dl_configs.connectors_settings import (
    ConnectorsConfigType,
    ConnectorSettingsBase,
)
from dl_configs.settings_loaders.meta_definition import s_attrib
from dl_core.connectors.settings.primitives import ConnectorSettingsDefinition


@attr.s(frozen=True)
class MetricaConnectorSettings(ConnectorSettingsBase):
    COUNTER_ALLOW_MANUAL_INPUT: bool = s_attrib("COUNTER_ALLOW_MANUAL_INPUT", missing=False)  # type: ignore
    ALLOW_AUTO_DASH_CREATION: bool = s_attrib("ALLOW_AUTO_DASH_CREATION", missing=False)  # type: ignore


@attr.s(frozen=True)
class AppmetricaConnectorSettings(ConnectorSettingsBase):
    COUNTER_ALLOW_MANUAL_INPUT: bool = s_attrib("COUNTER_ALLOW_MANUAL_INPUT", missing=False)  # type: ignore
    ALLOW_AUTO_DASH_CREATION: bool = s_attrib("ALLOW_AUTO_DASH_CREATION", missing=False)  # type: ignore


def metrica_settings_fallback(full_cfg: ConnectorsConfigType) -> dict[str, ConnectorSettingsBase]:
    return dict(METRIKA_API=MetricaConnectorSettings())


class MetricaSettingDefinition(ConnectorSettingsDefinition):
    settings_class = MetricaConnectorSettings
    fallback = metrica_settings_fallback


def appmetrica_settings_fallback(full_cfg: ConnectorsConfigType) -> dict[str, ConnectorSettingsBase]:
    return dict(APPMETRICA_API=AppmetricaConnectorSettings())


class AppMetricaSettingDefinition(ConnectorSettingsDefinition):
    settings_class = AppmetricaConnectorSettings
    fallback = appmetrica_settings_fallback

lib/dl_connector_metrica/dl_connector_metrica/core/settings.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_core.us_manager.storage_schemas.connection import ConnectionBaseDataStorageSchema

from dl_connector_metrica.core.us_connection import (
    AppMetricaApiConnection,
    MetrikaApiConnection,
)


class ConnectionMetrikaApiDataStorageSchema(ConnectionBaseDataStorageSchema[MetrikaApiConnection.DataModel]):
    TARGET_CLS = MetrikaApiConnection.DataModel

    token = ma_fields.String(required=True, allow_none=False)
    counter_id = ma_fields.String(required=True, allow_none=False)
    counter_creation_date = ma_fields.Date(required=False, allow_none=True, load_default=None, dump_default=None)
    accuracy = ma_fields.Float(required=False, allow_none=True, load_default=None, dump_default=None)


class ConnectionAppMetricaApiDataStorageSchema(ConnectionMetrikaApiDataStorageSchema):
    TARGET_CLS = AppMetricaApiConnection.DataModel

lib/dl_connector_metrica/dl_connector_metrica/core/storage_schemas/connection.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO
from dl_core.utils import secrepr


@attr.s(frozen=True)
class BaseMetricaAPIConnTargetDTO(ConnTargetDTO):
    token: str = attr.ib(repr=secrepr)
    accuracy: float = attr.ib()

    def get_effective_host(self) -> Optional[str]:
        return None


@attr.s(frozen=True)
class MetricaAPIConnTargetDTO(BaseMetricaAPIConnTargetDTO):
    pass


@attr.s(frozen=True)
class AppMetricaAPIConnTargetDTO(BaseMetricaAPIConnTargetDTO):
    pass

lib/dl_connector_metrica/dl_connector_metrica/core/target_dto.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Any
import uuid

from dl_core.us_manager.us_manager_sync import SyncUSManager

from dl_connector_metrica.core.constants import (
    CONNECTION_TYPE_APPMETRICA_API,
    CONNECTION_TYPE_METRICA_API,
)
from dl_connector_metrica.core.us_connection import (
    AppMetricaApiConnection,
    MetrikaApiConnection,
)


def make_saved_metrika_api_connection(
    sync_usm: SyncUSManager,
    counter_id: str,
    token: str,
    **kwargs: Any,
) -> MetrikaApiConnection:
    conn_name = "metrica api test_revision_id conn %s" % uuid.uuid4()
    conn = MetrikaApiConnection.create_from_dict(
        MetrikaApiConnection.DataModel(
            token=token,
            counter_id=counter_id,
            name=conn_name,
        ),
        ds_key=conn_name,
        type_=CONNECTION_TYPE_METRICA_API.name,
        us_manager=sync_usm,
        **kwargs,
    )
    sync_usm.save(conn)
    return conn


def make_saved_appmetrica_api_connection(
    sync_usm: SyncUSManager, counter_id: str, token: str, **kwargs: Any
) -> AppMetricaApiConnection:
    conn_name = "appmetrica api test_revision_id conn %s" % uuid.uuid4()
    conn = AppMetricaApiConnection.create_from_dict(
        AppMetricaApiConnection.DataModel(
            token=token,
            counter_id=counter_id,
            name=conn_name,
        ),
        ds_key=conn_name,
        type_=CONNECTION_TYPE_APPMETRICA_API.name,
        us_manager=sync_usm,
        **kwargs,
    )
    sync_usm.save(conn)
    return conn

lib/dl_connector_metrica/dl_connector_metrica/core/testing/connection.py
datalens-tech/datalens-backend
from dl_constants.enums import UserDataType
from dl_core.db import TypeTransformer
from dl_core.db.conversion_base import make_native_type

from dl_connector_metrica.core.constants import CONNECTION_TYPE_METRICA_API


class MetrikaApiTypeTransformer(TypeTransformer):
    conn_type = CONNECTION_TYPE_METRICA_API

    native_to_user_map = {
        make_native_type(CONNECTION_TYPE_METRICA_API, "string"): UserDataType.string,
        make_native_type(CONNECTION_TYPE_METRICA_API, "integer"): UserDataType.integer,
        make_native_type(CONNECTION_TYPE_METRICA_API, "float"): UserDataType.float,
        make_native_type(CONNECTION_TYPE_METRICA_API, "date"): UserDataType.date,
        make_native_type(CONNECTION_TYPE_METRICA_API, "datetime"): UserDataType.genericdatetime,
    }

    user_to_native_map = {
        UserDataType.string: make_native_type(CONNECTION_TYPE_METRICA_API, "string"),
        UserDataType.integer: make_native_type(CONNECTION_TYPE_METRICA_API, "integer"),
        UserDataType.float: make_native_type(CONNECTION_TYPE_METRICA_API, "float"),
        UserDataType.date: make_native_type(CONNECTION_TYPE_METRICA_API, "date"),
        UserDataType.datetime: make_native_type(CONNECTION_TYPE_METRICA_API, "datetime"),
        UserDataType.genericdatetime: make_native_type(CONNECTION_TYPE_METRICA_API, "datetime"),
        UserDataType.datetimetz: make_native_type(CONNECTION_TYPE_METRICA_API, "datetime"),
    }

lib/dl_connector_metrica/dl_connector_metrica/core/type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from datetime import date
import logging
from typing import (
    TYPE_CHECKING,
    Callable,
    ClassVar,
    Optional,
    Sequence,
)

import attr

from dl_constants.enums import (
    ConnectionType,
    UserDataType,
)
from dl_core import exc
from dl_core.db import (
    SchemaColumn,
    get_type_transformer,
)
from dl_core.us_connection_base import (
    ConnectionBase,
    DataSourceTemplate,
    ExecutorBasedMixin,
)
from dl_core.utils import secrepr
import dl_sqlalchemy_metrica_api
from dl_sqlalchemy_metrica_api import api_client as metrika_api_client
from dl_sqlalchemy_metrica_api.api_info.appmetrica import AppMetricaFieldsNamespaces
from dl_sqlalchemy_metrica_api.api_info.metrika import MetrikaApiCounterSource
from dl_utils.utils import DataKey

from dl_connector_metrica.core.constants import (
    SOURCE_TYPE_APPMETRICA_API,
    SOURCE_TYPE_METRICA_API,
)
from dl_connector_metrica.core.dto import (
    AppMetricaAPIConnDTO,
    MetricaAPIConnDTO,
)


if TYPE_CHECKING:
    from dl_core.connection_executors import SyncConnExecutorBase
    from dl_core.services_registry.top_level import ServicesRegistry


LOGGER = logging.getLogger(__name__)


def parse_metrica_ids(ids_str: str) -> Sequence[str]:
    if not ids_str:
        return []
    return [id_.strip() for id_ in ids_str.split(",")]


class MetrikaBaseMixin(ConnectionBase):
    metrica_host: Optional[str] = None

    def __init__(self, *args, **kwargs):  # type: ignore  # TODO: fix
        super().__init__(*args, **kwargs)  # type: ignore  # TODO: fix
        self._initial_counter_id = self.data.counter_id if self._data is not None else None  # type: ignore  # TODO: fix

    @property
    def allow_public_usage(self) -> bool:
        return False

    @property
    def metrika_oauth(self):  # type: ignore  # TODO: fix
        return self.data.token  # type: ignore  # TODO: fix

    @property
    def counter_id(self):  # type: ignore  # TODO: fix
        return self.data.counter_id  # type: ignore  # TODO: fix

    def get_metrica_api_cli(self) -> metrika_api_client.MetrikaApiClient:
        return metrika_api_client.MetrikaApiClient(oauth_token=self.metrika_oauth, host=self.metrica_host)

    def get_counter_creation_date(self):  # type: ignore  # TODO: fix
        assert isinstance(self.counter_id, str)
        ids = list(filter(lambda t: t, parse_metrica_ids(self.counter_id)))
        min_date = min([self.get_metrica_api_cli().get_counter_creation_date(cid) for cid in ids])
        return min_date

    @property
    def counter_creation_date(self):  # type: ignore  # TODO: fix
        return self.data.counter_creation_date  # type: ignore  # TODO: fix

    def get_available_counters(self) -> list[dict]:
        return self.get_metrica_api_cli().get_available_counters()

    async def validate_new_data(
        self,
        services_registry: ServicesRegistry,
        changes: Optional[dict] = None,
        original_version: Optional[ConnectionBase] = None,
    ) -> None:
        await super().validate_new_data(  # type: ignore  # TODO: fix  # mixin
            services_registry=services_registry,
            changes=changes,
            original_version=original_version,
        )
        if original_version is None:
            return  # only validating edits here
        assert isinstance(changes, dict)
        data_changes = changes.get("data") or {}
        if data_changes.get("token"):
            return  # token provided, nothing to check
        current_counter_id = self.data.counter_id  # type: ignore  # TODO: fix  # mixin
        if str(data_changes.get("counter_id") or "") == str(current_counter_id):
            return  # no counter_id change
        raise exc.ConnectionConfigurationError('"token" must be specified if "counter_id" is changing.')


class MetrikaApiConnection(MetrikaBaseMixin, ExecutorBasedMixin, ConnectionBase):  # type: ignore  # TODO: fix
    is_always_internal_source: ClassVar[bool] = True
    allow_cache: ClassVar[bool] = True

    metrica_host = metrika_api_client.METRIKA_API_HOST
    source_type = SOURCE_TYPE_METRICA_API

    @attr.s(kw_only=True)
    class DataModel(ConnectionBase.DataModel):
        token: str = attr.ib(repr=secrepr)
        counter_id: str = attr.ib()  # single counter id or comma-separated counters list
        counter_creation_date: Optional[date] = attr.ib(default=None)  # minimal date in case of multiple counters
        accuracy: Optional[float] = attr.ib(default=None)  # sample share (0; 1]

        @classmethod
        def get_secret_keys(cls) -> set[DataKey]:
            return {
                *super().get_secret_keys(),
                DataKey(parts=("token",)),
            }

    @property
    def metrika_oauth(self):  # type: ignore  # TODO: fix
        return self.data.token

    @property
    def table_name(self):  # type: ignore  # TODO: fix
        return self.data.counter_id

    def get_conn_dto(self) -> MetricaAPIConnDTO:
        return MetricaAPIConnDTO(
            conn_id=self.uuid,
            token=self.metrika_oauth,
            accuracy=self.data.accuracy,
        )

    @property
    def cache_ttl_sec_override(self) -> Optional[int]:
        return None

    @classmethod
    def get_api_fields_info(cls):  # type: ignore  # TODO: fix
        return dl_sqlalchemy_metrica_api.api_info.metrika

    @classmethod
    def get_raw_schema(cls, metrica_namespace: str, actual_conn_type: ConnectionType) -> Sequence[SchemaColumn]:
        fields_info = cls.get_api_fields_info().fields_by_namespace.get(
            cls.get_api_fields_info().metrica_fields_namespaces[metrica_namespace], []
        )

        def user_type_converter(type_name: str) -> UserDataType:
            return UserDataType[type_name] if type_name != "datetime" else UserDataType.genericdatetime

        raw_schema = tuple(
            SchemaColumn(
                name=field["name"],
                title=field["title"],
                user_type=user_type_converter(field["type"]),
                nullable=True,
                native_type=get_type_transformer(actual_conn_type).type_user_to_native(
                    user_t=user_type_converter(field["type"])
                ),
                description=field.get("description", ""),
                has_auto_aggregation=not field["is_dim"],
                lock_aggregation=True,
            )
            for field in fields_info
        )

        return raw_schema

    def get_parameter_combinations(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[dict]:
        return [dict(db_name=item.name) for item in MetrikaApiCounterSource]

    def get_data_source_templates(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[DataSourceTemplate]:
        return [
            DataSourceTemplate(
                title=parameters["db_name"],
                group=[],
                source_type=self.source_type,
                connection_id=self.uuid,  # type: ignore  # TODO: fix
                parameters=parameters,
            )
            for parameters in self.get_parameter_combinations(conn_executor_factory=conn_executor_factory)
        ]

    @property
    def allow_public_usage(self) -> bool:
        return False


class AppMetricaApiConnection(MetrikaApiConnection):
    allow_cache: ClassVar[bool] = True

    metrica_host = metrika_api_client.APPMETRICA_API_HOST
    source_type = SOURCE_TYPE_APPMETRICA_API

    def get_conn_dto(self) -> AppMetricaAPIConnDTO:
        return AppMetricaAPIConnDTO(
            conn_id=self.uuid,
            token=self.metrika_oauth,
            accuracy=self.data.accuracy,
        )

    @classmethod
    def get_api_fields_info(cls):  # type: ignore  # TODO: fix
        return dl_sqlalchemy_metrica_api.api_info.appmetrica

    def get_parameter_combinations(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[dict]:
        return [dict(db_name=item.name) for item in AppMetricaFieldsNamespaces]

    @property
    def allow_public_usage(self) -> bool:
        return False

lib/dl_connector_metrica/dl_connector_metrica/core/us_connection.py
datalens-tech/datalens-backend
from dl_formula.connectors.base.column import UnprefixedColumnRenderer
from dl_formula.connectors.base.connector import FormulaConnector
from dl_sqlalchemy_metrica_api.base import MetrikaApiDialect as SAMetrikaApiDialect

from dl_connector_metrica.formula.constants import MetricaDialect as MetricaDialectNS
from dl_connector_metrica.formula.definitions.all import DEFINITIONS


class MetricaFormulaConnector(FormulaConnector):
    dialect_ns_cls = MetricaDialectNS
    dialects = MetricaDialectNS.METRIKAAPI
    default_dialect = MetricaDialectNS.METRIKAAPI
    op_definitions = DEFINITIONS
    sa_dialect = SAMetrikaApiDialect()
    column_renderer_cls = UnprefixedColumnRenderer

lib/dl_connector_metrica/dl_connector_metrica/formula/connector.py
datalens-tech/datalens-backend
from dl_formula.core.dialect import (
    DialectName,
    DialectNamespace,
    simple_combo,
)


DIALECT_NAME_METRICAAPI = DialectName.declare("METRICAAPI")


class MetricaDialect(DialectNamespace):
    METRIKAAPI = simple_combo(name=DIALECT_NAME_METRICAAPI)

lib/dl_connector_metrica/dl_connector_metrica/formula/constants.py
datalens-tech/datalens-backend
from dl_connector_metrica.formula.definitions.functions_datetime import DEFINITIONS_DATETIME
from dl_connector_metrica.formula.definitions.functions_markup import DEFINITIONS_MARKUP
from dl_connector_metrica.formula.definitions.functions_string import DEFINITIONS_STRING
from dl_connector_metrica.formula.definitions.functions_type import DEFINITIONS_TYPE
from dl_connector_metrica.formula.definitions.operators_binary import DEFINITIONS_BINARY
from dl_connector_metrica.formula.definitions.operators_ternary import DEFINITIONS_TERNARY
from dl_connector_metrica.formula.definitions.operators_unary import DEFINITIONS_UNARY


DEFINITIONS = [
    *DEFINITIONS_DATETIME,
    *DEFINITIONS_MARKUP,
    *DEFINITIONS_STRING,
    *DEFINITIONS_TYPE,
    *DEFINITIONS_UNARY,
    *DEFINITIONS_BINARY,
    *DEFINITIONS_TERNARY,
]

lib/dl_connector_metrica/dl_connector_metrica/formula/definitions/all.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_datetime as base

from dl_connector_metrica.formula.constants import MetricaDialect as D


DEFINITIONS_DATETIME = [
    # dateadd
    base.FuncDateadd1.for_dialect(D.METRIKAAPI),
    base.FuncDateadd2Unit.for_dialect(D.METRIKAAPI),
    base.FuncDateadd2Number.for_dialect(D.METRIKAAPI),
]

lib/dl_connector_metrica/dl_connector_metrica/formula/definitions/functions_datetime.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_markup as base

from dl_connector_metrica.formula.constants import MetricaDialect as D


DEFINITIONS_MARKUP = [
    # __str
    base.FuncInternalStrConst.for_dialect(D.METRIKAAPI),
    base.FuncInternalStr.for_dialect(D.METRIKAAPI),
]

lib/dl_connector_metrica/dl_connector_metrica/formula/definitions/functions_markup.py
datalens-tech/datalens-backend
from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_string as base

from dl_connector_metrica.formula.constants import MetricaDialect as D


V = TranslationVariant.make


DEFINITIONS_STRING = [
    # contains
    base.FuncContainsConst(
        variants=[
            V(D.METRIKAAPI, lambda x, y: x.contains(y)),
        ]
    ),
    # notcontains
    base.FuncNotContainsConst.for_dialect(D.METRIKAAPI),
    # endswith
    base.FuncEndswithConst(
        variants=[
            V(D.METRIKAAPI, lambda x, y: x.endswith(y)),
        ]
    ),
    # startswith
    base.FuncStartswithConst(
        variants=[
            V(D.METRIKAAPI, lambda x, y: x.startswith(y)),
        ]
    ),
]

lib/dl_connector_metrica/dl_connector_metrica/formula/definitions/functions_string.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_type as base

from dl_connector_metrica.formula.constants import MetricaDialect as D


DEFINITIONS_TYPE = [
    # datetime
    base.FuncDatetime1FromDatetime.for_dialect(D.METRIKAAPI),
    # datetimetz
    base.FuncDatetimeTZConst.for_dialect(D.METRIKAAPI),
    # genericdatetime
    base.FuncGenericDatetime1FromDatetime.for_dialect(D.METRIKAAPI),
]

lib/dl_connector_metrica/dl_connector_metrica/formula/definitions/functions_type.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_binary as base

from dl_connector_metrica.formula.constants import MetricaDialect as D


DEFINITIONS_BINARY = [
    # !=
    base.BinaryNotEqual.for_dialect(D.METRIKAAPI),
    # *
    base.BinaryMultNumbers.for_dialect(D.METRIKAAPI),
    # +
    base.BinaryPlusNumbers.for_dialect(D.METRIKAAPI),
    # -
    base.BinaryMinusNumbers.for_dialect(D.METRIKAAPI),
    # /
    base.BinaryDivInt.for_dialect(D.METRIKAAPI),
    base.BinaryDivFloat.for_dialect(D.METRIKAAPI),
    # <
    base.BinaryLessThan.for_dialect(D.METRIKAAPI),
    # <=
    base.BinaryLessThanOrEqual.for_dialect(D.METRIKAAPI),
    # ==
    base.BinaryEqual.for_dialect(D.METRIKAAPI),
    # >
    base.BinaryGreaterThan.for_dialect(D.METRIKAAPI),
    # >=
    base.BinaryGreaterThanOrEqual.for_dialect(D.METRIKAAPI),
    # _!=
    base.BinaryNotEqualInternal.for_dialect(D.METRIKAAPI),
    # _==
    base.BinaryEqualInternal.for_dialect(D.METRIKAAPI),
    # _dneq
    base.BinaryEqualDenullified.for_dialect(D.METRIKAAPI),
    # and
    base.BinaryAnd.for_dialect(D.METRIKAAPI),
    # in
    base.BinaryIn.for_dialect(D.METRIKAAPI),
    # notin
    base.BinaryNotIn.for_dialect(D.METRIKAAPI),
    # or
    base.BinaryOr.for_dialect(D.METRIKAAPI),
]

lib/dl_connector_metrica/dl_connector_metrica/formula/definitions/operators_binary.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_ternary as base

from dl_connector_metrica.formula.constants import MetricaDialect as D


DEFINITIONS_TERNARY = [
    # between
    base.TernaryBetween.for_dialect(D.METRIKAAPI),
]

lib/dl_connector_metrica/dl_connector_metrica/formula/definitions/operators_ternary.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_unary as base

from dl_connector_metrica.formula.constants import MetricaDialect as D


DEFINITIONS_UNARY = [
    # not
    base.UnaryNotBool.for_dialect(D.METRIKAAPI),
]

lib/dl_connector_metrica/dl_connector_metrica/formula/definitions/operators_unary.py
datalens-tech/datalens-backend
from dl_formula_ref.texts import StyledDialect

from dl_connector_metrica.formula.constants import MetricaDialect


HUMAN_DIALECTS = {
    MetricaDialect.METRIKAAPI: StyledDialect(
        "`Yandex Metrica`",
        "`Yandex Metrica`",
        "`Yandex Metrica`",
    ),
}

lib/dl_connector_metrica/dl_connector_metrica/formula_ref/human_dialects.py
datalens-tech/datalens-backend
from dl_formula_ref.plugins.base.plugin import FormulaRefPlugin

from dl_connector_metrica.formula.constants import MetricaDialect
from dl_connector_metrica.formula_ref.human_dialects import HUMAN_DIALECTS


class MetricaFormulaRefPlugin(FormulaRefPlugin):
    any_dialects = frozenset((*MetricaDialect.METRIKAAPI.to_list(),))
    human_dialects = HUMAN_DIALECTS

lib/dl_connector_metrica/dl_connector_metrica/formula_ref/plugin.py
datalens-tech/datalens-backend
pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'

lib/dl_connector_metrica/dl_connector_metrica_tests/conftest.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_api_lib_testing.data_api_base import (
    DataApiTestParams,
    StandardizedDataApiTestBase,
)
from dl_api_lib_testing.dataset_base import DatasetTestBase
from dl_sqlalchemy_metrica_api.api_info.appmetrica import AppMetricaFieldsNamespaces
from dl_sqlalchemy_metrica_api.api_info.metrika import MetrikaApiCounterSource

from dl_connector_metrica.core.constants import (
    CONNECTION_TYPE_APPMETRICA_API,
    CONNECTION_TYPE_METRICA_API,
    SOURCE_TYPE_APPMETRICA_API,
    SOURCE_TYPE_METRICA_API,
)
from dl_connector_metrica_tests.ext.config import (
    API_TEST_CONFIG,
    APPMETRICA_SAMPLE_COUNTER_ID,
    METRIKA_SAMPLE_COUNTER_ID,
)
from dl_connector_metrica_tests.ext.core.base import (
    BaseAppMetricaTestClass,
    BaseMetricaTestClass,
)


class MetricaConnectionTestBase(BaseMetricaTestClass, ConnectionTestBase):
    conn_type = CONNECTION_TYPE_METRICA_API
    bi_compeng_pg_on = False

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def connection_params(self, metrica_token: str) -> dict:
        return dict(
            counter_id=METRIKA_SAMPLE_COUNTER_ID,
            token=metrica_token,
            accuracy=0.01,
        )


class MetricaDatasetTestBase(MetricaConnectionTestBase, DatasetTestBase):
    @pytest.fixture(scope="class")
    def dataset_params(self) -> dict:
        return dict(
            source_type=SOURCE_TYPE_METRICA_API.name,
            parameters=dict(
                db_name=MetrikaApiCounterSource.hits.name,
            ),
        )


class MetricaDataApiTestBase(MetricaDatasetTestBase, StandardizedDataApiTestBase):
    mutation_caches_on = False

    @pytest.fixture(scope="class")
    def data_api_test_params(self) -> DataApiTestParams:
        return DataApiTestParams(
            two_dims=("Домен страницы", "Просмотров в минуту"),
            summable_field="Просмотров в минуту",
            range_field="Дата и время просмотра",
            distinct_field="Адрес страницы",
            date_field="Дата просмотра",
        )


class AppMetricaConnectionTestBase(BaseAppMetricaTestClass, ConnectionTestBase):
    conn_type = CONNECTION_TYPE_APPMETRICA_API
    bi_compeng_pg_on = False

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def connection_params(self, metrica_token: str) -> dict:
        return dict(
            counter_id=APPMETRICA_SAMPLE_COUNTER_ID,
            token=metrica_token,
            accuracy=0.01,
        )


class AppMetricaDatasetTestBase(AppMetricaConnectionTestBase, DatasetTestBase):
    @pytest.fixture(scope="class")
    def dataset_params(self) -> dict:
        return dict(
            source_type=SOURCE_TYPE_APPMETRICA_API.name,
            parameters=dict(
                db_name=AppMetricaFieldsNamespaces.installs.name,
            ),
        )

lib/dl_connector_metrica/dl_connector_metrica_tests/ext/api/base.py
datalens-tech/datalens-backend
from copy import deepcopy
import json
import uuid

import pytest

from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase
from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_metrica_tests.ext.api.base import (
    AppMetricaConnectionTestBase,
    MetricaConnectionTestBase,
)
from dl_connector_metrica_tests.ext.config import METRIKA_SAMPLE_COUNTER_ID


class TestMetricaConnection(MetricaConnectionTestBase, DefaultConnectorConnectionTestSuite):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultConnectorConnectionTestSuite.test_test_connection: "Doesn't work for Metrica",
            DefaultConnectorConnectionTestSuite.test_cache_ttl_sec_override: "Unavailable for Metrica",
        }
    )

    @pytest.mark.parametrize("bad_counter", ["44147844,-44147844", "44147844,asdf"])
    def test_invalid_counter(
        self, control_api_sync_client: SyncHttpClientBase, connection_params: dict, bad_counter: str
    ) -> None:
        params = deepcopy(connection_params)
        params["counter_id"] = bad_counter
        params["name"] = f"metrica_{uuid.uuid4().hex}"
        params["type"] = self.conn_type.name
        resp = control_api_sync_client.post(
            "/api/v1/connections", data=json.dumps(params), content_type="application/json"
        )
        assert resp.status_code == 400, resp.json

    def test_update_connection(
        self,
        control_api_sync_client: SyncHttpClientBase,
        saved_connection_id: str,
        metrica_token: str,
    ) -> None:
        new_counter_id = f"{METRIKA_SAMPLE_COUNTER_ID},{METRIKA_SAMPLE_COUNTER_ID}"
        update_resp = control_api_sync_client.put(
            f"/api/v1/connections/{saved_connection_id}",
            data=json.dumps({"counter_id": new_counter_id, "token": metrica_token}),
            content_type="application/json",
        )
        assert update_resp.status_code == 200, update_resp.json

        update_resp = control_api_sync_client.put(
            f"/api/v1/connections/{saved_connection_id}",
            data=json.dumps({"token": "asdf"}),
            content_type="application/json",
        )
        assert update_resp.status_code == 200, update_resp.json


class TestAppMetricaConnection(AppMetricaConnectionTestBase, DefaultConnectorConnectionTestSuite):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultConnectorConnectionTestSuite.test_test_connection: "Doesn't work for AppMetrica",
            DefaultConnectorConnectionTestSuite.test_cache_ttl_sec_override: "Unavailable for AppMetrica",
        }
    )

lib/dl_connector_metrica/dl_connector_metrica_tests/ext/api/test_connection.py
datalens-tech/datalens-backend
import datetime

from dl_api_client.dsmaker.api.data_api import SyncHttpDataApiV2
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_testing.connector.data_api_suites import (
    DefaultConnectorDataDistinctTestSuite,
    DefaultConnectorDataGroupByFormulaTestSuite,
    DefaultConnectorDataRangeTestSuite,
    DefaultConnectorDataResultTestSuite,
)
from dl_api_lib_testing.data_api_base import DataApiTestParams
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_metrica_tests.ext.api.base import MetricaDataApiTestBase


class TestMetricaDataResult(MetricaDataApiTestBase, DefaultConnectorDataResultTestSuite):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultConnectorDataResultTestSuite.test_basic_result: "Metrica doesn't support SUM",
            DefaultConnectorDataResultTestSuite.test_dates: "Metrica doesn't support DATE",
            DefaultConnectorDataResultTestSuite.test_get_result_with_string_filter_operations_for_numbers: "Metrica doesn't support ICONTAINS",
        },
        mark_features_skipped={
            DefaultConnectorDataResultTestSuite.array_support: "Metrica doesn't support arrays",
        },
    )

    def test_metrica_result(
        self,
        saved_dataset: Dataset,
        data_api_test_params: DataApiTestParams,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        ds = saved_dataset
        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="Дата просмотра"),
                ds.find_field(title="Просмотров в минуту"),
            ],
            filters=[
                ds.find_field(title="Дата просмотра").filter("BETWEEN", values=["2019-12-01", "2019-12-07 12:00:00"])
            ],
            order_by=[ds.find_field(title="Просмотров в минуту")],
            limit=3,
        )
        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 3


class TestMetricaDataGroupBy(MetricaDataApiTestBase, DefaultConnectorDataGroupByFormulaTestSuite):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultConnectorDataGroupByFormulaTestSuite.test_complex_result: "Metrica doesn't support LEN"
        }
    )


class TestMetricaDataRange(MetricaDataApiTestBase, DefaultConnectorDataRangeTestSuite):
    def test_basic_range(
        self,
        saved_dataset: Dataset,
        data_api_test_params: DataApiTestParams,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        ds = saved_dataset

        range_resp = self.get_range(ds, data_api, field_name=data_api_test_params.range_field)
        range_rows = get_data_rows(range_resp)
        min_val, max_val = map(datetime.datetime.fromisoformat, range_rows[0])
        assert min_val <= max_val
        assert max_val - min_val < datetime.timedelta(seconds=1)


class TestMetricaDataDistinct(MetricaDataApiTestBase, DefaultConnectorDataDistinctTestSuite):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultConnectorDataDistinctTestSuite.test_date_filter_distinct: "Metrica doesn't support ICONTAINS"
        }
    )


# preview is not available for Metrica

lib/dl_connector_metrica/dl_connector_metrica_tests/ext/api/test_data.py
datalens-tech/datalens-backend
import abc
from copy import deepcopy
from typing import ClassVar

from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite
from dl_constants.enums import (
    AggregationFunction,
    FieldType,
    UserDataType,
)
from dl_sqlalchemy_metrica_api.api_info.appmetrica.installs import installs_fields
from dl_sqlalchemy_metrica_api.api_info.metrika.hits import hits_fields

from dl_connector_metrica_tests.ext.api.base import (
    AppMetricaDatasetTestBase,
    MetricaDatasetTestBase,
)


class MetricaDatasetChecker(DefaultConnectorDatasetTestSuite, metaclass=abc.ABCMeta):
    expected_fields: ClassVar[dict]

    def check_basic_dataset(self, ds: Dataset) -> None:
        assert ds.id
        assert len(ds.result_schema)
        assert all(field.aggregation == AggregationFunction.none for field in ds.result_schema)

        fields = [
            {
                "is_dim": field.type == FieldType.DIMENSION,
                "name": field.source,
                "title": field.title,
                "type": (
                    UserDataType.datetime.name
                    if field.initial_data_type == UserDataType.genericdatetime
                    else field.initial_data_type.name
                ),
            }
            for field in ds.result_schema
        ]

        expected = deepcopy(self.expected_fields)
        for elem in expected:
            elem.pop("description")
            elem.pop("src_key", None)
        assert fields == expected


class TestMetricaDataset(MetricaDatasetTestBase, MetricaDatasetChecker):
    expected_fields = hits_fields

    def test_add_field_to_dataset(
        self,
        saved_dataset: Dataset,
        control_api: SyncHttpDatasetApiV1,
    ) -> None:
        ds = saved_dataset

        ds.result_schema["new field"] = ds.field(
            avatar_id=ds.source_avatars[0].id,
            source="ym:pv:pageviewsPerMinute",
            aggregation=AggregationFunction.none,
            title="new field",
        )
        ds_resp = control_api.apply_updates(dataset=ds)
        resp_result_schema = ds_resp.json["dataset"]["result_schema"]
        assert resp_result_schema[0]["title"] == "new field"
        assert resp_result_schema[0]["type"] == FieldType.MEASURE.name

    def test_concat_validation(
        self,
        saved_dataset: Dataset,
        control_api: SyncHttpDatasetApiV1,
    ) -> None:
        ds = saved_dataset

        ds.result_schema["Test"] = ds.field(formula='CONCAT("TEST1 ", "test2")')
        ds_resp = control_api.apply_updates(dataset=ds, fail_ok=True)
        assert ds_resp.status_code == 400, ds_resp.json
        assert ds_resp.json["code"] == "ERR.DS_API.VALIDATION.ERROR"


class TestAppMetricaDataset(AppMetricaDatasetTestBase, MetricaDatasetChecker):
    expected_fields = installs_fields

lib/dl_connector_metrica/dl_connector_metrica_tests/ext/api/test_dataset.py
datalens-tech/datalens-backend
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport


METRIKA_SAMPLE_COUNTER_ID = "44147844"
APPMETRICA_SAMPLE_COUNTER_ID = "1111"

# Infra settings
CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=51911).host,
    port_us_http=get_test_container_hostport("us", fallback_port=51911).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=51910).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=51910).port,
    us_master_token="AC1ofiek8coB",
    core_connector_ep_names=["metrica_api", "appmetrica_api", "testing"],
)

API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["metrica_api", "appmetrica_api"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
)

lib/dl_connector_metrica/dl_connector_metrica_tests/ext/config.py
datalens-tech/datalens-backend
import os

import pytest

from dl_api_lib_testing.initialization import initialize_api_lib_test
from dl_testing.env_params.generic import GenericEnvParamGetter

from dl_connector_metrica_tests.ext.config import API_TEST_CONFIG


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)


@pytest.fixture(scope="session")
def env_param_getter() -> GenericEnvParamGetter:
    filepath = os.path.join(os.path.dirname(__file__), "params.yml")
    return GenericEnvParamGetter.from_yaml_file(filepath)


@pytest.fixture(scope="session")
def metrica_token(env_param_getter: GenericEnvParamGetter) -> str:
    return env_param_getter.get_str_value("METRIKA_OAUTH")

lib/dl_connector_metrica/dl_connector_metrica_tests/ext/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
from typing import (
    Generator,
    TypeVar,
)

import pytest

from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import (
    CoreDbConfig,
    Db,
)
from dl_core_testing.engine_wrapper import TestingEngineWrapper
from dl_core_testing.testcases.connection import BaseConnectionTestClass
import dl_sqlalchemy_metrica_api

from dl_connector_metrica.core.constants import (
    CONNECTION_TYPE_APPMETRICA_API,
    CONNECTION_TYPE_METRICA_API,
)
from dl_connector_metrica.core.testing.connection import (
    make_saved_appmetrica_api_connection,
    make_saved_metrika_api_connection,
)
from dl_connector_metrica.core.us_connection import (
    AppMetricaApiConnection,
    MetrikaApiConnection,
)
import dl_connector_metrica_tests.ext.config as test_config


_CONN_TV = TypeVar("_CONN_TV", MetrikaApiConnection, AppMetricaApiConnection)


class MetricaTestSetup(BaseConnectionTestClass[_CONN_TV]):
    @pytest.fixture(autouse=True)
    # FIXME: This fixture is a temporary solution for failing core tests when they are run together with api tests
    def loop(self, event_loop: asyncio.AbstractEventLoop) -> Generator[asyncio.AbstractEventLoop, None, None]:
        asyncio.set_event_loop(event_loop)
        yield event_loop
        # Attempt to cover an old version of pytest-asyncio:
        # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
        asyncio.set_event_loop_policy(None)

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return ""

    @pytest.fixture(scope="class")
    def db(self, db_config: CoreDbConfig) -> Db:
        engine_wrapper = TestingEngineWrapper(config=db_config.engine_config)
        return Db(config=db_config, engine_wrapper=engine_wrapper)

    @pytest.fixture(scope="function", autouse=True)
    def shrink_metrika_default_date_period(self, monkeypatch):
        """
        To reduce load for Metrika API and tests run time.
        """
        monkeypatch.setattr(dl_sqlalchemy_metrica_api.base, "DEFAULT_DATE_PERIOD", 3)


class BaseMetricaTestClass(MetricaTestSetup[MetrikaApiConnection]):
    conn_type = CONNECTION_TYPE_METRICA_API
    core_test_config = test_config.CORE_TEST_CONFIG

    @pytest.fixture(scope="function")
    def connection_creation_params(self, metrica_token: str) -> dict:
        return dict(
            counter_id=test_config.METRIKA_SAMPLE_COUNTER_ID,
            token=metrica_token,
        )

    @pytest.fixture(scope="function")
    def saved_connection(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params: dict,
    ) -> MetrikaApiConnection:
        return make_saved_metrika_api_connection(sync_usm=sync_us_manager, **connection_creation_params)


class BaseAppMetricaTestClass(MetricaTestSetup[AppMetricaApiConnection]):
    conn_type = CONNECTION_TYPE_APPMETRICA_API
    core_test_config = test_config.CORE_TEST_CONFIG

    @pytest.fixture(scope="function")
    def connection_creation_params(self, metrica_token: str) -> dict:
        return dict(
            counter_id=test_config.APPMETRICA_SAMPLE_COUNTER_ID,
            token=metrica_token,
        )

    @pytest.fixture(scope="function")
    def saved_connection(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params: dict,
    ) -> AppMetricaApiConnection:
        return make_saved_appmetrica_api_connection(sync_usm=sync_us_manager, **connection_creation_params)

lib/dl_connector_metrica/dl_connector_metrica_tests/ext/core/base.py
datalens-tech/datalens-backend
from dl_core.us_connection_base import DataSourceTemplate
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.testcases.connection import DefaultConnectionTestClass
from dl_sqlalchemy_metrica_api.api_info.metrika import MetrikaApiCounterSource
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_metrica.core.constants import SOURCE_TYPE_METRICA_API
from dl_connector_metrica.core.us_connection import MetrikaApiConnection
from dl_connector_metrica_tests.ext.core.base import BaseMetricaTestClass


class TestMetricaConnection(BaseMetricaTestClass, DefaultConnectionTestClass[MetrikaApiConnection]):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultConnectionTestClass.test_connection_test: "in Metrica, you can't run constant queries like SELECT 1",
        }
    )

    def check_saved_connection(self, conn: MetrikaApiConnection, params: dict) -> None:
        pass

    def test_make_connection(
        self, saved_connection: MetrikaApiConnection, conn_default_sync_us_manager: SyncUSManager
    ) -> None:
        conn = saved_connection
        usm = conn_default_sync_us_manager
        assert conn.uuid is not None
        assert conn.data.counter_creation_date

        usm_conn = usm.get_by_id(conn.uuid)
        assert conn.data.counter_creation_date == usm_conn.data.counter_creation_date
        revision_id_after_save = usm_conn.revision_id
        assert isinstance(revision_id_after_save, str)
        assert revision_id_after_save

        conn.data.name = "{} (changed)".format(conn.data.name)
        usm.save(conn)
        revision_id_after_modify = conn.revision_id
        assert isinstance(revision_id_after_modify, str)
        assert revision_id_after_modify
        assert revision_id_after_modify != revision_id_after_save

        loaded_conn = usm.get_by_id(conn.uuid)
        assert loaded_conn.revision_id == revision_id_after_modify

    def check_data_source_templates(self, conn: MetrikaApiConnection, dsrc_templates: list[DataSourceTemplate]) -> None:
        expected_templates = sorted(
            [
                DataSourceTemplate(
                    title=ns.name,
                    group=[],
                    connection_id=conn.uuid,
                    source_type=SOURCE_TYPE_METRICA_API,
                    parameters={"db_name": ns.name},
                )
                for ns in MetrikaApiCounterSource
            ],
            key=lambda el: el.title,
        )
        assert expected_templates == sorted(dsrc_templates, key=lambda el: el.title)

lib/dl_connector_metrica/dl_connector_metrica_tests/ext/core/test_connection.py
datalens-tech/datalens-backend
import datetime

import pytest

from dl_core.data_source_spec.sql import StandardSQLDataSourceSpec
from dl_core_testing.testcases.data_source import BaseDataSourceTestClass
from dl_sqlalchemy_metrica_api.api_info.metrika import MetrikaApiCounterSource

from dl_connector_metrica.core.constants import SOURCE_TYPE_METRICA_API
from dl_connector_metrica.core.data_source import MetrikaApiDataSource
from dl_connector_metrica.core.us_connection import MetrikaApiConnection
from dl_connector_metrica_tests.ext.core.base import BaseMetricaTestClass


class TestMetricaDataSource(
    BaseMetricaTestClass,
    BaseDataSourceTestClass[
        MetrikaApiConnection,
        StandardSQLDataSourceSpec,
        MetrikaApiDataSource,
    ],
):
    DSRC_CLS = MetrikaApiDataSource

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self) -> StandardSQLDataSourceSpec:
        return StandardSQLDataSourceSpec(
            source_type=SOURCE_TYPE_METRICA_API,
            db_name=MetrikaApiCounterSource.hits.name,
        )

    def test_expression_value_range(self, data_source: MetrikaApiDataSource, saved_connection: MetrikaApiConnection):
        conn = saved_connection
        dsrc = data_source

        # datetime
        min_value, max_value = dsrc.get_expression_value_range(col_name="ym:pv:dateTime")
        now = datetime.datetime.utcnow()
        assert min_value.date() == conn.data.counter_creation_date
        assert max_value - now < datetime.timedelta(seconds=1)

        # date
        min_value, max_value = dsrc.get_expression_value_range(col_name="ym:pv:startOfQuarter")
        now = datetime.datetime.utcnow()
        assert min_value == conn.data.counter_creation_date
        assert max_value == now.date()

lib/dl_connector_metrica/dl_connector_metrica_tests/ext/core/test_data_source.py
datalens-tech/datalens-backend
from typing import Optional

import pytest
import sqlalchemy as sa

from dl_constants.enums import (
    DataSourceRole,
    UserDataType,
)
from dl_core.dataset_capabilities import DatasetCapabilities
from dl_core.query.bi_query import BIQuery
from dl_core.query.expression import (
    ExpressionCtx,
    OrderByExpressionCtx,
)
from dl_core.services_registry.top_level import ServicesRegistry
from dl_core.us_dataset import Dataset
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.connection import make_saved_connection
from dl_core_testing.connector import (
    CONNECTION_TYPE_TESTING,
    SOURCE_TYPE_TESTING,
)
from dl_core_testing.dataset_wrappers import DatasetTestWrapper
from dl_core_testing.testcases.dataset import DefaultDatasetTestSuite
from dl_sqlalchemy_metrica_api.api_info.appmetrica import AppMetricaFieldsNamespaces
from dl_sqlalchemy_metrica_api.api_info.metrika import MetrikaApiCounterSource
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_metrica.core.constants import (
    SOURCE_TYPE_APPMETRICA_API,
    SOURCE_TYPE_METRICA_API,
)
from dl_connector_metrica.core.us_connection import (
    AppMetricaApiConnection,
    MetrikaApiConnection,
)
from dl_connector_metrica_tests.ext.core.base import (
    BaseAppMetricaTestClass,
    BaseMetricaTestClass,
)


class TestMetricaDataset(BaseMetricaTestClass, DefaultDatasetTestSuite[MetrikaApiConnection]):
    source_type = SOURCE_TYPE_METRICA_API

    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultDatasetTestSuite.test_simple_select_from_subquery: "Not allowed for Metrica",
        },
        mark_tests_failed={
            DefaultDatasetTestSuite.test_get_param_hash: "",  # TODO: FIXME
        },
    )

    @pytest.fixture(scope="function")
    def dsrc_params(self) -> dict:
        return dict(
            db_name=MetrikaApiCounterSource.hits.name,
        )

    def _check_simple_select(
        self,
        dataset_wrapper: DatasetTestWrapper,
        saved_dataset: Dataset,
        async_service_registry: ServicesRegistry,
        sync_us_manager: SyncUSManager,
        result_cnt: int,
        limit: Optional[int] = None,
        from_subquery: bool = False,
        subquery_limit: Optional[int] = None,
    ) -> None:
        assert limit is not None or (from_subquery and subquery_limit is not None)
        avatar_id = dataset_wrapper.get_root_avatar_strict().id
        bi_query = BIQuery(
            select_expressions=[
                ExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:pv:startOfHour", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col1",
                    user_type=UserDataType.datetime,
                ),  # is a dimension
                ExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:pv:pageviewsPerMinute", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col2",
                    user_type=UserDataType.integer,
                ),  # is a measure
            ],
            group_by_expressions=[
                ExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:pv:startOfHour", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col1",
                    user_type=UserDataType.datetime,
                ),
            ],
            limit=limit,
        )

        data = self.fetch_data(
            saved_dataset=saved_dataset,
            service_registry=async_service_registry,
            sync_us_manager=sync_us_manager,
            bi_query=bi_query,
            from_subquery=from_subquery,
            subquery_limit=subquery_limit,
        )
        assert len(list(data.data)) == result_cnt

    def test_select_data_distinct(
        self,
        dataset_wrapper: DatasetTestWrapper,
        saved_dataset: Dataset,
        conn_async_service_registry: ServicesRegistry,
        sync_us_manager: SyncUSManager,
    ) -> None:
        avatar_id = dataset_wrapper.get_root_avatar_strict().id
        bi_query = BIQuery(
            select_expressions=[
                ExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:pv:startOfHour", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col1",
                    user_type=UserDataType.datetime,
                ),  # is a dimension
            ],
            order_by_expressions=[
                OrderByExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:pv:startOfHour", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col1",
                    user_type=UserDataType.datetime,
                ),
            ],
            distinct=True,
        )

        data = self.fetch_data(
            saved_dataset=saved_dataset,
            service_registry=conn_async_service_registry,
            sync_us_manager=sync_us_manager,
            bi_query=bi_query,
        )
        values = [row[0] for row in data.data]
        assert values == sorted(set(values))

    def test_select_with_quotes(
        self,
        dataset_wrapper: DatasetTestWrapper,
        saved_dataset: Dataset,
        conn_async_service_registry: ServicesRegistry,
        sync_us_manager: SyncUSManager,
    ) -> None:
        avatar_id = dataset_wrapper.get_root_avatar_strict().id
        bi_query = BIQuery(
            select_expressions=[
                ExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:pv:startOfHour", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col1",
                    user_type=UserDataType.datetime,
                ),  # is a dimension
                ExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:pv:pageviewsPerMinute", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col2",
                    user_type=UserDataType.integer,
                ),  # is a measure
            ],
            group_by_expressions=[
                ExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:pv:startOfHour", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col1",
                    user_type=UserDataType.datetime,
                ),
            ],
            dimension_filters=[
                ExpressionCtx(
                    expression=sa.literal_column("ym:pv:openstatCampaign").in_(
                        ("Nizhny Novgorod Oblast'", "'m'a'n'y'q'u'o't'e's'")
                    ),
                    avatar_ids=[avatar_id],
                    user_type=UserDataType.boolean,
                ),
            ],
        )

        data = self.fetch_data(
            saved_dataset=saved_dataset,
            service_registry=conn_async_service_registry,
            sync_us_manager=sync_us_manager,
            bi_query=bi_query,
        )
        assert not list(data.data)  # not expecting any data, just checking that the request was successful

    def test_source_cannot_be_added(
        self,
        dataset_wrapper: DatasetTestWrapper,
        saved_dataset: Dataset,
        saved_connection: MetrikaApiConnection,
        sync_us_manager: SyncUSManager,
    ) -> None:
        testing_conn = make_saved_connection(sync_us_manager, conn_type=CONNECTION_TYPE_TESTING)
        try:
            capabilities = DatasetCapabilities(
                dataset=saved_dataset, dsrc_coll_factory=dataset_wrapper.dsrc_coll_factory
            )
            assert not capabilities.source_can_be_added(
                connection_id=saved_connection.uuid, created_from=SOURCE_TYPE_METRICA_API
            )
            assert not capabilities.source_can_be_added(
                connection_id=testing_conn.uuid, created_from=SOURCE_TYPE_TESTING
            )
        finally:
            sync_us_manager.delete(testing_conn)


class TestAppMetricaDataset(BaseAppMetricaTestClass, DefaultDatasetTestSuite[AppMetricaApiConnection]):
    source_type = SOURCE_TYPE_APPMETRICA_API

    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultDatasetTestSuite.test_simple_select_from_subquery: "Not allowed for AppMetrica",
        },
        mark_tests_failed={
            DefaultDatasetTestSuite.test_get_param_hash: "",  # TODO: FIXME
        },
    )

    @pytest.fixture(scope="function")
    def dsrc_params(self) -> dict:
        return dict(
            db_name=AppMetricaFieldsNamespaces.installs.name,
        )

    def _check_simple_select(
        self,
        dataset_wrapper: DatasetTestWrapper,
        saved_dataset: Dataset,
        async_service_registry: ServicesRegistry,
        sync_us_manager: SyncUSManager,
        result_cnt: int,
        limit: Optional[int] = None,
        from_subquery: bool = False,
        subquery_limit: Optional[int] = None,
    ) -> None:
        assert limit is not None or (from_subquery and subquery_limit is not None)
        avatar_id = dataset_wrapper.get_root_avatar_strict().id
        bi_query = BIQuery(
            select_expressions=[
                ExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:ts:regionCityName", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col1",
                    user_type=UserDataType.string,
                ),  # is a dimension
                ExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:ts:advInstallDevices", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col2",
                    user_type=UserDataType.integer,
                ),  # is a measure
            ],
            group_by_expressions=[
                ExpressionCtx(
                    expression=sa.literal_column(
                        dataset_wrapper.quote("ym:ts:regionCityName", role=DataSourceRole.origin)
                    ),
                    avatar_ids=[avatar_id],
                    alias="col1",
                    user_type=UserDataType.string,
                ),
            ],
            limit=limit,
        )

        data = self.fetch_data(
            saved_dataset=saved_dataset,
            service_registry=async_service_registry,
            sync_us_manager=sync_us_manager,
            bi_query=bi_query,
            from_subquery=from_subquery,
            subquery_limit=subquery_limit,
        )
        assert len(list(data.data)) == result_cnt

lib/dl_connector_metrica/dl_connector_metrica_tests/ext/core/test_dataset.py
datalens-tech/datalens-backend
import pytest

from dl_api_connector.i18n.localizer import CONFIGS as BI_API_CONNECTOR_CONFIGS
from dl_api_lib_testing.connection_form_base import ConnectionFormTestBase

from dl_connector_metrica.api.connection_form.form_config import (
    AppMetricaAPIConnectionFormFactory,
    MetricaAPIConnectionFormFactory,
)
from dl_connector_metrica.api.i18n.localizer import CONFIGS as BI_CONNECTOR_METRICA_CONFIGS
from dl_connector_metrica.core.settings import (
    AppmetricaConnectorSettings,
    MetricaConnectorSettings,
)


class MetricaLikeConnectionFormTestBase(ConnectionFormTestBase):
    @pytest.fixture(
        params=(True, False),
        ids=("auto_dash_True", "auto_dash_False"),
    )
    def allow_auto_dash_creation(self, request) -> bool:
        return request.param

    @pytest.fixture(
        params=(True, False),
        ids=("manual_counter_True", "manual_counter_False"),
    )
    def allow_counter_manual_input(self, request) -> bool:
        return request.param


class TestMetricaAPIConnectionForm(MetricaLikeConnectionFormTestBase):
    CONN_FORM_FACTORY_CLS = MetricaAPIConnectionFormFactory
    TRANSLATION_CONFIGS = BI_API_CONNECTOR_CONFIGS + BI_CONNECTOR_METRICA_CONFIGS

    @pytest.fixture
    def connectors_settings(  # noqa
        self,
        allow_auto_dash_creation,
        allow_counter_manual_input,
    ) -> MetricaConnectorSettings:
        return MetricaConnectorSettings(
            COUNTER_ALLOW_MANUAL_INPUT=allow_counter_manual_input,
            ALLOW_AUTO_DASH_CREATION=allow_auto_dash_creation,
        )


class TestAppMetricaAPIConnectionForm(MetricaLikeConnectionFormTestBase):
    CONN_FORM_FACTORY_CLS = AppMetricaAPIConnectionFormFactory
    TRANSLATION_CONFIGS = BI_API_CONNECTOR_CONFIGS + BI_CONNECTOR_METRICA_CONFIGS

    @pytest.fixture
    def connectors_settings(  # noqa
        self,
        allow_auto_dash_creation,
        allow_counter_manual_input,
    ) -> AppmetricaConnectorSettings:
        return AppmetricaConnectorSettings(
            COUNTER_ALLOW_MANUAL_INPUT=allow_counter_manual_input,
            ALLOW_AUTO_DASH_CREATION=allow_auto_dash_creation,
        )

lib/dl_connector_metrica/dl_connector_metrica_tests/unit/test_connection_form.py
datalens-tech/datalens-backend
import sys

import oracledb


oracledb.version = "8.3.0"
sys.modules["cx_Oracle"] = oracledb

lib/dl_connector_oracle/dl_connector_oracle/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.connection_base import ConnectionMetaMixin
from dl_api_connector.api_schema.connection_mixins import (
    DataExportForbiddenMixin,
    RawSQLLevelMixin,
)
from dl_api_connector.api_schema.connection_sql import ClassicSQLConnectionSchema
from dl_api_connector.api_schema.extras import FieldExtra

from dl_connector_oracle.core.constants import OracleDbNameType
from dl_connector_oracle.core.us_connection import ConnectionSQLOracle


class OracleConnectionSchema(
    ConnectionMetaMixin, RawSQLLevelMixin, DataExportForbiddenMixin, ClassicSQLConnectionSchema
):
    TARGET_CLS = ConnectionSQLOracle
    ALLOW_MULTI_HOST = True

    db_connect_method = ma_fields.Enum(
        OracleDbNameType,
        attribute="data.db_name_type",
        required=True,
        bi_extra=FieldExtra(editable=True),
    )

lib/dl_connector_oracle/dl_connector_oracle/api/api_schema/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import unique
from typing import Optional

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormActionApiSchema,
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
)
from dl_api_connector.form_config.models.common import (
    CommonFieldName,
    FormFieldName,
)
import dl_api_connector.form_config.models.rows as C
from dl_api_connector.form_config.models.shortcuts.rows import RowConstructor
from dl_configs.connectors_settings import ConnectorSettingsBase

from dl_connector_oracle.api.connection_info import OracleConnectionInfoProvider
from dl_connector_oracle.api.i18n.localizer import Translatable
from dl_connector_oracle.core.constants import OracleDbNameType


@unique
class OracleFieldName(FormFieldName):
    db_connect_method = "db_connect_method"


class OracleConnectionFormFactory(ConnectionFormFactory):
    def get_form_config(
        self,
        connector_settings: Optional[ConnectorSettingsBase],
        tenant: Optional[TenantDef],
    ) -> ConnectionForm:
        rc = RowConstructor(self._localizer)

        common_api_schema_items: list[FormFieldApiSchema] = [
            FormFieldApiSchema(name=CommonFieldName.host, required=True),
            FormFieldApiSchema(name=CommonFieldName.port, required=True),
            FormFieldApiSchema(name=OracleFieldName.db_connect_method, required=True),
            FormFieldApiSchema(name=CommonFieldName.db_name, required=True),
            FormFieldApiSchema(name=CommonFieldName.username, required=True),
            FormFieldApiSchema(name=CommonFieldName.password, required=self.mode == ConnectionFormMode.create),
        ]

        edit_api_schema = FormActionApiSchema(
            items=[
                *common_api_schema_items,
                FormFieldApiSchema(name=CommonFieldName.cache_ttl_sec, nullable=True),
                FormFieldApiSchema(name=CommonFieldName.raw_sql_level),
                FormFieldApiSchema(name=CommonFieldName.data_export_forbidden),
            ]
        )

        create_api_schema = FormActionApiSchema(
            items=[
                *edit_api_schema.items,
                *self._get_top_level_create_api_schema_items(),
            ]
        )

        check_api_schema = FormActionApiSchema(
            items=[
                *common_api_schema_items,
                *self._get_top_level_check_api_schema_items(),
            ]
        )

        db_name_row = C.CustomizableRow(
            items=[
                *rc.db_name_row().items,
                C.RadioButtonRowItem(
                    name=OracleFieldName.db_connect_method,
                    options=[
                        C.SelectableOption(
                            text=self._localizer.translate(Translatable("value_db-connect-method-service-name")),
                            value=OracleDbNameType.service_name.value,
                        ),
                        C.SelectableOption(
                            text=self._localizer.translate(Translatable("value_db-connect-method-sid")),
                            value=OracleDbNameType.sid.value,
                        ),
                    ],
                    default_value=OracleDbNameType.service_name.value,
                ),
            ]
        )

        return ConnectionForm(
            title=OracleConnectionInfoProvider.get_title(self._localizer),
            rows=[
                rc.host_row(),
                rc.port_row(default_value="1521"),
                db_name_row,
                rc.username_row(),
                rc.password_row(self.mode),
                C.CacheTTLRow(name=CommonFieldName.cache_ttl_sec),
                rc.raw_sql_level_row(),
                rc.collapse_advanced_settings_row(),
                rc.data_export_forbidden_row(),
            ],
            api_schema=FormApiSchema(
                create=create_api_schema if self.mode == ConnectionFormMode.create else None,
                edit=edit_api_schema if self.mode == ConnectionFormMode.edit else None,
                check=check_api_schema,
            ),
        )

lib/dl_connector_oracle/dl_connector_oracle/api/connection_form/form_config.py
datalens-tech/datalens-backend
from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_oracle.api.i18n.localizer import Translatable


class OracleConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-oracle")

lib/dl_connector_oracle/dl_connector_oracle/api/connection_info.py
datalens-tech/datalens-backend
from dl_api_connector.api_schema.source_base import (
    SchematizedSQLDataSourceSchema,
    SchematizedSQLDataSourceTemplateSchema,
    SubselectDataSourceSchema,
    SubselectDataSourceTemplateSchema,
)
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)

from dl_connector_oracle.api.api_schema.connection import OracleConnectionSchema
from dl_connector_oracle.api.connection_form.form_config import OracleConnectionFormFactory
from dl_connector_oracle.api.connection_info import OracleConnectionInfoProvider
from dl_connector_oracle.api.dashsql import OracleDashSQLParamLiteralizer
from dl_connector_oracle.api.i18n.localizer import CONFIGS
from dl_connector_oracle.core.connector import (
    OracleCoreConnectionDefinition,
    OracleCoreConnector,
    OracleSubselectCoreSourceDefinition,
    OracleTableCoreSourceDefinition,
)
from dl_connector_oracle.formula.constants import DIALECT_NAME_ORACLE


class OracleApiTableSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = OracleTableCoreSourceDefinition
    api_schema_cls = SchematizedSQLDataSourceSchema
    template_api_schema_cls = SchematizedSQLDataSourceTemplateSchema


class OracleApiSubselectSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = OracleSubselectCoreSourceDefinition
    api_schema_cls = SubselectDataSourceSchema
    template_api_schema_cls = SubselectDataSourceTemplateSchema


class OracleApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = OracleCoreConnectionDefinition
    api_generic_schema_cls = OracleConnectionSchema
    info_provider_cls = OracleConnectionInfoProvider
    form_factory_cls = OracleConnectionFormFactory


class OracleApiConnector(ApiConnector):
    core_connector_cls = OracleCoreConnector
    connection_definitions = (OracleApiConnectionDefinition,)
    source_definitions = (
        OracleApiTableSourceDefinition,
        OracleApiSubselectSourceDefinition,
    )
    formula_dialect_name = DIALECT_NAME_ORACLE
    translation_configs = frozenset(CONFIGS)
    dashsql_literalizer_cls = OracleDashSQLParamLiteralizer

lib/dl_connector_oracle/dl_connector_oracle/api/connector.py
datalens-tech/datalens-backend
import sqlalchemy as sa
from sqlalchemy.types import TypeEngine

from dl_api_connector.dashsql import (
    DefaultDashSQLParamLiteralizer,
    TValueBase,
)
from dl_constants.enums import UserDataType


class OracleDashSQLParamLiteralizer(DefaultDashSQLParamLiteralizer):
    def get_sa_type(self, bi_type: UserDataType, value_base: TValueBase) -> TypeEngine:
        if bi_type == UserDataType.string:
            # See also: dl_formula/definitions/literals.py
            value_lst = [value_base] if isinstance(value_base, str) else value_base
            max_len = max(len(val) for val in value_lst)
            try:
                for val in value_lst:
                    val.encode("ascii")
            except UnicodeEncodeError:
                return sa.NCHAR(max_len)
            return sa.CHAR(max_len)

        return super().get_sa_type(bi_type=bi_type, value_base=value_base)

lib/dl_connector_oracle/dl_connector_oracle/api/dashsql.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_oracle as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_oracle/dl_connector_oracle/api/i18n/localizer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
    Optional,
    Tuple,
    Type,
)

import oracledb  # type: ignore  # TODO: fix
import sqlalchemy as sa
import sqlalchemy.dialects.oracle.base as sa_ora  # not all data types are imported in init in older SA versions
from sqlalchemy.sql.type_api import TypeEngine

from dl_core.connection_executors.adapters.adapters_base_sa_classic import (
    BaseClassicAdapter,
    ClassicSQLConnLineConstructor,
)
from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
from dl_core.connection_models import (
    DBIdent,
    SchemaIdent,
    TableIdent,
)
from dl_core.db.native_type import SATypeSpec

from dl_connector_oracle.core.constants import CONNECTION_TYPE_ORACLE
from dl_connector_oracle.core.target_dto import OracleConnTargetDTO


class OracleConnLineConstructor(ClassicSQLConnLineConstructor[OracleConnTargetDTO]):
    def _get_dsn_params(
        self,
        safe_db_symbols: Tuple[str, ...] = (),
        db_name: Optional[str] = None,
        standard_auth: Optional[bool] = True,
    ) -> dict:
        return dict(
            super()._get_dsn_params(
                safe_db_symbols=safe_db_symbols,
                db_name=db_name,
                standard_auth=standard_auth,
            ),
            db_name_type=self._target_dto.db_name_type.value.upper(),
        )


class OracleDefaultAdapter(BaseClassicAdapter[OracleConnTargetDTO]):
    conn_type = CONNECTION_TYPE_ORACLE
    conn_line_constructor_type: ClassVar[Type[OracleConnLineConstructor]] = OracleConnLineConstructor

    dsn_template = (
        "{dialect}://{user}:{passwd}@(DESCRIPTION="
        "(ADDRESS=(PROTOCOL=TCP)(HOST={host})(PORT={port}))"
        "(CONNECT_DATA=({db_name_type}={db_name})))"
    )

    def _test(self) -> None:
        self.execute(DBAdapterQuery("SELECT 1 FROM DUAL")).get_all()

    def _get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        return self.execute(DBAdapterQuery("SELECT * FROM V$VERSION", db_name=db_ident.db_name)).get_all()[0][0]

    def normalize_sa_col_type(self, sa_col_type: TypeEngine) -> TypeEngine:
        if isinstance(sa_col_type, sa_ora.NUMBER) and not sa_col_type.scale:
            return sa.Integer()

        return super().normalize_sa_col_type(sa_col_type)

    def _is_table_exists(self, table_ident: TableIdent) -> bool:
        sa_exists_result = super()._is_table_exists(table_ident)

        # Dialect does not check for views so we will check it manually
        if not sa_exists_result:
            assert "'" not in table_ident.table_name  # FIXME: quote the name properly

            rows = self.execute(
                DBAdapterQuery(
                    f"SELECT view_name FROM all_views WHERE UPPER(view_name) = '{table_ident.table_name.upper()}'"
                )
            ).get_all()

            return len(rows) > 0

        return sa_exists_result

    _type_code_to_sa = {
        oracledb.NUMBER: sa_ora.NUMBER,
        oracledb.STRING: sa_ora.VARCHAR,  # e.g. 'VARCHAR2(44)'
        oracledb.NATIVE_FLOAT: sa_ora.BINARY_FLOAT,  # ... or `sa_ora.BINARY_FLOAT`
        oracledb.FIXED_CHAR: sa_ora.CHAR,  # e.g. 'CHAR(1)'
        oracledb.FIXED_NCHAR: sa_ora.NCHAR,  # e.g. 'NCHAR(1)'
        oracledb.NCHAR: sa_ora.NVARCHAR,  # e.g. 'NVARCHAR(5)'
        oracledb.DATETIME: sa_ora.DATE,
        oracledb.TIMESTAMP: sa_ora.TIMESTAMP,  # e.g. 'TIMESTAMP(9)', 'TIMESTAMP(9) WITH TIME ZONE'
        # # For newer versions of oracledb, just in case.
        # # Listed are all types from 8.1.0,
        # # Enabled are types listed in `dl_core.db.conversions`.
        # getattr(oracledb, 'DB_TYPE_BFILE', None): sa_ora.NULL,
        getattr(oracledb, "DB_TYPE_BINARY_DOUBLE", None): sa_ora.BINARY_DOUBLE,
        getattr(oracledb, "DB_TYPE_BINARY_FLOAT", None): sa_ora.BINARY_FLOAT,
        # getattr(oracledb, 'DB_TYPE_BINARY_INTEGER', None): sa_ora.NULL,
        # getattr(oracledb, 'DB_TYPE_BLOB', None): sa_ora.NULL,
        # getattr(oracledb, 'DB_TYPE_BOOLEAN', None): sa_ora.NULL,
        getattr(oracledb, "DB_TYPE_CHAR", None): sa_ora.CHAR,
        # getattr(oracledb, 'DB_TYPE_CLOB', None): sa_ora.NULL,
        # getattr(oracledb, 'DB_TYPE_CURSOR', None): sa_ora.NULL,
        getattr(oracledb, "DB_TYPE_DATE", None): sa_ora.DATE,
        # getattr(oracledb, 'DB_TYPE_INTERVAL_DS', None): sa_ora.NULL,
        # getattr(oracledb, 'DB_TYPE_INTERVAL_YM', None): sa_ora.NULL,
        # getattr(oracledb, 'DB_TYPE_JSON', None): sa_ora.NULL,
        # getattr(oracledb, 'DB_TYPE_LONG', None): sa_ora.NULL,
        # getattr(oracledb, 'DB_TYPE_LONG_RAW', None): sa_ora.NULL,
        getattr(oracledb, "DB_TYPE_NCHAR", None): sa_ora.NCHAR,
        # getattr(oracledb, 'DB_TYPE_NCLOB', None): sa_ora.NULL,
        getattr(oracledb, "DB_TYPE_NUMBER", None): sa_ora.NUMBER,
        getattr(oracledb, "DB_TYPE_NVARCHAR", None): sa_ora.NVARCHAR,
        # getattr(oracledb, 'DB_TYPE_OBJECT', None): sa_ora.NULL,
        # getattr(oracledb, 'DB_TYPE_RAW', None): sa_ora.NULL,
        # getattr(oracledb, 'DB_TYPE_ROWID', None): sa_ora.NULL,
        getattr(oracledb, "DB_TYPE_TIMESTAMP", None): sa_ora.TIMESTAMP,
        # getattr(oracledb, 'DB_TYPE_TIMESTAMP_LTZ', None): sa_ora.NULL,
        # # NOTE: not `timestamptz` here; matches the view schema, somehow.
        getattr(oracledb, "DB_TYPE_TIMESTAMP_TZ", None): sa_ora.TIMESTAMP,
        getattr(oracledb, "DB_TYPE_VARCHAR", None): sa_ora.VARCHAR,
    }

    def _cursor_column_to_name(self, cursor_col, dialect=None) -> str:  # type: ignore  # TODO: fix
        assert dialect, "required in this case"
        # To match the `get_columns` result.
        # Generally just lowercases the name.
        # Notably, column names seem to be case-insensitive in oracle.
        return dialect.normalize_name(cursor_col[0])

    def _cursor_column_to_sa(self, cursor_col, require: bool = True) -> Optional[SATypeSpec]:  # type: ignore  # TODO: fix
        """
        cursor_col:

            name, type, display_size, internal_size, precision, scale, null_ok

        See also:
        https://cx-oracle.readthedocs.io/en/latest/api_manual/cursor.html#Cursor.description
        """
        type_code = cursor_col[1]
        if type_code is None:  # shouldn't really happen
            if require:
                raise ValueError(f"Cursor column has no type: {cursor_col!r}")
            return None
        sa_cls = self._type_code_to_sa.get(type_code)
        if sa_cls is None:
            if require:
                raise ValueError(f"Unknown cursor type: {type_code!r}")
            return None

        # It would be nice to distinguish timestamp/timestamptz here, but the
        # only observed difference is `display_size=23` for tz-*naive*
        # timestamp.

        if sa_cls is sa_ora.NUMBER:
            # See also: `self.normalize_sa_col_type`
            precision = cursor_col[4]
            scale = cursor_col[5]

            # Going by the comparison with the 'create view' -> SA logic.
            if scale == -127:
                scale = 0

            sa_type = sa_cls(precision, scale)
        else:
            sa_type = sa_cls

        return sa_type

    def _cursor_column_to_nullable(self, cursor_col) -> Optional[bool]:  # type: ignore  # TODO: fix
        return bool(cursor_col[6])

    def _make_cursor_info(self, cursor, db_session=None) -> dict:  # type: ignore  # TODO: fix
        return dict(
            super()._make_cursor_info(cursor, db_session=db_session),
            cxoracle_types=[self._cursor_type_to_str(column[1]) for column in cursor.description],
        )

    ORACLE_LIST_SOURCES_ALL_SCHEMA_SQL = """
        SELECT OWNER, TABLE_NAME FROM ALL_TABLES
        WHERE nvl(tablespace_name, 'no tablespace')
            NOT IN ('SYSTEM', 'SYSAUX')
            AND IOT_NAME IS NULL
            AND DURATION IS NULL
    """

    def _get_tables(self, schema_ident: SchemaIdent) -> list[TableIdent]:
        if schema_ident.schema_name is not None:
            return super()._get_tables(schema_ident)

        db_name = schema_ident.db_name
        db_engine = self.get_db_engine(db_name)

        query = self.ORACLE_LIST_SOURCES_ALL_SCHEMA_SQL
        result = db_engine.execute(sa.text(query))
        return [
            TableIdent(
                db_name=db_name,
                schema_name=owner,
                table_name=table_name,
            )
            for owner, table_name in result
        ]

    @staticmethod
    def _cursor_type_to_str(value: Any) -> str:
        return value.name.lower()

lib/dl_connector_oracle/dl_connector_oracle/core/adapters_oracle.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import List

import attr

from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor

from dl_connector_oracle.core.adapters_oracle import OracleDefaultAdapter
from dl_connector_oracle.core.dto import OracleConnDTO
from dl_connector_oracle.core.target_dto import OracleConnTargetDTO


@attr.s(cmp=False, hash=False)
class OracleDefaultConnExecutor(DefaultSqlAlchemyConnExecutor[OracleDefaultAdapter]):
    TARGET_ADAPTER_CLS = OracleDefaultAdapter

    _conn_dto: OracleConnDTO = attr.ib()

    async def _make_target_conn_dto_pool(self) -> List[OracleConnTargetDTO]:  # type: ignore  # TODO: fix
        dto_pool = []
        for host in self._conn_hosts_pool:
            dto_pool.append(
                OracleConnTargetDTO(
                    conn_id=self._conn_dto.conn_id,
                    pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                    pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                    host=host,
                    port=self._conn_dto.port,
                    db_name=self._conn_dto.db_name,
                    db_name_type=self._conn_dto.db_name_type,
                    username=self._conn_dto.username,
                    password=self._conn_dto.password,
                )
            )
        return dto_pool

lib/dl_connector_oracle/dl_connector_oracle/core/connection_executors.py
datalens-tech/datalens-backend
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
)
from dl_core.connectors.sql_base.connector import (
    SQLSubselectCoreSourceDefinitionBase,
    SQLTableCoreSourceDefinitionBase,
)

from dl_connector_oracle.core.adapters_oracle import OracleDefaultAdapter
from dl_connector_oracle.core.connection_executors import OracleDefaultConnExecutor
from dl_connector_oracle.core.constants import (
    BACKEND_TYPE_ORACLE,
    CONNECTION_TYPE_ORACLE,
    SOURCE_TYPE_ORACLE_SUBSELECT,
    SOURCE_TYPE_ORACLE_TABLE,
)
from dl_connector_oracle.core.data_source import (
    OracleDataSource,
    OracleSubselectDataSource,
)
from dl_connector_oracle.core.data_source_migration import OracleDataSourceMigrator
from dl_connector_oracle.core.sa_types import SQLALCHEMY_ORACLE_TYPES
from dl_connector_oracle.core.storage_schemas.connection import ConnectionSQLOracleDataStorageSchema
from dl_connector_oracle.core.type_transformer import OracleServerTypeTransformer
from dl_connector_oracle.core.us_connection import ConnectionSQLOracle


class OracleCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_ORACLE
    connection_cls = ConnectionSQLOracle
    us_storage_schema_cls = ConnectionSQLOracleDataStorageSchema
    type_transformer_cls = OracleServerTypeTransformer
    sync_conn_executor_cls = OracleDefaultConnExecutor
    async_conn_executor_cls = OracleDefaultConnExecutor
    dialect_string = "bi_oracle"
    data_source_migrator_cls = OracleDataSourceMigrator


class OracleTableCoreSourceDefinition(SQLTableCoreSourceDefinitionBase):
    source_type = SOURCE_TYPE_ORACLE_TABLE
    source_cls = OracleDataSource


class OracleSubselectCoreSourceDefinition(SQLSubselectCoreSourceDefinitionBase):
    source_type = SOURCE_TYPE_ORACLE_SUBSELECT
    source_cls = OracleSubselectDataSource


class OracleCoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_ORACLE
    connection_definitions = (OracleCoreConnectionDefinition,)
    source_definitions = (
        OracleTableCoreSourceDefinition,
        OracleSubselectCoreSourceDefinition,
    )
    rqe_adapter_classes = frozenset({OracleDefaultAdapter})
    sa_types = SQLALCHEMY_ORACLE_TYPES

lib/dl_connector_oracle/dl_connector_oracle/core/connector.py
datalens-tech/datalens-backend
from enum import (
    Enum,
    unique,
)

from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
    SourceBackendType,
)


BACKEND_TYPE_ORACLE = SourceBackendType.declare("ORACLE")
CONNECTION_TYPE_ORACLE = ConnectionType.declare("oracle")
SOURCE_TYPE_ORACLE_TABLE = DataSourceType.declare("ORACLE_TABLE")
SOURCE_TYPE_ORACLE_SUBSELECT = DataSourceType.declare("ORACLE_SUBSELECT")


@unique
class OracleDbNameType(Enum):
    sid = "sid"
    service_name = "service_name"

lib/dl_connector_oracle/dl_connector_oracle/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Any,
    Optional,
)

from dl_constants.enums import DataSourceType
from dl_core.data_source.sql import (
    BaseSQLDataSource,
    StandardSchemaSQLDataSource,
    SubselectDataSource,
    require_table_name,
)
from dl_core.utils import sa_plain_text

from dl_connector_oracle.core.constants import (
    CONNECTION_TYPE_ORACLE,
    SOURCE_TYPE_ORACLE_SUBSELECT,
    SOURCE_TYPE_ORACLE_TABLE,
)
from dl_connector_oracle.core.query_compiler import OracleQueryCompiler


LOGGER = logging.getLogger(__name__)


class OracleDataSourceMixin(BaseSQLDataSource):
    compiler_cls = OracleQueryCompiler

    conn_type = CONNECTION_TYPE_ORACLE

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type in (SOURCE_TYPE_ORACLE_TABLE, SOURCE_TYPE_ORACLE_SUBSELECT)


class OracleDataSource(OracleDataSourceMixin, StandardSchemaSQLDataSource):
    """Oracle table"""

    @require_table_name
    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        q = self.quote
        alias_str = "" if alias is None else f" {q(alias)}"
        schema_str = "" if self.schema_name is None else f"{q(self.schema_name)}."
        return sa_plain_text(f"{schema_str}{q(self.table_name)}{alias_str}")


class OracleSubselectDataSource(OracleDataSourceMixin, SubselectDataSource):
    """Oracle subselect"""

    # In oracle, `(select …) as source` doesn't work, only `(select …) source`.
    _subquery_alias_joiner = " "

lib/dl_connector_oracle/dl_connector_oracle/core/data_source.py
datalens-tech/datalens-backend
from dl_core.connectors.sql_base.data_source_migration import DefaultSQLDataSourceMigrator

from dl_connector_oracle.core.constants import (
    SOURCE_TYPE_ORACLE_SUBSELECT,
    SOURCE_TYPE_ORACLE_TABLE,
)


class OracleDataSourceMigrator(DefaultSQLDataSourceMigrator):
    table_source_type = SOURCE_TYPE_ORACLE_TABLE
    subselect_source_type = SOURCE_TYPE_ORACLE_SUBSELECT

lib/dl_connector_oracle/dl_connector_oracle/core/data_source_migration.py
datalens-tech/datalens-backend
import attr

from dl_core.connection_models.dto_defs import DefaultSQLDTO

from dl_connector_oracle.core.constants import (
    CONNECTION_TYPE_ORACLE,
    OracleDbNameType,
)


@attr.s(frozen=True)
class OracleConnDTO(DefaultSQLDTO):  # noqa
    conn_type = CONNECTION_TYPE_ORACLE

    db_name_type: OracleDbNameType = attr.ib(kw_only=True)

lib/dl_connector_oracle/dl_connector_oracle/core/dto.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_core.connectors.base.query_compiler import QueryCompiler


class OracleQueryCompiler(QueryCompiler):
    force_nulls_lower_than_values = True

lib/dl_connector_oracle/dl_connector_oracle/core/query_compiler.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa
from sqlalchemy.dialects.oracle import base as or_types  # not all data types are imported in init in older SA versions

from dl_core.db.sa_types_base import (
    lengthed_instantiator,
    make_native_type,
    simple_instantiator,
)

from dl_connector_oracle.core.constants import CONNECTION_TYPE_ORACLE


SQLALCHEMY_ORACLE_BASE_TYPES = (
    or_types.NUMBER,
    or_types.BINARY_FLOAT,
    or_types.BINARY_DOUBLE,
    or_types.DATE,
    or_types.TIMESTAMP,
)
SQLALCHEMY_ORACLE_LENGTHED_TYPES = (
    or_types.CHAR,
    sa.NCHAR,
    or_types.VARCHAR,
    or_types.NVARCHAR,
    or_types.VARCHAR2,
    or_types.NVARCHAR2,
)
SQLALCHEMY_ORACLE_TYPES = {
    **{
        make_native_type(CONNECTION_TYPE_ORACLE, typecls): simple_instantiator(typecls)
        for typecls in SQLALCHEMY_ORACLE_BASE_TYPES
    },
    # A tricky substitution (possibly in the wrong place):
    make_native_type(CONNECTION_TYPE_ORACLE, sa.Integer): simple_instantiator(or_types.NUMBER),
    **{
        make_native_type(CONNECTION_TYPE_ORACLE, typecls): lengthed_instantiator(typecls)
        for typecls in SQLALCHEMY_ORACLE_LENGTHED_TYPES
    },
}

lib/dl_connector_oracle/dl_connector_oracle/core/sa_types.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_core.us_manager.storage_schemas.connection import ConnectionSQLDataStorageSchema

from dl_connector_oracle.core.constants import OracleDbNameType
from dl_connector_oracle.core.us_connection import ConnectionSQLOracle


class ConnectionSQLOracleDataStorageSchema(
    ConnectionSQLDataStorageSchema[ConnectionSQLOracle.DataModel],
):
    TARGET_CLS = ConnectionSQLOracle.DataModel

    db_name_type = ma_fields.Enum(OracleDbNameType, required=True, allow_none=False)

lib/dl_connector_oracle/dl_connector_oracle/core/storage_schemas/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Type,
    TypeVar,
)

import attr

from dl_core.connection_executors.models.connection_target_dto_base import BaseSQLConnTargetDTO

from dl_connector_oracle.core.constants import OracleDbNameType


if TYPE_CHECKING:
    from dl_constants.types import TJSONLike


_CT_DTO_TV = TypeVar("_CT_DTO_TV", bound="OracleConnTargetDTO")


@attr.s(frozen=True)
class OracleConnTargetDTO(BaseSQLConnTargetDTO):
    db_name_type: OracleDbNameType = attr.ib()

    def to_jsonable_dict(self) -> dict[str, TJSONLike]:
        d = super().to_jsonable_dict()
        return {
            **d,
            "db_name_type": self.db_name_type.name,
        }

    @classmethod
    def _from_jsonable_dict(cls: Type[_CT_DTO_TV], data: dict) -> _CT_DTO_TV:
        prepared_data = {**data, "db_name_type": OracleDbNameType[data["db_name_type"]]}
        return cls(**prepared_data)  # type: ignore

lib/dl_connector_oracle/dl_connector_oracle/core/target_dto.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Optional,
)
import uuid

from dl_constants.enums import RawSQLLevel
from dl_core.us_manager.us_manager_sync import SyncUSManager

from dl_connector_oracle.core.constants import CONNECTION_TYPE_ORACLE
from dl_connector_oracle.core.us_connection import ConnectionSQLOracle


def make_oracle_saved_connection(
    sync_usm: SyncUSManager,
    db_name: Optional[str],
    host: str,
    port: Optional[int],
    username: Optional[str],
    password: Optional[str],
    raw_sql_level: RawSQLLevel = RawSQLLevel.off,
    **kwargs: Any,
) -> ConnectionSQLOracle:
    conn_name = "oracle test conn {}".format(uuid.uuid4())
    conn = ConnectionSQLOracle.create_from_dict(
        data_dict=ConnectionSQLOracle.DataModel(
            db_name=db_name,
            host=host,
            port=port,
            username=username,
            password=password,
            raw_sql_level=raw_sql_level,
        ),
        ds_key=conn_name,
        type_=CONNECTION_TYPE_ORACLE.name,
        us_manager=sync_usm,
        **kwargs,
    )
    sync_usm.save(conn)
    return conn

lib/dl_connector_oracle/dl_connector_oracle/core/testing/connection.py
datalens-tech/datalens-backend
import sqlalchemy as sa
from sqlalchemy.dialects.oracle import base as or_types  # not all data types are imported in init in older SA versions

from dl_constants.enums import UserDataType
from dl_core.db.conversion_base import (
    TypeTransformer,
    make_native_type,
)

from dl_connector_oracle.core.constants import CONNECTION_TYPE_ORACLE


class OracleServerTypeTransformer(TypeTransformer):
    conn_type = CONNECTION_TYPE_ORACLE
    native_to_user_map = {
        # No separate type for integer. Number acts as DECIMAL with customizable precision.
        **{
            make_native_type(CONNECTION_TYPE_ORACLE, t): UserDataType.integer
            for t in (sa.Integer,)  # pseudo type used if scale == 0
        },
        **{
            make_native_type(CONNECTION_TYPE_ORACLE, t): UserDataType.float  # type: ignore  # TODO: fix
            for t in (or_types.NUMBER, or_types.BINARY_FLOAT, or_types.BINARY_DOUBLE)
        },
        **{
            make_native_type(CONNECTION_TYPE_ORACLE, t): UserDataType.string
            for t in (
                or_types.CHAR,
                or_types.VARCHAR,
                or_types.VARCHAR2,
                sa.NCHAR,
                or_types.NVARCHAR,
                or_types.NVARCHAR2,
            )
        },
        **{
            make_native_type(CONNECTION_TYPE_ORACLE, t): UserDataType.genericdatetime  # type: ignore  # TODO: fix
            for t in (or_types.DATE, or_types.TIMESTAMP)
        },
        # No separate type for date, it's the same as for datetime
        make_native_type(CONNECTION_TYPE_ORACLE, sa.sql.sqltypes.NullType): UserDataType.unsupported,
    }
    user_to_native_map = {
        UserDataType.integer: make_native_type(CONNECTION_TYPE_ORACLE, sa.INTEGER),
        UserDataType.float: make_native_type(CONNECTION_TYPE_ORACLE, or_types.BINARY_DOUBLE),
        UserDataType.boolean: make_native_type(CONNECTION_TYPE_ORACLE, or_types.NUMBER),
        UserDataType.string: make_native_type(CONNECTION_TYPE_ORACLE, or_types.NVARCHAR2),
        UserDataType.date: make_native_type(CONNECTION_TYPE_ORACLE, or_types.DATE),
        UserDataType.datetime: make_native_type(CONNECTION_TYPE_ORACLE, or_types.DATE),
        UserDataType.genericdatetime: make_native_type(CONNECTION_TYPE_ORACLE, or_types.DATE),
        UserDataType.geopoint: make_native_type(CONNECTION_TYPE_ORACLE, or_types.NVARCHAR2),
        UserDataType.geopolygon: make_native_type(CONNECTION_TYPE_ORACLE, or_types.NVARCHAR2),
        UserDataType.uuid: make_native_type(CONNECTION_TYPE_ORACLE, or_types.NVARCHAR2),
        UserDataType.markup: make_native_type(CONNECTION_TYPE_ORACLE, or_types.NVARCHAR2),
        UserDataType.unsupported: make_native_type(CONNECTION_TYPE_ORACLE, sa.sql.sqltypes.NullType),
    }

lib/dl_connector_oracle/dl_connector_oracle/core/type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Callable,
    ClassVar,
)

import attr

from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.us_connection_base import (
    ClassicConnectionSQL,
    ConnectionBase,
    DataSourceTemplate,
)
from dl_i18n.localizer_base import Localizer

from dl_connector_oracle.core.constants import (
    SOURCE_TYPE_ORACLE_SUBSELECT,
    SOURCE_TYPE_ORACLE_TABLE,
    OracleDbNameType,
)
from dl_connector_oracle.core.dto import OracleConnDTO


class ConnectionSQLOracle(ClassicConnectionSQL):
    has_schema: ClassVar[bool] = True
    # Default schema is usually defined on a per-user basis,
    # so it's better to omit the schema if it isn't explicitly specified.
    default_schema_name = None
    source_type = SOURCE_TYPE_ORACLE_TABLE
    allowed_source_types = frozenset((SOURCE_TYPE_ORACLE_TABLE, SOURCE_TYPE_ORACLE_SUBSELECT))
    allow_dashsql: ClassVar[bool] = True
    allow_cache: ClassVar[bool] = True
    is_always_user_source: ClassVar[bool] = True

    @attr.s(kw_only=True)
    class DataModel(ClassicConnectionSQL.DataModel):
        db_name_type: OracleDbNameType = attr.ib(default=OracleDbNameType.service_name)

    def get_conn_dto(self) -> OracleConnDTO:
        return OracleConnDTO(
            conn_id=self.uuid,
            host=self.data.host,
            multihosts=self.parse_multihosts(),  # type: ignore  # TODO: fix
            port=self.data.port,
            db_name=self.data.db_name,
            db_name_type=self.data.db_name_type,
            username=self.data.username,
            password=self.password,
        )

    def get_data_source_template_templates(self, localizer: Localizer) -> list[DataSourceTemplate]:
        return self._make_subselect_templates(source_type=SOURCE_TYPE_ORACLE_SUBSELECT, localizer=localizer)

    def get_parameter_combinations(
        self, conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase]
    ) -> list[dict]:
        if not self.db_name:
            return []

        assert self.has_schema
        return [
            dict(schema_name=tid.schema_name, table_name=tid.table_name)
            for tid in self.get_tables(schema_name=None, conn_executor_factory=conn_executor_factory)
        ]

    @property
    def allow_public_usage(self) -> bool:
        return True

lib/dl_connector_oracle/dl_connector_oracle/core/us_connection.py
datalens-tech/datalens-backend
from dl_db_testing.connectors.base.connector import DbTestingConnector

from dl_connector_oracle.db_testing.engine_wrapper import OracleEngineWrapper


class OracleDbTestingConnector(DbTestingConnector):
    engine_wrapper_classes = (OracleEngineWrapper,)

lib/dl_connector_oracle/dl_connector_oracle/db_testing/connector.py
datalens-tech/datalens-backend
import re
from typing import (
    Any,
    Optional,
    Sequence,
)

import sqlalchemy as sa

from dl_db_testing.database.engine_wrapper import EngineWrapperBase


class OracleEngineWrapper(EngineWrapperBase):
    URL_PREFIX = "oracle"  # Not using the bi_* version because we only need basic functionality here

    def execute(self, query: Any, *multiparams: Any, **params: Any):  # type: ignore  # TODO: fix
        # FIXME: Note the following problem:
        #  https://stackoverflow.com/questions/54709396/incorrect-type-conversion-with-cx-oracle-and-sqlalchemy-queries

        return super().execute(query, *multiparams, **params)

    def count_sql_sessions(self) -> int:
        # noinspection SqlDialectInspection
        cur = self.execute("SELECT * FROM sys.V_$SESSION")
        try:
            lines = cur.fetchall()
            return len(lines)
        finally:
            cur.close()

    def get_conn_credentials(self, full: bool = False) -> dict:
        cred_prefix, oracle_dsn = str(self.url).split("@")
        username, password = cred_prefix.split("//")[1].split(":")
        match = re.search(
            r"HOST=(?P<host>[^)]+)\).*PORT=(?P<port>\d+)\).*(SERVICE_NAME|SID)=(?P<db_name>[^)]+)\)", oracle_dsn
        )
        return dict(
            host=match.group("host"),  # type: ignore  # TODO: fix
            port=int(match.group("port")),  # type: ignore  # TODO: fix
            username=username,
            password=password,
            db_name=match.group("db_name"),  # type: ignore  # TODO: fix
        )

    def insert_into_table(self, table: sa.Table, data: Sequence[dict]) -> None:
        # Multi-row insert doesn't work correctly
        for row in data:
            self.execute(table.insert(row))

    def create_schema(self, schema_name: str) -> None:
        self.execute(f"CREATE USER {self.quote(schema_name)} IDENTIFIED BY qwerty")
        self.execute(f"GRANT ALL PRIVILEGES TO {self.quote(schema_name)}")

    def get_version(self) -> Optional[str]:
        return self.execute("SELECT * FROM V$VERSION").scalar()

lib/dl_connector_oracle/dl_connector_oracle/db_testing/engine_wrapper.py
datalens-tech/datalens-backend
from sqlalchemy.dialects.oracle.base import OracleDialect as SAOracleDialect

from dl_formula.connectors.base.connector import FormulaConnector
from dl_formula.mutation.general import OptimizeUnaryBoolFunctions

from dl_connector_oracle.formula.constants import OracleDialect as OracleDialectNS
from dl_connector_oracle.formula.context_processor import OracleContextPostprocessor
from dl_connector_oracle.formula.definitions.all import DEFINITIONS
from dl_connector_oracle.formula.literal import OracleLiteralizer
from dl_connector_oracle.formula.type_constructor import OracleTypeConstructor


class OracleFormulaConnector(FormulaConnector):
    dialect_ns_cls = OracleDialectNS
    dialects = OracleDialectNS.ORACLE
    default_dialect = OracleDialectNS.ORACLE
    op_definitions = DEFINITIONS
    literalizer_cls = OracleLiteralizer
    context_processor_cls = OracleContextPostprocessor
    type_constructor_cls = OracleTypeConstructor
    sa_dialect = SAOracleDialect()

    @classmethod
    def registration_hook(cls) -> None:
        OptimizeUnaryBoolFunctions.register_dialect(
            "isnull", cls.dialects, f=lambda x: x is None or (isinstance(x, str) and x == "")
        )

lib/dl_connector_oracle/dl_connector_oracle/formula/connector.py
datalens-tech/datalens-backend
from dl_formula.core.dialect import (
    DialectName,
    DialectNamespace,
    simple_combo,
)


DIALECT_NAME_ORACLE = DialectName.declare("ORACLE")


class OracleDialect(DialectNamespace):
    ORACLE_12_0 = simple_combo(name=DIALECT_NAME_ORACLE, version=(12, 0))
    ORACLE = ORACLE_12_0

lib/dl_connector_oracle/dl_connector_oracle/formula/constants.py
datalens-tech/datalens-backend
from typing import cast

import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

from dl_formula.connectors.base.context_processor import BooleanlessContextPostprocessor
from dl_formula.core.datatype import DataType


class OracleContextPostprocessor(BooleanlessContextPostprocessor):
    def booleanize_expression(self, data_type: DataType, expression: ClauseElement) -> ClauseElement:
        if data_type in (
            DataType.CONST_BOOLEAN,
            DataType.CONST_INTEGER,
            DataType.CONST_FLOAT,
            DataType.BOOLEAN,
            DataType.INTEGER,
            DataType.FLOAT,
        ):
            expression = cast(ClauseElement, expression != 0)
        elif data_type in (DataType.CONST_STRING, DataType.STRING):
            expression = expression.isnot(None)
        elif data_type in (DataType.CONST_DATE, DataType.CONST_DATETIME, DataType.DATE, DataType.DATETIME):
            expression = cast(ClauseElement, sa.literal(1) == 1)
        else:
            expression = cast(ClauseElement, expression == 1)

        return expression

    def debooleanize_expression(self, data_type: DataType, expression: ClauseElement) -> ClauseElement:
        return sa.case(whens=[(expression, sa.literal(1))], else_=sa.literal(0))

lib/dl_connector_oracle/dl_connector_oracle/formula/context_processor.py
datalens-tech/datalens-backend
from dl_connector_oracle.formula.definitions.conditional_blocks import DEFINITIONS_COND_BLOCKS
from dl_connector_oracle.formula.definitions.functions_aggregation import DEFINITIONS_AGG
from dl_connector_oracle.formula.definitions.functions_datetime import DEFINITIONS_DATETIME
from dl_connector_oracle.formula.definitions.functions_logical import DEFINITIONS_LOGICAL
from dl_connector_oracle.formula.definitions.functions_markup import DEFINITIONS_MARKUP
from dl_connector_oracle.formula.definitions.functions_math import DEFINITIONS_MATH
from dl_connector_oracle.formula.definitions.functions_string import DEFINITIONS_STRING
from dl_connector_oracle.formula.definitions.functions_type import DEFINITIONS_TYPE
from dl_connector_oracle.formula.definitions.operators_binary import DEFINITIONS_BINARY
from dl_connector_oracle.formula.definitions.operators_ternary import DEFINITIONS_TERNARY
from dl_connector_oracle.formula.definitions.operators_unary import DEFINITIONS_UNARY


DEFINITIONS = [
    *DEFINITIONS_COND_BLOCKS,
    *DEFINITIONS_AGG,
    *DEFINITIONS_DATETIME,
    *DEFINITIONS_LOGICAL,
    *DEFINITIONS_MARKUP,
    *DEFINITIONS_MATH,
    *DEFINITIONS_STRING,
    *DEFINITIONS_TYPE,
    *DEFINITIONS_UNARY,
    *DEFINITIONS_BINARY,
    *DEFINITIONS_TERNARY,
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/all.py
datalens-tech/datalens-backend
from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.conditional_blocks as base

from dl_connector_oracle.formula.constants import OracleDialect as D


V = TranslationVariant.make


DEFINITIONS_COND_BLOCKS = [
    # _case_block_
    base.CaseBlock.for_dialect(D.ORACLE),
    # _if_block_
    base.IfBlock3.for_dialect(D.ORACLE),
    base.IfBlockMulti.for_dialect(D.ORACLE),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/conditional_blocks.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common import (
    quantile_value,
    within_group,
)
import dl_formula.definitions.functions_aggregation as base
from dl_formula.definitions.literals import un_literal

from dl_connector_oracle.formula.constants import OracleDialect as D


V = TranslationVariant.make


DEFINITIONS_AGG = [
    # avg
    base.AggAvgFromNumber.for_dialect(D.ORACLE),
    base.AggAvgFromDate.for_dialect(D.ORACLE),
    base.AggAvgFromDatetime.for_dialect(D.ORACLE),
    base.AggAvgFromDatetimeTZ.for_dialect(D.ORACLE),
    # avg_if
    base.AggAvgIf.for_dialect(D.ORACLE),
    # count
    base.AggCount0.for_dialect(D.ORACLE),
    base.AggCount1.for_dialect(D.ORACLE),
    # count_if
    base.AggCountIf.for_dialect(D.ORACLE),
    # countd
    base.AggCountd.for_dialect(D.ORACLE),
    # countd_approx
    base.AggCountdApprox(
        variants=[
            V(D.ORACLE_12_0, sa.func.APPROX_COUNT_DISTINCT),
        ]
    ),
    # countd_if
    base.AggCountdIf.for_dialect(D.ORACLE),
    # max
    base.AggMax.for_dialect(D.ORACLE),
    # median
    base.AggMedian(
        variants=[
            V(D.ORACLE, lambda expr: within_group(sa.func.PERCENTILE_DISC(0.5), expr)),
        ]
    ),
    # min
    base.AggMin.for_dialect(D.ORACLE),
    # quantile
    base.AggQuantile(
        variants=[
            V(
                D.ORACLE,
                lambda expr, quant: within_group(sa.func.percentile_disc(quantile_value(un_literal(quant))), expr),
            ),
        ]
    ),
    # stdev
    base.AggStdev.for_dialect(D.ORACLE),
    # stdevp
    base.AggStdevp.for_dialect(D.ORACLE),
    # sum
    base.AggSum.for_dialect(D.ORACLE),
    # sum_if
    base.AggSumIf.for_dialect(D.ORACLE),
    # var
    base.AggVar.for_dialect(D.ORACLE),
    # varp
    base.AggVarp.for_dialect(D.ORACLE),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/functions_aggregation.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common import raw_sql
from dl_formula.definitions.common_datetime import (
    EPOCH_START_D,
    EPOCH_START_DOW,
    datetime_interval,
)
import dl_formula.definitions.functions_datetime as base
from dl_formula.definitions.literals import (
    literal,
    un_literal,
)

from dl_connector_oracle.formula.constants import OracleDialect as D


V = TranslationVariant.make


class FuncDatetrunc2Oracle(base.FuncDatetrunc2):
    _oracle_fmt_map = {
        "minute": "MI",
        "hour": "HH",
        "day": "DDD",
        "week": "IW",
        "month": "MM",
        "quarter": "Q",
        "year": "YYYY",
    }

    variants = [
        V(
            D.ORACLE,
            lambda date, unit: (
                sa.func.TRUNC(
                    date, literal(FuncDatetrunc2Oracle._oracle_fmt_map[base.norm_datetrunc_unit(unit)], d=D.ORACLE)
                )
                if base.norm_datetrunc_unit(unit) in FuncDatetrunc2Oracle._oracle_fmt_map
                else date
            ),
        ),
    ]
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME},  # TODO: DataType.DATETIMETZ
                DataType.CONST_STRING,
            ]
        ),
    ]


DEFINITIONS_DATETIME = [
    # dateadd
    base.FuncDateadd1.for_dialect(D.ORACLE),
    base.FuncDateadd2Unit.for_dialect(D.ORACLE),
    base.FuncDateadd2Number.for_dialect(D.ORACLE),
    base.FuncDateadd3Legacy.for_dialect(D.ORACLE),
    base.FuncDateadd3DateConstNum(
        variants=[
            V(
                D.ORACLE,
                lambda date, what, num: (
                    sa.cast(date + datetime_interval(un_literal(what), un_literal(num), literal_mult=True), sa.Date)
                ),
            ),
        ]
    ),
    base.FuncDateadd3DatetimeConstNum(
        variants=[
            V(
                D.ORACLE,
                lambda dt, what, num: (
                    sa.cast(dt + datetime_interval(what.value, num.value, literal_mult=True), sa.Date)
                ),
            ),
        ]
    ),
    # datepart
    base.FuncDatepart2Legacy.for_dialect(D.ORACLE),
    base.FuncDatepart2.for_dialect(D.ORACLE),
    base.FuncDatepart3Const.for_dialect(D.ORACLE),
    base.FuncDatepart3NonConst.for_dialect(D.ORACLE),
    # datetrunc
    FuncDatetrunc2Oracle(),
    # day
    base.FuncDay(
        variants=[
            V(D.ORACLE, lambda date: sa.func.extract(raw_sql("DAY"), date)),
        ]
    ),
    # dayofweek
    base.FuncDayofweek1.for_dialect(D.ORACLE),
    base.FuncDayofweek2(
        variants=[
            V(
                D.ORACLE,
                lambda date, firstday: sa.cast(
                    date - EPOCH_START_D - (EPOCH_START_DOW + base.norm_fd(firstday) - 1), sa.Integer
                )
                % 7
                + 1,
            ),
        ]
    ),
    # genericnow
    base.FuncGenericNow(
        variants=[
            V(D.ORACLE, lambda: sa.type_coerce(sa.cast(raw_sql("SYSTIMESTAMP"), sa.Date), sa.DateTime)),
        ]
    ),
    # hour
    base.FuncHourDate.for_dialect(D.ORACLE),
    base.FuncHourDatetime(
        variants=[
            V(D.ORACLE, lambda date: sa.func.extract(raw_sql("HOUR"), sa.cast(date, sa.TIMESTAMP))),
        ]
    ),
    # minute
    base.FuncMinuteDate.for_dialect(D.ORACLE),
    base.FuncMinuteDatetime(
        variants=[
            V(D.ORACLE, lambda date: sa.func.extract(raw_sql("MINUTE"), sa.cast(date, sa.TIMESTAMP))),
        ]
    ),
    # month
    base.FuncMonth(
        variants=[
            V(D.ORACLE, lambda date: sa.func.extract(raw_sql("MONTH"), date)),
        ]
    ),
    # now
    base.FuncNow(
        variants=[
            V(D.ORACLE, lambda: sa.type_coerce(sa.cast(raw_sql("SYSTIMESTAMP"), sa.Date), sa.DateTime)),
        ]
    ),
    # quarter
    base.FuncQuarter(
        variants=[
            V(D.ORACLE, lambda date: sa.func.TRUNC((sa.func.extract(raw_sql("MONTH"), date) + 2) / 3)),
        ]
    ),
    # second
    base.FuncSecondDate.for_dialect(D.ORACLE),
    base.FuncSecondDatetime(
        variants=[
            V(D.ORACLE, lambda date: sa.func.extract(raw_sql("SECOND"), sa.cast(date, sa.TIMESTAMP))),
        ]
    ),
    # today
    base.FuncToday(
        variants=[
            V(D.ORACLE, lambda: sa.type_coerce(sa.func.TRUNC(sa.cast(raw_sql("SYSTIMESTAMP"), sa.Date)), sa.Date)),
        ]
    ),
    # week
    base.FuncWeek(
        variants=[
            V(D.ORACLE, lambda date: sa.cast(sa.func.TO_CHAR(date, "IW"), sa.Integer)),
        ]
    ),
    # year
    base.FuncYear(
        variants=[
            V(D.ORACLE, lambda date: sa.func.extract(raw_sql("YEAR"), date)),
        ]
    ),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/functions_datetime.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_logical as base

from dl_connector_oracle.formula.constants import OracleDialect as D


V = TranslationVariant.make


DEFINITIONS_LOGICAL = [
    # case
    base.FuncCase.for_dialect(D.ORACLE),
    # if
    base.FuncIf.for_dialect(D.ORACLE),
    # ifnull
    base.FuncIfnull(
        variants=[
            V(D.ORACLE, sa.func.NVL),
        ]
    ),
    # iif
    base.FuncIif3Legacy.for_dialect(D.ORACLE),
    # isnull
    base.FuncIsnull.for_dialect(D.ORACLE),
    # zn
    base.FuncZn(
        variants=[
            V(D.ORACLE, lambda x: sa.func.NVL(x, 0)),
        ]
    ),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/functions_logical.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_markup as base

from dl_connector_oracle.formula.constants import OracleDialect as D


DEFINITIONS_MARKUP = [
    # +
    base.BinaryPlusMarkup.for_dialect(D.ORACLE),
    # __str
    base.FuncInternalStrConst.for_dialect(D.ORACLE),
    base.FuncInternalStr.for_dialect(D.ORACLE),
    # bold
    base.FuncBold.for_dialect(D.ORACLE),
    # italic
    base.FuncItalics.for_dialect(D.ORACLE),
    # markup
    base.ConcatMultiMarkup.for_dialect(D.ORACLE),
    # url
    base.FuncUrl.for_dialect(D.ORACLE),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/functions_markup.py
datalens-tech/datalens-backend
import math

import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_math as base

from dl_connector_oracle.formula.constants import OracleDialect as D


V = TranslationVariant.make


DEFINITIONS_MATH = [
    # abs
    base.FuncAbs.for_dialect(D.ORACLE),
    # acos
    base.FuncAcos.for_dialect(D.ORACLE),
    # asin
    base.FuncAsin.for_dialect(D.ORACLE),
    # atan
    base.FuncAtan.for_dialect(D.ORACLE),
    # atan2
    base.FuncAtan2.for_dialect(D.ORACLE),
    # ceiling
    base.FuncCeiling.for_dialect(D.ORACLE),
    # cos
    base.FuncCos.for_dialect(D.ORACLE),
    # cot
    base.FuncCot(
        variants=[
            V(D.ORACLE, lambda x: sa.func.COS(x) / sa.func.SIN(x)),
        ]
    ),
    # degrees
    base.FuncDegrees(
        variants=[
            V(D.ORACLE, lambda x: x / math.pi * 180),
        ]
    ),
    # div
    base.FuncDivBasic(
        variants=[
            V(D.ORACLE, lambda x, y: sa.func.TRUNC(x / y)),
        ]
    ),
    # div_safe
    base.FuncDivSafe2(
        variants=[
            V(D.ORACLE, lambda x, y: sa.func.TRUNC(x / sa.func.nullif(y, 0))),
        ]
    ),
    base.FuncDivSafe3(
        variants=[
            V(D.ORACLE, lambda x, y, default: sa.func.coalesce(sa.func.TRUNC(x / sa.func.nullif(y, 0)), default)),
        ]
    ),
    # exp
    base.FuncExp.for_dialect(D.ORACLE),
    # fdiv_safe
    base.FuncFDivSafe2.for_dialect(D.ORACLE),
    base.FuncFDivSafe3.for_dialect(D.ORACLE),
    # floor
    base.FuncFloor.for_dialect(D.ORACLE),
    # greatest
    base.FuncGreatest1.for_dialect(D.ORACLE),
    base.FuncGreatestMain.for_dialect(D.ORACLE),
    base.GreatestMulti.for_dialect(D.ORACLE),
    # least
    base.FuncLeast1.for_dialect(D.ORACLE),
    base.FuncLeastMain.for_dialect(D.ORACLE),
    base.LeastMulti.for_dialect(D.ORACLE),
    # ln
    base.FuncLn.for_dialect(D.ORACLE),
    # log
    base.FuncLog.for_dialect(D.ORACLE),
    # log10
    base.FuncLog10(
        variants=[
            V(D.ORACLE, lambda x: sa.func.LOG(10, x)),
        ]
    ),
    # pi
    base.FuncPi(
        variants=[
            V(D.ORACLE, lambda: sa.literal(math.pi)),
        ]
    ),
    # power
    base.FuncPower.for_dialect(D.ORACLE),
    # radians
    base.FuncRadians(
        variants=[
            V(D.ORACLE, lambda x: x * math.pi / 180),
        ]
    ),
    # round
    base.FuncRound1.for_dialect(D.ORACLE),
    base.FuncRound2.for_dialect(D.ORACLE),
    # sign
    base.FuncSign.for_dialect(D.ORACLE),
    # sin
    base.FuncSin.for_dialect(D.ORACLE),
    # sqrt
    base.FuncSqrt.for_dialect(D.ORACLE),
    # square
    base.FuncSquare(
        variants=[
            V(D.ORACLE, lambda x: sa.func.POWER(x, 2)),
        ]
    ),
    # tan
    base.FuncTan.for_dialect(D.ORACLE),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/functions_math.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common import make_binary_chain
import dl_formula.definitions.functions_string as base

from dl_connector_oracle.formula.constants import OracleDialect as D


V = TranslationVariant.make


DEFINITIONS_STRING = [
    # ascii
    base.FuncAscii.for_dialect(D.ORACLE),
    # char
    base.FuncChar(
        variants=[
            V(D.ORACLE, sa.func.CHR),
        ]
    ),
    # concat
    base.Concat1.for_dialect((D.ORACLE)),
    base.ConcatMultiStrConst.for_dialect(D.ORACLE),
    base.ConcatMultiStr(
        variants=[
            V(
                D.ORACLE,
                lambda *args: make_binary_chain(
                    # chain of 2-arg CONCAT function calls
                    (lambda x, y: sa.func.CONCAT(x, y)),
                    *args,
                    wrap_as_nodes=False,
                ),
            ),
        ]
    ),
    base.ConcatMultiAny.for_dialect(D.ORACLE),
    # contains
    base.FuncContainsConst.for_dialect(D.ORACLE),
    base.FuncContainsNonConst(
        variants=[
            V(D.ORACLE, lambda x, y: sa.func.INSTR(x, y) > 0),
        ]
    ),
    base.FuncContainsNonString.for_dialect(D.ORACLE),
    # notcontains
    base.FuncNotContainsConst.for_dialect(D.ORACLE),
    base.FuncNotContainsNonConst.for_dialect(D.ORACLE),
    base.FuncNotContainsNonString.for_dialect(D.ORACLE),
    # endswith
    base.FuncEndswithConst.for_dialect(D.ORACLE),
    base.FuncEndswithNonConst(
        variants=[
            V(
                D.ORACLE,
                lambda x, y: (sa.func.SUBSTR(x, sa.func.LENGTH(x) - sa.func.LENGTH(y) + 1, sa.func.LENGTH(y)) == y),
            ),
        ]
    ),
    base.FuncEndswithNonString.for_dialect(D.ORACLE),
    # find
    base.FuncFind2(
        variants=[
            V(D.ORACLE, sa.func.INSTR),
        ]
    ),
    base.FuncFind3(
        variants=[
            V(D.ORACLE, sa.func.INSTR),
        ]
    ),
    # icontains
    base.FuncIContainsNonConst.for_dialect(D.ORACLE),
    base.FuncIContainsNonString.for_dialect(D.ORACLE),
    # iendswith
    base.FuncIEndswithNonConst.for_dialect(D.ORACLE),
    base.FuncIEndswithNonString.for_dialect(D.ORACLE),
    # istartswith
    base.FuncIStartswithNonConst.for_dialect(D.ORACLE),
    base.FuncIStartswithNonString.for_dialect(D.ORACLE),
    # left
    base.FuncLeft(
        variants=[
            V(D.ORACLE, lambda x, y: sa.func.SUBSTR(x, 1, y)),
        ]
    ),
    # len
    base.FuncLenString(
        variants=[
            V(D.ORACLE, sa.func.LENGTH),
        ]
    ),
    # lower
    base.FuncLowerConst.for_dialect(D.ORACLE),
    base.FuncLowerNonConst.for_dialect(D.ORACLE),
    # ltrim
    base.FuncLtrim.for_dialect(D.ORACLE),
    # regexp_extract
    base.FuncRegexpExtract(
        variants=[
            V(D.ORACLE, sa.func.REGEXP_SUBSTR),
        ]
    ),
    # regexp_extract_nth
    base.FuncRegexpExtractNth(
        variants=[
            V(D.ORACLE, lambda text, pattern, ind: sa.func.REGEXP_SUBSTR(text, pattern, 1, ind)),
        ]
    ),
    # regexp_match
    base.FuncRegexpMatch(
        variants=[
            V(D.ORACLE, lambda text, pattern: sa.func.REGEXP_SUBSTR(text, pattern).isnot(None)),
        ]
    ),
    # regexp_replace
    base.FuncRegexpReplace(
        variants=[
            V(D.ORACLE, sa.func.REGEXP_REPLACE),
        ]
    ),
    # replace
    base.FuncReplace.for_dialect(D.ORACLE),
    # right
    base.FuncRight(
        variants=[
            V(D.ORACLE, lambda x, y: sa.func.SUBSTR(x, sa.func.LENGTH(x) - y + 1, y)),
        ]
    ),
    # rtrim
    base.FuncRtrim.for_dialect(D.ORACLE),
    # space
    base.FuncSpaceConst.for_dialect(D.ORACLE),
    base.FuncSpaceNonConst(
        variants=[
            V(D.ORACLE, lambda size: sa.func.SUBSTR(sa.func.LPAD(".", size + 1, " "), 1, size)),
        ]
    ),
    # startswith
    base.FuncStartswithConst.for_dialect(D.ORACLE),
    base.FuncStartswithNonConst(
        variants=[
            V(D.ORACLE, lambda x, y: (sa.func.SUBSTR(x, 1, sa.func.LENGTH(y)) == y)),
        ]
    ),
    base.FuncStartswithNonString.for_dialect(D.ORACLE),
    # substr
    base.FuncSubstr2(
        variants=[
            V(D.ORACLE, lambda text, start: sa.func.SUBSTR(text, start, sa.func.LENGTH(text) - start + 1)),
        ]
    ),
    base.FuncSubstr3(
        variants=[
            V(D.ORACLE, sa.func.SUBSTR),
        ]
    ),
    # trim
    base.FuncTrim.for_dialect(D.ORACLE),
    # upper
    base.FuncUpperConst.for_dialect(D.ORACLE),
    base.FuncUpperNonConst.for_dialect(D.ORACLE),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/functions_string.py
datalens-tech/datalens-backend
import sqlalchemy as sa
import sqlalchemy.dialects.oracle.base as sa_oracle

from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common_datetime import (
    DAY_SEC,
    EPOCH_START_D,
)
import dl_formula.definitions.functions_type as base

from dl_connector_oracle.formula.constants import OracleDialect as D


V = TranslationVariant.make


DEFINITIONS_TYPE = [
    # bool
    base.FuncBoolFromNull(
        variants=[
            V(D.ORACLE, lambda _: sa.cast(sa.null(), sa_oracle.INTEGER())),
        ]
    ),
    base.FuncBoolFromNumber.for_dialect(D.ORACLE),
    base.FuncBoolFromBool.for_dialect(D.ORACLE),
    base.FuncBoolFromStrGeo(
        variants=[
            V(D.ORACLE, lambda value: value.isnot(None)),
        ]
    ),
    base.FuncBoolFromDateDatetime.for_dialect(D.ORACLE),
    # date
    base.FuncDate1FromNull.for_dialect(D.ORACLE),
    base.FuncDate1FromDatetime(
        variants=[
            V(D.ORACLE, lambda expr: expr),
        ]
    ),
    base.FuncDate1FromString(
        variants=[
            V(D.ORACLE, lambda expr: sa.func.TO_DATE(expr, "YYYY-MM-DD")),
        ]
    ),
    base.FuncDate1FromNumber(
        variants=[
            V(D.ORACLE, lambda expr: EPOCH_START_D + expr / DAY_SEC),
        ]
    ),
    # datetime
    base.FuncDatetime1FromNull(
        variants=[
            V(D.ORACLE, lambda _: sa.cast(sa.null(), sa.Date())),
        ]
    ),
    base.FuncDatetime1FromDatetime.for_dialect(D.ORACLE),
    base.FuncDatetime1FromDate(
        variants=[
            V(D.ORACLE, lambda expr: expr),
        ]
    ),
    base.FuncDatetime1FromNumber(
        variants=[
            V(D.ORACLE, lambda expr: EPOCH_START_D + expr / DAY_SEC),
        ]
    ),
    base.FuncDatetime1FromString(
        variants=[
            V(D.ORACLE, lambda expr: sa.func.TO_DATE(expr, "YYYY-MM-DD HH:MI:SS")),
        ]
    ),
    # datetimetz
    base.FuncDatetimeTZConst.for_dialect(D.ORACLE),
    # float
    base.FuncFloatNumber(
        variants=[
            V(D.ORACLE, lambda value: sa.cast(value, sa_oracle.BINARY_DOUBLE)),
        ]
    ),
    base.FuncFloatString(
        variants=[
            V(D.ORACLE, lambda value: sa.cast(value, sa_oracle.BINARY_DOUBLE)),
        ]
    ),
    base.FuncFloatFromBool(
        variants=[
            V(D.ORACLE, lambda value: sa.cast(value, sa_oracle.BINARY_DOUBLE)),
        ]
    ),
    base.FuncFloatFromDate(
        variants=[
            V(D.ORACLE, lambda value: sa.type_coerce((value - EPOCH_START_D) * DAY_SEC, sa_oracle.BINARY_DOUBLE)),
        ]
    ),
    base.FuncFloatFromDatetime(
        variants=[
            V(D.ORACLE, lambda value: sa.type_coerce((value - EPOCH_START_D) * DAY_SEC, sa_oracle.BINARY_DOUBLE)),
        ]
    ),
    base.FuncFloatFromGenericDatetime(
        variants=[
            V(D.ORACLE, lambda value: sa.type_coerce((value - EPOCH_START_D) * DAY_SEC, sa_oracle.BINARY_DOUBLE)),
        ]
    ),
    # genericdatetime
    base.FuncGenericDatetime1FromNull(
        variants=[
            V(D.ORACLE, lambda _: sa.cast(sa.null(), sa.Date())),
        ]
    ),
    base.FuncGenericDatetime1FromDatetime.for_dialect(D.ORACLE),
    base.FuncGenericDatetime1FromDate(
        variants=[
            V(D.ORACLE, lambda expr: expr),
        ]
    ),
    base.FuncGenericDatetime1FromNumber(
        variants=[
            V(D.ORACLE, lambda expr: EPOCH_START_D + expr / DAY_SEC),
        ]
    ),
    base.FuncGenericDatetime1FromString(
        variants=[
            V(D.ORACLE, lambda expr: sa.func.TO_DATE(expr, "YYYY-MM-DD HH:MI:SS")),
        ]
    ),
    # geopoint
    base.FuncGeopointFromStr.for_dialect(D.ORACLE),
    base.FuncGeopointFromCoords.for_dialect(D.ORACLE),
    # geopolygon
    base.FuncGeopolygon.for_dialect(D.ORACLE),
    # int
    base.FuncIntFromNull(
        variants=[
            V(D.ORACLE, lambda _: sa.cast(sa.null(), sa_oracle.INTEGER())),
        ]
    ),
    base.FuncIntFromInt.for_dialect(D.ORACLE),
    base.FuncIntFromFloat(
        variants=[
            V(D.ORACLE, lambda value: sa.cast(sa.func.FLOOR(value), sa.Integer)),
        ]
    ),
    base.FuncIntFromBool(
        variants=[
            V(D.ORACLE, lambda value: value),
        ]
    ),
    base.FuncIntFromStr(
        variants=[
            V(D.ORACLE, lambda value: sa.cast(value, sa.Integer)),
        ]
    ),
    base.FuncIntFromDate(
        variants=[
            V(D.ORACLE, lambda value: sa.cast((value - EPOCH_START_D) * DAY_SEC, sa.Integer)),
        ]
    ),
    base.FuncIntFromDatetime(
        variants=[
            V(D.ORACLE, lambda value: sa.cast((value - EPOCH_START_D) * DAY_SEC, sa.Integer)),
        ]
    ),
    base.FuncIntFromGenericDatetime(
        variants=[
            V(D.ORACLE, lambda value: sa.cast((value - EPOCH_START_D) * DAY_SEC, sa.Integer)),
        ]
    ),
    # str
    base.FuncStrFromUnsupported(
        variants=[
            V(D.ORACLE, sa.func.TO_CHAR),
        ]
    ),
    base.FuncStrFromInteger(
        variants=[
            V(D.ORACLE, sa.func.TO_CHAR),
        ]
    ),
    base.FuncStrFromFloat(
        variants=[
            V(D.ORACLE, sa.func.TO_CHAR),
        ]
    ),
    base.FuncStrFromBool(
        variants=[
            V(
                D.ORACLE,
                lambda value: sa.case(
                    whens=[(value.is_(None), sa.null()), (value != sa.literal(0), "True")], else_="False"
                ),
            ),
        ]
    ),
    base.FuncStrFromStrGeo.for_dialect(D.ORACLE),
    base.FuncStrFromDate(
        variants=[
            V(D.ORACLE, lambda value: sa.func.TO_CHAR(value, "YYYY-MM-DD")),
        ]
    ),
    base.FuncStrFromDatetime(
        variants=[
            V(D.ORACLE, lambda value: sa.func.TO_CHAR(value, "YYYY-MM-DD HH:MI:SS")),
        ]
    ),
    base.FuncStrFromString.for_dialect(D.ORACLE),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/functions_type.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.operators_binary as base

from dl_connector_oracle.formula.constants import OracleDialect as D


V = TranslationVariant.make


DEFINITIONS_BINARY = [
    # !=
    base.BinaryNotEqual.for_dialect(D.ORACLE),
    # %
    base.BinaryModInteger.for_dialect(D.ORACLE),
    base.BinaryModFloat.for_dialect(D.ORACLE),
    # *
    base.BinaryMultNumbers.for_dialect(D.ORACLE),
    base.BinaryMultStringConst.for_dialect(D.ORACLE),
    base.BinaryMultStringNonConst(
        variants=[
            V(
                D.ORACLE,
                lambda text, size: sa.func.SUBSTR(
                    sa.func.LPAD(".", size * sa.func.LENGTH(text) + 1, text), 1, size * sa.func.LENGTH(text)
                ),
            ),
        ]
    ),
    # +
    base.BinaryPlusNumbers.for_dialect(D.ORACLE),
    base.BinaryPlusStrings.for_dialect(D.ORACLE),
    base.BinaryPlusDateInt(
        variants=[
            V(D.ORACLE, lambda date, days: date + days),
        ]
    ),
    base.BinaryPlusDateFloat(
        variants=[
            V(D.ORACLE, lambda date, days: date + days),
        ]
    ),
    base.BinaryPlusDatetimeNumber(
        variants=[
            V(D.ORACLE, lambda dt, days: dt + days),
        ]
    ),
    base.BinaryPlusGenericDatetimeNumber(
        variants=[
            V(D.ORACLE, lambda dt, days: dt + days),
        ]
    ),
    # -
    base.BinaryMinusNumbers.for_dialect(D.ORACLE),
    base.BinaryMinusDateInt(
        variants=[
            V(D.ORACLE, lambda date, days: date - days),
        ]
    ),
    base.BinaryMinusDateFloat(
        variants=[
            V(D.ORACLE, lambda date, days: date - days),
        ]
    ),
    base.BinaryMinusDatetimeNumber(
        variants=[
            V(D.ORACLE, lambda dt, days: dt - days),
        ]
    ),
    base.BinaryMinusGenericDatetimeNumber(
        variants=[
            V(D.ORACLE, lambda dt, days: dt - days),
        ]
    ),
    base.BinaryMinusDates.for_dialect(D.ORACLE),
    base.BinaryMinusDatetimes(
        variants=[
            V(D.ORACLE, lambda left, right: left - right),
        ]
    ),
    base.BinaryMinusGenericDatetimes(
        variants=[
            V(D.ORACLE, lambda left, right: left - right),
        ]
    ),
    # /
    base.BinaryDivInt.for_dialect(D.ORACLE),
    base.BinaryDivFloat.for_dialect(D.ORACLE),
    # <
    base.BinaryLessThan.for_dialect(D.ORACLE),
    # <=
    base.BinaryLessThanOrEqual.for_dialect(D.ORACLE),
    # ==
    base.BinaryEqual.for_dialect(D.ORACLE),
    # >
    base.BinaryGreaterThan.for_dialect(D.ORACLE),
    # >=
    base.BinaryGreaterThanOrEqual.for_dialect(D.ORACLE),
    # ^
    base.BinaryPower.for_dialect(D.ORACLE),
    # _!=
    base.BinaryNotEqualInternal.for_dialect(D.ORACLE),
    # _==
    base.BinaryEqualInternal.for_dialect(D.ORACLE),
    # _dneq
    base.BinaryEqualDenullified.for_dialect(D.ORACLE),
    # and
    base.BinaryAnd.for_dialect(D.ORACLE),
    # in
    base.BinaryIn.for_dialect(D.ORACLE),
    # like
    base.BinaryLike.for_dialect(D.ORACLE),
    # notin
    base.BinaryNotIn.for_dialect(D.ORACLE),
    # notlike
    base.BinaryNotLike.for_dialect(D.ORACLE),
    # or
    base.BinaryOr.for_dialect(D.ORACLE),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/operators_binary.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_ternary as base

from dl_connector_oracle.formula.constants import OracleDialect as D


DEFINITIONS_TERNARY = [
    # between
    base.TernaryBetween.for_dialect(D.ORACLE),
    # notbetween
    base.TernaryNotBetween.for_dialect(D.ORACLE),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/operators_ternary.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.operators_unary as base

from dl_connector_oracle.formula.constants import OracleDialect as D


V = TranslationVariant.make


DEFINITIONS_UNARY = [
    # isfalse
    base.UnaryIsFalseStringGeo(
        variants=[
            V(D.ORACLE, lambda x: x.is_(None)),
        ]
    ),
    base.UnaryIsFalseNumbers.for_dialect(D.ORACLE),
    base.UnaryIsFalseDateTime(
        variants=[
            V(D.ORACLE, lambda x: sa.literal(0)),
        ]
    ),
    base.UnaryIsFalseBoolean(
        variants=[
            V(D.ORACLE, lambda x: x == 0),
        ]
    ),
    # istrue
    base.UnaryIsTrueStringGeo(
        variants=[
            V(D.ORACLE, lambda x: x.isnot(None)),
        ]
    ),
    base.UnaryIsTrueNumbers.for_dialect(D.ORACLE),
    base.UnaryIsTrueDateTime(
        variants=[
            V(D.ORACLE, lambda x: sa.literal(1)),
        ]
    ),
    base.UnaryIsTrueBoolean(
        variants=[
            V(D.ORACLE, lambda x: x != 0),
        ]
    ),
    # neg
    base.UnaryNegate.for_dialect(D.ORACLE),
    # not
    base.UnaryNotBool.for_dialect(D.ORACLE),
    base.UnaryNotNumbers.for_dialect(D.ORACLE),
    base.UnaryNotStringGeo(
        variants=[
            V(D.ORACLE, lambda x: x.is_(None)),
        ]
    ),
    base.UnaryNotDateDatetime(
        variants=[
            V(D.ORACLE, lambda x: sa.literal(0)),
        ]
    ),
]

lib/dl_connector_oracle/dl_connector_oracle/formula/definitions/operators_unary.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.connectors.base.literal import (
    Literal,
    Literalizer,
)
from dl_formula.core.dialect import DialectCombo


class OracleLiteralizer(Literalizer):
    __slots__ = ()

    def literal_bool(self, value: bool, dialect: DialectCombo) -> Literal:
        return sa.literal(int(value))

    def literal_str(self, value: str, dialect: DialectCombo) -> Literal:
        # SA uses NVARCHAR2 for string literals by default,
        # but CHAR has better compatibility with some of the functions, so try to use it
        # with a fallback to NCHAR for non-ASCII strings
        try:
            value.encode("ascii")
            type_ = sa.CHAR(len(value))
        except UnicodeEncodeError:
            type_ = sa.NCHAR(len(value))
        return sa.literal(value, type_)

lib/dl_connector_oracle/dl_connector_oracle/formula/literal.py
datalens-tech/datalens-backend
import sqlalchemy as sa
import sqlalchemy.dialects.oracle.base as sa_oracle
from sqlalchemy.types import TypeEngine

from dl_formula.connectors.base.type_constructor import DefaultSATypeConstructor
from dl_formula.core.datatype import DataType


class OracleTypeConstructor(DefaultSATypeConstructor):
    def get_sa_type(self, data_type: DataType) -> TypeEngine:
        type_map: dict[DataType, TypeEngine] = {
            DataType.BOOLEAN: sa_oracle.NUMBER(1, 0),
            DataType.DATETIME: sa.Date(),
            DataType.GENERICDATETIME: sa.Date(),
        }
        if (type_eng := type_map.get(data_type)) is not None:
            return type_eng
        else:
            return super().get_sa_type(data_type)

lib/dl_connector_oracle/dl_connector_oracle/formula/type_constructor.py
datalens-tech/datalens-backend
from dl_formula_ref.texts import StyledDialect

from dl_connector_oracle.formula.constants import OracleDialect
from dl_connector_oracle.formula_ref.i18n import Translatable


HUMAN_DIALECTS = {
    OracleDialect.ORACLE_12_0: StyledDialect(
        "`Oracle Database 12c (12.1)`",
        "`Oracle`<br/>`Database 12c`<br/>`(12.1)`",
        Translatable("`Oracle Database` version `12c (12.1)`"),
    ),
}

lib/dl_connector_oracle/dl_connector_oracle/formula_ref/human_dialects.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_oracle as package


DOMAIN = f"dl_formula_ref_{package.__name__}"

_LOCALE_DIR = os.path.join(os.path.dirname(__file__), "..", "locales")

CONFIGS = [
    TranslationConfig(path=_LOCALE_DIR, domain=DOMAIN, locale="en"),
    TranslationConfig(path=_LOCALE_DIR, domain=DOMAIN, locale="ru"),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_oracle/dl_connector_oracle/formula_ref/i18n.py
datalens-tech/datalens-backend
from dl_formula_ref.plugins.base.plugin import FormulaRefPlugin

from dl_connector_oracle.formula.constants import OracleDialect
from dl_connector_oracle.formula_ref.human_dialects import HUMAN_DIALECTS
from dl_connector_oracle.formula_ref.i18n import CONFIGS


class OracleFormulaRefPlugin(FormulaRefPlugin):
    any_dialects = frozenset((*OracleDialect.ORACLE.to_list(),))
    compeng_support_dialects = frozenset((*OracleDialect.ORACLE.to_list(),))
    human_dialects = HUMAN_DIALECTS
    translation_configs = frozenset(CONFIGS)

lib/dl_connector_oracle/dl_connector_oracle/formula_ref/plugin.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_api_lib_testing.data_api_base import StandardizedDataApiTestBase
from dl_api_lib_testing.dataset_base import DatasetTestBase
from dl_constants.enums import RawSQLLevel

from dl_connector_oracle.core.constants import (
    CONNECTION_TYPE_ORACLE,
    SOURCE_TYPE_ORACLE_TABLE,
    OracleDbNameType,
)
from dl_connector_oracle_tests.db.config import (
    API_TEST_CONFIG,
    CoreConnectionSettings,
)
from dl_connector_oracle_tests.db.core.base import BaseOracleTestClass


class OracleConnectionTestBase(BaseOracleTestClass, ConnectionTestBase):
    conn_type = CONNECTION_TYPE_ORACLE
    bi_compeng_pg_on = False

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def connection_params(self) -> dict:
        return dict(
            db_connect_method=OracleDbNameType.service_name.name,
            db_name=CoreConnectionSettings.DB_NAME,
            host=CoreConnectionSettings.HOST,
            port=CoreConnectionSettings.PORT,
            username=CoreConnectionSettings.USERNAME,
            password=CoreConnectionSettings.PASSWORD,
            **(dict(raw_sql_level=self.raw_sql_level.value) if self.raw_sql_level is not None else {}),
        )


class OracleDashSQLConnectionTest(OracleConnectionTestBase):
    raw_sql_level = RawSQLLevel.dashsql


class OracleDatasetTestBase(OracleConnectionTestBase, DatasetTestBase):
    @pytest.fixture(scope="class")
    def dataset_params(self, sample_table) -> dict:
        return dict(
            source_type=SOURCE_TYPE_ORACLE_TABLE.name,
            parameters=dict(
                db_name=sample_table.db.name,
                schema_name=sample_table.schema,
                table_name=sample_table.name,
            ),
        )


class OracleDataApiTestBase(OracleDatasetTestBase, StandardizedDataApiTestBase):
    mutation_caches_on = False

lib/dl_connector_oracle/dl_connector_oracle_tests/db/api/base.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.complex_queries import DefaultBasicComplexQueryTestSuite
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_oracle_tests.db.api.base import OracleDataApiTestBase


class TestOracleBasicComplexQueries(OracleDataApiTestBase, DefaultBasicComplexQueryTestSuite):
    test_params = RegulatedTestParams(
        mark_features_skipped={
            DefaultBasicComplexQueryTestSuite.feature_window_functions: "Native window functions are not implemented"
        }
    )

lib/dl_connector_oracle/dl_connector_oracle_tests/db/api/test_complex_queries.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite

from dl_connector_oracle_tests.db.api.base import OracleConnectionTestBase


class TestOracleConnection(OracleConnectionTestBase, DefaultConnectorConnectionTestSuite):
    pass

lib/dl_connector_oracle/dl_connector_oracle_tests/db/api/test_connection.py
datalens-tech/datalens-backend
from aiohttp.test_utils import TestClient
import pytest

from dl_api_lib_testing.connector.dashsql_suite import DefaultDashSQLTestSuite

from dl_connector_oracle_tests.db.api.base import OracleDashSQLConnectionTest
from dl_connector_oracle_tests.db.config import SUBSELECT_QUERY_FULL


class TestOracleDashSQL(OracleDashSQLConnectionTest, DefaultDashSQLTestSuite):
    @pytest.fixture(scope="class")
    def dashsql_basic_query(self) -> str:
        return "select 1, 2, 3 from dual"

    @pytest.mark.asyncio
    async def test_result(self, data_api_lowlevel_aiohttp_client: TestClient, saved_connection_id: str):
        resp = await self.get_dashsql_response(
            data_api_aio=data_api_lowlevel_aiohttp_client,
            conn_id=saved_connection_id,
            query=SUBSELECT_QUERY_FULL,
        )

        resp_data = await resp.json()
        assert resp_data[0]["event"] == "metadata", resp_data
        assert resp_data[0]["data"]["names"] == [
            "NUM",
            "NUM_STR",
            "NUM_INTEGER",
            "NUM_NUMBER",
            "NUM_BINARY_FLOAT",
            "NUM_BINARY_DOUBLE",
            "NUM_CHAR",
            "NUM_VARCHAR",
            "NUM_VARCHAR2",
            "NUM_NCHAR",
            "NUM_NVARCHAR2",
            "NUM_DATE",
            "NUM_TIMESTAMP",
            "NUM_TIMESTAMP_TZ",
        ]
        assert resp_data[0]["data"]["driver_types"] == [
            "db_type_number",
            "db_type_varchar",
            "db_type_number",
            "db_type_number",
            "db_type_binary_float",
            "db_type_binary_double",
            "db_type_char",
            "db_type_varchar",
            "db_type_varchar",
            "db_type_nchar",
            "db_type_nvarchar",
            "db_type_date",
            "db_type_timestamp",
            "db_type_timestamp_tz",
        ]
        assert resp_data[0]["data"]["db_types"] == [
            "integer",
            "varchar",
            "integer",
            "integer",
            "binary_float",
            "binary_double",
            "char",
            "varchar",
            "varchar",
            "nchar",
            "nvarchar",
            "date",
            "timestamp",
            "timestamp",
        ]
        assert resp_data[0]["data"]["bi_types"] == [
            "integer",
            "string",
            "integer",
            "integer",
            "float",
            "float",
            "string",
            "string",
            "string",
            "string",
            "string",
            "genericdatetime",
            "genericdatetime",
            "genericdatetime",
        ]

lib/dl_connector_oracle/dl_connector_oracle_tests/db/api/test_dashsql.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.data_api_suites import (
    DefaultConnectorDataDistinctTestSuite,
    DefaultConnectorDataGroupByFormulaTestSuite,
    DefaultConnectorDataPreviewTestSuite,
    DefaultConnectorDataRangeTestSuite,
    DefaultConnectorDataResultTestSuite,
)
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_oracle_tests.db.api.base import OracleDataApiTestBase


class TestOracleDataResult(OracleDataApiTestBase, DefaultConnectorDataResultTestSuite):
    test_params = RegulatedTestParams(
        mark_features_skipped={
            DefaultConnectorDataResultTestSuite.array_support: "Oracle doesn't support arrays",
        },
        mark_tests_failed={
            DefaultConnectorDataResultTestSuite.test_get_result_with_string_filter_operations_for_numbers: (
                "BI-4978: Need to ignore the exponent"
            )
        },
    )


class TestOracleDataGroupBy(OracleDataApiTestBase, DefaultConnectorDataGroupByFormulaTestSuite):
    pass


class TestOracleDataRange(OracleDataApiTestBase, DefaultConnectorDataRangeTestSuite):
    pass


class TestOracleDataDistinct(OracleDataApiTestBase, DefaultConnectorDataDistinctTestSuite):
    pass


class TestOracleDataPreview(OracleDataApiTestBase, DefaultConnectorDataPreviewTestSuite):
    pass

lib/dl_connector_oracle/dl_connector_oracle_tests/db/api/test_data.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite

from dl_connector_oracle_tests.db.api.base import OracleDatasetTestBase


class TestOracleDataset(OracleDatasetTestBase, DefaultConnectorDatasetTestSuite):
    pass

lib/dl_connector_oracle/dl_connector_oracle_tests/db/api/test_dataset.py
datalens-tech/datalens-backend
import os
from typing import ClassVar

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport

from dl_connector_oracle.formula.constants import OracleDialect as D


# Infra settings
CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=51811).host,
    port_us_http=get_test_container_hostport("us", fallback_port=51811).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=51810).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=51810).port,
    us_master_token="AC1ofiek8coB",
)

COMPOSE_PROJECT_NAME = os.environ.get("COMPOSE_PROJECT_NAME", "dl_connector_oracle")
ORACLE_CONTAINER_LABEL = "db-oracle"


class CoreConnectionSettings:
    DB_NAME: ClassVar[str] = "XEPDB1"
    HOST: ClassVar[str] = get_test_container_hostport("db-oracle", fallback_port=51800).host
    PORT: ClassVar[int] = get_test_container_hostport("db-oracle", fallback_port=51800).port
    USERNAME: ClassVar[str] = "datalens"
    PASSWORD: ClassVar[str] = "qwerty"


DEFAULT_ORACLE_SCHEMA_NAME = "DATALENS"


SUBSELECT_QUERY_FULL = r"""
select
    num,
    'test' || num as num_str,
    cast(num as integer) as num_integer,
    cast(num as number) as num_number,
    -- cast(num as number(9,9)) as num_number_9_9,
    cast(num as binary_float) as num_binary_float,
    cast(num as binary_double) as num_binary_double,
    cast(num as char) as num_char,
    cast(num as varchar(3)) as num_varchar,
    cast(num as varchar2(4)) as num_varchar2,
    cast(num as nchar) as num_nchar,
    -- cast(num as nvarchar(5)) as num_nvarchar,
    cast(num as nvarchar2(5)) as num_nvarchar2,
    DATE '2020-01-01' + num as num_date,
    TIMESTAMP '1999-12-31 23:59:59.10' + numToDSInterval(num, 'second') as num_timestamp,
    TIMESTAMP '1999-12-31 23:59:59.10-07:00' + numToDSInterval(num, 'second') as num_timestamp_tz
from (
    select 0 as num from dual
    union all
    select 1 as num from dual
    union all
    select 6 as num from dual
) sq
"""

_DB_URL = (
    f'oracle://datalens:qwerty@(DESCRIPTION=(ADDRESS=(PROTOCOL=TCP)(HOST={get_test_container_hostport("db-oracle", fallback_port=51800).host})'
    f'(PORT={get_test_container_hostport("db-oracle", fallback_port=51800).port}))(CONNECT_DATA=(SERVICE_NAME={CoreConnectionSettings.DB_NAME})))'
)
SYSDBA_URL = (
    f'oracle://sys:qwerty@(DESCRIPTION=(ADDRESS=(PROTOCOL=TCP)(HOST={get_test_container_hostport("db-oracle", fallback_port=51800).host})'
    f'(PORT={get_test_container_hostport("db-oracle", fallback_port=51800).port}))(CONNECT_DATA=(SERVICE_NAME={CoreConnectionSettings.DB_NAME})))?mode=sysdba'
)
DB_CORE_URL = _DB_URL
DB_URLS = {
    D.ORACLE_12_0: _DB_URL,
}

API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["oracle"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
)

lib/dl_connector_oracle/dl_connector_oracle_tests/db/config.py
datalens-tech/datalens-backend
from dl_api_lib_testing.initialization import initialize_api_lib_test
from dl_core_testing.database import (
    CoreDbConfig,
    CoreDbDispenser,
)
from dl_db_testing.database.engine_wrapper import DbEngineConfig
from dl_formula_testing.forced_literal import forced_literal_use

from dl_connector_oracle.core.constants import CONNECTION_TYPE_ORACLE
from dl_connector_oracle_tests.db.config import (
    API_TEST_CONFIG,
    SYSDBA_URL,
)


pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)
    initialize_db()


def initialize_db():
    # Granting of priveleges con only be done when connecting as sydba
    db_dispenser = CoreDbDispenser()
    sysdba_db_config = CoreDbConfig(engine_config=DbEngineConfig(url=SYSDBA_URL), conn_type=CONNECTION_TYPE_ORACLE)
    sysdba_db = db_dispenser.get_database(db_config=sysdba_db_config)
    sysdba_db.execute("GRANT SELECT ON sys.V_$SESSION TO datalens")


__all__ = (
    # auto-use fixtures:
    "forced_literal_use",
)

lib/dl_connector_oracle/dl_connector_oracle_tests/db/conftest.py
datalens-tech/datalens-backend
import asyncio
from typing import Generator
import uuid

import pytest

from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import Db
from dl_core_testing.testcases.connection import BaseConnectionTestClass

from dl_connector_oracle.core.constants import CONNECTION_TYPE_ORACLE
from dl_connector_oracle.core.testing.connection import make_oracle_saved_connection
from dl_connector_oracle.core.us_connection import ConnectionSQLOracle
import dl_connector_oracle_tests.db.config as test_config


class BaseOracleTestClass(BaseConnectionTestClass[ConnectionSQLOracle]):
    conn_type = CONNECTION_TYPE_ORACLE
    core_test_config = test_config.CORE_TEST_CONFIG

    @pytest.fixture(autouse=True)
    # FIXME: This fixture is a temporary solution for failing core tests when they are run together with api tests
    def loop(self, event_loop: asyncio.AbstractEventLoop) -> Generator[asyncio.AbstractEventLoop, None, None]:
        asyncio.set_event_loop(event_loop)
        yield event_loop
        # Attempt to cover an old version of pytest-asyncio:
        # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
        asyncio.set_event_loop_policy(None)

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return test_config.DB_CORE_URL

    @pytest.fixture(scope="function")
    def connection_creation_params(self) -> dict:
        return dict(
            db_name=test_config.CoreConnectionSettings.DB_NAME,
            host=test_config.CoreConnectionSettings.HOST,
            port=test_config.CoreConnectionSettings.PORT,
            username=test_config.CoreConnectionSettings.USERNAME,
            password=test_config.CoreConnectionSettings.PASSWORD,
            **(dict(raw_sql_level=self.raw_sql_level) if self.raw_sql_level is not None else {}),
        )

    @pytest.fixture(scope="class")
    def _empty_table(self, db: Db) -> None:
        # So that template listings are not empty
        table_name = f"t_{uuid.uuid4().hex[:6]}"
        db.execute(f'CREATE TABLE datalens.{table_name} ("value" VARCHAR2(255))')

    @pytest.fixture(scope="function")
    def saved_connection(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params: dict,
        _empty_table,
    ) -> ConnectionSQLOracle:
        conn = make_oracle_saved_connection(sync_usm=sync_us_manager, **connection_creation_params)
        return conn

lib/dl_connector_oracle/dl_connector_oracle_tests/db/core/base.py
datalens-tech/datalens-backend
from dl_core.us_connection_base import DataSourceTemplate
from dl_core_testing.testcases.connection import DefaultConnectionTestClass

from dl_connector_oracle.core.us_connection import ConnectionSQLOracle
from dl_connector_oracle_tests.db.config import DEFAULT_ORACLE_SCHEMA_NAME
from dl_connector_oracle_tests.db.core.base import BaseOracleTestClass


class TestOracleConnection(
    BaseOracleTestClass,
    DefaultConnectionTestClass[ConnectionSQLOracle],
):
    do_check_data_export_flag = True

    def check_saved_connection(self, conn: ConnectionSQLOracle, params: dict) -> None:
        assert conn.uuid is not None
        assert conn.data.db_name == params["db_name"]
        assert conn.data.host == params["host"]
        assert conn.data.port == params["port"]
        assert conn.data.username == params["username"]
        assert conn.data.password == params["password"]

    def check_data_source_templates(
        self,
        conn: ConnectionSQLOracle,
        dsrc_templates: list[DataSourceTemplate],
    ) -> None:
        assert dsrc_templates
        for dsrc_tmpl in dsrc_templates:
            assert dsrc_tmpl.title
            if dsrc_tmpl.parameters.get("schema_name") is not None:
                assert dsrc_tmpl.group == [dsrc_tmpl.parameters["schema_name"]]

        schema_names = {tmpl.parameters.get("schema_name") for tmpl in dsrc_templates}
        assert DEFAULT_ORACLE_SCHEMA_NAME in schema_names

lib/dl_connector_oracle/dl_connector_oracle_tests/db/core/test_connection.py
datalens-tech/datalens-backend
from typing import (
    Optional,
    Sequence,
)

import pytest
import sqlalchemy as sa
from sqlalchemy.dialects.oracle import base as oracle_types

from dl_constants.enums import UserDataType
from dl_core.connection_models.common_models import DBIdent
from dl_core_testing.testcases.connection_executor import (
    DefaultAsyncConnectionExecutorTestSuite,
    DefaultSyncAsyncConnectionExecutorCheckBase,
    DefaultSyncConnectionExecutorTestSuite,
)
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_oracle.core.us_connection import ConnectionSQLOracle
from dl_connector_oracle_tests.db.config import CoreConnectionSettings
from dl_connector_oracle_tests.db.core.base import BaseOracleTestClass


class OracleSyncAsyncConnectionExecutorCheckBase(
    BaseOracleTestClass,
    DefaultSyncAsyncConnectionExecutorCheckBase[ConnectionSQLOracle],
):
    test_params = RegulatedTestParams(
        mark_tests_failed={
            DefaultAsyncConnectionExecutorTestSuite.test_error_on_select_from_nonexistent_source: "",  # TODO: FIXME
            DefaultAsyncConnectionExecutorTestSuite.test_get_table_schema_info_for_nonexistent_table: (
                "Empty schema is returned instead of an error"
            ),  # FIXME
        },
    )

    @pytest.fixture(scope="function")
    def db_ident(self) -> DBIdent:
        return DBIdent(db_name=CoreConnectionSettings.DB_NAME)

    def check_db_version(self, db_version: Optional[str]) -> None:
        assert db_version is not None
        assert "." in db_version

    @pytest.fixture(scope="class")
    def query_for_session_check(self) -> str:
        return "SELECT 567 FROM DUAL"


class TestOracleSyncConnectionExecutor(
    OracleSyncAsyncConnectionExecutorCheckBase,
    DefaultSyncConnectionExecutorTestSuite[ConnectionSQLOracle],
):
    subselect_query_for_schema_test = "(SELECT 1 AS num FROM DUAL)"

    def get_schemas_for_type_recognition(self) -> dict[str, Sequence[DefaultSyncConnectionExecutorTestSuite.CD]]:
        return {
            "oracle_types_number": [
                self.CD(sa.Numeric(16, 0), UserDataType.integer, nt_name="integer"),
                self.CD(sa.Numeric(16, 8), UserDataType.float, nt_name="number"),
                self.CD(oracle_types.NUMBER(), UserDataType.integer, nt_name="integer"),
                self.CD(oracle_types.NUMBER(10, 5), UserDataType.float),
                self.CD(oracle_types.BINARY_FLOAT(), UserDataType.float),
                self.CD(oracle_types.BINARY_DOUBLE(), UserDataType.float),
            ],
            "oracle_types_string": [
                self.CD(sa.CHAR(), UserDataType.string),
                self.CD(sa.NCHAR(), UserDataType.string),
                self.CD(sa.String(length=256), UserDataType.string, nt_name="varchar"),
                self.CD(oracle_types.VARCHAR2(100), UserDataType.string, nt_name="varchar"),
                self.CD(oracle_types.NVARCHAR2(100), UserDataType.string),
            ],
            "oracle_types_date": [
                self.CD(oracle_types.DATE(), UserDataType.genericdatetime),
                self.CD(oracle_types.TIMESTAMP(), UserDataType.genericdatetime),
            ],
        }


class TestOracleAsyncConnectionExecutor(
    OracleSyncAsyncConnectionExecutorCheckBase,
    DefaultAsyncConnectionExecutorTestSuite[ConnectionSQLOracle],
):
    test_params = RegulatedTestParams(
        mark_tests_failed={
            DefaultAsyncConnectionExecutorTestSuite.test_closing_sql_sessions: "Sessions not closed",  # TODO: FIXME
        },
    )
lib/dl_connector_oracle/dl_connector_oracle_tests/db/core/test_connection_executor.py
datalens-tech/datalens-backend
import pytest

from dl_constants.enums import (
    RawSQLLevel,
    UserDataType,
)
from dl_core.data_source_spec.sql import (
    StandardSchemaSQLDataSourceSpec,
    SubselectDataSourceSpec,
)
from dl_core.db import SchemaColumn
from dl_core_testing.fixtures.sample_tables import TABLE_SPEC_SAMPLE_SUPERSTORE
from dl_core_testing.testcases.data_source import (
    DataSourceTestByViewClass,
    DefaultDataSourceTestClass,
)

from dl_connector_oracle.core.constants import (
    SOURCE_TYPE_ORACLE_SUBSELECT,
    SOURCE_TYPE_ORACLE_TABLE,
)
from dl_connector_oracle.core.data_source import (
    OracleDataSource,
    OracleSubselectDataSource,
)
from dl_connector_oracle.core.us_connection import ConnectionSQLOracle
from dl_connector_oracle_tests.db.config import SUBSELECT_QUERY_FULL
from dl_connector_oracle_tests.db.core.base import BaseOracleTestClass


def _update_utype_for_oracle(user_type: UserDataType) -> UserDataType:
    if user_type == UserDataType.date:
        return UserDataType.genericdatetime
    return user_type


SAMPLE_TABLE_SCHEMA_SUPERSTORE_ORACLIZED = [
    (name, _update_utype_for_oracle(user_type)) for name, user_type in TABLE_SPEC_SAMPLE_SUPERSTORE.table_schema
]


class TestOracleTableDataSource(
    BaseOracleTestClass,
    DefaultDataSourceTestClass[
        ConnectionSQLOracle,
        StandardSchemaSQLDataSourceSpec,
        OracleDataSource,
    ],
):
    DSRC_CLS = OracleDataSource

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self, sample_table) -> StandardSchemaSQLDataSourceSpec:
        dsrc_spec = StandardSchemaSQLDataSourceSpec(
            source_type=SOURCE_TYPE_ORACLE_TABLE,
            db_name=sample_table.db.name,
            schema_name=sample_table.schema,
            table_name=sample_table.name,
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return list(SAMPLE_TABLE_SCHEMA_SUPERSTORE_ORACLIZED)


class TestOracleSubselectDataSource(
    BaseOracleTestClass,
    DefaultDataSourceTestClass[
        ConnectionSQLOracle,
        SubselectDataSourceSpec,
        OracleSubselectDataSource,
    ],
):
    DSRC_CLS = OracleSubselectDataSource

    raw_sql_level = RawSQLLevel.subselect

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self, sample_table) -> SubselectDataSourceSpec:
        dsrc_spec = SubselectDataSourceSpec(
            source_type=SOURCE_TYPE_ORACLE_SUBSELECT,
            subsql=f'SELECT * FROM "{sample_table.name}"',
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return list(SAMPLE_TABLE_SCHEMA_SUPERSTORE_ORACLIZED)


class TestOracleSubselectByView(
    BaseOracleTestClass,
    DataSourceTestByViewClass[
        ConnectionSQLOracle,
        SubselectDataSourceSpec,
        OracleSubselectDataSource,
    ],
):
    DSRC_CLS = OracleSubselectDataSource

    raw_sql_level = RawSQLLevel.subselect

    @pytest.fixture(scope="session")
    def initial_data_source_spec(self) -> SubselectDataSourceSpec:
        dsrc_spec = SubselectDataSourceSpec(
            source_type=SOURCE_TYPE_ORACLE_SUBSELECT,
            subsql=SUBSELECT_QUERY_FULL,
        )
        return dsrc_spec

    def postprocess_view_schema(
        self, view_schema: list[SchemaColumn], cursor_schema: list[SchemaColumn]
    ) -> list[SchemaColumn]:
        result = super().postprocess_view_schema(view_schema, cursor_schema=cursor_schema)
        if len(view_schema) != len(cursor_schema):
            return result

        # Actual result seems to depend on the cx_Oracle version too much,
        # and DL uses these equivalently anyway.
        tnames = ("binary_double", "binary_float")
        for idx, schema_col in enumerate(result):
            if schema_col.native_type.name in tnames:
                cs_name = cursor_schema[idx].native_type.name
                if cs_name in tnames:
                    schema_col = schema_col.clone(native_type=schema_col.native_type.clone(name=cs_name))
                    result[idx] = schema_col
        return result

lib/dl_connector_oracle/dl_connector_oracle_tests/db/core/test_data_source.py
datalens-tech/datalens-backend
from dl_core_testing.testcases.dataset import DefaultDatasetTestSuite
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_oracle.core.constants import SOURCE_TYPE_ORACLE_TABLE
from dl_connector_oracle.core.us_connection import ConnectionSQLOracle
from dl_connector_oracle_tests.db.core.base import BaseOracleTestClass


class TestOracleDataset(BaseOracleTestClass, DefaultDatasetTestSuite[ConnectionSQLOracle]):
    source_type = SOURCE_TYPE_ORACLE_TABLE

    test_params = RegulatedTestParams(
        mark_tests_failed={
            DefaultDatasetTestSuite.test_get_param_hash: "db_name in dsrc",  # TODO: FIXME
        },
    )

lib/dl_connector_oracle/dl_connector_oracle_tests/db/core/test_dataset.py
datalens-tech/datalens-backend
import pytest

from dl_formula_testing.testcases.base import FormulaConnectorTestBase

from dl_connector_oracle.formula.constants import OracleDialect as D
from dl_connector_oracle_tests.db.config import DB_URLS


class OracleTestBase(FormulaConnectorTestBase):
    dialect = D.ORACLE_12_0
    supports_arrays = False
    supports_uuid = False
    bool_is_expression = True
    empty_str_is_null = True  # '' and NULL are the same thing
    null_casts_to_number = True

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return DB_URLS[self.dialect]

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/base.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.conditional_blocks import DefaultConditionalBlockFormulaConnectorTestSuite

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


class TestConditionalBlockOracle(OracleTestBase, DefaultConditionalBlockFormulaConnectorTestSuite):
    pass

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_conditional_blocks.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_aggregation import DefaultMainAggFunctionFormulaConnectorTestSuite

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


class TestMainAggFunctionOracle(OracleTestBase, DefaultMainAggFunctionFormulaConnectorTestSuite):
    supports_countd_approx = True
    supports_quantile = True
    supports_median = True

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_functions_aggregation.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_datetime import DefaultDateTimeFunctionFormulaConnectorTestSuite

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


class TestDateTimeFunctionOracle(OracleTestBase, DefaultDateTimeFunctionFormulaConnectorTestSuite):
    supports_addition_to_feb_29 = False
    supports_deprecated_dateadd = True

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_functions_datetime.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_logical import DefaultLogicalFunctionFormulaConnectorTestSuite

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


class TestLogicalFunctionOracle(OracleTestBase, DefaultLogicalFunctionFormulaConnectorTestSuite):
    supports_iif = True

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_functions_logical.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_markup import DefaultMarkupFunctionFormulaConnectorTestSuite

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


class TestMarkupFunctionOracle(OracleTestBase, DefaultMarkupFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_functions_markup.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_math import DefaultMathFunctionFormulaConnectorTestSuite

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


class TestMathFunctionOracle(OracleTestBase, DefaultMathFunctionFormulaConnectorTestSuite):
    supports_atan_2_in_origin = False

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_functions_math.py
datalens-tech/datalens-backend
import os

import sqlalchemy as sa

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.functions_string import DefaultStringFunctionFormulaConnectorTestSuite

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


class TestStringFunctionOracle(OracleTestBase, DefaultStringFunctionFormulaConnectorTestSuite):
    supports_split_3 = False

    def test_oracle_unicode(self, dbe: DbEvaluator) -> None:
        assert os.environ.get("NLS_LANG") == ".AL32UTF8"
        # assert dbe.db.scalar(sa.select([sa.literal('карл')])) == 'карл'  # not working.
        assert dbe.db.scalar(sa.literal("карл", type_=sa.sql.sqltypes.NCHAR())) == "карл"

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_functions_string.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_type_conversion import (
    DefaultBoolTypeFunctionFormulaConnectorTestSuite,
    DefaultDateTypeFunctionFormulaConnectorTestSuite,
    DefaultFloatTypeFunctionFormulaConnectorTestSuite,
    DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite,
    DefaultGeopointTypeFunctionFormulaConnectorTestSuite,
    DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite,
    DefaultIntTypeFunctionFormulaConnectorTestSuite,
    DefaultStrTypeFunctionFormulaConnectorTestSuite,
)

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


# STR


class TestStrTypeFunctionOracle(OracleTestBase, DefaultStrTypeFunctionFormulaConnectorTestSuite):
    zero_float_to_str_value = "0"
    skip_custom_tz = True


# FLOAT


class TestFloatTypeFunctionOracle(OracleTestBase, DefaultFloatTypeFunctionFormulaConnectorTestSuite):
    pass


# BOOL


class TestBoolTypeFunctionOracle(OracleTestBase, DefaultBoolTypeFunctionFormulaConnectorTestSuite):
    pass


# INT


class TestIntTypeFunctionOracle(OracleTestBase, DefaultIntTypeFunctionFormulaConnectorTestSuite):
    pass


# DATE


class TestDateTypeFunctionOracle(OracleTestBase, DefaultDateTypeFunctionFormulaConnectorTestSuite):
    pass


# GENERICDATETIME (& DATETIME)


class TestGenericDatetimeTypeFunctionOracle(
    OracleTestBase,
    DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite,
):
    pass


# GEOPOINT


class TestGeopointTypeFunctionOracle(
    OracleTestBase,
    DefaultGeopointTypeFunctionFormulaConnectorTestSuite,
):
    pass


# GEOPOLYGON


class TestGeopolygonTypeFunctionOracle(
    OracleTestBase,
    DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite,
):
    pass

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_functions_type_conversion.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.literals import DefaultLiteralFormulaConnectorTestSuite

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


class TestConditionalBlockOracle(OracleTestBase, DefaultLiteralFormulaConnectorTestSuite):
    supports_microseconds = False
    supports_utc = False
    supports_custom_tz = False

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_literals.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.misc_funcs import DefaultMiscFunctionalityConnectorTestSuite

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


class TestMiscFunctionalityOracle(OracleTestBase, DefaultMiscFunctionalityConnectorTestSuite):
    pass

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_misc_funcs.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.operators import DefaultOperatorFormulaConnectorTestSuite

from dl_connector_oracle_tests.db.formula.base import OracleTestBase


class TestOperatorOracle(OracleTestBase, DefaultOperatorFormulaConnectorTestSuite):
    pass

lib/dl_connector_oracle/dl_connector_oracle_tests/db/formula/test_operators.py
datalens-tech/datalens-backend
from dl_api_connector.i18n.localizer import CONFIGS as BI_API_CONNECTOR_CONFIGS
from dl_api_lib_testing.connection_form_base import ConnectionFormTestBase

from dl_connector_oracle.api.connection_form.form_config import OracleConnectionFormFactory
from dl_connector_oracle.api.i18n.localizer import CONFIGS as BI_CONNECTOR_ORACLE_CONFIGS


class TestOracleConnectionForm(ConnectionFormTestBase):
    CONN_FORM_FACTORY_CLS = OracleConnectionFormFactory
    TRANSLATION_CONFIGS = BI_CONNECTOR_ORACLE_CONFIGS + BI_API_CONNECTOR_CONFIGS

lib/dl_connector_oracle/dl_connector_oracle_tests/unit/test_connection_form.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import fields

from dl_api_connector.api_schema.connection_base import ConnectionMetaMixin
from dl_api_connector.api_schema.connection_mixins import (
    DataExportForbiddenMixin,
    RawSQLLevelMixin,
)
from dl_api_connector.api_schema.connection_sql import ClassicSQLConnectionSchema
from dl_api_connector.api_schema.extras import FieldExtra
import dl_core.marshmallow as core_ma_fields

from dl_connector_postgresql.core.postgresql.us_connection import ConnectionPostgreSQL
from dl_connector_postgresql.core.postgresql_base.constants import PGEnforceCollateMode


class PostgreSQLConnectionSchema(
    ConnectionMetaMixin, DataExportForbiddenMixin, RawSQLLevelMixin, ClassicSQLConnectionSchema
):
    TARGET_CLS = ConnectionPostgreSQL
    ALLOW_MULTI_HOST = True

    enforce_collate = fields.Enum(
        PGEnforceCollateMode,
        attribute="data.enforce_collate",
        required=False,
        load_default=PGEnforceCollateMode.auto,
        dump_default=PGEnforceCollateMode.auto,
        bi_extra=FieldExtra(editable=True),
    )  # Side note: previously this fields was not declared in connection testing schemas
    ssl_enable = core_ma_fields.OnOffField(
        attribute="data.ssl_enable",
        required=False,
        load_default=False,
        bi_extra=FieldExtra(editable=True),
    )
    ssl_ca = core_ma_fields.Base64StringField(
        attribute="data.ssl_ca",
        required=False,
        allow_none=True,
        load_only=True,
        load_default=None,
        bi_extra=FieldExtra(editable=True),
    )

lib/dl_connector_postgresql/dl_connector_postgresql/api/api_schema/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import unique
from typing import (
    Optional,
    Sequence,
)

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormActionApiSchema,
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
)
from dl_api_connector.form_config.models.common import (
    CommonFieldName,
    FormFieldName,
)
import dl_api_connector.form_config.models.rows as C
from dl_api_connector.form_config.models.rows.base import FormRow
from dl_api_connector.form_config.models.shortcuts.rows import RowConstructor
from dl_configs.connectors_settings import ConnectorSettingsBase

from dl_connector_postgresql.api.connection_info import PostgreSQLConnectionInfoProvider
from dl_connector_postgresql.api.i18n.localizer import Translatable
from dl_connector_postgresql.core.postgresql_base.constants import PGEnforceCollateMode


@unique
class PostgreSQLFieldName(FormFieldName):
    enforce_collate = "enforce_collate"


class PostgreSQLConnectionFormFactory(ConnectionFormFactory):
    def _get_base_edit_api_schema(self) -> FormActionApiSchema:
        return FormActionApiSchema(
            items=[
                FormFieldApiSchema(name=CommonFieldName.host, required=True),
                FormFieldApiSchema(name=CommonFieldName.port, required=True),
                FormFieldApiSchema(name=CommonFieldName.username, required=True),
                FormFieldApiSchema(name=CommonFieldName.db_name, required=True),
                FormFieldApiSchema(name=CommonFieldName.password, required=self.mode == ConnectionFormMode.create),
                FormFieldApiSchema(name=CommonFieldName.cache_ttl_sec, nullable=True),
                FormFieldApiSchema(name=CommonFieldName.raw_sql_level),
                FormFieldApiSchema(name=PostgreSQLFieldName.enforce_collate),
                FormFieldApiSchema(name=CommonFieldName.ssl_enable),
                FormFieldApiSchema(name=CommonFieldName.ssl_ca),
                FormFieldApiSchema(name=CommonFieldName.data_export_forbidden),
            ],
        )

    def _get_base_create_api_schema(self, edit_api_schema: FormActionApiSchema) -> FormActionApiSchema:
        return FormActionApiSchema(
            items=[
                *edit_api_schema.items,
                *self._get_top_level_create_api_schema_items(),
            ],
            conditions=edit_api_schema.conditions.copy(),
        )

    def _get_base_check_api_schema(self) -> FormActionApiSchema:
        return FormActionApiSchema(
            items=[
                FormFieldApiSchema(name=CommonFieldName.host, required=True),
                FormFieldApiSchema(name=CommonFieldName.port, required=True),
                FormFieldApiSchema(name=CommonFieldName.db_name, required=True),
                FormFieldApiSchema(name=CommonFieldName.username, required=True),
                FormFieldApiSchema(name=CommonFieldName.password, required=self.mode == ConnectionFormMode.create),
                FormFieldApiSchema(name=CommonFieldName.ssl_enable),
                FormFieldApiSchema(name=CommonFieldName.ssl_ca),
                *self._get_top_level_check_api_schema_items(),
            ]
        )

    def _get_base_form_config(
        self,
        host_section: Sequence[FormRow],
        username_section: Sequence[FormRow],
        db_name_section: Sequence[FormRow],
        create_api_schema: FormActionApiSchema,
        edit_api_schema: FormActionApiSchema,
        check_api_schema: FormActionApiSchema,
        rc: RowConstructor,
    ) -> ConnectionForm:
        return ConnectionForm(
            title=PostgreSQLConnectionInfoProvider.get_title(self._localizer),
            rows=self._filter_nulls(
                [
                    *host_section,
                    rc.port_row(default_value="6432"),
                    *db_name_section,
                    *username_section,
                    rc.password_row(mode=self.mode),
                    C.CacheTTLRow(name=CommonFieldName.cache_ttl_sec),
                    rc.raw_sql_level_row(),
                    rc.collapse_advanced_settings_row(),
                    C.CustomizableRow(
                        items=[
                            C.LabelRowItem(
                                text=self._localizer.translate(Translatable("field_enforce-collate")),
                                display_conditions={CommonFieldName.advanced_settings: "opened"},
                            ),
                            C.RadioButtonRowItem(
                                name=PostgreSQLFieldName.enforce_collate,
                                options=[
                                    C.SelectableOption(
                                        text=self._localizer.translate(Translatable("value_enforce-collate-auto")),
                                        value=PGEnforceCollateMode.auto.value,
                                    ),
                                    C.SelectableOption(
                                        text=self._localizer.translate(Translatable("value_enforce-collate-off")),
                                        value=PGEnforceCollateMode.off.value,
                                    ),
                                    C.SelectableOption(
                                        text=self._localizer.translate(Translatable("value_enforce-collate-on")),
                                        value=PGEnforceCollateMode.on.value,
                                    ),
                                ],
                                default_value=PGEnforceCollateMode.auto.value,
                                display_conditions={CommonFieldName.advanced_settings: "opened"},
                            ),
                        ]
                    ),
                    *rc.ssl_rows(
                        enabled_name=CommonFieldName.ssl_enable,
                        enabled_help_text=self._localizer.translate(Translatable("label_postgres-ssl-enabled-tooltip")),
                        enabled_default_value=False,
                    ),
                    rc.data_export_forbidden_row(),
                ]
            ),
            api_schema=FormApiSchema(
                create=create_api_schema if self.mode == ConnectionFormMode.create else None,
                edit=edit_api_schema if self.mode == ConnectionFormMode.edit else None,
                check=check_api_schema,
            ),
        )

    def get_form_config(
        self,
        connector_settings: Optional[ConnectorSettingsBase],
        tenant: Optional[TenantDef],
    ) -> ConnectionForm:
        rc = RowConstructor(localizer=self._localizer)

        host_section = [rc.host_row()]
        username_section = [rc.username_row()]
        db_name_section = [rc.db_name_row()]

        edit_api_schema = self._get_base_edit_api_schema()
        create_api_schema = self._get_base_create_api_schema(edit_api_schema)
        check_api_schema = self._get_base_check_api_schema()

        return self._get_base_form_config(
            host_section=host_section,
            username_section=username_section,
            db_name_section=db_name_section,
            create_api_schema=create_api_schema,
            edit_api_schema=edit_api_schema,
            check_api_schema=check_api_schema,
            rc=rc,
        )

lib/dl_connector_postgresql/dl_connector_postgresql/api/connection_form/form_config.py
datalens-tech/datalens-backend
from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_postgresql.api.i18n.localizer import Translatable


class PostgreSQLConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-postgres")

lib/dl_connector_postgresql/dl_connector_postgresql/api/connection_info.py
datalens-tech/datalens-backend
from dl_api_connector.api_schema.source_base import (
    SchematizedSQLDataSourceSchema,
    SchematizedSQLDataSourceTemplateSchema,
    SubselectDataSourceSchema,
    SubselectDataSourceTemplateSchema,
)
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)
from dl_api_lib.query.registry import MQMFactorySettingItem
from dl_constants.enums import QueryProcessingMode
from dl_query_processing.multi_query.factory import NoCompengMultiQueryMutatorFactory

from dl_connector_postgresql.api.api_schema.connection import PostgreSQLConnectionSchema
from dl_connector_postgresql.api.connection_form.form_config import PostgreSQLConnectionFormFactory
from dl_connector_postgresql.api.connection_info import PostgreSQLConnectionInfoProvider
from dl_connector_postgresql.api.i18n.localizer import CONFIGS
from dl_connector_postgresql.core.postgresql.connector import (
    PostgreSQLCoreConnectionDefinition,
    PostgreSQLCoreConnector,
    PostgreSQLSubselectCoreSourceDefinition,
    PostgreSQLTableCoreSourceDefinition,
)
from dl_connector_postgresql.formula.constants import (
    DIALECT_NAME_POSTGRESQL,
    PostgreSQLDialect,
)


class PostgreSQLApiTableSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = PostgreSQLTableCoreSourceDefinition
    api_schema_cls = SchematizedSQLDataSourceSchema
    template_api_schema_cls = SchematizedSQLDataSourceTemplateSchema


class PostgreSQLApiSubselectSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = PostgreSQLSubselectCoreSourceDefinition
    api_schema_cls = SubselectDataSourceSchema
    template_api_schema_cls = SubselectDataSourceTemplateSchema


class PostgreSQLApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = PostgreSQLCoreConnectionDefinition
    api_generic_schema_cls = PostgreSQLConnectionSchema
    info_provider_cls = PostgreSQLConnectionInfoProvider
    form_factory_cls = PostgreSQLConnectionFormFactory


class PostgreSQLApiConnector(ApiConnector):
    core_connector_cls = PostgreSQLCoreConnector
    connection_definitions = (PostgreSQLApiConnectionDefinition,)
    source_definitions = (
        PostgreSQLApiTableSourceDefinition,
        PostgreSQLApiSubselectSourceDefinition,
    )
    formula_dialect_name = DIALECT_NAME_POSTGRESQL
    translation_configs = frozenset(CONFIGS)
    compeng_dialect = PostgreSQLDialect.COMPENG
    multi_query_mutation_factories = ApiConnector.multi_query_mutation_factories + (
        MQMFactorySettingItem(
            query_proc_mode=QueryProcessingMode.native_wf,
            dialects=PostgreSQLDialect.and_above(PostgreSQLDialect.POSTGRESQL_9_4).to_list(),
            factory_cls=NoCompengMultiQueryMutatorFactory,
        ),
    )

lib/dl_connector_postgresql/dl_connector_postgresql/api/connector.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_postgresql as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_postgresql/dl_connector_postgresql/api/i18n/localizer.py
datalens-tech/datalens-backend
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
)
from dl_core.connectors.sql_base.connector import (
    SQLSubselectCoreSourceDefinitionBase,
    SQLTableCoreSourceDefinitionBase,
)

from dl_connector_postgresql.core.postgresql.constants import (
    BACKEND_TYPE_POSTGRES,
    CONNECTION_TYPE_POSTGRES,
    SOURCE_TYPE_PG_SUBSELECT,
    SOURCE_TYPE_PG_TABLE,
)
from dl_connector_postgresql.core.postgresql.data_source import (
    PostgreSQLDataSource,
    PostgreSQLSubselectDataSource,
)
from dl_connector_postgresql.core.postgresql.data_source_migration import PostgreSQLDataSourceMigrator
from dl_connector_postgresql.core.postgresql.storage_schemas.connection import ConnectionPostgreSQLDataStorageSchema
from dl_connector_postgresql.core.postgresql.us_connection import ConnectionPostgreSQL
from dl_connector_postgresql.core.postgresql_base.adapters_postgres import PostgresAdapter
from dl_connector_postgresql.core.postgresql_base.async_adapters_postgres import AsyncPostgresAdapter
from dl_connector_postgresql.core.postgresql_base.connection_executors import (
    AsyncPostgresConnExecutor,
    PostgresConnExecutor,
)
from dl_connector_postgresql.core.postgresql_base.sa_types import SQLALCHEMY_POSTGRES_TYPES
from dl_connector_postgresql.core.postgresql_base.type_transformer import PostgreSQLTypeTransformer


class PostgreSQLCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_POSTGRES
    connection_cls = ConnectionPostgreSQL
    us_storage_schema_cls = ConnectionPostgreSQLDataStorageSchema
    type_transformer_cls = PostgreSQLTypeTransformer
    sync_conn_executor_cls = PostgresConnExecutor
    async_conn_executor_cls = AsyncPostgresConnExecutor
    dialect_string = "bi_postgresql"
    data_source_migrator_cls = PostgreSQLDataSourceMigrator


class PostgreSQLTableCoreSourceDefinition(SQLTableCoreSourceDefinitionBase):
    source_type = SOURCE_TYPE_PG_TABLE
    source_cls = PostgreSQLDataSource


class PostgreSQLSubselectCoreSourceDefinition(SQLSubselectCoreSourceDefinitionBase):
    source_type = SOURCE_TYPE_PG_SUBSELECT
    source_cls = PostgreSQLSubselectDataSource


class PostgreSQLCoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_POSTGRES
    connection_definitions = (PostgreSQLCoreConnectionDefinition,)
    source_definitions = (
        PostgreSQLTableCoreSourceDefinition,
        PostgreSQLSubselectCoreSourceDefinition,
    )
    rqe_adapter_classes = frozenset({PostgresAdapter, AsyncPostgresAdapter})
    sa_types = SQLALCHEMY_POSTGRES_TYPES

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
    SourceBackendType,
)


BACKEND_TYPE_POSTGRES = SourceBackendType.declare("POSTGRES")
CONNECTION_TYPE_POSTGRES = ConnectionType.declare("postgres")
SOURCE_TYPE_PG_TABLE = DataSourceType.declare("PG_TABLE")
SOURCE_TYPE_PG_SUBSELECT = DataSourceType.declare("PG_SUBSELECT")

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    ClassVar,
    FrozenSet,
)

from dl_constants.enums import (
    DataSourceType,
    JoinType,
)
from dl_core.data_source.sql import (
    BaseSQLDataSource,
    StandardSchemaSQLDataSource,
    SubselectDataSource,
)

from dl_connector_postgresql.core.postgresql.constants import (
    CONNECTION_TYPE_POSTGRES,
    SOURCE_TYPE_PG_SUBSELECT,
    SOURCE_TYPE_PG_TABLE,
)
from dl_connector_postgresql.core.postgresql_base.query_compiler import PostgreSQLQueryCompiler


class PostgreSQLDataSourceMixin(BaseSQLDataSource):
    compiler_cls = PostgreSQLQueryCompiler
    supported_join_types: ClassVar[FrozenSet[JoinType]] = frozenset(
        {
            JoinType.inner,
            JoinType.left,
            JoinType.full,
            JoinType.right,
        }
    )

    conn_type = CONNECTION_TYPE_POSTGRES

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type in (SOURCE_TYPE_PG_TABLE, SOURCE_TYPE_PG_SUBSELECT)


class PostgreSQLDataSource(PostgreSQLDataSourceMixin, StandardSchemaSQLDataSource):
    """PG table"""


class PostgreSQLSubselectDataSource(PostgreSQLDataSourceMixin, SubselectDataSource):
    """PG subselect"""

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/data_source.py
datalens-tech/datalens-backend
from dl_core.connectors.sql_base.data_source_migration import DefaultSQLDataSourceMigrator

from dl_connector_postgresql.core.postgresql.constants import (
    SOURCE_TYPE_PG_SUBSELECT,
    SOURCE_TYPE_PG_TABLE,
)


class PostgreSQLDataSourceMigrator(DefaultSQLDataSourceMigrator):
    table_source_type = SOURCE_TYPE_PG_TABLE
    subselect_source_type = SOURCE_TYPE_PG_SUBSELECT
    default_schema_name = "public"

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/data_source_migration.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_connector_postgresql.core.postgresql.constants import CONNECTION_TYPE_POSTGRES
from dl_connector_postgresql.core.postgresql_base.dto import PostgresConnDTOBase


@attr.s(frozen=True)
class PostgresConnDTO(PostgresConnDTOBase):
    conn_type = CONNECTION_TYPE_POSTGRES

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/dto.py
datalens-tech/datalens-backend
from dl_connector_postgresql.core.postgresql.us_connection import ConnectionPostgreSQL
from dl_connector_postgresql.core.postgresql_base.storage_schemas.connection import (
    ConnectionPostgreSQLBaseDataStorageSchema,
)


class ConnectionPostgreSQLDataStorageSchema(ConnectionPostgreSQLBaseDataStorageSchema):
    TARGET_CLS = ConnectionPostgreSQL.DataModel

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/storage_schemas/connection.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Optional,
)
import uuid

from dl_constants.enums import RawSQLLevel
from dl_core.us_manager.us_manager_sync import SyncUSManager

from dl_connector_postgresql.core.postgresql.constants import CONNECTION_TYPE_POSTGRES
from dl_connector_postgresql.core.postgresql.us_connection import ConnectionPostgreSQL


def make_postgresql_saved_connection(
    sync_usm: SyncUSManager,
    db_name: Optional[str],
    host: str,
    port: Optional[int],
    username: Optional[str],
    password: Optional[str],
    raw_sql_level: RawSQLLevel = RawSQLLevel.off,
    ssl_enable: bool = False,
    ssl_ca: Optional[str] = None,
    **kwargs: Any,
) -> ConnectionPostgreSQL:
    conn_name = "postgres test conn {}".format(uuid.uuid4())
    conn = ConnectionPostgreSQL.create_from_dict(
        data_dict=ConnectionPostgreSQL.DataModel(
            db_name=db_name,
            host=host,
            port=port,
            username=username,
            password=password,
            raw_sql_level=raw_sql_level,
            ssl_enable=ssl_enable,
            ssl_ca=ssl_ca,
        ),
        ds_key=conn_name,
        type_=CONNECTION_TYPE_POSTGRES.name,
        us_manager=sync_usm,
        **kwargs,
    )
    sync_usm.save(conn)
    return conn

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/testing/connection.py
datalens-tech/datalens-backend
from typing import Type

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.models.connection_target_dto_base import BaseSQLConnTargetDTO
from dl_core_testing.executors import ExecutorFactoryBase

from dl_connector_postgresql.core.postgresql_base.adapters_postgres import PostgresAdapter
from dl_connector_postgresql.core.postgresql_base.target_dto import PostgresConnTargetDTO


class PostgresExecutorFactory(ExecutorFactoryBase):
    def get_dto_class(self) -> Type[BaseSQLConnTargetDTO]:
        return PostgresConnTargetDTO

    def get_dba_class(self) -> Type[CommonBaseDirectAdapter]:
        return PostgresAdapter

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/testing/exec_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import ClassVar

from dl_core.us_connection_base import DataSourceTemplate
from dl_i18n.localizer_base import Localizer

from dl_connector_postgresql.core.postgresql.constants import (
    SOURCE_TYPE_PG_SUBSELECT,
    SOURCE_TYPE_PG_TABLE,
)
from dl_connector_postgresql.core.postgresql.dto import PostgresConnDTO
from dl_connector_postgresql.core.postgresql_base.us_connection import ConnectionPostgreSQLBase


class ConnectionPostgreSQL(ConnectionPostgreSQLBase):
    source_type = SOURCE_TYPE_PG_TABLE
    allowed_source_types = frozenset((SOURCE_TYPE_PG_TABLE, SOURCE_TYPE_PG_SUBSELECT))
    allow_dashsql: ClassVar[bool] = True
    allow_cache: ClassVar[bool] = True
    is_always_user_source: ClassVar[bool] = True

    def get_conn_dto(self) -> PostgresConnDTO:
        return PostgresConnDTO(
            conn_id=self.uuid,
            host=self.data.host,
            multihosts=self.parse_multihosts(),  # type: ignore  # TODO: fix
            port=self.data.port,
            db_name=self.data.db_name,
            username=self.data.username,
            password=self.password,
            enforce_collate=self.data.enforce_collate,
            ssl_enable=self.data.ssl_enable,
            ssl_ca=self.data.ssl_ca,
        )

    def get_data_source_template_templates(self, localizer: Localizer) -> list[DataSourceTemplate]:
        return self._make_subselect_templates(
            source_type=SOURCE_TYPE_PG_SUBSELECT,
            field_doc_key="PG_SUBSELECT/subsql",
            localizer=localizer,
        )

    @property
    def allow_public_usage(self) -> bool:
        return True

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql/us_connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
    Optional,
)

import attr
import sqlalchemy.dialects.postgresql as sa_pg

from dl_core.connectors.ssl_common.adapter import BaseSSLCertAdapter
from dl_core.db.native_type import SATypeSpec

from dl_connector_postgresql.core.postgresql.constants import CONNECTION_TYPE_POSTGRES
from dl_connector_postgresql.core.postgresql_base.constants import PGEnforceCollateMode
from dl_connector_postgresql.core.postgresql_base.target_dto import PostgresConnTargetDTO


# One way to obtain this data:
# docker exec bi-api_db-postgres_1 psql -U datalens -d  bi_db -A -t -F, -c 'select oid, typname from pg_type order by oid'
OID_KNOWLEDGE_RAW = """
16,bool
17,bytea
18,char
19,name
20,int8
21,int2
22,int2vector
23,int4
24,regproc
25,text
26,oid
27,tid
28,xid
29,cid
30,oidvector
71,pg_type
75,pg_attribute
81,pg_proc
83,pg_class
114,json
142,xml
143,_xml
194,pg_node_tree
199,_json
210,smgr
600,point
601,lseg
602,path
603,box
604,polygon
628,line
629,_line
650,cidr
651,_cidr
700,float4
701,float8
702,abstime
703,reltime
704,tinterval
705,unknown
718,circle
719,_circle
790,money
791,_money
829,macaddr
869,inet
1000,_bool
1001,_bytea
1002,_char
1003,_name
1005,_int2
1006,_int2vector
1007,_int4
1008,_regproc
1009,_text
1010,_tid
1011,_xid
1012,_cid
1013,_oidvector
1014,_bpchar
1015,_varchar
1016,_int8
1017,_point
1018,_lseg
1019,_path
1020,_box
1021,_float4
1022,_float8
1023,_abstime
1024,_reltime
1025,_tinterval
1027,_polygon
1028,_oid
1033,aclitem
1034,_aclitem
1040,_macaddr
1041,_inet
1042,bpchar
1043,varchar
1082,date
1083,time
1114,timestamp
1115,_timestamp
1182,_date
1183,_time
1184,timestamptz
1185,_timestamptz
1186,interval
1187,_interval
1231,_numeric
1248,pg_database
1263,_cstring
1266,timetz
1270,_timetz
1560,bit
1561,_bit
1562,varbit
1563,_varbit
1700,numeric
1790,refcursor
2201,_refcursor
2202,regprocedure
2203,regoper
2204,regoperator
2205,regclass
2206,regtype
2207,_regprocedure
2208,_regoper
2209,_regoperator
2210,_regclass
2211,_regtype
2249,record
2275,cstring
2276,any
2277,anyarray
2278,void
2279,trigger
2280,language_handler
2281,internal
2282,opaque
2283,anyelement
2287,_record
2776,anynonarray
2842,pg_authid
2843,pg_auth_members
2949,_txid_snapshot
2950,uuid
2951,_uuid
2970,txid_snapshot
3115,fdw_handler
3500,anyenum
3614,tsvector
3615,tsquery
3642,gtsvector
3643,_tsvector
3644,_gtsvector
3645,_tsquery
3734,regconfig
3735,_regconfig
3769,regdictionary
3770,_regdictionary
3831,anyrange
3838,event_trigger
3904,int4range
3905,_int4range
3906,numrange
3907,_numrange
3908,tsrange
3909,_tsrange
3910,tstzrange
3911,_tstzrange
3912,daterange
3913,_daterange
3926,int8range
3927,_int8range
10000,pg_attrdef
10001,pg_constraint
10002,pg_inherits
10003,pg_index
10004,pg_operator
10005,pg_opfamily
10006,pg_opclass
10117,pg_am
10118,pg_amop
10522,pg_amproc
10814,pg_language
10815,pg_largeobject_metadata
10816,pg_largeobject
10817,pg_aggregate
10818,pg_statistic
10819,pg_rewrite
10820,pg_trigger
10821,pg_event_trigger
10822,pg_description
10823,pg_cast
11020,pg_enum
11021,pg_namespace
11022,pg_conversion
11023,pg_depend
11024,pg_db_role_setting
11025,pg_tablespace
11026,pg_pltemplate
11027,pg_shdepend
11028,pg_shdescription
11029,pg_ts_config
11030,pg_ts_config_map
11031,pg_ts_dict
11032,pg_ts_parser
11033,pg_ts_template
11034,pg_extension
11035,pg_foreign_data_wrapper
11036,pg_foreign_server
11037,pg_user_mapping
11038,pg_foreign_table
11039,pg_default_acl
11040,pg_seclabel
11041,pg_shseclabel
11042,pg_collation
11043,pg_range
11055,pg_roles
11058,pg_shadow
11061,pg_group
11064,pg_user
11067,pg_rules
11071,pg_views
11075,pg_tables
11079,pg_matviews
11083,pg_indexes
11087,pg_stats
11091,pg_locks
11094,pg_cursors
11097,pg_available_extensions
11100,pg_available_extension_versions
11103,pg_prepared_xacts
11107,pg_prepared_statements
11110,pg_seclabels
11114,pg_settings
11119,pg_timezone_abbrevs
11122,pg_timezone_names
11125,pg_stat_all_tables
11129,pg_stat_xact_all_tables
11133,pg_stat_sys_tables
11137,pg_stat_xact_sys_tables
11140,pg_stat_user_tables
11144,pg_stat_xact_user_tables
11147,pg_statio_all_tables
11151,pg_statio_sys_tables
11154,pg_statio_user_tables
11157,pg_stat_all_indexes
11161,pg_stat_sys_indexes
11164,pg_stat_user_indexes
11167,pg_statio_all_indexes
11171,pg_statio_sys_indexes
11174,pg_statio_user_indexes
11177,pg_statio_all_sequences
11181,pg_statio_sys_sequences
11184,pg_statio_user_sequences
11187,pg_stat_activity
11190,pg_stat_replication
11193,pg_stat_database
11196,pg_stat_database_conflicts
11199,pg_stat_user_functions
11203,pg_stat_xact_user_functions
11207,pg_stat_bgwriter
11210,pg_user_mappings
11496,cardinal_number
11498,character_data
11499,sql_identifier
11501,information_schema_catalog_name
11503,time_stamp
11504,yes_or_no
11507,applicable_roles
11511,administrable_role_authorizations
11514,attributes
11518,character_sets
11522,check_constraint_routine_usage
11526,check_constraints
11530,collations
11533,collation_character_set_applicability
11536,column_domain_usage
11540,column_privileges
11544,column_udt_usage
11548,columns
11552,constraint_column_usage
11556,constraint_table_usage
11560,domain_constraints
11564,domain_udt_usage
11567,domains
11571,enabled_roles
11574,key_column_usage
11578,parameters
11582,referential_constraints
11586,role_column_grants
11589,routine_privileges
11593,role_routine_grants
11596,routines
11600,schemata
11603,sequences
11607,sql_features
11612,sql_implementation_info
11617,sql_languages
11622,sql_packages
11627,sql_parts
11632,sql_sizing
11637,sql_sizing_profiles
11642,table_constraints
11646,table_privileges
11650,role_table_grants
11653,tables
11657,triggered_update_columns
11661,triggers
11665,udt_privileges
11669,role_udt_grants
11672,usage_privileges
11676,role_usage_grants
11679,user_defined_types
11683,view_column_usage
11687,view_routine_usage
11691,view_table_usage
11695,views
11699,data_type_privileges
11703,element_types
11707,_pg_foreign_table_columns
11711,column_options
11714,_pg_foreign_data_wrappers
11717,foreign_data_wrapper_options
11720,foreign_data_wrappers
11723,_pg_foreign_servers
11727,foreign_server_options
11730,foreign_servers
11733,_pg_foreign_tables
11737,foreign_table_options
11740,foreign_tables
11743,_pg_user_mappings
11746,user_mapping_options
11750,user_mappings
"""
OID_KNOWLEDGE = {  # type: ignore  # TODO: fix
    int(oid): typname for oid, typname in (item.split(",", 1) for item in OID_KNOWLEDGE_RAW.strip().splitlines())
}

# `relkind`:
# https://www.postgresql.org/docs/12/catalog-pg-class.html
# r = ordinary table, i = index, S = sequence, t = TOAST table, v = view, m =
# materialized view, c = composite type, f = foreign table, p = partitioned
# table, I = partitioned index
# The `nspname` filter matches sqlalchemy's `get_schema_names` method.
# TODO?: skip `information_schema` tables here already?
PG_LIST_SOURCES_ALL_SCHEMAS_SQL = """
SELECT
    pg_namespace.nspname as schema,
    pg_class.relname as name
FROM
    pg_class
    JOIN pg_namespace
    ON pg_namespace.oid = pg_class.relnamespace
WHERE
    pg_namespace.nspname not like 'pg_%'
    AND pg_namespace.nspname != 'information_schema'
    AND pg_class.relkind in ('m', 'p', 'r', 'v')
ORDER BY schema, name
""".strip().replace(
    "\n", "  "
)
# NOTE: there's also `AND NOT pg_class.relispartition` in postgresql>=10


@attr.s(cmp=False)
class BasePostgresAdapter(BaseSSLCertAdapter):
    conn_type = CONNECTION_TYPE_POSTGRES

    # The cursor description returns type `oid` values.
    # These values can be found in `pg_type.oid` system table,
    # and mapped to `pg_type.typname`.
    # They can also be found in `psycopg2.extensions.string_types` mapping
    # (which maps oid to a string-to-python-value conversion function).
    # It is potentially possible to plug into the sqlalchemy logic
    # (with a few extra raw sql queries),
    # but for now, a more simple mapping is preferred.
    _type_code_to_sa: Optional[Dict[Any, SATypeSpec]] = {
        16: sa_pg.BOOLEAN,
        17: sa_pg.BYTEA,  # unsupported
        # 19: 'name',
        20: sa_pg.BIGINT,
        21: sa_pg.SMALLINT,
        23: sa_pg.INTEGER,
        25: sa_pg.TEXT,
        26: sa_pg.OID,  # unsupported
        114: sa_pg.JSON,  # unsupported
        700: sa_pg.REAL,
        701: sa_pg.DOUBLE_PRECISION,
        1005: sa_pg.ARRAY(sa_pg.SMALLINT),
        1007: sa_pg.ARRAY(sa_pg.INTEGER),
        1014: sa_pg.ARRAY(sa_pg.CHAR),
        1015: sa_pg.ARRAY(sa_pg.VARCHAR),
        1016: sa_pg.ARRAY(sa_pg.BIGINT),
        1021: sa_pg.ARRAY(sa_pg.REAL),
        1022: sa_pg.ARRAY(sa_pg.DOUBLE_PRECISION),
        1025: sa_pg.ARRAY(sa_pg.TEXT),
        1031: sa_pg.ARRAY(sa_pg.NUMERIC),
        1042: sa_pg.CHAR,
        1043: sa_pg.VARCHAR,
        1082: sa_pg.DATE,
        1083: sa_pg.TIME,  # unsupported
        1114: sa_pg.TIMESTAMP,
        1184: sa_pg.TIMESTAMP(timezone=True),
        1186: sa_pg.INTERVAL,  # unsupported
        1700: sa_pg.NUMERIC,  # untested
        2950: sa_pg.UUID,  # untested
        11020: sa_pg.ENUM,  # untested
    }

    @staticmethod
    def _convert_bytea(value: memoryview) -> Optional[str]:
        str_value = bytes(value[:110]).decode("utf-8", errors="replace")
        if len(value) > 110 or len(str_value) > 100:
            str_value = str_value[:100]
            str_value = str_value + "…"
        return str_value

    # it's static because of typing hell in target_dto
    @staticmethod
    def _get_enforce_collate(target_dto: PostgresConnTargetDTO) -> Optional[str]:
        if target_dto.enforce_collate == PGEnforceCollateMode.auto:
            raise Exception("Resolution of PGEnforceCollateMode.auto is not allowed at this point")
        if target_dto.enforce_collate == PGEnforceCollateMode.on:
            return "en_US"
        return None

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/adapters_base_postgres.py
datalens-tech/datalens-backend
from __future__ import annotations

import contextlib
import typing
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Tuple,
)

import attr
import sqlalchemy as sa

from dl_core.connection_executors.adapters.adapters_base_sa_classic import BaseClassicAdapter
from dl_core.connection_models.common_models import TableIdent

from dl_connector_postgresql.core.postgresql_base.adapters_base_postgres import (
    OID_KNOWLEDGE,
    PG_LIST_SOURCES_ALL_SCHEMAS_SQL,
    BasePostgresAdapter,
)
from dl_connector_postgresql.core.postgresql_base.error_transformer import sync_pg_db_error_transformer
from dl_connector_postgresql.core.postgresql_base.target_dto import PostgresConnTargetDTO


if TYPE_CHECKING:
    from dl_core.connection_executors.models.db_adapter_data import ExecutionStepCursorInfo
    from dl_core.connection_models.common_models import SchemaIdent


@attr.s()
class PostgresAdapter(BasePostgresAdapter, BaseClassicAdapter[PostgresConnTargetDTO]):
    _error_transformer = sync_pg_db_error_transformer

    execution_options = {
        "stream_results": True,
    }

    def get_connect_args(self) -> dict:
        return {
            "sslmode": "require" if self._target_dto.ssl_enable else "prefer",
            "sslrootcert": self.get_ssl_cert_path(self._target_dto.ssl_ca),
        }

    def get_engine_kwargs(self) -> dict:
        result: Dict = {}
        enforce_collate = self._get_enforce_collate(self._target_dto)
        if enforce_collate:
            result.update(enforce_collate=enforce_collate)
        return result

    @contextlib.contextmanager
    def execution_context(self) -> typing.Generator[None, None, None]:
        contexts: list[typing.ContextManager[None]] = [super().execution_context()]

        if self._target_dto.ssl_ca is not None:
            contexts.append(self.ssl_cert_context(self._target_dto.ssl_ca))

        with contextlib.ExitStack() as stack:
            for context in contexts:
                stack.enter_context(context)
            try:
                yield
            finally:
                stack.close()

    def _get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        if schema_ident.schema_name is not None:
            # For a single schema, plug into the common SA code.
            # (might not be ever used)
            return super()._get_tables(schema_ident)

        assert schema_ident.schema_name is None
        db_name = schema_ident.db_name
        db_engine = self.get_db_engine(db_name)
        query = PG_LIST_SOURCES_ALL_SCHEMAS_SQL
        result = db_engine.execute(sa.text(query))
        return [
            TableIdent(
                db_name=db_name,
                schema_name=schema_name,
                table_name=name,
            )
            for schema_name, name in result
        ]

    def _make_cursor_info(self, cursor, db_session=None) -> dict:  # type: ignore  # TODO: fix
        return dict(
            super()._make_cursor_info(cursor, db_session=db_session),
            # Deprecating:
            columns=[
                dict(
                    name=str(column[0]),
                    postgresql_oid=column[1],
                    postgresql_typname=OID_KNOWLEDGE.get(column[1]),  # type: ignore  # TODO: fix
                )
                for column in cursor.description
            ],
            # dashsql convenience:
            postgresql_typnames=[OID_KNOWLEDGE.get(column[1]) for column in cursor.description],  # type: ignore  # TODO: fix
        )

    def _get_row_converters(self, cursor_info: ExecutionStepCursorInfo) -> Tuple[Optional[Callable[[Any], Any]], ...]:
        return tuple(
            self._convert_bytea if col.type_code == 17 else None  # `bytea`
            for col in cursor_info.raw_cursor_description
        )

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/adapters_postgres.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import contextlib
from contextlib import asynccontextmanager
from datetime import datetime
import logging
import typing
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncIterator,
    Callable,
    ClassVar,
    Iterable,
    Optional,
    Type,
    TypeVar,
)
from urllib.parse import quote_plus

import asyncpg
import asyncpg.exceptions
import attr
from dateutil import parser as dateutil_parser
import sqlalchemy as sa

from dl_app_tools.profiling_base import generic_profiler_async
from dl_constants.types import (
    TBIChunksGen,
    TBIDataRow,
)
from dl_core import exc
from dl_core.connection_executors.adapters.adapters_base_sa_classic import ClassicSQLConnLineConstructor
from dl_core.connection_executors.adapters.async_adapters_base import (
    AsyncCache,
    AsyncDirectDBAdapter,
    AsyncRawExecutionResult,
)
from dl_core.connection_executors.adapters.mixins import (
    SATypeTransformer,
    WithAsyncGetDBVersion,
    WithDatabaseNameOverride,
)
from dl_core.connection_executors.adapters.sa_utils import make_debug_query
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    ExecutionStep,
    ExecutionStepCursorInfo,
    ExecutionStepDataChunk,
    RawSchemaInfo,
)
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
from dl_core.connection_models import (
    DBIdent,
    TableDefinition,
)
from dl_core.connection_models.common_models import TableIdent
from dl_core.connectors.base.error_handling import ETBasedExceptionMaker
from dl_sqlalchemy_postgres import AsyncBIPGDialect
from dl_sqlalchemy_postgres.asyncpg import DBAPIMock

from dl_connector_postgresql.core.postgresql_base.adapters_base_postgres import (
    OID_KNOWLEDGE,
    PG_LIST_SOURCES_ALL_SCHEMAS_SQL,
    BasePostgresAdapter,
)
from dl_connector_postgresql.core.postgresql_base.error_transformer import make_async_pg_error_transformer
from dl_connector_postgresql.core.postgresql_base.target_dto import PostgresConnTargetDTO
from dl_connector_postgresql.core.postgresql_base.utils import compile_pg_query


if TYPE_CHECKING:
    from dl_core.connection_models.common_models import SchemaIdent

_DBA_ASYNC_POSTGRES_TV = TypeVar("_DBA_ASYNC_POSTGRES_TV", bound="AsyncPostgresAdapter")

LOGGER = logging.getLogger(__name__)


PG_LIST_SCHEMA_NAMES = """
SELECT nspname FROM pg_namespace
WHERE nspname NOT LIKE 'pg_%'
ORDER BY nspname
"""

PG_LIST_TABLE_NAMES = """
SELECT c.relname FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = :schema AND c.relkind in ('r', 'p')
"""

# https://github.com/sqlalchemy/sqlalchemy/blob/rel_1_4/lib/sqlalchemy/dialects/postgresql/base.py#L3802
PG_LIST_VIEW_NAMES = """
SELECT c.relname FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = :schema AND c.relkind IN 'v', 'm'
"""


# native SA asyncpg dialect is beta now
# we can use it after release
# look at https://docs.sqlalchemy.org/en/14/orm/extensions/asyncio.html
# right now we have to implement some logic in our code
@attr.s(cmp=False, kw_only=True)
class AsyncPostgresAdapter(
    WithAsyncGetDBVersion,
    WithDatabaseNameOverride,
    AsyncDirectDBAdapter,
    BasePostgresAdapter,
    SATypeTransformer,
    ETBasedExceptionMaker,
):
    dsn_template: ClassVar[str] = "{dialect}://{user}:{passwd}@{host}:{port}/{db_name}"
    _default_chunk_size: int = attr.ib()
    _target_dto: PostgresConnTargetDTO = attr.ib()
    _req_ctx_info: DBAdapterScopedRCI = attr.ib()
    _conn_pools: AsyncCache[asyncpg.Pool] = attr.ib(default=attr.Factory(AsyncCache))

    _error_transformer = make_async_pg_error_transformer()
    __dialect: Optional[AsyncBIPGDialect] = None

    EXTRA_EXC_CLS: ClassVar[tuple[Type[Exception], ...]] = (
        exc.DataStreamValidationError,
        # I don't have any idea why asyncpg doesn't provide common exception in public interface
        # but there are about 40 different types of exceptions for every case of error
        # I hope private type will work for a long time and we won't put all ~40 exceptions here
        asyncpg.exceptions._base.PostgresError,
        asyncio.TimeoutError,
        # invalid/not available cluster provokes it
        # not connection/timeout error
        OSError,
    )

    @property
    def _dialect(self) -> AsyncBIPGDialect:
        if self.__dialect is not None:
            return self.__dialect
        enforce_collate = self._get_enforce_collate(self._target_dto)
        if enforce_collate is not None:
            self.__dialect = AsyncBIPGDialect(enforce_collate=enforce_collate)
        else:
            self.__dialect = AsyncBIPGDialect()
        return self.__dialect

    def get_conn_line(self, db_name: Optional[str] = None, params: dict[str, Any] = None) -> str:
        params = params or {}
        params["sslrootcert"] = self.get_ssl_cert_path(self._target_dto.ssl_ca)
        return AsyncPGConnLineConstructor(
            dsn_template=self.dsn_template,
            target_dto=self._target_dto,
            dialect_name="postgres",
        ).make_conn_line(db_name=db_name, params=params)

    def get_default_db_name(self) -> Optional[str]:
        return self._target_dto.db_name

    def _convert_date(self, s: str, ignoretz: bool) -> datetime:
        try:
            d = dateutil_parser.parse(s, ignoretz=ignoretz)
        except (dateutil_parser.ParserError, OverflowError) as e:
            LOGGER.info(f"Can't parse date {s} by {ignoretz}")
            # its impossible to get extra info about the position in stream
            # because it's a callback for asyncpg
            raise exc.DataStreamValidationError(value=s) from e
        return d

    async def create_conn_pool(self, db_name_from_query: str) -> asyncpg.Pool:
        db_name = self.get_db_name_for_query(db_name_from_query)
        conn_line = self.get_conn_line(db_name=db_name)
        return await asyncpg.create_pool(conn_line, statement_cache_size=0)

    @asynccontextmanager
    async def _get_connection(self, db_name_from_query: str) -> AsyncIterator[asyncpg.Connection]:
        conn_pool = await self._conn_pools.get(db_name_from_query, generator=self.create_conn_pool)
        async with conn_pool.acquire() as connection:
            # Fix for the "asyncpg.exceptions.InvalidSQLStatementNameError: unnamed prepared statement does not exist"
            # We have already disabled cache by statement_cache_size=0, but it still works,
            # so let's try to use cache inside a transaction.
            async with connection.transaction():
                # There is some decimal-magic in asyncpg
                # and this magic is incompatible with magic in sqlalchemy-psycopg2
                # so lets disable it
                await connection.set_type_codec(
                    "numeric",
                    encoder=str,
                    decoder=lambda x: x,
                    schema="pg_catalog",
                    format="text",
                )
                # we set date-like values to params as strings
                # but asyncpg expects python-objects
                await connection.set_type_codec(
                    "date",
                    encoder=str,
                    decoder=lambda x: self._convert_date(x, ignoretz=True).date(),
                    schema="pg_catalog",
                    format="text",
                )
                await connection.set_type_codec(
                    "timestamp",
                    encoder=str,
                    decoder=lambda x: self._convert_date(x, ignoretz=True),
                    schema="pg_catalog",
                    format="text",
                )
                await connection.set_type_codec(
                    "timestamptz",
                    encoder=str,
                    decoder=lambda x: self._convert_date(x, ignoretz=False),
                    schema="pg_catalog",
                    format="text",
                )
            yield connection

    async def test(self) -> None:
        await self.execute(DBAdapterQuery("select 1"))

    def _make_cursor_info(self, query_attrs: Iterable[asyncpg.Attribute]) -> dict:
        return dict(
            names=[str(a.name) for a in query_attrs],
            driver_types=[self._cursor_type_to_str(a.type.oid) for a in query_attrs],
            db_types=[
                self._cursor_column_to_native_type(
                    (
                        a.name,
                        a.type.oid,
                    ),
                    require=False,
                )
                for a in query_attrs
            ],
            columns=[
                dict(
                    name=str(a.name),
                    postgresql_oid=a.type.oid,
                    postgresql_typname=OID_KNOWLEDGE.get(a.type.oid),  # type: ignore  # TODO: fix
                )
                for a in query_attrs
            ],
            # dashsql convenience:
            postgresql_typnames=[OID_KNOWLEDGE.get(a.type.oid) for a in query_attrs],
        )

    def _get_row_converters(
        self, query_attrs: Iterable[asyncpg.Attribute]
    ) -> tuple[Optional[Callable[[Any], Any]], ...]:
        return tuple(self._convert_bytea if a.type.oid == 17 else None for a in query_attrs)  # `bytea`

    @contextlib.contextmanager
    def execution_context(self) -> typing.Generator[None, None, None]:
        contexts = []

        if self._target_dto.ssl_ca is not None:
            contexts.append(self.ssl_cert_context(self._target_dto.ssl_ca))

        with contextlib.ExitStack() as stack:
            for context in contexts:
                stack.enter_context(context)
            yield

    async def _execute_by_step(self, query: DBAdapterQuery) -> AsyncIterator[ExecutionStep]:
        def make_record(raw_rec: asyncpg.Record, query_attrs: Iterable[asyncpg.Attribute]) -> TBIDataRow:
            row_converters = self._get_row_converters(query_attrs=query_attrs)
            row = tuple(raw_rec.values())
            return tuple(
                (col_converter(val) if col_converter is not None and val is not None else val)
                for val, col_converter in zip(row, row_converters)
            )

        q = query.query
        # exclude ENUM because of
        # https://github.com/sqlalchemy/sqlalchemy/blob/rel_1_4/lib/sqlalchemy/dialects/postgresql/asyncpg.py#L345
        # and exclude STRINGS because of our BI ENUMS (it's strings)
        # in asyncpg we can skip type annotations for strings, it should work like in psycopg
        compiled_query, params = compile_pg_query(q, self._dialect, exclude_types={DBAPIMock.ENUM, DBAPIMock.STRING})
        debug_query = None
        if self._target_dto.pass_db_query_to_user:
            debug_query = query.debug_compiled_query or make_debug_query(compiled_query, params)

        with self.handle_execution_error(debug_query), self.execution_context():
            async with self._get_connection(query.db_name) as conn:
                # prepare works only inside a transaction
                async with conn.transaction():
                    prepared_query = await conn.prepare(compiled_query)
                    cursor_info = self._make_cursor_info(prepared_query.get_attributes())
                    yield ExecutionStepCursorInfo(cursor_info=cursor_info)
                    cursor = await prepared_query.cursor(*params)
                    is_enough = False
                    while not is_enough:
                        result = await cursor.fetch(self._default_chunk_size)
                        if not result:
                            is_enough = True
                        else:
                            chunk = tuple(make_record(record, prepared_query.get_attributes()) for record in result)
                            yield ExecutionStepDataChunk(chunk=chunk)

    @generic_profiler_async("db-full")  # type: ignore  # TODO: fix
    async def execute(self, query: DBAdapterQuery) -> AsyncRawExecutionResult:
        LOGGER.info("Run by async postgres adapter")

        async def _process_chunk(steps: AsyncIterator[ExecutionStep]) -> TBIChunksGen:
            async for step in steps:
                assert isinstance(step, ExecutionStepDataChunk)
                yield step.chunk

        steps = self._execute_by_step(query)
        first_step = await steps.__anext__()
        assert isinstance(first_step, ExecutionStepCursorInfo)
        cursor_info = first_step.cursor_info
        return AsyncRawExecutionResult(
            raw_cursor_info=cursor_info,
            raw_chunk_generator=_process_chunk(steps),
        )

    async def get_schema_names(self, db_ident: DBIdent) -> list[str]:
        result = await self.execute(DBAdapterQuery(PG_LIST_SCHEMA_NAMES))
        schema_names = []
        async for row in result.get_all_rows():
            for value in row:
                schema_names.append(str(value))
        return schema_names

    async def _get_view_names(self, schema_ident: SchemaIdent) -> list[TableIdent]:
        query = sa.text(PG_LIST_VIEW_NAMES).bindparams(
            sa.bindparam(
                "schema",
                schema_ident.schema_name,
                type=sa.types.Unicode,
            )
        )
        result = await self.execute(DBAdapterQuery(query))
        views = []
        async for row in result.get_all_rows():
            views.append(str(row[0]))
        return [
            TableIdent(
                db_name=schema_ident.db_name,
                schema_name=schema_ident.schema_name,
                table_name=view,
            )
            for view in views
        ]

    async def _get_table_names(self, schema_ident: SchemaIdent) -> list[TableIdent]:
        query = sa.text(PG_LIST_TABLE_NAMES).bindparams(
            sa.bindparam(
                "schema",
                schema_ident.schema_name,
                type=sa.types.Unicode,
            )
        )
        result = await self.execute(DBAdapterQuery(query))
        tables = []
        async for row in result.get_all_rows():
            tables.append(str(row[0]))
        return [
            TableIdent(
                db_name=schema_ident.db_name,
                schema_name=schema_ident.schema_name,
                table_name=table,
            )
            for table in tables
        ]

    async def _get_tables_single_schema(self, schema_ident: SchemaIdent) -> list[TableIdent]:
        table_list = await self._get_table_names(schema_ident)
        view_list = await self._get_view_names(schema_ident)
        return table_list + view_list

    async def get_tables(self, schema_ident: SchemaIdent) -> list[TableIdent]:
        if schema_ident.schema_name is not None:
            # For a single schema, plug into the common SA code.
            # (might not be ever used)
            return await self._get_tables_single_schema(schema_ident)

        assert schema_ident.schema_name is None
        db_name = schema_ident.db_name
        result = await self.execute(DBAdapterQuery(PG_LIST_SOURCES_ALL_SCHEMAS_SQL))
        return [
            TableIdent(
                db_name=db_name,
                schema_name=str(schema_name),
                table_name=str(name),
            )
            async for schema_name, name in result.get_all_rows()
        ]

    async def get_table_info(self, table_def: TableDefinition, fetch_idx_info: bool) -> RawSchemaInfo:
        raise NotImplementedError()

    # magic copypaste from
    # https://github.com/sqlalchemy/sqlalchemy/blob/rel_1_4/lib/sqlalchemy/dialects/postgresql/base.py#L3589
    async def is_table_exists(self, table_ident: TableIdent) -> bool:
        if table_ident.schema_name is None:
            result = await self.execute(
                DBAdapterQuery(
                    sa.text(
                        "select relname from pg_class c join pg_namespace n on "
                        "n.oid=c.relnamespace where "
                        "pg_catalog.pg_table_is_visible(c.oid) "
                        "and relname=:name"
                    ).bindparams(
                        sa.bindparam(
                            "name",
                            table_ident.table_name,
                            type_=sa.types.Unicode,
                        )
                    )
                )
            )
        else:
            result = await self.execute(
                DBAdapterQuery(
                    sa.text(
                        "select relname from pg_class c join pg_namespace n on "
                        "n.oid=c.relnamespace where n.nspname=:schema and "
                        "relname=:name"
                    ).bindparams(
                        sa.bindparam(
                            "name",
                            table_ident.table_name,
                            type_=sa.types.Unicode,
                        ),
                        sa.bindparam(
                            "schema",
                            table_ident.schema_name,
                            type_=sa.types.Unicode,
                        ),
                    )
                )
            )
        async for _ in result.get_all_rows():
            return True
        return False

    async def close(self) -> None:
        await self._conn_pools.clear(finalizer=lambda pool: asyncio.wait_for(pool.close(), timeout=2.5))

    @property
    def default_chunk_size(self) -> int:
        return self._default_chunk_size

    @classmethod
    def create(
        cls: Type[_DBA_ASYNC_POSTGRES_TV],
        target_dto: PostgresConnTargetDTO,
        req_ctx_info: DBAdapterScopedRCI,
        default_chunk_size: int,
    ) -> _DBA_ASYNC_POSTGRES_TV:
        return cls(target_dto=target_dto, default_chunk_size=default_chunk_size, req_ctx_info=req_ctx_info)


class AsyncPGConnLineConstructor(ClassicSQLConnLineConstructor[PostgresConnTargetDTO]):
    def _get_dsn_query_params(self) -> dict:
        return {
            "sslmode": "require" if self._target_dto.ssl_enable else "prefer",
        }

    def _get_dsn_params(
        self,
        safe_db_symbols: tuple[str, ...] = (),
        db_name: Optional[str] = None,
        standard_auth: Optional[bool] = True,
    ) -> dict:
        return dict(
            dialect=self._dialect_name,
            user=quote_plus(self._target_dto.username) if standard_auth else None,
            passwd=quote_plus(self._target_dto.password) if standard_auth else None,
            host=quote_plus(self._target_dto.host),
            port=quote_plus(str(self._target_dto.port)),
            db_name=db_name or quote_plus(self._target_dto.db_name or "", safe="".join(safe_db_symbols)),
        )

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/async_adapters_postgres.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    Sequence,
    TypeVar,
)

import attr

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor

from dl_connector_postgresql.core.postgresql_base.adapters_postgres import PostgresAdapter
from dl_connector_postgresql.core.postgresql_base.async_adapters_postgres import AsyncPostgresAdapter
from dl_connector_postgresql.core.postgresql_base.constants import PGEnforceCollateMode
from dl_connector_postgresql.core.postgresql_base.dto import PostgresConnDTOBase
from dl_connector_postgresql.core.postgresql_base.target_dto import PostgresConnTargetDTO


_BASE_POSTGRES_ADAPTER_TV = TypeVar("_BASE_POSTGRES_ADAPTER_TV", bound=CommonBaseDirectAdapter)


@attr.s(cmp=False, hash=False)
class BasePostgresConnExecutor(DefaultSqlAlchemyConnExecutor[_BASE_POSTGRES_ADAPTER_TV]):
    _conn_dto: PostgresConnDTOBase = attr.ib()

    def _get_effective_enforce_collate(
        self,
        enforce_collate: PGEnforceCollateMode,
        multihosts: Sequence[str],
    ) -> PGEnforceCollateMode:
        if enforce_collate == PGEnforceCollateMode.auto:
            enforce_collate = PGEnforceCollateMode.off
        return enforce_collate

    async def _make_target_conn_dto_pool(self) -> list[PostgresConnTargetDTO]:  # type: ignore  # TODO: fix
        dto_pool = []
        effective_enforce_collate = self._get_effective_enforce_collate(
            enforce_collate=self._conn_dto.enforce_collate,
            multihosts=self._conn_dto.multihosts,
        )
        for host in self._conn_hosts_pool:
            dto_pool.append(
                PostgresConnTargetDTO(
                    conn_id=self._conn_dto.conn_id,
                    pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                    pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                    host=host,
                    port=self._conn_dto.port,
                    db_name=self._conn_dto.db_name,
                    username=self._conn_dto.username,
                    password=self._conn_dto.password,
                    enforce_collate=effective_enforce_collate,
                    ssl_enable=self._conn_dto.ssl_enable,
                    ssl_ca=self._conn_dto.ssl_ca,
                )
            )
        return dto_pool

    def mutate_for_dashsql(self, db_params: Optional[dict[str, str]] = None):  # type: ignore  # TODO: fix
        if db_params:
            # TODO: better exception class for HTTP 4xx response
            raise Exception("No db_params supported here at the moment")
        return self


@attr.s(cmp=False, hash=False)
class PostgresConnExecutor(BasePostgresConnExecutor[PostgresAdapter]):
    TARGET_ADAPTER_CLS = PostgresAdapter


@attr.s(cmp=False, hash=False)
class AsyncPostgresConnExecutor(BasePostgresConnExecutor[AsyncPostgresAdapter]):
    TARGET_ADAPTER_CLS = AsyncPostgresAdapter

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/connection_executors.py
datalens-tech/datalens-backend
from enum import (
    Enum,
    unique,
)


@unique
class PGEnforceCollateMode(Enum):
    auto = "auto"
    on = "on"
    off = "off"

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

import typing

import attr

from dl_core.connection_models.dto_defs import DefaultSQLDTO

from dl_connector_postgresql.core.postgresql_base.constants import PGEnforceCollateMode


@attr.s(frozen=True)
class PostgresConnDTOBase(DefaultSQLDTO):
    enforce_collate: PGEnforceCollateMode = attr.ib(kw_only=True, default=PGEnforceCollateMode.off)
    ssl_enable: bool = attr.ib(kw_only=True, default=False)
    ssl_ca: typing.Optional[str] = attr.ib(kw_only=True, default=None)

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/dto.py
datalens-tech/datalens-backend
import asyncpg.exceptions as asyncpg_exc
import psycopg2.errors

from dl_core.connectors.base.error_transformer import (
    ChainedDbErrorTransformer,
    DbErrorTransformer,
)
from dl_core.connectors.base.error_transformer import (
    make_rule_from_descr,
    orig_exc_is,
    wrapper_exc_is_and_matches_re,
)
from dl_core.connectors.base.error_transformer import ErrorTransformerRule as Rule
import dl_core.exc as exc

from dl_connector_postgresql.core.postgresql_base.exc import (
    PgDoublePrecisionRoundError,
    PostgresSourceDoesNotExistError,
)


sync_pg_db_error_transformer: DbErrorTransformer = ChainedDbErrorTransformer(
    [
        Rule(
            when=wrapper_exc_is_and_matches_re(
                wrapper_exc_cls=psycopg2.OperationalError, err_regex_str="Name or service not known"
            ),
            then_raise=exc.SourceHostNotKnownError,
        ),
        Rule(
            when=orig_exc_is(orig_exc_cls=psycopg2.errors.UndefinedTable),  # type: ignore
            then_raise=PostgresSourceDoesNotExistError,
        ),
    ]
)


def make_async_pg_error_transformer() -> DbErrorTransformer:
    rule_descriptions = (
        ((asyncpg_exc.DivisionByZeroError, None), exc.DivisionByZero),
        (
            (asyncpg_exc.UndefinedFunctionError, r"function round\(double precision"),
            PgDoublePrecisionRoundError,
        ),
        ((asyncpg_exc.UndefinedFunctionError, None), exc.UnknownFunction),
        ((OverflowError, "value out of .* range"), exc.NumberOutOfRange),
        ((asyncpg_exc.InvalidTextRepresentationError, None), exc.DataParseError),
        # ((asyncpg_exc.UndefinedTableError, None), exc.SourceDoesNotExist),
        ((asyncpg_exc.UndefinedTableError, None), PostgresSourceDoesNotExistError),
        # ((asyncpg_exc.FDWTableNotFoundError, None), exc.SourceDoesNotExist),
        ((asyncpg_exc.FDWTableNotFoundError, None), PostgresSourceDoesNotExistError),  # todo: test for this error
        ((asyncpg_exc.SyntaxOrAccessError, None), exc.DatabaseOperationalError),
        ((TimeoutError, None), exc.SourceConnectError),
        ((OSError, "Connect call failed"), exc.SourceConnectError),
        ((OSError, "Name or service not known"), exc.SourceHostNotKnownError),
    )
    return ChainedDbErrorTransformer([make_rule_from_descr(d) for d in rule_descriptions])

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/error_transformer.py
datalens-tech/datalens-backend
import re
from typing import (
    Any,
    Dict,
    Optional,
)

from dl_core import exc as exc


class PostgresSourceDoesNotExistError(exc.SourceDoesNotExist):
    ERR_RE = re.compile(r".*relation\s\"(?P<table>.*)\"\sdoes not exist.*")

    def __init__(
        self,
        db_message: Optional[str] = None,
        query: Optional[str] = None,
        message: Optional[str] = None,
        details: Optional[Dict[str, Any]] = None,
        orig: Optional[Exception] = None,
        debug_info: Optional[Dict[str, Any]] = None,
        params: Optional[Dict[str, Any]] = None,
    ):
        super(PostgresSourceDoesNotExistError, self).__init__(
            db_message=db_message,
            query=query,
            message=message,
            details=details,
            orig=orig,
            debug_info=debug_info,
            params=params,
        )
        if self.db_message:
            message = self.db_message.replace("\n", "")
            if message and (match := self.ERR_RE.match(message)):
                if table := match.group("table"):
                    self.params["table_definition"] = table


class PgDoublePrecisionRoundError(exc.UnknownFunction):
    err_code = exc.DatabaseQueryError.err_code + ["PG_DOUBLE_PRECISION_ROUND"]
    default_message = (
        "ROUND with precision parameter is not supported for `double precision` data type " "in PostgreSQL data source."
    )

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/exc.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import ClassVar

from dl_core.connectors.base.query_compiler import (
    QueryCompiler,
    SectionAliasMode,
)


class PostgreSQLQueryCompiler(QueryCompiler):
    force_nulls_lower_than_values = True

    groupby_alias_mode: ClassVar[SectionAliasMode] = SectionAliasMode.by_alias_in_section

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/query_compiler.py
datalens-tech/datalens-backend
from __future__ import annotations

from sqlalchemy.dialects import postgresql as pg_types

from dl_core.db.sa_types_base import (
    make_native_type,
    simple_instantiator,
    timezone_instantiator,
    typed_instantiator,
)

from dl_connector_postgresql.core.postgresql.constants import CONNECTION_TYPE_POSTGRES
from dl_connector_postgresql.core.postgresql_base.type_transformer import (
    PG_TYPES_FLOAT,
    PG_TYPES_INT,
    PG_TYPES_STRING,
)


SQLALCHEMY_POSTGRES_BASE_TYPES = (
    *PG_TYPES_INT,
    *PG_TYPES_FLOAT,
    *PG_TYPES_STRING,
    pg_types.BOOLEAN,
    pg_types.DATE,
    pg_types.UUID,
)
SQLALCHEMY_POSTGRES_TIMEZONE_TYPES = (pg_types.TIMESTAMP,)
SQLALCHEMY_POSTGRES_ARRAY_INNER_TYPES = (
    *PG_TYPES_INT,
    *PG_TYPES_FLOAT,
    *PG_TYPES_STRING,
)
SQLALCHEMY_POSTGRES_TYPES = {
    **{
        make_native_type(CONNECTION_TYPE_POSTGRES, typecls): simple_instantiator(typecls)
        for typecls in SQLALCHEMY_POSTGRES_BASE_TYPES
    },
    **{
        make_native_type(CONNECTION_TYPE_POSTGRES, typecls): timezone_instantiator(typecls)
        for typecls in SQLALCHEMY_POSTGRES_TIMEZONE_TYPES
    },
    **{
        make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.ARRAY(inner_typecls)): typed_instantiator(
            pg_types.ARRAY, inner_typecls
        )
        for inner_typecls in SQLALCHEMY_POSTGRES_ARRAY_INNER_TYPES
    },
}

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/sa_types.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_core.us_manager.storage_schemas.connection import ConnectionSQLDataStorageSchema

from dl_connector_postgresql.core.postgresql_base.constants import PGEnforceCollateMode
from dl_connector_postgresql.core.postgresql_base.us_connection import ConnectionPostgreSQLBase


class ConnectionPostgreSQLBaseDataStorageSchema(
    ConnectionSQLDataStorageSchema[ConnectionPostgreSQLBase.DataModel],
):
    enforce_collate = ma_fields.Enum(
        PGEnforceCollateMode,
        required=False,
        allow_none=False,
        dump_default=PGEnforceCollateMode.auto,
        load_default=PGEnforceCollateMode.auto,
    )
    ssl_enable = ma_fields.Boolean(
        required=False,
        allow_none=False,
        dump_default=False,
        load_default=False,
    )
    ssl_ca = ma_fields.String(
        required=False,
        allow_none=True,
        dump_default=None,
        load_default=None,
    )

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/storage_schemas/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Optional,
    Type,
    TypeVar,
)

import attr

from dl_core.connection_executors.models.connection_target_dto_base import BaseSQLConnTargetDTO

from dl_connector_postgresql.core.postgresql_base.constants import PGEnforceCollateMode


if TYPE_CHECKING:
    from dl_constants.types import TJSONLike


_CT_DTO_TV = TypeVar("_CT_DTO_TV", bound="PostgresConnTargetDTO")


@attr.s(frozen=True)
class PostgresConnTargetDTO(BaseSQLConnTargetDTO):
    enforce_collate: PGEnforceCollateMode = attr.ib(kw_only=True, default=PGEnforceCollateMode.off)
    ssl_enable: bool = attr.ib(kw_only=True, default=False)
    ssl_ca: Optional[str] = attr.ib(kw_only=True, default=None)

    def to_jsonable_dict(self) -> dict[str, TJSONLike]:
        return {
            **super().to_jsonable_dict(),
            "enforce_collate": self.enforce_collate.name,
            "ssl_enable": self.ssl_enable,
            "ssl_ca": self.ssl_ca,
        }

    @classmethod
    def _from_jsonable_dict(cls: Type[_CT_DTO_TV], data: dict) -> _CT_DTO_TV:
        return cls(
            **{  # type: ignore  # TODO: fix
                **data,
                "enforce_collate": PGEnforceCollateMode[data["enforce_collate"]],
            }
        )

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/target_dto.py
datalens-tech/datalens-backend
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql as pg_types

from dl_constants.enums import UserDataType
from dl_core.db.conversion_base import (
    TypeTransformer,
    UTCDatetimeTypeCaster,
    UTCTimezoneDatetimeTypeCaster,
    make_native_type,
)
from dl_sqlalchemy_postgres.base import CITEXT

from dl_connector_postgresql.core.postgresql.constants import CONNECTION_TYPE_POSTGRES


PG_TYPES_INT = frozenset((pg_types.SMALLINT, pg_types.INTEGER, pg_types.BIGINT))
PG_TYPES_FLOAT = frozenset((pg_types.REAL, pg_types.DOUBLE_PRECISION, pg_types.NUMERIC))
PG_TYPES_STRING = frozenset((pg_types.CHAR, pg_types.VARCHAR, pg_types.TEXT, CITEXT))


class PostgreSQLTypeTransformer(TypeTransformer):
    conn_type = CONNECTION_TYPE_POSTGRES
    casters = {
        **TypeTransformer.casters,  # type: ignore  # TODO: fix
        # Preliminary asyncpg-related hack: before inserting, make all datetimes UTC-naive.
        # A correct fix would require different BI-types for naive/aware datetimes.
        UserDataType.datetime: UTCDatetimeTypeCaster(),
        UserDataType.genericdatetime: UTCTimezoneDatetimeTypeCaster(),
    }
    native_to_user_map = {
        **{make_native_type(CONNECTION_TYPE_POSTGRES, t): UserDataType.integer for t in PG_TYPES_INT},  # type: ignore  # TODO: fix
        **{make_native_type(CONNECTION_TYPE_POSTGRES, t): UserDataType.float for t in PG_TYPES_FLOAT},
        make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.BOOLEAN): UserDataType.boolean,
        **{make_native_type(CONNECTION_TYPE_POSTGRES, t): UserDataType.string for t in PG_TYPES_STRING},
        make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.DATE): UserDataType.date,
        make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.TIMESTAMP): UserDataType.genericdatetime,
        make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.UUID): UserDataType.uuid,
        make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.ENUM): UserDataType.string,
        **{
            make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.ARRAY(typecls)): UserDataType.array_int
            for typecls in PG_TYPES_INT
        },
        **{
            make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.ARRAY(typecls)): UserDataType.array_float
            for typecls in PG_TYPES_FLOAT
        },
        **{
            make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.ARRAY(typecls)): UserDataType.array_str
            for typecls in PG_TYPES_STRING
        },
        make_native_type(CONNECTION_TYPE_POSTGRES, sa.sql.sqltypes.NullType): UserDataType.unsupported,
    }
    user_to_native_map = {
        UserDataType.integer: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.BIGINT),
        UserDataType.float: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.DOUBLE_PRECISION),
        UserDataType.boolean: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.BOOLEAN),
        UserDataType.string: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.TEXT),
        UserDataType.date: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.DATE),
        UserDataType.datetime: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.TIMESTAMP),
        UserDataType.genericdatetime: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.TIMESTAMP),
        UserDataType.geopoint: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.TEXT),
        UserDataType.geopolygon: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.TEXT),
        UserDataType.uuid: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.UUID),
        UserDataType.markup: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.TEXT),
        UserDataType.array_int: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.ARRAY(pg_types.BIGINT)),
        UserDataType.array_float: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.ARRAY(pg_types.DOUBLE_PRECISION)),
        UserDataType.array_str: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.ARRAY(pg_types.TEXT)),
        UserDataType.tree_str: make_native_type(CONNECTION_TYPE_POSTGRES, pg_types.ARRAY(pg_types.TEXT)),
        UserDataType.unsupported: make_native_type(CONNECTION_TYPE_POSTGRES, sa.sql.sqltypes.NullType),
    }

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Callable,
    Optional,
)

import attr

from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.us_connection_base import (
    ClassicConnectionSQL,
    ConnectionBase,
)

from dl_connector_postgresql.core.postgresql_base.constants import PGEnforceCollateMode


class ConnectionPostgreSQLBase(ClassicConnectionSQL):
    has_schema = True
    default_schema_name = "public"

    @attr.s(kw_only=True)
    class DataModel(ClassicConnectionSQL.DataModel):
        enforce_collate: PGEnforceCollateMode = attr.ib(default=PGEnforceCollateMode.auto)
        ssl_enable: bool = attr.ib(kw_only=True, default=False)
        ssl_ca: Optional[str] = attr.ib(kw_only=True, default=None)

    def get_parameter_combinations(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[dict]:
        if not self.db_name:
            return []

        assert self.has_schema
        return [
            dict(schema_name=tid.schema_name, table_name=tid.table_name)
            for tid in self.get_tables(conn_executor_factory=conn_executor_factory, schema_name=None)
        ]

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/us_connection.py
datalens-tech/datalens-backend
from typing import (
    AbstractSet,
    Any,
    Union,
)

import sqlalchemy as sa
from sqlalchemy import sql as sasql
from sqlalchemy.engine import Dialect


def compile_pg_query(
    query: Union[str, sasql.Select],
    dialect: Dialect,
    exclude_types: AbstractSet[Any] = frozenset(),
    add_types: bool = True,
) -> tuple[str, list]:
    """
    Compile query
    """
    if isinstance(query, str):
        return query, []

    if isinstance(query, sa.sql.ddl.DDLElement):
        compiled = query.compile(dialect=dialect)
        return compiled.string, []

    if isinstance(query, sa.sql.ClauseElement):
        if isinstance(query, (sa.sql.dml.Insert, sa.sql.dml.Update)):
            # query = execute_defaults(query)  # default values for Insert/Update
            raise Exception("Not supported here: insert / update SA queries. Use `copy_records_to_table`.")

        compiled = query.compile(dialect=dialect, compile_kwargs={"render_postcompile": True})
        compiled_params = sorted(compiled.params.items())
        if add_types and hasattr(dialect, "dbapi"):
            pg_types = dialect.dbapi.pg_types  # type: ignore  # TODO: fix
            input_sizes = compiled._get_set_input_sizes_lookup(exclude_types=exclude_types)
        else:
            pg_types = {}
            input_sizes = {}
        items = sorted([(key, input_sizes.get(bindparam)) for bindparam, key in compiled.bind_names.items()])
        mapping = {
            key: "$%d::%s" % (idx, typ) if typ else "$%d" % idx
            for idx, (key, typ) in enumerate(
                ((key, pg_types.get(typ)) for key, typ in items),
                1,
            )
        }
        new_query: str = compiled.string % mapping

        processors = compiled._bind_processors
        new_params = [processors[key](val) if key in processors else val for key, val in compiled_params]

        return new_query, new_params

    raise Exception("Unknown query object type", type(query))

lib/dl_connector_postgresql/dl_connector_postgresql/core/postgresql_base/utils.py
datalens-tech/datalens-backend
from dl_db_testing.connectors.base.connector import DbTestingConnector

from dl_connector_postgresql.db_testing.engine_wrapper import (
    BiPGEngineWrapper,
    PGEngineWrapper,
)


class PostgreSQLDbTestingConnector(DbTestingConnector):
    engine_wrapper_classes = (
        PGEngineWrapper,
        BiPGEngineWrapper,
    )

lib/dl_connector_postgresql/dl_connector_postgresql/db_testing/connector.py
datalens-tech/datalens-backend
from dl_db_testing.database.engine_wrapper import EngineWrapperBase


class PGEngineWrapperBase(EngineWrapperBase):
    def count_sql_sessions(self) -> int:
        # noinspection SqlDialectInspection
        cur = self.execute("select * from pg_stat_activity")
        try:
            lines = cur.fetchall()
            return len(lines)
        finally:
            cur.close()


class PGEngineWrapper(PGEngineWrapperBase):
    URL_PREFIX = "postgresql"


class BiPGEngineWrapper(PGEngineWrapperBase):
    URL_PREFIX = "bi_postgresql"

lib/dl_connector_postgresql/dl_connector_postgresql/db_testing/engine_wrapper.py
datalens-tech/datalens-backend
from sqlalchemy.dialects.postgresql.base import PGDialect

from dl_formula.connectors.base.connector import FormulaConnector

from dl_connector_postgresql.formula.constants import PostgreSQLDialect
from dl_connector_postgresql.formula.definitions.all import DEFINITIONS
from dl_connector_postgresql.formula.literal import GenericPostgreSQLLiteralizer
from dl_connector_postgresql.formula.type_constructor import PostgreSQLTypeConstructor


class PostgreSQLFormulaConnector(FormulaConnector):
    dialect_ns_cls = PostgreSQLDialect
    dialects = PostgreSQLDialect.POSTGRESQL
    default_dialect = PostgreSQLDialect.POSTGRESQL_9_4
    op_definitions = DEFINITIONS
    literalizer_cls = GenericPostgreSQLLiteralizer
    type_constructor_cls = PostgreSQLTypeConstructor
    sa_dialect = PGDialect()

lib/dl_connector_postgresql/dl_connector_postgresql/formula/connector.py
datalens-tech/datalens-backend
from dl_formula.core.dialect import (
    DialectName,
    DialectNamespace,
    simple_combo,
)


DIALECT_NAME_POSTGRESQL = DialectName.declare("POSTGRESQL")
DIALECT_NAME_COMPENG = DialectName.declare("COMPENG")  # computation engine (PostgreSQL-based)


class PostgreSQLDialect(DialectNamespace):
    COMPENG = simple_combo(name=DIALECT_NAME_COMPENG)
    POSTGRESQL_9_3 = simple_combo(name=DIALECT_NAME_POSTGRESQL, version=(9, 3))
    POSTGRESQL_9_4 = simple_combo(name=DIALECT_NAME_POSTGRESQL, version=(9, 4))
    NON_COMPENG_POSTGRESQL = POSTGRESQL_9_3 | POSTGRESQL_9_4
    POSTGRESQL = COMPENG | NON_COMPENG_POSTGRESQL

lib/dl_connector_postgresql/dl_connector_postgresql/formula/constants.py
datalens-tech/datalens-backend
from dl_connector_postgresql.formula.definitions.conditional_blocks import DEFINITIONS_COND_BLOCKS
from dl_connector_postgresql.formula.definitions.functions_aggregation import DEFINITIONS_AGG
from dl_connector_postgresql.formula.definitions.functions_array import DEFINITIONS_ARRAY
from dl_connector_postgresql.formula.definitions.functions_datetime import DEFINITIONS_DATETIME
from dl_connector_postgresql.formula.definitions.functions_logical import DEFINITIONS_LOGICAL
from dl_connector_postgresql.formula.definitions.functions_markup import DEFINITIONS_MARKUP
from dl_connector_postgresql.formula.definitions.functions_math import DEFINITIONS_MATH
from dl_connector_postgresql.formula.definitions.functions_special import DEFINITIONS_SPECIAL
from dl_connector_postgresql.formula.definitions.functions_string import DEFINITIONS_STRING
from dl_connector_postgresql.formula.definitions.functions_type import DEFINITIONS_TYPE
from dl_connector_postgresql.formula.definitions.functions_window import DEFINITIONS_WINDOW
from dl_connector_postgresql.formula.definitions.operators_binary import DEFINITIONS_BINARY
from dl_connector_postgresql.formula.definitions.operators_ternary import DEFINITIONS_TERNARY
from dl_connector_postgresql.formula.definitions.operators_unary import DEFINITIONS_UNARY


DEFINITIONS = [
    *DEFINITIONS_COND_BLOCKS,
    *DEFINITIONS_AGG,
    *DEFINITIONS_DATETIME,
    *DEFINITIONS_LOGICAL,
    *DEFINITIONS_MARKUP,
    *DEFINITIONS_MATH,
    *DEFINITIONS_SPECIAL,
    *DEFINITIONS_STRING,
    *DEFINITIONS_TYPE,
    *DEFINITIONS_WINDOW,
    *DEFINITIONS_UNARY,
    *DEFINITIONS_BINARY,
    *DEFINITIONS_TERNARY,
    *DEFINITIONS_ARRAY,
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/all.py
datalens-tech/datalens-backend
# FM: fill mode (suppress trailing zeroes and padding blanks)
# MI: minus sign in specified position (if number < 0)
# 20th "9" to cover 2**64
PG_INT_64_TO_CHAR_FMT = "FMMI99999999999999999999"

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/common.py
datalens-tech/datalens-backend
import dl_formula.definitions.conditional_blocks as base

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


DEFINITIONS_COND_BLOCKS = [
    # _case_block_
    base.CaseBlock.for_dialect(D.POSTGRESQL),
    # _if_block_
    base.IfBlock3.for_dialect(D.POSTGRESQL),
    base.IfBlockMulti.for_dialect(D.POSTGRESQL),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/conditional_blocks.py
datalens-tech/datalens-backend
import sqlalchemy as sa
from sqlalchemy.sql import ClauseElement

from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common import (
    quantile_value,
    within_group,
)
import dl_formula.definitions.functions_aggregation as base
from dl_formula.definitions.literals import un_literal

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


V = TranslationVariant.make


def _all_concat_pg(expr: ClauseElement, sep: str = ", ") -> ClauseElement:
    res = expr
    res = res.distinct()
    # This could be
    # `res = sa.func.string_agg(res, sep)`
    # but older postgres versions refuse to auto-cast into strings that way.
    # Also notable: `json_agg`, for unambiguous representation.
    res = sa.func.array_agg(res)
    res = sa.func.array_to_string(res, sep)
    return res


DEFINITIONS_AGG = [
    # all_concat
    base.AggAllConcat1(
        variants=[
            V(D.POSTGRESQL, _all_concat_pg),
        ]
    ),
    base.AggAllConcat2(
        variants=[
            V(D.POSTGRESQL, _all_concat_pg),
        ]
    ),
    # any
    base.AggAny(
        variants=[
            # Only in COMPENG.
            # For compatibility with CH and other dialects that support the ANY aggregation
            V(D.COMPENG, lambda value: sa.func.max(value)),
        ]
    ),
    # avg
    base.AggAvgFromNumber.for_dialect(D.POSTGRESQL),
    base.AggAvgFromDate.for_dialect(D.POSTGRESQL),
    base.AggAvgFromDatetime.for_dialect(D.POSTGRESQL),
    base.AggAvgFromDatetimeTZ.for_dialect(D.POSTGRESQL),
    # avg_if
    base.AggAvgIf.for_dialect(D.POSTGRESQL),
    # count
    base.AggCount0.for_dialect(D.POSTGRESQL),
    base.AggCount1.for_dialect(D.POSTGRESQL),
    # count_if
    base.AggCountIf.for_dialect(D.POSTGRESQL),
    # countd
    base.AggCountd.for_dialect(D.POSTGRESQL),
    # countd_if
    base.AggCountdIf.for_dialect(D.POSTGRESQL),
    # max
    base.AggMax.for_dialect(D.POSTGRESQL),
    # median
    base.AggMedian(
        variants=[
            V(
                D.and_above(D.POSTGRESQL_9_4) | D.COMPENG,
                lambda expr: (within_group(sa.func.PERCENTILE_DISC(0.5), expr)),
            ),
        ]
    ),
    # min
    base.AggMin.for_dialect(D.POSTGRESQL),
    # quantile
    base.AggQuantile(
        variants=[
            V(
                D.and_above(D.POSTGRESQL_9_4) | D.COMPENG,
                lambda expr, quant: (
                    # In MSSQLSRV window functions cannot be used as aggregations
                    within_group(
                        sa.func.percentile_disc(quantile_value(un_literal(quant))),
                        expr,
                    )
                ),
            ),
        ]
    ),
    # stdev
    base.AggStdev.for_dialect(D.POSTGRESQL),
    # stdevp
    base.AggStdevp.for_dialect(D.POSTGRESQL),
    # sum
    base.AggSum.for_dialect(D.POSTGRESQL),
    # sum_if
    base.AggSumIf.for_dialect(D.POSTGRESQL),
    # var
    base.AggVar.for_dialect(D.POSTGRESQL),
    # varp
    base.AggVarp.for_dialect(D.POSTGRESQL),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/functions_aggregation.py
datalens-tech/datalens-backend
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql as sa_postgresql
from sqlalchemy.sql import ClauseElement
from sqlalchemy.sql.elements import (
    Grouping,
    Null,
)
from sqlalchemy.sql.type_api import TypeEngine

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_array as base
from dl_formula.definitions.literals import (
    is_literal,
    un_literal,
)
from dl_formula.shortcuts import n

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D
from dl_connector_postgresql.formula.definitions.common import PG_INT_64_TO_CHAR_FMT


V = TranslationVariant.make


def _array_contains(array: ClauseElement, value: ClauseElement) -> ClauseElement:
    if isinstance(value, Null):
        return array != sa.func.array_remove(array, None)
    elif is_literal(value):
        return value == sa.func.ANY(array)
    else:
        return n.func.IF(
            n.func.ISNULL(value.self_group()), array != sa.func.array_remove(array, None), value == sa.func.ANY(array)
        )


def _array_notcontains(array: ClauseElement, value: ClauseElement) -> ClauseElement:
    if isinstance(value, Null):
        return array == sa.func.array_remove(array, None)
    elif is_literal(value):
        return value != sa.func.ALL(sa.func.array_remove(array, None))
    else:
        return n.func.IF(
            n.func.ISNULL(value.self_group()),
            array == sa.func.array_remove(array, None),
            value != sa.func.ALL(sa.func.array_remove(array, None)),
        )


DEFINITIONS_ARRAY = [
    # arr_remove
    base.FuncArrayRemoveLiteralNull(
        variants=[
            V(D.POSTGRESQL, sa.func.array_remove),
        ]
    ),
    base.FuncArrayRemoveDefault(
        variants=[
            V(D.POSTGRESQL, sa.func.array_remove),
        ]
    ),
    # arr_str
    base.FuncArrStr1(
        variants=[
            V(D.POSTGRESQL, lambda array: sa.func.array_to_string(array, ",")),
        ]
    ),
    base.FuncArrStr2(
        variants=[
            V(D.POSTGRESQL, lambda array, sep: sa.func.array_to_string(array, sep)),
        ]
    ),
    base.FuncArrStr3(
        variants=[
            V(D.POSTGRESQL, lambda array, sep, null_str: sa.func.array_to_string(array, sep, null_str)),
        ]
    ),
    # array
    base.FuncConstArrayFloat.for_dialect(D.POSTGRESQL),
    base.FuncConstArrayInt.for_dialect(D.POSTGRESQL),
    base.FuncConstArrayStr.for_dialect(D.POSTGRESQL),
    base.FuncNonConstArrayInt(
        variants=[
            V(D.POSTGRESQL, lambda *args: sa_postgresql.array(args)),
        ]
    ),
    base.FuncNonConstArrayFloat(
        variants=[
            V(D.POSTGRESQL, lambda *args: sa_postgresql.array(args)),
        ]
    ),
    base.FuncNonConstArrayStr(
        variants=[
            V(D.POSTGRESQL, lambda *args: sa_postgresql.array(args)),
        ]
    ),
    # cast_arr_float
    base.FuncFloatArrayFromIntArray(
        variants=[
            V(
                D.POSTGRESQL,
                lambda arr: sa.func.array(sa.select(sa.cast(sa.func.unnest(arr), sa_postgresql.DOUBLE_PRECISION))),
            ),
        ]
    ),
    base.FuncFloatArrayFromFloatArray.for_dialect(D.POSTGRESQL),
    base.FuncFloatArrayFromStringArray(
        variants=[
            V(D.POSTGRESQL, lambda arr: sa.func.array(sa.select(sa.cast(sa.func.unnest(arr), sa.FLOAT)))),
        ]
    ),
    # cast_arr_int
    base.FuncIntArrayFromIntArray.for_dialect(D.POSTGRESQL),
    base.FuncIntArrayFromFloatArray(
        variants=[
            V(
                D.POSTGRESQL,
                lambda arr: sa.func.array(sa.select(sa.cast(sa.func.floor(sa.func.unnest(arr)), sa.BIGINT))),
            ),
        ]
    ),
    base.FuncIntArrayFromStringArray(
        variants=[
            V(D.POSTGRESQL, lambda arr: sa.func.array(sa.select(sa.cast(sa.func.unnest(arr), sa.BIGINT)))),
        ]
    ),
    # cast_arr_str
    base.FuncStringArrayFromIntArray(
        variants=[
            V(
                D.POSTGRESQL,
                lambda arr: sa.func.array(sa.select(sa.func.to_char(sa.func.unnest(arr), PG_INT_64_TO_CHAR_FMT))),
            ),
        ]
    ),
    base.FuncStringArrayFromFloatArray(
        variants=[
            V(
                D.POSTGRESQL,
                lambda arr: sa.func.array(sa.select(sa.func.to_char(sa.func.unnest(arr), "FM999999999990.0999999999"))),
            ),
        ]
    ),
    base.FuncStringArrayFromStringArray.for_dialect(D.POSTGRESQL),
    # Array contains(_all, _any) functions comment:
    # Straight-forward `array_1.contains(array_2)` doesn't work due to postgres specifics in NULL and arrays interaction
    # [NULL, 1, 2].contains([NULL, 1]) = False in postgres, that is why branching on null containment needed
    # In PostgreSQL 9.5+ use sa.func.array_position(array, None) != None instead of array_remove comparison
    # contains
    base.FuncArrayContains(
        variants=[
            V(D.POSTGRESQL, _array_contains),
        ]
    ),
    # notcontains
    base.FuncArrayNotContains(
        variants=[
            V(D.POSTGRESQL, _array_notcontains),
        ]
    ),
    # contains_all
    base.FuncArrayContainsAll(
        variants=[
            V(
                D.POSTGRESQL,
                lambda array_1, array_2: n.func.IF(
                    array_2 != sa.func.array_remove(array_2, None),  # array_2 contains null
                    n.func.IF(
                        array_1 != sa.func.array_remove(array_1, None),  # array_1 contains null
                        sa.type_coerce(array_1, sa_postgresql.ARRAY(TypeEngine)).contains(
                            sa.type_coerce(sa.func.array_remove(array_2, None), sa_postgresql.ARRAY(TypeEngine))
                        ),
                        False,
                    ),
                    sa.type_coerce(array_1, sa_postgresql.ARRAY(TypeEngine)).contains(
                        sa.type_coerce(array_2, sa_postgresql.ARRAY(TypeEngine))
                    ),
                ),
            ),
        ]
    ),
    # contains_any
    base.FuncArrayContainsAny(
        variants=[
            V(
                D.POSTGRESQL,
                lambda array_1, array_2: n.func.IF(
                    array_2 != sa.func.array_remove(array_2, None),  # array_2 contains null
                    n.func.IF(
                        array_1 != sa.func.array_remove(array_1, None),  # array_1 contains null
                        True,
                        sa.type_coerce(array_1, sa_postgresql.ARRAY(TypeEngine)).overlap(
                            sa.type_coerce(sa.func.array_remove(array_2, None), sa_postgresql.ARRAY(TypeEngine))
                        ),
                    ),
                    sa.type_coerce(array_1, sa_postgresql.ARRAY(TypeEngine)).overlap(
                        sa.type_coerce(array_2, sa_postgresql.ARRAY(TypeEngine))
                    ),
                ),
            ),
        ]
    ),
    # count_item
    base.FuncArrayCountItemInt(
        variants=[
            V(
                D.POSTGRESQL,
                lambda array, value: (
                    sa.func.array_length(array, 1) - sa.func.array_length(sa.func.array_remove(array, value), 1)
                ),
            ),
        ]
    ),
    base.FuncArrayCountItemFloat(
        variants=[
            V(
                D.POSTGRESQL,
                lambda array, value: (
                    sa.func.array_length(array, 1)
                    - sa.func.array_length(
                        sa.func.array_remove(array, sa.cast(value, sa_postgresql.DOUBLE_PRECISION)), 1
                    )
                ),
            ),
        ]
    ),
    base.FuncArrayCountItemStr(
        variants=[
            V(
                D.POSTGRESQL,
                lambda array, value: (
                    sa.func.array_length(array, 1) - sa.func.array_length(sa.func.array_remove(array, value), 1)
                ),
            ),
        ]
    ),
    # get_item
    base.FuncGetArrayItemFloat(
        variants=[
            V(
                D.POSTGRESQL,
                # using Grouping as workaround to wrap expression before getitem with ()
                lambda array, index: Grouping(sa.type_coerce(array, sa_postgresql.ARRAY(TypeEngine)))[index],
            ),
        ]
    ),
    base.FuncGetArrayItemInt(
        variants=[
            V(
                D.POSTGRESQL,
                # using Grouping as workaround to wrap expression before getitem with ()
                lambda array, index: Grouping(sa.type_coerce(array, sa_postgresql.ARRAY(TypeEngine)))[index],
            ),
        ]
    ),
    base.FuncGetArrayItemStr(
        variants=[
            V(
                D.POSTGRESQL,
                # using Grouping as workaround to wrap expression before getitem with ()
                lambda array, index: Grouping(sa.type_coerce(array, sa_postgresql.ARRAY(TypeEngine)))[index],
            ),
        ]
    ),
    # len
    base.FuncLenArray(
        variants=[
            V(D.POSTGRESQL, lambda val: sa.func.ARRAY_LENGTH(val, 1)),
        ]
    ),
    # replace
    base.FuncReplaceArrayLiteralNull(
        variants=[
            V(D.POSTGRESQL, sa.func.array_replace),
        ]
    ),
    base.FuncReplaceArrayDefault(
        variants=[
            V(D.POSTGRESQL, sa.func.array_replace),
        ]
    ),
    # slice
    base.FuncArraySlice(
        variants=[
            V(
                D.POSTGRESQL,
                lambda array, offset, length: Grouping(sa.type_coerce(array, sa_postgresql.ARRAY(TypeEngine)))[
                    offset : offset + length - 1
                ],
            ),
        ]
    ),
    # startswith
    base.FuncStartswithArrayConst(
        variants=[
            V(
                D.POSTGRESQL,
                lambda x, y, _env: Grouping(sa.type_coerce(x, sa_postgresql.ARRAY(TypeEngine)))[1 : len(un_literal(y))]
                == y,
            ),
        ]
    ),
    base.FuncStartswithArrayNonConst(
        variants=[
            V(
                D.POSTGRESQL,
                lambda x, y, _env: Grouping(sa.type_coerce(x, sa_postgresql.ARRAY(TypeEngine)))[
                    1 : sa.func.array_length(y, 1)
                ]
                == y,
            ),
        ]
    ),
    # unnest
    base.FuncUnnestArrayFloat(
        variants=[
            V(D.POSTGRESQL, lambda arr: sa.func.unnest(arr)),
        ]
    ),
    base.FuncUnnestArrayInt(
        variants=[
            V(D.POSTGRESQL, lambda arr: sa.func.unnest(arr)),
        ]
    ),
    base.FuncUnnestArrayStr(
        variants=[
            V(D.POSTGRESQL, lambda arr: sa.func.unnest(arr)),
        ]
    ),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/functions_array.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import TYPE_CHECKING

import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

from dl_formula.definitions.base import (
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.common import raw_sql
from dl_formula.definitions.common_datetime import (
    datetime_interval,
    ensure_naive_first_arg,
)
import dl_formula.definitions.functions_datetime as base
from dl_formula.definitions.literals import un_literal
from dl_formula.shortcuts import n

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


if TYPE_CHECKING:
    from dl_formula.translation.context import TranslationCtx


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


def _datetrunc2_pg_tz_impl(date_ctx: TranslationCtx, unit_ctx: TranslationCtx) -> ClauseElement:
    unit = base.norm_datetrunc_unit(unit_ctx.expression)
    assert date_ctx.data_type_params is not None
    timezone = date_ctx.data_type_params.timezone
    date_expr = date_ctx.expression

    # Convert to local time (because postgresql doesn't have timezone argument
    # to date_trunc). `timestamptz` -> `timestamp`.
    date_expr = sa.func.timezone(timezone, date_expr)
    date_expr = sa.func.date_trunc(unit, date_expr)
    # Convert back to an aware (UTC) datetime. `timestamp` -> `timestamptz`.
    # NOTE: this is ambiguous-interpretation that might lead to incorrect
    # results within shifts from DST.
    date_expr = sa.func.timezone(timezone, date_expr)

    # not sure why mypy thinks there might be None
    assert date_expr is not None
    return date_expr


DEFINITIONS_DATETIME = [
    # dateadd
    base.FuncDateadd1.for_dialect(D.POSTGRESQL),
    base.FuncDateadd2Unit.for_dialect(D.POSTGRESQL),
    base.FuncDateadd2Number.for_dialect(D.POSTGRESQL),
    base.FuncDateadd3Legacy.for_dialect(D.POSTGRESQL),
    base.FuncDateadd3DateConstNum(
        variants=[
            V(
                D.POSTGRESQL,
                lambda date, what, num: (
                    sa.cast(
                        date
                        + datetime_interval(un_literal(what), un_literal(num), literal_type=True, literal_mult=True),
                        sa.Date,
                    )
                ),
            ),
        ]
    ),
    base.FuncDateadd3DatetimeConstNum(
        variants=[
            V(
                D.POSTGRESQL,
                lambda dt, what, num: (
                    dt + datetime_interval(what.value, num.value, literal_type=True, literal_mult=True)
                ),
            ),
        ]
    ),
    # datepart
    base.FuncDatepart2Legacy.for_dialect(D.POSTGRESQL),
    base.FuncDatepart2.for_dialect(D.POSTGRESQL),
    base.FuncDatepart3Const.for_dialect(D.POSTGRESQL),
    base.FuncDatepart3NonConst.for_dialect(D.POSTGRESQL),
    # datetrunc
    base.FuncDatetrunc2Date(
        variants=[
            V(
                D.POSTGRESQL,
                lambda date, unit: (
                    sa.cast(sa.func.date_trunc(base.norm_datetrunc_unit(unit), date), sa.Date)
                    if base.norm_datetrunc_unit(unit) in {"year", "quarter", "month", "week", "day"}
                    else date
                ),
            ),
        ]
    ),
    base.FuncDatetrunc2Datetime(
        variants=[
            VW(
                D.POSTGRESQL,
                lambda date, unit: (sa.func.date_trunc(base.norm_datetrunc_unit(unit.expression), date.expression)),
            ),
        ]
    ),
    base.FuncDatetrunc2DatetimeTZ(
        variants=[
            VW(D.POSTGRESQL, _datetrunc2_pg_tz_impl),
        ]
    ),
    # day
    base.FuncDay(
        variants=[
            VW(D.POSTGRESQL, lambda date: sa.cast(sa.func.DATE_PART("day", date.expression), sa.BIGINT)),
        ]
    ),
    base.FuncDayDatetimeTZ(
        variants=[
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.DAY)),
        ]
    ),
    # dayofweek
    base.FuncDayofweek1.for_dialect(D.POSTGRESQL),
    base.FuncDayofweek2(
        variants=[
            VW(
                D.POSTGRESQL,
                ensure_naive_first_arg(
                    lambda date_ctx, firstday_ctx: base.dow_firstday_shift(
                        # `date_part` returns a `double precision` value, so have to cast it;
                        # not casting to `SMALLINT` because pg doesn't expand typesize on e.g. multiplication.
                        sa.cast(sa.func.DATE_PART("isodow", date_ctx.expression), sa.BIGINT),
                        firstday_ctx.expression,
                    )
                ),
            ),
        ]
    ),
    base.FuncDayofweek2TZ(
        variants=[
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.DAYOFWEEK)),
        ]
    ),
    # genericnow
    base.FuncGenericNow(
        variants=[
            V(D.POSTGRESQL, sa.func.NOW),
        ]
    ),
    # hour
    base.FuncHourDate.for_dialect(D.POSTGRESQL),
    base.FuncHourDatetime(
        variants=[
            VW(D.POSTGRESQL, lambda date: sa.cast(sa.func.DATE_PART("hour", date.expression), sa.BIGINT)),
        ]
    ),
    base.FuncHourDatetimeTZ(
        variants=[
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.HOUR)),
        ]
    ),
    # minute
    base.FuncMinuteDate.for_dialect(D.POSTGRESQL),
    base.FuncMinuteDatetime(
        variants=[
            VW(D.POSTGRESQL, lambda date: sa.cast(sa.func.DATE_PART("minute", date.expression), sa.BIGINT)),
        ]
    ),
    base.FuncMinuteDatetimeTZ(
        variants=[
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.MINUTE)),
        ]
    ),
    # month
    base.FuncMonth(
        variants=[
            VW(D.POSTGRESQL, lambda date: sa.cast(sa.func.DATE_PART("month", date.expression), sa.BIGINT)),
        ]
    ),
    base.FuncMonthDatetimeTZ(
        variants=[
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.MONTH)),
        ]
    ),
    # now
    base.FuncNow(
        variants=[
            V(D.POSTGRESQL, sa.func.NOW),
        ]
    ),
    # quarter
    base.FuncQuarter(
        variants=[
            VW(D.POSTGRESQL, lambda date: sa.cast(sa.func.DATE_PART("quarter", date.expression), sa.BIGINT)),
        ]
    ),
    base.FuncQuarterDatetimeTZ(
        variants=[
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.QUARTER)),
        ]
    ),
    # second
    base.FuncSecondDate.for_dialect(D.POSTGRESQL),
    base.FuncSecondDatetime(
        variants=[
            VW(D.POSTGRESQL, lambda date: sa.cast(sa.func.DATE_PART("second", date.expression), sa.BIGINT)),
        ]
    ),
    base.FuncSecondDatetimeTZ(
        variants=[
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.SECOND)),
        ]
    ),
    # today
    base.FuncToday(
        variants=[
            V(D.POSTGRESQL, lambda: sa.type_coerce(raw_sql("CURRENT_DATE"), sa.Date)),
        ]
    ),
    # week
    base.FuncWeek(
        variants=[
            VW(D.POSTGRESQL, lambda date: sa.cast(sa.func.TO_CHAR(date.expression, "IW"), sa.BIGINT)),
        ]
    ),
    base.FuncWeekTZ(
        variants=[
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.WEEK)),
        ]
    ),
    # year
    base.FuncYear(
        variants=[
            VW(D.POSTGRESQL, lambda date: sa.cast(sa.func.DATE_PART("year", date.expression), sa.BIGINT)),
        ]
    ),
    base.FuncYearDatetimeTZ(
        variants=[
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.YEAR)),
        ]
    ),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/functions_datetime.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import (
    TranslationVariant,
    TranslationVariantWrapped,
)
import dl_formula.definitions.functions_logical as base
from dl_formula.shortcuts import n

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


DEFINITIONS_LOGICAL = [
    # case
    base.FuncCase.for_dialect(D.POSTGRESQL),
    # if
    base.FuncIf.for_dialect(D.POSTGRESQL),
    # ifnan
    base.FuncIfnan(
        variants=[
            VW(
                D.POSTGRESQL,
                lambda check_value, alt_value: n.func.IF(n.func.ISNAN(check_value), alt_value, check_value),
            ),
        ]
    ),
    # ifnull
    base.FuncIfnull(
        variants=[
            V(D.POSTGRESQL, sa.func.coalesce),
        ]
    ),
    # iif
    base.FuncIif3Legacy.for_dialect(D.POSTGRESQL),
    # isnan
    base.FuncIsnan(
        variants=[
            V(D.POSTGRESQL, lambda x: x == sa.text("double precision 'NaN'")),  # noqa
        ]
    ),
    # isnull
    base.FuncIsnull.for_dialect(D.POSTGRESQL),
    # zn
    base.FuncZn(
        variants=[
            V(D.POSTGRESQL, lambda x: sa.func.coalesce(x, 0)),
        ]
    ),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/functions_logical.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_markup as base

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


DEFINITIONS_MARKUP = [
    # +
    base.BinaryPlusMarkup.for_dialect(D.POSTGRESQL),
    # __str
    base.FuncInternalStrConst.for_dialect(D.POSTGRESQL),
    base.FuncInternalStr.for_dialect(D.POSTGRESQL),
    # bold
    base.FuncBold.for_dialect(D.POSTGRESQL),
    # italic
    base.FuncItalics.for_dialect(D.POSTGRESQL),
    # markup
    base.ConcatMultiMarkup.for_dialect(D.POSTGRESQL),
    # url
    base.FuncUrl.for_dialect(D.POSTGRESQL),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/functions_markup.py
datalens-tech/datalens-backend
from typing import ClassVar

import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

from dl_formula.core.datatype import DataType
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_math as base
from dl_formula.definitions.literals import (
    Literal,
    un_literal,
)

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


V = TranslationVariant.make


class FuncRound2ConstCompeng(base.FuncRound2):
    _DECIMAL_SIZE: ClassVar[int] = 20

    variants = [
        V(
            D.COMPENG,
            lambda number, precision: FuncRound2ConstCompeng._do_round_with_const_arg(number, precision),
        ),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.CONST_INTEGER]),
    ]

    @classmethod
    def _do_round_with_const_arg(self, number: Literal, precision: Literal) -> ClauseElement:
        precision_value = un_literal(precision)
        dec_precision = self._DECIMAL_SIZE
        dec_scale = precision_value + 5
        return sa.func.ROUND(sa.cast(number, sa.Numeric(precision=dec_precision, scale=dec_scale)), precision)


DEFINITIONS_MATH = [
    # abs
    base.FuncAbs.for_dialect(D.POSTGRESQL),
    # acos
    base.FuncAcos.for_dialect(D.POSTGRESQL),
    # asin
    base.FuncAsin.for_dialect(D.POSTGRESQL),
    # atan
    base.FuncAtan.for_dialect(D.POSTGRESQL),
    # atan2
    base.FuncAtan2.for_dialect(D.POSTGRESQL),
    # ceiling
    base.FuncCeiling.for_dialect(D.POSTGRESQL),
    # cos
    base.FuncCos.for_dialect(D.POSTGRESQL),
    # cot
    base.FuncCot.for_dialect(D.POSTGRESQL),
    # degrees
    base.FuncDegrees.for_dialect(D.POSTGRESQL),
    # div
    base.FuncDivBasic(
        variants=[
            V(D.COMPENG, lambda x, y: sa.func.div(x, sa.func.nullif(y, 0))),
        ]
    ),
    base.FuncDivBasic.for_dialect(D.NON_COMPENG_POSTGRESQL),
    # div_safe
    base.FuncDivSafe2.for_dialect(D.POSTGRESQL),
    base.FuncDivSafe3.for_dialect(D.POSTGRESQL),
    # exp
    base.FuncExp.for_dialect(D.POSTGRESQL),
    # fdiv_safe
    base.FuncFDivSafe2.for_dialect(D.POSTGRESQL),
    base.FuncFDivSafe3.for_dialect(D.POSTGRESQL),
    # floor
    base.FuncFloor.for_dialect(D.POSTGRESQL),
    # greatest
    base.FuncGreatest1.for_dialect(D.POSTGRESQL),
    base.FuncGreatestMain.for_dialect(D.POSTGRESQL),
    base.GreatestMulti.for_dialect(D.POSTGRESQL),
    # least
    base.FuncLeast1.for_dialect(D.POSTGRESQL),
    base.FuncLeastMain.for_dialect(D.POSTGRESQL),
    base.LeastMulti.for_dialect(D.POSTGRESQL),
    # ln
    base.FuncLn.for_dialect(D.POSTGRESQL),
    # log
    base.FuncLog(
        variants=[
            V(
                D.COMPENG,
                lambda x, base: (
                    sa.func.ln(sa.func.nullif(x, 0)) / sa.func.nullif(sa.func.ln(sa.func.nullif(base, 0)), 0)
                ),
            ),
        ]
    ),
    base.FuncLog(
        variants=[
            V(D.NON_COMPENG_POSTGRESQL, lambda x, base: sa.func.ln(x) / sa.func.ln(base)),
        ]
    ),
    # log10
    base.FuncLog10(
        variants=[
            V(D.POSTGRESQL, sa.func.LOG),
        ]
    ),
    # pi
    base.FuncPi.for_dialect(D.POSTGRESQL),
    # power
    base.FuncPower.for_dialect(D.POSTGRESQL),
    # radians
    base.FuncRadians.for_dialect(D.POSTGRESQL),
    # round
    FuncRound2ConstCompeng(),  # To avoid manual casts to numeric on user's side in COMPENG
    base.FuncRound1.for_dialect(D.POSTGRESQL),
    base.FuncRound2.for_dialect(D.POSTGRESQL),
    # sign
    base.FuncSign.for_dialect(D.POSTGRESQL),
    # sin
    base.FuncSin.for_dialect(D.POSTGRESQL),
    # sqrt
    base.FuncSqrt.for_dialect(D.POSTGRESQL),
    # square
    base.FuncSquare(
        variants=[
            V(D.POSTGRESQL, lambda x: sa.func.POW(x, 2)),
        ]
    ),
    # tan
    base.FuncTan.for_dialect(D.POSTGRESQL),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/functions_math.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_special as base

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


V = TranslationVariant.make


DEFINITIONS_SPECIAL = [
    # _make_nan
    base.FuncMakeNan(
        variants=[
            V(D.POSTGRESQL, lambda: sa.column("double precision 'NaN'", is_literal=True)),
        ]
    ),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/functions_special.py
datalens-tech/datalens-backend
import sqlalchemy as sa
import sqlalchemy.dialects.postgresql as sa_postgresql

from dl_formula.definitions.base import (
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.common import make_binary_chain
import dl_formula.definitions.functions_string as base
from dl_formula.shortcuts import n

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


DEFINITIONS_STRING = [
    # ascii
    base.FuncAscii.for_dialect(D.POSTGRESQL),
    # char
    base.FuncChar(
        variants=[
            V(D.POSTGRESQL, sa.func.CHR),
        ]
    ),
    # concat
    base.Concat1.for_dialect((D.POSTGRESQL)),
    base.ConcatMultiStrConst.for_dialect(D.POSTGRESQL),
    base.ConcatMultiStr(
        variants=[
            V(
                D.POSTGRESQL,
                lambda *args: make_binary_chain(
                    (lambda x, y: x.concat(y)),  # should result in `x || y` SQL.
                    *args,  # should result in `x || y || z || ...` SQL
                    wrap_as_nodes=False,
                ),
            ),
        ]
    ),
    base.ConcatMultiAny.for_dialect(D.POSTGRESQL),
    # contains
    base.FuncContainsConst.for_dialect(D.POSTGRESQL),
    base.FuncContainsNonConst(
        variants=[
            V(D.POSTGRESQL, lambda x, y: sa.func.STRPOS(x, y) > 0),
        ]
    ),
    base.FuncContainsNonString.for_dialect(D.POSTGRESQL),
    # notcontains
    base.FuncNotContainsConst.for_dialect(D.POSTGRESQL),
    base.FuncNotContainsNonConst.for_dialect(D.POSTGRESQL),
    base.FuncNotContainsNonString.for_dialect(D.POSTGRESQL),
    # endswith
    base.FuncEndswithConst.for_dialect(D.POSTGRESQL),
    base.FuncEndswithNonConst(
        variants=[
            V(
                D.POSTGRESQL,
                lambda x, y: (sa.func.SUBSTRING(x, sa.func.char_length(x) - sa.func.char_length(y) + 1) == y),
            ),
        ]
    ),
    base.FuncEndswithNonString.for_dialect(D.POSTGRESQL),
    # find
    base.FuncFind2(
        variants=[
            V(D.POSTGRESQL, sa.func.STRPOS),
        ]
    ),
    base.FuncFind3(
        variants=[
            V(
                D.POSTGRESQL,
                lambda x, y, z: n.if_(sa.func.STRPOS(sa.func.SUBSTRING(x, z), y) > 0)  # type: ignore  # TODO: fix
                .then(
                    sa.func.STRPOS(sa.func.SUBSTRING(x, z), y) + z - 1,
                )
                .else_(0),
            ),
        ]
    ),
    # icontains
    base.FuncIContainsConst.for_dialect(D.POSTGRESQL),
    base.FuncIContainsNonConst.for_dialect(D.POSTGRESQL),
    base.FuncIContainsNonString.for_dialect(D.POSTGRESQL),
    # iendswith
    base.FuncIEndswithConst.for_dialect(D.POSTGRESQL),
    base.FuncIEndswithNonConst.for_dialect(D.POSTGRESQL),
    base.FuncIEndswithNonString.for_dialect(D.POSTGRESQL),
    # istartswith
    base.FuncIStartswithConst.for_dialect(D.POSTGRESQL),
    base.FuncIStartswithNonConst.for_dialect(D.POSTGRESQL),
    base.FuncIStartswithNonString.for_dialect(D.POSTGRESQL),
    # left
    base.FuncLeft.for_dialect(D.POSTGRESQL),
    # len
    base.FuncLenString(
        variants=[
            V(D.POSTGRESQL, sa.func.LENGTH),
        ]
    ),
    # lower
    base.FuncLowerConst.for_dialect(D.POSTGRESQL),
    base.FuncLowerNonConst.for_dialect(D.POSTGRESQL),
    # ltrim
    base.FuncLtrim.for_dialect(D.POSTGRESQL),
    # regexp_extract
    base.FuncRegexpExtract(
        variants=[
            V(
                D.POSTGRESQL,
                lambda text, pattern: sa.type_coerce(
                    sa.select([sa.func.REGEXP_MATCHES(text, pattern, "g")]).limit(1).as_scalar(),
                    sa_postgresql.ARRAY(sa.String),
                )[1],
            ),
        ]
    ),
    # regexp_extract_nth
    base.FuncRegexpExtractNth(
        variants=[
            V(
                D.POSTGRESQL,
                lambda text, pattern, ind: sa.type_coerce(
                    sa.select([sa.func.REGEXP_MATCHES(text, pattern, "g")]).limit(1).offset(ind - 1).as_scalar(),
                    sa_postgresql.ARRAY(sa.String),
                )[1],
            ),
        ]
    ),
    # regexp_match
    base.FuncRegexpMatch(
        variants=[
            V(D.POSTGRESQL, lambda text, pattern: text.op("~")(pattern)),
        ]
    ),
    # regexp_replace
    base.FuncRegexpReplace(
        variants=[
            V(D.POSTGRESQL, lambda text, patt, repl: sa.func.REGEXP_REPLACE(text, patt, repl, "g")),
        ]
    ),
    # replace
    base.FuncReplace.for_dialect(D.POSTGRESQL),
    # right
    base.FuncRight.for_dialect(D.POSTGRESQL),
    # rtrim
    base.FuncRtrim.for_dialect(D.POSTGRESQL),
    # space
    base.FuncSpaceConst.for_dialect(D.POSTGRESQL),
    base.FuncSpaceNonConst(
        variants=[
            V(D.POSTGRESQL, lambda size: sa.func.REPEAT(" ", size)),
        ]
    ),
    # split
    base.FuncSplit1(
        variants=[
            V(D.POSTGRESQL, lambda text: sa.func.string_to_array(text, ",")),
        ]
    ),
    base.FuncSplit2(
        variants=[
            V(D.POSTGRESQL, lambda text, delim: sa.func.string_to_array(text, delim)),
        ]
    ),
    base.FuncSplit3(
        variants=[
            V(
                D.POSTGRESQL, lambda text, delim, ind: sa.func.SPLIT_PART(text, delim, sa.cast(ind, sa.INTEGER))
            ),  # FIXME: does not work with negative indices
        ]
    ),
    # startswith
    base.FuncStartswithConst.for_dialect(D.POSTGRESQL),
    base.FuncStartswithNonConst(
        variants=[
            V(D.POSTGRESQL, lambda x, y: sa.func.STRPOS(x, y) == 1),
        ]
    ),
    base.FuncStartswithNonString.for_dialect(D.POSTGRESQL),
    # substr
    base.FuncSubstr2.for_dialect(D.POSTGRESQL),
    base.FuncSubstr3.for_dialect(D.POSTGRESQL),
    # trim
    base.FuncTrim.for_dialect(D.POSTGRESQL),
    # upper
    base.FuncUpperConst.for_dialect(D.POSTGRESQL),
    base.FuncUpperNonConst.for_dialect(D.POSTGRESQL),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/functions_string.py
datalens-tech/datalens-backend
import sqlalchemy as sa
import sqlalchemy.dialects.postgresql as sa_postgresql

from dl_formula.core.datatype import DataType
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    Function,
    SingleVariantTranslationBase,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.common_datetime import ensure_naive_first_arg
import dl_formula.definitions.functions_type as base
from dl_formula.definitions.scope import Scope
from dl_formula.shortcuts import n

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D
from dl_connector_postgresql.formula.definitions.common import PG_INT_64_TO_CHAR_FMT


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class FuncTypeGenericDatetime2PGImpl(SingleVariantTranslationBase, base.FuncTypeGenericDatetime2Impl):
    dialects = D.POSTGRESQL
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.DATETIME, DataType.GENERICDATETIME, DataType.INTEGER, DataType.FLOAT, DataType.STRING},
                DataType.CONST_STRING,
            ]
        ),
    ]

    @classmethod
    def _translate_main(cls, value_ctx, tz_ctx):  # type: ignore
        """
        Equivalent to `dt at time zone tz`.
        Its semantics:

        1. For `timestamp` (naive) values, the value is interpreted
          at the specified timezone (with ambiguity),
          returning `timestamptz` (utc). Example:
          `select timezone('Europe/Moscow', '2010-10-31T02:00:00'::timestamp);`
          -> `2010-10-30 23:00:00+00`.
        2. For `timestamptz` (aware) values, the value is shifted to local time,
          returning `timestamp` (naive).
          Example:
          `select timezone('Europe/Moscow', timezone('Europe/Moscow', '2010-10-31T02:00:00'::timestamp));`
          -> `2010-10-31 02:00:00` (back to an ambiguous value).
        3. String / integer / float values are not accepted (however, some pg
          clients might autocast strings).
        4. `date` argument works, but it doesn't seem useful. Example:
          `select timezone('Europe/Moscow', '2010-10-31'::date);`
          -> `2010-10-31 03:00:00`

        This implementation provides a way to make datetime from string / integer / float values but it fully relies
        on database when it comes to transforming between naive and aware values.

        If input value is naive timestamp or string (will be casted to naive as well),
        the result will be tz-aware value (see 1. above)
        but in case of input being tz-aware timestamp or integer / float (casted to tz-aware utc value)
        the result will be naive (see 2. above)
        """
        value = value_ctx.expression
        value_type = value_ctx.data_type
        tz = tz_ctx.expression

        if value_type in (
            DataType.DATETIME,
            DataType.CONST_DATETIME,
            DataType.GENERICDATETIME,
            DataType.CONST_GENERICDATETIME,
        ):
            pass  # value = value
        elif value_type in (DataType.INTEGER, DataType.CONST_INTEGER, DataType.FLOAT, DataType.CONST_FLOAT):
            value = sa.func.to_timestamp(value)
        elif value_type in (DataType.STRING, DataType.CONST_STRING):
            value = sa.cast(value, sa.DateTime())
        else:
            raise Exception("Unexpected value type, likely a translation definition error", value_type)

        return sa.func.timezone(tz, value)


class FuncDatetimeTZPG(SingleVariantTranslationBase, base.FuncDatetimeTZ):
    dialects = D.POSTGRESQL
    argument_types = [
        ArgTypeSequence(
            [
                {
                    DataType.DATETIME,
                    DataType.GENERICDATETIME,
                    DataType.DATETIMETZ,
                    # Not even tried: `DataType.DATE`.
                    DataType.INTEGER,
                    DataType.FLOAT,
                    DataType.STRING,
                },
                DataType.CONST_STRING,
            ]
        ),
    ]

    @classmethod
    def _translate_main(cls, value_ctx, tz_ctx):  # type: ignore
        value = value_ctx.expression
        value_type = value_ctx.data_type
        tz = tz_ctx.expression

        if value_type == DataType.DATETIMETZ:
            # Only change the data_type_params, leave the value
            return value

        if value_type == DataType.GENERICDATETIME:
            return value

        if value_type == DataType.DATETIME:
            value = sa.cast(value, sa.TIMESTAMP())  # ensure naive dt
            return sa.func.timezone(tz, value)  # make `timestamptz`

        if value_type in (DataType.INTEGER, DataType.FLOAT):
            value = sa.func.to_timestamp(value)  # -> `timestamptz`
            # The timezone goes into the data_type_params.
            return value

        elif value_type == DataType.STRING:
            # Casting to `timestamptz` ensures tzoffset in ISO strings isn't ignored.
            value = sa.cast(value, sa.TIMESTAMP(timezone=True))
            # Then making the naive datetime of the value minus the offset (if it had one)
            value = sa.cast(value, sa.TIMESTAMP(timezone=False))
            # Then interpreting the naive datetime in the specified timezone
            return sa.func.timezone(tz, value)

        raise Exception("Unexpected value type, likely a translation definition error", value_type)


class FuncDatetime2PG(FuncTypeGenericDatetime2PGImpl):
    name = "datetime"


class FuncGenericDatetime2PG(FuncTypeGenericDatetime2PGImpl):
    name = "genericdatetime"
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncDatetimeTZToNaivePG(base.FuncDatetimeTZToNaive):
    dialects = D.POSTGRESQL | D.COMPENG

    @classmethod
    def _translate_main(cls, value_ctx):  # type: ignore
        expr = value_ctx.expression
        tz = value_ctx.data_type_params.timezone
        assert tz

        # Ensure the value is tz-aware in the postgresql:
        expr = sa.cast(expr, sa.TIMESTAMP(timezone=True))
        # Convert to postgresql's tz-naive value:
        expr = sa.func.timezone(tz, expr)
        return expr


class FuncDbCastPostgreSQLBase(base.FuncDbCastBase):
    WHITELISTS = {
        pg_dialect: {
            DataType.INTEGER: [
                base.WhitelistTypeSpec(name="smallint", sa_type=sa_postgresql.SMALLINT),
                base.WhitelistTypeSpec(name="integer", sa_type=sa_postgresql.INTEGER),
                base.WhitelistTypeSpec(name="bigint", sa_type=sa_postgresql.BIGINT),
            ],
            DataType.FLOAT: [
                base.WhitelistTypeSpec(name="double precision", sa_type=sa_postgresql.DOUBLE_PRECISION),
                base.WhitelistTypeSpec(name="real", sa_type=sa_postgresql.REAL),
                base.WhitelistTypeSpec(
                    name="numeric", sa_type=sa_postgresql.NUMERIC, arg_types=base.DECIMAL_CAST_ARG_T
                ),
            ],
            DataType.STRING: [
                base.WhitelistTypeSpec(name="text", sa_type=sa_postgresql.TEXT),
                base.WhitelistTypeSpec(name="character", sa_type=sa_postgresql.CHAR, arg_types=base.CHAR_CAST_ARG_T),
                base.WhitelistTypeSpec(
                    name="character varying", sa_type=sa_postgresql.VARCHAR, arg_types=base.CHAR_CAST_ARG_T
                ),
                # Here we'll make an exception and allow the usage of type aliases
                # just because in this case they are probably much more widely used
                # then the proper type names:
                base.WhitelistTypeSpec(name="char", sa_type=sa_postgresql.CHAR, arg_types=base.CHAR_CAST_ARG_T),
                base.WhitelistTypeSpec(name="varchar", sa_type=sa_postgresql.VARCHAR, arg_types=base.CHAR_CAST_ARG_T),
            ],
            DataType.ARRAY_STR: [
                base.WhitelistTypeSpec(name="text[]", sa_type=sa_postgresql.ARRAY, nested_sa_type=sa_postgresql.TEXT),
                base.WhitelistTypeSpec(
                    name="character varying[]", sa_type=sa_postgresql.ARRAY, nested_sa_type=sa_postgresql.VARCHAR
                ),
                base.WhitelistTypeSpec(
                    name="varchar[]", sa_type=sa_postgresql.ARRAY, nested_sa_type=sa_postgresql.VARCHAR
                ),
            ],
            DataType.ARRAY_INT: [
                base.WhitelistTypeSpec(
                    name="smallint[]", sa_type=sa_postgresql.ARRAY, nested_sa_type=sa_postgresql.SMALLINT
                ),
                base.WhitelistTypeSpec(
                    name="integer[]", sa_type=sa_postgresql.ARRAY, nested_sa_type=sa_postgresql.INTEGER
                ),
                base.WhitelistTypeSpec(
                    name="bigint[]", sa_type=sa_postgresql.ARRAY, nested_sa_type=sa_postgresql.BIGINT
                ),
            ],
            DataType.ARRAY_FLOAT: [
                base.WhitelistTypeSpec(
                    name="double precision[]",
                    sa_type=sa_postgresql.ARRAY,
                    nested_sa_type=sa_postgresql.DOUBLE_PRECISION,
                ),
                base.WhitelistTypeSpec(name="real[]", sa_type=sa_postgresql.ARRAY, nested_sa_type=sa_postgresql.REAL),
                base.WhitelistTypeSpec(
                    name="numeric[]",
                    sa_type=sa_postgresql.ARRAY,
                    nested_sa_type=sa_postgresql.NUMERIC,
                    arg_types=base.DECIMAL_CAST_ARG_T,
                ),
            ],
        }
        for pg_dialect in (D.COMPENG, D.NON_COMPENG_POSTGRESQL)
    }


class FuncDbCastPostgreSQL2(FuncDbCastPostgreSQLBase, base.FuncDbCast2):
    pass


class FuncDbCastPostgreSQL3(FuncDbCastPostgreSQLBase, base.FuncDbCast3):
    pass


class FuncDbCastPostgreSQL4(FuncDbCastPostgreSQLBase, base.FuncDbCast4):
    pass


DEFINITIONS_TYPE = [
    # bool
    base.FuncBoolFromNull.for_dialect(D.POSTGRESQL),
    base.FuncBoolFromNumber.for_dialect(D.POSTGRESQL),
    base.FuncBoolFromBool.for_dialect(D.POSTGRESQL),
    base.FuncBoolFromStrGeo.for_dialect(D.POSTGRESQL),
    base.FuncBoolFromDateDatetime.for_dialect(D.POSTGRESQL),
    # date
    base.FuncDate1FromNull.for_dialect(D.POSTGRESQL),
    base.FuncDate1FromDatetime.for_dialect(D.POSTGRESQL),
    base.FuncDate1FromDatetimeTZ(
        variants=[
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.DATE)),
        ]
    ),
    base.FuncDate1FromString.for_dialect(D.POSTGRESQL),
    base.FuncDate1FromNumber(
        variants=[
            V(D.POSTGRESQL, lambda expr: sa.cast(sa.func.TO_TIMESTAMP(expr), sa.Date())),
        ]
    ),
    # datetime
    base.FuncDatetime1FromNull.for_dialect(D.POSTGRESQL),
    base.FuncDatetime1FromDatetime.for_dialect(D.POSTGRESQL),
    base.FuncDatetime1FromDate.for_dialect(D.POSTGRESQL),
    base.FuncDatetime1FromNumber(
        variants=[
            V(D.POSTGRESQL, lambda expr: sa.func.TO_TIMESTAMP(expr)),
        ]
    ),
    base.FuncDatetime1FromString.for_dialect(D.POSTGRESQL),
    FuncDatetime2PG(),
    # datetimetz
    base.FuncDatetimeTZConst.for_dialect(D.POSTGRESQL),
    FuncDatetimeTZPG(),
    # datetimetz_to_naive
    FuncDatetimeTZToNaivePG(),
    # db_cast
    FuncDbCastPostgreSQL2(),
    FuncDbCastPostgreSQL3(),
    FuncDbCastPostgreSQL4(),
    # float
    base.FuncFloatNumber(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(value, sa_postgresql.DOUBLE_PRECISION)),
        ]
    ),
    base.FuncFloatString(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(value, sa_postgresql.DOUBLE_PRECISION)),
        ]
    ),
    base.FuncFloatFromBool(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(sa.cast(value, sa.INTEGER), sa_postgresql.DOUBLE_PRECISION)),
        ]
    ),
    base.FuncFloatFromDate(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(sa.func.extract("epoch", value), sa_postgresql.DOUBLE_PRECISION)),
        ]
    ),
    base.FuncFloatFromDatetime(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(sa.func.extract("epoch", value), sa_postgresql.DOUBLE_PRECISION)),
        ]
    ),
    base.FuncFloatFromGenericDatetime(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(sa.func.extract("epoch", value), sa_postgresql.DOUBLE_PRECISION)),
        ]
    ),
    # genericdatetime
    base.FuncGenericDatetime1FromNull.for_dialect(D.POSTGRESQL),
    base.FuncGenericDatetime1FromDatetime.for_dialect(D.POSTGRESQL),
    base.FuncGenericDatetime1FromDate.for_dialect(D.POSTGRESQL),
    base.FuncGenericDatetime1FromNumber(
        variants=[
            V(D.POSTGRESQL, lambda expr: sa.func.TO_TIMESTAMP(expr)),
        ]
    ),
    base.FuncGenericDatetime1FromString.for_dialect(D.POSTGRESQL),
    FuncGenericDatetime2PG(),
    # geopoint
    base.FuncGeopointFromStr.for_dialect(D.POSTGRESQL),
    base.FuncGeopointFromCoords.for_dialect(D.POSTGRESQL),
    # geopolygon
    base.FuncGeopolygon.for_dialect(D.POSTGRESQL),
    # int
    base.FuncIntFromNull(
        variants=[
            V(D.POSTGRESQL, lambda _: sa.cast(sa.null(), sa.BIGINT())),
        ]
    ),
    base.FuncIntFromInt.for_dialect(D.POSTGRESQL),
    base.FuncIntFromFloat(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(sa.func.FLOOR(value), sa.BIGINT)),
        ]
    ),
    base.FuncIntFromBool(
        variants=[
            # Direct cast bool->bigint is not supported; yet, casting to bigint for type consistency.
            V(D.POSTGRESQL, lambda value: sa.cast(sa.cast(value, sa.INTEGER), sa.BIGINT)),
        ]
    ),
    base.FuncIntFromStr(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(value, sa.BIGINT)),
        ]
    ),
    base.FuncIntFromDate(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(sa.func.extract("epoch", value), sa.BIGINT)),
        ]
    ),
    base.FuncIntFromDatetime(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(sa.func.extract("epoch", value), sa.BIGINT)),
        ]
    ),
    base.FuncIntFromGenericDatetime(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(sa.func.extract("epoch", value), sa.BIGINT)),
        ]
    ),
    base.FuncIntFromDatetimeTZ(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(sa.func.extract("epoch", value), sa.BIGINT)),
        ]
    ),
    # str
    base.FuncStrFromNull(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(sa.null(), sa.String())),
        ]
    ),
    base.FuncStrFromUnsupported(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(value, sa.TEXT)),  # supports more than `to_char`
        ]
    ),
    base.FuncStrFromInteger(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.func.TO_CHAR(value, PG_INT_64_TO_CHAR_FMT)),
        ]
    ),
    base.FuncStrFromFloat(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.func.TO_CHAR(value, "FM999999999990.0999999999")),
        ]
    ),
    base.FuncStrFromBool(
        variants=[
            V(
                D.POSTGRESQL,
                lambda value: sa.case(whens=[(value.is_(None), sa.null()), (value, "True")], else_="False"),
            ),
        ]
    ),
    base.FuncStrFromStrGeo.for_dialect(D.POSTGRESQL),
    base.FuncStrFromDate(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.func.TO_CHAR(value, "YYYY-MM-DD")),
        ]
    ),
    base.FuncStrFromDatetime(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.func.TO_CHAR(value, "YYYY-MM-DD HH24:MI:SS")),
        ]
    ),
    base.FuncStrFromDatetimeTZ(
        variants=[
            # There might be better options; but for now, just str() the naive version.
            # Particularly relevant because in PG and CH the aware value is actually in UTC
            VW(D.POSTGRESQL, ensure_naive_first_arg(n.func.STR)),
        ]
    ),
    base.FuncStrFromString.for_dialect(D.POSTGRESQL),
    base.FuncStrFromUUID(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(value, sa.TEXT)),
        ]
    ),
    base.FuncStrFromArray(
        variants=[
            V(D.POSTGRESQL, lambda value: sa.cast(value, sa.TEXT)),
        ]
    ),
    # tree
    base.FuncTreeStr.for_dialect(D.POSTGRESQL),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/functions_type.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_window as base

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


SUPPORTED_DIALECTS = D.COMPENG | D.and_above(D.POSTGRESQL_9_4)


DEFINITIONS_WINDOW = [
    # avg
    base.WinAvg.for_dialect(SUPPORTED_DIALECTS),
    # avg_if
    base.WinAvgIf.for_dialect(SUPPORTED_DIALECTS),
    # count
    base.WinCount0.for_dialect(SUPPORTED_DIALECTS),
    base.WinCount1.for_dialect(SUPPORTED_DIALECTS),
    # count_if
    base.WinCountIf.for_dialect(SUPPORTED_DIALECTS),
    # first
    base.WinFirst.for_dialect(SUPPORTED_DIALECTS),
    # lag
    base.WinLag1.for_dialect(SUPPORTED_DIALECTS),
    base.WinLag2.for_dialect(SUPPORTED_DIALECTS),
    base.WinLag3.for_dialect(SUPPORTED_DIALECTS),
    # last
    base.WinLast.for_dialect(SUPPORTED_DIALECTS),
    # mavg
    base.WinMAvg2.for_dialect(SUPPORTED_DIALECTS),
    base.WinMAvg3.for_dialect(SUPPORTED_DIALECTS),
    # max
    base.WinMax.for_dialect(SUPPORTED_DIALECTS),
    # mcount
    base.WinMCount2.for_dialect(SUPPORTED_DIALECTS),
    base.WinMCount3.for_dialect(SUPPORTED_DIALECTS),
    # min
    base.WinMin.for_dialect(SUPPORTED_DIALECTS),
    # mmax
    base.WinMMax2.for_dialect(SUPPORTED_DIALECTS),
    base.WinMMax3.for_dialect(SUPPORTED_DIALECTS),
    # mmin
    base.WinMMin2.for_dialect(SUPPORTED_DIALECTS),
    base.WinMMin3.for_dialect(SUPPORTED_DIALECTS),
    # msum
    base.WinMSum2.for_dialect(SUPPORTED_DIALECTS),
    base.WinMSum3.for_dialect(SUPPORTED_DIALECTS),
    # rank
    base.WinRank1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRank2.for_dialect(SUPPORTED_DIALECTS),
    # rank_dense
    base.WinRankDense1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRankDense2.for_dialect(SUPPORTED_DIALECTS),
    # rank_percentile
    base.WinRankPercentile1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRankPercentile2.for_dialect(SUPPORTED_DIALECTS),
    # rank_unique
    base.WinRankUnique1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRankUnique2.for_dialect(SUPPORTED_DIALECTS),
    # ravg
    base.WinRAvg1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRAvg2.for_dialect(SUPPORTED_DIALECTS),
    # rcount
    base.WinRCount1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRCount2.for_dialect(SUPPORTED_DIALECTS),
    # rmax
    base.WinRMax1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRMax2.for_dialect(SUPPORTED_DIALECTS),
    # rmin
    base.WinRMin1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRMin2.for_dialect(SUPPORTED_DIALECTS),
    # rsum
    base.WinRSum1.for_dialect(SUPPORTED_DIALECTS),
    base.WinRSum2.for_dialect(SUPPORTED_DIALECTS),
    # sum
    base.WinSum.for_dialect(SUPPORTED_DIALECTS),
    # sum_if
    base.WinSumIf.for_dialect(SUPPORTED_DIALECTS),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/functions_window.py
datalens-tech/datalens-backend
import sqlalchemy as sa
import sqlalchemy.dialects.postgresql as sa_postgresql

from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common_datetime import DAY_SEC
import dl_formula.definitions.operators_binary as base

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


V = TranslationVariant.make


DEFINITIONS_BINARY = [
    # !=
    base.BinaryNotEqual.for_dialect(D.POSTGRESQL),
    # %
    base.BinaryModInteger.for_dialect(D.POSTGRESQL),
    base.BinaryModFloat.for_dialect(D.POSTGRESQL),
    # *
    base.BinaryMultNumbers.for_dialect(D.POSTGRESQL),
    base.BinaryMultStringConst.for_dialect(D.POSTGRESQL),
    base.BinaryMultStringNonConst(
        variants=[
            V(D.POSTGRESQL, sa.func.REPEAT),
        ]
    ),
    # +
    base.BinaryPlusNumbers.for_dialect(D.POSTGRESQL),
    base.BinaryPlusStrings.for_dialect(D.POSTGRESQL),
    base.BinaryPlusDateInt(
        variants=[
            V(
                D.POSTGRESQL,
                lambda date, days: sa.cast(
                    sa.func.to_timestamp(sa.func.extract("epoch", date) + days * DAY_SEC), sa.Date
                ),
            ),
        ]
    ),
    base.BinaryPlusArray(
        variants=[
            V(D.POSTGRESQL, sa.func.array_cat),
        ]
    ),
    base.BinaryPlusDateFloat(
        variants=[
            V(
                D.POSTGRESQL,
                lambda date, days: sa.cast(
                    sa.func.to_timestamp(sa.func.extract("epoch", date) + base.as_bigint(days) * DAY_SEC), sa.Date
                ),
            ),
        ]
    ),
    base.BinaryPlusDatetimeNumber(
        variants=[
            V(D.POSTGRESQL, lambda dt, days: sa.func.to_timestamp(sa.func.extract("epoch", dt) + days * DAY_SEC)),
        ]
    ),
    base.BinaryPlusGenericDatetimeNumber(
        variants=[
            V(
                D.POSTGRESQL,
                lambda dt, days: dt + sa.cast(sa.cast(days, sa.TEXT) + " day", sa.Interval(second_precision=6)),
            ),
        ]
    ),
    # -
    base.BinaryMinusNumbers.for_dialect(D.POSTGRESQL),
    base.BinaryMinusDateInt(
        variants=[
            V(
                D.POSTGRESQL,
                lambda date, days: sa.cast(
                    sa.func.to_timestamp(sa.func.extract("epoch", date) - days * DAY_SEC), sa.Date
                ),
            ),
        ]
    ),
    base.BinaryMinusDateFloat(
        variants=[
            V(
                D.POSTGRESQL,
                lambda date, days: sa.cast(
                    sa.func.to_timestamp(sa.func.extract("epoch", date) - base.as_bigint(sa.func.CEIL(days)) * DAY_SEC),
                    sa.Date,
                ),
            ),
        ]
    ),
    base.BinaryMinusDatetimeNumber(
        variants=[
            V(D.POSTGRESQL, lambda dt, days: sa.func.to_timestamp(sa.func.extract("epoch", dt) - days * DAY_SEC)),
        ]
    ),
    base.BinaryMinusGenericDatetimeNumber(
        variants=[
            V(
                D.POSTGRESQL,
                lambda dt, days: dt - sa.cast(sa.cast(days, sa.TEXT) + " day", sa.Interval(second_precision=6)),
            ),
        ]
    ),
    base.BinaryMinusDates.for_dialect(D.POSTGRESQL),
    base.BinaryMinusDatetimes(
        variants=[
            V(
                D.POSTGRESQL,
                lambda left, right: ((sa.func.extract("epoch", left) - sa.func.extract("epoch", right)) / DAY_SEC),
            ),
        ]
    ),
    base.BinaryMinusGenericDatetimes(
        variants=[
            V(D.POSTGRESQL, lambda left, right: sa.func.extract("epoch", left - right) / DAY_SEC),
        ]
    ),
    # /
    base.BinaryDivInt(
        variants=[
            # Force safe division in COMPENG to make the behavior as permissive as possible
            V(D.COMPENG, lambda x, y: sa.cast(x, sa_postgresql.DOUBLE_PRECISION) / sa.func.nullif(y, 0)),
        ]
    ),
    base.BinaryDivFloat(
        variants=[
            # Force safe division in COMPENG to make the behavior as permissive as possible
            V(D.COMPENG, lambda x, y: x / sa.func.nullif(y, 0)),
        ]
    ),
    base.BinaryDivInt(
        variants=[
            V(D.NON_COMPENG_POSTGRESQL, lambda x, y: sa.cast(x, sa_postgresql.DOUBLE_PRECISION) / y),
        ]
    ),
    base.BinaryDivFloat.for_dialect(D.NON_COMPENG_POSTGRESQL),
    # <
    base.BinaryLessThan.for_dialect(D.POSTGRESQL),
    # <=
    base.BinaryLessThanOrEqual.for_dialect(D.POSTGRESQL),
    # ==
    base.BinaryEqual.for_dialect(D.POSTGRESQL),
    # >
    base.BinaryGreaterThan.for_dialect(D.POSTGRESQL),
    # >=
    base.BinaryGreaterThanOrEqual.for_dialect(D.POSTGRESQL),
    # ^
    base.BinaryPower.for_dialect(D.POSTGRESQL),
    # _!=
    base.BinaryNotEqualInternal.for_dialect(D.POSTGRESQL),
    # _==
    base.BinaryEqualInternal.for_dialect(D.POSTGRESQL),
    # _dneq
    base.BinaryEqualDenullified.for_dialect(D.POSTGRESQL),
    # and
    base.BinaryAnd.for_dialect(D.POSTGRESQL),
    # in
    base.BinaryIn.for_dialect(D.POSTGRESQL),
    # like
    base.BinaryLike.for_dialect(D.POSTGRESQL),
    # notin
    base.BinaryNotIn.for_dialect(D.POSTGRESQL),
    # notlike
    base.BinaryNotLike.for_dialect(D.POSTGRESQL),
    # or
    base.BinaryOr.for_dialect(D.POSTGRESQL),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/operators_binary.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_ternary as base

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


DEFINITIONS_TERNARY = [
    # between
    base.TernaryBetween.for_dialect(D.POSTGRESQL),
    # notbetween
    base.TernaryNotBetween.for_dialect(D.POSTGRESQL),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/operators_ternary.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_unary as base

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


DEFINITIONS_UNARY = [
    # isfalse
    base.UnaryIsFalseStringGeo.for_dialect(D.POSTGRESQL),
    base.UnaryIsFalseNumbers.for_dialect(D.POSTGRESQL),
    base.UnaryIsFalseDateTime.for_dialect(D.POSTGRESQL),
    base.UnaryIsFalseBoolean.for_dialect(D.POSTGRESQL),
    # istrue
    base.UnaryIsTrueStringGeo.for_dialect(D.POSTGRESQL),
    base.UnaryIsTrueNumbers.for_dialect(D.POSTGRESQL),
    base.UnaryIsTrueDateTime.for_dialect(D.POSTGRESQL),
    base.UnaryIsTrueBoolean.for_dialect(D.POSTGRESQL),
    # neg
    base.UnaryNegate.for_dialect(D.POSTGRESQL),
    # not
    base.UnaryNotBool.for_dialect(D.POSTGRESQL),
    base.UnaryNotNumbers.for_dialect(D.POSTGRESQL),
    base.UnaryNotStringGeo.for_dialect(D.POSTGRESQL),
    base.UnaryNotDateDatetime.for_dialect(D.POSTGRESQL),
]

lib/dl_connector_postgresql/dl_connector_postgresql/formula/definitions/operators_unary.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Type,
    Union,
)

import attr
import sqlalchemy as sa
import sqlalchemy.dialects.postgresql as sa_postgresql
from sqlalchemy.sql.sqltypes import Integer

from dl_formula.connectors.base.literal import (
    Literal,
    Literalizer,
    TypeDefiningCast,
)
from dl_formula.core.dialect import DialectCombo

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


PG_SMALLINT_MIN = -32768
PG_SMALLINT_MAX = 32767
PG_INTEGER_MIN = -2147483648
PG_INTEGER_MAX = 2147483647


def integer_to_pg_sa_type(value: int) -> Type[Integer]:
    # Some operations are not supported on longer ints, so optimize
    if PG_SMALLINT_MIN <= value <= PG_SMALLINT_MAX:
        return sa.SMALLINT
    if PG_INTEGER_MIN <= value <= PG_INTEGER_MAX:
        return sa.INTEGER
    return sa.BIGINT  # int64 in pg


class BasePostgreSQLLiteralizer(Literalizer):
    __slots__ = ()

    def literal_array(self, value: Union[tuple, list], dialect: DialectCombo) -> Literal:
        return sa_postgresql.array(value)


class MainPostgreSQLLiteralizer(BasePostgreSQLLiteralizer):
    __slots__ = ()

    def literal_int(self, value: int, dialect: DialectCombo) -> Literal:
        sa_type = integer_to_pg_sa_type(value)
        return sa.literal(value, sa_type)


class CompengPostgreSQLLiteralizer(BasePostgreSQLLiteralizer):
    __slots__ = ()

    def literal_float(self, value: float, dialect: DialectCombo) -> Literal:
        sa_type = sa_postgresql.DOUBLE_PRECISION
        return TypeDefiningCast(value, sa_type)

    def literal_int(self, value: int, dialect: DialectCombo) -> Literal:
        sa_type = integer_to_pg_sa_type(value)
        return TypeDefiningCast(value, sa_type)

    def literal_bool(self, value: bool, dialect: DialectCombo) -> Literal:
        sa_type = sa.BOOLEAN
        return TypeDefiningCast(value, sa_type)

    def literal_str(self, value: str, dialect: DialectCombo) -> Literal:
        sa_type = sa.TEXT
        return TypeDefiningCast(value, sa_type)


@attr.s(frozen=True)
class GenericPostgreSQLLiteralizer(Literalizer):
    lit_compeng = attr.ib(init=False, factory=CompengPostgreSQLLiteralizer)
    lit_main = attr.ib(init=False, factory=MainPostgreSQLLiteralizer)

    def literal(self, value: Any, dialect: DialectCombo) -> Literal:
        if dialect & D.COMPENG == dialect:
            return self.lit_compeng.literal(value, dialect=dialect)

        return self.lit_main.literal(value, dialect=dialect)

lib/dl_connector_postgresql/dl_connector_postgresql/formula/literal.py
datalens-tech/datalens-backend
import sqlalchemy.dialects.postgresql as sa_postgresql
from sqlalchemy.types import TypeEngine

from dl_formula.connectors.base.type_constructor import DefaultSATypeConstructor
from dl_formula.core.datatype import DataType


class PostgreSQLTypeConstructor(DefaultSATypeConstructor):
    def get_sa_type(self, data_type: DataType) -> TypeEngine:
        type_map: dict[DataType, TypeEngine] = {
            DataType.BOOLEAN: sa_postgresql.BOOLEAN(),
            DataType.UUID: sa_postgresql.UUID(),
        }
        if (type_eng := type_map.get(data_type)) is not None:
            return type_eng
        else:
            return super().get_sa_type(data_type)

lib/dl_connector_postgresql/dl_connector_postgresql/formula/type_constructor.py
datalens-tech/datalens-backend
from dl_formula_ref.texts import StyledDialect

from dl_connector_postgresql.formula.constants import PostgreSQLDialect
from dl_connector_postgresql.formula_ref.i18n import Translatable


HUMAN_DIALECTS = {
    PostgreSQLDialect.POSTGRESQL: StyledDialect(
        "`PostgreSQL`",
        "`PostgreSQL`",
        "`PostgreSQL`",
    ),
    PostgreSQLDialect.POSTGRESQL_9_3: StyledDialect(
        "`PostgreSQL 9.3`",
        "`PostgreSQL`<br/>`9.3`",
        Translatable("`PostgreSQL` version `9.3`"),
    ),
    PostgreSQLDialect.POSTGRESQL_9_4: StyledDialect(
        "`PostgreSQL 9.4`",
        "`PostgreSQL`<br/>`9.4`",
        Translatable("`PostgreSQL` version `9.4`"),
    ),
}

lib/dl_connector_postgresql/dl_connector_postgresql/formula_ref/human_dialects.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_postgresql as package


DOMAIN = f"dl_formula_ref_{package.__name__}"

_LOCALE_DIR = os.path.join(os.path.dirname(__file__), "..", "locales")

CONFIGS = [
    TranslationConfig(path=_LOCALE_DIR, domain=DOMAIN, locale="en"),
    TranslationConfig(path=_LOCALE_DIR, domain=DOMAIN, locale="ru"),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_postgresql/dl_connector_postgresql/formula_ref/i18n.py
datalens-tech/datalens-backend
from dl_formula_ref.functions.type_conversion import DbCastExtension
from dl_formula_ref.plugins.base.plugin import FormulaRefPlugin

from dl_connector_postgresql.formula.constants import PostgreSQLDialect
from dl_connector_postgresql.formula.definitions.functions_type import FuncDbCastPostgreSQLBase
from dl_connector_postgresql.formula_ref.human_dialects import HUMAN_DIALECTS
from dl_connector_postgresql.formula_ref.i18n import (
    CONFIGS,
    Translatable,
)


class PostgresSQLFormulaRefPlugin(FormulaRefPlugin):
    any_dialects = frozenset((*PostgreSQLDialect.NON_COMPENG_POSTGRESQL.to_list(),))
    compeng_support_dialects = frozenset((*PostgreSQLDialect.POSTGRESQL.to_list(),))  # FIXME: NON_COMPENG_ ?
    human_dialects = HUMAN_DIALECTS
    translation_configs = frozenset(CONFIGS)
    db_cast_extension = DbCastExtension(
        type_whitelists={
            dialect: whitelist
            for dialect, whitelist in FuncDbCastPostgreSQLBase.WHITELISTS.items()
            if dialect != PostgreSQLDialect.COMPENG
        },
        type_comments={
            (PostgreSQLDialect.NON_COMPENG_POSTGRESQL, "character"): Translatable("Alias: `char`"),
            (PostgreSQLDialect.NON_COMPENG_POSTGRESQL, "character varying"): Translatable("Alias: `varchar`"),
            (PostgreSQLDialect.NON_COMPENG_POSTGRESQL, "char"): Translatable("Alias for `character`"),
            (PostgreSQLDialect.NON_COMPENG_POSTGRESQL, "varchar"): Translatable("Alias for `character varying`"),
        },
    )

lib/dl_connector_postgresql/dl_connector_postgresql/formula_ref/plugin.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_api_lib_testing.data_api_base import StandardizedDataApiTestBase
from dl_api_lib_testing.dataset_base import DatasetTestBase
from dl_constants.enums import RawSQLLevel

from dl_connector_postgresql.core.postgresql.constants import (
    CONNECTION_TYPE_POSTGRES,
    SOURCE_TYPE_PG_TABLE,
)
from dl_connector_postgresql_tests.db.config import (
    API_TEST_CONFIG,
    CoreConnectionSettings,
)
from dl_connector_postgresql_tests.db.core.base import BasePostgreSQLTestClass


class PostgreSQLConnectionTestBase(BasePostgreSQLTestClass, ConnectionTestBase):
    conn_type = CONNECTION_TYPE_POSTGRES
    bi_compeng_pg_on = False

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def connection_params(self) -> dict:
        return dict(
            db_name=CoreConnectionSettings.DB_NAME,
            host=CoreConnectionSettings.HOST,
            port=CoreConnectionSettings.PORT,
            username=CoreConnectionSettings.USERNAME,
            password=CoreConnectionSettings.PASSWORD,
            **(dict(raw_sql_level=self.raw_sql_level.value) if self.raw_sql_level is not None else {}),
        )


class PostgreSQLDashSQLConnectionTest(PostgreSQLConnectionTestBase):
    raw_sql_level = RawSQLLevel.dashsql


class PostgreSQLDatasetTestBase(PostgreSQLConnectionTestBase, DatasetTestBase):
    @pytest.fixture(scope="class")
    def dataset_params(self, sample_table) -> dict:
        return dict(
            source_type=SOURCE_TYPE_PG_TABLE.name,
            parameters=dict(
                schema_name=sample_table.schema,
                table_name=sample_table.name,
            ),
        )


class PostgreSQLDataApiTestBase(PostgreSQLDatasetTestBase, StandardizedDataApiTestBase):
    mutation_caches_on = False

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/api/base.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.complex_queries import DefaultBasicComplexQueryTestSuite

from dl_connector_postgresql_tests.db.api.base import PostgreSQLDataApiTestBase


class TestPostgreSQLBasicComplexQueries(PostgreSQLDataApiTestBase, DefaultBasicComplexQueryTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/api/test_complex_queries.py
datalens-tech/datalens-backend
from dl_api_client.dsmaker.api.http_sync_base import SyncHttpClientBase
from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite

from dl_connector_postgresql_tests.db.api.base import PostgreSQLConnectionTestBase


class TestPostgreSQLConnection(PostgreSQLConnectionTestBase, DefaultConnectorConnectionTestSuite):
    def test_create_connection(self, control_api_sync_client: SyncHttpClientBase, saved_connection_id: str) -> None:
        assert saved_connection_id
        response = control_api_sync_client.get(
            url=f"/api/v1/connections/{saved_connection_id}",
        )
        assert response.status_code == 200
        assert response.json["db_type"] == "postgres"
        assert response.json["cache_ttl_sec"] is None

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/api/test_connection.py
datalens-tech/datalens-backend
from aiohttp.test_utils import TestClient
import pytest

from dl_api_lib_testing.connector.dashsql_suite import DefaultDashSQLTestSuite
from dl_testing.test_data.sql_queries import DASHSQL_EXAMPLE_PARAMS

from dl_connector_postgresql_tests.db.api.base import PostgreSQLDashSQLConnectionTest
from dl_connector_postgresql_tests.db.config import (
    DASHSQL_QUERY,
    QUERY_WITH_PARAMS,
    SUBSELECT_QUERY_FULL,
)


class TestPostgresDashSQL(PostgreSQLDashSQLConnectionTest, DefaultDashSQLTestSuite):
    @pytest.mark.asyncio
    async def test_result(self, data_api_lowlevel_aiohttp_client: TestClient, saved_connection_id: str):
        resp = await self.get_dashsql_response(
            data_api_aio=data_api_lowlevel_aiohttp_client,
            conn_id=saved_connection_id,
            query=DASHSQL_QUERY,
        )

        resp_data = await resp.json()
        assert resp_data[0]["event"] == "metadata", resp_data
        assert resp_data[0]["data"]["names"] == ["aa", "bb", "cc", "dd", "ee", "ff"]
        assert resp_data[0]["data"]["driver_types"] == [25, 1007, 23, 1114, 1184, 17]
        assert resp_data[0]["data"]["postgresql_typnames"] == [
            "text",
            "_int4",
            "int4",
            "timestamp",
            "timestamptz",
            "bytea",
        ]
        assert resp_data[0]["data"]["db_types"] == [
            "text",
            "array(integer)",
            "integer",
            "timestamp",
            "timestamp",
            "bytea",
        ]
        assert resp_data[0]["data"]["bi_types"] == [
            "string",
            "array_int",
            "integer",
            "genericdatetime",
            "genericdatetime",
            "unsupported",
        ]

        assert resp_data[-1]["event"] == "footer", resp_data[-1]

    @pytest.mark.asyncio
    async def test_result_with_error(self, data_api_lowlevel_aiohttp_client: TestClient, saved_connection_id: str):
        resp = await self.get_dashsql_response(
            data_api_aio=data_api_lowlevel_aiohttp_client,
            conn_id=saved_connection_id,
            query="select 1/0",
            fail_ok=True,
        )

        assert resp.status == 400
        resp_data = await resp.json()
        assert resp_data["code"] == "ERR.DS_API.DB.ZERO_DIVISION"
        assert resp_data["message"] == "Division by zero."
        assert resp_data["details"] == {"db_message": "division by zero"}

    @pytest.mark.asyncio
    async def test_result_extended(self, data_api_lowlevel_aiohttp_client: TestClient, saved_connection_id: str):
        resp = await self.get_dashsql_response(
            data_api_aio=data_api_lowlevel_aiohttp_client,
            conn_id=saved_connection_id,
            query=SUBSELECT_QUERY_FULL,
        )

        resp_data = await resp.json()
        assert resp_data[0]["data"]["names"] == [
            "number",
            "str",
            "num_bool",
            "num_bytea",
            "num_char",
            "num_int8",
            "num_int2",
            "num_int4",
            "num_text",
            "num_oid",
            "num_json",
            "num_float4",
            "num_float8",
            "num_numeric",
            "num_interval",
            "num_varchar",
            "num_date",
            "num_time",
            "num_timestamp",
            "num_timestamptz",
            "num_array",
            "some_nan",
        ]
        assert resp_data[0]["data"]["bi_types"] == [
            "integer",
            "string",
            "boolean",
            "unsupported",
            "string",
            "integer",
            "integer",
            "integer",
            "string",
            "unsupported",
            "unsupported",
            "float",
            "float",
            "float",
            "unsupported",
            "string",
            "date",
            "unsupported",
            "genericdatetime",
            "genericdatetime",
            "array_int",
            "float",
        ]
        assert resp_data[1]["data"][21] == "nan", "must not be a float('nan')"

    @pytest.mark.asyncio
    async def test_result_with_params(self, data_api_lowlevel_aiohttp_client: TestClient, saved_connection_id: str):
        resp = await self.get_dashsql_response(
            data_api_aio=data_api_lowlevel_aiohttp_client,
            conn_id=saved_connection_id,
            query=QUERY_WITH_PARAMS,
            params=DASHSQL_EXAMPLE_PARAMS,
        )

        resp_data = await resp.json()
        assert len(resp_data) == 3, resp_data
        assert resp_data[0]["event"] == "metadata", resp_data
        assert resp_data[0]["data"]["bi_types"] == [
            "string",
            "string",
            "string",
            "integer",
            "float",
            "boolean",
            "boolean",
            "date",
            "genericdatetime",
            "string",
            "string",
            "boolean",
            "boolean",
            "boolean",
            "boolean",
            "integer",
        ], resp_data[0]["data"]
        assert resp_data[1]["event"] == "row"
        assert resp_data[1]["data"] == [
            "normal ':string'",
            "extended:string\nwith\nnewlines",
            "some\\:string\nwith\\stuff",
            562949953421312,
            "73786976294838210000",
            True,
            False,
            "2021-07-19",
            "2021-07-19T19:35:43",
            "11",
            "22",
            True,
            False,
            True,
            False,
            1,
        ], resp_data[1]["data"]
        assert resp_data[2]["event"] == "footer"

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/api/test_dashsql.py
datalens-tech/datalens-backend
import pytest
import sqlalchemy as sa

from dl_api_client.dsmaker.api.data_api import SyncHttpDataApiV2
from dl_api_client.dsmaker.api.dataset_api import SyncHttpDatasetApiV1
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_client.dsmaker.shortcuts.result_data import get_data_rows
from dl_api_lib_testing.connector.data_api_suites import (
    DefaultConnectorDataDistinctTestSuite,
    DefaultConnectorDataGroupByFormulaTestSuite,
    DefaultConnectorDataPreviewTestSuite,
    DefaultConnectorDataRangeTestSuite,
    DefaultConnectorDataResultTestSuite,
)
from dl_api_lib_testing.data_api_base import DataApiTestParams
from dl_constants.enums import (
    UserDataType,
    WhereClauseOperation,
)
from dl_core_testing.database import (
    C,
    Db,
    DbTable,
    make_pg_table_with_enums,
    make_table,
)
from dl_core_testing.dataset import data_source_settings_from_table
from dl_sqlalchemy_postgres.base import CITEXT
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_postgresql_tests.db.api.base import PostgreSQLDataApiTestBase


class TestPostgreSQLDataResult(PostgreSQLDataApiTestBase, DefaultConnectorDataResultTestSuite):
    def test_isnull(
        self, saved_dataset: Dataset, data_api_test_params: DataApiTestParams, data_api: SyncHttpDataApiV2
    ) -> None:
        """
        Postgres has a specific implementation for the ISNULL function
        with constant literals, we basically check that it just works
        """
        ds = saved_dataset
        ds.result_schema["is_null"] = ds.field(formula="ISNULL('3')")

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title=data_api_test_params.two_dims[0]),
                ds.find_field(title="is_null"),
            ],
            filters=[
                ds.find_field(title="is_null").filter(
                    op=WhereClauseOperation.EQ,
                    values=[False],
                ),
            ],
        )

        data_rows = get_data_rows(result_resp)
        assert data_rows

    def test_enums(
        self,
        db: Db,
        saved_connection_id: str,
        dataset_params: dict,
        control_api: SyncHttpDatasetApiV1,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        db_table = make_pg_table_with_enums(db)
        params = self.get_dataset_params(dataset_params, db_table)
        ds = self.make_basic_dataset(control_api, connection_id=saved_connection_id, dataset_params=params)

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="col1"),
                ds.find_field(title="col2"),
                ds.find_field(title="col3"),
            ],
            filters=[
                ds.find_field(title="col1").filter(WhereClauseOperation.IN, values=["var1"]),
                ds.find_field(title="col2").filter(WhereClauseOperation.IN, values=["str2"]),
                ds.find_field(title="col3").filter(WhereClauseOperation.IN, values=["2"]),
            ],
        )

        data_rows = get_data_rows(result_resp)
        assert data_rows == [["var1", "str2", "2"]]

    @pytest.fixture(scope="function")
    def enabled_citext_extension(self, db: Db) -> None:
        db.execute("CREATE EXTENSION IF NOT EXISTS CITEXT;")

    def test_pg_citext_column(
        self,
        control_api: SyncHttpDatasetApiV1,
        data_api: SyncHttpDataApiV2,
        db: Db,
        saved_connection_id: str,
        enabled_citext_extension,
    ) -> None:
        table = db.table_from_columns(
            [sa.Column(name="citext_value", type_=CITEXT)],
        )
        db.create_table(table)
        db_table = DbTable(db=db, table=table)  # type: ignore  # TODO: fix
        db_table.insert(
            [
                {"citext_value": "var1"},
                {"citext_value": "var2"},
                {"citext_value": "var3"},
            ]
        )
        ds = Dataset()
        ds.sources["source_1"] = ds.source(
            connection_id=saved_connection_id, **data_source_settings_from_table(db_table)
        )
        ds.source_avatars["avatar_1"] = ds.sources["source_1"].avatar()
        ds = control_api.apply_updates(dataset=ds).dataset
        ds = control_api.save_dataset(ds).dataset

        result_resp = data_api.get_result(
            dataset=ds,
            fields=[
                ds.find_field(title="citext_value"),
            ],
        )

        data_rows = get_data_rows(result_resp)
        assert len(data_rows) == 3
        assert list(sorted(data_rows)) == [["var1"], ["var2"], ["var3"]]


class TestPostgreSQLDataGroupBy(PostgreSQLDataApiTestBase, DefaultConnectorDataGroupByFormulaTestSuite):
    pass


class TestPostgreSQLDataRange(PostgreSQLDataApiTestBase, DefaultConnectorDataRangeTestSuite):
    pass


class TestPostgreSQLDataDistinct(PostgreSQLDataApiTestBase, DefaultConnectorDataDistinctTestSuite):
    def test_bigint_literal(
        self,
        db: Db,
        saved_connection_id: str,
        dataset_params: dict,
        control_api: SyncHttpDatasetApiV1,
        data_api: SyncHttpDataApiV2,
    ) -> None:
        columns = [
            C(name="bigint_value", user_type=UserDataType.integer, vg=lambda rn, **kwargs: 10000002877 + rn),
        ]
        db_table = make_table(db, columns=columns)
        params = self.get_dataset_params(dataset_params, db_table)
        ds = self.make_basic_dataset(control_api, connection_id=saved_connection_id, dataset_params=params)

        lower_bound = 10000002879
        distinct_resp = data_api.get_distinct(
            dataset=ds,
            field=ds.find_field(title="bigint_value"),
            filters=[
                ds.find_field(title="bigint_value").filter(
                    op=WhereClauseOperation.GTE,
                    values=[str(lower_bound)],
                ),
            ],
        )

        data_rows = get_data_rows(distinct_resp)
        assert data_rows
        assert all(int(row[0]) >= lower_bound for row in data_rows)


class TestPostgreSQLDataPreview(PostgreSQLDataApiTestBase, DefaultConnectorDataPreviewTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/api/test_data.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite

from dl_connector_postgresql_tests.db.api.base import PostgreSQLDatasetTestBase


class TestPostgreSQLDataset(PostgreSQLDatasetTestBase, DefaultConnectorDatasetTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/api/test_dataset.py
datalens-tech/datalens-backend
import os
from typing import ClassVar

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


# Infra settings
CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=52311).host,
    port_us_http=get_test_container_hostport("us", fallback_port=52311).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=52310).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=52310).port,
    us_master_token="AC1ofiek8coB",
    core_connector_ep_names=["postgresql"],
)


def get_postgres_ssl_ca() -> str:
    path = os.path.join(os.path.dirname(__file__), "../../docker-compose/db-postgres/ssl", "root.crt")

    with open(path) as f:
        return f.read()


class CoreConnectionSettings:
    DB_NAME: ClassVar[str] = "test_data"
    HOST: ClassVar[str] = get_test_container_hostport("db-postgres-13", fallback_port=52301).host
    PORT: ClassVar[int] = get_test_container_hostport("db-postgres-13", fallback_port=52301).port
    USERNAME: ClassVar[str] = "datalens"
    PASSWORD: ClassVar[str] = "qwerty"


class CoreSslConnectionSettings:
    DB_NAME: ClassVar[str] = "test_data"
    HOST: ClassVar[str] = "127.0.0.1"
    PORT: ClassVar[int] = 52303
    USERNAME: ClassVar[str] = "datalens"
    PASSWORD: ClassVar[str] = "qwerty"


SUBSELECT_QUERY_FULL = r"""
with base as (
    select generate_series(0, 6) as number
)
select
    number,
    'test' || number::text as str,
    number::bool as num_bool,
    number::text::bytea as num_bytea,
    number::char as num_char,
    -- number::name as num_name,
    number::int8 as num_int8,
    number::int2 as num_int2,
    number::int4 as num_int4,
    number::text as num_text,
    number::oid as num_oid,
    number::text::json as num_json,
    -- number::text::jsonb as num_jsonb, -- tests-pg too old
    number::float4 as num_float4,
    number::float8 as num_float8,
    number::numeric as num_numeric,
    (number || ' second')::interval as num_interval,
    number::varchar(12) as num_varchar,
    ('2020-01-0' || (number + 1))::date as num_date,
    ('00:01:0' || number)::time as num_time,
    ('2020-01-01T00:00:0' || number)::timestamp as num_timestamp,
    ('2020-01-01T00:00:0' || number)::timestamptz as num_timestamptz,
    ARRAY[number, 2, 3] as num_array,
    'nan'::double precision + number as some_nan
from base
limit 10
"""
DASHSQL_QUERY = """select
    '1' as aa,
    ARRAY[2, 3] as bb,
    unnest(ARRAY[4, 5]) as cc,
    '2020-01-01 01:02:03'::timestamp as dd,
    '2020-01-01 02:03:0'::timestamptz as ee,
    'zxc'::bytea as ff
"""
QUERY_WITH_PARAMS = r"""
select
  'normal '':string'''::text as v1_normal_string,
  E'extended:string\nwith\nnewlines'::text as v2_ext_string,
  {{some_string}}::text as v3_param_string,
  {{some_integer}} as v4_param_integer,
  {{some_float}} as v5_param_float,
  {{some_boolean}} as v6_param_boolean,
  {{some_other_boolean}} as v7_param_boolean,
  {{some_date}} as v8_param_date,
  {{some_datetime}} as v9_param_datetime,
  {{3xtr4 ше1гd param}}::text as v10_weird_name,
  {{3xtr4 же1гd param}}::text as v11_weird_name,
  1 in {{intvalues}} as v12_int_in,
  0 in {{intvalues}} as v13_int_in_2,
  'a' in {{strvalues}} as v14_str_in,
  'z' in {{strvalues}} as v14_str_in_2,
  1 as stuff
"""

DB_URLS = {
    D.POSTGRESQL_9_3: (
        f'bi_postgresql://datalens:qwerty@{get_test_container_hostport("db-postgres-9-3", fallback_port=52300).as_pair()}/test_data'
    ),
    D.POSTGRESQL_9_4: (
        f'bi_postgresql://datalens:qwerty@{get_test_container_hostport("db-postgres-13", fallback_port=52301).as_pair()}/test_data'
    ),
    D.COMPENG: (
        f'bi_postgresql://datalens:qwerty@{get_test_container_hostport("db-postgres-13", fallback_port=52301).as_pair()}/test_data'
    ),
}
DB_CORE_URL = DB_URLS[D.POSTGRESQL_9_4]
DB_CORE_SSL_URL = "bi_postgresql://datalens:qwerty@localhost:52303/test_data"

API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["postgresql"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
)

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/config.py
datalens-tech/datalens-backend
from dl_api_lib_testing.initialization import initialize_api_lib_test
from dl_formula_testing.forced_literal import forced_literal_use

from dl_connector_postgresql_tests.db.config import API_TEST_CONFIG


pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)


__all__ = (
    # auto-use fixtures:
    "forced_literal_use",
)

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/conftest.py
datalens-tech/datalens-backend
import asyncio
import os
from typing import Generator

import pytest
import shortuuid

import dl_configs.utils as bi_configs_utils
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import Db
from dl_core_testing.testcases.connection import BaseConnectionTestClass

from dl_connector_postgresql.core.postgresql.constants import CONNECTION_TYPE_POSTGRES
from dl_connector_postgresql.core.postgresql.testing.connection import make_postgresql_saved_connection
from dl_connector_postgresql.core.postgresql.us_connection import ConnectionPostgreSQL
import dl_connector_postgresql_tests.db.config as test_config


class BasePostgreSQLTestClass(BaseConnectionTestClass[ConnectionPostgreSQL]):
    conn_type = CONNECTION_TYPE_POSTGRES
    core_test_config = test_config.CORE_TEST_CONFIG

    @pytest.fixture(autouse=True)
    # FIXME: This fixture is a temporary solution for failing core tests when they are run together with api tests
    def loop(self, event_loop: asyncio.AbstractEventLoop) -> Generator[asyncio.AbstractEventLoop, None, None]:
        asyncio.set_event_loop(event_loop)
        yield event_loop
        # Attempt to cover an old version of pytest-asyncio:
        # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
        asyncio.set_event_loop_policy(None)

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return test_config.DB_CORE_URL

    @pytest.fixture(scope="function")
    def connection_creation_params(self) -> dict:
        return dict(
            db_name=test_config.CoreConnectionSettings.DB_NAME,
            host=test_config.CoreConnectionSettings.HOST,
            port=test_config.CoreConnectionSettings.PORT,
            username=test_config.CoreConnectionSettings.USERNAME,
            password=test_config.CoreConnectionSettings.PASSWORD,
            **(dict(raw_sql_level=self.raw_sql_level) if self.raw_sql_level is not None else {}),
        )

    @pytest.fixture(scope="function")
    def saved_connection(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params: dict,
    ) -> ConnectionPostgreSQL:
        conn = make_postgresql_saved_connection(
            sync_usm=sync_us_manager,
            **connection_creation_params,
        )
        return conn

    @pytest.fixture(scope="class")
    def pg_partitioned_table_name(self, db: Db) -> str:
        name = f"test_partitioned_table_{shortuuid.uuid().lower()}"
        queries = [
            f"""
                create table {name}
                (ts timestamptz not null, value text)
                partition by range(ts)
            """,
            f"""
                create table {name}_01
                partition of {name}
                for values from ('2020-01-01 00:00:00') to ('2021-01-01 00:00:00')
            """,
            f"""
                insert into {name} (ts, value)
                values ('2020-01-01 01:02:03', 'a')
            """,
        ]
        for query in queries:
            db.execute(query)
        yield name
        db.execute(f"drop table if exists {name}")


class BaseSslPostgreSQLTestClass(BasePostgreSQLTestClass):
    @pytest.fixture(autouse=True)
    def clear_ssl_folder(self):
        folder_path = bi_configs_utils.get_temp_root_certificates_folder_path()

        for filename in os.listdir(folder_path):
            file_path = os.path.join(folder_path, filename)
            os.unlink(file_path)

    def check_ssl_directory_is_empty(self) -> None:
        assert not os.listdir(bi_configs_utils.get_temp_root_certificates_folder_path())

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return test_config.DB_CORE_SSL_URL

    @pytest.fixture(scope="function")
    def connection_creation_params(self) -> dict:
        return dict(
            db_name=test_config.CoreSslConnectionSettings.DB_NAME,
            host=test_config.CoreSslConnectionSettings.HOST,
            port=test_config.CoreSslConnectionSettings.PORT,
            username=test_config.CoreSslConnectionSettings.USERNAME,
            password=test_config.CoreSslConnectionSettings.PASSWORD,
            **(dict(raw_sql_level=self.raw_sql_level) if self.raw_sql_level is not None else {}),
            ssl_enable=True,
            ssl_ca=test_config.get_postgres_ssl_ca(),
        )

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/core/base.py
datalens-tech/datalens-backend
from dl_core_testing.testcases.adapter import BaseAsyncAdapterTestClass
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_postgresql.core.postgresql_base.async_adapters_postgres import AsyncPostgresAdapter
from dl_connector_postgresql.core.postgresql_base.target_dto import PostgresConnTargetDTO
from dl_connector_postgresql_tests.db.core.base import BasePostgreSQLTestClass


class TestAsyncPostgreSQLAdapter(
    BasePostgreSQLTestClass,
    BaseAsyncAdapterTestClass[PostgresConnTargetDTO],
):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            BaseAsyncAdapterTestClass.test_default_pass_db_query_to_user: "Not relevant",
        },
    )

    ASYNC_ADAPTER_CLS = AsyncPostgresAdapter

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/core/test_adapter.py
datalens-tech/datalens-backend
from __future__ import annotations

import os
from typing import TypeVar

import pytest

from dl_core.us_connection_base import (
    ConnectionSQL,
    DataSourceTemplate,
)
from dl_core_testing.testcases.connection import DefaultConnectionTestClass

from dl_connector_postgresql.core.postgresql.us_connection import ConnectionPostgreSQL
from dl_connector_postgresql_tests.db.core.base import (
    BasePostgreSQLTestClass,
    BaseSslPostgreSQLTestClass,
)


_CONN_TV = TypeVar("_CONN_TV", bound=ConnectionSQL)


class TestPostgreSQLConnection(
    BasePostgreSQLTestClass,
    DefaultConnectionTestClass[ConnectionPostgreSQL],
):
    do_check_data_export_flag = True

    def check_saved_connection(self, conn: ConnectionPostgreSQL, params: dict) -> None:
        assert conn.uuid is not None
        assert conn.data.db_name == params["db_name"]
        assert conn.data.host == params["host"]
        assert conn.data.port == params["port"]
        assert conn.data.username == params["username"]
        assert conn.data.password == params["password"]
        assert conn.data.ssl_enable is False
        assert conn.data.ssl_ca is None

    def check_data_source_templates(self, conn: ConnectionPostgreSQL, dsrc_templates: list[DataSourceTemplate]) -> None:
        assert dsrc_templates
        for dsrc_tmpl in dsrc_templates:
            assert dsrc_tmpl.title
            if dsrc_tmpl.parameters.get("schema_name") is not None:
                assert dsrc_tmpl.group == [dsrc_tmpl.parameters["schema_name"]]

    def test_get_data_source_templates_pg_partitioned(
        self, saved_connection, pg_partitioned_table_name, conn_default_service_registry
    ):
        connection = saved_connection
        service_registry = conn_default_service_registry
        templates = connection.get_data_source_templates(
            conn_executor_factory=service_registry.get_conn_executor_factory().get_sync_conn_executor,
        )
        names = [tpl.title for tpl in templates]
        assert pg_partitioned_table_name in names
        # # To reconsider later:
        # assert pg_partitioned_table + '_01' not in names


@pytest.mark.skipif(os.environ.get("WE_ARE_IN_CI"), reason="can't use localhost")
class TestSslPostgreSQLConnection(
    BaseSslPostgreSQLTestClass,
    DefaultConnectionTestClass[ConnectionPostgreSQL],
):
    def check_saved_connection(self, conn: ConnectionPostgreSQL, params: dict) -> None:
        assert conn.uuid is not None
        assert conn.data.db_name == params["db_name"]
        assert conn.data.host == params["host"]
        assert conn.data.port == params["port"]
        assert conn.data.username == params["username"]
        assert conn.data.password == params["password"]
        assert conn.data.ssl_enable is True
        assert conn.data.ssl_ca is params["ssl_ca"]

    def check_data_source_templates(self, conn: ConnectionPostgreSQL, dsrc_templates: list[DataSourceTemplate]) -> None:
        assert dsrc_templates
        for dsrc_tmpl in dsrc_templates:
            assert dsrc_tmpl.title

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/core/test_connection.py
datalens-tech/datalens-backend
import asyncio
import os
from typing import (
    Optional,
    Sequence,
)
import uuid

import pytest
from sqlalchemy.dialects import postgresql as pg_types

from dl_constants.enums import UserDataType
from dl_core.connection_executors import AsyncConnExecutorBase
from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.connection_models.common_models import DBIdent
from dl_core_testing.database import Db
from dl_core_testing.testcases.connection_executor import (
    DefaultAsyncConnectionExecutorTestSuite,
    DefaultSyncAsyncConnectionExecutorCheckBase,
    DefaultSyncConnectionExecutorTestSuite,
)
from dl_sqlalchemy_postgres.base import CITEXT
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_postgresql.core.postgresql.us_connection import ConnectionPostgreSQL
from dl_connector_postgresql_tests.db.config import CoreConnectionSettings
from dl_connector_postgresql_tests.db.core.base import (
    BasePostgreSQLTestClass,
    BaseSslPostgreSQLTestClass,
)


class PostgreSQLSyncAsyncConnectionExecutorCheckBase(
    BasePostgreSQLTestClass,
    DefaultSyncAsyncConnectionExecutorCheckBase[ConnectionPostgreSQL],
):
    test_params = RegulatedTestParams(
        mark_tests_failed={
            DefaultAsyncConnectionExecutorTestSuite.test_closing_sql_sessions: "",  # TODO: FIXME
        },
    )

    @pytest.fixture(scope="function")
    def db_ident(self) -> DBIdent:
        return DBIdent(db_name=CoreConnectionSettings.DB_NAME)

    def check_db_version(self, db_version: Optional[str]) -> None:
        assert db_version is not None
        assert "." in db_version

    @pytest.fixture(scope="function")
    def enabled_citext_extension(self, db: Db) -> None:
        db.execute("CREATE EXTENSION IF NOT EXISTS CITEXT;")


@pytest.mark.usefixtures("enabled_citext_extension")
class TestPostgreSQLSyncConnectionExecutor(
    PostgreSQLSyncAsyncConnectionExecutorCheckBase,
    DefaultSyncConnectionExecutorTestSuite[ConnectionPostgreSQL],
):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultAsyncConnectionExecutorTestSuite.test_table_exists: "",  # TODO: FIXME
        },
    )

    def get_schemas_for_type_recognition(self) -> dict[str, Sequence[DefaultSyncConnectionExecutorTestSuite.CD]]:
        return {
            "types_postgres_number": [
                self.CD(pg_types.SMALLINT(), UserDataType.integer),
                self.CD(pg_types.INTEGER(), UserDataType.integer),
                self.CD(pg_types.BIGINT(), UserDataType.integer),
                self.CD(pg_types.REAL(), UserDataType.float),
                self.CD(pg_types.DOUBLE_PRECISION(), UserDataType.float),
                self.CD(pg_types.NUMERIC(), UserDataType.float),
            ],
            "types_postgres_string": [
                self.CD(pg_types.CHAR(), UserDataType.string),
                self.CD(pg_types.VARCHAR(100), UserDataType.string),
                self.CD(pg_types.TEXT(), UserDataType.string),
                self.CD(CITEXT(), UserDataType.string),
            ],
            "types_postgres_date": [
                self.CD(pg_types.DATE(), UserDataType.date),
                self.CD(pg_types.TIMESTAMP(timezone=False), UserDataType.genericdatetime),
                self.CD(pg_types.TIMESTAMP(timezone=True), UserDataType.genericdatetime),
            ],
            "types_postgres_other": [
                self.CD(pg_types.BOOLEAN(), UserDataType.boolean),
                self.CD(pg_types.ENUM("var1", "var2", name=str(uuid.uuid4())), UserDataType.string),
            ],
        }


class TestPostgreSQLAsyncConnectionExecutor(
    PostgreSQLSyncAsyncConnectionExecutorCheckBase,
    DefaultAsyncConnectionExecutorTestSuite[ConnectionPostgreSQL],
):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultAsyncConnectionExecutorTestSuite.test_table_exists: "Not implemented",
            DefaultAsyncConnectionExecutorTestSuite.test_table_not_exists: "Not implemented",
            DefaultAsyncConnectionExecutorTestSuite.test_get_table_schema_info: "Not implemented",
            DefaultAsyncConnectionExecutorTestSuite.test_get_table_schema_info_for_nonexistent_table: "Not implemented",
        },
    )


@pytest.mark.skipif(os.environ.get("WE_ARE_IN_CI"), reason="can't use localhost")
class TestSslPostgreSQLSyncConnectionExecutor(
    BaseSslPostgreSQLTestClass,
    TestPostgreSQLSyncConnectionExecutor,
):
    def test_test(self, sync_connection_executor: SyncConnExecutorBase) -> None:
        super().test_test(sync_connection_executor)
        self.check_ssl_directory_is_empty()


@pytest.mark.skipif(os.environ.get("WE_ARE_IN_CI"), reason="can't use localhost")
class TestSslPostgreSQLAsyncConnectionExecutor(
    BaseSslPostgreSQLTestClass,
    TestPostgreSQLAsyncConnectionExecutor,
):
    async def test_test(self, async_connection_executor: AsyncConnExecutorBase) -> None:
        await super().test_test(async_connection_executor)
        await asyncio.sleep(0.1)
        self.check_ssl_directory_is_empty()

    async def test_multiple_connection_test(self, async_connection_executor: AsyncConnExecutorBase) -> None:
        await super().test_multiple_connection_test(async_connection_executor)
        await asyncio.sleep(0.1)
        self.check_ssl_directory_is_empty()

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/core/test_connection_executor.py
datalens-tech/datalens-backend
import pytest

from dl_constants.enums import (
    RawSQLLevel,
    UserDataType,
)
from dl_core.data_source_spec.sql import (
    StandardSchemaSQLDataSourceSpec,
    SubselectDataSourceSpec,
)
from dl_core_testing.fixtures.sample_tables import TABLE_SPEC_SAMPLE_SUPERSTORE
from dl_core_testing.testcases.data_source import (
    DataSourceTestByViewClass,
    DefaultDataSourceTestClass,
)

from dl_connector_postgresql.core.postgresql.constants import (
    SOURCE_TYPE_PG_SUBSELECT,
    SOURCE_TYPE_PG_TABLE,
)
from dl_connector_postgresql.core.postgresql.data_source import (
    PostgreSQLDataSource,
    PostgreSQLSubselectDataSource,
)
from dl_connector_postgresql.core.postgresql.us_connection import ConnectionPostgreSQL
from dl_connector_postgresql_tests.db.config import SUBSELECT_QUERY_FULL
from dl_connector_postgresql_tests.db.core.base import BasePostgreSQLTestClass


class TestPostgreSQLTableDataSource(
    BasePostgreSQLTestClass,
    DefaultDataSourceTestClass[
        ConnectionPostgreSQL,
        StandardSchemaSQLDataSourceSpec,
        PostgreSQLDataSource,
    ],
):
    DSRC_CLS = PostgreSQLDataSource

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self, sample_table) -> StandardSchemaSQLDataSourceSpec:
        dsrc_spec = StandardSchemaSQLDataSourceSpec(
            source_type=SOURCE_TYPE_PG_TABLE,
            db_name=sample_table.db.name,
            schema_name=sample_table.schema,
            table_name=sample_table.name,
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return list(TABLE_SPEC_SAMPLE_SUPERSTORE.table_schema)


class TestPostgreSQLSubselectDataSource(
    BasePostgreSQLTestClass,
    DefaultDataSourceTestClass[
        ConnectionPostgreSQL,
        SubselectDataSourceSpec,
        PostgreSQLSubselectDataSource,
    ],
):
    DSRC_CLS = PostgreSQLSubselectDataSource

    raw_sql_level = RawSQLLevel.subselect

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self, sample_table) -> SubselectDataSourceSpec:
        dsrc_spec = SubselectDataSourceSpec(
            source_type=SOURCE_TYPE_PG_SUBSELECT,
            subsql=f'SELECT * FROM "{sample_table.name}"',
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return list(TABLE_SPEC_SAMPLE_SUPERSTORE.table_schema)


class TestPostgreSQLSubselectByView(
    BasePostgreSQLTestClass,
    DataSourceTestByViewClass[
        ConnectionPostgreSQL,
        SubselectDataSourceSpec,
        PostgreSQLSubselectDataSource,
    ],
):
    DSRC_CLS = PostgreSQLSubselectDataSource

    raw_sql_level = RawSQLLevel.subselect

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self) -> SubselectDataSourceSpec:
        dsrc_spec = SubselectDataSourceSpec(
            source_type=SOURCE_TYPE_PG_SUBSELECT,
            subsql=SUBSELECT_QUERY_FULL,
        )
        return dsrc_spec

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/core/test_data_source.py
datalens-tech/datalens-backend
from dl_core_testing.testcases.dataset import DefaultDatasetTestSuite
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_postgresql.core.postgresql.constants import SOURCE_TYPE_PG_TABLE
from dl_connector_postgresql.core.postgresql.us_connection import ConnectionPostgreSQL
from dl_connector_postgresql_tests.db.core.base import BasePostgreSQLTestClass


class TestPostgreSQLDataset(BasePostgreSQLTestClass, DefaultDatasetTestSuite[ConnectionPostgreSQL]):
    source_type = SOURCE_TYPE_PG_TABLE

    test_params = RegulatedTestParams(
        mark_tests_failed={
            DefaultDatasetTestSuite.test_get_param_hash: "db_name in dsrc",  # TODO: FIXME
        },
    )

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/core/test_dataset.py
datalens-tech/datalens-backend
import pytest

from dl_formula_testing.testcases.base import FormulaConnectorTestBase

from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D
from dl_connector_postgresql_tests.db.config import DB_URLS


class PostgreSQLTestBase(FormulaConnectorTestBase):
    supports_arrays = True
    supports_uuid = True

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return DB_URLS[self.dialect]


class PostgreSQL_9_3TestBase(PostgreSQLTestBase):
    dialect = D.POSTGRESQL_9_3


class PostgreSQL_9_4TestBase(PostgreSQLTestBase):
    dialect = D.POSTGRESQL_9_4


class CompengTestBase(PostgreSQLTestBase):
    dialect = D.COMPENG

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/base.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.conditional_blocks import DefaultConditionalBlockFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class TestConditionalBlockPostgreSQL_9_3(PostgreSQL_9_3TestBase, DefaultConditionalBlockFormulaConnectorTestSuite):
    pass


class TestConditionalBlockPostgreSQL_9_4(PostgreSQL_9_4TestBase, DefaultConditionalBlockFormulaConnectorTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_conditional_blocks.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_aggregation import DefaultMainAggFunctionFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class MainAggFunctionPostgreTestSuite(DefaultMainAggFunctionFormulaConnectorTestSuite):
    supports_all_concat = True


class TestMainAggFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, MainAggFunctionPostgreTestSuite):
    pass


class TestMainAggFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, MainAggFunctionPostgreTestSuite):
    supports_quantile = True
    supports_median = True

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_functions_aggregation.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.functions_array import DefaultArrayFunctionFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class ArrayFunctionPostgreTestSuite(DefaultArrayFunctionFormulaConnectorTestSuite):
    make_decimal_cast = "numeric"
    make_float_cast = "double precision"

    def test_startswith_string_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("STARTSWITH([arr_str_value], [arr_str_value])", from_=data_table)
        assert not dbe.eval(
            'STARTSWITH([arr_str_value], DB_CAST(ARRAY("", "cde", NULL), "varchar[]"))', from_=data_table
        )

    def test_array_contains_all_string_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval('CONTAINS_ALL([arr_str_value], DB_CAST(ARRAY("cde"), "varchar[]"))', from_=data_table)
        assert dbe.eval('CONTAINS_ALL([arr_str_value], DB_CAST(ARRAY("cde", NULL), "varchar[]"))', from_=data_table)
        assert not dbe.eval('CONTAINS_ALL(DB_CAST(ARRAY("cde"), "varchar[]"), [arr_str_value])', from_=data_table)

    def test_array_contains_any_string_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval('CONTAINS_ANY([arr_str_value], DB_CAST(ARRAY("cde"), "varchar[]"))', from_=data_table)
        assert dbe.eval('CONTAINS_ANY([arr_str_value], DB_CAST(ARRAY("123", NULL), "varchar[]"))', from_=data_table)
        assert dbe.eval('CONTAINS_ANY(DB_CAST(ARRAY("cde"), "varchar[]"), [arr_str_value])', from_=data_table)


class TestArrayFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, ArrayFunctionPostgreTestSuite):
    pass


class TestArrayFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, ArrayFunctionPostgreTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_functions_array.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_datetime import DefaultDateTimeFunctionFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class DateTimeFunctionPostgreSQLTestSuite(DefaultDateTimeFunctionFormulaConnectorTestSuite):
    supports_deprecated_dateadd = True
    supports_deprecated_datepart_2 = True
    supports_datetimetz = True


class TestDateTimeFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, DateTimeFunctionPostgreSQLTestSuite):
    pass


class TestDateTimeFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, DateTimeFunctionPostgreSQLTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_functions_datetime.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_logical import DefaultLogicalFunctionFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class LogicalFunctionPostgreTestSuite(DefaultLogicalFunctionFormulaConnectorTestSuite):
    supports_nan_funcs = True
    supports_iif = True


class TestLogicalFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, LogicalFunctionPostgreTestSuite):
    pass


class TestLogicalFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, LogicalFunctionPostgreTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_functions_logical.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_markup import DefaultMarkupFunctionFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class TestMarkupFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, DefaultMarkupFunctionFormulaConnectorTestSuite):
    pass


class TestMarkupFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, DefaultMarkupFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_functions_markup.py
datalens-tech/datalens-backend
import pytest

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.functions_math import DefaultMathFunctionFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class PostgreSQLMatchTestMixin:
    def test_log_data_type_compatibility(self, dbe: DbEvaluator) -> None:
        assert dbe.eval('LOG(DB_CAST(4.0, "numeric", 10, 5), DB_CAST(2.0, "double precision"))') == pytest.approx(2)
        assert dbe.eval('LOG(DB_CAST(4.0, "double precision"), DB_CAST(2.0, "numeric", 10, 5))') == pytest.approx(2)


class TestMathFunctionPostgreSQL_9_3(
    PostgreSQL_9_3TestBase,
    DefaultMathFunctionFormulaConnectorTestSuite,
    PostgreSQLMatchTestMixin,
):
    pass


class TestMathFunctionPostgreSQL_9_4(
    PostgreSQL_9_4TestBase,
    DefaultMathFunctionFormulaConnectorTestSuite,
    PostgreSQLMatchTestMixin,
):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_functions_math.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_string import DefaultStringFunctionFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class TestStringFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, DefaultStringFunctionFormulaConnectorTestSuite):
    pass


class TestStringFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, DefaultStringFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_functions_string.py
datalens-tech/datalens-backend
import datetime
from decimal import Decimal

import pytest
import pytz
import sqlalchemy as sa
import sqlalchemy.exc as sa_exc

import dl_formula.core.exc as exc
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.functions_type_conversion import (
    DefaultBoolTypeFunctionFormulaConnectorTestSuite,
    DefaultDateTypeFunctionFormulaConnectorTestSuite,
    DefaultDbCastTypeFunctionFormulaConnectorTestSuite,
    DefaultFloatTypeFunctionFormulaConnectorTestSuite,
    DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite,
    DefaultGeopointTypeFunctionFormulaConnectorTestSuite,
    DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite,
    DefaultIntTypeFunctionFormulaConnectorTestSuite,
    DefaultStrTypeFunctionFormulaConnectorTestSuite,
    DefaultTreeTypeFunctionFormulaConnectorTestSuite,
)
from dl_formula_testing.util import (
    as_tz,
    dt_strip,
    to_str,
)

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


# STR


class StrTypeFunctionPostgreSQLTestSuite(DefaultStrTypeFunctionFormulaConnectorTestSuite):
    zero_float_to_str_value = "0.0"
    skip_custom_tz = True

    def test_str_from_array(self, dbe: DbEvaluator, data_table) -> None:
        assert to_str(dbe.eval("STR([arr_int_value])", from_=data_table)) == "{0,23,456,NULL}"
        assert to_str(dbe.eval("STR([arr_float_value])", from_=data_table)) == "{0,45,0.123,NULL}"
        assert to_str(dbe.eval("STR([arr_str_value])", from_=data_table)) == '{"","",cde,NULL}'


class TestStrTypeFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, StrTypeFunctionPostgreSQLTestSuite):
    pass


class TestStrTypeFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, StrTypeFunctionPostgreSQLTestSuite):
    pass


# FLOAT


class TestFloatTypeFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, DefaultFloatTypeFunctionFormulaConnectorTestSuite):
    pass


class TestFloatTypeFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, DefaultFloatTypeFunctionFormulaConnectorTestSuite):
    pass


# BOOL


class TestBoolTypeFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, DefaultBoolTypeFunctionFormulaConnectorTestSuite):
    pass


class TestBoolTypeFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, DefaultBoolTypeFunctionFormulaConnectorTestSuite):
    pass


# INT


class TestIntTypeFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, DefaultIntTypeFunctionFormulaConnectorTestSuite):
    pass


class TestIntTypeFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, DefaultIntTypeFunctionFormulaConnectorTestSuite):
    pass


# DATE


class TestDateTypeFunctionPostgreSQL_9_3(PostgreSQL_9_3TestBase, DefaultDateTypeFunctionFormulaConnectorTestSuite):
    pass


class TestDateTypeFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, DefaultDateTypeFunctionFormulaConnectorTestSuite):
    pass


# GENERICDATETIME (& DATETIME)


class GenericDatetimeTypeFunctionPostgreSQLTestSuite(
    DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite,
):
    @pytest.mark.parametrize("func_name", ("GENERICDATETIME",))
    def test_genericdatetime2_pg(self, dbe: DbEvaluator, func_name: str) -> None:
        mos_tz = "Europe/Moscow"
        dt_naive = datetime.datetime(2019, 1, 3, 2, 4, 5)
        dt_naive_iso_str = dt_naive.isoformat()
        dt_as_mos = as_tz(dt_naive, tzinfo=pytz.timezone(mos_tz))
        dt_as_mos_ts = dt_as_mos.timestamp()

        dt_as_mos_to_utc = dt_as_mos.astimezone(pytz.UTC)
        dt_as_mos_to_utc_stripped = dt_strip(dt_as_mos_to_utc)

        assert dt_strip(dbe.eval(f'{func_name}({int(dt_as_mos_ts)}, "{mos_tz}")')) == dt_naive
        assert dt_strip(dbe.eval(f'{func_name}({float(dt_as_mos_ts)}, "{mos_tz}")')) == dt_naive

        # Notably, the string and genericdatetime are interpreted at the specified timezone
        # (with undefined resolution of ambiguities),
        # the number / datetime is interpreted as UTC

        # from string
        assert dt_strip(dbe.eval(f"{func_name}('{dt_naive_iso_str}', '{mos_tz}')")) == dt_as_mos_to_utc_stripped
        # from literal
        assert dt_strip(dbe.eval(f"{func_name}(#{dt_naive_iso_str}#, '{mos_tz}')")) == dt_as_mos_to_utc_stripped

        # Double-wrap (from string)
        assert dt_strip(dbe.eval(f"{func_name}({func_name}('{dt_naive_iso_str}', '{mos_tz}'), '{mos_tz}')")) == dt_naive


class TestGenericDatetimeTypeFunctionPostgreSQL_9_3(
    PostgreSQL_9_3TestBase,
    GenericDatetimeTypeFunctionPostgreSQLTestSuite,
):
    pass


class TestGenericDatetimeTypeFunctionPostgreSQL_9_4(
    PostgreSQL_9_4TestBase,
    GenericDatetimeTypeFunctionPostgreSQLTestSuite,
):
    pass


# GEOPOINT


class TestGeopointTypeFunctionPostgreSQL_9_3(
    PostgreSQL_9_3TestBase,
    DefaultGeopointTypeFunctionFormulaConnectorTestSuite,
):
    pass


class TestGeopointTypeFunctionPostgreSQL_9_4(
    PostgreSQL_9_4TestBase,
    DefaultGeopointTypeFunctionFormulaConnectorTestSuite,
):
    pass


# GEOPOLYGON


class TestGeopolygonTypeFunctionPostgreSQL_9_3(
    PostgreSQL_9_3TestBase,
    DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite,
):
    pass


class TestGeopolygonTypeFunctionPostgreSQL_9_4(
    PostgreSQL_9_4TestBase,
    DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite,
):
    pass


# DB_CAST


class DbCastTypeFunctionPostgreSQLTestSuite(
    DefaultDbCastTypeFunctionFormulaConnectorTestSuite,
):
    def test_db_cast_postgresql(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        value = dbe.eval("[int_value]", from_=data_table)
        assert dbe.eval('DB_CAST(FLOAT([int_value]), "double precision")', from_=data_table) == pytest.approx(
            float(value)
        )
        assert dbe.eval('DB_CAST(FLOAT([int_value]), "numeric", 5, 0)', from_=data_table) == value
        assert dbe.eval('ROUND(DB_CAST(FLOAT([int_value]), "numeric", 5, 0), 0)', from_=data_table) == value
        with pytest.raises(sa_exc.DatabaseError):
            dbe.eval('ROUND(DB_CAST(FLOAT([int_value]), "double precision"), 0)', from_=data_table)

        with pytest.raises(exc.TranslationError):
            dbe.eval('DB_CAST(FLOAT([int_value]), "double precision", 1)', from_=data_table)

        with pytest.raises(exc.TranslationError):
            dbe.eval('DB_CAST(FLOAT([int_value]), "numeric", 5)', from_=data_table)

        with pytest.raises(exc.TranslationError):
            dbe.eval('DB_CAST(FLOAT([int_value]), "numeric", "5", "3")', from_=data_table)

    def test_db_cast_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        array_int = [0, 1, -2, None]
        array_int_string = ",".join("NULL" if item is None else str(item) for item in array_int)
        assert dbe.eval(f'DB_CAST(ARRAY({array_int_string}), "smallint[]")', from_=data_table) == array_int
        assert dbe.eval(f'DB_CAST(ARRAY({array_int_string}), "integer[]")', from_=data_table) == array_int
        assert dbe.eval(f'DB_CAST(ARRAY({array_int_string}), "bigint[]")', from_=data_table) == array_int
        assert dbe.eval('DB_CAST([arr_int_value], "bigint[]")', from_=data_table) == [0, 23, 456, None]

        array_float = [0.0, 1.0, -2.0, None]
        array_float_string = ",".join("NULL" if item is None else str(item) for item in array_float)
        assert dbe.eval(f'DB_CAST(ARRAY({array_float_string}), "double precision[]")', from_=data_table) == array_float
        assert dbe.eval(f'DB_CAST(ARRAY({array_float_string}), "real[]")', from_=data_table) == array_float
        assert dbe.eval(f'DB_CAST(ARRAY({array_float_string}), "numeric[]", 5, 0)', from_=data_table) == array_float
        assert dbe.eval('DB_CAST([arr_float_value], "numeric[]", 5, 0)', from_=data_table) == [
            Decimal("0"),
            Decimal("45"),
            Decimal("0"),
            None,
        ]

        array_str = ["", "a", "NULL", None]
        array_str_string = ",".join("NULL" if item is None else f"'{item}'" for item in array_str)
        assert dbe.eval(f'DB_CAST(ARRAY({array_str_string}), "text[]")', from_=data_table) == array_str
        assert dbe.eval(f'DB_CAST(ARRAY({array_str_string}), "character varying[]")', from_=data_table) == array_str
        assert dbe.eval(f'DB_CAST(ARRAY({array_str_string}), "varchar[]")', from_=data_table) == array_str
        assert dbe.eval('DB_CAST([arr_str_value], "varchar[]")', from_=data_table) == ["", "", "cde", None]


class TestDbCastTypeFunctionPostgreSQL_9_3(
    PostgreSQL_9_3TestBase,
    DbCastTypeFunctionPostgreSQLTestSuite,
):
    pass


class TestDbCastTypeFunctionPostgreSQL_9_4(
    PostgreSQL_9_4TestBase,
    DbCastTypeFunctionPostgreSQLTestSuite,
):
    pass


# TREE


class TestTreeTypeFunctionPostgreSQL_9_3(
    PostgreSQL_9_3TestBase,
    DefaultTreeTypeFunctionFormulaConnectorTestSuite,
):
    pass


class TestTreeTypeFunctionPostgreSQL_9_4(
    PostgreSQL_9_4TestBase,
    DefaultTreeTypeFunctionFormulaConnectorTestSuite,
):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_functions_type_conversion.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_window import DefaultWindowFunctionFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    CompengTestBase,
    PostgreSQL_9_4TestBase,
)


class TestWindowFunctionPostgreSQL_9_4(PostgreSQL_9_4TestBase, DefaultWindowFunctionFormulaConnectorTestSuite):
    pass


class TestWindowFunctionCompeng(CompengTestBase, DefaultWindowFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_functions_window.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.literals import DefaultLiteralFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class ConditionalBlockPostgreSQLTestSuite(DefaultLiteralFormulaConnectorTestSuite):
    supports_microseconds = True
    supports_utc = True
    supports_custom_tz = True


class TestConditionalBlockPostgreSQL_9_3(PostgreSQL_9_3TestBase, ConditionalBlockPostgreSQLTestSuite):
    pass


class TestConditionalBlockPostgreSQL_9_4(PostgreSQL_9_4TestBase, ConditionalBlockPostgreSQLTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_literals.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.misc_funcs import DefaultMiscFunctionalityConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class TestMiscFunctionalityPostgreSQL_9_3(PostgreSQL_9_3TestBase, DefaultMiscFunctionalityConnectorTestSuite):
    pass


class TestMiscFunctionalityPostgreSQL_9_4(PostgreSQL_9_4TestBase, DefaultMiscFunctionalityConnectorTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_misc_funcs.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.operators import DefaultOperatorFormulaConnectorTestSuite

from dl_connector_postgresql_tests.db.formula.base import (
    PostgreSQL_9_3TestBase,
    PostgreSQL_9_4TestBase,
)


class OperatorPostgreSQLTestSuite(DefaultOperatorFormulaConnectorTestSuite):
    subtraction_round_dt = False
    make_float_array_cast = "double precision[]"


class TestOperatorPostgreSQL_9_3(PostgreSQL_9_3TestBase, OperatorPostgreSQLTestSuite):
    pass


class TestOperatorPostgreSQL_9_4(PostgreSQL_9_4TestBase, OperatorPostgreSQLTestSuite):
    pass

lib/dl_connector_postgresql/dl_connector_postgresql_tests/db/formula/test_operators.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.dialect import DefaultDialectFormulaConnectorTestSuite

from dl_connector_postgresql.formula.constants import DIALECT_NAME_POSTGRESQL
from dl_connector_postgresql.formula.constants import PostgreSQLDialect as D


class DialectPostgreSQLTestSuite(DefaultDialectFormulaConnectorTestSuite):
    dialect_name = DIALECT_NAME_POSTGRESQL
    default_dialect = D.POSTGRESQL_9_4
    dialect_matches = (
        ("9.3.4", D.POSTGRESQL_9_3),
        ("9.6.1", D.POSTGRESQL_9_4),
    )

lib/dl_connector_postgresql/dl_connector_postgresql_tests/unit/formula/test_dialect.py
datalens-tech/datalens-backend
from dl_api_connector.i18n.localizer import CONFIGS as BI_API_CONNECTOR_CONFIGS
from dl_api_lib_testing.connection_form_base import ConnectionFormTestBase

from dl_connector_postgresql.api.connection_form.form_config import PostgreSQLConnectionFormFactory
from dl_connector_postgresql.api.i18n.localizer import CONFIGS as BI_CONNECTOR_POSTGRESQL_CONFIGS


class TestPostgresConnectionForm(ConnectionFormTestBase):
    CONN_FORM_FACTORY_CLS = PostgreSQLConnectionFormFactory
    TRANSLATION_CONFIGS = BI_API_CONNECTOR_CONFIGS + BI_CONNECTOR_POSTGRESQL_CONFIGS

lib/dl_connector_postgresql/dl_connector_postgresql_tests/unit/test_connection_form.py
datalens-tech/datalens-backend
from __future__ import annotations

from _socket import gaierror
import psycopg2

from dl_core.exc import SourceHostNotKnownError

from dl_connector_postgresql.core.postgresql_base.error_transformer import (
    make_async_pg_error_transformer,
    sync_pg_db_error_transformer,
)


NAME_OR_SERVICE_NOT_KNOWN_MSG = """
        could not translate host name
        "c-someclusterid.ro.mdb.yandexcloud.net" to address:
        Name or service not known
        """


def test_name_or_service_not_known_sync():
    transformer = sync_pg_db_error_transformer

    parameters = transformer.make_bi_error_parameters(
        wrapper_exc=psycopg2.OperationalError(NAME_OR_SERVICE_NOT_KNOWN_MSG)
    )

    assert parameters[0] == SourceHostNotKnownError


def test_name_or_service_not_known_async():
    transformer = make_async_pg_error_transformer()

    parameters = transformer.make_bi_error_parameters(wrapper_exc=gaierror(NAME_OR_SERVICE_NOT_KNOWN_MSG))

    assert parameters[0] == SourceHostNotKnownError

lib/dl_connector_postgresql/dl_connector_postgresql_tests/unit/test_error_transformer.py
datalens-tech/datalens-backend
import pytest

from dl_constants.enums import UserDataType

from dl_connector_postgresql.core.postgresql_base.type_transformer import PostgreSQLTypeTransformer


@pytest.mark.parametrize("array_type", (UserDataType.array_int, UserDataType.array_float, UserDataType.array_str))
def test_null_array_conversion(array_type):
    tt = PostgreSQLTypeTransformer
    assert tt.cast_for_output(None, user_t=array_type) is None

lib/dl_connector_postgresql/dl_connector_postgresql_tests/unit/test_type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.connection_base import ConnectionMetaMixin
from dl_api_connector.api_schema.connection_base_fields import secret_string_field
from dl_api_connector.api_schema.connection_sql import ClassicSQLConnectionSchema
from dl_api_connector.api_schema.extras import FieldExtra

from dl_connector_promql.core.us_connection import PromQLConnection


class PromQLConnectionSchema(ConnectionMetaMixin, ClassicSQLConnectionSchema):
    TARGET_CLS = PromQLConnection  # type: ignore

    secure = ma_fields.Boolean(attribute="data.secure", bi_extra=FieldExtra(editable=True))
    username = ma_fields.String(
        attribute="data.username",
        required=False,
        allow_none=True,
        dump_default=None,
        bi_extra=FieldExtra(editable=True),
    )
    password = secret_string_field(
        attribute="data.password",
        required=False,
        allow_none=True,
        bi_extra=FieldExtra(editable=True),
    )

lib/dl_connector_promql/dl_connector_promql/api/api_schema/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormActionApiSchema,
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
    FormUIOverride,
)
from dl_api_connector.form_config.models.common import CommonFieldName
import dl_api_connector.form_config.models.rows as C
from dl_api_connector.form_config.models.shortcuts.rows import RowConstructor
from dl_configs.connectors_settings import ConnectorSettingsBase

from dl_connector_promql.api.connection_info import PromQLConnectionInfoProvider


class PromQLConnectionFormFactory(ConnectionFormFactory):
    def get_form_config(
        self,
        connector_settings: Optional[ConnectorSettingsBase],
        tenant: Optional[TenantDef],
    ) -> ConnectionForm:
        rc = RowConstructor(localizer=self._localizer)

        edit_api_schema = FormActionApiSchema(
            items=[
                FormFieldApiSchema(name=CommonFieldName.host, required=True),
                FormFieldApiSchema(name=CommonFieldName.port, required=True),
                FormFieldApiSchema(name=CommonFieldName.username),
                FormFieldApiSchema(name=CommonFieldName.password),
                FormFieldApiSchema(name=CommonFieldName.cache_ttl_sec, nullable=True),
                FormFieldApiSchema(name=CommonFieldName.secure, type="boolean"),
            ]
        )

        create_api_schema = FormActionApiSchema(
            items=[
                *edit_api_schema.items,
                *self._get_top_level_create_api_schema_items(),
            ]
        )

        check_api_schema = FormActionApiSchema(
            items=[
                FormFieldApiSchema(name=CommonFieldName.host, required=True),
                FormFieldApiSchema(name=CommonFieldName.port, required=True),
                FormFieldApiSchema(name=CommonFieldName.username),
                FormFieldApiSchema(name=CommonFieldName.password),
                FormFieldApiSchema(name=CommonFieldName.secure, type="boolean"),
                *self._get_top_level_check_api_schema_items(),
            ]
        )

        return ConnectionForm(
            title=PromQLConnectionInfoProvider.get_title(self._localizer),
            rows=[
                rc.host_row(),
                rc.port_row(),
                rc.username_row(),
                rc.password_row(self.mode),
                C.CacheTTLRow(name=CommonFieldName.cache_ttl_sec),
                C.CustomizableRow(
                    items=[
                        C.CheckboxRowItem(name=CommonFieldName.secure, text="HTTPS", default_value=True),
                    ]
                ),
            ],
            api_schema=FormApiSchema(
                create=create_api_schema if self.mode == ConnectionFormMode.create else None,
                edit=edit_api_schema if self.mode == ConnectionFormMode.edit else None,
                check=check_api_schema,
            ),
            form_ui_override=FormUIOverride(
                show_create_dataset_btn=False,
                show_create_ql_chart_btn=True,
            ),
        )

lib/dl_connector_promql/dl_connector_promql/api/connection_form/form_config.py
datalens-tech/datalens-backend
from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_promql.api.i18n.localizer import Translatable


class PromQLConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-promql")

lib/dl_connector_promql/dl_connector_promql/api/connection_info.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_connector.api_schema.source_base import (
    SimpleDataSourceSchema,
    SimpleDataSourceTemplateSchema,
)
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)

from dl_connector_promql.api.api_schema.connection import PromQLConnectionSchema
from dl_connector_promql.api.connection_form.form_config import PromQLConnectionFormFactory
from dl_connector_promql.api.connection_info import PromQLConnectionInfoProvider
from dl_connector_promql.api.i18n.localizer import CONFIGS
from dl_connector_promql.core.connector import (
    PromQLCoreConnectionDefinition,
    PromQLCoreConnector,
    PromQLCoreSourceDefinition,
)


class PromQLApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = PromQLCoreConnectionDefinition
    api_generic_schema_cls = PromQLConnectionSchema
    form_factory_cls = PromQLConnectionFormFactory
    info_provider_cls = PromQLConnectionInfoProvider


class PromQLApiSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = PromQLCoreSourceDefinition
    api_schema_cls = SimpleDataSourceSchema
    template_api_schema_cls = SimpleDataSourceTemplateSchema


class PromQLApiConnector(ApiConnector):
    core_connector_cls = PromQLCoreConnector
    connection_definitions = (PromQLApiConnectionDefinition,)
    source_definitions = (PromQLApiSourceDefinition,)
    translation_configs = frozenset(CONFIGS)

lib/dl_connector_promql/dl_connector_promql/api/connector.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_promql as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_promql/dl_connector_promql/api/i18n/localizer.py
datalens-tech/datalens-backend
from __future__ import annotations

from datetime import datetime
import logging
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Optional,
    Type,
)
from urllib.parse import (
    quote_plus,
    urljoin,
)

from aiohttp.client import ClientResponse
from aiohttp.helpers import BasicAuth
from aiohttp.web import HTTPBadRequest
import attr
import sqlalchemy as sa

from dl_app_tools.profiling_base import generic_profiler_async
from dl_constants.enums import ConnectionType
from dl_core.connection_executors.adapters.adapters_base_sa_classic import (
    BaseClassicAdapter,
    ClassicSQLConnLineConstructor,
)
from dl_core.connection_executors.adapters.async_adapters_aiohttp import AiohttpDBAdapter
from dl_core.connection_executors.adapters.async_adapters_base import AsyncRawExecutionResult
from dl_core.db.native_type import GenericNativeType
from dl_core.exc import DatabaseQueryError
from dl_core.utils import make_url

from dl_connector_promql.core.constants import CONNECTION_TYPE_PROMQL


if TYPE_CHECKING:
    from dl_constants.types import TBIChunksGen
    from dl_core.connection_executors.models.db_adapter_data import (
        DBAdapterQuery,
        RawSchemaInfo,
    )
    from dl_core.connection_models import (
        DBIdent,
        SchemaIdent,
        TableDefinition,
        TableIdent,
    )

    from dl_connector_promql.core.target_dto import PromQLConnTargetDTO

LOGGER = logging.getLogger(__name__)


class PromQLConnLineConstructor(ClassicSQLConnLineConstructor["PromQLConnTargetDTO"]):
    def _get_dsn_params(
        self,
        safe_db_symbols: tuple[str, ...] = (),
        db_name: Optional[str] = None,
        standard_auth: Optional[bool] = True,
    ) -> dict:
        return dict(
            dialect=self._dialect_name,
            user=quote_plus(self._target_dto.username) if self._target_dto.username is not None else None,
            passwd=quote_plus(self._target_dto.password) if self._target_dto.password is not None else None,
            host=quote_plus(self._target_dto.host),
            port=quote_plus(str(self._target_dto.port)),
            db_name=db_name or quote_plus(self._target_dto.db_name or "", safe="".join(safe_db_symbols)),
        )

    def _get_dsn_query_params(self) -> dict:
        return {
            "protocol": self._target_dto.protocol,
        }


@attr.s
class PromQLAdapter(BaseClassicAdapter["PromQLConnTargetDTO"]):
    conn_type: ClassVar[ConnectionType] = CONNECTION_TYPE_PROMQL
    conn_line_constructor_type: ClassVar[Type[PromQLConnLineConstructor]] = PromQLConnLineConstructor

    _type_code_to_sa = {
        None: sa.TEXT,  # fallback
        "float64": sa.FLOAT,
        "string": sa.TEXT,
        "unix_timestamp": sa.DATETIME,
    }

    def _test(self) -> None:
        engine = self.get_db_engine(db_name="")
        connection = engine.raw_connection()
        try:
            connection.cli.test_connection()
        finally:
            connection.close()


class AsyncPromQLAdapter(AiohttpDBAdapter):
    conn_type: ClassVar[ConnectionType] = CONNECTION_TYPE_PROMQL

    _target_dto: PromQLConnTargetDTO = attr.ib()

    _url: str = attr.ib(init=False, default=None)

    def __attrs_post_init__(self) -> None:
        super().__attrs_post_init__()
        self._url = make_url(
            protocol=self._target_dto.protocol,
            host=self._target_dto.host,
            port=self._target_dto.port,
        )

    def get_session_auth(self) -> Optional[BasicAuth]:
        if self._target_dto.username and self._target_dto.password:
            return BasicAuth(
                login=self._target_dto.username,
                password=self._target_dto.password,
                encoding="utf-8",
            )
        return None

    async def _run_query(self, dba_query: DBAdapterQuery) -> ClientResponse:
        req_params = {"from", "to", "step"}
        conn_params = dba_query.connector_specific_params
        if conn_params is None or not req_params <= set(conn_params):
            db_exc = self.make_exc(
                status_code=HTTPBadRequest.status_code,
                err_body="'step', 'from', 'to' must be in parameters",
                debug_compiled_query=dba_query.debug_compiled_query,
            )
            raise db_exc

        for param in ("from", "to"):
            conn_param = conn_params[param]
            if isinstance(conn_param, datetime):
                conn_params[param] = int(conn_param.timestamp())

        query_text = self.compile_query_for_execution(dba_query.query)
        resp = await self._session.post(
            url=urljoin(self._url, "api/v1/query_range"),
            data={
                "query": query_text,
                "start": conn_params["from"],
                "end": conn_params["to"],
                "step": conn_params["step"],
            },
        )
        return resp

    @staticmethod
    def _parse_response_body_data(data: dict) -> dict:
        rows = []
        schema: list[tuple[str, str]] = []
        for chunk in data["result"]:
            chunk_schema = [("timestamp", "unix_timestamp"), ("value", "float64")] + [
                (label, "string") for label in chunk["metric"]
            ]

            if not schema:
                schema = chunk_schema
            elif set(schema) != set(chunk_schema):
                raise ValueError("Different schemas are not supported")

            label_values = [value for _, value in chunk["metric"].items()]
            for ts, v in chunk["values"]:
                row = [datetime.fromtimestamp(ts), float(v)] + label_values
                rows.append(row)

        return dict(rows=rows, schema=schema)

    async def _parse_response_body(self, resp: ClientResponse, dba_query: DBAdapterQuery) -> dict:
        data = await resp.json()
        try:
            return self._parse_response_body_data(data["data"])
        except ValueError as err:
            raise DatabaseQueryError(
                message=f"Unexpected API response body: {err.args[0]}",
                db_message=data["data"]["result"][:100],
                query=dba_query.debug_compiled_query,
            )

    @staticmethod
    def make_exc(  # TODO:  Move to ErrorTransformer
        status_code: int,  # noqa
        err_body: str,
        debug_compiled_query: Optional[str] = None,
    ) -> DatabaseQueryError:
        exc_cls = DatabaseQueryError
        return exc_cls(
            db_message=err_body,
            query=debug_compiled_query,
            orig=None,
            details={},
        )

    @generic_profiler_async("db-full")  # type: ignore  # TODO: fix
    async def execute(self, dba_query: DBAdapterQuery) -> AsyncRawExecutionResult:
        with self.wrap_execute_excs(query=dba_query, stage="request"):
            resp = await self._run_query(dba_query)

        if resp.status != 200:
            body = await resp.text()
            body_piece = body[:100] + ("…" if len(body) > 100 else "")  # TODO: depends on db ownership
            db_exc = self.make_exc(
                status_code=resp.status,
                err_body=body_piece,
                debug_compiled_query=dba_query.debug_compiled_query,
            )
            raise db_exc

        rd = await self._parse_response_body(resp, dba_query=dba_query)

        async def chunk_gen(
            chunk_size: int = dba_query.chunk_size or self._default_chunk_size,
        ) -> TBIChunksGen:
            data = rd["rows"]
            while data:
                chunk = data[:chunk_size]
                data = data[chunk_size:]
                yield chunk

        return AsyncRawExecutionResult(
            raw_cursor_info=dict(
                schema=rd["schema"],
                names=[name for name, _ in rd["schema"]],
                driver_types=[driver_type for _, driver_type in rd["schema"]],
                db_types=[self._promql_type_name_to_native_type(driver_type) for _, driver_type in rd["schema"]],
            ),
            raw_chunk_generator=chunk_gen(),
        )

    def _promql_type_name_to_native_type(self, type_name: str) -> GenericNativeType:
        return GenericNativeType.normalize_name_and_create(
            conn_type=self.conn_type,
            name=type_name,
        )

    async def test(self) -> None:
        await self._session.get(
            url=urljoin(self._url, "-/ready"),
        )

    async def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        raise NotImplementedError()

    async def get_schema_names(self, db_ident: DBIdent) -> list[str]:
        raise NotImplementedError()

    async def get_tables(self, schema_ident: SchemaIdent) -> list[TableIdent]:
        raise NotImplementedError()

    async def get_table_info(self, table_def: TableDefinition, fetch_idx_info: bool) -> RawSchemaInfo:
        raise NotImplementedError()

    async def is_table_exists(self, table_ident: TableIdent) -> bool:
        raise NotImplementedError()

lib/dl_connector_promql/dl_connector_promql/core/adapter.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Sequence,
)

import attr

from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor

from dl_connector_promql.core.adapter import (
    AsyncPromQLAdapter,
    PromQLAdapter,
)
from dl_connector_promql.core.target_dto import PromQLConnTargetDTO


if TYPE_CHECKING:
    from dl_connector_promql.core.dto import PromQLConnDTO


@attr.s(cmp=False, hash=False)
class PromQLConnExecutor(DefaultSqlAlchemyConnExecutor[PromQLAdapter]):
    TARGET_ADAPTER_CLS = PromQLAdapter

    _conn_dto: PromQLConnDTO = attr.ib()

    async def _make_target_conn_dto_pool(self) -> Sequence[PromQLConnTargetDTO]:
        return [
            PromQLConnTargetDTO(
                conn_id=self._conn_dto.conn_id,
                pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                host=self._conn_dto.host,
                port=self._conn_dto.port,
                username=self._conn_dto.username,
                password=self._conn_dto.password,
                protocol=self._conn_dto.protocol,
                db_name=self._conn_dto.db_name,
            )
        ]


@attr.s(cmp=False, hash=False)
class PromQLAsyncAdapterConnExecutor(DefaultSqlAlchemyConnExecutor[AsyncPromQLAdapter]):
    TARGET_ADAPTER_CLS = AsyncPromQLAdapter

    _conn_dto: PromQLConnDTO = attr.ib()

    async def _make_target_conn_dto_pool(self) -> Sequence[PromQLConnTargetDTO]:
        return [
            PromQLConnTargetDTO(
                conn_id=self._conn_dto.conn_id,
                pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                host=self._conn_dto.host,
                port=self._conn_dto.port,
                username=self._conn_dto.username,
                password=self._conn_dto.password,
                protocol=self._conn_dto.protocol,
                db_name=self._conn_dto.db_name,
            )
        ]

lib/dl_connector_promql/dl_connector_promql/core/connection_executors.py
datalens-tech/datalens-backend
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
    CoreSourceDefinition,
)

from dl_connector_promql.core.adapter import (
    AsyncPromQLAdapter,
    PromQLAdapter,
)
from dl_connector_promql.core.connection_executors import (
    PromQLAsyncAdapterConnExecutor,
    PromQLConnExecutor,
)
from dl_connector_promql.core.constants import (
    BACKEND_TYPE_PROMQL,
    CONNECTION_TYPE_PROMQL,
    SOURCE_TYPE_PROMQL,
)
from dl_connector_promql.core.data_source import PromQLDataSource
from dl_connector_promql.core.storage_schemas.connection import PromQLConnectionDataStorageSchema
from dl_connector_promql.core.type_transformer import PromQLTypeTransformer
from dl_connector_promql.core.us_connection import PromQLConnection


class PromQLCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_PROMQL
    connection_cls = PromQLConnection
    us_storage_schema_cls = PromQLConnectionDataStorageSchema
    type_transformer_cls = PromQLTypeTransformer
    sync_conn_executor_cls = PromQLConnExecutor
    async_conn_executor_cls = PromQLAsyncAdapterConnExecutor
    dialect_string = "bi_promql"
    custom_dashsql_key_names = frozenset(("from", "to", "step"))


class PromQLCoreSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_PROMQL
    source_cls = PromQLDataSource


class PromQLCoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_PROMQL
    connection_definitions = (PromQLCoreConnectionDefinition,)
    source_definitions = (PromQLCoreSourceDefinition,)
    rqe_adapter_classes = frozenset({AsyncPromQLAdapter, PromQLAdapter})

lib/dl_connector_promql/dl_connector_promql/core/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
    SourceBackendType,
)


BACKEND_TYPE_PROMQL = SourceBackendType.declare("PROMQL")
CONNECTION_TYPE_PROMQL = ConnectionType.declare("promql")
SOURCE_TYPE_PROMQL = DataSourceType.declare("PROMQL")

lib/dl_connector_promql/dl_connector_promql/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    ClassVar,
    Optional,
)

from dl_core.data_source.sql import PseudoSQLDataSource

from dl_connector_promql.core.constants import CONNECTION_TYPE_PROMQL


class PromQLDataSource(PseudoSQLDataSource):
    preview_enabled: ClassVar[bool] = False
    supports_offset: ClassVar[bool] = False

    conn_type = CONNECTION_TYPE_PROMQL

    @property
    def db_name(self) -> Optional[str]:
        return ""

lib/dl_connector_promql/dl_connector_promql/core/data_source.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_core.connection_models.dto_defs import DefaultSQLDTO

from dl_connector_promql.core.constants import CONNECTION_TYPE_PROMQL


@attr.s(frozen=True)
class PromQLConnDTO(DefaultSQLDTO):
    conn_type = CONNECTION_TYPE_PROMQL

    protocol: str = attr.ib(kw_only=True)

lib/dl_connector_promql/dl_connector_promql/core/dto.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_core.us_manager.storage_schemas.connection import ConnectionSQLDataStorageSchema

from dl_connector_promql.core.us_connection import PromQLConnection


class PromQLConnectionDataStorageSchema(ConnectionSQLDataStorageSchema[PromQLConnection.DataModel]):
    TARGET_CLS = PromQLConnection.DataModel

    secure = ma_fields.Boolean(required=False, allow_none=False, dump_default=False, load_default=False)

lib/dl_connector_promql/dl_connector_promql/core/storage_schemas/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_core.connection_executors.models.connection_target_dto_base import BaseSQLConnTargetDTO


@attr.s
class PromQLConnTargetDTO(BaseSQLConnTargetDTO):
    protocol: str = attr.ib()

lib/dl_connector_promql/dl_connector_promql/core/target_dto.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.enums import UserDataType
from dl_core.db.conversion_base import (
    TypeTransformer,
    make_native_type,
)

from dl_connector_promql.core.constants import CONNECTION_TYPE_PROMQL


class PromQLTypeTransformer(TypeTransformer):
    conn_type = CONNECTION_TYPE_PROMQL
    native_to_user_map = {
        make_native_type(CONNECTION_TYPE_PROMQL, "unix_timestamp"): UserDataType.genericdatetime,
        make_native_type(CONNECTION_TYPE_PROMQL, "float64"): UserDataType.float,
        make_native_type(CONNECTION_TYPE_PROMQL, "string"): UserDataType.string,
    }
    user_to_native_map = dict(
        [(bi_type, native_type) for native_type, bi_type in native_to_user_map.items()]
        + [
            (UserDataType.datetime, make_native_type(CONNECTION_TYPE_PROMQL, "unix_timestamp")),
        ]
    )

lib/dl_connector_promql/dl_connector_promql/core/type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import ClassVar

import attr

from dl_core.us_connection_base import ClassicConnectionSQL

from dl_connector_promql.core.dto import PromQLConnDTO


class PromQLConnection(ClassicConnectionSQL):
    allow_cache: ClassVar[bool] = True
    is_always_user_source: ClassVar[bool] = True
    allow_dashsql: ClassVar[bool] = True

    @attr.s(kw_only=True)
    class DataModel(ClassicConnectionSQL.DataModel):
        secure: bool = attr.ib(default=False)

    def get_conn_dto(self) -> PromQLConnDTO:
        return PromQLConnDTO(
            conn_id=self.uuid,
            host=self.data.host,
            port=self.data.port,
            username=self.data.username,
            password=self.data.password,
            db_name=self.data.db_name or "",
            protocol="https" if self.data.secure else "http",
            multihosts=(),
        )

    @property
    def is_dashsql_allowed(self) -> bool:
        return True

    @property
    def is_subselect_allowed(self) -> bool:
        return False

    @property
    def allow_public_usage(self) -> bool:
        return True

lib/dl_connector_promql/dl_connector_promql/core/us_connection.py
datalens-tech/datalens-backend
import pytest

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_constants.enums import (
    ConnectionType,
    RawSQLLevel,
)
from dl_core_testing.database import (
    CoreDbConfig,
    Db,
)
from dl_core_testing.engine_wrapper import TestingEngineWrapper

from dl_connector_promql.core.constants import CONNECTION_TYPE_PROMQL
from dl_connector_promql_tests.db.config import (
    API_CONNECTION_SETTINGS,
    API_TEST_CONFIG,
)


class PromQLConnectionTestBase(ConnectionTestBase):
    bi_compeng_pg_on = False
    conn_type = CONNECTION_TYPE_PROMQL

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return ""

    @pytest.fixture(scope="class")
    def db(self, db_config: CoreDbConfig) -> Db:
        engine_wrapper = TestingEngineWrapper(config=db_config.engine_config)
        return Db(config=db_config, engine_wrapper=engine_wrapper)

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def connection_params(self) -> dict:
        return API_CONNECTION_SETTINGS


class PromQLDashSQLConnectionTest(PromQLConnectionTestBase):
    @pytest.fixture(scope="class")
    def connection_params(self) -> dict:
        return API_CONNECTION_SETTINGS | dict(
            raw_sql_level=RawSQLLevel.dashsql.value,
        )

lib/dl_connector_promql/dl_connector_promql_tests/db/api/base.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite

from dl_connector_promql_tests.db.api.base import PromQLConnectionTestBase


class TestPromQLConnection(PromQLConnectionTestBase, DefaultConnectorConnectionTestSuite):
    pass

lib/dl_connector_promql/dl_connector_promql_tests/db/api/test_connection.py
datalens-tech/datalens-backend
from aiohttp.test_utils import TestClient
import pytest

from dl_api_lib_testing.connector.dashsql_suite import DefaultDashSQLTestSuite

from dl_connector_promql_tests.db.api.base import PromQLConnectionTestBase


class TestPromQLDashSQL(PromQLConnectionTestBase, DefaultDashSQLTestSuite):
    @pytest.mark.asyncio
    async def test_basic_select(self, data_api_lowlevel_aiohttp_client: TestClient, saved_connection_id: str) -> None:
        resp = await self.get_dashsql_response(
            data_api_aio=data_api_lowlevel_aiohttp_client,
            conn_id=saved_connection_id,
            query="select 1, 2, 3",
            fail_ok=True,
        )

        assert resp.status == 400
        resp_data = await resp.json()
        assert resp_data["code"] == "ERR.DS_API.DB"
        assert resp_data["details"]["db_message"] == "'step', 'from', 'to' must be in parameters"

    @pytest.mark.asyncio
    @pytest.mark.parametrize(
        "dashsql_params_type",
        [
            "string",
            "datetime",
        ],
    )
    async def test_result(
        self,
        data_api_lowlevel_aiohttp_client: TestClient,
        saved_connection_id: str,
        dashsql_params_type: str,
    ) -> None:
        await self.get_dashsql_response(
            data_api_aio=data_api_lowlevel_aiohttp_client,
            conn_id=saved_connection_id,
            query='{job="nodeexporter", type={{data_type}}}',
            params={
                "from": {"type_name": dashsql_params_type, "value": "2021-09-28T17:00:00Z"},
                "to": {"type_name": dashsql_params_type, "value": "2021-09-28T18:00:00Z"},
                "step": {"type_name": "string", "value": "5m"},
                "data_type": {"type_name": "string", "value": "test_data"},
            },
        )

    async def test_result_with_connector_specific_params(
        self,
        data_api_lowlevel_aiohttp_client: TestClient,
        saved_connection_id: str,
    ) -> None:
        await self.get_dashsql_response(
            data_api_aio=data_api_lowlevel_aiohttp_client,
            conn_id=saved_connection_id,
            query='{job="nodeexporter", type={{data_type}}}',
            params={"data_type": {"type_name": "string", "value": "test_data"}},
            connector_specific_params={
                "step": "5m",
                "from": "2021-09-28T17:00:00Z",
                "to": "2021-09-28T18:00:00Z",
            },
        )

lib/dl_connector_promql/dl_connector_promql_tests/db/api/test_dashsql.py
datalens-tech/datalens-backend
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport


# Infra settings
CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=51911).host,
    port_us_http=get_test_container_hostport("us", fallback_port=51911).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=51910).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=51910).port,
    us_master_token="AC1ofiek8coB",
    core_connector_ep_names=["promql"],
)


API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["promql"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
)

PROMETHEUS_CONTAINER_HOSTPORT = get_test_container_hostport("db-prometheus", fallback_port=50518)
API_CONNECTION_SETTINGS = dict(
    host=PROMETHEUS_CONTAINER_HOSTPORT.host,
    port=PROMETHEUS_CONTAINER_HOSTPORT.port,
    username="admin",
    password="admin",
)

lib/dl_connector_promql/dl_connector_promql_tests/db/config.py
datalens-tech/datalens-backend
from dl_api_lib_testing.initialization import initialize_api_lib_test

from dl_connector_promql_tests.db.config import API_TEST_CONFIG


pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)

lib/dl_connector_promql/dl_connector_promql_tests/db/conftest.py
datalens-tech/datalens-backend
from dl_api_connector.i18n.localizer import CONFIGS as BI_API_CONNECTOR_CONFIGS
from dl_api_lib_testing.connection_form_base import ConnectionFormTestBase

from dl_connector_promql.api.connection_form.form_config import PromQLConnectionFormFactory
from dl_connector_promql.api.i18n.localizer import CONFIGS as BI_CONNECTOR_PROMQL_CONFIGS


class TestPromQLConnectionForm(ConnectionFormTestBase):
    CONN_FORM_FACTORY_CLS = PromQLConnectionFormFactory
    TRANSLATION_CONFIGS = BI_API_CONNECTOR_CONFIGS + BI_CONNECTOR_PROMQL_CONFIGS

lib/dl_connector_promql/dl_connector_promql_tests/unit/test_connection_form.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.connection_base import ConnectionSchema
from dl_api_connector.api_schema.connection_base_fields import secret_string_field
from dl_api_connector.api_schema.connection_mixins import RawSQLLevelMixin
from dl_api_connector.api_schema.extras import FieldExtra

from dl_connector_snowflake.core.us_connection import ConnectionSQLSnowFlake


class SnowFlakeConnectionSchema(ConnectionSchema, RawSQLLevelMixin):
    TARGET_CLS = ConnectionSQLSnowFlake

    account_name = ma_fields.String(
        attribute="data.account_name",
        required=True,
        bi_extra=FieldExtra(editable=True),
    )
    user_name = ma_fields.String(
        attribute="data.user_name",
        required=True,
        bi_extra=FieldExtra(editable=True),
    )
    user_role = ma_fields.String(
        attribute="data.user_role",
        required=False,
        allow_none=True,
        bi_extra=FieldExtra(editable=True),
    )
    client_id = ma_fields.String(
        attribute="data.client_id",
        required=True,
        bi_extra=FieldExtra(editable=True),
    )
    client_secret = secret_string_field(
        attribute="data.client_secret",
        required=True,
        bi_extra=FieldExtra(editable=True),
    )
    schema = ma_fields.String(
        attribute="data.schema",
        required=True,
        bi_extra=FieldExtra(editable=True),
    )
    db_name = ma_fields.String(
        attribute="data.db_name",
        required=True,
        bi_extra=FieldExtra(editable=True),
    )
    warehouse = ma_fields.String(
        attribute="data.warehouse",
        required=True,
        bi_extra=FieldExtra(editable=True),
    )
    refresh_token = secret_string_field(
        attribute="data.refresh_token",
        required=False,
        bi_extra=FieldExtra(editable=True),
    )
    refresh_token_expire_time = ma_fields.DateTime(
        attribute="data.refresh_token_expire_time",
        required=False,
        allow_none=True,
        bi_extra=FieldExtra(editable=True),
    )

lib/dl_connector_snowflake/dl_connector_snowflake/api/api_schema/connection.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields

from dl_api_connector.api_schema.source_base import (
    SimpleDataSourceSchema,
    SimpleDataSourceTemplateSchema,
)
from dl_model_tools.schema.base import BaseSchema


class SnowFlakeTableParametersSchema(BaseSchema):
    db_name = ma_fields.String(allow_none=True)
    schema = ma_fields.String(allow_none=True)
    table_name = ma_fields.String(allow_none=True)


class SnowFlakeTableDataSourceSchema(SimpleDataSourceSchema):
    parameters = ma_fields.Nested(SnowFlakeTableParametersSchema)


class SnowFlakeTableDataSourceTemplateSchema(SimpleDataSourceTemplateSchema):
    parameters = ma_fields.Nested(SnowFlakeTableParametersSchema)

lib/dl_connector_snowflake/dl_connector_snowflake/api/api_schema/source.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import unique
from typing import Optional

import attr

from dl_api_commons.base_models import TenantDef
from dl_api_connector.form_config.models.api_schema import (
    FormActionApiSchema,
    FormApiSchema,
    FormFieldApiSchema,
)
from dl_api_connector.form_config.models.base import (
    ConnectionForm,
    ConnectionFormFactory,
    ConnectionFormMode,
)
from dl_api_connector.form_config.models.common import (
    CommonFieldName,
    FormFieldName,
    MarkdownStr,
    remap_skip_if_null,
)
import dl_api_connector.form_config.models.rows as C
from dl_api_connector.form_config.models.rows.base import (
    DisplayConditionsMixin,
    FormRow,
    TDisplayConditions,
)
from dl_api_connector.form_config.models.rows.prepared.base import PreparedRow
from dl_configs.connectors_settings import ConnectorSettingsBase

from dl_connector_snowflake.api.connection_info import SnowflakeConnectionInfoProvider
from dl_connector_snowflake.api.i18n.localizer import Translatable


@unique
class SnowFlakeFieldName(FormFieldName):
    account_name = "account_name"
    client_id = "client_id"
    client_secret = "client_secret"
    schema = "schema"
    warehouse = "warehouse"
    user_name = "user_name"
    user_role = "user_role"

    snowflake_auth = "snowflake_auth"
    snowflake_db_details = "snowflake_db_details"
    # ^ actually internal fields that show if a certain section is opened, but it is easier to have them here for now


@attr.s(kw_only=True, frozen=True)
class SnowFlakeOAuthIntegrationRow(PreparedRow, DisplayConditionsMixin):
    type = "snowflake_oauth_integration"


@attr.s(kw_only=True, frozen=True)
class OAuthSnowFlakeRow(PreparedRow, DisplayConditionsMixin):
    type = "oauth_snowflake"

    class Inner(PreparedRow.Inner):
        refresh_token = "refresh_token"
        refresh_token_expire_time = "refresh_token_expire_time"

    fake_value: Optional[str] = attr.ib(default=None, metadata=remap_skip_if_null("fakeValue"))
    button_text: Optional[str] = attr.ib(default=None, metadata=remap_skip_if_null("buttonText"))


def _basic_input_row(
    label_text: str,
    field: FormFieldName,
    help_text: Optional[MarkdownStr],
    display_conditions: Optional[TDisplayConditions] = None,
    password_input: bool = False,
    fake_value: Optional[str] = None,
) -> C.CustomizableRow:
    return C.CustomizableRow(
        items=[
            C.LabelRowItem(
                text=label_text,
                help_text=help_text,
                display_conditions=display_conditions,
            ),
            C.InputRowItem(
                name=field,
                width="m",
                display_conditions=display_conditions,
                control_props=C.InputRowItem.Props(type="password") if password_input else None,
                fake_value=fake_value,
            ),
        ]
    )


class SnowFlakeConnectionFormFactory(ConnectionFormFactory):
    def get_form_config(
        self,
        connector_settings: Optional[ConnectorSettingsBase],
        tenant: Optional[TenantDef],
    ) -> ConnectionForm:
        on_create = self.mode == ConnectionFormMode.create

        common_api_schema_items: list[FormFieldApiSchema] = [
            FormFieldApiSchema(name=SnowFlakeFieldName.account_name, required=True),
            FormFieldApiSchema(name=SnowFlakeFieldName.client_id, required=on_create),
            FormFieldApiSchema(name=SnowFlakeFieldName.client_secret, required=on_create),
            FormFieldApiSchema(name=OAuthSnowFlakeRow.Inner.refresh_token, required=on_create),
            FormFieldApiSchema(name=SnowFlakeFieldName.user_name, required=True),
            FormFieldApiSchema(name=CommonFieldName.db_name, required=True),
            FormFieldApiSchema(name=SnowFlakeFieldName.schema, required=True),
            FormFieldApiSchema(name=SnowFlakeFieldName.warehouse, required=True),
        ]

        edit_api_schema = FormActionApiSchema(
            items=[
                *common_api_schema_items,
                FormFieldApiSchema(name=OAuthSnowFlakeRow.Inner.refresh_token_expire_time, nullable=True),
                FormFieldApiSchema(name=SnowFlakeFieldName.user_role, nullable=True),
            ]
        )

        create_api_schema = FormActionApiSchema(
            items=[
                *edit_api_schema.items,
                *self._get_top_level_create_api_schema_items(),
            ]
        )

        check_api_schema = FormActionApiSchema(
            items=[
                *common_api_schema_items,
                *self._get_top_level_check_api_schema_items(),
            ]
        )

        on_auth_opened: TDisplayConditions = {SnowFlakeFieldName.snowflake_auth: "opened"}
        on_db_opened: TDisplayConditions = {SnowFlakeFieldName.snowflake_db_details: "opened"}

        help_text_sf_docs = self._localizer.translate(Translatable("label_snowflake-doc-help"))
        sf_docs_url = "https://docs.snowflake.com/en"
        help_account_name = help_text_sf_docs.format(
            LINK=f"{sf_docs_url}/user-guide/client-redirect.html#snowsight-the-snowflake-web-interface"
        )
        help_client = help_text_sf_docs.format(
            LINK=f"{sf_docs_url}/sql-reference/functions/system_show_oauth_client_secrets.html#system-show-oauth-client-secrets"
        )
        help_user_name = help_text_sf_docs.format(LINK=f"{sf_docs_url}/sql-reference/sql/create-user.html#create-user")
        help_db = help_text_sf_docs.format(LINK=f"{sf_docs_url}/sql-reference/sql/create-database.html#create-database")
        help_schema = help_text_sf_docs.format(LINK=f"{sf_docs_url}/sql-reference/sql/create-schema.html#create-schema")
        help_warehouse = help_text_sf_docs.format(
            LINK=f"{sf_docs_url}/sql-reference/sql/create-warehouse.html#create-warehouse"
        )
        help_user_role = help_text_sf_docs.format(LINK=f"{sf_docs_url}/sql-reference/sql/create-role.html#create-role")

        auth_section: list[FormRow] = [
            SnowFlakeOAuthIntegrationRow(display_conditions=on_auth_opened),
            _basic_input_row(
                label_text=self._localizer.translate(Translatable("label_account-name")),
                field=SnowFlakeFieldName.account_name,
                help_text=help_account_name,
                display_conditions=on_auth_opened,
            ),
            _basic_input_row(
                label_text=self._localizer.translate(Translatable("label_client-id")),
                field=SnowFlakeFieldName.client_id,
                help_text=help_client,
                display_conditions=on_auth_opened,
            ),
            _basic_input_row(
                label_text=self._localizer.translate(Translatable("label_client-secret")),
                field=SnowFlakeFieldName.client_secret,
                help_text=help_client,
                display_conditions=on_auth_opened,
                password_input=True,
                fake_value="******" if self.mode == ConnectionFormMode.edit else None,
            ),
            OAuthSnowFlakeRow(
                button_text=None
                if self.mode == ConnectionFormMode.create
                else self._localizer.translate(Translatable("button_get-new-token")),
                display_conditions=on_auth_opened,
                fake_value="******" if self.mode == ConnectionFormMode.edit else None,
            ),
            C.CustomizableRow(
                items=[
                    C.HiddenRowItem(name=OAuthSnowFlakeRow.Inner.refresh_token_expire_time),
                ]
            ),
        ]

        db_section: list[FormRow] = [
            _basic_input_row(
                label_text=self._localizer.translate(Translatable("label_user-name")),
                field=SnowFlakeFieldName.user_name,
                help_text=help_user_name,
                display_conditions=on_db_opened,
            ),
            _basic_input_row(
                label_text=self._localizer.translate(Translatable("label_database")),
                field=CommonFieldName.db_name,
                help_text=help_db,
                display_conditions=on_db_opened,
            ),
            _basic_input_row(
                label_text=self._localizer.translate(Translatable("label_database-schema")),
                field=SnowFlakeFieldName.schema,
                help_text=help_schema,
                display_conditions=on_db_opened,
            ),
            _basic_input_row(
                label_text=self._localizer.translate(Translatable("label_warehouse")),
                field=SnowFlakeFieldName.warehouse,
                help_text=help_warehouse,
                display_conditions=on_db_opened,
            ),
            _basic_input_row(
                label_text=self._localizer.translate(Translatable("label_user-role")),
                field=SnowFlakeFieldName.user_role,
                help_text=help_user_role,
                display_conditions=on_db_opened,
            ),
        ]

        return ConnectionForm(
            title=SnowflakeConnectionInfoProvider.get_title(self._localizer),
            rows=[
                C.CollapseRow(
                    name=SnowFlakeFieldName.snowflake_auth,
                    text=self._localizer.translate(Translatable("label_showflake-auth-section-title")),
                    component_props=C.CollapseRow.Props(default_expanded=self.mode == ConnectionFormMode.create),
                ),
                *auth_section,
                C.CollapseRow(
                    name=SnowFlakeFieldName.snowflake_db_details,
                    text=self._localizer.translate(Translatable("label_showflake-db-details-section-title")),
                    component_props=C.CollapseRow.Props(default_expanded=self.mode == ConnectionFormMode.edit),
                ),
                *db_section,
            ],
            api_schema=FormApiSchema(
                create=create_api_schema if self.mode == ConnectionFormMode.create else None,
                edit=edit_api_schema if self.mode == ConnectionFormMode.edit else None,
                check=check_api_schema,
            ),
        )

lib/dl_connector_snowflake/dl_connector_snowflake/api/connection_form/form_config.py
datalens-tech/datalens-backend
from dl_api_connector.connection_info import ConnectionInfoProvider

from dl_connector_snowflake.api.i18n.localizer import Translatable


class SnowflakeConnectionInfoProvider(ConnectionInfoProvider):
    title_translatable = Translatable("label_connector-snowflake")

lib/dl_connector_snowflake/dl_connector_snowflake/api/connection_info.py
datalens-tech/datalens-backend
from dl_api_connector.connector import (
    ApiConnectionDefinition,
    ApiConnector,
    ApiSourceDefinition,
)

from dl_connector_snowflake.api.api_schema.connection import SnowFlakeConnectionSchema
from dl_connector_snowflake.api.api_schema.source import (  # type: ignore
    SnowFlakeTableDataSourceSchema,
    SnowFlakeTableDataSourceTemplateSchema,
)
from dl_connector_snowflake.api.connection_form.form_config import SnowFlakeConnectionFormFactory
from dl_connector_snowflake.api.connection_info import SnowflakeConnectionInfoProvider
from dl_connector_snowflake.api.i18n.localizer import CONFIGS
from dl_connector_snowflake.core.connector import (
    SnowFlakeCoreConnectionDefinition,
    SnowFlakeCoreConnector,
    SnowFlakeCoreTableSourceDefinition,
)
from dl_connector_snowflake.formula.constants import DIALECT_NAME_SNOWFLAKE


class SnowFlakeApiTableSourceDefinition(ApiSourceDefinition):
    core_source_def_cls = SnowFlakeCoreTableSourceDefinition
    api_schema_cls = SnowFlakeTableDataSourceSchema
    template_api_schema_cls = SnowFlakeTableDataSourceTemplateSchema


class SnowFlakeApiConnectionDefinition(ApiConnectionDefinition):
    core_conn_def_cls = SnowFlakeCoreConnectionDefinition
    api_generic_schema_cls = SnowFlakeConnectionSchema
    info_provider_cls = SnowflakeConnectionInfoProvider
    form_factory_cls = SnowFlakeConnectionFormFactory


class SnowFlakeApiConnector(ApiConnector):
    core_connector_cls = SnowFlakeCoreConnector
    connection_definitions = (SnowFlakeApiConnectionDefinition,)
    source_definitions = (SnowFlakeApiTableSourceDefinition,)
    formula_dialect_name = DIALECT_NAME_SNOWFLAKE
    translation_configs = frozenset(CONFIGS)

lib/dl_connector_snowflake/dl_connector_snowflake/api/connector.py
datalens-tech/datalens-backend
import os

import attr

from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig

import dl_connector_snowflake as package


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_connector_snowflake/dl_connector_snowflake/api/i18n/localizer.py
datalens-tech/datalens-backend
from __future__ import annotations

from datetime import (
    datetime,
    timedelta,
)
from typing import Optional
from urllib import parse

import attr
import requests

from dl_api_commons.aiohttp.aiohttp_client import (
    BIAioHTTPClient,
    PredefinedIntervalsRetrier,
    THeaders,
)

from dl_connector_snowflake.core.dto import SnowFlakeConnDTO
from dl_connector_snowflake.core.exc import SnowflakeGetAccessTokenError


@attr.s()
class AccessTokenRequestArgs:
    endpoint: str = attr.ib()
    params: dict[str, str] = attr.ib()
    basic_auth: tuple[str, str] = attr.ib()

    @property
    def auth_as_headers(self) -> THeaders:
        return {
            "Authorization": requests.auth._basic_auth_str(*self.basic_auth),
        }


@attr.s()
class SFAuthProvider:
    account_name: str = attr.ib()
    user_name: str = attr.ib()
    client_id: str = attr.ib()
    client_secret: str = attr.ib()
    refresh_token: str = attr.ib()
    refresh_token_expire_time: Optional[datetime] = attr.ib(default=None)

    _SNOWFLAKE_DOMAIN = "snowflakecomputing.com"
    _NOTIFICATION_DAYS_BEFORE_TOKEN_EXPIRE = 7

    @classmethod
    def from_dto(cls, dto: SnowFlakeConnDTO) -> SFAuthProvider:
        return cls(
            account_name=dto.account_name,
            user_name=dto.user_name,
            client_id=dto.client_id,
            client_secret=dto.client_secret,
            refresh_token=dto.refresh_token,
            refresh_token_expire_time=dto.refresh_token_expire_time,
        )

    def prepare_access_token_request_kwargs(self) -> AccessTokenRequestArgs:
        return AccessTokenRequestArgs(
            endpoint=f"https://{self.account_name}.{self._SNOWFLAKE_DOMAIN}/oauth/token-request",
            params=dict(
                refresh_token=self.refresh_token,
                grant_type="refresh_token",
                # snowflake documentation mentions redirect uri, but works fine as is
                # redirect_uri="http://localhost:8000",
            ),
            basic_auth=(self.client_id, parse.quote(self.client_secret)),
        )

    def _store_access_token(self, access_token: str, ttl: Optional[int] = None) -> None:
        pass

    async def async_get_access_token(self) -> str:
        ra = self.prepare_access_token_request_kwargs()
        async with BIAioHTTPClient(
            base_url=ra.endpoint,
            conn_timeout_sec=1,
            read_timeout_sec=2,
            raise_for_status=False,
            retrier=PredefinedIntervalsRetrier(retry_intervals=(0.2, 0.3, 0.5), retry_methods={"POST"}),
            session=None,
        ) as http_client:
            async with http_client.request(
                method="POST",
                data=ra.params,
                headers=ra.auth_as_headers,
                conn_timeout_sec=1,
            ) as response:
                try:
                    raw = await response.read()

                    if response.status == 200:
                        try:
                            json_data = await response.json()
                            return json_data["access_token"]
                        except Exception:  # noqa
                            pass

                    raise SnowflakeGetAccessTokenError(
                        "Failed to get access token, try to obtain new refresh token",
                        details={
                            "status_code": response.status,
                            "raw_response": repr(raw),
                        },
                    )
                except Exception:
                    raise SnowflakeGetAccessTokenError(
                        "Failed to get access token, try to obtain new refresh token",
                        details={
                            "status_code": response.status,
                            "raw_response": repr(raw),
                        },
                    )

    def should_notify_refresh_token_to_expire_soon(self) -> bool:
        if self.refresh_token_expire_time is None:
            return False

        td = timedelta(days=self._NOTIFICATION_DAYS_BEFORE_TOKEN_EXPIRE)
        return self.refresh_token_expire_time - datetime.now() < td

    def is_refresh_token_expired(self) -> bool:
        if self.refresh_token_expire_time is None:
            return False

        return self.refresh_token_expire_time < datetime.now()

lib/dl_connector_snowflake/dl_connector_snowflake/auth.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Callable,
    Optional,
)

import attr
from snowflake import sqlalchemy as ssa
from snowflake.connector import connect as sf_connect
import sqlalchemy as sa
from sqlalchemy.engine import Engine
from sqlalchemy.sql.type_api import TypeEngine

from dl_core.connection_executors.adapters.adapters_base_sa import BaseSAAdapter
from dl_core.connection_executors.adapters.adapters_base_sa_classic import BaseClassicAdapter
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    RawSchemaInfo,
)
from dl_core.connection_models.common_models import (
    DBIdent,
    SATextTableDefinition,
)
from dl_core.db.native_type import SATypeSpec

from dl_connector_snowflake.core.constants import CONNECTION_TYPE_SNOWFLAKE
from dl_connector_snowflake.core.error_transformer import snowflake_error_transformer
from dl_connector_snowflake.core.target_dto import SnowFlakeConnTargetDTO


def construct_creator_func(target_dto: SnowFlakeConnTargetDTO) -> Callable:
    def get_connection() -> Any:
        params = dict(
            user=target_dto.user_name,
            account=target_dto.account_name,
            authenticator="oauth",
            token=target_dto.access_token,
            cache_column_metadata=True,
            database=target_dto.db_name,
            schema=target_dto.schema,
            warehouse=target_dto.warehouse,
        )
        if target_dto.user_role:
            params["role"] = target_dto.user_role

        conn = sf_connect(**params)

        # note: let's keep following lines commented for some time,
        #   at one point SF adapter did not respect db/schema/wh passed to connect
        #   in the recent runs it seems fixed
        # conn.cursor().execute(f"use database {target_dto.db_name}")
        # conn.cursor().execute(f"use schema {target_dto.schema}")
        # conn.cursor().execute(f"use warehouse {target_dto.warehouse}")
        return conn

    return get_connection


@attr.s(kw_only=True)
class SnowFlakeDefaultAdapter(BaseClassicAdapter, BaseSAAdapter[SnowFlakeConnTargetDTO]):
    conn_type = CONNECTION_TYPE_SNOWFLAKE
    _error_transformer = snowflake_error_transformer

    # https://docs.snowflake.com/en/user-guide/python-connector-api#type-codes
    _type_code_to_sa = {
        0: ssa.DECIMAL,
        1: ssa.REAL,
        2: ssa.STRING,
        3: ssa.DATE,
        4: ssa.TIMESTAMP,
        5: ssa.VARIANT,
        6: ssa.TIMESTAMP_LTZ,
        7: ssa.TIMESTAMP_TZ,
        8: ssa.TIMESTAMP_TZ,
        9: ssa.OBJECT,
        10: ssa.ARRAY,
        11: ssa.BINARY,
        12: ssa.TIME,
        13: ssa.BOOLEAN,
    }

    def _cursor_column_to_sa(self, cursor_col, require: bool = True) -> Optional[SATypeSpec]:  # type: ignore  # TODO: fix
        """
        cursor_col:

            name, type_code, display_size, internal_size, precision, scale, is_nullable

            https://docs.snowflake.com/en/user-guide/python-connector-api#id4
        """
        type_code = cursor_col[1]
        if type_code is None:  # shouldn't really happen
            if require:
                raise ValueError(f"Cursor column has no type: {cursor_col!r}")
            return None
        sa_cls = self._type_code_to_sa.get(type_code)
        if sa_cls is None:
            if require:
                raise ValueError(f"Unknown cursor type: {type_code!r}")
            return None

        if sa_cls is ssa.NUMBER:
            # See also: `self.normalize_sa_col_type`
            precision = cursor_col[4]
            scale = cursor_col[5]

            # taken from Oracle adapter, not sure if applicapble for Snowflake
            # # Going by the comparison with the 'create view' -> SA logic.
            # if scale == -127:
            #     scale = 0
            sa_type = sa_cls(precision, scale)
        else:
            sa_type = sa_cls

        return sa_type

    def get_default_db_name(self) -> Optional[str]:
        return None

    def get_db_name_for_query(self, db_name_from_query: Optional[str]) -> str:
        return self._target_dto.db_name

    def _get_db_engine(self, db_name: str, disable_streaming: bool = False) -> Engine:
        if disable_streaming:
            raise Exception("`disable_streaming` is not applicable here")
        return sa.create_engine(
            "snowflake://not@used/db",
            creator=construct_creator_func(self._target_dto),
        ).execution_options(compiled_cache=None)

    def _get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        return self.execute(DBAdapterQuery("SELECT CURRENT_VERSION()", db_name=db_ident.db_name)).get_all()[0][0]

    def normalize_sa_col_type(self, sa_col_type: TypeEngine) -> TypeEngine:
        if isinstance(sa_col_type, ssa.DECIMAL) and sa_col_type.scale == 0:
            return sa.Integer()

        return super().normalize_sa_col_type(sa_col_type)

    def _get_subselect_table_info(self, subquery: SATextTableDefinition) -> RawSchemaInfo:
        raw_cursor_info, _ = self._get_subselect_raw_cursor_info_and_data(subquery.text)
        return self._raw_cursor_info_to_schema(raw_cursor_info)

lib/dl_connector_snowflake/dl_connector_snowflake/core/adapters.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import TypeVar

import attr

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor
from dl_core.reporting.notifications import get_notification_record

from dl_connector_snowflake.auth import SFAuthProvider
from dl_connector_snowflake.core.adapters import SnowFlakeDefaultAdapter
from dl_connector_snowflake.core.constants import NOTIF_TYPE_SF_REFRESH_TOKEN_SOON_TO_EXPIRE
from dl_connector_snowflake.core.dto import SnowFlakeConnDTO
from dl_connector_snowflake.core.exc import SnowflakeRefreshTokenInvalid
from dl_connector_snowflake.core.target_dto import SnowFlakeConnTargetDTO


LOGGER = logging.getLogger(__name__)


_BASE_SNOWFLAKE_ADAPTER_TV = TypeVar("_BASE_SNOWFLAKE_ADAPTER_TV", bound=CommonBaseDirectAdapter)


@attr.s(cmp=False, hash=False)
class BaseSnowFlakeConnExecutor(DefaultSqlAlchemyConnExecutor[_BASE_SNOWFLAKE_ADAPTER_TV]):
    _conn_dto: SnowFlakeConnDTO = attr.ib()

    async def _make_target_conn_dto_pool(self) -> list[SnowFlakeConnTargetDTO]:  # type: ignore  # TODO: fix
        sf_auth_provider = SFAuthProvider.from_dto(self._conn_dto)

        if sf_auth_provider.is_refresh_token_expired():
            raise SnowflakeRefreshTokenInvalid()
        elif sf_auth_provider.should_notify_refresh_token_to_expire_soon():
            reporting_service = None
            if self._services_registry is not None:
                try:
                    reporting_service = self._services_registry.get_reporting_registry()
                except AttributeError:
                    LOGGER.error("Failed to get reporting registry from service registry")

            if reporting_service:
                reporting_service.save_reporting_record(
                    get_notification_record(NOTIF_TYPE_SF_REFRESH_TOKEN_SOON_TO_EXPIRE)
                )

        access_token = await sf_auth_provider.async_get_access_token()

        return [
            SnowFlakeConnTargetDTO(
                conn_id=self._conn_dto.conn_id,
                pass_db_messages_to_user=self._conn_options.pass_db_messages_to_user,
                pass_db_query_to_user=self._conn_options.pass_db_query_to_user,
                account_name=self._conn_dto.account_name,
                user_name=self._conn_dto.user_name,
                user_role=self._conn_dto.user_role,
                schema=self._conn_dto.schema,
                db_name=self._conn_dto.db_name,
                warehouse=self._conn_dto.warehouse,
                access_token=access_token,
            )
        ]


@attr.s(cmp=False, hash=False)
class SnowFlakeSyncConnExecutor(BaseSnowFlakeConnExecutor[SnowFlakeDefaultAdapter]):
    TARGET_ADAPTER_CLS = SnowFlakeDefaultAdapter

lib/dl_connector_snowflake/dl_connector_snowflake/core/connection_executors.py
datalens-tech/datalens-backend
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
    CoreSourceDefinition,
)

from dl_connector_snowflake.core.adapters import SnowFlakeDefaultAdapter
from dl_connector_snowflake.core.connection_executors import SnowFlakeSyncConnExecutor
from dl_connector_snowflake.core.constants import (
    BACKEND_TYPE_SNOWFLAKE,
    CONNECTION_TYPE_SNOWFLAKE,
    SOURCE_TYPE_SNOWFLAKE_SUBSELECT,
    SOURCE_TYPE_SNOWFLAKE_TABLE,
)
from dl_connector_snowflake.core.data_source import (
    SnowFlakeSubselectDataSource,
    SnowFlakeTableDataSource,
)
from dl_connector_snowflake.core.data_source_spec import (
    SnowFlakeSubselectDataSourceSpec,
    SnowFlakeTableDataSourceSpec,
)
from dl_connector_snowflake.core.lifecycle import SnowFlakeConnectionLifecycleManager
from dl_connector_snowflake.core.notifications import SnowflakeRefreshTokenSoonToExpire
from dl_connector_snowflake.core.storage_schemas.connection import SnowFlakeConnectionDataStorageSchema
from dl_connector_snowflake.core.storage_schemas.data_source_spec import (
    SnowFlakeSubselectDataSourceSpecStorageSchema,
    SnowFlakeTableDataSourceSpecStorageSchema,
)
from dl_connector_snowflake.core.type_transformer import SnowFlakeTypeTransformer
from dl_connector_snowflake.core.us_connection import ConnectionSQLSnowFlake


class SnowFlakeCoreConnectionDefinition(CoreConnectionDefinition):
    conn_type = CONNECTION_TYPE_SNOWFLAKE
    connection_cls = ConnectionSQLSnowFlake
    us_storage_schema_cls = SnowFlakeConnectionDataStorageSchema
    type_transformer_cls = SnowFlakeTypeTransformer
    sync_conn_executor_cls = SnowFlakeSyncConnExecutor
    lifecycle_manager_cls = SnowFlakeConnectionLifecycleManager
    dialect_string = "snowflake"


class SnowFlakeCoreTableSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_SNOWFLAKE_TABLE
    source_cls = SnowFlakeTableDataSource
    source_spec_cls = SnowFlakeTableDataSourceSpec
    us_storage_schema_cls = SnowFlakeTableDataSourceSpecStorageSchema


class SnowFlakeCoreSubselectSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_SNOWFLAKE_SUBSELECT
    source_cls = SnowFlakeSubselectDataSource
    source_spec_cls = SnowFlakeSubselectDataSourceSpec
    us_storage_schema_cls = SnowFlakeSubselectDataSourceSpecStorageSchema


class SnowFlakeCoreConnector(CoreConnector):
    backend_type = BACKEND_TYPE_SNOWFLAKE
    connection_definitions = (SnowFlakeCoreConnectionDefinition,)
    source_definitions = (
        SnowFlakeCoreTableSourceDefinition,
        SnowFlakeCoreSubselectSourceDefinition,
    )
    rqe_adapter_classes = frozenset({SnowFlakeDefaultAdapter})
    notification_classes = (SnowflakeRefreshTokenSoonToExpire,)

lib/dl_connector_snowflake/dl_connector_snowflake/core/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
    NotificationType,
    SourceBackendType,
)


BACKEND_TYPE_SNOWFLAKE = SourceBackendType.declare("SNOWFLAKE")

CONNECTION_TYPE_SNOWFLAKE = ConnectionType.declare("snowflake")  # FIXME: Move the declaration here

SOURCE_TYPE_SNOWFLAKE_TABLE = DataSourceType.declare("SNOWFLAKE_TABLE")
SOURCE_TYPE_SNOWFLAKE_SUBSELECT = DataSourceType.declare("SNOWFLAKE_SUBSELECT")

NOTIF_TYPE_SF_REFRESH_TOKEN_SOON_TO_EXPIRE = NotificationType.declare("snowflake_refresh_token_soon_to_expire")

ACCOUNT_NAME_RE = r"^[a-zA-Z\d\-\_\.]+$"

lib/dl_connector_snowflake/dl_connector_snowflake/core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Optional,
    Type,
)

from dl_constants.enums import DataSourceType
from dl_core.connection_models import (
    TableDefinition,
    TableIdent,
)
from dl_core.data_source.sql import (
    BaseSQLDataSource,
    SQLDataSource,
    SubselectDataSource,
    TableSQLDataSourceMixin,
    require_table_name,
)
from dl_core.db import SchemaInfo
from dl_core.utils import sa_plain_text

from dl_connector_snowflake.core.constants import (
    CONNECTION_TYPE_SNOWFLAKE,
    SOURCE_TYPE_SNOWFLAKE_SUBSELECT,
    SOURCE_TYPE_SNOWFLAKE_TABLE,
)
from dl_connector_snowflake.core.data_source_spec import (
    SnowFlakeSubselectDataSourceSpec,
    SnowFlakeTableDataSourceSpec,
)
from dl_connector_snowflake.core.us_connection import ConnectionSQLSnowFlake


if TYPE_CHECKING:
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase


class SnowFlakeDataSourceMixin(BaseSQLDataSource):
    conn_type = CONNECTION_TYPE_SNOWFLAKE

    @classmethod
    def get_connection_cls(cls) -> Type[ConnectionSQLSnowFlake]:
        return ConnectionSQLSnowFlake

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return source_type in (SOURCE_TYPE_SNOWFLAKE_TABLE, SOURCE_TYPE_SNOWFLAKE_SUBSELECT)


class SnowFlakeTableDataSource(SnowFlakeDataSourceMixin, TableSQLDataSourceMixin, SQLDataSource):
    @property
    def spec(self) -> SnowFlakeTableDataSourceSpec:
        assert isinstance(self._spec, SnowFlakeTableDataSourceSpec)
        return self._spec

    @property
    def connection(self) -> ConnectionSQLSnowFlake:
        connection = super().connection
        assert isinstance(connection, ConnectionSQLSnowFlake)
        return connection

    @property
    def db_name(self) -> Optional[str]:
        return self.connection.db_name

    @property
    def schema_name(self) -> Optional[str]:
        return self.connection.schema_name

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            table_name=self.spec.table_name,
        )

    def get_table_definition(self) -> TableDefinition:
        assert self.table_name is not None
        return TableIdent(
            db_name=self.db_name,
            schema_name=self.schema_name,
            table_name=self.table_name,
        )

    def get_schema_info(
        self,
        conn_executor_factory: Callable[[], SyncConnExecutorBase],
    ) -> SchemaInfo:
        return super(SnowFlakeTableDataSource, self).get_schema_info(conn_executor_factory=conn_executor_factory)

    @require_table_name
    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        q = self.quote
        alias_str = "" if alias is None else f" AS {q(alias)}"
        return sa_plain_text(f"{q(self.db_name)}.{q(self.schema_name)}.{q(self.table_name)}{alias_str}")


class SnowFlakeSubselectDataSource(SnowFlakeDataSourceMixin, SubselectDataSource):
    @property
    def spec(self) -> SnowFlakeSubselectDataSourceSpec:
        assert isinstance(self._spec, SnowFlakeSubselectDataSourceSpec)
        return self._spec

    @property
    def db_version(self) -> Optional[str]:
        return None

lib/dl_connector_snowflake/dl_connector_snowflake/core/data_source.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_core.data_source_spec.sql import (
    DbSQLDataSourceSpec,
    SchemaSQLDataSourceSpec,
    SubselectDataSourceSpec,
    TableSQLDataSourceSpec,
)


@attr.s
class SnowFlakeTableDataSourceSpec(TableSQLDataSourceSpec, SchemaSQLDataSourceSpec, DbSQLDataSourceSpec):
    pass


@attr.s
class SnowFlakeSubselectDataSourceSpec(SubselectDataSourceSpec):
    pass

lib/dl_connector_snowflake/dl_connector_snowflake/core/data_source_spec.py
datalens-tech/datalens-backend
from __future__ import annotations

from datetime import datetime
from typing import Optional

import attr

from dl_core.connection_models.dto_defs import ConnDTO
from dl_core.utils import secrepr

from dl_connector_snowflake.core.constants import CONNECTION_TYPE_SNOWFLAKE


@attr.s(frozen=True)
class SnowFlakeConnDTO(ConnDTO):
    conn_type = CONNECTION_TYPE_SNOWFLAKE

    account_name: str = attr.ib()
    user_name: str = attr.ib()
    user_role: Optional[str] = attr.ib()
    client_id: str = attr.ib()
    client_secret: str = attr.ib(repr=secrepr)

    refresh_token: str = attr.ib(repr=secrepr)
    refresh_token_expire_time: Optional[datetime] = attr.ib()

    schema: str = attr.ib()
    db_name: str = attr.ib()
    warehouse: str = attr.ib()

lib/dl_connector_snowflake/dl_connector_snowflake/core/dto.py
datalens-tech/datalens-backend
import sqlalchemy

from dl_core.connectors.base.error_transformer import (
    ChainedDbErrorTransformer,
    DbErrorTransformer,
)
from dl_core.connectors.base.error_transformer import ErrorTransformerRule as Rule
from dl_core.connectors.base.error_transformer import ExcMatchCondition

from dl_connector_snowflake.core.exc import SnowflakeAccessTokenError


def is_access_token_error() -> ExcMatchCondition:
    def func(exc: Exception) -> bool:
        if isinstance(exc, sqlalchemy.exc.DatabaseError):
            orig = getattr(exc, "orig", None)
            if orig and (getattr(orig, "errno", None) == 250001):
                return True

        return False

    return func


snowflake_error_transformer: DbErrorTransformer = ChainedDbErrorTransformer(
    [Rule(when=is_access_token_error(), then_raise=SnowflakeAccessTokenError)]
)

lib/dl_connector_snowflake/dl_connector_snowflake/core/error_transformer.py
datalens-tech/datalens-backend
from dl_core.exc import (
    DatabaseUnavailable,
    SourceAccessInvalidToken,
    UserQueryAccessDenied,
)


class SnowflakeAccessTokenError(UserQueryAccessDenied):
    default_message = "Access token is expired or invalid"


class SnowflakeGetAccessTokenError(DatabaseUnavailable):
    default_message = "Could not get access toke for the snowflake connection"


class SnowflakeRefreshTokenInvalid(SourceAccessInvalidToken):
    default_message = "OAuth refresh token for the Snowflake connections is invalid or expired"

lib/dl_connector_snowflake/dl_connector_snowflake/core/exc.py
datalens-tech/datalens-backend
import logging

from dl_core.connectors.base.lifecycle import ConnectionLifecycleManager

from dl_connector_snowflake.core.us_connection import ConnectionSQLSnowFlake


LOGGER = logging.getLogger(__name__)


class SnowFlakeConnectionLifecycleManager(ConnectionLifecycleManager[ConnectionSQLSnowFlake]):
    ENTRY_CLS = ConnectionSQLSnowFlake

lib/dl_connector_snowflake/dl_connector_snowflake/core/lifecycle.py
datalens-tech/datalens-backend
from dl_constants.enums import NotificationLevel
from dl_core.reporting.notifications import BaseNotification

from dl_connector_snowflake.core.constants import NOTIF_TYPE_SF_REFRESH_TOKEN_SOON_TO_EXPIRE


class SnowflakeRefreshTokenSoonToExpire(BaseNotification):
    type = NOTIF_TYPE_SF_REFRESH_TOKEN_SOON_TO_EXPIRE
    _title = "Snowflake connection token is going to expire soon"
    _message = """Snowflake connection token is going to expire soon. \n\n
Please go to the connection details page and obtain a new refresh token. \n\n"""
    _level = NotificationLevel.warning

lib/dl_connector_snowflake/dl_connector_snowflake/core/notifications.py
datalens-tech/datalens-backend
from marshmallow import fields as ma_fields
from marshmallow.validate import Regexp

from dl_core.us_manager.storage_schemas.connection import (
    BaseConnectionDataStorageSchema,
    CacheableConnectionDataSchemaMixin,
    SubselectConnectionDataSchemaMixin,
)

from dl_connector_snowflake.core.constants import ACCOUNT_NAME_RE
from dl_connector_snowflake.core.us_connection import ConnectionSQLSnowFlake


class SnowFlakeConnectionDataStorageSchema(
    BaseConnectionDataStorageSchema[ConnectionSQLSnowFlake.DataModel],
    CacheableConnectionDataSchemaMixin,
    SubselectConnectionDataSchemaMixin,
):
    TARGET_CLS = ConnectionSQLSnowFlake.DataModel

    account_name = ma_fields.String(allow_none=False, required=True, validate=Regexp(regex=ACCOUNT_NAME_RE))
    user_name = ma_fields.String(allow_none=False, required=True)
    user_role = ma_fields.String(allow_none=True, required=False)
    client_id = ma_fields.String(allow_none=False, required=True)
    client_secret = ma_fields.String(allow_none=False, required=True)
    refresh_token = ma_fields.String(allow_none=False, required=True)
    refresh_token_expire_time = ma_fields.DateTime(allow_none=True, required=False)

    schema = ma_fields.String(allow_none=False, required=True)
    db_name = ma_fields.String(allow_none=False, required=True)
    warehouse = ma_fields.String(allow_none=False, required=True)

lib/dl_connector_snowflake/dl_connector_snowflake/core/storage_schemas/connection.py
datalens-tech/datalens-backend
from dl_core.us_manager.storage_schemas.data_source_spec_base import (
    BaseSQLDataSourceSpecStorageSchema,
    SubselectSQLDataSourceSpecStorageSchemaMixin,
    TableSQLDataSourceSpecStorageSchemaMixin,
)

from dl_connector_snowflake.core.data_source_spec import (
    SnowFlakeSubselectDataSourceSpec,
    SnowFlakeTableDataSourceSpec,
)


class SnowFlakeTableDataSourceSpecStorageSchema(
    TableSQLDataSourceSpecStorageSchemaMixin,
    BaseSQLDataSourceSpecStorageSchema,
):
    TARGET_CLS = SnowFlakeTableDataSourceSpec  # type: ignore


class SnowFlakeSubselectDataSourceSpecStorageSchema(
    SubselectSQLDataSourceSpecStorageSchemaMixin,
    BaseSQLDataSourceSpecStorageSchema,
):
    TARGET_CLS = SnowFlakeSubselectDataSourceSpec  # type: ignore

lib/dl_connector_snowflake/dl_connector_snowflake/core/storage_schemas/data_source_spec.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

import attr

from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO
from dl_core.utils import secrepr

from dl_connector_snowflake.core.dto import SnowFlakeConnDTO


@attr.s(frozen=True, kw_only=True)
class SnowFlakeConnTargetDTO(ConnTargetDTO):
    account_name: str = attr.ib()
    user_name: str = attr.ib()
    user_role: Optional[str] = attr.ib()

    schema: str = attr.ib()
    db_name: str = attr.ib()
    warehouse: str = attr.ib()

    access_token: str = attr.ib(repr=secrepr)

    @classmethod
    def from_dto(cls, dto: SnowFlakeConnDTO, access_token: str) -> SnowFlakeConnTargetDTO:
        return cls(
            access_token=access_token,
            account_name=dto.account_name,
            user_name=dto.user_name,
            user_role=dto.user_role,
            schema=dto.schema,
            db_name=dto.db_name,
            warehouse=dto.warehouse,
            conn_id=dto.conn_id,
        )

    def get_effective_host(self) -> Optional[str]:
        return None

lib/dl_connector_snowflake/dl_connector_snowflake/core/target_dto.py
datalens-tech/datalens-backend
from datetime import datetime
from typing import Optional
import uuid

from dl_constants.enums import RawSQLLevel
from dl_core.us_manager.us_manager_sync import SyncUSManager

from dl_connector_snowflake.core.constants import CONNECTION_TYPE_SNOWFLAKE
from dl_connector_snowflake.core.us_connection import ConnectionSQLSnowFlake


def make_snowflake_saved_connection(
    sync_usm: SyncUSManager,
    account_name: str,
    user_name: str,
    client_id: str,
    client_secret: str,
    refresh_token: str,
    refresh_token_expire_time: Optional[datetime],
    schema: Optional[str] = None,
    db_name: Optional[str] = None,
    warehouse: Optional[str] = None,
    user_role: Optional[str] = None,
    raw_sql_level: RawSQLLevel = RawSQLLevel.off,
) -> ConnectionSQLSnowFlake:
    conn_name = f"snowflake test conn: {uuid.uuid4()}"
    conn = ConnectionSQLSnowFlake.create_from_dict(
        data_dict=ConnectionSQLSnowFlake.DataModel(
            account_name=account_name,
            user_name=user_name,
            user_role=user_role,
            client_id=client_id,
            client_secret=client_secret,
            refresh_token=refresh_token,
            refresh_token_expire_time=refresh_token_expire_time,
            schema=schema,
            db_name=db_name,
            warehouse=warehouse,
            raw_sql_level=raw_sql_level,
        ),
        ds_key=conn_name,
        type_=CONNECTION_TYPE_SNOWFLAKE.name,
        us_manager=sync_usm,
    )
    sync_usm.save(conn)
    return conn

lib/dl_connector_snowflake/dl_connector_snowflake/core/testing/connection.py
datalens-tech/datalens-backend
import abc
from typing import ClassVar

import attr

from dl_testing.env_params.getter import EnvParamGetter


class SnowFlakeSecretReaderBase(abc.ABC):
    KEY_CONFIG: ClassVar[str] = "SNOWFLAKE_CONFIG"
    KEY_CLIENT_SECRET: ClassVar[str] = "SNOWFLAKE_CLIENT_SECRET"
    KEY_REFRESH_TOKEN_EXPIRED: ClassVar[str] = "SNOWFLAKE_REFRESH_TOKEN_EXPIRED"
    KEY_REFRESH_TOKEN_X: ClassVar[str] = "SNOWFLAKE_REFRESH_TOKEN_X"

    @property
    @abc.abstractmethod
    def project_config(self) -> dict:
        raise NotImplementedError

    def get_account_name(self) -> str:
        return self.project_config["account_name"]

    def get_user_name(self) -> str:
        return self.project_config["user_name"]

    def get_user_role(self) -> str:
        return self.project_config["user_role"]

    def get_login_name(self) -> str:
        return self.project_config["login_name"]

    def get_database(self) -> str:
        return self.project_config["database"]

    def get_schema(self) -> str:
        return self.project_config["schema"]

    def get_table_name(self) -> str:
        return self.project_config["table_name"]

    def get_warehouse(self) -> str:
        return self.project_config["warehouse"]

    def get_client_id(self) -> str:
        return self.project_config["client_id"]

    @abc.abstractmethod
    def get_client_secret(self) -> str:
        raise NotImplementedError

    @abc.abstractmethod
    def get_refresh_token_expired(self) -> str:
        raise NotImplementedError

    @abc.abstractmethod
    def get_refresh_token_x(self) -> str:
        raise NotImplementedError


@attr.s
class SnowFlakeSecretReader(SnowFlakeSecretReaderBase):
    _env_param_getter: EnvParamGetter = attr.ib()
    _project_config: dict = attr.ib(init=False)

    @_project_config.default
    def _make_project_config(self) -> dict:
        return self._env_param_getter.get_json_value(self.KEY_CONFIG)

    @property
    def project_config(self) -> dict:
        return self._project_config

    def get_client_secret(self) -> str:
        return self._env_param_getter.get_str_value(self.KEY_CLIENT_SECRET)

    def get_refresh_token_expired(self) -> str:
        return self._env_param_getter.get_str_value(self.KEY_REFRESH_TOKEN_EXPIRED)

    def get_refresh_token_x(self) -> str:
        return self._env_param_getter.get_str_value(self.KEY_REFRESH_TOKEN_X)

lib/dl_connector_snowflake/dl_connector_snowflake/core/testing/secrets.py
datalens-tech/datalens-backend
from snowflake import sqlalchemy as ssa

from dl_constants.enums import UserDataType
from dl_core.db.conversion_base import (
    TypeTransformer,
    make_native_type,
)

from dl_connector_snowflake.core.constants import CONNECTION_TYPE_SNOWFLAKE


SNOW_TYPES_INT = frozenset(
    (
        ssa.BIGINT,
        ssa.SMALLINT,
        ssa.INTEGER,
        ssa.TINYINT,
        ssa.SMALLINT,
    )
)
SNOW_TYPES_FLOAT = frozenset(
    (
        ssa.FLOAT,
        ssa.DECIMAL,
        ssa.DEC,
        ssa.DOUBLE,
        ssa.REAL,
    )
)
SNOW_TYPES_STRING = frozenset(
    (
        ssa.STRING,
        ssa.VARCHAR,
        ssa.CHAR,
        ssa.CHARACTER,
        ssa.TEXT,
    )
)


class SnowFlakeTypeTransformer(TypeTransformer):
    # https://docs.snowflake.com/en/sql-reference/data-types.html

    conn_type = CONNECTION_TYPE_SNOWFLAKE
    native_to_user_map = {
        make_native_type(CONNECTION_TYPE_SNOWFLAKE, ssa.DATE): UserDataType.date,
        make_native_type(CONNECTION_TYPE_SNOWFLAKE, ssa.DATETIME): UserDataType.genericdatetime,
        **{make_native_type(CONNECTION_TYPE_SNOWFLAKE, t): UserDataType.integer for t in SNOW_TYPES_INT},
        **{make_native_type(CONNECTION_TYPE_SNOWFLAKE, t): UserDataType.float for t in SNOW_TYPES_FLOAT},
        **{make_native_type(CONNECTION_TYPE_SNOWFLAKE, t): UserDataType.string for t in SNOW_TYPES_STRING},
        make_native_type(CONNECTION_TYPE_SNOWFLAKE, ssa.BOOLEAN): UserDataType.boolean,
        # todo: review datetime/genericdatime/datetimetz/timestamp
        # todo: array, geo
    }
    user_to_native_map = {
        UserDataType.date: make_native_type(CONNECTION_TYPE_SNOWFLAKE, ssa.DATE),
        UserDataType.genericdatetime: make_native_type(CONNECTION_TYPE_SNOWFLAKE, ssa.DATETIME),
        UserDataType.integer: make_native_type(CONNECTION_TYPE_SNOWFLAKE, ssa.INT),
        UserDataType.string: make_native_type(CONNECTION_TYPE_SNOWFLAKE, ssa.STRING),
        UserDataType.boolean: make_native_type(CONNECTION_TYPE_SNOWFLAKE, ssa.BOOLEAN),
    }
    casters = {
        **TypeTransformer.casters,
        # UserDataType.date: GSheetsDateTypeCaster(),
        # UserDataType.datetime: GSheetsDatetimeTypeCaster(),
        # UserDataType.genericdatetime: GSheetsGenericDatetimeTypeCaster(),
    }

lib/dl_connector_snowflake/dl_connector_snowflake/core/type_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

from datetime import datetime
import re
from typing import (
    TYPE_CHECKING,
    Callable,
    ClassVar,
    Optional,
)

import attr
import marshmallow as ma

from dl_api_commons.reporting.models import NotificationReportingRecord
from dl_core.base_models import (
    ConnCacheableDataModelMixin,
    ConnectionDataModelBase,
    ConnSubselectDataModelMixin,
)
from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.reporting.notifications import get_notification_record
from dl_core.us_connection_base import (
    ConnectionBase,
    ConnectionSQL,
    DataSourceTemplate,
)
from dl_i18n.localizer_base import Localizer
from dl_utils.utils import DataKey

from dl_connector_snowflake.auth import SFAuthProvider
from dl_connector_snowflake.core.constants import (
    ACCOUNT_NAME_RE,
    CONNECTION_TYPE_SNOWFLAKE,
    NOTIF_TYPE_SF_REFRESH_TOKEN_SOON_TO_EXPIRE,
    SOURCE_TYPE_SNOWFLAKE_SUBSELECT,
    SOURCE_TYPE_SNOWFLAKE_TABLE,
)
from dl_connector_snowflake.core.dto import SnowFlakeConnDTO


if TYPE_CHECKING:
    from dl_core.services_registry import ServicesRegistry


class ConnectionSQLSnowFlake(ConnectionSQL):
    conn_type = CONNECTION_TYPE_SNOWFLAKE
    has_schema: ClassVar[bool] = True
    default_schema_name = None
    source_type = SOURCE_TYPE_SNOWFLAKE_TABLE
    allowed_source_types = frozenset((SOURCE_TYPE_SNOWFLAKE_TABLE, SOURCE_TYPE_SNOWFLAKE_SUBSELECT))
    allow_dashsql: ClassVar[bool] = False
    allow_cache: ClassVar[bool] = True
    is_always_user_source: ClassVar[bool] = True

    @attr.s(kw_only=True)
    class DataModel(ConnCacheableDataModelMixin, ConnSubselectDataModelMixin, ConnectionDataModelBase):
        account_name: str = attr.ib()
        user_name: str = attr.ib()
        user_role: Optional[str] = attr.ib(default=None)
        client_id: str = attr.ib()
        client_secret: str = attr.ib()
        refresh_token: str = attr.ib()
        refresh_token_expire_time: Optional[datetime] = attr.ib(default=None)

        schema: Optional[str] = attr.ib()
        db_name: Optional[str] = attr.ib()
        warehouse: Optional[str] = attr.ib()

        @classmethod
        def get_secret_keys(cls) -> set[DataKey]:
            return {DataKey(parts=("client_secret",)), DataKey(parts=("refresh_token",))}

    async def validate_new_data(
        self,
        services_registry: ServicesRegistry,
        changes: Optional[dict] = None,
        original_version: Optional[ConnectionBase] = None,
    ) -> None:
        if not re.fullmatch(ACCOUNT_NAME_RE, self.data.account_name):
            raise ma.ValidationError(message="Field account_name is not valid")

    def check_for_notifications(self) -> list[Optional[NotificationReportingRecord]]:
        notifications = super().check_for_notifications()

        sf_auth_provider = SFAuthProvider.from_dto(self.get_conn_dto())
        if sf_auth_provider.should_notify_refresh_token_to_expire_soon():
            notifications.append(get_notification_record(NOTIF_TYPE_SF_REFRESH_TOKEN_SOON_TO_EXPIRE))

        return notifications

    @property
    def schema_name(self) -> Optional[str]:
        return self.data.schema

    def get_conn_dto(self) -> SnowFlakeConnDTO:
        return SnowFlakeConnDTO(
            conn_id=self.uuid,
            account_name=self.data.account_name,
            user_name=self.data.user_name,
            user_role=self.data.user_role,
            client_id=self.data.client_id,
            client_secret=self.data.client_secret,
            refresh_token=self.data.refresh_token,
            refresh_token_expire_time=self.data.refresh_token_expire_time,
            schema=self.data.schema,
            db_name=self.data.db_name,
            warehouse=self.data.warehouse,
        )

    def get_parameter_combinations(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[dict]:
        # todo: for initial implementation using only schema and db name from conn dto
        tables = self.get_tables(
            conn_executor_factory=conn_executor_factory,
            schema_name=self.data.schema,
            db_name=self.data.db_name,
        )
        return [dict(schema_name=tid.schema_name, table_name=tid.table_name) for tid in tables]

    def get_data_source_template_group(self, parameters: dict) -> list[str]:
        return [
            val
            for val in (
                parameters.get("account_name"),
                parameters.get("user_name"),
                parameters.get("user_role"),
                parameters.get("client_id"),
                parameters.get("client_secret"),
                parameters.get("refresh_token"),
                parameters.get("refresh_token_expire_time"),
                parameters.get("schema"),
                parameters.get("db_name"),
                parameters.get("warehouse"),
            )
            if val is not None
        ]

    @property
    def allow_public_usage(self) -> bool:
        return False

    def get_data_source_template_templates(self, localizer: Localizer) -> list[DataSourceTemplate]:
        return self._make_subselect_templates(source_type=SOURCE_TYPE_SNOWFLAKE_SUBSELECT, localizer=localizer)

lib/dl_connector_snowflake/dl_connector_snowflake/core/us_connection.py
datalens-tech/datalens-backend
from dl_db_testing.connectors.base.connector import DbTestingConnector

from dl_connector_snowflake.db_testing.engine_wrapper import SnowFlakeEngineWrapper


class SnowFlakeDbTestingConnector(DbTestingConnector):
    engine_wrapper_classes = (SnowFlakeEngineWrapper,)

lib/dl_connector_snowflake/dl_connector_snowflake/db_testing/connector.py
datalens-tech/datalens-backend
from typing import Optional

import attr
import sqlalchemy as sa

from dl_db_testing.database.engine_wrapper import (
    DbEngineConfig,
    EngineWrapperBase,
)


@attr.s(frozen=True)
class SnowFlakeDbEngineConfig(DbEngineConfig):
    db_name: str = attr.ib(kw_only=True)


class SnowFlakeEngineWrapper(EngineWrapperBase):
    CONFIG_CLS = SnowFlakeDbEngineConfig
    URL_PREFIX = "snowflake"

    @property
    def config(self) -> SnowFlakeDbEngineConfig:
        assert isinstance(self._config, SnowFlakeDbEngineConfig)
        return self._config

    @property
    def name(self) -> str:
        name = self.execute("SELECT CURRENT_DATABASE()").scalar()
        assert isinstance(name, str)
        return name

    def get_version(self) -> Optional[str]:
        return self.name

    def drop_table(self, db_name: str, table: sa.Table) -> None:
        table.drop(bind=self.engine, checkfirst=True)

    def get_conn_credentials(self, full: bool = False) -> dict:
        return dict(
            db_name=self.config.db_name,
        )

lib/dl_connector_snowflake/dl_connector_snowflake/db_testing/engine_wrapper.py
datalens-tech/datalens-backend
from snowflake.sqlalchemy.snowdialect import SnowflakeDialect as SASnowflakeDialect

from dl_formula.connectors.base.connector import FormulaConnector

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as SnowFlakeDialectNS
from dl_connector_snowflake.formula.definitions.all import DEFINITIONS
from dl_connector_snowflake.formula.literal import SnowFlakeLiteralizer


class SnowFlakeFormulaConnector(FormulaConnector):
    dialect_ns_cls = SnowFlakeDialectNS
    dialects = SnowFlakeDialectNS.SNOWFLAKE
    default_dialect = SnowFlakeDialectNS.SNOWFLAKE
    op_definitions = DEFINITIONS
    literalizer_cls = SnowFlakeLiteralizer
    sa_dialect = SASnowflakeDialect()

lib/dl_connector_snowflake/dl_connector_snowflake/formula/connector.py
datalens-tech/datalens-backend
from dl_formula.core.dialect import (
    DialectName,
    DialectNamespace,
    simple_combo,
)


DIALECT_NAME_SNOWFLAKE = DialectName.declare("SNOWFLAKE")


class SnowFlakeDialect(DialectNamespace):
    SNOWFLAKE = simple_combo(name=DIALECT_NAME_SNOWFLAKE)

lib/dl_connector_snowflake/dl_connector_snowflake/formula/constants.py
datalens-tech/datalens-backend
from dl_connector_snowflake.formula.definitions.conditional_blocks import DEFINITIONS_COND_BLOCKS
from dl_connector_snowflake.formula.definitions.functions_aggregation import DEFINITIONS_AGG
from dl_connector_snowflake.formula.definitions.functions_datetime import DEFINITIONS_DATETIME
from dl_connector_snowflake.formula.definitions.functions_logical import DEFINITIONS_LOGICAL
from dl_connector_snowflake.formula.definitions.functions_markup import DEFINITIONS_MARKUP
from dl_connector_snowflake.formula.definitions.functions_math import DEFINITIONS_MATH
from dl_connector_snowflake.formula.definitions.functions_string import DEFINITIONS_STRING
from dl_connector_snowflake.formula.definitions.functions_type import DEFINITIONS_TYPE
from dl_connector_snowflake.formula.definitions.operators_binary import DEFINITIONS_BINARY
from dl_connector_snowflake.formula.definitions.operators_ternary import DEFINITIONS_TERNARY
from dl_connector_snowflake.formula.definitions.operators_unary import DEFINITIONS_UNARY


DEFINITIONS = [
    *DEFINITIONS_UNARY,
    *DEFINITIONS_BINARY,
    *DEFINITIONS_TERNARY,
    *DEFINITIONS_MARKUP,
    *DEFINITIONS_MATH,
    *DEFINITIONS_TYPE,
    *DEFINITIONS_LOGICAL,
    *DEFINITIONS_STRING,
    *DEFINITIONS_COND_BLOCKS,
    *DEFINITIONS_AGG,
    *DEFINITIONS_DATETIME,
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/all.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.conditional_blocks as base

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


V = TranslationVariant.make


DEFINITIONS_COND_BLOCKS = [
    # _case_block_
    base.CaseBlock.for_dialect(D.SNOWFLAKE),
    # _if_block_
    base.IfBlock3(
        variants=[
            V(D.SNOWFLAKE, sa.func.IFF),
        ]
    ),
    base.IfBlockMulti.for_dialect(D.SNOWFLAKE),
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/conditional_blocks.py
datalens-tech/datalens-backend
import sqlalchemy as sa
from sqlalchemy import within_group

from dl_formula.core.datatype import DataType
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common import quantile_value
import dl_formula.definitions.functions_aggregation as base
from dl_formula.definitions.literals import un_literal

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


V = TranslationVariant.make


class AggMedianIntSF(base.AggMedian):
    name = "median"
    argument_types = [
        ArgTypeSequence([{DataType.INTEGER}]),
    ]
    variants = [
        V(
            D.SNOWFLAKE,
            lambda expr: sa.func.TO_NUMBER(
                within_group(sa.func.PERCENTILE_DISC(0.5), expr),
            ),
        ),
    ]


class AggMedianFloatSF(base.AggMedian):
    name = "median"
    argument_types = [
        ArgTypeSequence([{DataType.FLOAT}]),
    ]
    variants = [
        V(
            D.SNOWFLAKE,
            lambda expr: within_group(sa.func.PERCENTILE_DISC(0.5), expr),
        ),
    ]


class AggMedianDateSF(base.AggMedian):
    name = "median"
    argument_types = [ArgTypeSequence([{DataType.DATE}])]
    variants = [
        V(
            D.SNOWFLAKE,
            lambda expr: sa.func.TO_DATE(
                sa.func.TO_TIMESTAMP(
                    within_group(
                        sa.func.PERCENTILE_DISC(0.5),
                        sa.func.DATE_PART("EPOCH_SECOND", expr),
                    )
                )
            ),
        ),
    ]


class AggMedianDateTimeSF(base.AggMedian):
    name = "median"
    argument_types = [ArgTypeSequence([{DataType.DATETIME, DataType.GENERICDATETIME}])]
    variants = [
        V(
            D.SNOWFLAKE,
            lambda expr: sa.func.TO_TIMESTAMP(
                within_group(
                    sa.func.PERCENTILE_DISC(0.5),
                    sa.func.DATE_PART("EPOCH_SECOND", expr),
                )
            ),
        ),
    ]


DEFINITIONS_AGG = [
    # any
    base.AggAny(
        variants=[
            V(D.SNOWFLAKE, sa.func.ANY_VALUE),
        ]
    ),
    # avg
    base.AggAvgFromNumber.for_dialect(D.SNOWFLAKE),
    # base.AggAvgFromDate.for_dialect(D.SNOWFLAKE),
    base.AggAvgFromDate(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda date_val: sa.func.TO_DATE(
                    sa.func.TO_TIMESTAMP(
                        sa.func.AVG(sa.func.DATE_PART("EPOCH_SECOND", date_val)),
                    )
                ),
            )
        ]
    ),
    base.AggAvgFromDatetime.for_dialect(D.SNOWFLAKE),
    base.AggAvgFromDatetimeTZ.for_dialect(D.SNOWFLAKE),
    # avg_if
    base.AggAvgIf.for_dialect(D.SNOWFLAKE),
    # count
    base.AggCount0.for_dialect(D.SNOWFLAKE),
    base.AggCount1.for_dialect(D.SNOWFLAKE),
    # count_if
    base.AggCountIf(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda expr: sa.func.IFNULL(sa.func.COUNT_IF(expr), 0),
            )
        ]
    ),
    # countd
    base.AggCountd.for_dialect(D.SNOWFLAKE),
    # countd_approx
    base.AggCountdApprox(
        variants=[
            V(D.SNOWFLAKE, sa.func.APPROX_COUNT_DISTINCT),
        ]
    ),
    # countd_if
    base.AggCountdIf.for_dialect(D.SNOWFLAKE),
    # max
    base.AggMax.for_dialect(D.SNOWFLAKE),
    # median
    AggMedianIntSF.for_dialect(D.SNOWFLAKE),
    AggMedianFloatSF.for_dialect(D.SNOWFLAKE),
    AggMedianDateSF.for_dialect(D.SNOWFLAKE),
    AggMedianDateTimeSF.for_dialect(D.SNOWFLAKE),
    # quantile
    base.AggQuantile(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda expr, quant: within_group(sa.func.percentile_disc(quantile_value(un_literal(quant))), expr),
            ),
        ]
    ),
    # min
    base.AggMin.for_dialect(D.SNOWFLAKE),
    # stdev
    base.AggStdev.for_dialect(D.SNOWFLAKE),
    # stdevp
    base.AggStdevp.for_dialect(D.SNOWFLAKE),
    # sum
    base.AggSum.for_dialect(D.SNOWFLAKE),
    # sum_if
    base.AggSumIf.for_dialect(D.SNOWFLAKE),
    # var
    base.AggVar.for_dialect(D.SNOWFLAKE),
    # varp
    base.AggVarp.for_dialect(D.SNOWFLAKE),
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/functions_aggregation.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common import raw_sql
import dl_formula.definitions.functions_datetime as base

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


V = TranslationVariant.make


DEFINITIONS_DATETIME = [
    # dateadd
    base.FuncDateadd3Legacy.for_dialect(D.SNOWFLAKE),
    base.FuncDateadd1.for_dialect(D.SNOWFLAKE),
    base.FuncDateadd2Unit.for_dialect(D.SNOWFLAKE),
    base.FuncDateadd2Number.for_dialect(D.SNOWFLAKE),
    base.FuncDateadd3DateConstNum(variants=[V(D.SNOWFLAKE, lambda date, what, num: sa.func.DATEADD(what, num, date))]),
    base.FuncDateadd3DatetimeConstNum(
        variants=[V(D.SNOWFLAKE, lambda date, what, num: sa.func.DATEADD(what, num, date))]
    ),
    # datepart
    base.FuncDatepart2.for_dialect(D.SNOWFLAKE),
    base.FuncDatepart3Const.for_dialect(D.SNOWFLAKE),
    base.FuncDatepart3NonConst.for_dialect(D.SNOWFLAKE),
    # datetrunc
    base.FuncDatetrunc2Date(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda date, unit: (
                    sa.func.DATE_TRUNC(
                        raw_sql(base.norm_datetrunc_unit(unit).replace("week", "week")),
                        date,
                    )
                    if base.norm_datetrunc_unit(unit) in {"year", "quarter", "month", "week", "day"}
                    else date
                ),
            ),
        ]
    ),
    base.FuncDatetrunc2Datetime(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda date, unit: sa.func.DATE_TRUNC(
                    raw_sql(base.norm_datetrunc_unit(unit).replace("week", "week")),
                    date,
                ),
            ),
        ]
    ),
    # day
    base.FuncDay(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.EXTRACT("DAY", value)),
        ]
    ),
    # dayofweek
    base.FuncDayofweek1(
        variants=[
            V(
                D.SNOWFLAKE,
                sa.func.DAYOFWEEKISO,
            )
        ]
    ),
    base.FuncDayofweek2(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda date, firstday: sa.func.MOD(sa.func.DAYOFWEEKISO(date) - base.norm_fd(firstday), 7) + 1,
            ),
        ]
    ),
    # genericnow
    base.FuncGenericNow(
        variants=[
            V(D.SNOWFLAKE, lambda: sa.func.CONVERT_TIMEZONE("UTC", sa.func.LOCALTIMESTAMP())),
        ]
    ),
    # hour
    base.FuncHourDate.for_dialect(D.SNOWFLAKE),
    base.FuncHourDatetime(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.EXTRACT("HOUR", value)),
        ]
    ),
    # minute
    base.FuncMinuteDate.for_dialect(D.SNOWFLAKE),
    base.FuncMinuteDatetime(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.EXTRACT("MINUTE", value)),
        ]
    ),
    # month
    base.FuncMonth(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.EXTRACT("MONTH", value)),
        ]
    ),
    # now
    base.FuncNow(
        variants=[
            V(D.SNOWFLAKE, lambda: sa.func.CONVERT_TIMEZONE("UTC", sa.func.LOCALTIMESTAMP())),
        ]
    ),
    # quarter
    base.FuncQuarter(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.EXTRACT("QUARTER", value)),
        ]
    ),
    # second
    base.FuncSecondDate.for_dialect(D.SNOWFLAKE),
    base.FuncSecondDatetime(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.EXTRACT("SECOND", value)),
        ]
    ),
    # today
    base.FuncToday(
        variants=[
            V(D.SNOWFLAKE, sa.func.CURRENT_DATE),
        ]
    ),
    # week
    base.FuncWeek(
        variants=[
            V(D.SNOWFLAKE, sa.func.WEEKISO),
        ]
    ),
    # year
    base.FuncYear(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.EXTRACT("YEAR", value)),
        ]
    ),
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/functions_datetime.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_logical as base

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


V = TranslationVariant.make


DEFINITIONS_LOGICAL = [
    # case
    base.FuncCase.for_dialect(D.SNOWFLAKE),
    # if
    base.FuncIf(
        variants=[
            V(D.SNOWFLAKE, sa.func.IFF),
        ]
    ),
    # ifnull
    base.FuncIfnull(
        variants=[
            V(D.SNOWFLAKE, sa.func.IFNULL),
        ]
    ),
    # isnan
    base.FuncIsnan(
        variants=[
            V(D.SNOWFLAKE, sa.func.IS_NAN),
        ]
    ),
    # isnull
    base.FuncIsnull.for_dialect(D.SNOWFLAKE),
    # zn
    base.FuncZn(
        variants=[
            V(D.SNOWFLAKE, lambda x: sa.func.IFNULL(x, 0)),
        ]
    ),
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/functions_logical.py
datalens-tech/datalens-backend
import dl_formula.definitions.functions_markup as base

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


DEFINITIONS_MARKUP = [
    # +
    base.BinaryPlusMarkup.for_dialect(D.SNOWFLAKE),
    # __str
    base.FuncInternalStrConst.for_dialect(D.SNOWFLAKE),
    base.FuncInternalStr.for_dialect(D.SNOWFLAKE),
    # bold
    base.FuncBold.for_dialect(D.SNOWFLAKE),
    # italic
    base.FuncItalics.for_dialect(D.SNOWFLAKE),
    # markup
    base.ConcatMultiMarkup.for_dialect(D.SNOWFLAKE),
    # url
    base.FuncUrl.for_dialect(D.SNOWFLAKE),
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/functions_markup.py
datalens-tech/datalens-backend
import math

import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.functions_math as base

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


V = TranslationVariant.make


class FuncDivSF(base.FuncDiv):
    variants = [V(D.SNOWFLAKE, lambda x, y: sa.cast(sa.func.trunc(x / y, 0), sa.BIGINT))]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER, DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.INTEGER, DataType.FLOAT]),
        ArgTypeSequence([DataType.FLOAT, DataType.INTEGER]),
    ]


class FuncDivSafe2SF(base.FuncDivSafe2):
    variants = [
        V(
            D.SNOWFLAKE,
            lambda x, y: sa.func.IFF(
                y != 0,
                sa.cast(sa.func.trunc(x / y, 0), sa.BIGINT),
                None,
            ),
        )
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER, DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.INTEGER, DataType.FLOAT]),
        ArgTypeSequence([DataType.FLOAT, DataType.INTEGER]),
    ]


class FuncDivSafe3SF(base.FuncDivSafe3):
    variants = [
        V(
            D.SNOWFLAKE,
            lambda x, y, default: sa.func.IFF(
                y != 0,
                sa.cast(sa.func.trunc(x / y, 0), sa.BIGINT),
                default,
            ),
        )
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER, DataType.INTEGER, DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT, DataType.INTEGER]),
        ArgTypeSequence([DataType.INTEGER, DataType.FLOAT, DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT, DataType.INTEGER, DataType.INTEGER]),
    ]


DEFINITIONS_MATH = [
    # abs
    base.FuncAbs.for_dialect(D.SNOWFLAKE),
    # acos
    base.FuncAcos.for_dialect(D.SNOWFLAKE),
    # asin
    base.FuncAsin.for_dialect(D.SNOWFLAKE),
    # atan
    base.FuncAtan.for_dialect(D.SNOWFLAKE),
    # atan2
    base.FuncAtan2.for_dialect(D.SNOWFLAKE),
    # ceiling
    base.FuncCeiling.for_dialect(D.SNOWFLAKE),
    # cos
    base.FuncCos.for_dialect(D.SNOWFLAKE),
    # cot
    base.FuncCot.for_dialect(D.SNOWFLAKE),
    # degrees
    base.FuncDegrees(
        variants=[
            V(D.SNOWFLAKE, lambda x: x / math.pi * 180),
        ]
    ),
    # div
    FuncDivSF(),
    # div_safe
    FuncDivSafe2SF(),
    FuncDivSafe3SF(),
    # exp
    base.FuncExp.for_dialect(D.SNOWFLAKE),
    # fdiv_safe
    base.FuncFDivSafe2.for_dialect(D.SNOWFLAKE),
    base.FuncFDivSafe3.for_dialect(D.SNOWFLAKE),
    # floor
    base.FuncFloor.for_dialect(D.SNOWFLAKE),
    # greatest
    base.FuncGreatest1.for_dialect(D.SNOWFLAKE),
    base.FuncGreatestMain.for_dialect(D.SNOWFLAKE),
    # base.GreatestMulti.for_dialect(D.SNOWFLAKE),  # FIXME
    base.GreatestMulti(
        variants=[
            V(D.SNOWFLAKE, sa.func.GREATEST),
        ]
    ),
    # least
    base.FuncLeast1.for_dialect(D.SNOWFLAKE),
    base.FuncLeastMain.for_dialect(D.SNOWFLAKE),
    # base.LeastMulti.for_dialect(D.SNOWFLAKE),  # FIXME
    base.LeastMulti(
        variants=[
            V(D.SNOWFLAKE, sa.func.LEAST),
        ]
    ),
    # ln
    base.FuncLn.for_dialect(D.SNOWFLAKE),
    # log
    base.FuncLog(
        variants=[
            V(D.SNOWFLAKE, lambda x, base: sa.func.LOG(base, x)),
        ]
    ),
    # log10
    base.FuncLog10(
        variants=[
            V(D.SNOWFLAKE, lambda x: sa.func.LOG(10, x)),
        ]
    ),
    # pi
    base.FuncPi(
        variants=[
            V(D.SNOWFLAKE, lambda: sa.literal(math.pi)),
        ]
    ),
    # power
    base.FuncPower(
        variants=[
            V(D.SNOWFLAKE, sa.func.POWER),
        ]
    ),
    # radians
    base.FuncRadians(
        variants=[
            V(D.SNOWFLAKE, lambda x: x * math.pi / 180),
        ]
    ),
    # round
    base.FuncRound1.for_dialect(D.SNOWFLAKE),
    base.FuncRound2.for_dialect(D.SNOWFLAKE),
    # sign
    base.FuncSign.for_dialect(D.SNOWFLAKE),
    # sin
    base.FuncSin.for_dialect(D.SNOWFLAKE),
    # sqrt
    base.FuncSqrt.for_dialect(D.SNOWFLAKE),
    # square
    base.FuncSquare(
        variants=[
            V(D.SNOWFLAKE, lambda x: sa.func.POW(x, 2)),
        ]
    ),
    # tan
    base.FuncTan.for_dialect(D.SNOWFLAKE),
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/functions_math.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import (
    TranslationVariant,
    TranslationVariantWrapped,
)
import dl_formula.definitions.functions_string as base
from dl_formula.shortcuts import n

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class FuncIEndswithNonConstSF(base.FuncIEndswithNonConst):
    variants = [
        VW(
            D.SNOWFLAKE,
            lambda s, substr: n.func.ENDSWITH(n.func.LOWER(s), n.func.LOWER(substr)),
        ),
    ]


DEFINITIONS_STRING = [
    # ascii
    base.FuncAscii.for_dialect(D.SNOWFLAKE),
    # char
    base.FuncChar(
        variants=[
            V(D.SNOWFLAKE, sa.func.CHR),
        ]
    ),
    # concat
    base.Concat1.for_dialect(D.SNOWFLAKE),
    base.ConcatMultiStrConst.for_dialect(D.SNOWFLAKE),
    base.ConcatMultiStr.for_dialect(D.SNOWFLAKE),
    base.ConcatMultiAny.for_dialect(D.SNOWFLAKE),
    # contains
    base.FuncContainsNonConst(
        variants=[
            V(D.SNOWFLAKE, lambda x, y: sa.func.CONTAINS(x, y)),
        ]
    ),
    base.FuncContainsNonString.for_dialect(D.SNOWFLAKE),
    # notcontains
    base.FuncNotContainsNonConst.for_dialect(D.SNOWFLAKE),
    base.FuncNotContainsNonString.for_dialect(D.SNOWFLAKE),
    # endswith
    base.FuncEndswithNonString.for_dialect(D.SNOWFLAKE),
    base.FuncEndswithNonConst(variants=[V(D.SNOWFLAKE, sa.func.ENDSWITH)]),
    # find
    base.FuncFind2(variants=[V(D.SNOWFLAKE, lambda string, substring: sa.func.POSITION(substring, string))]),
    base.FuncFind3(
        variants=[V(D.SNOWFLAKE, lambda string, substring, start: sa.func.POSITION(substring, string, start))]
    ),
    # icontains
    base.FuncIContainsNonConst.for_dialect(D.SNOWFLAKE),
    base.FuncIContainsNonString.for_dialect(D.SNOWFLAKE),
    # iendswith
    base.FuncIEndswithNonConst.for_dialect(D.SNOWFLAKE),
    base.FuncIEndswithNonString.for_dialect(D.SNOWFLAKE),
    # istartswith
    base.FuncIStartswithNonConst.for_dialect(D.SNOWFLAKE),
    base.FuncIStartswithNonString.for_dialect(D.SNOWFLAKE),
    # left
    base.FuncLeft.for_dialect(D.SNOWFLAKE),
    # len
    base.FuncLenString(variants=[V(D.SNOWFLAKE, sa.func.LENGTH)]),
    # lower
    base.FuncLowerConst.for_dialect(D.SNOWFLAKE),
    base.FuncLowerNonConst.for_dialect(D.SNOWFLAKE),
    # ltrim
    base.FuncLtrim(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.LTRIM(value, " ")),
        ]
    ),
    # regexp_extract
    base.FuncRegexpExtract(variants=[V(D.SNOWFLAKE, sa.func.REGEXP_SUBSTR)]),
    # regexp_extract_nth
    base.FuncRegexpExtractNth(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda subj, pattern, nth: sa.func.REGEXP_SUBSTR(subj, pattern, 1, nth),
            ),
        ]
    ),
    # regexp_match
    base.FuncRegexpMatch(
        variants=[
            # todo: has a native method, but it should be a bin op, but SA seems to translate into func call
            #   maybe dig in and fix at some point
            V(D.SNOWFLAKE, lambda text, pattern: sa.func.REGEXP_SUBSTR(text, pattern).isnot(None)),
        ]
    ),
    # regexp_replace
    base.FuncRegexpReplace(variants=[V(D.SNOWFLAKE, sa.func.REGEXP_REPLACE)]),
    # replace
    base.FuncReplace.for_dialect(D.SNOWFLAKE),
    # right
    base.FuncRight.for_dialect(D.SNOWFLAKE),
    # rtrim
    base.FuncRtrim(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.RTRIM(value, " ")),
        ]
    ),
    # space
    base.FuncSpaceConst.for_dialect(D.SNOWFLAKE),
    base.FuncSpaceNonConst(variants=[V(D.SNOWFLAKE, lambda value: sa.func.REPEAT(" ", value))]),
    # split
    # base.FuncSplit1,  # FIXME: Add array support
    # base.FuncSplit2,  # FIXME: Add array support
    base.FuncSplit3(
        variants=[
            V(D.SNOWFLAKE, sa.func.SPLIT_PART),
        ]
    ),
    # startswith
    base.FuncStartswithNonConst(variants=[V(D.SNOWFLAKE, sa.func.STARTSWITH)]),
    base.FuncStartswithNonString.for_dialect(D.SNOWFLAKE),
    # substr
    base.FuncSubstr2.for_dialect(D.SNOWFLAKE),
    base.FuncSubstr3.for_dialect(D.SNOWFLAKE),
    # trim
    base.FuncTrim(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TRIM(value, " ")),
        ]
    ),
    # upper
    base.FuncUpperConst.for_dialect(D.SNOWFLAKE),
    base.FuncUpperNonConst.for_dialect(D.SNOWFLAKE),
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/functions_string.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    Function,
    SingleVariantTranslationBase,
    TranslationVariant,
)
import dl_formula.definitions.functions_type as base
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import Fixed

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


V = TranslationVariant.make


# todo: review datetime * timezone * snowflake logic
#   copied from postgresql connector


class FuncDatetime2SF(SingleVariantTranslationBase, base.FuncTypeGenericDatetime2Impl):
    dialects = D.SNOWFLAKE
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.DATETIME, DataType.GENERICDATETIME, DataType.INTEGER, DataType.FLOAT, DataType.STRING},
                DataType.CONST_STRING,
            ]
        ),
    ]
    name = "datetime"

    @classmethod
    def _translate_main(cls, value_ctx, tz_ctx):  # type: ignore
        """
        Equivalent to `dt at time zone tz`.
        Its semantics:

        1. For `timestamp` (naive) values, the value is interpreted
          at the specified timezone (with ambiguity),
          returning `timestamptz` (utc). Example:
          `select timezone('Europe/Moscow', '2010-10-31T02:00:00'::timestamp);`
          -> `2010-10-30 23:00:00+00`.
        2. For `timestamptz` (aware) values, the value is shifted to local time,
          returning `timestamp` (naive).
          Example:
          `select timezone('Europe/Moscow', timezone('Europe/Moscow', '2010-10-31T02:00:00'::timestamp));`
          -> `2010-10-31 02:00:00` (back to an ambiguous value).
        3. String / integer / float values are not accepted (however, some pg
          clients might autocast strings).
        4. `date` argument works, but it doesn't seem useful. Example:
          `select timezone('Europe/Moscow', '2010-10-31'::date);`
          -> `2010-10-31 03:00:00`

        This implementation provides a way to make datetime from string / integer / float values but it fully relies
        on database when it comes to transforming between naive and aware values.

        If input value is naive timestamp or string (will be casted to naive as well),
        the result will be tz-aware value (see 1. above)
        but in case of input being tz-aware timestamp or integer / float (casted to tz-aware utc value)
        the result will be naive (see 2. above)
        """
        value = value_ctx.expression
        value_type = value_ctx.data_type
        tz = tz_ctx.expression

        if value_type in (
            DataType.DATETIME,
            DataType.CONST_DATETIME,
            DataType.GENERICDATETIME,
            DataType.CONST_GENERICDATETIME,
        ):
            pass  # value = value
        elif value_type in (DataType.INTEGER, DataType.CONST_INTEGER, DataType.FLOAT, DataType.CONST_FLOAT):
            value = sa.func.to_timestamp(value)
        elif value_type in (DataType.STRING, DataType.CONST_STRING):
            value = sa.cast(value, sa.DateTime())
        else:
            raise Exception("Unexpected value type, likely a translation definition error", value_type)

        return sa.func.convert_timezone(tz, value)


class FuncGenericDatetime2SF(FuncDatetime2SF):
    name = "genericdatetime"
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncBoolFromNullSF(base.FuncBoolFromNull):
    variants = [
        V(D.SNOWFLAKE, lambda _: sa.null()),
    ]
    return_type = Fixed(DataType.NULL)


class FuncFloatFromNullSF(base.FuncFloat):
    variants = [
        V(D.SNOWFLAKE, lambda _: sa.null()),
    ]
    argument_types = [
        ArgTypeSequence([DataType.NULL]),
    ]


DEFINITIONS_TYPE = [
    # bool
    FuncBoolFromNullSF(),
    base.FuncBoolFromNumber.for_dialect(D.SNOWFLAKE),
    base.FuncBoolFromBool.for_dialect(D.SNOWFLAKE),
    base.FuncBoolFromStrGeo.for_dialect(D.SNOWFLAKE),
    base.FuncBoolFromDateDatetime.for_dialect(D.SNOWFLAKE),
    # date
    base.FuncDate1FromNull.for_dialect(D.SNOWFLAKE),
    base.FuncDate1FromDatetime(
        variants=[
            V(D.SNOWFLAKE, lambda expr: sa.func.TO_DATE(expr)),
        ]
    ),
    base.FuncDate1FromString(
        variants=[
            V(D.SNOWFLAKE, lambda expr: sa.func.TO_DATE(expr, "YYYY-MM-DD")),
        ]
    ),
    base.FuncDate1FromNumber(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda expr: sa.func.TO_DATE(sa.func.TO_CHAR(expr)),
            ),
        ]
    ),
    # datetime
    base.FuncDatetime1FromNull.for_dialect(D.SNOWFLAKE),
    base.FuncDatetime1FromDatetime.for_dialect(D.SNOWFLAKE),
    base.FuncDatetime1FromDate.for_dialect(D.SNOWFLAKE),
    base.FuncDatetime1FromNumber(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda expr: sa.func.TO_TIMESTAMP(expr),
            )
        ]
    ),
    base.FuncDatetime1FromString(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda value: sa.func.to_timestamp(value),
            )
        ]
    ),
    FuncDatetime2SF.for_dialect(D.SNOWFLAKE),
    # datetimetz
    base.FuncDatetimeTZConst.for_dialect(D.SNOWFLAKE),
    # float
    FuncFloatFromNullSF(),
    base.FuncFloatNumber(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TO_DOUBLE(value)),
        ]
    ),
    base.FuncFloatString(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TO_DOUBLE(value)),
        ]
    ),
    base.FuncFloatFromBool(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda value: sa.case(
                    whens=[
                        (value.is_(None), sa.null()),
                        (value == True, 1.0),  # noqa: E712
                        (value == False, 0.0),  # noqa: E712
                    ],
                ),
            )
        ]
    ),
    base.FuncFloatFromDate(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TO_DOUBLE(sa.func.DATE_PART("epoch_second", value))),
        ]
    ),
    base.FuncFloatFromDatetime(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TO_DOUBLE(sa.func.DATE_PART("epoch_second", value))),
        ]
    ),
    base.FuncFloatFromGenericDatetime(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TO_DOUBLE(sa.func.DATE_PART("epoch_second", value))),
        ]
    ),
    # genericdatetime
    base.FuncGenericDatetime1FromNull.for_dialect(D.SNOWFLAKE),
    base.FuncGenericDatetime1FromDatetime.for_dialect(D.SNOWFLAKE),
    base.FuncGenericDatetime1FromDate.for_dialect(D.SNOWFLAKE),
    base.FuncGenericDatetime1FromNumber(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda expr: sa.func.TO_TIMESTAMP(expr),
            )
        ]
    ),
    base.FuncGenericDatetime1FromString.for_dialect(D.SNOWFLAKE),
    FuncGenericDatetime2SF.for_dialect(D.SNOWFLAKE),
    # geopoint
    base.FuncGeopointFromStr.for_dialect(D.SNOWFLAKE),  # todo:
    base.FuncGeopointFromCoords.for_dialect(D.SNOWFLAKE),  # todo:
    # geopolygon
    base.FuncGeopolygon.for_dialect(D.SNOWFLAKE),  # todo:
    # int
    base.FuncIntFromNull(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TO_NUMBER(value)),
        ]
    ),
    base.FuncIntFromInt(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TO_NUMBER(value)),
        ]
    ),
    base.FuncIntFromFloat(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TO_NUMBER(sa.func.FLOOR(value))),
        ]
    ),
    base.FuncIntFromBool(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TO_NUMBER(value)),
        ]
    ),
    base.FuncIntFromStr(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.TO_NUMBER(value)),
        ]
    ),
    base.FuncIntFromDate(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.DATE_PART("epoch_second", value)),
        ]
    ),
    base.FuncIntFromDatetime(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.DATE_PART("epoch_second", value)),
        ]
    ),
    base.FuncIntFromGenericDatetime(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.DATE_PART("epoch_second", value)),
        ]
    ),
    base.FuncIntFromDatetimeTZ(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.func.DATE_PART("epoch_second", value)),
        ]
    ),
    # str
    base.FuncStrFromNull(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.cast(sa.null(), sa.String())),
        ]
    ),
    base.FuncStrFromUnsupported(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    base.FuncStrFromInteger(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    base.FuncStrFromFloat(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    base.FuncStrFromBool(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.case(whens=[(value.is_(None), sa.null()), (value, "True")], else_="False")),
        ]
    ),
    base.FuncStrFromStrGeo.for_dialect(D.SNOWFLAKE),
    base.FuncStrFromDate(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda value: sa.func.TO_CHAR(value),
            ),
        ]
    ),
    base.FuncStrFromDatetime(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda value: sa.func.LEFT(sa.func.TO_CHAR(sa.func.to_timestamp_ltz(value)), 19),
            ),
        ]
    ),
    base.FuncStrFromDatetimeTZ(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda value: sa.func.LEFT(sa.func.TO_CHAR(sa.func.to_timestamp_ltz(value)), 19),
            ),
        ]
    ),
    base.FuncStrFromString.for_dialect(D.SNOWFLAKE),
    base.FuncStrFromUUID(
        variants=[
            V(D.SNOWFLAKE, lambda value: sa.cast(value, sa.String())),
        ]
    ),
    # base.FuncStrFromArray,  # FIXME
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/functions_type.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.definitions.base import TranslationVariant
from dl_formula.definitions.common_datetime import DAY_SEC
import dl_formula.definitions.operators_binary as base

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


V = TranslationVariant.make


DEFINITIONS_BINARY = [
    # !=
    base.BinaryNotEqual.for_dialect(D.SNOWFLAKE),
    # %
    base.BinaryModInteger(
        variants=[
            V(D.SNOWFLAKE, sa.func.MOD),
        ]
    ),
    base.BinaryModFloat(
        variants=[
            V(D.SNOWFLAKE, lambda left, right: left - sa.func.FLOOR(left / right) * right),
        ]
    ),
    # *
    base.BinaryMultNumbers.for_dialect(D.SNOWFLAKE),
    base.BinaryMultStringConst.for_dialect(D.SNOWFLAKE),
    base.BinaryMultStringNonConst(
        variants=[
            V(D.SNOWFLAKE, sa.func.REPEAT),
        ]
    ),
    # +
    base.BinaryPlusNumbers.for_dialect(D.SNOWFLAKE),
    base.BinaryPlusStrings.for_dialect(D.SNOWFLAKE),
    base.BinaryPlusDateInt(
        variants=[
            V(D.SNOWFLAKE, lambda date, days: date + days),
        ]
    ),
    base.BinaryPlusDateFloat(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda date, days: sa.func.TO_DATE(sa.func.DATEADD("second", DAY_SEC * days, date)),
            )
        ]
    ),
    base.BinaryPlusDatetimeNumber(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda date, days: sa.func.DATEADD("second", DAY_SEC * days, date),
            ),
        ]
    ),
    base.BinaryPlusGenericDatetimeNumber(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda date, days: sa.func.DATEADD("second", DAY_SEC * days, date),
            ),
        ]
    ),
    # -
    base.BinaryMinusNumbers.for_dialect(D.SNOWFLAKE),
    base.BinaryMinusDateInt(
        variants=[
            V(D.SNOWFLAKE, lambda date, days: date - days),
        ]
    ),
    base.BinaryMinusDateFloat(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda date, days: sa.func.TO_DATE(sa.func.DATEADD("second", -DAY_SEC * days, date)),
            ),
        ]
    ),
    base.BinaryMinusDatetimeNumber(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda date, days: sa.func.DATEADD("second", -DAY_SEC * days, date),
            ),
        ]
    ),
    base.BinaryMinusGenericDatetimeNumber(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda date, days: sa.func.DATEADD("second", -DAY_SEC * days, date),
            ),
        ]
    ),
    base.BinaryMinusDates(
        variants=[
            V(D.SNOWFLAKE, lambda left, right: sa.func.DATEDIFF("second", right, left) / DAY_SEC),
        ]
    ),
    base.BinaryMinusDatetimes(
        variants=[
            V(D.SNOWFLAKE, lambda left, right: sa.func.DATEDIFF("second", right, left) / DAY_SEC),
        ]
    ),
    base.BinaryMinusGenericDatetimes(
        variants=[
            V(
                D.SNOWFLAKE,
                lambda left, right: sa.func.DATEDIFF("second", right, left) / DAY_SEC,
            ),
        ]
    ),
    # /
    base.BinaryDivInt.for_dialect(D.SNOWFLAKE),
    base.BinaryDivFloat.for_dialect(D.SNOWFLAKE),
    # <
    base.BinaryLessThan.for_dialect(D.SNOWFLAKE),
    # <=
    base.BinaryLessThanOrEqual.for_dialect(D.SNOWFLAKE),
    # ==
    base.BinaryEqual.for_dialect(D.SNOWFLAKE),
    # >
    base.BinaryGreaterThan.for_dialect(D.SNOWFLAKE),
    # >=
    base.BinaryGreaterThanOrEqual.for_dialect(D.SNOWFLAKE),
    # ^
    base.BinaryPower.for_dialect(D.SNOWFLAKE),
    # _!=
    base.BinaryNotEqualInternal.for_dialect(D.SNOWFLAKE),
    # _==
    base.BinaryEqualInternal.for_dialect(D.SNOWFLAKE),
    # _dneq
    base.BinaryEqualDenullified.for_dialect(D.SNOWFLAKE),
    # and
    base.BinaryAnd.for_dialect(D.SNOWFLAKE),
    # in
    base.BinaryIn.for_dialect(D.SNOWFLAKE),
    # like
    base.BinaryLike.for_dialect(D.SNOWFLAKE),
    # notin
    base.BinaryNotIn.for_dialect(D.SNOWFLAKE),
    # notlike
    base.BinaryNotLike.for_dialect(D.SNOWFLAKE),
    # or
    base.BinaryOr.for_dialect(D.SNOWFLAKE),
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/operators_binary.py
datalens-tech/datalens-backend
import dl_formula.definitions.operators_ternary as base

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


DEFINITIONS_TERNARY = [
    # between
    base.TernaryBetween.for_dialect(D.SNOWFLAKE),
    # notbetween
    base.TernaryNotBetween.for_dialect(D.SNOWFLAKE),
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/operators_ternary.py
datalens-tech/datalens-backend
from dl_formula.definitions.base import TranslationVariant
import dl_formula.definitions.operators_unary as base

from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D


V = TranslationVariant.make


DEFINITIONS_UNARY = [
    # isfalse
    base.UnaryIsFalseStringGeo.for_dialect(D.SNOWFLAKE),
    base.UnaryIsFalseNumbers.for_dialect(D.SNOWFLAKE),
    base.UnaryIsFalseDateTime.for_dialect(D.SNOWFLAKE),
    base.UnaryIsFalseBoolean(
        variants=[
            V(D.SNOWFLAKE, lambda x: x == False),  # noqa: E712
        ]
    ),
    # istrue
    base.UnaryIsTrueStringGeo.for_dialect(D.SNOWFLAKE),
    base.UnaryIsTrueNumbers.for_dialect(D.SNOWFLAKE),
    base.UnaryIsTrueDateTime.for_dialect(D.SNOWFLAKE),
    base.UnaryIsTrueBoolean(
        variants=[
            V(D.SNOWFLAKE, lambda x: x == True),  # noqa: E712
        ]
    ),
    # neg
    base.UnaryNegate.for_dialect(D.SNOWFLAKE),
    # not
    base.UnaryNotBool.for_dialect(D.SNOWFLAKE),
    base.UnaryNotNumbers.for_dialect(D.SNOWFLAKE),
    base.UnaryNotStringGeo.for_dialect(D.SNOWFLAKE),
    base.UnaryNotDateDatetime.for_dialect(D.SNOWFLAKE),
]

lib/dl_connector_snowflake/dl_connector_snowflake/formula/definitions/operators_unary.py
datalens-tech/datalens-backend
import datetime

import sqlalchemy as sa

from dl_formula.connectors.base.literal import (
    Literal,
    Literalizer,
)
from dl_formula.core.dialect import DialectCombo


class SnowFlakeLiteralizer(Literalizer):
    __slots__ = ()

    def literal_datetime(self, value: datetime.datetime, dialect: DialectCombo) -> Literal:
        return sa.func.TO_TIMESTAMP(value)

    def literal_date(self, value: datetime.date, dialect: DialectCombo) -> Literal:
        return sa.func.TO_DATE(value)

lib/dl_connector_snowflake/dl_connector_snowflake/formula/literal.py
datalens-tech/datalens-backend
import datetime

import pytest

from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_api_lib_testing.connection_base import ConnectionTestBase
from dl_api_lib_testing.data_api_base import StandardizedDataApiTestBase
from dl_api_lib_testing.dataset_base import DatasetTestBase

from dl_connector_snowflake.core.constants import (
    CONNECTION_TYPE_SNOWFLAKE,
    SOURCE_TYPE_SNOWFLAKE_TABLE,
)
from dl_connector_snowflake_tests.ext.config import API_TEST_CONFIG
from dl_connector_snowflake_tests.ext.core.base import BaseSnowFlakeTestClass


class SnowFlakeConnectionTestBase(BaseSnowFlakeTestClass, ConnectionTestBase):
    conn_type = CONNECTION_TYPE_SNOWFLAKE

    @pytest.fixture(scope="class")
    def bi_test_config(self) -> ApiTestEnvironmentConfiguration:
        return API_TEST_CONFIG

    @pytest.fixture(scope="class")
    def connection_params_native(self, sf_secrets) -> dict:
        return dict(
            account_name=sf_secrets.get_account_name(),
            user_name=sf_secrets.get_user_name(),
            user_role=sf_secrets.get_user_role(),
            client_id=sf_secrets.get_client_id(),
            client_secret=sf_secrets.get_client_secret(),
            db_name=sf_secrets.get_database(),
            schema=sf_secrets.get_schema(),
            warehouse=sf_secrets.get_warehouse(),
            refresh_token=sf_secrets.get_refresh_token_x(),
            refresh_token_expire_time=datetime.datetime.now() + datetime.timedelta(days=80),
        )

    @pytest.fixture(scope="class")
    def connection_params(self, connection_params_native) -> dict:
        to_serialize = {k: v for k, v in connection_params_native.items()}
        # 2014-12-22T03:12:58
        if connection_params_native["refresh_token_expire_time"]:
            dt = connection_params_native["refresh_token_expire_time"]
            to_serialize["refresh_token_expire_time"] = dt.strftime("%Y-%m-%dT%H:%M:%S")

        return to_serialize


class SnowFlakeDatasetTestBase(SnowFlakeConnectionTestBase, DatasetTestBase):
    @pytest.fixture(scope="class")
    def dataset_params(self, sf_secrets) -> dict:
        return dict(
            source_type=SOURCE_TYPE_SNOWFLAKE_TABLE.name,
            parameters=dict(
                table_name=sf_secrets.get_table_name(),
            ),
        )


class SnowFlakeDataApiTestBase(SnowFlakeDatasetTestBase, StandardizedDataApiTestBase):
    pass

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/api/base.py
datalens-tech/datalens-backend
import json
import uuid

from dl_api_lib_testing.connector.connection_suite import DefaultConnectorConnectionTestSuite
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_snowflake_tests.ext.api.base import SnowFlakeConnectionTestBase


class TestSnowFlakeConnection(SnowFlakeConnectionTestBase, DefaultConnectorConnectionTestSuite):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            DefaultConnectorConnectionTestSuite.test_cache_ttl_sec_override: "Unavailable for SnowFlake",
        }
    )

    def test_connection_parameters_tester(self, client, connection_params):
        params = {
            "type": "snowflake",
            "name": uuid.uuid4().hex,
        }
        params.update(connection_params)
        r = client.post(
            "/api/v1/connections/test_connection_params",
            data=json.dumps(params),
            content_type="application/json",
        )
        assert r.status_code == 200

    def test_create_connection_get_tester(self, client, connection_params):
        params = {
            "type": "snowflake",
            "name": uuid.uuid4().hex,
        }
        params.update(connection_params)
        resp = client.post(
            "/api/v1/connections",
            data=json.dumps(params),
            content_type="application/json",
        )
        assert resp.status_code == 200
        conn_id = resp.json["id"]

        resp = client.get(f"/api/v1/connections/{conn_id}")
        assert resp.status_code == 200

        data = resp.json
        for k, v in connection_params.items():
            if k in ("client_secret", "refresh_token"):
                continue

            assert v == data[k]

        resp = client.post(
            f"/api/v1/connections/test_connection/{conn_id}",
            data=json.dumps({}),
            content_type="application/json",
        )
        assert resp.status_code == 200

    def test_connection_malicious_account_name(self, client, connection_params):
        bad_names = ["foo;", "https://fake", "your-site.com/oauth/token-request?"]
        params = {
            "type": "snowflake",
            "name": uuid.uuid4().hex,
        }
        params.update(connection_params)
        for bad_name in bad_names:
            params["account_name"] = bad_name
            resp = client.post(
                "/api/v1/connections/test_connection_params",
                data=json.dumps(params),
                content_type="application/json",
            )
            assert resp.status_code == 400

            resp = client.post(
                "/api/v1/connections",
                data=json.dumps(params),
                content_type="application/json",
            )
            assert resp.status_code == 400
            assert "account_name" in resp.json["message"]
            assert "id" not in resp.json

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/api/test_connection.py
datalens-tech/datalens-backend
from dl_api_lib_testing.connector.data_api_suites import (
    DefaultConnectorDataDistinctTestSuite,
    DefaultConnectorDataGroupByFormulaTestSuite,
    DefaultConnectorDataPreviewTestSuite,
    DefaultConnectorDataRangeTestSuite,
    DefaultConnectorDataResultTestSuite,
)
from dl_testing.regulated_test import RegulatedTestParams

from dl_connector_snowflake_tests.ext.api.base import SnowFlakeDataApiTestBase


class TestSnowFlakeDataResult(SnowFlakeDataApiTestBase, DefaultConnectorDataResultTestSuite):
    test_params = RegulatedTestParams(
        mark_features_skipped={
            DefaultConnectorDataResultTestSuite.array_support: "SnowFlake doesn't support arrays",
        }
    )


class TestSnowFlakeDataRange(SnowFlakeDataApiTestBase, DefaultConnectorDataRangeTestSuite):
    pass


class TestSnowFlakeDataDistinct(SnowFlakeDataApiTestBase, DefaultConnectorDataDistinctTestSuite):
    pass


class TestSnowFlakeDataGroupBy(SnowFlakeDataApiTestBase, DefaultConnectorDataGroupByFormulaTestSuite):
    pass


class TestSnowFlakeDataPreview(SnowFlakeDataApiTestBase, DefaultConnectorDataPreviewTestSuite):
    pass

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/api/test_data.py
datalens-tech/datalens-backend
from dl_api_client.dsmaker.primitives import Dataset
from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite

from dl_connector_snowflake_tests.ext.api.base import SnowFlakeDatasetTestBase


class TestSnowFlakeDataset(SnowFlakeDatasetTestBase, DefaultConnectorDatasetTestSuite):
    def check_basic_dataset(self, ds: Dataset) -> None:
        field_names = {field.title for field in ds.result_schema}
        assert {"Category", "City"}.issubset(field_names)

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/api/test_dataset.py
datalens-tech/datalens-backend
from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration
from dl_constants.enums import UserDataType
from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport


# Infra settings
CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=51700).host,
    port_us_http=get_test_container_hostport("us", fallback_port=51700).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=51709).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=51709).port,
    us_master_token="AC1ofiek8coB",
    core_connector_ep_names=["snowflake"],
)

# Connector settings
DB_DSN = "snowflake://does@not/matter"
# todo: maybe support also certificates and store them in sec storage, to avoid refresh token updates

SAMPLE_TABLE_SIMPLIFIED_SCHEMA = [
    ("Category", UserDataType.string),
    ("City", UserDataType.string),
    ("Country", UserDataType.string),
    ("Customer ID", UserDataType.string),
    ("Customer Name", UserDataType.string),
    ("Discount", UserDataType.float),
    ("Order Date", UserDataType.date),
    ("Order ID", UserDataType.string),
    ("Postal Code", UserDataType.integer),
    ("Product ID", UserDataType.string),
    ("Product Name", UserDataType.string),
    ("Profit", UserDataType.float),
    ("Quantity", UserDataType.integer),
    ("Region", UserDataType.string),
    ("Row ID", UserDataType.integer),
    ("Sales", UserDataType.float),
    ("Segment", UserDataType.string),
    ("Ship Date", UserDataType.date),
    ("Ship Mode", UserDataType.string),
    ("State", UserDataType.string),
    ("Sub-Category", UserDataType.string),
]


API_TEST_CONFIG = ApiTestEnvironmentConfiguration(
    api_connector_ep_names=["snowflake"],
    core_test_config=CORE_TEST_CONFIG,
    ext_query_executer_secret_key="_some_test_secret_key_",
)

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/config.py
datalens-tech/datalens-backend
import os

import pytest

from dl_api_lib.loader import load_api_lib
from dl_api_lib_testing.initialization import initialize_api_lib_test
from dl_formula_testing.forced_literal import forced_literal_use
from dl_testing.env_params.generic import GenericEnvParamGetter

from dl_connector_snowflake.core.testing.secrets import SnowFlakeSecretReader
from dl_connector_snowflake_tests.ext.config import API_TEST_CONFIG


pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'


def pytest_configure(config):  # noqa
    initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG)


@pytest.fixture(scope="session")
def env_param_getter() -> GenericEnvParamGetter:
    filepath = os.path.join(os.path.dirname(__file__), "params.yml")
    return GenericEnvParamGetter.from_yaml_file(filepath)


@pytest.fixture(scope="session")
def sf_secrets(env_param_getter):
    return SnowFlakeSecretReader(env_param_getter)


__all__ = (
    # auto-use fixtures:
    "forced_literal_use",
)

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/conftest.py
datalens-tech/datalens-backend
import datetime
from typing import Any

import pytest

from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import CoreDbConfig
from dl_core_testing.testcases.connection import BaseConnectionTestClass

from dl_connector_snowflake.core.constants import CONNECTION_TYPE_SNOWFLAKE
from dl_connector_snowflake.core.testing.connection import make_snowflake_saved_connection
from dl_connector_snowflake.core.us_connection import ConnectionSQLSnowFlake
from dl_connector_snowflake.db_testing.engine_wrapper import SnowFlakeDbEngineConfig
import dl_connector_snowflake_tests.ext.config as test_config  # noqa


class BaseSnowFlakeTestClass(BaseConnectionTestClass[ConnectionSQLSnowFlake]):
    conn_type = CONNECTION_TYPE_SNOWFLAKE
    core_test_config = test_config.CORE_TEST_CONFIG
    engine_config_cls = SnowFlakeDbEngineConfig

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return test_config.DB_DSN

    @pytest.fixture(scope="class")
    def db(self, db_config: CoreDbConfig) -> Any:
        return None

    @pytest.fixture(scope="function")
    def connection_creation_params(self, sf_secrets) -> dict:
        return dict(
            account_name=sf_secrets.get_account_name(),
            user_name=sf_secrets.get_user_name(),
            user_role=sf_secrets.get_user_role(),
            client_id=sf_secrets.get_client_id(),
            client_secret=sf_secrets.get_client_secret(),
            db_name=sf_secrets.get_database(),
            schema=sf_secrets.get_schema(),
            warehouse=sf_secrets.get_warehouse(),
            refresh_token=sf_secrets.get_refresh_token_x(),
            refresh_token_expire_time=None,
            **(dict(raw_sql_level=self.raw_sql_level) if self.raw_sql_level is not None else {}),
        )

    @pytest.fixture(scope="function")
    def saved_connection(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params: dict,
    ) -> ConnectionSQLSnowFlake:
        conn = make_snowflake_saved_connection(sync_usm=sync_us_manager, **connection_creation_params)
        return conn

    @pytest.fixture(scope="function")
    def connection_creation_params_with_expired_refresh_token(self, sf_secrets) -> dict:
        # note: in fact any bad string could be used for this test as well ...
        return dict(
            account_name=sf_secrets.get_account_name(),
            user_name=sf_secrets.get_user_name(),
            client_id=sf_secrets.get_client_id(),
            client_secret=sf_secrets.get_client_secret(),
            db_name=sf_secrets.get_database(),
            schema=sf_secrets.get_schema(),
            warehouse=sf_secrets.get_warehouse(),
            refresh_token=sf_secrets.get_refresh_token_expired(),
            refresh_token_expire_time=None,
        )

    @pytest.fixture(scope="function")
    def saved_connection_with_expired_refresh_token(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params_with_expired_refresh_token: dict,
    ) -> ConnectionSQLSnowFlake:
        conn = make_snowflake_saved_connection(
            sync_usm=sync_us_manager, **connection_creation_params_with_expired_refresh_token
        )
        return conn

    @pytest.fixture(scope="function")
    def connection_creation_params_with_refresh_token_soon_to_expire(self, sf_secrets) -> dict:
        # note: in fact any bad string could be used for this test as well ...
        return dict(
            account_name=sf_secrets.get_account_name(),
            user_name=sf_secrets.get_user_name(),
            client_id=sf_secrets.get_client_id(),
            client_secret=sf_secrets.get_client_secret(),
            db_name=sf_secrets.get_database(),
            schema=sf_secrets.get_schema(),
            warehouse=sf_secrets.get_warehouse(),
            refresh_token=sf_secrets.get_refresh_token_x(),
            refresh_token_expire_time=datetime.datetime.now() + datetime.timedelta(days=2),
        )

    @pytest.fixture(scope="function")
    def saved_connection_with_refresh_token_soon_to_expire(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params_with_refresh_token_soon_to_expire: dict,
    ) -> ConnectionSQLSnowFlake:
        conn = make_snowflake_saved_connection(
            sync_usm=sync_us_manager, **connection_creation_params_with_refresh_token_soon_to_expire
        )
        return conn

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/core/base.py
datalens-tech/datalens-backend
import asyncio

from mock import MagicMock
import pytest

from dl_core.connection_executors.common_base import ConnExecutorQuery
from dl_core.services_registry.top_level import ServicesRegistry
from dl_core.us_connection_base import DataSourceTemplate
from dl_core_testing.testcases.connection import DefaultConnectionTestClass

from dl_connector_snowflake.auth import SFAuthProvider
from dl_connector_snowflake.core.exc import (
    SnowflakeAccessTokenError,
    SnowflakeGetAccessTokenError,
)
from dl_connector_snowflake.core.us_connection import ConnectionSQLSnowFlake
from dl_connector_snowflake_tests.ext.core.base import BaseSnowFlakeTestClass  # noqa


class TestSnowFlakeConnection(
    BaseSnowFlakeTestClass,
    DefaultConnectionTestClass[ConnectionSQLSnowFlake],
):
    def check_saved_connection(self, conn: ConnectionSQLSnowFlake, params: dict) -> None:
        assert conn.uuid is not None

    def check_data_source_templates(
        self,
        conn: ConnectionSQLSnowFlake,
        dsrc_templates: list[DataSourceTemplate],
    ) -> None:
        pass  # TODO

    def test_connection_with_correct_refresh_token(
        self,
        saved_connection: ConnectionSQLSnowFlake,
        conn_sync_service_registry: ServicesRegistry,
    ):
        ce_factory = conn_sync_service_registry.get_conn_executor_factory()
        conn_executor = ce_factory.get_sync_conn_executor(saved_connection)
        conn_executor.execute(ConnExecutorQuery(query="select 1"))

    def test_connection_test_with_expired_refresh_token(
        self,
        saved_connection_with_expired_refresh_token: ConnectionSQLSnowFlake,
        conn_sync_service_registry: ServicesRegistry,
    ):
        ce_factory = conn_sync_service_registry.get_conn_executor_factory()
        conn_executor = ce_factory.get_sync_conn_executor(saved_connection_with_expired_refresh_token)
        # note: currently could not get proper error message with aiohttp client
        # with pytest.raises(SnowflakeRefreshTokenInvalid):
        with pytest.raises(SnowflakeGetAccessTokenError):
            conn_executor.execute(ConnExecutorQuery(query="select 1"))

    def test_bad_access_token(
        self,
        saved_connection,
        conn_sync_service_registry: ServicesRegistry,
        monkeypatch,
    ):
        ce_factory = conn_sync_service_registry.get_conn_executor_factory()
        conn_executor = ce_factory.get_sync_conn_executor(saved_connection)

        future = asyncio.Future()
        future.set_result("bad_token")
        mck = MagicMock()
        mck.return_value = future

        monkeypatch.setattr(SFAuthProvider, "async_get_access_token", mck)
        with pytest.raises(SnowflakeAccessTokenError):
            conn_executor.execute(ConnExecutorQuery(query="select 1"))

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/core/test_connection.py
datalens-tech/datalens-backend
import pytest

from dl_constants.enums import (
    RawSQLLevel,
    UserDataType,
)
from dl_core_testing.testcases.data_source import DefaultDataSourceTestClass

from dl_connector_snowflake.core.constants import (
    SOURCE_TYPE_SNOWFLAKE_SUBSELECT,
    SOURCE_TYPE_SNOWFLAKE_TABLE,
)
from dl_connector_snowflake.core.data_source import (
    SnowFlakeSubselectDataSource,
    SnowFlakeTableDataSource,
)
from dl_connector_snowflake.core.data_source_spec import (
    SnowFlakeSubselectDataSourceSpec,
    SnowFlakeTableDataSourceSpec,
)
from dl_connector_snowflake.core.us_connection import ConnectionSQLSnowFlake
from dl_connector_snowflake_tests.ext.config import SAMPLE_TABLE_SIMPLIFIED_SCHEMA
from dl_connector_snowflake_tests.ext.core.base import BaseSnowFlakeTestClass


class TestSnowFlakeTableDataSource(
    BaseSnowFlakeTestClass,
    DefaultDataSourceTestClass[
        ConnectionSQLSnowFlake,
        SnowFlakeTableDataSourceSpec,
        SnowFlakeTableDataSource,
    ],
):
    DSRC_CLS = SnowFlakeTableDataSource

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self, sf_secrets) -> SnowFlakeTableDataSourceSpec:
        dsrc_spec = SnowFlakeTableDataSourceSpec(
            source_type=SOURCE_TYPE_SNOWFLAKE_TABLE,
            table_name=sf_secrets.get_table_name(),
            db_name=sf_secrets.get_database(),
            schema_name=sf_secrets.get_schema(),
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return SAMPLE_TABLE_SIMPLIFIED_SCHEMA


class TestSnowFlakeSubselectDataSoure(
    BaseSnowFlakeTestClass,
    DefaultDataSourceTestClass[
        ConnectionSQLSnowFlake,
        SnowFlakeSubselectDataSourceSpec,
        SnowFlakeSubselectDataSource,
    ],
):
    DSRC_CLS = SnowFlakeSubselectDataSource

    raw_sql_level = RawSQLLevel.subselect

    @pytest.fixture(scope="class")
    def initial_data_source_spec(self, sf_secrets) -> SnowFlakeSubselectDataSourceSpec:
        dsrc_spec = SnowFlakeSubselectDataSourceSpec(
            source_type=SOURCE_TYPE_SNOWFLAKE_SUBSELECT,
            subsql=f"SELECT * FROM {sf_secrets.get_database()}.{sf_secrets.get_schema()}.{sf_secrets.get_table_name()}",
        )
        return dsrc_spec

    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        return SAMPLE_TABLE_SIMPLIFIED_SCHEMA

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/core/test_data_source.py
datalens-tech/datalens-backend
import pytest as pytest

from dl_connector_snowflake_tests.ext.core.base import BaseSnowFlakeTestClass


class TestSnowflakeConnectionNotification(BaseSnowFlakeTestClass):
    @pytest.mark.asyncio
    async def test_refresh_token_notification_soon_to_expire(self, saved_connection_with_refresh_token_soon_to_expire):
        notifications = saved_connection_with_refresh_token_soon_to_expire.check_for_notifications()
        record = notifications[0]
        assert record.locator == "snowflake_refresh_token_soon_to_expire"

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/core/test_token_expire_notification.py
datalens-tech/datalens-backend
import asyncio
from typing import Callable

import pytest

from dl_formula_testing.testcases.base import FormulaConnectorTestBase

from dl_connector_snowflake.auth import SFAuthProvider
from dl_connector_snowflake.core.adapters import construct_creator_func
from dl_connector_snowflake.core.dto import SnowFlakeConnDTO
from dl_connector_snowflake.core.target_dto import SnowFlakeConnTargetDTO
from dl_connector_snowflake.db_testing.engine_wrapper import SnowFlakeDbEngineConfig
from dl_connector_snowflake.formula.constants import SnowFlakeDialect as D
import dl_connector_snowflake_tests.ext.config as test_config  # noqa


class SnowFlakeTestBase(FormulaConnectorTestBase):
    dialect = D.SNOWFLAKE
    eval_attempts = 2
    engine_config_cls = SnowFlakeDbEngineConfig

    supports_arrays = False
    supports_uuid = False
    null_casts_to_false = False

    # skip_custom_tz = True  # todo: probably could support with adjustments

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return test_config.DB_DSN

    @pytest.fixture(scope="class")
    def loop(self):
        return asyncio.get_event_loop()

    @pytest.fixture(scope="class")
    def _sf_target_dto(self, loop, sf_secrets):
        conn_dto = SnowFlakeConnDTO(
            conn_id=None,
            account_name=sf_secrets.get_account_name(),
            user_name=sf_secrets.get_user_name(),
            user_role=sf_secrets.get_user_role(),
            client_id=sf_secrets.get_client_id(),
            client_secret=sf_secrets.get_client_secret(),
            db_name=sf_secrets.get_database(),
            schema=sf_secrets.get_schema(),
            warehouse=sf_secrets.get_warehouse(),
            refresh_token=sf_secrets.get_refresh_token_x(),
            refresh_token_expire_time=None,
        )
        sf_auth_provider = SFAuthProvider.from_dto(conn_dto)
        access_token = loop.run_until_complete(sf_auth_provider.async_get_access_token())

        return SnowFlakeConnTargetDTO.from_dto(dto=conn_dto, access_token=access_token)

    @pytest.fixture(scope="class")
    def _sf_creator_func(self, _sf_target_dto):
        return construct_creator_func(_sf_target_dto)

    @pytest.fixture(scope="class")
    def engine_params(self, _sf_creator_func: Callable) -> dict:
        return dict(creator=_sf_creator_func)

    @pytest.fixture(scope="class")
    def engine_config(self, db_url: str, engine_params: dict, _sf_target_dto) -> SnowFlakeDbEngineConfig:
        return self.engine_config_cls(url=db_url, engine_params=engine_params, db_name=_sf_target_dto.db_name)

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/base.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.conditional_blocks import DefaultConditionalBlockFormulaConnectorTestSuite

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


class TestConditionalBlockSnowFlake(SnowFlakeTestBase, DefaultConditionalBlockFormulaConnectorTestSuite):
    pass

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_conditional_blocks.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_aggregation import DefaultMainAggFunctionFormulaConnectorTestSuite

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


class TestMainAggFunctionSnowFlake(SnowFlakeTestBase, DefaultMainAggFunctionFormulaConnectorTestSuite):
    supports_countd_approx = True
    supports_quantile = True
    supports_median = True

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_functions_aggregation.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_datetime import DefaultDateTimeFunctionFormulaConnectorTestSuite

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


class TestDateTimeFunctionSnowFlake(SnowFlakeTestBase, DefaultDateTimeFunctionFormulaConnectorTestSuite):
    supports_addition_to_feb_29 = True
    supports_deprecated_dateadd = True

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_functions_datetime.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_logical import DefaultLogicalFunctionFormulaConnectorTestSuite

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


class TestLogicalFunctionSnowFlake(SnowFlakeTestBase, DefaultLogicalFunctionFormulaConnectorTestSuite):
    supports_nan_funcs = False
    supports_iif = False

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_functions_logical.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_markup import DefaultMarkupFunctionFormulaConnectorTestSuite

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


class TestMarkupFunctionSnowFlake(SnowFlakeTestBase, DefaultMarkupFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_functions_markup.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_math import DefaultMathFunctionFormulaConnectorTestSuite

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


class TestMathFunctionSnowFlake(SnowFlakeTestBase, DefaultMathFunctionFormulaConnectorTestSuite):
    supports_float_div = True

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_functions_math.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_string import DefaultStringFunctionFormulaConnectorTestSuite

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


class TestStringFunctionSnowFlake(SnowFlakeTestBase, DefaultStringFunctionFormulaConnectorTestSuite):
    pass

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_functions_string.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.functions_type_conversion import (
    DefaultBoolTypeFunctionFormulaConnectorTestSuite,
    DefaultDateTypeFunctionFormulaConnectorTestSuite,
    DefaultFloatTypeFunctionFormulaConnectorTestSuite,
    DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite,
    DefaultIntTypeFunctionFormulaConnectorTestSuite,
    DefaultStrTypeFunctionFormulaConnectorTestSuite,
)

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


# STR


class TestStrTypeFunctionSnowFlake(SnowFlakeTestBase, DefaultStrTypeFunctionFormulaConnectorTestSuite):
    pass


# FLOAT


class TestFloatTypeFunctionSnowFlake(SnowFlakeTestBase, DefaultFloatTypeFunctionFormulaConnectorTestSuite):
    pass


# BOOL


class TestBoolTypeFunctionSnowFlake(SnowFlakeTestBase, DefaultBoolTypeFunctionFormulaConnectorTestSuite):
    pass


# INT


class TestIntTypeFunctionSnowFlake(SnowFlakeTestBase, DefaultIntTypeFunctionFormulaConnectorTestSuite):
    pass


# DATE


class TestDateTypeFunctionSnowFlake(SnowFlakeTestBase, DefaultDateTypeFunctionFormulaConnectorTestSuite):
    pass


# GENERICDATETIME (& DATETIME)


class TestGenericDatetimeTypeFunctionSnowFlake(
    SnowFlakeTestBase,
    DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite,
):
    pass

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_functions_type_conversion.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.literals import DefaultLiteralFormulaConnectorTestSuite

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


class TestLiteralSnowFlake(SnowFlakeTestBase, DefaultLiteralFormulaConnectorTestSuite):
    supports_microseconds = True
    supports_utc = False
    supports_custom_tz = False

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_literals.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.misc_funcs import DefaultMiscFunctionalityConnectorTestSuite

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


class TestMiscFunctionalitySnowFlake(SnowFlakeTestBase, DefaultMiscFunctionalityConnectorTestSuite):
    pass

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_misc_funcs.py
datalens-tech/datalens-backend
from dl_formula_testing.testcases.operators import DefaultOperatorFormulaConnectorTestSuite

from dl_connector_snowflake_tests.ext.formula.base import SnowFlakeTestBase  # noqa


class TestOperatorSnowFlake(SnowFlakeTestBase, DefaultOperatorFormulaConnectorTestSuite):
    pass

lib/dl_connector_snowflake/dl_connector_snowflake_tests/ext/formula/test_operators.py
datalens-tech/datalens-backend
from dl_api_connector.i18n.localizer import CONFIGS as BI_API_CONNECTOR_CONFIGS
from dl_api_lib_testing.connection_form_base import ConnectionFormTestBase

from dl_connector_snowflake.api.connection_form.form_config import SnowFlakeConnectionFormFactory
from dl_connector_snowflake.api.i18n.localizer import CONFIGS as BI_CONNECTOR_SNOWFLAKE_CONFIGS


class TestSnowFlakeConnectionForm(ConnectionFormTestBase):
    CONN_FORM_FACTORY_CLS = SnowFlakeConnectionFormFactory
    TRANSLATION_CONFIGS = BI_CONNECTOR_SNOWFLAKE_CONFIGS + BI_API_CONNECTOR_CONFIGS

lib/dl_connector_snowflake/dl_connector_snowflake_tests/unit/test_connection_form.py
datalens-tech/datalens-backend
from __future__ import annotations

import enum


class DLHeaders(str, enum.Enum):
    """
    Base marker class for all headers enumeration.
    Should be empty because enums with members can not be extended by other members.
    """


class DLCookies(enum.Enum):
    """"""


class DLContextKey(str, enum.Enum):
    DASH_ID = "dashId"
    CHART_ID = "chartId"
    CHART_KIND = "chartKind"
    DASH_TAB_ID = "dashTabId"


class DLHeadersCommon(DLHeaders):
    AUTHORIZATION_TOKEN = "Authorization"
    REQUEST_ID = "X-Request-ID"
    PUBLIC_API_KEY = "X-DL-Back-Public-API-Token"
    FILE_UPLOADER_MASTER_TOKEN = "X-DL-File-Uploader-Master-Token"
    TENANT_ID = "X-DL-TenantId"  # Tenant in the UnitedStorage format.
    PROJECT_ID = "X-DC-ProjectId"
    FORWARDED_FOR = "X-Forwarded-For"
    CSRF_TOKEN = "X-CSRF-Token"
    CONTENT_TYPE = "Content-Type"
    ACCEPT_LANGUAGE = "Accept-Language"
    ORIGIN = "Origin"

    ACCESS_CTRL_REQ_HEADERS = "Access-Control-Request-Headers"
    ACCESS_CTRL_REQ_METH = "Access-Control-Request-Method"

    # US access-control
    SUDO = "X-DL-Sudo"
    ALLOW_SUPERUSER = "X-DL-Allow-Superuser"

    API_KEY = "X-DL-API-Key"

    DL_COMPONENT = "X-DL-Component"
    DL_CONTEXT = "X-DL-Context"
    US_MASTER_TOKEN = "X-US-Master-Token"
    US_PUBLIC_TOKEN = "X-US-Public-API-Token"

    EMBED_TOKEN = "X-DL-Embed-Token"

    # TODO: BI-4918 drop after all usages moved to bi_api_commons_ya_cloud.constants.DLHeadersYC
    IAM_TOKEN = "X-YaCloud-SubjectToken"
    FOLDER_ID = "X-YaCloud-FolderId"
    ORG_ID = "X-YaCloud-OrgId"


class DLCookiesCommon(DLCookies):
    ...

lib/dl_constants/dl_constants/api_constants.py
datalens-tech/datalens-backend
"""
Common enums used in API and core internals of BI projects
"""

from __future__ import annotations

from enum import (
    Enum,
    IntEnum,
    auto,
    unique,
)
from typing import (
    Generic,
    Optional,
    Type,
    TypeVar,
    Union,
    cast,
)

from dynamic_enum import (
    AutoEnumValue,
    DynamicEnum,
)


_ENUM_TYPE = TypeVar("_ENUM_TYPE", bound=Enum)


class _Normalizable(Generic[_ENUM_TYPE]):
    @classmethod
    def normalize(cls, value: Union[_ENUM_TYPE, str, None]) -> Optional[_ENUM_TYPE]:
        if isinstance(value, str):
            value = cast(Type[_ENUM_TYPE], cls)[value]
        return value


@unique
class UserDataType(_Normalizable["UserDataType"], Enum):
    string = auto()
    integer = auto()
    float = auto()
    date = auto()
    datetime = auto()
    boolean = auto()
    geopoint = auto()
    geopolygon = auto()
    uuid = auto()
    markup = auto()
    datetimetz = auto()
    unsupported = auto()
    array_str = auto()
    array_int = auto()
    array_float = auto()
    tree_str = auto()
    genericdatetime = auto()


class SourceBackendType(DynamicEnum):
    NONE = AutoEnumValue()


class ConnectionType(DynamicEnum):
    unknown = AutoEnumValue()


@unique
class QueryBlockPlacementType(Enum):
    root = "root"
    after = "after"
    dispersed_after = "dispersed_after"


@unique
class FieldRole(Enum):
    row = "row"
    measure = "measure"  # TODO: Deprecate
    info = "info"
    order_by = "order_by"
    filter = "filter"
    parameter = "parameter"
    distinct = "distinct"
    range = "range"
    total = "total"
    template = "template"
    tree = "tree"


@unique
class FieldVisibility(Enum):
    visible = "visible"
    hidden = "hidden"


@unique
class PivotRole(Enum):
    pivot_row = "pivot_row"
    pivot_column = "pivot_column"
    pivot_measure = "pivot_measure"
    pivot_annotation = "pivot_annotation"
    pivot_info = "pivot_info"


@unique
class PivotHeaderRole(Enum):
    data = "data"
    total = "total"


@unique
class PivotItemType(Enum):
    stream_item = "stream_item"
    measure_name = "measure_name"
    dimension_name = "dimension_name"


@unique
class LegendItemType(Enum):
    field = "field"
    measure_name = "measure_name"
    dimension_name = "dimension_name"
    placeholder = "placeholder"


class QueryItemRefType(Enum):
    id = "id"
    title = "title"
    measure_name = "measure_name"
    dimension_name = "dimension_name"
    placeholder = "placeholder"


@unique
class ParameterValueConstraintType(Enum):
    all = "all"
    range = "range"
    set = "set"


@unique
class CalcMode(_Normalizable["CalcMode"], Enum):
    direct = "direct"
    formula = "formula"
    parameter = "parameter"


@unique
class ConditionPartCalcMode(Enum):
    direct = "direct"
    formula = "formula"
    result_field = "result_field"


@unique
class AggregationFunction(_Normalizable["AggregationFunction"], Enum):
    none = "none"
    sum = "sum"
    avg = "avg"
    min = "min"
    max = "max"
    count = "count"
    countunique = "countunique"
    # TODO: 'any'


@unique
class BinaryJoinOperator(Enum):
    gt = "gt"
    lt = "lt"
    gte = "gte"
    lte = "lte"
    eq = "eq"
    ne = "ne"


@unique
class FieldType(_Normalizable["FieldType"], Enum):
    DIMENSION = "DIMENSION"
    MEASURE = "MEASURE"


@unique
class WhereClauseOperation(Enum):
    # unary
    ISNULL = "isnull"
    ISNOTNULL = "isnotnull"

    # binary
    GT = "gt"
    LT = "lt"
    GTE = "gte"
    LTE = "lte"
    EQ = "eq"
    NE = "ne"
    STARTSWITH = "startswith"
    ISTARTSWITH = "istartswith"
    ENDSWITH = "endswith"
    IENDSWITH = "iendswith"
    CONTAINS = "contains"
    ICONTAINS = "icontains"
    NOTCONTAINS = "notcontains"
    NOTICONTAINS = "noticontains"
    LENEQ = "leneq"
    LENNE = "lenne"
    LENGT = "lengt"
    LENGTE = "lengte"
    LENLT = "lenlt"
    LENLTE = "lenlte"

    # binary with list
    IN = "in"
    NIN = "nin"

    # ternary
    BETWEEN = "between"


class DataSourceType(DynamicEnum):
    @classmethod
    def normalize(cls, value: DataSourceType | str | None) -> Optional[DataSourceType]:
        # FIXME: Remove this hack (used only in dsmaker)
        if isinstance(value, str):
            value = DataSourceType(value)
        return value


@unique
class ConnectionState(IntEnum):
    new = 0
    parsing_error = 1
    unvalidated = 2
    saved = 3
    error = 4


class DataSourceRole(Enum):
    origin = "origin"
    sample = "sample"
    materialization = "materialization"


class DataSourceCollectionType(Enum):
    collection = "collection"


class DataSourceCreatedVia(DynamicEnum):
    user = AutoEnumValue()


class RLSSubjectType(Enum):
    user = "user"  # `{value}: {login}` meaning `value allowed for user {login}`
    all = "all"  # `{value}: *` meaning 'value allowed for all subjects'
    userid = "userid"  # 'userid: userid' meaning 'filter by `field = {userid}`'
    unknown = "unknown"  # placeholder, should not be normally used.
    notfound = "notfound"  # subject that couldn't be found in subject resolver


class RLSPatternType(Enum):
    value = "user"
    all = "all"
    userid = "userid"


class ReportingQueryType(Enum):
    internal = "internal"
    external = "external"


class JoinConditionType(Enum):
    binary = "binary"


@unique
class JoinType(Enum):
    inner = "inner"
    left = "left"
    right = "right"
    full = "full"


@unique
class ManagedBy(_Normalizable["ManagedBy"], Enum):
    user = "user"
    feature = "feature"
    compiler_runtime = "compiler_runtime"


@unique
class USAuthMode(Enum):
    master = "master"
    public = "public"
    regular = "regular"


@unique
class ComponentType(Enum):
    data_source = "data_source"
    source_avatar = "source_avatar"
    avatar_relation = "avatar_relation"
    field = "field"
    obligatory_filter = "obligatory_filter"
    result_schema = "result_schema"


@unique
class TopLevelComponentId(Enum):
    result_schema = "__result_schema__"


@unique
class ComponentErrorLevel(Enum):
    error = "error"
    warning = "warning"


@unique
class OrderDirection(Enum):
    asc = "asc"
    desc = "desc"


@unique
class RangeType(Enum):
    min = "min"
    max = "max"


class SelectorType(Enum):
    CACHED = auto()
    CACHED_LAZY = auto()


class ProcessorType(Enum):
    SOURCE_DB = auto()
    ASYNCPG = auto()
    AIOPG = auto()


@unique
class RawSQLLevel(Enum):
    off = "off"  # no raw sql allowed
    subselect = "subselect"  # wrapped raw SQL with `edit` permissions
    dashsql = "dashsql"  # unwrapped raw SQL with `execute` permissions


@unique
class NotificationStatus(Enum):
    OK = "OK"
    WARN = "WARN"
    CRIT = "CRIT"


class IndexKind(Enum):
    table_sorting = auto()


class RedisInstanceKind(Enum):
    caches = auto()
    persistent = auto()
    mutations = auto()
    arq = auto()


class FileProcessingStatus(Enum):
    in_progress = "in_progress"
    ready = "ready"
    failed = "failed"
    expired = "expired"


@unique
class NotificationLevel(Enum):
    info = "info"
    warning = "warning"
    critical = "critical"


class NotificationType(DynamicEnum):
    totals_removed_due_to_measure_filter = AutoEnumValue()


class ConnectorAvailability(Enum):
    free = "free"
    whitelist = "whitelist"


class QueryProcessingMode(DynamicEnum):
    basic = AutoEnumValue()
    no_compeng = AutoEnumValue()
    native_wf = AutoEnumValue()

lib/dl_constants/dl_constants/enums.py
datalens-tech/datalens-backend
from typing import (
    Any,
    ClassVar,
    Optional,
    Type,
)

from typing_extensions import final

from dl_constants.types import TJSONLike


GLOBAL_ERR_PREFIX = "ERR"
DEFAULT_ERR_CODE_API_PREFIX = "DS_API"
CODE_OK = "OK"


# TODO FIX: Simplify exc data structures to simplify serialization
#  Make only 2 arguments in constructor: data & orig
class DLBaseException(Exception):
    # code parts joining by dots on api response preparing.
    err_code: ClassVar[list[str]] = []  # TODO: Implement automatic hierarchial code inheritance
    _message: str

    default_message = "Internal Server Error"
    formatting_messages: Optional[dict[frozenset[str], str]] = None

    # Auxiliary error info. Will not be shown to user (or will be shown as is only in intranet)
    debug_info: dict[str, Any]

    # Dictionary with some fixed structure per each err_code.
    # Can be used by front to show user some additional info.
    details: dict[str, Any]

    params: dict[str, Any]

    def __init__(
        self,
        message: Optional[str] = None,
        details: Optional[dict[str, Any]] = None,
        orig: Optional[Exception] = None,
        debug_info: Optional[dict[str, Any]] = None,
        params: Optional[dict[str, Any]] = None,
    ):
        self._message = message or self.default_message
        self.details = details or {}
        self.debug_info = debug_info or {}
        self.orig = orig or None
        self.params = params or {}

    @property
    def message(self) -> str:
        """
        Message will be shown directly to user.
        """
        if (
            self.params
            and self.formatting_messages
            and frozenset(self.params.keys()) in self.formatting_messages.keys()
        ):
            return self.formatting_messages[frozenset(self.params.keys())].format(**self.params)
        else:
            return self._message

    def __str__(self) -> str:
        return self.message or super().__str__()

    # ##
    # Serialization logic for passing from QE
    # ##
    _MAP_CLASS_NAME_CLASS: ClassVar[dict[str, Type["DLBaseException"]]] = {}

    def __init_subclass__(cls, **kwargs):  # type: ignore
        cls._MAP_CLASS_NAME_CLASS[cls.__qualname__] = cls

    @classmethod
    @final
    def from_jsonable_dict(cls, data: dict) -> "DLBaseException":
        data = {**data}
        cls_name = data.pop("cls_name")
        target_cls = cls._MAP_CLASS_NAME_CLASS[cls_name]
        # noinspection PyProtectedMember
        return target_cls._from_jsonable_dict(data)

    @classmethod
    def _from_jsonable_dict(cls, data: dict) -> "DLBaseException":
        new_exc = cls(message=data.pop("message", None))
        new_exc.details = data.pop("details")
        new_exc.debug_info = data.pop("debug_info")
        new_exc.params = data.pop("params", {})
        return new_exc

    def to_jsonable_dict(self) -> dict[str, TJSONLike]:
        return dict(
            cls_name=type(self).__qualname__,
            message=None if self._message == self.default_message else self._message,
            details=self.details,
            debug_info=self.debug_info,
            params=self.params,
        )

lib/dl_constants/dl_constants/exc.py
datalens-tech/datalens-backend
MEASURE_NAME_ID = "__@measure_name__"  # Pseudo-dimension ID
MEASURE_NAME_TITLE = "Measure Names"
DIMENSION_NAME_ID = "__@dimension_name__"  # Pseudo-dimension ID
DIMENSION_NAME_TITLE = "Dimension Names"
MEASURE_VALUE_ID = "__@measure_value__"  # Pseudo-dimension ID
PLACEHOLDER_ID = "__@placeholder__"  # Pseudo-dimension ID
PLACEHOLDER_TITLE = "Placeholder"

lib/dl_constants/dl_constants/internal_constants.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import decimal
from typing import (
    AsyncGenerator,
    AsyncIterable,
    Sequence,
    Union,
)
import uuid


TJSONScalar = Union[str, float, int, bool, None]
TBIDataValue = Union[
    # All the python types that connection executors are known to spit out.
    # (does not include e.g. `memoryview` which should always be processed in the executor).
    datetime.date,
    datetime.datetime,
    datetime.time,
    datetime.timedelta,
    decimal.Decimal,
    uuid.UUID,
    bytes,
    TJSONScalar,
]


# Allowing `tuple` as it is "jsonable"
# # Recursive (not supported by mypy):
# TJSONLike = Union[
#     TJSONScalar,
#     dict[str, 'TJSONLike'],
#     list['TJSONLike'],
#     tuple['TJSONLike', ...],
# ]
# # Limited recursion:
TJSONLikeUnchecked = Union[
    dict,
    list,
    tuple,
    TJSONScalar,
]
TJSONLike = Union[
    dict[str, TJSONLikeUnchecked],
    list[TJSONLikeUnchecked],
    tuple[TJSONLikeUnchecked, ...],
    TJSONScalar,
]


# Types supported by the extended JSON serializer (`bi_core.serialization`)
# # Recursive (not supported by mypy):
# TJSONExt = Union[
#     dict[str, 'TJSONExt'],
#     list['TJSONExt']
#     tuple['TJSONExt', ...],
#     TBIDataValue,
# ]
# # Limited recursion:
TJSONExtUnchecked = Union[
    dict,
    list,
    tuple,
    TBIDataValue,
]
TJSONExt = Union[
    dict[str, TJSONExtUnchecked],
    list[TJSONExtUnchecked],
    tuple[TJSONExtUnchecked, ...],
    TBIDataValue,
]


# More convenience shortcuts
TJSONLikeAIter = AsyncIterable[TJSONLike]
TJSONExtAIter = AsyncIterable[TJSONExt]
TBIDataRow = Sequence[TBIDataValue]  # sequence of cells -> row
TBIDataTable = Sequence[TBIDataRow]  # sequence of rows -> table
TBIChunksGen = AsyncGenerator[TBIDataTable, None]  # sequence of tables -> chunks
# NOTE: ^ often wrapped in an AsyncChunked

lib/dl_constants/dl_constants/types.py
datalens-tech/datalens-backend
from typing import TYPE_CHECKING

import redis.asyncio

from dl_api_commons.aiohttp.aiohttp_wrappers import (
    DLRequestBase,
    DLRequestView,
)
from dl_constants.enums import (
    ProcessorType,
    RedisInstanceKind,
)
from dl_core.aio.web_app_services.data_processing.data_processor import DataProcessorService
from dl_core.aio.web_app_services.data_processing.factory import get_data_processor_service_class
from dl_core.aio.web_app_services.redis import RedisBaseService
from dl_core.us_manager.us_manager_async import AsyncUSManager


if TYPE_CHECKING:
    from dl_core.services_registry import ServicesRegistry  # noqa


class DLRequestDataCore(DLRequestBase):
    KEY_US_MANAGER = "us_manager"
    KEY_SERVICE_US_MANAGER = "service_us_manager"
    KEY_SERVICES_REGISTRY = "services_registry"

    @property
    def us_manager(self) -> AsyncUSManager:
        """US manager associated with request's user"""
        return self.request.get(self.KEY_US_MANAGER)  # type: ignore  # TODO: fix

    @us_manager.setter
    def us_manager(self, value: AsyncUSManager) -> None:
        self._set_attr_once(self.KEY_US_MANAGER, value)

    @property
    def service_us_manager(self) -> AsyncUSManager:
        """US manager with master token"""
        return self.request.get(self.KEY_SERVICE_US_MANAGER)  # type: ignore  # TODO: fix

    @service_us_manager.setter
    def service_us_manager(self, value: AsyncUSManager) -> None:
        self._set_attr_once(self.KEY_SERVICE_US_MANAGER, value)

    @property
    def services_registry(self) -> "ServicesRegistry":
        return self.request.get(self.KEY_SERVICES_REGISTRY)  # type: ignore  # TODO: fix

    @services_registry.setter
    def services_registry(self, value: "ServicesRegistry") -> None:
        self._set_attr_once(self.KEY_SERVICES_REGISTRY, value)

    def get_caches_redis(self, allow_slave: bool = False) -> redis.asyncio.Redis:
        return RedisBaseService.get_app_instance(
            self.request.app,
            RedisInstanceKind.caches,
        ).get_redis(allow_slave=allow_slave)

    def get_persistent_redis(self, allow_slave: bool = False) -> redis.asyncio.Redis:
        return RedisBaseService.get_app_instance(
            self.request.app,
            RedisInstanceKind.persistent,
        ).get_redis(allow_slave=allow_slave)

    def get_mutations_redis(self, allow_slave: bool = False) -> redis.asyncio.Redis:
        return RedisBaseService.get_app_instance(
            self.request.app,
            RedisInstanceKind.mutations,
        ).get_redis(allow_slave=allow_slave)

    def get_data_processor_service(self, processor_type: ProcessorType) -> DataProcessorService:
        return get_data_processor_service_class(processor_type).get_app_instance(self.request.app)


class DLRequestDataCoreView(DLRequestView[DLRequestDataCore]):
    dl_request_cls = DLRequestDataCore

lib/dl_core/dl_core/aio/aiohttp_wrappers_data_core.py
datalens-tech/datalens-backend
from __future__ import annotations

import os
from typing import ClassVar

from aiohttp import web
from statcommons.unistat.common import dump_for_prometheus
from statcommons.unistat.uwsgi import uwsgi_prometheus

from dl_api_commons.aiohttp.aiohttp_wrappers import (
    DLRequestView,
    RequiredResource,
    RequiredResourceCommon,
)


class MetricsView(DLRequestView):
    REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset({RequiredResourceCommon.SKIP_AUTH})

    async def get(self) -> web.Response:
        result = []
        if os.environ.get("UWSGI_STATS"):
            uwsgi_metrics = uwsgi_prometheus(label_prefix="uwsgi_")
            result.extend(uwsgi_metrics)

        body = "".join(dump_for_prometheus(result))
        return web.Response(text=body)

lib/dl_core/dl_core/aio/metrics_view.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging

from aiohttp import web
from aiohttp.typedefs import Handler

from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_api_commons.aiohttp.aiohttp_wrappers import (
    DLRequestBase,
    RequiredResourceCommon,
)
from dl_constants.api_constants import (
    DLHeaders,
    DLHeadersCommon,
)


LOGGER = logging.getLogger(__name__)


def master_key_middleware(
    master_key: str,
    header: DLHeaders = DLHeadersCommon.API_KEY,
) -> AIOHTTPMiddleware:
    @web.middleware
    @DLRequestBase.use_dl_request
    async def actual_master_key_middleware(dl_request: DLRequestBase, handler: Handler) -> web.StreamResponse:
        handler_required_resources = dl_request.required_resources

        if RequiredResourceCommon.MASTER_KEY in handler_required_resources:
            LOGGER.info("Master key is required for this handler")
            request_api_key = dl_request.get_single_header(header)

            if request_api_key != master_key:
                LOGGER.warning("Attempt to access private resource without master key")
                raise web.HTTPForbidden()

            LOGGER.info("Master key is OK")

        return await handler(dl_request.request)

    return actual_master_key_middleware

lib/dl_core/dl_core/aio/middlewares/master_key.py
datalens-tech/datalens-backend
from __future__ import annotations

from aiohttp import web
from aiohttp.typedefs import Handler

from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_core.aio.aiohttp_wrappers_data_core import DLRequestDataCore
from dl_core.services_registry.sr_factories import SRFactory
from dl_core.us_manager.mutation_cache.usentry_mutation_cache_factory import DefaultUSEntryMutationCacheFactory


def services_registry_middleware(
    services_registry_factory: SRFactory,
    use_query_cache: bool = True,
    use_mutation_cache: bool = False,
    mutation_cache_default_ttl: float = 60,
) -> AIOHTTPMiddleware:
    """

    :param services_registry_factory:
    :param use_query_cache: If `True` - Redis client factory will be passed in service registry, Otherwise - `None`.
        Warning: if Redis service is not configured for aiohttp application and this flag is `True` -
        attempt to call `services_registry.get_caches_redis_client()` will cause an exception.
    :return: Configured middleware that creates service registry for each request.
    """

    @web.middleware
    @DLRequestDataCore.use_dl_request
    async def actual_services_registry_middleware(
        dl_request: DLRequestDataCore, handler: Handler
    ) -> web.StreamResponse:
        mutations_cache_factory, mutations_redis_client_factory = (
            (DefaultUSEntryMutationCacheFactory(default_ttl=mutation_cache_default_ttl), dl_request.get_mutations_redis)
            if use_mutation_cache
            else (None, None)
        )
        sr = services_registry_factory.make_service_registry(
            request_context_info=dl_request.rci,
            caches_redis_client_factory=dl_request.get_caches_redis if use_query_cache else None,
            mutations_cache_factory=mutations_cache_factory,
            mutations_redis_client_factory=mutations_redis_client_factory,
            data_processor_service_factory=dl_request.get_data_processor_service,
            reporting_registry=dl_request.reporting_registry,
        )
        try:
            dl_request.services_registry = sr
            return await handler(dl_request.request)
        finally:
            await sr.close_async()

    return actual_services_registry_middleware

lib/dl_core/dl_core/aio/middlewares/services_registry.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import contextlib
import logging
from typing import (
    AsyncGenerator,
    Optional,
)

from aiohttp import web
from aiohttp.typedefs import Handler
import attr

from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_api_commons.aiohttp import aiohttp_wrappers
from dl_api_commons.tenant_resolver import TenantResolver
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_core.aio.aiohttp_wrappers_data_core import DLRequestDataCore
from dl_core.services_registry.top_level import DummyServiceRegistry
from dl_utils.aio import shield_wait_for_complete

from ... import exc
from ...us_manager.factory import USMFactory
from ...us_manager.us_manager_async import AsyncUSManager


LOGGER = logging.getLogger(__name__)


def public_usm_workaround_middleware(
    us_base_url: str,
    us_public_token: str,
    crypto_keys_config: CryptoKeysConfig,
    dataset_id_match_info_code: str,
    conn_id_match_info_code: str,
    us_master_token: Optional[str],
    tenant_resolver: TenantResolver,
) -> AIOHTTPMiddleware:
    """
    This is workaround for the following public API issues:
     1. We have no explicit tenant ID in request
     2. Materialization connection is not public

    For issue 1: middleware fetches dataset or connection from US and pick tenant ID from response
    For issue 2: middleware fetches materialization config (with US master token) for given tenant
      and stores it in DLRequest

    This middleware works with temp (uncommitted) RCI.
    :param us_base_url:
    :param us_master_token:
    :param us_public_token:
    :param crypto_keys_config:
    :param dataset_id_match_info_code: name of URL match param for dataset ID
    :param conn_id_match_info_code: name of URL match param for connection ID
    :param tenant_resolver:
    :return: Actual middleware
    """
    usm_factory = USMFactory(
        us_base_url=us_base_url,
        crypto_keys_config=crypto_keys_config,
        us_master_token=us_master_token,
        us_public_token=us_public_token,
    )

    @web.middleware
    @DLRequestDataCore.use_dl_request
    async def actual_public_usm_workaround_middleware(
        dl_request: DLRequestDataCore, handler: Handler
    ) -> web.StreamResponse:
        if aiohttp_wrappers.RequiredResourceCommon.US_MANAGER not in dl_request.required_resources:
            return await handler(dl_request.request)

        rci = dl_request.temp_rci

        # Fetching dataset or connection & determining tenant id
        # Here we know that we need US manager for request
        entry_id = next(
            dl_request.request.match_info[match_info_code]
            for match_info_code in (dataset_id_match_info_code, conn_id_match_info_code)
            if match_info_code in dl_request.request.match_info
        )

        sr = dl_request.services_registry
        if sr is None:
            # It won't really be needed for anything.
            # We just need to initialize a USM and fetch an entry.
            # Might be used in pre/post-init hooks though.
            # But we expect nothing serious - it just won't work with the dummy SR otherwise
            sr = DummyServiceRegistry(rci=rci)
        assert sr is not None

        public_usm = usm_factory.get_public_async_usm(rci=rci, services_registry=sr)
        try:
            try:
                entry = await public_usm.get_by_id(entry_id)
            except exc.USObjectNotFoundException:
                raise web.HTTPNotFound()
            else:
                tenant = tenant_resolver.resolve_tenant_def_by_tenant_id(entry.raw_tenant_id)
        finally:
            await public_usm.close()

        rci = attr.evolve(rci, tenant=tenant)
        dl_request.replace_temp_rci(rci)

        return await handler(dl_request.request)

    return actual_public_usm_workaround_middleware


@contextlib.asynccontextmanager
async def _usm_close_cm(usm: AsyncUSManager, label: str) -> AsyncGenerator[None, None]:
    try:
        yield
    finally:
        try:
            await shield_wait_for_complete(usm.close())
        except asyncio.CancelledError:
            raise
        except Exception:  # noqa
            LOGGER.warning(f"Error during closing {label} USManager", exc_info=True)


def us_manager_middleware(
    us_base_url: str,
    crypto_keys_config: CryptoKeysConfig,
    embed: bool = False,
) -> AIOHTTPMiddleware:
    usm_factory = USMFactory(
        us_base_url=us_base_url,
        crypto_keys_config=crypto_keys_config,
    )

    @web.middleware
    @DLRequestDataCore.use_dl_request
    async def actual_us_manager_middleware(dl_request: DLRequestDataCore, handler: Handler) -> web.StreamResponse:
        if aiohttp_wrappers.RequiredResourceCommon.US_MANAGER not in dl_request.required_resources:
            return await handler(dl_request.request)

        if embed:
            usm = usm_factory.get_embed_async_usm(
                rci=dl_request.rci,
                services_registry=dl_request.services_registry,
            )
        else:
            usm = usm_factory.get_regular_async_usm(
                rci=dl_request.rci,
                services_registry=dl_request.services_registry,
            )

        async with _usm_close_cm(usm, label="user"):
            dl_request.us_manager = usm
            return await handler(dl_request.request)

    return actual_us_manager_middleware


def public_us_manager_middleware(
    us_base_url: str,
    us_public_token: str,
    crypto_keys_config: CryptoKeysConfig,
) -> AIOHTTPMiddleware:
    """
    Middleware to create public US manager. Works with committed RCI.
    Must be used after `public_usm_workaround_middleware`.
    Can not be used together with `us_manager_middleware` because public USM treated as user USM.
    """
    usm_factory = USMFactory(
        us_base_url=us_base_url,
        crypto_keys_config=crypto_keys_config,
        us_public_token=us_public_token,
    )

    @web.middleware
    @DLRequestDataCore.use_dl_request
    async def actual_public_us_manager_middleware(
        dl_request: DLRequestDataCore, handler: Handler
    ) -> web.StreamResponse:
        if aiohttp_wrappers.RequiredResourceCommon.US_MANAGER not in dl_request.required_resources:
            return await handler(dl_request.request)

        usm = usm_factory.get_public_async_usm(
            rci=dl_request.rci,
            services_registry=dl_request.services_registry,
        )
        async with _usm_close_cm(usm, label="public"):
            dl_request.us_manager = usm
            return await handler(dl_request.request)

    return actual_public_us_manager_middleware


def service_us_manager_middleware(
    us_base_url: str,
    us_master_token: str,
    crypto_keys_config: CryptoKeysConfig,
    as_user_usm: bool = False,
) -> AIOHTTPMiddleware:
    usm_factory = USMFactory(
        us_base_url=us_base_url,
        crypto_keys_config=crypto_keys_config,
        us_master_token=us_master_token,
    )

    @web.middleware
    @DLRequestDataCore.use_dl_request
    async def actual_service_us_manager_middleware(
        dl_request: DLRequestDataCore, handler: Handler
    ) -> web.StreamResponse:
        target_resource = (
            aiohttp_wrappers.RequiredResourceCommon.US_MANAGER
            if as_user_usm
            else aiohttp_wrappers.RequiredResourceCommon.SERVICE_US_MANAGER
        )

        if target_resource not in dl_request.required_resources:
            return await handler(dl_request.request)

        usm = usm_factory.get_master_async_usm(
            rci=dl_request.rci,
            services_registry=dl_request.services_registry,
        )

        async with _usm_close_cm(usm, label="service"):
            if not as_user_usm:
                dl_request.service_us_manager = usm
            else:
                # For tests only!!!
                LOGGER.info("Setting service USM as user USM")
                dl_request.us_manager = usm

            return await handler(dl_request.request)

    return actual_service_us_manager_middleware

lib/dl_core/dl_core/aio/middlewares/us_manager.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import ClassVar

from aiohttp import web

from dl_api_commons.aiohttp.aiohttp_wrappers import (
    DLRequestView,
    RequiredResource,
    RequiredResourceCommon,
)


class PingView(DLRequestView):
    REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset({RequiredResourceCommon.SKIP_AUTH})

    async def get(self) -> web.StreamResponse:
        return web.json_response(
            {
                "request_id": self.dl_request.rci.request_id,
            }
        )

lib/dl_core/dl_core/aio/ping_view.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    ClassVar,
    Type,
    TypeVar,
    cast,
)

from aiohttp import web
import attr

from dl_core.data_processing.processing.processor import OperationProcessorAsyncBase


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class DataProcessorConfig:
    pass


_DATA_PROC_SRV_TV = TypeVar("_DATA_PROC_SRV_TV", bound="DataProcessorService")


@attr.s
class DataProcessorService(metaclass=abc.ABCMeta):
    APP_KEY: ClassVar[str]

    @classmethod
    def get_full_app_key(cls) -> str:
        return cls.APP_KEY

    async def init_hook(self, target_app: web.Application) -> None:
        assert self.APP_KEY not in target_app, f"{type(self).__name__} already initiated for app"
        LOGGER.info(f"Initializing {type(self).__name__} data processor service")
        target_app[self.APP_KEY] = self
        await self.initialize()

    async def tear_down_hook(self, target_app: web.Application) -> None:
        LOGGER.info(f"Tear down {type(self).__name__} data processor service...")
        await self.tear_down(target_app)
        LOGGER.info(f"Tear down {type(self).__name__} data processor service: done.")

    async def tear_down(self, target_app: web.Application) -> None:
        await self.finalize()

    @abc.abstractmethod
    async def initialize(self) -> None:
        raise NotImplementedError

    @abc.abstractmethod
    async def finalize(self) -> None:
        raise NotImplementedError

    @classmethod
    def get_app_instance(cls: Type[_DATA_PROC_SRV_TV], app: web.Application) -> _DATA_PROC_SRV_TV:
        service = cast(_DATA_PROC_SRV_TV, app.get(cls.get_full_app_key(), None))
        if service is None:
            raise ValueError(f"{cls.__name__} was not initiated for application")

        return service

    @abc.abstractmethod
    def get_data_processor(self) -> OperationProcessorAsyncBase:
        raise NotImplementedError

    @classmethod
    def from_config(cls: Type[_DATA_PROC_SRV_TV], config: DataProcessorConfig) -> _DATA_PROC_SRV_TV:
        return cls()

lib/dl_core/dl_core/aio/web_app_services/data_processing/data_processor.py
datalens-tech/datalens-backend
from typing import Type

from dl_constants.enums import ProcessorType
from dl_core.aio.web_app_services.data_processing.data_processor import (
    DataProcessorConfig,
    DataProcessorService,
)


def make_compeng_service(processor_type: ProcessorType, config: DataProcessorConfig) -> DataProcessorService:
    data_proc_cls = get_data_processor_service_class(processor_type=processor_type)
    return data_proc_cls.from_config(config)


_DATA_PROCESSOR_SERVICE_CLASSES: dict[ProcessorType, Type[DataProcessorService]] = {}


def register_data_processor_service_class(
    processor_type: ProcessorType,
    data_processor_srv_cls: Type[DataProcessorService],
) -> None:
    try:
        assert _DATA_PROCESSOR_SERVICE_CLASSES[processor_type] is data_processor_srv_cls
    except KeyError:
        _DATA_PROCESSOR_SERVICE_CLASSES[processor_type] = data_processor_srv_cls


def get_data_processor_service_class(processor_type: ProcessorType) -> Type[DataProcessorService]:
    try:
        return _DATA_PROCESSOR_SERVICE_CLASSES[processor_type]
    except KeyError:
        # TODO: Support non-pg data processor here too
        raise ValueError(f"Data processor type {processor_type} is not supported")

lib/dl_core/dl_core/aio/web_app_services/data_processing/factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import enum
import logging
from typing import (
    ClassVar,
    Optional,
    Union,
)

from aiohttp import web
import attr


LOGGER = logging.getLogger(__name__)


class NumberFormatType(str, enum.Enum):
    """
    According to Google API Docs + several custom types
    https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets/cells#CellFormat
    """

    NUMBER_FORMAT_TYPE_UNSPECIFIED = "NUMBER_FORMAT_TYPE_UNSPECIFIED"
    TEXT = "TEXT"

    NUMBER = "NUMBER"
    PERCENT = "PERCENT"
    CURRENCY = "CURRENCY"
    SCIENTIFIC = "SCIENTIFIC"

    DATE = "DATE"
    TIME = "TIME"
    DATE_TIME = "DATE_TIME"

    # not google types
    INTEGER = "INTEGER"
    FLOAT = "FLOAT"
    BOOLEAN = "BOOLEAN"


@attr.s(auto_attribs=True)
class Cell:
    value: Optional[Union[str, int, float, bool]]
    number_format_type: NumberFormatType = NumberFormatType.NUMBER_FORMAT_TYPE_UNSPECIFIED
    empty: bool = False  # means this is an actually unfilled cell, not just clear


@attr.s(auto_attribs=True)
class Range:
    sheet_title: str

    col_from: int
    row_from: int
    col_to: int
    row_to: int

    def __str__(self) -> str:
        return "{sheet_title}!R{row_from}C{col_from}:R{row_to}C{col_to}".format(**self.__dict__)


@attr.s(auto_attribs=True)
class Sheet:
    """
    id: number from /edit#gid=<sheetId>
    index: order of the sheet e.g. 0, 1, 2, 3, ...
    title: title of the sheet
    row_count: row count from sheet properties returned by sheets api
    column_count: column count from sheet properties returned by sheet api and reduced to cover data only in sample
    batch_size_rows: how many rows to request in `spreadsheets.get.values` in every batch to make responses around 4 MB
    data: actual grid data with values and formats
    """

    id: int
    index: int
    title: str
    row_count: int
    column_count: int
    batch_size_rows: int = 50
    data: Optional[list[list[Cell]]] = None

    def col_is_time(self, idx: int, has_header: bool) -> bool:
        """
        Check if column with given idx contains only TIME values (or TIME and nulls)
        """

        if not self.data or has_header and len(self.data) < 2:
            return False
        data_iter = iter(self.data)
        if has_header:
            next(data_iter)
        time_rows = 0
        for row in data_iter:
            if row[idx].number_format_type != NumberFormatType.TIME and not row[idx].empty:
                return False
            if row[idx].number_format_type == NumberFormatType.TIME:
                time_rows += 1
        return time_rows > 0


@attr.s(auto_attribs=True)
class Spreadsheet:
    id: str
    url: str  # TODO unused
    title: str
    sheets: list[Sheet]


@attr.s
class GSheetsSettings:
    APP_KEY: ClassVar[str] = "GSHEETS_OAUTH2"

    api_key: str = attr.ib(repr=False)
    client_id: str = attr.ib(repr=False)
    client_secret: str = attr.ib(repr=False)

    def __attrs_post_init__(self):  # type: ignore  # monkeypatch point
        pass

    def bind_to_app(self, app: web.Application) -> None:
        app[self.APP_KEY] = self

    @classmethod
    def get_for_app(cls, app: web.Application) -> GSheetsSettings:
        return app[cls.APP_KEY]

lib/dl_core/dl_core/aio/web_app_services/gsheets.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    ClassVar,
    Optional,
    Sequence,
)

from aiohttp import web
import attr
import redis.asyncio
from redis.asyncio.sentinel import Sentinel as RedisSentinel
from redis.asyncio.sentinel import SentinelConnectionPool

from dl_constants.enums import RedisInstanceKind
from dl_core.utils import secrepr


LOGGER = logging.getLogger(__name__)


@attr.s(kw_only=True)
class RedisConnParams:
    host: str = attr.ib()
    port: int = attr.ib()
    db: int = attr.ib()
    password: str = attr.ib(repr=secrepr)
    ssl: Optional[bool] = attr.ib()


async def extract_redis_conn_params(redis_cli: Optional[redis.asyncio.Redis]) -> Optional[RedisConnParams]:
    if redis_cli is None:
        return None
    redis_conn_params = redis_cli.connection_pool.connection_kwargs
    host = redis_conn_params.get("host")
    port = redis_conn_params.get("port")

    if "connection_pool" in redis_conn_params:
        conn_pool: SentinelConnectionPool = redis_conn_params["connection_pool"]
        master_address = await conn_pool.get_master_address()
        host, port = master_address

    assert host is not None
    assert port is not None

    return RedisConnParams(
        host=host,
        port=port,
        db=redis_conn_params["db"],
        password=redis_conn_params["password"],
        ssl=redis_conn_params.get("ssl"),
    )


@attr.s
class RedisBaseService(metaclass=abc.ABCMeta):
    APP_KEY: ClassVar[str] = "REDIS_SERVICE"
    _instance_kind: RedisInstanceKind = attr.ib()
    _ssl: Optional[bool] = attr.ib(default=None)

    async def init_hook(self, target_app: web.Application) -> None:
        LOGGER.info(f"Initializing Redis {self._instance_kind.name}")
        target_app[self.get_full_app_key(self._instance_kind)] = self
        await self.initialize()

    async def tear_down_hook(self, target_app: web.Application) -> None:
        await self.tear_down()

    @abc.abstractmethod
    def get_redis(self, allow_slave: bool = False) -> redis.asyncio.Redis:
        pass

    @classmethod
    def get_full_app_key(cls, instance_kind: RedisInstanceKind) -> str:
        return f"{cls.APP_KEY}_{instance_kind.name}"

    @classmethod
    def get_app_instance(cls, app: web.Application, instance_kind: RedisInstanceKind) -> "RedisBaseService":
        service = app.get(cls.get_full_app_key(instance_kind), None)
        if service is None:
            raise ValueError("Redis was not initiated for application")

        return service

    @abc.abstractmethod
    async def initialize(self) -> None:
        pass

    @abc.abstractmethod
    async def tear_down(self) -> None:
        pass


@attr.s(kw_only=True)
class RedisSentinelService(RedisBaseService):
    _sentinel_hosts: Sequence[str] = attr.ib()
    _sentinel_port: int = attr.ib()
    _namespace: str = attr.ib()
    _db: int = attr.ib()
    _password: str = attr.ib(repr=False)
    _cached_master: Optional[redis.asyncio.Redis] = attr.ib(init=False, default=None)
    _cached_slave: Optional[redis.asyncio.Redis] = attr.ib(init=False, default=None)

    _sentinel_client: RedisSentinel = attr.ib(init=False, repr=False, hash=False, cmp=False)

    async def initialize(self) -> None:
        LOGGER.info("Internal initializing Redis Sentinels")
        self._sentinel_client = RedisSentinel(
            [(h, self._sentinel_port) for h in self._sentinel_hosts],
            db=self._db,
            password=self._password,
            ssl=self._ssl,
        )
        await self._sentinel_client.discover_master(self._namespace)

    async def tear_down(self) -> None:
        if self._cached_slave:
            await self._cached_slave.close()
        if self._cached_master:
            await self._cached_master.close()

    def get_redis(self, allow_slave: bool = False) -> redis.asyncio.Redis:
        sentinel_cli = self._sentinel_client
        if allow_slave:
            cached_redis = self._cached_slave
            if cached_redis is None:
                cached_redis = sentinel_cli.slave_for(self._namespace)
                self._cached_slave = cached_redis
        else:
            cached_redis = self._cached_master
            if cached_redis is None:
                cached_redis = sentinel_cli.master_for(self._namespace)
                self._cached_master = cached_redis
        return cached_redis


@attr.s(kw_only=True)
class SingleHostSimpleRedisService(RedisBaseService):
    _url: str = attr.ib()
    _password: Optional[str] = attr.ib(default=None, repr=False)

    _redis_pool: redis.asyncio.ConnectionPool = attr.ib(init=False, repr=False, hash=False, cmp=False)

    async def initialize(self) -> None:
        LOGGER.info("Internal initializing Redis Single Host")
        self._redis_pool = redis.asyncio.ConnectionPool.from_url(self._url, password=self._password)

    async def tear_down(self) -> None:
        LOGGER.info("Tear down redis")
        LOGGER.info("Waiting for redis pool to become closed")
        await self._redis_pool.disconnect()
        LOGGER.info("Redis pool become closed")

    def get_redis(self, allow_slave: bool = False) -> redis.asyncio.Redis:
        return redis.asyncio.Redis(connection_pool=self._redis_pool)

lib/dl_core/dl_core/aio/web_app_services/redis.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
)

from aiobotocore.client import AioBaseClient
from aiobotocore.session import get_session
from aiohttp import web
import attr
import boto3


if TYPE_CHECKING:
    from mypy_boto3_s3.client import S3Client as SyncS3Client
    from types_aiobotocore_s3 import S3Client as AsyncS3Client


LOGGER = logging.getLogger(__name__)


@attr.s
class S3Service:
    APP_KEY: ClassVar[str] = "S3_SERVICE"

    _access_key_id: str = attr.ib(repr=False)
    _secret_access_key: str = attr.ib(repr=False)
    _endpoint_url: str = attr.ib()

    tmp_bucket_name: str = attr.ib()
    persistent_bucket_name: str = attr.ib()

    _client: AioBaseClient = attr.ib(init=False, repr=False, hash=False, cmp=False)
    _client_init_params: dict[str, Any] = attr.ib(init=False, repr=False, hash=False, cmp=False)

    @property
    def client(self) -> AsyncS3Client:
        return self._client  # type: ignore

    @classmethod
    def get_full_app_key(cls) -> str:
        return cls.APP_KEY

    async def init_hook(self, target_app: web.Application) -> None:
        target_app[self.APP_KEY] = self
        await self.initialize()

    async def tear_down_hook(self, target_app: web.Application) -> None:
        await self.tear_down()

    async def initialize(self) -> None:
        LOGGER.info("Initializing S3 service")
        self._client_init_params = dict(
            service_name="s3",
            aws_access_key_id=self._access_key_id,
            aws_secret_access_key=self._secret_access_key,
            endpoint_url=self._endpoint_url,
        )

        session = get_session()
        client = session.create_client(**self._client_init_params)
        self._client = await client.__aenter__()

    async def tear_down(self) -> None:
        LOGGER.info("Tear down S3 service")
        await self._client.close()

    @classmethod
    def get_app_instance(cls, app: web.Application) -> "S3Service":
        service = app.get(cls.APP_KEY, None)
        if service is None:
            raise ValueError("S3BucketService was not initiated for application")

        return service

    def get_client(self) -> AsyncS3Client:
        return self.client

    def get_sync_client(self) -> SyncS3Client:
        session = boto3.session.Session()
        client = session.client(**self._client_init_params)
        return client

lib/dl_core/dl_core/aio/web_app_services/s3.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    SourceBackendType,
)


_CONNECTION_BACKEND_TYPES: dict[ConnectionType, SourceBackendType] = {}


def register_connection_backend_type(conn_type: ConnectionType, backend_type: SourceBackendType) -> None:
    if (registered_backend_type := _CONNECTION_BACKEND_TYPES.get(conn_type)) is not None:
        assert registered_backend_type == backend_type
    else:
        _CONNECTION_BACKEND_TYPES[conn_type] = backend_type


def get_backend_type(conn_type: ConnectionType) -> SourceBackendType:
    """Return class for given connection type"""
    return _CONNECTION_BACKEND_TYPES[conn_type]

lib/dl_core/dl_core/backend_types.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    Any,
    Optional,
)

import attr

from dl_constants.enums import (
    ManagedBy,
    RawSQLLevel,
    WhereClauseOperation,
)
from dl_utils.utils import DataKey


LOGGER = logging.getLogger(__name__)


# Connection references


@attr.s(frozen=True, slots=True)
class ConnectionRef:
    pass


@attr.s(frozen=True, slots=True)
class DefaultConnectionRef(ConnectionRef):
    conn_id: str = attr.ib(kw_only=True)


@attr.s(frozen=True, slots=True)
class InternalMaterializationConnectionRef(ConnectionRef):
    pass


def connection_ref_from_id(connection_id: Optional[str]) -> ConnectionRef:
    if connection_id is None:
        # TODO REMOVE: some sample source code still relies on mat con ref
        return InternalMaterializationConnectionRef()
    else:
        return DefaultConnectionRef(conn_id=connection_id)


@attr.s()
class DefaultWhereClause:
    operation: WhereClauseOperation = attr.ib()
    values: list = attr.ib(default=[])


@attr.s()
class ObligatoryFilter:
    id: str = attr.ib()
    field_guid: str = attr.ib()
    default_filters: list[DefaultWhereClause] = attr.ib()
    managed_by: ManagedBy = attr.ib(default=ManagedBy.user)
    valid: bool = attr.ib(default=True)


@attr.s()
class SourceFilterSpec:
    name: str = attr.ib()
    operation: WhereClauseOperation = attr.ib()
    values: list = attr.ib(default=[])


class EntryLocation(metaclass=abc.ABCMeta):
    def to_short_string(self) -> str:
        raise NotImplementedError()

    def to_us_req_api_params(self) -> dict[str, Any]:
        raise NotImplementedError()

    def to_us_resp_api_params(self, key_from_us_resp: Optional[str]) -> dict[str, str]:
        raise NotImplementedError()


@attr.s(frozen=True)
class PathEntryLocation(EntryLocation):
    path: str = attr.ib()

    def to_short_string(self) -> str:
        return self.path

    def to_us_req_api_params(self) -> dict[str, Any]:
        return {"key": self.path}

    def to_us_resp_api_params(self, key_from_us_resp: Optional[str]) -> dict[str, str]:
        return {"key": self.path}


@attr.s(frozen=True)
class WorkbookEntryLocation(EntryLocation):
    workbook_id: str = attr.ib()
    entry_name: str = attr.ib()

    def to_short_string(self) -> str:
        return f"WB({self.workbook_id}/{self.entry_name})"

    def to_us_req_api_params(self) -> dict[str, Any]:
        return {
            "workbookId": self.workbook_id,
            "name": self.entry_name,
        }

    def to_us_resp_api_params(self, key_from_us_resp: Optional[str]) -> dict[str, str]:
        restored_key: str
        if key_from_us_resp is None:
            restored_key = f"dummy_workbook_path_repr/{self.entry_name}"
        else:
            restored_key = key_from_us_resp

        return {"key": restored_key, "workbookId": self.workbook_id}


@attr.s
class BaseAttrsDataModel:
    """Marker class for attrs-based data"""

    @classmethod
    def get_secret_keys(cls) -> set[DataKey]:
        return set()


@attr.s
class ConnectionDataModelBase(BaseAttrsDataModel):
    pass


@attr.s
class ConnCacheableDataModelMixin:
    cache_ttl_sec: Optional[int] = attr.ib(kw_only=True, default=None)  # Override for default cache TTL


@attr.s
class ConnSubselectDataModelMixin:
    raw_sql_level: RawSQLLevel = attr.ib(kw_only=True, default=RawSQLLevel.off)

lib/dl_core/dl_core/base_models.py
datalens-tech/datalens-backend
#!/usr/bin/env python3

from __future__ import annotations

from dl_core.connection_executors.remote_query_executor import app_sync as app_sync_mod


app = app_sync_mod.create_sync_app()


app_sync_mod.hook_init_logging(app)

lib/dl_core/dl_core/bin/query_executor_sync.py
datalens-tech/datalens-backend
# TODO: Move to dl_core.components.*

from __future__ import annotations

from typing import (
    Dict,
    List,
    Optional,
    Sequence,
    Type,
)

import attr

from dl_constants.enums import (
    ComponentErrorLevel,
    ComponentType,
)


@attr.s
class ComponentError:
    level: ComponentErrorLevel = attr.ib()
    message: str = attr.ib()
    code: List[str] = attr.ib(factory=list)
    details: dict = attr.ib(factory=dict)


@attr.s
class ComponentErrorPack:
    id: str = attr.ib()
    type: ComponentType = attr.ib()
    errors: List[ComponentError] = attr.ib(factory=list)

    def get_errors(
        self,
        code: Optional[Sequence[str]] = None,
        code_prefix: Optional[Sequence[str]] = None,
    ) -> List[ComponentError]:
        if code is not None and code_prefix is not None:
            raise ValueError("Cannot specify both code and code_prefix")

        if code is not None:
            code = list(code)
            cond = lambda err: err.code == code
        elif code_prefix is not None:
            code_prefix = list(code_prefix)
            cond = lambda err: err.code[: len(code_prefix)] == code_prefix  # type: ignore  # TODO: fix
        else:
            cond = lambda err: True
        return [err for err in self.errors if cond(err)]

    def remove_errors(
        self,
        code: Optional[Sequence[str]] = None,
        code_prefix: Optional[Sequence[str]] = None,
    ) -> None:
        for err in self.get_errors(code=code, code_prefix=code_prefix):
            self.errors.remove(err)


# FIXME: switch to subclass for fields (with row, column, token)
ERROR_CLS_BY_COMP_TYPE: Dict[ComponentType, Type[ComponentError]] = {
    ComponentType.data_source: ComponentError,
    ComponentType.source_avatar: ComponentError,
    ComponentType.avatar_relation: ComponentError,
    ComponentType.field: ComponentError,
    ComponentType.obligatory_filter: ComponentError,
    ComponentType.result_schema: ComponentError,
}


@attr.s
class ComponentErrorRegistry:
    items: List[ComponentErrorPack] = attr.ib(factory=list)

    def get_pack(self, id: str) -> Optional[ComponentErrorPack]:  # type: ignore  # TODO: fix
        for item in self.items:
            if item.id == id:
                return item

    def remove_errors(
        self,
        id: str,
        code: Optional[Sequence[str]] = None,
        code_prefix: Optional[Sequence[str]] = None,
    ) -> None:
        """
        Remove errors for object with given ID.
        if code is given, remove only errors having the given error code.
        """
        item = self.get_pack(id=id)
        if item is None:
            return

        item.remove_errors(code=code, code_prefix=code_prefix)
        if not item.errors:
            self.items.remove(item)

    def add_error(  # type: ignore  # TODO: fix
        self,
        id: str,
        type: ComponentType,
        message: str,
        code: Sequence[str],
        level: ComponentErrorLevel = ComponentErrorLevel.error,
        details: Optional[dict] = None,
    ):
        details = details or {}
        error_cls = ERROR_CLS_BY_COMP_TYPE.get(type)
        if error_cls is None:
            raise ValueError(f"Invalid value for type: {type}")

        item = self.get_pack(id=id)
        if item is None:
            item = ComponentErrorPack(id=id, type=type)
            self.items.append(item)

        item.errors.append(error_cls(message=message, code=list(code), level=level, details=details))

    def for_type(self, type: ComponentType) -> List[ComponentErrorPack]:
        return [item for item in self.items if item.type == type]

    def rename_pack(self, old_id: str, new_id: str) -> None:
        for item in self.items:
            if item.id == old_id:
                item.id = new_id

lib/dl_core/dl_core/component_errors.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Optional,
)

import attr

import dl_core.exc as exc


if TYPE_CHECKING:
    from dl_core.base_models import ObligatoryFilter
    from dl_core.data_source_spec.collection import DataSourceCollectionSpecBase
    from dl_core.multisource import (
        AvatarRelation,
        SourceAvatar,
    )
    from dl_core.us_dataset import Dataset


@attr.s(frozen=True)
class DatasetComponentAccessor:
    _dataset: Dataset = attr.ib(kw_only=True)

    def get_avatar_list(self, source_id: Optional[str] = None) -> list[SourceAvatar]:
        """
        Return list of source avatars
        :param source_id: limit result to avatars with this source ID
        :return: list of ``SourceAvatar`` objects
        """
        return [
            avatar for avatar in self._dataset.data.source_avatars if (not source_id or avatar.source_id == source_id)
        ]

    def has_avatar(self, avatar_id: str) -> bool:
        return self.get_avatar_opt(avatar_id=avatar_id) is not None

    def get_avatar_opt(self, avatar_id: str) -> Optional[SourceAvatar]:
        for avatar in self._dataset.data.source_avatars:
            if avatar.id == avatar_id:
                return avatar
        return None

    def get_avatar_strict(self, avatar_id: str) -> SourceAvatar:
        avatar = self.get_avatar_opt(avatar_id=avatar_id)
        if avatar is None:
            raise exc.SourceAvatarNotFound(f"Unknown source avatar: {avatar_id}")
        return avatar

    def get_root_avatar_opt(self) -> Optional[SourceAvatar]:
        for avatar in self._dataset.data.source_avatars:
            if avatar.is_root:
                return avatar
        return None

    def get_root_avatar_strict(self) -> SourceAvatar:
        root_avatar = self.get_root_avatar_opt()
        if root_avatar is None:
            raise exc.DatasetConfigurationError("No avatars found")
        return root_avatar

    def get_data_source_id_list(self) -> list[str]:
        return [dsrc_coll_spec.id for dsrc_coll_spec in self.get_data_source_coll_spec_list()]

    def get_data_source_coll_spec_list(self) -> list[DataSourceCollectionSpecBase]:
        return [dsrc_coll_spec for dsrc_coll_spec in self._dataset.data.source_collections]

    def get_data_source_coll_spec_opt(self, source_id: str) -> Optional[DataSourceCollectionSpecBase]:
        for dsrc_coll_spec in self._dataset.data.source_collections:
            if dsrc_coll_spec.id == source_id:
                return dsrc_coll_spec
        return None

    def get_data_source_coll_spec_strict(self, source_id: str) -> DataSourceCollectionSpecBase:
        dsrc_coll_spec = self.get_data_source_coll_spec_opt(source_id=source_id)
        if dsrc_coll_spec is None:
            raise exc.DataSourceNotFound(f"Unknown data source: {source_id}")
        return dsrc_coll_spec

    def get_avatar_relation_list(
        self, left_avatar_id: Optional[str] = None, right_avatar_id: Optional[str] = None
    ) -> list[AvatarRelation]:
        """
        Return list of relations
        :param left_avatar_id: limit result to relations with this left ID
        :param right_avatar_id: limit result to relations with this left ID
        :return: list of ``AvatarRelation`` objects
        """
        result = []
        for relation in self._dataset.data.avatar_relations:
            if (not left_avatar_id or relation.left_avatar_id == left_avatar_id) and (
                not right_avatar_id or relation.right_avatar_id == right_avatar_id
            ):
                result.append(relation)
        return result

    def has_avatar_relation(
        self,
        relation_id: Optional[str] = None,
        left_avatar_id: Optional[str] = None,
        right_avatar_id: Optional[str] = None,
    ) -> bool:
        return (
            self.get_avatar_relation_opt(
                relation_id=relation_id,
                left_avatar_id=left_avatar_id,
                right_avatar_id=right_avatar_id,
            )
            is not None
        )

    def get_avatar_relation_opt(
        self,
        relation_id: Optional[str] = None,
        left_avatar_id: Optional[str] = None,
        right_avatar_id: Optional[str] = None,
    ) -> Optional[AvatarRelation]:
        for relation in self._dataset.data.avatar_relations:
            if (
                relation_id is not None
                and relation.id == relation_id
                or (relation.left_avatar_id == left_avatar_id and relation.right_avatar_id == right_avatar_id)
            ):
                return relation
        return None

    def get_avatar_relation_strict(
        self,
        relation_id: Optional[str] = None,
        left_avatar_id: Optional[str] = None,
        right_avatar_id: Optional[str] = None,
    ) -> AvatarRelation:
        relation = self.get_avatar_relation_opt(
            relation_id=relation_id,
            left_avatar_id=left_avatar_id,
            right_avatar_id=right_avatar_id,
        )
        if relation is None:
            raise exc.AvatarRelationNotFound(f"Unknown avatar relation: {relation_id, left_avatar_id, right_avatar_id}")
        return relation

    def get_obligatory_filter_list(self) -> list[ObligatoryFilter]:
        return self._dataset.data.obligatory_filters

    def has_obligatory_filter(
        self,
        obfilter_id: Optional[str] = None,
        field_guid: Optional[str] = None,
    ) -> bool:
        return (
            self.get_obligatory_filter_opt(
                obfilter_id=obfilter_id,
                field_guid=field_guid,
            )
            is not None
        )

    def get_obligatory_filter_opt(
        self,
        obfilter_id: Optional[str] = None,
        field_guid: Optional[str] = None,
    ) -> Optional[ObligatoryFilter]:
        if obfilter_id is not None:
            for filter_object in self._dataset.data.obligatory_filters:
                if filter_object.id == obfilter_id:
                    return filter_object
        elif field_guid is not None:
            for filter_object in self._dataset.data.obligatory_filters:
                if filter_object.field_guid == field_guid:
                    return filter_object
        else:
            raise ValueError('"id" and "field_guid" parameters can not be None together')

        return None

    def get_obligatory_filter_strict(
        self,
        obfilter_id: Optional[str] = None,
        field_guid: Optional[str] = None,
    ) -> ObligatoryFilter:
        obfilter = self.get_obligatory_filter_opt(obfilter_id=obfilter_id, field_guid=field_guid)
        if obfilter is None:
            raise exc.ObligatoryFilterNotFound(f"Unknown obligatory filter: {obfilter_id, field_guid}")
        return obfilter

lib/dl_core/dl_core/components/accessor.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Collection,
    Dict,
    Optional,
    Set,
    Tuple,
)

import attr

from dl_constants.enums import ManagedBy
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.dependencies.avatar_tree_base import AvatarTreeResolverBase
from dl_core.components.dependencies.relation_avatar_base import RelationAvatarDependencyManagerBase
from dl_core.components.ids import (
    AvatarId,
    RelationId,
)
import dl_core.exc as exc
from dl_core.us_dataset import Dataset


LOGGER = logging.getLogger(__name__)


@attr.s
class AvatarTreeResolver(AvatarTreeResolverBase):
    _dataset: Dataset = attr.ib(kw_only=True)
    _relation_avatar_dep_mgr: RelationAvatarDependencyManagerBase = attr.ib(kw_only=True)
    _ds_accessor: DatasetComponentAccessor = attr.ib(init=False)

    @_ds_accessor.default
    def _make_ds_accessor(self) -> DatasetComponentAccessor:
        return DatasetComponentAccessor(dataset=self._dataset)

    def rank_avatars(self) -> Dict[AvatarId, int]:
        def populate_recursively(avatar_id: str, rank: int) -> None:
            ranks[avatar_id] = rank
            for relation in self._ds_accessor.get_avatar_relation_list(left_avatar_id=avatar_id):
                if relation.managed_by == ManagedBy.feature:
                    continue
                elif relation.managed_by == ManagedBy.user:
                    populate_recursively(avatar_id=relation.right_avatar_id, rank=rank + 1)
                else:
                    raise ValueError(f"Unsupported managed_by value in relation: {relation.managed_by}")

        root_avatar = self._ds_accessor.get_root_avatar_strict()
        ranks: Dict[str, int] = {}
        populate_recursively(avatar_id=root_avatar.id, rank=0)
        return ranks

    def expand_required_avatar_ids(
        self, required_avatar_ids: Collection[str]
    ) -> Tuple[Optional[AvatarId], Set[AvatarId], Set[RelationId]]:
        if len(required_avatar_ids) == 1:
            # Single avatar -> nothing to resolve
            return next(iter(required_avatar_ids)), set(required_avatar_ids), set()

        required_avatar_ids = set(required_avatar_ids)
        ranks = self.rank_avatars()
        if len(required_avatar_ids) > 1:
            LOGGER.info(f"Got avatar ranks: {ranks}")

        required_relation_ids: Set[RelationId] = set()

        # first add avatars used by required feature-managed relations
        for iteration in range(10):
            updated_required_avatar_ids = required_avatar_ids.copy()
            for relation in self._ds_accessor.get_avatar_relation_list():
                if relation.right_avatar_id in required_avatar_ids and relation.managed_by == ManagedBy.feature:
                    updated_required_avatar_ids |= self._relation_avatar_dep_mgr.get_relation_avatar_references(
                        relation_id=relation.id
                    )
                    required_relation_ids.add(relation.id)
            if updated_required_avatar_ids == required_avatar_ids:
                LOGGER.info(f"Finished resolving feature-managed avatars on iteration {iteration}")
                break
            else:
                LOGGER.info(
                    "Found additional avatars in required feature-managed relations: "
                    f"{updated_required_avatar_ids - required_avatar_ids}"
                )
            required_avatar_ids = updated_required_avatar_ids
        else:
            # It means that for every iteration we are still getting new avatars
            raise RuntimeError("Failed to resolve required avatars")  # Should not happen

        user_avatar_ids = {
            avatar_id
            for avatar_id in required_avatar_ids
            if self._ds_accessor.get_avatar_strict(avatar_id=avatar_id).managed_by == ManagedBy.user
        }
        assert user_avatar_ids, "Must have at least one user-managed source"

        def get_relation_and_parent_id(avatar_id: AvatarId) -> Tuple[RelationId, AvatarId]:
            relations = self._ds_accessor.get_avatar_relation_list(right_avatar_id=avatar_id)
            assert len(relations) == 1
            relation = relations[0]
            return relation.id, relation.left_avatar_id

        try:
            min_rank = min(rank for avatar_id, rank in ranks.items() if avatar_id in user_avatar_ids)
        except ValueError:
            raise exc.UnboundAvatarError(
                "Detected the usage of an unbound avatar. " "Dataset is configured incorrectly."
            )

        # add parent_ids until single common ancestor remains
        common_root_ids = user_avatar_ids
        while len(common_root_ids) > 1:
            new_common_root_ids = set()
            for avatar_id in common_root_ids:
                # now fill in the gaps so that all `new_common_root_ids` are at the same level in the avatar tree
                while ranks[avatar_id] > min_rank:
                    new_relation_id, new_avatar_id = get_relation_and_parent_id(avatar_id)
                    # add this parent to required IDs
                    if new_avatar_id not in required_avatar_ids:
                        LOGGER.info(
                            f"Implicitly adding avatar {new_avatar_id} " f"so that avatar {avatar_id} can be joined to"
                        )
                        required_avatar_ids.add(new_avatar_id)
                    avatar_id = new_avatar_id
                    required_relation_ids.add(new_relation_id)

                new_common_root_ids.add(avatar_id)

            common_root_ids = new_common_root_ids
            if len(common_root_ids) > 1:
                min_rank -= 1

        return next(iter(common_root_ids)), required_avatar_ids, required_relation_ids

lib/dl_core/dl_core/components/dependencies/avatar_tree.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    Collection,
    Dict,
    Optional,
    Set,
    Tuple,
)

from dl_core.components.ids import (
    AvatarId,
    RelationId,
)


class AvatarTreeResolverBase(abc.ABC):
    @abc.abstractmethod
    def rank_avatars(self) -> Dict[AvatarId, int]:
        """
        Return mapping: ``{avatar_id: avatar_rank}``,
        where:
        - the root avatar has rank 0,
        - avatars joined directly to the root have rank 1,
        - and so on.

        Feature-managed avatars are excluded because their formal rank is always 1,
        but in reality their relation expressions may require non-root avatars.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def expand_required_avatar_ids(
        self, required_avatar_ids: Collection[str]
    ) -> Tuple[Optional[AvatarId], Set[AvatarId], Set[RelationId]]:
        """
        Complement required avatars wit ones that are required to apply JOINs of the given avatars

        :param required_avatar_ids: avatars that must be used in the query
        :param relation_expressions: compiled relation expressions that contain info
            about avatars required by each JOIN
        :return: root avatar ID, the complemented set of avatar IDs, required relation IDs

        Illustration:
        ::

            Avatar tree:

             A0 -->  A1  -->  A2  --> [A4]
               |         --> [A3]
               |         -->  A5
               |
                --> [F1]  # feature-managed avatar (joined to via an expression that uses a field from A5)

            Parameters:
                required_avatar_ids: {A3, A4, F1}  (emphasized with brackets in the diagram)
                relation_expressions (req avatars by each relation):
                    {
                        # user-managed:
                        A0-A1: {A0, A1},
                        A1-A2: {A1, A2},
                        A1-A3: {A1, A3},
                        A1-A5: {A1, A5},
                        A2-A4: {A2, A4},
                        # feature-managed:
                        A0-F1: {A5, F1},  # via result_field, so it might not really need A0
                    }

        1. First we inspect feature-managed relations that are needed by required avatars (``A0-F1`` in this case)
        and add avatars required by their expressions to ``required_avatar_ids``.

        After this ``required_avatar_ids`` becomes ``{A3, A4, F1, A5}``;
        ``required_relation_ids``: ``{A0-F1}``.

        2. Now, we exclude feature managed avatars (``F1``): ``user_avatar_ids = {A3, A4, A5}``,
        traverse the tree in search of their latest common ancestor (LUCA)
        and add all avatars we encounter on the way (``A2``, ``A1``)

        Finally ``required_avatar_ids = {A3, A4, F1, A5, A2, A1}``;
        ``required_relation_ids``: ``{A0-F1, A1-A2, A1-A3, A1-A5, A2-A4}``.

        3. Return LUCA's ID (``A1``), updated ``required_avatar_ids`` and ``required_relation_ids``.
        """
        raise NotImplementedError

lib/dl_core/dl_core/components/dependencies/avatar_tree_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_core.components.dependencies.avatar_tree import AvatarTreeResolver
from dl_core.components.dependencies.factory_base import ComponentDependencyManagerFactoryBase
from dl_core.components.dependencies.field_avatar import FieldAvatarDependencyManager
from dl_core.components.dependencies.field_deep import FieldDeepInterDependencyManager
from dl_core.components.dependencies.field_shallow import FieldShallowInterDependencyManager
from dl_core.components.dependencies.relation_avatar import RelationAvatarDependencyManager
from dl_core.us_dataset import Dataset


@attr.s
class ComponentDependencyManagerFactory(ComponentDependencyManagerFactoryBase):  # noqa
    _dataset: Dataset = attr.ib(kw_only=True)

    def _make_field_shallow_inter_dependency_manager(self) -> FieldShallowInterDependencyManager:
        return FieldShallowInterDependencyManager(
            inter_dep_info=self._dataset.data.result_schema_aux.inter_dependencies,
        )

    def _make_field_deep_inter_dependency_manager(self) -> FieldDeepInterDependencyManager:
        return FieldDeepInterDependencyManager(
            shallow=self.get_field_shallow_inter_dependency_manager(),
        )

    def _make_field_avatar_dependency_manager(self) -> FieldAvatarDependencyManager:
        return FieldAvatarDependencyManager(
            result_schema=self._dataset.result_schema,
            deep=self.get_field_deep_inter_dependency_manager(),
        )

    def _make_relation_avatar_dependency_manager(self) -> RelationAvatarDependencyManager:
        return RelationAvatarDependencyManager(
            avatar_relations=self._dataset.data.avatar_relations,
            field_avatar_mgr=self.get_field_avatar_dependency_manager(),
        )

    def _make_avatar_tree_resolver(self) -> AvatarTreeResolver:
        return AvatarTreeResolver(
            dataset=self._dataset, relation_avatar_dep_mgr=self.get_relation_avatar_dependency_manager()
        )

lib/dl_core/dl_core/components/dependencies/factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import Optional

import attr

from dl_core.components.dependencies.avatar_tree_base import AvatarTreeResolverBase
from dl_core.components.dependencies.field_avatar_base import FieldAvatarDependencyManagerBase
from dl_core.components.dependencies.field_deep_base import FieldDeepInterDependencyManagerBase
from dl_core.components.dependencies.field_shallow_base import FieldShallowInterDependencyManagerBase
from dl_core.components.dependencies.relation_avatar_base import RelationAvatarDependencyManagerBase


@attr.s
class ComponentDependencyManagerFactoryBase(abc.ABC):
    _field_shallow_inter_dependency_manager: Optional[FieldShallowInterDependencyManagerBase] = attr.ib(
        kw_only=True, default=None
    )
    _field_deep_inter_dependency_manager: Optional[FieldDeepInterDependencyManagerBase] = attr.ib(
        kw_only=True, default=None
    )
    _field_avatar_dependency_manager: Optional[FieldAvatarDependencyManagerBase] = attr.ib(kw_only=True, default=None)
    _relation_avatar_dependency_manager: Optional[RelationAvatarDependencyManagerBase] = attr.ib(
        kw_only=True, default=None
    )
    _avatar_tree_resolver: Optional[AvatarTreeResolverBase] = attr.ib(kw_only=True, default=None)

    @abc.abstractmethod
    def _make_field_shallow_inter_dependency_manager(self) -> FieldShallowInterDependencyManagerBase:
        raise NotImplementedError

    @abc.abstractmethod
    def _make_field_deep_inter_dependency_manager(self) -> FieldDeepInterDependencyManagerBase:
        raise NotImplementedError

    @abc.abstractmethod
    def _make_field_avatar_dependency_manager(self) -> FieldAvatarDependencyManagerBase:
        raise NotImplementedError

    @abc.abstractmethod
    def _make_relation_avatar_dependency_manager(self) -> RelationAvatarDependencyManagerBase:
        raise NotImplementedError

    @abc.abstractmethod
    def _make_avatar_tree_resolver(self) -> AvatarTreeResolverBase:
        raise NotImplementedError

    def get_field_shallow_inter_dependency_manager(self) -> FieldShallowInterDependencyManagerBase:
        if self._field_shallow_inter_dependency_manager is None:
            self._field_shallow_inter_dependency_manager = self._make_field_shallow_inter_dependency_manager()
        assert self._field_shallow_inter_dependency_manager is not None
        return self._field_shallow_inter_dependency_manager

    def get_field_deep_inter_dependency_manager(self) -> FieldDeepInterDependencyManagerBase:
        if self._field_deep_inter_dependency_manager is None:
            self._field_deep_inter_dependency_manager = self._make_field_deep_inter_dependency_manager()
        assert self._field_deep_inter_dependency_manager is not None
        return self._field_deep_inter_dependency_manager

    def get_field_avatar_dependency_manager(self) -> FieldAvatarDependencyManagerBase:
        if self._field_avatar_dependency_manager is None:
            self._field_avatar_dependency_manager = self._make_field_avatar_dependency_manager()
        assert self._field_avatar_dependency_manager is not None
        return self._field_avatar_dependency_manager

    def get_relation_avatar_dependency_manager(self) -> RelationAvatarDependencyManagerBase:
        if self._relation_avatar_dependency_manager is None:
            self._relation_avatar_dependency_manager = self._make_relation_avatar_dependency_manager()
        assert self._relation_avatar_dependency_manager is not None
        return self._relation_avatar_dependency_manager

    def get_avatar_tree_resolver(self) -> AvatarTreeResolverBase:
        if self._avatar_tree_resolver is None:
            self._avatar_tree_resolver = self._make_avatar_tree_resolver()
        assert self._avatar_tree_resolver is not None
        return self._avatar_tree_resolver

lib/dl_core/dl_core/components/dependencies/factory_base.py
datalens-tech/datalens-backend
# FIXME: Remove after transition to US-saved field depoendencies

from __future__ import annotations

from typing import Dict

import attr

from dl_core.components.dependencies.factory import ComponentDependencyManagerFactory
from dl_core.components.dependencies.relation_avatar_legacy import LegacyRelationAvatarDependencyManager
from dl_core.components.ids import RelationId
from dl_core.query.expression import ExpressionCtx


@attr.s
class LegacyComponentDependencyManagerFactory(ComponentDependencyManagerFactory):
    _relation_expressions: Dict[RelationId, ExpressionCtx] = attr.ib(factory=list)

    def set_relation_expressions(self, _relation_expressions: Dict[RelationId, ExpressionCtx]) -> None:
        self._relation_expressions = _relation_expressions

    def _make_relation_avatar_dependency_manager(self) -> LegacyRelationAvatarDependencyManager:  # type: ignore  # TODO: fix
        return LegacyRelationAvatarDependencyManager(
            relation_expressions=self._relation_expressions,
        )

lib/dl_core/dl_core/components/dependencies/factory_legacy.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    AbstractSet,
    Set,
)

import attr

from dl_constants.enums import CalcMode
from dl_core.components.dependencies.field_avatar_base import FieldAvatarDependencyManagerBase
from dl_core.components.dependencies.field_deep_base import FieldDeepInterDependencyManagerBase
from dl_core.components.ids import (
    AvatarId,
    FieldId,
)
from dl_core.fields import ResultSchema


@attr.s
class FieldAvatarDependencyManager(FieldAvatarDependencyManagerBase):
    _result_schema: ResultSchema = attr.ib(kw_only=True)
    _deep: FieldDeepInterDependencyManagerBase = attr.ib(kw_only=True)

    def get_field_avatar_references(self, dep_field_id: FieldId) -> AbstractSet[AvatarId]:
        deep_refs = self._deep.get_field_deep_references(dep_field_id=dep_field_id) | {dep_field_id}

        avatar_refs: Set[AvatarId] = set()
        for ref_field_if in deep_refs:
            field = self._result_schema.by_guid(ref_field_if)
            if field.calc_mode == CalcMode.direct:
                assert field.avatar_id is not None
                avatar_refs.add(field.avatar_id)

        return avatar_refs

lib/dl_core/dl_core/components/dependencies/field_avatar.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import AbstractSet

from dl_core.components.ids import (
    AvatarId,
    FieldId,
)


class FieldAvatarDependencyManagerBase:
    @abc.abstractmethod
    def get_field_avatar_references(self, dep_field_id: FieldId) -> AbstractSet[AvatarId]:
        """
        Return a set of all avatars referenced directly or indirectly by the given field.
        """
        raise NotImplementedError

lib/dl_core/dl_core/components/dependencies/field_avatar_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    AbstractSet,
    Set,
)

import attr

from dl_core.components.dependencies.field_deep_base import FieldDeepInterDependencyManagerBase
from dl_core.components.dependencies.field_shallow_base import FieldShallowInterDependencyManagerBase
from dl_core.components.ids import FieldId


@attr.s
class FieldDeepInterDependencyManager(FieldDeepInterDependencyManagerBase):
    _shallow: FieldShallowInterDependencyManagerBase = attr.ib(kw_only=True)

    def get_field_deep_references(self, dep_field_id: FieldId) -> AbstractSet[FieldId]:
        refs: Set[FieldId] = set()

        def _populate_refs_recursively(field_id: FieldId) -> AbstractSet[FieldId]:  # type: ignore  # TODO: fix
            if field_id in refs:
                return set()
            refs.add(field_id)
            new_refs = self._shallow.get_field_direct_references(dep_field_id=field_id)
            for _id in new_refs:
                _populate_refs_recursively(_id)

        _populate_refs_recursively(dep_field_id)
        refs.remove(dep_field_id)
        return refs

lib/dl_core/dl_core/components/dependencies/field_deep.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import AbstractSet

from dl_core.components.ids import FieldId


class FieldDeepInterDependencyManagerBase(abc.ABC):
    """Builds deep field dependencies"""

    @abc.abstractmethod
    def get_field_deep_references(self, dep_field_id: FieldId) -> AbstractSet[FieldId]:
        """
        Walk field dependencies and collect directly referenced fields,
        fields referenced by directly referenced fields and so on.
        Return a set of these referenced fields
        """
        raise NotImplementedError

lib/dl_core/dl_core/components/dependencies/field_deep_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    AbstractSet,
    Collection,
    List,
    Optional,
)

import attr

from dl_core.components.dependencies.field_shallow_base import FieldShallowInterDependencyManagerBase
from dl_core.components.dependencies.primitives import (
    FieldInterDependencyInfo,
    FieldInterDependencyItem,
)
from dl_core.components.ids import FieldId


@attr.s
class FieldShallowInterDependencyManager(FieldShallowInterDependencyManagerBase):
    _inter_dep_info: FieldInterDependencyInfo = attr.ib(kw_only=FieldInterDependencyInfo)  # type: ignore  # TODO: fix

    @property
    def _direct_dependencies(self) -> List[FieldInterDependencyItem]:
        return self._inter_dep_info.deps

    def _get_item_for_field(self, dep_field_id: FieldId) -> Optional[FieldInterDependencyItem]:  # type: ignore  # TODO: fix
        for item in self._direct_dependencies:
            if item.dep_field_id == dep_field_id:
                return item

    def set_field_direct_references(self, dep_field_id: FieldId, ref_field_ids: Collection[FieldId]) -> None:
        item = self._get_item_for_field(dep_field_id=dep_field_id)
        if item is None:
            item = FieldInterDependencyItem(dep_field_id=dep_field_id)
            self._direct_dependencies.append(item)

        assert item is not None
        item.ref_field_ids = set(ref_field_ids)

        if not item.ref_field_ids:
            self._direct_dependencies.remove(item)

    def clear_field_direct_references(self, dep_field_id: FieldId) -> None:
        item = self._get_item_for_field(dep_field_id=dep_field_id)
        if item is not None:
            self._direct_dependencies.remove(item)

    def get_field_direct_references(self, dep_field_id: FieldId) -> AbstractSet[FieldId]:
        item = self._get_item_for_field(dep_field_id=dep_field_id)
        if item is not None:
            ref_field_ids = item.ref_field_ids
        else:
            ref_field_ids = set()
        return ref_field_ids

lib/dl_core/dl_core/components/dependencies/field_shallow.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    AbstractSet,
    Collection,
)

from dl_core.components.ids import FieldId


class FieldShallowInterDependencyManagerBase(abc.ABC):
    @abc.abstractmethod
    def set_field_direct_references(self, dep_field_id: FieldId, ref_field_ids: Collection[FieldId]) -> None:
        raise NotImplementedError

    @abc.abstractmethod
    def get_field_direct_references(self, dep_field_id: FieldId) -> AbstractSet[FieldId]:
        raise NotImplementedError

lib/dl_core/dl_core/components/dependencies/field_shallow_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    List,
    Set,
)

import attr

from dl_core.components.ids import FieldId


@attr.s
class FieldInterDependencyItem:
    dep_field_id: FieldId = attr.ib(kw_only=True)
    ref_field_ids: Set[FieldId] = attr.ib(kw_only=True, factory=set, converter=set)


@attr.s
class FieldInterDependencyInfo:
    """Stores direct (shallow) field-interdependencies."""

    deps: List[FieldInterDependencyItem] = attr.ib(factory=list)

lib/dl_core/dl_core/components/dependencies/primitives.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    AbstractSet,
    List,
    Set,
)

import attr

from dl_core.components.dependencies.field_avatar_base import FieldAvatarDependencyManagerBase
from dl_core.components.dependencies.relation_avatar_base import RelationAvatarDependencyManagerBase
from dl_core.components.ids import (
    AvatarId,
    RelationId,
)
from dl_core.multisource import (
    AvatarRelation,
    BinaryCondition,
    ConditionPart,
    ConditionPartDirect,
    ConditionPartResultField,
)


@attr.s
class RelationAvatarDependencyManager(RelationAvatarDependencyManagerBase):
    _avatar_relations: List[AvatarRelation] = attr.ib(kw_only=True)  # FIXME: Replace with RelationManager or smth?
    _field_avatar_mgr: FieldAvatarDependencyManagerBase = attr.ib(kw_only=True)

    def _get_relation(self, relation_id: RelationId) -> AvatarRelation:
        for relation in self._avatar_relations:
            if relation.id == relation_id:
                return relation
        raise KeyError(relation_id)

    def _get_condition_part_avatar_ids(
        self,
        default_avatar_id: AvatarId,
        part: ConditionPart,
    ) -> AbstractSet[AvatarId]:
        if isinstance(part, ConditionPartDirect):
            return {default_avatar_id}
        elif isinstance(part, ConditionPartResultField):
            return self._field_avatar_mgr.get_field_avatar_references(part.field_id)
        raise TypeError(f"Unsupported type for part: {type(part)}")

    def _get_condition_avatar_ids(
        self,
        relation: AvatarRelation,
        condition: BinaryCondition,
    ) -> AbstractSet[AvatarId]:
        return self._get_condition_part_avatar_ids(
            part=condition.left_part,
            default_avatar_id=relation.left_avatar_id,
        ) | self._get_condition_part_avatar_ids(
            part=condition.right_part,
            default_avatar_id=relation.right_avatar_id,
        )

    def get_relation_avatar_references(self, relation_id: RelationId) -> AbstractSet[AvatarId]:
        relation = self._get_relation(relation_id)
        avatar_refs: Set[AvatarId] = set()
        for condition in relation.conditions:
            assert isinstance(condition, BinaryCondition)
            avatar_refs |= self._get_condition_avatar_ids(relation=relation, condition=condition)

        return avatar_refs

lib/dl_core/dl_core/components/dependencies/relation_avatar.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import AbstractSet

from dl_core.components.ids import (
    AvatarId,
    RelationId,
)


class RelationAvatarDependencyManagerBase:
    @abc.abstractmethod
    def get_relation_avatar_references(self, relation_id: RelationId) -> AbstractSet[AvatarId]:
        """
        Return a set of all avatars referenced directly or indirectly by the given relation.
        """
        raise NotImplementedError

lib/dl_core/dl_core/components/dependencies/relation_avatar_base.py
datalens-tech/datalens-backend
# FIXME: remove this after full transition to saved field dependencies in favor of RelationAvatarDependencyManager

from __future__ import annotations

from typing import (
    AbstractSet,
    Dict,
)

import attr

from dl_core.components.dependencies.relation_avatar_base import RelationAvatarDependencyManagerBase
from dl_core.components.ids import (
    AvatarId,
    RelationId,
)
from dl_core.query.expression import ExpressionCtx


@attr.s
class LegacyRelationAvatarDependencyManager(RelationAvatarDependencyManagerBase):
    _relation_expressions: Dict[RelationId, ExpressionCtx] = attr.ib(kw_only=True)

    def get_relation_avatar_references(self, relation_id: RelationId) -> AbstractSet[AvatarId]:
        """
        Return a set of all avatars referenced directly or indirectly by the given relation.
        """
        return frozenset(self._relation_expressions[relation_id].avatar_ids)  # type: ignore  # TODO: fix

lib/dl_core/dl_core/components/dependencies/relation_avatar_legacy.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    FrozenSet,
    Iterable,
    Optional,
)

import attr

from dl_constants.enums import (
    DataSourceCreatedVia,
    DataSourceRole,
    DataSourceType,
    JoinType,
    ManagedBy,
)
from dl_core.base_models import (
    DefaultConnectionRef,
    DefaultWhereClause,
    ObligatoryFilter,
    connection_ref_from_id,
)
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.connectors.base.data_source_migration import get_data_source_migrator
from dl_core.data_source.type_mapping import get_data_source_class
from dl_core.data_source_merge_tools import (
    make_spec_from_dict,
    update_spec_from_dict,
)
from dl_core.data_source_spec.collection import (
    DataSourceCollectionSpec,
    DataSourceCollectionSpecBase,
)
from dl_core.data_source_spec.sql import StandardSQLDataSourceSpec
from dl_core.db.elements import (
    IndexInfo,
    SchemaColumn,
)
import dl_core.exc as exc
from dl_core.fields import (
    BIField,
    ResultSchema,
)
from dl_core.multisource import (
    AvatarRelation,
    BinaryCondition,
    SourceAvatar,
)


if TYPE_CHECKING:
    from dl_core.us_connection_base import ConnectionBase
    from dl_core.us_dataset import Dataset


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class DatasetComponentEditor:
    _dataset: Dataset = attr.ib(kw_only=True)
    _ds_accessor: DatasetComponentAccessor = attr.ib(init=False)

    @_ds_accessor.default
    def _make_ds_accessor(self) -> DatasetComponentAccessor:
        return DatasetComponentAccessor(dataset=self._dataset)

    def ensure_data_source_exists(self, source_id: str) -> None:
        self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)  # It checks for existence

    def add_data_source_collection(
        self,
        *,
        source_id: str,
        title: Optional[str],
        managed_by: ManagedBy = ManagedBy.user,
        valid: bool = True,
    ) -> DataSourceCollectionSpecBase:
        """Add a new data source collection configuration entry. Return its ID"""

        dsrc_coll_spec = DataSourceCollectionSpec(
            id=source_id,
            title=title,
            managed_by=managed_by,
            valid=valid,
        )
        self._dataset.data.source_collections.append(dsrc_coll_spec)
        return dsrc_coll_spec

    def update_data_source_collection(
        self, source_id: str, title: Optional[str] = None, valid: Optional[bool] = None
    ) -> None:
        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)

        if title is not None:
            dsrc_coll_spec.title = title

        if valid is not None:
            dsrc_coll_spec.valid = valid

    def remove_data_source_collection(self, source_id: str, ignore_avatars: bool = False) -> None:
        """
        Remove data source collection configuration from the dataset.
        If ``ignore_avatars`` is true, then avatar dependency check is skipped;
        this can be used if the source will be immediately recreated; use with caution.
        """

        if not ignore_avatars:
            bound_avatars = self._ds_accessor.get_avatar_list(source_id=source_id)
            if bound_avatars:
                raise exc.DatasetConfigurationError(
                    "Can't delete source because it is bound by avatars: {}".format([ava.id for ava in bound_avatars])
                )
        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
        assert dsrc_coll_spec is not None
        self._dataset.data.source_collections.remove(dsrc_coll_spec)

    def add_data_source(
        self,
        *,
        source_id: str,
        role: DataSourceRole = DataSourceRole.origin,
        created_from: DataSourceType,
        connection_id: Optional[str] = None,
        title: Optional[str] = None,
        raw_schema: Optional[list[SchemaColumn]] = None,
        index_info_set: Optional[FrozenSet[IndexInfo]] = None,
        managed_by: Optional[ManagedBy] = None,
        parameters: Optional[dict[str, Any]] = None,
    ) -> None:
        """Add a new data source to the dataset"""

        if not created_from:
            raise ValueError("created_from is required")
        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_opt(source_id=source_id)
        if dsrc_coll_spec is None:
            managed_by = managed_by or ManagedBy.user
            assert managed_by is not None
            if role != DataSourceRole.origin:
                raise exc.DatasetConfigurationError(f"No data source collection for ID {source_id}")
            dsrc_coll_spec = self.add_data_source_collection(
                source_id=source_id,
                managed_by=managed_by,
                title=title,
                valid=True,
            )

        assert isinstance(dsrc_coll_spec, DataSourceCollectionSpec)
        dsrc_cls = get_data_source_class(created_from)
        if not dsrc_cls.store_raw_schema:
            raw_schema = None

        connection_ref = connection_ref_from_id(connection_id=connection_id)
        parameters = parameters or {}
        parameters["connection_ref"] = connection_ref
        parameters["raw_schema"] = raw_schema
        parameters["index_info_set"] = index_info_set

        dsrc_spec = make_spec_from_dict(source_type=created_from, data=parameters)

        dsrc_coll_spec.set_for_role(role, dsrc_spec)

    def update_data_source(
        self,
        source_id: str,
        role: Optional[DataSourceRole] = None,
        connection_id: Optional[str] = None,
        created_from: Optional[DataSourceType] = None,
        raw_schema: Optional[list] = None,
        index_info_set: Optional[FrozenSet[IndexInfo]] = None,
        **parameters: Any,
    ) -> None:
        """Update data source config data"""
        if role is None:
            role = DataSourceRole.origin
        assert role is not None

        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
        assert isinstance(dsrc_coll_spec, DataSourceCollectionSpec)
        dsrc_spec = dsrc_coll_spec.get_for_role(role)
        assert dsrc_spec is not None

        if connection_id is not None:
            parameters["connection_ref"] = DefaultConnectionRef(conn_id=connection_id)
        if raw_schema is not None:
            parameters["raw_schema"] = raw_schema
        # If index info was not provided during update we should clear it
        parameters["index_info_set"] = index_info_set

        new_source_type = created_from if created_from is not None else dsrc_spec.source_type
        dsrc_spec = update_spec_from_dict(source_type=new_source_type, data=parameters, old_spec=dsrc_spec)
        dsrc_coll_spec.set_for_role(role=role, value=dsrc_spec)

    def remove_data_source(self, source_id: str, role: DataSourceRole, delete_mat_table: bool = True) -> None:
        """Remove data source configuration from collection"""

        if role == DataSourceRole.origin:
            self.remove_data_source_collection(source_id=source_id)
        else:
            dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)

            if dsrc_coll_spec.managed_by != ManagedBy.user:
                LOGGER.info(f"dsrc_coll_spec.managed_by = {dsrc_coll_spec.managed_by}. Skipping datasource deletion.")
                return

            # We can delete only sources from non-ref collections
            assert isinstance(dsrc_coll_spec, DataSourceCollectionSpec)

            dsrc_spec = dsrc_coll_spec.get_for_role(role=role)
            if dsrc_spec is None:
                return

            # Sample and mat sources must be SQL
            assert isinstance(dsrc_spec, StandardSQLDataSourceSpec)

            dsrc_coll_spec.set_for_role(role, None)

    def _elect_new_root_avatar(self) -> str:
        """Choose among avatars that have no left relations"""
        relations_by_right = defaultdict(list)
        for rel in self._dataset.data.avatar_relations:
            relations_by_right[rel.right_avatar_id].append(rel)
        for avatar in self._dataset.data.source_avatars:
            if not relations_by_right[avatar.id]:
                return avatar.id

        raise RuntimeError("No avatars left to elect from")

    def add_avatar(
        self,
        avatar_id: str,
        source_id: str,
        title: str,
        managed_by: Optional[ManagedBy] = None,
        valid: bool = True,
    ) -> None:
        is_root = not self._dataset.data.source_avatars  # first avatar is always root
        self.ensure_data_source_exists(source_id)
        self._dataset.data.source_avatars.append(
            SourceAvatar(
                id=avatar_id,
                source_id=source_id,
                title=title,
                is_root=is_root,
                managed_by=managed_by or ManagedBy.user,
                valid=valid,
            )
        )

    def update_avatar(
        self,
        avatar_id: str,
        source_id: Optional[str] = None,
        title: Optional[str] = None,
        valid: Optional[bool] = None,
    ) -> None:
        avatar = self._ds_accessor.get_avatar_opt(avatar_id=avatar_id)
        if avatar is not None:
            if source_id is not None:
                self.ensure_data_source_exists(source_id=source_id)
                avatar.source_id = source_id
            if title is not None:
                avatar.title = title
            if valid is not None:
                avatar.valid = valid

    def remove_avatar(self, avatar_id: str) -> None:
        bound_relations = self._ds_accessor.get_avatar_relation_list(
            left_avatar_id=avatar_id
        ) + self._ds_accessor.get_avatar_relation_list(right_avatar_id=avatar_id)
        if bound_relations:
            raise exc.DatasetConfigurationError(
                "Can't delete avatar because it is bound by relations: {}".format([rel.id for rel in bound_relations])
            )

        avatar = self._ds_accessor.get_avatar_strict(avatar_id=avatar_id)
        self._dataset.data.source_avatars.remove(avatar)

        if avatar.is_root and self._dataset.data.source_avatars:
            self.set_root_avatar(avatar_id=self._elect_new_root_avatar(), rebuild_relations=False)

    def set_root_avatar(self, avatar_id: str, rebuild_relations: bool = True) -> None:
        for avatar in self._dataset.data.source_avatars:
            if avatar.id == avatar_id:
                avatar.is_root = True
            elif avatar.is_root:
                avatar.is_root = False

        if rebuild_relations:
            # root must have no left relations
            relations_to_reverse = []
            current_avatar_id = avatar_id
            while True:
                left_relations = self._ds_accessor.get_avatar_relation_list(right_avatar_id=current_avatar_id)
                if not left_relations:
                    # it is already the leftmost avatar
                    break

                if len(left_relations) != 1:
                    raise exc.DatasetConfigurationError("Can't have more than one left relation")

                relations_to_reverse.append(left_relations[0].id)
                current_avatar_id = left_relations[0].left_avatar_id

            for relation_id in reversed(relations_to_reverse):
                self.reverse_relation(relation_id=relation_id)

    def ensure_avatar_exists(self, avatar_id: str) -> None:
        self._ds_accessor.get_avatar_strict(avatar_id)  # It checks for existence

    def add_avatar_relation(
        self,
        relation_id: str,
        left_avatar_id: str,
        right_avatar_id: str,
        conditions: list[BinaryCondition],
        join_type: Optional[JoinType] = None,
        managed_by: Optional[ManagedBy] = None,
        valid: bool = True,
    ) -> None:
        # validate
        if self._ds_accessor.get_avatar_relation_list(right_avatar_id=right_avatar_id):
            raise exc.DatasetConfigurationError("Avatar {} already has a left relation".format(right_avatar_id))
        for avatar_id in (left_avatar_id, right_avatar_id):
            if avatar_id:  # isn't it always true?
                self.ensure_avatar_exists(avatar_id=avatar_id)

        right_avatar = self._ds_accessor.get_avatar_strict(right_avatar_id)
        if right_avatar.is_root:
            self.set_root_avatar(avatar_id=left_avatar_id, rebuild_relations=False)

        self._dataset.data.avatar_relations.append(
            AvatarRelation(
                id=relation_id,
                left_avatar_id=left_avatar_id,
                right_avatar_id=right_avatar_id,
                conditions=conditions[:],
                join_type=join_type or JoinType.inner,
                managed_by=managed_by or ManagedBy.user,
                valid=valid,
            )
        )

    def update_avatar_relation(
        self,
        relation_id: str,
        conditions: Optional[list[BinaryCondition]] = None,
        join_type: Optional[JoinType] = None,
        valid: Optional[bool] = None,
    ) -> None:
        # left/right avatar IDs cannot be changed. Delete/create new relation to do that
        relation = self._ds_accessor.get_avatar_relation_opt(relation_id=relation_id)
        if relation is not None:
            if conditions is not None:
                relation.conditions = conditions[:]
            if join_type is not None:
                relation.join_type = join_type
            if valid is not None:
                relation.valid = valid

    def reverse_relation(self, relation_id: str) -> None:
        """Swap left and right parts of the relation. This might be needed if root avatar is reset"""
        # TODO
        raise RuntimeError("Relation rebuilding is not supported")

    def remove_avatar_relation(self, relation_id: str) -> None:
        found_ind: Optional[int] = None
        for i, relation_config in enumerate(self._dataset.data.avatar_relations):
            if relation_config.id == relation_id:
                found_ind = i
        if found_ind is not None:
            del self._dataset.data.avatar_relations[found_ind]

    def add_obligatory_filter(
        self,
        obfilter_id: str,
        field_guid: str,
        default_filters: list[DefaultWhereClause],
        managed_by: Optional[ManagedBy] = None,
        valid: bool = True,
    ) -> None:
        for filter_object in self._dataset.data.obligatory_filters:
            if filter_object.field_guid == field_guid:
                raise exc.DatasetConfigurationError(f"Obligatory filter for field {field_guid} already exists.")

        filter_object = ObligatoryFilter(
            id=obfilter_id,
            field_guid=field_guid,
            default_filters=default_filters,
            managed_by=managed_by or ManagedBy.user,
            valid=valid,
        )
        self._dataset.data.obligatory_filters.append(filter_object)

    def update_obligatory_filter(
        self,
        obfilter_id: str,
        default_filters: Optional[list[DefaultWhereClause]] = None,
        valid: Optional[bool] = None,
    ) -> None:
        for filter_object in self._dataset.data.obligatory_filters:
            if filter_object.id == obfilter_id:
                if default_filters is not None:
                    filter_object.default_filters = default_filters
                if valid is not None:
                    filter_object.valid = valid
                return

        raise exc.DatasetConfigurationError(f"Obligatory filter with id {obfilter_id} not found.")

    def remove_obligatory_filter(self, obfilter_id: str) -> None:
        new_filters = []
        for filter_object in self._dataset.data.obligatory_filters:
            if filter_object.id != obfilter_id:
                new_filters.append(filter_object)
        self._dataset.data.obligatory_filters = new_filters

    def set_result_schema(self, result_schema: ResultSchema | Iterable[BIField]) -> None:
        if not isinstance(result_schema, ResultSchema):
            result_schema = ResultSchema(fields=list(result_schema))
        self._dataset.data.result_schema = result_schema

    def set_revision_id(self, revision_id: Optional[str]) -> None:
        self._dataset.data.revision_id = revision_id

    def set_created_via(self, created_via: DataSourceCreatedVia) -> None:
        self._dataset.meta["created_via"] = created_via.name

    def replace_connection(self, old_connection: ConnectionBase, new_connection: ConnectionBase) -> None:
        old_migrator = get_data_source_migrator(old_connection.conn_type)
        new_migrator = get_data_source_migrator(new_connection.conn_type)

        role = DataSourceRole.origin
        for source_id in self._ds_accessor.get_data_source_id_list():
            old_source_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
            assert isinstance(old_source_coll_spec, DataSourceCollectionSpec)
            old_source_spec = old_source_coll_spec.get_for_role(role)
            assert old_source_spec is not None
            migration_dtos = old_migrator.export_migration_dtos(data_source_spec=old_source_spec)
            new_connection_ref = new_connection.conn_ref
            assert new_connection_ref is not None
            new_source_spec = new_migrator.import_migration_dtos(
                migration_dtos=migration_dtos, connection_ref=new_connection_ref
            )
            old_source_coll_spec.set_for_role(role=role, value=new_source_spec)

lib/dl_core/dl_core/components/editor.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from enum import (
    Enum,
    unique,
)
import random
import string
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Optional,
    Set,
    Type,
)
import uuid

from anyascii import anyascii
import attr


if TYPE_CHECKING:
    from dl_core.us_dataset import Dataset


ID_LENGTH = 36
ID_VALID_SYMBOLS = string.ascii_lowercase + string.digits + "_-"


FieldId = str
SourceId = str
AvatarId = str
RelationId = str


def generate_random_str(length: int = 4) -> str:
    return "".join(random.choices(string.ascii_lowercase + string.digits, k=length))


def make_field_id() -> FieldId:
    return str(uuid.uuid4())


def make_avatar_id() -> AvatarId:
    return str(uuid.uuid4())


def make_source_id() -> SourceId:
    return str(uuid.uuid4())


def resolve_id_collisions(item: str, existing_items: Set[str], formatter: str = "{} ({})") -> str:
    idx = 1
    orig_item = item
    while item in existing_items:
        item = formatter.format(orig_item, idx)
        idx += 1
    return item


@attr.s
class FieldIdValidator:
    _id_valid_symbols: str = ID_VALID_SYMBOLS
    _id_length: int = ID_LENGTH

    def is_valid(self, field_id: FieldId) -> bool:
        if not len(field_id) or len(field_id) > self._id_length:
            return False
        if any((symbol not in self._id_valid_symbols for symbol in field_id)):
            return False
        return True


@attr.s
class FieldIdGenerator(metaclass=abc.ABCMeta):
    dataset: Dataset = attr.ib(kw_only=True)

    @abc.abstractmethod
    def make_field_id(self, *args: Any, **kwargs: Any) -> FieldId:
        raise NotImplementedError()


@attr.s
class DefaultFieldIdGenerator(FieldIdGenerator):
    def make_field_id(self, *args: Any, **kwargs: Any) -> FieldId:
        return str(uuid.uuid4())


def make_readable_field_id(title: str, valid_symbols: str, max_length: int) -> str:
    field_id = "".join(symbol for symbol in "_".join(anyascii(title).lower().split()) if symbol in valid_symbols)
    field_id = field_id[:max_length]
    return field_id


@attr.s
class ReadableFieldIdGenerator(FieldIdGenerator):
    _id_valid_symbols: str = ID_VALID_SYMBOLS
    _id_length: int = ID_LENGTH
    _id_formatter: str = "{}_{}"

    def make_field_id(self, *args: Any, title: Optional[str] = None, **kwargs: Any) -> FieldId:
        if title is None:
            return str(uuid.uuid4())
        field_id = make_readable_field_id(title, self._id_valid_symbols, self._id_length)
        field_id = self._resolve_id_collisions(field_id)
        return field_id

    def _resolve_id_collisions(self, item: FieldId) -> FieldId:
        existing_items = {f.guid for f in self.dataset.result_schema.fields}
        idx = 1
        orig_item = item
        while item in existing_items:
            item = self._id_formatter.format(orig_item, idx)
            idx += 1
        return item


@attr.s
class ReadableFieldIdGeneratorWithPrefix(ReadableFieldIdGenerator):
    def make_field_id(self, *args: Any, title: Optional[str] = None, **kwargs: Any) -> FieldId:
        field_id = super().make_field_id(title=title)
        return "_".join([generate_random_str(), field_id])


@attr.s
class ReadableFieldIdGeneratorWithSuffix(ReadableFieldIdGenerator):
    def make_field_id(self, *args: Any, title: Optional[str] = None, **kwargs: Any) -> FieldId:
        field_id = super().make_field_id(title=title)
        return "_".join([field_id, generate_random_str()])


@unique
class FieldIdGeneratorType(Enum):
    default = "default"
    readable = "readable"
    prefix = "prefix"
    suffix = "suffix"


DEFAULT_FIELD_ID_GENERATOR_TYPE: FieldIdGeneratorType = FieldIdGeneratorType.readable


FIELD_ID_GENERATOR_MAP: Dict[FieldIdGeneratorType, Type[FieldIdGenerator]] = {
    FieldIdGeneratorType.default: DefaultFieldIdGenerator,
    FieldIdGeneratorType.readable: ReadableFieldIdGenerator,
    FieldIdGeneratorType.prefix: ReadableFieldIdGeneratorWithPrefix,
    FieldIdGeneratorType.suffix: ReadableFieldIdGeneratorWithSuffix,
}

lib/dl_core/dl_core/components/ids.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_core.connection_executors.async_base import AsyncConnExecutorBase
from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor
from dl_core.connection_executors.common_base import (
    ConnExecutorQuery,
    ExecutionMode,
)
from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.connection_executors.sync_executor_wrapper import SyncWrapperForAsyncConnExecutor


__all__ = (
    "ExecutionMode",
    "ConnExecutorQuery",
    "AsyncConnExecutorBase",
    "DefaultSqlAlchemyConnExecutor",
    "SyncConnExecutorBase",
    "SyncWrapperForAsyncConnExecutor",
)

lib/dl_core/dl_core/connection_executors/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    Generator,
    Iterable,
    List,
    Optional,
    Sequence,
    TypeVar,
)

import attr
from typing_extensions import final

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    ExecutionStep,
    ExecutionStepCursorInfo,
    ExecutionStepDataChunk,
    RawSchemaInfo,
)
from dl_core.connection_models import (
    DBIdent,
    SchemaIdent,
    TableDefinition,
    TableIdent,
)


if TYPE_CHECKING:
    from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO  # noqa: F401


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class DBAdapterQueryResult:
    cursor_info: dict = attr.ib()
    data_chunks: Iterable[Sequence] = attr.ib()
    # Notable difference from `cursor_info`: raw value is not meant to be serialized.
    raw_cursor_info: Optional[ExecutionStepCursorInfo] = attr.ib(default=None)

    def get_all(self) -> Sequence[Sequence]:
        """
        :return: All fetched data
        """
        result: List[Sequence] = []
        for chunk in self.data_chunks:
            result.extend(chunk)

        return result


_TARGET_DTO_TV = TypeVar("_TARGET_DTO_TV", bound="ConnTargetDTO")


class SyncDirectDBAdapter(CommonBaseDirectAdapter[_TARGET_DTO_TV], metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def execute_by_steps(self, db_adapter_query: DBAdapterQuery) -> Generator[ExecutionStep, None, None]:
        pass

    @abc.abstractmethod
    def test(self) -> None:
        pass

    # TODO CONSIDER: Make version with simplified interface execute(query: str, db_name: Optional[str] = None)
    @final
    def execute(self, query: DBAdapterQuery) -> DBAdapterQueryResult:
        steps_generator = self.execute_by_steps(query)

        cursor_info_step = next(steps_generator)
        if not isinstance(cursor_info_step, ExecutionStepCursorInfo):
            raise RuntimeError(f"Unexpected type of first step from database: {cursor_info_step}")

        def real_generator() -> Generator[Sequence[Sequence], None, None]:
            for step in steps_generator:
                if not isinstance(step, ExecutionStepDataChunk):
                    raise RuntimeError(f"Unexpected type of non-first step from database: {step}")
                yield step.chunk

        return DBAdapterQueryResult(
            cursor_info=cursor_info_step.cursor_info,
            data_chunks=real_generator(),
            raw_cursor_info=cursor_info_step,
        )

    @abc.abstractmethod
    def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        pass

    @abc.abstractmethod
    def get_schema_names(self, db_ident: DBIdent) -> List[str]:
        pass

    @abc.abstractmethod
    def get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        pass

    @abc.abstractmethod
    def get_table_info(self, table_def: TableDefinition, fetch_idx_info: bool) -> RawSchemaInfo:
        pass

    @abc.abstractmethod
    def is_table_exists(self, table_ident: TableIdent) -> bool:
        pass

    @abc.abstractmethod
    def close(self) -> None:
        pass

lib/dl_core/dl_core/connection_executors/adapters/adapters_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import contextlib
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Dict,
    Generator,
    List,
    Optional,
    Tuple,
    Type,
    TypeVar,
    Union,
)

import attr
import sqlalchemy as sa
from sqlalchemy import event
from sqlalchemy.engine.base import Engine
from typing_extensions import final

from dl_app_tools.profiling_base import (
    GenericProfiler,
    generic_profiler,
)
from dl_core import exc
from dl_core.connection_executors.adapters.adapters_base import SyncDirectDBAdapter
from dl_core.connection_executors.adapters.mixins import (
    SAColumnTypeNormalizer,
    WithCursorInfo,
    WithDatabaseNameOverride,
    WithNoneRowConverters,
)
from dl_core.connection_executors.adapters.sa_utils import (
    CursorLogger,
    compile_query_for_debug,
    get_db_version_query,
)
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    ExecutionStep,
    ExecutionStepCursorInfo,
    ExecutionStepDataChunk,
    RawColumnInfo,
    RawIndexInfo,
    RawSchemaInfo,
)
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
from dl_core.connection_models import (
    DBIdent,
    SATextTableDefinition,
    SchemaIdent,
    TableDefinition,
    TableIdent,
)
from dl_core.connectors.base.error_handling import ETBasedExceptionMaker
from dl_core.connectors.base.error_transformer import (
    DbErrorTransformer,
    make_default_transformer_with_custom_rules,
)
from dl_core.db.native_type import CommonNativeType
from dl_core.db_session_utils import db_session_context
from dl_utils.utils import get_type_full_name


if TYPE_CHECKING:
    from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO  # noqa: F401


LOGGER = logging.getLogger(__name__)

_DBA_SA_TV = TypeVar("_DBA_SA_TV", bound="BaseSAAdapter")
_DBA_SA_DTO_TV = TypeVar("_DBA_SA_DTO_TV", bound="ConnTargetDTO")


@attr.s()
class BaseSAAdapter(
    WithCursorInfo,
    WithDatabaseNameOverride,
    WithNoneRowConverters,
    ETBasedExceptionMaker,
    SyncDirectDBAdapter[_DBA_SA_DTO_TV],
    SAColumnTypeNormalizer,
):
    allow_sa_text_as_columns_source: ClassVar[bool] = False

    # Instance attributes
    _default_chunk_size: int = attr.ib()
    _target_dto: _DBA_SA_DTO_TV = attr.ib()
    _req_ctx_info: DBAdapterScopedRCI = attr.ib()
    # Internals
    _engine_cache: Dict[Tuple[Optional[str], bool], Engine] = attr.ib(init=False, factory=lambda: {})
    _error_transformer: ClassVar[DbErrorTransformer] = make_default_transformer_with_custom_rules()

    @classmethod
    def create(
        cls: Type[_DBA_SA_TV], target_dto: _DBA_SA_DTO_TV, req_ctx_info: DBAdapterScopedRCI, default_chunk_size: int
    ) -> _DBA_SA_TV:
        return cls(target_dto=target_dto, default_chunk_size=default_chunk_size, req_ctx_info=req_ctx_info)

    @property
    def conn_id(self) -> Optional[str]:
        return self._target_dto.conn_id

    @property
    def default_chunk_size(self) -> int:
        return self._default_chunk_size

    @final
    def get_db_engine(self, db_name: Optional[str], disable_streaming: bool = False) -> Engine:
        effective_db_name = self.get_db_name_for_query(db_name_from_query=db_name)

        cache_key: Tuple[Optional[str], bool] = effective_db_name, disable_streaming

        if cache_key not in self._engine_cache:
            new_engine = self._get_db_engine(db_name=effective_db_name, disable_streaming=disable_streaming)
            # Adding event listeners
            # TODO CONSIDER: May be add adapters by flag?
            event.listen(new_engine, "before_cursor_execute", CursorLogger.before_cursor_execute_handler)
            event.listen(new_engine, "after_cursor_execute", CursorLogger.after_cursor_execute_handler)

            self._engine_cache[cache_key] = new_engine

        return self._engine_cache[cache_key]

    @abc.abstractmethod
    def _get_db_engine(self, db_name: str, disable_streaming: bool = False) -> Engine:
        pass

    @contextlib.contextmanager
    def execution_context(self) -> Generator[None, None, None]:
        yield

    @generic_profiler("db-full")
    def execute_by_steps(self, db_adapter_query: DBAdapterQuery) -> Generator[ExecutionStep, None, None]:
        """Generator that yielding messages with data chunks and execution meta-info"""
        chunk_size = db_adapter_query.get_effective_chunk_size(self.default_chunk_size)
        query = db_adapter_query.query

        engine = self.get_db_engine(
            db_name=db_adapter_query.db_name,
            disable_streaming=db_adapter_query.disable_streaming,
        )

        # TODO FIX: Delegate query compilation for debug to error handler or make method of debug compilation
        compiled_query = query if isinstance(query, str) else compile_query_for_debug(query, engine.dialect)

        with db_session_context(
            backend_type=self.get_backend_type(), db_engine=engine
        ) as db_session, self.handle_execution_error(compiled_query), self.execution_context():
            with GenericProfiler("db-exec"):
                result = db_session.execute(
                    query,
                    # *args,
                    # **kwargs,
                )

            cursor_info = ExecutionStepCursorInfo(
                cursor_info=self._make_cursor_info(result.cursor, db_session=db_session),
                raw_cursor_description=list(result.cursor.description),
                raw_engine=engine,
            )
            yield cursor_info

            row_converters = self._get_row_converters(cursor_info=cursor_info)
            while True:
                LOGGER.info("Fetching %s rows (conn %s)", chunk_size, self.conn_id)
                with GenericProfiler("db-fetch"):
                    rows = result.fetchmany(chunk_size)

                if not rows:
                    LOGGER.info("No rows remaining")
                    break

                LOGGER.info("Rows fetched, yielding")
                yield ExecutionStepDataChunk(
                    tuple(
                        tuple(
                            (col_converter(val) if col_converter is not None and val is not None else val)
                            # for val, col_converter in zip_longest(row, row_converters)
                            for val, col_converter in zip(row, row_converters)
                        )
                        for row in rows
                    )
                )

    @final
    def test(self) -> None:
        with self.handle_execution_error(debug_compiled_query="<test()>"):
            self._test()

    @final
    def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        with self.handle_execution_error(debug_compiled_query=f"<get_db_version({db_ident})>"):
            return self._get_db_version(db_ident)

    @final
    def get_schema_names(self, db_ident: DBIdent) -> List[str]:
        with self.handle_execution_error(
            debug_compiled_query=f"<get_schema_names({db_ident})>"
        ), self.execution_context():
            return self._get_schema_names(db_ident)

    @final
    def get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        with self.handle_execution_error(
            debug_compiled_query=f"<get_tables({schema_ident})>"
        ), self.execution_context():
            return self._get_tables(schema_ident)

    @final
    def get_table_info(self, table_def: TableDefinition, fetch_idx_info: bool) -> RawSchemaInfo:
        debug_compiled_query: str
        exc_post_processor: Optional[Callable[[exc.DatabaseQueryError], None]] = None

        if isinstance(table_def, TableIdent):
            debug_compiled_query = f"<get_columns({table_def})>"

        elif isinstance(table_def, SATextTableDefinition):
            debug_compiled_query = "<get_columns(<subselect>)>"

            # TODO FIX: Do not fill details at such low level
            #  Better way is to set flag that this was a user subselect and depending on user permissions
            #  fill details at bi-api level
            # A catch-point intended for adding `query` and `db_message` only for subselect-schema errors.
            def exc_post_processor(db_exc: exc.DatabaseQueryError) -> None:
                assert isinstance(db_exc, exc.DatabaseQueryError)
                db_exc.query = "SELECT * FROM {}".format(table_def.text)  # type: ignore  # TODO: fix
                db_exc.details.setdefault("query", db_exc.query)
                db_exc.details.setdefault("db_message", db_exc.db_message)

        else:
            raise TypeError(
                f"Unsupported type of table table definition to get info: {get_type_full_name(type(table_def))}"
            )

        with self.handle_execution_error(
            debug_compiled_query=debug_compiled_query,
            exc_post_processor=exc_post_processor,
        ), self.execution_context():
            if isinstance(table_def, TableIdent):
                columns = self._get_raw_columns_info(table_def)
                indexes: Optional[Tuple[RawIndexInfo, ...]] = None

                if fetch_idx_info and isinstance(table_def, TableIdent):
                    try:
                        with GenericProfiler("db-idx-fetch"):
                            indexes = self._get_table_indexes(table_def)

                    except Exception:  # noqa
                        LOGGER.exception("Indexes fetching fail for %s", table_def)

                return RawSchemaInfo(columns=columns, indexes=indexes)

            elif isinstance(table_def, SATextTableDefinition):
                return self._get_subselect_table_info(table_def)

            else:
                raise TypeError(
                    f"Unsupported type of table table definition to get info: {get_type_full_name(type(table_def))}"
                )

        raise AssertionError("Non handled case of get_table_info()")

    @final
    def is_table_exists(self, table_ident: TableIdent) -> bool:
        with self.handle_execution_error(f"<exists({table_ident})>"), self.execution_context():
            return self._is_table_exists(table_ident)

    def _test(self) -> None:
        self.execute(DBAdapterQuery("select 1")).get_all()

    def _get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        """Return database version string for this connection"""
        return self.execute(get_db_version_query(db_ident)).get_all()[0][0]

    def _get_schema_names(self, db_ident: DBIdent) -> List[str]:
        db_engine = self.get_db_engine(db_ident.db_name)
        return sa.inspect(db_engine).get_schema_names()

    def _get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        db_name = schema_ident.db_name
        schema_name = schema_ident.schema_name

        db_engine = self.get_db_engine(db_name)
        table_list = sa.inspect(db_engine).get_table_names(schema=schema_name)
        view_list = sa.inspect(db_engine).get_view_names(schema=schema_name)
        return [
            TableIdent(
                db_name=db_name,
                schema_name=schema_name,
                table_name=name,
            )
            for name in table_list + view_list
        ]

    def _get_sa_table_columns(self, table_def: TableDefinition) -> List[Dict[str, Any]]:
        db_engine: sa.engine.Engine
        columns_source: Union[str, sa.sql.elements.TextClause]
        if isinstance(table_def, TableIdent):
            columns_source = table_def.table_name
            db_schema = table_def.schema_name
            db_engine = self.get_db_engine(table_def.db_name)
        elif isinstance(table_def, SATextTableDefinition) and self.allow_sa_text_as_columns_source:
            columns_source = table_def.text
            db_schema = None
            db_engine = self.get_db_engine(None)
        else:
            raise AssertionError(
                f"{get_type_full_name(type(self))}._get_sa_table_columns() does not support"
                f" table_def {get_type_full_name(type(table_def))}"
            )

        inspection = sa.inspect(db_engine)
        table_columns = inspection.get_columns(columns_source, schema=db_schema)
        table_columns = list(table_columns)
        return table_columns

    def _get_raw_columns_info(self, table_def: TableDefinition) -> Tuple[RawColumnInfo, ...]:
        table_columns = self._get_sa_table_columns(table_def)

        return tuple(
            RawColumnInfo(
                name=column["name"],
                title=column.get("title"),
                nullable=column.get("nullable", True),  # To be deprecated.
                native_type=CommonNativeType.normalize_name_and_create(
                    conn_type=self.conn_type,
                    name=self.normalize_sa_col_type(column["type"]),
                    nullable=column.get("nullable", True),
                ),
            )
            for column in table_columns
        )

    def _get_table_indexes(self, table_ident: TableIdent) -> Tuple[RawIndexInfo, ...]:
        db_engine = self.get_db_engine(table_ident.db_name)
        inspection = sa.inspect(db_engine)

        idx_list: List[Dict[str, Any]] = inspection.get_indexes(table_ident.table_name, schema=table_ident.schema_name)

        return tuple(
            RawIndexInfo(
                columns=tuple(idx["column_names"]),
                unique=idx["unique"],
                kind=None,
            )
            for idx in idx_list
        )

    def _get_subselect_table_info(self, subquery: SATextTableDefinition) -> RawSchemaInfo:
        raise Exception(f"`_get_subselect_table_info` is not implemented on SA adapter {self!r}")

    def _is_table_exists(self, table_ident: TableIdent) -> bool:
        eng = self.get_db_engine(table_ident.db_name)
        return sa.inspect(eng).has_table(table_ident.table_name, schema=table_ident.schema_name)

    def close(self) -> None:
        LOGGER.info("Method .close() of SA DB adapter was called")
        # TODO FIX: Find all created generators and close it
        for eng in self._engine_cache.values():
            eng.dispose()

lib/dl_core/dl_core/connection_executors/adapters/adapters_base_sa.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Generic,
    Optional,
    Type,
    TypeVar,
)
from urllib.parse import (
    quote_plus,
    urlencode,
)

import attr
import sqlalchemy as sa
from sqlalchemy.engine.base import Engine

from dl_core.connection_executors.adapters.adapters_base_sa import BaseSAAdapter
from dl_core.connection_executors.adapters.common_base import get_dialect_string
from dl_core.connection_executors.adapters.mixins import WithMinimalCursorInfo
from dl_core.connection_executors.models.connection_target_dto_base import (
    BaseSQLConnTargetDTO,
    ConnTargetDTO,
)
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    ExecutionStepCursorInfo,
    RawColumnInfo,
    RawSchemaInfo,
)
from dl_core.db.native_type import CommonNativeType


if TYPE_CHECKING:
    from dl_core.connection_models import SATextTableDefinition

LOGGER = logging.getLogger(__name__)


_CONN_DTO_TV = TypeVar("_CONN_DTO_TV", bound=ConnTargetDTO)


@attr.s(cmp=False)
class BaseConnLineConstructor(abc.ABC, Generic[_CONN_DTO_TV]):
    _target_dto: _CONN_DTO_TV = attr.ib()
    _dsn_template: str = attr.ib()
    _dialect_name: str = attr.ib()

    @abc.abstractmethod
    def _get_dsn_params(
        self,
        safe_db_symbols: tuple[str, ...] = (),
        db_name: Optional[str] = None,
        standard_auth: Optional[bool] = True,
    ) -> dict:
        raise NotImplementedError

    def _get_dsn_query_params(self) -> dict:
        return {}

    def make_conn_line(self, db_name: Optional[str] = None, params: Optional[dict[str, Any]] = None) -> str:
        # TODO?: replace with `get_conn_line`, delete `dsn_template`.
        conn_line = self._dsn_template.format(**self._get_dsn_params(db_name=db_name))

        params = (params or {}).copy()
        params.update(self._get_dsn_query_params())

        if params:
            params = {key: val for key, val in params.items() if val is not None}
            conn_line += "?" + urlencode(params)

        return conn_line


_DBA_CLASSIC_SA_DTO_TV = TypeVar("_DBA_CLASSIC_SA_DTO_TV", bound="BaseSQLConnTargetDTO")


@attr.s(cmp=False)
class ClassicSQLConnLineConstructor(
    BaseConnLineConstructor[_DBA_CLASSIC_SA_DTO_TV],
    Generic[_DBA_CLASSIC_SA_DTO_TV],
):
    def _get_dsn_params(
        self,
        safe_db_symbols: tuple[str, ...] = (),
        db_name: Optional[str] = None,
        standard_auth: Optional[bool] = True,
    ) -> dict:
        return dict(
            dialect=self._dialect_name,
            user=quote_plus(self._target_dto.username) if standard_auth else None,
            passwd=quote_plus(self._target_dto.password) if standard_auth else None,
            host=quote_plus(self._target_dto.host),
            port=quote_plus(str(self._target_dto.port)),
            db_name=db_name or quote_plus(self._target_dto.db_name or "", safe="".join(safe_db_symbols)),
        )


@attr.s(cmp=False)
class BaseClassicAdapter(WithMinimalCursorInfo, BaseSAAdapter[_CONN_DTO_TV]):
    dsn_template: ClassVar[str] = "{dialect}://{user}:{passwd}@{host}:{port}/{db_name}"
    execution_options: ClassVar[dict[str, Any]] = {}
    conn_line_constructor_type: ClassVar[Type[BaseConnLineConstructor]] = ClassicSQLConnLineConstructor

    # Instance attributes
    _target_dto: _CONN_DTO_TV = attr.ib()

    def get_connect_args(self) -> dict:
        return {}

    def get_engine_kwargs(self) -> dict:
        return {}

    def get_default_db_name(self) -> Optional[str]:
        if isinstance(self._target_dto, BaseSQLConnTargetDTO):
            return self._target_dto.db_name
        raise NotImplementedError

    def _get_dsn_query_params(self) -> dict:
        return {}

    def get_conn_line(self, db_name: Optional[str] = None, params: Optional[dict[str, Any]] = None) -> str:
        return self.conn_line_constructor_type(
            dsn_template=self.dsn_template,
            dialect_name=get_dialect_string(self.conn_type),
            target_dto=self._target_dto,
        ).make_conn_line(db_name=db_name, params=params)

    def _get_db_engine(self, db_name: str, disable_streaming: bool = False) -> Engine:
        conn_line = self.get_conn_line(db_name=db_name)

        execution_options = self.execution_options
        if disable_streaming:
            execution_options = execution_options.copy()
            execution_options.pop("stream_results", None)

        engine = sa.create_engine(
            conn_line,
            connect_args=self.get_connect_args(),
            execution_options=execution_options,
            # Apparently, this is the only way to pass something to the dialect's `__init__`.
            **self.get_engine_kwargs(),
        )
        engine = engine.execution_options(compiled_cache=None)
        return engine

    _subselect_cursor_info_where_false: ClassVar[bool] = True

    def _get_subselect_raw_cursor_info_and_data(
        self,
        subselect: sa.sql.elements.TextClause,
        limit: Optional[int] = 1,
        where_false: Optional[bool] = None,
    ) -> dict:
        """
        Run a `select * limit 1` query, return cursor info.

        Primarily meant for getting schema of subselects.

        NOTE: `subselect` must be already aliased (in most cases).
        """
        sa_query = sa.select([sa.literal_column("*")]).select_from(subselect)
        if where_false is None:
            where_false = self._subselect_cursor_info_where_false
        if where_false:
            sa_query = sa_query.where(False)
        if limit is not None:
            sa_query = sa_query.limit(limit)
        dba_query = DBAdapterQuery(query=sa_query)
        query_res = self.execute(dba_query)

        assert query_res.raw_cursor_info
        data = list(query_res.data_chunks)
        return query_res.raw_cursor_info, data  # type: ignore  # TODO: fix

    def _get_subselect_table_info(self, subquery: SATextTableDefinition) -> RawSchemaInfo:
        """Will not work without non-empty `self._type_code_to_sa`"""
        raw_cursor_info, _ = self._get_subselect_raw_cursor_info_and_data(subquery.text)
        return self._raw_cursor_info_to_schema(raw_cursor_info)

    def _raw_cursor_info_to_schema(self, raw_cursor_info: ExecutionStepCursorInfo) -> RawSchemaInfo:
        """
        See also: `_get_raw_columns_info`, `_get_sa_table_columns`,
        `create_schema_info_from_raw_schema_info`.
        """
        description = raw_cursor_info.raw_cursor_description
        dialect = None
        if raw_cursor_info.raw_engine is not None:
            dialect = raw_cursor_info.raw_engine.dialect

        columns = []
        for cursor_col in description:
            name = self._cursor_column_to_name(cursor_col, dialect=dialect)
            native_type = self._cursor_column_to_native_type(cursor_col, require=False)
            if native_type is None:
                native_type = CommonNativeType.normalize_name_and_create(
                    conn_type=self.conn_type,
                    name=self.normalize_sa_col_type(sa.sql.sqltypes.NullType),  # type: ignore  # TODO: fix
                    nullable=True,
                )
            columns.append(
                RawColumnInfo(
                    name=name,
                    title=name,
                    nullable=native_type.nullable,
                    native_type=native_type,
                )
            )

        return RawSchemaInfo(columns=tuple(columns))

lib/dl_core/dl_core/connection_executors/adapters/adapters_base_sa_classic.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
import contextlib
import ssl
from typing import (
    TYPE_CHECKING,
    Dict,
    Generator,
    Optional,
    Sequence,
    Tuple,
    Type,
    TypeVar,
)

from aiohttp.client import (
    ClientSession,
    ClientTimeout,
)
import aiohttp.client_exceptions
import attr

from dl_configs.utils import get_root_certificates_path
from dl_core import exc
from dl_core.connection_executors.adapters.async_adapters_base import AsyncDirectDBAdapter


if TYPE_CHECKING:
    from aiohttp import BasicAuth

    from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO
    from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
    from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI


_DBA_TV = TypeVar("_DBA_TV", bound="AiohttpDBAdapter")


@attr.s
class AiohttpDBAdapter(AsyncDirectDBAdapter, metaclass=abc.ABCMeta):
    """Common base for adapters that primarily use an aiohttp client"""

    # TODO?: commonize some of the CTDTO attributes such as connect_timeout?
    _target_dto: ConnTargetDTO = attr.ib()
    _req_ctx_info: DBAdapterScopedRCI = attr.ib()

    _http_read_chunk_size: int = attr.ib(init=False, default=(1024 * 64))
    _default_chunk_size: int = attr.ib()

    _session: ClientSession = attr.ib(init=False, default=None)

    def __attrs_post_init__(self) -> None:
        self._session = ClientSession(
            timeout=self.get_session_timeout(),
            auth=self.get_session_auth(),
            headers=self.get_session_headers(),
            connector=self.create_aiohttp_connector(
                ssl_context=ssl.create_default_context(cafile=get_root_certificates_path())
            ),
        )

    def create_aiohttp_connector(self, ssl_context: Optional[ssl.SSLContext]) -> aiohttp.TCPConnector:
        return aiohttp.TCPConnector(
            ssl_context=ssl_context,
        )

    @classmethod
    def create(
        cls: Type[_DBA_TV],
        target_dto: ConnTargetDTO,
        req_ctx_info: DBAdapterScopedRCI,
        default_chunk_size: int,
    ) -> _DBA_TV:
        return cls(
            target_dto=target_dto,
            req_ctx_info=req_ctx_info,
            default_chunk_size=default_chunk_size,
        )

    def get_session_timeout(self) -> ClientTimeout:
        # Sime default large-ish timeout.
        return ClientTimeout(total=3600.0)

    def get_session_auth(self) -> Optional[BasicAuth]:
        return None

    def get_session_headers(self) -> Dict[str, str]:
        return {
            # TODO: bi_constants / dl_configs.constants
            "user-agent": "DataLens",
        }

    async def close(self) -> None:
        await self._session.close()

    execute_err_map: Sequence[Tuple[Type[Exception], Type[exc.DatabaseQueryError]]] = (
        (aiohttp.client_exceptions.ClientConnectorError, exc.SourceConnectError),
        (exc.AIOHttpConnTimeoutError, exc.SourceConnectError),
        (aiohttp.client_exceptions.ServerTimeoutError, exc.SourceTimeout),
        (asyncio.exceptions.TimeoutError, exc.SourceTimeout),
        (aiohttp.client_exceptions.ServerDisconnectedError, exc.SourceClosedPrematurely),
        (exc.RSTError, exc.SourceClosedPrematurely),
        (aiohttp.client_exceptions.ClientPayloadError, exc.SourceProtocolError),
    )

    @contextlib.contextmanager
    def wrap_execute_excs(self, query: DBAdapterQuery, stage: Optional[str]) -> Generator[None, None, None]:
        try:
            try:
                yield None
            except aiohttp.client_exceptions.ClientOSError as pre_err:
                # need to check errno, so hard to plug into the main exc handler nicely.
                if pre_err.errno == 104:  # 104 == errno.ECONNRESET
                    raise exc.RSTError() from pre_err
                raise
            except aiohttp.client_exceptions.ServerTimeoutError as aiohttp_timeout_err:
                if aiohttp_timeout_err.args[0].startswith("Connection timeout"):
                    # https://github.com/aio-libs/aiohttp/blob/3.8/aiohttp/client.py#L541
                    raise exc.AIOHttpConnTimeoutError() from aiohttp_timeout_err
                raise
        except Exception as err:
            for orig_err_type, mapped_err_type in self.execute_err_map:
                if isinstance(err, orig_err_type):
                    # TODO CONSIDER: Add messages for users' DBs (not internal materialization CH)
                    debug_compiled_query = None
                    if self._target_dto.pass_db_query_to_user:
                        debug_compiled_query = query.debug_compiled_query

                    new_err = mapped_err_type(
                        db_message=None,
                        query=debug_compiled_query,
                        orig=err,
                        details={},
                    )
                    new_err.debug_info.update(stage=stage)
                    raise new_err from err
            raise

lib/dl_core/dl_core/connection_executors/adapters/async_adapters_aiohttp.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
import logging
from typing import (
    TYPE_CHECKING,
    AsyncIterable,
    Awaitable,
    Callable,
    Generic,
    Optional,
    TypeVar,
)

import attr

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    RawSchemaInfo,
)


if TYPE_CHECKING:
    from dl_constants.types import (
        TBIChunksGen,
        TBIDataRow,
    )
    from dl_core.connection_models.common_models import (
        DBIdent,
        SchemaIdent,
        TableDefinition,
        TableIdent,
    )


LOGGER = logging.getLogger(__name__)


@attr.s
class AsyncRawExecutionResult:
    raw_cursor_info: dict = attr.ib()
    raw_chunk_generator: TBIChunksGen = attr.ib()

    async def get_all_rows(self) -> AsyncIterable[TBIDataRow]:
        async for chunk in self.raw_chunk_generator:
            for row in chunk:
                yield row


_CACHE_TV = TypeVar("_CACHE_TV")


@attr.s
class AsyncCache(Generic[_CACHE_TV]):
    _cache: dict[str, _CACHE_TV] = attr.ib(init=False, default=attr.Factory(dict))
    _lock: asyncio.Lock = attr.ib(init=False, default=attr.Factory(asyncio.Lock))

    async def get(self, key: str, generator: Callable[[str], Awaitable[_CACHE_TV]]) -> _CACHE_TV:
        async with self._lock:
            if key not in self._cache:
                self._cache[key] = await generator(key)
            return self._cache[key]

    async def clear(self, finalizer: Callable[[_CACHE_TV], Awaitable[None]]) -> None:
        async with self._lock:
            for elem in self._cache.values():
                await finalizer(elem)
        self._cache.clear()


_DBA_TV = TypeVar("_DBA_TV", bound="AsyncDBAdapter")


class AsyncDBAdapter(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    async def test(self) -> None:
        pass

    @abc.abstractmethod
    async def execute(self, query: DBAdapterQuery) -> AsyncRawExecutionResult:
        pass

    @abc.abstractmethod
    async def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        pass

    @abc.abstractmethod
    async def get_schema_names(self, db_ident: DBIdent) -> list[str]:
        pass

    @abc.abstractmethod
    async def get_tables(self, schema_ident: SchemaIdent) -> list[TableIdent]:
        pass

    @abc.abstractmethod
    async def get_table_info(self, table_def: TableDefinition, fetch_idx_info: bool) -> RawSchemaInfo:
        pass

    @abc.abstractmethod
    async def is_table_exists(self, table_ident: TableIdent) -> bool:
        pass

    @abc.abstractmethod
    async def close(self) -> None:
        pass

    async def __aenter__(self: _DBA_TV) -> _DBA_TV:
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):  # type: ignore
        await self.close()


class AsyncDirectDBAdapter(AsyncDBAdapter, CommonBaseDirectAdapter, metaclass=abc.ABCMeta):
    pass

lib/dl_core/dl_core/connection_executors/adapters/async_adapters_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import json
import logging
import pickle
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncGenerator,
    List,
    Optional,
    Sequence,
    Tuple,
    Type,
    TypeVar,
)

import aiohttp
from aiohttp import (
    ClientTimeout,
    TCPConnector,
)
from aiohttp.client_exceptions import ServerTimeoutError
import attr

from dl_api_commons.headers import (
    HEADER_LOGGING_CONTEXT,
    INTERNAL_HEADER_PROFILING_STACK,
)
from dl_api_commons.tracing import get_current_tracing_headers
from dl_app_tools import log
from dl_app_tools.profiling_base import (
    GenericProfiler,
    generic_profiler_async,
)
from dl_core import exc as common_exc
from dl_core import utils
from dl_core.connection_executors.adapters.adapters_base import SyncDirectDBAdapter
from dl_core.connection_executors.adapters.async_adapters_base import (
    AsyncDBAdapter,
    AsyncDirectDBAdapter,
    AsyncRawExecutionResult,
)
from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.models.common import RemoteQueryExecutorData
from dl_core.connection_executors.models.constants import (
    HEADER_BODY_SIGNATURE,
    HEADER_REQUEST_ID,
)
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    RawSchemaInfo,
)
from dl_core.connection_executors.models.exc import QueryExecutorException
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
from dl_core.connection_executors.qe_serializer import (
    ActionSerializer,
    ResponseTypes,
)
from dl_core.connection_executors.qe_serializer import dba_actions as dba_actions
from dl_core.connection_executors.remote_query_executor.crypto import get_hmac_hex_digest
from dl_core.connection_models.conn_options import ConnectOptions
from dl_core.enums import RQEEventType


if TYPE_CHECKING:
    from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO
    from dl_core.connection_models.common_models import (
        DBIdent,
        SchemaIdent,
        TableDefinition,
        TableIdent,
    )


LOGGER = logging.getLogger(__name__)

_RESP_TV = TypeVar("_RESP_TV", bound=ResponseTypes)


@attr.s()
class RemoteAsyncAdapter(AsyncDBAdapter):
    _target_dto: ConnTargetDTO = attr.ib()
    _req_ctx_info: DBAdapterScopedRCI = attr.ib()
    _dba_cls: Type[CommonBaseDirectAdapter] = attr.ib()
    _rqe_data: RemoteQueryExecutorData = attr.ib()
    _conn_options: ConnectOptions = attr.ib(factory=ConnectOptions)
    _session: aiohttp.ClientSession = attr.ib(init=False)
    _force_async_rqe: bool = attr.ib(default=False)
    _serializer: ActionSerializer = attr.ib(init=False, factory=ActionSerializer)
    _use_sync_rqe: bool = attr.ib(init=False, default=None)

    DEFAULT_REL_PATH = "/execute_action"

    def __attrs_post_init__(self) -> None:
        if self._force_async_rqe:
            self._use_sync_rqe = False
        else:
            if issubclass(self._dba_cls, AsyncDirectDBAdapter):
                self._use_sync_rqe = False
            elif issubclass(self._dba_cls, SyncDirectDBAdapter):
                self._use_sync_rqe = True
            else:
                raise QueryExecutorException(f"Unexpected DBA class: {self._dba_cls}")

        self._session = aiohttp.ClientSession(
            # TODO CONSIDER: May be turn on keepalive on sync QE?
            # uWSGI at this moment fails on second request with keepalive=True
            connector=TCPConnector(force_close=True),
            timeout=ClientTimeout(
                total=(5 * 60 if (override := self._conn_options.rqe_total_timeout) is None else override),
                connect=5,
                sock_read=(600 if (override := self._conn_options.rqe_sock_read_timeout) is None else override),
                sock_connect=5,
            ),
        )

    # TODO FIX: Handle exceptions
    async def _make_request(
        self,
        req_obj: dba_actions.RemoteDBAdapterAction,
        rel_path: Optional[str] = None,
    ) -> aiohttp.ClientResponse:
        if rel_path is None:
            rel_path = self.DEFAULT_REL_PATH

        logging_context = log.context.get_log_context()
        serialized_context = json.dumps(logging_context)

        qe = self._rqe_data
        body_bytes = json.dumps(self._serializer.serialize_action(req_obj)).encode()
        signature = get_hmac_hex_digest(body_bytes, qe.hmac_key)

        tracing_headers = get_current_tracing_headers()

        headers = {
            HEADER_BODY_SIGNATURE: signature,
            HEADER_LOGGING_CONTEXT: serialized_context,
            INTERNAL_HEADER_PROFILING_STACK: GenericProfiler.get_current_stages_stack_str(),
            "Content-Type": "application/json",
            **tracing_headers,
        }

        if self._req_ctx_info.request_id:
            headers[HEADER_REQUEST_ID] = self._req_ctx_info.request_id

        url: str
        if self._use_sync_rqe:
            url = utils.make_url(
                protocol=qe.sync_protocol,
                host=qe.sync_host,
                port=qe.sync_port,
                path=rel_path,
            )
        else:
            url = utils.make_url(
                protocol=qe.async_protocol,
                host=qe.async_host,
                port=qe.async_port,
                path=rel_path,
            )

        LOGGER.info("Making request to QE: %s", url)
        try:
            resp = await self._session.post(
                url=url,
                headers=headers,
                data=body_bytes,
            )
        except ServerTimeoutError as err:
            # we want to distinguish between RQE and source timeouts; unfortunately,
            # aiohttp doesn't provide separate ConnectionTimeout and RequestTimeout
            # exceptions like requests does, so we should check it manually
            if str(err).startswith("Connection timeout"):
                raise

            query = (
                req_obj.db_adapter_query.debug_compiled_query
                if isinstance(req_obj, dba_actions.ActionExecuteQuery)
                else ""
            )
            raise common_exc.SourceTimeout(db_message="Source timed out", query=query) from err
        return resp

    @staticmethod
    async def _read_body_json(resp: aiohttp.ClientResponse) -> dict:
        try:
            body = await resp.read()
            try:
                return json.loads(body)
            except Exception as json_parsing_exc:
                raise QueryExecutorException("Response is not a valid JSON") from json_parsing_exc
        except Exception as unexpected_exc:
            raise QueryExecutorException("Unexpected exception during QE response parsing") from unexpected_exc
        finally:
            resp.release()

    def _parse_exception(self, body_json: dict) -> Exception:
        try:
            return self._serializer.deserialize_exc(body_json)
        except Exception as exc_deserialization_exc:
            raise QueryExecutorException("Unexpected error JSON schema") from exc_deserialization_exc

    async def _make_request_parse_response(
        self,
        req_obj: dba_actions.NonStreamAction[_RESP_TV],
        rel_path: Optional[str] = None,
    ) -> _RESP_TV:
        resp = await self._make_request(
            rel_path=rel_path,
            req_obj=req_obj,
        )

        async with resp:
            body_json = await self._read_body_json(resp)

        if resp.status != 200:
            exc = self._parse_exception(body_json)
            raise exc

        try:
            return req_obj.deserialize_response(body_json)
        except Exception as resp_deserialization_exc:
            raise QueryExecutorException("Unexpected response JSON schema") from resp_deserialization_exc

    # Here we have some async problem:
    #  "qe" stage will be finished before some nested stages in RQE (e.g. "qe/fetch")
    @generic_profiler_async("qe")  # type: ignore  # TODO: fix
    async def execute(self, query: DBAdapterQuery) -> AsyncRawExecutionResult:
        resp = await self._make_request(
            dba_actions.ActionExecuteQuery(
                db_adapter_query=query,
                target_conn_dto=self._target_dto,
                dba_cls=self._dba_cls,
                req_ctx_info=self._req_ctx_info,
            ),
        )

        if resp.status != 200:
            resp_body_json = await self._read_body_json(resp)
            exc = self._parse_exception(resp_body_json)
            raise exc

        async def event_gen() -> AsyncGenerator[Tuple[str, Any], None]:
            buf = b""

            while True:
                try:
                    raw_chunk, end_of_chunk = await resp.content.readchunk()
                except asyncio.CancelledError as err:
                    raise common_exc.SourceTimeout(
                        db_message="Source timed out", query=query.debug_compiled_query
                    ) from err
                buf += raw_chunk

                # This isn't a very correct way, but it's hard to use pickle in async differently.
                if end_of_chunk and buf:
                    try:
                        parsed_event = pickle.loads(buf)
                    except Exception as err:
                        raise QueryExecutorException("QE parse: failed to unpickle") from err

                    buf = b""
                    if not isinstance(parsed_event, tuple):
                        raise QueryExecutorException(f"QE parse: unexpected event type: {type(parsed_event)}")
                    if len(parsed_event) != 2:
                        raise QueryExecutorException(f"QE parse: event is not a pair: length={len(parsed_event)}")
                    event_type_name, event_data = parsed_event
                    if not isinstance(event_type_name, str):
                        raise QueryExecutorException(f"QE parse: event_type is not a str: {type(event_type_name)}")
                    try:
                        event_type = RQEEventType[event_type_name]
                    except KeyError:
                        raise QueryExecutorException(f"QE parse: unknown event_type: {event_type_name!r}")

                    yield event_type, event_data  # type: ignore  # TODO: fix

                if raw_chunk == b"" and not end_of_chunk:
                    return

        events = event_gen()
        ev_type, ev_data = await events.__anext__()
        if ev_type == RQEEventType.raw_cursor_info:
            raw_cursor_info = ev_data
        else:
            raise QueryExecutorException(f"QE parse: first event type is not 'raw_cursor_info': {ev_type}")

        # 3-layer iterable: generator of chunks, chunks is a list of rows, row
        # is a list of column values.
        async def data_generator() -> AsyncGenerator[Sequence[Sequence[Any]], None]:
            ev_type = None

            async for ev_type, ev_data in events:
                if ev_type == RQEEventType.raw_chunk:
                    chunk = ev_data
                    if not isinstance(chunk, (list, tuple)):
                        raise QueryExecutorException(f"QE parse: unexpected chunk type: {type(chunk)}")
                    yield chunk
                elif ev_type == RQEEventType.error_dump:
                    try:
                        exc = self._parse_exception(ev_data)
                    except Exception:
                        raise QueryExecutorException(f"QE parse: failed to parse an error event: {ev_data!r}")
                    raise exc
                elif ev_type == RQEEventType.finished:
                    return

            raise QueryExecutorException(f"QE parse: finish event was not received (last: {ev_type})")

        return AsyncRawExecutionResult(
            raw_cursor_info=raw_cursor_info,
            raw_chunk_generator=data_generator(),
        )

    async def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        return await self._make_request_parse_response(
            dba_actions.ActionGetDBVersion(
                target_conn_dto=self._target_dto,
                dba_cls=self._dba_cls,
                req_ctx_info=self._req_ctx_info,
                db_ident=db_ident,
            ),
        )

    async def get_schema_names(self, db_ident: DBIdent) -> List[str]:
        return await self._make_request_parse_response(
            dba_actions.ActionGetSchemaNames(
                target_conn_dto=self._target_dto,
                dba_cls=self._dba_cls,
                req_ctx_info=self._req_ctx_info,
                db_ident=db_ident,
            ),
        )

    async def get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        return await self._make_request_parse_response(  # type: ignore  # TODO: fix
            dba_actions.ActionGetTables(
                target_conn_dto=self._target_dto,
                dba_cls=self._dba_cls,
                req_ctx_info=self._req_ctx_info,
                schema_ident=schema_ident,
            ),
        )

    async def get_table_info(self, table_def: TableDefinition, fetch_idx_info: bool) -> RawSchemaInfo:
        return await self._make_request_parse_response(
            dba_actions.ActionGetTableInfo(
                target_conn_dto=self._target_dto,
                dba_cls=self._dba_cls,
                req_ctx_info=self._req_ctx_info,
                table_def=table_def,
                fetch_idx_info=fetch_idx_info,
            ),
        )

    async def is_table_exists(self, table_ident: TableIdent) -> bool:
        return await self._make_request_parse_response(
            dba_actions.ActionIsTableExists(
                target_conn_dto=self._target_dto,
                dba_cls=self._dba_cls,
                req_ctx_info=self._req_ctx_info,
                table_ident=table_ident,
            ),
        )

    async def test(self) -> None:
        return await self._make_request_parse_response(
            dba_actions.ActionTest(
                target_conn_dto=self._target_dto,
                dba_cls=self._dba_cls,
                req_ctx_info=self._req_ctx_info,
            ),
        )

    async def close(self) -> None:
        await self._session.close()

lib/dl_core/dl_core/connection_executors/adapters/async_adapters_remote.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import logging
from typing import (
    TYPE_CHECKING,
    Generator,
    List,
    Optional,
    Set,
)
import weakref

import attr

from dl_api_commons.aio.async_wrapper_for_sync_generator import (
    EndOfStream,
    Job,
)
from dl_core.connection_executors.adapters.adapters_base import SyncDirectDBAdapter
from dl_core.connection_executors.adapters.async_adapters_base import (
    AsyncDBAdapter,
    AsyncRawExecutionResult,
)
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    ExecutionStep,
    ExecutionStepCursorInfo,
    ExecutionStepDataChunk,
    RawSchemaInfo,
)
from dl_utils.aio import ContextVarExecutor


if TYPE_CHECKING:
    from dl_constants.types import TBIChunksGen
    from dl_core.connection_models.common_models import (
        DBIdent,
        SchemaIdent,
        TableDefinition,
        TableIdent,
    )


LOGGER = logging.getLogger(__name__)


@attr.s()
class AsyncWrapperForSyncAdapter(AsyncDBAdapter):
    _sync_adapter: SyncDirectDBAdapter = attr.ib()
    _tpe: ContextVarExecutor = attr.ib()
    _loop: asyncio.AbstractEventLoop = attr.ib(init=False, factory=asyncio.get_running_loop)
    _gen_wrappers_weak_set: Set[Job] = attr.ib(init=False, factory=weakref.WeakSet)

    @attr.s(cmp=False, hash=False)
    class AdapterExecuteJob(Job[ExecutionStep]):  # noqa
        """Class to be instantiated on connection executor"""

        _adapter: SyncDirectDBAdapter = attr.ib(kw_only=True)
        _query: DBAdapterQuery = attr.ib(kw_only=True)

        def make_generator(self) -> Generator[ExecutionStep, None, None]:
            return self._adapter.execute_by_steps(self._query)

    async def execute(self, query: DBAdapterQuery) -> AsyncRawExecutionResult:
        # TODO FIX: Separate TPE
        gen_wrapper = self.AdapterExecuteJob(
            service_tpe=self._tpe,
            workers_tpe=self._tpe,
            query=query,
            adapter=self._sync_adapter,
        )
        self._gen_wrappers_weak_set.add(gen_wrapper)
        await gen_wrapper.run()

        first_msg = await gen_wrapper.get_next()
        if not isinstance(first_msg, ExecutionStepCursorInfo):
            # TODO FIX: Custom exception
            raise ValueError(f"Unexpected type of first execution message in queue: {type(first_msg)}")

        async def data_generator() -> TBIChunksGen:
            try:
                while True:
                    msg = await gen_wrapper.get_next()

                    if isinstance(msg, ExecutionStepDataChunk):
                        yield msg.chunk
                    elif msg is None:
                        break
                    else:
                        # TODO FIX: Custom exception
                        raise ValueError(f"Unexpected type of execution message in queue: {type(msg)}")
            except EndOfStream:
                return
            finally:
                await gen_wrapper.cancel()

        return AsyncRawExecutionResult(
            raw_cursor_info=first_msg.cursor_info,
            raw_chunk_generator=data_generator(),
        )

    async def test(self) -> None:
        await self._loop.run_in_executor(self._tpe, self._sync_adapter.test)

    async def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        return await self._loop.run_in_executor(self._tpe, self._sync_adapter.get_db_version, db_ident)

    async def get_schema_names(self, db_ident: DBIdent) -> List[str]:
        return await self._loop.run_in_executor(self._tpe, self._sync_adapter.get_schema_names, db_ident)

    async def get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        return await self._loop.run_in_executor(self._tpe, self._sync_adapter.get_tables, schema_ident)

    async def get_table_info(self, table_def: TableDefinition, fetch_idx_info: bool) -> RawSchemaInfo:
        return await self._loop.run_in_executor(
            self._tpe,
            self._sync_adapter.get_table_info,
            table_def,
            fetch_idx_info,
        )

    async def is_table_exists(self, table_ident: TableIdent) -> bool:
        return await self._loop.run_in_executor(self._tpe, self._sync_adapter.is_table_exists, table_ident)

    async def close(self) -> None:
        LOGGER.info("Starting async executor wrapper close procedure")
        fut_list = []
        for wr in self._gen_wrappers_weak_set:
            fut_list.append(wr.cancel())
        if fut_list:
            LOGGER.info("Waiting for %s running generators to be closed", len(fut_list))
            await asyncio.gather(*fut_list, return_exceptions=True)
        else:
            LOGGER.info("No running generator found")
        # TODO FIX: add timeout
        LOGGER.info("Waiting for sync adapter to be closed")
        await self._loop.run_in_executor(self._tpe, self._sync_adapter.close)
        LOGGER.info("Sync adapter was closed")

lib/dl_core/dl_core/connection_executors/adapters/async_adapters_sync_wrapper.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Generic,
    Optional,
    Type,
    TypeVar,
)

import sqlalchemy as sa
from sqlalchemy.engine.default import DefaultDialect

from dl_constants.enums import (
    ConnectionType,
    SourceBackendType,
)
from dl_core.backend_types import get_backend_type
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI


if TYPE_CHECKING:
    from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO  # noqa: F401


_DIALECT_STRINGS: dict[ConnectionType, str] = {}  # Filled from connectors


def register_dialect_string(conn_type: ConnectionType, dialect_str: str) -> None:
    if (registered_dialect_string := _DIALECT_STRINGS.get(conn_type)) is not None:
        assert registered_dialect_string == dialect_str
    else:
        _DIALECT_STRINGS[conn_type] = dialect_str


def get_dialect_string(conn_type: ConnectionType) -> str:
    return _DIALECT_STRINGS[conn_type]


def get_dialect_for_conn_type(conn_type: ConnectionType) -> DefaultDialect:
    engine = sa.create_engine(f"{get_dialect_string(conn_type)}://", strategy="mock", executor=lambda *_, **__: None)
    engine = engine.execution_options(compiled_cache=None)
    return engine.dialect


_DBA_TV = TypeVar("_DBA_TV", bound="CommonBaseDirectAdapter")
_TARGET_DTO_TV = TypeVar("_TARGET_DTO_TV", bound="ConnTargetDTO")


class CommonBaseDirectAdapter(Generic[_TARGET_DTO_TV], metaclass=abc.ABCMeta):
    # TODO FIX: May be use dialect strings right here???
    conn_type: ClassVar[ConnectionType]  # Kostyl to use `get_dialect_string`  # FIXME: Switch to backend_type here

    def get_backend_type(self) -> SourceBackendType:
        return get_backend_type(self.conn_type)

    @classmethod
    def get_dialect_str(cls) -> str:
        return get_dialect_string(cls.conn_type)

    @classmethod
    @abc.abstractmethod
    def create(
        cls: Type[_DBA_TV], target_dto: _TARGET_DTO_TV, req_ctx_info: DBAdapterScopedRCI, default_chunk_size: int
    ) -> _DBA_TV:
        """Generic way to create"""

    @classmethod
    def get_dialect(cls) -> DefaultDialect:
        return get_dialect_for_conn_type(cls.conn_type)

    def compile_query_for_execution(
        self, query: sa.sql.Select | str, dialect: Optional[sa.engine.Dialect] = None
    ) -> str:
        """Should not be used unless it is impossible to avoid"""
        if isinstance(query, str):
            return query

        if dialect is None:
            dialect = self.get_dialect()

        compiled_query = query.compile(
            dialect=dialect,
            compile_kwargs={"literal_binds": True},
        )

        # TODO: check the `dialect dbapi paramstyle`
        query_text = str(compiled_query)

        return query_text

lib/dl_core/dl_core/connection_executors/adapters/common_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    Any,
    Callable,
    ClassVar,
    Dict,
    Optional,
    Tuple,
)

from sqlalchemy.sql.type_api import TypeEngine

from dl_constants.enums import ConnectionType
from dl_core.connection_executors.adapters.async_adapters_base import AsyncRawExecutionResult
from dl_core.connection_executors.adapters.sa_utils import get_db_version_query
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    ExecutionStepCursorInfo,
)
from dl_core.connection_models import DBIdent
from dl_core.db.native_type import (
    CommonNativeType,
    SATypeSpec,
)


LOGGER = logging.getLogger(__name__)


class SAColumnTypeNormalizer:
    def normalize_sa_col_type(self, sa_col_type: TypeEngine) -> TypeEngine:
        return sa_col_type


class SATypeTransformer(SAColumnTypeNormalizer):
    _type_code_to_sa: Optional[Dict[Any, SATypeSpec]] = None
    conn_type: ClassVar[ConnectionType]

    @staticmethod
    def _cursor_type_to_str(value: Any) -> str:
        if isinstance(value, type):
            return value.__name__.lower()
        if isinstance(value, int):
            return value  # type: ignore  # TODO: fix
        return str(value).lower()

    def _cursor_column_to_name(self, cursor_col: Tuple[Any, ...], dialect: Any = None) -> str:
        return cursor_col[0]

    def _cursor_column_to_sa(self, cursor_col: Tuple[Any, ...], require: bool = True) -> Optional[SATypeSpec]:
        type_code_to_sa = self._type_code_to_sa
        if not type_code_to_sa:
            if not require:
                return None
            raise Exception(f"No defined `self._type_code_to_sa` on {self.__class__!r}")

        type_code = cursor_col[1]
        sa_type = type_code_to_sa.get(type_code)
        return sa_type

    def _cursor_column_to_nullable(self, cursor_col: Tuple[Any, ...]) -> Optional[bool]:
        # No known `nullable=False` cases for subselects in PG and MySQL and
        # Oracle. But that might change.
        return True

    def _cursor_column_to_native_type(
        self, cursor_col: Tuple[Any, ...], require: bool = True
    ) -> Optional[CommonNativeType]:
        sa_type = self._cursor_column_to_sa(cursor_col, require=require)
        if sa_type is None:
            if not require:
                return None
            raise ValueError(f"Unknown/unsupported type_code for {self}, column info {cursor_col}", self, cursor_col)

        nullable = self._cursor_column_to_nullable(cursor_col)
        if nullable is None:
            nullable = True

        return CommonNativeType.normalize_name_and_create(
            conn_type=self.conn_type,
            name=self.normalize_sa_col_type(sa_type),  # type: ignore  # TODO: fix
            # no idea whether it can be False here; `create view` version
            # returns `nullable=True` for all known cases (for PG and MySQL).
            nullable=nullable,
        )


class WithAsyncGetDBVersion:
    @abc.abstractmethod
    async def execute(self, query: DBAdapterQuery) -> AsyncRawExecutionResult:
        pass

    async def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        result = await self.execute(get_db_version_query(db_ident))
        async for row in result.get_all_rows():
            return str(row[0])
        return None


class WithCursorInfo:
    def _make_cursor_info(self, cursor, db_session=None) -> dict:  # type: ignore  # TODO: fix
        return {}


class WithMinimalCursorInfo(WithCursorInfo, SATypeTransformer):
    def _make_cursor_info(self, cursor, db_session=None) -> dict:  # type: ignore  # TODO: fix
        """
        Minimal information from the cursor,
        will likely get processed through JSON at some point,
        therefore should only contain jsonable primitives.
        """
        return dict(
            super()._make_cursor_info(cursor, db_session),
            names=[str(cursor_col[0]) for cursor_col in cursor.description],
            driver_types=[self._cursor_type_to_str(cursor_col[1]) for cursor_col in cursor.description],
            db_types=[
                self._cursor_column_to_native_type(cursor_col, require=False) for cursor_col in cursor.description
            ],
        )


class WithDatabaseNameOverride:
    warn_on_default_db_name_override: ClassVar[bool] = True

    @abc.abstractmethod
    def get_default_db_name(self) -> Optional[str]:
        pass

    def get_db_name_for_query(self, db_name_from_query: Optional[str]) -> str:
        return self._get_db_name_for_query(
            default=self.get_default_db_name(),
            from_query=db_name_from_query,
            warn_override=self.warn_on_default_db_name_override,
        )

    @staticmethod
    def _get_db_name_for_query(default: Optional[str], from_query: Optional[str], warn_override: bool) -> str:
        if default is not None and from_query is not None:
            if default != from_query and warn_override:
                LOGGER.warning(f"Divergence in DB names: default='{default}' from_query='{from_query}'")
            return from_query
        elif default is not None:
            return default
        elif from_query is not None:
            return from_query
        else:
            raise ValueError("Can not determine DB name for engine construction. Both are None")


class WithNoneRowConverters:
    def _get_row_converters(self, cursor_info: ExecutionStepCursorInfo) -> Tuple[Optional[Callable[[Any], Any]], ...]:
        return tuple(None for _ in cursor_info.raw_cursor_description)

lib/dl_core/dl_core/connection_executors/adapters/mixins.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import time
from typing import Union

import opentracing
import sqlalchemy as sa
from sqlalchemy import sql as sasql
from sqlalchemy.engine import Dialect

from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
from dl_core.connection_models import DBIdent


LOGGER = logging.getLogger(__name__)


def get_db_version_query(db_ident: DBIdent) -> DBAdapterQuery:
    """Return database version string query for this connection"""
    return DBAdapterQuery(sa.select([sa.func.version()]), db_name=db_ident.db_name)


def compile_query_for_debug(query: sasql.Select, dialect: Dialect) -> str:
    """
    Compile query to string.
    This function is only suitable for logging and not execution of the result.
    Its result might not be valid SQL if query contains date/datetime literals.
    """
    # noinspection PyBroadException
    try:
        try:
            return str(query.compile(dialect=dialect, compile_kwargs={"literal_binds": True}))
        except NotImplementedError:
            compiled = query.compile(dialect=dialect)
            return make_debug_query(str(query), compiled.params)
    except Exception:
        return "-"


def make_debug_query(query: str, params: Union[list, dict]) -> str:
    return f"{query} {params!r}"


class CursorLogger:
    @staticmethod
    def before_cursor_execute_handler(conn, cursor, statement, parameters, context, executemany):  # type: ignore
        tracer = opentracing.global_tracer()
        # Scope was not created due to decoupled span closing procedure
        #  which may cause broken parent-child relationships in case when `after_cursor_execute_handler()`
        #  will not be called
        span = tracer.start_span("sa-cursor-execute")

        engine_url = context.engine.url
        span.log_kv(
            dict(
                sa_dialect=engine_url.drivername,
            )
        )

        conn.info.setdefault("query_start_time", []).append(time.monotonic())
        conn.info.setdefault("ot_span_scope_stack", []).append(span)

    @staticmethod
    def after_cursor_execute_handler(conn, cursor, statement, parameters, context, executemany):  # type: ignore
        engine_url = context.engine.url
        query_start_time = conn.info["query_start_time"].pop(-1)

        span: opentracing.Span = conn.info["ot_span_scope_stack"].pop()
        span.finish()

        execution_time_seconds = time.monotonic() - query_start_time

        extra = dict(
            event_code="db_exec",
            username=engine_url.username,
            query_id=str(getattr(cursor, "_query_id", None)),
            # missing: folder_id
            execution_time=int(round(execution_time_seconds * 1000)),
            database=engine_url.database,
            host=engine_url.host,
            drivername=engine_url.drivername,
            statement=statement,
            parameters_size=len(parameters),
            # missing: dataset_id
            # missing: connection_id
        )
        LOGGER.info("Cursor executed", extra=extra)

lib/dl_core/dl_core/connection_executors/adapters/sa_utils.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
import functools
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncIterable,
    Awaitable,
    Callable,
    List,
    Optional,
    Sequence,
    TypeVar,
)

import attr
from typing_extensions import final

from dl_core.connection_executors.common_base import (
    ConnExecutorBase,
    ConnExecutorQuery,
)
from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
from dl_core.db import SchemaInfo


if TYPE_CHECKING:
    from dl_constants.enums import UserDataType
    from dl_constants.types import TBIDataTable
    from dl_core.connection_models.common_models import (
        DBIdent,
        SchemaIdent,
        TableDefinition,
        TableIdent,
    )


LOGGER = logging.getLogger(__name__)


@attr.s
class AsyncExecutionResult:
    cursor_info: dict = attr.ib()
    result: AsyncIterable[TBIDataTable] = attr.ib()  # iterable of tables (chunks)
    # for `autodetect_user_types` result
    user_types: Optional[List[UserDataType]] = attr.ib(default=None)
    # DB-specific result. Should be mutable, and get filled after `result` is consumed.
    result_footer: dict = attr.ib(factory=dict)

    async def get_all(self) -> Sequence[Sequence]:
        """
        :return: All fetched data
        """
        result: List[Sequence] = []
        async for chunk in self.result:
            result.extend(chunk)

        return result


_RET_TV = TypeVar("_RET_TV")


def init_required(wrapped: Callable[..., Awaitable[_RET_TV]]) -> Callable[..., Awaitable[_RET_TV]]:
    @functools.wraps(wrapped)
    async def wrapper(self: "AsyncConnExecutorBase", *args: Any, **kwargs: Any) -> _RET_TV:
        if not self._is_initialized:
            await self.initialize()

        return await wrapped(self, *args, **kwargs)

    return wrapper


@attr.s(cmp=False, hash=False)
class AsyncConnExecutorBase(ConnExecutorBase, metaclass=abc.ABCMeta):
    @classmethod
    @abc.abstractmethod
    def is_pure_async(cls) -> bool:
        pass

    async def initialize(self) -> None:
        pass

    @abc.abstractmethod
    def executor_query_to_db_adapter_query(self, conn_exec_query: ConnExecutorQuery) -> DBAdapterQuery:
        pass

    # DB Interaction method wrappers
    @final
    @init_required
    async def execute(self, query: ConnExecutorQuery) -> AsyncExecutionResult:
        """
        This method should not be overridden!
        :param query:
        :return: Chunks of result
        """
        return await self._execute(query)

    @final
    @init_required
    async def test(self) -> None:
        await self._test()

    @final
    @init_required
    async def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        return await self._get_db_version(db_ident)

    @final
    @init_required
    async def get_schema_names(self, db_ident: DBIdent) -> List[str]:
        return await self._get_schema_names(db_ident)

    @final
    @init_required
    async def get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        return await self._get_tables(schema_ident)

    @final
    @init_required
    async def get_table_schema_info(self, table_def: TableDefinition) -> SchemaInfo:
        return await self._get_table_schema_info(table_def)

    @final
    @init_required
    async def is_table_exists(self, table_ident: TableIdent) -> bool:
        return await self._is_table_exists(table_ident)

    @final
    async def close(self) -> None:
        LOGGER.info("Method close() of connection executor was called")
        return await self._close()

    # Methods to implements
    @abc.abstractmethod
    async def _execute(self, query: ConnExecutorQuery) -> AsyncExecutionResult:
        pass

    @abc.abstractmethod
    async def _test(self) -> None:
        pass

    @abc.abstractmethod
    async def _get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        pass

    @abc.abstractmethod
    async def _get_schema_names(self, db_ident: DBIdent) -> List[str]:
        pass

    @abc.abstractmethod
    async def _get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        pass

    @abc.abstractmethod
    async def _get_table_schema_info(self, table_def: TableDefinition) -> SchemaInfo:
        pass

    @abc.abstractmethod
    async def _is_table_exists(self, table_ident: TableIdent) -> bool:
        pass

    @abc.abstractmethod
    async def _close(self) -> None:
        pass

lib/dl_core/dl_core/connection_executors/async_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
import functools
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Dict,
    Generic,
    List,
    Optional,
    Sequence,
    Type,
    TypeVar,
)

import attr
from typing_extensions import final

from dl_api_commons.base_models import RequestContextInfo
from dl_constants.enums import UserDataType
from dl_core import exc
from dl_core.connection_executors.adapters.adapters_base import SyncDirectDBAdapter
from dl_core.connection_executors.adapters.async_adapters_base import (
    AsyncDBAdapter,
    AsyncDirectDBAdapter,
    AsyncRawExecutionResult,
)
from dl_core.connection_executors.adapters.async_adapters_remote import RemoteAsyncAdapter
from dl_core.connection_executors.adapters.async_adapters_sync_wrapper import AsyncWrapperForSyncAdapter
from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.async_base import (
    AsyncConnExecutorBase,
    AsyncExecutionResult,
)
from dl_core.connection_executors.common_base import (
    ConnExecutorQuery,
    ExecutionMode,
)
from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
from dl_core.connection_models.common_models import TableDefinition
from dl_core.db import SchemaInfo
from dl_utils.aio import ContextVarExecutor
from dl_utils.utils import maybe_postmortem


if TYPE_CHECKING:
    from dl_constants.types import TBIChunksGen
    from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO
    from dl_core.connection_models.common_models import (
        DBIdent,
        SchemaIdent,
        TableIdent,
    )


LOGGER = logging.getLogger(__name__)

_DBA_TV = TypeVar("_DBA_TV", bound=CommonBaseDirectAdapter)


def _dba_pool_revolver_wrapper(func: Callable) -> Callable:
    @functools.wraps(func)
    async def method_wrapper(self: DefaultSqlAlchemyConnExecutor, *args, **kwargs):  # type: ignore  # TODO: fix
        # Not sure do we have to reset index on every call
        # self._dba_attempt_index = 0

        while self._dba_attempt_index < len(self._async_dba_pool):  # type: ignore  # TODO: fix
            try:
                return await func(self, *args, **kwargs)
            except (exc.SourceConnectError, exc.DatabaseOperationalError) as ex:
                LOGGER.info("Retriable exception caught: %s", type(ex), exc_info=True)

                if self._conn_hosts_pool:
                    self._host_fail_callback(self._conn_hosts_pool[self._dba_attempt_index])

                self._dba_attempt_index += 1
                if self._dba_attempt_index >= len(self._async_dba_pool):  # type: ignore  # TODO: fix
                    raise
            except Exception as ex:
                LOGGER.info("Non-retriable dba exception %s: %s", type(ex), ex, exc_info=True)
                raise

    return method_wrapper


def _postprocess_db_excs_wrapper(func: Callable) -> Callable:
    @functools.wraps(func)
    async def method_wrapper(self: DefaultSqlAlchemyConnExecutor, *args, **kwargs):  # type: ignore  # TODO: fix
        with self._postprocess_db_excs():
            return await func(self, *args, **kwargs)

    return method_wrapper


def _common_exec_wrapper(func: Callable) -> Callable:
    func = _dba_pool_revolver_wrapper(func)
    func = _postprocess_db_excs_wrapper(func)
    return func


# TODO FIX: Encapsulate STRAIGHTFORWARD/WRAPPED state to single object to simplify code
@attr.s(cmp=False, hash=False)
class DefaultSqlAlchemyConnExecutor(AsyncConnExecutorBase, Generic[_DBA_TV], metaclass=abc.ABCMeta):
    """
    Base class that use SA adapter. To execute queries.
    Executor has 2 use cases in DIRECT execution mode (in RQE execution mode this is not true):
      STRAIGHTFORWARD: Direct execution in async environment
      WRAPPED: Wrapping by sync wrapper in sync environment
    To use STRAIGHTFORWARD - pass TPE in constructor. It will be used to execute sync operations.
    To use WRAPPED - do not pass TPE. Assumed that executor will be used by sync wrapper.
      All direct operations will cause an exception in this case.
      In this mode async database adapter will not be created.
      Adapter close process is delegated to

    Why do we need it: to keep initialization creation of DBA in one place for sync and async environments.
    """

    TARGET_ADAPTER_CLS: ClassVar[Type[_DBA_TV]]

    # Constructor attributes
    _tpe: Optional[ContextVarExecutor] = attr.ib()

    # Internals
    _async_dba_pool: Optional[List[AsyncDBAdapter]] = attr.ib(init=False, default=None)
    _sync_dba: Optional[SyncDirectDBAdapter] = attr.ib(init=False, default=None)
    _initialization_lock = attr.ib(init=False, factory=asyncio.Lock)
    _dba_attempt_index: int = attr.ib(init=False, default=0)

    @classmethod
    def is_pure_async(cls) -> bool:
        return issubclass(cls.TARGET_ADAPTER_CLS, AsyncDirectDBAdapter)

    @property
    def _target_dba(self) -> AsyncDBAdapter:
        if self._async_dba_pool is None:
            raise ValueError("Executor created in WRAPPED mode (no TPE was provided). Async DBA was not created.")
        return self._async_dba_pool[self._dba_attempt_index]

    def _get_sync_sa_adapter(self) -> Optional[SyncDirectDBAdapter]:
        """Method for sync conn executor wrapper. Please do not use it in other places."""
        if not self._is_initialized:
            raise ValueError("Connection executor was not initiated")
        return self._sync_dba

    @abc.abstractmethod
    async def _make_target_conn_dto_pool(self) -> Sequence[ConnTargetDTO]:
        """
        This method will be used for creating target connection DTO to pass it to RQE.
        This method also should be used in `_make_sa_adapter()` to create adapter for DIRECT mode.
        :return: Target connection DTO
        """

    @final
    async def initialize(self) -> None:
        async with self._initialization_lock:
            if self._is_initialized:
                return
            else:
                await self._initialize()
                self._is_initialized = True

    @final
    async def _initialize(self) -> None:
        target_conn_dto_pool = await self._make_target_conn_dto_pool()
        assert target_conn_dto_pool

        req_ctx_info = self._req_ctx_info or RequestContextInfo.create_empty()

        if self._exec_mode == ExecutionMode.DIRECT:
            if issubclass(self.TARGET_ADAPTER_CLS, SyncDirectDBAdapter):
                sync_dba = self.TARGET_ADAPTER_CLS.create(
                    target_dto=target_conn_dto_pool[0],
                    req_ctx_info=DBAdapterScopedRCI.from_full_rci(req_ctx_info),
                    default_chunk_size=self.default_chunk_size,
                )
                self._sync_dba = sync_dba

                # Normal use case
                if self._tpe is not None:
                    self._async_dba_pool = [
                        AsyncWrapperForSyncAdapter(
                            sync_adapter=sync_dba,
                            tpe=self._tpe,
                        )
                    ]
                # Sync wrapper use-case
                else:
                    LOGGER.info("TPE was not provided async wrapper for DB adapter will not be created")
            elif issubclass(self.TARGET_ADAPTER_CLS, AsyncDirectDBAdapter):
                self._async_dba_pool = [
                    self.TARGET_ADAPTER_CLS.create(
                        target_dto=target_conn_dto,
                        req_ctx_info=DBAdapterScopedRCI.from_full_rci(req_ctx_info),
                        default_chunk_size=self.default_chunk_size,
                    )
                    for target_conn_dto in target_conn_dto_pool
                ]
            else:
                raise TypeError(f"Can not create DBA TARGET_ADAPTER_CLS={self.TARGET_ADAPTER_CLS}")

        elif self._exec_mode == ExecutionMode.RQE:
            if self._remote_qe_data is None:
                raise ValueError("Can not initialize: no remote QE data was provided")
            self._async_dba_pool = [
                RemoteAsyncAdapter(
                    target_dto=target_conn_dto,
                    rqe_data=self._remote_qe_data,
                    req_ctx_info=DBAdapterScopedRCI.from_full_rci(req_ctx_info),
                    dba_cls=self.TARGET_ADAPTER_CLS,
                    conn_options=self._conn_options,
                )
                for target_conn_dto in target_conn_dto_pool
            ]

        else:
            raise NotImplementedError(f"Unsupported execution mode {self._exec_mode}")

    # noinspection PyMethodMayBeStatic
    def executor_query_to_db_adapter_query(self, conn_exec_query: ConnExecutorQuery) -> DBAdapterQuery:
        return DBAdapterQuery(
            query=conn_exec_query.query,
            chunk_size=conn_exec_query.chunk_size,
            db_name=conn_exec_query.db_name,
            debug_compiled_query=conn_exec_query.debug_compiled_query,
            connector_specific_params=conn_exec_query.connector_specific_params,
            trusted_query=conn_exec_query.trusted_query,
            is_ddl_dml_query=conn_exec_query.is_ddl_dml_query,
            is_dashsql_query=conn_exec_query.is_dashsql_query,
        )

    @_common_exec_wrapper
    async def _execute_query(self, query: DBAdapterQuery) -> AsyncRawExecutionResult:
        return await self._target_dba.execute(query)

    def _autodetect_user_types(self, raw_cursor_info: dict) -> Optional[List[UserDataType]]:
        db_types = raw_cursor_info.get("db_types")
        if not db_types:
            return None
        tt = self.type_transformer

        result = []
        for native_type in db_types:
            bi_type = UserDataType.unsupported

            if native_type:
                try:
                    bi_type = tt.type_native_to_user(native_type)
                except exc.UnsupportedNativeTypeError:
                    pass

            result.append(bi_type)

        return result

    async def _execute(self, query: ConnExecutorQuery) -> AsyncExecutionResult:
        raw_result = await self._execute_query(self.executor_query_to_db_adapter_query(query))

        user_types = query.user_types
        if query.autodetect_user_types:
            user_types = self._autodetect_user_types(raw_result.raw_cursor_info)

        result_footer: Dict[str, Any] = {}

        async def data_generator() -> TBIChunksGen:
            async for chunk in raw_result.raw_chunk_generator:
                yield [self.cast_row_to_output(row, query.user_types) for row in chunk]

        return AsyncExecutionResult(
            cursor_info=raw_result.raw_cursor_info,
            result=data_generator(),
            user_types=user_types,
            result_footer=result_footer,
        )

    @_common_exec_wrapper
    async def _test(self) -> None:
        await self._target_dba.test()

    @_common_exec_wrapper
    async def _get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        return await self._target_dba.get_db_version(db_ident)

    @_common_exec_wrapper
    async def _get_schema_names(self, db_ident: DBIdent) -> List[str]:
        return await self._target_dba.get_schema_names(db_ident)

    @_common_exec_wrapper
    async def _get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        return await self._target_dba.get_tables(schema_ident)

    @_common_exec_wrapper
    async def _get_table_schema_info(self, table_def: TableDefinition) -> SchemaInfo:
        raw_table_info = await self._target_dba.get_table_info(table_def, self._conn_options.fetch_table_indexes)
        return self.create_schema_info_from_raw_schema_info(raw_table_info)

    @_common_exec_wrapper
    async def _is_table_exists(self, table_ident: TableIdent) -> bool:
        return await self._target_dba.is_table_exists(table_ident)

    async def _close(self) -> None:
        if self._async_dba_pool is not None:
            # noinspection PyBroadException
            for dba in self._async_dba_pool:
                try:
                    # TODO FIX: May be make timeout
                    await dba.close()
                except Exception:
                    LOGGER.exception("Error during async adapter closing attempt")

lib/dl_core/dl_core/connection_executors/async_sa_executors.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import contextlib
import enum
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Dict,
    FrozenSet,
    Generator,
    List,
    Optional,
    Sequence,
    TypeVar,
    Union,
)

import attr
from sqlalchemy import sql as sasql

from dl_api_commons.base_models import RequestContextInfo
from dl_constants.enums import UserDataType
from dl_constants.exc import DLBaseException
from dl_constants.types import TBIDataValue
from dl_core.connection_executors.models.common import RemoteQueryExecutorData
from dl_core.connection_executors.models.db_adapter_data import RawSchemaInfo
from dl_core.connection_models import (
    ConnDTO,
    ConnectOptions,
)
from dl_core.db import (
    SAMPLE_ID_COLUMN_NAME,
    IndexInfo,
    SchemaColumn,
    SchemaInfo,
    TypeTransformer,
    get_type_transformer,
)
from dl_core.exc import UnsupportedNativeTypeError


if TYPE_CHECKING:
    from dl_constants.types import TJSONExt
    from dl_core.connections_security.base import ConnectionSecurityManager  # noqa
    from dl_core.services_registry.top_level import ServicesRegistry


LOGGER = logging.getLogger(__name__)


class ExecutionMode(enum.Enum):
    """Will be turned into something like (is_safe: bool, is_async: bool, via_qe: bool)"""

    DIRECT = enum.auto()
    RQE = enum.auto()


@attr.s
class ConnExecutorQuery:
    query: Union[sasql.Select, str] = attr.ib()
    user_types: Optional[List[UserDataType]] = attr.ib(default=None)
    debug_compiled_query: Optional[str] = attr.ib(default=None)
    chunk_size: Optional[int] = attr.ib(default=None)
    connector_specific_params: Optional[Dict[str, TJSONExt]] = attr.ib(default=None)
    # TODO FIX: We really need it in query?
    db_name: Optional[str] = attr.ib(default=None)
    autodetect_user_types: bool = attr.ib(default=False)
    trusted_query: bool = attr.ib(default=False)
    is_ddl_dml_query: bool = attr.ib(default=False)
    is_dashsql_query: bool = attr.ib(default=False)


_CONN_EXEC_TV = TypeVar("_CONN_EXEC_TV", bound="ConnExecutorBase")


@attr.s(cmp=False, hash=False)
class ConnExecutorBase(metaclass=abc.ABCMeta):
    default_chunk_size: ClassVar[int] = 100
    supported_exec_mode: ClassVar[FrozenSet[ExecutionMode]] = frozenset(ExecutionMode)

    _conn_dto: ConnDTO = attr.ib()
    _conn_options: ConnectOptions = attr.ib()
    _conn_hosts_pool: Sequence[str] = attr.ib()
    _host_fail_callback: Callable = attr.ib()
    _req_ctx_info: Optional[RequestContextInfo] = attr.ib()
    _exec_mode: ExecutionMode = attr.ib()
    _sec_mgr: "ConnectionSecurityManager" = attr.ib()
    _remote_qe_data: Optional[RemoteQueryExecutorData] = attr.ib()
    _services_registry: Optional[ServicesRegistry] = attr.ib(
        kw_only=True, default=None
    )  # Do not use. To be deprecated. Somehow.
    _is_initialized: bool = attr.ib(init=False, default=False)

    @_exec_mode.validator
    def _validate_exc_mode(self, _: Any, value: Any) -> None:
        if value not in self.supported_exec_mode:
            raise ValueError(f"Unsupported execution mode {value} for {type(self).__qualname__}")

    @property
    def is_initialized(self) -> bool:
        return self._is_initialized

    @property
    def type_transformer(self) -> TypeTransformer:
        return get_type_transformer(self._conn_dto.conn_type)

    def is_conn_dto_equals(self, another: ConnDTO) -> bool:
        try:
            return self._conn_dto == another
        except Exception:  # noqa
            LOGGER.exception("Exception during connection DTO comparision")
            return False

    def is_context_info_equals(self, another: RequestContextInfo) -> bool:
        try:
            return self._req_ctx_info == another
        except Exception:  # noqa
            LOGGER.exception("Exception during request context info comparision")
            return False

    def cast_row_to_output(self, row: Sequence, user_types: Optional[Sequence[UserDataType]]) -> Sequence[TBIDataValue]:
        if user_types is None:
            return row

        if len(user_types) != len(row):
            raise ValueError(f"Length of user_types {len(user_types)} != length of row {len(row)}")
        return tuple(
            self.type_transformer.cast_for_output(val, user_t=user_type) for val, user_type in zip(row, user_types)
        )

    def create_schema_info_from_raw_schema_info(
        self,
        raw_schema_info: RawSchemaInfo,
        require_all: bool = False,
    ) -> SchemaInfo:
        schema = []
        type_transformer = self.type_transformer

        for raw_column_info in raw_schema_info.columns:
            if raw_column_info.name == SAMPLE_ID_COLUMN_NAME:
                continue

            try:
                user_type = type_transformer.type_native_to_user(native_t=raw_column_info.native_type)
            except UnsupportedNativeTypeError:
                if require_all:
                    raise
                LOGGER.warning("Unable to detect type of field: %s", raw_column_info)
                user_type = UserDataType.unsupported

            schema_col = SchemaColumn(
                name=raw_column_info.name,
                title=raw_column_info.title or raw_column_info.name,
                user_type=user_type,
                nullable=raw_column_info.nullable,
                native_type=raw_column_info.native_type,
            )
            schema.append(schema_col)

        index_info_set: Optional[FrozenSet[IndexInfo]] = None

        if raw_schema_info.indexes is not None:
            index_info_set = frozenset(
                IndexInfo(
                    columns=raw_idx_info.columns,
                    kind=raw_idx_info.kind,
                )
                for raw_idx_info in raw_schema_info.indexes
            )

        return SchemaInfo(
            schema=schema,
            indexes=index_info_set,
        )

    @contextlib.contextmanager
    def _postprocess_db_excs(self) -> Generator[None, None, None]:
        try:
            yield
        except DLBaseException as exc:
            if self._conn_options.pass_db_messages_to_user:
                for key in ("query", "db_message"):
                    value = exc.debug_info.get(key)
                    if exc.details.get(key) is None and value is not None:
                        exc.details[key] = value
            raise

    def mutate_for_dashsql(self, db_params: Optional[Dict[str, str]] = None):  # type: ignore  # TODO: fix
        """
        A place to do CE-specific alterations for DashSQL support.
        """
        return self

    def clone(self: _CONN_EXEC_TV, **kwargs: Any) -> _CONN_EXEC_TV:
        return attr.evolve(self, **kwargs)

lib/dl_core/dl_core/connection_executors/common_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr


@attr.s(frozen=True)
class RemoteQueryExecutorData:
    hmac_key: bytes = attr.ib(repr=False)
    # async QE data
    async_protocol: str = attr.ib()
    async_host: str = attr.ib()
    async_port: int = attr.ib()
    # sync QE data
    sync_protocol: str = attr.ib()
    sync_host: str = attr.ib()
    sync_port: int = attr.ib()

lib/dl_core/dl_core/connection_executors/models/common.py
datalens-tech/datalens-backend
"""
DTOs representing connection targets for SA adapters.
Should be constructed by connection executors based on static connection data stored in connections storage
and possibly augmented with data calculated in runtime or fetched from cache e.g. DB server version
"""

from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Dict,
    Optional,
    Type,
    TypeVar,
)

import attr
from typing_extensions import final

from dl_core.utils import secrepr


if TYPE_CHECKING:
    from dl_constants.types import TJSONLike


_CT_DTO_TV = TypeVar("_CT_DTO_TV", bound="ConnTargetDTO")


@attr.s(frozen=True)
class ConnTargetDTO(metaclass=abc.ABCMeta):
    conn_id: Optional[str] = attr.ib()

    pass_db_messages_to_user: bool = attr.ib(default=False, kw_only=True)
    pass_db_query_to_user: bool = attr.ib(default=False, kw_only=True)

    @abc.abstractmethod
    def get_effective_host(self) -> Optional[str]:
        pass

    _MAP_CLASS_NAME_CLASS: ClassVar[Dict[str, Type[ConnTargetDTO]]] = {}

    def __init_subclass__(cls, **kwargs):  # type: ignore
        cls._MAP_CLASS_NAME_CLASS[cls.__qualname__] = cls

    @classmethod
    @final
    def from_polymorphic_jsonable_dict(cls, data: Dict) -> ConnTargetDTO:
        data = {**data}
        cls_name = data.pop("cls_name")
        target_cls = cls._MAP_CLASS_NAME_CLASS[cls_name]
        return target_cls.from_jsonable_dict(data)

    @classmethod
    @final
    def from_jsonable_dict(cls: Type[_CT_DTO_TV], data: dict) -> _CT_DTO_TV:
        if "cls_name" in data:
            raise ValueError("Property 'cls_name' in data dict. Use from_polymorphic_jsonable_dict() to unmarshal it.")
        return cls._from_jsonable_dict(data)

    @classmethod
    def _from_jsonable_dict(cls: Type[_CT_DTO_TV], data: dict) -> _CT_DTO_TV:
        return cls(**data)

    def to_jsonable_dict(self) -> Dict[str, TJSONLike]:
        return dict(
            **attr.asdict(self),
            cls_name=type(self).__qualname__,
        )

    def clone(self: _CT_DTO_TV, **kwargs: Any) -> _CT_DTO_TV:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True)
class BaseSQLConnTargetDTO(ConnTargetDTO):
    host: str = attr.ib()
    port: int = attr.ib()
    username: str = attr.ib()
    password: str = attr.ib(repr=secrepr)
    db_name: Optional[str] = attr.ib()

    def get_effective_host(self) -> Optional[str]:
        return self.host

lib/dl_core/dl_core/connection_executors/models/connection_target_dto_base.py
datalens-tech/datalens-backend
HEADER_BODY_SIGNATURE = "X-DL-QEBodySignature"
HEADER_REQUEST_ID = "X-Request-Id"

lib/dl_core/dl_core/connection_executors/models/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Optional,
    Sequence,
    Tuple,
    TypeVar,
    Union,
)

import attr
from sqlalchemy import sql as sasql


if TYPE_CHECKING:
    from dl_constants.enums import IndexKind
    from dl_constants.types import TJSONExt
    from dl_core.db.native_type import GenericNativeType


class ExecutionStep:
    pass


@attr.s(frozen=True)
class ExecutionStepCursorInfo(ExecutionStep):  # type: ignore  # TODO: fix
    cursor_info: Dict = attr.ib(factory=lambda: {})
    # Definitely not to be serialized:
    raw_cursor_description = attr.ib(default=None)  # type: ignore  # TODO: fix
    raw_engine = attr.ib(default=None)  # type: ignore  # TODO: fix


@attr.s(frozen=True)
class ExecutionStepDataChunk(ExecutionStep):
    chunk: Sequence[Sequence] = attr.ib()


_DB_ADAPTER_QUERY_TV = TypeVar("_DB_ADAPTER_QUERY_TV", bound="DBAdapterQuery")


@attr.s(frozen=True)
class DBAdapterQuery:
    query: Union[sasql.Select, str] = attr.ib()
    db_name: Optional[str] = attr.ib(default=None)
    debug_compiled_query: Optional[str] = attr.ib(default=None)
    chunk_size: Optional[int] = attr.ib(default=None)
    connector_specific_params: Optional[Dict[str, TJSONExt]] = attr.ib(default=None)
    # Use-case: `explain ...` queries are incompatible with streaming (and
    # server-side cursors in general) because psycopg2 prepends 'DECLARE ...
    # CURSOR WITHOUT HOLD FOR ...' to the statement.
    # NOTE: if there are more flags, replace the whole `disable_streaming` with a FrozenDict.
    disable_streaming: bool = attr.ib(default=False)
    trusted_query: bool = attr.ib(default=False)
    is_ddl_dml_query: bool = attr.ib(default=False)
    is_dashsql_query: bool = attr.ib(default=False)

    def get_effective_chunk_size(self, default_chunk_size: int) -> int:
        if self.chunk_size is None:
            return default_chunk_size
        else:
            return self.chunk_size

    def clone(self: _DB_ADAPTER_QUERY_TV, **kwargs: Any) -> _DB_ADAPTER_QUERY_TV:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True, slots=True, auto_attribs=True)
class RawColumnInfo:
    name: str
    title: Optional[str]
    nullable: bool
    native_type: "GenericNativeType"


@attr.s(frozen=True, slots=True, auto_attribs=True)
class RawIndexInfo:
    columns: Tuple[str, ...]
    unique: bool
    kind: Optional[IndexKind]


@attr.s(frozen=True, slots=True, auto_attribs=True)
class RawSchemaInfo:
    columns: Tuple[RawColumnInfo, ...]
    indexes: Optional[Tuple[RawIndexInfo, ...]] = None

lib/dl_core/dl_core/connection_executors/models/db_adapter_data.py
datalens-tech/datalens-backend
class QueryExecutorException(Exception):
    pass

lib/dl_core/dl_core/connection_executors/models/exc.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_api_commons.base_models import RequestContextInfo


@attr.s(frozen=True)
class DBAdapterScopedRCI:
    request_id: Optional[str] = attr.ib(default=None)
    x_dl_debug_mode: Optional[bool] = attr.ib(default=None)
    user_name: Optional[str] = attr.ib(default=None)
    client_ip: Optional[str] = attr.ib(default=None, repr=False)

    @classmethod
    def from_full_rci(cls, full_rci: RequestContextInfo) -> "DBAdapterScopedRCI":
        return cls(
            request_id=full_rci.request_id,
            x_dl_debug_mode=full_rci.x_dl_debug_mode,
            user_name=full_rci.user_name,
            client_ip=full_rci.client_ip,
        )

lib/dl_core/dl_core/connection_executors/models/scoped_rci.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_core.connection_executors.qe_serializer.serializer import (
    ActionSerializer,
    ResponseTypes,
)


__all__ = (
    "ActionSerializer",
    "ResponseTypes",
)

lib/dl_core/dl_core/connection_executors/qe_serializer/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Dict,
    Generic,
    List,
    Optional,
    Type,
    TypeVar,
)

import attr
from marshmallow import fields

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.models.db_adapter_data import (
    DBAdapterQuery,
    RawSchemaInfo,
)
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
from dl_core.connection_executors.qe_serializer.schema_base import BaseQEAPISchema
from dl_core.connection_executors.qe_serializer.schemas_common import TableIdentSchema
from dl_core.connection_executors.qe_serializer.schemas_responses import (
    PrimitivesResponseSchema,
    RawSchemaInfoSchema,
)
from dl_core.connection_models import (
    DBIdent,
    SchemaIdent,
    TableDefinition,
    TableIdent,
)


if TYPE_CHECKING:
    from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO


@attr.s(frozen=True)
class RemoteDBAdapterAction:
    target_conn_dto: ConnTargetDTO = attr.ib()
    req_ctx_info: DBAdapterScopedRCI = attr.ib()
    dba_cls: Type[CommonBaseDirectAdapter] = attr.ib()


@attr.s(frozen=True)
class ActionExecuteQuery(RemoteDBAdapterAction):
    db_adapter_query: DBAdapterQuery = attr.ib()


_RES_SCHEMA_TV = TypeVar("_RES_SCHEMA_TV")


class NonStreamAction(Generic[_RES_SCHEMA_TV], RemoteDBAdapterAction, metaclass=abc.ABCMeta):
    ResultSchema = ClassVar[Type[BaseQEAPISchema]]

    def serialize_response(self, val: _RES_SCHEMA_TV) -> Dict:
        return self.ResultSchema().dump(val)  # type: ignore  # TODO: fix

    def deserialize_response(self, data: Dict) -> _RES_SCHEMA_TV:
        return self.ResultSchema().load(data)  # type: ignore  # TODO: fix


@attr.s(frozen=True)
class ActionTest(NonStreamAction[None]):
    class ResultSchema(BaseQEAPISchema):
        def to_object(self, data: Dict[str, Any]) -> Any:
            return None


@attr.s(frozen=True)
class ActionGetDBVersion(NonStreamAction[Optional[str]]):
    db_ident: DBIdent = attr.ib()

    class ResultSchema(PrimitivesResponseSchema):
        value = fields.String(allow_none=True)  # type: ignore  # TODO: fix


@attr.s(frozen=True)
class ActionGetSchemaNames(NonStreamAction[List[str]]):
    db_ident: DBIdent = attr.ib()

    class ResultSchema(PrimitivesResponseSchema):
        value = fields.List(fields.String())  # type: ignore  # TODO: fix


@attr.s(frozen=True)
class ActionGetTables(NonStreamAction[List[TableIdent]]):
    schema_ident: SchemaIdent = attr.ib()

    class ResultSchema(PrimitivesResponseSchema):
        value = fields.Nested(TableIdentSchema, many=True)  # type: ignore  # TODO: fix


@attr.s(frozen=True)
class ActionGetTableInfo(NonStreamAction[RawSchemaInfo]):
    table_def: TableDefinition = attr.ib()
    fetch_idx_info: bool = attr.ib()

    ResultSchema = RawSchemaInfoSchema


@attr.s(frozen=True)
class ActionIsTableExists(NonStreamAction[bool]):
    table_ident: TableIdent = attr.ib()

    class ResultSchema(PrimitivesResponseSchema):
        value = fields.Boolean()  # type: ignore  # TODO: fix

lib/dl_core/dl_core/connection_executors/qe_serializer/dba_actions.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
    FrozenSet,
    Type,
    Union,
)

from marshmallow import fields

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO
from dl_core.connection_executors.qe_serializer import dba_actions as dba_actions
from dl_core.connection_executors.qe_serializer.schema_base import BaseQEAPISchema
from dl_core.connection_executors.qe_serializer.schemas_common import (
    DBAdapterScopedRCISchema,
    DBIdentSchema,
    GenericDBAQuerySchema,
    SchemaIdentSchema,
    TableDefinitionSchema,
    TableIdentSchema,
)


class DBAdapterActionBaseSchema(BaseQEAPISchema):
    target_conn_dto = fields.Method(serialize="dump_target_conn_dto", deserialize="load_target_conn_dto")
    dba_cls = fields.Method(serialize="dump_dba_cls", deserialize="load_dba_cls")
    req_ctx_info = fields.Nested(DBAdapterScopedRCISchema)

    @property
    def allowed_dba_classes(self) -> FrozenSet[Type[CommonBaseDirectAdapter]]:
        # TODO FIX: Ensure no classes with same qualname or use FQDN
        return self.context["allowed_dba_classes"]

    def dump_dba_cls(self, act: dba_actions.ActionExecuteQuery) -> str:
        return act.dba_cls.__qualname__

    def load_dba_cls(self, value: str) -> Union[Type[CommonBaseDirectAdapter]]:
        candidate = next(filter(lambda clz: clz.__qualname__ == value, self.allowed_dba_classes), None)
        if candidate is None:
            raise ValueError(f"Can not restore DBA class from string '{value}'")
        return candidate

    def dump_target_conn_dto(self, act: dba_actions.ActionExecuteQuery) -> Dict:
        return act.target_conn_dto.to_jsonable_dict()

    def load_target_conn_dto(self, value: Dict) -> ConnTargetDTO:
        return ConnTargetDTO.from_polymorphic_jsonable_dict(value)

    def to_object(self, data: Dict[str, Any]) -> Any:
        raise NotImplementedError()


class ActionExecuteQuerySchema(DBAdapterActionBaseSchema):
    db_adapter_query = fields.Nested(GenericDBAQuerySchema)

    def to_object(self, data: Dict[str, Any]) -> dba_actions.ActionExecuteQuery:
        return dba_actions.ActionExecuteQuery(**data)


class ActionGetDBVersionSchema(DBAdapterActionBaseSchema):
    db_ident = fields.Nested(DBIdentSchema)

    def to_object(self, data: Dict[str, Any]) -> Any:
        return dba_actions.ActionGetDBVersion(**data)


class ActionTestSchema(DBAdapterActionBaseSchema):
    def to_object(self, data: Dict[str, Any]) -> dba_actions.ActionTest:
        return dba_actions.ActionTest(**data)


class ActionGetSchemaNamesSchema(DBAdapterActionBaseSchema):
    db_ident = fields.Nested(DBIdentSchema)

    def to_object(self, data: Dict[str, Any]) -> dba_actions.ActionGetSchemaNames:
        return dba_actions.ActionGetSchemaNames(**data)


class ActionGetTablesSchema(DBAdapterActionBaseSchema):
    schema_ident = fields.Nested(SchemaIdentSchema)

    def to_object(self, data: Dict[str, Any]) -> dba_actions.ActionGetTables:
        return dba_actions.ActionGetTables(**data)


class ActionGetTableInfoSchema(DBAdapterActionBaseSchema):
    table_def = fields.Nested(TableDefinitionSchema)
    fetch_idx_info = fields.Bool()

    def to_object(self, data: Dict[str, Any]) -> dba_actions.ActionGetTableInfo:
        return dba_actions.ActionGetTableInfo(**data)


class ActionIsTableExistsSchema(DBAdapterActionBaseSchema):
    table_ident = fields.Nested(TableIdentSchema)

    def to_object(self, data: Dict[str, Any]) -> dba_actions.ActionIsTableExists:
        return dba_actions.ActionIsTableExists(**data)

lib/dl_core/dl_core/connection_executors/qe_serializer/schema_actions.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    Any,
    Dict,
)

from marshmallow import (
    Schema,
    post_load,
)


class BaseQEAPISchema(Schema):
    @post_load(pass_many=False)
    def post_load(self, data: Dict, **_: Any) -> Any:
        return self.to_object(data)

    @abc.abstractmethod
    def to_object(self, data: Dict[str, Any]) -> Any:
        raise NotImplementedError()

lib/dl_core/dl_core/connection_executors/qe_serializer/schema_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from base64 import (
    b64decode,
    b64encode,
)
import json
from typing import (
    Any,
    Dict,
    List,
    Optional,
    Sequence,
    Tuple,
)

from marshmallow import fields
from marshmallow_oneofschema import OneOfSchema
from multidict import CIMultiDict
import sqlalchemy as sa
from sqlalchemy import (
    Column,
    MetaData,
    Table,
)
from sqlalchemy.ext.serializer import dumps as sa_dumps
from sqlalchemy.ext.serializer import loads

from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
from dl_core.connection_executors.qe_serializer.schema_base import BaseQEAPISchema
from dl_core.connection_models import (
    DBIdent,
    SATextTableDefinition,
    SchemaIdent,
    TableDefinition,
    TableIdent,
)
from dl_core.enums import QueryExecutorMode
from dl_core.serialization import (
    RedisDatalensDataJSONDecoder,
    RedisDatalensDataJSONEncoder,
)
from dl_utils.utils import get_type_full_name


class SATextField(fields.String):
    def _serialize(self, value: Any, attr: Any, obj: Any, **kwargs: Any) -> Optional[str]:
        assert isinstance(value, sa.sql.elements.TextClause)
        text_value = value.text
        return super()._serialize(text_value, attr, obj, **kwargs)

    def _deserialize(self, value: Any, attr: Any, data: Any, **kwargs: Any) -> Any:
        text_value = super()._deserialize(value, attr, data, **kwargs)
        assert text_value is not None
        return sa.sql.elements.TextClause(text_value)


class DBAdapterQueryStrSchema(BaseQEAPISchema):
    chunk_size = fields.Integer(allow_none=True)
    query = fields.String()
    db_name = fields.String(allow_none=True)
    disable_streaming = fields.Boolean()
    connector_specific_params = fields.Method(
        allow_none=True, dump_default=None, serialize="dump_conn_params", deserialize="load_conn_params"
    )
    is_dashsql_query = fields.Boolean()

    def to_object(self, data: Dict) -> DBAdapterQuery:
        return DBAdapterQuery(
            query=data["query"],
            chunk_size=data["chunk_size"],
            db_name=data["db_name"],
            disable_streaming=data["disable_streaming"],
            connector_specific_params=data["connector_specific_params"],
            is_dashsql_query=data["is_dashsql_query"],
        )

    def dump_conn_params(self, dba_query: DBAdapterQuery) -> Optional[dict]:
        conn_params = dba_query.connector_specific_params
        if conn_params is not None:
            for k, v in conn_params.items():
                conn_params[k] = json.dumps(v, cls=RedisDatalensDataJSONEncoder)
        return conn_params

    def load_conn_params(self, conn_params: Optional[dict]) -> Optional[Dict[str, Any]]:
        if conn_params is not None:
            for k, v in conn_params.items():
                conn_params[k] = json.loads(v, cls=RedisDatalensDataJSONDecoder)
        return conn_params


# noinspection PyMethodMayBeStatic
class DBAdapterQuerySQLASchema(DBAdapterQueryStrSchema):
    query = fields.Method(serialize="dump_query", deserialize="load_query")  # type: ignore  # TODO: fix
    tables = fields.Method(serialize="dump_tables", deserialize="load_tables")

    def dump_query(self, dba_query: DBAdapterQuery) -> str:
        query = dba_query.query
        return b64encode(sa_dumps(query)).decode()

    def load_query(self, data: str) -> str:
        return data

    def dump_tables(self, dba_query: DBAdapterQuery) -> Dict[str, List[str]]:
        tables_dict = {}

        query = dba_query.query
        if isinstance(query, str):
            raise ValueError("DBAdapterQuerySQLASchema should be used only for plain-text queries serialization.")

        for table in getattr(query, "froms", []):
            if isinstance(table, sa.Table):
                tables_dict[table.name] = [c.name for c in table.c]

        return tables_dict

    def load_tables(self, data: dict) -> MetaData:
        metadata = MetaData()

        for table_name in data:
            Table(table_name, metadata, *[Column(col_name) for col_name in data[table_name]])

        return metadata

    def to_object(self, data: Dict[str, Any]) -> DBAdapterQuery:
        patched_data = dict(data)
        metadata = patched_data.pop("tables")
        dumped_sa_query = patched_data.pop("query")
        patched_data["query"] = loads(b64decode(dumped_sa_query), metadata=metadata)
        return super().to_object(patched_data)


class GenericDBAQuerySchema(OneOfSchema):
    type_field = "mode"
    type_schemas = {
        t.name: s
        for t, s in {
            QueryExecutorMode.sqla_dump: DBAdapterQuerySQLASchema,
            QueryExecutorMode.text: DBAdapterQueryStrSchema,
        }.items()
    }

    def get_obj_type(self, obj: DBAdapterQuery) -> str:
        if isinstance(obj.query, str):
            return QueryExecutorMode.text.name
        else:
            return QueryExecutorMode.sqla_dump.name


# TODO FIX: Validation
class CIMultiDictField(fields.Field):
    def _serialize(self, value: CIMultiDict, attr: Any, obj: Any, **kwargs: Any) -> Tuple[Tuple[str, str], ...]:
        return tuple((item[0], item[1]) for item in value.items())

    def _deserialize(self, value: Sequence[Tuple[str, str]], attr: Any, data: Any, **kwargs: Any) -> CIMultiDict:
        return CIMultiDict(value)


class DBAdapterScopedRCISchema(BaseQEAPISchema):
    request_id = fields.String(allow_none=True)
    user_name = fields.String(allow_none=True)
    x_dl_debug_mode = fields.Boolean(allow_none=True)
    client_ip = fields.String(allow_none=True)

    def to_object(self, data: Dict[str, Any]) -> DBAdapterScopedRCI:
        return DBAdapterScopedRCI(
            request_id=data["request_id"],
            user_name=data["user_name"],
            x_dl_debug_mode=data["x_dl_debug_mode"],
            client_ip=data["client_ip"],
        )


class DBIdentSchema(BaseQEAPISchema):
    db_name = fields.String(allow_none=True)

    def to_object(self, data: Dict[str, Any]) -> DBIdent:
        return DBIdent(**data)


class SchemaIdentSchema(BaseQEAPISchema):
    schema_name = fields.String(allow_none=True)
    db_name = fields.String(allow_none=True)

    def to_object(self, data: Dict[str, Any]) -> Any:
        return SchemaIdent(**data)


class TableIdentSchema(BaseQEAPISchema):
    table_name = fields.String()
    db_name = fields.String(allow_none=True)
    schema_name = fields.String(allow_none=True)

    def to_object(self, data: Dict[str, Any]) -> TableIdent:
        return TableIdent(**data)


class SATextTableDefinitionSchema(BaseQEAPISchema):
    text = SATextField()

    def to_object(self, data: Dict[str, Any]) -> SATextTableDefinition:
        return SATextTableDefinition(**data)


class TableDefinitionSchema(OneOfSchema):
    type_field = "mode"
    supported_types = {
        TableIdent: TableIdentSchema,
        SATextTableDefinition: SATextTableDefinitionSchema,
    }
    type_schemas = {type_obj.def_type.name: schema_cls for type_obj, schema_cls in supported_types.items()}  # type: ignore  # TODO: fix

    def get_obj_type(self, obj: TableDefinition) -> str:
        if type(obj) in self.supported_types:
            return obj.def_type.name

        raise TypeError(f"Class is not supported by TableDefinitionSchema: {get_type_full_name(type(obj))}")

lib/dl_core/dl_core/connection_executors/qe_serializer/schemas_common.py
datalens-tech/datalens-backend
from __future__ import annotations

import traceback
from typing import Any

from marshmallow import (
    Schema,
    fields,
    post_load,
)
from marshmallow_oneofschema import OneOfSchema

from dl_constants.exc import DLBaseException
from dl_core.connection_executors.models.exc import QueryExecutorException


class DLExcSchema(Schema):
    dl_exec_dict = fields.Function(serialize=lambda exc: exc.to_jsonable_dict(), deserialize=lambda o: o)

    @post_load(pass_many=False)
    def post_load(self, data: dict, **_) -> DLBaseException:  # type: ignore
        return DLBaseException.from_jsonable_dict(data["dl_exec_dict"])


class OtherExcSchema(Schema):
    exc_type_name = fields.Function(serialize=lambda err: type(err).__qualname__, deserialize=lambda obj: obj)
    exc_msg = fields.Function(serialize=lambda err: str(err), deserialize=lambda obj: obj)
    exc_tb = fields.Function(
        serialize=lambda err: (traceback.format_tb(err.__traceback__) if getattr(err, "__traceback__", None) else None),
        deserialize=lambda obj: obj,
    )

    @post_load(pass_many=False)
    def post_load(self, data: dict, **_: dict) -> QueryExecutorException:  # type: ignore
        tb = "".join(data.get("exc_tb") or ())
        return QueryExecutorException(f"{data['exc_type_name']}: {data['exc_msg']}\n{tb}")


class GenericExcSchema(OneOfSchema):
    type_schemas = {
        "dl_exc": DLExcSchema,
        "other_exc": OtherExcSchema,
    }

    def get_obj_type(self, obj: Any) -> str:
        if isinstance(obj, DLBaseException):
            return "dl_exc"
        elif isinstance(obj, Exception):
            return "other_exc"
        raise TypeError("Can not serialize non-exception")

lib/dl_core/dl_core/connection_executors/qe_serializer/schemas_exc.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
)

from marshmallow import (
    fields,
    pre_dump,
)

from dl_constants.enums import IndexKind
from dl_core.connection_executors.models.db_adapter_data import (
    RawColumnInfo,
    RawIndexInfo,
    RawSchemaInfo,
)
from dl_core.connection_executors.qe_serializer.schema_base import BaseQEAPISchema
from dl_core.db.native_type_schema import OneOfNativeTypeSchemaBase


class RawColumnInfoSchema(BaseQEAPISchema):
    name = fields.String()
    title = fields.String(allow_none=True)
    nullable = fields.Boolean(allow_none=True)  # deprecated in favor of native_type.nullable
    # Note: using the non-transition schema in the hopes that RQE and app codebases always match.
    native_type = fields.Nested(OneOfNativeTypeSchemaBase)

    def to_object(self, data: Dict[str, Any]) -> RawColumnInfo:
        return RawColumnInfo(**data)


class IndexInfoSchema(BaseQEAPISchema):
    columns = fields.List(fields.String())
    kind = fields.Enum(IndexKind, allow_none=True)
    unique = fields.Boolean()

    def to_object(self, data: Dict[str, Any]) -> RawIndexInfo:
        columns = tuple(data.pop("columns"))
        return RawIndexInfo(columns=columns, **data)


class RawSchemaInfoSchema(BaseQEAPISchema):
    columns = fields.Nested(RawColumnInfoSchema, many=True)
    indexes = fields.Nested(IndexInfoSchema, many=True, allow_none=True)

    def to_object(self, data: Dict[str, Any]) -> RawSchemaInfo:
        return RawSchemaInfo(**data)


class PrimitivesResponseSchema(BaseQEAPISchema):
    value = fields.Raw()

    @pre_dump(pass_many=False)
    def wrap(self, data: Any, **__: Any) -> Dict:
        return {"value": data}

    def to_object(self, data: Dict[str, Any]) -> Any:
        return data["value"]

lib/dl_core/dl_core/connection_executors/qe_serializer/schemas_responses.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    ClassVar,
    Dict,
    FrozenSet,
    List,
    Optional,
    Type,
    Union,
)

from marshmallow import Schema

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.models.db_adapter_data import RawSchemaInfo
from dl_core.connection_executors.qe_serializer import (
    dba_actions,
    schema_actions,
    schemas_exc,
)


ResponseTypes = Union[RawSchemaInfo, List[str], Optional[str], None, bool, int]


class ActionSerializer:
    MAP_ACT_TYPE_SCHEMA_CLS: ClassVar[Dict[Type[dba_actions.RemoteDBAdapterAction], Type[Schema]]] = {
        dba_actions.ActionExecuteQuery: schema_actions.ActionExecuteQuerySchema,
        dba_actions.ActionTest: schema_actions.ActionTestSchema,
        dba_actions.ActionGetDBVersion: schema_actions.ActionGetDBVersionSchema,
        dba_actions.ActionGetSchemaNames: schema_actions.ActionGetSchemaNamesSchema,
        dba_actions.ActionGetTables: schema_actions.ActionGetTablesSchema,
        dba_actions.ActionGetTableInfo: schema_actions.ActionGetTableInfoSchema,
        dba_actions.ActionIsTableExists: schema_actions.ActionIsTableExistsSchema,
    }

    EXC_SCHEMA_CLS = schemas_exc.GenericExcSchema

    def serialize_action(self, obj: dba_actions.RemoteDBAdapterAction) -> Dict:
        schema = self.MAP_ACT_TYPE_SCHEMA_CLS[type(obj)]()
        result = schema.dump(obj)
        result["type"] = type(obj).__qualname__
        return result

    def deserialize_action(
        self,
        data: Dict,
        allowed_dba_classes: FrozenSet[Type[CommonBaseDirectAdapter]],
    ) -> dba_actions.RemoteDBAdapterAction:
        action_cls_qualname = data.pop("type", None)
        action_cls = next(
            filter(lambda clz: clz.__qualname__ == action_cls_qualname, self.MAP_ACT_TYPE_SCHEMA_CLS.keys()), None
        )
        if action_cls is None:
            raise TypeError(f"Action {action_cls_qualname} is not listed in serializable actions")

        schema = self.MAP_ACT_TYPE_SCHEMA_CLS[action_cls](context=dict(allowed_dba_classes=allowed_dba_classes))
        obj = schema.load(data)

        if not isinstance(obj, action_cls):
            raise TypeError(f"Unexpected type of action after deserialization: {type(obj)} instead of {action_cls}")

        return obj

    def serialize_exc(self, exc: Exception) -> Dict:
        return self.EXC_SCHEMA_CLS().dump(exc)

    def deserialize_exc(self, data: Dict) -> Exception:
        # TODO FIX: Validate type
        return self.EXC_SCHEMA_CLS().load(data)

lib/dl_core/dl_core/connection_executors/qe_serializer/serializer.py
datalens-tech/datalens-backend
from __future__ import annotations

import argparse
import logging
import pickle
import sys
from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
    Type,
    Union,
)

from aiohttp import web
from aiohttp.typedefs import Handler

from dl_api_commons.aio.middlewares.request_bootstrap import RequestBootstrap
from dl_api_commons.aio.middlewares.request_id import RequestId
from dl_api_commons.aio.server_header import ServerHeader
from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_configs.env_var_definitions import (
    jaeger_service_name_env_aware,
    use_jaeger_tracer,
)
from dl_configs.settings_loaders.loader_env import load_settings_from_env_with_fallback
from dl_core.connection_executors.adapters.adapters_base import SyncDirectDBAdapter
from dl_core.connection_executors.adapters.async_adapters_base import (
    AsyncDBAdapter,
    AsyncDirectDBAdapter,
)
from dl_core.connection_executors.adapters.async_adapters_sync_wrapper import AsyncWrapperForSyncAdapter
from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.models.constants import (
    HEADER_BODY_SIGNATURE,
    HEADER_REQUEST_ID,
)
from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
from dl_core.connection_executors.qe_serializer import (
    ActionSerializer,
    ResponseTypes,
)
from dl_core.connection_executors.qe_serializer import dba_actions as act
from dl_core.connection_executors.remote_query_executor.commons import (
    DEFAULT_CHUNK_SIZE,
    SUPPORTED_ADAPTER_CLS,
)
from dl_core.connection_executors.remote_query_executor.crypto import get_hmac_hex_digest
from dl_core.connection_executors.remote_query_executor.error_handler_rqe import RQEErrorHandler
from dl_core.connection_executors.remote_query_executor.settings import RQESettings
from dl_core.enums import RQEEventType
from dl_core.loader import (
    CoreLibraryConfig,
    load_core_lib,
)
from dl_core.logging_config import configure_logging
from dl_utils.aio import ContextVarExecutor


if TYPE_CHECKING:
    from dl_core.connection_executors.models.connection_target_dto_base import ConnTargetDTO


LOGGER = logging.getLogger(__name__)


class BaseView(web.View):
    @property
    def tpe(self) -> ContextVarExecutor:
        return self.request.app["tpe"]


def adapter_factory(
    target_conn_dto: ConnTargetDTO,
    dba_cls: Type[CommonBaseDirectAdapter],
    req_ctx_info: DBAdapterScopedRCI,
    tpe: ContextVarExecutor,
) -> AsyncDBAdapter:
    default_chunk_size = DEFAULT_CHUNK_SIZE

    if issubclass(dba_cls, SyncDirectDBAdapter):
        sync_dba = dba_cls.create(
            target_dto=target_conn_dto, req_ctx_info=req_ctx_info, default_chunk_size=default_chunk_size
        )

        return AsyncWrapperForSyncAdapter(
            tpe=tpe,
            sync_adapter=sync_dba,
        )

    elif issubclass(dba_cls, AsyncDirectDBAdapter):
        return dba_cls.create(
            target_dto=target_conn_dto, req_ctx_info=req_ctx_info, default_chunk_size=default_chunk_size
        )

    else:
        raise ValueError(f"DBA class {dba_cls} is not supported in EQQ")


class PingView(BaseView):
    async def get(self) -> web.Response:
        return web.json_response(
            dict(
                result="PONG",
            )
        )


class ActionHandlingView(BaseView):
    async def get_action(self) -> act.RemoteDBAdapterAction:
        raw_body = await self.request.json()
        action = ActionSerializer().deserialize_action(raw_body, allowed_dba_classes=SUPPORTED_ADAPTER_CLS)  # type: ignore
        return action

    @staticmethod
    def serialize_event(event: RQEEventType, data: Any) -> bytes:
        return pickle.dumps((event.value, data))

    async def handle_query_action(
        self,
        dba: AsyncDBAdapter,
        dba_query: DBAdapterQuery,
    ) -> Union[web.StreamResponse, web.Response]:
        try:
            result = await dba.execute(dba_query)
        except Exception:
            LOGGER.exception("Exception during execution")
            raise
        else:
            response = web.StreamResponse()
            response.enable_chunked_encoding()
            # TODO FIX: Use schema/custom JSON encoder

            await response.prepare(self.request)
            # After this point, `error_handling_middleware` does not work, so
            # send errors into the events stream:
            try:
                await response.write(self.serialize_event(RQEEventType.raw_cursor_info, result.raw_cursor_info))

                async for raw_chunk in result.raw_chunk_generator:
                    await response.write(self.serialize_event(RQEEventType.raw_chunk, raw_chunk))

            except Exception as err:
                await response.write(
                    self.serialize_event(RQEEventType.error_dump, ActionSerializer().serialize_exc(err))
                )
            # Proper chunked close should also indicate this, but making an explicit end-of-stream at this
            await response.write(self.serialize_event(RQEEventType.finished, None))

        return response

    async def execute_non_streamed_action(
        self,
        dba: AsyncDBAdapter,
        action: act.NonStreamAction,
    ) -> ResponseTypes:
        if isinstance(action, act.ActionTest):
            await dba.test()
            return None

        elif isinstance(action, act.ActionGetDBVersion):
            return await dba.get_db_version(db_ident=action.db_ident)

        elif isinstance(action, act.ActionGetSchemaNames):
            return await dba.get_schema_names(db_ident=action.db_ident)

        elif isinstance(action, act.ActionGetTables):
            return await dba.get_tables(schema_ident=action.schema_ident)  # type: ignore

        elif isinstance(action, act.ActionGetTableInfo):
            return await dba.get_table_info(table_def=action.table_def, fetch_idx_info=action.fetch_idx_info)

        elif isinstance(action, act.ActionIsTableExists):
            return await dba.is_table_exists(table_ident=action.table_ident)

        else:
            raise NotImplementedError(f"Action {action} is not implemented in QE")

    async def post(self) -> Union[web.Response, web.StreamResponse]:
        action = await self.get_action()
        LOGGER.info("Got QE action request: %s", action)

        # Creating adapter
        LOGGER.info("Creating DBA")
        adapter = adapter_factory(
            target_conn_dto=action.target_conn_dto,
            req_ctx_info=action.req_ctx_info,
            tpe=self.tpe,
            dba_cls=action.dba_cls,
        )
        async with adapter:
            LOGGER.info("DBA for action was created: %s", adapter)

            if isinstance(action, act.ActionExecuteQuery):
                return await self.handle_query_action(adapter, action.db_adapter_query)
            elif isinstance(action, act.NonStreamAction):
                result = await self.execute_non_streamed_action(adapter, action)
                resp_data = action.serialize_response(result)
                return web.json_response(resp_data)
            else:
                raise NotImplementedError(f"Action {action} is not implemented in QE")


def body_signature_validation_middleware(hmac_key: bytes) -> AIOHTTPMiddleware:
    @web.middleware
    async def actual_middleware(request: web.Request, handler: Handler) -> web.StreamResponse:
        if not hmac_key:  # do not consider an empty hmac key as valid.
            raise Exception("body_signature_validation_middleware: no hmac_key.")

        if request.method in ("HEAD", "OPTIONS", "GET"):
            return await handler(request)

        body_bytes = await request.read()
        expected_signature = get_hmac_hex_digest(body_bytes, secret_key=hmac_key)
        signature_str_from_header = request.headers.get(HEADER_BODY_SIGNATURE)

        if expected_signature != signature_str_from_header:
            raise web.HTTPForbidden(reason="Invalid signature")

        return await handler(request)

    return actual_middleware


def create_async_qe_app(hmac_key: bytes) -> web.Application:
    req_id_service = RequestId(
        header_name=HEADER_REQUEST_ID,
        accept_logging_ctx=True,
    )
    error_handler = RQEErrorHandler(
        sentry_app_name_tag=None,
    )
    app = web.Application(
        middlewares=[
            RequestBootstrap(
                req_id_service=req_id_service,
                error_handler=error_handler,
            ).middleware,
            # TODO FIX: Add profiling middleware.
            body_signature_validation_middleware(hmac_key=hmac_key),
        ]
    )
    app.on_response_prepare.append(req_id_service.on_response_prepare)
    ServerHeader("DataLens QE").add_signal_handlers(app)

    # TODO FIX: Close on app exit
    app["tpe"] = ContextVarExecutor()

    app.router.add_route("get", "/ping", PingView)
    app.router.add_route("*", "/execute_action", ActionHandlingView)

    return app


def async_qe_main() -> None:
    configure_logging(
        app_name="rqe-async",
        app_prefix=None,  # not useful with `append_local_req_id=False`.
        use_jaeger_tracer=use_jaeger_tracer(),
        jaeger_service_name=jaeger_service_name_env_aware("bi-rqe-async"),
    )

    settings = load_settings_from_env_with_fallback(RQESettings)
    load_core_lib(core_lib_config=CoreLibraryConfig(core_connector_ep_names=settings.CORE_CONNECTOR_WHITELIST))

    try:
        parser = argparse.ArgumentParser(description="Process some integers.")
        parser.add_argument("--host", type=str)
        parser.add_argument("--port", type=int)
        args = parser.parse_args()
    except Exception as err:
        LOGGER.exception("rqe-async args error: %r", err)
        sys.exit(-1)
        raise  # just-in-case

    hmac_key = settings.RQE_SECRET_KEY
    if hmac_key is None:
        raise Exception("No `hmac_key` set.")
    web.run_app(create_async_qe_app(hmac_key.encode()), host=args.host, port=args.port, print=LOGGER.info)

lib/dl_core/dl_core/connection_executors/remote_query_executor/app_async.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import pickle
from typing import (
    TYPE_CHECKING,
    Any,
    Iterable,
    Optional,
    Tuple,
    Union,
)

import attr
import flask.views
from werkzeug.exceptions import (
    Forbidden,
    HTTPException,
)

from dl_api_commons.flask.middlewares.aio_event_loop_middleware import AIOEventLoopMiddleware
from dl_api_commons.flask.middlewares.context_var_middleware import ContextVarMiddleware
from dl_api_commons.flask.middlewares.logging_context import RequestLoggingContextControllerMiddleWare
from dl_api_commons.flask.middlewares.request_id import RequestIDService
from dl_api_commons.flask.middlewares.tracing import TracingMiddleware
from dl_configs.env_var_definitions import (
    jaeger_service_name_env_aware,
    use_jaeger_tracer,
)
from dl_configs.settings_loaders.loader_env import load_settings_from_env_with_fallback
from dl_core import profiling_middleware
from dl_core.connection_executors.adapters.adapters_base import SyncDirectDBAdapter
from dl_core.connection_executors.models.constants import HEADER_BODY_SIGNATURE
from dl_core.connection_executors.qe_serializer import (
    ActionSerializer,
    ResponseTypes,
)
from dl_core.connection_executors.qe_serializer import dba_actions as act
from dl_core.connection_executors.remote_query_executor.commons import (
    DEFAULT_CHUNK_SIZE,
    SUPPORTED_ADAPTER_CLS,
)
from dl_core.connection_executors.remote_query_executor.crypto import get_hmac_hex_digest
from dl_core.connection_executors.remote_query_executor.settings import RQESettings
from dl_core.enums import RQEEventType
from dl_core.loader import (
    CoreLibraryConfig,
    load_core_lib,
)
from dl_core.logging_config import hook_configure_logging as _hook_configure_logging


if TYPE_CHECKING:
    from dl_core.connection_executors.adapters.adapters_base import DBAdapterQueryResult


LOGGER = logging.getLogger(__name__)


def chunked_wrap(iterable: Iterable[Union[bytes, str]]) -> Iterable[bytes]:
    """
    `Transfer-Encoding: Chunked` wrap.

    See also: https://rhodesmill.org/brandon/2013/chunked-wsgi/
    """
    for item in iterable:
        if not item:
            continue
        if isinstance(item, str):
            item = item.encode("utf-8")
        yield b"%x" % (len(item),)
        yield b"\r\n"
        yield item
        yield b"\r\n"
    yield b"0\r\n\r\n"


class ActionHandlingView(flask.views.View):
    methods = ["POST"]

    def get_action(self) -> act.RemoteDBAdapterAction:
        return ActionSerializer().deserialize_action(flask.request.json, allowed_dba_classes=SUPPORTED_ADAPTER_CLS)  # type: ignore

    def execute_non_streamed_action(
        self,
        dba: SyncDirectDBAdapter,
        action: act.NonStreamAction,
    ) -> ResponseTypes:
        if isinstance(action, act.ActionTest):
            return dba.test()

        elif isinstance(action, act.ActionGetDBVersion):
            return dba.get_db_version(db_ident=action.db_ident)

        elif isinstance(action, act.ActionGetSchemaNames):
            return dba.get_schema_names(db_ident=action.db_ident)

        elif isinstance(action, act.ActionGetTables):
            return dba.get_tables(schema_ident=action.schema_ident)  # type: ignore

        elif isinstance(action, act.ActionGetTableInfo):
            return dba.get_table_info(table_def=action.table_def, fetch_idx_info=action.fetch_idx_info)

        elif isinstance(action, act.ActionIsTableExists):
            return dba.is_table_exists(table_ident=action.table_ident)

        else:
            raise NotImplementedError(f"Action {action} is not implemented in QE")

    @staticmethod
    def try_close_dba(dba: SyncDirectDBAdapter) -> None:
        try:
            dba.close()
        except Exception:
            LOGGER.exception("Exception during DBA closing")

    @staticmethod
    def serialize_event(event: RQEEventType, data: Any) -> bytes:
        return pickle.dumps((event.value, data))

    def response_events_gen(self, db_result: "DBAdapterQueryResult", dba: SyncDirectDBAdapter) -> Iterable[bytes]:
        try:
            yield self.serialize_event(RQEEventType.raw_cursor_info, db_result.cursor_info)

            for raw_chunk in db_result.data_chunks:
                yield self.serialize_event(RQEEventType.raw_chunk, raw_chunk)

        except Exception as err:
            yield self.serialize_event(RQEEventType.error_dump, ActionSerializer().serialize_exc(err))

        finally:
            dba.close()

        yield self.serialize_event(RQEEventType.finished, None)

    def execute_execute_action(
        self,
        action: act.ActionExecuteQuery,
    ) -> flask.Response:
        dba = self.create_dba_for_action(action)
        try:
            db_result = dba.execute(action.db_adapter_query)
        except Exception:
            # noinspection PyBroadException
            self.try_close_dba(dba)
            raise

        return flask.Response(
            chunked_wrap(self.response_events_gen(db_result=db_result, dba=dba)),
            headers={
                "Transfer-Encoding": "Chunked",
            },
        )

    @staticmethod
    def create_dba_for_action(action: act.RemoteDBAdapterAction) -> SyncDirectDBAdapter:
        LOGGER.info("Creating DBA")
        default_chunk_size = DEFAULT_CHUNK_SIZE
        dba_cls = action.dba_cls
        if not issubclass(dba_cls, SyncDirectDBAdapter):
            raise TypeError(f"Can not create DBA type '{dba_cls}' in sync RQE")

        dba = dba_cls.create(
            target_dto=action.target_conn_dto,
            req_ctx_info=action.req_ctx_info,
            default_chunk_size=default_chunk_size,
        )
        LOGGER.info("DBA for action was created: %s", dba)
        return dba

    def dispatch_request(self) -> flask.Response:  # type: ignore
        action = self.get_action()
        LOGGER.info("Got QE action request: %s", action)

        if isinstance(action, act.NonStreamAction):
            dba = self.create_dba_for_action(action)
            try:
                result = self.execute_non_streamed_action(dba, action)
                return flask.jsonify(action.serialize_response(result))
            finally:
                dba.close()
        elif isinstance(action, act.ActionExecuteQuery):
            return self.execute_execute_action(action)
        else:
            raise NotImplementedError(f"Action {action} is not implemented in QE")


@attr.s
class BodySignatureValidator:
    hmac_key: bytes = attr.ib()

    def validate_request_body(self) -> None:
        if flask.request.method in ("HEAD", "OPTIONS", "GET"):  # no body to validate.
            return

        # For import-test reasons, can't verify this when getting it;
        # but allowing requests when the key is empty is too dangerous.
        if not self.hmac_key:
            raise Exception("validate_request_body: no hmac_key")

        body_bytes = flask.request.get_data()
        expected_signature = get_hmac_hex_digest(body_bytes, secret_key=self.hmac_key)
        signature_str_from_header = flask.request.headers.get(HEADER_BODY_SIGNATURE)

        if expected_signature != signature_str_from_header:
            raise Forbidden("Invalid signature")

    def set_up(self, app: flask.Flask) -> None:
        app.before_request(self.validate_request_body)


def ping_view() -> flask.Response:
    return flask.jsonify(dict(result="PONG"))


def hook_init_logging(
    app: Any,
    app_name: str = "rqe-sync",
    jaeger_service_name: str = "bi-rqe-sync",
    app_prefix: Optional[str] = None,
    **kwargs: Any,
) -> None:
    return _hook_configure_logging(
        app=app,
        app_name=app_name,
        app_prefix=app_prefix,
        use_jaeger_tracer=use_jaeger_tracer(),
        jaeger_service_name=jaeger_service_name_env_aware(jaeger_service_name),
        **kwargs,
    )


def _handle_exception(err: Exception) -> Tuple[flask.Response, int]:
    LOGGER.exception("Exception occurred in sync RQE")
    if isinstance(err, HTTPException):
        return flask.jsonify(ActionSerializer().serialize_exc(err)), err.code or 500
    else:
        return flask.jsonify(ActionSerializer().serialize_exc(err)), 500


def create_sync_app() -> flask.Flask:
    settings = load_settings_from_env_with_fallback(RQESettings)
    hmac_key = settings.RQE_SECRET_KEY
    if hmac_key is None:
        raise Exception("No `hmac_key` set.")

    load_core_lib(core_lib_config=CoreLibraryConfig(core_connector_ep_names=settings.CORE_CONNECTOR_WHITELIST))

    app = flask.Flask(__name__)
    TracingMiddleware(
        url_prefix_exclude=(
            "/ping",
            "/unistat",
            "/metrics",
        ),
    ).wrap_flask_app(app)
    ContextVarMiddleware().wrap_flask_app(app)
    AIOEventLoopMiddleware().wrap_flask_app(app)

    RequestLoggingContextControllerMiddleWare().set_up(app)
    RequestIDService(
        request_id_app_prefix=None,
        append_local_req_id=False,
        accept_logging_context=True,
    ).set_up(app)
    profiling_middleware.set_up(app, accept_outer_stages=True)
    BodySignatureValidator(
        hmac_key=hmac_key.encode(),
    ).set_up(app)

    app.add_url_rule("/ping", view_func=ping_view)
    app.add_url_rule("/execute_action", view_func=ActionHandlingView.as_view("execute_action"))

    app.errorhandler(Exception)(_handle_exception)

    return app

lib/dl_core/dl_core/connection_executors/remote_query_executor/app_sync.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Set,
    Type,
)

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter


DEFAULT_CHUNK_SIZE: int = 1_000

SUPPORTED_ADAPTER_CLS: Set[Type[CommonBaseDirectAdapter]] = set()


def register_adapter_class(adapter_cls: Type[CommonBaseDirectAdapter]) -> None:
    SUPPORTED_ADAPTER_CLS.add(adapter_cls)

lib/dl_core/dl_core/connection_executors/remote_query_executor/commons.py
datalens-tech/datalens-backend
from __future__ import annotations

import hashlib
import hmac
from typing import Union


def get_hmac_hex_digest(target: Union[bytes, str], secret_key: bytes) -> str:
    return hmac.new(
        key=secret_key, msg=target if isinstance(target, bytes) else target.encode(), digestmod=hashlib.sha256
    ).hexdigest()

lib/dl_core/dl_core/connection_executors/remote_query_executor/crypto.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging

from aiohttp import web

from dl_api_commons.aio.middlewares.error_handling_outer import (
    AIOHTTPErrorHandler,
    ErrorData,
    ErrorLevel,
)
from dl_core.connection_executors.qe_serializer import ActionSerializer


LOGGER = logging.getLogger(__name__)


class RQEErrorHandler(AIOHTTPErrorHandler):
    def _classify_error(self, err: Exception, request: web.Request) -> ErrorData:
        if isinstance(err, web.HTTPException):
            return ErrorData(
                err.status,
                http_reason=err.reason,
                response_body=ActionSerializer().serialize_exc(err),
                level=ErrorLevel.info,
            )
        else:
            LOGGER.exception("RQE-async handler error: %r", err)
            return ErrorData(
                500,
                http_reason="Internal server error",
                response_body=ActionSerializer().serialize_exc(err),
                level=ErrorLevel.error,
            )

lib/dl_core/dl_core/connection_executors/remote_query_executor/error_handler_rqe.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_configs.settings_loaders.meta_definition import s_attrib
from dl_configs.utils import split_by_comma


@attr.s(frozen=True)
class RQESettings:
    CORE_CONNECTOR_WHITELIST: Optional[list[str]] = s_attrib(  # type: ignore
        "CORE_CONNECTOR_WHITELIST",
        env_var_converter=lambda s: list(split_by_comma(s)),
        missing=None,
    )
    RQE_SECRET_KEY: Optional[str] = s_attrib("RQE_SECRET_KEY", missing=None, sensitive=True)  # type: ignore

lib/dl_core/dl_core/connection_executors/remote_query_executor/settings.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    Iterable,
    List,
    Optional,
    Sequence,
)

import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_core.connection_executors import ConnExecutorQuery
from dl_core.db import SchemaInfo


if TYPE_CHECKING:
    from dl_core.connection_models.common_models import (
        DBIdent,
        SchemaIdent,
        TableDefinition,
        TableIdent,
    )
    from dl_core.connection_models.dto_defs import ConnDTO


LOGGER = logging.getLogger(__name__)


@attr.s
class SyncExecutionResult:
    cursor_info: dict = attr.ib()
    result: Iterable[Sequence] = attr.ib()

    def get_all(self) -> Sequence[Sequence]:
        """
        :return: All fetched data
        """
        result: List[Sequence] = []
        for chunk in self.result:
            result.extend(chunk)

        return result


class SyncConnExecutorBase(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def initialize(self) -> None:
        pass

    @abc.abstractmethod
    def close(self) -> None:
        pass

    @abc.abstractmethod
    def execute(self, query: ConnExecutorQuery) -> SyncExecutionResult:
        pass

    @abc.abstractmethod
    def test(self) -> None:
        pass

    @abc.abstractmethod
    def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        pass

    @abc.abstractmethod
    def get_schema_names(self, db_ident: DBIdent) -> List[str]:
        pass

    @abc.abstractmethod
    def get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        pass

    @abc.abstractmethod
    def get_table_schema_info(self, table_def: TableDefinition) -> SchemaInfo:
        pass

    @abc.abstractmethod
    def is_table_exists(self, table_ident: TableIdent) -> bool:
        pass

    @abc.abstractmethod
    def is_conn_dto_equals(self, another: ConnDTO) -> bool:
        pass

    @abc.abstractmethod
    def is_context_info_equals(self, another: RequestContextInfo) -> bool:
        pass

lib/dl_core/dl_core/connection_executors/sync_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import functools
from typing import (
    TYPE_CHECKING,
    Any,
    Awaitable,
    Callable,
    Generator,
    List,
    Optional,
    Type,
    TypeVar,
    overload,
)

import attr
from typing_extensions import Literal

from dl_api_commons.base_models import RequestContextInfo
from dl_core.connection_executors import ConnExecutorQuery
from dl_core.connection_executors.adapters.adapters_base import SyncDirectDBAdapter
from dl_core.connection_executors.async_base import (
    AsyncConnExecutorBase,
    AsyncExecutionResult,
)
from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor
from dl_core.connection_executors.models.db_adapter_data import (
    ExecutionStepCursorInfo,
    ExecutionStepDataChunk,
)
from dl_core.connection_executors.sync_base import (
    SyncConnExecutorBase,
    SyncExecutionResult,
)
from dl_core.db import SchemaInfo


if TYPE_CHECKING:
    from dl_constants.types import TBIDataTable
    from dl_core.connection_models.common_models import (
        DBIdent,
        SchemaIdent,
        TableDefinition,
        TableIdent,
    )
    from dl_core.connection_models.dto_defs import ConnDTO


_RET_TV = TypeVar("_RET_TV")


def init_required(wrapped: Callable[..., _RET_TV]) -> Callable[..., _RET_TV]:
    @functools.wraps(wrapped)
    def wrapper(self: "SyncWrapperForAsyncConnExecutor", *args: Any, **kwargs: Any) -> _RET_TV:
        if not self.is_initialized:
            self.initialize()
        return wrapped(self, *args, **kwargs)

    return wrapper


_CORO_RET_TV = TypeVar("_CORO_RET_TV")


@attr.s(cmp=False, hash=False)
class SyncWrapperForAsyncConnExecutor(SyncConnExecutorBase):
    _async_conn_executor: AsyncConnExecutorBase = attr.ib()
    _loop: asyncio.AbstractEventLoop = attr.ib()

    def _await_sync(self, coro: Awaitable[_CORO_RET_TV]) -> _CORO_RET_TV:
        return self._loop.run_until_complete(coro)

    def is_conn_dto_equals(self, another: ConnDTO) -> bool:
        return self._async_conn_executor.is_conn_dto_equals(another)

    def is_context_info_equals(self, another: RequestContextInfo) -> bool:
        return self._async_conn_executor.is_context_info_equals(another)

    @property
    def is_initialized(self) -> bool:
        return self._async_conn_executor.is_initialized

    def initialize(self) -> None:
        self._await_sync(self._async_conn_executor.initialize())

    def close(self) -> None:
        sa_adapter = self._extract_sync_sa_adapter(raise_on_not_exists=False)

        if sa_adapter is not None:
            sa_adapter.close()

        self._await_sync(self._async_conn_executor.close())

    @overload
    def _extract_sync_sa_adapter(self, raise_on_not_exists: Literal[False]) -> Optional[SyncDirectDBAdapter]:
        pass

    @overload  # noqa
    def _extract_sync_sa_adapter(self, raise_on_not_exists: Literal[True]) -> SyncDirectDBAdapter:
        pass

    def _extract_sync_sa_adapter(self, raise_on_not_exists: bool = False) -> Optional[SyncDirectDBAdapter]:  # noqa
        actual_executor = self._async_conn_executor
        if not actual_executor.is_initialized:
            self.initialize()

        if isinstance(actual_executor, DefaultSqlAlchemyConnExecutor):
            # noinspection PyProtectedMember
            sa_adapter = actual_executor._get_sync_sa_adapter()
            if sa_adapter is None and raise_on_not_exists:
                raise ValueError(f"Async conn executor {actual_executor} has no sync SA adapter.")
            return sa_adapter

        if raise_on_not_exists:
            raise TypeError("Direct sync execution allowed only if wrapping DefaultSqlAlchemyConnExecutor")

        return None

    def _execute_on_sync_adapter_from_wrapped_executor(self, query: ConnExecutorQuery) -> SyncExecutionResult:
        sa_adapter = self._extract_sync_sa_adapter(raise_on_not_exists=True)
        db_adapter_query = self._async_conn_executor.executor_query_to_db_adapter_query(query)
        generator = sa_adapter.execute_by_steps(db_adapter_query)
        cursor_msg = next(generator)

        if not isinstance(cursor_msg, ExecutionStepCursorInfo):
            raise ValueError(f"Unexpected first message from SA adapter: {cursor_msg}")

        def data_generator() -> Generator[TBIDataTable, None, None]:
            for msg in generator:
                if not isinstance(msg, ExecutionStepDataChunk):
                    raise ValueError(f"Unexpected message instead of ExecutionStepDataChunk from SA adapter : {msg}")
                # TODO CONSIDER: May be convert to user types in SAAdapter
                yield [self._async_conn_executor.cast_row_to_output(row, query.user_types) for row in msg.chunk]

        return SyncExecutionResult(cursor_info=cursor_msg.cursor_info, result=data_generator())

    def _execute_in_loop(self, query: ConnExecutorQuery) -> SyncExecutionResult:
        async_result: AsyncExecutionResult = self._await_sync(self._async_conn_executor.execute(query))

        def data_generator() -> Generator[TBIDataTable, None, None]:
            chunk_iter = async_result.result.__aiter__()
            while True:
                try:
                    yield self._await_sync(chunk_iter.__anext__())
                except StopAsyncIteration:
                    return

        return SyncExecutionResult(cursor_info=async_result.cursor_info, result=data_generator())

    @init_required
    def execute(self, query: ConnExecutorQuery) -> SyncExecutionResult:
        sa_adapter = self._extract_sync_sa_adapter(raise_on_not_exists=False)
        if sa_adapter is not None:
            return self._execute_on_sync_adapter_from_wrapped_executor(query)
        return self._execute_in_loop(query)

    @init_required
    def test(self) -> None:
        sa_adapter = self._extract_sync_sa_adapter(raise_on_not_exists=False)
        if sa_adapter is not None:
            return sa_adapter.test()
        return self._await_sync(self._async_conn_executor.test())

    @init_required
    def get_db_version(self, db_ident: DBIdent) -> Optional[str]:
        sa_adapter = self._extract_sync_sa_adapter(raise_on_not_exists=False)
        if sa_adapter is not None:
            return sa_adapter.get_db_version(db_ident)
        return self._await_sync(self._async_conn_executor.get_db_version(db_ident))

    @init_required
    def get_schema_names(self, db_ident: DBIdent) -> List[str]:
        sa_adapter = self._extract_sync_sa_adapter(raise_on_not_exists=False)
        if sa_adapter is not None:
            return sa_adapter.get_schema_names(db_ident)
        return self._await_sync(self._async_conn_executor.get_schema_names(db_ident))

    @init_required
    def get_tables(self, schema_ident: SchemaIdent) -> List[TableIdent]:
        sa_adapter = self._extract_sync_sa_adapter(raise_on_not_exists=False)
        if sa_adapter is not None:
            return sa_adapter.get_tables(schema_ident)
        return self._await_sync(self._async_conn_executor.get_tables(schema_ident))

    @init_required
    def get_table_schema_info(self, table_def: TableDefinition) -> SchemaInfo:
        sa_adapter = self._extract_sync_sa_adapter(raise_on_not_exists=False)

        if sa_adapter is not None:
            raw_schema_info = sa_adapter.get_table_info(
                table_def,
                self._async_conn_executor._conn_options.fetch_table_indexes,
            )
            return self._async_conn_executor.create_schema_info_from_raw_schema_info(raw_schema_info)

        return self._await_sync(self._async_conn_executor.get_table_schema_info(table_def))

    @init_required
    def is_table_exists(self, table_ident: TableIdent) -> bool:
        sa_adapter = self._extract_sync_sa_adapter(raise_on_not_exists=False)
        if sa_adapter is not None:
            return sa_adapter.is_table_exists(table_ident)
        return self._await_sync(self._async_conn_executor.is_table_exists(table_ident))

    @init_required
    def __enter__(self) -> SyncWrapperForAsyncConnExecutor:
        return self

    def __exit__(self, exc_type: Optional[Type[Exception]], exc_val: Optional[Exception], exc_tb: Any) -> None:
        self.close()

lib/dl_core/dl_core/connection_executors/sync_executor_wrapper.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_core.connection_models.common_models import (
    DBIdent,
    SATextTableDefinition,
    SchemaIdent,
    TableDefinition,
    TableIdent,
)
from dl_core.connection_models.conn_options import ConnectOptions
from dl_core.connection_models.dto_defs import (
    ConnDTO,
    DefaultSQLDTO,
)


__all__ = (
    "DBIdent",
    "TableIdent",
    "SchemaIdent",
    "TableDefinition",
    "SATextTableDefinition",
    "ConnectOptions",
    "ConnDTO",
    "DefaultSQLDTO",
)

lib/dl_core/dl_core/connection_models/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
    Optional,
)

import attr
import sqlalchemy as sa

from dl_core.enums import TableDefinitionType


@attr.s(frozen=True, auto_attribs=True)
class DBIdent:
    db_name: Optional[str]

    def __str__(self) -> str:
        return self.db_name or ""


@attr.s(frozen=True, auto_attribs=True)
class SchemaIdent:
    db_name: Optional[str]
    schema_name: Optional[str]

    def __str__(self) -> str:
        return f"{self.db_name or ''}.{self.schema_name or ''}"

    def clone(self, **kwargs: Any) -> SchemaIdent:
        return attr.evolve(self, **kwargs)


class TableDefinition:
    def_type: ClassVar[TableDefinitionType]


@attr.s(frozen=True, auto_attribs=True)
class TableIdent(TableDefinition):
    def_type = TableDefinitionType.table_ident

    db_name: Optional[str]
    schema_name: Optional[str]
    table_name: str

    def __str__(self) -> str:
        return f"{self.db_name or ''}.{self.schema_name or ''}.{self.table_name or ''}"

    def clone(self, **kwargs: Any) -> TableIdent:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True, auto_attribs=True)
class SATextTableDefinition(TableDefinition):
    def_type = TableDefinitionType.sa_text

    text: sa.sql.elements.TextClause

lib/dl_core/dl_core/connection_models/common_models.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Optional,
    Type,
    TypeVar,
)

import attr


_CONNECT_OPTIONS_TV = TypeVar("_CONNECT_OPTIONS_TV", bound="ConnectOptions")


@attr.s(frozen=True, hash=True)
class ConnectOptions:
    rqe_total_timeout: Optional[int] = attr.ib(default=None)
    rqe_sock_read_timeout: Optional[int] = attr.ib(default=None)
    fetch_table_indexes: bool = attr.ib(default=False)
    pass_db_messages_to_user: bool = attr.ib(default=False)
    pass_db_query_to_user: bool = attr.ib(default=False)

    def to_subclass(self, subcls: Type[_CONNECT_OPTIONS_TV], **kwargs: Any) -> _CONNECT_OPTIONS_TV:
        assert issubclass(subcls, type(self))
        full_kwargs = {
            **attr.asdict(self),
            **kwargs,
        }
        # `error: Incompatible return value type (got "ConnectOptions", expected "_CONNECT_OPTIONS_TV")`
        return subcls(**full_kwargs)  # type: ignore  # TODO: fix

    def clone(self: _CONNECT_OPTIONS_TV, **kwargs: Any) -> _CONNECT_OPTIONS_TV:
        return attr.evolve(self, **kwargs)

lib/dl_core/dl_core/connection_models/conn_options.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    ClassVar,
    List,
    Optional,
)

import attr

from dl_constants.enums import ConnectionType


@attr.s(frozen=True)
class ConnDTO:
    conn_type: ClassVar[ConnectionType]

    conn_id: Optional[str] = attr.ib(kw_only=True)

    def conn_reporting_data(self) -> dict:
        return dict(
            connection_id=self.conn_id,
        )


@attr.s(frozen=True)
class DefaultSQLDTO(ConnDTO):  # noqa
    host: str = attr.ib(kw_only=True)
    multihosts: tuple[str, ...] = attr.ib(kw_only=True, converter=tuple)
    port: int = attr.ib(kw_only=True)
    db_name: str = attr.ib(kw_only=True)
    username: str = attr.ib(kw_only=True)
    password: str = attr.ib(repr=False, kw_only=True)

    def get_all_hosts(self) -> List[str]:
        return list(self.multihosts) if self.multihosts else [self.host] if self.host else []

    def conn_reporting_data(self) -> dict:
        return super().conn_reporting_data() | dict(
            host=self.host,
        )

lib/dl_core/dl_core/connection_models/dto_defs.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    AbstractSet,
    ClassVar,
    Type,
)

import attr

from dl_core.connection_models import (
    ConnDTO,
    ConnectOptions,
)


LOGGER = logging.getLogger(__name__)


@attr.s
class DBDomainManager:
    async def normalize_db_host(self, original_host: str) -> str:
        return original_host


@attr.s
class ConnectionSecurityManager(metaclass=abc.ABCMeta):
    """
    This class performs security checks against connection executions.
    Assumed that instance of this class will be created on each request.
    Individual checks should be performed in context of instance fields that are representing current request context.
    """

    db_domain_manager: DBDomainManager = attr.ib(factory=DBDomainManager)

    @abc.abstractmethod
    def is_safe_connection(self, conn_dto: ConnDTO, conn_options: ConnectOptions) -> bool:
        """
        Must return False if connection is potentially unsafe and should be executed in isolated environment.
        """

    # TODO FIX: determine if we need dedicated method
    @abc.abstractmethod
    def is_internal_connection(self, conn_dto: ConnDTO) -> bool:
        """
        Mostly a kludge for dl_core/data_processing/selectors/utils.py:get_query_type()
        Potentially it's behaviour should be same as .is_safe_connection() (TBD: check it),
         but we need exactly the same data to determine if connection is internal.
        """


class ConnectionSafetyChecker(metaclass=abc.ABCMeta):
    _DTO_TYPES: ClassVar[set[Type[ConnDTO]]]

    @classmethod
    def register_dto_types(cls, dto_classes: AbstractSet[Type[ConnDTO]]) -> None:
        cls._DTO_TYPES.update(dto_classes)

    @abc.abstractmethod
    def is_safe_connection(self, conn_dto: ConnDTO, conn_options: ConnectOptions) -> bool:
        raise NotImplementedError


@attr.s(kw_only=True)
class InsecureConnectionSafetyChecker(ConnectionSafetyChecker):
    """Always safe"""

    _DTO_TYPES: ClassVar[set[Type[ConnDTO]]] = set()

    def is_safe_connection(self, conn_dto: ConnDTO, conn_options: ConnectOptions) -> bool:
        return True


@attr.s(kw_only=True)
class NonUserInputConnectionSafetyChecker(ConnectionSafetyChecker):
    """Hosts are not entered by user"""

    _DTO_TYPES: ClassVar[set[Type[ConnDTO]]] = set()

    def is_safe_connection(self, conn_dto: ConnDTO, conn_options: ConnectOptions) -> bool:
        if type(conn_dto) in self._DTO_TYPES:
            LOGGER.info("%r in safe DTO types", type(conn_dto))
            return True
        return False


@attr.s(kw_only=True)
class GenericConnectionSecurityManager(ConnectionSecurityManager, metaclass=abc.ABCMeta):
    conn_sec_checkers: list[ConnectionSafetyChecker] = attr.ib()

    def is_safe_connection(self, conn_dto: ConnDTO, conn_options: ConnectOptions) -> bool:
        return any(
            conn_sec_checker.is_safe_connection(conn_dto, conn_options) for conn_sec_checker in self.conn_sec_checkers
        )


@attr.s(kw_only=True)
class InsecureConnectionSecurityManager(GenericConnectionSecurityManager):
    conn_sec_checkers: list[ConnectionSafetyChecker] = attr.ib(default=[InsecureConnectionSafetyChecker()])

    def is_internal_connection(self, conn_dto: ConnDTO) -> bool:
        return True


@attr.s(frozen=True)
class ConnSecuritySettings:
    security_checker_cls: Type[ConnectionSafetyChecker] = attr.ib()
    dtos: AbstractSet[Type[ConnDTO]] = attr.ib()

lib/dl_core/dl_core/connections_security/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    AbstractSet,
    Callable,
    ClassVar,
    Optional,
    Type,
)

from sqlalchemy.orm import Query

from dl_constants.enums import SourceBackendType
from dl_core.connections_security.base import ConnSecuritySettings
from dl_core.connectors.base.data_source_migration import (
    DataSourceMigrator,
    DefaultDataSourceMigrator,
)
from dl_core.connectors.base.lifecycle import (
    ConnectionLifecycleManager,
    DefaultConnectionLifecycleManager,
)
from dl_core.connectors.base.query_compiler import QueryCompiler
from dl_core.connectors.settings.primitives import ConnectorSettingsDefinition
from dl_core.data_source.base import DataSource
from dl_core.data_source_spec.base import DataSourceSpec
from dl_core.us_manager.storage_schemas.data_source_spec_base import DataSourceSpecStorageSchema


if TYPE_CHECKING:
    from marshmallow import Schema
    from sqlalchemy.types import TypeEngine

    from dl_constants.enums import (
        ConnectionType,
        DataSourceType,
    )
    from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
    from dl_core.connection_executors.async_base import AsyncConnExecutorBase
    from dl_core.connection_executors.common_base import ConnExecutorBase
    from dl_core.db.conversion_base import TypeTransformer
    from dl_core.db.native_type import GenericNativeType
    from dl_core.reporting.notifications import BaseNotification
    from dl_core.us_connection_base import ConnectionBase


class CoreSourceDefinition(abc.ABC):
    source_type: ClassVar[DataSourceType]
    source_cls: ClassVar[Type[DataSource]] = DataSource  # type: ignore
    source_spec_cls: ClassVar[Type[DataSourceSpec]] = DataSourceSpec
    us_storage_schema_cls: ClassVar[Type[DataSourceSpecStorageSchema]] = DataSourceSpecStorageSchema


class CoreConnectionDefinition(abc.ABC):
    conn_type: ClassVar[ConnectionType]
    connection_cls: ClassVar[Type[ConnectionBase]]
    us_storage_schema_cls: ClassVar[Optional[Type[Schema]]] = None
    type_transformer_cls: ClassVar[Type[TypeTransformer]]
    sync_conn_executor_cls: ClassVar[Optional[Type[ConnExecutorBase]]] = None
    async_conn_executor_cls: ClassVar[Optional[Type[AsyncConnExecutorBase]]] = None
    lifecycle_manager_cls: ClassVar[Type[ConnectionLifecycleManager]] = DefaultConnectionLifecycleManager
    dialect_string: ClassVar[str]
    data_source_migrator_cls: ClassVar[Type[DataSourceMigrator]] = DefaultDataSourceMigrator
    settings_definition: ClassVar[Optional[Type[ConnectorSettingsDefinition]]] = None
    custom_dashsql_key_names: frozenset[str] = frozenset()


class CoreConnector(abc.ABC):
    # backend_type-bound properties - TODO: move to a separate entity
    backend_type: ClassVar[SourceBackendType] = SourceBackendType.NONE
    compiler_cls: ClassVar[Type[QueryCompiler]] = QueryCompiler
    query_cls: ClassVar[Type[Query]] = Query
    # others
    connection_definitions: ClassVar[tuple[Type[CoreConnectionDefinition], ...]] = ()
    source_definitions: ClassVar[tuple[Type[CoreSourceDefinition], ...]] = ()
    sa_types: ClassVar[Optional[dict[GenericNativeType, Callable[[GenericNativeType], TypeEngine]]]] = None
    rqe_adapter_classes: ClassVar[AbstractSet[Type[CommonBaseDirectAdapter]]] = frozenset()
    conn_security: ClassVar[AbstractSet[ConnSecuritySettings]] = frozenset()
    query_fail_exceptions: frozenset[Type[Exception]] = frozenset()
    notification_classes: ClassVar[tuple[Type[BaseNotification], ...]] = ()

    @classmethod
    def registration_hook(cls) -> None:
        """Do some non-standard stuff here. Think twice before implementing."""

lib/dl_core/dl_core/connectors/base/connector.py
datalens-tech/datalens-backend
from dl_constants.enums import ConnectionType


_CONNECTION_DASHSQL_KEYS: dict[ConnectionType, frozenset[str]] = {}


def register_custom_dash_sql_key_names(conn_type: ConnectionType, key_names: frozenset[str]) -> None:
    if (registered_key_names := _CONNECTION_DASHSQL_KEYS.get(conn_type)) is not None:
        assert registered_key_names == key_names
    else:
        _CONNECTION_DASHSQL_KEYS[conn_type] = key_names


def get_custom_dash_sql_key_names(conn_type: ConnectionType) -> frozenset[str]:
    return _CONNECTION_DASHSQL_KEYS[conn_type]

lib/dl_core/dl_core/connectors/base/dashsql.py
datalens-tech/datalens-backend
import abc
from typing import (
    Any,
    Callable,
    Optional,
    Sequence,
    Type,
)

import attr

from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
)
from dl_core.base_models import ConnectionRef
from dl_core.data_source_spec.base import DataSourceSpec
import dl_core.exc as exc


@attr.s(frozen=True)
class DataSourceMigrationInterface:
    """Base class for all data source migration interfaces"""


class DataSourceMigrator(abc.ABC):
    @abc.abstractmethod
    def export_migration_dtos(
        self,
        data_source_spec: DataSourceSpec,
    ) -> Sequence[DataSourceMigrationInterface]:
        raise NotImplementedError

    def import_migration_dtos(
        self,
        migration_dtos: Sequence[DataSourceMigrationInterface],
        connection_ref: ConnectionRef,
    ) -> DataSourceSpec:
        chosen_dto = self._choose_migration_dto(migration_dtos=migration_dtos)
        return self._load_migration_dto(migration_dto=chosen_dto, connection_ref=connection_ref)

    @abc.abstractmethod
    def _choose_migration_dto(
        self,
        migration_dtos: Sequence[DataSourceMigrationInterface],
    ) -> DataSourceMigrationInterface:
        raise NotImplementedError

    @abc.abstractmethod
    def _load_migration_dto(
        self,
        migration_dto: DataSourceMigrationInterface,
        connection_ref: ConnectionRef,
    ) -> DataSourceSpec:
        raise NotImplementedError


class DefaultDataSourceMigrator(DataSourceMigrator):
    def export_migration_dtos(
        self,
        data_source_spec: DataSourceSpec,
    ) -> Sequence[DataSourceMigrationInterface]:
        return ()

    def _choose_migration_dto(
        self,
        migration_dtos: Sequence[DataSourceMigrationInterface],
    ) -> DataSourceMigrationInterface:
        raise exc.DataSourceMigrationImpossible()

    def _load_migration_dto(
        self,
        migration_dto: DataSourceMigrationInterface,
        connection_ref: ConnectionRef,
    ) -> DataSourceSpec:
        raise NotImplementedError


@attr.s(frozen=True)
class MigrationKeyMappingItem:
    migration_dto_key: str = attr.ib(kw_only=True)
    source_spec_key: str = attr.ib(kw_only=True)
    custom_export_resolver: Optional[Callable[[DataSourceSpec, str], Any]] = attr.ib(kw_only=True, default=None)
    custom_import_resolver: Optional[Callable[[DataSourceMigrationInterface, str], Any]] = attr.ib(
        kw_only=True, default=None
    )


@attr.s(frozen=True)
class MigrationSpec:
    source_type: DataSourceType = attr.ib(kw_only=True)
    dto_cls: Type[DataSourceMigrationInterface] = attr.ib(kw_only=True)
    dsrc_spec_cls: Type[DataSourceSpec] = attr.ib(kw_only=True)
    migration_mapping_items: tuple[MigrationKeyMappingItem, ...] = attr.ib(kw_only=True)


class SpecBasedSourceMigrator(DataSourceMigrator):
    def get_migration_specs(self) -> list[MigrationSpec]:
        return []

    def _export_for_migration_spec(
        self,
        data_source_spec: DataSourceSpec,
        migration_spec: MigrationSpec,
    ) -> DataSourceMigrationInterface:
        assert data_source_spec.source_type == migration_spec.source_type
        params: dict[str, Any] = {}

        for mapping_item in migration_spec.migration_mapping_items:
            if mapping_item.custom_export_resolver is not None:
                params[mapping_item.migration_dto_key] = mapping_item.custom_export_resolver(
                    data_source_spec, mapping_item.source_spec_key
                )
            elif hasattr(data_source_spec, mapping_item.source_spec_key):
                params[mapping_item.migration_dto_key] = getattr(data_source_spec, mapping_item.source_spec_key)

        return migration_spec.dto_cls(**params)  # type: ignore

    def export_migration_dtos(
        self,
        data_source_spec: DataSourceSpec,
    ) -> Sequence[DataSourceMigrationInterface]:
        result: list[DataSourceMigrationInterface] = []
        for migration_spec in self.get_migration_specs():
            if migration_spec.source_type == data_source_spec.source_type:
                result.append(
                    self._export_for_migration_spec(
                        data_source_spec=data_source_spec,
                        migration_spec=migration_spec,
                    )
                )

        return result

    def _choose_migration_dto(
        self,
        migration_dtos: Sequence[DataSourceMigrationInterface],
    ) -> DataSourceMigrationInterface:
        migration_specs = self.get_migration_specs()
        for dto in migration_dtos:
            for migration_spec in migration_specs:
                if type(dto) is migration_spec.dto_cls:
                    return dto

        raise exc.DataSourceMigrationImpossible()

    def _load_migration_dto(
        self,
        migration_dto: DataSourceMigrationInterface,
        connection_ref: ConnectionRef,
    ) -> DataSourceSpec:
        for migration_spec in self.get_migration_specs():
            if type(migration_dto) is not migration_spec.dto_cls:
                continue

            params: dict[str, Any] = {"source_type": migration_spec.source_type, "connection_ref": connection_ref}
            for mapping_item in migration_spec.migration_mapping_items:
                if mapping_item.custom_import_resolver is not None:
                    params[mapping_item.source_spec_key] = mapping_item.custom_import_resolver(
                        migration_dto, mapping_item.migration_dto_key
                    )
                elif mapping_item.source_spec_key in migration_spec.dsrc_spec_cls.get_public_fields():
                    params[mapping_item.source_spec_key] = getattr(migration_dto, mapping_item.migration_dto_key)

            params = {key: value for key, value in params.items() if value is not None}
            dsrc_spec = migration_spec.dsrc_spec_cls(**params)
            return dsrc_spec

        raise TypeError(f"Invalid dto class: {type(migration_dto)}")


_SOURCE_MIGRATORS: dict[ConnectionType, Type[DataSourceMigrator]] = {
    ConnectionType.unknown: DefaultDataSourceMigrator,
}


def register_data_source_migrator(conn_type: ConnectionType, migrator_cls: Type[DataSourceMigrator]) -> None:
    try:
        assert _SOURCE_MIGRATORS[conn_type] is migrator_cls
    except KeyError:
        _SOURCE_MIGRATORS[conn_type] = migrator_cls


def get_data_source_migrator(conn_type: ConnectionType) -> DataSourceMigrator:
    return _SOURCE_MIGRATORS[conn_type]()

lib/dl_core/dl_core/connectors/base/data_source_migration.py
datalens-tech/datalens-backend
import contextlib
import logging
from typing import (
    Callable,
    ClassVar,
    Generator,
    Optional,
    Type,
)

from sqlalchemy import exc as sa_exc

from dl_core.connectors.base.error_transformer import (
    DbErrorTransformer,
    DBExcKWArgs,
)
import dl_core.exc as exc


LOGGER = logging.getLogger(__name__)


class ErrorHandlerMixin:
    # TODO: Switch to nested mode from mixin
    EXTRA_EXC_CLS: ClassVar[tuple[Type[Exception], ...]] = ()

    # TODO FIX: Try to implement overrides in more functional style
    @classmethod
    def make_exc(
        cls, wrapper_exc: Exception, orig_exc: Optional[Exception], debug_compiled_query: Optional[str]
    ) -> tuple[Type[exc.DatabaseQueryError], DBExcKWArgs]:
        raise NotImplementedError()

    # TODO CONSIDER: Pass DBAdapterQuery to be able to add some context logging
    @contextlib.contextmanager
    def handle_execution_error(
        self,
        debug_compiled_query: Optional[str],
        # This function will be called after converting DatabaseQueryError
        # or in case if DatabaseQueryError will be caught by this handler
        exc_post_processor: Optional[Callable[[exc.DatabaseQueryError], None]] = None,
    ) -> Generator[None, None, None]:
        exc_clses_to_catch: tuple[Type[Exception], ...] = (
            sa_exc.DatabaseError,
            sa_exc.InvalidRequestError,
        ) + self.EXTRA_EXC_CLS

        def post_process_exc_ignore_errors(exc_to_post_process: exc.DatabaseQueryError) -> None:
            if exc_post_processor is not None:
                try:
                    exc_post_processor(exc_to_post_process)
                except Exception:  # noqa
                    LOGGER.exception("Error during postprocessing DatabaseQueryError exception")

        try:
            yield

        except exc.DatabaseQueryError as already_converted_err:
            post_process_exc_ignore_errors(already_converted_err)
            raise

        except exc_clses_to_catch as err:
            LOGGER.info("Got DB exception in DBA.handle_execution_error()", exc_info=True)

            orig_exc = getattr(err, "orig", None)

            exc_cls, exc_kwargs = self.make_exc(
                wrapper_exc=err, orig_exc=orig_exc, debug_compiled_query=debug_compiled_query
            )
            try:
                raise exc_cls(**exc_kwargs) from err
            except exc.DatabaseQueryError as just_raised:
                post_process_exc_ignore_errors(just_raised)
                raise

        except Exception:
            LOGGER.info("Got unexpected exception in DBA.handle_execution_error()", exc_info=True)
            raise


class ETBasedExceptionMaker(ErrorHandlerMixin):
    """
    ErrorTransformer-based exception handling
    """

    _error_transformer: ClassVar[DbErrorTransformer]

    @classmethod
    def make_exc(
        cls, wrapper_exc: Exception, orig_exc: Optional[Exception], debug_compiled_query: Optional[str]
    ) -> tuple[Type[exc.DatabaseQueryError], DBExcKWArgs]:
        trans_exc_cls, kw = cls._error_transformer.make_bi_error_parameters(
            wrapper_exc=wrapper_exc,
            debug_compiled_query=debug_compiled_query,
        )

        return trans_exc_cls, kw

lib/dl_core/dl_core/connectors/base/error_handling.py
datalens-tech/datalens-backend
import abc
import re
from typing import (
    Any,
    Callable,
    ClassVar,
    Iterable,
    Optional,
    Type,
    TypedDict,
)

import attr
from sqlalchemy import exc as sa_exc

import dl_core.exc as exc
from dl_core.exc import DatabaseQueryError


class DBExcKWArgs(TypedDict, total=False):
    db_message: Optional[str]
    query: Optional[str]
    orig: Optional[Exception]
    details: dict[str, Any]


class DbErrorTransformer(abc.ABC):
    _DEFAULT_EXC_CLS: ClassVar[Type[exc.DatabaseQueryError]] = exc.DatabaseQueryError

    def make_bi_error_parameters(
        self,
        wrapper_exc: Exception,
        debug_compiled_query: Optional[str] = None,
    ) -> tuple[Type[exc.DatabaseQueryError], DBExcKWArgs]:
        kw: DBExcKWArgs = DBExcKWArgs(
            db_message=str(wrapper_exc),
            query=debug_compiled_query,
            orig=wrapper_exc,
            details={},
        )

        return self._DEFAULT_EXC_CLS, kw

    def make_bi_error(
        self,
        wrapper_exc: Exception,
        orig_exc: Optional[Exception] = None,
        debug_compiled_query: Optional[str] = None,
    ) -> exc.DatabaseQueryError:
        trans_exc_cls, kw = self.make_bi_error_parameters(
            wrapper_exc=wrapper_exc,
            debug_compiled_query=debug_compiled_query,
        )
        return trans_exc_cls(**kw)


ExcMatchCondition = Callable[[Exception], bool]


@attr.s(frozen=True, kw_only=True)
class ErrorTransformerRule(abc.ABC):
    """Definition of a rule for transforming exceptions raised by a connector
    to their BI-domain exception counterparts. See `ChainedDbErrorTransformer` for more details.

    `when` - is a condition that is evaluated for the exception raised by a connector.
    If a condition is matched then `then_raise` exception will be raised

    `then_raise` - is a BI-domain exception that will be raised if `when` condition is matched
    """

    when: ExcMatchCondition = attr.ib()
    then_raise: Type[exc.DatabaseQueryError] = attr.ib()

    def get_bi_error_class(
        self,
        wrapper_exc: Exception,
    ) -> Optional[Type[exc.DatabaseQueryError]]:
        if self.when(wrapper_exc):
            return self.then_raise
        else:
            return None


def wrapper_exc_is(wrapper_exc_cls: Type[Exception]) -> ExcMatchCondition:
    def _(wrapper_exc: Exception) -> bool:
        return isinstance(wrapper_exc, wrapper_exc_cls)

    return _


def orig_exc_is(orig_exc_cls: Type[Exception]) -> ExcMatchCondition:
    def _(wrapper_exc: Exception) -> bool:
        orig_exc = getattr(wrapper_exc, "orig", None)
        return isinstance(orig_exc, orig_exc_cls)

    return _


def wrapper_exc_is_and_matches_re(wrapper_exc_cls: Type[Exception], err_regex_str: str) -> ExcMatchCondition:
    pattern = re.compile(err_regex_str)

    def _(wrapper_exc: Exception) -> bool:
        return isinstance(wrapper_exc, wrapper_exc_cls) and pattern.search(str(wrapper_exc)) is not None

    return _


@attr.s(frozen=True)
class ChainedDbErrorTransformer(DbErrorTransformer):
    """Transformer for exceptions raised by a connector
    to their BI-domain exception counterparts that works as a part of a chain of rules.
    The rules are applied one-by-one in a chain.
    The first rule in a chain with a matching condition will result in raising
    corresponding BI-domain exception.
    If none of the rules match the condition a fallback exception will be raised.

    `rule_chain` - is a sequence of `ErrorTransformerRule` that are applied one-by-one.
    If a rule condition is matched it produces a BI-domain exception that will be raised.

    `fallback_exc` - is a BI-domain exception that will be raised if none of the rule conditions were matched
    """

    rule_chain: Iterable[ErrorTransformerRule] = attr.ib()
    fallback_exc: Type[exc.DatabaseQueryError] = attr.ib(default=exc.DatabaseQueryError)

    def make_bi_error(
        self,
        wrapper_exc: Exception,
        orig_exc: Optional[Exception] = None,
        debug_compiled_query: Optional[str] = None,
    ) -> DatabaseQueryError:
        transformed_exc_cls: Type[DatabaseQueryError] = self._get_bi_error_cls(wrapper_exc)
        kw: DBExcKWArgs = self._get_error_kw(debug_compiled_query, orig_exc, wrapper_exc)
        return transformed_exc_cls(**kw)

    def make_bi_error_parameters(
        self,
        wrapper_exc: Exception,
        debug_compiled_query: Optional[str] = None,
    ) -> tuple[Type[exc.DatabaseQueryError], DBExcKWArgs]:
        transformed_exc_cls: Type[DatabaseQueryError] = self._get_bi_error_cls(wrapper_exc)
        orig_exc = getattr(wrapper_exc, "orig", None)
        kw: DBExcKWArgs = self._get_error_kw(debug_compiled_query, orig_exc, wrapper_exc)

        return transformed_exc_cls, kw

    def _get_bi_error_cls(
        self,
        wrapper_exc: Exception,
    ) -> Type[exc.DatabaseQueryError]:
        transformed_exc_class_generator = (t.get_bi_error_class(wrapper_exc) for t in self.rule_chain)
        return next(
            (exc_cls for exc_cls in transformed_exc_class_generator if exc_cls is not None),
            DbErrorTransformer._DEFAULT_EXC_CLS,
        )

    @staticmethod
    def _get_error_kw(
        debug_compiled_query: Optional[str], orig_exc: Optional[Exception], wrapper_exc: Exception
    ) -> DBExcKWArgs:
        return dict(
            db_message=str(orig_exc) if orig_exc else str(wrapper_exc),
            query=debug_compiled_query,
            orig=orig_exc,
            details={},
        )


default_error_transformer_rules = (
    ErrorTransformerRule(
        when=orig_exc_is(sa_exc.NoSuchTableError),  # FIXME: not an exc.DatabaseQueryError
        then_raise=exc.SourceDoesNotExist,
    ),
    ErrorTransformerRule(when=wrapper_exc_is(sa_exc.OperationalError), then_raise=exc.DatabaseOperationalError),
)


def make_default_transformer_with_custom_rules(*custom_rules: ErrorTransformerRule) -> ChainedDbErrorTransformer:
    return ChainedDbErrorTransformer(custom_rules + default_error_transformer_rules)


def make_rule_from_descr(
    descr: tuple[
        tuple[Type[Exception], Optional[str]],
        Type[exc.DatabaseQueryError],
    ]
) -> ErrorTransformerRule:
    (wrapper_exc, re_pattern), then_raise = descr
    if re_pattern is None:
        return ErrorTransformerRule(when=wrapper_exc_is(wrapper_exc), then_raise=then_raise)
    else:
        return ErrorTransformerRule(when=wrapper_exc_is_and_matches_re(wrapper_exc, re_pattern), then_raise=then_raise)

lib/dl_core/dl_core/connectors/base/error_transformer.py
datalens-tech/datalens-backend
from typing import (
    Generic,
    TypeVar,
)

from dl_core.lifecycle.base import EntryLifecycleManager
from dl_core.us_connection_base import ConnectionBase


_CONNECTION_TV = TypeVar("_CONNECTION_TV", bound=ConnectionBase)


class ConnectionLifecycleManager(EntryLifecycleManager[_CONNECTION_TV], Generic[_CONNECTION_TV]):
    pass


class DefaultConnectionLifecycleManager(ConnectionLifecycleManager[ConnectionBase]):
    ENTRY_CLS = ConnectionBase

lib/dl_core/dl_core/connectors/base/lifecycle.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import (
    Enum,
    unique,
)
from typing import (
    TYPE_CHECKING,
    ClassVar,
    List,
    Optional,
    Sequence,
)

import attr
import sqlalchemy as sa
from sqlalchemy.sql.expression import (
    nullsfirst,
    nullslast,
)

from dl_constants.enums import OrderDirection
from dl_core import exc
from dl_core.query.expression import ExpressionCtx


if TYPE_CHECKING:
    from sqlalchemy.engine.default import DefaultDialect
    from sqlalchemy.sql.elements import ClauseElement
    from sqlalchemy.sql.selectable import Select

    from dl_core.query.bi_query import (
        BIQuery,
        SqlSourceType,
    )
    from dl_core.query.expression import OrderByExpressionCtx


@unique
class SectionAliasMode(Enum):
    """
    Compiler configuration that determines how the expressions should be placed
    in e.g. a "group by" section.

    * `unaliased`: e.g. `select {expr} … group by {expr}`
    * `by_alias_in_section`: `select {expr} as col … group by col` (ClickHouse, PostgreSQL, MySQL)
    * `by_alias_in_select`: `select col … group by {expr} as col` (YQL)

    In a more general case, this would be 'section priority order for placing
    the column expression with alias', but there's only 3 known cases anyway.
    """

    aliased = "aliased"
    unaliased = "unaliased"
    by_alias_in_section = "by_alias_in_section"
    by_alias_in_select = "by_alias_in_select"


@attr.s(frozen=True)
class QueryCompiler:
    dialect: DefaultDialect = attr.ib()

    orderby_alias_mode: ClassVar[SectionAliasMode] = SectionAliasMode.by_alias_in_section
    groupby_alias_mode: ClassVar[SectionAliasMode] = SectionAliasMode.unaliased

    force_nulls_lower_than_values: ClassVar[bool] = False

    def quote(self, name: str) -> str:  # TODO: one str is `identifier`, another str is `SQL`
        return self.dialect.identifier_preparer.quote(name)

    def aliased_column(self, expr_ctx: ExpressionCtx) -> ClauseElement:
        alias = expr_ctx.alias
        assert alias, "cannot refer to unaliased expr_ctx with an alias"
        return sa.literal_column(self.quote(alias))

    def make_select_expression(self, expr_ctx: ExpressionCtx, bi_query: BIQuery) -> ClauseElement:
        """Make expression for the SELECT clause"""
        expr = expr_ctx.expression
        if not expr_ctx.alias:
            return expr
        expr_with_alias = expr.label(expr_ctx.alias)
        # TODO: put this separation into subclasses (if the logic proves to be correct).
        if self.groupby_alias_mode == SectionAliasMode.unaliased:
            return expr_with_alias
        if self.groupby_alias_mode == SectionAliasMode.by_alias_in_section:
            return expr_with_alias
        if self.groupby_alias_mode == SectionAliasMode.by_alias_in_select:
            is_in_group_by = any(sel_expr_ctx.alias == expr_ctx.alias for sel_expr_ctx in bi_query.group_by_expressions)
            if is_in_group_by:
                return self.aliased_column(expr_ctx)
            return expr_with_alias
        raise Exception(f"Unexpected {self.groupby_alias_mode=!r}")

    def make_group_by_expression(self, expr_ctx: ExpressionCtx, bi_query: BIQuery) -> ClauseElement:
        """Make expression for the GROUP BY clause"""
        expr = expr_ctx.expression
        if not expr_ctx.alias:
            return expr
        expr_with_alias = expr.label(expr_ctx.alias)
        # TODO: put this separation into subclasses (if the logic proves to be correct).
        if self.groupby_alias_mode == SectionAliasMode.unaliased:
            return expr
        if self.groupby_alias_mode == SectionAliasMode.by_alias_in_section:
            is_selected = any(sel_expr_ctx.alias == expr_ctx.alias for sel_expr_ctx in bi_query.select_expressions)
            if is_selected:
                return self.aliased_column(expr_ctx)
            return expr
        if self.groupby_alias_mode == SectionAliasMode.by_alias_in_select:
            return expr_with_alias
        raise Exception(f"Unexpected {self.groupby_alias_mode=!r}")

    def should_order_by_alias(self, expr_ctx: ExpressionCtx, bi_query: BIQuery) -> bool:
        if not expr_ctx.alias:
            return False
        if self.orderby_alias_mode == SectionAliasMode.unaliased:
            return False
        if self.orderby_alias_mode == SectionAliasMode.by_alias_in_section:
            prior_sections = (
                (bi_query.select_expressions, bi_query.group_by_expressions)
                if self.groupby_alias_mode == SectionAliasMode.by_alias_in_select
                else (bi_query.select_expressions,)
            )
            return any(
                prior_expr_ctx.alias == expr_ctx.alias for section in prior_sections for prior_expr_ctx in section
            )
        raise Exception(f"Unexpected {self.orderby_alias_mode=!r}")

    def make_order_by_expression(self, order_by_ctx: OrderByExpressionCtx, bi_query: BIQuery) -> ClauseElement:
        """Make expression for the ORDER BY clause"""
        expr: ClauseElement
        if self.should_order_by_alias(expr_ctx=order_by_ctx, bi_query=bi_query):
            expr = self.aliased_column(order_by_ctx)
        else:
            expr = order_by_ctx.expression

        dir_wrapper = {
            OrderDirection.asc: sa.asc,
            OrderDirection.desc: sa.desc,
        }[order_by_ctx.direction]
        expr = dir_wrapper(expr)

        if self.force_nulls_lower_than_values:
            nulls_wrapper = {
                OrderDirection.asc: nullsfirst,
                OrderDirection.desc: nullslast,
            }[order_by_ctx.direction]
            expr = nulls_wrapper(expr)

        return expr

    def compile_select(self, bi_query: BIQuery, sql_source: Optional[SqlSourceType]) -> Select:
        """Compile a ``BIQuery`` object into an SQLAlchemy ``Select``"""

        def _unwrap_expressions(ex_list: Sequence) -> List[ClauseElement]:
            """Unwrap ``ExpressionCtx`` objects to nested SA expressions"""
            return [ex.expression if isinstance(ex, ExpressionCtx) else ex for ex in ex_list]

        query = sa.select(
            [self.make_select_expression(expr_ctx, bi_query=bi_query) for expr_ctx in bi_query.select_expressions]
        )

        if sql_source is not None:  # support queries without `from`; mostly for gsheets
            query = query.select_from(sql_source)

        for expr in _unwrap_expressions(bi_query.dimension_filters):
            query = query.where(expr)

        for expr in _unwrap_expressions(bi_query.measure_filters):
            query = query.having(expr)

        if bi_query.distinct:
            query = query.distinct()

        query = query.group_by(
            *[self.make_group_by_expression(expr_ctx, bi_query=bi_query) for expr_ctx in bi_query.group_by_expressions]
        )

        for order_by_ctx in bi_query.order_by_expressions:
            query = query.order_by(self.make_order_by_expression(order_by_ctx, bi_query=bi_query))

        if bi_query.limit is not None:
            query = query.limit(bi_query.limit)
        if bi_query.offset is not None:
            if not bi_query.order_by_expressions:
                raise exc.QueryConstructorError("Offset cannot be used without sorting")
            query = query.offset(bi_query.offset)

        return query

lib/dl_core/dl_core/connectors/base/query_compiler.py
datalens-tech/datalens-backend
from typing import (
    TYPE_CHECKING,
    Type,
)

from dl_core.backend_types import register_connection_backend_type
from dl_core.connection_executors.adapters.common_base import register_dialect_string
from dl_core.connection_executors.remote_query_executor.commons import register_adapter_class
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
    CoreSourceDefinition,
)
from dl_core.connectors.base.dashsql import register_custom_dash_sql_key_names
from dl_core.connectors.base.data_source_migration import register_data_source_migrator
from dl_core.connectors.settings.registry import register_connector_settings_class
from dl_core.data_source.type_mapping import register_data_source_class
from dl_core.data_source_spec.type_mapping import register_data_source_spec_class
from dl_core.db.conversion_base import register_type_transformer_class
from dl_core.db.sa_types import register_sa_types
from dl_core.db_session_utils import (
    register_query_fail_exceptions,
    register_sa_query_cls,
)
from dl_core.reporting.notifications import register_notification
from dl_core.services_registry.conn_executor_factory import (
    register_async_conn_executor_class,
    register_sync_conn_executor_class,
)
from dl_core.us_connection import register_connection_class
from dl_core.us_manager.storage_schemas.connection_schema_registry import register_connection_schema
from dl_core.us_manager.storage_schemas.data_source_spec import register_data_source_schema


if TYPE_CHECKING:
    from dl_core.connections_security.base import ConnSecuritySettings


class CoreConnectorRegistrator:
    @classmethod
    def register_source_definition(cls, source_def: Type[CoreSourceDefinition]) -> None:
        register_data_source_spec_class(
            source_type=source_def.source_type,  # type: ignore
            spec_cls=source_def.source_spec_cls,
        )
        register_data_source_class(
            source_type=source_def.source_type,  # type: ignore
            source_cls=source_def.source_cls,
        )
        register_data_source_schema(
            source_type=source_def.source_type,  # type: ignore
            schema_cls=source_def.us_storage_schema_cls,
        )

    @classmethod
    def register_connection_definition(
        cls,
        conn_def: Type[CoreConnectionDefinition],
        connector: Type[CoreConnector],
    ) -> None:
        register_connection_class(
            new_conn_cls=conn_def.connection_cls,
            allow_ct_override=True,
            conn_type=conn_def.conn_type,
            lifecycle_manager_cls=conn_def.lifecycle_manager_cls,
        )
        register_connection_backend_type(conn_type=conn_def.conn_type, backend_type=connector.backend_type)
        register_connection_schema(conn_cls=conn_def.connection_cls, schema_cls=conn_def.us_storage_schema_cls)  # type: ignore
        register_type_transformer_class(conn_type=conn_def.conn_type, tt_cls=conn_def.type_transformer_cls)  # type: ignore
        if conn_def.sync_conn_executor_cls is not None:
            register_sync_conn_executor_class(
                conn_cls=conn_def.connection_cls, sync_ce_cls=conn_def.sync_conn_executor_cls
            )
        if conn_def.async_conn_executor_cls is not None:
            register_async_conn_executor_class(
                conn_cls=conn_def.connection_cls, async_ce_cls=conn_def.async_conn_executor_cls
            )
        register_dialect_string(conn_type=conn_def.conn_type, dialect_str=conn_def.dialect_string)
        if conn_def.settings_definition is not None:
            register_connector_settings_class(
                conn_type=conn_def.conn_type,
                settings_class=conn_def.settings_definition.settings_class,
                fallback=conn_def.settings_definition.fallback,
            )
        register_data_source_migrator(conn_type=conn_def.conn_type, migrator_cls=conn_def.data_source_migrator_cls)
        register_custom_dash_sql_key_names(conn_type=conn_def.conn_type, key_names=conn_def.custom_dashsql_key_names)

    @classmethod
    def register_connector(cls, connector: Type[CoreConnector]) -> None:
        for source_def in connector.source_definitions:
            cls.register_source_definition(source_def=source_def)
        for conn_def in connector.connection_definitions:
            cls.register_connection_definition(conn_def=conn_def, connector=connector)
        register_sa_types(connector.sa_types or {})
        for conn_sec_settings in connector.conn_security:  # type: ConnSecuritySettings
            dto_types = conn_sec_settings.dtos
            conn_sec_settings.security_checker_cls.register_dto_types(dto_types)
        for adapter_cls in connector.rqe_adapter_classes:
            register_adapter_class(adapter_cls=adapter_cls)
        register_sa_query_cls(backend_type=connector.backend_type, query_cls=connector.query_cls)
        register_query_fail_exceptions(exception_classes=connector.query_fail_exceptions)
        for notification_cls in connector.notification_classes:
            register_notification()(notification_cls)  # it is a parameterized decorator

        connector.registration_hook()  # for custom actions


CONN_REG_CORE = CoreConnectorRegistrator

lib/dl_core/dl_core/connectors/base/registrator.py
datalens-tech/datalens-backend
from typing import Type

import attr

import dl_core as package
from dl_core.connectors.base.connector import CoreConnector
from dl_utils.entrypoints import EntrypointClassManager


_CONNECTOR_EP_GROUP = f"{package.__name__}.connectors"


@attr.s
class CoreConnectorEntrypointManager(EntrypointClassManager[CoreConnector]):
    entrypoint_group_name = attr.ib(init=False, default=_CONNECTOR_EP_GROUP)


def get_all_connectors() -> dict[str, Type[CoreConnector]]:
    ep_mgr = CoreConnectorEntrypointManager()
    return ep_mgr.get_all_ep_classes()


def get_connector_cls(name: str) -> Type[CoreConnector]:
    ep_mgr = CoreConnectorEntrypointManager()
    return ep_mgr.get_ep_class(name=name)

lib/dl_core/dl_core/connectors/registry.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Optional,
    Type,
    TypeVar,
)

import attr

from dl_configs.connectors_data import ConnectorsDataBase
from dl_configs.connectors_settings import (
    ConnectorsConfigType,
    ConnectorSettingsBase,
    SettingsFallbackType,
)
from dl_configs.settings_loaders.fallback_cfg_resolver import ObjectLikeConfig


_CONN_DATA_TV = TypeVar("_CONN_DATA_TV", bound=ConnectorsDataBase)


def get_connectors_settings_config(
    full_cgf: ConnectorsConfigType,
    object_like_config_key: str,
    connector_data_class: Type[_CONN_DATA_TV],
) -> Optional[ObjectLikeConfig | _CONN_DATA_TV]:
    if isinstance(full_cgf, connector_data_class):  # installation-style settings
        return full_cgf
    if isinstance(full_cgf, ObjectLikeConfig):  # yaml-style settings
        assert hasattr(full_cgf, "CONNECTORS")
        settings: Optional[ObjectLikeConfig] = getattr(full_cgf.CONNECTORS, object_like_config_key, None)
        if settings is not None:
            for key, setting in settings.items():  # type: str, Any  # converts left ObjectLikeConfigs to dicts
                if isinstance(setting, (tuple, list)):
                    settings._data[key] = [  # noqa  # TODO avoid changing cfg data directly
                        item.to_dict() if isinstance(item, ObjectLikeConfig) else item for item in setting
                    ]
        return settings
    return None  # no settings from a given connector


@attr.s(frozen=True)
class ConnectorSettingsDefinition:
    settings_class: Type[ConnectorSettingsBase] = attr.ib()
    fallback: SettingsFallbackType = attr.ib()

lib/dl_core/dl_core/connectors/settings/primitives.py
datalens-tech/datalens-backend
from typing import Type

from dl_configs.connectors_settings import (
    ConnectorSettingsBase,
    SettingsFallbackType,
)
from dl_constants.enums import ConnectionType


CONNECTORS_SETTINGS_CLASSES: dict[ConnectionType, Type[ConnectorSettingsBase]] = {}
CONNECTORS_SETTINGS_FALLBACKS: dict[ConnectionType, SettingsFallbackType] = {}


def register_connector_settings_class(
    conn_type: ConnectionType,
    settings_class: Type[ConnectorSettingsBase],
    fallback: SettingsFallbackType,
) -> None:
    if (registered_settings_class := CONNECTORS_SETTINGS_CLASSES.get(conn_type)) is not None:
        assert registered_settings_class == settings_class
        assert CONNECTORS_SETTINGS_FALLBACKS[conn_type] == fallback
    else:
        assert conn_type not in CONNECTORS_SETTINGS_FALLBACKS
        CONNECTORS_SETTINGS_CLASSES[conn_type] = settings_class
        CONNECTORS_SETTINGS_FALLBACKS[conn_type] = fallback

lib/dl_core/dl_core/connectors/settings/registry.py
datalens-tech/datalens-backend
from dl_core.connectors.base.connector import CoreSourceDefinition
from dl_core.data_source_spec.sql import (
    StandardSchemaSQLDataSourceSpec,
    SubselectDataSourceSpec,
)
from dl_core.us_manager.storage_schemas.data_source_spec_base import (
    SchemaSQLDataSourceSpecStorageSchema,
    SubselectDataSourceSpecStorageSchema,
)


class SQLTableCoreSourceDefinitionBase(CoreSourceDefinition):
    source_spec_cls = StandardSchemaSQLDataSourceSpec
    us_storage_schema_cls = SchemaSQLDataSourceSpecStorageSchema


class SQLSubselectCoreSourceDefinitionBase(CoreSourceDefinition):
    source_spec_cls = SubselectDataSourceSpec
    us_storage_schema_cls = SubselectDataSourceSpecStorageSchema

lib/dl_core/dl_core/connectors/sql_base/connector.py
datalens-tech/datalens-backend
from typing import (
    ClassVar,
    Optional,
    Type,
)

import attr

from dl_constants.enums import DataSourceType
from dl_core.connectors.base.data_source_migration import (
    DataSourceMigrationInterface,
    MigrationKeyMappingItem,
    MigrationSpec,
    SpecBasedSourceMigrator,
)
from dl_core.data_source_spec.base import DataSourceSpec
from dl_core.data_source_spec.sql import (
    StandardSQLDataSourceSpec,
    SubselectDataSourceSpec,
)


@attr.s(frozen=True)
class SQLTableDSMI(DataSourceMigrationInterface):
    schema_name: Optional[str] = attr.ib(kw_only=True, default=None)
    table_name: Optional[str] = attr.ib(kw_only=True)


@attr.s(frozen=True)
class SQLTableWDbDSMI(SQLTableDSMI):
    db_name: Optional[str] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class SQLSubselectDSMI(DataSourceMigrationInterface):
    subsql: Optional[str] = attr.ib(kw_only=True)


class DefaultSQLDataSourceMigrator(SpecBasedSourceMigrator):
    table_source_type: ClassVar[Optional[DataSourceType]] = None
    table_dsrc_spec_cls: ClassVar[Optional[Type[DataSourceSpec]]] = StandardSQLDataSourceSpec
    with_db_name: ClassVar[bool] = False

    subselect_source_type: ClassVar[Optional[DataSourceType]] = None
    subselect_dsrc_spec_cls: ClassVar[Optional[Type[DataSourceSpec]]] = SubselectDataSourceSpec

    default_schema_name: ClassVar[Optional[str]] = None

    def _resolve_schema_name_for_export(
        self,
        source_spec: DataSourceSpec,
        attr_name: str,
    ) -> Optional[str]:
        assert attr_name == "schema_name"
        schema_name: Optional[str] = getattr(source_spec, attr_name, None)
        if schema_name == self.default_schema_name:
            schema_name = None
        return schema_name

    def _resolve_schema_name_for_import(
        self,
        migration_dto: DataSourceMigrationInterface,
        attr_name: str,
    ) -> Optional[str]:
        assert attr_name == "schema_name"
        schema_name: Optional[str] = getattr(migration_dto, attr_name, None)
        if schema_name == self.default_schema_name:
            schema_name = None
        return schema_name

    def get_migration_specs(self) -> list[MigrationSpec]:
        result: list[MigrationSpec] = []
        if self.table_source_type is not None:
            assert self.table_dsrc_spec_cls is not None
            result.append(
                MigrationSpec(
                    source_type=self.table_source_type,
                    dto_cls=SQLTableDSMI,
                    dsrc_spec_cls=self.table_dsrc_spec_cls,
                    migration_mapping_items=(
                        MigrationKeyMappingItem(
                            migration_dto_key="schema_name",
                            source_spec_key="schema_name",
                            custom_export_resolver=self._resolve_schema_name_for_export,
                            custom_import_resolver=self._resolve_schema_name_for_import,
                        ),
                        MigrationKeyMappingItem(migration_dto_key="table_name", source_spec_key="table_name"),
                    ),
                )
            )

            if self.with_db_name:
                result.append(
                    MigrationSpec(
                        source_type=self.table_source_type,
                        dto_cls=SQLTableWDbDSMI,
                        dsrc_spec_cls=self.table_dsrc_spec_cls,
                        migration_mapping_items=(
                            MigrationKeyMappingItem(migration_dto_key="db_name", source_spec_key="db_name"),
                            MigrationKeyMappingItem(
                                migration_dto_key="schema_name",
                                source_spec_key="schema_name",
                                custom_export_resolver=self._resolve_schema_name_for_export,
                                custom_import_resolver=self._resolve_schema_name_for_import,
                            ),
                            MigrationKeyMappingItem(migration_dto_key="table_name", source_spec_key="table_name"),
                        ),
                    )
                )

        if self.subselect_source_type is not None:
            assert self.subselect_dsrc_spec_cls is not None
            result.append(
                MigrationSpec(
                    source_type=self.subselect_source_type,
                    dto_cls=SQLSubselectDSMI,
                    dsrc_spec_cls=self.subselect_dsrc_spec_cls,
                    migration_mapping_items=(
                        MigrationKeyMappingItem(migration_dto_key="subsql", source_spec_key="subsql"),
                    ),
                )
            )

        return result

lib/dl_core/dl_core/connectors/sql_base/data_source_migration.py
datalens-tech/datalens-backend
import contextlib
import logging
import os
import typing
import uuid

import attr

import dl_configs.utils as bi_config_utils


@attr.s()
class BaseSSLCertAdapter:
    __context_counter: int = 0
    __file_name: str | None = None

    def __get_filename(self) -> str:
        if self.__file_name is None:
            self.__file_name = f"{uuid.uuid4()}.crt"

        return self.__file_name

    def get_ssl_cert_path(self, ssl_ca: typing.Optional[str]) -> str:
        if ssl_ca is None:
            return bi_config_utils.get_root_certificates_path()

        return os.path.join(
            bi_config_utils.get_temp_root_certificates_folder_path(),
            self.__get_filename(),
        )

    @contextlib.contextmanager
    def ssl_cert_context(self, ssl_ca: typing.Optional[str]) -> typing.Generator[None, None, None]:
        assert ssl_ca is not None, "Root CA can't be None for ssl connection context"

        crt_path = self.get_ssl_cert_path(ssl_ca=ssl_ca)

        if self.__context_counter == 0:
            with open(crt_path, "x", encoding="utf-8") as f:
                logging.debug(f"Writing root CA to {crt_path}")
                f.write(ssl_ca)
        self.__context_counter += 1

        try:
            yield
        finally:
            self.__context_counter -= 1
            if self.__context_counter == 0:
                logging.debug(f"Removing root CA from {crt_path}")
                os.remove(crt_path)

lib/dl_core/dl_core/connectors/ssl_common/adapter.py
datalens-tech/datalens-backend
class ConstContainer:
    """
    Base class for all const containers.
    """


class DatasetConstraints(ConstContainer):
    """
    Various dataset limits
    """

    # TODO: Figure out the best place to put this stuff
    FIELD_COUNT_LIMIT_SOFT = 1200
    FIELD_COUNT_LIMIT_HARD = 1250


class DataAPILimits(ConstContainer):
    """
    Data limits
    """

    DEFAULT_SUBQUERY_LIMIT = 100_000
    DEFAULT_SOURCE_DB_LIMIT = 1_000_000
    PREVIEW_ROW_LIMIT = 1000

    # endpoint-specific
    DATA_API_DEFAULT_ROW_COUNT_HARD_LIMIT = 1_000_000
    PIVOT_API_DEFAULT_ROW_COUNT_HARD_LIMIT = 100_000

lib/dl_core/dl_core/constants.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import re
from typing import (
    Any,
    Iterable,
    Optional,
)

import dateutil.parser


spaces_in_numbers_re = re.compile(r"(?<=\d) (?=\d)")


def _to_int(value: Any) -> int:
    if isinstance(value, str):
        if " " in value:
            value = value.replace(" ", "")
    return int(value)


def _to_float(value: Any) -> float:
    if isinstance(value, str):
        if "," in value:
            value = value.replace(",", ".")
        if " " in value:
            value = spaces_in_numbers_re.sub("", value)
    return float(value)


def _to_date(value: str, formats: Optional[Iterable[str]] = None) -> datetime.date:
    # TODO?: try `ciso8601.parse_datetime`, for correct ISO-8601, timezones, speed, etc.
    formats = formats or (
        "%Y-%m-%d",
        # weird stuff
        "%d.%m.%y",
        "%d.%m.%Y",
        "%d/%m/%Y",
        "%d/%m/%y",
        # day-start:
        "%Y-%m-%d 00:00:00",
        "%Y-%m-%d 00:00:00.000000",
    )
    dt = None
    for fmt in formats:
        try:
            dt = datetime.datetime.strptime(value, fmt)
        except ValueError:
            pass
        else:
            break
    if dt is None:
        raise ValueError(value)
    return dt.date()


def _to_datetime(value: str) -> datetime.datetime:
    dt = dateutil.parser.parse(value)
    return dt


def _to_boolean(value: str | bool) -> bool:
    if isinstance(value, str):
        value = value.lower()
        if value in ("true", "1"):
            return True
        elif value in ("false", "0"):
            return False
        raise ValueError(value)
    return bool(value)

lib/dl_core/dl_core/converter_types_cast.py
datalens-tech/datalens-backend
from typing import (
    Collection,
    Optional,
    Type,
)

from dl_core.connectors.base.connector import CoreConnector
from dl_core.connectors.base.registrator import CONN_REG_CORE
from dl_core.connectors.registry import get_all_connectors


def _register_connector(connector_cls: Type[CoreConnector]) -> None:
    CONN_REG_CORE.register_connector(connector_cls)


def load_all_connectors() -> None:
    get_all_connectors()


def register_all_connectors(connector_ep_names: Optional[Collection[str]] = None) -> None:
    for ep_name, connector_cls in sorted(get_all_connectors().items()):
        if connector_ep_names is not None and ep_name not in connector_ep_names:
            continue
        _register_connector(connector_cls)

lib/dl_core/dl_core/core_connectors.py
datalens-tech/datalens-backend
from typing import Type

from dl_core.data_processors.base.plugin import DataProcessorPlugin
from dl_core.data_processors.ep_loader import get_all_data_processor_plugins
from dl_core.data_processors.registrator import DATA_PROC_REG


def _register_plugin(plugin_cls: Type[DataProcessorPlugin]) -> None:
    DATA_PROC_REG.register_data_processor_plugin(plugin_cls)


def register_all_data_processor_plugins() -> None:
    for _, plugin_cls in sorted(get_all_data_processor_plugins().items()):
        _register_plugin(plugin_cls)

lib/dl_core/dl_core/core_data_processors.py
datalens-tech/datalens-backend
"""

Terminology note: to avoid confusion, using `entry` as in `cache entry`, and
`entity` as in `united-storage entity` (dataset/connection).

Logging fields:

  * Save, Read:
    * cache_key
    * connection_id
    * connection_revision
    * cache_serialized_bytesize  # serialized, without compression
    * cache_dumped_bytesize  # 'fully dumped' as in serialized + maybe-compressed
    * cache_was_compressed
    * cache_compress_alg
    * cache_row_count
    * cache_serializer_timing  # for both serialize and deserialize
    * cache_compressor_timing  # for both compress and decompress
    * cache_ttl_sec  # required on save, optional update on read

  * Save: INFO "Cache saved at {key}"

  * Read: INFO "Cache read from {key}"
    * cache_refresh_on_read

  * Read miss: "Cache miss: …: …"
    * cache_miss_reason

  * Pre-del (sync, async): INFO "Invalidating cache for entity %s"


(note that there's also GenericProfiler logging)
TODO?: gather all the profilings into a log field?


Logging fields in earlier versions:

  * Pre-save: INFO "Saving key … …"
  * Pre-save: INFO "Redis key is %s. Requesting redis"
    * cache_key
    * connection_id
    * connection_revision
    * ttl_sec
    * compress_alg
  * Pre-save: DEBUG ("Preparing cache entry %s"
  * Pre-save: Data dump (serialization + compression): INFO "Cache data was dumped"
    * rows_count
    * encoded_bytes_count
    * was_compressed
    * final_bytes_count
  * Pre-save (sync, async): INFO "Saving to redis..."

  * Pre-load: INFO "Looking for key … …"
  * Pre-load: INFO "Redis key is %s. Requesting redis..."
    * cache_key
    * connection_id
    * connection_revision
    * new_ttl_sec
    * refresh_on_read

  * Cache miss: "BI common version was updated, ignoring cache entry"
  * Cache miss: "Query hash collision, ignoring cache entry"
  * Cache miss: "Connection was updated, ignoring cache entry"

  * Pre-del (sync, async): INFO "Invalidating cache for entity %s"

"""

from __future__ import annotations

import asyncio
import collections
import contextlib
import enum
import gzip
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncGenerator,
    Deque,
    Dict,
    Optional,
    Tuple,
    Type,
    TypeVar,
)

import attr
import lz4.frame
from redis_cache_lock.main import RedisCacheLock
from redis_cache_lock.utils import HistoryHolder

from dl_app_tools.profiling_base import (
    GenericProfiler,
    generic_profiler_async,
)
from dl_constants.types import TJSONExt
import dl_core
from dl_core.data_processing.cache.exc import CachedEntryPackageVersionMismatch
from dl_core.data_processing.cache.primitives import LocalKeyRepresentation
from dl_core.serialization import (
    CacheMetadataSerialization,
    common_dumps,
    common_loads,
)


if TYPE_CHECKING:
    from redis.asyncio import Redis


LOGGER = logging.getLogger(__name__)


@enum.unique
class CompressAlg(enum.Enum):
    GZIP = "gzip"
    LZ4 = "lz4"


@attr.s(frozen=True, auto_attribs=True, slots=True)
class _CacheUpdateRequest:
    full_key: str
    ttl_sec: float
    entry_data: bytes
    details: Dict[str, Any]  # debug/stats, to be dumped into logging.


@attr.s(auto_attribs=True, slots=True)
class ResultCacheEntry:
    """
    ...

    Reminder: `metadata` is saved with the cache, `details` are for the logging.
    """

    key_parts_str: str
    metadata: Optional[Dict[str, Any]] = None
    details: Optional[Dict[str, Any]] = None

    _result_data: Optional[TJSONExt] = None
    _result_data_redis_repr: Optional[bytes] = None
    _is_compressed: Optional[bool] = None
    _compress_alg: Optional[CompressAlg] = None

    def __attrs_post_init__(self) -> None:
        if self._result_data is not None:
            assert self._result_data_redis_repr is None
            assert self._is_compressed is None
            self._is_compressed = False
        elif self._result_data_redis_repr is not None:
            assert self._result_data is None
            assert self._is_compressed is not None

        metadata = self.metadata
        self.metadata = metadata.copy() if metadata else {}
        details = self.details
        self.details = details.copy() if details else {}

    @staticmethod
    def bi_c_version_str() -> str:
        return "12.99.0"

    @property
    def is_error(self) -> bool:
        assert self.metadata is not None
        return self.metadata.get("error") is not None
        # TODO: require the key presence:
        # return self.metadata['error'] is not None

    def make_result_data(self) -> TJSONExt:
        """Deserialize the result data (if not saved) and update the details"""
        if self._result_data is not None:
            return self._result_data

        if not self._result_data_redis_repr:
            raise ValueError("Cache entry contain no data")

        prof_decompress = None
        if self._is_compressed:
            with GenericProfiler("qcache-decompress") as prof_decompress:
                if self._compress_alg == CompressAlg.LZ4:
                    encoded_data = lz4.frame.decompress(self._result_data_redis_repr)
                elif self._compress_alg == CompressAlg.GZIP:
                    encoded_data = gzip.decompress(self._result_data_redis_repr)
                else:
                    encoded_data = gzip.decompress(self._result_data_redis_repr)
        else:
            encoded_data = self._result_data_redis_repr

        with GenericProfiler("qcache-deserialize") as prof_deserialize:
            self._result_data = common_loads(encoded_data)

        details = self.details
        assert details is not None
        result_data = self._result_data
        assert isinstance(result_data, list)
        details.update(
            cache_dumped_bytesize=len(self._result_data_redis_repr),
            cache_was_compressed=self._is_compressed,
            cache_compress_alg=(self._compress_alg.name if self._compress_alg is not None else None),
            cache_compressor_timing=(prof_decompress.exec_time_sec if prof_decompress is not None else None),
            cache_serialized_bytesize=len(encoded_data),
            cache_serializer_timing=prof_deserialize.exec_time_sec,
            cache_row_count=len(result_data),
        )
        return self._result_data

    @classmethod
    def from_redis_data(
        cls,
        redis_data: bytes,
        details: Optional[Dict[str, Any]] = None,
    ) -> ResultCacheEntry:
        ser_obj = CacheMetadataSerialization(redis_data)
        metadata = ser_obj.metadata
        assert isinstance(metadata, dict)
        initial_bi_c_version_str = metadata["bi_common_version"]

        if initial_bi_c_version_str != cls.bi_c_version_str():
            raise CachedEntryPackageVersionMismatch(
                "Package version mismatch. In cache: {}. Actual: {}".format(
                    initial_bi_c_version_str, cls.bi_c_version_str()
                )
            )

        is_compressed = metadata["is_compressed"]
        assert isinstance(is_compressed, bool)
        compress_alg_name = metadata.get("compress_alg")
        compress_alg: Optional[CompressAlg] = None

        if is_compressed:
            if not compress_alg_name:
                compress_alg = CompressAlg.GZIP
            else:
                compress_alg = CompressAlg(compress_alg_name)

        data = ser_obj.data_bytes
        key_parts_str = metadata["key_parts_str"]
        assert isinstance(key_parts_str, str)
        return cls(
            key_parts_str=key_parts_str,
            result_data_redis_repr=data,
            is_compressed=is_compressed,
            compress_alg=compress_alg,
            metadata=metadata,
            details=details,
        )

    def to_redis_data_and_log_details(
        self, compress: bool = False, compress_alg: Optional[CompressAlg] = None
    ) -> Tuple[bytes, Dict[str, Any]]:
        result_data = self._result_data
        assert result_data is not None
        assert isinstance(result_data, list)

        # Note: "qcache-serialize" used to be on this entire method rather than only on dumps.
        with GenericProfiler("qcache-serialize") as prof_serialize:
            serialized_result_data = common_dumps(result_data)

        result_data_to_store = serialized_result_data

        prof_compress = None
        if compress:
            if compress_alg is None:
                compress_alg = CompressAlg.GZIP

            with GenericProfiler("qcache-compress") as prof_compress:
                if compress_alg == CompressAlg.LZ4:
                    result_data_to_store = lz4.frame.compress(serialized_result_data)
                elif compress_alg == CompressAlg.GZIP:
                    result_data_to_store = gzip.compress(serialized_result_data)

        details: Dict[str, Any] = dict(
            cache_row_count=len(result_data),
            cache_serialized_bytesize=len(serialized_result_data),
            cache_serializer_timing=prof_serialize.exec_time_sec,
            cache_was_compressed=compress,
            cache_compressor_timing=(prof_compress.exec_time_sec if prof_compress is not None else None),
            cache_dumped_bytesize=len(result_data_to_store),
        )

        metadata = dict(  # TODO?: attrs-object? or protobuf?
            self.metadata or {},
            bi_common_version=self.bi_c_version_str(),
            key_parts_str=self.key_parts_str,
            is_compressed=bool(compress),
            compress_alg=compress_alg.value if compress_alg else None,
        )
        data_bytes = CacheMetadataSerialization.serialize(metadata, result_data_to_store)
        return data_bytes, details

    def to_redis_data(
        self,
        compress: bool = False,
        compress_alg: Optional[CompressAlg] = None,
    ) -> bytes:
        """
        Convenience wrapper, primarily for tests.
        """
        result, _ = self.to_redis_data_and_log_details(
            compress=compress,
            compress_alg=compress_alg,
        )
        return result


_ECEMA_TV = TypeVar("_ECEMA_TV", bound="EntityCacheEntryManagerAsyncBase")


@attr.s(auto_attribs=True, slots=True)
class EntityCacheEntryManagerAsyncBase:
    """
    Single-cache-entry manager.

    Primarily intended to support cache-entry locking,
    i.e. to provide a protocol that makes it possible.
    """

    local_key_rep: LocalKeyRepresentation
    cache_engine: EntityCacheEngineAsync
    write_ttl_sec: float
    read_extend_ttl_sec: Optional[float] = None
    allow_cache_read: bool = True

    _starting: bool = False
    _started: bool = False

    def clone(self: _ECEMA_TV, **updates: Any) -> _ECEMA_TV:
        return attr.evolve(self, **updates)

    async def initialize(self) -> Optional[TJSONExt]:
        if self._starting or self._started:
            raise Exception("Already initialized and not finalized")
        self._starting = True
        result = await self._initialize()
        self._started = True
        self._starting = False
        return result

    async def _initialize(self) -> Optional[TJSONExt]:
        raise NotImplementedError

    async def finalize(
        self,
        result: Optional[TJSONExt],
        error: Optional[Any] = None,
        ttl_sec: Optional[float] = None,
    ) -> None:
        if not self._starting and not self._started:
            raise Exception("Not initialized")
        try:
            return await self._finalize(result=result, error=error, ttl_sec=ttl_sec)
        finally:
            self._starting = False
            self._started = False

    async def _finalize(
        self,
        result: Optional[TJSONExt],
        error: Optional[Any] = None,
        ttl_sec: Optional[float] = None,
    ) -> None:
        raise NotImplementedError


@attr.s(auto_attribs=True, slots=True)
class EntityCacheEntryManagerAsync(EntityCacheEntryManagerAsyncBase):
    async def _initialize(self) -> Optional[TJSONExt]:
        if not self.allow_cache_read:
            return None
        return await self.cache_engine._get_from_cache(
            local_key_rep=self.local_key_rep,
            new_ttl_sec=self.read_extend_ttl_sec,
        )

    async def _finalize(
        self,
        result: Optional[TJSONExt],
        error: Optional[Any] = None,
        ttl_sec: Optional[float] = None,
    ) -> None:
        if result is None:
            if error is None:
                return None  # nothing to save at all
            result = []  # TODO: allow an actually null result (metadata-only)
        return await self.cache_engine._update_cache(
            local_key_rep=self.local_key_rep,
            result=result,
            metadata=dict(error=error),
            ttl_sec=ttl_sec or self.write_ttl_sec,
        )


class RedisCacheLockWrapped(RedisCacheLock):
    """Local profiling additions to the RCL"""

    @generic_profiler_async("qcache-locked-get-data-slave")  # type: ignore  # TODO: fix
    async def get_data_slave(self) -> Optional[bytes]:
        return await super().get_data_slave()

    @generic_profiler_async("qcache-locked-get-data-main")  # type: ignore  # TODO: fix
    async def _get_data(self) -> Any:
        return await super()._get_data()

    @generic_profiler_async("qcache-locked-get-data-wait")  # type: ignore  # TODO: fix
    async def _wait_for_result(self, sub: Any) -> Any:
        return await super()._wait_for_result(sub)


@attr.s(auto_attribs=True, slots=True)
class EntityCacheEntryLockedManagerAsync(EntityCacheEntryManagerAsyncBase):
    _rcl: Optional[RedisCacheLock] = None
    _history_holder: Optional[HistoryHolder] = None

    @contextlib.asynccontextmanager
    async def _with_redis(self, *, master: bool = True, exclusive: bool = True) -> AsyncGenerator[Redis, None]:
        cache_engine = self.cache_engine
        # This does not allow retries at this layer; would need the source
        # sentinel object for that.
        cli = cache_engine.rc if master else (cache_engine.rc_slave or cache_engine.rc)
        assert cli is not None
        try:
            yield cli
        finally:
            await cli.close()

    def _make_rcl(self) -> Tuple[RedisCacheLock, HistoryHolder]:
        cache_engine = self.cache_engine
        local_key_rep = self.local_key_rep
        local_key_rep.validate()
        key = local_key_rep.key_parts_hash
        key_root = cache_engine._get_key_root()
        # Keys structure: `resource_tag + ('/data:'|'/notif:'|'/lock:') + key`
        # e.g. 'bic_dp_ce_l__iyirei5oqew0c__query_cache//data:657f18518eaa2f41307895e18c3ba0d12d97b8a23c6de3966f52c6ba39a07ee4'
        # NOTE: not using `cache_engine._get_key_query_cache_entry` at all
        history_holder = HistoryHolder()
        rcl = RedisCacheLockWrapped(
            key=key,
            client_acm=self._with_redis,
            resource_tag=key_root,
            lock_ttl_sec=15.0,
            data_ttl_sec=self.write_ttl_sec,
            # Not exactly the same semantics, but generally the same order of magnitude.
            # If a more precise timeout is desired, `rcl.initialize` can be
            # timeouted safely (as long as `finalize` is called).
            network_call_timeout_sec=self.cache_engine.CACHE_GET_TIMEOUT_SEC,
            enable_background_tasks=self.cache_engine.CACHE_SAVE_BACKGROUND,
            enable_slave_get=self.cache_engine.CACHE_CHECK_SLAVE,
            debug_log=history_holder,
        )
        return rcl, history_holder

    @property
    def rcl(self) -> RedisCacheLock:
        rcl = self._rcl
        if rcl is None:
            rcl, history_holder = self._make_rcl()
            self._rcl = rcl
            self._history_holder = history_holder
        return rcl

    @generic_profiler_async("qcache-locked-initialize")  # type: ignore  # TODO: fix
    async def _initialize(self) -> Optional[TJSONExt]:
        """
        WARNING: partially reimplements `EntityCacheEngineAsync._get_from_cache`.
        """
        local_key_rep = self.local_key_rep
        cache_engine = self.cache_engine
        rcl = self.rcl

        full_key, details = cache_engine._make_full_key_and_log_details(
            local_key_rep=local_key_rep,
            new_ttl_sec=self.read_extend_ttl_sec,
        )

        result = await rcl.initialize()

        situation = rcl._situation
        LOGGER.debug("Cache-locked situation=%r, have result: %r", situation, result is not None)

        cache_entry = self.cache_engine._make_result_cache_entry(
            local_key_rep=self.local_key_rep,
            cache_entry_redis_data=result,
            details=details,
        )
        if cache_entry is None:
            return None

        result_data = cache_entry.make_result_data()
        cache_engine._log_after_read(full_key=full_key, cache_entry=cache_entry)
        return result_data

    @generic_profiler_async("qcache-locked-finalize")  # type: ignore  # TODO: fix
    async def _finalize(
        self,
        result: Optional[TJSONExt],
        error: Optional[Any] = None,
        ttl_sec: Optional[float] = None,
    ) -> None:
        rcl = self._rcl
        assert rcl is not None, "should have been created by now"

        ttl_sec = ttl_sec or self.write_ttl_sec
        if result is None:
            result = []  # TODO: allow an actually null result (metadata-only)

        update_request = self.cache_engine._make_cache_update_request(
            local_key_rep=self.local_key_rep,
            result=result,
            compress=False,
            compress_alg=None,  # `compress=False` to prevent CPU-bound operation in event loop
            metadata=dict(error=error),
            ttl_sec=ttl_sec,
        )
        try:
            await rcl.finalize(update_request.entry_data, ttl_sec=ttl_sec)
        except BaseException as err:
            self.cache_engine._log_finalize_failed(update_request=update_request, err=err)
            self.save_history(error=error, save_error=err)
            raise

        self.cache_engine._log_after_finalize(update_request=update_request)
        self.save_history(error=error, save_error=None)

    def save_history(self, error: Any = None, save_error: Any = None) -> None:
        history_holder = self._history_holder
        assert history_holder
        events = history_holder.history
        history_holder.history = []
        min_ts = min((ts for ts, _, _ in events), default=0)
        events_readable = ";  ".join(f"{round(ts - min_ts, 3)}: {msg}" for ts, msg, _ in events)
        LOGGER.debug(
            "RedisCacheLock logs: %d events, err=%r, save_err=%r: %s",
            len(events),
            error,
            save_error,
            events_readable,
            extra=dict(redis_cache_lock_logs=events),
        )


_ECE_TV = TypeVar("_ECE_TV", bound="EntityCacheEngineBase")


@attr.s(auto_attribs=True)
class EntityCacheEngineBase:
    entity_id: str

    DEFAULT_TTL_SEC: float = 60.0
    DELETE_SCAN_CHUNK_SIZE: int = 100
    DEFAULT_COMPRESS_ALG: CompressAlg = CompressAlg.GZIP
    MIN_BYTES_TO_COMPRESS: int = 120

    # Old keys, effectively, were 'dataset:{entity_id}:query_cache:'
    key_prefix_tpl: str = "bic_dp_ce__{entity_id}__query_cache/"  # dl_core data_processing cache_engine
    data_keys_suffix: str = ""

    def clone(self: _ECE_TV, **updates: Any) -> _ECE_TV:
        return attr.evolve(self, **updates)

    # Keys section
    def _get_key_root(self) -> str:
        return self.key_prefix_tpl.format(entity_id=self.entity_id)

    def _get_key_query_cache_entry(self, local_key_rep: LocalKeyRepresentation) -> str:
        local_key_rep.validate()
        # noinspection PyProtectedMember
        return "{entity_root_key}/{local_key_rep_str}".format(
            entity_root_key=self._get_key_root(),
            local_key_rep_str=local_key_rep.key_parts_hash,
        )

    def _get_all_keys_pattern(self) -> str:
        """Returns pattern for all cached keys for the entity"""
        return self._get_key_root() + self.data_keys_suffix + "*"

    def _make_cache_update_request(
        self,
        local_key_rep: LocalKeyRepresentation,
        result: TJSONExt,
        compress: bool = False,
        compress_alg: Optional[CompressAlg] = None,
        ttl_sec: Optional[float] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> _CacheUpdateRequest:
        if ttl_sec is None:
            ttl_sec = self.DEFAULT_TTL_SEC

        if compress_alg is None:
            compress_alg = self.DEFAULT_COMPRESS_ALG

        full_key = self._get_key_query_cache_entry(local_key_rep=local_key_rep)

        # noinspection PyProtectedMember
        entry_data, details = ResultCacheEntry(
            key_parts_str=local_key_rep.key_parts_str,
            result_data=result,
            metadata=metadata,
        ).to_redis_data_and_log_details(
            compress,
            compress_alg=compress_alg,
        )

        # Enrich the debug/stats details for after-success log:
        details = dict(
            details,
            cache_key=full_key,
            cache_ttl_sec=ttl_sec,
            cache_compress_alg=compress_alg.name,
        )

        return _CacheUpdateRequest(
            full_key=full_key,
            ttl_sec=ttl_sec,
            entry_data=entry_data,
            details=details,
        )

    def _make_full_key_and_log_details(
        self,
        local_key_rep: LocalKeyRepresentation,
        new_ttl_sec: Optional[float] = None,
    ) -> Tuple[str, dict]:
        full_key = self._get_key_query_cache_entry(local_key_rep)

        details = dict(
            cache_key=full_key,
            cache_refresh_on_read=new_ttl_sec is not None,
            cache_ttl_sec=new_ttl_sec,
        )

        return full_key, details

    def _make_result_cache_entry(
        self,
        local_key_rep: LocalKeyRepresentation,
        cache_entry_redis_data: Optional[bytes],
        details: Dict[str, Any],
        allow_error_entry: bool = False,
    ) -> Optional[ResultCacheEntry]:
        if not cache_entry_redis_data:
            self._log_cache_miss(reason="no_data", details=details)
            return None

        try:
            cache_entry = ResultCacheEntry.from_redis_data(
                cache_entry_redis_data,
                details=details,
            )
        except CachedEntryPackageVersionMismatch as err:
            self._log_cache_miss(
                reason="bi_common_was_updated",
                details=details,
                message=str(err),
            )
            return None

        if cache_entry.key_parts_str != local_key_rep.key_parts_str:
            self._log_cache_miss(
                reason="query_hash_collision",
                details=details,
                message="Query hash collision: {!r} != {!r}".format(
                    cache_entry.key_parts_str, local_key_rep.key_parts_str
                ),
            )
            return None

        if not allow_error_entry:
            if cache_entry.is_error:
                metadata = cache_entry.metadata
                assert metadata is not None
                self._log_cache_miss(
                    reason="error_mark",
                    details=dict(details, error=metadata["error"]),
                )
                return None

        return cache_entry

    def _log_after_save(self, update_request: _CacheUpdateRequest) -> None:
        LOGGER.info("Cache saved at %r", update_request.full_key, extra=update_request.details)

    def _log_save_failed(self, update_request: _CacheUpdateRequest, err: BaseException) -> None:
        LOGGER.info("Cache save failed at %r: %r", update_request.full_key, err, extra=update_request.details)

    def _log_after_finalize(self, update_request: _CacheUpdateRequest) -> None:
        LOGGER.info("Cache finalized at %r", update_request.full_key, extra=update_request.details)

    def _log_finalize_failed(self, update_request: _CacheUpdateRequest, err: BaseException) -> None:
        LOGGER.info("Cache finalize failed at %r: %r", update_request.full_key, err, extra=update_request.details)

    def _log_after_read(self, full_key: str, cache_entry: ResultCacheEntry) -> None:
        LOGGER.info("Cache read from %r", full_key, extra=cache_entry.details)

    def _log_cache_miss(self, reason: str, details: Dict[str, Any], message: Optional[str] = None) -> None:
        full_key = details["cache_key"]
        details = dict(details, cache_miss_reason=reason)
        LOGGER.info("Cache miss at %r: %r: %r", full_key, reason, message or "-", extra=details)

    def _log_cache_timeout(self, timeout: float, details: Dict[str, Any], message: Optional[str] = None) -> None:
        full_key = details["cache_key"]
        details = dict(
            details,
            cache_read_timeout=timeout,
        )
        LOGGER.info("Cache timeout at %r: %r", full_key, message or "-", extra=details)


@attr.s(auto_attribs=True)
class EntityCacheEngineAsync(EntityCacheEngineBase):
    """
    ...

    WARNING: this is, mostly, an async+await copy of the `EntityCacheEngine`,
    requires a different class of the redis client,
    does not compress on save.
    """

    rc: Redis = attr.ib(kw_only=True)
    rc_slave: Optional[Redis] = None

    CACHE_GET_TIMEOUT_SEC: float = 5.0
    CACHE_SAVE_TIMEOUT_SEC: float = 30.0  # should not be a problem since it happens in background
    CACHE_SAVE_BACKGROUND: bool = True
    CACHE_CHECK_SLAVE: bool = False  # does not go well when the slave is in another country.

    entity_cache_entry_manager_cls: Type[EntityCacheEntryManagerAsyncBase] = EntityCacheEntryManagerAsync

    def as_lockable(self) -> EntityCacheEngineAsync:
        """
        Use the entry manager that synchronizes the data generation over a lock,
        and a different key prefix just in case.
        """
        return self.clone(
            entity_cache_entry_manager_cls=EntityCacheEntryLockedManagerAsync,
            # dl_core data_processing cache_engine locked
            key_prefix_tpl="bic_dp_ce_l__{entity_id}__query_cache/",
            data_keys_suffix="/data:*",
        )

    # Redis helpers
    @generic_profiler_async("qcache-invalidate")  # type: ignore  # TODO: fix
    async def _delete_keys_by_pattern(self, pattern: str) -> None:
        """SCAN is used to prevent redis lockout with KEYS command."""
        LOGGER.info("Invalidating cache for entity %s", self.entity_id)
        all_keys: Deque[bytes] = collections.deque()
        cur = 0
        started = False
        while cur or not started:
            started = True
            cur, keys = await self.rc.scan(cur, match=pattern, count=self.DELETE_SCAN_CHUNK_SIZE)
            for single_key in keys:
                all_keys.append(single_key)

        pipeline = self.rc.pipeline()
        for key in all_keys:
            pipeline.delete(key)
        await pipeline.execute()

    @generic_profiler_async("qcache-update")  # type: ignore  # TODO: fix
    async def _update_cache(
        self,
        local_key_rep: LocalKeyRepresentation,
        result: TJSONExt,
        metadata: Optional[Dict[str, Any]] = None,
        ttl_sec: Optional[float] = None,
    ) -> None:
        update_request = self._make_cache_update_request(
            local_key_rep=local_key_rep,
            result=result,
            compress=False,
            compress_alg=None,  # `compress=False` to prevent CPU-bound operation in event loop
            ttl_sec=ttl_sec,
        )

        if self.CACHE_SAVE_BACKGROUND:
            asyncio.create_task(
                self._redis_set(update_request),
            )
        else:
            try:
                await asyncio.wait_for(
                    self._redis_set(update_request),
                    timeout=self.CACHE_SAVE_TIMEOUT_SEC,
                )
            except asyncio.TimeoutError as err:
                self._log_save_failed(update_request=update_request, err=err)

    @generic_profiler_async("qcache-write-redis-exec")  # type: ignore  # TODO: fix
    async def _redis_set(self, update_request: _CacheUpdateRequest) -> None:
        try:
            rcli = self.rc
            pipe = rcli.pipeline()
            pipe.set(update_request.full_key, update_request.entry_data)
            pipe.pexpire(update_request.full_key, int(update_request.ttl_sec * 1000))
            await pipe.execute()
        except Exception as err:
            self._log_save_failed(update_request=update_request, err=err)
        else:
            self._log_after_save(update_request=update_request)

    @generic_profiler_async("qcache-get")  # type: ignore  # TODO: fix
    async def _get_from_cache(
        self,
        local_key_rep: LocalKeyRepresentation,
        new_ttl_sec: Optional[float] = None,
    ) -> Optional[TJSONExt]:
        full_key, details = self._make_full_key_and_log_details(local_key_rep=local_key_rep, new_ttl_sec=new_ttl_sec)

        read_timeout = self.CACHE_GET_TIMEOUT_SEC
        try:
            cache_entry_redis_data = await asyncio.wait_for(
                self._redis_get(full_key, new_ttl_sec=new_ttl_sec),
                timeout=read_timeout,
            )
        except asyncio.TimeoutError:
            self._log_cache_timeout(timeout=read_timeout, details=details)
            return None

        cache_entry = self._make_result_cache_entry(
            local_key_rep=local_key_rep,
            cache_entry_redis_data=cache_entry_redis_data,
            details=details,
        )

        if cache_entry is None:  # cache miss, supposed to be logged by now.
            return None

        # deserializes data and also mutates `cache_entry.details`.
        result = cache_entry.make_result_data()
        self._log_after_read(full_key=full_key, cache_entry=cache_entry)
        return result

    @generic_profiler_async("qcache-read-redis-exec")  # type: ignore  # TODO: fix
    async def _redis_get(
        self,
        full_key: str,
        new_ttl_sec: Optional[float] = None,
    ) -> Optional[bytes]:
        rcli = self.rc
        if new_ttl_sec is not None:
            pipe = rcli.pipeline()
            pipe.pexpire(full_key, int(new_ttl_sec * 1000))
            pipe.get(full_key)
            _, result = await pipe.execute()
            return result

        rc_slave = self.rc_slave
        if self.CACHE_CHECK_SLAVE and rc_slave is not None:
            result = await rc_slave.get(full_key)
            if result is not None:
                return result

        result = await rcli.get(full_key)
        return result

    async def invalidate_all(self) -> None:
        await self._delete_keys_by_pattern(self._get_all_keys_pattern())

    def get_cache_entry_manager(
        self,
        local_key_rep: LocalKeyRepresentation,
        **kwargs: Any,
    ) -> Optional[EntityCacheEntryManagerAsyncBase]:
        return self.entity_cache_entry_manager_cls(
            local_key_rep=local_key_rep,
            cache_engine=self,
            **kwargs,
        )

lib/dl_core/dl_core/data_processing/cache/engine.py
datalens-tech/datalens-backend
from __future__ import annotations


class CacheException(Exception):
    pass


class CachedEntryPackageVersionMismatch(CacheException):
    pass


class CacheKeyValidationError(CacheException):
    pass


class CachePreparationFailed(CacheException):
    pass

lib/dl_core/dl_core/data_processing/cache/exc.py
datalens-tech/datalens-backend
from __future__ import annotations

import hashlib
from typing import (
    Any,
    Hashable,
    NamedTuple,
    Optional,
    Tuple,
)

import attr

from dl_core.data_processing.cache.exc import CacheKeyValidationError


@attr.s(frozen=True)
class CacheTTLConfig:
    ttl_sec_direct: int = attr.ib(default=60)
    ttl_sec_materialized: int = attr.ib(default=3600)  # TODO remove

    def clone(self, **kwargs: Any) -> CacheTTLConfig:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True)
class CacheTTLInfo:
    ttl_sec: int = attr.ib()
    refresh_ttl_on_read: bool = attr.ib()

    def clone(self, **kwargs: Any) -> CacheTTLInfo:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True, auto_attribs=True)
class BIQueryCacheOptions:
    cache_enabled: bool
    key: Optional[LocalKeyRepresentation]
    ttl_sec: int
    refresh_ttl_on_read: bool


class DataKeyPart(NamedTuple):
    part_type: str
    part_content: Hashable

    def stringify(self) -> str:
        if isinstance(self.part_content, DataKeyPart):
            part_content_str = self.part_content.stringify()
        else:
            part_content_str = str(self.part_content)
        return f"{self.part_type}:{part_content_str}"

    def validate(self) -> None:
        if not self.part_type:
            raise CacheKeyValidationError("Data key part type is empty")
        if self.part_content is None:
            raise CacheKeyValidationError(f"Data key part content is None for part type {self.part_type}")


@attr.s(slots=True)
class LocalKeyRepresentation:
    _key_parts: Tuple[DataKeyPart, ...] = attr.ib(default=())
    _key_parts_str: Optional[str] = attr.ib(init=False, default=None)
    _key_parts_hash: Optional[str] = attr.ib(init=False, default=None)

    def validate(self) -> None:
        if len(self.key_parts) == 0:
            raise CacheKeyValidationError("key_parts is empty")
        for key_part in self.key_parts:
            key_part.validate()

    @property
    def key_parts(self) -> Tuple[DataKeyPart, ...]:
        return self._key_parts

    @property
    def key_parts_str(self) -> str:
        if self._key_parts_str is None:
            self._key_parts_str = ";".join(part.stringify() for part in self.key_parts)
        return self._key_parts_str

    @property
    def key_parts_hash(self) -> str:
        if self._key_parts_hash is None:
            self._key_parts_hash = hashlib.sha256(self.key_parts_str.encode()).hexdigest()
        return self._key_parts_hash

    def extend(self, part_type: str, part_content: Hashable) -> LocalKeyRepresentation:
        new_part = DataKeyPart(part_type=part_type, part_content=part_content)
        return LocalKeyRepresentation(key_parts=self.key_parts + (new_part,))

    def multi_extend(self, *parts: DataKeyPart) -> LocalKeyRepresentation:
        return LocalKeyRepresentation(key_parts=self.key_parts + parts)

lib/dl_core/dl_core/data_processing/cache/primitives.py
datalens-tech/datalens-backend
from __future__ import annotations

import enum
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Awaitable,
    Callable,
    ClassVar,
    Optional,
    Tuple,
)

import attr

from dl_core.data_processing.streaming import AsyncChunked


if TYPE_CHECKING:
    from dl_core.data_processing.cache.engine import (
        EntityCacheEngineAsync,
        EntityCacheEntryManagerAsyncBase,
    )
    from dl_core.data_processing.cache.primitives import BIQueryCacheOptions
    from dl_core.data_processing.types import TJSONExtChunkStream
    import dl_core.data_source  # noqa
    from dl_core.services_registry import ServicesRegistry  # noqa
    import dl_core.us_dataset  # noqa


LOGGER = logging.getLogger(__name__)


@enum.unique
class CacheSituation(enum.IntEnum):
    cache_disabled = enum.auto()
    cache_error = enum.auto()
    full_hit = enum.auto()
    generated = enum.auto()
    # TODO: discern some other cases:
    #  * now handled as cache miss -> `generated`;
    #    * cache error-marker
    #    * cache read timeout
    #    * cache lock timeout


@attr.s
class CacheProcessingHelper:
    entity_id: str = attr.ib(kw_only=True)
    _service_registry: ServicesRegistry = attr.ib(kw_only=True)
    _cache_engine: Optional[EntityCacheEngineAsync] = attr.ib(init=False, default=None)

    error_ttl_sec: ClassVar[float] = 1.5

    def __attrs_post_init__(self) -> None:
        cache_engine_factory = self._service_registry.get_cache_engine_factory()
        if cache_engine_factory is not None:
            self._cache_engine = cache_engine_factory.get_cache_engine(entity_id=self.entity_id)

    async def get_cache_entry_manager(
        self,
        *,
        cache_options: Optional[BIQueryCacheOptions],
        allow_cache_read: bool = True,
        locked_cache: bool = False,
    ) -> Optional[EntityCacheEntryManagerAsyncBase]:
        if cache_options is None:
            return None
        if not cache_options.cache_enabled:
            return None

        cache_engine = self._cache_engine
        if cache_engine is None:
            return None

        assert cache_options.key is not None
        local_key_rep = cache_options.key

        if locked_cache:
            cache_engine = cache_engine.as_lockable()

        return cache_engine.get_cache_entry_manager(
            local_key_rep=local_key_rep,
            write_ttl_sec=cache_options.ttl_sec,
            read_extend_ttl_sec=(cache_options.ttl_sec if cache_options.refresh_ttl_on_read else None),
            allow_cache_read=allow_cache_read,
        )

    def _dump_error_for_cache(self, err: BaseException) -> Any:
        return str(err)

    async def run_with_cache(
        self,
        *,
        generate_func: Callable[[], Awaitable[Optional[TJSONExtChunkStream]]],
        cache_options: BIQueryCacheOptions,
        allow_cache_read: bool = True,
        use_locked_cache: bool = False,
    ) -> Tuple[CacheSituation, Optional[TJSONExtChunkStream]]:
        cem = await self.get_cache_entry_manager(
            cache_options=cache_options,
            allow_cache_read=allow_cache_read,
            locked_cache=use_locked_cache,
        )
        if cem is None:
            result = await generate_func()
            return CacheSituation.cache_disabled, result

        # As a reference, the `cem` usage should be equivalent to this:
        #     result = None
        #     try:
        #         cache_result = await cem.initialize()
        #         if cache_result is None:
        #             result = await generate()
        #     finally:
        #         await cem.finalize(result=result, error=dump_error(sys.exc_info()))
        # The rest is just for verbosity.

        result_iter: Optional[TJSONExtChunkStream]
        sync_result_iter = None
        try:
            sync_result_iter = await cem.initialize()
        except BaseException as err:  # noqa
            LOGGER.error("Error during checking cache", exc_info=True)
            try:
                err_serializable = self._dump_error_for_cache(err)
                await cem.finalize(result=None, error=err_serializable)
            except Exception:  # Not skipping `CancelledError` here
                LOGGER.error("Error during finalizing cache (after an error during checking cache)", exc_info=True)
            LOGGER.debug("Going to db selector due to cache read error")
            result = await generate_func()
            return CacheSituation.cache_error, result

        if sync_result_iter is not None:
            result_data_inmem = list(sync_result_iter)  # type: ignore  # TODO: fix
            result_iter = AsyncChunked.from_chunked_iterable([result_data_inmem])
            try:
                await cem.finalize(result=None)
            except Exception:  # Not skipping `CancelledError` here
                LOGGER.error("Error during finalizing cache (after a cache hit)", exc_info=True)
            return CacheSituation.full_hit, result_iter

        LOGGER.info("Got selector result from cache engine: not found")

        result_as_list = None
        try:
            result_iter = await generate_func()
            if result_iter is not None:
                # Cache is not (currently) streamed, so have to collect the data.
                # And just in case, do this within the db-query try-block.
                result_as_list = await result_iter.all()
                result_iter = AsyncChunked.from_chunked_iterable([result_as_list])
        except BaseException as err:
            LOGGER.info("Error during generate_func, saving an error marker to cache")
            err_serializable = self._dump_error_for_cache(err)
            try:
                await cem.finalize(
                    result=None,
                    # short-lived error-result cache,
                    # to make it possible to make an informed decision,
                    # and to unlock the cache lock (if necessary)
                    error=err_serializable,
                    ttl_sec=self.error_ttl_sec,
                )
            except Exception:
                LOGGER.error("Error during finalizing cache (after a generate error)", exc_info=True)
            raise

        try:
            await cem.finalize(
                result=result_as_list,
            )
        except Exception:
            LOGGER.error("Error during finalizing cache (after a generate success)", exc_info=True)

        return CacheSituation.generated, result_iter

lib/dl_core/dl_core/data_processing/cache/processing_helper.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    Collection,
    Optional,
)

import attr
from sqlalchemy.exc import DatabaseError

from dl_constants.enums import DataSourceRole
from dl_core.data_processing.cache.exc import CachePreparationFailed
from dl_core.data_processing.cache.primitives import (
    BIQueryCacheOptions,
    CacheTTLConfig,
    CacheTTLInfo,
    DataKeyPart,
    LocalKeyRepresentation,
)
from dl_core.serialization import hashable_dumps
from dl_core.us_connection_base import (
    ConnectionBase,
    ExecutorBasedMixin,
)


if TYPE_CHECKING:
    from sqlalchemy.engine.default import DefaultDialect
    from sqlalchemy.sql import Select

    from dl_constants.enums import UserDataType
    from dl_constants.types import TJSONExt
    from dl_core.data_processing.prepared_components.primitives import PreparedMultiFromInfo
    from dl_core.data_processing.stream_base import DataStreamBase
    from dl_core.data_source.base import DataSource
    from dl_core.us_dataset import Dataset
    from dl_core.us_manager.local_cache import USEntryBuffer


LOGGER = logging.getLogger(__name__)


@attr.s
class CacheOptionsBuilderBase:
    default_ttl_config: CacheTTLConfig = attr.ib(factory=CacheTTLConfig)
    _is_bleeding_edge_user: bool = attr.ib(default=False)

    def get_actual_ttl_config(
        self,
        connection: ConnectionBase,
        dataset: Optional[Dataset] = None,
    ) -> CacheTTLConfig:
        ctc = self.default_ttl_config

        if isinstance(connection, ExecutorBasedMixin):
            override = connection.cache_ttl_sec_override
            if override is not None:
                ctc = ctc.clone(ttl_sec_direct=override)

        return ctc

    @staticmethod
    def get_query_str_for_cache(query: Select, dialect: DefaultDialect) -> str:
        try:
            compiled_query = query.compile(dialect=dialect)
        except DatabaseError as err:
            raise CachePreparationFailed from err

        if isinstance(compiled_query.params, dict):
            ordered_params = sorted(
                compiled_query.params.items(),
                key=lambda item: item[0],
            )
        else:
            ordered_params = compiled_query.params

        return ";".join(
            (
                str(compiled_query),
                str(ordered_params),
            )
        )

    @staticmethod
    def config_to_ttl_info(
        ttl_config: CacheTTLConfig,
        is_materialized: bool,
        data_source_list: Optional[Collection[DataSource]] = None,
    ) -> CacheTTLInfo:
        ttl_info = CacheTTLInfo(
            ttl_sec=ttl_config.ttl_sec_direct,
            refresh_ttl_on_read=False,
        )

        if not is_materialized:
            return ttl_info
        if data_source_list is None:
            return ttl_info

        data_dump_id_list = [dsrc.data_dump_id for dsrc in data_source_list]
        if not all(data_dump_id_list):
            return ttl_info

        # Materialized
        return ttl_info.clone(
            ttl_sec=ttl_config.ttl_sec_materialized,
            refresh_ttl_on_read=True,
        )

    def get_cache_ttl_info(
        self,
        is_materialized: bool,
        data_source_list: Collection[DataSource],
        # For future use
        dataset: Optional[Dataset] = None,  # noqa
    ) -> CacheTTLInfo:
        # TODO FIX: Assert that there is no connection divergence or migrate to joint data source info
        assert data_source_list, "Cannot generate cache options for empty source list"
        actual_connection = next(iter({dsrc.connection for dsrc in data_source_list}))

        ttl_config = self.get_actual_ttl_config(
            connection=actual_connection,
            dataset=dataset,
        )
        return self.config_to_ttl_info(
            ttl_config=ttl_config,
            is_materialized=is_materialized,
            data_source_list=data_source_list,
        )


@attr.s
class CacheOptionsBuilderDataProcessor(CacheOptionsBuilderBase):
    def get_cache_options_for_stream(
        self,
        stream: DataStreamBase,
        dataset: Optional[Dataset] = None,
    ) -> BIQueryCacheOptions:
        ttl_info = self.get_cache_ttl_info(
            is_materialized=stream.meta.is_materialized,
            data_source_list=stream.meta.data_source_list,
            dataset=dataset,
        )
        key = stream.data_key
        cache_enabled = key is not None
        return BIQueryCacheOptions(
            cache_enabled=cache_enabled,
            key=key,
            ttl_sec=ttl_info.ttl_sec,
            refresh_ttl_on_read=ttl_info.refresh_ttl_on_read,
        )


@attr.s
class SelectorCacheOptionsBuilder(CacheOptionsBuilderBase):
    _is_bleeding_edge_user: bool = attr.ib(default=False)
    _us_entry_buffer: USEntryBuffer = attr.ib(kw_only=True)

    def get_cache_enabled(self, joint_dsrc_info: PreparedMultiFromInfo) -> bool:
        assert joint_dsrc_info.data_source_list is not None
        cache_enabled = all(dsrc.cache_enabled for dsrc in joint_dsrc_info.data_source_list)
        return cache_enabled

    def get_cache_options(
        self,
        role: DataSourceRole,
        joint_dsrc_info: PreparedMultiFromInfo,
        query: Select,
        user_types: list[UserDataType],
        dataset: Dataset,
    ) -> BIQueryCacheOptions:
        """Returns cache key, TTL for new entries, refresh TTL flag"""

        merged_data_dump_id = "N/A"
        assert joint_dsrc_info.data_source_list is not None
        if role != DataSourceRole.origin:
            data_dump_id_list = [dsrc.data_dump_id for dsrc in joint_dsrc_info.data_source_list]
            if all(data_dump_id_list):
                merged_data_dump_id = "+".join(data_dump_id_list)  # type: ignore

        compiled_query = self.get_query_str_for_cache(
            query=query,
            dialect=joint_dsrc_info.query_compiler.dialect,
        )
        local_key_rep: Optional[LocalKeyRepresentation] = self.make_data_select_cache_key(
            joint_dsrc_info=joint_dsrc_info,
            compiled_query=compiled_query,
            user_types=user_types,
            data_dump_id=merged_data_dump_id,
            is_bleeding_edge_user=self._is_bleeding_edge_user,
        )
        ttl_info = self.get_cache_ttl_info(
            is_materialized=role != DataSourceRole.origin,
            data_source_list=joint_dsrc_info.data_source_list,
            dataset=dataset,
        )
        cache_enabled = self.get_cache_enabled(joint_dsrc_info=joint_dsrc_info)
        if not cache_enabled:
            local_key_rep = None

        return BIQueryCacheOptions(
            cache_enabled=cache_enabled,
            key=local_key_rep,
            ttl_sec=ttl_info.ttl_sec,
            refresh_ttl_on_read=ttl_info.refresh_ttl_on_read,
        )

    def make_data_select_cache_key(
        self,
        joint_dsrc_info: PreparedMultiFromInfo,
        compiled_query: str,
        user_types: list[UserDataType],
        data_dump_id: str,
        is_bleeding_edge_user: bool,
    ) -> LocalKeyRepresentation:
        assert joint_dsrc_info.target_connection_ref is not None
        target_connection = self._us_entry_buffer.get_entry(joint_dsrc_info.target_connection_ref)
        assert isinstance(target_connection, ConnectionBase)
        connection_id = target_connection.uuid
        assert connection_id is not None

        local_key_rep = target_connection.get_cache_key_part()
        if joint_dsrc_info.db_name is not None:
            # FIXME: Replace with key parts for every participating dsrc
            # For now db_name will be duplicated in some source types
            # (one from connection, one from source)
            local_key_rep = local_key_rep.extend(part_type="db_name", part_content=joint_dsrc_info.db_name)
        local_key_rep = local_key_rep.extend(part_type="query", part_content=str(compiled_query))
        local_key_rep = local_key_rep.extend(part_type="user_types", part_content=tuple(user_types or ()))
        local_key_rep = local_key_rep.extend(part_type="data_dump_id", part_content=data_dump_id or "N/A")
        local_key_rep = local_key_rep.extend(
            part_type="is_bleeding_edge_user",
            part_content=is_bleeding_edge_user,
        )

        return local_key_rep


@attr.s
class DashSQLCacheOptionsBuilder(CacheOptionsBuilderBase):
    _is_bleeding_edge_user: bool = attr.ib(default=False)

    def get_cache_enabled(self, conn: ConnectionBase) -> bool:
        return conn.allow_cache

    def get_cache_options(
        self,
        conn: ConnectionBase,
        query_text: str,
        params: TJSONExt,
        db_params: dict[str, str],
        connector_specific_params: TJSONExt,
    ) -> BIQueryCacheOptions:
        cache_enabled = self.get_cache_enabled(conn=conn)
        ttl_config = self.get_actual_ttl_config(connection=conn, dataset=None)
        ttl_info = self.config_to_ttl_info(ttl_config, is_materialized=False, data_source_list=None)

        local_key_rep: Optional[LocalKeyRepresentation] = conn.get_cache_key_part()
        local_key_rep = local_key_rep.multi_extend(
            DataKeyPart(part_type="query", part_content=query_text),
            DataKeyPart(part_type="params", part_content=hashable_dumps(params)),
            DataKeyPart(part_type="db_params", part_content=tuple(db_params.items())),
            DataKeyPart(part_type="connector_specific_params", part_content=hashable_dumps(connector_specific_params)),
            DataKeyPart(part_type="is_bleeding_edge_user", part_content=self._is_bleeding_edge_user),
        )
        if not cache_enabled:
            local_key_rep = None

        return BIQueryCacheOptions(
            cache_enabled=cache_enabled,
            key=local_key_rep,
            ttl_sec=ttl_info.ttl_sec,
            refresh_ttl_on_read=ttl_info.refresh_ttl_on_read,
        )

lib/dl_core/dl_core/data_processing/cache/utils.py
datalens-tech/datalens-backend
from __future__ import annotations

import enum
import logging
import re
import sys
import time
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncIterable,
    Dict,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
    Union,
)

import attr
import sqlalchemy as sa

from dl_api_commons.reporting.models import (
    QueryExecutionCacheInfoReportingRecord,
    QueryExecutionEndReportingRecord,
    QueryExecutionStartReportingRecord,
)
from dl_constants.types import TJSONExt  # not under `TYPE_CHECKING`, need to define new type aliases.
from dl_core import exc
from dl_core.base_models import WorkbookEntryLocation
from dl_core.connection_executors.common_base import ConnExecutorQuery
from dl_core.connectors.base.dashsql import get_custom_dash_sql_key_names
from dl_core.data_processing.cache.processing_helper import (
    CacheProcessingHelper,
    CacheSituation,
)
from dl_core.data_processing.cache.utils import DashSQLCacheOptionsBuilder
from dl_core.data_processing.streaming import (
    AsyncChunked,
    chunkify_by_one,
)
from dl_core.us_connection_base import ExecutorBasedMixin
from dl_core.utils import (
    compile_query_for_debug,
    make_id,
    sa_plain_text,
)


if TYPE_CHECKING:
    from dl_core.connection_executors.async_base import (
        AsyncConnExecutorBase,
        AsyncExecutionResult,
    )
    from dl_core.data_processing.cache.primitives import BIQueryCacheOptions
    from dl_core.data_processing.types import TJSONExtChunkStream
    from dl_core.services_registry import ServicesRegistry
    from dl_core.us_connection_base import ConnectionBase


LOGGER = logging.getLogger(__name__)


@enum.unique
class DashSQLEvent(enum.Enum):
    metadata = "metadata"
    row = "row"  # to be deprecated
    rowchunk = "rowchunk"
    error = "error"
    footer = "footer"


TRow = Tuple[TJSONExt, ...]
TMeta = Dict[str, TJSONExt]
# # More correct but mymy couldn't:
# TResultEvent = Union[
#     Tuple[Literal['metadata'], TMeta],
#     Tuple[Literal['row'], TRow],
#     Tuple[Literal['rowchunk'], Tuple[TRow, ...]],
#     Tuple[Literal['error'], TMeta],
#     Tuple[Literal['footer'], TMeta],
# ]
TResultEvent = Tuple[str, Union[TMeta, TRow]]
TResultEvents = AsyncIterable[TResultEvent]


@attr.s(frozen=True, auto_attribs=True)
class BindParamInfo:
    name: str
    type_name: str
    sa_type: sa.sql.type_api.TypeEngine
    value: Any
    expanding: bool = False


@attr.s
class ParamsReformatter:
    query: str = attr.ib()
    existing_params: Set[str] = attr.ib()

    found_params: Set[str] = attr.ib(factory=set)  # for optionally skipping bindparams
    known_params: Set[str] = attr.ib(factory=set)  # for the uniqueness checks
    params_namemap: Dict[str, str] = attr.ib(factory=dict)
    pcounter: int = attr.ib(default=1)  # for making new unique names
    nameprefix: str = attr.ib(default="p")

    def __attrs_post_init__(self) -> None:
        self.known_params.update(self.existing_params)

    def prenormalize_param_name(self, name: str) -> str:
        """To make sure there're no catches with `:param`, normalize param names"""
        name = re.sub(r"[^A-Za-z0-9]", "_", name)
        if not re.search(r"^[A-Za-z]", name):
            name = f"{self.nameprefix}{name}"
        return name

    def normalize_param_name(self, name: str) -> str:
        mapped = self.params_namemap.get(name)
        if mapped is not None:
            return mapped
        mapped = self.prenormalize_param_name(name)
        if mapped == name:
            return name

        if mapped in self.known_params:
            while True:
                candidate = f"{mapped}_{self.pcounter}"
                if candidate not in self.known_params:
                    mapped = candidate
                    break
                self.pcounter += 1

        self.params_namemap[name] = mapped
        self.known_params.add(mapped)
        return mapped

    def param_repl(self, match: re.Match) -> str:
        param_name = match.group(1)
        if param_name not in self.existing_params:
            return match.group(0)  # hacky: leave as-is

        self.found_params.add(param_name)
        mapped_name = self.normalize_param_name(param_name)
        return f":{mapped_name}"

    def reformat(self) -> str:
        """Convert from `{{paramname}}` to `:paramname`"""
        query = self.query
        query = query.replace(":", r"\:")
        query = re.sub(r"\{\{([^}]+)\}\}", self.param_repl, query)
        return query


@attr.s(auto_attribs=True)
class DashSQLSelector:
    conn: ConnectionBase
    sql_query: str
    params: Optional[List[BindParamInfo]]
    db_params: Dict[str, str]
    _service_registry: ServicesRegistry
    connector_specific_params: Optional[Dict[str, TJSONExt]] = attr.ib(default=None)

    def __attrs_post_init__(self) -> None:
        specific_param_keys = get_custom_dash_sql_key_names(conn_type=self.conn.conn_type)
        if specific_param_keys and self.connector_specific_params is None and self.params is not None:
            self.connector_specific_params = {
                param.name: param.value for param in self.params if param.name in specific_param_keys
            }
            self.params = [param for param in self.params if param.name not in specific_param_keys]

    def make_ce(self) -> AsyncConnExecutorBase:
        ce_factory = self._service_registry.get_conn_executor_factory()
        ce = ce_factory.get_async_conn_executor(self.conn)
        ce = ce.mutate_for_dashsql(db_params=self.db_params)
        return ce

    def _make_sa_query(self) -> Tuple[sa.sql.elements.TextClause, str]:
        query = self.sql_query
        params = self.params
        if params is None:
            # No parameters substitution
            return sa_plain_text(query), query

        reformatter = ParamsReformatter(
            query=query,
            existing_params={param.name for param in params},
        )
        query = reformatter.reformat()
        try:
            sa_text = sa.text(query).bindparams(
                *[
                    sa.bindparam(
                        reformatter.normalize_param_name(param.name),
                        type_=param.sa_type,
                        value=param.value,
                        expanding=param.expanding,
                    )
                    for param in params
                    # dubious current behavior: skip unknown params
                    if param.name in reformatter.found_params
                ]
            )
        except sa.exc.ArgumentError:
            raise exc.WrongQueryParameterization()
        conn = self.conn
        # This should've been `dialect = conn.get_dialect() if isinstance(conn, ExecutorBasedMixin) else None`,
        # but mypy couldn't handle that.
        if isinstance(conn, ExecutorBasedMixin):
            dialect = conn.get_dialect()
        else:
            dialect = None  # type: ignore  # TODO: fix
        debug_text = compile_query_for_debug(self.sql_query, dialect=dialect)
        return sa_text, debug_text

    def make_ce_query(self) -> ConnExecutorQuery:
        sa_text, debug_compiled_query = self._make_sa_query()
        return ConnExecutorQuery(
            query=sa_text,  # type: ignore  # TODO: fixC
            # # TODO: try:
            # query=sql_query,
            debug_compiled_query=debug_compiled_query,
            connector_specific_params=self.connector_specific_params,
            autodetect_user_types=True,
            # connect_args=...,  # TODO: pass db_params to e.g. CH query arguments
            # db_name=None,
            # user_types=None,
            # chunk_size=None,
            is_dashsql_query=True,
        )

    @staticmethod
    async def event_gen(
        result_head: TMeta,
        result_chunks: AsyncIterable[Sequence[Any]],
        result_footer_holder: TMeta,
    ) -> TResultEvents:
        # Note: returning strings for easier jsonability.
        # Additionally, those strings are in the API output as-is.
        yield DashSQLEvent.metadata.value, result_head
        try:
            async for chunk in result_chunks:
                yield DashSQLEvent.rowchunk.value, tuple(chunk)
        except Exception as err:
            LOGGER.exception("Runtime error while fetching rows: %r", err)
            raise err
        else:
            yield DashSQLEvent.footer.value, result_footer_holder

    def process_result(self, exec_result: AsyncExecutionResult) -> TResultEvents:
        result_head = exec_result.cursor_info
        assert result_head
        db_types = result_head.get("db_types")
        user_types = exec_result.user_types
        result_head = dict(
            result_head,
            bi_types=[],
            db_types=[],
        )
        if user_types:
            result_head["bi_types"] = [bi_type.name if bi_type else None for bi_type in user_types]
        if db_types:
            result_head["db_types"] = [db_type.name if db_type else None for db_type in db_types]

        result_chunks = exec_result.result
        result_footer_holder: dict = {}  # mutable  # TODO: return footer from CE

        # Wrapping the rest in an additional function to have the code above
        # execute on call rather than on iteration.
        return self.event_gen(result_head, result_chunks, result_footer_holder)

    async def _execute(self, ce: AsyncConnExecutorBase, ce_query: ConnExecutorQuery) -> TResultEvents:
        """Simplified override point"""
        exec_result = await ce.execute(ce_query)
        return self.process_result(exec_result)

    async def execute(self) -> TResultEvents:
        if not self.conn.is_dashsql_allowed:
            raise exc.DashSQLNotAllowed()

        ce = self.make_ce()
        ce_query = self.make_ce_query()
        return await self._execute(ce, ce_query)


@attr.s
class DashSQLCachedSelector(DashSQLSelector):
    _is_bleeding_edge_user: bool = attr.ib(default=False)

    def make_query_id(self) -> str:
        return "dashsql_{}".format(make_id())

    def _get_jsonable_params(self) -> Optional[TJSONExt]:
        if self.params is None:
            return None
        return tuple(
            # `value` might contain e.g. datetimes
            (param.name, param.type_name, param.value, param.expanding)
            for param in self.params
        )

    def _get_jsonable_connector_specific_params(self) -> Optional[TJSONExt]:
        if self.connector_specific_params is None:
            return None
        return tuple((name, value) for name, value in self.connector_specific_params.items())

    def make_cache_options(self) -> BIQueryCacheOptions:
        cache_options_builder = DashSQLCacheOptionsBuilder(is_bleeding_edge_user=self._is_bleeding_edge_user)
        return cache_options_builder.get_cache_options(
            conn=self.conn,
            query_text=self.sql_query,
            params=self._get_jsonable_params(),
            db_params=self.db_params,
            connector_specific_params=self._get_jsonable_connector_specific_params(),
        )

    async def _execute(
        self,
        ce: AsyncConnExecutorBase,
        ce_query: ConnExecutorQuery,
    ) -> TResultEvents:
        async def _request_db() -> TResultEvents:  # basically just `super()`
            exec_result = await ce.execute(ce_query)
            return self.process_result(exec_result)

        async def _generate_func() -> Optional[TJSONExtChunkStream]:
            events = await _request_db()
            chunks = chunkify_by_one(events)  # not `await` for tricky reasons.
            # ^ At this point there's sometimes a chunk that contains one item
            # that is an Event with a chunk of rows in it.
            # TODO?: support type-annotated chunks in the `AsyncChunked`?
            return AsyncChunked(chunked_data=chunks)

        query_id = self.make_query_id()
        assert isinstance(self.conn, ExecutorBasedMixin), "connection must be derived from ExecutorBasedMixin"
        conn = self.conn
        conn_id = conn.uuid
        assert conn_id

        workbook_id = (
            self.conn.entry_key.workbook_id if isinstance(self.conn.entry_key, WorkbookEntryLocation) else None
        )

        service_registry = self._service_registry
        reporting_registry = service_registry.get_reporting_registry()

        reporting_registry.save_reporting_record(
            QueryExecutionStartReportingRecord(
                timestamp=time.time(),
                query_id=query_id,
                dataset_id=None,
                query_type=None,
                connection_type=conn.conn_type,
                conn_reporting_data=self.conn.get_conn_dto().conn_reporting_data(),
                query=str(ce_query.query),
                workbook_id=workbook_id,
            )
        )

        cache_helper = CacheProcessingHelper(
            entity_id=conn_id,
            service_registry=service_registry,
        )
        cache_options = self.make_cache_options()
        if not cache_options.cache_enabled:
            # cache not applicable
            try:
                return await _request_db()
            finally:
                _, exec_exception, _ = sys.exc_info()

                reporting_registry.save_reporting_record(
                    QueryExecutionEndReportingRecord(
                        timestamp=time.time(),
                        query_id=query_id,
                        exception=exec_exception,
                    )
                )

        # TODO: make this env-configurable through settings.
        use_locked_cache = self.conn.use_locked_cache

        cache_full_hit = None
        try:
            sit, result_stream = await cache_helper.run_with_cache(
                generate_func=_generate_func,
                cache_options=cache_options,
                use_locked_cache=use_locked_cache,
            )
            if sit == CacheSituation.full_hit:
                cache_full_hit = True
            elif sit == CacheSituation.generated:
                cache_full_hit = False
        finally:
            _, exec_exception, _ = sys.exc_info()

            reporting_registry.save_reporting_record(
                QueryExecutionCacheInfoReportingRecord(
                    query_id=query_id,
                    cache_full_hit=cache_full_hit,
                    timestamp=time.time(),
                )
            )

            reporting_registry.save_reporting_record(
                QueryExecutionEndReportingRecord(
                    timestamp=time.time(),
                    query_id=query_id,
                    exception=exec_exception,
                )
            )

        assert result_stream is not None
        return result_stream.items  # type: ignore  # TODO: fix

lib/dl_core/dl_core/data_processing/dashsql.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import contextmanager
from typing import (
    Generator,
    Optional,
)

import attr
import sqlalchemy as sa
from sqlalchemy.sql.selectable import FromClause

from dl_constants.enums import DataSourceRole
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.ids import AvatarId
from dl_core.constants import DataAPILimits
from dl_core.data_processing.prepared_components.manager_base import PreparedComponentManagerBase
from dl_core.data_processing.prepared_components.primitives import PreparedSingleFromInfo
from dl_core.data_source.collection import DataSourceCollectionFactory
import dl_core.data_source.sql
import dl_core.exc as exc
from dl_core.multisource import SourceAvatar
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_dataset import Dataset
from dl_core.us_manager.local_cache import USEntryBuffer


@attr.s
class DefaultPreparedComponentManager(PreparedComponentManagerBase):
    """
    Manages basic prepared components - ones formed from static dataset avatars and relations.
    """

    _dataset: Dataset = attr.ib(kw_only=True)
    _role: DataSourceRole = attr.ib(kw_only=True)
    _us_entry_buffer: USEntryBuffer = attr.ib(kw_only=True)
    _ds_accessor: DatasetComponentAccessor = attr.ib(init=False)
    _dsrc_coll_factory: DataSourceCollectionFactory = attr.ib(init=False)

    @_ds_accessor.default
    def _make_ds_accessor(self) -> DatasetComponentAccessor:
        return DatasetComponentAccessor(dataset=self._dataset)

    @_dsrc_coll_factory.default
    def _make_dsrc_coll_factory(self) -> DataSourceCollectionFactory:
        return DataSourceCollectionFactory(us_entry_buffer=self._us_entry_buffer)

    @contextmanager
    def _handle_incomplete_source(self, avatar: SourceAvatar) -> Generator[None, None, None]:
        try:
            yield
        except exc.TableNameNotConfiguredError as err:
            if self._role == DataSourceRole.materialization:
                raise exc.MaterializationNotFinished(
                    message="Materialization is not yet finished",
                    query="<get_from_clause>",
                    db_message="",
                    details={"avatar_id": avatar.id, "avatar_title": avatar.title},
                ) from err
            else:
                raise

    def get_prepared_source(
        self, avatar_id: AvatarId, alias: str, from_subquery: bool, subquery_limit: Optional[int]
    ) -> PreparedSingleFromInfo:
        avatar = self._ds_accessor.get_avatar_strict(avatar_id=avatar_id)
        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=avatar.source_id)
        dsrc_coll = self._dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
        dsrc = dsrc_coll.get_strict(role=self._role)
        if not isinstance(dsrc, dl_core.data_source.sql.BaseSQLDataSource):
            raise TypeError(f"Root data source has non-SQL type: {type(dsrc)}")

        def get_columns():  # type: ignore  # TODO: fix
            with self._handle_incomplete_source(avatar=avatar):  # type: ignore  # TODO: fix
                columns = dsrc.saved_raw_schema  # type: ignore  # TODO: fix
                if columns is None:
                    raise exc.TableNameNotConfiguredError()
                return columns

        columns = get_columns()
        col_names = [col.name for col in columns]
        from_subquery = from_subquery and dsrc.supports_preview_from_subquery
        query_compiler = dsrc.get_query_compiler()

        sql_source: FromClause
        if from_subquery:
            # Subquery mode: wrap the leftmost (root) source/table into a "SELECT ... FROM ..." subquery
            # to limit the number of entries before the GROUP BY clause is executed
            fields = [sa.literal_column(query_compiler.quote(col_name)) for col_name in col_names] or ["*"]
            sql_source = (
                sa.select(fields)
                .select_from(dsrc.get_sql_source())  # type: ignore  # TODO: fix
                .limit(subquery_limit or DataAPILimits.DEFAULT_SUBQUERY_LIMIT)
                .alias(alias)
            )
        else:
            sql_source = dsrc.get_sql_source(alias=alias)  # type: ignore  # TODO: fix

        target_connection_ref = dsrc.connection_ref
        target_connection = self._us_entry_buffer.get_entry(target_connection_ref)
        assert isinstance(target_connection, ConnectionBase)

        prep_src_info = PreparedSingleFromInfo(
            id=avatar_id,
            alias=alias,
            data_source_list=(dsrc,),
            col_names=col_names,
            user_types=[col.user_type for col in columns],
            sql_source=sql_source,
            query_compiler=query_compiler,
            supported_join_types=dsrc.supported_join_types,
            db_name=(dsrc.db_name if isinstance(dsrc, dl_core.data_source.sql.DbSQLDataSourceMixin) else None),
            connect_args=dsrc.get_connect_args(),
            pass_db_query_to_user=target_connection.pass_db_query_to_user,
            target_connection_ref=target_connection_ref,
        )
        return prep_src_info

lib/dl_core/dl_core/data_processing/prepared_components/default_manager.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import Optional

from dl_core.components.ids import AvatarId
from dl_core.data_processing.prepared_components.primitives import PreparedSingleFromInfo


class PreparedComponentManagerBase(abc.ABC):
    @abc.abstractmethod
    def get_prepared_source(
        self, avatar_id: AvatarId, alias: str, from_subquery: bool, subquery_limit: Optional[int]
    ) -> PreparedSingleFromInfo:
        raise NotImplementedError

lib/dl_core/dl_core/data_processing/prepared_components/manager_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Collection,
    Optional,
    Sequence,
)

import attr

from dl_constants.enums import (
    JoinType,
    UserDataType,
)
from dl_core.connectors.base.query_compiler import QueryCompiler
from dl_core.query.bi_query import SqlSourceType


if TYPE_CHECKING:
    from dl_core.base_models import ConnectionRef
    import dl_core.data_source


@attr.s(frozen=True)
class PreparedFromInfo:
    """
    Represents an entity that can be used in a ``FROM`` clause of a ``SELECT`` statement.
    In the simplest case - an SQL data source (database table),
    but can also itself be a ``SELECT`` statement.
    Can reflect a static ``SourceAvatar`` as well as a dynamic combination of avatars
    created during the formation of complex ``SELECT`` statements.
    """

    sql_source: Optional[SqlSourceType] = attr.ib(kw_only=True)
    query_compiler: QueryCompiler = attr.ib(kw_only=True)
    supported_join_types: Collection[JoinType] = attr.ib(kw_only=True)
    data_source_list: Optional[tuple[dl_core.data_source.BaseSQLDataSource, ...]] = attr.ib(kw_only=True)
    db_name: Optional[str] = attr.ib(kw_only=True)
    connect_args: dict[str, Any] = attr.ib(kw_only=True)
    pass_db_query_to_user: bool = attr.ib(kw_only=True)
    target_connection_ref: Optional[ConnectionRef] = attr.ib(kw_only=True)

    def supports_join_type(self, join_type: JoinType) -> bool:
        return join_type in self.supported_join_types

    @property
    def non_null_sql_source(self) -> SqlSourceType:
        assert self.sql_source is not None
        return self.sql_source


@attr.s(frozen=True)
class PreparedMultiFromInfo(PreparedFromInfo):
    pass


@attr.s(frozen=True)
class PreparedSingleFromInfo(PreparedFromInfo):
    """
    Represents an entity that can be used in a ``FROM`` clause of a ``SELECT`` statement.
    In the simplest case - an SQL data source (database table),
    but can also itself be a ``SELECT`` statement.
    Can reflect a static ``SourceAvatar`` as well as a dynamic combination of avatars
    created during the formation of complex ``SELECT`` statements.
    """

    id: str = attr.ib(kw_only=True)
    alias: str = attr.ib(kw_only=True)
    col_names: Sequence[str] = attr.ib(kw_only=True)
    user_types: Sequence[UserDataType] = attr.ib(kw_only=True)

lib/dl_core/dl_core/data_processing/prepared_components/primitives.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Collection,
    Dict,
    Optional,
    Set,
    Tuple,
)

from dl_core.data_processing.processing.operation import (
    BaseOp,
    MultiSourceOp,
    SingleSourceOp,
)
from dl_core.data_processing.stream_base import (
    AbstractStream,
    DataStreamAsync,
)


class OpExecutionContext:
    _processing_id: str
    _streams: Dict[str, AbstractStream]
    _operations: Tuple[BaseOp, ...]

    def __init__(self, processing_id: str, streams: Collection[AbstractStream], operations: Collection[BaseOp]):
        self._processing_id = processing_id
        self._streams = {stream.id: stream for stream in streams}
        self._operations = tuple(op for op in operations)
        self._operations_by_output: Dict[str, BaseOp] = {}

        for op in operations:
            self._operations_by_output[op.dest_stream_id] = op

    @property
    def operations(self) -> Tuple[BaseOp, ...]:
        return self._operations

    @property
    def streams(self) -> Collection[AbstractStream]:
        return [stream for stream in self._streams.values()]

    @property
    def data_streams(self) -> Collection[DataStreamAsync]:
        return [stream for stream in self._streams.values() if isinstance(stream, DataStreamAsync)]

    @property
    def processing_id(self) -> str:
        return self._processing_id

    def get_stream(self, stream_id: str) -> Optional[AbstractStream]:
        return self._streams.get(stream_id)

    def add_stream(self, stream: AbstractStream) -> None:
        self._streams[stream.id] = stream

    def remove_stream(self, stream: AbstractStream) -> None:
        del self._streams[stream.id]

    def get_stream_ids_for_operation(self, op: BaseOp) -> Set[str]:
        result = set()
        if isinstance(op, SingleSourceOp):
            result.add(op.source_stream_id)
        elif isinstance(op, MultiSourceOp):
            for src_id in op.source_stream_ids:
                result.add(src_id)

        return result

    def get_generating_operation(self, stream_id: str) -> BaseOp:
        """Return the operation that generates the data stream with the given ID"""
        return self._operations_by_output[stream_id]

    def clone(self, **kwargs: Any) -> OpExecutionContext:
        new_kwargs = dict(
            processing_id=self._processing_id,
            streams=self._streams.values(),
            operations=self._operations,
        )

        for key, val in kwargs.items():
            if key not in new_kwargs:
                raise ValueError(f"Unknown kwarg: {key!r}")
            new_kwargs[key] = val

        return self.__class__(**new_kwargs)  # type: ignore

lib/dl_core/dl_core/data_processing/processing/context.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Optional,
    Sequence,
    Union,
)

import attr
import sqlalchemy as sa
from sqlalchemy.sql.selectable import Select

from dl_constants.types import TBIDataValue
from dl_core.data_processing.types import TValuesChunkStream
from dl_core.utils import make_id


if TYPE_CHECKING:
    from dl_constants.enums import UserDataType
    from dl_core.data_processing.cache.primitives import LocalKeyRepresentation
    from dl_core.data_processing.prepared_components.primitives import PreparedMultiFromInfo


LOGGER = logging.getLogger(__name__)
DEFAULT_CHUNK_SIZE = 1000


@attr.s
class ProcessorDbExecAdapterBase(abc.ABC):
    _default_chunk_size: ClassVar[int] = 1000
    _log: ClassVar[logging.Logger] = LOGGER.getChild("ProcessorDbExecAdapterBase")

    @abc.abstractmethod
    async def _execute_and_fetch(
        self,
        *,
        query: Union[Select, str],
        user_types: Sequence[UserDataType],
        chunk_size: int,
        joint_dsrc_info: Optional[PreparedMultiFromInfo] = None,
        query_id: str,
    ) -> TValuesChunkStream:
        """
        Execute SELECT statement.
        Return data as ``AsyncChunked``.
        """
        raise NotImplementedError

    async def scalar(self, query: Union[str, Select], user_type: UserDataType) -> TBIDataValue:
        """Execute a statement returning a scalar value."""
        data_stream = await self._execute_and_fetch(
            query_id=make_id(),
            query=query,
            user_types=[user_type],
            chunk_size=5,
        )
        data = await data_stream.all()
        assert len(data) == 1, f"Expected 1 entry, got {len(data)}"
        assert len(data[0]) == 1, f"Expected 1 column, got {len(data[0])}"
        return data[0][0]

    async def fetch_data_from_select(
        self,
        *,
        query: Union[str, sa.sql.selectable.Select],
        user_types: Sequence[UserDataType],
        chunk_size: Optional[int] = None,
        joint_dsrc_info: Optional[PreparedMultiFromInfo] = None,
        query_id: str,
    ) -> TValuesChunkStream:
        """Fetch data from a table"""

        chunk_size = chunk_size or self._default_chunk_size
        self._log.info(f"Fetching data from query {query_id}")

        return await self._execute_and_fetch(
            query=query,
            user_types=user_types,
            chunk_size=chunk_size,
            joint_dsrc_info=joint_dsrc_info,
            query_id=query_id,
        )

    def get_data_key(
        self,
        *,
        query_id: str,
        query: Union[str, Select],
        user_types: Sequence[UserDataType],
        joint_dsrc_info: Optional[PreparedMultiFromInfo] = None,
    ) -> Optional[LocalKeyRepresentation]:
        return None

lib/dl_core/dl_core/data_processing/processing/db_base/exec_adapter_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import wraps
import logging
from typing import (
    Awaitable,
    Callable,
    List,
)

import attr
import sqlalchemy as sa
from sqlalchemy.sql.selectable import (
    Alias,
    Select,
    Subquery,
)

from dl_core.data_processing.prepared_components.primitives import (
    PreparedFromInfo,
    PreparedMultiFromInfo,
    PreparedSingleFromInfo,
)
from dl_core.data_processing.processing.context import OpExecutionContext
from dl_core.data_processing.processing.db_base.exec_adapter_base import ProcessorDbExecAdapterBase
from dl_core.data_processing.processing.operation import (
    BaseOp,
    CalcOp,
    DownloadOp,
    JoinOp,
    MultiSourceOp,
    SingleSourceOp,
)
from dl_core.data_processing.source_builder import SqlSourceBuilder
from dl_core.data_processing.stream_base import (
    AbstractStream,
    DataRequestMetaInfo,
    DataSourceVS,
    DataStreamAsync,
    JointDataSourceVS,
)
from dl_core.utils import (
    compile_query_for_debug,
    make_id,
)


LOGGER = logging.getLogger(__name__)


@attr.s
class OpExecutorAsync:
    ctx: OpExecutionContext = attr.ib()
    db_ex_adapter: ProcessorDbExecAdapterBase = attr.ib()

    async def execute(self, op: BaseOp) -> AbstractStream:
        raise NotImplementedError


def log_op(
    func: Callable[[OpExecutorAsync, BaseOp], Awaitable[AbstractStream]]
) -> Callable[[OpExecutorAsync, BaseOp], Awaitable[AbstractStream]]:
    """Log operation calls"""

    @wraps(func)
    async def wrapper(self: OpExecutorAsync, op: BaseOp) -> AbstractStream:
        assert isinstance(op, (SingleSourceOp, MultiSourceOp))
        stream_ids = [op.source_stream_id] if isinstance(op, SingleSourceOp) else op.source_stream_ids
        LOGGER.info(
            f"Running {type(op).__name__} "
            f"with source streams {stream_ids} "
            f"and destination stream {op.dest_stream_id}"
        )
        return await func(self, op)

    return wrapper


class DownloadOpExecutorAsync(OpExecutorAsync):
    """Loads data from a database table"""

    @log_op  # type: ignore  # TODO: fix
    async def execute(self, op: BaseOp) -> DataStreamAsync:  # type: ignore  # TODO: fix
        assert isinstance(op, DownloadOp)

        source_stream = self.ctx.get_stream(op.source_stream_id)
        assert isinstance(source_stream, DataSourceVS)
        query_compiler = source_stream.prep_src_info.query_compiler

        query = source_stream.prep_src_info.sql_source
        if isinstance(query, Alias):
            # Unwrap Alias:
            # 1. There is no reason for alias when executing directly.
            # 2. Serialization/deserialization of Alias is not supported for RQE
            query = query.original

        if isinstance(query, Subquery):
            # Unwrap Subquery for sake of sa 1.3 compatibility with 1.4.
            # TODO: Confirm that it is ok or adapt Subquery feature for our code.
            query = query.original

        assert isinstance(query, Select), f"Got type {type(query).__name__} for query, expected Select"

        query_debug_str = compile_query_for_debug(query=query, dialect=query_compiler.dialect)
        LOGGER.info(f"Going to database with SQL query:\n{query_debug_str}")

        joint_dsrc_info = PreparedMultiFromInfo(
            sql_source=query,
            data_source_list=source_stream.prep_src_info.data_source_list,  # type: ignore  # TODO: fix
            db_name=source_stream.prep_src_info.db_name,
            connect_args=source_stream.prep_src_info.connect_args,
            query_compiler=source_stream.prep_src_info.query_compiler,
            supported_join_types=source_stream.prep_src_info.supported_join_types,
            pass_db_query_to_user=source_stream.prep_src_info.pass_db_query_to_user,
            target_connection_ref=source_stream.prep_src_info.target_connection_ref,
        )  # FIXME: replace usage with prep_src_info

        query_id = make_id()
        data = await self.db_ex_adapter.fetch_data_from_select(
            query=query,
            user_types=source_stream.user_types,
            joint_dsrc_info=joint_dsrc_info,
            query_id=query_id,
        )

        if op.row_count_hard_limit is not None:
            data = data.limit(op.row_count_hard_limit)

        data_key = self.db_ex_adapter.get_data_key(
            query=query, user_types=source_stream.user_types, joint_dsrc_info=joint_dsrc_info, query_id=query_id
        )

        pass_db_query_to_user = joint_dsrc_info.pass_db_query_to_user
        data_source_list = source_stream.prep_src_info.data_source_list or ()
        assert data_source_list is not None

        return DataStreamAsync(
            id=op.dest_stream_id,
            names=source_stream.names,
            user_types=source_stream.user_types,
            data=data,
            meta=DataRequestMetaInfo(
                query=query_debug_str,
                data_source_list=data_source_list,
                query_id=query_id,
                pass_db_query_to_user=pass_db_query_to_user,
            ),
            data_key=data_key,
        )


class CalcOpExecutorAsync(OpExecutorAsync):
    """
    Applies query to stream (lazy operation).
    The real calculation is done when the data is fetched.
    """

    @log_op  # type: ignore  # TODO: fix
    async def execute(self, op: BaseOp) -> DataSourceVS:  # type: ignore  # TODO: fix
        assert isinstance(op, CalcOp)

        source_stream = self.ctx.get_stream(op.source_stream_id)

        from_info: PreparedFromInfo
        if isinstance(source_stream, DataSourceVS):
            from_info = source_stream.prep_src_info
        elif isinstance(source_stream, JointDataSourceVS):
            from_info = source_stream.joint_dsrc_info
        else:
            raise TypeError(f"Type {type(source_stream).__name__} is not supported as a source for CalcOp")

        query_compiler = from_info.query_compiler
        query = query_compiler.compile_select(
            bi_query=op.bi_query,
            sql_source=from_info.sql_source,
        )
        query_debug_str = compile_query_for_debug(query=query, dialect=query_compiler.dialect)

        LOGGER.info(f"Generated lazy query: {query_debug_str}")

        names = []
        user_types = []
        for expr_ctx in op.bi_query.select_expressions:
            assert expr_ctx.alias is not None
            assert expr_ctx.user_type is not None
            names.append(expr_ctx.alias)
            user_types.append(expr_ctx.user_type)

        alias = op.alias
        prep_src_info = PreparedSingleFromInfo(
            id=op.result_id,
            alias=alias,
            col_names=names,
            user_types=user_types,
            sql_source=sa.alias(query, alias),
            query_compiler=query_compiler,
            data_source_list=from_info.data_source_list,
            supported_join_types=from_info.supported_join_types,
            db_name=from_info.db_name,
            connect_args=from_info.connect_args,
            pass_db_query_to_user=from_info.pass_db_query_to_user,
            target_connection_ref=from_info.target_connection_ref,
        )
        return DataSourceVS(
            id=op.dest_stream_id,
            alias=op.alias,
            result_id=op.result_id,
            prep_src_info=prep_src_info,
            names=prep_src_info.col_names,
            user_types=user_types,
        )


class JoinOpExecutorAsync(OpExecutorAsync):
    @log_op  # type: ignore  # TODO: fix
    async def execute(self, op: BaseOp) -> JointDataSourceVS:  # type: ignore  # TODO: fix
        assert isinstance(op, JoinOp)

        prepared_sources: List[PreparedSingleFromInfo] = []
        for stream_id in op.source_stream_ids:
            stream = self.ctx.get_stream(stream_id=stream_id)
            assert isinstance(stream, DataSourceVS)
            prepared_sources.append(stream.prep_src_info)

        source_builder = SqlSourceBuilder()

        joint_dsrc_info = source_builder.build_source(
            root_avatar_id=op.root_avatar_id,
            prepared_sources=prepared_sources,
            join_on_expressions=op.join_on_expressions,
            use_empty_source=op.use_empty_source,
        )
        return JointDataSourceVS(
            id=op.dest_stream_id, names=[], user_types=[], joint_dsrc_info=joint_dsrc_info  # not used  # not used
        )

lib/dl_core/dl_core/data_processing/processing/db_base/op_executors.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    AbstractSet,
    ClassVar,
    Collection,
    Optional,
    Tuple,
    Type,
)

import attr

from dl_core.components.ids import AvatarId
from dl_core.data_processing.stream_base import (
    AbstractStream,
    DataSourceVS,
    DataStreamAsync,
    JointDataSourceVS,
)
from dl_core.query.bi_query import BIQuery
from dl_core.query.expression import JoinOnExpressionCtx


# Bases


@attr.s(frozen=True)
class BaseOp:
    """Base class for all operations"""

    supported_source_types: ClassVar[Tuple[Type[AbstractStream], ...]] = ()
    supported_dest_types: ClassVar[Tuple[Type[AbstractStream], ...]] = ()

    dest_stream_id: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class SingleSourceOp(BaseOp):
    source_stream_id: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class MultiSourceOp(BaseOp):
    source_stream_ids: AbstractSet[str] = attr.ib(kw_only=True)


# Actual ops


@attr.s(frozen=True)
class UploadOp(SingleSourceOp):
    """Upload data from a stream to a virtual one"""

    supported_source_types = (DataStreamAsync,)
    supported_dest_types = (DataSourceVS,)

    result_id: AvatarId = attr.ib(kw_only=True)
    alias: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class DownloadOp(SingleSourceOp):
    """Download data from a table (virtual stream) to a real data stream"""

    supported_source_types = (DataSourceVS,)
    supported_dest_types = (DataStreamAsync,)

    row_count_hard_limit: Optional[int] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class CalcOp(SingleSourceOp):
    """Apply query to the source"""

    supported_source_types = (DataSourceVS, JointDataSourceVS)
    supported_dest_types = (DataSourceVS,)

    result_id: AvatarId = attr.ib(kw_only=True)
    bi_query: BIQuery = attr.ib(kw_only=True)
    alias: str = attr.ib(kw_only=True)


# To be implemented:


@attr.s(frozen=True)
class JoinOp(MultiSourceOp):  # TODO: Implement me
    """Join several streams"""

    supported_source_types = (DataSourceVS,)
    supported_dest_types = (JointDataSourceVS,)

    join_on_expressions: Collection[JoinOnExpressionCtx] = attr.ib(kw_only=True)
    root_avatar_id: AvatarId = attr.ib(kw_only=True)
    # Workaround for sourceless subqueries in LODs
    use_empty_source: bool = attr.ib(kw_only=True, default=False)


@attr.s(frozen=True)
class UnionOp(MultiSourceOp):  # TODO: Implement me
    """Concatenate several streams"""

    supported_source_types = (DataSourceVS,)
    supported_dest_types = (DataSourceVS,)


@attr.s(frozen=True)
class MapOp(SingleSourceOp):  # TODO: Implement me
    supported_source_types = (DataStreamAsync,)
    supported_dest_types = (DataStreamAsync,)

    mapper_id: str = attr.ib(kw_only=True)

lib/dl_core/dl_core/data_processing/processing/operation.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
import time
from typing import (
    TYPE_CHECKING,
    Collection,
    Dict,
    List,
    Optional,
    TypeVar,
)

import attr
import shortuuid

from dl_api_commons.reporting.models import (
    DataProcessingEndReportingRecord,
    DataProcessingStartReportingRecord,
)
from dl_core.data_processing.processing.context import OpExecutionContext
from dl_core.data_processing.processing.operation import (
    BaseOp,
    MultiSourceOp,
    SingleSourceOp,
)
from dl_core.data_processing.stream_base import (
    AbstractStream,
    DataStreamAsync,
)


if TYPE_CHECKING:
    from dl_api_commons.reporting.registry import ReportingRegistry  # noqa
    from dl_core.services_registry import ServicesRegistry  # noqa


LOGGER = logging.getLogger(__name__)


_OP_PROC_TV = TypeVar("_OP_PROC_TV", bound="OperationProcessorAsyncBase")


class OperationProcessorAsyncBase(abc.ABC):
    @abc.abstractmethod
    async def ping(self) -> Optional[int]:
        """Check processor readiness"""

    async def execute_operation(self, op: BaseOp, ctx: OpExecutionContext) -> AbstractStream:
        raise NotImplementedError

    async def start(self) -> None:
        pass

    async def end(self) -> None:
        pass

    async def __aenter__(self: _OP_PROC_TV) -> _OP_PROC_TV:
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):  # type: ignore
        await self.end()

    async def prepare_output_streams(self, output_streams: List[AbstractStream]) -> List[DataStreamAsync]:
        result: List[DataStreamAsync] = []
        for stream in output_streams:
            assert isinstance(
                stream, DataStreamAsync
            ), f"Expected DataStreamAsync, got {type(stream).__name__} ({stream.id}) as out stream"
            result.append(stream)

        return result

    def pre_run(self, ctx: OpExecutionContext) -> None:
        pass

    def post_run(self, ctx: OpExecutionContext, exec_exception: Optional[Exception]) -> None:
        pass

    def _validate_input_stream(self, stream: AbstractStream, op: BaseOp) -> None:
        if not isinstance(stream, op.supported_source_types):
            raise TypeError(
                f"Requested stream {stream.id} has type {type(stream)} "
                f"that is not supported for input by operation {type(op)}"
            )

    def _validate_output_stream(self, stream: AbstractStream, op: BaseOp) -> None:
        if not isinstance(stream, op.supported_dest_types):
            raise TypeError(
                f"Returned stream {stream.id} has type {type(stream)} "
                f"that is not supported for output by operation {type(op)}"
            )

    async def execute_operations(
        self,
        ctx: OpExecutionContext,
        output_stream_ids: Collection[str],
    ) -> List[DataStreamAsync]:
        ready_streams: Dict[str, AbstractStream] = {stream.id: stream for stream in ctx.streams}
        assert all(ready_streams.values())

        async def _make_output_streams_ready_recursive(req_out_stream_ids: Collection[str]) -> None:
            for out_stream_id in sorted(req_out_stream_ids):
                op = ctx.get_generating_operation(stream_id=out_stream_id)

                if isinstance(op, MultiSourceOp):
                    source_stream_ids = op.source_stream_ids
                else:
                    assert isinstance(op, SingleSourceOp)
                    source_stream_ids = {op.source_stream_id}

                if not source_stream_ids.issubset(set(ready_streams)):  # type: ignore  # TODO: fix
                    # Make these streams ready
                    await _make_output_streams_ready_recursive(source_stream_ids)

                # All streams are ready for this operation
                for source_stream_id in source_stream_ids:
                    source_stream = ctx.get_stream(stream_id=source_stream_id)
                    self._validate_input_stream(source_stream, op)  # type: ignore  # TODO: fix
                    if isinstance(source_stream, DataStreamAsync):  # These streams are not reusable
                        del ready_streams[source_stream_id]

                new_stream = await self.execute_operation(op=op, ctx=ctx)
                self._validate_output_stream(new_stream, op)
                ctx.add_stream(new_stream)
                ready_streams[new_stream.id] = new_stream

        await _make_output_streams_ready_recursive(output_stream_ids)
        LOGGER.info("Done processing operations")
        return [ready_streams[stream_id] for stream_id in sorted(output_stream_ids)]  # type: ignore

    async def run(
        self,
        operations: Collection[BaseOp],
        streams: Collection[AbstractStream],
        output_stream_ids: Collection[str],
    ) -> List[DataStreamAsync]:
        processing_id = shortuuid.uuid()
        LOGGER.info(f"Initializing processing context {processing_id}")
        ctx = OpExecutionContext(
            processing_id=processing_id,
            streams=streams,
            operations=operations,
        )

        self.pre_run(ctx=ctx)

        LOGGER.info(f'Processor got operations: {", ".join([type(op).__name__ for op in operations])}')

        exec_exception: Optional[Exception] = None
        try:
            # Execute and get output streams
            stream_list = await self.execute_operations(ctx=ctx, output_stream_ids=output_stream_ids)
            actual_output_stream_ids = {st.id for st in stream_list}
            assert actual_output_stream_ids == set(
                output_stream_ids
            ), f"Destination streams do not match: expected {output_stream_ids}, got {actual_output_stream_ids}"
            # Finalize streams
            result = await self.prepare_output_streams(output_streams=stream_list)  # type: ignore  # TODO: fix

        except Exception as fired_exception:
            exec_exception = fired_exception
            raise
        finally:
            self.post_run(ctx=ctx, exec_exception=exec_exception)

        return result


@attr.s
class SROperationProcessorAsyncBase(OperationProcessorAsyncBase):
    _service_registry: "ServicesRegistry" = attr.ib(default=None, kw_only=True)  # Service registry override

    @property
    def service_registry(self) -> "ServicesRegistry":
        if self._service_registry is not None:
            return self._service_registry

    @property
    def _reporting_registry(self) -> "ReportingRegistry":
        return self.service_registry.get_reporting_registry()

    def _save_start_exec_reporting_record(self, ctx: OpExecutionContext) -> None:
        report = DataProcessingStartReportingRecord(
            timestamp=time.time(),
            processing_id=ctx.processing_id,
            source_query_ids=tuple(
                in_stream.meta.query_id for in_stream in ctx.data_streams if in_stream.meta.query_id is not None
            ),  # type: ignore
        )
        self._reporting_registry.save_reporting_record(report=report)

    def _save_end_exec_reporting_record(
        self,
        ctx: OpExecutionContext,
        exec_exception: Optional[Exception],
    ) -> None:
        report = DataProcessingEndReportingRecord(
            timestamp=time.time(),
            processing_id=ctx.processing_id,
            exception=exec_exception,
        )
        self._reporting_registry.save_reporting_record(report=report)

    def pre_run(self, ctx: OpExecutionContext) -> None:
        self._save_start_exec_reporting_record(ctx=ctx)

    def post_run(self, ctx: OpExecutionContext, exec_exception: Optional[Exception]) -> None:
        self._save_end_exec_reporting_record(ctx=ctx, exec_exception=exec_exception)

lib/dl_core/dl_core/data_processing/processing/processor.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import time
from typing import (
    TYPE_CHECKING,
    Collection,
    List,
    Optional,
)

import attr
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql.base import PGDialect

from dl_api_commons.reporting.models import DataProcessingCacheInfoReportingRecord
from dl_core.connectors.base.query_compiler import QueryCompiler
from dl_core.data_processing.cache.primitives import LocalKeyRepresentation
from dl_core.data_processing.cache.processing_helper import (
    CacheProcessingHelper,
    CacheSituation,
)
from dl_core.data_processing.cache.utils import CacheOptionsBuilderDataProcessor
from dl_core.data_processing.processing.operation import (
    BaseOp,
    CalcOp,
    MultiSourceOp,
    SingleSourceOp,
)
from dl_core.data_processing.processing.processor import SROperationProcessorAsyncBase
from dl_core.data_processing.stream_base import (
    CacheVirtualStream,
    DataRequestMetaInfo,
    DataStreamAsync,
    DataStreamBase,
)


if TYPE_CHECKING:
    from dl_core.data_processing.processing.context import OpExecutionContext
    from dl_core.data_processing.processing.processor import OperationProcessorAsyncBase
    from dl_core.data_processing.stream_base import AbstractStream
    from dl_core.data_processing.types import TValuesChunkStream
    from dl_core.us_dataset import Dataset


LOGGER = logging.getLogger(__name__)


@attr.s
class SingleStreamCacheOperationProcessor(SROperationProcessorAsyncBase):
    _dataset: Dataset = attr.ib(kw_only=True)
    _query_compiler: QueryCompiler = attr.ib(kw_only=True)
    _cache_options_builder: CacheOptionsBuilderDataProcessor = attr.ib(kw_only=True)
    _main_processor: OperationProcessorAsyncBase = attr.ib(kw_only=True)

    _cache_helper: Optional[CacheProcessingHelper] = attr.ib(init=False)

    async def ping(self) -> Optional[int]:
        return await self._main_processor.ping()  # TODO: also ping cache engine?

    @_query_compiler.default  # noqa
    def _default_query_compiler(self) -> QueryCompiler:
        return QueryCompiler(dialect=PGDialect())

    @_cache_options_builder.default  # noqa
    def _cache_options_builder_default(self) -> CacheOptionsBuilderDataProcessor:
        return CacheOptionsBuilderDataProcessor(default_ttl_config=self._service_registry.default_cache_ttl_config)

    def __attrs_post_init__(self):  # type: ignore  # TODO: fix
        # cache_engine_factory = self.service_registry.get_cache_engine_factory()
        ds_id = self._dataset.uuid
        self._cache_helper = CacheProcessingHelper(
            entity_id=ds_id,  # type: ignore  # TODO: fix
            service_registry=self.service_registry,
        )

    # TODO FIX: @altvod Move data key building to CacheOptionsBuilderDataProcessor
    def _make_op_data_key(self, op: BaseOp, ctx: OpExecutionContext) -> Optional[LocalKeyRepresentation]:
        """Generate new data key for operation"""

        base_key: Optional[LocalKeyRepresentation]
        if isinstance(op, SingleSourceOp):
            # inherit key from single source stream
            stream = ctx.get_stream(op.source_stream_id)
            assert isinstance(stream, DataStreamBase)
            base_key = stream.data_key
        elif isinstance(op, MultiSourceOp):
            # combine key from multiple source streams
            base_key = LocalKeyRepresentation()
            for stream_id in op.source_stream_ids:
                stream = ctx.get_stream(stream_id)
                assert isinstance(stream, DataStreamBase)
                if stream.data_key is None:
                    base_key = None
                    break
                base_key = base_key.multi_extend(*stream.data_key.key_parts)
        else:
            raise TypeError(f"Invalid operation type {type(op)}")

        if base_key is None:
            return None

        data_key = base_key.extend(part_type="operation", part_content=op.__class__.__name__)  # type: ignore  # TODO: fix
        if isinstance(op, CalcOp):
            data_key = data_key.extend(
                part_type="query",
                part_content=CacheOptionsBuilderDataProcessor.get_query_str_for_cache(
                    query=self._query_compiler.compile_select(
                        bi_query=op.bi_query,
                        # The real table name doesn't really matter here
                        sql_source=sa.table("table"),
                    ),
                    dialect=self._query_compiler.dialect,
                ),
            )

        return data_key

    async def execute_operation(self, op: BaseOp, ctx: OpExecutionContext) -> CacheVirtualStream:
        """
        Imitate execution by constructing a new virtual stream
        with an updated data key for the cache engine
        """

        if isinstance(op, SingleSourceOp):
            source_stream = ctx.get_stream(op.source_stream_id)
            assert isinstance(source_stream, DataStreamBase)
            user_types = source_stream.user_types
            names = source_stream.names
            stream_meta = source_stream.meta
        elif isinstance(op, MultiSourceOp):
            source_streams = [ctx.get_stream(source_stream_id) for source_stream_id in sorted(op.source_stream_ids)]
            names = []  # not used
            user_types = []  # not used
            stream_meta = DataRequestMetaInfo(
                query_id=None,  # FIXME
                query=None,
                is_materialized=all(
                    stream.meta.is_materialized for stream in source_streams if isinstance(stream, DataStreamBase)
                ),
                data_source_list=[
                    dsrc for s in source_streams if isinstance(s, DataStreamBase) for dsrc in s.meta.data_source_list
                ],
            )
        else:
            raise TypeError(f"Invalid operation type {type(op)}")

        data_key = self._make_op_data_key(op=op, ctx=ctx)

        return CacheVirtualStream(
            id=op.dest_stream_id,
            names=names,
            user_types=user_types,
            data_key=data_key,
            meta=stream_meta,
        )

    def _save_cache_info_reporting_record(self, ctx: OpExecutionContext, cache_full_hit: bool) -> None:
        report = DataProcessingCacheInfoReportingRecord(
            timestamp=time.time(),
            processing_id=ctx.processing_id,
            cache_full_hit=cache_full_hit,
        )
        self._reporting_registry.save_reporting_record(report=report)

    async def execute_operations(
        self,
        ctx: OpExecutionContext,
        output_stream_ids: Collection[str],
    ) -> List[DataStreamAsync]:
        """
        Turn all virtual output streams into real data-containing ones
        by getting data from the cache or the source.
        """
        if len(output_stream_ids) != 1:
            raise ValueError("This cacher is only applicable for single-output streams")

        original_ctx = ctx.clone()

        # Uses the recursive logic to call `self.execute_operation` which converts
        # the streams into `VirtualStream` objects (and likely mutates the `ctx`).

        # TODO: refactor to use an explicit self-method call instead of super()
        output_streams = await super().execute_operations(ctx=ctx, output_stream_ids=output_stream_ids)
        assert len(output_streams) == 1
        (virtual_stream,) = output_streams
        assert isinstance(virtual_stream, CacheVirtualStream)
        # Resolve TTL info and save BIQueryCacheOptions object
        cache_options_builder = self._cache_options_builder
        cache_options = cache_options_builder.get_cache_options_for_stream(virtual_stream, self._dataset)

        ds_id = self._dataset.uuid
        cache_helper = CacheProcessingHelper(
            entity_id=ds_id,  # type: ignore  # TODO: fix
            service_registry=self.service_registry,
        )

        async def _get_from_source() -> Optional[TValuesChunkStream]:
            streams = await self._main_processor.execute_operations(
                ctx=original_ctx,
                output_stream_ids=output_stream_ids,
            )
            assert len(streams) == 1
            (stream,) = streams
            return stream.data

        cache_full_hit = None
        try:
            sit, result_iter = await cache_helper.run_with_cache(
                generate_func=_get_from_source,  # type: ignore  # TODO: fix
                cache_options=cache_options,
            )
            if sit == CacheSituation.full_hit:
                cache_full_hit = True
            elif sit == CacheSituation.generated:
                cache_full_hit = False
        finally:
            self._save_cache_info_reporting_record(
                ctx=ctx,
                cache_full_hit=cache_full_hit,  # type: ignore  # TODO: fix
            )

        result_stream = DataStreamAsync(
            id=virtual_stream.id,
            names=virtual_stream.names,
            user_types=virtual_stream.user_types,
            data=result_iter,  # type: ignore  # TODO: fix
            meta=virtual_stream.meta,
            data_key=cache_options.key,
        )
        return [result_stream]

    def _validate_input_stream(self, stream: AbstractStream, op: BaseOp) -> None:
        if not isinstance(stream, CacheVirtualStream):
            super()._validate_input_stream(stream, op)

    def _validate_output_stream(self, stream: AbstractStream, op: BaseOp) -> None:
        if not isinstance(stream, CacheVirtualStream):
            super()._validate_output_stream(stream, op)


@attr.s
class CachedDatasetProcessor(SROperationProcessorAsyncBase):  # noqa
    _dataset: Dataset = attr.ib(kw_only=True)
    _main_processor: OperationProcessorAsyncBase = attr.ib(kw_only=True)

    _cache_processor: SingleStreamCacheOperationProcessor = attr.ib(init=False)

    def __attrs_post_init__(self) -> None:
        self._cache_processor = SingleStreamCacheOperationProcessor(
            dataset=self._dataset,
            service_registry=self.service_registry,
            main_processor=self._main_processor,
        )

    async def ping(self) -> Optional[int]:
        return await self._main_processor.ping()  # TODO: also ping cache engine?

    async def execute_operations(
        self,
        ctx: OpExecutionContext,
        output_stream_ids: Collection[str],
    ) -> List[DataStreamAsync]:
        if len(output_stream_ids) != 1:
            LOGGER.warning("Cannot currently use cache for multiple (%d) output streams", len(output_stream_ids))
            return await self._main_processor.execute_operations(
                ctx=ctx,
                output_stream_ids=output_stream_ids,
            )

        cache_processor = self._cache_processor
        return await cache_processor.execute_operations(
            ctx=ctx,
            output_stream_ids=output_stream_ids,
        )

    async def start(self) -> None:
        await self._main_processor.start()
        await self._cache_processor.start()

    async def end(self) -> None:
        await self._cache_processor.end()
        await self._main_processor.end()

lib/dl_core/dl_core/data_processing/processing/processor_dataset_cached.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    ClassVar,
    Dict,
    Optional,
    Type,
)

import attr

from dl_constants.enums import DataSourceRole
from dl_core.data_processing.processing.context import OpExecutionContext
from dl_core.data_processing.processing.db_base.exec_adapter_base import ProcessorDbExecAdapterBase
from dl_core.data_processing.processing.db_base.op_executors import (
    CalcOpExecutorAsync,
    DownloadOpExecutorAsync,
    JoinOpExecutorAsync,
    OpExecutorAsync,
)
from dl_core.data_processing.processing.operation import (
    BaseOp,
    CalcOp,
    DownloadOp,
    JoinOp,
)
from dl_core.data_processing.processing.processor import OperationProcessorAsyncBase
from dl_core.data_processing.processing.source_db.selector_exec_adapter import SourceDbExecAdapter
from dl_core.data_processing.selectors.base import DataSelectorAsyncBase
from dl_core.data_processing.stream_base import AbstractStream
from dl_core.us_dataset import Dataset
from dl_core.us_manager.local_cache import USEntryBuffer


@attr.s
class SourceDbOperationProcessor(OperationProcessorAsyncBase):
    _role: DataSourceRole = attr.ib(kw_only=True)
    _dataset: Dataset = attr.ib(kw_only=True)
    _selector: DataSelectorAsyncBase = attr.ib(kw_only=True)
    _row_count_hard_limit: Optional[int] = attr.ib(kw_only=True, default=None)
    _us_entry_buffer: USEntryBuffer = attr.ib(kw_only=True)

    _db_ex_adapter: Optional[ProcessorDbExecAdapterBase] = attr.ib(init=False, default=None)

    _executors: ClassVar[Dict[Type[BaseOp], Type[OpExecutorAsync]]] = {
        DownloadOp: DownloadOpExecutorAsync,
        CalcOp: CalcOpExecutorAsync,
        JoinOp: JoinOpExecutorAsync,
    }

    async def ping(self) -> Optional[int]:
        return 1

    async def start(self) -> None:
        self._db_ex_adapter = SourceDbExecAdapter(
            role=self._role,
            dataset=self._dataset,
            selector=self._selector,
            row_count_hard_limit=self._row_count_hard_limit,
            us_entry_buffer=self._us_entry_buffer,
        )

    async def stop(self) -> None:
        self._db_ex_adapter = None

    async def execute_operation(self, op: BaseOp, ctx: OpExecutionContext) -> AbstractStream:
        opex_cls: Type[OpExecutorAsync] = self._executors[type(op)]
        opex = opex_cls(db_ex_adapter=self._db_ex_adapter, ctx=ctx)  # type: ignore  # TODO: fix
        return await opex.execute(op)

lib/dl_core/dl_core/data_processing/processing/source_db/processor.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Optional,
    Sequence,
    Union,
)

import attr

from dl_constants.enums import UserDataType
from dl_core.data_processing.prepared_components.default_manager import DefaultPreparedComponentManager
from dl_core.data_processing.processing.db_base.exec_adapter_base import ProcessorDbExecAdapterBase
from dl_core.data_processing.selectors.dataset_base import DatasetDataSelectorAsyncBase
from dl_core.query.bi_query import QueryAndResultInfo


if TYPE_CHECKING:
    from sqlalchemy.sql.selectable import Select

    from dl_constants.enums import DataSourceRole
    from dl_core.data_processing.cache.primitives import LocalKeyRepresentation
    from dl_core.data_processing.prepared_components.manager_base import PreparedComponentManagerBase
    from dl_core.data_processing.prepared_components.primitives import PreparedMultiFromInfo
    from dl_core.data_processing.selectors.base import DataSelectorAsyncBase
    from dl_core.data_processing.types import TValuesChunkStream
    from dl_core.us_dataset import Dataset
    from dl_core.us_manager.local_cache import USEntryBuffer


@attr.s
class SourceDbExecAdapter(ProcessorDbExecAdapterBase):  # noqa
    _role: DataSourceRole = attr.ib(kw_only=True)
    _dataset: Dataset = attr.ib(kw_only=True)
    _selector: DataSelectorAsyncBase = attr.ib(kw_only=True)
    _prep_component_manager: Optional[PreparedComponentManagerBase] = attr.ib(kw_only=True, default=None)
    _row_count_hard_limit: Optional[int] = attr.ib(kw_only=True, default=None)
    _us_entry_buffer: USEntryBuffer = attr.ib(kw_only=True)

    def __attrs_post_init__(self):  # type: ignore  # TODO: fix
        if self._prep_component_manager is None:
            self._prep_component_manager = DefaultPreparedComponentManager(
                dataset=self._dataset,
                role=self._role,
                us_entry_buffer=self._us_entry_buffer,
            )

    def get_prep_component_manager(self) -> PreparedComponentManagerBase:
        assert self._prep_component_manager is not None
        return self._prep_component_manager

    def _make_query_res_info(
        self,
        query: Union[str, Select],
        user_types: Sequence[UserDataType],
    ) -> QueryAndResultInfo:
        query_res_info = QueryAndResultInfo(
            query=query,  # type: ignore  # TODO: fix
            user_types=list(user_types),
            # This is basically legacy and will be removed.
            # col_names are not really used anywhere, just passed around a lot.
            # So we generate random ones here
            col_names=[f"col_{i}" for i in range(len(user_types))],
        )
        return query_res_info

    async def _execute_and_fetch(
        self,
        *,
        query: Union[str, Select],
        user_types: Sequence[UserDataType],
        chunk_size: int,
        joint_dsrc_info: Optional[PreparedMultiFromInfo] = None,
        query_id: str,
    ) -> TValuesChunkStream:
        assert not isinstance(query, str), "String queries are not supported by source DB processor"
        assert joint_dsrc_info is not None, "joint_dsrc_info is required for source DB processor"

        query_res_info = self._make_query_res_info(query=query, user_types=user_types)
        data_stream = await self._selector.get_data_stream(
            query_id=query_id,
            role=self._role,
            joint_dsrc_info=joint_dsrc_info,
            query_res_info=query_res_info,
            row_count_hard_limit=self._row_count_hard_limit,
        )
        return data_stream.data

    def get_data_key(
        self,
        *,
        query_id: str,
        query: Union[str, Select],
        user_types: Sequence[UserDataType],
        joint_dsrc_info: Optional[PreparedMultiFromInfo] = None,
    ) -> Optional[LocalKeyRepresentation]:
        selector = self._selector
        assert isinstance(selector, DatasetDataSelectorAsyncBase)
        query_res_info = self._make_query_res_info(query=query, user_types=user_types)
        query_execution_ctx = selector.build_query_execution_ctx(
            query_res_info=query_res_info,
            joint_dsrc_info=joint_dsrc_info,  # type: ignore  # TODO: fix
            role=self._role,
            query_id=query_id,
        )
        data_key = selector.get_data_key(query_execution_ctx=query_execution_ctx)
        return data_key

lib/dl_core/dl_core/data_processing/processing/source_db/selector_exec_adapter.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    List,
    Optional,
    Sequence,
)

import attr


if TYPE_CHECKING:
    from sqlalchemy.sql.selectable import Select

    from dl_constants.enums import (
        DataSourceRole,
        UserDataType,
    )
    from dl_core.data_processing.cache.primitives import BIQueryCacheOptions
    from dl_core.data_processing.prepared_components.primitives import PreparedMultiFromInfo
    from dl_core.data_processing.stream_base import DataStreamAsync
    from dl_core.query.bi_query import QueryAndResultInfo
    from dl_core.us_connection_base import ExecutorBasedMixin


class NoData(Exception):
    pass


@attr.s(frozen=True, auto_attribs=True)
class BIQueryExecutionContext:
    query_id: str
    query: Select
    compiled_query: str  # for logs only
    target_connection: ExecutorBasedMixin
    requested_bi_types: List[UserDataType]
    result_col_names: Sequence[str]
    target_db_name: Optional[str] = attr.ib(default=None)
    cache_options: Optional[BIQueryCacheOptions] = attr.ib(default=None)
    connect_args: dict = attr.ib(default=None)


class DataSelectorAsyncBase(abc.ABC):
    @abc.abstractmethod
    async def get_data_stream(
        self,
        *,
        query_id: Optional[str] = None,
        role: DataSourceRole,
        query_res_info: QueryAndResultInfo,
        joint_dsrc_info: PreparedMultiFromInfo,
        row_count_hard_limit: Optional[int] = None,
        stream_id: Optional[str] = None,
    ) -> DataStreamAsync:
        """
        Fetch data from the database.
        Return SQL query compiled as a string and an iterable or result rows.

        :param role: data source role to use
        :param query_res_info: an object containing a compiled SA ``Select`` and info about result columns
        :param joint_dsrc_info: represents the combined SQL data source (joined tables)
        :param row_count_hard_limit: Result rows count limit. Will raise ResultRowCountLimitExceeded if receive more.
        :return: ``DataStreamAsync`` instance containing an async iterable of lists
            and some metadata such as compiled query, column names and data types
        """

        raise NotImplementedError

    async def close(self) -> None:
        pass

lib/dl_core/dl_core/data_processing/selectors/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from contextlib import contextmanager
import logging
import time
from typing import (
    TYPE_CHECKING,
    Generator,
    Optional,
)

import attr

from dl_api_commons.reporting.models import (
    QueryExecutionEndReportingRecord,
    QueryExecutionStartReportingRecord,
)
from dl_constants.enums import DataSourceRole
from dl_core import utils
from dl_core.base_models import WorkbookEntryLocation
from dl_core.data_processing.selectors.base import (
    BIQueryExecutionContext,
    DataSelectorAsyncBase,
    NoData,
)
from dl_core.data_processing.selectors.utils import get_query_type
from dl_core.data_processing.stream_base import (
    DataRequestMetaInfo,
    DataStreamAsync,
)
import dl_core.exc as exc
from dl_core.query.bi_query import QueryAndResultInfo
from dl_core.us_connection_base import ExecutorBasedMixin
from dl_core.utils import make_id


if TYPE_CHECKING:
    from dl_api_commons.reporting.registry import ReportingRegistry
    from dl_core.data_processing.cache.primitives import LocalKeyRepresentation
    from dl_core.data_processing.prepared_components.primitives import PreparedMultiFromInfo
    from dl_core.data_processing.types import TValuesChunkStream
    from dl_core.services_registry import ServicesRegistry
    from dl_core.us_dataset import Dataset
    from dl_core.us_manager.local_cache import USEntryBuffer


LOGGER = logging.getLogger(__name__)


# noinspection PyDataclass
@attr.s
class DatasetDataSelectorAsyncBase(DataSelectorAsyncBase, metaclass=abc.ABCMeta):
    dataset: Dataset = attr.ib(kw_only=True)
    _us_entry_buffer: USEntryBuffer = attr.ib(kw_only=True)
    _service_registry: ServicesRegistry = attr.ib(kw_only=True)  # Service registry override

    """
    Base class for dataset-dependent asynchronous data selectors
    :param dataset: the dataset to operate on
    """

    def _save_start_exec_reporting_record(
        self,
        query_execution_ctx: BIQueryExecutionContext,
    ) -> None:
        connection = query_execution_ctx.target_connection
        workbook_id = (
            connection.entry_key.workbook_id if isinstance(connection.entry_key, WorkbookEntryLocation) else None
        )
        report = QueryExecutionStartReportingRecord(
            timestamp=time.time(),
            query_id=query_execution_ctx.query_id,
            dataset_id=self.dataset.uuid,  # type: ignore  # TODO: fix
            query_type=get_query_type(
                connection=connection,
                conn_sec_mgr=self.service_registry.get_conn_executor_factory().conn_security_manager,
            ),
            connection_type=connection.conn_type,
            conn_reporting_data=connection.get_conn_dto().conn_reporting_data(),
            query=query_execution_ctx.compiled_query,
            workbook_id=workbook_id,
        )
        self.reporting_registry.save_reporting_record(report=report)

    def _save_end_exec_reporting_record(
        self,
        query_execution_ctx: BIQueryExecutionContext,
        exec_exception: Optional[Exception],
    ) -> None:
        report = QueryExecutionEndReportingRecord(
            timestamp=time.time(),
            query_id=query_execution_ctx.query_id,
            exception=exec_exception,
        )
        self.reporting_registry.save_reporting_record(report=report)

    def pre_exec(self, query_execution_ctx: BIQueryExecutionContext) -> None:
        self._save_start_exec_reporting_record(query_execution_ctx)

    def post_exec(
        self,
        query_execution_ctx: BIQueryExecutionContext,
        exec_exception: Optional[Exception],
    ) -> None:
        self._save_end_exec_reporting_record(query_execution_ctx, exec_exception)

    @contextmanager
    def _execution_cm(
        self,
        query_execution_ctx: BIQueryExecutionContext,
    ) -> Generator[None, None, None]:
        self.pre_exec(query_execution_ctx=query_execution_ctx)
        exec_exception: Optional[Exception] = None
        try:
            yield
        except Exception as fired_exception:
            exec_exception = fired_exception
            raise
        finally:
            self.post_exec(query_execution_ctx=query_execution_ctx, exec_exception=exec_exception)

    def get_data_key(
        self,
        query_execution_ctx: BIQueryExecutionContext,
    ) -> Optional[LocalKeyRepresentation]:
        return None

    async def get_data_stream(
        self,
        *,
        query_id: Optional[str] = None,
        role: DataSourceRole,
        query_res_info: QueryAndResultInfo,
        joint_dsrc_info: PreparedMultiFromInfo,
        row_count_hard_limit: Optional[int] = None,
        stream_id: Optional[str] = None,
    ) -> DataStreamAsync:
        """Generate data stream from a data source"""

        query_id = query_id or make_id()
        query_execution_ctx = self.build_query_execution_ctx(
            query_id=query_id,
            query_res_info=query_res_info,
            role=role,
            joint_dsrc_info=joint_dsrc_info,
        )

        if not isinstance(query_execution_ctx.target_connection, ExecutorBasedMixin):
            raise exc.NotAvailableError(
                f"Connection {type(query_execution_ctx.target_connection).__qualname__}"
                f" does not support async data selection"
            )

        with self._execution_cm(query_execution_ctx):
            result_iter = await self.execute_query_context(
                role=role,
                query_execution_ctx=query_execution_ctx,
                row_count_hard_limit=row_count_hard_limit,
            )
            if result_iter is None:
                raise NoData("Got no data from selector")

            if row_count_hard_limit is not None:
                result_iter = result_iter.limit(max_count=row_count_hard_limit)

            data_key = self.get_data_key(query_execution_ctx=query_execution_ctx)
            stream_id = stream_id or make_id()
            LOGGER.info("Making data stream %s,", stream_id, extra=dict(data_key=str(data_key)))
            assert stream_id is not None

            data_source_list = joint_dsrc_info.data_source_list
            assert data_source_list is not None

            return DataStreamAsync(
                id=stream_id,
                user_types=query_execution_ctx.requested_bi_types,
                names=query_execution_ctx.result_col_names,
                data=result_iter,
                meta=DataRequestMetaInfo(
                    query_id=query_id,
                    query=query_execution_ctx.compiled_query,
                    is_materialized=role != DataSourceRole.origin,
                    data_source_list=data_source_list,
                ),
                data_key=data_key,
            )

    @property
    def service_registry(self) -> ServicesRegistry:
        return self._service_registry

    @property
    def reporting_registry(self) -> ReportingRegistry:
        return self.service_registry.get_reporting_registry()

    def build_query_execution_ctx(
        self,
        *,
        query_id: str,
        query_res_info: QueryAndResultInfo,
        role: DataSourceRole,
        joint_dsrc_info: PreparedMultiFromInfo,
    ) -> BIQueryExecutionContext:
        compiled_query = utils.compile_query_for_debug(query_res_info.query, joint_dsrc_info.query_compiler.dialect)
        LOGGER.info(f"SQL query for dataset: {compiled_query}")

        assert joint_dsrc_info.target_connection_ref is not None
        target_connection = self._us_entry_buffer.get_entry(joint_dsrc_info.target_connection_ref)
        assert isinstance(target_connection, ExecutorBasedMixin)

        return BIQueryExecutionContext(
            query_id=query_id,
            query=query_res_info.query,
            compiled_query=compiled_query,
            target_connection=target_connection,
            target_db_name=joint_dsrc_info.db_name,
            requested_bi_types=query_res_info.user_types,
            result_col_names=query_res_info.col_names,
            cache_options=None,
            connect_args=joint_dsrc_info.connect_args,
        )

    @abc.abstractmethod
    async def execute_query_context(
        self,
        role: DataSourceRole,
        query_execution_ctx: BIQueryExecutionContext,
        row_count_hard_limit: Optional[int] = None,
    ) -> Optional[TValuesChunkStream]:
        """Get data using info in ``query_execution_ctx``"""
        raise NotImplementedError

lib/dl_core/dl_core/data_processing/selectors/dataset_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import Optional

import attr

from dl_constants.enums import DataSourceRole
from dl_core.data_processing.cache.exc import CachePreparationFailed
from dl_core.data_processing.cache.primitives import LocalKeyRepresentation
from dl_core.data_processing.cache.utils import SelectorCacheOptionsBuilder
from dl_core.data_processing.prepared_components.primitives import PreparedMultiFromInfo
from dl_core.data_processing.selectors.base import BIQueryExecutionContext
from dl_core.data_processing.selectors.dataset_base import DatasetDataSelectorAsyncBase
from dl_core.query.bi_query import QueryAndResultInfo


LOGGER = logging.getLogger(__name__)


@attr.s
class DatasetCacheCommonDataSelectorAsyncBase(DatasetDataSelectorAsyncBase, metaclass=abc.ABCMeta):
    """
    Abstract class that implement base cache-related operations:
    - Cache key calculation
    - TTL determination

    If cache usage is allowed in class, it will build cache options during building query execution context.
    """

    _allow_cache_usage: bool = attr.ib(default=True, kw_only=True)
    _is_bleeding_edge_user: bool = attr.ib(default=False)
    _cache_options_builder: SelectorCacheOptionsBuilder = attr.ib(kw_only=True)

    @_cache_options_builder.default  # noqa
    def _cache_options_builder_default(self) -> SelectorCacheOptionsBuilder:
        return SelectorCacheOptionsBuilder(
            default_ttl_config=self._service_registry.default_cache_ttl_config,
            is_bleeding_edge_user=self._is_bleeding_edge_user,
            us_entry_buffer=self._us_entry_buffer,
        )

    def get_data_key(self, query_execution_ctx: BIQueryExecutionContext) -> Optional[LocalKeyRepresentation]:
        cache_options = query_execution_ctx.cache_options
        if cache_options is not None:
            return cache_options.key
        return super().get_data_key(query_execution_ctx=query_execution_ctx)

    def build_query_execution_ctx(  # type: ignore  # TODO: fix
        self,
        *,
        query_id: str,
        query_res_info: QueryAndResultInfo,
        role: DataSourceRole,
        joint_dsrc_info: PreparedMultiFromInfo,
    ) -> BIQueryExecutionContext:
        q_exec_ctx: BIQueryExecutionContext = super().build_query_execution_ctx(
            query_id=query_id,
            query_res_info=query_res_info,
            role=role,
            joint_dsrc_info=joint_dsrc_info,
        )
        if self._allow_cache_usage:
            try:
                cache_options = self._cache_options_builder.get_cache_options(
                    joint_dsrc_info=joint_dsrc_info,
                    role=role,
                    query=q_exec_ctx.query,
                    user_types=q_exec_ctx.requested_bi_types,
                    dataset=self.dataset,
                )

                return attr.evolve(
                    q_exec_ctx,
                    cache_options=cache_options,
                )
            except CachePreparationFailed:
                LOGGER.exception("Cache preparation failed")
                # Do not fail the request though (very likely it will still fail)

        return q_exec_ctx

lib/dl_core/dl_core/data_processing/selectors/dataset_cache_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import time
from typing import (
    TYPE_CHECKING,
    Optional,
)

import attr

from dl_api_commons.reporting.models import QueryExecutionCacheInfoReportingRecord
from dl_constants.enums import DataSourceRole
from dl_core.data_processing.cache.processing_helper import (
    CacheProcessingHelper,
    CacheSituation,
)
from dl_core.data_processing.selectors.base import BIQueryExecutionContext
from dl_core.data_processing.selectors.dataset_cache_base import DatasetCacheCommonDataSelectorAsyncBase
from dl_core.data_processing.selectors.db import DatasetDbDataSelectorAsync


if TYPE_CHECKING:
    from dl_core.data_processing.types import TValuesChunkStream
    from dl_core.us_connection_base import ExecutorBasedMixin

LOGGER = logging.getLogger(__name__)


@attr.s
class CachedDatasetDataSelectorAsync(DatasetCacheCommonDataSelectorAsyncBase):
    """
    Asynchronous cached dataset data selector

    :param allow_cache_usage: if ``True``, cache **might** be used. Otherwise,
    neither read nor write will be executed.
    """

    _allow_cache_usage: bool = attr.ib(default=True, kw_only=True)
    # cache-wrapped selector
    _db_selector: DatasetDbDataSelectorAsync = attr.ib(init=False)
    _cache_helper: Optional[CacheProcessingHelper] = attr.ib(init=False, default=None)

    def __attrs_post_init__(self) -> None:
        self._db_selector = DatasetDbDataSelectorAsync(
            dataset=self.dataset,
            service_registry=self._service_registry,
            us_entry_buffer=self._us_entry_buffer,
        )
        ds_id = self.dataset.uuid
        if ds_id:
            self._cache_helper = CacheProcessingHelper(
                entity_id=ds_id,
                service_registry=self.service_registry,
            )

    async def execute_query_context(
        self,
        role: DataSourceRole,
        query_execution_ctx: BIQueryExecutionContext,
        row_count_hard_limit: Optional[int] = None,
    ) -> Optional[TValuesChunkStream]:
        async def _request_db() -> Optional[TValuesChunkStream]:
            result_iter = await self._db_selector.execute_query_context(
                role=role, query_execution_ctx=query_execution_ctx, row_count_hard_limit=row_count_hard_limit
            )

            if result_iter is None:
                return None

            # Is this still necessary, after passing `row_count_hard_limit` to the `db_selector`?
            if row_count_hard_limit is not None:
                result_iter = result_iter.limit(max_count=row_count_hard_limit)

            return result_iter

        cache_helper = self._cache_helper
        cache_options = query_execution_ctx.cache_options
        if cache_helper is None or cache_options is None:
            # cache not applicable, call db selector directly
            return await _request_db()

        target_connection: ExecutorBasedMixin = query_execution_ctx.target_connection
        # TODO: make this env-configurable through settings.
        use_locked_cache = target_connection.use_locked_cache

        cache_full_hit = None
        try:
            sit, result_iter = await cache_helper.run_with_cache(
                generate_func=_request_db,  # type: ignore  # TODO: fix
                cache_options=cache_options,
                use_locked_cache=use_locked_cache,
            )
            if sit == CacheSituation.full_hit:
                cache_full_hit = True
            elif sit == CacheSituation.generated:
                cache_full_hit = False
        finally:
            query_cache_rr = QueryExecutionCacheInfoReportingRecord(
                query_id=query_execution_ctx.query_id,
                cache_full_hit=cache_full_hit,
                timestamp=time.time(),
            )
            self.reporting_registry.save_reporting_record(query_cache_rr)

        return result_iter  # type: ignore  # TODO: fix

lib/dl_core/dl_core/data_processing/selectors/dataset_cached.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    List,
    Optional,
    Sequence,
)

import attr

from dl_constants.enums import DataSourceRole
from dl_core.data_processing.selectors.base import BIQueryExecutionContext
from dl_core.data_processing.selectors.dataset_cached import CachedDatasetDataSelectorAsync
from dl_core.data_processing.streaming import (
    AsyncChunked,
    AsyncChunkedBase,
    LazyAsyncChunked,
)


if TYPE_CHECKING:
    from dl_constants.types import TBIDataValue


LOGGER = logging.getLogger(__name__)


@attr.s
class LazyCachedDatasetDataSelectorAsync(CachedDatasetDataSelectorAsync):
    """
    Lazy asynchronous cached dataset data selector
    """

    _active_queries: List[BIQueryExecutionContext] = attr.ib(init=False, factory=list)

    # pre_exec is not redefined so that a reporting record is created
    # even if the selector doesn't fetch any data in the end

    def post_exec(
        self,
        query_execution_ctx: BIQueryExecutionContext,
        exec_exception: Optional[Exception],
    ) -> None:
        """Lazy behaviour: since the query has not even been executed at this point, do nothing."""

    async def _close_and_remove_active_query(self, query_execution_ctx: BIQueryExecutionContext):  # type: ignore  # TODO: fix
        if query_execution_ctx in self._active_queries:
            self._active_queries.remove(query_execution_ctx)
            super().post_exec(query_execution_ctx=query_execution_ctx, exec_exception=None)

    async def execute_query_context(
        self,
        role: DataSourceRole,
        query_execution_ctx: BIQueryExecutionContext,
        row_count_hard_limit: Optional[int] = None,
    ) -> Optional[AsyncChunkedBase[Sequence[TBIDataValue]]]:
        result_iter_awaitable = super().execute_query_context(
            role=role, query_execution_ctx=query_execution_ctx, row_count_hard_limit=row_count_hard_limit
        )

        async def initialize_data_stream() -> AsyncChunkedBase[List[TBIDataValue]]:
            wrapped_result_iter = await result_iter_awaitable
            if wrapped_result_iter is None:
                wrapped_result_iter = AsyncChunked.from_chunked_iterable([[]])
            self._active_queries.append(query_execution_ctx)
            return wrapped_result_iter  # type: ignore  # TODO: fix

        async def finalize_data_stream() -> None:
            await self._close_and_remove_active_query(query_execution_ctx=query_execution_ctx)

        return LazyAsyncChunked(initializer=initialize_data_stream, finalizer=finalize_data_stream)  # type: ignore  # TODO: fix

    async def close(self) -> None:
        """Close all active queries"""
        while self._active_queries:
            await self._close_and_remove_active_query(query_execution_ctx=self._active_queries[-1])

lib/dl_core/dl_core/data_processing/selectors/dataset_cached_lazy.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Optional,
    Sequence,
)

import attr

from dl_constants.enums import DataSourceRole
from dl_core.connection_executors import ConnExecutorQuery
from dl_core.data_processing.selectors.dataset_base import DatasetDataSelectorAsyncBase
from dl_core.data_processing.selectors.utils import select_data_context
from dl_core.data_processing.streaming import (
    AsyncChunked,
    AsyncChunkedBase,
)
import dl_core.exc as exc
from dl_core.us_connection_base import ExecutorBasedMixin


if TYPE_CHECKING:
    from dl_constants.types import TBIDataValue
    from dl_core.data_processing.selectors.base import BIQueryExecutionContext


@attr.s
class DatasetDbDataSelectorAsync(DatasetDataSelectorAsyncBase):
    """Async selector that fetches data from the database"""

    async def execute_query_context(
        self,
        role: DataSourceRole,
        query_execution_ctx: BIQueryExecutionContext,
        row_count_hard_limit: Optional[int] = None,
    ) -> Optional[AsyncChunkedBase[Sequence[TBIDataValue]]]:
        if not isinstance(query_execution_ctx.target_connection, ExecutorBasedMixin):
            raise exc.NotAvailableError(
                f"Connection {type(query_execution_ctx.target_connection).__qualname__}"
                f" does not support async data selection"
            )

        with select_data_context(role=role):
            ce_factory = self.service_registry.get_conn_executor_factory()
            ce = ce_factory.get_async_conn_executor(query_execution_ctx.target_connection)

            exec_result = await ce.execute(
                ConnExecutorQuery(
                    query=query_execution_ctx.query,
                    db_name=query_execution_ctx.target_db_name,
                    user_types=query_execution_ctx.requested_bi_types,
                    # connect_args=query_execution_ctx.connect_args,
                    debug_compiled_query=query_execution_ctx.compiled_query,
                    chunk_size=None,
                )
            )
            return AsyncChunked(chunked_data=exec_result.result)

lib/dl_core/dl_core/data_processing/selectors/db.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import contextmanager
import logging
from typing import (
    TYPE_CHECKING,
    Generator,
    Sequence,
)

import sqlalchemy as sa

from dl_constants.enums import (
    DataSourceRole,
    ReportingQueryType,
)
import dl_core.exc as exc
from dl_core.query.bi_query import BIQuery
from dl_core.query.expression import ExpressionCtx
from dl_core.us_connection_base import ClassicConnectionSQL


if TYPE_CHECKING:
    from dl_core.connections_security.base import ConnectionSecurityManager
    from dl_core.us_connection_base import ConnectionBase


LOGGER = logging.getLogger(__name__)


@contextmanager
def select_data_context(role: DataSourceRole) -> Generator[None, None, None]:
    """
    Re-raise database-related errors taking dataset settings into account
    """
    try:
        yield
    except exc.SourceDoesNotExist as err:
        if role == DataSourceRole.materialization:
            raise exc.MaterializationNotFinished(db_message=err.db_message, query=err.query)
        raise


def get_value_range_query(expression: ExpressionCtx, dimension_filters: Sequence[ExpressionCtx] = ()) -> BIQuery:
    return BIQuery(
        select_expressions=[
            ExpressionCtx(
                expression=sa.func.min(expression.expression),
                alias="min_val",
                user_type=expression.user_type,
                avatar_ids=expression.avatar_ids,
            ),
            ExpressionCtx(
                expression=sa.func.max(expression.expression),
                alias="max_val",
                user_type=expression.user_type,
                avatar_ids=expression.avatar_ids,
            ),
        ],
        dimension_filters=dimension_filters,
    )


def get_query_type(connection: ConnectionBase, conn_sec_mgr: ConnectionSecurityManager) -> ReportingQueryType:
    if connection.is_always_internal_source:
        return ReportingQueryType.internal

    if isinstance(connection, ClassicConnectionSQL):
        if conn_sec_mgr.is_internal_connection(connection.get_conn_dto()):
            return ReportingQueryType.internal

    return ReportingQueryType.external

lib/dl_core/dl_core/data_processing/selectors/utils.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from itertools import chain
import logging
from typing import (
    Any,
    ClassVar,
    Collection,
    Dict,
    List,
    NamedTuple,
    Optional,
    Set,
    Tuple,
    Union,
)

import attr
import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

from dl_constants.enums import JoinType
from dl_core import exc
from dl_core.components.ids import AvatarId
from dl_core.data_processing.prepared_components.primitives import (
    PreparedMultiFromInfo,
    PreparedSingleFromInfo,
)
from dl_core.query.bi_query import SqlSourceType
from dl_core.query.expression import JoinOnExpressionCtx


LOGGER = logging.getLogger(__name__)

JoinExpressionType = Union[ClauseElement, str]

_MAX_AVATAR_REUSAGE_CNT = 10


class AvatarJoinInfo(NamedTuple):
    avatar_id: AvatarId
    join_type: JoinType
    on_clause: JoinExpressionType


@attr.s
class SqlSourceBuilder:
    """
    :param avatar_alias_mapper: a string-to-string function that should be used for converting avatar aliases
    :param prep_component_manager: an ``PreparedComponentManagerBase`` instance to provide access to data sources
        via avatar IDs. These may be real, straight from the dataset, or may be virtual, such as nested queries
    """

    # Constants
    _COMMON_JOIN_PARAMS: ClassVar[Dict[str, bool]] = dict(all=False, any=False, global_=False)
    _JOIN_PARAMS_BY_TYPE: ClassVar[Dict[JoinType, Dict[str, Any]]] = {
        JoinType.inner: {"type": "inner"},
        JoinType.left: {"isouter": True, "type": "left"},
        JoinType.right: {"type": "right"},
        JoinType.full: {"isouter": True, "full": True, "type": "full"},
    }

    def validate_prep_src_infos(self, required_prep_src_infos: Collection[PreparedSingleFromInfo]) -> None:
        if len({prep_src_info.target_connection_ref for prep_src_info in required_prep_src_infos}) != 1:
            raise ValueError("There is a divergence in data source connections")

    def make_join_clause(
        self,
        left: SqlSourceType,
        right: SqlSourceType,
        join_type: JoinType,
        on_clause: JoinExpressionType,
    ) -> SqlSourceType:
        join_params = dict(self._COMMON_JOIN_PARAMS, **self._JOIN_PARAMS_BY_TYPE[join_type])
        result = sa.join(left, right, on_clause)  # type: ignore
        for name, value in join_params.items():
            setattr(result, name, value)
        return result

    def collect_avatar_join_info_list(
        self,
        root_avatar_id: AvatarId,
        required_avatar_ids: Collection[AvatarId],
        join_on_expressions: Collection[JoinOnExpressionCtx] = (),
    ) -> Tuple[Set[AvatarId], List[AvatarJoinInfo]]:
        avatar_join_info_list: List[AvatarJoinInfo] = []
        used_avatar_ids: Set[AvatarId] = {root_avatar_id}
        used_avatar_ids_as_list: List[AvatarId] = []

        relations_by_left: Dict[Optional[AvatarId], List[JoinOnExpressionCtx]] = defaultdict(list)
        for rel in join_on_expressions:
            relations_by_left[rel.left_id].append(rel)

        def _recursive_collect(avatar_id: Optional[AvatarId], ambiguous_relations: bool) -> None:
            avatar_user_relations = [
                rel
                for rel in relations_by_left[avatar_id]
                if ((rel.left_id is not None or ambiguous_relations) and rel.right_id in required_avatar_ids)
            ]
            for prep_relation in avatar_user_relations:
                child_avatar_id = prep_relation.right_id
                if used_avatar_ids_as_list.count(child_avatar_id) > _MAX_AVATAR_REUSAGE_CNT:
                    raise exc.DatasetConfigurationError("Avatar re-usage limit has been reached")
                used_avatar_ids.add(child_avatar_id)
                used_avatar_ids_as_list.append(child_avatar_id)

                relation_expr_ctx = prep_relation.expression
                relation_expr: JoinExpressionType = relation_expr_ctx.expression
                LOGGER.info(f"Including relation between avatars {avatar_id} and {child_avatar_id}")
                avatar_join_info_list.append(
                    AvatarJoinInfo(
                        avatar_id=child_avatar_id, join_type=prep_relation.join_type, on_clause=relation_expr
                    )
                )
                if not ambiguous_relations:
                    # feature-managed relations are not recursive
                    _recursive_collect(avatar_id=child_avatar_id, ambiguous_relations=ambiguous_relations)

        # feature-managed relations can have references to any of the other tables
        # in their formula expressions, so they must be joined to last
        # (note that they all have root as their left avatar)
        # so first collect user joins, and then feature joins
        _recursive_collect(root_avatar_id, ambiguous_relations=False)
        # feature-managed relations are all formally bound to root and have no real left ID (are ambiguous)
        _recursive_collect(None, ambiguous_relations=True)

        return used_avatar_ids, avatar_join_info_list

    def build_from_avatar_join_info(
        self,
        *,
        root_avatar_id: AvatarId,
        avatar_join_info_list: List[AvatarJoinInfo],
        prepared_sources: Collection[PreparedSingleFromInfo],
    ) -> SqlSourceType:
        prep_sources_as_dict: Dict[AvatarId, PreparedSingleFromInfo] = {psi.id: psi for psi in prepared_sources}

        root_prep_src_info = prep_sources_as_dict[root_avatar_id]
        source = root_prep_src_info.non_null_sql_source
        for avatar_join_info in avatar_join_info_list:
            avatar_id = avatar_join_info.avatar_id
            prep_src_info = prep_sources_as_dict[avatar_id]
            join_type = avatar_join_info.join_type
            if not prep_src_info.supports_join_type(join_type):
                raise exc.DatasetConfigurationError(f"Data source does not support {join_type.name.upper()} JOIN")
            right = prep_src_info.non_null_sql_source
            source = self.make_join_clause(
                left=source,
                right=right,
                join_type=join_type,
                on_clause=avatar_join_info.on_clause,
            )

        return source

    def build_source(
        self,
        *,
        root_avatar_id: AvatarId,
        prepared_sources: Collection[PreparedSingleFromInfo],
        join_on_expressions: Collection[JoinOnExpressionCtx] = (),
        use_empty_source: bool = False,
    ) -> PreparedMultiFromInfo:
        """
        Compile source for the FROM-clause of the future SELECT statements.
        Validate joint info of all data sources for role in dataset
         (e.g. dialect, database name, query compiler)

        :param root_avatar_id: ID of the avatar that is the root of the required avatar tree
        :param prepared_sources: SQL sources to be joined
        :param join_on_expressions: The relations to be used for joining
        :param use_empty_source: force usage of an empty sql source
        """

        required_avatar_ids = {psi.id for psi in prepared_sources}
        self.validate_prep_src_infos(prepared_sources)

        prepared_sources = sorted(prepared_sources, key=lambda prep_src_info: prep_src_info.id)
        assert prepared_sources
        first_prep_source = prepared_sources[0]
        data_source_list = tuple(
            chain.from_iterable((prep_src_info.data_source_list or ()) for prep_src_info in prepared_sources)
        )
        pass_db_query_to_user = all(prep_src_info.pass_db_query_to_user for prep_src_info in prepared_sources)

        used_avatar_ids, avatar_join_info_list = self.collect_avatar_join_info_list(
            root_avatar_id=root_avatar_id,
            required_avatar_ids=required_avatar_ids,
            join_on_expressions=join_on_expressions,
        )
        if used_avatar_ids != required_avatar_ids:
            LOGGER.warning(
                f"Used avatar IDs are different from required: required: {required_avatar_ids}, "
                f"used: {used_avatar_ids}"
            )

        joint_source: Optional[SqlSourceType] = None
        if not use_empty_source:
            joint_source = self.build_from_avatar_join_info(
                root_avatar_id=root_avatar_id,
                avatar_join_info_list=avatar_join_info_list,
                prepared_sources=prepared_sources,
            )

        joint_dsrc_info = PreparedMultiFromInfo(
            sql_source=joint_source,
            data_source_list=data_source_list,
            query_compiler=first_prep_source.query_compiler,
            db_name=first_prep_source.db_name,
            connect_args=first_prep_source.connect_args,
            supported_join_types=frozenset(JoinType),  # TODO: determine this honestly
            pass_db_query_to_user=pass_db_query_to_user,
            target_connection_ref=first_prep_source.target_connection_ref,
        )
        return joint_dsrc_info

lib/dl_core/dl_core/data_processing/source_builder.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Collection,
    Iterable,
    Optional,
    Sequence,
)

import attr


if TYPE_CHECKING:
    from dl_constants.enums import UserDataType
    from dl_constants.types import TBIDataRow
    from dl_core.components.ids import AvatarId
    from dl_core.data_processing.cache.primitives import LocalKeyRepresentation
    from dl_core.data_processing.prepared_components.primitives import (
        PreparedMultiFromInfo,
        PreparedSingleFromInfo,
    )
    from dl_core.data_processing.types import TValuesChunkStream
    import dl_core.data_source


@attr.s
class AbstractStream:
    id: str = attr.ib(kw_only=True)
    names: Sequence[str] = attr.ib(kw_only=True)
    user_types: Sequence[UserDataType] = attr.ib(kw_only=True)


@attr.s
class DataRequestMetaInfo:
    query_id: Optional[str] = attr.ib(default=None, kw_only=True)
    query: Optional[str] = attr.ib(default=None, kw_only=True)
    is_materialized: bool = attr.ib(default=False, kw_only=True)
    data_source_list: Collection[dl_core.data_source.DataSource] = attr.ib(default=(), kw_only=True)
    pass_db_query_to_user: bool = attr.ib(default=True, kw_only=True)


@attr.s
class DataStreamBase(AbstractStream):
    meta: DataRequestMetaInfo = attr.ib(kw_only=True)
    data_key: Optional[LocalKeyRepresentation] = attr.ib(kw_only=True)
    # data: either sync or async, defined in subclasses

    @property
    def debug(self) -> DataRequestMetaInfo:  # TODO: Remove
        return self.meta


@attr.s
class DataStream(DataStreamBase):
    data: Iterable[TBIDataRow] = attr.ib(kw_only=True)


@attr.s
class DataStreamAsync(DataStreamBase):
    data: TValuesChunkStream = attr.ib(kw_only=True)


@attr.s
class VirtualStream(AbstractStream):
    """A representation of data that is being streamed in an external system (database)"""


@attr.s
class DataSourceVS(VirtualStream):
    """A data source avatar"""

    result_id: AvatarId = attr.ib(kw_only=True)
    alias: str = attr.ib(kw_only=True)
    prep_src_info: PreparedSingleFromInfo = attr.ib(kw_only=True)


@attr.s
class JointDataSourceVS(VirtualStream):
    """Joint data source info"""

    joint_dsrc_info: PreparedMultiFromInfo = attr.ib(kw_only=True)


@attr.s
class CacheVirtualStream(DataStreamBase):
    """Represents a virtual stream for cache processor. Has no real data"""

lib/dl_core/dl_core/data_processing/stream_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Any,
    AsyncGenerator,
    AsyncIterable,
    Awaitable,
    Callable,
    Generic,
    Iterable,
    List,
    Optional,
    Sequence,
    Type,
    TypeVar,
)

import attr

import dl_core.exc as exc


LOGGER = logging.getLogger(__name__)


_ENTRY_TV = TypeVar("_ENTRY_TV")
# Alias for readability
TChunk = Sequence  # `Sequence[_ENTRY_TV]`  # Chunk of "rows" (which might be anything)


class AsyncChunkedBase(Generic[_ENTRY_TV]):
    """A wrapper for an iterable providing both chunked and pre-item interfaces"""

    @property
    def items(self) -> AsyncIterable[_ENTRY_TV]:
        raise NotImplementedError

    @property
    def chunks(self) -> AsyncIterable[TChunk[_ENTRY_TV]]:
        raise NotImplementedError

    async def all(self) -> List[_ENTRY_TV]:
        result: List[_ENTRY_TV] = []
        async for chunk in self.chunks:
            result += chunk
        return result

    def limit(self, max_count: int) -> AsyncChunkedLimited:
        return AsyncChunkedLimited(chunked=self, max_count=max_count)


@attr.s
class AsyncChunked(AsyncChunkedBase[_ENTRY_TV]):
    _chunked_data: AsyncIterable[TChunk[_ENTRY_TV]] = attr.ib(repr=False)

    @classmethod
    def from_chunked_iterable(
        cls,
        sync_chunked_data: Iterable[TChunk[_ENTRY_TV]],
    ) -> AsyncChunked[_ENTRY_TV]:
        async def async_chunk_gen() -> AsyncGenerator[TChunk[_ENTRY_TV], None]:
            for chunk in sync_chunked_data:
                yield chunk

        return cls(chunked_data=async_chunk_gen())

    @property
    def items(self) -> AsyncIterable[_ENTRY_TV]:
        async def item_gen() -> AsyncGenerator[_ENTRY_TV, None]:
            async for chunk in self._chunked_data:
                for item in chunk:
                    yield item

        return item_gen()

    @property
    def chunks(self) -> AsyncIterable[TChunk[_ENTRY_TV]]:
        return self._chunked_data


@attr.s
class AsyncChunkedLimited(AsyncChunkedBase[_ENTRY_TV]):
    _chunked: AsyncChunkedBase = attr.ib()
    _max_count: int = attr.ib()
    limit_exc_to_raise: Type[Exception] = exc.ResultRowCountLimitExceeded

    @property
    def items(self) -> AsyncIterable[_ENTRY_TV]:
        if self._max_count is None:
            return self._chunked.items

        async def async_item_gen() -> AsyncGenerator[_ENTRY_TV, None]:
            cnt = 0
            async for item in self._chunked.items:
                cnt += 1
                if cnt > self._max_count:
                    raise self.limit_exc_to_raise()
                yield item

            LOGGER.info(f"Received {cnt} data rows.")

        return async_item_gen()

    @property
    def chunks(self) -> AsyncIterable[TChunk[_ENTRY_TV]]:
        if self._max_count is None:
            return self._chunked.chunks

        async def async_chunk_gen() -> AsyncGenerator[TChunk[_ENTRY_TV], None]:
            cnt = 0
            async for chunk in self._chunked.chunks:
                cnt += len(chunk)
                if cnt > self._max_count:
                    raise exc.ResultRowCountLimitExceeded()
                yield chunk

            LOGGER.info(f"Received {cnt} data rows.")

        return async_chunk_gen()


@attr.s
class LazyAsyncChunked(AsyncChunkedBase[_ENTRY_TV]):
    """
    ``AsyncChunkedBase`` implementation that executes an initialization callable only when iteration starts.
    """

    _initializer: Callable[[], Awaitable[AsyncChunkedBase[_ENTRY_TV]]] = attr.ib(repr=False)
    _finalizer: Callable[[], Awaitable[Any]] = attr.ib(repr=False)
    _chunked: Optional[Awaitable[AsyncChunkedBase[_ENTRY_TV]]] = attr.ib(init=False, default=None)

    @property
    def items(self) -> AsyncIterable[_ENTRY_TV]:
        async def item_gen() -> AsyncGenerator[_ENTRY_TV, None]:
            if self._chunked is None:
                self._chunked = await self._initializer()  # type: ignore  # TODO: fix
            try:
                async for item in self._chunked.items:  # type: ignore  # TODO: fix
                    yield item
            finally:
                await self._finalizer()

        return item_gen()

    @property
    def chunks(self) -> AsyncIterable[TChunk[_ENTRY_TV]]:
        async def chunk_gen() -> AsyncGenerator[TChunk[_ENTRY_TV], None]:
            if self._chunked is None:
                self._chunked = await self._initializer()  # type: ignore  # TODO: fix
            try:
                async for chunk in self._chunked.chunks:  # type: ignore  # TODO: fix
                    yield chunk
            finally:
                await self._finalizer()

        return chunk_gen()


_CBO_ITEM_TV = TypeVar("_CBO_ITEM_TV")


async def chunkify_by_one(
    items: AsyncIterable[_CBO_ITEM_TV],
) -> AsyncIterable[Sequence[_CBO_ITEM_TV]]:
    """Helper to wrap an iterable into a chunked iterable with one item per chunk"""
    async for item in items:
        yield (item,)

lib/dl_core/dl_core/data_processing/streaming.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.types import (
    TBIDataRow,
    TJSONExt,
)
from dl_core.data_processing.streaming import AsyncChunkedBase


TValuesChunkStream = AsyncChunkedBase[TBIDataRow]  # wrapper of async iterable of chunks of rows
TJSONExtChunkStream = AsyncChunkedBase[TJSONExt]  # wrapper of async iterable of chunks of jsonables

lib/dl_core/dl_core/data_processing/types.py
datalens-tech/datalens-backend
from typing import (
    ClassVar,
    Type,
)

from dl_constants.enums import ProcessorType
from dl_core.aio.web_app_services.data_processing.data_processor import DataProcessorService


class DataProcessorPlugin:
    data_processor_type: ProcessorType
    data_processor_service_cls: ClassVar[Type[DataProcessorService]]

lib/dl_core/dl_core/data_processors/base/plugin.py
datalens-tech/datalens-backend
from typing import Type

import attr

import dl_core as package
from dl_core.data_processors.base.plugin import DataProcessorPlugin
from dl_utils.entrypoints import EntrypointClassManager


_DATA_PROC_SRV_EP_GROUP = f"{package.__name__}.data_processor_plugins"


@attr.s
class DataProcessorPluginEntrypointManager(EntrypointClassManager[DataProcessorPlugin]):
    entrypoint_group_name = attr.ib(init=False, default=_DATA_PROC_SRV_EP_GROUP)


def get_all_data_processor_plugins() -> dict[str, Type[DataProcessorPlugin]]:
    ep_mgr = DataProcessorPluginEntrypointManager()
    return ep_mgr.get_all_ep_classes()

lib/dl_core/dl_core/data_processors/ep_loader.py
datalens-tech/datalens-backend
from typing import Type

from dl_core.aio.web_app_services.data_processing.factory import register_data_processor_service_class
from dl_core.data_processors.base.plugin import DataProcessorPlugin


class DataProcessorPluginRegistrator:
    def register_data_processor_plugin(self, data_processor_plugin_cls: Type[DataProcessorPlugin]) -> None:
        register_data_processor_service_class(
            processor_type=data_processor_plugin_cls.data_processor_type,
            data_processor_srv_cls=data_processor_plugin_cls.data_processor_service_cls,
        )


DATA_PROC_REG = DataProcessorPluginRegistrator()

lib/dl_core/dl_core/data_processors/registrator.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    Any,
    Generic,
    List,
    Optional,
    Type,
    TypeVar,
)

from dl_core.db.elements import SchemaColumn
from dl_core.raw_data_streaming.stream import DataStreamBase


LOGGER = logging.getLogger(__name__)


class DataSink:
    """Object that accepts data in batches and stores it somewhere"""

    def __init__(self, bi_schema: List[SchemaColumn]):
        self._bi_schema = bi_schema

    def initialize(self) -> None:
        """Create table, start external service, etc. - whatever is needed to get it running"""

    def dump_data_stream(self, data_stream: DataStreamBase) -> None:
        """Send stream of entries (list of dicts) to storage"""

    def finalize(self) -> None:
        """Finalize materialization"""

    def cleanup(self) -> None:
        """Cleanup tmp data"""

    def close(self) -> None:
        """Close conn executor and all closable entities"""

    def __enter__(self) -> DataSink:
        self.initialize()
        return self

    def __exit__(self, exc_type: Optional[Type[Exception]], exc_val: Optional[Exception], exc_tb: Any) -> None:
        if exc_type is None:
            self.finalize()
        else:
            LOGGER.exception("Exception occurred in DataSink")
        self.cleanup()
        self.close()


_DATA_STREAM_TV = TypeVar("_DATA_STREAM_TV")


class DataSinkAsync(Generic[_DATA_STREAM_TV], metaclass=abc.ABCMeta):
    """Object that accepts data in batches and stores it somewhere"""

    async def initialize(self) -> None:
        """Create table, start external service, etc. - whatever is needed to get it running"""

    async def dump_data_stream(self, data_stream: _DATA_STREAM_TV) -> None:
        """Send stream of entries (list of dicts) to storage"""

    async def finalize(self) -> None:
        """Finalize materialization"""

    async def cleanup(self) -> None:
        """Cleanup tmp data"""

    async def close(self) -> None:
        """Close conn executor and all closable entities"""

    async def __aenter__(self) -> DataSinkAsync:
        await self.initialize()
        return self

    async def __aexit__(self, exc_type: Optional[Type[Exception]], exc_val: Optional[Exception], exc_tb: Any) -> None:
        if exc_type is None:
            await self.finalize()
        await self.cleanup()
        await self.close()

lib/dl_core/dl_core/data_sink.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_core.data_source.base import (
    DataSource,
    DbInfo,
)
from dl_core.data_source.sql import (
    BaseSQLDataSource,
    SQLDataSource,
    StandardSQLDataSource,
)
from dl_core.data_source.type_mapping import get_data_source_class
from dl_core.data_source.utils import get_parameters_hash
from dl_core.raw_data_streaming.stream import DataStreamBase

from .collection import (
    DataSourceCollection,
    DataSourceCollectionBase,
    DataSourceCollectionFactory,
)


__all__ = (
    "DataSource",
    "DbInfo",
    "SQLDataSource",
    "StandardSQLDataSource",
    "BaseSQLDataSource",
    "get_data_source_class",
    "DataSourceCollectionBase",
    "DataSourceCollection",
    "DataSourceCollectionFactory",
    "DataStreamBase",
    "get_parameters_hash",
)

lib/dl_core/dl_core/data_source/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Iterable,
    NamedTuple,
    Optional,
    Type,
)

import attr

from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
    JoinType,
)
from dl_core.base_models import (
    ConnectionRef,
    SourceFilterSpec,
)
from dl_core.data_source_spec.base import DataSourceSpec
from dl_core.db import (
    IndexInfo,
    SchemaColumn,
    SchemaInfo,
    get_type_transformer,
)
from dl_core.us_connection import get_connection_class
from dl_core.us_connection_base import ConnectionBase


if TYPE_CHECKING:
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase
    from dl_core.db.conversion_base import TypeTransformer
    from dl_core.services_registry.top_level import ServicesRegistry
    from dl_core.us_manager.local_cache import USEntryBuffer


LOGGER = logging.getLogger(__name__)


class DbInfo(NamedTuple):
    version: Optional[str]


@attr.s()
class DataSource(metaclass=abc.ABCMeta):
    """
    Abstraction layer for all data sources:
    - various kinds of DBs (SQL and non-SQL),
    - CSV,
    - HTTP API-based data sources,
    etc.
    """

    supported_join_types: ClassVar[frozenset[JoinType]] = frozenset()
    supports_schema_update: ClassVar[bool] = True
    supports_sample_usage: ClassVar[bool] = False
    supports_limit: ClassVar[bool] = True
    supports_offset: ClassVar[bool] = True
    supports_preview_from_subquery: ClassVar[bool] = True
    store_raw_schema: ClassVar[bool] = True
    preview_enabled: ClassVar[bool] = True
    default_chunk_row_count: ClassVar[int] = 10000
    chunk_size_bytes: ClassVar[int] = 3 * 1024**2

    conn_type: ClassVar[ConnectionType]  # TODO unbind DataSource and Connection classes BI-4083

    # TODO FIX: Remove ASAP
    _id: str = attr.ib(default=None)
    _us_entry_buffer: USEntryBuffer = attr.ib(default=None, eq=False)
    _spec: Optional[DataSourceSpec] = attr.ib(kw_only=True, default=None)

    _connection: Optional[ConnectionBase] = attr.ib(default=None)

    _cache_enabled: bool = attr.ib(default=True)

    @property
    def spec(self) -> DataSourceSpec:
        assert self._spec is not None
        return self._spec

    def _validate_connection(self) -> None:
        if self._connection is not None and self._spec.connection_ref is None:  # type: ignore  # TODO: fix
            # TODO CONSIDER: extraction of connection ref
            pass
        elif self._spec.connection_ref is not None and self._connection is None:  # type: ignore  # TODO: fix
            pass
        else:
            raise ValueError(
                f"Unexpected combination of 'connection' and 'connection_ref':"  # type: ignore  # TODO: fix
                f" {self._connection} and {self._spec.connection_ref}"
            )

        if self._connection is not None:
            self._validate_connection_cls(self._connection)

    def __attrs_post_init__(self):  # type: ignore  # TODO: fix
        self._validate_connection()

    @classmethod
    def get_connection_cls(cls) -> Type[ConnectionBase]:
        """
        Method should return connection class for this data source.
        Class var not used to do not import connection on module level.
        TODO CONSIDER: May be data sources should directly import connections???
        """
        assert cls.conn_type is not None
        return get_connection_class(cls.conn_type)

    def _validate_connection_cls(self, connection: ConnectionBase) -> None:
        expected_cls = self.get_connection_cls()

        if expected_cls != type(connection):
            raise TypeError(f"Unexpected connection class {type(connection)} for {type(self)}, expected {expected_cls}")

    def initialize(self) -> None:
        """All IO-dependent initialization goes here (whatever can later be made awaitable)"""
        pass

    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        raise NotImplementedError

    def get_parameters(self) -> dict:
        return {}

    @property
    def id(self) -> Optional[str]:
        return self._id

    @property
    def connection(self) -> ConnectionBase:
        return self._get_connection()

    def _get_connection(self) -> ConnectionBase:
        if self._connection is None:
            connection_ref = self.spec.connection_ref
            assert connection_ref is not None
            loaded_conn = self._us_entry_buffer.get_entry(connection_ref)
            assert isinstance(loaded_conn, ConnectionBase)
            self._validate_connection_cls(loaded_conn)
            self._connection = loaded_conn

        return self._connection

    @property
    def connection_ref(self) -> ConnectionRef:
        if self.spec.connection_ref is None:
            # TODO CONSIDER: May raise exception if connection has no ref? Or raise on save attempt?
            self.spec.connection_ref = self._connection.conn_ref  # type: ignore  # TODO: fix

        connection_ref = self.spec.connection_ref
        assert connection_ref is not None
        return connection_ref

    @property
    def data_dump_id(self) -> Optional[str]:
        return self.spec.data_dump_id

    @property
    def cache_enabled(self) -> bool:
        conn = self._connection
        if conn is not None:
            if not conn.allow_cache:
                return False
        return self._cache_enabled

    def get_schema_info(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> SchemaInfo:
        """Return structure of the data source. Implementation-specific."""
        raise NotImplementedError()

    @property
    def saved_raw_schema(self) -> Optional[list[SchemaColumn]]:
        if self.spec.raw_schema is not None:
            return self.spec.raw_schema
        return None

    @property
    def saved_index_info_set(self) -> Optional[frozenset[IndexInfo]]:
        return None

    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        """
        Return something that can be used in a ``select_from`` ``SQLAlchemy`` clause for fetching data.
        Optionally assign the source an alias that can be used in the query to refer to it.
        """
        return None

    def source_exists(
        self,
        conn_executor_factory: Callable[[], SyncConnExecutorBase],
        force_refresh: bool = False,
    ) -> bool:
        return True

    def get_db_info(
        self,
        conn_executor_factory: Callable[[], SyncConnExecutorBase],
    ) -> DbInfo:
        return self.get_cached_db_info()

    def get_cached_db_info(self) -> DbInfo:
        return DbInfo(
            version=None,
        )

    @property
    def type_transformer(self) -> TypeTransformer:
        return get_type_transformer(self.conn_type)

    def get_expression_value_range(self, col_name: str) -> tuple[Any, Any]:
        """Return MIN and MAX values for given expression"""
        return None, None

    def supports_join_type(self, join_type: JoinType) -> bool:
        return join_type in self.supported_join_types

    @property
    def default_title(self) -> str:
        raise NotImplementedError

    @property
    def is_configured(self) -> bool:
        return True

    def get_filters(self, service_registry: ServicesRegistry) -> Iterable[SourceFilterSpec]:
        return []

    @property
    def data_export_forbidden(self) -> bool:
        return self.connection.data_export_forbidden


class IncompatibleDataSourceMixin(DataSource):
    @classmethod
    def is_compatible_with_type(cls, source_type: DataSourceType) -> bool:
        return False

lib/dl_core/dl_core/data_source/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Callable,
    ClassVar,
    Optional,
    Union,
)

import attr

from dl_constants.enums import (
    DataSourceCollectionType,
    DataSourceRole,
    DataSourceType,
    JoinType,
    ManagedBy,
)
from dl_core.base_models import (
    ConnectionRef,
    DefaultConnectionRef,
    InternalMaterializationConnectionRef,
)
from dl_core.data_source import (
    base,
    type_mapping,
)
from dl_core.data_source.utils import get_parameters_hash
from dl_core.data_source_spec.base import DataSourceSpec
from dl_core.data_source_spec.collection import (
    DataSourceCollectionSpec,
    DataSourceCollectionSpecBase,
)
from dl_core.db import SchemaColumn
from dl_core.enums import RoleReason
import dl_core.exc as exc


if TYPE_CHECKING:
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase
    from dl_core.us_connection_base import ConnectionBase
    from dl_core.us_manager.local_cache import USEntryBuffer


@attr.s(slots=True)
class RoleResolutionInfo:
    origin: RoleReason = attr.ib(kw_only=True, default=RoleReason.not_needed)
    materialization: RoleReason = attr.ib(kw_only=True, default=RoleReason.not_needed)
    sample: RoleReason = attr.ib(kw_only=True, default=RoleReason.not_needed)
    priorities: list[DataSourceRole] = attr.ib(kw_only=True, factory=list)


LazyDataSourceType = Union[base.DataSource, dict]  # data source instance or param dict


@attr.s
class DataSourceCollectionBase:
    dsrc_coll_type: ClassVar[DataSourceCollectionType]
    _us_entry_buffer: USEntryBuffer = attr.ib(kw_only=True)

    _spec: DataSourceCollectionSpecBase = attr.ib(default=None)

    @property
    def spec(self) -> DataSourceCollectionSpecBase:
        assert isinstance(self._spec, DataSourceCollectionSpecBase)
        return self._spec

    @property
    def id(self) -> str:
        return self.spec.id

    @property
    def valid(self) -> bool:
        return self.spec.valid

    @property
    def title(self) -> Optional[str]:
        return self.spec.title

    @property
    def managed_by(self) -> ManagedBy:
        return self.spec.managed_by

    def get_connection_id(self, role: Optional[DataSourceRole] = None) -> Optional[str]:
        conn_ref = self.get_strict(role=role).connection_ref
        if isinstance(conn_ref, DefaultConnectionRef):
            return conn_ref.conn_id
        elif isinstance(conn_ref, InternalMaterializationConnectionRef):
            return None
        else:
            raise TypeError(f"Unexpected conn_ref class: {type(conn_ref)}")

    @property
    def effective_connection_id(self) -> Optional[str]:
        return self.get_connection_id()

    @property
    def source_type(self) -> DataSourceType:
        return self.get_strict().spec.source_type

    def supports_join_type(
        self,
        join_type: JoinType,
        role: Optional[DataSourceRole] = None,
        for_preview: bool = False,
    ) -> bool:
        if role is None:
            role = self.resolve_role(for_preview=for_preview)
        dsrc = self.get_strict(role=role)
        return join_type in dsrc.supported_join_types

    def _get_connection_by_ref(self, connection_ref: ConnectionRef) -> ConnectionBase:
        from dl_core.us_connection_base import ConnectionBase  # to avoid circular imports

        connection = self._us_entry_buffer.get_entry(connection_ref)
        assert isinstance(connection, ConnectionBase)
        return connection

    def __contains__(self, role: DataSourceRole) -> bool:
        raise NotImplementedError

    def exists(self, role: DataSourceRole) -> bool:
        raise NotImplementedError

    def invalidate(self, role: DataSourceRole) -> None:
        pass

    def are_schemas_compatible(self, role_1: DataSourceRole, role_2: DataSourceRole) -> bool:
        """
        Compare raw schemas of two sources and decide whether they are compatible for data selection
        """

        assert role_1 is DataSourceRole.origin, "First role in comparison must be origin"
        assert role_2 in (
            DataSourceRole.materialization,
            DataSourceRole.sample,
        ), "Second role must be some kind of materialization (materialization or sample)"

        # Origin never uses the materialization connection,
        # so we can safely load it
        src_1 = self.get_opt(role=role_1)
        # materialization/sample sources, however, MUST use it
        # and they must be SQL sources, which means we can get the schema from the spec
        src_2_spec = self._get_spec_for_role(role_2)
        if src_1 is None or src_2_spec is None:
            return False

        raw_schema_1 = sorted(src_1.saved_raw_schema or (), key=lambda col: col.name)
        raw_schema_2 = sorted(src_2_spec.raw_schema or (), key=lambda col: col.name)
        if len(raw_schema_1) != len(raw_schema_2):
            return False

        for col_1, col_2 in zip(raw_schema_1, raw_schema_2):
            if col_1.name != col_2.name or col_1.user_type != col_2.user_type:
                return False

        return True

    def resolve_role_priorities(self, for_preview: bool = False) -> RoleResolutionInfo:
        """
        Decide, which source roles can be used for main source or preview
        taking into account the current configuration and states
        of the collection and individual data sources.
        Return them as a list with items at the beginning having top priority
        """

        resolution_info = RoleResolutionInfo()
        feature_managed = self.managed_by == ManagedBy.feature

        if for_preview:
            sample_spec = self._get_spec_for_role(role=DataSourceRole.sample)
            if sample_spec is None:
                resolution_info.sample = RoleReason.missing_source
            elif not self.get_strict(DataSourceRole.origin).supports_sample_usage:
                resolution_info.sample = RoleReason.not_supported
            elif not sample_spec.is_configured:
                resolution_info.sample = RoleReason.not_configured
            # At this point sample is materialized and ready for use.
            elif not feature_managed and not self.are_schemas_compatible(DataSourceRole.origin, DataSourceRole.sample):
                resolution_info.sample = RoleReason.schema_mismatch
            else:
                # structure is the same, so we can use the existing sample data source
                resolution_info.sample = RoleReason.selected
                resolution_info.priorities.append(DataSourceRole.sample)

            # either there is no sample that is ready for usage or source's structure has changed in origin

        # all other cases (including main source for direct mode)
        if feature_managed:
            resolution_info.origin = RoleReason.forbidden_for_features
        else:
            resolution_info.origin = RoleReason.selected
            resolution_info.priorities.append(DataSourceRole.origin)

        return resolution_info

    def resolve_role(self, for_preview: bool = False) -> DataSourceRole:
        """
        Decide, which source role to use for main source or preview
        taking into account the current configuration and states
        of the collection and individual data sources.
        """
        role_priorities = self.resolve_role_priorities(for_preview=for_preview).priorities
        if not role_priorities:
            raise exc.NoCommonRoleError()
        return role_priorities[0]

    def get_opt(self, role: Optional[DataSourceRole] = None, for_preview: bool = False) -> Optional["base.DataSource"]:
        raise NotImplementedError

    def get_strict(self, role: Optional[DataSourceRole] = None, for_preview: bool = False) -> base.DataSource:
        dsrc = self.get_opt(role=role, for_preview=for_preview)
        assert dsrc is not None
        return dsrc

    def _get_spec_for_role(self, role: DataSourceRole) -> Optional[DataSourceSpec]:
        raise NotImplementedError

    def get_cached_raw_schema(
        self,
        role: Optional[DataSourceRole] = None,
        for_preview: bool = False,
    ) -> Optional[list[SchemaColumn]]:
        dsrc = self.get_strict(role=role, for_preview=for_preview)
        raw_schema: Optional[list[SchemaColumn]] = dsrc.saved_raw_schema or []
        if not raw_schema and role != DataSourceRole.origin:
            # we need to at least have BI types, so deduce them from the original raw schema
            origin_dsrc = self.get_strict(role=DataSourceRole.origin)
            origin_columns = origin_dsrc.saved_raw_schema
            if origin_columns is None:
                return None

            # "import" raw schema from origin dsrc to current one
            tt = origin_dsrc.type_transformer
            raw_schema = [
                col.clone(  # type: ignore  # TODO: fix
                    # convert data type to match data source type
                    user_type=col.user_type
                    or tt.type_native_to_user(native_t=col.native_type),
                )
                for col in origin_columns
            ]

        raw_schema = self._patch_raw_schema(raw_schema)
        return raw_schema

    def get_raw_schema(
        self,
        conn_executor_factory: Callable[[], SyncConnExecutorBase],
        role: Optional[DataSourceRole] = None,
        for_preview: bool = False,
    ) -> Optional[list[SchemaColumn]]:
        """
        Return raw schema of data source.
        Since the underlying data source might be a reference to an external one,
        its ID may be different to ``self._id``, so all columns need to be patched
        with the correct ID (see end of method).

        :param from_cache: return raw_schema saved in config
        :param role: which data source in collection
        :param for_preview: preview mode
        :return:
        """

        dsrc = self.get_strict(role=role, for_preview=for_preview)
        raw_schema: Optional[list[SchemaColumn]] = dsrc.get_schema_info(
            conn_executor_factory=conn_executor_factory
        ).schema
        raw_schema = self._patch_raw_schema(raw_schema)
        return raw_schema

    def _patch_raw_schema(self, raw_schema: Optional[list[SchemaColumn]]) -> Optional[list[SchemaColumn]]:
        """Patch all raw_schema columns with own source_id"""
        if raw_schema is None:
            return None

        return [col.clone(source_id=self.id) for col in raw_schema]

    def get_param_hash(self) -> str:
        raise NotImplementedError


@attr.s
class DataSourceCollection(DataSourceCollectionBase):
    dsrc_coll_type = DataSourceCollectionType.collection

    _loaded_sources: dict[DataSourceRole, Optional[base.DataSource]] = attr.ib(init=False, factory=dict)

    @property
    def spec(self) -> DataSourceCollectionSpec:
        assert isinstance(self._spec, DataSourceCollectionSpec)
        return self._spec

    def _initialize_dsrc(self, dsrc_spec: DataSourceSpec) -> "base.DataSource":
        dsrc_cls = type_mapping.get_data_source_class(dsrc_spec.source_type)
        dsrc = dsrc_cls(id=self.id, us_entry_buffer=self._us_entry_buffer, spec=dsrc_spec)
        return dsrc

    def _get_spec_for_role(self, role: DataSourceRole) -> Optional[DataSourceSpec]:
        if role == DataSourceRole.origin:
            return self.spec.origin
        if role == DataSourceRole.materialization:
            return self.spec.materialization
        if role == DataSourceRole.sample:
            return self.spec.sample
        raise ValueError(f"Invalid role: {role}")

    def __contains__(self, role: DataSourceRole) -> bool:
        return self._get_spec_for_role(role) is not None

    def get_opt(self, role: Optional[DataSourceRole] = None, for_preview: bool = False) -> Optional["base.DataSource"]:
        """Return data source for role. Initialize it if it isn't initialized yet"""
        if role is None:
            role = self.resolve_role(for_preview=for_preview)
        assert role is not None
        if role not in self:
            return None
        if role not in self._loaded_sources:
            dsrc_spec = self._get_spec_for_role(role=role)
            if dsrc_spec is None:
                self._loaded_sources[role] = None
            else:
                self._loaded_sources[role] = self._initialize_dsrc(dsrc_spec=dsrc_spec)
        return self._loaded_sources[role]

    def exists(self, role: DataSourceRole) -> bool:
        return self.get_opt(role) is not None

    def invalidate(self, role: DataSourceRole) -> None:
        try:
            del self._loaded_sources[role]
        except KeyError:
            pass

    def get_param_hash(self) -> str:
        role = DataSourceRole.origin
        dsrc = self.get_strict(role)
        return get_parameters_hash(
            source_type=dsrc.spec.source_type,
            connection_id=self.get_connection_id(role=role),
            **dsrc.get_parameters(),
        )


@attr.s
class DataSourceCollectionFactory:  # TODO: Move to service_registry
    _us_entry_buffer: USEntryBuffer = attr.ib(default=None)

    def get_data_source_collection(self, spec: DataSourceCollectionSpecBase) -> DataSourceCollectionBase:
        if isinstance(spec, DataSourceCollectionSpec):
            return DataSourceCollection(us_entry_buffer=self._us_entry_buffer, spec=spec)
        raise TypeError(f"Invalid spec type: {spec})")

lib/dl_core/dl_core/data_source/collection.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from functools import wraps
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Optional,
    Type,
)

import attr
import sqlalchemy as sa
from sqlalchemy.engine.default import DefaultDialect

from dl_constants.enums import JoinType
from dl_core import exc
from dl_core.connection_models import (
    DBIdent,
    SATextTableDefinition,
    TableDefinition,
    TableIdent,
)
from dl_core.connectors.base.query_compiler import QueryCompiler
from dl_core.data_source.base import (
    DataSource,
    DbInfo,
    IncompatibleDataSourceMixin,
)
from dl_core.data_source_spec.sql import (
    DbSQLDataSourceSpec,
    IndexedSQLDataSourceSpec,
    SchemaSQLDataSourceSpec,
    SQLDataSourceSpecBase,
    StandardSchemaSQLDataSourceSpec,
    StandardSQLDataSourceSpec,
    SubselectDataSourceSpec,
    TableSQLDataSourceSpec,
)
from dl_core.db import (
    IndexInfo,
    SchemaInfo,
)
from dl_core.utils import sa_plain_text


if TYPE_CHECKING:
    from dl_core.connection_executors.async_base import AsyncConnExecutorBase
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase


LOGGER = logging.getLogger(__name__)


def require_table_name(func):  # type: ignore  # TODO: fix
    @wraps(func)
    def wrapper(self, *args, **kwargs):  # type: ignore  # TODO: fix
        if not self.table_name:
            raise exc.TableNameNotConfiguredError

        return func(self, *args, **kwargs)

    return wrapper


@attr.s
class BaseSQLDataSource(DataSource):
    """Data source for SQL database"""

    default_raw_data_batch_size = 10000
    supported_join_types = frozenset(
        {
            JoinType.inner,
            JoinType.left,
        }
    )
    compiler_cls: Type[QueryCompiler] = QueryCompiler

    # Instance attrs
    _exists: Optional[bool] = attr.ib(default=None, init=False, eq=False, hash=False)

    @property
    def spec(self) -> SQLDataSourceSpecBase:
        assert isinstance(self._spec, SQLDataSourceSpecBase)
        return self._spec

    @property
    def db_version(self) -> Optional[str]:
        return self.spec.db_version

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            db_version=self.db_version,
        )

    def get_connect_args(self) -> dict:
        """Customize connection's ``connect_args``"""
        return {}

    def _postprocess_raw_schema_from_db(self, schema_info: SchemaInfo) -> SchemaInfo:
        columns = tuple(col._replace(source_id=self.id) for col in schema_info.schema)

        if not columns:
            raise exc.FailedToLoadSchema()

        return schema_info.clone(schema=columns)

    @property
    def default_title(self) -> str:
        return "SQL"

    def _check_existence(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> bool:
        return True

    async def _check_existence_async(self, conn_executor_factory: Callable[[], AsyncConnExecutorBase]) -> bool:
        return True

    def source_exists(
        self,
        conn_executor_factory: Callable[[], SyncConnExecutorBase],
        force_refresh: bool = False,
    ) -> bool:
        if self._exists is None or force_refresh:
            self._exists = self._check_existence(conn_executor_factory=conn_executor_factory)
        return self._exists

    def get_dialect(self) -> DefaultDialect:
        return self.connection.get_dialect()

    def quote(self, value) -> sa.sql.quoted_name:  # type: ignore  # TODO: fix  # subclass of str
        return self.connection.quote(value)

    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        raise NotImplementedError()

    def get_table_definition(self) -> TableDefinition:
        """Should returns TableDefinition which will be used to fetch schema info from DB"""
        raise NotImplementedError()

    def get_schema_info(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> SchemaInfo:
        table_def = self.get_table_definition()
        conn_executor = conn_executor_factory()
        schema_info = conn_executor.get_table_schema_info(table_def)
        return self._postprocess_raw_schema_from_db(schema_info)

    def get_query_compiler(self) -> QueryCompiler:
        return self.compiler_cls(
            dialect=self.get_dialect(),
        )

    def _get_db_version(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> Optional[str]:
        conn_executor = conn_executor_factory()
        return conn_executor.get_db_version(
            db_ident=DBIdent(db_name=None),
        )

    def get_db_info(
        self,
        conn_executor_factory: Callable[[], SyncConnExecutorBase],
    ) -> DbInfo:
        # self.connection
        LOGGER.info("Fetching version info from the database")
        db_version = self._get_db_version(conn_executor_factory=conn_executor_factory)
        return DbInfo(
            version=db_version,
        )

    def get_cached_db_info(self) -> DbInfo:
        return DbInfo(
            version=self.db_version,
        )


@attr.s
class SubselectDataSource(BaseSQLDataSource):
    """
    `select … from (select …) …` source (for user-input inner-select).
    """

    _subsql: Optional[str] = attr.ib(default=None)

    @property
    def spec(self) -> SubselectDataSourceSpec:
        assert isinstance(self._spec, SubselectDataSourceSpec)
        return self._spec

    @property
    def subsql(self) -> Optional[str]:
        return self.spec.subsql

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            subsql=self.subsql,
        )

    _subquery_alias_joiner = " AS "
    _subquery_auto_alias = "source"

    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        if not self.connection.is_subselect_allowed:
            raise exc.SubselectNotAllowed()

        subsql = self.subsql
        if not subsql:
            raise exc.TableNameNotConfiguredError

        from_sql = "(\n{}\n)".format(subsql)

        # e.g. PG doesn't allow unaliased subqueries at all.
        if alias is None:
            alias = self._subquery_auto_alias

        if alias is not None:
            from_sql = "{}{}{}".format(from_sql, self._subquery_alias_joiner, self.quote(alias))
        return sa_plain_text(from_sql)

    @property
    def default_title(self) -> str:
        return "(select ...)"

    def get_table_definition(self) -> TableDefinition:
        subselect = self.get_sql_source(alias="source")
        return SATextTableDefinition(text=subselect)


@attr.s
class SQLDataSource(abc.ABC, BaseSQLDataSource):
    """Data source for SQL database"""

    @property
    @abc.abstractmethod
    def db_name(self) -> Optional[str]:
        raise NotImplementedError

    @property
    @abc.abstractmethod
    def table_name(self) -> Optional[str]:
        raise NotImplementedError

    @property
    def default_title(self) -> str:
        return self.table_name  # type: ignore  # TODO: fix

    def _get_db_version(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> Optional[str]:
        conn_executor = conn_executor_factory()
        return conn_executor.get_db_version(
            db_ident=DBIdent(db_name=self.db_name),
        )

    def _check_existence(self, conn_executor_factory: Callable[[], SyncConnExecutorBase]) -> bool:
        table_def = self.get_table_definition()
        conn_executor = conn_executor_factory()
        return conn_executor.is_table_exists(table_def)

    async def _check_existence_async(self, conn_executor_factory: Callable[[], AsyncConnExecutorBase]) -> bool:
        table_def = self.get_table_definition()
        async_conn_executor = conn_executor_factory()
        return await async_conn_executor.is_table_exists(table_def)

    def source_exists(
        self,
        conn_executor_factory: Callable[[], SyncConnExecutorBase],
        force_refresh: bool = False,
    ) -> bool:
        if not self.table_name:
            return False

        return super().source_exists(conn_executor_factory=conn_executor_factory, force_refresh=force_refresh)

    @require_table_name
    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        q = self.quote
        alias_str = "" if alias is None else f" AS {q(alias)}"
        return sa_plain_text(f"{q(self.db_name)}.{q(self.table_name)}{alias_str}")

    @require_table_name
    def get_table_definition(self) -> TableDefinition:
        table_name = self.table_name
        assert table_name is not None
        return TableIdent(db_name=self.db_name, schema_name=None, table_name=table_name)

    @property
    def is_configured(self) -> bool:
        return self.db_name is not None and self.table_name is not None


@attr.s
class DbSQLDataSourceMixin(BaseSQLDataSource):
    """SQL data source with db_name"""

    @property
    def db_name(self) -> Optional[str]:
        assert isinstance(self._spec, DbSQLDataSourceSpec)
        return self._spec.db_name or getattr(self.connection, "db_name", None)

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            db_name=self.db_name,
        )


@attr.s
class TableSQLDataSourceMixin(BaseSQLDataSource):
    """SQL data source with table_name"""

    @property
    def table_name(self) -> Optional[str]:
        assert isinstance(self._spec, TableSQLDataSourceSpec)
        return self._spec.table_name or getattr(self.connection, "table_name", None)

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            table_name=self.table_name,
        )


@attr.s
class IndexedSQLDataSourceMixin(BaseSQLDataSource):
    """SQL data source with index info"""

    @property
    def saved_index_info_set(self) -> Optional[frozenset[IndexInfo]]:
        assert isinstance(self._spec, IndexedSQLDataSourceSpec)
        return self._spec.index_info_set


@attr.s
class StandardSQLDataSource(
    DbSQLDataSourceMixin,
    TableSQLDataSourceMixin,
    IndexedSQLDataSourceMixin,
    SQLDataSource,  # <-- the main logic is in this one
):
    """SQL data source with db_name, table_name and indexes"""

    @property
    def spec(self) -> StandardSQLDataSourceSpec:
        assert isinstance(self._spec, StandardSQLDataSourceSpec)
        return self._spec

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            db_name=self.spec.db_name,  # FIXME: Remove (db_name is part of the data source only in CH)
            table_name=self.spec.table_name,
        )


class PseudoSQLDataSource(StandardSQLDataSource, IncompatibleDataSourceMixin):
    supported_join_types: ClassVar[frozenset[JoinType]] = frozenset()
    supports_schema_update: ClassVar[bool] = False

    @require_table_name
    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        # ignore alias
        return sa.table(self.table_name)


class SchemaSQLDataSourceMixin(BaseSQLDataSource):
    @property
    def schema_name(self) -> Optional[str]:
        assert isinstance(self._spec, SchemaSQLDataSourceSpec)
        # TODO FIX: DO NOT DO THIS IN PROPERTY!!!!!
        #  USE METHOD get_effective schema name or dump from connection on initial data source creation
        return self._spec.schema_name if self._spec.schema_name else self.get_connection_cls().default_schema_name  # type: ignore  # TODO: fix


@attr.s
class StandardSchemaSQLDataSource(StandardSQLDataSource, SchemaSQLDataSourceMixin):
    _schema_name: Optional[str] = attr.ib(default=None)

    @property
    def spec(self) -> StandardSchemaSQLDataSourceSpec:
        assert isinstance(self._spec, StandardSchemaSQLDataSourceSpec)
        return self._spec

    def get_parameters(self) -> dict:
        return dict(
            super().get_parameters(),
            schema_name=self.schema_name,
        )

    @require_table_name
    def get_sql_source(self, alias: Optional[str] = None) -> Any:
        if not self.schema_name:
            return super().get_sql_source(alias=alias)
        q = self.quote
        alias_str = "" if alias is None else f" AS {q(alias)}"
        schema_str = "" if self.schema_name is None else f"{q(self.schema_name)}."
        return sa_plain_text(f"{q(self.db_name)}.{schema_str}{q(self.table_name)}{alias_str}")

    @require_table_name
    def get_table_definition(self) -> TableDefinition:
        table_name = self.table_name
        assert table_name is not None
        return TableIdent(db_name=self.db_name, schema_name=self.schema_name, table_name=table_name)

lib/dl_core/dl_core/data_source/sql.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Collection,
    Type,
)


if TYPE_CHECKING:
    from dl_constants.enums import DataSourceType
    from dl_core.data_source.base import DataSource


_DSRC_TYPES: dict[DataSourceType, Type[DataSource]] = {}


def list_registered_source_types() -> Collection[DataSourceType]:
    return set(_DSRC_TYPES)


def get_data_source_class(ds_type: DataSourceType) -> Type[DataSource]:
    """Return ``DataSource`` subclass to be used for given dataset type."""
    return _DSRC_TYPES[ds_type]


def register_data_source_class(source_type: DataSourceType, source_cls: Type[DataSource]) -> None:
    """Register ``DataSource`` subclass in the mapping."""
    _DSRC_TYPES[source_type] = source_cls

lib/dl_core/dl_core/data_source/type_mapping.py
datalens-tech/datalens-backend
from __future__ import annotations

import itertools
from typing import (
    Any,
    Optional,
)

import xxhash

from dl_constants.enums import DataSourceType


_IGNORE_IN_HASH = frozenset(("db_version",))


def get_parameters_hash(source_type: DataSourceType, connection_id: Optional[str], **parameters: Any) -> str:
    data = (
        source_type,
        connection_id,
        *itertools.chain.from_iterable(
            (key, value)
            for key, value in sorted(parameters.items())
            if value is not None and key not in _IGNORE_IN_HASH
        ),
    )
    return xxhash.xxh64(str(data), seed=0).hexdigest()

lib/dl_core/dl_core/data_source/utils.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
)

import attr

from dl_constants.enums import DataSourceType
from dl_core.data_source_spec.base import DataSourceSpec
from dl_core.data_source_spec.type_mapping import get_data_source_spec_class


def make_spec_from_dict(source_type: DataSourceType, data: Dict[str, Any]) -> DataSourceSpec:
    dsrc_spec_cls = get_data_source_spec_class(ds_type=source_type)
    field_names = {
        field.name.lstrip("_")
        for field in attr.fields(dsrc_spec_cls)
        # Allow only fields with init=True
        if field.init
    }

    filtered_data = {k: v for k, v in data.items() if k in field_names}
    filtered_data["source_type"] = source_type
    return dsrc_spec_cls(**filtered_data)


def update_spec_from_dict(
    source_type: DataSourceType, data: Dict[str, Any], old_spec: DataSourceSpec
) -> DataSourceSpec:
    merged_data = {
        # collect old attributes
        field.name.lstrip("_"): getattr(old_spec, field.name)
        for field in attr.fields(type(old_spec))
    }
    merged_data.update(data)
    return make_spec_from_dict(source_type=source_type, data=merged_data)

lib/dl_core/dl_core/data_source_merge_tools.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

import attr

from dl_constants.enums import DataSourceType
from dl_core.base_models import ConnectionRef
from dl_core.db.elements import SchemaColumn


@attr.s
class DataSourceSpec:
    source_type: DataSourceType = attr.ib(kw_only=True)
    connection_ref: Optional[ConnectionRef] = attr.ib(kw_only=True, default=None)
    raw_schema: Optional[list[SchemaColumn]] = attr.ib(kw_only=True, default=None)
    data_dump_id: Optional[str] = attr.ib(kw_only=True, default=None)

    @property
    def is_configured(self) -> bool:
        return True

    @classmethod
    def get_public_fields(cls) -> list[str]:
        return [field.name for field in attr.fields(cls) if not field.name.startswith("_")]

lib/dl_core/dl_core/data_source_spec/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    ClassVar,
    Optional,
)

import attr

from dl_constants.enums import (
    DataSourceCollectionType,
    DataSourceRole,
    ManagedBy,
)
from dl_core.base_models import DefaultConnectionRef
from dl_core.data_source_spec.base import DataSourceSpec


@attr.s
class DataSourceCollectionSpecBase(abc.ABC):
    dsrc_coll_type: ClassVar[DataSourceCollectionType]

    id: str = attr.ib(kw_only=True)
    title: Optional[str] = attr.ib(kw_only=True, default=None)
    managed_by: ManagedBy = attr.ib(kw_only=True, default=None)
    valid: bool = attr.ib(kw_only=True, default=True)

    def __attrs_post_init__(self) -> None:
        self.managed_by = self.managed_by or ManagedBy.user

    @abc.abstractmethod
    def collect_links(self) -> dict[str, str]:
        raise NotImplementedError


@attr.s
class DataSourceCollectionSpec(DataSourceCollectionSpecBase):  # noqa
    dsrc_coll_type = DataSourceCollectionType.collection

    origin: Optional[DataSourceSpec] = attr.ib(kw_only=True, default=None)
    materialization: Optional[DataSourceSpec] = attr.ib(kw_only=True, default=None)
    sample: Optional[DataSourceSpec] = attr.ib(kw_only=True, default=None)

    def get_for_role(self, role: DataSourceRole) -> Optional[DataSourceSpec]:
        return getattr(self, role.name)

    def set_for_role(self, role: DataSourceRole, value: Optional[DataSourceSpec]) -> None:
        setattr(self, role.name, value)

    def has_role(self, role: DataSourceRole) -> bool:
        return bool(getattr(self, role.name, None))

    def collect_links(self) -> dict[str, str]:
        assert self.origin is not None
        result: dict[str, str] = {}
        connection_ref = self.origin.connection_ref
        if isinstance(connection_ref, DefaultConnectionRef):
            result[self.id] = connection_ref.conn_id
        return result

lib/dl_core/dl_core/data_source_spec/collection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    FrozenSet,
    Optional,
)

import attr

from dl_core.data_source_spec.base import DataSourceSpec
from dl_core.db import IndexInfo
from dl_utils.utils import get_type_full_name


@attr.s
class SQLDataSourceSpecBase(DataSourceSpec):
    db_version: Optional[str] = attr.ib(kw_only=True, default=None)


@attr.s
class SubselectDataSourceSpec(SQLDataSourceSpecBase):
    subsql: Optional[str] = attr.ib(kw_only=True, default=None)


@attr.s
class DbSQLDataSourceSpec(SQLDataSourceSpecBase):
    db_name: Optional[str] = attr.ib(kw_only=True, default=None)


@attr.s
class TableSQLDataSourceSpec(SQLDataSourceSpecBase):
    table_name: Optional[str] = attr.ib(kw_only=True, default=None)

    @property
    def is_configured(self) -> bool:
        return self.table_name is not None


@attr.s
class IndexedSQLDataSourceSpec(SQLDataSourceSpecBase):
    index_info_set: Optional[FrozenSet[IndexInfo]] = attr.ib(kw_only=True, default=None)

    @index_info_set.validator
    def check_indexes(self, attribute: Any, value: Any) -> None:
        if value is None or isinstance(value, frozenset):
            pass
        else:
            raise TypeError(f"Indexes must be a frozen set or non, not {get_type_full_name(type(value))!r}")


@attr.s
class StandardSQLDataSourceSpec(
    DbSQLDataSourceSpec,
    TableSQLDataSourceSpec,
    IndexedSQLDataSourceSpec,
):
    pass


@attr.s
class SchemaSQLDataSourceSpec(SQLDataSourceSpecBase):
    schema_name: Optional[str] = attr.ib(kw_only=True, default=None)


@attr.s
class StandardSchemaSQLDataSourceSpec(
    StandardSQLDataSourceSpec,
    SchemaSQLDataSourceSpec,
):
    pass

lib/dl_core/dl_core/data_source_spec/sql.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Type

from dl_constants.enums import DataSourceType
from dl_core.data_source_spec.base import DataSourceSpec


_DSRC_SPEC_CLASSES: dict[DataSourceType, Type[DataSourceSpec]] = {}


def get_data_source_spec_class(ds_type: DataSourceType) -> Type[DataSourceSpec]:
    """Return ``DataSourceSpec`` subclass to be used for given dataset type."""
    return _DSRC_SPEC_CLASSES[ds_type]


def register_data_source_spec_class(source_type: DataSourceType, spec_cls: Type[DataSourceSpec]) -> None:
    """Register ``DataSourceSpec`` subclass in the mapping."""
    _DSRC_SPEC_CLASSES[source_type] = spec_cls

lib/dl_core/dl_core/data_source_spec/type_mapping.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import lru_cache
from itertools import chain
import logging
from typing import (
    TYPE_CHECKING,
    Collection,
    Dict,
    FrozenSet,
    Optional,
    Set,
    Tuple,
)

import attr

from dl_constants.enums import (
    ConnectionType,
    DataSourceRole,
    DataSourceType,
    JoinType,
    SourceBackendType,
)
from dl_core.backend_types import get_backend_type
from dl_core.components.accessor import DatasetComponentAccessor
import dl_core.data_source as data_source
from dl_core.data_source.type_mapping import list_registered_source_types
import dl_core.exc as exc
from dl_core.us_connection import CONNECTION_TYPES
from dl_core.us_dataset import Dataset


if TYPE_CHECKING:
    from dl_core.data_source.collection import (
        DataSourceCollectionBase,
        DataSourceCollectionFactory,
    )


LOGGER = logging.getLogger(__name__)


@lru_cache(maxsize=200)
def get_compatible_source_types(source_type: DataSourceType) -> FrozenSet[DataSourceType]:
    """Return frozen set of data source types compatible with ``ds_type``"""

    raw_comp_types = frozenset(list_registered_source_types())

    dsrc_cls = data_source.get_data_source_class(ds_type=source_type)
    compat_types = {
        comp_source_type
        for comp_source_type in raw_comp_types
        if dsrc_cls.is_compatible_with_type(source_type=comp_source_type)
    }
    return frozenset(compat_types)


_SOURCE_CONNECTION_COMPATIBILITY: Dict[DataSourceType, FrozenSet[ConnectionType]] = {}


def _populate_compatibility_map():  # type: ignore  # TODO: fix
    for conn_type, conn_cls in CONNECTION_TYPES.items():
        for dsrc_type in conn_cls.get_provided_source_types():
            _SOURCE_CONNECTION_COMPATIBILITY[dsrc_type] = _SOURCE_CONNECTION_COMPATIBILITY.get(
                dsrc_type, frozenset()
            ) | frozenset([conn_type])


@lru_cache(maxsize=100)
def get_conn_types_compatible_with_src_types(source_types: FrozenSet[DataSourceType]) -> FrozenSet[ConnectionType]:
    if not _SOURCE_CONNECTION_COMPATIBILITY:
        _populate_compatibility_map()
    assert _SOURCE_CONNECTION_COMPATIBILITY
    return frozenset(
        chain.from_iterable(_SOURCE_CONNECTION_COMPATIBILITY.get(dsrc_type, frozenset()) for dsrc_type in source_types)
    )


@attr.s
class DatasetCapabilities:
    _dataset: Dataset = attr.ib(kw_only=True)
    _dsrc_coll_factory: DataSourceCollectionFactory = attr.ib(kw_only=True)
    _ds_accessor: DatasetComponentAccessor = attr.ib(init=False)

    @_ds_accessor.default
    def _make_ds_accessor(self) -> DatasetComponentAccessor:
        return DatasetComponentAccessor(dataset=self._dataset)

    def _get_data_source_strict(self, source_id: str, role: DataSourceRole) -> data_source.DataSource:
        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
        dsrc_coll = self._dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
        dsrc = dsrc_coll.get_strict(role=role)
        return dsrc

    def _get_first_dsrc_collection(
        self,
        ignore_source_ids: Optional[Collection[str]] = None,
    ) -> Optional["data_source.DataSourceCollectionBase"]:
        source_id = self._dataset.get_single_data_source_id(ignore_source_ids=ignore_source_ids)
        if source_id is not None:
            dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_opt(source_id)
            if dsrc_coll_spec is not None:
                return self._dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)

        return None

    def get_effective_connection_id(  # type: ignore  # TODO: fix
        self,
        ignore_source_ids: Optional[Collection[str]] = None,
    ) -> Optional[str]:
        dsrc_coll = self._get_first_dsrc_collection(ignore_source_ids=ignore_source_ids)
        if dsrc_coll is not None:
            return dsrc_coll.effective_connection_id

    def get_supported_join_types(
        self,
        ignore_source_ids: Optional[Collection[str]] = None,
    ) -> Set[JoinType]:
        ignore_source_ids = ignore_source_ids or ()
        role = self.resolve_source_role(ignore_source_ids=ignore_source_ids)
        result = set(jt for jt in JoinType)
        for source_id in self._ds_accessor.get_data_source_id_list():
            if source_id in ignore_source_ids:
                continue
            result &= self._get_data_source_strict(source_id=source_id, role=role).supported_join_types
        return result

    def source_can_be_added(
        self,
        connection_id: Optional[str],
        created_from: DataSourceType,
        ignore_source_ids: Optional[Collection[str]] = None,
    ) -> bool:
        """
        Check whether a data source with given connection and type can be added to the dataset.

        ``ignore_source_ids`` can be specified to exclude certain exising sources from all checks.
        This can be useful for checking whether a source can be replaced with a new one
        """

        ignore_source_ids = ignore_source_ids or ()
        dsrc_cls = data_source.get_data_source_class(ds_type=created_from)
        existing_source_id = self._dataset.get_single_data_source_id(ignore_source_ids=ignore_source_ids)
        effective_connection_id = self.get_effective_connection_id(ignore_source_ids=ignore_source_ids)
        supported_join_types = self.get_supported_join_types(ignore_source_ids=ignore_source_ids)

        if existing_source_id is None:  # dataset is empty
            return True

        if not supported_join_types & dsrc_cls.supported_join_types:
            return False

        return connection_id == effective_connection_id

    def get_compatible_source_types(
        self,
        ignore_source_ids: Optional[Collection[str]] = None,
    ) -> FrozenSet[DataSourceType]:
        """Return a frozen set of source types compatible with dataset's current state"""

        ignore_source_ids = ignore_source_ids or ()
        # All available source types
        comp_types = frozenset(list_registered_source_types())
        # Now check them for compatibility with existing sources
        for source_id in self._ds_accessor.get_data_source_id_list():
            if source_id in ignore_source_ids:
                continue
            dsrc = self._get_data_source_strict(source_id=source_id, role=DataSourceRole.origin)
            dsrc_type = dsrc.spec.source_type
            comp_types &= get_compatible_source_types(dsrc_type)
        return comp_types

    def _get_data_source_collections(self) -> dict[str, DataSourceCollectionBase]:
        result = {}
        for dsrc_coll_spec in self._ds_accessor.get_data_source_coll_spec_list():
            result[dsrc_coll_spec.id] = self._dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
        return result

    def get_compatible_connection_types(
        self,
        ignore_connection_ids: Optional[Collection[str]] = None,
    ) -> FrozenSet[ConnectionType]:
        """Return a frozen set of connection types compatible with dataset's current state"""

        ignore_connection_ids = ignore_connection_ids or ()
        ignore_source_ids = {
            # all sources that use the ignored connections
            dsrc_coll.id
            for dsrc_coll in self._get_data_source_collections().values()
            if dsrc_coll.get_connection_id(role=DataSourceRole.origin) in ignore_connection_ids
        }

        # No two connections can be used together in a single non-materialized dataset
        existing_source_ids = {
            source_id for source_id in self._ds_accessor.get_data_source_id_list() if source_id not in ignore_source_ids
        }
        if existing_source_ids:
            # A connection is already present in the dataset. Can't add a second one, no matter what type.
            return frozenset()

        # Dataset is empty (taking into account ignore_source_ids)
        return get_conn_types_compatible_with_src_types(frozenset(list_registered_source_types()))

    def supports_offset(self, role: DataSourceRole) -> bool:
        for _source_id, dsrc_coll in self._get_data_source_collections().items():
            dsrc = dsrc_coll.get_strict(role=role)
            if not dsrc.supports_offset:
                return False
        return True

    def get_backend_type(self, role: DataSourceRole) -> SourceBackendType:
        source_id = self._dataset.get_single_data_source_id()
        assert source_id is not None
        dsrc = self._get_data_source_strict(source_id=source_id, role=role)
        assert dsrc is not None
        return get_backend_type(conn_type=dsrc.conn_type)

    def resolve_source_role(
        self,
        for_preview: bool = False,
        ignore_source_ids: Optional[Collection[str]] = None,
        log_reasons: bool = False,
    ) -> DataSourceRole:
        """
        Resolve role that should be used for all source collections
        """
        collections = list(self._get_data_source_collections().values())
        if not collections:
            return DataSourceRole.origin

        # dict: {role: sum_of_priority_values}
        common_priorities = {role: 0 for role in DataSourceRole}
        for coll in collections:
            if ignore_source_ids is not None and coll.id in ignore_source_ids:
                continue
            role_resolution_info = coll.resolve_role_priorities(for_preview=for_preview)
            coll_priorities = role_resolution_info.priorities

            if log_reasons:
                LOGGER.info(
                    f"Role resolution_info for data source {coll.id}",
                    extra=dict(
                        role_resolution_info=dict(
                            source_id=coll.id,
                            source_type=coll.source_type.name,
                            origin=role_resolution_info.origin.name,
                            sample=role_resolution_info.sample.name,
                            materialization=role_resolution_info.materialization.name,
                        )
                    ),
                )

            for role in common_priorities.copy():
                if role not in coll_priorities:
                    del common_priorities[role]
                else:
                    common_priorities[role] += coll_priorities.index(role)

        if not common_priorities:
            raise exc.NoCommonRoleError()

        def role_sorting_key(role_priority_pair: Tuple[DataSourceRole, int]) -> Tuple[int, str]:
            # first priority, then name (so that it's deterministic for matching priorities)
            return role_priority_pair[1], role_priority_pair[0].name

        sorted_role_priorities = sorted(common_priorities.items(), key=role_sorting_key)
        # at this point we know it's not empty
        role, _ = sorted_role_priorities[0]
        return role

    def supports_preview(self) -> bool:
        """
        Check whether dataset supports preview.
        Used in options to tell the UI (or some other client) whether preview should/could be shown.
        """
        for _source_id, dsrc_coll in self._get_data_source_collections().items():
            if not dsrc_coll.get_strict().preview_enabled:
                return False

        return True

lib/dl_core/dl_core/dataset_capabilities.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_core.db.conversion_base import (
    TypeTransformer,
    get_type_transformer,
)
from dl_core.db.elements import (
    IndexInfo,
    SchemaColumn,
    SchemaInfo,
)
from dl_core.db.sa_types import make_sa_type
from dl_core.db.schema import (
    SAMPLE_ID_COLUMN_NAME,
    are_raw_schemas_same,
)


__all__ = (
    "get_type_transformer",
    "TypeTransformer",
    "SchemaColumn",
    "SchemaInfo",
    "IndexInfo",
    "SAMPLE_ID_COLUMN_NAME",
    "make_sa_type",
    "are_raw_schemas_same",
)

lib/dl_core/dl_core/db/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import contextmanager
import datetime
from functools import (
    lru_cache,
    partial,
)
import math
from typing import (
    Any,
    Callable,
    ClassVar,
    Dict,
    Generator,
    Iterable,
    Optional,
    Type,
    TypeVar,
    final,
)

import pytz

from dl_constants.enums import (
    ConnectionType,
    UserDataType,
)
from dl_core import (
    converter_types_cast,
    exc,
)
from dl_core.db.native_type import (
    CommonNativeType,
    GenericNativeType,
)


make_native_type = GenericNativeType.normalize_name_and_create


_INN_TV = TypeVar("_INN_TV")


def _if_not_none(value: _INN_TV, func: Callable[[_INN_TV], Any]) -> Any:
    if value is not None:
        try:
            return func(value)
        except ValueError:
            raise exc.DataParseError(f"Cannot convert {value!r} to {func.__name__}", query=None)
    return None


def make_date(value: Any) -> Optional[datetime.date]:
    if isinstance(value, datetime.datetime):
        return value.date()
    if isinstance(value, datetime.date):
        return value
    if isinstance(value, str):
        return converter_types_cast._to_date(value)
    return None


def make_datetime(value: Any) -> Optional[datetime.datetime]:
    if isinstance(value, datetime.datetime):
        return value
    if isinstance(value, datetime.date):
        return datetime.datetime(value.year, value.month, value.day)
    if isinstance(value, str):
        return converter_types_cast._to_datetime(value)
    return None


def make_int(value: Any) -> Optional[int]:
    if value is None:
        return None
    if isinstance(value, float) and (math.isinf(value) or math.isnan(value)):
        return None
    return int(value)


@contextmanager
def _handle_type_cast_errors() -> Generator[None, None, None]:
    try:
        yield
    except (TypeError, ValueError):
        raise exc.TypeCastFailed("Type casting failed for value")


class TypeCaster:
    cast_func: Callable[[Any], Any] = lambda value: value
    custom_null_value: Any = None

    def _cast_for_input(self, value: Any) -> Any:
        if self.custom_null_value is not None:
            if value is None or value == self.custom_null_value:
                return self.custom_null_value
        return self.__class__.cast_func(value)

    def _cast_for_output(self, value: Any) -> Any:
        if self.custom_null_value is not None:
            if value is None or value == self.custom_null_value:
                return None
        return self.__class__.cast_func(value)

    @final
    def cast_for_input(self, value: Any) -> Any:
        with _handle_type_cast_errors():
            return self._cast_for_input(value)

    @final
    def cast_for_output(self, value: Any) -> Any:
        with _handle_type_cast_errors():
            return self._cast_for_output(value)


class IntegerTypeCaster(TypeCaster):
    cast_func = make_int


class FloatTypeCaster(TypeCaster):
    cast_func = lambda value: _if_not_none(value, float)  # noqa


class BooleanTypeCaster(TypeCaster):
    cast_func = lambda value: _if_not_none(value, bool)  # noqa


class YTBooleanTypeCaster(BooleanTypeCaster):
    def _cast_for_output(self, value: Any) -> Optional[bool]:
        if value == "0":
            return False
        elif value == "1":
            return True
        else:
            return super()._cast_for_output(value)


class StringTypeCaster(TypeCaster):
    cast_func = lambda value: _if_not_none(value, str)  # noqa


class DateTypeCaster(TypeCaster):
    cast_func = make_date


class DatetimeTypeCaster(TypeCaster):
    cast_func = make_datetime


class DatetimeTZTypeCaster(TypeCaster):
    cast_func = make_datetime  # TODO: re-check edgy cases


class GenericDatetimeTypeCaster(TypeCaster):
    cast_func = make_datetime


class UTCDatetimeTypeCaster(TypeCaster):
    def _cast_for_input(self, value: Any) -> Optional[datetime.datetime]:
        result = super()._cast_for_output(value)
        if isinstance(result, datetime.datetime) and result.tzinfo is not None:
            result = result.astimezone(datetime.timezone.utc).replace(tzinfo=None)
        return result


class UTCTimezoneDatetimeTypeCaster(TypeCaster):
    def _cast_for_input(self, value: Any) -> Optional[datetime.datetime]:
        result = super()._cast_for_output(value)
        if isinstance(result, datetime.datetime) and result.tzinfo is None:
            result = result.replace(tzinfo=pytz.UTC)
        return result


class UnsupportedCaster(TypeCaster):
    """
    Only allows `NULL`s.
    See also: `dl_core.data_source.sql.BaseSQLDataSource._make_raw_column_select`
    """

    def _cast_for_input(self, value: Any) -> None:
        if value is None:
            return None
        raise exc.TypeCastUnsupported("Asked `cast_for_input` for an Unsupported type for a non-null")

    def _cast_for_output(self, value: Any) -> None:
        if value is None:
            return None
        raise exc.TypeCastUnsupported("Asked `cast_for_output` for an Unsupported type for a non-null")


def _cast_array(value: Optional[Iterable[_INN_TV]], f: Callable[[_INN_TV], Any]) -> Optional[tuple]:
    if value is None:
        return None
    return tuple(map(f, value))


class ArrayIntTypeCaster(TypeCaster):
    cast_func = partial(_cast_array, f=IntegerTypeCaster.cast_func)


class ArrayFloatTypeCaster(TypeCaster):
    cast_func = partial(_cast_array, f=FloatTypeCaster.cast_func)


class ArrayStrTypeCaster(TypeCaster):
    cast_func = partial(_cast_array, f=StringTypeCaster.cast_func)


def _make_lowercase(value: Optional[str]) -> Optional[str]:
    if not value:
        return value
    return value.lower()


class LowercaseTypeCaster(TypeCaster):
    cast_func = _make_lowercase


class TypeTransformer:
    conn_type: ClassVar[ConnectionType]
    native_to_user_map: ClassVar[dict[GenericNativeType, UserDataType]] = {}
    user_to_native_map: ClassVar[dict[UserDataType, GenericNativeType]] = {}
    casters: ClassVar[dict[UserDataType, TypeCaster]] = {
        UserDataType.integer: IntegerTypeCaster(),
        UserDataType.float: FloatTypeCaster(),
        UserDataType.boolean: BooleanTypeCaster(),
        UserDataType.string: StringTypeCaster(),
        UserDataType.date: DateTypeCaster(),
        UserDataType.datetime: DatetimeTypeCaster(),
        UserDataType.datetimetz: DatetimeTZTypeCaster(),
        UserDataType.genericdatetime: GenericDatetimeTypeCaster(),
        UserDataType.geopoint: StringTypeCaster(),
        UserDataType.geopolygon: StringTypeCaster(),
        UserDataType.uuid: StringTypeCaster(),
        UserDataType.markup: StringTypeCaster(),
        UserDataType.array_int: ArrayIntTypeCaster(),
        UserDataType.array_float: ArrayFloatTypeCaster(),
        UserDataType.array_str: ArrayStrTypeCaster(),
        UserDataType.tree_str: ArrayStrTypeCaster(),  # Same as array
        UserDataType.unsupported: UnsupportedCaster(),
    }

    def type_native_to_user(
        self,
        native_t: GenericNativeType,
        user_t: Optional[UserDataType] = None,
    ) -> UserDataType:
        if user_t is not None:
            # original UT is given, try to validate against NT.
            # read as 'native type might have been made from the provided user type'.
            # Sample case: native 'uint8' might be a 'boolean' user type.
            if native_t.as_generic == self.user_to_native_map[user_t].as_generic:
                return user_t

        try:
            return self.native_to_user_map.get(native_t) or self.native_to_user_map[native_t.as_generic]
        except KeyError:
            raise exc.UnsupportedNativeTypeError(native_t)

    def make_foreign_native_type_conversion(self, native_t: GenericNativeType) -> GenericNativeType:
        """
        Attempt to make a native type for the current database that corresponds
        to a native type of another database.

        If there is no known direct conversion, should return `native_t` as-is
        (with non-matching `conn_type`).
        """
        return native_t  # no known conversion

    def type_user_to_native(
        self, user_t: UserDataType, native_t: Optional[GenericNativeType] = None
    ) -> GenericNativeType:
        if native_t is not None:
            # original NT is given, try to do a direct conversion

            if native_t.conn_type != self.conn_type:
                # attempt to translate to own native type
                native_t = self.make_foreign_native_type_conversion(native_t)

            if native_t.conn_type == self.conn_type:
                # it is from the same DB type, so try to validate against UT
                if user_t == self.native_to_user_map.get(native_t) or user_t == self.native_to_user_map.get(
                    native_t.as_generic
                ):
                    return native_t

        try:
            result = self.user_to_native_map[user_t]
        except KeyError:
            raise exc.UnsupportedBITypeError(user_t)

        # Assume all databases support `nullable` in a similar way, so pass it
        # along if possible.
        if native_t is not None and isinstance(native_t, CommonNativeType):
            result = result.as_common().clone(nullable=native_t.nullable)

        return result

    @classmethod
    def cast_for_input(cls, value: Any, user_t: UserDataType) -> Any:
        """Prepare value for insertion into the database"""
        return cls.casters[user_t].cast_for_input(value=value)

    @classmethod
    def cast_for_output(cls, value: Any, user_t: Optional[UserDataType] = None) -> Any:
        """Convert value from DB to Python value conforming to given ``user_t``"""
        if user_t is None:
            return value
        return cls.casters[user_t].cast_for_output(value=value)


TYPE_TRANSFORMER_MAP: Dict[ConnectionType, TypeTransformer] = {}


def register_type_transformer_class(conn_type: ConnectionType, tt_cls: Type[TypeTransformer]) -> None:
    TYPE_TRANSFORMER_MAP[conn_type] = tt_cls()


@lru_cache(maxsize=32)
def get_type_transformer(conn_type: ConnectionType) -> TypeTransformer:
    try:
        return TYPE_TRANSFORMER_MAP[conn_type]
    except KeyError:
        raise exc.UnsupportedDatabaseError(conn_type)

lib/dl_core/dl_core/db/conversion_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import namedtuple
from typing import (
    Any,
    FrozenSet,
    List,
    Optional,
    Sequence,
    Union,
)

import attr

from dl_constants.enums import (
    IndexKind,
    UserDataType,
)
from dl_core.db.native_type import GenericNativeType
from dl_utils.utils import get_type_full_name


_SchemaColumn = namedtuple(
    "_SchemaColumn",
    (
        "name",
        "title",
        "user_type",
        "nullable",  # Deprecated in favor of `native_type.nullable`
        "native_type",
        "source_id",
        "has_auto_aggregation",
        "lock_aggregation",
        "description",
    ),
)


class SchemaColumn(_SchemaColumn):
    def __new__(  # type: ignore  # TODO: fix
        cls,
        name: str,
        title: Optional[str] = None,
        user_type: Optional[Union[UserDataType, str]] = None,
        nullable: Optional[bool] = True,
        native_type: Optional[GenericNativeType] = None,
        source_id: Any = None,
        has_auto_aggregation: Optional[bool] = None,
        lock_aggregation: bool = False,
        description: Optional[str] = None,
    ):
        title = title or name
        if isinstance(user_type, str):
            user_type = UserDataType[user_type]
        has_auto_aggregation = has_auto_aggregation if has_auto_aggregation is not None else False
        return super().__new__(
            cls,
            name=name,
            title=title,
            user_type=user_type,
            nullable=nullable,
            native_type=native_type,
            source_id=source_id,
            has_auto_aggregation=has_auto_aggregation,
            lock_aggregation=lock_aggregation,
            description=description or "",
        )

    def clone(self, **kwargs) -> "SchemaColumn":  # type: ignore  # TODO: fix
        return SchemaColumn(**dict(self._asdict(), **kwargs))


@attr.s(slots=True, frozen=True)
class IndexInfo:
    columns: Sequence[str] = attr.ib()
    kind: Optional[IndexKind] = attr.ib()


@attr.s(slots=True)
class SchemaInfo:
    # TODO FIX: May be make a Tuple/Sequence instead of list
    schema: List[SchemaColumn] = attr.ib()
    skipped: List[SchemaColumn] = attr.ib(default=None)
    converted: List[SchemaColumn] = attr.ib(default=None)

    # None value should be treated as "unknown/no info"
    indexes: Optional[FrozenSet[IndexInfo]] = attr.ib(default=None)

    @indexes.validator
    def check_indexes(self, attribute: Any, value: Any) -> None:
        if value is None or isinstance(value, frozenset):
            pass
        else:
            raise TypeError(f"Indexes must be a frozen set or non, not {get_type_full_name(type(value))!r}")

    def clone(self, **kwargs) -> SchemaInfo:  # type: ignore  # TODO: fix
        return attr.evolve(self, **kwargs)

    @classmethod
    def from_schema(cls, schema: Sequence[SchemaColumn]) -> SchemaInfo:
        # TODO FIX: May be make a Tuple/Sequence instead of list
        return SchemaInfo(schema=list(schema))

lib/dl_core/dl_core/db/elements.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
    Optional,
    Type,
    TypeVar,
    Union,
)

import attr
from sqlalchemy.types import TypeEngine

from dl_constants.enums import ConnectionType


SATypeSpec = Union[Type[TypeEngine], TypeEngine, str, None]


def norm_native_type(native_t: SATypeSpec) -> Optional[str]:
    """Normalize native type to internal name"""
    if native_t is None:
        return None

    if hasattr(native_t, "__visit_name__"):
        native_t_name = native_t.__visit_name__
    elif not isinstance(native_t, str):
        native_t_name = type(native_t).__name__
    else:
        native_t_name = native_t

    if hasattr(native_t, "item_type"):
        item_type_name = norm_native_type(native_t.item_type)
        native_t_name = f"{native_t_name}({item_type_name})"

    return native_t_name.lower()


_GENERIC_NATIVE_TYPE_TV = TypeVar("_GENERIC_NATIVE_TYPE_TV", bound="GenericNativeType")


@attr.s(frozen=True)
class GenericNativeType:
    """
    Basic NativeType information usable for type-conversion maps.
    """

    native_type_class_name: ClassVar[str] = "generic_native_type"
    conn_type: ConnectionType = attr.ib()
    name: str = attr.ib()

    @classmethod
    def normalize_name_and_create(
        cls: Type[_GENERIC_NATIVE_TYPE_TV],
        conn_type: ConnectionType,
        name: SATypeSpec,
    ) -> _GENERIC_NATIVE_TYPE_TV:
        normalized_name = norm_native_type(name)
        return cls(
            conn_type=conn_type,
            name=normalized_name,  # type: ignore  # TODO: fix
        )

    @property
    def as_generic(self) -> GenericNativeType:  # for subclasses
        return GenericNativeType(conn_type=self.conn_type, name=self.name)

    def as_common(self, default_nullable=True) -> CommonNativeType:  # type: ignore  # TODO: fix
        """
        Helper method that converts a GenericNativeType object to a
        CommonNativeType object with the specified `nullable` value; for
        CommonNativeType objects does nothing.
        """
        return CommonNativeType(conn_type=self.conn_type, name=self.name, nullable=default_nullable)

    def clone(self: _GENERIC_NATIVE_TYPE_TV, **kwargs: Any) -> _GENERIC_NATIVE_TYPE_TV:
        return attr.evolve(self, **kwargs)


_COMMON_NATIVE_TYPE_TV = TypeVar("_COMMON_NATIVE_TYPE_TV", bound="CommonNativeType")


@attr.s(frozen=True)
class CommonNativeType(GenericNativeType):
    native_type_class_name = "common_native_type"
    nullable: bool = attr.ib(default=True)

    @classmethod
    def normalize_name_and_create(
        cls: Type[_COMMON_NATIVE_TYPE_TV],
        conn_type: ConnectionType,
        name: SATypeSpec,
        nullable: bool = True,
    ) -> _COMMON_NATIVE_TYPE_TV:
        normalized_name = norm_native_type(name)
        return cls(
            conn_type=conn_type,
            name=normalized_name,  # type: ignore  # TODO: fix
            nullable=nullable,
        )

    def as_common(self: _COMMON_NATIVE_TYPE_TV, default_nullable=None) -> _COMMON_NATIVE_TYPE_TV:  # type: ignore  # TODO: fix
        return self


@attr.s(frozen=True)
class LengthedNativeType(CommonNativeType):
    native_type_class_name = "lengthed_native_type"
    length: Optional[int] = attr.ib(default=None)


@attr.s(frozen=True)
class ClickHouseNativeType(CommonNativeType):
    native_type_class_name = "clickhouse_native_type"
    lowcardinality: bool = attr.ib(default=False)


@attr.s(frozen=True)
class ClickHouseDateTime64NativeType(ClickHouseNativeType):  # noqa
    native_type_class_name = "clickhouse_datetime64_native_type"
    precision: int = attr.ib(kw_only=True)


@attr.s(frozen=True)
class ClickHouseDateTimeWithTZNativeType(ClickHouseNativeType):
    native_type_class_name = "clickhouse_datetimewithtz_native_type"
    # It is kind-of non-optional, but hard to do this after all the inheritance.
    timezone_name: str = attr.ib(default="UTC")
    # Whether the timezone was specified explicitly in the type,
    # or was inferred from the system timezone.
    explicit_timezone: bool = attr.ib(default=True)


@attr.s(frozen=True)
class ClickHouseDateTime64WithTZNativeType(ClickHouseDateTime64NativeType):
    native_type_class_name = "clickhouse_datetime64withtz_native_type"
    # It is kind-of non-optional, but hard to do this after all the inheritance.
    timezone_name: str = attr.ib(default="UTC")
    # Whether the timezone was specified explicitly in the type,
    # or was inferred from the system timezone.
    explicit_timezone: bool = attr.ib(default=True)

lib/dl_core/dl_core/db/native_type.py
datalens-tech/datalens-backend
"""
Marshmallow schema for the NativeType classes.

Shared between storage (dl_core.us_manager.storage_schemas) and RQE api.

TODO: refactor together with `dl_core.base_models`.
"""

from __future__ import annotations

import logging
from typing import (
    ClassVar,
    Generic,
    Type,
    TypeVar,
)

from marshmallow import (
    Schema,
    fields,
    post_load,
)
from marshmallow_oneofschema import OneOfSchema

from dl_constants.enums import ConnectionType
from dl_core.db.native_type import (
    ClickHouseDateTime64NativeType,
    ClickHouseDateTime64WithTZNativeType,
    ClickHouseDateTimeWithTZNativeType,
    ClickHouseNativeType,
    CommonNativeType,
    GenericNativeType,
    LengthedNativeType,
)
from dl_model_tools.schema.dynamic_enum_field import DynamicEnumField


LOGGER = logging.getLogger(__name__)


_TARGET_TV = TypeVar("_TARGET_TV")


class NativeTypeSchemaBase(Schema, Generic[_TARGET_TV]):
    """(Shared ((Native Type) Storage Schema)), common base class for NT schemas."""

    TARGET_CLS: ClassVar[Type[_TARGET_TV]]

    @post_load(pass_many=False)
    def to_object(self, data: dict, **_):  # type: ignore  # TODO: fix
        return self.TARGET_CLS(**data)  # type: ignore  # TODO: fix


class GenericNativeTypeSchema(NativeTypeSchemaBase[GenericNativeType]):
    """
    ((Generic (Native Type)) Storage Schema),
    (not to be confused with (Generic ((Native Type) Storage Schema))).
    """

    TARGET_CLS = GenericNativeType
    conn_type = DynamicEnumField(ConnectionType)
    name = fields.String()


class CommonNativeTypeSchema(GenericNativeTypeSchema):
    TARGET_CLS = CommonNativeType
    nullable = fields.Boolean(required=False, load_default=True)


class LengthedNativeTypeSchema(CommonNativeTypeSchema):
    TARGET_CLS = LengthedNativeType
    length = fields.Integer(required=False, allow_none=True, load_default=None)


class ClickHouseNativeTypeSchema(CommonNativeTypeSchema):
    TARGET_CLS = ClickHouseNativeType
    lowcardinality = fields.Boolean(required=False, load_default=False)


class ClickHouseDateTimeWithTZNativeTypeSchema(ClickHouseNativeTypeSchema):
    TARGET_CLS = ClickHouseDateTimeWithTZNativeType
    timezone_name = fields.String(required=False, load_default="UTC")
    explicit_timezone = fields.Boolean(required=False, load_default=True)


class ClickHouseDateTime64NativeTypeSchema(ClickHouseNativeTypeSchema):
    TARGET_CLS = ClickHouseDateTime64NativeType  # type: ignore  # TODO: fix
    precision = fields.Integer(required=True)


class ClickHouseDateTime64WithTZNativeTypeSchema(ClickHouseDateTime64NativeTypeSchema):
    TARGET_CLS = ClickHouseDateTime64WithTZNativeType  # type: ignore  # TODO: fix
    timezone_name = fields.String(required=False, load_default="UTC")
    explicit_timezone = fields.Boolean(required=False, load_default=True)


class OneOfNativeTypeSchemaBase(OneOfSchema):
    """(OneOf (Native Type) Storage Schema)"""

    type_field = "native_type_class_name"
    type_schemas = {
        schema.TARGET_CLS.native_type_class_name: schema
        for schema in (
            GenericNativeTypeSchema,
            CommonNativeTypeSchema,
            LengthedNativeTypeSchema,
            ClickHouseNativeTypeSchema,
            ClickHouseDateTimeWithTZNativeTypeSchema,
            ClickHouseDateTime64NativeTypeSchema,
            ClickHouseDateTime64WithTZNativeTypeSchema,
        )
    }

    def get_obj_type(self, obj):  # type: ignore  # TODO: fix
        return getattr(obj, self.type_field)

    def _load(self, data, *, partial=None, unknown=None):  # type: ignore  # TODO: fix
        data.setdefault(self.type_field, "common_native_type")
        return super()._load(data, partial=partial, unknown=unknown)


class OneOfNativeTypeSchema(OneOfNativeTypeSchemaBase):
    """Transition/compatibility layer. Should eventually become empty."""

    def dump(self, value, *args, **kwargs):  # type: ignore  # TODO: fix
        if value is None:
            return None

        if isinstance(value, dict):
            raise Exception("OneOfNativeTypeSchema dumping a dict")

        if isinstance(value, str):
            raise Exception("OneOfNativeTypeSchema dumping an str")

        return super().dump(value, *args, **kwargs)

    def load(self, value, *args, **kwargs):  # type: ignore  # TODO: fix
        if value is None:
            return None

        if isinstance(value, str):
            # probably a normal case until the transition is done.
            # LOGGER.info("OneOfNativeTypeSchema loading from an str")

            from dl_core.us_manager.storage_schemas.base import CtxKey

            ds_conn_type = self.context.get("ds_conn_type") or self.context[CtxKey.ds_conn_type]
            return GenericNativeType(conn_type=ds_conn_type, name=value)

        if isinstance(value, GenericNativeType):
            raise Exception("OneOfNativeTypeSchema loading an obj")

        return super().load(value, *args, **kwargs)

lib/dl_core/dl_core/db/native_type_schema.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Callable,
    Optional,
)

from sqlalchemy.types import TypeEngine

from dl_core.db.native_type import GenericNativeType


SQLALCHEMY_TYPES: dict[GenericNativeType, Callable[[GenericNativeType], TypeEngine]] = {}


def make_sa_type(native_type: GenericNativeType, nullable: Optional[bool] = None) -> TypeEngine:
    if nullable is not None:
        # For a type without `nullable`, specifies it.
        native_type = native_type.as_common(default_nullable=nullable)
    type_gen = SQLALCHEMY_TYPES.get(native_type) or SQLALCHEMY_TYPES[native_type.as_generic]
    return type_gen(native_type)


def register_sa_types(sa_type_dict: dict[GenericNativeType, Callable[[GenericNativeType], TypeEngine]]) -> None:
    SQLALCHEMY_TYPES.update(sa_type_dict)

lib/dl_core/dl_core/db/sa_types.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Callable,
    Type,
)

from dl_core.db.native_type import (
    GenericNativeType,
    LengthedNativeType,
)


if TYPE_CHECKING:
    from sqlalchemy.types import TypeEngine


make_native_type = GenericNativeType.normalize_name_and_create


def simple_instantiator(typecls: Type[TypeEngine]) -> Callable[[GenericNativeType], TypeEngine]:
    def type_gen(nt: GenericNativeType) -> TypeEngine:
        return typecls()

    return type_gen


def lengthed_instantiator(typecls: Type[TypeEngine], default_length: int = 255) -> Callable:
    def type_gen(nt: GenericNativeType) -> TypeEngine:
        if isinstance(nt, LengthedNativeType):
            length = nt.length or default_length
        else:
            length = default_length
        return typecls(length)  # type: ignore  # TODO: fix

    return type_gen


def timezone_instantiator(typecls: Type[TypeEngine]) -> Callable:
    def type_gen(nt: GenericNativeType) -> TypeEngine:
        return typecls(timezone=True)  # type: ignore  # TODO: fix

    return type_gen


def typed_instantiator(typecls: Type[TypeEngine], inner_type: Type[TypeEngine]) -> Callable:
    def type_gen(nt: GenericNativeType) -> TypeEngine:
        return typecls(inner_type)  # type: ignore  # TODO: fix

    return type_gen

lib/dl_core/dl_core/db/sa_types_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    List,
    Optional,
)

from dl_core.db.elements import SchemaColumn


SAMPLE_ID_COLUMN_NAME = "__dl_sample_id"  # TODO: Move somewhere else - it has to do with sampling, not schemas


def are_raw_schemas_same(first: Optional[List[SchemaColumn]], second: Optional[List[SchemaColumn]]) -> bool:
    if first is None or second is None:
        return first == second
    if len(first) != len(second):
        return False
    for col_1, col_2 in zip(first, second):
        data_1, data_2 = col_1._asdict(), col_2._asdict()
        # FIXME:
        # API doesn't have the full NativeType definition yet, so we can't compare them directly
        exclude_from_comparison = ("native_type", "source_id")
        for col_data in (data_1, data_2):
            for attr in exclude_from_comparison:
                col_data.pop(attr)

        ntypes_arent_none = all((col_1.native_type, col_2.native_type))
        if (
            # main attrs are different
            data_1 != data_2
            # or exactly one of the two native types is None
            or not ntypes_arent_none
            and col_1.native_type != col_2.native_type
            # or native types have different names (only for different db (conn) types)
            or (
                ntypes_arent_none
                and col_1.native_type.conn_type == col_2.native_type.conn_type
                and col_1.native_type.name != col_2.native_type.name
            )
        ):
            return False

    return True

lib/dl_core/dl_core/db/schema.py
datalens-tech/datalens-backend
# TODO FIX: Move this file to connection_executor package after migration to connection executors will be done

from __future__ import annotations

from contextlib import contextmanager
import logging
from typing import (
    Collection,
    Generator,
    Type,
)

from sqlalchemy.engine import Engine
import sqlalchemy.exc
from sqlalchemy.orm import (
    Query,
    Session,
    sessionmaker,
)

from dl_app_tools.profiling_base import GenericProfiler
from dl_constants.enums import SourceBackendType
from dl_core import exc


LOGGER = logging.getLogger(__name__)


_SA_QUERY_CLASSES: dict[SourceBackendType, Type[Query]] = {}


def register_sa_query_cls(backend_type: SourceBackendType, query_cls: Type[Query]) -> None:
    if backend_type in _SA_QUERY_CLASSES:
        assert _SA_QUERY_CLASSES[backend_type] is query_cls
    else:
        _SA_QUERY_CLASSES[backend_type] = query_cls


def get_sa_query_cls(backend_type: SourceBackendType) -> Type[Query]:
    return _SA_QUERY_CLASSES[backend_type]


class CustomSession(Session):
    __closed = False

    def close(self):  # type: ignore
        super().close()
        self.__closed = True

    def execute(self, *args, **kwargs):  # type: ignore
        if self.__closed:
            raise exc.DBSessionError("Cannot execute on closed session")
        return super().execute(*args, **kwargs)


def get_db_session(db_engine: Engine, backend_type: SourceBackendType) -> Session:
    session_kwargs = {}
    session_kwargs["query_cls"] = get_sa_query_cls(backend_type=backend_type)
    SessionCls = sessionmaker(bind=db_engine, class_=CustomSession)
    return SessionCls(**session_kwargs)


_QUERY_FAIL_EXCEPTIONS: set[Type[Exception]] = {
    sqlalchemy.exc.InvalidRequestError,
    sqlalchemy.exc.DBAPIError,
}


def register_query_fail_exceptions(exception_classes: Collection[Type[Exception]]) -> None:
    _QUERY_FAIL_EXCEPTIONS.update(exception_classes)


def get_query_fail_exceptions() -> tuple[Type[Exception], ...]:
    return tuple(_QUERY_FAIL_EXCEPTIONS)


@contextmanager
def db_session_context(db_engine: Engine, backend_type: SourceBackendType) -> Generator[Session, None, None]:
    LOGGER.info("Going to create DB session")
    with GenericProfiler("db-create-session"):
        session = get_db_session(db_engine=db_engine, backend_type=backend_type)

    try:
        yield session
    except Exception as err:
        try:
            session.rollback()
        except get_query_fail_exceptions():
            LOGGER.info(
                "Rollback failed because transaction has already been rolled back " "due to an internal exception"
            )
        LOGGER.error("Rolling back transaction because of error: %s", err)
        raise
    else:
        session.commit()
    finally:
        session.close()
        db_engine.dispose()

lib/dl_core/dl_core/db_session_utils.py
datalens-tech/datalens-backend
from enum import (
    Enum,
    unique,
)


@unique
class QueryExecutorMode(Enum):
    text = "text"
    sqla_dump = "sqla_dump"


@unique
class TableDefinitionType(Enum):
    table_ident = "table_ident"
    sa_text = "sa_text"


@unique
class RQEEventType(Enum):
    raw_cursor_info = "raw_cursor_info"
    raw_chunk = "raw_chunk"
    error_dump = "error_dump"
    finished = "finished"


@unique
class USApiType(Enum):
    v1 = "v1"
    private = "private"
    public = "public"
    embeds = "embeds"


@unique
class RoleReason(Enum):
    # Role was selected
    selected = "selected"
    # Role was nort requested
    not_needed = "not_needed"
    # Role is not supported for source
    not_supported = "not_supported"
    # Role is not yet configured
    not_configured = "not_configured"
    # Missing source for default role
    missing_source = "missing_source"
    # Schema mismatch between default role and origin
    schema_mismatch = "schema_mismatch"
    # This role is not allowed for feature-managed sources
    forbidden_for_features = "forbidden_for_features"

lib/dl_core/dl_core/enums.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
    Optional,
)

import requests

from dl_constants.exc import DLBaseException
from dl_constants.types import TJSONLike


class RSTError(Exception):
    """
    An intermediary error that wraps the "Connection reset by peer"
    ClientOSError cases.
    Made for skipping the `errno` check in a main exception handler.
    """


class AIOHttpConnTimeoutError(Exception):
    """
    An intermediary error that wraps the "Connection timeout"
    aiohttp.client_exceptions.ServerTimeoutError.
    Made for skipping the exc message check in a main exception handler.
    """


class UnsupportedNativeTypeError(Exception):
    pass


class UnsupportedBITypeError(Exception):
    pass


class UnsupportedDatabaseError(Exception):
    pass


class TypeCastError(DLBaseException):
    err_code = DLBaseException.err_code + ["TYPE_CAST"]


class TypeCastUnsupported(TypeCastError):
    err_code = TypeCastError.err_code + ["UNSUPPORTED"]


class TypeCastFailed(TypeCastError):
    err_code = TypeCastError.err_code + ["FAILED"]
    default_message = "Type casting failed for value"


class SourceAccessDenied(DLBaseException):
    err_code = DLBaseException.err_code + ["SOURCE_ACCESS_DENIED"]


class SourceAccessInvalidToken(SourceAccessDenied):
    err_code = SourceAccessDenied.err_code + ["INVALID_TOKEN"]
    default_message = "Invalid user token"


class DatabaseUnavailable(DLBaseException):
    err_code = DLBaseException.err_code + ["DATABASE_UNAVAILABLE"]
    default_message = "Data source is unavailable"


class DataSourceConfigurationError(DLBaseException):
    err_code = DLBaseException.err_code + ["SOURCE_CONFIG"]


class TableNameNotConfiguredError(DataSourceConfigurationError):
    err_code = DataSourceConfigurationError.err_code + ["TABLE_NOT_CONFIGURED"]
    default_message = "Table is not ready yet"
    details = {"info": "No table for data source"}


class TableNameInvalidError(DataSourceConfigurationError):
    err_code = DataSourceConfigurationError.err_code + ["TABLE_NAME_INVALID"]


class DatasetConfigurationError(DLBaseException):
    err_code = DLBaseException.err_code + ["DS_CONFIG"]


class NoCommonRoleError(DLBaseException):
    err_code = DatasetConfigurationError.err_code + ["NO_COMMON_ROLE"]


# Data sources


class DataSourceError(DLBaseException):
    err_code = DLBaseException.err_code + ["SOURCE"]
    default_message = "Data source error"


class DataSourceNotFound(DataSourceError):
    err_code = DataSourceError.err_code + ["NOT_FOUND"]
    default_message = "Data source not found"


class DataSourceTitleError(DataSourceError):
    err_code = DataSourceError.err_code + ["TITLE"]
    default_message = "Invalid data source title"


class DataSourceTitleConflict(DataSourceTitleError):
    err_code = DataSourceTitleError.err_code + ["CONFLICT"]
    default_message = "Data source title conflicts with another source"


class DataSourcesInconsistent(DataSourceError):
    err_code = DataSourceError.err_code + ["INCONSISTENT"]
    default_message = "Data sources are inconsistent"


class DataSourceErrorFromComponentError(DataSourceError):
    err_code = DataSourceError.err_code


# Source avatars


class SourceAvatarError(DLBaseException):
    err_code = DLBaseException.err_code + ["AVATAR"]
    default_message = "Source avatar error"


class SourceAvatarNotFound(SourceAvatarError):
    err_code = SourceAvatarError.err_code + ["NOT_FOUND"]
    default_message = "Source avatar not found"


class SourceAvatarTitleError(SourceAvatarError):
    err_code = SourceAvatarError.err_code + ["TITLE"]
    default_message = "Invalid source avatar title"


class SourceAvatarTitleConflict(SourceAvatarTitleError):
    err_code = SourceAvatarTitleError.err_code + ["CONFLICT"]
    default_message = "Source avatar title conflicts with another avatar"


class UnknownReferencedAvatar(SourceAvatarNotFound):
    err_code = SourceAvatarNotFound.err_code + ["FIELD_REF"]
    default_message = "Field references an unknown source avatar"


class UnboundAvatarError(DatasetConfigurationError):
    err_code = DatasetConfigurationError.err_code + ["UNBOUND_AVATAR"]


# Avatar relations


class AvatarRelationError(DLBaseException):
    err_code = DLBaseException.err_code + ["RELATION"]
    default_message = "Avatar relation error"


class AvatarRelationNotFound(AvatarRelationError):
    err_code = AvatarRelationError.err_code + ["NOT_FOUND"]
    default_message = "Avatar relation not found"


class AvatarRelationJoinTypeError(AvatarRelationError):
    err_code = AvatarRelationError.err_code + ["JOIN_TYPE"]
    default_message = "Invalid join type for referenced data sources"


# Obligatory filters


class ObligatoryFilterError(DLBaseException):
    err_code = DLBaseException.err_code + ["OBLIG_FILTER"]
    default_message = "Obligatory filter error"


class ObligatoryFilterNotFound(ObligatoryFilterError):
    err_code = ObligatoryFilterError.err_code + ["NOT_FOUND"]
    default_message = "Obligatory filter not found"


# Fields


class FieldError(DLBaseException):
    err_code = DLBaseException.err_code + ["FIELD"]
    default_message = "Field error"


class InvalidFieldError(FieldError):
    err_code = FieldError.err_code + ["INVALID"]
    default_message = "Field is invalid"


class FieldNotFound(FieldError):
    err_code = FieldError.err_code + ["NOT_FOUND"]
    default_message = "Unknown field"


class FieldTitleError(FieldError):
    err_code = FieldError.err_code + ["TITLE"]
    default_message = "Invalid field title"


class FieldTitleConflict(FieldTitleError):
    err_code = FieldTitleError.err_code + ["CONFLICT"]
    default_message = "Field title conflicts with another field"


# US Client errors


class USReqException(DLBaseException):
    err_code = DLBaseException.err_code + ["US"]

    def __init__(
        self,
        message: Optional[str] = None,
        details: Optional[Dict[str, Any]] = None,
        orig: Optional[Exception] = None,
        orig_exc: Optional[Exception] = None,
        debug_info: Optional[Dict[str, Any]] = None,
        params: Optional[Dict[str, Any]] = None,
    ):
        self.orig_exc = orig_exc  # TODO: stop using '.orig_exc' and '.orig
        super().__init__(
            message=message,
            details=details,
            orig=orig if orig is not None else orig_exc,
            debug_info=debug_info,
            params=params,
        )


class USObjectNotFoundException(USReqException):
    default_message = "Object not found"
    err_code = USReqException.err_code + ["OBJ_NOT_FOUND"]


class USAccessDeniedException(USReqException):
    default_message = "Access denied"
    err_code = USReqException.err_code + ["ACCESS_DENIED"]


class USLockUnacquiredException(USReqException):
    default_message = "Object locked"
    err_code = USReqException.err_code + ["ENTRY_LOCKED"]


class USPermissionCheckError(USReqException):
    err_code = USReqException.err_code + ["PERMISSIONS_ERROR"]

    @property
    def message(self) -> str:
        return "US permission check error"


class USBadRequestException(USReqException):
    err_code = USReqException.err_code + ["BAD_REQUEST"]

    @property
    def message(self) -> str:
        try:
            if self.orig_exc and isinstance(self.orig_exc.response, requests.Response):
                return self.orig_exc.response.json()["message"]
        except (ValueError, KeyError, AttributeError):
            pass
        return super().message


class USAlreadyExistsException(USBadRequestException):
    err_code = USBadRequestException.err_code + ["ALREADY_EXISTS"]


class USIncorrectTenantIdException(USReqException):
    @property
    def message(self) -> str:
        return "Incorrect tenant_id"


class USInteractionDisabled(USReqException):
    pass


class USInvalidResponse(USReqException):
    pass


# TODO FIX: Validate err_code/message (and think about should this exception be thrown up to user)
class UnexpectedUSEntryType(DLBaseException):
    """Should be thrown by USM if expected type does not match actual US entry type"""

    err_code = DLBaseException.err_code + ["UNEXPECTED_ENTRY_TYPE"]
    default_message = "Unexpected entry type"


# Other


class FailedToLoadSchema(DLBaseException):
    err_code = DLBaseException.err_code + ["COLUMN_SCHEMA_FAILED"]
    default_message = "Failed to load description of table columns."


class DatabaseQueryError(DLBaseException):
    err_code = DLBaseException.err_code + ["DB"]
    default_message = "Database error."

    db_message: Optional[str] = None
    query: Optional[str] = None

    def __init__(
        self,
        db_message: Optional[str] = None,
        query: Optional[str] = None,
        message: Optional[str] = None,
        details: Optional[Dict[str, Any]] = None,
        orig: Optional[Exception] = None,
        debug_info: Optional[Dict[str, Any]] = None,
        params: Optional[Dict[str, Any]] = None,
    ):
        super().__init__(message=message, details=details, orig=orig, debug_info=debug_info, params=params)
        self.db_message = db_message
        if db_message is not None:
            self.details.update(db_message=db_message)
        self.query = query
        # Future note about the `self.debug_info` / `self.details`:
        # Currently (at least on the internal installation),
        # both debug and details should be printed to the user.
        # But `details` are to-be-done, and should contain templatable
        # i18n-able something.
        # Therefore, for now, there's no use case for `details`,
        # and they're not printed to the user.
        self.debug_info.setdefault("db_message", db_message)
        self.debug_info.setdefault("query", query)

    @classmethod
    def _from_jsonable_dict(cls, data: Dict) -> "DLBaseException":
        new_exc = cls(
            db_message=data.pop("db_message"),
            query=data.pop("query"),
            message=data.pop("message", None),
        )
        new_exc.details = data.pop("details")
        new_exc.debug_info = data.pop("debug_info")
        new_exc.params = data.pop("params", dict())
        return new_exc

    def to_jsonable_dict(self) -> Dict[str, TJSONLike]:
        d = super().to_jsonable_dict()
        d.update(
            db_message=self.db_message,
            query=self.query,
        )
        return d

    def __str__(self) -> str:
        return "{0.db_message} (DB query: {0.query})".format(self)


class ResultRowCountLimitExceeded(DLBaseException):
    err_code = DLBaseException.err_code + ["ROW_COUNT_LIMIT"]
    default_message = "Received too many result data rows."


class SourceDoesNotExist(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["SOURCE_DOES_NOT_EXIST"]
    default_message = "Data source (table) does not exist."
    formatting_messages = {
        frozenset({"table_definition"}): "Data source (table) identified by: `{table_definition}` does not exist."
    }

    def __init__(
        self,
        db_message: Optional[str] = None,
        query: Optional[str] = None,
        message: Optional[str] = None,
        details: Optional[Dict[str, Any]] = None,
        orig: Optional[Exception] = None,
        debug_info: Optional[Dict[str, Any]] = None,
        params: Optional[Dict[str, Any]] = None,
    ):
        super(SourceDoesNotExist, self).__init__(
            db_message=db_message,
            query=query,
            message=message,
            details=details,
            orig=orig,
            debug_info=debug_info,
            params=params or {},
        )

        if self.params.get("table_definition"):
            return


class DatabaseDoesNotExist(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["DATABASE_DOES_NOT_EXIST"]
    default_message = "Data source database does not exist."


class ColumnDoesNotExist(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["COLUMN_DOES_NOT_EXIST"]
    default_message = "Requested database column does not exist."


class InvalidQuery(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["INVALID_QUERY"]
    default_message = "Invalid SQL query to the database."


class SourceResponseError(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["SOURCE_ERROR"]
    default_message = "Data source failed to respond correctly."


class SourceTimeout(SourceResponseError):
    err_code = SourceResponseError.err_code + ["TIMEOUT"]
    default_message = "Data source timed out."


class SourceClosedPrematurely(SourceResponseError):
    err_code = SourceResponseError.err_code + ["CLOSED_PREMATURELY"]
    default_message = "Data source ended the response prematurely (disconnected or dropped the connection)."


class SourceProtocolError(SourceResponseError):
    err_code = SourceResponseError.err_code + ["INVALID_RESPONSE"]
    default_message = "Data source failed to respond correctly (invalid HTTP or higher protocol response)."


class InvalidArgumentType(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["INVALID_ARGUMENT_TYPE"]
    default_message = "Invalid argument type."


class UnknownFunction(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["INVALID_FUNCTION"]
    default_message = "Unknown function."


class SourceConnectError(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["SOURCE_CONNECT_ERROR"]
    default_message = "Data source refused connection."


class SourceHostNotKnownError(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["SOURCE_HOST_NOT_KNOWN_ERROR"]
    default_message = "Data source host name or service not known."


class DatabaseOperationalError(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["DB_OPERATIONAL_ERROR"]


class MaterializationNotFinished(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["MATERIALIZATION_NOT_FINISHED"]
    default_message = "Data is not available because materialization is not yet complete."


class DbMemoryLimitExceeded(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["MEMORY_LIMIT_EXCEEDED"]
    default_message = "Memory limit has been exceeded during query execution."


class DbAuthenticationFailed(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["AUTHENTICATION_FAILED"]
    default_message = "Database authentication failed."


class DBIndexNotUsed(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["INDEX_NOT_USED"]
    default_message = "Filtration by any of indexed columns required."


class DataParseError(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["CANNOT_PARSE"]
    default_message = "Cannot parse value."


class CannotParseNumber(DataParseError):
    err_code = DataParseError.err_code + ["NUMBER"]
    default_message = "Cannot parse number."


class CannotParseDateTime(DataParseError):
    err_code = DataParseError.err_code + ["DATETIME"]
    default_message = "Cannot parse datetime."


class NumberOutOfRange(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["NUMBER_OUT_OF_RANGE"]
    default_message = "Number out of range."


class UnexpectedInfOrNan(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["UNEXPECTED_INF_OR_NAN"]
    default_message = "Unexpected inf or nan to integer conversion."


class JoinColumnTypeMismatch(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["JOIN_COLUMN_TYPE_MISMATCH"]
    default_message = "Columns in JOIN have different types."


class NoSpaceLeft(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["NO_SPACE_LEFT"]
    default_message = "No space left on device."


class DivisionByZero(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["ZERO_DIVISION"]
    default_message = "Division by zero."


class UserQueryAccessDenied(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["USER_QUERY_ACCESS_DENIED"]
    default_message = "Query access denied to user"


class WrongQueryParameterization(DatabaseQueryError):
    err_code = DatabaseQueryError.err_code + ["WRONG_QUERY_PARAMETERIZATION"]
    default_message = "Wrong query parameterization. Parameter was not found"


class QueryConstructorError(Exception):
    pass


class DBSessionError(Exception):
    pass


class EntryValidationError(Exception):
    def __init__(self, message: str, model_field: Optional[str] = None):
        self.message = message
        self.model_field = model_field


class DataStreamValidationError(Exception):
    line_idx: Optional[int]
    line_value: Optional[str]
    field_idx: Optional[int]
    field_name: Optional[str]
    value: Optional[Any]

    def __init__(
        self,
        *args: Any,
        line_idx: Optional[int] = None,
        line_value: Optional[str] = None,
        field_idx: Optional[int] = None,
        field_name: Optional[str] = None,
        value: Any = None,
    ):
        self.line_idx = line_idx
        self.line_value = line_value
        self.field_idx = field_idx
        self.field_name = field_name
        self.value = value
        super().__init__(*args)

    # TODO FIX: Move to error schemas
    @staticmethod
    def jsonify_value(val: Any) -> TJSONLike:
        if isinstance(val, bytes):
            return val.decode("utf-8", "backslashreplace")
        elif isinstance(val, (str, int, float)):
            return val
        elif val is None:
            return val
        else:
            # TODO CONSIDER: different variants for logging & error response
            return {"unexpected_value_type": str(type(val)), "str_val": str(val)}

    # TODO FIX: Replace with schema
    def get_context_dict(self) -> Dict[str, TJSONLike]:
        return dict(
            line_idx=self.jsonify_value(self.line_idx),
            line_value=self.jsonify_value(self.line_value),
            field_idx=self.jsonify_value(self.field_idx),
            field_name=self.jsonify_value(self.field_name),
            value=self.jsonify_value(self.value),
        )


class ReferencedUSEntryNotFound(DLBaseException):
    err_code = DLBaseException.err_code + ["REFERENCED_ENTRY_NOT_FOUND"]
    default_message = "Some referenced entries not found"


class ReferencedUSEntryAccessDenied(DLBaseException):
    err_code = DLBaseException.err_code + ["REFERENCED_ENTRY_ACCESS_DENIED"]
    default_message = "Some referenced entries cannot be loaded: access denied"


class ConnectionConfigurationError(DLBaseException):
    err_code = DLBaseException.err_code + ["CONNECTION_CONFIG"]


class SubselectNotAllowed(ConnectionConfigurationError):
    err_code = ConnectionConfigurationError.err_code + ["SUBSELECT_NOT_ALLOWED"]
    default_message = "Subquery source is disallowed in the connection settings"


class DashSQLNotAllowed(ConnectionConfigurationError):
    err_code = ConnectionConfigurationError.err_code + ["DASHSQL_NOT_ALLOWED"]
    default_message = "DashSQL API is disallowed in the connection settings"


class NotAvailableError(DLBaseException):
    err_code = DLBaseException.err_code + ["NOT_AVAILABLE"]
    default_message = "Not available"


class InvalidColumnError(DLBaseException):
    err_code = DLBaseException.err_code + ["INVALID_COLUMN"]


class PlatformPermissionRequired(DLBaseException):
    err_code = DLBaseException.err_code + ["PLATFORM_PERMISSION_REQUIRED"]


class EntityUsageNotAllowed(DLBaseException):
    err_code = DLBaseException.err_code + ["ENTITY_USAGE_NOT_ALLOWED"]
    default_message = "Some of the entities cannot be used in this context"


class USPermissionRequired(USReqException):
    err_code = USReqException.err_code + ["PERMISSION_REQUIRED"]

    entry_id: str
    permission: str

    def __init__(self, entry_id: str, permission: str):
        super().__init__()
        self.details.update(entry_id=entry_id, permission=permission)
        self.entry_id = entry_id
        self.permission = permission

    @property
    def message(self) -> str:
        return f"No permission {self.permission} for entry {self.entry_id}"


class DataSourceMigrationImpossible(DLBaseException):
    err_code = DLBaseException.err_code + ["DSRC_MIGRATION_IMPOSSIBLE"]


class RLSSubjectNotFound(DLBaseException):
    err_code = DLBaseException.err_code + ["RLS_SUBJECT_NOT_FOUND"]


class DashSQLError(DLBaseException):
    err_code = DLBaseException.err_code + ["DASHSQL"]

lib/dl_core/dl_core/exc.py
datalens-tech/datalens-backend
# TODO: Move to dl_core.components.* (separate into 2 files: field and result_schema)

from __future__ import annotations

import logging
import re
from typing import (
    Any,
    ClassVar,
    Collection,
    Dict,
    List,
    NamedTuple,
    Optional,
    Type,
    TypeVar,
    Union,
)

import attr

from dl_constants.enums import (
    AggregationFunction,
    CalcMode,
    FieldType,
    ManagedBy,
    ParameterValueConstraintType,
    TopLevelComponentId,
    UserDataType,
)
from dl_core.components.ids import FieldId
from dl_core.exc import FieldNotFound
from dl_core.values import BIValue


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class ParameterValueConstraint:
    type: ParameterValueConstraintType = attr.ib(default=ParameterValueConstraintType.all)

    def is_valid(self, value: Any) -> bool:
        if isinstance(value, BIValue):
            value = value.value
        try:
            return self._is_valid(value)
        except TypeError:
            return False

    def _is_valid(self, value: Any) -> bool:
        return True


@attr.s(frozen=True)
class RangeParameterValueConstraint(ParameterValueConstraint):
    min: Optional[BIValue] = attr.ib(default=None)
    max: Optional[BIValue] = attr.ib(default=None)
    type: ParameterValueConstraintType = attr.ib(default=ParameterValueConstraintType.range)

    def _is_valid(self, value: Any) -> bool:
        return (self.min is None or self.min.value <= value) and (self.max is None or value <= self.max.value)


@attr.s(frozen=True)
class SetParameterValueConstraint(ParameterValueConstraint):
    values: List[BIValue] = attr.ib(factory=list)
    type: ParameterValueConstraintType = attr.ib(default=ParameterValueConstraintType.set)

    def _is_valid(self, value: Any) -> bool:
        return any([v.value == value for v in self.values])


_CALC_SPEC_TV = TypeVar("_CALC_SPEC_TV", bound="CalculationSpec")


@attr.s(frozen=True)
class CalculationSpec:
    """Defines how the field's expression is constructed"""

    mode: ClassVar[CalcMode]  # Different enum value for each subclass

    def clone(self: _CALC_SPEC_TV, **kwargs: Any) -> _CALC_SPEC_TV:
        return attr.evolve(self, **kwargs)

    def depends_on(self, field: BIField) -> bool:
        return False


@attr.s(frozen=True)
class DirectCalculationSpec(CalculationSpec):
    """
    Describes direct fields.
    Such fields are created from data source (db table) columns
    and inherit some of their properties.
    """

    mode = CalcMode.direct

    # ID of the avatar (table) this field refers to
    avatar_id: Optional[str] = attr.ib(kw_only=True, default=None)
    # Name of the table column this field refers to
    # (name attribute of a raw_Schema column).
    # Use empty string (`''`) for non-direct fields.
    source: str = attr.ib(kw_only=True, default="")


@attr.s(frozen=True)
class FormulaCalculationSpec(CalculationSpec):
    """
    Field is defined by a formula.
    """

    mode = CalcMode.formula

    # The formula itself. Parsed and handled mostly by the dl_formula package.
    # In this formula other fields are referred to exclusively by title,
    # Use empty string (`''`) for non-formula fields.
    formula: str = attr.ib(kw_only=True, default="")

    # In this formula other fields are referred to exclusively by guid,
    # Use empty string (`''`) for non-formula fields.
    guid_formula: str = attr.ib(kw_only=True, default="")

    def depends_on(self, field: BIField) -> bool:
        return f"[{field.title}]" in self.formula or f"[{field.guid}]" in self.guid_formula


@attr.s(frozen=True)
class ParameterCalculationSpec(CalculationSpec):
    """
    Field is a parameter.
    The parameter value is specified in data API.
    After value substitution, the field acts as a formula consisting of a single literal.
    """

    mode = CalcMode.parameter

    # Default value of the parameter.
    default_value: Optional[BIValue] = attr.ib(kw_only=True, default=None)
    # Value constraint of the parameter
    # (defines the restrictions and origins of possible values).
    value_constraint: Optional[ParameterValueConstraint] = attr.ib(kw_only=True, default=None)


_CALCULATION_SPECS_BY_MODE = {
    CalcMode.direct.name: DirectCalculationSpec,
    CalcMode.formula.name: FormulaCalculationSpec,
    CalcMode.parameter.name: ParameterCalculationSpec,
}


def filter_calc_spec_kwargs(mode: str, kwargs: Dict[str, Any]) -> Dict[str, Any]:
    spec_cls = _CALCULATION_SPECS_BY_MODE[mode]
    field_keys = {f.name for f in attr.fields(spec_cls)}
    return {key: value for key, value in kwargs.items() if key in field_keys}


def create_calc_spec_from(kwargs: Dict[str, Any]) -> CalculationSpec:
    mode = kwargs["calc_mode"]
    spec_cls = _CALCULATION_SPECS_BY_MODE[mode.name]
    return spec_cls(**filter_calc_spec_kwargs(mode.name, kwargs))


def del_calc_spec_kwargs_from(
    kwargs: Dict[str, Any], spec_cls: Optional[Type[CalculationSpec]] = None
) -> Dict[str, Any]:
    spec_classes = _CALCULATION_SPECS_BY_MODE.values() if spec_cls is None else [spec_cls]
    keys = [f.name for s in spec_classes for f in attr.fields(s)] + ["calc_mode"]
    return {key: value for key, value in kwargs.items() if key not in keys}


class BIField(NamedTuple):  # TODO: Convert to attr.s
    """Dataset field, contained in dataset's result_schema"""

    # The ID of the field. The main means of identifying
    # or referring to the field.
    guid: str

    # A human-readable name of the field (used in UI and formulas).
    title: str

    # Describes how the field was created
    # (directly by the user or various parts of the backend)
    # and is basically a primitive implementation of access control.
    managed_by: ManagedBy

    # Flag indicates whether this field is always shown in UI or not.
    hidden: bool

    # A plain-text human-readable description of the field.
    # Not used for any internal purposes.
    description: str

    # Flag indicates whether the field contains any errors.
    valid: bool

    # The earliest identifiable data type of the field
    # in the expression construction pipeline.
    # Works a bit differently depending on the calc mode:
    # - direct: it corresponds to the user_type of the referenced raw_schema column;
    # - formula: automatically derived from the formula;
    # - parameter: defined by the user.
    initial_data_type: Optional[UserDataType]

    # redefines the data type in `initial_data_type`, is set by the user.
    # For parameter calc_mode, it is the same as `initial_data_type`.
    cast: Optional[UserDataType]

    # An explicitly set aggregation (via UI).
    # Value "none" corresponds to no aggregation.
    # Is not the final data type (see `data_type`)!
    aggregation: AggregationFunction

    # The data type automatically determined after the aggregation is applied
    data_type: Optional[UserDataType]

    # Flag indicates that the field is automatically aggregated
    # and an explicit aggregation (`aggregation` attribute) is not applicable.
    # This can be in one of the following cases:
    # - field is a formula that contains an aggregation
    # - field is in direct mode and references a column
    #   that is pre-aggregated in the data source
    has_auto_aggregation: bool

    # Flag that enables or disables the explicit aggregation in UI
    # for direct fields if their (pre)aggregation is controlled by the data source.
    # FIXME: Looks like this should be part of `DirectCalculationSpec`,
    #   but its usage is complicated and a bit unclear,
    #   so it requires some research and (probably) preliminary fixes
    lock_aggregation: bool

    # Type of field in terms of measure/dimension.
    # If the field contains an aggregation
    # (aggregation attribute, formula, or direct field with has_auto_aggregation),
    # then it is a measure. In all other cases it id a dimension.
    type: FieldType

    # Specified
    calc_spec: CalculationSpec

    @classmethod
    def make(
        cls,
        guid: FieldId,
        title: Optional[str],
        calc_spec: CalculationSpec,
        aggregation: Union[AggregationFunction, str, None] = None,
        type: Union[FieldType, str, None] = None,
        hidden: bool = False,
        description: str = "",
        cast: Union[UserDataType, str, None] = None,
        initial_data_type: Union[UserDataType, str, None] = None,
        data_type: Union[UserDataType, str, None] = None,
        valid: Optional[bool] = None,
        has_auto_aggregation: bool = False,
        lock_aggregation: bool = False,
        managed_by: Optional[ManagedBy] = None,
    ) -> BIField:
        """
        Normalize attributes and create a BIField object.
        """

        if isinstance(calc_spec, DirectCalculationSpec):
            title = title or calc_spec.source or ""
            calc_spec = calc_spec.clone(source=calc_spec.source or title or "")
        assert title is not None

        aggregation = AggregationFunction.normalize(aggregation)
        if aggregation is None:
            aggregation = AggregationFunction.none
        assert aggregation is not None

        type = FieldType.normalize(type)
        if type is None:
            type = FieldType.DIMENSION
        assert type is not None

        cast = UserDataType.normalize(cast)
        initial_data_type = UserDataType.normalize(initial_data_type)
        data_type = UserDataType.normalize(data_type)
        valid = valid if valid is not None else True
        managed_by = ManagedBy.normalize(managed_by) or ManagedBy.user

        return cls(
            title=title,
            guid=guid,
            aggregation=aggregation,
            hidden=hidden,
            description=description,
            type=type,
            cast=cast,
            initial_data_type=initial_data_type,
            data_type=data_type,
            valid=valid,
            has_auto_aggregation=has_auto_aggregation,
            lock_aggregation=lock_aggregation,
            managed_by=managed_by,
            calc_spec=calc_spec,
        )

    @property
    def calc_mode(self) -> CalcMode:
        return self.calc_spec.mode

    @property
    def avatar_id(self) -> Optional[str]:
        assert isinstance(self.calc_spec, DirectCalculationSpec)
        return self.calc_spec.avatar_id

    @property
    def source(self) -> str:
        assert isinstance(self.calc_spec, DirectCalculationSpec)
        return self.calc_spec.source

    @property
    def formula(self) -> str:
        assert isinstance(self.calc_spec, FormulaCalculationSpec)
        return self.calc_spec.formula

    @property
    def guid_formula(self) -> str:
        assert isinstance(self.calc_spec, FormulaCalculationSpec)
        return self.calc_spec.guid_formula

    @property
    def default_value(self) -> Optional[BIValue]:
        assert isinstance(self.calc_spec, ParameterCalculationSpec)
        return self.calc_spec.default_value

    @property
    def value_constraint(self) -> Optional[ParameterValueConstraint]:
        assert isinstance(self.calc_spec, ParameterCalculationSpec)
        return self.calc_spec.value_constraint

    @property
    def autoaggregated(self) -> bool:
        return (
            self.has_auto_aggregation
            # TODO: Remove this second condition after dataset migration #11 is applied
            or self.type == FieldType.MEASURE
            and self.aggregation == AggregationFunction.none
        )

    @property
    def aggregation_locked(self) -> bool:
        return (
            # TODO: Remove the first condition after aggregation handling is enabled for formula fields in bi-api
            self.calc_mode == CalcMode.formula
            or self.calc_mode == CalcMode.parameter
            or self.has_auto_aggregation
            or self.lock_aggregation
        )

    @staticmethod
    def rename_in_formula(formula: str, key_map: Dict[str, str]) -> str:
        found_keys = FIELD_RE.findall(formula)

        for key in found_keys:
            try:
                value = key_map[key]
            except KeyError:
                LOGGER.warning("Unknown field: %s", key)
                continue

            formula = formula.replace("[{}]".format(key), "[{}]".format(value))

        return formula

    def depends_on(self, field: BIField) -> bool:
        return self.calc_spec.depends_on(field)

    def clone(self, **kwargs: Any) -> BIField:
        if "calc_spec" not in kwargs:
            new_calc_spec = None
            mode = kwargs.get("calc_mode", self.calc_mode)
            if mode != self.calc_mode:
                new_calc_spec = create_calc_spec_from(kwargs)
            else:
                calc_spec_kwargs = filter_calc_spec_kwargs(mode.name, kwargs)
                if calc_spec_kwargs:
                    new_calc_spec = self.calc_spec.clone(**calc_spec_kwargs)
            kwargs = del_calc_spec_kwargs_from(kwargs)
            if new_calc_spec is not None:
                kwargs["calc_spec"] = new_calc_spec

        return self._replace(**kwargs)


FIELD_RE = re.compile(r"\[([^]]*)\]")


@attr.s
class ResultSchema:
    fields: List[BIField] = attr.ib(default=attr.Factory(list))
    _guid_cache: Dict[FieldId, BIField] = attr.ib(init=False, eq=False)
    _title_cache: Dict[str, BIField] = attr.ib(init=False, eq=False)
    valid: bool = attr.ib(default=True)

    def __attrs_post_init__(self):  # type: ignore  # TODO: fix
        self._guid_cache = {}
        self._title_cache = {}
        self.reload_caches()

    @property
    def id(self) -> str:
        return TopLevelComponentId.result_schema.value

    def clear_caches(self) -> None:
        self._guid_cache.clear()
        self._title_cache.clear()

    def reload_caches(self) -> None:
        self.clear_caches()
        for field in self.fields:
            self._guid_cache[field.guid] = field
            self._title_cache[field.title] = field

    def by_guid(self, guid: FieldId) -> BIField:
        """Find field by guid"""
        if not self._guid_cache:
            self.reload_caches()
        try:
            return self._guid_cache[guid]
        except KeyError:
            raise FieldNotFound(f"Unknown field {guid}")

    def by_title(self, title: str) -> BIField:
        """Find field by title"""
        if not self._title_cache:
            self.reload_caches()
        try:
            return self._title_cache[title]
        except KeyError:
            raise FieldNotFound(f"Unknown field {title}")

    @property
    def titles_to_guids(self) -> Dict[str, FieldId]:
        return {f.title: f.guid for f in self.fields}

    @property
    def guids_to_titles(self) -> Dict[FieldId, str]:
        return {f.guid: f.title for f in self.fields}

    def __iter__(self):  # type: ignore  # TODO: fix
        return self.fields.__iter__()

    def __bool__(self):  # type: ignore  # TODO: fix
        return self.fields.__iter__()

    def __len__(self):  # type: ignore  # TODO: fix
        return self.fields.__len__()

    def __getitem__(self, ind):  # type: ignore  # TODO: fix
        return self.fields.__getitem__(ind)

    def add(self, field: BIField, idx: Optional[int]) -> None:
        if idx is not None:
            self.fields.insert(idx, field)
        else:
            self.fields.append(field)
        self.clear_caches()

    def remove(self, field_id: FieldId) -> None:
        self.fields.remove(self.by_guid(field_id))
        self.clear_caches()

    def update_field(self, idx: int, field: BIField) -> None:
        self.fields[idx] = field
        self.clear_caches()

    def index(self, field_id: FieldId) -> int:
        return next((idx for idx, f in enumerate(self.fields) if f.guid == field_id))

    def remove_multiple(self, field_ids: Collection[str]) -> None:
        fields_to_remove = [self.by_guid(field_id) for field_id in field_ids]
        for field in fields_to_remove:
            self.fields.remove(field)
        self.clear_caches()

lib/dl_core/dl_core/fields.py
datalens-tech/datalens-backend
"""
Everything common for `flask`s.
"""

lib/dl_core/dl_core/flask_utils/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    Optional,
)

import attr
import flask

from dl_api_commons.flask.middlewares.commit_rci_middleware import ReqCtxInfoMiddleware

from ..services_registry.sr_factories import SRFactory


if TYPE_CHECKING:
    from dl_core.services_registry import ServicesRegistry  # noqa


LOGGER = logging.getLogger(__name__)


class NoServiceRegistryForRequest(Exception):
    pass


@attr.s
class ServicesRegistryMiddleware:
    services_registry_factory: SRFactory = attr.ib()

    def set_up(self, app: flask.Flask) -> None:
        app.before_request(self.bind_services_registry_to_request)
        app.teardown_request(self.cleanup_request_services_registry)

    def cleanup_request_services_registry(self, _: Optional[BaseException] = None) -> None:
        try:
            services_registry = self.get_request_services_registry()
        except NoServiceRegistryForRequest:
            LOGGER.debug("Service registry was not created. Close procedure does not required.")
        else:
            try:
                LOGGER.info("Closing services registry...")
                services_registry.close()
            except Exception:  # noqa
                LOGGER.exception("Error during services registry cleanup")
            LOGGER.info("Closed services registry.")

    def bind_services_registry_to_request(self) -> None:
        rci = ReqCtxInfoMiddleware.get_request_context_info()
        flask.g.bi_services_registry = self.services_registry_factory.make_service_registry(rci)

    @classmethod
    def get_request_services_registry(cls) -> "ServicesRegistry":
        try:
            services_registry = flask.g.bi_services_registry
        except AttributeError:
            raise NoServiceRegistryForRequest()
        else:
            if services_registry is None:
                raise NoServiceRegistryForRequest()
            return services_registry

    @classmethod
    def enabled_for_request(cls) -> bool:
        return "bi_services_registry" in flask.g

lib/dl_core/dl_core/flask_utils/services_registry_middleware.py
datalens-tech/datalens-backend
from typing import Optional

import attr
import flask

from dl_api_commons.base_models import (
    NoAuthData,
    TenantCommon,
    TenantDef,
)
from dl_api_commons.flask.middlewares.commit_rci_middleware import ReqCtxInfoMiddleware


@attr.s(frozen=True, auto_attribs=True)
class TrustAuthService:
    """
    Service that unconditionally injects user/tenant info into request context.
    Should be used only in tests
    """

    fake_user_id: Optional[str] = None
    fake_user_name: Optional[str] = None
    fake_tenant: Optional[TenantDef] = None

    def _before_request(self) -> None:
        fake_user_id = self.fake_user_id
        fake_user_name = self.fake_user_name
        fake_tenant = self.fake_tenant

        temp_rci = ReqCtxInfoMiddleware.get_temp_rci().clone(
            auth_data=NoAuthData(),
            tenant=TenantCommon() if fake_tenant is None else fake_tenant,
        )
        if fake_user_id is not None:
            temp_rci = temp_rci.clone(user_id=fake_user_id)
        if fake_user_name is not None:
            temp_rci = temp_rci.clone(user_name=fake_user_name)

        ReqCtxInfoMiddleware.replace_temp_rci(temp_rci)

    def set_up(self, app: flask.Flask) -> None:
        app.before_request(self._before_request)

lib/dl_core/dl_core/flask_utils/trust_auth.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import Optional

import attr
import flask

from dl_api_commons.flask.middlewares.commit_rci_middleware import ReqCtxInfoMiddleware
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_constants.enums import USAuthMode
from dl_core.flask_utils.services_registry_middleware import ServicesRegistryMiddleware
from dl_core.us_manager.factory import USMFactory
from dl_core.us_manager.us_manager_sync import SyncUSManager


LOGGER = logging.getLogger(__name__)


# TODO FIX: Add tests
@attr.s
class USManagerFlaskMiddleware:
    us_base_url: str = attr.ib()
    us_auth_mode: USAuthMode = attr.ib(
        validator=attr.validators.in_(
            (
                USAuthMode.regular,
                USAuthMode.master,
            )
        )
    )
    crypto_keys_config: CryptoKeysConfig = attr.ib(repr=False)
    us_master_token: Optional[str] = attr.ib(default=None, repr=False)
    us_public_token: Optional[str] = attr.ib(default=None, repr=False)

    _usm_factory: USMFactory = attr.ib(init=False, default=None)

    def __attrs_post_init__(self) -> None:
        self._usm_factory = USMFactory(
            us_base_url=self.us_base_url,
            crypto_keys_config=self.crypto_keys_config,
            us_master_token=self.us_master_token,
            us_public_token=self.us_public_token,
        )

    def bind_us_managers_to_request(self) -> None:
        bi_context = ReqCtxInfoMiddleware.get_request_context_info()
        services_registry = (
            ServicesRegistryMiddleware.get_request_services_registry()
            if ServicesRegistryMiddleware.enabled_for_request()
            else None
        )
        assert services_registry is not None, "Service registry must be enabled!"

        if bi_context.user_id is None:
            LOGGER.info("User US manager will not be created due to no user info in RCI")
        else:
            LOGGER.info("Creating user US manager")
            usm: SyncUSManager

            if self.us_auth_mode == USAuthMode.regular:
                usm = self._usm_factory.get_regular_sync_usm(rci=bi_context, services_registry=services_registry)
            elif self.us_auth_mode == USAuthMode.master:
                usm = self._usm_factory.get_master_sync_usm(rci=bi_context, services_registry=services_registry)
            else:
                raise AssertionError(f"Mode is not supported by USManagerFlaskMiddleware: {self.us_auth_mode!r}")

            flask.g.us_manager = usm

        if self.us_master_token is not None:
            LOGGER.info("Creating service US manager")
            flask.g.service_us_manager = self._usm_factory.get_master_sync_usm(
                rci=bi_context,
                services_registry=services_registry,
            )
        else:
            LOGGER.info("US master token was not provided. Service USM will not be created")
            flask.g.service_us_manager = None

    def set_up(self, app: flask.Flask) -> USManagerFlaskMiddleware:
        app.before_request(self.bind_us_managers_to_request)
        return self

    def run(self):  # type: ignore  # TODO: fix
        return self.bind_us_managers_to_request()

    @classmethod
    def get_request_us_manager(cls) -> SyncUSManager:
        return flask.g.us_manager

    @classmethod
    def get_request_service_us_manager(cls) -> SyncUSManager:
        usm = flask.g.service_us_manager
        if usm is not None:
            return flask.g.service_us_manager
        else:
            raise ValueError("Service USM was not created for current request")

lib/dl_core/dl_core/flask_utils/us_manager_middleware.py
datalens-tech/datalens-backend
import os

import attr

import dl_core as package
from dl_i18n.localizer_base import Translatable as BaseTranslatable
from dl_i18n.localizer_base import TranslationConfig


DOMAIN = f"{package.__name__}"
CONFIGS = [
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../locales")),
        domain=DOMAIN,
        locale="en",
    ),
    TranslationConfig(
        path=os.path.relpath(os.path.join(os.path.dirname(__file__), "../locales")),
        domain=DOMAIN,
        locale="ru",
    ),
]


@attr.s
class Translatable(BaseTranslatable):
    domain: str = attr.ib(default=DOMAIN)

lib/dl_core/dl_core/i18n/localizer.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Generic,
    Type,
    TypeVar,
)

import attr

from dl_core.us_entry import USEntry


if TYPE_CHECKING:
    from dl_core.services_registry.top_level import ServicesRegistry
    from dl_core.us_manager.us_manager import USManagerBase


_US_ENTRY_TV = TypeVar("_US_ENTRY_TV", bound=USEntry)


@attr.s
class EntryLifecycleManager(abc.ABC, Generic[_US_ENTRY_TV]):
    ENTRY_CLS: ClassVar[Type[_US_ENTRY_TV]]

    _entry: _US_ENTRY_TV = attr.ib(kw_only=True)
    _us_manager: USManagerBase = attr.ib(kw_only=True)
    _service_registry: ServicesRegistry = attr.ib(kw_only=True)

    @property
    def entry(self) -> _US_ENTRY_TV:
        assert isinstance(self._entry, self.ENTRY_CLS)
        return self._entry

    def pre_save_hook(self) -> None:
        pass

    def post_save_hook(self) -> None:
        pass

    def post_copy_hook(self) -> None:
        pass

    def pre_delete_hook(self) -> None:
        pass

    def post_delete_hook(self) -> None:
        pass

    async def post_init_async_hook(self) -> None:
        pass

    async def post_exec_async_hook(self) -> None:
        pass

    def collect_links(self) -> dict[str, str]:
        return {}

lib/dl_core/dl_core/lifecycle/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging

from dl_constants.enums import DataSourceRole
from dl_core.components.editor import DatasetComponentEditor
from dl_core.data_source_spec.collection import DataSourceCollectionSpec
from dl_core.lifecycle.base import EntryLifecycleManager
from dl_core.us_dataset import Dataset


LOGGER = logging.getLogger(__name__)


class DatasetLifecycleManager(EntryLifecycleManager[Dataset]):
    ENTRY_CLS = Dataset

    def pre_save_hook(self) -> None:
        super().pre_save_hook()

        self.entry._dump_rls()

        try:
            self.entry.links = self.collect_links()
        except Exception:
            LOGGER.exception("Failed to collect links for ds %s", self.entry.uuid)

    def post_copy_hook(self) -> None:
        super().post_copy_hook()

        own_collection_specs = [
            dsrc_coll_spec
            for dsrc_coll_spec in self.entry.data.source_collections
            if isinstance(dsrc_coll_spec, DataSourceCollectionSpec)
        ]

        ds_editor = DatasetComponentEditor(dataset=self.entry)
        for dsrc_coll_spec in own_collection_specs:
            ds_editor.remove_data_source(
                source_id=dsrc_coll_spec.id, role=DataSourceRole.materialization, delete_mat_table=False
            )
            ds_editor.remove_data_source(
                source_id=dsrc_coll_spec.id, role=DataSourceRole.sample, delete_mat_table=False
            )

    def collect_links(self) -> dict[str, str]:
        links = {}
        for dsrc_coll_spec in self.entry.data.source_collections:
            links.update(dsrc_coll_spec.collect_links())

        return links

lib/dl_core/dl_core/lifecycle/dataset.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import singledispatchmethod
from typing import TYPE_CHECKING

from dl_core.connectors.base.lifecycle import ConnectionLifecycleManager
from dl_core.lifecycle.base import EntryLifecycleManager
from dl_core.lifecycle.dataset import DatasetLifecycleManager
from dl_core.lifecycle.factory_base import EntryLifecycleManagerFactoryBase
from dl_core.us_connection import get_lifecycle_manager_cls
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_dataset import Dataset


if TYPE_CHECKING:
    from dl_core.services_registry import ServicesRegistry
    from dl_core.us_entry import USEntry
    from dl_core.us_manager.us_manager import USManagerBase


class DummyEntryLifecycleManagerFactory(EntryLifecycleManagerFactoryBase):
    def get_lifecycle_manager(
        self,
        entry: USEntry,
        us_manager: USManagerBase,
        service_registry: ServicesRegistry,
    ) -> EntryLifecycleManager:
        return EntryLifecycleManager(us_manager=us_manager, entry=entry, service_registry=service_registry)


class DefaultEntryLifecycleManagerFactory(EntryLifecycleManagerFactoryBase):
    def get_lifecycle_manager(
        self,
        entry: USEntry,
        us_manager: USManagerBase,
        service_registry: ServicesRegistry,
    ) -> EntryLifecycleManager:
        return self._get_lifecycle_manager(entry, us_manager=us_manager, service_registry=service_registry)

    @singledispatchmethod
    def _get_lifecycle_manager(
        self,
        entry: USEntry,
        us_manager: USManagerBase,
        service_registry: ServicesRegistry,
    ) -> EntryLifecycleManager:
        return EntryLifecycleManager(us_manager=us_manager, entry=entry, service_registry=service_registry)

    @_get_lifecycle_manager.register(Dataset)
    def _get_dataset_lifecycle_manager(
        self,
        entry: Dataset,
        us_manager: USManagerBase,
        service_registry: ServicesRegistry,
    ) -> DatasetLifecycleManager:
        return DatasetLifecycleManager(us_manager=us_manager, entry=entry, service_registry=service_registry)

    @_get_lifecycle_manager.register(ConnectionBase)
    def _get_connection_lifecycle_manager(
        self,
        entry: ConnectionBase,
        us_manager: USManagerBase,
        service_registry: ServicesRegistry,
    ) -> ConnectionLifecycleManager:
        lifecycle_manager_cls = get_lifecycle_manager_cls(conn_type=entry.conn_type)
        return lifecycle_manager_cls(us_manager=us_manager, entry=entry, service_registry=service_registry)

lib/dl_core/dl_core/lifecycle/factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import TYPE_CHECKING


if TYPE_CHECKING:
    from dl_core.lifecycle.base import EntryLifecycleManager
    from dl_core.services_registry import ServicesRegistry
    from dl_core.us_entry import USEntry
    from dl_core.us_manager.us_manager import USManagerBase


class EntryLifecycleManagerFactoryBase(abc.ABC):
    @abc.abstractmethod
    def get_lifecycle_manager(
        self,
        entry: USEntry,
        us_manager: USManagerBase,
        service_registry: ServicesRegistry,
    ) -> EntryLifecycleManager:
        pass

lib/dl_core/dl_core/lifecycle/factory_base.py
datalens-tech/datalens-backend
from typing import (
    Collection,
    Optional,
)

import attr

from dl_core.core_connectors import (
    load_all_connectors,
    register_all_connectors,
)
from dl_core.core_data_processors import register_all_data_processor_plugins


@attr.s(frozen=True)
class CoreLibraryConfig:
    # Whitelist of connector entrypoints to be loaded
    core_connector_ep_names: Optional[Collection[str]] = attr.ib(kw_only=True, default=None)


def preload_bi_core() -> None:
    """Loads all entrypoint connectors without registering them"""

    load_all_connectors()


def load_core_lib(core_lib_config: CoreLibraryConfig = CoreLibraryConfig()) -> None:
    """Initialize the library"""
    register_all_connectors(connector_ep_names=core_lib_config.core_connector_ep_names)
    register_all_data_processor_plugins()

lib/dl_core/dl_core/loader.py
datalens-tech/datalens-backend
""" A common logging configuration makers """

from __future__ import annotations

import logging
import os
from typing import (
    Callable,
    Optional,
    Sequence,
)

import jaeger_client
from opentracing.scope_managers.contextvars import ContextVarsScopeManager
import statcommons.log_config
import statcommons.logs

from dl_api_commons.logging_config import add_log_context
import dl_app_tools.log.context as log_context
from dl_app_tools.log.format import JsonFormatter


# Same check as in `dl_app_tools.log.format.JsonFormatter`

IS_DEPLOY = "DEPLOY_BOX_ID" in os.environ

LOGMUTATORS = statcommons.logs.LOGMUTATORS
HOSTNAME = statcommons.log_config.HOSTNAME
ENV_CONTEXT = statcommons.log_config.ENV_CONTEXT


class FastlogsFilter(logging.Filter):
    def filter(self, record):  # type: ignore  # TODO: fix
        event_code = getattr(record, "event_code", None)
        if event_code:
            if isinstance(event_code, str) and event_code.startswith("_"):
                pass
            else:
                return True

        # ...

        return False


_LOG_FORMATTER = JsonFormatter()
_JSON_FORMATTER = statcommons.logs.JsonExtFormatter()


class StdoutFormatter(logging.Formatter):
    """
    Write logs in YDeploy format in YDeploy,
    and as top-level JSON otherwise.
    """

    def format(self, record: logging.LogRecord) -> str:
        if IS_DEPLOY:
            return _LOG_FORMATTER.format(record)
        return _JSON_FORMATTER.format(record)


def _make_logging_config(
    for_development: bool,
    logcfg_processors: Optional[Sequence[Callable]] = None,
) -> dict:
    """
    ...

    `logcfg_processors`: see description in `configure_logging`.

    `for_development`: if True - will apply simple plain text stderr/stdout config without any integrations.
      Should be used only for running tests
    """
    if for_development:
        base = statcommons.log_config.DEV_LOGGING_CONFIG
        common_handlers = []
        logging_config = {
            **base,
            "loggers": {
                **base.get("loggers", {}),  # type: ignore  # TODO: fix
                "dl_core": {
                    "handlers": ["stream"],
                    "level": "INFO",
                    "propagate": False,
                },
            },
            "root": {
                **base.get("root", {}),  # type: ignore  # TODO: fix
                "level": "INFO",
            },
        }

    else:
        base = statcommons.log_config.BASE_LOGGING_CONFIG
        # Handlers for root and for every non-propagated logger
        common_handlers = base["root"]["handlers"]  # type: ignore  # TODO: fix
        # ^ ['debug_log', 'fast_log', 'event_log']

        default_handlers = ["stream"] + common_handlers  # everything to stdout
        logging_config = {
            **base,
            "filters": {
                **base.get("filters", {}),  # type: ignore  # TODO: fix
                "fastlogs": {"()": "dl_core.logging_config.FastlogsFilter"},
            },
            "formatters": {
                **base.get("formatters", {}),  # type: ignore  # TODO: fix
                "json": {"()": "dl_core.logging_config.StdoutFormatter"},
            },
            "handlers": {
                **(base.get("handlers") or {}),  # type: ignore  # TODO: fix
            },
            "loggers": {
                **(base.get("loggers") or {}),  # type: ignore  # TODO: fix
                # Set minimal level to some unhelpful libraries' logging:
                "jaeger_tracing": {"level": "WARNING", "propagate": False, "handlers": default_handlers},
                "asyncio": {"level": "INFO", "propagate": False, "handlers": default_handlers},
                "kikimr": {"level": "WARNING", "propagate": False, "handlers": default_handlers},
                "botocore": {"level": "INFO", "propagate": False, "handlers": default_handlers},
                "ydb": {"level": "WARNING", "propagate": False, "handlers": default_handlers},
            },
            "root": {
                **(base.get("root") or {}),  # type: ignore  # TODO: fix
                "handlers": default_handlers,
                "level": "DEBUG",
            },
        }

    for logcfg_process in logcfg_processors or ():
        logging_config = logcfg_process(
            logging_config,
            common_handlers=common_handlers,
        )

    return logging_config


# TODO FIX: Remove after all tests will be refactored to use unscoped log context
def add_log_context_scoped(record):  # type: ignore  # TODO: fix
    context = log_context.get_log_context()
    record.log_context = context


def update_tags(record):  # type: ignore  # TODO: fix
    current = getattr(record, "tags", None)
    if current is None:
        current = {}
        record.tags = current
    # NOTE: this will add a copy of request_id in the other logs (stdout / files);
    # TODO: transition to sentry_sdk, make a better common solution for pulling
    # data from record into sentry tags.
    current["request_id"] = getattr(record, "request_id", "unkn")


def logcfg_process_enable_handler(logger_name, handler_name="stream_info"):  # type: ignore  # TODO: fix
    """
    Make a `logcfg_process` mapper
    that enables a specified handler for the specified logger
    (in addition to the common handlers).

    Useful for enabling INFO or DEBUG logs.
    """

    def logcfg_process(cfg, common_handlers, **kwargs):  # type: ignore  # TODO: fix
        if handler_name in cfg["handlers"]:
            cfg["loggers"][logger_name] = dict(
                handlers=[handler_name] + common_handlers,
                propagate=False,
                level="DEBUG",
            )
        return cfg

    return logcfg_process


def logcfg_process_stream_human_readable(cfg, common_handlers, **kwargs):  # type: ignore  # TODO: fix  # noqa
    """
    Set human-readable logs to stdout/stderr.
    For maintenance procedures/migrations that are launched manually from shell.
    """
    if "verbose" in cfg.get("formatters", {}):
        for handler_name, handler in cfg.get("handlers", {}).items():
            if handler_name.startswith("stream_") or handler_name == "stream":
                handler.update(
                    formatter="verbose",
                    level="INFO",
                )

    return cfg


def setup_jaeger_client(service_name: str):  # type: ignore  # TODO: fix
    config = jaeger_client.Config(
        config={  # usually read from some yaml config
            "sampler": {
                "type": "const",
                "param": 1,
            },
            "logging": True,
        },
        service_name=service_name,
        scope_manager=ContextVarsScopeManager(),
        validate=True,
    )
    config.initialize_tracer()


def configure_logging(  # type: ignore  # TODO: fix
    app_name,
    for_development: Optional[bool] = None,
    app_prefix: Optional[str] = None,
    logcfg_processors: Optional[Sequence[Callable]] = None,
    use_jaeger_tracer: bool = False,
    jaeger_service_name: Optional[str] = None,
) -> None:
    """
    Make sure the global logging state is configured.

    Mostly idempotent but does some checks to ensure the configuration has not changed.

    `app_prefix` is not currently used; see `request_id_app_prefix` instead.

    `logcfg_processors`: convenient (but dangerous) processing of the logging config.
    Iterable of Callables `(log_cfg, **context) -> log_cfg`.
    Context includes `common_handlers`.
    Example: see `logcfg_process_enable_handler`.
    """
    if for_development is None:
        for_development = os.environ.get("DEV_LOGGING", "0").lower() in ("1", "true")

    cfg = _make_logging_config(
        for_development=for_development,
        logcfg_processors=logcfg_processors,
    )
    statcommons.log_config.configure_logging(app_name=app_name, cfg=cfg)
    LOGMUTATORS.add_mutator("log_context", add_log_context)
    LOGMUTATORS.add_mutator("update_tags", update_tags)

    if use_jaeger_tracer:
        effective_service_name = app_name if jaeger_service_name is None else jaeger_service_name
        setup_jaeger_client(effective_service_name)


def hook_configure_logging(app, *args, **kwargs):  # type: ignore  # TODO: fix
    """
    Try to configure logging in uwsgi `postfork` if possible,
    but ensure it is configured in `before_first_request` (flask app).
    """
    try:
        import uwsgidecorators  # type: ignore  # TODO: fix  # noqa
    except Exception:
        pass
    else:

        @uwsgidecorators.postfork
        def _init_logging_in_uwsgi_postfork():  # type: ignore  # TODO: fix
            configure_logging(*args, **kwargs)

    @app.before_first_request
    def _init_logging_in_before_first_request():  # type: ignore  # TODO: fix
        configure_logging(*args, **kwargs)

lib/dl_core/dl_core/logging_config.py
datalens-tech/datalens-backend
from __future__ import annotations

import base64
from typing import (
    Any,
    FrozenSet,
    Iterable,
    Mapping,
    Optional,
)

from marshmallow import fields
import marshmallow.utils as ma_utils


class FrozenSetField(fields.List):
    def _deserialize(self, value: Any, attr: Optional[str], data: Optional[Mapping[str, Any]], **kwargs: Any) -> FrozenSet[Any]:  # type: ignore
        return frozenset(super()._deserialize(value, attr, data, **kwargs))

    def _serialize(self, value: Any, attr: Optional[str], obj: Any, **kwargs: Any) -> Optional[list[Any]]:
        do_sort = self.metadata.get("sort_output", False)
        if do_sort and value is not None:
            value = sorted(value)
        return super()._serialize(value, attr, obj, **kwargs)


class ErrorCodeField(fields.Field):
    def __init__(self, *args: Any, prefix: Optional[Iterable[str]] = None, **kwargs: Any):
        self._prefix = list(prefix or [])
        super().__init__(*args, **kwargs)

    def _serialize(self, value: Any, attr: Optional[str], obj: Any, **kwargs: Any) -> str:
        return ".".join(self._prefix + list(value))

    def _deserialize(
        self, value: Any, attr: Optional[str], data: Optional[Mapping[str, Any]], **kwargs: Any
    ) -> list[str]:
        parts = value.split(".")
        if parts[: len(self._prefix)] == self._prefix:
            parts = parts[len(self._prefix) :]
        return parts


class OnOffField(fields.Field):
    ON = "on"
    OFF = "off"

    default_error_messages = {"invalid": "Not a valid boolean."}

    def _serialize(self, value: Any, attr: Optional[str], obj: Any, **kwargs: Any) -> str:
        if value:
            return self.ON

        return self.OFF

    def _deserialize(self, value: Any, attr: Optional[str], data: Optional[Mapping[str, Any]], **kwargs: Any) -> bool:
        if value == self.ON:
            return True
        if value == self.OFF:
            return False

        raise self.make_error("invalid", input=value)


class Base64StringField(fields.Field):
    default_error_messages = {
        "invalid": "Not a valid string.",
        "invalid_utf8": "Not a valid utf-8 string.",
        "invalid_format": "Not a valid file data string",
    }

    def _deserialize(
        self, value: Any, attr: Optional[str], data: Optional[Mapping[str, Any]], **kwargs: Any
    ) -> str | None:
        if isinstance(value, bytes):
            value = value.decode(encoding="utf-8")

        if not isinstance(value, str):
            raise self.make_error("invalid")

        if value == "":
            return None

        # stripping metadata
        try:
            value = value.split(";base64,", 1)[1]
        except IndexError as exc:
            raise self.make_error("invalid_format") from exc

        value = base64.b64decode(value)
        try:
            return ma_utils.ensure_text_type(value)
        except UnicodeDecodeError as error:
            raise self.make_error("invalid_utf8") from error

lib/dl_core/dl_core/marshmallow.py
datalens-tech/datalens-backend
# TODO: Move to dl_core.components.*

from __future__ import annotations

from typing import List

import attr

from dl_constants.enums import (
    BinaryJoinOperator,
    ConditionPartCalcMode,
    JoinConditionType,
    JoinType,
    ManagedBy,
)
from dl_core.components.ids import (
    AvatarId,
    FieldId,
    RelationId,
    SourceId,
)


@attr.s(frozen=True)
class ConditionPart:
    calc_mode: ConditionPartCalcMode = attr.ib()


@attr.s(frozen=True)
class ConditionPartDirect(ConditionPart):
    """Condition part that references a column from avatar's raw schema"""

    source: str = attr.ib()
    calc_mode: ConditionPartCalcMode = attr.ib(default=ConditionPartCalcMode.direct)


@attr.s(frozen=True)
class ConditionPartFormula(ConditionPart):  # TODO: Remove
    """Condition part defined by a formula that references avatar's raw schema columns"""

    formula: str = attr.ib()
    calc_mode: ConditionPartCalcMode = attr.ib(default=ConditionPartCalcMode.formula)


@attr.s(frozen=True)
class ConditionPartResultField(ConditionPart):
    """Condition part that references a field from dataset's result schema"""

    field_id: FieldId = attr.ib()
    calc_mode: ConditionPartCalcMode = attr.ib(default=ConditionPartCalcMode.result_field)


@attr.s(frozen=True)
class JoinCondition:
    pass


@attr.s(frozen=True)
class BinaryCondition(JoinCondition):
    operator: BinaryJoinOperator = attr.ib()
    left_part: ConditionPart = attr.ib()
    right_part: ConditionPart = attr.ib()
    condition_type: JoinConditionType = attr.ib(default=JoinConditionType.binary)


@attr.s(slots=True)
class SourceAvatar:
    id: AvatarId = attr.ib()
    source_id: SourceId = attr.ib()
    title: str = attr.ib()
    is_root: bool = attr.ib()
    managed_by: ManagedBy = attr.ib(default=ManagedBy.user)
    valid: bool = attr.ib(default=True)


@attr.s(slots=True)
class AvatarRelation:
    id: RelationId = attr.ib()
    left_avatar_id: AvatarId = attr.ib()
    right_avatar_id: AvatarId = attr.ib()
    conditions: List[BinaryCondition] = attr.ib(factory=list)
    join_type: JoinType = attr.ib(default=JoinType.inner)
    managed_by: ManagedBy = attr.ib(default=ManagedBy.user)
    valid: bool = attr.ib(default=True)

lib/dl_core/dl_core/multisource.py
datalens-tech/datalens-backend
from __future__ import annotations

import flask


def _ping_handler():  # type: ignore  # TODO: fix
    return ""


def _ping_handler_hax():  # type: ignore  # TODO: fix
    if flask.request.method == "GET" and flask.request.path.rstrip("/") == "/ping":
        return _ping_handler()
    return None  # flask: fall through


def register_ping_handler_hax(app):  # type: ignore  # TODO: fix
    # HAX to avoid bumping into inappropriate authorization added in 'before_request'.
    app.before_request(_ping_handler_hax)

lib/dl_core/dl_core/ping.py
datalens-tech/datalens-backend
"""
Prepare profiling stack for flask application
"""

from __future__ import annotations

from flask import (
    Flask,
    request,
)

from dl_api_commons import headers
from dl_app_tools.log import context
from dl_app_tools.profiling_base import GenericProfiler


def set_up(app: Flask, accept_outer_stages=False, **kwargs):  # type: ignore  # TODO: fix
    def pre_log_profiling_stack():  # type: ignore  # TODO: fix
        if accept_outer_stages:
            profiling_stack = request.headers.get(headers.INTERNAL_HEADER_PROFILING_STACK)
            GenericProfiler.reset_all(profiling_stack if profiling_stack else None)
        else:
            GenericProfiler.reset_all()

    def post_cleanup_context(exception=None):  # type: ignore  # TODO: fix
        context.pop_from_context("folder_id")

    app.before_request(pre_log_profiling_stack)
    app.teardown_request(post_cleanup_context)

lib/dl_core/dl_core/profiling_middleware.py
datalens-tech/datalens-backend
from __future__ import annotations

from itertools import chain
from typing import (
    List,
    Optional,
    Sequence,
    Set,
    Union,
)

import attr
from sqlalchemy.sql.elements import TextClause
from sqlalchemy.sql.selectable import (
    FromClause,
    Select,
)

from dl_constants.enums import UserDataType
from dl_core.query.expression import (
    ExpressionCtx,
    OrderByExpressionCtx,
)


SqlSourceType = Union[FromClause, TextClause]


@attr.s(frozen=True, auto_attribs=True)
class BIQuery:
    select_expressions: Sequence[ExpressionCtx]
    group_by_expressions: Sequence[ExpressionCtx] = attr.ib(factory=tuple)
    order_by_expressions: Sequence[OrderByExpressionCtx] = attr.ib(factory=tuple)
    dimension_filters: Sequence[ExpressionCtx] = attr.ib(factory=tuple)
    measure_filters: Sequence[ExpressionCtx] = attr.ib(factory=tuple)
    distinct: bool = False
    limit: Optional[int] = None
    offset: Optional[int] = None

    def get_required_avatar_ids(self) -> Set[str]:
        """Collect source avatar references from all expressions in the query."""

        return set(
            chain.from_iterable(
                [
                    ctx.avatar_ids
                    for expr_list in (
                        self.select_expressions,
                        self.group_by_expressions,
                        self.order_by_expressions,
                        self.dimension_filters,
                        self.measure_filters,
                    )
                    if expr_list
                    for ctx in expr_list
                    if isinstance(ctx, ExpressionCtx) and ctx.avatar_ids
                ]
            )
        )


@attr.s
class QueryAndResultInfo:
    query: Select = attr.ib(kw_only=True)
    user_types: List[UserDataType] = attr.ib(kw_only=True)
    col_names: List[str] = attr.ib(kw_only=True)

lib/dl_core/dl_core/query/bi_query.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
    Sequence,
    TypeVar,
)

import attr

from dl_constants.enums import OrderDirection


if TYPE_CHECKING:
    from sqlalchemy.sql.elements import ClauseElement

    from dl_constants.enums import (
        JoinType,
        UserDataType,
    )
    from dl_core.components.ids import AvatarId


_EXPRESSION_CTX_TV = TypeVar("_EXPRESSION_CTX_TV", bound="ExpressionCtx")


@attr.s(auto_attribs=True, frozen=True)
class ExpressionCtx:
    expression: ClauseElement
    avatar_ids: Optional[Sequence[str]] = None  # TODO: make required
    user_type: Optional[UserDataType] = None
    alias: Optional[str] = None
    original_field_id: Optional[Any] = None

    def clone(self: _EXPRESSION_CTX_TV, **kwargs: Any) -> _EXPRESSION_CTX_TV:
        return attr.evolve(self, **kwargs)


@attr.s(auto_attribs=True, frozen=True)
class OrderByExpressionCtx(ExpressionCtx):
    direction: OrderDirection = OrderDirection.asc


@attr.s(auto_attribs=True, frozen=True)
class JoinOnExpressionCtx(ExpressionCtx):  # noqa
    left_id: Optional[AvatarId] = attr.ib(kw_only=True)  # can be None for feature-managed relations
    right_id: AvatarId = attr.ib(kw_only=True)
    join_type: JoinType = attr.ib(kw_only=True)

lib/dl_core/dl_core/query/expression.py
datalens-tech/datalens-backend
from __future__ import annotations

from abc import ABCMeta
from typing import (
    AsyncIterator,
    Iterator,
    Optional,
    TypeVar,
    Union,
)


_ITER_TV = TypeVar("_ITER_TV")


class DataStreamBase(Iterator[_ITER_TV], metaclass=ABCMeta):
    max_percent = 99

    def __iter__(self):  # type: ignore  # TODO: fix
        return self

    def get_progress_percent(self) -> int:
        """Return current progress in percent"""
        return self.max_percent

    def __next__(self) -> _ITER_TV:
        raise StopIteration


class AsyncDataStreamBase(AsyncIterator[_ITER_TV], metaclass=ABCMeta):
    max_percent = 99

    def __aiter__(self) -> AsyncDataStreamBase:
        return self

    def get_progress_percent(self) -> int:
        """Return current progress in percent"""
        return self.max_percent

    async def __anext__(self) -> _ITER_TV:
        raise StopIteration


class SimpleDataStream(DataStreamBase):
    """
    Data stream that calculates progress percentage based on total data size.
    """

    def __init__(self, data_iter: Iterator[dict], rows_to_copy: Optional[int] = None):
        self._data_iter = data_iter
        self._rows_to_copy = rows_to_copy

        self._rows_read = 0

    def get_progress_percent(self) -> int:
        if self._rows_to_copy is None:
            return 0
        if self._rows_to_copy == 0:
            return 100
        return min(self.max_percent, int((self._rows_read / self._rows_to_copy) * 100))

    def __next__(self) -> dict:
        row = next(self._data_iter)
        self._rows_read += 1
        return row


class LazyDataStream(DataStreamBase):
    """
    Data stream with progress percentage set from the outside.
    Does no calculations by itself.
    """

    def __init__(self, data_iter: Iterator[dict]):
        self._data_iter = data_iter

        self._progress = 0

    def set_progress_percent(self, progress: Union[int, float]) -> None:
        self._progress = int(progress)

    def get_progress_percent(self) -> int:
        return self._progress

    def __next__(self) -> dict:
        row = next(self._data_iter)
        return row


class SimpleUntypedDataStream(DataStreamBase):
    """
    Data stream that calculates progress percentage based on total data size.
    """

    def __init__(self, data_iter: Iterator[list], rows_to_copy: Optional[int] = None):
        self._data_iter = data_iter
        self._rows_to_copy = rows_to_copy

        self._rows_read = 0

    def get_progress_percent(self) -> int:
        if self._rows_to_copy is None:
            return 0
        if self._rows_to_copy == 0:
            return 100
        return min(self.max_percent, int((self._rows_read / self._rows_to_copy) * 100))

    def __next__(self) -> list:
        row = next(self._data_iter)
        self._rows_read += 1
        return row


class SimpleUntypedAsyncDataStream(AsyncDataStreamBase):
    def __init__(self, data_iter: AsyncIterator[list], rows_to_copy: Optional[int] = None):
        self._data_iter = data_iter
        self._rows_to_copy = rows_to_copy

        self._rows_read = 0

    def get_progress_percent(self) -> int:
        if self._rows_to_copy is None:
            return 0
        if self._rows_to_copy == 0:
            return 100
        return min(self.max_percent, int((self._rows_read / self._rows_to_copy) * 100))

    async def __anext__(self) -> list:
        row = await self._data_iter.__anext__()
        self._rows_read += 1
        return row

lib/dl_core/dl_core/raw_data_streaming/stream.py
datalens-tech/datalens-backend
import time
from typing import (
    Any,
    Callable,
    Optional,
    Type,
)

from dl_api_commons.reporting.models import NotificationReportingRecord
from dl_constants.enums import (
    NotificationLevel,
    NotificationType,
)


class BaseNotification:
    type: NotificationType

    _title: str
    _message: str
    _level: NotificationLevel

    def __init__(self, **kwargs: Any) -> None:
        pass

    @property
    def title(self) -> str:
        return self._title

    @property
    def message(self) -> str:
        return self._message

    @property
    def level(self) -> NotificationLevel:
        return self._level

    @property
    def locator(self) -> str:
        return self.type.value


_NOTIFICATIONS: dict[NotificationType, Type[BaseNotification]] = {}


# TODO: add control (via config/env) over which notifications should not be displayed
def register_notification() -> Callable[[Type[BaseNotification]], Type[BaseNotification]]:
    def wrap(cls: Type[BaseNotification]) -> Type[BaseNotification]:
        _NOTIFICATIONS[cls.type] = cls
        return cls

    return wrap


@register_notification()
class TotalsRemovedDueToMeasureFilterNotification(BaseNotification):
    type = NotificationType.totals_removed_due_to_measure_filter
    _title = "No totals shown"
    _message = "Totals have been removed because there is filtering by measure"
    _level = NotificationLevel.info


def get_notification_record(
    notification_type: NotificationType, **kwargs: Any
) -> Optional[NotificationReportingRecord]:
    ntf_cls = _NOTIFICATIONS.get(notification_type)
    if ntf_cls is None:
        return None
    notification = ntf_cls(**kwargs)
    return NotificationReportingRecord(
        title=notification.title,
        message=notification.message,
        level=notification.level,
        locator=notification.locator,
        timestamp=time.time(),
    )

lib/dl_core/dl_core/reporting/notifications.py
datalens-tech/datalens-backend
from __future__ import annotations

import enum
from typing import (
    ClassVar,
    Dict,
    Optional,
    Union,
)

import attr
from typing_extensions import Literal

from dl_constants.enums import (
    ConnectionType,
    ReportingQueryType,
)


# TODO FIX: Use as intermediate entity (at this moment is not used)
@attr.s(frozen=True, auto_attribs=True)
class DbQueryExecutionReport:
    event_code: ClassVar[str] = "profile_db_request"

    query_id: str
    dataset_id: str
    user_id: Optional[str]
    billing_folder_id: Optional[str]
    connection_id: str
    connection_type: ConnectionType
    source: Optional[str]
    username: Optional[str]
    execution_time: int
    query: Optional[str]
    status: Union[Literal["success"], Literal["error"]]
    error: Optional[str]
    host: str

    cache_used: bool
    cache_full_hit: bool

    # not in action fields:
    query_type: ReportingQueryType
    is_public: bool

    def convert_for_logging_extras(self, value) -> Union[str, int, bool, None]:  # type: ignore  # TODO: fix
        if value is None:
            return None
        elif isinstance(value, (str, int, bool)):
            return value
        elif isinstance(value, enum.Enum):
            return value.name
        else:
            return repr(value)

    def to_logging_extras(self) -> Dict[str, Union[str, int, bool]]:
        return dict(
            {  # type: ignore  # TODO: fix
                k: self.convert_for_logging_extras(v) for k, v in attr.asdict(self)  # type: ignore  # TODO: fix
            },
            event_code=self.event_code,
        )

lib/dl_core/dl_core/reporting/reports.py
datalens-tech/datalens-backend
"""
RLS - Row-Level Security

Module implements RLS management functionality
"""

from __future__ import annotations

import abc
import copy
from typing import (
    Dict,
    List,
    Optional,
    Tuple,
)

import attr

from dl_constants.enums import (
    RLSPatternType,
    RLSSubjectType,
)
from dl_utils.utils import split_list


@attr.s(slots=True)
class RLSSubject:
    subject_type: RLSSubjectType = attr.ib()
    subject_id: str = attr.ib()
    subject_name: str = attr.ib()  # login, group name, etc


# Special type subject that denotes 'all subjects'.
RLS_ALL_SUBJECT_NAME = "*"
RLS_ALL_SUBJECT = RLSSubject(
    subject_type=RLSSubjectType.all, subject_id=RLS_ALL_SUBJECT_NAME, subject_name=RLS_ALL_SUBJECT_NAME
)
RLS_FAILED_USER_NAME_PREFIX = "!FAILED_"
# Special type subject that denotes 'insert userid'.
RLS_USERID_SUBJECT_NAME = "userid"
RLS_USERID_SUBJECT = RLSSubject(subject_type=RLSSubjectType.userid, subject_id="", subject_name=RLS_USERID_SUBJECT_NAME)


@attr.s(slots=True)
class RLSEntry:
    field_guid: str = attr.ib()
    allowed_value: Optional[str] = attr.ib()
    subject: RLSSubject = attr.ib()
    # Note: this is a bit of a hack to avoid very extensive splitting of the
    # RLSEntry into multiple classes.
    # For `pattern_type=RLSPatternType.all`, `allowed_value` must be `None`.
    # For `pattern_type=RLSPatternType.userid`, `allowed_value` must be `None`,
    # and `subject` must be `RLSSubjectType.userid`.
    pattern_type: RLSPatternType = attr.ib(default=RLSPatternType.value)

    def ensure_removed_failed_subject_prefix(self) -> RLSEntry:
        rls_entry = copy.deepcopy(self)
        rls_entry.subject.subject_name = rls_entry.subject.subject_name.removeprefix(RLS_FAILED_USER_NAME_PREFIX)
        return rls_entry


@attr.s
class RLS:
    items: List[RLSEntry] = attr.ib(factory=list)

    @property
    def has_restrictions(self) -> bool:
        return bool(self.items)

    @property
    def fields_with_rls(self) -> List[str]:
        return list(set(item.field_guid for item in self.items))

    def get_entries(
        self, field_guid: str, subject_type: RLSSubjectType, subject_id: str, add_userid_entry: bool = True
    ) -> List[RLSEntry]:
        return [
            item
            for item in self.items
            if item.field_guid == field_guid
            and (
                # Same subject
                (item.subject.subject_type == subject_type and item.subject.subject_id == subject_id)
                # 'all subjects' matches any subject.
                or item.subject.subject_type == RLSSubjectType.all
                # `userid: userid`
                or (add_userid_entry and item.pattern_type == RLSPatternType.userid)
            )
        ]

    def get_field_restriction_for_subject(
        self,
        field_guid: str,
        subject_type: RLSSubjectType,
        subject_id: str,
    ) -> Tuple[bool, bool, Optional[List[str]]]:
        """
        For subject and field, return `allow_all_values, allowed_values`.
        """
        rls_entries = self.get_entries(field_guid=field_guid, subject_type=subject_type, subject_id=subject_id)

        # There's a `*: {current_user}` entry, no need to filter.
        if any(rls_entry.pattern_type == RLSPatternType.all for rls_entry in rls_entries):
            return True, False, None

        # Pick out userid-entry, if any
        userid_entries, rls_entries = split_list(
            rls_entries, lambda rls_entry: rls_entry.pattern_type == RLSPatternType.userid
        )

        # normal values
        assert all(
            rls_entry.pattern_type == RLSPatternType.value for rls_entry in rls_entries
        ), "only simple values should remain at this point"
        allowed_values = [rls_entry.allowed_value for rls_entry in rls_entries]

        # `userid: userid` case
        allow_userid = False
        if userid_entries:
            assert len(userid_entries) == 1
            userid_entry = userid_entries[0]
            assert userid_entry.pattern_type == RLSPatternType.userid
            assert userid_entry.allowed_value is None
            # only `userid: userid` makes sense
            assert userid_entry.subject.subject_type == RLSSubjectType.userid
            allow_userid = True

        return False, allow_userid, allowed_values  # type: ignore  # TODO: fix

    def _should_add_entry(
        self,
        field_guid: str,
        subject_type: RLSSubjectType,
        subject_id: str,
        allowed_value: Optional[str],
        pattern_type: RLSPatternType,
    ) -> bool:
        existing_allow_all_values, allow_userid, existing_allowed_values = self.get_field_restriction_for_subject(
            field_guid=field_guid, subject_type=subject_type, subject_id=subject_id
        )
        if existing_allow_all_values:  # already wildcarded, skip
            return False
        if pattern_type == RLSPatternType.userid:
            return not allow_userid  # add if it's not already set.
        # Add if not existing:
        return allowed_value not in existing_allowed_values  # type: ignore  # TODO: fix

    def add_field_restriction_for_subject(
        self,
        field_guid: str,
        subject_type: RLSSubjectType,
        subject_id: str,
        subject_name: str,
        allowed_value: Optional[str],
        pattern_type: RLSPatternType = RLSPatternType.value,
        force: bool = False,
    ) -> None:
        """
        Register a new RLS restriction.

        Currently, only used in tests.
        """
        if force or self._should_add_entry(
            field_guid=field_guid,
            subject_type=subject_type,
            subject_id=subject_id,
            allowed_value=allowed_value,
            pattern_type=pattern_type,
        ):
            # TODO?: if allow_all_values, drop all prior items for subject+field?
            self.items.append(
                RLSEntry(
                    field_guid=field_guid,
                    allowed_value=allowed_value,
                    pattern_type=pattern_type,
                    subject=RLSSubject(
                        subject_type=subject_type,
                        subject_name=subject_name,
                        subject_id=subject_id,
                    ),
                )
            )

    def get_subject_restrictions(
        self,
        subject_type: RLSSubjectType,
        subject_id: str,
    ) -> Dict[str, List[str]]:
        result = {}
        for field_guid in self.fields_with_rls:
            allow_all_values, allow_userid, allowed_values = self.get_field_restriction_for_subject(
                field_guid=field_guid, subject_type=subject_type, subject_id=subject_id
            )
            if allow_all_values:
                # `*: subject` => 'not restricted'
                continue

            # For `userid: userid`, add the subject id to the values.
            if allow_userid:
                allowed_values = list(allowed_values) + [subject_id]  # type: ignore  # TODO: fix

            result[field_guid] = allowed_values

        return result  # type: ignore  # TODO: fix


class BaseSubjectResolver(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_subjects_by_names(self, names: List[str]) -> List[RLSSubject]:
        raise NotImplementedError

lib/dl_core/dl_core/rls.py
datalens-tech/datalens-backend
"""
RDL JSON - Redis DataLens JSON

Serialization with support for custom objects like ``date`` & ``datetime``.

>>> data = dict(some_dt=datetime.datetime(2019, 6, 17, 7, 1, 41, 79585, tzinfo=datetime.timezone(datetime.timedelta(seconds=-1320))))
>>> dumped = json.dumps(data, cls=RedisDatalensDataJSONEncoder)
>>> json.loads(dumped)['some_dt']
{'__dl_type__': 'datetime', 'value': [[2019, 6, 17, 7, 1, 41, 79585], -1320.0]}
>>> json.loads(dumped, cls=RedisDatalensDataJSONDecoder)['some_dt']
datetime.datetime(2019, 6, 17, 7, 1, 41, 79585, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=85080)))
>>> # ^ negative timedeltas' repr is actually this weird.
>>> tricky_data = dict(normal=data, abnormal=json.loads(dumped))
>>> tricky_data_dumped = json.dumps(tricky_data, cls=RedisDatalensDataJSONEncoder)
>>> tricky_roundtrip = json.loads(tricky_data_dumped, cls=RedisDatalensDataJSONDecoder)
>>> # WARNING: not currently viable: `assert tricky_roundtrip == tricky_data, tricky_roundtrip`
"""

from __future__ import annotations

import abc
import datetime
import decimal
import json
from typing import (
    Any,
    ClassVar,
    Generic,
    Type,
    TypeVar,
    Union,
)
import uuid

from dl_constants.types import (
    TJSONExt,
    TJSONLike,
)


_TS_TV = TypeVar("_TS_TV")


class TypeSerializer(Generic[_TS_TV]):
    typeobj: ClassVar[Type[_TS_TV]]  # type: ignore  # TODO: fix
    typename: ClassVar[str]

    @staticmethod
    @abc.abstractmethod
    def to_jsonable(value: _TS_TV) -> TJSONLike:
        raise NotImplementedError

    @staticmethod
    @abc.abstractmethod
    def from_jsonable(value: TJSONLike) -> _TS_TV:
        raise NotImplementedError


class DateSerializer(TypeSerializer[datetime.date]):
    typeobj = datetime.date
    typename = "date"

    @staticmethod
    def to_jsonable(obj: datetime.date) -> TJSONLike:
        return obj.isoformat()

    @staticmethod
    def from_jsonable(value: TJSONLike) -> datetime.date:
        assert isinstance(value, str)
        return datetime.date.fromisoformat(value)


class DatetimeSerializer(TypeSerializer[datetime.datetime]):
    typeobj = datetime.datetime
    typename = "datetime"

    @staticmethod
    def to_jsonable(obj: datetime.datetime) -> TJSONLike:
        tzinfo = (
            obj.tzinfo.utcoffset(obj).total_seconds() if obj.tzinfo is not None else None  # type: ignore  # TODO: fix
        )
        return (
            [obj.year, obj.month, obj.day, obj.hour, obj.minute, obj.second, obj.microsecond],
            tzinfo,
        )

    @staticmethod
    def from_jsonable(value: TJSONLike) -> datetime.datetime:
        assert isinstance(value, (list, tuple))
        assert len(value) == 2
        parts, offset_sec = value
        assert len(parts) == 7  # type: ignore  # TODO: fix
        tzinfo = (
            datetime.timezone(datetime.timedelta(seconds=offset_sec))  # type: ignore  # TODO: fix
            if offset_sec is not None
            else None
        )
        return datetime.datetime(*parts, tzinfo=tzinfo)  # type: ignore  # TODO: fix


class TimeSerializer(TypeSerializer[datetime.time]):
    typeobj = datetime.time
    typename = "time"

    @staticmethod
    def to_jsonable(obj: datetime.time) -> TJSONLike:
        # Not expecting this to be valid for non-constant time tzinfos.
        tzinfo = (
            obj.tzinfo.utcoffset(datetime.datetime(1970, 1, 1)).total_seconds()  # type: ignore  # TODO: fix
            if obj.tzinfo is not None
            else None
        )
        return ([obj.hour, obj.minute, obj.second, obj.microsecond], tzinfo)

    @staticmethod
    def from_jsonable(value: TJSONLike) -> datetime.time:
        assert isinstance(value, (list, tuple))
        assert len(value) == 2, value
        parts, offset_sec = value
        assert len(parts) == 4, parts  # type: ignore  # TODO: fix
        assert isinstance(offset_sec, (int, float)), offset_sec
        tzinfo = datetime.timezone(datetime.timedelta(seconds=offset_sec)) if offset_sec is not None else None
        return datetime.time(*parts, tzinfo=tzinfo)  # type: ignore  # TODO: fix


class TimedeltaSerializer(TypeSerializer[datetime.timedelta]):
    typeobj = datetime.timedelta
    typename = "timedelta"

    @staticmethod
    def to_jsonable(obj: datetime.timedelta) -> TJSONLike:
        return obj.total_seconds()

    @staticmethod
    def from_jsonable(value: TJSONLike) -> datetime.timedelta:
        assert isinstance(value, float)
        return datetime.timedelta(seconds=value)


class DecimalSerializer(TypeSerializer[decimal.Decimal]):
    typeobj = decimal.Decimal
    typename = "decimal"

    @staticmethod
    def to_jsonable(obj: decimal.Decimal) -> TJSONLike:
        return str(obj)

    @staticmethod
    def from_jsonable(value: TJSONLike) -> decimal.Decimal:
        assert isinstance(value, str)
        return decimal.Decimal(value)


class UUIDSerializer(TypeSerializer[uuid.UUID]):
    typeobj = uuid.UUID
    typename = "uuid"

    @staticmethod
    def to_jsonable(obj: uuid.UUID) -> TJSONLike:
        return str(obj)

    @staticmethod
    def from_jsonable(value: TJSONLike) -> uuid.UUID:
        assert isinstance(value, str)
        return uuid.UUID(value)


COMMON_SERIALIZERS = [
    DateSerializer,
    DatetimeSerializer,
    TimeSerializer,
    TimedeltaSerializer,
    DecimalSerializer,
    UUIDSerializer,
]
assert len(set(cls.typename for cls in COMMON_SERIALIZERS)) == len(COMMON_SERIALIZERS), "uniqueness check"


class RedisDatalensDataJSONEncoder(json.JSONEncoder):
    JSONABLERS_MAP = {cls.typeobj: cls for cls in COMMON_SERIALIZERS}

    def default(self, obj):  # type: ignore  # TODO: fix
        typeobj = type(obj)
        preprocessor = self.JSONABLERS_MAP.get(typeobj)
        if preprocessor is not None:
            return dict(__dl_type__=preprocessor.typename, value=preprocessor.to_jsonable(obj))

        return super().default(obj)  # effectively, raises `TypeError`


class RedisDatalensDataJSONDecoder(json.JSONDecoder):
    DEJSONABLERS_MAP = {cls.typename: cls for cls in COMMON_SERIALIZERS}
    # Transition
    COMPAT_CONVERTERS_MAP = {
        "date": lambda o: datetime.date(*[int(p) for p in o["value"].split("-")]),
        "datetime": lambda o: datetime.datetime(  # type: ignore  # TODO: fix
            *o["parts"],
            tzinfo=None if o["offset_sec"] is None else datetime.timezone(datetime.timedelta(seconds=o["offset_sec"])),
        ),
    }

    def __init__(self, *args, **kwargs):  # type: ignore  # TODO: fix
        super().__init__(*args, object_hook=self.object_hook, **kwargs)

    def object_hook(self, obj):  # type: ignore  # TODO: fix
        # WARNING: this might collide with some unexpected objects that have a `__dl_type__` key.
        # A correct roundtrip way would be to wrap all objects with a `__dl_type__` key into another layer.
        dl_type = obj.get("__dl_type__")
        if dl_type is not None:
            postprocessor = self.DEJSONABLERS_MAP.get(dl_type)
            if postprocessor is not None:
                return postprocessor.from_jsonable(obj["value"])

        # Transition
        compat_dl_type = obj.get("dl_type")
        if compat_dl_type is not None:
            compat_converter = self.COMPAT_CONVERTERS_MAP.get(compat_dl_type)
            if compat_converter is not None:
                return compat_converter(obj)

        return obj


def common_dumps(value: TJSONExt, **kwargs: Any) -> bytes:
    return json.dumps(
        value,
        cls=RedisDatalensDataJSONEncoder,
        separators=(",", ":"),
        ensure_ascii=False,
        check_circular=False,  # dangerous but faster
        **kwargs,
    ).encode("utf-8")


def hashable_dumps(value: TJSONExt, sort_keys: bool = True, check_circular: bool = False, **kwargs: Any) -> str:
    return json.dumps(
        value,
        cls=RedisDatalensDataJSONEncoder,
        separators=(",", ":"),
        ensure_ascii=False,
        check_circular=check_circular,
        sort_keys=sort_keys,
        **kwargs,
    )


def common_loads(value: Union[bytes, str], **kwargs: Any) -> TJSONExt:
    return json.loads(
        value,
        cls=RedisDatalensDataJSONDecoder,
        **kwargs,
    )


class CacheMetadataSerialization:
    r"""
    Serialization/deserialization of metadata+data pair.
    Avoids touching/copying of most of the data
    (e.g. if the metadata shows it is not necessary).

    Similar solutions: msgpack streaming; protobufs.

    >>> metadata = dict(success=True)
    >>> data = b'\nzxc'
    >>> full_data = CacheMetadataSerialization.serialize(metadata, data)
    >>> ser_obj = CacheMetadataSerialization(full_data)
    >>> assert ser_obj.metadata == metadata
    >>> assert set_obj.data_bytes == data
    """

    _sep = b"\n"

    @classmethod
    def serialize(cls, metadata: TJSONExt, data_bytes: bytes) -> bytes:
        metadata_bytes = common_dumps(metadata)
        assert cls._sep not in metadata_bytes
        return metadata_bytes + cls._sep + data_bytes

    def __init__(self, full_data: bytes):
        self.full_data = full_data
        sep_pos = full_data.index(self._sep)
        self.sep_pos = sep_pos
        self.metadata: TJSONExt = common_loads(full_data[:sep_pos])

    @property
    def data_bytes(self) -> bytes:
        return self.full_data[self.sep_pos + len(self._sep) :]

lib/dl_core/dl_core/serialization.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_core.services_registry.top_level import (
    DefaultServicesRegistry,
    ServicesRegistry,
)


__all__ = (
    "ServicesRegistry",
    "DefaultServicesRegistry",
)

lib/dl_core/dl_core/services_registry/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Optional,
)

import attr

from dl_core.data_processing.cache.engine import EntityCacheEngineAsync
from dl_core.utils import FutureRef


if TYPE_CHECKING:
    from dl_core.services_registry.top_level import ServicesRegistry


LOGGER = logging.getLogger(__name__)


class CacheEngineFactory(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_cache_engine(self, entity_id: Optional[str]) -> Optional[EntityCacheEngineAsync]:
        pass


@attr.s
class DefaultCacheEngineFactory(CacheEngineFactory):
    _services_registry_ref: FutureRef[ServicesRegistry] = attr.ib()
    cache_save_background: Optional[bool] = attr.ib(default=None)

    @property
    def service_registry(self) -> ServicesRegistry:
        return self._services_registry_ref.ref

    def get_cache_engine(self, entity_id: Optional[str]) -> Optional[EntityCacheEngineAsync]:
        if not entity_id:
            LOGGER.info("Can not create entity cache engine: no entity_id")
            return None

        redis_client = self.service_registry.get_caches_redis_client()
        if redis_client is None:
            LOGGER.info("Can not create entity cache engine: service registry did not return a Redis client")
            return None

        redis_client_slave = self.service_registry.get_caches_redis_client(allow_slave=True)

        # Avoid unexpectedly overriding the default of the `EntityCacheEngineAsync` itself.
        kwargs: Dict[str, Any] = {}
        c_s_bg = self.cache_save_background
        if c_s_bg is not None:
            kwargs.update(CACHE_SAVE_BACKGROUND=c_s_bg)

        return EntityCacheEngineAsync(
            entity_id=entity_id,
            rc=redis_client,
            rc_slave=redis_client_slave,
            **kwargs,
        )

lib/dl_core/dl_core/services_registry/cache_engine_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
from typing import (
    Callable,
    TypeVar,
)

from dl_utils.aio import ContextVarExecutor


_IDENTITY_TV = TypeVar("_IDENTITY_TV")


class ComputeExecutor:
    @abc.abstractmethod
    async def execute(self, func: Callable[[], _IDENTITY_TV]) -> _IDENTITY_TV:
        pass

    def close(self) -> None:
        pass


class ComputeExecutorTPE:
    def __init__(self, name_prefix: str = "COMPUTE_EXECUTOR_TPE_") -> None:
        self._tpe = ContextVarExecutor(thread_name_prefix=name_prefix)

    async def execute(self, func: Callable[[], _IDENTITY_TV]) -> _IDENTITY_TV:
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(self._tpe, func)

    def close(self) -> None:
        self._tpe.shutdown(wait=False, cancel_futures=True)

lib/dl_core/dl_core/services_registry/compute_executor.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import ChainMap
import logging
import random
from typing import (
    TYPE_CHECKING,
    ClassVar,
    FrozenSet,
    MutableMapping,
    Optional,
    Sequence,
    Tuple,
    Type,
)

import attr

from dl_configs.rqe import (
    RQEBaseURL,
    RQEConfig,
)
from dl_core import connection_executors
from dl_core.connection_executors import ExecutionMode
from dl_core.connection_executors.async_base import AsyncConnExecutorBase
from dl_core.connection_executors.models.common import RemoteQueryExecutorData
from dl_core.connection_models import (
    ConnDTO,
    ConnectOptions,
    DefaultSQLDTO,
)
from dl_core.services_registry.conn_executor_factory_base import (
    BaseClosableExecutorFactory,
    CEFactoryError,
    ConnExecutorRecipe,
)
from dl_core.us_connection_base import ConnectionBase
from dl_utils.aio import ContextVarExecutor


if TYPE_CHECKING:
    from dl_core.connection_executors.common_base import ConnExecutorBase
    from dl_core.connections_security.base import ConnectionSecurityManager
    from dl_core.services_registry.typing import (
        ConnectOptionsFactory,
        ConnectOptionsMutator,
    )
    from dl_core.us_connection_base import ExecutorBasedMixin


LOGGER = logging.getLogger(__name__)


# noinspection PyDataclass
@attr.s(frozen=True)
class DefaultConnExecutorFactory(BaseClosableExecutorFactory):
    rqe_config: Optional[RQEConfig] = attr.ib()
    tpe: Optional[ContextVarExecutor] = attr.ib()
    conn_sec_mgr: ConnectionSecurityManager = attr.ib()

    is_bleeding_edge_user: bool = attr.ib(default=False)
    conn_cls_whitelist: Optional[FrozenSet[Type[ExecutorBasedMixin]]] = attr.ib(default=None)
    # User function that can override connect options.
    #  If factory returns `None` default connection connect options will be used
    connect_options_factory: Optional[ConnectOptionsFactory] = attr.ib(default=None)
    # User function that can alter connect options (whether they were generated
    # by factory or from the conn).
    connect_options_mutator: Optional[ConnectOptionsMutator] = attr.ib(default=None)
    force_non_rqe_mode: bool = attr.ib(default=False)

    DEFAULT_MAP_CONN_TYPE_SYNC_CE_TYPE: dict[Type[ConnectionBase], Type[ConnExecutorBase]] = {}

    DEFAULT_MAP_CONN_TYPE_ASYNC_CE_TYPE: dict[Type[ConnectionBase], Type[AsyncConnExecutorBase]] = {}

    SINGLE_HOST_RETRY_ATTEMPTS: ClassVar[int] = 3
    MAX_HOST_RETRY_ATTEMPTS: ClassVar[int] = 3

    # TODO FIX: Make typing/implementation more accurate
    def get_async_conn_executor_cls(self, conn: ExecutorBasedMixin) -> Type[AsyncConnExecutorBase]:
        map_conn_type_ce_type: MutableMapping[Type[ConnectionBase], Type[ConnExecutorBase]] = ChainMap()

        if self.async_env:
            map_conn_type_ce_type.update(self.DEFAULT_MAP_CONN_TYPE_SYNC_CE_TYPE)
            map_conn_type_ce_type.update(self.DEFAULT_MAP_CONN_TYPE_ASYNC_CE_TYPE)
        else:
            map_conn_type_ce_type.update(self.DEFAULT_MAP_CONN_TYPE_SYNC_CE_TYPE)

        if self.conn_cls_whitelist is not None:
            map_conn_type_ce_type = {
                conn_cls: ce_cls
                for conn_cls, ce_cls in map_conn_type_ce_type.items()
                if conn_cls in self.conn_cls_whitelist
            }

        try:
            ce = map_conn_type_ce_type[type(conn)]
            assert issubclass(ce, AsyncConnExecutorBase)
            return ce
        except KeyError:
            raise CEFactoryError(f"No executor for connection type {type(conn)}")

    def _get_conn_hosts_pool(self, dto: ConnDTO) -> Sequence[str]:
        if isinstance(dto, DefaultSQLDTO):
            hosts = list(dto.multihosts)

            random.shuffle(hosts)

            if len(hosts) == 1:
                hosts = hosts * self.SINGLE_HOST_RETRY_ATTEMPTS
            elif len(hosts) > self.MAX_HOST_RETRY_ATTEMPTS:
                LOGGER.info(f"Hosts list truncated from {len(hosts)} to {self.MAX_HOST_RETRY_ATTEMPTS}")
                hosts = hosts[: self.MAX_HOST_RETRY_ATTEMPTS]

            return tuple(hosts)
        else:
            return ()

    def _get_async_conn_executor_recipe(
        self,
        conn: ExecutorBasedMixin,
    ) -> ConnExecutorRecipe:
        # noinspection PyProtectedMember
        assert conn._context is self.req_ctx_info, "Divergence in RCI between CE factory and US connection"

        dto = conn.get_conn_dto()
        conn_hosts_pool = self._get_conn_hosts_pool(dto)

        executor_cls = self.get_async_conn_executor_cls(conn)
        exec_mode, rqe_data = self._get_exec_mode_and_rqe_attrs(conn, executor_cls)
        LOGGER.info("Connection executor exec_mode = %s", exec_mode)

        overridden_connect_options: Optional[ConnectOptions] = None
        if self.connect_options_factory is not None:
            overridden_connect_options = self.connect_options_factory(conn)

        connect_options: ConnectOptions
        if overridden_connect_options is not None:
            connect_options = overridden_connect_options
        else:
            connect_options = conn.get_conn_options()

        if self.connect_options_mutator is not None:
            connect_options = self.connect_options_mutator(connect_options)

        return ConnExecutorRecipe(
            ce_cls=executor_cls,
            conn_dto=dto,
            connect_options=connect_options,
            rqe_data=rqe_data,
            exec_mode=exec_mode,
            conn_hosts_pool=conn_hosts_pool,  # type: ignore  # TODO: fix
        )

    def _cook_conn_executor(self, recipe: ConnExecutorRecipe, with_tpe: bool) -> AsyncConnExecutorBase:
        def _conn_host_fail_callback_func(host: str):  # type: ignore  # TODO: fix
            LOGGER.info("DB host %s unavailable", host)

        executor_cls = recipe.ce_cls
        if issubclass(executor_cls, connection_executors.DefaultSqlAlchemyConnExecutor):
            # TODO FIX: log connection info
            LOGGER.info("Creating CE %s", executor_cls)
            return executor_cls(
                conn_dto=recipe.conn_dto,
                conn_options=recipe.connect_options,
                req_ctx_info=self.req_ctx_info,
                remote_qe_data=recipe.rqe_data,
                tpe=self.tpe if with_tpe else None,
                exec_mode=recipe.exec_mode,
                sec_mgr=self.conn_sec_mgr,
                conn_hosts_pool=recipe.conn_hosts_pool,
                host_fail_callback=_conn_host_fail_callback_func,
                services_registry=self._services_registry_ref.ref,  # Do not use. To be deprecated. Somehow.
            )
        else:
            raise CEFactoryError(f"Can not instantiate {executor_cls}")

    def _get_exec_mode_and_rqe_attrs(
        self, conn: ExecutorBasedMixin, executor_cls: Type[AsyncConnExecutorBase]
    ) -> Tuple[ExecutionMode, Optional[RemoteQueryExecutorData]]:
        conn_dto = conn.get_conn_dto()
        conn_options = conn.get_conn_options()
        ce_cls = self.get_async_conn_executor_cls(conn)

        if not self.conn_sec_mgr.is_safe_connection(conn_dto, conn_options):
            # Only RQE mode with external RQE supported for unsafe connection
            if ExecutionMode.RQE not in executor_cls.supported_exec_mode:
                raise CEFactoryError(
                    f"Connection classified as insecure,"
                    f" but configured executor class {executor_cls} does not support RQE exec mode"
                )
            return ExecutionMode.RQE, self._get_rqe_data(external=True)

        else:
            if self.force_non_rqe_mode:
                # WARNING: debug-only flag, might not even work
                return ExecutionMode.DIRECT, None
            elif ce_cls.is_pure_async():
                return ExecutionMode.DIRECT, None
            elif self.async_env:
                return ExecutionMode.RQE, self._get_rqe_data(external=False)
            else:
                return ExecutionMode.RQE, self._get_rqe_data(external=False)

    def _get_rqe_data(self, external: bool) -> RemoteQueryExecutorData:
        if self.rqe_config is None:
            raise CEFactoryError("RQE data was requested, but RQE config was not passed to CE factory")

        async_rqe_netloc: RQEBaseURL
        sync_rqe_netloc: RQEBaseURL

        LOGGER.info('RQE mode is "%s"', "external" if external else "internal")
        if external:
            async_rqe_netloc = self.rqe_config.ext_async_rqe
            sync_rqe_netloc = self.rqe_config.ext_sync_rqe
        else:
            async_rqe_netloc = self.rqe_config.int_async_rqe
            sync_rqe_netloc = self.rqe_config.int_sync_rqe

        return RemoteQueryExecutorData(
            hmac_key=self.rqe_config.hmac_key,
            async_protocol=async_rqe_netloc.scheme,
            async_host=async_rqe_netloc.host,
            async_port=async_rqe_netloc.port,
            sync_protocol=sync_rqe_netloc.scheme,
            sync_host=sync_rqe_netloc.host,
            sync_port=sync_rqe_netloc.port,
        )

    @property
    def conn_security_manager(self) -> ConnectionSecurityManager:
        return self.conn_sec_mgr

    def clone(self, **kwargs):  # type: ignore  # TODO: fix
        return attr.evolve(self, **kwargs)


def register_sync_conn_executor_class(conn_cls: Type[ConnectionBase], sync_ce_cls: Type[ConnExecutorBase]) -> None:
    DefaultConnExecutorFactory.DEFAULT_MAP_CONN_TYPE_SYNC_CE_TYPE[conn_cls] = sync_ce_cls


def register_async_conn_executor_class(
    conn_cls: Type[ConnectionBase], async_ce_cls: Type[AsyncConnExecutorBase]
) -> None:
    DefaultConnExecutorFactory.DEFAULT_MAP_CONN_TYPE_ASYNC_CE_TYPE[conn_cls] = async_ce_cls

lib/dl_core/dl_core/services_registry/conn_executor_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
import functools
import itertools
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    List,
    Optional,
    Sequence,
    Type,
    TypeVar,
)

import attr
from typing_extensions import final

from dl_api_commons.base_models import RequestContextInfo
from dl_core.connection_executors import (
    ExecutionMode,
    SyncWrapperForAsyncConnExecutor,
)
from dl_core.connection_executors.models.common import RemoteQueryExecutorData
from dl_core.connection_models import (
    ConnDTO,
    ConnectOptions,
)
from dl_core.services_registry.entity_checker import EntityUsageChecker
from dl_core.us_connection_base import (
    ConnectionBase,
    ExecutorBasedMixin,
)
from dl_core.utils import FutureRef


if TYPE_CHECKING:
    from dl_core.connection_executors import (
        AsyncConnExecutorBase,
        SyncConnExecutorBase,
    )
    from dl_core.connections_security.base import ConnectionSecurityManager
    from dl_core.services_registry import ServicesRegistry


LOGGER = logging.getLogger(__name__)


class CEFactoryError(Exception):
    pass


class ConnExecutorFactory(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_async_conn_executor(self, conn: ConnectionBase) -> AsyncConnExecutorBase:
        pass

    @abc.abstractmethod
    def get_async_conn_executor_cls(self, conn: ExecutorBasedMixin) -> Type[AsyncConnExecutorBase]:
        pass

    @abc.abstractmethod
    def get_sync_conn_executor(self, conn: ConnectionBase) -> SyncConnExecutorBase:
        pass

    @abc.abstractmethod
    def close_sync(self) -> None:
        pass

    @abc.abstractmethod
    async def close_async(self) -> None:
        pass

    @property
    @abc.abstractmethod
    def req_ctx_info(self) -> RequestContextInfo:
        pass

    # TODO FIX: Move to ServiceRegistry
    @property
    @abc.abstractmethod
    def conn_security_manager(self) -> ConnectionSecurityManager:
        pass


_RESULT_TV = TypeVar("_RESULT_TV")


def ensure_env(async_env: bool = True) -> Callable[[Callable[..., _RESULT_TV]], Callable[..., _RESULT_TV]]:
    def real_deco(wrapped: Callable[..., _RESULT_TV]) -> Callable[..., _RESULT_TV]:
        @functools.wraps(wrapped)
        def wrapper(self: "BaseClosableExecutorFactory", *args: Any, **kwargs: Any) -> _RESULT_TV:
            if self._async_env != async_env:
                meth_env = "async" if async_env else "sync"
                factory_env = "async" if self._async_env else "sync"
                raise NotImplementedError(
                    f"Factory created for {factory_env} env "
                    f"but .{wrapped.__name__}() may be used only in {meth_env}"
                )
            return wrapped(self, *args, **kwargs)

        return wrapper

    return real_deco


@attr.s(frozen=True, auto_attribs=True)
class ConnExecutorRecipe:
    """Class that represent all data to create"""

    ce_cls: Type
    conn_dto: ConnDTO
    connect_options: ConnectOptions
    exec_mode: ExecutionMode
    rqe_data: Optional[RemoteQueryExecutorData]
    conn_hosts_pool: Sequence[str] = attr.ib(kw_only=True, converter=tuple)


@attr.s(frozen=True)
class BaseClosableExecutorFactory(ConnExecutorFactory, metaclass=abc.ABCMeta):
    """
    Base class for factories. Register all created conn executors.
    Mode depends on `async_env` flag. If it's on - all sync-related methods will throws NotImplementedError
      and vise-versa.
    If async_env=False - creates async CEs are not registered. Closing is delegate to sync wrappers.
    """

    # TODO CONSIDER: May be add some creation metadata here
    @attr.s(frozen=True, auto_attribs=True)
    class _CECreationResult:
        async_ce: AsyncConnExecutorBase
        sync_wrapper: Optional[SyncWrapperForAsyncConnExecutor]

    _map_recipe_created_ce_pair: Dict[ConnExecutorRecipe, List[_CECreationResult]] = attr.ib(init=False, factory=dict)

    _async_env: bool = attr.ib()
    _services_registry_ref: FutureRef[ServicesRegistry] = attr.ib()
    _entity_usage_checker: Optional[EntityUsageChecker] = attr.ib(default=None, kw_only=True)

    @property
    def req_ctx_info(self) -> RequestContextInfo:
        return self._services_registry_ref.ref.rci

    @property
    def async_env(self) -> bool:
        return self._async_env

    @final
    @ensure_env(async_env=True)
    def get_async_conn_executor(self, conn: ExecutorBasedMixin) -> AsyncConnExecutorBase:
        return self._get_or_create_conn_executor(conn, with_sync_wrapper=False).async_ce

    @final
    @ensure_env(async_env=False)
    def get_sync_conn_executor(self, conn: ExecutorBasedMixin) -> SyncConnExecutorBase:
        pair = self._get_or_create_conn_executor(conn, with_sync_wrapper=True)
        assert pair.sync_wrapper is not None
        return pair.sync_wrapper

    # TODO CONSIDER: May be add option to ignore cache?
    @final
    def _get_or_create_conn_executor(self, conn: ExecutorBasedMixin, with_sync_wrapper: bool) -> _CECreationResult:
        if self._entity_usage_checker is not None:
            assert isinstance(conn, ConnectionBase)
            self._entity_usage_checker.ensure_data_connection_can_be_used(rci=self.req_ctx_info, conn=conn)

        assert isinstance(conn, ExecutorBasedMixin)
        recipe = self._get_async_conn_executor_recipe(conn)
        created_list = self._map_recipe_created_ce_pair.setdefault(recipe, [])

        if len(created_list) > 0:
            LOGGER.info("CE by recipe was already created, returning cached one")
            return created_list[0]

        async_ce = self._cook_conn_executor(recipe, with_tpe=not with_sync_wrapper)
        sync_wrapper = (
            SyncWrapperForAsyncConnExecutor(
                # TODO CONSIDER: May be move to arguments
                loop=asyncio.get_event_loop(),
                async_conn_executor=async_ce,
            )
            if with_sync_wrapper
            else None
        )
        pair = self._CECreationResult(async_ce, sync_wrapper)
        created_list.append(pair)
        return pair

    @final
    @ensure_env(async_env=False)
    def close_sync(self) -> None:
        all_created_sync_wrappers = [
            pair.sync_wrapper
            for pair in itertools.chain(*self._map_recipe_created_ce_pair.values())
            if pair.sync_wrapper is not None
        ]
        for sync_wrapper in all_created_sync_wrappers:
            # noinspection PyBroadException
            try:
                sync_wrapper.close()
            except Exception:
                LOGGER.exception("Error during sync CE closing")

    @final
    @ensure_env(async_env=True)
    async def close_async(self) -> None:
        async def close_executor(executor: AsyncConnExecutorBase) -> None:
            # noinspection PyBroadException
            try:
                await executor.close()
            except Exception:
                LOGGER.exception("Error during async CE closing")

        all_created_ce = [pair.async_ce for pair in itertools.chain(*self._map_recipe_created_ce_pair.values())]
        await asyncio.gather(*[close_executor(ce) for ce in all_created_ce])

    @abc.abstractmethod
    def _get_async_conn_executor_recipe(
        self,
        conn: ExecutorBasedMixin,
    ) -> ConnExecutorRecipe:
        pass

    @abc.abstractmethod
    def _cook_conn_executor(self, recipe: ConnExecutorRecipe, with_tpe: bool) -> AsyncConnExecutorBase:
        pass

lib/dl_core/dl_core/services_registry/conn_executor_factory_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Any,
    Optional,
)

import attr

from dl_constants.enums import (
    DataSourceRole,
    ProcessorType,
    SelectorType,
)
from dl_core.data_processing.processing.processor import OperationProcessorAsyncBase
from dl_core.data_processing.processing.processor_dataset_cached import CachedDatasetProcessor
from dl_core.data_processing.processing.source_db.processor import SourceDbOperationProcessor
from dl_core.services_registry.data_processor_factory_base import (
    BaseClosableDataProcessorFactory,
    DataProcessorFactory,
)
from dl_core.us_dataset import Dataset
from dl_core.us_manager.local_cache import USEntryBuffer


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class SourceDataProcessorFactory(BaseClosableDataProcessorFactory):
    def _create_data_processor(
        self,
        dataset: Dataset,
        processor_type: ProcessorType,
        *,
        us_entry_buffer: USEntryBuffer,
        allow_cache_usage: bool = True,
        # SOURCE_DB-specific
        selector_type: Optional[SelectorType] = None,
        role: Optional[DataSourceRole] = None,
        **kwargs: Any,
    ) -> OperationProcessorAsyncBase:
        assert selector_type is not None
        assert role is not None
        selector_factory = self.services_registry.get_selector_factory()
        selector = selector_factory.get_dataset_selector(
            dataset=dataset,
            selector_type=selector_type,
            allow_cache_usage=allow_cache_usage,
            us_entry_buffer=us_entry_buffer,
        )
        processor = SourceDbOperationProcessor(
            dataset=dataset,
            selector=selector,
            role=role,
            us_entry_buffer=us_entry_buffer,
        )
        return processor


@attr.s(frozen=True)
class CompengDataProcessorFactory(BaseClosableDataProcessorFactory):
    def _create_data_processor(
        self,
        dataset: Dataset,
        processor_type: ProcessorType,
        *,
        us_entry_buffer: USEntryBuffer,
        allow_cache_usage: bool = True,
        **kwargs: Any,
    ) -> OperationProcessorAsyncBase:
        processor: OperationProcessorAsyncBase
        dproc_srv_factory = self.services_registry.get_data_processor_service_factory()
        if dproc_srv_factory is None:
            raise ValueError("Processor factory was created without a PG pool. Cannot create a PG processor")

        data_proc_service = dproc_srv_factory(processor_type)
        processor = data_proc_service.get_data_processor()

        if allow_cache_usage:
            processor = CachedDatasetProcessor(
                service_registry=self.services_registry,
                dataset=dataset,
                main_processor=processor,
            )

        return processor


@attr.s(frozen=True)
class DefaultDataProcessorFactory(DataProcessorFactory):
    _source_data_processor_factory: SourceDataProcessorFactory = attr.ib(init=False)
    _compeng_data_processor_factory: CompengDataProcessorFactory = attr.ib(init=False)

    @_source_data_processor_factory.default
    def _make_source_data_processor_factory(self) -> SourceDataProcessorFactory:
        return SourceDataProcessorFactory(
            services_registry_ref=self._services_registry_ref,  # type: ignore  # mypy bug
        )

    @_compeng_data_processor_factory.default
    def _make_compeng_data_processor_factory(self) -> CompengDataProcessorFactory:
        return CompengDataProcessorFactory(
            services_registry_ref=self._services_registry_ref,  # type: ignore  # mypy bug
        )

    async def get_data_processor(
        self,
        dataset: Dataset,
        processor_type: ProcessorType,
        *,
        us_entry_buffer: USEntryBuffer,
        allow_cache_usage: bool = True,
        **kwargs: Any,
    ) -> OperationProcessorAsyncBase:
        processor: OperationProcessorAsyncBase
        if processor_type == ProcessorType.SOURCE_DB:
            processor = await self._source_data_processor_factory.get_data_processor(
                dataset=dataset,
                processor_type=processor_type,
                us_entry_buffer=us_entry_buffer,
                allow_cache_usage=allow_cache_usage,
                **kwargs,
            )
        else:
            processor = await self._compeng_data_processor_factory.get_data_processor(
                dataset=dataset,
                processor_type=processor_type,
                us_entry_buffer=us_entry_buffer,
                allow_cache_usage=allow_cache_usage,
                **kwargs,
            )

        return processor

    async def close_async(self) -> None:
        await self._source_data_processor_factory.close_async()
        await self._compeng_data_processor_factory.close_async()

lib/dl_core/dl_core/services_registry/data_processor_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
import logging
from typing import (
    TYPE_CHECKING,
    Any,
)

import attr

from dl_constants.enums import ProcessorType
from dl_core.data_processing.processing.processor import OperationProcessorAsyncBase
from dl_core.us_dataset import Dataset
from dl_core.us_manager.local_cache import USEntryBuffer
from dl_core.utils import FutureRef


if TYPE_CHECKING:
    from dl_core.services_registry.top_level import ServicesRegistry  # noqa


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class DataProcessorFactory(metaclass=abc.ABCMeta):
    _services_registry_ref: FutureRef["ServicesRegistry"] = attr.ib(kw_only=True)

    @property
    def services_registry(self) -> "ServicesRegistry":
        return self._services_registry_ref.ref

    @abc.abstractmethod
    async def get_data_processor(
        self,
        dataset: Dataset,
        processor_type: ProcessorType,
        *,
        us_entry_buffer: USEntryBuffer,
        allow_cache_usage: bool = True,
        **kwargs: Any,
    ) -> OperationProcessorAsyncBase:
        pass

    @abc.abstractmethod
    async def close_async(self) -> None:
        pass


@attr.s(frozen=True)
class BaseClosableDataProcessorFactory(DataProcessorFactory):
    _created_data_processors: list[OperationProcessorAsyncBase] = attr.ib(factory=list, init=False)

    async def get_data_processor(
        self,
        dataset: Dataset,
        processor_type: ProcessorType,
        *,
        us_entry_buffer: USEntryBuffer,
        allow_cache_usage: bool = True,
        **kwargs: Any,
    ) -> OperationProcessorAsyncBase:
        processor = self._create_data_processor(
            dataset,
            processor_type,
            us_entry_buffer=us_entry_buffer,
            allow_cache_usage=allow_cache_usage,
            **kwargs,
        )
        self._created_data_processors.append(processor)

        await processor.start()
        return processor

    @abc.abstractmethod
    def _create_data_processor(  # type: ignore  # TODO: fix
        self,
        dataset: Dataset,
        processor_type: ProcessorType,
        *,
        us_entry_buffer: USEntryBuffer,
        allow_cache_usage: bool = True,
        **kwargs,
    ) -> OperationProcessorAsyncBase:
        pass

    async def close_async(self) -> None:
        async def close_processor(s: "OperationProcessorAsyncBase") -> None:
            # noinspection PyBroadException
            try:
                await s.end()
            except Exception:
                LOGGER.exception("Error during processor closing")

        await asyncio.gather(*[close_processor(processor) for processor in self._created_data_processors])

lib/dl_core/dl_core/services_registry/data_processor_factory_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
)


if TYPE_CHECKING:
    from dl_api_commons.base_models import RequestContextInfo
    from dl_core.us_connection_base import ConnectionBase
    from dl_core.us_dataset import Dataset
    from dl_core.us_manager.us_manager import USManagerBase


class EntityUsageNotAllowed(Exception):
    _user_reason: Optional[str]

    def __init__(self, *args: Any, user_reason: Optional[str] = None):
        super().__init__(*args)
        self._user_reason = user_reason

    @property
    def user_reason(self) -> str:
        return self._user_reason if self._user_reason else "Unknown reason"


class EntityUsageChecker(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def ensure_dataset_can_be_used(
        self,
        rci: RequestContextInfo,
        dataset: Dataset,
        us_manager: USManagerBase,
    ) -> None:
        pass

    @abc.abstractmethod
    def ensure_data_connection_can_be_used(self, rci: RequestContextInfo, conn: ConnectionBase):  # type: ignore  # TODO: fix
        pass

lib/dl_core/dl_core/services_registry/entity_checker.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import TYPE_CHECKING

import attr

from dl_core.connections_security.base import (
    ConnectionSecurityManager,
    InsecureConnectionSecurityManager,
)
from dl_core.services_registry.env_manager_factory_base import EnvManagerFactory


if TYPE_CHECKING:
    from dl_core.connection_models import ConnectOptions


@attr.s
class DefaultEnvManagerFactory(EnvManagerFactory):
    def make_security_manager(self) -> ConnectionSecurityManager:
        raise NotImplementedError()

    @classmethod
    def mutate_conn_opts(cls, conn_opts: ConnectOptions) -> ConnectOptions:
        return conn_opts


class InsecureEnvManagerFactory(DefaultEnvManagerFactory):
    def make_security_manager(self) -> ConnectionSecurityManager:
        return InsecureConnectionSecurityManager()

lib/dl_core/dl_core/services_registry/env_manager_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import TYPE_CHECKING

from dl_core.connections_security.base import ConnectionSecurityManager


if TYPE_CHECKING:
    from dl_core.connection_models import ConnectOptions


class EnvManagerFactory(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def make_security_manager(self) -> ConnectionSecurityManager:
        raise NotImplementedError

    @classmethod
    @abc.abstractmethod
    def mutate_conn_opts(cls, conn_opts: ConnectOptions) -> ConnectOptions:
        """
        Location for adding env-specific but non-app-specific alterations to
        the connect options.
        """
        raise NotImplementedError

lib/dl_core/dl_core/services_registry/env_manager_factory_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import logging
from types import TracebackType
from typing import (
    Any,
    ClassVar,
    Optional,
    Type,
    Union,
)

import aiohttp
import attr
import marshmallow as ma

from dl_api_commons.aiohttp.aiohttp_client import (
    BIAioHTTPClient,
    TCookies,
    THeaders,
)
from dl_constants.api_constants import DLHeadersCommon
from dl_constants.enums import UserDataType
from dl_core.db.elements import SchemaColumn
from dl_core.db.native_type_schema import OneOfNativeTypeSchema
from dl_utils.aio import await_sync


LOGGER = logging.getLogger(__name__)


RawSchemaType = list[SchemaColumn]


@attr.s(frozen=True, kw_only=True)
class FileSourceDesc:
    file_id: str = attr.ib()
    source_id: str = attr.ib()
    title: str = attr.ib()
    raw_schema: Optional[Union[RawSchemaType, list[dict[str, Any]]]] = attr.ib()
    preview_id: Optional[str] = attr.ib(default=None)


@attr.s(frozen=True, kw_only=True)
class GSheetsFileSourceDesc(FileSourceDesc):
    spreadsheet_id: Optional[str] = attr.ib()
    sheet_id: Optional[int] = attr.ib()
    first_line_is_header: Optional[bool] = attr.ib()


@attr.s(frozen=True)
class SourceInternalParams:
    preview_id: str = attr.ib()
    raw_schema: RawSchemaType = attr.ib()


@attr.s(frozen=True)
class SourcePreview:
    source_id: str = attr.ib()
    preview: list[list[Any]] = attr.ib()


class RawSchemaColumnSchema(ma.Schema):
    name = ma.fields.String()
    title = ma.fields.String()

    native_type = ma.fields.Nested(OneOfNativeTypeSchema, allow_none=True)

    user_type = ma.fields.Enum(UserDataType)
    description = ma.fields.String(dump_default="", allow_none=True)
    has_auto_aggregation = ma.fields.Boolean(dump_default=False, allow_none=True)
    lock_aggregation = ma.fields.Boolean(dump_default=False, allow_none=True)
    nullable = ma.fields.Boolean(dump_default=None, allow_none=True)

    @ma.post_load
    def make_column(self, data: dict, **kwargs: Any) -> SchemaColumn:
        return SchemaColumn(
            name=data["name"],
            title=data["title"],
            user_type=data["user_type"],
            native_type=data["native_type"],
            description=data.get("description", ""),
            has_auto_aggregation=data.get("has_auto_aggregation", False),
            lock_aggregation=data.get("lock_aggregation", False),
            nullable=data["nullable"],
        )


class SourcePreviewApiSchema(ma.Schema):
    preview_id = ma.fields.String()
    raw_schema = ma.fields.Nested(RawSchemaColumnSchema, many=True)


class SourceInternalParamsRequestSchema(ma.Schema):
    raw_schema = ma.fields.Nested(RawSchemaColumnSchema, many=True, allow_none=True)


class SourceInternalParamsResultSchema(SourceInternalParamsRequestSchema):
    preview_id = ma.fields.String()

    @ma.post_load
    def to_object(self, data: dict, **kwargs: Any) -> SourceInternalParams:
        return SourceInternalParams(
            preview_id=data["preview_id"],
            raw_schema=data["raw_schema"],
        )


class CleanupApiSchema(ma.Schema):
    tenant_id = ma.fields.String(required=True)


class UpdateConnectionDataRequestSchema(ma.Schema):
    class UpdateConnectionDataSourceSchema(ma.Schema):
        id = ma.fields.String(attribute="source_id")
        title = ma.fields.String()
        spreadsheet_id = ma.fields.String()
        sheet_id = ma.fields.Integer()
        first_line_is_header = ma.fields.Boolean()

    connection_id = ma.fields.String()
    save = ma.fields.Boolean()
    sources = ma.fields.Nested(UpdateConnectionDataSourceSchema, many=True)
    authorized = ma.fields.Boolean()
    tenant_id = ma.fields.String(allow_none=True)


@attr.s
class FileUploaderClient(BIAioHTTPClient):
    async def get_preview(self, src: FileSourceDesc) -> SourcePreview:
        path = f"api/v2/files/{src.file_id}/sources/{src.source_id}/preview"
        json_data = SourcePreviewApiSchema().dump(src)
        try:
            async with self.request("post", path=path, json_data=json_data, read_timeout_sec=5) as resp:
                resp_data = await resp.json()
                return SourcePreview(source_id=src.source_id, preview=resp_data["preview"])
        except aiohttp.ClientError:
            LOGGER.exception(f"Failed to get preview for file {src.file_id} source {src.source_id}")
            return SourcePreview(source_id=src.source_id, preview=[])

    async def get_preview_batch(self, file_sources: list[FileSourceDesc]) -> list[SourcePreview]:
        previews = await asyncio.gather(*[self.get_preview(src) for src in file_sources])
        return previews  # type: ignore

    def get_preview_batch_sync(self, file_sources: list[FileSourceDesc]) -> list[SourcePreview]:
        return await_sync(self.get_preview_batch(file_sources))

    async def get_internal_params(self, src: FileSourceDesc) -> SourceInternalParams:
        path = f"api/v2/files/{src.file_id}/sources/{src.source_id}/internal_params"
        json_data = SourceInternalParamsRequestSchema().dump(src)
        try:
            async with self.request("post", path=path, json_data=json_data, read_timeout_sec=15) as resp:
                resp_data = await resp.json()
                internal_params = SourceInternalParamsResultSchema().load(resp_data)
                return internal_params
        except aiohttp.ClientError:
            LOGGER.exception(f"Failed to get raw_schema for file {src.file_id} source {src.source_id}")
            raise

    async def get_internal_params_batch(self, file_sources: list[FileSourceDesc]) -> list[SourceInternalParams]:
        internal_params = await asyncio.gather(*[self.get_internal_params(src) for src in file_sources])
        return internal_params  # type: ignore

    async def cleanup_tenant(self, tenant_id: str) -> None:
        path = "api/v2/cleanup"
        json_data = CleanupApiSchema().dump(dict(tenant_id=tenant_id))
        try:
            async with self.request("post", path=path, json_data=json_data, read_timeout_sec=5):
                LOGGER.info(f"Scheduled cleanup for tenant id {tenant_id}")
        except aiohttp.ClientError:
            LOGGER.exception(f"Failed to call cleanup for tenant id {tenant_id}")
            raise

    def cleanup_tenant_sync(self, tenant_id: str) -> None:
        await_sync(self.cleanup_tenant(tenant_id))

    async def update_connection_data_internal(
        self,
        conn_id: str,
        sources: list[GSheetsFileSourceDesc],
        authorized: bool,
        tenant_id: Optional[str],
    ) -> None:
        path = "/api/v2/update_connection_data_internal"
        json_data = UpdateConnectionDataRequestSchema().dump(
            dict(
                connection_id=conn_id,
                save=True,
                sources=sources,
                authorized=authorized,
                tenant_id=tenant_id,
            )
        )
        try:
            async with self.request("post", path=path, json_data=json_data, read_timeout_sec=5):
                LOGGER.info(f"Scheduled update for connection id {conn_id}")
        except aiohttp.ClientError:
            LOGGER.exception(f"Failed to call update for connection id {conn_id}")
            raise

    def close_sync(self) -> None:
        return await_sync(self.close())

    def __enter__(self) -> FileUploaderClient:
        return self

    def __exit__(
        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
    ) -> None:
        self.close_sync()


@attr.s
class FileUploaderSettings:
    base_url: str = attr.ib()
    master_token: str = attr.ib()


@attr.s
class FileUploaderClientFactory:
    _file_uploader_settings: FileUploaderSettings = attr.ib()

    _file_uploader_client_cls: ClassVar[Type[FileUploaderClient]] = FileUploaderClient  # tests mockup point

    def get_client(self, headers: Optional[THeaders] = None, cookies: Optional[TCookies] = None) -> FileUploaderClient:
        full_headers = headers.copy() if headers is not None else {}
        full_headers.update(
            {
                DLHeadersCommon.FILE_UPLOADER_MASTER_TOKEN.value: self._file_uploader_settings.master_token,
            }
        )
        client_cls = self._file_uploader_client_cls
        return client_cls(
            base_url=self._file_uploader_settings.base_url,
            headers=full_headers,
            cookies=cookies.copy() if cookies is not None else {},
            session=None,
        )

lib/dl_core/dl_core/services_registry/file_uploader_client_factory.py
datalens-tech/datalens-backend
import abc
from typing import TYPE_CHECKING

import attr

from dl_core.rls import BaseSubjectResolver
from dl_core.utils import FutureRef


if TYPE_CHECKING:
    from dl_core.services_registry.top_level import ServicesRegistry  # noqa


@attr.s
class InstallationSpecificServiceRegistry(metaclass=abc.ABCMeta):
    _service_registry_ref: FutureRef["ServicesRegistry"] = attr.ib()

    @property
    def service_registry(self) -> "ServicesRegistry":
        return self._service_registry_ref.ref

    @abc.abstractmethod
    async def get_subject_resolver(self) -> BaseSubjectResolver:
        raise NotImplementedError()


class InstallationSpecificServiceRegistryFactory(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_inst_specific_sr(self, sr_ref: FutureRef["ServicesRegistry"]) -> InstallationSpecificServiceRegistry:
        raise NotImplementedError()

lib/dl_core/dl_core/services_registry/inst_specific_sr.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_configs.settings_submodels import RedisSettings


@attr.s
class RQECachesSetting:
    redis_settings: Optional[RedisSettings] = attr.ib(default=None)
    caches_ttl: Optional[int] = attr.ib(default=None)

lib/dl_core/dl_core/services_registry/rqe_caches.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
import logging
from typing import (
    TYPE_CHECKING,
    List,
)

import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_constants.enums import SelectorType
from dl_core.data_processing.selectors.dataset_base import DatasetDataSelectorAsyncBase
from dl_core.data_processing.selectors.dataset_cached import CachedDatasetDataSelectorAsync
from dl_core.data_processing.selectors.dataset_cached_lazy import LazyCachedDatasetDataSelectorAsync
from dl_core.us_dataset import Dataset
from dl_core.us_manager.local_cache import USEntryBuffer
from dl_core.utils import FutureRef


if TYPE_CHECKING:
    from dl_core.services_registry.top_level import ServicesRegistry


LOGGER = logging.getLogger(__name__)


class SelectorFactory(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_dataset_selector(
        self,
        dataset: Dataset,
        selector_type: SelectorType,
        *,
        us_entry_buffer: USEntryBuffer,
        allow_cache_usage: bool = True,
    ) -> DatasetDataSelectorAsyncBase:
        pass

    @abc.abstractmethod
    async def close_async(self) -> None:
        pass


@attr.s(frozen=True)
class BaseClosableSelectorFactory(SelectorFactory, metaclass=abc.ABCMeta):
    _services_registry_ref: FutureRef[ServicesRegistry] = attr.ib()
    _created_dataset_selectors: List[DatasetDataSelectorAsyncBase] = attr.ib(factory=list)

    @property
    def services_registry(self) -> ServicesRegistry:
        return self._services_registry_ref.ref

    @property
    def rci(self) -> RequestContextInfo:
        return self.services_registry.rci

    def get_dataset_selector(
        self,
        dataset: Dataset,
        selector_type: SelectorType,
        *,
        us_entry_buffer: USEntryBuffer,
        allow_cache_usage: bool = True,
    ) -> DatasetDataSelectorAsyncBase:
        selector = self._create_dataset_selector(dataset, selector_type, us_entry_buffer=us_entry_buffer)
        self._created_dataset_selectors.append(selector)

        return selector

    @abc.abstractmethod
    def _create_dataset_selector(
        self,
        dataset: Dataset,
        selector_type: SelectorType,
        *,
        us_entry_buffer: USEntryBuffer,
        allow_cache_usage: bool = True,
    ) -> DatasetDataSelectorAsyncBase:
        pass

    async def close_async(self) -> None:
        async def close_selector(s: "DatasetDataSelectorAsyncBase") -> None:
            # noinspection PyBroadException
            try:
                await s.close()
            except Exception:
                LOGGER.exception("Error during selector closing")

        await asyncio.gather(*[close_selector(selector) for selector in self._created_dataset_selectors])


@attr.s(frozen=True)
class DefaultSelectorFactory(BaseClosableSelectorFactory):
    _is_bleeding_edge_user: bool = attr.ib(default=False)

    def _create_dataset_selector(
        self,
        dataset: Dataset,
        selector_type: SelectorType,
        *,
        us_entry_buffer: USEntryBuffer,
        allow_cache_usage: bool = True,
    ) -> DatasetDataSelectorAsyncBase:
        if selector_type == SelectorType.CACHED:
            return CachedDatasetDataSelectorAsync(  # type: ignore  # TODO: fix
                dataset=dataset,
                service_registry=self.services_registry,
                allow_cache_usage=allow_cache_usage,
                is_bleeding_edge_user=self._is_bleeding_edge_user,
                us_entry_buffer=us_entry_buffer,
            )
        elif selector_type == SelectorType.CACHED_LAZY:
            return LazyCachedDatasetDataSelectorAsync(  # type: ignore  # TODO: fix
                dataset=dataset,
                service_registry=self.services_registry,
                allow_cache_usage=allow_cache_usage,
                is_bleeding_edge_user=self._is_bleeding_edge_user,
                us_entry_buffer=us_entry_buffer,
            )
        else:
            raise NotImplementedError(f"Creation of selector with type {selector_type} is not supported")

lib/dl_core/dl_core/services_registry/selector_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Dict,
    FrozenSet,
    Generic,
    Optional,
    Sequence,
    Type,
    TypeVar,
)

from arq.connections import RedisSettings as ArqRedisSettings
import attr
from typing_extensions import final

from dl_api_commons.reporting.registry import ReportingRegistry
from dl_configs.enums import RequiredService
from dl_constants.enums import ProcessorType
from dl_core.services_registry.conn_executor_factory import DefaultConnExecutorFactory
from dl_core.services_registry.file_uploader_client_factory import (
    FileUploaderClientFactory,
    FileUploaderSettings,
)
from dl_core.services_registry.rqe_caches import RQECachesSetting
from dl_core.services_registry.selector_factory import DefaultSelectorFactory
from dl_core.services_registry.top_level import (
    DefaultServicesRegistry,
    ServicesRegistry,
)
from dl_core.us_manager.mutation_cache.usentry_mutation_cache_factory import USEntryMutationCacheFactory
from dl_core.utils import FutureRef
from dl_task_processor.processor import ARQTaskProcessorFactory


if TYPE_CHECKING:
    from redis.asyncio import Redis

    from dl_api_commons.base_models import RequestContextInfo
    from dl_configs.connectors_settings import ConnectorSettingsBase
    from dl_configs.rqe import RQEConfig
    from dl_constants.enums import ConnectionType
    from dl_core.aio.web_app_services.data_processing.data_processor import DataProcessorService
    from dl_core.data_processing.cache.primitives import CacheTTLConfig
    from dl_core.services_registry.entity_checker import EntityUsageChecker
    from dl_core.services_registry.env_manager_factory_base import EnvManagerFactory
    from dl_core.services_registry.inst_specific_sr import InstallationSpecificServiceRegistryFactory
    from dl_core.services_registry.typing import ConnectOptionsFactory
    from dl_core.us_connection_base import ExecutorBasedMixin


LOGGER = logging.getLogger(__name__)
SERVICE_REGISTRY_TV = TypeVar("SERVICE_REGISTRY_TV", bound=ServicesRegistry)


class SRFactory(Generic[SERVICE_REGISTRY_TV], metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def make_service_registry(
        self,
        request_context_info: RequestContextInfo,
        mutations_redis_client_factory: Optional[Callable[[bool], Optional[Redis]]] = None,
        mutations_cache_factory: Optional[USEntryMutationCacheFactory] = None,
        reporting_registry: Optional[ReportingRegistry] = None,
        caches_redis_client_factory: Optional[Callable[[bool], Optional[Redis]]] = None,
        data_processor_service_factory: Optional[Callable[[ProcessorType], DataProcessorService]] = None,
    ) -> SERVICE_REGISTRY_TV:
        """
        :param request_context_info: RCI to pass to services registry.
        :param caches_redis_client_factory: Callable that will return Redis client for caches Redis.
            Moved to argument because Redis client factory may be instantiated after instantiation of SRFactory.
        :param pg_pool_factory: Factory that produces PG pools (BasePgPoolWrapper subclass instances).
        :param reporting_registry: ReportingRegistry associated with a registry
        :return: New services registry.
        """
        pass


@attr.s
class DefaultSRFactory(SRFactory[SERVICE_REGISTRY_TV]):  # type: ignore  # TODO: fix
    rqe_config: Optional[RQEConfig] = attr.ib()
    async_env: bool = attr.ib()
    env_manager_factory: EnvManagerFactory = attr.ib()
    # Config options
    default_cache_ttl_config: Optional[CacheTTLConfig] = attr.ib(default=None)
    bleeding_edge_users: Sequence[str] = attr.ib(default=())
    conn_cls_whitelist: Optional[FrozenSet[Type[ExecutorBasedMixin]]] = attr.ib(default=None)
    connect_options_factory: Optional[ConnectOptionsFactory] = attr.ib(default=None)
    entity_usage_checker: Optional[EntityUsageChecker] = attr.ib(default=None)
    connectors_settings: dict[ConnectionType, ConnectorSettingsBase] = attr.ib(factory=dict)
    file_uploader_settings: Optional[FileUploaderSettings] = attr.ib(default=None)
    redis_pool_settings: Optional[ArqRedisSettings] = attr.ib(default=None)
    force_non_rqe_mode: bool = attr.ib(default=False)
    rqe_caches_settings: Optional[RQECachesSetting] = attr.ib(default=None)
    required_services: set[RequiredService] = attr.ib(factory=set)
    inst_specific_sr_factory: Optional[InstallationSpecificServiceRegistryFactory] = attr.ib(default=None)

    service_registry_cls: ClassVar[Type[SERVICE_REGISTRY_TV]] = DefaultServicesRegistry  # type: ignore  # TODO: fix

    def is_bleeding_edge_user(self, request_context_info: RequestContextInfo) -> bool:
        return request_context_info.user_name in self.bleeding_edge_users

    def make_conn_executor_factory(
        self,
        request_context_info: RequestContextInfo,
        sr_ref: FutureRef[SERVICE_REGISTRY_TV],
    ) -> DefaultConnExecutorFactory:
        is_bleeding_edge_user = self.is_bleeding_edge_user(request_context_info)
        if is_bleeding_edge_user:
            LOGGER.info("ATTENTION! It's bleeding edge user")
        return DefaultConnExecutorFactory(
            async_env=self.async_env,
            tpe=None,
            conn_sec_mgr=self.env_manager_factory.make_security_manager(),
            rqe_config=self.rqe_config,
            services_registry_ref=sr_ref,  # type: ignore  # TODO: fix
            is_bleeding_edge_user=is_bleeding_edge_user,
            conn_cls_whitelist=self.conn_cls_whitelist,
            connect_options_factory=self.connect_options_factory,
            # Env-specific alterations of connect_options:
            connect_options_mutator=self.env_manager_factory.mutate_conn_opts,
            entity_usage_checker=self.entity_usage_checker,
            force_non_rqe_mode=self.force_non_rqe_mode,
        )

    def additional_sr_constructor_kwargs(
        self,
        request_context_info: RequestContextInfo,
        sr_ref: FutureRef[ServicesRegistry],
    ) -> Dict[str, Any]:
        return dict()

    @final
    def make_service_registry(
        self,
        request_context_info: RequestContextInfo,
        mutations_redis_client_factory: Optional[Callable[[bool], Optional[Redis]]] = None,
        mutations_cache_factory: Optional[USEntryMutationCacheFactory] = None,
        reporting_registry: Optional[ReportingRegistry] = None,
        # TODO: refactor usage of redis and pg here
        #  (some kind of multi-purpose factory instead of separate getters)
        caches_redis_client_factory: Optional[Callable[[bool], Optional[Redis]]] = None,
        data_processor_service_factory: Optional[Callable[[ProcessorType], DataProcessorService]] = None,
    ) -> SERVICE_REGISTRY_TV:
        sr_ref: FutureRef[ServicesRegistry] = FutureRef()
        sr = self.service_registry_cls(  # type: ignore  # TODO: fix
            default_cache_ttl_config=self.default_cache_ttl_config,
            rci=request_context_info,
            conn_exec_factory=self.make_conn_executor_factory(request_context_info, sr_ref),  # type: ignore  # TODO: fix
            caches_redis_client_factory=caches_redis_client_factory,
            mutations_redis_client_factory=mutations_redis_client_factory,
            mutations_cache_factory=mutations_cache_factory,
            data_processor_service_factory=data_processor_service_factory,
            connectors_settings=self.connectors_settings,
            reporting_registry=reporting_registry,
            selector_factory=DefaultSelectorFactory(
                services_registry_ref=sr_ref,
                is_bleeding_edge_user=self.is_bleeding_edge_user(request_context_info),
            ),
            file_uploader_client_factory=FileUploaderClientFactory(
                self.file_uploader_settings,
            )
            if self.file_uploader_settings
            else None,
            task_processor_factory=ARQTaskProcessorFactory(
                redis_pool_settings=self.redis_pool_settings,
            )
            if self.redis_pool_settings
            else None,
            rqe_caches_settings=self.rqe_caches_settings,
            required_services=self.required_services,
            inst_specific_sr=(
                self.inst_specific_sr_factory.get_inst_specific_sr(sr_ref)
                if self.inst_specific_sr_factory is not None
                else None
            ),
            **self.additional_sr_constructor_kwargs(request_context_info, sr_ref),
        )
        sr_ref.fulfill(sr)
        return sr

lib/dl_core/dl_core/services_registry/sr_factories.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Callable,
    Optional,
    Type,
    TypeVar,
)

import attr
import redis.asyncio

from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.reporting.registry import ReportingRegistry
from dl_configs.enums import RequiredService
from dl_constants.enums import ProcessorType
from dl_core.data_processing.cache.primitives import CacheTTLConfig
from dl_core.data_source.collection import DataSourceCollectionFactory
from dl_core.services_registry.cache_engine_factory import (
    CacheEngineFactory,
    DefaultCacheEngineFactory,
)
from dl_core.services_registry.compute_executor import ComputeExecutorTPE
from dl_core.services_registry.conn_executor_factory_base import ConnExecutorFactory
from dl_core.services_registry.data_processor_factory import DefaultDataProcessorFactory
from dl_core.services_registry.data_processor_factory_base import BaseClosableDataProcessorFactory
from dl_core.services_registry.inst_specific_sr import InstallationSpecificServiceRegistry
from dl_core.services_registry.rqe_caches import RQECachesSetting
from dl_core.services_registry.selector_factory import (
    DefaultSelectorFactory,
    SelectorFactory,
)
from dl_core.us_manager.mutation_cache.engine_factory import (
    DefaultMutationCacheEngineFactory,
    MutationCacheEngineFactory,
)
from dl_core.us_manager.mutation_cache.usentry_mutation_cache import GenericCacheEngine
from dl_core.us_manager.mutation_cache.usentry_mutation_cache_factory import USEntryMutationCacheFactory
from dl_core.utils import FutureRef
from dl_task_processor.processor import TaskProcessorFactory


if TYPE_CHECKING:
    from dl_configs.connectors_settings import ConnectorSettingsBase
    from dl_constants.enums import ConnectionType
    from dl_core.aio.web_app_services.data_processing.data_processor import DataProcessorService
    from dl_core.services_registry.compute_executor import ComputeExecutor
    from dl_core.services_registry.file_uploader_client_factory import FileUploaderClientFactory
    from dl_core.us_manager.local_cache import USEntryBuffer


_ISSR_T = TypeVar("_ISSR_T", bound=InstallationSpecificServiceRegistry)


class ServicesRegistry(metaclass=abc.ABCMeta):
    """
    Composition of services/service factories that are used in US entries methods:
     * Connection executors
     * Cache engines
     * ...etc
    """

    @property
    @abc.abstractmethod
    def rci(self) -> RequestContextInfo:
        pass

    @property
    @abc.abstractmethod
    def default_cache_ttl_config(self) -> CacheTTLConfig:
        pass

    @abc.abstractmethod
    def get_conn_executor_factory(self) -> ConnExecutorFactory:
        pass

    @abc.abstractmethod
    def get_caches_redis_client(self, allow_slave: bool = False) -> Optional[redis.asyncio.Redis]:
        pass

    @abc.abstractmethod
    def get_mutations_redis_client(self, allow_slave: bool = False) -> Optional[redis.asyncio.Redis]:
        pass

    @abc.abstractmethod
    def get_mutation_cache_factory(self) -> Optional[USEntryMutationCacheFactory]:
        pass

    @abc.abstractmethod
    def get_reporting_registry(self) -> ReportingRegistry:
        pass

    @abc.abstractmethod
    def get_compute_executor(self) -> ComputeExecutor:
        pass

    @abc.abstractmethod
    def get_selector_factory(self) -> SelectorFactory:
        pass

    @abc.abstractmethod
    def get_data_processor_service_factory(self) -> Optional[Callable[[ProcessorType], DataProcessorService]]:
        pass

    @abc.abstractmethod
    def get_data_processor_factory(self) -> BaseClosableDataProcessorFactory:
        pass

    @abc.abstractmethod
    def close(self) -> None:
        pass

    @abc.abstractmethod
    async def close_async(self) -> None:
        pass

    @abc.abstractmethod
    def get_cache_engine_factory(self, allow_slave: bool = False) -> Optional[CacheEngineFactory]:
        pass

    @abc.abstractmethod
    def get_mutation_cache_engine_factory(
        self, cache_type: Type[GenericCacheEngine]
    ) -> MutationCacheEngineFactory:  # type: ignore  # TODO: fix
        pass

    @abc.abstractmethod
    def get_connectors_settings(self, conn_type: ConnectionType) -> Optional[ConnectorSettingsBase]:
        pass

    @abc.abstractmethod
    def get_data_source_collection_factory(self, us_entry_buffer: USEntryBuffer) -> DataSourceCollectionFactory:
        pass

    @abc.abstractmethod
    def get_file_uploader_client_factory(self) -> FileUploaderClientFactory:
        pass

    @abc.abstractmethod
    def get_task_processor_factory(self) -> TaskProcessorFactory:
        pass

    @abc.abstractmethod
    def get_rqe_caches_settings(self) -> Optional[RQECachesSetting]:
        pass

    @abc.abstractmethod
    def get_required_services(self) -> set[RequiredService]:
        pass

    @abc.abstractmethod
    def get_installation_specific_service_registry(self, issr_cls: Type[_ISSR_T]) -> _ISSR_T:
        pass


@attr.s(hash=False)
class DefaultServicesRegistry(ServicesRegistry):  # type: ignore  # TODO: fix
    _rci: RequestContextInfo = attr.ib()
    _reporting_registry: ReportingRegistry = attr.ib()
    _mutations_cache_factory: Optional[USEntryMutationCacheFactory] = attr.ib()
    _mutations_redis_client_factory: Optional[Callable[[bool], Optional[redis.asyncio.Redis]]] = attr.ib(default=None)
    _default_cache_ttl_config: Optional[CacheTTLConfig] = attr.ib(default=None)
    _conn_exec_factory: Optional[ConnExecutorFactory] = attr.ib(default=None)
    _caches_redis_client_factory: Optional[Callable[[bool], Optional[redis.asyncio.Redis]]] = attr.ib(default=None)
    _compute_executor: ComputeExecutor = attr.ib()
    _selector_factory: SelectorFactory = attr.ib()
    _cache_engine_factory: CacheEngineFactory = attr.ib()
    _mutation_cache_engine_factory: MutationCacheEngineFactory = attr.ib(default=None)
    _data_processor_service_factory: Optional[Callable[[ProcessorType], DataProcessorService]] = attr.ib(default=None)
    _data_processor_factory: BaseClosableDataProcessorFactory = attr.ib()
    _connectors_settings: dict[ConnectionType, ConnectorSettingsBase] = attr.ib(default=None)
    _file_uploader_client_factory: Optional[FileUploaderClientFactory] = attr.ib(default=None)
    _task_processor_factory: Optional[TaskProcessorFactory] = attr.ib(default=None)
    _rqe_caches_settings: Optional[RQECachesSetting] = attr.ib(default=None)
    _required_services: set[RequiredService] = attr.ib(factory=set)
    _inst_specific_sr: Optional[InstallationSpecificServiceRegistry] = attr.ib(default=None)

    @_compute_executor.default  # noqa
    def _default_compute_executor(self) -> ComputeExecutor:
        return ComputeExecutorTPE()  # type: ignore  # Incompatible return value type (got "ComputeExecutorTPE", expected "ComputeExecutor")

    @_selector_factory.default  # noqa
    def _default_selector_factory(self) -> SelectorFactory:
        return DefaultSelectorFactory(services_registry_ref=FutureRef.fulfilled(self))

    @_cache_engine_factory.default  # noqa
    def _default_cache_engine_factory(self) -> CacheEngineFactory:
        return DefaultCacheEngineFactory(services_registry_ref=FutureRef.fulfilled(self))

    @_data_processor_factory.default  # noqa
    def _default_data_processor_factory(self) -> BaseClosableDataProcessorFactory:
        return DefaultDataProcessorFactory(  # type: ignore  # TODO: fix
            services_registry_ref=FutureRef.fulfilled(self),
        )

    @property
    def rci(self) -> RequestContextInfo:
        return self._rci

    @property
    def default_cache_ttl_config(self) -> CacheTTLConfig:
        if self._default_cache_ttl_config is None:
            self._default_cache_ttl_config = CacheTTLConfig()
        return self._default_cache_ttl_config

    def get_conn_executor_factory(self) -> ConnExecutorFactory:
        if self._conn_exec_factory is None:
            raise ValueError("ConnExecutor factory was not injected in this service registry")
        return self._conn_exec_factory

    def get_caches_redis_client(self, allow_slave: bool = False) -> Optional[redis.asyncio.Redis]:
        if self._caches_redis_client_factory is not None:
            return self._caches_redis_client_factory(allow_slave)
        return None

    def get_mutations_redis_client(self, allow_slave: bool = False) -> Optional[redis.asyncio.Redis]:
        if self._mutations_redis_client_factory is not None:
            return self._mutations_redis_client_factory(allow_slave)
        return None

    def get_reporting_registry(self) -> ReportingRegistry:
        return self._reporting_registry

    def get_compute_executor(self) -> ComputeExecutor:
        return self._compute_executor

    def get_selector_factory(self) -> SelectorFactory:
        return self._selector_factory

    def get_cache_engine_factory(self) -> Optional[CacheEngineFactory]:  # type: ignore  # TODO: fix
        return self._cache_engine_factory

    def get_mutation_cache_factory(self) -> Optional[USEntryMutationCacheFactory]:
        return self._mutations_cache_factory

    def get_mutation_cache_engine_factory(
        self, cache_type: Type[GenericCacheEngine]
    ) -> MutationCacheEngineFactory:  # type: ignore  # TODO: fix
        # TODO: Save already created CacheEngine's?
        return DefaultMutationCacheEngineFactory(services_registry_ref=FutureRef.fulfilled(self), cache_type=cache_type)

    def get_data_processor_service_factory(self) -> Optional[Callable[[ProcessorType], DataProcessorService]]:
        return self._data_processor_service_factory

    def get_data_processor_factory(self) -> BaseClosableDataProcessorFactory:
        return self._data_processor_factory

    def get_rqe_caches_settings(self) -> Optional[RQECachesSetting]:
        return self._rqe_caches_settings

    def get_required_services(self) -> set[RequiredService]:
        return self._required_services

    def get_connectors_settings(self, conn_type: ConnectionType) -> Optional[ConnectorSettingsBase]:
        return self._connectors_settings.get(conn_type)

    def get_data_source_collection_factory(self, us_entry_buffer: USEntryBuffer) -> DataSourceCollectionFactory:
        return DataSourceCollectionFactory(us_entry_buffer=us_entry_buffer)

    def get_file_uploader_client_factory(self) -> FileUploaderClientFactory:
        if self._file_uploader_client_factory is not None:
            return self._file_uploader_client_factory
        raise ValueError("FileUploaderClientFactory hasn't been initialized.")

    def get_task_processor_factory(self) -> TaskProcessorFactory:
        if self._task_processor_factory is not None:
            return self._task_processor_factory
        raise ValueError("TaskProcessorFactory hasn't been initialized.")

    def get_installation_specific_service_registry(self, issr_cls: Type[_ISSR_T]) -> _ISSR_T:
        if isinstance(self._inst_specific_sr, issr_cls):
            return self._inst_specific_sr
        raise ValueError(f"Invalid ISST type: expected {issr_cls}, got {type(self._inst_specific_sr)}")

    def close(self) -> None:
        if self._conn_exec_factory is not None:
            self._conn_exec_factory.close_sync()
        # FIXME: others? (TPE if necessary)
        if self._compute_executor is not None:
            self._compute_executor.close()
        if self._task_processor_factory is not None:
            self._task_processor_factory.cleanup()

    async def close_async(self) -> None:
        if self._data_processor_factory is not None:
            # Processor factory must be closed before selectors
            # because processors might use selectors while closing
            await self._data_processor_factory.close_async()
        if self._selector_factory is not None:
            # Should not parallelize with selectors factory closing
            # because its selectors may use CE in closing procedure
            await self._selector_factory.close_async()
        if self._conn_exec_factory is not None:
            await self._conn_exec_factory.close_async()
        if self._task_processor_factory is not None:
            await self._task_processor_factory.cleanup_async()

    def clone(self, **kwargs):  # type: ignore  # TODO: fix
        return attr.evolve(self, **kwargs)


@attr.s
class DummyServiceRegistry(ServicesRegistry):
    """
    Service registry implementation that holds RCI for simple test cases (actual services are not required)
    """

    _rci: RequestContextInfo = attr.ib()
    _default_cache_ttl_config: Optional[CacheTTLConfig] = attr.ib(default=None)
    _inst_specific_sr: Optional[InstallationSpecificServiceRegistry] = attr.ib(default=None)

    NOT_IMPLEMENTED_MSG = "DummyServiceRegistry acts only as RCI container for underlying services"

    @property
    def rci(self) -> RequestContextInfo:
        return self._rci

    @property
    def default_cache_ttl_config(self) -> CacheTTLConfig:
        if self._default_cache_ttl_config is None:
            self._default_cache_ttl_config = CacheTTLConfig()
        return self._default_cache_ttl_config

    def get_conn_executor_factory(self) -> ConnExecutorFactory:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_caches_redis_client(self) -> Optional[redis.asyncio.Redis]:  # type: ignore  # TODO: fix
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_mutations_redis_client(self) -> Optional[redis.asyncio.Redis]:  # type: ignore  # TODO: fix
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_reporting_registry(self) -> ReportingRegistry:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_compute_executor(self) -> ComputeExecutor:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_selector_factory(self) -> SelectorFactory:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_cache_engine_factory(self) -> Optional[CacheEngineFactory]:  # type: ignore  # TODO: fix
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_mutation_cache_factory(self) -> Optional[USEntryMutationCacheFactory]:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_mutation_cache_engine_factory(
        self, cache_type: Type[GenericCacheEngine]
    ) -> MutationCacheEngineFactory:  # type: ignore  # TODO: fix
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_data_processor_service_factory(self) -> Optional[Callable[[ProcessorType], DataProcessorService]]:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_data_processor_factory(self) -> BaseClosableDataProcessorFactory:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_connectors_settings(self, conn_type: ConnectionType) -> Optional[ConnectorSettingsBase]:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_data_source_collection_factory(self, us_entry_buffer: USEntryBuffer) -> DataSourceCollectionFactory:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_file_uploader_client_factory(self) -> FileUploaderClientFactory:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_task_processor_factory(self) -> TaskProcessorFactory:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_rqe_caches_settings(self) -> Optional[RQECachesSetting]:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_required_services(self) -> set[RequiredService]:
        raise NotImplementedError(self.NOT_IMPLEMENTED_MSG)

    def get_installation_specific_service_registry(self, issr_cls: Type[_ISSR_T]) -> _ISSR_T:
        if isinstance(self._inst_specific_sr, issr_cls):
            return self._inst_specific_sr
        raise ValueError(f"Invalid ISST type: expected {issr_cls}, got {type(self._inst_specific_sr)}")

    def close(self) -> None:
        pass

    async def close_async(self) -> None:
        pass

lib/dl_core/dl_core/services_registry/top_level.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Callable,
    Optional,
)

from dl_core.connection_models import ConnectOptions
from dl_core.us_connection_base import ExecutorBasedMixin


ConnectOptionsFactory = Callable[[ExecutorBasedMixin], Optional[ConnectOptions]]
ConnectOptionsMutator = Callable[[ConnectOptions], ConnectOptions]

lib/dl_core/dl_core/services_registry/typing.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from contextlib import contextmanager
from json.decoder import JSONDecodeError
import logging
import re
import time
from typing import (
    Any,
    ClassVar,
    Generator,
    Iterable,
    NamedTuple,
    Optional,
    Type,
    Union,
    cast,
)

import attr
import requests
from requests.exceptions import HTTPError

from dl_api_commons.base_models import (
    AuthData,
    TenantCommon,
    TenantDef,
)
from dl_api_commons.logging import mask_sensitive_fields_by_name_in_json_recursive
from dl_api_commons.tracing import get_current_tracing_headers
from dl_api_commons.utils import (
    stringify_dl_cookies,
    stringify_dl_headers,
)
from dl_constants.api_constants import (
    DLCookies,
    DLHeaders,
    DLHeadersCommon,
)
from dl_core.base_models import EntryLocation
from dl_core.enums import USApiType
import dl_core.exc as exc
from dl_core.utils import get_retriable_requests_session


LOGGER = logging.getLogger(__name__)


class USClientHTTPExceptionWrapper(Exception):
    pass


@attr.s
class USAuthContextBase:
    DEFAULT_US_PREFIX: ClassVar[USApiType] = USApiType.v1
    IS_TENANT_ID_MUTABLE: ClassVar[bool] = False

    dl_component: str = attr.ib(default="backend", kw_only=True)

    def is_tenant_id_mutable(self) -> bool:
        return self.IS_TENANT_ID_MUTABLE

    def get_folder_id(self) -> Optional[str]:
        return None

    def get_default_prefix(self) -> str:
        val = self.DEFAULT_US_PREFIX.value
        if isinstance(val, str):
            return val

        raise ValueError(f"US prefix enum member {self.DEFAULT_US_PREFIX} is not a string")

    @abc.abstractmethod
    def get_tenant(self) -> TenantDef:
        raise NotImplementedError()

    @abc.abstractmethod
    def get_outbound_headers(self, include_tenancy: bool = True) -> dict[DLHeaders, str]:
        raise NotImplementedError()

    @abc.abstractmethod
    def get_outbound_cookies(self) -> dict[DLCookies, str]:
        raise NotImplementedError()


@attr.s
class USAuthContextRegular(USAuthContextBase):
    auth_data: AuthData = attr.ib(kw_only=True)
    tenant: TenantDef = attr.ib(kw_only=True)
    # Raw header value. None is interpreted as no header
    dl_sudo: Optional[str] = attr.ib(kw_only=True)
    # Raw header value. None is interpreted as no header
    dl_allow_super_user: Optional[str] = attr.ib(kw_only=True)

    def get_tenant(self) -> TenantDef:
        return self.tenant

    def get_outbound_headers(self, include_tenancy: bool = True) -> dict[DLHeaders, str]:
        flags: dict[DLHeaders, Optional[str]] = {
            DLHeadersCommon.SUDO: self.dl_sudo,
            DLHeadersCommon.ALLOW_SUPERUSER: self.dl_allow_super_user,
        }

        return {
            **(self.tenant.get_outbound_tenancy_headers() if include_tenancy else {}),
            **self.auth_data.get_headers(),
            **{name: val for name, val in flags.items() if val is not None},
        }

    def get_outbound_cookies(self) -> dict[DLCookies, str]:
        return self.auth_data.get_cookies()


@attr.s(frozen=True)
class USAuthContextPublic(USAuthContextBase):
    DEFAULT_US_PREFIX = USApiType.public
    DEFAULT_TENANT = TenantCommon()

    us_public_token: str = attr.ib(repr=False)

    def get_tenant(self) -> TenantDef:
        return self.DEFAULT_TENANT

    def get_outbound_headers(self, include_tenancy: bool = True) -> dict[DLHeaders, str]:
        return {
            DLHeadersCommon.US_PUBLIC_TOKEN: self.us_public_token,
        }

    def get_outbound_cookies(self) -> dict[DLCookies, str]:
        return {}


@attr.s(frozen=True)
class USAuthContextMaster(USAuthContextBase):
    DEFAULT_US_PREFIX = USApiType.private
    DEFAULT_TENANT = TenantCommon()
    IS_TENANT_ID_MUTABLE = True

    us_master_token: str = attr.ib(repr=False)

    def get_tenant(self) -> TenantDef:
        return self.DEFAULT_TENANT

    def get_outbound_headers(self, include_tenancy: bool = True) -> dict[DLHeaders, str]:
        return {DLHeadersCommon.US_MASTER_TOKEN: self.us_master_token}

    def get_outbound_cookies(self) -> dict[DLCookies, str]:
        return {}


@attr.s(frozen=True)
class USAuthContextNoAuth(USAuthContextBase):
    DEFAULT_TENANT = TenantCommon()

    def get_tenant(self) -> TenantDef:
        return self.DEFAULT_TENANT

    def get_outbound_headers(self, include_tenancy: bool = True) -> dict[DLHeaders, str]:
        return {}

    def get_outbound_cookies(self) -> dict[DLCookies, str]:
        return {}


@attr.s(frozen=True)
class USAuthContextEmbed(USAuthContextBase):
    DEFAULT_US_PREFIX = USApiType.embeds

    auth_data: AuthData = attr.ib(kw_only=True)
    tenant: TenantDef = attr.ib(kw_only=True)

    def get_tenant(self) -> TenantDef:
        return self.tenant

    def get_outbound_headers(self, include_tenancy: bool = True) -> dict[DLHeaders, str]:
        return {
            **(self.tenant.get_outbound_tenancy_headers() if include_tenancy else {}),
            **self.auth_data.get_headers(),
        }

    def get_outbound_cookies(self) -> dict[DLCookies, str]:
        return self.auth_data.get_cookies()


# noinspection PyCallByClass
class UStorageClientBase:
    class RequestAdapter(metaclass=abc.ABCMeta):
        @property
        @abc.abstractmethod
        def relative_url(self) -> str:
            pass

        @property
        @abc.abstractmethod
        def method(self) -> str:
            pass

        @abc.abstractmethod
        def get_header(self, name: str) -> Optional[str]:
            pass

        @property
        @abc.abstractmethod
        def json(self) -> dict:
            pass

    class ResponseAdapter(metaclass=abc.ABCMeta):
        @property
        @abc.abstractmethod
        def status_code(self) -> int:
            pass

        @abc.abstractmethod
        def get_header(self, name: str) -> Optional[str]:
            pass

        @property
        @abc.abstractmethod
        def elapsed_seconds(self) -> int:
            pass

        @property
        @abc.abstractmethod
        def content(self) -> bytes:
            pass

        @property
        @abc.abstractmethod
        def request(self) -> "UStorageClientBase.RequestAdapter":
            pass

        @abc.abstractmethod
        def raise_for_status(self) -> None:
            pass

        @abc.abstractmethod
        def json(self) -> dict:
            pass

    ERROR_MAP: list[tuple[int, Optional[re.Pattern], Type[exc.USReqException]]] = [
        (400, None, exc.USBadRequestException),
        (403, None, exc.USAccessDeniedException),
        (404, None, exc.USObjectNotFoundException),
        (409, re.compile("The entry already exists"), exc.USAlreadyExistsException),
        (409, None, exc.USIncorrectTenantIdException),
        (423, None, exc.USLockUnacquiredException),
        (530, None, exc.USPermissionCheckError),
    ]

    RequestData = NamedTuple(
        "RequestData",
        [
            ("method", str),
            ("relative_url", str),
            ("params", Optional[dict[str, str]]),
            ("json", Optional[dict]),
        ],
    )

    def __init__(
        self,
        host: str,
        auth_ctx: USAuthContextBase,
        prefix: Optional[str] = None,
        timeout: int = 30,
        context_request_id: Optional[str] = None,
        context_forwarded_for: Optional[str] = None,
    ):
        self.host = host
        self.prefix = auth_ctx.get_default_prefix() if prefix is None else prefix
        self.timeout = timeout

        self._disabled = False
        self._folder_id = auth_ctx.get_folder_id()
        self._auth_ctx = auth_ctx
        self._tenant_override: Optional[TenantDef] = None

        # Default headers for HTTP sessions
        self._default_headers = self._auth_ctx_to_default_headers(auth_ctx)
        if context_forwarded_for is not None:
            self._default_headers[DLHeadersCommon.FORWARDED_FOR.value] = context_forwarded_for
        # Initial cookies for HTTP session
        self._cookies = self._auth_ctx_to_cookies(auth_ctx)
        # Headers that might be changed during lifecycle e.g. Folder ID
        self._extra_headers = {}  # type: ignore  # TODO: fix

        if context_request_id is not None:
            self._default_headers["X-Request-Id"] = context_request_id

        self._ensure_extra_headers()

    def _ensure_extra_headers(self) -> None:
        effective_tenant = self.get_effective_tenant()
        expected_extra_headers = stringify_dl_headers(effective_tenant.get_outbound_tenancy_headers())
        if self._extra_headers != expected_extra_headers:
            self._extra_headers = expected_extra_headers

    def get_effective_tenant(self) -> TenantDef:
        if self._tenant_override is not None:
            return self._tenant_override
        return self._auth_ctx.get_tenant()

    def set_tenant_override(self, tenant: TenantDef) -> None:
        if self._auth_ctx.is_tenant_id_mutable():
            self._tenant_override = tenant
            self._ensure_extra_headers()
        else:
            raise AssertionError(
                f"US client folder ID is immutable with auth context {type(self._auth_ctx).__qualname__}"
            )

    @staticmethod
    def _auth_ctx_to_default_headers(ctx: USAuthContextBase) -> dict[str, str]:
        headers: dict[DLHeaders, str] = {
            **ctx.get_outbound_headers(include_tenancy=False),
            DLHeadersCommon.DL_COMPONENT: ctx.dl_component,
        }
        return stringify_dl_headers(headers)

    @staticmethod
    def _auth_ctx_to_cookies(ctx: USAuthContextBase) -> dict[str, str]:
        return stringify_dl_cookies(ctx.get_outbound_cookies())

    def _get_full_url(self, relative_url: str) -> str:
        return "/".join(map(lambda s: s.strip("/"), (self.host, self.prefix, relative_url)))

    @staticmethod
    def _log_request_start(request_data: RequestData):  # type: ignore  # TODO: fix
        LOGGER.info(
            "Asking UnitedStorage: method: %s, url: %s, params:(%s)",
            request_data.method,
            request_data.relative_url,
            request_data.params or {},
        )

    @classmethod
    def _get_us_json_from_response(cls, response: ResponseAdapter):  # type: ignore  # TODO: fix
        LOGGER.info(
            "Got %s from UnitedStorage (%s %s), content length: %s, request-id: %s, elapsed: %s",
            response.status_code,
            response.request.method,
            response.request.relative_url,
            len(response.content),
            response.get_header("X-Request-ID"),
            response.elapsed_seconds,
        )
        if response.status_code >= 400:
            LOGGER.info(
                "US error response on %s %s (%s): %s, folder_id: %s, org_id: %s, req_id: %s",
                response.request.method,
                response.request.relative_url,
                mask_sensitive_fields_by_name_in_json_recursive(response.request.json),
                response.content,
                # TODO: BI-4918 move to local injection and reuse bi_api_commons_ya_cloud.constants.DLHeadersYC
                response.request.get_header(DLHeadersCommon.FOLDER_ID.value),
                response.request.get_header(DLHeadersCommon.ORG_ID.value),
                response.get_header("X-Request-ID"),
            )
        try:
            response.raise_for_status()
        except USClientHTTPExceptionWrapper as http_err_ex_wrapper:
            http_err_ex = cast(Optional[Exception], http_err_ex_wrapper.__cause__)

            message: Optional[str] = response.json().get("message")
            for status_code, error_regex, exc_cls in cls.ERROR_MAP:
                if response.status_code == status_code:
                    if error_regex is None or message is not None and error_regex.match(message):
                        raise exc_cls(orig_exc=http_err_ex)

            raise exc.USReqException(orig_exc=http_err_ex)

        try:
            return response.json()
        except JSONDecodeError as ex:
            LOGGER.info("Got http status %s with invalid json %s from US", response.status_code, response.content)
            raise exc.USInvalidResponse from ex

    def _raise_for_disabled_interactions(self):  # type: ignore  # TODO: fix
        if self._disabled:
            raise exc.USInteractionDisabled("Interaction with US is forbidden")

    @contextmanager
    def interaction_disabled(self):  # type: ignore  # TODO: fix
        try:
            self._disabled = True
            yield
        except Exception:
            raise
        finally:
            self._disabled = False

    # Requests definitions
    @classmethod
    def _req_data_create_entry(  # type: ignore  # TODO: fix
        cls,
        key: EntryLocation,
        scope: str,
        meta=None,
        data=None,
        type_=None,
        hidden=None,
        links=None,
        mode="publish",
        unversioned_data=None,
        **kwargs,
    ) -> RequestData:
        meta = meta or {}
        data = data or {}
        unversioned_data = unversioned_data or {}
        type_ = type_ or ""
        links = links or {}

        return cls.RequestData(
            method="post",
            relative_url="/entries",
            params=None,
            json={
                "scope": scope,
                **key.to_us_req_api_params(),
                "meta": meta,
                "data": data,
                "unversionedData": unversioned_data,
                "type": type_,
                "recursion": True,
                "hidden": hidden,
                "links": links,
                "mode": mode,
                **kwargs,
            },
        )

    @classmethod
    def _req_data_get_entry(cls, entry_id, params=None, include_permissions=True, include_links=True):  # type: ignore  # TODO: fix
        params = params or {}
        if include_permissions:
            params["includePermissionsInfo"] = 1
        if include_links:
            params["includeLinks"] = 1

        return cls.RequestData(
            method="get",
            relative_url="/entries/{}".format(entry_id),
            params=params,
            json=None,
        )

    @classmethod
    def _req_data_move_entry(cls, entry_id: str, destination: str):  # type: ignore  # TODO: fix
        return cls.RequestData(
            method="post",
            relative_url="/entries/{}/move".format(entry_id),
            params=None,
            json={"destination": destination},
        )

    @classmethod
    def _req_data_update_entry(  # type: ignore  # TODO: fix
        cls,
        entry_id: str,
        data=None,
        unversioned_data=None,
        meta=None,
        mode="publish",
        lock=None,
        hidden=None,
        links=None,
    ):
        data = data or {}
        unversioned_data = unversioned_data or {}
        meta = meta or {}
        links = links or {}

        json_data = {
            "data": data,
            "unversionedData": unversioned_data,
            "meta": meta,
            "mode": mode,
            "links": links,
        }
        if hidden is not None:
            json_data.update(hidden=hidden)
        if lock is not None:
            json_data.update(lockToken=lock)
        return cls.RequestData(
            method="post",
            relative_url="/entries/{}".format(entry_id),
            params=None,
            json=json_data,
        )

    @classmethod
    def _req_data_delete_entry(cls, entry_id, lock=None):  # type: ignore  # TODO: fix
        params = None
        if lock is not None:
            params = {"lockToken": lock}
        return cls.RequestData(method="delete", relative_url="/entries/{}".format(entry_id), params=params, json=None)

    @classmethod
    def _req_data_iter_entries(  # type: ignore  # TODO: fix
        cls,
        scope: str,
        entry_type: Optional[str] = None,
        meta: Optional[dict] = None,
        all_tenants: bool = False,
        include_data: bool = False,
        ids: Optional[Iterable[str]] = None,
        page: int = 0,
        creation_time: Optional[dict[str, Union[str, int, None]]] = None,
    ) -> RequestData:
        req_params = dict(scope=scope, includeData=int(include_data))
        if entry_type:
            req_params.update(type=entry_type)
        meta = meta or {}
        if meta:
            req_params.update({"meta[{}]".format(k): v for k, v in meta.items()})
        if creation_time:
            req_params.update({"creationTime[{}]".format(k): v for k, v in creation_time.items()})
        if ids:
            req_params["ids"] = ids

        if all_tenants:
            assert include_data is False  # not supported for this endpoint
            endpoint = "/interTenant/entries"
        else:
            endpoint = "/entries"

        params: dict[Any, Any] = dict(page=page, **req_params)
        return cls.RequestData(
            method="get",
            relative_url=endpoint,
            params=params,
            json=None,
        )

    @classmethod
    def _req_data_acquire_lock(
        cls,
        entry_id: str,
        duration: Optional[int] = None,
        force: Optional[bool] = None,
    ) -> RequestData:
        """
        :param duration: in seconds. Default = 300 (5min)
        :return: lock token
        """
        params = {}  # type: ignore  # TODO: fix
        if duration is not None:
            params.update(duration=duration * 1000)  # US accepts duration in milliseconds
        if force is not None:
            params.update(force=force)

        return cls.RequestData(method="post", relative_url="/locks/{}".format(entry_id), params=None, json=params)

    @classmethod
    def _req_data_release_lock(cls, entry_id, lock):  # type: ignore  # TODO: fix
        return cls.RequestData(
            method="delete", relative_url="/locks/{}".format(entry_id), params={"lockToken": lock}, json=None
        )

    @classmethod
    def _req_data_entries_in_path(cls, *, us_path: str, page_size: int, page_idx: int) -> RequestData:
        return cls.RequestData(
            method="get",
            relative_url="/navigation",
            params=dict(
                page=str(page_idx),
                pageSize=str(page_size),
                path=us_path,
                includePermissionsInfo="false",
            ),
            json=None,
        )


class UStorageClient(UStorageClientBase):
    class RequestAdapter(UStorageClientBase.RequestAdapter):
        def __init__(self, request: requests.Request, request_data: UStorageClientBase.RequestData):
            self.req = request
            self._request_data = request_data

        @property
        def relative_url(self) -> str:
            return self._request_data.relative_url

        @property
        def method(self) -> str:
            return self.req.method.lower()

        def get_header(self, name: str) -> Optional[str]:
            return self.req.headers.get(name)

        @property
        def json(self) -> dict:
            return self._request_data.json  # type: ignore  # TODO: fix

    class ResponseAdapter(UStorageClientBase.ResponseAdapter):
        def __init__(self, response: requests.Response, request_data: "UStorageClient.RequestData"):
            self.resp = response
            self._req_adapter = UStorageClient.RequestAdapter(response.request, request_data)  # type: ignore  # TODO: fix

        @property
        def status_code(self) -> int:
            return self.resp.status_code

        def get_header(self, name: str) -> Optional[str]:
            return self.resp.headers.get(name)

        @property
        def elapsed_seconds(self) -> int:
            return self.resp.elapsed.total_seconds()  # type: ignore  # TODO: fix

        @property
        def content(self) -> bytes:
            return self.resp.content

        @property
        def request(self) -> "UStorageClientBase.RequestAdapter":
            return self._req_adapter

        def raise_for_status(self) -> None:
            try:
                self.resp.raise_for_status()
            except HTTPError as http_error:
                raise USClientHTTPExceptionWrapper(str(http_error)) from http_error

        def json(self) -> dict:
            return self.resp.json()

    def __init__(
        self,
        host: str,
        auth_ctx: USAuthContextBase,
        prefix: Optional[str] = None,
        timeout: int = 30,
        context_request_id: Optional[str] = None,
        context_forwarded_for: Optional[str] = None,
    ):
        super().__init__(
            host=host,
            auth_ctx=auth_ctx,
            prefix=prefix,
            timeout=timeout,
            context_request_id=context_request_id,
            context_forwarded_for=context_forwarded_for,
        )

        self._session = get_retriable_requests_session()
        self._session.headers.update(self._default_headers)
        self._session.cookies.update(self._cookies)

    def close(self) -> None:
        self._session.close()

    def _request(self, request_data: UStorageClientBase.RequestData) -> dict[str, Any]:
        self._raise_for_disabled_interactions()

        url = self._get_full_url(request_data.relative_url)

        self._log_request_start(request_data)

        request_kwargs = {"json": request_data.json} if request_data.json is not None else {}
        tracing_headers = get_current_tracing_headers()

        response = self._session.request(
            request_data.method,
            url,
            stream=False,
            timeout=self.timeout,
            params=request_data.params,
            headers={
                **self._extra_headers,
                **tracing_headers,
            },
            **request_kwargs,  # type: ignore  # TODO: fix
        )

        response_adapter = self.ResponseAdapter(response, request_data)
        return self._get_us_json_from_response(response_adapter)

    def create_entry(
        self,
        key: EntryLocation,
        scope: str,
        meta: Optional[dict[str, str]] = None,
        data: Optional[dict[str, Any]] = None,
        unversioned_data: Optional[dict[str, Any]] = None,
        type_: Optional[str] = None,
        hidden: Optional[bool] = None,
        links: Optional[dict[str, Any]] = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        rq_data = self._req_data_create_entry(
            key=key,
            scope=scope,
            meta=meta,
            data=data,
            unversioned_data=unversioned_data,
            type_=type_,
            hidden=hidden,
            links=links,
            **kwargs,
        )
        return self._request(rq_data)

    def get_entry(
        self,
        entry_id: str,
        params: Optional[dict[str, Any]] = None,
        include_permissions: bool = True,
        include_links: bool = True,
    ) -> dict[str, Any]:
        return self._request(
            self._req_data_get_entry(
                entry_id, params=params, include_permissions=include_permissions, include_links=include_links
            )
        )

    def move_entry(self, entry_id, destination):  # type: ignore  # TODO: fix
        return self._request(self._req_data_move_entry(entry_id, destination=destination))

    def update_entry(  # type: ignore  # TODO: fix
        self,
        entry_id: str,
        data: Optional[dict[str, Any]] = None,
        unversioned_data: Optional[dict[str, Any]] = None,
        meta: Optional[dict[str, str]] = None,
        lock: Optional[str] = None,
        hidden: Optional[bool] = None,
        links: Optional[dict[str, Any]] = None,
    ):
        return self._request(
            self._req_data_update_entry(
                entry_id,
                data=data,
                unversioned_data=unversioned_data,
                meta=meta,
                lock=lock,
                hidden=hidden,
                links=links,
            )
        )

    def entries_iterator(
        self,
        scope: str,
        entry_type: Optional[str] = None,
        meta: Optional[dict] = None,
        all_tenants: bool = False,
        include_data: bool = False,
        ids: Optional[Iterable[str]] = None,
        creation_time: Optional[dict[str, Union[str, int, None]]] = None,
    ) -> Generator[dict, None, None]:
        """
        :param scope:
        :param entry_type:
        :param meta: Filter entries by "meta" section values.
        :param all_tenants: Look up across all tenants. False by default.
        :param include_data: Return full US entry data. False by default.
        :param ids: Filter entries by uuid.
        :param creation_time: Filter entries by creation_time. Available filters: eq, ne, gt, gte, lt, lte
        :return:
        """
        page = 0
        while True:
            resp = self._request(
                self._req_data_iter_entries(
                    scope,
                    entry_type=entry_type,
                    meta=meta,
                    all_tenants=all_tenants,
                    include_data=include_data,
                    ids=ids,
                    page=page,
                    creation_time=creation_time,
                )
            )
            if resp.get("entries"):
                page_entries = resp["entries"]
            else:
                break

            for entr in page_entries:
                yield entr

            if resp.get("nextPageToken"):
                page = resp["nextPageToken"]
            else:
                break

    def delete_entry(self, entry_id, lock=None):  # type: ignore  # TODO: fix
        self._request(self._req_data_delete_entry(entry_id, lock=lock))

    def acquire_lock(self, entry_id, duration=None, wait_timeout=None, force=None):  # type: ignore  # TODO: fix
        """
        :param entry_id:
        :param force:
        :param duration: in seconds. Default = 300 (5min)
        :param wait_timeout: in seconds
        :return: lock token
        """
        req_data = self._req_data_acquire_lock(entry_id, duration=duration, force=force)
        start_ts = time.time()
        while True:
            try:
                resp = self._request(req_data)
                lock = resp["lockToken"]
                LOGGER.info('Acquired lock "%s" for object "%s"', lock, entry_id)
                return lock
            except exc.USLockUnacquiredException:
                if wait_timeout and time.time() - start_ts < wait_timeout:
                    time.sleep(0.25)
                else:
                    raise

    def release_lock(self, entry_id, lock):  # type: ignore  # TODO: fix
        try:
            self._request(self._req_data_release_lock(entry_id, lock=lock))
        except exc.USReqException:
            LOGGER.exception('Unable to release lock "%s"', lock)

    def get_entries_info_in_path(self, us_path: str) -> list[dict[str, Any]]:
        resp = self._request(
            self._req_data_entries_in_path(
                us_path=us_path,
                page_size=100,
                page_idx=0,
            )
        )
        return resp["entries"]

lib/dl_core/dl_core/united_storage_client.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import json
import logging
import ssl
import time
from typing import (
    AsyncGenerator,
    Iterable,
    NoReturn,
    Optional,
    Union,
)

import aiohttp

from dl_api_commons.aiohttp.aiohttp_client import (
    BIAioHTTPClient,
    PredefinedIntervalsRetrier,
)
from dl_api_commons.tracing import get_current_tracing_headers
from dl_configs.utils import get_root_certificates_path
from dl_core.base_models import EntryLocation
from dl_core.exc import (
    USLockUnacquiredException,
    USReqException,
)
from dl_core.united_storage_client import (
    USAuthContextBase,
    USClientHTTPExceptionWrapper,
    UStorageClientBase,
)


LOGGER = logging.getLogger(__name__)


class UStorageClientAIO(UStorageClientBase):
    class ResponseAdapter(UStorageClientBase.ResponseAdapter):
        def __init__(
            self,
            response: aiohttp.ClientResponse,
            content: bytes,
            request_data: UStorageClientBase.RequestData,
            elapsed_seconds: float,
        ):
            self._response = response
            self._content = content
            self._request_data = request_data
            self._elapsed_seconds = elapsed_seconds

            self._parsed_json_data = None
            self._request_adapter = UStorageClientAIO.RequestAdapter(response.request_info, request_data)

        @property
        def status_code(self) -> int:
            return self._response.status

        def get_header(self, name: str) -> Optional[str]:
            return self._response.headers.get(name)

        @property
        def elapsed_seconds(self) -> float:  # type: ignore  # TODO: fix
            return self._elapsed_seconds

        @property
        def content(self) -> bytes:
            return self._content

        @property
        def request(self) -> "UStorageClientBase.RequestAdapter":
            return self._request_adapter

        def raise_for_status(self) -> None:
            try:
                self._response.raise_for_status()
            except aiohttp.ClientResponseError as err:
                if err.status is not None:
                    raise USClientHTTPExceptionWrapper(str(err)) from err
                raise

        def json(self) -> dict:
            if self._parsed_json_data is None:
                self._parsed_json_data = json.loads(self._content.decode("utf-8"))
            return self._parsed_json_data  # type: ignore  # TODO: fix

    class RequestAdapter(UStorageClientBase.RequestAdapter):
        def __init__(self, request: aiohttp.RequestInfo, request_data: UStorageClientBase.RequestData):
            self._request = request
            self._request_data = request_data

        @property
        def relative_url(self) -> str:
            return self._request_data.relative_url

        @property
        def method(self) -> str:
            return self._request_data.method

        def get_header(self, name: str) -> Optional[str]:
            return self._request.headers.get(name)

        @property
        def json(self) -> dict:
            return self._request_data.json  # type: ignore  # TODO: fix

    _session: aiohttp.ClientSession
    DEFAULT_SSL_CONTEXT = ssl.create_default_context(cafile=get_root_certificates_path())

    def __init__(
        self,
        host: str,
        prefix: Optional[str],
        auth_ctx: USAuthContextBase,
        context_request_id: Optional[str] = None,
        context_forwarded_for: Optional[str] = None,
        timeout: int = 30,
    ):
        super().__init__(
            host=host,
            auth_ctx=auth_ctx,
            prefix=prefix,
            timeout=timeout,
            context_request_id=context_request_id,
            context_forwarded_for=context_forwarded_for,
        )

        self._retry_intervals = (0.5, 1.0, 1.1, 2.0, 2.2)
        self._retry_codes = {408, 429, 500, 502, 503, 504}

        self._session = aiohttp.ClientSession(
            cookies=self._cookies,
            connector=aiohttp.TCPConnector(ssl_context=self.DEFAULT_SSL_CONTEXT),
            headers={
                **self._default_headers,
            },
        )

    async def _request(self, request_data: UStorageClientBase.RequestData):  # type: ignore  # TODO: fix
        self._raise_for_disabled_interactions()
        self._log_request_start(request_data)
        tracing_headers = get_current_tracing_headers()
        start = time.monotonic()

        async with BIAioHTTPClient(
            base_url="/".join([self.host, self.prefix]),
            session=self._session,
            retrier=PredefinedIntervalsRetrier(
                retry_intervals=self._retry_intervals,
                retry_codes=self._retry_codes,
                retry_methods={"GET", "POST", "PUT", "DELETE"},  # TODO: really retry all of them?..
            ),
            raise_for_status=False,
            close_session_on_exit=False,
        ) as bi_http_client:
            async with bi_http_client.request(
                method=request_data.method,
                path=request_data.relative_url,
                params=request_data.params,
                json_data=request_data.json,
                read_timeout_sec=self.timeout,  # TODO: total timeout
                conn_timeout_sec=1,
                headers={
                    **self._extra_headers,
                    **tracing_headers,
                },
            ) as response:
                content = await response.read()

        end = time.monotonic()

        response_adapter = self.ResponseAdapter(
            response=response,
            content=content,
            request_data=request_data,
            elapsed_seconds=(end - start),
        )
        return self._get_us_json_from_response(response_adapter)

    async def get_entry(self, entry_id: str) -> dict:
        return await self._request(self._req_data_get_entry(entry_id=entry_id))

    async def create_entry(  # type: ignore  # TODO: fix
        self,
        key: EntryLocation,
        scope: str,
        meta=None,
        data=None,
        unversioned_data=None,
        type_=None,
        hidden=None,
        links=None,
        **kwargs,
    ):
        rq_data = self._req_data_create_entry(
            key=key,
            scope=scope,
            meta=meta,
            data=data,
            unversioned_data=unversioned_data,
            type_=type_,
            hidden=hidden,
            links=links,
            **kwargs,
        )
        return await self._request(rq_data)

    async def update_entry(  # type: ignore  # TODO: fix
        self, entry_id, data=None, unversioned_data=None, meta=None, lock=None, hidden=None, links=None
    ):
        return await self._request(
            self._req_data_update_entry(
                entry_id,
                data=data,
                unversioned_data=unversioned_data,
                meta=meta,
                lock=lock,
                hidden=hidden,
                links=links,
            )
        )

    async def delete_entry(self, entry_id, lock=None) -> NoReturn:  # type: ignore  # TODO: fix
        await self._request(self._req_data_delete_entry(entry_id, lock=lock))

    async def entries_iterator(
        self,
        scope: str,
        entry_type: Optional[str] = None,
        meta: Optional[dict] = None,
        all_tenants: bool = False,
        include_data: bool = False,
        ids: Optional[Iterable[str]] = None,
        creation_time: Optional[dict[str, Union[str, int, None]]] = None,
    ) -> AsyncGenerator[dict, None]:
        page = 0
        while True:
            resp = await self._request(
                self._req_data_iter_entries(
                    scope,
                    entry_type=entry_type,
                    meta=meta,
                    all_tenants=all_tenants,
                    page=page,
                    include_data=include_data,
                    ids=ids,
                    creation_time=creation_time,
                )
            )
            if resp.get("entries"):
                page_entries = resp["entries"]
            else:
                break

            for entr in page_entries:
                yield entr

            if resp.get("nextPageToken"):
                page = resp["nextPageToken"]
            else:
                break

    async def list_all_entries(
        self, scope: str, entry_type: Optional[str] = None, meta: Optional[dict] = None, all_tenants: bool = False
    ) -> list:
        ret = []
        async for e in self.entries_iterator(scope, entry_type, meta, all_tenants, include_data=False):
            ret.append(e)

        return ret

    async def acquire_lock(
        self,
        entry_id: str,
        duration: Optional[int] = None,
        wait_timeout: Optional[int] = None,
        force: Optional[bool] = None,
    ) -> str:
        """
        :param entry_id: US entry ID to lock
        :param duration: Duration of acquired lock
        :param wait_timeout: Acquire attempts timeout in seconds
        :param force: Force acquire lock without wait
        :return: Lock token
        """
        req_data = self._req_data_acquire_lock(entry_id, duration=duration, force=force)
        start_ts = time.time()
        while True:
            try:
                resp = await self._request(req_data)
                lock = resp["lockToken"]
                LOGGER.info('Acquired lock "%s" for object "%s"', lock, entry_id)
                return lock
            except USLockUnacquiredException:
                if wait_timeout and time.time() - start_ts < wait_timeout:
                    await asyncio.sleep(0.25)
                else:
                    raise

    async def release_lock(self, entry_id, lock):  # type: ignore  # TODO: fix
        try:
            await self._request(self._req_data_release_lock(entry_id, lock=lock))
        except USReqException:
            LOGGER.exception('Unable to release lock "%s"', lock)

    async def close(self):  # type: ignore  # TODO: fix
        await self._session.close()

lib/dl_core/dl_core/united_storage_client_aio.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Type

from dl_constants.enums import ConnectionType
from dl_core.connectors.base.lifecycle import (
    ConnectionLifecycleManager,
    DefaultConnectionLifecycleManager,
)
from dl_core.us_connection_base import (
    ConnectionBase,
    UnknownConnection,
)


CONNECTION_TYPES: dict[ConnectionType, Type[ConnectionBase]] = {
    ConnectionType.unknown: UnknownConnection,
}
CONNECTION_LIFECYCLE_MGR_CLASSES: dict[ConnectionType, Type[ConnectionLifecycleManager]] = {
    ConnectionType.unknown: DefaultConnectionLifecycleManager,
}


def get_connection_class(conn_type: ConnectionType) -> Type[ConnectionBase]:
    """Return class for given connection type"""
    return CONNECTION_TYPES[conn_type]


def get_lifecycle_manager_cls(conn_type: ConnectionType) -> Type[ConnectionLifecycleManager]:
    """Return class for given connection type"""
    return CONNECTION_LIFECYCLE_MGR_CLASSES[conn_type]


def register_connection_class(
    new_conn_cls: Type[ConnectionBase],
    conn_type: ConnectionType,
    lifecycle_manager_cls: Type[ConnectionLifecycleManager] = DefaultConnectionLifecycleManager,
    allow_ct_override: bool = False,
) -> None:
    if conn_type is None:
        raise ValueError(f"Connection type for {new_conn_cls} is None")

    if not allow_ct_override:
        if conn_type in CONNECTION_TYPES:
            raise ValueError(f"Attempt to duplicated connection type registration: {conn_type}")

        if new_conn_cls in CONNECTION_TYPES.values():
            raise ValueError(f"Attempt to duplicated connection class registration: {new_conn_cls}")

    CONNECTION_TYPES[conn_type] = new_conn_cls
    CONNECTION_LIFECYCLE_MGR_CLASSES[conn_type] = lifecycle_manager_cls

lib/dl_core/dl_core/us_connection.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Generic,
    NamedTuple,
    Optional,
    Type,
    TypeVar,
    Union,
)

import attr
import sqlalchemy as sa
from sqlalchemy.engine.default import DefaultDialect

from dl_api_commons.reporting.models import NotificationReportingRecord
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_constants.enums import (
    ConnectionState,
    ConnectionType,
    DataSourceRole,
    DataSourceType,
    RawSQLLevel,
)
from dl_core import connection_models
from dl_core.base_models import (
    ConnCacheableDataModelMixin,
    ConnectionDataModelBase,
    ConnectionRef,
    ConnSubselectDataModelMixin,
    DefaultConnectionRef,
    EntryLocation,
    WorkbookEntryLocation,
)
from dl_core.connection_executors.adapters.common_base import get_dialect_for_conn_type
from dl_core.connection_models import (
    DBIdent,
    SchemaIdent,
)
from dl_core.data_processing.cache.primitives import (
    DataKeyPart,
    LocalKeyRepresentation,
)
from dl_core.db import (
    TypeTransformer,
    get_type_transformer,
)
from dl_core.i18n.localizer import Translatable
from dl_core.us_entry import (
    BaseAttrsDataModel,
    USEntry,
)
from dl_core.utils import (
    parse_comma_separated_hosts,
    secrepr,
)
from dl_i18n.localizer_base import Localizer
from dl_utils.aio import await_sync
from dl_utils.utils import DataKey


if TYPE_CHECKING:
    from dl_core.connection_executors import SyncConnExecutorBase
    from dl_core.connection_models.common_models import TableIdent
    from dl_core.services_registry import ServicesRegistry
    from dl_core.us_manager.us_manager import USManagerBase
    from dl_core.us_manager.us_manager_sync import SyncUSManager


LOGGER = logging.getLogger(__name__)


class DataSourceTemplate(NamedTuple):
    # for visualization in UI
    title: str
    group: list[str]
    # main properties
    source_type: DataSourceType
    connection_id: str
    # type-specific
    parameters: dict
    tab_title: Optional[str] = None
    form: Optional[list[dict[str, Any]]] = None
    disabled: bool = False

    def get_param_hash(self) -> str:
        from dl_core import data_source  # noqa

        return data_source.get_parameters_hash(
            source_type=self.source_type,
            connection_id=self.connection_id,
            **self.parameters,
        )


_CB_TV = TypeVar("_CB_TV", bound="ConnectionBase")


class ConnectionBase(USEntry, metaclass=abc.ABCMeta):
    dir_name: ClassVar[str] = ""  # type: ignore  # TODO: fix
    scope: ClassVar[str] = "connection"  # type: ignore  # TODO: fix

    conn_type: ConnectionType
    source_type: ClassVar[Optional[DataSourceType]] = None
    allowed_source_types: ClassVar[Optional[frozenset[DataSourceType]]] = None
    allow_dashsql: ClassVar[bool] = False
    allow_cache: ClassVar[bool] = False
    is_always_internal_source: ClassVar[bool] = False
    is_always_user_source: ClassVar[bool] = False

    _preview_conn = None

    def __init__(
        self,
        uuid: Optional[str] = None,
        data: Optional[dict] = None,
        entry_key: Optional[EntryLocation] = None,
        type_: Optional[str] = None,
        meta: Optional[dict] = None,
        is_locked: Optional[bool] = None,
        is_favorite: Optional[bool] = None,
        permissions_mode: Optional[str] = None,
        initial_permissions: Optional[str] = None,
        permissions: Optional[dict[str, bool]] = None,
        links: Optional[dict] = None,
        hidden: bool = False,
        data_strict: bool = True,
        *,
        us_manager: USManagerBase,
    ):
        super().__init__(
            uuid=uuid,
            data=data,
            entry_key=entry_key,
            type_=type_,
            meta=meta,
            is_locked=is_locked,
            is_favorite=is_favorite,
            permissions_mode=permissions_mode,
            initial_permissions=initial_permissions,
            permissions=permissions,
            links=links,
            hidden=hidden,
            data_strict=data_strict,
            us_manager=us_manager,
        )

        try:
            assert self.type_ is not None
            self.conn_type = ConnectionType(self.type_)
        except ValueError:
            self.conn_type = ConnectionType.unknown

    @property
    def allow_public_usage(self) -> bool:
        return False

    @property
    def pass_db_query_to_user(self) -> bool:
        return self.is_always_user_source or not self.is_always_internal_source

    @attr.s(kw_only=True)
    class DataModel(ConnectionDataModelBase):
        table_name: Optional[str] = attr.ib(default=None)
        sample_table_name: Optional[str] = attr.ib(default=None)
        name: Optional[str] = attr.ib(default=None)
        data_export_forbidden: bool = attr.ib(default=False)

    @property
    def data_export_forbidden(self) -> bool:
        return self.data.data_export_forbidden if hasattr(self.data, "data_export_forbidden") else False

    @classmethod
    def get_provided_source_types(cls) -> frozenset[DataSourceType]:
        if cls.allowed_source_types is not None:
            return cls.allowed_source_types
        if cls.source_type is not None:
            return frozenset((cls.source_type,))
        return frozenset()

    @property
    def conn_ref(self) -> Optional[ConnectionRef]:
        if self.uuid is not None:
            return DefaultConnectionRef(conn_id=self.uuid)

        return None

    @property
    def type_transformer(self) -> TypeTransformer:
        return get_type_transformer(self.conn_type)

    @property
    def state(self):  # type: ignore  # TODO: fix
        return getattr(ConnectionState, self.meta.get("state", ""), None)

    @state.setter
    def state(self, value):  # type: ignore  # TODO: fix
        if isinstance(value, ConnectionState):
            self.meta["state"] = value.name
        else:
            raise Exception("Unknown state value {}".format(value))

    @property
    def sample_state(self):  # type: ignore  # TODO: fix
        return getattr(ConnectionState, self.meta.get("sample_state", ""), None)

    @sample_state.setter
    def sample_state(self, value):  # type: ignore  # TODO: fix
        if isinstance(value, ConnectionState):
            self.meta["sample_state"] = value.name
        else:
            raise Exception("Unknown sample_state value {}".format(value))

    @property
    def is_allow_dataset_creation(self):  # type: ignore  # TODO: fix
        if self.type_ not in ("csv", "xls"):
            return True
        if self.state and self.state.value >= ConnectionState.saved.value:
            return True
        return False

    @property
    def is_dashsql_allowed(self) -> bool:
        """Placeholder, should normally be overridden by `SubselectMixin`"""
        return False

    def as_dict(self, short=False):  # type: ignore  # TODO: fix
        resp = super().as_dict(short=short)
        if short:
            resp.update(
                {
                    "state": self.state.name if self.state is not None else None,
                    "is_allow_dataset_creation": self.is_allow_dataset_creation,
                    "type": self.conn_type.value,
                    "content_length": self.meta.get("content_length", None),
                }
            )
        else:
            resp.update(
                {
                    "db_type": self.conn_type.value,
                    "meta": self.meta,
                    "created_at": self.created_at,
                    "updated_at": self.updated_at,
                }
            )

        if not isinstance(self.entry_key, WorkbookEntryLocation):
            resp.pop("name", None)
        resp.pop("is_favorite", None)
        resp.pop("table_name", None)
        resp.pop("sample_table_name", None)
        return resp

    @classmethod
    def create_from_dict(  # type: ignore  # TODO: fix
        cls: Type[_CB_TV],
        data_dict: Union[dict, BaseAttrsDataModel],
        ds_key: Union[EntryLocation, str, None] = None,
        type_: Optional[str] = None,
        meta: Optional[dict] = None,
        **kwargs: Any,
    ) -> _CB_TV:
        if meta is None:
            meta = {}
        return super().create_from_dict(
            data_dict,
            ds_key,
            type_,
            meta,
            **kwargs,
        )

    def test(self, conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase]) -> None:
        raise NotImplementedError

    @property
    def table_name(self):  # type: ignore  # TODO: fix
        return self.data.table_name

    @property
    def sample_table_name(self):  # type: ignore  # TODO: fix
        return self.data.sample_table_name

    _dsrc_error = "Only materializable connections can have a data source"  # TODO remove?

    def update_data_source(
        self,
        id: str,
        role: Optional[DataSourceRole] = None,
        raw_schema: Optional[list] = None,
        remove_raw_schema: bool = False,
        **parameters: dict[str, Any],
    ) -> None:
        raise NotImplementedError(self._dsrc_error)

    def get_single_data_source_id(self) -> str:
        raise NotImplementedError(self._dsrc_error)

    def get_data_source_template_templates(self, localizer: Localizer) -> list[DataSourceTemplate]:
        """
        For user-input parametrized sources such as subselects.

        Must *not* query the user database for this.
        """
        # A safe default:
        return []

    def get_data_source_local_templates(self) -> Optional[list[DataSourceTemplate]]:
        """
        A list of available data sources that are stored in the metadata or
        code, i.e. don't require querying the actual database.

        Generally, nonempty local_templates should mean that full templates
        don't need to be requested.
        """
        return None

    def get_data_source_templates(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[DataSourceTemplate]:
        raise NotImplementedError

    async def validate_new_data(
        self,
        services_registry: ServicesRegistry,
        changes: Optional[dict] = None,
        original_version: Optional[ConnectionBase] = None,
    ) -> None:
        """Override point for subclasses"""

    def validate_new_data_sync(
        self,
        services_registry: ServicesRegistry,
        changes: Optional[dict] = None,
        original_version: Optional[ConnectionBase] = None,
    ) -> None:
        """Convenience wrapper"""
        return await_sync(
            self.validate_new_data(
                services_registry=services_registry,
                changes=changes,
                original_version=original_version,
            )
        )

    def check_for_notifications(self) -> list[Optional[NotificationReportingRecord]]:
        return []

    def get_cache_key_part(self) -> LocalKeyRepresentation:
        local_key_rep = LocalKeyRepresentation()
        local_key_rep = local_key_rep.multi_extend(
            DataKeyPart(
                part_type="connection_id",
                part_content=self.uuid,
            ),
            DataKeyPart(
                part_type="connection_revision_id",
                part_content=self.revision_id,
            ),
        )
        return local_key_rep


class ExecutorBasedMixin(ConnectionBase, metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_conn_dto(self) -> connection_models.ConnDTO:
        pass

    def get_conn_options(self) -> connection_models.ConnectOptions:
        return connection_models.ConnectOptions(
            pass_db_messages_to_user=self.is_always_user_source and not self.is_always_internal_source,
            pass_db_query_to_user=self.is_always_user_source and not self.is_always_internal_source,
        )

    def get_dialect(self) -> DefaultDialect:
        return get_dialect_for_conn_type(self.conn_type)

    def quote(self, value) -> sa.sql.quoted_name:  # type: ignore  # TODO: fix  # subclass of str
        return self.get_dialect().identifier_preparer.quote(value)

    @property
    @abc.abstractmethod
    def cache_ttl_sec_override(self) -> Optional[int]:
        pass

    @property
    def use_locked_cache(self) -> bool:
        """Transition configurable override; should eventually become all-`True`"""
        import os

        disabled_types = (os.environ.get("DL_NONLOCKED_CACHE_CONN_TYPES") or "").split(",")
        result = self.conn_type.name not in disabled_types
        LOGGER.debug(f"use_locked_cache = {self.conn_type.name!r} not in {disabled_types!r} = {result!r}")
        return result

    def test(self, conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase]) -> None:
        conn_executor = conn_executor_factory(self)
        conn_executor.test()

    def get_schema_names(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
        db_name: Optional[str] = None,
    ) -> list[str]:
        conn_executor = conn_executor_factory(self)
        return conn_executor.get_schema_names(DBIdent(db_name))

    def get_tables(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
        db_name: Optional[str] = None,
        schema_name: Optional[str] = None,
    ) -> list[TableIdent]:
        conn_executor = conn_executor_factory(self)
        return conn_executor.get_tables(SchemaIdent(db_name=db_name, schema_name=schema_name))


class SubselectMixin:
    # Not included: DataModel.raw_sql_level

    @property
    def is_subselect_allowed(self) -> bool:
        return self.data.raw_sql_level in (RawSQLLevel.subselect, RawSQLLevel.dashsql)  # type: ignore  # TODO: fix

    @property
    def is_dashsql_allowed(self) -> bool:
        if not self.allow_dashsql:  # type: ignore  # TODO: fix
            return False
        if self.data.raw_sql_level == RawSQLLevel.dashsql:  # type: ignore  # TODO: fix
            return True
        return False

    def _make_subselect_templates(
        self,
        source_type: DataSourceType,
        localizer: Localizer,
        title: str = "SQL",
        field_doc_key: str = "ANY_SUBSELECT/subsql",
    ) -> list[DataSourceTemplate]:
        allowed = self.is_subselect_allowed
        return [
            DataSourceTemplate(
                title=title,
                tab_title=localizer.translate(Translatable("source_templates-tab_title-sql")),
                source_type=source_type,
                form=[
                    {
                        "name": "subsql",
                        "input_type": "sql",
                        "default": "select 1 as a",
                        "required": True,  # Note: `required` means `disallow an empty value` (e.g. an empty string).
                        "title": localizer.translate(Translatable("source_templates-label-subquery")),
                        "field_doc_key": field_doc_key,
                    },
                ],
                disabled=not allowed,
                group=[],
                connection_id=self.uuid,  # type: ignore  # TODO: fix
                parameters={},
            ),
        ]


class ConnectionSQL(SubselectMixin, ExecutorBasedMixin, ConnectionBase):  # type: ignore  # TODO: fix
    has_schema: ClassVar[bool] = False
    default_schema_name: ClassVar[Optional[str]] = None

    @attr.s(kw_only=True)
    class DataModel(ConnCacheableDataModelMixin, ConnSubselectDataModelMixin, ConnectionBase.DataModel):
        host: Optional[str] = attr.ib(default=None)
        port: Optional[int] = attr.ib(default=None)
        db_name: Optional[str] = attr.ib(default=None)
        username: Optional[str] = attr.ib(default=None)
        password: Optional[str] = attr.ib(repr=secrepr, default=None)

        @classmethod
        def get_secret_keys(cls) -> set[DataKey]:
            return {DataKey(parts=("password",))}

    @property
    def cache_ttl_sec_override(self) -> Optional[int]:
        return self.data.cache_ttl_sec

    @property
    def db_name(self) -> Optional[str]:
        return self.data.db_name

    def get_parameter_combinations(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[dict]:
        if not self.db_name:
            return []

        # has db_name, so just list tables
        if not self.has_schema:
            return [
                dict(table_name=tid.table_name) for tid in self.get_tables(conn_executor_factory=conn_executor_factory)
            ]

        assert self.has_schema
        schemas = self.get_schema_names(conn_executor_factory=conn_executor_factory)
        return [
            dict(schema_name=schema_name, table_name=tid.table_name)
            for schema_name in schemas
            for tid in self.get_tables(conn_executor_factory=conn_executor_factory, schema_name=schema_name)
        ]

    def get_data_source_template_title(self, parameters: dict) -> str:
        return parameters["table_name"]

    def get_data_source_template_group(self, parameters: dict) -> list[str]:
        return [val for val in (parameters.get("db_name"), parameters.get("schema_name")) if val is not None]

    def get_data_source_templates(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[DataSourceTemplate]:
        """For listed sources such as db tables"""
        return [
            DataSourceTemplate(
                title=self.get_data_source_template_title(parameters),
                group=self.get_data_source_template_group(parameters),
                source_type=self.source_type,  # type: ignore  # TODO: fix
                connection_id=self.uuid,  # type: ignore  # TODO: fix
                parameters=parameters,
            )
            for parameters in self.get_parameter_combinations(conn_executor_factory=conn_executor_factory)
        ]

    def validate(self) -> None:
        pass


class ClassicConnectionSQL(ConnectionSQL):
    @property
    def password(self) -> Optional[str]:
        return self.data.password

    def parse_multihosts(self) -> tuple[str, ...]:
        return parse_comma_separated_hosts(self.data.host)


CONNECTOR_SETTINGS_TV = TypeVar("CONNECTOR_SETTINGS_TV", bound=ConnectorSettingsBase)


class ConnectionHardcodedDataMixin(Generic[CONNECTOR_SETTINGS_TV], metaclass=abc.ABCMeta):
    """Connector type specific data is loaded from dl_configs.connectors_settings"""

    conn_type: ConnectionType
    settings_type: Type[CONNECTOR_SETTINGS_TV]
    us_manager: SyncUSManager

    # TODO: remove
    @classmethod
    def _get_connector_settings(cls, usm: SyncUSManager) -> CONNECTOR_SETTINGS_TV:
        sr = usm.get_services_registry()
        connectors_settings = sr.get_connectors_settings(cls.conn_type)
        assert connectors_settings is not None
        assert isinstance(connectors_settings, cls.settings_type)
        return connectors_settings

    @property
    def _connector_settings(self) -> CONNECTOR_SETTINGS_TV:
        sr = self.us_manager.get_services_registry()
        connectors_settings = sr.get_connectors_settings(self.conn_type)
        assert connectors_settings is not None
        assert isinstance(connectors_settings, self.settings_type)
        return connectors_settings


class HiddenDatabaseNameMixin(ConnectionSQL):
    """Mixin that hides db name in UI"""

    def get_parameter_combinations(
        self,
        conn_executor_factory: Callable[[ConnectionBase], SyncConnExecutorBase],
    ) -> list[dict]:
        return [
            dict(combination, db_name=None)  # don't expose db name
            for combination in super().get_parameter_combinations(conn_executor_factory=conn_executor_factory)
            if combination["db_name"] == self.db_name
        ]


class UnknownConnection(ConnectionBase):
    """
    Special connection class for entries with unknown type and connection scope
    """

    conn_type = ConnectionType.unknown

    @attr.s(kw_only=True)
    class DataModel(ConnectionBase.DataModel):
        pass

lib/dl_core/dl_core/us_connection_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    Collection,
    Generator,
    Optional,
)

import attr

from dl_constants.enums import (
    AggregationFunction,
    DataSourceCreatedVia,
    DataSourceRole,
    DataSourceType,
    FieldType,
    ManagedBy,
    UserDataType,
)
from dl_core import multisource
from dl_core.base_models import (
    ObligatoryFilter,
    connection_ref_from_id,
)
from dl_core.component_errors import ComponentErrorRegistry
from dl_core.components.dependencies.primitives import FieldInterDependencyInfo
from dl_core.components.ids import (
    DEFAULT_FIELD_ID_GENERATOR_TYPE,
    FIELD_ID_GENERATOR_MAP,
    resolve_id_collisions,
)
from dl_core.data_source_spec.base import DataSourceSpec
from dl_core.data_source_spec.collection import (
    DataSourceCollectionSpec,
    DataSourceCollectionSpecBase,
)
from dl_core.data_source_spec.sql import StandardSQLDataSourceSpec
from dl_core.db import SchemaColumn
from dl_core.fields import (
    DirectCalculationSpec,
    ResultSchema,
)
from dl_core.rls import RLS
from dl_core.us_entry import (
    BaseAttrsDataModel,
    USEntry,
)


if TYPE_CHECKING:
    from dl_core.components.ids import FieldIdGenerator


LOGGER = logging.getLogger(__name__)


class Dataset(USEntry):
    dir_name = ""  # type: ignore  # TODO: fix
    scope = "dataset"  # type: ignore  # TODO: fix

    @attr.s
    class DataModel(BaseAttrsDataModel):
        name: str = attr.ib()
        revision_id: Optional[str] = attr.ib(default=None)
        result_schema: ResultSchema = attr.ib(factory=ResultSchema)
        source_collections: list[DataSourceCollectionSpecBase] = attr.ib(factory=list)
        source_avatars: list[multisource.SourceAvatar] = attr.ib(factory=list)
        avatar_relations: list[multisource.AvatarRelation] = attr.ib(factory=list)
        rls: RLS = attr.ib(factory=RLS)

        component_errors: ComponentErrorRegistry = attr.ib(factory=ComponentErrorRegistry)
        obligatory_filters: list[ObligatoryFilter] = attr.ib(factory=list)

        @attr.s
        class ResultSchemaAux:
            inter_dependencies: FieldInterDependencyInfo = attr.ib(factory=FieldInterDependencyInfo)

        result_schema_aux: ResultSchemaAux = attr.ib(factory=ResultSchemaAux)

    @property
    def rls(self) -> RLS:
        return self.data.rls

    @property
    def error_registry(self) -> ComponentErrorRegistry:
        return self.data.component_errors

    def get_single_data_source_id(self, ignore_source_ids: Optional[Collection[str]] = None) -> str:  # type: ignore  # TODO: fix
        # FIXME: remove in the future
        ignore_source_ids = ignore_source_ids or ()
        for dsrc_coll_spec in self.data.source_collections or ():
            managed_by = dsrc_coll_spec.managed_by or ManagedBy.user
            if dsrc_coll_spec.id in ignore_source_ids or managed_by != ManagedBy.user:
                continue
            return dsrc_coll_spec.id

    def get_own_materialized_tables(self, source_id: Optional[str] = None) -> Generator[str, None, None]:
        for dsrc_coll_spec in self.data.source_collections or ():
            if not dsrc_coll_spec or source_id and dsrc_coll_spec.id != source_id:
                continue
            if not isinstance(dsrc_coll_spec, DataSourceCollectionSpec):
                continue
            for role in (DataSourceRole.materialization, DataSourceRole.sample):
                dsrc_spec = dsrc_coll_spec.get_for_role(role)
                if dsrc_spec is None:
                    continue
                assert isinstance(dsrc_spec, StandardSQLDataSourceSpec)
                if dsrc_spec.table_name is not None:
                    yield dsrc_spec.table_name

    def find_data_source_configuration(  # type: ignore  # TODO: fix
        self,
        connection_id: Optional[str],
        created_from: Optional[DataSourceType] = None,
        title: Optional[str] = None,
        parameters: Optional[dict] = None,
    ) -> Optional[str]:
        """
        Check whether dataset already has a source with the given configuration.
        Return ``None`` if it doesn't exist, otherwise return the corresponding ``source_id``.
        """
        # TODO refactor using lazy datasets instead of configs

        if connection_id is None:
            return None

        if parameters is None:
            parameters = {}

        def spec_matches_parameters(existing_spec: DataSourceSpec) -> bool:
            # FIXME: Refactor
            for key, value in parameters.items():  # type: ignore  # TODO: fix
                if getattr(existing_spec, key, None) != value:
                    return False
            return True

        connection_ref = connection_ref_from_id(connection_id=connection_id)

        for dsrc_coll_spec in self.data.source_collections:
            if (
                isinstance(dsrc_coll_spec, DataSourceCollectionSpec)
                and dsrc_coll_spec.origin is not None
                and dsrc_coll_spec.origin.connection_ref == connection_ref
                and dsrc_coll_spec.origin.source_type == created_from
                and spec_matches_parameters(dsrc_coll_spec.origin)
                and (True if title is None else dsrc_coll_spec.title == title)
            ):
                # reference matches params
                return dsrc_coll_spec.id

    def create_result_schema_field(
        self,
        column: SchemaColumn,
        field_id_generator: Optional[FieldIdGenerator] = None,
        autofix_title: bool = True,
        avatar_id: Optional[str] = None,
    ) -> dict:
        title = column.title
        if autofix_title:
            existing_titles = {f.title for f in self.result_schema.fields}
            title = resolve_id_collisions(item=title, existing_items=existing_titles, formatter="{} ({})")

        if field_id_generator is None:
            FieldIdGeneratorClass = FIELD_ID_GENERATOR_MAP[DEFAULT_FIELD_ID_GENERATOR_TYPE]
            field_id_generator = FieldIdGeneratorClass(dataset=self)

        guid = field_id_generator.make_field_id(title=title)

        hidden = False
        if column.user_type == UserDataType.unsupported:
            # Auto-hide because it's unselectable.
            hidden = True

        column_data = {
            "title": title,
            "guid": guid,
            "hidden": hidden,
            "description": column.description,
            "aggregation": AggregationFunction.none.name,
            "cast": column.user_type.name,
            "type": FieldType.DIMENSION.name,
            "data_type": column.user_type.name,
            "initial_data_type": column.user_type.name,
            "valid": True,
            "has_auto_aggregation": column.has_auto_aggregation,
            "lock_aggregation": column.lock_aggregation,
            "managed_by": ManagedBy.user.name,
            "calc_spec": DirectCalculationSpec(
                avatar_id=avatar_id,
                source=column.name,
            ),
        }
        return column_data

    @property
    def result_schema(self) -> ResultSchema:
        return self.data.result_schema

    def _dump_rls(self):  # type: ignore  # TODO: fix
        """Remove rls entries for non-existing in result_schema fields"""
        # FIXME: this should be in the scope of result_schema management
        field_guids = [f.guid for f in self.result_schema.fields]
        self.rls.items = [rlse for rlse in self.rls.items if rlse.field_guid in field_guids]

    @property
    def created_via(self) -> DataSourceCreatedVia:
        created_via_name = self.meta.get("created_via") or DataSourceCreatedVia.user.name
        return DataSourceCreatedVia[created_via_name]

    @property
    def revision_id(self) -> Optional[str]:
        return self.data.revision_id

    def rename_field_id_usages(self, old_id: str, new_id: str) -> None:
        self.error_registry.rename_pack(old_id=old_id, new_id=new_id)

lib/dl_core/dl_core/us_dataset.py
datalens-tech/datalens-backend
""" ... """

from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Optional,
    Type,
    TypeVar,
    Union,
)

from dl_api_commons.base_models import RequestContextInfo
from dl_core.base_models import (
    BaseAttrsDataModel,
    EntryLocation,
    PathEntryLocation,
    WorkbookEntryLocation,
)
from dl_utils.utils import DotDict


if TYPE_CHECKING:
    from dl_core.us_manager.us_manager import USManagerBase
    from dl_core.us_manager.us_manager_sync import SyncUSManager

LOGGER = logging.getLogger(__name__)


_USENTRY_TV = TypeVar("_USENTRY_TV", bound="USEntry")


class USEntry:
    DataModel: ClassVar[Union[Type[BaseAttrsDataModel], None]] = None
    dir_name = None

    uuid: Optional[str] = None
    _data = None
    entry_key: Optional[EntryLocation] = None
    scope: Optional[str] = None
    type_: Optional[str] = None
    is_locked: Optional[bool] = None
    is_favorite: Optional[bool] = None
    permissions_mode: Optional[str] = None
    initial_permissions: Optional[str] = None
    permissions: Optional[dict[str, bool]] = None
    hidden: bool
    links: Optional[dict] = None

    _stored_in_db = False
    _us_resp: Optional[dict] = None
    _lock = None
    _us_manager: Optional[USManagerBase]

    @classmethod
    def create_from_dict(
        cls: Type[_USENTRY_TV],
        data_dict: Union[dict, BaseAttrsDataModel],
        ds_key: Union[EntryLocation, str, None] = None,
        type_: Optional[str] = None,
        meta: Any = None,
        *,
        us_manager: SyncUSManager,
        **kwargs: Any,
    ) -> _USENTRY_TV:
        from dl_core.us_manager.us_manager import USManagerBase

        if not isinstance(us_manager, USManagerBase):
            raise TypeError(f"us_manager must be USManagerBase, not {type(us_manager)}")

        if not (
            # dict
            isinstance(data_dict, dict)
            and cls.DataModel is None
            # otherwise types must match
            or type(data_dict) is cls.DataModel
        ):
            raise TypeError(f"Invalid object type for data_dict: {type(data_dict)}")

        effective_entry_key: Optional[EntryLocation]

        if ds_key is None:
            effective_entry_key = None
        elif isinstance(ds_key, str):
            effective_entry_key = PathEntryLocation(ds_key)
        elif isinstance(ds_key, EntryLocation):
            effective_entry_key = ds_key
        else:
            raise AssertionError(f"Unexpected type of entry key: {type(ds_key)}")

        obj = cls(
            data=data_dict,  # type: ignore  # TODO: fix
            entry_key=effective_entry_key,
            type_=type_,
            meta=meta,
            us_manager=us_manager,
            **kwargs,
        )
        return obj

    def __init__(
        self,
        uuid: Optional[str] = None,
        data: Optional[dict] = None,
        entry_key: Optional[EntryLocation] = None,
        type_: Optional[str] = None,
        meta: Optional[dict] = None,
        is_locked: Optional[bool] = None,
        is_favorite: Optional[bool] = None,
        permissions_mode: Optional[str] = None,
        initial_permissions: Optional[str] = None,
        permissions: Optional[dict[str, bool]] = None,
        links: Optional[dict] = None,
        hidden: bool = False,
        data_strict: bool = True,
        *,
        us_manager: USManagerBase,
    ):
        if entry_key is not None:
            assert isinstance(entry_key, EntryLocation), f"Unexpected type of entry key: {type(entry_key)}"

        self._us_manager = us_manager
        self.uuid = uuid
        self._data = self._load_data(data, strict=data_strict) if data is not None else None
        self.entry_key = entry_key
        self.type_ = type_
        self.meta = dict(meta or {})
        self.is_locked = is_locked
        self.is_favorite = is_favorite
        self.permissions_mode = permissions_mode
        self.initial_permissions = initial_permissions
        self.permissions = permissions
        self.links = links
        self.hidden = hidden

    @property
    def us_manager(self) -> USManagerBase:
        assert self._us_manager is not None
        return self._us_manager

    @property
    def _context(self) -> RequestContextInfo:
        return self._us_manager.bi_context  # type: ignore  # TODO: fix

    def on_updated(self):  # type: ignore  # TODO: fix
        """Post-update actions go here"""

    def _load_data(self, data: dict, strict: bool = True) -> Union[BaseAttrsDataModel, DotDict]:
        from dl_core.us_manager.us_manager import USManagerBase

        if self.DataModel is None:
            return DotDict(data)
        elif issubclass(self.DataModel, BaseAttrsDataModel):
            if isinstance(data, self.DataModel):
                return data

            schema = USManagerBase.get_load_storage_schema(self.DataModel)
            if schema is None:
                return DotDict(data)

            data = schema.load(data)
            return data  # type: ignore  # TODO: fix
        else:
            raise TypeError(f"Unexpected data type ({type(self.DataModel)}) for entry class {type(self)}")

    def has_data(self) -> bool:
        return self._data is not None

    @property
    def data(self):  # type: ignore  # TODO: fix
        assert self._data is not None
        return self._data

    @data.setter
    def data(self, value):  # type: ignore  # TODO: fix
        self._data = value

    @property
    def raw_us_key(self) -> Optional[str]:
        """
        For backward compatibility.
        Returns private US key if workbook location is used and actual key if path location is used.
        Must be removed after total migration to workbooks.
        """
        if self._us_resp:
            raw_key = self._us_resp.get("key")
            if raw_key is not None:
                assert isinstance(raw_key, str)
                return raw_key
        return None

    @property
    def location_short_string(self) -> Optional[str]:
        loc = self.entry_key
        if loc:
            return loc.to_short_string()
        return None

    def as_dict(self, short: bool = False) -> dict:
        if not short:
            if isinstance(self.data, BaseAttrsDataModel):
                # TODO FIX: implement private fields filtering in schema/USManagerBase
                data_dict = {}
            else:
                data_dict = self.data
            data = {k: v for k, v in data_dict.items()}
        else:
            data = {
                "type": self.type_,
                "created_by": self.created_by,
                "created_at": self.created_at,
                "is_locked": self.is_locked,
            }
        ret = {"id": self.uuid, "key": self.location_short_string, "is_favorite": self.is_favorite, **data}
        if isinstance(self.entry_key, WorkbookEntryLocation):
            ret.update(
                workbook_id=self.entry_key.workbook_id,
                name=self.entry_key.entry_name,
                key=self.raw_us_key,
            )

        return ret

    @property
    def created_by(self):  # type: ignore  # TODO: fix
        return self._us_resp.get("createdBy") if isinstance(self._us_resp, dict) else None  # type: ignore  # TODO: fix

    @property
    def created_at(self):  # type: ignore  # TODO: fix
        return self._us_resp.get("createdAt") if isinstance(self._us_resp, dict) else None  # type: ignore  # TODO: fix

    @property
    def updated_at(self):  # type: ignore  # TODO: fix
        return self._us_resp.get("updatedAt") if isinstance(self._us_resp, dict) else None  # type: ignore  # TODO: fix

    @property
    def revision_id(self):  # type: ignore  # TODO: fix
        return self._us_resp.get("revId") if isinstance(self._us_resp, dict) else None  # type: ignore  # TODO: fix

    @property
    def raw_tenant_id(self) -> Optional[str]:
        if not isinstance(self._us_resp, dict):
            raise ValueError("Can not get tenantId from US entry without ._us_resp.")

        may_be_tenant_id = self._us_resp.get("tenantId")

        if may_be_tenant_id is None or isinstance(may_be_tenant_id, str):
            return may_be_tenant_id
        raise ValueError(
            f"Gon unexpected type of tenantId in US entry {self.scope}/{self.type_}/{self.uuid}: {may_be_tenant_id!r}"
        )

    def validate(self) -> None:
        """
        Fast data validation method. Should not be called in constructor.
        It's to instantiate invalid object.
        Purposes:
         * application-level validation after create/update by end user
         * making decision if entry can be used
        Think about extracting to separate layer
        """
        pass


class USMigrationEntry(USEntry):
    def __init__(
        self,
        uuid: Optional[str] = None,
        data: Optional[dict] = None,
        entry_key: Optional[EntryLocation] = None,
        type_: Optional[str] = None,
        meta: Optional[dict] = None,
        is_locked: Optional[bool] = None,
        is_favorite: Optional[bool] = None,
        permissions_mode: Optional[str] = None,
        initial_permissions: Optional[str] = None,
        permissions: Optional[dict] = None,
        links: Optional[dict] = None,
        hidden: bool = False,
        data_strict: bool = True,
        *,
        us_manager: USManagerBase,
        unversioned_data: Optional[dict[str, Any]],
    ):
        super().__init__(
            uuid=uuid,
            data=data,
            entry_key=entry_key,
            type_=type_,
            meta=meta,
            is_locked=is_locked,
            is_favorite=is_favorite,
            permissions_mode=permissions_mode,
            initial_permissions=initial_permissions,
            permissions=permissions,
            links=links,
            hidden=hidden,
            data_strict=data_strict,
            us_manager=us_manager,
        )

        self._unversioned_data = DotDict(unversioned_data) if unversioned_data is not None else None

    @property
    def unversioned_data(self) -> DotDict:
        return self._unversioned_data  # type: ignore  # TODO: fix

lib/dl_core/dl_core/us_entry.py
datalens-tech/datalens-backend
from __future__ import annotations

import enum
from typing import TYPE_CHECKING

import attr


if TYPE_CHECKING:
    from dl_core.base_models import ConnectionRef


class BrokenUSLinkErrorKind(enum.Enum):
    NOT_FOUND = enum.auto()
    OTHER = enum.auto()
    ACCESS_DENIED = enum.auto()


@attr.s
class BrokenUSLink:
    error_kind: BrokenUSLinkErrorKind = attr.ib()
    reference: ConnectionRef = attr.ib()
    _referrer_id_set: set[str] = attr.ib(init=False, factory=set)

    def __attrs_post_init__(self):  # type: ignore  # TODO: fix
        self._referrer_id_set = set(self._referrer_id_set)

    @property
    def referrer_id_set(self) -> frozenset[str]:
        return frozenset(self._referrer_id_set)

    def add_referrer_id(self, referrer_id: str) -> None:
        self._referrer_id_set.add(referrer_id)

lib/dl_core/dl_core/us_manager/broken_link.py
datalens-tech/datalens-backend
from __future__ import annotations

import typing

import attr
from cryptography import fernet

from dl_configs.crypto_keys import CryptoKeysConfig


class EncryptedData(typing.TypedDict):
    key_id: str
    key_kind: str
    cypher_text: str


@attr.s
class CryptoController:
    key_kind = "local_fernet"

    _key_config: CryptoKeysConfig = attr.ib()
    _map_key_id_fernet_instance: typing.Dict[str, fernet.Fernet] = attr.ib(init=False)

    def __attrs_post_init__(self):  # type: ignore  # TODO: fix
        self._map_key_id_fernet_instance = {
            key_id: fernet.Fernet(key_val) for key_id, key_val in self._key_config.map_id_key.items()
        }

    def copy(self) -> "CryptoController":
        return attr.evolve(self)

    def encrypt(self, key_id: str, plain_text: typing.Optional[str]) -> typing.Optional[EncryptedData]:
        if plain_text is None:
            return None

        encoded_plain_text = plain_text.encode()
        cypher_text = self._map_key_id_fernet_instance[key_id].encrypt(encoded_plain_text).decode()
        return dict(
            key_id=key_id,
            key_kind=self.key_kind,
            cypher_text=cypher_text,
        )

    def decrypt(self, encrypted_data: typing.Optional[EncryptedData]) -> typing.Optional[str]:
        if encrypted_data is None:
            return None

        assert encrypted_data["key_kind"] == self.key_kind
        encoded_cypher_text = encrypted_data["cypher_text"].encode()
        key_id = encrypted_data["key_id"]
        plain_text = self._map_key_id_fernet_instance[key_id].decrypt(encoded_cypher_text).decode()
        return plain_text

    def encrypt_with_actual_key(self, plain_text: typing.Optional[str]) -> typing.Optional[EncryptedData]:
        return self.encrypt(key_id=self._key_config.actual_key_id, plain_text=plain_text)

    @property
    def actual_key_id(self) -> str:
        return self._key_config.actual_key_id

lib/dl_core/dl_core/us_manager/crypto/main.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_constants.api_constants import DLHeadersCommon
from dl_core.services_registry import ServicesRegistry
from dl_core.united_storage_client import (
    USAuthContextEmbed,
    USAuthContextMaster,
    USAuthContextPublic,
    USAuthContextRegular,
)
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_core.us_manager.us_manager_sync import SyncUSManager


@attr.s(frozen=True)
class USMFactory:
    us_base_url: str = attr.ib()
    crypto_keys_config: Optional[CryptoKeysConfig] = attr.ib()
    us_master_token: Optional[str] = attr.ib(default=None, repr=False)
    us_public_token: Optional[str] = attr.ib(default=None, repr=False)

    @classmethod
    def get_regular_us_auth_ctx_from_rci(cls, rci: RequestContextInfo) -> USAuthContextRegular:
        auth_data = rci.auth_data
        tenant = rci.tenant
        assert tenant is not None
        assert auth_data is not None

        return USAuthContextRegular(
            auth_data=auth_data,
            tenant=tenant,
            dl_allow_super_user=rci.plain_headers.get(DLHeadersCommon.ALLOW_SUPERUSER.value),
            dl_sudo=rci.plain_headers.get(DLHeadersCommon.SUDO.value),
        )

    def def_embed_us_auth_ctx_from_rci(self, rci: RequestContextInfo) -> USAuthContextEmbed:
        auth_data = rci.auth_data
        tenant = rci.tenant
        assert tenant is not None
        assert auth_data is not None

        return USAuthContextEmbed(
            auth_data=auth_data,
            tenant=tenant,
        )

    def get_master_auth_context(self) -> USAuthContextMaster:
        assert self.us_master_token is not None, "US master token must be set in factory to create USAuthContextMaster"
        return USAuthContextMaster(us_master_token=self.us_master_token)

    # Async
    def get_regular_async_usm(
        self,
        rci: RequestContextInfo,
        services_registry: ServicesRegistry,
    ) -> AsyncUSManager:
        return AsyncUSManager(
            us_auth_context=self.get_regular_us_auth_ctx_from_rci(rci),
            us_base_url=self.us_base_url,
            bi_context=rci,
            crypto_keys_config=self.crypto_keys_config,
            services_registry=services_registry,
        )

    def get_master_async_usm(
        self,
        rci: RequestContextInfo,
        services_registry: ServicesRegistry,
    ) -> AsyncUSManager:
        return AsyncUSManager(
            us_auth_context=self.get_master_auth_context(),
            us_base_url=self.us_base_url,
            bi_context=rci,
            crypto_keys_config=self.crypto_keys_config,
            services_registry=services_registry,
        )

    def get_public_async_usm(
        self,
        rci: RequestContextInfo,
        services_registry: ServicesRegistry,
    ) -> AsyncUSManager:
        assert self.us_public_token is not None, "US public token must be set in factory to create USAuthContextMaster"

        return AsyncUSManager(
            us_auth_context=USAuthContextPublic(
                us_public_token=self.us_public_token,
            ),
            us_base_url=self.us_base_url,
            bi_context=rci,
            crypto_keys_config=self.crypto_keys_config,
            services_registry=services_registry,
        )

    # Sync
    def get_regular_sync_usm(
        self,
        rci: RequestContextInfo,
        services_registry: ServicesRegistry,
    ) -> SyncUSManager:
        return SyncUSManager(
            us_auth_context=self.get_regular_us_auth_ctx_from_rci(rci),
            us_base_url=self.us_base_url,
            bi_context=rci,
            crypto_keys_config=self.crypto_keys_config,
            services_registry=services_registry,
        )

    def get_master_sync_usm(
        self,
        rci: RequestContextInfo,
        services_registry: ServicesRegistry,
    ) -> SyncUSManager:
        return SyncUSManager(
            us_auth_context=self.get_master_auth_context(),
            us_base_url=self.us_base_url,
            bi_context=rci,
            crypto_keys_config=self.crypto_keys_config,
            services_registry=services_registry,
        )

    def get_embed_async_usm(
        self,
        rci: RequestContextInfo,
        services_registry: ServicesRegistry,
    ) -> AsyncUSManager:
        return AsyncUSManager(
            us_auth_context=self.def_embed_us_auth_ctx_from_rci(rci),
            us_base_url=self.us_base_url,
            bi_context=rci,
            crypto_keys_config=self.crypto_keys_config,
            services_registry=services_registry,
        )

lib/dl_core/dl_core/us_manager/factory.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Iterable,
    Union,
)

import attr

from dl_core.base_models import (
    ConnectionRef,
    DefaultConnectionRef,
)
import dl_core.exc as exc
from dl_core.us_manager.broken_link import (
    BrokenUSLink,
    BrokenUSLinkErrorKind,
)


if TYPE_CHECKING:
    from dl_core.us_entry import USEntry


@attr.s
class USEntryBuffer:
    _data: dict[ConnectionRef, Union[USEntry, BrokenUSLink]] = attr.ib(factory=dict)

    def get_entry(self, entry_id: ConnectionRef) -> USEntry:
        entry = self._data.get(entry_id)

        if isinstance(entry, BrokenUSLink):
            if isinstance(entry.reference, DefaultConnectionRef):
                if entry.error_kind == BrokenUSLinkErrorKind.NOT_FOUND:
                    raise exc.ReferencedUSEntryNotFound(f"Referenced connection {entry.reference.conn_id} was deleted")
                elif entry.error_kind == BrokenUSLinkErrorKind.ACCESS_DENIED:
                    raise exc.ReferencedUSEntryAccessDenied(
                        f"Referenced connection {entry.reference.conn_id} cannot be loaded: access denied",
                        details=dict(scope="connection", entry_id=entry.reference.conn_id),
                    )
                else:
                    raise ValueError(f"Referenced connection {entry.reference} cannot be loaded: {entry.error_kind}")
            else:
                raise ValueError(f"Requested referenced US entry {entry_id} is broken: {entry}", entry)

        if entry is None:
            raise ValueError(f"Connection {entry_id} is not loaded")

        return entry

    def set_entry(self, entry_id: ConnectionRef, entry: Union[USEntry, BrokenUSLink]) -> None:
        self._data[entry_id] = entry

    def __getitem__(self, entry_id: ConnectionRef) -> USEntry:
        return self.get_entry(entry_id)

    def __setitem__(self, entry_id: ConnectionRef, entry: Union[USEntry, BrokenUSLink]) -> None:
        self.set_entry(entry_id, entry)

    def __contains__(self, entry_id: ConnectionRef) -> bool:
        return entry_id in self._data

    def keys(self) -> Iterable[ConnectionRef]:
        return self._data.keys()

    def clear(self) -> None:
        self._data.clear()

lib/dl_core/dl_core/us_manager/local_cache.py
datalens-tech/datalens-backend
import abc
import logging
from typing import (
    TYPE_CHECKING,
    Optional,
    Type,
)

import attr

from dl_core.us_manager.mutation_cache.usentry_mutation_cache import (
    GenericCacheEngine,
    MemoryCacheEngine,
    RedisCacheEngine,
)
from dl_core.utils import FutureRef


if TYPE_CHECKING:
    from dl_core.services_registry.top_level import ServicesRegistry  # noqa

LOGGER = logging.getLogger(__name__)


class CacheInitializationError(Exception):
    pass


class MutationCacheEngineFactory(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_cache_engine(self) -> GenericCacheEngine:
        pass


@attr.s
class DefaultMutationCacheEngineFactory(MutationCacheEngineFactory):
    _services_registry_ref: FutureRef["ServicesRegistry"] = attr.ib()
    cache_type: Type[GenericCacheEngine] = attr.ib()
    _saved_inmemory_engine: Optional[MemoryCacheEngine] = None

    @classmethod
    def _get_memory_cache_engine_singleton(cls) -> MemoryCacheEngine:
        if cls._saved_inmemory_engine is None:
            cls._saved_inmemory_engine = MemoryCacheEngine()
        return cls._saved_inmemory_engine

    @property
    def service_registry(self) -> "ServicesRegistry":
        return self._services_registry_ref.ref

    def _get_redis_cache_engine(self) -> Optional[RedisCacheEngine]:
        try:
            redis_client = self.service_registry.get_mutations_redis_client()
            if redis_client is None:
                return None
            return RedisCacheEngine(redis_client)
        except ValueError:
            return None

    def get_cache_engine(self) -> GenericCacheEngine:
        if self.cache_type == MemoryCacheEngine:
            return self._get_memory_cache_engine_singleton()
        elif self.cache_type == RedisCacheEngine:
            redis_cache_engine = self._get_redis_cache_engine()
            if redis_cache_engine:
                return redis_cache_engine
            LOGGER.info("Can not create mutation cache engine: service registry did not return a Redis client")
            raise CacheInitializationError("Cannot create mutation cache engine")
        else:
            raise CacheInitializationError("No initialization for this type of Cache Engine in factory")

lib/dl_core/dl_core/us_manager/mutation_cache/engine_factory.py
datalens-tech/datalens-backend
import abc
from typing import Any


class MutationKey(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def get_hash(self) -> str:
        """Return string representation of hash that will become part of cache key."""
        raise NotImplementedError()

    @abc.abstractmethod
    def get_collision_tier_breaker(self) -> Any:
        """Returns less collision-affected but serializable representation of key."""
        raise NotImplementedError()

    def __eq__(self, other: Any) -> bool:
        """Overriden comparation operator. Used in tests."""
        if other is self:
            return True
        if not isinstance(other, MutationKey):
            return False
        if self.get_hash() != other.get_hash():
            return False
        return self.get_collision_tier_breaker() == other.get_collision_tier_breaker()

lib/dl_core/dl_core/us_manager/mutation_cache/mutation_key_base.py
datalens-tech/datalens-backend
import abc
import datetime
import json
import logging
import time
from typing import (
    Dict,
    Optional,
    Tuple,
    Type,
)

import attr
from redis.asyncio import Redis
import redis.exceptions

from dl_core.us_entry import USEntry
from dl_core.us_manager.mutation_cache.mutation_key_base import MutationKey
from dl_core.us_manager.us_manager import USManagerBase


LOGGER = logging.getLogger(__name__)


class GenericCacheEngine(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    async def save(self, key: str, data: str, ttl: float) -> None:
        raise NotImplementedError()

    @abc.abstractmethod
    async def load(self, key: str) -> Optional[str]:
        raise NotImplementedError()


class MemoryCacheEngine(GenericCacheEngine):
    def __init__(self) -> None:
        # Dict[ Key: (expired_time, cached value) ]
        self._cache: Dict[str, Tuple[float, str]] = dict()

    async def save(self, key: str, data: str, ttl: float) -> None:
        self._cache[key] = (time.monotonic() + ttl, data)

    async def load(self, key: str) -> Optional[str]:
        cached = self._cache.get(key)
        if not cached:
            return None
        expires_in, cached_value = cached
        if expires_in < time.monotonic():
            return None
        return cached_value


class MutationCacheError(redis.exceptions.ConnectionError):
    pass


@attr.s
class RedisCacheEngine(GenericCacheEngine):
    redis: Redis = attr.ib()

    async def save(self, key: str, data: str, ttl: float) -> None:
        try:
            await self.redis.set(name=key, value=data, ex=datetime.timedelta(seconds=ttl))
        except redis.exceptions.ConnectionError as e:
            LOGGER.warning("Error while saving to redis cache: connection error")
            raise MutationCacheError() from e

    async def load(self, key: str) -> Optional[str]:
        try:
            return await self.redis.get(key)
        except redis.exceptions.ConnectionError as e:
            LOGGER.warning("Error while loading from redis cache: connection error")
            raise MutationCacheError() from e


@attr.s(auto_attribs=True, frozen=True, kw_only=True)
class USEntryMutationCacheKey:
    scope: str
    entry_id: str
    entry_revision_id: str
    mutation_key_hash: str

    def to_string(self) -> str:
        return f"{self.scope}:{self.entry_id}:{self.entry_revision_id}:{self.mutation_key_hash}"


@attr.s()
class USEntryMutationCache:
    _usm: USManagerBase = attr.ib()
    _cache_engine: GenericCacheEngine = attr.ib()
    _default_ttl: float = attr.ib()

    def _prepare_cache_data(self, entry: USEntry, mutation_key: MutationKey) -> str:
        serialized_entry = self._usm._get_entry_save_params(entry)
        serialized_entry["key"] = entry.raw_us_key
        serialized_entry["entryId"] = entry.uuid
        serialized_entry["unversionedData"] = serialized_entry.pop("unversioned_data")
        return self._dump_raw_cache_data(mutation_key.get_collision_tier_breaker(), serialized_entry)

    def _restore_entry_from_cache_representation(self, data: str) -> USEntry:
        entry_dict = json.loads(data)
        return self._usm._entry_dict_to_obj(entry_dict)

    async def save_mutation_cache(self, entry: USEntry, mutation_key: MutationKey) -> None:
        key = USEntryMutationCacheKey(
            scope=entry.scope,  # type: ignore  # TODO: Fix
            entry_id=entry.uuid,  # type: ignore  # TODO: Fix
            entry_revision_id=entry.revision_id,
            mutation_key_hash=mutation_key.get_hash(),
        )
        data = self._prepare_cache_data(entry, mutation_key)
        try:
            await self._cache_engine.save(key=key.to_string(), data=data, ttl=self._default_ttl)
        except MutationCacheError:
            LOGGER.warning("Error saving to cache", exc_info=True)

    @classmethod
    def _load_raw_cache_data(cls, data: str) -> Tuple[str, str]:
        all_data: dict = json.loads(data)
        return (
            all_data["collision_meta"],
            all_data["us_entry_data"],
        )

    @classmethod
    def _dump_raw_cache_data(cls, collision_meta: str, us_entry_data: dict) -> str:
        return json.dumps(
            {"us_entry_data": json.dumps(us_entry_data), "collision_meta": collision_meta}, ensure_ascii=True
        )

    # TODO FIX: Consider passing whole USEntry to decrease possibility of unauthorized access to cache
    async def get_mutated_entry_from_cache(
        self,
        expected_type: Type[USEntry],
        entry_id: str,
        revision_id: str,
        mutation_key: MutationKey,
    ) -> Optional[USEntry]:
        key = USEntryMutationCacheKey(
            scope=expected_type.scope,  # type: ignore  # TODO: Fix
            entry_id=entry_id,
            entry_revision_id=revision_id,
            mutation_key_hash=mutation_key.get_hash(),
        )

        try:
            data = await self._cache_engine.load(key.to_string())
        except MutationCacheError:
            LOGGER.warning("Error loading from cache", exc_info=True)
            return None

        if data is None:
            return None

        collision_meta, us_entry_data = self._load_raw_cache_data(data)
        if collision_meta != mutation_key.get_collision_tier_breaker():
            return None

        return self._restore_entry_from_cache_representation(us_entry_data)

lib/dl_core/dl_core/us_manager/mutation_cache/usentry_mutation_cache.py
datalens-tech/datalens-backend
import abc
import logging
from typing import Optional

import attr

from dl_core.us_manager.mutation_cache.usentry_mutation_cache import (
    GenericCacheEngine,
    USEntryMutationCache,
)
from dl_core.us_manager.us_manager import USManagerBase


LOGGER = logging.getLogger(__name__)


@attr.s
class USEntryMutationCacheFactory(metaclass=abc.ABCMeta):
    default_ttl: float = attr.ib(default=60)

    @abc.abstractmethod
    def get_mutation_cache(
        self,
        usm: USManagerBase,
        engine: GenericCacheEngine,
        ttl: Optional[float] = None,
    ) -> USEntryMutationCache:
        pass


@attr.s
class DefaultUSEntryMutationCacheFactory(USEntryMutationCacheFactory):
    def get_mutation_cache(
        self,
        usm: USManagerBase,
        engine: GenericCacheEngine,
        ttl: Optional[float] = None,
    ) -> USEntryMutationCache:
        if ttl is None:
            ttl = self.default_ttl
        return USEntryMutationCache(
            usm=usm,
            cache_engine=engine,
            default_ttl=ttl,
        )

lib/dl_core/dl_core/us_manager/mutation_cache/usentry_mutation_cache_factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import enum
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Dict,
    Generic,
    Type,
    TypeVar,
)

from cryptography.fernet import Fernet
from marshmallow import (
    Schema,
    post_dump,
    post_load,
    pre_load,
)


if TYPE_CHECKING:
    from dl_core.us_manager.us_manager import USManagerBase  # noqa


class CtxKey(enum.Enum):
    """Schema context keys"""

    fernet = enum.auto()
    us_manager = enum.auto()
    dsc_id = enum.auto()
    ds_conn_type = enum.auto()


_TARGET_TV = TypeVar("_TARGET_TV")


class BaseStorageSchema(Schema, Generic[_TARGET_TV]):
    TARGET_CLS: ClassVar[Type[_TARGET_TV]]

    @classmethod
    def get_target_cls(cls) -> Type[_TARGET_TV]:
        return cls.TARGET_CLS

    @property
    def fernet(self) -> Fernet:
        return self.context[CtxKey.fernet]

    @property
    def usm(self) -> "USManagerBase":
        return self.context[CtxKey.us_manager]

    # Lifecycle
    def pre_process_input_data(self, data: Dict[str, Any]) -> Dict[str, Any]:  # noqa
        return data

    def post_process_output_data(self, data: Dict[str, Any]) -> Dict[str, Any]:  # noqa
        return data

    def push_ctx(self, data: dict) -> None:
        """Used for pushing data upper-layer data to nested schemas. Note that data is raw dict after pre-processing"""
        pass

    def pop_ctx(self, data: dict) -> None:
        """Cleanup con"""
        pass

    def to_object(self, data: dict) -> _TARGET_TV:
        raise NotImplementedError("This schema is does not implement object deserialization")

    @pre_load(pass_many=False)
    def pre_load(self, data, **_):  # type: ignore  # TODO: fix
        normalized_data = self.pre_process_input_data(data)
        # storing unknown fields
        unknown_keys = normalized_data.keys() - self.fields.keys()
        # TODO FIX: Store unknown fields somewhere
        # unknown = {key: normalized_data.pop(key) for key in unknown_keys}
        for key in unknown_keys:
            normalized_data.pop(key)

        self.push_ctx(normalized_data)
        return normalized_data

    @post_load(pass_many=False)
    def post_load(self, data, **_):  # type: ignore  # TODO: fix
        obj = self.to_object(data)
        self.pop_ctx(data)
        return obj

    @post_dump(pass_many=False)
    def post_dump(self, data: Dict[str, Any], **_) -> Dict[str, Any]:  # type: ignore  # TODO: fix
        normalized_data = self.post_process_output_data(data)
        return normalized_data


class DefaultStorageSchema(BaseStorageSchema[_TARGET_TV], Generic[_TARGET_TV]):
    def to_object(self, data: dict) -> _TARGET_TV:
        return self.get_target_cls()(**data)  # type: ignore


# class CustomNested(fields.Nested):
#     """Nested Field with pre/post processing of value"""
#
#     def __init__(self, nested, pre_dump: str, post_load: str, **kwargs):
#         super().__init__(nested, **kwargs)
#         self._pre_dump_meth_name = pre_dump
#         self._post_load_meth_name = post_load
#
#     def _deserialize(self, value, attr, data, partial=None, **kwargs):
#         loaded_data = super()._deserialize(value, attr, data, partial=partial, **kwargs)
#         return getattr(self.parent, self._post_load_meth_name)(loaded_data)
#
#     def _serialize(self, nested_obj, attr, obj, **kwargs):
#         pre_processed_data = getattr(self.parent, self._pre_dump_meth_name)(nested_obj)
#         return super()._serialize(pre_processed_data, attr, obj, **kwargs)

lib/dl_core/dl_core/us_manager/storage_schemas/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Any

from marshmallow import (
    Schema,
    fields,
    post_load,
)

from dl_constants.enums import (
    ConnectionType,
    UserDataType,
)
from dl_core.db.elements import SchemaColumn
from dl_core.db.native_type import GenericNativeType
from dl_model_tools.schema.dynamic_enum_field import DynamicEnumField


class NativeTypeSchema(Schema):
    """(currently used for uploads / ProviderConnection)"""

    conn_type = DynamicEnumField(ConnectionType, by_value=True)
    name = fields.String()

    @post_load(pass_many=False)
    def to_object(self, data: dict[str, Any], **kwargs: Any) -> Any:
        return GenericNativeType(**data)


class SchemaColumnStorageSchema(Schema):
    """(currently used for uploads / ProviderConnection)"""

    name = fields.String(allow_none=False)
    title = fields.String(allow_none=True)
    user_type = fields.Enum(UserDataType, by_value=False)
    nullable = fields.Boolean()
    native_type = fields.Nested(NativeTypeSchema, allow_none=True)
    source_id = fields.String(allow_none=True)

    @post_load(pass_many=False)
    def to_object(self, data: dict[str, Any], **kwargs: Any) -> Any:
        return SchemaColumn(**data)

lib/dl_core/dl_core/us_manager/storage_schemas/base_types.py
datalens-tech/datalens-backend
from typing import (
    Generic,
    TypeVar,
)

from marshmallow import fields as ma_fields

from dl_constants.enums import RawSQLLevel
from dl_core.base_models import ConnectionDataModelBase
from dl_core.us_connection_base import (
    ConnectionBase,
    ConnectionSQL,
)
from dl_core.us_manager.storage_schemas.base import DefaultStorageSchema


_CONN_DATA_TV = TypeVar("_CONN_DATA_TV", bound=ConnectionDataModelBase)


class BaseConnectionDataStorageSchema(DefaultStorageSchema[_CONN_DATA_TV], Generic[_CONN_DATA_TV]):
    pass


_CB_DATA_TV = TypeVar("_CB_DATA_TV", bound=ConnectionBase.DataModel)


class ConnectionBaseDataStorageSchema(BaseConnectionDataStorageSchema[_CB_DATA_TV], Generic[_CB_DATA_TV]):
    table_name = ma_fields.String(required=False, allow_none=True, load_default=None, dump_default=None)
    sample_table_name = ma_fields.String(required=False, allow_none=True, load_default=None, dump_default=None)
    name = ma_fields.String(required=False, allow_none=True, load_default=None, dump_default=None)
    data_export_forbidden = ma_fields.Boolean(required=False, allow_none=False, load_default=False, dump_default=False)


class CacheableConnectionDataSchemaMixin:
    cache_ttl_sec = ma_fields.Integer(required=False, allow_none=True, load_default=None, dump_default=None)


class SubselectConnectionDataSchemaMixin:
    raw_sql_level = ma_fields.Enum(RawSQLLevel, dump_default=RawSQLLevel.off)


_SQL_CONN_DATA_TV = TypeVar("_SQL_CONN_DATA_TV", bound=ConnectionSQL.DataModel)


class ConnectionSQLDataStorageSchema(
    CacheableConnectionDataSchemaMixin,
    SubselectConnectionDataSchemaMixin,
    ConnectionBaseDataStorageSchema[_SQL_CONN_DATA_TV],
    Generic[_SQL_CONN_DATA_TV],
):
    host = ma_fields.String(required=False, allow_none=True, load_default=None, dump_default=None)
    port = ma_fields.Integer(required=False, allow_none=True, load_default=None, dump_default=None)
    db_name = ma_fields.String(required=False, allow_none=True, load_default=None, dump_default=None)
    username = ma_fields.String(required=False, allow_none=True, load_default=None, dump_default=None)
    password = ma_fields.String(required=False, allow_none=True, load_default=None, dump_default=None)

lib/dl_core/dl_core/us_manager/storage_schemas/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Mapping,
    Optional,
)

from marshmallow.fields import Field

from dl_core.base_models import (
    ConnectionRef,
    DefaultConnectionRef,
)


class ConnectionRefField(Field):
    def _serialize(self, value: Optional[ConnectionRef], attr: Optional[str], obj: Any, **kwargs: Any) -> Optional[str]:
        if value is None:
            raise ValueError("Can not serialize data source without connection reference")
        if isinstance(value, DefaultConnectionRef):
            return value.conn_id
        return None

    def _deserialize(
        self, value: Optional[str], attr: Optional[str], data: Optional[Mapping[str, Any]], **kwargs: Any
    ) -> ConnectionRef:
        """
        Wraps connection ID with ConnRef object
        `None` value will not be passed here, so we should handle it in `.to_object()`
        """
        return DefaultConnectionRef(conn_id=value)  # type: ignore

lib/dl_core/dl_core/us_manager/storage_schemas/connection_ref_field.py
datalens-tech/datalens-backend
from typing import Type

from marshmallow import Schema

from dl_core.us_connection_base import (
    ConnectionBase,
    UnknownConnection,
)
from dl_core.us_manager.storage_schemas.connection import ConnectionBaseDataStorageSchema


MAP_TYPE_TO_SCHEMA_MAP_TYPE_TO_SCHEMA = {}


def register_connection_schema(conn_cls: Type[ConnectionBase], schema_cls: Type[Schema]) -> None:
    MAP_TYPE_TO_SCHEMA_MAP_TYPE_TO_SCHEMA[conn_cls.DataModel] = schema_cls  # type: ignore


class UnknownConnectionStorageSchema(ConnectionBaseDataStorageSchema[UnknownConnection.DataModel]):
    TARGET_CLS = UnknownConnection.DataModel


register_connection_schema(UnknownConnection, UnknownConnectionStorageSchema)

lib/dl_core/dl_core/us_manager/storage_schemas/connection_schema_registry.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Any,
    Optional,
)

from marshmallow import fields
from marshmallow_oneofschema import OneOfSchema

from dl_constants.enums import (
    DataSourceCollectionType,
    ManagedBy,
)
from dl_core.data_source_spec.collection import (
    DataSourceCollectionSpec,
    DataSourceCollectionSpecBase,
)
from dl_core.us_manager.storage_schemas.base import (
    BaseStorageSchema,
    CtxKey,
)
from dl_core.us_manager.storage_schemas.data_source_spec import GenericDataSourceSpecStorageSchema


LOGGER = logging.getLogger(__name__)


class DataSourceCollectionSpecStorageSchema(BaseStorageSchema[DataSourceCollectionSpec]):
    TARGET_CLS = DataSourceCollectionSpec

    id = fields.String(required=False, allow_none=True)
    title = fields.String(required=False, allow_none=True, load_default=None)
    managed_by = fields.Enum(ManagedBy, required=False, allow_none=True)
    valid = fields.Boolean(load_default=True)

    # Sources
    origin = fields.Nested(GenericDataSourceSpecStorageSchema, required=False, allow_none=True)
    materialization = fields.Nested(GenericDataSourceSpecStorageSchema, required=False, allow_none=True)
    sample = fields.Nested(GenericDataSourceSpecStorageSchema, required=False, allow_none=True)

    def push_ctx(self, data: dict) -> None:
        self.context[CtxKey.dsc_id] = data.get("id")

    def pop_ctx(self, data: dict) -> None:
        self.context.pop(CtxKey.dsc_id, None)

    def to_object(self, data):  # type: ignore  # TODO: fix
        return self.get_target_cls()(
            id=data["id"],
            title=data["title"],
            managed_by=data["managed_by"],
            origin=data["origin"],
            materialization=data["materialization"],
            sample=data["sample"],
            valid=data["valid"],
        )


class GenericDataSourceCollectionStorageSchema(OneOfSchema):
    type_field = "type"
    type_field_remove = True
    type_schemas = {
        k.name: v
        for k, v in {
            DataSourceCollectionType.collection: DataSourceCollectionSpecStorageSchema,
        }.items()
    }

    def get_obj_type(self, obj: Any) -> Optional[DataSourceCollectionType]:
        if isinstance(obj, DataSourceCollectionSpecBase):
            return obj.dsrc_coll_type.name  # type: ignore  # TODO: fix
        return None

lib/dl_core/dl_core/us_manager/storage_schemas/data_source_collection.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Type,
)

from marshmallow_oneofschema import OneOfSchema

from dl_constants.enums import DataSourceType
from dl_core.data_source_spec.base import DataSourceSpec
from dl_core.us_manager.storage_schemas.data_source_spec_base import DataSourceSpecStorageSchema


class GenericDataSourceSpecStorageSchema(OneOfSchema):
    type_field = "created_from"
    type_field_remove = False
    type_schemas: dict[str, Type[DataSourceSpecStorageSchema]] = {}

    def get_obj_type(self, obj: Any) -> str:
        assert isinstance(obj, DataSourceSpec)
        return obj.source_type.name


def register_data_source_schema(source_type: DataSourceType, schema_cls: Type[DataSourceSpecStorageSchema]) -> None:
    GenericDataSourceSpecStorageSchema.type_schemas[source_type.name] = schema_cls

lib/dl_core/dl_core/us_manager/storage_schemas/data_source_spec.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Any

import marshmallow.fields as ma_fields

from dl_constants.enums import DataSourceType
from dl_core.base_models import InternalMaterializationConnectionRef
from dl_core.data_source.type_mapping import get_data_source_class
from dl_core.data_source_spec.base import DataSourceSpec
from dl_core.data_source_spec.sql import (
    SQLDataSourceSpecBase,
    StandardSchemaSQLDataSourceSpec,
    StandardSQLDataSourceSpec,
    SubselectDataSourceSpec,
)
from dl_core.marshmallow import FrozenSetField
from dl_core.us_manager.storage_schemas.base import (
    BaseStorageSchema,
    CtxKey,
)
from dl_core.us_manager.storage_schemas.connection_ref_field import ConnectionRefField
from dl_core.us_manager.storage_schemas.index_info import DataSourceIndexInfoStorageSchema
from dl_core.us_manager.storage_schemas.raw_schema import DataSourceRawSchemaEntryStorageSchema
from dl_model_tools.schema.dynamic_enum_field import DynamicEnumField


_DSRC_TOPLEVEL_PARAMS = ("connection_id", "created_from", "raw_schema")


class DataSourceSpecStorageSchema(BaseStorageSchema):  # noqa
    TARGET_CLS = DataSourceSpec

    connection_id = ConnectionRefField(allow_none=True, load_default=None, attribute="connection_ref")
    data_dump_id = ma_fields.String(allow_none=True)
    raw_schema = ma_fields.List(
        ma_fields.Nested(DataSourceRawSchemaEntryStorageSchema),
        load_default=list,
        allow_none=True,
    )
    created_from = DynamicEnumField(DataSourceType, attribute="source_type")

    def pre_process_input_data(self, data: dict[str, Any]) -> dict:
        data = data.copy()
        if "parameters" in data:
            data = {
                **{k: v for k, v in data.items() if k != "parameters"},
                **data["parameters"],
            }
        return data

    def post_process_output_data(self, data: dict[str, Any]) -> dict[str, Any]:  # noqa
        # TODO: Remove. Temporary hack for compatibility with old schemas:
        #  Pack all non-standard keys into "parameters" sub-dict
        parameters = {}
        for param_key in data.copy():
            if param_key not in _DSRC_TOPLEVEL_PARAMS:
                param_value = data.pop(param_key)
                if param_value is not None:  # FIXME: Remove
                    # Save only non-None values here
                    parameters[param_key] = param_value
        data["parameters"] = parameters
        return data

    def push_ctx(self, data: dict):  # type: ignore  # TODO: fix
        dsrc_cls = get_data_source_class(DataSourceType[data["created_from"]])
        self.context[CtxKey.ds_conn_type] = dsrc_cls.conn_type

    def pop_ctx(self, data: dict):  # type: ignore  # TODO: fix
        self.context.pop(CtxKey.ds_conn_type, None)

    def constructor_kwargs(self, data: dict[str, Any]) -> dict[str, Any]:
        data["connection_ref"] = data.pop("connection_ref", None)
        if data["connection_ref"] is None:
            data["connection_ref"] = InternalMaterializationConnectionRef()
        return data

    def to_object(self, data: dict) -> Any:
        kw = self.constructor_kwargs(data)
        return self.TARGET_CLS(**kw)


class BaseSQLDataSourceSpecStorageSchema(DataSourceSpecStorageSchema):  # noqa
    TARGET_CLS = SQLDataSourceSpecBase

    db_version = ma_fields.String(required=False, allow_none=True, load_default=None)


class SubselectSQLDataSourceSpecStorageSchemaMixin(BaseSQLDataSourceSpecStorageSchema):  # noqa
    subsql = ma_fields.String(required=False, allow_none=True, load_default=None)


class SubselectDataSourceSpecStorageSchema(
    SubselectSQLDataSourceSpecStorageSchemaMixin,
    BaseSQLDataSourceSpecStorageSchema,
):
    TARGET_CLS = SubselectDataSourceSpec


class DbSQLDataSourceSpecStorageSchemaMixin(BaseSQLDataSourceSpecStorageSchema):  # noqa
    db_name = ma_fields.String(required=False, allow_none=True, load_default=None)


class TableSQLDataSourceSpecStorageSchemaMixin(BaseSQLDataSourceSpecStorageSchema):  # noqa
    table_name = ma_fields.String(required=False, allow_none=True, load_default=None)


class IndexedSQLDataSourceSpecStorageSchemaMixin(BaseSQLDataSourceSpecStorageSchema):  # noqa
    index_info_set = FrozenSetField(
        ma_fields.Nested(DataSourceIndexInfoStorageSchema),
        sort_output=True,
        required=False,
        load_default=None,
        allow_none=True,
    )


class SQLDataSourceSpecStorageSchema(
    DbSQLDataSourceSpecStorageSchemaMixin,
    TableSQLDataSourceSpecStorageSchemaMixin,
    IndexedSQLDataSourceSpecStorageSchemaMixin,
    BaseSQLDataSourceSpecStorageSchema,
):  # noqa
    TARGET_CLS = StandardSQLDataSourceSpec


class SchemaSQLDataSourceSpecStorageSchemaMixin(BaseSQLDataSourceSpecStorageSchema):  # noqa
    schema_name = ma_fields.String(required=False, allow_none=True, load_default=None)


class SchemaSQLDataSourceSpecStorageSchema(SchemaSQLDataSourceSpecStorageSchemaMixin, SQLDataSourceSpecStorageSchema):
    TARGET_CLS = StandardSchemaSQLDataSourceSpec

lib/dl_core/dl_core/us_manager/storage_schemas/data_source_spec_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from copy import deepcopy
from typing import (
    Any,
    Dict,
)

from marshmallow import (
    post_dump,
    post_load,
    pre_load,
)
from marshmallow import fields as ma_fields
from marshmallow_oneofschema import OneOfSchema

from dl_constants.enums import (
    AggregationFunction,
    BinaryJoinOperator,
    CalcMode,
    ConditionPartCalcMode,
    FieldType,
    JoinConditionType,
    JoinType,
    ManagedBy,
    ParameterValueConstraintType,
    RLSPatternType,
    RLSSubjectType,
    UserDataType,
    WhereClauseOperation,
)
from dl_core import multisource
from dl_core import rls as rls_module
from dl_core.base_models import (
    DefaultWhereClause,
    ObligatoryFilter,
)
from dl_core.components.dependencies.primitives import (
    FieldInterDependencyInfo,
    FieldInterDependencyItem,
)
from dl_core.fields import (
    BIField,
    CalculationSpec,
    DirectCalculationSpec,
    FormulaCalculationSpec,
    ParameterCalculationSpec,
    ParameterValueConstraint,
    RangeParameterValueConstraint,
    ResultSchema,
    SetParameterValueConstraint,
    del_calc_spec_kwargs_from,
    filter_calc_spec_kwargs,
)
from dl_core.us_dataset import Dataset
from dl_core.us_manager.storage_schemas.base import DefaultStorageSchema
from dl_core.us_manager.storage_schemas.data_source_collection import GenericDataSourceCollectionStorageSchema
from dl_core.us_manager.storage_schemas.error_registry import ComponentErrorListSchema
from dl_core.values import (
    ArrayFloatValue,
    ArrayIntValue,
    ArrayStrValue,
    BIValue,
    BooleanValue,
    DateTimeTZValue,
    DateTimeValue,
    DateValue,
    FloatValue,
    GenericDateTimeValue,
    GeoPointValue,
    GeoPolygonValue,
    IntegerValue,
    MarkupValue,
    StringValue,
    TreeStrValue,
    UuidValue,
)


class SourceAvatarSchema(DefaultStorageSchema):
    TARGET_CLS = multisource.SourceAvatar

    id = ma_fields.String()
    title = ma_fields.String()
    source_id = ma_fields.String()
    is_root = ma_fields.Boolean()
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user)
    valid = ma_fields.Boolean(allow_none=True)


class ConditionPartBaseSchema(DefaultStorageSchema):
    calc_mode = ma_fields.Enum(ConditionPartCalcMode)


class ConditionPartDirectSchema(ConditionPartBaseSchema):
    TARGET_CLS = multisource.ConditionPartDirect
    source = ma_fields.String()


class ConditionPartResultFieldSchema(ConditionPartBaseSchema):
    TARGET_CLS = multisource.ConditionPartResultField
    field_id = ma_fields.String()


class ConditionPartFormulaSchema(ConditionPartBaseSchema):
    TARGET_CLS = multisource.ConditionPartFormula
    formula = ma_fields.String()


class ConditionPartSchema(OneOfSchema):
    type_field = "calc_mode"
    type_field_remove = False
    type_schemas = {
        k.name: v
        for k, v in {
            ConditionPartCalcMode.direct: ConditionPartDirectSchema,
            ConditionPartCalcMode.result_field: ConditionPartResultFieldSchema,
            ConditionPartCalcMode.formula: ConditionPartFormulaSchema,
        }.items()
    }

    def get_obj_type(self, obj):  # type: ignore  # TODO: fix
        return getattr(obj, self.type_field).name


class AvatarRelationSchema(DefaultStorageSchema):
    TARGET_CLS = multisource.AvatarRelation

    class JoinConditionSchema(DefaultStorageSchema):
        TARGET_CLS = multisource.BinaryCondition

        operator = ma_fields.Enum(BinaryJoinOperator)
        left_part = ma_fields.Nested(ConditionPartSchema)
        right_part = ma_fields.Nested(ConditionPartSchema)
        condition_type = ma_fields.Enum(JoinConditionType)

    id = ma_fields.String()
    left_avatar_id = ma_fields.String()
    right_avatar_id = ma_fields.String()
    conditions = ma_fields.List(ma_fields.Nested(JoinConditionSchema))
    join_type = ma_fields.Enum(JoinType)
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user)
    valid = ma_fields.Boolean(allow_none=True)


class RLSSchema(DefaultStorageSchema):
    TARGET_CLS = rls_module.RLS

    class RLSEntrySchema(DefaultStorageSchema):
        TARGET_CLS = rls_module.RLSEntry

        class RLSSubjectSchema(DefaultStorageSchema):
            TARGET_CLS = rls_module.RLSSubject

            subject_type = ma_fields.Enum(RLSSubjectType)
            subject_id = ma_fields.String()
            subject_name = ma_fields.String()  # login, group name, etc

        field_guid = ma_fields.String()
        # TODO: validate that either of these is correct for each instance:
        #   * pattern_type=RLSPatternType.value and allowed_value is not None
        #   * pattern_type=RLSPatternType.all and allowed_value is None
        allowed_value = ma_fields.String(allow_none=True)
        pattern_type = ma_fields.Enum(RLSPatternType, dump_default=RLSPatternType.value)
        subject = ma_fields.Nested(RLSSubjectSchema)

    items = ma_fields.List(ma_fields.Nested(RLSEntrySchema))

    def pre_process_input_data(self, data):  # type: ignore  # TODO: fix
        return {"items": data}

    @post_dump
    def flatten_items(self, data, **_):  # type: ignore  # TODO: fix
        return data.get("items")


class BaseValueSchema(DefaultStorageSchema):
    type = ma_fields.Enum(UserDataType)
    value = ma_fields.Field()


class StringValueSchema(BaseValueSchema):
    TARGET_CLS = StringValue
    value = ma_fields.String()


class IntegerValueSchema(BaseValueSchema):
    TARGET_CLS = IntegerValue
    value = ma_fields.Integer()


class FloatValueSchema(BaseValueSchema):
    TARGET_CLS = FloatValue
    value = ma_fields.Float()


class DateValueSchema(BaseValueSchema):
    TARGET_CLS = DateValue
    value = ma_fields.Date()


class DateTimeValueSchema(BaseValueSchema):
    TARGET_CLS = DateTimeValue
    value = ma_fields.DateTime()


class DateTimeTZValueSchema(BaseValueSchema):
    TARGET_CLS = DateTimeTZValue
    value = ma_fields.DateTime()


class GenericDateTimeValueSchema(BaseValueSchema):
    TARGET_CLS = GenericDateTimeValue
    value = ma_fields.DateTime()


class BooleanValueSchema(BaseValueSchema):
    TARGET_CLS = BooleanValue
    value = ma_fields.Boolean()


class GeoPointValueSchema(BaseValueSchema):
    TARGET_CLS = GeoPointValue
    value = ma_fields.List(ma_fields.Float())


class GeoPolygonValueSchema(BaseValueSchema):
    TARGET_CLS = GeoPolygonValue
    value = ma_fields.List(ma_fields.List(ma_fields.List(ma_fields.Float())))


class UuidValueSchema(BaseValueSchema):
    TARGET_CLS = UuidValue
    value = ma_fields.String()


class MarkupValueSchema(BaseValueSchema):
    TARGET_CLS = MarkupValue
    value = ma_fields.String()


class ArrayStrValueSchema(BaseValueSchema):
    TARGET_CLS = ArrayStrValue
    value = ma_fields.List(ma_fields.String())


class ArrayIntValueSchema(BaseValueSchema):
    TARGET_CLS = ArrayIntValue
    value = ma_fields.List(ma_fields.Integer())


class ArrayFloatValueSchema(BaseValueSchema):
    TARGET_CLS = ArrayFloatValue
    value = ma_fields.List(ma_fields.Float())


class TreeStrValueSchema(BaseValueSchema):
    TARGET_CLS = TreeStrValue
    value = ma_fields.List(ma_fields.String())


class ValueSchema(OneOfSchema):
    type_field = "type"
    type_schemas = {
        UserDataType.string.name: StringValueSchema,
        UserDataType.integer.name: IntegerValueSchema,
        UserDataType.float.name: FloatValueSchema,
        UserDataType.date.name: DateValueSchema,
        UserDataType.datetime.name: DateTimeValueSchema,
        UserDataType.datetimetz.name: DateTimeTZValueSchema,
        UserDataType.genericdatetime.name: GenericDateTimeValueSchema,
        UserDataType.boolean.name: BooleanValueSchema,
        UserDataType.geopoint.name: GeoPointValueSchema,
        UserDataType.geopolygon.name: GeoPolygonValueSchema,
        UserDataType.uuid.name: UuidValueSchema,
        UserDataType.markup.name: MarkupValueSchema,
        UserDataType.array_str.name: ArrayStrValueSchema,
        UserDataType.array_int.name: ArrayIntValueSchema,
        UserDataType.array_float.name: ArrayFloatValueSchema,
        UserDataType.tree_str.name: TreeStrValueSchema,
    }

    def get_obj_type(self, obj: BIValue) -> str:
        assert isinstance(obj, BIValue)
        return obj.type.name


class BaseParameterValueConstraintSchema(DefaultStorageSchema):
    type = ma_fields.Enum(ParameterValueConstraintType)


class AllParameterValueConstraintSchema(BaseParameterValueConstraintSchema):
    TARGET_CLS = ParameterValueConstraint


class RangeParameterValueConstraintSchema(BaseParameterValueConstraintSchema):
    TARGET_CLS = RangeParameterValueConstraint

    min = ma_fields.Nested(ValueSchema, allow_none=True)
    max = ma_fields.Nested(ValueSchema, allow_none=True)


class SetParameterValueConstraintSchema(BaseParameterValueConstraintSchema):
    TARGET_CLS = SetParameterValueConstraint

    values = ma_fields.List(ma_fields.Nested(ValueSchema))


class ParameterValueConstraintSchema(OneOfSchema):
    type_field = "type"
    type_schemas = {
        ParameterValueConstraintType.all.name: AllParameterValueConstraintSchema,
        ParameterValueConstraintType.range.name: RangeParameterValueConstraintSchema,
        ParameterValueConstraintType.set.name: SetParameterValueConstraintSchema,
    }

    def get_obj_type(self, obj: ParameterValueConstraint) -> str:
        return getattr(obj, self.type_field).name


class DirectCalculationSpecSchema(DefaultStorageSchema):
    TARGET_CLS = DirectCalculationSpec
    avatar_id = ma_fields.String(allow_none=True)
    source = ma_fields.String()


class FormulaCalculationSpecSchema(DefaultStorageSchema):
    TARGET_CLS = FormulaCalculationSpec
    formula = ma_fields.String()
    guid_formula = ma_fields.String()


class ParameterCalculationSpecSchema(DefaultStorageSchema):
    TARGET_CLS = ParameterCalculationSpec
    default_value = ma_fields.Nested(ValueSchema, allow_none=True)
    value_constraint = ma_fields.Nested(ParameterValueConstraintSchema, allow_none=True)


class ResultSchemaStorageSchema(DefaultStorageSchema):
    TARGET_CLS = ResultSchema

    class BIFieldSchema(DefaultStorageSchema):
        TARGET_CLS = BIField

        class CalculationSpecSchema(OneOfSchema):
            type_field = "mode"
            type_field_remove = False
            type_schemas = {
                CalcMode.direct.name: DirectCalculationSpecSchema,
                CalcMode.formula.name: FormulaCalculationSpecSchema,
                CalcMode.parameter.name: ParameterCalculationSpecSchema,
            }

            def get_obj_type(self, obj: CalculationSpec) -> str:
                return obj.mode.name

        title = ma_fields.String()
        guid = ma_fields.String()
        aggregation = ma_fields.Enum(AggregationFunction)
        type = ma_fields.Enum(FieldType)
        hidden = ma_fields.Boolean()
        description = ma_fields.String()
        cast = ma_fields.Enum(UserDataType, allow_none=True)
        initial_data_type = ma_fields.Enum(UserDataType, allow_none=True)
        data_type = ma_fields.Enum(UserDataType, allow_none=True)
        valid = ma_fields.Boolean(allow_none=True)
        has_auto_aggregation = ma_fields.Boolean(allow_none=True)
        lock_aggregation = ma_fields.Boolean(allow_none=True)
        managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user)

        # this will be flattened on dump and un-flattened before load
        # TODO: dump/load as is and migrate data in storage respectively
        calc_spec = ma_fields.Nested(CalculationSpecSchema)

        @post_dump(pass_many=False)
        def add_calc_spec(self, data: Dict[str, Any], **_: Any) -> Dict[str, Any]:
            data = deepcopy(data)
            calc_spec_data = data.pop("calc_spec")
            calc_spec_data["calc_mode"] = calc_spec_data.pop("mode")
            data.update(calc_spec_data)
            # For backward compatibility use '' for formula and source; avatar_id must be present even if None
            for key in ("formula", "guid_formula", "source"):
                data.setdefault(key, "")
            for key in ("avatar_id", "default_value", "value_constraint"):
                data.setdefault(key, None)
            return data

        @pre_load(pass_many=False)
        def extract_calc_spec(self, data: Dict[str, Any], **_: Any) -> Dict[str, Any]:
            data = deepcopy(data)
            mode = data["calc_mode"]
            data["calc_spec"] = dict(filter_calc_spec_kwargs(mode, data), mode=mode)
            return del_calc_spec_kwargs_from(data)

        def to_object(self, data: dict) -> BIField:
            return BIField.make(**data)

    fields = ma_fields.List(ma_fields.Nested(BIFieldSchema))  # type: ignore  # TODO: fix

    def pre_process_input_data(self, data):  # type: ignore  # TODO: fix
        return {"fields": data}

    @post_dump
    def flatten_fields(self, data, **_):  # type: ignore  # TODO: fix
        return data.get("fields")

    @post_load
    def make_missing_formulas(self, data: Dict[str, Any], **_: Any) -> Dict[str, Any]:
        field: BIField
        titles_to_guids = {field.title: field.guid for field in data["fields"]}
        guids_to_titles = {field.guid: field.title for field in data["fields"]}
        new_fields = []
        for field in data["fields"]:
            if field.calc_mode == CalcMode.formula:
                new_field = field
                if field.formula and not field.guid_formula:
                    new_field = field.clone(
                        guid_formula=BIField.rename_in_formula(formula=field.formula, key_map=titles_to_guids)
                    )
                elif field.guid_formula and not field.formula:
                    new_field = field.clone(
                        formula=BIField.rename_in_formula(formula=field.guid_formula, key_map=guids_to_titles)
                    )
                new_fields.append(new_field)
            else:
                new_fields.append(field)
        data["fields"] = new_fields
        return data


class ResultSchemaAuxSchema(DefaultStorageSchema):
    TARGET_CLS = Dataset.DataModel.ResultSchemaAux

    class FieldInterDependencyInfoSchema(DefaultStorageSchema):
        TARGET_CLS = FieldInterDependencyInfo

        class FieldInterDependencyItemSchema(DefaultStorageSchema):
            TARGET_CLS = FieldInterDependencyItem

            dep_field_id = ma_fields.String(allow_none=False)
            ref_field_ids = ma_fields.List(ma_fields.String(allow_none=False), allow_none=False)

        deps = ma_fields.List(ma_fields.Nested(FieldInterDependencyItemSchema()), allow_none=False)

    inter_dependencies = ma_fields.Nested(FieldInterDependencyInfoSchema(), allow_none=False)


class DefaultWhereClauseSchema(DefaultStorageSchema):
    TARGET_CLS = DefaultWhereClause

    operation = ma_fields.Enum(WhereClauseOperation)
    values = ma_fields.List(ma_fields.String())


class ObligatoryFilterSchema(DefaultStorageSchema):
    TARGET_CLS = ObligatoryFilter

    id = ma_fields.String()
    field_guid = ma_fields.String()
    default_filters = ma_fields.List(ma_fields.Nested(DefaultWhereClauseSchema))
    managed_by = ma_fields.Enum(ManagedBy, allow_none=True, dump_default=ManagedBy.user)
    valid = ma_fields.Boolean(allow_none=True)


class DatasetStorageSchema(DefaultStorageSchema):
    TARGET_CLS = Dataset.DataModel

    name = ma_fields.String(allow_none=True, load_default=None)
    revision_id = ma_fields.String(allow_none=True, dump_default=None, load_default=None)
    result_schema = ma_fields.Nested(ResultSchemaStorageSchema, allow_none=False)
    result_schema_aux = ma_fields.Nested(ResultSchemaAuxSchema, allow_none=False)
    source_collections = ma_fields.List(ma_fields.Nested(GenericDataSourceCollectionStorageSchema, allow_none=False))
    source_avatars = ma_fields.List(ma_fields.Nested(SourceAvatarSchema, allow_none=False))
    avatar_relations = ma_fields.List(ma_fields.Nested(AvatarRelationSchema, allow_none=False))
    rls = ma_fields.Nested(RLSSchema, allow_none=False)
    component_errors = ma_fields.Nested(ComponentErrorListSchema)
    obligatory_filters = ma_fields.List(ma_fields.Nested(ObligatoryFilterSchema))

lib/dl_core/dl_core/us_manager/storage_schemas/dataset.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import fields as ma_fields
from marshmallow_oneofschema import OneOfSchema

from dl_constants.enums import (
    ComponentErrorLevel,
    ComponentType,
)
from dl_core import component_errors
from dl_core.us_manager.storage_schemas.base import DefaultStorageSchema


class GenericComponentErrorPackSchema(DefaultStorageSchema):
    TARGET_CLS = component_errors.ComponentErrorPack

    class GenericComponentErrorSchema(DefaultStorageSchema):
        TARGET_CLS = component_errors.ComponentError

        message = ma_fields.String()
        level = ma_fields.Enum(ComponentErrorLevel)
        code = ma_fields.List(ma_fields.String())
        details = ma_fields.Dict()

    id = ma_fields.String()
    type = ma_fields.Enum(ComponentType)
    errors = ma_fields.List(ma_fields.Nested(GenericComponentErrorSchema))


class ComponentErrorPackSchema(OneOfSchema):
    type_field = "type"
    type_field_remove = False
    type_schemas = {
        k.name: v
        for k, v in {
            ComponentType.data_source: GenericComponentErrorPackSchema,
            ComponentType.source_avatar: GenericComponentErrorPackSchema,
            ComponentType.avatar_relation: GenericComponentErrorPackSchema,
            ComponentType.field: GenericComponentErrorPackSchema,
            ComponentType.obligatory_filter: GenericComponentErrorPackSchema,
        }.items()
    }

    def get_obj_type(self, obj):  # type: ignore  # TODO: fix
        return getattr(obj, self.type_field).name


class ComponentErrorListSchema(DefaultStorageSchema):
    TARGET_CLS = component_errors.ComponentErrorRegistry

    items = ma_fields.List(ma_fields.Nested(ComponentErrorPackSchema))

lib/dl_core/dl_core/us_manager/storage_schemas/error_registry.py
datalens-tech/datalens-backend
from __future__ import annotations

from marshmallow import fields

from dl_constants.enums import IndexKind
from dl_core.db import IndexInfo
from dl_core.us_manager.storage_schemas.base import BaseStorageSchema


class DataSourceIndexInfoStorageSchema(BaseStorageSchema[IndexInfo]):
    columns = fields.List(fields.String)
    kind = fields.Enum(IndexKind, allow_none=True)

    def to_object(self, data: dict) -> IndexInfo:
        columns = tuple(data.pop("columns"))
        return IndexInfo(
            columns=columns,
            **data,
        )

lib/dl_core/dl_core/us_manager/storage_schemas/index_info.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Any,
    Dict,
)

from marshmallow import fields

from dl_constants.enums import UserDataType
from dl_core.db import SchemaColumn
from dl_core.db.native_type_schema import OneOfNativeTypeSchema
from dl_core.us_manager.storage_schemas.base import (
    BaseStorageSchema,
    CtxKey,
)


LOGGER = logging.getLogger(__name__)


class DataSourceRawSchemaEntryStorageSchema(BaseStorageSchema[SchemaColumn]):
    name = fields.String(required=False, allow_none=True)
    title = fields.String(required=False, allow_none=True)
    description = fields.String(required=False, allow_none=True)

    type = fields.Enum(UserDataType, attribute="user_type")
    nullable = fields.Boolean(required=False, allow_none=True)

    lock_aggregation = fields.Boolean(required=False, allow_none=True)
    has_auto_aggregation = fields.Boolean(required=False, allow_none=True)

    native_type = fields.Nested(OneOfNativeTypeSchema, required=False, load_default=None)

    def to_object(self, data: Dict[str, Any]) -> SchemaColumn:
        return SchemaColumn(
            source_id=self.context.get(CtxKey.dsc_id),
            name=data["name"],
            title=data["title"],
            description=data["description"],
            user_type=data["user_type"],
            native_type=data["native_type"],
            nullable=data["nullable"],
            lock_aggregation=data["lock_aggregation"],
            has_auto_aggregation=data["has_auto_aggregation"],
        )

lib/dl_core/dl_core/us_manager/storage_schemas/raw_schema.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from collections import ChainMap
import copy
from functools import reduce
import json
import logging
from typing import (
    ClassVar,
    Optional,
    Type,
)
from typing import (
    TYPE_CHECKING,
    Any,
)
from typing import ChainMap as ChainMapGeneric

import attr
import marshmallow

from dl_core import exc
from dl_core.base_models import BaseAttrsDataModel
from dl_core.us_dataset import Dataset
from dl_core.us_entry import USEntry
from dl_core.us_manager.crypto.main import EncryptedData
from dl_core.us_manager.storage_schemas.connection_schema_registry import MAP_TYPE_TO_SCHEMA_MAP_TYPE_TO_SCHEMA
from dl_core.us_manager.storage_schemas.dataset import DatasetStorageSchema
from dl_utils.utils import (
    AddressableData,
    DataKey,
)


if TYPE_CHECKING:
    from dl_core.us_manager.us_manager import USManagerBase


LOGGER = logging.getLogger(__name__)


@attr.s()
class USDataPack:
    data: dict[str, Any] = attr.ib()
    secrets: dict[str, Optional[str | EncryptedData]] = attr.ib(repr=False, factory=dict)


class USEntrySerializer(abc.ABC):
    _MAP_TYPE_TO_SCHEMA: ClassVar[ChainMapGeneric[Type[BaseAttrsDataModel], Type[marshmallow.Schema]]] = ChainMap(
        MAP_TYPE_TO_SCHEMA_MAP_TYPE_TO_SCHEMA,  # type: ignore
        {
            Dataset.DataModel: DatasetStorageSchema,
        },
    )

    @classmethod
    def get_load_storage_schema(cls, data_cls: Type[BaseAttrsDataModel]) -> marshmallow.Schema:
        schema_cls = cls._MAP_TYPE_TO_SCHEMA[data_cls]
        return schema_cls()

    @classmethod
    def get_dump_storage_schema(cls, data_cls: Type[BaseAttrsDataModel]) -> marshmallow.Schema:
        schema_cls = cls._MAP_TYPE_TO_SCHEMA[data_cls]
        return schema_cls()

    @abc.abstractmethod
    def serialize_raw(self, entry: USEntry) -> dict[str, Any]:
        """Returns raw entry data dict"""

        raise NotImplementedError()

    @abc.abstractmethod
    def deserialize_raw(
        self,
        cls: Type[USEntry],
        raw_data: dict[str, Any],
        entry_id: str,
        us_manager: USManagerBase,
        common_properties: dict[str, Any],
        data_strict: bool = True,
    ) -> USEntry:
        """Initializes USEntry with data provided in raw_data"""

        raise NotImplementedError()

    @abc.abstractmethod
    def get_secret_keys(self, cls: Type[USEntry]) -> set[DataKey]:
        raise NotImplementedError()

    @abc.abstractmethod
    def set_data_attr(self, entry: USEntry, key: DataKey, value: Any) -> None:
        raise NotImplementedError()

    @abc.abstractmethod
    def get_data_attr(self, entry: USEntry, key: DataKey) -> Any:
        raise NotImplementedError()

    def serialize(self, obj: USEntry) -> USDataPack:
        raw_dict = self.serialize_raw(obj)

        secret_keys = self.get_secret_keys(type(obj))

        raw_addressable = AddressableData(raw_dict)
        secrets_addressable = AddressableData({})

        for sec_key in secret_keys:
            sec_val = raw_addressable.pop(sec_key)
            secrets_addressable.set(sec_key, sec_val)

        return USDataPack(
            data=raw_addressable.data,
            secrets=secrets_addressable.data,
        )

    def deserialize(
        self,
        cls: Type[USEntry],
        data_pack: USDataPack,
        entry_id: str,
        us_manager: USManagerBase,
        common_properties: dict[str, Any],
        data_strict: bool = True,
    ) -> USEntry:
        # Assumed that data_pack's dicts was decoupled with initial US response
        data_pack = copy.deepcopy(data_pack)
        secret_source_addressable = AddressableData(data_pack.secrets)
        raw_addressable = AddressableData(data_pack.data)

        declared_secret_keys = self.get_secret_keys(cls)

        for secret_key in declared_secret_keys:
            if secret_source_addressable.contains(secret_key):
                sec_val = secret_source_addressable.pop(secret_key)
                sec_val_str = json.dumps(sec_val)
                raw_addressable.set(secret_key, sec_val_str)

        if secret_source_addressable.data:
            LOGGER.warning("Undeclared secrets found")

        return self.deserialize_raw(cls, raw_addressable.data, entry_id, us_manager, common_properties, data_strict)


class USEntrySerializerMarshmallow(USEntrySerializer):
    def serialize_raw(self, entry: USEntry) -> dict[str, Any]:
        dump_schema = self.get_dump_storage_schema(type(entry.data))
        data_dict = dump_schema.dump(entry.data)

        # Preventing saving what we can not load
        load_schema = self.get_load_storage_schema(type(entry.data))
        try:
            load_schema.load(data_dict)
            # TODO FIX: Validate that data is equal!!!
            # reloaded_data = load_schema.load(data_dict)
            # assert reloaded_data == entry.data
        except Exception as err:
            LOGGER.exception("Deserialization of US serialized data failed. Preventing save...")
            # TODO FIX: Custom exception
            raise exc.DLBaseException() from err

        return data_dict

    def deserialize_raw(
        self,
        cls: Type[USEntry],
        raw_data: dict[str, Any],
        entry_id: str,
        us_manager: USManagerBase,
        common_properties: dict[str, Any],
        data_strict: bool = True,
    ) -> USEntry:
        data_cls = cls.DataModel
        assert data_cls is not None and issubclass(data_cls, BaseAttrsDataModel)
        schema = self.get_load_storage_schema(data_cls)

        data = schema.load(raw_data)
        entry = cls(
            uuid=entry_id,
            us_manager=us_manager,
            data=data,
            data_strict=data_strict,
            **common_properties,
        )

        return entry

    def get_secret_keys(self, cls: Type[USEntry]) -> set[DataKey]:
        data_cls = cls.DataModel
        assert data_cls is not None and issubclass(data_cls, BaseAttrsDataModel)
        return data_cls.get_secret_keys()

    def set_data_attr(self, entry: USEntry, key: DataKey, value: Any) -> None:
        key_head = DataKey(parts=key.parts[:-1])
        obj = self.get_data_attr(entry, key_head)
        obj.__setattr__(key.parts[-1], value)

    def get_data_attr(self, entry: USEntry, key: DataKey) -> Any:
        return reduce(lambda obj, attribute: obj.__getattribute__(attribute), key.parts, entry.data)

lib/dl_core/dl_core/us_manager/us_entry_serializer.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import ChainMap
from contextlib import contextmanager
import json
import logging
from typing import (
    ClassVar,
    Dict,
    Optional,
    Set,
    Type,
    TypeVar,
    Union,
)
from typing import (
    TYPE_CHECKING,
    Any,
)
from typing import ChainMap as ChainMapGeneric

import attr
import marshmallow

from dl_api_commons.base_models import (
    RequestContextInfo,
    TenantDef,
)
from dl_configs.crypto_keys import (
    CryptoKeysConfig,
    get_crypto_keys_config_from_env,
)
from dl_constants.enums import ConnectionType
from dl_core import exc
from dl_core.base_models import (
    ConnectionRef,
    DefaultConnectionRef,
    EntryLocation,
    PathEntryLocation,
    WorkbookEntryLocation,
)
from dl_core.lifecycle.factory import DefaultEntryLifecycleManagerFactory
from dl_core.united_storage_client import (
    USAuthContextBase,
    UStorageClientBase,
)
from dl_core.us_connection import get_connection_class
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_dataset import Dataset
from dl_core.us_entry import (
    BaseAttrsDataModel,
    USEntry,
    USMigrationEntry,
)
from dl_core.us_manager.crypto.main import CryptoController
from dl_core.us_manager.local_cache import USEntryBuffer
from dl_core.us_manager.storage_schemas.connection_schema_registry import MAP_TYPE_TO_SCHEMA_MAP_TYPE_TO_SCHEMA
from dl_core.us_manager.storage_schemas.dataset import DatasetStorageSchema
from dl_core.us_manager.us_entry_serializer import (
    USDataPack,
    USEntrySerializer,
    USEntrySerializerMarshmallow,
)
from dl_core.us_manager.utils.fake_us_client import FakeUSClient
from dl_utils.utils import AddressableData


if TYPE_CHECKING:
    from dl_core.lifecycle.base import EntryLifecycleManager
    from dl_core.lifecycle.factory_base import EntryLifecycleManagerFactoryBase
    from dl_core.services_registry import ServicesRegistry


LOGGER = logging.getLogger(__name__)

_ENTRY_TV = TypeVar("_ENTRY_TV", bound=USEntry)


@attr.s(frozen=True)
class CryptoKeyInfo:
    key_id: Optional[str] = attr.ib()  # None value is None
    key_kind: Optional[str] = attr.ib()


class USManagerBase:
    _MAP_TYPE_TO_SCHEMA: ClassVar[ChainMapGeneric[Type[BaseAttrsDataModel], Type[marshmallow.Schema]]] = ChainMap(
        MAP_TYPE_TO_SCHEMA_MAP_TYPE_TO_SCHEMA,  # type: ignore  # TODO: fix
        {
            Dataset.DataModel: DatasetStorageSchema,
        },
    )

    _bi_context: RequestContextInfo

    _loaded_entries: USEntryBuffer

    _us_auth_context: USAuthContextBase
    _us_client: UStorageClientBase
    _fake_us_client: FakeUSClient
    _lifecycle_manager_factory: EntryLifecycleManagerFactoryBase

    def __init__(
        self,
        bi_context: RequestContextInfo,
        crypto_keys_config: Optional[CryptoKeysConfig],
        us_base_url: str,
        us_api_prefix: Optional[str],
        us_auth_context: USAuthContextBase,
        services_registry: ServicesRegistry,
        lifecycle_manager_factory: Optional[EntryLifecycleManagerFactoryBase] = None,
    ):
        # TODO FIX: Try to connect it together to eliminate possible divergence
        if services_registry is not None:
            try:
                conn_factory = services_registry.get_conn_executor_factory()
            except NotImplementedError:
                # Got some kind of dummy SR.
                # Skip this check because we'll get NIError anyway if we try ue it
                pass
            else:
                assert bi_context is conn_factory.req_ctx_info

        if crypto_keys_config is None:
            actual_crypto_keys_config = get_crypto_keys_config_from_env()
        else:
            actual_crypto_keys_config = crypto_keys_config

        self._bi_context = bi_context
        self._us_auth_context = us_auth_context

        self._crypto_keys_config = actual_crypto_keys_config
        self._us_base_url = us_base_url
        self._us_api_prefix = us_api_prefix

        self._crypto_controller = CryptoController(self._crypto_keys_config)
        self._loaded_entries = USEntryBuffer()
        # To prevent usage of flask context
        self._fake_us_client = FakeUSClient()
        self._services_registry: ServicesRegistry = services_registry

        lifecycle_manager_factory = lifecycle_manager_factory or DefaultEntryLifecycleManagerFactory()
        assert lifecycle_manager_factory is not None
        self._lifecycle_manager_factory = lifecycle_manager_factory

    def get_entry_buffer(self) -> USEntryBuffer:
        return self._loaded_entries

    @property
    def bi_context(self) -> RequestContextInfo:
        return self._bi_context

    @property
    def actual_crypto_key_id(self) -> str:
        return self._crypto_keys_config.actual_key_id

    def get_lifecycle_manager(
        self, entry: USEntry, service_registry: Optional[ServicesRegistry] = None
    ) -> EntryLifecycleManager:
        if service_registry is None:
            service_registry = self.get_services_registry()
        return self._lifecycle_manager_factory.get_lifecycle_manager(
            entry=entry, us_manager=self, service_registry=service_registry
        )

    # TODO FIX: Prevent saving entries with tenant ID that doesn't match current tenant ID
    def set_tenant_override(self, tenant: TenantDef) -> None:
        if not self._us_auth_context.is_tenant_id_mutable():
            raise AssertionError("Tenant ID might be set only for US manager with master US auth context")
        self._us_client.set_tenant_override(tenant)

    @classmethod
    def get_load_storage_schema(cls, data_cls: Type[BaseAttrsDataModel]) -> marshmallow.Schema:
        schema_cls = cls._MAP_TYPE_TO_SCHEMA.get(data_cls)
        return schema_cls()  # type: ignore  # TODO: fix

    @classmethod
    def get_dump_storage_schema(cls, data_cls: Type[BaseAttrsDataModel]) -> marshmallow.Schema:
        schema_cls = cls._MAP_TYPE_TO_SCHEMA.get(data_cls)
        return schema_cls()  # type: ignore  # TODO: fix

    @classmethod
    def get_us_entry_serializer(cls, entry_cls: Type[USEntry]) -> USEntrySerializer:
        data_cls = entry_cls.DataModel

        if data_cls is not None:
            if issubclass(data_cls, BaseAttrsDataModel):
                return USEntrySerializerMarshmallow()

        raise ValueError(f"Unknown type for US entry data: {entry_cls.DataModel}")

    # TODO FIX: Replace with on-import collecting of USEntry inheritors type/scope
    @classmethod
    def _get_entry_class(
        cls, *, us_scope: str, us_type: str, entry_key: Optional[EntryLocation] = None
    ) -> Type[USEntry]:
        err_msg = (
            f"Unknown combination of scope/type: {us_scope}/{us_type}"
            f" key={'not provided' if entry_key is None else entry_key.to_short_string()}"
        )

        try:
            if us_scope == "dataset":
                return Dataset

            elif us_scope == "connection":
                conn_type: ConnectionType
                try:
                    conn_type = ConnectionType(us_type)
                except ValueError:
                    conn_type = ConnectionType.unknown
                return get_connection_class(conn_type)

        except Exception as e:
            raise TypeError(err_msg) from e

        else:
            raise exc.UnexpectedUSEntryType()

    def get_sensitive_fields_key_info(self, entry: USEntry) -> Dict[str, CryptoKeyInfo]:
        if isinstance(entry, USMigrationEntry):
            # noinspection PyProtectedMember
            us_resp = entry._us_resp
            assert us_resp is not None
            entry_cls = self._get_entry_class(
                us_type=us_resp["type"],
                us_scope=us_resp["scope"],
                entry_key=entry.entry_key,
            )
            data_model_cls = entry_cls.DataModel
        else:
            entry_cls = type(entry)
            data_model_cls = entry.DataModel

        if data_model_cls is not None and issubclass(data_model_cls, BaseAttrsDataModel):
            serializer = self.get_us_entry_serializer(entry_cls)
            sensitive_fields = serializer.get_secret_keys(entry_cls)
            us_resp = entry._us_resp
            assert us_resp is not None
            fields_info_addressable = AddressableData({})
            unversioned_addressable = AddressableData(us_resp["unversionedData"])

            for field_key in sensitive_fields:
                if unversioned_addressable.contains(field_key):
                    field_value = unversioned_addressable.get(field_key)
                    if field_value is None:
                        key_id = None
                        key_kind = None
                    else:
                        key_id = field_value["key_id"]
                        key_kind = field_value["key_kind"]
                    fields_info_addressable.set(field_key, CryptoKeyInfo(key_id=key_id, key_kind=key_kind))

                else:
                    fields_info_addressable.set(field_key, None)
            return fields_info_addressable.data
        return {}

    def actualize_crypto_keys(self, entry: USMigrationEntry) -> None:
        us_resp = entry._us_resp
        assert us_resp is not None

        entry_cls = self._get_entry_class(
            us_type=us_resp["type"],
            us_scope=us_resp["scope"],
            entry_key=entry.entry_key,
        )

        serializer = self.get_us_entry_serializer(entry_cls)
        secret_keys = serializer.get_secret_keys(entry_cls)
        for key in secret_keys:
            secret_addressable = AddressableData(entry.unversioned_data)
            decrypted_value = self._crypto_controller.decrypt(secret_addressable.get(key))
            encrypted_value = self._crypto_controller.encrypt_with_actual_key(decrypted_value)
            secret_addressable.set(key, encrypted_value)

    def _entry_dict_to_obj(self, us_resp: dict, expected_type: Optional[Type[USEntry]] = None) -> USEntry:
        """
        Deserialize US entry dict.
        :param us_resp: US response as-is
        :param expected_type: If not none - type/scope will be checked to that actual class is subclass of expected one.
         If check fails or if there is no mapping for type/scope `TypeError` will be thrown.
         `USEntry` is special value. Actual class will not be determined and base `USEntry` will be constructed
        :return: Deserialized US entry
        """

        entry_loc: EntryLocation
        raw_entry_key = us_resp["key"]
        assert isinstance(raw_entry_key, str), f"Got non-string US entry key: {raw_entry_key!r}"

        if us_resp.get("workbookId") is not None:
            entry_name = raw_entry_key.split("/")[-1]
            entry_loc = WorkbookEntryLocation(
                workbook_id=us_resp["workbookId"],
                entry_name=entry_name,
            )
        else:
            entry_loc = PathEntryLocation(path=raw_entry_key)

        entry_cls: Type[USEntry]
        if expected_type == USMigrationEntry:
            entry_cls = USMigrationEntry
        else:
            entry_cls = self._get_entry_class(  # type: ignore  # TODO: fix
                us_type=us_resp["type"],
                us_scope=us_resp["scope"],
                entry_key=entry_loc,
            )
            if expected_type is not None and not issubclass(entry_cls, expected_type):
                raise exc.UnexpectedUSEntryType(
                    f"Unexpected type of US entry fetched:"
                    f" {entry_cls.__qualname__} instead of {expected_type.__qualname__}"
                )

        common_properties: dict[str, Any] = dict(
            entry_key=entry_loc,
            type_=us_resp["type"],
            meta=us_resp["meta"],
            is_locked=us_resp.get("isLocked"),
            is_favorite=us_resp.get("isFavorite"),
            permissions=us_resp.get("permissions") or {},
            links=us_resp.get("links") or {},
            hidden=us_resp["hidden"],
        )

        entry: USEntry
        if entry_cls == USMigrationEntry:
            entry = USMigrationEntry(
                uuid=us_resp["entryId"],
                us_manager=self,
                data=us_resp.get("data"),
                unversioned_data=us_resp.get("unversionedData"),
                data_strict=False,
                **common_properties,
            )
        else:
            data = us_resp.get("data")
            serializer = self.get_us_entry_serializer(entry_cls)
            data_pack = USDataPack(
                data=data,  # type: ignore
                secrets=us_resp.get("unversionedData"),  # type: ignore
            )

            entry = serializer.deserialize(
                entry_cls,
                data_pack,
                us_resp["entryId"],
                us_manager=self,
                common_properties=common_properties,
                data_strict=False,
            )
            secret_keys = serializer.get_secret_keys(entry_cls)
            for key in secret_keys:
                old_data = serializer.get_data_attr(entry, key)
                decrypted_data = self._crypto_controller.decrypt(json.loads(old_data)) if old_data is not None else None
                serializer.set_data_attr(entry, key, decrypted_data)

        entry._stored_in_db = True
        entry._us_resp = us_resp

        return entry

    def clone_entry_instance(self, entry: USEntry) -> USEntry:
        if not entry.uuid:
            raise ValueError("Only saved entries are supported here (must have an .uuid)")
        entry_data = self._get_entry_save_params(entry)
        entry_loc = entry.entry_key

        assert entry_loc is not None, "Could not clone US entry without location"

        entry_data.update(
            entryId=entry.uuid,
            unversionedData=entry_data.pop("unversioned_data"),
            **entry_loc.to_us_resp_api_params(entry.raw_us_key),
        )
        return self._entry_dict_to_obj(entry_data)

    @staticmethod
    def _get_us_type_for_entry(entry: USEntry) -> Optional[str]:
        if isinstance(entry, ConnectionBase):
            us_type = entry.conn_type.name
        else:
            us_type = entry.type_  # type: ignore  # TODO: fix

        if us_type is not None and not isinstance(us_type, str):
            raise ValueError(f"Unexpected US type ({us_type}) for entry {entry}")

        return us_type

    def dump_data(self, entry: USEntry) -> dict[str, Any]:
        entry_cls = type(entry)
        if entry_cls.DataModel is not None and issubclass(entry_cls.DataModel, BaseAttrsDataModel):
            serializer = self.get_us_entry_serializer(entry_cls)
            data_dict = serializer.serialize_raw(entry)
        else:
            data_dict = entry.data
        return data_dict

    # TODO CONSIDER: Use DTO instead of dict
    def _get_entry_save_params(self, entry: USEntry) -> dict:
        us_type = self._get_us_type_for_entry(entry)

        # Validating type on update
        if entry._us_resp:
            prev_type = entry._us_resp.get("type")
            if prev_type == "" and us_type is None:
                pass
            elif prev_type != us_type:
                raise ValueError(f"Type mismatch on US entry update: was '{prev_type}'; become '{us_type}'")

        data_dict = self.dump_data(entry)

        save_params: dict[str, Any] = {}

        if not entry._stored_in_db:
            if entry.permissions_mode is not None:
                save_params.update(permissionsMode=entry.permissions_mode)
            if entry.initial_permissions is not None:
                save_params.update(initialPermissions=entry.initial_permissions)

        if isinstance(entry, USMigrationEntry):
            save_params.update(
                data=data_dict,
                unversioned_data=entry.unversioned_data,
            )
        else:
            entry_cls = type(entry)
            if entry_cls.DataModel is not None:
                serializer = self.get_us_entry_serializer(type(entry))
                data_pack = serializer.serialize(entry)
            else:
                data_pack = USDataPack(data=data_dict)

            for key, secret in data_pack.secrets.items():
                assert secret is None or isinstance(secret, str)
                data_pack.secrets[key] = self._crypto_controller.encrypt_with_actual_key(secret)

            save_params.update(
                data=data_pack.data,
                unversioned_data=data_pack.secrets,
            )

        save_params.update(
            meta=entry.meta,
            hidden=entry.hidden,
            links=entry.links,
            scope=entry.scope,
            type=us_type,
        )

        return save_params

    def copy_entry(self, source: _ENTRY_TV, key: Optional[EntryLocation] = None) -> _ENTRY_TV:
        if not isinstance(source, Dataset):
            raise ValueError("Only dataset can be copied at this time")

        # TODO FIX: Find another way to copy data to avoid issues with unversioned data
        source_save_params = self._get_entry_save_params(source)

        copied_entry = type(source)(
            uuid=None,
            data=source_save_params["data"],
            entry_key=key,
            type_=source.type_,
            is_locked=source.is_locked,
            is_favorite=source.is_favorite,
            # initial_permissions=???,
            # Links will be collected on pre-save
            # links={},
            hidden=source.hidden,
            us_manager=self,
            data_strict=False,
        )

        self.get_lifecycle_manager(entry=copied_entry).post_copy_hook()

        return copied_entry  # type: ignore  # TODO: fix

    @staticmethod
    @contextmanager
    def _enrich_us_exception(  # type: ignore  # TODO: fix
        entry_id: Optional[str] = None,
        entry_scope: Optional[str] = None,
    ):
        try:
            yield
        except exc.USReqException as err:
            err.details.update(
                entry_id=entry_id,
                scope=entry_scope,
            )
            raise

    def get_loaded_us_connection(self, identity: Union[str, ConnectionRef]) -> ConnectionBase:
        if isinstance(identity, str):
            identity = DefaultConnectionRef(conn_id=identity)

        assert isinstance(identity, ConnectionRef)
        entry = self._loaded_entries.get_entry(identity)

        if not isinstance(entry, ConnectionBase):
            raise ValueError(f"Entry {identity} is not a connection: {type(identity).__qualname__}")

        return entry

    def _get_entry_links(self, entry: USEntry) -> Set[ConnectionRef]:
        if isinstance(entry, Dataset):
            lifecycle_manager = self.get_lifecycle_manager(entry=entry)
            linked_entries_refs: Set[ConnectionRef] = {
                DefaultConnectionRef(conn_id=conn_id) for conn_id in lifecycle_manager.collect_links().values()
            }

            return linked_entries_refs
        elif isinstance(entry, ConnectionBase):
            pass
        return set()

    def get_services_registry(self) -> ServicesRegistry:
        if self._services_registry is not None:
            return self._services_registry
        raise ValueError("Services registry was not passed to US manager")

lib/dl_core/dl_core/us_manager/us_manager.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import asynccontextmanager
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncGenerator,
    AsyncIterable,
    Iterable,
    Optional,
    Type,
    TypeVar,
    Union,
    overload,
)

from dl_api_commons.base_models import RequestContextInfo
from dl_app_tools.profiling_base import generic_profiler_async
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_core import exc
from dl_core.base_models import (
    ConnectionRef,
    DefaultConnectionRef,
)
from dl_core.united_storage_client import USAuthContextBase
from dl_core.united_storage_client_aio import UStorageClientAIO
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_dataset import Dataset
from dl_core.us_entry import USEntry
from dl_core.us_manager.broken_link import (
    BrokenUSLink,
    BrokenUSLinkErrorKind,
)
from dl_core.us_manager.us_manager import USManagerBase
from dl_utils.aio import shield_wait_for_complete


if TYPE_CHECKING:
    from dl_core.lifecycle.factory_base import EntryLifecycleManagerFactoryBase
    from dl_core.services_registry import ServicesRegistry


LOGGER = logging.getLogger(__name__)

_ENTRY_TV = TypeVar("_ENTRY_TV", bound=USEntry)


class AsyncUSManager(USManagerBase):
    _us_client: UStorageClientAIO

    def __init__(
        self,
        us_auth_context: USAuthContextBase,
        us_base_url: str,
        bi_context: RequestContextInfo,
        services_registry: ServicesRegistry,
        crypto_keys_config: Optional[CryptoKeysConfig] = None,
        us_api_prefix: Optional[str] = None,
        lifecycle_manager_factory: Optional[EntryLifecycleManagerFactoryBase] = None,
    ):
        self._us_client = UStorageClientAIO(
            host=us_base_url,
            prefix=us_api_prefix,
            auth_ctx=us_auth_context,
            context_request_id=bi_context.request_id if bi_context is not None else None,
            context_forwarded_for=bi_context.forwarder_for,
        )

        super().__init__(
            bi_context=bi_context,
            crypto_keys_config=crypto_keys_config,
            us_base_url=us_base_url,
            us_api_prefix=us_api_prefix,
            us_auth_context=us_auth_context,
            services_registry=services_registry,
            lifecycle_manager_factory=lifecycle_manager_factory,
        )

    @property
    def us_client(self) -> UStorageClientAIO:
        return self._us_client

    def clone(self) -> AsyncUSManager:
        return AsyncUSManager(
            crypto_keys_config=self._crypto_keys_config,
            us_auth_context=self._us_auth_context,
            us_base_url=self._us_base_url,
            us_api_prefix=self._us_api_prefix,
            bi_context=self._bi_context,
            services_registry=self._services_registry,
            lifecycle_manager_factory=self._lifecycle_manager_factory,
        )

    async def close(self) -> None:
        await self._us_client.close()

    async def __aenter__(self) -> AsyncUSManager:
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        try:
            await self.close()
        except Exception:  # noqa
            LOGGER.warning("Error during closing AsyncUSManager", exc_info=True)

    @overload
    async def get_by_id(self, entry_id: str, expected_type: type(None) = None) -> USEntry:  # type: ignore  # TODO: fix
        pass

    @overload  # noqa
    async def get_by_id(self, entry_id: str, expected_type: Type[_ENTRY_TV] = None) -> _ENTRY_TV:  # type: ignore  # TODO: fix
        pass

    @generic_profiler_async("us-fetch-entity")  # type: ignore  # TODO: fix
    async def get_by_id(self, entry_id: str, expected_type: Type[_ENTRY_TV] = None) -> _ENTRY_TV:  # type: ignore  # TODO: fix
        with self._enrich_us_exception(
            entry_id=entry_id,
            entry_scope=expected_type.scope if expected_type is not None else None,
        ):
            us_resp = await self._us_client.get_entry(entry_id)

        obj: _ENTRY_TV = self._entry_dict_to_obj(us_resp, expected_type)  # type: ignore  # TODO: fix
        await self.get_lifecycle_manager(entry=obj).post_init_async_hook()

        return obj

    async def save(self, entry: USEntry) -> None:
        self.get_lifecycle_manager(entry=entry).pre_save_hook()

        save_params = self._get_entry_save_params(entry)
        us_scope = save_params.pop("scope")
        us_type = save_params.pop("type")
        assert "data" in save_params and "unversioned_data" in save_params

        # noinspection PyProtectedMember
        if not entry._stored_in_db:
            entry_loc = entry.entry_key
            assert entry_loc is not None, "Entry location must be set before saving US entry"

            resp = await self._us_client.create_entry(
                entry_loc,
                scope=us_scope,
                type_=us_type,
                **save_params,
            )
            entry.uuid = resp["entryId"]
            entry._stored_in_db = True
        else:
            # noinspection PyProtectedMember
            resp = await self._us_client.update_entry(entry.uuid, lock=entry._lock, **save_params)

        entry._us_resp = resp

    async def delete(self, entry: USEntry) -> None:
        # TODO FIX: Use pre_delete_async_hook!!!
        self.get_lifecycle_manager(entry=entry).pre_delete_hook()

        if entry.uuid is None:
            raise ValueError("Entry has no id")

        # noinspection PyProtectedMember
        await self._us_client.delete_entry(entry.uuid, lock=entry._lock)
        entry._stored_in_db = False

        # noinspection PyBroadException
        try:
            # TODO FIX: Use post_delete_async_hook!!!
            self.get_lifecycle_manager(entry=entry).post_delete_hook()
        except Exception:
            LOGGER.exception("Error during post-delete hook execution for entry %s", entry.uuid)

    async def reload_data(self, entry: USEntry) -> None:
        assert entry.uuid is not None
        us_resp = await self._us_client.get_entry(entry.uuid)
        reloaded_entry = self._entry_dict_to_obj(us_resp, expected_type=type(entry))
        entry.data = reloaded_entry.data
        entry._us_resp = us_resp

    @asynccontextmanager  # type: ignore  # TODO: fix
    async def locked_entry_cm(
        self,
        entry_id: str,
        expected_type: Type[_ENTRY_TV],
        wait_timeout_sec: int = 30,
        duration_sec: int = 300,
        force: bool = False,
    ) -> AsyncGenerator[_ENTRY_TV, None]:
        entry: Optional[_ENTRY_TV] = None
        lock_token = await self._us_client.acquire_lock(
            entry_id,
            wait_timeout=wait_timeout_sec,
            duration=duration_sec,
            force=force,
        )

        try:
            entry = await self.get_by_id(entry_id, expected_type)  # type: ignore
            entry._lock = lock_token  # type: ignore  # TODO: fix
            assert entry is not None
            yield entry

        finally:
            await shield_wait_for_complete(self._us_client.release_lock(entry_id, lock_token))
            if entry is not None and hasattr(entry, "_lock"):
                entry._lock = None

    async def ensure_entry_preloaded(self, conn_ref: ConnectionRef) -> None:
        await self._ensure_conn_in_cache(None, conn_ref)

    async def _ensure_conn_in_cache(
        self, referrer: Optional[USEntry], conn_ref: ConnectionRef
    ) -> Optional[ConnectionBase]:
        conn: Union[USEntry, BrokenUSLink]
        if conn_ref in self._loaded_entries:
            conn = self._loaded_entries[conn_ref]
        else:
            try:
                assert isinstance(conn_ref, DefaultConnectionRef)
                conn = await self.get_by_id(conn_ref.conn_id, ConnectionBase)
            except (exc.USObjectNotFoundException, exc.USAccessDeniedException) as err:
                r_scope = referrer.scope if referrer is not None else "unk"
                r_uuid = referrer.uuid if referrer is not None else "unk"
                LOGGER.info(
                    "Can not load dependency for parent=%s/%s: connection %s (%s)",
                    r_scope,
                    r_uuid,
                    conn_ref,
                    err.message,
                )
                err_kind = {
                    exc.USObjectNotFoundException: BrokenUSLinkErrorKind.NOT_FOUND,
                    exc.USAccessDeniedException: BrokenUSLinkErrorKind.ACCESS_DENIED,
                }.get(type(err), BrokenUSLinkErrorKind.OTHER)
                conn = BrokenUSLink(reference=conn_ref, error_kind=err_kind)
            self._loaded_entries[conn_ref] = conn

        if isinstance(conn, BrokenUSLink):
            if referrer is not None:
                conn.add_referrer_id(referrer.uuid)  # type: ignore  # TODO: fix
            return None
        elif isinstance(conn, ConnectionBase):
            return conn
        else:
            raise ValueError("Entry was in cache but it is not a connection: %s", type(conn))

    # TODO FIX: Think about cache control
    async def load_dependencies(self, entry: USEntry) -> None:
        if not isinstance(entry, Dataset):
            raise NotImplementedError("Links loading is supported only for dataset")

        processed_refs: set[ConnectionRef] = set()
        # TODO FIX: Find a way to track direct source of link
        refs_to_load_queue = self._get_entry_links(
            entry,
        )

        while refs_to_load_queue:
            ref = refs_to_load_queue.pop()
            processed_refs.add(ref)
            try:
                resolved_ref = await self._ensure_conn_in_cache(referrer=entry, conn_ref=ref)
            except Exception:
                LOGGER.exception("Can not load linked US entry %s for entry %s", ref, entry.uuid)
                raise

            refs_to_load_queue.update(
                self._get_entry_links(
                    resolved_ref,  # type: ignore  # TODO: fix
                )
                - processed_refs
            )

    def get_raw_collection(
        self,
        entry_scope: str,
        entry_type: Optional[str] = None,
        all_tenants: bool = False,
    ) -> AsyncIterable[dict]:
        return self._us_client.entries_iterator(
            scope=entry_scope,
            entry_type=entry_type,
            all_tenants=all_tenants,
            include_data=False,
        )

    async def get_collection(
        self,
        entry_cls: Optional[Type[USEntry]],
        entry_type: Optional[str] = None,
        entry_scope: Optional[str] = None,
        meta: Optional[dict[str, Union[str, int, None]]] = None,
        all_tenants: bool = False,
        include_data: bool = True,
        ids: Optional[Iterable[str]] = None,
        creation_time: Optional[dict[str, Union[str, int, None]]] = None,
        raise_on_broken_entry: bool = False,
    ) -> AsyncGenerator[USEntry, None]:
        if all_tenants and include_data:
            raise ValueError("all_tenants and include_data cannot both be True")

        effective_scope = entry_scope
        if entry_cls and entry_cls.scope:
            if effective_scope:
                raise ValueError("US scope can not be provided in entry class and in parameters simultaneously")
            effective_scope = entry_cls.scope
        if effective_scope is None:
            raise ValueError("US scope not provided neither in entry class nor in parameters")

        effective_type = entry_type
        if not effective_type and entry_cls and entry_cls.type_:
            effective_type = entry_cls.type_
        assert isinstance(effective_scope, str)

        us_entry_iterator = self._us_client.entries_iterator(
            scope=effective_scope,
            entry_type=effective_type,
            meta=meta,
            all_tenants=all_tenants,
            include_data=include_data,
            ids=ids,
            creation_time=creation_time,
        )

        async for us_resp in us_entry_iterator:
            if us_resp:
                # noinspection PyBroadException
                try:
                    obj: USEntry = self._entry_dict_to_obj(us_resp, entry_cls)
                    yield obj
                except Exception:
                    LOGGER.exception("Failed to load US object: %s", us_resp)
                    if raise_on_broken_entry:
                        raise

lib/dl_core/dl_core/us_manager/us_manager_async.py
datalens-tech/datalens-backend
from cryptography import fernet

from dl_api_commons.base_models import RequestContextInfo
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_core.services_registry.top_level import (
    DummyServiceRegistry,
    ServicesRegistry,
)
from dl_core.united_storage_client import USAuthContextMaster
from dl_core.us_manager.us_manager import USManagerBase


class DummyUSManager(USManagerBase):
    def __init__(
        self,
        bi_context: RequestContextInfo = RequestContextInfo.create_empty(),
        services_registry: ServicesRegistry = DummyServiceRegistry(rci=RequestContextInfo.create_empty()),
    ):
        super().__init__(
            bi_context=bi_context,
            us_base_url="http://localhost:66000",
            us_api_prefix="dummy",
            crypto_keys_config=CryptoKeysConfig(  # type: ignore  # TODO: fix
                map_id_key={"dummy_usm_key": fernet.Fernet.generate_key().decode("ascii")},
                actual_key_id="dummy_usm_key",
            ),
            us_auth_context=USAuthContextMaster("FakeKey"),
            services_registry=services_registry,
        )

lib/dl_core/dl_core/us_manager/us_manager_dummy.py
datalens-tech/datalens-backend
from __future__ import annotations

import contextlib
import logging
from typing import (
    TYPE_CHECKING,
    Any,
    ContextManager,
    Dict,
    Generator,
    Iterable,
    List,
    Optional,
    Set,
    Type,
    TypeVar,
    Union,
    overload,
)

from dl_api_commons.base_models import RequestContextInfo
from dl_app_tools.profiling_base import generic_profiler
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_core import exc
from dl_core.base_models import (
    ConnectionRef,
    DefaultConnectionRef,
)
from dl_core.united_storage_client import (
    USAuthContextBase,
    UStorageClient,
)
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_dataset import Dataset
from dl_core.us_entry import USEntry
from dl_core.us_manager.broken_link import (
    BrokenUSLink,
    BrokenUSLinkErrorKind,
)
from dl_core.us_manager.us_manager import USManagerBase
from dl_utils.aio import await_sync


if TYPE_CHECKING:
    from dl_core.lifecycle.factory_base import EntryLifecycleManagerFactoryBase
    from dl_core.services_registry import ServicesRegistry


LOGGER = logging.getLogger(__name__)

_ENTRY_TV = TypeVar("_ENTRY_TV", bound=USEntry)


class SyncUSManager(USManagerBase):
    _us_client: UStorageClient

    def __init__(
        self,
        us_auth_context: USAuthContextBase,
        us_base_url: str,
        bi_context: RequestContextInfo,
        services_registry: ServicesRegistry,
        crypto_keys_config: Optional[CryptoKeysConfig] = None,
        us_api_prefix: Optional[str] = None,
        # caches_redis: Optional[aioredis.Redis] = None,
        request_timeout_sec: int = 30,  # WARNING: unused.
        lifecycle_manager_factory: Optional[EntryLifecycleManagerFactoryBase] = None,
    ):
        super().__init__(
            bi_context=bi_context,
            crypto_keys_config=crypto_keys_config,
            us_base_url=us_base_url,
            us_api_prefix=us_api_prefix,
            us_auth_context=us_auth_context,
            services_registry=services_registry,
            lifecycle_manager_factory=lifecycle_manager_factory,
        )
        self._us_client = self._create_us_client()

    def _create_us_client(self) -> UStorageClient:
        return UStorageClient(
            host=self._us_base_url,
            prefix=self._us_api_prefix,
            auth_ctx=self._us_auth_context,
            context_request_id=self._bi_context.request_id if self._bi_context is not None else None,
            context_forwarded_for=self._bi_context.forwarder_for,
        )

    def clone(self, **kwargs):  # type: ignore  # TODO: fix
        """This should've been an `attr.evolve` wrapper"""
        base_kwargs = dict(
            us_auth_context=self._us_auth_context,
            us_base_url=self._us_base_url,
            bi_context=self._bi_context,
            crypto_keys_config=self._crypto_keys_config,
            us_api_prefix=self._us_api_prefix,
            # request_timeout_sec=self._request_timeout_sec,
            services_registry=self._services_registry,
            lifecycle_manager_factory=self._lifecycle_manager_factory,
        )
        kwargs = {
            **base_kwargs,
            **kwargs,
        }
        return self.__class__(**kwargs)

    def close(self) -> None:
        self._us_client.close()

    def __enter__(self) -> SyncUSManager:
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        try:
            self.close()
        except Exception:
            LOGGER.warning("Error during closing SyncUSManager", exc_info=True)

    # CRUD
    #
    def save(self, entry: USEntry) -> None:
        lifecycle_manager = self.get_lifecycle_manager(entry=entry, service_registry=self._services_registry)
        lifecycle_manager.pre_save_hook()

        save_params = self._get_entry_save_params(entry)
        us_scope = save_params.pop("scope")
        us_type = save_params.pop("type")
        entry_loc = entry.entry_key
        assert "data" in save_params and "unversioned_data" in save_params
        assert entry_loc is not None, "Can not save entry without key/workbook data"

        # noinspection PyProtectedMember
        if not entry._stored_in_db:
            resp = self._us_client.create_entry(
                key=entry_loc,
                scope=us_scope,
                type_=us_type,
                **save_params,
            )
            entry.uuid = resp["entryId"]
            entry._stored_in_db = True
        else:
            # noinspection PyProtectedMember
            resp = self._us_client.update_entry(
                entry.uuid, lock=entry._lock, **save_params  # type: ignore  # TODO: fix
            )

        entry._us_resp = resp  # type: ignore  # TODO: fix
        lifecycle_manager.post_save_hook()

    def delete(self, entry: USEntry) -> None:
        lifecycle_manager = self.get_lifecycle_manager(entry=entry)
        lifecycle_manager.pre_delete_hook()

        self._us_client.delete_entry(entry.uuid, lock=entry._lock)
        entry._stored_in_db = False
        # noinspection PyBroadException
        try:
            LOGGER.info("Executing post-delete hook %s", entry.uuid)
            lifecycle_manager.post_delete_hook()
        except Exception:
            LOGGER.exception("Error during post-delete hook execution for entry %s", entry.uuid)

    @overload
    def get_by_id(self, entry_id: str, expected_type: type(None) = None) -> USEntry:  # type: ignore  # TODO: fix
        pass

    @overload  # noqa
    def get_by_id(self, entry_id: str, expected_type: Type[_ENTRY_TV] = None) -> _ENTRY_TV:  # type: ignore  # TODO: fix
        pass

    @generic_profiler("us-fetch-entity")
    def get_by_id(self, entry_id: str, expected_type: Optional[Type[USEntry]] = None) -> USEntry:
        with self._enrich_us_exception(
            entry_id=entry_id,
            entry_scope=expected_type.scope if expected_type is not None else None,
        ):
            us_resp = self._us_client.get_entry(entry_id)
        obj = self._entry_dict_to_obj(us_resp, expected_type)
        await_sync(self.get_lifecycle_manager(entry=obj).post_init_async_hook())
        return obj

    @overload
    def get_by_key(self, entry_key: str, expected_type: type(None) = None) -> USEntry:  # type: ignore  # TODO: fix
        pass

    @overload  # noqa
    def get_by_key(self, entry_key: str, expected_type: Type[_ENTRY_TV] = None) -> _ENTRY_TV:  # type: ignore  # TODO: fix
        pass

    def get_by_key(self, entry_key: str, expected_type=None):  # type: ignore  # TODO: fix  # noqa
        raise NotImplementedError()

    def get_collection(
        self,
        entry_cls: Optional[Type[_ENTRY_TV]],
        entry_type: Optional[str] = None,
        entry_scope: Optional[str] = None,
        meta: Optional[Dict[str, Union[str, int, None]]] = None,
        all_tenants: bool = False,
        include_data: bool = True,
        ids: Optional[Iterable[str]] = None,
        creation_time: Optional[Dict[str, Union[str, int, None]]] = None,
        raise_on_broken_entry: bool = False,
    ) -> Generator[_ENTRY_TV, None, None]:
        if all_tenants and include_data:
            # Not supported; see _req_data_iter_entries
            raise ValueError("all_tenants and include_data cannot both be True")
        if entry_scope and entry_cls.scope:  # type: ignore  # TODO: fix
            raise ValueError("US scope can not be provided in entry class and in parameters simultaneously")
        if not entry_scope and not entry_cls.scope:  # type: ignore  # TODO: fix
            raise ValueError("US scope not provided neither in entry class nor in parameters")

        effective_scope = entry_scope or entry_cls.scope  # type: ignore  # TODO: fix
        effective_type = entry_type or entry_cls.type_  # type: ignore  # TODO: fix
        assert isinstance(effective_scope, str)

        us_entry_iterator = self._us_client.entries_iterator(
            scope=effective_scope,
            entry_type=effective_type,
            meta=meta,
            all_tenants=all_tenants,
            include_data=include_data,
            ids=ids,
            creation_time=creation_time,
        )

        for us_resp in us_entry_iterator:
            # noinspection PyBroadException
            try:
                yield self._entry_dict_to_obj(us_resp, expected_type=entry_cls)  # type: ignore  # TODO: fix
            except Exception:
                LOGGER.exception("Failed to load US object: %s", us_resp)
                if raise_on_broken_entry:
                    raise

    def get_raw_entry(self, entry_id: str, include_permissions: bool = True, include_links: bool = True) -> dict:
        return self._us_client.get_entry(entry_id, include_permissions=include_permissions, include_links=include_links)

    def get_raw_collection(
        self,
        entry_scope: str,
        entry_type: Optional[str] = None,
        all_tenants: bool = False,
    ) -> Iterable[dict]:
        return self._us_client.entries_iterator(
            scope=entry_scope,
            entry_type=entry_type,
            all_tenants=all_tenants,
            include_data=False,
        )

    def move(self, entry: USEntry, destination: str) -> None:
        self._us_client.move_entry(entry.uuid, destination)
        self.reload_data(entry)  # maybe we have to update only entry_key

    def reload_data(self, entry: USEntry) -> None:
        assert entry.uuid is not None
        us_resp = self._us_client.get_entry(entry.uuid)
        reloaded_entry = self._entry_dict_to_obj(us_resp, expected_type=type(entry))
        entry.data = reloaded_entry.data
        entry._us_resp = us_resp

    # Locks
    #
    def acquire_lock(self, entry: USEntry, duration: Optional[int] = None, wait_timeout: Optional[int] = None) -> str:
        lock_token = self._us_client.acquire_lock(entry.uuid, duration, wait_timeout)
        entry._lock = lock_token
        return lock_token

    def release_lock(self, entry: USEntry) -> None:
        # noinspection PyProtectedMember
        if entry._lock:
            self._us_client.release_lock(entry.uuid, entry._lock)
            entry._lock = None

    @contextlib.contextmanager  # type: ignore  # TODO: fix
    def locked_cm(  # type: ignore  # TODO: fix
        self, entry: USEntry, duration: Optional[int] = None, wait_timeout: Optional[int] = None
    ) -> ContextManager:
        self.acquire_lock(entry, duration=duration, wait_timeout=wait_timeout)
        try:
            yield
        finally:
            self.release_lock(entry)

    @contextlib.contextmanager  # type: ignore  # TODO: fix
    def get_locked_entry_cm(  # type: ignore  # TODO: fix
        self,
        expected_type: Type[_ENTRY_TV],
        entry_id: str,
        duration: Optional[int] = None,
        wait_timeout: Optional[int] = None,
    ) -> ContextManager[_ENTRY_TV]:
        lock_token = self._us_client.acquire_lock(entry_id, duration, wait_timeout)
        entry = None
        try:
            entry = self.get_by_id(entry_id, expected_type=expected_type)
            entry._lock = lock_token
            yield entry
        finally:
            self._us_client.release_lock(entry_id, lock_token)
            if entry is not None:
                entry._lock = None

    # Dependencies
    #
    def get_loaded_us_connection(self, identity: Union[str, ConnectionRef]) -> ConnectionBase:
        # Temporary workaround to mitigate forgotten .load_dependencies()
        if isinstance(identity, str):
            identity = DefaultConnectionRef(conn_id=identity)

        if identity not in self._loaded_entries:
            # TODO FIX: Uncomment after migration to dataset v7
            # _log.warning("SyncUSManager.get_loaded_us_connection() called without connection in cache")
            self._ensure_conn_in_cache(None, identity)
        return super().get_loaded_us_connection(identity)

    def ensure_entry_preloaded(self, conn_ref: ConnectionRef) -> None:
        self._ensure_conn_in_cache(None, conn_ref)

    def _ensure_conn_in_cache(self, referrer: Optional[USEntry], conn_ref: ConnectionRef) -> Optional[ConnectionBase]:
        conn: Union[USEntry, BrokenUSLink]
        if conn_ref in self._loaded_entries:
            conn = self._loaded_entries[conn_ref]
        else:
            try:
                assert isinstance(conn_ref, DefaultConnectionRef)
                conn = self.get_by_id(conn_ref.conn_id, ConnectionBase)
            except (exc.USObjectNotFoundException, exc.USAccessDeniedException) as err:
                r_scope = referrer.scope if referrer is not None else "unk"
                r_uuid = referrer.uuid if referrer is not None else "unk"
                LOGGER.info(
                    "Can not load dependency for parent=%s/%s: connection %s (%s)",
                    r_scope,
                    r_uuid,
                    conn_ref,
                    err.message,
                )
                err_kind = {
                    exc.USObjectNotFoundException: BrokenUSLinkErrorKind.NOT_FOUND,
                    exc.USAccessDeniedException: BrokenUSLinkErrorKind.ACCESS_DENIED,
                }.get(type(err), BrokenUSLinkErrorKind.OTHER)
                conn = BrokenUSLink(reference=conn_ref, error_kind=err_kind)
            self._loaded_entries[conn_ref] = conn

        if isinstance(conn, BrokenUSLink):
            if referrer is not None:
                conn.add_referrer_id(referrer.uuid)  # type: ignore  # TODO: fix
            return None
        elif isinstance(conn, ConnectionBase):
            return conn
        else:
            raise ValueError("Entry was in cache but it is not a connection: %s", type(conn))

    # TODO FIX: Think about cache control
    def load_dependencies(self, entry: USEntry) -> None:
        if not isinstance(entry, Dataset):
            raise NotImplementedError("Links loading is supported only for dataset")

        processed_refs: Set[ConnectionRef] = set()
        # TODO FIX: Find a way to track direct source of link
        refs_to_load_queue = self._get_entry_links(
            entry,
        )

        while refs_to_load_queue:
            ref = refs_to_load_queue.pop()
            processed_refs.add(ref)
            try:
                resolved_ref = self._ensure_conn_in_cache(referrer=entry, conn_ref=ref)
            except Exception:
                LOGGER.exception("Can not load linked US entry %s for entry %s", ref, entry.uuid)
                raise

            refs_to_load_queue.update(
                self._get_entry_links(
                    resolved_ref,  # type: ignore  # TODO: fix
                )
                - processed_refs
            )

    def load_get_entries_at_path(self, us_path: str) -> list[dict[str, Any]]:
        return self._us_client.get_entries_info_in_path(us_path)

lib/dl_core/dl_core/us_manager/us_manager_sync.py
datalens-tech/datalens-backend
from __future__ import annotations

import copy
from datetime import (
    datetime,
    timezone,
)
from typing import (
    Any,
    Dict,
    Optional,
)

from cryptography import fernet
import shortuuid

from dl_api_commons.base_models import RequestContextInfo
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_core.base_models import EntryLocation
from dl_core.exc import USObjectNotFoundException
from dl_core.services_registry.top_level import (
    DummyServiceRegistry,
    ServicesRegistry,
)
from dl_core.united_storage_client import (
    USAuthContextBase,
    USAuthContextMaster,
    UStorageClient,
    UStorageClientBase,
)
from dl_core.us_manager.us_manager_sync import SyncUSManager


class MockedUStorageClient(UStorageClient):
    def __init__(
        self,
        host: str,
        auth_ctx: USAuthContextBase,
        prefix: Optional[str] = None,
        timeout: int = 30,
        context_request_id: Optional[str] = None,
    ):
        super().__init__(host, auth_ctx, prefix, timeout, context_request_id)
        self._saved_entries: Dict[str, Dict[str, Any]] = {}

    @classmethod
    def format_dt(cls, dt: datetime) -> str:
        assert dt.tzinfo
        pre_formatted = (
            dt.replace(microsecond=dt.microsecond // 1000 * 1000)
            .astimezone(timezone.utc)
            .strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]
        )
        return f"{pre_formatted}Z"

    def _request(self, request_data: UStorageClientBase.RequestData) -> Dict[str, Any]:
        raise NotImplementedError("This is dummy US client")

    def create_entry(
        self,
        key: EntryLocation,
        scope: str,
        meta: Optional[Dict[str, str]] = None,
        data: Optional[Dict[str, Any]] = None,
        unversioned_data: Optional[Dict[str, Any]] = None,
        type_: Optional[str] = None,
        hidden: Optional[bool] = None,
        links: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Dict[str, Any]:
        assert not kwargs, "Not supported by dummy"
        entry_id = shortuuid.uuid()
        rev_id = shortuuid.uuid()
        subject = "unknown"
        created_at = self.format_dt(datetime.now(timezone.utc))

        resp = dict(
            entryId=entry_id,
            scope=scope,
            type=type_,
            unversionedData=copy.deepcopy(unversioned_data),
            createdBy=subject,
            createdAt=created_at,
            updatedBy=subject,
            updatedAt=created_at,
            savedId=rev_id,
            # publishedId="???",
            revId=rev_id,
            tenantId="common",
            data=copy.deepcopy(data),
            meta=copy.deepcopy(meta),
            hidden=hidden,
            # mirrored=???,
            # public=???,
            links=links,
            **key.to_us_resp_api_params(None),
        )

        self._saved_entries[entry_id] = resp
        return copy.deepcopy(resp)

    def update_entry(
        self,
        entry_id: str,
        data: Optional[Dict[str, Any]] = None,
        unversioned_data: Optional[Dict[str, Any]] = None,
        meta: Optional[Dict[str, str]] = None,
        lock: Optional[str] = None,
        hidden: Optional[bool] = None,
        links: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        previous_resp = self._saved_entries[entry_id]

        new_revision_id = shortuuid.uuid()
        new_updated_at = self.format_dt(datetime.now(timezone.utc))
        subject = "unknown"

        previous_resp.update(
            data=copy.deepcopy(data),
            unversionedData=copy.deepcopy(unversioned_data),
            meta=copy.deepcopy(meta),
            hidden=hidden,
            revId=new_revision_id,
            savedId=new_revision_id,
            updatedBy=subject,
            updatedAt=new_updated_at,
        )

        return copy.deepcopy(previous_resp)

    def get_entry(
        self,
        entry_id: str,
        params: Optional[Dict[str, Any]] = None,
        include_permissions: bool = True,
        include_links: bool = True,
    ) -> Dict[str, Any]:
        assert params is None
        try:
            previous_resp = self._saved_entries[entry_id]
        except KeyError:
            raise USObjectNotFoundException()
        else:
            return previous_resp


class MockedSyncUSManager(SyncUSManager):
    def __init__(
        self,
        bi_context: RequestContextInfo = RequestContextInfo.create_empty(),
        crypto_keys_config: Optional[CryptoKeysConfig] = None,
        services_registry: ServicesRegistry = DummyServiceRegistry(rci=RequestContextInfo.create_empty()),
    ):
        super().__init__(
            bi_context=bi_context,
            us_base_url="http://localhost:66000",
            us_api_prefix="dummy",
            crypto_keys_config=CryptoKeysConfig(
                map_id_key={"dummy_usm_key": fernet.Fernet.generate_key()},  # type: ignore
                actual_key_id="dummy_usm_key",
            )
            if crypto_keys_config is None
            else crypto_keys_config,
            us_auth_context=USAuthContextMaster("FakeKey"),
            services_registry=services_registry,
        )

    def _create_us_client(self) -> UStorageClient:
        return MockedUStorageClient(
            host=self._us_base_url,
            auth_ctx=self._us_auth_context,
            prefix=self._us_api_prefix,
        )

lib/dl_core/dl_core/us_manager/us_manager_sync_mock.py
datalens-tech/datalens-backend
from __future__ import annotations

import collections
from typing import Type


def inheritors(klass: Type):  # type: ignore  # TODO: fix
    subclasses = set()
    work = collections.deque([klass])
    while work:
        parent = work.pop()
        for child in parent.__subclasses__():
            if child not in subclasses:
                subclasses.add(child)
                work.append(child)
    return subclasses

lib/dl_core/dl_core/us_manager/utils/common.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_core.united_storage_client import (
    USAuthContextNoAuth,
    UStorageClient,
    UStorageClientBase,
)


# TODO FIX: Remove it
class FakeUSClient(UStorageClient):
    """
    This class is used by async US manager to prevent saving/updating entities with their BaseModel.*() methods
    """

    def __init__(self):  # type: ignore  # TODO: fix
        super().__init__(host="http://127.0.0.1:3030", auth_ctx=USAuthContextNoAuth())

    def _request(self, request_data: UStorageClientBase.RequestData):  # type: ignore  # TODO: fix
        raise NotImplementedError("US entries created by async manager can not communicate with US directly")

lib/dl_core/dl_core/us_manager/utils/fake_us_client.py
datalens-tech/datalens-backend
from __future__ import annotations

import ipaddress
import logging
import os
import re
from typing import (
    Any,
    Generic,
    Iterable,
    Optional,
    Type,
    TypeVar,
)
from urllib.parse import urlparse
import uuid

from aiohttp import (
    ClientResponse,
    ClientResponseError,
    RequestInfo,
)
import attr
import jaeger_client
from multidict import (
    CIMultiDict,
    CIMultiDictProxy,
)
import opentracing
import requests
import requests.adapters
import requests.exceptions
import shortuuid
import sqlalchemy as sa

from dl_api_commons import clean_secret_data_in_headers
from dl_api_commons.base_models import (
    AuthData,
    RequestContextInfo,
)
from dl_api_commons.utils import (
    stringify_dl_cookies,
    stringify_dl_headers,
)
from dl_configs.settings_loaders.env_remap import remap_env
from dl_constants.api_constants import DLHeadersCommon


LOGGER = logging.getLogger(__name__)


def make_url(
    protocol: str,
    host: str,
    port: int,
    path: Optional[str] = None,
) -> str:
    # TODO FIX: Sanitize/use urllib
    if path is None:
        path = ""
    return f"{protocol}://{host}:{port}/{path.lstrip('/')}"


def get_requests_session() -> requests.Session:
    session = requests.Session()
    ua = "{}, Datalens".format(requests.utils.default_user_agent())
    session.headers.update({"User-Agent": ua})
    return session


def get_retriable_requests_session() -> requests.Session:
    session = get_requests_session()

    retry_conf = requests.adapters.Retry(
        total=5,
        backoff_factor=0.5,
        status_forcelist=[500, 501, 502, 503, 504, 521],
        redirect=10,
        method_whitelist=frozenset(["HEAD", "TRACE", "GET", "PUT", "OPTIONS", "DELETE", "POST"]),
        # # TODO:
        # # (the good: will return a response when it's an error response)
        # # (the bad: need to raise_for_status() manually, same as without retry conf)
        # raise_on_status=False,
    )

    for schema in ("http://", "https://"):
        session.mount(
            schema,
            # noinspection PyUnresolvedReferences
            requests.adapters.HTTPAdapter(
                max_retries=retry_conf,
            ),
        )

    return session


def make_user_auth_headers(
    rci: RequestContextInfo,
    auth_data_override: Optional[AuthData] = None,
    sudo_override: Optional[bool] = None,
    allow_superuser_override: Optional[bool] = None,
) -> dict[str, str]:
    headers: dict[str, str] = {}

    effective_auth_data: Optional[AuthData] = rci.auth_data if auth_data_override is None else auth_data_override
    if effective_auth_data is not None:
        headers.update(stringify_dl_headers(effective_auth_data.get_headers()))

    req_id = rci.request_id
    if req_id is not None:
        headers.update({DLHeadersCommon.REQUEST_ID.value: req_id})

    tenant = rci.tenant
    assert tenant is not None
    headers.update(stringify_dl_headers(tenant.get_outbound_tenancy_headers()))

    def may_be_add_extra_header(header: DLHeadersCommon, override_value: Optional[Any]) -> None:
        incoming_value = rci.plain_headers.get(header.value)
        if override_value is not None:
            if isinstance(override_value, str):
                headers[header.value] = override_value
            if isinstance(override_value, bool):
                if override_value:
                    headers[header.value] = "1"
            else:
                raise AssertionError(f"Unsupported type of header override: {type(override_value)}")
        else:
            if incoming_value is not None:
                headers[header.value] = incoming_value

    may_be_add_extra_header(DLHeadersCommon.SUDO, sudo_override)
    may_be_add_extra_header(DLHeadersCommon.ALLOW_SUPERUSER, allow_superuser_override)
    may_be_add_extra_header(DLHeadersCommon.DL_CONTEXT, None)

    return headers


def make_user_auth_cookies(
    rci: RequestContextInfo,
    auth_data_override: Optional[AuthData] = None,
) -> dict[str, str]:
    effective_auth_data: Optional[AuthData] = rci.auth_data if auth_data_override is None else auth_data_override
    if effective_auth_data is not None:
        return stringify_dl_cookies(effective_auth_data.get_cookies())
    return {}


# TODO FIX: Remove after migration to Connection Executors
def compile_query_for_debug(query, dialect):  # type: ignore  # TODO: fix
    """
    Compile query to string.
    This function is only suitable for logging and not execution of the result.
    Its result might not be valid SQL if query contains date/datetime literals.
    """
    try:
        try:
            return str(query.compile(dialect=dialect, compile_kwargs={"literal_binds": True}))
        except NotImplementedError:
            compiled = query.compile(dialect=dialect)
            return "{0}; {1!r}".format(str(query), compiled.params)
    except Exception:
        return "-"


def parse_comma_separated_hosts(host: Optional[str]) -> tuple[str, ...]:
    if not host:
        return tuple()
    return tuple(h.strip() for h in host.split(","))


def validate_hostname_or_ip_address(hostname: str):  # type: ignore  # TODO: fix
    # IP address case
    try:
        ipaddress.ip_address(hostname)
    except ValueError:
        pass  # Not a valid IP address, continue to check if valid hostname
    else:
        return

    # DNS hostname case
    if hostname[-1] == ".":
        hostname = hostname[:-1]
    if len(hostname) > 253:
        raise ValueError("Hostname is too long")

    label_list = hostname.split(".")

    if re.match(r"[0-9]+$", label_list[-1]):
        raise ValueError("TLD must be not all-numeric")

    label_re = re.compile(r"(?!-)[a-z0-9-]{1,63}(?<!-)$", re.IGNORECASE)

    if not all(label_re.match(label) for label in label_list):
        raise ValueError("Not a valid domain name")

    # Ensure valid netloc (from previous hostname validation mechanism)
    parsed_url = urlparse("//{}".format(hostname))
    if parsed_url.netloc != hostname:
        raise ValueError("Not a valid netloc")


def shorten_uuid(some_uuid: str):  # type: ignore  # TODO: fix
    return shortuuid.encode(uuid.UUID(some_uuid))


_FUTURE_REF_TV = TypeVar("_FUTURE_REF_TV")


@attr.s(eq=False, hash=False, order=False)
class FutureRef(Generic[_FUTURE_REF_TV]):
    __ref: Optional[_FUTURE_REF_TV] = attr.ib(init=False, default=None)

    @property
    def ref(self) -> "_FUTURE_REF_TV":
        if self.__ref is None:
            raise ValueError("FutureRef was not fulfilled")
        return self.__ref

    def fulfill(self, obj: _FUTURE_REF_TV) -> None:
        if self.__ref is not None:
            raise ValueError("FutureRef already fulfilled")
        self.__ref = obj

    @classmethod
    def fulfilled(cls, obj: _FUTURE_REF_TV) -> "FutureRef[_FUTURE_REF_TV]":
        fr = cls()
        fr.fulfill(obj)
        return fr


def make_id() -> str:
    return shortuuid.uuid()


_MODEL_TYPE_TV = TypeVar("_MODEL_TYPE_TV", bound=attr.AttrsInstance)


def attrs_evolve_to_subclass(cls: Type[_MODEL_TYPE_TV], inst: Any, **kwargs) -> _MODEL_TYPE_TV:  # type: ignore  # TODO: fix
    """
    Evolve an attr.s instance to a subclass instance with additional attributes.
    Note that this works correctly only for attributes with ``init=True``.
    """

    super_cls = type(inst)
    if super_cls is cls:
        if kwargs:
            return attr.evolve(inst, **kwargs)
        else:
            return inst
    assert issubclass(cls, super_cls), f"Expected subclass of {super_cls.__name__}, got {cls.__name__}"
    all_attrs = {f.name: getattr(inst, f.name.lstrip("_")) for f in attr.fields(super_cls) if f.init}
    all_attrs.update(kwargs)
    return cls(**all_attrs)  # type: ignore  # TODO: fix


def attrs_evolve_to_superclass(cls: Type[_MODEL_TYPE_TV], inst: Any, **kwargs) -> _MODEL_TYPE_TV:  # type: ignore  # TODO: fix
    """
    Evolve an attr.s instance to a superclass instance with additional attributes.
    Note that this works correctly only for attributes with ``init=True``.
    """

    sub_cls = type(inst)
    if sub_cls is cls:
        if kwargs:
            return attr.evolve(inst, **kwargs)
        else:
            return inst
    assert issubclass(sub_cls, cls), f"Expected superclass of {sub_cls.__name__}, got {cls.__name__}"
    all_attrs = {f.name: getattr(inst, f.name.lstrip("_")) for f in attr.fields(cls) if f.init}
    all_attrs.update(kwargs)
    return cls(**all_attrs)  # type: ignore  # TODO: fix


def get_current_w3c_tracing_headers(
    tracer: Optional[opentracing.Tracer] = None,
    req_id: Optional[str] = None,
    override_sampled: Optional[bool] = None,
) -> dict[str, str]:
    actual_tracer = opentracing.global_tracer() if tracer is None else tracer
    tracing_headers: dict[str, str] = {}

    active_span = actual_tracer.active_span

    if active_span is not None:
        span_ctx = active_span.context

        if isinstance(span_ctx, jaeger_client.SpanContext):
            if override_sampled is None:
                flags = "01" if span_ctx.flags & 0x1 else "00"
            else:
                flags = "01" if override_sampled else "00"

            header_value = f"00-{span_ctx.trace_id:032x}-{span_ctx.span_id:016x}-{flags}"
            LOGGER.info("Traceparent header was generated for req_id %r: %r", req_id, header_value)
            tracing_headers.update(
                traceparent=header_value,
            )
        else:
            LOGGER.debug("Could not create W3C tracing headers (span is not Jaeger span) for req id: %s", req_id)

    else:
        LOGGER.debug("Could not create W3C tracing headers (no active span) for req id: %s", req_id)

    return tracing_headers


def generate_revision_id() -> str:
    return shortuuid.uuid()


def sa_plain_text(query: str) -> sa.sql.elements.TextClause:
    """sa.text without SA semantics (':'-parameters)"""
    # query = query.replace('\\', r'\\')  # weirdly, seems unnecessary
    query = query.replace(":", r"\:")
    return sa.text(query)


SECREPR_SIDE_SIZE = 0


def secrepr(value: Optional[str]) -> str:
    """Convenience function for attrs-repr of secrets"""
    if value is None:
        return repr(value)
    if not value:
        return repr(value)
    if not isinstance(value, str):
        # Really shouldn't raise in a `repr`
        return f"???{type(value)!r}???"

    side_size = SECREPR_SIDE_SIZE  # not a kwarg just to keep mypy happier
    if not side_size or len(value) <= side_size * 3:
        return "..."
    return repr(f"{value[:side_size]}...{value[-side_size:]}")


def _multidict_to_list(md: CIMultiDictProxy[str]) -> Iterable[tuple[str, str]]:
    return [(str(k), str(v)) for k, v in md.items()]


def raise_for_status_and_hide_secret_headers(response: ClientResponse) -> None:
    """
    Overriden function from aiohttp which raises ClientResponseError is response result is not ok,
    but hides secret data in headers for exception
    """
    if not response.ok:
        # reason should always be not None for a started response
        assert response.reason is not None
        response.release()
        new_request_info = RequestInfo(
            response.request_info.url,
            response.request_info.method,
            CIMultiDict(clean_secret_data_in_headers(_multidict_to_list(response.request_info.headers))),
        )
        raise ClientResponseError(
            new_request_info,
            response.history,
            status=response.status,
            message=response.reason,
            headers=CIMultiDict(clean_secret_data_in_headers(_multidict_to_list(response.headers))),
        )

lib/dl_core/dl_core/utils.py
datalens-tech/datalens-backend
from __future__ import annotations

from datetime import (
    date,
    datetime,
)
from typing import (
    ClassVar,
    Generic,
    List,
    TypeVar,
    Union,
)

import attr

from dl_constants.enums import UserDataType


_INNER_TYPE = TypeVar("_INNER_TYPE")


@attr.s(frozen=True)
class BIValue(Generic[_INNER_TYPE]):
    type: ClassVar[UserDataType]
    value: _INNER_TYPE = attr.ib()


@attr.s(frozen=True)
class StringValue(BIValue[str]):
    type: ClassVar[UserDataType] = UserDataType.string


@attr.s(frozen=True)
class IntegerValue(BIValue[int]):
    type: ClassVar[UserDataType] = UserDataType.integer


@attr.s(frozen=True)
class FloatValue(BIValue[float]):
    type: ClassVar[UserDataType] = UserDataType.float


@attr.s(frozen=True)
class DateValue(BIValue[date]):
    type: ClassVar[UserDataType] = UserDataType.date


@attr.s(frozen=True)
class DateTimeValue(BIValue[datetime]):
    type: ClassVar[UserDataType] = UserDataType.datetime


@attr.s(frozen=True)
class DateTimeTZValue(BIValue[datetime]):
    type: ClassVar[UserDataType] = UserDataType.datetimetz


@attr.s(frozen=True)
class GenericDateTimeValue(BIValue[datetime]):
    type: ClassVar[UserDataType] = UserDataType.genericdatetime


@attr.s(frozen=True)
class BooleanValue(BIValue[bool]):
    type: ClassVar[UserDataType] = UserDataType.boolean


@attr.s(frozen=True)
class GeoPointValue(BIValue[List[Union[int, float]]]):
    type: ClassVar[UserDataType] = UserDataType.geopoint


@attr.s(frozen=True)
class GeoPolygonValue(BIValue[List[List[List[Union[int, float]]]]]):
    type: ClassVar[UserDataType] = UserDataType.geopolygon


@attr.s(frozen=True)
class UuidValue(BIValue[str]):
    type: ClassVar[UserDataType] = UserDataType.uuid


@attr.s(frozen=True)
class MarkupValue(BIValue[str]):
    type: ClassVar[UserDataType] = UserDataType.markup


@attr.s(frozen=True)
class ArrayStrValue(BIValue[List[str]]):
    type: ClassVar[UserDataType] = UserDataType.array_str


@attr.s(frozen=True)
class TreeStrValue(BIValue[List[str]]):
    type: ClassVar[UserDataType] = UserDataType.tree_str


@attr.s(frozen=True)
class ArrayIntValue(BIValue[List[int]]):
    type: ClassVar[UserDataType] = UserDataType.array_int


@attr.s(frozen=True)
class ArrayFloatValue(BIValue[List[float]]):
    type: ClassVar[UserDataType] = UserDataType.array_float

lib/dl_core/dl_core/values.py
datalens-tech/datalens-backend
import asyncio
from typing import Generator

import pytest

from dl_core.services_registry.top_level import ServicesRegistry
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import Db
from dl_core_testing.dataset_builder import (
    DatasetBuilderFactory,
    DefaultDbDatasetBuilderFactory,
)
from dl_core_testing.testcases.dataset import BaseDatasetTestClass
import dl_core_tests.db.config as test_config

from dl_connector_clickhouse.core.clickhouse.constants import SOURCE_TYPE_CH_TABLE
from dl_connector_clickhouse.core.clickhouse.testing.connection import make_clickhouse_saved_connection
from dl_connector_clickhouse.core.clickhouse.us_connection import ConnectionClickhouse
from dl_connector_clickhouse.core.clickhouse_base.constants import CONNECTION_TYPE_CLICKHOUSE
from dl_connector_clickhouse.db_testing.engine_wrapper import ClickhouseDbEngineConfig


class DefaultCoreTestClass(BaseDatasetTestClass[ConnectionClickhouse]):
    """Base class for generic, non-connectorized core tests"""

    conn_type = CONNECTION_TYPE_CLICKHOUSE
    source_type = SOURCE_TYPE_CH_TABLE

    core_test_config = test_config.CORE_TEST_CONFIG

    @pytest.fixture(autouse=True)
    # FIXME: This fixture is a temporary solution for failing core tests when they are run together with api tests
    def loop(self, event_loop: asyncio.AbstractEventLoop) -> Generator[asyncio.AbstractEventLoop, None, None]:
        asyncio.set_event_loop(event_loop)
        yield event_loop
        # Attempt to cover an old version of pytest-asyncio:
        # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
        asyncio.set_event_loop_policy(None)

    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        return test_config.DB_CORE_URL

    @pytest.fixture(scope="class")
    def engine_config(self, db_url: str, engine_params: dict) -> ClickhouseDbEngineConfig:
        return ClickhouseDbEngineConfig(url=db_url, engine_params=engine_params)

    @pytest.fixture(scope="function")
    def connection_creation_params(self) -> dict:
        return dict(
            db_name=test_config.CoreConnectionSettings.DB_NAME,
            host=test_config.CoreConnectionSettings.HOST,
            port=test_config.CoreConnectionSettings.PORT,
            username=test_config.CoreConnectionSettings.USERNAME,
            password=test_config.CoreConnectionSettings.PASSWORD,
            **(dict(raw_sql_level=self.raw_sql_level) if self.raw_sql_level is not None else {}),
        )

    @pytest.fixture(scope="function")
    def saved_connection(
        self,
        sync_us_manager: SyncUSManager,
        connection_creation_params: dict,
    ) -> ConnectionClickhouse:
        conn = make_clickhouse_saved_connection(sync_usm=sync_us_manager, **connection_creation_params)
        return conn

    @pytest.fixture(scope="function")
    def dataset_builder_factory(
        self,
        sync_us_manager: SyncUSManager,
        conn_default_service_registry: ServicesRegistry,
        db: Db,
        saved_connection: ConnectionBase,
    ) -> DatasetBuilderFactory:
        return DefaultDbDatasetBuilderFactory(
            service_registry=conn_default_service_registry,
            sync_us_manager=sync_us_manager,
            db=db,
            connection=saved_connection,
        )

lib/dl_core/dl_core_tests/db/base.py
datalens-tech/datalens-backend
from dl_core.components.dependencies.factory import ComponentDependencyManagerFactory
from dl_core_tests.db.base import DefaultCoreTestClass


class TestAvatarTreeResolution(DefaultCoreTestClass):
    def test_resolve_req_avatars_with_gaps(
        self,
        empty_saved_dataset,
        sync_us_manager,
        conn_default_service_registry,
        dataset_builder_factory,
    ):
        dataset = empty_saved_dataset
        dataset_builder = dataset_builder_factory.get_dataset_builder(dataset=dataset)

        dsrc_1_proxy = dataset_builder.add_data_source()
        dsrc_2_proxy = dataset_builder.add_data_source()
        dsrc_3_proxy = dataset_builder.add_data_source()

        avatar_1_proxy = dsrc_1_proxy.add_avatar()
        avatar_2_proxy = dsrc_2_proxy.add_avatar()
        avatar_3_proxy = dsrc_3_proxy.add_avatar()

        relation_1_2_proxy = avatar_1_proxy.add_relation_simple_eq(
            right=avatar_2_proxy, left_col_name="int_value", right_col_name="int_value"
        )
        relation_2_3_proxy = avatar_2_proxy.add_relation_simple_eq(
            right=avatar_3_proxy, left_col_name="int_value", right_col_name="int_value"
        )

        dep_mgr_factory = ComponentDependencyManagerFactory(dataset=dataset)
        tree_resolver = dep_mgr_factory.get_avatar_tree_resolver()
        root_avatar, required_avatar_ids, required_relation_ids = tree_resolver.expand_required_avatar_ids(
            required_avatar_ids={
                avatar_1_proxy.avatar_id,
                avatar_3_proxy.avatar_id,
            },  # leave a gap in place of avatar_2_id
        )

        assert root_avatar == avatar_1_proxy.avatar_id
        assert required_avatar_ids == {avatar_1_proxy.avatar_id, avatar_2_proxy.avatar_id, avatar_3_proxy.avatar_id}
        assert required_relation_ids == {relation_1_2_proxy.relation_id, relation_2_3_proxy.relation_id}

lib/dl_core/dl_core_tests/db/components/test_avatar_tree_resolution.py
datalens-tech/datalens-backend
import uuid

from dl_constants.enums import ComponentType
from dl_core_tests.db.base import DefaultCoreTestClass


class TestComponentErrors(DefaultCoreTestClass):
    def test_component_errors(self, sync_us_manager, saved_dataset):
        dataset = saved_dataset
        source_id = dataset.get_single_data_source_id()
        erreg = dataset.error_registry
        erreg.add_error(id=source_id, type=ComponentType.data_source, message="This is an error", code=["ERR", "1"])
        assert len(erreg.get_pack(id=source_id).errors) == 1
        erreg.add_error(id=source_id, type=ComponentType.data_source, message="This is the end", code=["ERR", "2"])
        assert len(erreg.get_pack(id=source_id).errors) == 2

        field_id = "111111"
        erreg.add_error(
            id=field_id,
            type=ComponentType.field,
            message="This field needs ploughing",
            code=["ERR", "3"],
            details=dict(row=10, column=8),
        )
        erreg.add_error(
            id=field_id,
            type=ComponentType.field,
            message="Just an error",
            code=["ERR", "3", "1"],
            details=dict(row=100, column=80),
        )

        obligatory_filter_id = str(uuid.uuid4())
        erreg.add_error(
            id=obligatory_filter_id,
            type=ComponentType.obligatory_filter,
            message="obligatory_filter test err",
            code=["ERR", "OBLIG_FILT", "TEST"],
        )

        sync_us_manager.save(dataset)
        dataset = sync_us_manager.get_by_id(entry_id=dataset.uuid)

        erreg = dataset.error_registry
        field_errors = erreg.get_pack(id=field_id).errors
        assert len(field_errors) == 2
        assert field_errors[0].code == ["ERR", "3"]
        assert field_errors[0].details["row"] == 10
        assert field_errors[0].details["column"] == 8

        oblig_filter_errors = erreg.for_type(type=ComponentType.obligatory_filter)
        assert len(oblig_filter_errors) == 1
        assert oblig_filter_errors[0].id == obligatory_filter_id
        assert len(oblig_filter_errors[0].errors) == 1
        assert oblig_filter_errors[0].errors[0].message == "obligatory_filter test err"

        erreg.remove_errors(id=field_id, code=["ERR", "3"])

        sync_us_manager.save(dataset)
        dataset = sync_us_manager.get_by_id(entry_id=dataset.uuid)

        erreg = dataset.error_registry
        assert len(erreg.get_pack(id=field_id).errors) == 1

        erreg.remove_errors(id=field_id, code_prefix=["ERR", "3"])

        sync_us_manager.save(dataset)
        dataset = sync_us_manager.get_by_id(entry_id=dataset.uuid)

        erreg = dataset.error_registry
        assert erreg.get_pack(id=field_id) is None

lib/dl_core/dl_core_tests/db/components/test_component_errors.py
datalens-tech/datalens-backend
from typing import ClassVar

from dl_core_testing.configuration import DefaultCoreTestConfiguration
from dl_testing.containers import get_test_container_hostport


# Infra settings
CORE_TEST_CONFIG = DefaultCoreTestConfiguration(
    host_us_http=get_test_container_hostport("us", fallback_port=50300).host,
    port_us_http=get_test_container_hostport("us", fallback_port=50300).port,
    host_us_pg=get_test_container_hostport("pg-us", fallback_port=50309).host,
    port_us_pg_5432=get_test_container_hostport("pg-us", fallback_port=50309).port,
    us_master_token="AC1ofiek8coB",
)


class CoreConnectionSettings:
    DB_NAME: ClassVar[str] = "test_data"
    HOST: ClassVar[str] = get_test_container_hostport("db-clickhouse", fallback_port=50310).host
    PORT: ClassVar[int] = get_test_container_hostport("db-clickhouse", fallback_port=50310).port
    USERNAME: ClassVar[str] = "datalens"
    PASSWORD: ClassVar[str] = "qwerty"


DB_CORE_URL = (
    f"clickhouse://datalens:qwerty@"
    f'{get_test_container_hostport("db-clickhouse", fallback_port=50310).as_pair()}/test_data'
)

lib/dl_core/dl_core_tests/db/config.py
datalens-tech/datalens-backend
from dl_core_testing.initialization import initialize_core_test
from dl_core_tests.db.config import CORE_TEST_CONFIG


pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'


def pytest_configure(config):  # noqa
    initialize_core_test(pytest_config=config, core_test_config=CORE_TEST_CONFIG)

lib/dl_core/dl_core_tests/db/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    Tuple,
)

import attr
import shortuuid

from dl_constants.enums import (
    AggregationFunction,
    BinaryJoinOperator,
    FieldType,
    ManagedBy,
    UserDataType,
)
from dl_core.components.dependencies.field_avatar import FieldAvatarDependencyManager
from dl_core.components.dependencies.field_deep import FieldDeepInterDependencyManager
from dl_core.components.dependencies.field_shallow import FieldShallowInterDependencyManager
from dl_core.components.dependencies.primitives import FieldInterDependencyInfo
from dl_core.components.dependencies.relation_avatar import RelationAvatarDependencyManager
from dl_core.components.ids import (
    AvatarId,
    FieldId,
)
from dl_core.fields import (
    BIField,
    DirectCalculationSpec,
    FormulaCalculationSpec,
    ResultSchema,
)
from dl_core.multisource import (
    AvatarRelation,
    BinaryCondition,
    ConditionPartDirect,
    ConditionPartResultField,
)


@attr.s
class ShallowInfo:
    field_1_id: FieldId = attr.ib(factory=shortuuid.uuid)
    field_2_id: FieldId = attr.ib(factory=shortuuid.uuid)
    field_3_id: FieldId = attr.ib(factory=shortuuid.uuid)
    field_4_id: FieldId = attr.ib(factory=shortuuid.uuid)
    shallow: FieldShallowInterDependencyManager = attr.ib(
        factory=lambda: FieldShallowInterDependencyManager(inter_dep_info=FieldInterDependencyInfo())
    )


def make_shallow() -> ShallowInfo:
    shallow_info = ShallowInfo()

    shallow_info.shallow.set_field_direct_references(
        dep_field_id=shallow_info.field_1_id, ref_field_ids=[shallow_info.field_2_id, shallow_info.field_3_id]
    )
    shallow_info.shallow.set_field_direct_references(
        dep_field_id=shallow_info.field_4_id, ref_field_ids=[shallow_info.field_1_id]
    )
    return shallow_info


def make_direct_field(field_id: FieldId, avatar_id: Optional[AvatarId]) -> BIField:
    return BIField.make(
        guid=field_id,
        title=shortuuid.uuid(),
        type=FieldType.DIMENSION,
        aggregation=AggregationFunction.none,
        initial_data_type=UserDataType.integer,
        cast=UserDataType.integer,
        data_type=UserDataType.integer,
        has_auto_aggregation=False,
        lock_aggregation=False,
        valid=True,
        managed_by=ManagedBy.user,
        calc_spec=DirectCalculationSpec(
            avatar_id=avatar_id,
            source=shortuuid.uuid(),
        ),
    )


def make_formula_field(field_id: FieldId) -> BIField:
    return BIField.make(
        guid=field_id,
        title=shortuuid.uuid(),
        type=FieldType.DIMENSION,
        aggregation=AggregationFunction.none,
        initial_data_type=UserDataType.integer,
        cast=UserDataType.integer,
        data_type=UserDataType.integer,
        has_auto_aggregation=False,
        lock_aggregation=False,
        valid=True,
        managed_by=ManagedBy.user,
        calc_spec=FormulaCalculationSpec(
            formula="...",
        ),
    )


def make_result_schema(shallow_info) -> Tuple[ResultSchema, AvatarId, AvatarId]:
    avatar_1_id = shortuuid.uuid()
    avatar_2_id = shortuuid.uuid()
    return (
        ResultSchema(
            fields=[
                make_formula_field(field_id=shallow_info.field_1_id),
                make_direct_field(field_id=shallow_info.field_2_id, avatar_id=avatar_1_id),
                make_direct_field(field_id=shallow_info.field_3_id, avatar_id=avatar_2_id),
                make_formula_field(field_id=shallow_info.field_4_id),
            ]
        ),
        avatar_1_id,
        avatar_2_id,
    )


def test_shallow_manager():
    shallow_info = make_shallow()

    assert shallow_info.shallow.get_field_direct_references(dep_field_id=shallow_info.field_1_id) == {
        shallow_info.field_2_id,
        shallow_info.field_3_id,
    }
    assert shallow_info.shallow.get_field_direct_references(dep_field_id=shallow_info.field_2_id) == set()
    assert shallow_info.shallow.get_field_direct_references(dep_field_id=shallow_info.field_4_id) == {
        shallow_info.field_1_id,
    }


def test_deep_manager():
    shallow_info = make_shallow()
    deep = FieldDeepInterDependencyManager(shallow=shallow_info.shallow)

    assert deep.get_field_deep_references(dep_field_id=shallow_info.field_1_id) == {
        shallow_info.field_2_id,
        shallow_info.field_3_id,
    }
    assert deep.get_field_deep_references(dep_field_id=shallow_info.field_2_id) == set()
    assert deep.get_field_deep_references(dep_field_id=shallow_info.field_4_id) == {
        shallow_info.field_1_id,
        shallow_info.field_2_id,
        shallow_info.field_3_id,
    }


def test_field_avatar_manager():
    shallow_info = make_shallow()
    deep = FieldDeepInterDependencyManager(shallow=shallow_info.shallow)
    result_schema, avatar_1_id, avatar_2_id = make_result_schema(shallow_info=shallow_info)
    field_avatar_mgr = FieldAvatarDependencyManager(
        result_schema=result_schema,
        deep=deep,
    )

    assert field_avatar_mgr.get_field_avatar_references(dep_field_id=shallow_info.field_1_id) == {
        avatar_1_id,
        avatar_2_id,
    }
    assert field_avatar_mgr.get_field_avatar_references(dep_field_id=shallow_info.field_2_id) == {
        avatar_1_id,
    }
    assert field_avatar_mgr.get_field_avatar_references(dep_field_id=shallow_info.field_4_id) == {
        avatar_1_id,
        avatar_2_id,
    }


def test_relation_avatar_manager():
    shallow_info = make_shallow()
    deep = FieldDeepInterDependencyManager(shallow=shallow_info.shallow)
    result_schema, avatar_1_id, avatar_2_id = make_result_schema(shallow_info=shallow_info)
    avatar_3_id = shortuuid.uuid()
    field_avatar_mgr = FieldAvatarDependencyManager(
        result_schema=result_schema,
        deep=deep,
    )
    relation_1_id = shortuuid.uuid()
    relation_2_id = shortuuid.uuid()
    avatar_relations = [
        AvatarRelation(
            id=relation_1_id,
            left_avatar_id=avatar_1_id,
            right_avatar_id=avatar_2_id,
            conditions=[
                BinaryCondition(
                    operator=BinaryJoinOperator.eq,
                    left_part=ConditionPartDirect(source="col_1"),
                    right_part=ConditionPartDirect(source="col_2"),
                )
            ],
        ),
        AvatarRelation(
            id=relation_2_id,
            left_avatar_id=avatar_1_id,
            right_avatar_id=avatar_3_id,
            conditions=[
                BinaryCondition(
                    operator=BinaryJoinOperator.eq,
                    left_part=ConditionPartResultField(field_id=shallow_info.field_1_id),
                    right_part=ConditionPartDirect(source="col_2"),
                )
            ],
        ),
    ]
    relation_avatar_mgr = RelationAvatarDependencyManager(
        avatar_relations=avatar_relations,
        field_avatar_mgr=field_avatar_mgr,
    )

    assert relation_avatar_mgr.get_relation_avatar_references(relation_id=relation_1_id) == {
        avatar_1_id,
        avatar_2_id,
    }
    assert relation_avatar_mgr.get_relation_avatar_references(relation_id=relation_2_id) == {
        avatar_1_id,
        avatar_2_id,
        avatar_3_id,
    }

lib/dl_core/dl_core_tests/unit/components/test_dependencies.py
datalens-tech/datalens-backend
from dl_core.components.ids import (
    ID_LENGTH,
    ID_VALID_SYMBOLS,
    make_readable_field_id,
)


def test_make_readable_field_id():
    assert make_readable_field_id(title="hello!:*", valid_symbols=ID_VALID_SYMBOLS, max_length=ID_LENGTH) == "hello"
    assert (
        make_readable_field_id(
            title=" René François Lacôte   123", valid_symbols=ID_VALID_SYMBOLS, max_length=ID_LENGTH
        )
        == "rene_francois_lacote_123"
    )
    assert make_readable_field_id(title="Проверка", valid_symbols=ID_VALID_SYMBOLS, max_length=ID_LENGTH) == "proverka"

lib/dl_core/dl_core_tests/unit/components/test_make_readable_field_id.py
datalens-tech/datalens-backend
import asyncio

import pytest


pytest_plugins = ("aiohttp.pytest_plugin",)  # and it, in turn, includes 'pytest_asyncio.plugin'


@pytest.fixture
def loop(event_loop):
    asyncio.set_event_loop(event_loop)
    yield event_loop
    # Attempt to cover an old version of pytest-asyncio:
    # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
    asyncio.set_event_loop_policy(None)

lib/dl_core/dl_core_tests/unit/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Generator

import pytest

from dl_core.data_processing.streaming import (
    AsyncChunked,
    AsyncChunkedLimited,
    LazyAsyncChunked,
)
import dl_core.exc as exc


def chunked_range(size: int, chunk_size: int) -> Generator[Generator[int, None, None], None, None]:
    for chunk_start in range(0, size, chunk_size):
        yield range(chunk_start, min(chunk_start + chunk_size, size))


@pytest.mark.asyncio
async def test_async_chunked():
    # chunks
    chunked = AsyncChunked.from_chunked_iterable(chunked_range(97, 10))
    data = [item async for chunk in chunked.chunks for item in chunk]
    assert data == list(range(97))

    # items
    chunked = AsyncChunked.from_chunked_iterable(chunked_range(97, 10))
    data = [item async for item in chunked.items]
    assert data == list(range(97))

    # all
    chunked = AsyncChunked.from_chunked_iterable(chunked_range(97, 10))
    data = await chunked.all()
    assert data == list(range(97))


@pytest.mark.asyncio
async def test_async_chunked_limited():
    # chunks
    chunked = AsyncChunkedLimited(chunked=AsyncChunked.from_chunked_iterable(chunked_range(97, 10)), max_count=53)
    data = []
    with pytest.raises(exc.ResultRowCountLimitExceeded):
        async for chunk in chunked.chunks:
            for item in chunk:
                data.append(item)
    # Data is appended in whole chunks,
    # so the chunk that reaches the limit is never returned,
    # which is why we have only 50
    assert data == list(range(50))

    # items
    chunked = AsyncChunkedLimited(chunked=AsyncChunked.from_chunked_iterable(chunked_range(97, 10)), max_count=53)
    data = []
    with pytest.raises(exc.ResultRowCountLimitExceeded):
        async for item in chunked.items:
            data.append(item)
    assert data == list(range(53))

    # all
    chunked = AsyncChunkedLimited(chunked=AsyncChunked.from_chunked_iterable(chunked_range(97, 10)), max_count=53)
    with pytest.raises(exc.ResultRowCountLimitExceeded):
        await chunked.all()


@pytest.mark.asyncio
async def test_async_chunked_lazy():
    started_cnt = 0
    finished_cnt = 0

    async def make_range() -> AsyncChunked[int]:
        nonlocal started_cnt
        started_cnt += 1
        return AsyncChunked.from_chunked_iterable(chunked_range(97, 10))

    async def finalize() -> None:
        nonlocal finished_cnt
        finished_cnt += 1

    # chunks
    chunked = LazyAsyncChunked(initializer=make_range, finalizer=finalize)
    assert started_cnt == 0
    data = [item async for chunk in chunked.chunks for item in chunk]
    assert data == list(range(97))
    assert started_cnt == finished_cnt == 1

    # items
    chunked = LazyAsyncChunked(initializer=make_range, finalizer=finalize)
    assert started_cnt == 1
    data = [item async for item in chunked.items]
    assert data == list(range(97))
    assert started_cnt == finished_cnt == 2

    # all
    chunked = LazyAsyncChunked(initializer=make_range, finalizer=finalize)
    assert started_cnt == 2
    data = await chunked.all()
    assert data == list(range(97))
    assert started_cnt == finished_cnt == 3

lib/dl_core/dl_core_tests/unit/data_processing/test_streaming.py
datalens-tech/datalens-backend
import base64

import marshmallow
import pytest

import dl_core.marshmallow as bi_core_marshmallow


class Schema(marshmallow.Schema):
    field = bi_core_marshmallow.Base64StringField()


def test_load():
    schema = Schema()

    value = b"data:smt/smt;base64," + base64.b64encode(b"test")

    data = schema.load(data={"field": value})
    assert data["field"] == "test"

    data = schema.load(data={"field": value.decode(encoding="utf-8")})
    assert data["field"] == "test"

    data = schema.load(data={"field": ""})
    assert data["field"] is None

    with pytest.raises(marshmallow.exceptions.ValidationError):
        schema.load(data=base64.b64encode(b"test"))

    with pytest.raises(marshmallow.exceptions.ValidationError):
        schema.load(data={"field": 123})

lib/dl_core/dl_core_tests/unit/marshmallow/test_base64_string_field.py
datalens-tech/datalens-backend
import marshmallow
import pytest

import dl_core.marshmallow as bi_core_marshmallow


class Schema(marshmallow.Schema):
    field = bi_core_marshmallow.OnOffField()


def test_dump():
    schema = Schema()

    data = schema.dump(obj={"field": True})
    assert data["field"] == "on"

    data = schema.dump(obj={"field": False})
    assert data["field"] == "off"


def test_load():
    schema = Schema()

    data = schema.load(data={"field": "on"})
    assert data["field"] is True

    data = schema.load(data={"field": "off"})
    assert data["field"] is False

    with pytest.raises(marshmallow.exceptions.ValidationError):
        schema.load(data={"field": True})

    with pytest.raises(marshmallow.exceptions.ValidationError):
        schema.load(data={"field": "On"})

lib/dl_core/dl_core_tests/unit/marshmallow/test_on_off_field.py
datalens-tech/datalens-backend
from __future__ import annotations

from sqlalchemy import exc as sa_exc

from dl_core import exc
import dl_core.connectors.base.error_transformer as error_transformer
from dl_core.connectors.base.error_transformer import wrapper_exc_is_and_matches_re


def test_transform_with_custom_rules():
    transformer = error_transformer.make_default_transformer_with_custom_rules(
        error_transformer.ErrorTransformerRule(
            when=wrapper_exc_is_and_matches_re(Exception, "Hello"), then_raise=exc.InvalidQuery
        )
    )

    transformed_invalid_query = transformer.make_bi_error(
        sa_exc.OperationalError(orig=Exception("Well, Hello"), params={}, statement="")
    )
    assert isinstance(transformed_invalid_query, exc.InvalidQuery)

    transformed_operational = transformer.make_bi_error(
        sa_exc.OperationalError(orig=Exception(), params={}, statement="")
    )
    assert isinstance(transformed_operational, exc.DatabaseOperationalError)

    transformed_no_such_table = transformer.make_bi_error(
        sa_exc.OperationalError(orig=sa_exc.NoSuchTableError(), params={}, statement="")
    )
    assert isinstance(transformed_no_such_table, exc.SourceDoesNotExist)

    transformed_default = transformer.make_bi_error(Exception())
    assert isinstance(transformed_default, exc.DatabaseQueryError)

lib/dl_core/dl_core_tests/unit/test_error_transformer.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import decimal
import json

from dl_core.serialization import (
    RedisDatalensDataJSONDecoder,
    RedisDatalensDataJSONEncoder,
)


SAMPLE_DATA = dict(
    some_dt=datetime.datetime(
        2019, 6, 17, 7, 1, 41, 79585, tzinfo=datetime.timezone(datetime.timedelta(seconds=-1320))
    ),
    some_date=datetime.date(2020, 7, 18),
    some_time=datetime.time(11, 22, 33, 444555, tzinfo=datetime.timezone(datetime.timedelta(seconds=1320))),
    some_timedelta=datetime.timedelta(seconds=1320.0231),
    some_decimal=decimal.Decimal("12345" * 10 + "." + "54321" * 10),
)


def test_json_serialization():
    data = SAMPLE_DATA
    dumped = json.dumps(data, cls=RedisDatalensDataJSONEncoder)
    dumped_dt = json.loads(dumped)["some_dt"]
    assert dumped_dt == {"__dl_type__": "datetime", "value": [[2019, 6, 17, 7, 1, 41, 79585], -1320.0]}, dumped_dt
    roundtrip = json.loads(dumped, cls=RedisDatalensDataJSONDecoder)
    assert roundtrip == data


def test_json_tricky_serialization():
    data = SAMPLE_DATA
    dumped = json.dumps(data, cls=RedisDatalensDataJSONEncoder)
    tricky_data = dict(normal=data, abnormal=json.loads(dumped))
    tricky_data_dumped = json.dumps(tricky_data, cls=RedisDatalensDataJSONEncoder)
    tricky_roundtrip = json.loads(tricky_data_dumped, cls=RedisDatalensDataJSONDecoder)
    # WARNING: not currently viable: `assert tricky_roundtrip == tricky_data, tricky_roundtrip`
    assert tricky_roundtrip

lib/dl_core/dl_core_tests/unit/test_json_serializer.py
datalens-tech/datalens-backend
from __future__ import annotations

from datetime import datetime
import uuid

from multidict import CIMultiDict
import pytest

from dl_api_commons.base_models import (
    NoAuthData,
    RequestContextInfo,
    TenantCommon,
)
from dl_constants.api_constants import (
    DLContextKey,
    DLHeadersCommon,
)
from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
from dl_core.connection_executors.qe_serializer.schemas_common import (
    DBAdapterQueryStrSchema,
    DBAdapterScopedRCISchema,
)


# Full RCI is used in test-cases to prevent extra work when scope of DBAdapterScopedRCI will be extended
_TEST_REQ_CTX_INFO = (
    DBAdapterScopedRCI.from_full_rci(
        RequestContextInfo.create(
            request_id=None,
            tenant=None,
            user_id=None,
            user_name=None,
            x_dl_debug_mode=None,
            endpoint_code=None,
            x_dl_context=None,
            plain_headers=None,
            secret_headers=None,
            auth_data=None,
        )
    ),
    DBAdapterScopedRCI.from_full_rci(
        RequestContextInfo.create(
            request_id=str(uuid.uuid4()),
            tenant=TenantCommon(),
            user_id=str(uuid.uuid4()),
            user_name="vasya",
            x_dl_debug_mode=False,
            endpoint_code="some_endpoint",
            x_dl_context={},
            plain_headers={
                "uber-trace-id": "805324dfd85a80dc699a45d557ccf4a:418f6626fe7843af:c699a45d557ccf4a:1",
            },
            secret_headers={},
            auth_data=None,
        )
    ),
    DBAdapterScopedRCI.from_full_rci(
        RequestContextInfo.create(
            request_id=str(uuid.uuid4()),
            tenant=TenantCommon(),
            user_id=str(uuid.uuid4()),
            user_name="vasya",
            x_dl_debug_mode=True,
            endpoint_code="some_endpoint",
            x_dl_context={DLContextKey.CHART_ID.value: str(uuid.uuid4())},
            plain_headers={"a": "b", DLHeadersCommon.FORWARDED_FOR.value: "1.1.1.1,2.2.2.2,3.3.3.3"},
            secret_headers={"c": "d"},
            auth_data=NoAuthData(),
        )
    ),
    DBAdapterScopedRCI.from_full_rci(
        RequestContextInfo.create(
            request_id=str(uuid.uuid4()),
            tenant=TenantCommon(),
            user_id=str(uuid.uuid4()),
            user_name="vasya",
            x_dl_debug_mode=True,
            endpoint_code="some_endpoint",
            x_dl_context={k.value: str(uuid.uuid4()) for k in DLContextKey},
            plain_headers={"a": "b"},
            secret_headers={"c": "d"},
            auth_data=NoAuthData(),
        )
    ),
    DBAdapterScopedRCI.from_full_rci(
        RequestContextInfo.create(
            request_id=str(uuid.uuid4()),
            tenant=TenantCommon(),
            user_id=str(uuid.uuid4()),
            user_name="vasya",
            x_dl_debug_mode=True,
            endpoint_code="some_endpoint",
            x_dl_context={"asdf": "qwer", "None": None},
            plain_headers={"a": "b"},
            secret_headers={"c": "d"},
            auth_data=NoAuthData(),
        )
    ),
    DBAdapterScopedRCI.from_full_rci(
        RequestContextInfo.create(
            request_id=str(uuid.uuid4()),
            tenant=TenantCommon(),
            user_id=str(uuid.uuid4()),
            user_name="vasya",
            x_dl_debug_mode=True,
            endpoint_code="some_endpoint",
            x_dl_context={DLContextKey.CHART_ID.value: str(uuid.uuid4())},
            plain_headers=CIMultiDict([("a", "a1"), ("a", "a2")]),
            secret_headers=None,
            auth_data=None,
        )
    ),
    DBAdapterScopedRCI.from_full_rci(
        RequestContextInfo.create(
            request_id=str(uuid.uuid4()),
            tenant=TenantCommon(),
            user_id=str(uuid.uuid4()),
            user_name="vasya",
            x_dl_debug_mode=True,
            endpoint_code="some_endpoint",
            x_dl_context={DLContextKey.CHART_ID.value: str(uuid.uuid4()), "other": None},
            plain_headers=CIMultiDict([("a", "a1"), ("a", "a2")]),
            secret_headers=None,
            auth_data=None,
        )
    ),
)


_TEST_DB_ADAPTER_QUERY = (
    DBAdapterQuery(
        query="",
    ),
    DBAdapterQuery(
        query="",
        connector_specific_params={"from": "2020-01-01T00:00:00Z"},
    ),
    DBAdapterQuery(
        query="",
        connector_specific_params={"from": datetime(2020, 1, 1)},
    ),
)


@pytest.mark.parametrize(
    "req_ctx_info",
    _TEST_REQ_CTX_INFO,
)
def test_req_ctx_info_schema(req_ctx_info: DBAdapterScopedRCI):
    schema = DBAdapterScopedRCISchema()
    dumped_req_ctx_info = schema.dump(req_ctx_info)
    reloaded_req_ctx_info = schema.load(dumped_req_ctx_info)

    assert reloaded_req_ctx_info == req_ctx_info


@pytest.mark.parametrize(
    "db_adapter_query",
    _TEST_DB_ADAPTER_QUERY,
)
def test_db_adapter_query_schema(db_adapter_query: DBAdapterQuery):
    schema = DBAdapterQueryStrSchema()
    dumped_db_adapter_query = schema.dump(db_adapter_query)
    reloaded_db_adapter_query = schema.load(dumped_db_adapter_query)

    assert reloaded_db_adapter_query == db_adapter_query

lib/dl_core/dl_core_tests/unit/test_qe_schemas.py
datalens-tech/datalens-backend
from __future__ import annotations

from aiohttp import web
import pytest

from dl_core.base_models import PathEntryLocation
from dl_core.exc import USBadRequestException
from dl_core.united_storage_client import USAuthContextMaster
from dl_core.united_storage_client_aio import UStorageClientAIO


@pytest.mark.asyncio
async def test_fields_masking(aiohttp_client, caplog):
    caplog.set_level("INFO", logger="dl_core.united_storage_client")

    @web.middleware
    async def always_400_mw(request, handler):
        return web.json_response({}, status=400)

    app = web.Application(middlewares=[always_400_mw])

    mock = await aiohttp_client(app)

    us_client = UStorageClientAIO(
        auth_ctx=USAuthContextMaster("fake_token"), host=f"http://{mock.host}:{mock.port}", prefix="/api/private"
    )

    secret_value = "some_cypher_text_898"
    non_secret_value = "some_non_secret_value_989"

    with pytest.raises(USBadRequestException):
        await us_client.create_entry(
            key=PathEntryLocation("fake_key"),
            scope="connection",
            data=dict(
                password=secret_value,
                plain_val=non_secret_value,
            ),
            unversioned_data=dict(token=dict(cypher_text=secret_value)),
        )

    try:
        rec = next(r for r in caplog.records if r.message.startswith("US error response on"))
    # StopIteration causes RuntimeError in async test launcher
    except StopIteration as err:
        raise AssertionError() from err

    assert non_secret_value in rec.message, f"Expected non-secret {repr(secret_value)} not found in logs: {repr(rec)}"
    assert secret_value not in rec.message, f"Secret {repr(secret_value)} found in logs: {repr(rec)}"

lib/dl_core/dl_core_tests/unit/us_manager/test_us_client.py
datalens-tech/datalens-backend
"""
Helper functions and classes for usage in tests of all dependent projects
"""

lib/dl_core_testing/dl_core_testing/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_api_commons.base_models import RequestContextInfo
from dl_core.connections_security.base import (
    ConnectionSecurityManager,
    InsecureConnectionSecurityManager,
)
from dl_core.services_registry.env_manager_factory import DefaultEnvManagerFactory


class TestEnvManagerFactory(DefaultEnvManagerFactory):
    def make_security_manager(self) -> ConnectionSecurityManager:
        return InsecureConnectionSecurityManager()

lib/dl_core_testing/dl_core_testing/app_test_workarounds.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
from typing import (
    Collection,
    Optional,
)

import attr

from dl_configs.crypto_keys import (
    CryptoKeysConfig,
    get_single_key_crypto_keys_config,
)
from dl_core.loader import CoreLibraryConfig


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class UnitedStorageConfiguration:
    us_host: str = attr.ib(kw_only=True)
    us_master_token: str = attr.ib(kw_only=True)
    us_pg_dsn: str = attr.ib(kw_only=True)
    force: bool = attr.ib(kw_only=True, default=True)


@attr.s(frozen=True)
class CoreTestEnvironmentConfigurationBase(abc.ABC):
    @abc.abstractmethod
    def get_us_config(self) -> UnitedStorageConfiguration:
        raise NotImplementedError

    @abc.abstractmethod
    def get_crypto_keys_config(self) -> CryptoKeysConfig:
        raise NotImplementedError

    @abc.abstractmethod
    def get_core_library_config(self) -> CoreLibraryConfig:
        raise NotImplementedError


# These are used only for creation of local environments in tests, not actual external ones
DEFAULT_FERNET_KEY = "h1ZpilcYLYRdWp7Nk8X1M1kBPiUi8rdjz9oBfHyUKIk="


@attr.s(frozen=True)
class DefaultCoreTestConfiguration(CoreTestEnvironmentConfigurationBase):
    host_us_http: str = attr.ib(kw_only=True)
    port_us_http: int = attr.ib(kw_only=True)
    host_us_pg: str = attr.ib(kw_only=True)
    port_us_pg_5432: int = attr.ib(kw_only=True)
    us_master_token: str = attr.ib(kw_only=True)
    fernet_key: str = attr.ib(kw_only=True, default=DEFAULT_FERNET_KEY)
    core_connector_ep_names: Optional[Collection[str]] = attr.ib(kw_only=True, default=None)

    def get_us_config(self) -> UnitedStorageConfiguration:
        return UnitedStorageConfiguration(
            us_master_token=self.us_master_token,
            us_host=f"http://{self.host_us_http}:{self.port_us_http}",
            us_pg_dsn=f"host={self.host_us_pg} port={self.port_us_pg_5432} user=us password=us dbname=us-db-ci_purgeable",
        )

    def get_crypto_keys_config(self) -> CryptoKeysConfig:
        return get_single_key_crypto_keys_config(key_id="0", key_value=self.fernet_key)

    def get_core_library_config(self) -> CoreLibraryConfig:
        return CoreLibraryConfig(
            core_connector_ep_names=self.core_connector_ep_names,
        )

lib/dl_core_testing/dl_core_testing/configuration.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Optional,
    Sequence,
    Type,
)
import uuid

from dl_constants.enums import ConnectionType
from dl_core.base_models import (
    ConnectionDataModelBase,
    EntryLocation,
    PathEntryLocation,
)
from dl_core.us_connection import get_connection_class
from dl_core.us_manager.us_manager import USManagerBase
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import Db


if TYPE_CHECKING:
    from dl_core.us_connection_base import ConnectionBase


def make_conn_key(*args) -> str:  # type: ignore  # TODO: fix
    return "/".join(["tests", *args])


def _ensure_entry_location(loc: Optional[EntryLocation], gen_path_postfix: Sequence[str]) -> EntryLocation:
    if loc is not None:
        return loc
    else:
        return PathEntryLocation(path=make_conn_key(*gen_path_postfix))


def make_connection_base(
    us_manager: USManagerBase,
    conn_type: ConnectionType,
    conn_name: str = None,
    data_dict: Optional[dict] = None,
    conn_cls: Optional[Type[ConnectionBase]] = None,
) -> ConnectionBase:
    conn_name = conn_name or "{} test conn {}".format(conn_type.name, uuid.uuid4())
    data_dict = data_dict or {}
    if conn_cls is None:
        conn_cls = get_connection_class(conn_type)

    effective_data: dict | ConnectionDataModelBase
    data_model_cls = conn_cls.DataModel

    if issubclass(data_model_cls, ConnectionDataModelBase):
        effective_data = data_model_cls(**data_dict)  # type: ignore
    else:
        raise TypeError(f"Unexpected data model class: {data_model_cls}")

    return conn_cls.create_from_dict(
        effective_data,
        ds_key=make_conn_key("connections", conn_name),
        type_=conn_type.name,
        meta={"title": conn_name, "state": "saved"},
        us_manager=us_manager,
    )


def make_connection(
    us_manager: USManagerBase,
    db: Optional[Db] = None,
    conn_type: Optional[ConnectionType] = None,
    conn_name: str = None,
    data_dict: Optional[dict] = None,
) -> ConnectionBase:
    credentials: dict = {}
    if conn_type is None:
        assert db is not None
        conn_type = db.conn_type
        credentials = db.get_conn_credentials()
    assert conn_type is not None
    data_dict = {
        **credentials,
        **(data_dict or {}),
    }
    return make_connection_base(
        us_manager=us_manager,
        conn_type=conn_type,
        conn_name=conn_name,
        data_dict=data_dict,
    )


def make_saved_connection(  # type: ignore  # TODO: fix
    sync_usm: SyncUSManager,
    db: Optional[Db] = None,
    conn_type: Optional[ConnectionType] = None,
    conn_name: str = None,
    data_dict: Optional[dict] = None,
) -> ConnectionBase:
    conn = make_connection(us_manager=sync_usm, db=db, conn_type=conn_type, conn_name=conn_name, data_dict=data_dict)
    sync_usm.save(conn)
    return conn

lib/dl_core_testing/dl_core_testing/connection.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    ConnectionType,
    DataSourceType,
)
from dl_core.connectors.base.connector import (
    CoreConnectionDefinition,
    CoreConnector,
    CoreSourceDefinition,
)
from dl_core.db.conversion_base import TypeTransformer
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_manager.storage_schemas.connection import BaseConnectionDataStorageSchema


CONNECTION_TYPE_TESTING = ConnectionType.declare("testing")
SOURCE_TYPE_TESTING = DataSourceType.declare("TESTING")


class TestingConnection(ConnectionBase):
    pass


class TestingTypeTransformer(TypeTransformer):
    conn_type = CONNECTION_TYPE_TESTING


class TestingConnectionDataStorageSchema(BaseConnectionDataStorageSchema):
    TARGET_CLS = TestingConnection.DataModel


class TestingCoreSourceDefinition(CoreSourceDefinition):
    source_type = SOURCE_TYPE_TESTING


class TestingCoreConnectionDefinition(CoreConnectionDefinition):
    connection_cls = TestingConnection
    conn_type = CONNECTION_TYPE_TESTING
    type_transformer_cls = TestingTypeTransformer
    us_storage_schema_cls = TestingConnectionDataStorageSchema
    dialect_string = "testing"  # No such dialect


class TestingCoreConnector(CoreConnector):
    connection_definitions = (TestingCoreConnectionDefinition,)
    source_definitions = (TestingCoreSourceDefinition,)

lib/dl_core_testing/dl_core_testing/connector.py
datalens-tech/datalens-backend
from __future__ import annotations

import csv
import io
from typing import (
    Dict,
    List,
)


def make_csv(headers: List[str], rows: List[Dict[str, str]]) -> str:
    """Turn list of headers and list of data rows into a CSV string"""
    f = io.StringIO()
    writer = csv.DictWriter(f, fieldnames=headers)
    writer.writeheader()
    for row in rows:
        writer.writerow(row)
    return f.getvalue()

lib/dl_core_testing/dl_core_testing/converter.py
datalens-tech/datalens-backend
import csv
import datetime
import io
import random
from typing import (
    Any,
    Callable,
    ClassVar,
    Optional,
    Sequence,
)

import attr
import shortuuid

from dl_constants.enums import UserDataType
from dl_core_testing.database import (
    C,
    Db,
    DbTable,
    make_table,
)


def _int_or_none(v: Optional[str]) -> Optional[int]:
    return int(v) if v is not None else None


def _float_or_none(v: Optional[str]) -> Optional[float]:
    return float(v) if v is not None else None


def _date_or_none(v: Optional[str]) -> Optional[datetime.date]:
    return datetime.date.fromisoformat(v) if v is not None else None


def _datetime_or_none(v: Optional[str]) -> Optional[datetime.datetime]:
    return datetime.datetime.fromisoformat(v) if v is not None else None


@attr.s
class CsvTableDumper:
    db: Db = attr.ib(kw_only=True)

    _PY_CONVERTERS_BY_BITYPE: ClassVar[dict[UserDataType, Callable[[Optional[str]], Any]]] = {
        UserDataType.string: lambda v: v,
        UserDataType.integer: _int_or_none,
        UserDataType.float: _float_or_none,
        UserDataType.date: _date_or_none,
        UserDataType.datetime: _datetime_or_none,
    }

    def _convert_value(self, value: Optional[str], user_type: UserDataType) -> Any:
        return self._PY_CONVERTERS_BY_BITYPE[user_type](value)

    def _convert_row(self, row: Sequence[Optional[str]], type_schema: Sequence[UserDataType]) -> list[Any]:
        return [self._convert_value(v, t) for v, t in zip(row, type_schema)]

    def _load_table_data(self, raw_csv_data: str, type_schema: Sequence[UserDataType]) -> list[list[Any]]:
        reader = csv.reader(io.StringIO(raw_csv_data))
        return [self._convert_row(row=row, type_schema=type_schema) for row in reader]

    def make_table_from_csv(
        self,
        raw_csv_data: str,
        table_schema: Sequence[tuple[str, UserDataType]],
        table_name_prefix: Optional[str] = None,
        nullable: bool = True,
    ) -> DbTable:
        table_name_prefix = table_name_prefix or "table_"
        if not table_name_prefix.endswith("_"):
            table_name_prefix = f"{table_name_prefix}_"
        table_name = f"{table_name_prefix}{shortuuid.uuid()}"

        type_schema = [user_type for col_name, user_type in table_schema]
        data = self._load_table_data(raw_csv_data=raw_csv_data, type_schema=type_schema)

        def _value_gen_factory(_col_idx: int) -> Callable[[int, datetime.datetime, random.Random], Any]:
            def _value_gen(rn: int, ts: datetime.datetime, rnd: random.Random) -> Any:
                return data[rn][_col_idx]

            return _value_gen

        columns = [
            C(name=name, user_type=user_type, vg=_value_gen_factory(_col_idx=col_idx), nullable=nullable)  # type: ignore
            for col_idx, (name, user_type) in enumerate(table_schema)
        ]

        db_table = make_table(db=self.db, name=table_name, columns=columns, rows=len(data))
        return db_table

lib/dl_core_testing/dl_core_testing/csv_table_dumper.py
datalens-tech/datalens-backend
"""
Helpers for fetching data from datasets
"""

from __future__ import annotations

from typing import (
    Collection,
    Optional,
)

import attr

from dl_constants.enums import (
    DataSourceRole,
    ProcessorType,
    SelectorType,
)
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.ids import AvatarId
from dl_core.data_processing.prepared_components.default_manager import DefaultPreparedComponentManager
from dl_core.data_processing.processing.operation import (
    CalcOp,
    DownloadOp,
    JoinOp,
)
from dl_core.data_processing.selectors.base import DataSelectorAsyncBase
from dl_core.data_processing.stream_base import (
    DataSourceVS,
    DataStream,
    DataStreamAsync,
)
from dl_core.query.bi_query import BIQuery
from dl_core.query.expression import JoinOnExpressionCtx
from dl_core.services_registry import ServicesRegistry
from dl_core.us_dataset import Dataset
from dl_core.us_manager.local_cache import USEntryBuffer
from dl_core.us_manager.us_manager import USManagerBase
from dl_utils.aio import (
    await_sync,
    to_sync_iterable,
)


@attr.s
class DataFetcher:
    _dataset: Dataset = attr.ib(kw_only=True)
    _selector: Optional[DataSelectorAsyncBase] = attr.ib(kw_only=True, default=None)
    _service_registry: Optional[ServicesRegistry] = attr.ib(kw_only=True, default=None)
    _selector_type: SelectorType = attr.ib(kw_only=True, default=SelectorType.CACHED)
    _us_manager: Optional[USManagerBase] = attr.ib(kw_only=True, default=None)  # FIXME: Legacy; remove
    _us_entry_buffer: USEntryBuffer = attr.ib(kw_only=True)
    _ds_accessor: DatasetComponentAccessor = attr.ib(init=False)

    @_us_entry_buffer.default
    def _make_us_entry_buffer(self) -> USEntryBuffer:
        # TODO: Remove with self._us_manager
        assert self._us_manager is not None
        return self._us_manager.get_entry_buffer()

    @_ds_accessor.default
    def _make_ds_accessor(self) -> DatasetComponentAccessor:
        return DatasetComponentAccessor(dataset=self._dataset)

    def __attrs_post_init__(self) -> None:
        if self._selector is None:
            assert self._service_registry is not None
            sel_factory = self._service_registry.get_selector_factory()
            self._selector = sel_factory.get_dataset_selector(
                dataset=self._dataset,
                selector_type=self._selector_type,
                us_entry_buffer=self._us_entry_buffer,
            )

    def _get_avatar_virtual_data_stream(
        self,
        avatar_id: AvatarId,
        stream_id: str,
        role: DataSourceRole,
        from_subquery: bool = False,
        subquery_limit: Optional[int] = None,
    ) -> DataSourceVS:
        alias = avatar_id
        prep_component_manager = DefaultPreparedComponentManager(
            dataset=self._dataset,
            role=role,
            us_entry_buffer=self._us_entry_buffer,
        )
        prep_src_info = prep_component_manager.get_prepared_source(
            avatar_id=avatar_id, alias=alias, from_subquery=from_subquery, subquery_limit=subquery_limit
        )
        return DataSourceVS(
            id=stream_id,
            alias=alias,
            result_id=avatar_id,
            names=prep_src_info.col_names,
            user_types=prep_src_info.user_types,
            prep_src_info=prep_src_info,
        )

    async def get_data_stream_async(
        self,
        *,
        bi_query: BIQuery,
        role: DataSourceRole = DataSourceRole.origin,
        row_count_hard_limit: Optional[int] = None,
        root_avatar_id: Optional[AvatarId] = None,
        required_avatar_ids: Optional[Collection[AvatarId]] = None,
        join_on_expressions: Collection[JoinOnExpressionCtx] = (),
        from_subquery: bool = False,
        subquery_limit: Optional[int] = None,
    ) -> DataStreamAsync:
        if root_avatar_id is None:
            root_avatar_id = self._ds_accessor.get_root_avatar_strict().id
        if required_avatar_ids is None:
            required_avatar_ids = bi_query.get_required_avatar_ids()

        dp_factory = self._service_registry.get_data_processor_factory()  # type: ignore  # TODO: fix
        data_processor = await dp_factory.get_data_processor(
            dataset=self._dataset,
            processor_type=ProcessorType.SOURCE_DB,
            selector_type=self._selector_type,
            role=role,
            us_entry_buffer=self._us_entry_buffer,
        )
        streams = [
            self._get_avatar_virtual_data_stream(
                avatar_id=avatar_id,
                stream_id=f"ava_{avatar_id}",
                role=role,
                from_subquery=from_subquery,
                subquery_limit=subquery_limit,
            )
            for avatar_id in sorted(required_avatar_ids)
        ]
        operations = [
            JoinOp(
                source_stream_ids={f"ava_{avatar_id}" for avatar_id in sorted(required_avatar_ids)},
                dest_stream_id="join_0",
                join_on_expressions=join_on_expressions,
                root_avatar_id=root_avatar_id,
            ),
            CalcOp(
                source_stream_id="join_0",
                dest_stream_id="calc_0",
                result_id="res",
                bi_query=bi_query,
                alias="res",
            ),
            DownloadOp(
                source_stream_id="calc_0",
                dest_stream_id="down_0",
                row_count_hard_limit=row_count_hard_limit,
            ),
        ]
        output_streams = await data_processor.run(
            streams=streams,
            operations=operations,
            output_stream_ids={"down_0"},
        )
        assert len(output_streams) == 1, f"Expected 1 stream, got {len(output_streams)}"
        data_stream = output_streams[0]
        return data_stream

    def get_data_stream(
        self,
        *,
        bi_query: BIQuery,
        role: DataSourceRole = DataSourceRole.origin,
        row_count_hard_limit: Optional[int] = None,
        root_avatar_id: Optional[AvatarId] = None,
        required_avatar_ids: Optional[Collection[AvatarId]] = None,
        join_on_expressions: Collection[JoinOnExpressionCtx] = (),
        from_subquery: bool = False,
        subquery_limit: Optional[int] = None,
    ) -> DataStream:
        async_data_stream = await_sync(
            self.get_data_stream_async(
                role=role,
                bi_query=bi_query,
                row_count_hard_limit=row_count_hard_limit,
                root_avatar_id=root_avatar_id,
                required_avatar_ids=required_avatar_ids,
                join_on_expressions=join_on_expressions,
                from_subquery=from_subquery,
                subquery_limit=subquery_limit,
            )
        )
        return DataStream(
            id=async_data_stream.id,
            names=async_data_stream.names,
            user_types=async_data_stream.user_types,
            meta=async_data_stream.meta,
            data=to_sync_iterable(async_data_stream.data.items),
            data_key=async_data_stream.data_key,
        )

lib/dl_core_testing/dl_core_testing/data.py
datalens-tech/datalens-backend
from __future__ import annotations

import json
from typing import (
    Dict,
    Iterable,
    List,
    Optional,
    Sequence,
    Sized,
)

from aiohttp import web
import attr
from marshmallow import (
    Schema,
    fields,
)

from dl_constants.types import TJSONLike


@attr.s
class SC:
    name: str = attr.ib()
    type: str = attr.ib(
        validator=attr.validators.in_(
            [
                "string",
                "integer",
                "float",
                "boolean",
                "date",
                "datetime",
            ]
        )
    )
    nullable: bool = attr.ib(default=False)
    title: Optional[str] = attr.ib(default=None)


@attr.s
class Table:
    name: str = attr.ib()
    schema: List[SC] = attr.ib()
    partition_key: List[str] = attr.ib(factory=lambda: [])
    meta: Dict[str, str] = attr.ib(factory=lambda: {})
    data: Iterable[Sequence[TJSONLike]] = attr.ib(factory=lambda: [])
    _data_length: Optional[int] = attr.ib(default=None)

    @property
    def data_length(self) -> int:
        if self._data_length is not None:
            return self._data_length

        if isinstance(self.data, Sized):
            return len(self.data)

        raise ValueError("Data length is unknown")


@attr.s
class TablesRegistry:
    tbl_list: List[Table] = attr.ib(factory=lambda: [])

    def get_table(self, table_name) -> Optional[Table]:  # type: ignore  # TODO: fix
        return next((t for t in self.tbl_list if t.name == table_name), None)


#
# Schemas
#
class ColumnSchema(Schema):
    name = fields.String()
    type = fields.String()
    nullable = fields.Bool()
    title = fields.String()


class TablePropertiesSchema(Schema):
    name = fields.String()
    schema = fields.Nested(ColumnSchema, many=True)
    partition_key = fields.List(fields.String())
    meta = fields.Dict(keys=fields.String(), values=fields.String())


def create_app(table_registry: TablesRegistry):  # type: ignore  # TODO: fix
    app = web.Application()
    app["TABLE_REGISTRY"] = table_registry

    def get_table_registry() -> TablesRegistry:
        return app["TABLE_REGISTRY"]

    def get_table_for_request(request: web.Request):  # type: ignore  # TODO: fix
        tbl = get_table_registry().get_table(request.match_info["table_name"])
        if tbl is None:
            raise web.HTTPNotFound()
        return tbl

    async def rq_table_data(request: web.Request):  # type: ignore  # TODO: fix
        table = get_table_for_request(request)
        resp = web.StreamResponse(status=200)
        resp.headers["X-DL-Row-Count"] = str(table.data_length)
        resp.headers["X-DL-Table-Definition"] = TablePropertiesSchema().dumps(table)
        resp.enable_chunked_encoding()
        await resp.prepare(request)

        for idx, data_line in enumerate(table.data):
            bytes_line = json.dumps(data_line, ensure_ascii=True).encode()
            if idx != 0:
                bytes_line = b"\n" + bytes_line

            await resp.write(bytes_line)

        await resp.write_eof()
        return resp

    async def rq_list_tables(_):  # type: ignore  # TODO: fix
        tr = get_table_registry()

        return web.json_response({"tables": [t.name for t in tr.tbl_list]})

    async def rq_table_properties(request: web.Request):  # type: ignore  # TODO: fix
        table = get_table_for_request(request)
        return web.json_response(TablePropertiesSchema().dump(table))

    app.router.add_post("/list_tables", rq_list_tables)
    app.router.add_post("/table_data/{table_name}", rq_table_data)
    app.router.add_post("/table_properties/{table_name}", rq_table_properties)
    return app

lib/dl_core_testing/dl_core_testing/data_adapter_mock.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import math
import random
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    NamedTuple,
    Optional,
    Sequence,
)
import uuid

import attr
import shortuuid
import sqlalchemy as sa
from sqlalchemy.sql.type_api import TypeEngine

from dl_constants.enums import (
    ConnectionType,
    UserDataType,
)
from dl_core.db import (
    get_type_transformer,
    make_sa_type,
)
from dl_db_testing.database.base import (
    DbBase,
    DbConfig,
    DbTableBase,
)
from dl_db_testing.database.dispenser import (
    DbDispenserBase,
    ReInitableDbDispenser,
)
from dl_db_testing.database.engine_wrapper import get_engine_wrapper_cls_for_url


if TYPE_CHECKING:
    from dl_core.db.conversion_base import TypeTransformer


@attr.s(frozen=True)
class CoreDbConfig(DbConfig):
    conn_type: ConnectionType = attr.ib(kw_only=True)


@attr.s
class Db(DbBase[CoreDbConfig]):
    """An executable interface"""

    @property
    def conn_type(self) -> ConnectionType:
        return self.config.conn_type

    def eval(
        self,
        expr: sa.sql.ClauseElement,
        from_: Optional[sa.sql.ClauseElement] = None,
        user_t: Optional[UserDataType] = None,
    ) -> Any:
        value = self.base_eval(expr, from_=from_)
        if user_t is not None:
            value = get_type_transformer(self.conn_type).cast_for_output(value, user_t=user_t)
        return value


def make_db_config(
    conn_type: ConnectionType,
    url: str,
    cluster: Optional[str] = None,
) -> CoreDbConfig:
    engine_config_kwargs: dict[str, Any] = {}
    if cluster is not None:
        engine_config_kwargs["cluster"] = cluster
    db_eng_config_cls = get_engine_wrapper_cls_for_url(url).CONFIG_CLS
    db_eng_config = db_eng_config_cls(url=url, **engine_config_kwargs)

    db_config = CoreDbConfig(
        engine_config=db_eng_config,
        conn_type=conn_type,
    )
    return db_config


def make_db_from_config(db_config: CoreDbConfig) -> Db:
    ew_cls = get_engine_wrapper_cls_for_url(db_config.engine_config.url)
    engine_wrapper = ew_cls(config=db_config.engine_config)
    return Db(config=db_config, engine_wrapper=engine_wrapper)


def make_db(
    conn_type: ConnectionType,
    url: str,
    cluster: Optional[str] = None,
) -> Db:
    # FIXME: Switch to accepting a fully ready config here
    db_config = make_db_config(conn_type=conn_type, url=url, cluster=cluster)
    return make_db_from_config(db_config)


class CoreDbDispenser(DbDispenserBase[CoreDbConfig, Db]):
    def make_database(self, db_config: CoreDbConfig) -> Db:
        return make_db_from_config(db_config)


class CoreReInitableDbDispenser(ReInitableDbDispenser[CoreDbConfig, Db], CoreDbDispenser):
    pass


@attr.s(frozen=True)
class DbTable(DbTableBase):
    db: Db = attr.ib(kw_only=True)  # Redefine type to Db from DbBase


@attr.s()
class C:
    name: str = attr.ib()
    user_type: UserDataType = attr.ib()
    nullable: Optional[bool] = attr.ib(default=None)
    _sa_type: TypeEngine = attr.ib(default=None)
    _vg: Callable[[int, datetime.datetime], Any] = attr.ib(default=None)

    DEFAULT_VALUE_GENERATORS = {
        UserDataType.string: lambda rn, **kwargs: f"str_value_{rn}",
        UserDataType.integer: lambda rn, **kwargs: rn,
        UserDataType.float: lambda rn, **kwargs: rn + (rn / 10),
        UserDataType.date: lambda rn, ts, **kwargs: ts.date() + datetime.timedelta(days=rn),
        UserDataType.datetime: lambda rn, ts, **kwargs: ts + datetime.timedelta(days=rn / math.pi),
        UserDataType.genericdatetime: lambda rn, ts, **kwargs: ts + datetime.timedelta(days=rn / math.pi),
        UserDataType.boolean: lambda rn, **kwargs: bool(int(rn) % 2),
        UserDataType.uuid: lambda rn, **kwargs: str(uuid.UUID(int=rn)),
        UserDataType.array_int: lambda rn, **kwargs: [rn * idx for idx in range(5)],
        UserDataType.array_str: lambda rn, **kwargs: [f"str_{str(rn * idx)}" for idx in range(5)],
        UserDataType.array_float: lambda rn, **kwargs: [float(rn * idx) * 1.1 for idx in range(5)],
    }

    @attr.s(auto_attribs=True, frozen=True)
    class ArrayDataGetter:
        _data_container: Sequence[Sequence[Any]]

        def __getitem__(self, col_idx: int) -> Callable[[int, datetime.datetime], Any]:
            return lambda rn, **kwargs: self._data_container[rn][col_idx]

    @property
    def vg(self) -> Callable[[int, datetime.datetime], Any]:
        if self._vg is not None:
            return self._vg
        try:
            return self.DEFAULT_VALUE_GENERATORS[self.user_type]
        except KeyError as e:
            raise ValueError(f"No default value generator for user type {e}")

    def get_sa_type(self, tt: TypeTransformer) -> TypeEngine:
        if self._sa_type is not None:
            return self._sa_type
        native_type = tt.type_user_to_native(user_t=self.user_type)
        return make_sa_type(native_type, nullable=self.nullable)

    @classmethod
    def array_data_getter(cls, data_container) -> "C.ArrayDataGetter":  # type: ignore  # TODO: fix
        return cls.ArrayDataGetter(data_container)

    @classmethod
    def int_value(cls, name: str = "int_value"):  # type: ignore  # TODO: fix
        return cls(name, UserDataType.integer)

    @classmethod
    def datetime_value(cls, name: str = "datetime_value"):  # type: ignore  # TODO: fix
        return cls(name, UserDataType.datetime)

    @classmethod
    def full_house(cls) -> list[C]:
        return [
            cls("string_value", UserDataType.string, nullable=False),
            cls("n_string_value", UserDataType.string, nullable=True),
            cls("int_value", UserDataType.integer, nullable=False),
            cls("n_int_value", UserDataType.integer, nullable=True),
            cls("float_value", UserDataType.float),
            cls("datetime_value", UserDataType.genericdatetime, nullable=False),
            cls("n_datetime_value", UserDataType.genericdatetime, nullable=True),
            cls("date_value", UserDataType.date),
            cls("boolean_value", UserDataType.boolean),
            cls("uuid_value", UserDataType.uuid),
            # Not included: arrays, geopoint, geopolygon, markup (not db-types, generally).
        ]

    @classmethod
    def array_columns(cls):  # type: ignore  # TODO: fix
        return [
            cls("array_int_value", UserDataType.array_int),
            cls("array_str_value", UserDataType.array_str),
            cls("array_float_value", UserDataType.array_float),
        ]


TEST_SEED = random.getrandbits(128)


def make_sample_data(
    columns: list[C] = None, rows: int = 10, seed: int = TEST_SEED, start_value: int = 0
) -> list[dict[str, Any]]:
    rnd = random.Random(seed)

    # 2017-01-01 .. 2020-04-17
    ts_uts = rnd.randrange(1483228800000000, 1587137765807968)
    ts = datetime.datetime(1970, 1, 1) + datetime.timedelta(seconds=ts_uts / 1000000)
    if columns is None:
        columns = C.full_house()
    result = [
        {col.name: col.vg(rn=idx, ts=ts, rnd=rnd) for col in columns}  # type: ignore  # TODO: fix
        for idx in range(start_value, start_value + rows)
    ]
    return result


# TODO FIX: Implement option to provide data via arguments
def make_table(
    db: Db,
    schema: Optional[str] = None,
    rows: int = 10,
    start_value: int = 0,
    columns: Optional[list[C]] = None,
    name: Optional[str] = None,
    data: Optional[list[dict[str, Any]]] = None,
    create_in_db: bool = True,
) -> DbTable:
    columns = columns or C.full_house()
    tt = get_type_transformer(db.conn_type)
    table = db.table_from_columns(
        [sa.Column(name=col.name, type_=col.get_sa_type(tt)) for col in columns],
        schema=schema,
        table_name=name,
    )
    db_table = DbTable(db=db, table=table, schema=schema)  # type: ignore  # TODO: fix
    # now = datetime.datetime.utcnow().replace(tzinfo=None, microsecond=0)
    if data is None:
        data = make_sample_data(columns=columns, rows=rows, start_value=start_value)
        data_reproduce = make_sample_data(columns=columns, rows=rows, start_value=start_value)
        assert data == data_reproduce, "should be reproducible"

    if create_in_db:
        db.create_table(table)
        db_table.insert(data)

    return db_table


def make_pg_table_with_enums(  # TODO: move to pg connector package
    db: Db,
    schema: Optional[str] = None,
    name: Optional[str] = None,
) -> DbTable:
    table = db.table_from_columns(
        [
            sa.Column(name=name, type_=type_)
            for name, type_ in [
                ("col1", sa.Enum("var1", "var2", name=str(uuid.uuid4()))),
                ("col2", sa.TEXT),
                ("col3", sa.INTEGER),
            ]
        ],
        schema=schema,
        table_name=name,
    )
    db.create_table(table)
    db_table = DbTable(db=db, table=table, schema=schema)  # type: ignore  # TODO: fix
    db_table.insert(
        [
            {"col1": "var1", "col2": "str1", "col3": 1},
            {"col1": "var1", "col2": "str2", "col3": 2},
            {"col1": "var2", "col2": "str3", "col3": 3},
        ]
    )
    return db_table


def make_table_with_arrays(
    db: Db,
    schema: Optional[str] = None,
    rows: int = 10,
    start_value: int = 0,
    columns: Optional[list[C]] = None,
    name: Optional[str] = None,
    data: Optional[list[dict[str, Any]]] = None,
) -> DbTable:
    if columns is None:
        columns = C.full_house() + C.array_columns()
    return make_table(db=db, schema=schema, rows=rows, start_value=start_value, columns=columns, name=name, data=data)


def make_schema(db: Db, schema_name: Optional[str] = None) -> str:
    schema_name = schema_name or f"sch_{shortuuid.uuid()}"
    assert schema_name is not None
    db.create_schema(schema_name=schema_name)
    return schema_name


class DbView(NamedTuple):
    db: Db
    name: str
    query: str
    schema: str = None  # type: ignore  # TODO: fix

    @property
    def as_table(self) -> DbTable:
        return DbTable(
            db=self.db,
            table=self.db.load_table(schema=self.schema, table_name=self.name),
            schema=self.schema,
            can_insert=False,
        )


def make_view(db: Db, query: str, schema: str = None, name: str = None) -> DbView:
    name = name or "test_view_{}".format(uuid.uuid4().hex[:10])
    db.create_view(query=query, schema=schema, name=name)
    return DbView(db=db, schema=schema, name=name, query=query)  # type: ignore  # TODO: fix


def make_view_from_table(db_table: DbTable, name: str = None, schema: str = None) -> DbView:
    return make_view(db=db_table.db, query=db_table.select_all_query, schema=schema or db_table.schema, name=name)


class MultiTables(NamedTuple):
    users: DbTable
    posts: DbTable
    comments: DbTable


def make_multiple_db_tables(db, string_kwargs):  # type: ignore  # TODO: fix
    """Basic table for all db types"""
    sa_tables = [
        db.table_from_columns(
            [  # users
                sa.Column(name="id", type_=sa.Integer()),
                sa.Column(name="name", type_=sa.String(**string_kwargs)),
            ]
        ),
        db.table_from_columns(
            [  # posts
                sa.Column(name="id", type_=sa.Integer()),
                sa.Column(name="author_id", type_=sa.Integer()),
                sa.Column(name="text", type_=sa.String(**string_kwargs)),
            ]
        ),
        db.table_from_columns(
            [  # comments
                sa.Column(name="id", type_=sa.Integer()),
                sa.Column(name="post_id", type_=sa.Integer()),
                sa.Column(name="from_id", type_=sa.Integer()),
                sa.Column(name="to_id", type_=sa.Integer()),
                sa.Column(name="text", type_=sa.String(**string_kwargs)),
            ]
        ),
    ]

    def create(table):  # type: ignore  # TODO: fix
        db.create_table(table)
        return DbTable(db=db, table=table)

    tables = MultiTables(*[create(t) for t in sa_tables])
    tables.users.insert(
        [
            {"id": 1, "name": "Clark Kent"},
            {"id": 2, "name": "Gaius Marius"},
            {"id": 3, "name": "Charlie Chaplin"},
            {"id": 4, "name": "Rosalind Franklin"},
        ]
    )
    tables.posts.insert(
        [
            {"id": 1, "author_id": 2, "text": "Let it be"},
            {"id": 2, "author_id": 4, "text": "Help"},
        ]
    )
    tables.comments.insert(
        [
            {"id": 1, "post_id": 1, "from_id": 1, "to_id": 2, "text": "first"},
            {"id": 2, "post_id": 1, "from_id": 2, "to_id": 1, "text": "second"},
            {"id": 3, "post_id": 1, "from_id": 3, "to_id": 2, "text": "third"},
            {"id": 4, "post_id": 2, "from_id": 3, "to_id": 4, "text": "first"},
            {"id": 5, "post_id": 2, "from_id": 4, "to_id": 3, "text": "second"},
            {"id": 6, "post_id": 2, "from_id": 1, "to_id": 4, "text": "third"},
        ]
    )

    return tables

lib/dl_core_testing/dl_core_testing/database.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
)
import uuid

from dl_constants.enums import (
    DataSourceCreatedVia,
    DataSourceRole,
    DataSourceType,
)
from dl_core.us_connection import get_connection_class
from dl_core.us_connection_base import ExecutorBasedMixin
from dl_core.us_dataset import Dataset
from dl_core_testing.dataset_wrappers import EditableDatasetTestWrapper


if TYPE_CHECKING:
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase
    from dl_core.us_connection_base import ConnectionBase
    from dl_core.us_manager.us_manager_sync import SyncUSManager
    from dl_core_testing.database import (
        Db,
        DbTable,
    )


def make_ds_key(*args) -> str:  # type: ignore  # TODO: fix
    return "/".join(["tests", *args])


def make_dataset(  # type: ignore  # TODO: fix
    sync_usm: SyncUSManager,
    connection: Optional[ConnectionBase] = None,
    db_table: Optional[DbTable] = None,
    schema_name: Optional[str] = None,
    table_name: Optional[str] = None,
    created_from: Optional[DataSourceType] = None,
    db_name: Optional[str] = None,
    yt_path: Optional[str] = None,
    yt_cluster: Optional[str] = None,
    ds_info: Optional[dict] = None,
    dsrc_params: Optional[dict] = None,
    created_via: Optional[DataSourceCreatedVia] = None,
):
    service_registry = sync_usm.get_services_registry()
    ds_info = dict(ds_info or {})
    db = db_table.db if db_table else None
    table_name = table_name or (db_table.name if db_table else None)
    schema_name = schema_name or (db_table.schema if db_table else None)
    db_name = db_name or (db.name if db else None)

    name = "Dataset {}".format(str(uuid.uuid4()))
    created_from = created_from or (get_created_from(db) if db is not None else None)
    dataset = Dataset.create_from_dict(
        Dataset.DataModel(name=name, **ds_info),
        ds_key=make_ds_key("datasets", name),
        us_manager=sync_usm,
    )
    ds_wrapper = EditableDatasetTestWrapper(dataset=dataset, us_manager=sync_usm)

    if created_via:
        ds_wrapper.set_created_via(created_via=created_via)

    if connection:
        dsrc_params = {
            **dict(
                db_name=db_name,
                schema_name=schema_name,
                table_name=table_name,
                path=yt_path,
                cluster=yt_cluster,
            ),
            **(dsrc_params or {}),
        }
        dsrc_params = {key: value for key, value in dsrc_params.items() if value is not None}
        source_id = str(uuid.uuid4())

        def conn_executor_factory() -> SyncConnExecutorBase:
            assert isinstance(connection, ExecutorBasedMixin)
            return service_registry.get_conn_executor_factory().get_sync_conn_executor(conn=connection)

        assert created_from is not None
        ds_wrapper.add_data_source(
            source_id=source_id,
            role=DataSourceRole.origin,
            connection_id=connection.uuid,
            created_from=created_from,
            parameters=dsrc_params,
        )
        sync_usm.load_dependencies(dataset)
        dsrc = ds_wrapper.get_data_source_strict(source_id=source_id, role=DataSourceRole.origin)
        ds_wrapper.update_data_source(
            source_id=source_id,
            role=DataSourceRole.origin,
            raw_schema=dsrc.get_schema_info(conn_executor_factory=conn_executor_factory).schema,
        )

        ds_wrapper.add_avatar(avatar_id=str(uuid.uuid4()), source_id=source_id, title="Main Avatar")

    return dataset


def get_created_from(db: Db) -> DataSourceType:
    conn_cls = get_connection_class(conn_type=db.conn_type)
    source_type = conn_cls.source_type
    assert source_type is not None
    return source_type


def data_source_settings_from_table(table: DbTable, db_name: Optional[str] = None) -> dict:
    source_type = get_created_from(db=table.db)
    data: dict[str, Any] = {  # this still requires connection_id to be defined
        "source_type": source_type,
        "parameters": {
            "table_name": table.name,
            "db_name": db_name,
        },
    }

    if table.schema:
        data["parameters"]["schema_name"] = table.schema

    return data

lib/dl_core_testing/dl_core_testing/dataset.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import uuid

import attr

from dl_constants.enums import (
    BinaryJoinOperator,
    DataSourceRole,
    DataSourceType,
)
from dl_core.base_models import DefaultConnectionRef
from dl_core.data_source.base import DataSource
from dl_core.multisource import (
    AvatarRelation,
    BinaryCondition,
    ConditionPartDirect,
    SourceAvatar,
)
from dl_core.services_registry.top_level import ServicesRegistry
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_dataset import Dataset
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import (
    Db,
    DbTable,
    make_table,
)
from dl_core_testing.dataset_wrappers import EditableDatasetTestWrapper


@attr.s(frozen=True)
class DataSourceCreationSpec:
    connection: ConnectionBase = attr.ib(kw_only=True)
    source_type: DataSourceType = attr.ib(kw_only=True)
    dsrc_params: dict = attr.ib(kw_only=True)


@attr.s
class DatasetSourceGenerator(abc.ABC):
    @abc.abstractmethod
    def generate_source_params(self) -> DataSourceCreationSpec:
        raise NotImplementedError


@attr.s
class DefaultDbDatasetSourceGenerator(DatasetSourceGenerator):
    db: Db = attr.ib(kw_only=True)
    connection: ConnectionBase = attr.ib(kw_only=True)

    def _get_table_source_params(self, db_table: DbTable) -> dict:
        return dict(
            db_name=db_table.db.name,
            schema_name=db_table.schema,
            table_name=db_table.name,
        )

    def generate_source_params(self) -> DataSourceCreationSpec:
        db_table = make_table(self.db)
        source_type = self.connection.source_type
        assert source_type is not None
        return DataSourceCreationSpec(
            connection=self.connection, source_type=source_type, dsrc_params=self._get_table_source_params(db_table)
        )


@attr.s(frozen=True)
class DatasetBuilderComponentProxy:
    dataset_builder: DatasetBuilder = attr.ib(kw_only=True)


@attr.s(frozen=True)
class DataSourceProxy(DatasetBuilderComponentProxy):
    source_id: str = attr.ib(kw_only=True)

    @property
    def target(self) -> DataSource:
        return self.dataset_builder.ds_wrapper.get_data_source_strict(
            source_id=self.source_id, role=DataSourceRole.origin
        )

    def add_avatar(self) -> SourceAvatarProxy:
        return self.dataset_builder.add_source_avatar(source_id=self.source_id)


@attr.s(frozen=True)
class SourceAvatarProxy(DatasetBuilderComponentProxy):
    avatar_id: str = attr.ib(kw_only=True)

    @property
    def target(self) -> SourceAvatar:
        return self.dataset_builder.ds_wrapper.get_avatar_strict(avatar_id=self.avatar_id)

    def add_relation(
        self,
        *,
        right: SourceAvatarProxy | str,
        conditions: list[BinaryCondition],
    ) -> AvatarRelationProxy:
        if isinstance(right, SourceAvatarProxy):
            right = right.avatar_id
        assert isinstance(right, str)
        return self.dataset_builder.add_avatar_relation(
            left_avatar_id=self.avatar_id,
            right_avatar_id=right,
            conditions=conditions,
        )

    def add_relation_simple_eq(
        self,
        *,
        right: SourceAvatarProxy | str,
        left_col_name: str,
        right_col_name: str,
    ) -> AvatarRelationProxy:
        if isinstance(right, SourceAvatarProxy):
            right = right.avatar_id
        assert isinstance(right, str)
        return self.dataset_builder.add_avatar_relation_simple_eq(
            left_avatar_id=self.avatar_id,
            right_avatar_id=right,
            left_col_name=left_col_name,
            right_col_name=right_col_name,
        )


@attr.s(frozen=True)
class AvatarRelationProxy(DatasetBuilderComponentProxy):
    relation_id: str = attr.ib(kw_only=True)

    @property
    def target(self) -> AvatarRelation:
        return self.dataset_builder.ds_wrapper.get_avatar_relation_strict(relation_id=self.relation_id)


@attr.s
class DatasetBuilder:
    service_registry: ServicesRegistry = attr.ib(kw_only=True)
    sync_us_manager: SyncUSManager = attr.ib(kw_only=True)
    dataset: Dataset = attr.ib(kw_only=True)
    dsrc_generator: DatasetSourceGenerator = attr.ib(kw_only=True, default=None)

    ds_wrapper: EditableDatasetTestWrapper = attr.ib(init=False)

    @ds_wrapper.default
    def _make_ds_wrapper(self) -> EditableDatasetTestWrapper:
        return EditableDatasetTestWrapper(
            dataset=self.dataset,
            us_entry_buffer=self.sync_us_manager.get_entry_buffer(),
        )

    def add_data_source(self, *, fill_raw_schema: bool = True) -> DataSourceProxy:
        source_id: str = str(uuid.uuid4())
        role = DataSourceRole.origin

        dsrc_creation_spec = self.dsrc_generator.generate_source_params()
        connection_id = dsrc_creation_spec.connection.uuid
        assert connection_id is not None
        self.sync_us_manager.ensure_entry_preloaded(DefaultConnectionRef(conn_id=connection_id))
        self.ds_wrapper.add_data_source(
            source_id=source_id,
            role=role,
            created_from=dsrc_creation_spec.source_type,
            connection_id=connection_id,
            parameters=dsrc_creation_spec.dsrc_params,
        )

        if fill_raw_schema:
            dsrc = self.ds_wrapper.get_data_source_strict(source_id=source_id, role=role)
            ce_factory = self.service_registry.get_conn_executor_factory()
            conn_executor = ce_factory.get_sync_conn_executor(conn=dsrc_creation_spec.connection)
            raw_schema = dsrc.get_schema_info(conn_executor_factory=lambda: conn_executor).schema
            assert raw_schema is not None
            self.ds_wrapper.update_data_source(source_id=source_id, role=role, raw_schema=raw_schema)

        return DataSourceProxy(
            dataset_builder=self,
            source_id=source_id,
        )

    def add_source_avatar(self, source_id: str) -> SourceAvatarProxy:
        avatar_id = str(uuid.uuid4())
        avatar_title = str(uuid.uuid4())
        self.ds_wrapper.add_avatar(avatar_id=avatar_id, source_id=source_id, title=avatar_title)
        return SourceAvatarProxy(
            dataset_builder=self,
            avatar_id=avatar_id,
        )

    def add_avatar_relation(
        self,
        left_avatar_id: str,
        right_avatar_id: str,
        conditions: list[BinaryCondition],
    ) -> AvatarRelationProxy:
        relation_id = str(uuid.uuid4())
        self.ds_wrapper.add_avatar_relation(
            relation_id=relation_id,
            left_avatar_id=left_avatar_id,
            right_avatar_id=right_avatar_id,
            conditions=conditions,
        )
        return AvatarRelationProxy(
            dataset_builder=self,
            relation_id=relation_id,
        )

    def add_avatar_relation_simple_eq(
        self,
        left_avatar_id: str,
        right_avatar_id: str,
        left_col_name: str,
        right_col_name: str,
    ) -> AvatarRelationProxy:
        conditions = [
            BinaryCondition(
                operator=BinaryJoinOperator.eq,
                left_part=ConditionPartDirect(source=left_col_name),
                right_part=ConditionPartDirect(source=right_col_name),
            ),
        ]
        return self.add_avatar_relation(
            left_avatar_id=left_avatar_id,
            right_avatar_id=right_avatar_id,
            conditions=conditions,
        )


@attr.s
class DatasetBuilderFactory(abc.ABC):
    @abc.abstractmethod
    def get_dataset_builder(self, dataset: Dataset) -> DatasetBuilder:
        raise NotImplementedError


@attr.s
class DefaultDbDatasetBuilderFactory(DatasetBuilderFactory):
    db: Db = attr.ib(kw_only=True)
    service_registry: ServicesRegistry = attr.ib(kw_only=True)
    sync_us_manager: SyncUSManager = attr.ib(kw_only=True)
    connection: ConnectionBase = attr.ib(kw_only=True)

    def _get_dsrc_generator(self) -> DatasetSourceGenerator:
        return DefaultDbDatasetSourceGenerator(db=self.db, connection=self.connection)

    def get_dataset_builder(self, dataset: Dataset) -> DatasetBuilder:
        return DatasetBuilder(
            dataset=dataset,
            dsrc_generator=self._get_dsrc_generator(),
            service_registry=self.service_registry,
            sync_us_manager=self.sync_us_manager,
        )

lib/dl_core_testing/dl_core_testing/dataset_builder.py
datalens-tech/datalens-backend
from __future__ import annotations

from itertools import chain
from typing import (
    Any,
    Callable,
    FrozenSet,
    Iterable,
    List,
    Optional,
)

import attr

from dl_constants.enums import (
    DataSourceCreatedVia,
    DataSourceRole,
    DataSourceType,
    JoinType,
    ManagedBy,
)
from dl_core.base_models import (
    DefaultWhereClause,
    ObligatoryFilter,
)
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.editor import DatasetComponentEditor
from dl_core.connection_executors.sync_base import SyncConnExecutorBase
from dl_core.data_source.base import DataSource
from dl_core.data_source.collection import (
    DataSourceCollectionBase,
    DataSourceCollectionFactory,
)
from dl_core.data_source.sql import BaseSQLDataSource
from dl_core.data_source_spec.collection import DataSourceCollectionSpecBase
from dl_core.dataset_capabilities import DatasetCapabilities
from dl_core.db.elements import (
    IndexInfo,
    SchemaColumn,
)
from dl_core.fields import (
    BIField,
    ResultSchema,
)
from dl_core.multisource import (
    AvatarRelation,
    BinaryCondition,
    SourceAvatar,
)
from dl_core.us_dataset import Dataset
from dl_core.us_manager.local_cache import USEntryBuffer
from dl_core.us_manager.us_manager import USManagerBase


@attr.s
class DatasetTestWrapper:
    """
    A helper class for tests.
    It consolidates the interfaces of various dataset wrappers
    and factories.

    Fo use in tests only!
    """

    _dataset: Dataset = attr.ib()
    _us_manager: Optional[USManagerBase] = attr.ib(kw_only=True, default=None)  # FIXME: Legacy; remove
    _us_entry_buffer: USEntryBuffer = attr.ib(kw_only=True)

    _ds_accessor: DatasetComponentAccessor = attr.ib(init=False)
    _dsrc_coll_factory: DataSourceCollectionFactory = attr.ib(init=False)
    _ds_capabilities: DatasetCapabilities = attr.ib(init=False)

    @_us_entry_buffer.default
    def _make_us_entry_buffer(self) -> USEntryBuffer:
        # TODO: Remove with self._us_manager
        assert self._us_manager is not None
        return self._us_manager.get_entry_buffer()

    @_ds_accessor.default
    def _make_ds_accessor(self) -> DatasetComponentAccessor:
        return DatasetComponentAccessor(dataset=self._dataset)

    @_dsrc_coll_factory.default
    def _make_dsrc_coll_factory(self) -> DataSourceCollectionFactory:
        return DataSourceCollectionFactory(us_entry_buffer=self._us_entry_buffer)

    @_ds_capabilities.default
    def _make_ds_capabilities(self) -> DatasetCapabilities:
        return DatasetCapabilities(dataset=self._dataset, dsrc_coll_factory=self._dsrc_coll_factory)

    @property
    def dsrc_coll_factory(self) -> DataSourceCollectionFactory:
        return self._dsrc_coll_factory

    def get_root_avatar_opt(self) -> Optional[SourceAvatar]:
        return self._ds_accessor.get_root_avatar_opt()

    def get_root_avatar_strict(self) -> SourceAvatar:
        return self._ds_accessor.get_root_avatar_strict()

    def get_avatar_list(self, source_id: Optional[str] = None) -> list[SourceAvatar]:
        return self._ds_accessor.get_avatar_list(source_id=source_id)

    def get_avatar_opt(self, avatar_id: str) -> Optional[SourceAvatar]:
        return self._ds_accessor.get_avatar_opt(avatar_id=avatar_id)

    def get_avatar_strict(self, avatar_id: str) -> SourceAvatar:
        return self._ds_accessor.get_avatar_strict(avatar_id=avatar_id)

    def get_data_source_coll_spec_opt(self, source_id: str) -> Optional[DataSourceCollectionSpecBase]:
        return self._ds_accessor.get_data_source_coll_spec_opt(source_id=source_id)

    def get_data_source_coll_spec_strict(self, source_id: str) -> DataSourceCollectionSpecBase:
        return self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)

    def get_data_source_coll_list(self) -> list[DataSourceCollectionBase]:
        return [self.get_data_source_coll_strict(source_id=source_id) for source_id in self.get_data_source_id_list()]

    def get_data_source_coll_opt(self, source_id: str) -> Optional[DataSourceCollectionBase]:
        dsrc_coll_spec = self._ds_accessor.get_data_source_coll_spec_strict(source_id=source_id)
        dsrc_coll = self._dsrc_coll_factory.get_data_source_collection(spec=dsrc_coll_spec)
        return dsrc_coll

    def get_data_source_coll_strict(self, source_id: str) -> DataSourceCollectionBase:
        dsrc_coll = self.get_data_source_coll_opt(source_id=source_id)
        assert dsrc_coll is not None
        return dsrc_coll

    def get_data_source_id_list(self) -> list[str]:
        return self._ds_accessor.get_data_source_id_list()

    def get_data_source_list(self, role: DataSourceRole) -> list[DataSource]:
        return [
            self.get_data_source_strict(source_id=source_id, role=role) for source_id in self.get_data_source_id_list()
        ]

    def get_data_source_opt(self, source_id: str, role: DataSourceRole) -> Optional[DataSource]:
        assert role is not None
        dsrc_coll = self.get_data_source_coll_strict(source_id=source_id)
        dsrc = dsrc_coll.get_opt(role=role)
        return dsrc

    def get_data_source_strict(self, source_id: str, role: DataSourceRole) -> DataSource:
        assert role is not None
        dsrc_coll = self.get_data_source_coll_strict(source_id=source_id)
        dsrc = dsrc_coll.get_strict(role=role)
        return dsrc

    def get_sql_data_source_strict(self, source_id: str, role: DataSourceRole) -> BaseSQLDataSource:
        dsrc = self.get_data_source_strict(source_id=source_id, role=role)
        assert isinstance(dsrc, BaseSQLDataSource)
        return dsrc

    def get_avatar_relation_list(
        self, left_avatar_id: Optional[str] = None, right_avatar_id: Optional[str] = None
    ) -> list[AvatarRelation]:
        return self._ds_accessor.get_avatar_relation_list(
            left_avatar_id=left_avatar_id,
            right_avatar_id=right_avatar_id,
        )

    def get_avatar_relation_strict(
        self,
        relation_id: Optional[str] = None,
        left_avatar_id: Optional[str] = None,
        right_avatar_id: Optional[str] = None,
    ) -> AvatarRelation:
        return self._ds_accessor.get_avatar_relation_strict(
            relation_id=relation_id,
            left_avatar_id=left_avatar_id,
            right_avatar_id=right_avatar_id,
        )

    def get_obligatory_filter_list(self) -> list[ObligatoryFilter]:
        return self._ds_accessor.get_obligatory_filter_list()

    def get_obligatory_filter_opt(
        self,
        obfilter_id: Optional[str] = None,
        field_guid: Optional[str] = None,
    ) -> Optional[ObligatoryFilter]:
        return self._ds_accessor.get_obligatory_filter_opt(obfilter_id=obfilter_id, field_guid=field_guid)

    def get_obligatory_filter_strict(
        self,
        obfilter_id: Optional[str] = None,
        field_guid: Optional[str] = None,
    ) -> ObligatoryFilter:
        return self._ds_accessor.get_obligatory_filter_strict(obfilter_id=obfilter_id, field_guid=field_guid)

    def get_cached_raw_schema(self, role: DataSourceRole) -> list[SchemaColumn]:
        return list(
            chain.from_iterable(
                self.get_data_source_coll_strict(source_id=source_id).get_cached_raw_schema(role=role) or ()
                for source_id in self._ds_accessor.get_data_source_id_list()
            )
        )

    def get_new_raw_schema(
        self,
        role: DataSourceRole,
        conn_executor_factory: Callable[[], SyncConnExecutorBase],
    ) -> list[SchemaColumn]:
        """
        Return raw schema of all data sources
        """
        return list(
            chain.from_iterable(
                (
                    self.get_data_source_coll_strict(source_id=source_id).get_raw_schema(
                        role=role,
                        conn_executor_factory=conn_executor_factory,
                    )
                    or ()
                )
                for source_id in self._ds_accessor.get_data_source_id_list()
            )
        )

    def quote(self, value: str, role: DataSourceRole) -> str:
        sql_dsrc = self.get_sql_data_source_strict(source_id=self._dataset.get_single_data_source_id(), role=role)
        dialect = sql_dsrc.get_dialect()
        return dialect.identifier_preparer.quote(value)

    def resolve_source_role(self, for_preview: bool = False) -> DataSourceRole:
        return self._ds_capabilities.resolve_source_role(for_preview=for_preview)


@attr.s
class EditableDatasetTestWrapper(DatasetTestWrapper):
    _ds_editor: DatasetComponentEditor = attr.ib(init=False)

    @_ds_editor.default
    def _make_ds_editor(self) -> DatasetComponentEditor:
        return DatasetComponentEditor(dataset=self._dataset)

    def update_data_source_collection(
        self,
        source_id: str,
        title: Optional[str] = None,
        valid: Optional[bool] = None,
    ) -> None:
        self._ds_editor.update_data_source_collection(source_id=source_id, title=title, valid=valid)

    def remove_data_source_collection(self, source_id: str) -> None:
        self._ds_editor.remove_data_source_collection(source_id=source_id)

    def add_data_source(
        self,
        *,
        source_id: str,
        role: DataSourceRole = DataSourceRole.origin,
        created_from: DataSourceType,
        connection_id: Optional[str] = None,
        title: Optional[str] = None,
        raw_schema: Optional[list[SchemaColumn]] = None,
        index_info_set: FrozenSet[IndexInfo] = None,
        managed_by: Optional[ManagedBy] = None,
        parameters: Optional[dict[str, Any]] = None,
    ) -> None:
        self._ds_editor.add_data_source(
            source_id=source_id,
            role=role,
            created_from=created_from,
            connection_id=connection_id,
            title=title,
            raw_schema=raw_schema,
            index_info_set=index_info_set,
            managed_by=managed_by,
            parameters=parameters,
        )

    def update_data_source(
        self,
        source_id: str,
        role: Optional[DataSourceRole] = None,
        connection_id: Optional[str] = None,
        created_from: Optional[DataSourceType] = None,
        raw_schema: Optional[list] = None,
        index_info_set: FrozenSet[IndexInfo] = None,
        **parameters: Any,
    ) -> None:
        self._ds_editor.update_data_source(
            source_id=source_id,
            role=role,
            connection_id=connection_id,
            created_from=created_from,
            raw_schema=raw_schema,
            index_info_set=index_info_set,
            **parameters,
        )

    def remove_data_source(self, source_id: str, role: DataSourceRole) -> None:
        self._ds_editor.remove_data_source(source_id=source_id, role=role)

    def add_avatar(
        self,
        avatar_id: str,
        source_id: str,
        title: str,
        managed_by: Optional[ManagedBy] = None,
        valid: bool = True,
    ) -> None:
        return self._ds_editor.add_avatar(
            avatar_id=avatar_id,
            source_id=source_id,
            title=title,
            managed_by=managed_by,
            valid=valid,
        )

    def update_avatar(
        self,
        avatar_id: str,
        source_id: Optional[str] = None,
        title: Optional[str] = None,
        valid: Optional[bool] = None,
    ) -> None:
        return self._ds_editor.update_avatar(
            avatar_id=avatar_id,
            source_id=source_id,
            title=title,
            valid=valid,
        )

    def remove_avatar(self, avatar_id: str) -> None:
        self._ds_editor.remove_avatar(avatar_id=avatar_id)

    def add_avatar_relation(
        self,
        relation_id: str,
        left_avatar_id: str,
        right_avatar_id: str,
        conditions: List[BinaryCondition],
        join_type: JoinType = None,
        managed_by: ManagedBy = None,
        valid: bool = True,
    ) -> None:
        self._ds_editor.add_avatar_relation(
            relation_id=relation_id,
            left_avatar_id=left_avatar_id,
            right_avatar_id=right_avatar_id,
            conditions=conditions,
            join_type=join_type,
            managed_by=managed_by,
            valid=valid,
        )

    def update_avatar_relation(
        self,
        relation_id: str,
        conditions: Optional[List[BinaryCondition]] = None,
        join_type: Optional[JoinType] = None,
        valid: Optional[bool] = None,
    ) -> None:
        self._ds_editor.update_avatar_relation(
            relation_id=relation_id,
            conditions=conditions,
            join_type=join_type,
            valid=valid,
        )

    def remove_avatar_relation(self, relation_id: str) -> None:
        self._ds_editor.remove_avatar_relation(relation_id=relation_id)

    def add_obligatory_filter(
        self,
        obfilter_id: str,
        field_guid: str,
        default_filters: list[DefaultWhereClause],
        managed_by: Optional[ManagedBy] = None,
        valid: bool = True,
    ) -> None:
        self._ds_editor.add_obligatory_filter(
            obfilter_id=obfilter_id,
            field_guid=field_guid,
            default_filters=default_filters,
            managed_by=managed_by,
            valid=valid,
        )

    def update_obligatory_filter(
        self,
        obfilter_id: str,
        default_filters: Optional[list[DefaultWhereClause]] = None,
        valid: Optional[bool] = None,
    ) -> None:
        self._ds_editor.update_obligatory_filter(
            obfilter_id=obfilter_id,
            default_filters=default_filters,
            valid=valid,
        )

    def remove_obligatory_filter(self, obfilter_id: str) -> None:
        self._ds_editor.remove_obligatory_filter(obfilter_id=obfilter_id)

    def set_result_schema(self, result_schema: ResultSchema | Iterable[BIField]) -> None:
        self._ds_editor.set_result_schema(result_schema=result_schema)

    def set_created_via(self, created_via: DataSourceCreatedVia) -> None:
        self._ds_editor.set_created_via(created_via=created_via)

lib/dl_core_testing/dl_core_testing/dataset_wrappers.py
datalens-tech/datalens-backend
from dl_db_testing.database.engine_wrapper import EngineWrapperBase


class TestingEngineWrapper(EngineWrapperBase):
    """
    Dummy engine wrapper to use in tests without a proper DB and an sqlalchemy dialect
    """

    URL_PREFIX = "bi_testing"

    def __attrs_post_init__(self) -> None:
        pass

lib/dl_core_testing/dl_core_testing/engine_wrapper.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import os
import time
from typing import Optional

import requests

from dl_core.logging_config import setup_jaeger_client
from dl_core_testing.configuration import UnitedStorageConfiguration
from dl_db_testing.loader import load_db_testing_lib
from dl_utils.wait import wait_for


LOGGER = logging.getLogger(__name__)


def _wait_for_pg(dsn: str, timeout: int = 300, interval: float = 1.0) -> None:
    import psycopg2

    start_time = time.monotonic()
    max_time = start_time + timeout
    while True:
        loop_start_time = time.monotonic()
        try:
            conn = psycopg2.connect(dsn)
            cur = conn.cursor()
            cur.execute("select 1./2")
            res = cur.fetchall()
        except Exception as exc:
            now = time.monotonic()
            if now > max_time:
                raise  # timed out
            next_try_time = loop_start_time + interval
            sleep_time = next_try_time - now
            if sleep_time > 0.001:
                LOGGER.debug("Waiting for pg to come up (%r)", exc)
                time.sleep(sleep_time)
        else:
            if res == [(0.5,)]:
                return
            else:
                raise Exception("Unexpected result", res)


def restart_container(container_name: str) -> None:
    import docker

    docker_cli = docker.from_env(timeout=300)
    container = docker_cli.containers.list(filters=dict(name=container_name), all=True)[0]
    container.restart()


def restart_container_by_label(label: str, compose_project: str) -> None:
    import docker

    docker_cli = docker.from_env(timeout=300)
    container = docker_cli.containers.list(
        filters=dict(
            label=[
                f"datalens.ci.service={label}",
                f"com.docker.compose.project={compose_project}",
            ]
        ),
        all=True,
    )[0]
    container.restart()


def prepare_united_storage(
    *,
    us_host: str,
    us_master_token: str,
    tenant_id: str = "common",
    us_pg_dsn: Optional[str] = None,
    force: bool = False,
) -> None:
    if not force and not os.environ.get("CLEAR_US_DATABASE", ""):
        LOGGER.debug("prepare_united_storage: CLEAR_US_DATABASE env is disabled, skipping.")
        return

    if us_pg_dsn is not None:
        LOGGER.debug("prepare_united_storage: wait for pg-us to be up...")
        _wait_for_pg(us_pg_dsn)

    headers = {
        "X-US-Master-Token": us_master_token,
        "X-DL-TenantId": tenant_id,
    }

    def _wait_for_us() -> tuple[bool, str]:
        try:
            with requests.Session() as reqr:
                resp = reqr.get(f"{us_host}/ping-db", headers=headers)
                resp.raise_for_status()
                return True, ""
        except Exception as e:
            return False, str(e)

    max_wait_time = 160.0
    wait_pause = 0.5
    LOGGER.debug(f"prepare_united_storage: waiting for up to {max_wait_time}s for US to respond with ok...")
    wait_for("US startup", condition=_wait_for_us, timeout=max_wait_time, interval=wait_pause)

    LOGGER.debug("prepare_united_storage: done.")


def prepare_united_storage_from_config(us_config: UnitedStorageConfiguration) -> None:
    prepare_united_storage(
        us_host=us_config.us_host,
        us_master_token=us_config.us_master_token,
        us_pg_dsn=us_config.us_pg_dsn,
        force=us_config.force,
    )


def common_pytest_configure(
    use_jaeger_tracer: bool = False,
    tracing_service_name: str = "tests",
) -> None:
    load_db_testing_lib()

    if use_jaeger_tracer:
        setup_jaeger_client(tracing_service_name)

lib/dl_core_testing/dl_core_testing/environment.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Dict,
    Type,
)

import attr

from dl_core.connection_executors.adapters.common_base import CommonBaseDirectAdapter
from dl_core.connection_executors.models.connection_target_dto_base import BaseSQLConnTargetDTO
from dl_core_testing.database import Db


@attr.s
class ExecutorFactoryBase:
    db: Db = attr.ib()

    def get_dto_class(self) -> Type[BaseSQLConnTargetDTO]:
        raise NotImplementedError

    def get_dba_class(self) -> Type[CommonBaseDirectAdapter]:
        return CommonBaseDirectAdapter

    def get_dto_kwargs(self) -> Dict[str, Any]:
        return dict(
            conn_id=None,
            host=self.db.url.host,
            port=self.db.url.port,
            db_name=self.db.url.database,
            username=self.db.url.username,
            password=self.db.url.password,
        )

    def make_dto(self) -> BaseSQLConnTargetDTO:
        dto_cls = self.get_dto_class()
        return dto_cls(**self.get_dto_kwargs())

lib/dl_core_testing/dl_core_testing/executors.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import multiprocessing
import os
import signal
import socket
import subprocess
import sys
import time
from typing import (
    Dict,
    Optional,
    TypeVar,
)

import attr
import requests
from statcommons.log_config import deconfigure_logging


LOGGER = logging.getLogger()

_RUNNER_TV = TypeVar("_RUNNER_TV")


class ForkPopenHack(subprocess.Popen):
    """Execute `func` in a forked process"""

    def __init__(self, target, **kwargs):  # type: ignore  # TODO: fix
        self._target = target
        super().__init__(args=(), **kwargs)  # type: ignore  # TODO: fix

    def _execute_child(self, *args, **kwargs):  # type: ignore  # TODO: fix
        # Most of the arguments are ignored,
        # unix-only version
        # skipping: errpipe.
        pid = os.fork()
        if not pid:  # child
            return self._execute_child_code()
        # parent
        return self._manage_child(pid)

    def _manage_child(self, pid):  # type: ignore  # TODO: fix
        self.pid = pid

    def _execute_child_code(self):  # type: ignore  # TODO: fix
        try:
            self._target()
        finally:
            # HACK: Prevent cleanup of parent's state using suicide.
            os.kill(os.getpid(), signal.SIGKILL)


class ProcessPopenDuck(multiprocessing.Process):
    """`multiprocessing.Process` with some duck-compatibility with `subprocess.Popen`"""

    _retcode = 0

    def __init__(self, *args, **kwargs):  # type: ignore  # TODO: fix
        super().__init__(*args, **kwargs)
        self.start()

    def wait(self, timeout):  # type: ignore  # TODO: fix
        self.join(timeout)
        return self._retcode

    @property
    def returncode(self):  # type: ignore  # TODO: fix
        if self.is_alive():
            return None
        return self._retcode


# TODO FIX: Communicate from thread and fetch STDOUT/STDERR to logs
@attr.s
class WSGIRunner:
    _module: str = attr.ib()
    _callable: str = attr.ib()
    _ping_path: str = attr.ib()
    _bind_port: Optional[int] = attr.ib(default=None)
    _bind_addr: str = attr.ib(default="127.0.0.1")
    _wait_time: float = attr.ib(default=15.0)
    _poll_time: float = attr.ib(default=0.1)
    _wait_term_time: float = attr.ib(default=5.0)
    _env: Optional[Dict[str, str]] = attr.ib(default=None)

    _proc: subprocess.Popen = attr.ib(init=False, default=None)

    def _debug(self, message, *args):  # type: ignore  # TODO: fix
        sys.stderr.write("{} @ {}: ".format(self.__class__.__name__, os.getpid()))
        sys.stderr.write(message % args)
        sys.stderr.write("\n")
        sys.stderr.flush()

    @property
    def bind_addr(self) -> str:
        return self._bind_addr

    @property
    def bind_port(self) -> int:
        if self._bind_port is not None:
            return self._bind_port
        raise ValueError("Bind port was not provided and was not generated yet")

    def find_free_port(self) -> int:
        LOGGER.debug("Finding an available port...")
        # Binding on random port
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.bind((self._bind_addr, 0))
        port = s.getsockname()[1]
        s.close()
        LOGGER.debug("Found port %s", port)
        # And hope that nobody bind on it until Flask will bind
        return port

    # TODO FIX: Logging
    def wait_for_up(self) -> None:
        url = f"http://{self._bind_addr}:{self._bind_port}/{self._ping_path.lstrip('/')}"
        start_time = time.monotonic()
        max_time = start_time + self._wait_time
        last_error = None
        resp = None

        while True:
            if self._proc.returncode is not None:
                # TODO FIX: add output
                raise Exception("WSGI process down")

            next_attempt_time = time.monotonic() + self._poll_time
            try:
                resp = requests.get(url, timeout=self._poll_time)
            except Exception as err:
                last_error = err
            else:
                if resp.ok:
                    LOGGER.debug("WSGI app ok on url: %s", url)
                    break
            LOGGER.debug("Still waiting for WSGI app: err=%r / resp=%r", last_error, resp)
            now = time.monotonic()
            if now > max_time:
                raise Exception(
                    f"Timed out waiting for WSGI to come up at {url}",
                    dict(last_error=last_error, last_resp=resp, last_resp_text=resp.text if resp is not None else None),
                )
            sleep_time = next_attempt_time - now
            if sleep_time > 0.001:
                time.sleep(sleep_time)

    def run(self) -> None:
        if self._bind_port is None:
            self._bind_port = self.find_free_port()
        self._run_subproc()
        self.wait_for_up()

    def _make_uwsgi_params(self):  # type: ignore  # TODO: fix
        return (
            "--die-on-term",
            "--master",
            "--module",
            self._module,
            "--callable",
            self._callable,
            "--http",
            f"{self._bind_addr}:{self._bind_port}",
            "--http-timeout",
            "600",
            "--socket-timeout",
            "600",
            # TODO FIX: Figure out why server does not starts within 15 seconds with this settings
            # '--harakiri', '300',
            # '--harakiri-verbose', 'true',
            # '--workers', '16',
            # '--threads', '10',
            "--workers",
            "3",
            "--threads",
            "3",
        )

    def _run_subproc(self):  # type: ignore  # TODO: fix
        cmd = ["uwsgi"] + list(self._make_uwsgi_params())
        env = {
            **(self._env or {}),
            **os.environ,
        }
        self._proc = subprocess.Popen(cmd, env=env)

    def _run_fork_child_code(self):  # type: ignore  # TODO: fix
        self._debug("child: running uwsgi")
        os.environ.update(self._env)  # type: ignore  # TODO: fix
        import pyuwsgi  # noqa

        cmd = [sys.argv[0]] + list(self._make_uwsgi_params())
        sys.argv = cmd
        self._debug("child: cmd: %r", cmd)
        deconfigure_logging()  # so that postfork configure_logging doesn't fail
        try:
            pyuwsgi.run()
        finally:
            self._debug("child: uwsgi done.")

    def shutdown(self):  # type: ignore  # TODO: fix
        try:
            self._proc.terminate()
            exit_code = self._proc.wait(self._wait_term_time)
            LOGGER.debug("Test fixture worker on port %s was terminated. Exit code: %s", self._bind_port, exit_code)
        except Exception:  # noqa
            LOGGER.exception("Exception during test worker graceful shutdown. Going to kill...")
            self._proc.kill()

    def __enter__(self: _RUNNER_TV) -> _RUNNER_TV:
        self.run()  # type: ignore  # TODO: fix
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):  # type: ignore
        self.shutdown()

lib/dl_core_testing/dl_core_testing/fixture_server_runner.py
datalens-tech/datalens-backend
import pkgutil

import attr

from dl_core_testing.csv_table_dumper import CsvTableDumper
from dl_core_testing.database import (
    Db,
    DbConfig,
    DbTable,
)
from dl_core_testing.fixtures.primitives import FixtureTableSpec


@attr.s
class DbCsvTableDispenser:
    _tables: dict[DbConfig, dict[FixtureTableSpec, DbTable]] = attr.ib(init=False, factory=dict)

    def _get_raw_csv_data(self, path: str) -> str:
        byte_data = pkgutil.get_data(__name__, path)
        assert byte_data is not None
        return byte_data.decode()

    def _make_new_csv_table(self, db: Db, spec: FixtureTableSpec) -> DbTable:
        dumper = CsvTableDumper(db=db)
        db_table = dumper.make_table_from_csv(
            raw_csv_data=self._get_raw_csv_data(spec.csv_name),
            table_schema=spec.table_schema,
            nullable=spec.nullable,
        )
        if db.config not in self._tables:
            self._tables[db.config] = {}
        self._tables[db.config][spec] = db_table
        return db_table

    def get_csv_table(self, db: Db, spec: FixtureTableSpec) -> DbTable:
        try:
            return self._tables[db.config][spec]
        except KeyError:
            return self._make_new_csv_table(db=db, spec=spec)

lib/dl_core_testing/dl_core_testing/fixtures/dispenser.py
datalens-tech/datalens-backend
import attr

from dl_constants.enums import UserDataType


@attr.s(frozen=True)
class FixtureTableSpec:
    csv_name: str = attr.ib(kw_only=True)
    table_schema: tuple[tuple[str, UserDataType], ...] = attr.ib(kw_only=True)
    nullable: bool = attr.ib(kw_only=True, default=True)

    def get_user_type_for_col(self, col_name: str) -> UserDataType:
        return next(tbl_col[1] for tbl_col in self.table_schema if tbl_col[0] == col_name)

lib/dl_core_testing/dl_core_testing/fixtures/primitives.py
datalens-tech/datalens-backend
from dl_constants.enums import UserDataType
from dl_core_testing.fixtures.primitives import FixtureTableSpec


TABLE_SPEC_SAMPLE_SUPERSTORE = FixtureTableSpec(
    csv_name="sample_superstore.csv",
    table_schema=(
        ("category", UserDataType.string),
        ("city", UserDataType.string),
        ("country", UserDataType.string),
        ("customer_id", UserDataType.string),
        ("customer_name", UserDataType.string),
        ("discount", UserDataType.float),
        ("order_date", UserDataType.date),
        ("order_id", UserDataType.string),
        ("postal_code", UserDataType.integer),
        ("product_id", UserDataType.string),
        ("product_name", UserDataType.string),
        ("profit", UserDataType.float),
        ("quantity", UserDataType.integer),
        ("region", UserDataType.string),
        ("row_id", UserDataType.integer),
        ("sales", UserDataType.float),
        ("segment", UserDataType.string),
        ("ship_date", UserDataType.date),
        ("ship_mode", UserDataType.string),
        ("state", UserDataType.string),
        ("sub_category", UserDataType.string),
    ),
    nullable=False,
)

lib/dl_core_testing/dl_core_testing/fixtures/sample_tables.py
datalens-tech/datalens-backend
from __future__ import annotations

import json
from typing import (
    Dict,
    Optional,
)

from flask.testing import FlaskClient
import werkzeug

from dl_api_commons.tracing import get_current_tracing_headers


class FlaskTestResponse(werkzeug.wrappers.Response):
    def _json(self):  # type: ignore  # TODO: fix
        """Get the result of json.loads if possible."""
        if "json" not in self.mimetype:  # type: ignore  # TODO: fix
            raise AttributeError("Not a JSON response")
        return json.loads(self.data)

    json = werkzeug.utils.cached_property(_json)


class FlaskTestClient(FlaskClient):
    def get_us_home_directory(self) -> str:
        """To split US directories for different users to prevent permissions issues"""
        raise NotImplementedError()

    def get_default_headers(self) -> Dict[str, Optional[str]]:
        return {}

    def post_process_response(self, resp) -> None:  # type: ignore  # TODO: fix
        pass

    def open(self, *args, **kw):  # type: ignore  # TODO: fix
        kw["headers"] = {**self.get_default_headers(), **kw.get("headers", {}), **get_current_tracing_headers()}
        kw["headers"] = {key: val for key, val in kw["headers"].items() if val}
        resp = super().open(*args, **kw)
        self.post_process_response(resp)
        return resp

lib/dl_core_testing/dl_core_testing/flask_utils.py
datalens-tech/datalens-backend
from pytest import Config
from statcommons.logs import LOGMUTATORS

from dl_core.loader import load_core_lib
from dl_core.logging_config import add_log_context_scoped
from dl_core_testing.configuration import CoreTestEnvironmentConfigurationBase
from dl_core_testing.environment import prepare_united_storage_from_config
from dl_db_testing.loader import load_db_testing_lib


def initialize_core_test(pytest_config: Config, core_test_config: CoreTestEnvironmentConfigurationBase) -> None:
    # Configure pytest itself
    pytest_config.addinivalue_line("markers", "slow: ...")
    pytest_config.addinivalue_line("markers", "yt: ...")

    # Add Log context to logging records (not only in format phase)
    LOGMUTATORS.apply(require=False)
    # TODO FIX: Replace with add_log_context after all tests will be refactored to use unscoped log context
    LOGMUTATORS.add_mutator("log_context_scoped", add_log_context_scoped)

    # Prepare US
    prepare_united_storage_from_config(core_test_config.get_us_config())

    # Initialize this and other libraries
    load_db_testing_lib()
    load_core_lib(core_lib_config=core_test_config.get_core_library_config())

lib/dl_core_testing/dl_core_testing/initialization.py
datalens-tech/datalens-backend
import abc
from typing import (
    ClassVar,
    Generic,
    Optional,
    Type,
    TypeVar,
)

import pytest

from dl_api_commons.base_models import RequestContextInfo
from dl_core import exc
from dl_core.connection_executors import AsyncConnExecutorBase
from dl_core.connection_executors.adapters.async_adapters_base import AsyncDirectDBAdapter
from dl_core.connection_executors.models.connection_target_dto_base import BaseSQLConnTargetDTO
from dl_core.connection_executors.models.db_adapter_data import DBAdapterQuery
from dl_core.connection_executors.models.scoped_rci import DBAdapterScopedRCI
from dl_core_testing.testcases.connection_executor import BaseConnectionExecutorTestClass


_TARGET_DTO_TV = TypeVar("_TARGET_DTO_TV", bound=BaseSQLConnTargetDTO)


class BaseAsyncAdapterTestClass(BaseConnectionExecutorTestClass, Generic[_TARGET_DTO_TV], metaclass=abc.ABCMeta):
    """
    Most of the testable adapter behaviour is covered by conn executor tests
    Here are the tests, that are just easier to implement at the adapter level
    """

    ASYNC_ADAPTER_CLS: ClassVar[Type[AsyncDirectDBAdapter]]  # TODO add tests for other adapters

    @pytest.fixture
    async def target_conn_dto(self, async_connection_executor: AsyncConnExecutorBase) -> _TARGET_DTO_TV:
        target_conn_dto_pool = await async_connection_executor._make_target_conn_dto_pool()  # noqa
        return next(iter(target_conn_dto_pool))

    def _make_dba(self, target_dto: _TARGET_DTO_TV, rci: RequestContextInfo) -> AsyncDirectDBAdapter:
        return self.ASYNC_ADAPTER_CLS.create(
            req_ctx_info=DBAdapterScopedRCI.from_full_rci(rci),
            target_dto=target_dto,
            default_chunk_size=10,
        )

    async def _test_pass_db_query_to_user(
        self,
        pass_db_query_to_user: Optional[bool],
        query_to_send: str,
        expected_query: Optional[str],
        conn_bi_context: RequestContextInfo,
        target_conn_dto: _TARGET_DTO_TV,
    ):
        debug_query = expected_query

        target_conn_dto = target_conn_dto.clone(port="65535")  # not the actual db port to raise connect error
        if pass_db_query_to_user is not None:
            target_conn_dto = target_conn_dto.clone(pass_db_query_to_user=pass_db_query_to_user)

        dba = self._make_dba(target_conn_dto, conn_bi_context)

        with pytest.raises(exc.SourceConnectError) as exception_info:
            await dba.execute(DBAdapterQuery(query=query_to_send, debug_compiled_query=debug_query))

        assert exception_info.value.query == expected_query, exception_info.value.query

    @pytest.mark.parametrize(
        "pass_db_query_to_user, expected_query",
        (
            (False, None),
            (True, "select 1 from <hidden>"),
        ),
    )
    async def test_pass_db_query_to_user(
        self,
        pass_db_query_to_user: bool,
        expected_query: str,
        conn_bi_context: RequestContextInfo,
        target_conn_dto: _TARGET_DTO_TV,
    ) -> None:
        query_to_send = "select 1 from very_secret_table"
        await self._test_pass_db_query_to_user(
            pass_db_query_to_user, query_to_send, expected_query, conn_bi_context, target_conn_dto
        )

    async def test_default_pass_db_query_to_user(
        self,
        conn_bi_context: RequestContextInfo,
        target_conn_dto: _TARGET_DTO_TV,
    ) -> None:
        await self._test_pass_db_query_to_user(
            pass_db_query_to_user=None,
            query_to_send="select 1 from very_secret_table",
            expected_query=None,
            conn_bi_context=conn_bi_context,
            target_conn_dto=target_conn_dto,
        )

    async def test_timeout(
        self,
        conn_bi_context: RequestContextInfo,
        target_conn_dto: _TARGET_DTO_TV,
    ) -> None:
        target_conn_dto = target_conn_dto.clone(port="65535")
        dba = self._make_dba(target_conn_dto, conn_bi_context)

        with pytest.raises(exc.SourceConnectError):
            await dba.execute(DBAdapterQuery(query="select 1"))

lib/dl_core_testing/dl_core_testing/testcases/adapter.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Callable,
    ClassVar,
    Generic,
    Optional,
    TypeVar,
)

import pytest

from dl_constants.enums import RawSQLLevel
from dl_core.us_connection_base import (
    ConnectionBase,
    DataSourceTemplate,
)
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import (
    Db,
    DbTable,
)
from dl_core_testing.fixtures.primitives import FixtureTableSpec
from dl_core_testing.fixtures.sample_tables import TABLE_SPEC_SAMPLE_SUPERSTORE
from dl_core_testing.testcases.service_base import (
    DbServiceFixtureTextClass,
    ServiceFixtureTextClass,
)
from dl_testing.regulated_test import RegulatedTestCase


if TYPE_CHECKING:
    from dl_core.connection_executors.async_base import AsyncConnExecutorBase
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase
    from dl_core.services_registry.top_level import ServicesRegistry


_CONN_TV = TypeVar("_CONN_TV", bound=ConnectionBase)


class BaseConnectionTestClass(
    ServiceFixtureTextClass, DbServiceFixtureTextClass, Generic[_CONN_TV], metaclass=abc.ABCMeta
):
    raw_sql_level: ClassVar[Optional[RawSQLLevel]] = None

    @pytest.fixture(scope="function")
    def sync_us_manager(self, conn_default_sync_us_manager: SyncUSManager) -> SyncUSManager:
        return conn_default_sync_us_manager

    @pytest.fixture(scope="function")
    def async_us_manager(self, conn_default_async_us_manager: AsyncUSManager) -> AsyncUSManager:
        return conn_default_async_us_manager

    @abc.abstractmethod
    @pytest.fixture(scope="function")
    def connection_creation_params(self) -> dict:
        raise NotImplementedError

    @abc.abstractmethod
    @pytest.fixture(scope="function")
    def saved_connection(self, sync_us_manager: SyncUSManager, connection_creation_params: dict) -> _CONN_TV:
        raise NotImplementedError

    @pytest.fixture(scope="function")
    def sync_conn_executor_factory(
        self,
        saved_connection: _CONN_TV,
        conn_sync_service_registry: ServicesRegistry,
    ) -> Callable[[], SyncConnExecutorBase]:
        def factory() -> SyncConnExecutorBase:
            ce_factory = conn_sync_service_registry.get_conn_executor_factory()
            return ce_factory.get_sync_conn_executor(saved_connection)

        return factory

    @pytest.fixture(scope="function")
    def async_conn_executor_factory(
        self,
        saved_connection: _CONN_TV,
        conn_async_service_registry: ServicesRegistry,
    ) -> Callable[[], AsyncConnExecutorBase]:
        def factory() -> AsyncConnExecutorBase:
            ce_factory = conn_async_service_registry.get_conn_executor_factory()
            return ce_factory.get_async_conn_executor(saved_connection)

        return factory

    @pytest.fixture(scope="class")
    def sample_table_spec(self) -> FixtureTableSpec:
        return TABLE_SPEC_SAMPLE_SUPERSTORE

    @pytest.fixture(scope="class")
    def sample_table(self, sample_table_spec: FixtureTableSpec, db: Db) -> DbTable:
        return self.db_table_dispenser.get_csv_table(db=db, spec=sample_table_spec)


class DefaultConnectionTestClass(RegulatedTestCase, BaseConnectionTestClass[_CONN_TV], Generic[_CONN_TV]):
    do_check_data_export_flag: ClassVar[bool] = False

    @abc.abstractmethod
    def check_saved_connection(self, conn: _CONN_TV, params: dict) -> None:
        raise NotImplementedError

    def test_make_connection(self, saved_connection: _CONN_TV, connection_creation_params: dict) -> None:
        conn = saved_connection
        self.check_saved_connection(conn=conn, params=connection_creation_params)

    def test_connection_test(
        self,
        saved_connection: _CONN_TV,
        sync_conn_executor_factory: Callable[[], SyncConnExecutorBase],
    ) -> None:
        def sync_conn_executor_factory_for_conn(connection: ConnectionBase) -> SyncConnExecutorBase:
            return sync_conn_executor_factory()

        conn = saved_connection
        # if test is unsuccessful, an exception is raised by the connection
        conn.test(conn_executor_factory=sync_conn_executor_factory_for_conn)

    @abc.abstractmethod
    def check_data_source_templates(self, conn: _CONN_TV, dsrc_templates: list[DataSourceTemplate]) -> None:
        # In case of disabled templates, explicitly assert `not dsrc_templates`
        raise NotImplementedError

    def test_connection_get_data_source_templates(
        self,
        saved_connection: _CONN_TV,
        sync_conn_executor_factory: Callable[[], SyncConnExecutorBase],
    ) -> None:
        def sync_conn_executor_factory_for_conn(connection: ConnectionBase) -> SyncConnExecutorBase:
            return sync_conn_executor_factory()

        conn = saved_connection
        dsrc_templates = conn.get_data_source_templates(conn_executor_factory=sync_conn_executor_factory_for_conn)

        self.check_data_source_templates(conn=conn, dsrc_templates=dsrc_templates)

    def test_connection_data_export_flag(
        self,
        saved_connection: ConnectionBase,
        sync_us_manager: SyncUSManager,
    ) -> None:
        if not self.do_check_data_export_flag:
            pytest.skip()

        conn = saved_connection
        us_manager = sync_us_manager
        assert conn.data_export_forbidden is False  # default value
        conn.data.data_export_forbidden = True
        assert conn.data_export_forbidden is True
        us_manager.save(conn)
        loaded_conn = us_manager.get_by_id(conn.uuid)
        assert loaded_conn.data.data_export_forbidden is True
        assert loaded_conn.data_export_forbidden is True

lib/dl_core_testing/dl_core_testing/testcases/connection.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import contextlib
import functools
from typing import (
    TYPE_CHECKING,
    AsyncGenerator,
    Callable,
    ClassVar,
    Generator,
    Generic,
    Optional,
    Sequence,
    TypeVar,
)

import attr
import pytest
import shortuuid
import sqlalchemy as sa
from sqlalchemy.types import TypeEngine

from dl_constants.enums import (
    ConnectionType,
    UserDataType,
)
from dl_core.connection_executors.common_base import ConnExecutorQuery
from dl_core.connection_models.common_models import (
    DBIdent,
    SATextTableDefinition,
    SchemaIdent,
    TableIdent,
)
from dl_core.db.native_type import (
    CommonNativeType,
    GenericNativeType,
    norm_native_type,
)
import dl_core.exc as core_exc
from dl_core.us_connection_base import ConnectionBase
from dl_core_testing.database import (
    C,
    Db,
    DbTable,
    make_table,
)
from dl_core_testing.testcases.connection import BaseConnectionTestClass
from dl_testing.regulated_test import RegulatedTestCase


if TYPE_CHECKING:
    from dl_core.connection_executors.async_base import AsyncConnExecutorBase
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase


_CONN_TV = TypeVar("_CONN_TV", bound=ConnectionBase)


class BaseConnectionExecutorTestClass(RegulatedTestCase, BaseConnectionTestClass[_CONN_TV], Generic[_CONN_TV]):
    @pytest.fixture(scope="function")
    def sync_connection_executor(
        self,
        sync_conn_executor_factory: Callable[[], SyncConnExecutorBase],
    ) -> Generator[SyncConnExecutorBase, None, None]:
        sync_conn_executor = sync_conn_executor_factory()
        yield sync_conn_executor_factory()
        sync_conn_executor.close()

    @pytest.fixture(scope="function")
    async def async_connection_executor(
        self,
        async_conn_executor_factory: Callable[[], AsyncConnExecutorBase],
    ) -> AsyncGenerator[AsyncConnExecutorBase, None, None]:
        async_conn_executor = async_conn_executor_factory()
        yield async_conn_executor_factory()
        await async_conn_executor.close()


class DefaultSyncAsyncConnectionExecutorCheckBase(BaseConnectionExecutorTestClass[_CONN_TV], Generic[_CONN_TV]):
    @pytest.fixture(scope="function")
    def db_ident(self) -> DBIdent:
        raise NotImplementedError

    @pytest.fixture(scope="function")
    def existing_table_ident(self, sample_table: DbTable) -> TableIdent:
        # Using sample table by default, but this can be overridden by redefining this fixture
        return TableIdent(
            db_name=sample_table.db.name,
            schema_name=sample_table.schema,
            table_name=sample_table.name,
        )

    @pytest.fixture(scope="class")
    def db_table_columns(self, db: Db) -> list[C]:
        return C.full_house()

    @pytest.fixture(scope="function")
    def db_table(self, db: Db, db_table_columns: list[C]) -> DbTable:
        db_table = make_table(db, columns=db_table_columns)
        yield db_table
        db.drop_table(db_table.table)

    @pytest.fixture(scope="function")
    def nonexistent_table_ident(self, existing_table_ident: TableIdent) -> TableIdent:
        return existing_table_ident.clone(table_name=f"nonexistent_table_{shortuuid.uuid()}")

    def check_db_version(self, db_version: Optional[str]) -> None:
        pass

    @pytest.fixture(scope="class")
    def query_for_session_check(self) -> str:
        return "SELECT 567"

    @contextlib.contextmanager
    def check_closing_sql_sessions(self, db: Db) -> Generator[None, None, None]:
        old_cnt = db.count_sql_sessions()
        yield
        new_cnt = db.count_sql_sessions()
        assert old_cnt == new_cnt, f"Expected {old_cnt} sessions, got {new_cnt}"


class DefaultSyncConnectionExecutorTestSuite(DefaultSyncAsyncConnectionExecutorCheckBase[_CONN_TV], Generic[_CONN_TV]):
    @pytest.fixture(scope="session")
    def conn_exec_factory_async_env(self) -> bool:
        return False

    def test_get_db_version(self, sync_connection_executor: SyncConnExecutorBase, db_ident: DBIdent) -> None:
        db_version = sync_connection_executor.get_db_version(db_ident)
        self.check_db_version(db_version)

    def test_test(self, sync_connection_executor: SyncConnExecutorBase) -> None:
        sync_connection_executor.test()

    def test_get_table_names(
        self,
        sample_table: DbTable,
        db: Db,
        sync_connection_executor: SyncConnExecutorBase,
    ):
        # at the moment, checks that sample table is listed among the others

        tables = [sample_table]
        expected_table_names = set(table.name for table in tables)

        actual_tables = sync_connection_executor.get_tables(SchemaIdent(db_name=db.name, schema_name=None))
        actual_table_names = [tid.table_name for tid in actual_tables]

        assert set(actual_table_names).issuperset(expected_table_names)

    def test_table_exists(
        self,
        sync_connection_executor: SyncConnExecutorBase,
        existing_table_ident: TableIdent,
        db_table: DbTable,
    ) -> None:
        assert sync_connection_executor.is_table_exists(existing_table_ident)

    def test_table_not_exists(
        self,
        sync_connection_executor: SyncConnExecutorBase,
        nonexistent_table_ident: TableIdent,
    ) -> None:
        assert not sync_connection_executor.is_table_exists(nonexistent_table_ident)

    @attr.s(frozen=True)
    class CD:
        sa_type: TypeEngine = attr.ib()
        # Expected data
        user_type: UserDataType = attr.ib()
        nullable: bool = attr.ib(default=True)
        nt_name: Optional[str] = attr.ib(default=None)
        nt: Optional[GenericNativeType] = attr.ib(default=None)

        def get_expected_native_type(self, conn_type: ConnectionType) -> GenericNativeType:
            return self.nt or CommonNativeType(
                conn_type=conn_type,
                name=norm_native_type(self.nt_name if self.nt_name is not None else self.sa_type),
                nullable=self.nullable,
            )

    def get_schemas_for_type_recognition(self) -> dict[str, Sequence[CD]]:
        return {
            "standard_types": [
                self.CD(sa.Integer(), UserDataType.integer),
                self.CD(sa.Float(), UserDataType.float),
                self.CD(sa.String(length=256), UserDataType.string),
                self.CD(sa.Date(), UserDataType.date),
                self.CD(sa.DateTime(), UserDataType.genericdatetime),
            ],
        }

    def test_type_recognition(self, request, db: Db, sync_connection_executor: SyncConnExecutorBase) -> None:
        for schema_name, type_schema in self.get_schemas_for_type_recognition().items():
            columns = [
                sa.Column(name=f"c_{shortuuid.uuid().lower()}", type_=column_data.sa_type)
                for column_data in type_schema
            ]
            sa_table = db.table_from_columns(columns=columns)

            db.create_table(sa_table)
            request.addfinalizer(functools.partial(db.drop_table, sa_table))

            detected_columns = sync_connection_executor.get_table_schema_info(
                table_def=TableIdent(db_name=db.name, schema_name=sa_table.schema, table_name=sa_table.name)
            ).schema
            assert len(detected_columns) == len(type_schema), f"Incorrect number of columns in schema {schema_name}"
            for col_idx, (expected_col, detected_col) in enumerate(zip(type_schema, detected_columns)):
                assert detected_col.user_type == expected_col.user_type, (
                    f"Incorrect user type detected for schema {schema_name} col #{col_idx}: "
                    f"expected {expected_col.user_type.name}, got {detected_col.user_type.name}"
                )
                expected_native_type = expected_col.get_expected_native_type(self.conn_type)
                assert (
                    detected_col.native_type == expected_native_type
                ), f"Incorrect native type detected for schema {schema_name} col #{col_idx}: expected {repr(expected_native_type)}, got {repr(detected_col.native_type)}"

    def test_simple_select(self, sync_connection_executor: SyncConnExecutorBase) -> None:
        query = ConnExecutorQuery(query=sa.select([sa.literal(1)]))
        result = next(iter(sync_connection_executor.execute(query).result))
        assert len(result) == 1
        assert result[0] == (1,)

    def test_error_on_select_from_nonexistent_source(
        self,
        db: Db,
        sync_connection_executor: SyncConnExecutorBase,
        nonexistent_table_ident: TableIdent,
    ) -> None:
        query = ConnExecutorQuery(query=f"SELECT * from {db.quote(nonexistent_table_ident.table_name)}")
        with pytest.raises(core_exc.SourceDoesNotExist):
            sync_connection_executor.execute(query)

    def test_closing_sql_sessions(
        self, db: Db, sync_connection_executor: SyncConnExecutorBase, query_for_session_check: str
    ) -> None:
        with self.check_closing_sql_sessions(db=db):
            for _i in range(5):
                sync_connection_executor.execute(ConnExecutorQuery(query=query_for_session_check))

    subselect_query_for_schema_test: ClassVar[str] = "(SELECT 1 AS num) AS source"

    @pytest.mark.parametrize("case", ["table", "subselect"])
    def test_get_table_schema_info(
        self,
        case: str,
        sync_connection_executor: SyncConnExecutorBase,
        existing_table_ident: TableIdent,
        db_table: DbTable,
    ) -> None:
        # Just tests that the adapter can successfully retrieve the schema.
        # Data source tests check this in more detail
        table_def = {
            "table": existing_table_ident,
            "subselect": SATextTableDefinition(
                text=sa.sql.elements.TextClause(self.subselect_query_for_schema_test),
            ),
        }[case]
        detected_columns = sync_connection_executor.get_table_schema_info(table_def=table_def).schema
        assert len(detected_columns) > 0

    def test_get_table_schema_info_for_nonexistent_table(
        self,
        sync_connection_executor: SyncConnExecutorBase,
        nonexistent_table_ident: TableIdent,
    ) -> None:
        with pytest.raises(core_exc.DLBaseException):
            sync_connection_executor.get_table_schema_info(table_def=nonexistent_table_ident)


class DefaultAsyncConnectionExecutorTestSuite(DefaultSyncAsyncConnectionExecutorCheckBase[_CONN_TV], Generic[_CONN_TV]):
    @pytest.fixture(scope="session")
    def conn_exec_factory_async_env(self) -> bool:
        return True

    async def test_get_db_version(self, async_connection_executor: AsyncConnExecutorBase, db_ident: DBIdent) -> None:
        db_version = await async_connection_executor.get_db_version(db_ident)
        self.check_db_version(db_version)

    async def test_test(self, async_connection_executor: AsyncConnExecutorBase) -> None:
        await async_connection_executor.test()

    async def test_multiple_connection_test(self, async_connection_executor: AsyncConnExecutorBase) -> None:
        coros = (async_connection_executor.test() for _ in range(10))
        await asyncio.gather(*coros)

    async def test_table_exists(
        self,
        async_connection_executor: AsyncConnExecutorBase,
        existing_table_ident: TableIdent,
    ) -> None:
        assert await async_connection_executor.is_table_exists(existing_table_ident)

    async def test_table_not_exists(
        self,
        async_connection_executor: AsyncConnExecutorBase,
        nonexistent_table_ident: TableIdent,
    ) -> None:
        assert not await async_connection_executor.is_table_exists(nonexistent_table_ident)

    async def test_simple_select(self, async_connection_executor: AsyncConnExecutorBase) -> None:
        query = ConnExecutorQuery(query=sa.select([sa.literal(1)]))
        result = await anext(aiter((await async_connection_executor.execute(query)).result))
        assert len(result) == 1
        assert result[0] == (1,)

    async def test_select_data(self, sample_table: DbTable, async_connection_executor: AsyncConnExecutorBase) -> None:
        n_rows = 3
        result = await async_connection_executor.execute(
            ConnExecutorQuery(
                query=sa.select(columns=sample_table.table.columns)
                .select_from(sample_table.table)
                .order_by(sample_table.table.columns[0])
                .limit(n_rows),
                chunk_size=6,
            )
        )
        rows = await result.get_all()
        assert len(rows) == n_rows

    async def test_cast_row_to_output(
        self,
        sample_table: DbTable,
        async_connection_executor: AsyncConnExecutorBase,
    ) -> None:
        result = await async_connection_executor.execute(
            ConnExecutorQuery(
                sa.select(columns=[sa.literal(1), sa.literal(2), sa.literal(3)])
                .select_from(sample_table.table)
                .limit(1),
                user_types=[
                    UserDataType.boolean,
                    UserDataType.float,
                    UserDataType.integer,
                ],
            )
        )
        rows = await result.get_all()
        assert rows == [(True, 2.0, 3)], rows

    async def test_error_on_select_from_nonexistent_source(
        self,
        db: Db,
        async_connection_executor: AsyncConnExecutorBase,
        nonexistent_table_ident: TableIdent,
    ) -> None:
        query = ConnExecutorQuery(query=f"SELECT * from {db.quote(nonexistent_table_ident.table_name)}")
        with pytest.raises(core_exc.SourceDoesNotExist):
            await async_connection_executor.execute(query)

    async def test_closing_sql_sessions(
        self, db: Db, async_connection_executor: AsyncConnExecutorBase, query_for_session_check: str
    ) -> None:
        with self.check_closing_sql_sessions(db=db):
            for _i in range(5):
                await async_connection_executor.execute(ConnExecutorQuery(query=query_for_session_check))

    async def test_get_table_schema_info(
        self,
        async_connection_executor: AsyncConnExecutorBase,
        existing_table_ident: TableIdent,
        db_table: DbTable,
    ) -> None:
        # Just tests that the adapter can successfully retrieve the schema.
        # Data source tests check this in more detail
        detected_columns = (
            await async_connection_executor.get_table_schema_info(table_def=existing_table_ident)
        ).schema
        assert len(detected_columns) > 0

    async def test_get_table_schema_info_for_nonexistent_table(
        self,
        async_connection_executor: AsyncConnExecutorBase,
        nonexistent_table_ident: TableIdent,
    ) -> None:
        with pytest.raises(core_exc.DLBaseException):
            await async_connection_executor.get_table_schema_info(table_def=nonexistent_table_ident)

lib/dl_core_testing/dl_core_testing/testcases/connection_executor.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Callable,
    ClassVar,
    Generic,
    Sequence,
    Type,
    TypeVar,
)
import uuid

import pytest
import shortuuid

from dl_constants.enums import UserDataType
from dl_core.connection_models import TableIdent
from dl_core.connectors.base.data_source_migration import (
    DataSourceMigrationInterface,
    get_data_source_migrator,
)
from dl_core.connectors.sql_base.data_source_migration import (
    SQLSubselectDSMI,
    SQLTableDSMI,
)
from dl_core.data_source.base import DataSource
from dl_core.data_source.sql import (
    SubselectDataSource,
    TableSQLDataSourceMixin,
)
from dl_core.data_source_spec.base import DataSourceSpec
from dl_core.data_source_spec.sql import SubselectDataSourceSpec
from dl_core.db import SchemaColumn
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.database import Db
from dl_core_testing.testcases.connection import BaseConnectionTestClass


if TYPE_CHECKING:
    from dl_core.connection_executors.sync_base import SyncConnExecutorBase


_CONN_TV = TypeVar("_CONN_TV", bound=ConnectionBase)
_DSRC_SPEC_TV = TypeVar("_DSRC_SPEC_TV", bound=DataSourceSpec)
_DSRC_TV = TypeVar("_DSRC_TV", bound=DataSource)


class BaseDataSourceTestClass(
    BaseConnectionTestClass[_CONN_TV],
    Generic[_CONN_TV, _DSRC_SPEC_TV, _DSRC_TV],
):
    DSRC_CLS: ClassVar[Type[_DSRC_TV]]

    @abc.abstractmethod
    @pytest.fixture(scope="class")
    def initial_data_source_spec(self) -> _DSRC_SPEC_TV:
        raise NotImplementedError

    @pytest.fixture(scope="function")
    def data_source(
        self,
        sync_us_manager: SyncUSManager,
        saved_connection: _CONN_TV,
        initial_data_source_spec: _DSRC_SPEC_TV,
    ) -> _DSRC_TV:
        dsrc = self.DSRC_CLS(
            id=uuid.uuid4().hex,
            us_entry_buffer=sync_us_manager.get_entry_buffer(),
            connection=saved_connection,
            spec=initial_data_source_spec,
        )
        return dsrc


class DefaultDataSourceTestClass(
    BaseDataSourceTestClass[_CONN_TV, _DSRC_SPEC_TV, _DSRC_TV],
    Generic[_CONN_TV, _DSRC_SPEC_TV, _DSRC_TV],
):
    def test_data_source_exists(
        self,
        data_source: _DSRC_TV,
        sync_conn_executor_factory: Callable[[], SyncConnExecutorBase],
    ) -> None:
        dsrc = data_source
        result = dsrc.source_exists(conn_executor_factory=sync_conn_executor_factory)
        assert result

    @abc.abstractmethod
    def get_expected_simplified_schema(self) -> list[tuple[str, UserDataType]]:
        raise NotImplementedError

    def test_get_raw_schema(
        self,
        data_source: _DSRC_TV,
        sync_conn_executor_factory: Callable[[], SyncConnExecutorBase],
    ) -> None:
        raw_schema = data_source.get_schema_info(conn_executor_factory=sync_conn_executor_factory).schema
        simplified_schema = [(col.title, col.user_type) for col in raw_schema]
        expected_schema = self.get_expected_simplified_schema()
        print("schema:")
        for col in simplified_schema:
            print(col)
        assert len(simplified_schema) == len(expected_schema)
        for actual, expected in zip(simplified_schema, expected_schema):
            assert actual == expected, f"{expected=} {actual=}"

    def _check_migration_dtos(
        self,
        data_source: _DSRC_TV,
        migration_dtos: Sequence[DataSourceMigrationInterface],
    ) -> None:
        pass

    def _check_migration_dtos_table(  # TODO: Plug it in
        self,
        data_source: _DSRC_TV,
        migration_dtos: Sequence[DataSourceMigrationInterface],
    ) -> None:
        assert isinstance(data_source, TableSQLDataSourceMixin)
        assert migration_dtos[0] == SQLTableDSMI(
            db_name=getattr(data_source, "db_name", None),
            schema_name=getattr(data_source, "schema_name", None),
            table_name=data_source.table_name,
        )

    def _check_migration_dtos_subselect(  # TODO: Plug it in
        self,
        data_source: _DSRC_TV,
        migration_dtos: Sequence[DataSourceMigrationInterface],
    ) -> None:
        assert isinstance(data_source, SubselectDataSource)
        assert migration_dtos[0] == SQLSubselectDSMI(
            subsql=data_source.subsql,
        )

    def test_export_dsrc_migration_dtos(
        self,
        data_source: _DSRC_TV,
        sync_us_manager: SyncUSManager,
    ) -> None:
        migrator = get_data_source_migrator(self.conn_type)
        source_spec = data_source.spec
        migration_dtos = migrator.export_migration_dtos(data_source_spec=source_spec)
        self._check_migration_dtos(data_source=data_source, migration_dtos=migration_dtos)


_SUBSELECT_DSRC_SPEC_TV = TypeVar("_SUBSELECT_DSRC_SPEC_TV", bound=SubselectDataSourceSpec)
_SUBSELECT_DSRC_TV = TypeVar("_SUBSELECT_DSRC_TV", bound=SubselectDataSource)


class DataSourceTestByViewClass(
    BaseDataSourceTestClass[_CONN_TV, _SUBSELECT_DSRC_SPEC_TV, _SUBSELECT_DSRC_TV],
    Generic[_CONN_TV, _SUBSELECT_DSRC_SPEC_TV, _SUBSELECT_DSRC_TV],
):
    def postprocess_view_schema_column(self, schema_col: SchemaColumn) -> SchemaColumn:
        """
        Place for marking known inevitable discrepancies between subselect
        schema and view schema.
        """
        return schema_col

    def postprocess_view_schema(
        self, view_schema: list[SchemaColumn], cursor_schema: list[SchemaColumn]
    ) -> list[SchemaColumn]:
        return [self.postprocess_view_schema_column(schema_col) for schema_col in view_schema]

    @pytest.fixture(scope="function")
    def subselect_view(self, initial_data_source_spec: _SUBSELECT_DSRC_SPEC_TV, db: Db) -> str:
        name = f"test_view_from_subselect_{shortuuid.uuid().lower()}"
        db.execute(f"create view {name} as select * from ({initial_data_source_spec.subsql}) source")
        return name

    def test_view_from_subselect(
        self,
        data_source: _SUBSELECT_DSRC_TV,
        db: Db,
        sync_conn_executor_factory: Callable[[], SyncConnExecutorBase],
        subselect_view: str,
    ) -> None:
        """Compare the subselect schema to a view-based schema"""
        dsrc = data_source
        conn_executor = sync_conn_executor_factory()
        schema = dsrc.get_schema_info(conn_executor_factory=sync_conn_executor_factory).schema

        # A bit of a hack, for easier inter-db testing:
        view_schema_info = conn_executor.get_table_schema_info(
            TableIdent(db_name=None, schema_name=None, table_name=subselect_view),
        )
        view_schema = dsrc._postprocess_raw_schema_from_db(view_schema_info).schema
        view_schema = self.postprocess_view_schema(view_schema, cursor_schema=schema)

        # for a more readable diff:
        if len(schema) == len(view_schema):
            sc_mismatch = [sc for sc, vsc in zip(schema, view_schema) if sc != vsc]
            vsc_mismatch = [vsc for sc, vsc in zip(schema, view_schema) if sc != vsc]
            assert sc_mismatch == vsc_mismatch

        assert schema == tuple(view_schema)


class SQLDataSourceTestClass(
    BaseDataSourceTestClass[_CONN_TV, _DSRC_SPEC_TV, _DSRC_TV],
    Generic[_CONN_TV, _DSRC_SPEC_TV, _DSRC_TV],
):
    QUERY_PATTERN: ClassVar[str]

    def test_data_source_query(self, data_source: _DSRC_TV) -> None:
        dsrc = data_source
        sql_query = str(dsrc.get_sql_source())
        assert self.QUERY_PATTERN in sql_query

lib/dl_core_testing/dl_core_testing/testcases/data_source.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    ClassVar,
    Generic,
    Optional,
    TypeVar,
)

import pytest
import sqlalchemy as sa

from dl_constants.enums import (
    DataSourceRole,
    DataSourceType,
)
from dl_core.data_processing.stream_base import DataStream
from dl_core.query.bi_query import BIQuery
from dl_core.query.expression import ExpressionCtx
from dl_core.services_registry.top_level import ServicesRegistry
from dl_core.us_connection_base import ConnectionBase
from dl_core.us_dataset import Dataset
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.data import DataFetcher
from dl_core_testing.database import DbTable
from dl_core_testing.dataset import make_dataset
from dl_core_testing.dataset_wrappers import (
    DatasetTestWrapper,
    EditableDatasetTestWrapper,
)
from dl_core_testing.testcases.connection import BaseConnectionTestClass
from dl_testing.regulated_test import RegulatedTestCase


_CONN_TV = TypeVar("_CONN_TV", bound=ConnectionBase)


class BaseDatasetTestClass(BaseConnectionTestClass[_CONN_TV], Generic[_CONN_TV]):
    source_type: ClassVar[DataSourceType]

    @pytest.fixture(scope="function")
    def dataset_table(self, sample_table: DbTable) -> DbTable:
        """The table to be used for datasets. By default use the sample table"""
        return sample_table

    @pytest.fixture(scope="function")
    def dsrc_params(self, dataset_table: DbTable) -> dict:
        return dict(
            db_name=dataset_table.db.name,
            schema_name=dataset_table.schema,
            table_name=dataset_table.name,
        )

    @pytest.fixture(scope="function")
    def saved_dataset(
        self,
        sync_us_manager: SyncUSManager,
        saved_connection: _CONN_TV,
        dsrc_params: dict,
    ) -> Dataset:
        dataset = make_dataset(
            sync_usm=sync_us_manager,
            connection=saved_connection,
            created_from=self.source_type,
            dsrc_params=dsrc_params,
        )
        sync_us_manager.save(dataset)
        return dataset

    @pytest.fixture(scope="function")
    def empty_saved_dataset(self, sync_us_manager: SyncUSManager) -> Dataset:
        dataset = make_dataset(sync_usm=sync_us_manager)
        sync_us_manager.save(dataset)
        return dataset

    @pytest.fixture(scope="function")
    def dataset_wrapper(self, saved_dataset: Dataset, sync_us_manager: SyncUSManager) -> DatasetTestWrapper:
        return DatasetTestWrapper(dataset=saved_dataset, us_manager=sync_us_manager)

    @pytest.fixture(scope="function")
    def editable_dataset_wrapper(
        self,
        saved_dataset: Dataset,
        sync_us_manager: SyncUSManager,
    ) -> EditableDatasetTestWrapper:
        return EditableDatasetTestWrapper(dataset=saved_dataset, us_manager=sync_us_manager)

    def fetch_data(
        self,
        saved_dataset: Dataset,
        service_registry: ServicesRegistry,
        sync_us_manager: SyncUSManager,
        bi_query: BIQuery,
        from_subquery: bool = False,
        subquery_limit: Optional[int] = None,
    ) -> DataStream:
        data_fetcher = DataFetcher(
            service_registry=service_registry,
            dataset=saved_dataset,
            us_manager=sync_us_manager,
        )
        return data_fetcher.get_data_stream(
            role=DataSourceRole.origin,
            bi_query=bi_query,
            from_subquery=from_subquery,
            subquery_limit=subquery_limit,
        )


class DefaultDatasetTestSuite(RegulatedTestCase, BaseDatasetTestClass[_CONN_TV], Generic[_CONN_TV]):
    def _check_simple_select(
        self,
        dataset_wrapper: DatasetTestWrapper,
        saved_dataset: Dataset,
        async_service_registry: ServicesRegistry,
        sync_us_manager: SyncUSManager,
        result_cnt: int,
        limit: Optional[int] = None,
        from_subquery: bool = False,
        subquery_limit: Optional[int] = None,
    ) -> None:
        assert limit is not None or (from_subquery and subquery_limit is not None)
        avatar_id = dataset_wrapper.get_root_avatar_strict().id
        raw_schema = dataset_wrapper.get_cached_raw_schema(role=DataSourceRole.origin)
        col_expr = sa.literal_column(dataset_wrapper.quote(raw_schema[0].name, role=DataSourceRole.origin))
        bi_query = BIQuery(
            select_expressions=[
                ExpressionCtx(
                    expression=col_expr,
                    user_type=raw_schema[0].user_type,
                    avatar_ids=[avatar_id],
                    alias=raw_schema[0].name,
                ),
            ],
            limit=limit,
        )
        data = self.fetch_data(
            saved_dataset=saved_dataset,
            service_registry=async_service_registry,
            sync_us_manager=sync_us_manager,
            bi_query=bi_query,
            from_subquery=from_subquery,
            subquery_limit=subquery_limit,
        )
        assert len(list(data.data)) == result_cnt

    def test_simple_select(
        self,
        dataset_wrapper: DatasetTestWrapper,
        saved_dataset: Dataset,
        conn_async_service_registry: ServicesRegistry,
        sync_us_manager: SyncUSManager,
    ) -> None:
        self._check_simple_select(
            dataset_wrapper=dataset_wrapper,
            saved_dataset=saved_dataset,
            async_service_registry=conn_async_service_registry,
            sync_us_manager=sync_us_manager,
            limit=5,
            result_cnt=5,
        )

    def test_simple_select_from_subquery(
        self,
        dataset_wrapper: DatasetTestWrapper,
        saved_dataset: Dataset,
        conn_async_service_registry: ServicesRegistry,
        sync_us_manager: SyncUSManager,
    ) -> None:
        self._check_simple_select(
            dataset_wrapper=dataset_wrapper,
            saved_dataset=saved_dataset,
            async_service_registry=conn_async_service_registry,
            sync_us_manager=sync_us_manager,
            from_subquery=True,
            subquery_limit=5,
            result_cnt=5,
        )

    def test_get_param_hash(
        self,
        sample_table: DbTable,
        saved_connection: ConnectionBase,
        saved_dataset: Dataset,
        conn_default_service_registry: ServicesRegistry,
        dataset_wrapper: DatasetTestWrapper,
    ) -> None:
        dataset = saved_dataset
        service_registry = conn_default_service_registry
        source_id = dataset.get_single_data_source_id()
        dsrc_coll = dataset_wrapper.get_data_source_coll_strict(source_id=source_id)
        hash_from_dataset = dsrc_coll.get_param_hash()

        templates = saved_connection.get_data_source_templates(
            conn_executor_factory=service_registry.get_conn_executor_factory().get_sync_conn_executor,
        )
        found_template = False
        for template in templates:
            if template.parameters["table_name"] == sample_table.name:
                found_template = True
                hash_from_template = template.get_param_hash()
                assert hash_from_dataset == hash_from_template

        assert found_template

lib/dl_core_testing/dl_core_testing/testcases/dataset.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    Any,
    ClassVar,
    NamedTuple,
    Optional,
)

import pytest

from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.reporting.registry import DefaultReportingRegistry
from dl_configs.connectors_settings import ConnectorSettingsBase
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_configs.rqe import RQEConfig
from dl_constants.enums import ConnectionType
from dl_core.connections_security.base import InsecureConnectionSecurityManager
from dl_core.services_registry.conn_executor_factory import DefaultConnExecutorFactory
from dl_core.services_registry.inst_specific_sr import InstallationSpecificServiceRegistryFactory
from dl_core.services_registry.top_level import (
    DefaultServicesRegistry,
    ServicesRegistry,
)
from dl_core.united_storage_client import USAuthContextMaster
from dl_core.us_manager.mutation_cache.usentry_mutation_cache_factory import DefaultUSEntryMutationCacheFactory
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core.utils import FutureRef
from dl_core_testing.configuration import CoreTestEnvironmentConfigurationBase
from dl_core_testing.database import (
    CoreDbConfig,
    CoreReInitableDbDispenser,
    Db,
)
from dl_core_testing.fixtures.dispenser import DbCsvTableDispenser
from dl_db_testing.database.engine_wrapper import DbEngineConfig
from dl_utils.aio import ContextVarExecutor


class USConfig(NamedTuple):
    us_base_url: str
    us_auth_context: USAuthContextMaster
    us_crypto_keys_config: CryptoKeysConfig


class ServiceFixtureTextClass(metaclass=abc.ABCMeta):
    conn_type: ClassVar[ConnectionType]
    core_test_config: ClassVar[CoreTestEnvironmentConfigurationBase]
    connection_settings: ClassVar[Optional[ConnectorSettingsBase]] = None
    inst_specific_sr_factory: ClassVar[Optional[InstallationSpecificServiceRegistryFactory]] = None

    @pytest.fixture(scope="session")
    def conn_us_config(self) -> USConfig:
        us_env_config = self.core_test_config.get_us_config()
        return USConfig(
            us_base_url=us_env_config.us_host,
            us_auth_context=USAuthContextMaster(us_env_config.us_master_token),
            us_crypto_keys_config=self.core_test_config.get_crypto_keys_config(),
        )

    @pytest.fixture(scope="session")
    def conn_bi_context(self) -> RequestContextInfo:
        bi_context = RequestContextInfo.create_empty()
        return bi_context

    @pytest.fixture(scope="session")
    def conn_exec_factory_async_env(self) -> bool:
        return False

    def service_registry_factory(
        self,
        conn_exec_factory_async_env: bool,
        conn_bi_context: RequestContextInfo,
        **kwargs: Any,
    ) -> ServicesRegistry:
        sr_future_ref: FutureRef[ServicesRegistry] = FutureRef()
        service_registry = DefaultServicesRegistry(
            rci=conn_bi_context,
            mutations_cache_factory=DefaultUSEntryMutationCacheFactory(),
            reporting_registry=DefaultReportingRegistry(
                rci=conn_bi_context,
            ),
            conn_exec_factory=DefaultConnExecutorFactory(
                async_env=conn_exec_factory_async_env,
                force_non_rqe_mode=True,
                rqe_config=RQEConfig.get_default(),  # Not used because RQE is disabled
                services_registry_ref=sr_future_ref,
                conn_sec_mgr=InsecureConnectionSecurityManager(),
                tpe=ContextVarExecutor(),
            ),
            connectors_settings={self.conn_type: self.connection_settings} if self.connection_settings else {},
            inst_specific_sr=(
                self.inst_specific_sr_factory.get_inst_specific_sr(sr_future_ref)
                if self.inst_specific_sr_factory is not None
                else None
            ),
            **kwargs,
        )
        sr_future_ref.fulfill(service_registry)
        return service_registry

    @pytest.fixture(scope="session")
    def conn_sync_service_registry(
        self,
        conn_bi_context: RequestContextInfo,
    ) -> ServicesRegistry:
        return self.service_registry_factory(
            conn_exec_factory_async_env=False,
            conn_bi_context=conn_bi_context,
        )

    @pytest.fixture(scope="session")
    def conn_async_service_registry(
        self,
        conn_bi_context: RequestContextInfo,
    ) -> ServicesRegistry:
        return self.service_registry_factory(
            conn_exec_factory_async_env=True,
            conn_bi_context=conn_bi_context,
        )

    @pytest.fixture(scope="session")
    def conn_default_service_registry(
        self,
        conn_exec_factory_async_env: bool,
        conn_sync_service_registry: ServicesRegistry,
        conn_async_service_registry: ServicesRegistry,
    ) -> ServicesRegistry:
        if conn_exec_factory_async_env:
            return conn_async_service_registry
        return conn_sync_service_registry

    @pytest.fixture(scope="class")
    def conn_default_sync_us_manager(
        self,
        conn_us_config: USConfig,
        conn_bi_context: RequestContextInfo,
        conn_default_service_registry: ServicesRegistry,
    ) -> SyncUSManager:
        us_manager = SyncUSManager(
            bi_context=conn_bi_context,
            services_registry=conn_default_service_registry,
            us_base_url=conn_us_config.us_base_url,
            us_auth_context=conn_us_config.us_auth_context,
            crypto_keys_config=conn_us_config.us_crypto_keys_config,
        )
        return us_manager

    @pytest.fixture(scope="class")
    def conn_default_async_us_manager(
        self,
        conn_us_config: USConfig,
        conn_bi_context: RequestContextInfo,
        conn_async_service_registry: ServicesRegistry,
    ) -> AsyncUSManager:
        us_manager = AsyncUSManager(
            bi_context=conn_bi_context,
            services_registry=conn_async_service_registry,
            us_base_url=conn_us_config.us_base_url,
            us_auth_context=conn_us_config.us_auth_context,
            crypto_keys_config=conn_us_config.us_crypto_keys_config,
        )
        return us_manager


class DbServiceFixtureTextClass(metaclass=abc.ABCMeta):
    conn_type: ClassVar[ConnectionType]  # FIXME: Remove after conn_type is removed from Db
    db_dispenser: ClassVar[CoreReInitableDbDispenser] = CoreReInitableDbDispenser()
    db_table_dispenser = DbCsvTableDispenser()

    @abc.abstractmethod
    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        pass

    @pytest.fixture(scope="class")
    def engine_params(self) -> dict:
        return {}

    @pytest.fixture(scope="class")
    def engine_config(self, db_url: str, engine_params: dict) -> DbEngineConfig:
        return DbEngineConfig(url=db_url, engine_params=engine_params)

    @pytest.fixture(scope="class")
    def db_config(self, engine_config: DbEngineConfig) -> CoreDbConfig:
        return CoreDbConfig(
            engine_config=engine_config,
            conn_type=self.conn_type,
        )

    @pytest.fixture(scope="class", autouse=True)
    def _add_db_init_hook(self, db_config: CoreDbConfig) -> None:
        self.db_dispenser.add_reinit_hook(
            # This method is idempotent, so it's OK that this will get called for every subclass
            db_config=db_config,
            reinit_hook=self.db_reinit_hook,
        )

    def db_reinit_hook(self) -> None:
        pass

    @pytest.fixture(scope="class")
    def db(self, db_config: CoreDbConfig) -> Db:
        return self.db_dispenser.get_database(db_config)

    @pytest.fixture(scope="class", autouse=True)
    def db_is_ready(self, db: Db) -> str:
        pass  # Just making sure that the database is available

lib/dl_core_testing/dl_core_testing/testcases/service_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_core.data_processing.cache.primitives import CacheTTLConfig


@attr.s(auto_attribs=True, frozen=True)
class SROptions:
    rci: RequestContextInfo
    with_caches: bool = False
    cache_save_background: Optional[bool] = False
    default_caches_ttl_config: CacheTTLConfig = CacheTTLConfig()
    with_compeng_pg: bool = False

lib/dl_core_testing/dl_core_testing/utils.py
datalens-tech/datalens-backend
from typing import Type

from dl_db_testing.connectors.base.connector import DbTestingConnector
from dl_db_testing.database.engine_wrapper import register_engine_wrapper_cls


_REGISTERED: set[Type[DbTestingConnector]] = set()


class DbTestingConnectorRegistrator:
    @classmethod
    def register_connector(cls, connector: Type[DbTestingConnector]) -> None:
        if connector in _REGISTERED:
            return

        for engine_wrapper_cls in connector.engine_wrapper_classes:
            register_engine_wrapper_cls(engine_wrapper_cls)

        _REGISTERED.add(connector)


CONN_REG_DB_TESTING = DbTestingConnectorRegistrator

lib/dl_db_testing/dl_db_testing/connector_registration.py
datalens-tech/datalens-backend
from typing import (
    ClassVar,
    Type,
)

from dl_db_testing.database.engine_wrapper import EngineWrapperBase


class DbTestingConnector:
    engine_wrapper_classes: ClassVar[tuple[Type[EngineWrapperBase], ...]] = ()

lib/dl_db_testing/dl_db_testing/connectors/base/connector.py
datalens-tech/datalens-backend
from __future__ import annotations

import contextlib
from typing import (
    Any,
    Generator,
    Generic,
    List,
    Optional,
    Sequence,
    TypeVar,
    Union,
    cast,
)

import attr
import sqlalchemy as sa
from sqlalchemy.engine import Connection
import sqlalchemy.engine.url
import sqlalchemy.exc as sa_exc

import dl_db_testing.database.engine_wrapper as ew


@attr.s(frozen=True)
class DbConfig:
    engine_config: ew.DbEngineConfig = attr.ib(kw_only=True)


_DB_CONFIG_TV = TypeVar("_DB_CONFIG_TV", bound=DbConfig)


@attr.s
class DbBase(Generic[_DB_CONFIG_TV]):
    _engine_wrapper: ew.EngineWrapperBase = attr.ib(kw_only=True)
    _config: _DB_CONFIG_TV = attr.ib(kw_only=True)

    @property
    def engine(self) -> sqlalchemy.engine.Engine:
        return self._engine_wrapper.engine

    @property
    def engine_config(self) -> ew.DbEngineConfig:
        return self._engine_wrapper.config

    @property
    def config(self) -> _DB_CONFIG_TV:
        return cast(_DB_CONFIG_TV, self._config)

    @property
    def url(self) -> sqlalchemy.engine.url.URL:
        return self._engine_wrapper.url

    def get_conn_credentials(self, full: bool = False) -> dict:
        return self._engine_wrapper.get_conn_credentials(full=full)

    @property
    def name(self) -> str:
        return self.get_conn_credentials(full=True)["db_name"]

    def count_sql_sessions(self) -> int:
        return self._engine_wrapper.count_sql_sessions()

    def get_conn_hosts(self) -> tuple[str, ...]:
        return self._engine_wrapper.get_conn_hosts()

    def quote(self, value: str) -> str:
        return self._engine_wrapper.quote(value)

    def expr_as_str(self, expr: sa.sql.ClauseElement) -> str:
        try:
            return expr.compile(
                compile_kwargs={"literal_binds": True},
                dialect=self._engine_wrapper.dialect,
            ).string
        except (NotImplementedError, sa_exc.CompileError):
            return str(expr)

    def execute(self, query: Any, *multiparams: Any, **params: Any) -> Any:
        return self._engine_wrapper.execute(query, *multiparams, **params)

    def scalar(self, expr: Any, from_: Optional[sa.sql.ClauseElement] = None) -> Any:
        assert isinstance(expr, sa.sql.ClauseElement)
        return self.execute(sa.select([expr]), from_=from_).scalar()

    def has_table(self, table_name: str) -> bool:
        return self._engine_wrapper.has_table(table_name)

    def get_current_connection(self) -> Connection:
        # FIXME: Remove
        return self._engine_wrapper.connection

    @contextlib.contextmanager
    def connect(self) -> Generator[Connection, None, None]:
        with self._engine_wrapper.connect() as connection:
            yield connection

    def drop_table(self, table: sa.Table) -> None:
        self._engine_wrapper.drop_table(db_name=self.name, table=table)

    def create_table(self, table: sa.Table) -> None:
        self.drop_table(table)
        self._engine_wrapper.create_table(table)

    def create_view(self, query: str, name: str, schema: Optional[str] = None) -> None:
        name = ".".join([self.quote(part) for part in (schema, name) if part])
        self.execute(f"CREATE VIEW {name} AS {query}")

    def load_table(self, table_name: str, schema: Optional[str] = None) -> sa.Table:
        return self._engine_wrapper.load_table(schema=schema, table_name=table_name)

    def table_from_columns(
        self,
        columns: Sequence[sa.Column],
        *,
        schema: Optional[str] = None,
        table_name: Optional[str] = None,
    ) -> sa.Table:
        return self._engine_wrapper.table_from_columns(
            columns=columns,
            schema=schema,
            table_name=table_name,
        )

    def insert_into_table(self, table: sa.Table, data: Sequence[dict]) -> None:
        self._engine_wrapper.insert_into_table(table=table, data=data)

    def create_schema(self, schema_name: str) -> None:
        self._engine_wrapper.create_schema(schema_name=schema_name)

    def base_eval(self, expr: sa.sql.ClauseElement, from_: Optional[sa.sql.ClauseElement] = None) -> Any:
        query = sa.select([expr])
        if from_ is not None:
            query = query.select_from(from_)

        try:
            value = self.execute(query).scalar()
            return value
        except Exception:
            print("QUERY:", self.expr_as_str(query))
            raise

    def get_version(self) -> Optional[str]:
        return self._engine_wrapper.get_version()

    def test(self) -> bool:
        return self._engine_wrapper.test()


@attr.s(frozen=True)
class DbTableBase:
    db: DbBase = attr.ib(kw_only=True)
    table: sa.Table = attr.ib(kw_only=True)
    schema: Optional[str] = attr.ib(kw_only=True, default=None)
    can_insert: bool = attr.ib(kw_only=True, default=True)

    @property
    def name(self) -> str:
        return self.table.name

    def insert(self, data: Union[dict, List[dict]]) -> None:
        # TODO: Use insert_into_table
        if not self.can_insert:
            raise RuntimeError("Can't insert into table")
        if isinstance(data, dict):
            self.db.execute(self.table.insert(data))
        elif isinstance(data, list):
            for row in data:
                self.insert(row)
        else:
            raise TypeError(type(data))

    @property
    def select_all_query(self) -> str:
        name = ".".join([self.db.quote(part) for part in (self.schema, self.name) if part])
        return f"SELECT * FROM {name}"

lib/dl_db_testing/dl_db_testing/database/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from collections import defaultdict
from typing import (
    Callable,
    Dict,
    Generic,
    Optional,
    TypeVar,
)

import attr

from dl_db_testing.database.base import (
    DbBase,
    DbConfig,
)
from dl_utils.wait import wait_for


_DB_CONFIG_TV = TypeVar("_DB_CONFIG_TV", bound=DbConfig)
_DB_TV = TypeVar("_DB_TV", bound=DbBase)


@attr.s
class DbDispenserBase(abc.ABC, Generic[_DB_CONFIG_TV, _DB_TV]):
    _wait_on_init: bool = attr.ib(kw_only=True, default=True)
    _default_reconnect_timeout: int = attr.ib(kw_only=True, default=600)

    @abc.abstractmethod
    def make_database(self, db_config: _DB_CONFIG_TV) -> _DB_TV:
        raise NotImplementedError

    def wait_for_db(self, db: _DB_TV, reconnect_timeout: Optional[int] = None) -> None:
        if reconnect_timeout is None:
            reconnect_timeout = self._default_reconnect_timeout
        wait_for(name=f"test_db_{db.config}", condition=db.test, timeout=reconnect_timeout)

    def initialize_db(self, db: _DB_TV) -> None:
        if self._wait_on_init:
            if not db.test():
                self.wait_for_db(db=db)

    def get_database(self, db_config: _DB_CONFIG_TV) -> _DB_TV:
        db = self.make_database(db_config)
        self.initialize_db(db)
        return db


@attr.s
class ReInitableDbDispenser(DbDispenserBase[_DB_CONFIG_TV, _DB_TV], Generic[_DB_CONFIG_TV, _DB_TV]):
    _max_reinit_count: int = attr.ib(kw_only=True, default=4)
    _db_cache: Dict[_DB_CONFIG_TV, _DB_TV] = attr.ib(init=False, factory=dict)
    _reinit_hooks: Dict[_DB_CONFIG_TV, Callable[[], None]] = attr.ib(init=False, factory=dict)
    _db_reinit_counts: Dict[_DB_CONFIG_TV, int] = attr.ib(init=False, factory=lambda: defaultdict(lambda: 0))

    def get_database(self, db_config: _DB_CONFIG_TV) -> _DB_TV:
        if db_config not in self._db_cache:
            self._db_cache[db_config] = super().get_database(db_config)
        self._check_reinit_db(db_config=db_config)
        return self._db_cache[db_config]

    def add_reinit_hook(self, *, db_config: _DB_CONFIG_TV, reinit_hook: Callable[[], None]) -> None:
        self._reinit_hooks[db_config] = reinit_hook

    def _check_reinit_db(self, db_config: DbConfig, reconnect_timeout: Optional[int] = None) -> bool:
        """Check if DB is alive. If it is not, try to re-initialize it"""
        db = self._db_cache.get(db_config)
        if db is None:
            return False

        reinit_hook = self._reinit_hooks.get(db_config)
        if reinit_hook is None or self._db_reinit_counts[db_config] >= self._max_reinit_count:
            return False

        if not db.test():
            # DB is unavailable, so re-initialize the DB is possible
            reinit_hook()
            # Wait until it comes up
            self.wait_for_db(db=db, reconnect_timeout=reconnect_timeout)
            self._db_reinit_counts[db_config] += 1
            return True

        return False

    def check_reinit_all(self, reconnect_timeout: Optional[int] = None) -> None:
        for db_config, _db in self._db_cache.items():
            self._check_reinit_db(db_config=db_config, reconnect_timeout=reconnect_timeout)

lib/dl_db_testing/dl_db_testing/database/dispenser.py
datalens-tech/datalens-backend
from __future__ import annotations

import contextlib
from typing import (
    Any,
    ClassVar,
    Generator,
    Optional,
    Sequence,
    Type,
)
import urllib.parse

import attr
from frozendict import frozendict
import shortuuid
import sqlalchemy as sa
from sqlalchemy.engine.base import (
    Connection,
    Engine,
)
from sqlalchemy.engine.default import DefaultDialect
from sqlalchemy.engine.url import (
    URL,
    make_url,
)
import sqlalchemy.exc as sa_exc
from sqlalchemy.pool import (
    NullPool,
    Pool,
)

from dl_utils.wait import wait_for


@attr.s(frozen=True)
class DbEngineConfig:
    url: str | URL = attr.ib(kw_only=True)
    engine_params: frozendict[str, Any] = attr.ib(kw_only=True, converter=frozendict, factory=frozendict)


@attr.s(frozen=True)
class TableSpec:
    # db_name: str = attr.ib(kw_only=True)
    table_name: str = attr.ib(kw_only=True)


@attr.s
class EngineWrapperBase:
    URL_PREFIX: ClassVar[str]
    POOL_CLS: ClassVar[Type[Pool]] = NullPool
    TABLE_AVAILABILITY_TIMEOUT: ClassVar[float] = 0.0
    CONFIG_CLS: ClassVar[Type[DbEngineConfig]] = DbEngineConfig

    _config: DbEngineConfig = attr.ib(repr=False)

    # Internals
    _engine: Engine = attr.ib(init=False, default=None)
    _connection: Optional[Connection] = attr.ib(init=False, default=None, repr=False)

    def _db_connect_params(self) -> dict:
        return dict()

    def __attrs_post_init__(self) -> None:
        url = self._config.url
        connect_params = self._db_connect_params()
        if isinstance(url, str):
            if not url.startswith(self.URL_PREFIX):
                raise ValueError(
                    f"Unexpected URL prefix for {type(self)}: expected prefix {self.URL_PREFIX!r} in {url!r}"
                )
            if connect_params:
                url = f"{url}?{urllib.parse.urlencode(connect_params)}"
        else:
            assert isinstance(url, URL)
            if not url.drivername.startswith(self.URL_PREFIX):
                raise ValueError(f"Unexpected URL prefix for {type(self)}: {url}")
            if connect_params:
                url.query.update(connect_params)

        self._engine = sa.create_engine(
            url,
            poolclass=self.POOL_CLS,
            **self._config.engine_params,
        ).execution_options(compiled_cache=None)

    @property
    def config(self) -> DbEngineConfig:
        return self._config

    @property
    def engine(self) -> Engine:
        return self._engine

    @property
    def url(self) -> URL:
        return self._engine.url

    @property
    def dialect(self) -> DefaultDialect:
        return self._engine.dialect

    @property
    def connection(self) -> Connection:
        if self._connection is None:
            self._connection = self._engine.connect()

        assert self._connection is not None
        return self._connection

    def execute(self, query: Any, *multiparams: Any, **params: Any) -> sa.engine.CursorResult:
        return self._engine.execute(query, *multiparams, **params)

    def has_table(self, table_name: str, schema: Optional[str] = None) -> bool:
        return self.dialect.has_table(self._engine, table_name, schema=schema)

    def load_table(self, table_name: str, schema: Optional[str] = None) -> sa.Table:
        return sa.Table(table_name, sa.MetaData(bind=self.engine), schema=schema, autoload=True)

    def create_table(self, table: sa.Table) -> None:
        table.create(bind=self.engine)
        self.wait_for_table(table)

    def table_available(self, table: sa.Table) -> bool:
        """
        Check that the table is available for further actions.
        This might require checks different from `has_table`.
        """
        return self.has_table(table_name=table.name, schema=table.schema)

    def create_schema(self, schema_name: str) -> None:
        self.execute(f"CREATE SCHEMA {self.quote(schema_name)}")

    def wait_for_table(self, table: sa.Table) -> None:
        if self.TABLE_AVAILABILITY_TIMEOUT:
            wait_for(
                name=f"Table {table.name}",
                condition=lambda: self.table_available(table),
                timeout=self.TABLE_AVAILABILITY_TIMEOUT,
            )

    def drop_table(self, db_name: str, table: sa.Table) -> None:
        table.drop(bind=self.engine, checkfirst=True)

    def table_from_columns(
        self,
        columns: Sequence[sa.Column],
        *,
        schema: Optional[str] = None,
        table_name: Optional[str] = None,
    ) -> sa.Table:
        table_name = table_name or f"test_table_{shortuuid.uuid()[:10]}"
        table = sa.Table(table_name, sa.MetaData(), *columns, schema=schema)
        return table

    def count_sql_sessions(self) -> int:
        return 0

    def quote(self, value: str) -> str:
        return self.dialect.identifier_preparer.quote(value)

    def get_conn_credentials(self, full: bool = False) -> dict:
        return dict(
            host=self.url.host,
            port=int(self.url.port),
            username=self.url.username,
            password=self.url.password,
            db_name=self.url.database,
        )

    def get_conn_hosts(self) -> tuple[str, ...]:
        host = self.get_conn_credentials(full=True).get("host")
        if host:
            return tuple(h.strip() for h in host.split(","))
        else:
            return ()

    @contextlib.contextmanager
    def connect(self) -> Generator[Connection, None, None]:
        connection = self._engine.connect(close_with_result=True)
        yield connection
        # Don't close immediately because the cursor may still be in use

    def insert_into_table(self, table: sa.Table, data: Sequence[dict]) -> None:
        self.execute(table.insert(data))

    def get_version(self) -> Optional[str]:
        return self.execute(sa.select([sa.func.version()])).scalar()

    def test(self) -> bool:
        try:
            self.execute(sa.select(1))
        except sa_exc.DatabaseError:
            return False
        else:
            return True


_REGISTERED_ENGINE_WRAPPERS: dict[str, Type[EngineWrapperBase]] = {}


def register_engine_wrapper_cls(engine_wrapper_cls: Type[EngineWrapperBase]) -> None:
    _REGISTERED_ENGINE_WRAPPERS[engine_wrapper_cls.URL_PREFIX] = engine_wrapper_cls


def get_engine_wrapper_cls_for_url(url: str | URL) -> Type[EngineWrapperBase]:
    if isinstance(url, str):
        url = make_url(url)
    assert isinstance(url, URL)
    schema = url.get_backend_name()
    return _REGISTERED_ENGINE_WRAPPERS[schema]

lib/dl_db_testing/dl_db_testing/database/engine_wrapper.py
datalens-tech/datalens-backend
from importlib import metadata
from typing import Type

import dl_db_testing as package
from dl_db_testing.connector_registration import CONN_REG_DB_TESTING
from dl_db_testing.connectors.base.connector import DbTestingConnector


_CONNECTOR_EP_GROUP = f"{package.__name__}.connectors"


def _get_all_ep_connectors() -> dict[str, Type[DbTestingConnector]]:
    entrypoints = list(metadata.entry_points().select(group=_CONNECTOR_EP_GROUP))  # type: ignore
    return {ep.name: ep.load() for ep in entrypoints}


def register_all_connectors() -> None:
    connectors = _get_all_ep_connectors()
    for _, connector_cls in sorted(connectors.items()):
        CONN_REG_DB_TESTING.register_connector(connector_cls)

lib/dl_db_testing/dl_db_testing/db_testing_connectors.py
datalens-tech/datalens-backend
from dl_db_testing.db_testing_connectors import register_all_connectors


def load_db_testing_lib() -> None:
    """Initialize the library"""
    register_all_connectors()

lib/dl_db_testing/dl_db_testing/loader.py
datalens-tech/datalens-backend
from concurrent.futures import ThreadPoolExecutor
import logging
import os

from aiohttp import web

from dl_file_secure_reader_lib.resources import (
    ping,
    reader,
)


LOGGER = logging.getLogger(__name__)


def create_app() -> web.Application:
    app = web.Application()

    app.router.add_route("*", "/reader/ping", ping.PingView)
    app.router.add_route("*", "/reader/excel", reader.ReaderView)

    app["tpe"] = ThreadPoolExecutor(max_workers=min(32, (os.cpu_count() or 1) * 3 + 4))

    return app

lib/dl_file_secure_reader_lib/dl_file_secure_reader_lib/app.py
datalens-tech/datalens-backend
from __future__ import annotations

from aiohttp import web


class PingView(web.View):
    async def get(self) -> web.Response:
        return web.Response()

lib/dl_file_secure_reader_lib/dl_file_secure_reader_lib/resources/ping.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import io
import logging
from typing import (
    BinaryIO,
    List,
)

from aiohttp import web
from aiohttp.multipart import BodyPartReader
from openpyxl import load_workbook


LOGGER = logging.getLogger(__name__)


def parse_excel_data(data: BinaryIO) -> List:
    result = []
    wb = load_workbook(data, data_only=True, keep_links=False)
    for sheetname in wb.sheetnames:
        sheet = wb[sheetname]
        result.append(
            {
                "sheetname": sheetname,
                "data": [
                    [
                        {
                            "value": str(cell.value) if cell.data_type == "d" else cell.value,
                            "data_type": "i" if isinstance(cell.value, int) else cell.data_type,
                        }
                        for cell in row
                    ]
                    for row in sheet.rows
                ],
            }
        )
    return result


class ReaderView(web.View):
    async def post(self) -> web.StreamResponse:
        loop = asyncio.get_running_loop()
        reader = await self.request.multipart()
        field = await reader.next()
        assert isinstance(field, BodyPartReader)
        assert field.name == "file"

        data = io.BytesIO(bytes(await field.read()))

        tpe = self.request.app["tpe"]
        result = await loop.run_in_executor(tpe, parse_excel_data, data)

        return web.json_response(data=result)

lib/dl_file_secure_reader_lib/dl_file_secure_reader_lib/resources/reader.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio

import aiohttp.pytest_plugin
import aiohttp.test_utils
import pytest


pytest_plugins = ("aiohttp.pytest_plugin",)


try:
    del aiohttp.pytest_plugin.loop
except AttributeError:
    pass


@pytest.fixture
def loop(event_loop):
    asyncio.set_event_loop(event_loop)
    yield event_loop
    # Attept to cover an old version of pytest-asyncio:
    # https://github.com/pytest-dev/pytest-asyncio/commit/51d986cec83fdbc14fa08015424c79397afc7ad9
    asyncio.set_event_loop_policy(None)

lib/dl_file_secure_reader_lib/dl_file_secure_reader_lib_tests/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import os

import aiohttp.pytest_plugin
import aiohttp.test_utils
import pytest

from dl_file_secure_reader_lib.app import create_app


try:
    del aiohttp.pytest_plugin.loop
except AttributeError:
    pass


@pytest.fixture
def web_app(loop, aiohttp_client):
    return loop.run_until_complete(aiohttp_client(create_app()))


@pytest.fixture(scope="session")
def excel_data():
    filename = "data.xlsx"

    dirname = os.path.dirname(os.path.abspath(__file__))
    path = os.path.join(dirname, filename)

    with open(path, "rb") as fd:
        return fd.read()

lib/dl_file_secure_reader_lib/dl_file_secure_reader_lib_tests/db/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest


@pytest.mark.asyncio
async def test_ping(web_app):
    resp = await web_app.get("/reader/ping")
    assert resp.status == 200

lib/dl_file_secure_reader_lib/dl_file_secure_reader_lib_tests/db/test_ping.py
datalens-tech/datalens-backend
from __future__ import annotations

import aiohttp
import pytest


@pytest.mark.asyncio
async def test_reader(web_app, excel_data):
    with aiohttp.MultipartWriter() as mpwriter:
        part = mpwriter.append(excel_data)
        part.set_content_disposition("form-data", name="file")
        resp = await web_app.post("/reader/excel", data=mpwriter)

    assert resp.status == 200
    result = await resp.json()
    assert len(result) == 3
    spreadsheet = result[2]
    assert spreadsheet["sheetname"] == "People"
    assert spreadsheet["data"][0] == [
        {"value": "Person", "data_type": "s"},
        {"value": "Region", "data_type": "s"},
    ]

lib/dl_file_secure_reader_lib/dl_file_secure_reader_lib_tests/db/test_reader.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    ClassVar,
    Optional,
)

from aiohttp import web
import arq
from arq.connections import RedisSettings as ArqRedisSettings
import attr

from dl_task_processor.arq_wrapper import create_redis_pool


LOGGER = logging.getLogger(__name__)


@attr.s
class ArqRedisService:
    APP_KEY: ClassVar[str] = "ARQ_REDIS_SERVICE"
    _arq_settings: ArqRedisSettings = attr.ib()
    _arq_pool: Optional[arq.ArqRedis] = attr.ib(init=False, default=None)

    async def init_hook(self, target_app: web.Application) -> None:
        target_app[self.APP_KEY] = self
        await self.initialize()

    async def tear_down_hook(self, target_app: web.Application) -> None:
        await self.tear_down()

    def get_arq_pool(self) -> arq.ArqRedis:
        assert self._arq_pool is not None
        return self._arq_pool

    @classmethod
    def get_app_instance(cls, app: web.Application) -> ArqRedisService:
        service: Optional[ArqRedisService] = app.get(cls.APP_KEY, None)
        if service is None:
            raise ValueError("Redis was not initiated for application")

        return service

    async def initialize(self) -> None:
        LOGGER.info("Initializing Redis ARQ pool")
        self._arq_pool = await create_redis_pool(self._arq_settings)

    async def tear_down(self) -> None:
        LOGGER.info("Close Redis ARQ pool")
        assert self._arq_pool is not None
        await self._arq_pool.close()

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/aiohttp_services/arq_redis.py
datalens-tech/datalens-backend
from typing import ClassVar

from aiohttp import web
import attr

from dl_configs.crypto_keys import CryptoKeysConfig


@attr.s
class CryptoService:
    APP_KEY: ClassVar[str] = "CRYPTO"

    crypto_keys_config: CryptoKeysConfig = attr.ib(repr=False)

    def bind_to_app(self, app: web.Application) -> None:
        app[self.APP_KEY] = self

    @classmethod
    def get_config_for_app(cls, app: web.Application) -> CryptoKeysConfig:
        crypto_service = app[cls.APP_KEY]
        assert isinstance(crypto_service, CryptoService)
        return crypto_service.crypto_keys_config

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/aiohttp_services/crypto.py
datalens-tech/datalens-backend
from http import HTTPStatus

import aiogoogle
from aiohttp import web
from marshmallow import ValidationError as MValidationError

from dl_api_commons.aio.middlewares.error_handling_outer import (
    AIOHTTPErrorHandler,
    ErrorData,
    ErrorLevel,
)
from dl_file_uploader_lib import exc
from dl_file_uploader_lib.gsheets_client import google_api_error_to_file_uploader_exception
from dl_file_uploader_lib.redis_model.base import (
    RedisModelAccessDenied,
    RedisModelNotFound,
)
from dl_file_uploader_lib.redis_model.models import models


STATUS_CODES = {
    exc.CannotUpdateDataError: HTTPStatus.BAD_REQUEST,
    exc.DocumentNotFound: HTTPStatus.BAD_REQUEST,
    exc.DownloadFailed: HTTPStatus.INTERNAL_SERVER_ERROR,
    exc.EmptyDocument: HTTPStatus.BAD_REQUEST,
    exc.FileLimitError: HTTPStatus.REQUEST_ENTITY_TOO_LARGE,
    exc.InvalidFieldCast: HTTPStatus.BAD_REQUEST,
    exc.InvalidLink: HTTPStatus.BAD_REQUEST,
    exc.InvalidSource: HTTPStatus.BAD_REQUEST,
    exc.ParseFailed: HTTPStatus.INTERNAL_SERVER_ERROR,
    exc.PermissionDenied: HTTPStatus.BAD_REQUEST,
    exc.RemoteServerError: HTTPStatus.BAD_REQUEST,
    exc.UnsupportedDocument: HTTPStatus.BAD_REQUEST,
}


class FileUploaderErrorHandler(AIOHTTPErrorHandler):
    def _classify_error(self, err: Exception, request: web.Request) -> ErrorData:
        if isinstance(err, aiogoogle.HTTPError):
            err = google_api_error_to_file_uploader_exception(err)

        if isinstance(err, web.HTTPException):
            return ErrorData(
                status_code=err.status,
                http_reason=err.reason,
                response_body=dict(message=err.reason),
                level=ErrorLevel.info,
            )
        elif isinstance(err, MValidationError):
            return ErrorData(
                status_code=HTTPStatus.BAD_REQUEST,
                response_body=dict(message=str(err)),
                level=ErrorLevel.info,
            )
        elif isinstance(err, RedisModelAccessDenied):
            return ErrorData(HTTPStatus.FORBIDDEN, response_body={}, level=ErrorLevel.info)
        elif isinstance(err, RedisModelNotFound):
            return ErrorData(HTTPStatus.NOT_FOUND, response_body={}, level=ErrorLevel.info)
        elif isinstance(err, (models.EmptySourcesError, models.SourceNotFoundError)):
            return ErrorData(HTTPStatus.NOT_FOUND, response_body={}, level=ErrorLevel.info)
        elif isinstance(err, exc.DLFileUploaderBaseError):
            status_code = STATUS_CODES.get(err.__class__, HTTPStatus.INTERNAL_SERVER_ERROR)  # type: ignore
            body = dict(
                message=err.message,
                code=exc.make_err_code(err),
                debug=err.debug_info,
                details=err.details,
            )
            return ErrorData(
                status_code=status_code,
                response_body=body,
                level=ErrorLevel.info,
            )
        else:
            return ErrorData(
                HTTPStatus.INTERNAL_SERVER_ERROR,
                http_reason="Internal server error",
                response_body=dict(message="Internal server error"),
                level=ErrorLevel.error,
            )

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/aiohttp_services/error_handler.py
datalens-tech/datalens-backend
import abc
from typing import (
    Generic,
    Type,
    TypeVar,
)

from aiohttp import web
import attr

from dl_api_commons.aio.middlewares.commit_rci import commit_rci_middleware
from dl_api_commons.aio.middlewares.cors import cors_middleware
from dl_api_commons.aio.middlewares.csrf import CSRFMiddleware
from dl_api_commons.aio.middlewares.request_bootstrap import RequestBootstrap
from dl_api_commons.aio.middlewares.request_id import RequestId
from dl_api_commons.aio.middlewares.tracing import TracingService
from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_api_commons.sentry_config import (
    SentryConfig,
    configure_sentry_for_aiohttp,
)
from dl_constants.api_constants import DLHeadersCommon
from dl_core.aio.metrics_view import MetricsView
from dl_core.aio.middlewares.master_key import master_key_middleware
from dl_core.aio.ping_view import PingView
from dl_core.aio.web_app_services.s3 import S3Service
from dl_core.loader import load_core_lib
from dl_file_uploader_api_lib.aiohttp_services.arq_redis import ArqRedisService
from dl_file_uploader_api_lib.aiohttp_services.crypto import CryptoService
from dl_file_uploader_api_lib.aiohttp_services.error_handler import FileUploaderErrorHandler
from dl_file_uploader_api_lib.dl_request import FileUploaderDLRequest
from dl_file_uploader_api_lib.settings import FileUploaderAPISettings
from dl_file_uploader_api_lib.views import files as files_views
from dl_file_uploader_api_lib.views import misc as misc_views
from dl_file_uploader_api_lib.views import sources as sources_views
from dl_file_uploader_lib.settings_utils import init_redis_service
from dl_task_processor.arq_wrapper import create_arq_redis_settings


_TSettings = TypeVar("_TSettings", bound=FileUploaderAPISettings)


@attr.s(kw_only=True)
class FileUploaderApiAppFactory(Generic[_TSettings], abc.ABC):
    CSRF_MIDDLEWARE_CLS: Type[CSRFMiddleware] = CSRFMiddleware

    _settings: _TSettings = attr.ib()

    @abc.abstractmethod
    def get_auth_middlewares(self) -> list[AIOHTTPMiddleware]:
        raise NotImplementedError()

    def set_up_sentry(self, secret_sentry_dsn: str, release: str) -> None:
        configure_sentry_for_aiohttp(
            SentryConfig(
                dsn=secret_sentry_dsn,
                release=release,
            )
        )

    def create_app(self, app_version: str) -> web.Application:
        load_core_lib()

        if (secret_sentry_dsn := self._settings.SENTRY_DSN) is not None:
            self.set_up_sentry(secret_sentry_dsn, app_version)

        req_id_service = RequestId(
            append_own_req_id=True,
            app_prefix="f",
            dl_request_cls=FileUploaderDLRequest,
        )

        error_handler = FileUploaderErrorHandler(
            use_sentry=(secret_sentry_dsn is not None),
            sentry_app_name_tag="file-uploader",
        )

        middleware_list = [
            TracingService().middleware,
            cors_middleware(
                allow_origins=self._settings.CORS.ALLOWED_ORIGINS,
                expose_headers=self._settings.CORS.EXPOSE_HEADERS,
                allow_headers=self._settings.CORS.ALLOWED_HEADERS,
                allow_credentials=True,
            ),
            RequestBootstrap(
                req_id_service=req_id_service,
                error_handler=error_handler,
            ).middleware,
            master_key_middleware(
                master_key=self._settings.FILE_UPLOADER_MASTER_TOKEN,
                header=DLHeadersCommon.FILE_UPLOADER_MASTER_TOKEN,
            ),
            *self.get_auth_middlewares(),
            commit_rci_middleware(),
            self.CSRF_MIDDLEWARE_CLS(
                csrf_header_name=self._settings.CSRF.HEADER_NAME,
                csrf_time_limit=self._settings.CSRF.TIME_LIMIT,
                csrf_secret=self._settings.CSRF.SECRET,
                csrf_methods=self._settings.CSRF.METHODS,
            ).middleware,
            # TODO FIX: Add when json_body_middleware will be moved to dl_core
            # json_body_middleware(),
        ]

        app = web.Application(
            middlewares=middleware_list,  # type: ignore
        )
        app.on_response_prepare.append(req_id_service.on_response_prepare)

        redis_service = init_redis_service(self._settings)
        app.on_startup.append(redis_service.init_hook)
        app.on_shutdown.append(redis_service.tear_down_hook)

        s3_service = S3Service(
            access_key_id=self._settings.S3.ACCESS_KEY_ID,
            secret_access_key=self._settings.S3.SECRET_ACCESS_KEY,
            endpoint_url=self._settings.S3.ENDPOINT_URL,
            tmp_bucket_name=self._settings.S3_TMP_BUCKET_NAME,
            persistent_bucket_name=self._settings.S3_PERSISTENT_BUCKET_NAME,
        )
        app.on_startup.append(s3_service.init_hook)
        app.on_shutdown.append(s3_service.tear_down_hook)

        arq_redis_service = ArqRedisService(arq_settings=create_arq_redis_settings(self._settings.REDIS_ARQ))  # type: ignore
        app.on_startup.append(arq_redis_service.init_hook)
        app.on_shutdown.append(arq_redis_service.tear_down_hook)

        CryptoService(crypto_keys_config=self._settings.CRYPTO_KEYS_CONFIG).bind_to_app(app)

        app.router.add_route("get", "/api/v2/ping", PingView)
        app.router.add_route("get", "/api/v2/metrics", MetricsView)

        app.router.add_route("post", "/api/v2/files", files_views.FilesView)
        app.router.add_route("post", "/api/v2/links", files_views.LinksView)
        app.router.add_route("post", "/api/v2/update_connection_data", files_views.UpdateConnectionDataView)
        app.router.add_route(
            "post", "/api/v2/update_connection_data_internal", files_views.InternalUpdateConnectionDataView
        )
        app.router.add_route("get", "/api/v2/files/{file_id}/status", files_views.FileStatusView)
        app.router.add_route("get", "/api/v2/files/{file_id}/sources", files_views.FileSourcesView)

        app.router.add_route(
            "get", "/api/v2/files/{file_id}/sources/{source_id}/status", sources_views.SourceStatusView
        )
        app.router.add_route("post", "/api/v2/files/{file_id}/sources/{source_id}", sources_views.SourceInfoView)
        app.router.add_route(
            "post", "/api/v2/files/{file_id}/sources/{source_id}/preview", sources_views.SourcePreviewView
        )
        app.router.add_route(
            "post",
            "/api/v2/files/{file_id}/sources/{source_id}/internal_params",
            sources_views.SourceInternalParamsView,
        )
        app.router.add_route(
            "post", "/api/v2/files/{file_id}/sources/{source_id}/apply_settings", sources_views.SourceApplySettingsView
        )

        app.router.add_route("post", "/api/v2/cleanup", misc_views.CleanupTenantView)
        app.router.add_route("post", "/api/v2/rename_tenant_files", misc_views.RenameTenantFilesView)

        app["ALLOW_XLSX"] = self._settings.ALLOW_XLSX

        return app

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/app.py
datalens-tech/datalens-backend
import logging
from typing import Optional
import urllib.parse

from dl_constants.enums import FileProcessingStatus
from dl_file_uploader_lib.enums import FileType
from dl_file_uploader_lib.exc import InvalidLink
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models import DataFile
from dl_file_uploader_lib.redis_model.models.models import GSheetsUserSourceProperties


LOGGER = logging.getLogger(__name__)


async def gsheets_data_file_preparer(
    url: str,
    redis_model_manager: RedisModelManager,
    token: Optional[str],
) -> DataFile:
    default_host: str = "docs.google.com"
    allowed_hosts: frozenset[str] = frozenset({default_host})
    example_url: str = (
        "https://docs.google.com/spreadsheets/d/1zRPTxxLOQ08n0_MReIBbouAbPFw6pJ4ibNSVdFkZ3gs/edit#gid=1140182768"
    )
    example_url_message: str = f"example: {example_url!r}"

    parts = urllib.parse.urlparse(url)
    if parts.scheme not in ("http", "https"):
        raise InvalidLink(f"Invalid url scheme: {parts.scheme!r}, must be 'http' or 'https'; {example_url_message}")
    netloc_pieces = parts.netloc.split(":", 1)
    host = netloc_pieces[0]
    if len(netloc_pieces) == 2:
        port = netloc_pieces[1]
        if port not in ("80", "443"):
            raise InvalidLink(f"Invalid port: {port!r}. Should be omitted or default; {example_url_message}")
    if host not in allowed_hosts:
        raise InvalidLink(f"Invalid host: {host!r}; should be {default_host!r}; {example_url_message}")

    path = parts.path
    prefix = "/spreadsheets/d/"
    if not path.startswith(prefix):
        raise InvalidLink(f"Invalid URL path prefix: {path!r}; should be {prefix!r}; {example_url_message}")
    path = path[len(prefix) :]

    suffix = "/edit"
    spreadsheet_id = path[: -len(suffix)] if path.endswith(suffix) else path

    df = DataFile(
        manager=redis_model_manager,
        filename="Spreadsheet",
        status=FileProcessingStatus.in_progress,
        file_type=FileType.gsheets,
        user_source_properties=GSheetsUserSourceProperties(
            spreadsheet_id=spreadsheet_id,
            refresh_token=token,
        ),
    )

    return df

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/data_file_preparer.py
datalens-tech/datalens-backend
from dl_core.aio.aiohttp_wrappers_data_core import DLRequestDataCore
from dl_core.aio.web_app_services.s3 import S3Service
from dl_file_uploader_api_lib.aiohttp_services.arq_redis import ArqRedisService
from dl_file_uploader_api_lib.aiohttp_services.crypto import CryptoService
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_task_processor.processor import (
    TaskProcessor,
    make_task_processor,
)


class FileUploaderDLRequest(DLRequestDataCore):
    def get_s3_service(self) -> S3Service:
        return S3Service.get_app_instance(self.request.app)

    def get_task_processor(self) -> TaskProcessor:
        return make_task_processor(
            redis_pool=ArqRedisService.get_app_instance(self.request.app).get_arq_pool(),
            request_id=self.rci.request_id,
        )

    def get_redis_model_manager(self) -> RedisModelManager:
        return RedisModelManager(
            redis=self.get_persistent_redis(),
            rci=self.rci,
            crypto_keys_config=CryptoService.get_config_for_app(self.request.app),
        )

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/dl_request.py
datalens-tech/datalens-backend
import marshmallow as ma


class BaseRequestSchema(ma.Schema):
    class Meta:
        unknown = ma.EXCLUDE

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/schemas/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import marshmallow as ma

from dl_constants.exc import GLOBAL_ERR_PREFIX
from dl_core.marshmallow import ErrorCodeField
from dl_file_uploader_api_lib.schemas.base import BaseRequestSchema
from dl_file_uploader_lib.enums import (
    ErrorLevel,
    ErrorObjectKind,
)


class ErrorInfoSchema(BaseRequestSchema):
    code = ma.fields.String()
    object_kind = ma.fields.Enum(enum=ErrorObjectKind)
    object_id = ma.fields.String()
    message = ma.fields.String()
    details = ma.fields.Dict(keys=ma.fields.String())
    level = ma.fields.Enum(enum=ErrorLevel)


class FileProcessingErrorApiSchema(ma.Schema):
    level = ma.fields.Enum(ErrorLevel)
    message = ma.fields.String()
    code = ErrorCodeField(prefix=(GLOBAL_ERR_PREFIX,))
    details = ma.fields.Dict()

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/schemas/errors.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Any

import marshmallow as ma

from dl_constants.enums import FileProcessingStatus
from dl_file_uploader_api_lib.schemas.base import BaseRequestSchema
from dl_file_uploader_api_lib.schemas.errors import (
    ErrorInfoSchema,
    FileProcessingErrorApiSchema,
)
from dl_file_uploader_lib import exc
from dl_file_uploader_lib.enums import FileType


def validate_authorized(data: dict) -> None:
    if data["authorized"] and data["refresh_token"] is None and data["connection_id"] is None:
        raise ma.ValidationError(
            "Either refresh_token or connection_id must be provided when authorized is true",
            "authorized",
        )


class FileLinkRequestSchema(BaseRequestSchema):
    type = ma.fields.Enum(FileType, required=True)
    url = ma.fields.String(required=True)
    refresh_token = ma.fields.String(load_default=None, allow_none=True)
    connection_id = ma.fields.String(load_default=None, allow_none=True)
    authorized = ma.fields.Boolean(required=True)

    @ma.validates_schema(skip_on_field_errors=True)
    def validate_object(self, data: dict, **kwargs: Any) -> None:
        validate_authorized(data)


class FileUploadResponseSchema(ma.Schema):
    file_id = ma.fields.String()
    title = ma.fields.String()


class FileStatusRequestSchema(BaseRequestSchema):
    file_id = ma.fields.String(required=True)


class FileStatusResultSchema(ma.Schema):
    file_id = ma.fields.String(attribute="id")
    status = ma.fields.Enum(enum=FileProcessingStatus)
    errors = ma.fields.Nested(ErrorInfoSchema, many=True)
    error = ma.fields.Nested(FileProcessingErrorApiSchema, allow_none=True)
    # progress = ma.fields.Float()


class FileSourcesRequestSchema(BaseRequestSchema):
    file_id = ma.fields.String(required=True)


class SourceShortInfoSchema(ma.Schema):
    source_id = ma.fields.String(attribute="id")
    title = ma.fields.String()
    is_applicable = ma.fields.Boolean()
    error = ma.fields.Nested(FileProcessingErrorApiSchema, allow_none=True)


class FileSourcesResultSchema(ma.Schema):
    file_id = ma.fields.String(attribute="id")
    errors = ma.fields.Nested(ErrorInfoSchema, many=True)
    sources = ma.fields.Nested(SourceShortInfoSchema, many=True)


class UpdateConnectionDataRequestSchema(BaseRequestSchema):
    class UpdateConnectionDataSourceSchema(BaseRequestSchema):
        id = ma.fields.String()
        title = ma.fields.String()
        spreadsheet_id = ma.fields.String(load_default=None)
        sheet_id = ma.fields.Integer(load_default=None)
        first_line_is_header = ma.fields.Boolean(load_default=None)

    connection_id = ma.fields.String(allow_none=True, load_default=None)
    refresh_token = ma.fields.String(allow_none=True, load_default=None)
    authorized = ma.fields.Boolean(required=True)
    save = ma.fields.Boolean(load_default=False)
    sources = ma.fields.Nested(UpdateConnectionDataSourceSchema, many=True)
    tenant_id = ma.fields.String(allow_none=True, load_default=None)

    @ma.validates_schema(skip_on_field_errors=True)
    def validate_object(self, data: dict, **kwargs: Any) -> None:
        validate_authorized(data)

        incomplete_sources: list[dict[str, str]] = []
        for src in data["sources"]:
            if src["spreadsheet_id"] is None or src["sheet_id"] is None or src["first_line_is_header"] is None:
                incomplete_sources.append(
                    dict(
                        source_id=src["id"],
                        title=src["title"],
                    )
                )

        if incomplete_sources:
            raise exc.CannotUpdateDataError(
                details=dict(
                    incomplete_sources=incomplete_sources,
                )
            )


class UpdateConnectionDataResultSchema(ma.Schema):
    class FileSourcesSchema(ma.Schema):
        class SingleFileSourceSchema(ma.Schema):
            source_id = ma.fields.String(attribute="id")

        file_id = ma.fields.String(attribute="id")
        sources = ma.fields.Nested(SingleFileSourceSchema, many=True)

    files = ma.fields.Nested(FileSourcesSchema, many=True)

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/schemas/files.py
datalens-tech/datalens-backend
import marshmallow as ma

from dl_file_uploader_api_lib.schemas.base import BaseRequestSchema


class CleanupRequestSchema(BaseRequestSchema):
    tenant_id = ma.fields.String(required=True)


class RenameFilesRequestSchema(BaseRequestSchema):
    tenant_id = ma.fields.String(required=True)

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/schemas/misc.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Type,
)

import marshmallow as ma
from marshmallow_oneofschema import OneOfSchema

from dl_constants.enums import (
    FileProcessingStatus,
    UserDataType,
)
from dl_core.db.elements import SchemaColumn
from dl_core.db.native_type_schema import OneOfNativeTypeSchema
from dl_file_uploader_api_lib.schemas.base import BaseRequestSchema
from dl_file_uploader_api_lib.schemas.errors import (
    ErrorInfoSchema,
    FileProcessingErrorApiSchema,
)
from dl_file_uploader_lib.enums import (
    CSVDelimiter,
    CSVEncoding,
    FileType,
)


class SourceInfoBaseRequestSchema(BaseRequestSchema):
    file_id = ma.fields.String(required=True)
    source_id = ma.fields.String(required=True)


class SourceStatusResultSchema(ma.Schema):
    file_id = ma.fields.String()
    source_id = ma.fields.String()
    status = ma.fields.Enum(enum=FileProcessingStatus)
    errors = ma.fields.Nested(ErrorInfoSchema, many=True)
    error = ma.fields.Nested(FileProcessingErrorApiSchema, allow_none=True)
    # progress = ma.fields.Float()


class RawSchemaColumnSchema(ma.Schema):
    name = ma.fields.String()
    title = ma.fields.String(load_default=None, allow_none=True)

    native_type = ma.fields.Nested(OneOfNativeTypeSchema, allow_none=True, load_default=None)

    user_type = ma.fields.Enum(UserDataType)
    description = ma.fields.String(dump_default="", allow_none=True, load_default="")
    has_auto_aggregation = ma.fields.Boolean(dump_default=False, allow_none=True, load_default=False)
    lock_aggregation = ma.fields.Boolean(dump_default=False, allow_none=True, load_default=False)
    nullable = ma.fields.Boolean(dump_default=None, allow_none=True, load_default=None)

    @ma.post_load
    def make_column(self, data: dict, **kwargs: Any) -> SchemaColumn:
        return SchemaColumn(
            name=data["name"],
            title=data["title"],
            user_type=data["user_type"],
            native_type=data["native_type"],
            description=data.get("description", ""),
            has_auto_aggregation=data.get("has_auto_aggregation", False),
            lock_aggregation=data.get("lock_aggregation", False),
            nullable=data["nullable"],
        )


class SourceInfoRequestSchema(SourceInfoBaseRequestSchema):
    column_types = ma.fields.Nested(RawSchemaColumnSchema, many=True, load_default=[])


class FileTypeOneOfSchema(OneOfSchema):
    class Meta:
        unknown = ma.EXCLUDE

    type_field_remove = True
    type_field = "file_type"

    def get_obj_type(self, obj: dict[str, Any]) -> str:
        type_field = obj[self.type_field] if isinstance(obj, dict) else getattr(obj, self.type_field)
        assert isinstance(type_field, FileType)
        return type_field.name


class SourceInfoSchemaBase(ma.Schema):
    class RawSchemaColumnSchemaShorten(ma.Schema):
        name = ma.fields.String()
        title = ma.fields.String()
        user_type = ma.fields.Enum(UserDataType)

    source_id = ma.fields.String()
    title = ma.fields.String()
    is_valid = ma.fields.Boolean()
    raw_schema = ma.fields.Nested(RawSchemaColumnSchemaShorten, many=True)
    preview = ma.fields.List(ma.fields.List(ma.fields.Raw(allow_none=True)))
    error = ma.fields.Nested(FileProcessingErrorApiSchema, allow_none=True)


class SourceInfoSchemaGSheets(SourceInfoSchemaBase):
    sheet_id = ma.fields.Integer()
    spreadsheet_id = ma.fields.String()


class SourceInfoSchema(FileTypeOneOfSchema):
    type_schemas: dict[str, Type[SourceInfoSchemaBase]] = {
        FileType.csv.name: SourceInfoSchemaBase,
        FileType.gsheets.name: SourceInfoSchemaGSheets,
        FileType.xlsx.name: SourceInfoSchemaBase,
        FileType.yadocs.name: SourceInfoSchemaBase,
    }


class CSVSettingsSchema(ma.Schema):
    encoding = ma.fields.Enum(CSVEncoding)
    delimiter = ma.fields.Enum(CSVDelimiter)
    first_line_is_header = ma.fields.Boolean()


class CSVSettingsOptionsSchema(ma.Schema):
    encoding = ma.fields.List(ma.fields.Enum(CSVEncoding))
    delimiter = ma.fields.List(ma.fields.Enum(CSVDelimiter))


class OptionsSchema(ma.Schema):
    class ColumnsOptionsSchema(ma.Schema):
        name = ma.fields.String()
        user_type = ma.fields.List(ma.fields.Enum(UserDataType))

    data_settings = ma.fields.Nested(CSVSettingsOptionsSchema)
    columns = ma.fields.Nested(ColumnsOptionsSchema, many=True)


class SourceInfoResultSchema(ma.Schema):
    file_id = ma.fields.String()
    file_type = ma.fields.Enum(FileType)
    errors = ma.fields.Nested(ErrorInfoSchema, many=True)
    source = ma.fields.Nested(SourceInfoSchema)
    data_settings = ma.fields.Nested(CSVSettingsSchema)  # TODO: One of schema by file_type
    options = ma.fields.Nested(OptionsSchema, allow_none=True)


class SourceApplySettingsRequestSchema(SourceInfoBaseRequestSchema):
    data_settings = ma.fields.Nested(CSVSettingsSchema)
    title = ma.fields.String(required=False, load_default=None, allow_none=True)


class SourcePreviewRequestSchema(SourceInfoBaseRequestSchema):
    preview_id = ma.fields.String()
    raw_schema = ma.fields.Nested(RawSchemaColumnSchema, many=True, required=True)


class SourcePreviewResultSchema(ma.Schema):
    preview = ma.fields.List(ma.fields.List(ma.fields.Raw(allow_none=True)))


class SourceInternalParamsRequestSchema(SourceInfoBaseRequestSchema):
    raw_schema = ma.fields.Nested(RawSchemaColumnSchema, many=True, load_default=[], allow_none=True)


class SourceInternalParamsResultSchema(ma.Schema):
    preview_id = ma.fields.String()
    raw_schema = ma.fields.Nested(RawSchemaColumnSchema, many=True)

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/schemas/sources.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_configs.environments import is_setting_applicable
from dl_configs.settings_loaders.meta_definition import (
    required,
    s_attrib,
)
from dl_configs.settings_submodels import (
    CorsSettings,
    CsrfSettings,
)
from dl_file_uploader_lib.settings import FileUploaderBaseSettings


@attr.s(frozen=True)
class FileUploaderAPISettings(FileUploaderBaseSettings):
    CORS: CorsSettings = s_attrib(  # type: ignore
        "CORS",
        fallback_factory=(
            lambda cfg: CorsSettings(  # type: ignore
                ALLOWED_ORIGINS=cfg.CORS_ALLOWED_ORIGINS,
                ALLOWED_HEADERS=cfg.CORS_ALLOWED_HEADERS,
                EXPOSE_HEADERS=cfg.CORS_EXPOSE_HEADERS,
                # TODO: move this values to a separate key
            )
            if is_setting_applicable(cfg, "CORS_ALLOWED_ORIGINS")
            else None
        ),
    )

    CSRF: CsrfSettings = s_attrib(  # type: ignore
        "CSRF",
        fallback_factory=(
            lambda cfg: CsrfSettings(  # type: ignore
                METHODS=cfg.CSRF_METHODS,
                HEADER_NAME=cfg.CSRF_HEADER_NAME,
                TIME_LIMIT=cfg.CSRF_TIME_LIMIT,
                SECRET=required(str),
                # TODO: move this values to a separate key
            )
            if is_setting_applicable(cfg, "CSRF_METHODS")
            else None
        ),
    )

    SENTRY_DSN: Optional[str] = s_attrib(  # type: ignore
        "SENTRY_DSN",
        fallback_cfg_key="SENTRY_DSN_FILE_UPLOADER_API",
        missing=None,
    )

    FILE_UPLOADER_MASTER_TOKEN: str = s_attrib("FILE_UPLOADER_MASTER_TOKEN", sensitive=True)  # type: ignore

    ALLOW_XLSX: bool = s_attrib("ALLOW_XLSX", missing=False)  # type: ignore

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/settings.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
    Type,
)

from marshmallow import Schema

from dl_api_commons.aiohttp.aiohttp_wrappers import (
    DLRequestView,
    RequiredResource,
    RequiredResourceCommon,
)
from dl_file_uploader_api_lib.dl_request import FileUploaderDLRequest


class FileUploaderBaseView(DLRequestView[FileUploaderDLRequest]):
    dl_request_cls = FileUploaderDLRequest

    REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset({RequiredResourceCommon.SKIP_CSRF})

    async def _load_post_request_schema_data(self, schema_cls: Type[Schema]) -> dict[str, Any]:
        json_data = await self.request.json() if self.request.can_read_body else {}
        req_data = schema_cls().load(
            {
                **self.request.match_info,
                **json_data,
            }
        )
        return req_data

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/views/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from http import HTTPStatus
import logging
from typing import (
    AsyncGenerator,
    Awaitable,
    Callable,
    ClassVar,
    Optional,
)

from aiohttp import web
from aiohttp.multipart import BodyPartReader
from redis.asyncio.lock import Lock as RedisLock

from dl_api_commons.aiohttp.aiohttp_wrappers import (
    RequiredResource,
    RequiredResourceCommon,
)
from dl_constants.enums import FileProcessingStatus
from dl_file_uploader_api_lib.data_file_preparer import gsheets_data_file_preparer
from dl_file_uploader_api_lib.schemas import files as files_schemas
from dl_file_uploader_api_lib.views.base import FileUploaderBaseView
from dl_file_uploader_lib.common_locks import get_update_connection_source_lock
from dl_file_uploader_lib.data_sink.raw_bytes import (
    RawBytesAsyncDataStream,
    S3RawFileAsyncDataSink,
)
from dl_file_uploader_lib.enums import FileType
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    DataSource,
    GSheetsFileSourceSettings,
    GSheetsUserSourceDataSourceProperties,
    GSheetsUserSourceProperties,
)
from dl_file_uploader_task_interface.tasks import (
    DownloadGSheetTask,
    ParseFileTask,
    ProcessExcelTask,
    TaskExecutionMode,
)


LOGGER = logging.getLogger(__name__)


def get_file_type_from_name(
    filename: Optional[str],
    default_type: FileType = FileType.csv,
    allow_xlsx: bool = False,
) -> FileType:
    FILE_TYPE_MAPPER = {"csv": FileType.csv}
    if allow_xlsx:
        FILE_TYPE_MAPPER["xlsx"] = FileType.xlsx
    file_extension = "csv"
    if isinstance(filename, str) and "." in filename:
        file_extension = filename.split(".")[-1]
    file_type = FILE_TYPE_MAPPER.get(file_extension, default_type)
    return file_type


class FilesView(FileUploaderBaseView):
    REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset()  # Don't skip CSRF check

    async def post(self) -> web.StreamResponse:
        reader = await self.request.multipart()
        field = await reader.next()
        assert isinstance(field, BodyPartReader)
        assert field.name == "file"
        filename = field.filename
        file_type = get_file_type_from_name(filename=filename, allow_xlsx=self.request.app["ALLOW_XLSX"])

        s3 = self.dl_request.get_s3_service()

        rmm = self.dl_request.get_redis_model_manager()
        df = DataFile(
            manager=rmm,
            filename=filename,
            file_type=file_type,
            status=FileProcessingStatus.in_progress,
        )
        LOGGER.info(f"Data file id: {df.id}")

        async def _chunk_iter(chunk_size: int = 10 * 1024 * 1024) -> AsyncGenerator[bytes, None]:
            assert isinstance(field, BodyPartReader)
            while True:
                chunk = await field.read_chunk(size=chunk_size)
                if chunk:
                    LOGGER.debug(f"Received chunk of {len(chunk)} bytes.")
                    yield chunk
                else:
                    LOGGER.info("Empty chunk received.")
                    break

        data_stream = RawBytesAsyncDataStream(data_iter=_chunk_iter())
        async with S3RawFileAsyncDataSink(
            s3=s3.client,
            s3_key=df.s3_key,
            bucket_name=s3.tmp_bucket_name,
        ) as data_sink:
            await data_sink.dump_data_stream(data_stream)

        # df.size = size    # TODO

        await df.save()
        LOGGER.info(f'Uploaded file "{filename}".')

        task_processor = self.dl_request.get_task_processor()
        if file_type == FileType.xlsx:
            await task_processor.schedule(ProcessExcelTask(file_id=df.id))
            LOGGER.info(f"Scheduled ProcessExcelTask for file_id {df.id}")
        else:
            await task_processor.schedule(ParseFileTask(file_id=df.id))
            LOGGER.info(f"Scheduled ParseFileTask for file_id {df.id}")

        return web.json_response(
            files_schemas.FileUploadResponseSchema().dump({"file_id": df.id, "title": df.filename}),
            status=HTTPStatus.CREATED,
        )


class LinksView(FileUploaderBaseView):
    REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset()  # Don't skip CSRF check

    FILE_TYPE_TO_DATA_FILE_PREPARER_MAP: dict[
        FileType, Callable[[str, RedisModelManager, Optional[str]], Awaitable[DataFile]]
    ] = {
        FileType.gsheets: gsheets_data_file_preparer,
    }

    async def post(self) -> web.StreamResponse:
        req_data = await self._load_post_request_schema_data(files_schemas.FileLinkRequestSchema)

        file_type = FileType(req_data["type"])

        rmm = self.dl_request.get_redis_model_manager()

        refresh_token: Optional[str] = req_data["refresh_token"]
        connection_id: Optional[str] = req_data["connection_id"]
        authorized: bool = req_data["authorized"]
        df = await self.FILE_TYPE_TO_DATA_FILE_PREPARER_MAP[file_type](req_data["url"], rmm, refresh_token)

        LOGGER.info(f"Data file id: {df.id}")
        await df.save()

        task_processor = self.dl_request.get_task_processor()
        await task_processor.schedule(
            DownloadGSheetTask(
                file_id=df.id,
                authorized=authorized,
                connection_id=connection_id,
            )
        )
        LOGGER.info(f"Scheduled DownloadGSheetTask for file_id {df.id}")

        return web.json_response(
            files_schemas.FileUploadResponseSchema().dump({"file_id": df.id, "title": df.filename}),
            status=HTTPStatus.CREATED,
        )


class FileStatusView(FileUploaderBaseView):
    async def get(self) -> web.StreamResponse:
        req_data = files_schemas.FileStatusRequestSchema().load(self.request.match_info)
        rmm = self.dl_request.get_redis_model_manager()
        df = await DataFile.get_authorized(manager=rmm, obj_id=req_data["file_id"])

        return web.json_response(files_schemas.FileStatusResultSchema().dump(df))


class FileSourcesView(FileUploaderBaseView):
    async def get(self) -> web.StreamResponse:
        req_data = files_schemas.FileSourcesRequestSchema().load(self.request.match_info)
        rmm = self.dl_request.get_redis_model_manager()
        df = await DataFile.get_authorized(manager=rmm, obj_id=req_data["file_id"])

        return web.json_response(files_schemas.FileSourcesResultSchema().dump(df))


class UpdateConnectionDataView(FileUploaderBaseView):
    _is_internal: ClassVar[bool] = False

    async def post(self) -> web.StreamResponse:
        req_data = await self._load_post_request_schema_data(files_schemas.UpdateConnectionDataRequestSchema)
        sources = req_data["sources"]

        rmm = self.dl_request.get_redis_model_manager()
        redis = self.dl_request.get_persistent_redis()

        dfile_by_spreadsheet: dict[str, DataFile] = {}
        for src in sources:
            source_lock_key, source_lock_token = get_update_connection_source_lock(src["id"])
            LOGGER.info(f"Acquiring redis lock {source_lock_key}")
            # source locks will be released by file uploader worker when the update process is finished
            source_lock = RedisLock(redis, name=source_lock_key, timeout=120, blocking=False)
            if not await source_lock.acquire(token=source_lock_token):
                LOGGER.warning(f"Cannot acquire lock {source_lock_key}, skipping the source")
                continue
            LOGGER.info(f"Lock {source_lock_key} acquired")

            if src["spreadsheet_id"] not in dfile_by_spreadsheet:
                dfile_by_spreadsheet[src["spreadsheet_id"]] = DataFile(
                    user_id=DataFile.system_user_id if self._is_internal else None,  # filled from rci when None
                    manager=rmm,
                    filename="TITLE",
                    status=FileProcessingStatus.in_progress,
                    file_type=FileType.gsheets,
                    sources=[],
                    user_source_properties=GSheetsUserSourceProperties(
                        spreadsheet_id=src["spreadsheet_id"],
                        refresh_token=req_data["refresh_token"],
                    ),
                )
            LOGGER.info(f'Data file id: {dfile_by_spreadsheet[src["spreadsheet_id"]].id}')

            sheet_data_source = DataSource(
                id=src["id"],
                title=src["title"],
                raw_schema=[],
                file_source_settings=GSheetsFileSourceSettings(
                    first_line_is_header=src["first_line_is_header"],
                    raw_schema_header=[],
                    raw_schema_body=[],
                ),
                user_source_dsrc_properties=GSheetsUserSourceDataSourceProperties(
                    sheet_id=src["sheet_id"],
                ),
                status=FileProcessingStatus.in_progress,
                error=None,
            )

            dfile_by_spreadsheet[src["spreadsheet_id"]].sources.append(sheet_data_source)  # type: ignore

        task_processor = self.dl_request.get_task_processor()
        exec_mode = TaskExecutionMode.UPDATE_AND_SAVE if req_data["save"] else TaskExecutionMode.UPDATE_NO_SAVE
        if req_data.get("tenant_id") is not None:
            tenant_id = req_data["tenant_id"]
        else:
            assert self.dl_request.rci.tenant is not None
            tenant_id = self.dl_request.rci.tenant.get_tenant_id()
        assert tenant_id is not None
        for dfile in dfile_by_spreadsheet.values():
            await dfile.save()

            download_gsheet_task = DownloadGSheetTask(
                file_id=dfile.id,
                authorized=req_data["authorized"],
                tenant_id=tenant_id,
                connection_id=req_data["connection_id"],
                exec_mode=exec_mode,
            )
            await task_processor.schedule(download_gsheet_task)
            LOGGER.info(f"Scheduled DownloadGSheetTask for file_id {dfile.id} (update connection)")

        return web.json_response(
            files_schemas.UpdateConnectionDataResultSchema().dump(dict(files=dfile_by_spreadsheet.values()))
        )


class InternalUpdateConnectionDataView(UpdateConnectionDataView):
    REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset(
        {
            RequiredResourceCommon.SKIP_AUTH,
            RequiredResourceCommon.SKIP_CSRF,
            RequiredResourceCommon.MASTER_KEY,
        }
    )

    _is_internal: ClassVar[bool] = True

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/views/files.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import ClassVar

from aiohttp import web

from dl_api_commons.aiohttp.aiohttp_wrappers import (
    RequiredResource,
    RequiredResourceCommon,
)
from dl_file_uploader_api_lib.schemas import misc as misc_schemas
from dl_file_uploader_api_lib.views.base import FileUploaderBaseView
from dl_file_uploader_lib.enums import RenameTenantStatus
from dl_file_uploader_lib.redis_model.models import RenameTenantStatusModel
from dl_file_uploader_task_interface.tasks import (
    CleanupTenantFilePreviewsTask,
    CleanupTenantTask,
    RenameTenantFilesTask,
)


LOGGER = logging.getLogger(__name__)


class CleanupTenantView(FileUploaderBaseView):
    REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset(
        {
            RequiredResourceCommon.SKIP_AUTH,
            RequiredResourceCommon.SKIP_CSRF,
            RequiredResourceCommon.MASTER_KEY,
        }
    )

    async def post(self) -> web.StreamResponse:
        req_data = await self._load_post_request_schema_data(misc_schemas.CleanupRequestSchema)
        tenant_id = req_data["tenant_id"]

        task_processor = self.dl_request.get_task_processor()
        await task_processor.schedule(
            CleanupTenantTask(
                tenant_id=tenant_id,
            )
        )
        LOGGER.info(f"Scheduled CleanupTenantTask for tenant id {tenant_id}")

        await task_processor.schedule(
            CleanupTenantFilePreviewsTask(
                tenant_id=tenant_id,
            )
        )
        LOGGER.info(f"Scheduled CleanupTenantFilePreviewsTask for tenant id {tenant_id}")

        return web.Response()


class RenameTenantFilesView(FileUploaderBaseView):
    # TODO: Delete this API handler (and task) after migration to organizations

    REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset(
        {
            RequiredResourceCommon.SKIP_AUTH,
            RequiredResourceCommon.SKIP_CSRF,
            RequiredResourceCommon.MASTER_KEY,
        }
    )

    async def post(self) -> web.StreamResponse:
        req_data = await self._load_post_request_schema_data(misc_schemas.RenameFilesRequestSchema)
        tenant_id = req_data["tenant_id"]

        rmm = self.dl_request.get_redis_model_manager()
        await RenameTenantStatusModel(manager=rmm, id=tenant_id, status=RenameTenantStatus.scheduled).save()

        task_processor = self.dl_request.get_task_processor()
        await task_processor.schedule(RenameTenantFilesTask(tenant_id=tenant_id))
        LOGGER.info(f"Scheduled RenameTenantFilesTask for tenant_id {tenant_id}")

        return web.Response()

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/views/misc.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    Any,
    ClassVar,
    Optional,
)

from aiohttp import web

from dl_api_commons.aiohttp.aiohttp_wrappers import (
    RequiredResource,
    RequiredResourceCommon,
)
from dl_constants.enums import (
    FileProcessingStatus,
    UserDataType,
)
from dl_core.db import get_type_transformer
from dl_core.db.elements import SchemaColumn
from dl_file_uploader_api_lib.schemas import sources as sources_schemas
from dl_file_uploader_api_lib.views.base import FileUploaderBaseView
from dl_file_uploader_lib import exc
from dl_file_uploader_lib.enums import (
    CSVDelimiter,
    CSVEncoding,
    FileType,
)
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models.models import (
    CSVFileSettings,
    DataFile,
    DataSource,
    DataSourcePreview,
    FileSettings,
    GSheetsUserSourceDataSourceProperties,
    GSheetsUserSourceProperties,
    UserSourceProperties,
)
from dl_file_uploader_task_interface.tasks import ParseFileTask

from dl_connector_bundle_chs3.file.core.constants import CONNECTION_TYPE_FILE


LOGGER = logging.getLogger(__name__)


RawSchemaType = list[SchemaColumn]


class SourceStatusView(FileUploaderBaseView):
    async def get(self) -> web.StreamResponse:
        req_data = sources_schemas.SourceInfoBaseRequestSchema().load(self.request.match_info)
        rmm = self.dl_request.get_redis_model_manager()
        dfile = await DataFile.get_authorized(manager=rmm, obj_id=req_data["file_id"])
        assert isinstance(dfile, DataFile)
        source = dfile.get_source_by_id(req_data["source_id"])

        status: FileProcessingStatus
        if dfile.status == FileProcessingStatus.failed:
            # TODO: temporary hack. Should be fixed with properly parsing errors saving.
            status = dfile.status
        else:
            status = source.status

        return web.json_response(
            sources_schemas.SourceStatusResultSchema().dump(
                {
                    "file_id": dfile.id,
                    "source_id": source.id,
                    "status": status,
                    "error": source.error,
                }
            ),
        )


def get_compatible_user_type(user_type: Optional[UserDataType]) -> tuple[UserDataType, ...]:
    compatible_types: dict[Optional[UserDataType], tuple[UserDataType, ...]] = {
        UserDataType.string: (UserDataType.string,),
        UserDataType.integer: (
            UserDataType.integer,
            UserDataType.float,
            UserDataType.string,
        ),
        UserDataType.float: (
            UserDataType.float,
            UserDataType.string,
        ),
        UserDataType.date: (
            UserDataType.date,
            UserDataType.genericdatetime,
            UserDataType.string,
        ),
        UserDataType.genericdatetime: (
            UserDataType.genericdatetime,
            UserDataType.string,
        ),
    }
    return compatible_types.get(user_type, tuple())


def cast_preview_data(
    preview_data: list[list[Optional[str]]],
    raw_schema: RawSchemaType,
) -> list[list[Optional[str]]]:
    tt = get_type_transformer(CONNECTION_TYPE_FILE)  # FIXME
    for row in preview_data:
        for i, cell in enumerate(row):
            cast_value = tt.cast_for_input(cell, raw_schema[i].user_type)
            row[i] = str(cast_value) if cast_value is not None else None
    return preview_data


def get_raw_schema_with_overrides(raw_schema: RawSchemaType, schema_overrides: RawSchemaType) -> RawSchemaType:
    orig_column_types_by_name: dict[str, UserDataType] = {sch.name: sch.user_type for sch in raw_schema}
    column_type_overrides = dict()
    for col in schema_overrides:
        col_name = col.name
        new_type = col.user_type
        orig_type = orig_column_types_by_name.get(col_name)
        if new_type not in get_compatible_user_type(orig_type):
            raise exc.InvalidFieldCast(f"Unsupported cast {orig_type} to {new_type} for field {col_name}")
        column_type_overrides[col_name] = new_type

    raw_schema_with_overrides: RawSchemaType = [
        col.clone(user_type=column_type_overrides[col.name]) if col.name in column_type_overrides else col
        for col in raw_schema
    ]
    return raw_schema_with_overrides


class SourceInfoView(FileUploaderBaseView):
    async def _make_source_resp(
        self,
        file_type: FileType,
        user_source_properties: Optional[UserSourceProperties],
        source: DataSource,
        column_type_overrides: RawSchemaType,
        rmm: RedisModelManager,
    ) -> dict[Any, Any]:
        preview_id = source.preview_id
        if preview_id is not None:
            preview = await DataSourcePreview.get(manager=rmm, obj_id=preview_id)
            assert isinstance(preview, DataSourcePreview)
        else:
            preview = DataSourcePreview(preview_data=[])

        raw_schema = get_raw_schema_with_overrides(source.raw_schema, column_type_overrides)

        extra: dict[str, Any] = {}
        if file_type == FileType.gsheets:
            assert isinstance(user_source_properties, GSheetsUserSourceProperties)
            extra["spreadsheet_id"] = user_source_properties.spreadsheet_id
            assert isinstance(source.user_source_dsrc_properties, GSheetsUserSourceDataSourceProperties)
            extra["sheet_id"] = source.user_source_dsrc_properties.sheet_id

        source_resp = dict(
            file_type=file_type,
            source_id=source.id,
            title=source.title,
            is_valid=True,  # TODO
            raw_schema=raw_schema,
            preview=cast_preview_data(preview.preview_data, raw_schema),
            error=source.error,
            **extra,
        )

        return source_resp

    def _make_data_settings(self, source: DataSource, file_settings: Optional[FileSettings]) -> dict[str, Any]:
        data_settings: dict[str, Any] = dict()
        if source.file_source_settings is not None:
            data_settings["first_line_is_header"] = source.file_source_settings.first_line_is_header
        if isinstance(file_settings, CSVFileSettings):
            data_settings["encoding"] = file_settings.encoding
            data_settings["delimiter"] = file_settings.delimiter
        return data_settings

    def _make_options(self, file_type: FileType, source: DataSource) -> Optional[dict[str, Any]]:
        if file_type == FileType.csv:
            return dict(
                data_settings={
                    "encoding": sorted((enc for enc in CSVEncoding), key=lambda t: t.name),
                    "delimiter": sorted((delim for delim in CSVDelimiter), key=lambda t: t.name),
                },
                columns=[
                    {"name": sch.name, "user_type": get_compatible_user_type(sch.user_type)}
                    for sch in source.raw_schema
                ],
            )
        return None

    async def _prepare_source_info(self, req_data: dict[str, Any]) -> web.StreamResponse:
        rmm = self.dl_request.get_redis_model_manager()
        dfile = await DataFile.get_authorized(manager=rmm, obj_id=req_data["file_id"])
        assert isinstance(dfile, DataFile)
        source = dfile.get_source_by_id(req_data["source_id"])

        # TODO: check if source.status == ready ???

        assert dfile.file_type is not None
        source_resp = await self._make_source_resp(
            file_type=dfile.file_type,
            user_source_properties=dfile.user_source_properties,
            source=source,
            column_type_overrides=[] if dfile.file_type == FileType.gsheets else req_data["column_types"],
            rmm=rmm,
        )

        assert dfile.file_type is not None
        return web.json_response(
            sources_schemas.SourceInfoResultSchema().dump(
                dict(
                    file_id=dfile.id,
                    source=source_resp,
                    data_settings=self._make_data_settings(source, dfile.file_settings),
                    options=self._make_options(dfile.file_type, source),
                )
            ),
        )

    async def post(self) -> web.StreamResponse:
        req_data = await self._load_post_request_schema_data(sources_schemas.SourceInfoRequestSchema)
        return await self._prepare_source_info(req_data)


class SourceApplySettingsView(FileUploaderBaseView):
    async def post(self) -> web.StreamResponse:
        req_data = await self._load_post_request_schema_data(sources_schemas.SourceApplySettingsRequestSchema)
        rmm = self.dl_request.get_redis_model_manager()
        dfile = await DataFile.get_authorized(manager=rmm, obj_id=req_data["file_id"])
        assert isinstance(dfile, DataFile)
        source = dfile.get_source_by_id(req_data["source_id"])
        LOGGER.info(f"Applying data settings for file {dfile.id}, source {source.id}")
        dfile.status = FileProcessingStatus.in_progress
        source.status = FileProcessingStatus.in_progress
        if req_data["title"] is not None:
            source.title = req_data["title"]
        LOGGER.info("Switching file and source status to `in_progress`")
        await dfile.save()

        LOGGER.info(f'Settings to apply: {req_data["data_settings"]}')

        task_processor = self.dl_request.get_task_processor()
        await task_processor.schedule(
            ParseFileTask(
                file_id=dfile.id,
                source_id=source.id,
                file_settings=sources_schemas.CSVSettingsSchema().dump(req_data["data_settings"]),
                source_settings={},
            )
        )
        LOGGER.info(f"Scheduled ParseFileTask for file {dfile.id}, source {source.id}")
        return web.Response()


class SourcePreviewView(FileUploaderBaseView):
    REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset(
        {
            RequiredResourceCommon.SKIP_CSRF,
            RequiredResourceCommon.MASTER_KEY,
        }
    )

    async def post(self) -> web.StreamResponse:
        req_data = await self._load_post_request_schema_data(sources_schemas.SourcePreviewRequestSchema)
        rmm = self.dl_request.get_redis_model_manager()

        preview = await DataSourcePreview.get(manager=rmm, obj_id=req_data["preview_id"])
        assert isinstance(preview, DataSourcePreview)

        return web.json_response(
            sources_schemas.SourcePreviewResultSchema().dump(
                dict(
                    preview=cast_preview_data(preview.preview_data, req_data["raw_schema"]),
                )
            ),
        )


class SourceInternalParamsView(FileUploaderBaseView):
    REQUIRED_RESOURCES: ClassVar[frozenset[RequiredResource]] = frozenset(
        {
            RequiredResourceCommon.SKIP_CSRF,
            RequiredResourceCommon.MASTER_KEY,
        }
    )

    async def post(self) -> web.StreamResponse:
        req_data = await self._load_post_request_schema_data(sources_schemas.SourceInternalParamsRequestSchema)
        rmm = self.dl_request.get_redis_model_manager()
        dfile = await DataFile.get_authorized(manager=rmm, obj_id=req_data["file_id"])
        assert isinstance(dfile, DataFile)
        source = dfile.get_source_by_id(req_data["source_id"])
        preview_id = source.preview_id

        column_type_overrides: Optional[RawSchemaType] = req_data["raw_schema"]
        LOGGER.info(f"Got overrides for raw_schema: {column_type_overrides=}")
        if column_type_overrides is not None:
            raw_schema = get_raw_schema_with_overrides(source.raw_schema, column_type_overrides)
        else:
            raw_schema = source.raw_schema
        LOGGER.info(f"Raw schema for file id {req_data['file_id']}, source id {req_data['source_id']}: {raw_schema}")

        return web.json_response(
            sources_schemas.SourceInternalParamsResultSchema().dump(dict(raw_schema=raw_schema, preview_id=preview_id)),
        )

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib/views/sources.py
datalens-tech/datalens-backend
import attr

from dl_configs.crypto_keys import (
    CryptoKeysConfig,
    get_single_key_crypto_keys_config,
)


@attr.s(kw_only=True)
class TestingUSConfig:
    container_name: str = attr.ib(default="bi_file_uploader_api_us-tests")
    master_token: str = attr.ib(default="AC1ofiek8coB")
    base_url: str = attr.ib()
    crypto_keys_config: CryptoKeysConfig = attr.ib(
        default=get_single_key_crypto_keys_config(
            key_id="0",
            key_value="BRa9gGXGm9qyy1efiIH7OkRHALjAohYtxOz-wSvlARY=",
        )
    )
    psycopg2_pg_dns: str = attr.ib()
    force_clear_db_on_launch: bool = attr.ib(default=True)

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/config.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import logging
import os
import sys
from typing import (
    TYPE_CHECKING,
    Any,
)

import aiohttp.pytest_plugin
import aiohttp.test_utils
import aiohttp.web
import attr
import pytest
import redis.asyncio

from dl_api_commons.aio.middlewares.auth_trust_middleware import auth_trust_middleware
from dl_api_commons.aio.typing import AIOHTTPMiddleware
from dl_api_commons.base_models import (
    NoAuthData,
    RequestContextInfo,
    TenantCommon,
)
from dl_api_commons.client.common import DLCommonAPIClient
from dl_configs.crypto_keys import (
    CryptoKeysConfig,
    get_dummy_crypto_keys_config,
)
from dl_configs.enums import RedisMode
from dl_configs.settings_submodels import (
    CorsSettings,
    CsrfSettings,
    GoogleAppSettings,
    RedisSettings,
    S3Settings,
)
from dl_constants.api_constants import DLHeadersCommon
from dl_core.loader import load_core_lib
from dl_core.services_registry.top_level import DummyServiceRegistry
from dl_core.united_storage_client import USAuthContextMaster
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_core_testing.environment import (
    common_pytest_configure,
    prepare_united_storage,
)
from dl_file_secure_reader_lib.app import create_app as create_reader_app
from dl_file_uploader_api_lib.app import FileUploaderApiAppFactory
from dl_file_uploader_api_lib.dl_request import FileUploaderDLRequest
from dl_file_uploader_api_lib.settings import FileUploaderAPISettings
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_task_interface.context import FileUploaderTaskContext
from dl_file_uploader_worker_lib.settings import (
    FileUploaderConnectorsSettings,
    FileUploaderWorkerSettings,
    SecureReader,
)
from dl_file_uploader_worker_lib.testing.task_processor_client import get_task_processor_client
from dl_task_processor.processor import TaskProcessor
from dl_task_processor.state import (
    BITaskStateImpl,
    TaskState,
)
from dl_testing.containers import get_test_container_hostport
from dl_testing.s3_utils import (
    create_s3_bucket,
    create_s3_client,
)
from dl_testing.utils import wait_for_initdb

from dl_connector_bundle_chs3.chs3_base.core.settings import FileS3ConnectorSettings

from .config import TestingUSConfig


if TYPE_CHECKING:
    from types_aiobotocore_s3 import S3Client as AsyncS3Client


LOGGER = logging.getLogger(__name__)


pytest_plugins = ("aiohttp.pytest_plugin",)

try:
    del aiohttp.pytest_plugin.loop
except AttributeError:
    pass


def pytest_configure(config: Any) -> None:  # noqa
    common_pytest_configure(tracing_service_name="tests_bi_file_uploader")


@pytest.fixture(autouse=True)
def loop(event_loop):
    """
    Preventing creation of new loop by `aiohttp.pytest_plugin` loop fixture in favor of pytest-asyncio one
    And set loop pytest-asyncio created loop as default for thread
    """
    asyncio.set_event_loop(event_loop)
    return event_loop


@pytest.fixture(scope="session")
def initdb_ready():
    return wait_for_initdb(initdb_port=51408)


@pytest.fixture(scope="session")
def us_config():
    return TestingUSConfig(
        base_url=f"http://{get_test_container_hostport('us').as_pair()}",
        psycopg2_pg_dns=(
            f'host={get_test_container_hostport("pg-us").host}'
            f' port={get_test_container_hostport("pg-us").port}'
            " user=us"
            " password=us"
            " dbname=us-db-ci_purgeable"
        ),
    )


@pytest.fixture(scope="session", autouse=True)
def loaded_libraries():
    load_core_lib()


@pytest.fixture(scope="session")
def redis_app_settings() -> RedisSettings:
    return RedisSettings(
        MODE=RedisMode.single_host,
        CLUSTER_NAME="",
        HOSTS=(get_test_container_hostport("redis").host,),
        PORT=get_test_container_hostport("redis").port,
        DB=9,
    )


@pytest.fixture(scope="session")
def redis_arq_settings(redis_app_settings):
    return attr.evolve(redis_app_settings, DB=11)


@pytest.fixture(scope="session")
def s3_settings() -> S3Settings:
    return S3Settings(
        ENDPOINT_URL=f'http://{get_test_container_hostport("s3-storage").as_pair()}',
        ACCESS_KEY_ID="accessKey1",
        SECRET_ACCESS_KEY="verySecretKey1",
    )


@pytest.fixture(scope="session")
def secure_reader():
    socket_name = "reader.sock"
    if sys.platform == "darwin":
        path = "/var/tmp"
    else:
        path = "/var"

    return SecureReader(
        SOCKET=os.path.join(path, socket_name),
    )


@pytest.fixture(scope="session")
def crypto_keys_config() -> CryptoKeysConfig:
    return get_dummy_crypto_keys_config()


@pytest.fixture(scope="function")
def app_settings(monkeypatch, redis_app_settings, redis_arq_settings, s3_settings, crypto_keys_config):
    monkeypatch.setenv("EXT_QUERY_EXECUTER_SECRET_KEY", "dummy")

    settings = FileUploaderAPISettings(
        REDIS_APP=redis_app_settings,
        REDIS_ARQ=redis_arq_settings,
        CORS=CorsSettings(
            ALLOWED_ORIGINS=("https://foo.bar", "https://fee.bar"),
            ALLOWED_HEADERS=("Content-Type", "X-Request-ID", "x-csrf-token", "x-unauthorized"),
            EXPOSE_HEADERS=("X-Request-ID",),
        ),
        CSRF=CsrfSettings(
            METHODS=("POST", "PUT", "DELETE"),
            HEADER_NAME="x-csrf-token",
            TIME_LIMIT=3600 * 12,
            SECRET="123321",
        ),
        S3=S3Settings(
            ENDPOINT_URL=s3_settings.ENDPOINT_URL,
            ACCESS_KEY_ID=s3_settings.ACCESS_KEY_ID,
            SECRET_ACCESS_KEY=s3_settings.SECRET_ACCESS_KEY,
        ),
        S3_TMP_BUCKET_NAME="bi-file-uploader-tmp",
        S3_PERSISTENT_BUCKET_NAME="bi-file-uploader",
        FILE_UPLOADER_MASTER_TOKEN="valid-master-token",
        CRYPTO_KEYS_CONFIG=crypto_keys_config,
        ALLOW_XLSX=True,
    )
    yield settings


class TestingFileUploaderApiAppFactory(FileUploaderApiAppFactory[FileUploaderAPISettings]):
    def get_auth_middlewares(self) -> list[AIOHTTPMiddleware]:
        return [
            auth_trust_middleware(
                fake_user_id="_the_tests_file_uploader_api_user_id_",
                fake_user_name="_the_tests_file_uploader_api_user_name_",
            )
        ]


@pytest.fixture(scope="function")
def bi_file_uploader_app(loop, aiohttp_client, app_settings):
    app_factory = TestingFileUploaderApiAppFactory(settings=app_settings)
    app = app_factory.create_app("tests")
    return loop.run_until_complete(aiohttp_client(app))


@pytest.fixture(scope="function")
def fu_client(bi_file_uploader_app) -> DLCommonAPIClient:
    return DLCommonAPIClient(
        base_url=f"http://{bi_file_uploader_app.host}:{bi_file_uploader_app.port}",
        tenant=TenantCommon(),
        auth_data=NoAuthData(),
    )


@pytest.fixture(scope="function")
async def s3_client(s3_settings) -> AsyncS3Client:
    async with create_s3_client(s3_settings) as client:
        yield client


@pytest.fixture(scope="function")
async def s3_tmp_bucket(initdb_ready, s3_client, app_settings) -> str:
    bucket_name = app_settings.S3_TMP_BUCKET_NAME
    await create_s3_bucket(s3_client, bucket_name, max_attempts=1)
    return bucket_name


@pytest.fixture(scope="function")
async def s3_persistent_bucket(initdb_ready, s3_client, app_settings) -> str:
    bucket_name = app_settings.S3_PERSISTENT_BUCKET_NAME
    await create_s3_bucket(s3_client, bucket_name, max_attempts=1)
    return bucket_name


@pytest.fixture(scope="function")
def redis_cli(redis_app_settings) -> redis.asyncio.Redis:
    return redis.asyncio.Redis(
        host=redis_app_settings.HOSTS[0],
        port=redis_app_settings.PORT,
        db=redis_app_settings.DB,
        password=redis_app_settings.PASSWORD,
    )


@pytest.fixture(scope="function")
def rci() -> RequestContextInfo:
    return RequestContextInfo(user_id="_the_tests_asyncapp_user_id_")


@pytest.fixture(scope="session")
def bi_context() -> RequestContextInfo:
    return RequestContextInfo.create(
        request_id=None,
        tenant=TenantCommon(),
        user_name=None,
        user_id="_the_tests_fixture_user_id_",
        x_dl_debug_mode=None,
        secret_headers=None,
        auth_data=None,
        plain_headers=None,
        endpoint_code=None,
        x_dl_context=None,
    )


@pytest.fixture(scope="function")
def redis_model_manager(redis_cli, rci, crypto_keys_config) -> RedisModelManager:
    return RedisModelManager(redis=redis_cli, rci=rci, crypto_keys_config=crypto_keys_config)


@pytest.fixture(scope="function")
def master_token_header(app_settings) -> dict[DLHeadersCommon, str]:
    return {DLHeadersCommon.FILE_UPLOADER_MASTER_TOKEN: app_settings.FILE_UPLOADER_MASTER_TOKEN}


@pytest.fixture(scope="session")
def connectors_settings(s3_settings):
    return FileUploaderConnectorsSettings(
        FILE=FileS3ConnectorSettings(
            SECURE=False,
            HOST=get_test_container_hostport("db-clickhouse", original_port=8123).host,
            PORT=get_test_container_hostport("db-clickhouse", original_port=8123).port,
            USERNAME="datalens",
            PASSWORD="qwerty",
            ACCESS_KEY_ID=s3_settings.ACCESS_KEY_ID,
            SECRET_ACCESS_KEY=s3_settings.SECRET_ACCESS_KEY,
            BUCKET="bi-file-uploader",
            S3_ENDPOINT="http://s3-storage:8000",
        ),
    )


@pytest.fixture(scope="function")
def file_uploader_worker_settings(
    redis_app_settings,
    redis_arq_settings,
    s3_settings,
    connectors_settings,
    us_config,
    crypto_keys_config,
    secure_reader,
):
    settings = FileUploaderWorkerSettings(
        REDIS_APP=redis_app_settings,
        REDIS_ARQ=redis_arq_settings,
        S3=S3Settings(
            ENDPOINT_URL=s3_settings.ENDPOINT_URL,
            ACCESS_KEY_ID=s3_settings.ACCESS_KEY_ID,
            SECRET_ACCESS_KEY=s3_settings.SECRET_ACCESS_KEY,
        ),
        S3_TMP_BUCKET_NAME="bi-file-uploader-tmp",
        S3_PERSISTENT_BUCKET_NAME="bi-file-uploader",
        SENTRY_DSN=None,
        US_BASE_URL=us_config.base_url,
        US_MASTER_TOKEN=us_config.master_token,
        CONNECTORS=connectors_settings,
        GSHEETS_APP=GoogleAppSettings(
            API_KEY="dummy",
            CLIENT_ID="dummy",
            CLIENT_SECRET="dummy",
        ),
        CRYPTO_KEYS_CONFIG=crypto_keys_config,
        SECURE_READER=secure_reader,
    )
    yield settings


@pytest.fixture(scope="function", params=["local"], ids=["local_tp"])
async def use_local_task_processor(request, monkeypatch, loop, file_uploader_worker_settings):
    task_state = TaskState(BITaskStateImpl())
    async with get_task_processor_client(
        client_type=request.param,
        loop=loop,
        task_state=task_state,
        file_uploader_worker_settings=file_uploader_worker_settings,
    ) as task_processor_client:

        def new_get_task_processor(self: FileUploaderDLRequest) -> TaskProcessor:
            task_processor_client._request_id = self.rci.request_id
            return task_processor_client

        def new_make_task_processor(self, request_id: str) -> TaskProcessor:
            task_processor_client._request_id = request_id
            return task_processor_client

        monkeypatch.setattr(FileUploaderDLRequest, "get_task_processor", new_get_task_processor)
        monkeypatch.setattr(FileUploaderTaskContext, "make_task_processor", new_make_task_processor)

        yield


@pytest.fixture(scope="session")
def prepare_us(us_config):
    prepare_united_storage(
        us_host=us_config.base_url,
        us_master_token=us_config.master_token,
        us_pg_dsn=us_config.psycopg2_pg_dns,
        force=us_config.force_clear_db_on_launch,
    )

    return us_config


@pytest.fixture(scope="function")
@pytest.mark.usefixtures("loop")
def default_async_usm_per_test(bi_context, prepare_us, us_config):
    rci = RequestContextInfo.create_empty()
    return AsyncUSManager(
        us_base_url=us_config.base_url,
        us_auth_context=USAuthContextMaster(us_config.master_token),
        crypto_keys_config=us_config.crypto_keys_config,
        bi_context=bi_context,
        services_registry=DummyServiceRegistry(rci=rci),
    )


@pytest.fixture(scope="function")
def reader_app(loop, secure_reader):
    current_app = create_reader_app()
    runner = aiohttp.web.AppRunner(current_app)
    loop.run_until_complete(runner.setup())
    site = aiohttp.web.UnixSite(runner, path=secure_reader.SOCKET)
    return loop.run_until_complete(site.start())

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/conftest.py
datalens-tech/datalens-backend
import asyncio
import logging
import os

import pytest

from dl_api_commons.client.common import Req
from dl_file_uploader_api_lib_tests.req_builder import ReqBuilder
from dl_file_uploader_lib.testing.data_gen import generate_sample_csv_data_str


LOGGER = logging.getLogger(__name__)


@pytest.fixture(scope="function", autouse=True)
async def use_local_task_processor_auto(use_local_task_processor):
    yield


@pytest.fixture(scope="session")
def csv_data():
    return f"""f1,f2,f3,Дата,Дата и время
qwe,123,45.9,2021-02-04,2021-02-04 12:00:00
asd,345,47.9,2021-02-05,2021-02-05 14:01:00
zxc,456,"49,9",2021-02-06,2021-02-06 11:59:00
zxc,456,,,2021-02-06 11:59:00
{'zxc'*35000},456,49.9,2021-02-06,
"""


@pytest.fixture(scope="session")
def excel_data():
    filename = "data.xlsx"

    dirname = os.path.dirname(os.path.abspath(__file__))
    path = os.path.join(dirname, filename)

    with open(path, "rb") as fd:
        return fd.read()


@pytest.fixture(scope="session")
def upload_file_req(csv_data) -> Req:
    return ReqBuilder.upload_csv(csv_data)


@pytest.fixture(scope="session")
def upload_excel_req(excel_data) -> Req:
    return ReqBuilder.upload_xlsx(excel_data)


@pytest.fixture(scope="session")
def upload_file_req_10mb() -> Req:
    csv_data = generate_sample_csv_data_str(row_count=10000, str_cols_count=30)
    LOGGER.info(f'Generated csv data sample of {len(csv_data.encode("utf-8"))} bytes')
    return ReqBuilder.upload_csv(csv_data)


@pytest.fixture(scope="session")
def upload_file_req_12mb() -> Req:
    csv_data = generate_sample_csv_data_str(row_count=12000, str_cols_count=30)
    LOGGER.info(f'Generated csv data sample of {len(csv_data.encode("utf-8"))} bytes')
    return ReqBuilder.upload_csv(csv_data)


@pytest.fixture(scope="function")
async def uploaded_file_id(s3_tmp_bucket, fu_client, upload_file_req) -> str:
    resp = await fu_client.make_request(upload_file_req)
    assert resp.status == 201
    await asyncio.sleep(3)
    return resp.json["file_id"]


@pytest.fixture(scope="function")
async def uploaded_excel_id(
    s3_tmp_bucket,
    fu_client,
    upload_excel_req,
    reader_app,
) -> str:
    resp = await fu_client.make_request(upload_excel_req)
    assert resp.status == 201
    await asyncio.sleep(3)
    return resp.json["file_id"]

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/db/conftest.py
datalens-tech/datalens-backend
import pytest

from dl_api_commons.client.common import Req
from dl_constants.api_constants import (
    DLHeaders,
    DLHeadersCommon,
)
from dl_file_uploader_api_lib_tests.req_builder import ReqBuilder


class TempHeaders(DLHeaders):
    X_UNAUTHORIZED = "x-unauthorized"


@pytest.mark.asyncio
@pytest.mark.parametrize("view", ["ping", "metrics"])
async def test_basic(fu_client, view):
    resp = await fu_client.make_request(Req("get", f"/api/v2/{view}"))
    assert resp.status == 200


@pytest.mark.asyncio
async def test_cors(fu_client):
    cors_headers = (
        "Access-Control-Expose-Headers",
        "Access-Control-Allow-Origin",
        "Access-Control-Allow-Credentials",
        "Access-Control-Allow-Methods",
    )

    resp = await fu_client.make_request(
        Req(
            method="options",
            url="/api/v2/files",
            extra_headers={
                TempHeaders.X_UNAUTHORIZED: "1",
                DLHeadersCommon.ACCESS_CTRL_REQ_HEADERS: "x-csrf-token,x-request-id",
                DLHeadersCommon.ACCESS_CTRL_REQ_METH: "POST",
                DLHeadersCommon.ORIGIN: "https://foo.bar",
                DLHeadersCommon.REQUEST_ID: "qwerty",
            },
        ),
    )
    assert resp.status == 200
    for header in cors_headers:
        assert header in resp.headers

    resp = await fu_client.make_request(
        Req(
            method="options",
            url="/api/v2/files",
            extra_headers={
                TempHeaders.X_UNAUTHORIZED: "1",
                DLHeadersCommon.ACCESS_CTRL_REQ_HEADERS: "x-csrf-token,x-request-id",
                DLHeadersCommon.ACCESS_CTRL_REQ_METH: "POST",
                DLHeadersCommon.ORIGIN: "https://some.host",  # wrong host
            },
        ),
    )
    assert resp.status == 200
    for header in cors_headers:
        assert header not in resp.headers


@pytest.mark.asyncio
async def test_request_schema_validation_error(fu_client):
    resp = await fu_client.make_request(
        Req(
            method="post",
            url="/api/v2/links",
            data_json={
                "url": "asdf",
            },
            require_ok=False,
        )
    )

    assert resp.status == 400
    assert "Missing data for required field." in resp.json["message"]

    resp = await fu_client.make_request(ReqBuilder.upload_gsheet(url="asdf", authorized=True, require_ok=False))
    assert resp.status == 400
    assert "Either refresh_token or connection_id must be provided when authorized is true" in resp.json["message"]

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/db/test_app.py
datalens-tech/datalens-backend
import asyncio
import http
import json
import uuid

import attr
import pytest

from dl_api_commons.base_models import RequestContextInfo
from dl_configs.crypto_keys import get_dummy_crypto_keys_config
from dl_constants.enums import FileProcessingStatus
from dl_file_uploader_api_lib_tests.req_builder import ReqBuilder
from dl_file_uploader_lib.data_sink.raw_bytes import S3RawFileAsyncDataSink
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models import DataFile


@pytest.mark.asyncio
async def test_file_upload_cors(fu_client, s3_tmp_bucket, upload_file_req):
    resp = await fu_client.make_request(upload_file_req)
    assert resp.status == 201
    for cors_header in (
        "Access-Control-Expose-Headers",
        "Access-Control-Allow-Origin",
        "Access-Control-Allow-Credentials",
    ):
        assert cors_header in resp.headers


@pytest.mark.asyncio
async def test_file_upload(
    fu_client,
    s3_client,
    s3_tmp_bucket,
    redis_cli,
    upload_file_req,
    csv_data,
):
    resp = await fu_client.make_request(upload_file_req)
    assert resp.status == 201
    file_id = resp.json["file_id"]
    assert file_id
    title = resp.json["title"]
    assert title

    redis_data = await redis_cli.get(DataFile._generate_key_by_id(file_id))
    decoded_data = json.loads(redis_data.decode())
    assert decoded_data["id"] == file_id
    # assert decoded_data['size'] == len(csv_data.encode('utf-8'))

    s3_obj = await s3_client.get_object(
        Bucket=s3_tmp_bucket,
        Key=file_id,
    )
    assert s3_obj
    s3_data = await s3_obj["Body"].read()
    assert s3_data.decode() == csv_data


@pytest.mark.asyncio
async def test_excel_upload(
    fu_client,
    s3_client,
    s3_tmp_bucket,
    redis_cli,
    upload_excel_req,
    excel_data,
    reader_app,
):
    resp = await fu_client.make_request(upload_excel_req)
    assert resp.status == 201
    file_id = resp.json["file_id"]
    assert file_id
    title = resp.json["title"]
    assert title

    redis_data = await redis_cli.get(DataFile._generate_key_by_id(file_id))
    decoded_data = json.loads(redis_data.decode())
    assert decoded_data["id"] == file_id

    s3_obj = await s3_client.get_object(
        Bucket=s3_tmp_bucket,
        Key=file_id,
    )
    assert s3_obj
    s3_data = await s3_obj["Body"].read()
    assert s3_data == excel_data


@pytest.mark.asyncio
async def test_too_large_file_upload(
    monkeypatch,
    fu_client,
    s3_client,
    s3_tmp_bucket,
    redis_cli,
    upload_file_req_12mb,
    csv_data,
):
    monkeypatch.setattr(S3RawFileAsyncDataSink, "max_file_size_bytes", 8 * 1024**2)

    request = attr.evolve(upload_file_req_12mb, require_ok=False)
    resp = await fu_client.make_request(request)
    assert resp.status == http.HTTPStatus.REQUEST_ENTITY_TOO_LARGE, resp.content


@pytest.mark.skip("TODO")
@pytest.mark.asyncio
async def test_10MB_file_upload(
    fu_client,
    s3_client,
    s3_tmp_bucket,
    redis_model_manager,
    upload_file_req_10mb,
):
    resp = await fu_client.make_request(upload_file_req_10mb)
    assert resp.status == 201
    file_id = resp.json["file_id"]
    assert file_id

    await asyncio.sleep(3)

    df = await DataFile.get(manager=redis_model_manager, obj_id=file_id)
    assert df.sources


@pytest.mark.asyncio
async def test_file_status(fu_client, uploaded_file_id: str):
    file_id = uploaded_file_id
    resp = await fu_client.make_request(ReqBuilder.file_status(file_id))
    assert resp.status == 200
    assert resp.json["file_id"] == file_id
    assert resp.json["status"] in ("ready", "in_progress")
    assert resp.json["error"] is None


@pytest.mark.asyncio
async def test_file_sources(fu_client, uploaded_file_id):
    resp = await fu_client.make_request(ReqBuilder.file_sources(uploaded_file_id))
    assert resp.status == 200
    assert resp.json["file_id"] == uploaded_file_id
    assert isinstance(resp.json["sources"], list)
    source = resp.json["sources"][0]
    assert source["source_id"]
    assert source["title"]
    assert isinstance(source["is_applicable"], bool)


@pytest.mark.asyncio
async def test_file_404(fu_client):
    resp = await fu_client.make_request(ReqBuilder.file_sources(str(uuid.uuid4()), require_ok=False))
    assert resp.status == 404


@pytest.mark.asyncio
async def test_file_403(fu_client, redis_cli):
    rmm = RedisModelManager(
        redis=redis_cli,
        rci=RequestContextInfo(user_id="some_user_id_1"),
        crypto_keys_config=get_dummy_crypto_keys_config(),
    )
    df = DataFile(
        manager=rmm,
        filename="test_file.csv",
        status=FileProcessingStatus.in_progress,
    )
    await df.save()

    resp = await fu_client.make_request(ReqBuilder.file_sources(df.id, require_ok=False))
    assert resp.status == 403

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/db/test_files_api.py
datalens-tech/datalens-backend
import asyncio
import json
import time
import uuid

import pytest

from dl_constants.api_constants import DLHeadersCommon
from dl_file_uploader_api_lib_tests.req_builder import ReqBuilder
from dl_testing.s3_utils import get_lc_rules_number


@pytest.fixture(scope="function")
async def files_with_tenant_prefixes(s3_persistent_bucket, s3_client) -> list[str]:
    csv_data = """f1,f2,f3,Дата,Дата и время
qwe,123,45.9,2021-02-04,2021-02-04 12:00:00
asd,345,47.9,2021-02-05,2021-02-05 14:01:00""".encode(
        "utf-8"
    )

    tenant_ids = [f"tenant_{uuid.uuid4()}" for _ in range(3)]
    for tenant_id in tenant_ids:
        for idx in range(42):
            filename = f"file_{idx}"
            await s3_client.put_object(
                ACL="private",
                Bucket=s3_persistent_bucket,
                Key=f"{tenant_id}_{filename}",
                Body=csv_data,
            )

    return tenant_ids


@pytest.mark.asyncio
async def test_cleanup_tenant(
    master_token_header,
    fu_client,
    s3_client,
    s3_persistent_bucket: str,
    files_with_tenant_prefixes: list[str],
):
    await s3_client.delete_bucket_lifecycle(Bucket=s3_persistent_bucket)
    tenant_ids = files_with_tenant_prefixes

    n_lc_rules = await get_lc_rules_number(s3_client, s3_persistent_bucket)
    assert n_lc_rules == 0

    for tenant_id in tenant_ids:
        resp = await fu_client.make_request(ReqBuilder.cleanup(tenant_id, master_token_header))
        await asyncio.sleep(3)
        assert resp.status == 200, resp.json

        new_n_lc_rules = await get_lc_rules_number(s3_client, s3_persistent_bucket)
        assert new_n_lc_rules == n_lc_rules + 1
        n_lc_rules = new_n_lc_rules

    assert n_lc_rules == len(tenant_ids)


@pytest.mark.asyncio
async def test_cleanup_tenant_no_files(
    s3_client,
    master_token_header,
    fu_client,
    s3_persistent_bucket: str,
):
    await s3_client.delete_bucket_lifecycle(Bucket=s3_persistent_bucket)

    resp = await fu_client.make_request(ReqBuilder.cleanup("there are no files in this tenant", master_token_header))
    assert resp.status == 200, resp.json
    await asyncio.sleep(3)

    assert await get_lc_rules_number(s3_client, s3_persistent_bucket) == 1  # note: the LC rule is still created


@pytest.mark.asyncio
async def test_cleanup_tenant_invalid_master_token(fu_client):
    resp = await fu_client.make_request(
        ReqBuilder.cleanup(
            tenant_id="does not matter",
            master_token_header={DLHeadersCommon.FILE_UPLOADER_MASTER_TOKEN: "invalid-master-token"},
            require_ok=False,
        )
    )
    assert resp.status == 403


@pytest.mark.asyncio
async def test_cleanup_tenant_no_master_token(fu_client):
    resp = await fu_client.make_request(
        ReqBuilder.cleanup(
            tenant_id="does not matter",
            master_token_header=None,
            require_ok=False,
        )
    )
    assert resp.status == 403


@pytest.mark.asyncio
async def test_rename_tenant_files(master_token_header, fu_client, redis_cli):
    tenant_id = str(uuid.uuid4())
    start_time = time.time()
    resp = await fu_client.make_request(ReqBuilder.rename_tenant_files(master_token_header, tenant_id))
    assert resp.status == 200, resp.json

    await asyncio.sleep(3)
    redis_raw_data = await redis_cli.get(f"rename_tenant_status/{tenant_id}")
    redis_data = json.loads(redis_raw_data)
    assert redis_data["id"] == tenant_id
    assert redis_data["status"] == "success"
    assert start_time < redis_data["mtime"] < time.time()

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/db/test_misc_api.py
datalens-tech/datalens-backend
import pytest

from dl_constants.api_constants import DLHeadersCommon
from dl_file_uploader_api_lib_tests.req_builder import ReqBuilder
from dl_file_uploader_lib.redis_model.models import DataFile


@pytest.mark.asyncio
async def test_file_status(fu_client, redis_model_manager, uploaded_file_id: str):
    file_id = uploaded_file_id
    df = await DataFile.get(manager=redis_model_manager, obj_id=file_id)
    source_id = df.sources[0].id

    resp = await fu_client.make_request(ReqBuilder.source_status(file_id, source_id))
    assert resp.status == 200
    assert resp.json["file_id"] == file_id
    assert resp.json["source_id"] == source_id
    assert resp.json["status"] in ("ready", "in_progress")
    assert resp.json["error"] is None


@pytest.mark.asyncio
async def test_source_info(fu_client, uploaded_file_id: str):
    file_id = uploaded_file_id

    sources_resp = await fu_client.make_request(ReqBuilder.file_sources(uploaded_file_id))
    source_id = sources_resp.json["sources"][0]["source_id"]

    resp = await fu_client.make_request(ReqBuilder.source_info(file_id, source_id))
    assert resp.status == 200
    assert resp.json["file_id"] == file_id
    assert isinstance(resp.json["source"], dict)
    source = resp.json["source"]
    assert source["source_id"] == source_id
    assert source["title"]
    assert source["is_valid"]
    assert source["raw_schema"]
    assert source["preview"]
    assert "error" in source
    assert resp.json["options"]
    assert resp.json["data_settings"]
    assert all(setting in resp.json["data_settings"] for setting in ("encoding", "delimiter", "first_line_is_header"))

    resp1 = await fu_client.make_request(
        ReqBuilder.source_info(
            file_id=file_id,
            source_id=source_id,
            data_json={
                "column_types": [
                    {"name": "f2", "user_type": "float"},
                    {"name": "data", "user_type": "genericdatetime"},
                    {"name": "data_i_vremya", "user_type": "string"},
                ],
            },
        )
    )
    assert resp1.status == 200


@pytest.mark.asyncio
async def test_apply_settings(fu_client, redis_model_manager, uploaded_file_id: str):
    file_id = uploaded_file_id
    df = await DataFile.get(manager=redis_model_manager, obj_id=file_id)
    source_id = df.sources[0].id

    resp_info = await fu_client.make_request(
        ReqBuilder.source_info(
            file_id=file_id,
            source_id=source_id,
        )
    )
    previous_title = resp_info.json["source"]["title"]

    resp = await fu_client.make_request(
        ReqBuilder.apply_source_settings(
            file_id=file_id,
            source_id=source_id,
            data_json={
                "data_settings": {
                    "encoding": "utf8",
                    "delimiter": "comma",
                    "first_line_is_header": False,
                },
                "title": "new_title.csv",
            },
        )
    )
    assert resp.status == 200

    resp_info = await fu_client.make_request(
        ReqBuilder.source_info(
            file_id=file_id,
            source_id=source_id,
        )
    )
    new_title = resp_info.json["source"]["title"]
    assert new_title != previous_title
    assert new_title == "new_title.csv"

    resp = await fu_client.make_request(
        ReqBuilder.apply_source_settings(
            file_id=file_id,
            source_id=source_id,
            data_json={
                "data_settings": {
                    "encoding": "utf8",
                    "delimiter": "comma",
                    "first_line_is_header": False,
                },
            },
        )
    )
    assert resp.status == 200

    resp_info = await fu_client.make_request(
        ReqBuilder.source_info(
            file_id=file_id,
            source_id=source_id,
        )
    )
    newest_title = resp_info.json["source"]["title"]
    assert newest_title == new_title


@pytest.mark.asyncio
async def test_source_preview(fu_client, master_token_header, uploaded_file_id: str, redis_model_manager):
    file_id = uploaded_file_id

    sources_resp = await fu_client.make_request(ReqBuilder.file_sources(file_id))
    source_id = sources_resp.json["sources"][0]["source_id"]
    df = await DataFile.get(manager=redis_model_manager, obj_id=file_id)
    source = df.get_source_by_id(source_id)
    preview_id = source.preview_id

    resp = await fu_client.make_request(
        ReqBuilder.preview(
            file_id,
            source_id,
            master_token_header,
            data_json={
                "preview_id": preview_id,
                "raw_schema": [
                    {"user_type": "string", "name": "f1"},
                    {"user_type": "integer", "name": "f2"},
                    {"user_type": "string", "name": "f3"},
                    {"user_type": "date", "name": "data"},
                    {"user_type": "genericdatetime", "name": "data_i_vremya"},
                ],
            },
        )
    )
    assert resp.status == 200
    assert resp.json["preview"] == [
        ["qwe", "123", "45.9", "2021-02-04", "2021-02-04T12:00:00"],
        ["asd", "345", "47.9", "2021-02-05", "2021-02-05T14:01:00"],
        ["zxc", "456", "49,9", "2021-02-06", "2021-02-06T11:59:00"],
        ["zxc", "456", None, None, "2021-02-06T11:59:00"],
        ["zxczxczxcz...czxczxczxc", "456", "49.9", "2021-02-06", None],
    ]


@pytest.mark.asyncio
async def test_excel_source(
    master_token_header,
    fu_client,
    uploaded_excel_id: str,
    redis_model_manager,
    reader_app,
):
    file_id = uploaded_excel_id

    sources_resp = await fu_client.make_request(ReqBuilder.file_sources(file_id))
    source_id = sources_resp.json["sources"][0]["source_id"]
    df = await DataFile.get(manager=redis_model_manager, obj_id=file_id)
    source = df.get_source_by_id(source_id)
    preview_id = source.preview_id

    resp = await fu_client.make_request(
        ReqBuilder.preview(
            file_id,
            source_id,
            master_token_header,
            data_json={
                "preview_id": preview_id,
                "raw_schema": [
                    {"user_type": "integer", "name": "row"},
                    {"user_type": "string", "name": "order_id"},
                    {"user_type": "genericdatetime", "name": "order_date"},
                    {"user_type": "genericdatetime", "name": "ship_date"},
                    {"user_type": "string", "name": "ship_mode"},
                    {"user_type": "string", "name": "customer_id"},
                    {"user_type": "string", "name": "customer_name"},
                    {"user_type": "string", "name": "segment"},
                    {"user_type": "string", "name": "city"},
                    {"user_type": "string", "name": "state"},
                    {"user_type": "string", "name": "global_area"},
                    {"user_type": "string", "name": "postal_code"},
                    {"user_type": "string", "name": "market"},
                    {"user_type": "string", "name": "region"},
                    {"user_type": "string", "name": "product_id"},
                    {"user_type": "string", "name": "category"},
                    {"user_type": "string", "name": "sub-category"},
                    {"user_type": "string", "name": "product_name"},
                    {"user_type": "float", "name": "sales"},
                    {"user_type": "integer", "name": "quantity"},
                    {"user_type": "float", "name": "discount"},
                    {"user_type": "float", "name": "profit"},
                    {"user_type": "float", "name": "shipping_cost"},
                    {"user_type": "string", "name": "order_priority"},
                    {"user_type": "float", "name": "ammount"},
                ],
            },
        )
    )
    assert resp.status == 200
    assert resp.json["preview"][0] == [
        "1",
        "MX-2014-143658",
        "2014-10-02T00:00:00",
        "2014-10-06T00:00:00",
        "Standard Class",
        "SC-20575",
        "Sonia Cooley",
        "Consumer",
        "Mexico City",
        "Distrito Federal",
        "Mexico",
        None,
        "LATAM",
        "North",
        "OFF-LA-10002782",
        "Office Supplies",
        "Labels",
        "Hon File Folder Labels, Adjustable",
        "13.08",
        "3",
        "0.0",
        "4.56",
        "1.033",
        "Medium",
        "39.24",
    ]

    resp = await fu_client.make_request(ReqBuilder.source_info(file_id, source_id))
    assert resp.status == 200
    assert resp.json["file_id"] == file_id
    assert resp.json["options"] is None
    assert isinstance(resp.json["source"], dict)
    source = resp.json["source"]
    assert source["source_id"] == source_id
    assert source["title"]
    assert source["is_valid"]
    assert source["raw_schema"]
    assert source["preview"]
    assert "error" in source
    assert resp.json["data_settings"]
    assert "first_line_is_header" in resp.json["data_settings"]


@pytest.mark.asyncio
async def test_source_raw_schema(fu_client, master_token_header, uploaded_file_id: str):
    file_id = uploaded_file_id

    sources_resp = await fu_client.make_request(ReqBuilder.file_sources(file_id))
    source_id = sources_resp.json["sources"][0]["source_id"]

    resp = await fu_client.make_request(ReqBuilder.internal_params(file_id, source_id, master_token_header))
    assert resp.status == 200
    assert resp.json["raw_schema"][2]["user_type"] == "float"

    resp = await fu_client.make_request(
        ReqBuilder.internal_params(
            file_id,
            source_id,
            master_token_header,
            data_json={
                "raw_schema": [
                    {"user_type": "string", "name": "f3"},
                ],
            },
        )
    )
    assert resp.status == 200
    assert resp.json["raw_schema"][2]["user_type"] == "string"


@pytest.mark.asyncio
async def test_invalid_master_token(fu_client):
    resp = await fu_client.make_request(
        ReqBuilder.internal_params(
            "-",
            "-",
            master_token_header={DLHeadersCommon.FILE_UPLOADER_MASTER_TOKEN: "invalid-master-token"},
            require_ok=False,
        )
    )
    assert resp.status == 403


@pytest.mark.asyncio
async def test_no_master_token(fu_client):
    resp = await fu_client.make_request(
        ReqBuilder.internal_params(
            "-",
            "-",
            master_token_header=None,
            require_ok=False,
        )
    )
    assert resp.status == 403

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/db/test_sources_api.py
datalens-tech/datalens-backend
import os

import pytest

from dl_configs.settings_submodels import (
    GoogleAppSettings,
    S3Settings,
)
from dl_file_uploader_worker_lib.settings import FileUploaderWorkerSettings
from dl_testing.env_params.generic import GenericEnvParamGetter


@pytest.fixture(scope="function", autouse=True)
async def use_local_task_processor_auto(use_local_task_processor):
    yield


@pytest.fixture(scope="session")
def env_param_getter() -> GenericEnvParamGetter:
    filepath = os.path.join(os.path.dirname(__file__), "params.yml")
    return GenericEnvParamGetter.from_yaml_file(filepath)


# overriding top-level fixture to set google app settings that are only needed in ext tests
@pytest.fixture(scope="function")
def file_uploader_worker_settings(
    env_param_getter,
    redis_app_settings,
    redis_arq_settings,
    s3_settings,
    connectors_settings,
    us_config,
    crypto_keys_config,
    secure_reader,
):
    settings = FileUploaderWorkerSettings(
        REDIS_APP=redis_app_settings,
        REDIS_ARQ=redis_arq_settings,
        S3=S3Settings(
            ENDPOINT_URL=s3_settings.ENDPOINT_URL,
            ACCESS_KEY_ID=s3_settings.ACCESS_KEY_ID,
            SECRET_ACCESS_KEY=s3_settings.SECRET_ACCESS_KEY,
        ),
        S3_TMP_BUCKET_NAME="bi-file-uploader-tmp",
        S3_PERSISTENT_BUCKET_NAME="bi-file-uploader",
        SENTRY_DSN=None,
        US_BASE_URL=us_config.base_url,
        US_MASTER_TOKEN=us_config.master_token,
        CONNECTORS=connectors_settings,
        GSHEETS_APP=GoogleAppSettings(
            API_KEY=env_param_getter.get_str_value("GOOGLE_API_KEY"),
            CLIENT_ID="dummy",  # TODO test auth properly
            CLIENT_SECRET="dummy",
        ),
        SECURE_READER=secure_reader,
        CRYPTO_KEYS_CONFIG=crypto_keys_config,
    )
    yield settings

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/ext/conftest.py
datalens-tech/datalens-backend
import asyncio
import json
import logging
import os

import pytest

from dl_constants.enums import FileProcessingStatus
from dl_core.aio.web_app_services.gsheets import GSheetsSettings
from dl_file_uploader_api_lib_tests.req_builder import ReqBuilder
from dl_file_uploader_lib.gsheets_client import GSheetsClient
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    GSheetsUserSourceProperties,
)
from dl_utils.aio import ContextVarExecutor


LOGGER = logging.getLogger(__name__)


@pytest.mark.asyncio
async def test_gsheets(
    fu_client,
    s3_client,
    s3_tmp_bucket,
    redis_model_manager,
):
    spreadsheet_id = "1rnUFa7AiSKD5O80IKCvMy2cSZvLU1kRw9dxbtZbDMWc"
    resp = await fu_client.make_request(ReqBuilder.upload_gsheet(spreadsheet_id))
    assert resp.status == 201
    file_id = resp.json["file_id"]
    assert file_id

    await asyncio.sleep(5)

    rmm = redis_model_manager
    df = await DataFile.get(manager=rmm, obj_id=file_id)
    assert isinstance(df, DataFile)
    assert isinstance(df.user_source_properties, GSheetsUserSourceProperties)
    assert df.user_source_properties.spreadsheet_id == spreadsheet_id
    assert df.filename == "fu"
    assert df.status == FileProcessingStatus.ready
    assert df.sources[0].status == FileProcessingStatus.failed
    assert all(source.status == FileProcessingStatus.ready for source in df.sources[1:])

    resp = await fu_client.make_request(ReqBuilder.file_sources(file_id))
    assert resp.status == 200
    assert resp.json["sources"][0]["error"]["code"] == "ERR.FILE.NO_DATA"

    resp = await fu_client.make_request(ReqBuilder.source_info(file_id, df.sources[0].id))
    assert resp.status == 200
    assert resp.json["source"]["error"]["code"] == "ERR.FILE.NO_DATA"

    resp = await fu_client.make_request(ReqBuilder.source_info(file_id, df.sources[3].id))  # "elaborate" sheet
    assert resp.status == 200
    assert resp.json["options"] is None
    assert resp.json["data_settings"] == {"first_line_is_header": True}
    source = resp.json["source"]
    assert source["source_id"] == df.sources[3].id
    assert source["title"]
    assert source["is_valid"]
    assert source["raw_schema"]
    assert source["preview"]
    assert "error" in source
    assert source["spreadsheet_id"] == spreadsheet_id
    assert source["sheet_id"] == df.sources[3].user_source_dsrc_properties.sheet_id
    assert source["raw_schema"] == [
        {"name": "A", "user_type": "integer", "title": "int field"},
        {"name": "B", "user_type": "date", "title": "date field"},
        {"name": "C", "user_type": "float", "title": "float field"},
        {"name": "D", "user_type": "string", "title": "string field"},
        {"name": "E", "user_type": "boolean", "title": "bool field"},
        {"name": "F", "user_type": "genericdatetime", "title": "datetime field"},
        {"name": "G", "user_type": "string", "title": "time field"},
        {"name": "H", "user_type": "string", "title": "time period field"},
    ]
    assert source["preview"] == [
        ["1", "2001-07-31", "12.5", "asdf", "False", "2022-08-12T12:30:32", "12:00:12", "13:00:12"],
        ["9", "2022-08-11", None, "qwer", "True", "2022-08-13T12:30:32", "13:00:12", "14:00:12"],
        ["2", "1995-04-12", "12.4", "zxcv", "False", "2022-08-14T12:30:32", "14:00:12", "15:00:12"],
        ["8", "2008-02-29", None, "12", "True", "2022-08-15T12:30:32", "15:00:12", "16:00:12"],
        ["3", "2012-03-12", "12.3", "dfgh", "False", "2022-08-16T12:30:32", "16:00:12", "17:00:12"],
        [None, None, None, None, None, None, None, "18:00:12"],
    ]


@pytest.mark.asyncio
async def test_big_gsheet(
    use_local_task_processor,
    fu_client,
    s3_client,
    s3_tmp_bucket,
    redis_model_manager,
):
    spreadsheet_id = "1sPIpWUZa7wgUnUDa-MFPidJTSjQ40jS2OmD0g6V_wDw"
    resp = await fu_client.make_request(ReqBuilder.upload_gsheet(spreadsheet_id))
    assert resp.status == 201
    file_id = resp.json["file_id"]
    assert file_id

    await asyncio.sleep(5)

    rmm = redis_model_manager
    df = await DataFile.get(manager=rmm, obj_id=file_id)
    assert isinstance(df, DataFile)
    assert isinstance(df.user_source_properties, GSheetsUserSourceProperties)
    assert df.user_source_properties.spreadsheet_id == spreadsheet_id
    assert df.filename == "MoscowShops v8 (truncated)"
    assert df.status == FileProcessingStatus.ready
    assert len(df.sources) == 5
    assert all(source.status == FileProcessingStatus.ready for source in df.sources)

    resp = await fu_client.make_request(ReqBuilder.source_info(file_id, df.sources[0].id))
    assert resp.status == 200
    assert resp.json["options"] is None
    assert resp.json["data_settings"] == {"first_line_is_header": True}
    source = resp.json["source"]
    assert source["source_id"] == df.sources[0].id
    assert source["title"]
    assert source["is_valid"]
    assert source["raw_schema"]
    assert source["preview"]
    assert "error" in source
    assert source["spreadsheet_id"] == spreadsheet_id
    assert source["sheet_id"] == df.sources[0].user_source_dsrc_properties.sheet_id
    assert len(source["raw_schema"]) == 26


@pytest.mark.asyncio
async def test_gsheets_invalid_link(fu_client, s3_client, s3_tmp_bucket, redis_model_manager):
    resp = await fu_client.make_request(
        ReqBuilder.upload_gsheet(
            url="https://sheets.googleapis.com/v4/spreadsheets/1vFysKzPgy7Zw09K6wShVQPe6VXIoAzpGH3g_4X8dgyw",
            require_ok=False,
        )
    )
    assert resp.status == 400
    assert resp.json["code"] == "ERR.FILE.INVALID_LINK"


@pytest.mark.asyncio
async def test_gsheets_no_access_to_spreadsheet(
    fu_client,
    s3_client,
    s3_tmp_bucket,
    redis_model_manager,
):
    resp = await fu_client.make_request(ReqBuilder.upload_gsheet("1vFysKzPgy7Zw09K6wShVQPe6VXIoAzpGH3g_4X8dgyw"))
    assert resp.status == 201
    file_id = resp.json["file_id"]

    await asyncio.sleep(2)

    resp = await fu_client.make_request(ReqBuilder.file_status(file_id))
    assert resp.status == 200
    assert resp.json["file_id"] == file_id
    assert resp.json["status"] == "failed"
    assert resp.json["error"]["code"] == "ERR.FILE.PERMISSION_DENIED"


@pytest.mark.asyncio
async def test_gsheets_not_found(
    fu_client,
    s3_client,
    s3_tmp_bucket,
    redis_model_manager,
):
    resp = await fu_client.make_request(ReqBuilder.upload_gsheet("ffffffffffffffffffffffffffffffffffffffffffff"))
    assert resp.status == 201
    file_id = resp.json["file_id"]

    await asyncio.sleep(2)

    resp = await fu_client.make_request(ReqBuilder.file_status(file_id))
    assert resp.status == 200
    assert resp.json["file_id"] == file_id
    assert resp.json["status"] == "failed"
    assert resp.json["error"]["code"] == "ERR.FILE.NOT_FOUND"


@pytest.mark.asyncio
async def test_gsheets_unsupported_document(
    fu_client,
    s3_client,
    s3_tmp_bucket,
    redis_model_manager,
):
    resp = await fu_client.make_request(ReqBuilder.upload_gsheet("14wGijvJtpF6M3_S-gh-nctYsj50jfstm"))
    assert resp.status == 201
    file_id = resp.json["file_id"]

    await asyncio.sleep(2)

    resp = await fu_client.make_request(ReqBuilder.file_status(file_id))
    assert resp.status == 200
    assert resp.json["file_id"] == file_id
    assert resp.json["status"] == "failed"
    assert resp.json["error"]["code"] == "ERR.FILE.UNSUPPORTED_DOCUMENT"


@pytest.mark.asyncio
async def test_all_types(
    use_local_task_processor,
    fu_client,
    s3_client,
    s3_tmp_bucket,
    redis_model_manager,
    env_param_getter,
):
    """
    Note: the spreadsheet in this test contains test data on it's _TEST_DATA_ sheet,
    which is filled in manually or generated with subsequent manual verification
    """

    # use this flag o print actual test data to fill the _TEST_DATA_ sheet when something changes in parsing,
    # be sure to check generated data manually
    print_actual = False

    spreadsheet_id = "1vTeWZz2M8c3bvTdRIrqsNhN9m4YwjUguwAnAqBahxi0"
    resp = await fu_client.make_request(ReqBuilder.upload_gsheet(spreadsheet_id))
    assert resp.status == 201
    file_id = resp.json["file_id"]
    assert file_id

    await asyncio.sleep(5)

    rmm = redis_model_manager
    df: DataFile = await DataFile.get(manager=rmm, obj_id=file_id)
    assert isinstance(df.user_source_properties, GSheetsUserSourceProperties)
    assert df.status == FileProcessingStatus.ready
    assert all(source.status == FileProcessingStatus.ready for source in df.sources)

    loaded_test_data = {}
    actual_data = {}

    gsheets_settings = GSheetsSettings(
        api_key=env_param_getter.get_str_value("GOOGLE_API_KEY"),
        client_id="dummy",
        client_secret="dummy",
    )
    auth = None
    with ContextVarExecutor(max_workers=min(32, os.cpu_count() * 3 + 4)) as tpe:
        async with GSheetsClient(settings=gsheets_settings, auth=auth, tpe=tpe) as sheets_client:
            spreadsheet = await sheets_client.get_spreadsheet(spreadsheet_id=spreadsheet_id)
    test_data_sheet = spreadsheet.sheets[0]
    test_value_keys = [cell.value for cell in test_data_sheet.data[0][1:]]
    for sheet_test_data_cells in test_data_sheet.data[1:]:
        sheet_title = sheet_test_data_cells[0].value
        loaded_test_data[sheet_title] = {}
        for test_value_key, test_value_cell in zip(test_value_keys, sheet_test_data_cells[1:]):
            loaded_test_data[sheet_title][test_value_key] = json.loads(test_value_cell.value)

    if print_actual:
        for test_value_key in test_value_keys:
            actual_data[test_value_key] = []
    for src in df.sources[1:]:
        resp = await fu_client.make_request(ReqBuilder.source_info(file_id, src.id))
        assert resp.status == 200
        source = resp.json["source"]
        assert source["title"]
        assert source["is_valid"]

        gsheet_title = source["title"].split()[-1]
        loaded_sheet_test_data = loaded_test_data[gsheet_title]
        for test_value_key, test_value in loaded_sheet_test_data.items():
            if not print_actual:
                assert source[test_value_key] == test_value, f"Failed on sheet {gsheet_title} checking {test_value_key}"
                LOGGER.info(f"{test_value_key} for {gsheet_title} is OK")
            else:
                actual_data[test_value_key].append(json.dumps(source[test_value_key], separators=(", ", ": ")))

    if print_actual:
        for test_value_key, actual_test_data in actual_data.items():
            print(f"\n\n\n{test_value_key}", "\n".join(actual_test_data), sep="\n")

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/ext/test_links_api.py
datalens-tech/datalens-backend
import asyncio
import datetime
import itertools
import logging
import uuid

import pytest

from dl_constants.enums import (
    FileProcessingStatus,
    UserDataType,
)
from dl_core.db import SchemaColumn
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_core_testing.connection import make_conn_key
from dl_file_uploader_api_lib_tests.req_builder import ReqBuilder
from dl_file_uploader_lib import exc
from dl_testing.s3_utils import s3_file_exists

from dl_connector_bundle_chs3.chs3_gsheets.core.constants import CONNECTION_TYPE_GSHEETS_V2
from dl_connector_bundle_chs3.chs3_gsheets.core.lifecycle import GSheetsFileS3ConnectionLifecycleManager
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection


LOGGER = logging.getLogger(__name__)


@pytest.fixture(scope="function")
async def saved_gsheets_v2_connection(loop, bi_context, default_async_usm_per_test, s3_persistent_bucket, s3_client):
    us_manager = default_async_usm_per_test
    conn_name = "gsheets_v2 test conn {}".format(uuid.uuid4())
    long_long_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(
        seconds=GSheetsFileS3ConnectionLifecycleManager.STALE_THRESHOLD_SECONDS + 60,  # just in case
    )

    dummy_raw_schema = [SchemaColumn("dummy_column", user_type=UserDataType.string)]
    data = GSheetsFileS3Connection.DataModel(
        sources=[
            GSheetsFileS3Connection.FileDataSource(  # this is a valid source
                id=f"source_1_{uuid.uuid4()}",
                file_id=str(uuid.uuid4()),
                title="Source 1",
                status=FileProcessingStatus.ready,
                s3_filename=f"src_1_filename_{uuid.uuid4()}",
                raw_schema=dummy_raw_schema,
                spreadsheet_id="1rnUFa7AiSKD5O80IKCvMy2cSZvLU1kRw9dxbtZbDMWc",
                sheet_id=1239604412,
                first_line_is_header=True,
                data_updated_at=long_long_ago,
            ),
            GSheetsFileS3Connection.FileDataSource(  # this sheet has no data
                id=f"source_2_{uuid.uuid4()}",
                file_id=str(uuid.uuid4()),
                title="Source 2",
                status=FileProcessingStatus.ready,
                s3_filename=f"src_2_filename_{uuid.uuid4()}",
                raw_schema=dummy_raw_schema,
                spreadsheet_id="1rnUFa7AiSKD5O80IKCvMy2cSZvLU1kRw9dxbtZbDMWc",
                sheet_id=0,
                first_line_is_header=True,
                data_updated_at=long_long_ago,
            ),
            GSheetsFileS3Connection.FileDataSource(  # this sheet does not exist
                id=f"source_3_{uuid.uuid4()}",
                file_id=str(uuid.uuid4()),
                title="Source 3",
                status=FileProcessingStatus.ready,
                s3_filename=f"src_3_filename_{uuid.uuid4()}",
                raw_schema=dummy_raw_schema,
                spreadsheet_id="1rnUFa7AiSKD5O80IKCvMy2cSZvLU1kRw9dxbtZbDMWc",
                sheet_id=99999999999,
                first_line_is_header=True,
                data_updated_at=long_long_ago,
            ),
            GSheetsFileS3Connection.FileDataSource(  # this sheet's whole spreadsheet does not exist
                id=f"source_4_{uuid.uuid4()}",
                file_id=str(uuid.uuid4()),
                title="Source 4",
                status=FileProcessingStatus.ready,
                s3_filename=f"src_4_filename_{uuid.uuid4()}",
                raw_schema=dummy_raw_schema,
                spreadsheet_id="ffffffffffffffffffffffffffffffffffffffffffff",
                sheet_id=99999999999,
                first_line_is_header=True,
                data_updated_at=long_long_ago,
            ),
            GSheetsFileS3Connection.FileDataSource(  # this is a valid source, but it failed during the previous update
                id=f"source_5_{uuid.uuid4()}",
                file_id=str(uuid.uuid4()),
                title="Source 5",
                status=FileProcessingStatus.failed,
                s3_filename=f"src_5_filename_{uuid.uuid4()}",  # removed below after file upload
                raw_schema=dummy_raw_schema,
                spreadsheet_id="1rnUFa7AiSKD5O80IKCvMy2cSZvLU1kRw9dxbtZbDMWc",
                sheet_id=1239604412,
                first_line_is_header=True,
                data_updated_at=long_long_ago,
            ),
            GSheetsFileS3Connection.FileDataSource(  # no access
                id=f"source_6_{uuid.uuid4()}",
                file_id=str(uuid.uuid4()),
                title="Source 6",
                status=FileProcessingStatus.ready,
                s3_filename=f"src_6_filename_{uuid.uuid4()}",
                raw_schema=dummy_raw_schema,
                spreadsheet_id="16pNg5gR7WQ1Rs58AM8YrUrNTcUmYXYdtuscRSOYUjtM",
                sheet_id=0,
                first_line_is_header=True,
                data_updated_at=long_long_ago,
            ),
        ]
    )
    conn = GSheetsFileS3Connection.create_from_dict(
        data,
        ds_key=make_conn_key("connections", conn_name),
        type_=CONNECTION_TYPE_GSHEETS_V2.name,
        meta={"title": conn_name, "state": "saved"},
        us_manager=us_manager,
    )
    await us_manager.save(conn)

    csv_data = "f1,f2,f3\nqwe,123,45.9\nasd,345,47.9".encode("utf-8")
    for src in conn.data.sources:
        await s3_client.put_object(
            ACL="private",
            Bucket=s3_persistent_bucket,
            Key=src.s3_filename,
            Body=csv_data,
        )

    for src in conn.data.sources:  # make sure all files are intact before the update
        assert await s3_file_exists(s3_client, s3_persistent_bucket, src.s3_filename)

    conn.data.sources[4].s3_filename = None
    await us_manager.save(conn)

    yield conn

    updated_conn: GSheetsFileS3Connection = await us_manager.get_by_id(conn.uuid, expected_type=GSheetsFileS3Connection)
    for src in itertools.chain(
        conn.data.sources, updated_conn.data.sources
    ):  # cleanup original and updated files if any
        if src.s3_filename is not None:
            await s3_client.delete_object(
                Bucket=s3_persistent_bucket,
                Key=src.s3_filename,
            )
    await us_manager.delete(conn)


@pytest.mark.asyncio
async def test_update_connection_data_with_save(
    fu_client,
    redis_cli,
    default_async_usm_per_test: AsyncUSManager,
    s3_client,
    s3_tmp_bucket,
    s3_persistent_bucket,
    saved_gsheets_v2_connection: GSheetsFileS3Connection,
):
    conn = saved_gsheets_v2_connection
    usm = default_async_usm_per_test
    data_updated_at_orig = conn.data.oldest_data_update_time(exclude_statuses={FileProcessingStatus.in_progress})

    resp = await fu_client.make_request(ReqBuilder.update_conn_data(conn, save=True))
    assert resp.status == 200

    await asyncio.sleep(5)

    updated_conn: GSheetsFileS3Connection = await usm.get_by_id(conn.uuid, expected_type=GSheetsFileS3Connection)

    new_data_updated_at = updated_conn.data.oldest_data_update_time(exclude_statuses={FileProcessingStatus.in_progress})
    assert data_updated_at_orig != new_data_updated_at

    src_1 = updated_conn.get_file_source_by_id(conn.data.sources[0].id)  # this is a valid source
    assert src_1.file_id != conn.data.sources[0].file_id
    assert await s3_file_exists(s3_client, s3_persistent_bucket, conn.data.sources[0].s3_filename) is False
    assert await s3_file_exists(s3_client, s3_persistent_bucket, updated_conn.data.sources[0].s3_filename) is True
    assert src_1.status == FileProcessingStatus.ready
    assert src_1.data_updated_at != data_updated_at_orig
    assert src_1.raw_schema is not None

    src_2 = updated_conn.get_file_source_by_id(conn.data.sources[1].id)  # this sheet has no data
    assert src_2.file_id == conn.data.sources[1].file_id
    assert updated_conn.data.sources[1].s3_filename is None
    assert await s3_file_exists(s3_client, s3_persistent_bucket, conn.data.sources[1].s3_filename) is False
    assert src_2.status == FileProcessingStatus.failed
    assert src_2.data_updated_at != data_updated_at_orig
    assert src_2.raw_schema is not None

    src_3 = updated_conn.get_file_source_by_id(conn.data.sources[2].id)  # this sheet does not exist
    assert src_3.file_id == conn.data.sources[2].file_id
    assert updated_conn.data.sources[2].s3_filename is None
    assert await s3_file_exists(s3_client, s3_persistent_bucket, conn.data.sources[2].s3_filename) is False
    assert src_3.status == FileProcessingStatus.failed
    assert src_3.data_updated_at != data_updated_at_orig
    assert src_3.raw_schema is not None

    src_4 = updated_conn.get_file_source_by_id(conn.data.sources[3].id)  # this sheet's whole spreadsheet does not exist
    assert src_4.file_id == conn.data.sources[3].file_id
    assert updated_conn.data.sources[3].s3_filename is None
    assert await s3_file_exists(s3_client, s3_persistent_bucket, conn.data.sources[3].s3_filename) is False
    assert src_4.status == FileProcessingStatus.failed
    assert src_4.data_updated_at != data_updated_at_orig
    assert src_4.raw_schema is not None

    src_5 = updated_conn.get_file_source_by_id(
        conn.data.sources[4].id
    )  # this is a valid source, but it failed during the previous update
    assert src_5.status == FileProcessingStatus.ready
    assert src_5.data_updated_at != data_updated_at_orig
    assert src_5.raw_schema is not None

    src_6 = updated_conn.get_file_source_by_id(conn.data.sources[5].id)  # no access
    assert src_6.file_id == conn.data.sources[5].file_id
    assert updated_conn.data.sources[5].s3_filename is None
    assert await s3_file_exists(s3_client, s3_persistent_bucket, conn.data.sources[5].s3_filename) is False
    assert src_6.status == FileProcessingStatus.failed
    assert src_6.data_updated_at != data_updated_at_orig
    assert src_6.raw_schema is not None

    error_registry = updated_conn.data.component_errors
    assert len(error_registry.items) == 4
    for err_pack in error_registry.items:
        assert all("request-id" in err.details for err in err_pack.errors)


@pytest.mark.asyncio
async def test_update_connection_data_without_save(
    fu_client,
    redis_cli,
    s3_tmp_bucket,
    s3_persistent_bucket,
    saved_gsheets_v2_connection: GSheetsFileS3Connection,
):
    conn = saved_gsheets_v2_connection
    resp = await fu_client.make_request(ReqBuilder.update_conn_data(conn, save=False))
    assert resp.status == 200

    file_ids = [file["file_id"] for file in resp.json["files"]]
    assert len(file_ids) == 3

    for file_id in file_ids:
        resp = await fu_client.make_request(ReqBuilder.file_status(file_id))
        assert resp.status == 200
        assert resp.json["file_id"] == file_id

    await asyncio.sleep(5)

    # spreadsheet I
    resp = await fu_client.make_request(ReqBuilder.file_status(file_ids[0]))
    assert resp.status == 200
    assert resp.json["file_id"] == file_ids[0]
    assert resp.json["status"] == "ready"

    resp = await fu_client.make_request(ReqBuilder.file_sources(file_ids[0]))
    sources = resp.json["sources"]
    assert sources[0]["is_applicable"]
    assert sources[0]["error"] is None

    assert not sources[1]["is_applicable"]
    assert sources[1]["error"]["code"] == "ERR.FILE.NO_DATA"
    assert "request-id" in sources[1]["error"]["details"]

    assert not sources[2]["is_applicable"]
    assert sources[2]["error"]["code"] == "ERR.FILE.NOT_FOUND"

    assert sources[3]["is_applicable"]
    assert sources[3]["error"] is None

    # spreadsheet II
    resp = await fu_client.make_request(ReqBuilder.file_status(file_ids[1]))
    assert resp.status == 200
    assert resp.json["file_id"] == file_ids[1]
    assert resp.json["status"] == "failed"

    resp = await fu_client.make_request(ReqBuilder.file_sources(file_ids[1]))
    sources = resp.json["sources"]
    assert not sources[0]["is_applicable"]
    assert sources[0]["error"]["code"] == "ERR.FILE.NOT_FOUND"

    resp = await fu_client.make_request(ReqBuilder.source_status(file_ids[1], sources[0]["source_id"]))
    assert resp.json["file_id"] == file_ids[1]
    assert resp.json["status"] == "failed"
    assert resp.json["error"]["code"] == "ERR.FILE.NOT_FOUND"

    # spreadsheet III
    resp = await fu_client.make_request(ReqBuilder.file_status(file_ids[2]))
    assert resp.status == 200
    assert resp.json["file_id"] == file_ids[2]
    assert resp.json["status"] == "failed"

    resp = await fu_client.make_request(ReqBuilder.file_sources(file_ids[2]))
    sources = resp.json["sources"]
    assert not sources[0]["is_applicable"]
    assert sources[0]["error"]["code"] == "ERR.FILE.PERMISSION_DENIED"

    resp = await fu_client.make_request(ReqBuilder.source_status(file_ids[2], sources[0]["source_id"]))
    assert resp.json["file_id"] == file_ids[2]
    assert resp.json["status"] == "failed"
    assert resp.json["error"]["code"] == "ERR.FILE.PERMISSION_DENIED"


@pytest.mark.asyncio
async def test_update_in_progress_sources(
    fu_client,
    default_async_usm_per_test: AsyncUSManager,
    redis_cli,
    s3_tmp_bucket,
    s3_persistent_bucket,
    saved_gsheets_v2_connection: GSheetsFileS3Connection,
):
    usm = default_async_usm_per_test
    conn = saved_gsheets_v2_connection

    # making the source incomplete
    conn.data.sources[0].sheet_id = None
    conn.data.sources[0].status = FileProcessingStatus.in_progress
    await usm.save(conn)

    resp = await fu_client.make_request(
        ReqBuilder.update_conn_data(
            conn,
            save=False,
            require_ok=False,
        )
    )
    assert resp.status == 400

    assert resp.json["message"] == exc.CannotUpdateDataError.default_message
    assert resp.json["details"] == {
        "incomplete_sources": [{"source_id": conn.data.sources[0].id, "title": conn.data.sources[0].title}]
    }


@pytest.mark.asyncio
async def test_component_errors(
    fu_client,
    redis_cli,
    default_async_usm_per_test: AsyncUSManager,
    s3_client,
    s3_tmp_bucket,
    s3_persistent_bucket,
    saved_gsheets_v2_connection: GSheetsFileS3Connection,
    use_local_task_processor,
):
    conn = saved_gsheets_v2_connection
    usm = default_async_usm_per_test

    resp = await fu_client.make_request(ReqBuilder.update_conn_data(conn, save=True))
    assert resp.status == 200
    await asyncio.sleep(5)

    updated_conn: GSheetsFileS3Connection = await usm.get_by_id(conn.uuid, expected_type=GSheetsFileS3Connection)

    src_3 = updated_conn.get_file_source_by_id(conn.data.sources[2].id)  # this sheet does not exist
    assert src_3.file_id == conn.data.sources[2].file_id
    assert updated_conn.data.sources[2].s3_filename is None
    assert await s3_file_exists(s3_client, s3_persistent_bucket, conn.data.sources[2].s3_filename) is False
    assert src_3.status == FileProcessingStatus.failed
    assert src_3.raw_schema is not None

    error_registry = updated_conn.data.component_errors
    src_3_errors = error_registry.get_pack(id=src_3.id).get_errors()
    assert len(src_3_errors) == 1
    assert src_3_errors[0].code == exc.DocumentNotFound.err_code

    # make it valid
    updated_conn.data.sources[2].sheet_id = 1239604412
    await usm.save(updated_conn)

    resp = await fu_client.make_request(ReqBuilder.update_conn_data(updated_conn, save=True))
    assert resp.status == 200
    await asyncio.sleep(5)

    final_conn: GSheetsFileS3Connection = await usm.get_by_id(conn.uuid, expected_type=GSheetsFileS3Connection)

    src_3 = final_conn.get_file_source_by_id(updated_conn.data.sources[2].id)  # this sheet does now exist
    assert src_3.file_id != updated_conn.data.sources[2].file_id
    assert await s3_file_exists(s3_client, s3_persistent_bucket, final_conn.data.sources[2].s3_filename) is True
    assert src_3.status == FileProcessingStatus.ready
    assert src_3.raw_schema is not None

    error_registry = final_conn.data.component_errors
    src_3_error_pack = error_registry.get_pack(id=src_3.id)
    assert not src_3_error_pack

    src_6 = final_conn.get_file_source_by_id(updated_conn.data.sources[5].id)  # no access
    assert src_6.file_id == updated_conn.data.sources[5].file_id
    assert final_conn.data.sources[5].s3_filename is None
    assert await s3_file_exists(s3_client, s3_persistent_bucket, conn.data.sources[5].s3_filename) is False
    assert src_6.status == FileProcessingStatus.failed
    assert src_6.raw_schema is not None

    # make sure errors are not duplicating
    for src in final_conn.data.sources:
        src_err_pack = error_registry.get_pack(id=src.id)
        assert src_err_pack is None or len(src_err_pack.errors) <= 1
        if src_err_pack is not None:
            assert all("request-id" in err.details for err in src_err_pack.errors)

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/ext/test_update_data.py
datalens-tech/datalens-backend
from typing import (
    Any,
    ClassVar,
    Optional,
)

import aiohttp

from dl_api_commons.client.common import Req
from dl_constants.api_constants import (
    DLHeaders,
    DLHeadersCommon,
)

from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection


class ReqBuilder:
    ORIGIN: ClassVar[str] = "https://foo.bar"

    @classmethod
    def upload_csv(cls, csv_data: str) -> Req:
        upload_req_data = (
            f"--111111111111\n"
            f'Content-Disposition: form-data; name="file"; filename="sample.csv"\n'
            f"Content-Type: application/octet-stream\n"
            f"\n"
            f"{csv_data}"
        )

        return Req(
            method="post",
            url="/api/v2/files",
            extra_headers={
                DLHeadersCommon.CONTENT_TYPE: "multipart/form-data; boundary=111111111111",
                DLHeadersCommon.ORIGIN: cls.ORIGIN,
            },
            data=upload_req_data,
        )

    @classmethod
    def upload_xlsx(cls, xlsx_data: str) -> Req:
        with aiohttp.MultipartWriter() as mpwriter:
            part = mpwriter.append(xlsx_data)
            part.set_content_disposition("form-data", name="file", filename="data.xlsx")
            return Req(
                method="post",
                url="/api/v2/files",
                extra_headers={DLHeadersCommon.ORIGIN: cls.ORIGIN},
                data=mpwriter,
            )

    @classmethod
    def upload_gsheet(
        cls,
        spreadsheet_id: Optional[str] = None,
        url: Optional[str] = None,
        authorized: bool = False,
        *,
        require_ok: bool = True,
    ) -> Req:
        if (spreadsheet_id is None) ^ (url is not None):
            raise ValueError("Expected exactly one of [`spreadsheet_id`, `url`] to be specified")
        if url is None:
            url = f"https://docs.google.com/spreadsheets/d/{spreadsheet_id}"

        return Req(
            method="post",
            url="/api/v2/links",
            data_json={
                "type": "gsheets",
                "url": url,
                "authorized": authorized,
            },
            require_ok=require_ok,
        )

    @classmethod
    def file_status(cls, file_id: str) -> Req:
        return Req(
            method="get",
            url=f"/api/v2/files/{file_id}/status",
        )

    @classmethod
    def file_sources(cls, file_id, *, require_ok: bool = True) -> Req:
        return Req(method="get", url=f"/api/v2/files/{file_id}/sources", require_ok=require_ok)

    @classmethod
    def source_status(cls, file_id: str, source_id: str) -> Req:
        return Req(
            method="get",
            url=f"/api/v2/files/{file_id}/sources/{source_id}/status",
        )

    @classmethod
    def source_info(cls, file_id: str, source_id: str, data_json: Optional[dict[str, Any]] = None) -> Req:
        return Req(
            method="post",
            url=f"/api/v2/files/{file_id}/sources/{source_id}",
            data_json=data_json,
        )

    @classmethod
    def apply_source_settings(cls, file_id: str, source_id: str, data_json: Optional[dict[str, Any]] = None) -> Req:
        return Req(
            method="post",
            url=f"/api/v2/files/{file_id}/sources/{source_id}/apply_settings",
            data_json=data_json,
        )

    @classmethod
    def preview(
        cls,
        file_id: str,
        source_id: str,
        master_token_header: Optional[dict[DLHeaders, str]],
        data_json: Optional[dict[str, Any]] = None,
    ) -> Req:
        return Req(
            method="post",
            url=f"/api/v2/files/{file_id}/sources/{source_id}/preview",
            extra_headers=master_token_header,
            data_json=data_json,
        )

    @classmethod
    def internal_params(
        cls,
        file_id: str,
        source_id: str,
        master_token_header: Optional[dict[DLHeaders, str]],
        data_json: Optional[dict[str, Any]] = None,
        *,
        require_ok: bool = True,
    ) -> Req:
        return Req(
            method="post",
            url=f"/api/v2/files/{file_id}/sources/{source_id}/internal_params",
            extra_headers=master_token_header,
            data_json=data_json,
            require_ok=require_ok,
        )

    @classmethod
    def update_conn_data(
        cls,
        connection: GSheetsFileS3Connection,
        save: bool,
        *,
        require_ok: bool = True,
    ) -> Req:
        sources_desc = [src.get_desc() for src in connection.data.sources]
        return Req(
            method="post",
            url="/api/v2/update_connection_data",
            data_json={
                "connection_id": connection.uuid,
                "authorized": False,
                "save": save,
                "sources": [
                    dict(
                        id=src_desc.source_id,
                        title=src_desc.title,
                        spreadsheet_id=src_desc.spreadsheet_id,
                        sheet_id=src_desc.sheet_id,
                        first_line_is_header=src_desc.first_line_is_header,
                    )
                    for src_desc in sources_desc
                ],
            },
            require_ok=require_ok,
        )

    @classmethod
    def cleanup(
        cls,
        tenant_id: str,
        master_token_header: Optional[dict[DLHeaders, str]] = None,
        *,
        require_ok: bool = True,
    ) -> Req:
        return Req(
            method="post",
            url="/api/v2/cleanup",
            extra_headers=master_token_header,
            data_json={
                "tenant_id": tenant_id,
            },
            require_ok=require_ok,
        )

    @classmethod
    def rename_tenant_files(cls, master_token_header: Optional[dict[DLHeaders, str]], tenant_id: str) -> Req:
        return Req(
            method="post",
            url="/api/v2/rename_tenant_files",
            extra_headers=master_token_header,
            data_json={
                "tenant_id": tenant_id,
            },
        )

lib/dl_file_uploader_api_lib/dl_file_uploader_api_lib_tests/req_builder.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging

import redis
from redis.asyncio.lock import Lock as RedisLock


LOGGER = logging.getLogger(__name__)


def get_update_connection_source_lock(src_id: str) -> tuple[str, str]:
    """Returns redis lock key and token"""

    return f"UpdateConnectionSource/{src_id}", src_id


async def release_source_update_locks(redis_cli: redis.asyncio.Redis, *source_ids: str) -> None:
    for source_id in source_ids:
        source_update_lock_key, source_update_lock_token = get_update_connection_source_lock(src_id=source_id)
        source_lock = RedisLock(redis=redis_cli, name=source_update_lock_key)
        source_lock.local.token = source_update_lock_token  # type: ignore
        try:
            await source_lock.release()
        except redis.exceptions.LockError:
            LOGGER.info(f"Lock {source_update_lock_key} is already released")
        else:
            LOGGER.info(f"Released lock {source_update_lock_key}")

lib/dl_file_uploader_lib/dl_file_uploader_lib/common_locks.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Optional,
)


if TYPE_CHECKING:
    from types_aiobotocore_s3 import S3Client as AsyncS3Client
    from mypy_boto3_s3.client import S3Client as SyncS3Client

import ujson as json

from dl_constants.enums import ConnectionType
from dl_core.data_sink import (
    DataSink,
    DataSinkAsync,
)
from dl_core.db import (
    SchemaColumn,
    get_type_transformer,
)
from dl_core.raw_data_streaming.stream import (
    DataStreamBase,
    SimpleUntypedAsyncDataStream,
    SimpleUntypedDataStream,
)
from dl_file_uploader_lib import exc

from dl_connector_bundle_chs3.file.core.constants import CONNECTION_TYPE_FILE


LOGGER = logging.getLogger(__name__)


class S3JsonEachRowFileDataSink(DataSink):
    conn_type: ConnectionType = CONNECTION_TYPE_FILE
    batch_size_in_bytes: int = 30 * 1024**2
    max_batch_size: Optional[int] = None

    _upload_id: Optional[str] = None
    _part_tags: Optional[list[tuple[int, str]]] = None
    _part_number: int = 1
    _multipart_upload_started: bool = False

    def __init__(  # type: ignore
        self,
        bi_schema: list[SchemaColumn],
        s3: SyncS3Client,
        s3_key: str,
        bucket_name: str,
    ):
        super().__init__(bi_schema=bi_schema)

        self._s3 = s3
        self._s3_key = s3_key
        self._bucket_name = bucket_name

        self._row_index = 0
        self._rows_saved = 0
        self._tt = get_type_transformer(self.conn_type)
        self._batch_size = None

    def initialize(self) -> None:
        mp_upl_resp = self._s3.create_multipart_upload(
            ACL="private",
            Bucket=self._bucket_name,
            Key=self._s3_key,
        )
        self._multipart_upload_started = True
        self._upload_id = mp_upl_resp["UploadId"]
        LOGGER.info(f"Created multipart upload. S3 key: {self._s3_key}, UploadId: {self._upload_id}")
        self._part_tags = []
        self._part_number = 1

    def finalize(self) -> None:
        LOGGER.info(f"Completing S3 multipart upload. {self._part_number - 1} parts were uploaded.")
        assert self._upload_id is not None
        if self._multipart_upload_started:
            if self._part_tags:
                self._s3.complete_multipart_upload(
                    Bucket=self._bucket_name,
                    Key=self._s3_key,
                    UploadId=self._upload_id,
                    MultipartUpload={"Parts": [{"ETag": etag, "PartNumber": pnum} for pnum, etag in self._part_tags]},
                )
                LOGGER.info("Multipart upload competed.")
                self._multipart_upload_started = False

    def cleanup(self) -> None:
        if self._multipart_upload_started:
            LOGGER.info("Aborting S3 multipart upload,")
            assert self._upload_id is not None
            self._s3.abort_multipart_upload(
                Bucket=self._bucket_name,
                Key=self._s3_key,
                UploadId=self._upload_id,
            )
            LOGGER.info("Multipart upload aborted.")
            self._multipart_upload_started = False

    def _prepare_chunk_body(self, batch: list[bytes]) -> bytes:
        return b"\n".join(batch) + b"\n"

    def _dump_data_batch(self, batch: list[bytes], progress: int) -> None:
        LOGGER.info(f"Dumping {len(batch)} data rows into s3 file {self._s3_key}.")
        assert self._upload_id is not None
        part_resp = self._s3.upload_part(
            Bucket=self._bucket_name,
            Key=self._s3_key,
            UploadId=self._upload_id,
            PartNumber=self._part_number,
            Body=self._prepare_chunk_body(batch),
        )
        LOGGER.info(f"Part number {self._part_number} uploaded.")
        assert self._part_tags is not None
        self._part_tags.append((self._part_number, part_resp["ETag"]))
        self._part_number += 1

        self._rows_saved += len(batch)
        LOGGER.info(f"Copied: {self._rows_saved} rows ({progress}%%) to S3")

    def _process_row(self, row_data: dict) -> bytes:
        cast_row_data = [self._tt.cast_for_input(row_data[col.name], user_t=col.user_type) for col in self._bi_schema]
        return json.dumps(cast_row_data).encode("utf-8")

    def _should_dump_batch(self, batch: list[bytes], batch_size: int) -> bool:
        assert self.batch_size_in_bytes
        if not batch:
            return False
        return batch_size >= self.batch_size_in_bytes

    def _dump_data_stream_base(self, data_stream: DataStreamBase) -> None:
        batch: list[bytes] = []
        batch_size = 0
        for row_data in data_stream:
            processed_row = self._process_row(row_data)

            if self._should_dump_batch(batch, batch_size=batch_size):
                self._dump_data_batch(batch, progress=data_stream.get_progress_percent())
                batch.clear()
                batch_size = 0

            batch.append(processed_row)
            batch_size += len(processed_row)
            self._row_index += 1

        if batch:
            self._dump_data_batch(batch, progress=data_stream.get_progress_percent())

    def dump_data_stream(self, data_stream: DataStreamBase) -> None:
        self._dump_data_stream_base(data_stream)


class S3JsonEachRowUntypedFileDataSink(S3JsonEachRowFileDataSink):
    def __init__(  # type: ignore
        self,
        s3: SyncS3Client,
        s3_key: str,
        bucket_name: str,
    ):
        super().__init__(
            bi_schema=[],  # ignore bi_schema
            s3=s3,
            s3_key=s3_key,
            bucket_name=bucket_name,
        )

    def _process_row(self, row_data: list) -> bytes:  # type: ignore
        # skip type casts
        return json.dumps(row_data).encode("utf-8")

    def dump_data_stream(self, data_stream: SimpleUntypedDataStream) -> None:  # type: ignore
        self._dump_data_stream_base(data_stream)


class S3JsonEachRowUntypedFileAsyncDataSink(DataSinkAsync[SimpleUntypedAsyncDataStream]):
    conn_type: ConnectionType = CONNECTION_TYPE_FILE
    batch_size_in_bytes: int = 30 * 1024**2
    max_batch_size: Optional[int] = None
    max_file_size_bytes: ClassVar[int] = 200 * 1024**2

    _upload_id: Optional[str] = None
    _part_tags: Optional[list[tuple[int, str]]] = None
    _part_number: int = 1
    _multipart_upload_started: bool = False

    def __init__(self, s3: AsyncS3Client, s3_key: str, bucket_name: str):
        self._s3 = s3
        self._s3_key = s3_key
        self._bucket_name = bucket_name

        self._row_index = 0
        self._rows_saved = 0
        self._bytes_saved = 0
        self._batch_size = None

    async def initialize(self) -> None:
        mp_upl_resp = await self._s3.create_multipart_upload(
            ACL="private",
            Bucket=self._bucket_name,
            Key=self._s3_key,
        )
        self._multipart_upload_started = True
        self._upload_id = mp_upl_resp["UploadId"]
        LOGGER.info(f"Created multipart upload. S3 key: {self._s3_key}, UploadId: {self._upload_id}")
        self._part_tags = []
        self._part_number = 1

    async def finalize(self) -> None:
        if self._multipart_upload_started:
            LOGGER.info(f"Completing S3 multipart upload. {self._part_number - 1} parts were uploaded.")
            assert self._upload_id is not None
            await self._s3.complete_multipart_upload(
                Bucket=self._bucket_name,
                Key=self._s3_key,
                UploadId=self._upload_id,
                MultipartUpload={
                    "Parts": [{"ETag": etag, "PartNumber": pnum} for pnum, etag in (self._part_tags or [])]
                },
            )
            LOGGER.info("Multipart upload competed.")
            self._multipart_upload_started = False

    async def cleanup(self) -> None:
        if self._multipart_upload_started:
            LOGGER.exception("Aborting S3 multipart upload,")
            assert self._upload_id is not None
            await self._s3.abort_multipart_upload(
                Bucket=self._bucket_name,
                Key=self._s3_key,
                UploadId=self._upload_id,
            )
            LOGGER.info("Multipart upload aborted.")
            self._multipart_upload_started = False

    def _prepare_chunk_body(self, batch: list[bytes]) -> bytes:
        return b"\n".join(batch) + b"\n"

    async def _dump_data_batch(self, batch: list[bytes], progress: int) -> None:
        LOGGER.info(f"Dumping {len(batch)} data rows into s3 file {self._s3_key}.")
        assert self._upload_id is not None
        batch_to_write = self._prepare_chunk_body(batch)
        part_resp = await self._s3.upload_part(
            Bucket=self._bucket_name,
            Key=self._s3_key,
            UploadId=self._upload_id,
            PartNumber=self._part_number,
            Body=batch_to_write,
        )
        LOGGER.info(f"Part number {self._part_number} uploaded.")
        assert isinstance(self._part_tags, list)
        self._part_tags.append((self._part_number, part_resp["ETag"]))
        self._part_number += 1

        self._rows_saved += len(batch)
        self._bytes_saved += len(batch_to_write)
        LOGGER.info(f"Copied: {self._rows_saved} rows ({progress}%%) to S3")

    def _process_row(self, row_data: list) -> bytes:
        return json.dumps(row_data).encode("utf-8")

    def _should_dump_batch(self, batch: list[bytes], batch_size: int) -> bool:
        assert self.batch_size_in_bytes
        if not batch:
            return False
        return batch_size >= self.batch_size_in_bytes

    async def _dump_data_stream_base(self, data_stream: SimpleUntypedAsyncDataStream) -> None:
        batch: list[bytes] = []
        batch_size = 0
        async for row_data in data_stream:
            processed_row = self._process_row(row_data)

            if self._bytes_saved > self.max_file_size_bytes:
                raise exc.FileLimitError

            if self._should_dump_batch(batch, batch_size=batch_size):
                await self._dump_data_batch(batch, progress=data_stream.get_progress_percent())
                batch.clear()
                batch_size = 0

            batch.append(processed_row)
            batch_size += len(processed_row)
            self._row_index += 1

        if batch:
            actual_size = self._bytes_saved + len(batch)
            if actual_size > self.max_file_size_bytes:
                raise exc.FileLimitError
            await self._dump_data_batch(batch, progress=data_stream.get_progress_percent())

    async def dump_data_stream(self, data_stream: SimpleUntypedAsyncDataStream) -> None:
        await self._dump_data_stream_base(data_stream)

lib/dl_file_uploader_lib/dl_file_uploader_lib/data_sink/json_each_row.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    AsyncIterator,
    ClassVar,
    Optional,
)


if TYPE_CHECKING:
    from types_aiobotocore_s3 import S3Client

from dl_constants.enums import ConnectionType
from dl_core.data_sink import DataSinkAsync
from dl_core.raw_data_streaming.stream import AsyncDataStreamBase
from dl_file_uploader_lib import exc

from dl_connector_bundle_chs3.file.core.constants import CONNECTION_TYPE_FILE


LOGGER = logging.getLogger(__name__)


class RawBytesAsyncDataStream(AsyncDataStreamBase[bytes]):
    def __init__(self, data_iter: AsyncIterator[bytes], bytes_to_copy: Optional[int] = None):
        self._data_iter = data_iter
        self.bytes_to_copy = bytes_to_copy

        self._bytes_read = 0

    def get_progress_percent(self) -> int:
        if self.bytes_to_copy is None:
            return 0
        if self.bytes_to_copy == 0:
            return 100
        return min(self.max_percent, int((self._bytes_read / self.bytes_to_copy) * 100))

    async def __anext__(self) -> bytes:
        chunk = await self._data_iter.__anext__()
        self._bytes_read += len(chunk)
        return chunk


class S3RawFileAsyncDataSink(DataSinkAsync[RawBytesAsyncDataStream]):
    conn_type: ConnectionType = CONNECTION_TYPE_FILE
    batch_size_in_bytes: int = 10 * 1024**2
    max_batch_size: Optional[int] = None
    max_file_size_bytes: ClassVar[int] = 200 * 1024**2

    _upload_id: Optional[str] = None
    _part_tags: Optional[list[tuple[int, str]]] = None
    _part_number: int = 1
    _multipart_upload_started: bool = False
    _chunks_saved: int = 0
    _bytes_saved: int = 0

    def __init__(self, s3: S3Client, s3_key: str, bucket_name: str):
        self._s3 = s3
        self._s3_key = s3_key
        self._bucket_name = bucket_name

        self._chunks_saved = 0
        self._bytes_saved = 0

    async def initialize(self) -> None:
        mp_upl_resp = await self._s3.create_multipart_upload(
            ACL="private",
            Bucket=self._bucket_name,
            Key=self._s3_key,
        )
        self._multipart_upload_started = True
        self._upload_id = mp_upl_resp["UploadId"]
        LOGGER.info(f"Created multipart upload. S3 key: {self._s3_key}, UploadId: {self._upload_id}")
        self._part_tags = []
        self._part_number = 1

    async def finalize(self) -> None:
        if self._multipart_upload_started:
            LOGGER.info(f"Completing S3 multipart upload. {self._part_number - 1} parts were uploaded.")
            assert self._upload_id is not None
            await self._s3.complete_multipart_upload(
                Bucket=self._bucket_name,
                Key=self._s3_key,
                UploadId=self._upload_id,
                MultipartUpload={
                    "Parts": [{"ETag": etag, "PartNumber": pnum} for pnum, etag in (self._part_tags or [])]
                },
            )
            LOGGER.info("Multipart upload competed.")
            self._multipart_upload_started = False

    async def cleanup(self) -> None:
        if self._multipart_upload_started:
            LOGGER.exception("Aborting S3 multipart upload,")
            assert self._upload_id is not None
            await self._s3.abort_multipart_upload(
                Bucket=self._bucket_name,
                Key=self._s3_key,
                UploadId=self._upload_id,
            )
            LOGGER.info("Multipart upload aborted.")
            self._multipart_upload_started = False

    async def _dump_data_batch(self, batch: bytes, progress: int) -> None:
        LOGGER.info(f"Dumping {len(batch)} data rows into s3 file {self._s3_key}.")
        assert self._upload_id is not None
        part_resp = await self._s3.upload_part(
            Bucket=self._bucket_name,
            Key=self._s3_key,
            UploadId=self._upload_id,
            PartNumber=self._part_number,
            Body=batch,
        )
        LOGGER.info(f"Part number {self._part_number} uploaded.")
        assert isinstance(self._part_tags, list)
        self._part_tags.append((self._part_number, part_resp["ETag"]))
        self._part_number += 1

        self._bytes_saved += len(batch)
        LOGGER.info(f"Copied: {self._bytes_saved} bytes in {self._chunks_saved} chunks ({progress}%%) to S3")

    def _should_dump_batch(self, batch: bytes) -> bool:
        assert self.batch_size_in_bytes
        if not batch:
            return False
        return len(batch) >= self.batch_size_in_bytes

    async def _dump_data_stream_base(self, data_stream: RawBytesAsyncDataStream) -> None:
        batch = b""
        async for chunk in data_stream:
            if self._bytes_saved > self.max_file_size_bytes:
                raise exc.FileLimitError

            if self._should_dump_batch(batch):
                await self._dump_data_batch(batch, progress=data_stream.get_progress_percent())
                batch = b""

            batch += chunk
            self._chunks_saved += 1

        if batch:
            actual_size = self._bytes_saved + len(batch)
            if actual_size > self.max_file_size_bytes:
                raise exc.FileLimitError
            await self._dump_data_batch(batch, progress=data_stream.get_progress_percent())

    async def dump_data_stream(self, data_stream: RawBytesAsyncDataStream) -> None:
        await self._dump_data_stream_base(data_stream)

lib/dl_file_uploader_lib/dl_file_uploader_lib/data_sink/raw_bytes.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import Enum


class ErrorLevel(Enum):
    error = "error"
    warning = "warning"


class ErrorObjectKind(Enum):
    file = "file"
    source = "source"
    # system = 'system'


class FileType(Enum):
    csv = "csv"
    gsheets = "gsheets"
    xlsx = "xlsx"
    yadocs = "yadocs"


class CSVEncoding(Enum):
    utf8 = "utf-8"
    windows1251 = "windows-1251"
    utf8sig = "utf-8-sig"
    utf16 = "utf-16"


class CSVDelimiter(Enum):
    comma = ","
    semicolon = ";"
    tab = "\t"


class RenameTenantStatus(Enum):
    scheduled = "scheduled"
    started = "started"
    error = "error"
    success = "success"

lib/dl_file_uploader_lib/dl_file_uploader_lib/enums.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Type

from dl_constants.exc import (
    GLOBAL_ERR_PREFIX,
    DLBaseException,
)
from dl_file_uploader_lib.enums import ErrorLevel


def make_err_code(exc: Type[DLFileUploaderBaseError] | DLFileUploaderBaseError) -> str:
    return ".".join([GLOBAL_ERR_PREFIX] + exc.err_code)


class DLFileUploaderBaseError(DLBaseException):
    err_code = DLBaseException.err_code + ["FILE"]
    default_level = ErrorLevel.error


class CannotUpdateDataError(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["CANNOT_UPDATE_DATA"]
    default_message = "Source data cannot be updated, sources are most likely still being processed, try again later"


class FileLimitError(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["FILE_LIMIT_EXCEEDED"]
    default_message = "The file is too large"


class TooManyColumnsError(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["TOO_MANY_COLUMNS"]
    default_message = "Too many columns"


class TooFewRowsError(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["TOO_FEW_ROWS"]
    default_message = "Too few rows"


class InvalidFieldCast(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["INVALID_FIELD_CAST"]
    default_message = "Invalid field cast"


class PermissionDenied(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["PERMISSION_DENIED"]
    default_message = "The caller does not have permission"


class DocumentNotFound(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["NOT_FOUND"]
    default_message = "Requested document was not found"


class UnsupportedDocument(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["UNSUPPORTED_DOCUMENT"]
    default_message = "This document is not supported"


class InvalidLink(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["INVALID_LINK"]
    default_message = "Invalid format link provided"


class EmptyDocument(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["NO_DATA"]
    default_message = "Empty document"


class DownloadFailed(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["DOWNLOAD_FAILED"]
    default_message = "Download failed"


class ParseFailed(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["PARSE_FAILED"]
    default_message = "Parsing failed"


class InvalidSource(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["INVALID_SOURCE"]
    default_message = "Specified source is not applicable"


class RemoteServerError(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["REMOTE_SERVER_ERROR"]
    default_message = "Remote server error"
    default_level = ErrorLevel.warning


class QuotaExceeded(DLFileUploaderBaseError):
    err_code = DLFileUploaderBaseError.err_code + ["QUOTA_EXCEEDED"]
    default_message = "Quota exceeded, try again later. If the issue persists, please contact support"
    default_level = ErrorLevel.warning


class QuotaExceededReadRequestsPerMinutePerProject(QuotaExceeded):
    err_code = QuotaExceeded.err_code + ["PER_MINUTE_PER_PROJECT"]
    default_message = (
        "Quota exceeded (read requests per minute per project), try again later. "
        "If the issue persists, please contact support"
    )


class QuotaExceededReadRequestsPerMinutePerUser(QuotaExceeded):
    err_code = QuotaExceeded.err_code + ["PER_MINUTE_PER_USER"]
    default_message = (
        "Quota exceeded (read requests per minute per user), try again later. "
        "If the issue persists, please contact support"
    )

lib/dl_file_uploader_lib/dl_file_uploader_lib/exc.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
from asyncio import AbstractEventLoop
import datetime
import logging
import random
from ssl import SSLContext
import time
from typing import (
    Any,
    ClassVar,
    Optional,
    Tuple,
    Type,
    Union,
)

from aiogoogle import (
    Aiogoogle,
    GoogleAPI,
    HTTPError,
)
from aiogoogle.auth import (
    ApiKey,
    UserCreds,
)
from aiogoogle.auth.creds import ClientCreds
import aiogoogle.excs
import aiogoogle.models
from aiogoogle.models import Response
from aiogoogle.sessions.aiohttp_session import AiohttpSession
from aiohttp import (
    ClientResponse,
    ClientTimeout,
    Fingerprint,
)
from aiohttp.typedefs import StrOrURL
import attr
from yarl import URL

from dl_constants.enums import UserDataType
from dl_core.aio.web_app_services.gsheets import (
    Cell,
    GSheetsSettings,
    NumberFormatType,
    Range,
    Sheet,
    Spreadsheet,
)
from dl_file_uploader_lib import exc as file_upl_exc
from dl_utils.aio import ContextVarExecutor


LOGGER = logging.getLogger(__name__)


# this is the dt from which all DATE and DATE_TIME values are calculated in gsheets
# https://developers.google.com/sheets/api/guides/formats#about_date_time_values
GSHEETS_EPOCH = datetime.datetime(
    year=1899,
    month=12,
    day=30,
)
URL_SAFE_CHARS = ""  # we want to encode all symbols with no exceptions


def google_api_error_to_file_uploader_exception(err: HTTPError) -> file_upl_exc.DLFileUploaderBaseError:
    if not isinstance(err.res, Response):
        LOGGER.warning(f"Unknown aiogoogle http error: {err}")
        return file_upl_exc.DLFileUploaderBaseError(orig=err)

    orig_error_obj: dict[str, Any] = err.res.json.get("error", {})
    if isinstance(orig_error_obj, dict):
        orig_status_code = orig_error_obj.get("code", -1)
        orig_reason = orig_error_obj.get("status")
    else:
        orig_status_code = err.res.status_code
        orig_reason = None

    err_cls: Type[file_upl_exc.DLFileUploaderBaseError]
    if orig_status_code == 403:
        err_cls = file_upl_exc.PermissionDenied
    elif orig_status_code == 404:
        err_cls = file_upl_exc.DocumentNotFound
    elif orig_status_code == 400 and orig_reason == "FAILED_PRECONDITION":
        err_cls = file_upl_exc.UnsupportedDocument
    elif orig_status_code >= 500:
        err_cls = file_upl_exc.RemoteServerError
    else:
        LOGGER.warning(f"Unknown aiogoogle http error: {err}")
        err_cls = file_upl_exc.DLFileUploaderBaseError

    return err_cls(
        orig=err,
        debug_info=err.res.json,
        details=dict(
            orig_error=orig_error_obj,
        ),
    )


@attr.s
class GSheetsOAuth2:
    access_token: Optional[str] = attr.ib(repr=False)
    refresh_token: str = attr.ib(repr=False)


class AiohttpGSheetsSession(AiohttpSession):
    def __init__(
        self,
        *args: Any,
        proxy: Optional[StrOrURL] = None,
        proxy_headers: Optional[dict[str, str]] = None,
        ssl: Optional[Union[SSLContext, bool, Fingerprint]] = True,
        **kwargs: Any,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.proxy = proxy
        self.proxy_headers = proxy_headers
        self.ssl = ssl
        self.last_response_size_bytes = -1

    async def _request(
        self,
        method: str,
        str_or_url: StrOrURL,
        **kwargs: Any,
    ) -> ClientResponse:
        url = URL(str_or_url)
        headers = kwargs.pop("headers", {})
        if "key" in url.query:
            query_params = {**url.query}
            key = query_params.pop("key")
            url = url.with_query(query_params)
            headers["X-goog-api-key"] = key

        LOGGER.info(f"Sending request: {method} {url}")
        start_t = time.monotonic()
        resp = await super()._request(
            method,
            url,
            allow_redirects=False,
            proxy=self.proxy,
            proxy_headers=self.proxy_headers,
            ssl=self.ssl,
            headers=headers,
            **kwargs,
        )
        elapsed = time.monotonic() - start_t
        resp_text = await resp.read()
        self.last_response_size_bytes = len(resp_text)
        LOGGER.info(
            f"Received response: {method} {url},"
            f" status_code: {resp.status},"
            f" elapsed: {elapsed:.6f} s,"
            f" payload_size: {self.last_response_size_bytes} bytes"
        )
        return resp


def make_type(value: Any, user_type: UserDataType | str) -> Any:
    if value is None or value == "":
        return None
    if user_type == UserDataType.integer:
        if isinstance(value, str):  # overflow
            return None
        return int(value)
    if user_type == UserDataType.float:
        if isinstance(value, str):  # overflow
            return None
        return float(value)
    if user_type == UserDataType.boolean:
        return bool(value)
    if user_type in (UserDataType.date, UserDataType.genericdatetime, UserDataType.datetime, UserDataType.datetimetz):
        actual_dt = GSHEETS_EPOCH + datetime.timedelta(days=value)
        if user_type == UserDataType.date:
            dt_str = actual_dt.strftime("%Y-%m-%d")
        else:
            dt_str = actual_dt.strftime("%Y-%m-%d %H:%M:%S")
        year, *rest = dt_str.split("-")
        return f"{year:>04}-" + "-".join(rest)  # years <1000 are not padded properly in strftime
    if user_type == "time":
        time_value = datetime.timedelta(days=value)
        hours, remainder = divmod(time_value.total_seconds(), 3600)
        minutes, seconds = divmod(remainder, 60)
        return f"{int(hours):02}:{int(minutes):02}:{int(seconds):02}"
    if user_type == UserDataType.string:
        return str(value)

    raise ValueError(f"Type {user_type} is not supported here")


@attr.s
class GSheetsClient:
    settings: GSheetsSettings = attr.ib()
    _tpe: ContextVarExecutor = attr.ib()
    _loop: AbstractEventLoop = attr.ib(init=False, factory=asyncio.get_event_loop)
    auth: Optional[GSheetsOAuth2] = attr.ib(default=None)

    session_timeout: ClassVar[ClientTimeout] = ClientTimeout(total=180.0)

    last_request_size_bytes: int = attr.ib(init=False, default=-1)
    _aiogoogle: Optional[Aiogoogle] = attr.ib(init=False, default=None)
    _sheets_api: Optional[GoogleAPI] = attr.ib(init=False, default=None)

    async def __aenter__(self: GSheetsClient) -> GSheetsClient:
        self._aiogoogle = Aiogoogle(
            session_factory=lambda: AiohttpGSheetsSession(
                timeout=self.session_timeout,
                headers={
                    "user-agent": "DataLens",
                    "Accept": "application/json",
                },
            ),
            api_key=ApiKey(self.settings.api_key),
            client_creds=ClientCreds(
                client_id=self.settings.client_id,
                client_secret=self.settings.client_secret,
            ),
        )
        await self._aiogoogle.__aenter__()
        return self

    async def __aexit__(self, *args: Any) -> None:
        assert self._aiogoogle, "Aiogoogle is not initialized for the client"
        await self._aiogoogle.__aexit__(*args)

    def _process_values(
        self, raw_values: list[list[Any]], user_types: list[UserDataType | str]
    ) -> Tuple[list[list[Any]], list[UserDataType | str]]:
        """
        Tries to convert values to passed BITypes and falls back to string when fails to do so
        But the fallback happens only on return, i.e. it tries to convert all values to the original passed type
        Values are processed inplace, types are copied
        """

        new_user_types = user_types.copy()
        for row_idx, row in enumerate(raw_values):
            width_diff = len(user_types) - len(row)
            if width_diff > 0:
                row.extend([None] * width_diff)
            else:
                row = row[: len(user_types)]
            for col_idx, (value, user_type) in enumerate(zip(row, user_types)):
                try:
                    raw_values[row_idx][col_idx] = make_type(value, user_type)
                except (ValueError, TypeError):
                    new_user_types[col_idx] = UserDataType.string
                    raw_values[row_idx][col_idx] = str(value)
                except OverflowError:
                    raw_values[row_idx][col_idx] = None
        return raw_values, new_user_types

    def _get_cell_data(self, value: dict[str, Any]) -> Cell:
        top_level_value_keys = ("userEnteredValue", "effectiveValue", "formattedValue")
        if not value or not any(value_key in value for value_key in top_level_value_keys):
            return Cell(
                value=None,
                number_format_type=NumberFormatType.NUMBER_FORMAT_TYPE_UNSPECIFIED,
                empty=True,
            )
        number_format = NumberFormatType(
            value.get("effectiveFormat", {}).get(
                "numberFormat", {"type": NumberFormatType.NUMBER_FORMAT_TYPE_UNSPECIFIED}
            )["type"]
        )

        extended_value_key = (
            "formattedValue"
            if number_format
            in (
                NumberFormatType.TIME,
                NumberFormatType.NUMBER_FORMAT_TYPE_UNSPECIFIED,
                NumberFormatType.TEXT,
            )
            else "effectiveValue"
        )
        if "effectiveValue" in value and "boolValue" in value["effectiveValue"]:
            # bool values may be marked with any number format
            extended_value = value["effectiveValue"]
        else:
            extended_value = value.get(extended_value_key, None)

        if isinstance(extended_value, dict):  # formattedValue is a string
            value_key = list(extended_value.keys())[0]
            actual_value = extended_value[value_key]
            if value_key == "boolValue":
                number_format = NumberFormatType.BOOLEAN
        else:
            actual_value = extended_value

        NUMBER_FORMATS = (
            NumberFormatType.NUMBER,
            NumberFormatType.PERCENT,
            NumberFormatType.CURRENCY,
            NumberFormatType.SCIENTIFIC,
        )

        if number_format in NUMBER_FORMATS:
            if isinstance(actual_value, float) or actual_value is None:
                number_format = NumberFormatType.FLOAT
            elif isinstance(actual_value, int):
                number_format = NumberFormatType.INTEGER
            elif isinstance(actual_value, str):
                try:
                    _ = float(actual_value)
                except ValueError:
                    # not a number
                    number_format = NumberFormatType.TEXT
                else:
                    # if the conversion was successful, this means that the number is too big to be encoded in json
                    actual_value = None
                    number_format = NumberFormatType.FLOAT
            else:
                actual_value = str(actual_value)
                number_format = NumberFormatType.TEXT

        if number_format in (NumberFormatType.DATE, NumberFormatType.DATE_TIME):
            if isinstance(actual_value, (int, float)):
                gsheets_internal_value = actual_value

                try:
                    actual_dt = GSHEETS_EPOCH + datetime.timedelta(days=gsheets_internal_value)
                    if number_format == NumberFormatType.DATE:
                        actual_value = actual_dt.strftime("%Y-%m-%d")
                    elif number_format == NumberFormatType.DATE_TIME:
                        actual_value = actual_dt.strftime("%Y-%m-%d %H:%M:%S")
                except OverflowError:
                    actual_value = None
            else:
                number_format = NumberFormatType.TEXT

        return Cell(
            value=actual_value,
            number_format_type=number_format,
        )

    def _get_sheet_data(self, sheet: dict[str, Any], num_rows: Optional[int] = None) -> list[list[Cell]]:
        data: list[list[Cell]] = []
        rows_read = 0
        for rowdata in sheet["data"][0].get("rowData", []):
            if num_rows is not None and rows_read >= num_rows:
                break
            processed_rowdata = [self._get_cell_data(value) for value in rowdata.get("values", [])]
            actual_data_ends_at = len(processed_rowdata) - 1
            while actual_data_ends_at > 0 and processed_rowdata[actual_data_ends_at].empty:
                actual_data_ends_at -= 1
            processed_rowdata = processed_rowdata[: actual_data_ends_at + 1]
            data.append(processed_rowdata)
            rows_read += 1
        if data:
            n_cols = max(len(row) for row in data)
            for idx, row in enumerate(data):  # filling incomplete rows with NULLs and trimming too wide rows
                diff = n_cols - len(row)
                if diff > 0:
                    data[idx].extend([Cell(value=None, empty=True)] * diff)
                elif diff < 0:
                    data[idx] = row[:n_cols]
        return data

    async def _require_init(self) -> None:
        if self._aiogoogle is None:
            raise ValueError("Aiogoogle is not set up for the client. It must be used as an async context manager.")
        if self._sheets_api is None:
            self._sheets_api = await self._aiogoogle.discover("sheets", "v4")

    def _is_retryable_status(self, status_code: int) -> bool:
        return status_code in (408, 429) or status_code >= 500

    def _raise_retryable(self, err: aiogoogle.excs.HTTPError) -> None:
        """Makes an exception based on the passed error response and raises"""

        err_resp_json: dict[str, Any] = err.res.json or {}
        details: list[dict] = err_resp_json["error"].get("details", [])

        if err.res.status_code == 429:
            err_info: dict[str, Any] = next((item for item in details if item["reason"] == "RATE_LIMIT_EXCEEDED"), {})
            quota_err_cls = {
                "ReadRequestsPerMinutePerProject": file_upl_exc.QuotaExceededReadRequestsPerMinutePerProject,
                "ReadRequestsPerMinutePerUser": file_upl_exc.QuotaExceededReadRequestsPerMinutePerUser,
            }.get(
                err_info.get("metadata", {}).get("quota_limit"),
                file_upl_exc.QuotaExceeded,
            )
            raise quota_err_cls(debug_info=err_resp_json)
        if err.res.status_code >= 500:
            raise file_upl_exc.RemoteServerError(debug_info=err_resp_json)
        raise err

    async def _make_request(self, request: aiogoogle.models.Request) -> dict:
        await self._require_init()
        assert self._aiogoogle is not None

        # exponential backoff
        maximum_backoff = 64  # maximum number of seconds to wait between retries
        deadline = 220  # maximum number of seconds to keep sending requests
        current_backoff_sum = 0
        current_try = 0
        resp_json: Optional[dict] = None
        while True:
            try:
                if self.auth is not None:
                    # Aiogoogle saves user creds on refresh and reuses them
                    # But if we pass incomplete creds smth will go wrong, so it is better not to interfere
                    if self._aiogoogle.user_creds is None:
                        user_creds = UserCreds(
                            access_token=self.auth.access_token, refresh_token=self.auth.refresh_token
                        )
                    else:
                        user_creds = None
                    resp_json = await self._aiogoogle.as_user(request, user_creds=user_creds)  # type: ignore
                else:
                    resp_json = await self._aiogoogle.as_api_key(request)  # type: ignore
                break
            except aiogoogle.excs.HTTPError as err:
                if self._is_retryable_status(err.res.status_code):
                    LOGGER.error(err)
                    if current_backoff_sum > deadline:
                        self._raise_retryable(err)
                    backoff_seconds = min(2**current_try + random.uniform(0, 1), maximum_backoff)
                    current_backoff_sum += backoff_seconds
                    LOGGER.info(
                        f"Got status {err.res.status_code} on try {current_try},"
                        f" going to back off for {backoff_seconds:.3f} seconds"
                    )
                    current_try += 1
                    await asyncio.sleep(backoff_seconds)
                    continue
                raise

        assert isinstance(resp_json, dict), resp_json
        session = self._aiogoogle.session_context.get()
        assert isinstance(session, AiohttpGSheetsSession)
        self.last_response_size_bytes = session.last_response_size_bytes
        return resp_json

    async def _request_spreadsheet(self, spreadsheet_id: str, include_data: bool, range: Optional[str] = None) -> dict:
        await self._require_init()
        assert self._sheets_api is not None

        req_params = dict(
            spreadsheetId=spreadsheet_id,
            includeGridData=include_data,
        )
        if range is not None:
            req_params["ranges"] = range
            req_params["path_params_safe_chars"] = {"ranges": URL_SAFE_CHARS}

        return await self._make_request(self._sheets_api.spreadsheets.get(**req_params))

    async def _request_values(self, spreadsheet_id: str, range: str) -> dict:
        await self._require_init()
        assert self._sheets_api is not None

        req_params = dict(
            spreadsheetId=spreadsheet_id,
            range=range,
            dateTimeRenderOption="SERIAL_NUMBER",
            majorDimension="ROWS",
            valueRenderOption="UNFORMATTED_VALUE",
            path_params_safe_chars={"range": URL_SAFE_CHARS},
        )

        return await self._make_request(self._sheets_api.spreadsheets.values.get(**req_params))

    async def get_spreadsheet(
        self,
        spreadsheet_id: str,
        include_data: bool = True,
        num_rows: Optional[int] = None,
    ) -> Spreadsheet:
        """
        :param spreadsheet_id: ID of the spreadsheet to get, e.g. 1rnUFa7AiSKD5O80IKCvMy2cSZvLU1kRw9dxbtZbDMWc
        :param include_data: whether to include actual data or just save spreadsheet and sheet properties
        :param num_rows: number of rows to get, all rows if None (default), ignored when included_data is False
        :return: Spreadsheet object with properties and data if it is requested
        """

        resp_json = await self._request_spreadsheet(spreadsheet_id=spreadsheet_id, include_data=include_data)

        sheets = []
        for sheet in resp_json.get("sheets", []):
            sheet_properties = sheet["properties"]
            data: Optional[list[list[Cell]]]
            if include_data:
                data = await self._loop.run_in_executor(self._tpe, self._get_sheet_data, sheet, num_rows)
            else:
                data = None
            sheets.append(
                Sheet(
                    id=sheet_properties["sheetId"],
                    index=sheet_properties["index"],
                    title=sheet_properties["title"],
                    row_count=sheet_properties["gridProperties"]["rowCount"],
                    column_count=sheet_properties["gridProperties"]["columnCount"],
                    data=data,
                )
            )
        return Spreadsheet(
            id=resp_json["spreadsheetId"],
            url=resp_json["spreadsheetUrl"],
            title=resp_json["properties"]["title"],
            sheets=sheets,
        )

    async def get_spreadsheet_sample(self, spreadsheet_id: str, sample_rows: Optional[int] = None) -> Spreadsheet:
        """
        :param spreadsheet_id: spreadsheetId
        :param sample_rows: rows to request from every sheet (for values), default is `2600 / column count * 50` (~4 MB in each response, Google recommends 2 MB)
        :return: Spreadsheet instance
        """

        resp_json = await self._request_spreadsheet(spreadsheet_id=spreadsheet_id, include_data=False)

        sheets = []
        for sheet in resp_json.get("sheets", []):
            sheet_properties = sheet["properties"]
            sheets.append(
                Sheet(
                    id=sheet_properties["sheetId"],
                    index=sheet_properties["index"],
                    title=sheet_properties["title"],
                    row_count=sheet_properties["gridProperties"]["rowCount"],
                    column_count=sheet_properties["gridProperties"]["columnCount"],
                    data=None,
                )
            )
        spreadsheet_meta = Spreadsheet(
            id=resp_json["spreadsheetId"],
            url=resp_json["spreadsheetUrl"],
            title=resp_json["properties"]["title"],
            sheets=sheets,
        )

        for idx, sheet in enumerate(spreadsheet_meta.sheets):
            sample_rows_to_request = 2600 // sheet.column_count if sample_rows is None else sample_rows
            sample_range = Range(
                sheet_title=sheet.title,
                row_from=1,
                col_from=1,
                row_to=sample_rows_to_request,
                col_to=sheet.column_count,
            )
            sheet_sample = await self.get_single_range(
                spreadsheet_id=spreadsheet_id,
                range=sample_range,
            )

            sheet_sample.batch_size_rows = sample_rows_to_request * 50  # for values
            sheet_sample.column_count = len(sheet_sample.data[0]) if sheet_sample.data else 0

            spreadsheet_meta.sheets[idx] = sheet_sample

        return spreadsheet_meta

    async def get_single_range(self, spreadsheet_id: str, range: Range) -> Sheet:
        resp_json = await self._request_spreadsheet(spreadsheet_id=spreadsheet_id, include_data=True, range=str(range))

        sheet_json = resp_json["sheets"][0]
        sheet_properties = sheet_json["properties"]
        sheet_data = await self._loop.run_in_executor(self._tpe, self._get_sheet_data, sheet_json)
        sheet = Sheet(
            id=sheet_properties["sheetId"],
            index=sheet_properties["index"],
            title=sheet_properties["title"],
            row_count=sheet_properties["gridProperties"]["rowCount"],
            column_count=sheet_properties["gridProperties"]["columnCount"],
            data=sheet_data,
        )
        return sheet

    async def get_single_values_range(
        self, spreadsheet_id: str, range: Range, user_types: list[UserDataType | str]
    ) -> Tuple[list[list[Any]], list[UserDataType | str]]:
        resp_json = await self._request_values(spreadsheet_id=spreadsheet_id, range=str(range))

        raw_values = resp_json.get("values", [])
        values, updated_user_types = await self._loop.run_in_executor(
            self._tpe,
            self._process_values,
            raw_values,
            user_types,
        )

        return values, updated_user_types

lib/dl_file_uploader_lib/dl_file_uploader_lib/gsheets_client.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from datetime import timedelta
from functools import reduce
import json
import logging
from typing import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    ClassVar,
    Optional,
    Type,
    TypeVar,
    Union,
)
import uuid

import attr
import marshmallow as ma
import redis.asyncio

from dl_api_commons.base_models import RequestContextInfo
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_core.us_manager.crypto.main import CryptoController
from dl_utils.utils import (
    AddressableData,
    DataKey,
)


LOGGER = logging.getLogger(__name__)


class RedisModelException(Exception):
    pass


class RedisModelAccessDenied(RedisModelException):
    pass


class RedisModelNotFound(RedisModelException):
    pass


class SecretContainingMixin:
    def get_secret_keys(self) -> set[DataKey]:
        return set()


TRedisModel = TypeVar("TRedisModel", bound="RedisModel")


@attr.s(init=True, kw_only=True)
class RedisModel(SecretContainingMixin, metaclass=abc.ABCMeta):
    KEY_PREFIX: ClassVar[str]
    DEFAULT_TTL: ClassVar[Optional[int]] = None  # in seconds
    _manager: Optional[RedisModelManager] = attr.ib(default=None)

    id: str = attr.ib(factory=lambda: str(uuid.uuid4()))

    @classmethod
    def _generate_key_by_id(cls, obj_id: str) -> str:
        assert cls.KEY_PREFIX
        return "/".join((cls.KEY_PREFIX, obj_id))

    def generate_key(self) -> str:
        return self._generate_key_by_id(self.id)

    @classmethod
    async def get(cls: Type[TRedisModel], manager: RedisModelManager, obj_id: str) -> TRedisModel:
        key = cls._generate_key_by_id(obj_id)
        return await manager.get(key=key, target_cls=cls)  # type: ignore

    async def save(self, ttl: Union[int, timedelta, None] = None) -> None:
        assert self._manager
        await self._manager.save(self, ttl=ttl)

    async def delete(self) -> None:
        assert self._manager
        await self._manager.delete(self)


@attr.s(init=True, kw_only=True)
class RedisModelUserIdAuth(RedisModel, metaclass=abc.ABCMeta):
    system_user_id: ClassVar[str] = "system"

    user_id: Optional[str] = attr.ib(default=None)

    def __attrs_post_init__(self) -> None:
        if self.user_id is None:
            assert self._manager is not None and self._manager.rci is not None
            user_id = self._manager.rci.user_id
            assert user_id is not None
            self.user_id = user_id

    @classmethod
    async def get_authorized(cls: Type[TRedisModel], manager: RedisModelManager, obj_id: str) -> TRedisModel:
        """Load entry with authorization check by user id"""
        key = cls._generate_key_by_id(obj_id)
        obj = await manager.get(key=key, target_cls=cls, post_load=_check_obj_auth)
        assert isinstance(obj, RedisModelUserIdAuth)
        return obj  # type: ignore


def _check_obj_auth(obj: RedisModelUserIdAuth) -> None:
    if obj.user_id == obj.system_user_id:
        LOGGER.info("Attempt to get and object owned by system with authorization")
        raise RedisModelAccessDenied()

    manager = obj._manager  # noqa
    assert manager is not None and manager.rci is not None
    current_user_id = manager.rci.user_id

    if current_user_id is None:
        LOGGER.info("Current user_id not set. Unable to check authorization,")
        raise RedisModelAccessDenied()
    if obj.user_id != current_user_id:
        LOGGER.info(f"Current user_id and loaded object user_id don't match: {obj.user_id} != {current_user_id}")
        raise RedisModelAccessDenied()


@attr.s(init=True, kw_only=True)
class RedisSetManager(metaclass=abc.ABCMeta):
    KEY_PREFIX: ClassVar[str]
    _redis: redis.asyncio.Redis = attr.ib()

    id: str = attr.ib()

    @classmethod
    def _generate_key_by_id(cls, obj_id: str) -> str:
        return "/".join((cls.KEY_PREFIX, obj_id))

    @property
    def key(self) -> str:
        return self._generate_key_by_id(self.id)

    async def add(self, value: str) -> None:
        cmd = self._redis.sadd(self.key, value)
        assert isinstance(cmd, Awaitable)
        await cmd

    async def _sscan_iter_str(self) -> AsyncIterator:
        async for value in self._redis.sscan_iter(name=self.key):
            yield value.decode("utf-8")

    def sscan_iter(self) -> AsyncIterator:
        """
        Note: according to redis docs (https://redis.io/commands/scan/#scan-guarantees),
        may return some entries more than once, the client should handle this behaviour on its side.
        """

        return self._sscan_iter_str()

    async def delete(self) -> None:
        cmd = self._redis.delete(self.key)
        assert isinstance(cmd, Awaitable)
        await cmd

    async def size(self) -> int:
        cmd = self._redis.scard(self.key)
        assert isinstance(cmd, Awaitable)
        return await cmd

    async def sunion_by_id(self, *ids: str) -> None:
        if not ids:
            raise ValueError("ids can't be empty")
        keys = [self._generate_key_by_id(_id) for _id in ids]
        cmd = self._redis.sunionstore(self.key, keys)
        assert isinstance(cmd, Awaitable)
        await cmd


@attr.s
class RedisModelManager:
    _redis: redis.asyncio.Redis = attr.ib()
    _crypto_keys_config: CryptoKeysConfig = attr.ib()
    _crypto_controller: CryptoController = attr.ib(
        init=False,
        default=attr.Factory(
            lambda self: CryptoController(self._crypto_keys_config),
            takes_self=True,
        ),
    )
    rci: Optional[RequestContextInfo] = attr.ib(default=None)

    def _set_obj_attr(self, obj: RedisModel, key: DataKey, value: Any) -> None:
        key_head = DataKey(parts=key.parts[:-1])
        obj = self._get_obj_attr(obj, key_head)
        obj.__setattr__(key.parts[-1], value)

    def _get_obj_attr(self, obj: RedisModel, key: DataKey) -> Any:
        return reduce(lambda _obj, attribute: _obj.__getattribute__(attribute), key.parts, obj)

    def _serialize_object(self, obj: RedisModel) -> str:
        schema_cls = _get_model_schema_class(type(obj))
        schema = schema_cls()
        json_data_dict = schema.dump(obj)

        addressable_data = AddressableData(json_data_dict)
        secret_keys = obj.get_secret_keys()
        for key in secret_keys:
            plain_text = addressable_data.get(key)
            encrypted = self._crypto_controller.encrypt_with_actual_key(plain_text)
            addressable_data.set(key, json.dumps(encrypted))

        json_data = json.dumps(json_data_dict)
        return json_data

    def _deserialize_object(self, json_data: str, target_cls: Type[RedisModel]) -> RedisModel:
        schema_cls = _get_model_schema_class(target_cls)
        schema = schema_cls()
        obj = schema.loads(json_data)

        assert isinstance(obj, RedisModel)
        secret_keys = obj.get_secret_keys()
        for key in secret_keys:
            encrypted = json.loads(self._get_obj_attr(obj, key))
            decrypted = self._crypto_controller.decrypt(encrypted)
            self._set_obj_attr(obj, key, decrypted)

        return obj

    async def get(self, key: str, target_cls: Type[RedisModel], post_load: Optional[Callable] = None) -> RedisModel:
        json_data = await self._redis.get(key)
        if json_data is None:
            LOGGER.info(f"RedisModel object not found: {key}")
            raise RedisModelNotFound()

        obj = self._deserialize_object(json_data, target_cls)
        obj._manager = self  # type: ignore

        if post_load is not None:
            post_load(obj)

        return obj

    async def save(self, obj: RedisModel, ttl: Union[int, timedelta, None] = None) -> None:
        obj_key = obj.generate_key()
        data_json = self._serialize_object(obj)

        if ttl is None and obj.DEFAULT_TTL is not None:
            ttl = obj.DEFAULT_TTL

        if ttl:
            await self._redis.setex(obj_key, ttl, data_json)
        else:
            await self._redis.set(obj_key, data_json)

    async def delete(self, obj: RedisModel) -> None:
        obj_key = obj.generate_key()
        await self._redis.delete(obj_key)


class BaseSchema(ma.Schema):
    class Meta:
        unknown = ma.EXCLUDE

        target: Type

    @ma.post_load(pass_many=False)
    def to_object(self, data: dict[str, Any], **kwargs: Any) -> Any:
        return self.Meta.target(**data)


class BaseModelSchema(BaseSchema):
    id = ma.fields.String()
    user_id = ma.fields.String()


_MODEL_CLASS_TO_STORAGE_SCHEMA_MAP: dict[Type[RedisModel], Type[BaseModelSchema]] = dict()


def _get_model_schema_class(model_cls: Type[RedisModel]) -> Type[BaseModelSchema]:
    return _MODEL_CLASS_TO_STORAGE_SCHEMA_MAP[model_cls]


def register_redis_model_storage_schema(model_cls: Type[RedisModel], schema_cls: Type[BaseModelSchema]) -> None:
    _MODEL_CLASS_TO_STORAGE_SCHEMA_MAP[model_cls] = schema_cls

lib/dl_file_uploader_lib/dl_file_uploader_lib/redis_model/base.py
datalens-tech/datalens-backend
from dl_file_uploader_lib.redis_model.base import register_redis_model_storage_schema

from .models import (
    CSVFileSettings,
    CSVFileSourceSettings,
    DataFile,
    DataFileError,
    DataSource,
    DataSourcePreview,
    EmptySourcesError,
    ExcelFileSourceSettings,
    FileProcessingError,
    FileSettings,
    FileSourceSettings,
    GSheetsFileSourceSettings,
    GSheetsUserSourceDataSourceProperties,
    GSheetsUserSourceProperties,
    PreviewSet,
    RenameTenantStatusModel,
    SourceNotFoundError,
    SpreadsheetFileSourceSettings,
    YaDocsUserSourceProperties,
)
from .storage_schemas import (
    DataFileSchema,
    DataSourcePreviewSchema,
    RenameTenantStatusModelSchema,
)


__all__ = (
    "DataFile",
    "FileSettings",
    "FileSourceSettings",
    "CSVFileSettings",
    "CSVFileSourceSettings",
    "GSheetsFileSourceSettings",
    "ExcelFileSourceSettings",
    "SpreadsheetFileSourceSettings",
    "DataSource",
    "DataSourcePreview",
    "GSheetsUserSourceProperties",
    "GSheetsUserSourceDataSourceProperties",
    "FileProcessingError",
    "DataFileError",
    "SourceNotFoundError",
    "EmptySourcesError",
    "RenameTenantStatusModel",
    "PreviewSet",
    "YaDocsUserSourceProperties",
)


register_redis_model_storage_schema(DataFile, DataFileSchema)
register_redis_model_storage_schema(DataSourcePreview, DataSourcePreviewSchema)
register_redis_model_storage_schema(RenameTenantStatusModel, RenameTenantStatusModelSchema)

lib/dl_file_uploader_lib/dl_file_uploader_lib/redis_model/models/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import csv
import time
from typing import (
    ClassVar,
    Optional,
)
import uuid

import attr

from dl_constants.enums import FileProcessingStatus
from dl_core.db.elements import SchemaColumn
from dl_file_uploader_lib.enums import (
    CSVDelimiter,
    CSVEncoding,
    ErrorLevel,
    FileType,
    RenameTenantStatus,
)
from dl_file_uploader_lib.exc import DLFileUploaderBaseError
from dl_file_uploader_lib.redis_model.base import (
    RedisModel,
    RedisModelUserIdAuth,
    RedisSetManager,
    SecretContainingMixin,
)
from dl_utils.utils import DataKey


class DataFileError(Exception):
    pass


class SourceNotFoundError(DataFileError):
    pass


class EmptySourcesError(DataFileError):
    pass


@attr.s(init=True, kw_only=True)
class FileSettings:
    file_type: FileType = attr.ib()


@attr.s(init=True, kw_only=True)
class CSVFileSettings(FileSettings):
    file_type: FileType = attr.ib(default=FileType.csv)

    encoding: CSVEncoding = attr.ib()
    dialect: csv.Dialect = attr.ib()

    @property
    def delimiter(self) -> CSVDelimiter:
        return CSVDelimiter(self.dialect.delimiter)


@attr.s(init=True, kw_only=True)
class FileSourceSettings:
    file_type: FileType = attr.ib()

    first_line_is_header: bool = attr.ib()


@attr.s(init=True, kw_only=True)
class CSVFileSourceSettings(FileSourceSettings):
    file_type: FileType = attr.ib(default=FileType.csv)


@attr.s(init=True, kw_only=True)
class SpreadsheetFileSourceSettings(FileSourceSettings):
    raw_schema_header: list[SchemaColumn] = attr.ib()
    raw_schema_body: list[SchemaColumn] = attr.ib()


@attr.s(init=True, kw_only=True)
class GSheetsFileSourceSettings(SpreadsheetFileSourceSettings):
    file_type: FileType = attr.ib(default=FileType.gsheets)


@attr.s(init=True, kw_only=True)
class ExcelFileSourceSettings(SpreadsheetFileSourceSettings):
    file_type: FileType = attr.ib(default=FileType.xlsx)


@attr.s(init=True, kw_only=True)
class UserSourceProperties(SecretContainingMixin):
    file_type: FileType = attr.ib()


@attr.s(init=True, kw_only=True)
class GSheetsUserSourceProperties(UserSourceProperties):
    file_type: FileType = attr.ib(default=FileType.gsheets)

    spreadsheet_id: str = attr.ib()
    refresh_token: Optional[str] = attr.ib(default=None, repr=False)

    def get_secret_keys(self) -> set[DataKey]:
        return {DataKey(parts=("refresh_token",))}


@attr.s(init=True, kw_only=True)
class YaDocsUserSourceProperties(UserSourceProperties):
    file_type: FileType = attr.ib(default=FileType.yadocs)

    private_path: Optional[str] = attr.ib(default=None)
    public_link: Optional[str] = attr.ib(default=None)

    oauth_token: Optional[str] = attr.ib(default=None, repr=False)

    def get_secret_keys(self) -> set[DataKey]:
        return {DataKey(parts=("oauth_token",))}


@attr.s(init=True, kw_only=True)
class UserSourceDataSourceProperties:
    file_type: FileType = attr.ib()


@attr.s(init=True, kw_only=True)
class GSheetsUserSourceDataSourceProperties(UserSourceDataSourceProperties):
    file_type: FileType = attr.ib(default=FileType.gsheets)

    sheet_id: int = attr.ib()


@attr.s(init=True, kw_only=True)
class ParsingError:
    code: str = attr.ib()


@attr.s
class FileProcessingError:
    level: ErrorLevel = attr.ib()
    message: str = attr.ib()
    code: list[str] = attr.ib(factory=list)
    details: dict = attr.ib(factory=dict)

    @classmethod
    def from_exception(cls, exc: DLFileUploaderBaseError, level: Optional[ErrorLevel] = None) -> FileProcessingError:
        return cls(
            level=level or exc.default_level,
            message=exc.message,
            code=exc.err_code,
            details=exc.details,
        )


@attr.s(init=True, kw_only=True)
class DataSource:
    id: str = attr.ib(factory=lambda: str(uuid.uuid4()))
    s3_key: str = attr.ib(
        default=attr.Factory(
            lambda self: self.id + "_" + str(uuid.uuid4())[:8],
            takes_self=True,
        )
    )
    preview_id: Optional[str] = attr.ib(default=None)
    title: str = attr.ib()
    raw_schema: list[SchemaColumn] = attr.ib()
    file_source_settings: Optional[FileSourceSettings] = attr.ib(default=None)
    user_source_dsrc_properties: Optional[UserSourceDataSourceProperties] = attr.ib(default=None)
    status: FileProcessingStatus = attr.ib()
    error: Optional[FileProcessingError] = attr.ib(default=None)

    @property
    def is_applicable(self) -> bool:
        return not (self.error is not None or self.status == FileProcessingStatus.failed)


@attr.s(init=True, kw_only=True)
class DataFile(RedisModelUserIdAuth):
    filename: Optional[str] = attr.ib()
    file_type: Optional[FileType] = attr.ib(default=None)
    file_settings: Optional[FileSettings] = attr.ib(default=None)
    user_source_properties: Optional[UserSourceProperties] = attr.ib(default=None)
    size: Optional[int] = attr.ib(default=None)
    sources: Optional[list[DataSource]] = attr.ib(default=None)
    status: FileProcessingStatus = attr.ib()
    error: Optional[FileProcessingError] = attr.ib(default=None)

    KEY_PREFIX: ClassVar[str] = "df"
    DEFAULT_TTL: ClassVar[Optional[int]] = 12 * 60 * 60  # 12 hours

    @property
    def s3_key(self) -> str:
        return self.id

    def get_secret_keys(self) -> set[DataKey]:
        if self.user_source_properties is None:
            return set()

        lower_level_keys = self.user_source_properties.get_secret_keys()
        return {DataKey(parts=("user_source_properties",) + ll_key.parts) for ll_key in lower_level_keys}

    def get_source_by_id(self, source_id: str) -> DataSource:
        if not self.sources:
            raise EmptySourcesError("There are no sources for the file.")
        for src in self.sources:
            if src.id == source_id:
                return src
        raise SourceNotFoundError(f"Source with id {source_id} has not been found.")


@attr.s(init=True, kw_only=True)
class DataSourcePreview(RedisModel):
    preview_data: list[list[Optional[str]]] = attr.ib()

    KEY_PREFIX: ClassVar[str] = "df_preview"


@attr.s(init=True, kw_only=True)
class RenameTenantStatusModel(RedisModel):  # object id is tenant_id
    status: RenameTenantStatus = attr.ib()
    mtime: float = attr.ib(factory=time.time)

    KEY_PREFIX = "rename_tenant_status"


@attr.s(init=True, kw_only=True)
class PreviewSet(RedisSetManager):
    KEY_PREFIX = "tenant_previews"

lib/dl_file_uploader_lib/dl_file_uploader_lib/redis_model/models/models.py
datalens-tech/datalens-backend
import csv
from typing import (
    Any,
    Type,
)

from marshmallow import (
    EXCLUDE,
    fields,
    post_load,
)
from marshmallow_oneofschema import OneOfSchema

from dl_constants.enums import FileProcessingStatus
from dl_core.us_manager.storage_schemas.base_types import SchemaColumnStorageSchema
from dl_file_uploader_lib.enums import (
    CSVEncoding,
    ErrorLevel,
    FileType,
    RenameTenantStatus,
)
from dl_file_uploader_lib.redis_model.base import (
    BaseModelSchema,
    BaseSchema,
)
from dl_file_uploader_lib.redis_model.models import (
    CSVFileSettings,
    CSVFileSourceSettings,
    DataFile,
    DataSource,
    DataSourcePreview,
    ExcelFileSourceSettings,
    FileProcessingError,
    GSheetsFileSourceSettings,
    GSheetsUserSourceDataSourceProperties,
    GSheetsUserSourceProperties,
    RenameTenantStatusModel,
    YaDocsUserSourceProperties,
)


class CSVDialectSchema(BaseSchema):
    @post_load(pass_many=False)
    def to_object(self, data: dict[str, Any], **kwargs: Any) -> Any:
        class _Dialect(csv.Dialect):
            def __init__(self) -> None:
                for attr_name, val in data.items():
                    setattr(self, attr_name, val)
                super().__init__()

        return _Dialect()

    delimiter = fields.String()
    doublequote = fields.Boolean()
    escapechar = fields.String(allow_none=True)
    lineterminator = fields.String()
    quotechar = fields.String(allow_none=True)
    quoting = fields.Integer()
    skipinitialspace = fields.Boolean()


class FileProcessingErrorSchema(BaseSchema):
    class Meta(BaseSchema.Meta):
        target = FileProcessingError

    level = fields.Enum(ErrorLevel)
    message = fields.String()
    code = fields.List(fields.String())
    details = fields.Dict()


class FileTypeOneOfSchema(OneOfSchema):
    class Meta:
        unknown = EXCLUDE

    type_field_remove = False
    type_field = "file_type"

    def get_obj_type(self, obj: dict[str, Any]) -> str:
        return getattr(obj, self.type_field).name


class FileSettingsBaseSchema(BaseSchema):
    file_type = fields.Enum(FileType)


class CSVFileSettingsSchema(FileSettingsBaseSchema):
    class Meta(BaseSchema.Meta):
        target = CSVFileSettings

    encoding = fields.Enum(CSVEncoding)
    dialect = fields.Nested(CSVDialectSchema)


class ExcelFileSettingsSchema(FileSettingsBaseSchema):
    pass


class FileSettingsSchema(FileTypeOneOfSchema):
    type_schemas: dict[str, Type[FileSettingsBaseSchema]] = {
        FileType.csv.name: CSVFileSettingsSchema,
        FileType.xlsx.name: ExcelFileSettingsSchema,
    }


class FileSourceSettingsBaseSchema(BaseSchema):
    file_type = fields.Enum(FileType)


class CSVFileSourceSettingsSchema(FileSourceSettingsBaseSchema):
    class Meta(BaseSchema.Meta):
        target = CSVFileSourceSettings

    first_line_is_header = fields.Boolean()


class GSheetsFileSourceSettingsSchema(FileSourceSettingsBaseSchema):
    class Meta(BaseSchema.Meta):
        target = GSheetsFileSourceSettings

    first_line_is_header = fields.Boolean()
    raw_schema_header = fields.Nested(SchemaColumnStorageSchema, many=True)
    raw_schema_body = fields.Nested(SchemaColumnStorageSchema, many=True)


class ExcelFileSourceSettingsSchema(FileSourceSettingsBaseSchema):
    class Meta(BaseSchema.Meta):
        target = ExcelFileSourceSettings

    first_line_is_header = fields.Boolean()
    raw_schema_header = fields.Nested(SchemaColumnStorageSchema, many=True)
    raw_schema_body = fields.Nested(SchemaColumnStorageSchema, many=True)


class FileSourceSettingsSchema(FileTypeOneOfSchema):
    type_schemas: dict[str, Type[FileSourceSettingsBaseSchema]] = {
        FileType.csv.name: CSVFileSourceSettingsSchema,
        FileType.gsheets.name: GSheetsFileSourceSettingsSchema,
        FileType.xlsx.name: ExcelFileSourceSettingsSchema,
    }


class UserSourcePropertiesBaseSchema(BaseSchema):
    file_type = fields.Enum(FileType)


class GSheetsUserSourcePropertiesSchema(UserSourcePropertiesBaseSchema):
    class Meta(BaseSchema.Meta):
        target = GSheetsUserSourceProperties

    refresh_token = fields.String(allow_none=True)
    spreadsheet_id = fields.String()


class YaDocsUserSourcePropertiesSchema(UserSourcePropertiesBaseSchema):
    class Meta(BaseSchema.Meta):
        target = YaDocsUserSourceProperties

    private_path = fields.String(allow_none=True)
    public_link = fields.String(allow_none=True)
    oauth_token = fields.String(allow_none=True)


class UserSourcePropertiesSchema(FileTypeOneOfSchema):
    type_schemas: dict[str, Type[UserSourcePropertiesBaseSchema]] = {
        FileType.gsheets.name: GSheetsUserSourcePropertiesSchema,
        FileType.yadocs.name: YaDocsUserSourcePropertiesSchema,
    }


class UserSourceDataSourcePropertiesBaseSchema(BaseSchema):
    file_type = fields.Enum(FileType)


class GSheetsUserSourceDataSourcePropertiesSchema(UserSourceDataSourcePropertiesBaseSchema):
    class Meta(BaseSchema.Meta):
        target = GSheetsUserSourceDataSourceProperties

    sheet_id = fields.Integer()


class UserSourceDataSourcePropertiesSchema(FileTypeOneOfSchema):
    type_schemas: dict[str, Type[UserSourceDataSourcePropertiesBaseSchema]] = {
        FileType.gsheets.name: GSheetsUserSourceDataSourcePropertiesSchema,
    }


class DataSourceSchema(BaseSchema):
    class Meta(BaseSchema.Meta):
        target = DataSource

    id = fields.String()
    s3_key = fields.String()
    preview_id = fields.String(allow_none=True)
    title = fields.String()
    raw_schema = fields.Nested(SchemaColumnStorageSchema, many=True)
    file_source_settings = fields.Nested(FileSourceSettingsSchema, allow_none=True)
    user_source_dsrc_properties = fields.Nested(UserSourceDataSourcePropertiesSchema, allow_none=True)
    status = fields.Enum(FileProcessingStatus)
    error = fields.Nested(FileProcessingErrorSchema, allow_none=True)


class DataFileSchema(BaseModelSchema):
    class Meta(BaseModelSchema.Meta):
        target = DataFile

    filename = fields.String()
    file_type = fields.Enum(FileType, allow_none=True)
    file_settings = fields.Nested(FileSettingsSchema, allow_none=True)
    user_source_properties = fields.Nested(UserSourcePropertiesSchema, allow_none=True)
    size = fields.Integer(allow_none=True)
    sources = fields.Nested(DataSourceSchema, many=True, allow_none=True)
    status = fields.Enum(FileProcessingStatus)
    error = fields.Nested(FileProcessingErrorSchema, allow_none=True)


class DataSourcePreviewSchema(BaseModelSchema):
    class Meta(BaseModelSchema.Meta):
        target = DataSourcePreview

    preview_data = fields.List(fields.List(fields.Raw(allow_none=True)))


class RenameTenantStatusModelSchema(BaseModelSchema):
    class Meta(BaseModelSchema.Meta):
        target = RenameTenantStatusModel

    tenant_id = fields.String()
    status = fields.Enum(RenameTenantStatus)
    mtime = fields.Float()

lib/dl_file_uploader_lib/dl_file_uploader_lib/redis_model/models/storage_schemas.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Optional,
)

import attr

from dl_configs.crypto_keys import CryptoKeysConfig
from dl_configs.enums import RedisMode
from dl_configs.environments import is_setting_applicable
from dl_configs.settings_loaders.meta_definition import (
    required,
    s_attrib,
)
from dl_configs.settings_submodels import (
    RedisSettings,
    S3Settings,
)


def _make_redis_persistent_settings(cfg: Any, db: int) -> Optional[RedisSettings]:
    # TODO: move this values to a separate key
    return (
        RedisSettings(  # type: ignore
            MODE=RedisMode(cfg.REDIS_PERSISTENT_MODE),  # type: ignore
            CLUSTER_NAME=cfg.REDIS_PERSISTENT_CLUSTER_NAME,  # type: ignore
            HOSTS=cfg.REDIS_PERSISTENT_HOSTS,  # type: ignore
            PORT=cfg.REDIS_PERSISTENT_PORT,  # type: ignore
            SSL=cfg.REDIS_PERSISTENT_SSL,  # type: ignore
            PASSWORD=required(str),
            DB=db,
        )
        if is_setting_applicable(cfg, "REDIS_PERSISTENT_HOSTS")
        else None
    )


@attr.s(frozen=True)
class FileUploaderBaseSettings:
    REDIS_APP: RedisSettings = s_attrib(  # type: ignore
        "REDIS_APP",
        fallback_factory=(lambda cfg: _make_redis_persistent_settings(cfg=cfg, db=cfg.REDIS_FILE_UPLOADER_DATA_DB)),
    )

    REDIS_ARQ: RedisSettings = s_attrib(  # type: ignore
        "REDIS_ARQ",
        fallback_factory=(lambda cfg: _make_redis_persistent_settings(cfg=cfg, db=cfg.REDIS_FILE_UPLOADER_TASKS_DB)),
    )

    S3: S3Settings = s_attrib(  # type: ignore
        "S3",
        fallback_factory=(
            lambda cfg: S3Settings(  # type: ignore
                ACCESS_KEY_ID=required(str),
                SECRET_ACCESS_KEY=required(str),
                ENDPOINT_URL=cfg.S3_ENDPOINT_URL,
            )
            if is_setting_applicable(cfg, "S3_ENDPOINT_URL")
            else None
        ),
    )
    S3_TMP_BUCKET_NAME: str = s_attrib(  # type: ignore
        "S3_TMP_BUCKET_NAME",
        fallback_cfg_key="FILE_UPLOADER_S3_TMP_BUCKET_NAME",
    )
    S3_PERSISTENT_BUCKET_NAME: str = s_attrib(  # type: ignore
        "S3_PERSISTENT_BUCKET_NAME",
        fallback_cfg_key="FILE_UPLOADER_S3_PERSISTENT_BUCKET_NAME",
    )

    CRYPTO_KEYS_CONFIG: CryptoKeysConfig = s_attrib(  # type: ignore
        "DL_CRY",
        json_converter=CryptoKeysConfig.from_json,
        sensitive=True,
    )

lib/dl_file_uploader_lib/dl_file_uploader_lib/settings.py
datalens-tech/datalens-backend
from dl_configs.enums import RedisMode
from dl_constants.enums import RedisInstanceKind
from dl_core.aio.web_app_services.redis import (
    RedisBaseService,
    RedisSentinelService,
    SingleHostSimpleRedisService,
)
from dl_core.utils import make_url
from dl_file_uploader_lib.settings import FileUploaderBaseSettings


def init_redis_service(settings: FileUploaderBaseSettings) -> RedisBaseService:
    redis_service: RedisBaseService
    if settings.REDIS_APP.MODE == RedisMode.single_host:
        assert len(settings.REDIS_APP.HOSTS) == 1
        redis_service = SingleHostSimpleRedisService(
            url=make_url(
                protocol="rediss" if settings.REDIS_APP.SSL else "redis",
                host=settings.REDIS_APP.HOSTS[0],
                port=settings.REDIS_APP.PORT,
                path=str(settings.REDIS_APP.DB),
            ),
            password=settings.REDIS_APP.PASSWORD,
            instance_kind=RedisInstanceKind.persistent,
            ssl=settings.REDIS_APP.SSL,
        )
    elif settings.REDIS_APP.MODE == RedisMode.sentinel:
        redis_service = RedisSentinelService(
            sentinel_hosts=settings.REDIS_APP.HOSTS,
            sentinel_port=settings.REDIS_APP.PORT,
            namespace=settings.REDIS_APP.CLUSTER_NAME,
            password=settings.REDIS_APP.PASSWORD,
            db=settings.REDIS_APP.DB,
            instance_kind=RedisInstanceKind.persistent,
            ssl=settings.REDIS_APP.SSL,
        )
    else:
        raise ValueError(f"Unknown redis mode {settings.REDIS_APP.MODE}")
    return redis_service

lib/dl_file_uploader_lib/dl_file_uploader_lib/settings_utils.py
datalens-tech/datalens-backend
import csv
import datetime
import io
import random
import string


def generate_sample_csv_data_str(row_count: int = 10000, str_cols_count: int = 3) -> str:
    output = io.StringIO()
    csv_writer = csv.writer(output, dialect=csv.excel)
    csv_writer.writerow(
        ["some int", "some float", "some date", "some datetime"] + [f"some string {i}" for i in range(str_cols_count)]
    )
    today = datetime.date.today()
    now = datetime.datetime.now()
    printable = string.printable + "АаБбВвГгДдЕеЁёЖжЗзИиЙйКкЛлМмНнОоПпРрСсТтУуФфХхЦцЧчШшЩщЪъЫыЬьЭэЮюЯя"
    for _i in range(row_count):
        csv_writer.writerow(
            [
                random.randint(-100, 100),
                random.randint(1, 1000) * 1.0 / random.randint(1, 1000),
                (today - datetime.timedelta(days=random.randint(0, 100))).isoformat(),
                (now - datetime.timedelta(days=random.randint(0, 100), minutes=random.randint(0, 10000))).isoformat(),
            ]
            + ["".join(random.choices(printable, k=random.randint(25, 35)))] * str_cols_count
        )
    return output.getvalue()

lib/dl_file_uploader_lib/dl_file_uploader_lib/testing/data_gen.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Type,
)

from aiohttp.client import ClientSession

from dl_file_uploader_lib import exc as file_upl_exc


def yadocs_error_to_file_uploader_exception(status_code: int, resp_info: dict) -> file_upl_exc.DLFileUploaderBaseError:
    err_cls: Type[file_upl_exc.DLFileUploaderBaseError]

    if status_code == 401:
        err_cls = file_upl_exc.PermissionDenied
    elif status_code == 404:
        err_cls = file_upl_exc.DocumentNotFound
    elif status_code == 400:
        err_cls = file_upl_exc.UnsupportedDocument
    elif status_code >= 500:
        err_cls = file_upl_exc.RemoteServerError
    else:
        err_cls = file_upl_exc.DLFileUploaderBaseError

    return err_cls(
        details=resp_info,
    )


class YaDocsClient:
    headers: dict[str, Any] = {
        "Content-Type": "application/json",
        "Accept": "application/json",
    }
    hostname: str = "https://cloud-api.yandex.net/v1/disk"

    def __init__(self, session: ClientSession):
        self.session = session

    async def get_spreadsheet_public_ref(self, link: str) -> str:
        resp = await self.session.get(
            f"{self.hostname}/public/resources/download/?public_key={link}",
            headers=self.headers,
        )
        if resp.status != 200:
            raise yadocs_error_to_file_uploader_exception(resp.status, await resp.json())
        return (await resp.json())["href"]

    async def get_spreadsheet_public_meta(self, link: str) -> dict[str, Any]:
        resp = await self.session.get(
            f"{self.hostname}/public/resources/?public_key={link}",
            headers=self.headers,
        )
        if resp.status != 200:
            raise yadocs_error_to_file_uploader_exception(resp.status, await resp.json())
        return await resp.json()

    async def get_spreadsheet_private_ref(self, path: str, token: str) -> str:
        headers_with_token = self._create_headers_with_token(token)
        resp = await self.session.get(
            f"{self.hostname}/resources/download/?path={path}",
            headers=headers_with_token,
        )
        if resp.status != 200:
            raise yadocs_error_to_file_uploader_exception(resp.status, await resp.json())
        return (await resp.json())["href"]

    async def get_spreadsheet_private_meta(self, path: str, token: str) -> dict[str, Any]:
        headers_with_token = self._create_headers_with_token(token)
        resp = await self.session.get(
            f"{self.hostname}/resources/?path={path}",
            headers=headers_with_token,
        )
        if resp.status != 200:
            raise yadocs_error_to_file_uploader_exception(resp.status, await resp.json())
        return await resp.json()

    def _create_headers_with_token(self, token: str) -> dict[str, Any]:
        headers_with_token = self.headers.copy()
        headers_with_token.update({"Authorization": "OAuth " + token})
        return headers_with_token

lib/dl_file_uploader_lib/dl_file_uploader_lib/yadocs_client.py
datalens-tech/datalens-backend
import pytest
import redis.asyncio

from dl_api_commons.base_models import RequestContextInfo
from dl_configs.crypto_keys import (
    CryptoKeysConfig,
    get_dummy_crypto_keys_config,
)
from dl_configs.enums import RedisMode
from dl_configs.settings_submodels import RedisSettings
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_testing.containers import get_test_container_hostport


pytest_plugins = ("aiohttp.pytest_plugin",)


@pytest.fixture(scope="function")
def rci() -> RequestContextInfo:
    return RequestContextInfo(user_id="_the_tests_asyncapp_user_id_")


@pytest.fixture(scope="session")
def crypto_keys_config() -> CryptoKeysConfig:
    return get_dummy_crypto_keys_config()


@pytest.fixture(scope="session")
def redis_app_settings() -> RedisSettings:
    return RedisSettings(
        MODE=RedisMode.single_host,
        CLUSTER_NAME="",
        HOSTS=(get_test_container_hostport("redis", fallback_port=52404).host,),
        PORT=get_test_container_hostport("redis", fallback_port=52404).port,
        DB=9,
    )


@pytest.fixture(scope="function")
def redis_cli(redis_app_settings) -> redis.asyncio.Redis:
    return redis.asyncio.Redis(
        host=redis_app_settings.HOSTS[0],
        port=redis_app_settings.PORT,
        db=redis_app_settings.DB,
        password=redis_app_settings.PASSWORD,
    )


@pytest.fixture(scope="function")
def redis_model_manager(redis_cli, rci, crypto_keys_config) -> RedisModelManager:
    return RedisModelManager(redis=redis_cli, rci=rci, crypto_keys_config=crypto_keys_config)

lib/dl_file_uploader_lib/dl_file_uploader_lib_tests/conftest.py
datalens-tech/datalens-backend
from enum import Enum
import json
import uuid

import attr
from marshmallow import fields
import pytest

from dl_api_commons.base_models import RequestContextInfo
from dl_configs.crypto_keys import get_dummy_crypto_keys_config
from dl_file_uploader_lib.redis_model.base import (
    BaseModelSchema,
    BaseSchema,
    RedisModel,
    RedisModelAccessDenied,
    RedisModelManager,
    RedisModelNotFound,
    RedisModelUserIdAuth,
    RedisSetManager,
    SecretContainingMixin,
    register_redis_model_storage_schema,
)
from dl_utils.utils import DataKey


class SomeEnum(Enum):
    val1 = "val1"
    val2 = "val2"


@attr.s
class SomeSubModel:
    title: str = attr.ib()
    value: int = attr.ib()


@attr.s
class SomeSubModelWithSecret(SecretContainingMixin):
    secret: str = attr.ib()

    def get_secret_keys(self) -> set[DataKey]:
        return {DataKey(parts=("secret",))}


@attr.s
class SomeModel(RedisModel):
    title: str = attr.ib()
    value: int = attr.ib()
    some_list: list[str] = attr.ib()
    enum_field: SomeEnum = attr.ib()
    sub_model: SomeSubModel = attr.ib()
    sub_model_with_secret: SomeSubModelWithSecret = attr.ib()

    KEY_PREFIX: str = "some"

    def get_secret_keys(self) -> set[DataKey]:
        lower_level_keys = self.sub_model_with_secret.get_secret_keys()
        return {DataKey(parts=("sub_model_with_secret",) + ll_key.parts) for ll_key in lower_level_keys}


@attr.s
class SomeAuthorizedModel(RedisModelUserIdAuth):
    title: str = attr.ib()
    value: int = attr.ib()

    KEY_PREFIX: str = "some_auth"


class SomeSubModelSchema(BaseSchema):
    class Meta(BaseSchema.Meta):
        target = SomeSubModel

    title = fields.String()
    value = fields.Integer()


class SomeSubModelWithSecretSchema(BaseSchema):
    class Meta(BaseSchema.Meta):
        target = SomeSubModelWithSecret

    secret = fields.String()


class SomeModelSchema(BaseModelSchema):
    class Meta(BaseModelSchema.Meta):
        target = SomeModel

    title = fields.String()
    value = fields.Integer()
    some_list = fields.List(fields.String)
    enum_field = fields.Enum(SomeEnum)
    sub_model = fields.Nested(SomeSubModelSchema)
    sub_model_with_secret = fields.Nested(SomeSubModelWithSecretSchema)


class SomeAuthorizedModelSchema(BaseModelSchema):
    class Meta(BaseModelSchema.Meta):
        target = SomeAuthorizedModel

    title = fields.String()
    value = fields.Integer()


register_redis_model_storage_schema(SomeModel, SomeModelSchema)
register_redis_model_storage_schema(SomeAuthorizedModel, SomeAuthorizedModelSchema)


@pytest.mark.asyncio
async def test_redis_model(redis_cli):
    rmm = RedisModelManager(redis=redis_cli, crypto_keys_config=get_dummy_crypto_keys_config())

    with pytest.raises(RedisModelNotFound):
        await SomeModel.get(manager=rmm, obj_id="qwerty")

    obj_sub = SomeSubModel(title="qwe", value=12)
    obj_sub_w_secret = SomeSubModelWithSecret(secret="very very secret")
    obj = SomeModel(
        manager=rmm,
        title="some model instance",
        value=42,
        some_list=[str(i) for i in range(10)],
        enum_field=SomeEnum.val1,
        sub_model=obj_sub,
        sub_model_with_secret=obj_sub_w_secret,
    )
    assert obj.id

    await obj.save()

    redis_data = await redis_cli.get(f"some/{obj.id}")
    decoded_data = json.loads(redis_data.decode())
    assert decoded_data.pop("id")
    assert "cypher_text" in json.loads(decoded_data.pop("sub_model_with_secret")["secret"])
    assert decoded_data == {
        "some_list": ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"],
        "enum_field": "val1",
        "sub_model": {"value": 12, "title": "qwe"},
        "value": 42,
        "title": "some model instance",
    }
    obj1 = await SomeModel.get(manager=rmm, obj_id=obj.id)
    assert obj1.title == "some model instance"
    assert obj1.sub_model_with_secret.secret == "very very secret"
    assert obj == obj1

    await obj.delete()
    with pytest.raises(RedisModelNotFound):
        await SomeModel.get(manager=rmm, obj_id=obj.id)


@pytest.mark.asyncio
async def test_authorized_redis_model(redis_cli):
    crypto_keys_config = get_dummy_crypto_keys_config()
    rmm_no_auth = RedisModelManager(redis=redis_cli, crypto_keys_config=crypto_keys_config)
    rmm_auth_1 = RedisModelManager(
        redis=redis_cli, rci=RequestContextInfo(user_id="123"), crypto_keys_config=crypto_keys_config
    )
    rmm_auth_2 = RedisModelManager(
        redis=redis_cli, rci=RequestContextInfo(user_id="456"), crypto_keys_config=crypto_keys_config
    )

    obj = SomeAuthorizedModel(manager=rmm_auth_1, title="qwe", value=12)
    initial_obj_id = obj.id
    assert obj.user_id == "123"

    await obj.save()

    obj1 = await SomeAuthorizedModel.get(manager=rmm_no_auth, obj_id=initial_obj_id)
    assert obj1.user_id == "123"

    obj1 = await SomeAuthorizedModel.get(manager=rmm_auth_1, obj_id=initial_obj_id)
    assert obj1.user_id == "123"

    obj1 = await SomeAuthorizedModel.get_authorized(manager=rmm_auth_1, obj_id=initial_obj_id)
    assert obj1.user_id == "123"

    with pytest.raises(RedisModelAccessDenied):
        await SomeAuthorizedModel.get_authorized(manager=rmm_auth_2, obj_id=initial_obj_id)


class SomeRedisSet(RedisSetManager):
    KEY_PREFIX = "some_set_prefix"


@pytest.mark.asyncio
async def test_redis_set(redis_cli):
    some_set = SomeRedisSet(redis=redis_cli, id=str(uuid.uuid4()))

    expected_values = {"a", "b", "c"}

    for value in expected_values:
        await some_set.add(value)
    set_size = await some_set.size()
    assert set_size == len(expected_values)

    actual_values = set()
    async for value in some_set.sscan_iter():
        actual_values.add(value)

    assert actual_values == expected_values

    await some_set.delete()
    c = 0
    async for _ in some_set.sscan_iter():
        c += 1
    assert c == 0

lib/dl_file_uploader_lib/dl_file_uploader_lib_tests/db/test_redis_models.py
datalens-tech/datalens-backend
import os

import pytest

from dl_testing.env_params.generic import GenericEnvParamGetter


@pytest.fixture(scope="session")
def env_param_getter() -> GenericEnvParamGetter:
    filepath = os.path.join(os.path.dirname(__file__), "params.yml")
    return GenericEnvParamGetter.from_yaml_file(filepath)

lib/dl_file_uploader_lib/dl_file_uploader_lib_tests/ext/conftest.py
datalens-tech/datalens-backend
import logging
import os

import pytest

from dl_core.aio.web_app_services.gsheets import (
    GSheetsSettings,
    NumberFormatType,
)
from dl_file_uploader_lib.gsheets_client import GSheetsClient
from dl_utils.aio import ContextVarExecutor


LOGGER = logging.getLogger(__name__)


@pytest.mark.asyncio
async def test_number_format_types(env_param_getter, redis_model_manager):
    gsheets_settings = GSheetsSettings(
        api_key=env_param_getter.get_str_value("GOOGLE_API_KEY"),
        client_id="dummy",
        client_secret="dummy",
    )
    spreadsheet_id = "1vTeWZz2M8c3bvTdRIrqsNhN9m4YwjUguwAnAqBahxi0"
    auth = None

    with ContextVarExecutor(max_workers=min(32, os.cpu_count() * 3 + 4)) as tpe:
        async with GSheetsClient(settings=gsheets_settings, auth=auth, tpe=tpe) as sheets_client:
            spreadsheet = await sheets_client.get_spreadsheet_sample(
                spreadsheet_id=spreadsheet_id,
            )

    for sheet in spreadsheet.sheets[1:]:  # the first sheet contains test data
        data = sheet.data
        header = data[0]  # the first row contains expected number format types
        expected_types = [NumberFormatType(cell.value.upper()) for cell in header]
        for row_idx, row in enumerate(data[1:]):
            for col_idx, (cell, expected_fmt) in enumerate(zip(row, expected_types)):
                assert cell.number_format_type == expected_fmt, f"Failed at [{row_idx + 1}, {col_idx}], {cell=}"
        LOGGER.info(f"{sheet.title} is OK")

lib/dl_file_uploader_lib/dl_file_uploader_lib_tests/ext/test_gsheets_client.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Optional,
)

import arq
import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_api_commons.tenant_resolver import TenantResolver
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_core.aio.web_app_services.gsheets import GSheetsSettings
from dl_core.aio.web_app_services.redis import RedisBaseService
from dl_core.aio.web_app_services.s3 import S3Service
from dl_core.services_registry.top_level import ServicesRegistry
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_file_uploader_task_interface.utils_service_registry import (
    create_sr_factory_from_env_vars,
    get_async_service_us_manager,
)
from dl_task_processor.context import BaseContext
from dl_task_processor.processor import (
    TaskProcessor,
    make_task_processor,
)
from dl_utils.aio import ContextVarExecutor


@attr.s
class SecureReaderSettings:
    socket: str = attr.ib()
    endpoint: Optional[str] = attr.ib(default=None)
    cafile: Optional[str] = attr.ib(default=None)


@attr.s
class FileUploaderTaskContext(BaseContext):
    settings: Optional[Any] = attr.ib()
    tpe: ContextVarExecutor = attr.ib()
    redis_service: RedisBaseService = attr.ib()
    s3_service: S3Service = attr.ib()
    gsheets_settings: GSheetsSettings = attr.ib()
    redis_pool: arq.ArqRedis = attr.ib()
    crypto_keys_config: CryptoKeysConfig = attr.ib()
    secure_reader_settings: SecureReaderSettings = attr.ib()
    tenant_resolver: TenantResolver = attr.ib()

    def get_rci(self) -> RequestContextInfo:
        return RequestContextInfo.create_empty()

    def get_service_registry(self, rci: Optional[RequestContextInfo] = None) -> ServicesRegistry:
        rci = rci or RequestContextInfo.create_empty()
        return create_sr_factory_from_env_vars(
            self.settings.CONNECTORS,
        ).make_service_registry(rci)

    def get_async_usm(self, rci: Optional[RequestContextInfo] = None) -> AsyncUSManager:
        rci = rci or RequestContextInfo.create_empty()
        services_registry = self.get_service_registry(rci=rci)
        return get_async_service_us_manager(
            us_host=self.settings.US_BASE_URL,
            us_master_token=self.settings.US_MASTER_TOKEN,
            services_registry=services_registry,
            bi_context=rci,
            crypto_keys_config=self.crypto_keys_config,
        )

    def make_task_processor(self, request_id: Optional[str]) -> TaskProcessor:
        return make_task_processor(
            redis_pool=self.redis_pool,
            request_id=request_id,
        )

lib/dl_file_uploader_task_interface/dl_file_uploader_task_interface/context.py
datalens-tech/datalens-backend
import enum
from typing import Optional

import attr

from dl_task_processor.task import (
    BaseTaskMeta,
    TaskName,
)


class TaskExecutionMode(enum.Enum):
    BASIC = enum.auto()
    UPDATE_NO_SAVE = enum.auto()
    UPDATE_AND_SAVE = enum.auto()


@attr.s
class DownloadGSheetTask(BaseTaskMeta):
    name = TaskName("download_gsheet")

    file_id: str = attr.ib()
    authorized: bool = attr.ib()
    schedule_parsing: bool = attr.ib(default=True)

    tenant_id: Optional[str] = attr.ib(default=None)
    connection_id: Optional[str] = attr.ib(default=None)
    exec_mode: TaskExecutionMode = attr.ib(default=TaskExecutionMode.BASIC)


@attr.s
class DownloadYaDocsTask(BaseTaskMeta):
    name = TaskName("download_yadocs")

    file_id: str = attr.ib()


@attr.s
class ParseFileTask(BaseTaskMeta):
    name = TaskName("parse_file")

    file_id: str = attr.ib()
    source_id: Optional[str] = attr.ib(default=None)
    file_settings: Optional[dict] = attr.ib(factory=dict)
    source_settings: Optional[dict] = attr.ib(factory=dict)

    tenant_id: Optional[str] = attr.ib(default=None)
    connection_id: Optional[str] = attr.ib(default=None)
    exec_mode: TaskExecutionMode = attr.ib(default=TaskExecutionMode.BASIC)


@attr.s
class ProcessExcelTask(BaseTaskMeta):
    name = TaskName("process_excel")

    file_id: str = attr.ib()


@attr.s
class SaveSourceTask(BaseTaskMeta):
    """
    src_source_id -- the one that is supposed to be found in the DataFile.
    dst_source_id -- the one that is supposed to be saved in the Connection.

    This plays its role when replacing file sources since at this point (on save)
    source IDs in the DataFile and in the Connection are not the same for the same file.

    src_source_id == dst_source_id when the source is being saved for the 1st time
    """

    name = TaskName("save_source")

    tenant_id: str = attr.ib()
    file_id: str = attr.ib()
    src_source_id: str = attr.ib()
    dst_source_id: str = attr.ib()
    connection_id: str = attr.ib()

    exec_mode: TaskExecutionMode = attr.ib(default=TaskExecutionMode.BASIC)


@attr.s
class DeleteFileTask(BaseTaskMeta):
    name = TaskName("delete_file")

    s3_filename: str = attr.ib()
    preview_id: Optional[str] = attr.ib(default=None)


@attr.s
class CleanupTenantTask(BaseTaskMeta):
    name = TaskName("cleanup")

    tenant_id: str = attr.ib()


@attr.s
class CleanS3LifecycleRulesTask(BaseTaskMeta):
    name = TaskName("regular_bucket_lifecycle_cleanup")


@attr.s
class CleanupTenantFilePreviewsTask(BaseTaskMeta):
    name = TaskName("cleanup_previews")

    tenant_id: str = attr.ib()


@attr.s
class RenameTenantFilesTask(BaseTaskMeta):
    name = TaskName("rename_tenant_files_task")

    tenant_id: str = attr.ib()

lib/dl_file_uploader_task_interface/dl_file_uploader_task_interface/tasks.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import os
from typing import (
    TYPE_CHECKING,
    Optional,
)

import attr

from dl_api_commons.base_models import RequestContextInfo
from dl_configs.crypto_keys import CryptoKeysConfig
from dl_configs.rqe import rqe_config_from_env
from dl_core.services_registry.env_manager_factory import InsecureEnvManagerFactory
from dl_core.services_registry.sr_factories import DefaultSRFactory
from dl_core.services_registry.top_level import ServicesRegistry
from dl_core.united_storage_client import USAuthContextMaster
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_file_uploader_worker_lib.settings import FileUploaderConnectorsSettings

from dl_connector_bundle_chs3.chs3_gsheets.core.constants import CONNECTION_TYPE_GSHEETS_V2
from dl_connector_bundle_chs3.file.core.constants import CONNECTION_TYPE_FILE
from dl_connector_clickhouse.core.clickhouse_base.conn_options import CHConnectOptions


if TYPE_CHECKING:
    from dl_core.connection_models import ConnectOptions
    from dl_core.us_connection_base import ExecutorBasedMixin


LOGGER = logging.getLogger(__name__)


def create_sr_factory_from_env_vars(connectors_settings: FileUploaderConnectorsSettings) -> DefaultSRFactory:
    def get_conn_options(conn: ExecutorBasedMixin) -> Optional[ConnectOptions]:
        opts = conn.get_conn_options()

        if isinstance(opts, CHConnectOptions):
            opts = attr.evolve(opts, max_execution_time=int(os.environ.get("CH_MAX_EXECUTION_TIME", 1 * 60 * 60)))

        return attr.evolve(
            opts,
            rqe_total_timeout=int(os.environ.get("RQE_TOTAL_TIMEOUT", 1 * 60 * 60)),
            rqe_sock_read_timeout=int(os.environ.get("RQE_SOCK_READ_TIMEOUT", 30 * 60)),
        )

    connectors_settings = {
        CONNECTION_TYPE_FILE: connectors_settings.FILE,
        CONNECTION_TYPE_GSHEETS_V2: connectors_settings.FILE,
    }
    return DefaultSRFactory(
        rqe_config=rqe_config_from_env(),
        async_env=True,
        connect_options_factory=get_conn_options,
        env_manager_factory=InsecureEnvManagerFactory(),
        connectors_settings=connectors_settings,
    )


def get_async_service_us_manager(
    us_host: str,
    us_master_token: str,
    crypto_keys_config: CryptoKeysConfig,
    bi_context: Optional[RequestContextInfo] = None,
    services_registry: Optional[ServicesRegistry] = None,
) -> AsyncUSManager:
    usm = AsyncUSManager(
        us_api_prefix="private",
        us_base_url=us_host,
        us_auth_context=USAuthContextMaster(us_master_token=us_master_token),
        crypto_keys_config=crypto_keys_config,
        bi_context=bi_context or RequestContextInfo.create_empty(),
        services_registry=services_registry,
    )

    return usm

lib/dl_file_uploader_task_interface/dl_file_uploader_task_interface/utils_service_registry.py
datalens-tech/datalens-backend
import abc
import os
from typing import (
    Generic,
    Optional,
    TypeVar,
)

import attr

from dl_api_commons.tenant_resolver import (
    CommonTenantResolver,
    TenantResolver,
)
from dl_core.aio.web_app_services.gsheets import GSheetsSettings
from dl_core.aio.web_app_services.s3 import S3Service
from dl_core.loader import (
    CoreLibraryConfig,
    load_core_lib,
)
from dl_file_uploader_lib.settings_utils import init_redis_service
from dl_file_uploader_task_interface.context import (
    FileUploaderTaskContext,
    SecureReaderSettings,
)
from dl_file_uploader_task_interface.tasks import CleanS3LifecycleRulesTask
from dl_file_uploader_worker_lib.settings import FileUploaderWorkerSettings
from dl_file_uploader_worker_lib.tasks import REGISTRY
from dl_task_processor.arq_wrapper import (
    CronSchedule,
    create_arq_redis_settings,
    create_redis_pool,
    make_cron_task,
)
from dl_task_processor.context import BaseContextFabric
from dl_task_processor.executor import ExecutorFabric
from dl_task_processor.state import (
    DummyStateImpl,
    TaskState,
)
from dl_task_processor.worker import (
    ArqWorker,
    WorkerSettings,
)
from dl_utils.aio import ContextVarExecutor


_TSettings = TypeVar("_TSettings", bound=FileUploaderWorkerSettings)


@attr.s
class FileUploaderContextFab(BaseContextFabric):
    _settings: FileUploaderWorkerSettings = attr.ib()
    _tenant_resolver: TenantResolver = attr.ib(factory=lambda: CommonTenantResolver())

    async def make(self) -> FileUploaderTaskContext:
        core_conn_whitelist = ["clickhouse", "file", "gsheets_v2"]
        load_core_lib(core_lib_config=CoreLibraryConfig(core_connector_ep_names=core_conn_whitelist))

        redis_service = init_redis_service(self._settings)
        s3_service = S3Service(
            access_key_id=self._settings.S3.ACCESS_KEY_ID,
            secret_access_key=self._settings.S3.SECRET_ACCESS_KEY,
            endpoint_url=self._settings.S3.ENDPOINT_URL,
            tmp_bucket_name=self._settings.S3_TMP_BUCKET_NAME,
            persistent_bucket_name=self._settings.S3_PERSISTENT_BUCKET_NAME,
        )
        await redis_service.initialize()
        await s3_service.initialize()
        redis_pool = await create_redis_pool(create_arq_redis_settings(self._settings.REDIS_ARQ))
        return FileUploaderTaskContext(
            settings=self._settings,
            s3_service=s3_service,
            redis_service=redis_service,
            tpe=ContextVarExecutor(max_workers=min(32, (os.cpu_count() or 0) * 3 + 4)),
            gsheets_settings=GSheetsSettings(
                api_key=self._settings.GSHEETS_APP.API_KEY,
                client_id=self._settings.GSHEETS_APP.CLIENT_ID,
                client_secret=self._settings.GSHEETS_APP.CLIENT_SECRET,
            ),
            redis_pool=redis_pool,
            crypto_keys_config=self._settings.CRYPTO_KEYS_CONFIG,
            secure_reader_settings=SecureReaderSettings(
                socket=self._settings.SECURE_READER.SOCKET,
                endpoint=self._settings.SECURE_READER.ENDPOINT,
                cafile=self._settings.SECURE_READER.CAFILE,
            ),
            tenant_resolver=self._tenant_resolver,
        )

    async def tear_down(self, inst: FileUploaderTaskContext) -> None:  # type: ignore
        await inst.s3_service.tear_down()
        await inst.redis_service.tear_down()
        inst.tpe.shutdown()


@attr.s(kw_only=True)
class FileUploaderWorkerFactory(Generic[_TSettings], abc.ABC):
    _settings: _TSettings = attr.ib()

    @abc.abstractmethod
    def _get_tenant_resolver(self) -> TenantResolver:
        raise NotImplementedError()

    def create_worker(self, state: Optional[TaskState] = None) -> ArqWorker:
        if state is None:
            state = TaskState(DummyStateImpl())
        cron_tasks = []
        if self._settings.ENABLE_REGULAR_S3_LC_RULES_CLEANUP:
            cron_tasks.append(
                make_cron_task(
                    CleanS3LifecycleRulesTask(),
                    schedule=CronSchedule(hour={23}, minute={0}, second={0}),
                )
            )
        worker = ArqWorker(
            redis_settings=create_arq_redis_settings(self._settings.REDIS_ARQ),
            executor_fab=ExecutorFabric(
                registry=REGISTRY,
                state=state,
            ),
            context_fab=FileUploaderContextFab(self._settings, tenant_resolver=self._get_tenant_resolver()),
            worker_settings=WorkerSettings(),
            cron_tasks=cron_tasks,
        )
        return worker

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/app.py
datalens-tech/datalens-backend
from typing import Optional

import attr

from dl_configs.connectors_settings import ConnectorsConfigType
from dl_configs.settings_loaders.meta_definition import s_attrib
from dl_configs.settings_loaders.settings_obj_base import SettingsBase
from dl_configs.settings_submodels import GoogleAppSettings
from dl_file_uploader_lib.settings import FileUploaderBaseSettings

from dl_connector_bundle_chs3.chs3_base.core.settings import FileS3ConnectorSettings
from dl_connector_bundle_chs3.file.core.settings import file_s3_settings_fallback


@attr.s(frozen=True)
class SecureReader(SettingsBase):
    SOCKET: str = s_attrib("SOCKET")  # type: ignore
    ENDPOINT: Optional[str] = s_attrib("ENDPOINT", missing=None)  # type: ignore
    CAFILE: Optional[str] = s_attrib("CAFILE", missing=None)  # type: ignore


@attr.s(frozen=True)
class FileUploaderConnectorsSettings(SettingsBase):
    FILE: Optional[FileS3ConnectorSettings] = s_attrib("FILE", missing=None)  # type: ignore


def file_uploader_connectors_settings_fallback(full_cfg: ConnectorsConfigType) -> FileUploaderConnectorsSettings:
    settings = file_s3_settings_fallback(full_cfg)
    return FileUploaderConnectorsSettings(**settings)


@attr.s(frozen=True)
class FileUploaderWorkerSettings(FileUploaderBaseSettings):
    SENTRY_DSN: Optional[str] = s_attrib(  # type: ignore
        "DL_SENTRY_DSN",
        fallback_cfg_key="SENTRY_DSN_FILE_UPLOADER_WORKER",
        missing=None,
    )

    US_BASE_URL: str = s_attrib("US_HOST", fallback_cfg_key="US_BASE_URL")  # type: ignore
    US_MASTER_TOKEN: str = s_attrib("US_MASTER_TOKEN", sensitive=True)  # type: ignore

    CONNECTORS: Optional[FileUploaderConnectorsSettings] = s_attrib(  # type: ignore
        "CONNECTORS",
        fallback_factory=file_uploader_connectors_settings_fallback,
    )

    GSHEETS_APP: GoogleAppSettings = s_attrib(  # type: ignore
        # TODO: check gsheets connector availability
        "GSHEETS_APP",
    )

    ENABLE_REGULAR_S3_LC_RULES_CLEANUP: bool = s_attrib("ENABLE_REGULAR_S3_LC_RULES_CLEANUP", missing=False)  # type: ignore
    SECURE_READER: SecureReader = s_attrib(  # type: ignore
        "SECURE_READER",
        fallback_factory=(
            lambda: SecureReader(  # type: ignore
                SOCKET="/var/reader.sock",
                ENDPOINT=None,
                CAFILE=None,
            )
        ),
    )

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/settings.py
datalens-tech/datalens-backend
from dl_task_processor.task import TaskRegistry

from .cleanup import (
    CleanS3LifecycleRulesTask,
    CleanupTenantFilePreviewsTask,
    CleanupTenantTask,
    RenameTenantFilesTask,
)
from .delete import DeleteFileTask
from .download_gsheets import DownloadGSheetTask
from .download_yadocs import DownloadYaDocsTask
from .excel import ProcessExcelTask
from .parse import ParseFileTask
from .save import SaveSourceTask


REGISTRY: TaskRegistry = TaskRegistry.create(
    [
        DownloadGSheetTask,
        DownloadYaDocsTask,
        ParseFileTask,
        ProcessExcelTask,
        SaveSourceTask,
        DeleteFileTask,
        CleanupTenantTask,
        CleanS3LifecycleRulesTask,
        CleanupTenantFilePreviewsTask,
        RenameTenantFilesTask,
    ]
)

__all__ = ("REGISTRY",)

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/tasks/__init__.py
datalens-tech/datalens-backend
from datetime import (
    datetime,
    timedelta,
)
import logging
from typing import Any

import attr
from botocore.exceptions import ClientError
from redis.asyncio.lock import Lock as RedisLock

from dl_constants.enums import DataSourceRole
from dl_core.us_entry import USMigrationEntry
from dl_file_uploader_lib.enums import RenameTenantStatus
from dl_file_uploader_lib.redis_model.base import (
    RedisModelManager,
    RedisModelNotFound,
)
from dl_file_uploader_lib.redis_model.models import (
    DataSourcePreview,
    PreviewSet,
    RenameTenantStatusModel,
)
from dl_file_uploader_task_interface.context import FileUploaderTaskContext
import dl_file_uploader_task_interface.tasks as task_interface
from dl_file_uploader_worker_lib.tasks.save import _make_source_s3_filename
from dl_task_processor.task import (
    BaseExecutorTask,
    Retry,
    Success,
    TaskResult,
)

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3.chs3_gsheets.core.constants import CONNECTION_TYPE_GSHEETS_V2
from dl_connector_bundle_chs3.file.core.constants import CONNECTION_TYPE_FILE


LOGGER = logging.getLogger(__name__)


def make_s3_bucket_lifecycle_update_lock_key(bucket_name: str) -> str:
    return f"s3_bucket_lifecycle_update/{bucket_name}"


@attr.s
class CleanupTenantTask(BaseExecutorTask[task_interface.CleanupTenantTask, FileUploaderTaskContext]):
    cls_meta = task_interface.CleanupTenantTask

    async def run(self) -> TaskResult:
        try:
            tenant_id = self.meta.tenant_id
            obj_prefix = f"{tenant_id}_"
            LOGGER.info(f"CleanupTenantTask. tenant ID: {tenant_id}")

            s3_service = self._ctx.s3_service
            s3_client = s3_service.get_client()

            redis = self._ctx.redis_service.get_redis()
            lock_key = make_s3_bucket_lifecycle_update_lock_key(s3_service.persistent_bucket_name)
            LOGGER.info(f"Acquiring redis lock {lock_key}")
            async with RedisLock(redis, name=lock_key, timeout=120, blocking_timeout=120):
                LOGGER.info(f"Lock {lock_key} acquired")
                # lc == lifecycle
                try:
                    lc_config = await s3_client.get_bucket_lifecycle_configuration(
                        Bucket=s3_service.persistent_bucket_name,
                    )
                except ClientError as ex:
                    if ex.response["Error"]["Code"] == "NoSuchLifecycleConfiguration":
                        lc_config = {"Rules": []}
                    else:
                        raise
                lc_rules = lc_config["Rules"]

                new_rule = {
                    "ID": obj_prefix,  # note: currently unused
                    "Expiration": {
                        "Date": datetime.today().replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=2),
                    },
                    "Filter": {
                        "Prefix": obj_prefix,
                    },
                    "Status": "Enabled",
                }
                await s3_client.put_bucket_lifecycle_configuration(
                    Bucket=s3_service.persistent_bucket_name,
                    LifecycleConfiguration=dict(
                        Rules=lc_rules + [new_rule],
                    ),
                )

                LOGGER.info(f'Updated LC rules to remove objects with prefix "{obj_prefix}"')

            LOGGER.info(f"Released lock {lock_key}")

        except Exception as ex:
            LOGGER.exception(ex)
            return Retry(attempts=5)
        return Success()


@attr.s
class CleanS3LifecycleRulesTask(BaseExecutorTask[task_interface.CleanS3LifecycleRulesTask, FileUploaderTaskContext]):
    cls_meta = task_interface.CleanS3LifecycleRulesTask

    async def run(self) -> TaskResult:
        try:
            LOGGER.info("CleanS3LifecycleRulesTask")

            s3_service = self._ctx.s3_service
            s3_client = s3_service.get_client()

            async def prefix_exists(prefix: str) -> bool:
                list_resp = await s3_client.list_objects_v2(
                    Bucket=s3_service.persistent_bucket_name,
                    Prefix=prefix,
                    MaxKeys=1,
                )
                return bool(list_resp.get("Contents"))

            async def remove_old_rules(rules: list[dict[str, Any]]) -> list[dict[str, Any]]:
                indices_to_remove = set()
                for idx, rule in enumerate(rules):
                    rule_prefix = rule["Filter"]["Prefix"]
                    if not await prefix_exists(rule_prefix):
                        LOGGER.info(f"Going to remove lifecycle rule for prefix {rule_prefix}")
                        indices_to_remove.add(idx)
                return [rule for idx, rule in enumerate(rules) if idx not in indices_to_remove]

            redis = self._ctx.redis_service.get_redis()
            lock_key = make_s3_bucket_lifecycle_update_lock_key(s3_service.persistent_bucket_name)
            LOGGER.info(f"Acquiring redis lock {lock_key}")
            async with RedisLock(redis, name=lock_key, timeout=120, blocking_timeout=120):
                LOGGER.info(f"Lock {lock_key} acquired")
                # lc == lifecycle
                try:
                    lc_config = await s3_client.get_bucket_lifecycle_configuration(
                        Bucket=s3_service.persistent_bucket_name
                    )
                except ClientError as ex:
                    if ex.response["Error"]["Code"] == "NoSuchLifecycleConfiguration":
                        lc_config = {"Rules": []}
                    else:
                        raise
                lc_rules = lc_config["Rules"]

                new_lc_rules = await remove_old_rules(lc_rules)

                if len(new_lc_rules) > 750:
                    LOGGER.warning(f"Number of LC rules: {len(new_lc_rules)}")

                if new_lc_rules:
                    if len(new_lc_rules) < len(lc_rules):
                        await s3_client.put_bucket_lifecycle_configuration(
                            Bucket=s3_service.persistent_bucket_name,
                            LifecycleConfiguration=dict(
                                Rules=new_lc_rules,
                            ),
                        )
                else:
                    await s3_client.delete_bucket_lifecycle(Bucket=s3_service.persistent_bucket_name)

                LOGGER.info(f"Removed {len(lc_rules) - len(new_lc_rules)} LC rule(s)")

            LOGGER.info(f"Released lock {lock_key}")

        except Exception as ex:
            LOGGER.exception(ex)
            return Retry(attempts=5)
        return Success()


@attr.s
class CleanupTenantFilePreviewsTask(
    BaseExecutorTask[task_interface.CleanupTenantFilePreviewsTask, FileUploaderTaskContext]
):
    cls_meta = task_interface.CleanupTenantFilePreviewsTask

    async def run(self) -> TaskResult:
        try:
            tenant_id = self.meta.tenant_id
            LOGGER.info(f"CleanupTenantFilePreviewsTask. tenant ID: {tenant_id}")

            redis = self._ctx.redis_service.get_redis()
            rmm = RedisModelManager(
                redis=self._ctx.redis_service.get_redis(),
                crypto_keys_config=self._ctx.crypto_keys_config,
            )
            preview_set = PreviewSet(redis=redis, id=tenant_id)
            async for preview_id in preview_set.sscan_iter():
                try:
                    preview = await DataSourcePreview.get(manager=rmm, obj_id=preview_id)
                    await preview.delete()
                    LOGGER.info(f"Successfully deleted preview id={preview_id} for tenant {tenant_id}")
                except RedisModelNotFound:
                    LOGGER.info(f"Preview id={preview_id} not found for tenant {tenant_id}")

            LOGGER.info(f"Done with previews for tenant {tenant_id}, going to delete Redis Set {preview_set.key}")
            await preview_set.delete()

        except Exception as ex:
            LOGGER.exception(ex)
            return Retry(attempts=5)
        return Success()


@attr.s
class RenameTenantFilesTask(BaseExecutorTask[task_interface.RenameTenantFilesTask, FileUploaderTaskContext]):
    cls_meta = task_interface.RenameTenantFilesTask

    async def run(self) -> TaskResult:
        tenant_id = self.meta.tenant_id
        LOGGER.info(f"RenameTenantFilesTask. Moving to {tenant_id}")

        rmm = RedisModelManager(
            redis=self._ctx.redis_service.get_redis(),
            crypto_keys_config=self._ctx.crypto_keys_config,
        )
        await RenameTenantStatusModel(manager=rmm, id=tenant_id, status=RenameTenantStatus.started).save()
        try:
            usm = self._ctx.get_async_usm()
            usm.set_tenant_override(self._ctx.tenant_resolver.resolve_tenant_def_by_tenant_id(self.meta.tenant_id))
            s3_service = self._ctx.s3_service
            s3_client = s3_service.get_client()

            s3_file_based_conn_types = (
                CONNECTION_TYPE_FILE,
                CONNECTION_TYPE_GSHEETS_V2,
            )

            redis = self._ctx.redis_service.get_redis()
            lock_key = f"rename_tenant_files_task/{tenant_id}"
            LOGGER.info(f"Acquiring redis lock {lock_key}")
            async with RedisLock(redis, name=lock_key, timeout=120, blocking_timeout=120):
                LOGGER.info(f"Lock {lock_key} acquired")

                old_tenants: set[str] = set()

                async with usm:
                    for conn_type in s3_file_based_conn_types:
                        conn_iter = usm.get_collection(
                            entry_cls=USMigrationEntry,
                            entry_type=conn_type.name,
                            entry_scope="connection",
                            raise_on_broken_entry=True,
                        )
                        async for conn_placeholder in conn_iter:
                            async with usm.locked_entry_cm(
                                conn_placeholder.uuid,
                                expected_type=BaseFileS3Connection,
                                wait_timeout_sec=10,
                                duration_sec=300,
                                force=True,
                            ) as conn:
                                assert isinstance(conn, BaseFileS3Connection)
                                conn_changed = False
                                for source in conn.data.sources:
                                    old_s3_filename = source.s3_filename
                                    if old_s3_filename is None:
                                        LOGGER.info(
                                            f"Cannot rename file for source_id {source.id} - s3_filename not set"
                                        )
                                        continue
                                    if old_s3_filename.startswith(tenant_id):
                                        LOGGER.info(f"File is already tagged with the new tenant: {old_s3_filename}")
                                        continue

                                    old_fname_parts = old_s3_filename.split("_")
                                    if len(old_fname_parts) == 2 and old_fname_parts[0] and old_fname_parts[1]:
                                        # assume that first part is old tenant id
                                        old_tenants.add(old_fname_parts[0])

                                    new_s3_filename = _make_source_s3_filename(tenant_id)
                                    await s3_client.copy_object(
                                        CopySource=dict(
                                            Bucket=s3_service.persistent_bucket_name,
                                            Key=old_s3_filename,
                                        ),
                                        Bucket=s3_service.persistent_bucket_name,
                                        Key=new_s3_filename,
                                    )
                                    await s3_client.delete_object(
                                        Bucket=s3_service.persistent_bucket_name,
                                        Key=old_s3_filename,
                                    )
                                    LOGGER.info(f"Moved s3 file {old_s3_filename} -> {new_s3_filename}")

                                    conn_changed = True
                                    conn.update_data_source(
                                        source.id,
                                        role=DataSourceRole.origin,
                                        s3_filename=new_s3_filename,
                                    )

                                if conn_changed:
                                    await usm.save(conn)
                if old_tenants:
                    LOGGER.info(f"Moving redis old preview set[s]. From {old_tenants} to {tenant_id}")
                    await PreviewSet(redis=redis, id=tenant_id).sunion_by_id(*old_tenants)
                    for t_id in old_tenants:
                        await PreviewSet(redis=redis, id=t_id).delete()
        except Exception as ex:
            LOGGER.exception(ex)
            await RenameTenantStatusModel(manager=rmm, id=tenant_id, status=RenameTenantStatus.error).save()
            return Retry(attempts=10)

        await RenameTenantStatusModel(manager=rmm, id=tenant_id, status=RenameTenantStatus.success).save()
        return Success()

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/tasks/cleanup.py
datalens-tech/datalens-backend
import logging

import attr
from botocore.exceptions import ClientError

from dl_file_uploader_lib.redis_model.base import (
    RedisModelManager,
    RedisModelNotFound,
)
from dl_file_uploader_lib.redis_model.models import DataSourcePreview
from dl_file_uploader_task_interface.context import FileUploaderTaskContext
import dl_file_uploader_task_interface.tasks as task_interface
from dl_task_processor.task import (
    BaseExecutorTask,
    Retry,
    Success,
    TaskResult,
)


LOGGER = logging.getLogger(__name__)


@attr.s
class DeleteFileTask(BaseExecutorTask[task_interface.DeleteFileTask, FileUploaderTaskContext]):
    cls_meta = task_interface.DeleteFileTask

    async def run(self) -> TaskResult:
        try:
            s3_filename = self.meta.s3_filename
            preview_id = self.meta.preview_id
            LOGGER.info(f"DeleteFileTask. Filename: {s3_filename}")

            if preview_id is None:
                LOGGER.info(f"Unable to delete preview for file {s3_filename} since {preview_id=}")
            else:
                rmm = RedisModelManager(
                    redis=self._ctx.redis_service.get_redis(),
                    crypto_keys_config=self._ctx.crypto_keys_config,
                )
                try:
                    preview = await DataSourcePreview.get(manager=rmm, obj_id=preview_id)
                    await preview.delete()
                    LOGGER.info(f"Successfully deleted preview id={preview_id} for file {s3_filename}")
                except RedisModelNotFound:
                    LOGGER.info(f"Preview id={preview_id} not found for file {s3_filename}")

            s3_service = self._ctx.s3_service
            s3_client = s3_service.get_client()
            try:
                await s3_client.delete_object(Bucket=s3_service.persistent_bucket_name, Key=s3_filename)
                LOGGER.info(f"Successfully deleted file {s3_filename} from the persistent bucket")
            except ClientError as ex:
                if ex.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
                    LOGGER.info(f"File {s3_filename} was not found in the persistent bucket.")
                else:
                    raise

        except Exception as ex:
            LOGGER.exception(ex)
            return Retry(attempts=5)
        return Success()

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/tasks/delete.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
from collections import defaultdict
import itertools
import logging
from typing import (
    AsyncIterator,
    Iterable,
    Optional,
)

import aiogoogle
import attr

from dl_constants.enums import (
    FileProcessingStatus,
    UserDataType,
)
from dl_core.aio.web_app_services.gsheets import (
    Range,
    Sheet,
)
from dl_core.db import SchemaColumn
from dl_core.raw_data_streaming.stream import SimpleUntypedAsyncDataStream
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_file_uploader_lib import exc
from dl_file_uploader_lib.data_sink.json_each_row import S3JsonEachRowUntypedFileAsyncDataSink
from dl_file_uploader_lib.gsheets_client import (
    GSheetsClient,
    GSheetsOAuth2,
    google_api_error_to_file_uploader_exception,
)
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    DataSource,
    FileProcessingError,
    GSheetsFileSourceSettings,
    GSheetsUserSourceDataSourceProperties,
    GSheetsUserSourceProperties,
)
from dl_file_uploader_task_interface.context import FileUploaderTaskContext
import dl_file_uploader_task_interface.tasks as task_interface
from dl_file_uploader_task_interface.tasks import TaskExecutionMode
from dl_file_uploader_worker_lib.utils.connection_error_tracker import FileConnectionDataSourceErrorTracker
from dl_file_uploader_worker_lib.utils.parsing_utils import guess_header_and_schema_gsheet
from dl_task_processor.task import (
    BaseExecutorTask,
    Fail,
    Retry,
    Success,
    TaskResult,
)

from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection


LOGGER = logging.getLogger(__name__)


class NoToken(Exception):
    pass


async def _get_gsheets_auth(
    dfile_token: Optional[str],
    conn_id: Optional[str],
    usm: AsyncUSManager,
) -> Optional[GSheetsOAuth2]:
    if dfile_token is not None:  # if there is a token in dfile, then use it
        refresh_token = dfile_token
    elif conn_id is not None:  # otherwise, use the one from the connection
        conn: GSheetsFileS3Connection = await usm.get_by_id(conn_id, GSheetsFileS3Connection)
        if conn.data.refresh_token is None:
            raise NoToken()
        refresh_token = conn.data.refresh_token
    else:
        raise NoToken()

    return GSheetsOAuth2(access_token=None, refresh_token=refresh_token) if refresh_token is not None else None


async def _values_data_iter(
    sheets_client: GSheetsClient,
    spreadsheet_id: str,
    sheet_sample: Sheet,
    user_types: list[UserDataType | str],
    raw_schema_body: list[SchemaColumn],
) -> AsyncIterator[list]:
    """
    Trying to convert values to the originally guessed schema
    Fallback to strings after all data is read, for which
    raw_schema_body is updated inplace on each portion
    """

    acceptable_resp_size_bytes = 4e6
    acceptable_deviation = 0.25
    max_acceptable_resp_size = acceptable_resp_size_bytes * (1 + acceptable_deviation)

    # dump the first row separately not to bother about its types later
    assert sheet_sample.data is not None
    yield [cell.value for cell in sheet_sample.data[0]]

    current_range = Range(
        sheet_title=sheet_sample.title,
        row_from=2,  # start reading from the second line, because the first one is dumped above
        col_from=1,
        row_to=sheet_sample.batch_size_rows,
        col_to=sheet_sample.column_count,
    )
    body_rows_total = 0
    while True:
        sheet_portion_values, new_user_types = await sheets_client.get_single_values_range(
            spreadsheet_id=spreadsheet_id,
            range=current_range,
            user_types=user_types,
        )
        body_rows_total += len(sheet_portion_values)
        if not sheet_portion_values:
            # no data received, we assume that we are done
            if body_rows_total == 0:
                raise exc.TooFewRowsError()
            return

        resp_size = sheets_client.last_response_size_bytes
        if resp_size > max_acceptable_resp_size:
            prev_batch_size_rows = sheet_sample.batch_size_rows
            adjust_rows_coeff = acceptable_resp_size_bytes / resp_size
            sheet_sample.batch_size_rows = int(sheet_sample.batch_size_rows * adjust_rows_coeff)
            LOGGER.info(
                f"Got a response with size {resp_size}, which is too large, going to adjust"
                f" batch size rows from {prev_batch_size_rows}"
                f" to {sheet_sample.batch_size_rows}"
            )

        for idx, (col, new_user_type) in enumerate(zip(raw_schema_body, new_user_types)):
            # fall back to string
            if new_user_type == "time":
                new_user_type = UserDataType.string
            if new_user_type != col.user_type and new_user_type == UserDataType.string:
                raw_schema_body[idx] = raw_schema_body[idx].clone(user_type=new_user_type)

        for row in sheet_portion_values:
            yield row

        current_range.row_from = current_range.row_to + 1
        current_range.row_to = min(current_range.row_to + sheet_sample.batch_size_rows, sheet_sample.row_count)
        if current_range.row_from > sheet_sample.row_count:
            if body_rows_total == 0:
                raise exc.TooFewRowsError()
            return


@attr.s
class DownloadGSheetTask(BaseExecutorTask[task_interface.DownloadGSheetTask, FileUploaderTaskContext]):
    """Loads a spreadsheet into s3, building raw_schema based solely on gsheet cell types, schedules ParseFileTask"""

    cls_meta = task_interface.DownloadGSheetTask

    async def run(self) -> TaskResult:
        dfile: Optional[DataFile] = None
        sources_to_update_by_sheet_id: dict[int, list[DataSource]] = defaultdict(list)
        usm = self._ctx.get_async_usm()
        task_processor = self._ctx.make_task_processor(self._request_id)
        redis = self._ctx.redis_service.get_redis()
        connection_error_tracker = FileConnectionDataSourceErrorTracker(usm, task_processor, redis, self._request_id)
        try:
            LOGGER.info(f"DownloadGSheetTask. Mode: {self.meta.exec_mode.name}. File: {self.meta.file_id}")
            rmm = RedisModelManager(redis=redis, crypto_keys_config=self._ctx.crypto_keys_config)
            dfile = await DataFile.get(manager=rmm, obj_id=self.meta.file_id)
            assert dfile is not None

            assert isinstance(dfile.user_source_properties, GSheetsUserSourceProperties)
            spreadsheet_id = dfile.user_source_properties.spreadsheet_id

            if self.meta.exec_mode == TaskExecutionMode.BASIC:
                dfile.sources = []
            assert dfile.sources is not None

            auth: Optional[GSheetsOAuth2] = None
            if self.meta.authorized:
                try:
                    dfile_token = dfile.user_source_properties.refresh_token
                    auth = await _get_gsheets_auth(dfile_token, self.meta.connection_id, usm)
                except NoToken:
                    LOGGER.error("Authorized call but no token found in either DataFile or connection, failing task")
                    return Fail()

            async with GSheetsClient(self._ctx.gsheets_settings, self._ctx.tpe, auth) as sheets_client:
                try:
                    spreadsheet_sample = await sheets_client.get_spreadsheet_sample(spreadsheet_id)
                except (aiogoogle.HTTPError, exc.DLFileUploaderBaseError) as e:
                    if isinstance(e, aiogoogle.HTTPError):
                        e = google_api_error_to_file_uploader_exception(e)
                    download_error = FileProcessingError.from_exception(e)
                    dfile.status = FileProcessingStatus.failed
                    dfile.error = download_error

                    for src in dfile.sources:
                        src.status = FileProcessingStatus.failed
                        src.error = download_error
                        connection_error_tracker.add_error(src.id, src.error)
                    await dfile.save()

                    await connection_error_tracker.finalize(self.meta.exec_mode, self.meta.connection_id)
                    return Success()

                s3 = self._ctx.s3_service

                dfile.filename = spreadsheet_sample.title
                for src in dfile.sources:
                    assert isinstance(src.user_source_dsrc_properties, GSheetsUserSourceDataSourceProperties)
                    sources_to_update_by_sheet_id[src.user_source_dsrc_properties.sheet_id].append(src)
                for sheet_sample in spreadsheet_sample.sheets:
                    source_status = FileProcessingStatus.in_progress
                    has_header: Optional[bool] = None
                    raw_schema: list[SchemaColumn] = []
                    raw_schema_header: list[SchemaColumn] = []
                    raw_schema_body: list[SchemaColumn] = []
                    sheet_properties = GSheetsUserSourceDataSourceProperties(sheet_id=sheet_sample.id)

                    if self.meta.exec_mode == TaskExecutionMode.BASIC:
                        source_title = spreadsheet_sample.title + " – " + sheet_sample.title
                        sheet_data_sources = [
                            DataSource(
                                title=source_title,
                                raw_schema=raw_schema,
                                status=source_status,
                                error=None,
                            )
                        ]
                        dfile.sources.extend(sheet_data_sources)
                    else:
                        if sheet_sample.id not in sources_to_update_by_sheet_id:
                            continue
                        sheet_data_sources = sources_to_update_by_sheet_id.pop(sheet_sample.id)

                    loop = asyncio.get_running_loop()
                    if not sheet_sample.data:
                        for src in sheet_data_sources:
                            src.error = FileProcessingError.from_exception(exc.EmptyDocument())
                            src.status = FileProcessingStatus.failed
                            connection_error_tracker.add_error(src.id, src.error)
                    else:
                        try:
                            has_header, raw_schema, raw_schema_header, raw_schema_body = await loop.run_in_executor(
                                self._ctx.tpe, guess_header_and_schema_gsheet, sheet_sample
                            )
                        except exc.DLFileUploaderBaseError as e:
                            for src in sheet_data_sources:
                                src.status = FileProcessingStatus.failed
                                src.error = FileProcessingError.from_exception(e)
                                connection_error_tracker.add_error(src.id, src.error)

                    sheet_settings = None

                    for src in sheet_data_sources:
                        if src.is_applicable:
                            if self.meta.exec_mode != TaskExecutionMode.BASIC:
                                assert src.file_source_settings is not None
                                has_header = src.file_source_settings.first_line_is_header
                            assert has_header is not None

                            orig_user_types: list[UserDataType | str] = []
                            for idx, col in enumerate(raw_schema_body):
                                if col.user_type == UserDataType.string and sheet_sample.col_is_time(idx, has_header):
                                    orig_user_types.append("time")
                                else:
                                    orig_user_types.append(col.user_type)

                            data_stream = SimpleUntypedAsyncDataStream(
                                data_iter=_values_data_iter(
                                    sheets_client,
                                    spreadsheet_id,
                                    sheet_sample,
                                    orig_user_types,
                                    raw_schema_body,
                                ),
                                rows_to_copy=None,  # TODO
                            )
                            try:
                                async with S3JsonEachRowUntypedFileAsyncDataSink(
                                    s3=s3.get_client(),
                                    s3_key=src.s3_key,
                                    bucket_name=s3.tmp_bucket_name,
                                ) as data_sink:
                                    await data_sink.dump_data_stream(data_stream)
                            except exc.DLFileUploaderBaseError as e:
                                src.status = FileProcessingStatus.failed
                                src.error = FileProcessingError.from_exception(e)
                                connection_error_tracker.add_error(src.id, src.error)

                            sheet_settings = GSheetsFileSourceSettings(
                                first_line_is_header=has_header,
                                raw_schema_header=raw_schema_header,
                                raw_schema_body=raw_schema_body,
                            )

                        src.raw_schema = raw_schema
                        src.file_source_settings = sheet_settings
                        src.user_source_dsrc_properties = sheet_properties

            not_found_sources: Iterable[DataSource] = itertools.chain(*sources_to_update_by_sheet_id.values())
            for src in not_found_sources:
                src.error = FileProcessingError.from_exception(exc.DocumentNotFound())
                src.status = FileProcessingStatus.failed
                connection_error_tracker.add_error(src.id, src.error)

            await connection_error_tracker.finalize(self.meta.exec_mode, self.meta.connection_id)
            await dfile.save()
            LOGGER.info("DataFile object saved.")

            if self.meta.schedule_parsing:
                parse_file_task = task_interface.ParseFileTask(
                    file_id=dfile.id,
                    tenant_id=self.meta.tenant_id,
                    connection_id=self.meta.connection_id,
                    exec_mode=self.meta.exec_mode,
                )
                await task_processor.schedule(parse_file_task)
                LOGGER.info(f"Scheduled ParseFileTask for file_id {dfile.id}")
            else:
                LOGGER.info(f"Skipping ParseFileTask for file_id {dfile.id} because {self.meta.schedule_parsing=}")
        except Exception as ex:
            LOGGER.exception(ex)
            if dfile is None:
                return Retry(attempts=3)
            else:
                dfile.status = FileProcessingStatus.failed
                exc_to_save = ex if isinstance(ex, exc.DLFileUploaderBaseError) else exc.DownloadFailed()
                dfile.error = FileProcessingError.from_exception(exc_to_save)
                await dfile.save()

                for src in dfile.sources or ():
                    connection_error_tracker.add_error(src.id, dfile.error)
                await connection_error_tracker.finalize(self.meta.exec_mode, self.meta.connection_id)

                return Fail()
        finally:
            await usm.close()
        return Success()

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/tasks/download_gsheets.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import (
    AsyncGenerator,
    Optional,
)

import aiohttp
import attr

from dl_constants.enums import FileProcessingStatus
from dl_file_uploader_lib import exc
from dl_file_uploader_lib.data_sink.raw_bytes import (
    RawBytesAsyncDataStream,
    S3RawFileAsyncDataSink,
)
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    FileProcessingError,
    YaDocsUserSourceProperties,
)
from dl_file_uploader_lib.yadocs_client import (
    YaDocsClient,
    yadocs_error_to_file_uploader_exception,
)
from dl_file_uploader_task_interface.context import FileUploaderTaskContext
import dl_file_uploader_task_interface.tasks as task_interface
from dl_task_processor.task import (
    BaseExecutorTask,
    Fail,
    Retry,
    Success,
    TaskResult,
)


LOGGER = logging.getLogger(__name__)


@attr.s
class DownloadYaDocsTask(BaseExecutorTask[task_interface.DownloadYaDocsTask, FileUploaderTaskContext]):
    cls_meta = task_interface.DownloadYaDocsTask

    async def run(self) -> TaskResult:
        dfile: Optional[DataFile] = None
        redis = self._ctx.redis_service.get_redis()

        try:
            rmm = RedisModelManager(redis=redis, crypto_keys_config=self._ctx.crypto_keys_config)
            dfile = await DataFile.get(manager=rmm, obj_id=self.meta.file_id)
            assert dfile is not None

            assert isinstance(dfile.user_source_properties, YaDocsUserSourceProperties)

            async with aiohttp.ClientSession() as session:
                yadocs_client = YaDocsClient(session)
                try:
                    if dfile.user_source_properties.public_link is not None:
                        spreadsheet_ref = await yadocs_client.get_spreadsheet_public_ref(
                            link=dfile.user_source_properties.public_link
                        )
                        spreadsheet_meta = await yadocs_client.get_spreadsheet_public_meta(
                            link=dfile.user_source_properties.public_link
                        )

                    elif (
                        dfile.user_source_properties.private_path is not None
                        and dfile.user_source_properties.oauth_token is not None
                    ):
                        spreadsheet_ref = await yadocs_client.get_spreadsheet_private_ref(
                            path=dfile.user_source_properties.private_path,
                            token=dfile.user_source_properties.oauth_token,
                        )
                        spreadsheet_meta = await yadocs_client.get_spreadsheet_private_meta(
                            path=dfile.user_source_properties.private_path,
                            token=dfile.user_source_properties.oauth_token,
                        )
                    else:
                        raise exc.DLFileUploaderBaseError()
                except exc.DLFileUploaderBaseError as e:
                    LOGGER.exception(e)
                    download_error = FileProcessingError.from_exception(e)
                    dfile.status = FileProcessingStatus.failed
                    dfile.error = download_error
                    await dfile.save()
                    return Success()

            dfile.filename = spreadsheet_meta["name"]
            s3 = self._ctx.s3_service

            async def _chunk_iter(chunk_size: int = 10 * 1024 * 1024) -> AsyncGenerator[bytes, None]:
                async with aiohttp.ClientSession() as session:
                    async with session.get(spreadsheet_ref) as resp:
                        if resp.status != 200:
                            raise yadocs_error_to_file_uploader_exception(resp.status, await resp.json())
                        while True:
                            chunk = await resp.content.read(chunk_size)
                            if chunk:
                                LOGGER.debug(f"Received chunk of {len(chunk)} bytes.")
                                yield chunk
                            else:
                                LOGGER.info("Empty chunk received.")
                                break

            data_stream = RawBytesAsyncDataStream(data_iter=_chunk_iter())
            async with S3RawFileAsyncDataSink(
                s3=s3.client,
                s3_key=dfile.s3_key,
                bucket_name=s3.tmp_bucket_name,
            ) as data_sink:
                await data_sink.dump_data_stream(data_stream)

            await dfile.save()
            LOGGER.info(f'Uploaded file "{dfile.filename}".')

        except Exception as ex:
            LOGGER.exception(ex)
            if dfile is None:
                return Retry(attempts=3)
            else:
                dfile.status = FileProcessingStatus.failed
                exc_to_save = ex if isinstance(ex, exc.DLFileUploaderBaseError) else exc.DownloadFailed()
                dfile.error = FileProcessingError.from_exception(exc_to_save)
                await dfile.save()

                return Fail()
        return Success()

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/tasks/download_yadocs.py
datalens-tech/datalens-backend
import asyncio
import logging
import ssl
from typing import (
    Iterator,
    Optional,
)
import urllib.parse

import aiohttp
import attr

from dl_constants.enums import FileProcessingStatus
from dl_core.db import SchemaColumn
from dl_core.raw_data_streaming.stream import SimpleUntypedDataStream
from dl_file_uploader_lib import exc
from dl_file_uploader_lib.data_sink.json_each_row import S3JsonEachRowUntypedFileDataSink
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    DataSource,
    ExcelFileSourceSettings,
    FileProcessingError,
)
from dl_file_uploader_task_interface.context import FileUploaderTaskContext
import dl_file_uploader_task_interface.tasks as task_interface
from dl_file_uploader_worker_lib.utils.parsing_utils import guess_header_and_schema_excel
from dl_task_processor.task import (
    BaseExecutorTask,
    Fail,
    Retry,
    Success,
    TaskResult,
)


LOGGER = logging.getLogger(__name__)


@attr.s
class ProcessExcelTask(BaseExecutorTask[task_interface.ProcessExcelTask, FileUploaderTaskContext]):
    """Loads a xlsx file from s3, runs secure-reader, schedules ParseFileTask"""

    cls_meta = task_interface.ProcessExcelTask

    async def run(self) -> TaskResult:
        dfile: Optional[DataFile] = None
        try:
            LOGGER.info(f"ProcessExcelTask. File: {self.meta.file_id}")
            loop = asyncio.get_running_loop()

            redis = self._ctx.redis_service.get_redis()
            rmm = RedisModelManager(redis=redis, crypto_keys_config=self._ctx.crypto_keys_config)
            dfile = await DataFile.get(manager=rmm, obj_id=self.meta.file_id)
            assert dfile is not None

            s3 = self._ctx.s3_service

            dfile.sources = []

            s3_resp = await s3.client.get_object(
                Bucket=s3.tmp_bucket_name,
                Key=dfile.s3_key,
            )
            file_obj = await s3_resp["Body"].read()

            conn: aiohttp.BaseConnector
            if self._ctx.secure_reader_settings.endpoint is not None:
                secure_reader_endpoint = self._ctx.secure_reader_settings.endpoint
                ssl_context: Optional[ssl.SSLContext] = None
                if self._ctx.secure_reader_settings.cafile is not None:
                    ssl_context = ssl.SSLContext(protocol=ssl.PROTOCOL_TLS)
                    ssl_context.load_verify_locations(cafile=self._ctx.secure_reader_settings.cafile)
                conn = aiohttp.TCPConnector(ssl=ssl_context)
            else:
                socket_path = self._ctx.secure_reader_settings.socket
                secure_reader_endpoint = f"http+unix://{urllib.parse.quote_plus(socket_path)}"
                conn = aiohttp.UnixConnector(path=socket_path)

            async with aiohttp.ClientSession(connector=conn, loop=loop) as session:
                with aiohttp.MultipartWriter() as mpwriter:
                    part = mpwriter.append(file_obj)
                    part.set_content_disposition("form-data", name="file")
                    async with session.post(
                        url=f"{secure_reader_endpoint}/reader/excel",
                        data=mpwriter,
                    ) as resp:
                        file_data = await resp.json()

            for spreadsheet in file_data:
                sheetname = spreadsheet["sheetname"]
                sheetdata = spreadsheet["data"]
                source_status = FileProcessingStatus.in_progress
                has_header: Optional[bool] = None
                raw_schema: list[SchemaColumn] = []
                raw_schema_header: list[SchemaColumn] = []
                raw_schema_body: list[SchemaColumn] = []
                source_title = f"{dfile.filename} – {sheetname}"
                sheet_data_source = DataSource(
                    title=source_title,
                    raw_schema=raw_schema,
                    status=source_status,
                    error=None,
                )
                dfile.sources.append(sheet_data_source)

                if not sheetdata:
                    sheet_data_source.error = FileProcessingError.from_exception(exc.EmptyDocument())
                    sheet_data_source.status = FileProcessingStatus.failed
                else:
                    has_header, raw_schema, raw_schema_header, raw_schema_body = guess_header_and_schema_excel(
                        sheetdata
                    )
                sheet_settings = None

                if sheet_data_source.is_applicable:

                    def data_iter() -> Iterator[list]:
                        row_iter = iter(sheetdata)
                        for row in row_iter:
                            values = [cell["value"] for cell in row]
                            yield values

                    data_stream = SimpleUntypedDataStream(
                        data_iter=data_iter(),
                        rows_to_copy=None,  # TODO
                    )
                    with S3JsonEachRowUntypedFileDataSink(
                        s3=s3.get_sync_client(),
                        s3_key=sheet_data_source.s3_key,
                        bucket_name=s3.tmp_bucket_name,
                    ) as data_sink:
                        data_sink.dump_data_stream(data_stream)

                    assert has_header is not None
                    sheet_settings = ExcelFileSourceSettings(
                        first_line_is_header=has_header,
                        raw_schema_header=raw_schema_header,
                        raw_schema_body=raw_schema_body,
                    )

                sheet_data_source.raw_schema = raw_schema
                sheet_data_source.file_source_settings = sheet_settings

            await dfile.save()
            LOGGER.info("DataFile object saved.")

            task_processor = self._ctx.make_task_processor(self._request_id)
            parse_file_task = task_interface.ParseFileTask(file_id=dfile.id)
            await task_processor.schedule(parse_file_task)
            LOGGER.info(f"Scheduled ParseFileTask for file_id {dfile.id}")

        except Exception as ex:
            LOGGER.exception(ex)
            if dfile is None:
                return Retry(attempts=3)
            else:
                dfile.status = FileProcessingStatus.failed
                exc_to_save = ex if isinstance(ex, exc.DLFileUploaderBaseError) else exc.ParseFailed()
                dfile.error = FileProcessingError.from_exception(exc_to_save)
                await dfile.save()
                return Fail()
        return Success()

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/tasks/excel.py
datalens-tech/datalens-backend
import logging
from typing import (
    Optional,
    Type,
)

import attr

from dl_constants.enums import FileProcessingStatus
from dl_file_uploader_lib import exc
from dl_file_uploader_lib.common_locks import release_source_update_locks
from dl_file_uploader_lib.enums import FileType
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    FileProcessingError,
)
from dl_file_uploader_task_interface.context import FileUploaderTaskContext
import dl_file_uploader_task_interface.tasks as task_interface
from dl_file_uploader_task_interface.tasks import TaskExecutionMode
from dl_file_uploader_worker_lib.utils.connection_error_tracker import FileConnectionDataSourceErrorTracker
from dl_file_uploader_worker_lib.utils.file_parser import (
    CSVFileParser,
    FileParser,
    SpreadsheetFileParser,
)
from dl_task_processor.task import (
    BaseExecutorTask,
    Fail,
    Retry,
    Success,
    TaskResult,
)


LOGGER = logging.getLogger(__name__)


@attr.s
class ParseFileTask(BaseExecutorTask[task_interface.ParseFileTask, FileUploaderTaskContext]):
    cls_meta = task_interface.ParseFileTask

    FILE_TYPE_TO_PARSER_MAP: dict[FileType, Type[FileParser]] = {
        FileType.csv: CSVFileParser,
        FileType.gsheets: SpreadsheetFileParser,
        FileType.xlsx: SpreadsheetFileParser,
    }

    async def run(self) -> TaskResult:
        dfile: Optional[DataFile] = None
        usm = self._ctx.get_async_usm()
        task_processor = self._ctx.make_task_processor(self._request_id)
        redis = self._ctx.redis_service.get_redis()
        connection_error_tracker = FileConnectionDataSourceErrorTracker(usm, task_processor, redis, self._request_id)
        try:
            LOGGER.info(f"ParseFileTask. Mode: {self.meta.exec_mode.name}. File: {self.meta.file_id}")

            redis = self._ctx.redis_service.get_redis()
            rmm = RedisModelManager(redis=redis, crypto_keys_config=self._ctx.crypto_keys_config)
            dfile = await DataFile.get(manager=rmm, obj_id=self.meta.file_id)
            assert dfile is not None

            assert dfile.file_type is not None
            file_parser_cls = self.FILE_TYPE_TO_PARSER_MAP[dfile.file_type]

            file_parser = file_parser_cls(
                dfile=dfile,
                s3=self._ctx.s3_service,
                tpe=self._ctx.tpe,
                file_settings=self.meta.file_settings,
                source_id=self.meta.source_id,
            )

            for dsrc in file_parser.ensure_sources_and_return():
                if not dsrc.is_applicable:
                    if dsrc.error is None:  # all existing errors should've been saved by now
                        dsrc.error = FileProcessingError.from_exception(exc.InvalidSource())
                        connection_error_tracker.add_error(dsrc.id, dsrc.error)
                    continue

                try:
                    has_header, raw_schema, file_settings, dsrc_settings = await file_parser.guess_header_and_schema(
                        dsrc
                    )
                except exc.DLFileUploaderBaseError as e:
                    dsrc.status = FileProcessingStatus.failed
                    dsrc.error = FileProcessingError.from_exception(e)
                    connection_error_tracker.add_error(dsrc.id, dsrc.error)
                    continue

                dfile.file_settings = file_settings
                dsrc.raw_schema = raw_schema
                dsrc.file_source_settings = dsrc_settings

                preview = await file_parser.prepare_preview(dsrc, rmm)
                await preview.save(
                    ttl=12 * 12 * 60
                )  # save preview temporarily, so it is deleted if source never gets saved
                dsrc.preview_id = preview.id

                LOGGER.info("DataSourcePreview object saved.")

                dsrc.status = FileProcessingStatus.ready

            dfile.status = FileProcessingStatus.ready
            if dfile.file_type == FileType.csv and (csv_src := dfile.sources[0]).status == FileProcessingStatus.failed:  # type: ignore
                # there is guaranteed a single source in a csv file, and the error must be saved on the top level
                dfile.status = csv_src.status
                dfile.error = csv_src.error

            await connection_error_tracker.finalize(self.meta.exec_mode, self.meta.connection_id)
            await dfile.save()
            LOGGER.info("DataFile object saved.")

            if self.meta.exec_mode == TaskExecutionMode.UPDATE_AND_SAVE:
                for dsrc in dfile.sources or ():
                    if not dsrc.is_applicable:
                        continue

                    assert self.meta.tenant_id is not None and self.meta.connection_id is not None
                    save_source_task = task_interface.SaveSourceTask(
                        tenant_id=self.meta.tenant_id,
                        file_id=dfile.id,
                        src_source_id=dsrc.id,
                        dst_source_id=dsrc.id,
                        connection_id=self.meta.connection_id,
                        exec_mode=self.meta.exec_mode,
                    )
                    await task_processor.schedule(save_source_task)
                    LOGGER.info(f"Scheduled SaveSourceTask for file_id {dfile.id}")
            elif self.meta.exec_mode == TaskExecutionMode.UPDATE_NO_SAVE and dfile.sources:
                await release_source_update_locks(redis, *[dsrc.id for dsrc in dfile.sources])

        except Exception as ex:
            LOGGER.exception(ex)
            if dfile is None:
                return Retry(attempts=3)
            else:
                dfile.status = FileProcessingStatus.failed
                exc_to_save = ex if isinstance(ex, exc.DLFileUploaderBaseError) else exc.ParseFailed()
                dfile.error = FileProcessingError.from_exception(exc_to_save)
                await dfile.save()

                for src in dfile.sources or ():
                    connection_error_tracker.add_error(src.id, dfile.error)
                await connection_error_tracker.finalize(self.meta.exec_mode, self.meta.connection_id)
                return Fail()
        finally:
            await usm.close()
        return Success()

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/tasks/parse.py
datalens-tech/datalens-backend
import asyncio
import datetime
import logging
from typing import (
    Any,
    Optional,
)

import attr
from redis.asyncio.lock import Lock as RedisLock
import shortuuid

from dl_constants.enums import (
    DataSourceRole,
    FileProcessingStatus,
)
from dl_core import exc as core_exc
from dl_core.aio.web_app_services.s3 import S3Service
from dl_core.connection_executors import ConnExecutorQuery
from dl_core.db import SchemaColumn
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_file_uploader_lib.common_locks import release_source_update_locks
from dl_file_uploader_lib.enums import FileType
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    DataSource,
    DataSourcePreview,
    GSheetsFileSourceSettings,
    GSheetsUserSourceDataSourceProperties,
    GSheetsUserSourceProperties,
    PreviewSet,
)
from dl_file_uploader_task_interface.context import FileUploaderTaskContext
import dl_file_uploader_task_interface.tasks as task_interface
from dl_file_uploader_task_interface.tasks import TaskExecutionMode
from dl_file_uploader_worker_lib.utils.s3_utils import (
    S3Object,
    copy_from_s3_to_s3,
    make_s3_table_func_sql_source,
)
from dl_task_processor.task import (
    BaseExecutorTask,
    Fail,
    Retry,
    Success,
    TaskResult,
)
from dl_utils.aio import ContextVarExecutor

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection


LOGGER = logging.getLogger(__name__)


def _make_source_s3_filename(tenant_id: str) -> str:
    return f"{tenant_id}_{shortuuid.uuid()}"


def _make_tmp_source_s3_filename(source_id: str) -> str:
    return f"{source_id}-{shortuuid.uuid()}.json"


async def _prepare_file(
    s3_service: S3Service,
    tpe: ContextVarExecutor,
    dfile: DataFile,
    src_source: DataSource,
    dst_source_id: str,
    conn_raw_schema: list[SchemaColumn],
) -> str:
    src_filename = dfile.s3_key if dfile.file_type == FileType.csv else src_source.s3_key

    tmp_s3_filename = _make_tmp_source_s3_filename(dst_source_id)

    def do_s3_copy_sync() -> None:
        assert src_source.file_source_settings is not None and dfile.file_type is not None
        copy_from_s3_to_s3(
            s3_sync_cli=s3_service.get_sync_client(),
            src_file=S3Object(s3_service.tmp_bucket_name, src_filename),
            dst_file=S3Object(s3_service.tmp_bucket_name, tmp_s3_filename),
            file_type=dfile.file_type,
            file_settings=dfile.file_settings,
            file_source_settings=src_source.file_source_settings,
            raw_schema=conn_raw_schema,
        )

    loop = asyncio.get_running_loop()
    LOGGER.info("Copying data from user file to tmp s3 file (in JSONCompactEachRow format)")
    await loop.run_in_executor(tpe, do_s3_copy_sync)

    return tmp_s3_filename


def _get_conn_specific_dsrc_params(dfile: DataFile, src: DataSource) -> dict[str, Any]:
    kwargs = {}
    if dfile.file_type == FileType.gsheets:
        file_source_settings = src.file_source_settings
        assert isinstance(file_source_settings, GSheetsFileSourceSettings)
        user_source_properties = dfile.user_source_properties
        assert isinstance(user_source_properties, GSheetsUserSourceProperties)
        user_source_dsrc_properties = src.user_source_dsrc_properties
        assert isinstance(user_source_dsrc_properties, GSheetsUserSourceDataSourceProperties)

        kwargs.update(
            dict(
                spreadsheet_id=user_source_properties.spreadsheet_id,
                sheet_id=user_source_dsrc_properties.sheet_id,
                first_line_is_header=file_source_settings.first_line_is_header,
            )
        )
    return kwargs


@attr.s
class SaveSourceTask(BaseExecutorTask[task_interface.SaveSourceTask, FileUploaderTaskContext]):
    cls_meta = task_interface.SaveSourceTask

    async def run(self) -> TaskResult:
        try:
            LOGGER.info(f"SaveSourceTask. Mode: {self.meta.exec_mode.name}. File: {self.meta.file_id}")
            if self.meta.exec_mode == TaskExecutionMode.UPDATE_NO_SAVE:
                LOGGER.warning(f"Cannot run SaveSourceTask in {self.meta.exec_mode.name} mode, failing the task")
                return Fail()

            redis = self._ctx.redis_service.get_redis()
            rmm = RedisModelManager(redis=redis, crypto_keys_config=self._ctx.crypto_keys_config)
            dfile = await DataFile.get(manager=rmm, obj_id=self.meta.file_id)

            src_source_id, dst_source_id = self.meta.src_source_id, self.meta.dst_source_id
            src_source = dfile.get_source_by_id(src_source_id)

            # TODO: init all this stuff in a proper place, not in task
            rci = self._ctx.get_rci()
            usm = self._ctx.get_async_usm(rci=rci)
            service_registry = self._ctx.get_service_registry(rci=rci)
            release_update_source_lock_flag = False
            async with usm:
                source_lock_key = f"SaveSourceTask/{dst_source_id}"
                LOGGER.info(f"Acquiring redis lock {source_lock_key}")
                async with RedisLock(redis, name=source_lock_key, timeout=120, blocking_timeout=120):
                    LOGGER.info(f"Lock {source_lock_key} acquired")
                    assert isinstance(usm, AsyncUSManager)
                    conn = await usm.get_by_id(self.meta.connection_id, expected_type=BaseFileS3Connection)
                    assert isinstance(conn, BaseFileS3Connection)

                    preview = await DataSourcePreview.get(manager=rmm, obj_id=str(src_source.preview_id))
                    await preview.save(ttl=None)  # now that the source is saved the preview can be saved without ttl
                    preview_set = PreviewSet(redis=redis, id=self.meta.tenant_id)
                    await preview_set.add(preview.id)

                    try:
                        conn_file_source = conn.get_file_source_by_id(dst_source_id)
                        if (
                            conn_file_source.status == FileProcessingStatus.ready
                            and self.meta.exec_mode != TaskExecutionMode.UPDATE_AND_SAVE
                        ):
                            LOGGER.info(
                                f'Source {dst_source_id} in connection {conn.uuid} has status "ready". Finishing task.'
                            )
                            return Success()
                        conn_raw_schema: list[SchemaColumn] = conn_file_source.raw_schema or []
                        orig_file_id = conn_file_source.file_id
                        orig_s3_filename = conn_file_source.s3_filename
                        orig_preview_id = conn_file_source.preview_id
                    except core_exc.SourceDoesNotExist:
                        LOGGER.info(f"Source {dst_source_id} not present in connection {conn.uuid}. Finishing task.")
                        return Success()

                    s3_service = self._ctx.s3_service

                    raw_schema_override: Optional[list[SchemaColumn]]
                    if self.meta.exec_mode == TaskExecutionMode.UPDATE_AND_SAVE:
                        raw_schema_override = src_source.raw_schema
                    else:
                        raw_schema_override = None

                    tmp_s3_filename = await _prepare_file(
                        s3_service,
                        self._ctx.tpe,
                        dfile,
                        src_source,
                        dst_source_id,
                        raw_schema_override if raw_schema_override is not None else conn_raw_schema,
                    )

                    new_s3_filename = _make_source_s3_filename(tenant_id=self.meta.tenant_id)

                    def _construct_insert_from_select_query(for_debug: bool = False) -> str:
                        src_sql = make_s3_table_func_sql_source(
                            conn=conn,
                            source_id=dst_source_id,
                            bucket=s3_service.tmp_bucket_name,
                            filename=tmp_s3_filename,
                            file_fmt="JSONCompactEachRow",
                            for_debug=for_debug,
                            raw_schema_override=raw_schema_override,
                        )
                        dst_sql = make_s3_table_func_sql_source(
                            conn=conn,
                            source_id=dst_source_id,
                            bucket=s3_service.persistent_bucket_name,
                            filename=new_s3_filename,
                            file_fmt="Native",
                            for_debug=for_debug,
                            raw_schema_override=raw_schema_override,
                        )
                        return f"insert into FUNCTION {dst_sql} select * from {src_sql}"

                    ce_query = ConnExecutorQuery(
                        query=_construct_insert_from_select_query(),
                        debug_compiled_query=_construct_insert_from_select_query(for_debug=True),
                        trusted_query=True,
                        is_ddl_dml_query=True,
                    )

                    LOGGER.info(
                        "Running INSERT FROM SELECT clickhouse query " "to transform S3 data to CH Native format."
                    )
                    conn_executor_factory = service_registry.get_conn_executor_factory()
                    async_conn_executor = conn_executor_factory.get_async_conn_executor(conn)
                    exec_result = await async_conn_executor.execute(ce_query)
                    LOGGER.info(exec_result)

                    async with usm.locked_entry_cm(
                        self.meta.connection_id,
                        expected_type=BaseFileS3Connection,
                        wait_timeout_sec=60,
                        duration_sec=10,
                    ) as conn:
                        assert isinstance(conn, BaseFileS3Connection)
                        try:
                            conn.get_file_source_by_id(id=dst_source_id)
                            conn_file_source = conn.get_file_source_by_id(id=dst_source_id)
                        except core_exc.SourceDoesNotExist:
                            LOGGER.info(
                                f"Source {dst_source_id} not present in connection {conn.uuid}. Finishing task."
                            )
                            conn.data.component_errors.remove_errors(id=dst_source_id)
                            return Success()

                        if (
                            conn_file_source.status == FileProcessingStatus.ready
                            and self.meta.exec_mode != TaskExecutionMode.UPDATE_AND_SAVE
                        ):
                            LOGGER.info(
                                f'Source {dst_source_id} in connection {conn.uuid} already has status "ready". '
                                "Something strange, considering present lock. "
                                "Failing task."
                            )
                            return Fail()

                        extra_dsrc_params = _get_conn_specific_dsrc_params(dfile, src_source)
                        if self.meta.exec_mode == TaskExecutionMode.UPDATE_AND_SAVE:
                            extra_dsrc_params["data_updated_at"] = datetime.datetime.now(datetime.timezone.utc)
                            release_update_source_lock_flag = True
                            extra_dsrc_params["raw_schema"] = raw_schema_override
                            extra_dsrc_params["file_id"] = dfile.id

                        conn.update_data_source(
                            dst_source_id,
                            role=DataSourceRole.origin,
                            s3_filename=new_s3_filename,
                            status=FileProcessingStatus.ready,
                            preview_id=preview.id,
                            **extra_dsrc_params,
                        )
                        conn.data.component_errors.remove_errors(id=dst_source_id)
                        await usm.save(conn)

                    # sync source id with the connection to enable consistent work with dfile (e.g. polling)
                    src_source.id = dst_source_id
                    await dfile.save()

                    if orig_file_id != self.meta.file_id and orig_s3_filename is not None:
                        if self.meta.exec_mode == TaskExecutionMode.UPDATE_AND_SAVE:
                            delete_file_task = task_interface.DeleteFileTask(
                                s3_filename=orig_s3_filename,
                                preview_id=orig_preview_id,
                            )
                            task_processor = self._ctx.make_task_processor(self._request_id)
                            await task_processor.schedule(delete_file_task)
                        else:
                            LOGGER.warning(
                                f"file_id in connection source ({orig_file_id}) and passed file_id"
                                f" ({self.meta.file_id}) are different and execution mode is"
                                f" {self.meta.exec_mode.name}, expected it to be"
                                f" {TaskExecutionMode.UPDATE_AND_SAVE.name}"
                            )

                LOGGER.info(f"Released lock {source_lock_key}")

                if release_update_source_lock_flag:
                    await release_source_update_locks(redis, dst_source_id)

        except Exception as ex:
            LOGGER.exception(ex)
            return Retry(attempts=5)
        return Success()

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/tasks/save.py
datalens-tech/datalens-backend
import attr

from dl_api_commons.tenant_resolver import (
    CommonTenantResolver,
    TenantResolver,
)
from dl_file_uploader_worker_lib.app import FileUploaderWorkerFactory
from dl_file_uploader_worker_lib.settings import FileUploaderWorkerSettings


@attr.s(kw_only=True)
class TestingFileUploaderWorkerFactory(FileUploaderWorkerFactory[FileUploaderWorkerSettings]):
    def _get_tenant_resolver(self) -> TenantResolver:
        return CommonTenantResolver()

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/testing/app_factory.py
datalens-tech/datalens-backend
import asyncio
import contextlib
import logging
from typing import AsyncGenerator

import arq

from dl_file_uploader_worker_lib.app import FileUploaderContextFab
from dl_file_uploader_worker_lib.settings import FileUploaderWorkerSettings
from dl_file_uploader_worker_lib.tasks import REGISTRY
from dl_file_uploader_worker_lib.testing.app_factory import TestingFileUploaderWorkerFactory
from dl_task_processor.arq_wrapper import (
    create_arq_redis_settings,
    create_redis_pool,
)
from dl_task_processor.executor import Executor
from dl_task_processor.processor import (
    ARQProcessorImpl,
    LocalProcessorImpl,
    TaskProcessor,
)
from dl_task_processor.state import TaskState
from dl_task_processor.worker import (
    ArqWorker,
    ArqWorkerTestWrapper,
)


LOGGER = logging.getLogger(__name__)


@contextlib.asynccontextmanager
async def task_processor_arq_worker(
    loop: asyncio.AbstractEventLoop,
    task_state: TaskState,
    file_uploader_worker_settings: FileUploaderWorkerSettings,
) -> AsyncGenerator[ArqWorker, None]:
    LOGGER.info("Set up worker")
    worker = TestingFileUploaderWorkerFactory(settings=file_uploader_worker_settings).create_worker(state=task_state)
    wrapper = ArqWorkerTestWrapper(loop=loop, worker=worker)
    try:
        yield await wrapper.start()
    finally:
        await wrapper.stop()


@contextlib.asynccontextmanager
async def task_processor_arq_client(
    loop: asyncio.AbstractEventLoop,
    redis_pool: arq.ArqRedis,
    task_state: TaskState,
    file_uploader_worker_settings: FileUploaderWorkerSettings,
) -> AsyncGenerator[TaskProcessor, None]:
    async with task_processor_arq_worker(loop, task_state, file_uploader_worker_settings):
        impl = ARQProcessorImpl(redis_pool)
        p = TaskProcessor(impl=impl, state=task_state)
        LOGGER.info("Arq TP is ready")
        yield p


@contextlib.asynccontextmanager
async def task_processor_local_client(
    task_state: TaskState,
    file_uploader_worker_settings: FileUploaderWorkerSettings,
) -> AsyncGenerator[TaskProcessor, None]:
    context_fab = FileUploaderContextFab(file_uploader_worker_settings)
    context = await context_fab.make()
    executor = Executor(context=context, state=task_state, registry=REGISTRY)
    impl = LocalProcessorImpl(executor)
    processor = TaskProcessor(impl=impl, state=task_state)
    LOGGER.info("Local TP is ready")
    try:
        yield processor
    finally:
        await context_fab.tear_down(context)


@contextlib.asynccontextmanager
async def get_task_processor_client(
    client_type: str,
    loop: asyncio.AbstractEventLoop,
    task_state: TaskState,
    file_uploader_worker_settings: FileUploaderWorkerSettings,
) -> AsyncGenerator[TaskProcessor, None]:
    if client_type == "arq":
        arq_redis_settings = create_arq_redis_settings(file_uploader_worker_settings.REDIS_ARQ)
        pool = await create_redis_pool(arq_redis_settings)
        async with task_processor_arq_client(loop, pool, task_state, file_uploader_worker_settings) as tp:
            try:
                yield tp
            finally:
                await pool.close()
    elif client_type == "local":
        async with task_processor_local_client(task_state, file_uploader_worker_settings) as tp:
            yield tp

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/testing/task_processor_client.py
datalens-tech/datalens-backend
import datetime
import logging
from typing import Optional

import attr
import redis.asyncio

from dl_constants.enums import (
    ComponentErrorLevel,
    ComponentType,
    DataSourceRole,
    FileProcessingStatus,
)
from dl_core import exc as core_exc
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_file_uploader_lib.common_locks import release_source_update_locks
from dl_file_uploader_lib.redis_model.models import FileProcessingError
from dl_file_uploader_task_interface.tasks import (
    DeleteFileTask,
    TaskExecutionMode,
)
from dl_task_processor.processor import TaskProcessor

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection


LOGGER = logging.getLogger(__name__)


@attr.s
class FileConnectionDataSourceErrorTracker:
    """
    Stores errors that need to be dumped into connection, fails connection sources on `finalize`
    """

    _usm: AsyncUSManager = attr.ib()
    _task_processor: TaskProcessor = attr.ib()
    _redis: redis.asyncio.Redis = attr.ib()
    _request_id: Optional[str] = attr.ib(default=None)

    _error_registry: dict[str, FileProcessingError] = attr.ib(init=False, factory=dict)

    def add_error(self, source_id: str, err: FileProcessingError) -> None:
        """
        :param source_id: id of a source in a connection
        :param err: error to store, if an error for this source is already stored, do nothing
        """

        if (existing_error := self._error_registry.get(source_id)) is not None:
            LOGGER.warning(
                f"Attempt to overwrite existing error for source {source_id}"
                f"({existing_error.code}) with a new one: {err.code}"
            )
            return
        err.details["request-id"] = self._request_id
        self._error_registry[source_id] = err

    def clear(self) -> None:
        self._error_registry.clear()

    async def _release_source_update_locks(self) -> None:
        await release_source_update_locks(self._redis, *self._error_registry.keys())

    async def _fail_connection_sources(self, connection_id: str) -> None:
        delete_tasks: list[DeleteFileTask] = []

        async with self._usm.locked_entry_cm(
            connection_id,
            expected_type=BaseFileS3Connection,
            wait_timeout_sec=60,
            duration_sec=10,
        ) as conn:
            assert isinstance(conn, BaseFileS3Connection)
            for source_id, error in self._error_registry.items():
                reason = ".".join(error.code)
                try:
                    err_level = ComponentErrorLevel(error.level.name)
                except ValueError:
                    err_level = ComponentErrorLevel.error

                try:
                    src = conn.get_file_source_by_id(source_id)
                except core_exc.SourceDoesNotExist:
                    LOGGER.info(f"Source {source_id} not present in connection {conn.uuid}. No need to fail source.")
                    conn.data.component_errors.remove_errors(id=source_id)
                    continue

                LOGGER.info(
                    f"Going to fail source id={source_id} in connection id={connection_id}"
                    f' due to reason "{reason}" with level "{err_level}"'
                )

                if err_level == ComponentErrorLevel.error:
                    if src.s3_filename is not None:
                        delete_tasks.append(
                            DeleteFileTask(
                                s3_filename=src.s3_filename,
                                preview_id=src.preview_id,
                            )
                        )

                    conn.update_data_source(
                        source_id,
                        role=DataSourceRole.origin,
                        s3_filename=None,
                        status=FileProcessingStatus.failed,
                        preview_id=None,
                        data_updated_at=datetime.datetime.now(datetime.timezone.utc),
                    )
                else:
                    conn.update_data_source(
                        source_id,
                        role=DataSourceRole.origin,
                        status=FileProcessingStatus.ready,
                        data_updated_at=datetime.datetime.now(datetime.timezone.utc),
                    )

                conn.data.component_errors.remove_errors(id=source_id)  # there can only be one error at a time
                conn.data.component_errors.add_error(
                    id=source_id,
                    type=ComponentType.data_source,
                    message=error.message,
                    code=error.code,
                    level=err_level,
                    details=error.details,
                )

            await self._usm.save(conn)

        for delete_file_task in delete_tasks:
            await self._task_processor.schedule(delete_file_task)
            LOGGER.info(f"Scheduled task DeleteFileTask for source_id {src.id}, filename {src.s3_filename}")

    async def finalize(self, mode: TaskExecutionMode, connection_id: Optional[str] = None) -> None:
        """
        Deals with errors according to `mode`:
            - BASIC - do nothing
            - UPDATE_NO_SAVE - release locks (does not require connection_id)
            - UPDATE_AND_SAVE - update connection with new errors and release locks
        """

        if mode not in (TaskExecutionMode.UPDATE_AND_SAVE, TaskExecutionMode.UPDATE_NO_SAVE):
            LOGGER.info(f"Nothing to do with errors for mode {mode}")
            return
        if mode == TaskExecutionMode.UPDATE_AND_SAVE:
            assert connection_id is not None, "No connection_id to fail sources in"
            await self._fail_connection_sources(connection_id)
        await self._release_source_update_locks()
        self.clear()

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/utils/connection_error_tracker.py
datalens-tech/datalens-backend
from __future__ import annotations

import copy
import datetime
from functools import lru_cache
from itertools import islice
import logging
import re
import string
from typing import (
    Any,
    Callable,
    Iterator,
    Optional,
    Sequence,
    Union,
)
from uuid import uuid4

import attr

from dl_app_tools.profiling_base import (
    GenericProfiler,
    generic_profiler,
)
from dl_constants.enums import UserDataType
from dl_core import converter_types_cast
from dl_core.aio.web_app_services.gsheets import (
    Cell,
    NumberFormatType,
)
from dl_core.db import SchemaColumn


LOGGER = logging.getLogger(__name__)


class CsvParsingException(Exception):
    pass


DATE_RE = r"""
(?:
  [1-9]\d{3}-
  (?:
    (?:0[1-9]|1[0-2])
    -
    (?:0[1-9]|1\d|2[0-8])
    |
    (?:0[13-9]|1[0-2])
    -
    (?:29|30)
    |
    (?:0[13578]|1[02])
    -31
  )
|
  (?:
    [1-9]\d
    (?:0[48]|[2468][048]|[13579][26])
    |
    (?:[2468][048]|[13579][26])
    00
  )
  -02-29
)
"""


DATE_RU_FMT_RE = r"""
(?:
  (?:
    (?:0[1-9]|1\d|2[0-8])
    \.
    (?:0[1-9]|1[0-2])
    |
    (?:29|30)
    \.
    (?:0[13-9]|1[0-2])
    |
    31\.
    (?:0[13578]|1[02])
  )
  \.[1-9]\d{3}
|
  29\.02\.
  (?:
    [1-9]\d
    (?:0[48]|[2468][048]|[13579][26])
    |
    (?:[2468][048]|[13579][26])
    00
  )
)
"""


TIME_RE = r"""
(?:[01]\d|2[0-3])
:
[0-5]\d
:
[0-5]\d
(\.\d{1,6})?
"""


DATETIME_TZ_RE = r"""
(?:
  (Z|UTC)
|
  [+-][01]\d
  (:?[0-5]\d)?
)
"""


DATETIME_REGEXES = {
    "datetime__mid_space_without_tz": re.compile(rf"^{DATE_RE}\s{TIME_RE}$", flags=re.VERBOSE),
    "datetime__mid_space_with_tz": re.compile(rf"^{DATE_RE}\s{TIME_RE}\s?{DATETIME_TZ_RE}$", flags=re.VERBOSE),
    "datetime__full_with_tz": re.compile(rf"^{DATE_RE}T{TIME_RE}\s?{DATETIME_TZ_RE}$", flags=re.VERBOSE),
    "datetime__mid_t_without_tz": re.compile(rf"^{DATE_RE}T{TIME_RE}$", flags=re.VERBOSE),
    "datetime__ru_fmt_mid_space_without_tz": re.compile(rf"^{DATE_RU_FMT_RE}\s{TIME_RE}$", flags=re.VERBOSE),
}


@lru_cache(2**15)
def _check_datetime_re(value):  # type: ignore  # TODO: fix
    if not value:
        return None
    for desc, re_expr in DATETIME_REGEXES.items():
        try:
            if re_expr.match(value):
                return desc
        except TypeError:
            return None
    return None


date_re = {
    "date__iso_date": re.compile(rf"^{DATE_RE}$", flags=re.VERBOSE),
    "date__reversed_full_year_slashes": re.compile(
        r"""
^
(
  (
    29
    /
    0?2
    /
    (?:
      [1-9]\d
      (?:0[48]|[2468][048]|[13579][26])
    |
      (?:[2468][048]|[13579][26])
      00
    )
  )
|
  (
    (?:0?[1-9]|1\d|2[0-8])
    /
    (?:0?[1-9]|1[0-2])
  |
    (?:29|30)
    /
    (?:0?[13-9]|1[0-2])
  |
    31
    /
    (?:0?[13578]|1[02])
  )
  /
  [1-9]\d{3}
)
$
""",
        flags=re.VERBOSE,
    ),
    "date__reversed_full_year_dots": re.compile(
        r"""
^
(
  (
    29
    \.
    0?2
    \.
    (?:
      [1-9]\d
      (?:0[48]|[2468][048]|[13579][26])
    |
      (?:[2468][048]|[13579][26])
      00
    )
  )
|
  (
    (?:0?[1-9]|1\d|2[0-8])
    \.
    (?:0?[1-9]|1[0-2])
  |
    (?:29|30)
    \.
    (?:0?[13-9]|1[0-2])
  |
    31
    \.
    (?:0?[13578]|1[02])
  )
  \.
  [1-9]\d{3}
)
$
""",
        flags=re.VERBOSE,
    ),
    "date__reversed_short_year_slashes": re.compile(
        r"""
^
(
  (
    29
    \.
    02
    \.
    (?:
      (?:00|0[48]|[2468][048]|[13579][26])
    )
  )
|
  (
    (?:0?[1-9]|1\d|2[0-8])
    \.
    (?:0?[1-9]|1[0-2])
  |
    (?:29|30)
    \.
    (?:0?[13-9]|1[0-2])
  |
    31
    \.
    (?:0?[13578]|1[02])
  )
  \.
  \d{2}
)
$
""",
        flags=re.VERBOSE,
    ),
    "date__reversed_short_year_dots": re.compile(
        r"""
^
(
  (
    29
    /
    02
    /
    (?:
      (?:00|0[48]|[2468][048]|[13579][26])
    )
  )
|
  (
    (?:0?[1-9]|1\d|2[0-8])
    /
    (?:0?[1-9]|1[0-2])
  |
    (?:29|30)
    /
    (?:0?[13-9]|1[0-2])
  |
    31
    /
    (?:0?[13578]|1[02])
  )
  /
  \d{2}
)
$
""",
        flags=re.VERBOSE,
    ),
}


@lru_cache(2**15)
def _check_date_re(value):  # type: ignore  # TODO: fix
    if not value:
        return None
    for desc, re_expr in date_re.items():
        try:
            if re_expr.match(value):
                return desc
        except TypeError:
            return None
    return None


_TYPE_GROUP_NONE = 0
_TYPE_GROUP_BOOLEAN = 1
_TYPE_GROUP_NUMBER = 2
_TYPE_GROUP_DATETIME = 3
_TYPE_GROUP_STRING = 4


@attr.s(frozen=True)
class ParsingDataType:
    bi_type: UserDataType = attr.ib()
    type: Any = attr.ib()
    order: int = attr.ib()
    group: int = attr.ib()

    cast_func: Callable[[Any], Any] = attr.ib(default=None)
    check_func: Callable[[Any], Any] = attr.ib(default=None)
    format_desc: str = attr.ib(default=None)


_NONE_PARSING_DATA_TYPE = ParsingDataType(bi_type=UserDataType.string, type=str, order=-2, group=_TYPE_GROUP_NONE)
_BOOLEAN_PARSING_DATA_TYPE = ParsingDataType(
    bi_type=UserDataType.boolean,
    type=bool,
    order=-1,
    cast_func=converter_types_cast._to_boolean,
    group=_TYPE_GROUP_BOOLEAN,
)
_INTEGER_PARSING_DATA_TYPE = ParsingDataType(
    bi_type=UserDataType.integer, type=int, order=0, cast_func=converter_types_cast._to_int, group=_TYPE_GROUP_NUMBER
)
_FLOAT_PARSING_DATA_TYPE = ParsingDataType(
    bi_type=UserDataType.float, type=float, order=1, cast_func=converter_types_cast._to_float, group=_TYPE_GROUP_NUMBER
)
_DATE_PARSING_DATA_TYPE = ParsingDataType(
    bi_type=UserDataType.date,
    order=2,
    type=datetime.date,
    cast_func=converter_types_cast._to_date,
    check_func=_check_date_re,
    group=_TYPE_GROUP_DATETIME,
)
_DATETIME_PARSING_DATA_TYPE = ParsingDataType(
    bi_type=UserDataType.genericdatetime,
    order=3,
    type=datetime.datetime,
    cast_func=converter_types_cast._to_datetime,
    check_func=_check_datetime_re,
    group=_TYPE_GROUP_DATETIME,
)
_STRING_PARSING_DATA_TYPE = ParsingDataType(bi_type=UserDataType.string, type=str, order=4, group=_TYPE_GROUP_STRING)

ALLOWED_DATA_TYPES = (
    _INTEGER_PARSING_DATA_TYPE,
    _FLOAT_PARSING_DATA_TYPE,
    _DATE_PARSING_DATA_TYPE,
    _DATETIME_PARSING_DATA_TYPE,
    _STRING_PARSING_DATA_TYPE,
)
ALLOWED_DATA_TYPES_LEN = len(ALLOWED_DATA_TYPES)

TColumnTypes = dict[int, ParsingDataType]
TResultColumn = dict[str, Union[str, int]]
TResultTypes = list[TResultColumn]


@lru_cache(2**20)
def guess_cell_type(value, start_order: int = 0) -> ParsingDataType:  # type: ignore  # TODO: fix
    if value is None or value == "":
        return _NONE_PARSING_DATA_TYPE

    result = None
    start_order = max(start_order, 0)
    for cell_type in islice(ALLOWED_DATA_TYPES, start_order, ALLOWED_DATA_TYPES_LEN):
        if cell_type.check_func is not None:
            format_desc = cell_type.check_func(value)
            if format_desc:
                return attr.evolve(cell_type, format_desc=format_desc)
        else:
            try:
                func = cell_type.cast_func or cell_type.type
                # call validation func. Exception ==> wrong type
                func(value)
            except (ValueError, TypeError, OverflowError):
                pass
            else:
                return cell_type

    if not result:
        raise CsvParsingException("Unknown data type")  # TODO: ignore unknown cell types


def guess_cell_type_gsheet(cell: Cell, *args: Any, **kwargs: Any) -> ParsingDataType:
    value = cell.value
    number_format = cell.number_format_type
    if value is None or value == "":
        return _NONE_PARSING_DATA_TYPE

    number_format_to_parsing_data_type_map = {
        NumberFormatType.NUMBER_FORMAT_TYPE_UNSPECIFIED: _STRING_PARSING_DATA_TYPE,
        NumberFormatType.TEXT: _STRING_PARSING_DATA_TYPE,
        NumberFormatType.TIME: _STRING_PARSING_DATA_TYPE,
        NumberFormatType.DATE: _DATE_PARSING_DATA_TYPE,
        NumberFormatType.DATE_TIME: _DATETIME_PARSING_DATA_TYPE,
        NumberFormatType.INTEGER: _INTEGER_PARSING_DATA_TYPE,
        NumberFormatType.FLOAT: _FLOAT_PARSING_DATA_TYPE,
        NumberFormatType.BOOLEAN: _BOOLEAN_PARSING_DATA_TYPE,
    }
    return number_format_to_parsing_data_type_map[number_format]


def _choose_new_cell_type(new_type: ParsingDataType, prev_type: ParsingDataType) -> Optional[ParsingDataType]:
    if prev_type.group == _TYPE_GROUP_NONE:
        return new_type
    if new_type.group == _TYPE_GROUP_NONE:
        return None
    if new_type.group == prev_type.group:
        if new_type.format_desc is not None and new_type.format_desc != prev_type.format_desc:
            # If date/datetime values have different formats, assume cell type is string
            return _STRING_PARSING_DATA_TYPE
        if new_type.order > prev_type.order:
            return new_type
        return None
    else:
        return _STRING_PARSING_DATA_TYPE


def generate_column_hash_name():  # type: ignore  # TODO: fix
    return "f" + uuid4().hex[:16]


def guess_has_header(header_types: TColumnTypes, column_types: TColumnTypes) -> bool:
    has_header = False
    if len(header_types) >= len(column_types):
        header_type_gt_col = 0
        header_type_lt_col = 0
        for col_index, header_type in header_types.items():
            if len(column_types) <= col_index:
                continue

            if header_type.order < column_types[col_index].order:
                header_type_lt_col += 1
            elif header_type.order > column_types[col_index].order:
                header_type_gt_col += 1

        if header_type_lt_col == 0 and header_type_gt_col > 0:
            has_header = True
    return has_header


def merge_column_types(header_types: TColumnTypes, column_types: TColumnTypes, has_header: bool) -> TColumnTypes:
    if has_header:
        return column_types
    new_column_types = copy.deepcopy(column_types)
    for col_index, header_type in header_types.items():
        col_type = new_column_types[col_index]
        new_col_type = _choose_new_cell_type(new_type=header_type, prev_type=col_type)
        if new_col_type:
            new_column_types[col_index] = new_col_type
    return new_column_types


def raw_schema_to_column_types(raw_schema: list[SchemaColumn]) -> TColumnTypes:
    bi_type_to_parsing_data_type_map = {
        UserDataType.integer: _INTEGER_PARSING_DATA_TYPE,
        UserDataType.float: _FLOAT_PARSING_DATA_TYPE,
        UserDataType.date: _DATE_PARSING_DATA_TYPE,
        UserDataType.genericdatetime: _DATETIME_PARSING_DATA_TYPE,
        UserDataType.datetime: _DATETIME_PARSING_DATA_TYPE,
        UserDataType.string: _STRING_PARSING_DATA_TYPE,
        UserDataType.boolean: _BOOLEAN_PARSING_DATA_TYPE,
    }
    column_types = {}
    for col_index, col in enumerate(raw_schema):
        column_types[col_index] = bi_type_to_parsing_data_type_map.get(col.user_type, _STRING_PARSING_DATA_TYPE)
    return column_types


def idx_to_alphabet_notation(idx: int) -> str:
    letters = string.ascii_uppercase

    col_name = letters[idx % 26]
    idx = (idx // 26) - 1
    while idx >= 0:
        col_name = letters[idx % 26] + col_name
        idx = (idx // 26) - 1

    return col_name


def idx_to_file_notation(idx: int) -> str:
    return f"field{idx + 1}"


def make_result_types(
    header_values: Sequence[str | Cell],
    types: TColumnTypes,
    has_header: bool,
    missing_title_generator: Callable[[int], str] = idx_to_file_notation,
) -> TResultTypes:
    def get_value(value):  # type: ignore
        if isinstance(value, Cell):
            return str(value.value) if not value.empty else None
        return str(value)

    result_types = []
    for index, t in types.items():
        header_value = get_value(header_values[index]) if has_header and len(header_values) > index else ""
        result_types.append(
            dict(
                name=generate_column_hash_name(),
                index=index,
                cast=t.bi_type.name,
                title=header_value or missing_title_generator(index),
            )
        )
    result_types.sort(key=lambda t: t["index"])
    return result_types


def guess_column_types(
    data_iter: Iterator,
    sample_lines_count: Optional[int] = None,
    cell_type_guesser: Callable[..., ParsingDataType] = guess_cell_type,
) -> TColumnTypes:
    column_types: TColumnTypes = {}
    with GenericProfiler("guess_all_rows_types"):
        rows_seen = 0
        for row in data_iter:
            if sample_lines_count and rows_seen >= sample_lines_count:
                break
            rows_seen += 1

            for col_index, value in enumerate(row):
                prev_type = column_types.get(col_index)
                if prev_type:
                    new_type = cell_type_guesser(value, prev_type.order)
                    new_cell_type = _choose_new_cell_type(prev_type=prev_type, new_type=new_type)
                else:
                    new_cell_type = cell_type_guesser(value)

                if new_cell_type:
                    column_types[col_index] = new_cell_type

    if rows_seen == 0:
        raise CsvParsingException("Too few CSV rows received.")

    return column_types


@generic_profiler("guess_types_and_header")
def guess_types_and_header(
    data_iter: Iterator,
    has_header: Optional[bool] = None,
    sample_lines_count: Optional[int] = None,
    **kwargs: Any,
) -> tuple[bool, TResultTypes]:
    if has_header is None or has_header:
        header = next(data_iter)  # assume first row is header
    else:
        header = []

    header_types = {col_index: guess_cell_type(value) for col_index, value in enumerate(header)}
    column_types = guess_column_types(data_iter, sample_lines_count, guess_cell_type)

    if has_header is None:
        has_header = guess_has_header(header_types, column_types)

    merged_column_types = merge_column_types(header_types, column_types, has_header)
    result_merged_types = make_result_types(header, merged_column_types, has_header)

    return has_header, result_merged_types


@generic_profiler("guess_types_and_header_gsheets")
def guess_types_and_header_gsheets(
    data_iter: Iterator,
    sample_lines_count: Optional[int] = None,
) -> tuple[bool, TResultTypes, TResultTypes, TResultTypes]:
    header = next(data_iter)  # assume first row is header

    header_types = {col_index: guess_cell_type_gsheet(value) for col_index, value in enumerate(header)}
    column_types = guess_column_types(data_iter, sample_lines_count, guess_cell_type_gsheet)

    has_header = guess_has_header(header_types, column_types)

    merged_column_types = merge_column_types(header_types, column_types, has_header)

    result_merged_types = make_result_types(header, merged_column_types, has_header, idx_to_alphabet_notation)
    result_header_types = make_result_types(header, header_types, True, idx_to_alphabet_notation)
    result_column_types = make_result_types(header, column_types, False, idx_to_alphabet_notation)

    return has_header, result_merged_types, result_header_types, result_column_types


@generic_profiler("guess_types_gsheets")
def guess_types_gsheets(
    data_iter: Iterator,
    sample_lines_count: Optional[int] = None,
) -> TResultTypes:
    column_types = guess_column_types(data_iter, sample_lines_count, guess_cell_type_gsheet)

    header: list[str] = []
    result_column_types = make_result_types(header, column_types, False, idx_to_alphabet_notation)

    return result_column_types


def guess_cell_type_excel(cell: dict, *args: Any, **kwargs: Any) -> ParsingDataType:
    value = cell.get("value")
    number_format = cell.get("data_type")
    assert isinstance(number_format, str)

    number_format_to_parsing_data_type_map = {
        "s": _STRING_PARSING_DATA_TYPE,
        "d": _DATETIME_PARSING_DATA_TYPE,
        "i": _INTEGER_PARSING_DATA_TYPE,
        "n": _FLOAT_PARSING_DATA_TYPE,
    }

    if value is None or value == "":
        return _NONE_PARSING_DATA_TYPE

    # Fallback to string in case of type recognition errors
    return number_format_to_parsing_data_type_map.get(number_format, _STRING_PARSING_DATA_TYPE)


def make_excel_result_types(
    header_values: Sequence[dict],
    types: TColumnTypes,
    has_header: bool,
) -> TResultTypes:
    def get_value(value):  # type: ignore
        return value.get("value")

    result_types = []
    for index, t in types.items():
        header_value = get_value(header_values[index]) if has_header and len(header_values) > index else ""
        result_types.append(
            dict(
                name=generate_column_hash_name(),
                index=index,
                cast=t.bi_type.name,
                title=header_value or "field{}".format(index + 1),
            )
        )
    result_types.sort(key=lambda t: t["index"])
    return result_types


@generic_profiler("guess_types_and_header_excel")
def guess_types_and_header_excel(
    data_iter: Iterator,
    sample_lines_count: Optional[int] = None,
) -> tuple[bool, TResultTypes, TResultTypes, TResultTypes]:
    header = next(data_iter)  # assume first row is header

    header_types = {col_index: guess_cell_type_excel(value) for col_index, value in enumerate(header)}
    column_types = guess_column_types(data_iter, sample_lines_count, guess_cell_type_excel)

    has_header = guess_has_header(header_types, column_types)

    merged_column_types = merge_column_types(header_types, column_types, has_header)

    result_merged_types = make_excel_result_types(header, merged_column_types, has_header)
    result_header_types = make_excel_result_types(header, header_types, True)
    result_column_types = make_excel_result_types(header, column_types, False)

    return has_header, result_merged_types, result_header_types, result_column_types

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/utils/converter_parsing_utils.py
datalens-tech/datalens-backend
import abc
import asyncio
import concurrent.futures
import csv
import logging
from typing import (
    Any,
    BinaryIO,
    Optional,
)

import attr

from dl_constants.enums import FileProcessingStatus
from dl_core.aio.web_app_services.s3 import S3Service
from dl_core.db import SchemaColumn
from dl_file_uploader_lib.enums import (
    CSVDelimiter,
    CSVEncoding,
)
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models import (
    CSVFileSettings,
    CSVFileSourceSettings,
    DataFile,
    DataSource,
    DataSourcePreview,
    FileSettings,
    FileSourceSettings,
    SpreadsheetFileSourceSettings,
)
from dl_file_uploader_worker_lib.utils.parsing_utils import (
    detect_dialect,
    detect_encoding,
    guess_header_and_schema,
    make_upcropped_text_sample,
    prepare_preview,
    prepare_preview_from_json_each_row,
    reguess_header_and_schema_spreadsheet,
)
from dl_utils.aio import ContextVarExecutor


LOGGER = logging.getLogger(__name__)


class FileParser(metaclass=abc.ABCMeta):
    def __init__(
        self,
        dfile: DataFile,
        s3: S3Service,
        tpe: ContextVarExecutor,
        file_settings: Optional[dict[str, Any]] = None,
        sample_size: int = 1 * 1024 * 1024,
        source_id: Optional[str] = None,
    ):
        self.dfile = dfile
        self.s3 = s3
        self.tpe = tpe
        self.file_settings = file_settings or {}
        self.sample_size = sample_size
        self.source_id = source_id

    def _get_sync_s3_data_stream(self, key: str) -> BinaryIO:
        s3_sync_cli = self.s3.get_sync_client()
        s3_sync_resp = s3_sync_cli.get_object(Bucket=self.s3.tmp_bucket_name, Key=key)
        return s3_sync_resp["Body"]

    def ensure_sources_and_return(self) -> list[DataSource]:
        raise NotImplementedError

    async def guess_header_and_schema(
        self,
        dsrc: DataSource,
    ) -> tuple[bool, list[SchemaColumn], Optional[FileSettings], FileSourceSettings]:
        raise NotImplementedError

    async def prepare_preview(self, dsrc: DataSource, rmm: RedisModelManager) -> DataSourcePreview:
        raise NotImplementedError


class CSVFileParser(FileParser):
    def ensure_sources_and_return(self) -> list[DataSource]:
        if self.source_id is not None:
            return [self.dfile.get_source_by_id(self.source_id)]

        self.dfile.sources = [
            DataSource(
                title=self.dfile.filename or self.dfile.id,
                raw_schema=[],
                status=FileProcessingStatus.in_progress,
            )
        ]

        return self.dfile.sources

    async def _make_sample_text(self) -> str:
        loop = asyncio.get_running_loop()

        s3_resp = await self.s3.client.get_object(
            Bucket=self.s3.tmp_bucket_name,
            Key=self.dfile.s3_key,
            Range=f"bytes=0-{self.sample_size}",
        )
        sample_bytes = await s3_resp["Body"].read()
        LOGGER.info(f"Downloaded sample from s3. Sample length: {len(sample_bytes)} bytes.")

        if self.file_settings.get("encoding") is not None:
            encoding = getattr(CSVEncoding, self.file_settings["encoding"])
            LOGGER.info(f"Overriding encoding with user defined: {encoding}")
        else:
            encoding = await loop.run_in_executor(self.tpe, detect_encoding, sample_bytes)
            LOGGER.info(f"Detected encoding: {encoding}")

        sample_text = await loop.run_in_executor(
            self.tpe,
            make_upcropped_text_sample,
            sample_bytes,
            self.sample_size,
            encoding,
        )

        self.encoding = encoding
        self.sample_text = sample_text
        return sample_text

    async def ensure_sample_text_loaded(self) -> None:
        if not hasattr(self, "sample_text") or not hasattr(self, "encoding"):
            await self._make_sample_text()

    async def guess_header_and_schema(
        self,
        dsrc: DataSource,
    ) -> tuple[bool, list[SchemaColumn], Optional[FileSettings], FileSourceSettings]:
        loop = asyncio.get_running_loop()
        await self.ensure_sample_text_loaded()

        dialect_detection_retries = 50
        dialect_detection_retry_timeout_sec = 0.2
        with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
            future = executor.submit(detect_dialect, self.sample_text)

            while dialect_detection_retries > 0:
                dialect_detection_retries -= 1
                await asyncio.sleep(dialect_detection_retry_timeout_sec)

                if not future.running():
                    dialect = future.result()
                    break
            else:
                dialect = csv.excel()
                executor.shutdown(wait=False, cancel_futures=True)

        LOGGER.info(f"Detected dialect: {dialect}")
        if self.file_settings.get("delimiter") is not None:
            delimiter = getattr(CSVDelimiter, self.file_settings["delimiter"])
            dialect.delimiter = delimiter.value
            LOGGER.info(f"Overriding dialect delimiter with user defined: {delimiter}")

        has_header = None
        if self.file_settings.get("first_line_is_header") is not None:
            has_header = self.file_settings["first_line_is_header"]
            LOGGER.info(f"Overriding `has_header` with user defined: has_header={has_header}")

        data_stream = await loop.run_in_executor(self.tpe, self._get_sync_s3_data_stream, self.dfile.s3_key)
        has_header, raw_schema = await loop.run_in_executor(
            self.tpe,
            guess_header_and_schema,
            data_stream,
            self.encoding,
            dialect,
            has_header,
        )
        LOGGER.info(f"has_header: {has_header}, raw_schema: {raw_schema}")

        file_settings = CSVFileSettings(
            encoding=self.encoding,
            dialect=dialect,
        )
        file_source_settings = CSVFileSourceSettings(
            first_line_is_header=has_header,
        )

        return has_header, raw_schema, file_settings, file_source_settings

    async def prepare_preview(self, dsrc: DataSource, rmm: RedisModelManager) -> DataSourcePreview:
        loop = asyncio.get_running_loop()

        await self.ensure_sample_text_loaded()

        file_settings = self.dfile.file_settings
        assert isinstance(file_settings, CSVFileSettings)

        file_source_settings = dsrc.file_source_settings
        assert isinstance(file_source_settings, CSVFileSourceSettings)

        preview_data = await loop.run_in_executor(
            self.tpe,
            prepare_preview,
            self.sample_text,
            file_settings.dialect,
            file_source_settings.first_line_is_header,
        )
        return DataSourcePreview(manager=rmm, preview_data=preview_data)


class SpreadsheetFileParser(FileParser):
    def ensure_sources_and_return(self) -> list[DataSource]:
        if self.source_id is not None:
            return [self.dfile.get_source_by_id(self.source_id)]

        return self.dfile.sources or []

    async def guess_header_and_schema(
        self,
        dsrc: DataSource,
    ) -> tuple[bool, list[SchemaColumn], Optional[FileSettings], FileSourceSettings]:
        loop = asyncio.get_running_loop()

        file_source_settings = dsrc.file_source_settings
        assert isinstance(file_source_settings, SpreadsheetFileSourceSettings)

        has_header = file_source_settings.first_line_is_header
        if self.file_settings.get("first_line_is_header") is not None:
            has_header = self.file_settings["first_line_is_header"]
            LOGGER.info(f"Overriding `has_header` with user defined: has_header={has_header}")

        raw_schema_header = file_source_settings.raw_schema_header
        raw_schema_body = file_source_settings.raw_schema_body
        file_type = file_source_settings.file_type
        has_header, raw_schema = await loop.run_in_executor(
            self.tpe, reguess_header_and_schema_spreadsheet, raw_schema_header, raw_schema_body, has_header, file_type
        )
        LOGGER.info(f"Re-guessed string columns. {has_header=}, {raw_schema=}")

        file_settings = None
        file_source_settings = attr.evolve(file_source_settings, first_line_is_header=has_header)

        return has_header, raw_schema, file_settings, file_source_settings

    async def prepare_preview(self, dsrc: DataSource, rmm: RedisModelManager) -> DataSourcePreview:
        loop = asyncio.get_running_loop()

        s3_resp = await self.s3.client.get_object(
            Bucket=self.s3.tmp_bucket_name,
            Key=dsrc.s3_key,
            Range=f"bytes=0-{self.sample_size}",
        )
        sample_bytes = await s3_resp["Body"].read()
        LOGGER.info(f"Downloaded sample from s3. Sample length: {len(sample_bytes)} bytes.")

        file_source_settings = dsrc.file_source_settings
        assert isinstance(file_source_settings, SpreadsheetFileSourceSettings)
        has_header = file_source_settings.first_line_is_header

        preview_data = await loop.run_in_executor(
            self.tpe, prepare_preview_from_json_each_row, sample_bytes, has_header
        )

        return DataSourcePreview(manager=rmm, preview_data=preview_data)

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/utils/file_parser.py
datalens-tech/datalens-backend
import abc
import csv
import io
import json
import logging
import sys
from typing import (
    BinaryIO,
    Iterator,
    Optional,
    Type,
    final,
)

import attr
import cchardet as chardet

from dl_app_tools.profiling_base import generic_profiler
from dl_constants.enums import (
    ConnectionType,
    UserDataType,
)
from dl_core.aio.web_app_services.gsheets import Sheet
from dl_core.components.ids import (
    ID_LENGTH,
    ID_VALID_SYMBOLS,
    make_readable_field_id,
)
from dl_core.db import get_type_transformer
from dl_core.db.elements import SchemaColumn
from dl_file_uploader_lib import exc
from dl_file_uploader_lib.enums import (
    CSVDelimiter,
    CSVEncoding,
    FileType,
)
from dl_file_uploader_worker_lib.utils import converter_parsing_utils
from dl_file_uploader_worker_lib.utils.converter_parsing_utils import (
    idx_to_alphabet_notation,
    make_result_types,
    merge_column_types,
    raw_schema_to_column_types,
)

from dl_connector_bundle_chs3.chs3_gsheets.core.constants import CONNECTION_TYPE_GSHEETS_V2
from dl_connector_bundle_chs3.file.core.constants import CONNECTION_TYPE_FILE


LOGGER = logging.getLogger(__name__)


MAX_PREVIEW_ROWS = 30
MAX_PREVIEW_SIZE = 100 * 1024  # 100 KByte
MAX_PREVIEW_CELL_LEN = 40
MAX_COLUMNS_COUNT = 300


@attr.s(auto_attribs=True, kw_only=True)
class CSVDialect:
    delimiter: str
    doublequote: bool
    skipinitialspace: bool
    escapechar: Optional[str] = None
    lineterminator: str = "\r\n"
    quotechar: Optional[str] = '"'
    quoting: int = csv.QUOTE_MINIMAL


@generic_profiler("detect-encoding")
def detect_encoding(sample: bytes) -> CSVEncoding:
    encoding_resp = chardet.detect(sample).get("encoding")

    LOGGER.info(f'Run chardet.detect. Detected encoding "{encoding_resp}".')

    if not encoding_resp:
        LOGGER.info('Can\'t detect file encoding. Falling back to "utf-8".')
        encoding = CSVEncoding.utf8
    else:
        encoding_resp = encoding_resp.lower()
        try:
            encoding = CSVEncoding(encoding_resp)
        except ValueError:
            if encoding_resp == "ascii":
                LOGGER.info('replaced encoding "ascii" to "utf-8"')
                encoding = CSVEncoding.utf8
            elif "windows" in encoding_resp:
                LOGGER.info(f'replaced encoding "{encoding_resp}" to "windows1251"')
                encoding = CSVEncoding.windows1251
            else:
                LOGGER.info(f'Unexpected encoding: {encoding_resp}. Falling back to "utf-8".')
                encoding = CSVEncoding.utf8

    return encoding


def make_upcropped_text_sample(sample: bytes, expected_sample_len: int, encoding: CSVEncoding) -> str:
    sample_text = sample.decode(encoding.value, errors="replace")
    if len(sample) == expected_sample_len:
        # That means, the last line may be cropped in the middle. It's better to drop it.
        sample_lines = sample_text.splitlines(keepends=True)[:-1]
        sample_text = "".join(sample_lines)
    return sample_text


@generic_profiler("detect-dialect")
def detect_dialect(sample: str) -> CSVDialect:
    allowed_delimiters = "".join((d.value for d in CSVDelimiter))

    try:
        dialect = csv.Sniffer().sniff(sample, delimiters=allowed_delimiters)
    except csv.Error:
        LOGGER.info("Could not determine csv delimiter. Falling back to excel csv dialect.", exc_info=True)
        dialect = csv.excel

    local_dialect = CSVDialect(
        doublequote=dialect.doublequote,
        delimiter=dialect.delimiter,
        quotechar=dialect.quotechar,
        skipinitialspace=dialect.skipinitialspace,
    )

    return local_dialect


@attr.s
class FileUploaderFieldIdGenerator(metaclass=abc.ABCMeta):
    _id_formatter: str = "{}_{}"

    _seen_ids: set[str] = attr.ib(factory=set)

    @abc.abstractmethod
    def _make_id(self, field: converter_parsing_utils.TResultColumn) -> str:
        raise NotImplementedError

    def _resolve_id_collisions(self, item: str) -> str:
        idx = 1
        orig_item = item
        while item in self._seen_ids:
            item = self._id_formatter.format(orig_item, idx)
            idx += 1
        return item

    @final
    def make_field_id(self, field: converter_parsing_utils.TResultColumn) -> str:
        field_id = self._make_id(field)
        field_id = self._resolve_id_collisions(field_id)
        self._seen_ids.add(field_id)
        return field_id


@attr.s
class FileFieldIdGenerator(FileUploaderFieldIdGenerator):
    def _make_id(self, field: converter_parsing_utils.TResultColumn) -> str:
        return make_readable_field_id(field["title"], ID_VALID_SYMBOLS, ID_LENGTH)  # type: ignore  # TODO: FIX


@attr.s
class GSheetsFieldIdGenerator(FileUploaderFieldIdGenerator):
    def _make_id(self, field: converter_parsing_utils.TResultColumn) -> str:
        return idx_to_alphabet_notation(field["index"])  # type: ignore  # TODO: FIX


def get_field_id_generator(conn_type: ConnectionType) -> FileUploaderFieldIdGenerator:
    field_id_gen_cls_map: dict[ConnectionType, Type[FileUploaderFieldIdGenerator]] = {
        CONNECTION_TYPE_FILE: FileFieldIdGenerator,
        CONNECTION_TYPE_GSHEETS_V2: GSheetsFieldIdGenerator,
    }

    if conn_type not in field_id_gen_cls_map:
        raise ValueError(f'No field id generator specified for connection type "{conn_type.name}"')

    return field_id_gen_cls_map[conn_type]()


def result_column_types_to_raw_schema(
    column_types: converter_parsing_utils.TResultTypes,
    conn_type: ConnectionType,
) -> list[SchemaColumn]:
    field_id_gen = get_field_id_generator(conn_type)
    type_transformer = get_type_transformer(conn_type)

    raw_schema: list[SchemaColumn] = []
    for col in column_types:
        user_type = getattr(UserDataType, col["cast"])  # type: ignore  # TODO: FIX
        title: str = col["title"]  # type: ignore  # TODO: FIX
        sch_col = SchemaColumn(
            name=field_id_gen.make_field_id(col),
            title=title,
            user_type=user_type,
            native_type=type_transformer.type_user_to_native(user_type),
        )
        raw_schema.append(sch_col)

    return raw_schema


@generic_profiler("guess-header-and-schema")
def guess_header_and_schema(
    binary_data_stream: BinaryIO,
    encoding: CSVEncoding,
    dialect: csv.Dialect,
    has_header: Optional[bool] = None,
) -> tuple[bool, list[SchemaColumn]]:
    text_io_wrapper = io.TextIOWrapper(binary_data_stream, encoding=encoding.value, newline="")

    csv.field_size_limit(sys.maxsize)  # TODO: wrap csv.Error exception if limit exceeded

    csv_reader = csv.reader(text_io_wrapper, dialect)
    new_has_header, legacy_column_types = converter_parsing_utils.guess_types_and_header(
        csv_reader, has_header=has_header
    )

    raw_schema = result_column_types_to_raw_schema(legacy_column_types, CONNECTION_TYPE_FILE)

    if len(raw_schema) > MAX_COLUMNS_COUNT:
        raise exc.TooManyColumnsError()

    return new_has_header, raw_schema


@generic_profiler("guess-header-and-schema-gsheet")
def guess_header_and_schema_gsheet(
    sheet: Sheet,
) -> tuple[bool, list[SchemaColumn], list[SchemaColumn], list[SchemaColumn]]:
    (
        has_header,
        legacy_column_types,
        col_types_header,
        col_types_body,
    ) = converter_parsing_utils.guess_types_and_header_gsheets(
        data_iter=iter(sheet.data or ()),
        sample_lines_count=None,
    )

    raw_schema = result_column_types_to_raw_schema(legacy_column_types, CONNECTION_TYPE_GSHEETS_V2)

    if len(raw_schema) > MAX_COLUMNS_COUNT:
        raise exc.TooManyColumnsError()

    raw_schema_header = result_column_types_to_raw_schema(col_types_header, CONNECTION_TYPE_GSHEETS_V2)
    raw_schema_body = result_column_types_to_raw_schema(col_types_body, CONNECTION_TYPE_GSHEETS_V2)

    return has_header, raw_schema, raw_schema_header, raw_schema_body


@generic_profiler("guess-schema-gsheet")
def guess_schema_gsheet(sheet: Sheet) -> list[SchemaColumn]:
    col_types = converter_parsing_utils.guess_types_gsheets(
        data_iter=iter(sheet.data or ()),
        sample_lines_count=None,
    )

    raw_schema = result_column_types_to_raw_schema(col_types, CONNECTION_TYPE_GSHEETS_V2)

    if len(raw_schema) > MAX_COLUMNS_COUNT:
        raise exc.TooManyColumnsError()

    return raw_schema


def merge_raw_schemas_spreadsheet(
    header_rs: list[SchemaColumn],
    body_rs: list[SchemaColumn],
    has_header: bool,
    file_type: FileType,
) -> list[SchemaColumn]:
    conn_type_map: dict[FileType, ConnectionType] = {
        FileType.xlsx: CONNECTION_TYPE_FILE,
        FileType.gsheets: CONNECTION_TYPE_GSHEETS_V2,
    }
    col_types_header = raw_schema_to_column_types(header_rs)
    col_types_body = raw_schema_to_column_types(body_rs)

    if has_header:
        header = [col.title for col in header_rs]
        result_types = make_result_types(header, col_types_body, has_header, idx_to_alphabet_notation)
    else:
        header = [col.title for col in body_rs]
        merged_col_types = merge_column_types(col_types_header, col_types_body, has_header)
        result_types = make_result_types(header, merged_col_types, has_header, idx_to_alphabet_notation)

    return result_column_types_to_raw_schema(result_types, conn_type_map[file_type])


@generic_profiler("reguess-header-and-schema-gsheet")
def reguess_header_and_schema_spreadsheet(
    raw_schema_header: list[SchemaColumn],
    raw_schema_body: list[SchemaColumn],
    has_header: bool,
    file_type: FileType,
) -> tuple[bool, list[SchemaColumn]]:
    raw_schema = merge_raw_schemas_spreadsheet(raw_schema_header, raw_schema_body, has_header, file_type)
    return has_header, raw_schema


def reduce_preview_size_inplace(preview: list[list[Optional[str]]]) -> None:
    LOGGER.info("Trying to shorten cells to decrease preview size.")
    new_size = 0
    for row in preview:
        for cell_index, cell in enumerate(row):
            if cell is None:
                continue
            if len(cell) > MAX_PREVIEW_CELL_LEN:
                cropped_cell = f"{cell[:10]}...{cell[-10:]}"
                row[cell_index] = cropped_cell
                new_size += len(cropped_cell.encode("utf-8"))
            else:
                new_size += len(cell.encode("utf-8"))
    LOGGER.info(f"Preview size after shorten attempt: {new_size} bytes.")


@generic_profiler("prepare-preview")
def prepare_preview(sample: str, dialect: csv.Dialect, has_header: bool) -> list[list[Optional[str]]]:
    csv.field_size_limit(sys.maxsize)  # TODO: wrap csv.Error exception if limit exceeded

    csv_reader = csv.reader(io.StringIO(sample, newline=""), dialect=dialect)

    if has_header:
        # Skip header from preview
        next(csv_reader)

    preview: list[list[Optional[str]]] = []
    preview_size = 0
    for row in csv_reader:
        preview.append(row)  # type: ignore
        if len(preview) >= MAX_PREVIEW_ROWS:
            break

    for preview_row in preview:
        for cell_index, cell in enumerate(preview_row):
            cell_str = cell if cell is not None else ""
            preview_size += len(cell_str.encode("utf-8"))

            if cell_str == "":
                preview_row[cell_index] = None

    LOGGER.info(f"Collected {len(preview)} preview rows. Preview size: {preview_size} bytes.")

    if preview_size > MAX_PREVIEW_SIZE:
        reduce_preview_size_inplace(preview)

    return preview


def prepare_preview_from_json_each_row(sample: bytes, has_header: bool) -> list[list[Optional[str]]]:
    sample_text = sample.decode("utf-8", errors="replace")
    sample_lines = iter(sample_text.splitlines())
    preview: list[list[Optional[str]]] = []
    preview_size = 0
    if has_header:
        next(sample_lines)
    for line in sample_lines:
        try:
            row_data = json.loads(line)
        except json.JSONDecodeError:
            continue
        raw_row_data = [str(value) if value is not None else "" for value in row_data]
        preview_size += sum(len(value.encode("utf-8")) for value in raw_row_data)
        preview.append([cell if cell != "" else None for cell in raw_row_data])
        if len(preview) >= MAX_PREVIEW_ROWS:
            break

    LOGGER.info(f"Collected {len(preview)} preview rows. Preview size: {preview_size} bytes.")

    if preview_size > MAX_PREVIEW_SIZE:
        reduce_preview_size_inplace(preview)

    return preview


def get_csv_raw_data_iterator(
    binary_data_stream: BinaryIO,
    encoding: CSVEncoding,
    dialect: csv.Dialect,
    first_line_is_header: bool,
    raw_schema: list[SchemaColumn],
) -> Iterator:
    text_io_wrapper = io.TextIOWrapper(binary_data_stream, encoding=encoding.value, newline="")

    csv.field_size_limit(sys.maxsize)  # TODO: wrap csv.Error exception if limit exceeded

    csv_reader = csv.DictReader(
        text_io_wrapper,
        fieldnames=tuple(sch.name for sch in raw_schema),
        dialect=dialect,
    )

    if first_line_is_header:
        next(csv_reader)

    return csv_reader


@generic_profiler("guess_header_and_schema_excel")
def guess_header_and_schema_excel(
    data: list,
) -> tuple[bool, list[SchemaColumn], list[SchemaColumn], list[SchemaColumn]]:
    (
        has_header,
        legacy_column_types,
        col_types_header,
        col_types_body,
    ) = converter_parsing_utils.guess_types_and_header_excel(
        data_iter=iter(data),
        sample_lines_count=None,
    )

    raw_schema = result_column_types_to_raw_schema(legacy_column_types, CONNECTION_TYPE_FILE)

    if len(raw_schema) > MAX_COLUMNS_COUNT:
        raise ValueError("Too many columns")  # TODO: proper exc type

    raw_schema_header = result_column_types_to_raw_schema(col_types_header, CONNECTION_TYPE_FILE)
    raw_schema_body = result_column_types_to_raw_schema(col_types_body, CONNECTION_TYPE_FILE)

    return has_header, raw_schema, raw_schema_header, raw_schema_body

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/utils/parsing_utils.py
datalens-tech/datalens-backend
from __future__ import annotations

import io
import json
from typing import (
    TYPE_CHECKING,
    BinaryIO,
    Iterator,
    NamedTuple,
    Optional,
)

from clickhouse_sqlalchemy.quoting import Quoter

from dl_core.db import (
    SchemaColumn,
    get_type_transformer,
)
from dl_core.raw_data_streaming.stream import SimpleDataStream
from dl_file_uploader_lib.data_sink.json_each_row import S3JsonEachRowFileDataSink
from dl_file_uploader_lib.enums import FileType
from dl_file_uploader_lib.redis_model.models import CSVFileSettings
from dl_file_uploader_lib.redis_model.models.models import (
    CSVFileSourceSettings,
    FileSettings,
    FileSourceSettings,
    SpreadsheetFileSourceSettings,
)
from dl_file_uploader_worker_lib.utils.parsing_utils import get_csv_raw_data_iterator

from dl_connector_bundle_chs3.chs3_base.core.us_connection import BaseFileS3Connection
from dl_connector_bundle_chs3.file.core.adapter import AsyncFileS3Adapter
from dl_connector_clickhouse.core.clickhouse_base.ch_commons import create_column_sql


if TYPE_CHECKING:
    from mypy_boto3_s3.client import S3Client as SyncS3Client


def make_s3_table_func_sql_source(
    conn: BaseFileS3Connection,
    source_id: str,
    bucket: str,
    filename: str,
    file_fmt: str,
    for_debug: bool = False,
    raw_schema_override: Optional[list[SchemaColumn]] = None,
) -> str:
    source_raw_schema: list[SchemaColumn] = conn.get_file_source_by_id(source_id).raw_schema or []
    q = Quoter().quote_str

    s3_path = q("{}/{}/{}".format(conn.s3_endpoint.strip("/"), bucket.strip("/"), filename))
    file_fmt = q(file_fmt)

    dialect = AsyncFileS3Adapter.get_dialect()
    raw_schema = raw_schema_override if raw_schema_override is not None else source_raw_schema
    type_transformer = get_type_transformer(conn.conn_type)
    schema_str = q(", ".join(create_column_sql(dialect, col, type_transformer) for col in raw_schema))

    if for_debug:
        access_key_id, secret_access_key = "<hidden>", "<hidden>"
    else:
        access_key_id = q(conn.s3_access_key_id)
        secret_access_key = q(conn.s3_secret_access_key)
    return f"s3({s3_path}, {access_key_id}, {secret_access_key}, {file_fmt}, {schema_str})"


class S3Object(NamedTuple):
    bucket: str
    key: str


def copy_from_s3_to_s3(
    s3_sync_cli: SyncS3Client,
    src_file: S3Object,
    dst_file: S3Object,
    file_type: FileType,
    file_settings: Optional[FileSettings],
    file_source_settings: FileSourceSettings,
    raw_schema: list[SchemaColumn],
) -> None:
    s3_sync_resp = s3_sync_cli.get_object(Bucket=src_file.bucket, Key=src_file.key)
    s3_data_stream: BinaryIO = s3_sync_resp["Body"]  # type: ignore  # TODO: fix

    def spreadsheet_data_iter() -> Iterator[dict]:
        fieldnames = tuple(sch.name for sch in raw_schema)
        text_io_wrapper = io.TextIOWrapper(s3_data_stream, encoding="utf-8", newline="")
        assert isinstance(file_source_settings, SpreadsheetFileSourceSettings)
        if file_source_settings.first_line_is_header:
            next(text_io_wrapper)
        for line in text_io_wrapper:
            line_data = json.loads(line)
            yield dict(zip(fieldnames, line_data))

    if file_type == FileType.csv:
        assert isinstance(file_settings, CSVFileSettings)
        assert isinstance(file_source_settings, CSVFileSourceSettings)
        data_iter = get_csv_raw_data_iterator(
            binary_data_stream=s3_data_stream,
            encoding=file_settings.encoding,
            dialect=file_settings.dialect,
            first_line_is_header=file_source_settings.first_line_is_header,
            raw_schema=raw_schema,
        )
    else:
        data_iter = spreadsheet_data_iter()

    data_stream = SimpleDataStream(
        data_iter=data_iter,
        rows_to_copy=None,  # TODO
    )
    with S3JsonEachRowFileDataSink(
        bi_schema=raw_schema,
        s3=s3_sync_cli,
        s3_key=dst_file.key,
        bucket_name=dst_file.bucket,
    ) as data_sink:
        data_sink.dump_data_stream(data_stream)

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib/utils/s3_utils.py
datalens-tech/datalens-backend
import attr

from dl_configs.crypto_keys import (
    CryptoKeysConfig,
    get_single_key_crypto_keys_config,
)


@attr.s(kw_only=True)
class TestingUSConfig:
    master_token: str = attr.ib(default="AC1ofiek8coB")
    base_url: str = attr.ib()
    crypto_keys_config: CryptoKeysConfig = attr.ib(
        default=get_single_key_crypto_keys_config(
            key_id="0",
            key_value="BRa9gGXGm9qyy1efiIH7OkRHALjAohYtxOz-wSvlARY=",
        )
    )
    psycopg2_pg_dns: str = attr.ib()
    force_clear_db_on_launch: bool = attr.ib(default=True)

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/config.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import logging
import os
import sys
from typing import TYPE_CHECKING

import attr
import pytest
import redis.asyncio

from dl_api_commons.base_models import (
    RequestContextInfo,
    TenantCommon,
)
from dl_configs.crypto_keys import get_dummy_crypto_keys_config
from dl_configs.enums import RedisMode
from dl_configs.settings_submodels import (
    GoogleAppSettings,
    RedisSettings,
    S3Settings,
)
from dl_core.loader import load_core_lib
from dl_core.services_registry.top_level import DummyServiceRegistry
from dl_core.united_storage_client import USAuthContextMaster
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_core.us_manager.us_manager_sync import SyncUSManager
from dl_core_testing.environment import (
    common_pytest_configure,
    prepare_united_storage,
)
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_worker_lib.app import FileUploaderContextFab
from dl_file_uploader_worker_lib.settings import (
    FileUploaderConnectorsSettings,
    FileUploaderWorkerSettings,
    SecureReader,
)
from dl_file_uploader_worker_lib.tasks import REGISTRY
from dl_file_uploader_worker_lib.testing.app_factory import TestingFileUploaderWorkerFactory
from dl_task_processor.arq_wrapper import (
    create_arq_redis_settings,
    create_redis_pool,
)
from dl_task_processor.executor import Executor
from dl_task_processor.processor import (
    ARQProcessorImpl,
    LocalProcessorImpl,
    TaskProcessor,
)
from dl_task_processor.state import (
    BITaskStateImpl,
    TaskState,
)
from dl_task_processor.worker import ArqWorkerTestWrapper
from dl_testing.containers import get_test_container_hostport
from dl_testing.s3_utils import (
    create_s3_bucket,
    create_s3_client,
    create_sync_s3_client,
)
from dl_testing.utils import wait_for_initdb

from dl_connector_bundle_chs3.chs3_base.core.settings import FileS3ConnectorSettings

from .config import TestingUSConfig


if TYPE_CHECKING:
    from mypy_boto3_s3.client import S3Client as SyncS3Client
    from types_aiobotocore_s3 import S3Client as AsyncS3Client


LOGGER = logging.getLogger(__name__)


pytest_plugins = ("aiohttp.pytest_plugin",)


def pytest_configure(config: Any) -> None:  # noqa
    common_pytest_configure(tracing_service_name="tests_bi_file_uploader_worker")

    os.environ["EXT_QUERY_EXECUTER_SECRET_KEY"] = "qwerty"


@pytest.fixture(autouse=True)
def loop(event_loop):
    """
    Preventing creation of new loop by `aiohttp.pytest_plugin` loop fixture in favor of pytest-asyncio one
    And set loop pytest-asyncio created loop as default for thread
    """
    asyncio.set_event_loop(event_loop)
    return event_loop


@pytest.fixture(scope="session")
def initdb_ready():
    return wait_for_initdb(initdb_port=get_test_container_hostport("init-db", fallback_port=51508).port)


@pytest.fixture(scope="function")
def rci() -> RequestContextInfo:
    return RequestContextInfo(user_id="_the_tests_asyncapp_user_id_")


@pytest.fixture(scope="session", autouse=True)
def loaded_libraries():
    load_core_lib()


@pytest.fixture(scope="session")
def us_config():
    return TestingUSConfig(
        base_url=f"http://{get_test_container_hostport('us').as_pair()}",
        psycopg2_pg_dns=(
            f'host={get_test_container_hostport("pg-us").host}'
            f' port={get_test_container_hostport("pg-us").port}'
            " user=us"
            " password=us"
            " dbname=us-db-ci_purgeable"
        ),
    )


@pytest.fixture(scope="session")
def redis_app_settings():
    return RedisSettings(
        MODE=RedisMode.single_host,
        CLUSTER_NAME="",
        HOSTS=(get_test_container_hostport("redis").host,),
        PORT=get_test_container_hostport("redis").port,
        DB=9,
    )


@pytest.fixture(scope="session")
def redis_arq_settings(redis_app_settings):
    return attr.evolve(redis_app_settings, DB=11)


@pytest.fixture(scope="session")
def s3_settings() -> S3Settings:
    return S3Settings(
        ENDPOINT_URL=f'http://{get_test_container_hostport("s3-storage").as_pair()}',
        # default access and secrets for zenko/cloudserver
        ACCESS_KEY_ID="accessKey1",
        SECRET_ACCESS_KEY="verySecretKey1",
    )


@pytest.fixture(scope="session")
def secure_reader():
    socket_name = "reader.sock"
    if sys.platform == "darwin":
        path = "/var/tmp"
    else:
        path = "/var"

    return SecureReader(
        SOCKET=os.path.join(path, socket_name),
    )


@pytest.fixture(scope="session")
def connectors_settings(s3_settings):
    return FileUploaderConnectorsSettings(
        FILE=FileS3ConnectorSettings(
            SECURE=False,
            HOST=get_test_container_hostport("db-clickhouse", original_port=8123).host,
            PORT=get_test_container_hostport("db-clickhouse", original_port=8123).port,
            USERNAME="datalens",
            PASSWORD="qwerty",
            ACCESS_KEY_ID=s3_settings.ACCESS_KEY_ID,
            SECRET_ACCESS_KEY=s3_settings.SECRET_ACCESS_KEY,
            BUCKET="bi-file-uploader",
            S3_ENDPOINT="http://s3-storage:8000",
        ),
    )


@pytest.fixture(scope="session")
def file_uploader_worker_settings(
    redis_app_settings,
    redis_arq_settings,
    s3_settings,
    connectors_settings,
    us_config,
    secure_reader,
):
    settings = FileUploaderWorkerSettings(
        REDIS_APP=redis_app_settings,
        REDIS_ARQ=redis_arq_settings,
        S3=S3Settings(
            ENDPOINT_URL=s3_settings.ENDPOINT_URL,
            ACCESS_KEY_ID=s3_settings.ACCESS_KEY_ID,
            SECRET_ACCESS_KEY=s3_settings.SECRET_ACCESS_KEY,
        ),
        S3_TMP_BUCKET_NAME="bi-file-uploader-tmp",
        S3_PERSISTENT_BUCKET_NAME="bi-file-uploader",
        SENTRY_DSN=None,
        US_BASE_URL=us_config.base_url,
        US_MASTER_TOKEN=us_config.master_token,
        CONNECTORS=connectors_settings,
        GSHEETS_APP=GoogleAppSettings(
            API_KEY="dummy",
            CLIENT_ID="dummy",
            CLIENT_SECRET="dummy",
        ),
        CRYPTO_KEYS_CONFIG=get_dummy_crypto_keys_config(),
        SECURE_READER=secure_reader,
    )
    yield settings


@pytest.fixture(scope="function")
async def redis_pool(file_uploader_worker_settings):
    pool = await create_redis_pool(
        create_arq_redis_settings(file_uploader_worker_settings.REDIS_ARQ),
    )
    yield pool
    await pool.close()


@pytest.fixture(scope="function")
def task_state():
    return TaskState(BITaskStateImpl())


@pytest.fixture(scope="function")
async def task_processor_arq_worker(loop, task_state, file_uploader_worker_settings):
    LOGGER.info("Set up worker")
    worker = TestingFileUploaderWorkerFactory(settings=file_uploader_worker_settings).create_worker(state=task_state)
    wrapper = ArqWorkerTestWrapper(loop=loop, worker=worker)
    yield await wrapper.start()
    await wrapper.stop()


@pytest.fixture(scope="function")
def task_processor_arq_client(loop, task_processor_arq_worker, redis_pool, task_state):
    impl = ARQProcessorImpl(redis_pool)
    p = TaskProcessor(impl=impl, state=task_state)
    LOGGER.info("Arq TP is ready")
    return p


@pytest.fixture(scope="function")
async def task_processor_local_client(loop, task_state, file_uploader_worker_settings):
    context_fab = FileUploaderContextFab(file_uploader_worker_settings)
    context = await context_fab.make()
    executor = Executor(context=context, state=task_state, registry=REGISTRY)
    impl = LocalProcessorImpl(executor)
    processor = TaskProcessor(impl=impl, state=task_state)
    LOGGER.info("Local TP is ready")
    yield processor
    await context_fab.tear_down(context)


@pytest.fixture(scope="function", params=["local"])
def task_processor_client(request, task_processor_arq_client, task_processor_local_client):
    if request.param == "arq":
        return task_processor_arq_client
    elif request.param == "local":
        return task_processor_local_client


@pytest.fixture(scope="function")
async def s3_client(s3_settings) -> AsyncS3Client:
    async with create_s3_client(s3_settings) as client:
        yield client


@pytest.fixture(scope="function")
def s3_client_sync(s3_settings) -> SyncS3Client:
    return create_sync_s3_client(s3_settings)


@pytest.fixture(scope="function")
def redis_cli(redis_app_settings) -> redis.asyncio.Redis:
    return redis.asyncio.Redis(
        host=redis_app_settings.HOSTS[0],
        port=redis_app_settings.PORT,
        db=redis_app_settings.DB,
        password=redis_app_settings.PASSWORD,
    )


@pytest.fixture(scope="function")
def redis_model_manager(redis_cli, rci) -> RedisModelManager:
    return RedisModelManager(redis=redis_cli, rci=rci, crypto_keys_config=get_dummy_crypto_keys_config())


@pytest.fixture(scope="function")
async def s3_tmp_bucket(initdb_ready, s3_client, file_uploader_worker_settings) -> str:
    bucket_name = file_uploader_worker_settings.S3_TMP_BUCKET_NAME
    await create_s3_bucket(s3_client, bucket_name, max_attempts=1)
    return bucket_name


@pytest.fixture(scope="function")
async def s3_persistent_bucket(initdb_ready, s3_client, file_uploader_worker_settings) -> str:
    bucket_name = file_uploader_worker_settings.S3_PERSISTENT_BUCKET_NAME
    await create_s3_bucket(s3_client, bucket_name, max_attempts=1)
    return bucket_name


@pytest.fixture(scope="session")
def bi_context() -> RequestContextInfo:
    return RequestContextInfo.create(
        request_id=None,
        tenant=TenantCommon(),
        user_name=None,
        user_id="_the_tests_fixture_user_id_",
        x_dl_debug_mode=None,
        secret_headers=None,
        auth_data=None,
        plain_headers=None,
        endpoint_code=None,
        x_dl_context=None,
    )


@pytest.fixture(scope="session")
def prepare_us(us_config):
    prepare_united_storage(
        us_host=us_config.base_url,
        us_master_token=us_config.master_token,
        us_pg_dsn=us_config.psycopg2_pg_dns,
        force=us_config.force_clear_db_on_launch,
    )

    return us_config


@pytest.fixture(scope="session")
def default_sync_usm(bi_context, prepare_us, us_config):
    rci = RequestContextInfo.create_empty()
    return SyncUSManager(
        us_base_url=us_config.base_url,
        us_auth_context=USAuthContextMaster(us_config.master_token),
        bi_context=bi_context,
        crypto_keys_config=get_dummy_crypto_keys_config(),
        services_registry=DummyServiceRegistry(rci=rci),
    )


@pytest.fixture(scope="function")
@pytest.mark.usefixtures("loop")
def default_async_usm_per_test(bi_context, prepare_us, us_config):
    rci = RequestContextInfo.create_empty()
    return AsyncUSManager(
        us_base_url=us_config.base_url,
        us_auth_context=USAuthContextMaster(us_config.master_token),
        crypto_keys_config=us_config.crypto_keys_config,
        bi_context=bi_context,
        services_registry=DummyServiceRegistry(rci=rci),
    )

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
import os

import aiohttp.web
import pytest

from dl_constants.enums import FileProcessingStatus
from dl_file_secure_reader_lib.app import create_app as create_reader_app
from dl_file_uploader_lib.enums import FileType
from dl_file_uploader_lib.redis_model.models import DataFile
from dl_file_uploader_lib.testing.data_gen import generate_sample_csv_data_str
from dl_file_uploader_task_interface.tasks import (
    ParseFileTask,
    SaveSourceTask,
)
from dl_task_processor.state import wait_task

from .utils import create_file_connection


LOGGER = logging.getLogger(__name__)


@pytest.fixture(scope="function")
async def uploaded_file(s3_tmp_bucket, s3_persistent_bucket, s3_client, redis_model_manager) -> DataFile:
    csv_data = f"""f1,f2,f3,Дата,Дата и время
qwe,123,45.9,2021-02-04,2021-02-04 12:00:00
asd,345,47.9,2021-02-05,2021-02-05 14:01:00
zxc,456,"49,9",2021-02-06,2021-02-06 11:59:00
zxc,456,,,2021-02-06 11:59:00
{'zxc'*35000},456,49.9,2021-02-06,""".encode(
        "utf-8"
    )

    data_file_desc = DataFile(
        manager=redis_model_manager,
        filename="test_file.csv",
        file_type=FileType.csv,
        size=len(csv_data),
        status=FileProcessingStatus.in_progress,
    )
    await s3_client.put_object(
        ACL="private",
        Bucket=s3_tmp_bucket,
        Key=data_file_desc.s3_key,
        Body=csv_data,
    )

    await data_file_desc.save()
    yield data_file_desc


another_uploaded_file = uploaded_file


@pytest.fixture(scope="function")
async def uploaded_file_id(uploaded_file) -> str:
    yield uploaded_file.id


@pytest.fixture(scope="function")
async def another_uploaded_file_id(another_uploaded_file) -> str:
    yield another_uploaded_file.id


@pytest.fixture(scope="function")
async def uploaded_10mb_file_id(s3_tmp_bucket, s3_persistent_bucket, s3_client, redis_model_manager) -> str:
    csv_data = generate_sample_csv_data_str(row_count=10000, str_cols_count=30).encode("utf-8")

    data_file_desc = DataFile(
        manager=redis_model_manager,
        filename="test_file_10mb.csv",
        file_type=FileType.csv,
        size=len(csv_data),
        status=FileProcessingStatus.in_progress,
    )
    await s3_client.put_object(
        ACL="private",
        Bucket=s3_tmp_bucket,
        Key=data_file_desc.s3_key,
        Body=csv_data,
    )

    await data_file_desc.save()
    yield data_file_desc.id


@pytest.fixture(scope="function")
async def uploaded_excel(s3_tmp_bucket, s3_persistent_bucket, s3_client, redis_model_manager) -> DataFile:
    filename = "data.xlsx"
    data_file_desc = DataFile(
        manager=redis_model_manager,
        filename=filename,
        file_type=FileType.xlsx,
        status=FileProcessingStatus.in_progress,
    )

    dirname = os.path.dirname(os.path.abspath(__file__))
    path = os.path.join(dirname, filename)

    with open(path, "rb") as fd:
        await s3_client.put_object(
            ACL="private",
            Bucket=s3_tmp_bucket,
            Key=data_file_desc.s3_key,
            Body=fd.read(),
        )

    await data_file_desc.save()
    yield data_file_desc


@pytest.fixture(scope="function")
async def uploaded_excel_id(uploaded_excel) -> str:
    yield uploaded_excel.id


@pytest.fixture(scope="function")
async def uploaded_excel_with_one_row(s3_tmp_bucket, s3_persistent_bucket, s3_client, redis_model_manager) -> DataFile:
    filename = "one_row_table.xlsx"
    data_file_desc = DataFile(
        manager=redis_model_manager,
        filename=filename,
        file_type=FileType.xlsx,
        status=FileProcessingStatus.in_progress,
    )

    dirname = os.path.dirname(os.path.abspath(__file__))
    path = os.path.join(dirname, filename)

    with open(path, "rb") as fd:
        await s3_client.put_object(
            ACL="private",
            Bucket=s3_tmp_bucket,
            Key=data_file_desc.s3_key,
            Body=fd.read(),
        )

    await data_file_desc.save()
    yield data_file_desc


@pytest.fixture(scope="function")
async def uploaded_excel_with_one_row_id(uploaded_excel_with_one_row) -> str:
    yield uploaded_excel_with_one_row.id


@pytest.fixture(scope="function")
def reader_app(loop, secure_reader):
    current_app = create_reader_app()
    runner = aiohttp.web.AppRunner(current_app)
    loop.run_until_complete(runner.setup())
    site = aiohttp.web.UnixSite(runner, path=secure_reader.SOCKET)
    return loop.run_until_complete(site.start())


@pytest.fixture(scope="function")
async def saved_file_connection_id(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    uploaded_file_id,
    default_async_usm_per_test,
) -> str:
    task_parse = await task_processor_client.schedule(ParseFileTask(file_id=uploaded_file_id))
    result = await wait_task(task_parse, task_state)
    assert result[-1] == "success"

    df = await DataFile.get(manager=redis_model_manager, obj_id=uploaded_file_id)
    source = df.sources[0]
    assert source.status == FileProcessingStatus.ready

    conn = await create_file_connection(default_async_usm_per_test, df.id, source.id, source.raw_schema)

    task_save = await task_processor_client.schedule(
        SaveSourceTask(
            tenant_id="common",
            file_id=uploaded_file_id,
            src_source_id=source.id,
            dst_source_id=source.id,
            connection_id=conn.uuid,
        )
    )
    await wait_task(task_save, task_state)

    return conn.uuid

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/db/conftest.py
datalens-tech/datalens-backend
from asyncio import sleep
import logging

import pytest

from dl_constants.enums import (
    FileProcessingStatus,
    UserDataType,
)
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    DataSourcePreview,
)
from dl_file_uploader_task_interface.tasks import ProcessExcelTask
from dl_task_processor.state import wait_task


LOGGER = logging.getLogger(__name__)


@pytest.mark.asyncio
async def test_parse_excel_task(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    uploaded_excel_id,
    reader_app,
):
    rmm = redis_model_manager
    df = await DataFile.get(manager=rmm, obj_id=uploaded_excel_id)
    assert df.status == FileProcessingStatus.in_progress

    task = await task_processor_client.schedule(ProcessExcelTask(file_id=uploaded_excel_id))
    result = await wait_task(task, task_state)
    await sleep(60)

    assert result[-1] == "success"

    df = await DataFile.get(manager=rmm, obj_id=uploaded_excel_id)
    assert df.id == uploaded_excel_id
    assert df.status == FileProcessingStatus.ready

    assert len(df.sources) == 3
    dsrc = df.sources[0]
    assert dsrc.status == FileProcessingStatus.ready
    assert dsrc.title == "data.xlsx – Orders"
    assert [sch.user_type for sch in dsrc.raw_schema] == [
        UserDataType.integer,
        UserDataType.string,
        UserDataType.genericdatetime,
        UserDataType.genericdatetime,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.string,
        UserDataType.float,
        UserDataType.integer,
        UserDataType.float,
        UserDataType.float,
        UserDataType.float,
        UserDataType.string,
        UserDataType.float,
    ]
    assert [sch.name for sch in dsrc.raw_schema] == [
        "row",
        "order_id",
        "order_date",
        "ship_date",
        "ship_mode",
        "customer_id",
        "customer_name",
        "segment",
        "city",
        "state",
        "global_area",
        "postal_code",
        "market",
        "region",
        "product_id",
        "category",
        "sub-category",
        "product_name",
        "sales",
        "quantity",
        "discount",
        "profit",
        "shipping_cost",
        "order_priority",
        "ammount",
    ]
    assert [sch.title for sch in dsrc.raw_schema] == [
        "Row",
        "Order ID",
        "Order Date",
        "Ship Date",
        "Ship Mode",
        "Customer ID",
        "Customer Name",
        "Segment",
        "City",
        "State",
        "Global Area",
        "Postal Code",
        "Market",
        "Region",
        "Product ID",
        "Category",
        "Sub-Category",
        "Product Name",
        "Sales",
        "Quantity",
        "Discount",
        "Profit",
        "Shipping Cost",
        "Order Priority",
        "Ammount",
    ]

    preview = await DataSourcePreview.get(manager=rmm, obj_id=dsrc.preview_id)
    assert preview.id == dsrc.preview_id
    assert preview.preview_data[0] == [
        "1",
        "MX-2014-143658",
        "2014-10-02 00:00:00",
        "2014-10-06 00:00:00",
        "Standard Class",
        "SC-20575",
        "Sonia Cooley",
        "Consumer",
        "Mexico City",
        "Distrito Federal",
        "Mexico",
        None,
        "LATAM",
        "North",
        "OFF-LA-10002782",
        "Office Supplies",
        "Labels",
        "Hon File Folder Labels, Adjustable",
        "13.08",
        "3",
        "0",
        "4.56",
        "1.033",
        "Medium",
        "39.24",
    ]


@pytest.mark.asyncio
async def test_parse_excel_with_one_row_task(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    uploaded_excel_with_one_row_id,
    reader_app,
):
    uploaded_excel_id = uploaded_excel_with_one_row_id
    rmm = redis_model_manager
    df = await DataFile.get(manager=rmm, obj_id=uploaded_excel_id)
    assert df.status == FileProcessingStatus.in_progress

    task = await task_processor_client.schedule(ProcessExcelTask(file_id=uploaded_excel_id))
    result = await wait_task(task, task_state)
    await sleep(60)

    assert result[-1] == "failed"

    df = await DataFile.get(manager=rmm, obj_id=uploaded_excel_id)
    assert df.status == FileProcessingStatus.failed

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/db/test_excel.py
datalens-tech/datalens-backend
from datetime import (
    datetime,
    timedelta,
)
import logging
from typing import Any
import uuid

from botocore.exceptions import ClientError
import pytest

from dl_constants.enums import (
    FileProcessingStatus,
    UserDataType,
)
from dl_file_uploader_lib import exc
from dl_file_uploader_lib.enums import (
    CSVEncoding,
    FileType,
    RenameTenantStatus,
)
from dl_file_uploader_lib.redis_model.base import RedisModelNotFound
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    DataSourcePreview,
    PreviewSet,
    RenameTenantStatusModel,
)
from dl_file_uploader_lib.testing.data_gen import generate_sample_csv_data_str
from dl_file_uploader_task_interface.tasks import (
    CleanS3LifecycleRulesTask,
    CleanupTenantFilePreviewsTask,
    CleanupTenantTask,
    DeleteFileTask,
    ParseFileTask,
    RenameTenantFilesTask,
    SaveSourceTask,
)
from dl_file_uploader_worker_lib.utils import parsing_utils
from dl_task_processor.state import wait_task

from .utils import create_file_connection


LOGGER = logging.getLogger(__name__)


@pytest.mark.asyncio
async def test_parse_file_task(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    uploaded_file_id,
):
    rmm = redis_model_manager
    df = await DataFile.get(manager=rmm, obj_id=uploaded_file_id)
    assert df.status == FileProcessingStatus.in_progress

    task = await task_processor_client.schedule(ParseFileTask(file_id=uploaded_file_id))
    await wait_task(task, task_state)

    df = await DataFile.get(manager=rmm, obj_id=uploaded_file_id)
    assert df.id == uploaded_file_id
    assert df.status == FileProcessingStatus.ready

    assert len(df.sources) == 1
    dsrc = df.sources[0]
    assert dsrc.status == FileProcessingStatus.ready
    assert dsrc.title == "test_file.csv"
    assert [sch.user_type for sch in dsrc.raw_schema] == [
        UserDataType.string,
        UserDataType.integer,
        UserDataType.float,
        UserDataType.date,
        UserDataType.genericdatetime,
    ]
    assert [sch.name for sch in dsrc.raw_schema] == ["f1", "f2", "f3", "data", "data_i_vremya"]
    assert [sch.title for sch in dsrc.raw_schema] == ["f1", "f2", "f3", "Дата", "Дата и время"]

    preview = await DataSourcePreview.get(manager=rmm, obj_id=dsrc.preview_id)
    assert preview.id == dsrc.preview_id
    assert preview.preview_data == [
        ["qwe", "123", "45.9", "2021-02-04", "2021-02-04 12:00:00"],
        ["asd", "345", "47.9", "2021-02-05", "2021-02-05 14:01:00"],
        ["zxc", "456", "49,9", "2021-02-06", "2021-02-06 11:59:00"],
        ["zxc", "456", None, None, "2021-02-06 11:59:00"],
        ["zxczxczxcz...czxczxczxc", "456", "49.9", "2021-02-06", None],
    ]


@pytest.mark.asyncio
async def test_parse_file_task_with_file_settings(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    uploaded_file_id,
):
    rmm = redis_model_manager
    df = await DataFile.get(manager=rmm, obj_id=uploaded_file_id)
    assert df.status == FileProcessingStatus.in_progress
    task = await task_processor_client.schedule(ParseFileTask(file_id=uploaded_file_id))
    await wait_task(task, task_state)

    df = await DataFile.get(manager=rmm, obj_id=uploaded_file_id)
    assert df.status == FileProcessingStatus.ready
    dsrc = df.sources[0]
    assert dsrc.status == FileProcessingStatus.ready

    source_id = dsrc.id
    task = await task_processor_client.schedule(
        ParseFileTask(
            file_id=uploaded_file_id,
            source_id=source_id,
            file_settings=dict(encoding="utf8", delimiter="tab", first_line_is_header=False),
            source_settings={},
        )
    )
    await wait_task(task, task_state)

    df = await DataFile.get(manager=rmm, obj_id=uploaded_file_id)
    assert df.status == FileProcessingStatus.ready
    dsrc = df.sources[0]
    assert dsrc.id == source_id
    assert dsrc.status == FileProcessingStatus.ready
    assert len(dsrc.raw_schema) == 1
    assert dsrc.file_source_settings.first_line_is_header is False
    assert df.file_settings.dialect.delimiter == "\t"
    assert df.file_settings.encoding == CSVEncoding.utf8

    preview = await DataSourcePreview.get(manager=rmm, obj_id=dsrc.preview_id)
    assert preview.id == dsrc.preview_id
    assert len(preview.preview_data) == 6


@pytest.mark.asyncio
async def test_parse_10mb_file_task(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    uploaded_10mb_file_id,
):
    rmm = redis_model_manager
    df = await DataFile.get(manager=rmm, obj_id=uploaded_10mb_file_id)
    assert df.status == FileProcessingStatus.in_progress

    task = await task_processor_client.schedule(ParseFileTask(file_id=uploaded_10mb_file_id))
    await wait_task(task, task_state)

    df = await DataFile.get(manager=rmm, obj_id=uploaded_10mb_file_id)
    assert df.status == FileProcessingStatus.ready


@pytest.mark.asyncio
async def test_save_source_task(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    uploaded_file_id,
    default_async_usm_per_test,
):
    usm = default_async_usm_per_test
    task = await task_processor_client.schedule(ParseFileTask(file_id=uploaded_file_id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    df = await DataFile.get(manager=redis_model_manager, obj_id=uploaded_file_id)
    source = df.sources[0]
    assert source.status == FileProcessingStatus.ready

    conn = await create_file_connection(usm, df.id, source.id, source.raw_schema)
    assert conn.get_file_source_by_id(source.id).status == FileProcessingStatus.in_progress

    task_save = await task_processor_client.schedule(
        SaveSourceTask(
            tenant_id="common",
            file_id=uploaded_file_id,
            src_source_id=source.id,
            dst_source_id=source.id,
            connection_id=conn.uuid,
        )
    )
    await wait_task(task_save, task_state)

    conn1 = await usm.get_by_id(conn.uuid)
    assert conn1.get_file_source_by_id(source.id).status == FileProcessingStatus.ready


@pytest.mark.asyncio
async def test_save_source_task_on_replace(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    uploaded_file_id,
    another_uploaded_file_id,
    default_async_usm_per_test,
):
    usm = default_async_usm_per_test
    df = []
    source = []

    for file_id in (uploaded_file_id, another_uploaded_file_id):
        task = await task_processor_client.schedule(ParseFileTask(file_id=file_id))
        await wait_task(task, task_state)
        new_df = await DataFile.get(manager=redis_model_manager, obj_id=file_id)
        df.append(new_df)
        new_source = new_df.sources[0]
        source.append(new_source)
        assert new_source.status == FileProcessingStatus.ready

    conn = await create_file_connection(usm, df[0].id, source[1].id, source[0].raw_schema)
    assert conn.get_file_source_by_id(source[1].id).status == FileProcessingStatus.in_progress

    task_save = await task_processor_client.schedule(
        SaveSourceTask(
            tenant_id="common",
            file_id=uploaded_file_id,
            src_source_id=source[0].id,
            dst_source_id=source[1].id,
            connection_id=conn.uuid,
        )
    )
    await wait_task(task_save, task_state)

    conn = await usm.get_by_id(conn.uuid)
    assert conn.get_file_source_by_id(source[1].id).status == FileProcessingStatus.ready


@pytest.mark.asyncio
async def test_delete_file_task(
    task_processor_client,
    task_state,
    s3_client,
    s3_persistent_bucket,
    redis_model_manager,
    uploaded_file,
):
    filename = uploaded_file.filename

    task = await task_processor_client.schedule(ParseFileTask(file_id=uploaded_file.id))
    await wait_task(task, task_state)

    df = await DataFile.get(manager=redis_model_manager, obj_id=uploaded_file.id)
    source = df.sources[0]

    task = await task_processor_client.schedule(DeleteFileTask(filename, source.preview_id))
    await wait_task(task, task_state)

    with pytest.raises(ClientError) as ex:
        await s3_client.head_object(
            Bucket=s3_persistent_bucket,
            Key=filename,
        )
        assert ex["ResponseMetadata"]["HTTPStatusCode"] == 404

    # and now try with not existing file
    task = await task_processor_client.schedule(DeleteFileTask(filename, source.preview_id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"


@pytest.fixture(scope="function")
async def files_with_tenant_prefixes(s3_persistent_bucket, s3_client) -> list[str]:
    csv_data = """f1,f2,f3,Дата,Дата и время
qwe,123,45.9,2021-02-04,2021-02-04 12:00:00
asd,345,47.9,2021-02-05,2021-02-05 14:01:00""".encode(
        "utf-8"
    )

    tenant_ids = [f"tenant_{uuid.uuid4()}" for _ in range(3)]
    for tenant_id in tenant_ids:
        for idx in range(42):
            filename = f"file_{idx}"
            await s3_client.put_object(
                ACL="private",
                Bucket=s3_persistent_bucket,
                Key=f"{tenant_id}_{filename}",
                Body=csv_data,
            )

    return tenant_ids


async def get_lc_rules_number(s3_client, bucket) -> int:
    try:
        lc_config = await s3_client.get_bucket_lifecycle_configuration(Bucket=bucket)
    except ClientError as ex:
        if ex.response["Error"]["Code"] == "NoSuchLifecycleConfiguration":
            lc_config = {"Rules": []}
        else:
            raise
    return len(lc_config["Rules"])


@pytest.mark.asyncio
async def test_cleanup_tenant_task(
    task_processor_client,
    task_state,
    s3_client,
    s3_persistent_bucket,
    files_with_tenant_prefixes,
):
    await s3_client.delete_bucket_lifecycle(Bucket=s3_persistent_bucket)
    tenant_ids = files_with_tenant_prefixes

    n_lc_rules = await get_lc_rules_number(s3_client, s3_persistent_bucket)
    assert n_lc_rules == 0

    for tenant_id in tenant_ids:
        task = await task_processor_client.schedule(CleanupTenantTask(tenant_id=tenant_id))
        result = await wait_task(task, task_state)
        assert result[-1] == "success"

        new_n_lc_rules = await get_lc_rules_number(s3_client, s3_persistent_bucket)
        assert new_n_lc_rules == n_lc_rules + 1
        n_lc_rules = new_n_lc_rules

    assert n_lc_rules == len(tenant_ids)


def make_cleanup_rule(prefix: str) -> dict[str, Any]:
    return {
        "ID": prefix,
        "Expiration": {
            "Date": datetime.today().replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=2),
        },
        "Filter": {
            "Prefix": prefix,
        },
        "Status": "Enabled",
    }


@pytest.mark.asyncio
async def test_cleanup_tenant_task_no_files(
    task_processor_client,
    task_state,
    s3_client,
    s3_persistent_bucket,
):
    await s3_client.delete_bucket_lifecycle(Bucket=s3_persistent_bucket)

    task = await task_processor_client.schedule(CleanupTenantTask(tenant_id="there are no files in this tenant"))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    assert await get_lc_rules_number(s3_client, s3_persistent_bucket) == 1  # note: the LC rule is still created


@pytest.mark.asyncio
async def test_regular_bucket_lifecycle_cleanup_task(
    task_processor_client,
    task_state,
    s3_client,
    s3_persistent_bucket,
    files_with_tenant_prefixes,
):
    tenant_ids = files_with_tenant_prefixes
    await s3_client.delete_bucket_lifecycle(Bucket=s3_persistent_bucket)

    async def put_lifecycle_config(lc_rules: list[dict[str, Any]]) -> None:
        await s3_client.put_bucket_lifecycle_configuration(
            Bucket=s3_persistent_bucket,
            LifecycleConfiguration=dict(
                Rules=lc_rules,
            ),
        )

    rules = [make_cleanup_rule(str(uuid.uuid4())) for _ in range(100)]
    await put_lifecycle_config(rules)

    task = await task_processor_client.schedule(CleanS3LifecycleRulesTask())
    result = await wait_task(task, task_state)
    assert result[-1] == "success"
    new_n_lc_rules = await get_lc_rules_number(s3_client, s3_persistent_bucket)
    assert new_n_lc_rules == 0

    rules = [make_cleanup_rule(str(uuid.uuid4())) for _ in range(100)] + [make_cleanup_rule(tenant_ids[0])]
    await put_lifecycle_config(rules)

    task = await task_processor_client.schedule(CleanS3LifecycleRulesTask())
    result = await wait_task(task, task_state)
    assert result[-1] == "success"
    new_n_lc_rules = await get_lc_rules_number(s3_client, s3_persistent_bucket)
    assert new_n_lc_rules == 1


@pytest.mark.asyncio
async def test_datetime64(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    default_async_usm_per_test,
):
    usm = default_async_usm_per_test
    csv_data = """utc_start_dttm
2022-06-01 00:36:55.845368
2022-06-01 00:39:00.91446
2022-06-01 00:52:03.14199
2022-06-01 10:48:45.272827""".encode(
        "utf-8"
    )

    rmm = redis_model_manager
    df = DataFile(
        manager=rmm,
        filename="test_file.csv",
        file_type=FileType.csv,
        size=len(csv_data),
        status=FileProcessingStatus.in_progress,
    )
    await s3_client.put_object(
        ACL="private",
        Bucket=s3_tmp_bucket,
        Key=df.s3_key,
        Body=csv_data,
    )
    await df.save()

    df = await DataFile.get(manager=rmm, obj_id=df.id)
    assert df.status == FileProcessingStatus.in_progress

    task = await task_processor_client.schedule(ParseFileTask(file_id=df.id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    df = await DataFile.get(manager=rmm, obj_id=df.id)
    assert df.status == FileProcessingStatus.ready

    assert len(df.sources) == 1
    source = df.sources[0]
    assert source.status == FileProcessingStatus.ready

    conn = await create_file_connection(usm, df.id, source.id, source.raw_schema)
    assert conn.get_file_source_by_id(source.id).status == FileProcessingStatus.in_progress

    task_save = await task_processor_client.schedule(
        SaveSourceTask(
            tenant_id="common",
            file_id=df.id,
            src_source_id=source.id,
            dst_source_id=source.id,
            connection_id=conn.uuid,
        )
    )
    result = await wait_task(task_save, task_state)
    assert result[-1] == "success"

    conn = await usm.get_by_id(conn.uuid)
    assert conn.get_file_source_by_id(source.id).status == FileProcessingStatus.ready


@pytest.mark.asyncio
async def test_datetime_tz(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    default_async_usm_per_test,
):
    csv_data = """dt
2022-07-01 13:52:07+03
2022-07-01 14:34:36+03
2022-07-01 14:38:47+03
2022-07-02 12:33:20+03""".encode(
        "utf-8"
    )

    usm = default_async_usm_per_test
    rmm = redis_model_manager
    df = DataFile(
        manager=rmm,
        filename="test_file.csv",
        file_type=FileType.csv,
        size=len(csv_data),
        status=FileProcessingStatus.in_progress,
    )
    await s3_client.put_object(
        ACL="private",
        Bucket=s3_tmp_bucket,
        Key=df.s3_key,
        Body=csv_data,
    )
    await df.save()

    df = await DataFile.get(manager=rmm, obj_id=df.id)
    assert df.status == FileProcessingStatus.in_progress

    task = await task_processor_client.schedule(ParseFileTask(file_id=df.id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    df = await DataFile.get(manager=rmm, obj_id=df.id)
    assert df.status == FileProcessingStatus.ready

    assert len(df.sources) == 1
    source = df.sources[0]
    assert source.status == FileProcessingStatus.ready

    conn = await create_file_connection(usm, df.id, source.id, source.raw_schema)
    assert conn.get_file_source_by_id(source.id).status == FileProcessingStatus.in_progress

    task_save = await task_processor_client.schedule(
        SaveSourceTask(
            tenant_id="common",
            file_id=df.id,
            src_source_id=source.id,
            dst_source_id=source.id,
            connection_id=conn.uuid,
        )
    )
    result = await wait_task(task_save, task_state)
    assert result[-1] == "success"

    conn = await usm.get_by_id(conn.uuid)
    assert conn.get_file_source_by_id(source.id).status == FileProcessingStatus.ready


@pytest.mark.asyncio
async def test_cleanup_tenant_file_previews_task(
    redis_model_manager,
    task_processor_client,
    task_state,
    s3_client,
    s3_persistent_bucket,
):
    rmm = redis_model_manager

    tenant_id = str(uuid.uuid4())[:20]
    preview_set = PreviewSet(redis=rmm._redis, id=tenant_id)
    n_previews = 5
    for _ in range(n_previews):
        preview = DataSourcePreview(manager=rmm, preview_data=[])
        await preview.save()
        await preview_set.add(preview.id)

    # just making sure the previews are there and saving their ids
    set_size = await preview_set.size()
    assert set_size == n_previews
    preview_ids = []
    async for preview_id in preview_set.sscan_iter():
        preview = await DataSourcePreview.get(manager=rmm, obj_id=preview_id)
        preview_ids.append(preview.id)

    task = await task_processor_client.schedule(CleanupTenantFilePreviewsTask(tenant_id=tenant_id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    set_size = await preview_set.size()
    assert set_size == 0

    for preview_id in preview_ids:
        with pytest.raises(RedisModelNotFound):
            await DataSourcePreview.get(manager=rmm, obj_id=preview_id)


@pytest.mark.asyncio
async def test_too_many_columns_csv(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    monkeypatch,
):
    monkeypatch.setattr(parsing_utils, "MAX_COLUMNS_COUNT", 10)
    csv_data = generate_sample_csv_data_str(3, 11).encode("utf-8")

    rmm = redis_model_manager
    df = DataFile(
        manager=rmm,
        filename="too_many_columns.csv",
        file_type=FileType.csv,
        size=len(csv_data),
        status=FileProcessingStatus.in_progress,
    )
    await s3_client.put_object(
        ACL="private",
        Bucket=s3_tmp_bucket,
        Key=df.s3_key,
        Body=csv_data,
    )
    await df.save()

    df = await DataFile.get(manager=rmm, obj_id=df.id)
    assert df.status == FileProcessingStatus.in_progress

    task = await task_processor_client.schedule(ParseFileTask(file_id=df.id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    df = await DataFile.get(manager=rmm, obj_id=df.id)
    assert df.status == FileProcessingStatus.failed
    assert df.error is not None and df.error.code == exc.TooManyColumnsError.err_code

    assert len(df.sources) == 1
    assert df.sources[0].error.code == exc.TooManyColumnsError.err_code


@pytest.mark.asyncio
async def test_rename_tenant_files(
    saved_file_connection_id,
    task_processor_client,
    task_state,
    s3_client,
    s3_persistent_bucket,
    default_async_usm_per_test,
    redis_model_manager,
):
    usm = default_async_usm_per_test
    conn = await usm.get_by_id(saved_file_connection_id)
    source = conn.data.sources[0]
    assert source.s3_filename
    s3_obj = await s3_client.get_object(Bucket=s3_persistent_bucket, Key=source.s3_filename)
    s3_data = await s3_obj["Body"].read()
    preview_set = PreviewSet(redis=redis_model_manager._redis, id=conn.raw_tenant_id)
    ps_vals = {_ async for _ in preview_set.sscan_iter()}
    assert len(ps_vals) >= 1

    new_tenant_id = str(uuid.uuid4())
    task = await task_processor_client.schedule(RenameTenantFilesTask(tenant_id=new_tenant_id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    updated_conn = await usm.get_by_id(saved_file_connection_id)
    updated_source = updated_conn.get_file_source_by_id(source.id)
    assert updated_source.s3_filename
    assert updated_source.s3_filename.startswith(new_tenant_id)
    updated_s3_obj = await s3_client.get_object(Bucket=s3_persistent_bucket, Key=updated_source.s3_filename)
    updated_s3_obj_data = await updated_s3_obj["Body"].read()
    assert s3_data == updated_s3_obj_data

    status_obj = await RenameTenantStatusModel.get(manager=redis_model_manager, obj_id=new_tenant_id)
    assert status_obj.status == RenameTenantStatus.success

    new_preview_set = PreviewSet(redis=redis_model_manager._redis, id=new_tenant_id)
    nps_vals = {_ async for _ in new_preview_set.sscan_iter()}
    assert nps_vals >= ps_vals
    assert await preview_set.size() == 0

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/db/test_tasks.py
datalens-tech/datalens-backend
import uuid

from dl_constants.enums import (
    DataSourceRole,
    FileProcessingStatus,
)
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_core_testing.connection import make_conn_key

from dl_connector_bundle_chs3.file.core.constants import CONNECTION_TYPE_FILE
from dl_connector_bundle_chs3.file.core.us_connection import FileS3Connection


async def create_file_connection(us_manager: AsyncUSManager, file_id, source_id, raw_schema, src_title="Source 1"):
    conn_name = "file test conn {}".format(uuid.uuid4())
    data = FileS3Connection.DataModel(
        sources=[
            FileS3Connection.FileDataSource(
                id=source_id,
                file_id=file_id,
                title=src_title,
                status=FileProcessingStatus.in_progress,
            )
        ]
    )
    conn = FileS3Connection.create_from_dict(
        data,
        ds_key=make_conn_key("connections", conn_name),
        type_=CONNECTION_TYPE_FILE.name,
        meta={"title": conn_name, "state": "saved"},
        us_manager=us_manager,
    )
    conn.update_data_source(
        source_id,
        role=DataSourceRole.origin,
        raw_schema=raw_schema,
    )
    await us_manager.save(conn)
    return conn

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/db/utils.py
datalens-tech/datalens-backend
import os

import pytest

from dl_configs.crypto_keys import get_dummy_crypto_keys_config
from dl_configs.settings_submodels import (
    GoogleAppSettings,
    S3Settings,
)
from dl_file_uploader_worker_lib.settings import FileUploaderWorkerSettings
from dl_testing.env_params.generic import GenericEnvParamGetter


@pytest.fixture(scope="session")
def env_param_getter() -> GenericEnvParamGetter:
    filepath = os.path.join(os.path.dirname(__file__), "params.yml")
    return GenericEnvParamGetter.from_yaml_file(filepath)


# overriding top-level fixture to set google app settings that are only needed in ext tests
@pytest.fixture(scope="session")
def file_uploader_worker_settings(
    env_param_getter,
    redis_app_settings,
    redis_arq_settings,
    s3_settings,
    connectors_settings,
    us_config,
    secure_reader,
):
    settings = FileUploaderWorkerSettings(
        REDIS_APP=redis_app_settings,
        REDIS_ARQ=redis_arq_settings,
        S3=S3Settings(
            ENDPOINT_URL=s3_settings.ENDPOINT_URL,
            ACCESS_KEY_ID=s3_settings.ACCESS_KEY_ID,
            SECRET_ACCESS_KEY=s3_settings.SECRET_ACCESS_KEY,
        ),
        S3_TMP_BUCKET_NAME="bi-file-uploader-tmp",
        S3_PERSISTENT_BUCKET_NAME="bi-file-uploader",
        SENTRY_DSN=None,
        US_BASE_URL=us_config.base_url,
        US_MASTER_TOKEN=us_config.master_token,
        CONNECTORS=connectors_settings,
        GSHEETS_APP=GoogleAppSettings(
            API_KEY=env_param_getter.get_str_value("GOOGLE_API_KEY"),
            CLIENT_ID="dummy",  # TODO test auth properly
            CLIENT_SECRET="dummy",
        ),
        CRYPTO_KEYS_CONFIG=get_dummy_crypto_keys_config(),
        SECURE_READER=secure_reader,
    )
    yield settings

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/ext/conftest.py
datalens-tech/datalens-backend
import asyncio
import logging
import uuid

import pytest

from dl_constants.enums import (
    FileProcessingStatus,
    UserDataType,
)
from dl_core.us_manager.us_manager_async import AsyncUSManager
from dl_core_testing.connection import make_conn_key
from dl_file_uploader_lib import exc
from dl_file_uploader_lib.data_sink.json_each_row import S3JsonEachRowUntypedFileAsyncDataSink
from dl_file_uploader_lib.enums import FileType
from dl_file_uploader_lib.redis_model.base import RedisModelManager
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    DataSource,
    DataSourcePreview,
    GSheetsFileSourceSettings,
    GSheetsUserSourceDataSourceProperties,
    GSheetsUserSourceProperties,
)
from dl_file_uploader_task_interface.tasks import (
    DownloadGSheetTask,
    ParseFileTask,
    SaveSourceTask,
)
from dl_file_uploader_worker_lib.utils import parsing_utils
from dl_file_uploader_worker_lib.utils.converter_parsing_utils import idx_to_alphabet_notation
from dl_task_processor.state import wait_task

from dl_connector_bundle_chs3.chs3_gsheets.core.constants import CONNECTION_TYPE_GSHEETS_V2
from dl_connector_bundle_chs3.chs3_gsheets.core.us_connection import GSheetsFileS3Connection


LOGGER = logging.getLogger(__name__)


# UserDataType aliases
STR = UserDataType.string
INT = UserDataType.integer
F = UserDataType.float
D = UserDataType.date
DT = UserDataType.genericdatetime
B = UserDataType.boolean

SPREADHEET_ID = "1rnUFa7AiSKD5O80IKCvMy2cSZvLU1kRw9dxbtZbDMWc"

TEST_SHEET_TITLES_INDICES = {
    title: idx
    for idx, title in enumerate(
        (
            "image",
            "experiment",
            "corners",
            "elaborate",
            "empty titles",
            "elaborate (no number format)",
        )
    )
}

EXPECTED_BODY_TYPES = {  # { <sheet_title>: [ <user_type1>, <user_type2>, ... ] }
    "image": [],
    "experiment": [STR] * 10,
    "corners": [STR] * 5,
    "elaborate": [INT, D, F, STR, B, DT, STR, STR],
    "empty titles": [INT, STR, STR, STR, STR],
    "elaborate (no number format)": [STR] * 10,
}

EXPECTED_HEADERS = {  # { <sheet_title>: { <has_header>: [ <title1>, <title2>, ... ] } }
    "elaborate": {
        True: [
            "int field",
            "date field",
            "float field",
            "string field",
            "bool field",
            "datetime field",
            "time field",
            "time period field",
        ],
        False: ["A", "B", "C", "D", "E", "F", "G", "H"],
    },
    "empty titles": {
        True: ["int field", "B", "another field to test empty titles", "D", "E"],
        False: ["A", "B", "C", "D", "E"],
    },
}

_MOSCOWSHOPS_V8_TEST_PARAMS = [
    [True, True, False, True, False],
    [
        [
            "OrderID",
            "ClientID",
            "ProductID",
            "ShopID",
            "ClientAdress",
            "ShopAdress",
            "ShopName",
            "ProductBrend",
            "ProductCategory",
            "DayOfBirth",
            "OrderDatetime",
            "DeliveryDatetime",
            "Discount",
            "ProductCount",
            "DeliveryAdressCoord",
            "ShopAdressCoord",
            "ProductName",
            "ClientName",
            "DeliveryType",
            "PaymentType",
            "ProductPrice",
            "FinelAmount",
            "Gender",
            "ProductSubcategory",
            "Amount",
            "ClientStatus",
        ],
        [
            "OrderID",
            "ClientID",
            "ProductID",
            "ShopID",
            "OrderDatetime",
            "DeliveryDatetime",
            "Discount",
            "ProductCount",
            "DeliveryAdressCoord",
            "DeliveryType",
            "PaymentType",
            "Price",
            "FinelAmount",
            "Amount",
        ],
        [idx_to_alphabet_notation(i) for i in range(4)],  # ['ShopID', 'ShopName', 'ShopAdress', 'DeliveryAdressCoord'],
        ["ClientID", "ClientName", "Gender", "ClientStatus", "DayOfBirth", "ClientAdress", "BornBefore1980"],
        [
            idx_to_alphabet_notation(i) for i in range(5)
        ],  # ['ProductID', 'ProductCategory', 'ProductSubcategory', 'ProductBrend', 'ProductName'],
    ],
    [
        [STR] * 9 + [D, DT, STR, INT, INT] + [STR] * 6 + [F, INT, STR, STR, INT, STR],
        [STR, STR, STR, STR, DT, STR, INT, INT, STR, STR, STR, F, INT, INT],
        [STR, STR, STR, STR],
        [STR, STR, STR, STR, D, STR, B],
        [STR, STR, STR, STR, STR],
    ],
]

_big_gsheets_test_params = dict(
    argvalues=[
        (
            "1SEwxMwOqieNAgybrR_tGbt_N60mPp00T3LawutJ83RA",
            *_MOSCOWSHOPS_V8_TEST_PARAMS,
        ),
        (
            "1sPIpWUZa7wgUnUDa-MFPidJTSjQ40jS2OmD0g6V_wDw",
            *_MOSCOWSHOPS_V8_TEST_PARAMS,
        ),
        (
            "1OK6beKLQx8JqioiNTRN90OK1A9wN3bL6Xn0cwaWp0rw",
            *_MOSCOWSHOPS_V8_TEST_PARAMS,
        ),
        (
            "1G3BFi2Ai5jsNfazpwYCzZ5OUUiCMNpG_5BBMlFnpYy4",
            [False],
            [[idx_to_alphabet_notation(i) for i in range(17)]],
            [[D] + [STR] * 6 + [INT] * 8 + [F, B]],
        ),
    ],
    ids=[
        "MoscowShopsV8_tiny",
        "MoscowShopsV8_truncated",
        "MoscowShopsV8_full",
        "no_header",
    ],
)


@pytest.fixture(scope="function")
async def downloaded_gsheet_file_id(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
):
    df = DataFile(
        manager=redis_model_manager,
        filename="fu",
        status=FileProcessingStatus.in_progress,
        file_type=FileType.gsheets,
        user_source_properties=GSheetsUserSourceProperties(spreadsheet_id=SPREADHEET_ID),
    )
    await df.save()

    task = await task_processor_client.schedule(
        DownloadGSheetTask(file_id=df.id, authorized=False, schedule_parsing=False)
    )
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    yield df.id

    df = await DataFile.get(manager=redis_model_manager, obj_id=df.id)
    for source in df.sources:
        await s3_client.delete_object(
            Bucket=s3_tmp_bucket,
            Key=source.s3_key,
        )


@pytest.mark.asyncio
async def test_download_gsheet_task(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
):
    df = DataFile(
        manager=redis_model_manager,
        filename="fu",
        status=FileProcessingStatus.in_progress,
        file_type=FileType.gsheets,
        user_source_properties=GSheetsUserSourceProperties(spreadsheet_id=SPREADHEET_ID),
    )
    await df.save()

    task = await task_processor_client.schedule(
        DownloadGSheetTask(file_id=df.id, authorized=False, schedule_parsing=False)
    )
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    df = await DataFile.get(manager=redis_model_manager, obj_id=df.id)
    assert df.status == FileProcessingStatus.in_progress

    assert not df.sources[TEST_SHEET_TITLES_INDICES["image"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["experiment"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["corners"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["elaborate"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["empty titles"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["elaborate (no number format)"]].is_applicable

    # test elaborate sheet types
    elaborate_source = df.sources[TEST_SHEET_TITLES_INDICES["elaborate"]]
    actual_user_types = [col.user_type for col in elaborate_source.raw_schema]
    expected_user_types = [
        UserDataType.integer,
        UserDataType.date,
        UserDataType.float,
        UserDataType.string,
        UserDataType.boolean,
        UserDataType.genericdatetime,
        UserDataType.string,
        UserDataType.string,
    ]
    assert actual_user_types == expected_user_types

    for source in df.sources:
        await s3_client.delete_object(
            Bucket=s3_tmp_bucket,
            Key=source.s3_key,
        )


@pytest.mark.asyncio
async def test_parse_gsheet(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    downloaded_gsheet_file_id,
):
    file_id = downloaded_gsheet_file_id
    task = await task_processor_client.schedule(ParseFileTask(file_id=file_id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    df = await DataFile.get(manager=redis_model_manager, obj_id=file_id)
    assert df.status == FileProcessingStatus.ready

    elaborate_source_no_types = df.sources[TEST_SHEET_TITLES_INDICES["elaborate (no number format)"]]
    assert elaborate_source_no_types.status == FileProcessingStatus.ready

    actual_user_types = [col.user_type for col in elaborate_source_no_types.raw_schema]
    expected_user_types = [UserDataType.string] * 10
    assert actual_user_types == expected_user_types

    preview = await DataSourcePreview.get(manager=redis_model_manager, obj_id=elaborate_source_no_types.preview_id)
    assert preview.id == elaborate_source_no_types.preview_id
    assert preview.preview_data == [  # No number formats specified => everything is string => no header
        [
            "int field",
            "date normal",
            "date gsheet",
            "float field",
            "string field",
            "bool field",
            "dt normal",
            "dt gsheet",
            "time field",
            "time period field",
        ],
        [
            "1",
            "2001-07-31",
            "31.07.2001",
            "12,50",
            "asdf",
            "True",
            "2022-08-12 12:30:32",
            "12.08.2022 12:30:32",
            "12:00:12",
            "13:00:12",
        ],
        [None, None, None, None, None, None, None, None, None, None],
        [
            "2",
            "1995-04-12",
            "12.04.1995",
            "12,40",
            "zxcv",
            "True",
            "2022-08-14 12:30:34",
            "14.08.2022 12:30:34",
            "14:00:12",
            "15:00:12",
        ],
        [
            "8",
            "2008-02-29",
            "29.02.2008",
            None,
            "12,00",
            "True",
            "2022-08-15 12:30:35",
            "15.08.2022 12:30:35",
            "15:00:12",
            "16:00:12",
        ],
        [
            "3",
            "2012-03-12",
            "12.03.2012",
            "12,30",
            "dfgh",
            "True",
            "2022-08-16 12:30:36",
            "16.08.2022 12:30:36",
            "16:00:12",
            "17:00:12",
        ],
        [None, None, None, None, None, "StRiNg", None, None, None, "18:00:12"],
    ]


async def assert_parsing_results(
    file_id: str,
    has_header_expected: bool,
    rmm: RedisModelManager,
    dsrc_title: str,
    sheet_len: int,
) -> DataSource:
    df = await DataFile.get(manager=rmm, obj_id=file_id)
    assert df.status == FileProcessingStatus.ready

    dsrc = df.sources[TEST_SHEET_TITLES_INDICES[dsrc_title]]
    assert dsrc.status == FileProcessingStatus.ready

    file_source_settings = dsrc.file_source_settings
    assert isinstance(file_source_settings, GSheetsFileSourceSettings)

    preview = await DataSourcePreview.get(manager=rmm, obj_id=dsrc.preview_id)
    assert preview.id == dsrc.preview_id

    assert file_source_settings.first_line_is_header == has_header_expected
    if not has_header_expected:
        # mixed header => everything is string
        assert all(col.user_type == UserDataType.string for col in dsrc.raw_schema)
        assert len(preview.preview_data) == sheet_len
    else:
        # set header => correct type
        assert [sch_col.user_type for sch_col in dsrc.raw_schema] == EXPECTED_BODY_TYPES[dsrc_title]
        assert len(preview.preview_data) == sheet_len - 1  # preview becomes one line shorter

    assert [sch_col.title for sch_col in dsrc.raw_schema] == EXPECTED_HEADERS[dsrc_title][has_header_expected]

    return dsrc


@pytest.mark.asyncio
async def test_parse_gsheet_with_file_settings(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    downloaded_gsheet_file_id,
):
    """Testing first_line_is_header override: None (which is False) -> True -> False"""

    file_id = downloaded_gsheet_file_id
    rmm = redis_model_manager
    sheet_len = 7
    sheet_title = "empty titles"

    task = await task_processor_client.schedule(ParseFileTask(file_id=file_id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"
    empty_titles_dsrc = await assert_parsing_results(file_id, False, rmm, sheet_title, sheet_len)

    task = await task_processor_client.schedule(
        ParseFileTask(
            file_id=file_id,
            source_id=empty_titles_dsrc.id,
            file_settings=dict(first_line_is_header=True),
            source_settings={},
        )
    )
    result = await wait_task(task, task_state)
    assert result[-1] == "success"
    empty_titles_dsrc = await assert_parsing_results(file_id, True, rmm, sheet_title, sheet_len)

    task = await task_processor_client.schedule(
        ParseFileTask(
            file_id=file_id,
            source_id=empty_titles_dsrc.id,
            file_settings=dict(first_line_is_header=False),
            source_settings={},
        )
    )
    result = await wait_task(task, task_state)
    assert result[-1] == "success"
    await assert_parsing_results(file_id, False, rmm, sheet_title, sheet_len)


@pytest.mark.asyncio
async def test_parse_gsheet_with_file_settings_reverse(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    downloaded_gsheet_file_id,
):
    """Testing first_line_is_header override: None (which is True) -> False -> True"""

    file_id = downloaded_gsheet_file_id
    rmm = redis_model_manager
    sheet_len = 7
    sheet_title = "elaborate"

    task = await task_processor_client.schedule(ParseFileTask(file_id=file_id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"
    elaborate_no_number_format_dsrc = await assert_parsing_results(file_id, True, rmm, sheet_title, sheet_len)

    task = await task_processor_client.schedule(
        ParseFileTask(
            file_id=file_id,
            source_id=elaborate_no_number_format_dsrc.id,
            file_settings=dict(first_line_is_header=False),
            source_settings={},
        )
    )
    result = await wait_task(task, task_state)
    assert result[-1] == "success"
    elaborate_no_number_format_dsrc = await assert_parsing_results(file_id, False, rmm, sheet_title, sheet_len)

    task = await task_processor_client.schedule(
        ParseFileTask(
            file_id=file_id,
            source_id=elaborate_no_number_format_dsrc.id,
            file_settings=dict(first_line_is_header=True),
            source_settings={},
        )
    )
    result = await wait_task(task, task_state)
    assert result[-1] == "success"
    await assert_parsing_results(file_id, True, rmm, sheet_title, sheet_len)


@pytest.mark.asyncio
async def test_download_and_parse_gsheet(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
):
    df = DataFile(
        manager=redis_model_manager,
        filename="fu",
        status=FileProcessingStatus.in_progress,
        file_type=FileType.gsheets,
        user_source_properties=GSheetsUserSourceProperties(spreadsheet_id=SPREADHEET_ID),
    )
    await df.save()

    task = await task_processor_client.schedule(DownloadGSheetTask(file_id=df.id, authorized=False))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    await asyncio.sleep(15)

    df = await DataFile.get(manager=redis_model_manager, obj_id=df.id)
    assert df.status == FileProcessingStatus.ready

    assert not df.sources[TEST_SHEET_TITLES_INDICES["image"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["experiment"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["corners"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["elaborate"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["empty titles"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["elaborate (no number format)"]].is_applicable

    for source in df.sources:
        await s3_client.delete_object(
            Bucket=s3_tmp_bucket,
            Key=source.s3_key,
        )


async def create_gsheets_v2_connection(us_manager: AsyncUSManager, dfile: DataFile):
    conn_name = "gsheets_v2 test conn {}".format(uuid.uuid4())
    conn_sources = []
    for src in dfile.sources:
        if not src.is_applicable:
            continue
        user_source_properties = dfile.user_source_properties
        assert isinstance(user_source_properties, GSheetsUserSourceProperties)
        file_source_settings = src.file_source_settings
        assert isinstance(file_source_settings, GSheetsFileSourceSettings)
        user_source_dsrc_properties = src.user_source_dsrc_properties
        assert isinstance(user_source_dsrc_properties, GSheetsUserSourceDataSourceProperties)

        conn_sources.append(
            GSheetsFileS3Connection.FileDataSource(
                id=src.id,
                file_id=dfile.id,
                title=src.title,
                status=FileProcessingStatus.in_progress,
                spreadsheet_id=user_source_properties.spreadsheet_id,
                sheet_id=user_source_dsrc_properties.sheet_id,
                first_line_is_header=file_source_settings.first_line_is_header,
                raw_schema=src.raw_schema,
            )
        )
    data = GSheetsFileS3Connection.DataModel(sources=conn_sources, refresh_enabled=False)
    conn = GSheetsFileS3Connection.create_from_dict(
        data,
        ds_key=make_conn_key("connections", conn_name),
        type_=CONNECTION_TYPE_GSHEETS_V2.name,
        meta={"title": conn_name, "state": "saved"},
        us_manager=us_manager,
    )
    await us_manager.save(conn)
    return conn


@pytest.mark.asyncio
async def test_save_source_task(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    downloaded_gsheet_file_id,
    default_async_usm_per_test,
    s3_tmp_bucket,
):
    file_id = downloaded_gsheet_file_id
    usm = default_async_usm_per_test
    task = await task_processor_client.schedule(ParseFileTask(file_id=downloaded_gsheet_file_id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    df = await DataFile.get(manager=redis_model_manager, obj_id=file_id)
    assert df.sources[0].status == FileProcessingStatus.failed
    assert all(source.status == FileProcessingStatus.ready for source in df.sources[1:])

    conn = await create_gsheets_v2_connection(usm, df)
    assert all(
        conn.get_file_source_by_id(source.id).status == FileProcessingStatus.in_progress for source in df.sources[1:]
    )

    for src in df.sources:
        if not src.is_applicable:
            continue

        task_save = await task_processor_client.schedule(
            SaveSourceTask(
                tenant_id="common",
                file_id=file_id,
                src_source_id=src.id,
                dst_source_id=src.id,
                connection_id=conn.uuid,
            )
        )
        await wait_task(task_save, task_state)

        conn1 = await usm.get_by_id(conn.uuid)
        assert conn1.get_file_source_by_id(src.id).status == FileProcessingStatus.ready


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "spreadsheet_id, expected_has_header, expected_headers, expected_user_types",
    **_big_gsheets_test_params,
)
async def test_download_and_parse_big_gsheets(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    spreadsheet_id,
    expected_has_header,
    expected_headers,
    expected_user_types,
):
    assert len(expected_has_header) == len(expected_headers) == len(expected_user_types), "malformed test data"

    df = DataFile(
        manager=redis_model_manager,
        filename="fu",
        status=FileProcessingStatus.in_progress,
        file_type=FileType.gsheets,
        user_source_properties=GSheetsUserSourceProperties(spreadsheet_id=spreadsheet_id),
    )
    await df.save()

    # Download and Parse
    task = await task_processor_client.schedule(DownloadGSheetTask(file_id=df.id, authorized=False))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    await asyncio.sleep(15)  # let parse finish

    df = await DataFile.get(manager=redis_model_manager, obj_id=df.id)
    assert df.status == FileProcessingStatus.ready

    assert len(df.sources) == len(expected_headers)

    for src in df.sources:
        assert src.is_applicable

    for src, has_header, header_values, user_types in zip(
        df.sources, expected_has_header, expected_headers, expected_user_types
    ):
        assert src.file_source_settings.first_line_is_header == has_header, src.title

        assert len(src.raw_schema) == len(header_values), src.title

        actual_source_columns = [col.name for col in src.raw_schema]
        assert actual_source_columns == [idx_to_alphabet_notation(i) for i in range(len(actual_source_columns))]

        actual_header_values = [col.title for col in src.raw_schema]
        assert actual_header_values == header_values, src.title

        actual_user_types = [col.user_type for col in src.raw_schema]
        assert actual_user_types == user_types, src.title

    for source in df.sources:
        await s3_client.delete_object(
            Bucket=s3_tmp_bucket,
            Key=source.s3_key,
        )


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "spreadsheet_id, expected_has_header, expected_headers, expected_user_types",
    **_big_gsheets_test_params,
)
async def test_gsheets_full_pipeline(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    s3_persistent_bucket,
    spreadsheet_id,
    expected_has_header,
    expected_headers,
    expected_user_types,
    default_async_usm_per_test,
):
    assert len(expected_has_header) == len(expected_headers) == len(expected_user_types), "malformed test data"

    df = DataFile(
        manager=redis_model_manager,
        filename="fu",
        status=FileProcessingStatus.in_progress,
        file_type=FileType.gsheets,
        user_source_properties=GSheetsUserSourceProperties(spreadsheet_id=spreadsheet_id),
    )
    await df.save()

    # Download and Parse
    task = await task_processor_client.schedule(DownloadGSheetTask(file_id=df.id, authorized=False))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    await asyncio.sleep(5)  # let parse finish

    df = await DataFile.get(manager=redis_model_manager, obj_id=df.id)
    assert df.status == FileProcessingStatus.ready

    assert len(df.sources) == len(expected_headers)

    for src in df.sources:
        assert src.is_applicable

    for src, has_header, header_values, user_types in zip(
        df.sources, expected_has_header, expected_headers, expected_user_types
    ):
        assert src.file_source_settings.first_line_is_header == has_header, src.title

        assert len(src.raw_schema) == len(header_values), src.title

        actual_source_columns = [col.name for col in src.raw_schema]
        assert actual_source_columns == [idx_to_alphabet_notation(i) for i in range(len(actual_source_columns))]

        actual_header_values = [col.title for col in src.raw_schema]
        assert actual_header_values == header_values, src.title

        actual_user_types = [col.user_type for col in src.raw_schema]
        assert actual_user_types == user_types, src.title

    # Save
    file_id = df.id
    usm = default_async_usm_per_test
    task = await task_processor_client.schedule(ParseFileTask(file_id=file_id))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    df = await DataFile.get(manager=redis_model_manager, obj_id=file_id)
    assert all(source.status == FileProcessingStatus.ready for source in df.sources)

    conn = await create_gsheets_v2_connection(usm, df)
    assert all(
        conn.get_file_source_by_id(source.id).status == FileProcessingStatus.in_progress for source in df.sources
    )

    for src in df.sources:
        task_save = await task_processor_client.schedule(
            SaveSourceTask(
                tenant_id="common",
                file_id=file_id,
                src_source_id=src.id,
                dst_source_id=src.id,
                connection_id=conn.uuid,
            )
        )
        result = await wait_task(task_save, task_state)
        assert result[-1] == "success"

        updated_conn: GSheetsFileS3Connection = await usm.get_by_id(conn.uuid)
        assert updated_conn.get_file_source_by_id(src.id).status == FileProcessingStatus.ready

        for conn_src, has_header, header_values, user_types in zip(
            updated_conn.data.sources, expected_has_header, expected_headers, expected_user_types
        ):
            assert conn_src.first_line_is_header == has_header, src.title

            assert len(conn_src.raw_schema) == len(header_values), src.title

            actual_source_columns = [col.name for col in conn_src.raw_schema]
            assert actual_source_columns == [idx_to_alphabet_notation(i) for i in range(len(actual_source_columns))]

            actual_header_values = [col.title for col in conn_src.raw_schema]
            assert actual_header_values == header_values, src.title

            actual_user_types = [col.user_type for col in conn_src.raw_schema]
            assert actual_user_types == user_types, src.title

    for source in df.sources:
        await s3_client.delete_object(
            Bucket=s3_tmp_bucket,
            Key=source.s3_key,
        )


@pytest.mark.asyncio
async def test_too_many_columns_gsheets(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    monkeypatch,
):
    monkeypatch.setattr(parsing_utils, "MAX_COLUMNS_COUNT", 5)

    df = DataFile(
        manager=redis_model_manager,
        filename="too_many_columns",
        status=FileProcessingStatus.in_progress,
        file_type=FileType.gsheets,
        user_source_properties=GSheetsUserSourceProperties(spreadsheet_id=SPREADHEET_ID),
    )
    await df.save()

    task = await task_processor_client.schedule(DownloadGSheetTask(file_id=df.id, authorized=False))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    await asyncio.sleep(5)

    df = await DataFile.get(manager=redis_model_manager, obj_id=df.id)
    assert df.status == FileProcessingStatus.ready

    assert not df.sources[TEST_SHEET_TITLES_INDICES["image"]].is_applicable
    assert not df.sources[TEST_SHEET_TITLES_INDICES["experiment"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["corners"]].is_applicable
    assert not df.sources[TEST_SHEET_TITLES_INDICES["elaborate"]].is_applicable
    assert df.sources[TEST_SHEET_TITLES_INDICES["empty titles"]].is_applicable
    assert not df.sources[TEST_SHEET_TITLES_INDICES["elaborate (no number format)"]].is_applicable

    too_many_cols_sources_indices = (
        TEST_SHEET_TITLES_INDICES["experiment"],
        TEST_SHEET_TITLES_INDICES["elaborate"],
        TEST_SHEET_TITLES_INDICES["elaborate (no number format)"],
    )

    assert df.sources[TEST_SHEET_TITLES_INDICES["image"]].error.code == exc.EmptyDocument.err_code
    for idx in too_many_cols_sources_indices:
        assert df.sources[idx].error.code == exc.TooManyColumnsError.err_code, df.sources[idx].title

    for source in df.sources:
        await s3_client.delete_object(
            Bucket=s3_tmp_bucket,
            Key=source.s3_key,
        )


@pytest.mark.asyncio
async def test_too_large_sheet(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    monkeypatch,
):
    monkeypatch.setattr(S3JsonEachRowUntypedFileAsyncDataSink, "max_file_size_bytes", 8 * 1024**2)

    df = DataFile(
        manager=redis_model_manager,
        filename="too_large_sheet",
        status=FileProcessingStatus.in_progress,
        file_type=FileType.gsheets,
        user_source_properties=GSheetsUserSourceProperties(
            spreadsheet_id="1dfUajkRxGTxGXrwjXtCRWhEjHlZX9tMVEvla5fio93w"
        ),
    )
    await df.save()

    task = await task_processor_client.schedule(DownloadGSheetTask(file_id=df.id, authorized=False))
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

    await asyncio.sleep(5)

    df = await DataFile.get(manager=redis_model_manager, obj_id=df.id)
    assert df.status == FileProcessingStatus.ready

    assert not df.sources[0].is_applicable
    assert df.sources[1].is_applicable

    assert df.sources[0].error.code == exc.FileLimitError.err_code
    assert df.sources[1].error is None

    for source in df.sources:
        await s3_client.delete_object(
            Bucket=s3_tmp_bucket,
            Key=source.s3_key,
        )


@pytest.mark.asyncio
async def test_url_params_encoding(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
    monkeypatch,
):
    df = DataFile(
        manager=redis_model_manager,
        filename="slashes and stuff",
        status=FileProcessingStatus.in_progress,
        file_type=FileType.gsheets,
        user_source_properties=GSheetsUserSourceProperties(
            spreadsheet_id="1dE5-owygpjK-F65qlCwUa1YHGzh4dzBhyajURSLMvYU"
        ),
    )
    await df.save()

    task = await task_processor_client.schedule(
        DownloadGSheetTask(file_id=df.id, authorized=False, schedule_parsing=False)
    )
    result = await wait_task(task, task_state)
    assert result[-1] == "success"

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/ext/test_gsheets.py
datalens-tech/datalens-backend
import pytest

from dl_constants.enums import FileProcessingStatus
from dl_file_uploader_lib.enums import FileType
from dl_file_uploader_lib.redis_model.models import (
    DataFile,
    YaDocsUserSourceProperties,
)
from dl_file_uploader_task_interface.tasks import DownloadYaDocsTask
from dl_task_processor.state import wait_task
from dl_testing.s3_utils import s3_file_exists


@pytest.mark.asyncio
async def test_download_yadocs_task(
    task_processor_client,
    task_state,
    s3_client,
    redis_model_manager,
    s3_tmp_bucket,
):
    df = DataFile(
        filename="",
        file_type=FileType.yadocs,
        manager=redis_model_manager,
        status=FileProcessingStatus.in_progress,
        user_source_properties=YaDocsUserSourceProperties(public_link="https://disk.yandex.lt/i/OyzdmFI0MUEEgA"),
    )
    await df.save()

    task = await task_processor_client.schedule(DownloadYaDocsTask(file_id=df.id))
    result = await wait_task(task, task_state)

    assert result[-1] == "success"
    assert await s3_file_exists(key=df.s3_key, bucket=s3_tmp_bucket, s3_client=s3_client)

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/ext/test_yadoсs.py
datalens-tech/datalens-backend
import csv
import datetime
import io
import random
import string

import pytest


@pytest.fixture(scope="session")
def sample_data_str() -> str:
    output = io.StringIO()
    csv_writer = csv.writer(output, dialect=csv.excel)
    csv_writer.writerow(["some string", "some int", "some float", "some date", "some datetime"])
    today = datetime.date.today()
    now = datetime.datetime.now()
    for _i in range(10000):
        csv_writer.writerow(
            [
                "".join(random.choices(string.ascii_uppercase + string.digits, k=random.randint(3, 1000))),
                random.randint(-100, 100),
                random.randint(1, 1000) * 1.0 / random.randint(1, 1000),
                (today - datetime.timedelta(days=random.randint(0, 100))).isoformat(),
                (now - datetime.timedelta(days=random.randint(0, 100), minutes=random.randint(0, 10000))).isoformat(),
            ]
        )
    return output.getvalue()


@pytest.fixture(scope="session")
def sample_data_bytes_utf8(sample_data_str) -> bytes:
    return sample_data_str.encode("utf-8")

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/unit/conftest.py
datalens-tech/datalens-backend
import csv

from dl_file_uploader_lib.enums import CSVEncoding
from dl_file_uploader_worker_lib.utils.parsing_utils import (
    make_upcropped_text_sample,
    prepare_preview,
)


def test_prepare_preview(sample_data_bytes_utf8: bytes):
    orig_rows = len(sample_data_bytes_utf8.decode("utf-8").splitlines())
    sample_str = make_upcropped_text_sample(sample_data_bytes_utf8, len(sample_data_bytes_utf8), CSVEncoding.utf8)
    assert len(sample_str.splitlines()) == orig_rows - 1  # without last row
    preview_data = prepare_preview(sample_str, csv.excel, True)
    assert len(preview_data) == 30

lib/dl_file_uploader_worker_lib/dl_file_uploader_worker_lib_tests/unit/test_parsing_utils.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Dict,
    Generic,
    Iterable,
    Optional,
    Set,
    TypeVar,
)

import attr

import dl_formula.core.exc as exc
from dl_formula.core.extract import NodeExtract
import dl_formula.core.nodes as nodes


def validate_node_is_extractable(node: nodes.FormulaItem) -> nodes.NodeExtract:
    if node.extract is None:
        raise exc.CacheError("Node has no extract")
    return node.extract


class NodeSet:
    __slots__ = ("_data_set",)

    def __init__(self, nodes: Iterable[nodes.FormulaItem] = ()):
        self._data_set: Set[NodeExtract] = set()
        for node in nodes:
            self.add(node)

    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return False
        return self._data_set == other._data_set

    def add(self, node: nodes.FormulaItem) -> None:
        extract = validate_node_is_extractable(node)
        self._data_set.add(extract)

    def remove(self, node: nodes.FormulaItem) -> None:
        if node.extract is None:
            return

        self._data_set.remove(node.extract)

    def __contains__(self, node: nodes.FormulaItem) -> bool:
        return node.extract in self._data_set

    def clear(self) -> None:
        self._data_set.clear()

    def __or__(self, other: NodeSet) -> NodeSet:
        result = NodeSet()
        result._data_set = self._data_set | other._data_set
        return result

    def __and__(self, other: NodeSet) -> NodeSet:
        result = NodeSet()
        result._data_set = self._data_set & other._data_set
        return result

    def __len__(self) -> int:
        return len(self._data_set)


_MAP_VALUE_TV = TypeVar("_MAP_VALUE_TV")


@attr.s
class NodeValueMap(Generic[_MAP_VALUE_TV]):
    _data: Dict[NodeExtract, _MAP_VALUE_TV] = attr.ib(factory=dict)

    def add(self, node: nodes.FormulaItem, value: _MAP_VALUE_TV) -> None:
        extract = validate_node_is_extractable(node)
        self._data[extract] = value

    def get(self, node: nodes.FormulaItem) -> Optional[_MAP_VALUE_TV]:
        if node.extract is None:
            return None
        return self._data.get(node.extract)

    def __contains__(self, node: nodes.FormulaItem) -> bool:
        return node.extract in self._data

    def clear(self) -> None:
        self._data.clear()

    def __len__(self) -> int:
        return len(self._data)

lib/dl_formula/dl_formula/collections.py
datalens-tech/datalens-backend
import abc
from typing import Callable

import attr
import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

from dl_formula.core.dialect import DialectCombo
from dl_formula.translation.sa_dialects import get_sa_dialect


@attr.s
class ColumnRenderer(abc.ABC):
    _dialect: DialectCombo = attr.ib(kw_only=True)
    _field_names: dict[str, tuple[str, ...]] = attr.ib(factory=dict)
    _quoter: Callable[[str], ClauseElement] = attr.ib(init=False)

    @_quoter.default
    def _make_quoter(self) -> Callable[[str], ClauseElement]:
        return get_sa_dialect(self._dialect).identifier_preparer.quote

    @abc.abstractmethod
    def make_column(self, name: str) -> sa.sql.ClauseElement:
        raise NotImplementedError


@attr.s
class DefaultColumnRenderer(ColumnRenderer):
    def make_column(self, name: str) -> sa.sql.ClauseElement:
        full_name_parts = self._field_names.get(name) or (name,)
        return sa.literal_column(".".join([self._quoter(part) for part in full_name_parts]))


@attr.s
class UnprefixedColumnRenderer(ColumnRenderer):
    def make_column(self, name: str) -> sa.sql.ClauseElement:
        full_name_parts = self._field_names.get(name) or (name,)
        return sa.literal_column(self._quoter(full_name_parts[-1]))

lib/dl_formula/dl_formula/connectors/base/column.py
datalens-tech/datalens-backend
from typing import (
    ClassVar,
    Iterable,
    Optional,
    Sequence,
    Type,
)

from sqlalchemy.engine.default import DefaultDialect

from dl_formula.connectors.base.column import (
    ColumnRenderer,
    DefaultColumnRenderer,
)
from dl_formula.connectors.base.context_processor import ContextPostprocessor
from dl_formula.connectors.base.literal import Literalizer
from dl_formula.connectors.base.type_constructor import (
    DefaultSATypeConstructor,
    SATypeConstructor,
)
from dl_formula.core.dialect import (
    DialectCombo,
    DialectName,
    DialectNamespace,
    StandardDialect,
)
from dl_formula.definitions.base import NodeTranslation


class FormulaConnector:
    dialect_ns_cls: ClassVar[Type[DialectNamespace]] = StandardDialect
    dialects: ClassVar[DialectCombo]
    default_dialect: ClassVar[Optional[DialectCombo]] = None
    op_definitions: ClassVar[Iterable[NodeTranslation]] = []
    literalizer_cls: ClassVar[Type[Literalizer]] = Literalizer
    column_renderer_cls: ClassVar[Type[ColumnRenderer]] = DefaultColumnRenderer
    context_processor_cls: ClassVar[Type[ContextPostprocessor]] = ContextPostprocessor
    type_constructor_cls: ClassVar[Type[SATypeConstructor]] = DefaultSATypeConstructor
    sa_dialect: ClassVar[DefaultDialect] = DefaultDialect()

    @classmethod
    def get_dialect_names(cls) -> Sequence[DialectName]:
        dialect_names = sorted({bit.name for bit in cls.dialects.bits}, key=lambda item: item.name)
        return dialect_names

    @classmethod
    def registration_hook(cls) -> None:
        """Do some non-standard stuff here. Think twice before implementing."""

lib/dl_formula/dl_formula/connectors/base/connector.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc

from sqlalchemy.sql.elements import ClauseElement

from dl_formula.core.datatype import DataType
from dl_formula.definitions.flags import (
    ContextFlag,
    ContextFlags,
)


class ContextPostprocessor(abc.ABC):
    def get_warnings(self, data_type: DataType, flags: ContextFlags) -> list[str]:
        result: list[str] = []
        if flags & ContextFlag.DEPRECATED:
            result.append("Usage of deprecated function")

        return result

    def postprocess_expression(
        self,
        data_type: DataType,
        expression: ClauseElement,
        flags: ContextFlags,
    ) -> tuple[ClauseElement, ContextFlags]:
        return expression, flags


class BooleanlessContextPostprocessor(ContextPostprocessor):
    @abc.abstractmethod
    def booleanize_expression(self, data_type: DataType, expression: ClauseElement) -> ClauseElement:
        raise NotImplementedError

    @abc.abstractmethod
    def debooleanize_expression(self, data_type: DataType, expression: ClauseElement) -> ClauseElement:
        raise NotImplementedError

    def postprocess_expression(
        self,
        data_type: DataType,
        expression: ClauseElement,
        flags: ContextFlags,
    ) -> tuple[ClauseElement, ContextFlags]:
        if flags & ContextFlag.IS_CONDITION and not flags & ContextFlag.REQ_CONDITION:
            # expression is a boolean condition, which cannot be used as a value
            expression = self.debooleanize_expression(data_type=data_type, expression=expression)
            flags = flags | ContextFlag.REQ_CONDITION

        elif flags & ContextFlag.REQ_CONDITION and not flags & ContextFlag.IS_CONDITION:
            # function requires a boolean condition, which the given expression is not, so booleanize!
            expression = self.booleanize_expression(data_type=data_type, expression=expression)
            flags = flags | ContextFlag.IS_CONDITION

        return expression, flags

lib/dl_formula/dl_formula/connectors/base/context_processor.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from functools import singledispatchmethod
from typing import (
    Any,
    Union,
)

import sqlalchemy as sa
import sqlalchemy.dialects.postgresql as sa_postgresql
from sqlalchemy.sql.elements import (
    BindParameter,
    Cast,
    Null,
)

from dl_formula.core.dialect import DialectCombo
from dl_formula.core.nodes import BaseLiteral


class TypeDefiningCast(Cast):
    """
    A mostly straightforward Cast subclass
    that marks the 'actual type was not changed' cases.

    Intended for telling the bindparameter type to postgres
    (in compeng asyncpg postgres prepared statements).
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        assert isinstance(self.clause, sa.sql.elements.BindParameter), "This Cast is only meant for constants"

    @property
    def value(self):
        """
        Compatibility wrapper (mimicks the BindParameter).

        TODO: Should use `un_literal` instead of `….value` everywhere.
        """
        return self.clause.value


Literal = Union[
    BaseLiteral,  # unwrapped formula node
    # various SA wrappers and types:
    BindParameter,
    TypeDefiningCast,
    sa.sql.functions.Function,  # CH cast via toDatTime
    Null,
    sa_postgresql.array,
]


class Literalizer:
    __slots__ = ()

    @singledispatchmethod
    def _literal(self, value: Any, dialect: DialectCombo) -> Literal:
        return sa.literal(value)

    @_literal.register
    def _literal_int(self, value: int, dialect: DialectCombo) -> Literal:
        return self.literal_int(value, dialect=dialect)

    @_literal.register
    def _literal_float(self, value: float, dialect: DialectCombo) -> Literal:
        return self.literal_float(value, dialect=dialect)

    @_literal.register
    def _literal_datetime(self, value: datetime.datetime, dialect: DialectCombo) -> Literal:
        return self.literal_datetime(value, dialect=dialect)

    @_literal.register
    def _literal_date(self, value: datetime.date, dialect: DialectCombo) -> Literal:
        return self.literal_date(value, dialect=dialect)

    @_literal.register
    def _literal_bool(self, value: bool, dialect: DialectCombo) -> Literal:
        return self.literal_bool(value, dialect=dialect)

    @_literal.register
    def _literal_str(self, value: str, dialect: DialectCombo) -> Literal:
        return self.literal_str(value, dialect=dialect)

    @_literal.register(list)
    @_literal.register(tuple)
    def _literal_array(self, value: Union[tuple, list], dialect: DialectCombo) -> Literal:
        return self.literal_array(value, dialect=dialect)

    def literal_int(self, value: int, dialect: DialectCombo) -> Literal:
        return sa.literal(value)

    def literal_float(self, value: float, dialect: DialectCombo) -> Literal:
        return sa.literal(value)

    def literal_datetime(self, value: datetime.datetime, dialect: DialectCombo) -> Literal:
        return sa.literal(value)

    def literal_date(self, value: datetime.date, dialect: DialectCombo) -> Literal:
        return sa.literal(value)

    def literal_bool(self, value: bool, dialect: DialectCombo) -> Literal:
        return sa.literal(value)

    def literal_str(self, value: str, dialect: DialectCombo) -> Literal:
        return sa.literal(value)

    def literal_array(self, value: Union[tuple, list], dialect: DialectCombo) -> Literal:
        return sa.literal(value)

    def literal(self, value: Any, dialect: DialectCombo) -> Literal:
        return self._literal(value, dialect=dialect)

lib/dl_formula/dl_formula/connectors/base/literal.py
datalens-tech/datalens-backend
from typing import Type

import sqlalchemy as sa
from sqlalchemy.types import TypeEngine

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectName


class SATypeConstructor:
    def get_sa_type(self, data_type: DataType) -> TypeEngine:
        raise NotImplementedError


class DefaultSATypeConstructor(SATypeConstructor):
    def get_sa_type(self, data_type: DataType) -> TypeEngine:
        type_map: dict[DataType, TypeEngine] = {
            **{
                str_type: sa.String(length=255)
                for str_type in (
                    DataType.STRING,
                    DataType.GEOPOINT,
                    DataType.GEOPOLYGON,
                )
            },
            DataType.INTEGER: sa.Integer(),
            DataType.FLOAT: sa.Float(),
            DataType.BOOLEAN: sa.Boolean(),
            DataType.DATE: sa.Date(),
            DataType.DATETIME: sa.DateTime(),
            DataType.GENERICDATETIME: sa.DateTime(),
            DataType.UUID: sa.String(length=255),
            DataType.ARRAY_INT: sa.ARRAY(sa.Integer),
            DataType.ARRAY_FLOAT: sa.ARRAY(sa.Float),
            DataType.ARRAY_STR: sa.ARRAY(sa.String),
        }
        return type_map[data_type]


_REGISTRY: dict[DialectName, Type[SATypeConstructor]] = {}


def get_type_constructor(dialect_name: DialectName) -> SATypeConstructor:
    type_constructor_cls = _REGISTRY[dialect_name]
    return type_constructor_cls()


def register_type_constructor(
    dialect_name: DialectName,
    type_constructor_cls: Type[SATypeConstructor],
) -> None:
    if dialect_name in _REGISTRY:
        assert _REGISTRY[dialect_name] is type_constructor_cls
    else:
        _REGISTRY[dialect_name] = type_constructor_cls

lib/dl_formula/dl_formula/connectors/base/type_constructor.py
datalens-tech/datalens-backend
from typing import (
    Iterable,
    Type,
)

from dl_formula.connectors.base.connector import FormulaConnector
from dl_formula.connectors.base.type_constructor import register_type_constructor
from dl_formula.core.dialect import (
    register_basic_dialects,
    register_default_dialect,
    register_dialect_namespace,
)
from dl_formula.definitions.base import NodeTranslation
from dl_formula.definitions.literals import register_literalizer
from dl_formula.definitions.registry import OPERATION_REGISTRY
from dl_formula.translation.columns import register_column_renderer_cls
from dl_formula.translation.context_processing import register_context_processor
from dl_formula.translation.sa_dialects import register_sa_dialect


_REGISTERED: set[Type[FormulaConnector]] = set()


class FormulaConnectorRegistrator:
    @classmethod
    def _register_op_definitions(cls, op_definitions: Iterable[NodeTranslation]) -> None:
        for def_item in op_definitions:
            OPERATION_REGISTRY.register(def_item)

    @classmethod
    def register_connector(cls, connector: Type[FormulaConnector]) -> None:
        if connector in _REGISTERED:
            return

        cls._register_op_definitions(connector.op_definitions)
        register_literalizer(connector.dialects, connector.literalizer_cls)
        register_context_processor(connector.dialects, connector.context_processor_cls)
        register_column_renderer_cls(connector.dialects, connector.column_renderer_cls)
        register_basic_dialects(connector.dialects)
        register_dialect_namespace(connector.dialect_ns_cls)
        for dialect_name in connector.get_dialect_names():
            if connector.default_dialect is not None:
                register_default_dialect(dialect_name, connector.default_dialect)
            register_type_constructor(
                dialect_name=dialect_name,
                type_constructor_cls=connector.type_constructor_cls,
            )

        for dialect in connector.dialects.to_list(with_self=True):
            register_sa_dialect(dialect=dialect, sa_dialect=connector.sa_dialect)

        connector.registration_hook()  # for some weird custom actions

        _REGISTERED.add(connector)


CONN_REG_FORMULA = FormulaConnectorRegistrator

lib/dl_formula/dl_formula/connectors/registration.py
datalens-tech/datalens-backend
"""
Non-"natural" nodes that can be used to facilitate translation, validation, slicing, etc.
"""

from __future__ import annotations

from typing import (
    Hashable,
    Optional,
    Sequence,
    cast,
)

import dl_formula.core.nodes as nodes


class ErrorNode(nodes.Null):
    """Node for setting up an error inside a formula to be raised by the translator"""

    __slots__ = ()
    show_names = nodes.FormulaItem.show_names + ("err_code", "message")

    @property
    def err_code(self) -> tuple[str, ...]:
        return cast(tuple[str, ...], self.internal_value[0])

    @property
    def message(self) -> str:
        return cast(str, self.internal_value[1])

    @classmethod
    def make(
        cls,
        *,
        err_code: tuple[str, ...],
        message: str,
        meta: Optional[nodes.NodeMeta] = None,
    ) -> ErrorNode:
        children = ()
        internal_value = (err_code, message)
        return cls(*children, internal_value=internal_value, meta=meta)

    @classmethod
    def validate_children(cls, children: Sequence[nodes.FormulaItem]) -> None:
        assert not children

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert len(internal_value) == 2
        assert isinstance(internal_value[0], tuple)
        assert isinstance(internal_value[1], str)

lib/dl_formula/dl_formula/core/aux_nodes.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import OrderedDict
from enum import (
    Enum,
    unique,
)
from typing import (
    FrozenSet,
    Optional,
)

import attr

from dl_formula.core import exc


@unique
class DataType(Enum):
    NULL = "null"
    INTEGER = "integer"
    CONST_INTEGER = "const_integer"
    FLOAT = "float"
    CONST_FLOAT = "const_float"
    STRING = "string"
    CONST_STRING = "const_string"
    DATE = "date"
    CONST_DATE = "const_date"
    DATETIME = "datetime"
    CONST_DATETIME = "const_datetime"
    DATETIMETZ = "datetimetz"  # timezone-aware datetime: UTC datetime + timezone name, normally.
    CONST_DATETIMETZ = "const_datetimetz"
    GENERICDATETIME = "genericdatetime"  # datetime with no special timezone logic
    CONST_GENERICDATETIME = "const_genericdatetime"
    BOOLEAN = "boolean"
    CONST_BOOLEAN = "const_boolean"
    GEOPOINT = "geopoint"
    CONST_GEOPOINT = "const_geopoint"
    GEOPOLYGON = "geopolygon"
    CONST_GEOPOLYGON = "const_geopolygon"
    # "MARKUP": a specially-structured string that contains a markup tree.
    MARKUP = "markup"
    CONST_MARKUP = "const_markup"
    UUID = "uuid"
    CONST_UUID = "const_uuid"
    ARRAY_FLOAT = "array_float"
    CONST_ARRAY_FLOAT = "const_array_float"
    ARRAY_INT = "array_int"
    CONST_ARRAY_INT = "const_array_int"
    ARRAY_STR = "array_str"
    CONST_ARRAY_STR = "const_array_str"
    TREE_STR = "tree_str"
    CONST_TREE_STR = "const_tree_str"
    # 'unsupported type', for all unlisted database types. Should be
    # *explicitly* stringifiable. Should not ever be a constant.
    # Somewhat similar to DataType.NULL, but makes it possible to extract some value.
    UNSUPPORTED = "unsupported"

    @classmethod
    def get_common_cast_type(cls, *types) -> DataType:
        set_of_types = set(types)
        if len(set_of_types) == 1:
            return set_of_types.pop()

        for to_type, from_types in _AUTOCAST_FROM_TYPES.items():
            if set_of_types.issubset(from_types):
                return to_type

        raise exc.TypeConflictError("Type conflict: [{}]".format(", ".join(sorted(t.name for t in types))))

    def casts_to(self, type_: DataType) -> bool:
        """Check whether ``self`` can be automatically cast to ``type_``"""
        return self in _AUTOCAST_FROM_TYPES[type_]

    @property
    def is_const(self) -> bool:
        """Flag indicating whether the type is constant"""
        return self.value.startswith("const_")

    @property
    def non_const_type(self) -> DataType:
        """Return a non-constant version of the type"""
        return self.__class__[self.name.replace("CONST_", "")]

    @property
    def const_type(self) -> DataType:
        """Return a constant version of the type"""
        if self.is_const or self is DataType.NULL or self is DataType.UNSUPPORTED:
            return self
        return DataType["CONST_" + self.name]

    @property
    def autocast_types(self) -> FrozenSet[DataType]:
        return frozenset(_AUTOCAST_FROM_TYPES.get(self, ()))


@attr.s(frozen=True)
class DataTypeParams:
    """Mix of all possible parameters for parametrized data types"""

    timezone: Optional[str] = attr.ib(default=None)
    # Other possible cases: decimal precision, datetime sub-second precision, nullable, enum values.


_AUTOCAST_FROM_TYPES = OrderedDict(
    (
        # <to_type>: set(<from_types>)
        # ordered from more specific (fewer from_types) to less specific (more from_types) type casts
        (DataType.NULL, {DataType.NULL}),
        (DataType.CONST_INTEGER, {DataType.CONST_INTEGER}),
        (DataType.CONST_DATE, {DataType.CONST_DATE}),
        (DataType.CONST_STRING, {DataType.CONST_STRING}),
        (DataType.CONST_BOOLEAN, {DataType.CONST_BOOLEAN}),
        (DataType.CONST_DATETIME, {DataType.CONST_DATETIME}),
        (DataType.CONST_DATETIMETZ, {DataType.CONST_DATETIMETZ}),
        (DataType.CONST_GENERICDATETIME, {DataType.CONST_GENERICDATETIME}),
        (DataType.CONST_FLOAT, {DataType.CONST_FLOAT, DataType.CONST_INTEGER}),
        (DataType.CONST_GEOPOINT, {DataType.CONST_GEOPOINT}),
        (DataType.CONST_GEOPOLYGON, {DataType.CONST_GEOPOLYGON}),
        (DataType.CONST_MARKUP, {DataType.CONST_MARKUP}),
        (DataType.CONST_UUID, {DataType.CONST_UUID}),
        (DataType.CONST_ARRAY_FLOAT, {DataType.CONST_ARRAY_FLOAT}),
        (DataType.CONST_ARRAY_INT, {DataType.CONST_ARRAY_INT}),
        (DataType.CONST_ARRAY_STR, {DataType.CONST_ARRAY_STR}),
        (DataType.CONST_TREE_STR, {DataType.CONST_TREE_STR}),
        (DataType.INTEGER, {DataType.INTEGER, DataType.CONST_INTEGER, DataType.NULL}),
        (DataType.DATE, {DataType.DATE, DataType.CONST_DATE, DataType.NULL}),
        (DataType.STRING, {DataType.STRING, DataType.CONST_STRING, DataType.NULL}),
        (DataType.GEOPOINT, {DataType.GEOPOINT, DataType.CONST_GEOPOINT, DataType.NULL}),
        (DataType.GEOPOLYGON, {DataType.GEOPOLYGON, DataType.CONST_GEOPOLYGON, DataType.NULL}),
        (DataType.MARKUP, {DataType.MARKUP, DataType.CONST_MARKUP, DataType.NULL}),
        (DataType.BOOLEAN, {DataType.BOOLEAN, DataType.CONST_BOOLEAN, DataType.NULL}),
        (DataType.DATETIME, {DataType.DATETIME, DataType.CONST_DATETIME, DataType.NULL}),
        (DataType.DATETIMETZ, {DataType.DATETIMETZ, DataType.CONST_DATETIMETZ, DataType.NULL}),
        (DataType.GENERICDATETIME, {DataType.GENERICDATETIME, DataType.CONST_GENERICDATETIME, DataType.NULL}),
        (
            DataType.FLOAT,
            {DataType.FLOAT, DataType.INTEGER, DataType.CONST_FLOAT, DataType.CONST_INTEGER, DataType.NULL},
        ),
        (DataType.UUID, {DataType.UUID, DataType.CONST_UUID, DataType.NULL}),
        (DataType.TREE_STR, {DataType.TREE_STR, DataType.CONST_TREE_STR, DataType.NULL}),
        (DataType.ARRAY_FLOAT, {DataType.ARRAY_FLOAT, DataType.CONST_ARRAY_FLOAT, DataType.NULL}),
        (DataType.ARRAY_INT, {DataType.ARRAY_INT, DataType.CONST_ARRAY_INT, DataType.NULL}),
        (
            DataType.ARRAY_STR,
            {
                DataType.ARRAY_STR,
                DataType.CONST_ARRAY_STR,
                DataType.TREE_STR,
                DataType.CONST_TREE_STR,
                DataType.NULL,
            },
        ),
        (DataType.UNSUPPORTED, {DataType.UNSUPPORTED}),
    )
)

lib/dl_formula/dl_formula/core/datatype.py
datalens-tech/datalens-backend
""" ... """

from __future__ import annotations

import inspect
from itertools import chain
import logging
import re
from typing import (
    Any,
    Generator,
    NamedTuple,
    Optional,
    Type,
)

from dynamic_enum import (
    AutoEnumValue,
    DynamicEnum,
)


LOGGER = logging.getLogger(__name__)


VersionType = tuple[int, ...]

_NO_VERSION: VersionType = ()


class DialectName(DynamicEnum):
    UNKNOWN = AutoEnumValue()
    DUMMY = AutoEnumValue()
    ANY = AutoEnumValue()
    SQLITE = AutoEnumValue()  # TODO: Remove


class DialectBit(NamedTuple):
    name: DialectName
    version: VersionType = ()

    @property
    def orderable(self) -> tuple[str, VersionType]:
        return self.name.value, self.version


class DialectCombo(NamedTuple):
    bits: frozenset[DialectBit]

    @property
    def ambiguity(self) -> int:
        return len(self.bits)

    @property
    def deterministic(self) -> bool:
        return self.ambiguity == 1

    @property
    def single_bit(self) -> DialectBit:
        assert self.deterministic
        return next(iter(self.bits))

    @property
    def common_name(self) -> DialectName:
        names = {bit.name for bit in self.bits}
        assert len(names) == 1, str(names)
        return next(iter(names))

    @property
    def common_version(self) -> VersionType:
        versions = {bit.version for bit in self.bits}
        assert len(versions) == 1, str(versions)
        return next(iter(versions))

    @property
    def common_name_and_version(self) -> str:
        return f'{self.common_name.name}_{"_".join((str(part) for part in self.common_version))}'.strip("_")

    def __sub__(self, other: Any) -> DialectCombo:
        if isinstance(other, DialectBit):
            return DialectCombo(
                bits=self.bits - frozenset((other,)),
            )
        if isinstance(other, DialectCombo):
            return DialectCombo(
                bits=self.bits - other.bits,
            )
        else:
            raise TypeError(type(other))

    def __or__(self, other: Any) -> DialectCombo:
        if isinstance(other, DialectBit):
            return DialectCombo(
                bits=self.bits | frozenset((other,)),
            )
        if isinstance(other, DialectCombo):
            return DialectCombo(
                bits=self.bits | other.bits,
            )
        else:
            raise TypeError(type(other))

    def __and__(self, other: Any) -> DialectCombo:
        if isinstance(other, DialectBit):
            return DialectCombo(
                bits=self.bits & frozenset((other,)),
            )
        if isinstance(other, DialectCombo):
            return DialectCombo(
                bits=self.bits & other.bits,
            )
        else:
            raise TypeError(type(other))

    def __bool__(self) -> bool:
        return bool(self.bits)

    def to_list(self, with_self: bool = False) -> list[DialectCombo]:
        result = [DialectCombo(bits=frozenset((bit,))) for bit in self.bits]
        result.sort(key=lambda dialect: (dialect.single_bit.name.value, dialect.single_bit.version))
        if with_self and not self.deterministic:
            result.append(self)
        return result


def simple_combo(name: DialectName, version: VersionType = ()) -> DialectCombo:
    return DialectCombo(bits=frozenset((DialectBit(name=name, version=version),)))


class DialectNamespace:
    @classmethod
    def by_name(cls, item: str) -> DialectCombo:
        assert isinstance(item, str)
        return getattr(cls, item)

    @classmethod
    def iter_named_combos(cls) -> Generator[tuple[str, DialectCombo], None, None]:
        for name, item in inspect.getmembers(cls):
            if isinstance(item, DialectCombo):
                yield name, item

    @classmethod
    def all_iter(cls) -> Generator[DialectCombo, None, None]:
        for _, item in inspect.getmembers(cls):
            if isinstance(item, DialectCombo):
                yield item

    @classmethod
    def all_basic_iter(cls) -> Generator[DialectCombo, None, None]:
        for item in cls.all_iter():
            if item.deterministic:
                yield item

    @classmethod
    def all(cls) -> list[DialectCombo]:
        return list(cls.all_iter())

    @classmethod
    def and_above(cls, dialect: DialectCombo) -> DialectCombo:
        if not dialect.deterministic:
            raise ValueError("Only deterministic dialects are allowed")
        bit = dialect.single_bit
        return DialectCombo(
            bits=frozenset(
                item.single_bit
                for item in cls.all_basic_iter()
                if (item.single_bit.name == bit.name and item.single_bit.version >= bit.version)
            )
        )

    @classmethod
    def and_below(cls, dialect: DialectCombo) -> DialectCombo:
        if not dialect.deterministic:
            raise ValueError("Only deterministic dialects are allowed")
        bit = dialect.single_bit
        return DialectCombo(
            bits=frozenset(
                item.single_bit
                for item in cls.all_basic_iter()
                if (item.single_bit.name == bit.name and () < item.single_bit.version <= bit.version)
            )
        )


class StandardDialect(DialectNamespace):
    # Generic
    EMPTY = DialectCombo(bits=frozenset())
    DUMMY = simple_combo(name=DialectName.DUMMY)
    ANY = simple_combo(name=DialectName.ANY)

    # Others
    SQLITE = simple_combo(name=DialectName.SQLITE)


dialect_defaults: dict[DialectName, DialectCombo] = {
    DialectName.UNKNOWN: StandardDialect.EMPTY,
    DialectName.DUMMY: StandardDialect.DUMMY,
    DialectName.SQLITE: StandardDialect.SQLITE,
}


def register_default_dialect(dialect_name: DialectName, dialect_combo: DialectCombo):
    if dialect_name in dialect_defaults:
        assert dialect_defaults[dialect_name] is dialect_combo
    else:
        dialect_defaults[dialect_name] = dialect_combo


def from_name_and_version(
    dialect_name: DialectName,
    dialect_version: Optional[str],
) -> DialectCombo:
    """Return ``DialectCombo`` instance corresponding to the version string (``SELECT VERSION()``)"""

    parsed_version: Optional[VersionType]
    if dialect_version:
        version_match = re.search(r"(?P<version>\d+(\.\d+)+)", dialect_version)
        if version_match is None:
            parsed_version = None
        else:
            raw_version = version_match.group("version")
            parsed_version = tuple(int(part) for part in raw_version.split("."))
    else:
        parsed_version = None

    LOGGER.info(f"Parsed dialect version {dialect_version} as {parsed_version}")

    if parsed_version is None and dialect_name in dialect_defaults:
        LOGGER.info(f"Falling back to {dialect_defaults[dialect_name]}")
        return dialect_defaults[dialect_name]

    matches_by_name = sorted(
        chain.from_iterable(
            [
                item.to_list()
                for item in _ALL_BASIC_DIALECTS
                if item.deterministic and item.single_bit.name == dialect_name
            ],
        ),
        key=lambda item: item.single_bit.version,
    )
    matches_by_version = [
        item for item in matches_by_name if parsed_version is None or item.single_bit.version <= parsed_version
    ]
    if matches_by_version:
        # perfect match. Return the latest dialect version possible
        matched_dialect = matches_by_version[-1]
        LOGGER.info(f"Matched dialect found: {matched_dialect.single_bit}")
        return matched_dialect

    LOGGER.warning('Failed to get dialect for DB "%s" and version "%s"', dialect_name.name, dialect_version)
    if matches_by_name:
        # no matching dialect found. Fall back to the earliest supported for this DB type
        LOGGER.warning("Falling back to `%s`", matches_by_name[0].single_bit.name.name)
        return matches_by_name[0]

    # nothing. Just nothing
    raise ValueError("Invalid dialect parameters")


_ALL_BASIC_DIALECTS: set[DialectCombo] = {
    StandardDialect.DUMMY,
}


def register_basic_dialects(dialect: DialectCombo) -> None:
    for basic_dialect in dialect.to_list():
        _ALL_BASIC_DIALECTS.add(basic_dialect)


def get_all_basic_dialects() -> list[DialectCombo]:
    return sorted(_ALL_BASIC_DIALECTS, key=lambda x: x.single_bit.orderable)


_NAMED_DIALECT_COMBOS: dict[str, DialectCombo] = {}


def register_dialect_namespace(dialect_ns: Type[DialectNamespace]) -> None:
    for name, combo in dialect_ns.iter_named_combos():
        try:
            assert _NAMED_DIALECT_COMBOS[name] == combo
        except KeyError:
            _NAMED_DIALECT_COMBOS[name] = combo


def get_dialect_combos() -> list[DialectCombo]:
    return list(_NAMED_DIALECT_COMBOS.values())


def get_dialect_from_str(dialect_str: str) -> DialectCombo:
    return _NAMED_DIALECT_COMBOS[dialect_str]


register_dialect_namespace(StandardDialect)

lib/dl_formula/dl_formula/core/dialect.py
datalens-tech/datalens-backend
""" ... """

from __future__ import annotations

from typing import (
    ClassVar,
    List,
    Optional,
    Sequence,
    Tuple,
    Union,
)

from dl_formula.core.message_ctx import (
    FormulaErrorCtx,
    MessageLevel,
)
from dl_formula.core.position import Position


class ParserNotFoundError(Exception):
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.description = (
            "Cannot import parser files,"
            " check if you've generated them with"
            " `${PROJECT_ROOT}/docker_build/run-project-bake gen_antlr`"
        )

    def __str__(self):
        return self.description


class CacheError(Exception):
    pass


class InspectionError(Exception):
    pass


def _ifnone(first, second):
    return first if first is not None else second


class FormulaError(Exception):
    default_code: ClassVar[Tuple[str, ...]] = ("FORMULA",)

    def __init__(
        self,
        *errors: Union[str, FormulaErrorCtx],
        position: Optional[Position] = None,
        token: Optional[str] = None,
        code: Optional[Sequence[str]] = None,
    ):
        super().__init__()

        self.errors: List[FormulaErrorCtx] = []
        for error in errors:
            if not isinstance(error, FormulaErrorCtx):
                error = FormulaErrorCtx(
                    message=error,
                    level=MessageLevel.ERROR,
                    position=position or Position(),
                    token=token,
                    code=tuple(code or self.default_code or ()),
                )
            elif position is not None or token is not None or code is not None:
                error = FormulaErrorCtx(
                    message=error.message,
                    level=MessageLevel.ERROR,
                    position=_ifnone(position, error.position),
                    token=_ifnone(token, error.token),
                    code=tuple(_ifnone(code, error.code) or self.default_code or ()),
                )

            self.errors.append(error)

        self.error = self.errors[0] if len(errors) == 1 else None
        self.description = "; ".join([str(err) for err in self.errors])

    def __str__(self):
        return self.description


# Parse Errors


class ParseError(FormulaError):
    default_code = FormulaError.default_code + ("PARSE",)


class ParseUnexpectedEOFError(ParseError):
    default_code = ParseError.default_code + ("UNEXPECTED_EOF",)


class ParseEmptyFormulaError(ParseUnexpectedEOFError):
    default_code = ParseUnexpectedEOFError.default_code + ("EMPTY_FORMULA",)


class ParseUnexpectedTokenError(ParseError):
    default_code = ParseError.default_code + ("UNEXPECTED_TOKEN",)


class ParseValueError(ParseError):
    default_code = ParseError.default_code + ("VALUE",)


class ParseDateValueError(ParseValueError):
    default_code = ParseValueError.default_code + ("DATE",)


class ParseDatetimeValueError(ParseValueError):
    default_code = ParseValueError.default_code + ("DATETIME",)


class ParseClauseError(ParseError):
    default_code = FormulaError.default_code + ("CLAUSE",)


# Validation Errors


class ValidationError(FormulaError):
    default_code = FormulaError.default_code + ("VALIDATION",)


# Aggregation errors


class AggregationError(ValidationError):
    default_code = ValidationError.default_code + ("AGG",)


class InconsistentAggregationError(AggregationError):
    default_code = AggregationError.default_code + ("INCONSISTENT",)


class DoubleAggregationError(AggregationError):
    default_code = AggregationError.default_code + ("DOUBLE",)


# Window function errors


class WindowFunctionError(ValidationError):
    default_code = ValidationError.default_code + ("WIN_FUNC",)


class NestedWindowFunctionError(WindowFunctionError):
    default_code = WindowFunctionError.default_code + ("NESTED",)


class WindowFunctionWOAggregationError(WindowFunctionError):
    default_code = WindowFunctionError.default_code + ("NO_AGG",)


class WindowFunctionUnselectedDimensionError(WindowFunctionError):
    default_code = WindowFunctionError.default_code + ("BFB_UNSELECTED_DIMENSION",)


# Lookup function errors


class LookupFunctionError(ValidationError):
    default_code = ValidationError.default_code + ("LOOKUP_FUNC",)


class LookupFunctionArgNumberError(LookupFunctionError):
    default_code = LookupFunctionError.default_code + ("ARG_NUM",)


class LookupFunctionUnselectedDimensionError(LookupFunctionError):
    default_code = LookupFunctionError.default_code + ("UNSELECTED_DIMENSION",)


class LookupFunctionIgnoredLookupDimensionError(LookupFunctionError):
    default_code = LookupFunctionError.default_code + ("IGNORED_LOOKUP_DIMENSION",)


class LookupFunctionWOAggregationError(LookupFunctionError):
    default_code = LookupFunctionError.default_code + ("NO_AGG",)


class LookupFunctionConstantLookupDimensionError(LookupFunctionError):
    default_code = LookupFunctionError.default_code + ("CONSTANT_LOOKUP_DIMENSION",)


class LookupFunctionAggregatedDimensionError(LookupFunctionError):
    default_code = LookupFunctionError.default_code + ("AGGREGATED_DIMENSION",)


# LOD errors


class LodError(ValidationError):
    default_code = ValidationError.default_code + ("LOD",)


class LodIncompatibleDimensionsError(LodError):
    default_code = LodError.default_code + ("INCOMPATIBLE_DIMENSIONS",)


class LodInvalidTopLevelDimensionsError(LodError):
    default_code = LodError.default_code + ("INVALID_TOPLEVEL_DIMENSIONS",)


class LodMeasureDimensionsError(LodError):
    default_code = LodError.default_code + ("MEASURE_AS_DIMENSION",)


class UnknownSourceColumnError(FormulaError):
    default_code = FormulaError.default_code + ("UNKNOWN_SOURCE_COLUMN",)


class UnknownFieldInFormulaError(FormulaError):
    default_code = FormulaError.default_code + ("UNKNOWN_FIELD_IN_FORMULA",)


# Translation Errors


class TranslationError(FormulaError):
    default_code = FormulaError.default_code + ("TRANSLATION",)


class DataTypeError(TranslationError):
    default_code = TranslationError.default_code + ("DATA_TYPE",)


class TypeConflictError(DataTypeError):
    default_code = DataTypeError.default_code + ("CONFLICT",)


class TranslationUnknownFunctionError(TranslationError):
    default_code = TranslationError.default_code + ("UNKNOWN_FUNCTION",)


class TranslationUnknownFieldError(TranslationError):
    # An internal error, should not appear to users.
    # If it does, investigate.
    default_code = TranslationError.default_code + ("UNKNOWN_FIELD",)


# Other Errors


class UnknownWindowDimensionError(FormulaError):
    default_code = FormulaError.default_code + ("UNKNOWN_WINDOW_DIMENSION",)

lib/dl_formula/dl_formula/core/exc.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Hashable,
    NamedTuple,
    Optional,
    Tuple,
)


class NodeExtract(NamedTuple):
    """
    A hashable syntax-insensitive representation of the node
    that contains all of its essential attributes.
    For formula inspection and optimization purposes.
    """

    type_name: str
    value: Optional[Hashable] = None  # Python primitive for internal non-FormulaItem attributes
    children: Tuple["NodeExtract", ...] = ()  # type: ignore
    complexity: int = 1  # A simple and fast way to know the complexity of the expression

    def __repr__(self) -> str:
        return tuple.__repr__(self)

    def __str__(self) -> str:
        return self.__repr__()

lib/dl_formula/dl_formula/core/extract.py
datalens-tech/datalens-backend
"""
Auxiliary nodes specific to the forked subquery joining mechanism.
"""

from __future__ import annotations

from enum import Enum
from typing import (
    Hashable,
    Optional,
    Sequence,
    cast,
)

import dl_formula.core.nodes as nodes


class JoinType(Enum):
    inner = "inner"
    left = "left"


class JoinConditionBase(nodes.FormulaItem):
    """Base class for join conditions"""

    __slots__ = ()
    autonomous = False

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert not internal_value


class SelfEqualityJoinCondition(JoinConditionBase):
    """
    A join condition that indicates that the values
    of the same dimension equal each other in main subquery and its fork.
    """

    __slots__ = ()
    show_names = JoinConditionBase.show_names + ("expr",)

    expr: nodes.Child[nodes.FormulaItem] = nodes.Child(0)

    @classmethod
    def validate_children(cls, children: Sequence[nodes.FormulaItem]) -> None:
        assert len(children) == 1

    @classmethod
    def make(cls, expr: nodes.FormulaItem, *, meta: Optional[nodes.NodeMeta] = None) -> SelfEqualityJoinCondition:
        children = (expr,)
        return cls(*children, meta=meta)


class BinaryJoinCondition(JoinConditionBase):
    """
    A join condition that indicates that
    ``expr`` from the main subquery equals ``fork_expr`` from the fork subquery.
    """

    __slots__ = ()
    show_names = JoinConditionBase.show_names + ("expr", "fork_expr")

    expr: nodes.Child[nodes.FormulaItem] = nodes.Child(0)
    fork_expr: nodes.Child[nodes.FormulaItem] = nodes.Child(1)

    @classmethod
    def validate_children(cls, children: Sequence[nodes.FormulaItem]) -> None:
        assert len(children) == 2

    @classmethod
    def make(
        cls,
        expr: nodes.FormulaItem,
        fork_expr: nodes.FormulaItem,
        *,
        meta: Optional[nodes.NodeMeta] = None,
    ) -> BinaryJoinCondition:
        children = (expr, fork_expr)
        return cls(*children, meta=meta)


class QueryForkJoiningBase(nodes.FormulaItem):
    __slots__ = ()
    autonomous = False

    @property
    def is_self_eq_join(self) -> bool:
        raise NotImplementedError


class QueryForkJoiningWithList(QueryForkJoiningBase):
    """
    Represents a join using all dimensions with custom condition overrides
    """

    __slots__ = ()

    show_names = QueryForkJoiningBase.show_names + ("condition_list",)

    condition_list: nodes.MultiChild[JoinConditionBase] = nodes.MultiChild(slice(0, None))

    @classmethod
    def make(
        cls,
        condition_list: list[JoinConditionBase],
        *,
        meta: Optional[nodes.NodeMeta] = None,
    ) -> QueryForkJoiningWithList:
        children = condition_list
        return cls(*children, meta=meta)

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert not internal_value

    @property
    def is_self_eq_join(self) -> bool:
        """All join conditions are ``SelfEqualityJoinCondition``"""
        return all(isinstance(child, SelfEqualityJoinCondition) for child in self.children)


class QueryFork(nodes.FormulaItem):
    """
    Represents a point where the query should be forked in two:
    the main subquery and the forked subquery.
    The ``joining`` child node describes how the two subqueries should be joined
    """

    __slots__ = ()

    show_names = nodes.FormulaItem.show_names + ("join_type", "joining", "result_expr", "lod")

    joining: nodes.Child[QueryForkJoiningBase] = nodes.Child(0)
    result_expr: nodes.Child[nodes.FormulaItem] = nodes.Child(1)
    before_filter_by: nodes.Child[nodes.BeforeFilterBy] = nodes.Child(2)
    lod: nodes.Child[nodes.LodSpecifier] = nodes.Child(3)

    @classmethod
    def make(
        cls,
        join_type: JoinType,
        joining: QueryForkJoiningBase,
        result_expr: nodes.FormulaItem,
        before_filter_by: Optional[nodes.BeforeFilterBy] = None,
        lod: Optional[nodes.LodSpecifier] = None,
        meta: Optional[nodes.NodeMeta] = None,
    ) -> QueryFork:
        if before_filter_by is None:
            before_filter_by = nodes.BeforeFilterBy.make()
        if lod is None:
            lod = nodes.InheritedLodSpecifier()

        children = (joining, result_expr, before_filter_by, lod)
        internal_value = (join_type,)
        return cls(*children, internal_value=internal_value, meta=meta)

    @classmethod
    def validate_children(cls, children: Sequence[nodes.FormulaItem]) -> None:
        assert len(children) == 4

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert len(internal_value) == 1
        assert isinstance(internal_value[0], JoinType)

    @property
    def join_type(self) -> JoinType:
        return cast(JoinType, self.internal_value[0])

lib/dl_formula/dl_formula/core/fork_nodes.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    NamedTuple,
    Optional,
    Tuple,
)


class NodeHierarchyIndex(NamedTuple):
    indices: Tuple[int, ...] = ()

    @staticmethod
    def make(*args: int) -> NodeHierarchyIndex:
        return NodeHierarchyIndex(indices=args)

    def __bool__(self) -> bool:
        return bool(self.indices)

    def __add__(self, other: Any) -> NodeHierarchyIndex:
        if isinstance(other, int):
            return NodeHierarchyIndex(indices=self.indices + (other,))
        if isinstance(other, NodeHierarchyIndex):
            return NodeHierarchyIndex(indices=self.indices + other.indices)
        if isinstance(other, (tuple, list)):
            return NodeHierarchyIndex(indices=self.indices + tuple(other))
        raise TypeError(type(other))

    def lsplit(self) -> Tuple[Optional[int], NodeHierarchyIndex]:
        if not self:
            return None, NodeHierarchyIndex()
        return self.indices[0], NodeHierarchyIndex(indices=self.indices[1:])

    def rsplit(self) -> Tuple[NodeHierarchyIndex, Optional[int]]:
        if not self:
            return NodeHierarchyIndex(), None
        return NodeHierarchyIndex(indices=self.indices[:-1]), self.indices[-1]

    def startswith(self, other: NodeHierarchyIndex) -> bool:
        return self.indices[: len(other.indices)] == other.indices

lib/dl_formula/dl_formula/core/index.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import (
    Enum,
    unique,
)
from typing import (
    NamedTuple,
    Optional,
    Sequence,
    Tuple,
)

from dl_formula.core.position import Position


@unique
class MessageLevel(Enum):
    ERROR = "error"
    WARNING = "warning"


class FormulaErrorCtx(NamedTuple):
    message: str
    level: MessageLevel
    position: Position = Position()
    token: Optional[str] = None
    code: Tuple[str, ...] = ()

    def is_error(self) -> bool:
        return self.level is MessageLevel.ERROR

    def is_warning(self) -> bool:
        return self.level is MessageLevel.WARNING

    def __str__(self) -> str:
        return self.message

    @property
    def coords(self) -> Tuple[Optional[int], Optional[int]]:
        """Convert position to 2-dimensional coords"""
        return self.position.start_row, self.position.start_col

    def is_sub_error(self, code: Sequence[str]) -> bool:
        if len(self.code) < len(code):
            return False
        return tuple(code) == self.code[: len(code)]

lib/dl_formula/dl_formula/core/message_ctx.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import copy
import datetime
from typing import (
    Any,
    Callable,
    ClassVar,
    Collection,
    Generator,
    Generic,
    Hashable,
    Iterable,
    Mapping,
    Optional,
    Sequence,
    Type,
    TypeVar,
    Union,
    cast,
)
import uuid

from dl_formula.core.extract import NodeExtract
from dl_formula.core.index import NodeHierarchyIndex
from dl_formula.core.position import Position
from dl_formula.core.tag import LevelTag


class NodeMeta:
    __slots__ = ("position", "original_text", "level_tag")

    def __init__(
        self,
        position: Optional[Position] = None,
        original_text: Optional[str] = None,
        level_tag: Optional[LevelTag] = None,
    ):
        self.position = position or Position()
        self.original_text = original_text
        self.level_tag = level_tag

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(position={self.position!r})"

    def __str__(self) -> str:
        return repr(self)

    def with_tag(self, level_tag: LevelTag) -> NodeMeta:
        """
        Return a copy of meta with replaced ``self.level_tag``
        """
        return self.__class__(
            position=self.position,
            original_text=self.original_text,
            level_tag=level_tag,
        )


_FORMULA_ITEM_TV = TypeVar("_FORMULA_ITEM_TV", bound="FormulaItem")


class FormulaItem(abc.ABC):
    """Abstract class. Base class for all formula building blocks."""

    __slots__ = ("__children", "__internal_value", "__meta", "__extract")
    show_names: tuple[str, ...] = ()
    autonomous: ClassVar[bool] = True
    # If True, then the node can represent a standalone expression.
    # However, some nodes are can only exist as an integral part of a parent node (sub-clauses).

    @property
    def children(self) -> tuple[FormulaItem, ...]:
        return tuple(self.__children)

    @property
    def internal_value(self) -> tuple[Optional[Hashable], ...]:
        return self.__internal_value

    @property
    def position(self) -> Position:
        return self.__meta.position

    @property
    def original_text(self) -> Optional[str]:
        return self.__meta.original_text

    @property
    def level_tag(self) -> Optional[LevelTag]:
        return self.__meta.level_tag

    @property
    def pos_range(self) -> tuple[Optional[int], Optional[int]]:
        return self.__meta.position.start, self.__meta.position.end

    @property
    def meta(self) -> NodeMeta:
        return self.__meta

    def __init__(
        self,
        *children: FormulaItem,
        internal_value: tuple[Optional[Hashable], ...] = (),
        meta: Optional[NodeMeta] = None,
    ):
        for child in children:
            if not isinstance(child, FormulaItem):
                raise TypeError(f"Invalid type {type(child)} for child node")

        self.validate_children(children)
        assert isinstance(internal_value, tuple)
        self.validate_internal_value(internal_value)

        self.__children: list[FormulaItem] = [child for child in children]
        self.__internal_value = internal_value
        self.__meta: NodeMeta = meta or NodeMeta()
        self.__extract = self._make_extract()

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        pass

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert internal_value == ()

    def get_scalar_value(self) -> Optional[Hashable]:
        return self.__internal_value

    def stringify(self, with_meta: bool = False) -> str:
        show_names = self.show_names if not with_meta else (self.show_names + ("meta",))
        kv_pairs = []
        for name in show_names:
            if name == "meta":
                attr_value = self.meta
            else:
                attr_value = getattr(self, name)
            if isinstance(attr_value, FormulaItem):
                attr_value_s = attr_value.stringify(with_meta=with_meta)
            else:
                attr_value_s = repr(attr_value)
            pair = f"{name}={attr_value_s}"
            kv_pairs.append(pair)
        param_str = ", ".join(kv_pairs)
        return "{}({})".format(self.__class__.__name__, param_str)

    def __repr__(self) -> str:
        return self.stringify()

    def __str__(self) -> str:
        return repr(self)

    def __eq__(self, other: Any) -> bool:
        return other.__class__ is self.__class__ and self.__extract == other.__extract

    def pretty(
        self,
        indent: str = "  ",
        initial_indent: str = "",
        inline_len: int = 70,
        do_leading_indent: bool = False,
        with_meta: bool = False,
    ) -> str:
        """Return a "pretty" str version of the Node tree: multiple lines with indentations"""
        str_version = self.stringify(with_meta=with_meta)
        if len(str_version) <= inline_len:
            return str_version

        def prettify_child(_child: Any, _initial_indent: str) -> str:
            if isinstance(_child, FormulaItem):
                return _child.pretty(
                    initial_indent=_initial_indent,
                    indent=indent,
                    inline_len=inline_len,
                    do_leading_indent=False,
                    with_meta=with_meta,
                )
            if isinstance(_child, tuple):
                if not _child:
                    return "()"
                res = "(\n"
                item_indent = _initial_indent + indent
                for item in _child:
                    res += "{}{},\n".format(item_indent, prettify_child(item, item_indent))
                res += _initial_indent + ")"
                return res
            else:
                return repr(_child)

        lines = []
        lines.append("{}{}(".format(initial_indent if do_leading_indent else "", self.__class__.__name__))
        show_names = self.show_names if not with_meta else (self.show_names + ("meta",))
        for name in show_names:
            if name == "meta":
                child = str(self.__meta)
            else:
                child = prettify_child(getattr(self, name), initial_indent + indent)
            lines.append("{}{}={},".format(initial_indent + indent, name, child))
        lines.append("{})".format(initial_indent))

        return "\n".join(lines)

    def visit_node_type(self, node_type: Type[FormulaItem], visit_func: Callable[[FormulaItem], Any]) -> None:
        """Walk the whole ``FormulaItem`` (sub)tree and call ``visit_func`` for all nodes of given type"""
        if isinstance(self, node_type):
            visit_func(self)

        for child in self.__children:
            child.visit_node_type(node_type=node_type, visit_func=visit_func)

    def list_node_type(self, node_type: Type[_FORMULA_ITEM_TV]) -> list[_FORMULA_ITEM_TV]:
        res: list[_FORMULA_ITEM_TV] = []
        self.visit_node_type(node_type=node_type, visit_func=res.append)  # type: ignore
        return res

    def get_by_pos(
        self,
        pos: int,
        node_types: Optional[tuple[Type[FormulaItem], ...]] = None,
    ) -> Optional[FormulaItem]:
        """
        Return the innermost node at position ``pos``.
        If ``node_types`` is given, then choose only among nodes of the listed types.
        """
        if self.position.start is None or self.position.end is None:
            return None

        if self.position.start <= pos <= self.position.end:  # include both ends
            # self fits the position
            # try to narrow down the result  by inspecting children
            matching_children = []
            for child in self.__children:
                res = child.get_by_pos(pos, node_types=node_types)
                if res is not None:
                    matching_children.append(res)

            if matching_children:
                # always prefer the rightmost of the matches
                return matching_children[-1]

            if node_types is None or isinstance(self, node_types):
                return self

        return None

    @staticmethod
    def __inplace_replace_at_index_no_copy(
        top_node: _FORMULA_ITEM_TV, index: NodeHierarchyIndex, expr: FormulaItem
    ) -> None:
        parent_index, child_index = index.rsplit()
        assert child_index is not None
        parent = top_node.get(parent_index)

        parent.__children[child_index] = expr  # type: ignore
        for node in reversed(list(top_node.iter_index(parent_index))):
            node.__extract = node._make_extract()

    def substitute_batch(
        self: _FORMULA_ITEM_TV, to_substitute: Mapping[NodeHierarchyIndex, FormulaItem]
    ) -> _FORMULA_ITEM_TV:
        """
        Substitutes sub nodes by given indexes
        Note that substituent nodes are not copied before "insertion"
        """
        self_copy: _FORMULA_ITEM_TV = copy.copy(self)
        for index, expr in to_substitute.items():
            self.__inplace_replace_at_index_no_copy(self_copy, index=index, expr=expr)

        self_copy.__extract = self_copy._make_extract()
        return self_copy

    def replace_at_index(self: _FORMULA_ITEM_TV, index: NodeHierarchyIndex, expr: FormulaItem) -> _FORMULA_ITEM_TV:
        """
        Replace node at index ``index`` with a new one specified by ``expr``.
        Creating new nodes while descending along the `index` in recursion.
        """
        child_in_index, index_tail = index.lsplit()

        # final replacement
        if child_in_index is None:
            return expr  # type: ignore

        child = self.__children[child_in_index]
        new_child = child.replace_at_index(index_tail, expr)

        children = self.__children[:]
        children[child_in_index] = new_child  # type: ignore

        return self.light_copy(children)

    def light_copy(
        self: _FORMULA_ITEM_TV, children: Sequence[FormulaItem], *, meta: Optional[NodeMeta] = None
    ) -> _FORMULA_ITEM_TV:
        """
        Creates a copy of self with new children nodes
        """

        return self.__class__(
            *children,  # type: ignore
            internal_value=self.__internal_value,
            meta=meta or self.__meta,
        )

    def replace_nodes(
        self: _FORMULA_ITEM_TV,
        match_func: Callable[[FormulaItem, tuple[FormulaItem, ...]], bool],
        replace_func: Callable[[FormulaItem, tuple[FormulaItem, ...]], FormulaItem],
        parent_stack: tuple[FormulaItem, ...] = (),
    ) -> _FORMULA_ITEM_TV:
        """
        Walk the whole ``FormulaItem`` (sub)tree and replace child nodes
        for which ``replace_func(<child_node>)`` is ``True`` with``replace_func(<child_node>)``.
        Return a copy of the original node with replacements if any occurred, otherwise returns self.
        """
        if not self.__children:
            return self

        is_modified = False

        parent_stack_w_self = parent_stack + (self,)
        to_replace: dict[int, FormulaItem] = {}

        for idx, child in enumerate(self.__children):
            if match_func(child, parent_stack_w_self):
                modified_child = replace_func(child, parent_stack_w_self)
                if modified_child is not child or modified_child != child:
                    child = to_replace[idx] = modified_child  # type: ignore
                    is_modified = True

            modified_child = child.replace_nodes(match_func, replace_func, parent_stack_w_self)
            if modified_child is not child or modified_child != child:
                to_replace[idx] = modified_child  # type: ignore
                is_modified = True

        if is_modified:
            children = cast(
                list[FormulaItem],
                [to_replace[idx] if idx in to_replace else child for idx, child in enumerate(self.__children)],
            )
            return self.light_copy(children)

        return self

    def __copy__(self: _FORMULA_ITEM_TV) -> _FORMULA_ITEM_TV:
        """Create a copy of self"""

        return self.__class__(
            *[copy.copy(c) for c in self.__children],
            internal_value=self.__internal_value,
            meta=self.__meta,
        )

    def __inplace_replace_meta(self, meta: NodeMeta) -> None:
        self.__meta = meta
        self.__extract = self._make_extract()

    def with_tag(self: _FORMULA_ITEM_TV, level_tag: LevelTag) -> _FORMULA_ITEM_TV:
        """
        Return a copy of node with replaced ``self.meta.level_tag``
        """
        return self.light_copy(self.children[:], meta=self.meta.with_tag(level_tag=level_tag))

    def get(self, index: NodeHierarchyIndex) -> FormulaItem:
        """Find and return sub-node identified by ``index``"""
        head, tail = index.lsplit()
        if head is None:
            return self
        return self.__children[head].get(tail)

    def __getitem__(self, item: NodeHierarchyIndex) -> FormulaItem:
        return self.get(item)

    def enumerate(
        self,
        prefix: NodeHierarchyIndex = NodeHierarchyIndex(),
        max_depth: Optional[int] = None,
    ) -> Generator[tuple[NodeHierarchyIndex, FormulaItem], None, None]:
        """
        Similar to the standard ``enumerate`` function,
        but iterates over the node hierarchy
        and uses ``NodeHierarchyIndex`` instead of an integer index.

        Recursion can be limited by ``max_depth``.
        """
        yield prefix, self
        if max_depth == 0:
            return
        child_max_depth = max_depth - 1 if max_depth is not None else None
        for child_i, child in enumerate(self.children):  # type: ignore
            yield from child.enumerate(prefix=prefix + child_i, max_depth=child_max_depth)

    def resolve_index(self, node: FormulaItem, pointer_eq: bool = True) -> Optional[NodeHierarchyIndex]:
        """Get index for given ``node``."""
        for index, child in self.enumerate():
            if pointer_eq:
                if child is node:
                    return index
            elif child == node:  # (and not pointer_eq)
                return index

        raise ValueError("Cannot resolve index for node")

    def iter_index(self, index: NodeHierarchyIndex, exclude_last: bool = False) -> Generator[FormulaItem, None, None]:
        """Iterate over all ancestor nodes (and, optionally, the target child node)."""
        node = self
        yield self
        max_i = len(index.indices) - 1
        for hierarchy_level, flat_ind in enumerate(index.indices):
            node = node.__children[flat_ind]
            if not exclude_last or hierarchy_level < max_i:
                yield node

    @property
    def start_pos(self) -> Optional[int]:
        return self.position.start

    @property
    def end_pos(self) -> Optional[int]:
        return self.position.end

    @property
    def complexity(self) -> Optional[int]:
        return self.__extract.complexity if self.__extract is not None else None

    def _make_extract(self) -> Optional[NodeExtract]:
        """
        Return a hashable immutable object that represents this node and its children.
        It should not contain any syntactic info so that can be used for optimization of data calculations.
        """
        optional_child_extracts = tuple(child.extract for child in self.children)
        if None in optional_child_extracts:
            return None
        child_extracts = cast(tuple[NodeExtract, ...], optional_child_extracts)
        return NodeExtract(
            type_name=type(self).__name__,
            children=child_extracts,
            complexity=sum(ce.complexity if ce is not None else 0 for ce in child_extracts) + 1,
            value=self.get_scalar_value(),
        )

    @property
    def extract(self) -> Optional[NodeExtract]:
        return self.__extract

    @property
    def extract_not_none(self) -> NodeExtract:
        assert self.__extract is not None
        return self.__extract


class Child(Generic[_FORMULA_ITEM_TV]):
    """Descriptor for node properties that refer to child nodes"""

    __slots__ = ("ind",)

    def __init__(self, ind: int):
        self.ind = ind

    def __get__(self, instance: FormulaItem, owner: Type[FormulaItem]) -> _FORMULA_ITEM_TV:
        if instance is None:
            raise TypeError("Cannot be used on class")
        return instance.children[self.ind]  # type: ignore  # Cannot validate type here

    def __set__(self, instance: FormulaItem, value: Any) -> None:
        raise RuntimeError("Cannot set children")


class MultiChild(Generic[_FORMULA_ITEM_TV]):
    """Descriptor for node properties that refer to child nodes"""

    __slots__ = ("slice",)

    def __init__(self, slice: slice):
        self.slice = slice

    def __get__(self, instance: FormulaItem, owner: Type[FormulaItem]) -> tuple[_FORMULA_ITEM_TV, ...]:
        if instance is None:
            raise TypeError("Cannot be used on class")
        return instance.children[self.slice]  # type: ignore  # Cannot validate type here

    def __set__(self, instance: FormulaItem, value: Any) -> None:
        raise RuntimeError("Cannot set children")


class Null(FormulaItem):
    """Database NULL value"""

    __slots__ = ()

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert not children, "Null does not accept children"


class ExprWrapper(FormulaItem):
    __slots__ = ()
    show_names = FormulaItem.show_names + ("expr",)

    expr: Child[FormulaItem] = Child(0)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert len(children) == 1

    @classmethod
    def make(cls: Type[_FORMULA_ITEM_TV], expr: FormulaItem, *, meta: Optional[NodeMeta] = None) -> _FORMULA_ITEM_TV:
        return cls(expr, meta=meta)


class ParenthesizedExpr(ExprWrapper):
    """Represents the wrapping of an expression into parentheses"""

    __slots__ = ()


_LITERAL_TV = TypeVar("_LITERAL_TV", bound="BaseLiteral")


class BaseLiteral(FormulaItem):
    """Represents a constant value"""

    __slots__ = ()
    show_names = FormulaItem.show_names + ("value",)

    @property
    def value(self) -> Any:
        return self.internal_value[0]

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert len(children) == 0

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert len(internal_value) == 1
        cls.validate_literal_value(internal_value[0])

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        pass

    @classmethod
    def make(cls: Type[_LITERAL_TV], value: Any, *, meta: Optional[NodeMeta] = None) -> _LITERAL_TV:
        internal_value = (value,)
        return cls(internal_value=internal_value, meta=meta)


class Array:
    __slots__ = ("item_cls",)

    def __init__(self, item_cls: type):
        self.item_cls = item_cls

    def __str__(self) -> str:
        return f"{self.__class__.__name__}({self.item_cls})"

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}({self.item_cls})"


def does_type_match(
    value: Any,
    expected_type: Union[type, tuple[Union[type, Array], ...], Array],
) -> bool:
    if isinstance(expected_type, Array):
        if not isinstance(value, list):
            return False
        for item in value:
            if item is None or isinstance(item, Null):
                # null values are supported in arrays
                continue
            if not does_type_match(item, expected_type=expected_type.item_cls):
                return False
        return True
    elif isinstance(expected_type, tuple):
        return any(does_type_match(value, expected_type=sub_type) for sub_type in expected_type)
    elif isinstance(expected_type, type):
        return isinstance(value, expected_type)

    raise TypeError(f"Unexpected type: {expected_type}")


def check_type(value: Any, expected_type: Union[type, tuple[Union[type, Array], ...], Array]) -> None:
    if not does_type_match(value, expected_type):
        raise TypeError(f"Expected type {expected_type}, got {type(value)}")


class LiteralInteger(BaseLiteral):
    """Represents an INTEGER literal"""

    __slots__ = ()

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        assert isinstance(literal_value, int)


class LiteralFloat(BaseLiteral):
    """Represents a FLOAT literal"""

    __slots__ = ()

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        assert isinstance(literal_value, float)


class LiteralBoolean(BaseLiteral):
    """Represents a BOOLEAN literal"""

    __slots__ = ()

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        assert isinstance(literal_value, bool)


class LiteralString(BaseLiteral):
    """Represents a STRING literal"""

    __slots__ = ()

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        assert isinstance(literal_value, str)


class LiteralDate(BaseLiteral):
    """Represents a DATE literal"""

    __slots__ = ()

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        assert isinstance(literal_value, datetime.date)


class LiteralDatetime(BaseLiteral):
    """Represents a DATETIME literal"""

    __slots__ = ()

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        assert isinstance(literal_value, datetime.datetime)


class LiteralDatetimeTZ(BaseLiteral):
    """Represents a DATETIMETZ literal"""

    __slots__ = ()

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        assert isinstance(literal_value, datetime.datetime)
        if literal_value.tzinfo is None:
            raise TypeError("Expected non-empty datetime tzinfo")


class LiteralGenericDatetime(BaseLiteral):
    """Represents a GENERICDATETIME literal"""

    __slots__ = ()

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        assert isinstance(literal_value, datetime.datetime)


class LiteralGeopoint(LiteralString):
    """Represents a GEOPOINT literal"""

    __slots__ = ()


class LiteralGeopolygon(LiteralString):
    """Represents a GEOPOLYGON literal"""

    __slots__ = ()


# # Should *probably* not be done: LiteralMarkup


class LiteralUuid(BaseLiteral):
    """Represents a UUID literal"""

    __slots__ = ()

    @classmethod
    def make(cls, value: Union[uuid.UUID, str], *, meta: Optional[NodeMeta] = None) -> LiteralUuid:
        value = value if isinstance(value, uuid.UUID) else uuid.UUID(value)
        return super().make(value=value, meta=meta)

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        assert isinstance(literal_value, uuid.UUID)


class BaseLiteralArray(BaseLiteral):
    """Base class for array literals"""

    __slots__ = ()


def _validate_array(literal_value: Optional[Hashable], item_type: type | tuple[type, ...]) -> None:
    assert isinstance(literal_value, tuple)
    for el in literal_value:
        assert el is None or isinstance(el, item_type)


class LiteralArrayInteger(BaseLiteralArray):
    """Represents an ARRAY_INT literal"""

    __slots__ = ()

    @classmethod
    def make(cls, value: list[int], *, meta: Optional[NodeMeta] = None) -> LiteralArrayInteger:
        return super().make(value=tuple(value), meta=meta)

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        _validate_array(literal_value, int)


class LiteralArrayFloat(BaseLiteralArray):
    """Represents an ARRAY_FLOAT literal"""

    __slots__ = ()

    @classmethod
    def make(cls, value: list[float], *, meta: Optional[NodeMeta] = None) -> LiteralArrayFloat:
        return super().make(value=tuple(value), meta=meta)

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        _validate_array(literal_value, float)


class LiteralArrayString(BaseLiteralArray):
    """Represents an ARRAY_STR literal"""

    __slots__ = ()

    @classmethod
    def make(cls, value: list[str], *, meta: Optional[NodeMeta] = None) -> LiteralArrayString:
        return super().make(value=tuple(value), meta=meta)

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        _validate_array(literal_value, str)


class BaseLiteralTree(BaseLiteral):
    """Base class for tree literals"""

    __slots__ = ()


class LiteralTreeString(BaseLiteralArray):
    """Represents a TREE_STR literal"""

    __slots__ = ()

    @classmethod
    def make(cls, value: list[str], *, meta: Optional[NodeMeta] = None) -> LiteralTreeString:
        return super().make(value=tuple(value), meta=meta)

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        _validate_array(literal_value, str)


class ExpressionList(FormulaItem):
    """A list of expressions (for use with the IN operator)"""

    __slots__ = ()
    autonomous = False

    @classmethod
    def make(cls: Type[_FORMULA_ITEM_TV], *children: FormulaItem, meta: Optional[NodeMeta] = None) -> _FORMULA_ITEM_TV:
        return cls(*children, meta=meta)


_NAMED_ITEM_TV = TypeVar("_NAMED_ITEM_TV", bound="NamedItem")


class NamedItem(FormulaItem):
    __slots__ = ()
    show_names = FormulaItem.show_names + ("name",)

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert len(internal_value) == 1
        assert isinstance(internal_value[0], str)

    @property
    def name(self) -> str:
        return cast(str, self.internal_value[0])


class DimListNodeBase(FormulaItem):
    __slots__ = ()
    show_names = FormulaItem.show_names + ("dim_list",)

    dim_list: MultiChild[FormulaItem] = MultiChild(slice(0, None))  # TODO: rename to expr_list

    @classmethod
    def make(
        cls: Type[_FORMULA_ITEM_TV],
        dim_list: Optional[Sequence[FormulaItem]],
        meta: Optional[NodeMeta] = None,
    ) -> _FORMULA_ITEM_TV:
        return cls(*(dim_list or ()), meta=meta)


class LodSpecifier(DimListNodeBase):
    __slots__ = ()
    autonomous = False


class InheritedLodSpecifier(LodSpecifier):
    """Represents LOD clause that inherits the original subquery's dimensions"""

    __slots__ = ()

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert not children, "Dimension list is not allowed for this type of LOD"


class FixedLodSpecifier(LodSpecifier):
    """Represents the ``FIXED <dimensions>`` LOD specifier clause"""

    __slots__ = ()


class IncludeLodSpecifier(LodSpecifier):
    """Represents the ``INCLUDE <dimensions>`` LOD specifier clause"""

    __slots__ = ()


class DefaultAggregationLodSpecifier(IncludeLodSpecifier):
    """
    Represents the default aggregation LOD specifier clause.
    Basically acts as an empty INCLUDE
    """

    __slots__ = ()


class ExcludeLodSpecifier(LodSpecifier):
    """Represents the ``EXCLUDE <dimensions>`` LOD specifier clause"""

    __slots__ = ()


_FUNC_CALL_TV = TypeVar("_FUNC_CALL_TV", bound="FuncCall")


class OperationCall(NamedItem):
    """Base class for functions and operators"""

    __slots__ = ()

    args: MultiChild[FormulaItem] = MultiChild(slice(0, None))


class FuncCall(OperationCall):
    """Generic function call"""

    __slots__ = ()
    show_names = OperationCall.show_names + ("args",)
    # Omit before_filter_by, ignore_dimensions and lod because they will be very rarely used explicitly

    args: MultiChild[FormulaItem] = MultiChild(slice(0, -3))
    lod: Child[LodSpecifier] = Child(-3)
    ignore_dimensions: Child[IgnoreDimensions] = Child(-2)
    before_filter_by: Child[BeforeFilterBy] = Child(-1)

    @classmethod
    def make(
        cls: Type[_FUNC_CALL_TV],
        name: str,
        args: Iterable[FormulaItem],
        *,
        lod: Optional[LodSpecifier] = None,
        ignore_dimensions: Optional[IgnoreDimensions] = None,
        before_filter_by: Optional[BeforeFilterBy] = None,
        meta: Optional[NodeMeta] = None,
    ) -> _FUNC_CALL_TV:
        if lod is None:
            lod = DefaultAggregationLodSpecifier()
        if ignore_dimensions is None:
            ignore_dimensions = IgnoreDimensions()
        if before_filter_by is None:
            before_filter_by = BeforeFilterBy.make()

        children = (*args, lod, ignore_dimensions, before_filter_by)
        internal_value = (name,)
        return cls(*children, internal_value=internal_value, meta=meta)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert len(children) >= 3
        assert isinstance(children[-3], LodSpecifier)
        assert isinstance(children[-2], IgnoreDimensions)
        assert isinstance(children[-1], BeforeFilterBy)


class Unary(OperationCall):
    """Unary (arithmetic or logical) operation"""

    __slots__ = ()
    show_names = OperationCall.show_names + ("expr",)

    expr: Child[FormulaItem] = Child(0)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert len(children) == 1

    @classmethod
    def make(
        cls,
        name: str,
        expr: FormulaItem,
        *,
        meta: Optional[NodeMeta] = None,
    ) -> Unary:
        children = (expr,)
        internal_value = (name,)
        return cls(*children, internal_value=internal_value, meta=meta)


class Binary(OperationCall):
    """Binary (arithmetic or logical) operation"""

    __slots__ = ()
    show_names = OperationCall.show_names + ("left", "right")

    left: Child[FormulaItem] = Child(0)
    right: Child[FormulaItem] = Child(1)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert len(children) == 2

    @classmethod
    def make(
        cls,
        name: str,
        left: FormulaItem,
        right: FormulaItem,
        *,
        meta: Optional[NodeMeta] = None,
    ) -> Binary:
        children = (left, right)
        internal_value = (name,)
        return cls(*children, internal_value=internal_value, meta=meta)


class Ternary(OperationCall):
    """Ternary (arithmetic or logical) operation"""

    __slots__ = ()
    show_names = OperationCall.show_names + ("first", "second", "third")

    first: Child[FormulaItem] = Child(0)
    second: Child[FormulaItem] = Child(1)
    third: Child[FormulaItem] = Child(2)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert len(children) == 3

    @classmethod
    def make(
        cls,
        name: str,
        first: FormulaItem,
        second: FormulaItem,
        third: FormulaItem,
        *,
        meta: Optional[NodeMeta] = None,
    ) -> Ternary:
        children = (first, second, third)
        internal_value = (name,)
        return cls(*children, internal_value=internal_value, meta=meta)


class Field(NamedItem):
    """Database table field (column)"""

    __slots__ = ()

    @classmethod
    def make(cls, name: str, meta: Optional[NodeMeta] = None) -> Field:
        children = ()
        internal_value = (name,)
        return cls(*children, internal_value=internal_value, meta=meta)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert not children


class WindowGrouping(DimListNodeBase):
    """Represents the grouping clause of window functions"""

    __slots__ = ()
    autonomous = False


class WindowGroupingTotal(WindowGrouping):
    """Represents the ``TOTAL`` grouping"""

    __slots__ = ()

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert not children, f"{cls.__name__} does not accept children"


class WindowGroupingWithin(WindowGrouping):
    """Represents the ``WITHIN <dimensions>`` grouping"""

    __slots__ = ()


class WindowGroupingAmong(WindowGrouping):
    """Represents the ``AMONG <dimensions>`` grouping"""

    __slots__ = ()


class Ordering(FormulaItem):
    """Represents an ORDER BY clause"""

    __slots__ = ()
    autonomous = False
    show_names = FormulaItem.show_names + ("expr_list",)

    expr_list: MultiChild[FormulaItem] = MultiChild(slice(0, None))

    @classmethod
    def make(
        cls,
        expr_list: Optional[list[FormulaItem]] = None,
        meta: Optional[NodeMeta] = None,
    ) -> Ordering:
        children = expr_list or ()
        internal_value = ()
        return cls(*children, internal_value=internal_value, meta=meta)


_FIELD_NAME_LIST_NODE_TV = TypeVar("_FIELD_NAME_LIST_NODE_TV", bound="FieldNameListNode")


class FieldNameListNode(FormulaItem):
    """Represents clause containing a field name list"""

    __slots__ = ()
    show_names = FormulaItem.show_names + ("field_names",)

    @classmethod
    def make(
        cls: Type[_FIELD_NAME_LIST_NODE_TV],
        field_names: Optional[Collection[str]] = None,
        meta: Optional[NodeMeta] = None,
    ) -> _FIELD_NAME_LIST_NODE_TV:
        children = ()
        internal_value = (frozenset(field_names or ()),)
        return cls(*children, internal_value=internal_value, meta=meta)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert not children

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert isinstance(internal_value, tuple)
        assert len(internal_value) == 1
        name_list = internal_value[0]
        assert isinstance(name_list, frozenset)
        for name in name_list:
            assert isinstance(name, str)

    @property
    def field_names(self) -> frozenset[str]:
        return cast(frozenset[str], self.internal_value[0])


class BeforeFilterBy(FieldNameListNode):
    """Represents a BEFORE FILTER BY clause"""

    __slots__ = ()
    autonomous = False


class IgnoreDimensions(DimListNodeBase):
    """Represents an IGNORE DIMENSIONS clause"""

    __slots__ = ()
    autonomous = False


class OrderingDirectionBase(ParenthesizedExpr):
    """Base class for ASC and DESC representations"""

    __slots__ = ()
    autonomous = False


class OrderAscending(OrderingDirectionBase):
    """Represents an ASC clause modifier"""

    __slots__ = ()


class OrderDescending(OrderingDirectionBase):
    """Represents an DESC clause modifier"""

    __slots__ = ()


_WINDOW_FUNC_CALL_TV = TypeVar("_WINDOW_FUNC_CALL_TV", bound="WindowFuncCall")


class WindowFuncCall(FuncCall):
    """Represents window function calls"""

    __slots__ = ()
    show_names = FuncCall.show_names + ("grouping", "ordering", "before_filter_by")

    args: MultiChild[FormulaItem] = MultiChild(slice(0, -5))
    ordering: Child[Ordering] = Child(-5)
    grouping: Child[WindowGrouping] = Child(-4)
    # lod is -3; ignore_dimensions  is -2; before_filter_by is -1.

    @classmethod
    def make(
        cls,
        name: str,
        args: Iterable[FormulaItem],
        *,
        ordering: Optional[Ordering] = None,
        grouping: Optional[WindowGrouping] = None,
        lod: Optional[LodSpecifier] = None,
        ignore_dimensions: Optional[IgnoreDimensions] = None,
        before_filter_by: Optional[BeforeFilterBy] = None,
        meta: Optional[NodeMeta] = None,
    ) -> WindowFuncCall:
        if grouping is None:
            grouping = WindowGroupingTotal()
        if ordering is None:
            ordering = Ordering()

        return cast(
            WindowFuncCall,
            super().make(
                name=name,
                args=[*args, ordering, grouping],
                lod=lod,
                before_filter_by=before_filter_by,
                ignore_dimensions=ignore_dimensions,
                meta=meta,
            ),
        )

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        super().validate_children(children=children)

        assert len(children) >= 5
        assert isinstance(children[-5], Ordering)
        assert isinstance(children[-4], WindowGrouping)


class IfPart(FormulaItem):
    __slots__ = ()
    autonomous = False
    show_names = FormulaItem.show_names + ("cond", "expr")

    cond: Child[FormulaItem] = Child(0)
    expr: Child[FormulaItem] = Child(1)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert len(children) == 2

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert internal_value == ()

    @classmethod
    def make(
        cls,
        cond: FormulaItem,
        expr: FormulaItem,
        *,
        meta: Optional[NodeMeta] = None,
    ) -> IfPart:
        children = [cond, expr]
        return cls(*children, internal_value=(), meta=meta)


class IfBlock(FormulaItem):
    """
    Represents expressions like:
    IF condition_1 THEN expr_1 ELSEIF condition_2 THEN expr_2 ELSE expr_3 END
    """

    __slots__ = ()
    show_names = FormulaItem.show_names + ("if_list", "else_expr")

    if_list: MultiChild[IfPart] = MultiChild(slice(0, -1))
    else_expr: Child[FormulaItem] = Child(-1)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert len(children) >= 2
        for child in children[:-1]:
            assert isinstance(child, IfPart)

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert internal_value == ()

    @classmethod
    def make(
        cls,
        if_list: list[IfPart],
        else_expr: Optional[FormulaItem] = None,
        *,
        meta: Optional[NodeMeta] = None,
    ) -> IfBlock:
        for part in if_list:
            check_type(part, IfPart)
        if else_expr is None:
            else_expr = Null()
        assert else_expr is not None
        children = [*if_list, else_expr]
        return cls(*children, internal_value=(), meta=meta)


class WhenPart(FormulaItem):
    __slots__ = ()
    autonomous = False
    show_names = FormulaItem.show_names + ("val", "expr")

    val: Child[FormulaItem] = Child(0)
    expr: Child[FormulaItem] = Child(1)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert len(children) == 2

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert internal_value == ()

    @classmethod
    def make(
        cls,
        val: FormulaItem,
        expr: FormulaItem,
        meta: Optional[NodeMeta] = None,
    ) -> WhenPart:
        children = [val, expr]
        return cls(*children, meta=meta)


class CaseBlock(FormulaItem):
    """
    Represents expressions like:
    CASE expr WHEN value_1 THEN expr_1 WHEN value_2 THEN expr_2 ELSE expr_3 END
    """

    __slots__ = ()
    show_names = FormulaItem.show_names + ("case_expr", "when_list", "else_expr")

    case_expr: Child[FormulaItem] = Child(0)
    when_list: MultiChild[WhenPart] = MultiChild(slice(1, -1))
    else_expr: Child[FormulaItem] = Child(-1)

    @classmethod
    def validate_children(cls, children: Sequence[FormulaItem]) -> None:
        assert len(children) >= 2
        for child in children[1:-1]:
            assert isinstance(child, WhenPart)

    @classmethod
    def validate_internal_value(cls, internal_value: tuple[Optional[Hashable], ...]) -> None:
        assert internal_value == ()

    @classmethod
    def make(
        cls,
        case_expr: FormulaItem,
        when_list: list[WhenPart],
        else_expr: Optional[FormulaItem] = None,
        *,
        meta: Optional[NodeMeta] = None,
    ) -> CaseBlock:
        for part in when_list:
            check_type(part, WhenPart)
        if else_expr is None:
            else_expr = Null()
        assert else_expr is not None
        children = [case_expr, *when_list, else_expr]
        return cls(*children, meta=meta)


class Formula(ExprWrapper):
    """
    Top level wrapper for formula expressions.
    Provides some additional accessors for child nodes.
    """

    __slots__ = ()

lib/dl_formula/dl_formula/core/nodes.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    List,
    NamedTuple,
    Optional,
    Tuple,
)


class Position(NamedTuple):
    start: Optional[int] = None
    end: Optional[int] = None
    start_row: Optional[int] = None
    end_row: Optional[int] = None
    start_col: Optional[int] = None
    end_col: Optional[int] = None


class PositionConverter:
    def __init__(self, text: str):
        self._text: str = text
        self._pos_by_line: List[int] = [0]  # a list of starting positions for each line number
        for line in text.split("\n"):
            self._pos_by_line.append(self._pos_by_line[-1] + len(line) + 1)

    def idx_to_line_and_row(self, idx: int) -> Tuple[int, int]:
        cropped = self._text[:idx]
        row = cropped.count("\n")
        col = len(cropped.rsplit("\n", 1)[-1])
        return row, col

    def idx_to_position(self, idx: int):
        start = idx
        end = idx
        start_row, start_col = self.idx_to_line_and_row(idx)
        end_row, end_col = self.idx_to_line_and_row(idx)
        return Position(
            start=start,
            end=end,
            start_row=start_row,
            end_row=end_row,
            start_col=start_col,
            end_col=end_col,
        )

    def row_and_col_to_idx(self, line: int, column: int) -> int:
        return self._pos_by_line[line - 1] + column - 1

    def merge_positions(self, start_position: Position, end_position: Position):
        return Position(
            start=start_position.start,
            end=end_position.end,
            start_row=start_position.start_row,
            end_row=end_position.end_row,
            start_col=start_position.start_col,
            end_col=end_position.end_col,
        )

lib/dl_formula/dl_formula/core/position.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    FrozenSet,
    NamedTuple,
)


class LevelTag(NamedTuple):
    bfb_names: FrozenSet
    func_nesting: int
    qfork_nesting: int

    def __gt__(self, other: Any) -> bool:
        if type(other) is not type(self):
            raise TypeError(type(other))
        return (
            (self.bfb_names != other.bfb_names and self.bfb_names.issuperset(other.bfb_names))
            or (self.bfb_names == other.bfb_names and self.func_nesting > other.func_nesting)
            or (
                self.bfb_names == other.bfb_names
                and self.func_nesting == other.func_nesting
                and self.qfork_nesting > other.qfork_nesting
            )
        )

    def __ge__(self, other: Any) -> bool:
        if type(other) is not type(self):
            raise TypeError(type(other))
        return (
            (self.bfb_names != other.bfb_names and self.bfb_names.issuperset(other.bfb_names))
            or (self.bfb_names == other.bfb_names and self.func_nesting > other.func_nesting)
            or (
                self.bfb_names == other.bfb_names
                and self.func_nesting == other.func_nesting
                and self.qfork_nesting >= other.qfork_nesting
            )
        )

    def __lt__(self, other: Any) -> bool:
        if type(other) is not type(self):
            raise TypeError(type(other))
        return (
            (self.bfb_names != other.bfb_names and self.bfb_names.issubset(other.bfb_names))
            or (self.bfb_names == other.bfb_names and self.func_nesting < other.func_nesting)
            or (
                self.bfb_names == other.bfb_names
                and self.func_nesting == other.func_nesting
                and self.qfork_nesting < other.qfork_nesting
            )
        )

    def __le__(self, other: Any) -> bool:
        if type(other) is not type(self):
            raise TypeError(type(other))
        return (
            (self.bfb_names != other.bfb_names and self.bfb_names.issubset(other.bfb_names))
            or (self.bfb_names == other.bfb_names and self.func_nesting < other.func_nesting)
            or (
                self.bfb_names == other.bfb_names
                and self.func_nesting == other.func_nesting
                and self.qfork_nesting <= other.qfork_nesting
            )
        )

    # use default eq & ne

lib/dl_formula/dl_formula/core/tag.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    AbstractSet,
    Optional,
    Sequence,
    Set,
    Union,
)

from dl_formula.core.datatype import DataType
from dl_formula.definitions.flags import (
    ContextFlag,
    ContextFlags,
)


class ArgTypeMatcher:
    __slots__ = ()

    def match_arg_types(self, arg_types: Sequence[DataType]) -> bool:
        raise NotImplementedError

    def get_possible_arg_types_at_pos(self, pos: int, total: int) -> Set[DataType]:
        raise NotImplementedError


class ArgTypeSequence(ArgTypeMatcher):
    __slots__ = ("_exp_arg_types",)

    def __init__(self, arg_types: Sequence[Union[DataType, AbstractSet[DataType]]]):
        self._exp_arg_types = arg_types

    def match_arg_types(self, arg_types: Sequence[DataType]) -> bool:
        for expected_arg_type, real_arg_type in zip(self._exp_arg_types, arg_types):
            if isinstance(expected_arg_type, DataType):
                # a single type
                if not real_arg_type.casts_to(expected_arg_type):
                    return False
            else:  # set, frozenset
                if not any(
                    real_arg_type.casts_to(expected_arg_sub_type) for expected_arg_sub_type in expected_arg_type
                ):
                    return False
        return True

    def get_possible_arg_types_at_pos(self, pos: int, total: int) -> Set[DataType]:
        expected_arg_type = self._exp_arg_types[pos]
        if isinstance(expected_arg_type, DataType):
            return {expected_arg_type}
        else:  # set, frozenset
            return set(expected_arg_type)


class ArgTypeForAll(ArgTypeMatcher):
    __slots__ = ("_exp_arg_types", "_require_type_match")

    def __init__(
        self,
        expected_types: Union[DataType, AbstractSet[DataType]],
        require_type_match: Optional[Union[DataType, AbstractSet[DataType]]] = None,
    ):
        self._exp_arg_types = {expected_types} if isinstance(expected_types, DataType) else expected_types
        self._require_type_match = (
            {require_type_match} if isinstance(require_type_match, DataType) else require_type_match
        )

    def match_arg_types(self, arg_types: Sequence[DataType]) -> bool:
        if self._require_type_match is not None and not any(
            real_arg_type in self._require_type_match for real_arg_type in arg_types
        ):
            return False
        for real_arg_type in arg_types:
            if not any(real_arg_type.casts_to(expected_arg_sub_type) for expected_arg_sub_type in self._exp_arg_types):
                return False
        return True

    def get_possible_arg_types_at_pos(self, pos: int, total: int) -> Set[DataType]:
        return set(self._exp_arg_types)


class ArgFlagDispenser:
    """Defines which arguments get marked with context flags"""

    __slots__ = ()

    def get_flags_for_pos(self, pos: int, total: int) -> Optional[ContextFlags]:
        raise NotImplementedError


class ArgFlagSequence(ArgFlagDispenser):
    __slots__ = ("_arg_flags",)

    def __init__(self, arg_flags: Sequence[Optional[ContextFlags]]):
        self._arg_flags = arg_flags

    def get_flags_for_pos(self, pos: int, total: int) -> Optional[ContextFlags]:
        return self._arg_flags[pos]


class IfFlagDispenser(ArgFlagDispenser):
    __slots__ = ()

    def get_flags_for_pos(self, pos: int, total: int) -> Optional[ContextFlags]:
        if pos < total - 1 and pos % 2 == 0:
            # all even positions, except for last item -> condition
            return ContextFlag.REQ_CONDITION
        return None

lib/dl_formula/dl_formula/definitions/args.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import inspect
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Generic,
    Iterable,
    NamedTuple,
    Optional,
    Type,
    TypeVar,
    Union,
)

import attr
from sqlalchemy.sql.elements import (
    ClauseElement,
    ClauseList,
    FunctionFilter,
    TypeCoerce,
)
from sqlalchemy.sql.functions import Function as SAFunction

from dl_formula.core import (
    exc,
    nodes,
)
from dl_formula.core.datatype import (
    DataType,
    DataTypeParams,
)
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import (
    ArgFlagDispenser,
    ArgTypeMatcher,
)
from dl_formula.definitions.flags import ContextFlags
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import (
    FromArgs,
    ParamsEmpty,
)


if TYPE_CHECKING:
    from dl_formula.definitions.type_strategy import (
        TypeParamsStrategy,
        TypeStrategy,
    )
    from dl_formula.translation.context import TranslationCtx
    from dl_formula.translation.env import TranslationEnvironment


# FIXME: uncomment commented __slots__ if and when Python version is updated (bug in 3.5)


_VARIANT_OF_TV = TypeVar("_VARIANT_OF_TV")
_VARIANT_TV = TypeVar("_VARIANT_TV", bound="ValueVariant")


class ValueVariant(Generic[_VARIANT_OF_TV]):
    """
    Wrapper for a value to be chosen from a list by dialect.

    Method ``.match(dialect)`` should be called to check whether this value can be used for given dialect
    """

    __slots__ = ("dialects", "_value")

    def __init__(self, dialects: DialectCombo, value: _VARIANT_OF_TV):
        self.dialects: DialectCombo = dialects
        self._value: _VARIANT_OF_TV = value

    def clone(self: _VARIANT_TV, **kwargs) -> _VARIANT_TV:
        copy_kwargs = dict(
            dialects=self.dialects,
            value=self._value,
        )
        copy_kwargs.update(kwargs)
        return type(self)(**copy_kwargs)

    @property
    def value(self) -> _VARIANT_OF_TV:
        return self._value

    def match(self, dialect) -> bool:
        """Return True if this variant matches the given dialect"""
        return self.dialects == D.ANY or self.dialects & dialect == dialect


TransCallResult = Union[ClauseElement, nodes.FormulaItem]
TranslateCallback = Callable[[nodes.FormulaItem], TransCallResult]


_TRANS_IMPL_TV = TypeVar("_TRANS_IMPL_TV", bound="FuncTranslationImplementationBase")


@attr.s  # FIXME: frozen=True ?
class FuncTranslationImplementationBase(abc.ABC):
    unwrap_args: bool = attr.ib(default=True)  # TODO: make an args-unwrapping subclass instead

    def clone(self: _TRANS_IMPL_TV, **kwargs: Any) -> _TRANS_IMPL_TV:
        return attr.evolve(**kwargs)

    @abc.abstractmethod
    def translate(
        self,
        *raw_args: TranslationCtx,
        translator_cb: TranslateCallback,
        partition_by: Optional[ClauseList] = None,
        default_order_by: Optional[ClauseList] = None,
        translation_ctx: Optional[TranslationCtx] = None,
        translation_env: Optional[TranslationEnvironment] = None,
    ) -> ClauseElement:
        """
        Attempts to translate the node by calling it with the given ``args``.
        If the resulting value is a ``nodes.FormulaItem``,
        then call the translator callback to re-translate the node.
        Keep doing this until a ``ClauseElement`` is received.

        This callback can also be used for compiling complex SQLA objects
        that might require re-translation of intermediate results (e.g. window functions).
        """
        raise NotImplementedError

    @staticmethod
    def _await_conversion_to_sa(obj: TransCallResult, translator_cb: TranslateCallback) -> ClauseElement:
        """
        Wait until a node object becaomes an SQLAlchemy ``ClauseElement``
        by calling the translator callback in a cycle.
        """

        max_recursion_depth = 100
        for _ in range(max_recursion_depth):
            if isinstance(obj, nodes.FormulaItem):
                # got a nodes.FormulaItem that has to be translated again
                obj = translator_cb(obj)
            elif isinstance(obj, ClauseElement):
                # got sa object -> time to exit
                return obj
            else:
                raise TypeError(f"Invalid object type {type(obj)} for translation cycle")

        raise RuntimeError(f"Maximum translation recursion depth {max_recursion_depth} reached")

    @classmethod
    def _unwrap_translation_ctx_args(cls, args: Iterable[TranslationCtx]) -> list[ClauseElement]:
        result = []
        for arg in args:
            assert arg.expression is not None
            result.append(arg.expression)
        return result

    def _handle_args(
        self,
        args: tuple[TranslationCtx, ...],
        *,
        translation_func: Callable,
        translation_ctx: Optional[TranslationCtx] = None,
        translation_env: Optional[TranslationEnvironment] = None,
    ) -> tuple[Iterable[Union[TranslationCtx, ClauseElement]], dict[str, Any]]:
        handled_args: Iterable[Union[TranslationCtx, ClauseElement]] = args
        if self.unwrap_args:
            handled_args = self._unwrap_translation_ctx_args(args)
        # some mandatory unwrapping:
        # remove TypeCoerce wrapper because of bug in SA
        # that does not put parentheses around such expressions
        handled_args = tuple(arg.clause if isinstance(arg, TypeCoerce) else arg for arg in handled_args)

        handled_kwargs = {}
        possible_kwargs = {"_ctx": translation_ctx, "_env": translation_env}
        # NOTE: using `unwrap` is slightly dubious (it unwraps the
        # `functools.wraps` through `func.__wrapped__`)
        argspec = inspect.getfullargspec(inspect.unwrap(translation_func))
        for key, value in possible_kwargs.items():
            if value is None:
                continue
            # NOTE: passing to all `**kwargs`-functions too.
            if key in argspec.args or key in argspec.kwonlyargs or argspec.varkw:
                handled_kwargs[key] = value

        return handled_args, handled_kwargs


@attr.s
class FuncTranslationImplementation(FuncTranslationImplementationBase):  # noqa
    """
    Basic function implementation class.
    Consists of a single "main" translatable part that is generated
    by a ``self.translation_main(*args)`` call.
    """

    translation_main: Callable[..., TransCallResult] = attr.ib(kw_only=True)

    def translate(
        self,
        *raw_args: TranslationCtx,
        translator_cb: TranslateCallback,
        partition_by: Optional[ClauseList] = None,
        default_order_by: Optional[ClauseList] = None,
        translation_ctx: Optional[TranslationCtx] = None,
        translation_env: Optional[TranslationEnvironment] = None,
    ) -> ClauseElement:
        translation_main = self.translation_main
        assert translation_main is not None
        args, kwargs = self._handle_args(
            raw_args,
            translation_func=translation_main,
            translation_ctx=translation_ctx,
            translation_env=translation_env,
        )
        # generate the result
        return self._await_conversion_to_sa(translation_main(*args, **kwargs), translator_cb=translator_cb)


WinRangeTuple = tuple[Optional[int], Optional[int]]


@attr.s
class WindowFunctionImplementation(FuncTranslationImplementation):
    """
    Class for window function implementations.
    Here the sa clause is compiled of the "main" and "order_by" parts
    that are generated separately.
    """

    translation_order_by: Optional[Callable[..., TransCallResult]] = attr.ib(default=None)
    translation_range: Optional[Callable[..., WinRangeTuple]] = attr.ib(default=None)
    translation_rows: Optional[Callable[..., WinRangeTuple]] = attr.ib(default=None)

    def translate(
        self,
        *raw_args: TranslationCtx,
        translator_cb: TranslateCallback,
        partition_by: Optional[ClauseList] = None,
        default_order_by: Optional[ClauseList] = None,
        translation_ctx: Optional[TranslationCtx] = None,
        translation_env: Optional[TranslationEnvironment] = None,
    ) -> ClauseElement:
        translation_main = self.translation_main
        assert translation_main is not None
        # Constructing args for `translation_main`; for simplicity, for now,
        # not reconstructing them for other parts of the window expression.
        args, kwargs = self._handle_args(
            raw_args,
            translation_func=translation_main,
            translation_ctx=translation_ctx,
            translation_env=translation_env,
        )

        # generate the result
        func_part = self._await_conversion_to_sa(translation_main(*args, **kwargs), translator_cb=translator_cb)
        assert isinstance(func_part, (SAFunction, FunctionFilter))

        order_by_part = default_order_by
        if self.translation_order_by is not None:
            order_by_part = self._await_conversion_to_sa(
                self.translation_order_by(*args),  # NOTE: not passing the `kwargs` here at the moment.
                translator_cb=translator_cb,
            )
            assert isinstance(order_by_part, ClauseList)

        range_part = None
        if self.translation_range is not None:
            range_part = self.translation_range(*args)  # NOTE: not passing the `kwargs` here at the moment.

        rows_part = None
        if self.translation_rows is not None:
            rows_part = self.translation_rows(*args)  # NOTE: not passing the `kwargs` here at the moment.

        # Note that an `Over` object cannot be simply reconstructed from its parts
        # as attributes of the existing `Over` instance.
        # RANGE_UNBOUNDED has to be replaced with None and RANGE_CURRENT with 0.
        return func_part.over(partition_by=partition_by, order_by=order_by_part, range_=range_part, rows=rows_part)


_TRANS_VAR_TV = TypeVar("_TRANS_VAR_TV", bound="TranslationVariant")


class TranslationVariant(ValueVariant[FuncTranslationImplementationBase]):
    """
    A ``ValueVariant`` for function translations.
    Ties together dialect and function implementation
    """

    __slots__ = ()
    _repr_attrs = ("dialects",)

    @classmethod
    def make(
        cls: Type[_TRANS_VAR_TV],
        dialects: DialectCombo,
        translation: Callable[..., TransCallResult],
        unwrap_args: bool = True,
        as_winfunc: bool = False,
        translation_order_by: Optional[Callable[..., TransCallResult]] = None,
        translation_range: Optional[Callable[..., WinRangeTuple]] = None,
        translation_rows: Optional[Callable[..., WinRangeTuple]] = None,
    ) -> _TRANS_VAR_TV:
        translation_impl: FuncTranslationImplementationBase  # noqa
        if as_winfunc:
            translation_impl = WindowFunctionImplementation(
                translation_main=translation,
                unwrap_args=unwrap_args,
                translation_order_by=translation_order_by,
                translation_range=translation_range,
                translation_rows=translation_rows,
            )
        else:
            translation_impl = FuncTranslationImplementation(
                translation_main=translation,
                unwrap_args=unwrap_args,
            )

        instance = cls(dialects=dialects, value=translation_impl)
        return instance

    def __repr__(self):
        _c0 = "\x1b[0m"
        _c1 = "\x1b[038;5;118m"
        return "<{}({})>{}".format(
            self.__class__.__name__,
            ",     ".join(f"{_c1}{key}{_c0}={getattr(self, key)!r}" for key in self._repr_attrs),
            _c0,
        )

    def translate(
        self,
        *args,
        translator_cb: TranslateCallback,
        partition_by: Optional[ClauseList] = None,
        default_order_by: Optional[ClauseList] = None,
        translation_ctx: Optional[TranslationCtx] = None,
        translation_env: Optional[TranslationEnvironment] = None,
    ) -> ClauseElement:
        return self.value.translate(
            *args,
            translator_cb=translator_cb,
            partition_by=partition_by,
            default_order_by=default_order_by,
            translation_ctx=translation_ctx,
            translation_env=translation_env,
        )

    def match(self, dialect) -> bool:
        return super().match(dialect)


_TRANS_VAR_WR_TV = TypeVar("_TRANS_VAR_WR_TV", bound="TranslationVariantWrapped")


# TODO: similar class for `as_winfunc`
# TODO: Do this more properly.
class TranslationVariantWrapped(TranslationVariant):
    """
    Variant subclass that passes `TranslationCtx` rather than SA
    `ClauseElement` to the translation.
    """

    @classmethod
    def make(cls: Type[_TRANS_VAR_WR_TV], *args: Any, **kwargs: Any) -> _TRANS_VAR_WR_TV:
        return super().make(*args, unwrap_args=False, **kwargs)  # type: ignore  # TODO: fix


class TranslationResult(NamedTuple):
    impl_callable: Callable[..., ClauseElement]
    transformed_args: list[TranslationCtx]
    data_type: DataType
    data_type_params: DataTypeParams
    context_flags: int
    postprocess_args: bool


class PartialTranslationResult(NamedTuple):
    impl_callable: Optional[Callable[..., ClauseElement]] = None
    transformed_args: Optional[list[TranslationCtx]] = None
    data_type: Optional[DataType] = None
    data_type_params: Optional[DataTypeParams] = None
    context_flags: Optional[int] = None


class ArgTransformer:
    def transform_args(
        self,
        env: TranslationEnvironment,
        args: list[TranslationCtx],
        arg_types: list[DataType],
    ) -> list[TranslationCtx]:
        """Perform custom transformations of arguments"""
        return args


_NODE_TRANS_TV = TypeVar("_NODE_TRANS_TV", bound="NodeTranslation")


class NodeTranslation:
    __slots__ = ()

    # the following pair acts as a sort of identifier for the function/operator translation
    name: ClassVar[Optional[str]] = None
    arg_cnt: ClassVar[Optional[int]] = None

    is_function: ClassVar[bool] = True
    is_aggregation: ClassVar[bool] = False
    is_window: ClassVar[bool] = False
    uses_default_ordering: ClassVar[bool] = False
    supports_grouping: ClassVar[bool] = False
    supports_ordering: ClassVar[bool] = False
    supports_lod: ClassVar[bool] = False
    supports_ignore_dimensions: ClassVar[bool] = False
    supports_bfb: ClassVar[bool] = False

    scopes: ClassVar[int] = Scope.STABLE | Scope.EXPLICIT_USAGE | Scope.SUGGESTED | Scope.DOCUMENTED

    arg_names: ClassVar[Optional[list[str]]] = None

    @classmethod
    def match_types(cls, arg_types: list[DataType]) -> bool:
        raise NotImplementedError

    @classmethod
    def get_return_type(cls, arg_types: list[DataType]) -> DataType:
        raise NotImplementedError

    @classmethod
    def _get_return_type_info(cls, arg_types: list[DataType]) -> DataType:
        raise NotImplementedError

    def match_dialect(self, dialect: DialectCombo) -> bool:
        raise NotImplementedError

    def transform_args(
        self,
        env: TranslationEnvironment,
        args: list[TranslationCtx],
        arg_types: list[DataType],
    ) -> list[TranslationCtx]:
        """Perform custom transformations of arguments"""
        return args

    def translate(
        self,
        env: TranslationEnvironment,
        args: list[TranslationCtx],
    ) -> TranslationResult:
        raise NotImplementedError

    def supported_dialects(self) -> DialectCombo:
        raise NotImplementedError

    @classmethod
    def for_trans_var(cls: Type[_NODE_TRANS_TV], variant: TranslationVariant) -> _NODE_TRANS_TV:
        raise NotImplementedError

    @classmethod
    def for_dialect(
        cls: Type[_NODE_TRANS_TV],
        dialects: DialectCombo,
        arg_transformer: Optional[ArgTransformer] = None,
    ) -> _NODE_TRANS_TV:
        raise NotImplementedError


_MULTI_NODE_TRANS_TV = TypeVar("_MULTI_NODE_TRANS_TV", bound="MultiVariantTranslation")


class MultiVariantTranslation(NodeTranslation):
    __slots__ = ("_inst_variants", "_inst_arg_transformer")

    variants: ClassVar[list[TranslationVariant]] = []
    argument_types: ClassVar[Optional[list[ArgTypeMatcher]]] = None
    argument_flags: ClassVar[Optional[ArgFlagDispenser]] = None
    return_type: ClassVar[TypeStrategy] = FromArgs()
    return_type_params: ClassVar[TypeParamsStrategy] = ParamsEmpty()
    return_flags: ClassVar[ContextFlags] = 0
    postprocess_args: ClassVar[bool] = True
    arg_transformer: ClassVar[ArgTransformer] = ArgTransformer()

    # Instance vars
    _inst_variants: list[TranslationVariant]
    _inst_arg_transformer: ArgTransformer

    def __init__(
        self,
        variants: Optional[list[TranslationVariant]] = None,
        arg_transformer: Optional[ArgTransformer] = None,
    ):
        if variants is None:
            variants = self.variants
        assert variants is not None
        self._inst_variants = variants

        if arg_transformer is None:
            arg_transformer = self.arg_transformer
        assert arg_transformer is not None
        self._inst_arg_transformer = arg_transformer

        super().__init__()

    def get_variants(self) -> list[TranslationVariant]:
        """Override point (to avoid making a class property or a metaclass)"""
        return self._inst_variants

    def transform_args(
        self,
        env: TranslationEnvironment,
        args: list[TranslationCtx],
        arg_types: list[DataType],
    ) -> list[TranslationCtx]:
        """Perform custom transformations of arguments"""
        return self._inst_arg_transformer.transform_args(env=env, args=args, arg_types=arg_types)

    @classmethod
    def match_types(cls, arg_types: list[DataType]) -> bool:
        """
        Check whether data types of arguments match the required ones.
        Raise error if they don't.
        """
        if cls.argument_types is None:
            return True

        return any(arg_matcher.match_arg_types(arg_types) for arg_matcher in cls.argument_types)

    # TODO?: rename.
    @classmethod
    def get_return_type(cls, arg_types: list[DataType]) -> DataType:
        """NOTE: only for auto-documentation"""
        return cls.return_type.get_from_args(arg_types=arg_types)

    def match_dialect(self, dialect: DialectCombo) -> bool:
        return any(variant.match(dialect=dialect) for variant in self.get_variants())

    @classmethod
    def _get_argument_flag_dispenser(cls, dialect: DialectCombo) -> Optional[ArgFlagDispenser]:
        return cls.argument_flags

    @classmethod
    def _get_return_flags(cls, dialect: DialectCombo) -> ContextFlags:
        return cls.return_flags

    @classmethod
    def _get_return_type_info(cls, args):
        arg_types = [arg.data_type for arg in args]
        data_type = cls.return_type.get_from_args(arg_types)
        data_type_params = cls.return_type_params.get_from_arg_values(args)
        return data_type, data_type_params

    def translate(
        self,
        env: TranslationEnvironment,
        args: list[TranslationCtx],
    ) -> TranslationResult:
        """
        Find a translation variant that matches function name and argument types.
        Return a ``TranslationResult`` instance that consists of:
            0. result of the translation, which is either:
                a) a callable that should be called with transformed arguments
                b) a ``nodes.FormulaItem`` instance that has to be translated once more
            1. transformed list of arguments (most common transformation is changing of their order
                to precisely match the translation's signature in case of commutative operations)
            2. function result's data type (``datatype.DataType`` member)
            3. function result's data type parameters (e.g. timezone, precision, ...)
            4. function result's context flags (a bitwise combination of ``flags.ContextFlag`` members)
        """

        arg_types = []
        for arg in args:
            assert arg.data_type is not None
            arg_types.append(arg.data_type)

        args = self.transform_args(env=env, args=args, arg_types=arg_types)

        for variant in self.get_variants():
            if variant.match(dialect=env.dialect):
                # apply argument flags to args
                flag_dispenser = self._get_argument_flag_dispenser(env.dialect)
                if flag_dispenser is not None:
                    # found a match for dialect
                    for pos, arg in enumerate(args):
                        arg_flags = flag_dispenser.get_flags_for_pos(pos, total=len(args))
                        if arg_flags:
                            arg.set_flags(arg_flags)

                # result
                return_flags = self._get_return_flags(env.dialect)
                data_type, data_type_params = self._get_return_type_info(args)
                return TranslationResult(
                    impl_callable=variant.translate,
                    transformed_args=args,
                    data_type=data_type,
                    data_type_params=data_type_params,
                    context_flags=return_flags,
                    postprocess_args=self.postprocess_args,
                )

        dialect_bits = "|".join([basic_dialect.single_bit.name.name for basic_dialect in env.dialect.to_list()])
        raise exc.TranslationUnknownFunctionError(
            f'Function "{self.name}" is not implemented '
            f"for {dialect_bits} dialect and given arguments "
            f'({"|".join(t.name for t in arg_types)})'
        )

    def supported_dialects(self) -> DialectCombo:
        dialects: DialectCombo = D.EMPTY
        for variant in self.get_variants():
            dialects |= variant.dialects
        return dialects

    @classmethod
    def for_trans_var(cls: Type[_MULTI_NODE_TRANS_TV], variant: TranslationVariant) -> _MULTI_NODE_TRANS_TV:
        return cls(variants=[variant])

    @classmethod
    def for_dialect(
        cls: Type[_MULTI_NODE_TRANS_TV],
        dialects: DialectCombo,
        arg_transformer: Optional[ArgTransformer] = None,
    ) -> _MULTI_NODE_TRANS_TV:
        assert len(cls.variants) == 1
        variant = next(iter(cls.variants))
        patched_variant = variant.clone(dialects=dialects)
        return cls(variants=[patched_variant], arg_transformer=arg_transformer)


_SINGLE_NODE_TRANS_TV = TypeVar("_SINGLE_NODE_TRANS_TV", bound="SingleVariantTranslationBase")


class SingleVariantTranslationBase(MultiVariantTranslation):
    """
    'One' is a subset of 'many';
    this class implements a more convenient and extended form of defining an
    implementation for only one variant,
    for cases where databases behave too differently.
    """

    __slots__ = ("_inst_dialects",)

    dialects: ClassVar[Optional[DialectCombo]] = None

    # Instance vars
    _inst_dialects: Optional[DialectCombo]

    def __init__(
        self,
        dialects: Optional[DialectCombo] = None,
        arg_transformer: Optional[ArgTransformer] = None,
    ):
        if dialects is None:
            dialects = self.dialects

        assert dialects is not None
        self._inst_dialects = dialects
        super().__init__(arg_transformer=arg_transformer)

    def _translate_main(cls, *args):
        """Variant implementation"""
        raise NotImplementedError

    def get_variants(self) -> list[TranslationVariant]:
        """
        Semi-mockup for all supporting methods
        (e.g `match_dialect`, `variant.match`, `supported_dialects`).
        """
        return [TranslationVariantWrapped.make(self._inst_dialects, self._translate_main)]

    @classmethod
    def for_trans_var(cls: Type[_SINGLE_NODE_TRANS_TV], variant: TranslationVariant) -> _SINGLE_NODE_TRANS_TV:
        raise NotImplementedError

    @classmethod
    def for_dialect(
        cls: Type[_SINGLE_NODE_TRANS_TV],
        dialects: DialectCombo,
        arg_transformer: Optional[ArgTransformer] = None,
    ) -> _SINGLE_NODE_TRANS_TV:
        return cls(dialects=dialects)


class SingleVariantFullOverrideTranslationBase(SingleVariantTranslationBase):
    """
    A version of SingleVariantTranslationBase that allows overriding all parts of
    TranslationResult in the variant implementation.

    In the best case, this will never be used, and will only ever be an example.
    """

    __slots__ = ()

    @classmethod
    def _translate_all(cls, *args, env=None) -> PartialTranslationResult:
        """Main overridable method"""
        raise NotImplementedError

    @classmethod
    def _translate_main(cls, *args, **kwargs):
        raise Exception(
            "SingleVariantFullOverrideTranslation._translate_main dummy method was called",
            dict(cls=cls, args=args, kwargs=kwargs),
        )

    def translate(self, env, args) -> TranslationResult:
        sup_result: TranslationResult = super().translate(env=env, args=args)
        var_result: PartialTranslationResult = self._translate_all(*args, env=env)
        return sup_result._replace(**{key: val for key, val in var_result._asdict().items() if val is not None})


class Function(MultiVariantTranslation):
    is_function = True

lib/dl_formula/dl_formula/definitions/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import TYPE_CHECKING

import sqlalchemy as sa
from sqlalchemy.sql.elements import (
    UnaryExpression,
    WithinGroup,
)
from sqlalchemy.sql.operators import desc_op

from dl_formula.core import (
    exc,
    nodes,
)
from dl_formula.shortcuts import n


if TYPE_CHECKING:
    from sqlalchemy.sql.elements import ClauseElement


class _TextClauseHack(sa.sql.elements.TextClause):
    def __bool__(self):
        # Possibly a bug in sqlalchemy 1.4: query caching fails as, normally,
        # 'Boolean value of this clause is not defined'.
        # Not that the cache should be used here anyway, most likely.
        return True


def raw_sql(sql_text):
    """Basically just `sa.text(sql_text)` with hacks"""
    return _TextClauseHack(sql_text)


def make_binary_chain(binary_callable, *args, wrap_as_nodes: bool = True):
    if len(args) == 0:
        return None

    expr = args[0]
    for next_arg in args[1:]:
        binary_args = expr, next_arg
        if wrap_as_nodes:
            expr = binary_callable(*binary_args)
        else:
            expr = binary_callable(
                *[
                    n.lit(arg) if wrap_as_nodes and not isinstance(arg, nodes.FormulaItem) else arg
                    for arg in binary_args
                ]
            )
    return expr


def quantile_value(quant: float):
    """
    Validate and normalize the value for a quantile.
    """
    # Note: disallowing the `0` and `1` for some extra unambiguity
    # (for those who might think that `1` is `1%`)
    if quant <= 0:
        raise exc.TranslationError("Invalid quantile value: value must be greater than 0")
    if quant >= 1:
        raise exc.TranslationError("Invalid quantile value: value must be less than 1")
    return quant


def desc(value: ClauseElement) -> ClauseElement:
    if isinstance(value, UnaryExpression) and value.modifier is desc_op:
        # already wrapped in desc, so just unwrap
        value = value.element
    else:
        # not wrapped in desc, so wrap it
        value = sa.desc(value)

    return value


def ifnotnull(value, expr):
    """Wrap `expr` in an `sa.case` that results in a NULL for a NULL `value`"""
    return sa.case([(value.is_(None), sa.null())], else_=expr)


class _PatchedWithinGroup(WithinGroup):
    def __reduce__(self) -> tuple:
        return self.__class__, (self.element, *self.order_by)  # type: ignore


def within_group(clause_el: ClauseElement, *order_by) -> _PatchedWithinGroup:
    return _PatchedWithinGroup(clause_el, *order_by)

lib/dl_formula/dl_formula/definitions/common.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import functools
from typing import (
    TYPE_CHECKING,
    Any,
    Tuple,
)

import sqlalchemy as sa

from dl_formula.core import (
    exc,
    nodes,
)
from dl_formula.core.datatype import DataType
from dl_formula.definitions.literals import un_literal
from dl_formula.shortcuts import n


if TYPE_CHECKING:
    from sqlalchemy.sql.elements import ClauseElement


DAY_SEC = 3600 * 24
DAY_USEC = DAY_SEC * 1_000_000  # microseconds

EPOCH_START_S = "1970-01-01"
EPOCH_START_D = datetime.date(1970, 1, 1)
EPOCH_START_DOW = 4  # it was a Thursday

SUPPORTED_INTERVAL_TYPES = ("second", "minute", "hour", "day", "week", "month", "quarter", "year")


def normalize_and_validate_datetime_interval_type(type_name: str) -> str:
    type_name = type_name.lower()
    if type_name not in SUPPORTED_INTERVAL_TYPES:
        raise exc.TranslationError("Invalid interval type: '{}'".format(type_name))
    return type_name


YQL_INTERVAL_FUNCS = {
    "second": "IntervalFromSeconds",
    "minute": "IntervalFromMinutes",
    "hour": "IntervalFromHours",
    "day": "IntervalFromDays",
}
YQL_SHIFT_FUNCS = {
    "month": "ShiftMonths",
    "quarter": "ShiftQuarters",
    "year": "ShiftYears",
}


def _date_datetime_add_yql(value_expr, type_expr, mult_expr, *, const_mult: bool) -> ClauseElement:
    type_name = un_literal(type_expr)
    type_name = normalize_and_validate_datetime_interval_type(type_name)

    if not const_mult:
        # YQL requires a non-nullable mult;
        # so ensure it is such.
        mult_expr = sa.func.coalesce(mult_expr, 0)

    if type_name == "week":
        type_name = "day"
        mult_expr = mult_expr * 7

    func_name = YQL_INTERVAL_FUNCS.get(type_name)
    if func_name is not None:
        func = getattr(sa.func.DateTime, func_name)
        return value_expr + func(mult_expr)

    func_name = YQL_SHIFT_FUNCS.get(type_name)
    if func_name is not None:
        func = getattr(sa.func.DateTime, func_name)
        return func(value_expr, mult_expr)

    raise ValueError("Unexpectedly unsupported YQL datetime shift", type_name)


def date_add_yql(value_expr, type_expr, mult_expr, *, const_mult: bool) -> ClauseElement:
    expr = _date_datetime_add_yql(value_expr, type_expr, mult_expr, const_mult=const_mult)
    return sa.func.DateTime.MakeDate(expr)


def datetime_add_yql(value_expr, type_expr, mult_expr, *, const_mult: bool) -> ClauseElement:
    expr = _date_datetime_add_yql(value_expr, type_expr, mult_expr, const_mult=const_mult)
    return sa.func.DateTime.MakeDatetime(expr)


def datetime_interval_ch(type_name: str, mult: int) -> ClauseElement:
    type_name = normalize_and_validate_datetime_interval_type(type_name)
    func_name = "toInterval{}".format(type_name.capitalize())
    return getattr(sa.func, func_name)(mult)


def datetime_interval(
    type_name: str,
    mult: int,
    caps: bool = True,
    literal_mult: bool = False,
    literal_type: bool = False,
    ch_func: bool = False,
) -> ClauseElement:
    if ch_func:
        return datetime_interval_ch(type_name, mult)

    type_name = normalize_and_validate_datetime_interval_type(type_name)

    if literal_mult:
        if literal_type:
            sql = "INTERVAL '{} {}'".format(mult, type_name)
        else:
            sql = "INTERVAL '{}' {}".format(mult, type_name)
    else:
        sql = "INTERVAL {} {}".format(mult, type_name)

    if caps:
        sql = sql.upper()
    else:
        sql = sql.lower()

    return sa.text(sql)


def ensure_naive_datetime(datetime_ctx):
    """Convert DATETIMETZ values to DATETIME values; leave the other values as-is"""
    if isinstance(datetime_ctx, nodes.FuncCall) and datetime_ctx.name.lower() == "datetimetz_to_naive":
        # double-ensure
        return datetime_ctx
    if datetime_ctx.data_type == DataType.DATETIMETZ:
        datetime_ctx = n.func.DATETIMETZ_TO_NAIVE(datetime_ctx)
    return datetime_ctx


def ensure_naive_first_arg(func):
    """Make sure the first `func` argument is a naive DATETIME if it was a DATETIMETZ"""
    # TODO: this should have probably been an autocast-by-function (same with str->markup)

    @functools.wraps(func)
    def _ensured_naive_first_arg(datetime_ctx, *args, **kwargs):
        datetime_ctx = ensure_naive_datetime(datetime_ctx)
        return func(datetime_ctx, *args, **kwargs)

    return _ensured_naive_first_arg


UTC_CH_TZ_ARGS = ("UTC",)  # TODO later: make empty (backwards-incompatible change)


def make_ch_tz_args(date_ctx) -> Tuple[Tuple[Any, ...], Tuple[Any, ...]]:
    """
    For a datetime TranslationCtx argument,
    return a pair of CH function `*args`:
      * for datetrunc-functions
      * for output `toDateTime` cast

    Intended for all cases, but especially for cases that return a datetime,
    such as `datetrunc`.
    """
    return (
        # processing-related TZ args
        (
            (date_ctx.data_type_params.timezone,)
            if date_ctx.data_type in (DataType.DATETIMETZ, DataType.CONST_DATETIMETZ)
            else ()
            if date_ctx.data_type in (DataType.GENERICDATETIME, DataType.CONST_GENERICDATETIME)
            else UTC_CH_TZ_ARGS
        ),
        # output-case TZ args
        (() if date_ctx.data_type in (DataType.GENERICDATETIME, DataType.CONST_GENERICDATETIME) else ("UTC",)),
    )


def ch_date_with_tz(date_ctx) -> Tuple[ClauseElement, ...]:
    """Primarily intended for functions that return a non-datetime, such as toSecond()"""
    date_expr = date_ctx.expression
    data_type = date_ctx.data_type
    if data_type in (
        DataType.NULL,
        DataType.DATE,
        DataType.CONST_DATE,
        DataType.GENERICDATETIME,
        DataType.CONST_GENERICDATETIME,
    ):
        return (date_expr,)

    tz_args, _ = make_ch_tz_args(date_ctx)
    if data_type == DataType.DATETIMETZ:
        return (date_expr,) + tz_args
    if data_type in (DataType.DATETIME, DataType.CONST_DATETIME):
        return (date_expr,) + tuple(sa.literal(arg) for arg in tz_args)

    raise ValueError("Unexpected data type in ch_date_with_tz", data_type)

lib/dl_formula/dl_formula/definitions/common_datetime.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa
from sqlalchemy.sql.elements import Null

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import (
    ArgFlagSequence,
    IfFlagDispenser,
)
from dl_formula.definitions.base import (
    MultiVariantTranslation,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import (
    CaseTypeStrategy,
    FromArgs,
    IfTypeStrategy,
)


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class CondBlock(MultiVariantTranslation):
    is_function = False
    scopes = MultiVariantTranslation.scopes & ~Scope.SUGGESTED


class CaseBlock(CondBlock):
    name = "_case_block_"
    arg_names = ["expression", "value_1", "result_1", "value_2", "result_2", "default_result"]

    _null_replace_map = {
        DataType.INTEGER: sa.literal(0),
        DataType.FLOAT: sa.literal(0.0),
        DataType.BOOLEAN: sa.literal(0),
        DataType.STRING: sa.literal(""),
        DataType.MARKUP: sa.literal(""),
        # Note that datetime types are not supported in this case
    }

    @staticmethod
    def _ifnull(value, data_type: DataType) -> bool:
        if value is None or isinstance(value, Null):
            return CaseBlock._null_replace_map.get(data_type, sa.null())

        return value

    @classmethod
    def translation(cls, *args, replace_nulls: bool = False, as_multiif: bool = False):
        value_arg, args = args[0], args[1:]
        else_expr, else_type = None, DataType.NULL
        if len(args) % 2 == 1:
            # block has an ELSE
            else_expr, else_type = args[-1].expression, args[-1].data_type
            args = args[:-1]
        when_args = [(args[ind], args[ind + 1]) for ind in range(0, len(args), 2)]
        # validate value types (must all be castable to same type)
        # all 'WHEN' values must be the same type as 'CASE' value
        DataType.get_common_cast_type(value_arg.data_type, *[when_value.data_type for when_value, _ in when_args])

        return_type = DataType.get_common_cast_type(else_type, *[then_expr.data_type for _, then_expr in when_args])

        def _then_else_wrap(expr):
            if replace_nulls:
                # TODO: make sure it's not needed anymore and remove
                # for ClickHouse only, a workaround for https://github.com/ClickHouse/ClickHouse/issues/7237
                expr = cls._ifnull(expr, data_type=return_type.non_const_type)
            return expr

        if as_multiif and not return_type.is_const:
            # ClickHouse does not support non-constant THENs
            # So use the multiIf function
            return sa.func.multiIf(
                *[
                    # when_1, then_1, when_2, then_2, ...
                    expr
                    for when_value, then_expr in when_args
                    for expr in (value_arg.expression == when_value.expression, _then_else_wrap(then_expr.expression))
                ],
                _then_else_wrap(else_expr if else_expr is not None else sa.null()),
            )

        # Use CASE syntax in SQL
        whens = {when_value.expression: _then_else_wrap(then_expr.expression) for when_value, then_expr in when_args}

        return sa.case(
            value=value_arg.expression,
            whens=whens,
            else_=_then_else_wrap(else_expr if else_expr is not None else sa.null()),
        )

    arg_cnt = None
    variants = [
        VW(D.DUMMY, lambda *args: CaseBlock.translation(*args)),
    ]
    return_type = CaseTypeStrategy()


class IfBlock(CondBlock):
    name = "_if_block_"
    arg_names = ["condition_1", "result_1", "condition_2", "result_2", "default_result"]


class IfBlock3(IfBlock):
    arg_cnt = 3
    variants = [
        V(D.DUMMY, lambda cond, expr, else_expr: sa.case(whens=[(cond, expr)], else_=else_expr)),
    ]
    argument_flags = ArgFlagSequence([ContextFlag.REQ_CONDITION, None, None])
    return_type = FromArgs(1, 2)


class IfBlockMulti(IfBlock):
    arg_cnt = None
    variants = [
        V(
            D.DUMMY | D.SQLITE,
            lambda *args: sa.case(
                whens=[(args[ind], args[ind + 1]) for ind in range(0, len(args) - 1, 2)], else_=args[-1]
            ),
        )
    ]
    argument_flags = IfFlagDispenser()
    return_type = IfTypeStrategy()


DEFINITIONS_COND_BLOCKS = [
    # _case_block_
    CaseBlock,
    # _if_block_
    IfBlock3,
    IfBlockMulti,
]

lib/dl_formula/dl_formula/definitions/conditional_blocks.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import (
    IntEnum,
    unique,
)


ContextFlags = int  # bitwise combination of ContextFlag values


@unique
class ContextFlag(IntEnum):
    # for correct handling of boolean expressions in MS SQL Server
    REQ_CONDITION = 0x0001  # the expression should be translated into a boolean condition
    IS_CONDITION = 0x0002  # the expression is a boolean condition
    DEPRECATED = 0x0004  # function implementation is deprecated

lib/dl_formula/dl_formula/definitions/flags.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import (
    ArgFlagSequence,
    ArgTypeSequence,
)
from dl_formula.definitions.base import (
    Function,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.type_strategy import (
    Fixed,
    FromArgs,
    ParamsFromArgs,
)
from dl_formula.shortcuts import n


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class AggregationFunction(Function):
    is_aggregation = True
    supports_lod = True
    supports_bfb = True
    arg_cnt = 1


class AggSum(AggregationFunction):
    name = "sum"
    variants = [
        V(D.DUMMY | D.SQLITE, sa.func.sum),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = FromArgs()


class AggSumIf(AggregationFunction):
    name = "sum_if"
    variants = [
        VW(
            D.DUMMY,
            lambda expr, cond: n.func.SUM(n.func.IF(cond, expr, 0)),
        ),
    ]
    arg_cnt = 2
    arg_names = ["expression", "condition"]
    # TODO: Check if really all types for conditions are allowed
    # TODO: Check for case: bi-formula-cli translate --dialect ORACLE "SUM_IF([int_value], [int_value])"
    # argument_types = [
    #     *[[DataType.INTEGER, cond_type] for cond_type in DataType],
    #     *[[DataType.FLOAT, cond_type] for cond_type in DataType],
    # ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER, DataType.BOOLEAN]),
        ArgTypeSequence([DataType.FLOAT, DataType.BOOLEAN]),
    ]
    argument_flags = ArgFlagSequence([None, ContextFlag.REQ_CONDITION])
    return_type = FromArgs(0)


class AggAvg(AggregationFunction):
    name = "avg"
    arg_names = ["value"]


class AggAvgFromNumber(AggAvg):
    variants = [
        V(D.DUMMY | D.SQLITE, sa.func.avg),
    ]
    return_type = Fixed(DataType.FLOAT)
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
    ]


class AggAvgFromDate(AggAvg):
    variants = [
        VW(
            D.DUMMY,
            lambda date_val: n.func.DATE(n.func.AVG(n.func.INT(date_val))),
        ),
    ]
    return_type = Fixed(DataType.DATE)
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
    ]


class AggAvgFromDatetime(AggAvg):
    variants = [
        VW(
            D.DUMMY,
            lambda dt_val: n.func.DATETIME(n.func.AVG(n.func.INT(dt_val))),
        ),
    ]
    return_type = FromArgs(0)
    argument_types = [
        ArgTypeSequence([DataType.DATETIME]),
        ArgTypeSequence([DataType.GENERICDATETIME]),
    ]


class AggAvgFromDatetimeTZ(AggAvg):
    variants = [
        VW(
            D.DUMMY,
            lambda dt: n.func.DATETIMETZ(n.func.AVG(n.func.INT(dt)), dt.data_type_params.timezone),
        ),
    ]
    argument_types = [ArgTypeSequence([DataType.DATETIMETZ])]
    return_type = FromArgs(0)
    return_type_params = ParamsFromArgs(0)


class AggAvgIf(AggregationFunction):
    name = "avg_if"
    variants = [
        VW(
            D.DUMMY,
            lambda expr, cond: n.func.AVG(n.func.IF(cond, expr, None)),
        ),
    ]

    arg_cnt = 2
    arg_names = ["expression", "condition"]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER, DataType.BOOLEAN]),
        ArgTypeSequence([DataType.FLOAT, DataType.BOOLEAN]),
    ]
    argument_flags = ArgFlagSequence([None, ContextFlag.REQ_CONDITION])
    return_type = Fixed(DataType.FLOAT)


class AggMax(AggregationFunction):
    name = "max"
    arg_names = ["value"]
    variants = [
        V(D.DUMMY | D.SQLITE, sa.func.max),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
        ArgTypeSequence([DataType.BOOLEAN]),
        ArgTypeSequence([DataType.STRING]),
        ArgTypeSequence([DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME]),
        ArgTypeSequence([DataType.DATETIMETZ]),
        ArgTypeSequence([DataType.GENERICDATETIME]),
        ArgTypeSequence([DataType.UUID]),
        ArgTypeSequence([DataType.INTEGER]),
    ]
    return_type = FromArgs(0)
    return_type_params = ParamsFromArgs(0)


class AggMin(AggregationFunction):
    name = "min"
    arg_names = ["value"]
    variants = [
        V(D.DUMMY | D.SQLITE, sa.func.min),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
        ArgTypeSequence([DataType.BOOLEAN]),
        ArgTypeSequence([DataType.STRING]),
        ArgTypeSequence([DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME]),
        ArgTypeSequence([DataType.DATETIMETZ]),
        ArgTypeSequence([DataType.GENERICDATETIME]),
        ArgTypeSequence([DataType.UUID]),
        ArgTypeSequence([DataType.INTEGER]),
    ]
    return_type = FromArgs(0)
    return_type_params = ParamsFromArgs(0)


class AggCount(AggregationFunction):
    name = "count"
    arg_names = ["value"]
    return_type = Fixed(DataType.INTEGER)


class AggCount0(AggCount):
    arg_cnt = 0
    variants = [
        V(D.DUMMY | D.SQLITE, lambda: sa.func.count(1)),
    ]


class AggCount1(AggCount):
    variants = [
        V(D.DUMMY | D.SQLITE, sa.func.count),
    ]


class AggCountIf(AggregationFunction):
    name = "count_if"
    variants = [
        VW(
            D.DUMMY,
            lambda cond: n.func.COUNT(n.func.IF(cond, 1, None)),
        )
    ]

    arg_names = ["condition"]
    argument_types = [
        ArgTypeSequence([DataType.BOOLEAN]),
    ]
    argument_flags = ArgFlagSequence([ContextFlag.REQ_CONDITION])
    return_type = Fixed(DataType.INTEGER)


class AggCountd(AggregationFunction):
    name = "countd"
    variants = [
        V(D.DUMMY | D.SQLITE, lambda x: sa.func.count(sa.distinct(x))),
    ]
    return_type = Fixed(DataType.INTEGER)


class AggCountdIf(AggregationFunction):
    name = "countd_if"
    variants = [
        VW(
            D.DUMMY,
            lambda expr, cond: n.func.COUNTD(n.func.IF(cond, expr, None)),
        ),
    ]

    arg_cnt = 2
    arg_names = ["expression", "condition"]
    argument_types = [ArgTypeSequence([expr_type, DataType.BOOLEAN]) for expr_type in DataType]
    argument_flags = ArgFlagSequence([None, ContextFlag.REQ_CONDITION])
    return_type = Fixed(DataType.INTEGER)


class AggCountdApprox(AggregationFunction):
    name = "countd_approx"
    return_type = Fixed(DataType.INTEGER)


class AggStdev(AggregationFunction):
    name = "stdev"
    variants = [
        V(D.DUMMY, sa.func.STDDEV_SAMP),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class AggStdevp(AggregationFunction):
    name = "stdevp"
    variants = [
        V(D.DUMMY, sa.func.STDDEV_POP),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class AggVar(AggregationFunction):
    name = "var"
    variants = [
        V(D.DUMMY, sa.func.VAR_SAMP),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class AggVarp(AggregationFunction):
    name = "varp"
    variants = [
        V(D.DUMMY, sa.func.VAR_POP),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class AggQuantileBase(AggregationFunction):
    arg_cnt = 2
    arg_names = ["value", "quant"]
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.INTEGER, DataType.FLOAT, DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME},
                DataType.FLOAT,
            ]
        ),
    ]
    return_type = FromArgs(0)


class AggQuantile(AggQuantileBase):
    name = "quantile"


class AggQuantileApproximate(AggQuantileBase):
    name = "quantile_approx"


class AggMedian(AggregationFunction):
    name = "median"
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.INTEGER, DataType.FLOAT, DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME},
            ]
        ),
    ]
    return_type = FromArgs()


class AggAttr(AggregationFunction):
    name = "attr"
    # TODO: Implement me


class AggAny(AggregationFunction):
    name = "any"
    return_type = FromArgs()


class AggArgMinMax(AggregationFunction):
    return_type = FromArgs(0)


class AggArgMin(AggArgMinMax):
    name = "arg_min"
    arg_cnt = 2
    arg_names = ["value", "comp"]


class AggArgMax(AggArgMinMax):
    name = "arg_max"
    arg_cnt = 2
    arg_names = ["value", "comp"]


# Blacklist of types,
# because most types should get stringified and concatenated successfully,
# but markup, while internally a string, shouldn't become a string
# (and requires a bit more specific concatenation).
# TODO: support markup agg-concat.
# `DataType.UNSUPPORTED` could also work, but the idea is that the user should
# add `STR()` explicitly for those.
AGG_CONCAT_INPUT_TYPES = set(DataType) - set((DataType.MARKUP, DataType.CONST_MARKUP, DataType.UNSUPPORTED))


class AggAllConcatBase(AggregationFunction):
    name = "all_concat"
    arg_names = ["expression", "separator"]
    # (Any, ConstString) -> String
    argument_types = [
        ArgTypeSequence(
            [
                AGG_CONCAT_INPUT_TYPES,
                DataType.CONST_STRING,
            ]
        ),
    ]
    return_type = Fixed(DataType.STRING)


class AggAllConcat1(AggAllConcatBase):
    arg_cnt = 1


class AggAllConcat2(AggAllConcatBase):
    arg_cnt = 2


class AggTopConcatBase(AggregationFunction):
    name = "top_concat"
    # Doable for postgresql, but not performant, and very convoluted
    arg_names = ["expression", "amount", "separator"]
    # (Any, ConstInteger, ConstString) -> String
    argument_types = [
        ArgTypeSequence(
            [
                AGG_CONCAT_INPUT_TYPES,
                DataType.CONST_INTEGER,
                DataType.CONST_STRING,
            ]
        ),
    ]
    return_type = Fixed(DataType.STRING)


class AggTopConcat1(AggTopConcatBase):
    arg_cnt = 2


class AggTopConcat2(AggTopConcatBase):
    arg_cnt = 3


DEFINITIONS_AGGREGATION = [
    # all_concat
    AggAllConcat1,
    AggAllConcat2,
    # any
    AggAny,
    # arg_max
    AggArgMax,
    # arg_min
    AggArgMin,
    # avg
    AggAvgFromNumber,
    AggAvgFromDate,
    AggAvgFromDatetime,
    AggAvgFromDatetimeTZ,
    # avg_if
    AggAvgIf,
    # count
    AggCount0,
    AggCount1,
    # count_if
    AggCountIf,
    # countd
    AggCountd,
    # countd_approx
    AggCountdApprox,
    # countd_if
    AggCountdIf,
    # max
    AggMax,
    # median
    AggMedian,
    # min
    AggMin,
    # quantile
    AggQuantile,
    # quantile_approx
    AggQuantileApproximate,
    # stdev
    AggStdev,
    # stdevp
    AggStdevp,
    # sum
    AggSum,
    # sum_if
    AggSumIf,
    # top_concat
    AggTopConcat1,
    AggTopConcat2,
    # var
    AggVar,
    # varp
    AggVarp,
]

lib/dl_formula/dl_formula/definitions/functions_aggregation.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    List,
    Type,
)

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
import dl_formula.core.nodes as nodes
from dl_formula.definitions.args import (
    ArgTypeForAll,
    ArgTypeSequence,
)
from dl_formula.definitions.base import (
    Function,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.functions_string import (
    FuncContains,
    FuncLen,
    FuncNotContains,
    FuncStartswith,
)
from dl_formula.definitions.literals import un_literal
from dl_formula.definitions.type_strategy import (
    Fixed,
    FromArgs,
)
from dl_formula.shortcuts import n
from dl_formula.translation.context import TranslationCtx


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class ArrayFunction(Function):
    pass


class FuncCreateArray(ArrayFunction):
    name = "array"
    arg_names = ["value_1", "value_2", "value_3"]
    arg_cnt = None


def _normalize_arr_items(val: List[TranslationCtx], item_type: Type) -> list:
    """Unliteralize and convert array items"""
    result = []
    for item in val:
        un_lit_item = un_literal(item.expression)
        if un_lit_item is not None:
            un_lit_item = item_type(un_lit_item)
        result.append(un_lit_item)
    return result


class FuncConstArrayBase(FuncCreateArray):
    pass


class FuncConstArrayFloat(FuncConstArrayBase):
    variants = [
        VW(D.DUMMY, lambda *args: nodes.LiteralArrayFloat.make(_normalize_arr_items(args, float))),
    ]
    argument_types = [
        ArgTypeForAll({DataType.CONST_FLOAT, DataType.NULL}, require_type_match=DataType.CONST_FLOAT),
    ]
    return_type = Fixed(DataType.CONST_ARRAY_FLOAT)


class FuncConstArrayInt(FuncConstArrayBase):
    variants = [
        VW(D.DUMMY, lambda *args: nodes.LiteralArrayInteger.make(_normalize_arr_items(args, int))),
    ]
    argument_types = [
        ArgTypeForAll({DataType.CONST_INTEGER, DataType.NULL}, require_type_match=DataType.CONST_INTEGER),
    ]
    return_type = Fixed(DataType.CONST_ARRAY_INT)


class FuncConstArrayStr(FuncConstArrayBase):
    variants = [
        VW(D.DUMMY, lambda *args: nodes.LiteralArrayString.make(_normalize_arr_items(args, str))),
    ]
    argument_types = [
        ArgTypeForAll({DataType.CONST_STRING, DataType.NULL}, require_type_match=DataType.CONST_STRING),
    ]
    return_type = Fixed(DataType.CONST_ARRAY_STR)


class FuncNonConstArrayBase(FuncCreateArray):
    pass


class FuncNonConstArrayInt(FuncNonConstArrayBase):
    argument_types = [
        ArgTypeForAll(DataType.INTEGER, require_type_match={DataType.INTEGER, DataType.CONST_INTEGER}),
    ]
    return_type = Fixed(DataType.ARRAY_INT)


class FuncNonConstArrayFloat(FuncNonConstArrayBase):
    argument_types = [
        ArgTypeForAll(DataType.FLOAT, require_type_match={DataType.FLOAT, DataType.CONST_FLOAT}),
    ]
    return_type = Fixed(DataType.ARRAY_FLOAT)


class FuncNonConstArrayStr(FuncNonConstArrayBase):
    argument_types = [
        ArgTypeForAll(DataType.STRING, require_type_match={DataType.STRING, DataType.CONST_STRING}),
    ]
    return_type = Fixed(DataType.ARRAY_STR)


class FuncUnnestArray(ArrayFunction):
    name = "unnest"
    arg_names = ["array"]
    arg_cnt = 1


class FuncUnnestArrayFloat(FuncUnnestArray):
    argument_types = [ArgTypeSequence([DataType.ARRAY_FLOAT])]
    return_type = Fixed(DataType.FLOAT)


class FuncUnnestArrayInt(FuncUnnestArray):
    argument_types = [ArgTypeSequence([DataType.ARRAY_INT])]
    return_type = Fixed(DataType.INTEGER)


class FuncUnnestArrayStr(FuncUnnestArray):
    argument_types = [ArgTypeSequence([DataType.ARRAY_STR])]
    return_type = Fixed(DataType.STRING)


class FuncLenArray(FuncLen):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT]),
        ArgTypeSequence([DataType.ARRAY_INT]),
        ArgTypeSequence([DataType.ARRAY_STR]),
    ]


class FuncGetArrayItem(ArrayFunction):
    name = "get_item"
    arg_names = ["array", "index"]
    arg_cnt = 2


class FuncGetArrayItemFloat(FuncGetArrayItem):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncGetArrayItemInt(FuncGetArrayItem):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.INTEGER)


class FuncGetArrayItemStr(FuncGetArrayItem):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_STR, DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncArrStr(ArrayFunction):
    name = "arr_str"
    arg_names = ["array", "delimiter", "null_str"]
    return_type = Fixed(DataType.STRING)


class FuncArrStr1(FuncArrStr):
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT]),
        ArgTypeSequence([DataType.ARRAY_INT]),
        ArgTypeSequence([DataType.ARRAY_STR]),
    ]


class FuncArrStr2(FuncArrStr):
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.CONST_STRING]),
        ArgTypeSequence([DataType.ARRAY_INT, DataType.CONST_STRING]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.CONST_STRING]),
    ]


class FuncArrStr3(FuncArrStr):
    arg_cnt = 3
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.CONST_STRING, DataType.CONST_STRING]),
        ArgTypeSequence([DataType.ARRAY_INT, DataType.CONST_STRING, DataType.CONST_STRING]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.CONST_STRING, DataType.CONST_STRING]),
    ]


class FuncArrayCountItem(ArrayFunction):
    name = "count_item"
    arg_names = ["array", "value"]
    arg_cnt = 2

    return_type = Fixed(DataType.INTEGER)


class FuncArrayCountItemInt(FuncArrayCountItem):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.INTEGER]),
    ]


class FuncArrayCountItemFloat(FuncArrayCountItem):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.FLOAT]),
    ]


class FuncArrayCountItemStr(FuncArrayCountItem):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_STR, DataType.STRING]),
    ]


class FuncArrayContains(FuncContains):
    arg_names = ["array", "value"]
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.INTEGER]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.STRING]),
    ]


class FuncArrayNotContains(FuncNotContains):
    arg_names = ["array", "value"]
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.INTEGER]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.STRING]),
    ]
    variants = [
        VW(
            D.DUMMY,
            lambda arr, val: n.not_(n.func.CONTAINS(arr, val)),
        ),
    ]


class FuncArrayContainsAll(ArrayFunction):
    name = "contains_all"
    arg_names = ["array_1", "array_2"]
    arg_cnt = 2
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION

    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.ARRAY_INT]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.ARRAY_FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.ARRAY_STR]),
    ]


class FuncArrayContainsAny(ArrayFunction):
    name = "contains_any"
    arg_names = ["array_1", "array_2"]
    arg_cnt = 2
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION

    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.ARRAY_INT]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.ARRAY_FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.ARRAY_STR]),
    ]


class FuncArrayContainsSubsequence(ArrayFunction):
    name = "contains_subsequence"
    arg_names = ["array_1", "array_2"]
    arg_cnt = 2
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION

    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.ARRAY_INT]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.ARRAY_FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.ARRAY_STR]),
    ]


class FuncStartswithArrayBase(FuncStartswith):
    arg_names = ["array_1", "array_2"]


class FuncStartswithArrayConst(FuncStartswithArrayBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.CONST_ARRAY_INT]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.CONST_ARRAY_FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.CONST_ARRAY_STR]),
    ]


class FuncStartswithArrayNonConst(FuncStartswithArrayBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.ARRAY_INT]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.ARRAY_FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.ARRAY_STR]),
    ]


class FuncArraySlice(ArrayFunction):
    name = "slice"
    arg_names = ["array", "offset", "length"]
    arg_cnt = 3
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.ARRAY_INT, DataType.ARRAY_FLOAT, DataType.ARRAY_STR},
                DataType.CONST_INTEGER,
                DataType.CONST_INTEGER,
            ]
        ),
    ]

    return_type = FromArgs(0)


class FuncReplaceArrayBase(ArrayFunction):
    name = "replace"
    arg_cnt = 3
    arg_names = ["array", "old", "new"]

    return_type = FromArgs(0)


class FuncReplaceArrayLiteralNull(FuncReplaceArrayBase):
    """
    This implementation is needed because the default one would return an array
    of Nullables making it impossible to use array aggregation functions
    """

    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.NULL, DataType.INTEGER]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.NULL, DataType.FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.NULL, DataType.STRING]),
    ]


class FuncReplaceArrayDefault(FuncReplaceArrayBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.INTEGER, DataType.INTEGER]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.STRING, DataType.STRING]),
    ]


class FuncArrayCast(ArrayFunction):
    arg_cnt = 1
    arg_names = ["array"]


class FuncIntArray(FuncArrayCast):
    name = "cast_arr_int"
    return_type = Fixed(DataType.ARRAY_INT)


class FuncIntArrayFromIntArray(FuncIntArray):
    variants = [
        V(D.DUMMY, lambda arr: arr),
    ]
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT]),
    ]


class FuncIntArrayFromFloatArray(FuncIntArray):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT]),
    ]


class FuncIntArrayFromStringArray(FuncIntArray):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_STR]),
    ]


class FuncFloatArray(FuncArrayCast):
    name = "cast_arr_float"
    return_type = Fixed(DataType.ARRAY_FLOAT)


class FuncFloatArrayFromIntArray(FuncFloatArray):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT]),
    ]


class FuncFloatArrayFromFloatArray(FuncFloatArray):
    variants = [
        V(D.DUMMY, lambda arr: arr),
    ]
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT]),
    ]


class FuncFloatArrayFromStringArray(FuncFloatArray):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_STR]),
    ]


class FuncStrArray(FuncArrayCast):
    name = "cast_arr_str"
    return_type = Fixed(DataType.ARRAY_STR)


class FuncStringArrayFromIntArray(FuncStrArray):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT]),
    ]


class FuncStringArrayFromFloatArray(FuncStrArray):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT]),
    ]


class FuncStringArrayFromStringArray(FuncStrArray):
    variants = [
        V(D.DUMMY, lambda arr: arr),
    ]
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_STR]),
    ]


class FuncArrayMinBase(ArrayFunction):
    name = "arr_min"
    arg_names = ["array"]
    arg_cnt = 1


class FuncArrayMinFloat(FuncArrayMinBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncArrayMinInt(FuncArrayMinBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT]),
    ]
    return_type = Fixed(DataType.INTEGER)


class FuncArrayMaxBase(ArrayFunction):
    name = "arr_max"
    arg_names = ["array"]
    arg_cnt = 1


class FuncArrayMaxFloat(FuncArrayMaxBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncArrayMaxInt(FuncArrayMaxBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT]),
    ]
    return_type = Fixed(DataType.INTEGER)


class FuncArraySumBase(ArrayFunction):
    name = "arr_sum"
    arg_names = ["array"]
    arg_cnt = 1


class FuncArraySumFloat(FuncArraySumBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncArraySumInt(FuncArraySumBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT]),
    ]
    return_type = Fixed(DataType.INTEGER)


class FuncArrayAvg(ArrayFunction):
    name = "arr_avg"
    arg_names = ["array"]
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([{DataType.ARRAY_INT, DataType.ARRAY_FLOAT}]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncArrayProduct(ArrayFunction):
    name = "arr_product"
    arg_names = ["array"]
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([{DataType.ARRAY_INT, DataType.ARRAY_FLOAT}]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncArrayRemoveBase(ArrayFunction):
    name = "arr_remove"
    arg_names = ["array", "value"]
    arg_cnt = 2

    return_type = FromArgs(0)


class FuncArrayRemoveLiteralNull(FuncArrayRemoveBase):
    """
    This implementation is needed because the default one would return an array
    of Nullables making it impossible to use array aggregation functions
    """

    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.NULL]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.NULL]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.NULL]),
    ]


class FuncArrayRemoveDefault(FuncArrayRemoveBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.INTEGER]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.STRING]),
    ]


DEFINITIONS_ARRAY = [
    # arr_avg
    FuncArrayAvg,
    # arr_max
    FuncArrayMaxFloat,
    FuncArrayMaxInt,
    # arr_min
    FuncArrayMinFloat,
    FuncArrayMinInt,
    # arr_product
    FuncArrayProduct,
    # arr_remove
    FuncArrayRemoveLiteralNull,
    FuncArrayRemoveDefault,
    # arr_str
    FuncArrStr1,
    FuncArrStr2,
    FuncArrStr3,
    # arr_sum
    FuncArraySumFloat,
    FuncArraySumInt,
    # array
    FuncConstArrayFloat,
    FuncConstArrayInt,
    FuncConstArrayStr,
    FuncNonConstArrayInt,
    FuncNonConstArrayFloat,
    FuncNonConstArrayStr,
    # cast_arr_float
    FuncFloatArrayFromIntArray,
    FuncFloatArrayFromFloatArray,
    FuncFloatArrayFromStringArray,
    # cast_arr_int
    FuncIntArrayFromIntArray,
    FuncIntArrayFromFloatArray,
    FuncIntArrayFromStringArray,
    # cast_arr_str
    FuncStringArrayFromIntArray,
    FuncStringArrayFromFloatArray,
    FuncStringArrayFromStringArray,
    # contains
    FuncArrayContains,
    # notcontains
    FuncArrayNotContains,
    # contains_all
    FuncArrayContainsAll,
    # contains_any
    FuncArrayContainsAny,
    # contains_subsequence
    FuncArrayContainsSubsequence,
    # count_item
    FuncArrayCountItemInt,
    FuncArrayCountItemFloat,
    FuncArrayCountItemStr,
    # get_item
    FuncGetArrayItemFloat,
    FuncGetArrayItemInt,
    FuncGetArrayItemStr,
    # len
    FuncLenArray,
    # replace
    FuncReplaceArrayLiteralNull,
    FuncReplaceArrayDefault,
    # slice
    FuncArraySlice,
    # startswith
    FuncStartswithArrayConst,
    FuncStartswithArrayNonConst,
    # unnest
    FuncUnnestArrayFloat,
    FuncUnnestArrayInt,
    FuncUnnestArrayStr,
]

lib/dl_formula/dl_formula/definitions/functions_array.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
import dl_formula.core.exc as exc
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    Function,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.common_datetime import (
    ensure_naive_first_arg,
    normalize_and_validate_datetime_interval_type,
)
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.literals import un_literal
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import (
    Fixed,
    FromArgs,
    ParamsFromArgs,
)
from dl_formula.shortcuts import n


V = TranslationVariant.make
VW = TranslationVariantWrapped.make

_DATETIME_TYPES = {DataType.DATETIME, DataType.DATETIMETZ, DataType.GENERICDATETIME}


class DateFunction(Function):
    pass


class FuncDateaddBase(DateFunction):
    name = "dateadd"
    arg_names = ["datetime", "unit", "number"]
    return_type = FromArgs(0)
    return_type_params = ParamsFromArgs(0)


class FuncDateadd1Base(FuncDateaddBase):
    arg_cnt = 1


class FuncDateadd1(FuncDateadd1Base):
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda date: n.func.DATEADD(date, "day", 1)),
    ]
    argument_types = [
        ArgTypeSequence([{DataType.DATE, *_DATETIME_TYPES}]),
    ]


class FuncDateadd2Base(FuncDateaddBase):
    arg_cnt = 2


class FuncDateadd2Unit(FuncDateadd2Base):
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda date, unit: n.func.DATEADD(date, unit, 1)),
    ]
    argument_types = [
        ArgTypeSequence([{DataType.DATE, *_DATETIME_TYPES}, DataType.CONST_STRING]),
    ]


class FuncDateadd2Number(FuncDateadd2Base):
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda date, number: n.func.DATEADD(date, "day", number)),
    ]
    argument_types = [
        ArgTypeSequence([{DataType.DATE, *_DATETIME_TYPES}, DataType.CONST_INTEGER]),
    ]


class FuncDateadd3Base(FuncDateaddBase):
    arg_cnt = 3


class FuncDateadd3Legacy(FuncDateadd3Base):
    """Legacy version of the function with `date` as the last argument"""

    variants = [VW(D.DUMMY | D.SQLITE, lambda what, num, date: n.func.DATEADD(date, what, num))]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.INTEGER, DataType.DATE]),
        ArgTypeSequence([DataType.STRING, DataType.INTEGER, DataType.DATETIME]),
        ArgTypeSequence([DataType.STRING, DataType.INTEGER, DataType.GENERICDATETIME]),
    ]
    return_type = FromArgs(2)
    return_flags = ContextFlag.DEPRECATED


class FuncDateadd3DateConstNum(FuncDateadd3Base):
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.CONST_STRING, DataType.CONST_INTEGER]),
    ]


class FuncDateadd3DateNonConstNum(FuncDateadd3Base):
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.CONST_STRING, DataType.INTEGER]),
    ]


class FuncDateadd3DatetimeConstNum(FuncDateadd3Base):
    argument_types = [
        ArgTypeSequence([_DATETIME_TYPES, DataType.CONST_STRING, DataType.CONST_INTEGER]),
    ]


class FuncDateadd3DatetimeNonConstNum(FuncDateadd3Base):
    argument_types = [
        ArgTypeSequence([DataType.DATETIME, DataType.CONST_STRING, DataType.INTEGER]),
    ]


class FuncDateadd3GenericDatetimeNonConstNum(FuncDateadd3Base):
    argument_types = [
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.CONST_STRING, DataType.INTEGER]),
    ]


class FuncDateadd3DatetimeTZNonConstNum(FuncDateadd3Base):
    argument_types = [
        ArgTypeSequence([DataType.DATETIMETZ, DataType.CONST_STRING, DataType.INTEGER]),
    ]


class FuncDatepart(DateFunction):
    name = "datepart"
    arg_names = ["datetime", "unit", "firstday"]
    return_type = Fixed(DataType.INTEGER)


class FuncDatepart2Legacy(FuncDatepart):
    """Legacy version of the function with `date` as the last argument"""

    arg_cnt = 2
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda what, date: n.func.DATEPART(date, what)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.DATE]),
        ArgTypeSequence([DataType.STRING, DataType.DATETIME]),
        ArgTypeSequence([DataType.STRING, DataType.GENERICDATETIME]),
    ]
    return_flags = ContextFlag.DEPRECATED


class FuncDatepart2(FuncDatepart):
    arg_cnt = 2
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda date, part: n.func.DATEPART(date, part, "mon")),
    ]
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME, DataType.DATETIMETZ},
                DataType.STRING,
            ]
        ),
    ]


@ensure_naive_first_arg
def _datepart_const_impl(date_ctx, part_ctx, firstday_ctx):
    """`DATEPART` for `(date|datetime, str, str)`"""
    part: str = un_literal(part_ctx.expression)

    if part == "second":
        return n.func.SECOND(date_ctx)
    if part == "minute":
        return n.func.MINUTE(date_ctx)
    if part == "hour":
        return n.func.HOUR(date_ctx)
    if part == "day":
        return n.func.DAY(date_ctx)
    if part == "month":
        return n.func.MONTH(date_ctx)
    if part == "quarter":
        return n.func.QUARTER(date_ctx)
    if part == "year":
        return n.func.YEAR(date_ctx)
    if part in ("dow", "dayofweek"):
        firstday: str = un_literal(firstday_ctx.expression)
        return n.func.DAYOFWEEK(date_ctx, firstday)
    if part == "week":
        return n.func.WEEK(date_ctx)

    # TODO?: raise an exception to user?
    return n.null()


class FuncDatepart3Const(FuncDatepart):
    arg_cnt = 3
    variants = [
        VW(D.DUMMY | D.SQLITE, _datepart_const_impl),
    ]
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME, DataType.DATETIMETZ},
                DataType.CONST_STRING,
                DataType.CONST_STRING,
            ]
        ),
    ]


@ensure_naive_first_arg
def _datepart_nonconst_impl(date_ctx, part_ctx, firstday_ctx):
    firstday: str = un_literal(firstday_ctx.expression)
    return (
        n.case(part_ctx)
        .whens(
            n.when("second").then(n.func.SECOND(date_ctx)),
            n.when("minute").then(n.func.MINUTE(date_ctx)),
            n.when("hour").then(n.func.HOUR(date_ctx)),
            n.when("day").then(n.func.DAY(date_ctx)),
            n.when("month").then(n.func.MONTH(date_ctx)),
            n.when("quarter").then(n.func.QUARTER(date_ctx)),
            n.when("year").then(n.func.YEAR(date_ctx)),
            n.when("dow").then(n.func.DAYOFWEEK(date_ctx, firstday)),
            n.when("dayofweek").then(n.func.DAYOFWEEK(date_ctx, firstday)),
            n.when("week").then(n.func.WEEK(date_ctx)),
            n.when("second").then(n.func.SECOND(date_ctx)),
        )
        .else_(None)
    )


class FuncDatepart3NonConst(FuncDatepart):
    arg_cnt = 3
    variants = [
        VW(D.DUMMY | D.SQLITE, _datepart_nonconst_impl),
    ]
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME, DataType.DATETIMETZ},
                DataType.STRING,
                DataType.CONST_STRING,
            ]
        ),
    ]


def norm_datetrunc_unit(unit):
    return normalize_and_validate_datetime_interval_type(un_literal(unit))


class FuncDatetrunc(DateFunction):
    name = "datetrunc"
    arg_names = ["datetime", "unit", "number"]
    return_type = FromArgs(0)
    return_type_params = ParamsFromArgs(0)


class FuncDatetrunc2(FuncDatetrunc):
    arg_cnt = 2


class FuncDatetrunc2Date(FuncDatetrunc2):
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.CONST_STRING]),
    ]


class FuncDatetrunc2Datetime(FuncDatetrunc2):
    argument_types = [
        ArgTypeSequence([{DataType.DATETIME, DataType.GENERICDATETIME}, DataType.CONST_STRING]),
    ]


class FuncDatetrunc2DatetimeTZ(FuncDatetrunc2Datetime):
    argument_types = [ArgTypeSequence([DataType.DATETIMETZ, DataType.CONST_STRING])]


class FuncDatetrunc3(FuncDatetrunc):
    arg_cnt = 3


class FuncDatetrunc3Date(FuncDatetrunc3):
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.CONST_STRING, DataType.CONST_INTEGER]),
    ]


class FuncDatetrunc3Datetime(FuncDatetrunc3):
    argument_types = [
        ArgTypeSequence(
            [
                {DataType.DATETIME, DataType.GENERICDATETIME, DataType.DATETIMETZ},
                DataType.CONST_STRING,
                DataType.CONST_INTEGER,
            ]
        ),
    ]


class SpecificDatepartFunc(DateFunction):
    arg_cnt = 1
    _zero_variants = [V(D.DUMMY | D.SQLITE, lambda x: sa.literal(0))]
    return_type = Fixed(DataType.INTEGER)


class FuncSecond(SpecificDatepartFunc):
    name = "second"
    arg_names = ["datetime"]


class FuncSecondDate(FuncSecond):
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
    ]
    variants = SpecificDatepartFunc._zero_variants


class FuncSecondDatetime(FuncSecond):
    arg_cnt = 1
    argument_types = [ArgTypeSequence([{DataType.DATETIME, DataType.GENERICDATETIME}])]
    return_type = Fixed(DataType.INTEGER)


class FuncSecondDatetimeTZ(FuncSecondDatetime):
    argument_types = [ArgTypeSequence([DataType.DATETIMETZ])]
    return_type = Fixed(DataType.INTEGER)


class FuncMinute(SpecificDatepartFunc):
    name = "minute"
    arg_names = ["datetime"]


class FuncMinuteDate(FuncMinute):
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
    ]
    variants = SpecificDatepartFunc._zero_variants


class FuncMinuteDatetime(FuncMinute):
    arg_cnt = 1
    argument_types = [ArgTypeSequence([{DataType.DATETIME, DataType.GENERICDATETIME}])]
    return_type = Fixed(DataType.INTEGER)


class FuncMinuteDatetimeTZ(FuncMinuteDatetime):
    argument_types = [ArgTypeSequence([DataType.DATETIMETZ])]
    return_type = Fixed(DataType.INTEGER)


class FuncHour(SpecificDatepartFunc):
    name = "hour"
    arg_names = ["datetime"]


class FuncHourDate(FuncHour):
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
    ]
    variants = SpecificDatepartFunc._zero_variants


class FuncHourDatetime(FuncHour):
    arg_cnt = 1
    argument_types = [ArgTypeSequence([{DataType.DATETIME, DataType.GENERICDATETIME}])]
    return_type = Fixed(DataType.INTEGER)


class FuncHourDatetimeTZ(FuncHourDatetime):
    argument_types = [ArgTypeSequence([DataType.DATETIMETZ])]
    return_type = Fixed(DataType.INTEGER)


class FuncDay(SpecificDatepartFunc):
    name = "day"
    arg_cnt = 1
    arg_names = ["datetime"]
    argument_types = [ArgTypeSequence([{DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME}])]
    return_type = Fixed(DataType.INTEGER)


class FuncDayDatetimeTZ(FuncDay):
    argument_types = [ArgTypeSequence([DataType.DATETIMETZ])]
    return_type = Fixed(DataType.INTEGER)


class FuncMonth(SpecificDatepartFunc):
    name = "month"
    arg_names = ["datetime"]
    arg_cnt = 1
    argument_types = [ArgTypeSequence([{DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME}])]
    return_type = Fixed(DataType.INTEGER)


class FuncMonthDatetimeTZ(FuncMonth):
    argument_types = [ArgTypeSequence([DataType.DATETIMETZ])]
    return_type = Fixed(DataType.INTEGER)


class FuncQuarter(SpecificDatepartFunc):
    name = "quarter"
    arg_names = ["datetime"]
    arg_cnt = 1
    argument_types = [ArgTypeSequence([{DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME}])]
    return_type = Fixed(DataType.INTEGER)


class FuncQuarterDatetimeTZ(FuncQuarter):
    argument_types = [ArgTypeSequence([DataType.DATETIMETZ])]
    return_type = Fixed(DataType.INTEGER)


class FuncYear(SpecificDatepartFunc):
    name = "year"
    arg_names = ["datetime"]
    arg_cnt = 1
    argument_types = [ArgTypeSequence([{DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME}])]
    return_type = Fixed(DataType.INTEGER)


class FuncYearDatetimeTZ(FuncYear):
    argument_types = [ArgTypeSequence([DataType.DATETIMETZ])]
    return_type = Fixed(DataType.INTEGER)


class FuncDayofweek(SpecificDatepartFunc):
    name = "dayofweek"
    arg_names = ["datetime", "firstday"]
    return_type = Fixed(DataType.INTEGER)

    _fd_map = {
        "mon": 1,
        "monday": 1,
        "tue": 2,
        "tuesday": 2,
        "wed": 3,
        "wednesday": 3,
        "thu": 4,
        "thursday": 4,
        "fri": 5,
        "friday": 5,
        "sat": 6,
        "saturday": 6,
        "sun": 7,
        "sunday": 7,
    }

    @classmethod
    def _norm_fd(cls, firstday):
        """Normalize first-day-of-the-week value (convert to int (1-7))"""
        try:
            return cls._fd_map[un_literal(firstday).lower()]
        except KeyError:
            raise exc.TranslationError("Unknown time unit")


norm_fd = FuncDayofweek._norm_fd


class FuncDayofweek1(FuncDayofweek):
    arg_cnt = 1
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda date: n.func.DAYOFWEEK(date.expression, "mon")),
    ]
    argument_types = [
        ArgTypeSequence([{DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME, DataType.DATETIMETZ}]),
    ]


def dow_firstday_shift(base_1_dow, firstday_expr):
    """
    `base_1_dow` should be `1` for monday, ..., `6` for saturday, `7` for sunday
    """
    firstday = norm_fd(firstday_expr)
    if firstday == 1:
        return base_1_dow
    return (base_1_dow - firstday + 7) % 7 + 1


class FuncDayofweek2(FuncDayofweek):
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([{DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME}, DataType.CONST_STRING]),
    ]


class FuncDayofweek2TZ(FuncDayofweek):
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([DataType.DATETIMETZ, DataType.CONST_STRING]),
    ]


class FuncWeek(SpecificDatepartFunc):
    name = "week"
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([{DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME}]),
    ]
    return_type = Fixed(DataType.INTEGER)


class FuncWeekTZ(FuncWeek):
    argument_types = [ArgTypeSequence([DataType.DATETIMETZ])]


class FuncTypeGenericDatetimeNowImpl(DateFunction):
    arg_cnt = 0
    return_type = Fixed(DataType.GENERICDATETIME)


class FuncNow(FuncTypeGenericDatetimeNowImpl):
    name = "now"


class FuncGenericNow(FuncTypeGenericDatetimeNowImpl):
    name = "genericnow"
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncToday(DateFunction):
    name = "today"
    arg_cnt = 0
    return_type = Fixed(DataType.DATE)


# TODO: UTCTODAY() through DataType.DATETIMETZ / DATE(UTCNOW()), and remove the 'UTC' cast from TODAY().


DEFINITIONS_DATETIME = [
    # dateadd
    FuncDateadd1,
    FuncDateadd2Unit,
    FuncDateadd2Number,
    FuncDateadd3Legacy,
    FuncDateadd3DateConstNum,
    FuncDateadd3DateNonConstNum,
    FuncDateadd3DatetimeConstNum,
    FuncDateadd3DatetimeNonConstNum,
    FuncDateadd3GenericDatetimeNonConstNum,
    FuncDateadd3DatetimeTZNonConstNum,
    # datepart
    FuncDatepart2Legacy,
    FuncDatepart2,
    FuncDatepart3Const,
    FuncDatepart3NonConst,
    # datetrunc
    FuncDatetrunc2Date,
    FuncDatetrunc2Datetime,
    FuncDatetrunc2DatetimeTZ,
    FuncDatetrunc3Date,
    FuncDatetrunc3Datetime,
    # day
    FuncDay,
    FuncDayDatetimeTZ,
    # dayofweek
    FuncDayofweek1,
    FuncDayofweek2,
    FuncDayofweek2TZ,
    # genericnow
    FuncGenericNow,
    # hour
    FuncHourDate,
    FuncHourDatetime,
    FuncHourDatetimeTZ,
    # minute
    FuncMinuteDate,
    FuncMinuteDatetime,
    FuncMinuteDatetimeTZ,
    # month
    FuncMonth,
    FuncMonthDatetimeTZ,
    # now
    FuncNow,
    # quarter
    FuncQuarter,
    FuncQuarterDatetimeTZ,
    # second
    FuncSecondDate,
    FuncSecondDatetime,
    FuncSecondDatetimeTZ,
    # today
    FuncToday,
    # week
    FuncWeek,
    FuncWeekTZ,
    # year
    FuncYear,
    FuncYearDatetimeTZ,
]

lib/dl_formula/dl_formula/definitions/functions_datetime.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    Function,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import (
    CaseTypeStrategy,
    Fixed,
    FromArgs,
    IfTypeStrategy,
)
from dl_formula.shortcuts import n


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class LogicalFunction(Function):
    pass


class FuncIsnull(LogicalFunction):
    name = "isnull"
    arg_cnt = 1
    arg_names = ["expression"]
    variants = [
        V(D.DUMMY, lambda x: x.is_(None)),
    ]
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


class FuncIfnull(LogicalFunction):
    name = "ifnull"
    arg_cnt = 2
    arg_names = ["check_value", "alt_value"]
    return_type = FromArgs()


class FuncIsnan(LogicalFunction):
    name = "isnan"
    scopes = LogicalFunction.scopes & ~Scope.DOCUMENTED  # FIXME: add to doc
    arg_cnt = 1
    arg_names = ["expression"]
    return_type = Fixed(DataType.BOOLEAN)


class FuncIfnan(LogicalFunction):
    name = "ifnan"
    scopes = LogicalFunction.scopes & ~Scope.DOCUMENTED  # FIXME: add to doc
    arg_cnt = 2
    arg_names = ["check_value", "alt_value"]
    return_type = FromArgs()


class FuncZn(LogicalFunction):
    name = "zn"
    arg_cnt = 1
    arg_names = ["expression"]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = FromArgs()


class FuncIfBase(LogicalFunction):
    scopes = LogicalFunction.scopes & ~Scope.DOCUMENTED  # Documented in _if_block_
    variants = [VW(D.DUMMY, n.func._if_block_)]
    # Disable postprocessing of args so that it can be done when _if_block_ is being applied.
    # Otherwise we will lose the original context flags of these arguments
    postprocess_args = False
    return_type = IfTypeStrategy()


class FuncIf(FuncIfBase):
    name = "if"
    arg_cnt = None


class FuncIif3Legacy(FuncIfBase):
    name = "iif"
    arg_cnt = 3
    return_flags = ContextFlag.DEPRECATED


class FuncCase(LogicalFunction):
    name = "case"
    arg_cnt = None
    scopes = LogicalFunction.scopes & ~Scope.DOCUMENTED  # Documented in _case_block_
    variants = [VW(D.DUMMY, n.func._case_block_)]  # noqa
    return_type = CaseTypeStrategy()


DEFINITIONS_LOGICAL = [
    # case
    FuncCase,
    # if
    FuncIf,
    # ifnan
    FuncIfnan,
    # ifnull
    FuncIfnull,
    # iif
    FuncIif3Legacy,
    # isnan
    FuncIsnan,
    # isnull
    FuncIsnull,
    # zn
    FuncZn,
]

lib/dl_formula/dl_formula/definitions/functions_logical.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Sequence,
    Set,
)

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import (
    ArgTypeForAll,
    ArgTypeSequence,
)
from dl_formula.definitions.base import (
    Function,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.operators_binary import BinaryPlus
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import Fixed
from dl_formula.shortcuts import n


V = TranslationVariant.make
VW = TranslationVariantWrapped.make

MARKUP_EFFECTIVELY = DataType.MARKUP.autocast_types | DataType.STRING.autocast_types

MARK_LPAR = "("
MARK_RPAR = ")"
MARK_SEP = " "
MARK_QUOT = '"'


class StrictMarkupCompatibleArgTypes(ArgTypeForAll):
    """Markup-compatible arguments with at least one markup argument"""

    __slots__ = ()

    _compatible_types = MARKUP_EFFECTIVELY
    _required_types = DataType.MARKUP.autocast_types

    def __init__(self):
        super().__init__(expected_types=MARKUP_EFFECTIVELY)

    def match_arg_types(self, arg_types: Sequence[DataType]) -> bool:
        if not super().match_arg_types(arg_types):
            return False

        return bool(self._required_types & set(arg_types))

    def get_possible_arg_types_at_pos(self, pos: int, total: int) -> Set[DataType]:
        return set(MARKUP_EFFECTIVELY)


# TODO: make this definitely unnecessary.
def concat_strings(items, sep=""):
    result = []
    literals_buf = []  # type: ignore  # TODO: fix

    def flush():
        if literals_buf:
            result.append(sep.join(literals_buf))
            literals_buf[:] = []

    for item in items:
        if isinstance(item, str):
            literals_buf.append(item)
        else:
            flush()
            result.append(item)
    flush()
    return result


class FuncInternalStrBase(Function):
    """
    Base class for a helper for converting markup to str internally.
    """

    name = "__str"
    scopes = Function.scopes & ~Scope.EXPLICIT_USAGE & ~Scope.SUGGESTED & ~Scope.DOCUMENTED

    arg_names = ["expression"]
    arg_cnt = 1

    # Markup is already a string value, so nothing to do.
    variants = [
        V(D.DUMMY | D.SQLITE, lambda value: value),
    ]


class FuncInternalStrConst(FuncInternalStrBase):
    argument_types = [
        ArgTypeSequence([{DataType.CONST_MARKUP, DataType.CONST_STRING}]),
    ]
    return_type = Fixed(DataType.CONST_STRING)


class FuncInternalStr(FuncInternalStrBase):
    argument_types = [
        ArgTypeSequence([MARKUP_EFFECTIVELY]),
    ]
    return_type = Fixed(DataType.STRING)


def concat(nodes):
    nodes = concat_strings(nodes)
    nodes_processed = [n.func.__STR(node) for node in nodes]

    if len(nodes_processed) == 1:
        return nodes_processed[0]

    return n.func.CONCAT(*nodes_processed)


def process_markup_child(node):
    if node.data_type == DataType.MARKUP:
        return node
    if node.data_type == DataType.CONST_MARKUP:
        return node
    if node.data_type == DataType.STRING:
        pieces = [MARK_QUOT, n.func.REPLACE(node, MARK_QUOT, MARK_QUOT + MARK_QUOT), MARK_QUOT]
        return concat(pieces)
    # XXXX: Should probaly be handled automatically by the `STRING` case above.
    if node.data_type == DataType.CONST_STRING:
        value = node.expression.value
        return "".join((MARK_QUOT, str(value).replace(MARK_QUOT, MARK_QUOT + MARK_QUOT), MARK_QUOT))
    if node.data_type == DataType.NULL:
        return node
    raise Exception("Unexpected markup child type", node)


def markup_node(func_name, *children):
    """
    Using `concat` and `replace`, return an expression that results in
    something parseable by the special parser.
    """
    pieces = [
        MARK_LPAR,
        func_name,
    ]
    for child in children:
        child_res = process_markup_child(child)
        pieces += [MARK_SEP, child_res]
    pieces += [MARK_RPAR]
    return concat(pieces)


def make_variants(func_name):
    return [
        VW(
            D.DUMMY | D.SQLITE,
            lambda *children: markup_node(func_name, *children),
        ),
    ]


class MarkupTypeStrategy(Fixed):
    """
    Consts -> CONST_MARKUP, otherwise non-const MARKUP.

    Inherits `Fixed` for docgen.
    """

    def __init__(self):
        super().__init__(DataType.MARKUP)

    def get_from_args(self, arg_types):
        if all(arg.casts_to(DataType.CONST_MARKUP) or arg.casts_to(DataType.CONST_STRING) for arg in arg_types):
            return DataType.CONST_MARKUP
        return DataType.MARKUP


class FuncMarkup(Function):
    return_type = MarkupTypeStrategy()


class FuncMarkupUnary(FuncMarkup):
    # required: name
    # required: variants

    # optional: description
    # optional: notes
    # optional: examples

    # TODO: make all unary markup functions accept an arbitrary amount of
    # arguments, passing them to concat internally.

    arg_cnt = 1
    arg_names = ["text"]
    argument_types = [
        ArgTypeSequence([MARKUP_EFFECTIVELY]),
    ]


class FuncBold(FuncMarkupUnary):
    name = "bold"
    variants = make_variants("b")


class FuncItalics(FuncMarkupUnary):
    name = "italic"
    variants = make_variants("i")


class FuncUrl(FuncMarkup):
    name = "url"
    arg_cnt = 2
    arg_names = ["address", "text"]
    argument_types = [
        # url(str, markup|str)
        ArgTypeSequence([DataType.STRING, MARKUP_EFFECTIVELY]),
    ]
    variants = make_variants("a")


class BinaryPlusMarkup(BinaryPlus):
    variants = make_variants("c")
    argument_types = [
        StrictMarkupCompatibleArgTypes(),
    ]
    return_type = Fixed(DataType.MARKUP)


class ConcatMultiMarkup(FuncMarkup):
    """
    User-facing 'concat' of markups.
    Also useful for turning a string into markup,
    e.g. `if cond then italics(col) else markup(col) end` case.
    """

    name = "markup"
    arg_cnt = None
    argument_types = [
        ArgTypeForAll(MARKUP_EFFECTIVELY),
    ]
    return_type = Fixed(DataType.MARKUP)  # type: ignore  # TODO: fix

    variants = make_variants("c")


DEFINITIONS_MARKUP = [
    # +
    BinaryPlusMarkup,
    # __str
    FuncInternalStrConst,
    FuncInternalStr,
    # bold
    FuncBold,
    # italic
    FuncItalics,
    # markup
    ConcatMultiMarkup,
    # url
    FuncUrl,
]

lib/dl_formula/dl_formula/definitions/functions_markup.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import (
    ArgTypeForAll,
    ArgTypeSequence,
)
from dl_formula.definitions.base import (
    Function,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.common import make_binary_chain
from dl_formula.definitions.type_strategy import (
    Fixed,
    FromArgs,
    ParamsFromArgs,
)
from dl_formula.shortcuts import n


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class MathFunction(Function):
    pass


class FuncAbs(MathFunction):
    name = "abs"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY | D.SQLITE, sa.func.ABS),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = FromArgs()


class FuncAcos(MathFunction):
    name = "acos"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.ACOS),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncAsin(MathFunction):
    name = "asin"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.ASIN),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncAtan(MathFunction):
    name = "atan"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.ATAN),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncAtan2(MathFunction):
    name = "atan2"
    arg_cnt = 2
    arg_names = ["x", "y"]
    variants = [
        V(D.DUMMY, sa.func.ATAN2),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncCeiling(MathFunction):
    name = "ceiling"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.CEIL),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncFloor(MathFunction):
    name = "floor"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.FLOOR),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncCos(MathFunction):
    name = "cos"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.COS),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncCot(MathFunction):
    name = "cot"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.COT),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncDegrees(MathFunction):
    name = "degrees"
    arg_cnt = 1
    arg_names = ["radians"]
    variants = [
        V(D.DUMMY, sa.func.DEGREES),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncDiv(MathFunction):
    name = "div"
    arg_cnt = 2
    arg_names = ["number_1", "number_2"]
    return_type = Fixed(DataType.INTEGER)


class FuncDivBasic(FuncDiv):
    variants = [
        V(D.DUMMY, sa.func.DIV),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]


class FuncDivSafeBase(MathFunction):
    name = "div_safe"
    return_type = Fixed(DataType.INTEGER)


class FuncDivSafe2(FuncDivSafeBase):
    arg_names = ["numerator", "denominator"]
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]
    variants = [
        V(D.DUMMY, lambda x, y: sa.func.div(x, sa.func.nullif(y, 0))),
    ]


class FuncDivSafe3(FuncDivSafeBase):
    arg_names = ["numerator", "denominator", "fallback_value"]
    arg_cnt = 3
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT, DataType.INTEGER]),
    ]
    variants = [V(D.DUMMY, lambda x, y, default: sa.func.coalesce(sa.func.div(x, sa.func.nullif(y, 0)), default))]


class FuncFDivSafeBase(MathFunction):
    name = "fdiv_safe"
    return_type = Fixed(DataType.FLOAT)


class FuncFDivSafe2(FuncFDivSafeBase):
    arg_names = ["numerator", "denominator"]
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]
    variants = [
        V(D.DUMMY, lambda x, y: x / sa.func.nullif(y, 0)),
    ]


class FuncFDivSafe3(FuncFDivSafeBase):
    arg_names = ["numerator", "denominator", "fallback_value"]
    arg_cnt = 3
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT, DataType.FLOAT]),
    ]
    variants = [V(D.DUMMY, lambda x, y, default: sa.func.coalesce(x / sa.func.nullif(y, 0), default))]


class FuncExp(MathFunction):
    name = "exp"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.EXP),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncLn(MathFunction):
    name = "ln"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.LN),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncLog10(MathFunction):
    name = "log10"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.LOG10),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncLog(MathFunction):
    name = "log"
    arg_cnt = 2
    arg_names = ["value", "base"]
    variants = [
        V(D.DUMMY, lambda x, base: sa.func.LOG(base, x)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


_COMPARABLE_TYPES = [
    DataType.INTEGER,
    DataType.FLOAT,
    DataType.BOOLEAN,
    DataType.STRING,
    DataType.DATE,
    DataType.DATETIME,
    DataType.DATETIMETZ,
    DataType.GENERICDATETIME,
]


class FuncGreatestBase(MathFunction):
    name = "greatest"
    arg_names = ["value_1", "value_2", "value_3"]
    arg_cnt = 2
    # TODO: sqlite3: just `max(a, b, ...)` works
    return_type = FromArgs()
    return_type_params = ParamsFromArgs(0)


class FuncGreatest1(FuncGreatestBase):
    arg_cnt = 1
    variants = [
        V(D.DUMMY | D.SQLITE, lambda value: value),
    ]
    argument_types = [
        ArgTypeSequence([set(_COMPARABLE_TYPES)]),
    ]


class FuncGreatestMain(FuncGreatestBase):
    variants = [V(D.DUMMY, sa.func.GREATEST)]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.DATE, DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME, DataType.DATETIME]),
        ArgTypeSequence([DataType.DATETIMETZ, DataType.DATETIMETZ]),
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.GENERICDATETIME]),
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
        ArgTypeSequence([DataType.BOOLEAN, DataType.BOOLEAN]),
    ]


class GreatestMulti(FuncGreatestBase):
    arg_cnt = None  # type: ignore  # TODO: fix
    variants = [
        VW(
            D.DUMMY | D.SQLITE,
            lambda *args: make_binary_chain(
                # chain of GREATEST function calls
                (lambda x, y: n.func.GREATEST(x, y)),
                *args,
            ),
        ),
    ]
    argument_types = [
        # any type is allowed, but all arguments must be of the same type
        ArgTypeForAll({_t})
        for _t in _COMPARABLE_TYPES
    ]
    return_type_params = ParamsFromArgs(0)  # particularly for datetimetz


class FuncLeastBase(MathFunction):
    name = "least"
    arg_names = ["value_1", "value_2", "value_3"]
    arg_cnt = 2
    return_type = FromArgs()
    return_type_params = ParamsFromArgs(0)


class FuncLeast1(FuncLeastBase):
    arg_cnt = 1
    variants = [
        V(D.SQLITE, lambda value: value),
    ]
    argument_types = [
        ArgTypeSequence([set(_COMPARABLE_TYPES)]),
    ]


class FuncLeastMain(FuncLeastBase):
    variants = [
        V(D.DUMMY, sa.func.LEAST),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.DATE, DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME, DataType.DATETIME]),
        ArgTypeSequence([DataType.DATETIMETZ, DataType.DATETIMETZ]),
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.GENERICDATETIME]),
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
        ArgTypeSequence([DataType.BOOLEAN, DataType.BOOLEAN]),
    ]


class LeastMulti(FuncLeastBase):
    arg_cnt = None  # type: ignore  # TODO: fix
    variants = [
        VW(
            D.DUMMY | D.SQLITE,
            lambda *args: make_binary_chain(
                # chain of LEAST function calls
                (lambda x, y: n.func.LEAST(x, y)),
                *args,
            ),
        ),
    ]
    argument_types = [
        # any type is allowed, but all arguments must be of the same type
        ArgTypeForAll({_t})
        for _t in _COMPARABLE_TYPES
    ]
    return_type_params = ParamsFromArgs(0)  # particularly for datetimetz


class FuncPi(MathFunction):
    name = "pi"
    arg_cnt = 0
    variants = [
        V(D.DUMMY, sa.func.PI),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncPower(MathFunction):
    name = "power"
    arg_cnt = 2
    arg_names = ["base", "power"]
    variants = [
        V(D.DUMMY, sa.func.POWER),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncRadians(MathFunction):
    name = "radians"
    arg_cnt = 1
    arg_names = ["degrees"]
    variants = [
        V(D.DUMMY, sa.func.RADIANS),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncRoundBase(MathFunction):
    name = "round"
    arg_names = ["number", "precision"]


class FuncRound1(FuncRoundBase):
    arg_cnt = 1
    variants = [
        V(D.DUMMY | D.SQLITE, sa.func.ROUND),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.INTEGER)


class FuncRound2(FuncRoundBase):
    arg_cnt = 2
    variants = [
        V(D.DUMMY | D.SQLITE, sa.func.ROUND),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.INTEGER]),
    ]
    return_type = FromArgs(0)


class FuncSign(MathFunction):
    name = "sign"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.SIGN),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.INTEGER)


class FuncSin(MathFunction):
    name = "sin"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.SIN),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncSqrt(MathFunction):
    name = "sqrt"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.SQRT),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncSquare(MathFunction):
    name = "square"
    arg_names = ["number"]
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = FromArgs()


class FuncTan(MathFunction):
    name = "tan"
    arg_names = ["number"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.TAN),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class FuncCompare(MathFunction):
    name = "compare"
    arg_names = ["left", "right", "epsilon"]
    arg_cnt = 3
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT, DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.INTEGER)


DEFINITIONS_MATH = [
    # abs
    FuncAbs,
    # acos
    FuncAcos,
    # asin
    FuncAsin,
    # atan
    FuncAtan,
    # atan2
    FuncAtan2,
    # ceiling
    FuncCeiling,
    # compare
    FuncCompare,
    # cos
    FuncCos,
    # cot
    FuncCot,
    # degrees
    FuncDegrees,
    # div
    FuncDivBasic,
    # div_safe
    FuncDivSafe2,
    FuncDivSafe3,
    # fdiv_safe
    FuncFDivSafe2,
    FuncFDivSafe3,
    # exp
    FuncExp,
    # floor
    FuncFloor,
    # greatest
    FuncGreatest1,
    FuncGreatestMain,
    GreatestMulti,
    # least
    FuncLeast1,
    FuncLeastMain,
    LeastMulti,
    # ln
    FuncLn,
    # log
    FuncLog,
    # log10
    FuncLog10,
    # pi
    FuncPi,
    # power
    FuncPower,
    # radians
    FuncRadians,
    # round
    FuncRound1,
    FuncRound2,
    # sign
    FuncSign,
    # sin
    FuncSin,
    # sqrt
    FuncSqrt,
    # square
    FuncSquare,
    # tan
    FuncTan,
]

lib/dl_formula/dl_formula/definitions/functions_math.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_formula.core.datatype import DataType
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    Function,
    TranslationVariant,
)
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import Fixed


V = TranslationVariant.make


class SpecialFunction(Function):
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncMakeNan(SpecialFunction):
    name = "_make_nan"
    arg_cnt = 0
    argument_types = [ArgTypeSequence([])]
    return_type = Fixed(DataType.FLOAT)


DEFINITIONS_SPECIAL = [
    # _make_nan
    FuncMakeNan,
]

lib/dl_formula/dl_formula/definitions/functions_special.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import (
    ArgTypeForAll,
    ArgTypeSequence,
)
from dl_formula.definitions.base import (
    Function,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.common import make_binary_chain
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.literals import literal
from dl_formula.definitions.type_strategy import Fixed
from dl_formula.shortcuts import n


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class StringFunction(Function):
    pass


class FuncAscii(StringFunction):
    name = "ascii"
    arg_names = ["string"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.ASCII),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]
    return_type = Fixed(DataType.INTEGER)


class FuncChar(StringFunction):
    name = "char"
    arg_names = ["string"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.CHAR),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.STRING)


class ConcatBase(StringFunction):
    name = "concat"
    return_type = Fixed(DataType.STRING)


class ConcatMulti(ConcatBase):
    arg_cnt = None


_CONCAT_TYPES = {
    DataType.STRING,
    DataType.INTEGER,
    DataType.FLOAT,
    DataType.BOOLEAN,
    DataType.DATE,
    DataType.DATETIME,
    DataType.DATETIMETZ,
    DataType.GENERICDATETIME,
    DataType.GEOPOINT,
    DataType.GEOPOLYGON,
    DataType.UUID,
    # Should not be implemented: MARKUP:
    # markup should not be castable to string,
    # and, for consistency, `CONCAT` should always return a string.
}


class ConcatMultiStrConst(ConcatMulti):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda *args, _env: literal("".join(arg.value for arg in args), d=_env.dialect)),
    ]
    argument_types = [
        ArgTypeForAll(DataType.CONST_STRING),
    ]
    return_type = Fixed(DataType.CONST_STRING)


class Concat1(ConcatBase):
    arg_cnt = 1
    variants = [
        VW(D.DUMMY | D.SQLITE, n.func.STR),
    ]
    argument_types = [
        ArgTypeSequence([_CONCAT_TYPES]),
    ]


class ConcatMultiStr(ConcatMulti):
    """Simplified version that can only concatenate strings"""

    # TODO: join adjacent const strings, remove single-argument concat.
    variants = [
        V(D.DUMMY, lambda *args: sa.func.CONCAT(*args)),
    ]
    argument_types = [
        ArgTypeForAll(DataType.STRING),
    ]


class ConcatMultiAny(ConcatMulti):
    """Concatenation of any-type arguments"""

    variants = [
        VW(
            D.DUMMY | D.SQLITE,
            lambda *args: make_binary_chain(
                # convert all to string and then concatenate
                (lambda x, y: n.func.CONCAT(n.func.STR(x), n.func.STR(y))),
                *args,
            ),
        ),
    ]
    argument_types = [
        ArgTypeForAll(_CONCAT_TYPES),
    ]


class FuncContains(StringFunction):
    name = "contains"
    arg_names = ["string", "substring"]
    arg_cnt = 2
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


def quote_like(value, escape="\\", chars=("%", "_")):
    for char in (escape,) + tuple(chars):
        value = value.replace(char, escape + char)
    return value


def quote_like_metrica(value):
    value = value.replace("\\", r"\\")
    value = value.replace("*", r"\*")
    return value


NON_STR_CONTAINMENT_TYPES = {
    DataType.BOOLEAN,
    DataType.INTEGER,
    DataType.FLOAT,
    DataType.DATE,
    DataType.DATETIME,
    DataType.DATETIMETZ,
    DataType.GENERICDATETIME,
    DataType.UUID,
    DataType.GEOPOINT,
    DataType.GEOPOLYGON,
}


class FuncContainsConst(FuncContains):
    variants = [
        V(D.DUMMY, lambda x, y: x.like("%{}%".format(quote_like(y.value)), escape="\\")),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING]),
    ]


class FuncContainsNonConst(FuncContains):
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]


class FuncContainsNonString(FuncContains):
    variants = [
        VW(D.DUMMY, lambda x, y: n.func.CONTAINS(n.func.STR(x), y)),
    ]
    argument_types = [
        ArgTypeSequence([NON_STR_CONTAINMENT_TYPES, DataType.STRING]),
    ]


class FuncNotContains(StringFunction):
    name = "notcontains"
    arg_names = ["string", "substring"]
    arg_cnt = 2
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION
    variants = [
        VW(D.DUMMY, lambda x, y: n.not_(n.func.CONTAINS(x, y))),
    ]


class FuncNotContainsConst(FuncNotContains):
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING]),
    ]


class FuncNotContainsNonConst(FuncNotContains):
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]


class FuncNotContainsNonString(FuncNotContains):
    argument_types = [
        ArgTypeSequence([NON_STR_CONTAINMENT_TYPES, DataType.STRING]),
    ]


class FuncIContains(StringFunction):
    name = "icontains"
    arg_names = ["string", "substring"]
    arg_cnt = 2
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


class FuncIContainsConst(FuncIContains):
    variants = [
        V(
            D.DUMMY,
            lambda x, y: x.ilike(
                # WARNING: even with a proper collation, the python's
                # `val.lower()` might not exactly match postgresql's `lower(val)`
                "%{}%".format(quote_like(y.value.lower())),
                escape="\\",
            ),
        ),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING]),
    ]


class FuncIContainsNonConst(FuncIContains):
    variants = [
        VW(
            D.DUMMY | D.SQLITE,
            lambda s, substr: n.func.CONTAINS(n.func.LOWER(s), n.func.LOWER(substr)),
        ),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]
    return_flags = ContextFlag.IS_CONDITION


class FuncIContainsNonString(FuncIContains):
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda x, y: n.func.ICONTAINS(n.func.STR(x), y)),
    ]
    argument_types = [
        ArgTypeSequence([NON_STR_CONTAINMENT_TYPES, DataType.STRING]),
    ]


class FuncEndswith(StringFunction):
    name = "endswith"
    arg_names = ["string", "substring"]
    arg_cnt = 2
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


class FuncEndswithConst(FuncEndswith):
    variants = [
        V(D.DUMMY, lambda x, y: x.like("%{}".format(quote_like(y.value)), escape="\\")),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING]),
    ]


class FuncEndswithNonConst(FuncEndswith):
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]


class FuncEndswithNonString(FuncEndswith):
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda x, y: n.func.ENDSWITH(n.func.STR(x), y)),
    ]
    argument_types = [
        ArgTypeSequence([NON_STR_CONTAINMENT_TYPES, DataType.STRING]),
    ]


class FuncIEndswith(StringFunction):
    name = "iendswith"
    arg_names = ["string", "substring"]
    arg_cnt = 2
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


class FuncIEndswithConst(FuncIEndswith):
    variants = [
        V(D.DUMMY, lambda x, y: x.ilike("%{}".format(quote_like(y.value.lower())), escape="\\")),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING]),
    ]


class FuncIEndswithNonConst(FuncIEndswith):
    variants = [
        VW(
            D.DUMMY | D.SQLITE,
            lambda s, substr: n.func.ENDSWITH(n.func.LOWER(s), n.func.LOWER(substr)),
        ),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]


class FuncIEndswithNonString(FuncIEndswith):
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda x, y: n.func.IENDSWITH(n.func.STR(x), y)),
    ]
    argument_types = [
        ArgTypeSequence([NON_STR_CONTAINMENT_TYPES, DataType.STRING]),
    ]


class FuncStartswith(StringFunction):
    name = "startswith"
    arg_names = ["string", "substring"]
    arg_cnt = 2
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


class FuncStartswithConst(FuncStartswith):
    variants = [
        V(D.DUMMY, lambda x, y: x.like("{}%".format(quote_like(y.value)), escape="\\")),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING]),
    ]


class FuncStartswithNonConst(FuncStartswith):
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]


class FuncStartswithNonString(FuncStartswith):
    variants = [
        VW(D.DUMMY, lambda x, y: n.func.STARTSWITH(n.func.STR(x), y)),
    ]
    argument_types = [
        ArgTypeSequence([NON_STR_CONTAINMENT_TYPES, DataType.STRING]),
    ]


class FuncIStartswith(StringFunction):
    name = "istartswith"
    arg_names = ["string", "substring"]
    arg_cnt = 2
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


class FuncIStartswithConst(FuncIStartswith):
    variants = [
        V(D.DUMMY, lambda x, y: x.ilike("{}%".format(quote_like(y.value.lower())), escape="\\")),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING]),
    ]


class FuncIStartswithNonConst(FuncIStartswith):
    variants = [
        VW(
            D.DUMMY | D.SQLITE,
            lambda s, substr: n.func.STARTSWITH(n.func.LOWER(s), n.func.LOWER(substr)),
        ),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]


class FuncIStartswithNonString(FuncIStartswith):
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda x, y: n.func.ISTARTSWITH(n.func.STR(x), y)),
    ]
    argument_types = [
        ArgTypeSequence([NON_STR_CONTAINMENT_TYPES, DataType.STRING]),
    ]


class FuncFind(StringFunction):
    name = "find"
    arg_names = ["string", "substring", "start_index"]
    return_type = Fixed(DataType.INTEGER)


class FuncFind2(FuncFind):
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]


class FuncFind3(FuncFind):
    arg_cnt = 3
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING, DataType.INTEGER]),
    ]


class FuncLeft(StringFunction):
    name = "left"
    arg_cnt = 2
    arg_names = ["string", "number"]
    variants = [
        V(D.DUMMY, sa.func.LEFT),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncRight(StringFunction):
    name = "right"
    arg_cnt = 2
    arg_names = ["string", "number"]
    variants = [
        V(D.DUMMY, sa.func.RIGHT),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncLen(StringFunction):
    name = "len"
    arg_names = ["value"]
    arg_cnt = 1
    return_type = Fixed(DataType.INTEGER)


class FuncLenString(FuncLen):
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]


class FuncLower(StringFunction):
    name = "lower"
    arg_names = ["string"]
    arg_cnt = 1


class FuncLowerConst(FuncLower):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda s, _env: literal(s.value.lower(), d=_env.dialect)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.CONST_STRING]),
    ]
    return_type = Fixed(DataType.CONST_STRING)


class FuncLowerNonConst(FuncLower):
    variants = [
        V(D.DUMMY, sa.func.LOWER),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncUpper(StringFunction):
    name = "upper"
    arg_names = ["string"]
    arg_cnt = 1


class FuncUpperConst(FuncUpper):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda s, _env: literal(s.value.upper(), d=_env.dialect)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.CONST_STRING]),
    ]
    return_type = Fixed(DataType.CONST_STRING)


class FuncUpperNonConst(FuncUpper):
    variants = [
        V(D.DUMMY, sa.func.UPPER),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncLtrim(StringFunction):
    name = "ltrim"
    arg_names = ["string"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.LTRIM),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncRtrim(StringFunction):
    name = "rtrim"
    arg_names = ["string"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.RTRIM),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncTrim(StringFunction):
    name = "trim"
    arg_names = ["string"]
    arg_cnt = 1
    variants = [
        V(D.DUMMY, sa.func.TRIM),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncSubstr(StringFunction):
    name = "substr"
    arg_names = ["string", "from_index", "length"]


class FuncSubstr2(FuncSubstr):
    arg_cnt = 2
    variants = [
        V(D.DUMMY, sa.func.SUBSTRING),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncSubstr3(FuncSubstr):
    arg_cnt = 3
    variants = [
        V(D.DUMMY, sa.func.SUBSTRING),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.INTEGER, DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncRegexpExtract(StringFunction):
    name = "regexp_extract"
    arg_cnt = 2
    arg_names = ["string", "pattern"]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING]),
    ]
    return_type = Fixed(DataType.STRING)


# In MS SQL Server regular expressions can only be used in user-defined functions


class FuncRegexpExtractNth(StringFunction):
    name = "regexp_extract_nth"
    arg_cnt = 3
    arg_names = ["string", "pattern", "match_index"]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING, DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncRegexpMatch(StringFunction):
    name = "regexp_match"
    arg_cnt = 2
    arg_names = ["string", "pattern"]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


class FuncRegexpReplace(StringFunction):
    name = "regexp_replace"
    arg_cnt = 3
    arg_names = ["string", "pattern", "replace_with"]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING, DataType.STRING]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncReplace(StringFunction):
    name = "replace"
    arg_cnt = 3
    arg_names = ["string", "substring", "replace_with"]
    variants = [
        V(D.DUMMY, sa.func.REPLACE),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING, DataType.STRING]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncSpace(StringFunction):
    name = "space"
    arg_cnt = 1
    return_type = Fixed(DataType.STRING)


class FuncSpaceConst(FuncSpace):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda size, _env: literal(" " * size.value, d=_env.dialect)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.CONST_INTEGER]),
    ]
    return_type = Fixed(DataType.CONST_STRING)


class FuncSpaceNonConst(FuncSpace):
    variants = [
        V(D.DUMMY, sa.func.SPACE),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
    ]


class FuncSplit(StringFunction):
    name = "split"
    arg_names = ["orig_string", "delimiter", "part_index"]


class FuncSplit1(FuncSplit):
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]
    return_type = Fixed(DataType.ARRAY_STR)


class FuncSplit2(FuncSplit):
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING]),
    ]
    return_type = Fixed(DataType.ARRAY_STR)


class FuncSplit3(FuncSplit):
    arg_cnt = 3
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.CONST_STRING, DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.STRING)


class FuncUtf8(StringFunction):
    name = "utf8"
    arg_cnt = 2
    arg_names = ["string", "old_encoding"]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]
    return_type = Fixed(DataType.STRING)


DEFINITIONS_STRING = [
    # ascii
    FuncAscii,
    # char
    FuncChar,
    # concat
    Concat1,
    ConcatMultiStrConst,
    ConcatMultiStr,
    ConcatMultiAny,
    # contains
    FuncContainsConst,
    FuncContainsNonConst,
    FuncContainsNonString,
    # notcontains
    FuncNotContainsConst,
    FuncNotContainsNonConst,
    FuncNotContainsNonString,
    # endswith
    FuncEndswithConst,
    FuncEndswithNonConst,
    FuncEndswithNonString,
    # find
    FuncFind2,
    FuncFind3,
    # icontains
    FuncIContainsConst,
    FuncIContainsNonConst,
    FuncIContainsNonString,
    # iendswith
    FuncIEndswithConst,
    FuncIEndswithNonConst,
    FuncIEndswithNonString,
    # istartswith
    FuncIStartswithConst,
    FuncIStartswithNonConst,
    FuncIStartswithNonString,
    # left
    FuncLeft,
    # len
    FuncLenString,
    # lower
    FuncLowerConst,
    FuncLowerNonConst,
    # ltrim
    FuncLtrim,
    # regexp_extract
    FuncRegexpExtract,
    # regexp_extract_nth
    FuncRegexpExtractNth,
    # regexp_match
    FuncRegexpMatch,
    # regexp_replace
    FuncRegexpReplace,
    # replace
    FuncReplace,
    # right
    FuncRight,
    # rtrim
    FuncRtrim,
    # space
    FuncSpaceConst,
    FuncSpaceNonConst,
    # split
    FuncSplit1,
    FuncSplit2,
    FuncSplit3,
    # startswith
    FuncStartswithConst,
    FuncStartswithNonConst,
    FuncStartswithNonString,
    # substr
    FuncSubstr2,
    FuncSubstr3,
    # trim
    FuncTrim,
    # upper
    FuncUpperConst,
    FuncUpperNonConst,
    # utf8
    FuncUtf8,
]

lib/dl_formula/dl_formula/definitions/functions_string.py
datalens-tech/datalens-backend
from __future__ import annotations

import dl_formula.core.aux_nodes as aux_nodes
from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
import dl_formula.core.exc as exc
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    Function,
    TranslationVariantWrapped,
)
from dl_formula.definitions.type_strategy import FromArgs


VW = TranslationVariantWrapped.make


class TimeSeriesFunction(Function):
    supports_ignore_dimensions = True
    supports_bfb = True
    variants = [
        # These functions are not translatable,
        # so raise an error if we ever reach the translator here.
        VW(
            D.ANY,
            lambda *args: aux_nodes.ErrorNode.make(
                message="Lookup function should have been mutated",
                err_code=exc.TranslationError.default_code,
            ),
        )
    ]


class FuncAgoBase(TimeSeriesFunction):
    name = "ago"
    arg_names = ["measure", "date_dimension", "unit", "number"]
    argument_types = [
        ArgTypeSequence(
            [
                set(DataType),
                {DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME},
                DataType.CONST_STRING,
                DataType.CONST_INTEGER,
            ]
        ),
    ]
    return_type = FromArgs(0)


class FuncAgo2(FuncAgoBase):
    arg_cnt = 2


class FuncAgo3Unit(FuncAgoBase):
    arg_cnt = 3


class FuncAgo3Number(FuncAgoBase):
    """A variant with the number of days as the third argument with the units omitted"""

    arg_cnt = 3
    argument_types = [
        ArgTypeSequence(
            [
                set(DataType),
                {DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME},
                DataType.CONST_INTEGER,
            ]
        ),
    ]


class FuncAgo4(FuncAgoBase):
    arg_cnt = 4


class FuncAtDateBase(TimeSeriesFunction):
    name = "at_date"
    arg_names = ["measure", "date_dimension", "date_expr"]
    argument_types = [
        ArgTypeSequence([set(DataType), DataType.DATE, DataType.DATE]),
        ArgTypeSequence([set(DataType), DataType.DATETIME, DataType.DATETIME]),
        ArgTypeSequence([set(DataType), DataType.GENERICDATETIME, DataType.GENERICDATETIME]),
    ]
    return_type = FromArgs(0)


class FuncAtDate3(FuncAtDateBase):
    arg_cnt = 3


DEFINITIONS_TIME_SERIES = [
    # ago
    FuncAgo2,
    FuncAgo3Unit,
    FuncAgo3Number,
    FuncAgo4,
    # at_date
    FuncAtDate3,
]

lib/dl_formula/dl_formula/definitions/functions_time_series.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    ClassVar,
    NamedTuple,
    Optional,
    Sequence,
    Set,
    Tuple,
    Type,
    Union,
)

import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement
from sqlalchemy.types import TypeEngine

from dl_formula.core import exc
from dl_formula.core.datatype import (
    DataType,
    DataTypeParams,
)
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    Function,
    SingleVariantTranslationBase,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.literals import (
    literal,
    un_literal,
)
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import (
    Fixed,
    FromArgs,
    ParamsCustom,
    ParamsEmpty,
)
from dl_formula.shortcuts import n
from dl_formula.utils.datetime import make_datetimetz_value


if TYPE_CHECKING:
    from dl_formula.core.dialect import DialectCombo
    from dl_formula.translation.context import TranslationCtx
    from dl_formula.translation.env import TranslationEnvironment


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


class TypeConvFunction(Function):
    pass


class FuncDate(TypeConvFunction):
    name = "date"
    arg_names = ["expression", "timezone"]
    return_type = Fixed(DataType.DATE)


class FuncDate1(FuncDate):
    arg_cnt = 1


class FuncDate1FromNull(FuncDate1):
    variants = [
        V(D.DUMMY, lambda _: sa.cast(sa.null(), sa.Date())),
    ]
    argument_types = [
        ArgTypeSequence([DataType.NULL]),
    ]


class FuncDate1FromDatetime(FuncDate1):
    variants = [
        VW(D.DUMMY, lambda expr: sa.cast(expr.expression, sa.Date())),
    ]
    argument_types = [
        ArgTypeSequence([{DataType.DATE, DataType.DATETIME, DataType.GENERICDATETIME}]),
    ]


class FuncDate1FromDatetimeTZ(FuncDate1):
    # CLICKHOUSE & POSTGRESQL
    argument_types = [
        ArgTypeSequence([DataType.DATETIMETZ]),
    ]


class FuncDate1FromString(FuncDate1):
    variants = [
        # TODO: sqlite3: sa.func.date
        V(D.DUMMY, lambda expr: sa.cast(expr, sa.Date())),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]


class FuncDate1FromNumber(FuncDate1):
    # Custom implementation for each dialect
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
        ArgTypeSequence([DataType.INTEGER]),
    ]


class FuncDate2(FuncDate):
    arg_cnt = 2


class FuncDatetimeImpl(TypeConvFunction):
    # WARNING: the `{arg:1}` in this description requires there to be at least
    # one `arg_cnt = 2` registered subclass translation, such as
    # `FuncDatetime2FromCHStuff`.
    arg_names = ["expression", "timezone"]


class FuncTypeGenericDatetimeImpl(FuncDatetimeImpl):
    return_type = Fixed(DataType.GENERICDATETIME)


class FuncTypeGenericDatetime1Impl(FuncTypeGenericDatetimeImpl):
    arg_cnt = 1


class FuncTypeGenericDatetime1FromNullImpl(FuncTypeGenericDatetime1Impl):
    variants = [
        V(D.DUMMY, lambda _: sa.cast(sa.null(), sa.DateTime())),
    ]
    argument_types = [
        ArgTypeSequence([DataType.NULL]),
    ]


class FuncTypeGenericDatetime1FromDatetimeImpl(FuncTypeGenericDatetime1Impl):
    variants = [V(D.DUMMY | D.SQLITE, lambda expr: expr)]
    argument_types = [
        ArgTypeSequence([{DataType.DATETIME, DataType.GENERICDATETIME}]),
    ]


class FuncTypeGenericDatetime1FromDateImpl(FuncTypeGenericDatetime1Impl):
    variants = [
        V(D.DUMMY, lambda expr: sa.cast(expr, sa.DateTime())),
    ]
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
    ]


class FuncTypeGenericDatetime1FromNumberImpl(FuncTypeGenericDatetime1Impl):
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
        ArgTypeSequence([DataType.INTEGER]),
    ]


class FuncTypeGenericDatetime1FromStringImpl(FuncTypeGenericDatetime1Impl):
    variants = [
        V(D.DUMMY, lambda expr: sa.cast(expr, sa.DateTime())),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]


class FuncTypeGenericDatetime2Impl(FuncTypeGenericDatetimeImpl):
    arg_cnt = 2


class FuncDatetime1FromNull(FuncTypeGenericDatetime1FromNullImpl):
    name = "datetime"


class FuncDatetime1FromDatetime(FuncTypeGenericDatetime1FromDatetimeImpl):
    name = "datetime"


class FuncDatetime1FromDate(FuncTypeGenericDatetime1FromDateImpl):
    name = "datetime"


class FuncDatetime1FromNumber(FuncTypeGenericDatetime1FromNumberImpl):
    name = "datetime"


class FuncDatetime1FromString(FuncTypeGenericDatetime1FromStringImpl):
    name = "datetime"


class FuncGenericDatetime1FromNull(FuncTypeGenericDatetime1FromNullImpl):
    name = "genericdatetime"
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncGenericDatetime1FromDatetime(FuncTypeGenericDatetime1FromDatetimeImpl):
    name = "genericdatetime"
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncGenericDatetime1FromDate(FuncTypeGenericDatetime1FromDateImpl):
    name = "genericdatetime"
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncGenericDatetime1FromNumber(FuncTypeGenericDatetime1FromNumberImpl):
    name = "genericdatetime"
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncGenericDatetime1FromString(FuncTypeGenericDatetime1FromStringImpl):
    name = "genericdatetime"
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncDatetimeTZ(TypeConvFunction):
    name = "datetimetz"

    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED

    arg_cnt = 2
    arg_names = ["expression", "timezone"]

    return_type = Fixed(DataType.DATETIMETZ)
    return_type_params = ParamsCustom(
        lambda args: DataTypeParams(timezone=un_literal(args[1].expression))  # type: ignore  # TODO: fix
    )


class FuncDatetimeTZConst(SingleVariantTranslationBase, FuncDatetimeTZ):
    dialects = D.DUMMY | D.SQLITE
    argument_types = [
        ArgTypeSequence(
            [
                {
                    DataType.CONST_DATETIME,
                    DataType.CONST_GENERICDATETIME,
                    DataType.CONST_DATETIMETZ,
                    # Dubious: `DataType.CONST_DATE`.
                    DataType.CONST_INTEGER,
                    DataType.CONST_FLOAT,
                    DataType.CONST_STRING,
                },
                DataType.CONST_STRING,
            ]
        ),
    ]

    @classmethod
    def _translate_main(cls, value_ctx, tz_ctx, *, _env: TranslationEnvironment):  # type: ignore  # TODO: fix
        value = un_literal(value_ctx.expression, value_ctx=value_ctx)
        # TODO?: re-check the `type(value)` by `value_ctx.data_type`?

        # TODO: error-handling
        tzname = un_literal(tz_ctx.expression)
        # TODO: error-handling
        try:
            dt = make_datetimetz_value(value, tzname)
        except ValueError as err:
            raise exc.ParseValueError(str(err))
        return literal(dt, d=_env.dialect)


class FuncDatetimeTZToNaive(SingleVariantTranslationBase, TypeConvFunction):
    name = "datetimetz_to_naive"  # TODO: consider having `datetime(datetimetz_value)` for this.
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED
    arg_cnt = 1
    arg_names = ["datetimetz"]
    argument_types = [
        ArgTypeSequence([DataType.DATETIMETZ, DataType.CONST_STRING]),
    ]

    return_type = Fixed(DataType.DATETIME)
    return_type_params = ParamsEmpty()  # this is default, but making it explicit


class FuncFloat(TypeConvFunction):
    name = "float"
    arg_names = ["expression"]
    arg_cnt = 1
    return_type = Fixed(DataType.FLOAT)


class FuncFloatNumber(FuncFloat):
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
    ]


class FuncFloatString(FuncFloat):
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]


class FuncFloatFromBool(FuncFloat):
    argument_types = [
        ArgTypeSequence([DataType.BOOLEAN]),
    ]


class FuncFloatFromDate(FuncFloat):
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
    ]


class FuncFloatFromDatetime(FuncFloat):
    argument_types = [
        ArgTypeSequence([DataType.DATETIME]),
    ]


class FuncFloatFromGenericDatetime(FuncFloat):
    argument_types = [
        ArgTypeSequence([DataType.GENERICDATETIME]),
    ]


class FuncInt(TypeConvFunction):
    name = "int"
    arg_names = ["expression"]
    arg_cnt = 1
    return_type = Fixed(DataType.INTEGER)


class FuncIntFromNull(FuncInt):
    argument_types = [
        ArgTypeSequence([DataType.NULL]),
    ]


class FuncIntFromInt(FuncInt):
    variants = [V(D.DUMMY | D.SQLITE, lambda value: value)]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
    ]


class FuncIntFromFloat(FuncInt):
    variants = [
        # FIXME: use something like `NUMBER(*, 0)` for all Oracle `sa.Integer` cases.
        V(D.SQLITE, lambda value: sa.cast(value, sa.Integer)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]


class FuncIntFromBool(FuncInt):
    argument_types = [
        ArgTypeSequence([DataType.BOOLEAN]),
    ]


class FuncIntFromStr(FuncInt):
    variants = [
        V(D.SQLITE, lambda value: sa.cast(value, sa.Integer)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]


class FuncIntFromDate(FuncInt):
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
    ]


class FuncIntFromDatetime(FuncInt):
    argument_types = [
        ArgTypeSequence([DataType.DATETIME]),
    ]


class FuncIntFromGenericDatetime(FuncInt):
    argument_types = [
        ArgTypeSequence([DataType.GENERICDATETIME]),
    ]


class FuncIntFromDatetimeTZ(FuncInt):
    argument_types = [
        ArgTypeSequence([DataType.DATETIMETZ]),
    ]


class FuncBool(TypeConvFunction):
    name = "bool"
    arg_names = ["expression"]
    arg_cnt = 1
    return_type = Fixed(DataType.BOOLEAN)


class FuncBoolFromNull(FuncBool):
    variants = [
        V(D.DUMMY, lambda _: sa.cast(sa.null(), sa.Boolean())),
    ]
    argument_types = [ArgTypeSequence([DataType.NULL])]


class FuncBoolFromNumber(FuncBool):
    variants = [
        V(D.DUMMY, lambda value: value != 0),
    ]
    argument_types = [ArgTypeSequence([DataType.FLOAT])]
    return_flags = ContextFlag.IS_CONDITION


class FuncBoolFromBool(FuncBool):
    variants = [V(D.DUMMY | D.SQLITE, lambda value: value)]
    argument_types = [
        ArgTypeSequence([DataType.BOOLEAN]),
    ]


class FuncBoolFromStrGeo(FuncBool):
    variants = [
        V(D.DUMMY, lambda value: value != ""),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
        ArgTypeSequence([DataType.GEOPOINT]),
        ArgTypeSequence([DataType.GEOPOLYGON]),
    ]
    return_flags = ContextFlag.IS_CONDITION


class FuncBoolFromDateDatetime(FuncBool):
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda value: n.func.IF(n.func.ISNULL(value), None, True)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME]),
        ArgTypeSequence([DataType.GENERICDATETIME]),
    ]


class FuncStr(TypeConvFunction):
    name = "str"
    arg_names = ["expression"]
    arg_cnt = 1
    return_type = Fixed(DataType.STRING)


class FuncStrFromNull(FuncStr):
    argument_types = [
        ArgTypeSequence([DataType.NULL]),
    ]


class FuncStrFromUnsupported(FuncStr):
    variants = [
        V(D.SQLITE, lambda value: sa.cast(value, sa.TEXT)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.UNSUPPORTED]),
    ]


class FuncStrFromInteger(FuncStr):
    variants = [
        V(D.SQLITE, lambda value: sa.cast(value, sa.TEXT)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
    ]


class FuncStrFromFloat(FuncStr):
    variants = [
        V(D.SQLITE, lambda value: sa.cast(value, sa.TEXT)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]


class FuncStrFromBool(FuncStr):
    argument_types = [
        ArgTypeSequence([DataType.BOOLEAN]),
    ]


class FuncStrFromStrGeo(FuncStr):
    variants = [V(D.DUMMY | D.SQLITE, lambda value: value)]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
        ArgTypeSequence([DataType.GEOPOINT]),
        ArgTypeSequence([DataType.GEOPOLYGON]),
    ]


# Should *not* be done: FuncStrFromMarkup


class FuncStrFromDate(FuncStr):
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
    ]


class FuncStrFromDatetime(FuncStr):
    argument_types = [
        ArgTypeSequence([{DataType.DATETIME, DataType.GENERICDATETIME}]),
    ]


class FuncStrFromDatetimeTZ(FuncStr):
    argument_types = [
        ArgTypeSequence([DataType.DATETIMETZ]),
    ]


class FuncStrFromString(FuncStr):
    variants = [V(D.DUMMY | D.SQLITE, lambda value: value)]
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
    ]


class FuncStrFromUUID(FuncStr):
    argument_types = [
        ArgTypeSequence([DataType.UUID]),
    ]


class FuncStrFromArray(FuncStr):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_FLOAT]),
        ArgTypeSequence([DataType.ARRAY_INT]),
        ArgTypeSequence([DataType.ARRAY_STR]),
    ]


class FuncDatetimeParseImpl(TypeConvFunction):
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
    ]


class FuncTypeGenericDatetimeParseImpl(FuncDatetimeParseImpl):
    return_type = Fixed(DataType.GENERICDATETIME)


class FuncDatetimeParse(FuncTypeGenericDatetimeParseImpl):
    name = "datetime_parse"


class FuncGenericDatetimeParse(FuncTypeGenericDatetimeParseImpl):
    name = "genericdatetime_parse"
    scopes = Function.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class FuncDateParse(TypeConvFunction):
    name = "date_parse"
    arg_cnt = 1
    # CLICKHOUSE
    argument_types = [ArgTypeSequence([DataType.STRING])]
    return_type = Fixed(DataType.DATE)


class FuncGeopoint(TypeConvFunction):
    name = "geopoint"
    arg_names = ["value_1", "value_2"]
    return_type = Fixed(DataType.GEOPOINT)


class FuncGeopointFromStr(FuncGeopoint):
    arg_cnt = 1
    variants = [V(D.DUMMY | D.SQLITE, lambda val: val)]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
        ArgTypeSequence([DataType.GEOPOINT]),
    ]


class FuncGeopointFromCoords(FuncGeopoint):
    arg_cnt = 2
    variants = [
        VW(D.DUMMY | D.SQLITE, lambda lat, lon: n.func.CONCAT("[", lat, ",", lon, "]")),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]


class FuncGeopolygonBase(TypeConvFunction):
    name = "geopolygon"
    return_type = Fixed(DataType.GEOPOLYGON)
    arg_cnt = 1


class FuncGeopolygon(FuncGeopolygonBase):
    variants = [V(D.DUMMY | D.SQLITE, lambda val: val)]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
        ArgTypeSequence([DataType.GEOPOLYGON]),
    ]


class FuncTreeBase(TypeConvFunction):
    arg_cnt = 1
    name = "tree"
    arg_names = ["array"]
    variants = [
        V(D.DUMMY, lambda val: val),
    ]


class FuncTreeStr(FuncTreeBase):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_STR]),
    ]
    return_type = Fixed(DataType.TREE_STR)


class DbCastArgTypes(ArgTypeSequence):
    __slots__ = ()

    def __init__(self):
        super().__init__(
            arg_types=[
                # 1. expression
                {
                    DataType.INTEGER,
                    DataType.FLOAT,
                    DataType.STRING,
                    DataType.DATE,
                    DataType.ARRAY_INT,
                    DataType.ARRAY_FLOAT,
                    DataType.ARRAY_STR,
                    # Expand as other types become supported
                },
                # 2. native_type
                DataType.CONST_STRING,
                # Followed by custom type params, but they are all handled in the translation function,
                # not passed down to the arg matcher superclass.
            ]
        )

    def match_arg_types(self, arg_types: Sequence[DataType]) -> bool:
        # Check the first two arguments using superclass
        # The rest can have any types and will be validated by the translation
        return super().match_arg_types(arg_types[: len(self._exp_arg_types)])

    def get_possible_arg_types_at_pos(self, pos: int, total: int) -> Set[DataType]:
        if pos < len(self._exp_arg_types):
            return super().get_possible_arg_types_at_pos(pos=pos, total=total)
        return set(DataType)


DataTypeSpec = Union[DataType, Tuple[DataType, ...]]


class WhitelistTypeSpec(NamedTuple):
    name: str
    sa_type: Type[TypeEngine]
    nested_sa_type: Optional[Type[TypeEngine]] = None
    arg_types: Tuple[DataTypeSpec, ...] = ()


DECIMAL_CAST_ARG_T = (DataType.CONST_INTEGER, DataType.CONST_INTEGER)
CHAR_CAST_ARG_T = (DataType.CONST_INTEGER,)


class FuncDbCastBase(TypeConvFunction):
    name = "db_cast"
    arg_names = ["expression", "native_type", "param_1", "param_2", "param_3"]
    argument_types = [
        DbCastArgTypes(),
    ]
    return_type = FromArgs(0)

    WHITELISTS: ClassVar[dict[DialectCombo, dict[DataType, list[WhitelistTypeSpec]]]] = {}

    @classmethod
    def generate_cast_type(
        cls, dialect: DialectCombo, wr_name: TranslationCtx, value: TranslationCtx, type_args: Sequence[TranslationCtx]
    ) -> TypeEngine:
        whitelist_by_type = cls.WHITELISTS[dialect]
        name = un_literal(wr_name.expression)  # type: ignore  # TODO: fix
        assert isinstance(name, str)
        name = name.upper()
        data_type = value.data_type.non_const_type  # type: ignore  # TODO: fix

        try:
            whitelist = whitelist_by_type[data_type]
        except KeyError:
            raise exc.TranslationError(f"Data type {data_type} is not supported by function {cls.name.upper()}")

        # Find the correct spec by data type name
        for each_spec in whitelist:
            if each_spec.name.upper() == name:
                type_spec = each_spec
                break
        else:
            raise exc.TranslationError(f"Native data type {name} is not supported by function {cls.name.upper()}")

        if len(type_args) != len(type_spec.arg_types):
            raise exc.TranslationError(
                f"Native type cast to {name} requires {type_spec.arg_types}, got {len(type_args)}"
            )

        def _raise_type_mismath(_arg_num: int):
            dtype = type_args[_arg_num].data_type
            assert dtype is not None
            raise exc.TranslationError(
                f"Argument {_arg_num+1} of native type cast to {name} " f"has invalid type {dtype.name}"
            )

        unwrapped_native_type_args = []
        for arg_num, (arg, req_arg_type) in enumerate(zip(type_args, type_spec.arg_types)):
            if isinstance(req_arg_type, DataType):
                if not arg.data_type.casts_to(req_arg_type):  # type: ignore  # TODO: fix
                    _raise_type_mismath(arg_num)
            elif isinstance(req_arg_type, tuple):
                for sub_req_arg_type in req_arg_type:
                    if not arg.data_type.casts_to(sub_req_arg_type):  # type: ignore  # TODO: fix
                        _raise_type_mismath(arg_num)
            else:
                raise TypeError(type(req_arg_type))

            unwrapped_native_type_args.append(un_literal(arg.expression))  # type: ignore  # TODO: fix

        # All validation checks are done
        # Create the type and make the cast
        if type_spec.nested_sa_type is not None:
            cast_to_type = type_spec.sa_type(type_spec.nested_sa_type(*unwrapped_native_type_args))
        else:
            cast_to_type = type_spec.sa_type(*unwrapped_native_type_args)

        return cast_to_type

    @classmethod
    def apply_cast_wrapper(
        cls,
        dialect: DialectCombo,
        expression: ClauseElement,
        type_: TypeEngine,
    ) -> ClauseElement:
        return sa.cast(expression, type_)

    @classmethod
    def perform_cast(
        cls, dialect: DialectCombo, wr_name: TranslationCtx, value: TranslationCtx, type_args: Sequence[TranslationCtx]
    ) -> ClauseElement:
        type_ = cls.generate_cast_type(dialect=dialect, wr_name=wr_name, value=value, type_args=type_args)
        assert value.expression is not None
        return cls.apply_cast_wrapper(dialect=dialect, expression=value.expression, type_=type_)

    def __init__(self):
        super().__init__()

        self._inst_variants = [
            VW(
                dialects=dialect,
                translation=(
                    lambda value, wr_name, *args, _dialect=dialect: self.perform_cast(
                        dialect=_dialect,
                        wr_name=wr_name,
                        value=value,
                        type_args=args,
                    )
                ),
            )
            for dialect in self.WHITELISTS.keys()
        ]


class FuncDbCast2(FuncDbCastBase):
    arg_cnt = 2


class FuncDbCast3(FuncDbCastBase):
    arg_cnt = 3


class FuncDbCast4(FuncDbCastBase):
    arg_cnt = 4


DEFINITIONS_TYPE = [
    # bool
    FuncBoolFromNull,
    FuncBoolFromNumber,
    FuncBoolFromBool,
    FuncBoolFromStrGeo,
    FuncBoolFromDateDatetime,
    # date
    FuncDate1FromNull,
    FuncDate1FromDatetime,
    FuncDate1FromDatetimeTZ,
    FuncDate1FromString,
    FuncDate1FromNumber,
    # date_parse
    FuncDateParse,  # only CLICKHOUSE
    # datetime
    FuncDatetime1FromNull,
    FuncDatetime1FromDatetime,
    FuncDatetime1FromDate,
    FuncDatetime1FromNumber,
    FuncDatetime1FromString,
    # datetime_parse
    FuncDatetimeParse,  # only CLICKHOUSE
    # datetimetz
    FuncDatetimeTZConst,
    # datetimetz_to_naive
    # -> only for some dialects
    # db_cast
    FuncDbCast2,
    FuncDbCast3,
    FuncDbCast4,
    # float
    FuncFloatNumber,
    FuncFloatString,
    FuncFloatFromBool,
    FuncFloatFromDate,
    FuncFloatFromDatetime,
    FuncFloatFromGenericDatetime,
    # genericdatetime
    FuncGenericDatetime1FromNull,
    FuncGenericDatetime1FromDatetime,
    FuncGenericDatetime1FromDate,
    FuncGenericDatetime1FromNumber,
    FuncGenericDatetime1FromString,
    # genericdatetime_parse
    FuncGenericDatetimeParse,  # only CLICKHOUSE
    # geopoint
    FuncGeopointFromStr,
    FuncGeopointFromCoords,
    # geopolygon
    FuncGeopolygon,
    # int
    FuncIntFromNull,
    FuncIntFromInt,
    FuncIntFromFloat,
    FuncIntFromBool,
    FuncIntFromStr,
    FuncIntFromDate,
    FuncIntFromDatetime,
    FuncIntFromGenericDatetime,
    FuncIntFromDatetimeTZ,
    # str
    FuncStrFromNull,
    FuncStrFromUnsupported,
    FuncStrFromInteger,
    FuncStrFromFloat,
    FuncStrFromBool,
    FuncStrFromStrGeo,
    FuncStrFromDate,
    FuncStrFromDatetime,
    FuncStrFromDatetimeTZ,
    FuncStrFromString,
    FuncStrFromUUID,
    FuncStrFromArray,
    # tree
    FuncTreeStr,
]

lib/dl_formula/dl_formula/definitions/functions_type.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Optional,
)

import sqlalchemy as sa
from sqlalchemy.sql.elements import (
    ClauseElement,
    ClauseList,
)
from sqlalchemy.sql.functions import Function as SAFunction

from dl_formula.core import exc
from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    Function,
    TranslationVariant,
    WinRangeTuple,
)
from dl_formula.definitions.common import desc
from dl_formula.definitions.literals import (
    Literal,
    is_literal,
    un_literal,
)
from dl_formula.definitions.type_strategy import (
    Fixed,
    FromArgs,
)


V = TranslationVariant.make

NUMERIC_TYPES = {DataType.INTEGER, DataType.FLOAT}
RANKED_TYPES = {
    DataType.BOOLEAN,
    DataType.INTEGER,
    DataType.FLOAT,
    DataType.STRING,
    DataType.DATE,
    DataType.DATETIME,
    DataType.GENERICDATETIME,
    DataType.UUID,
}
ALL_TYPES = frozenset(DataType)


class WindowFunction(Function):
    is_function = True
    is_window = True
    supports_grouping = True
    supports_bfb = True


def _order_by_from_args(*args: ClauseElement) -> ClauseList:
    """Defined the ORDER BY clause that is generated from function arguments"""
    if len(args) not in (1, 2):
        raise ValueError(f"Invalid number of arguments: {len(args)}")

    param = args[0]
    if len(args) == 2:
        assert is_literal(args[1])
        dir_value = args[1].value.lower()
    else:
        # Default ranking is from greatest to least
        dir_value = "desc"

    if dir_value in ("asc", "ascending"):
        pass
    elif dir_value in ("desc", "descending"):
        param = desc(param)
    else:
        raise ValueError(dir_value)

    return ClauseList(param)


def _rows_full_window(*_: Any) -> WinRangeTuple:
    """Defines range tuple for functions that need the whole window."""
    return None, None


def _rows_stretching_window(_: ClauseElement, direction: ClauseElement = sa.literal("asc")) -> WinRangeTuple:
    """Defines range tuple for functions that need a stretching or shrinking window."""
    assert is_literal(direction)
    assert isinstance(un_literal(direction), str)
    dir_value = direction.value.lower()
    if dir_value == "asc":
        return None, 0
    elif dir_value == "desc":
        return 0, None
    raise exc.TranslationError("Invalid value for direction parameter")


class WinGenericRankBase(WindowFunction):
    """
    Base class for all RANK* window functions.

    Some translation examples:

    RANK([Profit], "asc")  ->

        RANK() OVER (
            ORDER BY "Profit"
            ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
        )

    RANK([Profit] TOTAL)  ->

        RANK() OVER (
            ORDER BY "Profit" DESC
            ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
        )

    RANK([Profit] WITHIN [Country])  ->

        RANK() OVER (
            PARTITION BY "Country"
            ORDER BY "Profit" DESC
            ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
        )

    RANK([Profit] AMONG [Date])  ->

        RANK() OVER (
            PARTITION BY <all dimensions except "Date">
            ORDER BY "Profit" DESC
            ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
        )
    """

    arg_names = ["value", "direction"]
    return_type = Fixed(DataType.INTEGER)


class WinRankBase(WinGenericRankBase):
    name = "rank"
    variants = [
        V(
            D.DUMMY,
            translation=lambda *args: sa.func.RANK(),
            translation_order_by=_order_by_from_args,
            as_winfunc=True,
        ),
    ]


class WinRank1(WinRankBase):
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([RANKED_TYPES]),
    ]


class WinRank2(WinRankBase):
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([RANKED_TYPES, DataType.CONST_STRING]),
    ]


class WinRankDenseBase(WinGenericRankBase):
    name = "rank_dense"
    variants = [
        V(
            D.DUMMY,
            translation=lambda *args: sa.func.DENSE_RANK(),
            translation_order_by=_order_by_from_args,
            as_winfunc=True,
        ),
    ]


class WinRankDense1(WinRankDenseBase):
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([RANKED_TYPES]),
    ]


class WinRankDense2(WinRankDenseBase):
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([RANKED_TYPES, DataType.CONST_STRING]),
    ]


class WinRankUniqueBase(WinGenericRankBase):
    name = "rank_unique"
    variants = [
        V(
            D.DUMMY,
            translation=lambda *args: sa.func.ROW_NUMBER(),
            translation_order_by=_order_by_from_args,
            as_winfunc=True,
        ),
    ]


class WinRankUnique1(WinRankUniqueBase):
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([RANKED_TYPES]),
    ]


class WinRankUnique2(WinRankUniqueBase):
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([RANKED_TYPES, DataType.CONST_STRING]),
    ]


class WinRankPercentileBase(WinGenericRankBase):
    name = "rank_percentile"
    return_type = Fixed(DataType.FLOAT)
    variants = [
        V(
            D.DUMMY,
            translation=lambda *args: sa.func.PERCENT_RANK(),
            translation_order_by=_order_by_from_args,
            as_winfunc=True,
        ),
    ]


class WinRankPercentile1(WinRankPercentileBase):
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([RANKED_TYPES]),
    ]


class WinRankPercentile2(WinRankPercentileBase):
    arg_cnt = 2
    argument_types = [
        ArgTypeSequence([RANKED_TYPES, DataType.CONST_STRING]),
    ]


class WinAggBase(WindowFunction):
    """
    Base for all window functions based on "classical" aggregation functions.

    Translation examples:

    SUM([Profit] TOTAL)  ->

        SUM("Profit") OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)

    SUM([Profit] WITHIN [Date])  ->

        SUM("Profit") OVER (
            PARTITION BY "Date"
            ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
        )

    SUM([Profit] AMONG [Date])  ->

        SUM("Profit") OVER (
            PARTITION BY <all dimensions except "Date">
            ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
        )
    """


class WinSum(WinAggBase):
    name = "sum"
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([NUMERIC_TYPES]),
    ]
    variants = [
        V(
            D.DUMMY,
            translation=lambda x: sa.func.SUM(x),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]
    return_type = FromArgs(0)


class WinMinMaxBase(WinAggBase):
    arg_cnt = 1
    argument_types = [
        ArgTypeSequence([RANKED_TYPES]),
    ]
    return_type = FromArgs(0)


class WinMin(WinMinMaxBase):
    name = "min"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x: sa.func.MIN(x),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]


class WinMax(WinMinMaxBase):
    name = "max"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x: sa.func.MAX(x),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]


class WinCountBase(WinAggBase):
    name = "count"
    return_type = Fixed(DataType.INTEGER)


class WinCount0(WinCountBase):
    arg_cnt = 0
    variants = [
        V(
            D.DUMMY,
            translation=lambda: sa.func.COUNT(sa.text("*")),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]


class WinCount1(WinCountBase):
    arg_cnt = 1
    variants = [
        V(
            D.DUMMY,
            translation=lambda x: sa.func.COUNT(x),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]


class WinAvg(WinAggBase):  # TODO: support DATE & DATETIME
    name = "avg"
    return_type = Fixed(DataType.FLOAT)
    arg_cnt = 1
    variants = [
        V(
            D.DUMMY,
            translation=lambda x: sa.func.AVG(x),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]


class WinAggIf(WindowFunction):
    """
    Base class for "classical" conditional window functions

    Translation examples:

    SUM_IF([Profit], [Category] = 'Office Supplies' TOTAL)  ->

        SUM("Profit") FILTER (WHERE "Category" = 'Office Supplies')
        OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)

    SUM_IF([Profit], [Category] = 'Office Supplies' WITHIN [Date])  ->

        SUM("Profit") FILTER (WHERE "Category" = 'Office Supplies')
        OVER (
            PARTITION BY "Date"
            ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
        )

    SUM_IF([Profit], [Category] = 'Office Supplies' AMONG [Date])  ->
        SUM("Profit") FILTER (WHERE "Category" = 'Office Supplies')
        OVER (
            PARTITION BY <all dimensions except "Date">
            ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
        )
    """

    arg_cnt = 2
    arg_names = ["expression", "condition"]
    argument_types = [
        ArgTypeSequence([NUMERIC_TYPES, DataType.BOOLEAN]),
    ]


class WinSumIf(WinAggIf):
    name = "sum_if"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, condition: sa.func.SUM(x).filter(condition),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]
    return_type = FromArgs(0)


class WinCountIf(WinAggIf):
    name = "count_if"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, condition: sa.func.COUNT(x).filter(condition),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]
    return_type = FromArgs(0)


class WinAvgIf(WinAggIf):  # TODO: support DATE & DATETIME
    name = "avg_if"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, condition: sa.func.AVG(x).filter(condition),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]
    return_type = FromArgs(0)


class OrderedWinFuncBase(WindowFunction):
    uses_default_ordering = True
    supports_ordering = True


class WinRFuncBase(OrderedWinFuncBase):
    """
    Base class for cumulative window functions.

    Translation examples:

    RSUM([Profit])  ->

        SUM("Profit") OVER (
            ORDER BY <same as query's main ORDER BY clause>
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        )

    RSUM([Profit], "desc" WITHIN [Date])  ->

        SUM("Profit") OVER (
            PARTITION BY "Date"
            ORDER BY <same as query's main ORDER BY clause>
            ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING
        )

     RSUM([Profit] AMONG [Date])  ->

        SUM("Profit") OVER (
            PARTITION BY <all dimensions except "Date">
            ORDER BY <same as query's main ORDER BY clause>
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        )
    """

    arg_names = ["value", "direction"]


class WinRSumBase(WinRFuncBase):
    name = "rsum"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, *_: sa.func.sum(x),
            translation_rows=_rows_stretching_window,
            as_winfunc=True,
        ),
    ]
    argument_types = [
        ArgTypeSequence([NUMERIC_TYPES, DataType.CONST_STRING]),
    ]
    return_type = FromArgs(0)


class WinRSum1(WinRSumBase):
    arg_cnt = 1


class WinRSum2(WinRSumBase):
    arg_cnt = 2


class WinRCountBase(WinRFuncBase):
    name = "rcount"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, *_: sa.func.COUNT(x),
            translation_rows=_rows_stretching_window,
            as_winfunc=True,
        ),
    ]
    argument_types = [
        ArgTypeSequence([ALL_TYPES, DataType.CONST_STRING]),
    ]
    return_type = Fixed(DataType.INTEGER)


class WinRCount1(WinRCountBase):
    arg_cnt = 1


class WinRCount2(WinRCountBase):
    arg_cnt = 2


class WinRMinMaxBase(WinRFuncBase):
    argument_types = [
        ArgTypeSequence([RANKED_TYPES, DataType.CONST_STRING]),
    ]
    return_type = FromArgs(0)


class WinRMinBase(WinRMinMaxBase):
    name = "rmin"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, *_: sa.func.MIN(x),
            translation_rows=_rows_stretching_window,
            as_winfunc=True,
        ),
    ]


class WinRMin1(WinRMinBase):
    arg_cnt = 1


class WinRMin2(WinRMinBase):
    arg_cnt = 2


class WinRMaxBase(WinRMinMaxBase):
    name = "rmax"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, *_: sa.func.MAX(x),
            translation_rows=_rows_stretching_window,
            as_winfunc=True,
        ),
    ]


class WinRMax1(WinRMaxBase):
    arg_cnt = 1


class WinRMax2(WinRMaxBase):
    arg_cnt = 2


class WinRAvgBase(WinRFuncBase):  # TODO: Add support for DATE & DATETIME
    name = "ravg"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, *_: sa.func.AVG(x),
            translation_rows=_rows_stretching_window,
            as_winfunc=True,
        ),
    ]
    argument_types = [
        ArgTypeSequence([NUMERIC_TYPES, DataType.CONST_STRING]),
    ]
    return_type = Fixed(DataType.INTEGER)


class WinRAvg1(WinRAvgBase):
    arg_cnt = 1


class WinRAvg2(WinRAvgBase):
    arg_cnt = 2


def _resolve_mfunc_rows(raw_arg1: Literal, raw_arg2: Optional[Literal] = None) -> WinRangeTuple:
    """
    Create range tuple for the following cases:
    ::

        MSUM(SUM([Sales]), 5)

    Calculates the sum of the current row and the previous 5
    ::

        MSUM(SUM([Sales]), -5)

    Calculates the sum of the current row and the following 5
    ::

        MSUM(SUM([Sales]), 3, 8)

    Calculates the sum of the previous 3, the current row and the following 8
    """
    arg1 = raw_arg1.value
    arg2 = raw_arg2.value if raw_arg2 is not None else None

    if arg2 is not None:
        return (-abs(arg1), abs(arg2))
    if arg1 >= 0:
        return (-arg1, 0)
    return (0, abs(arg1))


class MFuncBase(OrderedWinFuncBase):
    """
    Base class for moving window functions

    MSUM([Profit], -2)  ->

        SUM("Profit") OVER (
            ORDER BY <same as query's main ORDER BY clause>
            ROWS BETWEEN CURRENT ROW AND 2 FOLLOWING
        )

    MSUM([Profit], 3)  ->

        SUM("Profit") OVER (
            ORDER BY <same as query's main ORDER BY clause>
            ROWS BETWEEN 3 PRECEDING AND CURRENT ROW
        )

    MSUM([Profit] 5, 5 TOTAL)  ->

        SUM("Profit") OVER (
            ORDER BY <same as query's main ORDER BY clause>
            ROWS BETWEEN 5 PRECEDING AND 5 FOLLOWING
        )

    MSUM([Profit], -5 WITHIN [Date])  ->

        SUM("Profit") OVER (
            PARTITION BY "Date"
            ORDER BY <same as query's main ORDER BY clause>
            ROWS BETWEEN CURRENT ROW AND 5 FOLLOWING
        )

    MSUM([Profit], -5 AMONG [Date])  ->

        SUM("Profit") OVER (
            PARTITION BY <all dimensions except "Date">
            ORDER BY <same as query's main ORDER BY clause>
            ROWS BETWEEN CURRENT ROW AND 5 FOLLOWING
        )
    """

    arg_names = ["value", "rows_1", "rows_2"]


class WinMSumBase(MFuncBase):
    name = "msum"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, *_: sa.func.sum(x),
            translation_rows=lambda _, *row_args: _resolve_mfunc_rows(*row_args),
            as_winfunc=True,
        ),
    ]
    argument_types = [
        ArgTypeSequence([NUMERIC_TYPES, DataType.CONST_INTEGER, DataType.CONST_INTEGER]),
    ]
    return_type = FromArgs(0)


class WinMSum2(WinMSumBase):
    arg_cnt = 2


class WinMSum3(WinMSumBase):
    arg_cnt = 3


class WinMCountBase(MFuncBase):
    name = "mcount"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, *_: sa.func.COUNT(x),
            translation_rows=lambda _, *row_args: _resolve_mfunc_rows(*row_args),
            as_winfunc=True,
        ),
    ]
    argument_types = [
        ArgTypeSequence([ALL_TYPES, DataType.CONST_INTEGER, DataType.CONST_INTEGER]),
    ]
    return_type = FromArgs(0)


class WinMCount2(WinMCountBase):
    arg_cnt = 2


class WinMCount3(WinMCountBase):
    arg_cnt = 3


class WinMMinMaxBase(MFuncBase):
    argument_types = [
        ArgTypeSequence([RANKED_TYPES, DataType.CONST_INTEGER, DataType.CONST_INTEGER]),
    ]
    return_type = FromArgs(0)


class WinMMinBase(WinMMinMaxBase):
    name = "mmin"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, *_: sa.func.MIN(x),
            translation_rows=lambda _, *row_args: _resolve_mfunc_rows(*row_args),
            as_winfunc=True,
        ),
    ]


class WinMMin2(WinMMinBase):
    arg_cnt = 2


class WinMMin3(WinMMinBase):
    arg_cnt = 3


class WinMMaxBase(WinMMinMaxBase):
    name = "mmax"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, *_: sa.func.MAX(x),
            translation_rows=lambda _, *row_args: _resolve_mfunc_rows(*row_args),
            as_winfunc=True,
        ),
    ]


class WinMMax2(WinMMaxBase):
    arg_cnt = 2


class WinMMax3(WinMMaxBase):
    arg_cnt = 3


class WinMAvgBase(MFuncBase):  # TODO: support DATE & DATETIME
    name = "mavg"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x, *_: sa.func.AVG(x),
            translation_rows=lambda _, *row_args: _resolve_mfunc_rows(*row_args),
            as_winfunc=True,
        ),
    ]
    argument_types = [
        ArgTypeSequence([NUMERIC_TYPES, DataType.CONST_INTEGER, DataType.CONST_INTEGER]),
    ]
    return_type = FromArgs(0)


class WinMAvg2(WinMAvgBase):
    arg_cnt = 2


class WinMAvg3(WinMAvgBase):
    arg_cnt = 3


def lag_implementation(
    x: Any,
    offset: Any = sa.literal(1),
    default: Any = sa.null(),
    lag_name: str = "LAG",
    lead_name: str = "LEAD",
) -> ClauseElement:
    result: ClauseElement
    if un_literal(offset) > 0:
        result = SAFunction(lag_name, x, un_literal(offset), un_literal(default))
    elif un_literal(offset) == 0:
        result = x
    else:
        result = SAFunction(lead_name, x, sa.literal(-un_literal(offset)), un_literal(default))
    return result


class WinLagBase(OrderedWinFuncBase):
    name = "lag"
    arg_names = ["value", "offset", "default"]
    variants = [
        V(
            D.DUMMY,
            # PostgreSQL requires offset and default to be constant,
            # so we must unwrap these values in case casts have been added to them.
            translation=lambda x, offset=sa.literal(1), default=sa.null(), *_: (
                lag_implementation(x, offset=offset, default=default)
            ),
            translation_rows=lambda x, offset=sa.literal(1), *_: (None, None),
            as_winfunc=True,
        ),
    ]
    argument_types = [
        ArgTypeSequence([set(DataType), DataType.CONST_INTEGER, {dtype for dtype in DataType if dtype.is_const}]),
    ]
    return_type = FromArgs(0)


class WinLag1(WinLagBase):
    arg_cnt = 1


class WinLag2(WinLagBase):
    arg_cnt = 2


class WinLag3(WinLagBase):
    arg_cnt = 3


class WinFirstLastBase(OrderedWinFuncBase):
    arg_cnt = 1
    arg_names = ["value"]
    argument_types = [
        ArgTypeSequence(
            [
                set(DataType),
            ]
        ),
    ]
    return_type = FromArgs(0)


class WinFirst(WinFirstLastBase):
    name = "first"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x: sa.func.first_value(x),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]


class WinLast(WinFirstLastBase):
    name = "last"
    variants = [
        V(
            D.DUMMY,
            translation=lambda x: sa.func.last_value(x),
            translation_rows=_rows_full_window,
            as_winfunc=True,
        ),
    ]


DEFINITIONS_WINDOW = [
    # avg
    WinAvg,
    # avg_if
    WinAvgIf,
    # count
    WinCount0,
    WinCount1,
    # count_if
    WinCountIf,
    # first
    WinFirst,
    # lag
    WinLag1,
    WinLag2,
    WinLag3,
    # last
    WinLast,
    # mavg
    WinMAvg2,
    WinMAvg3,
    # max
    WinMax,
    # mcount
    WinMCount2,
    WinMCount3,
    # min
    WinMin,
    # mmax
    WinMMax2,
    WinMMax3,
    # mmin
    WinMMin2,
    WinMMin3,
    # msum
    WinMSum2,
    WinMSum3,
    # rank
    WinRank1,
    WinRank2,
    # rank_dense
    WinRankDense1,
    WinRankDense2,
    # rank_percentile
    WinRankPercentile1,
    WinRankPercentile2,
    # rank_unique
    WinRankUnique1,
    WinRankUnique2,
    # ravg
    WinRAvg1,
    WinRAvg2,
    # rcount
    WinRCount1,
    WinRCount2,
    # rmax
    WinRMax1,
    WinRMax2,
    # rmin
    WinRMin1,
    WinRMin2,
    # rsum
    WinRSum1,
    WinRSum2,
    # sum
    WinSum,
    # sum_if
    WinSumIf,
]

lib/dl_formula/dl_formula/definitions/functions_window.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import lru_cache
from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
    Type,
)

import sqlalchemy as sa
import sqlalchemy.dialects.postgresql as sa_postgresql
from sqlalchemy.sql.elements import Null
from sqlalchemy.sql.type_api import TypeEngine

from dl_formula.connectors.base.literal import (
    Literal,
    Literalizer,
    TypeDefiningCast,
)
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.nodes import BaseLiteral
from dl_formula.utils.datetime import make_datetime_value


if TYPE_CHECKING:
    from dl_formula.translation.context import TranslationCtx


_LITERALIZER_REGISTRY: dict[DialectCombo, Literalizer] = {}


@lru_cache
def get_literalizer(dialect: DialectCombo) -> Literalizer:
    sorted_dialect_items = sorted(
        _LITERALIZER_REGISTRY.items(), key=lambda el: el[0].ambiguity
    )  # Most specific dialects go first, the most ambiguous ones go last
    for d, literalizer in sorted_dialect_items:
        if d & dialect == dialect:
            return literalizer
    return Literalizer()  # FIXME: remove fallback after all dialects become connectors


def register_literalizer(dialect: DialectCombo, literalizer_cls: Type[Literalizer]) -> None:
    if dialect in _LITERALIZER_REGISTRY:
        assert isinstance(_LITERALIZER_REGISTRY[dialect], literalizer_cls)
        return

    _LITERALIZER_REGISTRY[dialect] = literalizer_cls()


def literal(value: Any, type_: Optional[TypeEngine] = None, d: Optional[DialectCombo] = None) -> Literal:
    """A fine-tuned version of ``sqlalchemy.literal``"""
    dialect = d
    if type_ is None and dialect is not None:
        literalizer = get_literalizer(dialect=dialect)
        return literalizer.literal(value, dialect=dialect)

    return sa.literal(value, type_)


def un_literal(sa_obj: Optional[Literal], value_ctx: Optional[TranslationCtx] = None):
    """Given an SA object produced by `literal` (or compatible), return the original value"""

    assert sa_obj is not None

    if value_ctx is not None and isinstance(value_ctx.node, BaseLiteral):
        return value_ctx.node.value

    if isinstance(sa_obj, sa.sql.elements.BindParameter):
        return sa_obj.value

    if isinstance(sa_obj, TypeDefiningCast):
        return sa_obj.clause.value

    if isinstance(sa_obj, Null):
        return None

    if isinstance(sa_obj, sa_postgresql.array):
        return [un_literal(clause) for clause in sa_obj.clauses]

    # More complex expressions via functions
    if isinstance(sa_obj, sa.sql.functions.Function):
        clauses = sa_obj.clauses.clauses
        if all(isinstance(clause, (sa.sql.elements.BindParameter, Null)) for clause in clauses):
            if sa_obj.name == "toDateTime":
                # CH datetime un-wrap
                if len(clauses) == 2:
                    dt_str = un_literal(clauses[0])
                    return make_datetime_value(dt_str)
            if sa_obj.name == "array":
                return [un_literal(clause) for clause in clauses]

    raise ValueError(f"Not a literal-like SA object: {sa_obj!r}")


def is_literal(sa_obj) -> bool:
    try:
        un_literal(sa_obj)
    except ValueError:
        return False
    return True

lib/dl_formula/dl_formula/definitions/literals.py
datalens-tech/datalens-backend
# TODO: Remove the `ll_` prefix

from itertools import chain
from typing import (
    Type,
    Union,
)

import attr

from dl_formula.definitions.base import NodeTranslation
from dl_formula.definitions.conditional_blocks import DEFINITIONS_COND_BLOCKS
from dl_formula.definitions.functions_aggregation import DEFINITIONS_AGGREGATION
from dl_formula.definitions.functions_array import DEFINITIONS_ARRAY
from dl_formula.definitions.functions_datetime import DEFINITIONS_DATETIME
from dl_formula.definitions.functions_logical import DEFINITIONS_LOGICAL
from dl_formula.definitions.functions_markup import DEFINITIONS_MARKUP
from dl_formula.definitions.functions_math import DEFINITIONS_MATH
from dl_formula.definitions.functions_special import DEFINITIONS_SPECIAL
from dl_formula.definitions.functions_string import DEFINITIONS_STRING
from dl_formula.definitions.functions_time_series import DEFINITIONS_TIME_SERIES
from dl_formula.definitions.functions_type import DEFINITIONS_TYPE
from dl_formula.definitions.functions_window import DEFINITIONS_WINDOW
from dl_formula.definitions.operators_binary import DEFINITIONS_BINARY
from dl_formula.definitions.operators_ternary import DEFINITIONS_TERNARY
from dl_formula.definitions.operators_unary import DEFINITIONS_UNARY
from dl_formula.definitions.registry import OPERATION_REGISTRY
from dl_formula.inspect.registry.item import BasicOpItem
from dl_formula.inspect.registry.registry import (
    LOWLEVEL_OP_REGISTRY,
    norm_name,
)


@attr.s
class _Flag:
    is_set: bool = attr.ib(default=False)


REGISTERED_FLAG = _Flag()


def populate_translation_registry() -> None:
    if REGISTERED_FLAG.is_set:
        # Already loaded
        return

    definitions = chain(
        DEFINITIONS_COND_BLOCKS,
        DEFINITIONS_AGGREGATION,
        DEFINITIONS_ARRAY,
        DEFINITIONS_DATETIME,
        DEFINITIONS_LOGICAL,
        DEFINITIONS_MARKUP,
        DEFINITIONS_MATH,
        DEFINITIONS_SPECIAL,
        DEFINITIONS_STRING,
        DEFINITIONS_TIME_SERIES,
        DEFINITIONS_TYPE,
        DEFINITIONS_WINDOW,
        DEFINITIONS_BINARY,
        DEFINITIONS_TERNARY,
        DEFINITIONS_UNARY,
    )
    def_item: Union[NodeTranslation, Type[NodeTranslation]]
    # TODO: Load defs from connectors here
    for def_item in definitions:
        if isinstance(def_item, type):
            def_item = def_item()
        assert isinstance(def_item, NodeTranslation)
        OPERATION_REGISTRY.register(def_item)

    REGISTERED_FLAG.is_set = True


def populate_ll_op_registry() -> None:
    for _key, impl_cls in OPERATION_REGISTRY.items():
        name = norm_name(impl_cls.name)
        assert name is not None
        ll_item = BasicOpItem(
            name=name,
            is_window=impl_cls.is_window,
            arg_cnt=impl_cls.arg_cnt,
            is_aggregate=impl_cls.is_aggregation,
            is_function=impl_cls.is_function,
            supports_grouping=impl_cls.supports_grouping,
            supports_ordering=impl_cls.supports_ordering,
            supports_ignore_dimensions=impl_cls.supports_ignore_dimensions,
            supports_bfb=impl_cls.supports_bfb,
            supports_lod=impl_cls.supports_lod,
            uses_default_ordering=impl_cls.uses_default_ordering,
        )
        LOWLEVEL_OP_REGISTRY.add(ll_item)

lib/dl_formula/dl_formula/definitions/ll_registry_loader.py
datalens-tech/datalens-backend
from __future__ import annotations

import functools
from typing import TYPE_CHECKING

import sqlalchemy as sa

from dl_formula.core import exc
from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import (
    ArgFlagSequence,
    ArgTypeMatcher,
    ArgTypeSequence,
)
from dl_formula.definitions.base import (
    ArgTransformer,
    MultiVariantTranslation,
    TranslationVariant,
    TranslationVariantWrapped,
)
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.literals import (
    literal,
    un_literal,
)
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import (
    Fixed,
    FromArgs,
)
from dl_formula.shortcuts import n


if TYPE_CHECKING:
    from dl_formula.translation.context import TranslationCtx
    from dl_formula.translation.env import TranslationEnvironment


V = TranslationVariant.make
VW = TranslationVariantWrapped.make


def as_bigint(value):
    return sa.cast(value, sa.BIGINT())


class Binary(MultiVariantTranslation):
    arg_cnt = 2
    is_function = False
    arg_names = ["left", "right"]


class BinaryPower(Binary):
    name = "^"
    arg_names = ["base", "power"]
    variants = [VW(D.DUMMY, lambda base, power: n.func.POWER(base, power))]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.FLOAT)


class BinaryMult(Binary):
    name = "*"
    arg_names = ["value_1", "value_2"]


class BinaryMultNumbers(BinaryMult):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda left, right: left * right),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]
    return_type = FromArgs()


class BinaryMultStringConst(BinaryMult):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda text, size: literal(un_literal(text) * un_literal(size))),
    ]
    argument_types = [
        ArgTypeSequence([DataType.CONST_STRING, DataType.CONST_INTEGER]),
        ArgTypeSequence([DataType.CONST_INTEGER, DataType.CONST_STRING]),
    ]
    return_type = Fixed(DataType.CONST_STRING)


class BinaryMultStringNonConst(BinaryMult):
    class BinaryMultStringNonConstArgTransformer(ArgTransformer):
        def transform_args(
            self, env: TranslationEnvironment, args: list[TranslationCtx], arg_types: list[DataType]
        ) -> list[TranslationCtx]:
            """Reorder args so that str is always first and number of repeats is second"""
            if arg_types[0] not in (DataType.CONST_STRING, DataType.STRING):
                args = [args[1], args[0]]
            return args

    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.INTEGER]),
        ArgTypeSequence([DataType.INTEGER, DataType.STRING]),
    ]
    return_type = Fixed(DataType.STRING)
    arg_transformer = BinaryMultStringNonConstArgTransformer()


class BinaryDiv(Binary):
    name = "/"
    return_type = Fixed(DataType.FLOAT)
    arg_names = ["number_1", "number_2"]


class BinaryDivInt(BinaryDiv):
    variants = [
        V(D.DUMMY, lambda x, y: x / y),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER, DataType.INTEGER]),
    ]


class BinaryDivFloat(BinaryDiv):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda x, y: x / y),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]


class BinaryMod(Binary):
    name = "%"
    arg_names = ["number_1", "number_2"]


class BinaryModInteger(BinaryMod):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda left, right: left % right),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER, DataType.INTEGER]),
    ]
    return_type = FromArgs()


class BinaryModFloat(BinaryMod):
    variants = [
        V(D.DUMMY, lambda left, right: left % right),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]
    return_type = FromArgs()


class BinaryPlus(Binary):
    name = "+"
    arg_names = ["value_1", "value_2"]


class BinaryPlusNumbers(BinaryPlus):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda left, right: left + right),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]
    return_type = FromArgs()


class BinaryPlusStrings(BinaryPlus):
    variants = [
        # `[str1] + [str2]` is equivalent to `concat([str1], [str2])`
        VW(D.DUMMY | D.SQLITE, n.func.CONCAT),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]
    return_type = Fixed(DataType.STRING)


class BinaryPlusDateInt(BinaryPlus):
    class BinaryPlusDateIntArgTransformer(ArgTransformer):
        def transform_args(
            self, env: TranslationEnvironment, args: list[TranslationCtx], arg_types: list[DataType]
        ) -> list[TranslationCtx]:
            """Reorder args so that date is always first and number of days is second"""
            if arg_types[0] not in (DataType.CONST_DATE, DataType.DATE):
                args = [args[1], args[0]]
            return args

    arg_transformer = BinaryPlusDateIntArgTransformer()
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.INTEGER]),
        ArgTypeSequence([DataType.INTEGER, DataType.DATE]),
    ]
    return_type = Fixed(DataType.DATE)


class BinaryPlusArray(BinaryPlus):
    argument_types = [
        ArgTypeSequence([DataType.ARRAY_INT, DataType.ARRAY_INT]),
        ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.ARRAY_FLOAT]),
        ArgTypeSequence([DataType.ARRAY_STR, DataType.ARRAY_STR]),
    ]
    return_type = FromArgs()


class BinaryPlusDateFloat(BinaryPlus):
    class BinaryPlusDateFloatArgTransformer(ArgTransformer):
        def transform_args(
            self, env: TranslationEnvironment, args: list[TranslationCtx], arg_types: list[DataType]
        ) -> list[TranslationCtx]:
            """Reorder args so that date is always first and number of days is second"""
            if arg_types[0] not in (DataType.CONST_DATE, DataType.DATE):
                args = [args[1], args[0]]
            return args

    arg_transformer = BinaryPlusDateFloatArgTransformer()
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.FLOAT]),
        ArgTypeSequence([DataType.FLOAT, DataType.DATE]),
    ]
    return_type = Fixed(DataType.DATE)


class BinaryPlusDatetimeNumber(BinaryPlus):
    class BinaryPlusDatetimeNumberArgTransformer(ArgTransformer):
        def transform_args(
            self, env: TranslationEnvironment, args: list[TranslationCtx], arg_types: list[DataType]
        ) -> list[TranslationCtx]:
            """Reorder args so that date is always first and number of days is second"""
            if arg_types[0] not in (DataType.CONST_DATETIME, DataType.DATETIME):
                args = [args[1], args[0]]
            return args

    arg_transformer = BinaryPlusDatetimeNumberArgTransformer()
    argument_types = [
        ArgTypeSequence([DataType.DATETIME, DataType.FLOAT]),
        ArgTypeSequence([DataType.FLOAT, DataType.DATETIME]),
    ]
    return_type = Fixed(DataType.DATETIME)


class BinaryPlusGenericDatetimeNumber(BinaryPlus):
    class BinaryPlusGenericDatetimeNumberArgTransformer(ArgTransformer):
        def transform_args(
            self, env: TranslationEnvironment, args: list[TranslationCtx], arg_types: list[DataType]
        ) -> list[TranslationCtx]:
            """Reorder args so that date is always first and number of days is second"""
            if arg_types[0] not in (DataType.CONST_GENERICDATETIME, DataType.GENERICDATETIME):
                args = [args[1], args[0]]
            return args

    arg_transformer = BinaryPlusGenericDatetimeNumberArgTransformer()
    argument_types = [
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.FLOAT]),
        ArgTypeSequence([DataType.FLOAT, DataType.GENERICDATETIME]),
    ]
    return_type = Fixed(DataType.GENERICDATETIME)


class BinaryMinus(Binary):
    name = "-"
    arg_names = ["value_1", "value_2"]


class BinaryMinusNumbers(BinaryMinus):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda left, right: left - right),
    ]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ]
    return_type = FromArgs()


class BinaryMinusInts(BinaryMinus):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda left, right: left - right),
    ]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER, DataType.INTEGER]),
    ]
    return_type = FromArgs()


class BinaryMinusDateInt(BinaryMinus):
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.INTEGER]),
    ]
    return_type = Fixed(DataType.DATE)


class BinaryMinusDateFloat(BinaryMinus):
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.DATE)


class BinaryMinusDatetimeNumber(BinaryMinus):
    argument_types = [
        ArgTypeSequence([DataType.DATETIME, DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.DATETIME)


class BinaryMinusGenericDatetimeNumber(BinaryMinus):
    argument_types = [
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.FLOAT]),
    ]
    return_type = Fixed(DataType.GENERICDATETIME)


class BinaryMinusDates(BinaryMinus):
    variants = [
        V(D.DUMMY, lambda left, right: left - right),
    ]
    argument_types = [
        ArgTypeSequence([DataType.DATE, DataType.DATE]),
    ]
    return_type = Fixed(DataType.INTEGER)


class BinaryMinusDatetimes(BinaryMinus):
    argument_types = [
        ArgTypeSequence([DataType.DATETIME, DataType.DATETIME]),
    ]
    return_type = Fixed(DataType.FLOAT)


class BinaryMinusGenericDatetimes(BinaryMinus):
    argument_types = [
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.GENERICDATETIME]),
    ]
    return_type = Fixed(DataType.FLOAT)


class BinaryComparison(Binary):
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


class BinaryLike(BinaryComparison):
    name = "like"
    variants = [
        V(D.DUMMY | D.SQLITE, lambda x, y: x.like(y)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]
    arg_names = ["string_1", "string_2"]


class BinaryNotLike(BinaryComparison):
    name = "notlike"
    scopes = BinaryComparison.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED
    variants = [
        V(D.DUMMY | D.SQLITE, lambda x, y: x.notlike(y)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ]
    arg_names = ["string_1", "string_2"]


class BinaryEqNEqOrderedBase(BinaryComparison):
    class BinaryEqNEqOrderedBaseArgTransformer(ArgTransformer):
        def transform_args(
            self, env: TranslationEnvironment, args: list[TranslationCtx], arg_types: list[DataType]
        ) -> list[TranslationCtx]:
            if DataType.NULL in arg_types:
                raise exc.TranslationError("Invalid comparison with NULL (use `<expr> IS [NOT] NULL`)")
            return args

    arg_transformer = BinaryEqNEqOrderedBaseArgTransformer()


COMPARABLE_TYPES: list[ArgTypeMatcher] = [
    ArgTypeSequence([DataType.BOOLEAN, DataType.BOOLEAN]),
    ArgTypeSequence([DataType.STRING, DataType.STRING]),
    ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
    ArgTypeSequence([DataType.DATE, DataType.DATE]),
    ArgTypeSequence([DataType.DATETIME, DataType.DATETIME]),
    ArgTypeSequence([DataType.DATE, DataType.DATETIME]),
    ArgTypeSequence([DataType.DATETIME, DataType.DATE]),
    ArgTypeSequence([DataType.GENERICDATETIME, DataType.GENERICDATETIME]),
    ArgTypeSequence([DataType.DATE, DataType.GENERICDATETIME]),
    ArgTypeSequence([DataType.GENERICDATETIME, DataType.DATE]),
    ArgTypeSequence([DataType.GEOPOINT, DataType.GEOPOINT]),
    ArgTypeSequence([DataType.GEOPOLYGON, DataType.GEOPOLYGON]),
    ArgTypeSequence([DataType.UUID, DataType.UUID]),
    ArgTypeSequence([DataType.ARRAY_INT, DataType.ARRAY_INT]),
    ArgTypeSequence([DataType.ARRAY_FLOAT, DataType.ARRAY_FLOAT]),
    ArgTypeSequence([DataType.ARRAY_STR, DataType.ARRAY_STR]),
]


class BinaryEqNEq(BinaryEqNEqOrderedBase):
    argument_types = COMPARABLE_TYPES


class BinaryEqNEqInternal(BinaryEqNEq):
    argument_types = BinaryEqNEq.argument_types + [
        ArgTypeSequence([DataType.MARKUP, DataType.MARKUP]),
    ]
    # not intended to be used explicitly, only in internal JOINs
    # but JOINs are implemented in the same manner as other formulas, using formula_nodes.Binary,
    # so they won't work without Scope.EXPLICIT_USAGE
    scopes = BinaryEqNEq.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class BinaryEqImpl(MultiVariantTranslation):
    variants = [
        V(
            D.DUMMY | D.SQLITE,
            lambda left, right: left == right,
        ),
    ]


class BinaryEqual(BinaryEqImpl, BinaryEqNEq):
    name = "=="
    arg_names = ["value_1", "value_2"]


class BinaryEqualInternal(BinaryEqImpl, BinaryEqNEqInternal):
    name = "_=="


class BinaryEqualDenullified(BinaryEqNEqInternal):
    name = "_dneq"
    scopes = BinaryEqNEqInternal.scopes
    variants = [
        V(D.DUMMY | D.SQLITE, lambda left, right: sa.or_(left == right, sa.and_(left.is_(None), right.is_(None)))),
    ]
    argument_types = BinaryEqNEqInternal.argument_types + [
        ArgTypeSequence([DataType.MARKUP, DataType.MARKUP]),
    ]


class BinaryNEqImpl(MultiVariantTranslation):
    variants = [
        V(D.DUMMY | D.SQLITE, lambda left, right: left != right),
    ]


class BinaryNotEqual(BinaryNEqImpl, BinaryEqNEq):
    name = "!="
    scopes = BinaryEqNEq.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED


class BinaryNotEqualInternal(BinaryNEqImpl, BinaryEqNEqInternal):
    name = "_!="


class BinaryOrderedComparison(BinaryEqNEqOrderedBase):
    argument_types = [
        ArgTypeSequence([DataType.STRING, DataType.STRING]),
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.DATE, DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME, DataType.DATETIME]),
        ArgTypeSequence([DataType.DATE, DataType.DATETIME]),
        ArgTypeSequence([DataType.DATETIME, DataType.DATE]),
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.GENERICDATETIME]),
        ArgTypeSequence([DataType.DATE, DataType.GENERICDATETIME]),
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.DATE]),
        ArgTypeSequence([DataType.UUID, DataType.UUID]),
    ]


class BinaryLessThan(BinaryOrderedComparison):
    name = "<"
    scopes = BinaryOrderedComparison.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED
    variants = [V(D.DUMMY | D.SQLITE, lambda left, right: left < right)]


class BinaryLessThanOrEqual(BinaryOrderedComparison):
    name = "<="
    scopes = BinaryOrderedComparison.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED
    variants = [
        V(D.DUMMY | D.SQLITE, lambda left, right: left <= right),
    ]


class BinaryGreaterThan(BinaryOrderedComparison):
    name = ">"
    scopes = BinaryOrderedComparison.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED
    variants = [
        V(D.DUMMY | D.SQLITE, lambda left, right: left > right),
    ]
    return_type = Fixed(DataType.BOOLEAN)


class BinaryGreaterThanOrEqual(BinaryOrderedComparison):
    name = ">="
    scopes = BinaryOrderedComparison.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED
    variants = [
        V(D.DUMMY | D.SQLITE, lambda left, right: left >= right),
    ]
    return_type = Fixed(DataType.BOOLEAN)


class BinaryAnd(Binary):  # FIXME: support other types
    name = "and"
    arg_names = ["value_1", "value_2"]
    argument_types = [
        ArgTypeSequence([DataType.BOOLEAN, DataType.BOOLEAN]),
    ]
    argument_flags = ArgFlagSequence([ContextFlag.REQ_CONDITION, ContextFlag.REQ_CONDITION])
    variants = [
        V(D.DUMMY | D.SQLITE, sa.and_),
    ]
    return_type = FromArgs()
    return_flags = ContextFlag.IS_CONDITION


class BinaryOr(Binary):  # FIXME: support other types
    name = "or"
    arg_names = ["value_1", "value_2"]
    argument_types = [
        ArgTypeSequence([DataType.BOOLEAN, DataType.BOOLEAN]),
    ]
    argument_flags = ArgFlagSequence([ContextFlag.REQ_CONDITION, ContextFlag.REQ_CONDITION])
    variants = [
        V(D.DUMMY | D.SQLITE, sa.or_),
    ]
    return_type = FromArgs()
    return_flags = ContextFlag.IS_CONDITION


def _in_fix_null(left, right, stringify_values):
    new_left, new_right = _prepare_in_args(left, right, stringify_values)
    if len(new_right) == len(right):
        return new_left.in_(new_right)

    if len(new_right) > 0:
        return sa.or_(
            new_left.is_(None),
            new_left.in_(new_right),
        )

    return new_left.is_(None)


class BinaryIn(Binary):
    name = "in"
    arg_names = ["item", "list"]
    variants = [
        V(
            D.DUMMY | D.SQLITE,
            functools.partial(_in_fix_null, stringify_values=False),
        ),
    ]
    argument_types = COMPARABLE_TYPES
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


def _not_in_fix_null(left, right, stringify_values):
    new_left, new_right = _prepare_in_args(left, right, stringify_values)
    if len(new_right) == len(right):
        return sa.or_(
            new_left.notin_(new_right),
            new_left.is_(None),
        )

    if len(new_right) > 0:
        return sa.and_(
            new_left.notin_(new_right),
            new_left.isnot(None),
        )

    return new_left.isnot(None)


def _prepare_in_args(left, right, stringify_values: bool):
    right_nulls_filtered = [x for x in right if not isinstance(x, sa.sql.elements.Null)]
    if not stringify_values:
        return left, right_nulls_filtered

    return left, [sa.func.toString(x) for x in right_nulls_filtered]


class BinaryNotIn(Binary):
    name = "notin"
    scopes = Binary.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED
    arg_names = ["item", "list"]
    variants = [
        V(
            D.DUMMY | D.SQLITE,
            functools.partial(_not_in_fix_null, stringify_values=False),
        ),
    ]
    argument_types = COMPARABLE_TYPES
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


DEFINITIONS_BINARY = [
    # !=
    BinaryNotEqual,
    # %
    BinaryModInteger,
    BinaryModFloat,
    # *
    BinaryMultNumbers,
    BinaryMultStringConst,
    BinaryMultStringNonConst,
    # +
    BinaryPlusNumbers,
    BinaryPlusStrings,
    BinaryPlusDateInt,
    BinaryPlusArray,
    BinaryPlusDateFloat,
    BinaryPlusDatetimeNumber,
    BinaryPlusGenericDatetimeNumber,
    # -
    BinaryMinusNumbers,
    BinaryMinusDateInt,
    BinaryMinusDateFloat,
    BinaryMinusDatetimeNumber,
    BinaryMinusGenericDatetimeNumber,
    BinaryMinusDates,
    BinaryMinusDatetimes,
    BinaryMinusGenericDatetimes,
    # /
    BinaryDivInt,
    BinaryDivFloat,
    # <
    BinaryLessThan,
    # <=
    BinaryLessThanOrEqual,
    # ==
    BinaryEqual,
    # >
    BinaryGreaterThan,
    # >=
    BinaryGreaterThanOrEqual,
    # ^
    BinaryPower,
    # _!=
    BinaryNotEqualInternal,
    # _==
    BinaryEqualInternal,
    # _dneq
    BinaryEqualDenullified,
    # and
    BinaryAnd,
    # in
    BinaryIn,
    # like
    BinaryLike,
    # notin
    BinaryNotIn,
    # notlike
    BinaryNotLike,
    # or
    BinaryOr,
]

lib/dl_formula/dl_formula/definitions/operators_binary.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import ArgTypeSequence
from dl_formula.definitions.base import (
    MultiVariantTranslation,
    TranslationVariant,
)
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import Fixed


V = TranslationVariant.make


class Ternary(MultiVariantTranslation):
    arg_cnt = 3
    is_function = False


class TernaryBetweenBase(Ternary):
    arg_names = ["value", "low", "high"]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT, DataType.FLOAT, DataType.FLOAT]),
        ArgTypeSequence([DataType.DATE, DataType.DATE, DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME, DataType.DATETIME, DataType.DATETIME]),
        ArgTypeSequence([DataType.DATETIMETZ, DataType.DATETIMETZ, DataType.DATETIMETZ]),
        ArgTypeSequence([DataType.GENERICDATETIME, DataType.GENERICDATETIME, DataType.GENERICDATETIME]),
        # TODO later: ArgTypeSequence([DataType.DATETIME, DataType.DATETIMETZ, DataType.DATETIMETZ]),
        ArgTypeSequence([DataType.STRING, DataType.STRING, DataType.STRING]),
    ]
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = ContextFlag.IS_CONDITION


class TernaryBetween(TernaryBetweenBase):
    name = "between"
    variants = [
        V(D.DUMMY | D.SQLITE, lambda a, b, c: a.between(b, c)),
    ]


class TernaryNotBetween(TernaryBetweenBase):
    name = "notbetween"
    scopes = TernaryBetweenBase.scopes & ~Scope.SUGGESTED & ~Scope.DOCUMENTED
    variants = [
        V(D.DUMMY | D.SQLITE, lambda a, b, c: sa.not_(a.between(b, c))),
    ]


DEFINITIONS_TERNARY = [
    # between
    TernaryBetween,
    # notbetween
    TernaryNotBetween,
]

lib/dl_formula/dl_formula/definitions/operators_ternary.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.args import (
    ArgFlagSequence,
    ArgTypeSequence,
)
from dl_formula.definitions.base import (
    MultiVariantTranslation,
    TranslationVariant,
)
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.type_strategy import (
    Fixed,
    FromArgs,
)


V = TranslationVariant.make


class Unary(MultiVariantTranslation):
    arg_cnt = 1
    is_function = False


class UnaryNot(Unary):
    name = "not"
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = int(ContextFlag.IS_CONDITION)


class UnaryNotBool(UnaryNot):
    variants = [V(D.DUMMY | D.SQLITE, sa.not_)]
    argument_types = [
        ArgTypeSequence([DataType.BOOLEAN]),
    ]
    argument_flags = ArgFlagSequence([ContextFlag.REQ_CONDITION])


class UnaryNotNumbers(UnaryNot):
    variants = [V(D.DUMMY | D.SQLITE, lambda x: x == 0)]
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
    ]


class UnaryNotStringGeo(UnaryNot):
    variants = [
        V(D.DUMMY, lambda x: x == ""),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
        ArgTypeSequence([DataType.GEOPOINT]),
        ArgTypeSequence([DataType.GEOPOLYGON]),
    ]


class UnaryNotDateDatetime(UnaryNot):
    variants = [
        V(D.DUMMY, lambda x: sa.literal(False)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME]),
        ArgTypeSequence([DataType.GENERICDATETIME]),
    ]
    return_flags = 0


class UnaryNegate(Unary):
    name = "neg"
    variants = [V(D.DUMMY | D.SQLITE, lambda x: -x)]
    return_type = FromArgs()
    argument_types = [
        ArgTypeSequence([DataType.INTEGER]),
        ArgTypeSequence([DataType.FLOAT]),
    ]


class UnaryIsOpBase(Unary):
    return_type = Fixed(DataType.BOOLEAN)
    return_flags = int(ContextFlag.IS_CONDITION)


class UnaryIsTrue(UnaryIsOpBase):
    name = "istrue"
    arg_names = ["value"]


class UnaryIsTrueStringGeo(UnaryIsTrue):
    variants = [
        # TODO?: use `n.func.NOT`?
        V(D.DUMMY, lambda x: x != ""),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
        ArgTypeSequence([DataType.GEOPOINT]),
        ArgTypeSequence([DataType.GEOPOLYGON]),
    ]


class UnaryIsTrueNumbers(UnaryIsTrue):
    # TODO?: use `n.func.BOOL`?
    variants = [V(D.DUMMY | D.SQLITE, lambda x: x != 0.0)]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]


class UnaryIsTrueDateTime(UnaryIsTrue):
    # TODO?: use `n.func.NOT`?
    variants = [
        V(D.DUMMY, lambda x: sa.literal(True)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME]),
        ArgTypeSequence([DataType.GENERICDATETIME]),
    ]
    return_flags = 0


class UnaryIsTrueBoolean(UnaryIsTrue):
    variants = [
        V(D.DUMMY, lambda x: x.is_(True)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.BOOLEAN]),
    ]


class UnaryIsFalse(UnaryIsOpBase):
    name = "isfalse"
    arg_names = ["value"]


class UnaryIsFalseStringGeo(UnaryIsFalse):
    variants = [
        # TODO?: use `n.func.NOT`?
        V(D.DUMMY, lambda x: x == ""),
    ]
    argument_types = [
        ArgTypeSequence([DataType.STRING]),
        ArgTypeSequence([DataType.GEOPOINT]),
        ArgTypeSequence([DataType.GEOPOLYGON]),
    ]


class UnaryIsFalseNumbers(UnaryIsFalse):
    variants = [V(D.DUMMY | D.SQLITE, lambda x: x == 0.0)]
    argument_types = [
        ArgTypeSequence([DataType.FLOAT]),
    ]


class UnaryIsFalseDateTime(UnaryIsFalse):
    variants = [
        V(D.DUMMY, lambda x: sa.literal(False)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.DATE]),
        ArgTypeSequence([DataType.DATETIME]),
        ArgTypeSequence([DataType.GENERICDATETIME]),
    ]
    return_flags = 0


class UnaryIsFalseBoolean(UnaryIsFalse):
    variants = [
        V(D.DUMMY, lambda x: x.is_(False)),
    ]
    argument_types = [
        ArgTypeSequence([DataType.BOOLEAN]),
    ]


DEFINITIONS_UNARY = [
    # isfalse
    UnaryIsFalseStringGeo,
    UnaryIsFalseNumbers,
    UnaryIsFalseDateTime,
    UnaryIsFalseBoolean,
    # istrue
    UnaryIsTrueStringGeo,
    UnaryIsTrueNumbers,
    UnaryIsTrueDateTime,
    UnaryIsTrueBoolean,
    # neg
    UnaryNegate,
    # not
    UnaryNotBool,
    UnaryNotNumbers,
    UnaryNotStringGeo,
    UnaryNotDateDatetime,
]

lib/dl_formula/dl_formula/definitions/operators_unary.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from typing import (
    Generator,
    NamedTuple,
    Optional,
    Sequence,
)

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D
import dl_formula.core.exc as exc
import dl_formula.definitions.base as op_base
from dl_formula.definitions.scope import Scope


class FuncKey(NamedTuple):
    name: str
    arg_cnt: Optional[int]
    is_window: bool

    @property
    def unlimited_arg_version(self) -> FuncKey:
        return self._replace(arg_cnt=None)


class OperationRegistry:
    def __init__(self):
        self.ops: dict[FuncKey, list["op_base.NodeTranslation"]] = {}

    def __contains__(self, key: FuncKey) -> bool:
        return key in self.ops

    def __getitem__(self, key: FuncKey) -> list:
        return self.ops[key]

    def __bool__(self) -> bool:
        return bool(self.ops)

    def _func_key_from_def_item(self, def_item: op_base.NodeTranslation) -> FuncKey:
        name = def_item.name
        assert name is not None
        return FuncKey(
            name=name,
            arg_cnt=def_item.arg_cnt,
            is_window=def_item.is_window,
        )

    def register(self, def_item: "op_base.NodeTranslation") -> None:
        assert isinstance(def_item, op_base.NodeTranslation)
        if def_item.name is None:
            raise ValueError("Cannot register nameless functions")
        func_key = self._func_key_from_def_item(def_item)
        if func_key not in self.ops:
            self.ops[func_key] = []
        if def_item not in self.ops[func_key]:
            self.ops[func_key].append(def_item)

    def unregister(self, def_item: "op_base.NodeTranslation") -> None:
        if def_item.name is None:
            raise ValueError("Cannot unregister nameless functions")
        func_key = self._func_key_from_def_item(def_item)
        self.ops[func_key].remove(def_item)
        if not self.ops[func_key]:
            del self.ops[func_key]

    def get_definition(
        self,
        *,
        name: str,
        arg_types: Sequence[DataType],
        is_window: bool = False,
        dialect: Optional[DialectCombo] = None,
        for_any_dialect: bool = False,
        required_scopes: int = Scope.EXPLICIT_USAGE,
    ) -> "op_base.NodeTranslation":
        """Find translation that matches given name and arguments"""

        if dialect is not None and for_any_dialect or dialect is None and not for_any_dialect:
            raise ValueError(
                "Either dialect should be provided or for_any_dialect be set to True." " Cannot provide both."
            )

        arg_types = list(arg_types)
        func_key = FuncKey(name=name.lower(), arg_cnt=len(arg_types), is_window=is_window)

        # to also find versions of the F with unlimited number of args:
        inf_func_key = func_key.unlimited_arg_version
        translations = [
            func_translation
            for func_translation in self.ops.get(func_key, []) + self.ops.get(inf_func_key, [])
            if func_translation.scopes & required_scopes == required_scopes
        ]
        if not translations:
            # function not found
            raise exc.TranslationUnknownFunctionError(
                f"Unknown {len(arg_types)}-argument function or operator {name.upper()}"
            )

        some_type_match = False

        for func_translation in translations:
            if func_translation.match_types(arg_types):
                if dialect is None:
                    assert for_any_dialect
                    return func_translation
                elif func_translation.match_dialect(dialect):
                    return func_translation
                else:
                    some_type_match = True

        # such a function does exist, but there is no implementation that accepts the given arguments
        what = "function" if translations[0].is_function else "operator"
        what = f"window {what}" if is_window else what
        arg_types_repr = ", ".join(arg_type.name for arg_type in arg_types)
        if some_type_match:
            message_comment = f" the {what} is only defined for other databases"
        else:
            message_comment = f" no such {what} currently exists"
        message_parts = [
            f"Invalid argument types for {what} {name.upper()}",
            f"dialect: {dialect.common_name_and_version if dialect is not None else None}",
            f"types: ({arg_types_repr})",
            message_comment,
        ]
        raise exc.DataTypeError("; ".join(message_parts) + ".")

    def items(self) -> Generator[tuple[FuncKey, "op_base.NodeTranslation"], None, None]:
        """Like regular ``dict.items()``,  but flatten translation lists"""
        sorted_keys = sorted(
            self.ops.keys(), key=lambda key: (key[0], key[1] if key[1] is not None else float("inf"), *key[2:])
        )
        for func_key in sorted_keys:
            trans_list = self.ops[func_key]
            for translation in trans_list:
                yield func_key, translation

    def get_supported_functions(
        self,
        require_dialects: DialectCombo = D.EMPTY,
        only_functions: bool = True,
        function_scopes: int = Scope.EXPLICIT_USAGE,
    ) -> list[FuncKey]:
        """
        Return a list of supported functions for given combination of dialects (``require_dialects``).
        Result is a list of (<func_name>, <arg_cnt>) tuples
        (keys from ``OPERATION_REGISTRY``).
        """
        support_by_key = defaultdict(lambda: D.EMPTY)  # type: ignore  # TODO: fix
        for func_key, func_tr in self.items():
            if func_tr.scopes & function_scopes != function_scopes:
                continue
            if only_functions and not func_tr.is_function:
                continue
            support_by_key[func_key] |= func_tr.supported_dialects()

        result = []
        sorted_keys = sorted(
            support_by_key.keys(), key=lambda key: (key[0], key[1] if key[1] is not None else float("inf"), *key[2:])
        )
        for func_key in sorted_keys:
            supp_dialects = support_by_key[func_key]
            if supp_dialects & require_dialects == require_dialects:
                result.append(func_key)

        return result


OPERATION_REGISTRY = OperationRegistry()

lib/dl_formula/dl_formula/definitions/registry.py
datalens-tech/datalens-backend
from enum import (
    IntEnum,
    unique,
)


_BIT = 1


def _auto_bit() -> int:
    global _BIT
    old = _BIT
    _BIT <<= 1
    return old


@unique
class Scope(IntEnum):
    STABLE = _auto_bit()
    UNSTABLE = _auto_bit()

    # Functionality-bound scopes
    EXPLICIT_USAGE = _auto_bit()  # Can be used in formulas explicitly
    SUGGESTED = _auto_bit()  # Function is listed in suggestions
    DOCUMENTED = _auto_bit()  # Documentation is generated for function

lib/dl_formula/dl_formula/definitions/scope.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Callable,
    List,
    Sequence,
    Union,
)

from dl_formula.core.datatype import (
    DataType,
    DataTypeParams,
)
import dl_formula.core.exc as exc
from dl_formula.translation.type_strategy import (
    TypeParamsStrategy,
    TypeStrategy,
)


if TYPE_CHECKING:
    from dl_formula.translation.context import TranslationCtx


class Fixed(TypeStrategy):
    __slots__ = ("_type",)

    def __init__(self, type_: DataType):
        self._type = type_

    @property
    def type(self) -> DataType:
        return self._type

    def get_from_args(self, arg_types: List[DataType]) -> DataType:
        return self._type


class FromArgs(TypeStrategy):
    __slots__ = ("_indices",)

    def __init__(self, *indices: Union[int, slice]):
        self._indices = indices or [slice(0, None)]

    def get_from_args(self, arg_types: List[DataType]) -> DataType:
        return self._get_from_args_and_indices(arg_types, self._indices)  # type: ignore

    @staticmethod
    def _get_from_args_and_indices(
        arg_types: List[DataType],
        indices: Sequence[Union[int, slice]],
    ) -> DataType:
        use_arg_types = []
        for ind in indices:  # type: ignore
            if isinstance(ind, int):
                use_arg_types.append(arg_types[ind])
            elif isinstance(ind, slice):
                use_arg_types.extend(arg_types[ind])
            else:
                raise TypeError(type(ind))

        return DataType.get_common_cast_type(*use_arg_types).non_const_type


class DynamicIndexStrategy(FromArgs):
    __slots__ = ()

    def get_indices(self, arg_cnt: int) -> List[Union[int, slice]]:
        return list(range(arg_cnt))

    def get_from_args(self, arg_types: List[DataType]) -> DataType:
        return self._get_from_args_and_indices(arg_types, self.get_indices(len(arg_types))).non_const_type


class Combined(TypeStrategy):
    __slots__ = ("_substrats", "_replace_types")

    def __init__(self, *substrats: TypeStrategy, replace_types=None):
        self._substrats = []
        self._replace_types = replace_types or {}
        for substrat in substrats or []:
            if isinstance(substrat, DataType):
                substrat = Fixed(substrat)
            self._substrats.append(substrat)

    def get_from_args(self, arg_types: List[DataType]) -> DataType:
        types = set()
        for substrat in self._substrats:
            types.add(substrat.get_from_args(arg_types))

        for repl_t, new_t in self._replace_types.items():
            if repl_t in types:
                types.remove(repl_t)
                types.add(new_t)

        return DataType.get_common_cast_type(*types)


class CaseTypeStrategy(DynamicIndexStrategy):
    __slots__ = ()

    def get_indices(self, arg_cnt: int) -> List[Union[int, slice]]:
        if arg_cnt < 2 or arg_cnt % 2 != 0:
            raise exc.TranslationError(f"Invalid number of arguments for CASE: {arg_cnt}")
        num_of_whens = (arg_cnt - 1) // 2  # 1 for main CASE value arg, 2 for each WHEN part
        return [*[2 + 2 * i for i in range(num_of_whens)], 1 + 2 * num_of_whens]  # THENs  # ELSE


class IfTypeStrategy(DynamicIndexStrategy):
    __slots__ = ()

    def get_indices(self, arg_cnt: int) -> List[Union[int, slice]]:
        if arg_cnt < 3 or arg_cnt % 2 != 1:
            raise exc.TranslationError(f"Invalid number of arguments for IF: {arg_cnt}")
        num_of_if_parts = arg_cnt // 2
        return [*[2 * i + 1 for i in range(num_of_if_parts)], 2 * num_of_if_parts]  # THENs  # ELSE


class ParamsEmpty(TypeParamsStrategy):
    def get_from_arg_values(self, args: List[TranslationCtx]) -> DataTypeParams:
        return DataTypeParams()


class ParamsCustom(TypeParamsStrategy):
    def __init__(self, func: Callable[[List[TranslationCtx]], DataTypeParams]):
        self._func = func

    def get_from_arg_values(self, args: List[TranslationCtx]) -> DataTypeParams:
        return self._func(args)


class ParamsFromArgs(TypeParamsStrategy):
    def __init__(self, index: int):
        self._index = index

    def get_from_arg_values(self, args: List[TranslationCtx]) -> DataTypeParams:
        return args[self._index].data_type_params  # type: ignore

lib/dl_formula/dl_formula/definitions/type_strategy.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import singledispatchmethod
import textwrap
from typing import Union
import uuid

import graphviz

from dl_formula.core import nodes


def make_name() -> str:
    return "node_{}".format(uuid.uuid4().hex)


class DotTranslator:
    terminal_opts = dict(shape="rectangle", style="filled", margin="0")
    literal_colors = {
        str: "#fff6c1",
        int: "#b2dee6",
        float: "#c4b4e6",
        bool: "#d7e6df",
    }

    def __init__(self):
        pass

    def translate(self, formula: nodes.Formula, title: str = "Formula") -> graphviz.Digraph:
        """Translate ``Formula`` object into GraphViz DOT format"""
        dot = graphviz.Digraph(comment="Formula")
        self.translate_node(formula, dot)
        return dot

    def translate_node(self, node: nodes.FormulaItem, dot: graphviz.Digraph) -> str:
        return self._translate_node(node, dot)

    @singledispatchmethod
    def _translate_node(self, node: nodes.FormulaItem, dot: graphviz.Digraph) -> str:
        """Default action for methdispatch."""
        raise TypeError(type(node))

    @_translate_node.register(nodes.Formula)
    def _translate_node_formula(self, node: nodes.Formula, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        wrapped_formula = "\n".join(textwrap.wrap(node.original_text or "", 90))
        dot.node(node_name, label=f"Formula:\n{wrapped_formula}", shape="plaintext")
        child_name = self.translate_node(node.expr, dot)
        dot.edge(node_name, child_name, label="EXPRESSION")
        return node_name

    @_translate_node.register(nodes.Field)
    def _translate_node_field(self, node: nodes.Field, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        dot.node(node_name, label="[{}]".format(node.name), **self.terminal_opts, fillcolor="#e6cccc")
        return node_name

    @_translate_node.register(nodes.LiteralInteger)
    @_translate_node.register(nodes.LiteralFloat)
    @_translate_node.register(nodes.LiteralBoolean)
    @_translate_node.register(nodes.LiteralString)
    @_translate_node.register(nodes.LiteralDate)
    @_translate_node.register(nodes.LiteralDatetime)
    @_translate_node.register(nodes.LiteralGenericDatetime)
    @_translate_node.register(nodes.LiteralGeopoint)
    @_translate_node.register(nodes.LiteralGeopolygon)
    @_translate_node.register(nodes.LiteralUuid)
    def _translate_node_literal(self, node: nodes.BaseLiteral, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        dot.node(
            node_name,
            label="{!r}".format(node.value),
            **self.terminal_opts,
            fillcolor=self.literal_colors[type(node.value)],
        )
        return node_name

    @_translate_node.register(nodes.Null)
    def _translate_node_null(self, node: nodes.Null, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        dot.node(node_name, label="NULL", **self.terminal_opts, fillcolor="#959595")
        return node_name

    @_translate_node.register(nodes.ExpressionList)
    def _translate_node_expression_list(self, node: nodes.ExpressionList, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        dot.node(
            node_name,
            label="[0..{}]".format(len(node.children) - 1),
            shape="hexagon",
            style="filled",
            fillcolor="#bba7fc",
        )
        for i, child in enumerate(node.children):  # type: ignore  # TODO: fix
            child_name = self.translate_node(child, dot)
            dot.edge(node_name, child_name, label="[{}]".format(i))

        return node_name

    @_translate_node.register(nodes.FuncCall)
    @_translate_node.register(nodes.WindowFuncCall)
    def _translate_node_func_call(
        self, node: Union[nodes.FuncCall, nodes.WindowFuncCall], dot: graphviz.Digraph
    ) -> str:
        node_name = make_name()
        args_names = ["ARG{}".format(i) for i in range(len(node.args))]
        dot.node(
            node_name,
            label="{}({})".format(node.name.upper(), ", ".join(args_names)),
            shape="oval",
            style="filled",
            fillcolor="#bde6b1",
        )
        for i, arg in enumerate(node.args):
            child_name = self.translate_node(arg, dot)
            dot.edge(node_name, child_name, label=args_names[i])

        if isinstance(node, nodes.WindowFuncCall):
            grouping_name = self.translate_node(node.grouping, dot)
            dot.edge(node_name, grouping_name, label="GROUPING")

            if list(node.ordering.expr_list):
                ordering_name = self.translate_node(node.ordering, dot)
                dot.edge(node_name, ordering_name, label="ORDERING")

            if list(node.before_filter_by.field_names):
                before_filter_by_name = self.translate_node(node.before_filter_by, dot)
                dot.edge(node_name, before_filter_by_name, label="...")

        return node_name

    @_translate_node.register(nodes.Binary)
    def _translate_node_binary(self, node: nodes.Binary, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        op_name = node.name.upper()
        dot.node(node_name, label=op_name, shape="oval", style="filled", fillcolor="#c2c2c2")

        if op_name == "IN":
            left_label, right_label = "EL", "ARRAY"
        else:
            left_label, right_label = "LEFT", "RIGHT"

        left_name = self.translate_node(node.left, dot)
        dot.edge(node_name, left_name, label=left_label)
        right_name = self.translate_node(node.right, dot)
        dot.edge(node_name, right_name, label=right_label)
        return node_name

    @_translate_node.register(nodes.Ternary)
    def _translate_node_ternary(self, node: nodes.Ternary, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        dot.node(node_name, label=node.name.upper(), shape="oval", style="filled", fillcolor="#c2c2c2")
        first_name = self.translate_node(node.first, dot)
        dot.edge(node_name, first_name, label="1ST")
        second_name = self.translate_node(node.second, dot)
        dot.edge(node_name, second_name, label="2ND")
        third_name = self.translate_node(node.third, dot)
        dot.edge(node_name, third_name, label="3RD")
        return node_name

    @_translate_node.register(nodes.Unary)
    def _translate_node_unary(self, node: nodes.Unary, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        dot.node(node_name, label=node.name.upper(), shape="oval", style="filled", fillcolor="#c2c2c2")
        expr_name = self.translate_node(node.expr, dot)
        dot.edge(node_name, expr_name)
        return node_name

    @_translate_node.register(nodes.WindowGroupingTotal)
    def _translate_node_window_grouping_total(self, node: nodes.WindowGroupingTotal, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        dot.node(node_name, label="TOTAL", shape="octagon", style="filled", fillcolor="#bba7fc")
        return node_name

    @_translate_node.register(nodes.WindowGroupingWithin)
    @_translate_node.register(nodes.WindowGroupingAmong)
    def _translate_node_window_grouping_with_dimensions(
        self, node: Union[nodes.WindowGroupingWithin, nodes.WindowGroupingAmong], dot: graphviz.Digraph
    ) -> str:
        node_name = make_name()
        keyword = {
            nodes.WindowGroupingWithin: "WHITHIN",
            nodes.WindowGroupingAmong: "AMONG",
        }[type(node)]
        dot.node(node_name, label=keyword, shape="octagon", style="filled", fillcolor="#ff898b")
        for i, child in enumerate(node.children):
            child_name = self.translate_node(child, dot)  # type: ignore  # TODO: fix
            dot.edge(node_name, child_name, label="DIM {}".format(i))

        return node_name

    @_translate_node.register(nodes.Ordering)
    def _translate_node_ordering(self, node: nodes.Ordering, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        dot.node(node_name, label="ORDER BY", shape="octagon", style="filled", fillcolor="#ff898b")
        for i, child in enumerate(node.children):  # type: ignore  # TODO: fix
            child_name = self.translate_node(child, dot)
            dot.edge(node_name, child_name, label="EXPR {}".format(i))

        return node_name

    @_translate_node.register(nodes.BeforeFilterBy)
    def _translate_node_before_filter_by(self, node: nodes.BeforeFilterBy, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        dot.node(node_name, label="BEFORE FILTER BY", shape="octagon", style="filled", fillcolor="#ff898b")
        for i, field_name in enumerate(node.field_names):
            child_name = self.translate_node(nodes.Field.make(name=field_name), dot)
            dot.edge(node_name, child_name, label="FIELD {}".format(i))

        return node_name

    @_translate_node.register(nodes.IfBlock)
    def _translate_node_if_block(self, node: nodes.IfBlock, dot: graphviz.Digraph) -> str:
        root_name = make_name()
        dot.node(root_name, label="IF-BLOCK", shape="pentagon", style="filled", fillcolor="#c2c2c2")

        parent_name = None
        for if_part in node.if_list:
            if_name = make_name()
            dot.node(if_name, label="IF", shape="diamond", style="filled", fillcolor="#c2c2c2")
            dot.edge(root_name, if_name, label="PART")
            if parent_name:
                dot.edge(parent_name, if_name, label="ELSE", style="dotted")
            cond_name = self.translate_node(if_part.cond, dot)
            dot.edge(if_name, cond_name, label="TRUE?")
            expr_name = self.translate_node(if_part.expr, dot)
            dot.edge(if_name, expr_name, label="THEN")
            parent_name = if_name

        else_name = self.translate_node(node.else_expr, dot)
        dot.edge(parent_name, else_name, label="ELSE", style="dotted")
        dot.edge(root_name, else_name, label="PART")

        return root_name

    @_translate_node.register(nodes.CaseBlock)
    def _translate_node_case_block(self, node: nodes.CaseBlock, dot: graphviz.Digraph) -> str:
        root_name = make_name()
        dot.node(root_name, label="CASE", shape="pentagon", style="filled", fillcolor="#c2c2c2")
        expr_name = self.translate_node(node.case_expr, dot)
        dot.edge(root_name, expr_name, label="WHAT")

        for when_part in node.when_list:
            when_name = make_name()
            dot.node(when_name, label="WHEN", shape="diamond", style="filled", fillcolor="#c2c2c2")
            dot.edge(root_name, when_name)
            value_name = self.translate_node(when_part.val, dot)
            dot.edge(when_name, value_name, label="IS")
            expr_name = self.translate_node(when_part.expr, dot)
            dot.edge(when_name, expr_name, label="THEN")

        else_name = self.translate_node(node.else_expr, dot)
        dot.edge(root_name, else_name, label="ELSE")

        return root_name

    @_translate_node.register(nodes.ParenthesizedExpr)
    @_translate_node.register(nodes.OrderAscending)
    @_translate_node.register(nodes.OrderDescending)
    def _translate_node_parenthesized_expr(self, node: nodes.ParenthesizedExpr, dot: graphviz.Digraph) -> str:
        node_name = make_name()
        label = {
            nodes.ParenthesizedExpr: "(...)",
            nodes.OrderAscending: "ASC",
            nodes.OrderDescending: "DESC",
        }[type(node)]
        dot.node(node_name, label=label, shape="oval", style="filled", fillcolor="#e0e0e0")
        expr_name = self.translate_node(node.expr, dot)
        dot.edge(node_name, expr_name, label="EXPR")
        return node_name


def translate(formula: nodes.Formula, title: str = "Formula") -> graphviz.Digraph:
    """Translate formula into DOT representation"""
    translator = DotTranslator()
    return translator.translate(formula, title=title)

lib/dl_formula/dl_formula/dot.py
datalens-tech/datalens-backend
from importlib import metadata
from typing import (
    Collection,
    Optional,
    Type,
)

import dl_formula as package
from dl_formula.connectors.base.connector import FormulaConnector
from dl_formula.connectors.registration import CONN_REG_FORMULA


_CONNECTOR_EP_GROUP = f"{package.__name__}.connectors"


def _get_all_ep_connectors() -> dict[str, Type[FormulaConnector]]:
    entrypoints = list(metadata.entry_points().select(group=_CONNECTOR_EP_GROUP))  # type: ignore
    return {ep.name: ep.load() for ep in entrypoints}


def load_all_connectors() -> None:
    _get_all_ep_connectors()


def register_all_connectors(connector_ep_names: Optional[Collection[str]] = None) -> None:
    connectors = _get_all_ep_connectors()
    for ep_name, connector_cls in sorted(connectors.items()):
        if connector_ep_names is not None and ep_name not in connector_ep_names:
            continue
        CONN_REG_FORMULA.register_connector(connector_cls)

lib/dl_formula/dl_formula/formula_connectors.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_formula.collections import NodeValueMap
from dl_formula.core.datatype import DataType


@attr.s
class InspectionEnvironment:
    cache_is_const_expr: NodeValueMap[bool] = attr.ib(factory=NodeValueMap)
    cache_is_unbound_expr: NodeValueMap[bool] = attr.ib(factory=NodeValueMap)
    cache_is_aggregate_expr: NodeValueMap[bool] = attr.ib(factory=NodeValueMap)
    cache_is_window_expr: NodeValueMap[bool] = attr.ib(factory=NodeValueMap)
    cache_is_query_fork_expr: NodeValueMap[bool] = attr.ib(factory=NodeValueMap)
    cache_query_fork_nesting: NodeValueMap[int] = attr.ib(factory=NodeValueMap)
    cache_data_type: NodeValueMap[DataType] = attr.ib(factory=NodeValueMap)

lib/dl_formula/dl_formula/inspect/env.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    AbstractSet,
    Callable,
    Dict,
    Generator,
    Iterable,
    List,
    Optional,
    Set,
    Tuple,
    Type,
)

from dl_formula.collections import NodeSet
from dl_formula.core.datatype import DataType
import dl_formula.core.exc as exc
import dl_formula.core.fork_nodes as fork_nodes
from dl_formula.core.index import NodeHierarchyIndex
import dl_formula.core.nodes as nodes
from dl_formula.core.tag import LevelTag
import dl_formula.definitions.type_strategy
import dl_formula.inspect.env
import dl_formula.inspect.function
import dl_formula.inspect.literal
import dl_formula.inspect.node
import dl_formula.translation.ext_nodes as ext_nodes


def used_fields(item: nodes.FormulaItem) -> List[nodes.Field]:
    return item.list_node_type(nodes.Field)


def used_func_calls(item: nodes.FormulaItem) -> List[nodes.FuncCall]:
    return item.list_node_type(nodes.FuncCall)


def validate_not_compiled_expressions(node: nodes.FormulaItem) -> None:
    # TODO: Remove
    if isinstance(node, ext_nodes.CompiledExpression):
        raise exc.InspectionError("CompiledExpression nodes are not supported in inspection")


def is_constant_expression(
    node: nodes.FormulaItem,
    env: Optional[dl_formula.inspect.env.InspectionEnvironment] = None,
) -> bool:
    """
    Check whether the given formula node (subtree) defines a constant expression.
    Return ``True`` is it does, otherwise ``False``.
    """
    result: Optional[bool] = None
    if env is not None:
        result = env.cache_is_const_expr.get(node)
    if result is None:
        validate_not_compiled_expressions(node)
        if isinstance(node, nodes.BaseLiteral) or isinstance(node, nodes.Null):
            result = True
        elif isinstance(node, nodes.ParenthesizedExpr):
            result = is_constant_expression(node.expr, env=env)
        else:
            result = False

        if env is not None:
            env.cache_is_unbound_expr.add(node, value=result)

    assert result is not None
    return result


def is_bound_only_to(node: nodes.FormulaItem, allow_nodes: NodeSet) -> bool:
    """
    Check whether the given formula node (subtree) defines an unbound (field-independent) expression.
    Return ``True`` is it does, otherwise ``False``.
    """

    if node in allow_nodes:
        return True

    if isinstance(node, nodes.Field):
        return False

    for child in autonomous_children(node):
        if not is_bound_only_to(node=child, allow_nodes=allow_nodes):
            return False

    return True


def is_aggregate_expression(
    node: nodes.FormulaItem,
    env: "dl_formula.inspect.env.InspectionEnvironment",
) -> bool:
    """
    Check whether the given formula node (subtree) defines an aggregate expression.
    Return ``True`` is it does, otherwise ``False``.
    """

    result = env.cache_is_aggregate_expr.get(node)
    if result is None:
        validate_not_compiled_expressions(node)
        if dl_formula.inspect.node.is_aggregate_function(node):
            result = True
        else:
            result = False
            for child in autonomous_children(node):
                if is_aggregate_expression(node=child, env=env):
                    result = True
                    break
        env.cache_is_aggregate_expr.add(node, value=result)

    return result


def is_aggregated_above_sub_node(
    node: nodes.FormulaItem,
    index: Optional[NodeHierarchyIndex] = None,
    sub_node: Optional[nodes.FormulaItem] = None,
) -> bool:
    """
    Check whether the given formula node (subtree) wraps
    the given ``sub_node`` into an aggregation.
    """

    if index is not None and sub_node is not None:
        raise ValueError("Cannot specify both index and sub_node")

    if index is None:
        assert sub_node is not None
        index = node.resolve_index(sub_node, pointer_eq=True)

    # Check whether any of the nodes wrapping sub_node are aggregations
    assert index is not None
    for wrapping_node in node.iter_index(index, exclude_last=True):  # exclude the sub_node
        # ... because we only care about its wrappers
        validate_not_compiled_expressions(wrapping_node)
        if dl_formula.inspect.node.is_aggregate_function(wrapping_node):
            return True

    return False


def is_window_expression(
    node: nodes.FormulaItem,
    env: "dl_formula.inspect.env.InspectionEnvironment",
) -> bool:
    """
    Check whether the given formula node (subtree) defines a window function expression.
    Return ``True`` is it does, otherwise ``False``.
    """

    result = env.cache_is_window_expr.get(node)
    if result is None:
        validate_not_compiled_expressions(node)
        if isinstance(node, nodes.WindowFuncCall):
            result = True
        else:
            result = False
            for child in autonomous_children(node):
                if is_window_expression(node=child, env=env):
                    result = True
                    break
        env.cache_is_window_expr.add(node, value=result)

    return result


def is_query_fork_expression(
    node: nodes.FormulaItem,
    env: "dl_formula.inspect.env.InspectionEnvironment",
) -> bool:
    """
    Check whether the given formula node (subtree) contains any ``QueryFork`` nodes
    either directly as one of its children, or indirectly nested in the tree of its children.
    Return ``True`` is it does, otherwise ``False``.
    """

    result = env.cache_is_query_fork_expr.get(node)
    if result is None:
        validate_not_compiled_expressions(node)
        if isinstance(node, fork_nodes.QueryFork):
            result = True
        else:
            result = False
            for child in autonomous_children(node):
                if is_query_fork_expression(node=child, env=env):
                    result = True
                    break
        env.cache_is_query_fork_expr.add(node, value=result)

    return result


def get_query_fork_nesting_level(
    node: nodes.FormulaItem,
    env: "dl_formula.inspect.env.InspectionEnvironment",
) -> int:
    """
    TODO
    """

    result = env.cache_query_fork_nesting.get(node)
    if result is None:
        validate_not_compiled_expressions(node)

        self_qf_level = 0
        if isinstance(node, fork_nodes.QueryFork):
            self_qf_level = 1

        max_child_qf_level = 0
        for child in autonomous_children(node):
            max_child_qf_level = max(
                max_child_qf_level,
                get_query_fork_nesting_level(node=child, env=env),
            )

        result = self_qf_level + max_child_qf_level
        env.cache_query_fork_nesting.add(node, value=result)

    return result


def infer_data_type(
    node: nodes.FormulaItem,
    field_types: Dict[str, DataType],
    env: "dl_formula.inspect.env.InspectionEnvironment",
) -> DataType:
    """Calculate the data type that should be returned by the given expression"""

    def _infer_data_type() -> DataType:
        if isinstance(node, nodes.Field):
            result = field_types.get(node.name)
            if result is None:
                raise exc.TranslationUnknownFieldError(
                    f"Data type inference error: unknown field: {node.name}",
                )

        elif isinstance(node, (nodes.BaseLiteral, nodes.Null)):
            result = dl_formula.inspect.literal.get_data_type(node)

        elif isinstance(node, nodes.ExpressionList):
            child_types = [infer_data_type(child, field_types=field_types, env=env) for child in node.children]
            result = dl_formula.definitions.type_strategy.FromArgs().get_from_args(arg_types=child_types)

        elif isinstance(node, (nodes.OperationCall, nodes.IfBlock, nodes.CaseBlock)):
            if isinstance(node, nodes.IfBlock):
                func_name = "if"
                children = autonomous_children(node)
            elif isinstance(node, nodes.CaseBlock):
                func_name = "case"
                children = autonomous_children(node)
            else:
                func_name = node.name
                children = node.args

            child_types = [infer_data_type(child, field_types=field_types, env=env) for child in children]
            result = dl_formula.inspect.function.get_return_type(
                name=func_name, arg_types=child_types, is_window=isinstance(node, nodes.WindowFuncCall)
            )

        elif isinstance(node, (nodes.Formula, nodes.ParenthesizedExpr)):
            result = infer_data_type(node.expr, field_types=field_types, env=env)

        elif isinstance(node, fork_nodes.QueryFork):
            result = infer_data_type(node.result_expr, field_types=field_types, env=env)

        else:
            raise TypeError(type(node))

        assert result is not None
        return result

    data_type = env.cache_data_type.get(node)
    if data_type is None:
        data_type = _infer_data_type()
        env.cache_data_type.add(node, value=data_type)

    return data_type


def enumerate_fields(
    node: nodes.FormulaItem,
    prefix: NodeHierarchyIndex = NodeHierarchyIndex(),
) -> Generator[Tuple[NodeHierarchyIndex, nodes.Field], None, None]:
    """Just like ``enumerate``, but yields only field nodes."""
    for index, sub_node in node.enumerate(prefix=prefix):
        if isinstance(sub_node, nodes.Field):
            yield index, sub_node


def enumerate_autonomous_children(
    node: nodes.FormulaItem,
    prefix: NodeHierarchyIndex = NodeHierarchyIndex(),
    exclude_node_types: Tuple[Type[nodes.FormulaItem], ...] = (),
    parent_stack: Tuple[nodes.FormulaItem, ...] = (),
) -> Iterable[Tuple[NodeHierarchyIndex, nodes.FormulaItem, Tuple[nodes.FormulaItem, ...]]]:
    """
    Return relative hierarchy indexes and nodes for all logical children -
    nodes that can be interpreted as standalone expressions.
    (e.g. a ``WhenPart`` or ``WindowGrouping`` can't - they're just parts
    of larger constructs - ``CaseBlock`` and ``WindowFunction`` respectively)
    """

    stack = parent_stack + (node,)
    for index, sub_node in node.enumerate(max_depth=1, prefix=prefix):
        if sub_node is node:
            # Don't need self - only children
            continue

        if isinstance(sub_node, exclude_node_types):
            continue

        if not sub_node.autonomous:
            yield from enumerate_autonomous_children(sub_node, prefix=index, parent_stack=stack)
            continue

        yield index, sub_node, stack


def autonomous_children_w_stack(
    node: nodes.FormulaItem,
    exclude_node_types: Tuple[Type[nodes.FormulaItem], ...] = (),
    parent_stack: Tuple[nodes.FormulaItem, ...] = (),
) -> Iterable[Tuple[nodes.FormulaItem, Tuple[nodes.FormulaItem, ...]]]:
    for _, child, stack in enumerate_autonomous_children(
        node,
        exclude_node_types=exclude_node_types,
        parent_stack=parent_stack,
    ):
        yield child, stack


def autonomous_children(
    node: nodes.FormulaItem,
    exclude_node_types: Tuple[Type[nodes.FormulaItem], ...] = (),
) -> Iterable[nodes.FormulaItem]:
    for _, child, _ in enumerate_autonomous_children(node, exclude_node_types=exclude_node_types):
        yield child


def collect_tags(node: nodes.FormulaItem) -> AbstractSet[LevelTag]:
    result: Set[LevelTag] = set()
    for _, child in node.enumerate():
        level_tag = child.level_tag
        if level_tag is not None:
            result.add(level_tag)
    return result


def get_wrapping_level(
    node: nodes.FormulaItem,
    qualify_node_cb: Callable[[nodes.FormulaItem], bool],
    stop_at_node_cb: Optional[Callable[[nodes.FormulaItem], bool]] = None,
    stop_at_level: Optional[int] = None,
) -> int:
    """
    Calculate wrapping level for node.

    This is the maximum number of times ``qualify_node_cb`` returns ``True``
    in a sub-tree of the given node.
    Recursion stops whenever ``stop_at_node_cb`` (if specified) returns ``True``.
    """

    def _get_level_recursively(cur_node: nodes.FormulaItem, level: int) -> int:
        if stop_at_node_cb is not None and stop_at_node_cb(cur_node):
            return level
        if qualify_node_cb(cur_node):
            level += 1

        if stop_at_level is not None and level >= stop_at_level:
            return level

        child_max_level = max((_get_level_recursively(child, level) for child in cur_node.children), default=level)
        return max(level, child_max_level)

    return _get_level_recursively(node, 0)


def _match_bfb_names(node: nodes.FormulaItem, bfb_names: AbstractSet[str]) -> bool:
    if isinstance(node, (nodes.FuncCall, fork_nodes.QueryFork)):
        return node.before_filter_by.field_names == bfb_names
    # Assume True if BFB is not supported
    return True


def get_window_function_wrapping_level(node: nodes.FormulaItem, bfb_names: AbstractSet[str]) -> int:
    return get_wrapping_level(
        node,
        qualify_node_cb=lambda _node: isinstance(_node, nodes.WindowFuncCall),
        stop_at_node_cb=lambda _node: (not _match_bfb_names(_node, bfb_names)),
    )


def is_double_aggregated_expression(node: nodes.FormulaItem) -> bool:
    level = get_wrapping_level(
        node,
        qualify_node_cb=lambda _node: dl_formula.inspect.node.is_aggregate_function(_node),
        stop_at_level=2,
    )
    return level == 2


def get_qfork_wrapping_level_until_winfunc(node: nodes.FormulaItem, bfb_names: AbstractSet[str]) -> int:
    return get_wrapping_level(
        node,
        qualify_node_cb=lambda _node: isinstance(_node, fork_nodes.QueryFork),
        stop_at_node_cb=lambda _node: (
            isinstance(_node, nodes.WindowFuncCall) or not _match_bfb_names(_node, bfb_names)
        ),
    )


def contains_non_default_lod_dimensions(node: nodes.FormulaItem) -> bool:
    if dl_formula.inspect.node.has_non_default_lod_dimensions(node):
        return True
    for child in autonomous_children(node):
        if contains_non_default_lod_dimensions(child):
            return True

    return False


def contains_extended_aggregations(node: nodes.FormulaItem, include_double_agg: bool = False) -> bool:
    if dl_formula.inspect.node.is_extended_aggregation(node):
        return True
    for child in autonomous_children(node):
        if contains_extended_aggregations(child, include_double_agg=False):
            return True

    if include_double_agg:
        return is_double_aggregated_expression(node)

    return False


def contains_lookup_functions(node: nodes.FormulaItem) -> bool:
    if dl_formula.inspect.node.is_lookup_function(node):
        return True
    for child in autonomous_children(node):
        if contains_lookup_functions(child):
            return True

    return False


def resolve_dimensions(
    node_stack: Iterable[nodes.FormulaItem],
    dimensions: List[nodes.FormulaItem],  # TODO: rename to global_dimensions
    env: "dl_formula.inspect.env.InspectionEnvironment",
    exclude_constants: bool = True,
) -> Tuple[List[nodes.FormulaItem], NodeSet, NodeSet]:
    """
    Return dimension list for current node and set of dimension extracts of the parent node
    ("the surrounding environment")
    """

    dimension_set = parent_dimension_set = NodeSet(dimensions)

    # Generate dimension list recursively
    for node in node_stack:  # from top level to the node in question
        if not (
            isinstance(node, nodes.FuncCall)
            and dl_formula.inspect.node.is_aggregate_function(node)
            or isinstance(node, fork_nodes.QueryFork)
        ):
            continue

        lod = node.lod
        if isinstance(lod, nodes.InheritedLodSpecifier):
            continue

        parent_dimensions = dimensions
        parent_dimension_set = dimension_set

        if isinstance(lod, nodes.FixedLodSpecifier):
            dimensions = list(lod.dim_list)
        elif isinstance(lod, nodes.IncludeLodSpecifier):
            dimensions = parent_dimensions + [dim for dim in lod.dim_list if dim not in parent_dimension_set]
        elif isinstance(lod, nodes.ExcludeLodSpecifier):
            exclude_set = NodeSet(lod.dim_list)
            dimensions = [dim for dim in parent_dimensions if dim not in exclude_set]
        else:
            raise TypeError(f"Invalid type for lod: {type(lod)}")

        if exclude_constants:
            dimensions = [dim for dim in dimensions if not is_constant_expression(dim, env=env)]

        dimension_set = NodeSet(dimensions)

    return dimensions, dimension_set, parent_dimension_set


def iter_aggregation_functions(node: nodes.FormulaItem) -> Generator[nodes.FuncCall, None, None]:
    """
    Iterate over all (indirect) children that are aggregate functions.
    Do not go inside aggregate functions.
    """

    for _, child in node.enumerate(max_depth=1):
        if child is node:
            # Don't need self - only children
            continue

        if dl_formula.inspect.node.is_aggregate_function(child):
            assert isinstance(child, nodes.FuncCall)
            yield child

        else:
            yield from iter_aggregation_functions(child)


def iter_operation_calls(node: nodes.FormulaItem) -> Generator[nodes.OperationCall, None, None]:
    """
    Iterate over all (indirect) children that are operation calls (functions and operators).
    """

    for _, child in node.enumerate():
        if isinstance(child, nodes.OperationCall):
            yield child


def get_toplevel_dimension_set(node: nodes.FormulaItem) -> NodeSet:
    toplevel_dimensions: NodeSet = NodeSet()
    for agg_func in iter_aggregation_functions(node):
        lod = agg_func.lod
        assert isinstance(lod, nodes.FixedLodSpecifier)  # Dimensions must already be resolved
        for dim in lod.dim_list:
            if dim not in toplevel_dimensions:
                toplevel_dimensions.add(dim)

    return toplevel_dimensions

lib/dl_formula/dl_formula/inspect/expression.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import lru_cache
from typing import Sequence

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.registry import (
    OPERATION_REGISTRY,
    FuncKey,
)
from dl_formula.definitions.scope import Scope
from dl_formula.inspect.registry.registry import LOWLEVEL_OP_REGISTRY


@lru_cache(maxsize=1000)
def can_be_aggregate(name: str) -> bool:
    """
    Return ``True`` if the function identified by ``name`` can an aggregate function
    (there may be a window and non-window variants).
    """
    return LOWLEVEL_OP_REGISTRY.can_be_aggregate(name=name)


@lru_cache(maxsize=1000)
def can_be_window(name: str) -> bool:
    """Return ``True`` if the function identified by ``name`` can be a window function"""
    return LOWLEVEL_OP_REGISTRY.can_be_window(name=name)


@lru_cache(maxsize=1000)
def can_be_nonwindow(name: str) -> bool:
    """Return ``True`` if the function identified by ``name`` can be a non-window function"""
    return LOWLEVEL_OP_REGISTRY.can_be_nonwindow(name=name)


@lru_cache(maxsize=1000)
def supports_grouping(name: str, is_window: bool) -> bool:
    """Return ``True`` if the function identified by ``name`` supports grouping clauses"""
    return LOWLEVEL_OP_REGISTRY.supports_grouping(name=name, is_window=is_window)


@lru_cache(maxsize=1000)
def supports_ordering(name: str, is_window: bool) -> bool:
    """Return ``True`` if the function identified by ``name`` supports ``ORDER BY`` clauses"""
    return LOWLEVEL_OP_REGISTRY.supports_ordering(name=name, is_window=is_window)


@lru_cache(maxsize=1000)
def supports_lod(name: str, is_window: bool) -> bool:
    """Return ``True`` if the function identified by ``name`` supports level-of-detail (LOD) clauses"""
    return LOWLEVEL_OP_REGISTRY.supports_lod(name=name, is_window=is_window)


@lru_cache(maxsize=1000)
def supports_ignore_dimensions(name: str, is_window: bool) -> bool:
    """Return ``True`` if the function identified by ``name`` supports ``IGNORE DIMENSIONS`` clauses"""
    return LOWLEVEL_OP_REGISTRY.supports_ignore_dimensions(name=name, is_window=is_window)


@lru_cache(maxsize=1000)
def supports_bfb(name: str, is_window: bool) -> bool:
    """Return ``True`` if the function identified by ``name`` supports ``BEFORE FILTER BY`` clauses"""
    return LOWLEVEL_OP_REGISTRY.supports_bfb(name=name, is_window=is_window)


@lru_cache(maxsize=1000)
def supports_extensions(name: str, is_window: bool) -> bool:
    """Return ``True`` if the function identified by ``name`` supports any kind of extension clause"""
    return (
        supports_lod(name=name, is_window=is_window)
        or supports_grouping(name=name, is_window=is_window)
        or supports_ordering(name=name, is_window=is_window)
        or supports_bfb(name=name, is_window=is_window)
        or supports_ignore_dimensions(name=name, is_window=is_window)
    )


@lru_cache(maxsize=1000)
def requires_grouping(name: str, is_window: bool) -> bool:
    """Return ``True`` if the (window) function identified by ``name`` requires the grouping clause"""
    return is_window and can_be_nonwindow(name=name)


@lru_cache(maxsize=1000)
def uses_default_ordering(name: str) -> bool:
    """Return ``True`` if the window function identified by ``name`` uses the same ordering as the query"""
    return LOWLEVEL_OP_REGISTRY.uses_default_ordering(name=name, is_window=True)


def get_return_type(name: str, arg_types: Sequence[DataType], is_window: bool = False) -> DataType:
    """Return the data type of the function result if it is called with the given argument types"""
    # FIXME: Switch to LOWLEVEL_OP_REGISTRY
    func_definition = OPERATION_REGISTRY.get_definition(
        name=name, arg_types=arg_types, is_window=is_window, for_any_dialect=True
    )
    return func_definition.get_return_type(arg_types=list(arg_types))


def get_supported_functions(
    require_dialects: DialectCombo = D.EMPTY,
    only_functions: bool = True,
    function_scopes: int = Scope.EXPLICIT_USAGE,
) -> list[FuncKey]:
    # FIXME: Switch to LOWLEVEL_OP_REGISTRY
    return OPERATION_REGISTRY.get_supported_functions(
        require_dialects=require_dialects,
        only_functions=only_functions,
        function_scopes=function_scopes,
    )

lib/dl_formula/dl_formula/inspect/function.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Union

from dl_formula.core.datatype import DataType
import dl_formula.core.nodes as nodes


_TYPE_BY_CLS = {
    nodes.LiteralBoolean: DataType.CONST_BOOLEAN,
    nodes.LiteralDate: DataType.CONST_DATE,
    nodes.LiteralDatetime: DataType.CONST_DATETIME,
    nodes.LiteralDatetimeTZ: DataType.CONST_DATETIMETZ,
    nodes.LiteralGenericDatetime: DataType.CONST_GENERICDATETIME,
    nodes.LiteralFloat: DataType.CONST_FLOAT,
    nodes.LiteralGeopoint: DataType.CONST_GEOPOINT,
    nodes.LiteralGeopolygon: DataType.CONST_GEOPOLYGON,
    nodes.LiteralInteger: DataType.CONST_INTEGER,
    nodes.LiteralString: DataType.CONST_STRING,
    nodes.LiteralUuid: DataType.CONST_UUID,
    nodes.LiteralArrayInteger: DataType.CONST_ARRAY_INT,
    nodes.LiteralArrayFloat: DataType.CONST_ARRAY_FLOAT,
    nodes.LiteralArrayString: DataType.CONST_ARRAY_STR,
    nodes.LiteralTreeString: DataType.CONST_TREE_STR,
    nodes.Null: DataType.NULL,
}


def get_data_type(node: Union[nodes.BaseLiteral, nodes.Null]):
    if not isinstance(node, (nodes.BaseLiteral, nodes.Null)):
        raise TypeError(type(node))

    return _TYPE_BY_CLS[type(node)]

lib/dl_formula/dl_formula/inspect/literal.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Optional

import dl_formula.core.fork_nodes as fork_nodes
import dl_formula.core.nodes as nodes
from dl_formula.inspect.function import (
    can_be_aggregate,
    supports_bfb,
)


def get_token(node: nodes.FormulaItem) -> Optional[str]:
    if isinstance(node, nodes.NamedItem):
        return node.name
    if isinstance(
        node,
        (
            nodes.LiteralUuid,
            nodes.LiteralString,
            nodes.LiteralBoolean,
            nodes.LiteralInteger,
            nodes.LiteralFloat,
            nodes.LiteralGeopoint,
            nodes.LiteralGeopolygon,
        ),
    ):
        return str(node.value)
    if isinstance(node, nodes.Null):
        return "null"
    return None


def is_aggregate_function(node: nodes.FormulaItem) -> bool:
    return (
        isinstance(node, nodes.FuncCall) and not isinstance(node, nodes.WindowFuncCall) and can_be_aggregate(node.name)
    )


def is_window_function(node: nodes.FormulaItem) -> bool:
    return isinstance(node, nodes.WindowFuncCall)


def is_lookup_function(node: nodes.FormulaItem) -> bool:
    # FIXME: Add a separate function attribute for that
    if isinstance(node, nodes.FuncCall) and not isinstance(node, nodes.WindowFuncCall):
        if not can_be_aggregate(node.name) and supports_bfb(node.name, is_window=False):
            return True

    return False


def has_non_default_lod_dimensions(node: nodes.FormulaItem) -> bool:
    if isinstance(node, nodes.FuncCall):
        if is_aggregate_function(node):
            if not isinstance(node.lod, nodes.DefaultAggregationLodSpecifier):
                # Aggregations with custom LODs
                return True

    return False


def is_extended_aggregation(node: nodes.FormulaItem) -> bool:
    if isinstance(node, nodes.FuncCall):
        if is_aggregate_function(node):
            if not isinstance(node.lod, nodes.DefaultAggregationLodSpecifier):
                # Aggregations with custom LODs
                return True
            if node.before_filter_by.field_names:
                return True

    return False


def is_default_lod_aggregation(node: nodes.FormulaItem) -> bool:
    if is_aggregate_function(node):
        assert isinstance(node, nodes.FuncCall)
        return isinstance(node.lod, nodes.DefaultAggregationLodSpecifier)

    return False


def qfork_is_aggregation(node: fork_nodes.QueryFork) -> bool:
    return not isinstance(node.lod, nodes.InheritedLodSpecifier) and is_aggregate_function(  # Only lookups have these
        node.result_expr
    )


def qfork_is_window(node: fork_nodes.QueryFork) -> bool:
    return not isinstance(node.lod, nodes.InheritedLodSpecifier) and is_window_function(  # Only lookups have these
        node.result_expr
    )


def qfork_is_lookup(node: fork_nodes.QueryFork) -> bool:
    return isinstance(node.lod, nodes.InheritedLodSpecifier)  # Only lookups have these

lib/dl_formula/dl_formula/inspect/node.py
datalens-tech/datalens-backend
from typing import (
    NamedTuple,
    Optional,
)


class BasicOpItemKey(NamedTuple):
    name: str
    arg_cnt: Optional[int]
    is_function: bool
    is_aggregate: bool
    is_window: bool


class BasicOpItem(NamedTuple):
    """
    Low-level translator-independent function/operator registry
    for inspection and parsing purposes.
    """

    name: str
    arg_cnt: Optional[int]
    is_function: bool
    is_aggregate: bool
    is_window: bool
    uses_default_ordering: bool
    supports_grouping: bool
    supports_ordering: bool
    supports_lod: bool
    supports_ignore_dimensions: bool
    supports_bfb: bool


class NameIsWinOpItemKey(NamedTuple):
    name: str
    is_window: bool


class NameIsWinOpItem(NamedTuple):
    name: str
    is_window: bool
    uses_default_ordering: bool
    supports_grouping: bool
    supports_ordering: bool
    supports_lod: bool
    supports_ignore_dimensions: bool
    supports_bfb: bool


class NameOpItem(NamedTuple):
    name: str
    can_be_aggregate: bool
    can_be_window: bool
    can_be_nonwindow: bool

lib/dl_formula/dl_formula/inspect/registry/item.py
datalens-tech/datalens-backend
from functools import lru_cache
from typing import Dict

from dl_formula.inspect.registry.item import (
    BasicOpItem,
    BasicOpItemKey,
    NameIsWinOpItem,
    NameIsWinOpItemKey,
    NameOpItem,
)


@lru_cache(maxsize=10000)
def norm_name(func_name: str) -> str:
    return func_name.lower()


class LowlevelOpRegistry:
    def __init__(self):
        self._basic_op_registry: Dict[BasicOpItemKey, BasicOpItem] = {}
        self._name_is_win_op_registry: Dict[NameIsWinOpItemKey, NameIsWinOpItem] = {}
        self._name_op_registry: Dict[str, NameOpItem] = {}

    def add(self, basic_item: BasicOpItem) -> None:
        self._add_to_basic_op_registry(basic_item)
        self._add_to_name_is_win_op_registry(basic_item)
        self._add_to_name_op_registry(basic_item)

    def _add_to_basic_op_registry(self, basic_item: BasicOpItem) -> None:
        basic_key = BasicOpItemKey(
            name=basic_item.name,
            arg_cnt=basic_item.arg_cnt,
            is_window=basic_item.is_window,
            is_aggregate=basic_item.is_aggregate,
            is_function=basic_item.is_function,
        )
        ex_basic_item = self._basic_op_registry.get(basic_key)
        if ex_basic_item is not None:
            assert basic_item == ex_basic_item
        else:
            self._basic_op_registry[basic_key] = basic_item

    def _add_to_name_is_win_op_registry(self, basic_item: BasicOpItem) -> None:
        name_is_win_key = NameIsWinOpItemKey(
            name=basic_item.name,
            is_window=basic_item.is_window,
        )
        name_is_win_item = NameIsWinOpItem(
            name=basic_item.name,
            is_window=basic_item.is_window,
            uses_default_ordering=basic_item.uses_default_ordering,
            supports_grouping=basic_item.supports_grouping,
            supports_ordering=basic_item.supports_ordering,
            supports_lod=basic_item.supports_lod,
            supports_ignore_dimensions=basic_item.supports_ignore_dimensions,
            supports_bfb=basic_item.supports_bfb,
        )
        ex_name_is_win_item = self._name_is_win_op_registry.get(name_is_win_key)
        if ex_name_is_win_item is not None:
            assert name_is_win_item == ex_name_is_win_item
        else:
            self._name_is_win_op_registry[name_is_win_key] = name_is_win_item

    def _add_to_name_op_registry(self, basic_item: BasicOpItem) -> None:
        name_key = basic_item.name
        can_be_aggregate: bool = False
        can_be_window: bool = False
        can_be_nonwindow: bool = False
        ex_name_item = self._name_op_registry.get(name_key)
        if ex_name_item is not None:
            can_be_aggregate = ex_name_item.can_be_aggregate
            can_be_window = ex_name_item.can_be_window
            can_be_nonwindow = ex_name_item.can_be_nonwindow
        name_item = NameOpItem(
            name=basic_item.name,
            can_be_aggregate=basic_item.is_aggregate or can_be_aggregate,
            can_be_window=basic_item.is_window or can_be_window,
            can_be_nonwindow=(not basic_item.is_window) or can_be_nonwindow,
        )
        self._name_op_registry[name_key] = name_item

    def can_be_aggregate(self, name: str) -> bool:
        name = norm_name(name)
        item = self._name_op_registry.get(name)
        return item is not None and item.can_be_aggregate

    def can_be_window(self, name: str) -> bool:
        name = norm_name(name)
        item = self._name_op_registry.get(name)
        return item is not None and item.can_be_window

    def can_be_nonwindow(self, name: str) -> bool:
        name = norm_name(name)
        item = self._name_op_registry.get(name)
        return item is not None and item.can_be_nonwindow

    def supports_grouping(self, name: str, is_window: bool) -> bool:
        name = norm_name(name)
        key = NameIsWinOpItemKey(name=name, is_window=is_window)
        item = self._name_is_win_op_registry.get(key)
        return item is not None and item.supports_grouping

    def supports_ordering(self, name: str, is_window: bool) -> bool:
        name = norm_name(name)
        key = NameIsWinOpItemKey(name=name, is_window=is_window)
        item = self._name_is_win_op_registry.get(key)
        return item is not None and item.supports_ordering

    def supports_ignore_dimensions(self, name: str, is_window: bool) -> bool:
        name = norm_name(name)
        key = NameIsWinOpItemKey(name=name, is_window=is_window)
        item = self._name_is_win_op_registry.get(key)
        return item is not None and item.supports_ignore_dimensions

    def supports_lod(self, name: str, is_window: bool) -> bool:
        name = norm_name(name)
        key = NameIsWinOpItemKey(name=name, is_window=is_window)
        item = self._name_is_win_op_registry.get(key)
        return item is not None and item.supports_lod

    def supports_bfb(self, name: str, is_window: bool) -> bool:
        name = norm_name(name)
        key = NameIsWinOpItemKey(name=name, is_window=is_window)
        item = self._name_is_win_op_registry.get(key)
        return item is not None and item.supports_bfb

    def uses_default_ordering(self, name: str, is_window: bool) -> bool:
        name = norm_name(name)
        key = NameIsWinOpItemKey(name=name, is_window=is_window)
        item = self._name_is_win_op_registry.get(key)
        return item is not None and item.uses_default_ordering


LOWLEVEL_OP_REGISTRY = LowlevelOpRegistry()

lib/dl_formula/dl_formula/inspect/registry/registry.py
datalens-tech/datalens-backend
from typing import (
    Collection,
    Optional,
)

import attr

from dl_formula.definitions.ll_registry_loader import (
    populate_ll_op_registry,
    populate_translation_registry,
)
from dl_formula.formula_connectors import (
    load_all_connectors,
    register_all_connectors,
)


@attr.s(frozen=True)
class FormulaLibraryConfig:
    # Whitelist of connector entrypoints to be loaded
    formula_connector_ep_names: Optional[Collection[str]] = attr.ib(kw_only=True, default=None)


def preload_bi_formula() -> None:
    """Loads all entrypoint connectors without registering them"""

    load_all_connectors()


def load_formula_lib(formula_lib_config: FormulaLibraryConfig = FormulaLibraryConfig()) -> None:
    """Initialize the library"""
    register_all_connectors(connector_ep_names=formula_lib_config.formula_connector_ep_names)
    populate_translation_registry()
    populate_ll_op_registry()

lib/dl_formula/dl_formula/loader.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging
from typing import Optional

import dl_formula.core.nodes as nodes
from dl_formula.mutation.mutation import FormulaMutation


LOGGER = logging.getLogger(__name__)


class RemapBfbMutation(FormulaMutation):
    """
    Applies a name mapper to all BEFORE FILTER BY clauses.
    """

    def __init__(self, name_mapping: Optional[dict[str, str]] = None):
        self._name_mapping = name_mapping

    def remap_name(self, name: str) -> str:
        if self._name_mapping is not None:
            name = self._name_mapping.get(name, name)
        return name

    def match_node(self, node: nodes.FormulaItem, parent_stack: tuple[nodes.FormulaItem, ...]) -> bool:
        return isinstance(node, nodes.BeforeFilterBy)

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.BeforeFilterBy)

        remapped_names = frozenset({self.remap_name(name) for name in old.field_names})
        return nodes.BeforeFilterBy.make(field_names=remapped_names, meta=old.meta)

lib/dl_formula/dl_formula/mutation/bfb.py
datalens-tech/datalens-backend
from __future__ import annotations

from itertools import chain
from typing import (
    TYPE_CHECKING,
    List,
    Tuple,
)

import attr

import dl_formula.core.nodes as nodes
from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.inspect.expression import resolve_dimensions
from dl_formula.mutation.mutation import FormulaMutation


if TYPE_CHECKING:
    from dl_formula.collections import NodeSet


@attr.s
class DimensionResolvingMutationBase(FormulaMutation):
    """
    A mutation base class that can resolve the expression's dimensions at any depth.
    """

    _global_dimensions: List[nodes.FormulaItem] = attr.ib(kw_only=True)
    _inspect_env: InspectionEnvironment = attr.ib(kw_only=True, factory=InspectionEnvironment)

    def _generate_dimensions(
        self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> Tuple[List[nodes.FormulaItem], NodeSet, NodeSet]:
        """
        Return dimension list for ``old`` and dimensions of the parent node
        ("the surrounding environment")
        """

        dimensions, dimension_set, parent_dimension_set = resolve_dimensions(
            node_stack=chain(parent_stack, (node,)),
            dimensions=self._global_dimensions,
            env=self._inspect_env,
        )
        return dimensions, dimension_set, parent_dimension_set

lib/dl_formula/dl_formula/mutation/dim_resolution.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from functools import cached_property
from itertools import chain
import logging
from typing import (
    Any,
    Callable,
    Optional,
    Tuple,
)

from dl_formula.core.dialect import DialectCombo
import dl_formula.core.nodes as nodes
from dl_formula.mutation.mutation import FormulaMutation


LOGGER = logging.getLogger(__name__)


class IgnoreParenthesisWrapperMutation(FormulaMutation):
    """
    A mutation that ignores (removes) all parenthesis wrappers
    because they do not change the meaning of any of the expressions,
    but may prevent comparisons and recognition of dimensions
    from working correctly.
    """

    def match_node(
        self,
        node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> bool:
        return type(node) is nodes.ParenthesizedExpr and node.level_tag is None

    def make_replacement(
        self,
        old: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.ParenthesizedExpr)
        return old.expr


class ConvertBlocksToFunctionsMutation(FormulaMutation):
    """
    A mutation that converts IF- and CASE-blocks to corresponding functions.
    """

    def match_node(
        self,
        node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> bool:
        return isinstance(node, nodes.IfBlock) or isinstance(node, nodes.IfBlock)

    def _make_replacement_for_if_block(self, node: nodes.IfBlock) -> nodes.FuncCall:
        args: list[nodes.FormulaItem] = [
            *chain.from_iterable(if_part.children for if_part in node.if_list),  # type: ignore
            node.else_expr,
        ]
        return nodes.FuncCall.make("if", args=args, meta=node.meta)

    def _make_replacement_for_case_block(self, node: nodes.CaseBlock) -> nodes.FuncCall:
        args: list[nodes.FormulaItem] = [
            node.case_expr,
            *chain.from_iterable(when_part.children for when_part in node.when_list),  # type: ignore
        ]
        if node.else_expr is not None:
            args.append(node.else_expr)
        return nodes.FuncCall.make("case", args=args, meta=node.meta)

    def make_replacement(
        self,
        old: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> nodes.FormulaItem:
        if isinstance(old, nodes.IfBlock):
            return self._make_replacement_for_if_block(old)
        elif isinstance(old, nodes.CaseBlock):
            return self._make_replacement_for_case_block(old)

        raise TypeError(type(old))


class OptimizeConstComparisonMutation(FormulaMutation):
    """
    A mutation that evaluates comparison expressions for constants
    and replaces them with the in-Python evaluation result.
    """

    _opt_ops = {
        "==": lambda left, right: left == right,
        "_==": lambda left, right: left == right,
        "!=": lambda left, right: left != right,
        "_!=": lambda left, right: left != right,
    }

    def match_node(
        self,
        node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> bool:
        return (
            isinstance(node, nodes.Binary)
            # Operator is supported
            and node.name in self._opt_ops
            # Both operands are constants
            and isinstance(node.left, nodes.BaseLiteral)
            and isinstance(node.right, nodes.BaseLiteral)
            # To stay on the safe side, only do this for same-type constants
            and type(node.left) is type(node.right)  # noqa
        )

    def make_replacement(
        self,
        old: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.Binary)
        left, right = old.left, old.right
        assert isinstance(left, nodes.BaseLiteral)
        assert isinstance(right, nodes.BaseLiteral)
        value = self._opt_ops[old.name](left.value, right.value)
        return nodes.LiteralBoolean.make(value, meta=old.meta)


class OptimizeUnaryBoolFunctions(FormulaMutation):
    _opt_ops: dict[str, dict[Optional[DialectCombo], Callable[[Any], bool]]] = {
        "isnull": {None: lambda node: node is None},
    }

    def __init__(self, dialects: DialectCombo):
        self.dialects = dialects

    @cached_property
    def opt_ops(self) -> dict[str, Callable[[Any], bool]]:
        lst = self.dialects.to_list()
        opt_ops = {}
        for k, v in self._opt_ops.items():
            if len(v) == 1:
                opt_ops[k] = v[None]
            elif self.dialects in v:
                opt_ops[k] = v[self.dialects]
            else:
                for dialect in lst:
                    if dialect in v:
                        opt_ops[k] = v[self.dialects]
                        break
                else:
                    opt_ops[k] = v[None]
        return opt_ops

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        return (
            isinstance(node, nodes.OperationCall)
            and len(node.args) == 1
            and node.name in self.opt_ops
            and isinstance(node.args[0], nodes.BaseLiteral)
        )

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.OperationCall)
        f = self.opt_ops[old.name]
        const_node = old.args[0]
        assert isinstance(const_node, nodes.BaseLiteral)
        return nodes.LiteralBoolean.make(value=f(const_node.value), meta=old.meta)

    @classmethod
    def register_dialect(cls, func_name: str, dialects: DialectCombo, f: Callable[[Any], bool]) -> None:
        for dialect in dialects.to_list(with_self=True):
            cls._opt_ops[func_name][dialect] = f


class OptimizeConstAndOrMutation(FormulaMutation):
    """
    A mutation that simplifies expressions such as:
    - <expr> AND <const>
    - <expr> OR <const>
    """

    _opt_ops = {
        "and": {
            True: lambda node, meta: node,
            False: lambda node, meta: nodes.LiteralBoolean.make(value=False, meta=meta),
        },
        "or": {
            True: lambda node, meta: nodes.LiteralBoolean.make(value=True, meta=meta),
            False: lambda node, meta: node,
        },
    }

    def match_node(
        self,
        node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> bool:
        return (
            isinstance(node, nodes.Binary)
            # Operator is supported
            and node.name in self._opt_ops
            # One of the operands is a constant
            and (isinstance(node.left, nodes.BaseLiteral) or isinstance(node.right, nodes.BaseLiteral))
        )

    def make_replacement(
        self,
        old: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.Binary)

        const_node: nodes.BaseLiteral
        expr_node: nodes.FormulaItem
        if isinstance(old.left, nodes.BaseLiteral):
            const_node, expr_node = old.left, old.right
        else:
            assert isinstance(old.right, nodes.BaseLiteral)
            const_node, expr_node = old.right, old.left

        new_node = self._opt_ops[old.name][bool(const_node.value)](expr_node, old.meta)
        return new_node


class OptimizeConstFuncMutation(FormulaMutation):
    """
    A mutation that simplifies expressions such as:
    - <expr> AND <const>
    - <expr> OR <const>
    """

    class ConstFuncOptimizer(abc.ABC):
        @abc.abstractmethod
        def can_optimize(self, node: nodes.FuncCall) -> bool:
            raise NotImplementedError()

        @abc.abstractmethod
        def optimize(self, node: nodes.FuncCall) -> nodes.FuncCall:
            raise NotImplementedError()

    class IfFuncOptimizer(ConstFuncOptimizer):
        def can_optimize(self, node: nodes.FuncCall) -> bool:
            assert node.name == "if"
            assert len(node.args) % 2 == 1
            for cond_arg in node.args[:-1:2]:  # excluding the last one (the "else" part) with a step of 2
                if isinstance(cond_arg, nodes.BaseLiteral):
                    return True

            return False

        def optimize(self, node: nodes.FuncCall) -> nodes.FuncCall:
            new_args: list[nodes.FormulaItem] = []
            for cond_arg, then_arg in zip(node.args[:-1:2], node.args[1:-1:2]):
                if isinstance(cond_arg, nodes.BaseLiteral):
                    if cond_arg.value is False:
                        # The condition is false, so skip this branch (don't add it to the optimized IF)
                        pass
                    if cond_arg.value is True:
                        # This is the first true condition, so just return its "then"
                        return then_arg

                else:
                    # Add the args without changes
                    new_args.append(cond_arg)
                    new_args.append(then_arg)

            new_args.append(node.args[-1])  # the "else"

            if len(new_args) == len(node.args):
                # Nothing seems to have changed
                return node

            if len(new_args) == 1:  # All conditions turned out to be false (only "else" in new args)
                return node.args[-1]  # the "else"

            return nodes.FuncCall.make("if", args=new_args, meta=node.meta)

    class CaseFuncOptimizer(ConstFuncOptimizer):
        def can_optimize(self, node: nodes.FuncCall) -> bool:
            assert node.name == "case"
            assert len(node.args) % 2 == 0 and len(node.args) >= 2
            if isinstance(node.args[0], nodes.BaseLiteral):
                # the CASE expression needs to be a const for the optimization
                for when_arg in node.args[1:-1:2]:
                    # CASE(<case_expr_0>, <when_expr_1>, <then_expr_2>, ..., <else_expr>)
                    if isinstance(when_arg, nodes.BaseLiteral):
                        return True

            return False

        def optimize(self, node: nodes.FuncCall) -> nodes.FuncCall:
            case_expr_node = node.args[0]
            if not isinstance(case_expr_node, nodes.BaseLiteral):
                return node

            new_args: list[nodes.FormulaItem] = []
            new_args.append(case_expr_node)
            case_expr_value = case_expr_node.value
            for when_arg, then_arg in zip(node.args[1:-1:2], node.args[2:-1:2]):
                if isinstance(when_arg, nodes.BaseLiteral):
                    when_expr_value = when_arg.value
                    if when_expr_value != case_expr_value:
                        # The "when" is not the same aas "case", so skip this branch
                        # (don't add it to the optimized CASE)
                        pass
                    else:  # They are equal
                        # This is the first matching "when", so just return its "then"
                        return then_arg

                else:
                    # Add the args without changes
                    new_args.append(when_arg)
                    new_args.append(then_arg)

            new_args.append(node.args[-1])  # the "else"

            if len(new_args) == len(node.args):
                # Nothing seems to have changed
                return node

            if len(new_args) == 2:  # All conditions turned out to be false
                return node.args[-1]  # the "else"

            return nodes.FuncCall.make("case", args=new_args, meta=node.meta)

    _func_optimizers: dict[str, ConstFuncOptimizer] = {
        "if": IfFuncOptimizer(),
        "case": CaseFuncOptimizer(),
    }

    def match_node(
        self,
        node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> bool:
        return (
            isinstance(node, nodes.FuncCall)
            # Operator is supported
            and node.name in self._func_optimizers
            # One of the operands is a constant
            and self._func_optimizers[node.name].can_optimize(node)
        )

    def make_replacement(
        self,
        old: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.FuncCall)

        return self._func_optimizers[old.name].optimize(old)

lib/dl_formula/dl_formula/mutation/general.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    List,
    Tuple,
)

import attr

import dl_formula.core.aux_nodes as aux_nodes
import dl_formula.core.exc as exc
import dl_formula.core.fork_nodes as fork_nodes
import dl_formula.core.nodes as nodes
from dl_formula.inspect.expression import is_aggregate_expression
from dl_formula.inspect.node import (
    is_aggregate_function,
    qfork_is_aggregation,
)
from dl_formula.mutation.dim_resolution import DimensionResolvingMutationBase
from dl_formula.mutation.mutation import FormulaMutation
from dl_formula.shortcuts import n


@attr.s
class ExtAggregationToQueryForkMutation(DimensionResolvingMutationBase):
    """
    A mutation that adds ``QueryFork`` nodes for all aggregations
    for future LOD handling.
    """

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        is_agg = is_aggregate_function(node)
        direct_parent = parent_stack[-1]
        already_patched = (
            isinstance(direct_parent, fork_nodes.QueryFork)
            and direct_parent.result_expr is node
            and qfork_is_aggregation(direct_parent)  # is an aggregation fork (vs. AGO forks)
        )
        return is_agg and not already_patched

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.FuncCall)

        dimensions, _, parent_dimension_set = self._generate_dimensions(node=old, parent_stack=parent_stack)
        lod = nodes.FixedLodSpecifier.make(dim_list=dimensions)

        condition_list: List[fork_nodes.JoinConditionBase] = []
        for dimension_expr in dimensions:
            if is_aggregate_expression(dimension_expr, env=self._inspect_env):
                return aux_nodes.ErrorNode.make(
                    err_code=exc.LodMeasureDimensionsError.default_code,
                    message="A measure has been used as a LOD dimension.",
                    meta=old.meta,
                )
            dim_condition = fork_nodes.SelfEqualityJoinCondition.make(expr=dimension_expr)
            condition_list.append(dim_condition)

        joining = fork_nodes.QueryForkJoiningWithList.make(condition_list=condition_list)

        # Replace lod in the original function
        old_updated = nodes.FuncCall.make(
            name=old.name,
            args=old.args,
            ignore_dimensions=old.ignore_dimensions,
            before_filter_by=old.before_filter_by,
            lod=lod,
            meta=old.meta,
        )

        return fork_nodes.QueryFork.make(
            join_type=fork_nodes.JoinType.inner,
            result_expr=old_updated,
            joining=joining,
            lod=lod,
            before_filter_by=old.before_filter_by,
            meta=old.meta,
        )


@attr.s
class DoubleAggregationCollapsingMutation(FormulaMutation):
    """
    A mutation removes double aggregations where possible.
    """

    # TODO: Add warnings when such optimizations are made

    _AGG_TRANSFORM_MAP = {
        # <parent agg> -> what to replace with
        "sum": lambda child, meta: child,
        "any": lambda child, meta: child,
        "max": lambda child, meta: child,
        "min": lambda child, meta: child,
        "count": lambda child, meta: n.func.INT(n.not_(n.func.ISNULL(child, meta=meta), meta=meta), meta=meta),
        "countd": lambda child, meta: n.func.INT(n.not_(n.func.ISNULL(child, meta=meta), meta=meta), meta=meta),
        # 'min' is a bit tricky (converts int -> float)
    }

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        """
        Match if node is agg, and "main" argument is also an agg
        and they have the same LOD and BFB.
        """

        if not is_aggregate_function(node):
            return False
        assert isinstance(node, nodes.FuncCall)

        if node.name not in self._AGG_TRANSFORM_MAP:
            # No rule for this aggregation, so abort
            return False

        if len(node.args) != 1:
            # No child to inspect, or too many of them
            return False

        child_agg_node = node.args[0]
        if not is_aggregate_function(child_agg_node):
            # Child is not agg
            return False
        assert isinstance(child_agg_node, nodes.FuncCall)

        if child_agg_node.lod != node.lod or child_agg_node.before_filter_by != node.before_filter_by:
            # Mismatch of advanced params
            return False

        return True

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.FuncCall)

        child_agg_node = old.args[0]

        node_transform = self._AGG_TRANSFORM_MAP[old.name]
        new_node = node_transform(child_agg_node, old.meta)

        assert isinstance(new_node, nodes.FormulaItem)
        return new_node

lib/dl_formula/dl_formula/mutation/lod.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    ClassVar,
    Dict,
    List,
    Optional,
    Sequence,
    Tuple,
    Type,
)

import attr

from dl_formula.collections import NodeSet
import dl_formula.core.aux_nodes as aux_nodes
import dl_formula.core.exc as exc
import dl_formula.core.fork_nodes as fork_nodes
import dl_formula.core.nodes as nodes
from dl_formula.inspect.expression import (
    is_aggregate_expression,
    used_fields,
)
from dl_formula.mutation.dim_resolution import DimensionResolvingMutationBase
from dl_formula.mutation.mutation import FormulaMutation


class LookupFunctionMutatorBase(abc.ABC):
    name: ClassVar[str]
    supported_arg_counts: Tuple[int, ...]

    @classmethod
    @abc.abstractmethod
    def get_result_expression(
        cls,
        func_args: Sequence[nodes.FormulaItem],
    ) -> nodes.FormulaItem:
        raise NotImplementedError

    @classmethod
    @abc.abstractmethod
    def get_lookup_dimension(
        cls,
        func_args: Sequence[nodes.FormulaItem],
    ) -> nodes.FormulaItem:
        raise NotImplementedError

    @classmethod
    @abc.abstractmethod
    def get_lookup_conditions(
        cls,
        func_args: Sequence[nodes.FormulaItem],
    ) -> List[fork_nodes.JoinConditionBase]:
        raise NotImplementedError


LOOKUP_MUTATOR_REGISTRY: Dict[str, Type[LookupFunctionMutatorBase]] = {}


def register_lookup_mutator(mutator_cls: Type[LookupFunctionMutatorBase]) -> Type[LookupFunctionMutatorBase]:
    LOOKUP_MUTATOR_REGISTRY[mutator_cls.name] = mutator_cls
    return mutator_cls


class DateLookupFunctionMutatorBase(LookupFunctionMutatorBase):
    @classmethod
    def get_result_expression(
        cls,
        func_args: Sequence[nodes.FormulaItem],
    ) -> nodes.FormulaItem:
        return func_args[0]

    @classmethod
    def get_lookup_dimension(
        cls,
        func_args: Sequence[nodes.FormulaItem],
    ) -> nodes.FormulaItem:
        return func_args[1]


_MONTH_BASED_UNITS = frozenset(("month", "year", "quarter"))


@register_lookup_mutator
class AgoLookupFunctionMutator(DateLookupFunctionMutatorBase):
    name = "ago"
    supported_arg_counts = (2, 3, 4)

    @classmethod
    def get_lookup_conditions(
        cls,
        func_args: Sequence[nodes.FormulaItem],
    ) -> List[fork_nodes.JoinConditionBase]:
        lookup_dimension = cls.get_lookup_dimension(func_args)
        fork_join_expr = nodes.FuncCall.make(
            name="dateadd",
            args=[lookup_dimension, *func_args[2:]],
        )
        main_lookup_condition = fork_nodes.BinaryJoinCondition.make(expr=lookup_dimension, fork_expr=fork_join_expr)

        conditions: List[fork_nodes.JoinConditionBase] = [main_lookup_condition]

        # An additional condition is needed for cases when the unit is month-based
        # to avoid date duplication for months with different lengths
        if len(func_args) > 2:  # default unit is "day", so we're not interested
            unit_name = "day"
            try:
                string_arg: nodes.LiteralString = [
                    arg for arg in func_args[2:] if isinstance(arg, nodes.LiteralString)
                ][0]
                unit_name = string_arg.value.lower()
            except IndexError:
                pass

            if unit_name in _MONTH_BASED_UNITS:
                day_eq_condition = fork_nodes.BinaryJoinCondition.make(
                    expr=nodes.FuncCall.make("day", args=[lookup_dimension]),
                    fork_expr=nodes.FuncCall.make("day", args=[lookup_dimension]),
                )
                conditions.append(day_eq_condition)

        return conditions


@register_lookup_mutator
class AtDateLookupFunctionMutator(DateLookupFunctionMutatorBase):
    name = "at_date"
    supported_arg_counts = (3,)

    @classmethod
    def get_lookup_conditions(
        cls,
        func_args: Sequence[nodes.FormulaItem],
    ) -> List[fork_nodes.JoinConditionBase]:
        lookup_dimension = cls.get_lookup_dimension(func_args)
        lookup_condition = fork_nodes.BinaryJoinCondition.make(expr=func_args[2], fork_expr=lookup_dimension)
        return [lookup_condition]


def _get_arg_error_node(node: nodes.FuncCall) -> Optional[aux_nodes.ErrorNode]:
    func_name = node.name
    mutator = LOOKUP_MUTATOR_REGISTRY[func_name]

    if len(node.args) not in mutator.supported_arg_counts:
        return aux_nodes.ErrorNode.make(
            err_code=exc.LookupFunctionArgNumberError.default_code,
            message=(
                f"Invalid number of arguments for function {func_name.upper()}: "
                f"{len(node.args)}. Either 2, 3 or 4 are needed."
            ),
            meta=node.meta,
        )

    return None


@attr.s
class LookupFunctionToQueryForkMutation(DimensionResolvingMutationBase):
    """
    A mutation that converts all lookup (`AGO`) function calls to ``QueryFork`` nodes
    using a list of explicit joining conditions.
    """

    _allow_empty_dimensions: bool = attr.ib(kw_only=True)

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        return isinstance(node, nodes.FuncCall) and node.name in LOOKUP_MUTATOR_REGISTRY

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.FuncCall)
        func_name = old.name
        assert func_name in LOOKUP_MUTATOR_REGISTRY

        err_node = _get_arg_error_node(old)
        if err_node is not None:
            return err_node

        mutator = LOOKUP_MUTATOR_REGISTRY[func_name]

        # Prepare result_expr
        result_expr = mutator.get_result_expression(old.args)
        if not is_aggregate_expression(node=result_expr, env=self._inspect_env):
            return aux_nodes.ErrorNode.make(
                err_code=exc.LookupFunctionWOAggregationError.default_code,
                message=f"Result expression of function {func_name.upper()} is not aggregated.",
                meta=old.meta,
            )

        # Prepare joining
        lookup_dimension = mutator.get_lookup_dimension(old.args)
        # Args will be validated as part of the condition expression
        lookup_conditions = mutator.get_lookup_conditions(old.args)
        ignore_dimensions_set = NodeSet(nodes=old.ignore_dimensions.children)

        if is_aggregate_expression(node=lookup_dimension, env=self._inspect_env):
            return aux_nodes.ErrorNode.make(
                err_code=exc.LookupFunctionAggregatedDimensionError.default_code,
                message=f"The lookup dimension of function {func_name.upper()} is an aggregation.",
                meta=old.meta,
            )
        if lookup_dimension in ignore_dimensions_set:
            return aux_nodes.ErrorNode.make(
                err_code=exc.LookupFunctionIgnoredLookupDimensionError.default_code,
                message=f"Cannot ignore lookup dimension of function {func_name.upper()}",
                meta=old.meta,
            )

        lookup_condition_used_fields: List[nodes.Field] = []
        for lookup_condition in lookup_conditions:
            lookup_condition_used_fields.extend(used_fields(lookup_condition))
        if not lookup_condition_used_fields:
            return aux_nodes.ErrorNode.make(
                err_code=exc.LookupFunctionConstantLookupDimensionError.default_code,
                message=f"Cannot use a constant expression as lookup dimension of function {func_name.upper()}",
                meta=old.meta,
            )

        dimensions, _, _ = self._generate_dimensions(node=old, parent_stack=parent_stack)

        condition_list: List[fork_nodes.JoinConditionBase] = []
        dim_list: List[nodes.FormulaItem] = []
        found_conditional_dimension = False
        for dimension_expr in dimensions:
            if dimension_expr in ignore_dimensions_set:
                # This dimension is specified in the function's IGNORE DIMENSIONS clause
                # so, do just that - skip it in the JOIN condition list
                continue

            dim_condition: fork_nodes.JoinConditionBase
            if dimension_expr == lookup_dimension:
                found_conditional_dimension = True
                condition_list.extend(lookup_conditions)
            else:
                dim_condition = fork_nodes.SelfEqualityJoinCondition.make(expr=dimension_expr)
                condition_list.append(dim_condition)
            dim_list.append(dimension_expr)

        if (not self._allow_empty_dimensions or dimensions) and not found_conditional_dimension:
            # This specific one wasn't found.
            # Dimensions can be empty during single-formula validation,
            # so we must allow this case to go on without errors
            return aux_nodes.ErrorNode.make(
                err_code=exc.LookupFunctionUnselectedDimensionError.default_code,
                message=(
                    f"Invalid dimension for function {func_name.upper()}. "
                    "It must be explicitly used in the data request as a dimension."
                ),
                meta=old.meta,
            )

        joining = fork_nodes.QueryForkJoiningWithList.make(condition_list=condition_list)
        lod = nodes.InheritedLodSpecifier()

        new_node = fork_nodes.QueryFork.make(
            join_type=fork_nodes.JoinType.left,
            result_expr=result_expr,
            joining=joining,
            before_filter_by=old.before_filter_by,
            lod=lod,
            meta=old.meta,
        )
        return new_node


@attr.s
class LookupDefaultBfbMutation(FormulaMutation):
    """
    A mutation that adds the lookup dimension of fork functions (`AGO` and `AT_DATE`)
    to the `BEFORE FILTER BY` clause.
    """

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        return isinstance(node, nodes.FuncCall) and node.name in LOOKUP_MUTATOR_REGISTRY

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.FuncCall)
        func_name = old.name
        assert func_name in LOOKUP_MUTATOR_REGISTRY

        err_node = _get_arg_error_node(old)
        if err_node is not None:
            return err_node

        mutator = LOOKUP_MUTATOR_REGISTRY[func_name]
        lookup_dimension = mutator.get_lookup_dimension(old.args)

        if isinstance(lookup_dimension, nodes.Field) and lookup_dimension.name not in old.before_filter_by.field_names:
            # Lookup dimension is not yet present in BFB, so add it.
            new_before_filter_by = nodes.BeforeFilterBy.make(
                field_names=old.before_filter_by.field_names | {lookup_dimension.name}
            )

            new_node = nodes.FuncCall.make(
                name=old.name,
                args=old.args,
                lod=old.lod,
                ignore_dimensions=old.ignore_dimensions,
                before_filter_by=new_before_filter_by,
                meta=old.meta,
            )

        else:
            # No changes
            new_node = old

        return new_node

lib/dl_formula/dl_formula/mutation/lookup.py
datalens-tech/datalens-backend
"""
Mutations allow making modification of nodes of a formula expression
that meet certain criteria.

Can be used to prepare the formula for translation and also for optimization of certain node constructs
"""

from __future__ import annotations

import abc
from typing import (
    Sequence,
    Tuple,
    TypeVar,
)

import dl_formula.core.nodes as nodes


class FormulaMutation(abc.ABC):
    """Class described how a formula object is transformed via replacements of its nodes"""

    @abc.abstractmethod
    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        """Check whether the given node matches internal criteria for replacement"""

        raise NotImplementedError

    @abc.abstractmethod
    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        """Generate a new node that will replace the old one in the formula"""

        raise NotImplementedError


_NODE_TV = TypeVar("_NODE_TV", bound=nodes.FormulaItem)


def apply_mutations(tree: _NODE_TV, mutations: Sequence[FormulaMutation]) -> _NODE_TV:
    """Apply multiple mutations to formula node tree"""

    def match_func(node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        for mutation in mutations:
            if mutation.match_node(node, parent_stack=parent_stack):
                return True
        return False

    def replace_func(node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> nodes.FormulaItem:
        replacement_made = False
        for mutation in mutations:
            if mutation.match_node(node, parent_stack=parent_stack):
                replacement_made = True
                node = mutation.make_replacement(node, parent_stack=parent_stack)
        if replacement_made:
            return node
        raise RuntimeError("Could not perform node replacement")

    return tree.replace_nodes(match_func=match_func, replace_func=replace_func)

lib/dl_formula/dl_formula/mutation/mutation.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    List,
    Sequence,
    Tuple,
)

import attr

import dl_formula.core.exc as exc
import dl_formula.core.fork_nodes as fork_nodes
import dl_formula.core.nodes as nodes
from dl_formula.inspect.expression import (
    is_aggregate_expression,
    is_bound_only_to,
)
from dl_formula.inspect.function import uses_default_ordering
from dl_formula.inspect.node import qfork_is_window
from dl_formula.mutation.dim_resolution import DimensionResolvingMutationBase
from dl_formula.mutation.mutation import FormulaMutation


class AmongToWithinGroupingMutation(DimensionResolvingMutationBase):
    """
    A mutation that converts all (relative) ``WITHIN`` window function groupings
    to (absolute) ``AMONG`` groupings.
    """

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        return isinstance(node, nodes.WindowGroupingAmong)

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.WindowGroupingAmong)
        dimensions, _, _ = self._generate_dimensions(node=old, parent_stack=parent_stack)
        new_dimensions = dimensions[:]
        for dimension in old.dim_list:
            for i in range(len(new_dimensions)):
                if new_dimensions[i].extract == dimension.extract:
                    del new_dimensions[i]
                    break
            else:
                raise exc.UnknownWindowDimensionError("Unknown dimension for window")

        return nodes.WindowGroupingWithin.make(dim_list=new_dimensions, meta=old.meta)


@attr.s
class IgnoreExtraWithinGroupingMutation(DimensionResolvingMutationBase):
    """
    A mutation that ignores (removes) all expressions from the``WITHIN`` clause
    that are not listed among dimensions or are not aggregate expressions.
    """

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        return isinstance(node, nodes.WindowGroupingWithin)

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.WindowGroupingWithin)
        _, dimension_set, _ = self._generate_dimensions(node=old, parent_stack=parent_stack)
        new_dimensions: List[nodes.FormulaItem] = [
            dimension
            for dimension in old.dim_list
            if (
                # Allow only known dimensions
                is_bound_only_to(node=dimension, allow_nodes=dimension_set)
                # and aggregate expressions.
                or is_aggregate_expression(node=dimension, env=self._inspect_env)
                # Ignore everything else.
            )
        ]
        return nodes.WindowGroupingWithin.make(dim_list=new_dimensions, meta=old.meta)


class DefaultWindowOrderingMutation(FormulaMutation):
    """
    A mutation that adds default ordering to all window function calls.

    If a function is not ordered before the mutation,
    then add the default ordering to it.
    If a function is already ordered, then add items
    from the default ordering that are not present in the function's ordering.
    """

    def __init__(self, default_order_by: Sequence[nodes.FormulaItem]):
        self._default_order_by = list(default_order_by)

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        return isinstance(node, nodes.WindowFuncCall) and uses_default_ordering(node.name)

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.WindowFuncCall)

        existing_order_by = list(old.ordering.expr_list)
        if not existing_order_by:
            # Node has no ordering of its own, so add it as-is
            new_order_by_list = self._default_order_by

        else:
            # For nodes that are already ordered patch orderings with items they don't contain
            # (ignore ASC/DESC when comparing ORDER BY items)

            def _unwrap_order_by_item(node: nodes.FormulaItem) -> nodes.FormulaItem:
                if isinstance(node, nodes.OrderingDirectionBase):
                    node = node.expr
                return node

            existing_order_by_extracts = [_unwrap_order_by_item(expr).extract for expr in existing_order_by]
            new_order_by_list = [n for n in old.ordering.expr_list]
            for default_order_item in self._default_order_by:
                if _unwrap_order_by_item(default_order_item).extract not in existing_order_by_extracts:
                    new_order_by_list.append(default_order_item)

        return nodes.WindowFuncCall.make(
            name=old.name,
            args=old.args,
            ordering=nodes.Ordering.make(expr_list=new_order_by_list),
            grouping=old.grouping,
            before_filter_by=old.before_filter_by,
            ignore_dimensions=old.ignore_dimensions,
            lod=old.lod,
            meta=old.meta,
        )


class WindowFunctionToQueryForkMutation(DimensionResolvingMutationBase):
    """
    A mutation that wraps all window functions into `QueryFork` nodes
    """

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        is_winfunc = isinstance(node, nodes.WindowFuncCall)
        direct_parent = parent_stack[-1]
        already_patched = (
            isinstance(direct_parent, fork_nodes.QueryFork)
            and direct_parent.result_expr is node
            and qfork_is_window(direct_parent)
        )
        return is_winfunc and not already_patched

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.WindowFuncCall)

        dimensions, _, parent_dimension_set = self._generate_dimensions(node=old, parent_stack=parent_stack)
        lod = nodes.FixedLodSpecifier.make(dim_list=dimensions)

        condition_list: List[fork_nodes.JoinConditionBase] = []
        for dimension_expr in dimensions:
            dim_condition = fork_nodes.SelfEqualityJoinCondition.make(expr=dimension_expr)
            condition_list.append(dim_condition)

        joining = fork_nodes.QueryForkJoiningWithList.make(condition_list=condition_list)

        return fork_nodes.QueryFork.make(
            join_type=fork_nodes.JoinType.inner,
            result_expr=old,
            joining=joining,
            lod=lod,
            before_filter_by=old.before_filter_by,
            meta=old.meta,
        )

lib/dl_formula/dl_formula/mutation/window.py
datalens-tech/datalens-backend
import sys

from antlr4 import DFA, ATNDeserializer, Lexer, LexerATNSimulator, PredictionContextCache

if sys.version_info[1] > 5:
    from typing import TextIO
else:
    from typing.io import TextIO


def serializedATN():
    return [
        4, 0, 49, 594, 6, -1, 2, 0, 7, 0, 2, 1, 7, 1, 2, 2, 7, 2, 2, 3, 7, 3, 2, 4, 7, 4, 2, 5, 7, 5,
        2, 6, 7, 6, 2, 7, 7, 7, 2, 8, 7, 8, 2, 9, 7, 9, 2, 10, 7, 10, 2, 11, 7, 11, 2, 12, 7, 12, 2,
        13, 7, 13, 2, 14, 7, 14, 2, 15, 7, 15, 2, 16, 7, 16, 2, 17, 7, 17, 2, 18, 7, 18, 2, 19, 7,
        19, 2, 20, 7, 20, 2, 21, 7, 21, 2, 22, 7, 22, 2, 23, 7, 23, 2, 24, 7, 24, 2, 25, 7, 25, 2,
        26, 7, 26, 2, 27, 7, 27, 2, 28, 7, 28, 2, 29, 7, 29, 2, 30, 7, 30, 2, 31, 7, 31, 2, 32, 7,
        32, 2, 33, 7, 33, 2, 34, 7, 34, 2, 35, 7, 35, 2, 36, 7, 36, 2, 37, 7, 37, 2, 38, 7, 38, 2,
        39, 7, 39, 2, 40, 7, 40, 2, 41, 7, 41, 2, 42, 7, 42, 2, 43, 7, 43, 2, 44, 7, 44, 2, 45, 7,
        45, 2, 46, 7, 46, 2, 47, 7, 47, 2, 48, 7, 48, 2, 49, 7, 49, 2, 50, 7, 50, 2, 51, 7, 51, 2,
        52, 7, 52, 2, 53, 7, 53, 2, 54, 7, 54, 2, 55, 7, 55, 2, 56, 7, 56, 2, 57, 7, 57, 2, 58, 7,
        58, 2, 59, 7, 59, 2, 60, 7, 60, 2, 61, 7, 61, 2, 62, 7, 62, 2, 63, 7, 63, 2, 64, 7, 64, 2,
        65, 7, 65, 2, 66, 7, 66, 2, 67, 7, 67, 2, 68, 7, 68, 2, 69, 7, 69, 2, 70, 7, 70, 2, 71, 7,
        71, 2, 72, 7, 72, 2, 73, 7, 73, 2, 74, 7, 74, 2, 75, 7, 75, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1,
        2, 1, 2, 1, 3, 1, 3, 1, 3, 1, 3, 5, 3, 165, 8, 3, 10, 3, 12, 3, 168, 9, 3, 1, 3, 3, 3, 171,
        8, 3, 1, 3, 1, 3, 1, 4, 1, 4, 1, 4, 1, 4, 1, 4, 4, 4, 180, 8, 4, 11, 4, 12, 4, 181, 1, 4, 5,
        4, 185, 8, 4, 10, 4, 12, 4, 188, 9, 4, 1, 4, 5, 4, 191, 8, 4, 10, 4, 12, 4, 194, 9, 4, 1,
        4, 1, 4, 1, 4, 1, 4, 1, 4, 1, 5, 4, 5, 202, 8, 5, 11, 5, 12, 5, 203, 1, 5, 1, 5, 1, 6, 1, 6,
        1, 7, 1, 7, 1, 8, 1, 8, 1, 9, 1, 9, 1, 10, 1, 10, 1, 11, 1, 11, 1, 12, 1, 12, 1, 13, 1, 13,
        1, 14, 1, 14, 1, 15, 1, 15, 1, 16, 1, 16, 1, 17, 1, 17, 1, 18, 1, 18, 1, 19, 1, 19, 1, 20,
        1, 20, 1, 21, 1, 21, 1, 22, 1, 22, 1, 23, 1, 23, 1, 24, 1, 24, 1, 25, 1, 25, 1, 26, 1, 26,
        1, 27, 1, 27, 1, 27, 1, 27, 1, 27, 1, 27, 1, 28, 1, 28, 1, 28, 1, 28, 1, 29, 1, 29, 1, 29,
        1, 29, 1, 30, 1, 30, 1, 30, 1, 30, 1, 30, 1, 30, 1, 30, 4, 30, 271, 8, 30, 11, 30, 12, 30,
        272, 1, 30, 1, 30, 1, 30, 1, 30, 1, 30, 1, 30, 1, 30, 4, 30, 282, 8, 30, 11, 30, 12, 30,
        283, 1, 30, 1, 30, 1, 30, 1, 31, 1, 31, 1, 31, 1, 31, 1, 31, 1, 31, 1, 31, 1, 31, 1, 32,
        1, 32, 1, 32, 1, 32, 1, 32, 1, 33, 1, 33, 1, 33, 1, 33, 1, 33, 1, 34, 1, 34, 1, 34, 1, 34,
        1, 34, 1, 35, 1, 35, 1, 35, 1, 35, 1, 35, 1, 35, 1, 35, 1, 36, 1, 36, 1, 36, 1, 36, 1, 37,
        1, 37, 1, 37, 1, 37, 1, 37, 1, 37, 1, 37, 1, 37, 1, 38, 1, 38, 1, 38, 1, 38, 1, 38, 1, 38,
        1, 39, 1, 39, 1, 39, 1, 39, 1, 39, 1, 39, 1, 40, 1, 40, 1, 40, 1, 41, 1, 41, 1, 41, 1, 41,
        1, 41, 1, 41, 1, 41, 4, 41, 353, 8, 41, 11, 41, 12, 41, 354, 1, 41, 1, 41, 1, 41, 1, 41,
        1, 41, 1, 41, 1, 41, 1, 41, 1, 41, 1, 41, 1, 41, 1, 42, 1, 42, 1, 42, 1, 43, 1, 43, 1, 43,
        1, 43, 1, 43, 1, 43, 1, 43, 1, 43, 1, 44, 1, 44, 1, 44, 1, 45, 1, 45, 1, 45, 1, 45, 1, 45,
        1, 46, 1, 46, 1, 46, 1, 46, 1, 47, 1, 47, 1, 47, 1, 47, 1, 47, 1, 48, 1, 48, 1, 48, 1, 49,
        1, 49, 1, 49, 1, 49, 1, 49, 1, 49, 4, 49, 405, 8, 49, 11, 49, 12, 49, 406, 1, 49, 1, 49,
        1, 49, 1, 50, 1, 50, 1, 50, 1, 50, 1, 50, 1, 51, 1, 51, 1, 51, 1, 51, 1, 51, 1, 51, 1, 52,
        1, 52, 1, 52, 1, 52, 1, 52, 1, 53, 1, 53, 1, 53, 1, 53, 1, 53, 1, 54, 1, 54, 1, 54, 1, 54,
        1, 54, 1, 54, 1, 54, 1, 55, 1, 55, 1, 56, 1, 56, 1, 57, 1, 57, 1, 58, 1, 58, 1, 59, 1, 59,
        1, 59, 1, 59, 1, 59, 1, 59, 1, 59, 1, 59, 1, 59, 1, 59, 1, 59, 3, 59, 459, 8, 59, 1, 60,
        1, 60, 1, 61, 1, 61, 1, 62, 1, 62, 1, 63, 1, 63, 1, 64, 1, 64, 1, 64, 3, 64, 472, 8, 64,
        1, 64, 1, 64, 3, 64, 476, 8, 64, 1, 65, 1, 65, 3, 65, 480, 8, 65, 1, 65, 1, 65, 1, 66, 4,
        66, 485, 8, 66, 11, 66, 12, 66, 486, 1, 67, 1, 67, 1, 67, 1, 67, 1, 67, 3, 67, 494, 8,
        67, 3, 67, 496, 8, 67, 1, 68, 1, 68, 1, 68, 1, 68, 5, 68, 502, 8, 68, 10, 68, 12, 68, 505,
        9, 68, 1, 68, 1, 68, 1, 68, 1, 68, 1, 68, 5, 68, 512, 8, 68, 10, 68, 12, 68, 515, 9, 68,
        1, 68, 3, 68, 518, 8, 68, 1, 69, 1, 69, 4, 69, 522, 8, 69, 11, 69, 12, 69, 523, 1, 69,
        1, 69, 1, 70, 1, 70, 3, 70, 530, 8, 70, 1, 70, 1, 70, 1, 70, 5, 70, 535, 8, 70, 10, 70,
        12, 70, 538, 9, 70, 1, 71, 1, 71, 1, 71, 1, 72, 1, 72, 1, 72, 1, 72, 1, 72, 1, 73, 1, 73,
        1, 73, 1, 73, 1, 73, 1, 73, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74,
        1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74,
        1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74, 1, 74,
        1, 74, 1, 74, 3, 74, 591, 8, 74, 1, 75, 1, 75, 0, 0, 76, 1, 1, 3, 2, 5, 3, 7, 4, 9, 5, 11,
        6, 13, 0, 15, 0, 17, 0, 19, 0, 21, 0, 23, 0, 25, 0, 27, 0, 29, 0, 31, 0, 33, 0, 35, 0, 37,
        0, 39, 0, 41, 0, 43, 0, 45, 0, 47, 0, 49, 0, 51, 0, 53, 0, 55, 7, 57, 8, 59, 9, 61, 10, 63,
        11, 65, 12, 67, 13, 69, 14, 71, 15, 73, 16, 75, 17, 77, 18, 79, 19, 81, 20, 83, 21, 85,
        22, 87, 23, 89, 24, 91, 25, 93, 26, 95, 27, 97, 28, 99, 29, 101, 30, 103, 31, 105, 32,
        107, 33, 109, 34, 111, 35, 113, 36, 115, 37, 117, 38, 119, 39, 121, 40, 123, 41, 125,
        0, 127, 0, 129, 0, 131, 0, 133, 42, 135, 43, 137, 44, 139, 45, 141, 46, 143, 0, 145,
        0, 147, 47, 149, 48, 151, 49, 1, 0, 34, 1, 0, 10, 10, 1, 1, 10, 10, 1, 0, 42, 42, 2, 0,
        42, 42, 47, 47, 4, 0, 9, 10, 12, 13, 32, 32, 160, 160, 2, 0, 65, 65, 97, 97, 2, 0, 66,
        66, 98, 98, 2, 0, 67, 67, 99, 99, 2, 0, 68, 68, 100, 100, 2, 0, 69, 69, 101, 101, 2, 0,
        70, 70, 102, 102, 2, 0, 71, 71, 103, 103, 2, 0, 72, 72, 104, 104, 2, 0, 73, 73, 105,
        105, 2, 0, 75, 75, 107, 107, 2, 0, 76, 76, 108, 108, 2, 0, 77, 77, 109, 109, 2, 0, 78,
        78, 110, 110, 2, 0, 79, 79, 111, 111, 2, 0, 82, 82, 114, 114, 2, 0, 83, 83, 115, 115,
        2, 0, 84, 84, 116, 116, 2, 0, 85, 85, 117, 117, 2, 0, 87, 87, 119, 119, 2, 0, 88, 88,
        120, 120, 2, 0, 89, 89, 121, 121, 3, 0, 37, 37, 42, 42, 47, 47, 1, 0, 48, 57, 2, 0, 65,
        90, 97, 122, 2, 0, 43, 43, 45, 45, 1, 0, 39, 39, 1, 0, 34, 34, 4, 0, 91, 91, 93, 93, 123,
        123, 125, 125, 3, 0, 32, 32, 84, 84, 116, 116, 601, 0, 1, 1, 0, 0, 0, 0, 3, 1, 0, 0, 0,
        0, 5, 1, 0, 0, 0, 0, 7, 1, 0, 0, 0, 0, 9, 1, 0, 0, 0, 0, 11, 1, 0, 0, 0, 0, 55, 1, 0, 0, 0, 0,
        57, 1, 0, 0, 0, 0, 59, 1, 0, 0, 0, 0, 61, 1, 0, 0, 0, 0, 63, 1, 0, 0, 0, 0, 65, 1, 0, 0, 0, 0,
        67, 1, 0, 0, 0, 0, 69, 1, 0, 0, 0, 0, 71, 1, 0, 0, 0, 0, 73, 1, 0, 0, 0, 0, 75, 1, 0, 0, 0, 0,
        77, 1, 0, 0, 0, 0, 79, 1, 0, 0, 0, 0, 81, 1, 0, 0, 0, 0, 83, 1, 0, 0, 0, 0, 85, 1, 0, 0, 0, 0,
        87, 1, 0, 0, 0, 0, 89, 1, 0, 0, 0, 0, 91, 1, 0, 0, 0, 0, 93, 1, 0, 0, 0, 0, 95, 1, 0, 0, 0, 0,
        97, 1, 0, 0, 0, 0, 99, 1, 0, 0, 0, 0, 101, 1, 0, 0, 0, 0, 103, 1, 0, 0, 0, 0, 105, 1, 0, 0,
        0, 0, 107, 1, 0, 0, 0, 0, 109, 1, 0, 0, 0, 0, 111, 1, 0, 0, 0, 0, 113, 1, 0, 0, 0, 0, 115,
        1, 0, 0, 0, 0, 117, 1, 0, 0, 0, 0, 119, 1, 0, 0, 0, 0, 121, 1, 0, 0, 0, 0, 123, 1, 0, 0, 0,
        0, 133, 1, 0, 0, 0, 0, 135, 1, 0, 0, 0, 0, 137, 1, 0, 0, 0, 0, 139, 1, 0, 0, 0, 0, 141, 1,
        0, 0, 0, 0, 147, 1, 0, 0, 0, 0, 149, 1, 0, 0, 0, 0, 151, 1, 0, 0, 0, 1, 153, 1, 0, 0, 0, 3,
        155, 1, 0, 0, 0, 5, 158, 1, 0, 0, 0, 7, 160, 1, 0, 0, 0, 9, 174, 1, 0, 0, 0, 11, 201, 1, 0,
        0, 0, 13, 207, 1, 0, 0, 0, 15, 209, 1, 0, 0, 0, 17, 211, 1, 0, 0, 0, 19, 213, 1, 0, 0, 0,
        21, 215, 1, 0, 0, 0, 23, 217, 1, 0, 0, 0, 25, 219, 1, 0, 0, 0, 27, 221, 1, 0, 0, 0, 29, 223,
        1, 0, 0, 0, 31, 225, 1, 0, 0, 0, 33, 227, 1, 0, 0, 0, 35, 229, 1, 0, 0, 0, 37, 231, 1, 0,
        0, 0, 39, 233, 1, 0, 0, 0, 41, 235, 1, 0, 0, 0, 43, 237, 1, 0, 0, 0, 45, 239, 1, 0, 0, 0,
        47, 241, 1, 0, 0, 0, 49, 243, 1, 0, 0, 0, 51, 245, 1, 0, 0, 0, 53, 247, 1, 0, 0, 0, 55, 249,
        1, 0, 0, 0, 57, 255, 1, 0, 0, 0, 59, 259, 1, 0, 0, 0, 61, 263, 1, 0, 0, 0, 63, 288, 1, 0,
        0, 0, 65, 296, 1, 0, 0, 0, 67, 301, 1, 0, 0, 0, 69, 306, 1, 0, 0, 0, 71, 311, 1, 0, 0, 0,
        73, 318, 1, 0, 0, 0, 75, 322, 1, 0, 0, 0, 77, 330, 1, 0, 0, 0, 79, 336, 1, 0, 0, 0, 81, 342,
        1, 0, 0, 0, 83, 345, 1, 0, 0, 0, 85, 367, 1, 0, 0, 0, 87, 370, 1, 0, 0, 0, 89, 378, 1, 0,
        0, 0, 91, 381, 1, 0, 0, 0, 93, 386, 1, 0, 0, 0, 95, 390, 1, 0, 0, 0, 97, 395, 1, 0, 0, 0,
        99, 398, 1, 0, 0, 0, 101, 411, 1, 0, 0, 0, 103, 416, 1, 0, 0, 0, 105, 422, 1, 0, 0, 0, 107,
        427, 1, 0, 0, 0, 109, 432, 1, 0, 0, 0, 111, 439, 1, 0, 0, 0, 113, 441, 1, 0, 0, 0, 115,
        443, 1, 0, 0, 0, 117, 445, 1, 0, 0, 0, 119, 458, 1, 0, 0, 0, 121, 460, 1, 0, 0, 0, 123,
        462, 1, 0, 0, 0, 125, 464, 1, 0, 0, 0, 127, 466, 1, 0, 0, 0, 129, 475, 1, 0, 0, 0, 131,
        477, 1, 0, 0, 0, 133, 484, 1, 0, 0, 0, 135, 495, 1, 0, 0, 0, 137, 517, 1, 0, 0, 0, 139,
        519, 1, 0, 0, 0, 141, 529, 1, 0, 0, 0, 143, 539, 1, 0, 0, 0, 145, 542, 1, 0, 0, 0, 147,
        547, 1, 0, 0, 0, 149, 590, 1, 0, 0, 0, 151, 592, 1, 0, 0, 0, 153, 154, 5, 35, 0, 0, 154,
        2, 1, 0, 0, 0, 155, 156, 5, 35, 0, 0, 156, 157, 5, 35, 0, 0, 157, 4, 1, 0, 0, 0, 158, 159,
        5, 44, 0, 0, 159, 6, 1, 0, 0, 0, 160, 161, 5, 45, 0, 0, 161, 162, 5, 45, 0, 0, 162, 166,
        1, 0, 0, 0, 163, 165, 8, 0, 0, 0, 164, 163, 1, 0, 0, 0, 165, 168, 1, 0, 0, 0, 166, 164,
        1, 0, 0, 0, 166, 167, 1, 0, 0, 0, 167, 170, 1, 0, 0, 0, 168, 166, 1, 0, 0, 0, 169, 171,
        7, 1, 0, 0, 170, 169, 1, 0, 0, 0, 171, 172, 1, 0, 0, 0, 172, 173, 6, 3, 0, 0, 173, 8, 1,
        0, 0, 0, 174, 175, 5, 47, 0, 0, 175, 176, 5, 42, 0, 0, 176, 186, 1, 0, 0, 0, 177, 185,
        8, 2, 0, 0, 178, 180, 7, 2, 0, 0, 179, 178, 1, 0, 0, 0, 180, 181, 1, 0, 0, 0, 181, 179,
        1, 0, 0, 0, 181, 182, 1, 0, 0, 0, 182, 183, 1, 0, 0, 0, 183, 185, 8, 3, 0, 0, 184, 177,
        1, 0, 0, 0, 184, 179, 1, 0, 0, 0, 185, 188, 1, 0, 0, 0, 186, 184, 1, 0, 0, 0, 186, 187,
        1, 0, 0, 0, 187, 192, 1, 0, 0, 0, 188, 186, 1, 0, 0, 0, 189, 191, 7, 2, 0, 0, 190, 189,
        1, 0, 0, 0, 191, 194, 1, 0, 0, 0, 192, 190, 1, 0, 0, 0, 192, 193, 1, 0, 0, 0, 193, 195,
        1, 0, 0, 0, 194, 192, 1, 0, 0, 0, 195, 196, 5, 42, 0, 0, 196, 197, 5, 47, 0, 0, 197, 198,
        1, 0, 0, 0, 198, 199, 6, 4, 0, 0, 199, 10, 1, 0, 0, 0, 200, 202, 7, 4, 0, 0, 201, 200, 1,
        0, 0, 0, 202, 203, 1, 0, 0, 0, 203, 201, 1, 0, 0, 0, 203, 204, 1, 0, 0, 0, 204, 205, 1,
        0, 0, 0, 205, 206, 6, 5, 0, 0, 206, 12, 1, 0, 0, 0, 207, 208, 7, 5, 0, 0, 208, 14, 1, 0,
        0, 0, 209, 210, 7, 6, 0, 0, 210, 16, 1, 0, 0, 0, 211, 212, 7, 7, 0, 0, 212, 18, 1, 0, 0,
        0, 213, 214, 7, 8, 0, 0, 214, 20, 1, 0, 0, 0, 215, 216, 7, 9, 0, 0, 216, 22, 1, 0, 0, 0,
        217, 218, 7, 10, 0, 0, 218, 24, 1, 0, 0, 0, 219, 220, 7, 11, 0, 0, 220, 26, 1, 0, 0, 0,
        221, 222, 7, 12, 0, 0, 222, 28, 1, 0, 0, 0, 223, 224, 7, 13, 0, 0, 224, 30, 1, 0, 0, 0,
        225, 226, 7, 14, 0, 0, 226, 32, 1, 0, 0, 0, 227, 228, 7, 15, 0, 0, 228, 34, 1, 0, 0, 0,
        229, 230, 7, 16, 0, 0, 230, 36, 1, 0, 0, 0, 231, 232, 7, 17, 0, 0, 232, 38, 1, 0, 0, 0,
        233, 234, 7, 18, 0, 0, 234, 40, 1, 0, 0, 0, 235, 236, 7, 19, 0, 0, 236, 42, 1, 0, 0, 0,
        237, 238, 7, 20, 0, 0, 238, 44, 1, 0, 0, 0, 239, 240, 7, 21, 0, 0, 240, 46, 1, 0, 0, 0,
        241, 242, 7, 22, 0, 0, 242, 48, 1, 0, 0, 0, 243, 244, 7, 23, 0, 0, 244, 50, 1, 0, 0, 0,
        245, 246, 7, 24, 0, 0, 246, 52, 1, 0, 0, 0, 247, 248, 7, 25, 0, 0, 248, 54, 1, 0, 0, 0,
        249, 250, 3, 13, 6, 0, 250, 251, 3, 35, 17, 0, 251, 252, 3, 39, 19, 0, 252, 253, 3, 37,
        18, 0, 253, 254, 3, 25, 12, 0, 254, 56, 1, 0, 0, 0, 255, 256, 3, 13, 6, 0, 256, 257, 3,
        37, 18, 0, 257, 258, 3, 19, 9, 0, 258, 58, 1, 0, 0, 0, 259, 260, 3, 13, 6, 0, 260, 261,
        3, 43, 21, 0, 261, 262, 3, 17, 8, 0, 262, 60, 1, 0, 0, 0, 263, 264, 3, 15, 7, 0, 264, 265,
        3, 21, 10, 0, 265, 266, 3, 23, 11, 0, 266, 267, 3, 39, 19, 0, 267, 268, 3, 41, 20, 0,
        268, 270, 3, 21, 10, 0, 269, 271, 5, 32, 0, 0, 270, 269, 1, 0, 0, 0, 271, 272, 1, 0, 0,
        0, 272, 270, 1, 0, 0, 0, 272, 273, 1, 0, 0, 0, 273, 274, 1, 0, 0, 0, 274, 275, 3, 23, 11,
        0, 275, 276, 3, 29, 14, 0, 276, 277, 3, 33, 16, 0, 277, 278, 3, 45, 22, 0, 278, 279,
        3, 21, 10, 0, 279, 281, 3, 41, 20, 0, 280, 282, 5, 32, 0, 0, 281, 280, 1, 0, 0, 0, 282,
        283, 1, 0, 0, 0, 283, 281, 1, 0, 0, 0, 283, 284, 1, 0, 0, 0, 284, 285, 1, 0, 0, 0, 285,
        286, 3, 15, 7, 0, 286, 287, 3, 53, 26, 0, 287, 62, 1, 0, 0, 0, 288, 289, 3, 15, 7, 0, 289,
        290, 3, 21, 10, 0, 290, 291, 3, 45, 22, 0, 291, 292, 3, 49, 24, 0, 292, 293, 3, 21, 10,
        0, 293, 294, 3, 21, 10, 0, 294, 295, 3, 37, 18, 0, 295, 64, 1, 0, 0, 0, 296, 297, 3, 17,
        8, 0, 297, 298, 3, 13, 6, 0, 298, 299, 3, 43, 21, 0, 299, 300, 3, 21, 10, 0, 300, 66,
        1, 0, 0, 0, 301, 302, 3, 19, 9, 0, 302, 303, 3, 21, 10, 0, 303, 304, 3, 43, 21, 0, 304,
        305, 3, 17, 8, 0, 305, 68, 1, 0, 0, 0, 306, 307, 3, 21, 10, 0, 307, 308, 3, 33, 16, 0,
        308, 309, 3, 43, 21, 0, 309, 310, 3, 21, 10, 0, 310, 70, 1, 0, 0, 0, 311, 312, 3, 21,
        10, 0, 312, 313, 3, 33, 16, 0, 313, 314, 3, 43, 21, 0, 314, 315, 3, 21, 10, 0, 315, 316,
        3, 29, 14, 0, 316, 317, 3, 23, 11, 0, 317, 72, 1, 0, 0, 0, 318, 319, 3, 21, 10, 0, 319,
        320, 3, 37, 18, 0, 320, 321, 3, 19, 9, 0, 321, 74, 1, 0, 0, 0, 322, 323, 3, 21, 10, 0,
        323, 324, 3, 51, 25, 0, 324, 325, 3, 17, 8, 0, 325, 326, 3, 33, 16, 0, 326, 327, 3, 47,
        23, 0, 327, 328, 3, 19, 9, 0, 328, 329, 3, 21, 10, 0, 329, 76, 1, 0, 0, 0, 330, 331, 3,
        23, 11, 0, 331, 332, 3, 13, 6, 0, 332, 333, 3, 33, 16, 0, 333, 334, 3, 43, 21, 0, 334,
        335, 3, 21, 10, 0, 335, 78, 1, 0, 0, 0, 336, 337, 3, 23, 11, 0, 337, 338, 3, 29, 14, 0,
        338, 339, 3, 51, 25, 0, 339, 340, 3, 21, 10, 0, 340, 341, 3, 19, 9, 0, 341, 80, 1, 0,
        0, 0, 342, 343, 3, 29, 14, 0, 343, 344, 3, 23, 11, 0, 344, 82, 1, 0, 0, 0, 345, 346, 3,
        29, 14, 0, 346, 347, 3, 25, 12, 0, 347, 348, 3, 37, 18, 0, 348, 349, 3, 39, 19, 0, 349,
        350, 3, 41, 20, 0, 350, 352, 3, 21, 10, 0, 351, 353, 5, 32, 0, 0, 352, 351, 1, 0, 0, 0,
        353, 354, 1, 0, 0, 0, 354, 352, 1, 0, 0, 0, 354, 355, 1, 0, 0, 0, 355, 356, 1, 0, 0, 0,
        356, 357, 3, 19, 9, 0, 357, 358, 3, 29, 14, 0, 358, 359, 3, 35, 17, 0, 359, 360, 3, 21,
        10, 0, 360, 361, 3, 37, 18, 0, 361, 362, 3, 43, 21, 0, 362, 363, 3, 29, 14, 0, 363, 364,
        3, 39, 19, 0, 364, 365, 3, 37, 18, 0, 365, 366, 3, 43, 21, 0, 366, 84, 1, 0, 0, 0, 367,
        368, 3, 29, 14, 0, 368, 369, 3, 37, 18, 0, 369, 86, 1, 0, 0, 0, 370, 371, 3, 29, 14, 0,
        371, 372, 3, 37, 18, 0, 372, 373, 3, 17, 8, 0, 373, 374, 3, 33, 16, 0, 374, 375, 3, 47,
        23, 0, 375, 376, 3, 19, 9, 0, 376, 377, 3, 21, 10, 0, 377, 88, 1, 0, 0, 0, 378, 379, 3,
        29, 14, 0, 379, 380, 3, 43, 21, 0, 380, 90, 1, 0, 0, 0, 381, 382, 3, 33, 16, 0, 382, 383,
        3, 29, 14, 0, 383, 384, 3, 31, 15, 0, 384, 385, 3, 21, 10, 0, 385, 92, 1, 0, 0, 0, 386,
        387, 3, 37, 18, 0, 387, 388, 3, 39, 19, 0, 388, 389, 3, 45, 22, 0, 389, 94, 1, 0, 0, 0,
        390, 391, 3, 37, 18, 0, 391, 392, 3, 47, 23, 0, 392, 393, 3, 33, 16, 0, 393, 394, 3,
        33, 16, 0, 394, 96, 1, 0, 0, 0, 395, 396, 3, 39, 19, 0, 396, 397, 3, 41, 20, 0, 397, 98,
        1, 0, 0, 0, 398, 399, 3, 39, 19, 0, 399, 400, 3, 41, 20, 0, 400, 401, 3, 19, 9, 0, 401,
        402, 3, 21, 10, 0, 402, 404, 3, 41, 20, 0, 403, 405, 5, 32, 0, 0, 404, 403, 1, 0, 0, 0,
        405, 406, 1, 0, 0, 0, 406, 404, 1, 0, 0, 0, 406, 407, 1, 0, 0, 0, 407, 408, 1, 0, 0, 0,
        408, 409, 3, 15, 7, 0, 409, 410, 3, 53, 26, 0, 410, 100, 1, 0, 0, 0, 411, 412, 3, 45,
        22, 0, 412, 413, 3, 27, 13, 0, 413, 414, 3, 21, 10, 0, 414, 415, 3, 37, 18, 0, 415, 102,
        1, 0, 0, 0, 416, 417, 3, 45, 22, 0, 417, 418, 3, 39, 19, 0, 418, 419, 3, 45, 22, 0, 419,
        420, 3, 13, 6, 0, 420, 421, 3, 33, 16, 0, 421, 104, 1, 0, 0, 0, 422, 423, 3, 45, 22, 0,
        423, 424, 3, 41, 20, 0, 424, 425, 3, 47, 23, 0, 425, 426, 3, 21, 10, 0, 426, 106, 1,
        0, 0, 0, 427, 428, 3, 49, 24, 0, 428, 429, 3, 27, 13, 0, 429, 430, 3, 21, 10, 0, 430,
        431, 3, 37, 18, 0, 431, 108, 1, 0, 0, 0, 432, 433, 3, 49, 24, 0, 433, 434, 3, 29, 14,
        0, 434, 435, 3, 45, 22, 0, 435, 436, 3, 27, 13, 0, 436, 437, 3, 29, 14, 0, 437, 438,
        3, 37, 18, 0, 438, 110, 1, 0, 0, 0, 439, 440, 5, 43, 0, 0, 440, 112, 1, 0, 0, 0, 441, 442,
        5, 45, 0, 0, 442, 114, 1, 0, 0, 0, 443, 444, 5, 94, 0, 0, 444, 116, 1, 0, 0, 0, 445, 446,
        7, 26, 0, 0, 446, 118, 1, 0, 0, 0, 447, 459, 5, 61, 0, 0, 448, 449, 5, 33, 0, 0, 449, 459,
        5, 61, 0, 0, 450, 451, 5, 60, 0, 0, 451, 459, 5, 62, 0, 0, 452, 459, 5, 62, 0, 0, 453,
        454, 5, 62, 0, 0, 454, 459, 5, 61, 0, 0, 455, 459, 5, 60, 0, 0, 456, 457, 5, 60, 0, 0,
        457, 459, 5, 61, 0, 0, 458, 447, 1, 0, 0, 0, 458, 448, 1, 0, 0, 0, 458, 450, 1, 0, 0, 0,
        458, 452, 1, 0, 0, 0, 458, 453, 1, 0, 0, 0, 458, 455, 1, 0, 0, 0, 458, 456, 1, 0, 0, 0,
        459, 120, 1, 0, 0, 0, 460, 461, 5, 40, 0, 0, 461, 122, 1, 0, 0, 0, 462, 463, 5, 41, 0,
        0, 463, 124, 1, 0, 0, 0, 464, 465, 7, 27, 0, 0, 465, 126, 1, 0, 0, 0, 466, 467, 7, 28,
        0, 0, 467, 128, 1, 0, 0, 0, 468, 469, 3, 133, 66, 0, 469, 471, 5, 46, 0, 0, 470, 472,
        3, 133, 66, 0, 471, 470, 1, 0, 0, 0, 471, 472, 1, 0, 0, 0, 472, 476, 1, 0, 0, 0, 473, 474,
        5, 46, 0, 0, 474, 476, 3, 133, 66, 0, 475, 468, 1, 0, 0, 0, 475, 473, 1, 0, 0, 0, 476,
        130, 1, 0, 0, 0, 477, 479, 7, 9, 0, 0, 478, 480, 7, 29, 0, 0, 479, 478, 1, 0, 0, 0, 479,
        480, 1, 0, 0, 0, 480, 481, 1, 0, 0, 0, 481, 482, 3, 133, 66, 0, 482, 132, 1, 0, 0, 0, 483,
        485, 3, 125, 62, 0, 484, 483, 1, 0, 0, 0, 485, 486, 1, 0, 0, 0, 486, 484, 1, 0, 0, 0, 486,
        487, 1, 0, 0, 0, 487, 134, 1, 0, 0, 0, 488, 489, 3, 133, 66, 0, 489, 490, 3, 131, 65,
        0, 490, 496, 1, 0, 0, 0, 491, 493, 3, 129, 64, 0, 492, 494, 3, 131, 65, 0, 493, 492,
        1, 0, 0, 0, 493, 494, 1, 0, 0, 0, 494, 496, 1, 0, 0, 0, 495, 488, 1, 0, 0, 0, 495, 491,
        1, 0, 0, 0, 496, 136, 1, 0, 0, 0, 497, 503, 5, 39, 0, 0, 498, 499, 5, 92, 0, 0, 499, 502,
        5, 39, 0, 0, 500, 502, 8, 30, 0, 0, 501, 498, 1, 0, 0, 0, 501, 500, 1, 0, 0, 0, 502, 505,
        1, 0, 0, 0, 503, 501, 1, 0, 0, 0, 503, 504, 1, 0, 0, 0, 504, 506, 1, 0, 0, 0, 505, 503,
        1, 0, 0, 0, 506, 518, 5, 39, 0, 0, 507, 513, 5, 34, 0, 0, 508, 509, 5, 92, 0, 0, 509, 512,
        5, 34, 0, 0, 510, 512, 8, 31, 0, 0, 511, 508, 1, 0, 0, 0, 511, 510, 1, 0, 0, 0, 512, 515,
        1, 0, 0, 0, 513, 511, 1, 0, 0, 0, 513, 514, 1, 0, 0, 0, 514, 516, 1, 0, 0, 0, 515, 513,
        1, 0, 0, 0, 516, 518, 5, 34, 0, 0, 517, 497, 1, 0, 0, 0, 517, 507, 1, 0, 0, 0, 518, 138,
        1, 0, 0, 0, 519, 521, 5, 91, 0, 0, 520, 522, 8, 32, 0, 0, 521, 520, 1, 0, 0, 0, 522, 523,
        1, 0, 0, 0, 523, 521, 1, 0, 0, 0, 523, 524, 1, 0, 0, 0, 524, 525, 1, 0, 0, 0, 525, 526,
        5, 93, 0, 0, 526, 140, 1, 0, 0, 0, 527, 530, 3, 127, 63, 0, 528, 530, 5, 95, 0, 0, 529,
        527, 1, 0, 0, 0, 529, 528, 1, 0, 0, 0, 530, 536, 1, 0, 0, 0, 531, 535, 3, 127, 63, 0, 532,
        535, 3, 125, 62, 0, 533, 535, 5, 95, 0, 0, 534, 531, 1, 0, 0, 0, 534, 532, 1, 0, 0, 0,
        534, 533, 1, 0, 0, 0, 535, 538, 1, 0, 0, 0, 536, 534, 1, 0, 0, 0, 536, 537, 1, 0, 0, 0,
        537, 142, 1, 0, 0, 0, 538, 536, 1, 0, 0, 0, 539, 540, 7, 27, 0, 0, 540, 541, 7, 27, 0,
        0, 541, 144, 1, 0, 0, 0, 542, 543, 7, 27, 0, 0, 543, 544, 7, 27, 0, 0, 544, 545, 7, 27,
        0, 0, 545, 546, 7, 27, 0, 0, 546, 146, 1, 0, 0, 0, 547, 548, 3, 145, 72, 0, 548, 549,
        5, 45, 0, 0, 549, 550, 3, 143, 71, 0, 550, 551, 5, 45, 0, 0, 551, 552, 3, 143, 71, 0,
        552, 148, 1, 0, 0, 0, 553, 554, 3, 145, 72, 0, 554, 555, 5, 45, 0, 0, 555, 556, 3, 143,
        71, 0, 556, 557, 5, 45, 0, 0, 557, 558, 3, 143, 71, 0, 558, 559, 7, 33, 0, 0, 559, 560,
        3, 143, 71, 0, 560, 561, 5, 58, 0, 0, 561, 562, 3, 143, 71, 0, 562, 563, 5, 58, 0, 0,
        563, 564, 3, 143, 71, 0, 564, 591, 1, 0, 0, 0, 565, 566, 3, 145, 72, 0, 566, 567, 5,
        45, 0, 0, 567, 568, 3, 143, 71, 0, 568, 569, 5, 45, 0, 0, 569, 570, 3, 143, 71, 0, 570,
        571, 7, 33, 0, 0, 571, 572, 3, 143, 71, 0, 572, 573, 5, 58, 0, 0, 573, 574, 3, 143, 71,
        0, 574, 591, 1, 0, 0, 0, 575, 576, 3, 145, 72, 0, 576, 577, 5, 45, 0, 0, 577, 578, 3,
        143, 71, 0, 578, 579, 5, 45, 0, 0, 579, 580, 3, 143, 71, 0, 580, 581, 7, 33, 0, 0, 581,
        582, 3, 143, 71, 0, 582, 591, 1, 0, 0, 0, 583, 584, 3, 145, 72, 0, 584, 585, 5, 45, 0,
        0, 585, 586, 3, 143, 71, 0, 586, 587, 5, 45, 0, 0, 587, 588, 3, 143, 71, 0, 588, 589,
        7, 21, 0, 0, 589, 591, 1, 0, 0, 0, 590, 553, 1, 0, 0, 0, 590, 565, 1, 0, 0, 0, 590, 575,
        1, 0, 0, 0, 590, 583, 1, 0, 0, 0, 591, 150, 1, 0, 0, 0, 592, 593, 9, 0, 0, 0, 593, 152,
        1, 0, 0, 0, 29, 0, 166, 170, 181, 184, 186, 192, 203, 272, 283, 354, 406, 458, 471,
        475, 479, 486, 493, 495, 501, 503, 511, 513, 517, 523, 529, 534, 536, 590, 1, 6,
        0, 0
    ]


class DataLensLexer(Lexer):

    atn = ATNDeserializer().deserialize(serializedATN())

    decisionsToDFA = [DFA(ds, i) for i, ds in enumerate(atn.decisionToState)]

    T__0 = 1
    T__1 = 2
    T__2 = 3
    SINGLE_LINE_COMMENT = 4
    MULTI_LINE_COMMENT = 5
    WS = 6
    AMONG = 7
    AND = 8
    ASC = 9
    BEFORE_FILTER_BY = 10
    BETWEEN = 11
    CASE = 12
    DESC = 13
    ELSE = 14
    ELSEIF = 15
    END = 16
    EXCLUDE = 17
    FALSE = 18
    FIXED = 19
    IF = 20
    IGNORE_DIMENSIONS = 21
    IN = 22
    INCLUDE = 23
    IS = 24
    LIKE = 25
    NOT = 26
    NULL = 27
    OR = 28
    ORDER_BY = 29
    THEN = 30
    TOTAL = 31
    TRUE = 32
    WHEN = 33
    WITHIN = 34
    PLUS = 35
    MINUS = 36
    POWER = 37
    MULDIV = 38
    COMPARISON = 39
    OPENING_PAR = 40
    CLOSING_PAR = 41
    INT = 42
    FLOAT = 43
    ESCAPED_STRING = 44
    FIELD_NAME = 45
    FUNC_NAME = 46
    DATE_INNER = 47
    DATETIME_INNER = 48
    UNEXPECTED_CHARACTER = 49

    channelNames = [u"DEFAULT_TOKEN_CHANNEL", u"HIDDEN"]

    modeNames = ["DEFAULT_MODE"]

    literalNames = ["<INVALID>",
                    "'#'", "'##'", "','", "'+'", "'-'", "'^'", "'('", "')'"]

    symbolicNames = ["<INVALID>",
                     "SINGLE_LINE_COMMENT", "MULTI_LINE_COMMENT", "WS", "AMONG",
                     "AND", "ASC", "BEFORE_FILTER_BY", "BETWEEN", "CASE", "DESC",
                     "ELSE", "ELSEIF", "END", "EXCLUDE", "FALSE", "FIXED", "IF",
                     "IGNORE_DIMENSIONS", "IN", "INCLUDE", "IS", "LIKE", "NOT", "NULL",
                     "OR", "ORDER_BY", "THEN", "TOTAL", "TRUE", "WHEN", "WITHIN",
                     "PLUS", "MINUS", "POWER", "MULDIV", "COMPARISON", "OPENING_PAR",
                     "CLOSING_PAR", "INT", "FLOAT", "ESCAPED_STRING", "FIELD_NAME",
                     "FUNC_NAME", "DATE_INNER", "DATETIME_INNER", "UNEXPECTED_CHARACTER"]

    ruleNames = ["T__0", "T__1", "T__2", "SINGLE_LINE_COMMENT", "MULTI_LINE_COMMENT",
                 "WS", "A", "B", "C", "D", "E", "F", "G", "H", "I", "K",
                 "L", "M", "N", "O", "R", "S", "T", "U", "W", "X", "Y",
                 "AMONG", "AND", "ASC", "BEFORE_FILTER_BY", "BETWEEN",
                 "CASE", "DESC", "ELSE", "ELSEIF", "END", "EXCLUDE", "FALSE",
                 "FIXED", "IF", "IGNORE_DIMENSIONS", "IN", "INCLUDE", "IS",
                 "LIKE", "NOT", "NULL", "OR", "ORDER_BY", "THEN", "TOTAL",
                 "TRUE", "WHEN", "WITHIN", "PLUS", "MINUS", "POWER", "MULDIV",
                 "COMPARISON", "OPENING_PAR", "CLOSING_PAR", "DIGIT", "LETTER",
                 "DECIMAL", "EXP", "INT", "FLOAT", "ESCAPED_STRING", "FIELD_NAME",
                 "FUNC_NAME", "DD", "DDDD", "DATE_INNER", "DATETIME_INNER",
                 "UNEXPECTED_CHARACTER"]

    grammarFileName = "DataLens.g4"

    def __init__(self, input=None, output: TextIO = sys.stdout):
        super().__init__(input, output)
        self.checkVersion("4.11.1")
        self._interp = LexerATNSimulator(self, self.atn, self.decisionsToDFA, PredictionContextCache())
        self._actions = None
        self._predicates = None

lib/dl_formula/dl_formula/parser/antlr/gen/DataLensLexer.py
datalens-tech/datalens-backend
# encoding: utf-8
import sys

from antlr4 import (ATN, DFA, ATNDeserializer, NoViableAltException, Parser, ParserATNSimulator, ParserRuleContext, ParseTreeVisitor, PredictionContextCache, RecognitionException, RuleContext, Token,
                    TokenStream)

if sys.version_info[1] > 5:
    from typing import TextIO
else:
    from typing.io import TextIO


def serializedATN():
    return [
        4, 1, 49, 365, 2, 0, 7, 0, 2, 1, 7, 1, 2, 2, 7, 2, 2, 3, 7, 3, 2, 4, 7, 4, 2, 5, 7, 5, 2, 6, 7,
        6, 2, 7, 7, 7, 2, 8, 7, 8, 2, 9, 7, 9, 2, 10, 7, 10, 2, 11, 7, 11, 2, 12, 7, 12, 2, 13, 7, 13,
        2, 14, 7, 14, 2, 15, 7, 15, 2, 16, 7, 16, 2, 17, 7, 17, 2, 18, 7, 18, 2, 19, 7, 19, 2, 20,
        7, 20, 2, 21, 7, 21, 2, 22, 7, 22, 2, 23, 7, 23, 2, 24, 7, 24, 2, 25, 7, 25, 2, 26, 7, 26,
        2, 27, 7, 27, 2, 28, 7, 28, 1, 0, 1, 0, 1, 1, 1, 1, 1, 2, 1, 2, 1, 3, 1, 3, 1, 3, 1, 3, 1, 4,
        1, 4, 1, 4, 1, 4, 1, 5, 1, 5, 1, 5, 1, 5, 1, 6, 1, 6, 1, 6, 1, 6, 1, 7, 1, 7, 1, 8, 1, 8, 1, 9,
        1, 9, 1, 10, 1, 10, 1, 10, 1, 10, 1, 10, 1, 10, 1, 10, 3, 10, 94, 8, 10, 1, 11, 1, 11, 1,
        11, 1, 11, 5, 11, 100, 8, 11, 10, 11, 12, 11, 103, 9, 11, 1, 12, 1, 12, 1, 12, 1, 12, 5,
        12, 109, 8, 12, 10, 12, 12, 12, 112, 9, 12, 3, 12, 114, 8, 12, 1, 12, 1, 12, 1, 12, 1,
        12, 5, 12, 120, 8, 12, 10, 12, 12, 12, 123, 9, 12, 3, 12, 125, 8, 12, 1, 12, 1, 12, 1,
        12, 1, 12, 5, 12, 131, 8, 12, 10, 12, 12, 12, 134, 9, 12, 3, 12, 136, 8, 12, 3, 12, 138,
        8, 12, 1, 13, 1, 13, 1, 13, 1, 13, 1, 13, 5, 13, 145, 8, 13, 10, 13, 12, 13, 148, 9, 13,
        3, 13, 150, 8, 13, 1, 13, 1, 13, 1, 13, 1, 13, 5, 13, 156, 8, 13, 10, 13, 12, 13, 159,
        9, 13, 3, 13, 161, 8, 13, 3, 13, 163, 8, 13, 1, 14, 1, 14, 1, 14, 1, 14, 5, 14, 169, 8,
        14, 10, 14, 12, 14, 172, 9, 14, 3, 14, 174, 8, 14, 1, 15, 1, 15, 1, 15, 1, 15, 5, 15, 180,
        8, 15, 10, 15, 12, 15, 183, 9, 15, 3, 15, 185, 8, 15, 1, 16, 1, 16, 1, 16, 1, 16, 1, 16,
        5, 16, 192, 8, 16, 10, 16, 12, 16, 195, 9, 16, 3, 16, 197, 8, 16, 1, 16, 3, 16, 200, 8,
        16, 1, 16, 3, 16, 203, 8, 16, 1, 16, 3, 16, 206, 8, 16, 1, 16, 3, 16, 209, 8, 16, 1, 16,
        3, 16, 212, 8, 16, 1, 16, 1, 16, 1, 17, 1, 17, 1, 17, 1, 17, 1, 17, 1, 18, 1, 18, 1, 18,
        1, 19, 1, 19, 1, 19, 1, 19, 1, 19, 1, 20, 1, 20, 5, 20, 231, 8, 20, 10, 20, 12, 20, 234,
        9, 20, 1, 20, 3, 20, 237, 8, 20, 1, 20, 1, 20, 1, 21, 1, 21, 1, 21, 1, 21, 1, 21, 1, 22,
        1, 22, 1, 22, 4, 22, 249, 8, 22, 11, 22, 12, 22, 250, 1, 22, 3, 22, 254, 8, 22, 1, 22,
        1, 22, 1, 23, 1, 23, 1, 23, 1, 23, 1, 24, 1, 24, 1, 24, 1, 24, 1, 24, 1, 24, 1, 24, 1, 24,
        1, 24, 1, 24, 1, 24, 1, 24, 1, 24, 1, 24, 3, 24, 276, 8, 24, 1, 25, 1, 25, 1, 25, 1, 25,
        1, 25, 1, 25, 3, 25, 284, 8, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25,
        1, 25, 1, 25, 1, 25, 3, 25, 297, 8, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25,
        1, 25, 1, 25, 3, 25, 308, 8, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25, 3, 25,
        317, 8, 25, 1, 25, 1, 25, 1, 25, 3, 25, 322, 8, 25, 1, 25, 1, 25, 1, 25, 1, 25, 1, 25, 5,
        25, 329, 8, 25, 10, 25, 12, 25, 332, 9, 25, 1, 25, 3, 25, 335, 8, 25, 1, 25, 5, 25, 338,
        8, 25, 10, 25, 12, 25, 341, 9, 25, 1, 26, 1, 26, 1, 26, 1, 26, 1, 26, 1, 26, 1, 26, 1, 26,
        1, 26, 5, 26, 352, 8, 26, 10, 26, 12, 26, 355, 9, 26, 1, 27, 1, 27, 1, 28, 1, 28, 1, 28,
        1, 28, 3, 28, 363, 8, 28, 1, 28, 0, 2, 50, 52, 29, 0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20,
        22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 0, 4, 2, 0, 18,
        18, 32, 32, 6, 0, 8, 8, 11, 12, 20, 20, 25, 26, 28, 28, 46, 46, 1, 0, 35, 36, 3, 0, 18,
        18, 27, 27, 32, 32, 399, 0, 58, 1, 0, 0, 0, 2, 60, 1, 0, 0, 0, 4, 62, 1, 0, 0, 0, 6, 64, 1,
        0, 0, 0, 8, 68, 1, 0, 0, 0, 10, 72, 1, 0, 0, 0, 12, 76, 1, 0, 0, 0, 14, 80, 1, 0, 0, 0, 16,
        82, 1, 0, 0, 0, 18, 84, 1, 0, 0, 0, 20, 93, 1, 0, 0, 0, 22, 95, 1, 0, 0, 0, 24, 137, 1, 0,
        0, 0, 26, 162, 1, 0, 0, 0, 28, 164, 1, 0, 0, 0, 30, 175, 1, 0, 0, 0, 32, 186, 1, 0, 0, 0,
        34, 215, 1, 0, 0, 0, 36, 220, 1, 0, 0, 0, 38, 223, 1, 0, 0, 0, 40, 228, 1, 0, 0, 0, 42, 240,
        1, 0, 0, 0, 44, 245, 1, 0, 0, 0, 46, 257, 1, 0, 0, 0, 48, 275, 1, 0, 0, 0, 50, 283, 1, 0,
        0, 0, 52, 342, 1, 0, 0, 0, 54, 356, 1, 0, 0, 0, 56, 362, 1, 0, 0, 0, 58, 59, 5, 42, 0, 0,
        59, 1, 1, 0, 0, 0, 60, 61, 5, 43, 0, 0, 61, 3, 1, 0, 0, 0, 62, 63, 5, 44, 0, 0, 63, 5, 1, 0,
        0, 0, 64, 65, 5, 1, 0, 0, 65, 66, 5, 47, 0, 0, 66, 67, 5, 1, 0, 0, 67, 7, 1, 0, 0, 0, 68, 69,
        5, 1, 0, 0, 69, 70, 5, 48, 0, 0, 70, 71, 5, 1, 0, 0, 71, 9, 1, 0, 0, 0, 72, 73, 5, 2, 0, 0,
        73, 74, 5, 47, 0, 0, 74, 75, 5, 2, 0, 0, 75, 11, 1, 0, 0, 0, 76, 77, 5, 2, 0, 0, 77, 78, 5,
        48, 0, 0, 78, 79, 5, 2, 0, 0, 79, 13, 1, 0, 0, 0, 80, 81, 7, 0, 0, 0, 81, 15, 1, 0, 0, 0, 82,
        83, 5, 27, 0, 0, 83, 17, 1, 0, 0, 0, 84, 85, 5, 45, 0, 0, 85, 19, 1, 0, 0, 0, 86, 94, 3, 54,
        27, 0, 87, 88, 3, 54, 27, 0, 88, 89, 5, 9, 0, 0, 89, 94, 1, 0, 0, 0, 90, 91, 3, 54, 27, 0,
        91, 92, 5, 13, 0, 0, 92, 94, 1, 0, 0, 0, 93, 86, 1, 0, 0, 0, 93, 87, 1, 0, 0, 0, 93, 90, 1,
        0, 0, 0, 94, 21, 1, 0, 0, 0, 95, 96, 5, 29, 0, 0, 96, 101, 3, 20, 10, 0, 97, 98, 5, 3, 0,
        0, 98, 100, 3, 20, 10, 0, 99, 97, 1, 0, 0, 0, 100, 103, 1, 0, 0, 0, 101, 99, 1, 0, 0, 0,
        101, 102, 1, 0, 0, 0, 102, 23, 1, 0, 0, 0, 103, 101, 1, 0, 0, 0, 104, 113, 5, 19, 0, 0,
        105, 110, 3, 54, 27, 0, 106, 107, 5, 3, 0, 0, 107, 109, 3, 54, 27, 0, 108, 106, 1, 0,
        0, 0, 109, 112, 1, 0, 0, 0, 110, 108, 1, 0, 0, 0, 110, 111, 1, 0, 0, 0, 111, 114, 1, 0,
        0, 0, 112, 110, 1, 0, 0, 0, 113, 105, 1, 0, 0, 0, 113, 114, 1, 0, 0, 0, 114, 138, 1, 0,
        0, 0, 115, 124, 5, 23, 0, 0, 116, 121, 3, 54, 27, 0, 117, 118, 5, 3, 0, 0, 118, 120, 3,
        54, 27, 0, 119, 117, 1, 0, 0, 0, 120, 123, 1, 0, 0, 0, 121, 119, 1, 0, 0, 0, 121, 122,
        1, 0, 0, 0, 122, 125, 1, 0, 0, 0, 123, 121, 1, 0, 0, 0, 124, 116, 1, 0, 0, 0, 124, 125,
        1, 0, 0, 0, 125, 138, 1, 0, 0, 0, 126, 135, 5, 17, 0, 0, 127, 132, 3, 54, 27, 0, 128, 129,
        5, 3, 0, 0, 129, 131, 3, 54, 27, 0, 130, 128, 1, 0, 0, 0, 131, 134, 1, 0, 0, 0, 132, 130,
        1, 0, 0, 0, 132, 133, 1, 0, 0, 0, 133, 136, 1, 0, 0, 0, 134, 132, 1, 0, 0, 0, 135, 127,
        1, 0, 0, 0, 135, 136, 1, 0, 0, 0, 136, 138, 1, 0, 0, 0, 137, 104, 1, 0, 0, 0, 137, 115,
        1, 0, 0, 0, 137, 126, 1, 0, 0, 0, 138, 25, 1, 0, 0, 0, 139, 163, 5, 31, 0, 0, 140, 149,
        5, 7, 0, 0, 141, 146, 3, 54, 27, 0, 142, 143, 5, 3, 0, 0, 143, 145, 3, 54, 27, 0, 144,
        142, 1, 0, 0, 0, 145, 148, 1, 0, 0, 0, 146, 144, 1, 0, 0, 0, 146, 147, 1, 0, 0, 0, 147,
        150, 1, 0, 0, 0, 148, 146, 1, 0, 0, 0, 149, 141, 1, 0, 0, 0, 149, 150, 1, 0, 0, 0, 150,
        163, 1, 0, 0, 0, 151, 160, 5, 34, 0, 0, 152, 157, 3, 54, 27, 0, 153, 154, 5, 3, 0, 0, 154,
        156, 3, 54, 27, 0, 155, 153, 1, 0, 0, 0, 156, 159, 1, 0, 0, 0, 157, 155, 1, 0, 0, 0, 157,
        158, 1, 0, 0, 0, 158, 161, 1, 0, 0, 0, 159, 157, 1, 0, 0, 0, 160, 152, 1, 0, 0, 0, 160,
        161, 1, 0, 0, 0, 161, 163, 1, 0, 0, 0, 162, 139, 1, 0, 0, 0, 162, 140, 1, 0, 0, 0, 162,
        151, 1, 0, 0, 0, 163, 27, 1, 0, 0, 0, 164, 173, 5, 10, 0, 0, 165, 170, 3, 18, 9, 0, 166,
        167, 5, 3, 0, 0, 167, 169, 3, 18, 9, 0, 168, 166, 1, 0, 0, 0, 169, 172, 1, 0, 0, 0, 170,
        168, 1, 0, 0, 0, 170, 171, 1, 0, 0, 0, 171, 174, 1, 0, 0, 0, 172, 170, 1, 0, 0, 0, 173,
        165, 1, 0, 0, 0, 173, 174, 1, 0, 0, 0, 174, 29, 1, 0, 0, 0, 175, 184, 5, 21, 0, 0, 176,
        181, 3, 18, 9, 0, 177, 178, 5, 3, 0, 0, 178, 180, 3, 18, 9, 0, 179, 177, 1, 0, 0, 0, 180,
        183, 1, 0, 0, 0, 181, 179, 1, 0, 0, 0, 181, 182, 1, 0, 0, 0, 182, 185, 1, 0, 0, 0, 183,
        181, 1, 0, 0, 0, 184, 176, 1, 0, 0, 0, 184, 185, 1, 0, 0, 0, 185, 31, 1, 0, 0, 0, 186, 187,
        7, 1, 0, 0, 187, 196, 5, 40, 0, 0, 188, 193, 3, 54, 27, 0, 189, 190, 5, 3, 0, 0, 190, 192,
        3, 54, 27, 0, 191, 189, 1, 0, 0, 0, 192, 195, 1, 0, 0, 0, 193, 191, 1, 0, 0, 0, 193, 194,
        1, 0, 0, 0, 194, 197, 1, 0, 0, 0, 195, 193, 1, 0, 0, 0, 196, 188, 1, 0, 0, 0, 196, 197,
        1, 0, 0, 0, 197, 199, 1, 0, 0, 0, 198, 200, 3, 26, 13, 0, 199, 198, 1, 0, 0, 0, 199, 200,
        1, 0, 0, 0, 200, 202, 1, 0, 0, 0, 201, 203, 3, 22, 11, 0, 202, 201, 1, 0, 0, 0, 202, 203,
        1, 0, 0, 0, 203, 205, 1, 0, 0, 0, 204, 206, 3, 24, 12, 0, 205, 204, 1, 0, 0, 0, 205, 206,
        1, 0, 0, 0, 206, 208, 1, 0, 0, 0, 207, 209, 3, 28, 14, 0, 208, 207, 1, 0, 0, 0, 208, 209,
        1, 0, 0, 0, 209, 211, 1, 0, 0, 0, 210, 212, 3, 30, 15, 0, 211, 210, 1, 0, 0, 0, 211, 212,
        1, 0, 0, 0, 212, 213, 1, 0, 0, 0, 213, 214, 5, 41, 0, 0, 214, 33, 1, 0, 0, 0, 215, 216,
        5, 15, 0, 0, 216, 217, 3, 54, 27, 0, 217, 218, 5, 30, 0, 0, 218, 219, 3, 54, 27, 0, 219,
        35, 1, 0, 0, 0, 220, 221, 5, 14, 0, 0, 221, 222, 3, 54, 27, 0, 222, 37, 1, 0, 0, 0, 223,
        224, 5, 20, 0, 0, 224, 225, 3, 54, 27, 0, 225, 226, 5, 30, 0, 0, 226, 227, 3, 54, 27,
        0, 227, 39, 1, 0, 0, 0, 228, 232, 3, 38, 19, 0, 229, 231, 3, 34, 17, 0, 230, 229, 1, 0,
        0, 0, 231, 234, 1, 0, 0, 0, 232, 230, 1, 0, 0, 0, 232, 233, 1, 0, 0, 0, 233, 236, 1, 0,
        0, 0, 234, 232, 1, 0, 0, 0, 235, 237, 3, 36, 18, 0, 236, 235, 1, 0, 0, 0, 236, 237, 1,
        0, 0, 0, 237, 238, 1, 0, 0, 0, 238, 239, 5, 16, 0, 0, 239, 41, 1, 0, 0, 0, 240, 241, 5,
        33, 0, 0, 241, 242, 3, 54, 27, 0, 242, 243, 5, 30, 0, 0, 243, 244, 3, 54, 27, 0, 244,
        43, 1, 0, 0, 0, 245, 246, 5, 12, 0, 0, 246, 248, 3, 54, 27, 0, 247, 249, 3, 42, 21, 0,
        248, 247, 1, 0, 0, 0, 249, 250, 1, 0, 0, 0, 250, 248, 1, 0, 0, 0, 250, 251, 1, 0, 0, 0,
        251, 253, 1, 0, 0, 0, 252, 254, 3, 36, 18, 0, 253, 252, 1, 0, 0, 0, 253, 254, 1, 0, 0,
        0, 254, 255, 1, 0, 0, 0, 255, 256, 5, 16, 0, 0, 256, 45, 1, 0, 0, 0, 257, 258, 5, 40, 0,
        0, 258, 259, 3, 54, 27, 0, 259, 260, 5, 41, 0, 0, 260, 47, 1, 0, 0, 0, 261, 276, 3, 0,
        0, 0, 262, 276, 3, 2, 1, 0, 263, 276, 3, 14, 7, 0, 264, 276, 3, 16, 8, 0, 265, 276, 3,
        40, 20, 0, 266, 276, 3, 44, 22, 0, 267, 276, 3, 4, 2, 0, 268, 276, 3, 18, 9, 0, 269, 276,
        3, 6, 3, 0, 270, 276, 3, 8, 4, 0, 271, 276, 3, 10, 5, 0, 272, 276, 3, 12, 6, 0, 273, 276,
        3, 32, 16, 0, 274, 276, 3, 46, 23, 0, 275, 261, 1, 0, 0, 0, 275, 262, 1, 0, 0, 0, 275,
        263, 1, 0, 0, 0, 275, 264, 1, 0, 0, 0, 275, 265, 1, 0, 0, 0, 275, 266, 1, 0, 0, 0, 275,
        267, 1, 0, 0, 0, 275, 268, 1, 0, 0, 0, 275, 269, 1, 0, 0, 0, 275, 270, 1, 0, 0, 0, 275,
        271, 1, 0, 0, 0, 275, 272, 1, 0, 0, 0, 275, 273, 1, 0, 0, 0, 275, 274, 1, 0, 0, 0, 276,
        49, 1, 0, 0, 0, 277, 278, 6, 25, -1, 0, 278, 279, 5, 36, 0, 0, 279, 284, 3, 50, 25, 10,
        280, 281, 5, 26, 0, 0, 281, 284, 3, 50, 25, 2, 282, 284, 3, 48, 24, 0, 283, 277, 1, 0,
        0, 0, 283, 280, 1, 0, 0, 0, 283, 282, 1, 0, 0, 0, 284, 339, 1, 0, 0, 0, 285, 286, 10, 11,
        0, 0, 286, 287, 5, 37, 0, 0, 287, 338, 3, 50, 25, 12, 288, 289, 10, 9, 0, 0, 289, 290,
        5, 38, 0, 0, 290, 338, 3, 50, 25, 10, 291, 292, 10, 8, 0, 0, 292, 293, 7, 2, 0, 0, 293,
        338, 3, 50, 25, 9, 294, 296, 10, 6, 0, 0, 295, 297, 5, 26, 0, 0, 296, 295, 1, 0, 0, 0,
        296, 297, 1, 0, 0, 0, 297, 298, 1, 0, 0, 0, 298, 299, 5, 25, 0, 0, 299, 338, 3, 50, 25,
        7, 300, 301, 10, 5, 0, 0, 301, 302, 5, 39, 0, 0, 302, 338, 3, 50, 25, 6, 303, 307, 10,
        4, 0, 0, 304, 308, 5, 11, 0, 0, 305, 306, 5, 26, 0, 0, 306, 308, 5, 11, 0, 0, 307, 304,
        1, 0, 0, 0, 307, 305, 1, 0, 0, 0, 308, 309, 1, 0, 0, 0, 309, 310, 3, 50, 25, 0, 310, 311,
        5, 8, 0, 0, 311, 312, 3, 50, 25, 5, 312, 338, 1, 0, 0, 0, 313, 314, 10, 7, 0, 0, 314, 316,
        5, 24, 0, 0, 315, 317, 5, 26, 0, 0, 316, 315, 1, 0, 0, 0, 316, 317, 1, 0, 0, 0, 317, 318,
        1, 0, 0, 0, 318, 338, 7, 3, 0, 0, 319, 321, 10, 3, 0, 0, 320, 322, 5, 26, 0, 0, 321, 320,
        1, 0, 0, 0, 321, 322, 1, 0, 0, 0, 322, 323, 1, 0, 0, 0, 323, 324, 5, 22, 0, 0, 324, 334,
        5, 40, 0, 0, 325, 326, 3, 54, 27, 0, 326, 327, 5, 3, 0, 0, 327, 329, 1, 0, 0, 0, 328, 325,
        1, 0, 0, 0, 329, 332, 1, 0, 0, 0, 330, 328, 1, 0, 0, 0, 330, 331, 1, 0, 0, 0, 331, 333,
        1, 0, 0, 0, 332, 330, 1, 0, 0, 0, 333, 335, 3, 54, 27, 0, 334, 330, 1, 0, 0, 0, 334, 335,
        1, 0, 0, 0, 335, 336, 1, 0, 0, 0, 336, 338, 5, 41, 0, 0, 337, 285, 1, 0, 0, 0, 337, 288,
        1, 0, 0, 0, 337, 291, 1, 0, 0, 0, 337, 294, 1, 0, 0, 0, 337, 300, 1, 0, 0, 0, 337, 303,
        1, 0, 0, 0, 337, 313, 1, 0, 0, 0, 337, 319, 1, 0, 0, 0, 338, 341, 1, 0, 0, 0, 339, 337,
        1, 0, 0, 0, 339, 340, 1, 0, 0, 0, 340, 51, 1, 0, 0, 0, 341, 339, 1, 0, 0, 0, 342, 343, 6,
        26, -1, 0, 343, 344, 3, 50, 25, 0, 344, 353, 1, 0, 0, 0, 345, 346, 10, 3, 0, 0, 346, 347,
        5, 8, 0, 0, 347, 352, 3, 52, 26, 4, 348, 349, 10, 2, 0, 0, 349, 350, 5, 28, 0, 0, 350,
        352, 3, 52, 26, 3, 351, 345, 1, 0, 0, 0, 351, 348, 1, 0, 0, 0, 352, 355, 1, 0, 0, 0, 353,
        351, 1, 0, 0, 0, 353, 354, 1, 0, 0, 0, 354, 53, 1, 0, 0, 0, 355, 353, 1, 0, 0, 0, 356, 357,
        3, 52, 26, 0, 357, 55, 1, 0, 0, 0, 358, 359, 3, 54, 27, 0, 359, 360, 5, 0, 0, 1, 360, 363,
        1, 0, 0, 0, 361, 363, 5, 0, 0, 1, 362, 358, 1, 0, 0, 0, 362, 361, 1, 0, 0, 0, 363, 57, 1,
        0, 0, 0, 42, 93, 101, 110, 113, 121, 124, 132, 135, 137, 146, 149, 157, 160, 162,
        170, 173, 181, 184, 193, 196, 199, 202, 205, 208, 211, 232, 236, 250, 253, 275,
        283, 296, 307, 316, 321, 330, 334, 337, 339, 351, 353, 362
    ]


class DataLensParser (Parser):

    grammarFileName = "DataLens.g4"

    atn = ATNDeserializer().deserialize(serializedATN())

    decisionsToDFA = [DFA(ds, i) for i, ds in enumerate(atn.decisionToState)]

    sharedContextCache = PredictionContextCache()

    literalNames = ["<INVALID>", "'#'", "'##'", "','", "<INVALID>", "<INVALID>",
                    "<INVALID>", "<INVALID>", "<INVALID>", "<INVALID>",
                    "<INVALID>", "<INVALID>", "<INVALID>", "<INVALID>",
                    "<INVALID>", "<INVALID>", "<INVALID>", "<INVALID>",
                    "<INVALID>", "<INVALID>", "<INVALID>", "<INVALID>",
                    "<INVALID>", "<INVALID>", "<INVALID>", "<INVALID>",
                    "<INVALID>", "<INVALID>", "<INVALID>", "<INVALID>",
                    "<INVALID>", "<INVALID>", "<INVALID>", "<INVALID>",
                    "<INVALID>", "'+'", "'-'", "'^'", "<INVALID>", "<INVALID>",
                    "'('", "')'"]

    symbolicNames = ["<INVALID>", "<INVALID>", "<INVALID>", "<INVALID>",
                     "SINGLE_LINE_COMMENT", "MULTI_LINE_COMMENT", "WS",
                     "AMONG", "AND", "ASC", "BEFORE_FILTER_BY", "BETWEEN",
                     "CASE", "DESC", "ELSE", "ELSEIF", "END", "EXCLUDE",
                     "FALSE", "FIXED", "IF", "IGNORE_DIMENSIONS", "IN",
                     "INCLUDE", "IS", "LIKE", "NOT", "NULL", "OR", "ORDER_BY",
                     "THEN", "TOTAL", "TRUE", "WHEN", "WITHIN", "PLUS",
                     "MINUS", "POWER", "MULDIV", "COMPARISON", "OPENING_PAR",
                     "CLOSING_PAR", "INT", "FLOAT", "ESCAPED_STRING", "FIELD_NAME",
                     "FUNC_NAME", "DATE_INNER", "DATETIME_INNER", "UNEXPECTED_CHARACTER"]

    RULE_integerLiteral = 0
    RULE_floatLiteral = 1
    RULE_stringLiteral = 2
    RULE_dateLiteral = 3
    RULE_datetimeLiteral = 4
    RULE_genericDateLiteral = 5
    RULE_genericDatetimeLiteral = 6
    RULE_boolLiteral = 7
    RULE_nullLiteral = 8
    RULE_fieldName = 9
    RULE_orderingItem = 10
    RULE_ordering = 11
    RULE_lodSpecifier = 12
    RULE_winGrouping = 13
    RULE_beforeFilterBy = 14
    RULE_ignoreDimensions = 15
    RULE_function = 16
    RULE_elseifPart = 17
    RULE_elsePart = 18
    RULE_ifPart = 19
    RULE_ifBlock = 20
    RULE_whenPart = 21
    RULE_caseBlock = 22
    RULE_parenthesizedExpr = 23
    RULE_exprBasic = 24
    RULE_exprMain = 25
    RULE_exprSecondary = 26
    RULE_expression = 27
    RULE_parse = 28

    ruleNames = ["integerLiteral", "floatLiteral", "stringLiteral", "dateLiteral",
                 "datetimeLiteral", "genericDateLiteral", "genericDatetimeLiteral",
                 "boolLiteral", "nullLiteral", "fieldName", "orderingItem",
                 "ordering", "lodSpecifier", "winGrouping", "beforeFilterBy",
                 "ignoreDimensions", "function", "elseifPart", "elsePart",
                 "ifPart", "ifBlock", "whenPart", "caseBlock", "parenthesizedExpr",
                 "exprBasic", "exprMain", "exprSecondary", "expression",
                 "parse"]

    EOF = Token.EOF
    T__0 = 1
    T__1 = 2
    T__2 = 3
    SINGLE_LINE_COMMENT = 4
    MULTI_LINE_COMMENT = 5
    WS = 6
    AMONG = 7
    AND = 8
    ASC = 9
    BEFORE_FILTER_BY = 10
    BETWEEN = 11
    CASE = 12
    DESC = 13
    ELSE = 14
    ELSEIF = 15
    END = 16
    EXCLUDE = 17
    FALSE = 18
    FIXED = 19
    IF = 20
    IGNORE_DIMENSIONS = 21
    IN = 22
    INCLUDE = 23
    IS = 24
    LIKE = 25
    NOT = 26
    NULL = 27
    OR = 28
    ORDER_BY = 29
    THEN = 30
    TOTAL = 31
    TRUE = 32
    WHEN = 33
    WITHIN = 34
    PLUS = 35
    MINUS = 36
    POWER = 37
    MULDIV = 38
    COMPARISON = 39
    OPENING_PAR = 40
    CLOSING_PAR = 41
    INT = 42
    FLOAT = 43
    ESCAPED_STRING = 44
    FIELD_NAME = 45
    FUNC_NAME = 46
    DATE_INNER = 47
    DATETIME_INNER = 48
    UNEXPECTED_CHARACTER = 49

    def __init__(self, input: TokenStream, output: TextIO = sys.stdout):
        super().__init__(input, output)
        self.checkVersion("4.11.1")
        self._interp = ParserATNSimulator(self, self.atn, self.decisionsToDFA, self.sharedContextCache)
        self._predicates = None

    class IntegerLiteralContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def INT(self):
            return self.getToken(DataLensParser.INT, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_integerLiteral

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitIntegerLiteral"):
                return visitor.visitIntegerLiteral(self)
            else:
                return visitor.visitChildren(self)

    def integerLiteral(self):

        localctx = DataLensParser.IntegerLiteralContext(self, self._ctx, self.state)
        self.enterRule(localctx, 0, self.RULE_integerLiteral)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 58
            self.match(DataLensParser.INT)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class FloatLiteralContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def FLOAT(self):
            return self.getToken(DataLensParser.FLOAT, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_floatLiteral

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitFloatLiteral"):
                return visitor.visitFloatLiteral(self)
            else:
                return visitor.visitChildren(self)

    def floatLiteral(self):

        localctx = DataLensParser.FloatLiteralContext(self, self._ctx, self.state)
        self.enterRule(localctx, 2, self.RULE_floatLiteral)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 60
            self.match(DataLensParser.FLOAT)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class StringLiteralContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def ESCAPED_STRING(self):
            return self.getToken(DataLensParser.ESCAPED_STRING, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_stringLiteral

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitStringLiteral"):
                return visitor.visitStringLiteral(self)
            else:
                return visitor.visitChildren(self)

    def stringLiteral(self):

        localctx = DataLensParser.StringLiteralContext(self, self._ctx, self.state)
        self.enterRule(localctx, 4, self.RULE_stringLiteral)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 62
            self.match(DataLensParser.ESCAPED_STRING)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class DateLiteralContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def DATE_INNER(self):
            return self.getToken(DataLensParser.DATE_INNER, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_dateLiteral

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitDateLiteral"):
                return visitor.visitDateLiteral(self)
            else:
                return visitor.visitChildren(self)

    def dateLiteral(self):

        localctx = DataLensParser.DateLiteralContext(self, self._ctx, self.state)
        self.enterRule(localctx, 6, self.RULE_dateLiteral)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 64
            self.match(DataLensParser.T__0)
            self.state = 65
            self.match(DataLensParser.DATE_INNER)
            self.state = 66
            self.match(DataLensParser.T__0)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class DatetimeLiteralContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def DATETIME_INNER(self):
            return self.getToken(DataLensParser.DATETIME_INNER, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_datetimeLiteral

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitDatetimeLiteral"):
                return visitor.visitDatetimeLiteral(self)
            else:
                return visitor.visitChildren(self)

    def datetimeLiteral(self):

        localctx = DataLensParser.DatetimeLiteralContext(self, self._ctx, self.state)
        self.enterRule(localctx, 8, self.RULE_datetimeLiteral)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 68
            self.match(DataLensParser.T__0)
            self.state = 69
            self.match(DataLensParser.DATETIME_INNER)
            self.state = 70
            self.match(DataLensParser.T__0)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class GenericDateLiteralContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def DATE_INNER(self):
            return self.getToken(DataLensParser.DATE_INNER, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_genericDateLiteral

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitGenericDateLiteral"):
                return visitor.visitGenericDateLiteral(self)
            else:
                return visitor.visitChildren(self)

    def genericDateLiteral(self):

        localctx = DataLensParser.GenericDateLiteralContext(self, self._ctx, self.state)
        self.enterRule(localctx, 10, self.RULE_genericDateLiteral)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 72
            self.match(DataLensParser.T__1)
            self.state = 73
            self.match(DataLensParser.DATE_INNER)
            self.state = 74
            self.match(DataLensParser.T__1)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class GenericDatetimeLiteralContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def DATETIME_INNER(self):
            return self.getToken(DataLensParser.DATETIME_INNER, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_genericDatetimeLiteral

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitGenericDatetimeLiteral"):
                return visitor.visitGenericDatetimeLiteral(self)
            else:
                return visitor.visitChildren(self)

    def genericDatetimeLiteral(self):

        localctx = DataLensParser.GenericDatetimeLiteralContext(self, self._ctx, self.state)
        self.enterRule(localctx, 12, self.RULE_genericDatetimeLiteral)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 76
            self.match(DataLensParser.T__1)
            self.state = 77
            self.match(DataLensParser.DATETIME_INNER)
            self.state = 78
            self.match(DataLensParser.T__1)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class BoolLiteralContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def TRUE(self):
            return self.getToken(DataLensParser.TRUE, 0)

        def FALSE(self):
            return self.getToken(DataLensParser.FALSE, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_boolLiteral

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitBoolLiteral"):
                return visitor.visitBoolLiteral(self)
            else:
                return visitor.visitChildren(self)

    def boolLiteral(self):

        localctx = DataLensParser.BoolLiteralContext(self, self._ctx, self.state)
        self.enterRule(localctx, 14, self.RULE_boolLiteral)
        self._la = 0  # Token type
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 80
            _la = self._input.LA(1)
            if not (_la == 18 or _la == 32):
                self._errHandler.recoverInline(self)
            else:
                self._errHandler.reportMatch(self)
                self.consume()
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class NullLiteralContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def NULL(self):
            return self.getToken(DataLensParser.NULL, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_nullLiteral

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitNullLiteral"):
                return visitor.visitNullLiteral(self)
            else:
                return visitor.visitChildren(self)

    def nullLiteral(self):

        localctx = DataLensParser.NullLiteralContext(self, self._ctx, self.state)
        self.enterRule(localctx, 16, self.RULE_nullLiteral)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 82
            self.match(DataLensParser.NULL)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class FieldNameContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def FIELD_NAME(self):
            return self.getToken(DataLensParser.FIELD_NAME, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_fieldName

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitFieldName"):
                return visitor.visitFieldName(self)
            else:
                return visitor.visitChildren(self)

    def fieldName(self):

        localctx = DataLensParser.FieldNameContext(self, self._ctx, self.state)
        self.enterRule(localctx, 18, self.RULE_fieldName)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 84
            self.match(DataLensParser.FIELD_NAME)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class OrderingItemContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def expression(self):
            return self.getTypedRuleContext(DataLensParser.ExpressionContext, 0)

        def ASC(self):
            return self.getToken(DataLensParser.ASC, 0)

        def DESC(self):
            return self.getToken(DataLensParser.DESC, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_orderingItem

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitOrderingItem"):
                return visitor.visitOrderingItem(self)
            else:
                return visitor.visitChildren(self)

    def orderingItem(self):

        localctx = DataLensParser.OrderingItemContext(self, self._ctx, self.state)
        self.enterRule(localctx, 20, self.RULE_orderingItem)
        try:
            self.state = 93
            self._errHandler.sync(self)
            la_ = self._interp.adaptivePredict(self._input, 0, self._ctx)
            if la_ == 1:
                self.enterOuterAlt(localctx, 1)
                self.state = 86
                self.expression()

            elif la_ == 2:
                self.enterOuterAlt(localctx, 2)
                self.state = 87
                self.expression()
                self.state = 88
                self.match(DataLensParser.ASC)

            elif la_ == 3:
                self.enterOuterAlt(localctx, 3)
                self.state = 90
                self.expression()
                self.state = 91
                self.match(DataLensParser.DESC)

        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class OrderingContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def ORDER_BY(self):
            return self.getToken(DataLensParser.ORDER_BY, 0)

        def orderingItem(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.OrderingItemContext)
            else:
                return self.getTypedRuleContext(DataLensParser.OrderingItemContext, i)

        def getRuleIndex(self):
            return DataLensParser.RULE_ordering

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitOrdering"):
                return visitor.visitOrdering(self)
            else:
                return visitor.visitChildren(self)

    def ordering(self):

        localctx = DataLensParser.OrderingContext(self, self._ctx, self.state)
        self.enterRule(localctx, 22, self.RULE_ordering)
        self._la = 0  # Token type
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 95
            self.match(DataLensParser.ORDER_BY)
            self.state = 96
            self.orderingItem()
            self.state = 101
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            while _la == 3:
                self.state = 97
                self.match(DataLensParser.T__2)
                self.state = 98
                self.orderingItem()
                self.state = 103
                self._errHandler.sync(self)
                _la = self._input.LA(1)

        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class LodSpecifierContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def FIXED(self):
            return self.getToken(DataLensParser.FIXED, 0)

        def expression(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExpressionContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExpressionContext, i)

        def INCLUDE(self):
            return self.getToken(DataLensParser.INCLUDE, 0)

        def EXCLUDE(self):
            return self.getToken(DataLensParser.EXCLUDE, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_lodSpecifier

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitLodSpecifier"):
                return visitor.visitLodSpecifier(self)
            else:
                return visitor.visitChildren(self)

    def lodSpecifier(self):

        localctx = DataLensParser.LodSpecifierContext(self, self._ctx, self.state)
        self.enterRule(localctx, 24, self.RULE_lodSpecifier)
        self._la = 0  # Token type
        try:
            self.state = 137
            self._errHandler.sync(self)
            token = self._input.LA(1)
            if token in [19]:
                self.enterOuterAlt(localctx, 1)
                self.state = 104
                self.match(DataLensParser.FIXED)
                self.state = 113
                self._errHandler.sync(self)
                _la = self._input.LA(1)
                if ((_la) & ~0x3f) == 0 and ((1 << _la) & 137512472549638) != 0:
                    self.state = 105
                    self.expression()
                    self.state = 110
                    self._errHandler.sync(self)
                    _la = self._input.LA(1)
                    while _la == 3:
                        self.state = 106
                        self.match(DataLensParser.T__2)
                        self.state = 107
                        self.expression()
                        self.state = 112
                        self._errHandler.sync(self)
                        _la = self._input.LA(1)

            elif token in [23]:
                self.enterOuterAlt(localctx, 2)
                self.state = 115
                self.match(DataLensParser.INCLUDE)
                self.state = 124
                self._errHandler.sync(self)
                _la = self._input.LA(1)
                if ((_la) & ~0x3f) == 0 and ((1 << _la) & 137512472549638) != 0:
                    self.state = 116
                    self.expression()
                    self.state = 121
                    self._errHandler.sync(self)
                    _la = self._input.LA(1)
                    while _la == 3:
                        self.state = 117
                        self.match(DataLensParser.T__2)
                        self.state = 118
                        self.expression()
                        self.state = 123
                        self._errHandler.sync(self)
                        _la = self._input.LA(1)

            elif token in [17]:
                self.enterOuterAlt(localctx, 3)
                self.state = 126
                self.match(DataLensParser.EXCLUDE)
                self.state = 135
                self._errHandler.sync(self)
                _la = self._input.LA(1)
                if ((_la) & ~0x3f) == 0 and ((1 << _la) & 137512472549638) != 0:
                    self.state = 127
                    self.expression()
                    self.state = 132
                    self._errHandler.sync(self)
                    _la = self._input.LA(1)
                    while _la == 3:
                        self.state = 128
                        self.match(DataLensParser.T__2)
                        self.state = 129
                        self.expression()
                        self.state = 134
                        self._errHandler.sync(self)
                        _la = self._input.LA(1)

            else:
                raise NoViableAltException(self)

        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class WinGroupingContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def TOTAL(self):
            return self.getToken(DataLensParser.TOTAL, 0)

        def AMONG(self):
            return self.getToken(DataLensParser.AMONG, 0)

        def expression(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExpressionContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExpressionContext, i)

        def WITHIN(self):
            return self.getToken(DataLensParser.WITHIN, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_winGrouping

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitWinGrouping"):
                return visitor.visitWinGrouping(self)
            else:
                return visitor.visitChildren(self)

    def winGrouping(self):

        localctx = DataLensParser.WinGroupingContext(self, self._ctx, self.state)
        self.enterRule(localctx, 26, self.RULE_winGrouping)
        self._la = 0  # Token type
        try:
            self.state = 162
            self._errHandler.sync(self)
            token = self._input.LA(1)
            if token in [31]:
                self.enterOuterAlt(localctx, 1)
                self.state = 139
                self.match(DataLensParser.TOTAL)
            elif token in [7]:
                self.enterOuterAlt(localctx, 2)
                self.state = 140
                self.match(DataLensParser.AMONG)
                self.state = 149
                self._errHandler.sync(self)
                _la = self._input.LA(1)
                if ((_la) & ~0x3f) == 0 and ((1 << _la) & 137512472549638) != 0:
                    self.state = 141
                    self.expression()
                    self.state = 146
                    self._errHandler.sync(self)
                    _la = self._input.LA(1)
                    while _la == 3:
                        self.state = 142
                        self.match(DataLensParser.T__2)
                        self.state = 143
                        self.expression()
                        self.state = 148
                        self._errHandler.sync(self)
                        _la = self._input.LA(1)

            elif token in [34]:
                self.enterOuterAlt(localctx, 3)
                self.state = 151
                self.match(DataLensParser.WITHIN)
                self.state = 160
                self._errHandler.sync(self)
                _la = self._input.LA(1)
                if ((_la) & ~0x3f) == 0 and ((1 << _la) & 137512472549638) != 0:
                    self.state = 152
                    self.expression()
                    self.state = 157
                    self._errHandler.sync(self)
                    _la = self._input.LA(1)
                    while _la == 3:
                        self.state = 153
                        self.match(DataLensParser.T__2)
                        self.state = 154
                        self.expression()
                        self.state = 159
                        self._errHandler.sync(self)
                        _la = self._input.LA(1)

            else:
                raise NoViableAltException(self)

        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class BeforeFilterByContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def BEFORE_FILTER_BY(self):
            return self.getToken(DataLensParser.BEFORE_FILTER_BY, 0)

        def fieldName(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.FieldNameContext)
            else:
                return self.getTypedRuleContext(DataLensParser.FieldNameContext, i)

        def getRuleIndex(self):
            return DataLensParser.RULE_beforeFilterBy

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitBeforeFilterBy"):
                return visitor.visitBeforeFilterBy(self)
            else:
                return visitor.visitChildren(self)

    def beforeFilterBy(self):

        localctx = DataLensParser.BeforeFilterByContext(self, self._ctx, self.state)
        self.enterRule(localctx, 28, self.RULE_beforeFilterBy)
        self._la = 0  # Token type
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 164
            self.match(DataLensParser.BEFORE_FILTER_BY)
            self.state = 173
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            if _la == 45:
                self.state = 165
                self.fieldName()
                self.state = 170
                self._errHandler.sync(self)
                _la = self._input.LA(1)
                while _la == 3:
                    self.state = 166
                    self.match(DataLensParser.T__2)
                    self.state = 167
                    self.fieldName()
                    self.state = 172
                    self._errHandler.sync(self)
                    _la = self._input.LA(1)

        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class IgnoreDimensionsContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def IGNORE_DIMENSIONS(self):
            return self.getToken(DataLensParser.IGNORE_DIMENSIONS, 0)

        def fieldName(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.FieldNameContext)
            else:
                return self.getTypedRuleContext(DataLensParser.FieldNameContext, i)

        def getRuleIndex(self):
            return DataLensParser.RULE_ignoreDimensions

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitIgnoreDimensions"):
                return visitor.visitIgnoreDimensions(self)
            else:
                return visitor.visitChildren(self)

    def ignoreDimensions(self):

        localctx = DataLensParser.IgnoreDimensionsContext(self, self._ctx, self.state)
        self.enterRule(localctx, 30, self.RULE_ignoreDimensions)
        self._la = 0  # Token type
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 175
            self.match(DataLensParser.IGNORE_DIMENSIONS)
            self.state = 184
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            if _la == 45:
                self.state = 176
                self.fieldName()
                self.state = 181
                self._errHandler.sync(self)
                _la = self._input.LA(1)
                while _la == 3:
                    self.state = 177
                    self.match(DataLensParser.T__2)
                    self.state = 178
                    self.fieldName()
                    self.state = 183
                    self._errHandler.sync(self)
                    _la = self._input.LA(1)

        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class FunctionContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def OPENING_PAR(self):
            return self.getToken(DataLensParser.OPENING_PAR, 0)

        def CLOSING_PAR(self):
            return self.getToken(DataLensParser.CLOSING_PAR, 0)

        def FUNC_NAME(self):
            return self.getToken(DataLensParser.FUNC_NAME, 0)

        def CASE(self):
            return self.getToken(DataLensParser.CASE, 0)

        def IF(self):
            return self.getToken(DataLensParser.IF, 0)

        def NOT(self):
            return self.getToken(DataLensParser.NOT, 0)

        def AND(self):
            return self.getToken(DataLensParser.AND, 0)

        def OR(self):
            return self.getToken(DataLensParser.OR, 0)

        def LIKE(self):
            return self.getToken(DataLensParser.LIKE, 0)

        def BETWEEN(self):
            return self.getToken(DataLensParser.BETWEEN, 0)

        def expression(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExpressionContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExpressionContext, i)

        def winGrouping(self):
            return self.getTypedRuleContext(DataLensParser.WinGroupingContext, 0)

        def ordering(self):
            return self.getTypedRuleContext(DataLensParser.OrderingContext, 0)

        def lodSpecifier(self):
            return self.getTypedRuleContext(DataLensParser.LodSpecifierContext, 0)

        def beforeFilterBy(self):
            return self.getTypedRuleContext(DataLensParser.BeforeFilterByContext, 0)

        def ignoreDimensions(self):
            return self.getTypedRuleContext(DataLensParser.IgnoreDimensionsContext, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_function

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitFunction"):
                return visitor.visitFunction(self)
            else:
                return visitor.visitChildren(self)

    def function(self):

        localctx = DataLensParser.FunctionContext(self, self._ctx, self.state)
        self.enterRule(localctx, 32, self.RULE_function)
        self._la = 0  # Token type
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 186
            _la = self._input.LA(1)
            if not (((_la) & ~0x3f) == 0 and ((1 << _la) & 70369114331392) != 0):
                self._errHandler.recoverInline(self)
            else:
                self._errHandler.reportMatch(self)
                self.consume()
            self.state = 187
            self.match(DataLensParser.OPENING_PAR)
            self.state = 196
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            if ((_la) & ~0x3f) == 0 and ((1 << _la) & 137512472549638) != 0:
                self.state = 188
                self.expression()
                self.state = 193
                self._errHandler.sync(self)
                _la = self._input.LA(1)
                while _la == 3:
                    self.state = 189
                    self.match(DataLensParser.T__2)
                    self.state = 190
                    self.expression()
                    self.state = 195
                    self._errHandler.sync(self)
                    _la = self._input.LA(1)

            self.state = 199
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            if ((_la) & ~0x3f) == 0 and ((1 << _la) & 19327352960) != 0:
                self.state = 198
                self.winGrouping()

            self.state = 202
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            if _la == 29:
                self.state = 201
                self.ordering()

            self.state = 205
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            if ((_la) & ~0x3f) == 0 and ((1 << _la) & 9043968) != 0:
                self.state = 204
                self.lodSpecifier()

            self.state = 208
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            if _la == 10:
                self.state = 207
                self.beforeFilterBy()

            self.state = 211
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            if _la == 21:
                self.state = 210
                self.ignoreDimensions()

            self.state = 213
            self.match(DataLensParser.CLOSING_PAR)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class ElseifPartContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def ELSEIF(self):
            return self.getToken(DataLensParser.ELSEIF, 0)

        def expression(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExpressionContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExpressionContext, i)

        def THEN(self):
            return self.getToken(DataLensParser.THEN, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_elseifPart

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitElseifPart"):
                return visitor.visitElseifPart(self)
            else:
                return visitor.visitChildren(self)

    def elseifPart(self):

        localctx = DataLensParser.ElseifPartContext(self, self._ctx, self.state)
        self.enterRule(localctx, 34, self.RULE_elseifPart)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 215
            self.match(DataLensParser.ELSEIF)
            self.state = 216
            self.expression()
            self.state = 217
            self.match(DataLensParser.THEN)
            self.state = 218
            self.expression()
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class ElsePartContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def ELSE(self):
            return self.getToken(DataLensParser.ELSE, 0)

        def expression(self):
            return self.getTypedRuleContext(DataLensParser.ExpressionContext, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_elsePart

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitElsePart"):
                return visitor.visitElsePart(self)
            else:
                return visitor.visitChildren(self)

    def elsePart(self):

        localctx = DataLensParser.ElsePartContext(self, self._ctx, self.state)
        self.enterRule(localctx, 36, self.RULE_elsePart)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 220
            self.match(DataLensParser.ELSE)
            self.state = 221
            self.expression()
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class IfPartContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def IF(self):
            return self.getToken(DataLensParser.IF, 0)

        def expression(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExpressionContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExpressionContext, i)

        def THEN(self):
            return self.getToken(DataLensParser.THEN, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_ifPart

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitIfPart"):
                return visitor.visitIfPart(self)
            else:
                return visitor.visitChildren(self)

    def ifPart(self):

        localctx = DataLensParser.IfPartContext(self, self._ctx, self.state)
        self.enterRule(localctx, 38, self.RULE_ifPart)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 223
            self.match(DataLensParser.IF)
            self.state = 224
            self.expression()
            self.state = 225
            self.match(DataLensParser.THEN)
            self.state = 226
            self.expression()
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class IfBlockContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def ifPart(self):
            return self.getTypedRuleContext(DataLensParser.IfPartContext, 0)

        def END(self):
            return self.getToken(DataLensParser.END, 0)

        def elseifPart(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ElseifPartContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ElseifPartContext, i)

        def elsePart(self):
            return self.getTypedRuleContext(DataLensParser.ElsePartContext, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_ifBlock

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitIfBlock"):
                return visitor.visitIfBlock(self)
            else:
                return visitor.visitChildren(self)

    def ifBlock(self):

        localctx = DataLensParser.IfBlockContext(self, self._ctx, self.state)
        self.enterRule(localctx, 40, self.RULE_ifBlock)
        self._la = 0  # Token type
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 228
            self.ifPart()
            self.state = 232
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            while _la == 15:
                self.state = 229
                self.elseifPart()
                self.state = 234
                self._errHandler.sync(self)
                _la = self._input.LA(1)

            self.state = 236
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            if _la == 14:
                self.state = 235
                self.elsePart()

            self.state = 238
            self.match(DataLensParser.END)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class WhenPartContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def WHEN(self):
            return self.getToken(DataLensParser.WHEN, 0)

        def expression(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExpressionContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExpressionContext, i)

        def THEN(self):
            return self.getToken(DataLensParser.THEN, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_whenPart

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitWhenPart"):
                return visitor.visitWhenPart(self)
            else:
                return visitor.visitChildren(self)

    def whenPart(self):

        localctx = DataLensParser.WhenPartContext(self, self._ctx, self.state)
        self.enterRule(localctx, 42, self.RULE_whenPart)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 240
            self.match(DataLensParser.WHEN)
            self.state = 241
            self.expression()
            self.state = 242
            self.match(DataLensParser.THEN)
            self.state = 243
            self.expression()
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class CaseBlockContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def CASE(self):
            return self.getToken(DataLensParser.CASE, 0)

        def expression(self):
            return self.getTypedRuleContext(DataLensParser.ExpressionContext, 0)

        def END(self):
            return self.getToken(DataLensParser.END, 0)

        def whenPart(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.WhenPartContext)
            else:
                return self.getTypedRuleContext(DataLensParser.WhenPartContext, i)

        def elsePart(self):
            return self.getTypedRuleContext(DataLensParser.ElsePartContext, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_caseBlock

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitCaseBlock"):
                return visitor.visitCaseBlock(self)
            else:
                return visitor.visitChildren(self)

    def caseBlock(self):

        localctx = DataLensParser.CaseBlockContext(self, self._ctx, self.state)
        self.enterRule(localctx, 44, self.RULE_caseBlock)
        self._la = 0  # Token type
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 245
            self.match(DataLensParser.CASE)
            self.state = 246
            self.expression()
            self.state = 248
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            while True:
                self.state = 247
                self.whenPart()
                self.state = 250
                self._errHandler.sync(self)
                _la = self._input.LA(1)
                if not (_la == 33):
                    break

            self.state = 253
            self._errHandler.sync(self)
            _la = self._input.LA(1)
            if _la == 14:
                self.state = 252
                self.elsePart()

            self.state = 255
            self.match(DataLensParser.END)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class ParenthesizedExprContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def OPENING_PAR(self):
            return self.getToken(DataLensParser.OPENING_PAR, 0)

        def expression(self):
            return self.getTypedRuleContext(DataLensParser.ExpressionContext, 0)

        def CLOSING_PAR(self):
            return self.getToken(DataLensParser.CLOSING_PAR, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_parenthesizedExpr

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitParenthesizedExpr"):
                return visitor.visitParenthesizedExpr(self)
            else:
                return visitor.visitChildren(self)

    def parenthesizedExpr(self):

        localctx = DataLensParser.ParenthesizedExprContext(self, self._ctx, self.state)
        self.enterRule(localctx, 46, self.RULE_parenthesizedExpr)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 257
            self.match(DataLensParser.OPENING_PAR)
            self.state = 258
            self.expression()
            self.state = 259
            self.match(DataLensParser.CLOSING_PAR)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class ExprBasicContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def integerLiteral(self):
            return self.getTypedRuleContext(DataLensParser.IntegerLiteralContext, 0)

        def floatLiteral(self):
            return self.getTypedRuleContext(DataLensParser.FloatLiteralContext, 0)

        def boolLiteral(self):
            return self.getTypedRuleContext(DataLensParser.BoolLiteralContext, 0)

        def nullLiteral(self):
            return self.getTypedRuleContext(DataLensParser.NullLiteralContext, 0)

        def ifBlock(self):
            return self.getTypedRuleContext(DataLensParser.IfBlockContext, 0)

        def caseBlock(self):
            return self.getTypedRuleContext(DataLensParser.CaseBlockContext, 0)

        def stringLiteral(self):
            return self.getTypedRuleContext(DataLensParser.StringLiteralContext, 0)

        def fieldName(self):
            return self.getTypedRuleContext(DataLensParser.FieldNameContext, 0)

        def dateLiteral(self):
            return self.getTypedRuleContext(DataLensParser.DateLiteralContext, 0)

        def datetimeLiteral(self):
            return self.getTypedRuleContext(DataLensParser.DatetimeLiteralContext, 0)

        def genericDateLiteral(self):
            return self.getTypedRuleContext(DataLensParser.GenericDateLiteralContext, 0)

        def genericDatetimeLiteral(self):
            return self.getTypedRuleContext(DataLensParser.GenericDatetimeLiteralContext, 0)

        def function(self):
            return self.getTypedRuleContext(DataLensParser.FunctionContext, 0)

        def parenthesizedExpr(self):
            return self.getTypedRuleContext(DataLensParser.ParenthesizedExprContext, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_exprBasic

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitExprBasic"):
                return visitor.visitExprBasic(self)
            else:
                return visitor.visitChildren(self)

    def exprBasic(self):

        localctx = DataLensParser.ExprBasicContext(self, self._ctx, self.state)
        self.enterRule(localctx, 48, self.RULE_exprBasic)
        try:
            self.state = 275
            self._errHandler.sync(self)
            la_ = self._interp.adaptivePredict(self._input, 29, self._ctx)
            if la_ == 1:
                self.enterOuterAlt(localctx, 1)
                self.state = 261
                self.integerLiteral()

            elif la_ == 2:
                self.enterOuterAlt(localctx, 2)
                self.state = 262
                self.floatLiteral()

            elif la_ == 3:
                self.enterOuterAlt(localctx, 3)
                self.state = 263
                self.boolLiteral()

            elif la_ == 4:
                self.enterOuterAlt(localctx, 4)
                self.state = 264
                self.nullLiteral()

            elif la_ == 5:
                self.enterOuterAlt(localctx, 5)
                self.state = 265
                self.ifBlock()

            elif la_ == 6:
                self.enterOuterAlt(localctx, 6)
                self.state = 266
                self.caseBlock()

            elif la_ == 7:
                self.enterOuterAlt(localctx, 7)
                self.state = 267
                self.stringLiteral()

            elif la_ == 8:
                self.enterOuterAlt(localctx, 8)
                self.state = 268
                self.fieldName()

            elif la_ == 9:
                self.enterOuterAlt(localctx, 9)
                self.state = 269
                self.dateLiteral()

            elif la_ == 10:
                self.enterOuterAlt(localctx, 10)
                self.state = 270
                self.datetimeLiteral()

            elif la_ == 11:
                self.enterOuterAlt(localctx, 11)
                self.state = 271
                self.genericDateLiteral()

            elif la_ == 12:
                self.enterOuterAlt(localctx, 12)
                self.state = 272
                self.genericDatetimeLiteral()

            elif la_ == 13:
                self.enterOuterAlt(localctx, 13)
                self.state = 273
                self.function()

            elif la_ == 14:
                self.enterOuterAlt(localctx, 14)
                self.state = 274
                self.parenthesizedExpr()

        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class ExprMainContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def getRuleIndex(self):
            return DataLensParser.RULE_exprMain

        def copyFrom(self, ctx: ParserRuleContext):
            super().copyFrom(ctx)

    class ExprBasicAltContext(ExprMainContext):

        def __init__(self, parser, ctx: ParserRuleContext):  # actually a DataLensParser.ExprMainContext
            super().__init__(parser)
            self.copyFrom(ctx)

        def exprBasic(self):
            return self.getTypedRuleContext(DataLensParser.ExprBasicContext, 0)

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitExprBasicAlt"):
                return visitor.visitExprBasicAlt(self)
            else:
                return visitor.visitChildren(self)

    class UnaryPrefixContext(ExprMainContext):

        def __init__(self, parser, ctx: ParserRuleContext):  # actually a DataLensParser.ExprMainContext
            super().__init__(parser)
            self.copyFrom(ctx)

        def MINUS(self):
            return self.getToken(DataLensParser.MINUS, 0)

        def exprMain(self):
            return self.getTypedRuleContext(DataLensParser.ExprMainContext, 0)

        def NOT(self):
            return self.getToken(DataLensParser.NOT, 0)

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitUnaryPrefix"):
                return visitor.visitUnaryPrefix(self)
            else:
                return visitor.visitChildren(self)

    class InExprContext(ExprMainContext):

        def __init__(self, parser, ctx: ParserRuleContext):  # actually a DataLensParser.ExprMainContext
            super().__init__(parser)
            self.copyFrom(ctx)

        def exprMain(self):
            return self.getTypedRuleContext(DataLensParser.ExprMainContext, 0)

        def IN(self):
            return self.getToken(DataLensParser.IN, 0)

        def OPENING_PAR(self):
            return self.getToken(DataLensParser.OPENING_PAR, 0)

        def CLOSING_PAR(self):
            return self.getToken(DataLensParser.CLOSING_PAR, 0)

        def NOT(self):
            return self.getToken(DataLensParser.NOT, 0)

        def expression(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExpressionContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExpressionContext, i)

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitInExpr"):
                return visitor.visitInExpr(self)
            else:
                return visitor.visitChildren(self)

    class BinaryExprContext(ExprMainContext):

        def __init__(self, parser, ctx: ParserRuleContext):  # actually a DataLensParser.ExprMainContext
            super().__init__(parser)
            self.copyFrom(ctx)

        def exprMain(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExprMainContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExprMainContext, i)

        def POWER(self):
            return self.getToken(DataLensParser.POWER, 0)

        def MULDIV(self):
            return self.getToken(DataLensParser.MULDIV, 0)

        def PLUS(self):
            return self.getToken(DataLensParser.PLUS, 0)

        def MINUS(self):
            return self.getToken(DataLensParser.MINUS, 0)

        def LIKE(self):
            return self.getToken(DataLensParser.LIKE, 0)

        def NOT(self):
            return self.getToken(DataLensParser.NOT, 0)

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitBinaryExpr"):
                return visitor.visitBinaryExpr(self)
            else:
                return visitor.visitChildren(self)

    class ComparisonChainContext(ExprMainContext):

        def __init__(self, parser, ctx: ParserRuleContext):  # actually a DataLensParser.ExprMainContext
            super().__init__(parser)
            self.copyFrom(ctx)

        def exprMain(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExprMainContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExprMainContext, i)

        def COMPARISON(self):
            return self.getToken(DataLensParser.COMPARISON, 0)

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitComparisonChain"):
                return visitor.visitComparisonChain(self)
            else:
                return visitor.visitChildren(self)

    class UnaryPostfixContext(ExprMainContext):

        def __init__(self, parser, ctx: ParserRuleContext):  # actually a DataLensParser.ExprMainContext
            super().__init__(parser)
            self.copyFrom(ctx)

        def exprMain(self):
            return self.getTypedRuleContext(DataLensParser.ExprMainContext, 0)

        def IS(self):
            return self.getToken(DataLensParser.IS, 0)

        def TRUE(self):
            return self.getToken(DataLensParser.TRUE, 0)

        def FALSE(self):
            return self.getToken(DataLensParser.FALSE, 0)

        def NULL(self):
            return self.getToken(DataLensParser.NULL, 0)

        def NOT(self):
            return self.getToken(DataLensParser.NOT, 0)

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitUnaryPostfix"):
                return visitor.visitUnaryPostfix(self)
            else:
                return visitor.visitChildren(self)

    class BetweenExprContext(ExprMainContext):

        def __init__(self, parser, ctx: ParserRuleContext):  # actually a DataLensParser.ExprMainContext
            super().__init__(parser)
            self.copyFrom(ctx)

        def exprMain(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExprMainContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExprMainContext, i)

        def AND(self):
            return self.getToken(DataLensParser.AND, 0)

        def BETWEEN(self):
            return self.getToken(DataLensParser.BETWEEN, 0)

        def NOT(self):
            return self.getToken(DataLensParser.NOT, 0)

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitBetweenExpr"):
                return visitor.visitBetweenExpr(self)
            else:
                return visitor.visitChildren(self)

    def exprMain(self, _p: int = 0):
        _parentctx = self._ctx
        _parentState = self.state
        localctx = DataLensParser.ExprMainContext(self, self._ctx, _parentState)

        _startState = 50
        self.enterRecursionRule(localctx, 50, self.RULE_exprMain, _p)
        self._la = 0  # Token type
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 283
            self._errHandler.sync(self)
            la_ = self._interp.adaptivePredict(self._input, 30, self._ctx)
            if la_ == 1:
                localctx = DataLensParser.UnaryPrefixContext(self, localctx)
                self._ctx = localctx

                self.state = 278
                self.match(DataLensParser.MINUS)
                self.state = 279
                self.exprMain(10)

            elif la_ == 2:
                localctx = DataLensParser.UnaryPrefixContext(self, localctx)
                self._ctx = localctx

                self.state = 280
                self.match(DataLensParser.NOT)
                self.state = 281
                self.exprMain(2)

            elif la_ == 3:
                localctx = DataLensParser.ExprBasicAltContext(self, localctx)
                self._ctx = localctx

                self.state = 282
                self.exprBasic()

            self._ctx.stop = self._input.LT(-1)
            self.state = 339
            self._errHandler.sync(self)
            _alt = self._interp.adaptivePredict(self._input, 38, self._ctx)
            while _alt != 2 and _alt != ATN.INVALID_ALT_NUMBER:
                if _alt == 1:
                    if self._parseListeners is not None:
                        self.triggerExitRuleEvent()

                    self.state = 337
                    self._errHandler.sync(self)
                    la_ = self._interp.adaptivePredict(self._input, 37, self._ctx)
                    if la_ == 1:
                        localctx = DataLensParser.BinaryExprContext(self, DataLensParser.ExprMainContext(self, _parentctx, _parentState))
                        self.pushNewRecursionContext(localctx, _startState, self.RULE_exprMain)
                        self.state = 285
                        if not self.precpred(self._ctx, 11):
                            from antlr4.error.Errors import FailedPredicateException
                            raise FailedPredicateException(self, "self.precpred(self._ctx, 11)")
                        self.state = 286
                        self.match(DataLensParser.POWER)
                        self.state = 287
                        self.exprMain(12)

                    elif la_ == 2:
                        localctx = DataLensParser.BinaryExprContext(self, DataLensParser.ExprMainContext(self, _parentctx, _parentState))
                        self.pushNewRecursionContext(localctx, _startState, self.RULE_exprMain)
                        self.state = 288
                        if not self.precpred(self._ctx, 9):
                            from antlr4.error.Errors import FailedPredicateException
                            raise FailedPredicateException(self, "self.precpred(self._ctx, 9)")
                        self.state = 289
                        self.match(DataLensParser.MULDIV)
                        self.state = 290
                        self.exprMain(10)

                    elif la_ == 3:
                        localctx = DataLensParser.BinaryExprContext(self, DataLensParser.ExprMainContext(self, _parentctx, _parentState))
                        self.pushNewRecursionContext(localctx, _startState, self.RULE_exprMain)
                        self.state = 291
                        if not self.precpred(self._ctx, 8):
                            from antlr4.error.Errors import FailedPredicateException
                            raise FailedPredicateException(self, "self.precpred(self._ctx, 8)")
                        self.state = 292
                        _la = self._input.LA(1)
                        if not (_la == 35 or _la == 36):
                            self._errHandler.recoverInline(self)
                        else:
                            self._errHandler.reportMatch(self)
                            self.consume()
                        self.state = 293
                        self.exprMain(9)

                    elif la_ == 4:
                        localctx = DataLensParser.BinaryExprContext(self, DataLensParser.ExprMainContext(self, _parentctx, _parentState))
                        self.pushNewRecursionContext(localctx, _startState, self.RULE_exprMain)
                        self.state = 294
                        if not self.precpred(self._ctx, 6):
                            from antlr4.error.Errors import FailedPredicateException
                            raise FailedPredicateException(self, "self.precpred(self._ctx, 6)")
                        self.state = 296
                        self._errHandler.sync(self)
                        _la = self._input.LA(1)
                        if _la == 26:
                            self.state = 295
                            self.match(DataLensParser.NOT)

                        self.state = 298
                        self.match(DataLensParser.LIKE)
                        self.state = 299
                        self.exprMain(7)

                    elif la_ == 5:
                        localctx = DataLensParser.ComparisonChainContext(self, DataLensParser.ExprMainContext(self, _parentctx, _parentState))
                        self.pushNewRecursionContext(localctx, _startState, self.RULE_exprMain)
                        self.state = 300
                        if not self.precpred(self._ctx, 5):
                            from antlr4.error.Errors import FailedPredicateException
                            raise FailedPredicateException(self, "self.precpred(self._ctx, 5)")
                        self.state = 301
                        self.match(DataLensParser.COMPARISON)
                        self.state = 302
                        self.exprMain(6)

                    elif la_ == 6:
                        localctx = DataLensParser.BetweenExprContext(self, DataLensParser.ExprMainContext(self, _parentctx, _parentState))
                        self.pushNewRecursionContext(localctx, _startState, self.RULE_exprMain)
                        self.state = 303
                        if not self.precpred(self._ctx, 4):
                            from antlr4.error.Errors import FailedPredicateException
                            raise FailedPredicateException(self, "self.precpred(self._ctx, 4)")
                        self.state = 307
                        self._errHandler.sync(self)
                        token = self._input.LA(1)
                        if token in [11]:
                            self.state = 304
                            self.match(DataLensParser.BETWEEN)
                        elif token in [26]:
                            self.state = 305
                            self.match(DataLensParser.NOT)
                            self.state = 306
                            self.match(DataLensParser.BETWEEN)
                        else:
                            raise NoViableAltException(self)

                        self.state = 309
                        self.exprMain(0)
                        self.state = 310
                        self.match(DataLensParser.AND)
                        self.state = 311
                        self.exprMain(5)

                    elif la_ == 7:
                        localctx = DataLensParser.UnaryPostfixContext(self, DataLensParser.ExprMainContext(self, _parentctx, _parentState))
                        self.pushNewRecursionContext(localctx, _startState, self.RULE_exprMain)
                        self.state = 313
                        if not self.precpred(self._ctx, 7):
                            from antlr4.error.Errors import FailedPredicateException
                            raise FailedPredicateException(self, "self.precpred(self._ctx, 7)")
                        self.state = 314
                        self.match(DataLensParser.IS)
                        self.state = 316
                        self._errHandler.sync(self)
                        _la = self._input.LA(1)
                        if _la == 26:
                            self.state = 315
                            self.match(DataLensParser.NOT)

                        self.state = 318
                        _la = self._input.LA(1)
                        if not (((_la) & ~0x3f) == 0 and ((1 << _la) & 4429447168) != 0):
                            self._errHandler.recoverInline(self)
                        else:
                            self._errHandler.reportMatch(self)
                            self.consume()

                    elif la_ == 8:
                        localctx = DataLensParser.InExprContext(self, DataLensParser.ExprMainContext(self, _parentctx, _parentState))
                        self.pushNewRecursionContext(localctx, _startState, self.RULE_exprMain)
                        self.state = 319
                        if not self.precpred(self._ctx, 3):
                            from antlr4.error.Errors import FailedPredicateException
                            raise FailedPredicateException(self, "self.precpred(self._ctx, 3)")
                        self.state = 321
                        self._errHandler.sync(self)
                        _la = self._input.LA(1)
                        if _la == 26:
                            self.state = 320
                            self.match(DataLensParser.NOT)

                        self.state = 323
                        self.match(DataLensParser.IN)
                        self.state = 324
                        self.match(DataLensParser.OPENING_PAR)
                        self.state = 334
                        self._errHandler.sync(self)
                        _la = self._input.LA(1)
                        if ((_la) & ~0x3f) == 0 and ((1 << _la) & 137512472549638) != 0:
                            self.state = 330
                            self._errHandler.sync(self)
                            _alt = self._interp.adaptivePredict(self._input, 35, self._ctx)
                            while _alt != 2 and _alt != ATN.INVALID_ALT_NUMBER:
                                if _alt == 1:
                                    self.state = 325
                                    self.expression()
                                    self.state = 326
                                    self.match(DataLensParser.T__2)
                                self.state = 332
                                self._errHandler.sync(self)
                                _alt = self._interp.adaptivePredict(self._input, 35, self._ctx)

                            self.state = 333
                            self.expression()

                        self.state = 336
                        self.match(DataLensParser.CLOSING_PAR)

                self.state = 341
                self._errHandler.sync(self)
                _alt = self._interp.adaptivePredict(self._input, 38, self._ctx)

        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.unrollRecursionContexts(_parentctx)
        return localctx

    class ExprSecondaryContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def getRuleIndex(self):
            return DataLensParser.RULE_exprSecondary

        def copyFrom(self, ctx: ParserRuleContext):
            super().copyFrom(ctx)

    class BinaryExprSecContext(ExprSecondaryContext):

        def __init__(self, parser, ctx: ParserRuleContext):  # actually a DataLensParser.ExprSecondaryContext
            super().__init__(parser)
            self.copyFrom(ctx)

        def exprSecondary(self, i: int = None):
            if i is None:
                return self.getTypedRuleContexts(DataLensParser.ExprSecondaryContext)
            else:
                return self.getTypedRuleContext(DataLensParser.ExprSecondaryContext, i)

        def AND(self):
            return self.getToken(DataLensParser.AND, 0)

        def OR(self):
            return self.getToken(DataLensParser.OR, 0)

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitBinaryExprSec"):
                return visitor.visitBinaryExprSec(self)
            else:
                return visitor.visitChildren(self)

    class ExprMainAltContext(ExprSecondaryContext):

        def __init__(self, parser, ctx: ParserRuleContext):  # actually a DataLensParser.ExprSecondaryContext
            super().__init__(parser)
            self.copyFrom(ctx)

        def exprMain(self):
            return self.getTypedRuleContext(DataLensParser.ExprMainContext, 0)

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitExprMainAlt"):
                return visitor.visitExprMainAlt(self)
            else:
                return visitor.visitChildren(self)

    def exprSecondary(self, _p: int = 0):
        _parentctx = self._ctx
        _parentState = self.state
        localctx = DataLensParser.ExprSecondaryContext(self, self._ctx, _parentState)

        _startState = 52
        self.enterRecursionRule(localctx, 52, self.RULE_exprSecondary, _p)
        try:
            self.enterOuterAlt(localctx, 1)
            localctx = DataLensParser.ExprMainAltContext(self, localctx)
            self._ctx = localctx

            self.state = 343
            self.exprMain(0)
            self._ctx.stop = self._input.LT(-1)
            self.state = 353
            self._errHandler.sync(self)
            _alt = self._interp.adaptivePredict(self._input, 40, self._ctx)
            while _alt != 2 and _alt != ATN.INVALID_ALT_NUMBER:
                if _alt == 1:
                    if self._parseListeners is not None:
                        self.triggerExitRuleEvent()

                    self.state = 351
                    self._errHandler.sync(self)
                    la_ = self._interp.adaptivePredict(self._input, 39, self._ctx)
                    if la_ == 1:
                        localctx = DataLensParser.BinaryExprSecContext(self, DataLensParser.ExprSecondaryContext(self, _parentctx, _parentState))
                        self.pushNewRecursionContext(localctx, _startState, self.RULE_exprSecondary)
                        self.state = 345
                        if not self.precpred(self._ctx, 3):
                            from antlr4.error.Errors import FailedPredicateException
                            raise FailedPredicateException(self, "self.precpred(self._ctx, 3)")
                        self.state = 346
                        self.match(DataLensParser.AND)
                        self.state = 347
                        self.exprSecondary(4)

                    elif la_ == 2:
                        localctx = DataLensParser.BinaryExprSecContext(self, DataLensParser.ExprSecondaryContext(self, _parentctx, _parentState))
                        self.pushNewRecursionContext(localctx, _startState, self.RULE_exprSecondary)
                        self.state = 348
                        if not self.precpred(self._ctx, 2):
                            from antlr4.error.Errors import FailedPredicateException
                            raise FailedPredicateException(self, "self.precpred(self._ctx, 2)")
                        self.state = 349
                        self.match(DataLensParser.OR)
                        self.state = 350
                        self.exprSecondary(3)

                self.state = 355
                self._errHandler.sync(self)
                _alt = self._interp.adaptivePredict(self._input, 40, self._ctx)

        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.unrollRecursionContexts(_parentctx)
        return localctx

    class ExpressionContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def exprSecondary(self):
            return self.getTypedRuleContext(DataLensParser.ExprSecondaryContext, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_expression

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitExpression"):
                return visitor.visitExpression(self)
            else:
                return visitor.visitChildren(self)

    def expression(self):

        localctx = DataLensParser.ExpressionContext(self, self._ctx, self.state)
        self.enterRule(localctx, 54, self.RULE_expression)
        try:
            self.enterOuterAlt(localctx, 1)
            self.state = 356
            self.exprSecondary(0)
        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    class ParseContext(ParserRuleContext):
        __slots__ = 'parser'

        def __init__(self, parser, parent: ParserRuleContext = None, invokingState: int = -1):
            super().__init__(parent, invokingState)
            self.parser = parser

        def expression(self):
            return self.getTypedRuleContext(DataLensParser.ExpressionContext, 0)

        def EOF(self):
            return self.getToken(DataLensParser.EOF, 0)

        def getRuleIndex(self):
            return DataLensParser.RULE_parse

        def accept(self, visitor: ParseTreeVisitor):
            if hasattr(visitor, "visitParse"):
                return visitor.visitParse(self)
            else:
                return visitor.visitChildren(self)

    def parse(self):

        localctx = DataLensParser.ParseContext(self, self._ctx, self.state)
        self.enterRule(localctx, 56, self.RULE_parse)
        try:
            self.state = 362
            self._errHandler.sync(self)
            token = self._input.LA(1)
            if token in [1, 2, 8, 11, 12, 18, 20, 25, 26, 27, 28, 32, 36, 40, 42, 43, 44, 45, 46]:
                self.enterOuterAlt(localctx, 1)
                self.state = 358
                self.expression()
                self.state = 359
                self.match(DataLensParser.EOF)
            elif token in [-1]:
                self.enterOuterAlt(localctx, 2)
                self.state = 361
                self.match(DataLensParser.EOF)
            else:
                raise NoViableAltException(self)

        except RecognitionException as re:
            localctx.exception = re
            self._errHandler.reportError(self, re)
            self._errHandler.recover(self, re)
        finally:
            self.exitRule()
        return localctx

    def sempred(self, localctx: RuleContext, ruleIndex: int, predIndex: int):
        if self._predicates is None:
            self._predicates = dict()
        self._predicates[25] = self.exprMain_sempred
        self._predicates[26] = self.exprSecondary_sempred
        pred = self._predicates.get(ruleIndex, None)
        if pred is None:
            raise Exception("No predicate with index:" + str(ruleIndex))
        else:
            return pred(localctx, predIndex)

    def exprMain_sempred(self, localctx: ExprMainContext, predIndex: int):
        if predIndex == 0:
            return self.precpred(self._ctx, 11)

        if predIndex == 1:
            return self.precpred(self._ctx, 9)

        if predIndex == 2:
            return self.precpred(self._ctx, 8)

        if predIndex == 3:
            return self.precpred(self._ctx, 6)

        if predIndex == 4:
            return self.precpred(self._ctx, 5)

        if predIndex == 5:
            return self.precpred(self._ctx, 4)

        if predIndex == 6:
            return self.precpred(self._ctx, 7)

        if predIndex == 7:
            return self.precpred(self._ctx, 3)

    def exprSecondary_sempred(self, localctx: ExprSecondaryContext, predIndex: int):
        if predIndex == 8:
            return self.precpred(self._ctx, 3)

        if predIndex == 9:
            return self.precpred(self._ctx, 2)

lib/dl_formula/dl_formula/parser/antlr/gen/DataLensParser.py
datalens-tech/datalens-backend
from antlr4 import ParseTreeVisitor

if __name__ is not None and "." in __name__:
    from .DataLensParser import DataLensParser
else:
    from DataLensParser import DataLensParser

# This class defines a complete generic visitor for a parse tree produced by DataLensParser.


class DataLensVisitor(ParseTreeVisitor):

    # Visit a parse tree produced by DataLensParser#integerLiteral.
    def visitIntegerLiteral(self, ctx: DataLensParser.IntegerLiteralContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#floatLiteral.
    def visitFloatLiteral(self, ctx: DataLensParser.FloatLiteralContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#stringLiteral.
    def visitStringLiteral(self, ctx: DataLensParser.StringLiteralContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#dateLiteral.
    def visitDateLiteral(self, ctx: DataLensParser.DateLiteralContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#datetimeLiteral.
    def visitDatetimeLiteral(self, ctx: DataLensParser.DatetimeLiteralContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#genericDateLiteral.
    def visitGenericDateLiteral(self, ctx: DataLensParser.GenericDateLiteralContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#genericDatetimeLiteral.
    def visitGenericDatetimeLiteral(self, ctx: DataLensParser.GenericDatetimeLiteralContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#boolLiteral.
    def visitBoolLiteral(self, ctx: DataLensParser.BoolLiteralContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#nullLiteral.
    def visitNullLiteral(self, ctx: DataLensParser.NullLiteralContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#fieldName.
    def visitFieldName(self, ctx: DataLensParser.FieldNameContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#orderingItem.
    def visitOrderingItem(self, ctx: DataLensParser.OrderingItemContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#ordering.
    def visitOrdering(self, ctx: DataLensParser.OrderingContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#lodSpecifier.
    def visitLodSpecifier(self, ctx: DataLensParser.LodSpecifierContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#winGrouping.
    def visitWinGrouping(self, ctx: DataLensParser.WinGroupingContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#beforeFilterBy.
    def visitBeforeFilterBy(self, ctx: DataLensParser.BeforeFilterByContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#ignoreDimensions.
    def visitIgnoreDimensions(self, ctx: DataLensParser.IgnoreDimensionsContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#function.
    def visitFunction(self, ctx: DataLensParser.FunctionContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#elseifPart.
    def visitElseifPart(self, ctx: DataLensParser.ElseifPartContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#elsePart.
    def visitElsePart(self, ctx: DataLensParser.ElsePartContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#ifPart.
    def visitIfPart(self, ctx: DataLensParser.IfPartContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#ifBlock.
    def visitIfBlock(self, ctx: DataLensParser.IfBlockContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#whenPart.
    def visitWhenPart(self, ctx: DataLensParser.WhenPartContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#caseBlock.
    def visitCaseBlock(self, ctx: DataLensParser.CaseBlockContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#parenthesizedExpr.
    def visitParenthesizedExpr(self, ctx: DataLensParser.ParenthesizedExprContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#exprBasic.
    def visitExprBasic(self, ctx: DataLensParser.ExprBasicContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#exprBasicAlt.
    def visitExprBasicAlt(self, ctx: DataLensParser.ExprBasicAltContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#unaryPrefix.
    def visitUnaryPrefix(self, ctx: DataLensParser.UnaryPrefixContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#inExpr.
    def visitInExpr(self, ctx: DataLensParser.InExprContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#binaryExpr.
    def visitBinaryExpr(self, ctx: DataLensParser.BinaryExprContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#comparisonChain.
    def visitComparisonChain(self, ctx: DataLensParser.ComparisonChainContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#unaryPostfix.
    def visitUnaryPostfix(self, ctx: DataLensParser.UnaryPostfixContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#betweenExpr.
    def visitBetweenExpr(self, ctx: DataLensParser.BetweenExprContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#binaryExprSec.
    def visitBinaryExprSec(self, ctx: DataLensParser.BinaryExprSecContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#exprMainAlt.
    def visitExprMainAlt(self, ctx: DataLensParser.ExprMainAltContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#expression.
    def visitExpression(self, ctx: DataLensParser.ExpressionContext):
        return self.visitChildren(ctx)

    # Visit a parse tree produced by DataLensParser#parse.
    def visitParse(self, ctx: DataLensParser.ParseContext):
        return self.visitChildren(ctx)


del DataLensParser

lib/dl_formula/dl_formula/parser/antlr/gen/DataLensVisitor.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import _CacheInfo
from typing import Dict

import antlr4
from antlr4.Token import Token
from antlr4.error.Errors import ParseCancellationException

import dl_formula.core.exc as exc
import dl_formula.core.nodes as nodes
from dl_formula.core.position import PositionConverter
from dl_formula.parser.antlr.visitor import CustomDataLensVisitor
from dl_formula.parser.base import (
    FORMULA_CACHE_SIZE,
    FormulaParser,
    parser_cache_qualifier,
)
from dl_formula.utils.caching import multi_cached_with_errors


try:
    from dl_formula.parser.antlr.gen.DataLensLexer import DataLensLexer
    from dl_formula.parser.antlr.gen.DataLensParser import DataLensParser
except ImportError:
    raise exc.ParserNotFoundError()


@multi_cached_with_errors(
    FORMULA_CACHE_SIZE,
    cache_exceptions=(exc.ParseError,),
    cache_qualifier=parser_cache_qualifier,
)
def parse(formula: str) -> nodes.Formula:
    pos_conv = PositionConverter(text=formula)

    if not formula.strip():
        raise exc.ParseEmptyFormulaError("Empty formula", position=pos_conv.idx_to_position(0))

    input_stream = antlr4.InputStream(formula)
    lexer = DataLensLexer(input_stream)
    stream = antlr4.CommonTokenStream(lexer)
    parser = DataLensParser(stream)
    parser._errHandler = antlr4.BailErrorStrategy()
    try:
        tree = parser.parse()
    except ParseCancellationException as err:
        raw_token = err.args[0].offendingToken
        if raw_token.type == Token.EOF:
            raise exc.ParseUnexpectedEOFError(
                "Failed to parse: unexpected end of formula",
                position=pos_conv.idx_to_position(idx=raw_token.start),
            )
        else:
            raise exc.ParseUnexpectedTokenError(
                "Failed to parse formula: unexpected token",
                position=pos_conv.idx_to_position(idx=raw_token.start),
                token=raw_token.text,
            )

    formula_obj = CustomDataLensVisitor(text=formula).visitParse(tree)
    return formula_obj


class AntlrPyFormulaParser(FormulaParser):
    def _parse(self, formula: str) -> nodes.Formula:
        return parse(formula)

    def _get_global_stats(self) -> Dict[str, _CacheInfo]:
        return parse.cache_info()  # type: ignore

lib/dl_formula/dl_formula/parser/antlr/main.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from typing import (
    Iterable,
    List,
    Optional,
    Tuple,
    Type,
)

import antlr4
from antlr4.Token import Token
from antlr4.tree.Tree import TerminalNodeImpl

import dl_formula.core.exc as exc
import dl_formula.core.nodes as nodes
from dl_formula.core.position import (
    Position,
    PositionConverter,
)
from dl_formula.parser.base import resolve_function_capabilities
from dl_formula.utils.datetime import make_datetime_value


try:
    from dl_formula.parser.antlr.gen.DataLensParser import DataLensParser
    from dl_formula.parser.antlr.gen.DataLensVisitor import DataLensVisitor
except ImportError:
    raise exc.ParserNotFoundError()


EMPTY_FORMULA_ERROR = "Empty formula is invalid"

COMPARISON_OP_RAW_NAMES = ("=", "!=", "<>", ">", ">=", "<", "<=")


class CustomDataLensVisitor(DataLensVisitor):
    _BOOL = {"true": True, "false": False}

    def __init__(self, text: str):
        self._text = text
        self._pos_conv = PositionConverter(text=text)

    def _make_node_meta(self, *ctxes: antlr4.ParserRuleContext) -> Optional[nodes.NodeMeta]:  # type: ignore
        first_ctx = ctxes[0]
        last_ctx = ctxes[-1]
        return nodes.NodeMeta(
            position=self._make_position(*ctxes),
            original_text=self._text[first_ctx.start.start : last_ctx.stop.stop + 1],  # type: ignore  # TODO: fix
        )

    def _make_position(self, *ctxes: antlr4.ParserRuleContext) -> Optional[Position]:  # type: ignore  # TODO: fix
        first_ctx = ctxes[0]
        last_ctx = ctxes[-1]
        return self._pos_conv.merge_positions(
            start_position=self._pos_conv.idx_to_position(first_ctx.start.start),  # type: ignore  # TODO: fix
            end_position=self._pos_conv.idx_to_position(last_ctx.stop.stop + 1),  # type: ignore  # TODO: fix
        )

    def _separate_children(
        self,
        ctx: antlr4.ParserRuleContext,  # type: ignore
        exclude: Iterable[str] = (),
        lower_str: bool = False,
    ) -> Tuple[List[str], List[nodes.FormulaItem]]:
        str_children = []
        node_children = []
        for child in ctx.children:  # type: ignore  # TODO: fix
            if isinstance(child, TerminalNodeImpl):
                child_str = str(child)
                if lower_str:
                    child_str = child_str.lower()
                if child_str not in exclude:
                    str_children.append(child_str)
            else:
                node_children.append(self.visit(child))

        return str_children, node_children

    def visitIntegerLiteral(self, ctx: DataLensParser.IntegerLiteralContext):
        return nodes.LiteralInteger.make(value=int(str(ctx.children[0])), meta=self._make_node_meta(ctx))

    def visitFloatLiteral(self, ctx: DataLensParser.FloatLiteralContext):
        return nodes.LiteralFloat.make(value=float(str(ctx.children[0])), meta=self._make_node_meta(ctx))

    def visitStringLiteral(self, ctx: DataLensParser.StringLiteralContext):
        s = str(ctx.children[0])[1:-1]  # strip off the quotes
        s = s.replace("\\'", "'").replace('\\"', '"')
        s = s.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t")
        s = s.replace("\\\\", "\\")
        return nodes.LiteralString.make(value=s, meta=self._make_node_meta(ctx))

    def visitDateLiteral(self, ctx: DataLensParser.DateLiteralContext):
        date_str = str(ctx.children[1])
        try:
            value = datetime.datetime.strptime(date_str, "%Y-%m-%d").date()
        except ValueError:
            raise exc.ParseDateValueError(
                f"Invalid date value: {date_str}", token=date_str, position=self._make_position(ctx)
            )
        return nodes.LiteralDate.make(value=value, meta=self._make_node_meta(ctx))

    def visitDatetimeLiteral(self, ctx: DataLensParser.DatetimeLiteralContext):
        date_str = str(ctx.children[1])
        try:
            value = make_datetime_value(date_str)
            assert value is not None
        except ValueError:
            raise exc.ParseDatetimeValueError(
                f"Invalid datetime value: {date_str}", token=date_str, position=self._make_position(ctx)
            )
        node_cls = nodes.LiteralGenericDatetime
        return node_cls.make(value=value, meta=self._make_node_meta(ctx))

    def visitGenericDateLiteral(self, ctx: DataLensParser.GenericDateLiteralContext):
        date_str = str(ctx.children[1])
        try:
            value = datetime.datetime.strptime(date_str, "%Y-%m-%d").date()
        except ValueError:
            raise exc.ParseDateValueError(
                f"Invalid date value: {date_str}", token=date_str, position=self._make_position(ctx)
            )
        return nodes.LiteralDate.make(value=value, meta=self._make_node_meta(ctx))

    def visitGenericDatetimeLiteral(self, ctx: DataLensParser.GenericDatetimeLiteralContext):
        date_str = str(ctx.children[1])
        try:
            value = make_datetime_value(date_str)
            assert value is not None
        except ValueError:
            raise exc.ParseDatetimeValueError(
                f"Invalid datetime value: {date_str}", token=date_str, position=self._make_position(ctx)
            )
        node_cls = nodes.LiteralGenericDatetime
        return node_cls.make(value=value, meta=self._make_node_meta(ctx))

    def visitBoolLiteral(self, ctx: DataLensParser.BoolLiteralContext):
        return nodes.LiteralBoolean.make(value=self._BOOL[str(ctx.children[0]).lower()], meta=self._make_node_meta(ctx))

    def visitNullLiteral(self, ctx: DataLensParser.NullLiteralContext):
        return nodes.Null(meta=self._make_node_meta(ctx))

    def visitFieldName(self, ctx: DataLensParser.FieldNameContext):
        return nodes.Field.make(name=str(ctx.children[0])[1:-1], meta=self._make_node_meta(ctx))

    def visitOrderingItem(self, ctx: DataLensParser.OrderingItemContext):
        str_children, node_children = self._separate_children(ctx)
        expr = node_children[0]

        if str_children:
            direction = str_children[0].lower()
            if direction == "asc":
                expr = nodes.OrderAscending.make(expr=expr, meta=self._make_node_meta(ctx))
            elif direction == "desc":
                expr = nodes.OrderDescending.make(expr=expr, meta=self._make_node_meta(ctx))
            else:
                raise ValueError("Unexpected value for ordering item")

        return expr

    def visitOrdering(self, ctx: DataLensParser.OrderingContext):
        _, node_children = self._separate_children(ctx)
        return nodes.Ordering.make(expr_list=node_children, meta=self._make_node_meta(ctx))

    def visitLodSpecifier(self, ctx: DataLensParser.LodSpecifierContext):
        str_nodes, node_children = self._separate_children(ctx)
        kind = str_nodes[0].lower()

        lod_cls: Type[nodes.LodSpecifier] = {
            "fixed": nodes.FixedLodSpecifier,
            "include": nodes.IncludeLodSpecifier,
            "exclude": nodes.ExcludeLodSpecifier,
        }[kind]

        return lod_cls.make(dim_list=node_children, meta=self._make_node_meta(ctx))

    def visitWinGrouping(self, ctx: DataLensParser.WinGroupingContext):
        str_nodes, node_children = self._separate_children(ctx)
        kind = str_nodes[0].lower()

        if kind == "total":
            return nodes.WindowGroupingTotal(meta=self._make_node_meta(ctx))

        grouping_cls = {
            "among": nodes.WindowGroupingAmong,
            "within": nodes.WindowGroupingWithin,
        }[kind]

        return grouping_cls.make(dim_list=node_children, meta=self._make_node_meta(ctx))

    def visitBeforeFilterBy(self, ctx: DataLensParser.BeforeFilterByContext):
        _, node_children = self._separate_children(ctx)
        field_names: List[str] = []
        for field in node_children:
            assert isinstance(field, nodes.Field)
            field_names.append(field.name)
        return nodes.BeforeFilterBy.make(field_names=frozenset(field_names), meta=self._make_node_meta(ctx))

    def visitIgnoreDimensions(self, ctx: DataLensParser.IgnoreDimensionsContext):
        _, node_children = self._separate_children(ctx)
        return nodes.IgnoreDimensions.make(dim_list=node_children, meta=self._make_node_meta(ctx))

    def visitFunction(self, ctx: DataLensParser.FunctionContext):
        str_nodes, node_children = self._separate_children(ctx)
        name = str_nodes[0].lower()

        grouping = next((arg for arg in node_children if isinstance(arg, nodes.WindowGrouping)), None)
        ordering = next((arg for arg in node_children if isinstance(arg, nodes.Ordering)), None)
        ignore_dimensions = next((arg for arg in node_children if isinstance(arg, nodes.IgnoreDimensions)), None)
        before_filter_by = next((arg for arg in node_children if isinstance(arg, nodes.BeforeFilterBy)), None)
        lod = next((arg for arg in node_children if isinstance(arg, nodes.LodSpecifier)), None)

        args = [
            arg
            for arg in node_children
            if not isinstance(
                arg,
                (
                    str,
                    nodes.WindowGrouping,
                    nodes.Ordering,
                    nodes.BeforeFilterBy,
                    nodes.IgnoreDimensions,
                    nodes.LodSpecifier,
                ),
            )
        ]

        function_capabilities = resolve_function_capabilities(
            name=name,
            grouping=grouping,
            lod=lod,
            ordering=ordering,
            before_filter_by=before_filter_by,
            ignore_dimensions=ignore_dimensions,
        )

        if function_capabilities.is_window:
            return nodes.WindowFuncCall.make(
                name=name,
                args=args,
                grouping=grouping,
                ordering=ordering,
                ignore_dimensions=ignore_dimensions,
                before_filter_by=before_filter_by,
                meta=self._make_node_meta(ctx),
            )

        return nodes.FuncCall.make(
            name=name,
            args=args,
            lod=lod,
            ignore_dimensions=ignore_dimensions,
            before_filter_by=before_filter_by,
            meta=self._make_node_meta(ctx),
        )

    def visitElsePart(self, ctx: DataLensParser.ElsePartContext):
        return self.visit(ctx.children[1])

    def visitIfPart(self, ctx: DataLensParser.IfPartContext):
        return nodes.IfPart.make(
            cond=self.visit(ctx.children[1]),
            expr=self.visit(ctx.children[3]),
            meta=self._make_node_meta(ctx),
        )

    def visitElseifPart(self, ctx: DataLensParser.ElseifPartContext):
        return nodes.IfPart.make(
            cond=self.visit(ctx.children[1]),
            expr=self.visit(ctx.children[3]),
            meta=self._make_node_meta(ctx),
        )

    def visitIfBlock(self, ctx: DataLensParser.IfBlockContext):
        # args: [*if_parts, ?else_part, "END"]
        assert len(ctx.children) >= 2

        has_explicit_else = not isinstance(
            ctx.children[-2], (DataLensParser.IfPartContext, DataLensParser.ElseifPartContext)
        )

        if has_explicit_else:
            else_expr = self.visit(ctx.children[-2])
            if_list = [self.visit(ch) for ch in ctx.children[0:-2]]
        else:
            else_expr = nodes.Null()
            if_list = [self.visit(ch) for ch in ctx.children[0:-1]]

        assert all(isinstance(item, nodes.IfPart) for item in if_list)
        return nodes.IfBlock.make(
            if_list=if_list,
            else_expr=else_expr,
            meta=self._make_node_meta(ctx),
        )

    def visitWhenPart(self, ctx: DataLensParser.WhenPartContext):
        # args: ["WHEN", val, "THEN", expr]
        assert len(ctx.children) == 4
        return nodes.WhenPart.make(
            val=self.visit(ctx.children[1]),
            expr=self.visit(ctx.children[3]) if len(ctx.children) > 2 else None,
            meta=self._make_node_meta(ctx),
        )

    def visitCaseBlock(self, ctx: DataLensParser.CaseBlockContext):
        # args: ["CASE", case_expr, *when_parts, ?else_expr, "END"]
        assert len(ctx.children) >= 4

        case_expr = self.visit(ctx.children[1])

        has_explicit_else = not isinstance(ctx.children[-2], DataLensParser.WhenPartContext)

        if has_explicit_else:
            else_expr = self.visit(ctx.children[-2])
            when_list = [self.visit(ch) for ch in ctx.children[2:-2]]
        else:
            else_expr = nodes.Null()
            when_list = [self.visit(ch) for ch in ctx.children[2:-1]]

        assert all(isinstance(item, nodes.WhenPart) for item in when_list)

        return nodes.CaseBlock.make(
            case_expr=case_expr,
            when_list=when_list,
            else_expr=else_expr,
            meta=self._make_node_meta(ctx),
        )

    def visitParenthesizedExpr(self, ctx: DataLensParser.ParenthesizedExprContext):
        return nodes.ParenthesizedExpr.make(
            expr=self.visit(ctx.children[1]),
            meta=self._make_node_meta(ctx),
        )

    def visitUnaryPrefix(self, ctx: DataLensParser.UnaryPrefixContext):
        str_children, node_children = self._separate_children(ctx, lower_str=True)
        op_name = "".join(str_children)
        expr = node_children[0]

        if op_name == "-" and isinstance(expr, (nodes.LiteralInteger, nodes.LiteralFloat)):
            # negative number
            return expr.__class__.make(value=-expr.value, meta=self._make_node_meta(ctx))

        if op_name == "-":
            op_name = "neg"  # to tell apart from minus
        return nodes.Unary.make(name=op_name, expr=expr, meta=self._make_node_meta(ctx))

    def visitBinaryExpr(self, ctx: DataLensParser.BinaryExprContext):
        str_children, node_children = self._separate_children(ctx)
        op_name = "".join(str_children).lower()
        return nodes.Binary.make(
            name=op_name,
            left=node_children[0],
            right=node_children[1],
            meta=self._make_node_meta(ctx),
        )

    def visitBinaryExprSec(self, ctx: DataLensParser.BinaryExprContext):
        return self.visitBinaryExpr(ctx)

    def visitInExpr(self, ctx: DataLensParser.InExprContext):
        str_children, node_children = self._separate_children(ctx, exclude=",()")
        raw_exlist_nodes = [child for child in ctx.children if not isinstance(child, TerminalNodeImpl)][1:]
        op_name = "".join(str_children).lower()
        child_meta: Optional[nodes.NodeMeta] = None
        if raw_exlist_nodes:
            child_meta = self._make_node_meta(*raw_exlist_nodes)
        return nodes.Binary.make(
            name=op_name,
            left=node_children[0],
            right=nodes.ExpressionList.make(*node_children[1:], meta=child_meta),
            meta=self._make_node_meta(ctx),
        )

    def visitComparisonChain(self, ctx: DataLensParser.ComparisonChainContext):
        # First flatten comparison binary calls that get nested by the parser

        def is_cmp(_ctx) -> bool:
            return (
                len(_ctx.children) == 3
                and isinstance(_ctx.children[1], TerminalNodeImpl)
                and str(_ctx.children[1]) in COMPARISON_OP_RAW_NAMES
            )

        def _flatten_and_visit_cmp_ctx(
            _ctx: antlr4.ParserRuleContext,  # type: ignore
        ) -> List[antlr4.ParserRuleContext]:  # type: ignore
            if is_cmp(_ctx):
                return [
                    *_flatten_and_visit_cmp_ctx(_ctx.children[0]),  # type: ignore  # TODO: fix
                    str(_ctx.children[1]),  # type: ignore  # TODO: fix
                    self.visit(_ctx.children[2]),  # type: ignore  # TODO: fix
                ]
            return [self.visit(_ctx)]

        args = _flatten_and_visit_cmp_ctx(ctx)

        if len(args) > 3:  # is a chain of more than one operator
            left, *middle_list, last_op, right = args
            middle_list = [
                [middle_list[i], middle_list[i + 1]] for i in range(0, len(middle_list), 2)
            ]  # repartition list into <operator, operand> pairs
            parts = middle_list + [(last_op, right)]
        elif len(args) == 3:
            left, last_op, right = args
            parts = [(last_op, right)]
        else:
            left, last_op = args
            right = nodes.Null()
            parts = [(last_op, right)]

        latest_pair = None
        latest_left = left
        for op_name, right in parts:  # type: ignore  # TODO: fix
            if op_name == "=":  # type: ignore  # TODO: fix
                op_name = "=="
            if op_name == "<>":
                op_name = "!="
            new_pair = nodes.Binary.make(
                name=op_name,
                left=latest_left,
                right=right,
                meta=self._make_node_meta(ctx),
            )
            if latest_pair is not None:
                new_pair = nodes.Binary.make(
                    name="and",
                    left=latest_pair,
                    right=new_pair,
                    meta=self._make_node_meta(ctx),
                )
            latest_left = right
            latest_pair = new_pair

        assert latest_pair is not None, "At least two operands are required in comparison chain"
        return latest_pair

    def visitUnaryPostfix(self, ctx: DataLensParser.UnaryPostfixContext):
        str_children, node_children = self._separate_children(ctx, lower_str=True)
        op_name = "".join([ch for ch in str_children if ch != "not"])
        expr = nodes.Unary.make(name=op_name, expr=node_children[0], meta=self._make_node_meta(ctx))
        if "not" in str_children:
            expr = nodes.Unary.make(name="not", expr=expr, meta=self._make_node_meta(ctx))
        return expr

    def visitBetweenExpr(self, ctx: DataLensParser.BetweenExprContext):
        str_children, node_children = self._separate_children(ctx, lower_str=True, exclude=("and",))
        op_name = "".join(str_children)
        return nodes.Ternary.make(
            name=op_name,
            first=node_children[0],
            second=node_children[1],
            third=node_children[2],
            meta=self._make_node_meta(ctx),
        )

    def visitParse(self, ctx: DataLensParser.ParseContext):
        if (
            len(ctx.children) > 0
            and isinstance(ctx.children[0], TerminalNodeImpl)
            and ctx.children[0].symbol.type == Token.EOF
        ):
            raise exc.ParseEmptyFormulaError(EMPTY_FORMULA_ERROR)
        return nodes.Formula.make(expr=self.visit(ctx.children[0]), meta=self._make_node_meta(ctx))

lib/dl_formula/dl_formula/parser/antlr/visitor.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from collections import defaultdict
from contextlib import contextmanager
from functools import _CacheInfo
import os
import statistics
import time
from typing import (
    Dict,
    Generator,
    List,
    NamedTuple,
    Optional,
)

import attr

import dl_formula.core.exc as exc
import dl_formula.core.nodes as nodes
from dl_formula.inspect.function import (
    can_be_nonwindow,
    can_be_window,
    supports_bfb,
    supports_grouping,
    supports_ignore_dimensions,
    supports_lod,
    supports_ordering,
)


FORMULA_CACHE_SIZE = int(os.environ.get("FORMULA_CACHE_SIZE", 1000))


class ParserStats(NamedTuple):
    global_cache_hits: Dict[str, int]
    global_cache_misses: Dict[str, int]
    global_cache_maxsize: Dict[str, Optional[int]]
    global_cache_currsize: Dict[str, int]
    call_count: int
    call_time: float
    avg_length: Optional[float]
    length_counts: Dict[int, int]
    length_avg_times: Dict[int, float]
    length_median_times: Dict[int, float]


@attr.s
class FormulaParser(abc.ABC):
    call_count: int = attr.ib(init=False, default=0)
    call_time: float = attr.ib(init=False, default=0.0)
    length_sum: int = attr.ib(init=False, default=0)
    by_length: Dict[int, List[float]] = attr.ib(init=False, factory=lambda: defaultdict(list))

    @contextmanager
    def _parse_cm(self, formula: str) -> Generator[None, None, None]:
        started = time.thread_time()
        try:
            yield
        finally:
            call_time = time.thread_time() - started
            self.call_time += call_time
            self.call_count += 1
            self.by_length[len(formula)].append(call_time)

    def parse(self, formula: str) -> nodes.Formula:
        with self._parse_cm(formula=formula):
            return self._parse(formula)

    @abc.abstractmethod
    def _parse(self, formula: str) -> nodes.Formula:
        raise NotImplementedError

    @abc.abstractmethod
    def _get_global_stats(self) -> Dict[str, _CacheInfo]:
        raise NotImplementedError

    def get_usage_stats(self) -> ParserStats:
        global_stats = self._get_global_stats()
        return ParserStats(
            global_cache_hits={key: stats.hits for key, stats in global_stats.items()},
            global_cache_misses={key: stats.misses for key, stats in global_stats.items()},
            global_cache_maxsize={key: stats.maxsize for key, stats in global_stats.items()},
            global_cache_currsize={key: stats.currsize for key, stats in global_stats.items()},
            call_count=self.call_count,
            call_time=self.call_time,
            avg_length=self.length_sum / self.call_count if self.call_count else None,
            length_counts={length: len(times) for length, times in self.by_length.items()},
            length_avg_times={length: statistics.mean(times) for length, times in self.by_length.items()},
            length_median_times={length: statistics.median(times) for length, times in self.by_length.items()},
        )


def parser_cache_qualifier(formula: str) -> str:
    """
    When used as a cache qualifier,
    function generates 3 separate caches for formula parser
    and routes calls to these caches by formula length.
    """
    length = len(formula)
    if length > 500:
        return "f500_inf"
    if length > 100:
        return "f100_500"
    return "f0_100"


class FunctionCapabilities(NamedTuple):
    is_window: bool


def resolve_function_capabilities(
    name: str,
    grouping: Optional[nodes.WindowGrouping] = None,
    ordering: Optional[nodes.Ordering] = None,
    before_filter_by: Optional[nodes.BeforeFilterBy] = None,
    ignore_dimensions: Optional[nodes.IgnoreDimensions] = None,
    lod: Optional[nodes.LodSpecifier] = None,
) -> FunctionCapabilities:
    can_be_window_res = can_be_window(name)
    can_be_nonwindow_res = can_be_nonwindow(name)

    is_window = False
    if can_be_window_res:
        if can_be_nonwindow_res:
            # For functions that can be both window and non-window
            # the type is determined by the presence of a grouping clause
            is_window = grouping is not None
        else:
            # Can only be window
            is_window = True

    func_t_str = "window" if is_window else "non-window"
    u_name = name.upper

    if grouping is not None and not supports_grouping(name, is_window=is_window):
        raise exc.ParseClauseError(f"Grouping clause is not supported by {func_t_str} function {u_name()}")
    if ordering is not None and not supports_ordering(name, is_window=is_window):
        raise exc.ParseClauseError(f"ORDER BY clause is not supported by {func_t_str} function {u_name()}")
    if lod is not None and not supports_lod(name, is_window=is_window):
        raise exc.ParseClauseError(f"Level-of-detail clause is not supported by {func_t_str} function {u_name()}")
    if ignore_dimensions is not None and not supports_ignore_dimensions(name, is_window=is_window):
        raise exc.ParseClauseError(f"IGNORE DIMENSIONS clause is not supported by {func_t_str} function {u_name()}")
    if before_filter_by is not None and not supports_bfb(name, is_window=is_window):
        raise exc.ParseClauseError(f"BEFORE FILTER BY clause is not supported by {func_t_str} function {u_name()}")

    return FunctionCapabilities(
        is_window=is_window,
    )

lib/dl_formula/dl_formula/parser/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import (
    Enum,
    unique,
)

from dl_formula.parser.antlr.main import AntlrPyFormulaParser
from dl_formula.parser.base import FormulaParser


@unique
class ParserType(Enum):
    antlr_py = "antlr_py"


def get_parser(parser_type: ParserType = ParserType.antlr_py) -> FormulaParser:
    if parser_type == ParserType.antlr_py:
        return AntlrPyFormulaParser()

    raise ValueError(f"Unknown parser type: {parser_type}")

lib/dl_formula/dl_formula/parser/factory.py
datalens-tech/datalens-backend
from __future__ import annotations

import os
import tempfile
from typing import (
    Callable,
    List,
    TypeVar,
    Union,
)
import uuid

from graphviz.dot import Dot


_TV = TypeVar("_TV")


def csv_type(type_func: Callable[[str], _TV]) -> Callable[[str], Union[List[_TV], str]]:
    def wrapper(s: str) -> Union[List[_TV], str]:
        if isinstance(s, str):
            return [type_func(it) for it in s.split(",")]
        return s

    return wrapper


def make_graphviz_graph(dot: Dot, render_to: str, view: bool) -> None:
    if render_to:
        if not render_to.endswith(".png"):
            raise ValueError("--render-to must have .png extension")
        render_to = render_to[:-4]

    if render_to or view:
        if not render_to:
            render_to = os.path.join(tempfile.gettempdir(), "dotfile_{}".format(uuid.uuid4().hex))

        dot.format = "png"
        dot.render(render_to, view=view, cleanup=True)

    else:
        print(dot.source)

lib/dl_formula/dl_formula/scripts/common.py
datalens-tech/datalens-backend
from __future__ import annotations

import argparse
from collections import defaultdict
import inspect
import os
import sys
from typing import (
    Generator,
    List,
    Optional,
    Tuple,
    Union,
)

from dl_formula.core import exc
from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import (
    get_all_basic_dialects,
    get_dialect_from_str,
)
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.base import (
    MultiVariantTranslation,
    NodeTranslation,
)
from dl_formula.definitions.registry import OPERATION_REGISTRY
import dl_formula.dot
import dl_formula.inspect.expression
import dl_formula.inspect.function
from dl_formula.loader import load_formula_lib
from dl_formula.parser.factory import get_parser
from dl_formula.slicing.schema import (
    AggregateFunctionLevelBoundary,
    NonFieldsBoundary,
    SliceSchema,
    TopLevelBoundary,
    WindowFunctionLevelBoundary,
)
from dl_formula.slicing.slicer import FormulaSlicer
import dl_formula.translation.translator

from .common import (
    csv_type,
    make_graphviz_graph,
)


formula_parser = get_parser()


def dialect_type(s: Union[str, DialectCombo]) -> Union[DialectCombo, int]:
    if isinstance(s, str):
        res = D.EMPTY
        for it in s.split(","):
            res |= get_dialect_from_str(it.upper())
        return res
    return s


list_type = csv_type(str)


parser = argparse.ArgumentParser(prog="Formula command line tool")
subparsers = parser.add_subparsers(title="command", dest="command")

text_parser = argparse.ArgumentParser(add_help=False)
text_parser.add_argument("text", help="formula text")
optional_text_parser = argparse.ArgumentParser(add_help=False)
optional_text_parser.add_argument("optional_text", nargs="?", help="formula text")

dialect_parser = argparse.ArgumentParser(add_help=False)
dialect_parser.add_argument(
    "--dialect",
    required=True,
    type=dialect_type,
    help="SQL dialect (comma-separated list, use ALL for combination of all dialects)",
)

category_parser = argparse.ArgumentParser(add_help=False)
category_parser.add_argument("--category", help="Locale")

supperr_parser = argparse.ArgumentParser(add_help=False)
supperr_parser.add_argument("--suppress-errors", action="store_true", help="Print ERROR token instead of failing")

fields_parser = argparse.ArgumentParser(add_help=False)
fields_parser.add_argument("--fields", type=list_type, help="List of fields", default=[])

diff_parser = argparse.ArgumentParser(add_help=False)
diff_parser.add_argument("--diff", action="store_true", help="Diff mode")

trans_parser = subparsers.add_parser(
    "translate",
    parents=[optional_text_parser, dialect_parser, supperr_parser],
    help="Translate formula into an SQL expression",
)
trans_parser.add_argument("--table", help="Name of database table")
trans_parser.add_argument("--unknown-funcs", action="store_true", help="Allow usage of unknown functions")

parse_parser = subparsers.add_parser(
    "parse", parents=[optional_text_parser, supperr_parser], help="Parse formula and print its internal representation"
)
parse_parser.add_argument("--pretty", action="store_true", help="Print in pretty multiline format")
parse_parser.add_argument("--pos", type=int, help="Print node at the given position")
parse_parser.add_argument("--with-meta", action="store_true", help="Print meta for nodes")

split_parser = subparsers.add_parser(
    "split", parents=[text_parser, diff_parser], help="Split formula into nodes and print them by nesting level"
)

graph_parser = subparsers.add_parser("graph", parents=[text_parser], help="Generate GraphViz DOT graph")
graph_parser.add_argument("--view", action="store_true", help="Show rendered graph (requires GUI)")
graph_parser.add_argument("--render-to", help="Render graph to specified file", default=None)

dialect_parser = subparsers.add_parser("dialects", help="List available basic dialects")

goto_parser = subparsers.add_parser("goto", help="Open PyCharm at the function definition")
goto_parser.add_argument("func", help="function name")
goto_parser.add_argument("--show", action="store_true", help="just print file name and line number")

slice_parser = subparsers.add_parser(
    "slice", parents=[text_parser, diff_parser], help="Parse formula slice it into translation levels"
)
slice_parser.add_argument("--levels", type=list_type, help="Formula levels")

registry_parser = subparsers.add_parser("registry", help="List function definitions in registry")


ERROR_TOKEN = "#ERROR"


class FormulaCliTool:
    @staticmethod
    def _stdin_lines() -> Generator[str, None, None]:
        lines = sys.stdin.readlines()
        for line in lines:
            line = line.strip()
            if not line:
                continue

            yield line

    @classmethod
    def parse(cls, text: str, pretty: bool, suppress_errors: bool, pos: Optional[int] = None, with_meta: bool = False):
        if not text:
            # bulk mode, read from STDIN
            for line in cls._stdin_lines():
                cls.parse(text=line, pretty=False, suppress_errors=suppress_errors, pos=pos)
            return

        try:
            tree = formula_parser.parse(text)
        except exc.ParseError:
            if suppress_errors:
                print(ERROR_TOKEN)
                return
            else:
                raise

        if pos is not None:
            tree = tree.get_by_pos(pos)  # type: ignore  # TODO: fix

        if pretty:
            text = tree.pretty(with_meta=with_meta)
        else:
            text = tree.stringify(with_meta=with_meta)
        print(text)

    @classmethod
    def split(cls, text: str, diff: bool):
        formula = formula_parser.parse(text)

        chars = []  # type: ignore  # TODO: fix

        def _mark_positions(node, level=0):
            if len(chars) == level:
                chars.append([" " for _ in range(len(text))])
            if node.pos_range != (None, None):
                for i in range(*node.pos_range):
                    chars[level][i] = text[i]
                    if diff and level > 0:
                        chars[level - 1][i] = "▫"

            for child in node.children:
                _mark_positions(child, level + 1)

        _mark_positions(formula)
        for num, line in enumerate(chars):
            print("{0:>2}: {1}".format(num, "".join(line)))

    @staticmethod
    def graph(text: str, render_to: str, view: bool):
        formula = formula_parser.parse(text)
        dot = dl_formula.dot.translate(formula)
        make_graphviz_graph(dot, render_to=render_to, view=view)

    @classmethod
    def translate(
        cls, text: str, tablename: Optional[str], dialect: DialectCombo, unknown_funcs: bool, suppress_errors: bool
    ):
        if not text:
            # bulk mode, read from STDIN
            for line in cls._stdin_lines():
                cls.translate(
                    text=line,
                    tablename=tablename,
                    dialect=dialect,
                    unknown_funcs=unknown_funcs,
                    suppress_errors=suppress_errors,
                )
            return

        try:
            formula = formula_parser.parse(text)

            field_list = sorted(set(f.name for f in dl_formula.inspect.expression.used_fields(formula)))
            field_types = {f: DataType.NULL for f in field_list}
            if tablename:
                field_names = {f: (tablename, f) for f in field_list}
            else:
                field_names = {f: (f,) for f in field_list}  # type: ignore  # TODO: fix

            raw_sql = dl_formula.translation.translator.translate_and_compile(
                formula,
                dialect=dialect,
                restrict_funcs=not unknown_funcs,
                field_types=field_types,
                field_names=field_names,  # type: ignore  # TODO: fix
            )
            print(raw_sql)

        except (exc.ParseError, exc.TranslationError):
            if suppress_errors:
                print(ERROR_TOKEN)
                return
            else:
                raise

    @staticmethod
    def list_dialects():
        print(
            "\n".join(
                [
                    d.common_name_and_version
                    for d in sorted(get_all_basic_dialects(), key=lambda el: el.single_bit.orderable)
                ]
            )
        )

    @staticmethod
    def _get_func_base_class(name: str) -> Optional[type]:  # type: ignore  # TODO: fix
        """Find the first (base) class for the given function"""
        name = name.lower()
        for _i, definition in OPERATION_REGISTRY.items():
            if definition.name.lower() == name:  # type: ignore  # TODO: fix
                item = definition
                break
        else:
            return None

        mro = inspect.getmro(type(item))  # `cls` itself plus all its bases
        for earliest_super_cls in reversed(mro):
            if (
                issubclass(earliest_super_cls, MultiVariantTranslation)
                and (earliest_super_cls.name or "").lower() == name
            ):
                return earliest_super_cls

    @classmethod
    def _get_func_source_info(cls, name: str) -> Tuple[str, int]:
        """Get file name and line number for given function"""
        func_cls = cls._get_func_base_class(name)
        filename = inspect.getsourcefile(func_cls)  # type: ignore  # TODO: fix
        lineno = inspect.getsourcelines(func_cls)[1]  # type: ignore  # TODO: fix
        return filename, lineno  # type: ignore  # TODO: fix

    @classmethod
    def goto_func(cls, name: str, show: bool):
        editor_exe = os.environ.get("PYCHARM", "pycharm")
        filename, lineno = cls._get_func_source_info(name)
        if show:
            print(filename, lineno)
        else:
            os.system('{0} --line {2} "{1}"'.format(editor_exe, filename, lineno))

    @classmethod
    def slice(cls, text: str, levels: List[str], diff: bool = False) -> None:
        boundaries = {
            "aggregate": AggregateFunctionLevelBoundary(name="aggregate"),
            "window": WindowFunctionLevelBoundary(name="window"),
            "nonfields": NonFieldsBoundary(name="nonfields"),
            "toplevel": TopLevelBoundary(name="toplevel"),
        }

        formula = formula_parser.parse(text)
        slicer = FormulaSlicer(slice_schema=SliceSchema(levels=[boundaries[level_name] for level_name in levels]))
        sliced_formula = slicer.slice_formula(formula=formula)

        if diff:
            char_table = []  # type: ignore  # TODO: fix
            for i in reversed(range(len(levels))):
                char_line = [" " for _ in range(len(text))]
                for alias, node in sorted(sliced_formula.slices[i].aliased_nodes.items()):
                    char_line[node.start_pos : node.end_pos] = list(text[node.start_pos : node.end_pos])
                    if char_table:
                        for j in range(node.start_pos, node.end_pos):  # type: ignore  # TODO: fix
                            char_table[0][j] = "▫"
                char_table.insert(0, char_line)

            for i, char_line in reversed(list(enumerate(char_table))):
                print("{0:<15}{1}".format(levels[i], "".join(char_line)))
            print()

        else:
            for slice in reversed(sliced_formula.slices):
                print(f"------- {slice.name}")
                for alias, node in sorted(slice.aliased_nodes.items()):
                    print(f"[{alias}]:  {node}")
                print()

    @classmethod
    def print_registry(cls) -> None:
        items_by_module: dict[str, list[NodeTranslation]] = defaultdict(list)
        for _key, item in OPERATION_REGISTRY.items():
            filename = inspect.getsourcefile(type(item))
            assert filename is not None
            items_by_module[filename].append(item)

        prev_name: Optional[str]
        for filename, item_list in sorted(items_by_module.items()):
            print(f"{filename}:")
            prev_name = None
            for item in item_list:
                if item.name != prev_name:
                    prev_name = item.name
                    print(f"    # {item.name}")
                print(f"    {type(item).__name__},")

    @classmethod
    def run(cls, args):
        tool = cls()

        if args.command == "parse":
            tool.parse(
                text=args.optional_text,
                pretty=args.pretty,
                suppress_errors=args.suppress_errors,
                pos=args.pos,
                with_meta=args.with_meta,
            )

        elif args.command == "split":
            tool.split(text=args.text, diff=args.diff)

        elif args.command == "graph":
            tool.graph(text=args.text, render_to=args.render_to, view=args.view)

        elif args.command == "translate":
            tool.translate(
                text=args.optional_text,
                tablename=args.table,
                dialect=args.dialect,
                unknown_funcs=args.unknown_funcs,
                suppress_errors=args.suppress_errors,
            )

        elif args.command == "dialects":
            tool.list_dialects()

        elif args.command == "goto":
            tool.goto_func(name=args.func, show=args.show)

        elif args.command == "slice":
            tool.slice(text=args.text, levels=args.levels, diff=args.diff)

        elif args.command == "registry":
            tool.print_registry()


def main():
    load_formula_lib()
    FormulaCliTool.run(parser.parse_args())


if __name__ == "__main__":
    main()

lib/dl_formula/dl_formula/scripts/formula_cli.py
datalens-tech/datalens-backend
from __future__ import annotations

import builtins
import datetime as datetime_mod
from functools import wraps
from typing import (
    FrozenSet,
    List,
    Optional,
    Sequence,
    Tuple,
    Union,
)
import uuid as uuid_mod

from sqlalchemy.sql.elements import ClauseElement

from dl_formula.core import (
    fork_nodes,
    nodes,
)
from dl_formula.core.tag import LevelTag
from dl_formula.translation import ext_nodes
from dl_formula.translation.context import TranslationCtx


def _require_type(expected_type: Union[type, Tuple[Union[type, nodes.Array], ...], nodes.Array]):
    def decorator(func):
        @wraps(func)
        def wrapper(self, value):
            nodes.check_type(value, expected_type=expected_type)
            return func(self, value)

        return wrapper

    return decorator


def _norm(value) -> nodes.FormulaItem:
    if not isinstance(value, nodes.FormulaItem):
        value = n.lit(value)
    return value


class NodeShortcut:
    @_require_type(
        (
            builtins.str,
            builtins.int,
            builtins.float,
            builtins.bool,
            datetime_mod.date,
            datetime_mod.datetime,
            uuid_mod.UUID,
            nodes.Array(int),
            nodes.Array(float),
            nodes.Array(str),
            ClauseElement,
            TranslationCtx,
            type(None),
        )
    )
    def lit(self, value) -> Union[nodes.BaseLiteral, nodes.Null]:
        if isinstance(value, str):
            return self.str(value)
        if isinstance(value, bool):
            return self.bool(value)
        if isinstance(value, int):
            return self.int(value)
        if isinstance(value, float):
            return self.float(value)
        if isinstance(value, datetime_mod.datetime):
            return self.genericdatetime(value)
        if isinstance(value, datetime_mod.date):
            return self.date(value)
        if isinstance(value, uuid_mod.UUID):
            return self.uuid(value)
        if nodes.does_type_match(value, nodes.Array(int)):
            return nodes.LiteralArrayInteger.make(value)
        if nodes.does_type_match(value, nodes.Array(float)):
            return nodes.LiteralArrayFloat.make(value)
        if nodes.does_type_match(value, nodes.Array(str)):
            return nodes.LiteralArrayString.make(value)
        if isinstance(value, (ClauseElement, TranslationCtx)):
            return ext_nodes.CompiledExpression.make(value)
        if value is None:
            return self.null()
        raise TypeError(type(value))

    @_require_type(builtins.int)
    def int(self, value: builtins.int) -> nodes.LiteralInteger:
        return nodes.LiteralInteger.make(value)

    @_require_type(builtins.float)
    def float(self, value: builtins.float) -> nodes.LiteralFloat:
        return nodes.LiteralFloat.make(value)

    @_require_type(builtins.bool)
    def bool(self, value: builtins.bool) -> nodes.LiteralBoolean:
        return nodes.LiteralBoolean.make(value)

    @_require_type(datetime_mod.datetime)
    def datetime(self, value: datetime_mod.datetime) -> nodes.LiteralDatetime:
        return nodes.LiteralDatetime.make(value)

    @_require_type(datetime_mod.datetime)
    def genericdatetime(self, value: datetime_mod.datetime) -> nodes.LiteralGenericDatetime:
        return nodes.LiteralGenericDatetime.make(value)

    @_require_type(datetime_mod.date)
    def date(self, value: datetime_mod.date) -> nodes.LiteralDate:
        return nodes.LiteralDate.make(value)

    @_require_type(builtins.str)
    def geopt(self, value: builtins.str) -> nodes.LiteralGeopoint:
        return nodes.LiteralGeopoint.make(value)

    @_require_type(builtins.str)
    def geopoly(self, value: builtins.str) -> nodes.LiteralGeopolygon:
        return nodes.LiteralGeopolygon.make(value)

    @_require_type((builtins.str, uuid_mod.UUID))
    def uuid(self, value: Union[builtins.str, uuid_mod.UUID]) -> nodes.LiteralUuid:
        return nodes.LiteralUuid.make(value)

    @_require_type(builtins.str)
    def str(self, value: builtins.str) -> nodes.LiteralString:
        return nodes.LiteralString.make(value)

    @_require_type((ClauseElement, TranslationCtx))
    def expr(self, value: Union[ClauseElement, TranslationCtx]) -> ext_nodes.CompiledExpression:
        return ext_nodes.CompiledExpression.make(value)

    def null(self) -> nodes.Null:
        return nodes.Null()

    def total(self) -> nodes.WindowGroupingTotal:
        return nodes.WindowGroupingTotal()

    def within(self, *dims: nodes.FormulaItem) -> nodes.WindowGroupingWithin:
        return nodes.WindowGroupingWithin.make(dim_list=list(dims))

    def among(self, *dims: nodes.FormulaItem) -> nodes.WindowGroupingAmong:
        return nodes.WindowGroupingAmong.make(dim_list=list(dims))

    def fixed(self, *dims: nodes.FormulaItem) -> nodes.FixedLodSpecifier:
        return nodes.FixedLodSpecifier.make(dim_list=list(dims))

    def inherited(self) -> nodes.InheritedLodSpecifier:
        return nodes.InheritedLodSpecifier()

    def include(self, *dims: nodes.FormulaItem) -> nodes.IncludeLodSpecifier:
        return nodes.IncludeLodSpecifier.make(dim_list=list(dims))

    def exclude(self, *dims: nodes.FormulaItem) -> nodes.ExcludeLodSpecifier:
        return nodes.ExcludeLodSpecifier.make(dim_list=list(dims))

    def asc(self, expr: nodes.FormulaItem) -> nodes.OrderAscending:
        return nodes.OrderAscending.make(expr=expr)

    def desc(self, expr: nodes.FormulaItem) -> nodes.OrderDescending:
        return nodes.OrderDescending.make(expr=expr)

    def _normalize_raw_bfb(
        self,
        before_filter_by: Optional[List[Union[nodes.Field, builtins.str]]] = None,
    ) -> nodes.BeforeFilterBy:
        return nodes.BeforeFilterBy.make(
            field_names=frozenset([f.name if isinstance(f, nodes.Field) else f for f in (before_filter_by or ())]),
        )

    class FuncShortcut:
        """SQLAlchemy-like function shortcut provider: ````"""

        def __init__(self, assume_window: bool = False):
            self._assume_window = assume_window

        def __getattr__(self, name: str):
            return self._make_func(name)

        def __call__(self, name: str):
            return self._make_func(name)

        def _make_func(self, name: str):
            def wrapper(
                *pos_args: nodes.FormulaItem,
                grouping: Optional[nodes.WindowGrouping] = None,
                total: Optional[bool] = None,
                within: Optional[List[nodes.FormulaItem]] = None,
                among: Optional[List[nodes.FormulaItem]] = None,
                order_by: Optional[List[nodes.FormulaItem]] = None,
                before_filter_by: Optional[List[Union[nodes.Field, str]]] = None,
                ignore_dimensions: Optional[nodes.IgnoreDimensions] = None,
                lod: Optional[nodes.LodSpecifier] = None,
                args: Sequence[nodes.FormulaItem] = (),
                meta: Optional[nodes.NodeMeta] = None,
            ) -> Union[nodes.FuncCall, nodes.WindowFuncCall]:
                assert (
                    len([bool(spec) for spec in (grouping, total, within, among) if spec is not None]) <= 1
                ), "Only one grouping specifier is allowed"

                if total is not None:
                    assert total is True, "Invalid value for total"
                    grouping = n.total()
                elif within is not None:
                    grouping = n.within(*within)
                elif among is not None:
                    grouping = n.among(*among)

                if grouping is None and self._assume_window:
                    grouping = n.total()

                before_filter_by_node = n._normalize_raw_bfb(before_filter_by=before_filter_by)

                assert not (pos_args and args), "Can't specify pos args and `args` keyword simultaneously"
                args = pos_args or args or ()

                if grouping is None and order_by is None:
                    return nodes.FuncCall.make(
                        name=name.lower(),
                        args=[_norm(arg) for arg in args],
                        lod=lod,
                        ignore_dimensions=ignore_dimensions,
                        before_filter_by=before_filter_by_node,
                        meta=meta,
                    )

                return nodes.WindowFuncCall.make(
                    name=name.lower(),
                    args=[_norm(arg) for arg in args],
                    grouping=grouping,
                    ordering=nodes.Ordering.make(expr_list=order_by or []),
                    before_filter_by=before_filter_by_node,
                    meta=meta,
                )

            return wrapper

    func = FuncShortcut()  # n.func.SOME_FUNCTION(*args)
    wfunc = FuncShortcut(assume_window=True)  # same, but always results in a window function

    class WhenProxy:
        def __init__(self, val):
            self._val = val

        def then(self, expr) -> nodes.WhenPart:
            return nodes.WhenPart.make(val=_norm(self._val), expr=_norm(expr))

    def when(self, val) -> "NodeShortcut.WhenProxy":
        return self.WhenProxy(val)

    class CaseProxy:
        def __init__(self, case_expr):
            self._case_expr = case_expr
            self._when_list = []

        def whens(self, *when_list) -> "NodeShortcut.CaseProxy":
            self._when_list += when_list or []
            return self

        def else_(self, expr) -> nodes.CaseBlock:
            return nodes.CaseBlock.make(
                case_expr=_norm(self._case_expr), when_list=self._when_list, else_expr=_norm(expr)
            )

    def case(self, expr) -> "NodeShortcut.CaseProxy":
        """
        Usage:
        TODO
        """
        return self.CaseProxy(expr)

    class IfPartProxy:
        def __init__(self, cond):
            self._cond = cond
            self._expr = None

        def then(self, expr) -> "NodeShortcut.IfPartProxy":
            self._expr = expr
            return self

        def else_(self, expr) -> nodes.IfBlock:
            return n.IfBlockProxy([self]).else_(expr)

    def elseif(self, cond) -> "NodeShortcut.IfPartProxy":
        return self.IfPartProxy(cond)

    class IfBlockProxy:
        def __init__(self, if_list: list):
            self._if_list = if_list

        def else_(self, expr) -> nodes.IfBlock:
            return nodes.IfBlock.make(
                if_list=[nodes.IfPart.make(cond=_norm(ipp._cond), expr=_norm(ipp._expr)) for ipp in self._if_list],
                else_expr=_norm(expr),
            )

    def if_(self, *exprs) -> Union["NodeShortcut.IfBlockProxy", "NodeShortcut.IfPartProxy"]:
        """
        Usage:
        TODO
        """
        if len(exprs) > 1:
            return self.IfBlockProxy(list(exprs))
        if isinstance(exprs[0], self.IfPartProxy):
            return exprs[0]
        return self.IfPartProxy(exprs[0])

    def field(self, name: builtins.str) -> nodes.Field:
        return nodes.Field.make(name=name)

    def formula(self, expr: nodes.FormulaItem) -> nodes.Formula:
        return nodes.Formula.make(expr=_norm(expr))

    def p(self, expr: nodes.FormulaItem) -> nodes.ParenthesizedExpr:
        return nodes.ParenthesizedExpr.make(expr=_norm(expr))

    def tagged(self, tag: LevelTag, expr: nodes.FormulaItem) -> nodes.FormulaItem:
        return expr.with_tag(level_tag=tag)

    def level_tag(
        self,
        names: FrozenSet[builtins.str],
        nesting: builtins.int,
        qfork_nesting: builtins.int = 0,
    ) -> LevelTag:
        return LevelTag(bfb_names=names, func_nesting=nesting, qfork_nesting=qfork_nesting)

    def fork(
        self,
        joining: fork_nodes.QueryForkJoiningBase,
        result_expr: nodes.FormulaItem,
        before_filter_by: Optional[List[Union[nodes.Field, builtins.str]]] = None,
        lod: Optional[nodes.FixedLodSpecifier] = None,
        join_type: fork_nodes.JoinType = fork_nodes.JoinType.left,
    ) -> fork_nodes.QueryFork:
        before_filter_by_node = self._normalize_raw_bfb(before_filter_by=before_filter_by)
        return fork_nodes.QueryFork.make(
            join_type=join_type,
            joining=joining,
            result_expr=result_expr,
            before_filter_by=before_filter_by_node,
            lod=lod,
        )

    def bin_condition(self, expr: nodes.FormulaItem, fork_expr: nodes.FormulaItem) -> fork_nodes.BinaryJoinCondition:
        return fork_nodes.BinaryJoinCondition.make(expr=expr, fork_expr=fork_expr)

    def self_condition(self, expr: nodes.FormulaItem) -> fork_nodes.SelfEqualityJoinCondition:
        return fork_nodes.SelfEqualityJoinCondition.make(expr=expr)

    def joining(
        self,
        *,
        conditions: Optional[Sequence[fork_nodes.JoinConditionBase]] = None,
    ) -> fork_nodes.QueryForkJoiningBase:
        if conditions is not None:
            return fork_nodes.QueryForkJoiningWithList.make(
                condition_list=list(conditions),
            )
        raise ValueError("Invalid options for joining object")

    def unary(
        self,
        name: builtins.str,
        expr: nodes.FormulaItem,
        meta: Optional[nodes.NodeMeta] = None,
    ) -> nodes.Unary:
        return nodes.Unary.make(name=name, expr=expr, meta=meta)

    def binary(
        self,
        name: builtins.str,
        left: nodes.FormulaItem,
        right: nodes.FormulaItem,
        meta: Optional[nodes.NodeMeta] = None,
    ) -> nodes.Binary:
        return nodes.Binary.make(name=name, left=left, right=right, meta=meta)

    def ternary(
        self,
        name: builtins.str,
        first: nodes.FormulaItem,
        second: nodes.FormulaItem,
        third: nodes.FormulaItem,
        meta: Optional[nodes.NodeMeta] = None,
    ) -> nodes.Ternary:
        return nodes.Ternary.make(name=name, first=first, second=second, third=third, meta=meta)

    def not_(
        self,
        expr: nodes.FormulaItem,
        meta: Optional[nodes.NodeMeta] = None,
    ):
        return self.unary("not", expr, meta=meta)


n = NodeShortcut()

lib/dl_formula/dl_formula/shortcuts.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from typing import Dict

import attr

from dl_formula.collections import NodeValueMap
from dl_formula.inspect.env import InspectionEnvironment


@attr.s
class SlicerEnvironment:
    inspect_env: InspectionEnvironment = attr.ib(factory=InspectionEnvironment)
    # Autogenerated names for parts of the formula that are sliced off and replaced by fields
    # Sliced nodes are cached separately for each level identified by an int
    cache_sliced_node_name: Dict[int, NodeValueMap[str]] = attr.ib(factory=lambda: defaultdict(NodeValueMap))

lib/dl_formula/dl_formula/slicing/env.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from enum import Enum
from typing import (
    Callable,
    List,
    Optional,
    Tuple,
)
import uuid

import attr

import dl_formula.core.fork_nodes as fork_nodes
import dl_formula.core.nodes as nodes
from dl_formula.core.tag import LevelTag
from dl_formula.inspect.env import InspectionEnvironment
import dl_formula.inspect.expression
import dl_formula.inspect.function
import dl_formula.inspect.node


class BoundaryCheckResult(Enum):
    raise_level = "raise_level"
    neutral = "neutral"
    maintain_level = "maintain_level"


@attr.s
class LevelBoundary(abc.ABC):
    name: str = attr.ib()
    name_gen: Optional[Callable[[nodes.FormulaItem, str], str]] = attr.ib(default=None)

    def _default_name_gen(self, node: nodes.FormulaItem) -> str:
        """Default name generator for sliced nodes"""
        return str(uuid.uuid4())

    def generate_sliced_node_name(self, node: nodes.FormulaItem) -> str:
        if self.name_gen is not None:
            return self.name_gen(node, self.name)
        return self._default_name_gen(node)

    @abc.abstractmethod
    def check_node(
        self,
        node: nodes.FormulaItem,
        inspect_env: InspectionEnvironment,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> BoundaryCheckResult:
        raise NotImplementedError

    def should_raise_level(
        self,
        node: nodes.FormulaItem,
        inspect_env: InspectionEnvironment,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> bool:
        """Return ``True`` if the expression should be raised to the next execution level"""
        return (
            self.check_node(
                node,
                inspect_env=inspect_env,
                parent_stack=parent_stack,
            )
            == BoundaryCheckResult.raise_level
        )

    def should_maintain_level(
        self,
        node: nodes.FormulaItem,
        inspect_env: InspectionEnvironment,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> bool:
        """Return ``True`` if the expression should stay at the current execution level"""
        return (
            self.check_node(
                node,
                inspect_env=inspect_env,
                parent_stack=parent_stack,
            )
            == BoundaryCheckResult.maintain_level
        )


@attr.s
class AggregateFunctionLevelBoundary(LevelBoundary):
    """Matches aggregate functions"""

    constants_neutral: bool = attr.ib(default=True)

    def check_node(
        self,
        node: nodes.FormulaItem,
        inspect_env: InspectionEnvironment,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> BoundaryCheckResult:
        if dl_formula.inspect.node.is_aggregate_function(node):
            return BoundaryCheckResult.raise_level

        if self.constants_neutral:
            if dl_formula.inspect.expression.is_constant_expression(node, env=inspect_env):
                return BoundaryCheckResult.neutral

        return BoundaryCheckResult.maintain_level


@attr.s
class WindowFunctionLevelBoundary(LevelBoundary):
    """Matches window functions"""

    constants_neutral: bool = attr.ib(default=True)

    def check_node(
        self,
        node: nodes.FormulaItem,
        inspect_env: InspectionEnvironment,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> BoundaryCheckResult:
        if isinstance(node, nodes.WindowFuncCall):
            return BoundaryCheckResult.raise_level

        if self.constants_neutral:
            if dl_formula.inspect.expression.is_constant_expression(node, env=inspect_env):
                return BoundaryCheckResult.neutral

        return BoundaryCheckResult.maintain_level


class TopLevelBoundary(LevelBoundary):
    """Never matches any node"""

    def check_node(
        self,
        node: nodes.FormulaItem,
        inspect_env: InspectionEnvironment,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> BoundaryCheckResult:
        return BoundaryCheckResult.maintain_level


class NonFieldsBoundary(LevelBoundary):
    """Raises all non-simple-field nodes to the next level"""

    def check_node(
        self,
        node: nodes.FormulaItem,
        inspect_env: InspectionEnvironment,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> BoundaryCheckResult:
        if isinstance(node, nodes.Field):
            return BoundaryCheckResult.maintain_level
        return BoundaryCheckResult.raise_level


@attr.s
class NestedLevelTaggedBoundary(LevelBoundary):
    """
    Operates over level tags.
    Tags, whose ``bfb_names`` set part is a subset of ``self.tag.bfb_names`` are raised.
    if the ``bfb_names`` sets are equal, then lower or equal values
    of ``tag.nesting`` result in raising
    The behavior for constants is inherited from ``TaggedBoundary``
    All other nodes maintain their level.

    Example:
        self.tag = ({A, B}, 1)

        Hierarchy:
            T({A}, 0)
            T({A}, 1)
            T({A, B}, 0)
            T({A, B}, 1)      <-- The part below this tag is sliced off
            T({A, B}, 2)
            T({A, B, C}, 0)
    """

    tag: Optional[LevelTag] = attr.ib(default=None)
    constants_neutral: bool = attr.ib(default=True)

    def check_tag(self, is_own_tag: bool, level_tag: LevelTag) -> BoundaryCheckResult:
        if self.tag is None or (level_tag < self.tag or is_own_tag and level_tag <= self.tag):
            # For non-own tags use strict <
            return BoundaryCheckResult.raise_level
        return BoundaryCheckResult.maintain_level

    def get_node_tag(
        self,
        node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> Tuple[bool, Optional[LevelTag]]:
        # For QueryForks the node itself is tagged
        if isinstance(node, fork_nodes.QueryFork):
            return True, node.level_tag

        parent = parent_stack[-1]

        # For aggregations functions, the wrapping QueryFork contains the tag
        if (
            dl_formula.inspect.node.is_aggregate_function(node)
            and isinstance(parent, fork_nodes.QueryFork)
            and dl_formula.inspect.node.qfork_is_aggregation(parent)  # only LODs, avoid lookups here
        ):
            return False, parent.level_tag

        # For window functions, the wrapping parentheses contain the tag
        if (
            isinstance(node, nodes.WindowFuncCall)
            and dl_formula.inspect.function.supports_bfb(node.name, is_window=True)
            and isinstance(parent, nodes.ParenthesizedExpr)
        ):
            return False, parent.level_tag

        # Tagged parentheses go hand in hand with window functions
        if isinstance(node, nodes.ParenthesizedExpr):
            return True, node.level_tag

        return True, None

    def check_node(
        self,
        node: nodes.FormulaItem,
        inspect_env: InspectionEnvironment,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> BoundaryCheckResult:
        is_own_tag, level_tag = self.get_node_tag(node=node, parent_stack=parent_stack)
        if level_tag is not None:
            return self.check_tag(is_own_tag=is_own_tag, level_tag=level_tag)

        if self.constants_neutral:
            if dl_formula.inspect.expression.is_constant_expression(node, env=inspect_env):
                return BoundaryCheckResult.neutral

        return BoundaryCheckResult.maintain_level


@attr.s
class SliceSchema:
    levels: List[LevelBoundary] = attr.ib(factory=list)

    @property
    def max_level(self) -> int:
        return len(self.levels) - 1

lib/dl_formula/dl_formula/slicing/schema.py
datalens-tech/datalens-backend
"""
Formula slicing.

A formula (``Formula`` class instance) represents a tree of nodes (instances of ``FormulaItem`` subclasses).
Complex processing of formulas may require them to be split into several parts that will be processed separately
by different parts of the code, in different environments, etc..

``FormulaSlicer`` implements this procedure (slicing) via recursing the formula from the leaves, up to the root
and checking whether the current node still fits the current slice, or has to be moved up to the next one.

``SliceSchema`` defines the sequence of ``LevelBoundary`` subclass instances that do the node checking.
If a formula has exhausted all of its nodes (root has been reached), but there are still more slices to go,
then it is "pulled" through these slices as a field node, at each level representing its content

from __future__ import annotations

from the previous level.

Example:

Say, we want to split ``RANK(SUM([Sales]), MIN([Order Date])) - 1`` and window function level (RANK)
so that the part without window functions can be processed in a database that does not support them,
and the rest somewhere that allows their usage:

```
original:       RANK(SUM([Sales]), MIN([Order Date])) - 1
levels:
    toplevel    RANK(▫▫▫▫▫▫▫▫▫▫▫▫, ▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫) - 1
    window           SUM([Sales])  MIN([Order Date])
```

Here, after slicing we still have one formula at the top level, while the window function arguments have been
cut off as two separate sub-formulas ``SUM([Sales])`` and ``MIN([Order Date])`` at the lower level.
The top-level formula that has the ``RANK`` function call references them as fields with autogenerated names.

``FormulaSlicer`` returns the sliced formula as a ``SlicedFormulaInfo`` object, that stores all the separate parts
of the formula on a per-slice basis with info about the autogenerated fields at each level.
"""


import logging
from typing import (
    Dict,
    Iterable,
    List,
    Optional,
    Set,
    Tuple,
    TypeVar,
)

import attr

from dl_formula.core.index import NodeHierarchyIndex
import dl_formula.core.nodes as nodes
from dl_formula.inspect.expression import enumerate_autonomous_children
from dl_formula.slicing.env import SlicerEnvironment
from dl_formula.slicing.schema import (
    LevelBoundary,
    SliceSchema,
)


LOGGER = logging.getLogger(__name__)


@attr.s
class LevelSlice:
    name: str = attr.ib()
    aliased_nodes: Dict[str, nodes.Formula] = attr.ib(factory=dict)
    required_fields: Set[str] = attr.ib(factory=set)

    def is_lazy(self) -> bool:
        """
        Check whether this slice is lazy
        (every formula consists of a single child  that is a field referencing the previous slice).
        """

        for _alias, node in self.aliased_nodes.items():
            if not isinstance(node.expr, nodes.Field):
                # there's something more going on
                return False

        return True


@attr.s
class SlicedFormulaInfo:
    slices: List[LevelSlice] = attr.ib(factory=list)


_FORMULA_ITEM_TV = TypeVar("_FORMULA_ITEM_TV", bound=nodes.FormulaItem)


@attr.s
class FormulaSlicer:
    slice_schema: SliceSchema = attr.ib(factory=SliceSchema)
    env: SlicerEnvironment = attr.ib(factory=SlicerEnvironment)

    @staticmethod
    def _get_child_expression_nodes(node: nodes.FormulaItem) -> Iterable[nodes.FormulaItem]:
        if isinstance(node, nodes.WindowFuncCall):
            return node.args  # to avoid iterating over grouping node
        return node.children

    def _boundary_for_level(self, level: int) -> LevelBoundary:
        return self.slice_schema.levels[level]

    def _get_node_alias(self, node: nodes.FormulaItem, level: int) -> str:
        level_cache = self.env.cache_sliced_node_name[level]

        # Unwrap Formula objects
        if isinstance(node, nodes.Formula):
            node = node.expr

        alias = level_cache.get(node)
        if alias is None:
            alias = self._boundary_for_level(level).generate_sliced_node_name(node)
            level_cache.add(node, value=alias)

        return alias

    def _slice_node(
        self,
        sliced_formula_info: SlicedFormulaInfo,
        parent_node: _FORMULA_ITEM_TV,
        child_node: nodes.FormulaItem,
        child_node_idx: NodeHierarchyIndex,
        level: int,
    ) -> Tuple[_FORMULA_ITEM_TV, nodes.Field]:
        # generate replacement field
        alias = self._get_node_alias(child_node, level=level)
        replacement_field = nodes.Field.make(
            name=alias,
            meta=nodes.NodeMeta(position=child_node.position),
        )

        # register in result object
        sliced_formula_info.slices[level].aliased_nodes[alias] = nodes.Formula.make(
            expr=child_node,
            meta=nodes.NodeMeta(position=child_node.position),
        )
        if level < len(sliced_formula_info.slices) - 1:
            sliced_formula_info.slices[level + 1].required_fields.add(alias)

        # replace in parent
        parent_node = parent_node.replace_at_index(index=child_node_idx, expr=replacement_field)
        return parent_node, replacement_field

    def _pull_to_level_and_replace(
        self,
        sliced_formula_info: SlicedFormulaInfo,
        parent_node: _FORMULA_ITEM_TV,
        child_node: nodes.FormulaItem,
        child_node_idx: NodeHierarchyIndex,
        parent_stack: Tuple[nodes.FormulaItem, ...],
        from_level: int,
        to_level: int,
    ) -> _FORMULA_ITEM_TV:
        original_child_node = child_node
        while from_level < to_level:
            if self._boundary_for_level(from_level).should_maintain_level(
                node=original_child_node, inspect_env=self.env.inspect_env, parent_stack=parent_stack
            ):
                # expression must stay at current execution level
                # so it must be sliced off and represented as a field on higher levels
                parent_node, child_node = self._slice_node(
                    sliced_formula_info=sliced_formula_info,
                    parent_node=parent_node,
                    child_node=child_node,
                    child_node_idx=child_node_idx,
                    level=from_level,
                )
            from_level += 1

        return parent_node

    def _slice_node_recursively(
        self,
        sliced_formula_info: SlicedFormulaInfo,
        cur_node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> Tuple[int, nodes.FormulaItem]:
        original_cur_node = cur_node
        parent_stack_for_children = parent_stack + (cur_node,)  # FIXME: replace with stack from enumerate?

        # first calculate the level of children
        children_and_levels: List[Tuple[NodeHierarchyIndex, nodes.FormulaItem, int]] = []
        for child_node_idx, child_node, _ in enumerate_autonomous_children(cur_node, parent_stack=parent_stack):
            child_level, sliced_child_node = self._slice_node_recursively(
                sliced_formula_info=sliced_formula_info,
                cur_node=child_node,
                parent_stack=parent_stack_for_children,
            )
            if sliced_child_node is not child_node:
                cur_node = cur_node.replace_at_index(index=child_node_idx, expr=sliced_child_node)

            children_and_levels.append((child_node_idx, sliced_child_node, child_level))

        cur_level = max(child_level for _, _, child_level in children_and_levels) if children_and_levels else 0

        # now check if we've reached the level boundary
        while self._boundary_for_level(cur_level).should_raise_level(
            node=original_cur_node,
            inspect_env=self.env.inspect_env,
            parent_stack=parent_stack,
        ):
            # current expression must be moved up one level
            cur_level += 1
            # its children will be either sliced off (if they are incompatible with this next level)
            # or also moved up if they are neutral with respect to the level boundary
            # (the result of `boundary.check_level(child)`

        # slice off all the children that belong to one of the previous levels
        sliced_cur_node = cur_node
        for child_node_idx, sliced_child_node, child_level in children_and_levels:
            sliced_cur_node = self._pull_to_level_and_replace(
                sliced_formula_info=sliced_formula_info,
                parent_node=sliced_cur_node,
                child_node=sliced_child_node,
                child_node_idx=child_node_idx,
                parent_stack=parent_stack_for_children,
                from_level=child_level,
                to_level=cur_level,
            )

        return cur_level, sliced_cur_node

    def slice_formula(
        self,
        formula: nodes.Formula,
        alias: Optional[str] = None,
    ) -> SlicedFormulaInfo:
        """
        Slice node expression into parts corresponding to each slice level.
        """

        sliced_formula_info = SlicedFormulaInfo(
            slices=[LevelSlice(name=boundary.name) for boundary in self.slice_schema.levels],
        )

        formula_expr_idx = NodeHierarchyIndex.make(0)
        formula_expr = formula[formula_expr_idx]
        # Just a formality to show and make sure that we have the correct `formula_expr_idx`
        assert formula_expr is formula.expr

        # slice it into parts
        formula_expr_parent_stack = (formula,)
        level_after_slicing, sliced_expr = self._slice_node_recursively(
            sliced_formula_info=sliced_formula_info,
            cur_node=formula_expr,
            parent_stack=formula_expr_parent_stack,
        )
        sliced_formula = formula.replace_at_index(index=formula_expr_idx, expr=sliced_expr)

        # stretch it up as a field to the top level
        sliced_formula = self._pull_to_level_and_replace(
            sliced_formula_info=sliced_formula_info,
            parent_node=sliced_formula,
            child_node=sliced_expr,
            child_node_idx=formula_expr_idx,
            from_level=level_after_slicing,
            parent_stack=formula_expr_parent_stack,
            to_level=self.slice_schema.max_level,
        )

        # register formula in last level under an alias
        level = self.slice_schema.max_level
        alias = alias or self._get_node_alias(sliced_formula, level=level)
        sliced_formula_info.slices[level].aliased_nodes[alias] = sliced_formula
        return sliced_formula_info

lib/dl_formula/dl_formula/slicing/slicer.py
datalens-tech/datalens-backend
from typing import Type

from dl_formula.connectors.base.column import (
    ColumnRenderer,
    DefaultColumnRenderer,
)
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D


_COLUMN_RENDERER_TYPES: dict[DialectCombo, Type[ColumnRenderer]] = {
    D.DUMMY: DefaultColumnRenderer,
}


def get_column_renderer(dialect: DialectCombo, field_names: dict[str, tuple[str, ...]]) -> ColumnRenderer:
    sorted_dialect_items = sorted(
        _COLUMN_RENDERER_TYPES.items(), key=lambda el: el[0].ambiguity
    )  # Most specific dialects go first, the most ambiguous ones go last
    for d, context_processor in sorted_dialect_items:
        if d & dialect == dialect:
            return context_processor(dialect=dialect, field_names=field_names)
    raise ValueError(f"Couldn't find a column renderer for dialect {dialect}")


def register_column_renderer_cls(dialect: DialectCombo, column_renderer_cls: Type[ColumnRenderer]) -> None:
    if dialect in _COLUMN_RENDERER_TYPES:
        assert _COLUMN_RENDERER_TYPES[dialect] is column_renderer_cls
    else:
        _COLUMN_RENDERER_TYPES[dialect] = column_renderer_cls

lib/dl_formula/dl_formula/translation/columns.py
datalens-tech/datalens-backend
from __future__ import annotations

import copy
from typing import (
    Any,
    Dict,
    List,
    Optional,
    Tuple,
)

from sqlalchemy.sql.elements import ClauseElement

from dl_formula.core.datatype import (
    DataType,
    DataTypeParams,
)
import dl_formula.core.exc as exc
from dl_formula.core.extract import NodeExtract
from dl_formula.core.message_ctx import (
    FormulaErrorCtx,
    MessageLevel,
)
from dl_formula.core.nodes import FormulaItem
from dl_formula.core.position import Position
from dl_formula.definitions.scope import Scope


ContextFlags = int  # bitwise combination of ContextFlag values


class TranslationCtx:
    __slots__ = (
        "collect_errors",
        "base_token",
        "children",
        "_messages",
        "forked",
        "data_type",
        "data_type_params",
        "expression",
        "flags",
        "node",
        "required_scopes",
    )
    _repr_attrs = (
        "base_token",
        "data_type",
        "data_type_params",
        "flags",
        "forked",
        "required_scopes",
        "node",
        "_expression_str",
        "children",
        "collect_errors",
        "_messages",
    )
    default_collect_errors = False

    def __init__(
        self,
        collect_errors: Optional[bool] = None,
        base_token: Optional[str] = None,
        flags: Optional[ContextFlags] = None,
        expression: Optional[ClauseElement] = None,
        data_type: Optional[DataType] = None,
        data_type_params: Optional[DataTypeParams] = None,
        node: Optional[FormulaItem] = None,
        required_scopes: int = Scope.EXPLICIT_USAGE,
    ):
        self.collect_errors = collect_errors if collect_errors is not None else self.default_collect_errors
        self.base_token = base_token
        self.children: List[TranslationCtx] = []
        self.forked = False  # TODO: think about removing this completely
        self.data_type = data_type
        self.data_type_params = data_type_params
        self.expression = expression
        self.flags = flags or 0
        self.node = node
        self.required_scopes: int = required_scopes

        self._messages: Dict[MessageLevel, List[FormulaErrorCtx]] = {
            MessageLevel.ERROR: [],
            MessageLevel.WARNING: [],
        }

    @property
    def _expression_str(self):
        try:
            return str(self.expression)
        except Exception:
            return repr(self.expression)

    def __repr__(self):
        _c0 = "\x1b[0m"
        _c1 = "\x1b[038;5;118m"
        return "<{}({})>{}".format(
            self.__class__.__name__,
            ",     ".join(f"{_c1}{key}{_c0}={getattr(self, key)!r}" for key in self._repr_attrs),
            _c0,
        )

    def copy(self) -> "TranslationCtx":
        return copy.copy(self)

    @property
    def extract(self) -> Optional[NodeExtract]:
        return self.node.extract if self.node is not None else None

    @property
    def errors(self) -> List[FormulaErrorCtx]:
        """Return list of registered errors."""
        return self._messages[MessageLevel.ERROR]

    @property
    def warnings(self) -> List[FormulaErrorCtx]:
        """Return list of registered warnings."""
        return self._messages[MessageLevel.WARNING]

    def flush(self) -> "TranslationCtx":
        """
        Flatten child contexts into self and validate:
        - aggregation consistency and
        - data type consistency (via type strategy)
        among children
        """

        for child in self.children:
            child.flush()

        for child in self.children:
            # propagate errors, warnings, etc.
            for key, messages in child._messages.items():
                self._messages[key].extend(messages)

        if not self.data_type:
            if self.forked:
                raise RuntimeError("Data type must be set for forked contexts")

            elif self.children:
                # propagate directly because there were no calculations here
                # (single child, no operator)
                self.data_type = self.children[0].data_type
                self.data_type_params = self.children[0].data_type_params
                self.node = self.children[0].node

            else:
                # some kind of exceptional case, no data type here, fall back to any type (NULL)
                self.data_type = DataType.NULL
                self.data_type_params = DataTypeParams()

        if not self.forked and self.children:  # single child, no operator
            if self.expression is None:
                # inherit the only child's expression
                self.expression = self.children[0].expression

            self.flags |= self.children[0].flags

        self.children.clear()
        self.forked = False
        return self

    def _add_message(
        self,
        level: MessageLevel,
        message: str,
        token: Optional[str] = None,
        code: Optional[Tuple[str, ...]] = None,
    ) -> None:
        token = token if token is not None else self.base_token
        position = self.node.position if self.node is not None else Position()
        error = exc.FormulaErrorCtx(
            message=message,
            level=level,
            token=token,
            position=position,
            code=tuple(code or ()),
        )

        if level is MessageLevel.ERROR and not self.collect_errors:
            raise exc.TranslationError(error)

        self._messages[level].append(error)

    def add_error(self, message: str, token: Optional[str] = None, code: Optional[Tuple[str, ...]] = None) -> None:
        return self._add_message(level=MessageLevel.ERROR, message=message, token=token, code=code)

    def add_warning(self, message: str, token: Optional[str] = None, code: Optional[Tuple[str, ...]] = None) -> None:
        return self._add_message(level=MessageLevel.WARNING, message=message, token=token, code=code)

    def set_type(self, data_type: DataType, data_type_params: Optional[DataTypeParams] = None) -> None:
        self.data_type = data_type
        self.data_type_params = data_type_params if data_type_params is not None else DataTypeParams()

    def set_expression(self, expression: Any) -> None:
        self.expression = expression

    def set_token(self, token: Optional[str]) -> None:
        self.base_token = token

    def set_flags(self, flags: ContextFlags) -> None:
        self.flags |= flags

    def fork(self) -> None:
        """Switch to muti-child mode. Should be called for function/operator nodes"""
        self.forked = True
        self.expression = None

    def _validate_can_add_new_child(self) -> None:
        if self.forked:
            pass
        else:
            if self.children:
                raise RuntimeError("Cannot have more than one child in non-forked mode")

    def child(
        self,
        token: Optional[str] = None,
        flags: Optional[int] = None,
        node: Optional[FormulaItem] = None,
        required_scopes: Optional[int] = None,
    ) -> "TranslationCtx":
        """
        Spawn new child context.
        In `forked` mode only one child is permitted.
        If multiple children are needed, i.e. this node is a function/operator,
        and children will be its arguments, then call `fork()` before spawning children.
        """

        self._validate_can_add_new_child()
        if not self.forked and node is not None:
            raise ValueError("Cannot assign node to child in non-forked mode")

        child = self.__class__(
            collect_errors=self.collect_errors,
            base_token=token if token is not None else self.base_token,
            flags=flags,
            node=node,
            required_scopes=self.required_scopes if required_scopes is None else required_scopes,
        )
        self.children.append(child)
        return child

    def adopt(self, child: "TranslationCtx") -> None:
        """
        "Adopt" a child context by adding it to children of self as if it were created by `self.fork()`.
        Can be used for substituting formula nodes for already translated expressions.
        """

        self._validate_can_add_new_child()
        self.children.append(child)

lib/dl_formula/dl_formula/translation/context.py
datalens-tech/datalens-backend
from functools import lru_cache
from typing import Type

from dl_formula.connectors.base.context_processor import ContextPostprocessor
from dl_formula.core.dialect import DialectCombo


_CONTEXT_PROCESSORS: dict[DialectCombo, ContextPostprocessor] = {}


@lru_cache
def get_context_processor(dialect: DialectCombo) -> ContextPostprocessor:
    sorted_dialect_items = sorted(
        _CONTEXT_PROCESSORS.items(), key=lambda el: el[0].ambiguity
    )  # Most specific dialects go first, the most ambiguous ones go last
    for d, context_processor in sorted_dialect_items:
        if d & dialect == dialect:
            return context_processor

    return ContextPostprocessor()


def register_context_processor(dialect: DialectCombo, context_processor_cls: Type[ContextPostprocessor]) -> None:
    if dialect in _CONTEXT_PROCESSORS:
        assert isinstance(_CONTEXT_PROCESSORS[dialect], context_processor_cls)
        return

    _CONTEXT_PROCESSORS[dialect] = context_processor_cls()

lib/dl_formula/dl_formula/translation/context_processing.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from typing import (
    Any,
    NamedTuple,
)

import attr

from dl_formula.collections import NodeValueMap
from dl_formula.core.datatype import (
    DataType,
    DataTypeParams,
)
from dl_formula.core.dialect import DialectCombo
import dl_formula.translation.context as ctx_mod  # noqa


class FunctionStatsSignature(NamedTuple):
    name: str
    arg_types: tuple[str, ...]
    dialect: str
    is_window: bool

    def as_dict(self) -> dict[str, Any]:
        return self._asdict()


@attr.s
class TranslationStats:
    # {(<function_name>, (<arg_type_names>), <dialect_name>, <is_window>): <usage_count>}
    function_usage_weights: dict[FunctionStatsSignature, int] = attr.ib(factory=lambda: defaultdict(lambda: 0))
    cache_hits: int = attr.ib(kw_only=True, default=0)

    @classmethod
    def combine(cls, *stat_instances: TranslationStats) -> TranslationStats:
        combined_func_usage_weights: dict[FunctionStatsSignature, int] = defaultdict(lambda: 0)
        combined_cache_hits = 0
        for stat_ins in stat_instances:
            combined_cache_hits += stat_ins.cache_hits
            for func_sign, func_value in stat_ins.function_usage_weights.items():
                combined_func_usage_weights[func_sign] += func_value

        return TranslationStats(
            function_usage_weights=combined_func_usage_weights,
            cache_hits=combined_cache_hits,
        )

    def __add__(self, other: TranslationStats) -> TranslationStats:
        assert isinstance(other, TranslationStats)
        return self.combine(self, other)

    def add_function_hit(self, func_signature: FunctionStatsSignature, value: int = 1) -> None:
        self.function_usage_weights[func_signature] += value

    def add_cache_hit(self, value: int = 1) -> None:
        self.cache_hits += value


@attr.s
class TranslationEnvironment:
    dialect: DialectCombo = attr.ib()  # a combination of dialect bits
    required_scopes: int = attr.ib()

    # TODO: combine these into `field_info = {field_id: {name: ..., type: ..., type_params: ...}, ...}`
    field_names: dict[str, tuple[str, ...]] = attr.ib(factory=dict)
    field_types: dict[str, DataType] = attr.ib(factory=dict)
    field_type_params: dict[str, DataTypeParams] = attr.ib(factory=dict)

    translation_cache: NodeValueMap["ctx_mod.TranslationCtx"] = attr.ib(factory=NodeValueMap)
    replacements: NodeValueMap["ctx_mod.TranslationCtx"] = attr.ib(factory=NodeValueMap)

    translation_stats: TranslationStats = attr.ib(kw_only=True, factory=TranslationStats)

lib/dl_formula/dl_formula/translation/env.py
datalens-tech/datalens-backend
"""
Extended formula nodes
"""

from __future__ import annotations

from typing import (
    Hashable,
    Optional,
)

from sqlalchemy.sql import ClauseElement

from dl_formula.core import nodes
from dl_formula.core.extract import NodeExtract
from dl_formula.translation.context import TranslationCtx


class CompiledExpression(nodes.BaseLiteral):
    """Represents a pre-compiled expression that should be substituted as-is"""

    @classmethod
    def make(
        cls,
        value: ClauseElement | TranslationCtx,
        *,
        meta: Optional[nodes.NodeMeta] = None,
    ) -> CompiledExpression:
        return super().make(value=value, meta=meta)

    @classmethod
    def validate_literal_value(cls, literal_value: Optional[Hashable]) -> None:
        assert isinstance(literal_value, (ClauseElement, TranslationCtx))

    def _make_extract(self) -> Optional[NodeExtract]:
        if isinstance(self.value, TranslationCtx):
            return self.value.extract

        return None

lib/dl_formula/dl_formula/translation/ext_nodes.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict

from sqlalchemy.engine.default import DefaultDialect

from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D


def get_sa_dialect(dialect: DialectCombo) -> DefaultDialect:
    return TO_SA_DIALECTS[dialect]


def register_sa_dialect(dialect: DialectCombo, sa_dialect: DefaultDialect) -> None:
    TO_SA_DIALECTS[dialect] = sa_dialect


TO_SA_DIALECTS: dict[DialectCombo, DefaultDialect] = defaultdict(DefaultDialect)
register_sa_dialect(D.DUMMY, DefaultDialect())

lib/dl_formula/dl_formula/translation/sa_dialects.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from functools import singledispatchmethod
from itertools import chain
from typing import (
    Optional,
    Union,
)

import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement
from sqlalchemy.sql.functions import Function

import dl_formula.core.aux_nodes as aux_nodes
from dl_formula.core.datatype import (
    DataType,
    DataTypeParams,
)
from dl_formula.core.dialect import DialectCombo
import dl_formula.core.exc as exc
import dl_formula.core.nodes as nodes
from dl_formula.definitions.common import desc
from dl_formula.definitions.literals import literal
from dl_formula.definitions.registry import OPERATION_REGISTRY
from dl_formula.definitions.scope import Scope
from dl_formula.definitions.type_strategy import FromArgs
import dl_formula.inspect.literal
from dl_formula.translation import ext_nodes
from dl_formula.translation.columns import get_column_renderer
from dl_formula.translation.context import TranslationCtx
from dl_formula.translation.context_processing import get_context_processor
from dl_formula.translation.env import (
    FunctionStatsSignature,
    TranslationEnvironment,
)
from dl_formula.translation.sa_dialects import get_sa_dialect


class SqlAlchemyTranslator:
    default_restrict_funcs = True
    default_restrict_fields = True

    def __init__(
        self,
        env: TranslationEnvironment,
        restrict_funcs: Optional[bool] = None,
        restrict_fields: Optional[bool] = None,
        collect_stats: bool = False,
    ):
        self._env = env
        self._restrict_funcs = restrict_funcs if restrict_funcs is not None else self.default_restrict_funcs
        self._restrict_fields = restrict_fields if restrict_fields is not None else self.default_restrict_fields
        self._collect_stats = collect_stats

        self._postprocessor = get_context_processor(dialect=self._env.dialect)
        self._column_renderer = get_column_renderer(dialect=self._env.dialect, field_names=self._env.field_names)

    def _generate_default_function(self, node: nodes.OperationCall, translated_args: list[TranslationCtx]) -> Function:
        return Function(node.name, *translated_args)

    def simple_translate_func(self, node: nodes.OperationCall, ctx: TranslationCtx) -> None:
        """
        Translate function call using ``OPERATION_REGISTRY``:
        - check its name against a list of allowed functions
        - check argument types (optional)
        - check argument count
        """

        def translator_cb(obj: nodes.FormulaItem) -> Union[nodes.FormulaItem, ClauseElement]:
            """Callback that can translate a formula node for usage from function implementation"""
            return self.translate_node(  # type: ignore  # TODO: fix
                obj,
                ctx=TranslationCtx(node=obj, required_scopes=ctx.required_scopes & ~Scope.EXPLICIT_USAGE),
                postprocess=False,
            ).expression

        func_name = node.name
        ctx.set_token(func_name)
        ctx.fork()

        translated_dimensions = []
        translated_order_by = []
        if isinstance(node, nodes.WindowFuncCall):
            grouping = node.grouping
            if isinstance(grouping, nodes.WindowGroupingTotal):
                pass  # empty dimension list
            elif isinstance(grouping, nodes.WindowGroupingWithin):
                translated_dimensions = [
                    self.translate_node(dim, ctx=ctx.child(node=dim), postprocess=False).expression
                    for dim in grouping.dim_list
                ]
            else:
                raise TypeError(f"{type(grouping)} is not supported as a grouping type")

            translated_order_by = [
                self.translate_node(dim, ctx=ctx.child(node=dim), postprocess=False).expression
                for dim in node.ordering.children
            ]

        translated_args = [self.translate_node(arg, ctx=ctx.child(node=arg), postprocess=False) for arg in node.args]
        arg_types = [arg.data_type for arg in translated_args]

        if self._collect_stats:
            dialect_name = self._env.dialect.single_bit.name.name if self._env.dialect.deterministic else ""
            signature = FunctionStatsSignature(
                name=func_name,
                arg_types=tuple(at.name for at in arg_types),  # type: ignore  # TODO: fix
                dialect=dialect_name,  # type: ignore  # TODO: fix
                is_window=isinstance(node, nodes.WindowFuncCall),
            )
            self._env.translation_stats.add_function_hit(signature)

        # find matching function in the registry
        try:
            func_definition = OPERATION_REGISTRY.get_definition(
                name=func_name,
                arg_types=arg_types,  # type: ignore  # TODO: fix
                is_window=isinstance(node, nodes.WindowFuncCall),
                dialect=self._env.dialect,
                required_scopes=ctx.required_scopes,
            )
            translation_result = func_definition.translate(env=self._env, args=translated_args)
            translated_args = translation_result.transformed_args
            ctx.set_type(translation_result.data_type, translation_result.data_type_params)
            ctx.set_flags(translation_result.context_flags)
            # postprocess nodes after the translation may have updated argument context flags
            if translation_result.postprocess_args:
                # Postprocessing may be switched of if this function redirects to another
                for child_ctx in translated_args:
                    self._postprocess_context(child_ctx)

            translated_func = translation_result.impl_callable(
                *translated_args,
                translator_cb=translator_cb,
                partition_by=translated_dimensions,
                default_order_by=translated_order_by,
                translation_ctx=ctx,
                translation_env=self._env,
            )
            # Here `translated_func` can be a lot of different things (SQLA objects)
            # Even for window functions we cannot rely on it being an `Over` instance

        except exc.TranslationError as err:
            # One of:
            # - no such function,
            # - wrong argument types
            # - error during translation (possibly errors in nested elements)
            if isinstance(err, exc.TranslationUnknownFunctionError) and not self._restrict_funcs:
                translated_func = self._generate_default_function(node, translated_args=translated_args)
            else:
                for each_error in err.errors:
                    ctx.add_error(each_error.message, token=each_error.token or func_name, code=each_error.code)
                translated_func = sa.null()
            ctx.set_type(DataType.NULL)

        ctx.set_expression(translated_func)

    def translate(
        self,
        formula: nodes.Formula,
        collect_errors: Optional[bool] = None,
        context_flags: Optional[int] = None,
    ) -> TranslationCtx:
        """
        Translate ``Formula`` object into an SQLAlchemy selectable,
        which then can be compiled into valid RAW SQL.
        """
        context_flags = context_flags or 0
        ctx = TranslationCtx(
            required_scopes=self._env.required_scopes,
            collect_errors=collect_errors,
            flags=context_flags,
            node=formula,
        )
        self.translate_node(formula.expr, ctx)
        if ctx.errors:
            raise exc.TranslationError(*ctx.errors)

        if isinstance(ctx.expression, (int, float, bool, str, datetime.datetime, datetime.date, type(None))):
            ctx.set_expression(literal(ctx.expression))

        self._coerce_type(ctx)
        return ctx

    @staticmethod
    def _coerce_type(ctx: TranslationCtx):
        """
        Make sure value is cast to proper Python type
        (this is done outside the DB by ``sqlalchemy``)
        """
        if ctx.data_type in (DataType.CONST_INTEGER, DataType.INTEGER):
            ctx.set_expression(sa.type_coerce(ctx.expression, sa.Integer))
        elif ctx.data_type in (DataType.CONST_FLOAT, DataType.FLOAT):
            ctx.set_expression(sa.type_coerce(ctx.expression, sa.Float))
        elif ctx.data_type in (DataType.CONST_BOOLEAN, DataType.BOOLEAN):
            pass  # coercion to BOOL results in extra ' = 1', so don't do it
        elif ctx.data_type in (DataType.CONST_DATE, DataType.DATE):
            ctx.set_expression(sa.type_coerce(ctx.expression, sa.Date))
        elif ctx.data_type in (DataType.CONST_DATETIME, DataType.DATETIME):
            ctx.set_expression(sa.type_coerce(ctx.expression, sa.DateTime))
        elif ctx.data_type in (DataType.CONST_GENERICDATETIME, DataType.GENERICDATETIME):
            ctx.set_expression(sa.type_coerce(ctx.expression, sa.DateTime))

    def translate_node(self, node: nodes.FormulaItem, ctx: TranslationCtx, postprocess: bool = True) -> TranslationCtx:
        """Wrapper for internal implementation of translation methods"""

        # check replacements and cache first
        found_translation_ctx = self._env.replacements.get(node)
        if found_translation_ctx is None:
            found_translation_ctx = self._env.translation_cache.get(node)
            if self._collect_stats:
                self._env.translation_stats.add_cache_hit()
        if found_translation_ctx is not None:
            # found in cache, so reuse it
            ctx.adopt(found_translation_ctx)
        else:
            # not found, so translate it
            self._translate_node(node, ctx)

        ctx.flush()
        try:
            self._env.translation_cache.add(node=node, value=ctx)
        except exc.CacheError:
            pass  # node is not cacheable

        if postprocess:
            self._postprocess_context(ctx)

        return ctx

    def _postprocess_context(self, ctx: TranslationCtx):
        """Additional contextual postprocessing of expressions (dialect dependent)"""

        assert ctx.data_type is not None
        assert ctx.expression is not None

        expression, flags = self._postprocessor.postprocess_expression(
            data_type=ctx.data_type, expression=ctx.expression, flags=ctx.flags
        )
        if expression is not ctx.expression or flags != ctx.flags:
            ctx.set_expression(expression)
            ctx.set_flags(flags)

        for warning in self._postprocessor.get_warnings(data_type=ctx.data_type, flags=ctx.flags):
            ctx.add_warning(warning)

    @singledispatchmethod
    def _translate_node(self, node: nodes.FormulaItem, ctx: TranslationCtx) -> None:
        """
        Default action for singledispatchmethod.
        Is called if ``node`` doesn't match any of the registered types.
        In this case we assume that it has already been translated, so return it as-is
        """
        raise TypeError(type(node))

    @_translate_node.register(TranslationCtx)
    def _translate_translation_ctx(self, node: TranslationCtx, ctx: TranslationCtx) -> None:
        """Already translated node. Do nothing"""
        pass

    @_translate_node.register(nodes.LiteralInteger)
    @_translate_node.register(nodes.LiteralFloat)
    @_translate_node.register(nodes.LiteralBoolean)
    @_translate_node.register(nodes.LiteralString)
    @_translate_node.register(nodes.LiteralDate)
    @_translate_node.register(nodes.LiteralDatetime)
    @_translate_node.register(nodes.LiteralDatetimeTZ)
    @_translate_node.register(nodes.LiteralGenericDatetime)
    @_translate_node.register(nodes.LiteralGeopoint)
    @_translate_node.register(nodes.LiteralGeopolygon)
    @_translate_node.register(nodes.LiteralUuid)
    @_translate_node.register(nodes.LiteralArrayInteger)
    @_translate_node.register(nodes.LiteralArrayFloat)
    @_translate_node.register(nodes.LiteralArrayString)
    @_translate_node.register(nodes.LiteralTreeString)
    def _translate_node_literal(self, node: nodes.BaseLiteral, ctx: TranslationCtx) -> None:
        data_type_params: Optional[DataTypeParams] = None
        if isinstance(node, nodes.LiteralDatetimeTZ):
            data_type_params = DataTypeParams(timezone="UTC")  # not certain
        ctx.set_type(
            dl_formula.inspect.literal.get_data_type(node),
            data_type_params=data_type_params,
        )
        value = node.value
        if isinstance(node, nodes.LiteralUuid):
            value = str(value)
        ctx.set_expression(literal(value, d=self._env.dialect))

    @_translate_node.register(ext_nodes.CompiledExpression)
    def _translate_node_compiled_expression(self, node: ext_nodes.CompiledExpression, ctx: TranslationCtx) -> None:
        if isinstance(node.value, TranslationCtx):
            ctx.adopt(node.value)
        else:
            ctx.set_expression(node.value)

    @_translate_node.register(nodes.Null)
    def _translate_node_null(self, node: Optional[nodes.Null], ctx: TranslationCtx) -> None:
        ctx.set_type(DataType.NULL)
        ctx.set_expression(sa.null())

    @_translate_node.register(nodes.ExpressionList)
    def _translate_node_expression_list(self, node: nodes.ExpressionList, ctx: TranslationCtx) -> None:
        ctx.fork()
        translated_child_list = []
        data_types = []
        for item in node.children:
            translated_child = self.translate_node(item, ctx=ctx.child(node=item))
            translated_child_list.append(translated_child.expression)
            data_types.append(translated_child.data_type)

        ctx.set_expression(translated_child_list)
        ctx.set_type(FromArgs().get_from_args(data_types))  # type: ignore  # TODO: fix  # TODO?: data_type_params

    def _make_literal_column(self, name: str) -> sa.sql.ClauseElement:
        return self._column_renderer.make_column(name)

    @_translate_node.register(nodes.Field)
    def _translate_node_field(self, node: nodes.Field, ctx: TranslationCtx) -> None:
        field_id = node.name
        ctx.set_token(field_id)
        try:
            ctx.set_expression(self._make_literal_column(field_id))
            ctx.set_type(self._env.field_types[field_id], self._env.field_type_params.get(field_id))
        except KeyError:
            if self._restrict_fields:
                name_parts = self._env.field_names.get(field_id) or (field_id,)
                ctx.add_error(
                    f'Translator: no info about column: {".".join(name_parts)}',
                    code=exc.TranslationUnknownFieldError.default_code,
                )
            ctx.set_expression(sa.literal(field_id))
            ctx.set_type(DataType.NULL)

    @_translate_node.register(nodes.Unary)
    @_translate_node.register(nodes.Binary)
    @_translate_node.register(nodes.Ternary)
    @_translate_node.register(nodes.FuncCall)
    @_translate_node.register(nodes.WindowFuncCall)
    def _translate_node_func_call(self, node: nodes.FuncCall, ctx: TranslationCtx) -> None:
        self.simple_translate_func(node=node, ctx=ctx)

    @_translate_node.register(nodes.CaseBlock)
    def _translate_node_case_block(self, node: nodes.CaseBlock, ctx: TranslationCtx) -> None:
        args: list[nodes.FormulaItem] = [
            node.case_expr,
            *chain.from_iterable(when_part.children for when_part in node.when_list),  # type: ignore
        ]
        if node.else_expr is not None:
            args.append(node.else_expr)
        self.simple_translate_func(
            node=nodes.FuncCall.make("_case_block_", args=args), ctx=ctx.child(required_scopes=ctx.required_scopes)
        )

    @_translate_node.register(nodes.IfBlock)
    def _translate_node_if_block(self, node: nodes.IfBlock, ctx: TranslationCtx) -> None:
        args: list[nodes.FormulaItem] = [
            *chain.from_iterable(if_part.children for if_part in node.if_list),  # type: ignore
            node.else_expr,
        ]
        self.simple_translate_func(
            node=nodes.FuncCall.make("_if_block_", args=args),
            ctx=ctx.child(required_scopes=ctx.required_scopes),
        )

    @_translate_node.register(nodes.ParenthesizedExpr)
    def _translate_node_parenthesized_expr(self, node: nodes.ParenthesizedExpr, ctx: TranslationCtx) -> None:
        self.translate_node(node.expr, ctx=ctx, postprocess=False)

    @_translate_node.register(nodes.OrderAscending)
    @_translate_node.register(nodes.OrderDescending)
    def _translate_node_ordering_direction(self, node: nodes.OrderingDirectionBase, ctx: TranslationCtx) -> None:
        ctx.fork()
        child = self.translate_node(node.expr, ctx=ctx.child(node=node.expr), postprocess=True)
        expression = child.expression
        if isinstance(node, nodes.OrderDescending):
            expression = desc(expression)  # type: ignore  # TODO: fix
        ctx.set_expression(expression)
        ctx.set_type(child.data_type, child.data_type_params)  # type: ignore  # TODO: fix

    @_translate_node.register(aux_nodes.ErrorNode)
    def _translate_node_error(self, node: aux_nodes.ErrorNode, ctx: TranslationCtx) -> None:
        ctx.set_expression(sa.null())
        ctx.set_type(DataType.NULL)
        ctx.add_error(message=node.message, code=node.err_code)


def translate(
    formula: nodes.Formula,
    dialect: Optional[DialectCombo] = None,
    required_scopes: int = Scope.EXPLICIT_USAGE,
    restrict_funcs: Optional[bool] = None,
    restrict_fields: Optional[bool] = None,
    collect_errors: Optional[bool] = None,
    collect_stats: bool = False,
    field_types: Optional[dict[str, DataType]] = None,
    context_flags: Optional[int] = None,
    field_names: Optional[dict[str, tuple[str, ...]]] = None,
    env: Optional[TranslationEnvironment] = None,
) -> TranslationCtx:
    """Translate ``Formula`` tree object into an SQLAlchemy representation that can be used for queries"""

    if env is None:
        assert dialect is not None
        env = TranslationEnvironment(
            dialect=dialect,
            required_scopes=required_scopes,
            field_types=field_types or {},
            field_names=field_names or {},
        )
    assert env is not None

    translator = SqlAlchemyTranslator(
        env=env,
        restrict_funcs=restrict_funcs,
        restrict_fields=restrict_fields,
        collect_stats=collect_stats,
    )
    return translator.translate(
        formula,
        collect_errors=collect_errors,
        context_flags=context_flags,
    )


def translate_and_compile(
    formula: nodes.Formula,
    dialect: Optional[DialectCombo] = None,
    required_scopes: int = Scope.EXPLICIT_USAGE,
    restrict_funcs: Optional[bool] = None,
    restrict_fields: Optional[bool] = None,
    collect_errors: Optional[bool] = None,
    collect_stats: bool = False,
    field_types: Optional[dict[str, DataType]] = None,
    context_flags: Optional[int] = None,
    field_names: Optional[dict[str, tuple[str, ...]]] = None,
    env: Optional[TranslationEnvironment] = None,
) -> str:
    """Translate ``Formula`` tree object into an SQLAlchemy representation and compile it into raw SQL"""

    if env is None:
        assert dialect is not None
        env = TranslationEnvironment(
            dialect=dialect,
            required_scopes=required_scopes,
            field_types=field_types or {},
            field_names=field_names or {},
        )
    assert env is not None

    translator = SqlAlchemyTranslator(
        env=env,
        restrict_funcs=restrict_funcs,
        restrict_fields=restrict_fields,
        collect_stats=collect_stats,
    )
    result = translator.translate(
        formula,
        collect_errors=collect_errors,
        context_flags=context_flags,
    )
    sa_dialect = get_sa_dialect(env.dialect)
    sa_compiled_expr = result.expression.compile(  # type: ignore  # TODO: fix
        compile_kwargs={"literal_binds": True},
        dialect=sa_dialect,
    )
    compiled_text = sa_compiled_expr.string
    return compiled_text

lib/dl_formula/dl_formula/translation/translator.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_formula.core.datatype import (
    DataType,
    DataTypeParams,
)


class TypeStrategy:
    __slots__ = ()

    def get_from_args(self, arg_types) -> DataType:
        raise NotImplementedError


class TypeParamsStrategy:
    __slots__ = ()

    def get_from_arg_values(self, args) -> DataTypeParams:
        raise NotImplementedError

lib/dl_formula/dl_formula/translation/type_strategy.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import (
    _CacheInfo,
    lru_cache,
    wraps,
)
import threading
from typing import (
    Any,
    Callable,
    Dict,
    Generic,
    Optional,
    Tuple,
    Type,
    TypeVar,
)


_FUNC = Callable[..., Any]
_QUALIFIER_VALUE_TV = TypeVar("_QUALIFIER_VALUE_TV")


class MultiCacheManager(Generic[_QUALIFIER_VALUE_TV]):
    def __init__(
        self,
        wrapped_function: _FUNC,
        maxsize: int,
        cache_exceptions: Tuple[Type[Exception], ...] = (),
        cache_qualifier: Optional[Callable[..., _QUALIFIER_VALUE_TV]] = None,
    ):
        self._wrapped_function = wrapped_function
        self._cache_exceptions = cache_exceptions
        self._maxsize = maxsize
        self._cache_qualifier = cache_qualifier
        self._cached_wrappers: Dict[Optional[_QUALIFIER_VALUE_TV], _FUNC] = {}
        self._cached_wrappers_lock = threading.Lock()

    def _get_qualifier_value(self, *args: Any, **kwargs: Any) -> Optional[_QUALIFIER_VALUE_TV]:
        if self._cache_qualifier is None:
            return None
        return self._cache_qualifier(*args, **kwargs)

    def _make_cached_wrapper(self) -> _FUNC:
        @lru_cache(self._maxsize)
        def _cached_result_and_error_wrapper(*args: Any, **kwargs: Any) -> Tuple[Optional[Exception], Any]:
            error: Optional[Exception] = None
            result: Any = None
            try:
                result = self._wrapped_function(*args, **kwargs)
            except self._cache_exceptions as err:
                error = err

            return error, result

        @wraps(self._wrapped_function)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            error, result = _cached_result_and_error_wrapper(*args, **kwargs)
            if error:
                raise error
            return result

        wrapper.cache_info = _cached_result_and_error_wrapper.cache_info  # type: ignore
        wrapper.cache_clear = _cached_result_and_error_wrapper.cache_clear  # type: ignore
        return wrapper

    def get_cached_wrapper_for_args(self, *args: Any, **kwargs: Any) -> _FUNC:
        qvalue = self._get_qualifier_value(*args, **kwargs)
        with self._cached_wrappers_lock:
            if qvalue not in self._cached_wrappers:
                wrapper = self._make_cached_wrapper()
                self._cached_wrappers[qvalue] = wrapper
            else:
                wrapper = self._cached_wrappers[qvalue]

        return wrapper

    def cache_info(self) -> Dict[Optional[_QUALIFIER_VALUE_TV], _CacheInfo]:
        """Collect cache info objects for all existing cache qualifier values."""
        with self._cached_wrappers_lock:
            return {qvalue: wrapper.cache_info() for qvalue, wrapper in self._cached_wrappers.items()}  # type: ignore

    def cache_clear(self) -> None:
        """Clear all existing caches."""
        with self._cached_wrappers_lock:
            for wrapper in self._cached_wrappers.values():
                wrapper.cache_clear()  # type: ignore


def multi_cached_with_errors(
    maxsize: int,
    cache_exceptions: Tuple[Type[Exception], ...] = (),
    cache_qualifier: Optional[Callable[..., _QUALIFIER_VALUE_TV]] = None,
) -> Callable[[_FUNC], _FUNC]:
    """
    Parameterized decorator that works just like ``lru_cache``,
    but instead of just caching results of successful executions,
    it also caches raised exceptions
    """

    def decorator(func: _FUNC) -> _FUNC:
        cache_manager: MultiCacheManager[_QUALIFIER_VALUE_TV] = MultiCacheManager(
            wrapped_function=func,
            maxsize=maxsize,
            cache_exceptions=cache_exceptions,
            cache_qualifier=cache_qualifier,
        )

        @wraps(func)
        def generic_wrapper(*args: Any, **kwargs: Any) -> Any:
            wrapper = cache_manager.get_cached_wrapper_for_args(*args, **kwargs)
            return wrapper(*args, **kwargs)

        generic_wrapper.multi_cache_manager = cache_manager  # type: ignore
        generic_wrapper.cache_clear = cache_manager.cache_clear  # type: ignore
        generic_wrapper.cache_info = cache_manager.cache_info  # type: ignore

        return generic_wrapper

    return decorator

lib/dl_formula/dl_formula/utils/caching.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from typing import Union

import ciso8601
import pytz
import pytz.tzinfo


# TODO: re-check and document the useful differences between
# `datetime.timezone.utc` and `pytz.utc`.
UTC = datetime.timezone.utc


SOME_DT = datetime.datetime(2020, 1, 1, 0, 0, 0)


DT_FORMATS = (
    "%Y-%m-%d",
    "%Y-%m-%d %H",
    "%Y-%m-%d %H:%M",
    "%Y-%m-%d %H:%M:%S",
    # To support e.g. '#2019-01-01T#'. Dubious.
    "%Y-%m-%d ",
)


DT_FORMATS_BY_LENGTH = {
    # length: strptime_format
    # assumes preprocessed 'T'
    len(SOME_DT.strftime(fmt)): fmt
    for fmt in DT_FORMATS
}


def get_tzobj(tzname: str) -> datetime.tzinfo:
    try:
        # Should essentially be a hashmap lookup, no need to cache it further.
        return pytz.timezone(tzname)
    except pytz.UnknownTimeZoneError:
        raise ValueError(f"Invalid timezone name: {tzname!r}")


def parse_dt_string(value: str) -> datetime.datetime:
    simple_format = DT_FORMATS_BY_LENGTH.get(len(value))
    if simple_format is not None:
        value = value.replace("T", " ").replace("t", " ")
        result = datetime.datetime.strptime(value, simple_format)  # raises ValueError
    else:
        result = ciso8601.parse_datetime(value)  # raises ValueError

    return result


DTSourcesT = Union[None, str, int, float, datetime.date, datetime.datetime]


def make_datetime_value_base(value: DTSourcesT) -> datetime.datetime:
    """Convert a value into a datetime object, without tzinfo or with a fixed offset"""
    if isinstance(value, str):
        return parse_dt_string(value)
    if isinstance(value, (int, float)):
        return datetime.datetime.utcfromtimestamp(value)
    if isinstance(value, datetime.datetime):
        return value
    if isinstance(value, datetime.date):
        return datetime.datetime.combine(value, datetime.time())
    raise Exception("Unexpected value type", dict(value_type=type(value)))


def make_datetime_value(value: DTSourcesT) -> datetime.datetime:
    """Convert a value into a datetime object, without tzinfo or at UTC"""
    dt = make_datetime_value_base(value)
    if dt.tzinfo is None:
        return dt
    # There's no known use for an offset at this point,
    # so ensure the result is at UTC.
    return dt.astimezone(UTC)


def make_datetimetz_value(value: DTSourcesT, tzname: str = "UTC") -> datetime.datetime:
    """Convert a value into a datetime object, at UTC, using tzname to interpret naive datetimes"""
    tzobj = get_tzobj(tzname)  # also pre-validates `tzname`
    assert isinstance(tzobj, (pytz.tzinfo.DstTzInfo, pytz.tzinfo.StaticTzInfo))
    dt = make_datetime_value(value)
    assert dt is not None
    if dt.tzinfo is None:
        dt = tzobj.localize(dt)
    dt = dt.astimezone(UTC)
    return dt

lib/dl_formula/dl_formula/utils/datetime.py
datalens-tech/datalens-backend
from __future__ import annotations

from itertools import chain
from typing import (
    ClassVar,
    List,
    Sequence,
    Tuple,
    Type,
)

import attr

from dl_formula.collections import NodeSet
import dl_formula.core.exc as exc
import dl_formula.core.fork_nodes as fork_nodes
import dl_formula.core.nodes as nodes
from dl_formula.inspect.env import InspectionEnvironment
import dl_formula.inspect.expression
import dl_formula.inspect.function
import dl_formula.inspect.node
from dl_formula.validation.env import ValidationEnvironment
from dl_formula.validation.validator import (
    Checker,
    ValidatorProxy,
)


@attr.s
class AggregationChecker(Checker):
    _inspect_env: InspectionEnvironment = attr.ib()
    _valid_env: ValidationEnvironment = attr.ib()
    _allow_nested_agg: bool = attr.ib(default=False)
    _allow_inconsistent_inside_agg: bool = attr.ib(default=False)
    _global_dimensions: List[nodes.FormulaItem] = attr.ib(kw_only=True)

    _exclude_node_types: ClassVar[Tuple[Type[nodes.FormulaItem], ...]] = (
        fork_nodes.QueryForkJoiningBase,
        nodes.LodSpecifier,
    )

    """
    Check whether the usage of aggregations is consistent throughout the expression.
    The following are considered an error:
    - multiple aggregations in nested nodes
    - nodes with different aggregation levels are present at a single level as children of one node
    """

    def _check_is_dimension_bound(self, node: nodes.FormulaItem, dimension_set: NodeSet) -> bool:
        """
        Return ``True`` if the node tree is dependent only on dimensions.
        Cache the results in ``self._valid_env``.
        """

        dim_bound = dl_formula.inspect.expression.is_bound_only_to(node, allow_nodes=dimension_set)
        return dim_bound

    def perform_node_check(
        self,
        validator: ValidatorProxy,
        node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> None:
        if not dl_formula.inspect.expression.is_aggregate_expression(node, env=self._inspect_env):
            # node has no aggregations in it,
            # so by definition there's nothing wrong with its aggregation
            return

        def _aggregation_in_stack(node_stack: Sequence[nodes.FormulaItem]) -> bool:
            for wrapping_node in node_stack:
                if dl_formula.inspect.node.is_aggregate_function(wrapping_node):
                    return True
            return False

        not_agg_children = False
        agg_children = False
        children_w_stacks = dl_formula.inspect.expression.autonomous_children_w_stack(
            node,  # TODO: Find a solution more elegant than excludes:
            parent_stack=parent_stack,
            exclude_node_types=self._exclude_node_types,
        )

        for child, stack in children_w_stacks:
            with validator.handle_error(node=node):
                self.check_node(validator=validator, node=child, parent_stack=stack)

            if dl_formula.inspect.expression.is_aggregate_expression(child, env=self._inspect_env):
                agg_children = True
            else:
                _, dimension_set, _ = dl_formula.inspect.expression.resolve_dimensions(
                    node_stack=chain(stack, (child,)),
                    dimensions=self._global_dimensions,
                    env=self._inspect_env,
                )
                dim_bound = self._check_is_dimension_bound(child, dimension_set=dimension_set)
                if not dim_bound:
                    not_agg_children = True
            if agg_children and not_agg_children:
                with validator.handle_error(node=node):
                    do_raise = True
                    if self._allow_inconsistent_inside_agg:
                        # Allow agg and non-agg arguments inside another aggregation
                        # (find aggregation in parent_stack)
                        do_raise = not _aggregation_in_stack(parent_stack)

                    if do_raise:
                        raise exc.InconsistentAggregationError(
                            "Inconsistent aggregation among operands",
                            token=dl_formula.inspect.node.get_token(child),
                            position=child.position,
                        )

        if not self._allow_nested_agg and dl_formula.inspect.node.is_aggregate_function(node):
            # function is an aggregation and some of its arguments are already aggregated -> double aggregation
            if agg_children:
                with validator.handle_error(node=node):
                    raise exc.DoubleAggregationError(
                        "Double aggregation is forbidden",
                        token=dl_formula.inspect.node.get_token(node),
                        position=node.position,
                    )

lib/dl_formula/dl_formula/validation/aggregation.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from typing import (
    Dict,
    Optional,
)

import attr

from dl_formula.collections import NodeValueMap
import dl_formula.core.exc as exc


@attr.s
class ErrInfo:
    is_error: bool = attr.ib(default=False)
    exception: Optional[exc.ValidationError] = attr.ib(default=None)


@attr.s
class ValidationEnvironment:
    cache_dim_bound: NodeValueMap[bool] = attr.ib(factory=NodeValueMap)
    generic_cache_valid: Dict[type, NodeValueMap[ErrInfo]] = attr.ib(factory=lambda: defaultdict(NodeValueMap))

lib/dl_formula/dl_formula/validation/env.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from contextlib import contextmanager
from typing import (
    Iterable,
    List,
    Optional,
    Tuple,
    Type,
)

import dl_formula.core.exc as exc
from dl_formula.core.message_ctx import FormulaErrorCtx
import dl_formula.core.nodes as nodes
from dl_formula.validation.env import (
    ErrInfo,
    ValidationEnvironment,
)


class Validator:
    def __init__(self, env: ValidationEnvironment, collect_errors: bool):
        self._env = env
        self._collect_errors = collect_errors

    @contextmanager
    def handle_error(self, checker_cls: Type["Checker"], node: nodes.FormulaItem):
        try:
            yield
        except exc.ValidationError as err:
            info = self.get_from_cache(checker_cls=checker_cls, node=node)
            if info is not None and info.exception is not None:
                # There already is an exception with errors for this node. Combine them
                errors = []
                errors.extend(info.exception.errors)
                errors.extend(err.errors)
                err = exc.ValidationError(*errors)

            self._env.generic_cache_valid[checker_cls].add(node, value=ErrInfo(is_error=True, exception=err))

            if not self._collect_errors:
                raise err

    def get_from_cache(self, checker_cls: Type["Checker"], node: nodes.FormulaItem) -> Optional[ErrInfo]:
        return self._env.generic_cache_valid[checker_cls].get(node)

    def mark_as_ok_in_cache(self, checker_cls: Type["Checker"], node: nodes.FormulaItem) -> None:
        return self._env.generic_cache_valid[checker_cls].add(node, value=ErrInfo(is_error=False, exception=None))

    def proxy_for(self, checker: "Checker") -> "ValidatorProxy":
        return ValidatorProxy(validator=self, checker_cls=type(checker))

    def get_all_errors(self, node: nodes.FormulaItem) -> List[FormulaErrorCtx]:
        errors = []  # type: ignore  # TODO: fix
        for error_cache in self._env.generic_cache_valid.values():
            info = error_cache.get(node)
            if info is None:
                continue
            if info.is_error and info.exception is not None:  # type: ignore  # TODO: fix
                errors.extend(info.exception.errors)  # type: ignore  # TODO: fix
        return errors


class ValidatorProxy:
    def __init__(self, validator: Validator, checker_cls: Type["Checker"]):
        self._validator = validator
        self._checker_cls = checker_cls

    @contextmanager
    def handle_error(self, node: nodes.FormulaItem):
        with self._validator.handle_error(checker_cls=self._checker_cls, node=node):
            yield

    def get_from_cache(self, node: nodes.FormulaItem) -> Optional[ErrInfo]:
        return self._validator.get_from_cache(checker_cls=self._checker_cls, node=node)

    def mark_as_ok_in_cache(self, node: nodes.FormulaItem) -> None:
        self._validator.mark_as_ok_in_cache(checker_cls=self._checker_cls, node=node)


class Checker(abc.ABC):
    def check_node(
        self,
        validator: ValidatorProxy,
        node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> None:
        """
        Check the validity of a node recursively an either exit or raise an error.
        First check whether the node has already been checked by this ``Checker``.
        Is so, then skip the check and add errors to the list if any are available in the cache.

        This method does not implement the check itself - that is delegated to ``self.perform_node_check(...)``.
        """

        from_cache = validator.get_from_cache(node)
        if from_cache is not None:
            if from_cache.exception is not None:
                raise from_cache.exception
            return

        self.perform_node_check(validator=validator, node=node, parent_stack=parent_stack)

        from_cache = validator.get_from_cache(node)
        if from_cache is not None and from_cache.exception:
            raise from_cache.exception
        else:
            # no errors in this node and its children
            validator.mark_as_ok_in_cache(node)

    @abc.abstractmethod
    def perform_node_check(
        self,
        validator: ValidatorProxy,
        node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> None:
        """Method checks the validity of a node recursively and either exits or raises an error"""

        raise NotImplementedError


def validate(
    node: nodes.FormulaItem,
    env: ValidationEnvironment,
    checkers: Iterable[Checker] = (),
    collect_errors: bool = False,
) -> None:
    """
    Apply multiple checkers to a formula item tree.
    Either exit quietly or raise a ``ValidationError``.
    """

    validator = Validator(env=env, collect_errors=collect_errors)
    for checker in checkers:
        try:
            checker.check_node(validator=validator.proxy_for(checker), node=node, parent_stack=())
        except exc.ValidationError:
            if not collect_errors:
                raise

    errors = validator.get_all_errors(node)
    if errors:
        raise exc.ValidationError(*errors)

lib/dl_formula/dl_formula/validation/validator.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Collection,
    Tuple,
)

import attr

import dl_formula.core.exc as exc
import dl_formula.core.nodes as nodes
from dl_formula.inspect.env import InspectionEnvironment
import dl_formula.inspect.expression
import dl_formula.inspect.function
import dl_formula.inspect.node
from dl_formula.validation.validator import (
    Checker,
    ValidatorProxy,
)


@attr.s
class WindowFunctionChecker(Checker):
    _inspect_env: InspectionEnvironment = attr.ib()
    _allow_nested: bool = attr.ib(default=False)
    _filter_ids: Collection[str] = attr.ib(factory=frozenset)
    _unselected_dimension_ids: Collection[str] = attr.ib(factory=frozenset)

    def perform_node_check(
        self,
        validator: ValidatorProxy,
        node: nodes.FormulaItem,
        parent_stack: Tuple[nodes.FormulaItem, ...],
    ) -> None:
        if not dl_formula.inspect.expression.is_window_expression(node, env=self._inspect_env):
            # node has no window functions in it,
            return

        children_w_stacks = dl_formula.inspect.expression.autonomous_children_w_stack(
            node,
            parent_stack=parent_stack,
        )
        for child, stack in children_w_stacks:
            with validator.handle_error(node=node):
                self.check_node(validator=validator, node=child, parent_stack=stack)

        if isinstance(node, nodes.WindowFuncCall):
            # inspect args
            has_agg_children = False
            for child in node.args:
                if not self._allow_nested:
                    if dl_formula.inspect.expression.is_window_expression(child, env=self._inspect_env):
                        with validator.handle_error(node=node):
                            raise exc.NestedWindowFunctionError(
                                "Nested window functions are not allowed",
                                token=dl_formula.inspect.node.get_token(child),
                                position=child.position,
                            )

                if dl_formula.inspect.expression.is_aggregate_expression(child, env=self._inspect_env):
                    has_agg_children = True

            if not has_agg_children:
                with validator.handle_error(node=node):
                    raise exc.WindowFunctionWOAggregationError(
                        f"Window function {node.name.upper()} has no aggregated expressions among its arguments",
                        token=dl_formula.inspect.node.get_token(node),
                        position=node.position,
                    )

            bfb = node.before_filter_by
            for bfb_name in sorted(bfb.field_names):
                if bfb_name in self._filter_ids and bfb_name in self._unselected_dimension_ids:
                    with validator.handle_error(node=bfb):
                        raise exc.WindowFunctionUnselectedDimensionError(
                            "Window function's BEFORE FILTER BY clause refers to a field "
                            "that is neither an aggregation nor a dimension in the query.",
                            token=dl_formula.inspect.node.get_token(bfb),
                            position=bfb.position,
                        )

lib/dl_formula/dl_formula/validation/window.py
datalens-tech/datalens-backend
import pytest

from dl_formula.loader import load_formula_lib
from dl_formula.parser.base import FormulaParser
from dl_formula.parser.factory import (
    ParserType,
    get_parser,
)
from dl_formula_testing.forced_literal import forced_literal_use  # noqa


PARSERS = (ParserType.antlr_py,)


@pytest.fixture(
    scope="session",
    params=PARSERS,
    ids=[parser.name for parser in PARSERS],
)
def parser(request) -> FormulaParser:
    parser_type = request.param
    return get_parser(parser_type=parser_type)


@pytest.fixture(scope="session", autouse=True)
def loaded_bi_libraries() -> None:
    load_formula_lib()

lib/dl_formula/dl_formula_tests/unit/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_formula.parser.factory import get_parser


def test_extract():
    parser = get_parser()
    assert parser.parse('my_func(  123, "qwerty")').extract == parser.parse("MY_FUNC(123, 'qwerty')").extract
    assert parser.parse("MY_FUNC(123, 'qwerty')").extract != parser.parse("MY_FUNC(1, 'qwerty')").extract
    assert parser.parse("named_without_children()").extract != parser.parse("[named_without_children]").extract

lib/dl_formula/dl_formula_tests/unit/core/test_extracts.py
datalens-tech/datalens-backend
from __future__ import annotations

import copy
from typing import Dict

from dl_formula.core import nodes
from dl_formula.core.index import NodeHierarchyIndex
from dl_formula.core.position import Position


def test_equality():
    assert nodes.LiteralString.make("a") == nodes.LiteralString.make("a")
    assert nodes.Field.make(name="a") == nodes.Field.make(name="a")
    assert nodes.FuncCall.make(name="a", args=[nodes.LiteralInteger.make(123)]) == (
        nodes.FuncCall.make(name="a", args=[nodes.LiteralInteger.make(123)])
    )


def test_copy():
    def _check_node_copy(node: nodes.FormulaItem) -> None:
        node_copy = copy.copy(node)
        assert node_copy == node
        assert node_copy is not node
        assert len(node.children) == len(node.children)
        for i in range(len(node.children)):
            assert node_copy.children[i] == node.children[i]
            assert node_copy.children[i] is not node.children[i]
        assert node_copy.original_text == node.original_text
        assert node_copy.position == node.position

    _check_node_copy(nodes.LiteralInteger.make(123, meta=nodes.NodeMeta(position=Position(1, 2, 0, 0, 1, 2))))
    _check_node_copy(nodes.Field.make("qwerty"))
    _check_node_copy(nodes.FuncCall.make("func", args=[nodes.LiteralInteger.make(123)]))
    _check_node_copy(
        nodes.FuncCall.make("func", args=[nodes.LiteralInteger.make(123), nodes.LiteralString.make("qwerty")])
    )
    _check_node_copy(nodes.Unary.make("un", nodes.LiteralInteger.make(123)))
    _check_node_copy(nodes.Binary.make("bin", nodes.LiteralInteger.make(123), nodes.LiteralString.make("qwerty")))
    _check_node_copy(
        nodes.Ternary.make(
            "tern", nodes.LiteralInteger.make(123), nodes.LiteralString.make("qwerty"), nodes.LiteralInteger.make(34)
        )
    )
    _check_node_copy(
        nodes.WindowFuncCall.make(
            "wfunc",
            args=[nodes.LiteralInteger.make(123)],
            grouping=nodes.WindowGroupingWithin.make(dim_list=[nodes.Field.make(name="my_field")]),
        )
    )
    _check_node_copy(
        nodes.WindowFuncCall.make(
            "owfunc",
            args=[nodes.LiteralInteger.make(123)],
            grouping=nodes.WindowGroupingWithin.make(
                dim_list=[
                    nodes.Field.make(name="my_field"),
                ]
            ),
            ordering=nodes.Ordering.make(
                expr_list=[
                    nodes.Field.make("dim"),
                ]
            ),
            before_filter_by=nodes.BeforeFilterBy.make(
                field_names=[
                    "field_1",
                    "field_3",
                    "field_2",
                ]
            ),
        )
    )
    _check_node_copy(
        nodes.IfBlock.make(
            [
                nodes.IfPart.make(nodes.LiteralBoolean.make(True), nodes.LiteralString.make("qwerty")),
                nodes.IfPart.make(nodes.LiteralBoolean.make(True), nodes.LiteralString.make("qwerty")),
            ],
            nodes.LiteralString.make("abc"),
        )
    )
    _check_node_copy(
        nodes.CaseBlock.make(
            nodes.LiteralString.make("val"),
            [
                nodes.WhenPart.make(nodes.LiteralString.make("case1"), nodes.LiteralString.make("qwerty")),
                nodes.WhenPart.make(nodes.LiteralString.make("case2"), nodes.LiteralString.make("asdfg")),
            ],
            nodes.LiteralString.make("abc"),
        )
    )
    _check_node_copy(
        nodes.Formula.make(
            expr=nodes.LiteralInteger.make(123),
            meta=nodes.NodeMeta(position=Position(0, 3, 0, 0, 0, 3), original_text="123"),
        )
    )


def test_light_copy():
    def _check_node_copy(
        src: nodes.FormulaItem,
        to_replace_map: Dict[int, nodes.FormulaItem] = None,
        expected: nodes.FormulaItem = None,
    ):
        assert src.light_copy(src.children) == src

        if to_replace_map is None:
            return

        children = [to_replace_map[idx] if idx in to_replace_map else child for idx, child in enumerate(src.children)]
        actual = src.light_copy(children)
        assert actual == expected

        assert len(actual.children) == len(expected.children)
        for idx, expected_child in enumerate(expected.children):
            assert actual.children[idx] == expected_child
            if idx not in to_replace_map:
                assert actual.children[idx] is src.children[idx]

        assert actual.original_text == expected.original_text
        assert actual.position == expected.position

    _check_node_copy(nodes.LiteralInteger.make(123, meta=nodes.NodeMeta(position=Position(1, 2, 0, 0, 1, 2))))
    _check_node_copy(nodes.Field.make("qwerty"))

    lit_123 = nodes.LiteralInteger.make(123)
    lit_234 = nodes.LiteralInteger.make(234)
    lit_x = nodes.LiteralString.make("x")

    field_my = nodes.Field.make(name="my_field")
    win_grouping_within_my_field = nodes.WindowGroupingWithin.make(dim_list=[field_my])
    win_grouping_within_my_field_2 = nodes.WindowGroupingWithin.make(dim_list=[nodes.Field.make(name="my_field2")])

    if_lit_123 = nodes.IfPart.make(nodes.LiteralBoolean.make(True), lit_123)
    if_lit_x = nodes.IfPart.make(nodes.LiteralBoolean.make(True), lit_x)

    case_lit_1_123 = nodes.WhenPart.make(nodes.LiteralString.make("case1"), lit_123)
    case_lit_1_x = nodes.WhenPart.make(nodes.LiteralString.make("case1"), lit_x)
    case_lit_2_234 = nodes.WhenPart.make(nodes.LiteralString.make("case2"), lit_234)

    _check_node_copy(
        nodes.FuncCall.make("func", args=[lit_123, lit_x]),
        {0: lit_234},
        nodes.FuncCall.make("func", args=[lit_234, lit_x]),
    )
    _check_node_copy(
        nodes.Unary.make("un", lit_123),
        {0: lit_234},
        nodes.Unary.make("un", lit_234),
    )
    _check_node_copy(
        nodes.Binary.make("bin", lit_123, nodes.LiteralString.make("qwerty")),
        {1: lit_234},
        nodes.Binary.make("bin", lit_123, lit_234),
    )
    _check_node_copy(
        nodes.Ternary.make("tern", lit_123, nodes.LiteralString.make("qwerty"), lit_234),
        {1: lit_x, 2: lit_123},
        nodes.Ternary.make("tern", lit_123, lit_x, lit_123),
    )
    _check_node_copy(
        nodes.WindowFuncCall.make("wfunc", args=[lit_123], grouping=win_grouping_within_my_field),
        {2: win_grouping_within_my_field_2},  # 0 single arg, 1  ordering, 2 grouping ...
        nodes.WindowFuncCall.make(
            "wfunc",
            args=[lit_123],
            grouping=win_grouping_within_my_field_2,
        ),
    )
    _check_node_copy(
        nodes.IfBlock.make([if_lit_123, if_lit_x], lit_234),
        {1: if_lit_123},
        nodes.IfBlock.make([if_lit_123, if_lit_123], lit_234),
    )
    _check_node_copy(
        nodes.CaseBlock.make(field_my, [case_lit_1_123, case_lit_2_234], lit_x),
        {1: case_lit_1_x, 3: lit_123},
        nodes.CaseBlock.make(field_my, [case_lit_1_x, case_lit_2_234], lit_123),
    )
    meta_1 = nodes.NodeMeta(position=Position(0, 3, 0, 0, 0, 3), original_text="123")
    _check_node_copy(
        nodes.Formula.make(
            expr=lit_123,
            meta=meta_1,
        ),
        {0: lit_x},
        nodes.Formula.make(
            expr=lit_x,
            meta=meta_1,
        ),
    )


def test_search_with_none_pos_range():
    f = nodes.Formula.make(
        expr=nodes.LiteralString.make("op"),
        meta=nodes.NodeMeta(position=Position(0, 2, 0, 0, 0, 2), original_text="op"),
    )
    assert f.get_by_pos(1) is f


def test_get():
    the_node = nodes.Field.make("My Field")
    index = NodeHierarchyIndex.make(0, 1)
    assert (
        nodes.Formula.make(
            nodes.FuncCall.make(
                name="concat",
                args=[
                    nodes.Field.make("Other Field"),
                    the_node,
                    nodes.LiteralString.make("qwerty"),
                ],
            )
        )[index]
        is the_node
    )


def test_replace_at_index():
    index = NodeHierarchyIndex.make(0, 1)
    updated_formula = nodes.Formula.make(
        nodes.FuncCall.make(
            name="concat",
            args=[
                nodes.Field.make("Other Field"),
                nodes.Field.make("Old Field"),
                nodes.LiteralString.make("qwerty"),
            ],
        )
    ).replace_at_index(index, nodes.Field.make("New Field"))
    expected_formula = nodes.Formula.make(
        nodes.FuncCall.make(
            name="concat",
            args=[
                nodes.Field.make("Other Field"),
                nodes.Field.make("New Field"),
                nodes.LiteralString.make("qwerty"),
            ],
        )
    )
    assert updated_formula == expected_formula

    # Also make sure that extracts match
    for index, updated_child in updated_formula.enumerate():
        expected_child = expected_formula[index]
        assert updated_child.extract == expected_child.extract


def test_substitute_batch():
    nhi = NodeHierarchyIndex.make
    r1 = nodes.Field.make("R1")
    r2 = nodes.Field.make("R2")

    updated_formula = nodes.Formula.make(
        nodes.FuncCall.make(
            name="concat",
            args=[
                nodes.Field.make("F1"),
                nodes.Field.make("F1"),
                nodes.Field.make("F2"),
                nodes.Field.make("F3"),
            ],
        )
    ).substitute_batch(
        {
            nhi(0, 0): nodes.Field.make("R1"),
            nhi(0, 1): nodes.Field.make("R1"),
            nhi(0, 2): nodes.Field.make("R2"),
        }
    )
    expected_formula = nodes.Formula.make(
        nodes.FuncCall.make(
            name="concat",
            args=[
                r1,
                r1,
                r2,
                nodes.Field.make("F3"),
            ],
        )
    )
    assert updated_formula == expected_formula

    # Also make sure that extracts match
    for index, updated_child in updated_formula.enumerate():
        expected_child = expected_formula[index]
        assert updated_child.extract == expected_child.extract


def test_enumerate():
    nhi = NodeHierarchyIndex.make
    main_func = nodes.FuncCall.make(
        name="concat",
        args=[
            nodes.Field.make("Other Field"),
            nodes.Field.make("Green Field"),
            nodes.LiteralString.make("qwerty"),
        ],
    )
    formula = nodes.Formula.make(main_func)
    assert list(formula.enumerate()) == [
        (nhi(), formula),
        (nhi(0), main_func),
        (nhi(0, 0), nodes.Field.make("Other Field")),
        (nhi(0, 1), nodes.Field.make("Green Field")),
        (nhi(0, 2), nodes.LiteralString.make("qwerty")),
        (nhi(0, 3), nodes.DefaultAggregationLodSpecifier()),  # implicitly created
        (nhi(0, 4), nodes.IgnoreDimensions()),  # implicitly created
        (nhi(0, 5), nodes.BeforeFilterBy.make()),  # implicitly created
    ]


def test_iter_index():
    main_func = nodes.FuncCall.make(
        name="concat",
        args=[
            nodes.Field.make("Other Field"),
            nodes.Field.make("Green Field"),
            nodes.LiteralString.make("qwerty"),
        ],
    )
    formula = nodes.Formula.make(main_func)
    assert list(formula.iter_index(index=NodeHierarchyIndex.make(0, 1))) == [
        formula,
        main_func,
        nodes.Field.make("Green Field"),
    ]
    assert list(formula.iter_index(index=NodeHierarchyIndex.make(0, 1), exclude_last=True)) == [
        formula,
        main_func,
    ]

lib/dl_formula/dl_formula_tests/unit/core/test_nodes.py
datalens-tech/datalens-backend
from __future__ import annotations

import copy

from dl_formula.collections import NodeSet
from dl_formula.core import nodes
from dl_formula.core.datatype import DataType
from dl_formula.core.index import NodeHierarchyIndex
from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.inspect.expression import (
    enumerate_fields,
    get_query_fork_nesting_level,
    infer_data_type,
    is_aggregate_expression,
    is_aggregated_above_sub_node,
    is_bound_only_to,
    is_constant_expression,
    is_query_fork_expression,
    is_window_expression,
    iter_operation_calls,
    used_fields,
    used_func_calls,
)
from dl_formula.shortcuts import n


def test_used_func_calls():
    assert used_func_calls(n.formula(nodes.FuncCall.make(name="func", args=[nodes.LiteralString.make("qwe")]))) == [
        nodes.FuncCall.make(name="func", args=[nodes.LiteralString.make("qwe")]),
    ]
    assert used_func_calls(
        n.formula(
            nodes.FuncCall.make(
                name="func",
                args=[
                    nodes.Field.make("f1"),
                    nodes.FuncCall.make(
                        name="other_func",
                        args=[
                            nodes.Field.make("f2"),
                            n.p(nodes.Field.make("f3")),
                        ],
                    ),
                ],
            )
        )
    ) == [
        nodes.FuncCall.make(
            name="func",
            args=[
                nodes.Field.make("f1"),
                nodes.FuncCall.make(
                    name="other_func",
                    args=[
                        nodes.Field.make("f2"),
                        n.p(nodes.Field.make("f3")),
                    ],
                ),
            ],
        ),
        nodes.FuncCall.make(
            name="other_func",
            args=[
                nodes.Field.make("f2"),
                n.p(nodes.Field.make("f3")),
            ],
        ),
    ]


def test_used_fields():
    assert used_fields(n.formula(nodes.Field.make("f1"))) == [
        nodes.Field.make("f1"),
    ]
    assert used_fields(
        n.formula(
            nodes.FuncCall.make(
                name="func",
                args=[
                    nodes.Field.make("f1"),
                    nodes.FuncCall.make(
                        name="other_func",
                        args=[
                            nodes.Field.make("f2"),
                            n.p(nodes.Field.make("f3")),
                        ],
                    ),
                ],
            )
        )
    ) == [
        nodes.Field.make("f1"),
        nodes.Field.make("f2"),
        nodes.Field.make("f3"),
    ]


def test_is_constant_expression():
    env = InspectionEnvironment()
    assert is_constant_expression(nodes.LiteralString.make("qwe"), env=env)
    assert is_constant_expression(nodes.Null(), env=env)
    assert not is_constant_expression(nodes.FuncCall.make(name="func", args=[nodes.LiteralString.make("qwe")]), env=env)
    assert is_constant_expression(n.p(nodes.LiteralString.make("qwe")), env=env)
    assert not is_constant_expression(nodes.FuncCall.make(name="func", args=[nodes.LiteralString.make("qwe")]), env=env)
    assert not is_constant_expression(
        n.p(nodes.FuncCall.make(name="func", args=[nodes.LiteralString.make("qwe")])), env=env
    )


def test_is_is_bound_only_to():
    dim_1 = nodes.Field.make("Dim Field")
    dim_2 = nodes.FuncCall.make(
        name="+",
        args=[
            nodes.Field.make("Other Field"),
            nodes.FuncCall.make(
                name="dim_func",
                args=[
                    nodes.Field.make("Third Field"),
                ],
            ),
        ],
    )
    allow_nodes = NodeSet([dim_1, dim_2])
    assert is_bound_only_to(nodes.LiteralString.make("qwe"), allow_nodes=allow_nodes)
    assert is_bound_only_to(n.p(nodes.LiteralString.make("qwe")), allow_nodes=allow_nodes)
    assert is_bound_only_to(n.p(dim_1), allow_nodes=allow_nodes)
    assert is_bound_only_to(n.p(dim_2), allow_nodes=allow_nodes)
    assert is_bound_only_to(
        nodes.FuncCall.make(name="func", args=[nodes.LiteralString.make("qwe")]), allow_nodes=allow_nodes
    )
    assert is_bound_only_to(
        n.p(nodes.FuncCall.make(name="func", args=[nodes.LiteralString.make("qwe"), dim_1])),
        allow_nodes=allow_nodes,
    )
    assert is_bound_only_to(
        nodes.FuncCall.make(
            name="func",
            args=[
                nodes.LiteralString.make("qwe"),
                n.p(nodes.LiteralString.make("qwe")),
                dim_2,
            ],
        ),
        allow_nodes=allow_nodes,
    )
    assert not is_bound_only_to(nodes.Field.make(name="my field"), allow_nodes=allow_nodes)
    assert not is_bound_only_to(n.p(nodes.Field.make(name="my field")), allow_nodes=allow_nodes)
    assert not is_bound_only_to(
        nodes.FuncCall.make(
            name="func",
            args=[
                nodes.LiteralString.make("qwe"),
                nodes.Field.make(name="my field"),
                dim_2,
            ],
        ),
        allow_nodes=allow_nodes,
    )


def test_is_aggregate_expression():
    env = InspectionEnvironment()
    assert is_aggregate_expression(nodes.FuncCall.make(name="avg", args=[nodes.LiteralString.make("qwe")]), env=env)
    assert is_aggregate_expression(nodes.FuncCall.make(name="avg", args=[nodes.Field.make("my field")]), env=env)
    assert is_aggregate_expression(
        n.p(nodes.FuncCall.make(name="avg", args=[nodes.LiteralString.make("qwe")])), env=env
    )
    assert not is_aggregate_expression(nodes.LiteralString.make("qwe"), env=env)
    assert not is_aggregate_expression(n.p(nodes.LiteralString.make("qwe")), env=env)
    assert not is_aggregate_expression(
        nodes.FuncCall.make(name="func", args=[nodes.LiteralString.make("qwe")]), env=env
    )
    assert not is_aggregate_expression(
        nodes.WindowFuncCall.make(
            name="sum", args=[nodes.Field.make("my field")], grouping=nodes.WindowGroupingTotal()
        ),
        env=env,
    )


def test_is_aggregated_above_sub_node():
    sub_node = n.field("My Field")
    # Simple aggregation
    assert is_aggregated_above_sub_node(
        node=n.func.MAX(sub_node),
        sub_node=sub_node,
    )
    # Aggregation with custom LOD
    assert is_aggregated_above_sub_node(
        node=n.func.MAX(sub_node, lod=n.fixed()),
        sub_node=sub_node,
    )
    # Aggregation + wrapper
    assert is_aggregated_above_sub_node(
        node=n.func.MAX(n.func.CONCAT(sub_node)),
        sub_node=sub_node,
    )
    # Wrapper + aggregation + wrapper
    assert is_aggregated_above_sub_node(
        node=n.func.FLOOR(n.func.MAX(n.func.CONCAT(sub_node))),
        sub_node=sub_node,
    )
    # Same, but by index
    assert is_aggregated_above_sub_node(
        node=n.func.FLOOR(n.func.MAX(n.func.CONCAT(sub_node))),
        index=NodeHierarchyIndex.make(0, 0, 0),
    )
    # No aggregations
    assert not is_aggregated_above_sub_node(
        node=n.func.FLOOR(n.func.CONCAT(sub_node)),
        sub_node=sub_node,
    )
    # Same, but by index
    assert not is_aggregated_above_sub_node(
        node=n.func.FLOOR(n.func.CONCAT(sub_node)),
        index=NodeHierarchyIndex.make(0, 0),
    )
    # Window function in this case, not agg
    assert not is_aggregated_above_sub_node(
        node=n.func.FLOOR(n.wfunc.MAX(n.func.CONCAT(sub_node))),
        sub_node=sub_node,
    )
    # Formula contains an aggregation, but it doesn't wrap sub_node
    assert not is_aggregated_above_sub_node(
        node=n.func.CONCAT(sub_node, n.func.MAX(copy.copy(sub_node))),
        sub_node=sub_node,
    )


def test_is_window_expression():
    env = InspectionEnvironment()
    assert is_window_expression(
        nodes.WindowFuncCall.make(
            name="rank",
            args=[nodes.LiteralString.make("qwe")],
            grouping=nodes.WindowGroupingTotal(),
        ),
        env=env,
    )
    assert is_window_expression(
        nodes.FuncCall.make(
            name="sum",
            args=[
                nodes.WindowFuncCall.make(
                    name="rank",
                    args=[nodes.Field.make("my field")],
                    grouping=nodes.WindowGroupingTotal(),
                )
            ],
        ),
        env=env,
    )
    assert is_window_expression(
        n.p(nodes.WindowFuncCall.make(name="rank", args=[nodes.LiteralString.make("qwe")])), env=env
    )
    assert not is_window_expression(nodes.FuncCall.make(name="avg", args=[nodes.LiteralString.make("qwe")]), env=env)
    assert not is_window_expression(nodes.FuncCall.make(name="func", args=[nodes.Field.make("my field")]), env=env)
    assert not is_window_expression(
        n.p(nodes.FuncCall.make(name="avg", args=[nodes.LiteralString.make("qwe")])), env=env
    )


def test_is_query_fork_expression():
    env = InspectionEnvironment()
    assert not is_query_fork_expression(
        nodes.WindowFuncCall.make(
            name="rank",
            args=[nodes.LiteralString.make("qwe")],
            grouping=nodes.WindowGroupingTotal(),
        ),
        env=env,
    )
    assert is_query_fork_expression(
        n.func.OTHERFUNC(
            n.fork(
                joining=n.joining(conditions=[n.self_condition(n.field("city"))]),
                result_expr=n.func.MY_FUNC(n.field("My Field")),
            ),
            n.field("Other Field"),
        ),
        env=env,
    )


def test_get_query_fork_nesting_level():
    env = InspectionEnvironment()
    assert (
        get_query_fork_nesting_level(
            n.formula(
                n.func.SOMEFUNC(
                    n.fork(
                        joining=n.joining(conditions=[n.self_condition(n.field("date"))]),
                        result_expr=n.func.OTHERFUNC(
                            n.fork(
                                joining=n.joining(conditions=[n.self_condition(n.field("city"))]),
                                result_expr=n.func.MY_FUNC(n.field("My Field")),
                            ),
                            n.field("Other Field"),
                        ),
                    )
                )
            ),
            env=env,
        )
        == 2
    )


def test_infer_data_type():
    env = InspectionEnvironment()
    field_types = {
        "My Field": DataType.STRING,
        "One Field": DataType.INTEGER,
        "Other Field": DataType.INTEGER,
    }
    assert (
        infer_data_type(
            nodes.Field.make("Other Field"),
            field_types=field_types,
            env=env,
        )
        == DataType.INTEGER
    )
    assert (
        infer_data_type(
            n.case(
                n.if_(n.func("==")(n.field("My Field"), n.p("qwe"))).then(True).else_(False),
            )
            .whens(
                n.when(True).then(3.4),
            )
            .else_(5.6),
            field_types=field_types,
            env=env,
        )
        == DataType.FLOAT
    )
    assert infer_data_type(n.wfunc.RANK(n.field("Other Field")), field_types=field_types, env=env) == DataType.INTEGER
    assert (
        infer_data_type(n.binary("==", n.field("One Field"), n.field("Other Field")), field_types=field_types, env=env)
        == DataType.BOOLEAN
    )


def test_enumerate():
    nhi = NodeHierarchyIndex.make
    main_func = nodes.FuncCall.make(
        name="concat",
        args=[
            nodes.Field.make("Other Field"),
            nodes.Field.make("Green Field"),
            nodes.LiteralString.make("qwerty"),
        ],
    )
    formula = n.formula(main_func)
    assert list(enumerate_fields(formula)) == [
        (nhi(0, 0), nodes.Field.make("Other Field")),
        (nhi(0, 1), nodes.Field.make("Green Field")),
    ]


def test_iter_operation_calls():
    names = [
        func_node.name
        for func_node in iter_operation_calls(
            n.func.CONCAT(
                n.func.INT(n.field("qwerty")),
                n.binary(">", n.field("f1"), n.field("f2")),
            )
        )
    ]
    assert names == ["concat", "int", ">"]

lib/dl_formula/dl_formula_tests/unit/inspect/test_expression.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_formula.inspect.function import can_be_aggregate


def test_can_be_aggregate():
    assert can_be_aggregate(name="sum")
    assert not can_be_aggregate(name="greatest")

lib/dl_formula/dl_formula_tests/unit/inspect/test_function.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_formula.mutation.general import (
    ConvertBlocksToFunctionsMutation,
    IgnoreParenthesisWrapperMutation,
    OptimizeConstAndOrMutation,
    OptimizeConstComparisonMutation,
    OptimizeConstFuncMutation,
)
from dl_formula.mutation.mutation import apply_mutations
from dl_formula.shortcuts import n


def test_ignore_parenthesis_wrapper_mutation():
    formula_obj = n.formula(n.func.my_func(args=[n.p(n.func.FUNC(args=[]))]))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            IgnoreParenthesisWrapperMutation(),
        ],
    )
    assert formula_obj == n.formula(n.func.my_func(args=[n.func.FUNC(args=[])]))


def test_convert_blocks_to_functions_mutation():
    formula_obj = n.formula(n.if_(n.field("this")).then(n.field("that")).else_(n.field("smth_else")))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            ConvertBlocksToFunctionsMutation(),
        ],
    )
    assert formula_obj == n.formula(n.func.IF(args=[n.field("this"), n.field("that"), n.field("smth_else")]))


def test_optimize_const_comparison_mutation():
    formula_obj = n.formula(n.binary("==", left=n.lit(3), right=n.lit(3)))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstComparisonMutation(),
        ],
    )
    assert formula_obj == n.formula(n.lit(True))

    formula_obj = n.formula(n.binary("==", left=n.lit(2), right=n.lit(3)))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstComparisonMutation(),
        ],
    )
    assert formula_obj == n.formula(n.lit(False))

    formula_obj = n.formula(n.binary("!=", left=n.lit(2), right=n.lit(2)))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstComparisonMutation(),
        ],
    )
    assert formula_obj == n.formula(n.lit(False))

    formula_obj = n.formula(n.binary("!=", left=n.lit(2), right=n.lit(3)))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstComparisonMutation(),
        ],
    )
    assert formula_obj == n.formula(n.lit(True))


def test_optimize_const_and_or_mutation():
    formula_obj = n.formula(n.binary("and", left=n.lit(True), right=n.field("my field")))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstAndOrMutation(),
        ],
    )
    assert formula_obj == n.formula(n.field("my field"))

    formula_obj = n.formula(n.binary("and", left=n.lit(False), right=n.field("my field")))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstAndOrMutation(),
        ],
    )
    assert formula_obj == n.formula(n.lit(False))

    formula_obj = n.formula(n.binary("or", left=n.field("my field"), right=n.lit(True)))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstAndOrMutation(),
        ],
    )
    assert formula_obj == n.formula(n.lit(True))

    formula_obj = n.formula(n.binary("or", left=n.field("my field"), right=n.lit(False)))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstAndOrMutation(),
        ],
    )
    assert formula_obj == n.formula(n.field("my field"))


def test_optimize_if_mutation():
    # Check removal of false conditions
    formula_obj = n.formula(
        n.func.IF(
            n.field("cond 1"),
            n.field("then field 1"),
            n.lit(False),
            n.field("then field 2"),
            n.lit(False),
            n.field("then field 3"),
            n.field("cond 4"),
            n.field("then field 4"),
            n.field("else field"),
        )
    )
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstFuncMutation(),
        ],
    )
    assert formula_obj == n.formula(
        n.func.IF(
            n.field("cond 1"),
            n.field("then field 1"),
            n.field("cond 4"),
            n.field("then field 4"),
            n.field("else field"),
        )
    )

    # Check only else
    formula_obj = n.formula(
        n.func.IF(
            n.lit(False),
            n.field("then field 1"),
            n.lit(False),
            n.field("then field 2"),
            n.field("else field"),
        )
    )
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstFuncMutation(),
        ],
    )
    assert formula_obj == n.formula(n.field("else field"))

    # Check true condition
    formula_obj = n.formula(
        n.func.IF(
            n.field("cond 1"),
            n.field("then field 1"),
            n.lit(False),
            n.field("then field 2"),
            n.lit(True),
            n.field("then field 3"),
            n.field("cond 4"),
            n.field("then field 4"),
            n.field("else field"),
        )
    )
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstFuncMutation(),
        ],
    )
    assert formula_obj == n.formula(n.field("then field 3"))


def test_optimize_case_mutation():
    # Check removal of false conditions
    formula_obj = n.formula(
        n.func.CASE(
            n.lit("what"),
            n.field("when 1"),
            n.field("then field 1"),
            n.lit("qwerty"),
            n.field("then field 2"),
            n.lit("uiop"),
            n.field("then field 3"),
            n.field("when 4"),
            n.field("then field 4"),
            n.field("else field"),
        )
    )
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstFuncMutation(),
        ],
    )
    assert formula_obj == n.formula(
        n.func.CASE(
            n.lit("what"),
            n.field("when 1"),
            n.field("then field 1"),
            n.field("when 4"),
            n.field("then field 4"),
            n.field("else field"),
        )
    )

    # Check only else
    formula_obj = n.formula(
        n.func.CASE(
            n.lit("what"),
            n.lit("qwerty"),
            n.field("then field 1"),
            n.lit("uiop"),
            n.field("then field 2"),
            n.field("else field"),
        )
    )
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstFuncMutation(),
        ],
    )
    assert formula_obj == n.formula(n.field("else field"))

    # Check match
    formula_obj = n.formula(
        n.func.CASE(
            n.lit("what"),
            n.field("when 1"),
            n.field("then field 1"),
            n.lit("qwerty"),
            n.field("then field 2"),
            n.lit("what"),
            n.field("then field 3"),
            n.field("when 4"),
            n.field("then field 4"),
            n.field("else field"),
        )
    )
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            OptimizeConstFuncMutation(),
        ],
    )
    assert formula_obj == n.formula(n.field("then field 3"))

lib/dl_formula/dl_formula_tests/unit/mutation/test_general.py
datalens-tech/datalens-backend
from __future__ import annotations

import dl_formula.core.fork_nodes as fork_nodes
from dl_formula.mutation.lod import ExtAggregationToQueryForkMutation
from dl_formula.mutation.mutation import apply_mutations
from dl_formula.shortcuts import n


def test_lod_aggregation_to_query_fork_mutation():
    mutations = [
        ExtAggregationToQueryForkMutation(global_dimensions=[n.field("f1")]),
    ]

    formula_obj = n.formula(
        n.func.SUM(
            args=[n.func.AVG(args=[], lod=n.include(n.field("f2")))],
            lod=n.include(),
        ),
    )
    formula_obj = apply_mutations(formula_obj, mutations=mutations)
    assert formula_obj == n.formula(
        n.fork(
            join_type=fork_nodes.JoinType.inner,
            result_expr=n.func.SUM(
                args=[
                    n.fork(
                        join_type=fork_nodes.JoinType.inner,
                        result_expr=n.func.AVG(
                            args=[],
                            lod=n.fixed(n.field("f1"), n.field("f2")),
                        ),
                        joining=n.joining(
                            conditions=[
                                n.self_condition(expr=n.field("f1")),
                                n.self_condition(expr=n.field("f2")),
                            ]
                        ),
                        lod=n.fixed(n.field("f1"), n.field("f2")),
                    ),
                ],
                lod=n.fixed(n.field("f1")),
            ),
            joining=n.joining(
                conditions=[
                    n.self_condition(expr=n.field("f1")),
                ]
            ),
            lod=n.fixed(n.field("f1")),
        ),
    )

lib/dl_formula/dl_formula_tests/unit/mutation/test_lod.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Tuple

import dl_formula.core.nodes as nodes
from dl_formula.mutation.mutation import (
    FormulaMutation,
    apply_mutations,
)
from dl_formula.shortcuts import n


class PrefixFunctionMutation(FormulaMutation):
    """Prepend `s` to function names"""

    def __init__(self, prefix: str):
        self._prefix = prefix

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        return isinstance(node, nodes.FuncCall)

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.FuncCall)
        return nodes.FuncCall.make(name=self._prefix + old.name, args=old.args)


class PostfixFieldMutation(FormulaMutation):
    """Append `s` to field names"""

    def __init__(self, postfix: str):
        self._postfix = postfix

    def match_node(self, node: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]) -> bool:
        return isinstance(node, nodes.Field)

    def make_replacement(
        self, old: nodes.FormulaItem, parent_stack: Tuple[nodes.FormulaItem, ...]
    ) -> nodes.FormulaItem:
        assert isinstance(old, nodes.Field)
        return nodes.Field.make(name=old.name + self._postfix)


def test_multiple_mutations():
    assert apply_mutations(
        n.formula(n.func.MYFUNC(n.field("field"))),
        mutations=[
            # note that they are applied in the same order as they are listed,
            # so w need to reverse them here to get a sensible result
            PrefixFunctionMutation("c"),
            PrefixFunctionMutation("b"),
            PrefixFunctionMutation("a"),
            PostfixFieldMutation("1"),
            PostfixFieldMutation("2"),
        ],
    ) == n.formula(n.func.ABCMYFUNC(n.field("field12")))

lib/dl_formula/dl_formula_tests/unit/mutation/test_muation.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_formula.mutation.lookup import (
    LookupDefaultBfbMutation,
    LookupFunctionToQueryForkMutation,
)
from dl_formula.mutation.mutation import apply_mutations
from dl_formula.shortcuts import n


def test_ago_to_query_fork_mutation():
    formula_obj = n.formula(n.func.SOMEFUNC(n.func.AGO(n.func.SUM(n.field("f1")), n.field("date"))))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            LookupFunctionToQueryForkMutation(
                global_dimensions=[
                    n.field("date"),
                    n.field("city"),
                ],
                allow_empty_dimensions=False,
            ),
        ],
    )
    assert formula_obj == n.formula(
        n.func.SOMEFUNC(
            n.fork(
                result_expr=n.func.SUM(n.field("f1")),
                joining=n.joining(
                    conditions=[
                        n.bin_condition(
                            expr=n.field("date"),
                            fork_expr=n.func.DATEADD(n.field("date")),
                        ),
                        n.self_condition(
                            expr=n.field("city"),
                        ),
                    ],
                ),
                lod=n.inherited(),
            )
        ),
    )


def test_lookup_default_bfb_mutation():
    formula_obj = n.formula(n.func.SOMEFUNC(n.func.AGO(n.func.SUM(n.field("f1")), n.field("date"))))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[
            LookupDefaultBfbMutation(),
        ],
    )
    assert formula_obj == n.formula(
        n.func.SOMEFUNC(
            n.func.AGO(
                n.func.SUM(n.field("f1")),
                n.field("date"),
                before_filter_by=["date"],
            ),
        )
    )

lib/dl_formula/dl_formula_tests/unit/mutation/test_subquery_mutations.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.mutation.bfb import RemapBfbMutation
from dl_formula.mutation.mutation import apply_mutations
from dl_formula.mutation.window import (
    AmongToWithinGroupingMutation,
    DefaultWindowOrderingMutation,
    IgnoreExtraWithinGroupingMutation,
)
from dl_formula.shortcuts import n


def test_among_to_within_mutation():
    formula_obj = n.formula(n.wfunc.RANK(n.func.SUM(n.field("f3")), among=[n.field("f2")]))
    formula_obj = apply_mutations(
        formula_obj, mutations=[AmongToWithinGroupingMutation(global_dimensions=[n.field("f1"), n.field("f2")])]
    )
    assert formula_obj == n.formula(n.wfunc.RANK(n.func.SUM(n.field("f3")), within=[n.field("f1")]))


def test_ignore_extra_within_mutation():
    inspect_env = InspectionEnvironment()
    formula_obj = n.formula(n.wfunc.RANK(n.func.SUM(n.field("f1")), within=[n.field("f2"), n.field("f3")]))
    formula_obj = apply_mutations(
        formula_obj,
        mutations=[IgnoreExtraWithinGroupingMutation(global_dimensions=[n.field("f2")], inspect_env=inspect_env)],
    )
    assert formula_obj == n.formula(n.wfunc.RANK(n.func.SUM(n.field("f1")), within=[n.field("f2")]))


def test_default_window_ordering_mutation_basic():
    formula_obj = n.formula(n.wfunc.RSUM(n.func.SUM(n.field("f3")), total=True))  # RSUM requires default ordering
    formula_obj = apply_mutations(
        formula_obj, mutations=[DefaultWindowOrderingMutation(default_order_by=[n.field("f1"), n.field("f2")])]
    )
    assert formula_obj == n.formula(
        n.wfunc.RSUM(n.func.SUM(n.field("f3")), total=True, order_by=[n.field("f1"), n.field("f2")])
    )


def test_default_window_ordering_mutation_with_order_by_override():
    formula_obj = n.formula(
        n.wfunc.RSUM(n.func.SUM(n.field("f3")), total=True, order_by=[n.desc(n.field("f2")), n.field("f3")])
    )
    formula_obj = apply_mutations(
        formula_obj, mutations=[DefaultWindowOrderingMutation(default_order_by=[n.field("f1"), n.field("f2")])]
    )
    assert formula_obj == n.formula(
        n.wfunc.RSUM(
            n.func.SUM(n.field("f3")),
            total=True,
            order_by=[
                n.desc(n.field("f2")),
                n.field("f3"),
                n.field("f1"),
            ],
        ),
    )


def test_remap_before_filter_by_mutation():
    formula_obj = n.formula(
        n.wfunc.FUNC(
            before_filter_by=(),
            args=[
                n.wfunc.FUNC(
                    before_filter_by=("A",),
                    args=[
                        n.func.FUNC(before_filter_by=("Q", "R"), args=[]),
                    ],
                ),
            ],
        ),
    )
    formula_obj = apply_mutations(
        formula_obj, mutations=[RemapBfbMutation(name_mapping={"A": "a1", "Q": "a2", "R": "a3"})]
    )
    assert formula_obj == n.formula(
        n.wfunc.FUNC(
            before_filter_by=(),
            args=[
                n.wfunc.FUNC(
                    before_filter_by=("a1",),
                    args=[
                        n.func.FUNC(before_filter_by=("a2", "a3"), args=[]),
                    ],
                ),
            ],
        ),
    )

lib/dl_formula/dl_formula_tests/unit/mutation/test_window_mutations.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime

import pytest

from dl_formula.core import exc
from dl_formula.core.nodes import (
    BeforeFilterBy,
    Binary,
    CaseBlock,
    ExcludeLodSpecifier,
    ExpressionList,
    Field,
    FixedLodSpecifier,
    FuncCall,
    IfBlock,
    IfPart,
    IgnoreDimensions,
    IncludeLodSpecifier,
    LiteralBoolean,
    LiteralDate,
    LiteralFloat,
    LiteralGenericDatetime,
    LiteralInteger,
    LiteralString,
    Null,
    OrderAscending,
    OrderDescending,
    Ordering,
    Ternary,
    Unary,
    WhenPart,
    WindowFuncCall,
    WindowGroupingAmong,
    WindowGroupingTotal,
    WindowGroupingWithin,
)
from dl_formula.shortcuts import n


def test_keyword_case_insensitivity(parser):
    upper_formula = (
        'IF [n0]*5 THEN MY_FUNC(NULL, NOT [n2], "qwerty") ELSEIF ISNULL([n3]) THEN "result" '
        'ELSE CASE [n4] WHEN 5 THEN "five" WHEN 6 THEN "seven" ELSE "forty two" END END'
    )
    lower_formula = (
        'if [n0]*5 then MY_FUNC(null, not [n2], "qwerty") elseif ISNULL([n3]) then "result" '
        'else case [n4] when 5 then "five" when 6 then "seven" else "forty two" end end'
    )
    assert parser.parse(upper_formula) == parser.parse(lower_formula)


def test_position_getter(parser):
    formula = parser.parse(
        "IF [n0] * 5 " 'THEN MY_FUNC(NULL, NOT [n2], "qwerty") ' "ELSE [SOMETHING] AND [Q1] IN (12, 34) " "END"
    )
    assert formula.get_by_pos(1).__class__ is IfPart
    assert formula.get_by_pos(5).__class__ is Field
    assert formula.get_by_pos(8).__class__ is Binary
    assert formula.get_by_pos(20).__class__ is FuncCall
    assert formula.get_by_pos(30).__class__ is FuncCall
    assert formula.get_by_pos(33).__class__ is Unary
    assert formula.get_by_pos(47).__class__ is LiteralString
    assert formula.get_by_pos(84).__class__ is ExpressionList


def test_boolean_w_incomplete_grammar(parser):
    formula = parser.parse("true AND false")
    assert formula.get_by_pos(3).__class__ is LiteralBoolean
    assert formula.get_by_pos(11).__class__ is LiteralBoolean


def test_parse_multiline(parser):
    assert parser.parse("AVG(\n    1\n)") == n.formula(FuncCall.make("avg", [LiteralInteger.make(1)]))


def test_parse_boolean(parser):
    assert parser.parse("TRUE") == n.formula(LiteralBoolean.make(True))
    assert type(parser.parse("TRUE").expr.value) is bool
    assert parser.parse("FALSE") == n.formula(LiteralBoolean.make(False))
    assert type(parser.parse("FALSE").expr.value) is bool


def test_parse_integer(parser):
    assert parser.parse("123") == n.formula(LiteralInteger.make(123))
    assert type(parser.parse("123").expr.value) is int
    assert parser.parse("-123") == n.formula(LiteralInteger.make(-123))
    assert type(parser.parse("-123").expr.value) is int


def test_parse_float(parser):
    parse = parser.parse
    assert parse("123.0") == n.formula(LiteralFloat.make(123.0))
    assert type(parse("123.0").expr.value) is float
    assert parse("-123.0") == n.formula(LiteralFloat.make(-123.0))
    assert type(parse("-123.0").expr.value) is float
    assert parse("1e0") == n.formula(LiteralFloat.make(1.0))
    assert parse("1e2") == n.formula(LiteralFloat.make(100.0))
    assert parse("1e+2") == n.formula(LiteralFloat.make(100.0))
    assert parse("1e-2") == n.formula(LiteralFloat.make(0.01))
    assert parse("1.1e2") == n.formula(LiteralFloat.make(110.0))
    assert parse("1.1e+2") == n.formula(LiteralFloat.make(110.0))
    assert parse("1.1e-2") == n.formula(LiteralFloat.make(0.011))


def test_parse_null(parser):
    assert parser.parse("NULL") == n.formula(Null())


def test_parse_string(parser):
    assert parser.parse('"qwerty with spaces, \\", \\\' and -()"') == n.formula(
        LiteralString.make("qwerty with spaces, \", ' and -()")
    )
    assert parser.parse("'qwerty with spaces, \\\", \\', \\n, \\r, \\t and -()'") == n.formula(
        LiteralString.make("qwerty with spaces, \", ', \n, \r, \t and -()")
    )


def test_parse_field(parser):
    assert parser.parse("[my_field]") == n.formula(Field.make("my_field"))


def test_parse_parentheses(parser):
    assert parser.parse('("expr with parentheses")') == n.formula(n.p(LiteralString.make("expr with parentheses")))


def test_parse_date(parser):
    assert parser.parse("#2019-02-12#") == n.formula(LiteralDate.make(datetime.date(2019, 2, 12)))


def test_parse_datetime(parser):
    assert parser.parse("#2019-02-12T#") == n.formula(
        LiteralGenericDatetime.make(datetime.datetime(2019, 2, 12, 0, 0, 0))
    )
    assert parser.parse("#2019-02-12 12#") == n.formula(
        LiteralGenericDatetime.make(datetime.datetime(2019, 2, 12, 12, 0, 0))
    )
    assert parser.parse("#2019-02-12T12:34#") == n.formula(
        LiteralGenericDatetime.make(datetime.datetime(2019, 2, 12, 12, 34, 0))
    )
    assert parser.parse("#2019-02-12t12:34:56#") == n.formula(
        LiteralGenericDatetime.make(datetime.datetime(2019, 2, 12, 12, 34, 56))
    )


def test_parse_generic_datetime(parser):
    assert parser.parse("##2019-02-12##") == n.formula(LiteralDate.make(datetime.date(2019, 2, 12)))
    assert parser.parse("##2019-02-12T##") == n.formula(
        LiteralGenericDatetime.make(datetime.datetime(2019, 2, 12, 0, 0, 0))
    )
    assert parser.parse("##2019-02-12 12##") == n.formula(
        LiteralGenericDatetime.make(datetime.datetime(2019, 2, 12, 12, 0, 0))
    )
    assert parser.parse("##2019-02-12T12:34##") == n.formula(
        LiteralGenericDatetime.make(datetime.datetime(2019, 2, 12, 12, 34, 0))
    )
    assert parser.parse("##2019-02-12t12:34:56##") == n.formula(
        LiteralGenericDatetime.make(datetime.datetime(2019, 2, 12, 12, 34, 56))
    )


def test_parse_operations(parser):
    parse = parser.parse
    # ^
    assert parse("[left] ^ [right]") == n.formula(Binary.make("^", Field.make("left"), Field.make("right")))
    assert parse("[left]^-[right]") == n.formula(
        Binary.make("^", Field.make("left"), Unary.make("neg", Field.make("right")))
    )
    assert parse("[left] ^ NOT [right]") == n.formula(
        Binary.make("^", Field.make("left"), Unary.make("not", Field.make("right")))
    )
    assert parse("[left] IS NULL ^ [right]") == n.formula(
        Binary.make("^", Unary.make("isnull", Field.make("left")), Field.make("right"))
    )

    # - (negate)
    assert parse("-[expr]") == n.formula(Unary.make("neg", Field.make("expr")))
    assert parse(" -  - [expr]") == n.formula(Unary.make("neg", Unary.make("neg", Field.make("expr"))))
    assert parse("-[expr]^2") == n.formula(
        Unary.make("neg", Binary.make("^", Field.make("expr"), LiteralInteger.make(2)))
    )

    # *
    assert parse("[left] * [right]") == n.formula(Binary.make("*", Field.make("left"), Field.make("right")))
    assert parse("[left]*[right]") == n.formula(Binary.make("*", Field.make("left"), Field.make("right")))
    assert parse("123*123") == n.formula(Binary.make("*", LiteralInteger.make(123), LiteralInteger.make(123)))
    assert parse("[left]*-[right]") == n.formula(
        Binary.make("*", Field.make("left"), Unary.make("neg", Field.make("right")))
    )
    assert parse("[left] * NOT [right]") == n.formula(
        Binary.make("*", Field.make("left"), Unary.make("not", Field.make("right")))
    )
    assert parse("[left] IS NULL * [right]") == n.formula(
        Binary.make("*", Unary.make("isnull", Field.make("left")), Field.make("right"))
    )
    # /
    assert parse("[left] / [right]") == n.formula(Binary.make("/", Field.make("left"), Field.make("right")))
    assert parse("[left] / -[right]") == n.formula(
        Binary.make("/", Field.make("left"), Unary.make("neg", Field.make("right")))
    )
    assert parse("[left] / NOT [right]") == n.formula(
        Binary.make("/", Field.make("left"), Unary.make("not", Field.make("right")))
    )
    assert parse("[left] IS NULL / [right]") == n.formula(
        Binary.make("/", Unary.make("isnull", Field.make("left")), Field.make("right"))
    )
    # %
    assert parse("[left] % [right]") == n.formula(Binary.make("%", Field.make("left"), Field.make("right")))
    assert parse("[left] % -[right]") == n.formula(
        Binary.make("%", Field.make("left"), Unary.make("neg", Field.make("right")))
    )
    assert parse("[left] % NOT [right]") == n.formula(
        Binary.make("%", Field.make("left"), Unary.make("not", Field.make("right")))
    )
    assert parse("[left] IS NULL % [right]") == n.formula(
        Binary.make("%", Unary.make("isnull", Field.make("left")), Field.make("right"))
    )
    # chained
    assert parse("[n1] / [n2] * [n3] % [n1]") == n.formula(
        Binary.make(
            "%",
            Binary.make(
                "*",
                Binary.make("/", Field.make("n1"), Field.make("n2")),
                Field.make("n3"),
            ),
            Field.make("n1"),
        )
    )
    # priority: negation vs * and /
    assert parse("-[a] * [b]") == n.formula(
        Binary.make(
            "*",
            Unary.make("neg", Field.make("a")),
            Field.make("b"),
        )
    )
    assert parse("[c] / -[a] * [b]") == n.formula(
        Binary.make(
            "*",
            Binary.make("/", Field.make("c"), Unary.make("neg", Field.make("a"))),
            Field.make("b"),
        )
    )
    # priority: * vs ^
    assert parse("[n1] * [n2] ^ [n3]") == n.formula(
        Binary.make(
            "*",
            Field.make("n1"),
            Binary.make("^", Field.make("n2"), Field.make("n3")),
        )
    )

    # +
    assert parse("[left] + [right]") == n.formula(Binary.make("+", Field.make("left"), Field.make("right")))
    assert parse("[left] + NOT [right]") == n.formula(
        Binary.make("+", Field.make("left"), Unary.make("not", Field.make("right")))
    )
    assert parse("[left] IS NULL + [right]") == n.formula(
        Binary.make("+", Unary.make("isnull", Field.make("left")), Field.make("right"))
    )
    # -
    assert parse("[left] - [right]") == n.formula(Binary.make("-", Field.make("left"), Field.make("right")))
    assert parse("[left] - NOT [right]") == n.formula(
        Binary.make("-", Field.make("left"), Unary.make("not", Field.make("right")))
    )
    assert parse("[left] IS NULL - [right]") == n.formula(
        Binary.make("-", Unary.make("isnull", Field.make("left")), Field.make("right"))
    )
    # chained +/-
    assert parse("[n1] - [n2] + [n3]") == n.formula(
        Binary.make(
            "+",
            Binary.make("-", Field.make("n1"), Field.make("n2")),
            Field.make("n3"),
        )
    )
    # priority: + vs *
    assert parse("[n1] + [n2] * [n3]") == n.formula(
        Binary.make(
            "+",
            Field.make("n1"),
            Binary.make("*", Field.make("n2"), Field.make("n3")),
        )
    )

    # is ...
    assert parse("[left] is null") == n.formula(Unary.make("isnull", Field.make("left")))
    assert parse("[left] is not null") == n.formula(Unary.make("not", Unary.make("isnull", Field.make("left"))))
    assert parse("[left] is true") == n.formula(Unary.make("istrue", Field.make("left")))
    assert parse("[left] is not true") == n.formula(Unary.make("not", Unary.make("istrue", Field.make("left"))))
    assert parse("[left] is false") == n.formula(Unary.make("isfalse", Field.make("left")))
    assert parse("[left] is not false") == n.formula(Unary.make("not", Unary.make("isfalse", Field.make("left"))))
    # priority: is vs +
    assert parse("[n1] + [n2] is null") == n.formula(
        Unary.make("isnull", Binary.make("+", Field.make("n1"), Field.make("n2")))
    )

    # like
    assert parse("[left] like [right]") == n.formula(Binary.make("like", Field.make("left"), Field.make("right")))
    assert parse("[left] not like [right]") == n.formula(
        Binary.make("notlike", Field.make("left"), Field.make("right"))
    )
    # priority: like vs is
    assert parse("[n1] like [n2] is null") == n.formula(
        Binary.make("like", Field.make("n1"), Unary.make("isnull", Field.make("n2")))
    )

    # =
    assert parse("[left] = [right]") == n.formula(Binary.make("==", Field.make("left"), Field.make("right")))
    assert parse("[left]=[right]") == n.formula(Binary.make("==", Field.make("left"), Field.make("right")))
    assert parse("123=123") == n.formula(Binary.make("==", LiteralInteger.make(123), LiteralInteger.make(123)))
    assert parse("[left] = NOT [right]") == n.formula(
        Binary.make("==", Field.make("left"), Unary.make("not", Field.make("right")))
    )
    # !=
    assert parse("[left] != [right]") == n.formula(Binary.make("!=", Field.make("left"), Field.make("right")))
    assert parse("[left] <> [right]") == n.formula(Binary.make("!=", Field.make("left"), Field.make("right")))
    # >
    assert parse("[left] > [right]") == n.formula(Binary.make(">", Field.make("left"), Field.make("right")))
    # >=
    assert parse("[left] >= [right]") == n.formula(Binary.make(">=", Field.make("left"), Field.make("right")))
    # <
    assert parse("[left] < [right]") == n.formula(Binary.make("<", Field.make("left"), Field.make("right")))
    # <=
    assert parse("[left] <= [right]") == n.formula(Binary.make("<=", Field.make("left"), Field.make("right")))
    # chained eq/neq
    assert parse("[n1] < [n2] = [n3] > [n4]") == n.formula(
        Binary.make(
            "and",
            Binary.make(
                "and",
                Binary.make("<", Field.make("n1"), Field.make("n2")),
                Binary.make("==", Field.make("n2"), Field.make("n3")),
            ),
            Binary.make(">", Field.make("n3"), Field.make("n4")),
        )
    )
    # priority: <> vs or
    assert parse("[n1] < [n2] or [n3] > [n4]") == n.formula(
        Binary.make(
            "or",
            Binary.make("<", Field.make("n1"), Field.make("n2")),
            Binary.make(">", Field.make("n3"), Field.make("n4")),
        ),
    )
    # priority: = vs like
    assert parse("[n1] = [n2] like [n3]") == n.formula(
        Binary.make(
            "==",
            Field.make("n1"),
            Binary.make("like", Field.make("n2"), Field.make("n3")),
        )
    )

    # BETWEEN / NOT BETWEEN
    assert parse("[my_field] BETWEEN 123 AND 456") == n.formula(
        Ternary.make("between", Field.make("my_field"), LiteralInteger.make(123), LiteralInteger.make(456))
    )
    assert parse("[my_field]NOT BETWEEN(123)AND(456)") == n.formula(
        Ternary.make("notbetween", Field.make("my_field"), n.p(LiteralInteger.make(123)), n.p(LiteralInteger.make(456)))
    )
    # priority: BETWEEN vs comparison
    assert parse("45 BETWEEN [n2] = 123 AND 456") == n.formula(
        Ternary.make(
            "between",
            LiteralInteger.make(45),
            Binary.make("==", Field.make("n2"), LiteralInteger.make(123)),
            LiteralInteger.make(456),
        )
    )
    assert parse("[first] BETWEEN NOT [second] AND NOT [third]") == n.formula(
        Ternary.make(
            "between",
            Field.make("first"),
            Unary.make("not", Field.make("second")),
            Unary.make("not", Field.make("third")),
        )
    )
    # priority: BETWEEN vs AND
    assert parse("[first] BETWEEN [second] AND [third] AND [fourth]") == n.formula(
        Binary.make(
            "and",
            Ternary.make("between", Field.make("first"), Field.make("second"), Field.make("third")),
            Field.make("fourth"),
        ),
    )

    # IN / NOT IN
    assert parse("[my_field] IN (123, 456)") == n.formula(
        Binary.make(
            "in", Field.make("my_field"), ExpressionList.make(LiteralInteger.make(123), LiteralInteger.make(456))
        )
    )
    assert parse("[my_field]NOT IN(123, 456)") == n.formula(
        Binary.make(
            "notin", Field.make("my_field"), ExpressionList.make(LiteralInteger.make(123), LiteralInteger.make(456))
        )
    )
    # priority: IN vs BETWEEN
    assert parse("[n2] BETWEEN 1 AND 2 IN (2)") == n.formula(
        Binary.make(
            "in",
            Ternary.make("between", Field.make("n2"), LiteralInteger.make(1), LiteralInteger.make(2)),
            ExpressionList.make(LiteralInteger.make(2)),
        )
    )

    # NOT
    assert parse("NOT [my_field]") == n.formula(Unary.make("not", Field.make("my_field")))
    assert parse("NOT[my_field]") == n.formula(Unary.make("not", Field.make("my_field")))
    assert parse("NOT 123") == n.formula(Unary.make("not", LiteralInteger.make(123)))
    assert parse("NOT NOT [my_field]") == n.formula(Unary.make("not", Unary.make("not", Field.make("my_field"))))
    assert parse("NOT(123)") == n.formula(Unary.make("not", n.p(LiteralInteger.make(123))))
    # priority: NOT vs IN
    assert parse("NOT [n1] IN (1)") == n.formula(
        Unary.make("not", Binary.make("in", Field.make("n1"), ExpressionList.make(LiteralInteger.make(1))))
    )

    # AND
    assert parse("[left] AND [right]") == n.formula(Binary.make("and", Field.make("left"), Field.make("right")))
    assert parse("[left]AND[right]") == n.formula(Binary.make("and", Field.make("left"), Field.make("right")))
    assert parse("123 AND 123") == n.formula(Binary.make("and", LiteralInteger.make(123), LiteralInteger.make(123)))
    assert parse("[n1] AND [n2] AND [n3]") == n.formula(
        Binary.make(
            "and",
            Binary.make("and", Field.make("n1"), Field.make("n2")),
            Field.make("n3"),
        )
    )
    # priority: AND vs NOT
    assert parse("NOT [n1] AND [n2]") == n.formula(
        Binary.make("and", Unary.make("not", Field.make("n1")), Field.make("n2"))
    )

    # OR
    assert parse("[left] OR [right]") == n.formula(Binary.make("or", Field.make("left"), Field.make("right")))
    assert parse("[left]OR[right]") == n.formula(Binary.make("or", Field.make("left"), Field.make("right")))
    assert parse("123 OR 123") == n.formula(Binary.make("or", LiteralInteger.make(123), LiteralInteger.make(123)))
    assert parse("[n1] OR [n2] OR [n3]") == n.formula(
        Binary.make(
            "or",
            Binary.make("or", Field.make("n1"), Field.make("n2")),
            Field.make("n3"),
        )
    )

    # priority: OR vs AND
    assert parse("[n1] OR [n2] AND [n3]") == n.formula(
        Binary.make(
            "or",
            Field.make("n1"),
            Binary.make("and", Field.make("n2"), Field.make("n3")),
        )
    )


def test_parse_if_block(parser):
    assert parser.parse('IF "cond_1" THEN "res_1" ELSEIF "cond_2" THEN "res_2" ELSE "res_3" END') == n.formula(
        IfBlock.make(
            if_list=[
                IfPart.make(LiteralString.make("cond_1"), LiteralString.make("res_1")),
                IfPart.make(LiteralString.make("cond_2"), LiteralString.make("res_2")),
            ],
            else_expr=LiteralString.make("res_3"),
        )
    )

    # nested IF blocks
    # TODO: flatten ELSE-nested IF blocks?
    assert parser.parse('IF "cond_1" THEN "res_1" ELSE IF "cond_2" THEN "res_2" ELSE "res_3" END END') == n.formula(
        IfBlock.make(
            if_list=[
                IfPart.make(LiteralString.make("cond_1"), LiteralString.make("res_1")),
            ],
            else_expr=IfBlock.make(
                if_list=[IfPart.make(LiteralString.make("cond_2"), LiteralString.make("res_2"))],
                else_expr=LiteralString.make("res_3"),
            ),
        )
    )


def test_parse_case_block(parser):
    assert parser.parse(
        'CASE [expr_0] WHEN "val_1" THEN "res_1" WHEN "val_2" THEN "res_2" ELSE "res_3" END'
    ) == n.formula(
        CaseBlock.make(
            case_expr=Field.make("expr_0"),
            when_list=[
                WhenPart.make(LiteralString.make("val_1"), LiteralString.make("res_1")),
                WhenPart.make(LiteralString.make("val_2"), LiteralString.make("res_2")),
            ],
            else_expr=LiteralString.make("res_3"),
        )
    )


def test_parse_func_call(parser):
    assert parser.parse("my_func()") == n.formula(FuncCall.make(name="my_func", args=[]))
    assert parser.parse('my_func("s")') == n.formula(FuncCall.make(name="my_func", args=[LiteralString.make("s")]))
    assert parser.parse('my_func("s", 45.6, [my_field])') == n.formula(
        FuncCall.make(name="my_func", args=[LiteralString.make("s"), LiteralFloat.make(45.6), Field.make("my_field")])
    )


def test_parse_win_func_call(parser):
    assert parser.parse("mavg()") == n.formula(
        WindowFuncCall.make(name="mavg", args=[], grouping=WindowGroupingTotal())
    )
    assert parser.parse("mavg(TOTAL)") == n.formula(
        WindowFuncCall.make(name="mavg", args=[], grouping=WindowGroupingTotal())
    )
    assert parser.parse("mavg([my_field] AMONG [dim1], [dim2])") == n.formula(
        WindowFuncCall.make(
            name="mavg",
            args=[Field.make("my_field")],
            grouping=WindowGroupingAmong.make(dim_list=[Field.make("dim1"), Field.make("dim2")]),
        )
    )
    assert parser.parse("mavg([my_field] WITHIN [dim1], [dim2])") == n.formula(
        WindowFuncCall.make(
            name="mavg",
            args=[Field.make("my_field")],
            grouping=WindowGroupingWithin.make(
                dim_list=[
                    Field.make("dim1"),
                    Field.make("dim2"),
                ]
            ),
        )
    )


def test_agg_with_lod(parser):
    assert parser.parse("sum([my_field] FIXED [dim1], [dim2])") == n.formula(
        FuncCall.make(
            name="sum",
            args=[Field.make("my_field")],
            lod=FixedLodSpecifier.make(dim_list=[Field.make("dim1"), Field.make("dim2")]),
        )
    )
    assert parser.parse("sum([my_field] INCLUDE [dim1], [dim2])") == n.formula(
        FuncCall.make(
            name="sum",
            args=[Field.make("my_field")],
            lod=IncludeLodSpecifier.make(
                dim_list=[
                    Field.make("dim1"),
                    Field.make("dim2"),
                ]
            ),
        )
    )
    assert parser.parse("sum([my_field] EXCLUDE [dim1], [dim2])") == n.formula(
        FuncCall.make(
            name="sum",
            args=[Field.make("my_field")],
            lod=ExcludeLodSpecifier.make(
                dim_list=[
                    Field.make("dim1"),
                    Field.make("dim2"),
                ]
            ),
        )
    )


def test_parse_win_func_call_with_order_by(parser):
    assert parser.parse("mavg([my_field] WITHIN [dim1], [dim2] ORDER BY [dim3])") == n.formula(
        WindowFuncCall.make(
            name="mavg",
            args=[Field.make("my_field")],
            grouping=WindowGroupingWithin.make(
                dim_list=[
                    Field.make("dim1"),
                    Field.make("dim2"),
                ]
            ),
            ordering=Ordering.make(
                expr_list=[
                    Field.make("dim3"),
                ]
            ),
        )
    )
    assert parser.parse("mavg([my_field] ORDER BY [dim3])") == n.formula(
        WindowFuncCall.make(
            name="mavg",
            args=[Field.make("my_field")],
            ordering=Ordering.make(
                expr_list=[
                    Field.make("dim3"),
                ]
            ),
        )
    )
    assert parser.parse("mavg([my_field] ORDER BY [dim1], [dim2] DESC, [dim3] ASC)") == n.formula(
        WindowFuncCall.make(
            name="mavg",
            args=[Field.make("my_field")],
            ordering=Ordering.make(
                expr_list=[
                    Field.make("dim1"),
                    OrderDescending.make(expr=Field.make("dim2")),
                    OrderAscending.make(expr=Field.make("dim3")),
                ]
            ),
        )
    )


def test_parse_win_func_call_with_before_filter_by(parser):
    assert parser.parse(
        "mavg([my_field] WITHIN [dim1], [dim2] ORDER BY [dim3] BEFORE FILTER BY [dim4], [dim5])"
    ) == n.formula(
        WindowFuncCall.make(
            name="mavg",
            args=[Field.make("my_field")],
            grouping=WindowGroupingWithin.make(
                dim_list=[
                    Field.make("dim1"),
                    Field.make("dim2"),
                ]
            ),
            ordering=Ordering.make(
                expr_list=[
                    Field.make("dim3"),
                ]
            ),
            before_filter_by=BeforeFilterBy.make(field_names=frozenset(["dim4", "dim5"])),
        )
    )
    assert parser.parse("mavg([my_field] BEFORE FILTER BY [dim3])") == n.formula(
        WindowFuncCall.make(
            name="mavg",
            args=[Field.make("my_field")],
            before_filter_by=BeforeFilterBy.make(field_names=frozenset(["dim3"])),
        )
    )


def test_parse_func_call_with_ignore_dimensions(parser):
    assert parser.parse("ago([my_field] IGNORE DIMENSIONS [dim1], [dim2])") == n.formula(
        FuncCall.make(
            name="ago",
            args=[Field.make("my_field")],
            ignore_dimensions=IgnoreDimensions.make(dim_list=(Field.make("dim1"), Field.make("dim2"))),
        )
    )


def test_comments_basic(parser):
    parse = parser.parse
    # Comments
    assert parse("[abc] -- [zxc]") == parse("[abc]")
    assert parse("[abc]-- [zxc]") == parse("[abc]")
    assert parse("[abc]--[zxc]") == parse("[abc]")
    assert parse("[abc] /* [zxc] */") == parse("[abc]")
    assert parse("[abc] /* [zxc]/* */") == parse("[abc]")
    assert parse('[abc] /* zxc" /**/') == parse("[abc]")
    assert parse("[abc] /* /* -- */") == parse("[abc]")
    assert parse(' -- " "\n"abc"') == parse('"abc"')

    # Not comments
    assert parse('" /* " /* */') == n.formula(LiteralString.make(value=" /* "))
    assert parse('" -- "') == n.formula(LiteralString.make(value=" -- "))
    assert parse("[ab -- c]") == n.formula(Field.make(name="ab -- c"))
    assert parse("[ab/*c*/]") == n.formula(Field.make(name="ab/*c*/"))
    double_minus = n.formula(
        expr=Binary.make(
            name="-",
            left=Field.make(name="abc"),
            right=Unary.make(name="neg", expr=Field.make(name="zxc")),
        )
    )
    assert parse("[abc] - - [zxc]") == double_minus
    assert parse("[abc]- - [zxc]") == double_minus
    assert parse("[abc] - -[zxc]") == double_minus
    assert parse("[abc]- -[zxc]") == double_minus

    # Invalid
    with pytest.raises(exc.ParseError):
        parse('" /* " */"')
    with pytest.raises(exc.ParseError):
        parse('" -- " \\n"')
    with pytest.raises(exc.ParseError):
        parse('[abc] /* "zxc */"*/')


def test_comments_in_if(parser):
    expected = parser.parse('IF "hi" THEN "world" END')

    def check(formula: str) -> None:
        assert parser.parse(formula) == expected

    # Single-line comments
    check('IF -- comment\n"hi" THEN "world" END')
    check('IF "hi" THEN "world" END-- comment\n')
    check('IF "hi" THEN "world" END-- comment')
    check('IF "hi" -- first\n THEN-- second \n"world" END -- third')

    # Multi-line comments
    check('IF /* comment */ "hi" THEN "world" END')
    check('IF /* co\nmm\nent */ "hi" THEN "world" END')

    # Mixed
    check('IF "hi" --first\n THEN /*  second*/ "world" END -- third')


def test_unexpected_char(parser):
    with pytest.raises(exc.ParseError) as exc_info:
        parser.parse("[n1]\n+ )")

    err = exc_info.value
    assert err.error.coords == (1, 2)
    assert err.error.token == ")"
    assert err.error.code == ("FORMULA", "PARSE", "UNEXPECTED_TOKEN")


def test_unexpected_end(parser):
    with pytest.raises(exc.ParseError) as exc_info:
        parser.parse("[n1]\n+ ")

    err = exc_info.value
    assert err.error.coords == (1, 2)
    assert err.error.token is None
    assert err.error.code == ("FORMULA", "PARSE", "UNEXPECTED_EOF")


def test_invalid_date_value(parser):
    with pytest.raises(exc.ParseError) as exc_info:
        parser.parse("#2010-13-45#")

    err = exc_info.value
    assert err.error.coords == (0, 0)
    assert err.error.token == "2010-13-45"
    assert err.error.code == ("FORMULA", "PARSE", "VALUE", "DATE")


def test_invalid_datetime_value(parser):
    with pytest.raises(exc.ParseError) as exc_info:
        parser.parse("#2010-04-02 28:00:00#")

    err = exc_info.value
    assert err.error.coords == (0, 0)
    assert err.error.token == "2010-04-02 28:00:00"
    assert err.error.code == ("FORMULA", "PARSE", "VALUE", "DATETIME")


def test_inconceivably_long_if_block(parser):
    # Note: at around 500 it timeouts the whole unit test too easily. Might
    # need to move it to `medium` tests to test this much.
    num = 300
    formula = 'if contains("s", [v]) then "a"\n'
    if_list = [
        IfPart.make(
            cond=FuncCall.make("contains", [LiteralString.make("s"), Field.make("v")]), expr=LiteralString.make("a")
        )
    ]
    for if_num in range(num):
        formula += f'elseif contains("s_{if_num}", [v]) then "a_{if_num}"\n'
        if_list.append(
            IfPart.make(
                cond=FuncCall.make("contains", [LiteralString.make(f"s_{if_num}"), Field.make("v")]),
                expr=LiteralString.make(f"a_{if_num}"),
            )
        )
    formula += 'else "" end'
    expected_result = n.formula(IfBlock.make(if_list=if_list, else_expr=LiteralString.make("")))
    actual_result = parser.parse(formula)
    assert actual_result == expected_result

lib/dl_formula/dl_formula_tests/unit/parser/test_parser.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest

from dl_formula.scripts.formula_cli import (
    FormulaCliTool,
    parser,
)
from dl_formula_testing.tool_runner import ToolRunner


@pytest.fixture(scope="module")
def tool():
    return ToolRunner(parser=parser, tool_cls=FormulaCliTool)


def test_parse(tool):
    stdout, stderr = tool.run(["parse", "my_func(123)"])
    assert stdout == "Formula(expr=FuncCall(name='my_func', args=(LiteralInteger(value=123),)))\n"
    assert stderr == ""

    stdout, stderr = tool.run(["parse", "--pretty", "--with-meta", "very_very_long_func_name(nested_func())"])
    assert (
        stdout
        == """Formula(
  expr=FuncCall(
    name='very_very_long_func_name',
    args=(
      FuncCall(
        name='nested_func',
        args=(),
        meta=NodeMeta(position=Position(start=25, end=38, start_row=0, end_row=0, start_col=25, end_col=38)),
      ),
    ),
    meta=NodeMeta(position=Position(start=0, end=39, start_row=0, end_row=0, start_col=0, end_col=39)),
  ),
  meta=NodeMeta(position=Position(start=0, end=39, start_row=0, end_row=0, start_col=0, end_col=39)),
)
"""
    )
    assert stderr == ""

    stdout, stderr = tool.run(
        [
            "parse",
            "--suppress-errors",
        ],
        stdin=("SUM(1)\n" "AVG(2)\n" "INVALID_FORMULA([n1"),
    )
    assert stdout.strip() == (
        "Formula(expr=FuncCall(name='sum', args=(LiteralInteger(value=1),)))\n"
        "Formula(expr=FuncCall(name='avg', args=(LiteralInteger(value=2),)))\n"
        "#ERROR"
    )
    assert stderr == ""


def test_translate(tool):
    stdout, stderr = tool.run(
        [
            "translate",
            "--table",
            "my_table",
            "--dialect",
            "clickhouse",
            "--unknown-funcs",
            "SUM(IF(ISNULL([field]), 0, [field]))",
        ]
    )
    assert stdout == ("sum(if(isNull(my_table.field), 0, my_table.field))\n")
    assert stderr == ""

    stdout, stderr = tool.run(
        [
            "translate",
            "--dialect",
            "clickhouse",
            "--suppress-errors",
        ],
        stdin=("SUM([n1])\n" "AVG([n1])\n" "UNKNOWN_FUNC([n1])"),
    )
    assert stdout.strip() == ("sum(n1)\n" "avg(n1)\n" "#ERROR")
    assert stderr == ""


def test_split(tool):
    stdout, stderr = tool.run(["split", "--diff", "45 BETWEEN [n2] = 123 AND 456"])
    assert stdout == (
        " 0: ▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫\n"
        " 1: ▫▫ BETWEEN ▫▫▫▫▫▫▫▫▫▫ AND ▫▫▫\n"
        " 2: 45         ▫▫▫▫ = ▫▫▫     456\n"
        " 3:            [n2]   123        \n"
    )


def test_graph(tool):
    stdout, stderr = tool.run(
        [
            "graph",
            (
                'IF ([n0]) IN (5, 6) THEN MY_FUNC(NULL, NOT [n2], "qwerty") '
                'ELSEIF [n3] BETWEEN 1 AND 2 THEN "result" '
                'ELSE CASE [n4] WHEN 5 THEN "five" WHEN 6 THEN "seven" ELSE "forty two" END END'
            ),
        ]
    )
    assert "digraph {" in stdout
    assert stderr == ""


def test_dialects(tool):
    stdout, stderr = tool.run(
        [
            "dialects",
        ]
    )
    assert stderr == ""
    assert "DUMMY" in stdout


def test_goto(tool):
    stdout, stderr = tool.run(["goto", "--show", "avg"])
    assert stderr == ""
    assert "functions_aggregation.py" in stdout


def test_slice(tool):
    stdout, stderr = tool.run(
        [
            "slice",
            "--diff",
            "--levels",
            "aggregate,window,toplevel",
            "456 + SUM(FUNC(123)) + RMAX([coeff] * AVG(456 + [qwe] + [rty] + [uio]) - 1)",
        ]
    )
    assert (
        stdout.strip()
        == (
            "toplevel       ▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫ + RMAX(▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫)\n"
            "window         456 + SUM(▫▫▫▫▫▫▫▫▫)        ▫▫▫▫▫▫▫ * AVG(▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫▫) - 1 \n"
            "aggregate                FUNC(123)         [coeff]       456 + [qwe] + [rty] + [uio]      \n"
        ).strip()
    )

lib/dl_formula/dl_formula_tests/unit/scripts/test_formula_cli.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict

from dl_formula.core import nodes
from dl_formula.core.tag import LevelTag
from dl_formula.shortcuts import n
from dl_formula.slicing.schema import (
    AggregateFunctionLevelBoundary,
    NestedLevelTaggedBoundary,
    NonFieldsBoundary,
    SliceSchema,
    TopLevelBoundary,
    WindowFunctionLevelBoundary,
)
from dl_formula.slicing.slicer import (
    FormulaSlicer,
    LevelSlice,
    SlicedFormulaInfo,
)


class NameGen:
    """
    Minor helper for numbered field names.

    >>> ng = NameGen()
    >>> ng(None, 'fld1')
    'fld1_1'
    >>> ng(None, 'fld2')
    'fld2_1'
    >>> ng(None, 'fld1')
    'fld1_2'
    >>> NameGen()(None, 'fld2')
    'fld2_1'
    """

    def __init__(self):
        self.counters = defaultdict(lambda: 0)

    def __call__(self, node: nodes.FormulaItem, name: str) -> str:
        self.counters[name] += 1
        return f"{name}_{self.counters[name]}"


def test_aggregation_function_slicer():
    name_gen = NameGen()

    # slice aggregations
    agg_slicer = FormulaSlicer(
        slice_schema=SliceSchema(
            levels=[
                AggregateFunctionLevelBoundary(name="agg", name_gen=name_gen),
                TopLevelBoundary(name="top", name_gen=name_gen),
            ]
        ),
    )

    sliced_fla_info = agg_slicer.slice_formula(
        formula=n.formula(n.func.SUM(n.func.MY_FUNC(n.field("My Field")))),
    )
    expected = SlicedFormulaInfo(
        slices=[
            LevelSlice(
                name="agg",
                required_fields=set(),
                aliased_nodes={"agg_1": n.formula(n.func.MY_FUNC(n.field("My Field")))},
            ),
            LevelSlice(
                name="top",
                required_fields={"agg_1"},
                aliased_nodes={"top_1": n.formula(n.func.SUM(n.field("agg_1")))},
            ),
        ]
    )
    assert sliced_fla_info == expected

    # check laziness of slices
    assert not sliced_fla_info.slices[0].is_lazy()
    assert not sliced_fla_info.slices[1].is_lazy()

    # no aggregation, so the top slice should be lazy
    sliced_fla_info = agg_slicer.slice_formula(
        formula=n.formula(n.func.OTHER_FUNC(n.func.MY_FUNC(n.field("My Field")))),
    )
    assert len(sliced_fla_info.slices) == 2
    # check laziness of slices
    assert not sliced_fla_info.slices[0].is_lazy()
    assert sliced_fla_info.slices[1].is_lazy()


def test_window_function_slicer():
    name_gen = NameGen()

    win_slicer = FormulaSlicer(
        slice_schema=SliceSchema(
            levels=[
                WindowFunctionLevelBoundary(name="win", name_gen=name_gen),
                TopLevelBoundary(name="top", name_gen=name_gen),
            ]
        ),
    )

    sliced_fla_info = win_slicer.slice_formula(
        formula=n.formula(
            n.func("*")(
                n.func("/")(
                    n.wfunc.RSUM(
                        n.func.SUM(n.func.MY_FUNC(n.field("Sales"))),
                        grouping=n.among(n.field("Dim")),
                        order_by=[n.desc(n.field("Order Date"))],
                    ),
                    n.wfunc.SUM(n.func.SUM(n.func.MY_FUNC(n.field("Sales"))), grouping=n.total()),
                ),
                n.lit(100),
            )
        ),
    )
    expected = SlicedFormulaInfo(
        slices=[
            LevelSlice(
                name="win",
                required_fields=set(),
                aliased_nodes={
                    "win_1": n.formula(n.func.SUM(n.func.MY_FUNC(n.field("Sales")))),
                    "win_2": n.formula(n.field("Order Date")),
                    "win_3": n.formula(n.field("Dim")),
                },
            ),
            LevelSlice(
                name="top",
                required_fields={"win_1", "win_2", "win_3"},
                aliased_nodes={
                    "top_1": n.formula(
                        n.func("*")(
                            n.func("/")(
                                n.wfunc.RSUM(
                                    n.field("win_1"),
                                    grouping=n.among(n.field("win_3")),
                                    order_by=[n.desc(n.field("win_2"))],
                                ),
                                n.wfunc.SUM(n.field("win_1"), grouping=n.total()),
                            ),
                            n.lit(100),
                        )
                    )
                },
            ),
        ]
    )
    assert sliced_fla_info == expected

    # check laziness of slices
    assert not sliced_fla_info.slices[0].is_lazy()
    assert not sliced_fla_info.slices[1].is_lazy()

    # no window function, so the top slice should be lazy
    sliced_fla_info = win_slicer.slice_formula(
        formula=n.formula(n.func.OTHER_FUNC(n.func.MY_FUNC(n.field("My Field")))),
    )
    assert len(sliced_fla_info.slices) == 2
    # check laziness of slices
    assert not sliced_fla_info.slices[0].is_lazy()
    assert sliced_fla_info.slices[1].is_lazy()


def test_non_fields_slicer():
    name_gen = NameGen()

    slicer = FormulaSlicer(
        slice_schema=SliceSchema(
            levels=[
                NonFieldsBoundary(name="nf", name_gen=name_gen),
                TopLevelBoundary(name="top", name_gen=name_gen),
            ]
        ),
    )

    sliced_fla_info = slicer.slice_formula(
        formula=n.formula(
            n.func("*")(
                n.func("/")(
                    n.wfunc.RSUM(
                        n.func.SUM(n.func.MY_FUNC(n.field("Sales"))),
                        grouping=n.among(n.field("Dim")),
                        order_by=[n.field("Order Date")],
                    ),
                    n.wfunc.SUM(n.func.SUM(n.func.MY_FUNC(n.field("Sales"))), grouping=n.total()),
                ),
                n.lit(100),
            )
        ),
    )
    expected = SlicedFormulaInfo(
        slices=[
            LevelSlice(
                name="nf",
                required_fields=set(),
                aliased_nodes={
                    "nf_1": n.formula(n.field("Sales")),
                    "nf_2": n.formula(n.field("Order Date")),
                    "nf_3": n.formula(n.field("Dim")),
                },
            ),
            LevelSlice(
                name="top",
                required_fields={"nf_1", "nf_2", "nf_3"},
                aliased_nodes={
                    "top_1": n.formula(
                        n.func("*")(
                            n.func("/")(
                                n.wfunc.RSUM(
                                    n.func.SUM(n.func.MY_FUNC(n.field("nf_1"))),
                                    grouping=n.among(n.field("nf_3")),
                                    order_by=[n.field("nf_2")],
                                ),
                                n.wfunc.SUM(
                                    n.func.SUM(n.func.MY_FUNC(n.field("nf_1"))),
                                    grouping=n.total(),
                                ),
                            ),
                            n.lit(100),
                        )
                    ),
                },
            ),
        ]
    )
    assert sliced_fla_info == expected

    # check laziness of slices
    assert sliced_fla_info.slices[0].is_lazy()
    assert not sliced_fla_info.slices[1].is_lazy()

    # nothing except field, so both slices should be lazy
    sliced_fla_info = slicer.slice_formula(
        formula=n.formula(n.field("My Field")),
    )
    assert len(sliced_fla_info.slices) == 2
    # check laziness of slices
    assert sliced_fla_info.slices[0].is_lazy()
    assert sliced_fla_info.slices[1].is_lazy()


def test_tagged_slicer():
    name_gen = NameGen()

    slicer = FormulaSlicer(
        slice_schema=SliceSchema(
            levels=[
                NestedLevelTaggedBoundary(
                    name="tagged",
                    name_gen=name_gen,
                    tag=LevelTag(bfb_names=frozenset({"slice_here"}), func_nesting=-1, qfork_nesting=0),
                ),
                TopLevelBoundary(name="top", name_gen=name_gen),
            ]
        ),
    )

    sliced_fla_info = slicer.slice_formula(
        formula=n.formula(
            n.func.SUM(
                n.tagged(
                    tag=LevelTag(bfb_names=frozenset({"slice_here"}), func_nesting=-1, qfork_nesting=0),
                    expr=n.p(n.wfunc.RSUM(n.field("My Field"))),
                )
            ),
        ),
    )
    expected = SlicedFormulaInfo(
        slices=[
            LevelSlice(
                name="tagged",
                required_fields=set(),
                aliased_nodes={"tagged_1": n.formula(n.wfunc.RSUM(n.field("My Field")))},
            ),
            LevelSlice(
                name="top",
                required_fields={"tagged_1"},
                aliased_nodes={
                    "top_1": n.formula(
                        n.func.SUM(
                            n.tagged(
                                tag=LevelTag(bfb_names=frozenset({"slice_here"}), func_nesting=-1, qfork_nesting=0),
                                expr=n.p(n.field("tagged_1")),
                            )
                        )
                    )
                },
            ),
        ]
    )
    assert sliced_fla_info == expected

lib/dl_formula/dl_formula_tests/unit/slicing/test_slicer.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import uuid

from dl_formula.core import nodes
from dl_formula.shortcuts import n


def test_literal_shortcuts():
    assert n.lit("qwerty") == nodes.LiteralString.make("qwerty")
    assert n.lit(123) == nodes.LiteralInteger.make(123)
    assert n.lit(12.3) == nodes.LiteralFloat.make(12.3)
    assert n.lit(True) == nodes.LiteralBoolean.make(True)
    assert n.lit(None) == nodes.Null()
    dt = datetime.datetime.utcnow().replace(microsecond=0, tzinfo=None)
    assert n.lit(dt) == nodes.LiteralGenericDatetime.make(dt)
    assert n.lit(dt.date()) == nodes.LiteralDate.make(dt.date())
    my_uuid = uuid.uuid4()
    assert n.lit(my_uuid) == nodes.LiteralUuid.make(my_uuid)


def test_functions():
    assert n.func("+")(n.func.MY_FUNC(123, True, nodes.LiteralString.make("456")), None) == (
        nodes.FuncCall.make(
            name="+",
            args=[
                nodes.FuncCall.make(
                    name="my_func",
                    args=[
                        nodes.LiteralInteger.make(123),
                        nodes.LiteralBoolean.make(True),
                        nodes.LiteralString.make("456"),
                    ],
                ),
                nodes.Null(),
            ],
        )
    )


def test_window_functions():
    assert n.func.MY_FUNC(n.field("f0"), total=True) == (
        nodes.WindowFuncCall.make(
            name="my_func", args=[nodes.Field.make(name="f0")], grouping=nodes.WindowGroupingTotal()
        )
    )
    assert n.func.MY_FUNC(n.field("f0"), within=[n.field("f1")]) == (
        nodes.WindowFuncCall.make(
            name="my_func",
            args=[nodes.Field.make(name="f0")],
            grouping=nodes.WindowGroupingWithin.make(dim_list=[nodes.Field.make(name="f1")]),
        )
    )
    assert n.func.MY_FUNC(n.field("f0"), among=[n.field("f1")]) == (
        nodes.WindowFuncCall.make(
            name="my_func",
            args=[nodes.Field.make(name="f0")],
            grouping=nodes.WindowGroupingAmong.make(dim_list=[nodes.Field.make(name="f1")]),
        )
    )
    assert n.func.MY_FUNC(n.field("f0"), grouping=n.among(n.field("f1"))) == (
        nodes.WindowFuncCall.make(
            name="my_func",
            args=[nodes.Field.make(name="f0")],
            grouping=nodes.WindowGroupingAmong.make(dim_list=[nodes.Field.make(name="f1")]),
        )
    )
    assert n.wfunc.MY_FUNC(n.field("f0")) == (
        nodes.WindowFuncCall.make(
            name="my_func", args=[nodes.Field.make(name="f0")], grouping=nodes.WindowGroupingTotal()
        )
    )


def test_window_functions_order_by():
    assert n.func.MY_FUNC(n.field("f0"), total=True, order_by=[n.field("f1")]) == (
        nodes.WindowFuncCall.make(
            name="my_func",
            args=[nodes.Field.make(name="f0")],
            ordering=nodes.Ordering.make(
                expr_list=[
                    nodes.Field.make("f1"),
                ]
            ),
        )
    )


def test_window_functions_before_filter_by():
    assert n.func.MY_FUNC(n.field("f0"), total=True, before_filter_by=[n.field("f1")]) == (
        nodes.WindowFuncCall.make(
            name="my_func",
            args=[nodes.Field.make(name="f0")],
            before_filter_by=nodes.BeforeFilterBy.make(
                field_names=frozenset(
                    [
                        "f1",
                    ]
                )
            ),
        )
    )


def test_logical_blocks():
    assert n.if_(n.func.WHATEVER()).then("qwerty").else_("uiop") == (
        nodes.IfBlock.make(
            if_list=[
                nodes.IfPart.make(
                    cond=nodes.FuncCall.make(name="whatever", args=[]),
                    expr=nodes.LiteralString.make("qwerty"),
                ),
            ],
            else_expr=nodes.LiteralString.make("uiop"),
        )
    )
    assert n.if_(
        n.if_(n.func.WHATEVER()).then("qwerty"),
        n.if_(n.func.WHENEVER()).then("uiop"),
        n.if_(n.func.WHOEVER()).then("asdfg"),
    ).else_("xyz") == (
        nodes.IfBlock.make(
            if_list=[
                nodes.IfPart.make(
                    cond=nodes.FuncCall.make(name="whatever", args=[]),
                    expr=nodes.LiteralString.make("qwerty"),
                ),
                nodes.IfPart.make(
                    cond=nodes.FuncCall.make(name="whenever", args=[]),
                    expr=nodes.LiteralString.make("uiop"),
                ),
                nodes.IfPart.make(
                    cond=nodes.FuncCall.make(name="whoever", args=[]),
                    expr=nodes.LiteralString.make("asdfg"),
                ),
            ],
            else_expr=nodes.LiteralString.make("xyz"),
        )
    )
    assert n.case(n.field("my_field")).whens(
        n.when(n.func.WHATEVER()).then("qwerty"),
        n.when(n.func.WHENEVER()).then("uiop"),
    ).else_("xyz") == (
        nodes.CaseBlock.make(
            case_expr=nodes.Field.make(name="my_field"),
            when_list=[
                nodes.WhenPart.make(
                    val=nodes.FuncCall.make(name="whatever", args=[]),
                    expr=nodes.LiteralString.make("qwerty"),
                ),
                nodes.WhenPart.make(
                    val=nodes.FuncCall.make(name="whenever", args=[]),
                    expr=nodes.LiteralString.make("uiop"),
                ),
            ],
            else_expr=nodes.LiteralString.make("xyz"),
        )
    )

lib/dl_formula/dl_formula_tests/unit/test_shortcuts.py
datalens-tech/datalens-backend
from dl_formula.translation.env import (
    FunctionStatsSignature,
    TranslationStats,
)


def test_function_stats_combine():
    stats_1 = TranslationStats(
        cache_hits=2,
        function_usage_weights={
            FunctionStatsSignature(name="my_func", arg_types=("int",), dialect="qqqq", is_window=False): 1,
            FunctionStatsSignature(name="my_func", arg_types=("int", "str"), dialect="qqqq", is_window=False): 3,
        },
    )
    stats_2 = TranslationStats(
        cache_hits=5,
        function_usage_weights={
            FunctionStatsSignature(name="my_func", arg_types=("int", "str"), dialect="qqqq", is_window=False): 4,
            FunctionStatsSignature(name="other_func", arg_types=(), dialect="qqqq", is_window=False): 10,
        },
    )
    expected_combined_stats = TranslationStats(
        cache_hits=7,
        function_usage_weights={
            FunctionStatsSignature(name="my_func", arg_types=("int",), dialect="qqqq", is_window=False): 1,
            FunctionStatsSignature(name="my_func", arg_types=("int", "str"), dialect="qqqq", is_window=False): 7,
            FunctionStatsSignature(name="other_func", arg_types=(), dialect="qqqq", is_window=False): 10,
        },
    )
    assert TranslationStats.combine(stats_1, stats_2) == expected_combined_stats
    assert stats_1 + stats_2 == expected_combined_stats

lib/dl_formula/dl_formula_tests/unit/translation/test_stats.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Dict,
    Optional,
)

import pytest
import sqlalchemy as sa

from dl_formula.core.aux_nodes import ErrorNode
from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
import dl_formula.core.exc as exc
from dl_formula.core.nodes import (
    Binary,
    Field,
    Formula,
    FuncCall,
    LiteralInteger,
    LiteralString,
    NodeMeta,
    Null,
    OrderDescending,
    ParenthesizedExpr,
)
from dl_formula.core.position import Position
from dl_formula.definitions.registry import OPERATION_REGISTRY
from dl_formula.definitions.scope import Scope
from dl_formula.translation.context import TranslationCtx
from dl_formula.translation.env import TranslationEnvironment
from dl_formula.translation.translator import (
    translate,
    translate_and_compile,
)


FIELD_TYPES = {
    "f1": DataType.FLOAT,
    "n1": DataType.INTEGER,
    "n2": DataType.INTEGER,
    "n3": DataType.INTEGER,
    "field with spaces": DataType.INTEGER,
    "b1": DataType.BOOLEAN,
    "b2": DataType.BOOLEAN,
    "s1": DataType.STRING,
    "s2": DataType.STRING,
    "d1": DataType.DATE,
    "d2": DataType.DATE,
}


def T(
    x,
    dialect=D.DUMMY,
    restrict_funcs=False,
    field_types: Dict[str, DataType] = None,
    table_name: str = None,
    env: Optional[TranslationEnvironment] = None,
):
    field_types = field_types or FIELD_TYPES
    if table_name:
        field_names = {f: (table_name, f) for f in field_types}
    else:
        field_names = {f: (f,) for f in field_types}
    return translate_and_compile(
        Formula(x),
        env=env,
        dialect=dialect,
        restrict_funcs=restrict_funcs,
        field_types=field_types,
        field_names=field_names,
    )


def test_translate_error_node():
    with pytest.raises(exc.TranslationError) as exc_info:
        T(ErrorNode.make(err_code=exc.InconsistentAggregationError.default_code, message="Hello"))

    err_ctx = exc_info.value.errors[0]
    assert err_ctx.code == exc.InconsistentAggregationError.default_code
    assert err_ctx.message == "Hello"


def test_translate_parenthesized_expressions():
    assert T(ParenthesizedExpr.make(expr=LiteralInteger.make(123)), dialect=D.DUMMY) == "123"


def test_translate_descending():
    assert T(OrderDescending.make(expr=LiteralInteger.make(123)), dialect=D.DUMMY) == "123 DESC"
    assert T(OrderDescending.make(expr=OrderDescending.make(expr=LiteralInteger.make(123))), dialect=D.DUMMY) == "123"


def test_operator_null_error():
    with pytest.raises(exc.TranslationError):
        T(Binary.make("==", Field.make("n1"), Null()), dialect=D.DUMMY)


def test_translate_multiple_errors():
    try:
        translate(
            Formula(
                FuncCall.make(
                    name="GREATEST",
                    args=[
                        FuncCall.make(
                            name="UNKNOWN_1", args=[], meta=NodeMeta(position=Position(4, 100, 0, 0, 4, 100))
                        ),
                        Field.make(name="Unknown F", meta=NodeMeta(position=Position(104, 200, 0, 0, 104, 200))),
                    ],
                )
            ),
            dialect=D.DUMMY,
            collect_errors=True,
        )
        raised = False
    except exc.TranslationError as err:
        raised = True
        assert len(err.errors) == 2
        assert err.errors[0].token == "UNKNOWN_1"
        assert err.errors[0].position.start == 4
        assert err.errors[0].code == ("FORMULA", "TRANSLATION", "UNKNOWN_FUNCTION")
        assert err.errors[1].token == "Unknown F"
        assert err.errors[1].position.start == 104
        assert err.errors[1].code == ("FORMULA", "TRANSLATION", "UNKNOWN_FIELD")

    assert raised


def test_internal_functions():
    env = TranslationEnvironment(dialect=D.DUMMY, required_scopes=Scope.EXPLICIT_USAGE)
    with pytest.raises(exc.TranslationError) as exc_info:
        T(
            FuncCall.make(name="__str", args=[LiteralString.make("qwerty")]),
            env=env,
            restrict_funcs=True,
        )

    assert exc_info.value.errors[0].code == exc.TranslationUnknownFunctionError.default_code


def test_internal_function_registration():
    for func_key, func_tr in OPERATION_REGISTRY.items():
        name = func_key[0]
        if name.startswith("__"):
            assert (
                func_tr.scopes & ~Scope.EXPLICIT_USAGE == func_tr.scopes
            ), 'Functions with names starting with "__" must be registered as internal'


def test_translation_cache():
    sub_node = Field.make(name="qwerty")
    env = TranslationEnvironment(dialect=D.DUMMY, required_scopes=Scope.EXPLICIT_USAGE)
    env.translation_cache.add(
        node=sub_node,
        value=TranslationCtx(
            required_scopes=env.required_scopes,
            node=sub_node,
            data_type=DataType.INTEGER,
            expression=sa.literal(123),
        ),
    )
    assert T(FuncCall.make(name="+", args=[Field.make("qwerty"), LiteralInteger.make(456)]), env=env) == ("123 + 456")


def test_translation_replacements():
    sub_node = Field.make(name="qwerty")
    env = TranslationEnvironment(dialect=D.DUMMY, required_scopes=Scope.EXPLICIT_USAGE)
    env.replacements.add(
        node=sub_node,
        value=TranslationCtx(
            node=sub_node,
            data_type=DataType.INTEGER,
            expression=sa.literal(123),
        ),
    )
    assert T(FuncCall.make(name="+", args=[Field.make("qwerty"), LiteralInteger.make(456)]), env=env) == ("123 + 456")

lib/dl_formula/dl_formula_tests/unit/translation/test_translator.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_formula.core.datatype import DataType as DT
from dl_formula.definitions.type_strategy import (
    Fixed,
    FromArgs,
)


def test_strategy_fixed():
    strat = Fixed(DT.STRING)
    assert strat.get_from_args([DT.INTEGER, DT.DATETIME]) == DT.STRING
    assert strat.get_from_args([DT.STRING]) == DT.STRING


def test_strategy_from_args():
    strat = FromArgs(slice(1, None))
    assert strat.get_from_args([DT.DATETIME, DT.INTEGER]) == DT.INTEGER
    assert strat.get_from_args([DT.DATETIME, DT.INTEGER, DT.FLOAT, DT.INTEGER]) == DT.FLOAT

lib/dl_formula/dl_formula_tests/unit/translation/test_type_resolution.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest

from dl_formula.utils.caching import multi_cached_with_errors


def test_cached_with_errors():
    counter = 0

    @multi_cached_with_errors(maxsize=3, cache_exceptions=(TypeError,))
    def my_func(a: int, b: int) -> int:
        nonlocal counter
        counter += 1
        return a // b

    assert my_func(2, 2) == 1
    assert counter == 1
    assert my_func(2, 2) == 1
    assert counter == 1

    assert my_func(4, 2) == 2
    assert my_func(6, 2) == 3
    assert my_func(8, 2) == 4
    assert counter == 4

    assert my_func(2, 2) == 1  # cache overflow
    assert counter == 5
    assert my_func(6, 2) == 3
    assert counter == 5

    # cached exception
    assert pytest.raises(TypeError, my_func, None, None)
    assert counter == 6
    assert pytest.raises(TypeError, my_func, None, None)
    assert counter == 6
    # noncached exception
    assert pytest.raises(ZeroDivisionError, my_func, 2, 0)
    assert counter == 7
    assert pytest.raises(ZeroDivisionError, my_func, 2, 0)
    assert counter == 8

    my_func.cache_clear()
    assert pytest.raises(TypeError, my_func, None, None)
    assert counter == 9


def test_multi_cache():
    counter = 0

    def cache_qualifier(value: str) -> str:
        return value[0].lower()

    @multi_cached_with_errors(maxsize=2, cache_qualifier=cache_qualifier)
    def my_func(value: str) -> str:
        nonlocal counter
        counter += 1
        return value.upper()

    assert my_func("panda") == "PANDA"  # miss
    assert my_func("puma") == "PUMA"  # miss
    assert my_func("tiger") == "TIGER"  # miss
    assert my_func("parrot") == "PARROT"  # miss, overflow -> "panda" ejected
    assert my_func("tiger") == "TIGER"  # hit
    assert my_func("panda") == "PANDA"  # miss because of "parrot"

    assert counter == 5  # the number of misses

    cache_info = my_func.cache_info()
    assert set(cache_info.keys()) == {"p", "t"}
    assert cache_info["p"].misses == 4
    assert cache_info["p"].hits == 0
    assert cache_info["t"].misses == 1
    assert cache_info["t"].hits == 1

lib/dl_formula/dl_formula_tests/unit/utils/test_caching.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import List

import pytest

from dl_formula.core import (
    exc,
    nodes,
)
from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.shortcuts import n
from dl_formula.validation.aggregation import AggregationChecker
from dl_formula.validation.env import ValidationEnvironment
from dl_formula.validation.validator import validate


def get_dims() -> List[nodes.FormulaItem]:
    return [
        nodes.Field.make("Dim Field"),
        nodes.FuncCall.make(
            name="+",
            args=[
                nodes.Field.make("Other Field"),
                nodes.FuncCall.make(
                    name="dim_func",
                    args=[
                        nodes.Field.make("Third Field"),
                    ],
                ),
            ],
        ),
    ]


def validate_aggregations(
    node: nodes.FormulaItem,
    env: ValidationEnvironment,
    dimensions: List[nodes.FormulaItem],
    collect_errors: bool = False,
) -> None:
    validate(
        node,
        env=env,
        checkers=[AggregationChecker(valid_env=env, inspect_env=InspectionEnvironment(), global_dimensions=dimensions)],
        collect_errors=collect_errors,
    )


def test_no_errors():
    dims = get_dims()
    dim1, dim2 = dims
    env = ValidationEnvironment()

    # not aggregated
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="+",
                args=[
                    dim2,
                    nodes.LiteralFloat.make(1.1),
                ],
            ),
        ),
        env=env,
        dimensions=dims,
    )

    # aggregated
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(name="sum", args=[nodes.Field.make("Some Field")]),
        ),
        env=env,
        dimensions=dims,
    )
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(name="sum", args=[nodes.LiteralString.make("qwerty")]),
        ),
        env=env,
        dimensions=dims,
    )
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="sum",
                args=[
                    nodes.ParenthesizedExpr.make(
                        nodes.LiteralString.make("qwerty"),
                    ),
                ],
            ),
        ),
        env=env,
        dimensions=dims,
    )
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="sum",
                args=[
                    nodes.FuncCall.make(
                        name="+",
                        args=[
                            nodes.Field.make("Barley Field"),
                            nodes.Field.make("Other Field"),
                        ],
                    ),
                ],
            ),
        ),
        env=env,
        dimensions=dims,
    )
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="sum",
                args=[
                    nodes.FuncCall.make(
                        name="+",
                        args=[
                            dim1,
                            nodes.Field.make("Other Field"),
                        ],
                    ),
                ],
            ),
        ),
        env=env,
        dimensions=dims,
    )
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="sum",
                args=[
                    nodes.FuncCall.make(
                        name="+",
                        args=[
                            dim1,
                            dim2,
                        ],
                    ),
                ],
            ),
        ),
        env=env,
        dimensions=dims,
    )
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="sum",
                args=[
                    nodes.FuncCall.make(
                        name="+",
                        args=[
                            dim2,
                            nodes.LiteralFloat.make(1.1),
                        ],
                    ),
                ],
            ),
        ),
        env=env,
        dimensions=dims,
    )

    # nested aggregated
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="func",
                args=[
                    nodes.FuncCall.make(
                        name="sum",
                        args=[
                            nodes.FuncCall.make(
                                name="+",
                                args=[
                                    dim2,
                                    nodes.LiteralFloat.make(1.1),
                                ],
                            ),
                        ],
                    ),
                ],
            ),
        ),
        env=env,
        dimensions=dims,
    )
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="func",
                args=[
                    nodes.FuncCall.make(
                        name="sum",
                        args=[
                            nodes.FuncCall.make(
                                name="+",
                                args=[
                                    dim2,
                                    nodes.LiteralFloat.make(1.1),
                                ],
                            ),
                        ],
                    ),
                    dim2,
                ],
            ),
        ),
        env=env,
        dimensions=dims,
    )
    validate_aggregations(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="func",
                args=[
                    nodes.LiteralBoolean.make(True),
                    nodes.FuncCall.make(
                        name="sum",
                        args=[
                            nodes.FuncCall.make(
                                name="+",
                                args=[
                                    dim2,
                                    nodes.LiteralFloat.make(1.1),
                                ],
                            ),
                        ],
                    ),
                ],
            ),
        ),
        env=env,
        dimensions=dims,
    )


def test_double_aggregation():
    dims = get_dims()
    env = ValidationEnvironment()

    with pytest.raises(exc.ValidationError) as exc_info:
        validate_aggregations(
            nodes.Formula(
                nodes.FuncCall.make(
                    name="+",
                    args=[
                        nodes.LiteralInteger.make(8),
                        nodes.FuncCall.make(
                            name="sum",
                            args=[
                                nodes.FuncCall.make(
                                    name="avg",
                                    args=[
                                        nodes.Field.make("Barley Field"),
                                    ],
                                ),
                            ],
                        ),
                    ],
                ),
            ),
            env=env,
            dimensions=dims,
        )
    assert exc_info.value.errors[0].code == exc.DoubleAggregationError.default_code


def test_window_sum_is_not_aggregation():
    # here an avg aggregation is wrapped into a sum window function - should not raise an error
    dims = get_dims()
    env = ValidationEnvironment()

    validate_aggregations(
        nodes.Formula(
            nodes.FuncCall.make(
                name="+",
                args=[
                    nodes.LiteralInteger.make(8),
                    nodes.WindowFuncCall.make(
                        name="sum",
                        args=[
                            nodes.FuncCall.make(
                                name="avg",
                                args=[
                                    nodes.Field.make("Barley Field"),
                                ],
                            ),
                        ],
                        grouping=nodes.WindowGroupingTotal(),
                    ),
                ],
            ),
        ),
        env=env,
        dimensions=dims,
    )
    assert True  # exception was not raised from the previous statement


def test_inconsistent_aggregation():
    dims = get_dims()
    dim1, dim2 = dims
    env = ValidationEnvironment()

    with pytest.raises(exc.ValidationError) as exc_info:
        validate_aggregations(
            node=nodes.Formula(
                nodes.FuncCall.make(
                    name="+",
                    args=[
                        nodes.FuncCall.make(name="sum", args=[nodes.Field.make("Some Field")]),
                        nodes.Field.make("Rye Field"),
                    ],
                ),
            ),
            env=env,
            dimensions=dims,
        )
    assert exc_info.value.errors[0].code == exc.InconsistentAggregationError.default_code

    with pytest.raises(exc.ValidationError) as exc_info:
        validate_aggregations(
            node=nodes.Formula(
                nodes.FuncCall.make(
                    name="+",
                    args=[
                        nodes.FuncCall.make(name="sum", args=[nodes.Field.make("Some Field")]),
                        nodes.FuncCall.make(name="some_func", args=[nodes.Field.make("Rye Field")]),
                        dim1,
                    ],
                ),
            ),
            env=env,
            dimensions=dims,
        )
    assert exc_info.value.errors[0].code == exc.InconsistentAggregationError.default_code


def test_window_func_of_agg_with_dimensions():
    # check window function that has aggregated and not aggregated children:
    # 1) aggregation func as an argument
    # 2) non-aggregated expression as a dimension
    # this should not raise an inconsistent agg error
    dims = get_dims()
    env = ValidationEnvironment()

    validate_aggregations(
        node=nodes.Formula(
            nodes.WindowFuncCall.make(
                name="rank",
                args=[
                    nodes.FuncCall.make(name="sum", args=[nodes.Field.make("Some Field")]),
                ],
                grouping=nodes.WindowGroupingWithin.make(dim_list=[*dims]),
            ),
        ),
        env=env,
        dimensions=dims,
    )
    assert True  # exception was not raised from the previous statement


def test_multiple_errors():
    dims = get_dims()
    env = ValidationEnvironment()

    def _check():
        with pytest.raises(exc.ValidationError) as exc_info:
            validate_aggregations(
                n.formula(
                    n.func.DO_SOMETHING(
                        n.func.SUM(n.func.AVG(n.field("Electromagnetic Field"))),  # 1. double aggregation
                        n.func.MAX(n.func.MIN(n.field("Field of Gold"))),  # 2. double aggregation
                        # repeat an expr that was already checked
                        n.func.SUM(n.func.AVG(n.field("Electromagnetic Field"))),  # 3. double aggregation (from cache)
                        n.func.MIN(n.field("Field of Gold")),
                        n.field("Field of Gold"),  # 4. inconsistent aggregation
                    ),
                ),
                env=env,
                collect_errors=True,
                dimensions=dims,
            )

        assert len(exc_info.value.errors) == 4  # see above
        assert {err.code for err in exc_info.value.errors} == {
            exc.InconsistentAggregationError.default_code,
            exc.DoubleAggregationError.default_code,
        }

    _check()
    # and now check again re-using the caches
    _check()

lib/dl_formula/dl_formula_tests/unit/validation/test_aggregation.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Tuple

import pytest

from dl_formula.core import (
    exc,
    nodes,
)
from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.validation.env import ValidationEnvironment
from dl_formula.validation.validator import validate
from dl_formula.validation.window import WindowFunctionChecker


def get_dims() -> Tuple[nodes.FormulaItem, nodes.FormulaItem]:
    return (
        nodes.Field.make("Dim Field"),
        nodes.FuncCall.make(
            name="+",
            args=[
                nodes.Field.make("Other Field"),
                nodes.FuncCall.make(
                    name="dim_func",
                    args=[
                        nodes.Field.make("Third Field"),
                    ],
                ),
            ],
        ),
    )


def validate_window_functions(
    node: nodes.FormulaItem, env: ValidationEnvironment, collect_errors: bool = False
) -> None:
    validate(
        node,
        env=env,
        checkers=[WindowFunctionChecker(inspect_env=InspectionEnvironment())],
        collect_errors=collect_errors,
    )


def test_no_errors():
    dim1, dim2 = get_dims()
    env = ValidationEnvironment()

    # no window function
    validate_window_functions(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="+",
                args=[
                    dim2,
                    nodes.LiteralFloat.make(1.1),
                ],
            ),
        ),
        env=env,
    )
    validate_window_functions(
        node=nodes.Formula(
            nodes.FuncCall.make(name="sum", args=[nodes.LiteralString.make("qwerty")]),
        ),
        env=env,
    )

    # with window function
    validate_window_functions(
        node=nodes.Formula(
            nodes.WindowFuncCall.make(
                name="rank",
                args=[
                    nodes.FuncCall.make(name="sum", args=[nodes.LiteralString.make("qwerty")]),
                ],
                grouping=nodes.WindowGroupingTotal(),
            ),
        ),
        env=env,
    )
    validate_window_functions(
        node=nodes.Formula(
            nodes.FuncCall.make(
                name="+",
                args=[
                    nodes.WindowFuncCall.make(
                        name="rank",
                        args=[
                            nodes.FuncCall.make(name="sum", args=[nodes.LiteralString.make("qwerty")]),
                        ],
                        grouping=nodes.WindowGroupingTotal(),
                    ),
                    nodes.WindowFuncCall.make(
                        name="rank_dense",
                        args=[
                            nodes.FuncCall.make(name="sum", args=[nodes.LiteralString.make("qwerty")]),
                        ],
                        grouping=nodes.WindowGroupingTotal(),
                    ),
                ],
            ),
        ),
        env=env,
    )


def test_nested_window_functions():
    env = ValidationEnvironment()

    with pytest.raises(exc.ValidationError) as exc_info:
        validate_window_functions(
            nodes.Formula(
                nodes.FuncCall.make(
                    name="+",
                    args=[
                        nodes.LiteralInteger.make(8),
                        nodes.WindowFuncCall.make(
                            name="rank",
                            args=[
                                nodes.WindowFuncCall.make(
                                    name="rank_dense",
                                    args=[
                                        nodes.FuncCall.make(
                                            name="avg",
                                            args=[
                                                nodes.Field.make("Barley Field"),
                                            ],
                                        ),
                                    ],
                                    grouping=nodes.WindowGroupingTotal(),
                                ),
                            ],
                            grouping=nodes.WindowGroupingTotal(),
                        ),
                    ],
                ),
            ),
            env=env,
        )
    assert exc_info.value.errors[0].code == exc.NestedWindowFunctionError.default_code


def test_window_function_without_aggregation():
    env = ValidationEnvironment()

    with pytest.raises(exc.ValidationError) as exc_info:
        validate_window_functions(
            nodes.Formula(
                nodes.FuncCall.make(
                    name="+",
                    args=[
                        nodes.LiteralInteger.make(8),
                        nodes.WindowFuncCall.make(
                            name="rank",
                            args=[
                                nodes.Field.make("Barley Field"),
                            ],
                            grouping=nodes.WindowGroupingTotal(),
                        ),
                    ],
                ),
            ),
            env=env,
        )
    assert exc_info.value.errors[0].code == exc.WindowFunctionWOAggregationError.default_code

lib/dl_formula/dl_formula_tests/unit/validation/test_window.py
datalens-tech/datalens-backend
import attr


@attr.s(frozen=True)
class Audience:
    name: str = attr.ib(kw_only=True, default="")
    default: bool = attr.ib(kw_only=True, default=False)


DEFAULT_AUDIENCE = Audience(default=True)

lib/dl_formula_ref/dl_formula_ref/audience.py
datalens-tech/datalens-backend
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.aliased_res import (
    AliasedTextResource,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.base import FunctionDocCategory


_ = get_gettext()

CATEGORY_AGGREGATION = FunctionDocCategory(
    name="aggregation",
    description=_(
        "Aggregate functions (or aggregations) are functions that combine multiple "
        "values from a group of entries into one, thus collapsing the group into a "
        "single entry.\n"
        "\n"
        "If you add an aggregation to a dimension, it becomes a measure.\n"
        "\n"
        "{text:agg_syntax}\n"
        "{text:agg_usage_restrictions}"
    ),
    resources=SimpleAliasedResourceRegistry(
        resources={
            "agg_syntax": AliasedTextResource(
                body=_(
                    "## Syntax {#syntax}\n"
                    "\n"
                    "In most cases aggregate functions have the same syntax as regular "
                    "functions:\n"
                    "{text:agg_simple_signature}\n"
                    "\n"
                    "For advanced cases, extended syntax may be required to indicate a custom "
                    "level of detail (LOD):\n"
                    "{text:agg_extended_signature}\n"
                    "{text:agg_syntax_lod}\n"
                    "{text:agg_syntax_bfb}\n"
                )
            ),
            "agg_simple_signature": AliasedTextResource(
                body=_("```\n" "AGGREGATE_FUNCTION_NAME(arg1, [arg2, ...])\n" "```")
            ),
            "agg_extended_signature": AliasedTextResource(
                body=_(
                    "```\n"
                    "<AGGREGATE_FUNCTION_NAME>(\n"
                    "    arg1, [arg2, ...]\n"
                    "\n"
                    "    [ FIXED dim1, dim2, ...\n"
                    "    | INCLUDE dim1, dim2, ...\n"
                    "    | EXCLUDE dim1, dim2, ... ]\n"
                    "\n"
                    "    [ BEFORE FILTER BY filtered_field1, ... ]\n"
                    ")\n"
                    "```"
                )
            ),
            "agg_syntax_lod": AliasedTextResource(
                body=_(
                    "### Level of Detail (LOD) {#syntax-lod}\n"
                    "\n"
                    "Custom LOD make possible nested aggregations and aggregations over the "
                    "entire set of rows or groups that are different from the grouping at the "
                    "chart's level.\n"
                    "\n"
                    "LOD can be specified using one of three keywords:\n"
                    "- `FIXED` — data is grouped using the listed dimensions (`dim1, dim2, ...`) "
                    "regardless of the dimensions used by the chart;\n"
                    "- `INCLUDE` — the listed dimensions (`dim1, dim2, ...`) are combined with "
                    "the chart's dimensions;\n"
                    "- `EXCLUDE` — all of the chart's dimensions are used with the exception of "
                    "those listed (`dim1, dim2, ...`).\n"
                    "\n"
                    "For any of these keywords the list may have any number of dimensions, or "
                    "even be empty.\n"
                    "\n"
                    "{text:agg_syntax_lod_inheritance}\n"
                    "{text:agg_syntax_lod_examples}\n"
                    "{text:agg_syntax_lod_consistency}\n"
                )
            ),
            "agg_syntax_lod_inheritance": AliasedTextResource(
                body=_(
                    "#### Dimension Inheritance {#syntax-lod-inheritance}\n"
                    "Dimensions are inherited by nested aggregations from the ones they are "
                    "inside of. The expression\n"
                    "```\n"
                    "AVG(MAX(SUM([Sales] INCLUDE [City]) INCLUDE [Category]))\n"
                    "```\n"
                    "in a chart with the additional dimension `[Date]` is equivalent to\n"
                    "```\n"
                    "AVG(MAX(SUM([Sales] FIXED [City], [Category], [Date]) FIXED [Category], "
                    "[Date]) FIXED [Date])\n"
                    "```\n"
                    "`INCLUDE` or `EXCLUDE` without a list of dimensions results in the same "
                    "dimensions as the aggregation above or dimensions of the chart if it is the "
                    "topmost aggregation. `FIXED` without a list means that all data is "
                    "aggregated in a single group, which can be used to calculate total values.\n"
                )
            ),
            "agg_syntax_lod_examples": AliasedTextResource(
                body=_(
                    "#### LOD Examples {#syntax-lod-examples}\n"
                    "\n"
                    "- average daily sum of `[Sales]`: `AVG(SUM([Sales] INCLUDE [Date]))`;\n"
                    "- ratio of the (daily) sum of `[Sales]` to the total sum: `SUM([Sales]) / "
                    "SUM([Sales] FIXED)`;\n"
                    "- sum of `[Sales]` of all orders that are smaller than average: "
                    "`SUM_IF(SUM([Sales] INCLUDE [Order ID]), SUM([Sales] INCLUDE [Order ID]) < "
                    "AVG([Sales] FIXED))`.\n"
                )
            ),
            "agg_syntax_lod_consistency": AliasedTextResource(
                body=_(
                    "#### Dimension Compatibility {#syntax-lod-compatibility}\n"
                    "\n"
                    "If several aggregations with custom LODs are nested inside another, their "
                    "sets of dimensions must be compatible, i.e. one of them must contain all of "
                    "the others.\n"
                    "\n"
                    "Invalid expression:\n"
                    "```\n"
                    "SUM(AVG([Sales] INCLUDE [City]) - AVG([Sales] INCLUDE [Category]))\n"
                    "```\n"
                    "One of the nested aggregations has dimension `[City]`, while the other has "
                    "`[Category]`, and there is no other that would contain both of these.\n"
                    "\n"
                    "Valid expression:\n"
                    "```\n"
                    "SUM(\n"
                    "    AVG([Sales] INCLUDE [City], [Category])\n"
                    "    - (AVG([Sales] INCLUDE [City]) + AVG([Sales] INCLUDE [Category])) / 2\n"
                    ")\n"
                    "```\n"
                    "\n"
                    "One of the nested aggregations' set of dimensions contains all of the "
                    "others.\n"
                    "\n"
                    "Another important restriction for LOD is that top-level aggregations must "
                    "not contain any dimensions not used in the chart."
                )
            ),
            "agg_syntax_bfb": AliasedTextResource(
                body=_(
                    "### BEFORE FILTER BY {#syntax-before-filter-by}\n"
                    "\n"
                    "If any fields are listed in `BEFORE FILTER BY`, then this aggregate function "
                    "is calculated before data is filtered using these fields.\n"
                    "\n"
                    "`BEFORE FILTER BY` applies to all nested aggregate functions too.\n"
                    "Example:\n"
                    "- Formula — `AVG(SUM([Sales] INCLUDE [Date]) BEFORE FILTER BY [City])`.\n"
                    "- Equivalent — `AVG(SUM([Sales] INCLUDE [Date] BEFORE FILTER BY [City]) "
                    "BEFORE FILTER BY [City])`.\n"
                    "\n"
                    "Do not use conflicting `BEFORE FILTER BY` clauses:\n"
                    "- Valid: `AVG(SUM([Sales] INCLUDE [Date] BEFORE FILTER BY [City], "
                    "[Category]) BEFORE FILTER BY [City])` — functions are nested and (`[City]`) "
                    "is a subset of (`[City], [Category]`).\n"
                    "- Valid: `AVG(SUM([Sales] INCLUDE [Date] BEFORE FILTER BY [Category]) BEFORE "
                    "FILTER BY [City])` — functions are nested, so field lists are combined in "
                    "the second of the two functions.\n"
                    "- Valid: `SUM([Sales] BEFORE FILTER BY [City], [Category]) - SUM([Sales] "
                    "BEFORE FILTER BY [City])` — (`[City]`) is a subset of (`[City], "
                    "[Category]`).\n"
                    "- Not valid: `SUM([Sales] BEFORE FILTER BY [Category]) - SUM([Sales] BEFORE "
                    "FILTER BY [City])` — functions are not nested and neither of (`[Category]`) "
                    "and (`[City]`) is a subset of the other."
                )
            ),
            "agg_usage_restrictions": AliasedTextResource(
                body=_(
                    "## Usage Restrictions {#usage-restrictions}\n"
                    "\n"
                    "There are the following features of using aggregations:\n"
                    "1. A function or operator cannot have aggregate and non-aggregate "
                    "expressions as its arguments simultaneously. The following usage is "
                    "forbidden: `CONCAT([Profit], SUM([Profit]))`.\n"
                )
            ),
        }
    ),
    keywords=_("aggregate functions,aggregations,aggregate data,level of detail,lod"),
)

lib/dl_formula_ref/dl_formula_ref/categories/aggregation.py
datalens-tech/datalens-backend
from dl_formula_ref.registry.base import FunctionDocCategory


CATEGORY_ARRAY = FunctionDocCategory(
    name="array",
    description="",
    keywords="",
)

lib/dl_formula_ref/dl_formula_ref/categories/array.py
datalens-tech/datalens-backend
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.aliased_res import (
    AliasedLinkResource,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.base import FunctionDocCategory


_ = get_gettext()

CATEGORY_DATE = FunctionDocCategory(
    name="date",
    description="",
    resources=SimpleAliasedResourceRegistry(
        resources={
            "iso_8601": AliasedLinkResource(url=_("https://en.wikipedia.org/wiki/ISO_8601")),
        }
    ),
    keywords="",
)

lib/dl_formula_ref/dl_formula_ref/categories/date.py
datalens-tech/datalens-backend
from dl_formula_ref.registry.base import FunctionDocCategory


CATEGORY_LOGICAL = FunctionDocCategory(
    name="logical",
    description="",
    keywords="",
)

lib/dl_formula_ref/dl_formula_ref/categories/logical.py
datalens-tech/datalens-backend
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.base import FunctionDocCategory


_ = get_gettext()

CATEGORY_MARKUP = FunctionDocCategory(
    name="markup",
    description=_(
        "Markup functions are used for creating marked up text (hypertext) from "
        "string values and/or other marked up text.\n"
        "\n"
        "## Usage notes {#usage-notes}\n"
        "\n"
        "There are the following features of using markup:\n"
        "1. These functions return `NULL` when any argument is `NULL`. To get a non-"
        "NULL value, wrap argumens in `IFNULL()`. Example: `bold(ifnull([value], "
        "'NULL'))`.\n"
        "1. Converting markup to a normal string is not currently possible.\n"
        "1. Markup functions can be used within logic functions. Example: "
        "`IF(STARTSWITH([value], 'n'), BOLD([value]), MARKUP([value]))`.\n"
    ),
    keywords="",
)

lib/dl_formula_ref/dl_formula_ref/categories/markup.py
datalens-tech/datalens-backend
from dl_formula_ref.registry.base import FunctionDocCategory


CATEGORY_MATHEMATICAL = FunctionDocCategory(
    name="numeric",  # FIXME: change to 'mathematical'
    description="",
    keywords="",
)

lib/dl_formula_ref/dl_formula_ref/categories/mathematical.py
datalens-tech/datalens-backend
from dl_formula_ref.registry.base import FunctionDocCategory


CATEGORY_OPERATOR = FunctionDocCategory(
    name="operator",
    description="",
    keywords="",
)

lib/dl_formula_ref/dl_formula_ref/categories/operator.py
datalens-tech/datalens-backend
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.aliased_res import (
    AliasedLinkResource,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.base import FunctionDocCategory


_ = get_gettext()

CATEGORY_STRING = FunctionDocCategory(
    name="string",
    description="",
    keywords="",
    resources=SimpleAliasedResourceRegistry(
        resources={
            "ch_re_link": AliasedLinkResource(url=_("https://github.com/google/re2/wiki/Syntax")),
        }
    ),
)

lib/dl_formula_ref/dl_formula_ref/categories/string.py
datalens-tech/datalens-backend
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.aliased_res import (
    AliasedTextResource,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.base import FunctionDocCategory


_ = get_gettext()

CATEGORY_TIME_SERIES = FunctionDocCategory(
    name="time-series",
    description=_(
        "Time series functions provide various ways to look up values corresponding "
        "to a specific time or offset along a given time axis.\n"
        "\n"
        "In a certain way this functionality is similar to the window function {ref:"
        "LAG}.\n"
        "The main difference is that `LAG` is indifferent to the actual values of the "
        "dimensions being used, and operates over positional offsets specified in "
        "_rows_, while time series functions use specific values and value offsets in "
        "date/time units like days, hours or seconds. This makes them sensitive to "
        'missing values in data. As a result of this `AGO(SUM([Sales]), [Date], "year'
        '")` will return `NULL` if the same-date row for the previous year is '
        "missing.\n"
        "{text:time_series_syntax}\n"
    ),
    resources=SimpleAliasedResourceRegistry(
        resources={
            "time_series_syntax": AliasedTextResource(
                body=_(
                    "## Syntax {#syntax}\n"
                    "\n"
                    "Time series functions support extended syntax:\n"
                    "{text:time_series_signature}\n"
                    "### BEFORE FILTER BY {#syntax-before-filter-by}\n"
                    "\n"
                    "If any fields are listed in `BEFORE FILTER BY`, then the function is "
                    "calculated before data is filtered using these fields.\n"
                    "Let's say you are trying to calculate the value of `[Sales]` a year ago from "
                    "today's date:\n"
                    "```\n"
                    'AGO([Sales], [Date], "year", 1)\n'
                    "```\n"
                    "and also have a filter in the chart limiting the dates to a specific year "
                    "(`2018 < [Year] <= 2019`, where `Year` is calculated as `YEAR([Date])`). All "
                    "data pertaining to the year 2018 will be omitted from the result. Because of "
                    "this the function will return `NULL`.\n"
                    "\n"
                    "If `BEFORE FILTER BY` is added to the function:\n"
                    "```\n"
                    'AGO([Sales], [Date], "year", 1 BEFORE FILTER BY [Year])\n'
                    "```\n"
                    "it will return the value of `[Sales]`.\n"
                    "\n"
                    "The date/time dimension specified as the second argument of time series "
                    "functions is implicitly added to the `BEFORE FILTER BY` clause. For example, "
                    "the following to formulas are equivalent:\n"
                    "```\n"
                    'AGO([Sales], [Date], "week")\n'
                    "```\n"
                    "```\n"
                    'AGO([Sales], [Date], "week" BEFORE FILTER BY [Date])\n'
                    "```\n"
                    "\n"
                    "### IGNORE DIMENSIONS {#syntax-ignore-dimensions}\n"
                    "\n"
                    "The `IGNORE DIMENSIONS` clause allows the exclusion of dimensions from the "
                    "search criteria. When looking up a value for a certain date, the search is "
                    "done by matching dimension values in each row against the values of the same "
                    "dimensions in the original row. If any of the other dimensions correlate "
                    "with the date dimension, then the data query may return an empty result "
                    "(`NULL`).\n"
                    "\n"
                    "For example, if the data request contains `[Date]`, `[Month]` and "
                    '`AGO([Sales], [Date], "month")`, then it will be impossible to find a row '
                    "with `[Date]` being a month earlier, but the value of `[Month]` being the "
                    "same as in the current row. This will result in `AGO` always returning "
                    "`NULL`.\n"
                    "\n"
                    "To get the correct value of `[Sales]` exclude `[Month]` using the `IGNORE "
                    "DIMENSIONS` clause:\n"
                    "```\n"
                    'AGO([Sales], [Date], "month" IGNORE DIMENSIONS [Month])\n'
                    "```"
                )
            ),
            "time_series_signature": AliasedTextResource(
                body=_(
                    "```\n"
                    "<FUNCTION_NAME>(\n"
                    "    arg1, arg2, ...\n"
                    "\n"
                    "    [ BEFORE FILTER BY filtered_field_1, ... ]\n"
                    "    [ IGNORE DIMENSIONS dimension_1, ... ]\n"
                    ")\n"
                    "```\n"
                )
            ),
        }
    ),
    keywords="",
)

lib/dl_formula_ref/dl_formula_ref/categories/time_series.py
datalens-tech/datalens-backend
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.aliased_res import (
    AliasedLinkResource,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.base import FunctionDocCategory


_ = get_gettext()

CATEGORY_TYPE_CONVERSION = FunctionDocCategory(
    name="type-conversion",
    description="",
    keywords="",
    resources=SimpleAliasedResourceRegistry(
        resources={
            "unix_ts": AliasedLinkResource(url=_("https://en.wikipedia.org/wiki/Unix_time")),
        }
    ),
)

lib/dl_formula_ref/dl_formula_ref/categories/type_conversion.py
datalens-tech/datalens-backend
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.aliased_res import (
    AliasedLinkResource,
    AliasedTableResource,
    AliasedTextResource,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.base import FunctionDocCategory


_ = get_gettext()

CATEGORY_WINDOW = FunctionDocCategory(
    name="window",
    description=_(
        "Window functions are calculated in the same way as aggregations, but they do "
        "not merge multiple entries into one. In some cases, this leads to duplication "
        "of values among entries in the same group (for example, `SUM(... TOTAL)`).\n"
        "\n"
        "Aggregate functions are calculated from groups of values that are determined "
        "by the dimension fields used in a data query: entries with matching dimension "
        "values are grouped. Window functions are also calculated over groups of "
        "entries called _windows_. In this case, you should specify grouping "
        "parameters in the function call as a list of dimensions to be included "
        "(`WITHIN ...`) or excluded (`AMONG ...`) from the grouping.\n"
        "{text:window_usage_restrictions}\n"
        "{text:window_syntax}\n"
        "{text:agg_as_window}\n"
        "{text:window_examples}\n"
    ),
    resources=SimpleAliasedResourceRegistry(
        resources={
            "window_usage_restrictions": AliasedTextResource(
                body=_(
                    "## Usage Restrictions {#usage-restrictions}\n"
                    "\n"
                    "1. The first argument in window functions can only be {link:"
                    "field_concept_link:measures}.\n"
                    "For the `AVG_IF`, `COUNT_IF`, `SUM_IF` window functions, the first argument "
                    "(`expression` in the function description) must always be a measure.\n"
                    "    Example:\n"
                    "\n"
                    "    `AVG_IF([Profit], [Profit] > 5)`\n"
                    "\n"
                    "For other window functions, the first (and only) argument (`value` in the "
                    "function description) must always be a measure, too.\n"
                    "\n"
                    "    Examples:\n"
                    "    - Valid: `SUM(SUM([Profit]) TOTAL)`.\n"
                    "    - Not valid: `RANK([Profit] TOTAL)`, where `[Profit]` is a non-"
                    "aggregated expression.\n"
                    "\n"
                    "1. For grouping window functions, only the {link:field_concept_link:"
                    "dimensions} used to build the chart can be applied. Only the dimensions used "
                    "to build the chart set the {link:chart_aggregation_link:grouping when "
                    "calculating a measure}. These dimensions define how values are split into "
                    "groups and therefore have fixed values in each group.\n"
                    "\n"
                    "If you specify a dimension that was not used to build the chart, it won't "
                    "have a fixed value and the value can be different in each group row. As a "
                    "result, it will be impossible to determine which value of this dimension "
                    "must be used to calculate the measure. This limitation applies to the "
                    "`WITHIN` and `AMONG` grouping types.\n"
                    "\n"
                    "    Examples:\n"
                    "    - Valid: `RANK(SUM([Profit]) WITHIN [Category])` in the chart with "
                    "grouping by the `[Order Date]` and `[Category]` dimensions.\n"
                    "    - Allowed: `RANK(SUM([Profit]) WITHIN [City])` in the chart with "
                    "grouping by the `[Category]` dimension,       `[City]` does not participate "
                    "in the grouping.\n"
                    "    - Not valid: `RANK(SUM([Profit]) WITHIN [Category])` in a chart with "
                    "grouping by the `[Order Date]` and `[City]` dimensions.\n"
                    "    - Not valid: `RANK(SUM([Profit]) AMONG [City])` in a chart with grouping "
                    "by the `[Order Date]` and `[Category]` dimensions.\n"
                    "1. The **Filters** section doesn't affect the chart grouping, so if the "
                    "dimension is only in this chart section, you can't use it in the window "
                    "function.\n"
                    "\n"
                    "    Example:\n"
                    "    - Chart type: **Table**.\n"
                    "    - In the **Columns** section, the `Category` dimension and the "
                    "`SUM(SUM([Sales] BEFORE FILTER BY [Date])` expression are added.\n"
                    "    - The `Date` dimension is added to the **Filters** section.\n"
                    "\n"
                    "    This will result in an error because the `Date` dimension isn't used to "
                    "build the chart.\n"
                    "1. If a window function is used to build a **Table** chart, we don't "
                    "recommend enabling the display of **Total** in the settings. This may cause "
                    "an error.\n"
                )
            ),
            "window_func_signature": AliasedTextResource(
                body=_(
                    "```\n"
                    "<WINDOW_FUNCTION_NAME>(\n"
                    "    arg1, arg2, ...\n"
                    "\n"
                    "    [ TOTAL\n"
                    "    | WITHIN dim1, dim2, ...\n"
                    "    | AMONG dim1, dim2, ... ]\n"
                    "\n"
                    "    [ ORDER BY field1, field2, ... ]\n"
                    "\n"
                    "    [ BEFORE FILTER BY filtered_field1, ... ]\n"
                    ")\n"
                    "```\n"
                )
            ),
            "window_syntax": AliasedTextResource(
                body=_(
                    "## Syntax {#syntax}\n"
                    "\n"
                    "The general syntax for window functions is as follows:\n"
                    "{text:window_func_signature}\n"
                    "It starts off, just like a regular function call, with its name and "
                    "arguments (`arg1, arg2, ...` in this case).\n"
                    "\n"
                    "### Grouping {#syntax-grouping}\n"
                    "\n"
                    "The arguments are followed by a window grouping, which can be one of three "
                    "types:\n"
                    "- `TOTAL` (equivalent to `WITHIN` without dimensions): all query entries "
                    "fall into a single window.\n"
                    "- `WITHIN dim1, dim2, ...` : records are grouped by the dimensions `dim1, "
                    "dim2, ...`.\n"
                    "- `AMONG dim1, dim2, ...` : records are grouped by all dimensions from the "
                    "query, except those listed. For example, if we use formula "
                    "`RSUM(SUM([Sales]) AMONG dim1, dim2)` with dimensions `dim1`, `dim2`, "
                    "`dim3`, `dim4` in the data query, then the entries will be grouped by `dim3` "
                    "and `dim4`, so it will be equivalent to `RSUM([Sales] WITHIN dim3, dim4)`.\n"
                    "\n"
                    "The grouping clause is optional. `TOTAL` is used by default.\n"
                    "\n"
                    "{text:window_syntax_ordering}\n"
                    "{text:window_syntax_bfb}"
                )
            ),
            "window_syntax_ordering": AliasedTextResource(
                body=_(
                    "### Ordering {#syntax-order-by}\n"
                    "\n"
                    "After the grouping comes the ordering clause. It is only supported for order-"
                    "dependent functions:\n"
                    "\n"
                    "{table:order_dep_win_functions}\n"
                    "\n"
                    "The ordering clause is optional for these functions.\n"
                    "\n"
                    "See the descriptions of these functions for more information on how this "
                    "order affects the result value.\n"
                    "The `ORDER BY` clause accepts dimensions as well as measures. It also "
                    "supports the standard `ASC`/`DESC` syntax (`ASC` is assumed by default) to "
                    "specify ascending or descending order respectively:\n"
                    "`... ORDER BY [Date] ASC, SUM([Sales]) DESC, [Category] ...`\n"
                    "\n"
                    "Fields listed in `ORDER BY` are combined with fields listed in the chart's "
                    "sorting section.\n"
                    "Example:\n"
                    "- Function — `... ORDER BY [Date] DESC, [City]`.\n"
                    "- Chart — Sorted by `Date` and `Category`.\n"
                    "- Result — `Date` (descending), `City`, `Category`.\n"
                )
            ),
            "order_dep_win_functions": AliasedTableResource(
                table_body=[
                    ["`M*`", "`R*`", _("Positional functions")],
                    ["{ref:MAVG}", "{ref:RAVG}", "{ref:LAG}"],
                    ["{ref:MCOUNT}", "{ref:RCOUNT}", "{ref:FIRST}"],
                    ["{ref:MMAX}", "{ref:RMAX}", "{ref:LAST}"],
                    ["{ref:MMIN}", "{ref:RMIN}", ""],
                    ["{ref:MSUM}", "{ref:RSUM}", ""],
                ]
            ),
            "window_syntax_bfb": AliasedTextResource(
                body=_(
                    "### BEFORE FILTER BY {#syntax-before-filter-by}\n"
                    "\n"
                    "If any fields are listed in `BEFORE FILTER BY`, then this window function is "
                    "calculated before data is filtered using these fields.\n"
                    "\n"
                    "`BEFORE FILTER BY` applies to all nested window functions too.\n"
                    "Example:\n"
                    "- Formula — `MAVG(RSUM([Sales]), 10 BEFORE FILTER BY [Date])`.\n"
                    "- Equivalent — `MAVG(RSUM([Sales] BEFORE FILTER BY [Date]), 10 BEFORE FILTER "
                    "BY [Date])`.\n"
                    "\n"
                    "Do not use conflicting `BEFORE FILTER BY` clauses:\n"
                    "- Valid: `MAVG(RSUM([Sales] BEFORE FILTER BY [Date], [Category]), 10 BEFORE "
                    "FILTER BY [Date])` — functions are nested and (`[Date]`) is a subset of "
                    "(`[Date], [Category]`).\n"
                    "- Valid: `MAVG(RSUM([Sales] BEFORE FILTER BY [Category]), 10 BEFORE FILTER "
                    "BY [Date])` — functions are nested, so field lists are combined in the "
                    "second of the two functions.\n"
                    "- Valid: `RSUM([Sales] BEFORE FILTER BY [Date], [Category]) - RSUM([Sales] "
                    "BEFORE FILTER BY [Date])` — (`[Date]`) is a subset of (`[Date], "
                    "[Category]`).\n"
                    "- Not valid: `RSUM([Sales] BEFORE FILTER BY [Category]) - RSUM([Sales] "
                    "BEFORE FILTER BY [Date])` — functions are not nested and neither of "
                    "(`[Category]`) and (`[Date]`) is a subset of the other.\n"
                )
            ),
            "agg_as_window": AliasedTextResource(
                body=_(
                    "## Aggregate Functions as Window Functions {#aggregate-functions-as-window-functions}\n"
                    "\n"
                    "The following aggregations can also be used as window functions:\n"
                    "\n"
                    "{table:agg_as_window_table}\n"
                    "\n"
                    "To use the window version of the aggregate functions, you must explicitly "
                    "specify the grouping (unlike other window functions, where it is optional).\n"
                    "\n"
                    "Example:\n"
                    "- `SUM([Sales]) / SUM(SUM([Sales]) TOTAL)` can be used to calculate the "
                    "ratio of a group's sum of `[Sales]` to the sum of `[Sales]` among all "
                    "entries.\n"
                )
            ),
            "agg_as_window_table": AliasedTableResource(
                table_body=[
                    [_("Aggregations"), _("Conditional Aggregations")],
                    ["{ref:window/SUM}", "{ref:window/SUM_IF}"],
                    ["{ref:window/COUNT}", "{ref:window/COUNT_IF}"],
                    ["{ref:window/AVG}", "{ref:window/AVG_IF}"],
                    ["{ref:window/MAX}", ""],
                    ["{ref:window/MIN}", ""],
                ]
            ),
            "window_examples": AliasedTextResource(body=""),  # TODO
        }
    ),
    keywords="",
)

lib/dl_formula_ref/dl_formula_ref/categories/window.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import Enum
import os
from typing import (
    Any,
    Mapping,
    NamedTuple,
)

import attr
from dynamic_enum import (
    AutoEnumValue,
    DynamicEnum,
)

from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.scope import Scope
from dl_formula_ref.audience import (
    DEFAULT_AUDIENCE,
    Audience,
)
from dl_formula_ref.paths import (
    CatPathTemplate,
    FuncPathTemplate,
)
from dl_formula_ref.registry.aliased_res import (
    AliasedResourceRegistryBase,
    SimpleAliasedResourceRegistry,
)


class FuncDocTemplateConfig(NamedTuple):
    template_file: str
    func_file_path: FuncPathTemplate  # where the doc file should be generated
    cat_file_path: CatPathTemplate


class FuncDocConfigVersion(Enum):
    overview_shortcut = "overview_shortcut"
    overview = "overview"
    summary = "summary"


class ConfigVersion(DynamicEnum):
    default = AutoEnumValue()


@attr.s
class RefDocGeneratorConfig:
    template_dir_rel: str = attr.ib(kw_only=True, default=os.path.join("templates"))  # rel to dl_formula_ref pkg
    func_doc_configs: dict[FuncDocConfigVersion, FuncDocTemplateConfig] = attr.ib(kw_only=True)

    # Templates
    doc_list_template: str = attr.ib(kw_only=True, default="doc_list.md.jinja")
    doc_avail_template: str = attr.ib(kw_only=True, default="doc_availability.md.jinja")
    doc_example_template: str = attr.ib(kw_only=True, default="doc_example.md.jinja")

    # Destination paths
    doc_toc_filename: str = attr.ib(kw_only=True)
    doc_all_filename: str = attr.ib(kw_only=True)
    doc_avail_filename: str = attr.ib(kw_only=True)

    db_config_file: str = attr.ib(
        kw_only=True, default=os.path.join(os.path.dirname(__file__), "db_config.json")
    )  # FIXME: move somewhere else
    example_data_file: str = attr.ib(kw_only=True, default=os.path.join(os.path.dirname(__file__), "example_data.json"))
    gen_availability_table: bool = attr.ib(kw_only=True, default=True)

    function_scopes: dict[Audience, int] = attr.ib(kw_only=True, default={DEFAULT_AUDIENCE: Scope.DOCUMENTED})
    block_conditions: Mapping[str, bool] = attr.ib(kw_only=True, factory=dict)

    supported_locales: frozenset[str] = attr.ib(kw_only=True, default=frozenset({"en"}))
    supported_dialects: frozenset[DialectCombo] = attr.ib(kw_only=True, default=frozenset({D.DUMMY}))
    default_example_dialect: DialectCombo = attr.ib(kw_only=True, default=D.DUMMY)

    resource_overrides: AliasedResourceRegistryBase = attr.ib(kw_only=True, factory=SimpleAliasedResourceRegistry)

    def clone(self, **kwargs: Any) -> RefDocGeneratorConfig:
        return attr.evolve(self, **kwargs)


DOC_GEN_CONFIG_DEFAULT = RefDocGeneratorConfig(
    func_doc_configs={
        FuncDocConfigVersion.overview_shortcut: FuncDocTemplateConfig(
            template_file="doc_func_long.md.jinja",
            func_file_path=FuncPathTemplate("function-ref/{func_name}.md"),
            cat_file_path=CatPathTemplate("function-ref/{category_name}-functions.md"),
        ),
    },
    doc_toc_filename="toc.yaml",
    doc_all_filename="function-ref/all.md",
    doc_avail_filename="function-ref/availability.md",
)


_CONFIGS_BY_VERSION: dict[ConfigVersion, RefDocGeneratorConfig] = {}


def register_config_version(version: ConfigVersion, config: RefDocGeneratorConfig) -> None:
    try:
        assert _CONFIGS_BY_VERSION[version] is config
    except KeyError:
        _CONFIGS_BY_VERSION[version] = config


def get_generator_config(version: ConfigVersion) -> RefDocGeneratorConfig:
    return _CONFIGS_BY_VERSION[version]


register_config_version(ConfigVersion.default, DOC_GEN_CONFIG_DEFAULT)

lib/dl_formula_ref/dl_formula_ref/config.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    List,
    Optional,
    Sequence,
    Tuple,
)

import attr

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo
from dl_formula_ref.examples.data_table import DataTable
from dl_formula_ref.texts import EXAMPLE_TITLE


@attr.s(frozen=True)
class ExampleSource:
    columns: List[Tuple[str, DataType]] = attr.ib(kw_only=True)
    data: List[List[Any]] = attr.ib(kw_only=True)


@attr.s(frozen=True)
class ExampleConfig:
    source: ExampleSource = attr.ib(kw_only=True)
    # e.g. formula_fields = (('My Field', 'STR([col1])'),)
    formula_fields: List[Tuple[str, str]] = attr.ib(kw_only=True)
    dialect: Optional[DialectCombo] = attr.ib(kw_only=True, default=None)
    name: str = attr.ib(kw_only=True, default=EXAMPLE_TITLE)
    # group_by and order_by are lists of formulas
    group_by: Sequence[str] = attr.ib(kw_only=True, default=())
    order_by: Sequence[str] = attr.ib(kw_only=True, default=())
    show_source_table: bool = attr.ib(kw_only=True, default=False)
    formulas_as_names: bool = attr.ib(kw_only=True, default=True)
    float_format: str = attr.ib(kw_only=True, default="{:.02f}")
    # Sequence of transformations described by formulas just like in `formula_fields`
    # They are applied to the result table in the same sequence as they are defined.
    additional_transformations: Sequence[List[Tuple[str, str]]] = attr.ib(kw_only=True, default=())
    # If additional transformations are applied,
    # then one might want to disguise real formulas used for the result with these overrides
    override_formula_fields: Optional[List[Tuple[str, str]]] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class PredefinedExampleConfig(ExampleConfig):
    result_table: DataTable = attr.ib(kw_only=True)

lib/dl_formula_ref/dl_formula_ref/examples/config.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Dict,
)

import attr

from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.scope import Scope
from dl_formula_ref.examples.config import (
    ExampleConfig,
    ExampleSource,
)
from dl_formula_ref.examples.data_table import DataTable
from dl_formula_ref.examples.dumper import (
    DataDumper,
    get_dumper,
)
from dl_formula_ref.examples.query_gen import QueryGenerator
from dl_formula_ref.examples.result_storage import WritableDataStorage
from dl_formula_ref.examples.sa_compiler import SaQueryCompiler
from dl_formula_ref.examples.utils import (
    make_data_table_from_example,
    make_key_for_example,
)


if TYPE_CHECKING:
    from dl_formula.core.dialect import DialectCombo
    from dl_formula_testing.database import Db


@attr.s
class DataPreparer:
    _db_by_dialect: Dict[DialectCombo, Db] = attr.ib(kw_only=True)
    _storage_filename: str = attr.ib(kw_only=True)
    _function_scopes: int = attr.ib(kw_only=True, default=Scope.EXPLICIT_USAGE)
    _query_gen: QueryGenerator = attr.ib(kw_only=True, factory=QueryGenerator)
    _dumper_by_dialect: Dict[DialectCombo, DataDumper] = attr.ib(init=False, factory=dict)
    _storage: WritableDataStorage = attr.ib(init=False)
    _sa_compiler_by_dialect: Dict[DialectCombo, SaQueryCompiler] = attr.ib(init=False, factory=dict)
    _default_example_dialect: DialectCombo = attr.ib(kw_only=True, default=D.DUMMY)

    def __attrs_post_init__(self) -> None:
        self._storage = WritableDataStorage(filename=self._storage_filename)

    def _get_example_dialect(self, example: ExampleConfig) -> DialectCombo:
        if example.dialect is not None:
            return example.dialect
        return self._default_example_dialect

    def _get_db(self, dialect: DialectCombo) -> Db:
        return self._db_by_dialect[dialect]

    def _get_dumper(self, dialect: DialectCombo) -> DataDumper:
        if dialect not in self._dumper_by_dialect:
            self._dumper_by_dialect[dialect] = get_dumper(db=self._get_db(dialect=dialect))
        return self._dumper_by_dialect[dialect]

    def _get_sa_compiler(self, dialect: DialectCombo) -> SaQueryCompiler:
        if dialect not in self._sa_compiler_by_dialect:
            self._sa_compiler_by_dialect[dialect] = SaQueryCompiler(
                dialect=dialect,
                function_scopes=self._function_scopes,
            )
        return self._sa_compiler_by_dialect[dialect]

    def apply_data_transformation(self, example: ExampleConfig, data_table: DataTable) -> DataTable:
        dumper = self._get_dumper(dialect=self._get_example_dialect(example))
        sa_compiler = self._get_sa_compiler(dialect=self._get_example_dialect(example))
        with dumper.temporary_data_table(data_table=data_table) as table_ref:
            raw_query = self._query_gen.generate_query(example=example)
            compiled_query_ctx = sa_compiler.compile_query(raw_query=raw_query, table_ref=table_ref)
            result_data = dumper.execute_query(query_ctx=compiled_query_ctx)
        return result_data

    def generate_example_data(self, func_name: str, example: ExampleConfig) -> None:
        storage_key = make_key_for_example(func_name=func_name, example=example)
        data_table = make_data_table_from_example(example=example)
        data_table = self.apply_data_transformation(example=example, data_table=data_table)

        # Apply additional transformations
        for formula_fields in example.additional_transformations:
            source_from_result = ExampleSource(
                columns=[(col.name, col.data_type) for col in data_table.columns],
                data=data_table.rows,
            )
            example_from_result = ExampleConfig(
                source=source_from_result,
                formula_fields=formula_fields,
                dialect=self._get_example_dialect(example),
                order_by=example.order_by,
            )
            data_table = self.apply_data_transformation(example=example_from_result, data_table=data_table)

        self._storage.set_result_data(storage_key=storage_key, result_data=data_table)

    def save(self) -> None:
        self._storage.save()

lib/dl_formula_ref/dl_formula_ref/examples/data_prep.py
datalens-tech/datalens-backend
from itertools import chain
from typing import (
    Any,
    Dict,
    List,
)

import attr

from dl_formula.core.datatype import DataType


@attr.s(frozen=True, slots=True)
class DataColumn:
    name: str = attr.ib(kw_only=True)
    data_type: DataType = attr.ib(kw_only=True)


@attr.s(frozen=True, slots=True)
class DataTable:
    columns: List[DataColumn] = attr.ib(kw_only=True)
    rows: List[List[Any]] = attr.ib(kw_only=True)


def horizontal_join(*tables: DataTable) -> DataTable:
    lengths = {len(t.rows) for t in tables}
    assert len(lengths) == 1, f"All tables must have the same number of rows for horizontal join, got {lengths}"
    return DataTable(
        columns=[col for t in tables for col in t.columns],
        rows=[list(chain.from_iterable(concat_row_parts)) for concat_row_parts in zip(*(t.rows for t in tables))],
    )


def rename_columns(table: DataTable, name_mapping: Dict[str, str]) -> DataTable:
    return DataTable(
        columns=[DataColumn(name=name_mapping[col.name], data_type=col.data_type) for col in table.columns],
        rows=table.rows,
    )

lib/dl_formula_ref/dl_formula_ref/examples/data_table.py
datalens-tech/datalens-backend
from contextlib import contextmanager
from typing import (
    Any,
    Generator,
)
import uuid

import attr
import sqlalchemy as sa
from sqlalchemy.types import TypeEngine

from dl_formula.connectors.base.type_constructor import get_type_constructor
from dl_formula.core.datatype import DataType
from dl_formula_ref.examples.data_table import (
    DataColumn,
    DataTable,
)
from dl_formula_ref.examples.query import (
    CompiledQueryContext,
    TableReference,
)
from dl_formula_testing.database import Db


def _make_name() -> str:
    return f"t_{uuid.uuid4().hex[:10]}"


@attr.s
class DataDumper:
    _db: Db = attr.ib(kw_only=True)

    def get_sa_type(self, data_type: DataType) -> TypeEngine:
        type_constructor = get_type_constructor(dialect_name=self._db.dialect.common_name)
        return type_constructor.get_sa_type(data_type)

    def generate_sa_table(self, table_ref: TableReference) -> sa.Table:
        sa_columns = [sa.Column(name=col.name, type_=self.get_sa_type(col.data_type)) for col in table_ref.columns]
        sa_table = self._db._engine_wrapper.table_from_columns(columns=sa_columns, table_name=table_ref.name)  # noqa
        return sa_table

    def create_table(self, table: DataTable) -> TableReference:
        name = _make_name()
        table_ref = TableReference(name=name, columns=table.columns)
        sa_table = self.generate_sa_table(table_ref)
        self._db.create_table(sa_table)
        return table_ref

    def dump_data(self, table_ref: TableReference, table: DataTable) -> None:
        data_as_dicts = [{col.name: value for col, value in zip(table.columns, row)} for row in table.rows]
        sa_table = self.generate_sa_table(table_ref)
        self._db.insert_into_table(sa_table, data_as_dicts)

    def remove_table(self, table_ref: TableReference) -> None:
        sa_table = self.generate_sa_table(table_ref)
        self._db.drop_table(sa_table)

    def create_and_dump_table(self, data_table: DataTable) -> TableReference:
        table_ref = self.create_table(table=data_table)
        self.dump_data(table_ref=table_ref, table=data_table)
        return table_ref

    @contextmanager
    def temporary_data_table(self, data_table: DataTable) -> Generator[TableReference, None, None]:
        table_ref = self.create_and_dump_table(data_table=data_table)
        try:
            yield table_ref
        finally:
            self.remove_table(table_ref)

    def execute_query(self, query_ctx: CompiledQueryContext) -> DataTable:
        def convert(val: Any, col: DataColumn) -> Any:
            # a workaround for missing boolean type in DBs and no coercion to bool
            if col.data_type in (DataType.CONST_BOOLEAN, DataType.BOOLEAN) and val in (0, 1):
                val = bool(val)
            return val

        rows = [
            [convert(val, col) for val, col in zip(row, query_ctx.result_columns)]
            for row in self._db.execute(query_ctx.sa_query).fetchall()
        ]
        result_data_table = DataTable(columns=list(query_ctx.result_columns), rows=rows)
        return result_data_table


def get_dumper(db: Db) -> DataDumper:
    return DataDumper(db=db)

lib/dl_formula_ref/dl_formula_ref/examples/dumper.py
datalens-tech/datalens-backend
from typing import (
    Optional,
    Sequence,
)

import attr
from sqlalchemy.sql.selectable import Select

import dl_formula.core.nodes as nodes
from dl_formula_ref.examples.data_table import DataColumn


@attr.s
class TableReference:
    name: str = attr.ib(kw_only=True)
    columns: Sequence[DataColumn] = attr.ib(kw_only=True)


@attr.s
class FormulaContext:
    formula: nodes.Formula = attr.ib(kw_only=True)
    alias: Optional[str] = attr.ib(kw_only=True)


@attr.s
class RawQueryContext:
    result_columns: Sequence[DataColumn] = attr.ib(kw_only=True)
    select: Sequence[FormulaContext] = attr.ib(kw_only=True)
    group_by: Sequence[FormulaContext] = attr.ib(kw_only=True, default=())
    order_by: Sequence[FormulaContext] = attr.ib(kw_only=True, default=())


@attr.s
class CompiledQueryContext:
    result_columns: Sequence[DataColumn] = attr.ib(kw_only=True)
    sa_query: Select = attr.ib(kw_only=True)

lib/dl_formula_ref/dl_formula_ref/examples/query.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    List,
    Optional,
)

import attr

from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.inspect.expression import infer_data_type
from dl_formula.parser.factory import (
    ParserType,
    get_parser,
)
from dl_formula_ref.examples.config import ExampleConfig
from dl_formula_ref.examples.data_table import DataColumn
from dl_formula_ref.examples.query import (
    FormulaContext,
    RawQueryContext,
)


if TYPE_CHECKING:
    from dl_formula.parser.base import FormulaParser


@attr.s
class QueryGenerator:
    _parser_type: ParserType = attr.ib(kw_only=True, default=ParserType.antlr_py)
    _parser: FormulaParser = attr.ib(init=False)

    def __attrs_post_init__(self) -> None:
        self._parser = get_parser(self._parser_type)

    def _make_formula_ctx(self, formula: str, alias: Optional[str] = None) -> FormulaContext:
        formula_obj = self._parser.parse(formula)
        return FormulaContext(formula=formula_obj, alias=alias)

    def generate_query(self, example: ExampleConfig) -> RawQueryContext:
        inspect_env = InspectionEnvironment()

        columns = [DataColumn(name=name, data_type=data_type) for name, data_type in example.source.columns]

        all_formulas = {alias: formula for formula, alias in example.formula_fields}
        all_formulas.update({col.name: f"[{col.name}]" for col in columns})

        col_data_types = {col.name: col.data_type for col in columns}

        select = [self._make_formula_ctx(formula, alias) for alias, formula in example.formula_fields]
        group_by = [self._make_formula_ctx(formula) for formula in example.group_by]
        order_by = [self._make_formula_ctx(formula) for formula in example.order_by]
        result_columns: List[DataColumn] = []
        for ctx in select:
            alias = ctx.alias
            assert alias is not None
            result_columns.append(
                DataColumn(
                    name=alias,
                    data_type=infer_data_type(
                        node=ctx.formula,
                        field_types=col_data_types,
                        env=inspect_env,
                    ),
                )
            )

        query_ctx = RawQueryContext(
            result_columns=result_columns,
            select=select,
            group_by=group_by,
            order_by=order_by,
        )
        return query_ctx

lib/dl_formula_ref/dl_formula_ref/examples/query_gen.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from decimal import Decimal
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Optional,
    Sequence,
    Tuple,
)

import attr
import jinja2
import tabulate

from dl_formula_ref.examples.data_table import rename_columns
from dl_formula_ref.examples.result_storage import ReadableDataStorage
from dl_formula_ref.examples.utils import (
    make_data_table_from_example,
    make_key_for_example,
)
from dl_formula_ref.i18n.registry import get_localizer


if TYPE_CHECKING:
    from dl_formula_ref.examples.config import ExampleConfig
    from dl_formula_ref.examples.data_table import DataTable


@attr.s(frozen=True)
class ExampleResult:
    name: str = attr.ib(kw_only=True)
    result: str = attr.ib(kw_only=True)
    source: Optional[str] = attr.ib(kw_only=True, default=None)
    group_by: Sequence[str] = attr.ib(kw_only=True, factory=list)
    order_by: Sequence[str] = attr.ib(kw_only=True, factory=list)
    formula_fields: Sequence[Tuple[str, str]] = attr.ib(kw_only=True, factory=list)


@attr.s
class ExampleRenderer:
    _example_template_filename: str = attr.ib(kw_only=True)
    _jinja_env: jinja2.Environment = attr.ib(kw_only=True)
    _locale: str = attr.ib(kw_only=True)
    _storage_filename: str = attr.ib(kw_only=True)
    _storage: ReadableDataStorage = attr.ib(init=False)
    _trans: Callable[[str], str] = attr.ib(init=False)

    def __attrs_post_init__(self) -> None:
        self._storage = ReadableDataStorage(filename=self._storage_filename)
        self._storage.load()
        self._trans = get_localizer(self._locale).translate

    def _format_value(self, value: Any, example: ExampleConfig) -> str:
        if value is None:
            return "NULL"
        if isinstance(value, (float, Decimal)):
            return example.float_format.format(value)
        if isinstance(value, (datetime.datetime, datetime.date)):
            return value.isoformat()
        return repr(value)

    def _md_escape(self, text: str) -> str:
        return text.replace("*", "&ast;")

    def _render_table(self, table: DataTable, example: ExampleConfig) -> str:
        return tabulate.tabulate(
            [[f"`{self._format_value(cell, example=example)}`" for cell in row] for row in table.rows],
            headers=[f"**{self._md_escape(col.name)}**" for col in table.columns],
            tablefmt="pipe",
        )

    def render_example(self, example: ExampleConfig, result_table: DataTable, under_cut: bool) -> str:
        source_table = make_data_table_from_example(example=example)

        if example.formulas_as_names:
            if example.override_formula_fields is not None:
                name_mapping = {name: formula for name, formula in example.override_formula_fields}
            elif not example.additional_transformations:
                name_mapping = {name: formula for name, formula in example.formula_fields}
            else:
                name_mapping = {name: formula for name, formula in example.additional_transformations[-1]}
            result_table = rename_columns(result_table, name_mapping=name_mapping)

        example_result = ExampleResult(
            name=example.name,
            order_by=example.order_by,
            group_by=example.group_by,
            source=self._render_table(source_table, example=example) if example.show_source_table else None,
            result=self._render_table(result_table, example=example),
            formula_fields=example.formula_fields if not example.formulas_as_names else [],
        )
        template = self._jinja_env.get_template(self._example_template_filename)
        return template.render(example_result=example_result, under_cut=under_cut, _=self._trans)

    def _get_result_table_from_storage(self, func_name: str, example: ExampleConfig) -> DataTable:
        storage_key = make_key_for_example(func_name=func_name, example=example)
        result_table = self._storage.get_result_data(storage_key)
        return result_table

    def render_example_from_storage(self, func_name: str, example: ExampleConfig, under_cut: bool) -> str:
        result_table = self._get_result_table_from_storage(func_name=func_name, example=example)
        return self.render_example(example=example, result_table=result_table, under_cut=under_cut)

lib/dl_formula_ref/dl_formula_ref/examples/renderer.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from decimal import Decimal
import json
from typing import (
    Any,
    Dict,
)

import attr

from dl_formula.core.datatype import DataType
from dl_formula_ref.examples.data_table import (
    DataColumn,
    DataTable,
)


@attr.s(frozen=True)
class StorageKey:
    item_key: str = attr.ib(kw_only=True)
    example_key: str = attr.ib(kw_only=True)

    @property
    def as_str(self) -> str:
        return f"{self.item_key}.{self.example_key}"


class CustomJSONEncoder(json.JSONEncoder):
    def default(self, obj) -> Any:
        if isinstance(obj, datetime.datetime):
            return {"__type__": "datetime", "__value__": obj.isoformat()}
        if isinstance(obj, datetime.date):
            return {"__type__": "date", "__value__": obj.isoformat()}
        if isinstance(obj, Decimal):
            return {"__type__": "decimal", "__value__": str(obj)}
        return super().default(obj)


class CustomJSONDecoder(json.JSONDecoder):
    def __init__(self, *args, **kwargs):
        super().__init__(object_hook=self.object_hook, *args, **kwargs)

    def object_hook(self, obj_dict: Dict[str, Any]) -> Any:
        if obj_dict.get("__type__") == "datetime":
            return datetime.datetime.fromisoformat(obj_dict["__value__"])
        if obj_dict.get("__type__") == "date":
            return datetime.date.fromisoformat(obj_dict["__value__"])
        if obj_dict.get("__type__") == "decimal":
            return Decimal(obj_dict["__value__"])
        return obj_dict


class ReadableDataStorageBase:
    def load(self) -> None:
        raise NotImplementedError

    def get_result_data(self, storage_key: StorageKey) -> DataTable:
        raise NotImplementedError


class WritableDataStorageBase:
    def save(self) -> None:
        raise NotImplementedError

    def set_result_data(self, storage_key: StorageKey, result_data: DataTable) -> None:
        raise NotImplementedError


@attr.s
class ReadableDataStorage(ReadableDataStorageBase):
    _filename: str = attr.ib(kw_only=True)
    _data: Dict[str, DataTable] = attr.ib(factory=dict)

    def _load_data_table(self, table_data: dict) -> DataTable:
        return DataTable(
            columns=[
                DataColumn(name=col_data["name"], data_type=DataType[col_data["type"]])
                for col_data in table_data["columns"]
            ],
            rows=table_data["rows"],
        )

    def load(self) -> None:
        with open(self._filename) as storage_file:
            data_as_list = json.load(storage_file, cls=CustomJSONDecoder)
        self._data = {key: self._load_data_table(table_data) for key, table_data in data_as_list}

    def get_result_data(self, storage_key: StorageKey) -> DataTable:
        return self._data[storage_key.as_str]


@attr.s
class WritableDataStorage(WritableDataStorageBase):
    _filename: str = attr.ib(kw_only=True)
    _data: Dict[str, DataTable] = attr.ib(init=False, factory=dict)

    def _dump_data_table(self, data_table: DataTable) -> dict:
        return {
            "columns": [{"name": col.name, "type": col.data_type.name} for col in data_table.columns],
            "rows": data_table.rows,
        }

    def save(self) -> None:
        data_as_list = sorted([[key, self._dump_data_table(data_table)] for key, data_table in self._data.items()])
        with open(self._filename, "w") as storage_file:
            json.dump(data_as_list, storage_file, cls=CustomJSONEncoder)

    def set_result_data(self, storage_key: StorageKey, result_data: DataTable) -> None:
        self._data[storage_key.as_str] = result_data

lib/dl_formula_ref/dl_formula_ref/examples/result_storage.py
datalens-tech/datalens-backend
import attr
import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

from dl_formula.core.dialect import DialectCombo
from dl_formula.definitions.scope import Scope
from dl_formula.translation.env import TranslationEnvironment
from dl_formula.translation.translator import SqlAlchemyTranslator
from dl_formula_ref.examples.query import (
    CompiledQueryContext,
    FormulaContext,
    RawQueryContext,
    TableReference,
)


@attr.s
class SaQueryCompiler:
    dialect: DialectCombo = attr.ib(kw_only=True)
    _function_scopes: int = attr.ib(kw_only=True, default=Scope.EXPLICIT_USAGE)

    def _get_translator(self, table_ref: TableReference) -> SqlAlchemyTranslator:
        translator = SqlAlchemyTranslator(
            env=TranslationEnvironment(
                dialect=self.dialect,
                required_scopes=self._function_scopes,
                field_names={col.name: (table_ref.name, col.name) for col in table_ref.columns},
                field_types={col.name: col.data_type for col in table_ref.columns},
                field_type_params={},
            ),
        )
        return translator

    def compile_formula(self, formula_ctx: FormulaContext, translator: SqlAlchemyTranslator) -> ClauseElement:
        ctx = translator.translate(formula=formula_ctx.formula)
        expr = ctx.expression
        assert expr is not None
        return expr

    def compile_from(self, table_name: str) -> ClauseElement:
        return sa.table(name=table_name)

    def compile_query(self, raw_query: RawQueryContext, table_ref: TableReference) -> CompiledQueryContext:
        trans = self._get_translator(table_ref=table_ref)
        sa_select = [self.compile_formula(formula_ctx, translator=trans) for formula_ctx in raw_query.select]
        sa_group_by = [self.compile_formula(formula_ctx, translator=trans) for formula_ctx in raw_query.group_by]
        sa_order_by = [self.compile_formula(formula_ctx, translator=trans) for formula_ctx in raw_query.order_by]
        sa_from = self.compile_from(table_name=table_ref.name)
        sa_query = sa.select(sa_select).select_from(sa_from).group_by(*sa_group_by).order_by(*sa_order_by)
        compiled_query_ctx = CompiledQueryContext(
            result_columns=raw_query.result_columns,
            sa_query=sa_query,
        )
        return compiled_query_ctx

lib/dl_formula_ref/dl_formula_ref/examples/sa_compiler.py
datalens-tech/datalens-backend
from dl_formula_ref.examples.config import ExampleConfig
from dl_formula_ref.examples.data_table import (
    DataColumn,
    DataTable,
)
from dl_formula_ref.examples.result_storage import StorageKey


def make_key_for_example(func_name: str, example: ExampleConfig) -> StorageKey:
    return StorageKey(item_key=func_name, example_key=example.name)


def make_data_table_from_example(example: ExampleConfig) -> DataTable:
    return DataTable(
        columns=[DataColumn(name=name, data_type=data_type) for name, data_type in example.source.columns],
        rows=example.source.data,
    )

lib/dl_formula_ref/dl_formula_ref/examples/utils.py
datalens-tech/datalens-backend
from typing import (
    Collection,
    Optional,
    Type,
)

import attr

import dl_formula_ref as package
from dl_formula_ref.plugins.base.plugin import FormulaRefPlugin
from dl_formula_ref.plugins.registration import FORMULA_REF_PLUGIN_REG
from dl_utils.entrypoints import EntrypointClassManager


_PLUGIN_EP_GROUP = f"{package.__name__}.plugins"


@attr.s
class FormulaRefPluginEntrypointManager(EntrypointClassManager[FormulaRefPlugin]):
    entrypoint_group_name = attr.ib(init=False, default=_PLUGIN_EP_GROUP)


def get_all_formula_ref_plugins() -> dict[str, Type[FormulaRefPlugin]]:
    ep_mgr = FormulaRefPluginEntrypointManager()
    return ep_mgr.get_all_ep_classes()


def _register_plugin(plugin_cls: Type[FormulaRefPlugin]) -> None:
    FORMULA_REF_PLUGIN_REG.register_plugin(plugin_cls)


def register_all_plugins(plugin_ep_names: Optional[Collection[str]] = None) -> None:
    for ep_name, plugin_cls in sorted(get_all_formula_ref_plugins().items()):
        if plugin_ep_names is not None and ep_name not in plugin_ep_names:
            continue
        _register_plugin(plugin_cls)

lib/dl_formula_ref/dl_formula_ref/formula_ref_plugins.py
datalens-tech/datalens-backend
from dl_formula.core.datatype import DataType
from dl_formula_ref.categories.aggregation import CATEGORY_AGGREGATION
from dl_formula_ref.examples.config import (
    ExampleConfig,
    ExampleSource,
)
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import (
    DataExample,
    SimpleExample,
)


_ = get_gettext()

_SOURCE_INT_FLOAT_1 = ExampleSource(
    columns=[
        ("City", DataType.STRING),
        ("Category", DataType.STRING),
        ("Orders", DataType.INTEGER),
        ("Profit", DataType.FLOAT),
    ],
    data=[
        ["London", "Office Supplies", 8, 120.10],
        ["London", "Furniture", 1, 750.0],
        ["Moscow", "Furniture", 2, 1250.50],
        ["Moscow", "Office Supplies", 4, 85.34],
        ["San Francisco", "Office Supplies", 23, 723.0],
        ["Detroit", "Furniture", 5, 6205.87],
    ],
)


FUNCTION_SUM = FunctionDocRegistryItem(
    name="sum",
    category=CATEGORY_AGGREGATION,
    description=_("Returns the sum of all expression values. Applicable to numeric data types only."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_INT_FLOAT_1,
                group_by=["[City]"],
                order_by=["[City]"],
                show_source_table=True,
                formula_fields=[("City", "[City]"), ("Sum 1", "SUM([Orders])"), ("Sum 2", "SUM([Profit])")],
            ),
        ),
    ],
)

FUNCTION_SUM_IF = FunctionDocRegistryItem(
    name="sum_if",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns the sum of all the expression values that meet the {arg:1} condition. "
        "Applicable to numeric data types only."
    ),
    examples=[
        SimpleExample("SUM_IF([Profit], [Profit] > 15)"),
    ],
)

FUNCTION_AVG = FunctionDocRegistryItem(
    name="avg",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns the average of all values. Applicable to numeric data types as well as {type:DATE|DATETIME}."
    ),
    examples=[
        SimpleExample("AVG([Profit])"),
    ],
)

FUNCTION_AVG_IF = FunctionDocRegistryItem(
    name="avg_if",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns the average of all values that meet the {arg:1} condition. "
        "If the values don't exist, it returns `NULL`. Applicable to numeric data types only."
    ),
    examples=[
        SimpleExample("AVG_IF([Profit], [Profit] > 5)"),
    ],
)

FUNCTION_MAX = FunctionDocRegistryItem(
    name="max",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns the maximum value.\n\n"
        "If {arg:0}:\n"
        "- number — Returns the largest number.\n"
        "- date — Returns the latest date.\n"
        "- string — Returns the last value in the alphabetic order.\n"
    ),
    examples=[
        SimpleExample("MAX([Profit])"),
    ],
)

FUNCTION_MIN = FunctionDocRegistryItem(
    name="min",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns the minimum value.\n\n"
        "If {arg:0}:\n"
        "- number — Returns the smallest number.\n"
        "- date — Returns the earliest date.\n"
        "- string — Returns the first value in the alphabetic order.\n"
    ),
    examples=[
        SimpleExample("MIN([Profit])"),
    ],
)

FUNCTION_COUNT = FunctionDocRegistryItem(
    name="count",
    category=CATEGORY_AGGREGATION,
    description=_("Returns the number of items in the group."),
    examples=[
        SimpleExample("COUNT()"),
        SimpleExample("COUNT([OrderID])"),
    ],
)

FUNCTION_COUNT_IF = FunctionDocRegistryItem(
    name="count_if",
    category=CATEGORY_AGGREGATION,
    description=_("Returns the number of items in the group meeting the {arg:0} condition."),
    examples=[
        SimpleExample("COUNT_IF([Profit] > 5)"),
    ],
)

FUNCTION_COUNTD = FunctionDocRegistryItem(
    name="countd",
    category=CATEGORY_AGGREGATION,
    description=_("Returns the number of unique values in the group.\n" "\n" "See also {ref:COUNTD_APPROX}."),
    examples=[
        SimpleExample("COUNTD([ClientID])"),
    ],
)

FUNCTION_COUNTD_IF = FunctionDocRegistryItem(
    name="countd_if",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns the number of unique values in the group that meet the {arg:1} "
        "condition.\n"
        "\n"
        "See also {ref:COUNTD_APPROX}."
    ),
    examples=[
        SimpleExample("COUNTD_IF([ClientID], [Profit] > 5)"),
    ],
)

FUNCTION_COUNTD_APPROX = FunctionDocRegistryItem(
    name="countd_approx",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns the approximate number of unique values in the group. Faster than "
        "{ref:COUNTD}, but doesn't guarantee accuracy."
    ),
    examples=[
        SimpleExample("COUNTD_APPROX([ClienID])"),
    ],
)

FUNCTION_STDEV = FunctionDocRegistryItem(
    name="stdev",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns the statistical standard deviation of all values in the expression "
        "based on a selection from the population."
    ),
    examples=[
        SimpleExample("STDEV([Profit])"),
    ],
)

FUNCTION_STDEVP = FunctionDocRegistryItem(
    name="stdevp",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns the statistical standard deviation of all values in the expression " "based on the biased population."
    ),
    examples=[
        SimpleExample("STDEVP([Profit])"),
    ],
)

FUNCTION_VAR = FunctionDocRegistryItem(
    name="var",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns the statistical variance of all values in an expression based on a " "selection from the population."
    ),
    examples=[
        SimpleExample("VAR([Profit])"),
    ],
)

FUNCTION_VARP = FunctionDocRegistryItem(
    name="varp",
    category=CATEGORY_AGGREGATION,
    description=_("Returns the statistical variance of all values in an expression across the " "entire population."),
    examples=[
        SimpleExample("VARP([Profit])"),
    ],
)

FUNCTION_QUANTILE = FunctionDocRegistryItem(
    name="quantile",
    category=CATEGORY_AGGREGATION,
    description=_("Returns the precise {arg:1}-level quantile ({arg:1} should be in range " "from 0 to 1)."),
)

FUNCTION_QUANTILE_APPROX = FunctionDocRegistryItem(
    name="quantile_approx",
    category=CATEGORY_AGGREGATION,
    description=_("Returns the approximate {arg:1}-level quantile ({arg:1} should be in range from " "0 to 1)."),
)

FUNCTION_MEDIAN = FunctionDocRegistryItem(
    name="median",
    category=CATEGORY_AGGREGATION,
    description=_("Returns the median value."),
    examples=[
        SimpleExample("MEDIAN([Profit])"),
    ],
)

FUNCTION_ANY = FunctionDocRegistryItem(
    name="any",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns one of the values of {arg:0} from the group. This is a "
        "nondeterministic aggregation — the result may vary for the same data over "
        "multiple queries."
    ),
    examples=[
        SimpleExample("ANY([Profit])"),
    ],
)

FUNCTION_ARG_MIN = FunctionDocRegistryItem(
    name="arg_min",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns {arg:0} for the minimum value of {arg:1} in the group. If multiple "
        "values of {arg:0} match the minimum value of {arg:1}, then the first one "
        "encountered is returned. This makes the function non-deterministic."
    ),
    examples=[
        SimpleExample("ARG_MIN([Sales], [Profit])"),
    ],
)

FUNCTION_ARG_MAX = FunctionDocRegistryItem(
    name="arg_max",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns {arg:0} for the maximum value of {arg:1} in the group. If multiple "
        "values of {arg:0} match the maximum value of {arg:1}, then the first one "
        "encountered is returned. This makes the function non-deterministic."
    ),
    examples=[
        SimpleExample("ARG_MAX([Sales], [Profit])"),
    ],
)

FUNCTION_ALL_CONCAT = FunctionDocRegistryItem(
    name="all_concat",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns a string that contains all unique values of {arg:0} delimited by {arg:1} "
        "(if {arg:1} is not specified, a comma is used)."
    ),
    examples=[
        SimpleExample("ALL_CONCAT([Profit])"),
        SimpleExample("ALL_CONCAT([Profit], '; ')"),
    ],
)

FUNCTION_TOP_CONCAT = FunctionDocRegistryItem(
    name="top_concat",
    category=CATEGORY_AGGREGATION,
    description=_(
        "Returns a string that contains top {arg:1} unique values of {arg:0} delimited by {arg:2} "
        "(if {arg:2} is not specified, a comma is used)."
    ),
    examples=[
        SimpleExample("TOP_CONCAT([Profit], 3)"),
        SimpleExample("TOP_CONCAT([Profit], 3, '; ')"),
    ],
)

FUNCTIONS_AGGREGATION = [
    FUNCTION_SUM,
    FUNCTION_SUM_IF,
    FUNCTION_AVG,
    FUNCTION_AVG_IF,
    FUNCTION_MAX,
    FUNCTION_MIN,
    FUNCTION_COUNT,
    FUNCTION_COUNT_IF,
    FUNCTION_COUNTD,
    FUNCTION_COUNTD_IF,
    FUNCTION_COUNTD_APPROX,
    FUNCTION_STDEV,
    FUNCTION_STDEVP,
    FUNCTION_VAR,
    FUNCTION_VARP,
    FUNCTION_QUANTILE,
    FUNCTION_QUANTILE_APPROX,
    FUNCTION_MEDIAN,
    FUNCTION_ANY,
    FUNCTION_ARG_MIN,
    FUNCTION_ARG_MAX,
    FUNCTION_ALL_CONCAT,
    FUNCTION_TOP_CONCAT,
]

lib/dl_formula_ref/dl_formula_ref/functions/aggregation.py
datalens-tech/datalens-backend
from dl_formula.core.datatype import DataType
from dl_formula_ref.categories.array import CATEGORY_ARRAY
from dl_formula_ref.examples.config import (
    ExampleConfig,
    ExampleSource,
)
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import (
    DataExample,
    SimpleExample,
)
from dl_formula_ref.registry.impl_selector import ArgAwareImplementationSelector
from dl_formula_ref.registry.naming import CategoryPostfixNamingProvider
from dl_formula_ref.registry.note import Note


_ = get_gettext()

_EXAMPLE_SOURCE_UNNEST = ExampleSource(
    columns=[
        ("City", DataType.STRING),
        ("Category", DataType.ARRAY_STR),
    ],
    data=[
        ["Moscow", ["Office Supplies", "Furniture"]],
        ["London", ["Office Supplies"]],
    ],
)

_EXAMPLE_SOURCE_ARRAY_TYPE_CONVERSION = ExampleSource(
    columns=[
        ("string_array", DataType.ARRAY_STR),
        ("float_array", DataType.ARRAY_FLOAT),
        ("int_array", DataType.ARRAY_INT),
    ],
    data=[
        [["100", "300", "200"], [12.3, 0.43, 42], [12, 0, 42]],
        [["100", "300", None, None], [0, None, 12.0], [-3, None, 0]],
        [["150", None, "192"], [3.2, 2.3, 3.5], [132, 637, None]],
    ],
)

_EXAMPLE_SOURCE_REPLACE = ExampleSource(
    columns=[
        ("array", DataType.ARRAY_INT),
    ],
    data=[
        [[100, 300, 200, 100, 300]],
        [[100, 300, None, None]],
        [[150, None, 130, 192]],
    ],
)

_EXAMPLE_SOURCE_CONTAINS = ExampleSource(
    columns=[
        ("array", DataType.ARRAY_INT),
        ("value", DataType.INTEGER),
    ],
    data=[
        [[1, 2, 3], 2],
        [[1, 2, 3], 4],
    ],
)

_EXAMPLE_SOURCE_CONTAINS_ALL = ExampleSource(
    columns=[
        ("array1", DataType.ARRAY_INT),
        ("array2", DataType.ARRAY_INT),
    ],
    data=[
        [[1, 2, 3], [1, 2]],
        [[1, 2, 3], [1, 4]],
    ],
)

_EXAMPLE_SOURCE_CONTAINS_ANY = ExampleSource(
    columns=[
        ("array1", DataType.ARRAY_INT),
        ("array2", DataType.ARRAY_INT),
    ],
    data=[
        [[1, 2, 3], [1, 5]],
        [[1, 2, 3], [4, 5]],
    ],
)

_EXAMPLE_SOURCE_CONTAINS_SUBSEQUENCE = ExampleSource(
    columns=[
        ("array1", DataType.ARRAY_INT),
        ("array2", DataType.ARRAY_INT),
    ],
    data=[
        [[1, 2, 3], [1, 2]],
        [[1, 2, 3], [2, 1]],
        [[1, 2, 3], [1, 3]],
    ],
)

_EXAMPLE_SOURCE_ARRAY_AGGREGATION_FUNCTIONS = ExampleSource(
    columns=[
        ("float_array", DataType.ARRAY_FLOAT),
        ("int_array", DataType.ARRAY_INT),
    ],
    data=[
        [[14.3, 0.42, 15], [21, 12, 0]],
        [[0, -3, 12.0], [-4, 12, 0]],
        [[3.2, 2.3, 3.2], [5, 7, 9]],
    ],
)

FUNCTION_ARRAY = FunctionDocRegistryItem(
    name="array",
    category=CATEGORY_ARRAY,
    description=_("Returns an array containing the passed values."),
    notes=[
        Note(_("All passed values must be of the same type or `NULL`. At least one value must be non-`NULL`.")),
    ],
    examples=[
        SimpleExample("ARRAY(1, 2, NULL, 3)"),
        SimpleExample("ARRAY('a', 'b', NULL, 'c')"),
        SimpleExample("ARRAY(0, 2.3, NULL, 0.18)"),
    ],
)

FUNCTION_UNNEST = FunctionDocRegistryItem(
    name="unnest",
    category=CATEGORY_ARRAY,
    description=_("Expands the {arg:0} array expression to a set of rows."),
    notes=[
        Note(
            _(
                "{dialects: POSTGRESQL} doesn't allow filtering fields containing the UNNEST "
                "function. If the data source is {dialects: POSTGRESQL}, do not use such "
                "fields in selectors."
            )
        ),
    ],
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_UNNEST,
                show_source_table=True,
                formula_fields=[("City", "[City]"), ("unnest_category", "UNNEST([Category])")],
            ),
        ),
    ],
)

FUNCTION_GET_ITEM = FunctionDocRegistryItem(
    name="get_item",
    category=CATEGORY_ARRAY,
    description=_(
        "Returns the element with the index {arg:1} from the array {arg:0}. "
        "Index must be any integer. Indexes in an array begin with one."
    ),
    examples=[
        SimpleExample("GET_ITEM([array_field], 2)"),
    ],
)

FUNCTION_ARR_STR = FunctionDocRegistryItem(
    name="arr_str",
    category=CATEGORY_ARRAY,
    description=_(
        "Concatenates elements of the array {arg:0} using {arg:1} as a delimiter (comma by default) "
        "and {arg:2} as a `NULL` string (`NULL` items are skipped by default).\n"
        "\n"
        "See also {ref:STR}"
    ),
    examples=[
        SimpleExample("ARR_STR(ARRAY(1, 2, NULL, 4)) = 1,2,4"),
        SimpleExample("ARR_STR(ARRAY(1, 2, NULL, 4), ';') = 1;2;4"),
        SimpleExample("ARR_STR(ARRAY(1, 2, NULL, 4), ';', '*') = 1;2;*;4"),
        SimpleExample("ARR_STR(ARRAY('a', 'b', '', NULL), ';', '*') = a;b;;*"),
    ],
)

FUNCTION_COUNT_ITEM = FunctionDocRegistryItem(
    name="count_item",
    category=CATEGORY_ARRAY,
    description=_(
        "Returns the number of elements in the array {arg:0} equal to {arg:1}. "
        "The type of {arg:1} must match the type of the {arg:0} elements."
    ),
    examples=[
        SimpleExample("COUNT_ITEM(ARRAY(1, 2, 2, 3), 2) = 2"),
        SimpleExample("COUNT_ITEM(ARRAY(1, 2, 2, 3), 4) = 0"),
        SimpleExample("COUNT_ITEM(ARRAY(1, NULL, 2, NULL), NULL) = 2"),
    ],
)

FUNCTION_CONTAINS = FunctionDocRegistryItem(
    name="contains",
    category=CATEGORY_ARRAY,
    impl_selector=ArgAwareImplementationSelector(
        exp_arg_types={
            0: {DataType.ARRAY_INT, DataType.ARRAY_FLOAT, DataType.ARRAY_STR},
        },
    ),
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="contains_array",
    ),
    description=_("Returns `TRUE` if {arg:0} contains {arg:1}."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_CONTAINS,
                formula_fields=[
                    ("array", "[array]"),
                    ("value", "[value]"),
                    ("contains", "CONTAINS([array], [value])"),
                ],
            ),
        ),
    ],
)

FUNCTION_CONTAINS_ALL = FunctionDocRegistryItem(
    name="contains_all",
    category=CATEGORY_ARRAY,
    impl_selector=ArgAwareImplementationSelector(
        exp_arg_types={
            0: {DataType.ARRAY_INT, DataType.ARRAY_FLOAT, DataType.ARRAY_STR},
        },
    ),
    description=_("Returns `TRUE` if {arg:0} contains all elements of {arg:1}."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_CONTAINS_ALL,
                formula_fields=[
                    ("array1", "[array1]"),
                    ("array2", "[array2]"),
                    ("contains_all", "CONTAINS_ALL([array1], [array2])"),
                ],
            ),
        ),
    ],
)

FUNCTION_CONTAINS_ANY = FunctionDocRegistryItem(
    name="contains_any",
    category=CATEGORY_ARRAY,
    impl_selector=ArgAwareImplementationSelector(
        exp_arg_types={
            0: {DataType.ARRAY_INT, DataType.ARRAY_FLOAT, DataType.ARRAY_STR},
        },
    ),
    description=_("Returns `TRUE` if {arg:0} contains any elements of {arg:1}."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_CONTAINS_ANY,
                formula_fields=[
                    ("array1", "[array1]"),
                    ("array2", "[array2]"),
                    ("contains_any", "CONTAINS_ANY([array1], [array2])"),
                ],
            ),
        ),
    ],
)

FUNCTION_CONTAINS_SUBSEQUENCE = FunctionDocRegistryItem(
    name="contains_subsequence",
    category=CATEGORY_ARRAY,
    impl_selector=ArgAwareImplementationSelector(
        exp_arg_types={
            0: {DataType.ARRAY_INT, DataType.ARRAY_FLOAT, DataType.ARRAY_STR},
        },
    ),
    description=_(
        "Returns `TRUE` if {arg:1} is a continuous subsequence of {arg:0}. "
        "In other words, returns `TRUE` if and only if `array1 = prefix + array2 + suffix`."
    ),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_CONTAINS_SUBSEQUENCE,
                formula_fields=[
                    ("array1", "[array1]"),
                    ("array2", "[array2]"),
                    ("contains_subsequence", "CONTAINS_SUBSEQUENCE([array1], [array2])"),
                ],
            ),
        ),
    ],
)

FUNCTION_STARTSWITH = FunctionDocRegistryItem(
    name="startswith",
    category=CATEGORY_ARRAY,
    impl_selector=ArgAwareImplementationSelector(
        exp_arg_types={
            0: {DataType.ARRAY_INT, DataType.ARRAY_FLOAT, DataType.ARRAY_STR},
        },
    ),
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="startswith_array",
    ),
    description=_("Returns `TRUE` if {arg:0} starts with {arg:1}."),
    examples=[
        SimpleExample("STARTSWITH(ARRAY(1, 2, 3), ARRAY(1, 2)) = TRUE"),
        SimpleExample("STARTSWITH(ARRAY(1, 2, 3), ARRAY(2, 3)) = FALSE"),
    ],
)

FUNCTION_SLICE = FunctionDocRegistryItem(
    name="slice",
    category=CATEGORY_ARRAY,
    description=_(
        "Returns the part of array {arg:0} of length {arg:2} starting from index {arg:1}. Indexes in an array begin with one."
    ),
    examples=[
        SimpleExample("SLICE(ARRAY(1, 2, 3, 4, 5), 3, 2) = [3, 4]"),
        SimpleExample("SLICE(ARRAY(1, 2, 3, 4, 5), 3, 1) = [3]"),
    ],
)

FUNCTION_ARRAY_INT = FunctionDocRegistryItem(
    name="cast_arr_int",
    category=CATEGORY_ARRAY,
    description=_("Converts {arg:0} to an array of integers. The conversion rules are the same as for {ref:INT}."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_ARRAY_TYPE_CONVERSION,
                formula_fields=[
                    ("string_array", "[string_array]"),
                    ("float_array", "[float_array]"),
                    ("Array of integers from string", "CAST_ARR_INT([string_array])"),
                    ("Array of integers from float", "CAST_ARR_INT([float_array])"),
                ],
            ),
        ),
    ],
)

FUNCTION_ARRAY_STR = FunctionDocRegistryItem(
    name="cast_arr_str",
    category=CATEGORY_ARRAY,
    description=_("Converts {arg:0} to an array of strings. The conversion rules are the same as for {ref:STR}."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_ARRAY_TYPE_CONVERSION,
                formula_fields=[
                    ("int_array", "[int_array]"),
                    ("float_array", "[float_array]"),
                    ("Array of string from integer", "CAST_ARR_STR([int_array])"),
                    ("Array of integers from float", "CAST_ARR_STR([float_array])"),
                ],
            ),
        ),
    ],
)

FUNCTION_ARRAY_FLOAT = FunctionDocRegistryItem(
    name="cast_arr_float",
    category=CATEGORY_ARRAY,
    description=_(
        "Converts {arg:0} to an array of fractional numbers. The conversion rules are the same as for {ref:FLOAT}."
    ),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_ARRAY_TYPE_CONVERSION,
                formula_fields=[
                    ("string_array", "[string_array]"),
                    ("int_array", "[int_array]"),
                    ("Array of float from string", "CAST_ARR_FLOAT([string_array])"),
                    ("Array of floats from integer", "CAST_ARR_FLOAT([int_array])"),
                ],
            ),
        ),
    ],
)

FUNCTION_REPLACE = FunctionDocRegistryItem(
    name="replace",
    category=CATEGORY_ARRAY,
    impl_selector=ArgAwareImplementationSelector(
        exp_arg_types={
            0: {DataType.ARRAY_INT, DataType.ARRAY_FLOAT, DataType.ARRAY_STR},
        },
    ),
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="replace_array",
    ),
    description=_("Replaces each {arg:0} element equal to {arg:1} with {arg:2}."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_REPLACE,
                formula_fields=[
                    ("array", "[array]"),
                    ("replace", "REPLACE([array], 300, -1)"),
                    ("No NULLs", "REPLACE([array], NULL, 0)"),
                ],
            ),
        ),
    ],
)

ARRAY_AGGREGATION_FUNCTIONS_NOTE = Note(
    _(
        "This function cannot work with arrays with `Nullable` items. To remove `NULL` items from the array,"
        " use {ref:ARR_REMOVE} or {ref:array/REPLACE}."
    )
)


def _make_array_aggregation_function_example(func: str) -> list[DataExample]:
    func = func.upper()
    examples = [
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_ARRAY_AGGREGATION_FUNCTIONS,
                formula_fields=[
                    ("[int_array]", "[int_array]"),
                    ("[float_array]", "[float_array]"),
                    (f"{func}([int_array])", f"{func}(REPLACE([int_array], NULL, 0))"),
                    (f"{func}([float_array])", f"{func}(REPLACE([float_array], NULL, 0))"),
                ],
                override_formula_fields=[
                    ("[int_array]", "[int_array]"),
                    ("[float_array]", "[float_array]"),
                    (f"{func}([int_array])", f"{func}([int_array])"),
                    (f"{func}([float_array])", f"{func}([float_array])"),
                ],
            ),
        ),
    ]
    return examples


FUNCTION_ARR_MIN = FunctionDocRegistryItem(
    name="arr_min",
    category=CATEGORY_ARRAY,
    description=_("Returns the least of elements in the {arg:0}."),
    notes=[ARRAY_AGGREGATION_FUNCTIONS_NOTE],
    examples=[
        *_make_array_aggregation_function_example("arr_min"),
    ],
)

FUNCTION_ARR_MAX = FunctionDocRegistryItem(
    name="arr_max",
    category=CATEGORY_ARRAY,
    description=_("Returns the greatest of elements in the {arg:0}."),
    notes=[ARRAY_AGGREGATION_FUNCTIONS_NOTE],
    examples=[
        *_make_array_aggregation_function_example("arr_max"),
    ],
)

FUNCTION_ARR_SUM = FunctionDocRegistryItem(
    name="arr_sum",
    category=CATEGORY_ARRAY,
    description=_("Returns the sum of elements in the {arg:0}."),
    notes=[ARRAY_AGGREGATION_FUNCTIONS_NOTE],
    examples=[
        *_make_array_aggregation_function_example("arr_sum"),
    ],
)

FUNCTION_ARR_AVG = FunctionDocRegistryItem(
    name="arr_avg",
    category=CATEGORY_ARRAY,
    description=_("Returns the average of elements in the {arg:0}."),
    notes=[ARRAY_AGGREGATION_FUNCTIONS_NOTE],
    examples=[
        *_make_array_aggregation_function_example("arr_avg"),
    ],
)

FUNCTION_ARR_PRODUCT = FunctionDocRegistryItem(
    name="arr_product",
    category=CATEGORY_ARRAY,
    description=_("Returns the product of elements in the {arg:0}."),
    notes=[ARRAY_AGGREGATION_FUNCTIONS_NOTE],
    examples=[
        *_make_array_aggregation_function_example("arr_product"),
    ],
)

FUNCTION_ARR_REMOVE = FunctionDocRegistryItem(
    name="arr_remove",
    category=CATEGORY_ARRAY,
    description=_("Removes all {arg:0} elements equal to {arg:1}."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_EXAMPLE_SOURCE_REPLACE,
                formula_fields=[
                    ("array", "[array]"),
                    ("No NULLs", "ARR_REMOVE([array], NULL)"),
                ],
            ),
        ),
    ],
)

FUNCTIONS_ARRAY = [
    FUNCTION_ARRAY,
    FUNCTION_UNNEST,
    FUNCTION_GET_ITEM,
    FUNCTION_ARR_STR,
    FUNCTION_COUNT_ITEM,
    FUNCTION_CONTAINS,
    FUNCTION_CONTAINS_ALL,
    FUNCTION_CONTAINS_ANY,
    FUNCTION_CONTAINS_SUBSEQUENCE,
    FUNCTION_STARTSWITH,
    FUNCTION_SLICE,
    FUNCTION_ARRAY_INT,
    FUNCTION_ARRAY_STR,
    FUNCTION_ARRAY_FLOAT,
    FUNCTION_REPLACE,
    FUNCTION_ARR_MIN,
    FUNCTION_ARR_MAX,
    FUNCTION_ARR_SUM,
    FUNCTION_ARR_AVG,
    FUNCTION_ARR_PRODUCT,
    FUNCTION_ARR_REMOVE,
]

lib/dl_formula_ref/dl_formula_ref/functions/array.py
datalens-tech/datalens-backend
import datetime

from dl_formula.core.datatype import DataType
from dl_formula_ref.categories.date import CATEGORY_DATE
from dl_formula_ref.examples.config import (
    ExampleConfig,
    ExampleSource,
)
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import (
    DataExample,
    SimpleExample,
)
from dl_formula_ref.registry.note import Note


_ = get_gettext()

_d = datetime.date
_dt = datetime.datetime

_SOURCE_DATE_1 = ExampleSource(
    columns=[("Date", DataType.DATE)], data=[[_d(2014, 10, 6)], [_d(2014, 10, 7)], [_d(2017, 3, 8)], [_d(2024, 2, 12)]]
)
_SOURCE_DATETIME_1 = ExampleSource(
    columns=[("DateTime", DataType.GENERICDATETIME)],
    data=[
        [_dt(2014, 10, 6, 7, 45, 12)],
        [_dt(2014, 10, 7, 11, 10, 15)],
        [_dt(2017, 3, 8, 23, 59, 59)],
        [_dt(2024, 2, 12, 7, 40, 33)],
    ],
)


FUNCTION_DATEADD = FunctionDocRegistryItem(
    name="dateadd",
    category=CATEGORY_DATE,
    description=_(
        "Returns the date obtained by adding {arg:1} in the amount of {arg:2} to the "
        "specified date {arg:0}.\n"
        "\n"
        "The {arg:2} argument is an integer. It can be negative.\n"
        "The {arg:1} argument takes the following values:\n"
        '- `"year"`;\n'
        '- `"month"`;\n'
        '- `"day"`;\n'
        '- `"hour"`;\n'
        '- `"minute"`;\n'
        '- `"second"`.'
    ),
    notes=[
        Note(
            _("For all sources except {dialects:CLICKHOUSE}, {arg:2} takes " "only constant values."),
        ),
    ],
    examples=[
        SimpleExample(example_str)
        for example_str in (
            'DATEADD(#2018-01-12#, "day", 6) = #2018-01-18#',
            'DATEADD(#2018-01-12#, "month", 6) = #2018-07-12#',
            'DATEADD(#2018-01-12#, "year", 6) = #2024-01-12#',
            'DATEADD(#2018-01-12 01:02:03#, "second", 6) = #2018-01-12 01:02:09#',
            'DATEADD(#2018-01-12 01:02:03#, "minute", 6) = #2018-01-12 01:08:03#',
            'DATEADD(#2018-01-12 01:02:03#, "hour", 6) = #2018-01-12 07:02:03#',
            'DATEADD(#2018-01-12 01:02:03#, "day", 6) = #2018-01-18 01:02:03#',
            'DATEADD(#2018-01-12 01:02:03#, "month", 6) = #2018-07-12 01:02:03#',
            'DATEADD(#2018-01-12 01:02:03#, "year", 6) = #2024-01-12 01:02:03#',
        )
    ],
)

FUNCTION_DATEPART = FunctionDocRegistryItem(
    name="datepart",
    category=CATEGORY_DATE,
    description=_(
        "Returns a part of the date as an integer.\n"
        "\n"
        "Depending on the argument, {arg:1} returns the following values:\n"
        '- `"year"` — the year number (see {ref:YEAR});\n'
        '- `"quarter"` — the number of the quarter (from `1` to `4`) of the year '
        "(see {ref:QUARTER});\n"
        '- `"month"` — the number of the month in the year (see {ref:MONTH});\n'
        '- `"week"` — the number of the week in the year according to {link: '
        "iso_8601: ISO 8601} (see {ref:WEEK});\n"
        '- `"dayofweek"`, `"dow"` — the number of the day of the week (see {ref:'
        "DAYOFWEEK});\n"
        '- `"day"` — the number of the day in the month (see {ref:DAY});\n'
        '- `"hour"` — the number of the hour in the day (see {ref:HOUR});\n'
        '- `"minute"` — the number of the minute in the hour (see {ref:MINUTE});\n'
        '- `"second"` — the number of the second in the minute (see {ref:SECOND}).\n'
        "\n"
        'If you select `"dayofweek"`, you can use the additional parameter {arg:2} '
        "to specify the first day of the week (Monday by default). Learn more about "
        "this parameter in the {ref:DAYOFWEEK} function description.\n"
    ),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with date"),
                source=_SOURCE_DATE_1,
                formula_fields=[
                    ("Date", "[Date]"),
                    ("Year", 'DATEPART([Date], "year")'),
                    ("Month", 'DATEPART([Date], "month")'),
                    ("Day", 'DATEPART([Date], "day")'),
                    ("DayOfWeek", 'DATEPART([Date], "dayofweek")'),
                    ("DOW", 'DATEPART([Date], "dow")'),
                ],
                formulas_as_names=False,
            )
        ),
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with custom first day of the week"),
                source=_SOURCE_DATE_1,
                formula_fields=[
                    ("Date", "[Date]"),
                    ("DOW", 'DATEPART([Date], "dow")'),
                    ("DOW sun", 'DATEPART([Date], "dow", "sun")'),
                    ("DOW Monday", 'DATEPART([Date], "dow", "Monday")'),
                    ("DOW wed", 'DATEPART([Date], "dow", "wed")'),
                ],
                formulas_as_names=False,
            )
        ),
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with date and time"),
                source=_SOURCE_DATETIME_1,
                formula_fields=[
                    ("DateTime", "[DateTime]"),
                    ("Year", 'DATEPART([DateTime], "year")'),
                    ("Month", 'DATEPART([DateTime], "month")'),
                    ("Day", 'DATEPART([DateTime], "day")'),
                    ("Hour", 'DATEPART([DateTime], "hour")'),
                    ("Minute", 'DATEPART([DateTime], "minute")'),
                    ("Second", 'DATEPART([DateTime], "second")'),
                ],
                formulas_as_names=False,
            ),
        ),
    ],
)

FUNCTION_DATETRUNC = FunctionDocRegistryItem(
    name="datetrunc",
    category=CATEGORY_DATE,
    description=_(
        "Rounds {arg:0} down to the given {arg:1}. If optional {arg:2} is given, then "
        "the value is rounded down to a {arg:2} multiple of {arg:1} (omitting {arg:2} "
        "is the same as `{argn:2} = 1`).\n"
        "\n"
        "Supported units:\n"
        '- `"second"`;\n'
        '- `"minute"`;\n'
        '- `"hour"`;\n'
        '- `"day"` (acts as the day of the year if {arg:2} is specified);\n'
        '- `"week"`;\n'
        '- `"month"`;\n'
        '- `"quarter"`;\n'
        '- `"year"`.'
    ),
    notes=[
        Note(
            _(
                "The function with three arguments is only available for the sources "
                "{dialects:CLICKHOUSE_21_8} or higher."
            ),
        ),
    ],
    examples=[
        SimpleExample(example_str)
        for example_str in (
            'DATETRUNC(#2018-07-12 11:07:13#, "minute") = #2018-07-12 11:07:00#',
            'DATETRUNC(#2018-07-12#, "year", 5) = #2015-01-01#',
            'DATETRUNC(#2018-07-12 11:07:13#, "second", 5) = #2018-07-12 11:07:10#',
            'DATETRUNC(#2018-07-12 11:07:13#, "month", 4) = #2018-05-01 00:00:00#',
        )
    ],
)

FUNCTION_SECOND = FunctionDocRegistryItem(
    name="second",
    category=CATEGORY_DATE,
    description=_(
        "Returns the number of the second in the minute of the specified date "
        "{arg:0}. When the date is specified without time, it returns `0`."
    ),
    examples=[
        SimpleExample("SECOND(#2019-01-23 15:07:47#) = 47"),
    ],
)

FUNCTION_MINUTE = FunctionDocRegistryItem(
    name="minute",
    category=CATEGORY_DATE,
    description=_(
        "Returns the number of the minute in the hour of the specified date {arg:0}. "
        "When the date is specified without time, it returns `0`."
    ),
    examples=[
        SimpleExample("MINUTE(#2019-01-23 15:07:47#) = 7"),
    ],
)

FUNCTION_HOUR = FunctionDocRegistryItem(
    name="hour",
    category=CATEGORY_DATE,
    description=_(
        "Returns the number of the hour in the day of the specified date and time "
        "{arg:0}. When the date is specified without time, it returns `0`."
    ),
    examples=[
        SimpleExample("HOUR(#2019-01-23 15:07:47#) = 15"),
    ],
)

FUNCTION_DAY = FunctionDocRegistryItem(
    name="day",
    category=CATEGORY_DATE,
    description=_("Returns the number of the day in the month of the specified date {arg:0}."),
    examples=[
        SimpleExample("DAY(#2019-01-23#) = 23"),
    ],
)

FUNCTION_MONTH = FunctionDocRegistryItem(
    name="month",
    category=CATEGORY_DATE,
    description=_("Returns the number of the month in the year of the specified date {arg:0}."),
    examples=[
        SimpleExample("MONTH(#2019-01-23#) = 1"),
    ],
)

FUNCTION_QUARTER = FunctionDocRegistryItem(
    name="quarter",
    category=CATEGORY_DATE,
    description=_("Returns the number of the quarter (from `1` to `4`) of the year of the specified date {arg:0}."),
    examples=[
        SimpleExample("QUARTER(#2019-01-23#) = 1"),
    ],
)

FUNCTION_YEAR = FunctionDocRegistryItem(
    name="year",
    category=CATEGORY_DATE,
    description=_("Returns the year number in the specified date {arg:0}."),
    examples=[
        SimpleExample("YEAR(#2019-01-23#) = 2019"),
    ],
)

FUNCTION_DAYOFWEEK = FunctionDocRegistryItem(
    name="dayofweek",
    category=CATEGORY_DATE,
    description=_(
        "Returns the day of the week according to {link: iso_8601: ISO 8601}.\n"
        "- Monday — 1.\n"
        "- Sunday — 7.\n"
        "\n"
        "If the additional parameter {arg:1} is specified, this day is considered the "
        "first day of the week. Valid values:\n"
        '- `"Monday"`, `"Mon"` — Monday;\n'
        '- `"Tuesday"`. `"Tue"` — Tuesday;\n'
        '- `"Wednesday"`, `"Wed"` — Wednesday;\n'
        '- `"Thursday"`, `"Thu"` — Thursday;\n'
        '- `"Friday"`, `"Fri"` — Friday;\n'
        '- `"Saturday"`, ` "Sat"` — Saturday;\n'
        '- `"Sunday"`, `"Sun"` — Sunday.\n'
    ),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            "DAYOFWEEK(#1971-01-14 01:02:03#) = 4",
            'DAYOFWEEK(#1971-01-14#, "wed") = 2',
            'DAYOFWEEK(#1971-01-14#, "wednesday") = 2',
        )
    ],
)

FUNCTION_WEEK = FunctionDocRegistryItem(
    name="week",
    category=CATEGORY_DATE,
    description=_(
        "The number of the week according to {link: iso_8601: ISO 8601}. The first "
        "week is the week that contains the first Thursday of the year or January 4th."
    ),
    examples=[
        SimpleExample("WEEK(#1971-01-14 01:02:03#) = 2"),
    ],
)

FUNCTION_NOW = FunctionDocRegistryItem(
    name="now",
    category=CATEGORY_DATE,
    description=_("Returns the current date and time, depending on the data source and " "connection type."),
    notes=[
        Note(
            _("On {dialects:YQL}, the function always returns the UTC date and time."),
        ),
    ],
    examples=[
        SimpleExample("NOW() = #2019-01-23 12:53:07#"),
    ],
)

FUNCTION_TODAY = FunctionDocRegistryItem(
    name="today",
    category=CATEGORY_DATE,
    description=_("Returns the current date, depending on the data source and connection type."),
    examples=[
        SimpleExample("TODAY() = #2019-01-23#"),
    ],
)

FUNCTIONS_DATE = [
    FUNCTION_DATEADD,
    FUNCTION_DATEPART,
    FUNCTION_DATEPART,
    FUNCTION_DATETRUNC,
    FUNCTION_SECOND,
    FUNCTION_MINUTE,
    FUNCTION_HOUR,
    FUNCTION_DAY,
    FUNCTION_MONTH,
    FUNCTION_QUARTER,
    FUNCTION_YEAR,
    FUNCTION_DAYOFWEEK,
    FUNCTION_WEEK,
    FUNCTION_NOW,
    FUNCTION_TODAY,
]

lib/dl_formula_ref/dl_formula_ref/functions/date.py
datalens-tech/datalens-backend
from dl_formula.core.datatype import DataType
from dl_formula_ref.categories.logical import CATEGORY_LOGICAL
from dl_formula_ref.examples.config import (
    ExampleConfig,
    ExampleSource,
)
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import (
    DataExample,
    SimpleExample,
)
from dl_formula_ref.registry.naming import CustomNamingProvider
from dl_formula_ref.registry.signature_base import SignaturePlacement
from dl_formula_ref.registry.signature_gen import (
    SignatureTemplate,
    TemplatedSignatureGenerator,
)


_ = get_gettext()

_SOURCE_SALES_1 = ExampleSource(
    columns=[("sales", DataType.FLOAT)],
    data=[[432.40], [77.0], [12_000.0], [None], [34.25], [128.0], [0.0], [None]],
)
_SOURCE_TIME_1 = ExampleSource(
    columns=[("unit", DataType.STRING)],
    data=[["s"], ["m"], ["h"]],
)


FUNCTION_ISNULL = FunctionDocRegistryItem(
    name="isnull",
    category=CATEGORY_LOGICAL,
    description=_(
        "Returns `TRUE` if {arg:0} is `NULL`, otherwise returns `FALSE`.\n"
        "\n"
        "`{argn:0} IS NOT NULL` returns the opposite result."
    ),
    signature_gen=TemplatedSignatureGenerator(
        signature_templates=[
            SignatureTemplate(title=_("As a function"), body="ISNULL( {1} )"),
            SignatureTemplate(title=_("As an operator"), body="{1} IS [ NOT ] NULL"),
        ]
    ),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_SALES_1,
                formula_fields=[
                    ("sales", "[sales]"),
                    ("result_1", "ISNULL([sales])"),
                    ("result_2", "[sales] IS NULL"),
                    ("result_3", "[sales] IS NOT NULL"),
                ],
            ),
        ),
    ],
)

FUNCTION_IFNULL = FunctionDocRegistryItem(
    name="ifnull",
    category=CATEGORY_LOGICAL,
    description=_("Returns {arg:0} if it's not `NULL`. Otherwise returns {arg:1}."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_SALES_1,
                formula_fields=[
                    ("sales", "[sales]"),
                    ("result", "IFNULL([sales], -100.0)"),
                ],
            ),
        ),
    ],
)

FUNCTION_ISNAN = FunctionDocRegistryItem(
    # FIXME: Disabled
    name="isnan",
    category=CATEGORY_LOGICAL,
    description=_("Returns `TRUE` if {arg:0} is `NaN`, otherwise returns `FALSE`.\n"),
)

FUNCTION_IFNAN = FunctionDocRegistryItem(
    # FIXME: Disabled
    name="ifnan",
    category=CATEGORY_LOGICAL,
    description=_("Returns {arg:0} if it's not `NaN`. Otherwise returns {arg:1}."),
    examples=[
        SimpleExample("IFNAN([Profit], 0)"),
    ],
)

FUNCTION_ZN = FunctionDocRegistryItem(
    name="zn",
    category=CATEGORY_LOGICAL,
    description=_("Returns {arg:0} if it's not `NULL`. Otherwise returns 0."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_SALES_1,
                formula_fields=[
                    ("sales", "[sales]"),
                    ("result", "ZN([sales])"),
                ],
            ),
        ),
    ],
)

FUNCTION_CASE = FunctionDocRegistryItem(
    name="_case_block_",
    naming_provider=CustomNamingProvider(
        title="CASE",
        internal_name="case",
    ),
    category=CATEGORY_LOGICAL,
    description=_(
        "Compares {arg:0} to {arg:1}, {arg:3}, ... consecutively "
        "and returns the corresponding result for the first match. "
        "If no match is found, it returns {arg:5}."
    ),
    signature_gen=TemplatedSignatureGenerator(
        placement_mode=SignaturePlacement.tabbed,
        signature_templates=(
            SignatureTemplate(
                title=_("As a block"),
                body="""
CASE {1}
    WHEN {2} THEN {3}
    [ WHEN {4} THEN {5}
      ... ]
    ELSE {6}
END""",
            ),
            SignatureTemplate(
                title=_("As a function"),
                body="""
CASE(
    {1},
    {2}, {3},
  [ {4}, {5},
    ... ]
    {6}
)""",
            ),
        ),
    ),
    examples=[
        SimpleExample(
            """
CASE (
    [country],
    "AO", "Angola",
    "AU", "Australia",
    "BY", "Belarus",
    "CA", "Canada",
    "TT", "Trinidad and Tobago",
    "Other Country"
)"""
        ),
        SimpleExample(
            """
CASE [country]
    WHEN "AO" THEN "Angola"
    WHEN "AU" THEN "Australia"
    WHEN "BY" THEN "Belarus"
    WHEN "CA" THEN "Canada"
    WHEN "TT" THEN "Trinidad and Tobago"
    ELSE "Other Country"
END"""
        ),
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with data table"),
                source=_SOURCE_TIME_1,
                formula_fields=[
                    ("unit", "[unit]"),
                    ("case_function", 'CASE([unit], "s", 1, "m", 60, "h", 3600, 0)'),
                    ("case_statement", 'CASE [unit] WHEN "s" THEN 1 WHEN "m" THEN 60 WHEN "h" THEN 3600 ELSE 0 END'),
                ],
                formulas_as_names=False,
            ),
        ),
    ],
)

FUNCTION_IF = FunctionDocRegistryItem(
    name="_if_block_",
    naming_provider=CustomNamingProvider(
        title="IF",
        internal_name="if",
    ),
    category=CATEGORY_LOGICAL,
    description=_(
        "Checks conditional expressions {arg:0}, {arg:1}, ... and returns the "
        "matching result for the first condition found to be `TRUE`. IF all "
        "conditional expressions are `FALSE`, it returns {arg:4}."
    ),
    signature_gen=TemplatedSignatureGenerator(
        placement_mode=SignaturePlacement.tabbed,
        signature_templates=(
            SignatureTemplate(
                title=_("As a block"),
                body="""
IF {1}
    THEN {2}
[ ELSEIF {3}
    THEN {4}
  ... ]
ELSE
    {5}
END""",
            ),
            SignatureTemplate(
                title=_("As a function"),
                body="""
IF(
    {1}, {2},
  [ {3}, {4},
    ... ]
    {5}
)""",
            ),
        ),
    ),
    examples=[
        SimpleExample(
            """
IF
    [Profit] > 100
        THEN "High"
    ELSEIF [Profit] > 25
        THEN "Medium"
    ELSE "Low"
END"""
        ),
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with data table"),
                source=_SOURCE_SALES_1,
                formula_fields=[
                    ("sales", "[sales]"),
                    (
                        "if_function",
                        (
                            'IF(ZN([sales]) < 100, "Less than 100", '
                            '[sales] < 1000, "100 - 1000", '
                            '"1000 and greater")'
                        ),
                    ),
                    (
                        "if_statement",
                        (
                            'IF ZN([sales]) < 100 THEN "Less than 100" '
                            'ELSEIF [sales] < 1000 THEN "100 - 1000" '
                            'ELSE "1000 and greater" END'
                        ),
                    ),
                ],
                formulas_as_names=False,
            ),
        ),
    ],
)

FUNCTIONS_LOGICAL = [
    FUNCTION_ISNULL,
    FUNCTION_IFNULL,
    # FUNCTION_ISNAN,
    # FUNCTION_IFNAN,
    FUNCTION_ZN,
    FUNCTION_CASE,
    FUNCTION_IF,
]

lib/dl_formula_ref/dl_formula_ref/functions/logical.py
datalens-tech/datalens-backend
from dl_formula_ref.categories.markup import CATEGORY_MARKUP
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import SimpleExample


_ = get_gettext()


FUNCTION_BOLD = FunctionDocRegistryItem(
    name="bold",
    category=CATEGORY_MARKUP,
    description=_("Stylizes the passed text in bold font."),
)

FUNCTION_ITALIC = FunctionDocRegistryItem(
    name="italic",
    category=CATEGORY_MARKUP,
    description=_("Stylizes the passed text in cursive font."),
)

FUNCTION_URL = FunctionDocRegistryItem(
    name="url",
    category=CATEGORY_MARKUP,
    description=_("Wraps {arg:1} into a hyperlink to URL {arg:0}."),
    examples=[
        SimpleExample("URL('https://example.com/?value=' + [value], [value])"),
    ],
)

FUNCTION_MARKUP = FunctionDocRegistryItem(
    name="markup",
    category=CATEGORY_MARKUP,
    description=_("Merges marked up text pieces. " "Can also be used for converting strings to marked up text."),
    examples=[
        SimpleExample("MARKUP(BOLD([a]), ': ', [b])"),
    ],
)

FUNCTIONS_MARKUP = [
    FUNCTION_BOLD,
    FUNCTION_ITALIC,
    FUNCTION_URL,
    FUNCTION_MARKUP,
]

lib/dl_formula_ref/dl_formula_ref/functions/markup.py
datalens-tech/datalens-backend
from dl_formula.core.datatype import DataType
from dl_formula_ref.categories.mathematical import CATEGORY_MATHEMATICAL
from dl_formula_ref.examples.config import (
    ExampleConfig,
    ExampleSource,
)
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import (
    DataExample,
    SimpleExample,
)


_ = get_gettext()

_SOURCE_1_VAL_1 = ExampleSource(
    columns=[("value", DataType.FLOAT)],
    data=[[1.0], [0.1], [-2.0], [50.5], [0.0], [-3.5]],
)

_SOURCE_SQ_1 = ExampleSource(
    columns=[("int_value", DataType.INTEGER), ("float_value", DataType.FLOAT)],
    data=[[0, 0.0], [1, 1.0], [3, 3.0], [4, 4.0], [13, 13.0], [16, 16.0], [20, 20.20], [25, 25.0]],
)

_SOURCE_TRIG_1 = ExampleSource(
    columns=[("n", DataType.FLOAT)],
    data=[[-1.0], [-0.5], [-0.25], [0.0], [0.25], [0.5], [1.0]],
)

_SOURCE_TRIG_2 = ExampleSource(
    columns=[("n", DataType.FLOAT)],
    data=[[-1.0], [-0.86602540378], [-0.70710678118], [-0.5], [0.0], [0.5], [0.70710678118], [0.86602540378], [1.0]],
)


_DIV_COLUMNS = [("numerator", DataType.FLOAT), ("denominator", DataType.FLOAT)]

_SOURCE_DIV = ExampleSource(columns=_DIV_COLUMNS, data=[[4, 2], [5, 3], [5, 2.0], [2.5, 1.2]])

_SOURCE_DIV_SAFE = ExampleSource(columns=_DIV_COLUMNS, data=[[5, 2], [5, 0]])


FUNCTION_ABS = FunctionDocRegistryItem(
    name="abs",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the absolute value of {arg:0}."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_1_VAL_1,
                formula_fields=[("value", "[value]"), ("result", "ABS([value])")],
            ),
        ),
    ],
)

FUNCTION_ACOS = FunctionDocRegistryItem(
    name="acos",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the arccosine of {arg:0} in radians."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_TRIG_2,
                formula_fields=[("value", "[n]"), ("result", "ACOS([n])")],
            ),
        ),
    ],
)

FUNCTION_ASIN = FunctionDocRegistryItem(
    name="asin",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the arcsine of {arg:0} in radians."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_TRIG_2,
                formula_fields=[("value", "[n]"), ("result", "ASIN([n])")],
            ),
        ),
    ],
)

FUNCTION_ATAN = FunctionDocRegistryItem(
    name="atan",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the arctangent of {arg:0} in radians."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_TRIG_2,
                formula_fields=[("value", "[n]"), ("result", "ATAN([n])")],
            ),
        ),
    ],
)

FUNCTION_ATAN2 = FunctionDocRegistryItem(
    name="atan2",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the arctangent in radians for the specified coordinates {arg:0} and " "{arg:1}."),
    examples=[
        SimpleExample("ATAN2(5, 7) = 0.62024"),
    ],
)

FUNCTION_CEILING = FunctionDocRegistryItem(
    name="ceiling",
    category=CATEGORY_MATHEMATICAL,
    description=_("Rounds the value up to the nearest integer."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_1_VAL_1,
                formula_fields=[("value", "[value]"), ("result", "CEILING([value])")],
            ),
        ),
    ],
)

FUNCTION_FLOOR = FunctionDocRegistryItem(
    name="floor",
    category=CATEGORY_MATHEMATICAL,
    description=_("Rounds the value down to the nearest integer."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_1_VAL_1,
                formula_fields=[("value", "[value]"), ("result", "FLOOR([value])")],
            ),
        ),
    ],
)

FUNCTION_COS = FunctionDocRegistryItem(
    name="cos",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the cosine of {arg:0} in radians."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_TRIG_1,
                formula_fields=[("n", "[n]"), ("angle", "[n] * PI()"), ("result", "COS([n]*PI())")],
            ),
        ),
    ],
)

FUNCTION_COT = FunctionDocRegistryItem(
    name="cot",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the cotangent of {arg:0} in radians."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_TRIG_1,
                formula_fields=[("n", "[n]"), ("angle", "[n] * PI()"), ("result", "COT([n]*PI())")],
            ),
        ),
    ],
)

FUNCTION_DEGREES = FunctionDocRegistryItem(
    name="degrees",
    category=CATEGORY_MATHEMATICAL,
    description=_("Converts radians to degrees."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_TRIG_1,
                formula_fields=[("n", "[n]"), ("angle", "[n] * PI()"), ("result", "DEGREES([n]*PI())")],
            ),
        ),
    ],
)

FUNCTION_DIV = FunctionDocRegistryItem(
    name="div",
    category=CATEGORY_MATHEMATICAL,
    description=_("Divides {arg:0} by {arg:1}. The result is rounded down to the nearest " "integer."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_DIV,
                formula_fields=[
                    ("numerator", "[numerator]"),
                    ("denominator", "[denominator]"),
                    ("result", "DIV([numerator], [denominator])"),
                ],
            ),
        ),
    ],
)

FUNCTION_DIV_SAFE = FunctionDocRegistryItem(
    name="div_safe",
    category=CATEGORY_MATHEMATICAL,
    description=_(
        "Divides {arg:0} by {arg:1}. Returns {arg:2} if division by zero occurs. "
        "If the number {arg:2} is omitted, it is assumed to be `NULL`.\n"
        "The result is rounded down to the nearest integer."
    ),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_DIV_SAFE,
                formula_fields=[
                    ("numerator", "[numerator]"),
                    ("denominator", "[denominator]"),
                    ("result2", "DIV_SAFE([numerator], [denominator])"),
                    ("result3", "DIV_SAFE([numerator], [denominator], 42)"),
                ],
            ),
        ),
    ],
)

FUNCTION_FDIV_SAFE = FunctionDocRegistryItem(
    name="fdiv_safe",
    category=CATEGORY_MATHEMATICAL,
    description=_(
        "Divides {arg:0} by {arg:1}. Returns {arg:2} if division by zero occurs. "
        "If the number {arg:2} is omitted, it is assumed to be `NULL`."
    ),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_DIV_SAFE,
                formula_fields=[
                    ("numerator", "[numerator]"),
                    ("denominator", "[denominator]"),
                    ("result2", "FDIV_SAFE([numerator], [denominator])"),
                    ("result3", "FDIV_SAFE([numerator], [denominator], 42)"),
                ],
            ),
        ),
    ],
)

FUNCTION_EXP = FunctionDocRegistryItem(
    name="exp",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the result of raising the number 'e' to the power of {arg:0}."),
    examples=[
        SimpleExample("EXP(0) = 1.0"),
        SimpleExample("EXP(1) = 2.718282"),
        SimpleExample("EXP(3) = 20.08553"),
    ],
)

FUNCTION_LN = FunctionDocRegistryItem(
    name="ln",
    category=CATEGORY_MATHEMATICAL,
    description=_(
        "Returns the natural logarithm of the number {arg:0}. Returns `NULL` if the "
        "number is less than or equal to 0."
    ),
    examples=[
        SimpleExample("LN(1) = 0.0"),
        SimpleExample("LN(EXP(2)) = 2.0"),
    ],
)

FUNCTION_LOG10 = FunctionDocRegistryItem(
    name="log10",
    category=CATEGORY_MATHEMATICAL,
    description=_(
        "Returns the logarithm of the number {arg:0} to base 10. Returns `NULL` if "
        "the number is less than or equal to 0."
    ),
    examples=[
        SimpleExample("LOG10(1) = 0.0"),
        SimpleExample("LOG10(1000) = 3.0"),
        SimpleExample("LOG10(100) = 2.0"),
    ],
)

FUNCTION_LOG = FunctionDocRegistryItem(
    name="log",
    category=CATEGORY_MATHEMATICAL,
    description=_(
        "Returns the logarithm of {arg:0} to base {arg:1}. Returns `NULL` if the "
        "number {arg:0} is less than or equal to 0."
    ),
    examples=[
        SimpleExample("LOG(1, 2.6) = 0.0"),
        SimpleExample("LOG(1024, 2) = 10.0"),
        SimpleExample("LOG(100, 10) = 2.0"),
    ],
)

FUNCTION_GREATEST = FunctionDocRegistryItem(
    name="greatest",
    category=CATEGORY_MATHEMATICAL,
    description=_(
        "Returns the greatest value.\n"
        "\n"
        "See also {ref:LEAST}.\n"
        "\n"
        "Depending on the specified data type, it returns:\n"
        "- The greatest number.\n"
        "- The last string in alphabetical order.\n"
        "- The latest date.\n"
        "- `TRUE` when selecting between `TRUE` and `FALSE` for Boolean type."
    ),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            "GREATEST(3.4, 2.6) = 3.4",
            'GREATEST("3.4", "2.6") = "3.4"',
            "GREATEST(#2019-01-02#, #2019-01-17#) = #2019-01-17#",
            "GREATEST(#2019-01-02 04:03:02#, #2019-01-17 03:02:01#) = #2019-01-17 03:02:01#",
            "GREATEST(TRUE, FALSE) = TRUE",
            "GREATEST(34, 5, 7, 3, 99, 1, 2, 2, 56) = 99",
            "GREATEST(5.6, 1.2, 7.8, 3.4) = 7.8",
            "GREATEST(#2019-01-02#, #2019-01-17#, #2019-01-10#) = #2019-01-17#",
        )
    ],
)

FUNCTION_LEAST = FunctionDocRegistryItem(
    name="least",
    category=CATEGORY_MATHEMATICAL,
    description=_(
        "Returns the smallest value.\n"
        "\n"
        "See also {ref:GREATEST}.\n"
        "\n"
        "Depending on the specified data type, it returns:\n"
        "- The smallest number.\n"
        "- The first string in alphabetical order.\n"
        "- The earliest date.\n"
        "- `FALSE` when selecting between `TRUE` and `FALSE` for Boolean type."
    ),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            "LEAST(3.4, 2.6) = 2.6",
            'LEAST("3.4", "2.6") = "2.6"',
            "LEAST(#2019-01-02#, #2019-01-17#) = #2019-01-02#",
            "LEAST(#2019-01-02 04:03:02#, #2019-01-17 03:02:01#) = #2019-01-02 04:03:02#",
            "LEAST(TRUE, FALSE) = FALSE",
            "LEAST(34, 5, 7, 3, 99, 1, 2, 2, 56) = 1",
            "LEAST(5.6, 1.2, 7.8, 3.4) = 1.2",
            "LEAST(#2019-01-02#, #2019-01-17#, #2019-01-10#) = #2019-01-02#",
        )
    ],
)

FUNCTION_PI = FunctionDocRegistryItem(
    name="pi",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns PI. The accuracy depends on the data source."),
    examples=[
        SimpleExample("PI() = 3.14159"),
    ],
)

FUNCTION_POWER = FunctionDocRegistryItem(
    name="power",
    category=CATEGORY_MATHEMATICAL,
    description=_("Raises {arg:0} to the power of {arg:1}."),
    examples=[
        SimpleExample("POWER(2.3, 4.5) = 42.43998894277659"),
        SimpleExample("POWER(6, 2) = 36.0"),
    ],
)

FUNCTION_RADIANS = FunctionDocRegistryItem(
    name="radians",
    category=CATEGORY_MATHEMATICAL,
    description=_("Converts {arg:0} degrees to radians."),
    examples=[
        SimpleExample("RADIANS(0) = 0.0"),
        SimpleExample("RADIANS(180) = 3.14159"),
    ],
)

FUNCTION_ROUND = FunctionDocRegistryItem(
    name="round",
    category=CATEGORY_MATHEMATICAL,
    description=_(
        "Rounds the number {arg:0} to the number of decimal digits specified in "
        "{arg:1}.\n"
        "If the number {arg:1} is omitted, {arg:0} is rounded to the nearest integer."
    ),
    examples=[
        SimpleExample("ROUND(3.14159) = 3"),
        SimpleExample("ROUND(3.14159, 3) = 3.142"),
    ],
)

FUNCTION_SIGN = FunctionDocRegistryItem(
    name="sign",
    category=CATEGORY_MATHEMATICAL,
    description=_(
        "Returns the sign of the number {arg:0}:\n"
        "- `-1` if the number is negative.\n"
        "`0` if the number is zero.\n"
        "- `1` if the number is positive."
    ),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_1_VAL_1,
                formula_fields=[("value", "[value]"), ("result", "SIGN([value])")],
            ),
        ),
    ],
)

FUNCTION_SIN = FunctionDocRegistryItem(
    name="sin",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the sine of {arg:0} in radians."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_TRIG_1,
                formula_fields=[("n", "[n]"), ("angle", "[n] * PI()"), ("result", "SIN([n]*PI())")],
            ),
        ),
    ],
)

FUNCTION_SQRT = FunctionDocRegistryItem(
    name="sqrt",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the square root of the specified number."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_SQ_1,
                formula_fields=[
                    ("int_value", "[int_value]"),
                    ("float_value", "[float_value]"),
                    ("int_result", "SQRT([int_value])"),
                    ("float_result", "SQRT([float_value])"),
                ],
            ),
        ),
    ],
)

FUNCTION_SQUARE = FunctionDocRegistryItem(
    name="square",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the number {arg:0} raised to the power of 2."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_SQ_1,
                formula_fields=[
                    ("int_value", "[int_value]"),
                    ("float_value", "[float_value]"),
                    ("int_result", "SQUARE([int_value])"),
                    ("float_result", "SQUARE([float_value])"),
                ],
            ),
        ),
    ],
)

FUNCTION_TAN = FunctionDocRegistryItem(
    name="tan",
    category=CATEGORY_MATHEMATICAL,
    description=_("Returns the tangent of {arg:0} in radians."),
    examples=[
        DataExample(
            example_config=ExampleConfig(
                source=_SOURCE_TRIG_1,
                formula_fields=[("n", "[n]"), ("angle", "[n] * PI()"), ("result", "tan([n]*PI())")],
            ),
        ),
    ],
)

FUNCTION_COMPARE = FunctionDocRegistryItem(
    name="compare",
    category=CATEGORY_MATHEMATICAL,
    description=_(
        "Returns:\n"
        "* 0 if {arg:0} and {arg:1} differs by not more than {arg:2}.\n"
        "* -1 if {arg:0} is less than {arg:1} by more than {arg:2}.\n"
        "* 1 if {arg:0} is greater than {arg:1} by more than {arg:2}."
    ),
    examples=[
        SimpleExample("COMPARE(1.25, 1.26, 0.1) = 0"),
        SimpleExample("COMPARE(1.25, 1.26, 0.001) = -1"),
        SimpleExample("COMPARE(1.26, 1.25, 0.001) = 1"),
    ],
)

FUNCTIONS_MATHEMATICAL = [
    FUNCTION_ABS,
    FUNCTION_ACOS,
    FUNCTION_ASIN,
    FUNCTION_ATAN,
    FUNCTION_ATAN2,
    FUNCTION_CEILING,
    FUNCTION_FLOOR,
    FUNCTION_COS,
    FUNCTION_COT,
    FUNCTION_DEGREES,
    FUNCTION_DIV,
    FUNCTION_DIV_SAFE,
    FUNCTION_FDIV_SAFE,
    FUNCTION_EXP,
    FUNCTION_LN,
    FUNCTION_LOG10,
    FUNCTION_LOG,
    FUNCTION_GREATEST,
    FUNCTION_LEAST,
    FUNCTION_PI,
    FUNCTION_POWER,
    FUNCTION_RADIANS,
    FUNCTION_ROUND,
    FUNCTION_SIGN,
    FUNCTION_SIN,
    FUNCTION_SQRT,
    FUNCTION_SQUARE,
    FUNCTION_TAN,
    FUNCTION_COMPARE,
]

lib/dl_formula_ref/dl_formula_ref/functions/mathematical.py
datalens-tech/datalens-backend
from dl_formula_ref.categories.operator import CATEGORY_OPERATOR
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.aliased_res import (
    AliasedTableResource,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import SimpleExample
from dl_formula_ref.registry.naming import CustomNamingProvider
from dl_formula_ref.registry.note import (
    Note,
    NoteLevel,
)
from dl_formula_ref.registry.signature_gen import (
    SignatureTemplate,
    TemplatedSignatureGenerator,
)


_ = get_gettext()


_SIGNATURE_OP_UNARY = TemplatedSignatureGenerator(signature_templates=(SignatureTemplate(body="{0} {1}"),))
_SIGNATURE_OP_BINARY = TemplatedSignatureGenerator(signature_templates=(SignatureTemplate(body="{1} {0} {2}"),))


FUNCTION_OP_NOT = FunctionDocRegistryItem(
    name="not",
    naming_provider=CustomNamingProvider(
        internal_name="op_not",
    ),
    category=CATEGORY_OPERATOR,
    description=_("Inverts a Boolean value."),
    signature_gen=_SIGNATURE_OP_UNARY,
    examples=[
        SimpleExample(example_str)
        for example_str in (
            "NOT FALSE = TRUE",
            "NOT TRUE = FALSE",
            'NOT "" = TRUE',
            'NOT "text" = FALSE',
            "NOT 0 = TRUE",
            "NOT 1 = FALSE",
            "NOT #2019-01-01# = FALSE",
            "NOT #2019-01-01 03:00:00# = FALSE",
        )
    ],
)

FUNCTION_OP_NEGATION = FunctionDocRegistryItem(
    name="neg",
    naming_provider=CustomNamingProvider(
        title=_("Negation (-)"),
        internal_name="op_negation",
    ),
    category=CATEGORY_OPERATOR,
    description=_("Returns the number {arg:0} with the opposite sign."),
    signature_gen=TemplatedSignatureGenerator(signature_templates=(SignatureTemplate(body="-{1}"),)),
    examples=[
        SimpleExample("- (5) = -5"),
    ],
)

FUNCTION_OP_IS_TRUE = FunctionDocRegistryItem(
    name="istrue",
    naming_provider=CustomNamingProvider(
        title="IS TRUE",  # Not translatable
    ),
    category=CATEGORY_OPERATOR,
    description=_(
        "Checks whether the value of {arg:0} is true (`TRUE`).\n"
        "\n"
        "The `{argn:0} IS NOT TRUE` option returns the opposite value."
    ),
    signature_gen=TemplatedSignatureGenerator(signature_templates=(SignatureTemplate(body="{1} IS [ NOT ] TRUE"),)),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            '"qwerty" IS TRUE = TRUE',
            "123 IS TRUE = TRUE",
            "TRUE IS TRUE = TRUE",
            "FALSE IS TRUE = FALSE",
            "FALSE IS NOT TRUE = TRUE",
            "#2019-03-05# IS TRUE = TRUE",
            "#2019-03-05 01:02:03# IS TRUE = TRUE",
        )
    ],
)

FUNCTION_OP_IS_FALSE = FunctionDocRegistryItem(
    name="isfalse",
    naming_provider=CustomNamingProvider(
        title="IS FALSE",  # Not translatable
    ),
    category=CATEGORY_OPERATOR,
    description=_(
        "Checks whether the {arg:0} value is false (`FALSE`).\n"
        "\n"
        "The `{argn:0} IS NOT FALSE` option returns the opposite value."
    ),
    signature_gen=TemplatedSignatureGenerator(signature_templates=(SignatureTemplate(body="{1} IS [ NOT ] FALSE"),)),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            '"" IS FALSE = TRUE',
            "0 IS FALSE = TRUE",
            "FALSE IS FALSE = TRUE",
            "TRUE IS FALSE = FALSE",
            "TRUE IS NOT FALSE = TRUE",
            "#2019-03-05# IS FALSE = FALSE",
            "#2019-03-05 01:02:03# IS FALSE = FALSE",
        )
    ],
)

FUNCTION_OP_POWER = FunctionDocRegistryItem(
    name="^",
    naming_provider=CustomNamingProvider(
        title=_("Power (^)"),
        internal_name="op_power",
    ),
    category=CATEGORY_OPERATOR,
    description=_("Raises {arg:0} to the power of {arg:1}."),
    signature_gen=_SIGNATURE_OP_BINARY,
    examples=[
        SimpleExample("2 ^ 3 = 8.0"),
        SimpleExample("2.1 ^ -0.3 = 0.80045"),
    ],
)

FUNCTION_OP_MULT = FunctionDocRegistryItem(
    name="*",
    naming_provider=CustomNamingProvider(
        title=_("Multiplication (*)"),
        internal_name="op_mult",
    ),
    category=CATEGORY_OPERATOR,
    description=_(
        "If both arguments are numbers, it returns the result by multiplying {arg:0} "
        "by {arg:1}.\n"
        "\n"
        "If one of the arguments is {type:STRING} and the other is {type:INTEGER}, it "
        "returns the string repeated the specified number of times."
    ),
    signature_gen=_SIGNATURE_OP_BINARY,
    examples=[
        SimpleExample("2 * 3 = 6"),
        SimpleExample("2.1 * 3 = 6.3"),
        SimpleExample("\"Lorem\" * 3 = 'LoremLoremLorem'"),
        SimpleExample("3 * \"Lorem\" = 'LoremLoremLorem'"),
    ],
)

FUNCTION_OP_DIV = FunctionDocRegistryItem(
    name="/",
    naming_provider=CustomNamingProvider(
        title=_("Division (/)"),
        internal_name="op_div",
    ),
    category=CATEGORY_OPERATOR,
    description=_("Divides the number {arg:0} by the number {arg:1}."),
    signature_gen=_SIGNATURE_OP_BINARY,
    examples=[
        SimpleExample("4 / 2 = 2.0"),
        SimpleExample("5 / 2 = 2.5"),
        SimpleExample("5.0 / 2 = 2.5"),
        SimpleExample("5 / 2.0 = 2.5"),
    ],
)

FUNCTION_OP_MOD = FunctionDocRegistryItem(
    name="%",
    naming_provider=CustomNamingProvider(
        title=_("Modulo (%)"),
        internal_name="op_mod",
    ),
    category=CATEGORY_OPERATOR,
    description=_("Returns the remainder from dividing the first number {arg:0} by the second " "number {arg:1}."),
    signature_gen=_SIGNATURE_OP_BINARY,
    examples=[
        SimpleExample("2 % 3 = 2"),
        SimpleExample("2.1 % 3 = 2.1"),
        SimpleExample("3 % 2.1 = 0.9"),
    ],
)

FUNCTION_OP_PLUS = FunctionDocRegistryItem(
    name="+",
    naming_provider=CustomNamingProvider(
        title=_("Addition and concatenation (+)"),
        internal_name="op_plus",
    ),
    category=CATEGORY_OPERATOR,
    description=_(
        "Behaves differently depending on the argument types. Possible options are "
        "listed in the table:\n"
        "\n"
        "{table: desc_by_type}\n"
        "\n"
        "Changing the order of arguments does not affect the result."
    ),
    resources=SimpleAliasedResourceRegistry(
        {
            "desc_by_type": AliasedTableResource(
                table_body=[
                    [_("Type of {arg:0}"), _("Type of {arg:1}"), _("Return value")],
                    ["{type:FLOAT|INTEGER}", "{type:FLOAT|INTEGER}", _("The sum of the numbers {arg:0} and {arg:1}.")],
                    [
                        "{type:DATE}",
                        "{type:FLOAT|INTEGER}",
                        _(
                            "The date that is {arg:1} days greater than {arg:0} (rounded down to an "
                            "integer number of days)."
                        ),
                    ],
                    [
                        "{type:DATETIME|GENERICDATETIME}",
                        "{type:FLOAT|INTEGER}",
                        _(
                            "The date with time, {arg:1} days greater than {arg:0}. If {arg:1} contains a "
                            "fractional part, it is converted hours (`1/24`), minutes (`1/1440`), and "
                            "seconds (`1/86400`)."
                        ),
                    ],
                    [
                        "{type:STRING}",
                        "{type:STRING}",
                        _("The merging (concatenation) of strings {arg:0} and {arg:1}."),
                    ],
                    [
                        "{type:ARRAY_INT|ARRAY_FLOAT|ARRAY_STR}",
                        "{type:ARRAY_INT|ARRAY_FLOAT|ARRAY_STR}",
                        _("The merging (concatenation) of arrays {arg:0} and {arg:1}."),
                    ],
                ]
            )
        }
    ),
    signature_gen=_SIGNATURE_OP_BINARY,
    examples=[
        SimpleExample(example_str)
        for example_str in (
            "2 + 3 = 5",
            '"Lorem" + " ipsum" = "Lorem ipsum"',
            "#2019-01-06# + 2 = #2019-01-08#",
            "2.2 + #2019-01-06# = #2019-01-08#",
            "2 + #2019-01-06 03# = #2019-01-08 03:00:00#",
            "#2019-01-06 03# + 2.5 = #2019-01-08 15:00:00#",
        )
    ],
    notes=[
        Note(
            level=NoteLevel.warning,
            text=_("{dialects: METRIKAAPI} does not support string concatenation."),
        )
    ],
)

FUNCTION_OP_MINUS = FunctionDocRegistryItem(
    name="-",
    naming_provider=CustomNamingProvider(
        title=_("Subtraction (-)"),
        internal_name="op_minus",
    ),
    category=CATEGORY_OPERATOR,
    description=_(
        "Behaves differently depending on the argument types. Possible options are "
        "listed in the table:\n"
        "\n"
        "{table: desc_by_type}"
    ),
    resources=SimpleAliasedResourceRegistry(
        {
            "desc_by_type": AliasedTableResource(
                table_body=[
                    [_("Type of {arg:0}"), _("Type of {arg:1}"), _("Return value")],
                    [
                        "{type:FLOAT|INTEGER}",
                        "{type:FLOAT|INTEGER}",
                        _("The difference between the numbers {arg:0} and {arg:1}."),
                    ],
                    [
                        "{type:DATE}",
                        "{type:FLOAT|INTEGER}",
                        _(
                            "The date that is {arg:1} days smaller than {arg:0} (rounded down to an "
                            "integer number of days)."
                        ),
                    ],
                    [
                        "{type:DATETIME|GENERICDATETIME}",
                        "{type:FLOAT|INTEGER}",
                        _(
                            "The date with time, {arg:1} days smaller than {arg:0}. If {arg:1} contains a "
                            "fractional part, it is converted to hours (`1/24`), minutes (`1/1440`), and "
                            "seconds (`1/86400`)."
                        ),
                    ],
                    ["{type:DATE}", "{type:DATE}", _("The difference between two dates in days.")],
                    [
                        "{type:DATETIME}",
                        "{type:DATETIME}",
                        _(
                            "The difference between two dates in days: the integer part — the number of "
                            "whole days, the fractional part — the number of hours, minutes and seconds "
                            "expressed as a fraction of the whole day (1 hour is '1/24')."
                        ),
                    ],
                    [
                        "{type:GENERICDATETIME}",
                        "{type:GENERICDATETIME}",
                        _(
                            "The difference between two dates in days: the integer part — the number of "
                            "whole days, the fractional part — the number of hours, minutes and seconds "
                            "expressed as a fraction of the whole day (1 hour is '1/24')."
                        ),
                    ],
                ]
            )
        }
    ),
    signature_gen=_SIGNATURE_OP_BINARY,
    examples=[
        SimpleExample(example_str)
        for example_str in (
            "2 - 3 = -1",
            "2 - 0.5 = 1.5",
            "#2019-01-06# - 2 = #2019-01-04#",
            "#2019-01-06# - 2.2 = #2019-01-03#",
            "#2019-01-06 03:00:00# - 2 = #2019-01-04 03:00:00#",
            "#2019-01-06 03:00:00# - 2.5 = #2019-01-03 15:00:00#",
            "#2019-01-06# - #2019-01-02# = 4",
            "#2019-01-06 15:00:00# - #2019-01-02 03:00:00# = 4.5",
        )
    ],
)

FUNCTION_OP_LIKE = FunctionDocRegistryItem(
    name="like",
    category=CATEGORY_OPERATOR,
    description=_(
        "Matches the string {arg:0} to the template {arg: 1} and returns `TRUE` on "
        "match.\n"
        "You can specify the value in {arg:1} or use the `%` character to match a "
        "string of any length.\n"
        "\n"
        "The `{argn:0} NOT LIKE` option returns the opposite value."
    ),
    signature_gen=TemplatedSignatureGenerator(signature_templates=(SignatureTemplate(body="{1} [ NOT ] LIKE {2}"),)),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            '"raspberry" LIKE "%spb%" = TRUE',
            '"raspberry" NOT LIKE "%spb%" = FALSE',
            "IIF([Country] LIKE 'RU', 'Y', 'N')",
            "IIF([Phone] LIKE '+7%', 'RU', 'notRU')",
        )
    ],
)

FUNCTION_OP_COMPARISON = FunctionDocRegistryItem(
    name="==",
    category=CATEGORY_OPERATOR,
    naming_provider=CustomNamingProvider(
        title=_("Comparison"),
        internal_name="op_comparison",
    ),
    description=_("Compares the value {arg:0} with the value {arg:1}."),
    notes=[
        Note(
            _(
                "Due to implementation details of the {type:FLOAT} type in {dialects:CLICKHOUSE} sources "
                "it is recommended to use the {ref:COMPARE} function instead of comparison operators for this type."
            ),
        )
    ],
    signature_gen=TemplatedSignatureGenerator(
        signature_templates=(
            SignatureTemplate(title=_("Equality"), body="{1} = {2}"),
            SignatureTemplate(title=_("Inequality"), body="{1} != {2}"),
            SignatureTemplate(title=_("Less than"), body="{1} < {2}"),
            SignatureTemplate(title=_("Less than or equal"), body="{1} <= {2}"),
            SignatureTemplate(title=_("Greater than"), body="{1} > {2}"),
            SignatureTemplate(title=_("Greater than or equal"), body="{1} >= {2}"),
        )
    ),
    examples=[
        SimpleExample("1 = 1 = TRUE"),
        SimpleExample("7 > 2 > 1 = TRUE"),
    ],
)

FUNCTION_OP_AND = FunctionDocRegistryItem(
    name="and",
    category=CATEGORY_OPERATOR,
    description=_("Performs a Boolean join of two expressions with the `AND` condition."),
    signature_gen=_SIGNATURE_OP_BINARY,
    examples=[
        SimpleExample('[Profit] > 30 AND [Weekday] IN ("Saturday", "Sunday")'),
    ],
)

FUNCTION_OP_OR = FunctionDocRegistryItem(
    name="or",
    category=CATEGORY_OPERATOR,
    description=_("Performs a Boolean join of two expressions with the `OR` condition."),
    signature_gen=_SIGNATURE_OP_BINARY,
    examples=[
        SimpleExample('[Month] = "January" OR [Year] < 2019'),
    ],
)

FUNCTION_OP_IN = FunctionDocRegistryItem(
    name="in",
    category=CATEGORY_OPERATOR,
    description=_(
        "Checks whether the value matches at least one of the values listed in "
        "`IN(...)`.\n"
        "\n"
        "The option `{argn:0} NOT IN (<{arg:1}>)` returns the opposite value."
    ),
    signature_gen=TemplatedSignatureGenerator(signature_templates=(SignatureTemplate(body="{1} [ NOT ] IN (<{2}>)"),)),
    examples=[
        SimpleExample("3 IN (23, 5, 3, 67) = TRUE"),
        SimpleExample("3 NOT IN (23, 5, 3, 67) = FALSE"),
    ],
)

FUNCTION_OP_BETWEEN = FunctionDocRegistryItem(
    name="between",
    category=CATEGORY_OPERATOR,
    description=_(
        "Returns `TRUE` if {arg:0} is in the range from {arg:1} to {arg:2}.\n"
        "\n"
        "The option `{argn:0} NOT BETWEEN {argn:1} AND {argn:2}` returns the "
        "opposite value."
    ),
    signature_gen=TemplatedSignatureGenerator(
        signature_templates=(SignatureTemplate(body="{1} [ NOT ] BETWEEN {2} AND {3}"),)
    ),
    examples=[
        SimpleExample("3 BETWEEN 1 AND 100 = TRUE"),
        SimpleExample("3 NOT BETWEEN 1 AND 100 = FALSE"),
    ],
)

FUNCTIONS_OPERATOR = [
    FUNCTION_OP_NOT,
    FUNCTION_OP_NEGATION,
    FUNCTION_OP_IS_TRUE,
    FUNCTION_OP_IS_FALSE,
    FUNCTION_OP_POWER,
    FUNCTION_OP_MULT,
    FUNCTION_OP_DIV,
    FUNCTION_OP_MOD,
    FUNCTION_OP_PLUS,
    FUNCTION_OP_MINUS,
    FUNCTION_OP_LIKE,
    FUNCTION_OP_COMPARISON,
    FUNCTION_OP_AND,
    FUNCTION_OP_OR,
    FUNCTION_OP_IN,
    FUNCTION_OP_BETWEEN,
]

lib/dl_formula_ref/dl_formula_ref/functions/operator.py
datalens-tech/datalens-backend
from dl_formula.core.datatype import DataType
from dl_formula_ref.categories.string import CATEGORY_STRING
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import SimpleExample
from dl_formula_ref.registry.impl_selector import ArgAwareImplementationSelector
from dl_formula_ref.registry.naming import CategoryPostfixNamingProvider
from dl_formula_ref.registry.note import Note


_ = get_gettext()


FUNCTION_ASCII = FunctionDocRegistryItem(
    name="ascii",
    category=CATEGORY_STRING,
    description=_("Returns the numeric representation of the first character of the string."),
    examples=[
        SimpleExample('ASCII("N") = 78'),
        SimpleExample('ASCII("¹") = 185'),  # \u00b9 = \xb9 = 185;  utf-8 \xc2\xb9, i.e. 194 185
        SimpleExample('ASCII("¹") = 194'),  # binary utf-8 interpretation, e.g. mysql, clickhouse
        SimpleExample('ASCII("…") = 8230'),  # \u2026 = 8230;  utf-8 \xe2\x80\xa6 = 226 128 166
    ],
)

FUNCTION_CHAR = FunctionDocRegistryItem(
    name="char",
    category=CATEGORY_STRING,
    description=_("Converts the numeric representation of an ASCII character to a value."),
    examples=[
        SimpleExample('CHAR(78) = "N"'),
    ],
)

FUNCTION_CONCAT = FunctionDocRegistryItem(
    name="concat",
    category=CATEGORY_STRING,
    description=_(
        "Merges any number of strings. When non-string types are used, they're " "converted to strings and then merged."
    ),
    examples=[
        SimpleExample('CONCAT("Date of birth ", #2019-01-23#) = "Date of birth 2019-01-23"'),
        SimpleExample('CONCAT(2019, 01, 23) = "20190123"'),
    ],
)

FUNCTION_CONTAINS = FunctionDocRegistryItem(
    name="contains",
    category=CATEGORY_STRING,
    impl_selector=ArgAwareImplementationSelector(
        exp_arg_types={
            0: {
                DataType.BOOLEAN,
                DataType.DATE,
                DataType.DATETIME,
                DataType.GENERICDATETIME,
                DataType.FLOAT,
                DataType.GEOPOINT,
                DataType.GEOPOLYGON,
                DataType.INTEGER,
                DataType.STRING,
                DataType.UUID,
            },
        },
    ),
    naming_provider=CategoryPostfixNamingProvider(),
    description=_("Returns `TRUE` if {arg:0} contains {arg:1}. For case-insensitive searches, " "see {ref:ICONTAINS}."),
    examples=[
        SimpleExample('CONTAINS("RU0891923", "RU") = TRUE'),
        SimpleExample('CONTAINS("Lorem ipsum", "abc") = FALSE'),
    ],
)

FUNCTION_ICONTAINS = FunctionDocRegistryItem(
    name="icontains",
    category=CATEGORY_STRING,
    description=_("Case-insensitive version of {ref:string/CONTAINS}. Returns `TRUE` if {arg:0} " "contains {arg:1}."),
    examples=[
        SimpleExample('ICONTAINS("RU0891923", "ru") = TRUE'),
        SimpleExample('ICONTAINS("Lorem ipsum", "abc") = FALSE'),
    ],
)

FUNCTION_ENDSWITH = FunctionDocRegistryItem(
    name="endswith",
    category=CATEGORY_STRING,
    description=_("Returns `TRUE` if {arg:0} ends in {arg:1}. For case-insensitive searches, " "see {ref:IENDSWITH}."),
    examples=[
        SimpleExample('ENDSWITH("Petrov Ivan", "Ivan") = TRUE'),
        SimpleExample('ENDSWITH("Lorem ipsum", "sum") = TRUE'),
        SimpleExample('ENDSWITH("Lorem ipsum", "abc") = FALSE'),
    ],
)

FUNCTION_IENDSWITH = FunctionDocRegistryItem(
    name="iendswith",
    category=CATEGORY_STRING,
    description=_("Case-insensitive version of {ref:ENDSWITH}. Returns `TRUE` if {arg:0} ends " "in {arg:1}."),
    examples=[
        SimpleExample('IENDSWITH("PETROV IVAN", "Ivan") = TRUE'),
        SimpleExample('IENDSWITH("Lorem ipsum", "SUM") = TRUE'),
        SimpleExample('IENDSWITH("Lorem ipsum", "abc") = FALSE'),
    ],
)

FUNCTION_STARTSWITH = FunctionDocRegistryItem(
    name="startswith",
    category=CATEGORY_STRING,
    impl_selector=ArgAwareImplementationSelector(
        exp_arg_types={
            0: {
                DataType.BOOLEAN,
                DataType.DATE,
                DataType.DATETIME,
                DataType.GENERICDATETIME,
                DataType.FLOAT,
                DataType.GEOPOINT,
                DataType.GEOPOLYGON,
                DataType.INTEGER,
                DataType.STRING,
                DataType.UUID,
            },
        },
    ),
    naming_provider=CategoryPostfixNamingProvider(),
    description=_(
        "Returns `TRUE` if {arg:0} starts with {arg:1}. For case-insensitive " "searches, see {ref:ISTARTSWITH}."
    ),
    examples=[
        SimpleExample('STARTSWITH("Petrov Ivan", "Petrov") = TRUE'),
        SimpleExample('STARTSWITH("Lorem ipsum", "Lore") = TRUE'),
        SimpleExample('STARTSWITH("Lorem ipsum", "abc") = FALSE'),
    ],
)

FUNCTION_ISTARTSWITH = FunctionDocRegistryItem(
    name="istartswith",
    category=CATEGORY_STRING,
    description=_(
        "Case-insensitive version of {ref:string/STARTSWITH}. Returns `TRUE`  if {arg:0} " "starts with {arg:1}."
    ),
    examples=[
        SimpleExample('ISTARTSWITH("petrov ivan", "Petrov") = TRUE'),
        SimpleExample('ISTARTSWITH("Lorem ipsum", "LORE") = TRUE'),
        SimpleExample('ISTARTSWITH("Lorem ipsum", "abc") = FALSE'),
    ],
)

FUNCTION_FIND = FunctionDocRegistryItem(
    name="find",
    category=CATEGORY_STRING,
    description=_(
        "Returns the index of the position of the first character of the substring "
        "{arg:1} in the string {arg:0}.\n"
        "\n"
        "If the `start_index` option is specified, the search starts from the "
        "specified position."
    ),
    examples=[
        SimpleExample('FIND("Lorem ipsum dolor sit amet", "abc") = 0'),
        SimpleExample('FIND("Lorem ipsum dolor sit amet", "or") = 2'),
        SimpleExample('FIND("Lorem ipsum dolor sit amet", "or", 7) = 16'),
    ],
)

FUNCTION_LEFT = FunctionDocRegistryItem(
    name="left",
    category=CATEGORY_STRING,
    description=_(
        "Returns a string that contains the number of characters specified in {arg:1} "
        "from the beginning of the string {arg:0}."
    ),
    examples=[
        SimpleExample('LEFT("Computer", 4) = "Comp"'),
    ],
)

FUNCTION_RIGHT = FunctionDocRegistryItem(
    name="right",
    category=CATEGORY_STRING,
    description=_(
        "Returns a string that contains the number of characters specified in {arg:1} "
        "from the end of the string {arg:0}."
    ),
    examples=[
        SimpleExample('RIGHT("Computer", 3) = "ter"'),
    ],
)

FUNCTION_LEN = FunctionDocRegistryItem(
    name="len",
    category=CATEGORY_STRING,
    description=_("Returns the number of characters in the string or items in array {arg:0}."),
    examples=[
        SimpleExample('LEN("Computer") = 8'),
    ],
)

FUNCTION_LOWER = FunctionDocRegistryItem(
    name="lower",
    category=CATEGORY_STRING,
    description=_("Returns the string {arg:0} in lowercase."),
    examples=[
        SimpleExample('LOWER("Lorem ipsum") = "lorem ipsum"'),
        SimpleExample('LOWER("Карл у Клары") = "карл у клары"'),
    ],
)

FUNCTION_UPPER = FunctionDocRegistryItem(
    name="upper",
    category=CATEGORY_STRING,
    description=_("Returns the string {arg:0} in uppercase."),
    examples=[
        SimpleExample('UPPER("Lorem ipsum") = "LOREM IPSUM"'),
        SimpleExample('UPPER("Карл у Клары") = "КАРЛ У КЛАРЫ"'),
    ],
)

FUNCTION_LTRIM = FunctionDocRegistryItem(
    name="ltrim",
    category=CATEGORY_STRING,
    description=_("Returns the string {arg:0} without spaces at the beginning of the string."),
    examples=[
        SimpleExample('LTRIM(" Computer") = "Computer"'),
    ],
)

FUNCTION_RTRIM = FunctionDocRegistryItem(
    name="rtrim",
    category=CATEGORY_STRING,
    description=_("Returns the string {arg:0} without spaces at the end of the string."),
    examples=[
        SimpleExample('RTRIM("Computer ") = "Computer"'),
    ],
)

FUNCTION_TRIM = FunctionDocRegistryItem(
    name="trim",
    category=CATEGORY_STRING,
    description=_("Returns the string {arg:0} without spaces at the beginning or end of the " "string."),
    examples=[
        SimpleExample('TRIM(" Computer ") = "Computer"'),
    ],
)

FUNCTION_SUBSTR = FunctionDocRegistryItem(
    name="substr",
    category=CATEGORY_STRING,
    description=_(
        "Returns the substring {arg:0} starting from the index {arg:1}.\n"
        "\n"
        "If an additional argument `{argn:2}` is specified, a substring of the "
        "specified length is returned."
    ),
    examples=[
        SimpleExample('SUBSTR("Computer", 3) = "mputer"'),
    ],
)

_REGEXP_NOTES = [
    Note(_("See the documentation of the data source to clarify the regular expression syntax.")),
]

FUNCTION_REGEXP_EXTRACT = FunctionDocRegistryItem(
    name="regexp_extract",
    category=CATEGORY_STRING,
    description=_("Returns the substring {arg:0} that matches the regular expression pattern " "{arg:1}."),
    notes=_REGEXP_NOTES,
    examples=[
        SimpleExample('REGEXP_EXTRACT("RU 912873", "[A-Z]+\\s+(\\d+)") = "912873"'),
    ],
)

FUNCTION_REGEXP_EXTRACT_NTH = FunctionDocRegistryItem(
    name="regexp_extract_nth",
    category=CATEGORY_STRING,
    description=_(
        "Returns a substring {arg:0} that matches the regular expression pattern "
        "{arg:1} starting from the specified index."
    ),
    notes=_REGEXP_NOTES,
    examples=[
        SimpleExample('REGEXP_EXTRACT_NTH("RU 912 EN 873", "[A-Z]+\\s+(\\d+)", 2) = "873"'),
    ],
)

FUNCTION_REGEXP_MATCH = FunctionDocRegistryItem(
    name="regexp_match",
    category=CATEGORY_STRING,
    description=_(
        "Returns 'TRUE' if the string {arg:0} has a substring that matches the " "regular expression pattern {arg:1}."
    ),
    notes=_REGEXP_NOTES,
    examples=[
        SimpleExample('REGEXP_MATCH("RU 912873","\\w\\s\\d") = TRUE'),
    ],
)

FUNCTION_REGEXP_REPLACE = FunctionDocRegistryItem(
    name="regexp_replace",
    category=CATEGORY_STRING,
    description=_(
        "Searches for a substring in the string {arg:0} using the regular expression "
        "pattern {arg:1} and replaces it with the string {arg:2}.\n"
        "\n"
        "If the substring is not found, the string is not changed."
    ),
    notes=_REGEXP_NOTES,
    examples=[
        SimpleExample('REGEXP_REPLACE("123 456", "\\s", "-") = "123-456"'),
    ],
)

FUNCTION_REPLACE = FunctionDocRegistryItem(
    name="replace",
    category=CATEGORY_STRING,
    impl_selector=ArgAwareImplementationSelector(
        exp_arg_types={
            0: {DataType.STRING},
        },
    ),
    naming_provider=CategoryPostfixNamingProvider(),
    description=_(
        "Searches for the substring {arg:1} in the string {arg:0} and replaces it "
        "with the string {arg:2}.\n"
        "\n"
        "If the substring is not found, the string is not changed."
    ),
    examples=[
        SimpleExample('REPLACE("350 RUB", "RUB", "USD") = "350 USD"'),
    ],
)

FUNCTION_SPACE = FunctionDocRegistryItem(
    name="space",
    category=CATEGORY_STRING,
    description=_("Returns a string with the specified number of spaces."),
    examples=[
        SimpleExample('SPACE(5) = "     "'),
    ],
)

FUNCTION_SPLIT = FunctionDocRegistryItem(
    name="split",
    category=CATEGORY_STRING,
    description=_(
        "Returns a substring from {arg:0} using the {arg:1} delimiter character to "
        "divide the string into a sequence of {arg:2} parts. Delimiter is a comma by "
        "default. If {arg:2} is not passed, an array is returned (only for "
        "{dialects:CLICKHOUSE|POSTGRESQL} sources)"
    ),
    examples=[
        SimpleExample('SPLIT("192.168.0.1", ".", 1) = "192"'),
        SimpleExample('SPLIT("192.168.0.1", ".", -1) = "1"'),
    ],
)

FUNCTION_UTF8 = FunctionDocRegistryItem(
    name="utf8",
    category=CATEGORY_STRING,
    description=_("Converts the {arg:0} string encoding to `UTF8`."),
    examples=[
        SimpleExample('UTF8([Name], "CP-1251")'),
    ],
)

FUNCTIONS_STRING = [
    FUNCTION_ASCII,
    FUNCTION_CHAR,
    FUNCTION_CONCAT,
    FUNCTION_CONTAINS,
    FUNCTION_ICONTAINS,
    FUNCTION_ENDSWITH,
    FUNCTION_IENDSWITH,
    FUNCTION_STARTSWITH,
    FUNCTION_ISTARTSWITH,
    FUNCTION_FIND,
    FUNCTION_LEFT,
    FUNCTION_RIGHT,
    FUNCTION_LEN,
    FUNCTION_LOWER,
    FUNCTION_UPPER,
    FUNCTION_LTRIM,
    FUNCTION_RTRIM,
    FUNCTION_TRIM,
    FUNCTION_SUBSTR,
    FUNCTION_REGEXP_EXTRACT,
    FUNCTION_REGEXP_EXTRACT_NTH,
    FUNCTION_REGEXP_MATCH,
    FUNCTION_REGEXP_REPLACE,
    FUNCTION_REPLACE,
    FUNCTION_SPACE,
    FUNCTION_SPLIT,
    FUNCTION_UTF8,
]

lib/dl_formula_ref/dl_formula_ref/functions/string.py
datalens-tech/datalens-backend
from dl_formula_ref.categories.time_series import CATEGORY_TIME_SERIES
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import SimpleExample
from dl_formula_ref.registry.note import Note


_ = get_gettext()


_NOTES_TIME_SERIES = [
    Note(_("The first argument must be a measure (aggregated expression), otherwise an " "error will be raised.")),
]

FUNCTION_AGO = FunctionDocRegistryItem(
    name="ago",
    category=CATEGORY_TIME_SERIES,
    description=_(
        "Re-evaluate {arg:0} for a date/time with a given offset.\n"
        "The {arg:1} argument is the dimension along which the offset is made.\n"
        "The {arg:3} argument is an integer. It can be negative.\n"
        "The {arg:2} argument takes the following values:\n"
        '- `"year"`;\n'
        '- `"month"`;\n'
        '- `"day"`;\n'
        '- `"hour"`;\n'
        '- `"minute"`;\n'
        '- `"second"`.\n'
        "\n"
        "Can also be used as `AGO( measure, date_dimension, number )`. In this case, "
        "the third argument is interpreted as the number of days.\n"
        "\n"
        "See also {ref:AT_DATE}, {ref:LAG}."
    ),
    notes=_NOTES_TIME_SERIES,
    examples=[
        SimpleExample('AGO(SUM([Sales]), [Order Date], "month", 3)'),
        SimpleExample('AGO(SUM([Sales]), [Order Date], "year")'),
        SimpleExample("AGO(SUM([Sales]), [Order Date], 1)"),
    ],
)

FUNCTION_AT_DATE = FunctionDocRegistryItem(
    name="at_date",
    category=CATEGORY_TIME_SERIES,
    description=_(
        "Re-evaluate {arg:0} for a date/time specified by {arg:2}.\n"
        "The {arg:1} argument is the dimension along which the offset is made.\n"
        "\n"
        "See also {ref:AGO}, {ref:LAG}."
    ),
    notes=_NOTES_TIME_SERIES,
    examples=[
        SimpleExample("AT_DATE(SUM([Sales]), [Order Date], #2019-01-01#)"),
        SimpleExample('AT_DATE(SUM([Sales]), [Order Date], DATETRUNC([Order Date], "month"))'),
    ],
)

FUNCTIONS_TIME_SERIES = [
    FUNCTION_AGO,
    FUNCTION_AT_DATE,
]

lib/dl_formula_ref/dl_formula_ref/functions/time_series.py
datalens-tech/datalens-backend
from typing import MutableMapping

import attr

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo
from dl_formula.definitions.functions_type import (
    DataTypeSpec,
    WhitelistTypeSpec,
)
from dl_formula_ref.categories.type_conversion import CATEGORY_TYPE_CONVERSION
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.aliased_res import (
    AliasedResource,
    AliasedResourceRegistryBase,
    AliasedTableResource,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import SimpleExample
from dl_formula_ref.registry.note import (
    Note,
    NoteLevel,
)
from dl_i18n.localizer_base import Translatable


_ = get_gettext()


FUNCTION_DATE = FunctionDocRegistryItem(
    name="date",
    category=CATEGORY_TYPE_CONVERSION,
    description=_(
        "Converts the {arg:0} expression to date format.\n"
        "\n"
        "The date must be in the format `YYYY-MM-DD`.\n"
        "\n"
        "If {arg:0} is a number, then the {arg:1} option can be used to convert the "
        "date to the specified time zone."
    ),
    notes=[
        Note(
            _("Argument {arg:1} is available only for {dialects:CLICKHOUSE} " "sources."),
        ),
        Note(
            _(
                "For {dialects:CLICKHOUSE} data sources, numeric {arg:0} values less than or "
                "equal to `65535` are interpreted as the number of days (not seconds, like in "
                "all other cases) since January 1st 1970. This is the result of the behavior "
                "of available {dialects:CLICKHOUSE} functions.\n"
                "\n"
                "One way to surpass this is to use the following formula: "
                "`DATE(DATETIME([value]))`. The result is more consistent, but is likely to "
                "be much slower."
            ),
            level=NoteLevel.warning,
        ),
    ],
    examples=[
        SimpleExample('DATE("2019-01-23") = #2019-01-23#'),
    ],
)

FUNCTION_DATETIME = FunctionDocRegistryItem(
    name="datetime",
    category=CATEGORY_TYPE_CONVERSION,
    description=_(
        "Converts the {arg:0} expression to date and time format. When converting "
        "`Date` to `DateTime`, the time is set to '00:00:00'.\n"
        "The date must be in the format `YYYY-MM-DDThh:mm:ss` or `YYYY-MM-DD hh:mm:"
        "ss`.\n"
        "\n"
        "The date and time can be converted to the "
        "specified time zone when the {arg:1} option is available."
    ),
    notes=[
        Note(
            _("Argument {arg:1} is available only for {dialects:CLICKHOUSE} " "sources."),
        ),
    ],
    examples=[
        SimpleExample('DATETIME("2019-01-23 15:07:47") = #2019-01-23 15:07:47#'),
    ],
)


FUNCTION_FLOAT = FunctionDocRegistryItem(
    name="float",
    category=CATEGORY_TYPE_CONVERSION,
    description=_(
        "Converts the {arg:0} expression to fractional number format according to the "
        "following rules:\n"
        "\n"
        "{table: conversion}"
    ),
    resources=SimpleAliasedResourceRegistry(
        resources={
            "conversion": AliasedTableResource(
                table_body=[
                    [_("Type"), _("Value")],
                    ["{type:INTEGER|FLOAT}", _("Original value.")],
                    [
                        "{type:DATE|DATETIME|GENERICDATETIME}",
                        _(
                            "{link: unix_ts: Unix time} corresponding to the date and time. If the value "
                            "contains time zone data, it's used in the calculation. If the time zone is "
                            "unknown, the time is set in UTC."
                        ),
                    ],
                    ["{type:STRING}", _("A number from a decimal string.")],
                    ["{type:BOOLEAN}", _("`TRUE` — `1.0`, `FALSE` — `0.0`.")],
                ]
            ),
        }
    ),
    examples=[
        SimpleExample("FLOAT(7) = 7.0"),
        SimpleExample('FLOAT("34.567") = 34.567'),
        SimpleExample("FLOAT(TRUE) = 1.0"),
    ],
)

FUNCTION_INT = FunctionDocRegistryItem(
    name="int",
    category=CATEGORY_TYPE_CONVERSION,
    description=_(
        "Converts the {arg:0} expression to integer format according to the following "
        "rules:\n"
        "\n"
        "{table: conversion}"
    ),
    resources=SimpleAliasedResourceRegistry(
        resources={
            "conversion": AliasedTableResource(
                table_body=[
                    [_("Type"), _("Value")],
                    ["{type:INTEGER}", _("Original value.")],
                    ["{type:FLOAT}", _("Integer part of the number (rounded down).")],
                    [
                        "{type:DATE|DATETIME|GENERICDATETIME}",
                        _(
                            "{link: unix_ts: Unix time} corresponding to the date and time. If the value "
                            "contains time zone data, it's used in the calculation. If the time zone is "
                            "unknown, the time is set in UTC."
                        ),
                    ],
                    ["{type:STRING}", _("A number from a decimal string.")],
                    ["{type:BOOLEAN}", _("`TRUE` — `1`, `FALSE` — `0`.")],
                ]
            ),
        }
    ),
    examples=[
        SimpleExample("INT(7.7) = 7"),
        SimpleExample('INT("365") = 365'),
        SimpleExample("INT(TRUE) = 1"),
    ],
)

FUNCTION_BOOL = FunctionDocRegistryItem(
    name="bool",
    category=CATEGORY_TYPE_CONVERSION,
    description=_(
        "Converts the {arg:0} expression to Boolean type according to the following "
        "rules:\n"
        "\n"
        "{table:conversion}"
    ),
    resources=SimpleAliasedResourceRegistry(
        resources={
            "conversion": AliasedTableResource(
                table_body=[
                    [_("Type"), "`FALSE`", "`TRUE`"],
                    ["{type:INTEGER|FLOAT}", "`0`, `0.0`", _("All others")],
                    ["{type:STRING}", _('Empty string (`""`)'), _("All others")],
                    ["{type:BOOLEAN}", "`FALSE`", "`TRUE`"],
                    ["{type:DATE|DATETIME|GENERICDATETIME}", "-", "`TRUE`"],
                ]
            ),
        }
    ),
    examples=[
        SimpleExample("BOOL(0) = FALSE"),
        SimpleExample("BOOL(#2019-04-04#) = TRUE"),
        SimpleExample('BOOL("Lorem ipsum") = TRUE'),
    ],
)

FUNCTION_STR = FunctionDocRegistryItem(
    name="str",
    category=CATEGORY_TYPE_CONVERSION,
    description=_("Converts the {arg:0} expression to string type."),
    examples=[
        SimpleExample('STR(77) = "77"'),
        SimpleExample('STR(1 != 1) = "False"'),
        SimpleExample('STR(#2019-01-01#) = "2019-01-01"'),
        SimpleExample("STR(ARRAY(1, 2, 3, NULL)) = {1,2,3,NULL}  PostgreSQL"),
        SimpleExample("STR(ARRAY(1, 2, 3, NULL)) = [1,2,3,NULL]  ClickHouse"),
        SimpleExample("STR(ARRAY('a', 'b', '', NULL)) = {a,b," ",NULL}  PostgreSQL"),
        SimpleExample("STR(ARRAY('a', 'b', '', NULL)) = ['a','b','',NULL]  ClickHouse"),
    ],
    notes=[
        Note(
            _(
                "If an array is passed (only for {dialects:CLICKHOUSE|POSTGRESQL} sources), the conversion is "
                "performed by a function in the source database and results may vary for different data sources. "
                "For consistent results use {ref:ARR_STR}."
            )
        ),
    ],
)

FUNCTION_DATETIME_PARSE = FunctionDocRegistryItem(
    name="datetime_parse",
    category=CATEGORY_TYPE_CONVERSION,
    description=_(
        "Converts the {arg:0} expression to date and time format. Unlike {ref: "
        "DATETIME}, it supports multiple formats."
    ),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            'DATETIME_PARSE("2019-01-02 03:04:05") = #2019-01-02 03:04:05#',
            'DATETIME_PARSE("20190102030405") = #2019-01-02 03:04:05#',
            'DATETIME_PARSE("20190102 030405") = #2019-01-02 03:04:05#',
            'DATETIME_PARSE("2019.01.02 03:04:05") = #2019-01-02 03:04:05#',
            'DATETIME_PARSE("02/01/2019 03:04:05") = #2019-01-02 03:04:05#',
            'DATETIME_PARSE("2019-01-02 03:04") = #2019-01-02 03:04:00#',
            'DATETIME_PARSE("2019-01-02 030405") = #2019-01-02 03:04:05#',
            'DATETIME_PARSE("2019 Jan 02 03:04:05") = #2019-01-02 03:04:05#',
        )
    ],
)

FUNCTION_DATE_PARSE = FunctionDocRegistryItem(
    name="date_parse",
    category=CATEGORY_TYPE_CONVERSION,
    description=_(
        "Converts the {arg:0} expression to date format. Unlike {ref: DATE}, it " "supports multiple formats."
    ),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            'DATE_PARSE("2019-01-02 03:04:05") = #2019-01-02#',
            'DATE_PARSE("2019-01-02") = #2019-01-02#',
            'DATE_PARSE("20190102") = #2019-01-02#',
            'DATE_PARSE("2019.01.02") = #2019-01-02#',
            'DATE_PARSE("02/01/2019") = #2019-01-02#',
            'DATE_PARSE("02/01/19") = #2019-01-02#',
            'DATE_PARSE("2019 Jan 02") = #2019-01-02#',
            'DATE_PARSE("2019 Jan") = #2019-01-01#',
            'DATE_PARSE("201901") = #2019-01-01#',
            'DATE_PARSE("2019") = #2019-01-01#',
        )
    ],
)

FUNCTION_GEOPOINT = FunctionDocRegistryItem(
    name="geopoint",
    category=CATEGORY_TYPE_CONVERSION,
    description=_(
        "Generates a Geopoint type value. For the input, it accepts a string, a "
        '"geopoint" type value, or coordinates — latitude {arg:0} and longitude '
        "{arg:1}. If a single string is input, it must contain a list of two numbers "
        "(latitude and longitude) in JSON syntax."
    ),
    examples=[
        SimpleExample('GEOPOINT("[55.75222,37.61556]") = "[55.75222,37.61556]"'),
        SimpleExample('GEOPOINT(55.75222, 37.61556) = "[55.75222,37.61556]"'),
    ],
)

FUNCTION_GEOPOLYGON = FunctionDocRegistryItem(
    name="geopolygon",
    category=CATEGORY_TYPE_CONVERSION,
    description=_("Converts the {arg:0} expression to geopolygon format."),
    examples=[
        SimpleExample(
            'GEOPOLYGON("[[[55.79421,37.65046],[55.79594,37.6513],[55.79642,37.65133],'
            '[55.7969, 37.65114],[55.79783, 37.65098],[55.78871,37.75101]]]")'
        ),
        SimpleExample(
            'GEOPOLYGON("[[[55.75,37.52],[55.75,37.68],[55.65,37.60]],' '[[55.79,37.60],[55.76,37.57],[55.76,37.63]]]")'
        ),
    ],
)


def _make_type_macro_from_dtype_spec(data_type_spec: DataTypeSpec) -> str:
    if isinstance(data_type_spec, DataType):
        return f"{{type:{data_type_spec.name}}}"
    elif isinstance(data_type_spec, tuple):
        return f'{{type:{"|".join([sub_data_type_spec.name for sub_data_type_spec in data_type_spec])}}}'
    else:
        raise TypeError(type(data_type_spec))


def _get_comment_for_type(dialect: DialectCombo, native_type_name: str) -> str | Translatable:
    return _DB_CAST_TYPE_COMMENTS.get((dialect, native_type_name), "")


@attr.s
class DbCastExtension:
    type_whitelists: dict[DialectCombo, dict[DataType, list[WhitelistTypeSpec]]] = attr.ib(kw_only=True, factory=dict)
    type_comments: dict[tuple[DialectCombo, str], str | Translatable] = attr.ib(kw_only=True, factory=dict)


_DB_CAST_TYPE_COMMENTS: dict[tuple[DialectCombo, str], str | Translatable] = {}
_DB_CAST_WHITELIST: dict[DialectCombo, dict[DataType, list[WhitelistTypeSpec]]] = {}


def register_db_cast_extension(extension: DbCastExtension) -> None:
    _DB_CAST_WHITELIST.update(extension.type_whitelists)
    _DB_CAST_TYPE_COMMENTS.update(extension.type_comments)


class DbCastWhiteListAliasedResourceRegistry(AliasedResourceRegistryBase):
    def get_resources(self) -> MutableMapping[str, AliasedResource]:
        resources = {
            "supported_native_types": AliasedTableResource(
                table_body=[
                    [
                        _("Data source"),
                        _("Data type"),
                        _("Native data type"),
                        _("Parameters for native type"),
                        _("Comment"),
                    ],
                    *[
                        [
                            f"{{dialects:{dialect.common_name.name}}}",
                            f"{{type:{data_type.name}}}",
                            f"`{spec.name}`",
                            f'{", ".join([_make_type_macro_from_dtype_spec(arg_type) for arg_type in spec.arg_types])}',
                            _get_comment_for_type(dialect=dialect, native_type_name=spec.name),
                        ]
                        for dialect, by_data_type in sorted(
                            _DB_CAST_WHITELIST.items(), key=lambda item: item[0].common_name.value
                        )
                        for data_type, spec_list in sorted(by_data_type.items(), key=lambda item: item[0].name)
                        for spec in spec_list
                    ],
                ]
            ),
        }
        return resources


FUNCTION_DB_CAST = FunctionDocRegistryItem(
    name="db_cast",
    category=CATEGORY_TYPE_CONVERSION,
    description=_(
        "Converts the {arg:0} expression to database's native type {arg:1}.\n"
        "\n"
        "The following type casts are supported:\n"
        "\n"
        "{table:supported_native_types}\n"
    ),
    # Using a custom resource registry here
    # because it needs to read from global maps
    # that are filled from plugins after import time
    resources=DbCastWhiteListAliasedResourceRegistry(),
    examples=[
        SimpleExample('DB_CAST([float_value], "Decimal", 10, 5)'),
        SimpleExample('DB_CAST([float_value], "double precision")'),
        SimpleExample('DB_CAST([float_value], "numeric", 10, 5)'),
    ],
)

FUNCTIONS_TYPE_CONVERSION = [
    FUNCTION_DATE,
    FUNCTION_DATETIME,
    FUNCTION_FLOAT,
    FUNCTION_INT,
    FUNCTION_BOOL,
    FUNCTION_STR,
    FUNCTION_DATETIME_PARSE,
    FUNCTION_DATE_PARSE,
    FUNCTION_GEOPOINT,
    FUNCTION_GEOPOLYGON,
    FUNCTION_DB_CAST,
]

lib/dl_formula_ref/dl_formula_ref/functions/type_conversion.py
datalens-tech/datalens-backend
from typing import List

from dl_formula.core.datatype import DataType
from dl_formula_ref.categories.window import CATEGORY_WINDOW
from dl_formula_ref.examples.config import (
    ExampleConfig,
    ExampleSource,
)
from dl_formula_ref.localization import get_gettext
from dl_formula_ref.registry.aliased_res import (
    AliasedTableResource,
    AliasedTextResource,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.example import (
    DataExample,
    SimpleExample,
)
from dl_formula_ref.registry.naming import CategoryPostfixNamingProvider
from dl_formula_ref.registry.note import (
    Note,
    NoteLevel,
)


_ = get_gettext()

_SOURCE_SALES_DATA_1 = ExampleSource(
    columns=[
        ("Date", DataType.STRING),
        ("City", DataType.STRING),
        ("Category", DataType.STRING),
        ("Orders", DataType.INTEGER),
        ("Profit", DataType.FLOAT),
    ],
    data=[
        ["2019-03-01", "London", "Office Supplies", 8, 120.8],
        ["2019-03-04", "London", "Office Supplies", 2, 100.0],
        ["2019-03-05", "London", "Furniture", 1, 750.0],
        ["2019-03-02", "Moscow", "Furniture", 2, 1250.5],
        ["2019-03-03", "Moscow", "Office Supplies", 4, 85.0],
        ["2019-03-01", "San Francisco", "Office Supplies", 23, 723.0],
        ["2019-03-01", "San Francisco", "Furniture", 1, 1000.0],
        ["2019-03-03", "San Francisco", "Furniture", 4, 4000.0],
        ["2019-03-02", "Detroit", "Furniture", 5, 3700.0],
        ["2019-03-04", "Detroit", "Office Supplies", 25, 1200.0],
        ["2019-03-04", "Detroit", "Furniture", 2, 3500.00],
    ],
)


def _make_standard_window_examples(func: str) -> List[DataExample]:
    func = func.upper()
    examples = [
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with grouping"),
                source=_SOURCE_SALES_DATA_1,
                group_by=["[City]", "[Category]"],
                order_by=["[City]", "[Category]"],
                show_source_table=True,
                formula_fields=[("City", "[City]"), ("Category", "[Category]"), ("Order Sum", "SUM([Orders])")],
                additional_transformations=[
                    [
                        ("City", "[City]"),
                        ("Category", "[Category]"),
                        ("Order Sum", "[Order Sum]"),
                        (f"{func} 1", f"{func}([Order Sum] TOTAL)"),
                        (f"{func} 2", f"{func}([Order Sum] WITHIN [City])"),
                        (f"{func} 3", f"{func}([Order Sum] WITHIN [Category])"),
                    ],
                ],
                override_formula_fields=[
                    ("City", "[City]"),
                    ("Category", "[Category]"),
                    ("Order Sum", "SUM([Orders])"),
                    (f"{func} 1", f"{func}(SUM([Orders]) TOTAL)"),
                    (f"{func} 2", f"{func}(SUM([Orders]) WITHIN [City])"),
                    (f"{func} 3", f"{func}(SUM([Orders]) AMONG [City])"),
                ],
            ),
        ),
    ]
    return examples


def _make_rank_examples(func: str) -> List[DataExample]:
    func = func.upper()
    examples = [
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with two arguments"),
                source=_SOURCE_SALES_DATA_1,
                group_by=["[City]"],
                order_by=["[City]"],
                show_source_table=True,
                formula_fields=[("City", "[City]"), ("Order Sum", "SUM([Orders])")],
                additional_transformations=[
                    [
                        ("City", "[City]"),
                        ("Order Sum", "[Order Sum]"),
                        (f"{func} 1", f'{func}([Order Sum], "desc")'),
                        (f"{func} 2", f'{func}([Order Sum], "asc")'),
                    ],
                ],
                override_formula_fields=[
                    ("City", "[City]"),
                    ("Order Sum", "SUM([Orders])"),
                    (f"{func} 1", f'{func}(SUM([Orders]), "desc")'),
                    (f"{func} 2", f'{func}(SUM([Orders]), "asc")'),
                ],
            ),
        ),
    ]
    return examples


_RANK_RESOURCES = SimpleAliasedResourceRegistry(
    resources={
        "rank_direction_description": AliasedTextResource(
            body=_(
                'If {arg:1} is `"desc"` or omitted, then ranking is done from greatest to '
                'least, if `"asc"`, then from least to greatest.'
            ),
        ),
    }
)

FUNCTION_RANK = FunctionDocRegistryItem(
    name="rank",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the rank of the current row if ordered by the given argument. Rows "
        "corresponding to the same value used for sorting have the same rank. If the "
        "first two rows both have rank of `1`, then the next row (if it features a "
        "different value) will have rank `3`, so, in effect, it is rank with gaps.\n"
        "\n"
        "{text:rank_direction_description}\n"
        "\n"
        "See also {ref: RANK_DENSE}, {ref: RANK_UNIQUE}, {ref:RANK_PERCENTILE}."
    ),
    resources=_RANK_RESOURCES,
    examples=[
        *_make_rank_examples("rank"),
        *_make_standard_window_examples("rank"),
    ],
)

FUNCTION_RANK_DENSE = FunctionDocRegistryItem(
    name="rank_dense",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the rank of the current row if ordered by the given argument. Rows "
        "corresponding to the same value used for sorting have the same rank. If the "
        "first two rows both have rank of `1`, then the next row (if it features a "
        "different value) will have rank `2`, (rank without gaps).\n"
        "\n"
        "{text:rank_direction_description}\n"
        "\n"
        "See also {ref: RANK}, {ref: RANK_UNIQUE}, {ref:RANK_PERCENTILE}."
    ),
    resources=_RANK_RESOURCES,
    examples=[
        *_make_rank_examples("rank_dense"),
        *_make_standard_window_examples("rank_dense"),
    ],
)

FUNCTION_RANK_UNIQUE = FunctionDocRegistryItem(
    name="rank_unique",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the rank of the current row if ordered by the given argument. Rows "
        "corresponding to the same value have different rank values. This means that "
        "rank values are sequential and different for all rows, always increasing by "
        "`1` for the next row.\n"
        "\n"
        "{text:rank_direction_description}\n"
        "\n"
        "See also {ref: RANK}, {ref: RANK_DENSE}, {ref:RANK_PERCENTILE}."
    ),
    resources=_RANK_RESOURCES,
    examples=[
        *_make_rank_examples("rank_unique"),
        *_make_standard_window_examples("rank_unique"),
    ],
)

FUNCTION_RANK_PERCENTILE = FunctionDocRegistryItem(
    name="rank_percentile",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the relative rank (from `0` to `1`) of the current row if ordered by the "
        "given argument. Calculated as `(RANK(...) - 1) / (row count) `.\n"
        "\n"
        "{text:rank_direction_description}\n"
        "\n"
        "See also {ref: RANK}, {ref: RANK_DENSE}, {ref: RANK_UNIQUE}."
    ),
    resources=_RANK_RESOURCES,
    examples=[
        *_make_rank_examples("rank_percentile"),
        *_make_standard_window_examples("rank_percentile"),
    ],
)

FUNCTION_SUM_WINDOW = FunctionDocRegistryItem(
    name="sum",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="sum_window",
    ),
    category=CATEGORY_WINDOW,
    description=_("Returns the sum of all expression values. Applicable to numeric data types only."),
    examples=[*_make_standard_window_examples("sum")],
)

FUNCTION_MIN_WINDOW = FunctionDocRegistryItem(
    name="min",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="min_window",
    ),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the minimum value.\n\n"
        "If {arg:0}:\n"
        "- number — Returns the smallest number.\n"
        "- date — Returns the earliest date.\n"
        "- string — Returns the first value in the alphabetic order.\n"
    ),
    examples=[*_make_standard_window_examples("min")],
)

FUNCTION_MAX_WINDOW = FunctionDocRegistryItem(
    name="max",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="max_window",
    ),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the maximum value.\n\n"
        "If {arg:0}:\n"
        "- number — Returns the largest number.\n"
        "- date — Returns the latest date.\n"
        "- string — Returns the last value in the alphabetic order.\n"
    ),
    examples=[*_make_standard_window_examples("max")],
)

FUNCTION_COUNT_WINDOW = FunctionDocRegistryItem(
    name="count",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="count_window",
    ),
    category=CATEGORY_WINDOW,
    description=_("Returns the number of items in the specified window."),
    examples=[*_make_standard_window_examples("count")],
)

FUNCTION_AVG_WINDOW = FunctionDocRegistryItem(
    name="avg",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="avg_window",
    ),
    category=CATEGORY_WINDOW,
    description=_("Returns the average of all values. Applicable to numeric data types."),
    examples=[*_make_standard_window_examples("avg")],
)

FUNCTION_SUM_IF_WINDOW = FunctionDocRegistryItem(
    name="sum_if",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="sum_if_window",
    ),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the sum of all the expression values that meet the {arg:1} condition. "
        "Applicable to numeric data types only."
    ),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            "SUM_IF([Profit], [Category] = 'Office Supplies' TOTAL)",  # translates into:
            #    SUM("Profit") FILTER (WHERE "Category" = 'Office Supplies')
            #    OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)
            "SUM_IF([Profit], [Category] = 'Office Supplies' WITHIN [Date])",  # translates into:
            #    SUM("Profit") FILTER (WHERE "Category" = 'Office Supplies')
            #    OVER (
            #        PARTITION BY "Date"
            #        ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
            #    )
            "SUM_IF([Profit], [Category] = 'Office Supplies' AMONG [Date])",  # translates into:
            #    SUM("Profit") FILTER (WHERE "Category" = 'Office Supplies')
            #    OVER (
            #        PARTITION BY <all dimensions except "Date">
            #        ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
            #    )
        )
    ],
)

FUNCTION_COUNT_IF_WINDOW = FunctionDocRegistryItem(
    name="count_if",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="count_if_window",
    ),
    category=CATEGORY_WINDOW,
    description=_("Returns the number of items in the specified window meeting the {arg:0} condition."),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            "COUNT_IF([Profit], [Category] = 'Office Supplies' TOTAL)",
            "COUNT_IF([Profit], [Category] = 'Office Supplies' WITHIN [Date])",
            "COUNT_IF([Profit], [Category] = 'Office Supplies' AMONG [Date])",
        )
    ],
)

FUNCTION_AVG_IF_WINDOW = FunctionDocRegistryItem(
    name="avg_if",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(
        internal_name="avg_if_window",
    ),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the average of all values that meet the {arg:1} condition. "
        "If the values don't exist, it returns `NULL`. Applicable to numeric data types only."
    ),
    examples=[
        SimpleExample(example_str)
        for example_str in (
            "AVG_IF([Profit], [Category] = 'Office Supplies' TOTAL)",
            "AVG_IF([Profit], [Category] = 'Office Supplies' WITHIN [Date])",
            "AVG_IF([Profit], [Category] = 'Office Supplies' AMONG [Date])",
        )
    ],
)


def _make_simple_order_by_examples(func: str) -> List[DataExample]:
    func = func.upper()
    examples = [
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with ORDER BY"),
                source=_SOURCE_SALES_DATA_1,
                group_by=["[City]"],
                order_by=["[City]"],
                show_source_table=True,
                formula_fields=[("City", "[City]"), ("Order Sum", "SUM([Orders])")],
                additional_transformations=[
                    [
                        ("City", "[City]"),
                        ("Order Sum", "[Order Sum]"),
                        (f"{func} 1", f"{func}([Order Sum] ORDER BY [City] DESC)"),
                        (f"{func} 2", f"{func}([Order Sum] ORDER BY [Order Sum])"),
                    ],
                ],
                override_formula_fields=[
                    ("City", "[City]"),
                    ("Order Sum", "SUM([Orders])"),
                    (f"{func} 1", f"{func}(SUM([Orders]) ORDER BY [City] DESC)"),
                    (f"{func} 2", f"{func}(SUM([Orders]) ORDER BY [Order Sum])"),
                ],
            ),
        ),
    ]
    return examples


_ORDERED_WFUNC_NOTES = [
    Note(
        level=NoteLevel.warning,
        text=_(
            "The sorting order is based on the fields listed in the sorting section of "
            "the chart and in the `ORDER BY` clause. First, `ORDER BY` fields are used, "
            "and then they are complemented by the fields from the chart."
        ),
    ),
]


def _make_rfunc_examples(func: str) -> List[DataExample]:
    func = func.upper()
    examples = [
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with ORDER BY"),
                source=_SOURCE_SALES_DATA_1,
                group_by=["[City]"],
                order_by=["[City]"],
                show_source_table=True,
                formula_fields=[("City", "[City]"), ("Order Sum", "SUM([Orders])")],
                additional_transformations=[
                    [
                        ("City", "[City]"),
                        ("Order Sum", "[Order Sum]"),
                        (f"{func} 1", f'{func}([Order Sum], "desc")'),
                        (f"{func} 2", f'{func}([Order Sum], "asc" ORDER BY [City] DESC)'),
                        (f"{func} 3", f"{func}([Order Sum] ORDER BY [Order Sum])"),
                    ],
                ],
                override_formula_fields=[
                    ("City", "[City]"),
                    ("Order Sum", "SUM([Orders])"),
                    (f"{func} 1", f'{func}(SUM([Orders]), "desc")'),
                    (f"{func} 2", f'{func}(SUM([Orders]), "asc" ORDER BY [City] DESC)'),
                    (f"{func} 3", f"{func}(SUM([Orders]) ORDER BY [Order Sum])"),
                ],
            ),
        ),
    ]
    return examples


_RFUNC_RESOURCES = SimpleAliasedResourceRegistry(
    resources={
        "rwindow_description_table": AliasedTableResource(
            table_body=[
                ["{arg:1}", _("Window")],
                ['`"asc"`', _("Starts from the first row and ends at the current row.")],
                ['`"desc"`', _("Starts from the current row and ends at the last row.")],
            ]
        ),
        "rwindow_description": AliasedTextResource(
            body=_("\n" "{table:rwindow_description_table}\n" "\n" 'By default `"asc"` is used.\n')
        ),
    }
)

FUNCTION_RSUM = FunctionDocRegistryItem(
    name="rsum",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the sum of all values in a growing (or shrinking) window defined by "
        "the sort order and the value of {arg:1}:\n"
        "{text:rwindow_description}\n"
        "\n"
        "Window functions with a similar behavior: {ref:RCOUNT}, {ref:RMIN}, {ref:RMAX}, {ref:RAVG}.\n"
        "\n"
        "See also {ref:aggregation/SUM}, {ref:MSUM}."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    resources=_RFUNC_RESOURCES,
    examples=[
        *_make_standard_window_examples("rsum"),
        *_make_rfunc_examples("rsum"),
    ],
)

FUNCTION_RCOUNT = FunctionDocRegistryItem(
    name="rcount",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the count of all values in a growing (or shrinking) window defined "
        "by the sort order and the value of {arg:1}:\n"
        "{text:rwindow_description}\n"
        "\n"
        "Window functions with a similar behavior: {ref:RSUM}, {ref:RMIN}, {ref:RMAX}, {ref:RAVG}.\n"
        "\n"
        "See also {ref:aggregation/COUNT}, {ref:MCOUNT}."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    resources=_RFUNC_RESOURCES,
    examples=[
        *_make_standard_window_examples("rcount"),
        *_make_rfunc_examples("rcount"),
    ],
)

FUNCTION_RMIN = FunctionDocRegistryItem(
    name="rmin",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the minimum of all values in a growing (or shrinking) window defined "
        "by the sort order and the value of {arg:1}:\n"
        "{text:rwindow_description}\n"
        "\n"
        "Window functions with a similar behavior: {ref:RSUM}, {ref:RCOUNT}, {ref:RMAX}, {ref:RAVG}.\n"
        "\n"
        "See also {ref:aggregation/MIN}, {ref:MMIN}."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    resources=_RFUNC_RESOURCES,
    examples=[
        *_make_standard_window_examples("rmin"),
        *_make_rfunc_examples("rmin"),
    ],
)

FUNCTION_RMAX = FunctionDocRegistryItem(
    name="rmax",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the maximum of all values in a growing (or shrinking) window defined "
        "by the sort order and the value of {arg:1}:\n"
        "{text:rwindow_description}\n"
        "\n"
        "Window functions with a similar behavior: {ref:RSUM}, {ref:RCOUNT}, "
        "{ref:RMIN}, {ref:RAVG}.\n"
        "\n"
        "See also {ref:aggregation/MAX}, {ref:MMAX}."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    resources=_RFUNC_RESOURCES,
    examples=[
        *_make_standard_window_examples("rmax"),
        *_make_rfunc_examples("rmax"),
    ],
)

FUNCTION_RAVG = FunctionDocRegistryItem(
    name="ravg",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the average of all values in a growing (or shrinking) window defined "
        "by the sort order and the value of {arg:1}:\n"
        "{text:rwindow_description}\n"
        "\n"
        "Window functions with a similar behavior: {ref:RSUM}, {ref:RCOUNT}, {ref:RMIN}, {ref:RMAX}.\n"
        "\n"
        "See also {ref:aggregation/AVG}, {ref:MAVG}."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    resources=_RFUNC_RESOURCES,
    examples=[
        *_make_standard_window_examples("ravg"),
        *_make_rfunc_examples("ravg"),
    ],
)


def _make_mfunc_examples(func: str) -> List[DataExample]:
    func = func.upper()
    examples = [
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with two and three arguments"),
                source=_SOURCE_SALES_DATA_1,
                group_by=["[City]"],
                order_by=["[City]"],
                show_source_table=True,
                formula_fields=[("City", "[City]"), ("Order Sum", "SUM([Orders])")],
                additional_transformations=[
                    [
                        ("City", "[City]"),
                        ("Order Sum", "[Order Sum]"),
                        (f"{func} 1", f"{func}([Order Sum], 1)"),
                        (f"{func} 2", f"{func}([Order Sum], -2)"),
                        (f"{func} 3", f"{func}([Order Sum], 1, 1)"),
                    ],
                ],
                override_formula_fields=[
                    ("City", "[City]"),
                    ("Order Sum", "SUM([Orders])"),
                    (f"{func} 1", f"{func}(SUM([Orders]), 1)"),
                    (f"{func} 2", f"{func}(SUM([Orders]), -2)"),
                    (f"{func} 3", f"{func}(SUM([Orders]) 1, 1)"),
                ],
            ),
        ),
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with ORDER BY"),
                source=_SOURCE_SALES_DATA_1,
                group_by=["[City]"],
                order_by=["[City]"],
                show_source_table=True,
                formula_fields=[("City", "[City]"), ("Order Sum", "SUM([Orders])")],
                additional_transformations=[
                    [
                        ("City", "[City]"),
                        ("Order Sum", "[Order Sum]"),
                        (f"{func} 1", f"{func}([Order Sum], 1 ORDER BY [City] DESC)"),
                        (f"{func} 2", f"{func}([Order Sum], 1 ORDER BY [Order Sum])"),
                    ],
                ],
                override_formula_fields=[
                    ("City", "[City]"),
                    ("Order Sum", "SUM([Orders])"),
                    (f"{func} 1", f"{func}(SUM([Orders]), 1 ORDER BY [City] DESC)"),
                    (f"{func} 2", f"{func}(SUM([Orders]), 1 ORDER BY [Order Sum])"),
                ],
            ),
        ),
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with grouping"),
                source=_SOURCE_SALES_DATA_1,
                group_by=["[City]", "[Category]"],
                order_by=["[City]", "[Category]"],
                show_source_table=True,
                formula_fields=[("City", "[City]"), ("Category", "[Category]"), ("Order Sum", "SUM([Orders])")],
                additional_transformations=[
                    [
                        ("City", "[City]"),
                        ("Category", "[Category]"),
                        ("Order Sum", "[Order Sum]"),
                        (f"{func} 1", f"{func}([Order Sum], 1 TOTAL)"),
                        (f"{func} 2", f"{func}([Order Sum], 1 WITHIN [City])"),
                        (f"{func} 3", f"{func}([Order Sum], 1 WITHIN [Category])"),
                    ],
                ],
                override_formula_fields=[
                    ("City", "[City]"),
                    ("Category", "[Category]"),
                    ("Order Sum", "SUM([Orders])"),
                    (f"{func} 1", f"{func}(SUM([Orders]), 1 TOTAL)"),
                    (f"{func} 2", f"{func}(SUM([Orders]), 1 WITHIN [City])"),
                    (f"{func} 3", f"{func}(SUM([Orders]), 1 AMONG [City])"),
                ],
            ),
        ),
    ]
    return examples


_MFUNC_RESOURCES = SimpleAliasedResourceRegistry(
    resources={
        "mwindow_description_table": AliasedTableResource(
            table_body=[
                ["{arg:1}", "{arg:2}", _("Window")],
                [_("positive"), "-", _("The current row and {arg:1} preceding rows.")],
                [_("negative"), "-", _("The current row and -{arg:1} following rows.")],
                [
                    _("any sign"),
                    _("any sign"),
                    _("{arg:1} preceding rows, the current row and {arg:2} following rows."),
                ],
            ]
        ),
        "mwindow_description": AliasedTextResource(body=("\n" "{table:mwindow_description_table}\n" "\n")),
    }
)

FUNCTION_MSUM = FunctionDocRegistryItem(
    name="msum",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the moving sum of values in a fixed-size window defined by "
        "the sort order and arguments:\n"
        "{text:mwindow_description}\n"
        "Window functions with a similar behavior: {ref:MCOUNT}, {ref:MMIN}, {ref:MMAX}, {ref:MAVG}.\n"
        "\n"
        "See also {ref:aggregation/SUM}, {ref:RSUM}."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    resources=_MFUNC_RESOURCES,
    examples=[
        *_make_mfunc_examples("msum"),
    ],
)

FUNCTION_MCOUNT = FunctionDocRegistryItem(
    name="mcount",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the moving count of (non-`NULL`) values in a fixed-size window defined by "
        "the sort order and arguments:\n"
        "{text:mwindow_description}\n"
        "Window functions with a similar behavior: {ref:MSUM}, {ref:MMIN}, {ref:MMAX}, {ref:MAVG}.\n"
        "\n"
        "See also {ref:aggregation/COUNT}, {ref:RCOUNT}."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    resources=_MFUNC_RESOURCES,
    examples=[
        *_make_mfunc_examples("mcount"),
    ],
)

FUNCTION_MMIN = FunctionDocRegistryItem(
    name="mmin",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the moving minimum of values in a fixed-size window defined by "
        "the sort order and arguments:\n"
        "{text:mwindow_description}\n"
        "Window functions with a similar behavior: {ref:MSUM}, {ref:MCOUNT}, {ref:MMAX}, {ref:MAVG}.\n"
        "\n"
        "See also {ref:aggregation/MIN}, {ref:RMIN}."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    resources=_MFUNC_RESOURCES,
    examples=[
        *_make_mfunc_examples("mmin"),
    ],
)

FUNCTION_MMAX = FunctionDocRegistryItem(
    name="mmax",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the moving maximum of values in a fixed-size window defined by "
        "the sort order and arguments:\n"
        "{text:mwindow_description}\n"
        "Window functions with a similar behavior: {ref:MSUM}, {ref:MCOUNT}, {ref:MMIN}, {ref:MAVG}.\n"
        "\n"
        "See also {ref:aggregation/MAX}, {ref:RMAX}."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    resources=_MFUNC_RESOURCES,
    examples=[
        *_make_mfunc_examples("mmax"),
    ],
)

FUNCTION_MAVG = FunctionDocRegistryItem(
    name="mavg",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns the moving average of values in a fixed-size window defined by "
        "the sort order and arguments:\n"
        "{text:mwindow_description}\n"
        "Window functions with a similar behavior: {ref:MSUM}, {ref:MCOUNT}, {ref:MMIN}, {ref:MMAX}.\n"
        "\n"
        "See also {ref:aggregation/AVG}, {ref:RAVG}."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    resources=_MFUNC_RESOURCES,
    examples=[
        *_make_mfunc_examples("mavg"),
    ],
)


def _make_lag_examples(func: str) -> List[DataExample]:
    func = func.upper()
    examples = [
        DataExample(
            example_config=ExampleConfig(
                name=_("Example with the optional argument"),
                source=_SOURCE_SALES_DATA_1,
                group_by=["[City]"],
                order_by=["[City]"],
                show_source_table=True,
                formula_fields=[("City", "[City]"), ("Order Sum", "SUM([Orders])")],
                additional_transformations=[
                    [
                        ("City", "[City]"),
                        ("Order Sum", "[Order Sum]"),
                        (f"{func} 1", f"{func}([Order Sum], 1)"),
                        (f"{func} 2", f"{func}([Order Sum], -2)"),
                    ],
                ],
                override_formula_fields=[
                    ("City", "[City]"),
                    ("Order Sum", "SUM([Orders])"),
                    (f"{func} 1", f"{func}(SUM([Orders]), 1)"),
                    (f"{func} 2", f"{func}(SUM([Orders]), -2)"),
                ],
            ),
        ),
    ]
    return examples


FUNCTION_LAG = FunctionDocRegistryItem(
    name="lag",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_(
        "Returns {arg:0} re-evaluated against the row that is offset from the current "
        "row by {arg:1} within the specified window:\n"
        "- Positive {arg:1} seeks among preceding rows.\n"
        "- Negative {arg:1} seeks among following rows.\n"
        "\n"
        "By default {arg:1} is `1`.\n"
        "\n"
        "If there is no available value ({arg:1} reaches before the first row or "
        "after the last one), then {arg:2} is returned. If {arg:2} is not specified, "
        "then `NULL` is used.\n"
        "\n"
        "See also {ref:AGO} for a non-window function alternative."
    ),
    notes=_ORDERED_WFUNC_NOTES,
    examples=[
        *_make_standard_window_examples("lag"),
        *_make_lag_examples("lag"),
        *_make_simple_order_by_examples("lag"),
    ],
)

FUNCTION_FIRST = FunctionDocRegistryItem(
    name="first",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_("Returns the value of {arg:0} from the first row in the window. See also {ref:LAST}."),
    notes=_ORDERED_WFUNC_NOTES,
    examples=[
        *_make_standard_window_examples("first"),
        *_make_simple_order_by_examples("first"),
    ],
)

FUNCTION_LAST = FunctionDocRegistryItem(
    name="last",
    is_window=True,
    naming_provider=CategoryPostfixNamingProvider(),
    category=CATEGORY_WINDOW,
    description=_("Returns the value of {arg:0} from the last row in the window. See also {ref:FIRST}."),
    notes=_ORDERED_WFUNC_NOTES,
    examples=[
        *_make_standard_window_examples("last"),
        *_make_simple_order_by_examples("last"),
    ],
)

FUNCTIONS_WINDOW = [
    FUNCTION_RANK,
    FUNCTION_RANK_DENSE,
    FUNCTION_RANK_UNIQUE,
    FUNCTION_RANK_PERCENTILE,
    FUNCTION_SUM_WINDOW,
    FUNCTION_MIN_WINDOW,
    FUNCTION_MAX_WINDOW,
    FUNCTION_COUNT_WINDOW,
    FUNCTION_AVG_WINDOW,
    FUNCTION_SUM_IF_WINDOW,
    FUNCTION_COUNT_IF_WINDOW,
    FUNCTION_AVG_IF_WINDOW,
    FUNCTION_RSUM,
    FUNCTION_RCOUNT,
    FUNCTION_RMIN,
    FUNCTION_RMAX,
    FUNCTION_RAVG,
    FUNCTION_MSUM,
    FUNCTION_MCOUNT,
    FUNCTION_MMIN,
    FUNCTION_MMAX,
    FUNCTION_MAVG,
    FUNCTION_LAG,
    FUNCTION_FIRST,
    FUNCTION_LAST,
]

lib/dl_formula_ref/dl_formula_ref/functions/window.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from contextlib import redirect_stdout
import inspect
import json
import os
from typing import (
    Collection,
    Dict,
    Iterable,
    List,
    Optional,
    Sequence,
    Tuple,
)

import attr
import jinja2
from tabulate import tabulate

from dl_formula.core.dialect import (
    get_all_basic_dialects,
    get_dialect_from_str,
)
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D
import dl_formula.core.exc as exc
from dl_formula.definitions.base import MultiVariantTranslation
from dl_formula.definitions.registry import OPERATION_REGISTRY
from dl_formula_ref.audience import (
    DEFAULT_AUDIENCE,
    Audience,
)
from dl_formula_ref.config import (
    CatPathTemplate,
    ConfigVersion,
    FuncDocConfigVersion,
    FuncDocTemplateConfig,
    FuncPathTemplate,
    RefDocGeneratorConfig,
    get_generator_config,
)
from dl_formula_ref.i18n.registry import get_localizer
from dl_formula_ref.paths import PathRenderer
from dl_formula_ref.primitives import (
    RawFunc,
    RawMultiAudienceFunc,
)
from dl_formula_ref.reference import (
    FuncReference,
    load_func_reference_from_registry,
)
from dl_formula_ref.registry.dialect_extractor import COMPENG_SUPPORT
from dl_formula_ref.registry.example import (
    DataExample,
    ExampleBase,
)
from dl_formula_ref.registry.example_base import DataExampleRendererConfig
from dl_formula_ref.registry.registry import RefFunctionKey
from dl_formula_ref.rendered import (
    RenderedFunc,
    RenderedMultiAudienceFunc,
)
from dl_formula_ref.rendering import (
    FuncRenderer,
    human_category,
    human_dialect,
)
from dl_formula_ref.texts import (
    ANY_DIALECTS,
    DOC_ALL_TITLE,
    DOC_AVAIL_TITLE,
    DOC_OVERVIEW_TEXT,
    DialectStyle,
)
from dl_formula_testing.database import (
    Db,
    make_db_config,
    make_db_from_config,
)


try:
    from dl_formula_ref.examples.data_prep import DataPreparer
except exc.ParserNotFoundError:

    class DataPreparer:  # type: ignore
        def __init__(self, *args, **kwargs):
            raise exc.ParserNotFoundError()


def get_jinja_env(gen_config: RefDocGeneratorConfig) -> jinja2.Environment:
    import dl_formula_ref as top

    return jinja2.Environment(
        loader=jinja2.PackageLoader(top.__name__, gen_config.template_dir_rel),
    )


@attr.s
class ReferenceDocGenerator:
    _locale: str = attr.ib(kw_only=True)
    _config_version: ConfigVersion = attr.ib(kw_only=True)

    _gen_config: RefDocGeneratorConfig = attr.ib(init=False)
    _func_ref: FuncReference = attr.ib(init=False)
    _renderers_by_tmpl: Dict[Tuple[FuncPathTemplate, CatPathTemplate], FuncRenderer] = attr.ib(init=False, factory=dict)
    _jinja_env: jinja2.Environment = attr.ib(init=False)

    def __attrs_post_init__(self) -> None:
        self._gen_config = get_generator_config(self._config_version)
        self._func_ref = load_func_reference_from_registry(
            scopes_by_audience=self._gen_config.function_scopes, supported_dialects=self._gen_config.supported_dialects
        )
        self._jinja_env = get_jinja_env(self._gen_config)

    def _get_renderer(self, doc_config: FuncDocTemplateConfig) -> FuncRenderer:
        path_renderer = PathRenderer(
            func_ref=self._func_ref,
            func_path_template=doc_config.func_file_path,
            cat_path_template=doc_config.cat_file_path,
        )
        key = (doc_config.func_file_path, doc_config.cat_file_path)
        if key not in self._renderers_by_tmpl:
            self._renderers_by_tmpl[key] = FuncRenderer(
                func_ref=self._func_ref,
                locale=self._locale,
                example_rend_config=DataExampleRendererConfig(
                    jinja_env=self._jinja_env,
                    template_filename=self._gen_config.doc_example_template,
                    example_data_filename=self._gen_config.example_data_file,
                ),
                block_conditions=self._gen_config.block_conditions,
                path_renderer=path_renderer,
                resource_overrides=self._gen_config.resource_overrides,
            )
        return self._renderers_by_tmpl[key]

    def _get_single_rendered_func(
        self,
        func_key: RefFunctionKey,
        doc_config: FuncDocTemplateConfig,
    ) -> RenderedMultiAudienceFunc:
        renderer = self._get_renderer(doc_config=doc_config)
        return renderer.render_multi_func_func(func_key=func_key)

    def _render_funcs(
        self,
        raw_funcs: List[RawMultiAudienceFunc],
        doc_config: FuncDocTemplateConfig,
    ) -> List[RenderedMultiAudienceFunc]:
        func_doc_structs = []
        for raw_func in raw_funcs:
            func_key = RefFunctionKey.normalized(name=raw_func.name, category_name=raw_func.category.name)
            func_doc_structs.append(self._get_single_rendered_func(func_key=func_key, doc_config=doc_config))
        return func_doc_structs

    def generate_doc_func(self, outdir: str):
        raw_funcs = self._func_ref.as_list()
        for doc_config in self._gen_config.func_doc_configs.values():
            template = self._jinja_env.get_template(doc_config.template_file)
            rend_funcs = self._render_funcs(raw_funcs=raw_funcs, doc_config=doc_config)
            for multi_func in rend_funcs:
                text = template.render(multi_func=multi_func, _=get_localizer(self._locale).translate)
                full_path = os.path.join(outdir, multi_func.file_path)
                dirname = os.path.dirname(full_path)
                if not os.path.exists(dirname):
                    os.makedirs(dirname)
                with open(full_path, mode="w", encoding="utf-8") as outfile:
                    outfile.write(text)
                print(full_path)

    def _group_raw_funcs_by_category(
        self, raw_funcs: List[RawMultiAudienceFunc]
    ) -> Dict[str, List[RawMultiAudienceFunc]]:
        funcs_by_category = defaultdict(list)
        for func in raw_funcs:
            funcs_by_category[func.category.name].append(func)
        return funcs_by_category

    def generate_doc_toc(self, list_funcs: bool = False) -> None:
        trans = get_localizer(self._locale).translate
        doc_config = self._gen_config.func_doc_configs[FuncDocConfigVersion.overview_shortcut]
        raw_funcs = self._func_ref.as_list()
        renderer = self._get_renderer(doc_config=doc_config)
        funcs_by_category = self._group_raw_funcs_by_category(raw_funcs=raw_funcs)
        global_audiences = set(self._gen_config.function_scopes.keys())

        print("items:")
        print("  - name: {}\n    href: {}".format(trans(DOC_ALL_TITLE), self._gen_config.doc_all_filename))

        localized_categories = {
            category: human_category(category=funcs_by_category[category][0].category.name, locale=self._locale)
            for category in funcs_by_category.keys()
        }
        locale_sorted_categories = sorted(funcs_by_category.keys(), key=lambda category: localized_categories[category])
        for category in locale_sorted_categories:
            print(f"  - name: {localized_categories[category]}")
            print("    items:")
            print("      - name: {}".format(trans(DOC_OVERVIEW_TEXT)))
            print("        href: {}".format(renderer.path_renderer.get_cat_path(category_name=category)))
            if list_funcs:
                for multi_func in sorted(
                    funcs_by_category[category], key=lambda func: func.get_title(locale=self._locale)
                ):
                    func_key = RefFunctionKey.normalized(
                        name=multi_func.name,
                        category_name=multi_func.category.name,
                    )
                    audiences = set(multi_func.keys())
                    items: Iterable[tuple[Audience, RawFunc]]
                    if audiences == global_audiences:
                        # Consolidate entries if all audiences are present
                        # even if they have different content
                        items = [(DEFAULT_AUDIENCE, next(iter(multi_func.values())))]
                    else:
                        items = multi_func.items()

                    for audience, raw_func in items:
                        print("      - name: {}".format(raw_func.get_short_title(locale=self._locale)))
                        print("        href: {}".format(renderer.path_renderer.get_func_path(func_key=func_key)))
                        if not audience.default:
                            print(f'        when: audience == "{audience.name}"')

        if self._gen_config.gen_availability_table:
            print("  - name: {}\n    href: {}".format(trans(DOC_AVAIL_TITLE), self._gen_config.doc_avail_filename))

    def _generate_doc_list(
        self,
        context_path: str,
        title: str,
        description: str,
        rend_funcs: List[RenderedMultiAudienceFunc],
        in_category: bool,
        meta_title: str = "",
        meta_description: str = "",
        meta_keywords: Sequence[str] = (),
    ) -> None:
        trans = get_localizer(self._locale).translate
        template = self._jinja_env.get_template(self._gen_config.doc_list_template)
        doc_config = self._gen_config.func_doc_configs[FuncDocConfigVersion.overview_shortcut]
        renderer = self._get_renderer(doc_config=doc_config)
        relative_path_renderer = renderer.path_renderer.child(context_path)

        def make_link(func: RenderedFunc) -> str:
            func_key = RefFunctionKey.normalized(name=func.name, category_name=func.category_name)
            return relative_path_renderer.get_func_path(func_key=func_key)

        text = template.render(
            _=trans,
            functions=rend_funcs,
            make_link=make_link,
            title=trans(title),
            description=description,
            in_category=in_category,
            meta_title=meta_title,
            meta_description=meta_description,
            meta_keywords=meta_keywords,
        )
        print(text)

    def generate_doc_list_all(self, context_path: str):
        doc_config = self._gen_config.func_doc_configs[FuncDocConfigVersion.overview_shortcut]
        raw_funcs = self._func_ref.as_list()
        rend_funcs = self._render_funcs(raw_funcs=raw_funcs, doc_config=doc_config)
        self._generate_doc_list(
            context_path=context_path,
            rend_funcs=rend_funcs,
            title=DOC_ALL_TITLE,
            description="",
            in_category=False,
        )

    def generate_doc_list_category(self, category: str, context_path: str):
        doc_config = self._gen_config.func_doc_configs[FuncDocConfigVersion.overview_shortcut]
        raw_funcs = self._func_ref.filter(category=category)
        rend_funcs = self._render_funcs(raw_funcs=raw_funcs, doc_config=doc_config)
        rend_first_func = [
            # Some funcs have double categories, so find one whose primary category is the one we need
            # TODO: refactor this
            func
            for func in rend_funcs
            if func.category_name == category
        ][0]
        title = human_category(category=category, locale=self._locale)
        self._generate_doc_list(
            context_path=context_path,
            rend_funcs=rend_funcs,
            title=title,
            description=rend_first_func.category_description,
            in_category=True,
            meta_title=title,
            meta_description=rend_first_func.category_description_short,
            meta_keywords=rend_first_func.category_keywords,
        )

    def generate_doc_availability_list(self, context_path: str, audience: Audience) -> str:
        priority = {"==", "NEG", "*", "/", "%", "^", "+", "-"}
        raw_funcs = self._func_ref.as_list()
        raw_funcs = sorted(
            raw_funcs,
            # move symbols to the top of the list
            key=lambda func: (
                " " + func.get_title(locale=self._locale)
                if func.name in priority
                else func.get_title(locale=self._locale)
            ),
        )
        doc_config = self._gen_config.func_doc_configs[FuncDocConfigVersion.overview_shortcut]
        renderer = self._get_renderer(doc_config=doc_config)
        relative_path_renderer = renderer.path_renderer.child(context_path)

        scopes = self._gen_config.function_scopes.get(
            audience,
            self._gen_config.function_scopes.get(DEFAULT_AUDIENCE),
        )
        assert scopes is not None
        dialects = [d for d in get_all_basic_dialects() if d in self._gen_config.supported_dialects]
        table_data = []
        for multi_func in raw_funcs:
            func = multi_func.get(audience, multi_func.get(DEFAULT_AUDIENCE))
            if func is None:
                continue
            func_key = RefFunctionKey.normalized(name=func.name, category_name=func.category.name)
            supported_dialects: Collection[DialectCombo]
            if func.category.name == "window":  # FIXME CATEGORY_WINDOW.name
                supported_dialects = COMPENG_SUPPORT
            else:
                supported_dialects = func.dialects
                if len(supported_dialects) == 1 and next(iter(supported_dialects)) == D.ANY:
                    supported_dialects = ANY_DIALECTS
            func_data = [
                "[{}]({})".format(
                    func.get_title(locale=self._locale),
                    relative_path_renderer.get_func_path(func_key=func_key),
                ),
                *[("X" if dialect in supported_dialects else "") for dialect in dialects],
            ]
            table_data.append(func_data)
        headers = [
            "Function",
            *[human_dialect(dialect, locale=self._locale, style=DialectStyle.multiline) for dialect in dialects],
        ]
        table = tabulate(table_data, headers=headers, tablefmt="pipe")
        # now hack column alignment to be centered for all columns except the first one
        first_match = table.find("-|")
        table = table[: first_match + 2] + table[first_match + 2 :].replace("-|", ":|")
        return table

    def generate_doc_availability_list_for_all_audiences(self, context_path: str) -> None:
        table_by_audience: dict[Audience, str] = {}
        for audience in sorted(self._gen_config.function_scopes.keys()):
            avail_tbl = self.generate_doc_availability_list(context_path=context_path, audience=audience)
            table_by_audience[audience] = avail_tbl

        all_tables = set(table_by_audience.values())
        if len(all_tables) == 1:
            table_by_audience = {DEFAULT_AUDIENCE: next(iter(all_tables))}

        template = self._jinja_env.get_template(self._gen_config.doc_avail_template)
        text = template.render(
            table_by_audience=table_by_audience,
            _=get_localizer(self._locale).translate,
        )
        print(text)

    def generate_doc_full_dir(self, outdir: str):
        raw_funcs = self._func_ref.as_list()
        funcs_by_category = self._group_raw_funcs_by_category(raw_funcs=raw_funcs)

        func_dir = outdir
        self.generate_doc_func(outdir=func_dir)
        all_funcs_path = os.path.join(func_dir, self._gen_config.doc_all_filename)
        with open(all_funcs_path, "w") as outfile:
            with redirect_stdout(outfile):
                self.generate_doc_list_all(context_path=os.path.dirname(all_funcs_path))
        for doc_config in self._gen_config.func_doc_configs.values():
            for category in sorted(funcs_by_category):
                category_path = os.path.join(func_dir, doc_config.cat_file_path.format(category_name=category))
                with open(category_path, "w") as outfile:
                    with redirect_stdout(outfile):
                        self.generate_doc_list_category(context_path=os.path.dirname(category_path), category=category)
        if self._gen_config.gen_availability_table:
            availability_path = os.path.join(func_dir, self._gen_config.doc_avail_filename)
            with open(availability_path, "w") as outfile:
                with redirect_stdout(outfile):
                    self.generate_doc_availability_list_for_all_audiences(
                        context_path=os.path.dirname(availability_path),
                    )
        with open(os.path.join(outdir, self._gen_config.doc_toc_filename), "w") as outfile:
            with redirect_stdout(outfile):
                self.generate_doc_toc(list_funcs=True)

    @staticmethod
    def _get_func_base_class(name: str) -> Optional[type]:
        """Find the first (base) class for the given function"""
        name = name.lower()
        for _i, definition in OPERATION_REGISTRY.items():
            func_name = definition.name
            assert func_name is not None
            if func_name.lower() == name:
                def_cls = type(definition)
                break
        else:
            return None

        mro = inspect.getmro(def_cls)  # `def_cls` itself plus all its bases
        for earliest_super_cls in reversed(mro):
            if (
                issubclass(earliest_super_cls, MultiVariantTranslation)
                and (earliest_super_cls.name or "").lower() == name
            ):
                return earliest_super_cls

        return None

    def _get_func_source_info(self, name: str) -> Tuple[str, int]:
        """Get file name and line number for given function"""
        func_cls = self._get_func_base_class(name)
        assert func_cls is not None
        filename = inspect.getsourcefile(func_cls)
        assert filename is not None
        lineno = inspect.getsourcelines(func_cls)[1]
        return filename, lineno

    def _load_db_config(self) -> Dict[DialectCombo, Db]:
        with open(self._gen_config.db_config_file) as config_file:
            raw_config: dict = json.load(config_file)
        return {
            get_dialect_from_str(d_name): make_db_from_config(
                make_db_config(dialect=get_dialect_from_str(d_name), url=d_url)
            )
            for d_name, d_url in raw_config.items()
        }

    def generate_example_data(self) -> None:
        raw_funcs = self._func_ref.as_list()
        db_by_dialect = self._load_db_config()
        preparer = DataPreparer(
            storage_filename=self._gen_config.example_data_file,
            db_by_dialect=db_by_dialect,
            default_example_dialect=self._gen_config.default_example_dialect,
        )
        examples: list[ExampleBase]
        for multi_func in raw_funcs:
            examples = []
            for raw_func in multi_func.values():
                for example in raw_func.examples:
                    if example not in examples:
                        examples.append(example)
            for example in examples:
                if isinstance(example, DataExample):
                    preparer.generate_example_data(
                        func_name=multi_func.internal_name,
                        example=example.example_config,
                    )

        preparer.save()

lib/dl_formula_ref/dl_formula_ref/generator.py
datalens-tech/datalens-backend
from typing import Iterable

import attr

import dl_formula_ref as package
from dl_i18n.localizer_base import (
    Localizer,
    LocalizerLoader,
    Translatable,
    TranslationConfig,
)


_LOCALIZATION_CONFIGS: set[TranslationConfig] = set()

DOMAIN = f"{package.__name__}_{package.__name__}"


def register_translation_configs(config: Iterable[TranslationConfig]) -> None:
    _LOCALIZATION_CONFIGS.update(config)


@attr.s
class _StringLocalizer:
    """
    An adapter for `Localizer`, which accepts only `Translatable` objects.
    """

    _translatable_localizer: Localizer = attr.ib(kw_only=True)
    _default_domain: str = attr.ib(kw_only=True)

    def translate(self, text: str | Translatable) -> str:
        if isinstance(text, str):
            text = Translatable(text, domain=self._default_domain)
        assert isinstance(text, Translatable)
        return self._translatable_localizer.translate(text)


def get_localizer(locale: str):
    loc_loader = LocalizerLoader(configs=_LOCALIZATION_CONFIGS)
    loc_factory = loc_loader.load()
    localizer = _StringLocalizer(
        translatable_localizer=loc_factory.get_for_locale(locale=locale),
        default_domain=DOMAIN,
    )
    return localizer

lib/dl_formula_ref/dl_formula_ref/i18n/registry.py
datalens-tech/datalens-backend
from typing import (
    Collection,
    Optional,
)

import attr

from dl_formula.loader import load_formula_lib
from dl_formula_ref.formula_ref_plugins import register_all_plugins


@attr.s(frozen=True)
class FormulaRefLibraryConfig:
    plugin_ep_names: Optional[Collection[str]] = attr.ib(kw_only=True, default=None)


def load_formula_ref(formula_ref_lib_config: FormulaRefLibraryConfig = FormulaRefLibraryConfig()) -> None:
    load_formula_lib()
    register_all_plugins(plugin_ep_names=formula_ref_lib_config.plugin_ep_names)

lib/dl_formula_ref/dl_formula_ref/loader.py
datalens-tech/datalens-backend
# TODO: Remove or refactor

import gettext


DEFAULT_LOCALE = "en"


def get_gettext():
    if "_" not in globals():
        gettext.install("app")
    return _  # type: ignore  # TODO: fix  # noqa

lib/dl_formula_ref/dl_formula_ref/localization.py
datalens-tech/datalens-backend
from __future__ import annotations

import os.path
from typing import (
    NamedTuple,
    Optional,
)

import attr

from dl_formula_ref.reference import FuncReference
from dl_formula_ref.registry.registry import RefFunctionKey
from dl_formula_ref.texts import HUMAN_CATEGORIES


class FileLink(NamedTuple):
    name: str
    path: str


class FuncPathTemplate(NamedTuple):
    path: str

    def format(self, category_name: str, func_name: str):
        return self.path.format(func_name=func_name, category_name=category_name)


class CatPathTemplate(NamedTuple):
    path: str

    def format(self, category_name: str):
        return self.path.format(category_name=category_name)


@attr.s
class PathRenderer:
    _cat_path_template: CatPathTemplate = attr.ib(kw_only=True)
    _func_path_template: FuncPathTemplate = attr.ib(kw_only=True)
    _func_ref: FuncReference = attr.ib(kw_only=True)
    _base_path: str = attr.ib(kw_only=True, default="")

    def _append_path(self, name: str) -> str:
        if not self._base_path:
            return name
        return os.path.join(self._base_path, name)

    def _get_relative_path(self, path: str) -> str:
        if not self._base_path:
            return path
        base_parts = self._base_path.replace(os.sep, "/").split("/")
        child_parts = path.replace(os.sep, "/").split("/")
        common_part_cnt = 0
        while common_part_cnt < min(len(base_parts), len(child_parts)):
            if base_parts[-(common_part_cnt + 1) :] != child_parts[: common_part_cnt + 1]:
                break
            common_part_cnt += 1

        result = os.path.join(*child_parts[common_part_cnt:])
        return result

    def child(self, name: str) -> PathRenderer:
        return attr.evolve(self, base_path=self._append_path(name))

    def get_func_path(self, func_key: RefFunctionKey, anchor_name: Optional[str] = None) -> str:
        raw_func = self._func_ref.get_func(func_key=func_key)
        filename = raw_func.internal_name.upper()
        path = self._get_relative_path(
            self._func_path_template.format(
                func_name=filename,
                category_name=raw_func.category.name,
            )
        )
        if anchor_name:
            path = f"{path}#{anchor_name}"
        return path

    def get_cat_path(self, category_name: str, anchor_name: Optional[str] = None) -> str:
        path = self._get_relative_path(
            self._cat_path_template.format(
                category_name=category_name,
            )
        )
        if anchor_name:
            path = f"{path}#{anchor_name}"
        return path

    def get_func_link(
        self,
        func_name: str,
        category_name: Optional[str] = None,
        anchor_name: Optional[str] = None,
    ) -> FileLink:
        if category_name is None:
            found_in_categories = [raw_func.category.name for raw_func in self._func_ref.filter(name=func_name)]
            if len(found_in_categories) > 1:
                raise ValueError(
                    "Function with name {} found in following categories: {} - but category is not specified. "
                    'Use the following syntax to specify category: "{{ref:category_name/FUNC_NAME[:ref text]}}"'
                    "".format(func_name, found_in_categories)
                )
            category_name = found_in_categories[0]
        func_key = RefFunctionKey.normalized(name=func_name, category_name=category_name)
        return FileLink(
            name=func_name.upper(),
            path=self.get_func_path(func_key=func_key, anchor_name=anchor_name),
        )

    def get_cat_link(
        self,
        category_name: str,
        anchor_name: Optional[str] = None,
    ) -> FileLink:
        return FileLink(
            name=HUMAN_CATEGORIES.get(category_name, category_name),
            path=self.get_cat_path(category_name=category_name, anchor_name=anchor_name),
        )

lib/dl_formula_ref/dl_formula_ref/paths.py
datalens-tech/datalens-backend
from typing import ClassVar

from dl_formula.core.dialect import DialectCombo
from dl_formula_ref.config import (
    ConfigVersion,
    RefDocGeneratorConfig,
)
from dl_formula_ref.functions.type_conversion import DbCastExtension
from dl_formula_ref.texts import StyledDialect
from dl_i18n.localizer_base import TranslationConfig


class FormulaRefPlugin:
    configs: ClassVar[dict[ConfigVersion, RefDocGeneratorConfig]] = {}
    translation_configs: ClassVar[frozenset[TranslationConfig]] = frozenset()

    # Dialects to include whenever ANY is used instead of implicit dialect list
    any_dialects: ClassVar[frozenset[DialectCombo]] = frozenset()

    # Human-readable dialect names
    human_dialects: ClassVar[dict[DialectCombo, StyledDialect]] = {}

    # Dialects that work with compeng
    compeng_support_dialects: ClassVar[frozenset[DialectCombo]] = frozenset()

    # Connector-specific types and comments for the DB_CAST function
    db_cast_extension: ClassVar[DbCastExtension] = DbCastExtension()

lib/dl_formula_ref/dl_formula_ref/plugins/base/plugin.py
datalens-tech/datalens-backend
import os

from dl_formula_ref.i18n.registry import DOMAIN
from dl_i18n.localizer_base import TranslationConfig


_LOCALE_DIR = os.path.join(os.path.dirname(__file__), "../..", "locales")

CONFIGS = [
    TranslationConfig(path=_LOCALE_DIR, domain=DOMAIN, locale="en"),
    TranslationConfig(path=_LOCALE_DIR, domain=DOMAIN, locale="ru"),
]

lib/dl_formula_ref/dl_formula_ref/plugins/default/i18n.py
datalens-tech/datalens-backend
from dl_formula_ref.config import (
    DOC_GEN_CONFIG_DEFAULT,
    ConfigVersion,
)
from dl_formula_ref.plugins.base.plugin import FormulaRefPlugin
from dl_formula_ref.plugins.default.i18n import CONFIGS


class DefaultFormulaRefPlugin(FormulaRefPlugin):
    configs = {
        ConfigVersion.default: DOC_GEN_CONFIG_DEFAULT,
    }
    translation_configs = frozenset(CONFIGS)

lib/dl_formula_ref/dl_formula_ref/plugins/default/plugin.py
datalens-tech/datalens-backend
from typing import Type

from dl_formula_ref.config import register_config_version
from dl_formula_ref.functions.type_conversion import register_db_cast_extension
from dl_formula_ref.i18n.registry import register_translation_configs
from dl_formula_ref.plugins.base.plugin import FormulaRefPlugin
from dl_formula_ref.registry.dialect_extractor import register_compeng_support_dialects
from dl_formula_ref.texts import (
    register_any_dialects,
    register_human_dialects,
)


class FormulaRefPluginRegistrator:
    def register_plugin(self, plugin_cls: Type[FormulaRefPlugin]) -> None:
        for version, config in plugin_cls.configs.items():
            register_config_version(version=version, config=config)

        register_any_dialects(plugin_cls.any_dialects)
        register_human_dialects(plugin_cls.human_dialects)
        register_translation_configs(plugin_cls.translation_configs)
        register_compeng_support_dialects(plugin_cls.compeng_support_dialects)
        register_db_cast_extension(plugin_cls.db_cast_extension)


FORMULA_REF_PLUGIN_REG = FormulaRefPluginRegistrator()

lib/dl_formula_ref/dl_formula_ref/plugins/registration.py
datalens-tech/datalens-backend
"""
Generate structures used for creating documentation
"""

from __future__ import annotations

from typing import (
    Callable,
    Collection,
    Iterable,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
)

import attr

from dl_formula.core.dialect import DialectCombo
from dl_formula_ref.audience import Audience
from dl_formula_ref.registry.aliased_res import AliasedResourceRegistryBase
from dl_formula_ref.registry.arg_base import FuncArg
from dl_formula_ref.registry.base import FunctionDocCategory
from dl_formula_ref.registry.example_base import ExampleBase
from dl_formula_ref.registry.note import (
    NoteLevel,
    NoteType,
    ParameterizedNote,
)
from dl_formula_ref.registry.signature_base import FunctionSignatureCollection
from dl_formula_ref.registry.text import ParameterizedText


TOP_NOTE_LEVELS: Collection[Tuple[Optional[NoteType], Optional[NoteLevel]]] = [
    (NoteType.REGULAR, NoteLevel.alert),
    (NoteType.REGULAR, NoteLevel.warning),
]
BOTTOM_NOTE_LEVELS: Collection[Tuple[Optional[NoteType], Optional[NoteLevel]]] = [
    (NoteType.REGULAR, NoteLevel.info),
    (NoteType.ARG_RESTRICTION, None),
]


@attr.s(frozen=True)
class RawFunc:
    _name: str = attr.ib(kw_only=True)
    _title_factory: Callable[[str], str] = attr.ib(kw_only=True)
    _short_title_factory: Callable[[str], str] = attr.ib(kw_only=True)
    _internal_name: str = attr.ib(kw_only=True)
    _description: str = attr.ib(kw_only=True)
    _args: List[FuncArg] = attr.ib(kw_only=True)
    _signature_coll: FunctionSignatureCollection = attr.ib(kw_only=True)
    _notes: List[ParameterizedNote] = attr.ib(kw_only=True)
    _return_type: ParameterizedText = attr.ib(kw_only=True)
    _category: FunctionDocCategory = attr.ib(kw_only=True)
    _resources: AliasedResourceRegistryBase = attr.ib(kw_only=True)
    _examples: Sequence[ExampleBase] = attr.ib(kw_only=True)
    _dialects: Set[DialectCombo] = attr.ib(kw_only=True)

    @property
    def name(self) -> str:
        return self._name.upper()

    @property
    def description(self) -> str:
        return self._description

    @property
    def category(self) -> FunctionDocCategory:
        return self._category

    @property
    def resources(self) -> AliasedResourceRegistryBase:
        return self._resources

    @property
    def examples(self) -> Sequence[ExampleBase]:
        return self._examples

    @property
    def args(self) -> List[FuncArg]:
        return self._args

    @property
    def return_type(self) -> ParameterizedText:
        return self._return_type

    @property
    def internal_name(self) -> str:
        return self._internal_name

    def get_title(self, locale: str) -> str:
        return self._title_factory(locale)

    def get_short_title(self, locale: str) -> str:
        return self._short_title_factory(locale)

    @property
    def short_title_factory(self) -> Callable[[str], str]:
        return self._short_title_factory

    @property
    def title_factory(self) -> Callable[[str], str]:
        return self._title_factory

    @property
    def signature_coll(self) -> FunctionSignatureCollection:
        return self._signature_coll

    @property
    def dialects(self) -> Set[DialectCombo]:
        return self._dialects

    def _filter_notes(
        self,
        notes: Iterable[ParameterizedNote],
        levels: Collection[Tuple[Optional[NoteType], Optional[NoteLevel]]],
    ) -> List[ParameterizedNote]:
        result: List[ParameterizedNote] = []
        for note in notes or ():
            for note_type_mask, note_level_mask in levels:
                if (note_type_mask is None or note.type == note_type_mask) and (
                    note_level_mask is None or note.level == note_level_mask
                ):
                    result.append(note)
                    break

        return result

    @property
    def bottom_notes(self) -> List[ParameterizedNote]:
        return self._filter_notes(notes=self._notes, levels=BOTTOM_NOTE_LEVELS)

    @property
    def top_notes(self) -> List[ParameterizedNote]:
        return self._filter_notes(notes=self._notes, levels=TOP_NOTE_LEVELS)


@attr.s(frozen=True)
class RawMultiAudienceFunc:
    name: str = attr.ib(kw_only=True)
    category: FunctionDocCategory = attr.ib(kw_only=True)
    internal_name: str = attr.ib(kw_only=True)
    _title_factory: Callable[[str], str] = attr.ib(kw_only=True)
    _short_title_factory: Callable[[str], str] = attr.ib(kw_only=True)
    _aud_dict: dict[Audience, RawFunc] = attr.ib(kw_only=True)

    def items(self) -> Iterable[tuple[Audience, RawFunc]]:
        items: list[tuple[Audience, RawFunc]] = list(self._aud_dict.items())
        return sorted(items, key=lambda item: item[0].name)

    def keys(self) -> Iterable[Audience]:
        return self._aud_dict.keys()

    def values(self) -> Iterable[RawFunc]:
        return self._aud_dict.values()

    def get(self, audience: Audience, default: Optional[RawFunc] = None) -> Optional[RawFunc]:
        return self._aud_dict.get(audience, default)

    def get_title(self, locale: str) -> str:
        return self._title_factory(locale)

    def get_short_title(self, locale: str) -> str:
        return self._short_title_factory(locale)

    @classmethod
    def from_aud_dict(cls, aud_dict: dict[Audience, RawFunc]) -> RawMultiAudienceFunc:
        names: set[str] = set()
        internal_names: set[str] = set()
        for _audience, raw_func in aud_dict.items():
            names.add(raw_func.name)
            internal_names.add(raw_func.internal_name)

        assert len(names) == 1
        assert len(internal_names) == 1

        one_raw_func = next(iter(aud_dict.values()))

        return cls(
            name=next(iter(names)),
            category=one_raw_func.category,
            internal_name=next(iter(internal_names)),
            aud_dict=aud_dict,
            title_factory=one_raw_func.short_title_factory,
            short_title_factory=one_raw_func.short_title_factory,
        )

lib/dl_formula_ref/dl_formula_ref/primitives.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Collection,
    Optional,
    Type,
)

from dl_formula.core.dialect import (
    DialectCombo,
    get_all_basic_dialects,
)
from dl_formula.definitions.base import MultiVariantTranslation
from dl_formula.definitions.flags import ContextFlag
from dl_formula_ref.audience import Audience
from dl_formula_ref.primitives import (
    RawFunc,
    RawMultiAudienceFunc,
)
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.env import GenerationEnvironment
from dl_formula_ref.registry.registry import (
    FUNC_REFERENCE_REGISTRY,
    RefFunctionKey,
)
from dl_formula_ref.registry.tools import populate_registry_from_definitions


class FuncReference:  # TODO: Merge with FunctionReferenceRegistry
    def __init__(self, funcs_by_key: Optional[dict[RefFunctionKey, RawMultiAudienceFunc]] = None):
        self._funcs_by_key = funcs_by_key or {}

    def __contains__(self, key: RefFunctionKey) -> bool:
        return key in self._funcs_by_key

    def add_func(self, func: RawMultiAudienceFunc) -> None:
        func_key = RefFunctionKey.normalized(name=func.name, category_name=func.category.name)
        self._funcs_by_key[func_key] = func

    def get_func(self, func_key: RefFunctionKey) -> RawMultiAudienceFunc:
        return self._funcs_by_key[func_key]

    def filter(self, category: Optional[str] = None, name: Optional[str] = None) -> list[RawMultiAudienceFunc]:
        result = []
        for func in self._funcs_by_key.values():
            if category is not None and func.category.name != category or name is not None and func.name != name:
                continue
            result.append(func)
        return result

    def as_list(self) -> list[RawMultiAudienceFunc]:
        return self.filter()


def _is_deprecated(defn: Type[MultiVariantTranslation]) -> bool:
    if not defn.return_flags:
        return False
    for dialect in get_all_basic_dialects():
        if defn._get_return_flags(dialect) & ContextFlag.DEPRECATED:
            return True
    return False


def _make_raw_func(item: FunctionDocRegistryItem, env: GenerationEnvironment) -> RawFunc:
    raw_func = RawFunc(
        name=item.name,
        title_factory=item.get_title,
        short_title_factory=item.get_short_title,
        internal_name=item.internal_name.upper(),
        description=item.description,
        dialects=item.get_dialects(env=env),
        args=item.get_args(env=env),
        signature_coll=item.get_signatures(env=env),
        notes=item.get_notes(env=env),
        return_type=item.get_return_type(env=env),
        category=item.category,
        resources=item.all_resources,
        examples=item.get_examples(env=env),
    )
    return raw_func


def _all_same(items: Collection[Any]) -> bool:
    if not items:
        return True
    one = next(iter(items))
    return all(item == one for item in items)


def load_func_reference_from_registry(
    scopes_by_audience: dict[Audience, int],
    supported_dialects: frozenset[DialectCombo],
) -> FuncReference:
    populate_registry_from_definitions()

    env_by_audience: dict[Audience, GenerationEnvironment] = {
        audience: GenerationEnvironment(
            scopes=aud_scopes,
            supported_dialects=supported_dialects,
        )
        for audience, aud_scopes in scopes_by_audience.items()
    }

    func_ref = FuncReference()
    for _key, item in FUNC_REFERENCE_REGISTRY.items():
        aud_func_dict: dict[Audience, RawFunc] = {}
        for audience, env in env_by_audience.items():
            if item.is_supported(env=env):
                aud_func_dict[audience] = _make_raw_func(item, env=env)
        if not aud_func_dict:
            continue

        all_raw_func_versions = list(aud_func_dict.values())
        if len(all_raw_func_versions) > 1 and _all_same(all_raw_func_versions):
            # Funcs for all audiences are the same, so consolidate them
            # under a single default audience
            aud_func_dict = {Audience(name="", default=True): all_raw_func_versions[0]}
        # Else keep the dict as-is
        multi_aud_func = RawMultiAudienceFunc.from_aud_dict(aud_func_dict)
        func_ref.add_func(multi_aud_func)

    return func_ref

lib/dl_formula_ref/dl_formula_ref/reference.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from collections import ChainMap
from typing import (
    Any,
    ItemsView,
    Iterable,
    MutableMapping,
    TypeVar,
)

import attr

from dl_i18n.localizer_base import Translatable


_RES_TV = TypeVar("_RES_TV", bound="AliasedResource")


@attr.s(frozen=True)
class AliasedResource:
    def clone(self: _RES_TV, **kwargs: Any) -> _RES_TV:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True)
class AliasedTextResource(AliasedResource):
    body: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class AliasedLinkResource(AliasedResource):
    url: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class AliasedNullLinkResource(AliasedResource):
    """For overriding links"""

    url = None


@attr.s(frozen=True)
class AliasedTableResource(AliasedResource):
    table_body: list[list[str | Translatable]] = attr.ib(kw_only=True)


@attr.s(frozen=True)
class AliasedResourceRegistryBase(abc.ABC):
    @abc.abstractmethod
    def get_resources(self) -> MutableMapping[str, AliasedResource]:
        raise NotImplementedError

    def __getitem__(self, item: str) -> AliasedResource:
        return self.get_resources()[item]

    def __add__(self, other: AliasedResourceRegistryBase) -> AliasedResourceRegistryBase:
        return CompoundAliasedResourceRegistry(nested_registries=(self, other))

    def __contains__(self, item: str) -> bool:
        return item in self.get_resources()

    def items(self) -> ItemsView[str, AliasedResource]:
        return self.get_resources().items()

    def keys(self) -> Iterable[str]:
        return self.get_resources().keys()

    def get_link(self, alias: str) -> AliasedLinkResource | AliasedNullLinkResource:
        res = self[alias]
        assert isinstance(res, (AliasedLinkResource, AliasedNullLinkResource))
        return res

    def get_text(self, alias: str) -> AliasedTextResource:
        res = self[alias]
        assert isinstance(res, AliasedTextResource)
        return res

    def get_table(self, alias: str) -> AliasedTableResource:
        res = self[alias]
        assert isinstance(res, AliasedTableResource)
        return res


@attr.s(frozen=True)
class CompoundAliasedResourceRegistry(AliasedResourceRegistryBase):
    _nested_registries: tuple[AliasedResourceRegistryBase, ...] = attr.ib(default=())

    def get_resources(self) -> MutableMapping[str, AliasedResource]:
        return ChainMap(*(res_reg.get_resources() for res_reg in self._nested_registries))


@attr.s(frozen=True)
class SimpleAliasedResourceRegistry(AliasedResourceRegistryBase):
    _resources: MutableMapping[str, AliasedResource] = attr.ib(factory=dict)

    def get_resources(self) -> MutableMapping[str, AliasedResource]:
        return self._resources

lib/dl_formula_ref/dl_formula_ref/registry/aliased_res.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    List,
    NamedTuple,
    Set,
)

from dl_formula.core.datatype import DataType


if TYPE_CHECKING:
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment


class FuncArg(NamedTuple):
    types: Set[DataType]
    name: str
    optional_level: int

    @property
    def formal_type(self) -> str:
        return "|".join(sorted([t.name for t in self.types])) or "ANY"

    @property
    def is_const(self) -> bool:
        return any(t.is_const for t in self.types)


class ArgumentExtractorBase(abc.ABC):
    @abc.abstractmethod
    def get_args(self, item: _registry_base.FunctionDocRegistryItem, env: GenerationEnvironment) -> List[FuncArg]:
        raise NotImplementedError

lib/dl_formula_ref/dl_formula_ref/registry/arg_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    List,
    NamedTuple,
    Optional,
    Sequence,
    Union,
    cast,
)

from dl_formula.core.datatype import DataType
from dl_formula.definitions.type_strategy import (
    DynamicIndexStrategy,
    Fixed,
    FromArgs,
    TypeStrategy,
)
from dl_formula_ref.registry.text import ParameterizedText
from dl_formula_ref.texts import (
    COMMON_TYPE_NOTE,
    DEPENDS_ON_ARGS,
    FROM_ARGS,
)


if TYPE_CHECKING:
    from dl_formula_ref.registry.arg_base import FuncArg
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment


class TypeInfo(NamedTuple):
    ret_type_str: ParameterizedText
    arg_note: Optional[ParameterizedText]


def type_macro(*types: DataType) -> str:
    return f'{{type: {"|".join(t.name for t in types)}}}'


class TypeStrategyInspector:
    @classmethod
    def _extract_common_type_args(
        cls,
        ret_type_strat: TypeStrategy,
        args: List[FuncArg],
    ) -> List[FuncArg]:
        ct_args = []
        if isinstance(ret_type_strat, FromArgs):
            # in this case the return type is determined from types of arguments
            use_arg_indices: Sequence[Union[int, slice]]
            if isinstance(ret_type_strat, DynamicIndexStrategy):
                use_arg_indices = ret_type_strat.get_indices(len(args))
            else:
                use_arg_indices = ret_type_strat._indices
            for ind in use_arg_indices:
                ct_args.extend(args[ind] if isinstance(ind, slice) else [args[ind]])

        return ct_args

    @classmethod
    def _extract_common_type_args_str(
        cls,
        ret_type_strat: TypeStrategy,
        inf_args: bool,
        args: List[FuncArg],
    ) -> str:
        ct_args = cls._extract_common_type_args(ret_type_strat, args)
        arg_str = ", ".join("`{}`".format(a.name) for a in ct_args)
        if inf_args is None:
            arg_str += ", ..."
        return arg_str

    @classmethod
    def get_return_type_and_arg_type_note(
        cls,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> TypeInfo:
        """
        Return string representing the return type
        and, possibly, a note regarding the argument types
        """
        def_list = item.get_implementation_specs(env=env)
        args = item.get_args(env=env)

        strategies = list({id(defn.return_type): defn.return_type for defn in def_list}.values())
        inf_args = any(defn.arg_cnt is None for defn in def_list)
        note = None
        ret_type_str: ParameterizedText = ParameterizedText(text="")

        # deduplicate const and non-const types
        fixed_types = set()
        ind = 0
        while ind < len(strategies):
            ret_type_strat = strategies[ind]
            if isinstance(ret_type_strat, Fixed):
                if ret_type_strat._type.non_const_type not in fixed_types:
                    fixed_types.add(ret_type_strat._type.non_const_type)
                else:
                    del strategies[ind]
                    continue
            ind += 1

        if len(strategies) == 1:
            ret_type_strat = strategies[0]
            if isinstance(ret_type_strat, Fixed):
                ret_type_str = ParameterizedText(text=type_macro(ret_type_strat.type.non_const_type))
            elif isinstance(ret_type_strat, FromArgs):
                common_type_args_str = cls._extract_common_type_args_str(ret_type_strat, inf_args, args)
                ret_type_str = ParameterizedText(text=FROM_ARGS, params=dict(args=common_type_args_str))
                if "," in common_type_args_str:  # FromArgs strategy is used for more than 1 argument
                    # a note about args that they must have the same type
                    note = ParameterizedText(text=COMMON_TYPE_NOTE, params=dict(args=common_type_args_str))

        else:
            resolved_ret_type = False
            if all(isinstance(ret_type_strat, Fixed) for ret_type_strat in strategies):
                # All are `Fixed`. Let's inspect the types that they return
                ret_types = {cast(Fixed, ret_type_strat).type.non_const_type for ret_type_strat in strategies}
                ret_types = ret_types - {DataType.NULL}  # ignore NULL
                if len(ret_types) == 1:
                    # We still have just one possible return type (with NULL ignored)
                    ret_type_str = ParameterizedText(text=type_macro(next(iter(ret_types))))
                    resolved_ret_type = True

            if not resolved_ret_type:
                ret_type_str = ParameterizedText(text=DEPENDS_ON_ARGS)

        return TypeInfo(ret_type_str=ret_type_str, arg_note=note)

lib/dl_formula_ref/dl_formula_ref/registry/arg_common.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from typing import (
    TYPE_CHECKING,
    cast,
)

from dl_formula.core.datatype import DataType
from dl_formula_ref.registry.arg_base import (
    ArgumentExtractorBase,
    FuncArg,
)


if TYPE_CHECKING:
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment


INFINITE_ARG_COUNT = 3


def _get_expandable_types() -> dict[DataType, set[DataType]]:
    """
    Generate a mapping of types that other non-const types can autocast to
    (e.g. {DataType.FLOAT: {DataType.INTEGER}, ...}).

    This is needed to expand all suitable types in function signatures
    if they are not listed explicitly.
    """

    result: dict[DataType, set[DataType]] = {}
    for dtype in DataType:
        if dtype.is_const:
            continue
        autocast_dtypes = cast(
            set[DataType],
            {
                actype
                for actype in dtype.autocast_types
                if actype is not dtype and not actype.is_const and actype is not DataType.NULL
            },
        )
        if autocast_dtypes:
            result[dtype] = autocast_dtypes

    return result


EXPANDABLE_TYPES: dict[DataType, set[DataType]] = _get_expandable_types()


class DefaultArgumentExtractor(ArgumentExtractorBase):
    def get_args(self, item: _registry_base.FunctionDocRegistryItem, env: GenerationEnvironment) -> list[FuncArg]:
        def_list = item.get_implementation_specs(env=env)
        min_arg_cnt: int
        max_arg_cnt: int
        if any(defn.arg_cnt is None for defn in def_list):
            arg_names = max((defn.arg_names or () for defn in def_list), key=lambda names: len(names))
            min_arg_cnt = max(len(arg_names), INFINITE_ARG_COUNT)
            max_arg_cnt = max(len(arg_names), INFINITE_ARG_COUNT)
        else:
            min_arg_cnt = min(defn.arg_cnt for defn in def_list)  # type: ignore
            max_arg_cnt = max(defn.arg_cnt for defn in def_list)  # type: ignore
            arg_names = max((defn.arg_names or () for defn in def_list), key=lambda names: len(names))

        # generate arg types
        arg_type_lists: dict[int, set[DataType]] = defaultdict(set)  # {<ind>: [<possible types>]}
        for defn in def_list:
            tr_arg_cnt = defn.arg_cnt if defn.arg_cnt is not None else max_arg_cnt
            for arg_matcher in defn.argument_types or ():
                # each variant has its own possible set of arg types
                for ind in range(tr_arg_cnt):
                    poss_types = arg_matcher.get_possible_arg_types_at_pos(ind, total=tr_arg_cnt)
                    arg_type_lists[ind] |= poss_types

                    for arg_type in poss_types:
                        # do not show const types if non-const type is accepted for this argument
                        const_type = arg_type.const_type
                        non_const_type = arg_type.non_const_type
                        if const_type in arg_type_lists[ind] and non_const_type in arg_type_lists[ind]:
                            arg_type_lists[ind].remove(const_type)

        # generate arg names if they are not defined
        if not arg_names:  # single anonymous arg
            if max_arg_cnt == 1:
                arg_names = ["value"]
            else:  # multiple anonymous args
                arg_names = ["arg_{}".format(i + 1) for i in range(max_arg_cnt)]

        # patch all type sets with implicit INTEGERs if FLOAT is accepted
        for types in arg_type_lists.values():
            for dtype in types.copy():
                if dtype in EXPANDABLE_TYPES:
                    types |= EXPANDABLE_TYPES[dtype]

        return [
            FuncArg(
                name=arg_names[i],
                types=arg_type_lists[i],
                optional_level=max(0, i + 1 - min_arg_cnt),
            )
            for i in range(max_arg_cnt)
        ]

lib/dl_formula_ref/dl_formula_ref/registry/arg_extractor.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    List,
    Optional,
    Sequence,
    Set,
)

import attr

from dl_formula.core.dialect import DialectCombo
from dl_formula_ref.registry.aliased_res import (
    AliasedResourceRegistryBase,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.arg_extractor import DefaultArgumentExtractor
from dl_formula_ref.registry.dialect_extractor import DefaultDialectExtractor
from dl_formula_ref.registry.impl_selector import DefaultImplementationSelector
from dl_formula_ref.registry.naming import DefaultNamingProvider
from dl_formula_ref.registry.note_extractor import DefaultNoteExtractor
from dl_formula_ref.registry.return_extractor import DefaultReturnTypeExtractor
from dl_formula_ref.registry.scopes import SCOPES_DEFAULT
from dl_formula_ref.registry.signature_gen import DefaultSignatureGenerator
from dl_formula_ref.registry.text import ParameterizedText


if TYPE_CHECKING:
    import dl_formula_ref.registry.arg_base as _arg_base
    from dl_formula_ref.registry.dialect_base import DialectExtractorBase
    from dl_formula_ref.registry.env import GenerationEnvironment
    from dl_formula_ref.registry.example_base import ExampleBase
    from dl_formula_ref.registry.impl_selector_base import ImplementationSelectorBase
    from dl_formula_ref.registry.impl_spec import FunctionImplementationSpec
    from dl_formula_ref.registry.naming import NamingProviderBase
    from dl_formula_ref.registry.note import (
        Note,
        ParameterizedNote,
    )
    import dl_formula_ref.registry.note_extr_base as _note_extr_base
    import dl_formula_ref.registry.return_base as _return_base
    import dl_formula_ref.registry.signature_base as _signature_base


@attr.s(frozen=True)
class FunctionDocCategory:
    _name: str = attr.ib(kw_only=True)
    _description: str = attr.ib(kw_only=True)
    _keywords: str = attr.ib(kw_only=True)
    _resources: AliasedResourceRegistryBase = attr.ib(kw_only=True, factory=SimpleAliasedResourceRegistry)

    @property
    def name(self) -> str:
        return self._name

    @property
    def description(self) -> str:
        return self._description

    @property
    def keywords(self) -> str:
        return self._keywords

    @property
    def resources(self) -> AliasedResourceRegistryBase:
        return self._resources


@attr.s(frozen=True)
class FunctionDocRegistryItem:
    _category: FunctionDocCategory = attr.ib(kw_only=True)
    _name: str = attr.ib(kw_only=True)
    _is_window: bool = attr.ib(kw_only=True, default=False)
    _description: str = attr.ib(kw_only=True)
    _scopes: int = attr.ib(kw_only=True, default=SCOPES_DEFAULT)
    _notes: List[Note] = attr.ib(kw_only=True, factory=list)
    _examples: Sequence[ExampleBase] = attr.ib(kw_only=True, factory=list)
    _impl_selector: ImplementationSelectorBase = attr.ib(kw_only=True, factory=DefaultImplementationSelector)
    _signature_gen: _signature_base.SignatureGeneratorBase = attr.ib(kw_only=True, factory=DefaultSignatureGenerator)
    _resources: AliasedResourceRegistryBase = attr.ib(kw_only=True, factory=SimpleAliasedResourceRegistry)
    _arg_extractor: _arg_base.ArgumentExtractorBase = attr.ib(kw_only=True, factory=DefaultArgumentExtractor)
    _note_extractor: _note_extr_base.NoteExtractorBase = attr.ib(kw_only=True, factory=DefaultNoteExtractor)
    _return_type_extractor: _return_base.ReturnTypeExtractorBase = attr.ib(
        kw_only=True, factory=DefaultReturnTypeExtractor
    )
    _dialect_extractor: DialectExtractorBase = attr.ib(kw_only=True, factory=DefaultDialectExtractor)
    _naming_provider: NamingProviderBase = attr.ib(kw_only=True, factory=DefaultNamingProvider)

    @property
    def category(self) -> FunctionDocCategory:
        return self._category

    @property
    def category_name(self) -> str:
        return self._category.name

    @property
    def name(self) -> str:
        return self._name

    @property
    def is_window(self) -> bool:
        return self._is_window

    @property
    def internal_name(self) -> str:
        return self._naming_provider.get_internal_name(self)

    def get_title(self, locale: str) -> str:
        return self._naming_provider.get_title(self, locale=locale)

    def get_short_title(self, locale: str) -> str:
        return self._naming_provider.get_short_title(self, locale=locale)

    @property
    def description(self) -> str:
        return self._description

    @property
    def resources(self) -> AliasedResourceRegistryBase:
        return self._resources

    @property
    def all_resources(self) -> AliasedResourceRegistryBase:
        return self._category.resources + self._resources

    def get_explicit_notes(self, env: GenerationEnvironment) -> List[Note]:
        return [note for note in self._notes if note.scopes & env.scopes == env.scopes]

    def get_examples(self, env: GenerationEnvironment) -> Sequence[ExampleBase]:
        return [example for example in self._examples if example.scopes & env.scopes == env.scopes]

    def get_notes(self, env: GenerationEnvironment) -> List[ParameterizedNote]:
        return self._note_extractor.get_notes(self, env=env)

    def get_implementation_specs(self, env: GenerationEnvironment) -> List[FunctionImplementationSpec]:
        return self._impl_selector.get_implementations(item=self, env=env)

    def get_signatures(self, env: GenerationEnvironment) -> _signature_base.FunctionSignatureCollection:
        return self._signature_gen.get_signatures(item=self, env=env)

    def get_return_type(self, env: GenerationEnvironment) -> ParameterizedText:
        return self._return_type_extractor.get_return_type(item=self, env=env)

    def get_dialects(self, env: GenerationEnvironment) -> Set[DialectCombo]:
        return self._dialect_extractor.get_dialects(item=self, env=env)

    def _get_one_implementation_spec(self, env: GenerationEnvironment) -> Optional[FunctionImplementationSpec]:
        implementations = self.get_implementation_specs(env=env)
        if implementations:
            return implementations[0]
        return None

    def is_supported(self, env: GenerationEnvironment) -> bool:
        # Validate against own scopes
        if self._scopes & env.scopes != env.scopes:
            return False

        # Filter implementations by scopes
        impl = self._get_one_implementation_spec(env=env)
        scopes: int = 0
        if impl is not None:
            scopes = impl.scopes
        return scopes & env.scopes == env.scopes

    def get_args(self, env: GenerationEnvironment) -> List[_arg_base.FuncArg]:
        return self._arg_extractor.get_args(self, env=env)

lib/dl_formula_ref/dl_formula_ref/registry/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Set,
)


if TYPE_CHECKING:
    from dl_formula.core.dialect import DialectCombo
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment


class DialectExtractorBase(abc.ABC):
    @abc.abstractmethod
    def get_dialects(
        self,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> Set[DialectCombo]:
        raise NotImplementedError

lib/dl_formula_ref/dl_formula_ref/registry/dialect_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import TYPE_CHECKING

from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D
from dl_formula_ref.registry.dialect_base import DialectExtractorBase


if TYPE_CHECKING:
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment


# Lowest versions of all backends that are compatible with COMPENG. Filled from plugins
COMPENG_SUPPORT: set[DialectCombo] = set()
EXPAND_COMPENG = True


def register_compeng_support_dialects(any_dialects: frozenset[DialectCombo]) -> None:
    COMPENG_SUPPORT.update(any_dialects)


class DefaultDialectExtractor(DialectExtractorBase):
    def get_dialects(
        self,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> set[DialectCombo]:
        if item.category.name == "window":  # TODO: Maybe separate into two different classes?
            if EXPAND_COMPENG:
                return set(dialect for dialect in COMPENG_SUPPORT if dialect in env.supported_dialects)
            return set()

        def_list = item.get_implementation_specs(env=env)
        dialect_mask = D.EMPTY
        for defn in def_list:
            dialect_mask |= defn.dialects

        return {dialect for dialect in dialect_mask.to_list() if dialect in env.supported_dialects}

lib/dl_formula_ref/dl_formula_ref/registry/dialect_extractor.py
datalens-tech/datalens-backend
import attr

from dl_formula.core.dialect import DialectCombo


@attr.s(frozen=True)
class GenerationEnvironment:
    scopes: int = attr.ib(kw_only=True)
    supported_dialects: frozenset[DialectCombo] = attr.ib(kw_only=True)

lib/dl_formula_ref/dl_formula_ref/registry/env.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr

from dl_formula_ref.examples.config import (
    ExampleConfig,
    PredefinedExampleConfig,
)
from dl_formula_ref.examples.renderer import ExampleRenderer
from dl_formula_ref.registry.example_base import (
    DataExampleRendererConfig,
    ExampleBase,
)


@attr.s(frozen=True)
class SimpleExample(ExampleBase):
    _text: str = attr.ib()

    def render(self, func_name: str, locale: str, config: DataExampleRendererConfig, under_cut: bool) -> str:
        return "```\n" f"{self._text.strip()}\n" "```"


@attr.s(frozen=True)
class DataExampleBase(ExampleBase):
    _example_config: ExampleConfig = attr.ib(kw_only=True)

    @property
    def example_config(self) -> ExampleConfig:
        return self._example_config

    def _get_renderer(self, locale: str, config: DataExampleRendererConfig) -> ExampleRenderer:
        return ExampleRenderer(
            jinja_env=config.jinja_env,
            example_template_filename=config.template_filename,
            locale=locale,
            storage_filename=config.example_data_filename,
        )


@attr.s(frozen=True)
class DataExample(DataExampleBase):
    @property
    def example_config(self) -> ExampleConfig:
        return self._example_config

    def render(self, func_name: str, locale: str, config: DataExampleRendererConfig, under_cut: bool) -> str:
        example_renderer = self._get_renderer(locale=locale, config=config)
        text = example_renderer.render_example_from_storage(
            example=self._example_config,
            func_name=func_name,
            under_cut=under_cut,
        )
        return text


@attr.s(frozen=True)
class PredefinedDataExample(DataExampleBase):
    _example_config: PredefinedExampleConfig = attr.ib(kw_only=True)

    def render(self, func_name: str, locale: str, config: DataExampleRendererConfig, under_cut: bool) -> str:
        example_renderer = self._get_renderer(locale=locale, config=config)
        text = example_renderer.render_example(
            example=self._example_config,
            result_table=self._example_config.result_table,
            under_cut=under_cut,
        )
        return text

lib/dl_formula_ref/dl_formula_ref/registry/example.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc

import attr
import jinja2

from dl_formula_ref.registry.scopes import SCOPES_DEFAULT


@attr.s
class DataExampleRendererConfig:
    jinja_env: jinja2.Environment = attr.ib(kw_only=True)
    template_filename: str = attr.ib(kw_only=True)
    example_data_filename: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class ExampleBase(abc.ABC):
    scopes: int = attr.ib(kw_only=True, default=SCOPES_DEFAULT)

    @abc.abstractmethod
    def render(self, func_name: str, locale: str, config: DataExampleRendererConfig, under_cut: bool) -> str:
        raise NotImplementedError

lib/dl_formula_ref/dl_formula_ref/registry/example_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from functools import lru_cache
from typing import TYPE_CHECKING

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.base import MultiVariantTranslation
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.registry import OPERATION_REGISTRY
from dl_formula_ref.registry.arg_extractor import INFINITE_ARG_COUNT
from dl_formula_ref.registry.impl_selector_base import ImplementationSelectorBase
from dl_formula_ref.registry.impl_spec import FunctionImplementationSpec


if TYPE_CHECKING:
    from dl_formula_ref.registry.base import FunctionDocRegistryItem
    from dl_formula_ref.registry.env import GenerationEnvironment


def _is_deprecated(defn: FunctionImplementationSpec) -> bool:
    return bool(defn.return_flags & ContextFlag.DEPRECATED)


def _make_spec_from_implementation(impl: MultiVariantTranslation) -> FunctionImplementationSpec:
    dialects = D.EMPTY
    for variant in impl.get_variants():
        dialects |= variant.dialects

    assert impl.name is not None
    impl_spec = FunctionImplementationSpec(
        name=impl.name,
        arg_cnt=impl.arg_cnt,
        arg_names=impl.arg_names or (),
        argument_types=impl.argument_types or (),
        return_type=impl.return_type,
        scopes=impl.scopes,
        dialects=dialects,
        return_flags=impl.return_flags,
    )
    return impl_spec


@lru_cache
def _get_implementation_map(env: GenerationEnvironment) -> dict[tuple[str, bool], list[FunctionImplementationSpec]]:
    result: dict[tuple[str, bool], list[FunctionImplementationSpec]] = defaultdict(list)
    for key, impl in OPERATION_REGISTRY.items():
        assert isinstance(impl, MultiVariantTranslation)
        impl_spec = _make_spec_from_implementation(impl)
        if _is_deprecated(impl_spec) or impl_spec.scopes & env.scopes != env.scopes:
            continue
        result[(key.name, key.is_window)].append(impl_spec)
    return result


class EmptyImplementationSelector(ImplementationSelectorBase):
    def get_implementations(
        self,
        item: FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> list[FunctionImplementationSpec]:
        return []


class DefaultImplementationSelector(ImplementationSelectorBase):
    def get_implementations(
        self,
        item: FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> list[FunctionImplementationSpec]:
        return _get_implementation_map(env=env)[(item.name, item.is_window)]


class ArgAwareImplementationSelector(ImplementationSelectorBase):
    def __init__(self, exp_arg_types: dict[int, set[DataType]]):
        self.exp_arg_types = exp_arg_types

    def get_implementations(
        self, item: FunctionDocRegistryItem, env: GenerationEnvironment
    ) -> list[FunctionImplementationSpec]:
        result: list[FunctionImplementationSpec] = []
        for impl in _get_implementation_map(env=env)[(item.name, item.is_window)]:
            actual_arg_types: dict[int, set[DataType]] = defaultdict(set)
            if impl.argument_types is None:
                continue
            for arg_type_matcher in impl.argument_types:
                for pos in self.exp_arg_types.keys():
                    actual_arg_types[pos] |= arg_type_matcher.get_possible_arg_types_at_pos(
                        pos=pos, total=impl.arg_cnt or INFINITE_ARG_COUNT
                    )
            for pos, arg_types in self.exp_arg_types.items():
                if pos not in actual_arg_types or pos in actual_arg_types and not (actual_arg_types[pos] & arg_types):
                    break
            else:
                result.append(impl)
        return result

lib/dl_formula_ref/dl_formula_ref/registry/impl_selector.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    List,
)


if TYPE_CHECKING:
    from dl_formula_ref.registry.base import FunctionDocRegistryItem
    from dl_formula_ref.registry.env import GenerationEnvironment
    from dl_formula_ref.registry.impl_spec import FunctionImplementationSpec


class ImplementationSelectorBase(abc.ABC):
    @abc.abstractmethod
    def get_implementations(
        self,
        item: FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> List[FunctionImplementationSpec]:
        raise NotImplementedError

lib/dl_formula_ref/dl_formula_ref/registry/impl_selector_base.py
datalens-tech/datalens-backend
from typing import (
    Optional,
    Sequence,
)

import attr

from dl_formula.core.dialect import DialectCombo
from dl_formula.definitions.args import ArgTypeMatcher
from dl_formula.definitions.type_strategy import TypeStrategy


@attr.s(frozen=True)
class FunctionImplementationSpec:
    name: str = attr.ib(kw_only=True)
    arg_cnt: Optional[int] = attr.ib(kw_only=True)
    arg_names: Sequence[str] = attr.ib(kw_only=True)
    argument_types: Sequence[ArgTypeMatcher] = attr.ib(kw_only=True)
    return_type: TypeStrategy = attr.ib(kw_only=True)
    scopes: int = attr.ib(kw_only=True)
    dialects: DialectCombo = attr.ib(kw_only=True)
    return_flags: int = attr.ib(kw_only=True)

lib/dl_formula_ref/dl_formula_ref/registry/impl_spec.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    Optional,
)

import attr

from dl_formula_ref.i18n.registry import get_localizer
from dl_formula_ref.texts import FUNCTION_CATEGORY_TAG


if TYPE_CHECKING:
    import dl_formula_ref.registry.base as _registry_base


def translate(text: Optional[str], locale: str) -> str:
    # Note: for `text=None` returns an empty string `''`
    gtext = get_localizer(locale).translate
    return gtext(text) if text else ""


class NamingProviderBase(abc.ABC):
    @abc.abstractmethod
    def get_internal_name(self, item: _registry_base.FunctionDocRegistryItem) -> str:
        raise NotImplementedError()

    @abc.abstractmethod
    def get_title(self, item: _registry_base.FunctionDocRegistryItem, locale: str) -> str:
        raise NotImplementedError()

    @abc.abstractmethod
    def get_short_title(self, item: _registry_base.FunctionDocRegistryItem, locale: str) -> str:
        raise NotImplementedError()


class DefaultNamingProvider(NamingProviderBase):
    def get_internal_name(self, item: _registry_base.FunctionDocRegistryItem) -> str:
        return item.name.upper()

    def get_title(self, item: _registry_base.FunctionDocRegistryItem, locale: str) -> str:
        return item.name.upper()

    def get_short_title(self, item: _registry_base.FunctionDocRegistryItem, locale: str) -> str:
        return item.name.upper()


@attr.s
class CustomNamingProvider(DefaultNamingProvider):
    _internal_name: Optional[str] = attr.ib(kw_only=True, default=None)
    _title: Optional[str] = attr.ib(kw_only=True, default=None)
    _short_title: Optional[str] = attr.ib(kw_only=True, default=None)

    def get_internal_name(self, item: _registry_base.FunctionDocRegistryItem) -> str:
        if self._internal_name is not None:
            return self._internal_name
        return super().get_internal_name(item)

    def get_title(self, item: _registry_base.FunctionDocRegistryItem, locale: str) -> str:
        if self._title is not None:
            return translate(self._title, locale=locale)
        return super().get_title(item, locale=locale)

    def get_short_title(self, item: _registry_base.FunctionDocRegistryItem, locale: str) -> str:
        if self._short_title is not None:
            return translate(self._short_title, locale=locale)
        if self._title is not None:
            return translate(self._title, locale=locale)
        return super().get_short_title(item, locale=locale)


class CategoryPostfixNamingProvider(CustomNamingProvider):
    def get_title(self, item: _registry_base.FunctionDocRegistryItem, locale: str) -> str:
        if self._title is not None:
            return translate(self._title, locale=locale)
        return f"{item.name.upper()} ({translate(FUNCTION_CATEGORY_TAG[item.category_name], locale=locale)})"

lib/dl_formula_ref/dl_formula_ref/registry/naming.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import Enum
from typing import (
    Any,
    List,
    NamedTuple,
)

import attr

from dl_formula_ref.registry.scopes import SCOPES_DEFAULT
from dl_formula_ref.registry.text import ParameterizedText


class NoteLevel(Enum):
    info = "info"
    warning = "warning"
    alert = "alert"


class Note(NamedTuple):
    text: str
    level: NoteLevel = NoteLevel.info
    formatting: bool = True
    scopes: int = SCOPES_DEFAULT

    def format(self, *args: Any, **kwargs: Any) -> Note:
        return self._replace(text=self.text.format(*args, **kwargs))


class NoteType(Enum):
    REGULAR = "REGULAR"
    ARG_RESTRICTION = "ARG_RESTRICTION"


@attr.s(frozen=True)
class ParameterizedNote:
    param_text: ParameterizedText = attr.ib(kw_only=True)
    level: NoteLevel = attr.ib(kw_only=True, default=NoteLevel.info)
    formatting: bool = attr.ib(kw_only=True, default=True)
    type: NoteType = attr.ib(kw_only=True, default=NoteType.REGULAR)

    @property
    def text(self) -> str:
        return self.param_text.format()

    def clone(self, **kwargs: Any) -> ParameterizedNote:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True)
class CrosslinkNote(ParameterizedNote):
    ref_list: List[str] = attr.ib(kw_only=True, default=list())

lib/dl_formula_ref/dl_formula_ref/registry/note.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    TYPE_CHECKING,
    List,
)


if TYPE_CHECKING:
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment
    from dl_formula_ref.registry.note import ParameterizedNote


class NoteExtractorBase(abc.ABC):
    @abc.abstractmethod
    def get_notes(
        self,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> List[ParameterizedNote]:
        raise NotImplementedError

lib/dl_formula_ref/dl_formula_ref/registry/note_extr_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    List,
)

from dl_formula_ref.registry.arg_common import TypeStrategyInspector
from dl_formula_ref.registry.note import (
    NoteType,
    ParameterizedNote,
)
from dl_formula_ref.registry.note_extr_base import NoteExtractorBase
from dl_formula_ref.registry.text import ParameterizedText
from dl_formula_ref.texts import CONST_TYPE_NOTE


if TYPE_CHECKING:
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment


class DefaultNoteExtractor(NoteExtractorBase):
    def get_notes(
        self,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> List[ParameterizedNote]:
        # explicitly defined notes
        notes = [
            ParameterizedNote(
                param_text=ParameterizedText(text=note.text),
                level=note.level,
                formatting=note.formatting,
                type=NoteType.REGULAR,
            )
            for note in item.get_explicit_notes(env=env)
        ]
        args = item.get_args(env=env)

        ts_insp = TypeStrategyInspector()

        # note about args that have to be of the same type (deduced from the function's return type strategy)
        type_info = ts_insp.get_return_type_and_arg_type_note(item, env=env)
        arg_type_note = type_info.arg_note
        if arg_type_note is not None:
            notes.append(
                ParameterizedNote(
                    param_text=arg_type_note,
                    type=NoteType.ARG_RESTRICTION,
                )
            )

        # note about args that must be const
        if any(a.is_const for a in args):
            notes.insert(
                0,
                ParameterizedNote(
                    param_text=ParameterizedText(
                        text=CONST_TYPE_NOTE,
                        params=dict(args=", ".join([f"`{a.name}`" for a in args if a.is_const])),
                    ),
                    type=NoteType.ARG_RESTRICTION,
                ),
            )
        return notes

lib/dl_formula_ref/dl_formula_ref/registry/note_extractor.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
    Iterable,
    NamedTuple,
    Tuple,
)

import attr

from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.env import GenerationEnvironment


class RefFunctionKey(NamedTuple):
    name: str
    category_name: str

    @classmethod
    def normalized(cls, name: str, category_name: str) -> RefFunctionKey:
        return RefFunctionKey(name=name.lower(), category_name=category_name.lower())


@attr.s
class FunctionReferenceRegistry:
    _registry: Dict[RefFunctionKey, FunctionDocRegistryItem] = attr.ib(kw_only=True, factory=dict)

    def add_item(self, item: FunctionDocRegistryItem) -> None:
        key = RefFunctionKey.normalized(name=item.name, category_name=item.category.name)
        self._registry[key] = item

    def items(self) -> Iterable[Tuple[RefFunctionKey, FunctionDocRegistryItem]]:
        return sorted(self._registry.items())

    def __contains__(self, key: Any) -> bool:
        if not isinstance(key, RefFunctionKey):
            return False
        return key in self._registry

    def __getitem__(self, key: Any) -> FunctionDocRegistryItem:
        return self._registry[key]

    def limit(self, env: GenerationEnvironment) -> FunctionReferenceRegistry:
        limited_reg = FunctionReferenceRegistry()
        for _key, item in self.items():
            if item.is_supported(env=env):
                limited_reg.add_item(item)

        return limited_reg


FUNC_REFERENCE_REGISTRY = FunctionReferenceRegistry()

lib/dl_formula_ref/dl_formula_ref/registry/registry.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import TYPE_CHECKING


if TYPE_CHECKING:
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment
    from dl_formula_ref.registry.text import ParameterizedText


class ReturnTypeExtractorBase(abc.ABC):
    @abc.abstractmethod
    def get_return_type(
        self,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> ParameterizedText:
        raise NotImplementedError

lib/dl_formula_ref/dl_formula_ref/registry/return_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import TYPE_CHECKING

from dl_formula_ref.registry.arg_common import TypeStrategyInspector
from dl_formula_ref.registry.return_base import ReturnTypeExtractorBase


if TYPE_CHECKING:
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment
    from dl_formula_ref.registry.text import ParameterizedText


class DefaultReturnTypeExtractor(ReturnTypeExtractorBase):
    def get_return_type(
        self,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> ParameterizedText:
        ts_insp = TypeStrategyInspector()
        type_info = ts_insp.get_return_type_and_arg_type_note(item, env=env)
        return type_info.ret_type_str

lib/dl_formula_ref/dl_formula_ref/registry/return_extractor.py
datalens-tech/datalens-backend
from dl_formula.definitions.scope import Scope


SCOPES_BASE = Scope.DOCUMENTED | Scope.STABLE
SCOPES_DEFAULT = SCOPES_BASE

lib/dl_formula_ref/dl_formula_ref/registry/scopes.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from enum import (
    Enum,
    unique,
)
from typing import (
    TYPE_CHECKING,
    List,
)

import attr


if TYPE_CHECKING:
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment


@attr.s(frozen=True)
class FunctionSignature:
    title: str = attr.ib(kw_only=True, default="")
    body: str = attr.ib(kw_only=True)
    description: list[str] = attr.ib(kw_only=True, factory=list)


@unique
class SignaturePlacement(Enum):
    compact = "compact"
    tabbed = "tabbed"


@attr.s(frozen=True)
class FunctionSignatureCollection:
    signatures: List[FunctionSignature] = attr.ib(kw_only=True)
    placement_mode: SignaturePlacement = attr.ib(kw_only=True, default=SignaturePlacement.compact)


@attr.s(frozen=True)
class SignatureGeneratorBase(abc.ABC):
    @abc.abstractmethod
    def get_signatures(
        self,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> FunctionSignatureCollection:
        raise NotImplementedError

lib/dl_formula_ref/dl_formula_ref/registry/signature_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    List,
    Sequence,
)

import attr

from dl_formula.inspect.function import (
    requires_grouping,
    supports_bfb,
    supports_extensions,
    supports_grouping,
    supports_ignore_dimensions,
    supports_lod,
    supports_ordering,
)
from dl_formula_ref.registry.signature_base import (
    FunctionSignature,
    FunctionSignatureCollection,
    SignatureGeneratorBase,
    SignaturePlacement,
)
from dl_formula_ref.texts import (
    SIGNATURE_DESC_EXTENDED_HEADER,
    SIGNATURE_TITLE_EXTENDED,
    SIGNATURE_TITLE_STANDARD,
)


if TYPE_CHECKING:
    import dl_formula_ref.registry.arg_base as _arg_base
    import dl_formula_ref.registry.base as _registry_base
    from dl_formula_ref.registry.env import GenerationEnvironment


@attr.s(frozen=True)
class StaticSignatureGenerator(SignatureGeneratorBase):
    signatures: List[FunctionSignature] = attr.ib(kw_only=True)

    def get_signatures(
        self,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> FunctionSignatureCollection:
        return FunctionSignatureCollection(
            signatures=self.signatures,
        )


@attr.s(frozen=True)
class SignatureTemplate:
    title: str = attr.ib(kw_only=True, default="")
    body: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class TemplatedSignatureGenerator(SignatureGeneratorBase):
    """
    Generates signatures from templates and the names of the function's arguments
    """

    _signature_templates: Sequence[SignatureTemplate] = attr.ib(kw_only=True)
    _placement_mode: SignaturePlacement = attr.ib(kw_only=True, default=SignaturePlacement.compact)

    def get_signatures(
        self,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> FunctionSignatureCollection:
        func_name = item.name.upper()
        args = item.get_args(env=env)
        signatures = [
            FunctionSignature(
                title=template.title,
                body=template.body.strip().format(func_name, *(arg.name for arg in args)),
            )
            for template in self._signature_templates
        ]
        sig_collection = FunctionSignatureCollection(placement_mode=self._placement_mode, signatures=signatures)
        return sig_collection


@attr.s(frozen=True)
class DefaultSignatureGenerator(SignatureGeneratorBase):
    """
    Generates signatures from the properties of the function and its argument names
    """

    _placement_mode: SignaturePlacement = attr.ib(kw_only=True, default=SignaturePlacement.tabbed)

    def _get_signature_from_args(
        self,
        func_name: str,
        args: List[_arg_base.FuncArg],
        inf_args: bool,
        is_window: bool,
        is_extended_syntax: bool,
        category_name: str,
    ) -> FunctionSignature:
        if not args:  # to avoid extra spaces with 0-argument functions
            return FunctionSignature(body="{}()".format(func_name))

        args_str = ""
        cur_opt_level = 0
        for i, arg in enumerate(args):
            if arg.optional_level > cur_opt_level:
                args_str += " [ " * (arg.optional_level - cur_opt_level)
            elif arg.optional_level < cur_opt_level:
                args_str += " ] " * (cur_opt_level - arg.optional_level)
            cur_opt_level = arg.optional_level
            if i > 0:
                args_str += ", "
            args_str += arg.name
        if inf_args:
            args_str += " [ , ... ]"
        if cur_opt_level > 0:
            args_str += " ]" * cur_opt_level

        modifiers: List[str] = []
        description: list[str] = []
        if requires_grouping(func_name, is_window=is_window):
            modifiers.append("TOTAL | WITHIN ... | AMONG ...")
            description.append(f"- {{category:{category_name}#syntax-grouping:TOTAL, WITHIN, AMONG}}")
        title: str
        if is_extended_syntax:
            title = SIGNATURE_TITLE_EXTENDED
            if supports_lod(func_name, is_window=is_window):
                modifiers.append("[ FIXED ... | INCLUDE ... | EXCLUDE ... ]")
                description.append(f"- {{category:{category_name}#syntax-lod:FIXED, INCLUDE, EXCLUDE}}")
            if supports_grouping(func_name, is_window=is_window) and not requires_grouping(
                func_name, is_window=is_window
            ):
                # If grouping is required (e.g. window SUM),
                # its non-optional form has already been added to modifiers
                modifiers.append("[ TOTAL | WITHIN ... | AMONG ... ]")
                description.append(f"- {{category:{category_name}#syntax-grouping:TOTAL, WITHIN, AMONG}}")
            if supports_ordering(func_name, is_window=is_window):
                modifiers.append("[ ORDER BY ... ]")
                description.append(f"- {{category:{category_name}#syntax-order-by:ORDER BY}}")
            if supports_bfb(func_name, is_window=is_window):
                modifiers.append("[ BEFORE FILTER BY ... ]")
                description.append(f"- {{category:{category_name}#syntax-before-filter-by:BEFORE FILTER BY}}")
            if supports_ignore_dimensions(func_name, is_window=is_window):
                modifiers.append("[ IGNORE DIMENSIONS ... ]")
                description.append(f"- {{category:{category_name}#syntax-ignore-dimensions:IGNORE DIMENSIONS}}")

        else:
            title = SIGNATURE_TITLE_STANDARD

        if description:
            # Something was added to the description, so add the header
            description.insert(0, SIGNATURE_DESC_EXTENDED_HEADER)

        signature: str
        if not args_str and not modifiers:
            signature = f"{func_name}()"
        else:
            if not modifiers:
                signature = f"{func_name}( {args_str} )"
            else:
                indent = " " * (len(func_name) + 2)
                modifiers_str = "\n".join([f"{indent}{mod}" for mod in modifiers])
                signature = f"{func_name}( {args_str}\n{modifiers_str}\n{indent[:-2]})"

        return FunctionSignature(title=title, body=signature, description=description)

    def get_signatures(
        self,
        item: _registry_base.FunctionDocRegistryItem,
        env: GenerationEnvironment,
    ) -> FunctionSignatureCollection:
        func_name = item.name.upper()
        inf_args = False
        for defn in item.get_implementation_specs(env=env):
            inf_args = inf_args or (defn.arg_cnt is None)

        is_extended_syntax_values = [False]
        if supports_extensions(func_name, is_window=item.is_window):
            is_extended_syntax_values.append(True)

        args = item.get_args(env=env)
        signatures: List[FunctionSignature] = [
            self._get_signature_from_args(
                func_name=func_name,
                args=args,
                inf_args=inf_args,
                is_window=item.is_window,
                is_extended_syntax=is_extended_syntax,
                category_name=item.category_name,
            )
            for is_extended_syntax in is_extended_syntax_values
        ]
        sig_collection = FunctionSignatureCollection(
            placement_mode=self._placement_mode,
            signatures=signatures,
        )
        return sig_collection

lib/dl_formula_ref/dl_formula_ref/registry/signature_gen.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Dict,
)

import attr


@attr.s(frozen=True)
class ParameterizedText:
    text: str = attr.ib(kw_only=True)
    params: Dict[str, str] = attr.ib(kw_only=True, factory=dict)

    def format(self) -> str:
        if not self.params:
            return self.text
        return self.text.format(**self.params)

    def __str__(self) -> str:
        return self.format()

    def clone(self, **kwargs: Any) -> ParameterizedText:
        return attr.evolve(self, **kwargs)

lib/dl_formula_ref/dl_formula_ref/registry/text.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import List

from dl_formula_ref.functions.aggregation import FUNCTIONS_AGGREGATION
from dl_formula_ref.functions.array import FUNCTIONS_ARRAY
from dl_formula_ref.functions.date import FUNCTIONS_DATE
from dl_formula_ref.functions.logical import FUNCTIONS_LOGICAL
from dl_formula_ref.functions.markup import FUNCTIONS_MARKUP
from dl_formula_ref.functions.mathematical import FUNCTIONS_MATHEMATICAL
from dl_formula_ref.functions.operator import FUNCTIONS_OPERATOR
from dl_formula_ref.functions.string import FUNCTIONS_STRING
from dl_formula_ref.functions.time_series import FUNCTIONS_TIME_SERIES
from dl_formula_ref.functions.type_conversion import FUNCTIONS_TYPE_CONVERSION
from dl_formula_ref.functions.window import FUNCTIONS_WINDOW
from dl_formula_ref.registry.base import FunctionDocRegistryItem
from dl_formula_ref.registry.registry import FUNC_REFERENCE_REGISTRY


EXPLICITLY_DEFINED_FUNCTIONS: List[FunctionDocRegistryItem] = [
    *FUNCTIONS_AGGREGATION,
    *FUNCTIONS_ARRAY,
    *FUNCTIONS_DATE,
    *FUNCTIONS_LOGICAL,
    *FUNCTIONS_MARKUP,
    *FUNCTIONS_MATHEMATICAL,
    *FUNCTIONS_OPERATOR,
    *FUNCTIONS_STRING,
    *FUNCTIONS_TIME_SERIES,
    *FUNCTIONS_TYPE_CONVERSION,
    *FUNCTIONS_WINDOW,
]


def populate_registry_from_definitions() -> None:
    for expl_item in EXPLICITLY_DEFINED_FUNCTIONS:
        FUNC_REFERENCE_REGISTRY.add_item(expl_item)

lib/dl_formula_ref/dl_formula_ref/registry/tools.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    Iterable,
    List,
    NamedTuple,
    Optional,
    Set,
)

import attr


if TYPE_CHECKING:
    from dl_formula.core.dialect import DialectCombo
    from dl_formula_ref.audience import Audience
    from dl_formula_ref.primitives import ParameterizedNote
    from dl_formula_ref.registry.arg_base import FuncArg
    from dl_formula_ref.registry.signature_base import FunctionSignatureCollection


class LocalizedFuncArg(NamedTuple):
    arg: FuncArg
    locale: str
    human_type: str

    @property
    def name(self) -> str:
        return self.arg.name


class RenderedFunc(NamedTuple):
    # The content (e.g. cross-reference links) is path dependent,
    # and so is the rendered function, so tha t is why it contains its own path
    file_path: str
    name: str
    title: str
    short_title: str
    internal_name: str
    category_name: str
    human_category: str
    category_description: str
    category_keywords: List[str]
    args: List[LocalizedFuncArg]
    dialects: Set[DialectCombo]
    human_dialects: List[str]
    description: str
    top_notes: List[ParameterizedNote]
    bottom_notes: List[ParameterizedNote]
    return_type: str
    examples: List[str]
    signature_coll: FunctionSignatureCollection
    locale: str
    crosslink_note: Optional[ParameterizedNote]

    @property
    def const_args(self) -> List[LocalizedFuncArg]:
        return [a for a in self.args if a.arg.is_const]

    @property
    def multi_signature(self) -> bool:
        return len(self.signature_coll.signatures) > 1

    @property
    def category_description_short(self) -> str:
        first_sentence = self.category_description.split(".", 1)[0]
        if not first_sentence:
            return ""
        return f"{first_sentence}."


@attr.s(frozen=True)
class RenderedMultiAudienceFunc:
    name: str = attr.ib(kw_only=True)
    category_name: str = attr.ib(kw_only=True)
    category_description: str = attr.ib(kw_only=True)
    category_description_short: str = attr.ib(kw_only=True)
    category_keywords: list[str] = attr.ib(kw_only=True)
    file_path: str = attr.ib(kw_only=True)
    _aud_dict: dict[Audience, RenderedFunc] = attr.ib(kw_only=True)

    def items(self) -> Iterable[tuple[Audience, RenderedFunc]]:
        items: list[tuple[Audience, RenderedFunc]] = list(self._aud_dict.items())
        return sorted(items, key=lambda item: item[0].name)

lib/dl_formula_ref/dl_formula_ref/rendered.py
datalens-tech/datalens-backend
"""
Generate structures used for creating documentation
"""

from __future__ import annotations

from collections import defaultdict
import os.path
from typing import (
    TYPE_CHECKING,
    Collection,
    Mapping,
    Optional,
    Union,
)

import attr

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D
from dl_formula_ref.audience import Audience
from dl_formula_ref.i18n.registry import get_localizer
from dl_formula_ref.paths import (
    FileLink,
    PathRenderer,
)
from dl_formula_ref.primitives import RawFunc
from dl_formula_ref.reference import FuncReference
from dl_formula_ref.registry.aliased_res import (
    AliasedResourceRegistryBase,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.note import (
    CrosslinkNote,
    ParameterizedNote,
)
from dl_formula_ref.registry.registry import RefFunctionKey
from dl_formula_ref.registry.signature_base import (
    FunctionSignature,
    FunctionSignatureCollection,
)
from dl_formula_ref.registry.text import ParameterizedText
from dl_formula_ref.rendered import (
    LocalizedFuncArg,
    RenderedFunc,
    RenderedMultiAudienceFunc,
)
from dl_formula_ref.rich_text.expander import MacroExpander
from dl_formula_ref.rich_text.helpers import get_human_data_type_list
from dl_formula_ref.rich_text.renderer import (
    MdRichTextRenderer,
    RichTextRenderEnvironment,
)
from dl_formula_ref.texts import (
    ALSO_IN_OTHER_CATEGORIES,
    ANY_DIALECTS,
    HIDDEN_DIALECTS,
    HUMAN_CATEGORIES,
    HUMAN_DIALECTS,
    DialectStyle,
)


if TYPE_CHECKING:
    from dl_formula_ref.registry.example_base import DataExampleRendererConfig


def translate(text: str, locale: str) -> str:
    gtext = get_localizer(locale).translate
    return gtext(text) if text else ""


def human_dialect(dialect: Union[DialectCombo, str], locale: str, style: DialectStyle = DialectStyle.simple) -> str:
    if isinstance(dialect, str):
        dialect = D.by_name(dialect.upper())

    return "{}".format(translate(HUMAN_DIALECTS[dialect].for_style(style), locale=locale))


def human_data_type(types: Collection[DataType], locale: str) -> str:
    result_types = get_human_data_type_list(types=types)
    raw_type_str = " | ".join(translate(t, locale=locale) for t in sorted(result_types))
    return f"`{raw_type_str}`"


def human_category(category: str, locale: str) -> str:
    return translate(HUMAN_CATEGORIES.get(category, category), locale=locale)


def human_dialects(dialects: set[DialectCombo], locale: str) -> list[str]:
    dialect_mask = D.EMPTY
    for dialect in dialects:
        if dialect == D.ANY:
            for any_dialect in ANY_DIALECTS:
                dialect_mask |= any_dialect
        else:
            dialect_mask |= dialect

    dialects_by_db = defaultdict(list)
    for dialect in dialect_mask.to_list():
        dialects_by_db[dialect.single_bit.name].append(dialect)

    sorted_dialects = sorted(
        [sorted(dialects_for_db, key=lambda d: d.single_bit.version)[0] for dialects_for_db in dialects_by_db.values()],
        key=lambda d: d.common_name_and_version,
    )

    h_dialects: list[str] = [human_dialect(d, locale=locale) for d in sorted_dialects if d not in HIDDEN_DIALECTS]
    return h_dialects


@attr.s
class FuncRenderer:
    _func_ref: FuncReference = attr.ib(kw_only=True)
    _locale: str = attr.ib(kw_only=True)
    _example_rend_config: DataExampleRendererConfig = attr.attrib(kw_only=True)
    _path_renderer: PathRenderer = attr.ib(kw_only=True)
    _rt_renderer: MdRichTextRenderer = attr.ib(init=False, factory=MdRichTextRenderer)
    _block_conditions: Mapping[str, bool] = attr.ib(kw_only=True, factory=dict)
    _rt_env: RichTextRenderEnvironment = attr.ib(kw_only=True)
    _resource_overrides: AliasedResourceRegistryBase = attr.ib(kw_only=True, factory=SimpleAliasedResourceRegistry)

    @_rt_env.default
    def _rt_env_default(self) -> RichTextRenderEnvironment:
        return RichTextRenderEnvironment(block_conditions=self._block_conditions)

    @property
    def path_renderer(self) -> PathRenderer:
        return self._path_renderer

    def render_text(self, context_path: str, raw_func: RawFunc, text: Union[str, ParameterizedText]) -> str:
        if not text:
            return ""
        if isinstance(text, ParameterizedText):
            text = ParameterizedText(
                text=translate(text.text, locale=self._locale),
                params=text.params,
            ).format()
        translation_callable = get_localizer(self._locale).translate
        relative_path_renderer = self._path_renderer.child(context_path)
        resources = self._resource_overrides + raw_func.resources
        expander = MacroExpander(
            resources=resources,
            args=raw_func.args,
            translation_callable=translation_callable,
            func_link_provider=relative_path_renderer.get_func_link,
            cat_link_provider=relative_path_renderer.get_cat_link,
        )
        rich_text = expander.expand_text(text)
        result = self._rt_renderer.render(rich_text, env=self._rt_env)
        return result

    def get_counterpart_keys(self, func_key: RefFunctionKey) -> list[RefFunctionKey]:
        return [
            RefFunctionKey.normalized(name=key.name, category_name=key.category.name)
            for key in self._func_ref.filter(name=func_key.name.upper())
            if key.category.name != func_key.category_name
        ]

    def get_counterpart_crosslink_note_and_refs(self, func_key: RefFunctionKey) -> Optional[CrosslinkNote]:
        counterpart_keys = self.get_counterpart_keys(func_key=func_key)
        if not counterpart_keys:
            return None

        counterpart_ref = "{{ref:{category_name}/{func_name}:{human_category}}}"
        counterpart_refs_list = [
            counterpart_ref.format(
                category_name=key.category_name,
                func_name=func_key.name.upper(),
                human_category=HUMAN_CATEGORIES[key.category_name],
            )
            for key in counterpart_keys
        ]

        note = CrosslinkNote(
            param_text=ParameterizedText(
                text=ALSO_IN_OTHER_CATEGORIES,
                params={"func_name": func_key.name.upper()},
            ),
            ref_list=counterpart_refs_list,
        )

        return note

    def get_func_link(self, func_name: str, category_name: str, anchor_name: Optional[str] = None) -> FileLink:
        return self._path_renderer.get_func_link(
            func_name=func_name,
            category_name=category_name,
            anchor_name=anchor_name,
        )

    def get_cat_link(self, category_name: str, anchor_name: Optional[str] = None) -> FileLink:
        return self._path_renderer.get_cat_link(
            category_name=category_name,
            anchor_name=anchor_name,
        )

    def render_func(self, func_key: RefFunctionKey, raw_func: RawFunc) -> RenderedFunc:
        gtext = get_localizer(self._locale).translate
        func_file_path = self._path_renderer.get_func_path(func_key=func_key)
        context_path = os.path.dirname(func_file_path)

        def _render_text(text: Union[str, ParameterizedText]) -> str:
            return self.render_text(context_path=context_path, raw_func=raw_func, text=text)

        def _render_note(note: ParameterizedNote) -> ParameterizedNote:
            text = _render_text(note.param_text)
            return note.clone(param_text=ParameterizedText(text=text))

        def _render_crosslink_note(note: CrosslinkNote) -> ParameterizedNote:
            text = _render_text(note.param_text)
            refs_list = ",".join(_render_text(ref) for ref in note.ref_list)
            new_text = text + refs_list + "."
            return note.clone(param_text=ParameterizedText(text=new_text))

        def _make_keywords(keyword_list_str: str) -> list[str]:
            if not keyword_list_str:
                return []
            keyword_list_str = gtext(keyword_list_str)
            return [kw.strip() for kw in keyword_list_str.split(",")]

        crosslink_note = self.get_counterpart_crosslink_note_and_refs(func_key=func_key)

        # Place examples under cut only if there are more than one
        ex_under_cut = len(raw_func.examples) > 1
        examples = [
            ex.render(
                func_name=raw_func.internal_name,
                config=self._example_rend_config,
                locale=self._locale,
                under_cut=ex_under_cut,
            )
            for ex in raw_func.examples
        ]
        signature_coll = FunctionSignatureCollection(
            signatures=[
                FunctionSignature(
                    title=sign.title,
                    body=sign.body,
                    description=[_render_text(desc_part) for desc_part in sign.description],
                )
                for sign in raw_func.signature_coll.signatures
            ],
            placement_mode=raw_func.signature_coll.placement_mode,
        )

        return RenderedFunc(
            file_path=func_file_path,
            name=raw_func.name,
            internal_name=raw_func.internal_name,
            title=raw_func.get_title(locale=self._locale),
            short_title=raw_func.get_short_title(locale=self._locale),
            category_name=raw_func.category.name or "",
            human_category=human_category(category=raw_func.category.name, locale=self._locale),
            category_description=_render_text(raw_func.category.description),
            category_keywords=_make_keywords(raw_func.category.keywords),
            args=[
                LocalizedFuncArg(
                    arg=arg, locale=self._locale, human_type=human_data_type(types=arg.types, locale=self._locale)
                )
                for arg in raw_func.args
            ],
            description=_render_text(raw_func.description),
            dialects=raw_func.dialects,
            human_dialects=human_dialects(dialects=raw_func.dialects, locale=self._locale),
            top_notes=[_render_note(w) for w in raw_func.top_notes],
            bottom_notes=[_render_note(n) for n in raw_func.bottom_notes[:]],
            return_type=_render_text(raw_func.return_type),
            examples=examples,
            signature_coll=signature_coll,
            locale=self._locale,
            crosslink_note=_render_crosslink_note(crosslink_note) if crosslink_note is not None else None,
        )

    def render_multi_func_func(self, func_key: RefFunctionKey) -> RenderedMultiAudienceFunc:
        raw_multi_func = self._func_ref.get_func(func_key=func_key)
        func_file_path = self._path_renderer.get_func_path(func_key=func_key)

        rendered_aud_func_dict: dict[Audience, RenderedFunc] = {}
        for audience, raw_func in raw_multi_func.items():
            rendered_func = self.render_func(func_key=func_key, raw_func=raw_func)
            rendered_aud_func_dict[audience] = rendered_func

        one_rendered_func = next(iter(rendered_aud_func_dict.values()))

        rendered_multi_func = RenderedMultiAudienceFunc(
            aud_dict=rendered_aud_func_dict,
            name=raw_multi_func.name,
            category_name=one_rendered_func.category_name,
            category_description=one_rendered_func.category_description,
            category_description_short=one_rendered_func.category_description_short,
            category_keywords=one_rendered_func.category_keywords,
            file_path=func_file_path,
        )
        return rendered_multi_func

lib/dl_formula_ref/dl_formula_ref/rendering.py
datalens-tech/datalens-backend
from typing import (
    Dict,
    List,
    NamedTuple,
    Sequence,
)

import attr


class MacroReplacementKey(NamedTuple):
    start_pos: int
    end_pos: int

    @property
    def pos(self) -> slice:
        return slice(self.start_pos, self.end_pos)


@attr.s(frozen=True)
class BaseTextElement:
    pass


@attr.s(frozen=True)
class RichText(BaseTextElement):
    text: str = attr.ib()
    replacements: Dict[MacroReplacementKey, BaseTextElement] = attr.ib(kw_only=True, factory=dict)


@attr.s(frozen=True)
class ConditionalBlock(BaseTextElement):
    condition: str = attr.ib()
    rich_text: RichText = attr.ib()


@attr.s(frozen=True)
class NoteBlock(BaseTextElement):
    level: str = attr.ib()
    rich_text: RichText = attr.ib()


@attr.s(frozen=True)
class AudienceBlock(BaseTextElement):
    audience_types: list[str] = attr.ib()
    rich_text: RichText = attr.ib()


@attr.s(frozen=True)
class LinkTextElement(BaseTextElement):
    text: str = attr.ib(kw_only=True)
    url: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class TableTextElement(BaseTextElement):
    table_body: List[List[RichText]] = attr.ib(kw_only=True)


@attr.s(frozen=True)
class ListTextElement(BaseTextElement):
    items: Sequence[BaseTextElement] = attr.ib(kw_only=True)
    sep: str = attr.ib(kw_only=True)
    wrap: bool = attr.ib(kw_only=True, default=False)


@attr.s(frozen=True)
class TermTextElement(BaseTextElement):
    term: str = attr.ib(kw_only=True)
    wrap: bool = attr.ib(kw_only=True, default=True)


@attr.s(frozen=True)
class CodeSpanTextElement(BaseTextElement):
    text: str = attr.ib(kw_only=True)
    wrap: bool = attr.ib(kw_only=True)


@attr.s(frozen=True)
class ExtMacroTextElement(BaseTextElement):
    macro_name: str = attr.ib(kw_only=True)

lib/dl_formula_ref/dl_formula_ref/rich_text/elements.py
datalens-tech/datalens-backend
from functools import singledispatchmethod
import re
from typing import (
    Callable,
    Dict,
    List,
    Optional,
    Tuple,
)

import attr

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import get_dialect_from_str
from dl_formula_ref.registry.aliased_res import (
    AliasedResourceRegistryBase,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.arg_base import FuncArg
from dl_formula_ref.rich_text.elements import (
    AudienceBlock,
    BaseTextElement,
    CodeSpanTextElement,
    ConditionalBlock,
    ExtMacroTextElement,
    LinkTextElement,
    ListTextElement,
    MacroReplacementKey,
    NoteBlock,
    RichText,
    TableTextElement,
    TermTextElement,
)
from dl_formula_ref.rich_text.helpers import get_human_data_type_list
from dl_formula_ref.rich_text.macro import (
    DOUBLE_ARG_MACROS,
    LIST_MACROS,
    SINGLE_ARG_MACROS,
    ArgMacro,
    ArgNMacro,
    BaseMacro,
    CategoryMacro,
    DialectsMacro,
    ExtMacroMacro,
    LinkMacro,
    RefMacro,
    TableMacro,
    TextMacro,
    TypeMacro,
)
from dl_formula_ref.texts import (
    HUMAN_DIALECTS,
    DialectStyle,
)
from dl_i18n.localizer_base import Translatable


_MACRO_CHOICES = "|".join(("ref", "link", "text", "table", "dialects", "arg", "argn", "type", "macro", "category"))
# 3 kinds of macros:
#     {type: single_value}
#     {type: first_value: second_value}
#     {type: value_1 | value_2 | value_3 | ...}
SIMPLE_MACRO_RE = re.compile(
    r"\{\s*"
    r"(?P<type>" + _MACRO_CHOICES + r")"
    r"\s*:\s*"
    r"(?P<first_value>[^{}:|`]+)"
    r"(|:(?P<second_value>[^{}:|`]+)|(?P<list>(\|[^{}:|`]+)+))"
    r"\}"
)
_BLOCK_MACRO_CHOICES = "|".join(
    (
        "if",
        "note",
        "audience",
        "end",
    )
)
END_TAG = "end"
BLOCK_TAG_RE = re.compile(r"\{\s*(?P<tag_name>" + _BLOCK_MACRO_CHOICES + r")(\s+(?P<block_param>[^}]+))?\s*\}")


@attr.s
class MacroExpander:
    _resources: AliasedResourceRegistryBase = attr.ib(kw_only=True, factory=SimpleAliasedResourceRegistry)
    _func_link_provider: Optional[Callable[[str, Optional[str]], Tuple[str, str]]] = attr.ib(kw_only=True, default=None)
    _cat_link_provider: Optional[Callable[[str, Optional[str]], Tuple[str, str]]] = attr.ib(kw_only=True, default=None)
    _args: Optional[List[FuncArg]] = attr.ib(kw_only=True, default=None)
    _translation_callable: Callable[[str | Translatable], str] = attr.ib(kw_only=True, default=lambda s: s)

    def _get_raw_link_url_by_alias(self, alias: str) -> Optional[str]:
        if alias not in self._resources:
            return None

        return self._resources.get_link(alias).url

    def _get_raw_text_body_by_alias(self, alias: str) -> str:
        return self._resources.get_text(alias).body

    def _get_raw_table_body_by_alias(self, alias: str) -> list[list[str | Translatable]]:
        return self._resources.get_table(alias).table_body

    def _get_func_name_and_url(self, func_name: str, category_name: Optional[str] = None) -> Tuple[str, str]:
        assert self._func_link_provider is not None
        return self._func_link_provider(func_name, category_name)

    def _get_cat_name_and_url(self, category_name: str, anchor_name: Optional[str] = None) -> Tuple[str, str]:
        assert self._cat_link_provider is not None
        return self._cat_link_provider(category_name, anchor_name)

    def _translate_text(self, text: str | Translatable) -> str:
        if isinstance(text, str) and not text:
            return text
        return self._translation_callable(text)

    def expand_text(self, text: str | Translatable) -> RichText:
        trans_text = self._translate_text(text)
        replacements: Dict[MacroReplacementKey, BaseTextElement] = {}

        block_replacements = self._get_block_replacements(trans_text)
        replacements.update(block_replacements)
        macro_replacements = self._get_macro_replacements(trans_text, block_replacements)
        replacements.update(macro_replacements)

        return RichText(text=trans_text, replacements=replacements)

    def _make_block_macro(
        self,
        tag_name: str,
        block_param: Optional[str],
        block_text: str,
        nested_replacements: Dict[MacroReplacementKey, BaseTextElement],
    ) -> BaseTextElement:
        element: BaseTextElement

        if tag_name == "if":
            assert block_param is not None
            element = ConditionalBlock(
                condition=block_param,
                rich_text=RichText(text=block_text, replacements=nested_replacements),
            )
        elif tag_name == "note":
            block_param = block_param or "info"
            assert block_param is not None
            element = NoteBlock(
                level=block_param,
                rich_text=RichText(text=block_text, replacements=nested_replacements),
            )
        elif tag_name == "audience":
            block_param = block_param
            assert block_param is not None
            audiences = [item.strip() for item in block_param.split((","))]
            element = AudienceBlock(
                audience_types=audiences,
                rich_text=RichText(text=block_text, replacements=nested_replacements),
            )
        else:
            raise ValueError(f"Unknown block tag: {tag_name}")
        return element

    def _get_block_replacements(self, text: str) -> Dict[MacroReplacementKey, BaseTextElement]:
        block_replacements: Dict[MacroReplacementKey, BaseTextElement] = {}

        block_match: Optional[re.Match] = BLOCK_TAG_RE.search(text)
        while block_match is not None and block_match.group("tag_name") != END_TAG:
            tag_name: str = block_match.group("tag_name")
            block_param: Optional[str] = block_match.group("block_param")
            block_text_start_pos = block_match.end()
            nested_block_replacements = self._get_block_replacements(text[block_text_start_pos:])

            nested_end_pos: int = 0
            if nested_block_replacements:
                nested_end_pos = max([nest_repl_key.end_pos for nest_repl_key in nested_block_replacements])

            block_end_match = BLOCK_TAG_RE.search(text, block_text_start_pos + nested_end_pos)
            assert block_end_match is not None
            assert block_end_match.group("tag_name") == END_TAG
            repl_key = MacroReplacementKey(
                start_pos=block_match.start(),
                end_pos=block_end_match.end(),
            )
            nested_text = text[block_match.end() : block_end_match.start()]
            nested_macro_replacements = self._get_macro_replacements(nested_text, nested_block_replacements)

            block_replacement = self._make_block_macro(
                tag_name=tag_name,
                block_param=block_param,
                block_text=nested_text,
                nested_replacements={
                    **nested_block_replacements,
                    **nested_macro_replacements,
                },
            )
            block_replacements[repl_key] = block_replacement

            block_match = BLOCK_TAG_RE.search(text, block_end_match.end())

        return block_replacements

    def _get_macro_replacements(
        self,
        text: str,
        block_replacements: Dict[MacroReplacementKey, BaseTextElement],
    ) -> Dict[MacroReplacementKey, BaseTextElement]:
        replacements: Dict[MacroReplacementKey, BaseTextElement] = {}

        for macro_match in SIMPLE_MACRO_RE.finditer(text):
            start_pos = macro_match.start()
            for block_repl_key in block_replacements:
                if block_repl_key.start_pos <= start_pos < block_repl_key.end_pos:
                    # skip macros inside blocks that are already being replaced
                    continue
            repl_key = MacroReplacementKey(
                start_pos=macro_match.start(),
                end_pos=macro_match.end(),
            )
            macro = self._make_macro(macro_match)
            expanded_element = self.expand_macro(macro)
            replacements[repl_key] = expanded_element

        return replacements

    def _make_macro(self, macro_match: re.Match) -> BaseMacro:
        macro_type: str = macro_match.group("type").strip()
        first_value: str = macro_match.group("first_value").strip()
        list_str: Optional[str] = macro_match.group("list")
        second_value: Optional[str] = macro_match.group("second_value")

        if macro_type in LIST_MACROS:
            value_list = [first_value]
            if list_str is not None:
                value_list += [s.strip() for s in list_str.split("|")[1:]]
            return LIST_MACROS[macro_type](values=value_list)

        if macro_type in DOUBLE_ARG_MACROS:
            if second_value is not None:
                second_value = second_value.strip()
            return DOUBLE_ARG_MACROS[macro_type](arg=first_value, second_arg=second_value)

        return SINGLE_ARG_MACROS[macro_type](arg=first_value)

    @singledispatchmethod
    def expand_macro(self, macro: BaseMacro) -> BaseTextElement:
        raise TypeError(type(macro))

    @expand_macro.register
    def expand_macro_ref(self, macro: RefMacro) -> LinkTextElement:
        # In theory, we can easily support anchors here just like in category links
        assert macro.arg.count("/") <= 1, 'The format is "[category_name/]function_name", got "{}"'.format(macro.arg)
        if macro.arg.find("/") > -1:
            macro.category_name, macro.arg = macro.arg.split("/")
        func_name, url = self._get_func_name_and_url(func_name=macro.arg, category_name=macro.category_name)
        text: str = func_name
        if macro.second_arg is not None:
            text = self._translate_text(macro.second_arg)
        return LinkTextElement(text=text, url=url)

    @expand_macro.register
    def expand_macro_category(self, macro: CategoryMacro) -> LinkTextElement:
        assert macro.arg.count("#") <= 1, 'The format is "category_name[#anchor_name]", got "{}"'.format(macro.arg)
        if macro.arg.find("#") > -1:
            macro.arg, macro.anchor_name = macro.arg.split("#")
        cat_name, url = self._get_cat_name_and_url(category_name=macro.arg, anchor_name=macro.anchor_name)
        text: str = cat_name
        if macro.second_arg is not None:
            text = self._translate_text(macro.second_arg)
        return LinkTextElement(text=text, url=url)

    @expand_macro.register
    def expand_macro_link(self, macro: LinkMacro) -> BaseTextElement:
        text: str = macro.arg
        if macro.second_arg is not None:
            text = self._translate_text(macro.second_arg)

        url = self._get_raw_link_url_by_alias(macro.arg)
        if url is None:
            return RichText(text=text)

        return LinkTextElement(text=text, url=self._translate_text(url))

    @expand_macro.register
    def expand_macro_text(self, macro: TextMacro) -> RichText:
        return self.expand_text(self._get_raw_text_body_by_alias(macro.arg))

    @expand_macro.register
    def expand_macro_table(self, macro: TableMacro) -> TableTextElement:
        raw_table_body = self._get_raw_table_body_by_alias(macro.arg)
        return TableTextElement(table_body=[[self.expand_text(cell) for cell in row] for row in raw_table_body])

    @expand_macro.register
    def expand_macro_dialects(self, macro: DialectsMacro) -> ListTextElement:
        items: list[TermTextElement] = []
        for dialect_name in macro.values:
            dialect = get_dialect_from_str(dialect_name)
            human_dialect = HUMAN_DIALECTS[dialect].for_style(DialectStyle.split_version)
            # FIXME: Temporary hack for removing backticks:
            human_dialect = self._translate_text(human_dialect)[1:-1]
            items.append(TermTextElement(term=human_dialect))

        return ListTextElement(items=items, sep=", ")

    @expand_macro.register
    def expand_macro_type(self, macro: TypeMacro) -> ListTextElement:
        items: List[TermTextElement] = []
        type_list = [DataType[type_name.strip().upper()] for type_name in macro.values]
        for h_type_name in get_human_data_type_list(types=type_list):
            items.append(
                TermTextElement(
                    term=self._translate_text(h_type_name),
                    wrap=False,
                )
            )

        return ListTextElement(items=items, sep=" | ", wrap=True)

    @expand_macro.register
    def expand_macro_arg(self, macro: ArgMacro) -> CodeSpanTextElement:
        assert self._args
        return CodeSpanTextElement(
            text=self._args[int(macro.arg)].name,
            wrap=True,
        )

    @expand_macro.register
    def expand_macro_argn(self, macro: ArgNMacro) -> CodeSpanTextElement:
        assert self._args
        return CodeSpanTextElement(
            text=self._args[int(macro.arg)].name,
            wrap=False,
        )

    @expand_macro.register
    def expand_macro_ext_macro(self, macro: ExtMacroMacro) -> ExtMacroTextElement:
        return ExtMacroTextElement(macro_name=macro.arg)

lib/dl_formula_ref/dl_formula_ref/rich_text/expander.py
datalens-tech/datalens-backend
from typing import (
    Collection,
    List,
    Set,
)

from dl_formula.core.datatype import DataType
from dl_formula_ref.texts import (
    ANY_TYPE,
    HUMAN_DATA_TYPES,
)


HIDDEN_TYPES = {
    DataType.DATETIME,
    DataType.CONST_DATETIME,
    DataType.DATETIMETZ,
    DataType.CONST_DATETIMETZ,
    DataType.TREE_STR,
    DataType.CONST_TREE_STR,
}


def get_human_data_type_list(types: Collection[DataType]) -> List[str]:
    types_set: Set[DataType] = set(t.non_const_type for t in types)
    result_types = set()
    # try to find known type combinations within the given list of types
    type_combos = [type_combo for type_combo in HUMAN_DATA_TYPES if isinstance(type_combo, tuple)]
    type_combos = sorted(type_combos, key=len, reverse=True)
    for type_combo in type_combos:
        if set(type_combo).issubset(types_set):  # type: ignore  # TODO: fix
            result_types.add(HUMAN_DATA_TYPES[type_combo])
            types_set -= set(type_combo)  # type: ignore  # TODO: fix

    for hidden_t in HIDDEN_TYPES:
        types_set.discard(hidden_t)

    # for the rest use basic types
    result_types.update({HUMAN_DATA_TYPES[t] for t in types_set})
    result = sorted(result_types) if result_types else [ANY_TYPE]
    return result


def escape_cell(s: str) -> str:
    """
    Escape string ``s`` for usage inside a Markdown table cell.
    Pipe (``|``) characters must either be wrapped into ``<code></code>`` tags
    or replaced with HTML escape sequence ``&#124;`` so that the don't break cell structure
    """
    if "|" not in s:
        return s

    if "`" not in s:
        return s.replace("|", "&#124;")

    escaped_s = ""
    i = 0
    while True:
        # add the substring before next opening backtick
        bt_ind = s.find("`", i)
        if bt_ind == -1:
            bt_ind = len(s)
        escaped_s += s[i:bt_ind].replace("|", "&#124;")
        if bt_ind >= len(s):
            break

        use_tags = False
        # find closing backtick and check whether there are any pipes inside
        next_pipe_ind = s.find("|", bt_ind + 1)
        next_bt_ind = s.find("`", bt_ind + 1)
        if next_pipe_ind != -1 and next_pipe_ind < next_bt_ind:
            # found pipe inside backticks, so will have to use tags instead of backticks
            use_tags = True

        if use_tags:
            escaped_s = "{}{}{}{}".format(escaped_s, "<code>", s[bt_ind + 1 : next_bt_ind], "</code>")
        else:
            escaped_s = escaped_s + s[bt_ind : next_bt_ind + 1]

        i = next_bt_ind + 1

    escaped_s = escaped_s.replace("|", "&#124;")
    return escaped_s

lib/dl_formula_ref/dl_formula_ref/rich_text/helpers.py
datalens-tech/datalens-backend
from typing import (
    List,
    Optional,
)

import attr


@attr.s
class BaseMacro:
    pass


@attr.s
class SingleArgMacro(BaseMacro):
    arg: str = attr.ib(kw_only=True)


@attr.s
class DoubleArgMacro(SingleArgMacro):
    second_arg: Optional[str] = attr.ib(kw_only=True, default=None)


@attr.s
class ListMacro(BaseMacro):
    values: List[str] = attr.ib(kw_only=True)


class RefMacro(DoubleArgMacro):
    category_name: Optional[str] = None


class LinkMacro(DoubleArgMacro):
    pass


class TableMacro(SingleArgMacro):
    pass


class TextMacro(SingleArgMacro):
    pass


class DialectsMacro(ListMacro):
    pass


class ArgMacro(SingleArgMacro):
    pass


class ArgNMacro(SingleArgMacro):
    pass


class TypeMacro(ListMacro):
    pass


class ExtMacroMacro(SingleArgMacro):
    pass


class CategoryMacro(DoubleArgMacro):
    anchor_name: Optional[str] = None


SINGLE_ARG_MACROS = {
    "table": TableMacro,
    "text": TextMacro,
    "arg": ArgMacro,
    "argn": ArgNMacro,
    "macro": ExtMacroMacro,
}
DOUBLE_ARG_MACROS = {
    "ref": RefMacro,
    "link": LinkMacro,
    "category": CategoryMacro,
}
LIST_MACROS = {
    "dialects": DialectsMacro,
    "type": TypeMacro,
}

lib/dl_formula_ref/dl_formula_ref/rich_text/macro.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from functools import singledispatchmethod
from typing import (
    List,
    Mapping,
    Optional,
)

import attr
import tabulate

from dl_formula_ref.rich_text.elements import (
    AudienceBlock,
    CodeSpanTextElement,
    ConditionalBlock,
    ExtMacroTextElement,
    LinkTextElement,
    ListTextElement,
    NoteBlock,
    RichText,
    TableTextElement,
    TermTextElement,
)
from dl_formula_ref.rich_text.helpers import escape_cell


@attr.s
class RichTextRenderEnvironment:
    block_conditions: Mapping[str, bool] = attr.ib(kw_only=True, factory=dict)


class RichTextRenderer(abc.ABC):
    @abc.abstractmethod
    def render(self, rtext: RichText, env: Optional[RichTextRenderEnvironment] = None) -> str:
        raise NotImplementedError


class MdRichTextRenderer(RichTextRenderer):
    def render(self, rtext: RichText, env: Optional[RichTextRenderEnvironment] = None) -> str:
        if env is None:
            env = RichTextRenderEnvironment()
        assert env is not None
        return self._render(rtext, env=env)

    @singledispatchmethod
    def _render(self, value, env: RichTextRenderEnvironment) -> str:  # type: ignore
        raise TypeError

    @_render.register
    def _render_rich_text(self, value: RichText, env: RichTextRenderEnvironment) -> str:
        parts: List[str] = []
        last_idx = 0
        for key, element in sorted(value.replacements.items()):
            parts.append(value.text[last_idx : key.start_pos])
            parts.append(self._render(element, env=env))
            last_idx = key.end_pos

        parts.append(value.text[last_idx:])
        return "".join(parts)

    @_render.register
    def _render_term(self, value: TermTextElement, env: RichTextRenderEnvironment) -> str:
        if value.wrap:
            return f"`{value.term}`"
        return value.term

    @_render.register
    def _render_code_span(self, value: CodeSpanTextElement, env: RichTextRenderEnvironment) -> str:
        if value.wrap:
            return f"`{value.text}`"
        return value.text

    @_render.register
    def _render_ext_macro(self, value: ExtMacroTextElement, env: RichTextRenderEnvironment) -> str:
        return f"{{{{ {value.macro_name} }}}}"

    @_render.register
    def _render_list(self, value: ListTextElement, env: RichTextRenderEnvironment) -> str:
        list_str = value.sep.join([self._render(item, env=env) for item in value.items])
        if value.wrap:
            list_str = f"`{list_str}`"
        return list_str

    @_render.register
    def _render_table(self, value: TableTextElement, env: RichTextRenderEnvironment) -> str:
        table_data = value.table_body
        return tabulate.tabulate(
            [[escape_cell(self._render(cell, env=env)) for cell in row] for row in table_data[1:]],
            headers=[escape_cell(self._render(cell, env=env)) for cell in table_data[0]],
            tablefmt="pipe",
        )

    @_render.register
    def _render_link(self, value: LinkTextElement, env: RichTextRenderEnvironment) -> str:
        return f"[{value.text}]({value.url})"

    @_render.register
    def _render_conditional_block(self, value: ConditionalBlock, env: RichTextRenderEnvironment) -> str:
        if env.block_conditions.get(value.condition):
            return self._render(value.rich_text, env=env)
        return ""

    @_render.register
    def _render_note_block(self, value: NoteBlock, env: RichTextRenderEnvironment) -> str:
        content = self._render(value.rich_text, env=env)
        return f"{{% note {value.level} %}}{content}{{% endnote %}}"

    @_render.register
    def _render_audience(self, value: AudienceBlock, env: RichTextRenderEnvironment) -> str:
        content = self._render(value.rich_text, env=env)
        if len(value.audience_types) == 1:
            return f"{{% if audience == '{value.audience_types[0]}' %}}{content}{{% endif %}}"

        audience_types_str = "[{}]".format(", ".join(["'{}'".format(aud) for aud in value.audience_types]))
        return f"{{% if audience in {audience_types_str} %}}{content}{{% endif %}}"

lib/dl_formula_ref/dl_formula_ref/rich_text/renderer.py
datalens-tech/datalens-backend
from __future__ import annotations

import argparse

from dl_db_testing.loader import load_db_testing_lib
from dl_formula_ref.generator import (
    ConfigVersion,
    ReferenceDocGenerator,
)
from dl_formula_ref.loader import load_formula_ref


parser = argparse.ArgumentParser(prog="Example data management tool")
subparsers = parser.add_subparsers(title="command", dest="command")

subparsers.add_parser("generate", help="Generate data for examples")


class ExampleDataTool:
    @classmethod
    def generate_example_data(cls):
        """
        Requires a DB URL mapping in local file `dl_formula_ref/db_config.json`
        with the following format:

        {
            "CLICKHOUSE_21_8": "clickhouse://datalens:qwerty@localhost:50456/formula_test"
        }

        defining all the required database types:
        - CLICKHOUSE_21_8
        """
        ref_doc_generator = ReferenceDocGenerator(locale="en", config_version=ConfigVersion.default)
        ref_doc_generator.generate_example_data()

        print("Generated data successfully")

    @classmethod
    def run(cls, args):
        tool = cls()

        if args.command == "generate":
            tool.generate_example_data()


def main():
    load_formula_ref()
    load_db_testing_lib()
    ExampleDataTool.run(parser.parse_args())


if __name__ == "__main__":
    main()

lib/dl_formula_ref/dl_formula_ref/scripts/example_data.py
datalens-tech/datalens-backend
from __future__ import annotations

import argparse

from dl_formula_ref.generator import (
    ConfigVersion,
    ReferenceDocGenerator,
    get_generator_config,
)
from dl_formula_ref.loader import load_formula_ref
from dl_formula_ref.localization import DEFAULT_LOCALE


def conf_version_type(s: str) -> ConfigVersion:
    return ConfigVersion[s]


parser = argparse.ArgumentParser(prog="Formula documentation command line tool")
subparsers = parser.add_subparsers(title="command", dest="command")

config_version = argparse.ArgumentParser(add_help=False)
config_version.add_argument(
    "--config-version", type=conf_version_type, default=ConfigVersion.default.name, help="Configuration version"
)

locale_parser = argparse.ArgumentParser(add_help=False)
locale_parser.add_argument("--locale", help="Locale", default=DEFAULT_LOCALE)

outdir_parser = argparse.ArgumentParser(add_help=False)
outdir_parser.add_argument("outdir", help="Directory where to save generated files")

subparsers.add_parser("config-versions", help="List available config versions")

subparsers.add_parser("locales", parents=[config_version], help="List available locales")

subparsers.add_parser(
    "generate", parents=[outdir_parser, config_version, locale_parser], help="Generate function documentation directory"
)


class FormulaDocTool:
    @staticmethod
    def print_versions():
        print("\n".join(sorted([version.name for version in ConfigVersion])))

    @staticmethod
    def print_locales(config_version: ConfigVersion):
        gen_config = get_generator_config(version=config_version)
        print("\n".join(sorted(gen_config.supported_locales)))

    @classmethod
    def generate_doc(cls, outdir: str, locale: str, config_version: ConfigVersion):
        ref_doc_generator = ReferenceDocGenerator(locale=locale, config_version=config_version)
        ref_doc_generator.generate_doc_full_dir(outdir=outdir)

    @classmethod
    def run(cls, args):
        tool = cls()

        match args.command:
            case "config-versions":
                tool.print_versions()
            case "locales":
                tool.print_locales(config_version=args.config_version)
            case "generate":
                tool.generate_doc(outdir=args.outdir, locale=args.locale, config_version=args.config_version)


def main():
    load_formula_ref()
    FormulaDocTool.run(parser.parse_args())


if __name__ == "__main__":
    main()

lib/dl_formula_ref/dl_formula_ref/scripts/formula_doc.py
datalens-tech/datalens-backend
from enum import Enum
from typing import NamedTuple

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.dialect import StandardDialect as D
from dl_formula_ref.localization import get_gettext
from dl_i18n.localizer_base import Translatable


_ = get_gettext()

EXAMPLE_TITLE = _("Example")
DOC_TOC_TITLE = _("Function Reference")
DOC_ALL_TITLE = _("All Functions")
DOC_OVERVIEW_TEXT = _("Overview")
DOC_AVAIL_TITLE = _("Function Availability")
COMMON_TYPE_NOTE = _("Arguments ({args}) must be of the same type.")
CONST_TYPE_NOTE = _("Only constant values are accepted for the arguments ({args}).")
ALSO_IN_OTHER_CATEGORIES = _("Function `{func_name}` is also found in the following categories: ")
DEPENDS_ON_ARGS = _("Depends on argument types")
FROM_ARGS = _("Same type as ({args})")
ANY_TYPE = _("Any")
HUMAN_CATEGORIES = {
    "aggregation": _("Aggregate functions"),
    "string": _("String functions"),
    "numeric": _("Mathematical functions"),
    "logical": _("Logical functions"),
    "date": _("Date/Time functions"),
    "type-conversion": _("Type conversion functions"),
    "operator": _("Operators"),
    "geographical": _("Geographical functions"),
    "markup": _("Text markup functions"),
    "window": _("Window functions"),
    "time-series": _("Time series functions"),
    "array": _("Array functions"),
}
FUNCTION_CATEGORY_TAG = {
    "aggregation": _("aggregation"),
    "string": _("string"),
    "numeric": _("mathematical"),
    "logical": _("logical"),
    "date": _("date/time"),
    "type-conversion": _("type conversion"),
    "operator": _("operator"),
    "geographical": _("geographical"),
    "markup": _("text markup"),
    "window": _("window"),
    "time-series": _("time series"),
    "array": _("array"),
}
HUMAN_DATA_TYPES = {
    DataType.INTEGER: _("Integer"),
    DataType.FLOAT: _("Fractional number"),
    DataType.BOOLEAN: _("Boolean"),
    DataType.STRING: _("String"),
    DataType.DATE: _("Date"),
    DataType.DATETIME: _("Datetime (deprecated)"),
    DataType.DATETIMETZ: _("Datetime with timezone"),
    DataType.GENERICDATETIME: _("Datetime"),
    DataType.GEOPOINT: _("Geopoint"),
    DataType.GEOPOLYGON: _("Geopolygon"),
    DataType.MARKUP: _("Markup"),
    DataType.UUID: _("UUID"),
    DataType.ARRAY_FLOAT: _("Array of fractional numbers"),
    DataType.ARRAY_INT: _("Array of integers"),
    DataType.ARRAY_STR: _("Array of strings"),
    DataType.TREE_STR: _("Tree"),
    (
        DataType.FLOAT,
        DataType.INTEGER,
        DataType.BOOLEAN,
        DataType.STRING,
        DataType.DATE,
        DataType.DATETIME,
        DataType.GENERICDATETIME,
        DataType.GEOPOINT,
        DataType.GEOPOLYGON,
        DataType.MARKUP,
        DataType.UUID,
        DataType.ARRAY_FLOAT,
        DataType.ARRAY_INT,
        DataType.ARRAY_STR,
    ): ANY_TYPE,  # TODO: remove after DATETIMETZ is fully adopted
    (
        DataType.FLOAT,
        DataType.INTEGER,
        DataType.BOOLEAN,
        DataType.STRING,
        DataType.DATE,
        DataType.DATETIME,
        DataType.DATETIMETZ,
        DataType.GENERICDATETIME,
        DataType.GEOPOINT,
        DataType.GEOPOLYGON,
        DataType.MARKUP,
        DataType.UUID,
        DataType.ARRAY_FLOAT,
        DataType.ARRAY_INT,
        DataType.ARRAY_STR,
        DataType.TREE_STR,
    ): ANY_TYPE,
}


class DialectStyle(Enum):
    simple = "simple"
    multiline = "multiline"
    split_version = "split_version"


class StyledDialect(NamedTuple):
    simple: str | Translatable
    multiline: str | Translatable
    split_version: str | Translatable

    def for_style(self, item):
        if isinstance(item, DialectStyle):
            item = item.name
        if isinstance(item, str):
            return getattr(self, item)


# These should be filled from plugins
HIDDEN_DIALECTS: set[DialectCombo] = set()
ANY_DIALECTS: set[DialectCombo] = set()


HUMAN_DIALECTS: dict[DialectCombo, StyledDialect] = {
    D.SQLITE: StyledDialect(  # TODO: remove
        "`SQLite`",
        "`SQLite`",
        "`SQLite`",
    ),
}

SIGNATURE_TITLE_STANDARD = _("Standard")
SIGNATURE_TITLE_EXTENDED = _("Extended")
SIGNATURE_DESC_EXTENDED_HEADER = _("More info:")


def register_any_dialects(any_dialects: frozenset[DialectCombo]) -> None:
    ANY_DIALECTS.update(any_dialects)


def register_human_dialects(human_dialects: dict[DialectCombo, StyledDialect]) -> None:
    HUMAN_DIALECTS.update(human_dialects)

lib/dl_formula_ref/dl_formula_ref/texts.py
datalens-tech/datalens-backend
import pytest

from dl_db_testing.loader import load_db_testing_lib
from dl_formula_ref.loader import load_formula_ref


@pytest.fixture(scope="session", autouse=True)
def loaded_libraries() -> None:
    load_formula_ref()
    load_db_testing_lib()

lib/dl_formula_ref/dl_formula_ref_tests/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import time
from typing import MutableSet

import pytest

from dl_formula.core.dialect import DialectCombo
from dl_formula.definitions.literals import literal
from dl_formula_testing.database import (
    Db,
    make_db_config,
    make_db_from_config,
)
from dl_formula_testing.evaluator import DbEvaluator
from dl_testing.containers import get_test_container_hostport

from dl_connector_clickhouse.formula.constants import ClickHouseDialect


ALL_DB_CONFIGURATIONS = {
    ClickHouseDialect.CLICKHOUSE_22_10: (
        f'clickhouse://datalens:qwerty@{get_test_container_hostport("db-clickhouse-22-10", fallback_port=51110).as_pair()}/datalens_test'
    ),
}
EXCLUDE_DIALECTS = []
DB_CONFIGURATIONS = {
    dialect: url for dialect, url in ALL_DB_CONFIGURATIONS.items() if url and dialect not in EXCLUDE_DIALECTS
}


@pytest.fixture(scope="session")
def all_db_configurations():
    return ALL_DB_CONFIGURATIONS


def base_db_for(dialect: DialectCombo) -> Db:
    url = ALL_DB_CONFIGURATIONS[dialect]
    assert url
    db = make_db_from_config(make_db_config(url=url, dialect=dialect))
    return db


def alive_db_for(dialect, timelimit, poll_pause):
    check_value_in = 123
    while True:
        db = base_db_for(dialect)
        try:
            check_value_res = db.scalar(literal(check_value_in, d=dialect))
        except Exception as exc:
            if time.monotonic() > timelimit:
                raise Exception(f"Timed out while waiting for db {dialect!r} {db!r}", exc)
            print(f"Waiting for db to come up: {dialect!r}   (db={db!r}, exc={exc!r}).")
            time.sleep(poll_pause)
        else:
            assert check_value_res == check_value_in
            break
    return db


class DbStateTracking:
    """State-keeper for db liveness, to be used as a singleton"""

    def __init__(self):
        self.db_is_up: MutableSet[DialectCombo] = set()
        self.first_call_time = None

    def get_db(self, dialect, global_timeout, poll_pause):
        if dialect in self.db_is_up:
            return base_db_for(dialect)  # supposed to be alive already.

        # Share the 'max time to up' between all dialects:
        if self.first_call_time is None:
            self.first_call_time = time.monotonic()

        timelimit = self.first_call_time + global_timeout
        db = alive_db_for(dialect=dialect, timelimit=timelimit, poll_pause=poll_pause)
        self.db_is_up.add(dialect)
        return db


DB_STATE_TRACKING = DbStateTracking()


def dbe_for(dialect: DialectCombo, wait_for_up_timeout=180.0, wait_for_up_pause=0.7) -> DbEvaluator:
    db = DB_STATE_TRACKING.get_db(dialect=dialect, global_timeout=wait_for_up_timeout, poll_pause=wait_for_up_pause)
    dbe = DbEvaluator(db=db)
    return dbe


@pytest.fixture(
    scope="session",
    params=[dialect for dialect in sorted(DB_CONFIGURATIONS)],
    ids=[d.common_name_and_version for d in sorted(DB_CONFIGURATIONS)],
)
def any_dialect(request) -> DialectCombo:
    return request.param


@pytest.fixture(scope="session")
def dbe(any_dialect):
    return dbe_for(any_dialect)

lib/dl_formula_ref/dl_formula_ref_tests/db/conftest.py
datalens-tech/datalens-backend
import sqlalchemy as sa

from dl_formula.core.datatype import DataType
from dl_formula_ref.examples.data_table import (
    DataColumn,
    DataTable,
)
from dl_formula_ref.examples.dumper import get_dumper
from dl_formula_ref.examples.query import CompiledQueryContext


def test_select_original_data(dbe):
    db = dbe.db
    original_data_table = DataTable(
        columns=[
            DataColumn(name="int_value", data_type=DataType.INTEGER),
            DataColumn(name="str_value", data_type=DataType.STRING),
        ],
        rows=[
            [11, "qwe"],
            [45, "rty"],
        ],
    )
    dumper = get_dumper(db=db)
    with dumper.temporary_data_table(data_table=original_data_table) as table_ref:
        sa_query = (
            sa.select(
                [
                    sa.literal_column("int_value"),
                    sa.literal_column("str_value"),
                ]
            )
            .select_from(sa.table(name=table_ref.name))
            .order_by(sa.literal_column("int_value"))
        )
        query_ctx = CompiledQueryContext(
            sa_query=sa_query,
            result_columns=original_data_table.columns,
        )
        loaded_data_table = dumper.execute_query(query_ctx=query_ctx)

    assert loaded_data_table == original_data_table

lib/dl_formula_ref/dl_formula_ref_tests/db/examples/test_dumper.py
datalens-tech/datalens-backend
import json
import tempfile

import pytest

from dl_formula_ref.config import DOC_GEN_CONFIG_DEFAULT
from dl_formula_ref.generator import (
    ConfigVersion,
    ReferenceDocGenerator,
)

from dl_connector_clickhouse.formula.constants import ClickHouseDialect


@pytest.fixture(scope="function")
def example_db_conf_patch(monkeypatch, all_db_configurations, dbe):
    db_conf_data = {
        "CLICKHOUSE_22_10": all_db_configurations[ClickHouseDialect.CLICKHOUSE_22_10],
    }
    with tempfile.NamedTemporaryFile() as db_conf_f:
        # Prepare DB config
        with open(db_conf_f.name, "w") as db_conf_f_w:
            json.dump(db_conf_data, db_conf_f_w)

        monkeypatch.setattr(DOC_GEN_CONFIG_DEFAULT, "db_config_file", db_conf_f.name)
        monkeypatch.setattr(DOC_GEN_CONFIG_DEFAULT, "default_example_dialect", ClickHouseDialect.CLICKHOUSE_22_10)
        yield


@pytest.fixture(scope="function")
def example_data_file_patch(monkeypatch):
    with tempfile.NamedTemporaryFile() as ex_data_f:
        monkeypatch.setattr(DOC_GEN_CONFIG_DEFAULT, "example_data_file", ex_data_f.name)
        yield


def test_prepare_example_data(example_db_conf_patch, example_data_file_patch):
    gen = ReferenceDocGenerator(config_version=ConfigVersion.default, locale="en")
    gen.generate_example_data()

    with open(DOC_GEN_CONFIG_DEFAULT.example_data_file, "rb") as ex_data_f:
        data: list = json.load(ex_data_f)
        keys = {item[0] for item in data}
        assert "ABS.Example" in keys

lib/dl_formula_ref/dl_formula_ref_tests/db/examples/test_preparation.py
datalens-tech/datalens-backend
from dl_formula_ref_tests.conftest import loaded_libraries


__all__ = ("loaded_libraries",)

lib/dl_formula_ref/dl_formula_ref_tests/unit/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import TYPE_CHECKING

from dl_formula.core.datatype import DataType
from dl_formula.parser.factory import get_parser
from dl_formula.shortcuts import n
from dl_formula_ref.examples.config import (
    ExampleConfig,
    ExampleSource,
)
from dl_formula_ref.examples.data_table import DataColumn
from dl_formula_ref.examples.query import (
    FormulaContext,
    RawQueryContext,
)
from dl_formula_ref.examples.query_gen import QueryGenerator

from dl_connector_clickhouse.formula.constants import ClickHouseDialect


if TYPE_CHECKING:
    from dl_formula.core import nodes


def parse(formula: str) -> nodes.Formula:
    return get_parser().parse(formula)


def test_gen_query_from_example():
    example = ExampleConfig(
        dialect=ClickHouseDialect.CLICKHOUSE_21_8,
        source=ExampleSource(
            columns=[("int_value", DataType.INTEGER), ("str_value", DataType.STRING)],
            data=[],
        ),
        formula_fields=[
            ("int_value", "[int_value]"),
            ("str_value", "[str_value]"),
            ("test_1", "[int_value] * 2"),
            ("test_2", 'CONCAT([str_value], "AAA")'),
            ("test_3", "INT([str_value])"),
        ],
        group_by=["[int_value]", "[str_value]"],
        order_by=["[int_value]"],
    )
    query_gen = QueryGenerator()
    actual_query = query_gen.generate_query(example=example)
    expected_query = RawQueryContext(
        result_columns=[
            DataColumn(name="int_value", data_type=DataType.INTEGER),
            DataColumn(name="str_value", data_type=DataType.STRING),
            DataColumn(name="test_1", data_type=DataType.INTEGER),
            DataColumn(name="test_2", data_type=DataType.STRING),
            DataColumn(name="test_3", data_type=DataType.INTEGER),
        ],
        select=[
            FormulaContext(formula=n.formula(n.field("int_value")), alias="int_value"),
            FormulaContext(formula=n.formula(n.field("str_value")), alias="str_value"),
            FormulaContext(
                formula=n.formula(n.binary("*", n.field("int_value"), n.lit(2))),
                alias="test_1",
            ),
            FormulaContext(
                formula=n.formula(n.func.CONCAT(n.field("str_value"), n.lit("AAA"))),
                alias="test_2",
            ),
            FormulaContext(
                formula=n.formula(n.func.INT(n.field("str_value"))),
                alias="test_3",
            ),
        ],
        group_by=[
            FormulaContext(formula=n.formula(n.field("int_value")), alias=None),
            FormulaContext(formula=n.formula(n.field("str_value")), alias=None),
        ],
        order_by=[
            FormulaContext(formula=n.formula(n.field("int_value")), alias=None),
        ],
    )
    assert actual_query == expected_query

lib/dl_formula_ref/dl_formula_ref_tests/unit/examples/test_query_gen.py
datalens-tech/datalens-backend
import tempfile

from dl_formula.core.datatype import DataType
from dl_formula_ref.examples.data_table import (
    DataColumn,
    DataTable,
)
from dl_formula_ref.examples.result_storage import (
    ReadableDataStorage,
    StorageKey,
    WritableDataStorage,
)


def test_result_storage():
    original_data_table = DataTable(
        columns=[
            DataColumn(name="int_value", data_type=DataType.INTEGER),
            DataColumn(name="str_value", data_type=DataType.STRING),
        ],
        rows=[
            [11, "qwe"],
            [45, "rty"],
        ],
    )
    storage_key = StorageKey(
        item_key="smth",
        example_key="whatever",
    )

    with tempfile.NamedTemporaryFile() as st_file:
        st_filename = st_file.name
        w_storage = WritableDataStorage(filename=st_filename)
        r_storage = ReadableDataStorage(filename=st_filename)

        w_storage.set_result_data(storage_key=storage_key, result_data=original_data_table)
        w_storage.save()
        r_storage.load()
        loaded_data_table = r_storage.get_result_data(storage_key=storage_key)

    assert loaded_data_table == original_data_table

lib/dl_formula_ref/dl_formula_ref_tests/unit/examples/test_result_storage.py
datalens-tech/datalens-backend
from typing import (
    Optional,
    Tuple,
)

from dl_formula.core.datatype import DataType
from dl_formula_ref.registry.aliased_res import (
    AliasedLinkResource,
    AliasedTableResource,
    AliasedTextResource,
    SimpleAliasedResourceRegistry,
)
from dl_formula_ref.registry.arg_base import FuncArg
from dl_formula_ref.rich_text.elements import (
    AudienceBlock,
    CodeSpanTextElement,
    ConditionalBlock,
    ExtMacroTextElement,
    LinkTextElement,
    ListTextElement,
)
from dl_formula_ref.rich_text.elements import (
    NoteBlock,
    RichText,
    TableTextElement,
    TermTextElement,
)
from dl_formula_ref.rich_text.elements import MacroReplacementKey as MRK
from dl_formula_ref.rich_text.expander import MacroExpander


def test_expand_macro_dialects():
    expander = MacroExpander()
    text = "qwerty {dialects: CLICKHOUSE} uiop"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 29): ListTextElement(
                items=[TermTextElement(term="ClickHouse")],
                sep=", ",
            ),
        },
    )
    assert actual_res == exp_res


def test_expand_macro_type():
    expander = MacroExpander()
    text = "qwerty {type: INTEGER | STRING} uiop"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 31): ListTextElement(
                items=[TermTextElement(term="Integer", wrap=False), TermTextElement(term="String", wrap=False)],
                sep=" | ",
                wrap=True,
            ),
        },
    )
    assert actual_res == exp_res


def test_expand_macro_ext_macro():
    expander = MacroExpander()
    text = "qwerty {macro: some_cloud_macro} uiop"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 32): ExtMacroTextElement(macro_name="some_cloud_macro"),
        },
    )
    assert actual_res == exp_res


def test_expand_macro_text():
    resources = SimpleAliasedResourceRegistry(
        resources={
            "my_text": AliasedTextResource(body="Some {text: nested_text} Text"),
            "nested_text": AliasedTextResource(body="Nested"),
        }
    )
    expander = MacroExpander(resources=resources)
    text = "qwerty {text: my_text} uiop"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 22): RichText(
                "Some {text: nested_text} Text",
                replacements={
                    MRK(5, 24): RichText("Nested"),
                },
            ),
        },
    )
    assert actual_res == exp_res


def test_expand_macro_link():
    resources = SimpleAliasedResourceRegistry(
        resources={
            "my_link": AliasedLinkResource(url="http://some.url"),
        }
    )
    expander = MacroExpander(resources=resources)
    text = "qwerty {link: my_link: My Link} uiop"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 31): LinkTextElement(url="http://some.url", text="My Link"),
        },
    )
    assert actual_res == exp_res


def test_expand_macro_table():
    resources = SimpleAliasedResourceRegistry(
        resources={
            "my_table": AliasedTableResource(table_body=[["Some {text: nested_text} Text", "Something"]]),
            "nested_text": AliasedTextResource(body="Nested"),
        }
    )
    expander = MacroExpander(resources=resources)
    text = "qwerty {table: my_table} uiop"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 24): TableTextElement(
                table_body=[
                    [
                        RichText("Some {text: nested_text} Text", replacements={MRK(5, 24): RichText("Nested")}),
                        RichText("Something"),
                    ],
                ]
            ),
        },
    )
    assert actual_res == exp_res


def test_expand_macro_arg_argn():
    expander = MacroExpander(
        args=[
            FuncArg(name="first_arg", types={DataType.STRING}, optional_level=0),
            FuncArg(name="second_arg", types={DataType.STRING}, optional_level=0),
        ]
    )
    text = "qwerty {argn: 1} uiop {arg: 0}"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 16): CodeSpanTextElement(text="second_arg", wrap=False),
            MRK(22, 30): CodeSpanTextElement(text="first_arg", wrap=True),
        },
    )
    assert actual_res == exp_res


def test_expand_macro_ref():
    def func_link_provider(func_name: str, category_name: str) -> Tuple[str, str]:
        return (func_name, f'https://{func_name.lower()}{"-win" if category_name == "window" else ""}.funcs.com')

    def cat_link_provider(category_name: str, anchor_name: Optional[str]) -> Tuple[str, str]:
        return (category_name, f"https://{category_name.lower()}.cats.com")

    expander = MacroExpander(
        func_link_provider=func_link_provider,
        cat_link_provider=cat_link_provider,
    )
    text = "qwerty {ref: SUM} uiop {ref: window/AVG} asdf {ref: aggregation/AVG:aggregate function}"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 17): LinkTextElement(url="https://sum.funcs.com", text="SUM"),
            MRK(23, 40): LinkTextElement(url="https://avg-win.funcs.com", text="AVG"),
            MRK(46, 87): LinkTextElement(url="https://avg.funcs.com", text="aggregate function"),
        },
    )
    assert actual_res == exp_res


def test_conditional_block():
    expander = MacroExpander()
    text = "qwerty {if some_cond} uiop{end}"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 31): ConditionalBlock(
                condition="some_cond",
                rich_text=RichText(text=" uiop", replacements={}),
            ),
        },
    )
    assert actual_res == exp_res


def test_nested_conditional_block():
    expander = MacroExpander()
    text = "qwerty {if some_cond} uio{if smth_else}pp{end}p{end}"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 52): ConditionalBlock(
                condition="some_cond",
                rich_text=RichText(
                    text=" uio{if smth_else}pp{end}p",
                    replacements={
                        MRK(4, 25): ConditionalBlock(
                            condition="smth_else",
                            rich_text=RichText(text="pp", replacements={}),
                        ),
                    },
                ),
            ),
        },
    )
    assert actual_res == exp_res


def test_note_block():
    expander = MacroExpander()
    text = "qwerty {note warning} uiop{end}"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 31): NoteBlock(
                level="warning",
                rich_text=RichText(text=" uiop", replacements={}),
            ),
        },
    )
    assert actual_res == exp_res


def test_audience_block():
    expander = MacroExpander()
    text = "qwerty {audience internal} uiop{end}"
    actual_res = expander.expand_text(text)
    exp_res = RichText(
        text=text,
        replacements={
            MRK(7, 36): AudienceBlock(
                audience_types=["internal"],
                rich_text=RichText(text=" uiop", replacements={}),
            ),
        },
    )
    assert actual_res == exp_res

lib/dl_formula_ref/dl_formula_ref_tests/unit/rich_text/test_expander.py
datalens-tech/datalens-backend
from dl_formula_ref.rich_text.elements import (
    AudienceBlock,
    CodeSpanTextElement,
    ConditionalBlock,
    ExtMacroTextElement,
    LinkTextElement,
    ListTextElement,
)
from dl_formula_ref.rich_text.elements import (
    NoteBlock,
    RichText,
    TableTextElement,
    TermTextElement,
)
from dl_formula_ref.rich_text.elements import MacroReplacementKey as MRK
from dl_formula_ref.rich_text.renderer import (
    MdRichTextRenderer,
    RichTextRenderEnvironment,
)


def test_term():
    renderer = MdRichTextRenderer()
    text = renderer.render(RichText(text="Hi ", replacements={MRK(3, 3): TermTextElement(term="int")}))
    assert text == "Hi `int`"
    text = renderer.render(RichText(text="Hi ", replacements={MRK(3, 3): TermTextElement(term="int", wrap=False)}))
    assert text == "Hi int"


def test_code_span():
    renderer = MdRichTextRenderer()
    text = renderer.render(RichText(text="Hi ", replacements={MRK(3, 3): CodeSpanTextElement(text="int", wrap=True)}))
    assert text == "Hi `int`"
    text = renderer.render(RichText(text="Hi ", replacements={MRK(3, 3): CodeSpanTextElement(text="int", wrap=False)}))
    assert text == "Hi int"


def test_ext_macro():
    renderer = MdRichTextRenderer()
    text = renderer.render(RichText(text="Hi ", replacements={MRK(3, 3): ExtMacroTextElement(macro_name="the_macro")}))
    assert text == "Hi {{ the_macro }}"


def test_list():
    renderer = MdRichTextRenderer()
    text = renderer.render(
        RichText(
            text="Hi ",
            replacements={
                MRK(3, 3): ListTextElement(
                    items=[
                        TermTextElement(term="int"),
                        ExtMacroTextElement(macro_name="my_macro"),
                    ],
                    sep=" : ",
                )
            },
        )
    )
    assert text == "Hi `int` : {{ my_macro }}"

    text = renderer.render(
        RichText(
            text="Hi ",
            replacements={
                MRK(3, 3): ListTextElement(
                    items=[
                        TermTextElement(term="int", wrap=False),
                        ExtMacroTextElement(macro_name="my_macro"),
                    ],
                    sep=" : ",
                    wrap=True,
                )
            },
        )
    )
    assert text == "Hi `int : {{ my_macro }}`"


def test_link():
    renderer = MdRichTextRenderer()
    text = renderer.render(
        RichText(text="Hi ", replacements={MRK(3, 3): LinkTextElement(text="Whatever", url="site.com")})
    )
    assert text == "Hi [Whatever](site.com)"


def test_table():
    renderer = MdRichTextRenderer()
    text = renderer.render(
        RichText(
            text="Hi\n",
            replacements={
                MRK(3, 3): TableTextElement(
                    table_body=[
                        [RichText("My"), RichText("friend")],
                        [TermTextElement(term="smth"), RichText("")],
                    ],
                )
            },
        )
    )
    assert (
        text
        == """
Hi
| My     | friend   |
|:-------|:---------|
| `smth` |          |
""".strip()
    )


def test_conditional_block():
    rich_text = RichText(
        text="qwerty {if some_cond} uio{if smth_else}pp{end}p{end}",
        replacements={
            MRK(7, 52): ConditionalBlock(
                condition="some_cond",
                rich_text=RichText(
                    text=" uio{if smth_else}pp{end}p",
                    replacements={
                        MRK(4, 25): ConditionalBlock(
                            condition="smth_else",
                            rich_text=RichText(text="pp", replacements={}),
                        ),
                    },
                ),
            ),
        },
    )

    renderer = MdRichTextRenderer()
    env = RichTextRenderEnvironment(block_conditions={"some_cond": False, "smth_else": False})
    assert renderer.render(rich_text, env=env) == "qwerty "

    env = RichTextRenderEnvironment(block_conditions={"some_cond": True, "smth_else": False})
    assert renderer.render(rich_text, env=env) == "qwerty  uiop"

    env = RichTextRenderEnvironment(block_conditions={"some_cond": False, "smth_else": True})
    assert renderer.render(rich_text, env=env) == "qwerty "

    env = RichTextRenderEnvironment(block_conditions={"some_cond": True, "smth_else": True})
    assert renderer.render(rich_text, env=env) == "qwerty  uioppp"


def test_note_block():
    rich_text = RichText(
        text="qwerty {note} uio{end}{note warning}pp{end}",
        replacements={
            MRK(7, 21): NoteBlock(
                level="info",
                rich_text=RichText(text=" uio", replacements={}),
            ),
            MRK(21, 43): NoteBlock(
                level="warning",
                rich_text=RichText(text="pp", replacements={}),
            ),
        },
    )

    renderer = MdRichTextRenderer()
    env = RichTextRenderEnvironment()
    assert renderer.render(rich_text, env=env) == (
        "qwerty {% note info %} uio{% endnote %}{% note warning %}pp{% endnote %}"
    )


def test_audience_block():
    rich_text = RichText(
        text="qwerty {audience internal} uio{end}{audience this,that} pp{end}",
        replacements={
            MRK(7, 35): AudienceBlock(
                audience_types=["internal"],
                rich_text=RichText(text=" uio", replacements={}),
            ),
            MRK(35, 63): AudienceBlock(
                audience_types=["this", "that"],
                rich_text=RichText(text=" pp", replacements={}),
            ),
        },
    )

    renderer = MdRichTextRenderer()
    env = RichTextRenderEnvironment()
    assert renderer.render(rich_text, env=env) == (
        "qwerty {% if audience == 'internal' %} uio{% endif %}" "{% if audience in ['this', 'that'] %} pp{% endif %}"
    )

lib/dl_formula_ref/dl_formula_ref_tests/unit/rich_text/test_renderer.py
datalens-tech/datalens-backend
from __future__ import annotations

import os
import pkgutil
import shutil
import tempfile

import pytest

import dl_formula_ref
from dl_formula_ref.config import (
    _CONFIGS_BY_VERSION,
    ConfigVersion,
)
from dl_formula_ref.scripts.formula_doc import (
    FormulaDocTool,
    parser,
)
from dl_formula_testing.tool_runner import ToolRunner

from dl_connector_clickhouse.formula.constants import ClickHouseDialect


@pytest.fixture(scope="module")
def tool():
    return ToolRunner(parser=parser, tool_cls=FormulaDocTool)


@pytest.fixture(scope="function")
def example_data_patch(monkeypatch):
    # TODO: Maybe this should be the standard way to access this file
    # Workaround for accessing data file from the monobinary in CI tests
    with tempfile.NamedTemporaryFile() as ex_file:
        ex_filename = ex_file.name

        data = pkgutil.get_data(dl_formula_ref.__name__, os.path.join("example_data.json"))
        with open(ex_filename, "wb") as ex_file_w:
            ex_file_w.write(data)

        tested_version = ConfigVersion.default
        original_config = _CONFIGS_BY_VERSION[tested_version]
        config_patched = original_config.clone(
            example_data_file=ex_filename,
            supported_dialects=frozenset({ClickHouseDialect.CLICKHOUSE_22_10}),
        )
        _CONFIGS_BY_VERSION[tested_version] = config_patched
        try:
            yield
        finally:
            _CONFIGS_BY_VERSION[tested_version] = original_config


def test_locales(tool):
    stdout, stderr = tool.run(
        [
            "locales",
            "--config-version",
            "default",
        ]
    )
    assert stderr == ""
    assert "en" in stdout


def test_versions(tool):
    stdout, stderr = tool.run(
        [
            "config-versions",
        ]
    )
    assert stderr == ""
    assert "default" in stdout


def test_doc_full_dir(tool, example_data_patch):
    tmpdir = tempfile.mkdtemp()
    stdout, stderr = tool.run(["generate", tmpdir])
    assert stderr == ""
    assert os.path.exists(os.path.join(tmpdir, "function-ref", "CONCAT.md"))
    assert os.path.exists(os.path.join(tmpdir, "function-ref", "all.md"))
    assert os.path.exists(os.path.join(tmpdir, "function-ref", "availability.md"))
    assert os.path.exists(os.path.join(tmpdir, "toc.yaml"))
    shutil.rmtree(tmpdir)

lib/dl_formula_ref/dl_formula_ref_tests/unit/scripts/test_formula_doc.py
datalens-tech/datalens-backend
from dl_formula.definitions.scope import Scope
from dl_formula_ref.categories.aggregation import CATEGORY_AGGREGATION
from dl_formula_ref.categories.logical import CATEGORY_LOGICAL
from dl_formula_ref.categories.mathematical import CATEGORY_MATHEMATICAL
from dl_formula_ref.categories.time_series import CATEGORY_TIME_SERIES
from dl_formula_ref.categories.window import CATEGORY_WINDOW
from dl_formula_ref.registry.env import GenerationEnvironment
from dl_formula_ref.registry.registry import (
    FUNC_REFERENCE_REGISTRY,
    RefFunctionKey,
)
from dl_formula_ref.registry.tools import populate_registry_from_definitions

from dl_connector_clickhouse.formula.constants import ClickHouseDialect


def check_function(func_name: str, exp_signature: list, category_name: str) -> None:
    func_key = RefFunctionKey.normalized(name=func_name, category_name=category_name)
    func = FUNC_REFERENCE_REGISTRY[func_key]
    env = GenerationEnvironment(
        scopes=Scope.DOCUMENTED,
        supported_dialects=frozenset({ClickHouseDialect.CLICKHOUSE_22_10}),
    )
    signatures = [sig.body for sig in func.get_signatures(env=env).signatures]
    assert signatures == exp_signature


def test_signatures():
    populate_registry_from_definitions()
    # Function without args
    check_function("PI", ["PI()"], CATEGORY_MATHEMATICAL.name)
    # Function with one arg
    check_function("ABS", ["ABS( number )"], CATEGORY_MATHEMATICAL.name)
    # Function with multiple args
    check_function("IFNULL", ["IFNULL( check_value, alt_value )"], CATEGORY_LOGICAL.name)
    # Function with optional args
    check_function("ROUND", ["ROUND( number [ , precision ] )"], CATEGORY_MATHEMATICAL.name)
    # Infinite arg function
    check_function("GREATEST", ["GREATEST( value_1, value_2, value_3 [ , ... ] )"], CATEGORY_MATHEMATICAL.name)
    # Aggregation
    check_function(
        "SUM",
        [
            "SUM( value )",
            "SUM( value\n" "     [ FIXED ... | INCLUDE ... | EXCLUDE ... ]\n" "     [ BEFORE FILTER BY ... ]\n" "   )",
        ],
        CATEGORY_AGGREGATION.name,
    )
    # Window function without ORDER BY
    check_function(
        "SUM",
        [
            "SUM( value\n" "     TOTAL | WITHIN ... | AMONG ...\n" "   )",
            "SUM( value\n" "     TOTAL | WITHIN ... | AMONG ...\n" "     [ BEFORE FILTER BY ... ]\n" "   )",
        ],
        CATEGORY_WINDOW.name,
    )
    # Window function with ORDER BY
    check_function(
        "MSUM",
        [
            "MSUM( value, rows_1 [ , rows_2 ] )",
            "MSUM( value, rows_1 [ , rows_2 ]\n"
            "      [ TOTAL | WITHIN ... | AMONG ... ]\n"
            "      [ ORDER BY ... ]\n"
            "      [ BEFORE FILTER BY ... ]\n"
            "    )",
        ],
        CATEGORY_WINDOW.name,
    )
    # Lookup function
    check_function(
        "AGO",
        [
            "AGO( measure, date_dimension [ , unit [ , number ] ] )",
            "AGO( measure, date_dimension [ , unit [ , number ] ]\n"
            "     [ BEFORE FILTER BY ... ]\n"
            "     [ IGNORE DIMENSIONS ... ]\n"
            "   )",
        ],
        CATEGORY_TIME_SERIES.name,
    )

lib/dl_formula_ref/dl_formula_ref_tests/unit/test_signature_gen.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from functools import partial
import logging
import time
from typing import (
    Optional,
    Union,
)

import attr
import pytz
import sqlalchemy as sa

from dl_db_testing.database.base import (
    DbBase,
    DbConfig,
)
from dl_db_testing.database.dispenser import DbDispenserBase
from dl_db_testing.database.engine_wrapper import get_engine_wrapper_cls_for_url
from dl_formula.core.dialect import DialectCombo
from dl_utils.wait import wait_for


LOGGER = logging.getLogger(__name__)


@attr.s(frozen=True)
class FormulaDbConfig(DbConfig):
    dialect: DialectCombo = attr.ib(kw_only=True)
    _tzinfo: Optional[datetime.tzinfo] = attr.ib(kw_only=True, default=None)

    @property
    def tzinfo(self) -> datetime.tzinfo:
        if self._tzinfo is not None:
            return self._tzinfo
        return pytz.UTC


@attr.s
class Db(DbBase[FormulaDbConfig]):
    @property
    def dialect(self) -> DialectCombo:
        return self.config.dialect

    @property
    def tzinfo(self) -> datetime.tzinfo:
        return self.config.tzinfo


def make_db_config(dialect: DialectCombo, url: str, tzinfo: Optional[datetime.tzinfo] = None) -> FormulaDbConfig:
    db_eng_config_cls = get_engine_wrapper_cls_for_url(url).CONFIG_CLS
    db_eng_config = db_eng_config_cls(url=url)
    db_config = FormulaDbConfig(
        engine_config=db_eng_config,
        dialect=dialect,
        tzinfo=tzinfo,
    )
    return db_config


def make_db_from_config(db_config: FormulaDbConfig) -> Db:
    ew_cls = get_engine_wrapper_cls_for_url(db_config.engine_config.url)
    engine_wrapper = ew_cls(config=db_config.engine_config)
    return Db(config=db_config, engine_wrapper=engine_wrapper)


@attr.s
class FormulaDbDispenser(DbDispenserBase[FormulaDbConfig, Db]):
    global_timeout: Union[int, float] = attr.ib(kw_only=True, default=180.0)
    poll_pause: Union[int, float] = attr.ib(kw_only=True, default=0.7)
    # internal
    db_is_up: dict[FormulaDbConfig, Db] = attr.ib(init=False, factory=dict)
    first_call_time: Optional[float] = attr.ib(init=False, default=None)

    def ensure_db_is_up(self, db: Db) -> tuple[bool, str]:
        check_value_in = 123
        try:
            check_value_res = db.scalar(sa.literal(check_value_in))
        except Exception as exc:
            return False, str(exc)
        else:
            assert check_value_res == check_value_in
            return True, ""

    def make_database(self, db_config: FormulaDbConfig) -> Db:
        # Share the 'max time to up' between all dialects:
        if self.first_call_time is None:
            self.first_call_time = time.monotonic()

        assert self.first_call_time is not None

        timelimit = self.first_call_time + self.global_timeout
        db = make_db_from_config(db_config)
        wait_condition = partial(self.ensure_db_is_up, db=db)
        logging.info(f"Waiting for db to come up: {db_config!r}")
        wait_for("Make database", condition=wait_condition, timeout=timelimit, interval=self.poll_pause)

        return db

    def get_database(self, db_config: FormulaDbConfig) -> Db:
        if db_config in self.db_is_up:
            return self.db_is_up[db_config]  # supposed to be alive already.

        db = self.make_database(db_config)
        self.db_is_up[db_config] = db
        return db

lib/dl_formula_testing/dl_formula_testing/database.py
datalens-tech/datalens-backend
from __future__ import annotations

import copy
import datetime
import re
import time
from typing import (
    Collection,
    Optional,
    Sequence,
    Type,
    Union,
)

import attr
import sqlalchemy as sa
from sqlalchemy.sql.elements import ClauseElement

from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo
from dl_formula.core.nodes import Formula
from dl_formula.definitions.flags import ContextFlag
from dl_formula.definitions.scope import Scope
import dl_formula.inspect.expression
from dl_formula.mutation.mutation import apply_mutations
from dl_formula.mutation.window import (
    AmongToWithinGroupingMutation,
    DefaultWindowOrderingMutation,
)
from dl_formula.parser.factory import (
    ParserType,
    get_parser,
)
from dl_formula.translation.ext_nodes import CompiledExpression
from dl_formula.translation.translator import (
    TranslationCtx,
    translate,
)
from dl_formula_testing.database import Db
from dl_formula_testing.forced_literal import forced_literal_use  # noqa


FIELD_TYPES = {
    # for the above table
    "id": DataType.INTEGER,
    "int_value": DataType.INTEGER,
    "date_value": DataType.DATE,
    "datetime_value": DataType.DATETIME,
    "str_value": DataType.STRING,
    "str_null_value": DataType.STRING,
    "arr_int_value": DataType.ARRAY_INT,
    "arr_float_value": DataType.ARRAY_FLOAT,
    "arr_str_value": DataType.ARRAY_STR,
    # for NULL data
    "int_null": DataType.INTEGER,
    "float_null": DataType.FLOAT,
    "bool_null": DataType.BOOLEAN,
    "str_null": DataType.STRING,
    "date_null": DataType.DATE,
    "datetime_null": DataType.DATETIME,
    "geopoint_null": DataType.GEOPOINT,
    "geopolygon_null": DataType.GEOPOLYGON,
    "uuid_null": DataType.UUID,
}


@attr.s
class DbEvaluator:
    db: Db = attr.ib(kw_only=True)
    attempts: int = attr.ib(kw_only=True, default=1)
    retry_on_exceptions: Collection[tuple[Type[Exception], re.Pattern]] = attr.ib(kw_only=True, default=())
    retry_delay: int = attr.ib(kw_only=True, default=5)

    @property
    def dialect(self) -> DialectCombo:
        return self.db.dialect

    def translate_formula(
        self,
        formula: str | Formula,
        context_flags: Optional[int] = None,
        other_fields: Optional[dict] = None,
        collect_errors: Optional[bool] = None,
        field_types: Optional[Sequence[DataType]] = None,
        group_by: Optional[list[str | Formula]] = None,
        order_by: Optional[list[str | Formula]] = None,
        required_scopes: int = Scope.EXPLICIT_USAGE,
    ) -> TranslationCtx:
        other_fields = other_fields or {}
        parser = get_parser(ParserType.antlr_py)
        if field_types is None:
            field_types = FIELD_TYPES  # type: ignore
        if isinstance(formula, str):
            formula = parser.parse(formula)
        if isinstance(formula, Formula):
            # substitute other fields
            for field_node_idx, field_node in dl_formula.inspect.expression.enumerate_fields(copy.copy(formula)):
                if field_node.name in other_fields:
                    formula = formula.replace_at_index(
                        index=field_node_idx,
                        expr=CompiledExpression.make(self.translate_formula((other_fields[field_node.name]))),
                    )

            # mutate
            mutations = []
            group_by_objs = [parser.parse(expr).expr for expr in (group_by or ())]  # type: ignore
            mutations.append(AmongToWithinGroupingMutation(global_dimensions=group_by_objs))
            if order_by is not None:
                order_by_objs = [parser.parse(expr).expr for expr in order_by]  # type: ignore
                mutations.append(DefaultWindowOrderingMutation(default_order_by=order_by_objs))  # type: ignore
            formula = apply_mutations(formula, mutations=mutations)

            # translate
            formula = translate(  # type: ignore
                formula=formula,
                dialect=self.dialect,
                field_types=field_types,  # type: ignore
                context_flags=context_flags,
                collect_errors=collect_errors,
                required_scopes=required_scopes,
            )
        if not isinstance(formula, TranslationCtx):
            raise TypeError(type(formula))

        return formula

    @staticmethod
    def print_as_example(formula: Union[str, Formula], result) -> None:
        if isinstance(result, datetime.date):
            result_str = "#{}#".format(str(result))
        elif isinstance(result, bool):
            result_str = str(result).upper()
        else:
            result_str = repr(result)
        print("{},".format(repr("{} = {}".format(formula, result_str))))

    def eval(
        self,
        formula: Union[str, Formula],
        from_: Optional[ClauseElement] = None,
        where: str | Formula | None = None,
        many: bool = False,
        other_fields: Optional[dict] = None,
        order_by: Optional[list[str | Formula]] = None,
        group_by: Optional[list[str | Formula]] = None,
        first: bool = False,
        required_scopes: int = Scope.EXPLICIT_USAGE,
    ):
        select_ctx = self.translate_formula(
            formula,
            other_fields=other_fields,
            order_by=order_by,
            group_by=group_by,
            required_scopes=required_scopes,
        )

        def convert(val):
            # a workaround for missing boolean type in DBs and no coercion to bool
            if select_ctx.data_type in (DataType.CONST_BOOLEAN, DataType.BOOLEAN) and val in (0, 1):
                val = bool(val)
            return val

        query = sa.select([select_ctx.expression])
        if from_ is not None:
            query = query.select_from(from_)
        if order_by is not None:
            query = query.order_by(
                *[self.translate_formula(expr, required_scopes=required_scopes).expression for expr in order_by]
            )
        if group_by:
            query = query.group_by(
                *[self.translate_formula(expr, required_scopes=required_scopes).expression for expr in group_by]
            )
        if where is not None:
            where_ctx = self.translate_formula(
                where,
                context_flags=ContextFlag.REQ_CONDITION,
                required_scopes=required_scopes,
            )
            query = query.where(where_ctx.expression)
        if first:
            query = query.limit(1)

        try:
            print("QUERY:", self.db.expr_as_str(query))
            for attempt in range(self.attempts):
                try:
                    if many:
                        return [convert(row[0]) for row in self.db.execute(query).fetchall()]
                    else:
                        result = convert(self.db.execute(query).scalar())
                        return result
                except Exception as exc:
                    exc_str = str(exc)
                    if attempt < self.attempts - 1 and (
                        any(
                            isinstance(exc, match_exc) and match_patt.search(exc_str)
                            for match_exc, match_patt in self.retry_on_exceptions
                        )
                    ):
                        time.sleep(self.retry_delay)
                        continue
                    raise

        except Exception:
            print("QUERY:", self.db.expr_as_str(query))
            raise

lib/dl_formula_testing/dl_formula_testing/evaluator.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import contextmanager
from typing import Generator

import pytest

from dl_formula.core.dialect import StandardDialect as D
from dl_formula.definitions.base import (
    MultiVariantTranslation,
    TranslationVariant,
)
from dl_formula.definitions.registry import OPERATION_REGISTRY
from dl_formula.definitions.type_strategy import FromArgs


V = TranslationVariant.make


class ForcedLiteral(MultiVariantTranslation):
    """
    A hack wrapper function to avoid a value being recognized as a constant.
    Originally it wrapped the constant value in a ``Literal`` node, hence the name.
    Now it just wraps it in a function call, so the value loses its ``CONST_`` quality.
    TODO: rename
    """

    name = "__lit__"
    arg_cnt = 1
    variants = [V(D.ANY, lambda x: x)]
    return_type = FromArgs()  # will automatically convert from const to non-const type

    @classmethod
    @contextmanager
    def use(cls):
        def_item = cls()
        OPERATION_REGISTRY.register(def_item)
        try:
            yield
        finally:
            OPERATION_REGISTRY.unregister(def_item)


@pytest.yield_fixture
def forced_literal_use() -> Generator[None, None, None]:
    with ForcedLiteral.use():
        yield None

lib/dl_formula_testing/dl_formula_testing/forced_literal.py
datalens-tech/datalens-backend
import datetime
from typing import (
    Any,
    NamedTuple,
)

from sqlalchemy.types import TypeEngine

from dl_formula.connectors.base.type_constructor import get_type_constructor
from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo


DAY = datetime.timedelta(days=1)
START_DATE = datetime.date(2014, 10, 5)
START_DATETIME = datetime.datetime(2014, 10, 5)
DATETIME_DELTA = datetime.timedelta(days=1, hours=1, minutes=10, seconds=30)


class ColumnSpec(NamedTuple):
    name: str
    data_type: DataType
    primary_key: bool = False


TABLE_SPEC = (
    ColumnSpec(name="id", data_type=DataType.INTEGER, primary_key=True),
    ColumnSpec(name="int_value", data_type=DataType.INTEGER),
    ColumnSpec(name="date_value", data_type=DataType.DATE),
    ColumnSpec(name="datetime_value", data_type=DataType.DATETIME),
    ColumnSpec(name="str_value", data_type=DataType.STRING),
    ColumnSpec(name="str_null_value", data_type=DataType.STRING),
)


TABLE_SPEC_ARRAYS = (
    ColumnSpec(name="arr_int_value", data_type=DataType.ARRAY_INT),
    ColumnSpec(name="arr_float_value", data_type=DataType.ARRAY_FLOAT),
    ColumnSpec(name="arr_str_value", data_type=DataType.ARRAY_STR),
)


NULL_DATA_TABLE_SPEC = (
    ColumnSpec(name="int_null", data_type=DataType.INTEGER),
    ColumnSpec(name="float_null", data_type=DataType.FLOAT),
    ColumnSpec(name="bool_null", data_type=DataType.BOOLEAN),
    ColumnSpec(name="str_null", data_type=DataType.STRING),
    ColumnSpec(name="date_null", data_type=DataType.DATE),
    ColumnSpec(name="datetime_null", data_type=DataType.DATETIME),
    ColumnSpec(name="geopoint_null", data_type=DataType.GEOPOINT),
    ColumnSpec(name="geopolygon_null", data_type=DataType.GEOPOLYGON),
    ColumnSpec(name="uuid_null", data_type=DataType.UUID),
)


def generate_sample_data(
    start_date: datetime.date = START_DATE,
    start_datetime: datetime.datetime = START_DATETIME,
    datetime_delta: datetime.timedelta = DATETIME_DELTA,
    add_arrays: bool = False,
) -> list[dict[str, Any]]:
    data = []
    for num, idx in enumerate([*range(10), *range(10)]):
        row = dict(
            id=num + 1,
            int_value=idx * 10,
            date_value=start_date + DAY * idx,
            datetime_value=start_datetime + datetime_delta * idx,
            str_value="q" * (idx + 1),
            str_null_value=None,
        )
        if add_arrays:
            row.update(
                arr_int_value=[1 * idx, 23 + idx, 456, None],
                arr_float_value=[1.23 * idx, 45.0 + idx, 0.123, None],
                arr_str_value=["a" * idx, "bb" * idx, "cde", None],
            )
        data.append(row)
    return data


def get_column_sa_type(data_type: DataType, dialect: DialectCombo) -> TypeEngine:
    type_constructor = get_type_constructor(dialect_name=dialect.common_name)
    return type_constructor.get_sa_type(data_type)

lib/dl_formula_testing/dl_formula_testing/table.py
datalens-tech/datalens-backend
import abc
import contextlib
import datetime
import re
from typing import (
    Any,
    ClassVar,
    Generator,
    Sequence,
    Type,
)
import uuid

import pytest
import pytz
import sqlalchemy as sa
from sqlalchemy.types import TypeEngine

from dl_db_testing.database.base import DbBase
from dl_db_testing.database.dispenser import DbDispenserBase
from dl_db_testing.database.engine_wrapper import (
    DbEngineConfig,
    TableSpec,
)
from dl_formula.core.datatype import DataType
from dl_formula.core.dialect import DialectCombo
from dl_formula_testing.database import (
    FormulaDbConfig,
    FormulaDbDispenser,
)
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.table import (
    NULL_DATA_TABLE_SPEC,
    TABLE_SPEC,
    TABLE_SPEC_ARRAYS,
    ColumnSpec,
    generate_sample_data,
    get_column_sa_type,
)


class FormulaConnectorTestBase(metaclass=abc.ABCMeta):
    dialect: ClassVar[DialectCombo]
    eval_attempts: ClassVar[int] = 1
    retry_on_exceptions: ClassVar[tuple[tuple[Type[Exception], re.Pattern], ...]] = ()

    supports_arrays: ClassVar[bool]
    supports_uuid: ClassVar[bool]
    bool_is_expression: ClassVar[bool] = False
    empty_str_is_null: ClassVar[bool] = False
    null_casts_to_number: ClassVar[bool] = False
    null_casts_to_false: ClassVar[bool] = False
    db_dispenser: ClassVar[DbDispenserBase] = FormulaDbDispenser()
    engine_config_cls: ClassVar[Type[DbEngineConfig]] = DbEngineConfig

    @abc.abstractmethod
    @pytest.fixture(scope="class")
    def db_url(self) -> str:
        pass

    @pytest.fixture(scope="class")
    def engine_params(self) -> dict:
        return {}

    @pytest.fixture(scope="class")
    def tzinfo(self) -> datetime.tzinfo:
        return pytz.UTC

    @pytest.fixture(scope="class")
    def engine_config(self, db_url: str, engine_params: dict) -> DbEngineConfig:
        return self.engine_config_cls(url=db_url, engine_params=engine_params)

    @pytest.fixture(scope="class")
    def db_config(self, engine_config: DbEngineConfig, tzinfo: datetime.tzinfo) -> FormulaDbConfig:
        return FormulaDbConfig(
            engine_config=engine_config,
            dialect=self.dialect,
            tzinfo=tzinfo,
        )

    @pytest.fixture(scope="class")
    def dbe(self, db_config: FormulaDbConfig) -> DbEvaluator:
        db = self.db_dispenser.get_database(db_config)
        dbe = DbEvaluator(
            db=db,
            attempts=self.eval_attempts,
            retry_on_exceptions=self.retry_on_exceptions,
        )
        return dbe

    def lowlevel_make_sa_table(self, db: DbBase, table_spec: TableSpec, columns: Sequence[sa.Column]) -> sa.Table:
        table = db.table_from_columns(table_name=table_spec.table_name, columns=columns)
        return table

    def get_column_sa_type(self, data_type: DataType) -> TypeEngine:
        return get_column_sa_type(data_type=data_type, dialect=self.dialect)

    def make_columns(self, column_specs: Sequence[ColumnSpec]) -> list[sa.Column]:
        columns = [
            sa.Column(name=spec.name, type_=self.get_column_sa_type(data_type=spec.data_type)) for spec in column_specs
        ]
        return columns

    def generate_table_name(self, prefix: str) -> str:
        return f"{prefix}_{uuid.uuid4().hex[:6]}"

    def generate_table_spec(self, table_name_prefix: str) -> TableSpec:
        return TableSpec(
            table_name=self.generate_table_name(prefix=table_name_prefix),
        )

    def make_sa_table(self, db: DbBase, table_spec: TableSpec) -> sa.Table:
        column_specs = list(TABLE_SPEC)
        if self.supports_arrays:
            column_specs += TABLE_SPEC_ARRAYS

        columns = self.make_columns(column_specs)
        return self.lowlevel_make_sa_table(db=db, table_spec=table_spec, columns=columns)

    @pytest.fixture(scope="class")
    def data_table(self, dbe: DbEvaluator) -> Generator[sa.Table, None, None]:
        with self.make_data_table(dbe=dbe) as table:
            yield table

    @contextlib.contextmanager
    def make_data_table(self, dbe: DbEvaluator) -> Generator[sa.Table, None, None]:
        db = dbe.db
        table_spec = self.generate_table_spec(table_name_prefix="test_table")
        table = self.make_sa_table(db=dbe.db, table_spec=table_spec)
        db.create_table(table)

        table_data = generate_sample_data(add_arrays=self.supports_arrays)
        db.insert_into_table(table, table_data)

        table.int_values = [row["int_value"] for row in table_data]
        table.date_values = [row["date_value"] for row in table_data]
        table.datetime_values = [row["datetime_value"] for row in table_data]

        try:
            yield table
        finally:
            dbe.db.drop_table(table)

    @pytest.fixture(scope="class")
    def null_data_table(self, dbe: DbEvaluator) -> Generator[sa.Table, None, None]:
        with self.make_null_data_table(dbe=dbe) as table:
            yield table

    @contextlib.contextmanager
    def make_null_data_table(self, dbe: DbEvaluator) -> Generator[sa.Table, None, None]:
        column_specs = NULL_DATA_TABLE_SPEC
        table_spec = self.generate_table_spec("null_test_table")

        columns = self.make_columns(column_specs)
        table = self.lowlevel_make_sa_table(db=dbe.db, table_spec=table_spec, columns=columns)
        table_data = [{column.name: None for column in column_specs}]

        dbe.db.create_table(table)
        dbe.db.insert_into_table(table, table_data)

        try:
            yield table
        finally:
            dbe.db.drop_table(table)

    @contextlib.contextmanager
    def make_scalar_table(
        self, dbe: DbEvaluator, col_name: str, data_type: DataType, value: Any
    ) -> Generator[sa.Table, None, None]:
        table_spec = self.generate_table_spec(table_name_prefix="scalar_test_table")
        column_specs = [ColumnSpec(name=col_name, data_type=data_type)]
        columns = self.make_columns(column_specs)
        table = self.lowlevel_make_sa_table(db=dbe.db, table_spec=table_spec, columns=columns)
        table_data = [{col_name: value}]
        dbe.db.create_table(table)
        dbe.db.insert_into_table(table, table_data)

        try:
            yield table
        finally:
            dbe.db.drop_table(table)

lib/dl_formula_testing/dl_formula_testing/testcases/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_formula.core.datatype import DataType
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import to_str


class DefaultConditionalBlockFormulaConnectorTestSuite(FormulaConnectorTestBase):
    def test_if_block(self, dbe: DbEvaluator):
        # simple
        assert to_str(dbe.eval('IF 2 = 2 THEN "1st" ELSE "2nd" END')) == "1st"
        assert to_str(dbe.eval('IF 2 = 1 THEN "1st" ELSE "2nd" END')) == "2nd"

        # with ELSIF
        assert to_str(dbe.eval('IF 1 = 1 THEN "1st" ELSEIF 1 = 2 THEN "2nd" ELSE "3rd" END')) == "1st"
        assert to_str(dbe.eval('IF 2 = 1 THEN "1st" ELSEIF 2 = 2 THEN "2nd" ELSE "3rd" END')) == "2nd"
        assert to_str(dbe.eval('IF 3 = 1 THEN "1st" ELSEIF 3 = 2 THEN "2nd" ELSE "3rd" END')) == "3rd"

        # without ELSE
        assert to_str(dbe.eval('IF 2 = 2 THEN "1st" END')) == "1st"
        assert dbe.eval('IF 2 = 1 THEN "1st" END') is None
        assert to_str(dbe.eval('IF 1 = 1 THEN "1st" ELSEIF 2 = 2 THEN "2nd" END')) == "1st"
        assert to_str(dbe.eval('IF 1 = 2 THEN "1st" ELSEIF 2 = 2 THEN "2nd" END')) == "2nd"
        assert dbe.eval('IF 1 = 2 THEN "1st" ELSEIF 2 = 1 THEN "2nd" END') is None

    def test_case_block_basic(self, dbe: DbEvaluator):
        # "normal"
        assert to_str(dbe.eval('CASE 1 WHEN 1 THEN "1st" WHEN 2 THEN "2nd" ELSE "3rd" END')) == "1st"
        assert to_str(dbe.eval('CASE 2 WHEN 1 THEN "1st" WHEN 2 THEN "2nd" ELSE "3rd" END')) == "2nd"
        assert to_str(dbe.eval('CASE 3 WHEN 1 THEN "1st" WHEN 2 THEN "2nd" ELSE "3rd" END')) == "3rd"

        # without ELSE
        assert to_str(dbe.eval('CASE 1 WHEN 1 THEN "1st" WHEN 2 THEN "2nd" ELSE NULL END')) == "1st"
        assert to_str(dbe.eval('CASE 2 WHEN 1 THEN "1st" WHEN 2 THEN "2nd" END')) == "2nd"

    def test_case_block_returning_null(self, dbe: DbEvaluator):
        assert dbe.eval('CASE 3 WHEN 1 THEN "1st" WHEN 2 THEN "2nd" END') is None
        assert dbe.eval("CASE 3 WHEN 1 THEN 1 WHEN 2 THEN 2 END") is None
        assert dbe.eval("CASE 3 WHEN 1 THEN 1.1 WHEN 2 THEN 2.2 END") is None
        assert dbe.eval("CASE 3 WHEN 1 THEN TRUE WHEN 2 THEN TRUE END") is None
        # NULL in THEN
        assert to_str(dbe.eval('CASE 2 WHEN 1 THEN NULL WHEN 2 THEN "2nd" ELSE "3rd" END')) == "2nd"
        assert dbe.eval('CASE 1 WHEN 1 THEN NULL WHEN 2 THEN "2nd" ELSE "3rd" END') is None

    def test_case_non_const_then(self, dbe: DbEvaluator):
        with self.make_scalar_table(dbe, col_name="int_value", data_type=DataType.INTEGER, value=1) as scalar:
            assert dbe.eval("CASE [int_value] WHEN 1 THEN [int_value]*10 WHEN 2 THEN 8 ELSE 0 END", from_=scalar) == 10

lib/dl_formula_testing/dl_formula_testing/testcases/conditional_blocks.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    ClassVar,
    Sequence,
)

from dl_formula.core.dialect import (
    DialectCombo,
    DialectName,
    from_name_and_version,
)


class DefaultDialectFormulaConnectorTestSuite:
    dialect_name: ClassVar[DialectName]
    default_dialect: DialectCombo
    dialect_matches: ClassVar[Sequence[tuple[str, DialectCombo]]]

    def test_dialect_versions(self):
        for version_str, expected_dialect in self.dialect_matches:
            assert from_name_and_version(self.dialect_name, version_str) == expected_dialect

    def test_default_dialect(self):
        assert from_name_and_version(self.dialect_name, None) == self.default_dialect

lib/dl_formula_testing/dl_formula_testing/testcases/dialect.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from typing import (
    ClassVar,
    Collection,
    TypeVar,
    Union,
)

import pytest
import sqlalchemy as sa

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import utc_ts


class DefaultMainAggFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    supports_countd_approx: ClassVar[bool] = False
    supports_quantile: ClassVar[bool] = False
    supports_median: ClassVar[bool] = False
    supports_arg_min_max: ClassVar[bool] = False
    supports_any: ClassVar[bool] = False
    supports_all_concat: ClassVar[bool] = False
    supports_top_concat: ClassVar[bool] = False

    def test_basic_aggregation_functions(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values

        assert dbe.eval("SUM([int_value])", from_=data_table) == sum(values)
        assert dbe.eval("AVG([int_value])", from_=data_table) == sum(values) / len(values)
        assert dbe.eval("MAX([int_value])", from_=data_table) == max(values)
        assert dbe.eval("MIN([int_value])", from_=data_table) == min(values)
        assert dbe.eval("COUNT()", from_=data_table) == len(values)
        assert dbe.eval("COUNT([int_value])", from_=data_table) == len(values)
        assert dbe.eval("COUNTD([int_value])", from_=data_table) == len(set(values))

    def test_date_avg_function(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        date_values = data_table.date_values

        assert dbe.eval("AVG([date_value])", from_=data_table) == datetime.date.fromtimestamp(
            sum(map(lambda date: utc_ts(date), date_values)) / len(date_values)
        )

    def test_datetime_avg_function(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        datetime_values = data_table.datetime_values

        def dt_avg(values) -> datetime.datetime:
            tses = [utc_ts(dt) for dt in values]
            ts = sum(tses) / len(tses)
            return datetime.datetime.utcfromtimestamp(ts)

        # Sanity check:
        assert dt_avg(datetime_values[:1]) == datetime_values[0]
        result = dbe.eval("AVG([datetime_value])", from_=data_table)
        # XXX: shouldn't this have `astimezone` first?
        result = result.replace(tzinfo=None)
        assert result == dt_avg(datetime_values)

        # To make sure that the everchanging DT types all work:
        result = dbe.eval("AVG(DATETIME([datetime_value]))", from_=data_table)
        result = result.replace(tzinfo=None)
        assert result == dt_avg(datetime_values)
        result = dbe.eval("AVG(GENERICDATETIME([datetime_value]))", from_=data_table)
        result = result.replace(tzinfo=None)
        assert result == dt_avg(datetime_values)

    def test_statistical_aggregation_functions(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values

        avg = sum(values) / len(values)
        var_samp = sum([(v - avg) ** 2 for v in values]) / (len(values) - 1)
        assert dbe.eval("VAR([int_value])", from_=data_table) == pytest.approx(var_samp)
        assert dbe.eval("STDEV([int_value])", from_=data_table) == pytest.approx(var_samp**0.5)
        var_pop = sum([(v - avg) ** 2 for v in values]) / len(values)
        assert dbe.eval("VARP([int_value])", from_=data_table) == pytest.approx(var_pop)
        assert dbe.eval("STDEVP([int_value])", from_=data_table) == pytest.approx(var_pop**0.5)

    def test_approximate_aggregation_functions(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_countd_approx:
            pytest.skip()

        values = data_table.int_values
        assert dbe.eval("COUNTD_APPROX([int_value])", from_=data_table) == len(set(values))

    def test_conditional_aggregation_functions(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values

        assert dbe.eval("SUM_IF([int_value], [int_value] > 5)", from_=data_table) == sum([v for v in values if v > 5])
        assert dbe.eval("SUM_IF([int_value] + 2 * 2, [int_value] > 5)", from_=data_table) == sum(
            [v + 2 * 2 for v in values if v > 5]
        )

        assert dbe.eval("AVG_IF([int_value], [int_value] > 5)", from_=data_table) == (lambda x: sum(x) / len(x))(
            [v for v in values if v > 5]
        )
        assert dbe.eval("AVG_IF([int_value] + 2 * 2, [int_value] > 5)", from_=data_table) == (
            lambda x: sum(x) / len(x)
        )([v + 2 * 2 for v in values if v > 5])

        # Zero matching elements
        assert (
            dbe.eval(
                "AVG_IF([int_value], [int_value] > {max_val} and [int_value] < {min_val})".format(
                    max_val=max(values), min_val=min(values)
                ),
                from_=data_table,
            )
            is None
        )

        assert dbe.eval("COUNT_IF([int_value] > 5)", from_=data_table) == len([v for v in values if v > 5])
        assert dbe.eval("COUNT_IF([int_value] + 1 > 5)", from_=data_table) == len([v for v in values if (v + 1) > 5])
        assert dbe.eval("COUNT_IF([int_value] > 5)", from_=data_table, where="1 = 0") == 0

        assert dbe.eval("COUNTD_IF([int_value], [int_value] > 5)", from_=data_table) == len(
            {v for v in values if v > 5}
        )
        assert dbe.eval("COUNTD_IF([int_value] % 2, [int_value] - 1 > 5)", from_=data_table) == len(
            {v % 2 for v in values if (v - 1) > 5}
        )

    def test_quantile(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_quantile:
            pytest.skip()

        value = dbe.eval("QUANTILE([int_value], 0.9)", from_=data_table)
        assert 80 <= value <= 90

    def test_median(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_median:
            pytest.skip()

        date_values = data_table.date_values
        datetime_values = data_table.datetime_values

        VALUE_TV = TypeVar("VALUE_TV")

        def medians(values: Collection[Union[VALUE_TV]]) -> tuple[VALUE_TV, VALUE_TV]:
            values = sorted(values)  # type: ignore
            middle = len(values) // 2 - 1
            return (values[middle], values[middle + 1])

        assert dbe.eval("MEDIAN([int_value])", from_=data_table) in (40, 50)
        assert dbe.eval("MEDIAN([date_value])", from_=data_table) in medians(date_values)
        assert dbe.eval("MEDIAN([datetime_value])", from_=data_table) in medians(datetime_values)

    def test_any(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_any:
            pytest.skip()

        int_values = data_table.int_values
        assert dbe.eval("ANY([int_value])", from_=data_table) in int_values

    def test_arg_min_max(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_arg_min_max:
            pytest.skip()

        assert dbe.eval('ARG_MIN(CONCAT("q_", [int_value]), (50-[int_value])^2)', from_=data_table) == "q_50"
        assert dbe.eval('ARG_MAX(CONCAT("q_", [int_value]), 50-[int_value]^2)', from_=data_table) == "q_0"

    def test_all_concat(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_all_concat:
            pytest.skip()

        expected_values = sorted(str(val) for val in set(data_table.int_values))
        assert dbe.eval("ALL_CONCAT([int_value])", from_=data_table) == ", ".join(expected_values)
        assert dbe.eval("ALL_CONCAT([int_value], '')", from_=data_table) == "".join(expected_values)
        assert dbe.eval("ALL_CONCAT([int_value], ' – ')", from_=data_table) == " – ".join(expected_values)

    def test_top_concat(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_top_concat:
            pytest.skip()

        expected_values = set(str(val) for val in set(data_table.int_values))
        res = dbe.eval("TOP_CONCAT([int_value], 3)", from_=data_table)
        assert not set(res.split(", ")) - expected_values, "should be a subset"
        res = dbe.eval("TOP_CONCAT([int_value], 3, '; ')", from_=data_table)
        assert not set(res.split("; ")) - expected_values, "should be a subset"

lib/dl_formula_testing/dl_formula_testing/testcases/functions_aggregation.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    Any,
    ClassVar,
    Optional,
    Union,
)

import pytest
import sqlalchemy as sa

import dl_formula.core.exc as exc
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import to_str


SAMPLE_DATA_ARRAYS_LENGTH = 4


class DefaultArrayFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    make_decimal_cast: ClassVar[Optional[str]] = None
    make_float_cast: ClassVar[Optional[str]] = None

    def test_create_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        sample_array: tuple[Union[str, int, float, None], ...]
        expected: tuple[Union[str, int, float, None], ...]
        item: Any
        exp: Any
        for sample_array in (
            (0, 23, 456, None),
            (0, 43, 0.123, None),
        ):
            array_string = ",".join("NULL" if item is None else str(item) for item in sample_array)
            for i, item in enumerate(sample_array):
                assert dbe.eval(f"GET_ITEM(ARRAY({array_string}), {i + 1})", from_=data_table) == item

        # testing string separately because they must be quoted
        sample_array = ("", "NULL", "cde", None)
        array_string = ",".join("NULL" if item is None else f"'{item}'" for item in sample_array)
        for i, item in enumerate(sample_array):
            assert dbe.eval(f"GET_ITEM(ARRAY({array_string}), {i + 1})", from_=data_table) == item

        # it is not allowed to create array with undefined element type
        with pytest.raises(exc.TranslationError):
            dbe.eval("ARRAY(NULL, NULL, NULL)")

        expected = (0, 1, None)
        for i, exp in enumerate(expected):
            assert dbe.eval(f"GET_ITEM(ARRAY([int_value], 1, NULL), {i + 1})", from_=data_table) == exp

        expected = ("q", "1", None)
        for i, exp in enumerate(expected):
            assert dbe.eval(f"GET_ITEM(ARRAY([str_value], STR(1), NULL), {i + 1})", from_=data_table) == exp

    def test_unnest_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        for array_field, inner_type in (
            ("arr_int_value", int),
            ("arr_float_value", float),
            ("arr_str_value", str),
        ):
            count_before_unnest = len(dbe.eval(f"[{array_field}]", from_=data_table, many=True))
            res = dbe.eval(f"UNNEST([{array_field}])", from_=data_table, many=True)
            assert len(res) == count_before_unnest * SAMPLE_DATA_ARRAYS_LENGTH
            assert all(isinstance(value, inner_type) or value is None for value in res)

    def test_array_len(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("LEN([arr_int_value])", from_=data_table) == SAMPLE_DATA_ARRAYS_LENGTH
        assert dbe.eval("LEN([arr_float_value])", from_=data_table) == SAMPLE_DATA_ARRAYS_LENGTH
        assert dbe.eval("LEN([arr_str_value])", from_=data_table) == SAMPLE_DATA_ARRAYS_LENGTH

    def test_array_get_item(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("GET_ITEM([arr_int_value], 3)", from_=data_table) == 456
        assert dbe.eval("GET_ITEM([arr_float_value], 3)", from_=data_table) == 0.123
        assert dbe.eval("GET_ITEM([arr_str_value], 3)", from_=data_table) == "cde"

    def test_array_str(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert to_str(dbe.eval("ARR_STR([arr_int_value])", from_=data_table)) == "0,23,456"
        assert to_str(dbe.eval("ARR_STR([arr_float_value])", from_=data_table)) == "0,45,0.123"
        assert to_str(dbe.eval("ARR_STR([arr_str_value])", from_=data_table)) == ",,cde"

        assert to_str(dbe.eval("ARR_STR([arr_int_value], ';')", from_=data_table)) == "0;23;456"
        assert to_str(dbe.eval("ARR_STR([arr_float_value], ';')", from_=data_table)) == "0;45;0.123"
        assert to_str(dbe.eval("ARR_STR([arr_str_value], ';')", from_=data_table)) == ";;cde"

        assert to_str(dbe.eval("ARR_STR([arr_int_value], ';', '*')", from_=data_table)) == "0;23;456;*"
        assert to_str(dbe.eval("ARR_STR([arr_float_value], ';', '*')", from_=data_table)) == "0;45;0.123;*"
        assert to_str(dbe.eval("ARR_STR([arr_str_value], ';', '*')", from_=data_table)) == ";;cde;*"

    def test_array_count_item(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("COUNT_ITEM([arr_int_value], 23)", from_=data_table) == 1
        assert dbe.eval("COUNT_ITEM([arr_float_value], 45)", from_=data_table) == 1
        assert dbe.eval('COUNT_ITEM([arr_str_value], "cde")', from_=data_table) == 1
        assert dbe.eval("COUNT_ITEM([arr_str_value], NULL)", from_=data_table) == 1

        assert dbe.eval("COUNT_ITEM(ARRAY(1, 2, 3), NULL)", from_=data_table) == 0
        assert dbe.eval("COUNT_ITEM(ARRAY(1, NULL, NULL), NULL)", from_=data_table) == 2

    def test_array_contains(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("CONTAINS(ARRAY(1, 2, 3), 1)", from_=data_table)
        assert dbe.eval("CONTAINS(ARRAY(1.1, 2.2, 3.3), 3.3)", from_=data_table)
        assert dbe.eval('CONTAINS(ARRAY("a", "b", "c"), "a")', from_=data_table)
        assert not dbe.eval('CONTAINS(ARRAY("a", "b", "c"), "d")', from_=data_table)
        assert dbe.eval('CONTAINS(ARRAY("a", NULL, "c"), NULL)', from_=data_table)
        assert not dbe.eval("CONTAINS(ARRAY(1.1, 2.2, 3.3), NULL)", from_=data_table)

        assert dbe.eval("CONTAINS([arr_int_value], 23)", from_=data_table)
        assert not dbe.eval("CONTAINS([arr_int_value], 24)", from_=data_table)
        assert dbe.eval('CONTAINS([arr_str_value], "cde")', from_=data_table)
        assert dbe.eval("CONTAINS([arr_str_value], NULL)", from_=data_table)

        assert dbe.eval("CONTAINS([arr_str_value], GET_ITEM([arr_str_value], 1))", from_=data_table)
        assert dbe.eval("CONTAINS([arr_str_value], GET_ITEM([arr_str_value], 2))", from_=data_table)
        assert dbe.eval("CONTAINS([arr_str_value], GET_ITEM([arr_str_value], 3))", from_=data_table)
        assert dbe.eval("CONTAINS([arr_str_value], GET_ITEM([arr_str_value], 4))", from_=data_table)

    def test_array_contains_all(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("CONTAINS_ALL(ARRAY(1, 2, 3), ARRAY(1, 2))", from_=data_table)
        assert dbe.eval("CONTAINS_ALL(ARRAY(1.1, 2.2, 3.3), ARRAY(1.1, 2.2))", from_=data_table)
        assert dbe.eval('CONTAINS_ALL(ARRAY("a", "b", "c"), ARRAY("a", "b"))', from_=data_table)

        assert dbe.eval("CONTAINS_ALL(ARRAY(1, 2, 3), ARRAY(2, 1))", from_=data_table)
        assert dbe.eval("CONTAINS_ALL(ARRAY(1.1, 2.2, 3.3), ARRAY(3.3, 3.3, 3.3, 3.3, 1.1))", from_=data_table)
        assert dbe.eval('CONTAINS_ALL(ARRAY("a", "b", "c"), ARRAY("c"))', from_=data_table)

        assert not dbe.eval('CONTAINS_ALL(ARRAY("a", "b", "c"), ARRAY("a", "d"))', from_=data_table)
        assert dbe.eval('CONTAINS_ALL(ARRAY("a", NULL, "c"), ARRAY("a", NULL))', from_=data_table)
        assert dbe.eval('CONTAINS_ALL(ARRAY("a", NULL, "c"), ARRAY("a", "c"))', from_=data_table)
        assert not dbe.eval("CONTAINS_ALL(ARRAY(1.1, 2.2, 3.3), ARRAY(1.1, NULL))", from_=data_table)

        assert dbe.eval("CONTAINS_ALL([arr_int_value], ARRAY(23, 456))", from_=data_table)
        assert not dbe.eval("CONTAINS_ALL([arr_int_value], ARRAY(24, 456))", from_=data_table)

        assert dbe.eval("CONTAINS_ALL([arr_int_value], SLICE([arr_int_value], 2, 2))", from_=data_table)
        assert dbe.eval("CONTAINS_ALL(ARRAY(0, 23, 456, NULL, 123), [arr_int_value])", from_=data_table)

    def test_array_contains_all_string_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval('CONTAINS_ALL([arr_str_value], ARRAY("cde"))', from_=data_table)
        assert dbe.eval('CONTAINS_ALL([arr_str_value], ARRAY("cde", NULL))', from_=data_table)
        assert not dbe.eval('CONTAINS_ALL(ARRAY("cde"), [arr_str_value])', from_=data_table)

    def test_array_contains_any(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("CONTAINS_ANY(ARRAY(1, 2, 3), ARRAY(1, 5))", from_=data_table)
        assert dbe.eval("CONTAINS_ANY(ARRAY(1.1, 2.2, 3.3), ARRAY(1.1, 5.5))", from_=data_table)
        assert dbe.eval('CONTAINS_ANY(ARRAY("a", "b", "c"), ARRAY("a", "e"))', from_=data_table)

        assert not dbe.eval("CONTAINS_ANY(ARRAY(1, 2, 3), ARRAY(6, 5))", from_=data_table)
        assert not dbe.eval("CONTAINS_ANY(ARRAY(1.1, 2.2, 3.3), ARRAY(6.6, 5.5))", from_=data_table)
        assert not dbe.eval('CONTAINS_ANY(ARRAY("a", "b", "c"), ARRAY("f", "e"))', from_=data_table)

        assert dbe.eval('CONTAINS_ANY(ARRAY("a", NULL, "c"), ARRAY("e", NULL))', from_=data_table)
        assert dbe.eval('CONTAINS_ANY(ARRAY("a", NULL, "c"), ARRAY("e", "c"))', from_=data_table)
        assert not dbe.eval("CONTAINS_ANY(ARRAY(1.1, 2.2, 3.3), ARRAY(5.5, NULL))", from_=data_table)

        assert dbe.eval("CONTAINS_ANY([arr_int_value], ARRAY(24, 456))", from_=data_table)
        assert not dbe.eval("CONTAINS_ANY([arr_int_value], ARRAY(24, 457))", from_=data_table)

        assert dbe.eval("CONTAINS_ANY([arr_int_value], SLICE([arr_int_value], 2, 1))", from_=data_table)
        assert dbe.eval("CONTAINS_ANY(ARRAY(0, 23, 456, NULL, 123), [arr_int_value])", from_=data_table)

    def test_array_contains_any_string_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval('CONTAINS_ANY([arr_str_value], ARRAY("cde"))', from_=data_table)
        assert dbe.eval('CONTAINS_ANY([arr_str_value], ARRAY("123", NULL))', from_=data_table)
        assert dbe.eval('CONTAINS_ANY(ARRAY("cde"), [arr_str_value])', from_=data_table)

    def test_array_not_contains(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert not dbe.eval("NOTCONTAINS(ARRAY(1, 2, 3), 1)", from_=data_table)
        assert not dbe.eval("NOTCONTAINS(ARRAY(1.1, 2.2, 3.3), 3.3)", from_=data_table)
        assert not dbe.eval('NOTCONTAINS(ARRAY("a", "b", "c"), "a")', from_=data_table)
        assert dbe.eval('NOTCONTAINS(ARRAY("a", "b", "c"), "d")', from_=data_table)
        assert not dbe.eval('NOTCONTAINS(ARRAY("a", NULL, "c"), NULL)', from_=data_table)
        assert dbe.eval("NOTCONTAINS(ARRAY(1.1, 2.2, 3.3), NULL)", from_=data_table)
        assert not dbe.eval("NOTCONTAINS([arr_int_value], 23)", from_=data_table)
        assert dbe.eval("NOTCONTAINS([arr_int_value], 24)", from_=data_table)
        assert not dbe.eval('NOTCONTAINS([arr_str_value], "cde")', from_=data_table)
        assert not dbe.eval("NOTCONTAINS([arr_str_value], NULL)", from_=data_table)

        assert not dbe.eval("NOTCONTAINS([arr_str_value], GET_ITEM([arr_str_value], 1))", from_=data_table)
        assert not dbe.eval("NOTCONTAINS([arr_str_value], GET_ITEM([arr_str_value], 2))", from_=data_table)
        assert not dbe.eval("NOTCONTAINS([arr_str_value], GET_ITEM([arr_str_value], 3))", from_=data_table)
        assert not dbe.eval("NOTCONTAINS([arr_str_value], GET_ITEM([arr_str_value], 4))", from_=data_table)
        assert not dbe.eval("NOTCONTAINS([arr_str_value], GET_ITEM([arr_str_value], 4))", from_=data_table)

    def test_array_slice(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("SLICE([arr_int_value], 2, 2)", from_=data_table) == dbe.eval(
            "ARRAY(23, 456)", from_=data_table
        )
        assert dbe.eval("SLICE(ARRAY(1, 2, 3, 4), 2, 2)", from_=data_table) == dbe.eval("ARRAY(2, 3)", from_=data_table)
        assert dbe.eval("SLICE(ARRAY(1, 2, 3, 4), 2, 1)", from_=data_table) == dbe.eval("ARRAY(2)", from_=data_table)
        assert str(dbe.eval("SLICE(ARRAY(1, 2, 3, 4), 2, 0)", from_=data_table)) == "[]"

    def test_replace_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("REPLACE([arr_int_value], -2, 1000)", from_=data_table) == dbe.eval("ARRAY(0, 23, 456, NULL)")
        assert dbe.eval("REPLACE([arr_int_value], 456, 1000)", from_=data_table) == dbe.eval("ARRAY(0, 23, 1000, NULL)")
        assert dbe.eval("REPLACE(ARRAY(6, 6, 6), 6, 7)") == dbe.eval("ARRAY(7, 7, 7)")
        assert dbe.eval("REPLACE(ARRAY(NULL, NULL, 1), NULL, 1)") == dbe.eval("ARRAY(1, 1, 1)")
        assert dbe.eval('REPLACE(ARRAY("NULL", NULL, "NULL"), "NULL", "LUL")') == dbe.eval('ARRAY("LUL", NULL, "LUL")')
        assert dbe.eval("REPLACE(ARRAY(45, 1, 45), GET_ITEM(ARRAY(24, 45), 2), 22 + 2)") == dbe.eval("ARRAY(24, 1, 24)")

    def test_array_cast(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("CAST_ARR_INT([arr_float_value])", from_=data_table) == dbe.eval(
            "ARRAY(0, 45, 0, NULL)", from_=data_table
        )
        assert dbe.eval("CAST_ARR_INT([arr_int_value])", from_=data_table) == dbe.eval(
            "[arr_int_value]", from_=data_table
        )
        assert dbe.eval('CAST_ARR_INT(ARRAY("1", "3", "2", NULL))', from_=data_table) == dbe.eval(
            "ARRAY(1, 3, 2, NULL)", from_=data_table
        )

        assert dbe.eval("CAST_ARR_FLOAT([arr_float_value])", from_=data_table) == dbe.eval(
            "[arr_float_value]", from_=data_table
        )
        assert dbe.eval("CAST_ARR_FLOAT([arr_int_value])", from_=data_table) == dbe.eval(
            "ARRAY(0, 23.0, 456, NULL)", from_=data_table
        )
        assert dbe.eval('CAST_ARR_FLOAT(ARRAY("1", "3.3", "2", NULL))', from_=data_table) == dbe.eval(
            "ARRAY(1, 3.3, 2, NULL)", from_=data_table
        )

        assert dbe.eval("CAST_ARR_STR([arr_float_value])", from_=data_table) in (
            dbe.eval('ARRAY("0.0", "45.0", "0.123", NULL)', from_=data_table),
            dbe.eval('ARRAY("0", "45", "0.123", NULL)', from_=data_table),
        )
        assert dbe.eval("CAST_ARR_STR([arr_int_value])", from_=data_table) == dbe.eval(
            'ARRAY("0", "23", "456", NULL)', from_=data_table
        )
        assert dbe.eval("CAST_ARR_STR([arr_str_value])", from_=data_table) == dbe.eval(
            "[arr_str_value]", from_=data_table
        )

    def test_array_get_item_nested(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        src = "ARRAY('qwe', 'rty', 'uio')"

        assert dbe.eval(f"GET_ITEM({src}, 1)") == "qwe"
        assert dbe.eval("GET_ITEM([arr_int_value], 3)", from_=data_table) == 456

        assert dbe.eval(f"GET_ITEM(SLICE({src}, 2, 1), 1)") == "rty"
        assert dbe.eval("GET_ITEM(SLICE([arr_int_value], 2, 2), 2)", from_=data_table) == 456

        assert dbe.eval(f"SLICE(SLICE({src}, 2, 1), 1, 1)") == dbe.eval('ARRAY("rty")', from_=data_table)
        assert dbe.eval("SLICE(SLICE([arr_int_value], 2, 1), 1, 1)", from_=data_table) == dbe.eval(
            "ARRAY(23)", from_=data_table
        )

        assert dbe.eval("STARTSWITH(SLICE(ARRAY(1, 2, 3, 4, 5), 1, 4), SLICE(ARRAY(1, 2, 3), 1, 2))")
        assert not dbe.eval("STARTSWITH(SLICE([arr_int_value], 1, 4), ARRAY(1, 2, 3))", from_=data_table)
        assert dbe.eval("STARTSWITH(SLICE([arr_int_value], 1, 3), ARRAY(0, 23, 456))", from_=data_table)

    def test_startswith_number_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("STARTSWITH(ARRAY(1, 2, 3, 4, 5), ARRAY(1, 2, 3))", from_=data_table)
        assert dbe.eval("STARTSWITH(ARRAY(1, 2, 3, 4, 5), ARRAY(1, 2, 3, 4, 5))", from_=data_table)
        assert dbe.eval("STARTSWITH(ARRAY(1, NULL, 3, NULL, 5), ARRAY(1, NULL, 3, NULL))", from_=data_table)
        assert not dbe.eval("STARTSWITH(ARRAY(1, 2, 3, 4, 5), ARRAY(1, 2, 3, 4, 5, 6, 7))", from_=data_table)
        assert not dbe.eval("STARTSWITH(ARRAY(1, 2, 3, 4, 5), ARRAY(5, 4, 3, 2, 1))", from_=data_table)
        assert not dbe.eval("STARTSWITH(ARRAY(1, 2, 3, 4, 5), ARRAY(3, 4))", from_=data_table)
        assert not dbe.eval("STARTSWITH(ARRAY(1, 2, NULL, 4, 5), ARRAY(2, NULL))", from_=data_table)
        assert dbe.eval("STARTSWITH(ARRAY(1.1, NULL, 3.3, 4.4, 5.5), ARRAY(1.1, NULL, 3.3))", from_=data_table)
        assert dbe.eval('STARTSWITH(ARRAY("", "NULL", NULL, "12"), ARRAY("", "NULL"))', from_=data_table)

        assert dbe.eval("STARTSWITH([arr_int_value], [arr_int_value])", from_=data_table)
        assert dbe.eval("STARTSWITH([arr_float_value], [arr_float_value])", from_=data_table)

    @abc.abstractmethod
    def test_startswith_string_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        pass  # Defined in a dialect-specific way

    def test_array_remove(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        # with ARRAY func | remove literal not NULL
        assert dbe.eval("ARR_REMOVE(ARRAY(1, 2, NULL), 1)", from_=data_table) == dbe.eval("ARRAY(2, NULL)")
        assert dbe.eval('ARR_REMOVE(ARRAY("a", "b", NULL), "a")', from_=data_table) == dbe.eval('ARRAY("b", NULL)')
        if self.make_decimal_cast:
            assert dbe.eval("ARR_REMOVE(ARRAY(1.1, 2.2, NULL), 1.1)", from_=data_table) == dbe.eval(
                f'ARRAY(DB_CAST(2.2, "{self.make_decimal_cast}", 2, 1), NULL)'
            )
        else:
            assert dbe.eval("ARR_REMOVE(ARRAY(1.1, 2.2, NULL), 1.1)", from_=data_table) == dbe.eval("ARRAY(2.2, NULL)")

        # with ARRAY func | remove literal NULL
        assert dbe.eval("ARR_REMOVE(ARRAY(1, 2, NULL), NULL)", from_=data_table) == dbe.eval("ARRAY(1, 2)")
        assert dbe.eval('ARR_REMOVE(ARRAY("a", "b", NULL), NULL)', from_=data_table) == dbe.eval('ARRAY("a", "b")')
        if self.make_decimal_cast:
            assert dbe.eval("ARR_REMOVE(ARRAY(1.1, 2.2, NULL), NULL)", from_=data_table) == dbe.eval(
                f'ARRAY(DB_CAST(1.1, "{self.make_decimal_cast}", 2, 1), DB_CAST(2.2, "{self.make_decimal_cast}", 2, 1))'
            )
        else:
            assert dbe.eval("ARR_REMOVE(ARRAY(1.1, 2.2, NULL), NULL)", from_=data_table) == dbe.eval("ARRAY(1.1, 2.2)")

        # with array in DB | remove literal not NULL
        assert dbe.eval("ARR_REMOVE([arr_int_value], 0)", from_=data_table) == dbe.eval("ARRAY(23, 456, NULL)")
        assert dbe.eval('ARR_REMOVE([arr_str_value], "")', from_=data_table) == dbe.eval('ARRAY("cde", NULL)')
        if self.make_float_cast:
            assert dbe.eval(
                f'ARR_REMOVE([arr_float_value], DB_CAST(0.0, "{self.make_float_cast}"))', from_=data_table
            ) == dbe.eval("ARRAY(45, 0.123, NULL)")
        else:
            assert dbe.eval("ARR_REMOVE([arr_float_value], 0.0)", from_=data_table) == dbe.eval(
                "ARRAY(45, 0.123, NULL)"
            )

        # with array in DB | remove literal NULL
        assert dbe.eval("ARR_REMOVE([arr_int_value], NULL)", from_=data_table) == dbe.eval("ARRAY(0, 23, 456)")
        assert dbe.eval("ARR_REMOVE([arr_str_value], NULL)", from_=data_table) == dbe.eval('ARRAY("", "", "cde")')
        assert dbe.eval("ARR_REMOVE([arr_float_value], NULL)", from_=data_table) == dbe.eval("ARRAY(0, 45, 0.123)")

        # with array in DB | remove non const
        assert dbe.eval("ARR_REMOVE([arr_int_value], GET_ITEM([arr_int_value], 1))", from_=data_table) == dbe.eval(
            "ARRAY(23, 456, NULL)"
        )
        assert dbe.eval("ARR_REMOVE([arr_str_value], GET_ITEM([arr_str_value], 1))", from_=data_table) == dbe.eval(
            'ARRAY("cde", NULL)'
        )
        assert dbe.eval("ARR_REMOVE([arr_float_value], GET_ITEM([arr_float_value], 1))", from_=data_table) == dbe.eval(
            "ARRAY(45, 0.123, NULL)"
        )

lib/dl_formula_testing/dl_formula_testing/testcases/functions_array.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from typing import (
    Any,
    ClassVar,
)

import pytest
import pytz
import sqlalchemy.exc as sa_exc

from dl_formula.core.exc import TranslationError
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import (
    approx_datetime,
    dt_strip,
    now,
)


def int_value(value: Any) -> int:
    assert isinstance(value, int)
    return value


class DefaultDateTimeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    supports_addition_to_feb_29: ClassVar[bool] = True
    supports_dateadd_non_const_unit_num: ClassVar[bool] = False
    supports_deprecated_dateadd: ClassVar[bool] = False
    supports_deprecated_datepart_2: ClassVar[bool] = False
    supports_datepart_2_non_const: ClassVar[bool] = True
    supports_datetrunc_3: ClassVar[bool] = False
    supports_datetimetz: ClassVar[bool] = False

    @pytest.mark.parametrize("func_name", ("NOW", "GENERICNOW"))
    def test_now(self, dbe: DbEvaluator, func_name: str) -> None:
        assert dt_strip(dbe.eval(f"{func_name}()")) == approx_datetime(
            now().replace(tzinfo=pytz.UTC).astimezone(dbe.db.tzinfo).replace(tzinfo=None)
        )

    def test_today(self, dbe: DbEvaluator) -> None:
        assert (
            dbe.eval("TODAY()")
            == (now().replace(tzinfo=pytz.UTC).astimezone(dbe.db.tzinfo).replace(tzinfo=None)).date()
        )

    def test_datetime_dateadd_with_uneven_month_lengths(self, dbe: DbEvaluator) -> None:
        if not self.supports_addition_to_feb_29:
            with pytest.raises(sa_exc.DatabaseError):
                assert dbe.eval('DATEADD(#2020-02-29#, "year", 1)')

        else:
            # Leap year
            assert dbe.eval('DATEADD(#2020-02-29#, "year", 1)') == datetime.date(2021, 2, 28)
            assert dbe.eval('DATEADD(#2020-02-29#, "year", -1)') == datetime.date(2019, 2, 28)
            # Short and long months
            assert dbe.eval('DATEADD(#2021-01-30#, "month", 1)') == datetime.date(2021, 2, 28)
            assert dbe.eval('DATEADD(#2021-03-30#, "month", -1)') == datetime.date(2021, 2, 28)

    @pytest.mark.parametrize("lit", (pytest.param("##", id="generic"), pytest.param("#", id="regular")))
    def test_datetime_dateadd(self, dbe: DbEvaluator, lit: str) -> None:
        def _dt_lit(s: str) -> str:
            return f"{lit}{s}{lit}"

        # 1 arg
        assert dbe.eval("DATEADD(#2018-01-12#)") == datetime.date(2018, 1, 13)

        # 2 args (unit)
        assert dbe.eval('DATEADD(#2018-01-12#, "day")') == datetime.date(2018, 1, 13)
        assert dbe.eval('DATEADD(#2018-01-12#, "month")') == datetime.date(2018, 2, 12)
        assert dbe.eval('DATEADD(#2018-01-12#, "year")') == datetime.date(2019, 1, 12)

        # 2 args (number)
        assert dbe.eval("DATEADD(#2018-01-12#, 6)") == datetime.date(2018, 1, 18)

        # 3 args
        assert dbe.eval('DATEADD(#2018-01-12#, "day", 6)') == datetime.date(2018, 1, 18)
        assert dbe.eval('DATEADD(#2018-01-12#, "month", 6)') == datetime.date(2018, 7, 12)
        assert dbe.eval('DATEADD(#2018-01-12#, "year", 6)') == datetime.date(2024, 1, 12)
        assert dt_strip(dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "second", 6)')) == datetime.datetime(
            2018, 1, 12, 1, 2, 9
        )
        assert dt_strip(dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "minute", 6)')) == datetime.datetime(
            2018, 1, 12, 1, 8, 3
        )
        assert dt_strip(dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "hour", 6)')) == datetime.datetime(
            2018, 1, 12, 7, 2, 3
        )
        assert dt_strip(dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "day", 6)')) == datetime.datetime(
            2018, 1, 18, 1, 2, 3
        )
        assert dt_strip(dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "month", 6)')) == datetime.datetime(
            2018, 7, 12, 1, 2, 3
        )
        assert dt_strip(dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "year", 6)')) == datetime.datetime(
            2024, 1, 12, 1, 2, 3
        )
        if self.supports_dateadd_non_const_unit_num:
            # non-const number
            assert dbe.eval('DATEADD(#2018-01-12#, "day", INT("6"))') == datetime.date(2018, 1, 18)
            assert dbe.eval('DATEADD(#2018-01-12#, "month", INT("6"))') == datetime.date(2018, 7, 12)
            assert dbe.eval('DATEADD(#2018-01-12#, "year", INT("6"))') == datetime.date(2024, 1, 12)
            assert dt_strip(
                dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "second", INT("6"))')
            ) == datetime.datetime(2018, 1, 12, 1, 2, 9)
            assert dt_strip(
                dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "minute", INT("6"))')
            ) == datetime.datetime(2018, 1, 12, 1, 8, 3)
            assert dt_strip(
                dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "hour", INT("6"))')
            ) == datetime.datetime(2018, 1, 12, 7, 2, 3)
            assert dt_strip(
                dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "day", INT("6"))')
            ) == datetime.datetime(2018, 1, 18, 1, 2, 3)
            assert dt_strip(
                dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "month", INT("6"))')
            ) == datetime.datetime(2018, 7, 12, 1, 2, 3)
            assert dt_strip(
                dbe.eval(f'DATEADD({_dt_lit("2018-01-12 01:02:03")}, "year", INT("6"))')
            ) == datetime.datetime(2024, 1, 12, 1, 2, 3)

    @pytest.mark.parametrize("lit", (pytest.param("##", id="generic"), pytest.param("#", id="regular")))
    def test_datetime_dateadd_deprecated(self, dbe: DbEvaluator, lit: str) -> None:
        if not self.supports_deprecated_dateadd:
            pytest.skip()

        def _dt_lit(s: str) -> str:
            return f"{lit}{s}{lit}"

        assert dbe.eval('DATEADD("day", 6, #2018-01-12#)') == datetime.date(2018, 1, 18)
        assert dbe.eval('DATEADD("month", 6, #2018-01-12#)') == datetime.date(2018, 7, 12)
        assert dbe.eval('DATEADD("year", 6, #2018-01-12#)') == datetime.date(2024, 1, 12)
        assert dt_strip(dbe.eval(f'DATEADD("second", 6, {_dt_lit("2018-01-12 01:02:03")})')) == datetime.datetime(
            2018, 1, 12, 1, 2, 9
        )
        assert dt_strip(dbe.eval(f'DATEADD("minute", 6, {_dt_lit("2018-01-12 01:02:03")})')) == datetime.datetime(
            2018, 1, 12, 1, 8, 3
        )
        assert dt_strip(dbe.eval(f'DATEADD("hour", 6, {_dt_lit("2018-01-12 01:02:03")})')) == datetime.datetime(
            2018, 1, 12, 7, 2, 3
        )
        assert dt_strip(dbe.eval(f'DATEADD("day", 6, {_dt_lit("2018-01-12 01:02:03")})')) == datetime.datetime(
            2018, 1, 18, 1, 2, 3
        )
        assert dt_strip(dbe.eval(f'DATEADD("month", 6, {_dt_lit("2018-01-12 01:02:03")})')) == datetime.datetime(
            2018, 7, 12, 1, 2, 3
        )
        assert dt_strip(dbe.eval(f'DATEADD("year", 6, {_dt_lit("2018-01-12 01:02:03")})')) == datetime.datetime(
            2024, 1, 12, 1, 2, 3
        )

    def test_datepart_2_const(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        # const version
        # date
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, "second")')) == 0
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, "minute")')) == 0
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, "hour")')) == 0
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, "day")')) == 12
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, "month")')) == 1
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, "quarter")')) == 1
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, "year")')) == 2018
        assert int_value(dbe.eval('DATEPART(#1971-01-14#, "dayofweek")')) == 4
        assert int_value(dbe.eval('DATEPART(#1971-01-14#, "dow")')) == 4
        assert int_value(dbe.eval('DATEPART(#1971-01-14#, "week")')) == 2

        # datetime
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, "second")')) == 3
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, "minute")')) == 2
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, "hour")')) == 1
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, "day")')) == 12
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, "month")')) == 1
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, "quarter")')) == 1
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, "year")')) == 2018
        assert int_value(dbe.eval('DATEPART(#1971-01-14 01:02:03#, "dayofweek")')) == 4
        assert int_value(dbe.eval('DATEPART(#1971-01-14 01:02:03#, "dow")')) == 4
        assert int_value(dbe.eval('DATEPART(#1971-01-14 01:02:03#, "week")')) == 2

    def test_datepart_3(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval('DATEPART(#1971-01-14 01:02:03#, "dayofweek", "mon")') == 4
        assert dbe.eval('DATEPART(#1971-01-14 01:02:03#, "dow", "mon")') == 4
        assert dbe.eval('DATEPART(#1971-01-14 01:02:03#, "dayofweek", "wed")') == 2
        assert dbe.eval('DATEPART(#1971-01-14 01:02:03#, "dow", "wed")') == 2
        assert dbe.eval('DATEPART(#1971-01-14 01:02:03#, "dayofweek", "wednesday")') == 2

    def test_datepart_2_non_const(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        if not self.supports_datepart_2_non_const:
            pytest.skip()

        # date
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, __LIT__("second"))')) == 0
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, __LIT__("minute"))')) == 0
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, __LIT__("hour"))')) == 0
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, __LIT__("day"))')) == 12
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, __LIT__("month"))')) == 1
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, __LIT__("quarter"))')) == 1
        assert int_value(dbe.eval('DATEPART(#2018-01-12#, __LIT__("year"))')) == 2018
        assert int_value(dbe.eval('DATEPART(#1971-01-14#, __LIT__("dayofweek"))')) == 4
        assert int_value(dbe.eval('DATEPART(#1971-01-14#, __LIT__("dow"))')) == 4
        assert int_value(dbe.eval('DATEPART(#1971-01-14#, __LIT__("week"))')) == 2

        # datetime
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, __LIT__("second"))')) == 3
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, __LIT__("minute"))')) == 2
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, __LIT__("hour"))')) == 1
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, __LIT__("day"))')) == 12
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, __LIT__("month"))')) == 1
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, __LIT__("quarter"))')) == 1
        assert int_value(dbe.eval('DATEPART(#2018-01-12 01:02:03#, __LIT__("year"))')) == 2018
        assert int_value(dbe.eval('DATEPART(#1971-01-14 01:02:03#, __LIT__("dayofweek"))')) == 4
        assert int_value(dbe.eval('DATEPART(#1971-01-14 01:02:03#, __LIT__("dow"))')) == 4
        assert int_value(dbe.eval('DATEPART(#1971-01-14 01:02:03#, __LIT__("week"))')) == 2

    def test_datepart_2_deprecated(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        if not self.supports_deprecated_datepart_2:
            pytest.skip()

        # const
        assert dbe.eval('DATEPART("second", #2018-01-12#)') == 0
        assert dbe.eval('DATEPART("minute", #2018-01-12#)') == 0
        assert dbe.eval('DATEPART("hour", #2018-01-12#)') == 0
        assert dbe.eval('DATEPART("day", #2018-01-12#)') == 12
        assert dbe.eval('DATEPART("month", #2018-01-12#)') == 1
        assert dbe.eval('DATEPART("year", #2018-01-12#)') == 2018
        assert dbe.eval('DATEPART("dayofweek", #1971-01-14#)') == 4
        assert dbe.eval('DATEPART("dow", #1971-01-14#)') == 4
        assert dbe.eval('DATEPART("week", #1971-01-14#)') == 2
        assert dbe.eval('DATEPART("second", #2018-01-12 01:02:03#)') == 3
        assert dbe.eval('DATEPART("minute", #2018-01-12 01:02:03#)') == 2
        assert dbe.eval('DATEPART("hour", #2018-01-12 01:02:03#)') == 1
        assert dbe.eval('DATEPART("day", #2018-01-12 01:02:03#)') == 12
        assert dbe.eval('DATEPART("month", #2018-01-12 01:02:03#)') == 1
        assert dbe.eval('DATEPART("year", #2018-01-12 01:02:03#)') == 2018
        assert dbe.eval('DATEPART("dayofweek", #1971-01-14 01:02:03#)') == 4
        assert dbe.eval('DATEPART("dow", #1971-01-14 01:02:03#)') == 4
        assert dbe.eval('DATEPART("week", #1971-01-14 01:02:03#)') == 2

        # non-const
        assert dbe.eval('DATEPART(__LIT__("second"), #2018-01-12#)') == 0
        assert dbe.eval('DATEPART(__LIT__("minute"), #2018-01-12#)') == 0
        assert dbe.eval('DATEPART(__LIT__("hour"), #2018-01-12#)') == 0
        assert dbe.eval('DATEPART(__LIT__("day"), #2018-01-12#)') == 12
        assert dbe.eval('DATEPART(__LIT__("month"), #2018-01-12#)') == 1
        assert dbe.eval('DATEPART(__LIT__("year"), #2018-01-12#)') == 2018
        assert dbe.eval('DATEPART(__LIT__("dayofweek"), #1971-01-14#)') == 4
        assert dbe.eval('DATEPART(__LIT__("dow"), #1971-01-14#)') == 4
        assert dbe.eval('DATEPART(__LIT__("week"), #1971-01-14#)') == 2
        assert dbe.eval('DATEPART(__LIT__("second"), #2018-01-12 01:02:03#)') == 3
        assert dbe.eval('DATEPART(__LIT__("minute"), #2018-01-12 01:02:03#)') == 2
        assert dbe.eval('DATEPART(__LIT__("hour"), #2018-01-12 01:02:03#)') == 1
        assert dbe.eval('DATEPART(__LIT__("day"), #2018-01-12 01:02:03#)') == 12
        assert dbe.eval('DATEPART(__LIT__("month"), #2018-01-12 01:02:03#)') == 1
        assert dbe.eval('DATEPART(__LIT__("year"), #2018-01-12 01:02:03#)') == 2018
        assert dbe.eval('DATEPART(__LIT__("dayofweek"), #1971-01-14 01:02:03#)') == 4
        assert dbe.eval('DATEPART(__LIT__("dow"), #1971-01-14 01:02:03#)') == 4
        assert dbe.eval('DATEPART(__LIT__("week"), #1971-01-14 01:02:03#)') == 2

    def test_specific_date_part_functions(self, dbe: DbEvaluator) -> None:
        assert int_value(dbe.eval("SECOND(#2018-01-12#)")) == 0
        assert int_value(dbe.eval("MINUTE(#2018-01-12#)")) == 0
        assert int_value(dbe.eval("HOUR(#2018-01-12#)")) == 0
        assert int_value(dbe.eval("DAY(#2018-01-12#)")) == 12
        assert int_value(dbe.eval("MONTH(#2018-01-12#)")) == 1
        assert int_value(dbe.eval("QUARTER(#2018-01-12#)")) == 1
        assert int_value(dbe.eval("YEAR(#2018-01-12#)")) == 2018
        assert int_value(dbe.eval("DAYOFWEEK(#1971-01-14#)")) == 4
        assert int_value(dbe.eval("WEEK(#1971-01-14#)")) == 2
        assert int_value(dbe.eval("SECOND(#2018-01-12 01:02:03#)")) == 3
        assert int_value(dbe.eval("MINUTE(#2018-01-12 01:02:03#)")) == 2
        assert int_value(dbe.eval("HOUR(#2018-01-12 01:02:03#)")) == 1
        assert int_value(dbe.eval("DAY(#2018-01-12 01:02:03#)")) == 12
        assert int_value(dbe.eval("MONTH(#2018-01-12 01:02:03#)")) == 1
        assert int_value(dbe.eval("QUARTER(#2018-01-12 01:02:03#)")) == 1
        assert int_value(dbe.eval("YEAR(#2018-01-12 01:02:03#)")) == 2018
        assert int_value(dbe.eval("DAYOFWEEK(#1971-01-14 01:02:03#)")) == 4
        assert int_value(dbe.eval("WEEK(#1971-01-14 01:02:03#)")) == 2

        assert int_value(dbe.eval('DAYOFWEEK(#1971-01-14#, "wed")')) == 2
        assert int_value(dbe.eval('DAYOFWEEK(#1971-01-14#, "wednesday")')) == 2

    def test_datetrunc_2(self, dbe: DbEvaluator) -> None:
        assert dbe.eval('DATETRUNC(#2018-07-12#, "second")') == datetime.date(2018, 7, 12)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "minute")') == datetime.date(2018, 7, 12)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "hour")') == datetime.date(2018, 7, 12)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "day")') == datetime.date(2018, 7, 12)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "week")') == datetime.date(2018, 7, 9)
        # Sic! This is exactly how ISO8601 works for week
        assert dbe.eval('DATETRUNC(#2015-01-01#, "week")') == datetime.date(2014, 12, 29)
        assert dbe.eval('DATETRUNC(#2017-01-01#, "week")') == datetime.date(2016, 12, 26)
        assert dbe.eval('DATETRUNC(#2017-01-02#, "week")') == datetime.date(2017, 1, 2)
        assert dbe.eval('DATETRUNC(#2019-01-01#, "week")') == datetime.date(2018, 12, 31)
        assert dbe.eval('DATETRUNC(#2019-12-31#, "week")') == datetime.date(2019, 12, 30)
        assert dbe.eval('DATETRUNC(#2020-01-01#, "week")') == datetime.date(2019, 12, 30)
        assert dbe.eval('DATETRUNC(#2021-01-01#, "week")') == datetime.date(2020, 12, 28)
        assert dbe.eval('DATETRUNC(#2021-01-02#, "week")') == datetime.date(2020, 12, 28)
        assert dbe.eval('DATETRUNC(#2021-01-03#, "week")') == datetime.date(2020, 12, 28)
        assert dbe.eval('DATETRUNC(#2021-01-04#, "week")') == datetime.date(2021, 1, 4)
        assert dbe.eval('DATETRUNC(#2021-01-07#, "week")') == datetime.date(2021, 1, 4)
        assert dbe.eval('DATETRUNC(#2022-01-01#, "week")') == datetime.date(2021, 12, 27)
        assert dbe.eval('DATETRUNC(#2023-01-01#, "week")') == datetime.date(2022, 12, 26)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "month")') == datetime.date(2018, 7, 1)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "quarter")') == datetime.date(2018, 7, 1)
        assert dbe.eval('DATETRUNC(#2018-08-12#, "quarter")') == datetime.date(2018, 7, 1)
        assert dbe.eval('DATETRUNC(#2018-09-12#, "quarter")') == datetime.date(2018, 7, 1)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "year")') == datetime.date(2018, 1, 1)
        with pytest.raises(TranslationError):
            dbe.eval('DATETRUNC(#2018-07-12#, "unexpected")')
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "second")')) == datetime.datetime(
            2018, 7, 12, 11, 7, 13
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-01-02 03:04:05#, "second")')) == datetime.datetime(
            2018, 1, 2, 3, 4, 5
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "minute")')) == datetime.datetime(
            2018, 7, 12, 11, 7, 0
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-01-02 03:04:05#, "minute")')) == datetime.datetime(
            2018, 1, 2, 3, 4, 0
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "hour")')) == datetime.datetime(
            2018, 7, 12, 11, 0, 0
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-01-02 03:04:05#, "hour")')) == datetime.datetime(2018, 1, 2, 3, 0, 0)
        assert dt_strip(dbe.eval('DATETRUNC(#2022-10-01 15:21:40#, "hour")')) == datetime.datetime(
            2022, 10, 1, 15, 0, 0
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2022-09-19 15:21:40#, "hour")')) == datetime.datetime(
            2022, 9, 19, 15, 0, 0
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "day")')) == datetime.datetime(2018, 7, 12, 0, 0, 0)
        assert dt_strip(dbe.eval('DATETRUNC(#2018-01-02 03:04:05#, "day")')) == datetime.datetime(2018, 1, 2, 0, 0, 0)
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "week")')) == datetime.datetime(2018, 7, 9, 0, 0, 0)
        assert dt_strip(dbe.eval('DATETRUNC(#2018-01-02 03:04:05#, "week")')) == datetime.datetime(2018, 1, 1, 0, 0, 0)
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "month")')) == datetime.datetime(2018, 7, 1, 0, 0, 0)
        assert dt_strip(dbe.eval('DATETRUNC(#2018-01-02 03:04:05#, "month")')) == datetime.datetime(2018, 1, 1, 0, 0, 0)
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "quarter")')) == datetime.datetime(
            2018, 7, 1, 0, 0, 0
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "year")')) == datetime.datetime(2018, 1, 1, 0, 0, 0)
        with pytest.raises(TranslationError):
            dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "unexpected")')

    def test_datetrunc_3(self, dbe: DbEvaluator) -> None:
        if not self.supports_datetrunc_3:
            pytest.skip()

        assert dbe.eval('DATETRUNC(#2018-07-12#, "second", 5)') == datetime.date(2018, 7, 12)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "minute", 5)') == datetime.date(2018, 7, 12)
        # assert dbe.eval('DATETRUNC(#2018-07-12#, "hour", 5)') == datetime.date(2018, 7, 12)  # FIXME: https://github.com/ClickHouse/ClickHouse/issues/10124
        assert dbe.eval('DATETRUNC(#2018-07-12#, "day", 5)') == datetime.date(2018, 7, 8)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "week", 5)') == datetime.date(2018, 7, 2)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "month", 4)') == datetime.date(2018, 5, 1)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "quarter", 2)') == datetime.date(2018, 7, 1)
        assert dbe.eval('DATETRUNC(#2018-07-12#, "year", 5)') == datetime.date(2015, 1, 1)
        with pytest.raises(TranslationError):
            dbe.eval('DATETRUNC(#2018-07-12#, "unexpected", 5)')
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "second", 5)')) == datetime.datetime(
            2018, 7, 12, 11, 7, 10
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "minute", 5)')) == datetime.datetime(
            2018, 7, 12, 11, 5, 0
        )
        # assert dt_strip(  # FIXME: https://github.com/ClickHouse/ClickHouse/issues/10124
        #     dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "hour", 5)')) == datetime.datetime(2018, 7, 12, 10, 0, 0)
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "day", 5)')) == datetime.datetime(
            2018, 7, 8, 0, 0, 0
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "week", 5)')) == datetime.datetime(
            2018, 7, 2, 0, 0, 0
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "month", 4)')) == datetime.datetime(
            2018, 5, 1, 0, 0, 0
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "quarter", 2)')) == datetime.datetime(
            2018, 7, 1, 0, 0, 0
        )
        assert dt_strip(dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "year", 5)')) == datetime.datetime(
            2015, 1, 1, 0, 0, 0
        )
        with pytest.raises(TranslationError):
            dbe.eval('DATETRUNC(#2018-07-12 11:07:13#, "unexpected", 5)')

    @pytest.mark.parametrize(
        "value_expr",
        (
            "__LIT__('2019-01-01T00:01:02')",
            "'2019-01-01T00:01:02'",
            "#2019-01-01T00:01:02#",
        ),
    )
    def test_datetimetz_make(self, dbe: DbEvaluator, forced_literal_use: Any, value_expr: str) -> None:
        if not self.supports_datetimetz:
            pytest.skip()

        tz = "America/New_York"
        expr = f"DATETIMETZ({value_expr}, '{tz}')"
        resp = dbe.eval(expr)

        assert isinstance(resp, datetime.datetime), resp
        assert resp.tzinfo, resp

        # -05:00, so +5hrs at UTC
        expected = datetime.datetime(2019, 1, 1, 5, 1, 2)
        assert resp.astimezone(datetime.timezone.utc).replace(tzinfo=None) == expected

    @pytest.fixture(params=["const", "lit"])
    def dttz_expr(self, request: Any, dbe: DbEvaluator, forced_literal_use: Any) -> str:
        if not self.supports_datetimetz:
            pytest.skip()

        value_iso = "2019-01-01T00:01:02"
        tz = "America/New_York"
        expr = f"'{value_iso}'"
        if request.param == "lit":
            expr = f"__LIT__({expr})"
        expr = f"DATETIMETZ({expr}, '{tz}')"
        return expr

    def test_datetimetz_funcs(self, dbe: DbEvaluator, dttz_expr: str, forced_literal_use: Any) -> None:
        resp = dbe.eval(f"AVG({dttz_expr})")
        assert isinstance(resp, datetime.datetime) and resp.tzinfo
        resp = dbe.eval(f"MAX({dttz_expr})")
        assert isinstance(resp, datetime.datetime) and resp.tzinfo
        resp = dbe.eval(f"MIN({dttz_expr})")
        assert isinstance(resp, datetime.datetime) and resp.tzinfo
        resp = dbe.eval(f"DATEADD({dttz_expr}, 'month', 2)")
        assert isinstance(resp, datetime.datetime) and resp.tzinfo
        assert dbe.eval(f"DATEPART({dttz_expr}, 'dayofweek', 'tue')") == 1
        assert dbe.eval(f"DATEPART({dttz_expr}, __LIT__('dayofweek'), 'wed')") == 7
        assert dbe.eval(f"DAYOFWEEK({dttz_expr})") == 2
        assert dbe.eval(f"DAYOFWEEK({dttz_expr}, 'thu')") == 6
        resp = dbe.eval(f"DATETRUNC({dttz_expr}, 'hour')")
        assert isinstance(resp, datetime.datetime) and resp.tzinfo
        resp = dbe.eval(f"DATETRUNC({dttz_expr}, 'month')")
        assert isinstance(resp, datetime.datetime) and resp.tzinfo

        if self.supports_datetrunc_3:
            resp = dbe.eval(f"DATETRUNC({dttz_expr}, 'month', 3)")
            assert isinstance(resp, datetime.datetime) and resp.tzinfo

        resp = dbe.eval(f"DATETRUNC({dttz_expr}, 'quarter')")
        assert isinstance(resp, datetime.datetime) and resp.tzinfo
        assert dbe.eval(f"DATE({dttz_expr})") == datetime.date(2019, 1, 1)

        assert dbe.eval(f"SECOND({dttz_expr})") == 2
        assert dbe.eval(f"MINUTE({dttz_expr})") == 1
        assert dbe.eval(f"HOUR({dttz_expr})") == 0
        assert dbe.eval(f"DAY({dttz_expr})") == 1
        assert dbe.eval(f"WEEK({dttz_expr})") == 1
        assert dbe.eval(f"MONTH({dttz_expr})") == 1
        assert dbe.eval(f"QUARTER({dttz_expr})") == 1
        assert dbe.eval(f"YEAR({dttz_expr})") == 2019

        # TODO: UTCNOW()
        # TODO: UTCTODAY()
        resp = dbe.eval(f"GREATEST({dttz_expr}, {dttz_expr})")
        assert isinstance(resp, datetime.datetime) and resp.tzinfo
        resp = dbe.eval(f"GREATEST({dttz_expr}, {dttz_expr}, {dttz_expr}, {dttz_expr})")
        assert isinstance(resp, datetime.datetime) and resp.tzinfo
        resp = dbe.eval(f"LEAST({dttz_expr}, {dttz_expr})")
        assert isinstance(resp, datetime.datetime) and resp.tzinfo
        with pytest.raises(TranslationError):
            dbe.eval(f"GREATEST({dttz_expr}, #2018-07-12 11:07:13#)")
        with pytest.raises(TranslationError):
            dbe.eval(f"LEAST(#2018-07-12 11:07:13#, {dttz_expr})")
        with pytest.raises(TranslationError):
            dbe.eval(f"LEAST({dttz_expr}, {dttz_expr}, #2018-07-12 11:07:13#, {dttz_expr}, {dttz_expr})")
        resp = dbe.eval(f"CONCAT('at ', {dttz_expr}, ' there')")
        assert isinstance(resp, str)

lib/dl_formula_testing/dl_formula_testing/testcases/functions_datetime.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
)

import pytest

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import to_str


class DefaultLogicalFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    supports_nan_funcs: ClassVar[bool] = False
    supports_iif: ClassVar[bool] = False

    def test_isnull(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("ISNULL(NULL)")
        assert not dbe.eval("ISNULL(2)")

    def test_ifnull(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("IFNULL(1, 2)") == 1
        assert dbe.eval("IFNULL(NULL, 2)") == 2
        assert dbe.eval("IFNULL(NULL, 0)") == 0

    def test_zn(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("ZN(1)") == 1
        assert dbe.eval("ZN(NULL)") == 0

    def test_isnan(self, dbe: DbEvaluator) -> None:
        if not self.supports_nan_funcs:  # FIXME: PG and CLICKHOUSe
            pytest.skip()

        assert dbe.eval("ISNAN(_MAKE_NAN())")

    def test_ifnan(self, dbe: DbEvaluator) -> None:
        if not self.supports_nan_funcs:  # FIXME: PG and CLICKHOUSe
            pytest.skip()

        assert dbe.eval("IFNAN(_MAKE_NAN(), 2)") == 2

    def test_if(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("IF(TRUE, 1, 2)") == 1
        assert dbe.eval("IF(FALSE, 1, 2)") == 2

    def test_if_with_operators(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("IF(1 = 1, 1, 2)") == 1
        assert dbe.eval("IF(1 != 1, 1, 2)") == 2

    def test_iif(self, dbe: DbEvaluator) -> None:
        if not self.supports_iif:  # FIXME: PG and CLICKHOUSe
            pytest.skip()
        assert dbe.eval("IIF(TRUE, 1, 2)") == 1
        assert dbe.eval("IIF(FALSE, 1, 2)") == 2

    def test_case(self, dbe: DbEvaluator) -> None:
        assert to_str(dbe.eval('CASE(1, 1, "1st", 2, "2nd", "3rd")')) == "1st"

    def test_expressions_with_conditional_flags(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        """Check correct usage of context flags"""

        assert to_str(dbe.eval('CASE(__LIT__(1) = __LIT__(2),  TRUE, "1st", FALSE, "2nd", "3rd")')) == "2nd"
        assert to_str(dbe.eval('CASE(__LIT__(TRUE),  TRUE, "1st", FALSE, "2nd", "3rd")')) == "1st"

lib/dl_formula_testing/dl_formula_testing/testcases/functions_logical.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Any

import pytest

import dl_formula.core.exc as exc
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import to_str


class DefaultMarkupFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    def test_markup_functions_simple(self, dbe: DbEvaluator) -> None:
        # markup value is a string, but it shouldn't normally leak to the user.
        with pytest.raises(exc.TranslationError):
            dbe.eval("str(italic('text00'))")

        with pytest.raises(exc.TranslationError):
            dbe.eval("str(italic('text00') + 'text01')")

        with pytest.raises(exc.TranslationError):
            dbe.eval("__str(italic('text00') + 'text01')")

        assert to_str(dbe.eval("markup(italic('text00'), bold('text01'))")) == '(c (i "text00") (b "text01"))'
        assert to_str(dbe.eval("markup(italic('text00'), 'text01')")) == '(c (i "text00") "text01")'

        # TODO: Decide the desired behavior.
        assert (dbe.eval("'abc' + NULL + 'qwe'") is None) == (dbe.eval("italic(NULL) + '...'") is None)

        assert to_str(dbe.eval("italic('text00') + 'text01'")) == '(c (i "text00") "text01")'
        assert to_str(dbe.eval("italic('text0')")) == '(i "text0")'

        assert to_str(dbe.eval("italic('text1') + bold('text2')")) == '(c (i "text1") (b "text2"))'
        assert to_str(dbe.eval("""italic('text "3"')""")) == '(i "text ""3""")'
        assert to_str(dbe.eval("url('http://example.com', 'text4')")) == '(a "http://example.com" "text4")'
        assert to_str(dbe.eval("url('javascript:alert(5)', bold('text5'))")) == '(a "javascript:alert(5)" (b "text5"))'

    def test_markup_functions_lit(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        expr = """
        italic(__LIT__('text "6"') + ' ... ') +
        __LIT__('text "8"') + ', text "9"'
        """
        expected_best_case = "(c " '(i "text ""6"" ...") ' '"text ""8""" ", text ""9""")'
        expected_acceptable = "(c " '(c (i "text ""6"" ... ") "text ""8""") ' '", text ""9""")'
        assert to_str(dbe.eval(expr)) in (expected_best_case, expected_acceptable)

lib/dl_formula_testing/dl_formula_testing/testcases/functions_markup.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import math
from typing import ClassVar

import pytest

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import (
    dt_strip,
    to_str,
)


class DefaultMathFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    supports_atan_2_in_origin: ClassVar[bool] = True
    supports_float_div: ClassVar[bool] = True

    def test_trigonometric_functions(self, dbe: DbEvaluator):
        assert dbe.eval("ACOS(1)") == 0
        assert dbe.eval("ACOS(0)") == pytest.approx(math.pi / 2)

        assert dbe.eval("ASIN(0)") == 0
        assert dbe.eval("ASIN(-1)") == pytest.approx(-math.pi / 2)

        assert dbe.eval("ATAN(0)") == 0
        assert dbe.eval("ATAN(-1)") == pytest.approx(-math.pi / 4)

        if self.supports_atan_2_in_origin:
            assert dbe.eval("ATAN2(0, 0)") == 0
        assert dbe.eval("ATAN2(1, 1)") == pytest.approx(math.pi / 4)
        assert dbe.eval("ATAN2(-1, -1)") == pytest.approx(-3 * math.pi / 4)
        assert dbe.eval("ATAN2(1, -1)") == pytest.approx(3 * math.pi / 4)
        assert dbe.eval("ATAN2(-1, 1)") == pytest.approx(-math.pi / 4)
        assert dbe.eval("ATAN2(1, 0)") == pytest.approx(math.pi / 2)
        assert dbe.eval("ATAN2(-1, 0)") == pytest.approx(-math.pi / 2)

        assert dbe.eval("COS(0)") == pytest.approx(1)
        assert dbe.eval("COS(PI()/2)") == pytest.approx(0)

        assert dbe.eval("COT(PI()/4)") == pytest.approx(1)
        assert dbe.eval("COT(PI()/2)") == pytest.approx(0)

        assert dbe.eval("DEGREES(0)") == 0
        assert dbe.eval("DEGREES(PI())") == pytest.approx(180)

        # MySQL uses double precision for PI internally, but displays only 7 decimal places
        assert dbe.eval("PI()") == pytest.approx(math.pi, abs=1e-06)

        assert dbe.eval("RADIANS(0)") == 0
        assert dbe.eval("RADIANS(180)") == pytest.approx(math.pi)

        assert dbe.eval("SIN(0)") == pytest.approx(0)
        assert dbe.eval("SIN(PI()/2)") == pytest.approx(1)

        assert dbe.eval("TAN(0)") == 0
        assert dbe.eval("TAN(PI()/4)") == pytest.approx(1)

    def test_exponential_functions(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("EXP(0)") == pytest.approx(1)
        assert dbe.eval("EXP(1)") == pytest.approx(math.e)

        assert dbe.eval("LN(1)") == pytest.approx(0)
        assert dbe.eval("LN(EXP(2))") == pytest.approx(2)

        assert dbe.eval("LOG10(1)") == pytest.approx(0)
        assert dbe.eval("LOG10(1000)") == pytest.approx(3)

        assert dbe.eval("LOG(1, 2.6)") == pytest.approx(0)
        assert dbe.eval("LOG(1024, 2)") == pytest.approx(10)

        assert dbe.eval("POWER(2.3, 4.5)") == pytest.approx(2.3**4.5)

        assert dbe.eval("SQRT(9)") == 3

        assert dbe.eval("SQUARE(9)") == 81

    def test_rounding_functions(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("CEILING(2.7)") == 3
        assert dbe.eval("CEILING(-2.7)") == -2

        assert dbe.eval("FLOOR(2.7)") == 2
        assert dbe.eval("FLOOR(-2.7)") == -3

        assert dbe.eval("ROUND(2.3)") == 2
        assert float(dbe.eval("ROUND(2.1234, 2)")) == 2.12

    def test_least(self, dbe: DbEvaluator) -> None:
        # single arg
        assert dbe.eval("LEAST(3.4)") == pytest.approx(3.4)
        # multiple args
        assert dbe.eval("LEAST(3.4, 2.6)") == pytest.approx(2.6)
        assert to_str(dbe.eval('LEAST("3.4", "2.6")')) == "2.6"
        assert dbe.eval("LEAST(#2019-01-02#, #2019-01-17#)") == datetime.date(2019, 1, 2)
        assert dt_strip(dbe.eval("LEAST(#2019-01-02 04:03:02#, #2019-01-17 03:02:01#)")) == datetime.datetime(
            2019, 1, 2, 4, 3, 2
        )
        assert dbe.eval("LEAST(TRUE, FALSE)") is False
        assert dbe.eval("LEAST(34, 5, 7, 3, 99, 1, 2, 2, 56)") == 1
        assert dbe.eval("LEAST(5.6, 1.2, 7.8, 3.4)") == pytest.approx(1.2)
        assert dbe.eval("LEAST(#2019-01-02#, #2019-01-17#, #2019-01-10#)") == datetime.date(2019, 1, 2)

    def test_greatest(self, dbe: DbEvaluator) -> None:
        # single arg
        assert dbe.eval("GREATEST(3.4)") == pytest.approx(3.4)
        # multiple args
        assert dbe.eval("GREATEST(3.4, 2.6)") == pytest.approx(3.4)
        assert to_str(dbe.eval('GREATEST("3.4", "2.6")')) == "3.4"
        assert dbe.eval("GREATEST(#2019-01-02#, #2019-01-17#)") == datetime.date(2019, 1, 17)
        assert dt_strip(dbe.eval("GREATEST(#2019-01-02 04:03:02#, #2019-01-17 03:02:01#)")) == datetime.datetime(
            2019, 1, 17, 3, 2, 1
        )
        assert dbe.eval("GREATEST(TRUE, FALSE)") is True
        assert dbe.eval("GREATEST(34, 5, 7, 3, 99, 1, 2, 2, 56)") == 99
        assert dbe.eval("GREATEST(5.6, 1.2, 7.8, 3.4)") == pytest.approx(7.8)
        assert dbe.eval("GREATEST(#2019-01-02#, #2019-01-17#, #2019-01-10#)") == datetime.date(2019, 1, 17)

    def test_arithmetic_functions(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("DIV(4, 2)") == 2
        assert dbe.eval("DIV(5, 3)") == 1
        assert dbe.eval("DIV_SAFE(5, 2)") == 2
        assert dbe.eval("DIV_SAFE(5, 0)") is None
        assert dbe.eval("DIV_SAFE(5, 0, 42)") == 42
        if self.supports_float_div:
            assert dbe.eval("DIV(5.0, 2)") == 2
            assert dbe.eval("DIV(5, 2.0)") == 2
            assert dbe.eval("DIV_SAFE(5.0, 2)") == 2
            assert dbe.eval("DIV_SAFE(5.0, 0)") is None
            assert dbe.eval("DIV_SAFE(5.0, 0, 42)") == 42
            assert dbe.eval("FDIV_SAFE(5.0, 2)") == 2.5
            assert dbe.eval("FDIV_SAFE(1, 0.0)") is None
            assert dbe.eval("FDIV_SAFE(1, 0.0, 42)") == 42

    def test_abs_sign(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("ABS(-3)") == 3

        assert dbe.eval("SIGN(-7)") == -1
        assert dbe.eval("SIGN(0)") == 0
        assert dbe.eval("SIGN(7)") == 1

lib/dl_formula_testing/dl_formula_testing/testcases/functions_math.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
)

import pytest
import sqlalchemy as sa

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import (
    to_str,
    to_unicode,
)


CONTAINS_TESTS = [
    ('"a b c % d e"', '""', True),
    ('"a b c % d e"', '"c % d"', True),
    ('"a b c % d e"', '"c \\\\% d"', False),
    ('"a b c % d e"', '"b \\\\% d"', False),
    ('"a b c %% d e"', '""', True),
    ('"a b c %% d e"', '"c % d"', False),
    ('"a b c %% d e"', '"c %% d"', True),
    ('"a b c %% d e"', '"c \\\\% d"', False),
    ('"a b c %% d e"', '"b \\\\% d"', False),
    ('"a b c \\\\% d e"', '""', True),
    ('"a b c \\\\% d e"', '"c % d"', False),
    ('"a b c \\\\% d e"', '"c \\\\% d"', True),
    ('"a b c \\\\% d e"', '"b \\\\% d"', False),
    ('"a b c \\\\\\\\% d e"', '""', True),
    ('"a b c \\\\\\\\% d e"', '"c % d"', False),
    ('"a b c \\\\\\\\% d e"', '"c \\\\% d"', False),
    ('"a b c \\\\\\\\% d e"', '"c \\\\\\\\% d"', True),
    ('"a b c \\\\\\\\% d e"', '"b \\\\% d"', False),
]


class DefaultStringFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    datetime_str_separator: ClassVar[str] = " "
    datetime_str_ending: ClassVar[str] = ""
    supports_trimming_funcs: ClassVar[bool] = True
    supports_regex_extract: ClassVar[bool] = True
    supports_regex_extract_nth: ClassVar[bool] = True
    supports_regex_replace: ClassVar[bool] = True
    supports_regex_match: ClassVar[bool] = True
    supports_split_3: ClassVar[bool] = True
    supports_non_const_percent_escape: ClassVar[bool] = True

    def test_char_conversions(self, dbe: DbEvaluator) -> None:
        assert dbe.eval('ASCII("W")') == 87
        assert to_str(dbe.eval("CHAR(87)")) == "W"

    def test_concat(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert to_str(dbe.eval('CONCAT("abc", "def", "qwe")')) == "abcdefqwe"
        sep = self.datetime_str_separator
        end = self.datetime_str_ending
        exp_value = f"1abcFalse2019-01-012019-01-01{sep}01:02:03{end}"
        assert to_str(dbe.eval('CONCAT(1, "abc", FALSE, #2019-01-01#, #2019-01-01 01:02:03#)')) == exp_value
        # const str
        assert to_str(dbe.eval('CONCAT("abc", "def")')) == "abcdef"

        # Single argument case
        # const
        assert to_str(dbe.eval('CONCAT("abc")')) == "abc"
        # non-const (different implementation)
        assert to_str(dbe.eval("CONCAT([str_value])", from_=data_table, order_by=["[str_value]"])) == "q"
        assert to_str(dbe.eval("CONCAT([int_value])", from_=data_table, order_by=["[int_value]"])) == "0"

    def test_find(self, dbe: DbEvaluator) -> None:
        assert dbe.eval('FIND("Lorem ipsum dolor sit amet", "psu")') == 8
        assert dbe.eval('FIND("Lorem ipsum dolor sit amet", "abc")') == 0
        assert dbe.eval('FIND("Карл у Клары украл кораллы", "рал")') == 16
        assert dbe.eval('FIND("Lorem ipsum dolor sit amet", "or", 7)') == 16
        assert dbe.eval('FIND("Карл у Клары украл кораллы", "рал", 18)') == 22

    def test_left_right(self, dbe: DbEvaluator) -> None:
        assert to_str(dbe.eval('LEFT("Lorem ipsum", 4)')) == "Lore"
        assert to_str(dbe.eval('LEFT("Карл у Клары", 4)')) == "Карл"

        assert to_str(dbe.eval('RIGHT("Lorem ipsum", 4)')) == "psum"
        assert to_str(dbe.eval('RIGHT("Карл у Клары", 4)')) == "лары"

    def test_len(self, dbe: DbEvaluator) -> None:
        assert dbe.eval('LEN("Lorem ipsum")') == 11
        assert dbe.eval('LEN("Карл у Клары")') == 12

    def test_lower_upper(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert to_str(dbe.eval('LOWER("Lorem ipsum")')) == "lorem ipsum"
        assert to_str(dbe.eval('LOWER(__LIT__("Lorem ipsum"))')) == "lorem ipsum"
        assert to_unicode(dbe.eval('LOWER("Карл у Клары")')) == "карл у клары"
        assert to_unicode(dbe.eval('LOWER(__LIT__("Карл у Клары"))')) == "карл у клары"

        assert to_str(dbe.eval('UPPER("Lorem ipsum")')) == "LOREM IPSUM"
        assert to_str(dbe.eval('UPPER(__LIT__("Lorem ipsum"))')) == "LOREM IPSUM"
        assert to_unicode(dbe.eval('UPPER("Карл у Клары")')) == "КАРЛ У КЛАРЫ"
        assert to_unicode(dbe.eval('UPPER(__LIT__("Карл у Клары"))')) == "КАРЛ У КЛАРЫ"

    def test_trimming(self, dbe: DbEvaluator) -> None:
        if not self.supports_trimming_funcs:
            pytest.skip()

        assert to_str(dbe.eval('LTRIM("  Lorem ipsum   ")')) == "Lorem ipsum   "
        assert to_str(dbe.eval('LTRIM(" \\n Lorem ipsum   ")')) == "\n Lorem ipsum   "

        assert to_str(dbe.eval('RTRIM("  Lorem ipsum   ")')) == "  Lorem ipsum"
        assert to_str(dbe.eval('RTRIM("  Lorem ipsum  \\n ")')) == "  Lorem ipsum  \n"

        assert to_str(dbe.eval('TRIM("  Lorem ipsum   ")')) == "Lorem ipsum"
        assert to_str(dbe.eval('TRIM(" \\n Lorem ipsum   ")')) == "\n Lorem ipsum"

    def test_substr(self, dbe: DbEvaluator) -> None:
        assert to_str(dbe.eval('SUBSTR("Lorem ipsum", 3)')) == "rem ipsum"
        assert to_str(dbe.eval('SUBSTR("Карл у Клары", 3)')) == "рл у Клары"
        assert to_str(dbe.eval('SUBSTR("Lorem ipsum", 3, 6)')) == "rem ip"
        assert to_str(dbe.eval('SUBSTR("Карл у Клары", 3, 6)')) == "рл у К"

    def test_regexp_extract(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_regex_extract:
            pytest.skip()

        assert to_str(dbe.eval('REGEXP_EXTRACT("Lorem ipsum dolor sit amet", "or..")')) == "orem"
        assert to_str(dbe.eval('REGEXP_EXTRACT("Lorem ipsum dolor sit amet", "abc")')) in (None, "")
        assert to_str(dbe.eval('REGEXP_EXTRACT("Карл у Клары украл кораллы", ".лары")')) == "Клары"
        assert to_str(dbe.eval('REGEXP_EXTRACT([str_null_value], "or..")', from_=data_table)) in (None, "")

    def test_regexp_extract_nth(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_regex_extract_nth:
            pytest.skip()

        assert to_str(dbe.eval('REGEXP_EXTRACT_NTH("Lorem ipsum dolor sit amet", "or..", 2)')) == "or s"
        assert to_str(dbe.eval('REGEXP_EXTRACT_NTH("Карл у Клары украл кораллы", ".рал", 2)')) == "орал"
        assert to_str(dbe.eval('REGEXP_EXTRACT_NTH([str_null_value], "or..", 2)', from_=data_table)) in (None, "")

    def test_regexp_extract_replace(self, dbe: DbEvaluator) -> None:
        if not self.supports_regex_replace:
            pytest.skip()

        assert to_str(dbe.eval('REGEXP_REPLACE("Lorem ipsum dolor sit amet", "or..", "____")')) == (
            "L____ ipsum dol____it amet"
        )
        assert to_str(dbe.eval('REGEXP_REPLACE("Карл у Клары украл кораллы", "рал.", "роп")')) == (
            "Карл у Клары укропкоропы"
        )

    def test_regexp_match(self, dbe: DbEvaluator) -> None:
        if not self.supports_regex_match:
            pytest.skip()

        assert dbe.eval('REGEXP_MATCH("Lorem ipsum", ".*ips.*")')
        assert not dbe.eval('REGEXP_MATCH("Lorem ipsum", ".*ipss.*")')
        assert dbe.eval('REGEXP_MATCH("Карл у Клары", ".*арл.*")')
        assert not dbe.eval('REGEXP_MATCH("Карл у Клары", ".*марл.*")')

    def test_replace(self, dbe: DbEvaluator) -> None:
        assert to_str(dbe.eval('REPLACE("Lorem ipsum dolor sit amet", "or", "__")')) == "L__em ipsum dol__ sit amet"
        assert to_str(dbe.eval('REPLACE("Карл у Клары украл кораллы", "ра", "ар")')) == "Карл у Клары укарл коарллы"

    def test_space(self, dbe: DbEvaluator) -> None:
        assert (
            to_str(dbe.eval("SPACE(4)")) == "    "
        )  # SPACE(<int literal>) is evaluated in Python rather than in the DB,
        assert to_str(dbe.eval("SPACE(2 + 2)")) == "    "  # while this version actually uses the DB function(s)

    def test_split_3(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_split_3:
            pytest.skip()

        assert to_str(dbe.eval('SPLIT("Lorem ipsum dolor sit amet", " ", 3)')) == "dolor"
        assert to_str(dbe.eval('SPLIT([str_null_value], " ", 3)', from_=data_table)) in (None, "")

    def test_split_1_2(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_arrays:
            pytest.skip()

        assert dbe.eval('ARR_STR(SPLIT(",One,two,three,"), "+")', from_=data_table) == "+One+two+three+"
        assert dbe.eval('ARR_STR(SPLIT(";One;two,three,", ";"), "+")', from_=data_table) == "+One+two,three,"
        assert dbe.eval('ARR_STR(SPLIT(""))', from_=data_table) == ""
        assert dbe.eval('ARR_STR(SPLIT("", "@"))', from_=data_table) == ""

    def test_contains_simple(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval('CONTAINS("Lorem ipsum", "ips")')
        assert not dbe.eval('CONTAINS("Lorem ipsum", "abc")')
        assert dbe.eval('CONTAINS(__LIT__("Lorem ipsum"), __LIT__("ips"))')
        assert not dbe.eval('CONTAINS(__LIT__("Lorem ipsum"), __LIT__("abc"))')
        assert dbe.eval('CONTAINS("Lorem %ipsum", "em %ip")')
        assert dbe.eval('CONTAINS(__LIT__("Lorem %ipsum"), __LIT__("em %ip"))')
        assert dbe.eval('CONTAINS("Карл у Клары украл кораллы", "Клары")')
        assert dbe.eval('CONTAINS(__LIT__("Карл у Клары украл кораллы"), __LIT__("Клары"))')
        # Non-string
        assert dbe.eval('CONTAINS(123456, "234")')
        assert not dbe.eval('CONTAINS(123456, "432")')
        assert dbe.eval('CONTAINS(TRUE, "ru")')
        assert dbe.eval('CONTAINS(#2019-03-04#, "019")')
        assert dbe.eval('CONTAINS(#2019-03-04T12:34:56#, "019")')

    def test_notcontains_simple(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert not dbe.eval('NOTCONTAINS("Lorem ipsum", "ips")')
        assert dbe.eval('NOTCONTAINS("Lorem ipsum", "abc")')
        assert not dbe.eval('NOTCONTAINS(__LIT__("Lorem ipsum"), __LIT__("ips"))')
        assert dbe.eval('NOTCONTAINS(__LIT__("Lorem ipsum"), __LIT__("abc"))')
        assert not dbe.eval('NOTCONTAINS("Lorem %ipsum", "em %ip")')
        assert not dbe.eval('NOTCONTAINS(__LIT__("Lorem %ipsum"), __LIT__("em %ip"))')
        assert not dbe.eval('NOTCONTAINS("Карл у Клары украл кораллы", "Клары")')
        assert not dbe.eval('NOTCONTAINS(__LIT__("Карл у Клары украл кораллы"), __LIT__("Клары"))')
        # Non-string
        assert not dbe.eval('NOTCONTAINS(123456, "234")')
        assert dbe.eval('NOTCONTAINS(123456, "432")')
        assert not dbe.eval('NOTCONTAINS(TRUE, "ru")')
        assert not dbe.eval('NOTCONTAINS(#2019-03-04#, "019")')
        assert not dbe.eval('NOTCONTAINS(#2019-03-04T12:34:56#, "019")')

    @pytest.mark.parametrize("value_fl,pattern_fl,expected", CONTAINS_TESTS)
    def test_contains_extended(
        self,
        dbe: DbEvaluator,
        data_table: sa.Table,
        forced_literal_use: Any,
        value_fl: str,
        pattern_fl: str,
        expected: bool,
    ) -> None:
        if self.empty_str_is_null and pattern_fl == '""':
            # hopeless?
            return

        # const:
        statement = "CONTAINS(__LIT__({}), {})".format(value_fl, pattern_fl)
        assert dbe.eval(statement) is expected, (statement, expected)

        # var:
        statement = "CONTAINS(__LIT__({}), __LIT__({}))".format(value_fl, pattern_fl)
        assert dbe.eval(statement) is expected, (statement, expected)

    @pytest.mark.parametrize("value_fl,pattern_fl,expected", CONTAINS_TESTS)
    def test_notcontains_extended(
        self,
        dbe: DbEvaluator,
        data_table: sa.Table,
        forced_literal_use: Any,
        value_fl: str,
        pattern_fl: str,
        expected: bool,
    ) -> None:
        if self.empty_str_is_null and pattern_fl == '""':
            # hopeless?
            return

        # const:
        statement = "NOTCONTAINS(__LIT__({}), {})".format(value_fl, pattern_fl)
        a = dbe.eval(statement)
        assert dbe.eval(statement) is not expected, (statement, expected)

        # var:
        statement = "NOTCONTAINS(__LIT__({}), __LIT__({}))".format(value_fl, pattern_fl)
        assert dbe.eval(statement) is not expected, (statement, expected)

    def test_icontains_simple(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval('ICONTAINS("Lorem ipsum", "IPS")')
        assert not dbe.eval('ICONTAINS("Lorem ipsum", "ABC")')
        assert dbe.eval('ICONTAINS(__LIT__("Lorem ipsum"), __LIT__("IPS"))')
        assert not dbe.eval('ICONTAINS(__LIT__("Lorem ipsum"), __LIT__("ABC"))')
        assert dbe.eval('ICONTAINS("Lorem %ipsum", "EM %IP")')
        assert dbe.eval('ICONTAINS(__LIT__("Lorem %ipsum"), __LIT__("EM %IP"))')
        assert dbe.eval('ICONTAINS("Карл у Клары украл кораллы", "КЛАРЫ")')
        assert dbe.eval('ICONTAINS(__LIT__("Карл у Клары украл кораллы"), __LIT__("КЛАРЫ"))')
        # Quoting note: non-raw python string (backslash-escaped),
        # formulas string contents (backslash-escaped),
        # `like` pattern (backslash-escaped at a later point).
        assert dbe.eval('ICONTAINS("Lorem\\\\n %ipsum", "EM\\\\n %IP")')
        assert not dbe.eval('ICONTAINS("Lorem .ipsum", "EM %IP")')
        assert not dbe.eval('ICONTAINS("Lorem .ipsum", "EM \\%IP")')
        assert not dbe.eval('ICONTAINS("Lorem .ipsum", "EM \\\\%IP")')

        if self.supports_non_const_percent_escape:  # No proper escaping of a non-const value for MSSQL.
            assert dbe.eval('ICONTAINS(__LIT__("Lorem\\\\n\\\\% %ipsum"), __LIT__("EM\\\\n\\\\% %IP"))')
            assert not dbe.eval('ICONTAINS(__LIT__("Lorem .ipsum"), __LIT__("EM %IP"))')
            assert not dbe.eval('ICONTAINS(__LIT__("Lorem .ipsum"), __LIT__("EM \\%IP"))')
            assert not dbe.eval('ICONTAINS(__LIT__("Lorem .ipsum"), __LIT__("EM \\\\%IP"))')

        # Non-string
        assert dbe.eval('ICONTAINS(123456, "234")')
        assert not dbe.eval('ICONTAINS(123456, "432")')
        assert dbe.eval('ICONTAINS(123456.34, "234")')
        assert dbe.eval('ICONTAINS(TRUE, "ru")')
        assert dbe.eval('ICONTAINS(#2019-03-04#, "019")')
        assert dbe.eval('ICONTAINS(#2019-03-04T12:34:56#, "019")')

    def test_startswith_simple(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval('STARTSWITH("Lorem ipsum", "Lore")')
        assert not dbe.eval('STARTSWITH("Lorem ipsum", "abc")')
        assert dbe.eval('STARTSWITH(__LIT__("Lorem ipsum"), __LIT__("Lore"))')
        assert not dbe.eval('STARTSWITH(__LIT__("Lorem ipsum"), __LIT__("abc"))')
        assert dbe.eval('STARTSWITH("Lorem% ipsum", "Lorem% ")')
        assert dbe.eval('STARTSWITH(__LIT__("Lorem% ipsum"), __LIT__("Lorem% "))')
        assert dbe.eval('STARTSWITH("Карл у Клары украл кораллы", "Карл")')
        assert dbe.eval('STARTSWITH(__LIT__("Карл у Клары украл кораллы"), __LIT__("Карл"))')
        # Non-string
        assert dbe.eval('STARTSWITH(123456, "123")')
        assert not dbe.eval('STARTSWITH(123456, "456")')
        assert dbe.eval('STARTSWITH(123456.34, "123")')
        assert dbe.eval('STARTSWITH(TRUE, "T")')
        assert dbe.eval('STARTSWITH(#2019-03-04#, "2019")')
        assert dbe.eval('STARTSWITH(#2019-03-04T12:34:56#, "2019")')

    def test_istartswith_simple(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval('ISTARTSWITH("Lorem ipsum", "LORE")')
        assert not dbe.eval('ISTARTSWITH("Lorem ipsum", "ABC")')
        assert dbe.eval('ISTARTSWITH(__LIT__("Lorem ipsum"), __LIT__("LORE"))')
        assert not dbe.eval('ISTARTSWITH(__LIT__("Lorem ipsum"), __LIT__("ABC"))')
        assert dbe.eval('ISTARTSWITH("Lorem% ipsum", "LOREM% ")')
        assert dbe.eval('ISTARTSWITH(__LIT__("Lorem% ipsum"), __LIT__("LOREM% "))')
        assert dbe.eval('ISTARTSWITH("Карл у Клары украл кораллы", "КАРЛ")')
        assert dbe.eval('ISTARTSWITH(__LIT__("Карл у Клары украл кораллы"), __LIT__("КАРЛ"))')
        # Non-string
        assert dbe.eval('ISTARTSWITH(123456, "123")')
        assert not dbe.eval('ISTARTSWITH(123456, "456")')
        assert dbe.eval('ISTARTSWITH(123456.34, "123")')
        assert dbe.eval('ISTARTSWITH(TRUE, "T")')
        assert dbe.eval('ISTARTSWITH(#2019-03-04#, "2019")')
        assert dbe.eval('ISTARTSWITH(#2019-03-04T12:34:56#, "2019")')

    def test_endswith_simple(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval('ENDSWITH("Lorem ipsum", "sum")')
        assert not dbe.eval('ENDSWITH("Lorem ipsum", "abc")')
        assert dbe.eval('ENDSWITH(__LIT__("Lorem ipsum"), __LIT__("sum"))')
        assert not dbe.eval('ENDSWITH(__LIT__("Lorem ipsum"), __LIT__("abc"))')
        assert dbe.eval('ENDSWITH("Lorem %ipsum", " %ipsum")')
        assert dbe.eval('ENDSWITH(__LIT__("Lorem %ipsum"), __LIT__(" %ipsum"))')
        assert dbe.eval('ENDSWITH("Карл у Клары украл кораллы", "аллы")')
        assert dbe.eval('ENDSWITH(__LIT__("Карл у Клары украл кораллы"), __LIT__("аллы"))')
        # Non-string
        assert dbe.eval('ENDSWITH(123456, "456")')
        assert not dbe.eval('ENDSWITH(123456, "123")')
        assert dbe.eval('ENDSWITH(123.456, "456")')
        assert dbe.eval('ENDSWITH(TRUE, "ue")')
        assert dbe.eval('ENDSWITH(#2019-03-04#, "04")')
        dt_end = self.datetime_str_ending
        assert dbe.eval(f'ENDSWITH(#2019-03-04T12:34:56#, "56{dt_end}")')

    def test_iendswith_simple(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval('IENDSWITH("Lorem ipsum", "SUM")')
        assert not dbe.eval('IENDSWITH("Lorem ipsum", "ABC")')
        assert dbe.eval('IENDSWITH(__LIT__("Lorem ipsum"), __LIT__("SUM"))')
        assert not dbe.eval('IENDSWITH(__LIT__("Lorem ipsum"), __LIT__("ABC"))')
        assert dbe.eval('IENDSWITH("Lorem %ipsum", " %IPSUM")')
        assert dbe.eval('IENDSWITH(__LIT__("Lorem %ipsum"), __LIT__(" %IPSUM"))')
        assert dbe.eval('IENDSWITH("Карл у Клары украл кораллы", "АЛЛЫ")')
        assert dbe.eval('IENDSWITH(__LIT__("Карл у Клары украл кораллы"), __LIT__("АЛЛЫ"))')
        # Non-string
        assert dbe.eval('IENDSWITH(123456, "456")')
        assert not dbe.eval('IENDSWITH(123456, "123")')
        assert dbe.eval('IENDSWITH(123.456, "456")')
        assert dbe.eval('IENDSWITH(TRUE, "ue")')
        assert dbe.eval('IENDSWITH(#2019-03-04#, "04")')
        dt_end = self.datetime_str_ending
        assert dbe.eval(f'IENDSWITH(#2019-03-04T12:34:56#, "56{dt_end}")')

lib/dl_formula_testing/dl_formula_testing/testcases/functions_string.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import re
from typing import ClassVar

import pytest
import pytz
import sqlalchemy as sa

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import (
    as_tz,
    dt_strip,
    to_datetime,
    to_str,
    utc_ts,
)


class DefaultStrTypeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    zero_float_to_str_value: ClassVar[str] = "0"
    bool_of_null_is_false: bool = False
    skip_custom_tz: ClassVar[bool] = False
    skip_max_ints: ClassVar[bool] = False

    def test_str_numbers(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("STR(0)") == "0"
        assert to_str(dbe.eval("STR(0.0)")) == self.zero_float_to_str_value
        assert to_str(dbe.eval("STR(48.2)")) == "48.2"
        assert to_str(dbe.eval("STR(-48.2)")) == "-48.2"
        assert to_str(dbe.eval("STR(4887.2752789)")) == "4887.2752789"
        assert to_str(dbe.eval("STR(-4887.2752789)")) == "-4887.2752789"
        assert to_str(dbe.eval("STR(8010)")) == "8010"
        assert to_str(dbe.eval("STR(-8010)")) == "-8010"

    def test_str_from_bool(self, dbe: DbEvaluator) -> None:
        assert to_str(dbe.eval("STR(TRUE)")) == "True"
        assert to_str(dbe.eval("STR(FALSE)")) == "False"

    def test_str_from_str(self, dbe: DbEvaluator) -> None:
        assert to_str(dbe.eval('STR("qwe")')) == "qwe"

    def test_str_from_date(self, dbe: DbEvaluator) -> None:
        assert to_str(dbe.eval("STR(#2019-01-02#)")) == "2019-01-02"

    def test_str_from_datetime(self, dbe: DbEvaluator) -> None:
        assert re.match(r"2019-01-02 03:04:05(\.\d+)?", to_str(dbe.eval("STR(#2019-01-02 03:04:05#)")))

    def test_str_from_datetime_custom_tz(self, dbe: DbEvaluator) -> None:
        if self.skip_custom_tz:
            pytest.skip("Custom TZ is not supported in function STR for this dialect")

        assert to_str(dbe.eval('STR(DATETIME("2019-01-02 03:04:05", "Pacific/Chatham"))')) == "2019-01-02 03:04:05"
        assert (
            to_str(dbe.eval('STR(GENERICDATETIME("2019-01-02 03:04:05", "Pacific/Chatham"))')) == "2019-01-02 03:04:05"
        )

    def test_str_from_max_ints(self, dbe: DbEvaluator) -> None:
        if self.skip_custom_tz:
            pytest.skip("Max int test is not supported in function STR for this dialect")

        # FIXME: Except for MSSQLSRV
        max_signed_bigint = 2**63 - 1
        min_signed_bigint = -(2**63)
        assert to_str(dbe.eval(f"STR({max_signed_bigint})")) == str(max_signed_bigint)
        assert to_str(dbe.eval(f"STR({min_signed_bigint})")) == str(min_signed_bigint)

    def test_str_from_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_arrays:
            pytest.skip("Arrays are not supported")

        raise NotImplementedError("This test is dialect-dependent")

    def test_str_from_const_null(self, dbe: DbEvaluator):
        if self.empty_str_is_null:
            pytest.skip()

        assert dbe.eval("STR(NULL)") is None

        if self.bool_of_null_is_false:
            # TODO: in any fresh CH, apparently, `if(NULL, 1, 0) == 0`.
            # Might want to reconsider the used formulas.
            assert dbe.eval("STR(BOOL(NULL))") == "False"
        else:
            assert dbe.eval("STR(BOOL(NULL))") is None

        assert dbe.eval("STR(INT(NULL))") is None
        assert dbe.eval("STR(FLOAT(NULL))") is None
        assert dbe.eval("STR(DATE(NULL))") is None
        assert dbe.eval("STR(DATETIME(NULL))") is None
        assert dbe.eval("STR(GENERICDATETIME(NULL))") is None
        assert dbe.eval("STR(GEOPOINT(NULL))") is None
        assert dbe.eval("STR(GEOPOLYGON(NULL))") is None

    def test_str_from_data_null(self, dbe: DbEvaluator, null_data_table: sa.Table):
        if self.empty_str_is_null:
            pytest.skip()

        if self.bool_of_null_is_false:
            # See `test_str_from_const_null` note above.
            assert dbe.eval("STR([bool_null])", from_=null_data_table) == "False"
        else:
            assert dbe.eval("STR([bool_null])", from_=null_data_table) is None

        assert dbe.eval("STR([int_null])", from_=null_data_table) is None
        assert dbe.eval("STR([float_null])", from_=null_data_table) is None
        assert dbe.eval("STR([date_null])", from_=null_data_table) is None
        assert dbe.eval("STR([datetime_null])", from_=null_data_table) is None
        assert dbe.eval("STR([geopoint_null])", from_=null_data_table) is None
        assert dbe.eval("STR([geopolygon_null])", from_=null_data_table) is None

        if self.supports_uuid:
            assert dbe.eval("STR([uuid_null])", from_=null_data_table) is None


class DefaultFloatTypeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    def test_float(self, dbe: DbEvaluator) -> None:
        assert dbe.eval('FLOAT("1.2")') == 1.2
        assert dbe.eval("FLOAT(8)") == 8.0
        assert type(dbe.eval("FLOAT(8)")) == float
        assert dbe.eval("FLOAT(TRUE)") == 1.0
        assert dbe.eval("FLOAT(FALSE)") == 0.0
        assert dbe.eval("FLOAT(#2019-01-02#)") == utc_ts(2019, 1, 2)
        assert type(dbe.eval("FLOAT(#2019-01-02#)")) == float
        assert dbe.eval("FLOAT(#2019-01-02 03:04:05#)") == (
            as_tz(datetime.datetime(2019, 1, 2, 3, 4, 5), tzinfo=dbe.db.tzinfo).timestamp()
        )
        assert type(dbe.eval("FLOAT(#2019-01-02 03:04:05#)")) == float

    def test_float_from_null(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("FLOAT(NULL)") is None

        if self.null_casts_to_false:
            assert dbe.eval("FLOAT(BOOL(NULL))") == 0.0
        else:
            assert dbe.eval("FLOAT(BOOL(NULL))") is None

        if not self.empty_str_is_null:
            assert dbe.eval("FLOAT(STR(NULL))") is None

        assert dbe.eval("FLOAT(INT(NULL))") is None
        assert dbe.eval("FLOAT(FLOAT(NULL))") is None
        assert dbe.eval("FLOAT(DATE(NULL))") is None
        assert dbe.eval("FLOAT(DATETIME(NULL))") is None
        assert dbe.eval("FLOAT(GENERICDATETIME(NULL))") is None

    def test_float_from_data_null(self, dbe: DbEvaluator, null_data_table: sa.Table) -> None:
        if self.null_casts_to_false:
            assert dbe.eval("FLOAT([bool_null])", from_=null_data_table) == 0.0
        else:
            assert dbe.eval("FLOAT([bool_null])", from_=null_data_table) is None

        if not self.empty_str_is_null:
            assert dbe.eval("FLOAT([str_null])", from_=null_data_table) is None

        assert dbe.eval("FLOAT([int_null])", from_=null_data_table) is None
        assert dbe.eval("FLOAT([float_null])", from_=null_data_table) is None
        assert dbe.eval("FLOAT([date_null])", from_=null_data_table) is None
        assert dbe.eval("FLOAT([datetime_null])", from_=null_data_table) is None


class DefaultBoolTypeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    def test_bool(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("BOOL(1.2)") is True
        assert dbe.eval("BOOL(0.0)") is False
        assert dbe.eval("BOOL(8)") is True
        assert dbe.eval("BOOL(0)") is False
        assert dbe.eval("BOOL(TRUE)") is True
        assert dbe.eval("BOOL(FALSE)") is False
        assert dbe.eval('BOOL("qwe")') is True
        assert dbe.eval('BOOL("")') is False
        assert type(dbe.eval("BOOL(1)")) == bool
        assert dbe.eval("BOOL(#2019-01-02#)") is True
        assert dbe.eval("BOOL(#2019-01-02 03:04:05#)") is True

    def test_bool_from_null(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("BOOL(NULL)") is None
        assert dbe.eval("BOOL(BOOL(NULL))") is None
        if not self.bool_is_expression:
            assert dbe.eval("BOOL(STR(NULL))") is None
        if not self.null_casts_to_number:
            assert dbe.eval("BOOL(INT(NULL))") is None
            assert dbe.eval("BOOL(FLOAT(NULL))") is None
        assert dbe.eval("BOOL(DATE(NULL))") is None
        assert dbe.eval("BOOL(DATETIME(NULL))") is None
        assert dbe.eval("BOOL(GENERICDATETIME(NULL))") is None

    def test_bool_from_data_null(self, dbe: DbEvaluator, null_data_table: sa.Table):
        assert dbe.eval("BOOL([bool_null])", from_=null_data_table) is None
        if not self.bool_is_expression:
            assert dbe.eval("BOOL([str_null])", from_=null_data_table) is None
        if not self.null_casts_to_number:
            assert dbe.eval("BOOL([int_null])", from_=null_data_table) is None
            assert dbe.eval("BOOL([float_null])", from_=null_data_table) is None
        assert dbe.eval("BOOL([date_null])", from_=null_data_table) is None
        assert dbe.eval("BOOL([datetime_null])", from_=null_data_table) is None


class DefaultIntTypeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    def test_int(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("INT(432.65)") == 432
        assert dbe.eval("INT(432.35)") == 432
        assert dbe.eval("INT(432)") == 432
        assert dbe.eval("INT(TRUE)") == 1
        assert dbe.eval("INT(FALSE)") == 0
        assert dbe.eval('INT("83")') == 83
        assert dbe.eval("INT(#2019-01-02#)") == int(utc_ts(2019, 1, 2))
        assert type(dbe.eval("INT(#2019-01-02#)")) == int
        assert dbe.eval("INT(#2019-01-02 03:04:05#)") == (
            int(as_tz(datetime.datetime(2019, 1, 2, 3, 4, 5), tzinfo=dbe.db.tzinfo).timestamp())
        )
        assert type(dbe.eval("INT(#2019-01-02 03:04:05#)")) == int

    def test_int_from_null(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("INT(NULL)") is None

        if self.null_casts_to_false:
            assert dbe.eval("INT(BOOL(NULL))") == 0
        else:
            assert dbe.eval("INT(BOOL(NULL))") is None

        if not self.empty_str_is_null:  # NULL can be interpreted as ''
            assert dbe.eval("INT(STR(NULL))") is None
        assert dbe.eval("INT(INT(NULL))") is None
        assert dbe.eval("INT(FLOAT(NULL))") is None
        assert dbe.eval("INT(DATE(NULL))") is None
        assert dbe.eval("INT(DATETIME(NULL))") is None
        assert dbe.eval("INT(GENERICDATETIME(NULL))") is None

    def test_int_from_data_null(self, dbe: DbEvaluator, null_data_table: sa.Table) -> None:
        if self.null_casts_to_false:
            assert dbe.eval("INT([bool_null])", from_=null_data_table) == 0
        else:
            assert dbe.eval("INT([bool_null])", from_=null_data_table) is None

        if not self.empty_str_is_null:  # NULL can be interpreted as ''
            assert dbe.eval("INT([str_null])", from_=null_data_table) is None
        assert dbe.eval("INT([int_null])", from_=null_data_table) is None
        assert dbe.eval("INT([float_null])", from_=null_data_table) is None
        assert dbe.eval("INT([date_null])", from_=null_data_table) is None
        assert dbe.eval("INT([datetime_null])", from_=null_data_table) is None


class DefaultDateTypeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    def test_date(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("DATE(#2019-01-02#)") == datetime.date(2019, 1, 2)
        assert dbe.eval("DATE(#2019-01-02 03:04:05#)") == datetime.date(2019, 1, 2)
        assert dbe.eval('DATE("2019-01-02")') == datetime.date(2019, 1, 2)
        expected_date = (
            datetime.datetime(2019, 1, 2, 3, 4, 5, tzinfo=pytz.UTC).astimezone(dbe.db.tzinfo).date()
        )  # Note that the result depends on db's TZ
        assert dbe.eval("DATE({})".format(int(utc_ts(2019, 1, 2, 3, 4, 5)))) == expected_date
        assert dbe.eval("DATE({})".format(float(utc_ts(2019, 1, 2, 3, 4, 5)))) == expected_date

    def test_date_from_null(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("DATE(NULL)") is None
        if not self.empty_str_is_null:
            assert dbe.eval("DATE(STR(NULL))") is None
        assert dbe.eval("DATE(INT(NULL))") is None
        assert dbe.eval("DATE(FLOAT(NULL))") is None
        assert dbe.eval("DATE(DATE(NULL))") is None
        assert dbe.eval("DATE(DATETIME(NULL))") is None
        assert dbe.eval("DATE(GENERICDATETIME(NULL))") is None

    def test_date_from_data_null(self, dbe: DbEvaluator, null_data_table: sa.Table) -> None:
        if not self.empty_str_is_null:  # NULL can be interpreted as ''
            assert dbe.eval("DATE([str_null])", from_=null_data_table) is None
        assert dbe.eval("DATE([int_null])", from_=null_data_table) is None
        assert dbe.eval("DATE([float_null])", from_=null_data_table) is None
        assert dbe.eval("DATE([date_null])", from_=null_data_table) is None
        assert dbe.eval("DATE([datetime_null])", from_=null_data_table) is None


class DefaultGenericDatetimeTypeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    @pytest.mark.parametrize("func_name", ("GENERICDATETIME", "DATETIME"))
    def test_genericdatetime(self, dbe: DbEvaluator, func_name: str) -> None:
        assert dt_strip(dbe.eval("##2019-01-02 03:04:05##")) == datetime.datetime(2019, 1, 2, 3, 4, 5)
        assert dbe.eval("##2019-01-02##") == datetime.date(2019, 1, 2)
        assert dt_strip(dbe.eval(f"{func_name}(#2019-01-02 03:04:05#)")) == datetime.datetime(2019, 1, 2, 3, 4, 5)
        assert dt_strip(dbe.eval(f"{func_name}(#2019-01-02#)")) == to_datetime(datetime.date(2019, 1, 2))
        expected_values = (
            datetime.datetime(2019, 1, 2, 3, 4, 5),  # UTC CH
            datetime.datetime(2019, 1, 1, 22, 4, 5),  # NYC CH
        )
        assert dt_strip(dbe.eval(f'{func_name}("2019-01-02 03:04:05")')) == datetime.datetime(2019, 1, 2, 3, 4, 5)
        assert dt_strip(dbe.eval(f"{func_name}({int(utc_ts(2019, 1, 2, 3, 4, 5))})")) in expected_values
        assert dt_strip(dbe.eval(f"{func_name}({float(utc_ts(2019, 1, 2, 3, 4, 5))})")) in expected_values

    @pytest.mark.parametrize("func_name", ("GENERICDATETIME", "DATETIME"))
    def test_genericdatetime_from_null(self, dbe: DbEvaluator, func_name: str) -> None:
        assert dbe.eval(f"{func_name}(NULL)") is None
        if not self.empty_str_is_null:  # NULL can be interpreted as ''
            assert dbe.eval(f"{func_name}(STR(NULL))") is None
        assert dbe.eval(f"{func_name}(DATE(NULL))") is None
        assert dbe.eval(f"{func_name}({func_name}(NULL))") is None

    @pytest.mark.parametrize("func_name", ("GENERICDATETIME", "DATETIME"))
    def test_genericdatetime_from_data_null(self, dbe: DbEvaluator, null_data_table: sa.Table, func_name: str) -> None:
        if not self.empty_str_is_null:  # NULL can be interpreted as ''
            assert dbe.eval(f"{func_name}([str_null])", from_=null_data_table) is None
        assert dbe.eval(f"{func_name}([date_null])", from_=null_data_table) is None
        assert dbe.eval(f"{func_name}([datetime_null])", from_=null_data_table) is None


class DefaultGeopointTypeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    def test_geopoint(self, dbe: DbEvaluator) -> None:
        assert to_str(dbe.eval('GEOPOINT("[55.75222,37.61556]")')) == "[55.75222,37.61556]"
        assert to_str(dbe.eval("GEOPOINT(GEOPOINT(55.75222, 37.61556))")) == "[55.75222,37.61556]"
        assert to_str(dbe.eval("GEOPOINT(55.75222, 37.61556)")) == "[55.75222,37.61556]"
        assert to_str(dbe.eval("GEOPOINT(55, 37)")) == "[55,37]"
        assert to_str(dbe.eval('GEOPOINT("55.75222", "37")')) == "[55.75222,37]"
        # Important case with values < 1
        assert to_str(dbe.eval('GEOPOINT("0.75222", "-0.37")')) == "[0.75222,-0.37]"

    def test_geopoint_from_null(self, dbe: DbEvaluator) -> None:
        if not self.empty_str_is_null:
            pytest.skip()

        assert dbe.eval("GEOPOINT(NULL)") is None
        assert dbe.eval("GEOPOINT(STR(NULL))") is None

    def test_geopoint_from_data_null(self, dbe: DbEvaluator, null_data_table: sa.Table) -> None:
        if not self.empty_str_is_null:
            pytest.skip()

        assert dbe.eval("GEOPOINT([str_null])", from_=null_data_table) is None


class DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    def test_geopolygon(self, dbe: DbEvaluator) -> None:
        assert (
            to_str(dbe.eval('GEOPOLYGON("[[[55,37],[56,37],[56,38],[55,38]],[[50,37],[51,37],[51,38],[50,38]]]")'))
            == "[[[55,37],[56,37],[56,38],[55,38]],[[50,37],[51,37],[51,38],[50,38]]]"
        )

    def test_geopolygon_from_null(self, dbe: DbEvaluator) -> None:
        if not self.empty_str_is_null:
            pytest.skip()

        assert dbe.eval("GEOPOLYGON(NULL)") is None
        assert dbe.eval("GEOPOLYGON(STR(NULL))") is None

    def test_geopolygon_from_data_null(self, dbe: DbEvaluator, null_data_table: sa.Table) -> None:
        if not self.empty_str_is_null:
            pytest.skip()

        assert dbe.eval("GEOPOLYGON([str_null])", from_=null_data_table) is None


class DefaultDbCastTypeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    pass


class DefaultTreeTypeFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    def test_tree_str(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_arrays:
            raise NotImplementedError("Trees are not supported")

        assert to_str(dbe.eval('ARR_STR(TREE([arr_str_value]), "+")', from_=data_table)) == "++cde"

lib/dl_formula_testing/dl_formula_testing/testcases/functions_type_conversion.py
datalens-tech/datalens-backend
from __future__ import annotations

import decimal
from typing import (
    Callable,
    ClassVar,
    List,
    Optional,
    Sequence,
    TypeVar,
    Union,
)

import pytest
import sqlalchemy as sa

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase


NUMERIC = Union[decimal.Decimal, int, float]


def approx(value):
    if isinstance(value, list):
        result = []
        for subvalue in value:
            if isinstance(subvalue, decimal.Decimal):
                subvalue = float(subvalue)
            result.append(subvalue)
        value = result
    return pytest.approx(value, rel=1e-4)


def avg(values: List[NUMERIC]) -> NUMERIC:
    return sum(values) / len(values)


class DefaultWindowFunctionFormulaConnectorTestSuite(FormulaConnectorTestBase):
    supports_rank_percentile: ClassVar[bool] = True

    def test_rank(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))

        assert dbe.eval("RANK([int_value])", from_=data_table, order_by=["[id]"], many=True) == (
            [19 - val * 2 for val in range(10)] + [19 - val * 2 for val in range(10)]
        )
        assert dbe.eval('RANK([int_value], "asc")', from_=data_table, order_by=["[id]"], many=True) == (
            [val * 2 + 1 for val in range(10)] + [val * 2 + 1 for val in range(10)]
        )
        assert dbe.eval('RANK(ABS(30 - [int_value]), "asc")', from_=data_table, order_by=["[id]"], many=True) == (
            [11, 7, 3, 1, 3, 7, 11, 15, 17, 19, 11, 7, 3, 1, 3, 7, 11, 15, 17, 19]
        )
        assert dbe.eval('RANK([id], "asc" WITHIN [int_value])', from_=data_table, order_by=["[id]"], many=True) == (
            [1] * 10 + [2] * 10
        )

    def test_rank_dense(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))

        assert dbe.eval("RANK_DENSE([int_value])", from_=data_table, order_by=["[id]"], many=True) == (
            [10 - val for val in range(10)] + [10 - val for val in range(10)]
        )
        assert dbe.eval('RANK_DENSE([int_value], "asc")', from_=data_table, order_by=["[id]"], many=True) == (
            [val + 1 for val in range(10)] + [val + 1 for val in range(10)]
        )
        assert dbe.eval('RANK_DENSE(ABS(30 - [int_value]), "asc")', from_=data_table, order_by=["[id]"], many=True) == (
            [4, 3, 2, 1, 2, 3, 4, 5, 6, 7, 4, 3, 2, 1, 2, 3, 4, 5, 6, 7]
        )
        assert dbe.eval(
            'RANK_DENSE([id], "asc" WITHIN [int_value])', from_=data_table, order_by=["[id]"], many=True
        ) == ([1] * 10 + [2] * 10)

    def test_rank_unique(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))

        assert dbe.eval("RANK_UNIQUE([id])", from_=data_table, order_by=["[id]"], many=True) == list(
            reversed(range(1, 21))
        )
        assert dbe.eval('RANK_UNIQUE([id], "asc")', from_=data_table, order_by=["[id]"], many=True) == list(
            range(1, 21)
        )
        assert set(dbe.eval("RANK_UNIQUE([int_value])", from_=data_table, order_by=["[id]"], many=True)) == set(
            range(1, 21)
        )
        assert set(dbe.eval('RANK_UNIQUE([int_value], "asc")', from_=data_table, order_by=["[id]"], many=True)) == set(
            range(1, 21)
        )
        assert set(
            dbe.eval('RANK_UNIQUE(ABS(30 - [int_value]), "asc")', from_=data_table, order_by=["[id]"], many=True)
        ) == set(range(1, 21))
        assert dbe.eval(
            'RANK_UNIQUE([id], "asc" WITHIN [int_value])', from_=data_table, order_by=["[id]"], many=True
        ) == ([1] * 10 + [2] * 10)

    def test_rank_percentile(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_rank_percentile:
            pytest.skip()

        values = data_table.int_values
        cnt = len(values)
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))

        assert approx(
            dbe.eval('RANK_PERCENTILE([int_value], "asc")', from_=data_table, order_by=["[id]"], many=True)
        ) == ([(val * 2) / (cnt - 1) for val in range(10)] + [(val * 2) / (cnt - 1) for val in range(10)])

    def test_sum(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        assert dbe.eval("SUM([int_value] TOTAL)", from_=data_table, order_by=["[id]"], many=True) == (
            [sum(values)] * 20
        )
        assert dbe.eval("SUM([id] WITHIN [int_value])", from_=data_table, order_by=["[id]"], many=True) == (
            [12 + val * 2 for val in range(10)] + [12 + val * 2 for val in range(10)]
        )

    def test_min(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        assert dbe.eval("MIN([int_value] TOTAL)", from_=data_table, order_by=["[id]"], many=True) == (
            [min(values)] * 20
        )
        assert dbe.eval("MIN([id] WITHIN [int_value])", from_=data_table, order_by=["[id]"], many=True) == (
            [1 + val for val in range(10)] + [1 + val for val in range(10)]
        )

    def test_max(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        assert dbe.eval("MAX([int_value] TOTAL)", from_=data_table, order_by=["[id]"], many=True) == (
            [max(values)] * 20
        )
        assert dbe.eval("MAX([id] WITHIN [int_value])", from_=data_table, order_by=["[id]"], many=True) == (
            [11 + val for val in range(10)] + [11 + val for val in range(10)]
        )

    def test_count(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        # 1 arg
        assert dbe.eval("COUNT([int_value] TOTAL)", from_=data_table, order_by=["[id]"], many=True) == (
            [len(values)] * 20
        )
        assert dbe.eval("COUNT([id] WITHIN [int_value])", from_=data_table, order_by=["[id]"], many=True) == (
            [2 for _ in range(10)] + [2 for _ in range(10)]
        )
        # 0 args
        assert dbe.eval("COUNT(TOTAL)", from_=data_table, order_by=["[id]"], many=True) == ([len(values)] * 20)
        assert dbe.eval("COUNT(WITHIN [int_value])", from_=data_table, order_by=["[id]"], many=True) == (
            [2 for _ in range(10)] + [2 for _ in range(10)]
        )

    def test_avg(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        assert approx(dbe.eval("AVG([int_value] TOTAL)", from_=data_table, order_by=["[id]"], many=True)) == (
            [sum(values) / 20] * 20
        )
        assert approx(dbe.eval("AVG([id] WITHIN [int_value])", from_=data_table, order_by=["[id]"], many=True)) == (
            [6 + val for val in range(10)] + [6 + val for val in range(10)]
        )

    def test_sum_if(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        assert dbe.eval("SUM_IF([int_value], [id] < 11 TOTAL)", from_=data_table, order_by=["[id]"], many=True) == (
            [sum(v for id_0, v in enumerate(values) if id_0 + 1 < 11)] * 20  # (id_0 = id - 1)
        )
        assert dbe.eval(
            "SUM_IF([id], [id] < 11 WITHIN [int_value])", from_=data_table, order_by=["[id]"], many=True
        ) == ([val for val in range(1, 11)] + [val for val in range(1, 11)])

    def test_count_if(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        assert dbe.eval("COUNT_IF([int_value], [id] < 11 TOTAL)", from_=data_table, order_by=["[id]"], many=True) == (
            [10] * 20
        )
        assert dbe.eval(
            "COUNT_IF([id], [id] < 11 WITHIN [int_value])", from_=data_table, order_by=["[id]"], many=True
        ) == ([1] * 10 + [1] * 10)

    def test_avg_if(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        assert dbe.eval(
            "AVG_IF([int_value], [id] < 11 TOTAL)", from_=data_table, order_by=["[id]"], many=True
        ) == pytest.approx(
            [sum(v for id_0, v in enumerate(values) if id_0 + 1 < 11) / 10.0] * 20  # (id_0 = id - 1)
        )
        assert dbe.eval(
            "AVG_IF([id], [id] < 11 WITHIN [int_value])", from_=data_table, order_by=["[id]"], many=True
        ) == pytest.approx([val / 1.0 for val in range(1, 11)] + [val for val in range(1, 11)])

    def _check_rfunc(
        self, dbe: DbEvaluator, data_table: sa.Table, func_name: str, py_agg_func: Callable[[List[NUMERIC]], NUMERIC]
    ) -> None:
        ids = list(range(1, 21))
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        assert approx(dbe.eval(f"{func_name}([int_value] TOTAL)", from_=data_table, order_by=["[id]"], many=True)) == (
            [py_agg_func(values[: i + 1]) for i in range(20)]
        )
        assert approx(
            dbe.eval(f"{func_name}([id] WITHIN [int_value])", from_=data_table, order_by=["[id]"], many=True)
        ) == ([py_agg_func([ids[i]]) for i in range(10)] + [py_agg_func([ids[i], ids[i + 10]]) for i in range(10)])
        assert approx(
            dbe.eval(f'{func_name}([int_value], "desc" TOTAL)', from_=data_table, order_by=["[id]"], many=True)
        ) == ([py_agg_func(values[i:]) for i in range(20)])
        assert approx(
            dbe.eval(
                f'{func_name}([int_value], "asc" TOTAL ORDER BY [id] DESC)',
                from_=data_table,
                order_by=["[id]"],
                many=True,
            )
        ) == ([py_agg_func(values[i:]) for i in range(20)])
        assert approx(
            dbe.eval(f'{func_name}([id], "desc" WITHIN [int_value])', from_=data_table, order_by=["[id]"], many=True)
        ) == ([py_agg_func([ids[i], ids[i + 10]]) for i in range(10)] + [py_agg_func([ids[i]]) for i in range(10, 20)])

    def test_rsum(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        self._check_rfunc(dbe, data_table, func_name="RSUM", py_agg_func=sum)

    def test_rcount(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        self._check_rfunc(dbe, data_table, func_name="RCOUNT", py_agg_func=len)

    def test_rmin(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        self._check_rfunc(dbe, data_table, func_name="RMIN", py_agg_func=min)

    def test_rmax(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        self._check_rfunc(dbe, data_table, func_name="RMAX", py_agg_func=max)

    def test_ravg(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        self._check_rfunc(dbe, data_table, func_name="RAVG", py_agg_func=avg)

    def _check_mfunc(
        self, dbe: DbEvaluator, data_table: sa.Table, func_name: str, py_agg_func: Callable[[List[NUMERIC]], NUMERIC]
    ) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        assert approx(dbe.eval(f"{func_name}([int_value], 1)", from_=data_table, order_by=["[id]"], many=True)) == (
            [py_agg_func(values[max(0, i - 1) : i + 1]) for i in range(20)]
        )
        assert approx(
            dbe.eval(f"{func_name}([int_value], 1 ORDER BY [id] DESC)", from_=data_table, order_by=["[id]"], many=True)
        ) == ([py_agg_func(values[i : i + 2]) for i in range(20)])
        assert approx(dbe.eval(f"{func_name}([int_value], -1)", from_=data_table, order_by=["[id]"], many=True)) == (
            [py_agg_func(values[i : i + 2]) for i in range(20)]
        )
        assert approx(dbe.eval(f"{func_name}([int_value], 5, 5)", from_=data_table, order_by=["[id]"], many=True)) == (
            [py_agg_func(values[max(0, i - 5) : i + 6]) for i in range(20)]
        )

    def test_msum(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        self._check_mfunc(dbe, data_table, func_name="MSUM", py_agg_func=sum)

    def test_mcount(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        self._check_mfunc(dbe, data_table, func_name="MCOUNT", py_agg_func=len)

    def test_mmin(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        self._check_mfunc(dbe, data_table, func_name="MMIN", py_agg_func=min)

    def test_mmax(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        self._check_mfunc(dbe, data_table, func_name="MMAX", py_agg_func=max)

    def test_mavg(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        self._check_mfunc(dbe, data_table, func_name="MAVG", py_agg_func=avg)

    def test_lag(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        VALUE_TV = TypeVar("VALUE_TV")

        def lag(
            values: Sequence[VALUE_TV],
            idx: int,
            offset: int = 1,
            default: Optional[VALUE_TV] = None,
        ) -> Optional[VALUE_TV]:
            offset_idx = idx - offset
            if offset_idx < 0 or offset_idx >= len(values):
                return default
            return values[offset_idx]

        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))
        assert dbe.eval("LAG([int_value])", from_=data_table, order_by=["[id]"], many=True) == (
            [lag(values, i, offset=1) for i in range(20)]
        )
        assert dbe.eval("LAG([int_value] ORDER BY [id] DESC)", from_=data_table, order_by=["[id]"], many=True) == (
            [lag(values, i, offset=-1) for i in range(20)]
        )
        assert dbe.eval("LAG([int_value], -2)", from_=data_table, order_by=["[id]"], many=True) == (
            [lag(values, i, offset=-2) for i in range(20)]
        )
        assert dbe.eval("LAG([int_value], 2)", from_=data_table, order_by=["[id]"], many=True) == (
            [lag(values, i, offset=2) for i in range(20)]
        )
        # With default
        assert dbe.eval("LAG([int_value], -3, 8)", from_=data_table, order_by=["[id]"], many=True) == (
            [lag(values, i, offset=-3, default=8) for i in range(20)]
        )
        assert dbe.eval("LAG([int_value], 3, 8)", from_=data_table, order_by=["[id]"], many=True) == (
            [lag(values, i, offset=3, default=8) for i in range(20)]
        )

    def test_first(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))

        assert dbe.eval("FIRST([int_value] TOTAL ORDER BY [id])", from_=data_table, many=True) == [values[0]] * len(
            values
        )
        assert dbe.eval("FIRST([int_value] TOTAL ORDER BY [id] DESC)", from_=data_table, many=True) == [
            values[-1]
        ] * len(values)

    def test_last(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        values = data_table.int_values
        assert values == list(range(0, 100, 10)) + list(range(0, 100, 10))

        assert dbe.eval("LAST([int_value] TOTAL ORDER BY [id])", from_=data_table, many=True) == [values[-1]] * len(
            values
        )
        assert dbe.eval("LAST([int_value] TOTAL ORDER BY [id] DESC)", from_=data_table, many=True) == [values[0]] * len(
            values
        )

lib/dl_formula_testing/dl_formula_testing/testcases/functions_window.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from typing import (
    ClassVar,
    Optional,
    Union,
)

import dateutil
import pytest
import pytz

from dl_formula.definitions.literals import literal
from dl_formula.shortcuts import n
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import (
    dt_strip,
    to_str,
)


def as_dt(value: Union[str, datetime.datetime]) -> datetime.datetime:
    if isinstance(value, str):
        value = dateutil.parser.parse(value)
    assert isinstance(value, datetime.datetime)
    return value


class DefaultLiteralFormulaConnectorTestSuite(FormulaConnectorTestBase):
    supports_microseconds: ClassVar[bool]
    supports_utc: ClassVar[bool]
    supports_custom_tz: ClassVar[bool]
    default_tz: ClassVar[Optional[datetime.tzinfo]] = None
    recognizes_datetime_type: ClassVar[bool] = True

    def test_number(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("1") == 1
        assert type(dbe.eval("1")) is int
        assert dbe.eval("1.2") == 1.2

    def test_string(self, dbe: DbEvaluator) -> None:
        assert to_str(dbe.eval('"qwerty"')) == "qwerty"
        assert to_str(dbe.eval('"Lorem\\nipsum"')) == "Lorem\nipsum"

    def test_null(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("NULL") is None

    def test_bool(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("TRUE") is True
        assert type(dbe.eval("TRUE")) is bool
        assert dbe.eval("FALSE") is False
        assert type(dbe.eval("FALSE")) is bool

    def test_date(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("#2019-01-02#") == datetime.date(2019, 1, 2)
        assert type(dbe.eval("#2019-01-02#")) is datetime.date

    def test_datetime(self, dbe: DbEvaluator) -> None:
        assert dt_strip(dbe.eval("#2019-01-02 12:34:56#")) == datetime.datetime(2019, 1, 2, 12, 34, 56)
        if self.recognizes_datetime_type:
            assert type(dbe.eval("#2019-01-02 12:34:56#")) is datetime.datetime

    def test_datetime_extended(self, dbe: DbEvaluator) -> None:
        # Without microseconds
        assert as_dt(
            dbe.db.scalar(literal(datetime.datetime(2020, 4, 8, 12, 34, 56), d=dbe.dialect))
        ) == datetime.datetime(2020, 4, 8, 12, 34, 56, tzinfo=self.default_tz)

        # With microseconds
        assert as_dt(
            dbe.db.scalar(literal(datetime.datetime(2020, 4, 8, 12, 34, 56, microsecond=7890), d=dbe.dialect))
        ) == datetime.datetime(
            2020, 4, 8, 12, 34, 56, microsecond=7890 if self.supports_microseconds else 0, tzinfo=self.default_tz
        )

        # With UTC without microseconds
        if self.supports_utc:
            assert as_dt(
                dbe.db.scalar(
                    literal(datetime.datetime(2020, 4, 8, 12, 34, 56, tzinfo=datetime.timezone.utc), d=dbe.dialect)
                )
            ) == datetime.datetime(2020, 4, 8, 12, 34, 56, tzinfo=datetime.timezone.utc)

        # With UTC and microseconds
        if self.supports_utc:
            assert as_dt(
                dbe.db.scalar(
                    literal(
                        datetime.datetime(2020, 4, 8, 12, 34, 56, microsecond=7890, tzinfo=datetime.timezone.utc),
                        d=dbe.dialect,
                    )
                )
            ) == datetime.datetime(
                2020,
                4,
                8,
                12,
                34,
                56,
                microsecond=7890 if self.supports_microseconds else 0,
                tzinfo=datetime.timezone.utc,
            )

        # With custom tz
        if self.supports_custom_tz:
            assert as_dt(
                dbe.db.scalar(
                    literal(
                        pytz.timezone("Europe/Moscow").localize(datetime.datetime(2020, 4, 8, 15, 34, 56, tzinfo=None)),
                        d=dbe.dialect,
                    )
                )
            ) == datetime.datetime(2020, 4, 8, 12, 34, 56, tzinfo=datetime.timezone.utc)

    def test_arrays(self, dbe) -> None:
        if not self.supports_arrays:
            pytest.skip("Not supported")
        assert dbe.eval(n.formula(n.func.GET_ITEM(n.lit([1, 4, 6]), 2))) == 4
        assert dbe.eval(n.formula(n.func.GET_ITEM(n.lit([1.1, 4.2, 6.3]), 2))) == 4.2
        assert dbe.eval(n.formula(n.func.GET_ITEM(n.lit(["qwe", "rty", "uio"]), 2))) == "rty"

lib/dl_formula_testing/dl_formula_testing/testcases/literals.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest
import sqlalchemy as sa

from dl_formula.core.dialect import from_name_and_version
import dl_formula.core.exc as exc
from dl_formula.shortcuts import n
from dl_formula.translation.ext_nodes import CompiledExpression
from dl_formula.translation.translator import translate
from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase


class DefaultMiscFunctionalityConnectorTestSuite(FormulaConnectorTestBase):
    def test_dialect_resolution(self, dbe: DbEvaluator) -> None:
        assert from_name_and_version(self.dialect.common_name, dbe.db.get_version()) == dbe.dialect

    def test_translate_for_where(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        # WHERE already is a condition
        assert dbe.eval("MIN([int_value])", from_=data_table, where="[int_value] > 30") == 40
        # WHERE is a boolean value, might have to be converted to a condition
        assert dbe.eval("MIN([int_value])", from_=data_table, where="IF([int_value] > 30, TRUE, FALSE)") == 40

    def test_translation_recombination(self, dbe: DbEvaluator) -> None:
        # In regular circumstances such a formula raises an error,
        # but here its parts are translated separately and then combined into a single expression.
        # Should still raise the same error
        dialect = self.dialect
        with pytest.raises(exc.TranslationError):
            translate(
                n.formula(
                    n.func.IF(
                        CompiledExpression.make(
                            translate(
                                n.formula(n.func.WHAA(n.field("n1"))),
                                dialect=dialect,
                            )
                        ),
                        CompiledExpression.make(
                            translate(
                                n.formula(n.field("n2")),
                                dialect=dialect,
                            )
                        ),
                        n.lit(42),
                    )
                ),
                dialect=dialect,
            )

lib/dl_formula_testing/dl_formula_testing/testcases/misc_funcs.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from typing import (
    Any,
    ClassVar,
    Optional,
)

import pytest
import sqlalchemy as sa

from dl_formula_testing.evaluator import DbEvaluator
from dl_formula_testing.testcases.base import FormulaConnectorTestBase
from dl_formula_testing.util import (
    dt_strip,
    to_str,
)


class DefaultOperatorFormulaConnectorTestSuite(FormulaConnectorTestBase):
    subtraction_round_dt: ClassVar[bool] = True
    supports_string_int_multiplication: ClassVar[bool] = True
    supports_null_in: ClassVar[bool] = True
    make_float_array_cast: Optional[str] = None

    def test_negation(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval("-__LIT__(2)") == -2

    def test_addition_number_number(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval("__LIT__(2) + 3") == 5

    def test_addition_string_string(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert to_str(dbe.eval('__LIT__("Lorem") + " ipsum"')) == "Lorem ipsum"

    def test_addition_date_number(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("#2019-01-06# + 2") == datetime.date(2019, 1, 8)
        assert dbe.eval("2 + #2019-01-06#") == datetime.date(2019, 1, 8)
        assert dbe.eval("#2019-01-06# + 2.2") == datetime.date(2019, 1, 8)
        assert dbe.eval("2.2 + #2019-01-06#") == datetime.date(2019, 1, 8)

    def test_addition_datetime_number(self, dbe: DbEvaluator) -> None:
        assert dt_strip(dbe.eval("#2019-01-06 03# + 2")) == datetime.datetime(2019, 1, 8, 3)
        assert dt_strip(dbe.eval("2 + #2019-01-06 03#")) == datetime.datetime(2019, 1, 8, 3)
        assert dt_strip(dbe.eval("#2019-01-06 03# + 2.5")) == datetime.datetime(2019, 1, 8, 15)
        assert dt_strip(dbe.eval("2.5 + #2019-01-06 03#")) == datetime.datetime(2019, 1, 8, 15)
        assert dt_strip(dbe.eval("2.0005 + #2019-01-06 03#")) == datetime.datetime(2019, 1, 8, 3, 0, 43)
        assert dt_strip(dbe.eval("#2019-01-06 03# + 2.0005")) == datetime.datetime(2019, 1, 8, 3, 0, 43)

    def test_addition_genericdatetime_number(self, dbe: DbEvaluator) -> None:
        assert dt_strip(dbe.eval("##2019-01-06 03## + 2")) == datetime.datetime(2019, 1, 8, 3)
        assert dt_strip(dbe.eval("2 + ##2019-01-06 03##")) == datetime.datetime(2019, 1, 8, 3)
        assert dt_strip(dbe.eval("##2019-01-06 03## + 2.5")) == datetime.datetime(2019, 1, 8, 15)
        assert dt_strip(dbe.eval("2.5 + ##2019-01-06 03##")) == datetime.datetime(2019, 1, 8, 15)
        assert dt_strip(dbe.eval("2.0005 + ##2019-01-06 03##")) == datetime.datetime(2019, 1, 8, 3, 0, 43)
        assert dt_strip(dbe.eval("##2019-01-06 03## + 2.0005")) == datetime.datetime(2019, 1, 8, 3, 0, 43)

    def test_subtraction_number_number(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval("__LIT__(2) - 3") == -1

    def test_subtraction_date_number(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("#2019-01-06# - 2") == datetime.date(2019, 1, 4)
        assert dbe.eval("#2019-01-06# - 2.2") == datetime.date(2019, 1, 3)

    def test_subtraction_datetime_number(self, dbe: DbEvaluator) -> None:
        assert dt_strip(dbe.eval("#2019-01-06 03# - 2")) == datetime.datetime(2019, 1, 4, 3)
        assert dt_strip(dbe.eval("#2019-01-06 03# - 2.5")) == datetime.datetime(2019, 1, 3, 15)
        if not self.subtraction_round_dt:
            assert dt_strip(dbe.eval("#2019-01-06 03# - 2.0005")) == datetime.datetime(2019, 1, 4, 2, 59, 16)
        else:
            assert dt_strip(dbe.eval("#2019-01-06 03# - 2.0005")) == datetime.datetime(2019, 1, 4, 2, 59, 17)

    def test_subtraction_genericdatetime_number(self, dbe: DbEvaluator) -> None:
        assert dt_strip(dbe.eval("##2019-01-06 03## - 2")) == datetime.datetime(2019, 1, 4, 3)
        assert dt_strip(dbe.eval("##2019-01-06 03## - 2.5")) == datetime.datetime(2019, 1, 3, 15)
        if not self.subtraction_round_dt:
            assert dt_strip(dbe.eval("##2019-01-06 03## - 2.0005")) == datetime.datetime(2019, 1, 4, 2, 59, 16)
        else:
            assert dt_strip(dbe.eval("##2019-01-06 03## - 2.0005")) == datetime.datetime(2019, 1, 4, 2, 59, 17)

    def test_subtraction_date_time_date_time(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("#2019-01-06# - #2019-01-02#") == 4
        assert dbe.eval("#2019-01-06 15# - #2019-01-02 03#") == pytest.approx(4.5)

    def test_subtraction_genericdatetime_genericdatetime(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("##2019-01-06## - ##2019-01-02##") == 4
        assert dbe.eval("##2019-01-06 15## - ##2019-01-02 03##") == pytest.approx(4.5)

    def test_multiplication(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval("__LIT__(2) * 3") == 6
        assert dbe.eval("__LIT__(2.1) * 3") == pytest.approx(6.3)
        assert to_str(dbe.eval('"Lorem" * 3')) == "LoremLoremLorem"
        assert to_str(dbe.eval('3 * "Lorem"')) == "LoremLoremLorem"
        if self.supports_string_int_multiplication:
            assert to_str(dbe.eval('__LIT__("Lorem") * 3')) == "LoremLoremLorem"
            assert to_str(dbe.eval('__LIT__(3) * "Lorem"')) == "LoremLoremLorem"

    def test_division(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval("__LIT__(4) / 2") == 2
        assert dbe.eval("__LIT__(5) / 2") == pytest.approx(2.5)
        assert dbe.eval("__LIT__(5.0) / 2") == pytest.approx(2.5)
        assert dbe.eval("__LIT__(5) / 2.0") == pytest.approx(2.5)

    def test_modulo(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval("__LIT__(2) % 3") == 2
        assert dbe.eval("__LIT__(2.1) % 3") == pytest.approx(2.1)
        assert dbe.eval("__LIT__(3) % 2.1") == pytest.approx(0.9)

    def test_power(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval("__LIT__(2) ^ 3") == 8
        assert dbe.eval("__LIT__(2.1) ^ -0.3") == pytest.approx(2.1**-0.3)

    def test_grouping_with_operators(self, dbe: DbEvaluator) -> None:
        """Make sure that operand grouping is performed correctly with field substitution"""
        assert (
            dbe.eval(
                "[f1] / [f2]",
                other_fields={
                    "f1": "ROUND(1000)",
                    "f2": "(ROUND(100) / 10) * 10",
                },
            )
            == 10
        )

    def test_is_null(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert not dbe.eval('__LIT__("qwerty") IS NULL')
        assert dbe.eval("__LIT__(NULL) IS NULL")
        assert dbe.eval('__LIT__("qwerty") IS NOT NULL')
        assert not dbe.eval("__LIT__(NULL) IS NOT NULL")

    def test_is_bool(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval('__LIT__("qwerty") IS TRUE')
        assert dbe.eval("__LIT__(123) IS TRUE")
        assert dbe.eval("__LIT__(TRUE) IS TRUE")
        assert not dbe.eval("__LIT__(FALSE) IS TRUE")
        assert dbe.eval("__LIT__(FALSE) IS NOT TRUE")
        assert dbe.eval("#2019-03-05# IS TRUE")
        assert dbe.eval("#2019-03-05 01:02:03# IS TRUE")
        assert dbe.eval("##2019-03-05 01:02:03## IS TRUE")
        assert dbe.eval('__LIT__("") IS FALSE')
        assert dbe.eval("__LIT__(0) IS FALSE")
        assert dbe.eval("__LIT__(FALSE) IS FALSE")
        assert not dbe.eval("__LIT__(TRUE) IS FALSE")
        assert dbe.eval("__LIT__(TRUE) IS NOT FALSE")
        assert not dbe.eval("#2019-03-05# IS FALSE")
        assert not dbe.eval("#2019-03-05 01:02:03# IS FALSE")
        assert not dbe.eval("##2019-03-05 01:02:03## IS FALSE")

    def test_like(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval('"raspberry pi" LIKE "%spb%"')
        assert not dbe.eval('"raspberry pi" NOT LIKE "%spb%"')

    def test_not(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval("NOT __LIT__(FALSE)")
        assert not dbe.eval("NOT __LIT__(TRUE)")
        assert dbe.eval('NOT __LIT__("")')
        assert not dbe.eval('NOT __LIT__("text")')
        assert dbe.eval("NOT __LIT__(0)")
        assert not dbe.eval("NOT __LIT__(1)")
        assert not dbe.eval("NOT __LIT__(#2019-01-01#)")
        assert not dbe.eval("NOT __LIT__(#2019-01-01 03:00:00#)")
        assert not dbe.eval("NOT ##2019-01-01 03:00:00##")

    def test_in(self, dbe: DbEvaluator, forced_literal_use: Any) -> None:
        assert dbe.eval("3 IN (23, 5, 3, 67)")
        assert not dbe.eval("3 NOT IN (23, 5, 3, 67)")

        assert not dbe.eval("3 in (1, 2, NULL)")
        assert dbe.eval("3 not in (1, 2, NULL)")

        if self.supports_null_in:
            assert dbe.eval("NULL in (1, 2, NULL)")
            assert not dbe.eval("NULL not in (1, 2, NULL)")

        assert not dbe.eval("3 in (NULL)")
        assert dbe.eval("3 not in (NULL)")

        assert dbe.eval("NULL in (NULL)")
        assert not dbe.eval("NULL not in (NULL)")

    def test_in_date(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        assert dbe.eval("[date_value] in (#2014-10-05#)", order_by=["[date_value]"], from_=data_table)

    def test_between(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("3 BETWEEN 1 AND 100")
        assert not dbe.eval("3 NOT BETWEEN 1 AND 100")

    def test_comparison_operators(self, dbe: DbEvaluator) -> None:
        assert dbe.eval("1 < 3 <= 4 = 4 != 10 > 7 >= 6")
        assert dbe.eval('"qwerty" = "qwerty"')
        assert dbe.eval('"qwerty" != "abc"')
        assert dbe.eval('"qwerty" > "abc"')
        assert dbe.eval('"qwerty" >= "abc"')
        assert dbe.eval('"abc" < "qwerty"')
        assert dbe.eval('"abc" <= "qwerty"')
        if self.supports_arrays:
            assert dbe.eval("ARRAY(1, 2, NULL, 4) = ARRAY(1, 2, NULL, 4)")

    @pytest.mark.parametrize("lit", (pytest.param("##", id="generic"), pytest.param("#", id="regular")))
    def test_comparison_operators_datetimes(self, dbe: DbEvaluator, lit: str) -> None:
        def _dt_lit(s: str) -> str:
            return f"{lit}{s}{lit}"

        def _assert_eq(first: str, second: str) -> None:
            assert dbe.eval(f"{_dt_lit(first)} = {_dt_lit(second)}")
            assert dbe.eval(f"{_dt_lit(first)} >= {_dt_lit(second)}")
            assert dbe.eval(f"{_dt_lit(first)} <= {_dt_lit(second)}")

        def _assert_neq(first: str, second: str) -> None:
            assert dbe.eval(f"{_dt_lit(first)} != {_dt_lit(second)}")
            assert dbe.eval(f"{_dt_lit(first)} < {_dt_lit(second)}")
            assert dbe.eval(f"{_dt_lit(first)} <= {_dt_lit(second)}")
            assert dbe.eval(f"{_dt_lit(second)} > {_dt_lit(first)}")
            assert dbe.eval(f"{_dt_lit(second)} >= {_dt_lit(first)}")

        # dates
        _assert_eq("2019-08-08", "2019-08-08")
        _assert_neq("2019-08-08", "2019-08-09")

        # datetime
        _assert_eq("2019-08-08 05:00:00", "2019-08-08 05:00:00")
        _assert_neq("2019-08-08 05:00:00", "2019-08-08 05:00:30")

        # date/datetime
        _assert_eq("2019-08-08", "2019-08-08 00:00:00")
        _assert_eq("2019-08-08 00:00:00", "2019-08-08")
        _assert_neq("2019-08-08", "2019-08-08 05:00:30")
        _assert_neq("2019-08-08 05:00:30", "2019-08-09")

    def test_addition_array_array(self, dbe: DbEvaluator, data_table: sa.Table) -> None:
        if not self.supports_arrays:
            pytest.skip()

        assert dbe.eval("[arr_int_value] + [arr_int_value]", from_=data_table) == dbe.eval(
            "ARRAY(0, 23, 456, NULL, 0, 23, 456, NULL)"
        )
        assert dbe.eval("ARRAY(1, 2, 3) + ARRAY(6, NULL, 0)", from_=data_table) == dbe.eval(
            "ARRAY(1, 2, 3, 6, NULL, 0)"
        )

        assert dbe.eval("[arr_str_value] + [arr_str_value]", from_=data_table) == dbe.eval(
            'ARRAY("", "", "cde", NULL, "", "", "cde", NULL)'
        )
        assert dbe.eval('ARRAY("a", "", "3") + ARRAY("c", NULL, "e")', from_=data_table) == dbe.eval(
            'ARRAY("a", "", "3", "c", NULL, "e")'
        )

        if self.make_float_array_cast is not None:
            assert dbe.eval("[arr_float_value] + [arr_float_value]", from_=data_table) == dbe.eval(
                f'DB_CAST(ARRAY(0, 45, 0.123, NULL, 0, 45, 0.123, NULL), "{self.make_float_array_cast}")'
            )
            assert dbe.eval(
                f'DB_CAST(ARRAY(1.1, 2.2, 3.3) + ARRAY(6.6, NULL, 0), "{self.make_float_array_cast}")', from_=data_table
            ) == dbe.eval("ARRAY(1.1, 2.2, 3.3, 6.6, NULL, 0)")
        else:
            assert dbe.eval("[arr_float_value] + [arr_float_value]", from_=data_table) == dbe.eval(
                "ARRAY(0, 45, 0.123, NULL, 0, 45, 0.123, NULL)"
            )
            assert dbe.eval("ARRAY(1.1, 2.2, 3.3) + ARRAY(6.6, NULL, 0)", from_=data_table) == dbe.eval(
                "ARRAY(1.1, 2.2, 3.3, 6.6, NULL, 0)"
            )

lib/dl_formula_testing/dl_formula_testing/testcases/operators.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import (
    contextmanager,
    redirect_stderr,
    redirect_stdout,
)
import io
import sys
from typing import (
    TYPE_CHECKING,
    Generator,
    Optional,
    TextIO,
    Type,
)


if TYPE_CHECKING:
    from argparse import ArgumentParser


@contextmanager
def redirect_stdin(stream: Optional[TextIO] = None) -> Generator[None, None, None]:
    if stream is None:
        yield

    else:
        old_stdin = sys.stdin
        try:
            sys.stdin = stream
            yield
        finally:
            sys.stdin = old_stdin


class ToolRunner:
    def __init__(self, parser: ArgumentParser, tool_cls: Type):
        self.parser = parser
        self.tool_cls = tool_cls

    def run(self, args: list[str], stdin: Optional[str] = None):
        stdin_str = stdin
        stdout = io.StringIO()
        stderr = io.StringIO()
        stdin_stream: Optional[TextIO] = None
        if stdin_str:
            stdin_stream = io.StringIO()
            assert stdin_stream is not None
            stdin_stream.write(stdin_str)
            stdin_stream.seek(0)
        with redirect_stdout(stdout), redirect_stderr(stderr), redirect_stdin(stdin_stream):
            self.tool_cls.run(self.parser.parse_args(args))  # noqa

        return stdout.getvalue(), stderr.getvalue()

lib/dl_formula_testing/dl_formula_testing/tool_runner.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from typing import (
    Optional,
    Union,
)

import dateutil.parser
import dateutil.relativedelta


def now():
    return datetime.datetime.utcnow().replace(microsecond=0)


def today():
    return datetime.datetime.utcnow().date()


hour = datetime.timedelta(hours=1)
day = datetime.timedelta(days=1)
month = dateutil.relativedelta.relativedelta(months=1)
year = dateutil.relativedelta.relativedelta(years=1)


class approximately:
    def __init__(self, value, rel_tol=1e-09, abs_tol=1e-09):
        self._value = value
        self._rel_tol = rel_tol
        self._abs_tol = abs_tol

    def __eq__(self, other):
        other = float(other)
        return abs(self._value - other) <= max(self._rel_tol * max(abs(self._value), abs(other)), self._abs_tol)


class approx_datetime:
    def __init__(self, value):
        self._value = value

    def __eq__(self, other):
        return abs(self._value - other) <= datetime.timedelta(seconds=2)


class approx_timestamp(approximately):
    def __init__(self, value):
        super().__init__(value, abs_tol=1)


def to_str(v) -> str:
    if isinstance(v, bytes):
        return v.decode("ascii")
    return v


def to_unicode(v) -> str:
    if isinstance(v, bytes):
        return v.decode("utf-8")
    return v


def to_datetime(date: datetime.date) -> datetime.datetime:
    return datetime.datetime(date.year, date.month, date.day)


def dt_strip(value: Union[datetime.datetime, str]) -> datetime.datetime:
    if isinstance(value, str):
        value = dateutil.parser.parse(value)
    assert isinstance(value, datetime.datetime)
    return value.replace(microsecond=0, tzinfo=None)


def utcize(dt: datetime.datetime) -> datetime.datetime:
    return dt.replace(tzinfo=datetime.timezone.utc)


def utc_ts(*args, tzinfo: Optional[datetime.tzinfo] = None) -> float:
    if len(args) == 1:
        dt = args[0]
        if not isinstance(dt, datetime.datetime):
            dt = to_datetime(dt)
    elif len(args) >= 3:
        dt = datetime.datetime(*args)
    else:
        raise ValueError(args)

    return utcize(dt).timestamp()


def as_tz(dt: datetime.datetime, tzinfo: datetime.tzinfo) -> datetime.datetime:
    assert hasattr(tzinfo, "localize"), "non-pytz time zones are not supported"
    # Sadly pytz has no single base class defining the `localize` interface
    return tzinfo.localize(dt)  # type: ignore

lib/dl_formula_testing/dl_formula_testing/util.py
datalens-tech/datalens-backend
class I18nException(Exception):
    pass


class UnknownLocale(I18nException):
    pass


class UnknownDomain(I18nException):
    pass


class InvalidLocaleConfig(I18nException):
    pass

lib/dl_i18n/dl_i18n/exc.py
datalens-tech/datalens-backend
from collections import defaultdict
import gettext
import logging
from typing import (
    Iterable,
    Optional,
)

import attr

from dl_i18n.exc import (
    UnknownDomain,
    UnknownLocale,
)


LOGGER = logging.getLogger(__name__)


@attr.s
class Translatable:
    _s: str = attr.ib()
    domain: str = attr.ib()

    def to_text(self) -> str:
        return self._s


@attr.s(frozen=True)
class TranslationConfig:
    path: str = attr.ib()
    locale: str = attr.ib()
    domain: str = attr.ib()

    @property
    def key(self) -> tuple:
        return self.domain, self.locale


@attr.s
class _GNULocalizer:
    domain: str = attr.ib()
    locale: str = attr.ib()
    _localizer: gettext.GNUTranslations = attr.ib()

    def translate(self, message: Translatable) -> str:
        return self._localizer.gettext(message.to_text())


class Localizer:
    def __init__(self, localizers: Iterable[_GNULocalizer]):
        self._localizers = {}
        for localizer in localizers:
            self._localizers[localizer.domain] = localizer

    def translate(self, message: Translatable) -> str:
        if message.domain not in self._localizers:
            raise UnknownDomain(f"Unknown domain {message.domain}")
        return self._localizers[message.domain].translate(message)


class LocalizerFactory:
    def __init__(self, localizers: Iterable[_GNULocalizer]):
        self._localizers: dict[str, list[_GNULocalizer]] = defaultdict(list)
        for localizer in localizers:
            self._localizers[localizer.locale].append(localizer)

    def get_for_locale(self, locale: str, fallback: Optional[Localizer] = None) -> Localizer:
        localizers = self._localizers.get(locale)
        if localizers is None:
            LOGGER.info("Unknown locale %s", locale)
            if fallback:
                return fallback
            raise UnknownLocale(f"Unknown locale {locale}")
        return Localizer(localizers)


@attr.s
class LocalizerLoader:
    _configs: Iterable[TranslationConfig] = attr.ib()

    def load(self) -> LocalizerFactory:
        localizers: list[_GNULocalizer] = []
        domain_locales: dict[tuple, int] = defaultdict(int)
        for config in self._configs:
            try:
                gnu_localizer = gettext.translation(
                    config.domain,
                    localedir=config.path,
                    languages=[config.locale],
                )
                assert isinstance(gnu_localizer, gettext.GNUTranslations)
            except FileNotFoundError:
                raise UnknownDomain(f"Unknown domain {config.domain}")
            localizers.append(
                _GNULocalizer(
                    domain=config.domain,
                    locale=config.locale,
                    localizer=gnu_localizer,
                ),
            )
            domain_locales[config.key] += 1
        duplicated_domain_locales = ["-".join(key) for key, value in domain_locales.items() if value > 1]
        if duplicated_domain_locales:
            LOGGER.warning("Duplicated domain-locale pairs: %s", ";".join(duplicated_domain_locales))
        return LocalizerFactory(
            localizers=localizers,
        )

lib/dl_i18n/dl_i18n/localizer_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
    Dict,
    Generic,
    Type,
    TypeVar,
)

import attr
from marshmallow import (
    EXCLUDE,
    Schema,
    post_load,
)


class BaseSchema(Schema):
    class Meta:
        unknown = EXCLUDE


_TARGET_OBJECT_TV = TypeVar("_TARGET_OBJECT_TV")


class DefaultSchema(BaseSchema, Generic[_TARGET_OBJECT_TV]):
    TARGET_CLS: ClassVar[Type[_TARGET_OBJECT_TV]]

    @classmethod
    def get_target_cls(cls) -> Type[_TARGET_OBJECT_TV]:
        return cls.TARGET_CLS

    def to_object(self, data: dict) -> _TARGET_OBJECT_TV:
        return self.get_target_cls()(**data)  # type: ignore[call-arg]

    @post_load(pass_many=False)
    def post_load(self, data: Dict[str, Any], **_: Any) -> _TARGET_OBJECT_TV:
        obj = self.to_object(data)
        return obj


class DefaultValidateSchema(DefaultSchema[_TARGET_OBJECT_TV], Generic[_TARGET_OBJECT_TV]):
    @post_load(pass_many=False)
    def post_load(self, data: Dict[str, Any], **_: Any) -> _TARGET_OBJECT_TV:
        data_with_orig = data.copy()
        obj = self.to_object(data_with_orig)
        as_dict = attr.asdict(obj, recurse=False)
        as_dict = {k: v for k, v in as_dict.items() if v is not None}
        data = {k: v for k, v in data.items() if v is not None}
        assert data == as_dict
        return obj

lib/dl_model_tools/dl_model_tools/schema/base.py
datalens-tech/datalens-backend
from typing import (
    Any,
    Mapping,
    Optional,
    Type,
)

from dynamic_enum import DynamicEnum
from marshmallow import ValidationError
from marshmallow import fields as ma_fields


class DynamicEnumField(ma_fields.Field):
    def __init__(self, dyn_enum_cls: Type[DynamicEnum], **kwargs: Any):
        super().__init__(**kwargs)
        self._dyn_enum_cls = dyn_enum_cls

    def _serialize(self, value: Any, attr: Optional[str], obj: Any, **kwargs: Any) -> Optional[str]:
        if value is None and self.allow_none:
            return None

        if not isinstance(value, self._dyn_enum_cls):
            raise ValidationError(f"Invalid type {type(value).__name__} for {attr}")

        return value.value

    def _deserialize(self, value: Any, attr: Optional[str], data: Optional[Mapping[str, Any]], **kwargs: Any) -> Any:
        if value is None and self.allow_none:
            return None

        if not isinstance(value, str):
            raise ValidationError(f"Invalid type {type(value).__name__} for {attr}")

        if not self._dyn_enum_cls.is_declared(value):
            raise ValidationError(f"Invalid value {value} for {attr}")

        deserialized_value = self._dyn_enum_cls(value)
        return deserialized_value

lib/dl_model_tools/dl_model_tools/schema/dynamic_enum_field.py
datalens-tech/datalens-backend
from dynamic_enum import (
    AutoEnumValue,
    DynamicEnum,
)
from marshmallow import (
    Schema,
    ValidationError,
)
import pytest

from dl_model_tools.schema.dynamic_enum_field import DynamicEnumField


def test_dynamic_enum_field():
    class Color(DynamicEnum):
        blue = AutoEnumValue()

    class ColorSchema(Schema):
        color = DynamicEnumField(Color, attribute="color_attr", allow_none=False)
        other_color = DynamicEnumField(Color, allow_none=True)

    schema = ColorSchema()
    original_data = {"color_attr": Color.blue, "other_color": None}

    serialized_data = schema.dump(original_data)
    assert serialized_data == {"color": "blue", "other_color": None}

    deserialized_data = schema.load(serialized_data)
    assert deserialized_data == original_data

    with pytest.raises(ValidationError):
        schema.dump({"color_attr": "smth", "other_color": None})

    with pytest.raises(ValidationError):
        schema.dump({"color_attr": None, "other_color": None})

    with pytest.raises(ValidationError):
        schema.load({"color_attr": "whatever", "other_color": None})

lib/dl_model_tools/dl_model_tools_tests/unit/test_dyn_enum_field.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Union

import attr

from dl_constants.types import TBIDataValue


CellValue = Union[TBIDataValue, dict]


@attr.s(frozen=True)
class DimensionSpec:
    legend_item_id: int = attr.ib(kw_only=True)


@attr.s(frozen=True)
class DimensionValueSpec:
    legend_item_id: int = attr.ib(kw_only=True)
    value: CellValue = attr.ib(kw_only=True)

lib/dl_query_processing/dl_query_processing/base_specs/dimensions.py
datalens-tech/datalens-backend
from __future__ import annotations

from itertools import chain
import logging
from typing import (
    Callable,
    Collection,
    Dict,
    NamedTuple,
    Optional,
    Set,
    Tuple,
)

from dl_core.components.ids import (
    AvatarId,
    SourceId,
)
from dl_core.db.elements import SchemaColumn
from dl_core.utils import make_id
from dl_formula.core.datatype import DataType
import dl_formula.core.exc as formula_exc
import dl_formula.core.nodes as formula_nodes
from dl_formula.inspect.expression import used_fields
from dl_query_processing.compilation.type_mapping import BI_TO_FORMULA_TYPES


LOGGER = logging.getLogger(__name__)


ColumnId = str


class AvatarColumn(NamedTuple):
    id: ColumnId
    avatar_id: AvatarId
    column: SchemaColumn


class ColumnRegistry:
    def __init__(
        self,
        db_columns: Optional[Collection[SchemaColumn]] = None,
        avatar_source_map: Optional[Dict[AvatarId, SourceId]] = None,
    ):
        self._db_columns = list(db_columns) if db_columns is not None else []
        self._avatar_source_map = avatar_source_map.copy() if avatar_source_map is not None else {}

        self._column_ids: Dict[Tuple[AvatarId, str], ColumnId] = {}
        self._columns: Dict[ColumnId, AvatarColumn] = {}

        for avatar_id, source_id in self._avatar_source_map.items():
            self.register_avatar(avatar_id=avatar_id, source_id=source_id)

    def __contains__(self, item: str) -> bool:
        return item in self._columns

    def __add__(self, other: ColumnRegistry) -> ColumnRegistry:
        if other is self:
            return self

        if not self._db_columns and not self._avatar_source_map:
            return other
        if not other._db_columns and not other._avatar_source_map:
            return self

        result = ColumnRegistry(
            avatar_source_map={**self._avatar_source_map, **other._avatar_source_map},
        )
        for column_id, avatar_column in chain(self._columns.items(), other._columns.items()):
            result.register_column(
                column_id=column_id, db_column=avatar_column.column, avatar_id=avatar_column.avatar_id
            )
        return result

    def register_avatar(self, avatar_id: AvatarId, source_id: SourceId) -> None:
        self._avatar_source_map[avatar_id] = source_id
        for col in self._db_columns:
            if col.source_id == source_id:
                column_id = make_id()
                self._column_ids[(avatar_id, col.name)] = column_id
                self._columns[column_id] = AvatarColumn(
                    id=column_id,
                    avatar_id=avatar_id,
                    column=col,
                )

    def register_column(self, column_id: ColumnId, avatar_id: AvatarId, db_column: SchemaColumn) -> None:
        if db_column not in self._db_columns:
            self._db_columns.append(db_column)
        avatar_column = AvatarColumn(
            id=column_id,
            avatar_id=avatar_id,
            column=db_column,
        )
        self._column_ids[(avatar_id, db_column.name)] = column_id
        self._columns[column_id] = avatar_column

    def unregister_avatar(self, avatar_id: AvatarId) -> None:
        source_id = self._avatar_source_map[avatar_id]
        for col in self._db_columns:
            if col.source_id == source_id:
                column_id = self._column_ids[(avatar_id, col.name)]
                del self._columns[column_id]
                del self._column_ids[(avatar_id, col.name)]
        del self._avatar_source_map[avatar_id]

    def get_used_avatar_ids_for_formula_obj(self, formula_obj: formula_nodes.Formula) -> Set[str]:
        column_ids = {field_node.name for field_node in used_fields(formula_obj)}
        return {self._columns.get(column_id).avatar_id for column_id in column_ids}  # type: ignore  # TODO: fix

    def get_multipart_column_names(self, avatar_alias_mapper: Callable[[AvatarId], str]) -> Dict[str, Tuple[str, str]]:
        return {
            column_id: (avatar_alias_mapper(av_column.avatar_id), av_column.column.name)
            for column_id, av_column in self._columns.items()
        }

    def get_column_formula_types(self) -> Dict[str, DataType]:
        return {
            column_id: BI_TO_FORMULA_TYPES[av_column.column.user_type] for column_id, av_column in self._columns.items()
        }

    def get(self, column_id: str) -> AvatarColumn:
        return self._columns.get(column_id)  # type: ignore  # TODO: fix

    def get_avatar_column(self, avatar_id: str, name: str) -> AvatarColumn:
        try:
            column_id = self._column_ids[(avatar_id, name)]
            return self._columns[column_id]
        except KeyError:
            raise formula_exc.FormulaError(
                f"Unknown referenced source column: {name}",
                code=formula_exc.UnknownSourceColumnError.default_code,
            )

    def get_column_ids(self) -> Dict[Tuple[str, str], str]:
        return {
            (av_column.avatar_id, av_column.column.name): column_id for column_id, av_column in self._columns.items()
        }

    def get_column_avatar_ids(self) -> Dict[Tuple[str, str], str]:
        return {
            column_id: av_column.avatar_id  # type: ignore  # TODO: fix
            for column_id, av_column in self._columns.items()
        }

    def get_columns_for_avatar(self, avatar_id: str) -> list[AvatarColumn]:
        return sorted(
            [av_column for av_column in self._columns.values() if av_column.avatar_id == avatar_id],
            key=lambda col: col.id,
        )

lib/dl_query_processing/dl_query_processing/column_registry.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc

from dl_query_processing.compilation.primitives import CompiledQuery
from dl_query_processing.compilation.specs import QuerySpec


class RawQueryCompilerBase(abc.ABC):
    @abc.abstractmethod
    def make_compiled_query(self, query_spec: QuerySpec) -> CompiledQuery:
        raise NotImplementedError

lib/dl_query_processing/dl_query_processing/compilation/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Collection,
    Iterator,
    Optional,
)

import attr

from dl_core.fields import BIField


@attr.s
class FieldRegistry:
    """A collection of fields that provides easy access to them by id, title, iteration, etc."""

    _field_list: list[BIField] = attr.ib(init=False, factory=list)
    _fields_by_id: dict[str, BIField] = attr.ib(init=False, factory=dict)
    _fields_by_title: dict[str, BIField] = attr.ib(init=False, factory=dict)

    def get(self, id: Optional[str] = None, title: Optional[str] = None) -> BIField:
        if id is not None and title is not None:
            raise ValueError("Cannot specify both id and title")

        if id is not None:
            return self._fields_by_id[id]

        if title is not None:
            return self._fields_by_title[title]

        raise ValueError("Either id or title is required")

    def add(self, field: BIField) -> None:
        self._field_list.append(field)
        self._fields_by_id[field.guid] = field
        self._fields_by_title[field.title] = field

    def remove(self, field: BIField) -> None:
        self._fields_by_title.pop(field.title, None)
        self._fields_by_id.pop(field.guid, None)
        self._field_list.remove(field)

    @property
    def titles_to_guids(self) -> dict[str, str]:
        return {title: field.guid for title, field in self._fields_by_title.items()}

    @property
    def titles(self) -> Collection[str]:
        return set(self._fields_by_title.keys())

    @property
    def ids(self) -> Collection[str]:
        return set(self._fields_by_id.keys())

    def __iter__(self) -> Iterator[BIField]:
        return iter(self._field_list)

lib/dl_query_processing/dl_query_processing/compilation/field_registry.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import logging
import os
from typing import (
    Any,
    Callable,
    ClassVar,
    Dict,
    Optional,
    TypeVar,
)

import attr

from dl_constants.enums import WhereClauseOperation
from dl_core.fields import BIField
from dl_formula.core.datatype import DataType
import dl_formula.core.nodes as formula_nodes
from dl_formula.shortcuts import n
from dl_query_processing.compilation.formula_compiler import (
    FORMULA_TO_BI_TYPES,
    FormulaCompiler,
)
from dl_query_processing.compilation.helpers import (
    ARRAY_TYPES,
    TREE_TYPES,
    make_literal_node,
)
from dl_query_processing.compilation.primitives import CompiledFormulaInfo
from dl_query_processing.compilation.specs import (
    FilterFieldSpec,
    FilterSourceColumnSpec,
)
import dl_query_processing.exc
from dl_query_processing.utils.datetime import parse_datetime


LOGGER = logging.getLogger(__name__)


_FILTER_PARAMS_TV = TypeVar("_FILTER_PARAMS_TV", bound="FilterParams")

CONTAINMENT_OPS = {
    WhereClauseOperation.STARTSWITH,
    WhereClauseOperation.ISTARTSWITH,
    WhereClauseOperation.CONTAINS,
    WhereClauseOperation.ICONTAINS,
    WhereClauseOperation.NOTCONTAINS,
    WhereClauseOperation.NOTICONTAINS,
    WhereClauseOperation.ENDSWITH,
    WhereClauseOperation.IENDSWITH,
}
ARRAY_LEN_OPS = {
    WhereClauseOperation.LENEQ,
    WhereClauseOperation.LENNE,
    WhereClauseOperation.LENGT,
    WhereClauseOperation.LENGTE,
    WhereClauseOperation.LENLT,
    WhereClauseOperation.LENLTE,
}


@attr.s(frozen=True)
class FilterDefinition:
    arg_cnt: Optional[int] = attr.ib(kw_only=True)
    callable: Callable[..., formula_nodes.FormulaItem] = attr.ib(kw_only=True)


@attr.s(auto_attribs=True)
class FilterParams:
    """
    Single-field filter state,
    primarily for mangling in subclasses.
    """

    field: BIField
    operation: WhereClauseOperation
    filter_args: list = None  # type: ignore  # TODO: fix
    data_type: DataType = None  # type: ignore  # TODO: fix
    field_cast_type: DataType = None  # type: ignore  # TODO: fix
    arg_cast_type: DataType = None  # type: ignore  # TODO: fix

    def clone(self: _FILTER_PARAMS_TV, **updates: Any) -> _FILTER_PARAMS_TV:
        """Convenience method so that callers don't need to know about `attr`"""
        return attr.evolve(self, **updates)


@attr.s
class FilterFormulaCompiler:
    _formula_compiler: FormulaCompiler = attr.ib(kw_only=True)

    # structure: {WhereClauseOperation: (add_arg_cnt, expr_callable)}
    FILTER_OPERATIONS: ClassVar[Dict[WhereClauseOperation, FilterDefinition]] = {
        # unary (arg_cnt = 0)
        WhereClauseOperation.ISNULL: FilterDefinition(arg_cnt=0, callable=n.func.ISNULL),
        WhereClauseOperation.ISNOTNULL: FilterDefinition(arg_cnt=0, callable=lambda f, *args: n.not_(n.func.ISNULL(f))),
        # binary (arg_cnt = 1)
        WhereClauseOperation.EQ: FilterDefinition(arg_cnt=1, callable=lambda f, val: n.binary("==", f, val)),
        WhereClauseOperation.NE: FilterDefinition(arg_cnt=1, callable=lambda f, val: n.binary("!=", f, val)),
        WhereClauseOperation.GT: FilterDefinition(arg_cnt=1, callable=lambda f, val: n.binary(">", f, val)),
        WhereClauseOperation.GTE: FilterDefinition(arg_cnt=1, callable=lambda f, val: n.binary(">=", f, val)),
        WhereClauseOperation.LT: FilterDefinition(arg_cnt=1, callable=lambda f, val: n.binary("<", f, val)),
        WhereClauseOperation.LTE: FilterDefinition(arg_cnt=1, callable=lambda f, val: n.binary("<=", f, val)),
        WhereClauseOperation.STARTSWITH: FilterDefinition(arg_cnt=1, callable=n.func.STARTSWITH),
        WhereClauseOperation.ISTARTSWITH: FilterDefinition(arg_cnt=1, callable=n.func.ISTARTSWITH),
        WhereClauseOperation.ENDSWITH: FilterDefinition(arg_cnt=1, callable=n.func.ENDSWITH),
        WhereClauseOperation.IENDSWITH: FilterDefinition(arg_cnt=1, callable=n.func.IENDSWITH),
        WhereClauseOperation.CONTAINS: FilterDefinition(arg_cnt=1, callable=n.func.CONTAINS),
        WhereClauseOperation.ICONTAINS: FilterDefinition(arg_cnt=1, callable=n.func.ICONTAINS),
        WhereClauseOperation.NOTCONTAINS: FilterDefinition(arg_cnt=1, callable=n.func.NOTCONTAINS),
        WhereClauseOperation.NOTICONTAINS: FilterDefinition(
            arg_cnt=1, callable=lambda f, val: n.not_(n.func.ICONTAINS(f, val))
        ),
        WhereClauseOperation.LENEQ: FilterDefinition(
            arg_cnt=1, callable=lambda f, val: n.binary("==", n.func.LEN(f), val)
        ),
        WhereClauseOperation.LENNE: FilterDefinition(
            arg_cnt=1, callable=lambda f, val: n.binary("!=", n.func.LEN(f), val)
        ),
        WhereClauseOperation.LENGT: FilterDefinition(
            arg_cnt=1, callable=lambda f, val: n.binary(">", n.func.LEN(f), val)
        ),
        WhereClauseOperation.LENGTE: FilterDefinition(
            arg_cnt=1, callable=lambda f, val: n.binary(">=", n.func.LEN(f), val)
        ),
        WhereClauseOperation.LENLT: FilterDefinition(
            arg_cnt=1, callable=lambda f, val: n.binary("<", n.func.LEN(f), val)
        ),
        WhereClauseOperation.LENLTE: FilterDefinition(
            arg_cnt=1, callable=lambda f, val: n.binary("<=", n.func.LEN(f), val)
        ),
        # binary with list (arg_cnt = None)
        # None for arg_cnt means that all args are converted to a list,
        # which is used as a single argument for binary operation
        WhereClauseOperation.IN: FilterDefinition(arg_cnt=None, callable=lambda f, val: n.binary("in", f, val)),
        WhereClauseOperation.NIN: FilterDefinition(arg_cnt=None, callable=lambda f, val: n.binary("notin", f, val)),
        # ternary (arg_cnt = 2)
        WhereClauseOperation.BETWEEN: FilterDefinition(
            arg_cnt=2, callable=lambda f, second, third: n.ternary("between", f, second, third)
        ),
    }

    def _custom_filter_cast(self, filter_params: FilterParams) -> FilterParams:
        """
        Customize the parameters for a field filtering, e.g. change data type
        to which the field and filter arg values need to be cast.
        """
        return filter_params

    def compile_abstract_field_filter_formula(
        self,
        field: BIField,
        operation: WhereClauseOperation,
        filter_args: Optional[list] = None,
        original_field_id: Optional[str] = None,
        anonymous: bool = False,
    ) -> CompiledFormulaInfo:
        field_formula_obj = self._formula_compiler.compile_field_formula(field, collect_errors=False).formula_obj

        data_type = self._formula_compiler.get_field_final_formula_data_type(field=field)
        assert data_type is not None
        LOGGER.info(f"Filtered field {field.title!r} has data type {data_type.name} and is a {field.type.name}")

        filter_params = FilterParams(
            field=field,
            operation=operation,
            filter_args=filter_args,  # type: ignore  # TODO: fix
            data_type=data_type,
            # defaults:
            field_cast_type=data_type,
            arg_cast_type=data_type,
        )
        mangled_filter_params = self._custom_filter_cast(filter_params)
        assert mangled_filter_params.field == field, "not meant to be changed for now"
        operation = mangled_filter_params.operation
        filter_args = mangled_filter_params.filter_args
        assert mangled_filter_params.data_type == data_type, "not meant to be changed for now"
        field_cast_type = mangled_filter_params.field_cast_type
        arg_cast_type = mangled_filter_params.arg_cast_type

        LOGGER.info(
            f"Will cast field {field.title!r} to {field_cast_type.name} " f"and {filter_args!r} to {arg_cast_type.name}"
        )
        if field_cast_type != data_type:
            field_formula_obj = self._formula_compiler.apply_cast_to_formula(
                formula_obj=field_formula_obj,
                current_dtype=data_type,
                cast=FORMULA_TO_BI_TYPES[field_cast_type],
            )

        def _remove_timezone(arg: Any) -> Any:
            if isinstance(arg, str):
                return arg.removesuffix("Z")
            if isinstance(arg, datetime.datetime):
                return arg.replace(tzinfo=None)
            return arg

        # Determine and, if necessary, convert data types of the filter arguments
        args_nodes = []
        remove_timezone = data_type in (DataType.GENERICDATETIME, DataType.CONST_GENERICDATETIME)
        for arg in filter_args or ():
            if remove_timezone:
                arg = _remove_timezone(arg)
            # we assume that filter arguments all have the same data type as the filtered field
            try:
                arg_lit_node = make_literal_node(val=arg, data_type=arg_cast_type)
            except dl_query_processing.exc.InvalidLiteralError as e:
                raise dl_query_processing.exc.FilterValueError(
                    f"Invalid filter value {arg!r} for type {arg_cast_type.name}"
                ) from e
            args_nodes.append(arg_lit_node)

        # Translate the whole expression
        # 1. Get filter translation info
        try:
            filter_def = self.FILTER_OPERATIONS[operation]
            add_arg_cnt = filter_def.arg_cnt
            expr_callable = filter_def.callable
        except KeyError:
            raise ValueError(operation)

        # 2. Make list of args for function call
        if add_arg_cnt is None:  # for list lookup operations (IN, NOT IN)
            args = [formula_nodes.ExpressionList.make(*args_nodes)]
        else:
            args = args_nodes[:add_arg_cnt]  # type: ignore  # TODO: fix

        # 3. Create formula object
        formula_obj = formula_nodes.Formula.make(expr=expr_callable(field_formula_obj.expr, *args))

        original_field_id = original_field_id or field.guid
        if anonymous:
            original_field_id = None

        formula_info = CompiledFormulaInfo(
            formula_obj=formula_obj,
            avatar_ids=self._formula_compiler._columns.get_used_avatar_ids_for_formula_obj(formula_obj),  # FIXME
            original_field_id=original_field_id,
            alias=None,
        )
        return formula_info

    def compile_filter_formula(
        self,
        filter_spec: FilterFieldSpec,
        original_field_id: Optional[str] = None,
    ) -> CompiledFormulaInfo:
        """
        Prepare a filter expression for given ``field_id`` using given filter ``operation``.
        The result can be used in a ``WHERE`` or ``HAVING`` clause of a ``SELECT`` statement.

        :param field_id: ID of field to create filter for
        :param operation: one of the filter operations/functions defined in ``WhereClauseOperation``
        :param filter_args: additional arguments for the filter function
        :return:
        """
        field = self._formula_compiler._fields.get(id=filter_spec.field_id)  # FIXME: access to protected member
        return self.compile_abstract_field_filter_formula(
            field=field,
            operation=filter_spec.operation,
            filter_args=filter_spec.values,
            original_field_id=original_field_id,
            anonymous=filter_spec.anonymous,
        )

    def compile_source_column_filter_formula(
        self,
        source_column_filter_spec: FilterSourceColumnSpec,
        original_field_id: Optional[str] = None,
    ) -> CompiledFormulaInfo:
        """
        Prepare a formula based on a source column filter spec
        """
        field = self._formula_compiler.make_direct_field(
            avatar_id=source_column_filter_spec.avatar_id,
            source=source_column_filter_spec.column_name,
        )
        return self.compile_abstract_field_filter_formula(
            field=field,
            operation=source_column_filter_spec.operation,
            filter_args=source_column_filter_spec.values,
            original_field_id=original_field_id,
        )


class MainFilterFormulaCompiler(FilterFormulaCompiler):
    def _mangle_date_filter(self, filter_params: FilterParams) -> FilterParams:
        """
        Generally, date filters receive datetime arguments like '2017-01-01T00:00:00';

        casting them to date might change equality
        (`2017-01-01 < 2017-01-01T00:00:01`; but, `2017-01-01 == date(2017-01-01T00:00:01)`)

        Easy fix is to cast the *field* to datetime
        (`2017-01-01T00:00:00 < 2017-01-01T00:00:01`);

        However, this hurts database performance in some cases

        so, if the filter arg dt is at midnight, cast it to date,
        without cast-wrapping the field itself.
        """

        if filter_params.data_type != DataType.DATE or not filter_params.filter_args:
            return filter_params

        try:
            parsed_args = [parse_datetime(value) for value in filter_params.filter_args]
        except ValueError:
            # not a valid value, cannot check midnight-ness, and it should probably fail later anyway.
            parsed_args = []
            is_midnight = [False for _ in filter_params.filter_args]
        else:
            # Precision note: False for '2017-01-01T00:00:00.000001', True for '2017-01-01T00:00:00.0000001'
            # Timezone note: intentionally ignoring the tzinfo, i.e.
            # '2017-01-01T00:00:00' and '2017-01-01T00:00:00+03:00' are treated the
            # same for date (and naive-datetime) filtering.
            is_midnight = [
                arg.hour == 0 and arg.minute == 0 and arg.second == 0 and arg.microsecond == 0 for arg in parsed_args
            ]

        filter_as_date = False

        if all(is_midnight):
            # `date_col <cmp> 2020-01-01T00:00:00` is equivalent to
            # `date_col <cmp> 2020-01-01`
            filter_as_date = True
        elif filter_params.operation == WhereClauseOperation.GT and len(parsed_args) == 1:
            # `date_col > 2020-01-01T01:00:00` is equivalent to
            # `date_col > 2020-01-01`
            filter_as_date = True
        elif filter_params.operation == WhereClauseOperation.GTE and len(parsed_args) == 1:
            # `date_col >= 2020-01-01T01:00:00` is equivalent to
            # `date_col >  2020-01-01`
            filter_as_date = True
            filter_params = filter_params.clone(operation=WhereClauseOperation.GT)
        elif filter_params.operation == WhereClauseOperation.LTE and len(parsed_args) == 1:
            # `date_col <= 2020-01-01T01:00:00` is equivalent to
            # `date_col <= 2020-01-01`
            filter_as_date = True
        elif filter_params.operation == WhereClauseOperation.LT and len(parsed_args) == 1:
            # `date_col <  2020-01-01T01:00:00` is equivalent to
            # `date_col <= 2020-01-01`
            filter_as_date = True
            filter_params = filter_params.clone(operation=WhereClauseOperation.LTE)
        elif filter_params.operation == WhereClauseOperation.BETWEEN and len(parsed_args) == 2:
            # equivalent to `date_col >= arg_from and date_col <= arg_to`,
            # but prefer to keep it as `between`
            arg_from, arg_to = parsed_args
            midnight_from, _ = is_midnight
            if not midnight_from:
                # `date_col between 2020-01-01T01:00:00 and ...` equivalent to `date_col between 2020-01-02 and ...`
                arg_from += datetime.timedelta(days=1)
            parsed_args = [arg_from, arg_to]
            filter_as_date = True

        if filter_as_date:
            new_args = [arg.date().isoformat() for arg in parsed_args]
            filter_params = filter_params.clone(
                filter_args=new_args,
                arg_cast_type=DataType.DATE,  # actually unchanged, just making it explicit.
            )
        else:
            # Cast both sides to datetime.
            filter_params = filter_params.clone(
                field_cast_type=DataType.DATETIME,
                arg_cast_type=DataType.DATETIME,
            )
        if filter_params.operation in CONTAINMENT_OPS:
            # Cast both sides to string.
            filter_params = filter_params.clone(
                field_cast_type=DataType.STRING,
                arg_cast_type=DataType.STRING,
            )
        return filter_params

    def _mangle_containment_filter(self, filter_params: FilterParams) -> FilterParams:
        if filter_params.operation in CONTAINMENT_OPS:
            if filter_params.data_type in ARRAY_TYPES or filter_params.data_type in TREE_TYPES:
                if filter_params.operation in (WhereClauseOperation.CONTAINS, WhereClauseOperation.NOTCONTAINS):
                    arg_cast_type = {
                        DataType.ARRAY_FLOAT: DataType.FLOAT,
                        DataType.ARRAY_INT: DataType.INTEGER,
                        DataType.ARRAY_STR: DataType.STRING,
                        DataType.TREE_STR: DataType.STRING,
                    }[filter_params.data_type]
                    filter_params = filter_params.clone(arg_cast_type=arg_cast_type)
                else:
                    pass  # Fall back to default behavior - cast to the same type as the field
            else:
                filter_params = filter_params.clone(arg_cast_type=DataType.STRING)
        return filter_params

    def _mangle_array_filter(self, filter_params: FilterParams) -> FilterParams:
        if filter_params.operation in ARRAY_LEN_OPS:
            filter_params = filter_params.clone(arg_cast_type=DataType.INTEGER)
        return filter_params

    def _custom_filter_cast(self, filter_params: FilterParams) -> FilterParams:
        filter_params = self._mangle_containment_filter(filter_params)
        filter_params = self._mangle_array_filter(filter_params)
        filter_params = self._mangle_date_filter(filter_params)
        return filter_params

lib/dl_query_processing/dl_query_processing/compilation/filter_compiler.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from contextlib import contextmanager
from functools import (
    partial,
    wraps,
)
from itertools import chain
import logging
from typing import (
    AbstractSet,
    Callable,
    Collection,
    Dict,
    Generator,
    Iterable,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
    Union,
    cast,
)
import uuid

from dl_constants.enums import (
    AggregationFunction,
    BinaryJoinOperator,
    CalcMode,
    ConditionPartCalcMode,
    FieldType,
    ManagedBy,
    OrderDirection,
    UserDataType,
)
from dl_core.components.ids import (
    AvatarId,
    FieldId,
    make_field_id,
)
from dl_core.fields import (
    BIField,
    DirectCalculationSpec,
    FormulaCalculationSpec,
)
from dl_core.multisource import (
    AvatarRelation,
    BinaryCondition,
    ConditionPart,
    ConditionPartDirect,
    ConditionPartFormula,
    ConditionPartResultField,
)
from dl_formula.collections import NodeValueMap
import dl_formula.core.aux_nodes as formula_aux_nodes
from dl_formula.core.datatype import DataType
import dl_formula.core.exc as formula_exc
from dl_formula.core.index import NodeHierarchyIndex
from dl_formula.core.message_ctx import (
    FormulaErrorCtx,
    MessageLevel,
)
import dl_formula.core.nodes as formula_nodes
from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.inspect.expression import (
    enumerate_fields,
    infer_data_type,
    is_aggregate_expression,
    is_aggregated_above_sub_node,
    is_window_expression,
    used_fields,
    used_func_calls,
)
from dl_formula.mutation.bfb import RemapBfbMutation
from dl_formula.mutation.general import (
    ConvertBlocksToFunctionsMutation,
    IgnoreParenthesisWrapperMutation,
)
from dl_formula.mutation.lookup import LookupDefaultBfbMutation
from dl_formula.mutation.mutation import (
    FormulaMutation,
    apply_mutations,
)
from dl_formula.mutation.window import (
    AmongToWithinGroupingMutation,
    DefaultWindowOrderingMutation,
    IgnoreExtraWithinGroupingMutation,
)
from dl_formula.parser.base import FormulaParser
from dl_formula.validation.aggregation import AggregationChecker
from dl_formula.validation.env import ValidationEnvironment
from dl_formula.validation.validator import (
    Checker,
    validate,
)
from dl_formula.validation.window import WindowFunctionChecker
from dl_query_processing.column_registry import ColumnRegistry
from dl_query_processing.compilation.field_registry import FieldRegistry
from dl_query_processing.compilation.helpers import make_literal_node
from dl_query_processing.compilation.primitives import (
    CompiledFormulaInfo,
    CompiledJoinOnFormulaInfo,
)
from dl_query_processing.compilation.specs import (
    OrderByFieldSpec,
    ParameterValueSpec,
    SelectWrapperSpec,
)
from dl_query_processing.compilation.type_mapping import (
    BI_TO_FORMULA_TYPES,
    DEFAULT_DATA_TYPE,
    FORMULA_TO_BI_TYPES,
)
from dl_query_processing.compilation.wrapper_applicator import ExpressionWrapperApplicator
from dl_query_processing.enums import ProcessingStage
import dl_query_processing.exc
from dl_utils.utils import enum_not_none


LOGGER = logging.getLogger(__name__)


class FieldProcessingStageManager:
    """
    Provides handling and storage of:
    - field errors
    - processing results
    - expression data type
    from different processing stages
    """

    default_type_on_error = DEFAULT_DATA_TYPE
    save_type_stages = {
        ProcessingStage.substitution,
        ProcessingStage.aggregation,
        ProcessingStage.final,
    }

    def __init__(self, columns: ColumnRegistry, inspect_env: InspectionEnvironment):
        self._columns = columns
        self._inspect_env = inspect_env

        self._errors: Dict[str, Dict[ProcessingStage, List[FormulaErrorCtx]]] = defaultdict(lambda: defaultdict(list))
        self._exprs: Dict[str, Dict[ProcessingStage, Optional[formula_nodes.Formula]]] = defaultdict(
            lambda: defaultdict(lambda: None)
        )
        self._data_types: Dict[str, Dict[ProcessingStage, DataType]] = defaultdict(
            lambda: defaultdict(lambda: DEFAULT_DATA_TYPE)
        )

    def _get_formula_obj_data_type(self, formula_obj: formula_nodes.Formula) -> DataType:
        try:
            try:
                return infer_data_type(
                    node=formula_obj,
                    field_types=self._columns.get_column_formula_types(),
                    env=self._inspect_env,
                )
            except KeyError:
                raise formula_exc.DataTypeError("Unknown type for column")
        except formula_exc.FormulaError:
            return self.default_type_on_error

    def add_errors(self, *errors: FormulaErrorCtx, field: BIField, stage: ProcessingStage) -> None:
        registered_errors: List[FormulaErrorCtx] = self._errors[field.guid][stage]
        for error in errors:
            if error not in registered_errors:
                registered_errors.append(error)

    def get_errors(
        self,
        field: BIField,
        stage: Union[ProcessingStage, Iterable[ProcessingStage], None] = None,
    ) -> List[FormulaErrorCtx]:
        """Return errors from cache for the given field (and type)."""

        field_errors = self._errors[field.guid]
        if stage is not None and isinstance(stage, ProcessingStage):
            return field_errors[stage]

        return list(
            chain.from_iterable(
                [
                    errors
                    for st, errors in sorted(field_errors.items(), key=lambda el: el[0].name)
                    if stage is None or st in stage
                ]
            )
        )

    def set_result(
        self, formula_obj: Optional[formula_nodes.Formula], *, field: BIField, stage: ProcessingStage
    ) -> None:
        self._exprs[field.guid][stage] = formula_obj
        if stage in self.save_type_stages:
            self._data_types[field.guid][stage] = self._get_formula_obj_data_type(formula_obj)  # type: ignore  # TODO: fix

    def get_result(self, field: BIField, stage: ProcessingStage) -> Optional[formula_nodes.Formula]:
        return self._exprs[field.guid][stage]

    def get_data_type(self, field: BIField, stage: ProcessingStage) -> DataType:
        return self._data_types[field.guid][stage]

    def get_user_type(self, field: BIField, stage: ProcessingStage) -> UserDataType:
        return FORMULA_TO_BI_TYPES[self.get_data_type(field=field, stage=stage)]

    def raise_if_any(self, field: BIField, stage: ProcessingStage) -> None:
        """Raise ``FormulaHandlingError`` if there are any errors in cache for the given field (and type)."""

        errors = self.get_errors(field=field, stage=stage)
        if errors:
            raise dl_query_processing.exc.FormulaHandlingError(*errors, stage=stage, field=field)

    def clear(self, field: BIField, stage: Optional[ProcessingStage] = None) -> None:
        """Clear error cache for the given field (and type)."""

        if stage is not None:
            self._errors[field.guid][stage].clear()
            self._exprs[field.guid][stage].clear()  # type: ignore  # TODO: fix
            self._data_types[field.guid][stage].clear()  # type: ignore  # TODO: fix
        else:
            self._errors[field.guid].clear()
            self._exprs[field.guid].clear()
            self._data_types[field.guid].clear()

    @contextmanager
    def handle(
        self,
        field: BIField,
        stage: ProcessingStage,
        collect: bool = True,
        raise_: bool = True,
    ) -> Generator[None, None, None]:
        try:
            yield
        except formula_exc.FormulaError as err:
            if collect:
                # Save error for later use
                self.add_errors(*err.errors, field=field, stage=stage)
            if raise_:
                pass_on_errors = err.errors
                if collect:
                    # Do not add error messages to the re-raised exception
                    # to avoid their duplication in the error storage
                    # when the error is caught in the next-level wraper
                    pass_on_errors = []
                if pass_on_errors:
                    LOGGER.error(
                        f"Formula handling errors found on stage {stage.name} "
                        f"for field {field.title!r} ({field.guid}): "
                        f"{[str(e) for e in pass_on_errors]}"
                    )
                raise dl_query_processing.exc.FormulaHandlingError(*pass_on_errors, stage=stage, field=field) from err


StageProcType = Callable[["FormulaCompiler", BIField, bool], formula_nodes.Formula]


def implements_stage(stage: ProcessingStage) -> Callable[[StageProcType], StageProcType]:
    """
    A parameterized decorator.

    Mark a ``FormulaCompiler`` method as an implementation of a certain formula processing stage.
    This adds the enabling of various caches and error handlers needed to optimize
    formula generation and validation.
    """

    def decorator(func: StageProcType) -> StageProcType:
        @wraps(func)
        def wrapper(self: "FormulaCompiler", field: BIField, collect_errors: bool = False) -> formula_nodes.Formula:
            # First check the error cache
            if not collect_errors:
                # Error collection is disabled, so we can just raise the first error we encounter
                self._stage_manager.raise_if_any(field=field, stage=stage)

            # Errors were either disabled or not found in cache,
            # so move on to the result cache
            formula_obj = self._stage_manager.get_result(field=field, stage=stage)
            if formula_obj is None:
                # No cached result was found, so generate it,
                # add it (and any errors encountered on the way) to cache
                # (all the cache stuff is done inside ``self._stage_manager.handle(...)``)
                with self._stage_manager.handle(field=field, stage=stage, collect=collect_errors):
                    try:
                        formula_obj = func(self, field, collect_errors)
                    except Exception:
                        LOGGER.debug(f"Stage {stage.name} failed for field {field.title!r} ({field.guid})")
                        raise
                assert formula_obj is not None
                self._stage_manager.set_result(formula_obj, field=field, stage=stage)
            return formula_obj

        return wrapper

    return decorator


def _unsupported_cast(typename):  # type: ignore  # TODO: fix
    return f"__UNSUPPORTED_CAST_TO_{typename.upper()}__"


_SUPPORTED_CASTS_FUNCTIONS = {
    UserDataType.boolean: "bool",
    UserDataType.date: "date",
    UserDataType.datetime: "datetime",
    UserDataType.genericdatetime: "genericdatetime",
    UserDataType.float: "float",
    UserDataType.integer: "int",
    UserDataType.geopoint: "geopoint",
    UserDataType.geopolygon: "geopolygon",
    UserDataType.string: "str",
    UserDataType.markup: "markup",
}
_CAST_FUNCTIONS = {
    bi_type: _SUPPORTED_CASTS_FUNCTIONS.get(bi_type) or _unsupported_cast(bi_type.name) for bi_type in UserDataType
}
_ALLOWED_PARAMETER_TYPES = {
    UserDataType.string,
    UserDataType.integer,
    UserDataType.float,
    UserDataType.boolean,
    UserDataType.date,
    UserDataType.datetime,
    UserDataType.genericdatetime,
}


class FormulaCompiler:
    _agg_functions = {
        AggregationFunction.countunique: "countd",
        AggregationFunction.count: "count",
        AggregationFunction.sum: "sum",
        AggregationFunction.avg: "avg",
        AggregationFunction.max: "max",
        AggregationFunction.min: "min",
    }
    _join_condition_operators = {
        BinaryJoinOperator.eq: "_==",
        BinaryJoinOperator.ne: "_!=",
        BinaryJoinOperator.gt: ">",
        BinaryJoinOperator.gte: ">=",
        BinaryJoinOperator.lt: "<",
        BinaryJoinOperator.lte: "<=",
    }

    def __init__(
        self,
        *,
        all_fields: Iterable[BIField],
        columns: ColumnRegistry,
        formula_parser: FormulaParser,
        group_by_ids: Collection[str] = (),
        filter_ids: Collection[str] = (),
        order_by_specs: Sequence[OrderByFieldSpec] = (),
        mock_among_dimensions: bool = False,  # mock dimensions for AMONG clauses (in validation mode)
        inspect_env: Optional[InspectionEnvironment] = None,
        suppress_double_aggregations: bool = True,
        allow_nested_window_functions: bool = False,
        parameter_value_specs: Sequence[ParameterValueSpec] = (),
        field_wrappers: Optional[Dict[str, SelectWrapperSpec]] = None,
        validate_aggregations: bool = True,
    ):
        self._fields = FieldRegistry()
        self._columns = columns
        self._formula_parser = formula_parser
        self._mock_among_dimensions = mock_among_dimensions
        self._suppress_double_aggregations = suppress_double_aggregations
        self._allow_nested_window_functions = allow_nested_window_functions
        self._field_wrappers = cast(Dict[str, SelectWrapperSpec], field_wrappers or {})
        self._validate_aggregations = validate_aggregations

        self._inspect_env = inspect_env if inspect_env is not None else InspectionEnvironment()
        self._valid_env = ValidationEnvironment()
        self._wrapper_applicator = ExpressionWrapperApplicator()

        for field in all_fields:
            self.register_field(field)

        self._group_by_ids = set(group_by_ids or ())
        self._filter_ids = set(filter_ids or ())
        self._order_by_specs = list(order_by_specs or ())
        self._parameter_value_specs = list(parameter_value_specs or ())

        # initialize caches
        # node/expression caches
        self._formula_parsed_cache: Dict[str, formula_nodes.Formula] = {}
        self._substituted_formula_fields: NodeValueMap[formula_nodes.FormulaItem] = NodeValueMap()
        # error caches
        self._stage_manager = FieldProcessingStageManager(columns=self._columns, inspect_env=self._inspect_env)
        self._formula_error_cache: Dict[str, List[FormulaErrorCtx]] = defaultdict(list)  # FIXME: ???
        # other caches
        self._field_types: Dict[str, FieldType] = {}
        self._field_dependencies: Dict[str, Set[str]] = defaultdict(set)  # fields this field depends on
        self._field_inverse_dependencies: Dict[str, Set[str]] = defaultdict(set)  # fields that depend on this one

        self.update_environments(group_by_ids=self._group_by_ids, order_by_specs=self._order_by_specs)

    @property
    def columns(self) -> ColumnRegistry:
        return self._columns

    @property
    def inspect_env(self) -> InspectionEnvironment:
        return self._inspect_env

    def register_avatar(self, avatar_id: str, source_id: str) -> None:
        self._columns.register_avatar(avatar_id=avatar_id, source_id=source_id)

    def unregister_avatar(self, avatar_id: str) -> None:
        # Clear all cached values that depend on this avatar
        for field in self._fields:
            if field.calc_mode == CalcMode.direct and field.avatar_id == avatar_id:
                self.uncache_field(field=field)
        self._columns.unregister_avatar(avatar_id=avatar_id)

    def uncache_field(self, field: BIField, visited_guids: AbstractSet[str] = frozenset()) -> None:
        """Clear all field-related caches and caches of dependent fields too"""

        if field.guid in visited_guids:
            # recursion detected
            return
        visited_guids |= frozenset([field.guid])

        # must be fetched before the cleanup starts
        dep_field_guids = self._get_dependent_fields(field).copy()

        # clear own caches
        self._stage_manager.clear(field=field)
        this_field_dependencies = self._field_dependencies.pop(field.guid, set())
        self._field_types.pop(field.guid, None)

        # clear own dependencies
        for child_field_id in this_field_dependencies:
            self._field_inverse_dependencies.pop(child_field_id, None)

        # clear cache for dependent fields
        for dep_field_guid in dep_field_guids:
            dependent_field = self._fields.get(id=dep_field_guid)
            self.uncache_field(dependent_field, visited_guids=visited_guids)

    def unregister_field(self, field: BIField) -> None:
        """Remove field from all collections, mappings and caches"""

        # `field` may have been updated, so get the registered version
        field = self._fields.get(id=field.guid)
        if field is None:
            return

        self.uncache_field(field)  # this uncaches dependencies also

        # clear direct usages
        self._fields.remove(field)
        self._group_by_ids.discard(field.guid)
        found_order_bys = [spec for spec in self._order_by_specs if spec.field_id == field.guid]
        for found_spec in found_order_bys:
            self._order_by_specs.remove(found_spec)

    def register_field(self, field: BIField) -> None:
        """Register field in translation environment"""
        self._fields.add(field)

        if field.calc_mode == CalcMode.direct:
            avatar_id = field.avatar_id
            assert avatar_id is not None
            try:
                av_column = self._columns.get_avatar_column(avatar_id=avatar_id, name=field.source)
            except formula_exc.FormulaError:
                LOGGER.warning(f"Unknown column {field.source} for avatar {avatar_id}")
            else:
                if av_column.column.has_auto_aggregation:
                    # Better to check this property in the column, than in the field
                    # because the field may have been updated, and not all of its attributes are synchronized yet.
                    self._inspect_env.cache_is_aggregate_expr.add(
                        formula_nodes.Field.make(name=av_column.id), value=True
                    )

    def update_field(self, field: BIField) -> None:
        is_in_group_by: bool = False
        if field.guid in self._group_by_ids:
            is_in_group_by = True
        found_order_bys = [  # order is crucial here, so save index along with each entry
            (ind, spec) for ind, spec in enumerate(self._order_by_specs) if spec.field_id == field.guid
        ]

        self.unregister_field(field)
        self.register_field(field)

        # add it back to group by
        if is_in_group_by:
            self._group_by_ids.add(field.guid)
        # put it back in the same positions in order by
        for ind, spec in found_order_bys:
            self._order_by_specs.insert(ind, spec)

    def update_environments(
        self,
        group_by_ids: Collection[str],
        order_by_specs: Sequence[OrderByFieldSpec],
    ) -> None:
        self._group_by_ids = set(group_by_ids)
        self._order_by_specs = list(order_by_specs)

    def get_expression_hash(self, expr: str) -> str:
        expr = expr.strip()
        # TODO: replace with formula_obj.extract
        expr = BIField.rename_in_formula(formula=expr, key_map=self._fields.titles_to_guids)
        return expr

    def _try_parse_formula(
        self, field: BIField, collect_errors: bool = False
    ) -> Tuple[Optional[formula_nodes.Formula], List[FormulaErrorCtx]]:
        """Attempt to parse formula. If `collect_errors`"""

        formula = field.formula
        formula_obj = self._formula_parsed_cache.get(formula)
        errors = self._formula_error_cache.get(formula, [])

        if formula_obj is None and not errors:  # has not been parsed yet
            try:
                formula_obj = self._formula_parser.parse(formula)
            except formula_exc.ParseError as err:
                if not collect_errors:
                    raise dl_query_processing.exc.FormulaHandlingError(*err.errors)
                errors.extend(err.errors)

            self._formula_parsed_cache[formula] = formula_obj  # type: ignore  # TODO: fix
            self._formula_error_cache[formula] = errors

        return formula_obj, errors

    def _is_field_recursive(self, field: BIField) -> bool:
        """Check whether field is recursive or not"""

        def _has_recursion(for_guid: str, visited_guids: AbstractSet[str] = frozenset()) -> bool:
            if for_guid in visited_guids:
                return True
            visited_guids |= frozenset([for_guid])
            for child in self._field_dependencies[for_guid]:
                if _has_recursion(child, visited_guids):
                    return True
            return False

        return _has_recursion(field.guid)

    def _make_dependencies_for_field(
        self,
        *,
        field: BIField,
        formula_obj: formula_nodes.Formula,
        collect_errors: bool = False,
    ) -> None:
        """Build dependency graph"""

        if field.calc_mode != CalcMode.formula:
            return

        if formula_obj is None:
            return

        child_fields = []
        for child_field_node in used_fields(formula_obj):
            if child_field_node.name not in self._fields.titles:
                continue

            child_field = self._fields.get(title=child_field_node.name)
            # register dependencies
            self._field_dependencies[field.guid].add(child_field.guid)
            self._field_inverse_dependencies[child_field.guid].add(field.guid)
            child_fields.append(child_field)

        if self._is_field_recursive(field):
            raise dl_query_processing.exc.FormulaHandlingError(
                FormulaErrorCtx("Recursion detected in field", level=MessageLevel.ERROR),
                field=field,
            )

        # generate dependencies of children
        for child_field in child_fields:
            self._generate_base_formula_obj_for_field(field=child_field, collect_errors=collect_errors)

    def get_referenced_fields(self, field: BIField) -> Set[FieldId]:
        return self._field_dependencies[field.guid]

    def _substitute_fields_in_formula(
        self, field: BIField, formula_obj: formula_nodes.Formula, collect_errors: bool = False
    ) -> formula_nodes.Formula:
        to_substitute: Dict[NodeHierarchyIndex, formula_nodes.FormulaItem] = {}

        sub_node: formula_nodes.FormulaItem
        for field_node_idx, child_field_node in enumerate_fields(formula_obj):
            if child_field_node.name not in self._fields.titles:
                # Not recognized as a field
                if child_field_node.name in self._columns:
                    # A DB column -> already substituted field
                    continue
                sub_node = formula_aux_nodes.ErrorNode.make(
                    err_code=formula_exc.UnknownFieldInFormulaError.default_code,
                    message=f"Unknown field found in formula: {child_field_node.name}",
                )

            else:
                # Recognized as an unsubstituted field, so substitute it
                child_field = self._fields.get(title=child_field_node.name)

                # Check for and handle double aggregations
                if (
                    self._suppress_double_aggregations
                    # Child has explicit aggregation
                    and child_field.aggregation != AggregationFunction.none
                    # And current field has formula (auto) aggregation
                    and is_aggregated_above_sub_node(node=formula_obj, index=field_node_idx)
                ):
                    # Double aggregation detected,
                    # so generate sub-node only up to the aggregation stage
                    # (`casting` is the previous one)
                    sub_node = self._process_field_stage_casting(child_field, collect_errors=collect_errors).expr  # type: ignore  # TODO: fix

                else:
                    # No ignorable double aggregations, proceed as usual
                    sub_node = self._process_field_stage_aggregation(
                        child_field, collect_errors=collect_errors
                    ).expr  # type: ignore  # TODO: fix

                child_wrapper = self._field_wrappers.get(child_field.guid)
                if child_wrapper is not None:
                    sub_node = self._wrapper_applicator.apply_wrapper(sub_node, wrapper=child_wrapper)

            self._substituted_formula_fields.add(sub_node, value=child_field_node)

            to_substitute[field_node_idx] = sub_node

        formula_obj = formula_obj.substitute_batch(to_substitute)
        return formula_obj

    def _generate_formula_obj_for_formula_field(
        self, field: BIField, collect_errors: bool = False
    ) -> formula_nodes.Formula:
        """Generate expression for field in formula mode"""
        errors: List[FormulaErrorCtx] = []

        formula_obj, parse_errors = self._try_parse_formula(field=field, collect_errors=collect_errors)
        errors.extend(parse_errors)

        if formula_obj is None:
            assert errors, "Got no formula and no errors. Should have either of the two"
            raise dl_query_processing.exc.FormulaHandlingError(*errors, field=field)

        return formula_obj

    def _generate_formula_obj_for_direct_field(self, field: BIField) -> formula_nodes.Formula:
        """Generate expression for field with direct type"""
        sub_node: formula_nodes.FormulaItem
        avatar_id = field.avatar_id
        assert avatar_id is not None
        try:
            av_column = self._columns.get_avatar_column(avatar_id=avatar_id, name=field.source)
        except formula_exc.FormulaError as err:
            error_info = err.errors[0]
            sub_node = formula_aux_nodes.ErrorNode.make(err_code=error_info.code, message=error_info.message)
        else:
            sub_node = formula_nodes.Field.make(name=av_column.id)
        return formula_nodes.Formula.make(expr=sub_node)

    def _generate_formula_obj_for_parameter_field(self, field: BIField) -> formula_nodes.Formula:
        """Generate expression for field with parameter type"""
        if field.cast not in _ALLOWED_PARAMETER_TYPES:
            raise dl_query_processing.exc.ParameterUnsupportedTypeError(
                f"Unsupported type {field.cast} for parameter {field.title}"
            )

        value = None
        try:
            value = next(spec.value for spec in self._parameter_value_specs if spec.field_id == field.guid)
        except StopIteration:
            if field.default_value is not None:
                value = field.default_value.value
        if value is None:
            raise dl_query_processing.exc.ParameterValueError(
                f"No value can be found for parameter field {field.title}"
            )

        value_constraint = field.value_constraint
        if value_constraint is not None and not value_constraint.is_valid(value):
            raise dl_query_processing.exc.ParameterValueError(
                f"Invalid parameter value {value!r} for constraint {value_constraint}"
            )

        try:
            return formula_nodes.Formula.make(
                expr=make_literal_node(val=value, data_type=BI_TO_FORMULA_TYPES[enum_not_none(field.cast)])
            )
        except dl_query_processing.exc.InvalidLiteralError as e:
            raise dl_query_processing.exc.ParameterValueError(
                f"Invalid parameter value {value!r} for type {field.cast}"
            ) from e

    def _generate_base_formula_obj_for_field(
        self, field: BIField, collect_errors: bool = False
    ) -> formula_nodes.Formula:
        """Generate raw (without casts, aggregations, substitutions, etc.) formula object for field."""

        if field.calc_mode == CalcMode.direct:
            formula_obj = self._generate_formula_obj_for_direct_field(field=field)
        elif field.calc_mode == CalcMode.parameter:
            formula_obj = self._generate_formula_obj_for_parameter_field(field=field)
        else:
            formula_obj = self._generate_formula_obj_for_formula_field(field=field, collect_errors=collect_errors)

        return formula_obj

    def _get_mock_dimensions_for_among(self, formula_obj: formula_nodes.Formula) -> List[formula_nodes.FormulaItem]:
        """
        Collect dimensions from the AMONG clauses of the formula expression.
        """

        func_list = used_func_calls(formula_obj)
        required_dims: List[formula_nodes.FormulaItem] = []
        for func in func_list:
            if isinstance(func, formula_nodes.WindowFuncCall):
                if isinstance(func.grouping, formula_nodes.WindowGroupingAmong):
                    required_dims.extend(func.grouping.dim_list)

        return required_dims

    def apply_pre_sub_mutations(
        self,
        field: BIField,
        formula_obj: formula_nodes.Formula,
        collect_errors: bool = False,
    ) -> formula_nodes.Formula:
        """
        Apply pre-substitution mutations required for window functions to be translated correctly.
        """

        # prepare mutations
        mutations: List[FormulaMutation] = []

        # prepare default ordering (for patching RSUM, MSUM functions and the like)
        default_order_by = []
        ob_expr_obj: formula_nodes.FormulaItem
        for order_by_spec in self._order_by_specs:
            ob_expr_obj = formula_nodes.Field.make(name=self._fields.get(id=order_by_spec.field_id).title)
            if order_by_spec.direction == OrderDirection.desc:
                ob_expr_obj = formula_nodes.OrderDescending.make(expr=ob_expr_obj)
            default_order_by.append(ob_expr_obj)

        # Only measures can contain BFB clauses
        title_id_map = {f.title: f.guid for f in self._fields}
        mutations.extend(
            [
                IgnoreParenthesisWrapperMutation(),
                ConvertBlocksToFunctionsMutation(),
                DefaultWindowOrderingMutation(default_order_by=default_order_by),
                LookupDefaultBfbMutation(),
                RemapBfbMutation(name_mapping=title_id_map),
            ]
        )

        formula_obj = apply_mutations(formula_obj, mutations=mutations)
        return formula_obj

    def _apply_function_by_name(self, formula_obj: formula_nodes.Formula, func_name: str) -> formula_nodes.Formula:
        """Apply a single-argument function to given expression"""

        return formula_nodes.Formula.make(expr=formula_nodes.FuncCall.make(name=func_name, args=[formula_obj.expr]))

    def apply_mutations(
        self,
        field: BIField,
        formula_obj: formula_nodes.Formula,
        collect_errors: bool = False,
    ) -> formula_nodes.Formula:
        """
        Apply the mutations required for window functions to be translated correctly.
        """

        # prepare mutations
        mutation_lists: List[List[FormulaMutation]] = [[]]

        if self._field_types[field.guid] != FieldType.DIMENSION:
            # prepare global dimensions (for patching AMONG clauses)
            if self._mock_among_dimensions:
                # in mock mode real dimensions are unavailable, so we just use the ones listed in AMONG clauses
                global_dimensions = self._get_mock_dimensions_for_among(formula_obj)
            else:
                global_dimensions = [
                    # We will need to process some other fields to a more advanced stage,
                    # but since they are all dimensions, there will be no recursion loop
                    self._process_field_stage_aggregation(  # type: ignore  # TODO: fix
                        self._fields.get(id=dim_id), collect_errors=collect_errors
                    ).expr
                    for dim_id in sorted(self._group_by_ids)
                ]

            if is_window_expression(node=formula_obj, env=self._inspect_env):
                mutation_lists[-1].extend(
                    [
                        AmongToWithinGroupingMutation(global_dimensions=global_dimensions),
                        IgnoreExtraWithinGroupingMutation(
                            global_dimensions=global_dimensions, inspect_env=self._inspect_env
                        ),
                    ]
                )

        # Apply mutation lists in the specified order
        for mutations in mutation_lists:
            formula_obj = apply_mutations(formula_obj, mutations=mutations)

        return formula_obj

    def _apply_aggregation(
        self, formula_obj: formula_nodes.Formula, aggregation: AggregationFunction
    ) -> formula_nodes.Formula:
        """Apply an explicit aggregation to given expression"""

        if aggregation != AggregationFunction.none:
            formula_obj = self._apply_function_by_name(formula_obj, func_name=self._agg_functions[aggregation])

        return formula_obj

    def _apply_cast(
        self,
        formula_obj: formula_nodes.Formula,
        current_dtype: DataType,
        cast: Optional[UserDataType],
    ) -> formula_nodes.Formula:
        """Apply a type cast to given expression"""

        if current_dtype != DataType.NULL:
            expr_type = FORMULA_TO_BI_TYPES[current_dtype]
            if cast is not None and cast != expr_type:
                formula_obj = self._apply_function_by_name(formula_obj, func_name=_CAST_FUNCTIONS[cast])
        return formula_obj

    def apply_cast_to_formula(
        self,
        formula_obj: formula_nodes.Formula,
        current_dtype: DataType,
        cast: Optional[UserDataType],
    ) -> formula_nodes.Formula:
        return self._apply_cast(formula_obj=formula_obj, current_dtype=current_dtype, cast=cast)

    def _validate_field_formula(
        self,
        formula_obj: formula_nodes.Formula,
        field_id: FieldId,
        collect_errors: bool = False,
    ) -> None:
        unselected_dimension_ids = {
            f.guid for f in self._fields if f.type == FieldType.DIMENSION and f.guid not in self._group_by_ids
        }
        checkers: List[Checker] = []
        if self._field_types[field_id] != FieldType.DIMENSION and self._validate_aggregations:
            global_dimensions = []
            for dim_field_id in self._group_by_ids:
                dim_field = self._fields.get(id=dim_field_id)
                try:
                    dim_formula_obj = self._process_field_stage_aggregation(
                        dim_field, collect_errors=collect_errors
                    )  # type: ignore  # TODO: fix
                except formula_exc.FormulaError:
                    # error has been registered, so when the field will be rendered explicitly the error will be raised
                    # here we can just skip it since there's nothing to register as a dimension
                    continue
                global_dimensions.append(dim_formula_obj.expr)

            checkers.append(
                AggregationChecker(
                    allow_nested_agg=True,
                    allow_inconsistent_inside_agg=False,
                    valid_env=self._valid_env,
                    inspect_env=self._inspect_env,
                    global_dimensions=global_dimensions,
                )
            )
        checkers.append(
            WindowFunctionChecker(
                inspect_env=self._inspect_env,
                allow_nested=self._allow_nested_window_functions,
                filter_ids=self._filter_ids,
                unselected_dimension_ids=unselected_dimension_ids,
            )
        )
        validate(node=formula_obj, env=self._valid_env, collect_errors=collect_errors, checkers=checkers)

    @implements_stage(ProcessingStage.base)
    def _process_field_stage_base(self, field: BIField, collect_errors: bool = False) -> formula_nodes.Formula:
        return self._generate_base_formula_obj_for_field(field=field, collect_errors=collect_errors)

    @implements_stage(ProcessingStage.pre_sub_mutation)
    def _process_field_stage_pre_sub_mutation(
        self, field: BIField, collect_errors: bool = False
    ) -> formula_nodes.Formula:
        formula_obj = self._process_field_stage_base(field, collect_errors=collect_errors)  # type: ignore  # TODO: fix
        return self.apply_pre_sub_mutations(field=field, formula_obj=formula_obj, collect_errors=collect_errors)

    @implements_stage(ProcessingStage.dep_generation)
    def _process_field_stage_dep_generation(
        self, field: BIField, collect_errors: bool = False
    ) -> formula_nodes.Formula:
        formula_obj = self._process_field_stage_pre_sub_mutation(field, collect_errors=collect_errors)  # type: ignore  # TODO: fix
        self._make_dependencies_for_field(field=field, formula_obj=formula_obj)
        return formula_obj

    @implements_stage(ProcessingStage.substitution)
    def _process_field_stage_substitution(self, field: BIField, collect_errors: bool = False) -> formula_nodes.Formula:
        formula_obj = self._process_field_stage_dep_generation(field, collect_errors=collect_errors)  # type: ignore  # TODO: fix
        return self._substitute_fields_in_formula(field=field, formula_obj=formula_obj, collect_errors=collect_errors)

    @implements_stage(ProcessingStage.casting)
    def _process_field_stage_casting(self, field: BIField, collect_errors: bool = False) -> formula_nodes.Formula:
        formula_obj = self._process_field_stage_substitution(field, collect_errors=collect_errors)  # type: ignore  # TODO: fix
        return self._apply_cast(
            formula_obj=formula_obj,
            current_dtype=self._stage_manager.get_data_type(field=field, stage=ProcessingStage.substitution),
            cast=field.cast,
        )

    @implements_stage(ProcessingStage.aggregation)
    def _process_field_stage_aggregation(self, field: BIField, collect_errors: bool = False) -> formula_nodes.Formula:
        formula_obj = self._process_field_stage_casting(field, collect_errors=collect_errors)  # type: ignore  # TODO: fix
        formula_obj = self._apply_aggregation(formula_obj=formula_obj, aggregation=field.aggregation)
        is_agg = is_aggregate_expression(formula_obj, env=self._inspect_env)
        self._field_types[field.guid] = FieldType.MEASURE if is_agg else FieldType.DIMENSION
        return formula_obj

    @implements_stage(ProcessingStage.mutation)
    def _process_field_stage_mutation(self, field: BIField, collect_errors: bool = False) -> formula_nodes.Formula:
        formula_obj = self._process_field_stage_aggregation(field, collect_errors=collect_errors)  # type: ignore  # TODO: fix
        return self.apply_mutations(field=field, formula_obj=formula_obj, collect_errors=collect_errors)

    @implements_stage(ProcessingStage.validation)
    def _process_field_stage_validation(self, field: BIField, collect_errors: bool = False) -> formula_nodes.Formula:
        formula_obj = self._process_field_stage_mutation(field, collect_errors=collect_errors)  # type: ignore  # TODO: fix
        self._validate_field_formula(formula_obj=formula_obj, field_id=field.guid, collect_errors=collect_errors)
        return formula_obj

    @implements_stage(ProcessingStage.final)
    def _process_field_stage_final(self, field: BIField, collect_errors: bool = False) -> formula_nodes.Formula:
        return self._process_field_stage_validation(field, collect_errors=collect_errors)  # type: ignore  # TODO: fix

    def _compile_field_formula(self, field: BIField, collect_errors: bool = False) -> formula_nodes.Formula:
        return self._process_field_stage_final(field=field, collect_errors=collect_errors)  # type: ignore  # TODO: fix

    @contextmanager
    def handle_formula_error(self, field_id: FieldId) -> Generator[None, None, None]:
        try:
            yield
        except formula_exc.FormulaError as err:
            field = self._fields.get(id=field_id)
            raise dl_query_processing.exc.DLFormulaError(field=field, formula_errors=err.errors)

    def _require_field_formula_preparation(self, field: BIField) -> None:
        """Make sure that the field's fully prepared and ready for translation formula is in the cache"""
        try:
            self._compile_field_formula(field, collect_errors=True)
        except formula_exc.FormulaError:
            pass  # ignore errors

    def get_field_validity(self, field: BIField) -> bool:
        """Return boolean flag indicating whether the field is valid."""
        return not self.get_field_errors(field)

    def get_field_initial_data_type(self, field: BIField) -> Optional[UserDataType]:
        """Return automatically determined data type of given field before cast and aggregation"""
        self._require_field_formula_preparation(field)
        return self._stage_manager.get_user_type(field=field, stage=ProcessingStage.substitution)

    def get_field_final_data_type(self, field: BIField) -> Optional[UserDataType]:
        """
        Return automatically determined user data type (``UserDataType``)
        of given field after cast and aggregation
        """
        self._require_field_formula_preparation(field)
        return self._stage_manager.get_user_type(field=field, stage=ProcessingStage.aggregation)

    def get_field_final_formula_data_type(self, field: BIField) -> Optional[DataType]:
        """
        Return automatically determined formula data type (``DataType``)
        of given field after cast and aggregation
        """
        self._require_field_formula_preparation(field)
        return self._stage_manager.get_data_type(field=field, stage=ProcessingStage.aggregation)

    def get_field_type(self, field: BIField) -> Optional[UserDataType]:
        """Return automatically determined field type"""
        self._require_field_formula_preparation(field)
        return self._field_types.get(field.guid, FieldType.DIMENSION)  # type: ignore  # TODO: fix

    def get_field_errors(self, field: BIField) -> Optional[List[FormulaErrorCtx]]:
        """Return list of errors found for given field"""
        self._require_field_formula_preparation(field)
        return self._stage_manager.get_errors(field)

    def _get_dependent_fields(self, field: BIField) -> Set[str]:
        """Return GUIDs of dependent fields"""
        return (self._field_inverse_dependencies.get(field.guid) or set()).copy()

    def get_node_text(self, node: formula_nodes.FormulaItem) -> str:
        text = node.original_text
        if text is None:
            original_node = self._substituted_formula_fields.get(node)
            if original_node:
                text = original_node.original_text
        if text is None:
            raise ValueError(f"Not able to get node text for node {node}")
        return text

    def _compile_binary_relation_condition_formula(
        self, relation: AvatarRelation, condition: BinaryCondition, collect_errors: bool = False
    ) -> formula_nodes.Binary:
        """
        Create a ``FormulaItem`` (to be used as part of a ``Formula``)
        from description of join condition
        """

        def _condition_part_as_field(avatar_id: str, part: ConditionPart) -> BIField:
            """
            Create an auxiliary ``BIField`` object from join condition part
            so that it can be used in ``self._generate_field_expression``
            """
            new_field = True
            if part.calc_mode == ConditionPartCalcMode.result_field:
                assert isinstance(part, ConditionPartResultField)
                field = self._fields.get(id=part.field_id)
                new_field = False
            elif part.calc_mode == ConditionPartCalcMode.direct:
                assert isinstance(part, ConditionPartDirect)
                field = self.make_direct_field(avatar_id=avatar_id, source=part.source)
            elif part.calc_mode == ConditionPartCalcMode.formula:
                # TODO: Remove
                assert isinstance(part, ConditionPartFormula)
                field = self.make_formula_field(formula=part.formula)
            else:
                raise ValueError(f"Invalid part calc_mode {part.calc_mode}")

            # JOINs over aggregated fields are not supported yet
            if new_field and field.has_auto_aggregation or not new_field and field.type == FieldType.MEASURE:
                raise dl_query_processing.exc.DatasetError("Joining over aggregated expressions is not supported")

            return field

        left, right = [
            self._compile_field_formula(
                _condition_part_as_field(avatar_id=avatar_id, part=part), collect_errors=collect_errors
            )
            for (avatar_id, part) in (
                (relation.left_avatar_id, condition.left_part),
                (relation.right_avatar_id, condition.right_part),
            )
        ]
        return formula_nodes.Binary.make(
            name=self._join_condition_operators[condition.operator],
            left=left.expr,
            right=right.expr,
        )

    def compile_relation_formula(
        self,
        relation: AvatarRelation,
        collect_errors: bool = False,
    ) -> CompiledJoinOnFormulaInfo:
        if len(relation.conditions) < 1:
            raise dl_query_processing.exc.FormulaHandlingError(
                FormulaErrorCtx("Relation requires at least one condition", level=MessageLevel.ERROR),
            )
        compile_condition = partial(
            self._compile_binary_relation_condition_formula, relation=relation, collect_errors=collect_errors
        )
        result = compile_condition(condition=relation.conditions[0])
        for condition in relation.conditions[1:]:
            result = formula_nodes.Binary.make(name="and", left=result, right=compile_condition(condition=condition))
        formula_obj = formula_nodes.Formula.make(expr=result)

        left_avatar_id: Optional[AvatarId] = relation.left_avatar_id
        if relation.managed_by == ManagedBy.feature:
            left_avatar_id = None

        formula_info = CompiledJoinOnFormulaInfo(
            formula_obj=formula_obj,
            avatar_ids=self._columns.get_used_avatar_ids_for_formula_obj(formula_obj),
            original_field_id=None,
            left_id=left_avatar_id,  # type: ignore  # TODO: fix
            right_id=relation.right_avatar_id,
            join_type=relation.join_type,
            alias=None,
        )
        return formula_info

    def compile_field_formula(  # type: ignore  # TODO: fix
        self,
        field: BIField,
        collect_errors: bool = False,
        original_field_id: Optional[str] = None,
    ) -> CompiledFormulaInfo:
        formula_obj = self._compile_field_formula(field=field, collect_errors=collect_errors)
        formula_info = CompiledFormulaInfo(
            formula_obj=formula_obj,
            avatar_ids=self._columns.get_used_avatar_ids_for_formula_obj(formula_obj),
            alias=field.guid,
            original_field_id=original_field_id or field.guid,
        )
        return formula_info

    def make_formula_field(self, formula: str) -> BIField:
        field = BIField.make(
            guid=make_field_id(),
            title=str(uuid.uuid4()),
            managed_by=ManagedBy.compiler_runtime,
            calc_spec=FormulaCalculationSpec(
                formula=formula,
            ),
        )
        self.register_field(field)
        return field

    def make_direct_field(self, avatar_id: AvatarId, source: str) -> BIField:
        column = self._columns.get_avatar_column(avatar_id=avatar_id, name=source).column
        field = BIField.make(
            guid=make_field_id(),
            title=str(uuid.uuid4()),
            type=FieldType.MEASURE if column.has_auto_aggregation else FieldType.DIMENSION,
            has_auto_aggregation=column.has_auto_aggregation,
            lock_aggregation=column.lock_aggregation,
            cast=column.user_type,
            initial_data_type=column.user_type,
            data_type=column.user_type,
            managed_by=ManagedBy.compiler_runtime,
            calc_spec=DirectCalculationSpec(
                avatar_id=avatar_id,
                source=source,
            ),
        )
        self.register_field(field)
        return field

    def compile_text_formula(self, formula: str, collect_errors: bool = False) -> CompiledFormulaInfo:
        field = self.make_formula_field(formula=formula)
        formula_obj = self._compile_field_formula(field=field, collect_errors=collect_errors)
        formula_info = CompiledFormulaInfo(
            formula_obj=formula_obj,
            avatar_ids=self._columns.get_used_avatar_ids_for_formula_obj(formula_obj),
            alias=field.guid,
            original_field_id=None,
        )
        return formula_info

    def get_formula_errors(self, formula: str) -> List[FormulaErrorCtx]:
        field = self.make_formula_field(formula=formula)
        errors = self.get_field_errors(field)
        return errors  # type: ignore  # TODO: fix

    def field_has_auto_aggregation(self, field: BIField) -> bool:
        try:
            formula_obj = self._process_field_stage_casting(field=field, collect_errors=False)  # type: ignore  # TODO: fix
            return is_aggregate_expression(formula_obj, env=self._inspect_env)
        except formula_exc.FormulaError:
            return False

lib/dl_query_processing/dl_query_processing/compilation/formula_compiler.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import json
import logging
from typing import (
    Any,
    Collection,
    Optional,
)

from dl_core.components.accessor import DatasetComponentAccessor
from dl_formula.core.datatype import DataType
import dl_formula.core.nodes as formula_nodes
import dl_formula.shortcuts
from dl_query_processing.column_registry import ColumnRegistry
from dl_query_processing.compilation.primitives import (
    BASE_QUERY_ID,
    AvatarFromObject,
    CompiledFormulaInfo,
    CompiledQuery,
    FromColumn,
    FromObject,
    JoinedFromObject,
)
from dl_query_processing.enums import ExecutionLevel
import dl_query_processing.exc


LOGGER = logging.getLogger(__name__)


def make_joined_from_for_avatars(
    used_avatar_ids: Collection[str],
    ds_accessor: DatasetComponentAccessor,
    column_reg: ColumnRegistry,
    root_avatar_id: Optional[str] = None,
) -> JoinedFromObject:
    sorted_avatar_ids = sorted(used_avatar_ids)
    if root_avatar_id is None and sorted_avatar_ids:
        root_avatar_id = sorted_avatar_ids[0]

    froms: list[FromObject] = []
    added_avatar_ids: set[str] = set()
    for avatar_id in sorted_avatar_ids:
        assert avatar_id not in added_avatar_ids, f"Avatars should not be duplicated, but got {avatar_id} twice"
        avatar = ds_accessor.get_avatar_strict(avatar_id=avatar_id)
        columns = tuple(
            FromColumn(id=column.id, name=column.column.name)
            for column in column_reg.get_columns_for_avatar(avatar_id=avatar_id)
        )
        froms.append(
            AvatarFromObject(
                id=avatar_id,
                alias=avatar_id,
                avatar_id=avatar_id,
                source_id=avatar.source_id,
                columns=columns,
            )
        )
        added_avatar_ids.add(avatar_id)

    return JoinedFromObject(root_from_id=root_avatar_id, froms=froms)


def single_formula_comp_query_for_validation(
    formula: CompiledFormulaInfo,
    ds_accessor: DatasetComponentAccessor,
    column_reg: ColumnRegistry,
) -> CompiledQuery:
    return CompiledQuery(
        id=BASE_QUERY_ID,
        level_type=ExecutionLevel.source_db,
        select=[formula],
        group_by=[],
        order_by=[],
        filters=[],
        join_on=[],
        joined_from=make_joined_from_for_avatars(
            used_avatar_ids=formula.avatar_ids,
            ds_accessor=ds_accessor,
            column_reg=column_reg,
        ),
        limit=None,
        offset=None,
    )


def _datetime_fromisoformat(val: str) -> datetime.datetime:
    val = val.replace(" ", "T")
    if val.endswith("Z"):
        val = val[:-1] + "+00:00"
    return datetime.datetime.fromisoformat(val)


ARRAY_TYPES = (
    DataType.ARRAY_STR,
    DataType.CONST_ARRAY_STR,
    DataType.ARRAY_INT,
    DataType.CONST_ARRAY_INT,
    DataType.ARRAY_FLOAT,
    DataType.CONST_ARRAY_FLOAT,
)
TREE_TYPES = (
    DataType.TREE_STR,
    DataType.CONST_TREE_STR,
)


def make_literal_node(val: Any, data_type: DataType) -> formula_nodes.BaseLiteral:
    """
    Make a ``formula_nodes.Literal`` from given value ``val`` converted to given type ``data_type``.
    For use in filters and parameter fields
    """
    node: Optional[formula_nodes.BaseLiteral] = None
    try:
        # strings can contain any of the types, so handle them separately
        if isinstance(val, str):
            if data_type in (DataType.DATE, DataType.CONST_DATE):
                dt_val = _datetime_fromisoformat(val).replace(tzinfo=None)
                if dt_val.hour or dt_val.minute or dt_val.second or dt_val.microsecond:
                    LOGGER.warning("Truncating datetime with nonzero time to date: %s", val)
                node = formula_nodes.LiteralDate.make(dt_val.date())
            elif data_type in (DataType.DATETIME, DataType.CONST_DATETIME):
                # NOTE: the value might have non-empty tzinfo.
                node = formula_nodes.LiteralDatetime.make(_datetime_fromisoformat(val))
            elif data_type in (DataType.DATETIMETZ, DataType.CONST_DATETIMETZ):
                val_dt = _datetime_fromisoformat(val)
                # Incoming offset-less datetimes are interpreted as UTC
                if val_dt.tzinfo is None:
                    val_dt = val_dt.replace(tzinfo=datetime.timezone.utc)
                node = formula_nodes.LiteralDatetimeTZ.make(val_dt)
            elif data_type in (DataType.GENERICDATETIME, DataType.CONST_GENERICDATETIME):
                node = formula_nodes.LiteralGenericDatetime.make(_datetime_fromisoformat(val))
            elif data_type in (DataType.INTEGER, DataType.CONST_INTEGER):
                node = formula_nodes.LiteralInteger.make(int(val))
            elif data_type in (DataType.FLOAT, DataType.CONST_FLOAT):
                node = formula_nodes.LiteralFloat.make(float(val))
            elif data_type in (DataType.BOOLEAN, DataType.CONST_BOOLEAN):
                bool_val = {"true": True, "false": False}.get(val.lower())
                if bool_val is None:
                    raise ValueError("Invalid value for bool")
                node = formula_nodes.LiteralBoolean.make(bool_val)
            elif data_type in (DataType.STRING, DataType.CONST_STRING):
                node = formula_nodes.LiteralString.make(val)
            elif data_type in (DataType.GEOPOINT, DataType.CONST_GEOPOINT):
                node = formula_nodes.LiteralGeopoint.make(val)
            elif data_type in (DataType.GEOPOLYGON, DataType.CONST_GEOPOLYGON):
                node = formula_nodes.LiteralGeopolygon.make(val)
            elif data_type in (DataType.UUID, DataType.CONST_UUID):
                node = formula_nodes.LiteralUuid.make(val)
            elif data_type in ARRAY_TYPES or data_type in TREE_TYPES:
                try:
                    val = json.loads(val)
                except json.decoder.JSONDecodeError:
                    raise ValueError("Invalid value for array")
            # No known use-cases: DataType.MARKUP
            else:
                raise ValueError("Unexpected data_type value")

        if data_type in ARRAY_TYPES or data_type in TREE_TYPES:
            if not isinstance(val, list):
                raise ValueError("Got non-list for array")
            if data_type in (DataType.ARRAY_STR, DataType.CONST_ARRAY_STR):
                node = formula_nodes.LiteralArrayString.make(val)
            elif data_type in (DataType.ARRAY_INT, DataType.CONST_ARRAY_INT):
                node = formula_nodes.LiteralArrayInteger.make(val)
            elif data_type in (DataType.ARRAY_FLOAT, DataType.CONST_ARRAY_FLOAT):
                node = formula_nodes.LiteralArrayFloat.make(val)
            elif data_type in (DataType.TREE_STR, DataType.CONST_TREE_STR):
                node = formula_nodes.LiteralTreeString.make(val)

        if node is None:
            node = dl_formula.shortcuts.n.lit(val)  # guess literal type

    except (ValueError, TypeError):
        raise dl_query_processing.exc.InvalidLiteralError(f"Invalid literal value {val!r} for type {data_type.name}")

    assert node is not None
    return node

lib/dl_query_processing/dl_query_processing/compilation/helpers.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    AbstractSet,
    Any,
    ClassVar,
    Generator,
    Iterable,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
    TypeVar,
)

import attr

from dl_constants.enums import (
    JoinType,
    OrderDirection,
)
import dl_formula.core.nodes as formula_nodes
from dl_query_processing.compilation.query_meta import QueryMetaInfo
from dl_query_processing.enums import (
    ExecutionLevel,
    QueryPart,
)


_COMPILED_FLA_TV = TypeVar("_COMPILED_FLA_TV", bound="CompiledFormulaInfo")


@attr.s(slots=True, frozen=True)
class CompiledFormulaInfo:
    show_names: ClassVar[Tuple[str, ...]] = (
        "formula_obj",
        "alias",
        "avatar_ids",
        "original_field_id",
    )

    formula_obj: formula_nodes.Formula = attr.ib()
    alias: Optional[str] = attr.ib()
    avatar_ids: Set[str] = attr.ib(factory=set)
    original_field_id: Optional[str] = attr.ib(default=None)

    @property
    def not_none_alias(self) -> str:
        assert self.alias is not None
        return self.alias

    @property
    def complexity(self) -> int:
        assert self.formula_obj.complexity is not None
        return self.formula_obj.complexity

    def pretty(
        self,
        indent: str = "  ",
        initial_indent: str = "",
        do_leading_indent: bool = False,
    ) -> str:
        content_item_strs: List[str] = []
        item_indent = f"{initial_indent}{indent}"
        for item_name in self.show_names:
            item = getattr(self, item_name)
            if isinstance(item, formula_nodes.Formula):
                item_str = item.pretty(
                    indent=indent,
                    initial_indent=item_indent,
                    do_leading_indent=False,
                )
            else:
                item_str = str(item)

            content_item_strs.append(f"{item_name}={item_str}")

        content_str = "".join([f"{item_indent}{item_str},\n" for item_str in content_item_strs])
        text = "{lead_idt}{cls}(\n{cont}\n{ini_idt})".format(
            lead_idt=initial_indent if do_leading_indent else "",
            cls=self.__class__.__name__,
            cont=content_str,
            ini_idt=initial_indent,
        )
        return text

    def clone(self: _COMPILED_FLA_TV, **updates: Any) -> _COMPILED_FLA_TV:
        return attr.evolve(self, **updates)


@attr.s(slots=True, frozen=True)
class CompiledOrderByFormulaInfo(CompiledFormulaInfo):  # noqa
    show_names = CompiledFormulaInfo.show_names + ("direction",)

    direction: OrderDirection = attr.ib(kw_only=True)


@attr.s(slots=True, frozen=True)
class CompiledJoinOnFormulaInfo(CompiledFormulaInfo):  # noqa
    show_names = CompiledFormulaInfo.show_names + (
        "left_id",
        "right_id",
        "join_type",
    )

    left_id: str = attr.ib(kw_only=True)  # is root for feature-managed relations
    right_id: str = attr.ib(kw_only=True)
    join_type: JoinType = attr.ib(kw_only=True)


_FROM_OBJ_TV = TypeVar("_FROM_OBJ_TV", bound="FromObject")


@attr.s(frozen=True)
class FromColumn:
    id: str = attr.ib(kw_only=True)
    name: str = attr.ib(kw_only=True)

    def clone(self, **updates: Any) -> FromColumn:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class FromObject:
    id: str = attr.ib(kw_only=True)
    alias: str = attr.ib(kw_only=True)
    columns: tuple[FromColumn, ...] = attr.ib(kw_only=True)

    def clone(self: _FROM_OBJ_TV, **updates: Any) -> _FROM_OBJ_TV:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class JoinedFromObject:
    root_from_id: Optional[str] = attr.ib(kw_only=True, default=None)
    froms: Sequence[FromObject] = attr.ib(kw_only=True, default=())

    def iter_ids(self) -> Iterable[str]:
        return (from_obj.id for from_obj in self.froms)

    def clone(self, **updates: Any) -> JoinedFromObject:
        return attr.evolve(self, **updates)


# Use double letter as the base query ID to avoid collisions
# with column aliases that use single letters
BASE_QUERY_ID = "qq"


@attr.s(slots=True, frozen=True)
class CompiledQuery:
    show_names = (
        "id",
        "select",
        "group_by",
        "filters",
        "order_by",
        "join_on",
        "limit",
        "offset",
    )

    id: str = attr.ib(kw_only=True)
    level_type: ExecutionLevel = attr.ib(kw_only=True)

    select: List[CompiledFormulaInfo] = attr.ib(kw_only=True, factory=list)
    group_by: List[CompiledFormulaInfo] = attr.ib(kw_only=True, factory=list)
    filters: List[CompiledFormulaInfo] = attr.ib(kw_only=True, factory=list)
    order_by: List[CompiledOrderByFormulaInfo] = attr.ib(kw_only=True, factory=list)
    join_on: List[CompiledJoinOnFormulaInfo] = attr.ib(kw_only=True, factory=list)
    joined_from: JoinedFromObject = attr.ib(kw_only=True, factory=JoinedFromObject)
    limit: Optional[int] = attr.ib(kw_only=True, default=None)
    offset: Optional[int] = attr.ib(kw_only=True, default=None)
    meta: QueryMetaInfo = attr.ib(kw_only=True, factory=QueryMetaInfo)

    def get_complexity(self) -> int:
        return sum(formula.complexity for formula in self.all_formulas)

    def is_empty(self) -> bool:
        return not self.select

    def pretty(
        self,
        indent: str = "  ",
        initial_indent: str = "",
        do_leading_indent: bool = False,
    ) -> str:
        content_item_strs: List[str] = []
        item_indent = f"{initial_indent}{indent}"
        subitem_indent = f"{item_indent}{indent}"
        for item_name in self.show_names:
            item = getattr(self, item_name)
            if isinstance(item, list):
                if len(item) == 0:
                    item_str = "[]"
                else:
                    item_str = "[\n"
                    for subitem in item:
                        assert isinstance(subitem, CompiledFormulaInfo)
                        subitem_str = subitem.pretty(
                            indent=indent,
                            initial_indent=subitem_indent,
                            do_leading_indent=True,
                        )
                        item_str += f"{subitem_str},\n"

                    item_str += f"{item_indent}]"
            else:
                item_str = str(item)

            content_item_strs.append(f"{item_name}={item_str}")

        content_str = "".join([f"{item_indent}{item_str},\n" for item_str in content_item_strs])
        text = "{lead_idt}{cls}(\n{cont}\n{ini_idt})".format(
            lead_idt=initial_indent if do_leading_indent else "",
            cls=self.__class__.__name__,
            cont=content_str,
            ini_idt=initial_indent,
        )
        return text

    @property
    def all_formulas(self) -> Generator[CompiledFormulaInfo, None, None]:
        yield from self.select
        yield from self.group_by
        yield from self.filters
        yield from self.order_by
        yield from self.join_on

    def get_formula_list(self, query_part: QueryPart) -> Sequence[CompiledFormulaInfo]:
        if query_part == QueryPart.select:
            return self.select
        if query_part == QueryPart.group_by:
            return self.group_by
        if query_part == QueryPart.filters:
            return self.filters
        if query_part == QueryPart.order_by:
            return self.order_by
        if query_part == QueryPart.join_on:
            return self.join_on
        raise ValueError(query_part)

    def clone(self, **updates: Any) -> CompiledQuery:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class AvatarFromObject(FromObject):
    avatar_id: str = attr.ib(kw_only=True)
    source_id: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class SubqueryFromObject(FromObject):
    query_id: str = attr.ib(kw_only=True)


_MULTI_QUERY_TV = TypeVar("_MULTI_QUERY_TV", bound="CompiledMultiQueryBase")


@attr.s(frozen=True)
class CompiledMultiQueryBase(abc.ABC):
    @abc.abstractmethod
    def iter_queries(self) -> Iterable[CompiledQuery]:
        raise NotImplementedError

    @abc.abstractmethod
    def get_top_queries(self) -> list[CompiledQuery]:
        raise NotImplementedError

    @abc.abstractmethod
    def get_bottom_queries(self) -> list[CompiledQuery]:
        raise NotImplementedError

    def get_single_top_query(self) -> CompiledQuery:
        queries = self.get_top_queries()
        assert len(queries) == 1
        return queries[0]

    @abc.abstractmethod
    def get_query_by_id(self, id: str) -> CompiledQuery:
        raise NotImplementedError

    @abc.abstractmethod
    def for_level_type(self: _MULTI_QUERY_TV, level_type: ExecutionLevel) -> _MULTI_QUERY_TV:
        raise NotImplementedError

    @abc.abstractmethod
    def get_base_froms(self) -> dict[str, FromObject]:
        raise NotImplementedError

    @abc.abstractmethod
    def is_empty(self) -> bool:
        raise NotImplementedError

    @abc.abstractmethod
    def get_base_root_from_ids(self) -> AbstractSet[str]:
        raise NotImplementedError

    def clone(self: _MULTI_QUERY_TV, **updates: Any) -> _MULTI_QUERY_TV:
        return attr.evolve(self, **updates)

    def get_complexity(self) -> int:
        return sum(query.get_complexity() for query in self.iter_queries())

    def query_count(self) -> int:
        return sum(1 for _ in self.iter_queries())


@attr.s(frozen=True)
class CompiledMultiQuery(CompiledMultiQueryBase):
    queries: list[CompiledQuery] = attr.ib(kw_only=True)

    # internal caches
    _queries_by_id: dict[str, CompiledQuery] = attr.ib(init=False)
    _unreferenced_ids: list[str] = attr.ib(init=False)

    @_queries_by_id.default
    def _make_queries_by_id(self) -> dict[str, CompiledQuery]:
        return {query.id: query for query in self.queries}

    @_unreferenced_ids.default
    def _make_unreferenced_ids(self) -> list[str]:
        referenced_ids = {from_obj.id for query in self.queries for from_obj in query.joined_from.froms}
        query_ids = {query.id for query in self.queries}
        return sorted(query_ids - referenced_ids)

    def iter_queries(self) -> Iterable[CompiledQuery]:
        return self.queries

    def get_top_queries(self) -> list[CompiledQuery]:
        return [self._queries_by_id[query_id] for query_id in self._unreferenced_ids]

    def get_bottom_queries(self) -> list[CompiledQuery]:
        raise NotImplementedError

    def get_query_by_id(self, id: str) -> CompiledQuery:
        return self._queries_by_id[id]

    def for_level_type(self, level_type: ExecutionLevel) -> CompiledMultiQuery:
        queries = [query for query in self.queries if query.level_type == level_type]
        return self.clone(queries=queries)

    def get_base_froms(self) -> dict[str, FromObject]:
        result: dict[str, FromObject] = {}
        for query in self.queries:
            for from_obj in query.joined_from.froms:
                if from_obj.id not in self._queries_by_id and from_obj.id not in result:
                    result[from_obj.id] = from_obj
        return result

    def is_empty(self) -> bool:
        return not self.queries

    def get_base_root_from_ids(self) -> AbstractSet[str]:
        result: set[str] = set()
        for query in self.queries:
            root_from_id = query.joined_from.root_from_id
            if root_from_id is not None and root_from_id not in self._queries_by_id:
                result.add(root_from_id)
        return result

lib/dl_query_processing/dl_query_processing/compilation/primitives.py
datalens-tech/datalens-backend
from __future__ import annotations

from itertools import chain
from itertools import count as it_count
from typing import (
    Optional,
    TypeVar,
)

import attr

from dl_constants.enums import OrderDirection
from dl_core.components.accessor import DatasetComponentAccessor
from dl_core.components.ids import (
    FieldId,
    RelationId,
)
from dl_core.us_dataset import Dataset
from dl_core.utils import attrs_evolve_to_subclass
from dl_formula.collections import (
    NodeSet,
    NodeValueMap,
)
import dl_formula.core.nodes as formula_nodes
from dl_query_processing.column_registry import ColumnRegistry
from dl_query_processing.compilation.base import RawQueryCompilerBase
from dl_query_processing.compilation.filter_compiler import FilterFormulaCompiler
from dl_query_processing.compilation.formula_compiler import FormulaCompiler
from dl_query_processing.compilation.helpers import make_joined_from_for_avatars
from dl_query_processing.compilation.primitives import (
    BASE_QUERY_ID,
    CompiledFormulaInfo,
    CompiledJoinOnFormulaInfo,
    CompiledOrderByFormulaInfo,
    CompiledQuery,
)
from dl_query_processing.compilation.specs import (
    QuerySpec,
    SelectWrapperSpec,
)
from dl_query_processing.compilation.wrapper_applicator import ExpressionWrapperApplicator
from dl_query_processing.enums import (
    ExecutionLevel,
    SelectValueType,
)


_COMPILED_FORMULA_INFO_TV = TypeVar("_COMPILED_FORMULA_INFO_TV", bound=CompiledFormulaInfo)


@attr.s
class DefaultQueryCompiler(RawQueryCompilerBase):
    _dataset: Dataset = attr.ib(kw_only=True)
    _column_reg: ColumnRegistry = attr.ib(kw_only=False)
    _formula_compiler: FormulaCompiler = attr.ib(kw_only=False)
    _filter_compiler: FilterFormulaCompiler = attr.ib(kw_only=True)
    _wrapper_applicator: ExpressionWrapperApplicator = attr.ib(kw_only=True, factory=ExpressionWrapperApplicator)

    _alias_cnt: it_count = attr.ib(init=False, factory=it_count)  # type: ignore  # TODO: fix
    _alias_node_mapping: NodeValueMap[str] = attr.ib(init=False, factory=NodeValueMap)
    _ds_accessor: DatasetComponentAccessor = attr.ib(init=False)

    @_ds_accessor.default
    def _make_ds_accessor(self) -> DatasetComponentAccessor:
        return DatasetComponentAccessor(dataset=self._dataset)

    def _get_expression_alias(self, expr: formula_nodes.Formula, force: bool = False) -> str:
        if force or expr not in self._alias_node_mapping:
            self._alias_node_mapping.add(node=expr, value=f"res_{next(self._alias_cnt)}")
        alias = self._alias_node_mapping.get(node=expr)
        assert alias is not None
        return alias

    def _update_alias(self, formula: _COMPILED_FORMULA_INFO_TV, force: bool = False) -> _COMPILED_FORMULA_INFO_TV:
        alias = self._get_expression_alias(expr=formula.formula_obj, force=force)
        formula = formula.clone(alias=alias)
        return formula

    def _replace_wrapped_formula(
        self,
        formula: CompiledFormulaInfo,
        expr: formula_nodes.FormulaItem,
        override_alias: Optional[str] = None,
    ) -> CompiledFormulaInfo:
        alias_kwarg = dict(alias=override_alias) if override_alias else {}
        return formula.clone(
            formula_obj=formula_nodes.Formula.make(expr=expr),
            **alias_kwarg,
        )

    def _get_override_alias(self, select_wrapper: SelectWrapperSpec) -> Optional[str]:
        if select_wrapper.type == SelectValueType.min:
            return "min_value"
        elif select_wrapper.type == SelectValueType.max:
            return "max_value"
        return None

    def _make_compiled_formula(
        self,
        field_id: FieldId,
        select_wrapper: SelectWrapperSpec = SelectWrapperSpec(type=SelectValueType.plain),
    ) -> CompiledFormulaInfo:
        formula: CompiledFormulaInfo
        # All of these types require the actual field object
        field = self._dataset.result_schema.by_guid(field_id)
        formula = self._formula_compiler.compile_field_formula(field=field)
        expr = formula.formula_obj.expr
        wrapped_expr = self._wrapper_applicator.apply_wrapper(expr, wrapper=select_wrapper)
        override_alias = self._get_override_alias(select_wrapper=select_wrapper)
        if wrapped_expr is not expr or override_alias is not None:
            formula = self._replace_wrapped_formula(formula, expr=wrapped_expr, override_alias=override_alias)

        # Override aliases with unique short ones
        formula = self._update_alias(formula)
        return formula

    def _make_compiled_order_by_formula(
        self,
        field_id: FieldId,
        direction: OrderDirection,
        select_wrapper: SelectWrapperSpec = SelectWrapperSpec(type=SelectValueType.plain),
    ) -> CompiledOrderByFormulaInfo:
        formula = self._make_compiled_formula(field_id=field_id, select_wrapper=select_wrapper)
        formula = attrs_evolve_to_subclass(
            cls=CompiledOrderByFormulaInfo,
            inst=formula,
            direction=direction,
        )
        return formula

    def _make_compiled_join_on_formula(self, relation_id: RelationId) -> CompiledJoinOnFormulaInfo:
        relation = self._ds_accessor.get_avatar_relation_strict(relation_id=relation_id)
        formula = self._formula_compiler.compile_relation_formula(relation=relation)
        return formula

    def _make_select(self, query_spec: QuerySpec) -> list[CompiledFormulaInfo]:
        result: list[CompiledFormulaInfo] = []
        used_aliases: set[str] = set()
        copy_counter = it_count()
        for field_spec in query_spec.select_specs:
            formula = self._make_compiled_formula(
                field_id=field_spec.field_id,
                select_wrapper=field_spec.wrapper,
            )
            alias = formula.not_none_alias
            if alias in used_aliases:
                alias = f"{alias}_cp{next(copy_counter)}"
                formula = formula.clone(alias=alias)

            assert alias not in used_aliases
            used_aliases.add(alias)
            result.append(formula)

        return result

    def _make_group_by(self, query_spec: QuerySpec) -> list[CompiledFormulaInfo]:
        return [
            self._make_compiled_formula(
                field_id=field_spec.field_id,
                select_wrapper=field_spec.wrapper,
            )
            for field_spec in query_spec.group_by_specs
        ]

    def _make_order_by(self, query_spec: QuerySpec) -> list[CompiledOrderByFormulaInfo]:
        result: list[CompiledOrderByFormulaInfo] = []
        already_used_ob_exprs = NodeSet()
        for order_by_spec in query_spec.order_by_specs:
            ob_expr = self._make_compiled_order_by_formula(
                field_id=order_by_spec.field_id,
                select_wrapper=order_by_spec.wrapper,
                direction=order_by_spec.direction,
            )
            if ob_expr.formula_obj in already_used_ob_exprs:
                # Deduplicate ORDER BY items
                continue
            result.append(ob_expr)
            already_used_ob_exprs.add(ob_expr.formula_obj)
        return result

    def _make_join_on(self, query_spec: QuerySpec) -> list[CompiledJoinOnFormulaInfo]:
        return [
            self._update_alias(
                self._make_compiled_join_on_formula(
                    relation_id=relation_spec.relation_id,
                )
            )
            for relation_spec in query_spec.relation_specs
        ]

    def _make_filters(self, query_spec: QuerySpec) -> list[CompiledFormulaInfo]:
        result: list[CompiledFormulaInfo] = []

        # Source filters
        for scf_spec in query_spec.source_column_filter_specs:
            formula = self._filter_compiler.compile_source_column_filter_formula(source_column_filter_spec=scf_spec)
            formula = self._update_alias(formula)
            result.append(formula)

        # Regular filters
        for filter_spec in query_spec.filter_specs:
            formula = self._filter_compiler.compile_filter_formula(filter_spec=filter_spec)
            formula = self._update_alias(formula)
            result.append(formula)

        return result

    def make_compiled_query(self, query_spec: QuerySpec) -> CompiledQuery:
        select = self._make_select(query_spec)
        group_by = self._make_group_by(query_spec)
        order_by = self._make_order_by(query_spec)
        filters = self._make_filters(query_spec)
        join_on = self._make_join_on(query_spec)

        used_avatar_ids = set(
            chain.from_iterable(
                formula.avatar_ids
                for expr_list in (select, group_by, order_by, filters, join_on)
                for formula in expr_list  # type: ignore  # TODO: fix
            )
        )

        joined_from = make_joined_from_for_avatars(
            used_avatar_ids=used_avatar_ids,
            ds_accessor=self._ds_accessor,
            root_avatar_id=query_spec.root_avatar_id,
            column_reg=self._column_reg,
        )

        compiled_query = CompiledQuery(
            id=BASE_QUERY_ID,
            level_type=ExecutionLevel.source_db,
            select=select,
            group_by=group_by,
            order_by=order_by,
            filters=filters,
            join_on=join_on,
            joined_from=joined_from,
            limit=query_spec.limit,
            offset=query_spec.offset,
            meta=query_spec.meta,
        )
        return compiled_query

lib/dl_query_processing/dl_query_processing/compilation/query_compiler.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Optional,
    TypeVar,
)

import attr

from dl_core.constants import DataAPILimits
from dl_query_processing.enums import (
    EmptyQueryMode,
    QueryType,
)


_QUERY_META_TV = TypeVar("_QUERY_META_TV", bound="QueryMetaInfo")


@attr.s
class QueryMetaInfo:
    query_type: QueryType = attr.ib(kw_only=True, default=QueryType.result)
    phantom_select_ids: list[str] = attr.ib(kw_only=True, factory=list)
    field_order: Optional[list[tuple[int, str]]] = attr.ib(kw_only=True, default=None)
    row_count_hard_limit: int = attr.ib(kw_only=True, default=DataAPILimits.DEFAULT_SOURCE_DB_LIMIT)
    from_subquery: bool = attr.ib(kw_only=True, default=False)
    subquery_limit: int = attr.ib(kw_only=True, default=None)
    empty_query_mode: EmptyQueryMode = attr.ib(kw_only=True, default=EmptyQueryMode.error)

    @property
    def result_field_ids(self) -> list[str]:
        assert self.field_order is not None
        return [field_id for idx, field_id in self.field_order]

    def clone(self: _QUERY_META_TV, **updates: Any) -> _QUERY_META_TV:
        return attr.evolve(self, **updates)

lib/dl_query_processing/dl_query_processing/compilation/query_meta.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    Callable,
    List,
    Sequence,
    Set,
    TypeVar,
)

import attr

from dl_core.components.ids import FieldId
from dl_formula.collections import NodeSet
from dl_formula.core.dialect import DialectCombo
import dl_formula.core.exc as formula_exc
import dl_formula.core.nodes as formula_nodes
from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.inspect.expression import (
    contains_extended_aggregations,
    contains_lookup_functions,
    get_toplevel_dimension_set,
    is_window_expression,
)
from dl_formula.mutation.general import (
    OptimizeConstAndOrMutation,
    OptimizeConstComparisonMutation,
    OptimizeConstFuncMutation,
    OptimizeUnaryBoolFunctions,
)
from dl_formula.mutation.lod import (
    DoubleAggregationCollapsingMutation,
    ExtAggregationToQueryForkMutation,
)
from dl_formula.mutation.lookup import LookupFunctionToQueryForkMutation
from dl_formula.mutation.mutation import (
    FormulaMutation,
    apply_mutations,
)
from dl_formula.mutation.window import WindowFunctionToQueryForkMutation
from dl_formula.validation.aggregation import AggregationChecker
from dl_formula.validation.env import ValidationEnvironment
from dl_formula.validation.validator import validate
from dl_query_processing.compilation.primitives import (
    CompiledFormulaInfo,
    CompiledQuery,
)
from dl_query_processing.enums import (
    QueryPart,
    QueryType,
)


_COMPILED_FLA_TV = TypeVar("_COMPILED_FLA_TV", bound=CompiledFormulaInfo)


def get_toplevel_dimension_set_for_query(compiled_query: CompiledQuery) -> NodeSet:
    all_dimensions = NodeSet()
    for formula in compiled_query.all_formulas:
        all_dimensions |= get_toplevel_dimension_set(formula.formula_obj.expr)
    return all_dimensions


class AtomicQueryMutatorBase(abc.ABC):
    @abc.abstractmethod
    def match_query(self, compiled_query: CompiledQuery) -> bool:
        raise NotImplementedError

    @abc.abstractmethod
    def mutate_query(self, compiled_query: CompiledQuery) -> CompiledQuery:
        raise NotImplementedError


class AtomicQueryFormulaListMutatorBase(AtomicQueryMutatorBase):
    @abc.abstractmethod
    def mutate_formula_list(
        self,
        formula_list: List[_COMPILED_FLA_TV],
        query_part: QueryPart,
    ) -> List[_COMPILED_FLA_TV]:
        raise NotImplementedError

    def mutate_query(self, compiled_query: CompiledQuery) -> CompiledQuery:
        new_select = self.mutate_formula_list(compiled_query.select, query_part=QueryPart.select)
        new_group_by = self.mutate_formula_list(compiled_query.group_by, query_part=QueryPart.group_by)
        new_filters = self.mutate_formula_list(compiled_query.filters, query_part=QueryPart.filters)
        new_order_by = self.mutate_formula_list(compiled_query.order_by, query_part=QueryPart.order_by)
        new_join_on = self.mutate_formula_list(compiled_query.join_on, query_part=QueryPart.join_on)

        changed = (
            new_select is not compiled_query.select
            or new_group_by is not compiled_query.group_by
            or new_filters is not compiled_query.filters
            or new_order_by is not compiled_query.order_by
            or new_join_on is not compiled_query.join_on
        )

        if changed:
            compiled_query = compiled_query.clone(
                select=new_select,
                group_by=new_group_by,
                filters=new_filters,
                order_by=new_order_by,
                join_on=new_join_on,
            )

        return compiled_query


class AtomicQueryFormulaMutatorBase(AtomicQueryFormulaListMutatorBase):
    @abc.abstractmethod
    def mutate_formula(self, formula: _COMPILED_FLA_TV) -> _COMPILED_FLA_TV:
        raise NotImplementedError

    def mutate_formula_list(
        self,
        formula_list: List[_COMPILED_FLA_TV],
        query_part: QueryPart,
    ) -> List[_COMPILED_FLA_TV]:
        new_formula_list: List[_COMPILED_FLA_TV] = []
        changed = False
        for formula in formula_list:
            new_formula = self.mutate_formula(formula)
            if new_formula is not formula:
                changed = True
            new_formula_list.append(new_formula)

        if changed:
            formula_list = new_formula_list

        return formula_list


@attr.s
class IgnoreFormulaAtomicQueryMutator(AtomicQueryFormulaListMutatorBase):
    _ignore_formula_checks: Sequence[Callable[[formula_nodes.Formula, QueryPart], bool]] = attr.ib(kw_only=True)

    def match_query(self, compiled_query: CompiledQuery) -> bool:
        return True  # Apply to all

    def should_ignore_formula(self, formula: CompiledFormulaInfo, query_part: QueryPart) -> bool:
        for checker in self._ignore_formula_checks:
            if checker(formula.formula_obj, query_part):
                return True
        return False

    def mutate_formula_list(
        self,
        formula_list: List[_COMPILED_FLA_TV],
        query_part: QueryPart,
    ) -> List[_COMPILED_FLA_TV]:
        new_formula_list = [
            formula for formula in formula_list if not self.should_ignore_formula(formula, query_part=query_part)
        ]
        if len(new_formula_list) == len(formula_list):
            return formula_list

        return new_formula_list


@attr.s
class NullifyFormulaAtomicQueryMutator(AtomicQueryFormulaListMutatorBase):
    _nullify_formula_checks: Sequence[Callable[[formula_nodes.Formula, QueryPart], bool]] = attr.ib(kw_only=True)

    def match_query(self, compiled_query: CompiledQuery) -> bool:
        return True  # Apply to all

    def should_nullify_formula(self, formula: CompiledFormulaInfo, query_part: QueryPart) -> bool:
        for checker in self._nullify_formula_checks:
            if checker(formula.formula_obj, query_part):
                return True
        return False

    def nullify_formula(self, formula: _COMPILED_FLA_TV) -> _COMPILED_FLA_TV:
        return formula.clone(
            formula_obj=formula_nodes.Formula.make(
                expr=formula_nodes.Null(meta=formula.formula_obj.expr.meta),
                meta=formula.formula_obj.meta,
            )
        )

    def mutate_formula_list(
        self,
        formula_list: List[_COMPILED_FLA_TV],
        query_part: QueryPart,
    ) -> List[_COMPILED_FLA_TV]:
        new_formula_list = [
            self.nullify_formula(formula) if self.should_nullify_formula(formula, query_part=query_part) else formula
            for formula in formula_list
        ]
        return new_formula_list


def contains_inconsistent_aggregations(
    node: formula_nodes.FormulaItem,
    dimensions: list[formula_nodes.FormulaItem],
    env: ValidationEnvironment,
) -> bool:
    try:
        validate(
            node,
            env=env,
            checkers=[
                AggregationChecker(valid_env=env, inspect_env=InspectionEnvironment(), global_dimensions=dimensions)
            ],
            collect_errors=False,
        )
    except formula_exc.ValidationError:
        return True

    return False


@attr.s
class DefaultAtomicQueryMutator(AtomicQueryFormulaMutatorBase):
    _mutations: Sequence[FormulaMutation] = attr.ib(kw_only=True)

    def match_query(self, compiled_query: CompiledQuery) -> bool:
        return True  # Apply to all

    def mutate_formula(self, formula: _COMPILED_FLA_TV) -> _COMPILED_FLA_TV:
        formula_obj = apply_mutations(tree=formula.formula_obj, mutations=self._mutations)
        if formula_obj is formula.formula_obj:
            # No changes
            return formula
        formula = formula.clone(formula_obj=formula_obj)
        return formula


def formula_is_true(formula: formula_nodes.Formula, query_part: QueryPart) -> bool:
    return (
        query_part == QueryPart.filters
        and isinstance(formula.expr, formula_nodes.LiteralBoolean)
        and formula.expr.value is True
    )


@attr.s
class QueryMutator(abc.ABC):
    """
    Applies mutations to formulas of the query after it has been more or less assembled.
    They are applied conditionally after the whole query is examined (or not).
    """

    @abc.abstractmethod
    def mutate_query(self, compiled_query: CompiledQuery) -> CompiledQuery:
        raise NotImplementedError


@attr.s
class OptimizingQueryMutator(QueryMutator):
    """Mutator applies various optimizations and (in specific cases) "fixes" to the query"""

    _dialect: DialectCombo = attr.ib(kw_only=True)
    _disable_optimizations: bool = attr.ib(kw_only=True, default=False)

    def mutate_query(self, compiled_query: CompiledQuery) -> CompiledQuery:
        # Apply "preliminary" mutation
        # (must be applied before we scan for nested aggregations)
        if not self._disable_optimizations:
            mutator: AtomicQueryMutatorBase = DefaultAtomicQueryMutator(
                mutations=[
                    DoubleAggregationCollapsingMutation(),
                    OptimizeConstComparisonMutation(),
                    OptimizeConstAndOrMutation(),
                    OptimizeUnaryBoolFunctions(self._dialect),
                    OptimizeConstFuncMutation(),
                ]
            )
            compiled_query = mutator.mutate_query(compiled_query)

        filter_mutator = IgnoreFormulaAtomicQueryMutator(ignore_formula_checks=[formula_is_true])
        compiled_query = filter_mutator.mutate_query(compiled_query)

        has_extaggs = any(
            contains_extended_aggregations(node=formula.formula_obj, include_double_agg=True)
            for formula in compiled_query.all_formulas
        )

        # Disable all the query forking and window functions for (sub)totals.
        if compiled_query.meta.query_type == QueryType.totals:
            validation_env = ValidationEnvironment()
            inspect_env = InspectionEnvironment()
            dimensions = [formula.formula_obj.expr for formula in compiled_query.group_by]
            has_incaggs = any(
                contains_inconsistent_aggregations(
                    node=formula.formula_obj,
                    dimensions=dimensions,
                    env=validation_env,
                )
                for formula in compiled_query.all_formulas
            )
            has_winfuncs = any(
                is_window_expression(node=formula.formula_obj, env=inspect_env)
                for formula in compiled_query.all_formulas
            )

            if has_extaggs or has_incaggs or has_winfuncs:
                null_mutator = NullifyFormulaAtomicQueryMutator(
                    nullify_formula_checks=[
                        lambda f, qpart: contains_lookup_functions(f),
                        lambda f, qpart: contains_extended_aggregations(f, include_double_agg=True),
                        lambda f, qpart: contains_inconsistent_aggregations(
                            f,
                            dimensions=dimensions,
                            env=validation_env,
                        ),
                        lambda f, qpart: is_window_expression(f, env=inspect_env),
                    ]
                )
                compiled_query = null_mutator.mutate_query(compiled_query)

        return compiled_query


@attr.s
class ExtendedAggregationQueryMutator(QueryMutator):
    """Mutator specifically for handling LODs and lookup functions"""

    _allow_empty_dimensions_for_forks: bool = attr.ib(kw_only=True)
    _allow_arbitrary_toplevel_lod_dimensions: bool = attr.ib(kw_only=True)
    _new_subquery_mode: bool = attr.ib(kw_only=True, default=False)

    def mutate_query(self, compiled_query: CompiledQuery) -> CompiledQuery:
        global_dimensions = [formula.formula_obj.expr for formula in compiled_query.group_by]

        filter_ids: Set[FieldId] = set()
        for formula in compiled_query.filters:
            if formula.original_field_id:
                filter_ids.add(formula.original_field_id)
            # Otherwise it is an anonymous filter (RLS) that cannot interact with BFB

        has_extaggs = any(
            contains_extended_aggregations(node=formula.formula_obj, include_double_agg=True)
            for formula in compiled_query.all_formulas
        )

        has_lookups = any(
            contains_lookup_functions(node=formula.formula_obj) for formula in compiled_query.all_formulas
        )

        inspect_env = InspectionEnvironment()
        has_winfuncs = any(
            is_window_expression(node=formula.formula_obj, env=inspect_env) for formula in compiled_query.all_formulas
        )
        will_mutate_winfuncs = has_winfuncs

        fork_mutators: List[AtomicQueryMutatorBase] = []

        if has_lookups:
            fork_mutators += [
                DefaultAtomicQueryMutator(
                    mutations=[
                        LookupFunctionToQueryForkMutation(
                            global_dimensions=global_dimensions,
                            allow_empty_dimensions=self._allow_empty_dimensions_for_forks,
                        ),
                    ]
                ),
            ]

        if has_extaggs or will_mutate_winfuncs:
            # Window functions also require aggregations to be wrapped in QF
            # so that they can be split-off to sub-queries separate from the ones with the WF
            fork_mutators += [
                DefaultAtomicQueryMutator(
                    mutations=[
                        ExtAggregationToQueryForkMutation(global_dimensions=global_dimensions),
                    ]
                ),
            ]

        if will_mutate_winfuncs:
            fork_mutators += [
                DefaultAtomicQueryMutator(
                    mutations=[
                        WindowFunctionToQueryForkMutation(global_dimensions=global_dimensions),
                    ]
                ),
            ]

        atomic_mutators: list[AtomicQueryMutatorBase]
        atomic_mutators = fork_mutators

        for mutator in atomic_mutators:
            if mutator.match_query(compiled_query):
                compiled_query = mutator.mutate_query(compiled_query)

        if has_extaggs and not self._allow_arbitrary_toplevel_lod_dimensions:
            dim_set_from_group_by = NodeSet(global_dimensions)
            dim_set_from_top_lods = get_toplevel_dimension_set_for_query(compiled_query)
            if (dim_set_from_group_by | dim_set_from_top_lods) != dim_set_from_group_by:
                # LODs contain some dimensions that are not in `group_by`
                raise formula_exc.LodInvalidTopLevelDimensionsError(
                    "Invalid top-level LOD dimension found in expression"
                )

        if has_extaggs and not self._new_subquery_mode:
            # The real group_by's of each sub-query will be managed by forker
            # while it is inspecting QueryFork's.
            # The original group_by can be removed because they will only get in the way
            compiled_query = compiled_query.clone(group_by=[])

        return compiled_query

lib/dl_query_processing/dl_query_processing/compilation/query_mutator.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    AbstractSet,
    Any,
    Optional,
    TypeVar,
    Union,
)

import attr

from dl_constants.enums import (
    OrderDirection,
    WhereClauseOperation,
)
from dl_core.components.ids import (
    AvatarId,
    FieldId,
    RelationId,
)
from dl_query_processing.compilation.query_meta import QueryMetaInfo
from dl_query_processing.enums import SelectValueType


FilterArgType = Union[str, int, float, None]

_QUERY_ITEM_SPEC_TV = TypeVar("_QUERY_ITEM_SPEC_TV", bound="QueryItemSpec")


@attr.s(frozen=True)
class QueryItemSpec:
    def clone(self: _QUERY_ITEM_SPEC_TV, **kwargs: Any) -> _QUERY_ITEM_SPEC_TV:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True)
class QueryItemFieldSpec(QueryItemSpec):
    """
    Base class for field-oriented specs
    """

    field_id: FieldId = attr.ib(kw_only=True)


@attr.s(frozen=True)
class SelectWrapperSpec:
    """
    SELECT clause item wrapper specifier.
    """

    type: SelectValueType = attr.ib(kw_only=True, default=SelectValueType.plain)


@attr.s(frozen=True)
class ArrayPrefixSelectWrapperSpec(SelectWrapperSpec):
    """
    Array prefix selection wrapper (for trees).
    """

    type: SelectValueType = attr.ib(kw_only=True, default=SelectValueType.array_prefix)
    length: int = attr.ib(kw_only=True)


@attr.s(frozen=True)
class SelectFieldSpec(QueryItemFieldSpec):
    """
    SELECT clause item specifier.
    """

    wrapper: SelectWrapperSpec = attr.ib(kw_only=True, factory=SelectWrapperSpec)


@attr.s(frozen=True)
class GroupByFieldSpec(QueryItemFieldSpec):
    """
    GROUP BY clause item specifier.
    """

    wrapper: SelectWrapperSpec = attr.ib(kw_only=True, factory=SelectWrapperSpec)


@attr.s(frozen=True)
class OrderByFieldSpec(QueryItemFieldSpec):
    """
    ORDER BY field clause specifier.
    """

    wrapper: SelectWrapperSpec = attr.ib(kw_only=True, factory=SelectWrapperSpec)
    direction: OrderDirection = attr.ib(kw_only=True)


@attr.s(frozen=True)
class FilterFieldSpec(QueryItemFieldSpec):
    """
    Field-based filter specifier (for WHERE and HAVING)
    """

    operation: WhereClauseOperation = attr.ib(kw_only=True)
    values: list[FilterArgType] = attr.ib(kw_only=True)
    # Indicates that the filter formula should not have an `original_field_id`.
    # For usage in RLS filters so that BFB clauses cannot be used to avoid RLS restrictions
    anonymous: bool = attr.ib(kw_only=True, default=False)


@attr.s(frozen=True)
class FilterSourceColumnSpec(QueryItemSpec):
    """
    Table column-based filter specifier (for WHERE and HAVING)
    """

    avatar_id: FieldId = attr.ib(kw_only=True)
    column_name: str = attr.ib(kw_only=True)
    operation: WhereClauseOperation = attr.ib(kw_only=True)
    values: list[FilterArgType] = attr.ib(kw_only=True)


@attr.s(frozen=True)
class RelationSpec(QueryItemSpec):
    """
    Relation (ON clause) specifier
    """

    relation_id: RelationId = attr.ib(kw_only=True)


@attr.s(frozen=True)
class ParameterValueSpec(QueryItemFieldSpec):
    """
    Parameter value specifier
    """

    value: Any = attr.ib(kw_only=True)


@attr.s(slots=True, frozen=True)
class QuerySpec:
    select_specs: list[SelectFieldSpec] = attr.ib(kw_only=True)
    group_by_specs: list[GroupByFieldSpec] = attr.ib(kw_only=True)
    order_by_specs: list[OrderByFieldSpec] = attr.ib(kw_only=True)
    filter_specs: list[FilterFieldSpec] = attr.ib(kw_only=True)
    relation_specs: list[RelationSpec] = attr.ib(kw_only=True)
    source_column_filter_specs: list[FilterSourceColumnSpec] = attr.ib(kw_only=True)
    parameter_value_specs: list[ParameterValueSpec] = attr.ib(kw_only=True)
    limit: Optional[int] = attr.ib(kw_only=True)
    offset: Optional[int] = attr.ib(kw_only=True)
    root_avatar_id: Optional[AvatarId] = attr.ib(kw_only=True, default=None)
    required_avatar_ids: AbstractSet[AvatarId] = attr.ib(kw_only=True, factory=frozenset)
    meta: QueryMetaInfo = attr.ib(kw_only=True, factory=QueryMetaInfo)

lib/dl_query_processing/dl_query_processing/compilation/specs.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_constants.enums import UserDataType
from dl_formula.core.datatype import DataType


BI_TO_FORMULA_TYPES = {
    UserDataType.integer: DataType.INTEGER,
    UserDataType.float: DataType.FLOAT,
    UserDataType.boolean: DataType.BOOLEAN,
    UserDataType.string: DataType.STRING,
    UserDataType.date: DataType.DATE,
    UserDataType.datetime: DataType.DATETIME,
    UserDataType.datetimetz: DataType.DATETIMETZ,
    UserDataType.genericdatetime: DataType.GENERICDATETIME,
    UserDataType.geopoint: DataType.GEOPOINT,
    UserDataType.geopolygon: DataType.GEOPOLYGON,
    UserDataType.uuid: DataType.UUID,
    UserDataType.markup: DataType.MARKUP,
    UserDataType.unsupported: DataType.UNSUPPORTED,
    UserDataType.array_float: DataType.ARRAY_FLOAT,
    UserDataType.array_int: DataType.ARRAY_INT,
    UserDataType.array_str: DataType.ARRAY_STR,
    UserDataType.tree_str: DataType.TREE_STR,
}
FORMULA_TO_BI_TYPES = {
    **{ft: bit for bit, ft in BI_TO_FORMULA_TYPES.items()},
    **{ft.const_type: bit for bit, ft in BI_TO_FORMULA_TYPES.items()},
    DataType.NULL: UserDataType.string,  # NULL can in theory be any type, but we need to choose one
}
DEFAULT_DATA_TYPE = DataType.STRING

lib/dl_query_processing/dl_query_processing/compilation/type_mapping.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import Sequence

import dl_formula.core.nodes as formula_nodes
from dl_query_processing.compilation.specs import (
    ArrayPrefixSelectWrapperSpec,
    SelectWrapperSpec,
)
from dl_query_processing.enums import SelectValueType


class ExpressionWrapperApplicator:
    def _apply_func_wrapper(
        self,
        expr: formula_nodes.FormulaItem,
        func_name: str,
        additional_args: Sequence[formula_nodes.FormulaItem] = (),
    ) -> formula_nodes.FormulaItem:
        func_name = func_name.lower()
        return formula_nodes.FuncCall.make(
            name=func_name,
            args=[expr, *additional_args],
        )

    def _apply_wrapper_min(self, expr: formula_nodes.FormulaItem) -> formula_nodes.FormulaItem:
        return self._apply_func_wrapper(expr, func_name="min")

    def _apply_wrapper_max(self, expr: formula_nodes.FormulaItem) -> formula_nodes.FormulaItem:
        return self._apply_func_wrapper(expr, func_name="max")

    def _apply_wrapper_array_prefix(
        self,
        expr: formula_nodes.FormulaItem,
        length: int,
    ) -> formula_nodes.FormulaItem:
        return self._apply_func_wrapper(
            expr,
            func_name="slice",
            additional_args=[
                formula_nodes.LiteralInteger.make(1),
                formula_nodes.LiteralInteger.make(length),
            ],
        )

    def apply_wrapper(self, expr: formula_nodes.FormulaItem, wrapper: SelectWrapperSpec) -> formula_nodes.FormulaItem:
        if wrapper.type == SelectValueType.plain:
            pass
        elif wrapper.type == SelectValueType.min:
            expr = self._apply_wrapper_min(expr)
        elif wrapper.type == SelectValueType.max:
            expr = self._apply_wrapper_max(expr)
        elif wrapper.type == SelectValueType.array_prefix:
            assert isinstance(wrapper, ArrayPrefixSelectWrapperSpec)
            expr = self._apply_wrapper_array_prefix(expr, length=wrapper.length)
        else:
            raise ValueError(f"Unsupported wrapper.type value: {wrapper.type}")

        return expr

lib/dl_query_processing/dl_query_processing/compilation/wrapper_applicator.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import (
    Enum,
    unique,
)


@unique
class QueryType(Enum):
    result = "result"
    preview = "preview"
    distinct = "distinct"  # FIXME: value_distinct
    totals = "totals"
    value_range = "value_range"
    pivot = "pivot"


@unique
class SelectValueType(Enum):
    plain = "plain"  # Select as-is, without any modifiers (default for most cases)
    max = "max"  # Wrap in MAX
    min = "min"  # Wrap in MIN
    literal = "literal"  # Select an arbitrary literal instead
    null = "null"  # Select `NULL` instead
    array_prefix = "array_prefix"  # Select `NULL` instead


@unique
class EmptyQueryMode(Enum):
    error = "error"  # Raise error on empty query
    empty = "empty"  # Return empty data on empty query
    empty_row = "empty_row"  # Return single empty row on empty query


@unique
class GroupByPolicy(Enum):
    force = "force"  # Enable and force GROUP BY
    disable = "disable"  # Disable Group BY
    if_measures = "if_measures"  # Enable only if the query contains measures, otherwise disable


@unique
class ProcessingStage(Enum):
    """
    Stages of field formula processing
    """

    base = "base"
    dep_generation = "dep_generation"
    pre_sub_mutation = "pre_sub_mutation"
    substitution = "substitution"
    casting = "casting"
    aggregation = "aggregation"
    mutation = "mutation"
    validation = "validation"
    final = "final"


@unique
class QueryPart(Enum):
    select = "select"
    group_by = "group_by"
    filters = "filters"
    order_by = "order_by"
    join_on = "join_on"


@unique
class ExecutionLevel(Enum):
    source_db = "source_db"
    compeng = "compeng"

lib/dl_query_processing/dl_query_processing/enums.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Optional,
    Sequence,
    Union,
)

from dl_constants.exc import DLBaseException
from dl_core.fields import BIField
from dl_formula.core import exc as formula_exc
from dl_formula.core.message_ctx import FormulaErrorCtx
from dl_query_processing.enums import ProcessingStage


class EmptyQuery(DLBaseException):
    err_code = DLBaseException.err_code + ["EMPTY_QUERY"]
    default_message = "Attempted to execute an empty query."


class LogicError(DLBaseException):
    pass


class InvalidGroupByConfiguration(DLBaseException):
    err_code = DLBaseException.err_code + ["INVALID_GROUP_BY_CONFIGURATION"]
    default_message = "Invalid GROUP BY configuration."


class FormulaHandlingError(formula_exc.FormulaError):
    # TODO: Cut off inheritance from FormulaError and patch all errors with field info
    def __init__(
        self,
        *errors: Union[FormulaErrorCtx],
        stage: Optional[ProcessingStage] = None,
        field: Optional[BIField] = None,
    ):
        super().__init__(*errors)
        self.stage = stage
        self.field = field


class InvalidLiteralError(DLBaseException):
    err_code = DLBaseException.err_code + ["INVALID_VALUE"]
    default_message = "Invalid value for literal"


class FilterError(DLBaseException):
    err_code = DLBaseException.err_code + ["FILTER"]
    default_message = "Invalid filter"


class FilterValueError(FilterError):
    err_code = FilterError.err_code + ["INVALID_VALUE"]
    default_message = "Invalid value for filter"


class MeasureFilterUnsupportedError(FilterError):
    err_code = FilterError.err_code + ["MEASURE_UNSUPPORTED"]
    default_message = "Measure filter is unsupported for this type of query"


class ParameterError(DLBaseException):
    err_code = DLBaseException.err_code + ["PARAMETER"]
    default_message = "Invalid parameter"


class ParameterValueError(ParameterError):
    err_code = ParameterError.err_code + ["INVALID_VALUE"]
    default_message = "Invalid value for parameter"


class ParameterUnsupportedTypeError(ParameterError):
    err_code = ParameterError.err_code + ["UNSUPPORTED"]
    default_message = "Unsupported type for parameter"


class InvalidQueryStructure(DLBaseException):
    err_code = DLBaseException.err_code + ["INVALID_QUERY_STRUCTURE"]
    default_message = "Failed to compile data query."


class BlockSpecError(DLBaseException):
    err_code = DLBaseException.err_code + ["BLOCK"]
    default_message = "Block spec error."


class MultipleBlocksUnsupportedError(BlockSpecError):
    err_code = BlockSpecError.err_code + ["MULTIPLE_UNSUPPORTED"]
    default_message = "This API does not support multiple query blocks."


class UnevenBlockColumnCountError(BlockSpecError):
    err_code = BlockSpecError.err_code + ["UNEVEN_COLUMN_COUNT"]
    default_message = "Blocks have different column count."


class NoRootBlockError(BlockSpecError):
    err_code = BlockSpecError.err_code + ["NO_ROOT"]
    default_message = "Got no blocks with root placement."


class MultipleRootBlockError(BlockSpecError):
    err_code = BlockSpecError.err_code + ["MULTIPLE_ROOTS"]
    default_message = "Got more than one block with root placement."


class BlockItemCompatibilityError(BlockSpecError):
    err_code = BlockSpecError.err_code + ["ITEM_COMPATIBILITY"]
    default_message = "Got items with incompatible roles in block."


class TreeError(DLBaseException):
    err_code = DLBaseException.err_code + ["TREE"]


class MultipleTreeError(TreeError):
    err_code = TreeError.err_code + ["MULTIPLE"]


class LegendError(DLBaseException):
    err_code = DLBaseException.err_code + ["LEGEND"]


class LegendItemReferenceError(LegendError):
    err_code = LegendError.err_code + ["ITEM_REFERENCE"]


class UnsopportedRoleInLegend(LegendError):
    err_code = LegendError.err_code + ["UNSUPPORTED_ROLE"]


class RoleDataTypeMismatch(LegendError):
    err_code = LegendError.err_code + ["ROLE_DATA_TYPE_MISMATCH"]


class NonUniqueLegendIdsError(LegendError):
    err_code = LegendError.err_code + ["NON_UNIQUE_IDS"]


class LegendMeasureNameError(LegendError):
    err_code = LegendError.err_code + ["MEASURE_NAME"]
    default_message = "Measure names misconfigured."


class MeasureNameUnsupported(LegendMeasureNameError):
    err_code = LegendMeasureNameError.err_code + ["UNSUPPORTED"]
    default_message = "Measure names is not supported."


class GenericInvalidRequestError(DLBaseException):
    err_code = DLBaseException.err_code + ["INVALID_REQUEST"]
    default_message = "Invalid request"


class PivotError(DLBaseException):
    err_code = DLBaseException.err_code + ["PIVOT"]
    default_message = "Pivot error."


class PivotLegendError(PivotError):
    err_code = PivotError.err_code + ["LEGEND"]


class PivotLegendItemReferenceError(PivotLegendError):
    err_code = PivotLegendError.err_code + ["ITEM_REFERENCE"]


class PivotSortingError(PivotError):
    err_code = PivotError.err_code + ["SORTING"]
    default_message = "Pivot sorting error."


class PivotSortingMultipleColumnsOrRows(PivotSortingError):
    err_code = PivotSortingError.err_code + ["MULTIPLE_COLUMNS_OR_ROWS"]
    default_message = "Requested sort by multiple columns or multiple rows."


class PivotSortingAgainstMultipleMeasures(PivotSortingError):
    err_code = PivotSortingError.err_code + ["AGAINST_MULTIPLE_MEASURES"]
    default_message = "If there are multiple measures, sorting can only be done along them."


class PivotSortingRowOrColumnNotFound(PivotSortingError):
    err_code = PivotSortingError.err_code + ["ROW_OR_COLUMN_NOT_FOUND"]
    default_message = "Requested sorting row or column not found."


class PivotSortingRowOrColumnIsAmbiguous(PivotSortingError):
    err_code = PivotSortingError.err_code + ["ROW_OR_COLUMN_IS_AMBIGUOUS"]
    default_message = "More than one row or column match requested sorting settings."


class PivotSortingWithSubtotalsIsNotAllowed(PivotSortingError):
    err_code = PivotSortingError.err_code + ["SUBTOTALS_ARE_NOT_ALLOWED"]
    default_message = "Measure sorting with subtotals is not allowed."


class PivotMeasureNameError(PivotError):
    err_code = PivotError.err_code + ["MEASURE_NAME"]
    default_message = "Pivot measure names error."


class PivotMeasureNameRequired(PivotMeasureNameError):
    err_code = PivotMeasureNameError.err_code + ["REQUIRED"]
    default_message = "Measure names must be used when pivoting multiple measures."


class PivotMeasureNameForbidden(PivotMeasureNameError):
    err_code = PivotMeasureNameError.err_code + ["FORBIDDEN"]
    default_message = "Measure names are not allowed in pivot table without measures."


class PivotMeasureNameDuplicate(PivotMeasureNameError):
    err_code = PivotMeasureNameError.err_code + ["DUPLICATE"]
    default_message = "Measure names cannot be used twice in pivot table."


class PivotInvalidRoleLegendError(PivotLegendError):
    err_code = PivotLegendError.err_code + ["INVALID_ROLE"]


class PivotDimensionError(PivotError):
    err_code = PivotError.err_code + ["DIMENSION"]
    default_message = "Pivot dimension error."


class PivotDimensionDuplicate(PivotDimensionError):
    err_code = PivotDimensionError.err_code + ["DUPLICATE"]
    default_message = "Pivot dimension cannot be used twice."


class PivotDuplicateDimensionValue(PivotDimensionError):
    err_code = PivotDimensionError.err_code + ["DUPLICATE_VALUE"]
    default_message = "Duplicate dimension values were encountered in pivot table."


class PivotOverflowError(PivotError):
    err_code = PivotError.err_code + ["OVERFLOW"]
    default_message = "Pivot table size overflow error."


class PivotUnevenDataColumnsError(PivotError):
    err_code = PivotError.err_code + ["UNEVEN_DATA_COLUMNS"]
    default_message = "Data columns for pivot table have uneven sizes."


class PivotItemsIncompatibleError(PivotLegendError):
    err_code = PivotLegendError.err_code + ["ITEMS_INCOMPATIBLE"]


class UnresolvableSlicingTagOrder(DLBaseException):
    err_code = DLBaseException.err_code + ["UNRESOLVABLE_TAG_ORDER"]
    default_message = "Function nesting order could not be resolved"  # FIXME


class DLFormulaError(DLBaseException):
    # TODO: catch all cases of FormulaError in DatasetView and re-raise as this exception
    err_code = DLBaseException.err_code + ["FORMULA"]

    def __init__(
        self,
        message: Optional[str] = None,
        details: Optional[dict] = None,
        orig: Exception = None,
        field: Optional[BIField] = None,
        formula_errors: Optional[Sequence[FormulaErrorCtx]] = None,
    ):
        details = {} if details is None else details.copy()
        if field is not None:
            details["field"] = {"guid": field.guid, "title": field.title, "errors": []}
        sub_errors = []
        for err in formula_errors or ():
            sub_errors.append(
                {
                    "code": err.code,
                    "message": err.message,
                    "row": err.position.start_row,
                    "col": err.position.start_col,
                    "token": err.token,
                }
            )
        details.update(formula_errors=sub_errors)

        super().__init__(message=message, details=details, orig=orig)


class DatasetError(DLBaseException):
    err_code = DLBaseException.err_code + ["DATASET"]


class DatasetTooManyFieldsFatal(DatasetError):
    err_code = DatasetError.err_code + ["TOO_MANY_FIELDS"]
    default_message = "There are too many fields in the dataset"

lib/dl_query_processing/dl_query_processing/exc.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    TYPE_CHECKING,
    List,
)

import attr


if TYPE_CHECKING:
    from dl_constants.enums import DataSourceRole
    from dl_core.us_connection_base import ClassicConnectionSQL
    from dl_query_processing.translation.primitives import TranslatedMultiQueryBase


@attr.s
class QueryExecutionInfo:
    role: DataSourceRole = attr.ib(kw_only=True)
    translated_multi_query: TranslatedMultiQueryBase = attr.ib(kw_only=True)
    target_connections: List[ClassicConnectionSQL] = attr.ib(kw_only=True)

lib/dl_query_processing/dl_query_processing/execution/exec_info.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import logging
from typing import (
    AbstractSet,
    Callable,
    Collection,
    Dict,
    List,
    Optional,
    Tuple,
)

import attr

from dl_constants.enums import (
    DataSourceRole,
    ProcessorType,
    SelectorType,
)
from dl_constants.types import TBIDataRow
from dl_core.components.ids import AvatarId
from dl_core.data_processing.prepared_components.default_manager import DefaultPreparedComponentManager
from dl_core.data_processing.processing.operation import (
    BaseOp,
    CalcOp,
    DownloadOp,
    JoinOp,
    UploadOp,
)
from dl_core.data_processing.processing.processor import OperationProcessorAsyncBase
from dl_core.data_processing.stream_base import (
    AbstractStream,
    DataSourceVS,
    DataStreamAsync,
)
from dl_core.services_registry import ServicesRegistry
from dl_core.us_dataset import Dataset
from dl_core.us_manager.us_manager import USManagerBase
from dl_core.utils import make_id
from dl_query_processing import exc
from dl_query_processing.enums import (
    EmptyQueryMode,
    ExecutionLevel,
    QueryType,
)
from dl_query_processing.execution.exec_info import QueryExecutionInfo
from dl_query_processing.execution.executor_base import QueryExecutorBase
from dl_query_processing.execution.primitives import (
    ExecutedQuery,
    ExecutedQueryMetaInfo,
)
from dl_query_processing.translation.primitives import TranslatedMultiQueryBase


LOGGER = logging.getLogger(__name__)


@attr.s
class SourceDbBaseSourceInfo:
    root_avatar_id: AvatarId = attr.ib(kw_only=True)
    required_avatar_ids: Collection[AvatarId] = attr.ib(kw_only=True)


@attr.s
class QueryExecutor(QueryExecutorBase):
    _dataset: Dataset = attr.ib(kw_only=True)
    _selector_type: SelectorType = attr.ib(kw_only=True)
    _compeng_processor_type: ProcessorType = attr.ib(kw_only=True)
    _source_db_processor_type: ProcessorType = attr.ib(kw_only=True)
    _allow_cache_usage: bool = attr.ib(kw_only=True)
    _avatar_alias_mapper: Optional[Callable[[str], str]] = attr.ib(kw_only=True, default=None)  # noqa
    _us_manager: USManagerBase = attr.ib(kw_only=True)
    _compeng_semaphore: asyncio.Semaphore = attr.ib(kw_only=True)

    @property
    def _service_registry(self) -> ServicesRegistry:
        return self._us_manager.get_services_registry()

    async def _get_compeng_data_processor(self) -> OperationProcessorAsyncBase:
        factory = self._service_registry.get_data_processor_factory()
        return await factory.get_data_processor(
            dataset=self._dataset,
            us_entry_buffer=self._us_manager.get_entry_buffer(),
            processor_type=self._compeng_processor_type,
            allow_cache_usage=self._allow_cache_usage,
        )

    async def _get_source_db_data_processor(self, role: DataSourceRole) -> OperationProcessorAsyncBase:
        factory = self._service_registry.get_data_processor_factory()
        return await factory.get_data_processor(
            dataset=self._dataset,
            us_entry_buffer=self._us_manager.get_entry_buffer(),
            processor_type=self._source_db_processor_type,
            allow_cache_usage=self._allow_cache_usage,
            role=role,
            selector_type=self._selector_type,
        )

    async def _get_data_processor(
        self,
        level_type: ExecutionLevel,
        role: DataSourceRole,
    ) -> OperationProcessorAsyncBase:
        if level_type == ExecutionLevel.source_db:
            return await self._get_source_db_data_processor(role=role)
        if level_type == ExecutionLevel.compeng:
            return await self._get_compeng_data_processor()
        raise ValueError(f"Unsupported level_type: {level_type}")

    async def _process_multi_query(
        self,
        *,
        level_type: ExecutionLevel,
        streams_by_result_id: Dict[AvatarId, AbstractStream],
        translated_multi_query: TranslatedMultiQueryBase,
        stream_aliases: Dict[str, str],
        upload_inputs: bool,
        distinct: bool,
        role: DataSourceRole,
        row_count_hard_limit: int,
    ) -> Tuple[Dict[AvatarId, DataStreamAsync], Dict[AvatarId, str]]:
        if translated_multi_query.is_empty():
            for stream in streams_by_result_id.values():
                assert isinstance(stream, DataStreamAsync)
            return streams_by_result_id, stream_aliases  # type: ignore

        op: BaseOp  # For usage in various parts of this function

        LOGGER.info(
            f"Executing level type {level_type}. " f"Got source streams with result IDs: {list(streams_by_result_id)}"
        )

        streams: List[DataStreamAsync] = [stream for _, stream in sorted(streams_by_result_id.items())]  # type: ignore  # TODO: fix
        operations: List[BaseOp] = []

        # 1. Make operations for uploading data (if necessary).
        # For every incoming data stream
        if upload_inputs:
            result_to_stream_id_map: Dict[AvatarId, str] = {}
            for result_id, stream in sorted(streams_by_result_id.items()):
                op = UploadOp(
                    result_id=result_id,
                    source_stream_id=stream.id,
                    alias=stream_aliases[stream.id],
                    dest_stream_id=make_id(),
                )
                operations.append(op)
                result_to_stream_id_map[result_id] = op.dest_stream_id
        else:
            result_to_stream_id_map = {result_id: stream.id for result_id, stream in streams_by_result_id.items()}

        required_avatar_ids: AbstractSet[str]

        # 2. Make operations for queries.
        # These can have many levels
        all_query_ids: set[str] = set()
        selected_avatar_ids: set[str] = set()  # avatar IDs used by at least one query
        for translated_flat_query in translated_multi_query.iter_queries():
            all_query_ids.add(translated_flat_query.id)
            selected_avatar_ids |= set(translated_flat_query.joined_from.iter_ids())
            dest_stream_id = make_id()
            # Save ID relationship for future resolution when creating CalcOps
            result_to_stream_id_map[translated_flat_query.id] = dest_stream_id

        top_level_query_ids = all_query_ids - selected_avatar_ids  # IDs of top-level queries (not in any FROMs)

        top_calc_stream_to_result_and_alias_map: Dict[str, Tuple[AvatarId, str]] = {}
        for translated_flat_query in translated_multi_query.iter_queries():
            query_froms = translated_flat_query.joined_from.froms
            is_top_level = translated_flat_query.id in top_level_query_ids
            is_bottom_level = bool(query_froms) and not any(from_obj.id in all_query_ids for from_obj in query_froms)
            if translated_flat_query.is_empty():
                raise exc.EmptyQuery()
            result_id = translated_flat_query.id
            required_avatar_ids = frozenset(translated_flat_query.joined_from.iter_ids())

            is_bottom_source_db = is_bottom_level and level_type == ExecutionLevel.source_db
            is_empty_source = len(required_avatar_ids) == 0
            is_multi_source = len(required_avatar_ids) > 1

            if is_bottom_source_db or is_empty_source or is_multi_source:
                # All of these cases require a `JoinOp`:
                # 1. JoinOp is required before the first-level query in source_db.
                #    The SELECT is made directly from the inter-joined source avatars.
                # 2. SELECT containing only constants or data-independent functions (see below)
                # 3. Multiple FROMs - real JOIN is required here

                root_avatar_id: Optional[str] = translated_flat_query.joined_from.root_from_id

                if is_empty_source:
                    # FIXME: This is quite a dirty hack and should be refactored in the long term
                    # This is a special case when the query selects only constants,
                    # such as with zero-dimensional LODs,
                    # and is used a a sort of glue that binds these LOD sub-queries.
                    # In this case the query must return exactly one line
                    # (it should not have a FROM clause), and yet it still needs
                    # a data source for all the meta stuff
                    # (connection instance, SQL compiler, etc.).
                    # For this we use the `use_empty_source` flag to indicate that,
                    # even though we have passed a list of sources to the JoinOp,
                    # the real FROM clause should be empty.

                    root_avatar_id = next(iter(translated_multi_query.get_base_root_from_ids()))
                    required_avatar_ids = {root_avatar_id}

                assert root_avatar_id is not None
                source_stream_ids = {result_to_stream_id_map[str_res_id] for str_res_id in required_avatar_ids}
                op = JoinOp(
                    source_stream_ids=source_stream_ids,
                    dest_stream_id=make_id(),
                    join_on_expressions=translated_flat_query.join_on,
                    root_avatar_id=root_avatar_id,
                    use_empty_source=is_empty_source,
                )
                operations.append(op)
                calc_op_source_stream_id = op.dest_stream_id

            else:
                # Query requires a single source
                source_result_id = next(iter(required_avatar_ids))
                calc_op_source_stream_id = result_to_stream_id_map[source_result_id]

            bi_query = self.make_bi_query(
                translated_flat_query=translated_flat_query,
                is_top_level=is_top_level,
                distinct=distinct,
                row_count_hard_limit=row_count_hard_limit,
            )
            op = CalcOp(
                result_id=result_id,
                source_stream_id=calc_op_source_stream_id,
                alias=translated_flat_query.alias,
                bi_query=bi_query,
                dest_stream_id=result_to_stream_id_map[result_id],
            )
            assert isinstance(op, CalcOp)  # for typing
            operations.append(op)
            # Save stream IDs for top-level queries
            if is_top_level:
                top_calc_stream_to_result_and_alias_map[op.dest_stream_id] = (op.result_id, op.alias)

        # 3. Operations for downloading data from the top-level queries
        output_stream_to_result_map: Dict[str, AvatarId] = {}
        output_stream_aliases: Dict[str, str] = {}
        for top_calc_stream_id, (result_id, alias) in top_calc_stream_to_result_and_alias_map.items():
            op = DownloadOp(
                source_stream_id=top_calc_stream_id,
                dest_stream_id=make_id(),
                row_count_hard_limit=row_count_hard_limit,
            )
            assert isinstance(op, DownloadOp)  # for typing
            LOGGER.info(
                f"Adding DownloadOp for query {result_id} with "
                f"input stream {op.source_stream_id} and output stream {op.dest_stream_id}"
            )
            operations.append(op)
            output_stream_to_result_map[op.dest_stream_id] = result_id
            output_stream_aliases[op.dest_stream_id] = alias

        processor = await self._get_data_processor(level_type=level_type, role=role)

        output_stream_list = await processor.run(
            operations=operations,
            streams=streams,
            output_stream_ids=set(output_stream_to_result_map),
        )
        output_streams_by_result_id = {output_stream_to_result_map[stream.id]: stream for stream in output_stream_list}
        return output_streams_by_result_id, output_stream_aliases

    def _make_source_db_input_streams(
        self,
        *,
        role: DataSourceRole,
        translated_multi_query: TranslatedMultiQueryBase,
        from_subquery: bool,
        subquery_limit: Optional[int],
    ) -> Tuple[Dict[AvatarId, AbstractStream], Dict[AvatarId, str]]:
        base_avatar_ids = translated_multi_query.get_base_root_from_ids()
        required_avatar_ids: list[str] = [from_id for from_id in translated_multi_query.get_base_froms()]
        required_avatar_ids = sorted(set(required_avatar_ids + list(base_avatar_ids)))

        prep_component_manager = DefaultPreparedComponentManager(
            dataset=self._dataset,
            role=role,
            us_entry_buffer=self._us_manager.get_entry_buffer(),
        )

        streams_by_result_id: Dict[AvatarId, AbstractStream] = {}
        stream_aliases: dict[str, str] = {}
        for avatar_id in required_avatar_ids:
            alias = self._avatar_alias_mapper(avatar_id)  # type: ignore  # TODO: fix
            prep_src_info = prep_component_manager.get_prepared_source(
                avatar_id=avatar_id,
                alias=alias,
                from_subquery=from_subquery,
                subquery_limit=subquery_limit,
            )
            result_id = avatar_id
            stream = DataSourceVS(
                id=make_id(),
                alias=alias,
                result_id=result_id,
                names=prep_src_info.col_names,
                user_types=prep_src_info.user_types,
                prep_src_info=prep_src_info,
            )
            streams_by_result_id[result_id] = stream
            stream_aliases[stream.id] = stream.alias

        id_map_for_log = {result_id: stream.id for result_id, stream in streams_by_result_id.items()}
        LOGGER.info(f"Generated base streams for result IDs: {id_map_for_log}")
        LOGGER.info(f"Using aliases: {stream_aliases}")
        assert len(streams_by_result_id) == len(stream_aliases)
        return streams_by_result_id, stream_aliases

    async def execute_query(self, exec_info: QueryExecutionInfo) -> ExecutedQuery:
        translated_multi_query: TranslatedMultiQueryBase = exec_info.translated_multi_query

        top_queries = translated_multi_query.get_top_queries()
        assert len(top_queries) == 1
        top_query = top_queries[0]

        distinct = top_query.meta.query_type == QueryType.distinct
        empty_query_mode = top_query.meta.empty_query_mode

        source_db_multi_query = translated_multi_query.for_level_type(ExecutionLevel.source_db)
        compeng_multi_query = translated_multi_query.for_level_type(ExecutionLevel.compeng)

        compeng_is_top = not compeng_multi_query.is_empty()

        from_subquery = top_query.meta.from_subquery  # FIXME
        subquery_limit = top_query.meta.subquery_limit  # FIXME
        row_count_hard_limit = top_query.meta.row_count_hard_limit

        streams_by_result_id, stream_aliases = self._make_source_db_input_streams(
            role=exec_info.role,
            translated_multi_query=translated_multi_query,
            from_subquery=from_subquery,
            subquery_limit=subquery_limit,
        )

        # Process data in source_db
        try:
            streams_by_result_id, stream_aliases = await self._process_multi_query(  # type: ignore  # TODO: fix
                level_type=ExecutionLevel.source_db,
                translated_multi_query=source_db_multi_query,
                streams_by_result_id=streams_by_result_id,
                stream_aliases=stream_aliases,
                distinct=False if compeng_is_top else distinct,
                upload_inputs=False,
                role=exec_info.role,
                row_count_hard_limit=row_count_hard_limit,
            )
        except exc.EmptyQuery:
            if empty_query_mode == EmptyQueryMode.error:
                raise
            empty_rows: List[TBIDataRow]
            if empty_query_mode == EmptyQueryMode.empty:
                empty_rows = []
            elif empty_query_mode == EmptyQueryMode.empty_row:
                empty_rows = [[]]
            else:
                raise ValueError(empty_query_mode)
            return ExecutedQuery(
                rows=empty_rows,
                meta=ExecutedQueryMetaInfo.from_trans_meta(
                    trans_meta=top_query.meta,
                    debug_query=None,
                ),
            )

        source_db_queries = [
            out_s.meta.query  # type: ignore  # TODO: fix
            for _, out_s in sorted(streams_by_result_id.items())
            if out_s.meta.query is not None and out_s.meta.pass_db_query_to_user  # type: ignore  # TODO: fix
        ]
        query_for_response = "\n;\n\n".join(source_db_queries) or None

        # Process data in compeng
        async with self._compeng_semaphore:
            streams_by_result_id, stream_aliases = await self._process_multi_query(  # type: ignore  # TODO: fix
                level_type=ExecutionLevel.compeng,
                translated_multi_query=compeng_multi_query,
                streams_by_result_id=streams_by_result_id,
                stream_aliases=stream_aliases,
                distinct=distinct if compeng_is_top else False,
                upload_inputs=True,
                role=exec_info.role,
                row_count_hard_limit=row_count_hard_limit,
            )

        streams = [stream for stream in streams_by_result_id.values()]
        assert len(streams) == 1, f"There must be exactly one output data stream, got {len(streams)}"
        one_stream = streams[0]

        assert isinstance(one_stream, DataStreamAsync)
        rows = await one_stream.data.all()

        executed_query = ExecutedQuery(
            rows=rows,
            meta=ExecutedQueryMetaInfo.from_trans_meta(
                trans_meta=top_query.meta,
                debug_query=query_for_response,
                target_connection_ids=set(
                    target_conn.uuid for target_conn in exec_info.target_connections if target_conn.uuid is not None
                ),
            ),
        )
        return executed_query

lib/dl_query_processing/dl_query_processing/execution/executor.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc

import attr

from dl_core.query.bi_query import BIQuery
from dl_query_processing.execution.exec_info import QueryExecutionInfo
from dl_query_processing.execution.primitives import ExecutedQuery
from dl_query_processing.translation.primitives import TranslatedFlatQuery


@attr.s
class QueryExecutorBase(abc.ABC):
    def make_bi_query(
        self,
        *,
        translated_flat_query: TranslatedFlatQuery,
        is_top_level: bool = False,
        distinct: bool = False,
        row_count_hard_limit: int,
    ) -> BIQuery:
        limit = translated_flat_query.limit

        # Limit query size if hard limit is set to minimize the amount of downloaded data
        if is_top_level:  # only for queries that will be downloaded
            if limit is None or limit > row_count_hard_limit:
                limit = row_count_hard_limit + 1

        bi_query = BIQuery(
            select_expressions=translated_flat_query.select,
            group_by_expressions=translated_flat_query.group_by,
            order_by_expressions=translated_flat_query.order_by,
            dimension_filters=translated_flat_query.where,
            measure_filters=translated_flat_query.having,
            limit=limit,
            offset=translated_flat_query.offset,
            distinct=distinct if is_top_level else False,
        )

        return bi_query

    @abc.abstractmethod
    async def execute_query(self, exec_info: QueryExecutionInfo) -> ExecutedQuery:
        raise NotImplementedError

lib/dl_query_processing/dl_query_processing/execution/executor_base.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    List,
    Optional,
    Type,
    TypeVar,
)

import attr

from dl_constants.types import TBIDataRow
from dl_core.utils import attrs_evolve_to_subclass
from dl_query_processing.translation.primitives import TranslatedQueryMetaInfo


_META_TV = TypeVar("_META_TV", bound="ExecutedQueryMetaInfo")


@attr.s
class ExecutedQueryMetaInfo(TranslatedQueryMetaInfo):
    debug_query: Optional[str] = attr.ib(kw_only=True, default=None)
    target_connection_ids: set[str] = attr.ib(kw_only=True, factory=set)

    @classmethod
    def from_trans_meta(
        cls: Type[_META_TV],
        trans_meta: TranslatedQueryMetaInfo,
        debug_query: Optional[str] = None,
        target_connection_ids: Optional[set[str]] = None,
    ) -> _META_TV:
        return attrs_evolve_to_subclass(
            cls=cls,
            inst=trans_meta,
            debug_query=debug_query,
            target_connection_ids=target_connection_ids or set(),
        )


@attr.s
class ExecutedQuery:
    rows: List[TBIDataRow] = attr.ib(kw_only=True)
    meta: ExecutedQueryMetaInfo = attr.ib(kw_only=True, factory=ExecutedQueryMetaInfo)

    def clone(self, **updates: Any) -> ExecutedQuery:
        return attr.evolve(self, **updates)

lib/dl_query_processing/dl_query_processing/execution/primitives.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
    List,
    Optional,
)

import attr

from dl_constants.enums import QueryBlockPlacementType
from dl_query_processing.base_specs.dimensions import (
    DimensionSpec,
    DimensionValueSpec,
)
from dl_query_processing.enums import (
    EmptyQueryMode,
    GroupByPolicy,
    QueryType,
)
from dl_query_processing.legend.field_legend import Legend


@attr.s(frozen=True)
class BlockPlacement:
    type: ClassVar[QueryBlockPlacementType]


@attr.s(frozen=True)
class RootBlockPlacement(BlockPlacement):
    type = QueryBlockPlacementType.root


@attr.s(frozen=True)
class AfterBlockPlacement(BlockPlacement):
    type = QueryBlockPlacementType.after

    dimension_values: Optional[List[DimensionValueSpec]] = attr.ib(default=None)


@attr.s(frozen=True)
class DispersedAfterBlockPlacement(BlockPlacement):
    type = QueryBlockPlacementType.dispersed_after

    parent_dimensions: List[DimensionSpec] = attr.ib(default=None)
    child_dimensions: List[DimensionSpec] = attr.ib(default=None)


@attr.s(frozen=True)
class BlockSpec:
    block_id: int = attr.ib(kw_only=True)
    parent_block_id: Optional[int] = attr.ib(kw_only=True)  # FIXME: Maybe remove this altogether
    placement: BlockPlacement = attr.ib(kw_only=True, factory=RootBlockPlacement)
    legend_item_ids: List[int] = attr.ib(kw_only=True, default=None)
    legend: Legend = attr.ib(kw_only=True)
    group_by_policy: GroupByPolicy = attr.ib(kw_only=True, default=GroupByPolicy.force)
    limit: Optional[int] = attr.ib(kw_only=True, default=None)
    offset: Optional[int] = attr.ib(kw_only=True, default=None)
    query_type: QueryType = attr.ib(kw_only=True, default=QueryType.result)
    row_count_hard_limit: Optional[int] = attr.ib(kw_only=True, default=None)
    disable_rls: bool = attr.ib(kw_only=True, default=False)
    ignore_nonexistent_filters: bool = attr.ib(kw_only=True, default=True)
    allow_measure_fields: bool = attr.ib(kw_only=True, default=True)
    empty_query_mode: EmptyQueryMode = attr.ib(kw_only=True, default=EmptyQueryMode.error)

    def clone(self, **updates: Any) -> BlockSpec:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class BlockLegendMeta:
    limit: Optional[int] = attr.ib(kw_only=True, default=None)
    offset: Optional[int] = attr.ib(kw_only=True, default=None)
    row_count_hard_limit: Optional[int] = attr.ib(kw_only=True, default=None)

    def clone(self, **updates: Any) -> BlockLegendMeta:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class BlockLegend:
    blocks: List[BlockSpec] = attr.ib(kw_only=True)
    meta: BlockLegendMeta = attr.ib(kw_only=True, factory=BlockLegendMeta)

    def clone(self, **updates: Any) -> BlockLegend:
        return attr.evolve(self, **updates)

lib/dl_query_processing/dl_query_processing/legend/block_legend.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from typing import (
    Any,
    Collection,
    Dict,
    Iterator,
    List,
    Optional,
    Sequence,
)

import attr

from dl_constants.enums import (
    FieldRole,
    FieldType,
    FieldVisibility,
    LegendItemType,
    OrderDirection,
    RangeType,
    UserDataType,
    WhereClauseOperation,
)
from dl_constants.internal_constants import (
    DIMENSION_NAME_ID,
    DIMENSION_NAME_TITLE,
    MEASURE_NAME_ID,
    MEASURE_NAME_TITLE,
    PLACEHOLDER_ID,
    PLACEHOLDER_TITLE,
)
from dl_query_processing.base_specs.dimensions import DimensionValueSpec
import dl_query_processing.exc


@attr.s(frozen=True)
class RoleSpec:
    role: FieldRole = attr.ib(kw_only=True, default=FieldRole.row)


@attr.s(frozen=True)
class OrderByRoleSpec(RoleSpec):
    direction: OrderDirection = attr.ib(kw_only=True)


@attr.s(frozen=True)
class FilterRoleSpec(RoleSpec):
    operation: WhereClauseOperation = attr.ib(kw_only=True)
    values: List[Any] = attr.ib(kw_only=True)


@attr.s(frozen=True)
class ParameterRoleSpec(RoleSpec):
    value: Any = attr.ib(kw_only=True)


@attr.s(frozen=True)
class RangeRoleSpec(RoleSpec):
    range_type: RangeType = attr.ib(kw_only=True)


@attr.s(frozen=True)
class DimensionRoleSpec(RoleSpec):
    visibility: FieldVisibility = attr.ib(kw_only=True, default=FieldVisibility.visible)


@attr.s(frozen=True)
class RowRoleSpec(DimensionRoleSpec):
    pass


@attr.s(frozen=True)
class TemplateRoleSpec(DimensionRoleSpec):
    template: Optional[str] = attr.ib(kw_only=True, default=None)


@attr.s(frozen=True)
class TreeRoleSpec(DimensionRoleSpec):
    level: int = attr.ib(kw_only=True)
    prefix: list = attr.ib(kw_only=True)
    dimension_values: List[DimensionValueSpec] = attr.ib(kw_only=True, factory=list)


@attr.s(frozen=True)
class ObjSpec:
    item_type: LegendItemType = attr.ib(kw_only=True)


@attr.s(frozen=True)
class MeasureNameObjSpec(ObjSpec):
    item_type: LegendItemType = attr.ib(kw_only=True, default=LegendItemType.measure_name)


@attr.s(frozen=True)
class DimensionNameObjSpec(ObjSpec):
    item_type: LegendItemType = attr.ib(kw_only=True, default=LegendItemType.dimension_name)


@attr.s(frozen=True)
class PlaceholderObjSpec(ObjSpec):
    item_type: LegendItemType = attr.ib(kw_only=True, default=LegendItemType.placeholder)


@attr.s(frozen=True)
class FieldObjSpec(ObjSpec):
    item_type: LegendItemType = attr.ib(kw_only=True, default=LegendItemType.field)
    id: str = attr.ib(kw_only=True)  # Field ID
    title: str = attr.ib(kw_only=True)  # Field title


@attr.s(frozen=True)
class LegendItem:
    legend_item_id: int = attr.ib(kw_only=True)
    obj: ObjSpec = attr.ib(kw_only=True)
    role_spec: RoleSpec = attr.ib(kw_only=True, factory=RoleSpec)
    data_type: UserDataType = attr.ib(kw_only=True)
    field_type: FieldType = attr.ib(kw_only=True)
    block_id: Optional[int] = attr.ib(kw_only=True, default=None)

    @property
    def item_type(self) -> LegendItemType:
        return self.obj.item_type

    @property
    def id(self) -> str:
        # TODO: Remove
        if self.obj.item_type == LegendItemType.measure_name:
            return MEASURE_NAME_ID
        if self.obj.item_type == LegendItemType.dimension_name:
            return DIMENSION_NAME_ID
        if self.obj.item_type == LegendItemType.placeholder:
            return PLACEHOLDER_ID
        assert isinstance(self.obj, FieldObjSpec)
        return self.obj.id

    @property
    def title(self) -> str:
        # TODO: Remove
        if self.obj.item_type == LegendItemType.measure_name:
            return MEASURE_NAME_TITLE
        if self.obj.item_type == LegendItemType.dimension_name:
            return DIMENSION_NAME_TITLE
        if self.obj.item_type == LegendItemType.placeholder:
            return PLACEHOLDER_TITLE
        assert isinstance(self.obj, FieldObjSpec)
        return self.obj.title

    def clone(self, **updates: Any) -> LegendItem:
        return attr.evolve(self, **updates)


NON_SELECTABLE_ROLES = {
    FieldRole.info,
    FieldRole.order_by,
    FieldRole.filter,
    FieldRole.parameter,
}
SELECTABLE_ROLES = set(FieldRole) - NON_SELECTABLE_ROLES


@attr.s
class Legend:
    _items: List[LegendItem] = attr.ib(kw_only=True)
    _liid_to_item: Dict[int, LegendItem] = attr.ib(init=False)
    _liid_to_field_id_map: Dict[int, str] = attr.ib(init=False)
    _field_id_to_liid_list_map: Dict[str, List[int]] = attr.ib(init=False)

    def __attrs_post_init__(self) -> None:
        self._liid_to_item = {}
        self._liid_to_field_id_map = {}
        self._field_id_to_liid_list_map = defaultdict(list)
        self._repopulate_mappings()

    @property
    def items(self) -> Sequence[LegendItem]:
        return self._items

    def get_item(self, legend_item_id: int) -> LegendItem:
        try:
            return self._liid_to_item[legend_item_id]
        except KeyError:
            raise dl_query_processing.exc.LegendItemReferenceError(f"Unknown legend item: {legend_item_id}")

    def add_item(self, item: LegendItem) -> None:
        self._items.append(item)
        self._repopulate_mappings()

    def _repopulate_mappings(self) -> None:
        self._liid_to_item.clear()
        self._liid_to_field_id_map.clear()
        self._field_id_to_liid_list_map.clear()
        for item in self._items:
            self._liid_to_item[item.legend_item_id] = item
            self._liid_to_field_id_map[item.legend_item_id] = item.id
            self._field_id_to_liid_list_map[item.id].append(item.legend_item_id)

    def get_measure_name_legend_item_ids(self) -> List[int]:
        return self.field_id_to_leg_item_id_list(MEASURE_NAME_ID)

    def get_dimension_name_legend_item_ids(self) -> List[int]:
        return self.field_id_to_leg_item_id_list(DIMENSION_NAME_ID)

    def field_id_to_leg_item_id_list(self, field_id: str) -> List[int]:
        return self._field_id_to_liid_list_map[field_id]

    def __iter__(self) -> Iterator[LegendItem]:
        return iter(self._items)

    def list_for_role(self, role: FieldRole) -> Sequence[LegendItem]:
        return tuple(item for item in self._items if item.role_spec.role == role)

    def limit_to_roles(self, roles: Collection[FieldRole]) -> Legend:
        return Legend(items=[item for item in self.items if item.role_spec.role in roles])

    def limit_to_block(self, block_id: int) -> Legend:
        return Legend(
            items=[item for item in self.items if item.block_id is None or item.block_id == block_id],
        )

    def list_selectable_items(self) -> List[LegendItem]:
        """List the items that can be selected from the data source"""
        items = [
            item
            for item in self.items
            if item.role_spec.role in SELECTABLE_ROLES and isinstance(item.obj, FieldObjSpec)
        ]
        return items

    def list_streamable_items(self) -> Sequence[LegendItem]:
        """List streamable items + special derived ones like templates"""
        items = [
            item
            for item in self.items
            if item.role_spec.role in SELECTABLE_ROLES and isinstance(item.obj, (FieldObjSpec, PlaceholderObjSpec))
        ]
        return items

lib/dl_query_processing/dl_query_processing/legend/field_legend.py
datalens-tech/datalens-backend
import abc
from typing import (
    Any,
    Iterable,
    List,
    Mapping,
    Optional,
    Sequence,
    Set,
    Type,
)

import attr

from dl_constants.enums import QueryBlockPlacementType
from dl_query_processing.legend.block_legend import (
    AfterBlockPlacement,
    BlockPlacement,
    DimensionValueSpec,
    DispersedAfterBlockPlacement,
)
from dl_query_processing.merging.primitives import (
    MergedQueryDataRow,
    MergedQueryDataRowIterable,
)
from dl_query_processing.postprocessing.primitives import PostprocessedValue


@attr.s(frozen=True)
class BinaryStreamMerger(abc.ABC):
    """
    Base class for merger of two data streams: parent and child.
    Used to convert data streams from multiple DB queries into a single output stream.
    """

    placement: BlockPlacement = attr.ib(kw_only=True)

    @abc.abstractmethod
    def merge_two_streams(
        self,
        parent_stream: MergedQueryDataRowIterable,
        child_stream: MergedQueryDataRowIterable,
    ) -> MergedQueryDataRowIterable:
        raise NotImplementedError


@attr.s
class RowMatcher:
    """
    Helper object for optimizing the matching of data rows to specified dimension values
    """

    _dimension_values: List[DimensionValueSpec] = attr.ib(kw_only=True)
    # internal stuff
    _required_legend_item_ids: Set[int] = attr.ib(init=False)
    _match_dim_values: List[PostprocessedValue] = attr.ib(init=False)
    _prev_legend_item_ids: Optional[Iterable[int]] = attr.ib(init=False, default=None)
    _dimension_mask: List[int] = attr.ib(init=False, factory=list)
    _row_has_req_items: bool = attr.ib(init=False, default=False)

    @_required_legend_item_ids.default
    def _make_required_legend_item_ids(self) -> Set[int]:
        return {dim_spec.legend_item_id for dim_spec in self._dimension_values}

    @_match_dim_values.default
    def _make_match_dim_values(self) -> List[PostprocessedValue]:
        return [dim_spec.value for dim_spec in self._dimension_values]

    def match_row(self, row: MergedQueryDataRow) -> bool:
        if row.legend_item_ids is not self._prev_legend_item_ids:
            # Row configuration has changed from the previous row.
            # This is an optimization to avoid generation of dimension_mask for every row.
            if set(row.legend_item_ids).issuperset(self._required_legend_item_ids):
                self._row_has_req_items = True
                self._dimension_mask = [
                    row.legend_item_ids.index(dim_spec.legend_item_id) for dim_spec in self._dimension_values
                ]
            else:
                # Row (and subsequent rows) doesn't contain the required legend items,
                # so don't try to match it.
                self._row_has_req_items = False
                self._dimension_mask = []

            self._prev_legend_item_ids = row.legend_item_ids

        if self._row_has_req_items:  # possible match candidate
            dim_values = [row.data[idx] for idx in self._dimension_mask]
            return dim_values == self._match_dim_values

        return False


class AfterBinaryStreamMerger(BinaryStreamMerger):
    """
    Handle the "after" placement type:
    1. If `dimension_values` is None: simply concatenate the two streams.
    2. If `dimension_values` is not None: insert child stream after
       matching values are found in parent stream.
    """

    @property
    def after_placement(self) -> AfterBlockPlacement:
        assert isinstance(self.placement, AfterBlockPlacement)
        return self.placement

    def _merge_two_streams_no_dims(
        self,
        parent_stream: MergedQueryDataRowIterable,
        child_stream: MergedQueryDataRowIterable,
    ) -> MergedQueryDataRowIterable:
        """
        Just concatenate them.
        """
        yield from parent_stream
        yield from child_stream

    def _merge_two_streams_with_dims(
        self,
        parent_stream: MergedQueryDataRowIterable,
        child_stream: MergedQueryDataRowIterable,
    ) -> MergedQueryDataRowIterable:
        """
        Inspect parent stream values and insert after match.
        """

        dimension_values = self.after_placement.dimension_values
        assert dimension_values is not None
        row_matcher = RowMatcher(dimension_values=dimension_values)
        inserted = False
        previous_row_matched = False

        for row in parent_stream:
            if inserted:
                # Nothing to inspect anymore; yield the rest of the stream
                yield row
                continue

            current_row_matched = row_matcher.match_row(row)
            if previous_row_matched and not current_row_matched:
                yield from child_stream
                inserted = True

            yield row
            previous_row_matched = current_row_matched

        if previous_row_matched and not inserted:
            # We got to the end of the parent stream,
            # but didn't yield the child stream yet
            yield from child_stream

    def merge_two_streams(
        self,
        parent_stream: MergedQueryDataRowIterable,
        child_stream: MergedQueryDataRowIterable,
    ) -> MergedQueryDataRowIterable:
        if self.after_placement.dimension_values is None:
            yield from self._merge_two_streams_no_dims(parent_stream, child_stream)
        else:
            yield from self._merge_two_streams_with_dims(parent_stream, child_stream)


class DispersedAfterBinaryStreamMerger(BinaryStreamMerger):
    """
    In a way similar to the "after" placement type,
    but inserts data rows from `child_stream` separately, one by one
    after the matching rows from parent_stream.
    """

    @property
    def disp_after_placement(self) -> DispersedAfterBlockPlacement:
        assert isinstance(self.placement, DispersedAfterBlockPlacement)
        return self.placement

    def merge_two_streams(
        self,
        parent_stream: MergedQueryDataRowIterable,
        child_stream: MergedQueryDataRowIterable,
    ) -> MergedQueryDataRowIterable:
        dim_indices: Optional[tuple[int, ...]] = None  # This value will be redefined
        dim_values: tuple[Any, ...]
        # Create mapping of child stream values:
        # { <dimension_value_vector>: child_stream_row }
        child_row_map: dict[tuple[Any, ...], MergedQueryDataRow] = {}
        prev_legend_item_ids: Optional[Sequence[int]] = None
        for row in child_stream:
            if row.legend_item_ids is not prev_legend_item_ids:
                # Recreate indices because the legend has changed
                dim_indices = tuple(
                    row.legend_item_ids.index(dim_spec.legend_item_id)
                    for dim_spec in self.disp_after_placement.child_dimensions
                )
                prev_legend_item_ids = row.legend_item_ids

            assert dim_indices is not None
            dim_values = tuple(row.data[idx] for idx in dim_indices)
            child_row_map[dim_values] = row

        # Iterate over parent stream and append child value after each match
        prev_dim_values: Optional[tuple[Any, ...]] = None
        fast_forward_row = False
        expected_parent_dim_set = {dim_spec.legend_item_id for dim_spec in self.disp_after_placement.parent_dimensions}
        for row in parent_stream:
            if row.legend_item_ids is not prev_legend_item_ids:
                prev_legend_item_ids = row.legend_item_ids
                # Check if this row even matches the monitoring criteria
                if not set(row.legend_item_ids).issuperset(expected_parent_dim_set):
                    # fast-forward all rows with this legend
                    fast_forward_row = True

                else:
                    # Recreate indices and values because the legend has changed
                    fast_forward_row = False
                    dim_indices = tuple(
                        row.legend_item_ids.index(dim_spec.legend_item_id)
                        for dim_spec in self.disp_after_placement.parent_dimensions
                    )

            if fast_forward_row:
                yield row
                continue

            assert dim_indices is not None
            dim_values = tuple(row.data[idx] for idx in dim_indices)
            if dim_values != prev_dim_values and prev_dim_values is not None:
                # Values of monitored dimensions have changed.
                # This is where we add the corresponding value from the child stream
                child_stream_row = child_row_map.pop(prev_dim_values, None)
                if child_stream_row is not None:
                    yield child_stream_row

            yield row
            prev_dim_values = dim_values

        # Dump the rest of child rows
        for row in child_row_map.values():
            yield row


def make_binary_merger(placement: BlockPlacement) -> BinaryStreamMerger:
    """
    Factory of BinaryStreamMerger instances.
    """

    placement_type_map: Mapping[QueryBlockPlacementType, Type[BinaryStreamMerger]] = {
        QueryBlockPlacementType.after: AfterBinaryStreamMerger,
        QueryBlockPlacementType.dispersed_after: DispersedAfterBinaryStreamMerger,
    }
    binary_merger_cls = placement_type_map[placement.type]
    binary_merger = binary_merger_cls(placement=placement)
    return binary_merger

lib/dl_query_processing/dl_query_processing/merging/binary_merger.py
datalens-tech/datalens-backend
from typing import (
    List,
    Optional,
)

from dl_query_processing.legend.block_legend import RootBlockPlacement
from dl_query_processing.merging.binary_merger import make_binary_merger
from dl_query_processing.merging.primitives import (
    MergedQueryBlockMetaInfo,
    MergedQueryDataRow,
    MergedQueryDataRowIterable,
    MergedQueryDataStream,
    MergedQueryMetaInfo,
)
from dl_query_processing.postprocessing.primitives import (
    PostprocessedQueryBlock,
    PostprocessedQueryUnion,
)


class DataStreamMerger:
    """
    Merges several postprocessed data streams into one
    vertically (rows from blocks are intermixed)
    and/or horizontally (row from one stream is combined with a row from another stream).
    """

    def normalize_stream(self, block: PostprocessedQueryBlock) -> MergedQueryDataRowIterable:
        for postprocessed_row in block.postprocessed_query.postprocessed_data:
            yield MergedQueryDataRow(data=postprocessed_row, legend_item_ids=block.legend_item_ids)

    def merge(self, postprocessed_query_union: PostprocessedQueryUnion) -> MergedQueryDataStream:
        legend_item_ids: Optional[List[int]] = None
        if len(postprocessed_query_union.blocks) == 1:
            legend_item_ids = postprocessed_query_union.blocks[0].legend_item_ids

        # Find the root block/stream
        root_block = next(
            iter(block for block in postprocessed_query_union.blocks if isinstance(block.placement, RootBlockPlacement))
        )
        data_stream: MergedQueryDataRowIterable = self.normalize_stream(root_block)

        merged_target_connection_ids: set[str] = set()
        for block in postprocessed_query_union.blocks:
            merged_target_connection_ids |= block.postprocessed_query.meta.target_connection_ids
            if block is not root_block:
                binary_merger = make_binary_merger(placement=block.placement)
                child_data_stream = self.normalize_stream(block)
                data_stream = binary_merger.merge_two_streams(
                    parent_stream=data_stream,
                    child_stream=child_data_stream,
                )

        merged_data_stream = MergedQueryDataStream(
            legend=postprocessed_query_union.legend,
            rows=data_stream,
            legend_item_ids=legend_item_ids,
            meta=MergedQueryMetaInfo(
                blocks=[
                    MergedQueryBlockMetaInfo(
                        block_id=block.block_id,
                        query_type=block.postprocessed_query.meta.query_type,
                        debug_query=block.postprocessed_query.meta.debug_query,
                    )
                    for block in postprocessed_query_union.blocks
                ],
                offset=postprocessed_query_union.offset,
                limit=postprocessed_query_union.limit,
                target_connection_ids=merged_target_connection_ids,
            ),
        )
        return merged_data_stream

lib/dl_query_processing/dl_query_processing/merging/merger.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Iterable,
    List,
    NamedTuple,
    Optional,
    Sequence,
)

import attr

from dl_query_processing.enums import QueryType
from dl_query_processing.legend.field_legend import Legend
from dl_query_processing.postprocessing.primitives import PostprocessedRow


@attr.s(frozen=True)
class MergedQueryBlockMetaInfo:
    block_id: int = attr.ib(kw_only=True)
    query_type: QueryType = attr.ib(kw_only=True)
    debug_query: Optional[str] = attr.ib(kw_only=True)


@attr.s(frozen=True)
class MergedQueryMetaInfo:
    blocks: List[MergedQueryBlockMetaInfo] = attr.ib(kw_only=True)
    offset: Optional[int] = attr.ib(kw_only=True, default=None)
    limit: Optional[int] = attr.ib(kw_only=True, default=None)
    target_connection_ids: set[str] = attr.ib(kw_only=True, factory=set)


class MergedQueryDataRow(NamedTuple):
    """This class should be as simple and as small as possible, so use NamedTuple instead of attrs"""

    data: PostprocessedRow
    legend_item_ids: Sequence[int]


MergedQueryDataRowIterable = Iterable[MergedQueryDataRow]


@attr.s(frozen=True)
class MergedQueryDataStream:
    rows: MergedQueryDataRowIterable = attr.ib(kw_only=True)
    legend: Legend = attr.ib(kw_only=True)
    legend_item_ids: Optional[Sequence[int]] = attr.ib(kw_only=True)  # in case there is only one stream, for v1
    meta: MergedQueryMetaInfo = attr.ib(kw_only=True)

    def clone(self, **updates: Any) -> MergedQueryDataStream:
        return attr.evolve(self, **updates)

lib/dl_query_processing/dl_query_processing/merging/primitives.py
datalens-tech/datalens-backend
import abc

import attr

from dl_core.fields import ResultSchema
from dl_query_processing.enums import ExecutionLevel
from dl_query_processing.multi_query.mutators.base import MultiQueryMutatorBase
from dl_query_processing.multi_query.mutators.splitter_based import SplitterMultiQueryMutator
from dl_query_processing.multi_query.mutators.winfunc_compeng import DefaultCompengMultiQueryMutator
from dl_query_processing.multi_query.splitters.prefiltered import PrefilteredFieldMultiQuerySplitter
from dl_query_processing.multi_query.splitters.query_fork import QueryForkQuerySplitter


@attr.s
class MultiQueryMutatorFactoryBase(abc.ABC):
    result_schema: ResultSchema = attr.ib(kw_only=True)

    @abc.abstractmethod
    def get_mutators(self) -> list[MultiQueryMutatorBase]:
        raise NotImplementedError


@attr.s
class DefaultMultiQueryMutatorFactory(MultiQueryMutatorFactoryBase):
    def get_mutators(self) -> list[MultiQueryMutatorBase]:
        return [
            SplitterMultiQueryMutator(
                splitters=[
                    QueryForkQuerySplitter(),
                ]
            ),
            DefaultCompengMultiQueryMutator(),
        ]


@attr.s
class NoCompengMultiQueryMutatorFactory(MultiQueryMutatorFactoryBase):
    def get_mutators(self) -> list[MultiQueryMutatorBase]:
        return [
            SplitterMultiQueryMutator(
                splitters=[
                    QueryForkQuerySplitter(),
                ]
            ),
        ]


@attr.s
class SimpleFieldSplitterMultiQueryMutatorFactory(MultiQueryMutatorFactoryBase):
    def get_mutators(self) -> list[MultiQueryMutatorBase]:
        return [
            SplitterMultiQueryMutator(
                splitters=[
                    PrefilteredFieldMultiQuerySplitter(crop_to_level_type=ExecutionLevel.compeng),
                    QueryForkQuerySplitter(),
                ]
            ),
            DefaultCompengMultiQueryMutator(),
        ]

lib/dl_query_processing/dl_query_processing/multi_query/factory.py
datalens-tech/datalens-backend
import attr

from dl_query_processing.compilation.primitives import CompiledMultiQueryBase


@attr.s
class MultiQueryMutatorBase:
    """
    Applies mutations at multi-query level.
    """

    def mutate_multi_query(self, multi_query: CompiledMultiQueryBase) -> CompiledMultiQueryBase:
        raise NotImplementedError

lib/dl_query_processing/dl_query_processing/multi_query/mutators/base.py
datalens-tech/datalens-backend
import logging
from typing import (
    ClassVar,
    Optional,
    Sequence,
)

import attr

from dl_query_processing.compilation.primitives import CompiledMultiQueryBase
from dl_query_processing.multi_query.mutators.base import MultiQueryMutatorBase
from dl_query_processing.multi_query.splitters.base import MultiQuerySplitterBase
from dl_query_processing.multi_query.tools import (
    CompiledMultiQueryPatch,
    apply_query_patch,
    build_requirement_subtree,
)
from dl_query_processing.utils.name_gen import PrefixedIdGen


LOGGER = logging.getLogger(__name__)


@attr.s
class SplitterMultiQueryMutator(MultiQueryMutatorBase):
    """
    Mutator base on splitter objects that are applied to each flat query
    that is part of the multi-query.
    """

    _splitters: Sequence[MultiQuerySplitterBase] = attr.ib(kw_only=True)
    _do_log: bool = attr.ib(kw_only=True, default=True)

    _max_iterations: ClassVar[int] = 1000

    def log(self, msg: str, level: int = logging.INFO) -> None:
        if self._do_log:
            LOGGER.log(level=level, msg=msg)

    def mutate_multi_query(self, multi_query: CompiledMultiQueryBase) -> CompiledMultiQueryBase:
        query_id_gen = PrefixedIdGen("q")
        expr_id_gen = PrefixedIdGen("e")
        for splitter in self._splitters:
            skip_queries: set[str] = set()

            iteration_cnt = 0
            while True:
                # Prevent infinite loop
                if iteration_cnt >= self._max_iterations:
                    raise RuntimeError("Maximum multi-query iterations")

                # Iterate over the whole multi-query object each time
                # because we don't know, how its structure has been changed
                multi_query_patch: Optional[CompiledMultiQueryPatch] = None
                for query in multi_query.iter_queries():
                    # Skip queries that have already been handled
                    if query.id in skip_queries:
                        continue

                    requirement_subtree = build_requirement_subtree(multi_query, query.id)
                    multi_query_patch = splitter.split_query(
                        query=query,
                        requirement_subtree=requirement_subtree,
                        query_id_gen=query_id_gen,
                        expr_id_gen=expr_id_gen,
                    )
                    skip_queries.add(query.id)
                    if multi_query_patch is not None:
                        break

                if multi_query_patch is not None:
                    # We have a patch. This means that the multi-query has to be mutated
                    multi_query = apply_query_patch(multi_query=multi_query, patch=multi_query_patch)
                    iteration_cnt += 1
                else:
                    # No patch. This means that there is nothing left to mutate in all of the multi-query.
                    # Leave.
                    break

            if iteration_cnt > 0:
                self.log(
                    f"Applied mutation patches in {iteration_cnt} iterations "
                    f"for splitter {splitter.__class__.__name__}"
                )

        return multi_query

lib/dl_query_processing/dl_query_processing/multi_query/mutators/splitter_based.py
datalens-tech/datalens-backend
import attr

from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.inspect.expression import is_window_expression
from dl_query_processing.compilation.primitives import (
    CompiledMultiQuery,
    CompiledMultiQueryBase,
    CompiledQuery,
    SubqueryFromObject,
)
from dl_query_processing.enums import ExecutionLevel
from dl_query_processing.multi_query.mutators.base import MultiQueryMutatorBase
from dl_query_processing.multi_query.splitters.mask_based import MultiQuerySplitter
from dl_query_processing.multi_query.splitters.win_func import WinFuncQuerySplitter
from dl_query_processing.multi_query.tools import build_requirement_subtree
from dl_query_processing.utils.name_gen import PrefixedIdGen


@attr.s
class DefaultCompengMultiQueryMutator(MultiQueryMutatorBase):
    """
    Updates `level_type` of non-base sub-queries to `compeng`
    if `multi_query` contains window functions,
    """

    optimize_compeng_usage: bool = attr.ib(kw_only=True, default=True)

    def split_borderline_query(
        self,
        query: CompiledQuery,
        requirement_subtree: CompiledMultiQueryBase,
        query_id_gen: PrefixedIdGen,
        expr_id_gen: PrefixedIdGen,
    ) -> list[CompiledQuery]:
        """
        Split the borderline sub-query into two:
        the top one with the window functions,
        and the bottom one with everything else that the top one selects from
        """

        original_query_id = query.id
        splitter: MultiQuerySplitter = WinFuncQuerySplitter()
        multi_query_patch = splitter.split_query(
            query=query,
            requirement_subtree=requirement_subtree,
            query_id_gen=query_id_gen,
            expr_id_gen=expr_id_gen,
        )
        assert multi_query_patch is not None
        new_subqueries = list(multi_query_patch.patch_multi_query.iter_queries())
        assert len(new_subqueries) == 2

        top_query = next(q for q in new_subqueries if q.id == original_query_id)
        bottom_query = next(q for q in new_subqueries if q.id != original_query_id)
        assert bottom_query.id == top_query.joined_from.root_from_id
        top_query = top_query.clone(level_type=ExecutionLevel.compeng)
        return [top_query, bottom_query]

    def mutate_multi_query(self, multi_query: CompiledMultiQueryBase) -> CompiledMultiQueryBase:
        query_id_gen = PrefixedIdGen("qb")
        expr_id_gen = PrefixedIdGen("eb")

        # First determine if there is a need for compeng
        inspect_env = InspectionEnvironment()
        queries_with_win_funcs: set[str] = set()  # queries that contain WFs or depend on other queries that do
        borderline_queries: set[str] = set()  # queries with WF, whose sub-queries don't have WF
        for query in multi_query.iter_queries():
            query_has_winfuncs = any(
                is_window_expression(node=formula.formula_obj, env=inspect_env) for formula in query.all_formulas
            )
            if query_has_winfuncs:
                queries_with_win_funcs.add(query.id)

        if not queries_with_win_funcs:
            return multi_query

        def propagate_wf_dependency(query: CompiledQuery) -> bool:
            has_wf: bool = False
            if query.id in queries_with_win_funcs:
                has_wf = True

            any_children_with_wf = False
            for child_from in query.joined_from.froms:
                if isinstance(child_from, SubqueryFromObject):
                    child_query_id = child_from.query_id
                    child_has_wf = propagate_wf_dependency(multi_query.get_query_by_id(child_query_id))
                    if child_has_wf:
                        # If a query's child uses window functions, then the query itself uses window functions
                        any_children_with_wf = True
                        has_wf = True
                        queries_with_win_funcs.add(query.id)

            if has_wf and not any_children_with_wf:
                borderline_queries.add(query.id)

            return has_wf

        propagate_wf_dependency(multi_query.get_single_top_query())

        # Replace `level_type` in all of the bottom-level sub-queries
        updated_queries: list[CompiledQuery] = []
        base_from_ids = {from_obj.id for from_obj in multi_query.get_base_froms().values()}
        for query in multi_query.iter_queries():
            query_is_base = bool(set(query.joined_from.iter_ids()) & base_from_ids)
            this_iteration_updated_queries: list[CompiledQuery] = []
            if (
                # optimized version
                self.optimize_compeng_usage
                and query.id in queries_with_win_funcs
                # non-optimized version
                or not self.optimize_compeng_usage
                and not query_is_base
            ):
                if self.optimize_compeng_usage and query.id in borderline_queries:
                    # Split borderline queries so that only the part that starts with window functions
                    # goes to COMPENG. Everything below it should remain in source DB.
                    requirement_subtree = build_requirement_subtree(multi_query, query.id)
                    this_iteration_updated_queries.extend(
                        self.split_borderline_query(
                            query=query,
                            requirement_subtree=requirement_subtree,
                            query_id_gen=query_id_gen,
                            expr_id_gen=expr_id_gen,
                        )
                    )
                else:
                    query = query.clone(level_type=ExecutionLevel.compeng)
                    this_iteration_updated_queries.append(query)

            else:
                this_iteration_updated_queries.append(query)

            updated_queries.extend(this_iteration_updated_queries)

        multi_query = CompiledMultiQuery(queries=updated_queries)
        return multi_query

lib/dl_query_processing/dl_query_processing/multi_query/mutators/winfunc_compeng.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import Optional

from dl_query_processing.compilation.primitives import (
    CompiledMultiQueryBase,
    CompiledQuery,
)
from dl_query_processing.multi_query.tools import CompiledMultiQueryPatch
from dl_query_processing.utils.name_gen import PrefixedIdGen


class MultiQuerySplitterBase(abc.ABC):
    @abc.abstractmethod
    def split_query(
        self,
        query: CompiledQuery,
        requirement_subtree: CompiledMultiQueryBase,
        query_id_gen: PrefixedIdGen,
        expr_id_gen: PrefixedIdGen,
    ) -> Optional[CompiledMultiQueryPatch]:
        raise NotImplementedError

lib/dl_query_processing/dl_query_processing/multi_query/splitters/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from enum import Enum
import itertools
from typing import (
    Any,
    ClassVar,
    Collection,
    Optional,
    Sequence,
)

import attr

from dl_constants.enums import JoinType
import dl_formula.core.aux_nodes as formula_aux_nodes
import dl_formula.core.exc as formula_exc
from dl_formula.core.extract import NodeExtract
import dl_formula.core.fork_nodes as formula_fork_nodes
from dl_formula.core.index import NodeHierarchyIndex
import dl_formula.core.nodes as formula_nodes
import dl_formula.inspect.expression as inspect_expression
from dl_query_processing.compilation.primitives import (
    CompiledFormulaInfo,
    CompiledJoinOnFormulaInfo,
    CompiledMultiQuery,
    CompiledMultiQueryBase,
    CompiledQuery,
    FromColumn,
    FromObject,
    JoinedFromObject,
    SubqueryFromObject,
)
from dl_query_processing.compilation.query_meta import QueryMetaInfo
from dl_query_processing.enums import QueryPart
from dl_query_processing.multi_query.splitters.base import MultiQuerySplitterBase
from dl_query_processing.multi_query.tools import (
    CompiledMultiQueryPatch,
    remap_formula_obj_fields,
)
from dl_query_processing.utils.name_gen import PrefixedIdGen


class SubqueryType(Enum):
    window_func = "window_func"
    default = "default"
    generated_base = "generated_base"


@attr.s(frozen=True)
class FormulaSplitMask:
    """
    Base class for `AliasedFormulaSplitMask` without the `alias` field;
    not intended for widespread usage.
    See `AliasedFormulaSplitMask` for more info.
    """

    # Indicates which part (clause) of the original query the formula is an item of
    query_part: QueryPart = attr.ib(kw_only=True)

    # The index within list specified by `query_part`
    formula_list_idx: int = attr.ib(kw_only=True)

    # Specifies the node which should be cut out from the original query
    # and substituted for a subquery column
    outer_node_idx: NodeHierarchyIndex = attr.ib(kw_only=True)

    # Specifies the node which should be used in the new sub-query.
    # Can be the same as outer_*, but can point to its sub-node
    # If some auxiliary nodes in-between need to be removed
    # from the queries completely.
    inner_node_idx: NodeHierarchyIndex = attr.ib(kw_only=True)

    def __attrs_post_init__(self) -> None:
        assert self.inner_node_idx.startswith(self.outer_node_idx)


@attr.s(frozen=True)
class AliasedFormulaSplitMask(FormulaSplitMask):
    """
    Formula-splitting mask.
    Performs two functions:

    1. Is used to generate SELECT items for sub-queries being split off the main query.
       `query_part`, `formula_list_idx` and `inner_node_idx` indicate where to find the
       expression in the original query, while `alias` tells us how to label the new
       split-off formula.

    2. Is used to "cut out" an expression from the original query and replace it with
       a `Field` node that now refers to a subquery where the cut-out part
       of the original expression went.
       `query_part`, `formula_list_idx` and `outer_node_idx` indicate what to cut,
       and `alias` is the name of the new `Field` node.
    """

    alias: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class AddFormulaInfo:
    """
    Is used to specify a formula that is being added to a sub-query.
    Used for dimensions and for measures that are being added to be used in
    JOIN ON expressions.
    """

    alias: str = attr.ib(kw_only=True)
    expr: formula_nodes.FormulaItem = attr.ib(kw_only=True)
    from_ids: frozenset[str] = attr.ib(kw_only=True)
    is_group_by: bool = attr.ib(kw_only=True)

    def clone(self, **kwargs: Any) -> AddFormulaInfo:
        return attr.evolve(self, **kwargs)


@attr.s(frozen=True)
class QuerySplitMask:
    subquery_type: SubqueryType = attr.ib(kw_only=True)
    subquery_id: str = attr.ib(kw_only=True)
    formula_split_masks: tuple[AliasedFormulaSplitMask, ...] = attr.ib(kw_only=True)
    add_formulas: tuple[AddFormulaInfo, ...] = attr.ib(kw_only=True)
    filter_indices: frozenset[int] = attr.ib(kw_only=True)
    join_type: Optional[JoinType] = attr.ib(kw_only=True)
    joining_node: Optional[formula_fork_nodes.QueryForkJoiningBase] = attr.ib(kw_only=True)
    is_base: bool = attr.ib(kw_only=True, default=False)

    # calculated props
    group_by_count: int = attr.ib(init=False)

    @group_by_count.default
    def _make_group_by_count(self) -> int:
        non_const_cnt = 0
        for dim in self.add_formulas:
            if dim.is_group_by:
                non_const_cnt += 1
        return non_const_cnt

    @property
    def has_direct_equality_join(self) -> bool:
        """
        True if the join contains only simple equalities
        of column from one sub-query to the corresponding column from the other sub-query.
        """

        assert isinstance(self.joining_node, formula_fork_nodes.QueryForkJoiningWithList)
        return all(
            isinstance(condition, formula_fork_nodes.SelfEqualityJoinCondition)
            for condition in self.joining_node.condition_list
        )

    def clone(self, **kwargs: Any) -> QuerySplitMask:
        return attr.evolve(self, **kwargs)


@attr.s
class MultiQuerySplitter(MultiQuerySplitterBase):
    parent_query_parts_for_splitting: ClassVar[tuple[QueryPart, ...]] = (
        QueryPart.select,
        QueryPart.order_by,
        QueryPart.filters,
    )

    @abc.abstractmethod
    def get_split_masks(
        self,
        query: CompiledQuery,
        expr_id_gen: PrefixedIdGen,
        query_id_gen: PrefixedIdGen,
    ) -> list[QuerySplitMask]:
        raise NotImplementedError

    def mutate_split_node(self, node: formula_nodes.FormulaItem) -> formula_nodes.FormulaItem:
        """
        Optionally mutate the node that was split off by the mask.
        By default, return it without any changes.
        """

        return node

    def mutate_cropped_query(self, query: CompiledQuery) -> CompiledQuery:
        """
        Optionally mutate the original query after subqueries have been split off it.
        """

        return query

    def _get_field_to_from_id_map(
        self,
        requirement_subtree: CompiledMultiQueryBase,
    ) -> dict[str, str]:
        result: dict[str, str] = {}
        for top_query in requirement_subtree.get_top_queries():
            for select_formula in top_query.select:
                result[select_formula.not_none_alias] = top_query.id

        return result

    def _generate_subquery_for_mask(
        self,
        query: CompiledQuery,
        split_mask: QuerySplitMask,
    ) -> CompiledQuery:
        """
        Generate sub-query described by a `QuerySplitMask`.
        """

        select: list[CompiledFormulaInfo] = []
        group_by: list[CompiledFormulaInfo] = []
        select_by_alias: dict[str, CompiledFormulaInfo] = {}
        for formula_split_mask in split_mask.formula_split_masks:
            formula_list = query.get_formula_list(formula_split_mask.query_part)
            original_formula = formula_list[formula_split_mask.formula_list_idx]
            alias = formula_split_mask.alias
            # Make sure the index doesn't point to the top-level Formula object
            assert len(formula_split_mask.inner_node_idx.indices) > 0

            split_off_formula_obj = original_formula.formula_obj[formula_split_mask.inner_node_idx]
            split_off_formula_obj = self.mutate_split_node(split_off_formula_obj)
            new_formula_obj = formula_nodes.Formula.make(expr=split_off_formula_obj)

            avatar_ids: set[str] = original_formula.avatar_ids
            select_formula = CompiledFormulaInfo(
                formula_obj=new_formula_obj,
                alias=alias,
                original_field_id=original_formula.original_field_id,
                avatar_ids=avatar_ids,
            )

            # Deduplicate formulas that are the same
            existing_formula = select_by_alias.get(alias)
            if existing_formula is not None:
                # Already added a formula for this alias. Make sure they are identical.
                assert (
                    select_formula.formula_obj == existing_formula.formula_obj
                ), "Got different SELECT formulas for the same alias"

            else:
                # This formula has not been added yet.
                select.append(select_formula)
                select_by_alias[alias] = select_formula

        for add_formula in split_mask.add_formulas:
            formula = CompiledFormulaInfo(
                formula_obj=formula_nodes.Formula.make(expr=add_formula.expr),
                alias=add_formula.alias,
                original_field_id=None,
                avatar_ids=set(add_formula.from_ids),
            )
            select.append(formula)
            if add_formula.is_group_by:
                group_by.append(formula)

        filters: list[CompiledFormulaInfo] = [
            formula for filter_idx, formula in enumerate(query.filters) if filter_idx in split_mask.filter_indices
        ]

        join_on = query.join_on
        joined_from = query.joined_from

        subquery = CompiledQuery(
            id=split_mask.subquery_id,
            level_type=query.level_type,
            select=select,
            group_by=group_by,
            order_by=[],  # ORDER BY only in the top-level sub-query
            filters=filters,
            join_on=join_on,
            joined_from=joined_from,
            meta=QueryMetaInfo(
                query_type=query.meta.query_type,
                field_order=None,
                from_subquery=query.meta.from_subquery,
                subquery_limit=query.meta.subquery_limit,
                empty_query_mode=query.meta.empty_query_mode,
                row_count_hard_limit=query.meta.row_count_hard_limit,
            ),
        )
        return subquery

    def _make_join_on_expression(
        self,
        left_subquery_mask: QuerySplitMask,
        right_subquery_mask: QuerySplitMask,
    ) -> Optional[CompiledJoinOnFormulaInfo]:
        def and_part(condition: Optional[formula_nodes.Binary], part: formula_nodes.Binary) -> formula_nodes.Binary:
            if condition is None:
                return part
            return formula_nodes.Binary.make(name="and", left=condition, right=part)

        join_expr: Optional[formula_nodes.Binary] = None

        aliases_from_right_to_left: dict[str, str] = {}
        left_map = {add_formula.expr.extract: add_formula.alias for add_formula in left_subquery_mask.add_formulas}
        right_map = {add_formula.expr.extract: add_formula.alias for add_formula in right_subquery_mask.add_formulas}
        for node_extract, right_alias in right_map.items():
            if node_extract in left_map:
                aliases_from_right_to_left[right_alias] = left_map[node_extract]

        if isinstance(right_subquery_mask.joining_node, formula_fork_nodes.QueryForkJoiningWithList):
            # Joining node explicitly lists joining conditions.
            for condition in right_subquery_mask.joining_node.condition_list:
                right_expr: formula_nodes.FormulaItem
                left_expr: formula_nodes.FormulaItem
                if isinstance(condition, formula_fork_nodes.SelfEqualityJoinCondition):
                    expr = condition.expr
                    if not isinstance(expr, (formula_nodes.Field, formula_aux_nodes.ErrorNode)):
                        raise TypeError(f"Expected field or error node, got {type(expr)}")

                    left_expr = expr
                    right_expr = expr
                elif isinstance(condition, formula_fork_nodes.BinaryJoinCondition):
                    left_expr = condition.expr
                    right_expr = condition.fork_expr
                else:
                    raise TypeError(f"Type {type(condition).__name__} is not supported")

                # Remap columns in left_expr
                # because it has to be evaluated against the left sub-query,
                # and not the right one, whose field names are used in the expression
                left_expr = remap_formula_obj_fields(node=left_expr, field_name_map=aliases_from_right_to_left)
                part = formula_nodes.Binary.make(name="_dneq", left=left_expr, right=right_expr)
                join_expr = and_part(condition=join_expr, part=part)

        else:
            raise TypeError(f"Joining node type {type(right_subquery_mask.joining_node).__name__} is not supported")

        if join_expr is None:
            return None

        join_type = right_subquery_mask.join_type
        assert join_type is not None

        return CompiledJoinOnFormulaInfo(
            alias=None,  # Will not be used
            formula_obj=formula_nodes.Formula.make(expr=join_expr),
            avatar_ids={left_subquery_mask.subquery_id, right_subquery_mask.subquery_id},
            original_field_id=None,
            left_id=left_subquery_mask.subquery_id,
            right_id=right_subquery_mask.subquery_id,
            join_type=join_type,
        )

    def get_used_from_ids(
        self,
        formula_obj: formula_nodes.Formula,
        formula_alias_to_subquery_id_map: dict[str, str],
    ) -> set[str]:
        used_fields = inspect_expression.used_fields(formula_obj)
        avatar_ids: set[str] = {formula_alias_to_subquery_id_map[field_node.name] for field_node in used_fields}
        return avatar_ids

    def _replace_dimensions_with_subquery_fields(
        self,
        formula_obj: formula_nodes.Formula,
        base_gb_alias_by_extract: dict[NodeExtract, str],
    ) -> formula_nodes.Formula:
        """
        Replace all dimensions in a formula with field nodes
        referencing these dimensions in the sub-query.
        """

        def match_func(node: formula_nodes.FormulaItem, *_: Any) -> bool:
            return node.extract_not_none in base_gb_alias_by_extract

        def replace_func(node: formula_nodes.FormulaItem, *_: Any) -> formula_nodes.FormulaItem:
            dim_field_name = base_gb_alias_by_extract[node.extract_not_none]
            return formula_nodes.Field.make(name=dim_field_name, meta=node.meta)

        return formula_obj.replace_nodes(match_func=match_func, replace_func=replace_func)

    def _generate_cropped_formula_list(
        self,
        query: CompiledQuery,
        query_part: QueryPart,
        gb_aliases: set[str],
        formula_split_masks: list[AliasedFormulaSplitMask],
        base_gb_alias_by_extract: dict[NodeExtract, str],
        formula_alias_to_subquery_id_map: dict[str, str],
        substitute_error_node: Optional[formula_aux_nodes.ErrorNode] = None,
        exclude_indices: Collection[int] = frozenset(),
    ) -> list[CompiledFormulaInfo]:
        if exclude_indices:
            assert query_part == QueryPart.filters, "Only filter formulas can be excluded during query cropping"

        formula_list = query.get_formula_list(query_part)

        formula_masks_by_list_idx: dict[int, list[AliasedFormulaSplitMask]] = {}
        for formula_split_mask in formula_split_masks:
            if formula_split_mask.query_part != query_part:
                continue

            if formula_split_mask.formula_list_idx not in formula_masks_by_list_idx:
                formula_masks_by_list_idx[formula_split_mask.formula_list_idx] = []

            formula_masks_by_list_idx[formula_split_mask.formula_list_idx].append(formula_split_mask)

        new_formula_list: list[CompiledFormulaInfo] = []
        for formula_idx, formula in enumerate(formula_list):
            if formula_idx in exclude_indices:
                # Skip formula. Only for filters (see check above)
                continue

            formula_obj = formula.formula_obj
            is_constant = inspect_expression.is_constant_expression(formula_obj.expr)

            if is_constant:
                # Constant dimensions remain unchanged
                pass

            # Separate logic for expressions that are dimensions
            elif formula.not_none_alias in gb_aliases:
                dim_field_name = base_gb_alias_by_extract[formula_obj.expr.extract_not_none]
                formula_obj = formula_nodes.Formula.make(
                    expr=formula_nodes.Field.make(
                        name=dim_field_name,
                        meta=formula_obj.expr.meta,
                    ),
                    meta=formula_obj.meta,
                )

            # A compatibility error has to be thrown, so put that instead of the field reference
            elif substitute_error_node is not None:
                formula_obj = formula_nodes.Formula.make(expr=substitute_error_node)

            # And for other fields (measures or dimension-dependent filters)
            # substitute expressions for field references to sub-query
            else:
                for formula_split_mask in formula_masks_by_list_idx.get(formula_idx, ()):
                    formula_obj = formula_obj.replace_at_index(
                        index=formula_split_mask.outer_node_idx,
                        expr=formula_nodes.Field.make(
                            name=formula_split_mask.alias,
                            meta=formula_obj.expr.meta,
                        ),
                    )
                if query_part == QueryPart.filters:
                    # Only filters can depend on dimensions, but not be dimensions themselves
                    formula_obj = self._replace_dimensions_with_subquery_fields(
                        formula_obj=formula_obj,
                        base_gb_alias_by_extract=base_gb_alias_by_extract,
                    )

            has_been_updated = formula_obj is not formula.formula_obj
            if has_been_updated:
                formula = formula.clone(
                    formula_obj=formula_obj,
                    avatar_ids=self.get_used_from_ids(formula_obj, formula_alias_to_subquery_id_map),
                )

            elif not is_constant:
                # Can't add a formula that has not been split
                # (it references the original FROM (avatar/sub-query)
                # that should not be accessible from this sub-query).
                # In theory this should happen only to filters.
                # Filters might need to be excluded if all non-dimensions
                # have BFB for this filter.
                assert query_part is QueryPart.filters
                continue

            new_formula_list.append(formula)

        return new_formula_list

    def _are_subquery_dimensions_compatible(self, split_masks: Collection[QuerySplitMask]) -> bool:
        if not split_masks:
            return True

        dimension_sets = [
            frozenset(dim.expr.extract for dim in mask.add_formulas if dim.is_group_by) for mask in split_masks
        ]
        max_dim_set = frozenset(itertools.chain.from_iterable(dimension_sets))
        return max_dim_set in dimension_sets

    def _crop_original_query(
        self,
        query: CompiledQuery,
        split_masks: list[QuerySplitMask],
        subqueries_compatible: bool,
    ) -> CompiledQuery:
        # Combine split indices from all subqueries into one list
        formula_split_masks: list[AliasedFormulaSplitMask] = []
        for subquery_mask in split_masks:
            formula_split_masks += subquery_mask.formula_split_masks

        base_subquery_mask, other_subquery_masks = self._separate_base_and_other_masks(split_masks)
        base_gb_alias_by_extract = {
            dim.expr.extract_not_none: dim.alias for dim in base_subquery_mask.add_formulas if dim.is_group_by
        }

        # Build map {<formula alias>: <subquery id>}
        # from 1) formula split masks, 2) additional formulas
        formula_alias_to_subquery_id_map: dict[str, str] = {}
        for subquery_mask in split_masks:
            for formula_split_mask in subquery_mask.formula_split_masks:
                formula_alias_to_subquery_id_map[formula_split_mask.alias] = subquery_mask.subquery_id
            for add_formula in subquery_mask.add_formulas:
                formula_alias_to_subquery_id_map[add_formula.alias] = subquery_mask.subquery_id

        substitute_error_node: Optional[formula_aux_nodes.ErrorNode] = None
        if not subqueries_compatible:
            substitute_error_node = formula_aux_nodes.ErrorNode.make(
                message="LOD dimensions are incompatible",
                err_code=formula_exc.LodIncompatibleDimensionsError.default_code,
            )

        gb_aliases: set[str] = set()  # Aliases of GROUP BY expressions
        for formula in query.group_by:
            gb_aliases.add(formula.not_none_alias)
            formula_obj_expr = formula.formula_obj.expr
            # Validate that the base mask has this expression among its dimensions
            if not inspect_expression.is_constant_expression(formula_obj_expr):
                assert formula_obj_expr.extract_not_none in base_gb_alias_by_extract

        def generate_formula_list(
            query_part: QueryPart,
            exclude_indices: Collection[int] = frozenset(),
        ) -> list[CompiledFormulaInfo]:
            """A convenience wrapper for method"""
            return self._generate_cropped_formula_list(
                query=query,
                query_part=query_part,
                gb_aliases=gb_aliases,
                base_gb_alias_by_extract=base_gb_alias_by_extract,
                formula_alias_to_subquery_id_map=formula_alias_to_subquery_id_map,
                formula_split_masks=formula_split_masks,
                substitute_error_node=substitute_error_node,
                exclude_indices=exclude_indices,
            )

        # Apply replacements to the original query's formula lists
        new_select = generate_formula_list(query_part=QueryPart.select)
        new_group_by = generate_formula_list(query_part=QueryPart.group_by)
        new_order_by = generate_formula_list(query_part=QueryPart.order_by)

        # Remove filters that have already been applied a lower level.
        new_filters = generate_formula_list(
            query_part=QueryPart.filters, exclude_indices=base_subquery_mask.filter_indices
        )

        # Generate JOIN ON and joined FROM
        join_on: list[CompiledJoinOnFormulaInfo] = []
        froms: list[FromObject] = []
        for subquery_mask in split_masks:
            columns: list[FromColumn] = []
            added_col_aliases: set[str] = set()
            item: AddFormulaInfo | AliasedFormulaSplitMask
            for item in itertools.chain(subquery_mask.formula_split_masks, subquery_mask.add_formulas):  # type: ignore
                alias = item.alias
                if alias not in added_col_aliases:
                    columns.append(FromColumn(id=alias, name=alias))
                    added_col_aliases.add(alias)
            froms.append(
                SubqueryFromObject(
                    id=subquery_mask.subquery_id,
                    alias=subquery_mask.subquery_id,
                    query_id=subquery_mask.subquery_id,
                    columns=tuple(columns),
                )
            )
            if not subqueries_compatible:
                # Queries are incompatible, so no point in generating JOIN ON
                continue
            if subquery_mask is not base_subquery_mask:
                join_on_expr = self._make_join_on_expression(
                    left_subquery_mask=base_subquery_mask,
                    right_subquery_mask=subquery_mask,
                )
                if join_on_expr is not None:
                    join_on.append(join_on_expr)

        root_from_id = froms[0].id
        joined_from = JoinedFromObject(root_from_id=root_from_id, froms=froms)

        # Update the original query object
        query = query.clone(
            select=new_select,
            group_by=new_group_by,
            filters=new_filters,
            order_by=new_order_by,
            join_on=join_on,
            joined_from=joined_from,
        )
        query = self.mutate_cropped_query(query)
        return query

    def _get_formula_masks_as_dict(
        self,
        split_masks: Sequence[QuerySplitMask],
    ) -> dict[QueryPart, dict[int, set[NodeHierarchyIndex]]]:
        """
        reorganize all split-indices from `split_masks` into dict with the following structure:

        {
            <query_part>: {
                <list_index_in_query_part>: {<split_node_idx_1>, <spdsplit_node_idx_2>, ...}
            }
        }
        """

        formula_mask_dict: dict[QueryPart, dict[int, set[NodeHierarchyIndex]]] = {}
        for mask in split_masks:
            for formula_split_mask in mask.formula_split_masks:
                if formula_split_mask.query_part not in formula_mask_dict:
                    formula_mask_dict[formula_split_mask.query_part] = {}
                by_query_part = formula_mask_dict[formula_split_mask.query_part]
                if formula_split_mask.formula_list_idx not in by_query_part:
                    by_query_part[formula_split_mask.formula_list_idx] = set()
                by_query_part[formula_split_mask.formula_list_idx].add(formula_split_mask.outer_node_idx)

        return formula_mask_dict

    def _get_formula_mask_counterparts_for_formula(
        self,
        root_node: formula_nodes.Formula,
        incoming_masks: Sequence[NodeHierarchyIndex],
    ) -> list[NodeHierarchyIndex]:
        incoming_mask_idx = 0
        result: list[NodeHierarchyIndex] = []

        def collect_recursively(node: formula_nodes.FormulaItem, prefix: NodeHierarchyIndex) -> None:
            nonlocal incoming_mask_idx

            node_is_const = inspect_expression.is_constant_expression(node)

            while incoming_mask_idx < len(incoming_masks):
                mask = incoming_masks[incoming_mask_idx]
                if mask >= prefix:
                    break
                incoming_mask_idx += 1
            else:
                # `break` was never called;
                # there are no more masks.
                # So add this node since it doesn't match the mask
                if not node_is_const:  # ignore constants
                    result.append(prefix)
                # Nothing more to do here
                return

            # At this point `mask >= prefix`

            if prefix < mask:
                if not mask.startswith(prefix):
                    # Mask does not point to a child
                    if not node_is_const:  # ignore constants
                        result.append(prefix)
                else:
                    # Mask does point to a child node
                    autonomous_children = list(inspect_expression.enumerate_autonomous_children(node, prefix=prefix))
                    # `autonomous_children` is used instead of node.children because
                    # We don't want to iterate over auxiliary nodes like LOD, BFB, etc. here -
                    # anything that we can't split off.
                    if not autonomous_children:
                        # mask points to nonexistent children
                        raise ValueError("Invalid node index")

                    else:
                        for child_node_idx, child, _ in autonomous_children:
                            collect_recursively(child, child_node_idx)

            # If `prefix == mask`, then this node is blacklisted, skip it.

        collect_recursively(root_node.expr, prefix=NodeHierarchyIndex(indices=(0,)))
        return result

    def _get_formula_mask_counterparts(
        self,
        query: CompiledQuery,
        split_masks: Sequence[QuerySplitMask],
        expr_id_gen: PrefixedIdGen,
    ) -> list[AliasedFormulaSplitMask]:
        """
        Generate "counterpart" split-indices for the ones given.
        Basically what this means is that for every incoming split-index
        the result will include an index for each of its siblings (child nodes of the same parent node)
        not covered by other incoming indices. Ignore constants - don't generate split-indices for them.
        """

        formula_mask_dict = self._get_formula_masks_as_dict(split_masks=split_masks)

        counter_formula_split_masks: list[AliasedFormulaSplitMask] = []

        gb_aliases = [formula.not_none_alias for formula in query.group_by]  # Aliases of GROUP BY expressions

        for query_part in self.parent_query_parts_for_splitting:
            by_list_idx = formula_mask_dict.get(query_part, {})
            formula_list = query.get_formula_list(query_part)
            for formula_list_idx, formula in enumerate(formula_list):
                if formula.alias in gb_aliases:
                    # Do not split dimensions
                    continue

                if inspect_expression.is_constant_expression(formula.formula_obj.expr):
                    # Do not split constant expressions
                    continue

                node_index_list = by_list_idx.get(formula_list_idx, ())
                if query_part == QueryPart.filters and not node_index_list:
                    # If a filter has not yet been split, don't split it now
                    continue

                counter_formula_indices = self._get_formula_mask_counterparts_for_formula(
                    root_node=formula.formula_obj,
                    incoming_masks=sorted(node_index_list),
                )
                if counter_formula_indices:
                    for node_idx in counter_formula_indices:
                        counter_formula_split_masks.append(
                            AliasedFormulaSplitMask(
                                alias=expr_id_gen.get_id(),
                                query_part=query_part,
                                formula_list_idx=formula_list_idx,
                                inner_node_idx=node_idx,
                                outer_node_idx=node_idx,
                            )
                        )

        return counter_formula_split_masks

    def _find_base_query_candidate(
        self,
        query: CompiledQuery,
        split_masks: list[QuerySplitMask],
        base_formula_split_masks: list[AliasedFormulaSplitMask],
        base_filter_indices: set[int],
    ) -> Optional[str]:
        """
        Here we try to determine if there is a need to generate a "base" sub-query
        that all the explicitly generated ones will be joined to.

        If there is a sub-query that has the same dimensions and same filters
        as one of the generated sub-queries and has no other expressions
        that need to be selected, we can use this existing sub-query as the base.
        """

        if base_formula_split_masks:
            # There are parts of the original query's expressions
            # that need to be selected, but are not part of other sub-queries
            # In this case we will need to select them as part of the base sub-query
            return None

        base_group_by_extracts = {
            formula.formula_obj.expr.extract_not_none
            for formula in query.group_by
            if not inspect_expression.is_constant_expression(formula.formula_obj.expr)
        }

        # Count the number of dimensions the base sub-query needs to have
        base_group_by_count = len(base_group_by_extracts)

        # Try to find a candidate for being the base sub-query among the existing sub-query masks
        max_dimensions = max(mask.group_by_count for mask in split_masks)

        max_query_id: Optional[str] = None
        max_query_has_all_base_dimensions = False
        max_query_has_all_base_filters = True
        max_query_has_direct_join = True
        for mask in split_masks:
            if mask.group_by_count != max_dimensions:
                # We are looking for the sub-query with the most dimensions
                continue

            mask_group_by_extracts = {
                add_formula.expr.extract_not_none for add_formula in mask.add_formulas if add_formula.is_group_by
            }
            if not mask_group_by_extracts.issuperset(base_group_by_extracts):
                # The sub-query has to have all of the original query's dimensions
                continue

            # Otherwise this is the largest sub-query (by number of dimensions)
            max_query_id = mask.subquery_id
            max_query_has_all_base_dimensions = mask.group_by_count == base_group_by_count
            max_query_has_all_base_filters = set(mask.filter_indices) == base_filter_indices
            max_query_has_direct_join = mask.has_direct_equality_join

        if max_query_has_all_base_dimensions and max_query_has_all_base_filters and max_query_has_direct_join:
            # No selects left for the original query,
            # and subquery with the most dimensions has the same
            # dimensions and filters as would the base sub-query.
            # So we can just use this "largest" subquery as the base.
            assert max_query_id is not None
            return max_query_id

        if base_group_by_count == 0:
            # Base query doesn't have any own dimensions
            # (inherited non-const ones from the original query).
            # If there are supposed to be any filters,
            # they are useless in a 0-dimensional query with no selects.
            # So don't create a new sub-query, use the existing sub-query
            # with the most dimensions as base
            assert max_query_id is not None
            return max_query_id

        return None

    def _patch_query_masks_with_base(
        self,
        query: CompiledQuery,
        split_masks: list[QuerySplitMask],
        expr_id_gen: PrefixedIdGen,
        query_id_gen: PrefixedIdGen,
    ) -> list[QuerySplitMask]:
        base_formula_split_masks = self._get_formula_mask_counterparts(
            query=query, split_masks=split_masks, expr_id_gen=expr_id_gen
        )

        # Find the indices of filters that will be applied at the upper level
        split_filter_indices = {
            fla_mask.formula_list_idx
            for query_mask in split_masks
            for fla_mask in query_mask.formula_split_masks
            if fla_mask.query_part == QueryPart.filters
        }
        base_filter_indices = set(range(len(query.filters))) - split_filter_indices

        base_subquery_candidate_id = self._find_base_query_candidate(
            query=query,
            split_masks=split_masks,
            base_formula_split_masks=base_formula_split_masks,
            base_filter_indices=base_filter_indices,
        )
        if base_subquery_candidate_id is not None:
            # Patch the candidate with `is_base=True` and return the original mask list
            split_masks = [
                (mask if mask.subquery_id is not base_subquery_candidate_id else mask.clone(is_base=True))
                for mask in split_masks
            ]
            # Re-order them so that the base is the first one (False goes before True)
            split_masks = sorted(split_masks, key=lambda mask: (not mask.is_base, split_masks.index(mask)))
            return split_masks

        # Generate additional formulas from query's dimensions
        new_add_formulas: list[AddFormulaInfo] = []
        base_dimension_extracts: set[NodeExtract] = set()
        for formula in query.group_by:
            alias = expr_id_gen.get_id()
            add_formula = AddFormulaInfo(
                alias=alias,
                expr=formula.formula_obj.expr,
                from_ids=frozenset(formula.avatar_ids),
                is_group_by=not inspect_expression.is_constant_expression(formula.formula_obj.expr),
            )
            new_add_formulas.append(add_formula)
            if add_formula.is_group_by:
                base_dimension_extracts.add(add_formula.expr.extract_not_none)

        new_filter_indices = {
            filter_idx for filter_idx in range(len(query.filters)) if filter_idx in base_filter_indices
        }

        base_mask = QuerySplitMask(
            subquery_type=SubqueryType.generated_base,
            subquery_id=query_id_gen.get_id(),
            formula_split_masks=tuple(base_formula_split_masks),
            add_formulas=tuple(new_add_formulas),
            filter_indices=frozenset(new_filter_indices),
            join_type=None,
            joining_node=None,
            is_base=True,
        )
        return [base_mask] + split_masks

    def _separate_base_and_other_masks(
        self,
        split_masks: list[QuerySplitMask],
    ) -> tuple[QuerySplitMask, list[QuerySplitMask]]:
        """Separate list of sub-query masks into base and all others."""

        assert split_masks

        base_mask = split_masks[0]
        assert base_mask.is_base

        other_masks = split_masks[1:]
        assert all(not other_mask.is_base for other_mask in other_masks)

        return base_mask, other_masks

    def optimize_query_split_masks(self, split_masks: list[QuerySplitMask]) -> list[QuerySplitMask]:
        """Dummy method that can be redefined in subclasses to optimize the creation of sub-queries"""
        return split_masks

    def _patch_mask_dimensions(
        self,
        split_masks: list[QuerySplitMask],
        expr_id_gen: PrefixedIdGen,
    ) -> list[QuerySplitMask]:
        """
        Find the base sub-query in the list and patch it with additional select items
        so that join expressions can be built correctly.
        """

        base_mask, other_masks = self._separate_base_and_other_masks(split_masks)

        base_add_formula_extracts: set[NodeExtract] = {
            add_formula.expr.extract_not_none for add_formula in base_mask.add_formulas
        }

        # Add extra dimensions (and other formulas needed for joins)
        # if and only if they are either constant expressions or not dimensions (are aggregations)
        # (for joining to sub-queries with total grouping - with aggregation, but without GROUP BY)
        for other_mask in other_masks:
            for add_formula in other_mask.add_formulas:
                if add_formula.expr.extract_not_none in base_add_formula_extracts:
                    continue  # Skip the ones that the mask already has

                new_add_formula = AddFormulaInfo(
                    alias=expr_id_gen.get_id(),
                    expr=add_formula.expr,
                    from_ids=add_formula.from_ids,
                    is_group_by=add_formula.is_group_by,
                )
                base_mask = base_mask.clone(add_formulas=base_mask.add_formulas + (new_add_formula,))
                base_add_formula_extracts.add(add_formula.expr.extract_not_none)

        return [base_mask] + other_masks

    def _get_query_patch_from_split_masks(
        self,
        query: CompiledQuery,
        split_masks: list[QuerySplitMask],
        query_id_gen: PrefixedIdGen,
        expr_id_gen: PrefixedIdGen,
    ) -> CompiledMultiQueryPatch:
        # Perform compatibility check before the list of query split masks is modified
        subqueries_compatible = self._are_subquery_dimensions_compatible(split_masks)

        # Create a mask for select items unaffected by the given masks.
        # It will represent the cropped original query.
        split_masks = self._patch_query_masks_with_base(
            query=query, split_masks=split_masks, expr_id_gen=expr_id_gen, query_id_gen=query_id_gen
        )
        split_masks = self.optimize_query_split_masks(split_masks=split_masks)
        split_masks = self._patch_mask_dimensions(split_masks=split_masks, expr_id_gen=expr_id_gen)

        # Generate subquery for each split-mask
        result_queries: list[CompiledQuery] = []
        for mask in split_masks:
            subquery = self._generate_subquery_for_mask(query=query, split_mask=mask)
            # Here some keys may appear for several subqueries.
            # Retain the value from the earliest appearance of the key.
            result_queries.append(subquery)

        # Crop original query (replace expressions specified by split masks with subquery fields)
        updated_original_query = self._crop_original_query(
            query=query, split_masks=split_masks, subqueries_compatible=subqueries_compatible
        )
        result_queries.append(updated_original_query)

        # Put it all into a patch object and return it
        patch = CompiledMultiQueryPatch(
            patch_multi_query=CompiledMultiQuery(queries=result_queries),
        )
        return patch

    def split_query(
        self,
        query: CompiledQuery,
        requirement_subtree: CompiledMultiQueryBase,
        query_id_gen: PrefixedIdGen,
        expr_id_gen: PrefixedIdGen,
    ) -> Optional[CompiledMultiQueryPatch]:
        split_masks = self.get_split_masks(query, expr_id_gen=expr_id_gen, query_id_gen=query_id_gen)
        if not split_masks:
            return None

        patch = self._get_query_patch_from_split_masks(
            query=query,
            split_masks=split_masks,
            query_id_gen=query_id_gen,
            expr_id_gen=expr_id_gen,
        )
        return patch

lib/dl_query_processing/dl_query_processing/multi_query/splitters/mask_based.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    ClassVar,
    Optional,
)

import attr

import dl_formula.core.nodes as formula_nodes
from dl_query_processing.compilation.primitives import (
    CompiledFormulaInfo,
    CompiledJoinOnFormulaInfo,
    CompiledMultiQuery,
    CompiledMultiQueryBase,
    CompiledQuery,
    FromColumn,
    FromObject,
    JoinedFromObject,
    SubqueryFromObject,
)
from dl_query_processing.compilation.query_meta import QueryMetaInfo
from dl_query_processing.enums import (
    ExecutionLevel,
    QueryPart,
)
from dl_query_processing.multi_query.splitters.base import MultiQuerySplitterBase
from dl_query_processing.multi_query.tools import CompiledMultiQueryPatch
from dl_query_processing.utils.name_gen import PrefixedIdGen


@attr.s
class PrefilteredFieldMultiQuerySplitter(MultiQuerySplitterBase):
    """
    Splits off sub-queries at field level.
    """

    duplicate_prefilters: ClassVar[bool] = True

    crop_to_level_type: ExecutionLevel = attr.ib(kw_only=True)

    def is_pre_filter(self, formula: CompiledFormulaInfo) -> bool:
        """By default pre-filters are not supported"""
        return False

    def is_simple_query(self, query: CompiledQuery) -> bool:
        if query.group_by or query.order_by or query.join_on:
            return False
        if not all(self.is_pre_filter(formula) for formula in query.filters):
            return False
        for formula in query.select:
            # Check whether formula is a simple field expression
            if not isinstance(formula.formula_obj.expr, formula_nodes.Field):
                return False

        return True

    def split_query(
        self,
        query: CompiledQuery,
        requirement_subtree: CompiledMultiQueryBase,
        query_id_gen: PrefixedIdGen,
        expr_id_gen: PrefixedIdGen,
    ) -> Optional[CompiledMultiQueryPatch]:
        # First check if we really need to do anything
        if self.is_simple_query(query):
            return None

        field_to_new_alias_map: dict[str, str] = {}
        avatar_subq_map: dict[str, str] = {}

        old_fields_to_from_ids: dict[str, str] = {}
        used_fields_by_avatar: dict[str, list[str]] = {}
        for original_from in query.joined_from.froms:
            original_from_id = original_from.id
            new_from_id = query_id_gen.get_id()
            avatar_subq_map[original_from_id] = new_from_id
            used_fields_by_avatar[original_from_id] = []
            for col in original_from.columns:
                old_fields_to_from_ids[col.id] = original_from_id

        def match_func(node: formula_nodes.FormulaItem, *args: Any) -> bool:
            return isinstance(node, formula_nodes.Field)

        def replace_func(node: formula_nodes.FormulaItem, *args: Any) -> formula_nodes.FormulaItem:
            assert isinstance(node, formula_nodes.Field)
            original_field_name = node.name
            new_field_name = field_to_new_alias_map.get(original_field_name)
            if new_field_name is None:
                new_field_name = expr_id_gen.get_id()
                field_to_new_alias_map[original_field_name] = new_field_name
                original_from_id = old_fields_to_from_ids[original_field_name]
                used_fields_by_avatar[original_from_id].append(original_field_name)
            assert new_field_name is not None
            return formula_nodes.Field.make(name=new_field_name, meta=node.meta)

        query_parts = (
            QueryPart.select,
            QueryPart.group_by,
            QueryPart.filters,
            QueryPart.order_by,
            QueryPart.join_on,
        )

        cropped_formula_list_by_query_part: dict[QueryPart, list[CompiledFormulaInfo]] = {}
        prefilters_by_from_id: dict[str, list[CompiledFormulaInfo]] = {}
        for query_part in query_parts:
            formula_list = query.get_formula_list(query_part)
            cropped_formula_list: list[CompiledFormulaInfo] = []
            for formula in formula_list:
                if self.is_pre_filter(formula):
                    # Send pre-filters to the sub-queries
                    assert len(formula.avatar_ids) == 1
                    from_id = next(iter(formula.avatar_ids))
                    if from_id not in prefilters_by_from_id:
                        prefilters_by_from_id[from_id] = []
                    prefilters_by_from_id[from_id].append(formula)

                    if not self.duplicate_prefilters:
                        # Pre-filter duplication is off, so do not add the filter to the cropped query
                        continue

                formula_obj = formula.formula_obj
                # Replace all field names with sub-query-referencing aliases
                cropped_formula_obj = formula_obj.replace_nodes(match_func=match_func, replace_func=replace_func)
                used_subquery_ids = {avatar_subq_map[from_id] for from_id in formula.avatar_ids}
                formula = formula.clone(
                    formula_obj=cropped_formula_obj,
                    avatar_ids=used_subquery_ids,
                )
                if isinstance(formula, CompiledJoinOnFormulaInfo):
                    formula = formula.clone(
                        left_id=avatar_subq_map[formula.left_id],
                        right_id=avatar_subq_map[formula.right_id],
                    )
                cropped_formula_list.append(formula)

            cropped_formula_list_by_query_part[query_part] = cropped_formula_list

        subqueries: list[CompiledQuery] = []

        cropped_froms: list[FromObject] = []
        for original_from in query.joined_from.froms:
            original_from_id = original_from.id
            new_from_id = avatar_subq_map[original_from_id]

            columns: list[FromColumn] = []
            for col in original_from.columns:
                new_col_id = field_to_new_alias_map.get(col.id)
                if new_col_id is None:
                    continue  # This column is not used in the query, so skip it
                columns.append(FromColumn(id=new_col_id, name=new_col_id))

            cropped_froms.append(
                SubqueryFromObject(
                    id=new_from_id,
                    alias=new_from_id,
                    query_id=new_from_id,
                    columns=tuple(columns),
                )
            )

            subquery_select = [
                CompiledFormulaInfo(
                    formula_obj=formula_nodes.Formula.make(expr=formula_nodes.Field.make(name=field_name)),
                    alias=field_to_new_alias_map[field_name],
                    avatar_ids={original_from_id},
                    original_field_id=None,
                )
                for field_name in used_fields_by_avatar[original_from_id]
            ]

            subquery_filters: list[CompiledFormulaInfo] = []
            if original_from_id in prefilters_by_from_id:
                for prefilter_formula in prefilters_by_from_id[original_from_id]:
                    subquery_filters.append(prefilter_formula)

            subquery = CompiledQuery(
                id=new_from_id,
                select=subquery_select,
                filters=subquery_filters,
                level_type=query.level_type,
                joined_from=JoinedFromObject(
                    root_from_id=original_from_id,
                    froms=[original_from],
                ),
                meta=QueryMetaInfo(
                    query_type=query.meta.query_type,
                    field_order=None,
                    from_subquery=query.meta.from_subquery,
                    subquery_limit=query.meta.subquery_limit,
                    empty_query_mode=query.meta.empty_query_mode,
                    row_count_hard_limit=query.meta.row_count_hard_limit,
                ),
            )
            subqueries.append(subquery)

        new_root_from_id: Optional[str] = None
        if query.joined_from.root_from_id is not None:
            new_root_from_id = avatar_subq_map[query.joined_from.root_from_id]

        cropped_joined_from = JoinedFromObject(root_from_id=new_root_from_id, froms=cropped_froms)
        cropped_query = query.clone(
            select=cropped_formula_list_by_query_part[QueryPart.select],
            group_by=cropped_formula_list_by_query_part[QueryPart.group_by],
            filters=cropped_formula_list_by_query_part[QueryPart.filters],
            order_by=cropped_formula_list_by_query_part[QueryPart.order_by],
            join_on=cropped_formula_list_by_query_part[QueryPart.join_on],
            joined_from=cropped_joined_from,
            level_type=self.crop_to_level_type,
        )

        patch = CompiledMultiQueryPatch(patch_multi_query=CompiledMultiQuery(queries=[cropped_query, *subqueries]))
        return patch

lib/dl_query_processing/dl_query_processing/multi_query/splitters/prefiltered.py
datalens-tech/datalens-backend
from collections import OrderedDict
import itertools
from typing import Any

import attr

from dl_constants.enums import JoinType
from dl_formula.core.extract import NodeExtract
import dl_formula.core.fork_nodes as formula_fork_nodes
from dl_formula.core.index import NodeHierarchyIndex
import dl_formula.core.nodes as formula_nodes
from dl_formula.inspect.env import InspectionEnvironment
import dl_formula.inspect.expression as inspect_expression
import dl_formula.inspect.node as inspect_node
from dl_query_processing.compilation.primitives import CompiledQuery
from dl_query_processing.enums import QueryPart
from dl_query_processing.multi_query.splitters.mask_based import (
    AddFormulaInfo,
    AliasedFormulaSplitMask,
    FormulaSplitMask,
    MultiQuerySplitter,
    QuerySplitMask,
    SubqueryType,
)
from dl_query_processing.utils.name_gen import PrefixedIdGen


_JOIN_TYPE_MAP = {
    formula_fork_nodes.JoinType.inner: JoinType.inner,
    formula_fork_nodes.JoinType.left: JoinType.left,
}


@attr.s(frozen=True, auto_attribs=True)
class SubqueryForkSignature:
    lod_idx: int  # A switch to disable grouping of same-dimension LODs together
    joining_node_extract: NodeExtract
    bfb_names: frozenset[str]
    dim_extracts: frozenset[NodeExtract]
    child_lod_extracts: frozenset[NodeExtract]
    join_type: JoinType


@attr.s(frozen=True)
class QueryForkInfo:
    subquery_type: SubqueryType = attr.ib(kw_only=True)
    joining_node: formula_fork_nodes.QueryForkJoiningBase = attr.ib(kw_only=True)
    bfb_field_ids: frozenset[str] = attr.ib(kw_only=True)
    add_formulas: tuple[AddFormulaInfo, ...] = attr.ib(kw_only=True)
    join_type: JoinType = attr.ib(kw_only=True)
    aliases_by_extract: dict[NodeExtract, str] = attr.ib(kw_only=True, factory=dict)
    formula_split_masks: list[AliasedFormulaSplitMask] = attr.ib(kw_only=True, factory=list)


# Some type aliases
FMask_QFork = tuple[FormulaSplitMask, formula_fork_nodes.QueryFork]
FMask_QFork_BFB = tuple[FormulaSplitMask, formula_fork_nodes.QueryFork, frozenset[str]]


@attr.s
class QueryForkQuerySplitter(MultiQuerySplitter):
    def _get_child_lod_extracts(
        self,
        node: formula_nodes.FormulaItem,
    ) -> frozenset[NodeExtract]:
        """
        Get LODs of child QueryFork's (retrieved subqueries)
        for the given formula node (most likely a QueryFork itself).
        This is done to check compatibility of query forks - whether or not
        they can be put in the same sub-query.
        """
        child_lod_extracts: set[NodeExtract] = set()
        child_query_forks = self._get_query_forks_from_expression(
            node=node,
            index_prefix=NodeHierarchyIndex(),
            query_part=QueryPart.select,
            formula_idx=0,  # Fake values, but they don't matter
        )
        for _, child_qfork in child_query_forks:
            child_lod_extracts.add(child_qfork.lod.extract_not_none)

        return frozenset(child_lod_extracts)

    def _get_query_forks_from_expression(
        self,
        node: formula_nodes.FormulaItem,
        index_prefix: NodeHierarchyIndex,
        query_part: QueryPart,
        formula_idx: int,
    ) -> list[tuple[FormulaSplitMask, formula_fork_nodes.QueryFork]]:
        result: list[tuple[FormulaSplitMask, formula_fork_nodes.QueryFork]] = []
        if isinstance(node, formula_fork_nodes.QueryFork):
            result.append(
                (
                    FormulaSplitMask(
                        query_part=query_part,
                        formula_list_idx=formula_idx,
                        outer_node_idx=index_prefix,
                        inner_node_idx=index_prefix + (1,),  # QueryFork.result_expr = Child(1)
                    ),
                    node,
                )
            )
        else:
            for child_index, child, _stack in inspect_expression.enumerate_autonomous_children(
                node, prefix=index_prefix
            ):
                result += self._get_query_forks_from_expression(
                    node=child,
                    index_prefix=child_index,
                    query_part=query_part,
                    formula_idx=formula_idx,
                )

        return result

    def _prioritize_and_filter_query_forks(
        self,
        fmask_qfork_bfb_list: list[FMask_QFork_BFB],
    ) -> tuple[SubqueryType, list[FMask_QFork_BFB]]:
        """
        Filter the found query forks.

        It is possible to minimize the number of JOINs in a query with window functions
        if we prioritize them and do not mix aggregation and lookup QueryForks with window function ones.
        If at the top level of a query we find WF and non-WF query forks,
        then we ignore the non-WF ones at this iteration.
        This way we get one sub-query above the split line and one below. Without any joins.

        All of this also requires the BFB of the top-level WF to be a sub-set of all other BFBs,
        otherwise this nesting is not possible
        """

        # Attempt to prioritize query forks generated by window functions.
        wf_qforks_found: bool = False
        wf_bfb_set: set[frozenset[str]] = set()
        non_wf_bfb_set: set[frozenset[str]] = set()
        for formula_split_mask, qfork_node, normalized_bfb in fmask_qfork_bfb_list:
            if inspect_node.qfork_is_window(qfork_node):
                wf_qforks_found = True
                wf_bfb_set.add(normalized_bfb)
            else:
                non_wf_bfb_set.add(normalized_bfb)

        winfunc_mode: bool = False
        smallest_wf_bfb: frozenset[str] = frozenset()
        if wf_qforks_found:
            winfunc_mode = True
            # Check whether the smallest WF BFB is a subset of all other (WF and non-WF) BFBs
            smallest_wf_bfb = sorted(wf_bfb_set, key=lambda bfb: len(bfb))[0]
            for other_bfb in wf_bfb_set | non_wf_bfb_set:
                if not smallest_wf_bfb.issubset(other_bfb):
                    winfunc_mode = False
                    break

        if not winfunc_mode:
            return SubqueryType.default, fmask_qfork_bfb_list

        filtered_fmask_qfork_bfb_list: list[FMask_QFork_BFB] = []
        for formula_split_mask, qfork_node, normalized_bfb in fmask_qfork_bfb_list:
            if not inspect_node.qfork_is_window(qfork_node) or normalized_bfb != smallest_wf_bfb:
                # Window function-prioritization is on,
                # so all non-WF q. forks or WF q. forks with non-matching BFBs will be ignored
                continue

            filtered_fmask_qfork_bfb_list.append((formula_split_mask, qfork_node, normalized_bfb))

        return SubqueryType.window_func, filtered_fmask_qfork_bfb_list

    def _unify_masks_for_window_functions(self, split_masks: list[QuerySplitMask]) -> list[QuerySplitMask]:
        """
        Part of the optimization hack for window functions.
        If sub-query masks are in WF mode, then combine them all with the base mask
        so that there are no JOINs. The resulting sub-query's GROUP BY will be eventually removed
        by the mutate_cropped_query method
        """

        subq_types = {mask.subquery_type for mask in split_masks if mask.subquery_type != SubqueryType.generated_base}
        assert len(subq_types) == 1
        subqery_type = next(iter(subq_types))

        if subqery_type != SubqueryType.window_func:
            return split_masks

        base_mask = split_masks[0]
        assert base_mask.is_base
        other_masks = split_masks[1:]

        if len(other_masks) == 0:
            # There's only the base mask
            return split_masks

        if len(other_masks) > 1:
            # Can't optimize for multiple window func subqueries
            return split_masks

        # At this point we must have one base mask and one containing the query fork
        assert len(other_masks) == 1
        other_mask = other_masks[0]

        add_formulas: list[AddFormulaInfo] = []
        add_formula_extracts: set[NodeExtract] = set()
        formula_masks: list[AliasedFormulaSplitMask] = []
        for mask in split_masks:
            for add_formula in mask.add_formulas:
                if add_formula.expr.extract not in add_formula_extracts:
                    add_formulas.append(add_formula)
                    add_formula_extracts.add(add_formula.expr.extract_not_none)
            for formula_mask in mask.formula_split_masks:
                formula_masks.append(formula_mask)

        base_mask = base_mask.clone(
            add_formulas=add_formulas,
            formula_split_masks=formula_masks,
            # replace filters of the generated base query so that we don't ignore BFBs
            filter_indices=other_mask.filter_indices,
        )

        split_masks = [base_mask]
        return split_masks

    def _collect_query_forks(
        self,
        query: CompiledQuery,
        expr_id_gen: PrefixedIdGen,
        disable_grouping_lods: bool = False,
    ) -> list[QueryForkInfo]:
        inspect_env = InspectionEnvironment()

        result: list[QueryForkInfo] = []

        lod_idx_counter = itertools.count()

        def get_new_lod_idx() -> int:
            if disable_grouping_lods:
                # LODs shouldn't be grouped, so generate a new index every time
                return next(lod_idx_counter)
            return 0

        fmask_qfork_list: list[FMask_QFork] = []
        for query_part in self.parent_query_parts_for_splitting:
            formula_list = query.get_formula_list(query_part)
            for formula_idx, formula in enumerate(formula_list):
                fmask_qfork_list += self._get_query_forks_from_expression(
                    node=formula.formula_obj,
                    index_prefix=NodeHierarchyIndex(),
                    query_part=query_part,
                    formula_idx=formula_idx,
                )

        # Normalize BFBs
        available_filter_ids = frozenset(
            filter_formula.original_field_id
            for filter_formula in query.filters
            if filter_formula.original_field_id is not None
        )

        def _normalize_bfb(qfork_node: formula_fork_nodes.QueryFork) -> frozenset[str]:
            return frozenset(qfork_node.before_filter_by.field_names) & available_filter_ids

        fmask_qfork_bfb_list: list[FMask_QFork_BFB] = []
        for formula_split_mask, qfork_node in fmask_qfork_list:
            normalized_bfb = _normalize_bfb(qfork_node)
            fmask_qfork_bfb_list.append((formula_split_mask, qfork_node, normalized_bfb))

        subquery_type, fmask_qfork_bfb_list = self._prioritize_and_filter_query_forks(fmask_qfork_bfb_list)

        qforks_by_signature: OrderedDict[SubqueryForkSignature, QueryForkInfo] = OrderedDict()
        for formula_split_mask, qfork_node, normalized_bfb in fmask_qfork_bfb_list:
            join_type = _JOIN_TYPE_MAP[qfork_node.join_type]
            lod = qfork_node.lod
            dim_list: tuple[formula_nodes.FormulaItem, ...]
            if isinstance(lod, formula_nodes.FixedLodSpecifier):
                dim_list = tuple(lod.dim_list)
            elif isinstance(lod, formula_nodes.InheritedLodSpecifier):
                # Inherit (copy) original query's dimension list
                dim_list = tuple(gb_formula.formula_obj.expr for gb_formula in query.group_by)
            else:
                raise TypeError(f"Unsupported LodSpecifier type: {type(lod).__name__}")

            joining: formula_fork_nodes.QueryForkJoiningBase = qfork_node.joining
            if len(dim_list) == 0:
                # Add a dummy dimension.
                dummy_dim_node = formula_nodes.LiteralInteger.make(value=1)
                dim_list += (dummy_dim_node,)
                joining = formula_fork_nodes.QueryForkJoiningWithList.make(
                    condition_list=[
                        formula_fork_nodes.SelfEqualityJoinCondition.make(expr=dummy_dim_node),
                    ],
                )

            # Generate hashable QueryFork "signature"
            # to deduplicate query forks with the same structure in the same query
            qfork_signature = SubqueryForkSignature(
                dim_extracts=frozenset(dim.extract_not_none for dim in dim_list),
                child_lod_extracts=self._get_child_lod_extracts(node=qfork_node.result_expr),
                join_type=join_type,
                bfb_names=normalized_bfb,
                joining_node_extract=joining.extract_not_none,
                lod_idx=get_new_lod_idx(),
            )

            if qfork_signature not in qforks_by_signature:
                dim_add_formulas: list[AddFormulaInfo] = []
                for dim in dim_list:
                    dim_add_formulas.append(
                        AddFormulaInfo(
                            alias=expr_id_gen.get_id(),
                            expr=dim,
                            from_ids=frozenset(query.joined_from.iter_ids()),
                            is_group_by=not inspect_expression.is_constant_expression(dim),
                        )
                    )

                # Collect non-dimension expressions from the `joining` node.
                # They will be needed to construct JOIN ON conditions correctly
                non_dim_add_formulas: list[AddFormulaInfo] = []
                assert isinstance(joining, formula_fork_nodes.QueryForkJoiningWithList)
                for condition in joining.condition_list:
                    expressions: list[formula_nodes.FormulaItem] = []
                    if isinstance(condition, formula_fork_nodes.SelfEqualityJoinCondition):
                        expressions.append(condition.expr)
                    elif isinstance(condition, formula_fork_nodes.BinaryJoinCondition):
                        expressions.append(condition.expr)
                        expressions.append(condition.fork_expr)
                    else:
                        raise TypeError(type(condition))

                    for expr in expressions:
                        if inspect_expression.is_aggregate_expression(expr, env=inspect_env):
                            non_dim_add_formulas.append(
                                AddFormulaInfo(
                                    alias=expr_id_gen.get_id(),
                                    expr=expr,
                                    from_ids=frozenset(query.joined_from.iter_ids()),
                                    is_group_by=False,
                                )
                            )

                add_formulas = tuple(dim_add_formulas + non_dim_add_formulas)
                qfork_info = QueryForkInfo(
                    subquery_type=subquery_type,
                    add_formulas=add_formulas,
                    joining_node=joining,
                    join_type=join_type,
                    bfb_field_ids=normalized_bfb,
                )
                qforks_by_signature[qfork_signature] = qfork_info
                result.append(qfork_info)

            qfork_info = qforks_by_signature[qfork_signature]

            # Deduplicate expressions that are the same
            # (by using the same alias for their split masks)
            subnode_extract = qfork_node.result_expr.extract_not_none
            alias: str
            existing_alias = qfork_info.aliases_by_extract.get(subnode_extract)
            if existing_alias is None:
                alias = expr_id_gen.get_id()
                qfork_info.aliases_by_extract[subnode_extract] = alias
            else:
                alias = existing_alias

            formula_split_mask = AliasedFormulaSplitMask(
                alias=alias,
                query_part=formula_split_mask.query_part,
                formula_list_idx=formula_split_mask.formula_list_idx,
                outer_node_idx=formula_split_mask.outer_node_idx,
                inner_node_idx=formula_split_mask.inner_node_idx,
            )
            qfork_info.formula_split_masks.append(formula_split_mask)

        return result

    def _normalize_joining_node(
        self,
        joining_node: formula_fork_nodes.QueryForkJoiningBase,
        expr_extract_to_alias_map: dict[NodeExtract, str],
    ) -> formula_fork_nodes.QueryForkJoiningBase:
        def match_func(node: formula_nodes.FormulaItem, *args: Any) -> bool:
            return node.extract in expr_extract_to_alias_map

        def replace_func(node: formula_nodes.FormulaItem, *args: Any) -> formula_nodes.FormulaItem:
            name = expr_extract_to_alias_map[node.extract_not_none]
            return formula_nodes.Field.make(name=name, meta=node.meta)

        joining_node = joining_node.replace_nodes(match_func=match_func, replace_func=replace_func)
        assert isinstance(joining_node, formula_fork_nodes.QueryForkJoiningBase)
        return joining_node

    def optimize_query_split_masks(self, split_masks: list[QuerySplitMask]) -> list[QuerySplitMask]:
        """Optimize for window function sub-queries"""
        split_masks = self._unify_masks_for_window_functions(split_masks=split_masks)
        return split_masks

    def get_split_masks(
        self,
        query: CompiledQuery,
        expr_id_gen: PrefixedIdGen,
        query_id_gen: PrefixedIdGen,
    ) -> list[QuerySplitMask]:
        mask_list: list[QuerySplitMask] = []
        collected_query_fork_infos = self._collect_query_forks(query, expr_id_gen=expr_id_gen)

        # Filters that have query forks at this level cannot be applied to any other query forks
        # (that would require them to be nested, which would mean that they belong to a different level)
        split_filter_indices = {
            formula_split_mask.formula_list_idx
            for qfork_info in collected_query_fork_infos
            for formula_split_mask in qfork_info.formula_split_masks
            if formula_split_mask.query_part == QueryPart.filters
        }

        for qfork_info in collected_query_fork_infos:
            # Map of expressions to aliases.
            # To replace these expressions with their column aliases
            # in the JOIN ON expressions.
            expr_extract_to_alias_map = {
                add_formula.expr.extract_not_none: add_formula.alias for add_formula in qfork_info.add_formulas
            }
            # So apply this mapping to the `joining` node
            joining_node = self._normalize_joining_node(
                joining_node=qfork_info.joining_node, expr_extract_to_alias_map=expr_extract_to_alias_map
            )

            # Collect indices of filters that should be applied to the sub-query
            filter_indices: set[int] = set()
            for filter_idx, filter_formula in enumerate(query.filters):
                if filter_formula.original_field_id in qfork_info.bfb_field_ids:
                    # Filter field is in BFB, so exclude it.
                    continue
                if filter_idx in split_filter_indices:
                    # This filter can only be applied to a higher-level query
                    # Because its pieces are being split of as SELECTs at this level.
                    # Skip it.
                    continue

                filter_indices.add(filter_idx)

            # Put it all together in the sub-query mask
            mask = QuerySplitMask(
                subquery_type=qfork_info.subquery_type,
                subquery_id=query_id_gen.get_id(),
                formula_split_masks=tuple(qfork_info.formula_split_masks),
                filter_indices=frozenset(filter_indices),
                add_formulas=qfork_info.add_formulas,
                join_type=qfork_info.join_type,
                joining_node=joining_node,
            )
            mask_list.append(mask)

        return mask_list

    def mutate_split_node(self, node: formula_nodes.FormulaItem) -> formula_nodes.FormulaItem:
        if inspect_node.is_aggregate_function(node):
            # Remove all the extra children like LOD, BFB, ID, etc.
            assert isinstance(node, formula_nodes.FuncCall)
            node = formula_nodes.FuncCall.make(
                name=node.name,
                args=node.args,
                meta=node.meta,
            )

        return node

    def mutate_cropped_query(self, query: CompiledQuery) -> CompiledQuery:
        gb_aliases = {gb_formula.not_none_alias for gb_formula in query.group_by}
        inspect_env = InspectionEnvironment()
        agg_statuses: dict[bool, formula_nodes.FormulaItem] = {}  # Save a formula for each status for troubleshooting
        for select_formula in query.select:
            if select_formula.not_none_alias in gb_aliases:
                continue
            if inspect_expression.is_constant_expression(select_formula.formula_obj.expr):
                continue
            is_agg = inspect_expression.is_aggregate_expression(select_formula.formula_obj, env=inspect_env)
            agg_statuses[is_agg] = select_formula.formula_obj.expr

        if len(agg_statuses) == 0:
            # There aren't any non-dimension SELECTs.
            # Nothing to do in this case
            return query

        assert len(agg_statuses) == 1, f"Inconsistent aggregation status among SELECT items. Got: {agg_statuses}"
        agg_status = next(iter(agg_statuses))
        if not agg_status:
            # SELECT items are not aggregated
            # Remove GROUP BY
            query = query.clone(group_by=[])

        return query

lib/dl_query_processing/dl_query_processing/multi_query/splitters/query_fork.py
datalens-tech/datalens-backend
import attr

from dl_constants.enums import JoinType
import dl_formula.core.fork_nodes as formula_fork_nodes
from dl_formula.core.index import NodeHierarchyIndex
import dl_formula.core.nodes as formula_nodes
from dl_formula.inspect.env import InspectionEnvironment
import dl_formula.inspect.expression as inspect_expression
import dl_formula.inspect.node as inspect_node
from dl_query_processing.compilation.primitives import CompiledQuery
from dl_query_processing.multi_query.splitters.mask_based import (
    AliasedFormulaSplitMask,
    MultiQuerySplitter,
    QuerySplitMask,
    SubqueryType,
)
from dl_query_processing.utils.name_gen import PrefixedIdGen


@attr.s
class WinFuncQuerySplitter(MultiQuerySplitter):
    """
    Splitter that splits a query into two parts:
    one with window functions, and the other with everything else.
    All window functions are expected to be the topmost node in formulas.
    This is the way `QueryForkQuerySplitter` leaves them.
    """

    def _collect_win_func_split_indices_from_expression(
        self,
        node: formula_nodes.Formula,
        inspect_env: InspectionEnvironment,
    ) -> list[NodeHierarchyIndex]:
        """
        Collect indices of nodes that have to be split off.
        If the formula has a WF as its topmost node, then there will be an index for each of its children.
        Otherwise an index pointing to the whole expression is returned.
        """

        result: list[NodeHierarchyIndex] = []

        expr = node.expr
        expr_index = NodeHierarchyIndex(indices=(0,))
        if inspect_node.is_window_function(expr):
            for child_index, child_node, _ in inspect_expression.enumerate_autonomous_children(expr, prefix=expr_index):
                if inspect_expression.is_constant_expression(child_node, env=inspect_env):
                    # Ignore constants
                    continue

                assert not inspect_expression.is_window_expression(child_node, env=inspect_env)
                result.append(child_index)

        else:
            result.append(expr_index)
            assert not inspect_expression.is_window_expression(expr, env=inspect_env)

        return result

    def _collect_win_func_formula_masks(
        self,
        query: CompiledQuery,
        expr_id_gen: PrefixedIdGen,
    ) -> list[AliasedFormulaSplitMask]:
        """
        Iterate over all formulas in the query and collect indices of nodes that should be split off.
        Generate an `AliasedFormulaSplitMask` for each.
        """

        inspect_env = InspectionEnvironment()
        result: list[AliasedFormulaSplitMask] = []
        for query_part in self.parent_query_parts_for_splitting:
            formula_list = query.get_formula_list(query_part)
            for formula_idx, formula in enumerate(formula_list):
                indices = self._collect_win_func_split_indices_from_expression(
                    node=formula.formula_obj,
                    inspect_env=inspect_env,
                )
                result += [
                    AliasedFormulaSplitMask(
                        inner_node_idx=node_idx,
                        outer_node_idx=node_idx,
                        query_part=query_part,
                        formula_list_idx=formula_idx,
                        alias=expr_id_gen.get_id(),
                    )
                    for node_idx in indices
                ]

        return result

    def get_split_masks(
        self,
        query: CompiledQuery,
        expr_id_gen: PrefixedIdGen,
        query_id_gen: PrefixedIdGen,
    ) -> list[QuerySplitMask]:
        """
        Collect formula split masks for all formulas in the query and generate a single query split mask
        with all the formula masks in it.
        """

        formula_split_masks = self._collect_win_func_formula_masks(query, expr_id_gen=expr_id_gen)

        # Put it all together in the sub-query mask
        subquery_mask = QuerySplitMask(
            subquery_type=SubqueryType.default,
            subquery_id=query_id_gen.get_id(),
            formula_split_masks=tuple(formula_split_masks),
            filter_indices=frozenset(),
            add_formulas=(),
            join_type=JoinType.inner,
            joining_node=formula_fork_nodes.QueryForkJoiningWithList.make(condition_list=[]),
        )
        return [subquery_mask]

lib/dl_query_processing/dl_query_processing/multi_query/splitters/win_func.py
datalens-tech/datalens-backend
from typing import (
    Any,
    TypeVar,
)

import attr

import dl_formula.core.nodes as formula_nodes
from dl_query_processing.compilation.primitives import (
    CompiledMultiQuery,
    CompiledMultiQueryBase,
    CompiledQuery,
)


def build_requirement_subtree(multi_query: CompiledMultiQueryBase, top_query_id: str) -> CompiledMultiQueryBase:
    def _recursive_build(top_query: CompiledQuery) -> dict[str, CompiledQuery]:
        base_from_ids = set(multi_query.get_base_froms().keys())
        result: dict[str, CompiledQuery] = {}
        for from_obj in top_query.joined_from.froms:
            if from_obj.id not in base_from_ids:
                subquery = multi_query.get_query_by_id(from_obj.id)
                result[from_obj.id] = subquery
                result.update(_recursive_build(subquery))
        return result

    query_dict = _recursive_build(multi_query.get_query_by_id(top_query_id))
    queries = sorted(query_dict.values(), key=lambda item: item.id)
    result_multi_query = CompiledMultiQuery(queries=queries)  # FIXME: Replace with a lazy version
    return result_multi_query


@attr.s(frozen=True)
class CompiledMultiQueryPatch:
    patch_multi_query: CompiledMultiQueryBase = attr.ib(kw_only=True)


def apply_query_patch(multi_query: CompiledMultiQueryBase, patch: CompiledMultiQueryPatch) -> CompiledMultiQueryBase:
    if patch.patch_multi_query.is_empty():
        # Patch is empty
        return multi_query

    original_query_indices: dict[str, int] = {}
    queries: list[CompiledQuery] = []
    for idx, query in enumerate(multi_query.iter_queries()):
        original_query_indices[query.id] = idx
        queries.append(query)

    for query in patch.patch_multi_query.iter_queries():
        patch_idx = original_query_indices.get(query.id)
        if patch_idx is not None:
            queries[patch_idx] = query
        else:
            queries.append(query)

    return CompiledMultiQuery(queries=queries)


_FORMULA_NODE_TV = TypeVar("_FORMULA_NODE_TV", bound=formula_nodes.FormulaItem)


def remap_formula_obj_fields(
    node: _FORMULA_NODE_TV,
    field_name_map: dict[str, str],
) -> _FORMULA_NODE_TV:
    def remap_field(_node: formula_nodes.Field, *args: Any, **kwargs: Any) -> formula_nodes.Field:
        return formula_nodes.Field.make(name=field_name_map[_node.name])

    # In case the node is itself a Field
    if isinstance(node, formula_nodes.Field):
        return remap_field(node)  # type: ignore

    # For all other cases when it contains Fields
    updated_node = node.replace_nodes(
        match_func=lambda _node, parent_stack: isinstance(_node, formula_nodes.Field),
        replace_func=remap_field,  # type: ignore
    )
    return updated_node

lib/dl_query_processing/dl_query_processing/multi_query/tools.py
datalens-tech/datalens-backend
from dl_query_processing.pagination.post_paginator import QueryPostPaginator
from dl_query_processing.pagination.pre_paginator import QueryPrePaginator


class QueryPaginator:
    def get_pre_paginator(self) -> QueryPrePaginator:
        return QueryPrePaginator()

    def get_post_paginator(self) -> QueryPostPaginator:
        return QueryPostPaginator()

lib/dl_query_processing/dl_query_processing/pagination/paginator.py
datalens-tech/datalens-backend
from typing import (
    Generator,
    Iterable,
    Iterator,
)

from dl_query_processing.merging.primitives import (
    MergedQueryDataRow,
    MergedQueryDataStream,
)


class QueryPostPaginator:
    """
    Implements offset/limit pagination for `MergedQueryDataStream`
    """

    def _skip_to_offset(
        self,
        row_it: Iterator[MergedQueryDataRow],
        offset: int,
    ) -> None:
        if offset <= 0:
            return
        try:
            while offset > 0:
                next(row_it)
                offset -= 1
        except StopIteration:
            pass

    def _iter_until_limit(
        self, row_it: Iterator[MergedQueryDataRow], limit: int
    ) -> Generator[MergedQueryDataRow, None, None]:
        try:
            while limit > 0:
                yield next(row_it)
                limit -= 1
        except StopIteration:
            pass

    def post_paginate(self, merged_stream: MergedQueryDataStream) -> MergedQueryDataStream:
        row_it = iter(merged_stream.rows)
        overridden = False

        # Apply offset
        if merged_stream.meta.offset is not None:
            self._skip_to_offset(row_it, offset=merged_stream.meta.offset)
            overridden = True

        # Apply limit
        rows: Iterable[MergedQueryDataRow] = row_it
        if merged_stream.meta.limit is not None:
            rows = self._iter_until_limit(row_it, limit=merged_stream.meta.limit)
            overridden = True

        if overridden:
            merged_stream = merged_stream.clone(rows=rows)
        return merged_stream

lib/dl_query_processing/dl_query_processing/pagination/post_paginator.py
datalens-tech/datalens-backend
from typing import (
    Optional,
    TypeVar,
)

from dl_query_processing.legend.block_legend import BlockLegend


_VAL_TV = TypeVar("_VAL_TV")


def ifnull(value: _VAL_TV, null_value: _VAL_TV) -> _VAL_TV:
    if value is not None:
        return value
    return null_value


class QueryPrePaginator:
    """
    Prepares and optimizes query union and individual queries
    for pagination in SQL and postprocessing by tweaking the objects' attributes.
    """

    def _paginate_block(
        self,
        block_legend: BlockLegend,
        block_ind: int,
        limit: Optional[int],
        offset: Optional[int],
    ) -> BlockLegend:
        """
        Set pagination for block specified by index ``block_ind``
        and replace it in the block legend with the paginated copy.
        """
        updated_block = block_legend.blocks[block_ind].clone(
            limit=limit,
            offset=offset,
        )
        updated_blocks = block_legend.blocks[:block_ind] + [updated_block] + block_legend.blocks[block_ind + 1 :]
        return block_legend.clone(blocks=updated_blocks)

    def pre_paginate(self, block_legend: BlockLegend) -> BlockLegend:
        # Handle single-block case.
        # Just paginate the block in the source to avoid selecting too much data
        if len(block_legend.blocks) == 1:
            only_block = block_legend.blocks[0]
            if (
                # Block's own pagination is not specified
                (only_block.limit is None and only_block.offset is None)
                # But global pagination is
                and (block_legend.meta.limit is not None or block_legend.meta.offset is not None)
            ):
                # So apply global pagination to the block
                block_legend = self._paginate_block(
                    block_legend=block_legend,
                    block_ind=0,
                    limit=block_legend.meta.limit,
                    offset=block_legend.meta.offset,
                )
                # and remove global pagination from the legend
                block_legend = block_legend.clone(
                    meta=block_legend.meta.clone(limit=None, offset=None)  # disable global pagination
                )

        # Multi-block case.
        # We can't use the offset on the first block because it might be empty
        # and then the offset should be applied to the second block and so on,
        # and we don't know how much of it should have been consumed by each block.
        # But what we can do is apply the overall offset+limit as a limit to every single block.
        if len(block_legend.blocks) > 1 and block_legend.meta.limit is not None:
            master_limit = block_legend.meta.limit
            if block_legend.meta.offset is not None:
                master_limit += block_legend.meta.offset

            for block_ind in range(len(block_legend.blocks)):
                block_legend = self._paginate_block(
                    block_legend=block_legend,
                    block_ind=block_ind,
                    limit=master_limit,
                    offset=None,
                )

        return block_legend

lib/dl_query_processing/dl_query_processing/pagination/pre_paginator.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import logging
import re
from typing import (
    Any,
    ClassVar,
    Dict,
    Generator,
    Iterable,
    List,
    Optional,
    Sequence,
    Tuple,
)

import attr

from dl_app_tools.profiling_base import GenericProfiler
from dl_constants.enums import LegendItemType
from dl_query_processing.enums import QueryType
from dl_query_processing.execution.primitives import ExecutedQuery
from dl_query_processing.legend.block_legend import BlockSpec
from dl_query_processing.legend.field_legend import TemplateRoleSpec
from dl_query_processing.postprocessing.postprocessors.all import postprocess_data
from dl_query_processing.postprocessing.primitives import (
    PostprocessedQuery,
    PostprocessedQueryMetaInfo,
)


LOGGER = logging.getLogger(__name__)


# FIXME: Make output typing more specific - should comply with `PostprocessedValue`
#  (especially in MARKUP postprocessor).


@attr.s(frozen=True)
class ValueRestorerBase(abc.ABC):
    """
    Base class for value restorers.
    Main purpose is to get the value of a specific legend item
    from the given raw data row.
    """

    @abc.abstractmethod
    def restore_value(self, raw_row: Sequence[Any]) -> Any:
        raise NotImplementedError


@attr.s(frozen=True)
class IndexValueRestorer(ValueRestorerBase):
    """Restores value simply by retrieving it at the given index"""

    row_idx: int = attr.ib(kw_only=True)

    def restore_value(self, raw_row: Sequence[Any]) -> Any:
        return raw_row[self.row_idx]


@attr.s(frozen=True)
class TemplateValueRestorer(ValueRestorerBase):
    """Restores value simply by applying values to a string template"""

    template_re: ClassVar[re.Pattern] = re.compile(r"\{\{(?P<field>[^{}]+)\}\}")

    template: Optional[str] = attr.ib(kw_only=True)
    idx_by_field_id: Dict[str, int] = attr.ib(kw_only=True)

    def restore_value(self, raw_row: Sequence[Any]) -> Any:
        if self.template is None:
            return None

        def sub_value(field_match: re.Match) -> str:
            field_id = field_match.group("field").strip()
            idx = self.idx_by_field_id[field_id]
            return raw_row[idx]

        return self.template_re.sub(sub_value, self.template)


@attr.s
class DataPostprocessor:
    profiler_prefix: str = attr.ib(kw_only=True)

    def _make_value_restorers(
        self,
        executed_query: ExecutedQuery,
        block_spec: BlockSpec,
    ) -> List[ValueRestorerBase]:
        # Mapping of which field corresponds to which index in the data stream
        field_order = executed_query.meta.field_order
        assert field_order is not None
        idx_by_field_id = {field_id: idx for idx, field_id in field_order}
        result: List[ValueRestorerBase] = []
        for legend_item_id in block_spec.legend_item_ids:
            legend_item = block_spec.legend.get_item(legend_item_id)
            restorer: ValueRestorerBase
            if legend_item.obj.item_type == LegendItemType.placeholder:
                role_spec = legend_item.role_spec
                assert isinstance(role_spec, TemplateRoleSpec)
                restorer = TemplateValueRestorer(template=role_spec.template, idx_by_field_id=idx_by_field_id)
            else:
                restorer = IndexValueRestorer(row_idx=idx_by_field_id[legend_item.id])

            result.append(restorer)

        return result

    def get_postprocessed_data(
        self,
        executed_query: ExecutedQuery,
        block_spec: BlockSpec,
    ) -> PostprocessedQuery:
        query_meta = executed_query.meta
        data: Iterable[Sequence[Any]] = executed_query.rows

        ordered_value_restorers = self._make_value_restorers(executed_query=executed_query, block_spec=block_spec)

        def restore_order(data: Iterable[Sequence[Any]]) -> Generator[Tuple[Any, ...], None, None]:
            for row in data:
                yield tuple(restorer.restore_value(row) for restorer in ordered_value_restorers)

        with GenericProfiler(f"{self.profiler_prefix}-response-prepare"):
            result_fields_types = query_meta.detailed_types
            assert result_fields_types is not None

            postprocessed_data = postprocess_data(data, result_fields_types)

            # Filter out phantom fields, restore deduplicated columns, restore column order
            if executed_query.meta.query_type != QueryType.value_range:
                # FIXME: Dirty hack -> remove this if
                # If this is a value range query, then
                # a) restoration isn't necessary;
                # b) it doesn't work because two columns have the same field_id, but different wrappers (min and max)
                postprocessed_data = tuple(restore_order(postprocessed_data))

        LOGGER.info(
            f"Returning dataset data: {len(postprocessed_data)} rows " f"with {len(result_fields_types)} columns",
            extra=dict(
                fetched_data_statistics=dict(
                    row_count=len(postprocessed_data),
                    column_count=len(result_fields_types),
                ),
            ),
        )

        postprocessed_query = PostprocessedQuery(
            postprocessed_data=postprocessed_data, meta=PostprocessedQueryMetaInfo.from_exec_meta(exec_meta=query_meta)
        )

        return postprocessed_query

lib/dl_query_processing/dl_query_processing/postprocessing/postprocessor.py
datalens-tech/datalens-backend
from __future__ import annotations

import json
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Iterable,
    Optional,
    Sequence,
)

from dl_constants.enums import UserDataType
from dl_formula.core.datatype import DataType
from dl_query_processing.postprocessing.postprocessors.datetime import (
    make_postprocess_datetimetz,
    postprocess_datetime,
    postprocess_genericdatetime,
)
from dl_query_processing.postprocessing.postprocessors.geo import (
    postprocess_geopoint,
    postprocess_geopolygon,
)
from dl_query_processing.postprocessing.postprocessors.markup import postprocess_markup
from dl_query_processing.postprocessing.primitives import PostprocessedData


if TYPE_CHECKING:
    from dl_query_processing.compilation.primitives import DetailedType  # type: ignore  # TODO: fix


def stringify_or_null(value: Any) -> Optional[str]:
    if value is None:
        return value
    return str(value)


def postprocess_array(value: Optional[Iterable[Any]]) -> Optional[Iterable[Optional[str]]]:
    if value is None:
        return None
    return json.dumps(value, ensure_ascii=False)  # Frontend fully supports non-ASCII JSON


TYPE_PROCESSORS = {
    UserDataType.datetime: postprocess_datetime,
    # parametrized: UserDataType.datetimetz
    UserDataType.genericdatetime: postprocess_genericdatetime,
    UserDataType.geopoint: postprocess_geopoint,
    UserDataType.geopolygon: postprocess_geopolygon,
    UserDataType.markup: postprocess_markup,
    UserDataType.array_int: postprocess_array,
    UserDataType.array_float: postprocess_array,
    UserDataType.array_str: postprocess_array,
    UserDataType.tree_str: postprocess_array,
}


def get_type_processor(field_type_info: Optional[DetailedType]) -> Callable[[Any], Any]:
    if field_type_info is None:
        return lambda val: val

    # Basic
    result = TYPE_PROCESSORS.get(field_type_info.data_type)
    if result is not None:
        return result  # type: ignore  # TODO: fix

    # Parmetrized
    if field_type_info.data_type == UserDataType.datetimetz:
        assert field_type_info.formula_data_type == DataType.DATETIMETZ
        assert field_type_info.formula_data_type_params
        assert field_type_info.formula_data_type_params.timezone
        return make_postprocess_datetimetz(field_type_info.formula_data_type_params.timezone)

    # Default
    return stringify_or_null


# TODO FIX: Turn into generator when response streaming become real
def postprocess_data(
    data: Iterable[Sequence[Any]],
    field_types: Iterable[Optional[DetailedType]],
) -> PostprocessedData:
    columns_processors = [get_type_processor(ft) for ft in field_types]

    return tuple(tuple(processor(col) for processor, col in zip(columns_processors, row)) for row in data)

lib/dl_query_processing/dl_query_processing/postprocessing/postprocessors/all.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from typing import Optional

import pytz


def postprocess_datetime(value: datetime.datetime) -> Optional[str]:
    if value is None:
        return value

    if isinstance(value, str):
        value = datetime.datetime.fromisoformat(value)
    if value.tzinfo is not None:
        value = value.astimezone(datetime.timezone.utc).replace(tzinfo=None)

    return value.isoformat()


def postprocess_genericdatetime(value: datetime.datetime) -> Optional[str]:
    if value is None:
        return value

    if isinstance(value, str):
        value = datetime.datetime.fromisoformat(value)

    return value.replace(tzinfo=None).isoformat()


def make_postprocess_datetimetz(tzname: str):  # type: ignore  # TODO: fix
    tzobj = pytz.timezone(tzname)

    def _postprocess_datetimetz(value: datetime.datetime, tzobj=tzobj) -> Optional[str]:  # type: ignore  # TODO: fix
        if value is None:
            return value

        if value.tzinfo is None:
            value = value.replace(tzinfo=datetime.timezone.utc)
        return value.astimezone(tzobj).isoformat()

    return _postprocess_datetimetz

lib/dl_query_processing/dl_query_processing/postprocessing/postprocessors/datetime.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import lru_cache
import json
from typing import Optional


DUMMY_GEO_VALUE = "__DUMMY_GEO_VALUE__"


@lru_cache(maxsize=2**10)
def postprocess_geopoint(value) -> Optional[str]:  # type: ignore  # TODO: fix
    value = str(value)
    if value == DUMMY_GEO_VALUE:
        return None
    try:
        geopoint = json.loads(value)
        if not isinstance(geopoint, list) or not len(geopoint) == 2:
            return None
        if not isinstance(geopoint[0], (int, float)) or not (-90 <= geopoint[0] <= 90):
            return None
        if not isinstance(geopoint[1], (int, float)) or not (-180 <= geopoint[1] <= 180):
            return None
    except ValueError:
        return None
    return value


# geopoint_re = re.compile(r'\[-?([0-9]|[1-8][0-9]|90)(\.\d+)?,\s*-?([0-9]|[1-9][0-9]|1[1-7][0-9]|180)(\.\d+)?\]')
#
#
# def postprocess_geopoint_re(value) -> str:
#     """
#     In [61]: %timeit postprocess_geopoint_re('[55.832995, -37.634928]')
#     1.14 µs ± 33.3 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
#
#     In [63]: %timeit postprocess_geopoint('[55.832995, -37.634928]')
#     3.75 µs ± 252 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
#     """
#
#     value = str(value)
#     if geopoint_re.match(value):
#         return value
#     else:
#         return None


@lru_cache(maxsize=2**10)
def postprocess_geopolygon(value) -> Optional[str]:  # type: ignore  # TODO: fix
    value = str(value)
    if value == DUMMY_GEO_VALUE:
        return None
    try:
        multipolygon = json.loads(value)
        if not isinstance(multipolygon, list) or not multipolygon:
            return None
        for sub_polygon in multipolygon:
            if not isinstance(sub_polygon, list) or not sub_polygon:
                return None
            for geopoint in sub_polygon:
                if not isinstance(geopoint, list) or not len(geopoint) == 2:
                    return None
                if not isinstance(geopoint[0], (int, float)) or not (-90 <= geopoint[0] <= 90):
                    return None
                if not isinstance(geopoint[1], (int, float)) or not (-180 <= geopoint[1] <= 180):
                    return None
    except ValueError:
        return None
    return value


# geopolygon_re = re.compile(
#     r'\s*\[\s*(\[\s*(\[\s*'
#     r'-?([0-9]|[1-8][0-9]|90)(\.\d+)?\s*,\s*'
#     r'-?([0-9]|[1-9][0-9]|1[1-7][0-9]|180)(\.\d+)?\s*'
#     r'\]\s*,\s*)*'
#     r'\[\s*-?([0-9]|[1-8][0-9]|90)(\.\d+)?\s*,\s*'
#     r'-?([0-9]|[1-9][0-9]|1[1-7][0-9]|180)(\.\d+)?\s*'
#     r'\]\s*\]\s*,\s*)*'
#     r'\[\s*(\[\s*'
#     r'-?([0-9]|[1-8][0-9]|90)(\.\d+)?\s*,\s*'
#     r'-?([0-9]|[1-9][0-9]|1[1-7][0-9]|180)(\.\d+)?\s*'
#     r'\]\s*,\s*)*'
#     r'\[\s*'
#     r'-?([0-9]|[1-8][0-9]|90)(\.\d+)?\s*,\s*'
#     r'-?([0-9]|[1-9][0-9]|1[1-7][0-9]|180)(\.\d+)?\s*'
#     r'\]\s*\]\s*\]\s*'
# )
#
#
# def postprocess_geopolygon_re(value) -> str:
#     """
#     In [67]: %timeit postprocess_geopolygon(rus_poly)
#     1.02 ms ± 23.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
#
#     In [68]: %timeit postprocess_geopolygon_re(rus_poly)
#     738 µs ± 17.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
#     """
#
#     value = str(value)
#     if geopolygon_re.match(value):
#         return value
#     else:
#         return None

lib/dl_query_processing/dl_query_processing/postprocessing/postprocessors/geo.py
datalens-tech/datalens-backend
from __future__ import annotations

from dataclasses import dataclass
from typing import (
    Any,
    ClassVar,
    Generic,
    Optional,
    Tuple,
    Type,
    TypeVar,
    Union,
)


_NODE_TV = TypeVar("_NODE_TV")
NodeActual = Union[_NODE_TV, str]  # what can act as a node
NodeInput = Union[_NODE_TV, str]  # what can be passed to node-builders


class MarkupProcessingBase(Generic[_NODE_TV]):
    quot: ClassVar[str] = '"'
    lpar: ClassVar[str] = "("
    rpar: ClassVar[str] = ")"
    sep: ClassVar[str] = " "

    node_text: ClassVar[str] = ""
    node_concat: ClassVar[str] = "c"
    node_url: ClassVar[str] = "a"
    node_i: ClassVar[str] = "i"
    node_b: ClassVar[str] = "b"

    _node_cls: ClassVar[Type[_NODE_TV]]

    _dbg: bool = False

    empty_verbose_node = dict(type="concat", children=[])

    def _dbg_print(self, msg: str, *args: Any) -> None:
        if not self._dbg:
            return
        print("DBG:", msg % args)

    def _make_node(self, funcname: str, *funcargs: NodeActual) -> _NODE_TV:
        raise NotImplementedError

    def _unpack_node(self, node: _NODE_TV) -> Tuple[str, Tuple[NodeActual, ...]]:
        raise NotImplementedError

    class DumpError(ValueError):
        """..."""

    class ParseError(ValueError):
        """..."""

        # effective args: message, data, pos, name_of_expected_stuff

    # Helper functions

    def _quote_i(self, value: str) -> str:
        # `"` -> `""`
        if not isinstance(value, str):
            raise ValueError("should only quote strings")
        return value.replace(self.quot, self.quot + self.quot)

    def _quote(self, value: str) -> str:
        """
        ...

        >>> MarkupProcessing()._quote('ab"cd')
        '"ab""cd"'
        >>> MarkupProcessing()._quote('"') == '"' * 4
        True
        """
        return self.quot + self._quote_i(value) + self.quot

    def _proc_arg(self, value: NodeInput) -> NodeActual:
        if isinstance(value, self._node_cls):
            return value
        if isinstance(value, str):
            # return self._make_node(self.node_text, value)
            return value  # more compact, more db-expr-friendly.
        raise Exception("Unexpected value", value)

    # Node-building

    def n_concat(self, *children: NodeInput) -> _NODE_TV:
        return self._make_node(self.node_concat, *(self._proc_arg(child) for child in children))

    def n_url(self, addr: str, title: Optional[NodeInput] = None) -> _NODE_TV:
        return self._make_node(self.node_url, str(addr), self._proc_arg(title or addr))

    def n_i(self, child: NodeInput) -> _NODE_TV:
        return self._make_node(self.node_i, self._proc_arg(child))

    def n_b(self, child: NodeInput) -> _NODE_TV:
        return self._make_node(self.node_b, self._proc_arg(child))

    # Node-structure into a string

    def _dump_i(self, node: NodeActual) -> str:
        if isinstance(node, str):
            return self._quote(node)
        if not isinstance(node, self._node_cls):
            raise ValueError("Value should be a node or a str")
        funcname, funcargs = self._unpack_node(node)
        return (
            self.lpar
            + funcname
            + self.sep
            + self.sep.join(
                # recurse
                self._dump_i(arg)
                for arg in funcargs
            )
            + self.rpar
        )

    def dump(self, node: _NODE_TV) -> str:
        if not isinstance(node, self._node_cls):
            raise self.DumpError("Outer value should be a node", type(node))
        return self._dump_i(node)

    # string into a node-structure

    def _get_at(self, data: str, pos: int, size: int, require: bool = True) -> str:
        result = data[pos : pos + size]
        if require and not result:
            raise self.ParseError("Unexpected data end", data, pos)
        return result

    def _is_at(self, data: str, pos: int, substr: str) -> bool:
        return self._get_at(data, pos, len(substr)) == substr

    def _find_next(self, data: str, substr: str, start: int, end: Optional[int] = None) -> Optional[int]:
        try:
            return data.index(substr, start, end)
        except ValueError:
            return None

    def _parse_i(self, data: str, pos: int = 0) -> Tuple[NodeActual, int]:
        if self._dbg:
            self._dbg_print("_parse_i(%r <-%r-> %r)", data[:pos], pos, data[pos:])

        if self._is_at(data, pos, self.lpar):  # ({funcname} {arg1} {arg2} ...)
            if self._dbg:
                self._dbg_print("lpar @ %r: %r", pos, data[pos:])
            pos += len(self.lpar)

            next_rpar = self._find_next(data, self.rpar, pos)
            if next_rpar is None:
                raise self.ParseError("Unexpected data end", data, pos, "rpar")

            next_sep = self._find_next(data, self.sep, pos, next_rpar)
            if next_sep is None:  # '({funcname})'
                # Could do `next_sep = next_rpar` and fall through the `while`
                return self._make_node(data[pos:next_rpar]), next_rpar + len(self.rpar)

            # else: '({funcname} ...)'
            result = [data[pos:next_sep]]
            pos = next_sep + len(self.sep)
            while not self._is_at(data, pos, self.rpar):
                last_pos = pos
                node, pos = self._parse_i(data, pos)
                assert pos > last_pos, "should advance"
                result.append(node)
                if self._is_at(data, pos, self.rpar):
                    break
                if not self._is_at(data, pos, self.sep):
                    raise ValueError("Unexpected non-separator", data, pos, "sep")
                pos += len(self.sep)

            assert self._is_at(data, pos, self.rpar), "should have ended up at rpar"
            return self._make_node(*result), pos + len(self.rpar)

        if self._is_at(data, pos, self.quot):
            if self._dbg:
                self._dbg_print("quot @ %r: %r", pos, data[pos:])

            pos += len(self.quot)
            result_pieces = []
            last_pos = pos
            while True:
                next_quot = self._find_next(data, self.quot, pos)
                if next_quot is None:
                    raise self.ParseError("Unexpected data end", data, pos, "quot")
                if self._dbg:
                    self._dbg_print("next_quot @ %r: %r", next_quot, data[next_quot:])

                # check for quot2-after-quot
                quot2 = self.quot  # simple doubling
                if self._dbg:
                    self._dbg_print(
                        "_parse_i string %r -%r-> %r <-%r- %r",
                        data[:pos],
                        pos,
                        data[pos:next_quot],
                        next_quot,
                        data[next_quot:],
                    )

                result_pieces.append(data[pos:next_quot])
                if self._is_at(data, next_quot + len(self.quot), quot2):
                    result_pieces.append(quot2)
                    pos = next_quot + len(self.quot) + len(quot2)
                else:  # something else after the next_quot
                    self._dbg_print("end quot @ %r", next_quot)
                    return "".join(result_pieces), next_quot + len(self.quot)

                assert pos > last_pos, "should advance"
                last_pos = pos

        raise self.ParseError("Unexpected data", data, pos)

    def parse(self, data: Optional[str]) -> _NODE_TV:
        if data is None:
            return None  # type: ignore  # TODO: fix
        if not data.startswith(self.lpar) or not data.endswith(self.rpar):
            raise self.ParseError("Malformed data", data)
        node, pos = self._parse_i(data, 0)
        if pos != len(data):
            raise self.ParseError("Extra data at the end", data, pos)
        assert isinstance(node, self._node_cls)
        return node

    _verbalized_info = {
        node_text: dict(name="text", arg="content"),
        node_concat: dict(name="concat", args="children"),
        node_url: dict(name="url", argnames=("url", "content")),
        node_i: dict(name="italics", arg="content"),
        node_b: dict(name="bold", arg="content"),
    }

    def _argcount_mismatch(self, node, **kwargs):  # type: ignore  # TODO: fix
        # Would be nice to do `self.DumpError("Argcount mismatch", node)` here,
        # but NULL-related behavior makes it unfeasible.
        # Thus, return an empty node (equivalent to an empty string).
        return self.empty_verbose_node

    def _verbalize_i(self, node: NodeActual) -> Any:
        """A kind-of the other way from dump: more verbose serializable structure. Still a dump."""
        if isinstance(node, str):
            return dict(type="text", content=node)

        if not isinstance(node, self._node_cls):
            raise self.DumpError("Value should be a node", type(node))

        funcname, funcargs = self._unpack_node(node)

        if funcname == self.node_concat:
            return dict(type="concat", children=[self._verbalize_i(arg) for arg in funcargs])
        if funcname == self.node_url:
            if len(funcargs) != 2:
                return self._argcount_mismatch(node=node, funcname=funcname, funcargs=funcargs, expected_args=2)
            url, content = funcargs
            return dict(type="url", url=url, content=self._verbalize_i(content))
        if funcname in (self.node_i, self.node_b):
            if len(funcargs) != 1:
                return self._argcount_mismatch(node=node, funcname=funcname, funcargs=funcargs, expected_args=1)
            return dict(
                type={self.node_i: "italics", self.node_b: "bold"}[funcname], content=self._verbalize_i(funcargs[0])
            )
        raise self.DumpError("Unknown func", node)

    def verbalize(self, node: Optional[_NODE_TV]) -> Any:
        """A kind-of the other way from dump: more verbose serializable structure. Still a dump."""
        if node is None:
            return None
        if not isinstance(node, self._node_cls):
            raise self.DumpError("Outer value should be a node", type(node))
        return self._verbalize_i(node)

    def parse_and_verbalize(self, data: str) -> Any:
        return self.verbalize(self.parse(data))


_TNode = tuple
_TNodeActual = Union[_TNode, str]


class MarkupProcessing(MarkupProcessingBase[_TNode]):
    _node_cls: ClassVar[Type[_TNode]] = tuple  # Node  # for isinstance checks

    def _make_node(self, funcname: str, *funcargs: _TNodeActual) -> _TNode:
        # return self._node_cls(funcname, *funcargs)
        return self._node_cls((funcname,) + funcargs)

    def _unpack_node(self, node: _TNode) -> Tuple[str, Tuple[_TNodeActual, ...]]:
        return node[0], node[1:]


@dataclass
class DCNode:
    funcname: str
    funcargs: Tuple[NodeActual, ...]


_DCNode = DCNode
_DCNodeActual = Union[_DCNode, str]


class MarkupProcessingDC(MarkupProcessingBase[_DCNode]):
    _node_cls: ClassVar[Type[_DCNode]] = DCNode

    def _make_node(self, funcname: str, *funcargs: _DCNodeActual) -> _DCNode:
        return self._node_cls(funcname=funcname, funcargs=funcargs)

    def _unpack_node(self, node: _DCNode) -> Tuple[str, Tuple[_DCNodeActual, ...]]:
        return node.funcname, node.funcargs


MARKUP_PROCESSING = MarkupProcessing()


postprocess_markup = MARKUP_PROCESSING.parse_and_verbalize

lib/dl_query_processing/dl_query_processing/postprocessing/postprocessors/markup.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    List,
    Optional,
    Tuple,
    Type,
    TypeVar,
    Union,
)

import attr

from dl_constants.types import TBIDataValue
from dl_core.utils import attrs_evolve_to_subclass
import dl_query_processing.exc
from dl_query_processing.execution.primitives import ExecutedQueryMetaInfo
from dl_query_processing.legend.block_legend import BlockSpec
from dl_query_processing.legend.field_legend import Legend


PostprocessedValue = Union[TBIDataValue, dict]
PostprocessedRow = Tuple[PostprocessedValue, ...]
PostprocessedData = Tuple[PostprocessedRow, ...]


_META_TV = TypeVar("_META_TV", bound="ExecutedQueryMetaInfo")


@attr.s
class PostprocessedQueryMetaInfo(ExecutedQueryMetaInfo):
    @classmethod
    def from_exec_meta(
        cls: Type[_META_TV],
        exec_meta: ExecutedQueryMetaInfo,
    ) -> _META_TV:
        return attrs_evolve_to_subclass(
            cls=cls,
            inst=exec_meta,
        )


@attr.s
class PostprocessedQuery:
    postprocessed_data: PostprocessedData = attr.ib(kw_only=True, default=())
    meta: PostprocessedQueryMetaInfo = attr.ib(kw_only=True, factory=PostprocessedQueryMetaInfo)

    def clone(self, **updates: Any) -> PostprocessedQuery:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class PostprocessedQueryBlock(BlockSpec):
    postprocessed_query: PostprocessedQuery = attr.ib(kw_only=True)

    @classmethod
    def from_block_spec(cls, block_spec: BlockSpec, postprocessed_query: PostprocessedQuery) -> PostprocessedQueryBlock:
        return attrs_evolve_to_subclass(
            cls=cls,
            inst=block_spec,
            postprocessed_query=postprocessed_query,
        )


@attr.s(frozen=True)
class PostprocessedQueryUnion:
    blocks: List[PostprocessedQueryBlock] = attr.ib(kw_only=True)
    legend: Legend = attr.ib(kw_only=True)
    offset: Optional[int] = attr.ib(kw_only=True)
    limit: Optional[int] = attr.ib(kw_only=True)

    def clone(self, **updates: Any) -> PostprocessedQueryUnion:
        return attr.evolve(self, **updates)

    @property
    def single_postprocessed_query(self) -> PostprocessedQuery:
        if len(self.blocks) != 1:
            raise dl_query_processing.exc.MultipleBlocksUnsupportedError()
        return self.blocks[0].postprocessed_query

lib/dl_query_processing/dl_query_processing/postprocessing/primitives.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Any,
    Optional,
    Type,
)

import attr

import dl_formula.core.exc as formula_exc
from dl_formula.core.message_ctx import FormulaErrorCtx
import dl_query_processing.exc


FEATURE_ERROR_CODE = formula_exc.TranslationError.default_code + ("FEATURE_UNAVAILABLE",)


@attr.s
class FormulaErrorCollector:
    _errors: list[FormulaErrorCtx] = attr.ib(init=False, factory=list)

    def get_errors(self, feature_errors: bool = True) -> list[FormulaErrorCtx]:
        return [err for err in self._errors if feature_errors or err.code != FEATURE_ERROR_CODE]

    def __enter__(self) -> "FormulaErrorCollector":
        return self

    def __exit__(self, exc_type: Optional[Type[Exception]], exc_val: Optional[Exception], exc_tb: Any) -> bool:
        if isinstance(exc_val, dl_query_processing.exc.FormulaHandlingError):
            self._errors += exc_val.errors
            return True
        return False

lib/dl_query_processing/dl_query_processing/translation/error_collector.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    Callable,
    Optional,
)

import attr

from dl_core.db.elements import SchemaColumn
from dl_core.query.expression import (
    ExpressionCtx,
    JoinOnExpressionCtx,
    OrderByExpressionCtx,
)
from dl_core.utils import (
    attrs_evolve_to_subclass,
    attrs_evolve_to_superclass,
)
from dl_formula.core.dialect import DialectCombo
import dl_formula.core.exc as formula_exc
from dl_formula.definitions.flags import ContextFlag
from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.inspect.expression import is_aggregate_expression
from dl_formula.translation.env import (
    TranslationEnvironment,
    TranslationStats,
)
from dl_formula.translation.translator import translate
from dl_query_processing.column_registry import ColumnRegistry
from dl_query_processing.compilation.primitives import (
    CompiledFormulaInfo,
    CompiledJoinOnFormulaInfo,
    CompiledOrderByFormulaInfo,
    CompiledQuery,
)
from dl_query_processing.compilation.type_mapping import (
    DEFAULT_DATA_TYPE,
    FORMULA_TO_BI_TYPES,
)
import dl_query_processing.exc
from dl_query_processing.translation.primitives import (
    DetailedType,
    ExpressionCtxExt,
    TranslatedFlatQuery,
    TranslatedJoinedFromObject,
    TranslatedQueryMetaInfo,
)


@attr.s
class FlatQueryTranslator:
    _columns: ColumnRegistry = attr.ib(kw_only=True)
    _inspect_env: Optional[InspectionEnvironment] = attr.ib(kw_only=True, factory=InspectionEnvironment)  # noqa
    _function_scopes: int = attr.ib(kw_only=True)
    _dialect: DialectCombo = attr.ib(kw_only=True)
    # Generated attributes
    _trans_env: TranslationEnvironment = attr.ib(init=False)
    _avatar_alias_mapper: Callable[[str], str] = attr.ib(kw_only=True)
    _collect_stats: bool = attr.ib(kw_only=True, default=False)

    def __attrs_post_init__(self) -> None:
        self._trans_env = self._make_trans_env()

    def _make_trans_env(self) -> TranslationEnvironment:
        return TranslationEnvironment(
            dialect=self._dialect,
            required_scopes=self._function_scopes,
            field_types=self._columns.get_column_formula_types(),
            field_names=self._columns.get_multipart_column_names(  # type: ignore  # TODO: fix
                avatar_alias_mapper=self._avatar_alias_mapper
            ),
        )

    def translate_formula(
        self,
        comp_formula: CompiledFormulaInfo,
        collect_errors: bool = False,
        as_condition: bool = False,
    ) -> ExpressionCtxExt:
        context_flags = 0
        if as_condition:  # For MSSQL & ORACLE
            context_flags = ContextFlag.REQ_CONDITION

        try:
            translation_ctx = translate(
                formula=comp_formula.formula_obj,
                collect_errors=collect_errors,
                collect_stats=self._collect_stats,
                context_flags=context_flags,
                env=self._trans_env,
            )
        except formula_exc.TranslationError as err:
            raise dl_query_processing.exc.FormulaHandlingError(*err.errors) from err

        user_type = FORMULA_TO_BI_TYPES.get(translation_ctx.data_type, FORMULA_TO_BI_TYPES[DEFAULT_DATA_TYPE])  # type: ignore  # TODO: fix

        # Use ExpressionCtxExt because the formula might be used in SELECT
        return ExpressionCtxExt(
            expression=translation_ctx.expression,  # type: ignore  # TODO: fix
            formula_data_type=translation_ctx.data_type,
            formula_data_type_params=translation_ctx.data_type_params,
            user_type=user_type,
            alias=comp_formula.alias,
            avatar_ids=list(comp_formula.avatar_ids),
            original_field_id=comp_formula.original_field_id,
        )

    def translate_order_by_formula(
        self,
        comp_formula: CompiledOrderByFormulaInfo,
        collect_errors: bool = False,
    ) -> OrderByExpressionCtx:
        assert isinstance(comp_formula, CompiledOrderByFormulaInfo)
        translated_formula = self.translate_formula(comp_formula=comp_formula, collect_errors=collect_errors)
        return attrs_evolve_to_subclass(
            cls=OrderByExpressionCtx,
            inst=attrs_evolve_to_superclass(cls=ExpressionCtx, inst=translated_formula),
            direction=comp_formula.direction,
        )

    def translate_join_on_formula(
        self,
        comp_formula: CompiledJoinOnFormulaInfo,
        collect_errors: bool = False,
    ) -> JoinOnExpressionCtx:
        assert isinstance(comp_formula, CompiledJoinOnFormulaInfo)
        translated_formula = self.translate_formula(
            comp_formula=comp_formula,
            collect_errors=collect_errors,
            as_condition=True,
        )
        return attrs_evolve_to_subclass(
            cls=JoinOnExpressionCtx,
            inst=attrs_evolve_to_superclass(cls=ExpressionCtx, inst=translated_formula),
            left_id=comp_formula.left_id,
            right_id=comp_formula.right_id,
            join_type=comp_formula.join_type,
        )

    def _get_detailed_types(
        self, compiled_flat_query: CompiledQuery, translated_select: list[ExpressionCtxExt]
    ) -> Optional[list[Optional[DetailedType]]]:
        field_order = compiled_flat_query.meta.field_order
        if field_order is None:
            return None

        result: list[Optional[DetailedType]] = []
        for field_idx, field_id in field_order:
            detailed_type: Optional[DetailedType]
            if field_id is None:
                detailed_type = None
            else:
                field_expr = translated_select[field_idx]
                assert field_expr.user_type is not None
                detailed_type = DetailedType(
                    field_id=field_id,
                    data_type=field_expr.user_type,
                    formula_data_type=field_expr.formula_data_type,
                    formula_data_type_params=field_expr.formula_data_type_params,
                )

            result.append(detailed_type)

        return result

    def _generate_colummn_list_for_query(
        self,
        query_id: str,
        translated_select: list[ExpressionCtxExt],
    ) -> list[SchemaColumn]:
        """
        Generate list of columns that can be used when selecting from this query
        as if it were a table.
        """

        column_list: list[SchemaColumn] = []
        for translated_formula in translated_select:
            column = SchemaColumn(
                name=translated_formula.alias,  # type: ignore  # TODO: fix
                title=translated_formula.alias,
                source_id=query_id,
                native_type=None,
                user_type=translated_formula.user_type,
                has_auto_aggregation=False,
                lock_aggregation=False,
            )
            column_list.append(column)

        return column_list

    def translate_flat_query(
        self,
        compiled_flat_query: CompiledQuery,
        collect_errors: bool = False,
    ) -> TranslatedFlatQuery:
        where: list[CompiledFormulaInfo] = []
        having: list[CompiledFormulaInfo] = []
        for formula in compiled_flat_query.filters:
            if is_aggregate_expression(node=formula.formula_obj, env=self._inspect_env):  # type: ignore  # TODO: fix
                having.append(formula)
            else:
                where.append(formula)

        translated_select = [
            self.translate_formula(formula, collect_errors=collect_errors) for formula in compiled_flat_query.select
        ]

        column_list = self._generate_colummn_list_for_query(
            query_id=compiled_flat_query.id,
            translated_select=translated_select,
        )

        translated_meta = TranslatedQueryMetaInfo.from_comp_meta(
            compiled_flat_query.meta,
            detailed_types=self._get_detailed_types(
                compiled_flat_query=compiled_flat_query,
                translated_select=translated_select,
            ),
        )

        alias = self._avatar_alias_mapper(compiled_flat_query.id)
        return TranslatedFlatQuery(
            id=compiled_flat_query.id,
            alias=alias,
            level_type=compiled_flat_query.level_type,
            select=translated_select,
            group_by=[
                self.translate_formula(formula, collect_errors=collect_errors)
                for formula in compiled_flat_query.group_by
            ],
            order_by=[
                self.translate_order_by_formula(formula, collect_errors=collect_errors)
                for formula in compiled_flat_query.order_by
            ],
            where=[
                self.translate_formula(formula, collect_errors=collect_errors, as_condition=True) for formula in where
            ],
            having=[
                self.translate_formula(formula, collect_errors=collect_errors, as_condition=True) for formula in having
            ],
            join_on=[
                self.translate_join_on_formula(formula, collect_errors=collect_errors)
                for formula in compiled_flat_query.join_on
            ],
            joined_from=TranslatedJoinedFromObject(
                root_from_id=compiled_flat_query.joined_from.root_from_id,
                froms=compiled_flat_query.joined_from.froms,
            ),
            limit=compiled_flat_query.limit,
            offset=compiled_flat_query.offset,
            column_list=column_list,
            meta=translated_meta,
        )

    def get_collected_stats(self) -> TranslationStats:
        return self._trans_env.translation_stats

lib/dl_query_processing/dl_query_processing/translation/flat_translator.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
import logging
from typing import (
    Callable,
    Optional,
)

import attr

from dl_formula.core.dialect import DialectCombo
from dl_formula.core.message_ctx import FormulaErrorCtx
from dl_formula.inspect.env import InspectionEnvironment
from dl_formula.translation.env import TranslationStats
from dl_query_processing.column_registry import ColumnRegistry
from dl_query_processing.compilation.primitives import (
    AvatarFromObject,
    CompiledMultiQueryBase,
    CompiledQuery,
    FromObject,
    SubqueryFromObject,
)
from dl_query_processing.enums import ExecutionLevel
from dl_query_processing.translation.error_collector import FormulaErrorCollector
from dl_query_processing.translation.flat_translator import FlatQueryTranslator
from dl_query_processing.translation.primitives import (
    TranslatedFlatQuery,
    TranslatedMultiQuery,
    TranslatedMultiQueryBase,
)


LOGGER = logging.getLogger(__name__)


@attr.s
class MultiLevelQueryTranslator:
    _source_db_columns: ColumnRegistry = attr.ib(kw_only=True)
    _inspect_env: InspectionEnvironment = attr.ib(kw_only=True)
    _function_scopes: int = attr.ib(kw_only=True)
    _verbose_logging: bool = attr.ib(kw_only=True, default=False)  # noqa
    _avatar_alias_mapper: Callable[[str], str] = attr.ib(kw_only=True, default=lambda s: s)  # noqa
    _dialect: DialectCombo = attr.ib(kw_only=True)
    _compeng_dialect: Optional[DialectCombo] = attr.ib(kw_only=True)

    # Attributes for source_db
    _collect_stats: bool = attr.ib(kw_only=True, default=False)

    def make_source_db_flat_translator(self, columns: ColumnRegistry) -> FlatQueryTranslator:
        return FlatQueryTranslator(
            columns=columns,
            inspect_env=self._inspect_env,
            function_scopes=self._function_scopes,
            dialect=self._dialect,
            avatar_alias_mapper=self._avatar_alias_mapper,
            collect_stats=self._collect_stats,
        )

    def make_compeng_flat_translator(self, columns: ColumnRegistry) -> FlatQueryTranslator:
        assert self._compeng_dialect is not None
        return FlatQueryTranslator(
            columns=columns,
            inspect_env=self._inspect_env,
            function_scopes=self._function_scopes,
            avatar_alias_mapper=self._avatar_alias_mapper,
            collect_stats=self._collect_stats,
            dialect=self._compeng_dialect,
        )

    def _log_info(self, *args, **kwargs) -> None:  # type: ignore  # TODO: fix
        if self._verbose_logging:
            LOGGER.info(*args, **kwargs)

    def _log_collected_stats(
        self,
        stats_lists: dict[ExecutionLevel, list[TranslationStats]],
    ) -> None:
        """
        Log stats collected by all translators during translation
        """

        for level_type in sorted(stats_lists, key=lambda lt: lt.name):
            combined_stats = TranslationStats.combine(*stats_lists[level_type])
            data = {
                "functions": [
                    dict(signature.as_dict(), weight=weight)
                    for signature, weight in combined_stats.function_usage_weights.items()
                ],
                "level_type": level_type.name,
                "cache_hits": combined_stats,
            }
            LOGGER.info(
                f"Function translation statistics for {level_type.name}",
                extra=dict(function_translation_statistics=data),
            )

    def _log_query_complexity_stats(self, compiled_multi_query: CompiledMultiQueryBase) -> None:
        LOGGER.info(
            "Query structural info",
            extra=dict(
                query_struct_info=dict(
                    complexity=compiled_multi_query.get_complexity(),
                    subquery_count=compiled_multi_query.query_count(),
                )
            ),
        )

    def _get_flat_translator_for_level_type(
        self,
        level_type: ExecutionLevel,
        columns: ColumnRegistry,
    ) -> FlatQueryTranslator:
        flat_trans: FlatQueryTranslator
        if level_type == ExecutionLevel.source_db:
            flat_trans = self.make_source_db_flat_translator(columns=columns)
        elif level_type == ExecutionLevel.compeng:
            flat_trans = self.make_compeng_flat_translator(columns=columns)
        else:
            raise ValueError(level_type)
        return flat_trans

    def _make_column_reg_for_from_obj(
        self,
        from_obj: FromObject,
        translated_queries_by_id: dict[str, TranslatedFlatQuery],
    ) -> ColumnRegistry:
        """Generate pair <source_id, avatar_id> for using a FromObject in a ColumnRegistry"""
        if isinstance(from_obj, AvatarFromObject):
            return self._source_db_columns
        if isinstance(from_obj, SubqueryFromObject):
            trans_from_query = translated_queries_by_id[from_obj.id]
            column_reg = ColumnRegistry()
            avatar_id = from_obj.query_id
            source_id = from_obj.query_id
            column_reg.register_avatar(avatar_id=avatar_id, source_id=source_id)
            for column in trans_from_query.column_list:
                column_reg.register_column(db_column=column, avatar_id=avatar_id, column_id=column.name)
            return column_reg
        raise TypeError(type(from_obj))

    def translate_multi_query(
        self,
        compiled_multi_query: CompiledMultiQueryBase,
        collect_errors: bool = False,
    ) -> TranslatedMultiQueryBase:
        def _translate_subquery(compiled_flat_subquery: CompiledQuery) -> None:
            from_ids = set(compiled_flat_subquery.joined_from.iter_ids())
            self._log_info(
                f"Translating flat query with id {compiled_flat_subquery.id}, " f"using from-objects {from_ids}"
            )
            joined_from_column_reg = ColumnRegistry()
            for from_obj in compiled_flat_subquery.joined_from.froms:
                from_column_reg = self._make_column_reg_for_from_obj(
                    from_obj=from_obj,
                    translated_queries_by_id=translated_queries_by_id,
                )
                joined_from_column_reg += from_column_reg

            flat_trans = self._get_flat_translator_for_level_type(
                level_type=compiled_flat_subquery.level_type,
                columns=joined_from_column_reg,
            )
            trans_flat_query = flat_trans.translate_flat_query(
                compiled_flat_query=compiled_flat_subquery,
                collect_errors=collect_errors,
            )
            translated_queries_by_id[trans_flat_query.id] = trans_flat_query
            stats_lists_by_level_type[compiled_flat_subquery.level_type].append(flat_trans.get_collected_stats())

        def _recursively_translate_query(compiled_flat_query: CompiledQuery) -> None:
            required_subquery_ids = compiled_flat_query.joined_from.iter_ids()
            for req_id in required_subquery_ids:
                if req_id in base_from_ids:
                    continue  # Nothing to do for base FROMs (source avatars)
                if req_id not in translated_queries_by_id:
                    # Subquery is not yet translated, so go ahead and do it
                    compiled_flat_subquery = compiled_multi_query.get_query_by_id(req_id)
                    _recursively_translate_query(compiled_flat_subquery)
                    assert req_id in translated_queries_by_id

            _translate_subquery(compiled_flat_query)

        base_from_ids = set(compiled_multi_query.get_base_froms().keys())
        translated_queries_by_id: dict[str, TranslatedFlatQuery] = {}
        stats_lists_by_level_type: dict[ExecutionLevel, list[TranslationStats]] = defaultdict(list)
        for compiled_top_query in compiled_multi_query.get_top_queries():
            _recursively_translate_query(compiled_top_query)

        if self._collect_stats:
            self._log_collected_stats(stats_lists_by_level_type)
            self._log_query_complexity_stats(compiled_multi_query)

        return TranslatedMultiQuery(queries=list(translated_queries_by_id.values()))

    def collect_errors(
        self, compiled_multi_query: CompiledMultiQueryBase, feature_errors: bool = True
    ) -> list[FormulaErrorCtx]:
        with FormulaErrorCollector() as collector:
            self.translate_multi_query(compiled_multi_query=compiled_multi_query, collect_errors=True)
        return collector.get_errors(feature_errors=feature_errors)

lib/dl_query_processing/dl_query_processing/translation/multi_level_translator.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from typing import (
    AbstractSet,
    Any,
    Iterable,
    List,
    NamedTuple,
    Optional,
    Sequence,
    Type,
    TypeVar,
)

import attr

from dl_constants.enums import UserDataType
from dl_core.components.ids import (
    AvatarId,
    FieldId,
)
from dl_core.db.elements import SchemaColumn
from dl_core.query.expression import (
    ExpressionCtx,
    JoinOnExpressionCtx,
    OrderByExpressionCtx,
)
from dl_core.utils import attrs_evolve_to_subclass
from dl_formula.core.datatype import (
    DataType,
    DataTypeParams,
)
from dl_query_processing.compilation.primitives import FromObject
from dl_query_processing.compilation.query_meta import QueryMetaInfo
from dl_query_processing.enums import ExecutionLevel


class DetailedType(NamedTuple):
    field_id: str
    data_type: UserDataType
    # TODO: native_type: Optional[GenericNativeType] = None
    formula_data_type: Optional[DataType] = None
    formula_data_type_params: Optional[DataTypeParams] = None


_META_TV = TypeVar("_META_TV", bound="TranslatedQueryMetaInfo")


@attr.s
class TranslatedQueryMetaInfo(QueryMetaInfo):
    detailed_types: Optional[list[Optional[DetailedType]]] = attr.ib(kw_only=True, factory=list)  # type: ignore

    @classmethod
    def from_comp_meta(
        cls: Type[_META_TV],
        comp_meta: QueryMetaInfo,
        detailed_types: Optional[list[Optional[DetailedType]]] = None,
    ) -> _META_TV:
        detailed_types = detailed_types or []
        assert detailed_types is not None
        return attrs_evolve_to_subclass(
            cls=cls,
            inst=comp_meta,
            detailed_types=detailed_types,
        )


@attr.s(auto_attribs=True, frozen=True)
class ExpressionCtxExt(ExpressionCtx):
    """
    A variant of ExpressionCtx that knows about dl_formula.
    Should be used for the SELECT part of the query,
    as it contains info for postprocessing the data,
    which is not needed in other clauses.
    """

    formula_data_type: Optional[DataType] = None
    formula_data_type_params: Optional[DataTypeParams] = None
    original_field_id: Optional[FieldId] = None


@attr.s(frozen=True)
class TranslatedJoinedFromObject:
    root_from_id: Optional[str] = attr.ib(kw_only=True, default=None)
    froms: Sequence[FromObject] = attr.ib(kw_only=True, default=())  # FIXME: switch to translated froms

    def iter_ids(self) -> Iterable[str]:
        return (from_obj.id for from_obj in self.froms)


@attr.s(frozen=True)
class TranslatedFlatQuery:
    id: AvatarId = attr.ib(kw_only=True)
    alias: str = attr.ib(kw_only=True)
    level_type: ExecutionLevel = attr.ib(kw_only=True)

    select: List[ExpressionCtxExt] = attr.ib(kw_only=True)
    where: List[ExpressionCtx] = attr.ib(kw_only=True)
    group_by: List[ExpressionCtx] = attr.ib(kw_only=True)
    having: List[ExpressionCtx] = attr.ib(kw_only=True)
    order_by: List[OrderByExpressionCtx] = attr.ib(kw_only=True)
    join_on: List[JoinOnExpressionCtx] = attr.ib(kw_only=True)
    joined_from: TranslatedJoinedFromObject = attr.ib(kw_only=True)
    limit: Optional[int] = attr.ib(kw_only=True)
    offset: Optional[int] = attr.ib(kw_only=True)
    column_list: list[SchemaColumn] = attr.ib(kw_only=True)
    meta: TranslatedQueryMetaInfo = attr.ib(kw_only=True, factory=TranslatedQueryMetaInfo)

    def is_empty(self) -> bool:
        return not self.select


_TRANS_MULTI_QUERY_TV = TypeVar("_TRANS_MULTI_QUERY_TV", bound="TranslatedMultiQueryBase")


@attr.s(frozen=True)
class TranslatedMultiQueryBase(abc.ABC):
    @abc.abstractmethod
    def iter_queries(self) -> Iterable[TranslatedFlatQuery]:
        raise NotImplementedError

    @abc.abstractmethod
    def get_top_queries(self) -> list[TranslatedFlatQuery]:
        raise NotImplementedError

    def get_single_top_query(self) -> TranslatedFlatQuery:
        queries = self.get_top_queries()
        assert len(queries) == 1
        return queries[0]

    @abc.abstractmethod
    def get_query_by_id(self, id: str) -> TranslatedFlatQuery:
        raise NotImplementedError

    @abc.abstractmethod
    def for_level_type(self: _TRANS_MULTI_QUERY_TV, level_type: ExecutionLevel) -> _TRANS_MULTI_QUERY_TV:
        raise NotImplementedError

    @abc.abstractmethod
    def get_base_froms(self) -> dict[str, FromObject]:
        raise NotImplementedError

    @abc.abstractmethod
    def is_empty(self) -> bool:
        raise NotImplementedError

    @abc.abstractmethod
    def get_base_root_from_ids(self) -> AbstractSet[str]:
        raise NotImplementedError

    def clone(self: _TRANS_MULTI_QUERY_TV, **updates: Any) -> _TRANS_MULTI_QUERY_TV:
        return attr.evolve(self, **updates)


@attr.s(frozen=True)
class TranslatedLevel:
    level_type: ExecutionLevel = attr.ib(kw_only=True)
    queries: List[TranslatedFlatQuery] = attr.ib(factory=list)


@attr.s(frozen=True)
class TranslatedMultiLevelQuery(TranslatedMultiQueryBase):  # TODO: Remove this implementation
    levels: List[TranslatedLevel] = attr.ib(kw_only=True)

    def level_count(self) -> int:
        return len(self.levels)

    @property
    def top_level(self) -> TranslatedLevel:
        return self.levels[-1]

    def get_required_avatar_ids(self) -> set[str]:
        if not self.levels:
            return set()
        return {from_id for query in self.levels[0].queries for from_id in query.joined_from.iter_ids()}

    def iter_queries(self) -> Iterable[TranslatedFlatQuery]:
        for level in self.levels:
            for query in level.queries:
                yield query

    def get_top_queries(self) -> list[TranslatedFlatQuery]:
        return self.top_level.queries

    def get_query_by_id(self, id: str) -> TranslatedFlatQuery:
        for level in self.levels:
            for query in level.queries:
                if query.id == id:
                    return query
        raise KeyError(id)

    def for_level_type(self, level_type: ExecutionLevel) -> TranslatedMultiLevelQuery:
        levels = [level for level in self.levels if level.level_type == level_type]
        return self.clone(levels=levels)

    def is_empty(self) -> bool:
        return not self.levels or not self.levels[0].queries

    def get_base_root_from_ids(self) -> AbstractSet[str]:
        if not self.levels or not self.levels[0].queries:
            return frozenset()
        return {
            query.joined_from.root_from_id
            for query in self.levels[0].queries
            if query.joined_from.root_from_id is not None
        }

    def get_base_froms(self) -> dict[str, FromObject]:
        if not self.levels:
            return {}

        result: dict[str, FromObject] = {}
        for query in self.levels[0].queries:
            for from_obj in query.joined_from.froms:
                if from_obj.id not in result:
                    result[from_obj.id] = from_obj
        return result


@attr.s(frozen=True)
class TranslatedMultiQuery(TranslatedMultiQueryBase):
    queries: list[TranslatedFlatQuery] = attr.ib(kw_only=True)

    # internal caches
    _queries_by_id: dict[str, TranslatedFlatQuery] = attr.ib(init=False)
    _unreferenced_ids: list[str] = attr.ib(init=False)

    @_queries_by_id.default
    def _make_queries_by_id(self) -> dict[str, TranslatedFlatQuery]:
        return {query.id: query for query in self.queries}

    @_unreferenced_ids.default
    def _make_unreferenced_ids(self) -> list[str]:
        referenced_ids = {from_obj.id for query in self.queries for from_obj in query.joined_from.froms}
        query_ids = {query.id for query in self.queries}
        return sorted(query_ids - referenced_ids)

    def iter_queries(self) -> Iterable[TranslatedFlatQuery]:
        return self.queries

    def get_top_queries(self) -> list[TranslatedFlatQuery]:
        return [self._queries_by_id[query_id] for query_id in self._unreferenced_ids]

    def get_query_by_id(self, id: str) -> TranslatedFlatQuery:
        return self._queries_by_id[id]

    def for_level_type(self, level_type: ExecutionLevel) -> TranslatedMultiQuery:
        queries = [query for query in self.queries if query.level_type == level_type]
        return self.clone(queries=queries)

    def get_base_froms(self) -> dict[str, FromObject]:
        result: dict[str, FromObject] = {}
        for query in self.queries:
            for from_obj in query.joined_from.froms:
                if from_obj.id not in self._queries_by_id and from_obj.id not in result:
                    result[from_obj.id] = from_obj
        return result

    def is_empty(self) -> bool:
        return not self.queries

    def get_base_root_from_ids(self) -> AbstractSet[str]:
        result: set[str] = set()
        for query in self.queries:
            root_from_id = query.joined_from.root_from_id
            if root_from_id is not None and root_from_id not in self._queries_by_id:
                result.add(root_from_id)
        return result

lib/dl_query_processing/dl_query_processing/translation/primitives.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime

import ciso8601  # type: ignore


def parse_datetime(value: str) -> datetime.datetime:
    """
    Parse an ISO8601 datetime value, e.g. from API parameters.

    See also: `dl_formula.utils.datetime.parse_dt_string`

    >>> parse_datetime('2020-01-02T03:04:05')
    datetime.datetime(2020, 1, 2, 3, 4, 5)
    >>> parse_datetime('2020-01-02T03:04:05.678')
    datetime.datetime(2020, 1, 2, 3, 4, 5, 678000)
    >>> parse_datetime('2020-01-02T03:04:05+12:34')
    datetime.datetime(2020, 1, 2, 3, 4, 5, tzinfo=datetime.timezone(datetime.timedelta(seconds=45240)))
    >>> try:
    ...     parse_datetime('2020-01-02T03:04:65')
    ... except ValueError:
    ...     pass
    >>> parse_datetime('2020-01-02')
    datetime.datetime(2020, 1, 2, 0, 0)
    """
    return ciso8601.parse_datetime(value)

lib/dl_query_processing/dl_query_processing/utils/datetime.py
datalens-tech/datalens-backend
from __future__ import annotations

from collections import defaultdict
from itertools import count

import dl_formula.core.nodes as formula_nodes


class NameGen:  # TODO: Rename
    """
    Minor helper for numbered field names.

    >>> ng = NameGen()
    >>> ng(None, 'fld1')
    'fld1_1'
    >>> ng(None, 'fld2')
    'fld2_1'
    >>> ng(None, 'fld1')
    'fld1_2'
    >>> NameGen()(None, 'fld2')
    'fld2_1'
    """

    sep = "_"

    def __init__(self) -> None:
        self.counters = defaultdict(lambda: 0)

    def __call__(self, node: formula_nodes.FormulaItem, name: str) -> str:
        self.counters[name] += 1
        return f"{name}{self.sep}{self.counters[name]}"


class MappedNameGen(NameGen):
    sep = ""

    def __init__(self, namemap) -> None:
        super().__init__()
        self.namemap = namemap

    def __call__(self, node, name) -> str:
        return super().__call__(node, self.namemap.get(name, name))


class PrefixedIdGen:
    """
    Simplistic non-random prefixed ID generator.
    """

    def __init__(self, prefix: str) -> None:
        self._prefix = prefix
        self._cnt = count()

    def get_id(self) -> str:
        return f"{self._prefix}_{next(self._cnt)}"

lib/dl_query_processing/dl_query_processing/utils/name_gen.py
datalens-tech/datalens-backend
from dl_query_processing.legend.block_legend import (
    AfterBlockPlacement,
    DimensionSpec,
    DimensionValueSpec,
    DispersedAfterBlockPlacement,
)
from dl_query_processing.merging.binary_merger import (
    AfterBinaryStreamMerger,
    DispersedAfterBinaryStreamMerger,
)
from dl_query_processing.merging.primitives import MergedQueryDataRow


def test_after_binary_stream_merger_no_dims():
    parent_stream = [
        MergedQueryDataRow(data=(12, 34), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=(56, 78), legend_item_ids=(1, 2)),
    ]
    child_stream = [
        MergedQueryDataRow(data=("as", "df"), legend_item_ids=(5, 6)),
        MergedQueryDataRow(data=("zx", "cv"), legend_item_ids=(5, 6)),
    ]
    binary_merger = AfterBinaryStreamMerger(placement=AfterBlockPlacement())
    output_stream = list(binary_merger.merge_two_streams(parent_stream=parent_stream, child_stream=child_stream))
    assert output_stream == [
        MergedQueryDataRow(data=(12, 34), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=(56, 78), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=("as", "df"), legend_item_ids=(5, 6)),
        MergedQueryDataRow(data=("zx", "cv"), legend_item_ids=(5, 6)),
    ]


def test_after_binary_stream_merger_with_dims():
    parent_stream = [
        MergedQueryDataRow(data=(12, 34), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=(56, 78), legend_item_ids=(6, 7)),
        MergedQueryDataRow(data=(56, 78), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=(56, 78), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=(45, 76), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=(390, 2), legend_item_ids=(1, 2)),
    ]
    child_stream = [
        MergedQueryDataRow(data=("as", "df"), legend_item_ids=(5, 6)),
        MergedQueryDataRow(data=("zx", "cv"), legend_item_ids=(5, 6)),
    ]
    binary_merger = AfterBinaryStreamMerger(
        placement=AfterBlockPlacement(
            dimension_values=[
                DimensionValueSpec(legend_item_id=1, value=56),
                DimensionValueSpec(legend_item_id=2, value=78),
            ],
        ),
    )
    output_stream = list(binary_merger.merge_two_streams(parent_stream=parent_stream, child_stream=child_stream))
    assert output_stream == [
        MergedQueryDataRow(data=(12, 34), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=(56, 78), legend_item_ids=(6, 7)),
        MergedQueryDataRow(data=(56, 78), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=(56, 78), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=("as", "df"), legend_item_ids=(5, 6)),
        MergedQueryDataRow(data=("zx", "cv"), legend_item_ids=(5, 6)),
        MergedQueryDataRow(data=(45, 76), legend_item_ids=(1, 2)),
        MergedQueryDataRow(data=(390, 2), legend_item_ids=(1, 2)),
    ]


def test_dispersed_after_binary_stream_merger():
    parent_stream = [
        MergedQueryDataRow(data=(10, "abc", 1, 100), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "abc", 2, 213), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "abc", 3, 489), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "def", 1, 178), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "def", 2, 231), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "def", 3, 876), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(23, "abc", 1, 178), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(23, "abc", 2, 231), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(23, "abc", 3, 876), legend_item_ids=(1, 2, 3, 4)),
    ]
    child_stream = [
        MergedQueryDataRow(data=(10, "abc", "-", 1000), legend_item_ids=(5, 6, 7, 8)),
        MergedQueryDataRow(data=(10, "def", "-", 2000), legend_item_ids=(5, 6, 7, 8)),
        MergedQueryDataRow(data=(23, "abc", "-", 3000), legend_item_ids=(5, 6, 7, 8)),
    ]
    binary_merger = DispersedAfterBinaryStreamMerger(
        placement=DispersedAfterBlockPlacement(
            parent_dimensions=[
                DimensionSpec(legend_item_id=1),
                DimensionSpec(legend_item_id=2),
            ],
            child_dimensions=[
                DimensionSpec(legend_item_id=5),
                DimensionSpec(legend_item_id=6),
            ],
        ),
    )
    output_stream = list(binary_merger.merge_two_streams(parent_stream=parent_stream, child_stream=child_stream))
    assert output_stream == [
        MergedQueryDataRow(data=(10, "abc", 1, 100), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "abc", 2, 213), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "abc", 3, 489), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "abc", "-", 1000), legend_item_ids=(5, 6, 7, 8)),
        MergedQueryDataRow(data=(10, "def", 1, 178), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "def", 2, 231), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "def", 3, 876), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "def", "-", 2000), legend_item_ids=(5, 6, 7, 8)),
        MergedQueryDataRow(data=(23, "abc", 1, 178), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(23, "abc", 2, 231), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(23, "abc", 3, 876), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(23, "abc", "-", 3000), legend_item_ids=(5, 6, 7, 8)),
    ]

    parent_stream = output_stream
    child_stream = [
        MergedQueryDataRow(data=(10, "-", "-", 40000), legend_item_ids=(9, 10, 11, 12)),
        MergedQueryDataRow(data=(23, "-", "-", 70000), legend_item_ids=(9, 10, 11, 12)),
    ]
    binary_merger = DispersedAfterBinaryStreamMerger(
        placement=DispersedAfterBlockPlacement(
            parent_dimensions=[
                DimensionSpec(legend_item_id=1),
            ],
            child_dimensions=[
                DimensionSpec(legend_item_id=9),
            ],
        ),
    )
    output_stream = list(binary_merger.merge_two_streams(parent_stream=parent_stream, child_stream=child_stream))
    assert output_stream == [
        MergedQueryDataRow(data=(10, "abc", 1, 100), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "abc", 2, 213), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "abc", 3, 489), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "abc", "-", 1000), legend_item_ids=(5, 6, 7, 8)),
        MergedQueryDataRow(data=(10, "def", 1, 178), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "def", 2, 231), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "def", 3, 876), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(10, "def", "-", 2000), legend_item_ids=(5, 6, 7, 8)),
        MergedQueryDataRow(data=(10, "-", "-", 40000), legend_item_ids=(9, 10, 11, 12)),
        MergedQueryDataRow(data=(23, "abc", 1, 178), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(23, "abc", 2, 231), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(23, "abc", 3, 876), legend_item_ids=(1, 2, 3, 4)),
        MergedQueryDataRow(data=(23, "abc", "-", 3000), legend_item_ids=(5, 6, 7, 8)),
        MergedQueryDataRow(data=(23, "-", "-", 70000), legend_item_ids=(9, 10, 11, 12)),
    ]

lib/dl_query_processing/dl_query_processing_tests/unit/merging/test_binary_merger.py
datalens-tech/datalens-backend
from dl_constants.enums import (
    JoinType,
    OrderDirection,
)
import dl_formula.core.fork_nodes as formula_fork_nodes
from dl_formula.core.index import NodeHierarchyIndex
from dl_formula.shortcuts import n
from dl_query_processing.compilation.primitives import (
    AvatarFromObject,
    CompiledFormulaInfo,
    CompiledJoinOnFormulaInfo,
    CompiledMultiQuery,
    CompiledOrderByFormulaInfo,
    CompiledQuery,
    FromColumn,
    JoinedFromObject,
    SubqueryFromObject,
)
from dl_query_processing.enums import (
    ExecutionLevel,
    QueryPart,
)
from dl_query_processing.multi_query.mutators.splitter_based import SplitterMultiQueryMutator
from dl_query_processing.multi_query.splitters.mask_based import (
    AddFormulaInfo,
    AliasedFormulaSplitMask,
    MultiQuerySplitter,
    QuerySplitMask,
    SubqueryType,
)
from dl_query_processing.multi_query.splitters.prefiltered import PrefilteredFieldMultiQuerySplitter
from dl_query_processing.utils.name_gen import PrefixedIdGen


class MyMultiQuerySplitter(MultiQuerySplitter):
    def __init__(self):
        self._done_once = False

    def get_split_masks(
        self,
        query: CompiledQuery,
        expr_id_gen: PrefixedIdGen,
        query_id_gen: PrefixedIdGen,
    ) -> list[QuerySplitMask]:
        """
        A splitter that splits the `m_1` field between the `SUM` and `FUNC_2` calls
        and joins the sub-queries along using the `dim_2`/`custom_alias_2` dimension.
        """

        if self._done_once:
            return []

        self._done_once = True
        return [
            QuerySplitMask(
                subquery_type=SubqueryType.default,
                subquery_id=query_id_gen.get_id(),
                formula_split_masks=(
                    AliasedFormulaSplitMask(
                        alias="custom_alias_1",
                        query_part=QueryPart.select,
                        formula_list_idx=0,
                        outer_node_idx=NodeHierarchyIndex(indices=(0, 0)),
                        inner_node_idx=NodeHierarchyIndex(indices=(0, 0)),
                    ),
                    AliasedFormulaSplitMask(
                        alias="custom_alias_1",
                        query_part=QueryPart.order_by,
                        formula_list_idx=0,
                        outer_node_idx=NodeHierarchyIndex(indices=(0, 0)),
                        inner_node_idx=NodeHierarchyIndex(indices=(0, 0)),
                    ),
                ),
                filter_indices=frozenset(),
                join_type=JoinType.inner,
                joining_node=formula_fork_nodes.QueryForkJoiningWithList.make(
                    condition_list=[
                        formula_fork_nodes.SelfEqualityJoinCondition.make(n.field("custom_alias_2")),
                    ],
                ),
                add_formulas=(
                    AddFormulaInfo(
                        alias="custom_alias_2",
                        expr=n.func.FUNC_3(n.field("asd")),
                        from_ids=frozenset({"ava_1"}),
                        is_group_by=True,
                    ),
                ),
            )
        ]

    def mutate_cropped_query(self, query: CompiledQuery) -> CompiledQuery:
        return query.clone(group_by=[])


def test_query_splitter():
    original_query = CompiledMultiQuery(
        queries=[
            CompiledQuery(
                id="qq",
                level_type=ExecutionLevel.source_db,
                select=[
                    CompiledFormulaInfo(
                        alias="m_1",
                        formula_obj=n.formula(n.func.FUNC_2(n.func.AVG(n.field("qwe")))),
                        original_field_id="field_1",
                        avatar_ids={"ava_1"},
                    ),
                    CompiledFormulaInfo(
                        alias="dim_2",
                        formula_obj=n.formula(n.func.FUNC_3(n.field("asd"))),
                        original_field_id="field_2",
                        avatar_ids={"ava_1"},
                    ),
                    CompiledFormulaInfo(
                        alias="m_3",
                        formula_obj=n.formula(n.func.SUM(n.func.FUNC_4(n.field("zxc")))),
                        original_field_id="field_3",
                        avatar_ids={"ava_1"},
                    ),
                ],
                group_by=[
                    CompiledFormulaInfo(
                        alias="dim_2",
                        formula_obj=n.formula(n.func.FUNC_3(n.field("asd"))),
                        original_field_id="field_2",
                        avatar_ids={"ava_1"},
                    ),
                ],
                order_by=[
                    CompiledOrderByFormulaInfo(
                        alias="m_1",
                        formula_obj=n.formula(n.func.FUNC_2(n.func.AVG(n.field("qwe")))),
                        original_field_id="field_1",
                        avatar_ids={"ava_1"},
                        direction=OrderDirection.desc,
                    ),
                ],
                joined_from=JoinedFromObject(
                    root_from_id="ava_1",
                    froms=[
                        AvatarFromObject(
                            id="ava_1",
                            avatar_id="ava_1",
                            source_id="src_1",
                            alias="ava_1",
                            columns=(
                                FromColumn(id="qwe", name="qwe"),
                                FromColumn(id="asd", name="asd"),
                                FromColumn(id="zxc", name="zxc"),
                            ),
                        ),
                    ],
                ),
            ),
        ],
    )
    mutator = SplitterMultiQueryMutator(splitters=[MyMultiQuerySplitter()])
    actual_result = mutator.mutate_multi_query(original_query)
    expected_result = CompiledMultiQuery(
        queries=[
            CompiledQuery(
                id="qq",
                level_type=ExecutionLevel.source_db,
                select=[
                    CompiledFormulaInfo(
                        formula_obj=n.formula(n.func.FUNC_2(n.field("custom_alias_1"))),
                        alias="m_1",
                        avatar_ids={"q_0"},
                        original_field_id="field_1",
                    ),
                    CompiledFormulaInfo(
                        formula_obj=n.formula(n.field("e_1")),
                        alias="dim_2",
                        avatar_ids={"q_1"},
                        original_field_id="field_2",
                    ),
                    CompiledFormulaInfo(
                        formula_obj=n.formula(n.field("e_0")),
                        alias="m_3",
                        avatar_ids={"q_1"},
                        original_field_id="field_3",
                    ),
                ],
                order_by=[
                    CompiledOrderByFormulaInfo(
                        formula_obj=n.formula(n.func.FUNC_2(n.field("custom_alias_1"))),
                        alias="m_1",
                        avatar_ids={"q_0"},
                        original_field_id="field_1",
                        direction=OrderDirection.desc,
                    ),
                ],
                join_on=[
                    CompiledJoinOnFormulaInfo(
                        alias=None,
                        formula_obj=n.formula(n.binary("_dneq", n.field("e_1"), n.field("custom_alias_2"))),
                        avatar_ids={"q_0", "q_1"},
                        original_field_id=None,
                        left_id="q_1",
                        right_id="q_0",
                        join_type=JoinType.inner,
                    ),
                ],
                joined_from=JoinedFromObject(
                    root_from_id="q_1",
                    froms=[
                        SubqueryFromObject(
                            id="q_1",
                            alias="q_1",
                            query_id="q_1",
                            columns=(
                                FromColumn(id="e_0", name="e_0"),
                                FromColumn(id="e_1", name="e_1"),
                            ),
                        ),
                        SubqueryFromObject(
                            id="q_0",
                            alias="q_0",
                            query_id="q_0",
                            columns=(
                                FromColumn(id="custom_alias_1", name="custom_alias_1"),
                                FromColumn(id="custom_alias_2", name="custom_alias_2"),
                            ),
                        ),
                    ],
                ),
            ),
            CompiledQuery(
                id="q_1",
                level_type=ExecutionLevel.source_db,
                select=[
                    CompiledFormulaInfo(
                        alias="e_0",
                        formula_obj=n.formula(n.func.SUM(n.func.FUNC_4(n.field("zxc")))),
                        original_field_id="field_3",
                        avatar_ids={"ava_1"},
                    ),
                    CompiledFormulaInfo(
                        formula_obj=n.formula(n.func.FUNC_3(n.field("asd"))),
                        alias="e_1",
                        avatar_ids={"ava_1"},
                        original_field_id=None,
                    ),
                ],
                group_by=[
                    CompiledFormulaInfo(
                        formula_obj=n.formula(n.func.FUNC_3(n.field("asd"))),
                        alias="e_1",
                        avatar_ids={"ava_1"},
                        original_field_id=None,
                    ),
                ],
                joined_from=JoinedFromObject(
                    root_from_id="ava_1",
                    froms=[
                        AvatarFromObject(
                            id="ava_1",
                            avatar_id="ava_1",
                            source_id="src_1",
                            alias="ava_1",
                            columns=(
                                FromColumn(id="qwe", name="qwe"),
                                FromColumn(id="asd", name="asd"),
                                FromColumn(id="zxc", name="zxc"),
                            ),
                        ),
                    ],
                ),
            ),
            CompiledQuery(
                id="q_0",
                level_type=ExecutionLevel.source_db,
                select=[
                    CompiledFormulaInfo(
                        alias="custom_alias_1",
                        formula_obj=n.formula(n.func.AVG(n.field("qwe"))),
                        avatar_ids={"ava_1"},
                        original_field_id="field_1",
                    ),
                    CompiledFormulaInfo(
                        formula_obj=n.formula(n.func.FUNC_3(n.field("asd"))),
                        alias="custom_alias_2",
                        avatar_ids={"ava_1"},
                        original_field_id=None,
                    ),
                ],
                group_by=[
                    CompiledFormulaInfo(
                        formula_obj=n.formula(n.func.FUNC_3(n.field("asd"))),
                        alias="custom_alias_2",
                        avatar_ids={"ava_1"},
                        original_field_id=None,
                    ),
                ],
                joined_from=JoinedFromObject(
                    root_from_id="ava_1",
                    froms=[
                        AvatarFromObject(
                            id="ava_1",
                            avatar_id="ava_1",
                            source_id="src_1",
                            alias="ava_1",
                            columns=(
                                FromColumn(id="qwe", name="qwe"),
                                FromColumn(id="asd", name="asd"),
                                FromColumn(id="zxc", name="zxc"),
                            ),
                        ),
                    ],
                ),
            ),
        ],
    )
    assert expected_result == actual_result


def test_field_splitter():
    original_query = CompiledMultiQuery(
        queries=[
            CompiledQuery(
                id="qq",
                level_type=ExecutionLevel.source_db,
                select=[
                    CompiledFormulaInfo(
                        alias="f_1",
                        formula_obj=n.formula(n.func.FUNC_1(n.field("qwe"))),
                        original_field_id="field_1",
                        avatar_ids={"ava_1"},
                    ),
                    CompiledFormulaInfo(
                        alias="f_2",
                        formula_obj=n.formula(n.func.FUNC_2(n.field("asd"))),
                        original_field_id="field_2",
                        avatar_ids={"ava_1"},
                    ),
                    CompiledFormulaInfo(
                        alias="f_3",
                        formula_obj=n.formula(n.func.FUNC_3(n.field("zxc"))),
                        original_field_id="field_3",
                        avatar_ids={"ava_2"},
                    ),
                ],
                order_by=[
                    CompiledOrderByFormulaInfo(
                        alias="f_1",
                        formula_obj=n.formula(n.func.FUNC_1(n.field("qwe"))),
                        original_field_id="field_1",
                        avatar_ids={"ava_1"},
                        direction=OrderDirection.desc,
                    ),
                ],
                join_on=[
                    CompiledJoinOnFormulaInfo(
                        alias=None,
                        formula_obj=n.formula(n.binary("_dneq", n.field("asd"), n.field("zxc"))),
                        avatar_ids={"ava_1", "ava_2"},
                        original_field_id=None,
                        left_id="ava_1",
                        right_id="ava_2",
                        join_type=JoinType.inner,
                    ),
                ],
                joined_from=JoinedFromObject(
                    root_from_id="ava_1",
                    froms=[
                        AvatarFromObject(
                            id="ava_1",
                            avatar_id="ava_1",
                            source_id="src_1",
                            alias="ava_1",
                            columns=(
                                FromColumn(id="qwe", name="qwe"),
                                FromColumn(id="asd", name="asd"),
                            ),
                        ),
                        AvatarFromObject(
                            id="ava_2",
                            avatar_id="ava_2",
                            source_id="src_2",
                            alias="ava_2",
                            columns=(FromColumn(id="zxc", name="zxc"),),
                        ),
                    ],
                ),
            ),
        ],
    )
    mutator = SplitterMultiQueryMutator(
        splitters=[PrefilteredFieldMultiQuerySplitter(crop_to_level_type=ExecutionLevel.compeng)],
    )
    actual_result = mutator.mutate_multi_query(original_query)
    expected_result = CompiledMultiQuery(
        queries=[
            CompiledQuery(
                id="qq",
                level_type=ExecutionLevel.compeng,
                select=[
                    CompiledFormulaInfo(
                        alias="f_1",
                        formula_obj=n.formula(n.func.FUNC_1(n.field("e_0"))),
                        original_field_id="field_1",
                        avatar_ids={"q_0"},
                    ),
                    CompiledFormulaInfo(
                        alias="f_2",
                        formula_obj=n.formula(n.func.FUNC_2(n.field("e_1"))),
                        original_field_id="field_2",
                        avatar_ids={"q_0"},
                    ),
                    CompiledFormulaInfo(
                        alias="f_3",
                        formula_obj=n.formula(n.func.FUNC_3(n.field("e_2"))),
                        original_field_id="field_3",
                        avatar_ids={"q_1"},
                    ),
                ],
                order_by=[
                    CompiledOrderByFormulaInfo(
                        alias="f_1",
                        formula_obj=n.formula(n.func.FUNC_1(n.field("e_0"))),
                        original_field_id="field_1",
                        avatar_ids={"q_0"},
                        direction=OrderDirection.desc,
                    ),
                ],
                join_on=[
                    CompiledJoinOnFormulaInfo(
                        alias=None,
                        formula_obj=n.formula(n.binary("_dneq", n.field("e_1"), n.field("e_2"))),
                        avatar_ids={"q_0", "q_1"},
                        original_field_id=None,
                        left_id="q_0",
                        right_id="q_1",
                        join_type=JoinType.inner,
                    ),
                ],
                joined_from=JoinedFromObject(
                    root_from_id="q_0",
                    froms=[
                        SubqueryFromObject(
                            id="q_0",
                            query_id="q_0",
                            alias="q_0",
                            columns=(
                                FromColumn(id="e_0", name="e_0"),
                                FromColumn(id="e_1", name="e_1"),
                            ),
                        ),
                        SubqueryFromObject(
                            id="q_1",
                            query_id="q_1",
                            alias="q_1",
                            columns=(FromColumn(id="e_2", name="e_2"),),
                        ),
                    ],
                ),
            ),
            CompiledQuery(
                id="q_0",
                level_type=ExecutionLevel.source_db,
                select=[
                    CompiledFormulaInfo(
                        alias="e_0",
                        formula_obj=n.formula(n.field("qwe")),
                        original_field_id=None,
                        avatar_ids={"ava_1"},
                    ),
                    CompiledFormulaInfo(
                        alias="e_1",
                        formula_obj=n.formula(n.field("asd")),
                        original_field_id=None,
                        avatar_ids={"ava_1"},
                    ),
                ],
                joined_from=JoinedFromObject(
                    root_from_id="ava_1",
                    froms=[
                        AvatarFromObject(
                            id="ava_1",
                            avatar_id="ava_1",
                            source_id="src_1",
                            alias="ava_1",
                            columns=(
                                FromColumn(id="qwe", name="qwe"),
                                FromColumn(id="asd", name="asd"),
                            ),
                        ),
                    ],
                ),
            ),
            CompiledQuery(
                id="q_1",
                level_type=ExecutionLevel.source_db,
                select=[
                    CompiledFormulaInfo(
                        alias="e_2",
                        formula_obj=n.formula(n.field("zxc")),
                        original_field_id=None,
                        avatar_ids={"ava_2"},
                    ),
                ],
                joined_from=JoinedFromObject(
                    root_from_id="ava_2",
                    froms=[
                        AvatarFromObject(
                            id="ava_2",
                            avatar_id="ava_2",
                            source_id="src_2",
                            alias="ava_2",
                            columns=(FromColumn(id="zxc", name="zxc"),),
                        ),
                    ],
                ),
            ),
        ],
    )
    assert expected_result == actual_result

lib/dl_query_processing/dl_query_processing_tests/unit/multi_query/test_splitter.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_query_processing.postprocessing.postprocessors.markup import (
    MarkupProcessing,
    MarkupProcessingDC,
)


def _test_markup_postprocessing_i(cls, check_nodes=True):
    MPP = cls()
    formulated = MPP.n_concat(
        'url: "', MPP.n_url("col1", "col2"), '"; "', MPP.n_i(MPP.n_b(MPP.n_url("col3", "col4"))), '"', ""
    )
    if MPP._dbg:
        print("formulated:", formulated)
    assert formulated
    if check_nodes:
        expected = ("c", 'url: "', ("a", "col1", "col2"), '"; "', ("i", ("b", ("a", "col3", "col4"))), '"', "")
        assert formulated == expected
    dumped = MPP.dump(formulated)
    if MPP._dbg:
        print("dumped:", repr(dumped))
    assert dumped == '(c "url: """ (a "col1" "col2") """; """ (i (b (a "col3" "col4"))) """" "")'
    parsed = MPP.parse(dumped)
    if MPP._dbg:
        print("parsed:    ", parsed)
        print("formulated:", formulated)
    assert parsed == formulated
    verbalized = MPP.verbalize(parsed)
    if MPP._dbg:
        print("verbalized:", verbalized)
    expected = {
        "type": "concat",
        "children": [
            {"type": "text", "content": 'url: "'},
            {"type": "url", "url": "col1", "content": {"type": "text", "content": "col2"}},
            {"type": "text", "content": '"; "'},
            {
                "type": "italics",
                "content": {
                    "type": "bold",
                    "content": {"type": "url", "url": "col3", "content": {"type": "text", "content": "col4"}},
                },
            },
            {"type": "text", "content": '"'},
            {"type": "text", "content": ""},
        ],
    }
    assert verbalized == expected


def test_markup_postprocessing():
    _test_markup_postprocessing_i(MarkupProcessing, check_nodes=True)
    _test_markup_postprocessing_i(MarkupProcessingDC, check_nodes=False)

lib/dl_query_processing/dl_query_processing_tests/unit/postprocessing/test_markup.py
datalens-tech/datalens-backend
from typing import Sequence

from dl_query_processing.compilation.primitives import (
    FromColumn,
    FromObject,
    JoinedFromObject,
    SubqueryFromObject,
)


def joined_from_from_avatar_ids(
    from_ids: Sequence[str], base_avatars: dict[str, FromObject], cols_by_query: dict[str, Sequence[str]]
) -> JoinedFromObject:
    froms: list[FromObject] = []
    for from_id in from_ids:
        from_obj: FromObject
        if from_id in base_avatars:
            from_obj = base_avatars[from_id]
        else:
            from_obj = SubqueryFromObject(
                id=from_id,
                query_id=from_id,
                alias=from_id,
                columns=tuple(FromColumn(id=col_name, name=col_name) for col_name in cols_by_query[from_id]),
            )
        froms.append(from_obj)

    return JoinedFromObject(root_from_id=from_ids[0], froms=froms)

lib/dl_query_processing/dl_query_processing_tests/unit/utils.py
datalens-tech/datalens-backend
import datetime
from typing import Any

import sqlalchemy as sa
from sqlalchemy.engine import default


class BitrixCompiler(sa.sql.compiler.SQLCompiler):
    def visit_label(
        self,
        label,
        add_to_result_map=None,
        within_label_clause=False,
        within_columns_clause=False,
        render_label_as_label=None,
        **kw,
    ):
        # Labels are not supported at all
        return label.element._compiler_dispatch(
            self,
            within_columns_clause=True,
            within_label_clause=False,
            **kw,
        )

    def render_literal_value(self, value, type_):
        if isinstance(value, (datetime.datetime, datetime.date)):
            return value.isoformat()
        return super().render_literal_value(value, type_)


class BitrixIdentifierPreparer(sa.sql.compiler.IdentifierPreparer):
    def __init__(self, *args, **kwargs):
        quote = "`"
        kwargs = {
            **kwargs,
            "initial_quote": quote,
            "escape_quote": quote,
        }
        super().__init__(*args, **kwargs)


class BitrixDialect(default.DefaultDialect):
    name = "bitrix"
    poolclass = sa.pool.NullPool

    statement_compiler = BitrixCompiler
    ddl_compiler = sa.sql.compiler.DDLCompiler
    type_compiler = sa.sql.compiler.GenericTypeCompiler
    preparer = BitrixIdentifierPreparer

    supports_alter = False
    supports_comments = False
    inline_comments = False

    supports_views = False
    supports_sequences = False
    sequences_optional = False
    preexecute_autoincrement_sequences = False
    postfetch_lastrowid = True
    implicit_returning = False

    supports_right_nested_joins = False
    cte_follows_insert = False

    supports_native_enum = False
    supports_native_boolean = False  # Uncertain
    non_native_boolean_check_constraint = False

    supports_simple_order_by_label = False

    tuple_in_values = False

    supports_native_decimal = False

    supports_unicode_statements = True
    supports_unicode_binds = True
    returns_unicode_strings = True
    description_encoding = None

    supports_sane_rowcount = True
    supports_sane_multi_rowcount = True

    supports_default_values = False
    supports_empty_insert = False
    supports_multivalues_insert = False

    supports_is_distinct_from = True  # Uncertain

    supports_server_side_cursors = False

    supports_for_update_of = False

    # Clear out:
    ischema_names: dict[str, Any] = {}

    @classmethod
    def dbapi(cls):
        return None  # Not Applicable... if possible.

    def _check_unicode_returns(self, connection, additional_tests=None):
        return True

    def _check_unicode_description(self, connection):
        return True

    def do_rollback(self, dbapi_connection):
        pass

    def get_columns(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_foreign_keys(self, connection, table_name, schema=None, **kw):
        return []

    def get_indexes(self, connection, table_name, schema=None, **kw):
        return []

    def get_pk_constraint(self, connection, table_name, schema=None, **kw):
        return []

    def _get_default_schema_name(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def _get_server_version_info(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def denormalize_name(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def do_begin_twophase(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def do_commit_twophase(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def do_prepare_twophase(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def do_recover_twophase(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def do_rollback_twophase(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_check_constraints(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_isolation_level(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_primary_keys(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_table_comment(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_table_names(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_temp_table_names(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_temp_view_names(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_unique_constraints(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_view_definition(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def get_view_names(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def has_sequence(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def has_table(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def normalize_name(self, *args, **kwargs):
        raise Exception("Not Implemented")

    def set_isolation_level(self, *args, **kwargs):
        raise Exception("Not Implemented")


dialect = BitrixDialect


def register_dialect(
    name="bi_bitrix",
    module="dl_sqlalchemy_bitrix.base",
    cls="BitrixDialect",
):
    """
    Make sure the dialect is registered
    (normally should happen automagically because of the `entry_point`)
    """
    return sa.dialects.registry.register(name, module, cls)

lib/dl_sqlalchemy_bitrix/dl_sqlalchemy_bitrix/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest
import sqlalchemy as sa

from dl_sqlalchemy_bitrix.base import register_dialect


def pytest_configure(config):  # noqa
    register_dialect()


@pytest.fixture(scope="session")
def engine_url():
    return "bi_bitrix://"


@pytest.fixture(scope="session")
def sa_engine(engine_url):
    return sa.create_engine(engine_url)

lib/dl_sqlalchemy_bitrix/dl_sqlalchemy_bitrix_tests/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa

from dl_sqlalchemy_bitrix.base import BitrixCompiler


def test_engine(sa_engine):
    assert sa_engine
    assert sa_engine.dialect


def test_select_compiler(sa_engine):
    query = sa.select([sa.column("COLUMN_ONE"), sa.column("COLUMN_TWO")]).where(
        sa.column("DATE_CREATE") == "2000-01-01",
    )
    query_compiled = str(BitrixCompiler(sa_engine.dialect, query))

    assert not query_compiled.startswith("SELECT *")

lib/dl_sqlalchemy_bitrix/dl_sqlalchemy_bitrix_tests/unit/test_dialect.py
datalens-tech/datalens-backend
from __future__ import annotations

from .base import (
    BICHYTDialect,
    CHYTTableExpression,
    CHYTTablesConcat,
    CHYTTablesRange,
    CHYTTableSubselect,
)


__all__ = (
    "CHYTTableExpression",
    "CHYTTablesConcat",
    "CHYTTablesRange",
    "CHYTTableSubselect",
    "BICHYTDialect",
)

lib/dl_sqlalchemy_chyt/dl_sqlalchemy_chyt/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import sqlalchemy as sa

import dl_sqlalchemy_chyt.types as chyt_types
from dl_sqlalchemy_clickhouse.base import BIClickHouseDialect as UPSTREAM


quote_identifier = UPSTREAM().identifier_preparer.quote


# Type converters
ischema_names = {
    **UPSTREAM.ischema_names,
    "YtBoolean": chyt_types.YtBoolean,
}


class CHYTTableExpression(sa.sql.elements.TextClause):
    @staticmethod
    def _quote_identifier(value):
        return quote_identifier(value)

    @staticmethod
    def _quote_value(value):
        from clickhouse_sqlalchemy.drivers.http.escaper import Escaper

        return Escaper().escape_item(value)
        # return QUOTER.quote(value)

    def __init__(self, text, alias=None, **kwargs):
        if alias:
            text = "{} as {}".format(text, self._quote_identifier(alias))
        text = text.replace(":", r"\:")  # see also: dl_core.utils.sa_plain_text
        super(CHYTTableExpression, self).__init__(text, **kwargs)


class CHYTTablesConcat(CHYTTableExpression):
    def __init__(self, *tables, **kwargs):
        text = "concatYtTables({})".format(", ".join(self._quote_identifier(table) for table in tables if table))
        self._tables = tables
        super(CHYTTablesConcat, self).__init__(text, **kwargs)


class CHYTTablesRange(CHYTTableExpression):
    def __init__(self, directory, start=None, end=None, **kwargs):
        if end is not None:
            args = [directory, start or "", end]
        elif start is not None:
            args = [directory, start]
        else:
            args = [directory]

        text = "concatYtTablesRange({})".format(", ".join(self._quote_value(arg) for arg in args))
        self._directory = directory
        self._start = start
        self._end = end
        super(CHYTTablesRange, self).__init__(text, **kwargs)


class CHYTTableSubselect(CHYTTableExpression):
    def __init__(self, subsql, **kwargs):
        text = "(\n{}\n)".format(subsql)
        self._subsql = subsql
        super(CHYTTableSubselect, self).__init__(text, **kwargs)


class BICHYTTypeCompiler(UPSTREAM.type_compiler):
    def visit_ytboolean(self, type_, **kw):
        return "YtBoolean"


class BICHYTDialect(UPSTREAM):
    # Add YT-only types
    ischema_names = ischema_names
    type_compiler = BICHYTTypeCompiler

    def get_columns(self, connection, table_name, schema=None, **kw):
        if not isinstance(table_name, (CHYTTableExpression, sa.sql.elements.TextClause)):
            table_name = self.identifier_preparer.quote_identifier(table_name)

        query = "DESCRIBE TABLE {}".format(table_name)
        rows = self._execute(connection, query)

        return [self._get_column_info(row.name, row.type) for row in rows]


def register_dialect():
    sa.dialects.registry.register("bi_chyt", "dl_sqlalchemy_chyt.base", "BICHYTDialect")

lib/dl_sqlalchemy_chyt/dl_sqlalchemy_chyt/base.py
datalens-tech/datalens-backend
from clickhouse_sqlalchemy.types import Int


class YtBoolean(Int):
    __visit_name__ = "ytboolean"

lib/dl_sqlalchemy_chyt/dl_sqlalchemy_chyt/types.py
datalens-tech/datalens-backend
from __future__ import annotations

from clickhouse_sqlalchemy.drivers.http.base import ClickHouseDialect_http as UPSTREAM
import sqlalchemy as sa
from sqlalchemy import exc
from sqlalchemy.sql import elements
from sqlalchemy.sql.compiler import (
    COLLECT_CARTESIAN_PRODUCTS,
    WARN_LINTING,
    FromLinter,
)

from dl_sqlalchemy_common.base import CompilerPrettyMixin


class BIClickHouseTypeCompiler(UPSTREAM.type_compiler):
    """..."""

    # # Not currently viable.
    # def visit_datetime(self, type_, **kw):
    #     # for casts? TODO: dl_sqlalchemy_clickhouse_tests.
    #     return "DateTime('UTC')"


class BIClickHouseCompilerBasic(UPSTREAM.statement_compiler):
    """..."""


class BIClickHouseCompiler(BIClickHouseCompilerBasic, UPSTREAM.statement_compiler, CompilerPrettyMixin):
    """
    A bunch of copypaste from the `ClickHouseCompiler` but building query text
    through `self._pretty`
    """

    def visit_join(self, join, asfrom=False, **kwargs):
        # need to make a variable to prevent leaks in some debuggers
        join_type = getattr(join, "type", None)
        if join_type is None:
            if join.isouter:
                join_type = "LEFT OUTER"
            else:
                join_type = "INNER"
        elif join_type is not None:
            join_type = join_type.upper()
            if join.isouter and "INNER" in join_type:
                raise exc.CompileError("can't compile join with specified " "INNER type and isouter=True")
            # isouter=False by default, disable that checking
            # elif not join.isouter and 'OUTER' in join.type:
            #     raise exc.CompileError(
            #         "can't compile join with specified "
            #         "OUTER type and isouter=False"
            #     )
        if join.full and "FULL" not in join_type:
            join_type = "FULL " + join_type

        if getattr(join, "strictness", None):
            join_type = join.strictness.upper() + " " + join_type

        if getattr(join, "distribution", None):
            join_type = join.distribution.upper() + " " + join_type

        onclause = join.onclause
        return self._pretty.join_sql_join(
            left=join.left._compiler_dispatch(self, asfrom=asfrom, **kwargs),
            join_type=(join_type + " JOIN").upper() if join_type else "",
            right=join.right._compiler_dispatch(self, asfrom=asfrom, **kwargs),
            onclause=onclause._compiler_dispatch(self, **kwargs),
            ontext="USING" if isinstance(onclause, elements.Tuple) else "ON",
        )

    def _compose_select_body(
        self,
        text,
        select,
        compile_state,
        inner_columns,
        froms,
        byfrom,
        toplevel,
        kwargs,
    ):
        text_pieces = [self._pretty.join_block(text.rstrip(), self._pretty.join_group(inner_columns))]
        del text  # should not be used any further

        # TODO: confirm that we need this part of code in overridden method
        if self.linting & COLLECT_CARTESIAN_PRODUCTS:
            from_linter = FromLinter({}, set())
            warn_linting = self.linting & WARN_LINTING
            if toplevel:
                self.from_linter = from_linter
        else:
            from_linter = None  # type: ignore
            warn_linting = False

        if froms:
            if select._hints:
                from_sql = self._pretty.join_group(
                    [
                        f._compiler_dispatch(self, asfrom=True, fromhints=byfrom, from_linter=from_linter, **kwargs)
                        for f in froms
                    ]
                )
            else:
                from_sql = self._pretty.join_group(
                    [f._compiler_dispatch(self, asfrom=True, from_linter=from_linter, **kwargs) for f in froms]
                )
            text_pieces.append(self._pretty.join_block("FROM", from_sql, block_indent=False))
        else:
            text_pieces.append(self.default_from())

        if getattr(select, "_array_join", None) is not None:
            text_pieces.append(select._array_join._compiler_dispatch(self, **kwargs))

        sample_clause = getattr(select, "_sample_clause", None)

        if sample_clause is not None:
            text_pieces.append(self.sample_clause(select, **kwargs))

        final_clause = getattr(select, "_final_clause", None)

        if final_clause is not None:
            text_pieces.append(self.final_clause())

        if select._where_criteria:
            t = self._generate_delimited_and_list(select._where_criteria, from_linter=from_linter, **kwargs)
            if t:
                text_pieces.append(self._pretty.join_block("WHERE", t))

        if warn_linting:
            from_linter.warn()

        if select._group_by_clauses:
            text_pieces.append(self.group_by_clause(select, **kwargs))

        if select._having_criteria:
            t = self._generate_delimited_and_list(select._having_criteria, **kwargs)
            if t:
                text_pieces.append(self._pretty.join_block("HAVING", t))

        if select._order_by_clauses:
            t = self.order_by_clause(select, **kwargs)
            text_pieces.append(self._pretty.indent(t))

        if select._has_row_limiting_clause:
            text_pieces.append(self._row_limit_clause(select, **kwargs))

        if select._for_update_arg is not None:
            text_pieces.append(self.for_update_clause(select, **kwargs))

        return self._pretty.join(text_pieces)

    # Copypaste of `CompilerPrettyMixin`, just to confirm that the same
    # processing is okay for the overridden methods.
    def group_by_clause(self, select, **kw):
        sup = super().group_by_clause(select, **kw)
        return self._pretty.postprocess_block("GROUP BY", sup)

    def limit_clause(self, select, **kw):
        sup = super().limit_clause(select, **kw)
        return self._pretty.postprocess_block("LIMIT", sup)


class BIClickHouseDialectBasic(UPSTREAM):
    type_compiler = BIClickHouseTypeCompiler
    statement_compiler = BIClickHouseCompilerBasic


class BIClickHouseDialect(BIClickHouseDialectBasic):
    statement_compiler = BIClickHouseCompiler


def register_dialect():
    sa.dialects.registry.register("bi_clickhouse_basic", "dl_sqlalchemy_clickhouse.base", "BIClickHouseDialectBasic")
    sa.dialects.registry.register("bi_clickhouse", "dl_sqlalchemy_clickhouse.base", "BIClickHouseDialect")

lib/dl_sqlalchemy_clickhouse/dl_sqlalchemy_clickhouse/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest  # type: ignore
import sqlalchemy as sa

from dl_sqlalchemy_clickhouse.base import register_dialect


def pytest_configure(config):  # noqa
    register_dialect()


@pytest.fixture(scope="session")
def engine_url():
    return "bi_clickhouse://"


@pytest.fixture(scope="session")
def sa_engine(engine_url):
    return sa.create_engine(engine_url)

lib/dl_sqlalchemy_clickhouse/dl_sqlalchemy_clickhouse_tests/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations


def test_engine(sa_engine):
    assert sa_engine
    assert sa_engine.dialect

lib/dl_sqlalchemy_clickhouse/dl_sqlalchemy_clickhouse_tests/unit/test_dialect.py
datalens-tech/datalens-backend
from __future__ import annotations

import itertools
from typing import (
    Optional,
    Sequence,
)

import sqlalchemy as sa


class Formatter:
    indenter: str = "  "
    separator: str = "\n"
    inline_separator: str = " "
    preferred_width: int = 120

    def join(self, pieces: Sequence[str], cleanup: bool = True, separator: Optional[str] = None) -> str:
        if separator is None:
            separator = self.separator
        if cleanup:
            pieces = [item.rstrip() for item in pieces]  # type: ignore
            pieces = [item for item in pieces if item]  # type: ignore
        return separator.join(pieces)

    def indent(self, text: str) -> str:
        return self.indenter + text.replace("\n", "\n" + self.indenter)

    def join_ext(self, pieces: Sequence[str], extra_separator: str) -> str:
        """
        Join either into "piece1, piece2" or "piece1,\npiece2"
        (^ example for `extra_separator=","`)
        """
        pieces = list(pieces)
        pieces_len = sum(len(piece) for piece in pieces)
        # Note that this isn't necessarily correct as the further indentation
        # isn't taken into account.
        inline_len = pieces_len + (len(extra_separator) + len(self.inline_separator)) * (
            len(pieces) - 1 if pieces else 0
        )
        if inline_len <= self.preferred_width:
            return self.join(pieces, separator=(extra_separator + self.inline_separator))
        return self.join(pieces, separator=(extra_separator + self.separator))

    def join_group(self, pieces: Sequence[str]) -> str:
        return self.join_ext(pieces, extra_separator=",")

    def join_block(self, name: str, contents: str, block_indent: bool = True) -> str:
        inline_pieces = [name, self.inline_separator, contents]
        if (
            "\n" not in name
            and "\n" not in contents
            and sum(len(item) for item in inline_pieces) <= self.preferred_width
        ):
            return "".join(inline_pieces)
        return self.join((name, self.indent(contents) if block_indent else contents))

    def postprocess_block(self, sql_block_name: str, sql_text: str):
        sql_text_clean = sql_text.strip()
        if sql_text_clean.startswith(sql_block_name):
            sql_text_clean = sql_text_clean.removeprefix(sql_block_name).lstrip()
            return self.join_block(sql_block_name, sql_text_clean)
        return sql_text

    def join_sql_join(
        self,
        left: str,
        join_type: str,
        right: str,
        onclause: str,
        ontext: str = "ON",
    ):
        return self.join_ext(
            (
                left,
                join_type.strip(),
                right,
                self.join_block(ontext.strip(), onclause),
            ),
            extra_separator="",
        )

    def reparenthesize(self, sql_text: str) -> str:
        if not (sql_text.startswith("(") and sql_text.endswith(")")):
            return sql_text
        if "\n" not in sql_text:
            return sql_text
        return self.join(
            (
                "(",
                self.indent(sql_text.removeprefix("(").removesuffix(")")),
                ")",
            )
        )


def _cls_keys(cls: type) -> set[str]:
    return {key for key in cls.__dict__ if not key.startswith("__")}


class CompilerPrettyMixin(sa.sql.compiler.SQLCompiler):  # type: ignore  # TODO: fix
    """
    Compiler with SQL prettification.

    Primarily, adds newlines in some places.

    Additionally might add support for `right join` and other
    common-but-not-omnipresent features.

    The recommended MRO for this mixin is

    1. local customizations
    2. dialect-specific compiler
    3. CompilerPrettyMixin
    4. SQLCompiler

    This avoids overriding dialect-specific code with the prettification code,
    which would lead to potentially very unclear bugs.

    To cacth the inverse problem, i.e. dialect-specific code overriding the
    prettification code, on subclass creation, the `__init_subclass__` hook
    checks that all prettification methods were either not overridden,
    or were last overridden in a local customization (any `bi_…` python package).
    """

    _pretty: Formatter = Formatter()

    @staticmethod
    def _is_local_overrides_class(supcls):
        if supcls.__module__.startswith("dl_") or supcls.__module__.startswith("bi_"):
            return True
        if supcls.__module__.startswith("tests") or supcls.__module__.endswith("_tests"):
            return True
        if supcls.__module__.startswith("__tests__"):
            return True
        return False

    @classmethod
    def __init_subclass__(cls) -> None:
        prettifier_attrs = _cls_keys(CompilerPrettyMixin)
        full_mro = cls.__mro__
        prettifier_idx = full_mro.index(CompilerPrettyMixin)
        mro = full_mro[:prettifier_idx]
        non_bi_idx = next(
            (idx for idx, supcls in enumerate(mro) if not cls._is_local_overrides_class(supcls)),
            -1,
        )
        bi_clss = mro[:non_bi_idx]
        base_clss = mro[non_bi_idx:]
        bi_attrs = {key for cls in bi_clss for key in _cls_keys(cls)}
        base_attrs = {key for cls in base_clss for key in _cls_keys(cls)}
        uncovered_attrs = (base_attrs - bi_attrs) & prettifier_attrs
        if uncovered_attrs:
            raise Exception(
                "Dialect-specific compiler overrides methods that aren't re-overridden in `bi_...`",
                dict(uncovered_attrs=uncovered_attrs, bi_clss=bi_clss, base_clss=base_clss),
            )
        super().__init_subclass__()

    def visit_join(self, join, asfrom=False, from_linter=None, **kwargs):
        if from_linter:
            from_linter.edges.update(itertools.product(join.left._from_objects, join.right._from_objects))

        if join.full:
            join_type = " FULL OUTER JOIN "
        elif join.isouter:
            join_type = " LEFT OUTER JOIN "
        # NOTE: Right join is not in the base dialect; might not be supported
        # in some databases.
        elif getattr(join, "type", None) == "right":
            join_type = " RIGHT OUTER JOIN "
        else:
            join_type = " JOIN "
        return self._pretty.join_sql_join(
            left=join.left._compiler_dispatch(self, asfrom=True, **kwargs),
            join_type=join_type,
            right=join.right._compiler_dispatch(self, asfrom=True, **kwargs),
            onclause=join.onclause._compiler_dispatch(self, **kwargs),
        )

    def visit_select(
        self,
        select_stmt,
        asfrom=False,
        insert_into=False,
        fromhints=None,
        compound_index=None,
        select_wraps_for=None,
        lateral=False,
        from_linter=None,
        **kwargs,
    ):
        assert select_wraps_for is None, (
            "SQLAlchemy 1.4 requires use of "
            "the translate_select_structure hook for structural "
            "translations of SELECT objects"
        )

        # initial setup of SELECT.  the compile_state_factory may now
        # be creating a totally different SELECT from the one that was
        # passed in.  for ORM use this will convert from an ORM-state
        # SELECT to a regular "Core" SELECT.  other composed operations
        # such as computation of joins will be performed.
        compile_state = select_stmt._compile_state_factory(select_stmt, self, **kwargs)
        select_stmt = compile_state.statement

        toplevel = not self.stack

        if toplevel and not self.compile_state:
            self.compile_state = compile_state

        is_embedded_select = compound_index is not None or insert_into

        # translate step for Oracle, SQL Server which often need to
        # restructure the SELECT to allow for LIMIT/OFFSET and possibly
        # other conditions
        if self.translate_select_structure:
            new_select_stmt = self.translate_select_structure(select_stmt, asfrom=asfrom, **kwargs)

            # if SELECT was restructured, maintain a link to the originals
            # and assemble a new compile state
            if new_select_stmt is not select_stmt:
                compile_state_wraps_for = compile_state
                select_wraps_for = select_stmt
                select_stmt = new_select_stmt

                compile_state = select_stmt._compile_state_factory(select_stmt, self, **kwargs)
                select_stmt = compile_state.statement

        entry = self._default_stack_entry if toplevel else self.stack[-1]

        populate_result_map = need_column_expressions = (
            toplevel
            or entry.get("need_result_map_for_compound", False)
            or entry.get("need_result_map_for_nested", False)
        )

        # indicates there is a CompoundSelect in play and we are not the
        # first select
        if compound_index:
            populate_result_map = False

        # this was first proposed as part of #3372; however, it is not
        # reached in current tests and could possibly be an assertion
        # instead.
        if not populate_result_map and "add_to_result_map" in kwargs:
            del kwargs["add_to_result_map"]

        froms = self._setup_select_stack(select_stmt, compile_state, entry, asfrom, lateral, compound_index)

        column_clause_args = kwargs.copy()
        column_clause_args.update({"within_label_clause": False, "within_columns_clause": False})

        text_pieces = ["SELECT"]  # we're off to a good start !

        if select_stmt._hints:
            hint_text, byfrom = self._setup_select_hints(select_stmt)
            if hint_text:
                text_pieces.append(hint_text)
        else:
            byfrom = None

        if select_stmt._independent_ctes:
            for cte in select_stmt._independent_ctes:
                cte._compiler_dispatch(self, **kwargs)

        if select_stmt._prefixes:
            text_pieces.append(self._generate_prefixes(select_stmt, select_stmt._prefixes, **kwargs))

        subtext = self.get_select_precolumns(select_stmt, **kwargs)
        if subtext:
            text_pieces.append(subtext)
        # the actual list of columns to print in the SELECT column list.
        inner_columns = [
            c
            for c in [
                self._label_select_column(
                    select_stmt,
                    column,
                    populate_result_map,
                    asfrom,
                    column_clause_args,
                    name=name,
                    proxy_name=proxy_name,
                    fallback_label_name=fallback_label_name,
                    column_is_repeated=repeated,
                    need_column_expressions=need_column_expressions,
                )
                for (
                    name,
                    proxy_name,
                    fallback_label_name,
                    column,
                    repeated,
                ) in compile_state.columns_plus_names
            ]
            if c is not None
        ]

        if populate_result_map and select_wraps_for is not None:
            # if this select was generated from translate_select,
            # rewrite the targeted columns in the result map

            translate = dict(
                zip(
                    [
                        name
                        for (
                            key,
                            proxy_name,
                            fallback_label_name,
                            name,
                            repeated,
                        ) in compile_state.columns_plus_names
                    ],
                    [
                        name
                        for (
                            key,
                            proxy_name,
                            fallback_label_name,
                            name,
                            repeated,
                        ) in compile_state_wraps_for.columns_plus_names
                    ],
                )
            )

            self._result_columns = [
                (key, name, tuple(translate.get(o, o) for o in obj), type_)
                for key, name, obj, type_ in self._result_columns
            ]

        body_text = self._pretty.join(text_pieces)
        body_text = self._compose_select_body(
            body_text + " ",
            select_stmt,
            compile_state,
            inner_columns,
            froms,
            byfrom,
            toplevel,
            kwargs,
        )
        text_pieces = [body_text]

        if select_stmt._statement_hints:
            per_dialect = [
                ht for (dialect_name, ht) in select_stmt._statement_hints if dialect_name in ("*", self.dialect.name)
            ]
            if per_dialect:
                text_pieces.append(self.get_statement_hint_text(per_dialect))

        if self.ctes:
            # In compound query, CTEs are shared at the compound level
            if not is_embedded_select:
                nesting_level = len(self.stack) if not toplevel else None
                text_pieces = [self._render_cte_clause(nesting_level=nesting_level)] + text_pieces

        if select_stmt._suffixes:
            text_pieces.append(self._generate_prefixes(select_stmt, select_stmt._suffixes, **kwargs))

        self.stack.pop(-1)

        return self._pretty.join(text_pieces)

    def _compose_select_body(
        self,
        text,
        select,
        compile_state,
        inner_columns,
        froms,
        byfrom,
        toplevel,
        kwargs,
    ):
        text_pieces = [self._pretty.join_block(text.rstrip(), self._pretty.join_group(inner_columns))]
        del text  # should not be used any further

        if self.linting & sa.sql.compiler.COLLECT_CARTESIAN_PRODUCTS:
            from_linter = sa.sql.compiler.FromLinter({}, set())
            warn_linting = self.linting & sa.sql.compiler.WARN_LINTING
            if toplevel:
                self.from_linter = from_linter
        else:
            from_linter = None  # type: ignore
            warn_linting = False

        if froms:
            if select._hints:
                from_sql = self._pretty.join_group(
                    [
                        f._compiler_dispatch(self, asfrom=True, fromhints=byfrom, from_linter=from_linter, **kwargs)
                        for f in froms
                    ]
                )
            else:
                from_sql = self._pretty.join_group(
                    [f._compiler_dispatch(self, asfrom=True, from_linter=from_linter, **kwargs) for f in froms]
                )
            text_pieces.append(self._pretty.join_block("FROM", from_sql, block_indent=False))
        else:
            text_pieces.append(self.default_from())

        if select._where_criteria:
            t = self._generate_delimited_and_list(select._where_criteria, from_linter=from_linter, **kwargs)
            if t:
                text_pieces.append(self._pretty.join_block("WHERE", t))

        if warn_linting:
            from_linter.warn()

        if select._group_by_clauses:
            text_pieces.append(self.group_by_clause(select, **kwargs))

        if select._having_criteria:
            t = self._generate_delimited_and_list(select._having_criteria, **kwargs)
            if t:
                text_pieces.append(self._pretty.join_block("HAVING", t))

        if select._order_by_clauses:
            t = self.order_by_clause(select, **kwargs)
            text_pieces.append(t)

        if select._has_row_limiting_clause:
            text_pieces.append(self._row_limit_clause(select, **kwargs))

        if select._for_update_arg is not None:
            text_pieces.append(self.for_update_clause(select, **kwargs))

        return self._pretty.join(text_pieces)

    def visit_alias(self, *args, **kwargs):
        sup = super().visit_alias(*args, **kwargs)
        sup = self._pretty.reparenthesize(sup)
        return sup

    def group_by_clause(self, select, **kw):
        sup = super().group_by_clause(select, **kw)
        return self._pretty.postprocess_block("GROUP BY", sup)

    def order_by_clause(self, select, **kw):
        sup = super().order_by_clause(select, **kw)
        return self._pretty.postprocess_block("ORDER BY", sup)

    def limit_clause(self, select, **kw):
        sup = super().limit_clause(select, **kw)
        return self._pretty.postprocess_block("LIMIT", sup)

    def visit_clauselist(self, clauselist, *args, **kwargs):
        if clauselist.operator is sa.sql.operators.and_:
            assert sa.sql.compiler.OPERATORS[clauselist.operator].strip() == "AND"
            pieces = [el._compiler_dispatch(self, *args, **kwargs) for el in clauselist.clauses]
            separator = "AND "
            pieces = pieces[:1] + [separator + piece for piece in pieces[1:]]
            return self._pretty.join_ext(pieces, extra_separator="")

        sup = super().visit_clauselist(clauselist, *args, **kwargs)
        return sup

    def _generate_delimited_list(self, elements, separator, **kw):
        if separator in (", ", " "):
            children = [el._compiler_dispatch(self, **kw) for el in elements]
            return self._pretty.join_ext(children, extra_separator=separator.strip())
        return super()._generate_delimited_list(elements, separator, **kw)

lib/dl_sqlalchemy_common/dl_sqlalchemy_common/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest
import sqlalchemy as sa
from sqlalchemy import (
    Column,
    Integer,
    MetaData,
    String,
    Table,
    desc,
)
from sqlalchemy.sql import (
    column,
    table,
)


# Reference:
# https://github.com/sqlalchemy/sqlalchemy/blob/rel_1_4/test/sql/test_compiler.py#L107


table1 = table(
    "mytable",
    column("myid", Integer),
    column("name", String),
    column("description", String),
)

table2 = table("myothertable", column("otherid", Integer), column("othername", String))

table3 = table("thirdtable", column("userid", Integer), column("otherstuff", String))

metadata = MetaData()

# table with a schema
table4 = Table(
    "remotetable",
    metadata,
    Column("rem_id", Integer, primary_key=True),
    Column("datatype_id", Integer),
    Column("value", String(20)),
    schema="remote_owner",
)

# table with a "multipart" schema
table5 = Table(
    "remotetable",
    metadata,
    Column("rem_id", Integer, primary_key=True),
    Column("datatype_id", Integer),
    Column("value", String(20)),
    schema="dbo.remote_owner",
)


t1a = table1.alias("t1a")
t2a = table2.alias("t2a")
t2b = table2.alias("t2b")


query1 = (
    t1a.join(t2a, t1a.c.myid == t2a.c.otherid, full=True)
    .outerjoin(t2b, t1a.c.myid == t2b.c.otherid, full=True)
    .select()
    # Note: currently this doesn't get wrapped.
    .where(t1a.c.myid > -1)
    .where(t1a.c.myid < 9000)
    # ... and this does.
    .where(sa.and_((t1a.c.name != letter * 8) for letter in "abcde"))
    # Postgresql does `limit all`, other dialects raise: `.offset(123)`
)
sq1 = query1.subquery().alias("q_1")

t4a = table4.alias("t4a")
t5a = table5.alias("t5a")
t5b = table5.alias("t5b")
query2 = (
    sq1.join(t4a, sq1.c.myid == t4a.c.rem_id)
    .join(t5a, t4a.c.datatype_id == t5a.c.datatype_id)
    .join(t5b, sq1.c.name == t5b.c.value)
    .select()
    .where(sq1.c.name != "wxyz")
    .order_by(
        desc(sq1.c.myid),
        t4a.c.datatype_id,
        sq1.c.name,
        desc(sq1.c.description),
        desc(sq1.c.otherid),
        desc(sq1.c.othername),
        t4a.c.datatype_id,
        t4a.c.value,
        desc(t5a.c.rem_id),
        desc(t5a.c.datatype_id),
        desc(t5a.c.value),
        desc(t5b.c.rem_id),
        desc(t5b.c.datatype_id),
        desc(t5b.c.value),
    )
    .limit(int("12345" * 10))
    .offset(int("54321" * 10))
)


def get_dialects():
    try:
        from dl_sqlalchemy_chyt.base import BICHYTDialect

        yield BICHYTDialect()
    except ImportError:
        pass
    try:
        from dl_sqlalchemy_clickhouse.base import BIClickHouseDialect

        yield BIClickHouseDialect()
    except ImportError:
        pass
    try:
        from dl_sqlalchemy_mssql.base import BIMSSQLDialect

        yield BIMSSQLDialect()
    except ImportError:
        pass
    try:
        from dl_sqlalchemy_oracle.base import BIOracleDialect

        yield BIOracleDialect()
    except ImportError:
        pass
    try:
        from dl_sqlalchemy_postgres import BIPGDialect

        yield BIPGDialect()
    except ImportError:
        pass


DIALECTS = tuple(get_dialects())


@pytest.fixture(
    scope="session",
    params=DIALECTS,
    ids=[dialect.__class__.__name__ for dialect in DIALECTS],
)
def dialect(request):
    return request.param


@pytest.fixture(
    scope="session",
    params=[False, True],
    ids=["literal_binds=False", "literal_binds=True"],
)
def literal_binds(request):
    return request.param


@pytest.fixture(
    scope="session",
    params=[query1, query2],
    ids=["query1", "query2"],
)
def query(request):
    return request.param


def test_compile(dialect, query, literal_binds):
    sql_text = str(
        query.compile(
            dialect=dialect,
            compile_kwargs=dict(literal_binds=literal_binds),
        )
    )
    assert sql_text
    max_width = max(len(row) for row in sql_text.splitlines())
    assert max_width < 9000
    print()
    print(f"{max_width=!r}")
    print("Query:\n    " + sql_text.replace("\n", "\n    ") + "\n")

lib/dl_sqlalchemy_common/dl_sqlalchemy_common_tests/unit/test_prettify.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_sqlalchemy_common.base import CompilerPrettyMixin


def test_inherit():
    """Ensure mid-point inheritance works"""
    from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2 as UPSTREAM

    class CustomPGCompiler(UPSTREAM.statement_compiler):
        """Local customizations"""

    class ActualCustomPGCompiler(CustomPGCompiler, CompilerPrettyMixin):
        """Local customization on top of pretty-formatting"""

        def limit_clause(self, *args, **kwargs):
            """
            This method has to be overridden because otherwise
            `UPSTREAM.statement_compiler.limit_clause` overrides the
            `CompilerPrettyMixin.limit_clause`, and there's a check to catch
            that.
            """
            return super().limit_clause(*args, **kwargs)

    assert ActualCustomPGCompiler

lib/dl_sqlalchemy_common/dl_sqlalchemy_common_tests/unit/test_stuff.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_sqlalchemy_metrica_api import base


base.dialect = dialect = base.MetrikaApiDialect


__all__ = ("dialect",)

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from json.decoder import JSONDecodeError
import logging
from typing import (
    List,
    Optional,
)

import requests
from requests.exceptions import RequestException

import dl_sqlalchemy_metrica_api as package
from dl_sqlalchemy_metrica_api.exceptions import (  # noqa
    ConnectionClosedException,
    CursorClosedException,
    DatabaseError,
    DataError,
    Error,
    IntegrityError,
    InterfaceError,
    InternalError,
    MetrikaApiAccessDeniedException,
    MetrikaApiException,
    MetrikaApiObjectNotFoundException,
    MetrikaHttpApiException,
    NotSupportedError,
    OperationalError,
    ProgrammingError,
    Warning,
)


LOGGER = logging.getLogger(__name__)


METRIKA_API_HOST = "https://api-metrika.yandex.net"
APPMETRICA_API_HOST = "https://api.appmetrica.yandex.ru"


def _get_retriable_requests_session():
    session = requests.Session()
    for schema in ["http://", "https://"]:
        session.mount(
            schema,
            # noinspection PyUnresolvedReferences
            requests.adapters.HTTPAdapter(
                max_retries=requests.packages.urllib3.util.Retry(
                    total=5,
                    backoff_factor=0.5,
                    status_forcelist=[500, 501, 502, 504, 521],
                    redirect=10,
                    method_whitelist=frozenset(["HEAD", "TRACE", "GET", "PUT", "OPTIONS", "DELETE", "POST"]),
                ),
            ),
        )
    # TODO: allow to customize UA
    ua = "{}, {}".format(requests.utils.default_user_agent(), package.__name__)
    session.headers.update({"User-Agent": ua})
    return session


def _parse_metrika_error(response):
    msg = "Unknown error"
    try:
        resp_data = response.json()
        msg = resp_data["message"]
    except Exception:
        LOGGER.exception("Unable to fetch error message.")
    return msg


class MetrikaApiClient(object):
    """
    Simple HTTP client for Metrika API
    https://tech.yandex.ru/metrika/doc/api2/api_v1/intro-docpage/
    """

    host = METRIKA_API_HOST
    default_timeout = 60

    def __init__(self, oauth_token: str, host: Optional[str] = None, default_timeout=-1, **kwargs):
        if host is not None:
            self.host = host
        if default_timeout != -1:
            self.default_timeout = default_timeout
        self.oauth_token = oauth_token
        self._session = _get_retriable_requests_session()
        self._session.headers.update({"Authorization": "OAuth {}".format(oauth_token)})

    @property
    def _is_appmetrica(self):
        return self.host == APPMETRICA_API_HOST

    def _request(self, method: str, uri: str, timeout: int = -1, _raw_resp: bool = False, **kwargs):
        if timeout == -1:
            timeout = self.default_timeout
        full_url = "/".join(map(lambda s: s.strip("/"), (self.host, uri)))

        LOGGER.info(
            "Requesting Metrika API: method: %s, url: %s, params:(%s), json:(%s)",
            method,
            full_url,
            kwargs.get("params", {}),
            kwargs.get("json", {}),
        )

        response = None
        try:
            response = self._session.request(
                method,
                full_url,
                timeout=timeout,
                allow_redirects=False,
                **kwargs,
            )
            LOGGER.info(
                "Got %s from Metrika API (%s %s), content length: %s",
                response.status_code,
                method,
                uri,
                response.headers.get("Content-Length"),
            )
            if response.status_code >= 400:
                LOGGER.error(
                    "Metrika API error on %s %s (%s): %s",
                    method,
                    uri,
                    kwargs.get("json", {}),
                    response.text,
                )
            # TODO: wrap 429 and maybe retry
            if _raw_resp:
                return response

            response.raise_for_status()
        except RequestException as ex:
            msg = _parse_metrika_error(response)
            if response.status_code == 403:
                raise MetrikaApiAccessDeniedException(msg, orig_exc=ex) from ex
            elif response.status_code == 404:
                raise MetrikaApiObjectNotFoundException(msg, orig_exc=ex) from ex
            else:
                raise MetrikaHttpApiException(msg, orig_exc=ex) from ex

        try:
            parsed_resp = response.json()
        except JSONDecodeError as ex:
            raise MetrikaHttpApiException("Unable to parse response.", orig_exc=ex)
        return parsed_resp

    def get(self, uri, **kwargs):
        return self._request("GET", uri, **kwargs)

    def post(self, uri, **kwargs):
        return self._request("POST", uri, **kwargs)

    def _parse_data_resp(self, resp, result_columns=None, req_metrics=None):
        """
        https://tech.yandex.ru/metrika/doc/api2/api_v1/data-docpage/
        """
        LOGGER.info(
            "Received data response: total_rows: %s, sample_share: %s", resp["total_rows"], resp["sample_share"]
        )

        rows = []
        try:
            q_metrics = resp["query"]["metrics"]
            q_dims = resp["query"]["dimensions"]
            if req_metrics is not None and len(req_metrics) == 1:
                req_metrics = req_metrics[0].split(",")
                if len(q_metrics) != len(req_metrics):
                    raise MetrikaApiException("Unexpected response metrics count.")
                if q_metrics != req_metrics:
                    LOGGER.info(
                        "Response query metrics not matching requested metrics: %s, %s.",
                        q_metrics,
                        req_metrics,
                    )
                    q_metrics = req_metrics

            rc_dict = {col["name"]: col for col in result_columns}
            dims_src_keys = [rc_dict.get(dim, {}).get("src_key") or "name" for dim in q_dims]

            if not result_columns:
                result_columns = [dict(name=col_name, label=None) for col_name in (*q_dims, *q_metrics)]

            for slice in resp["data"]:
                # TODO: date dimensions values better retrieve from 'id' key instead of 'name'?
                row_map = {q_dims[i]: dim_item[dims_src_keys[i]] for i, dim_item in enumerate(slice["dimensions"])}
                row_map.update({q_metrics[i]: metr_val for i, metr_val in enumerate(slice["metrics"])})
                rows.append(
                    tuple(
                        col["cast_processor"](row_map[col["name"]]) if "cast_processor" in col else row_map[col["name"]]
                        for col in result_columns
                    )
                )
        except (KeyError, ValueError) as ex:
            raise MetrikaApiException(orig_exc=ex)

        return dict(
            fields=result_columns,
            data=rows,
        )

    def get_table_data(self, params, result_columns=None, **kwargs):
        resp = self.get("/stat/v1/data", params=params, **kwargs)
        return self._parse_data_resp(
            resp,
            result_columns=result_columns,
            req_metrics=params.get("metrics"),
        )

    def get_available_counters(self, **kwargs) -> List[dict]:
        obj_name = "applications" if self._is_appmetrica else "counters"
        uri = "/management/v1/{}".format(obj_name)
        resp = self.get(uri, **kwargs)
        return [dict(id=c_info["id"], name=c_info["name"]) for c_info in resp[obj_name]]

    def get_counter_info(self, counter_id):
        """
        https://tech.yandex.ru/metrika/doc/api2/management/counters/counter-docpage/
        """
        obj_name = "application" if self._is_appmetrica else "counter"
        uri = "/management/v1/{obj_name}/{counter_id}".format(obj_name=obj_name, counter_id=counter_id)
        resp = self.get(uri)
        return resp[obj_name]

    def get_counter_creation_date(self, counter_id) -> datetime.date:
        counter_info = self.get_counter_info(counter_id)
        try:
            date_str = counter_info.get("create_time", counter_info.get("create_date")).split("T")[0]
            creation_date = datetime.datetime.strptime(date_str, "%Y-%m-%d").date()
        except (ValueError, KeyError) as ex:
            raise MetrikaApiException(orig_exc=ex)
        return creation_date

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_client.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import Enum

from dl_sqlalchemy_metrica_api.api_info.appmetrica.audience import audience_fields
from dl_sqlalchemy_metrica_api.api_info.appmetrica.audience_socdem import audience_socdem_fields
from dl_sqlalchemy_metrica_api.api_info.appmetrica.client_events import client_events_fields
from dl_sqlalchemy_metrica_api.api_info.appmetrica.crash_events import crash_events_fields
from dl_sqlalchemy_metrica_api.api_info.appmetrica.installs import installs_fields
from dl_sqlalchemy_metrica_api.api_info.appmetrica.push_events import push_events_fields


class AppMetricaFieldsNamespaces(Enum):
    installs = "installs"
    push_events = "push_events"
    client_events = "client_events"
    audience = "audience"
    audience_socdem = "audience_socdem"
    crash_events = "crash_events"


metrica_fields_namespaces = AppMetricaFieldsNamespaces


fields_by_namespace = {
    AppMetricaFieldsNamespaces.installs: installs_fields,
    AppMetricaFieldsNamespaces.push_events: push_events_fields,
    AppMetricaFieldsNamespaces.client_events: client_events_fields,
    AppMetricaFieldsNamespaces.audience: audience_fields,
    AppMetricaFieldsNamespaces.audience_socdem: audience_socdem_fields,
    AppMetricaFieldsNamespaces.crash_events: crash_events_fields,
}


fields_by_name = {f["name"]: f for f in sum(fields_by_namespace.values(), [])}


metrics_by_namespace = {
    cs: list(filter(lambda f: not f["is_dim"], fields)) for cs, fields in fields_by_namespace.items()
}


namespace_by_prefix = {
    "ym:ts": AppMetricaFieldsNamespaces.installs,
    "ym:pc": AppMetricaFieldsNamespaces.push_events,
    "ym:ce": AppMetricaFieldsNamespaces.client_events,
    "ym:u": AppMetricaFieldsNamespaces.audience,
    "ym:cr2": AppMetricaFieldsNamespaces.crash_events,
}


def get_namespace_by_name(name):
    name_parts = name.split(":", 2)
    if len(name_parts) != 3:
        raise ValueError('Unexpected name format: "%s".' % name)
    prefix = ":".join(name_parts[:2])
    if prefix not in namespace_by_prefix:
        raise ValueError('Unknown field prefix: "%s"' % prefix)
    return namespace_by_prefix[prefix]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/appmetrica/__init__.py
datalens-tech/datalens-backend
audience_fields = [
    {
        "description": "",
        "is_dim": False,
        "name": "ym:u:sessions",
        "title": "Количество сессий",
        "type": "integer",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:u:newUsers",
        "title": "Новые пользователи",
        "type": "integer",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:u:activeUsers",
        "title": "Пользователи",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:isRooted",
        "title": "Root-статус",
        "type": "string",
    },
    {
        "description": "Идентификатор приложения",
        "is_dim": True,
        "name": "ym:u:apiKey",
        "title": "ID приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:UUID",
        "title": "UUID приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:clientKitVersion",
        "title": "Версия SDK клиента",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:operatingSystemVersionInfo",
        "title": "Версия платформы",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:appVersion",
        "title": "Версия приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:year",
        "title": "Год",
        "type": "integer",
    },
    {
        "description": "Название города, к которому принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:u:regionCityName",
        "title": "Город",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:date",
        "title": "Дата",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала года.",
        "is_dim": True,
        "name": "ym:u:startOfYear",
        "title": "Дата (год)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
        "name": "ym:u:startOfDekaminute",
        "title": "Дата (декаминута)",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:startOfQuarter",
        "title": "Дата (квартал)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала месяца.",
        "is_dim": True,
        "name": "ym:u:startOfMonth",
        "title": "Дата (месяц)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.",
        "is_dim": True,
        "name": "ym:u:startOfMinute",
        "title": "Дата (минута)",
        "type": "datetime",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
        "name": "ym:u:startOfHour",
        "title": "Дата (час)",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:dekaminute",
        "title": "Декаминута",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:dayOfMonth",
        "title": "День месяца",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:dayOfWeekName",
        "title": "День недели",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:appID",
        "title": "Идентификатор приложения (в магазине приложений)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:googleAID",
        "title": "Идентификатор рекламы Google",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:iosIFA",
        "title": "Идентификатор рекламы на iOS",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:deviceID",
        "title": "Идентификатор, полученный средствами системного API или выдаваемый Яндексом",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:osMajorVersionInfoName",
        "title": "Мажорная версия ОС",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:clientKitVersionDetails",
        "title": "Информация о версии SDK",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:month",
        "title": "Месяц",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:minute",
        "title": "Минута",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:mobileDeviceModel",
        "title": "Модель",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:mpcName",
        "title": "Название сотового оператора на основе кода сотового оператора (MPC)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:mccName",
        "title": "Название страны на основе мобильного кода страны оператора (MCC)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:buildNumber",
        "title": "Номер сборки приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:regionAreaName",
        "title": "Область",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:operatingSystemInfoName",
        "title": "Операционная система",
        "type": "string",
    },
    {
        "description": 'Возможные значения: "yes", "no", "undefined".',
        "is_dim": True,
        "name": "ym:u:limitAdTracking",
        "title": "Признак ограничения рекламного трекинга",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:mobileDeviceBranding",
        "title": "Производитель",
        "type": "string",
    },
    {
        "description": 'Подробная информация о разрешении экрана. Например, "2560x1440 px (xxxhdpi)".',
        "is_dim": True,
        "name": "ym:u:screenResolutionDetailed",
        "title": "Разрешение (подробно)",
        "type": "string",
    },
    {
        "description": "Название страны, к которой принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:u:regionCountryName",
        "title": "Страна",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:connectionType",
        "title": "Тип подключения: Cellular или Wi-Fi",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:networkType",
        "title": "Тип сети: 3G, EDGE и т.д.",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:deviceTypeName",
        "title": "Тип устройства",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:device",
        "title": "Устройство",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:hour",
        "title": "Час",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:hourMinute",
        "title": "Час и минута",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:locale",
        "title": "Язык интерфейса",
        "type": "string",
    },
]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/appmetrica/audience.py
datalens-tech/datalens-backend
audience_socdem_fields = [
    {"description": "", "is_dim": False, "name": "ym:u:sessions", "title": "Количество сессий", "type": "integer"},
    {"description": "", "is_dim": False, "name": "ym:u:newUsers", "title": "Новые пользователи", "type": "integer"},
    {"description": "", "is_dim": False, "name": "ym:u:activeUsers", "title": "Пользователи", "type": "integer"},
    {
        "description": 'Пол посетителя. Возможные значения: "мужской" или "женский".',
        "is_dim": True,
        "name": "ym:u:genderName",
        "title": "Пол",
        "type": "string",
    },
    {
        "description": "Возраст посетителя, интервал",
        "is_dim": True,
        "name": "ym:u:ageIntervalName",
        "title": "Возраст",
        "type": "string",
    },
    {"description": None, "is_dim": True, "name": "ym:u:isRooted", "title": "Root-статус", "type": "string"},
    {
        "description": "Идентификатор приложения",
        "is_dim": True,
        "name": "ym:u:apiKey",
        "title": "ID приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:clientKitVersion",
        "title": "Версия SDK клиента",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:operatingSystemVersionInfo",
        "title": "Версия платформы",
        "type": "string",
    },
    {"description": None, "is_dim": True, "name": "ym:u:appVersion", "title": "Версия приложения", "type": "string"},
    {"description": None, "is_dim": True, "name": "ym:u:year", "title": "Год", "type": "integer"},
    {
        "description": "Название города, к которому принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:u:regionCityName",
        "title": "Город",
        "type": "string",
    },
    {"description": None, "is_dim": True, "name": "ym:u:date", "title": "Дата", "type": "date"},
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала года.",
        "is_dim": True,
        "name": "ym:u:startOfYear",
        "title": "Дата (год)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
        "name": "ym:u:startOfDekaminute",
        "title": "Дата (декаминута)",
        "type": "datetime",
    },
    {"description": None, "is_dim": True, "name": "ym:u:startOfQuarter", "title": "Дата (квартал)", "type": "date"},
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала месяца.",
        "is_dim": True,
        "name": "ym:u:startOfMonth",
        "title": "Дата (месяц)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.",
        "is_dim": True,
        "name": "ym:u:startOfMinute",
        "title": "Дата (минута)",
        "type": "datetime",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
        "name": "ym:u:startOfHour",
        "title": "Дата (час)",
        "type": "datetime",
    },
    {"description": None, "is_dim": True, "name": "ym:u:dekaminute", "title": "Декаминута", "type": "string"},
    {"description": None, "is_dim": True, "name": "ym:u:dayOfMonth", "title": "День месяца", "type": "integer"},
    {"description": None, "is_dim": True, "name": "ym:u:dayOfWeekName", "title": "День недели", "type": "string"},
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:appID",
        "title": "Идентификатор приложения (в магазине приложений)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:osMajorVersionInfoName",
        "title": "Мажорная версия ОС",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:clientKitVersionDetails",
        "title": "Информация о версии SDK",
        "type": "string",
    },
    {"description": None, "is_dim": True, "name": "ym:u:month", "title": "Месяц", "type": "integer"},
    {"description": None, "is_dim": True, "name": "ym:u:minute", "title": "Минута", "type": "integer"},
    {"description": None, "is_dim": True, "name": "ym:u:mobileDeviceModel", "title": "Модель", "type": "string"},
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:mpcName",
        "title": "Название сотового оператора на основе кода сотового оператора (MPC)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:mccName",
        "title": "Название страны на основе мобильного кода страны оператора (MCC)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:buildNumber",
        "title": "Номер сборки приложения",
        "type": "string",
    },
    {"description": None, "is_dim": True, "name": "ym:u:regionAreaName", "title": "Область", "type": "string"},
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:operatingSystemInfoName",
        "title": "Операционная система",
        "type": "string",
    },
    {
        "description": 'Возможные значения: "yes", "no", "undefined".',
        "is_dim": True,
        "name": "ym:u:limitAdTracking",
        "title": "Признак ограничения рекламного трекинга",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:mobileDeviceBranding",
        "title": "Производитель",
        "type": "string",
    },
    {
        "description": 'Подробная информация о разрешении экрана. Например, "2560x1440 px (xxxhdpi)".',
        "is_dim": True,
        "name": "ym:u:screenResolutionDetailed",
        "title": "Разрешение (подробно)",
        "type": "string",
    },
    {
        "description": "Название страны, к которой принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:u:regionCountryName",
        "title": "Страна",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:connectionType",
        "title": "Тип подключения: Cellular или Wi-Fi",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:u:networkType",
        "title": "Тип сети: 3G, EDGE и т.д.",
        "type": "string",
    },
    {"description": None, "is_dim": True, "name": "ym:u:deviceTypeName", "title": "Тип устройства", "type": "string"},
    {"description": None, "is_dim": True, "name": "ym:u:hour", "title": "Час", "type": "string"},
    {"description": None, "is_dim": True, "name": "ym:u:hourMinute", "title": "Час и минута", "type": "string"},
    {"description": None, "is_dim": True, "name": "ym:u:locale", "title": "Язык интерфейса", "type": "string"},
]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/appmetrica/audience_socdem.py
datalens-tech/datalens-backend
client_events_fields = [
    {
        "description": "",
        "is_dim": False,
        "name": "ym:ce:users",
        "title": "Количество пользователей",
        "type": "integer",
    },
    {
        "description": "Количество клиентских событий",
        "is_dim": False,
        "name": "ym:ce:clientEvents",
        "title": "Количество событий",
        "type": "integer",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:ce:devices",
        "title": "Количество устройств",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:isRooted",
        "title": "Root-статус",
        "type": "string",
    },
    {
        "description": "Идентификатор приложения",
        "is_dim": True,
        "name": "ym:ce:apiKey",
        "title": "ID приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:UUID",
        "title": "UUID приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:clientKitVersion",
        "title": "Версия SDK клиента",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:operatingSystemVersionInfo",
        "title": "Версия платформы",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:appVersion",
        "title": "Версия приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:year",
        "title": "Год",
        "type": "integer",
    },
    {
        "description": "Название города, к которому принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:ce:regionCityName",
        "title": "Город",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:date",
        "title": "Дата",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала года.",
        "is_dim": True,
        "name": "ym:ce:startOfYear",
        "title": "Дата (год)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
        "name": "ym:ce:startOfDekaminute",
        "title": "Дата (декаминута)",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:startOfQuarter",
        "title": "Дата (квартал)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала месяца.",
        "is_dim": True,
        "name": "ym:ce:startOfMonth",
        "title": "Дата (месяц)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.",
        "is_dim": True,
        "name": "ym:ce:startOfMinute",
        "title": "Дата (минута)",
        "type": "datetime",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
        "name": "ym:ce:startOfHour",
        "title": "Дата (час)",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:dekaminute",
        "title": "Декаминута",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:dayOfMonth",
        "title": "День месяца",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:dayOfWeekName",
        "title": "День недели",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:appID",
        "title": "Идентификатор приложения (в магазине приложений)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:googleAID",
        "title": "Идентификатор рекламы Google",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:iosIFA",
        "title": "Идентификатор рекламы на iOS",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:deviceID",
        "title": "Идентификатор, полученный средствами системного API или выдаваемый Яндексом",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:osMajorVersionInfoName",
        "title": "Мажорная версия ОС",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:clientKitVersionDetails",
        "title": "Информация о версии SDK",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:month",
        "title": "Месяц",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:minute",
        "title": "Минута",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:mobileDeviceModel",
        "title": "Модель",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:eventLabel",
        "title": "Название события",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:mpcName",
        "title": "Название сотового оператора на основе кода сотового оператора (MPC)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:mccName",
        "title": "Название страны на основе мобильного кода страны оператора (MCC)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:buildNumber",
        "title": "Номер сборки приложения",
        "type": "string",
    },
    {
        "description": "Области, к которым принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:ce:regionAreaName",
        "title": "Область",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:operatingSystemInfoName",
        "title": "Операционная система",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:paramsValueNormalized",
        "title": "Параметры событий",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:paramsLevel1",
        "title": "Параметры событий: уровень 1",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:paramsLevel2",
        "title": "Параметры событий: уровень 2",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:paramsLevel3",
        "title": "Параметры событий: уровень 3",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:paramsLevel4",
        "title": "Параметры событий: уровень 4",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:paramsLevel5",
        "title": "Параметры событий: уровень 5",
        "type": "string",
    },
    {
        "description": 'Возможные значения: "yes", "no", "undefined".',
        "is_dim": True,
        "name": "ym:ce:limitAdTracking",
        "title": "Признак ограничения рекламного трекинга",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:mobileDeviceBranding",
        "title": "Производитель",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:profileOrigin",
        "title": "Профиль",
        "type": "string",
    },
    {
        "description": 'Подробная информация о разрешении экрана. Например, "2560x1440 px (xxxhdpi)".',
        "is_dim": True,
        "name": "ym:ce:screenResolutionDetailed",
        "title": "Разрешение (подробно)",
        "type": "string",
    },
    {
        "description": "Название страны, к которой принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:ce:regionCountryName",
        "title": "Страна",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:connectionType",
        "title": "Тип подключения: Cellular или Wi-Fi",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:networkType",
        "title": "Тип сети: 3G, EDGE и т.д.",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:deviceTypeName",
        "title": "Тип устройства",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:device",
        "title": "Устройство",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:hour",
        "title": "Час",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:hourMinute",
        "title": "Час и минута",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ce:locale",
        "title": "Язык интерфейса",
        "type": "string",
    },
]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/appmetrica/client_events.py
datalens-tech/datalens-backend
crash_events_fields = [
    {
        "name": "ym:cr2:crashes",
        "title": "Количество крэшей",
        "type": "integer",
        "description": None,
        "is_dim": False,
    },
    {
        "name": "ym:cr2:crashDevices",
        "title": "Количество устройств, на которых возникли аварийные остановки приложения",
        "type": "integer",
        "description": None,
        "is_dim": False,
    },
    {
        "name": "ym:cr2:crashesFreeDevicesPercentage",
        "title": "Процент устройств без крэшей",
        "type": "float",
        "description": None,
        "is_dim": False,
    },
    {
        "name": "ym:cr2:crashesDevicesPercentage",
        "title": "Процент устройств с крэшами",
        "type": "float",
        "description": None,
        "is_dim": False,
    },
    {
        "name": "ym:cr2:isRooted",
        "title": "Root-статус",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "description": "Идентификатор приложения",
        "is_dim": True,
        "name": "ym:cr2:apiKey",
        "title": "ID приложения",
        "type": "string",
    },
    {
        "name": "ym:cr2:UUID",
        "title": "UUID приложения",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:clientKitVersion",
        "title": "Версия SDK клиента",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:appVersion",
        "title": "Версия приложения",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:year",
        "title": "Год",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:regionCityName",
        "title": "Город",
        "type": "string",
        "description": "Название города, к которому принадлежат посетители сайта.",
        "is_dim": True,
    },
    {
        "name": "ym:cr2:date",
        "title": "Дата",
        "type": "date",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:startOfYear",
        "title": "Дата (год)",
        "type": "date",
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала года.",
        "is_dim": True,
    },
    {
        "name": "ym:cr2:startOfDekaminute",
        "title": "Дата (декаминута)",
        "type": "datetime",
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
    },
    {
        "name": "ym:cr2:startOfQuarter",
        "title": "Дата (квартал)",
        "type": "date",
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала квартала.",
        "is_dim": True,
    },
    {
        "name": "ym:cr2:startOfMonth",
        "title": "Дата (месяц)",
        "type": "date",
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала месяца.",
        "is_dim": True,
    },
    {
        "name": "ym:cr2:startOfMinute",
        "title": "Дата (минута)",
        "type": "datetime",
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.",
        "is_dim": True,
    },
    {
        "name": "ym:cr2:startOfHour",
        "title": "Дата (час)",
        "type": "datetime",
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
    },
    {
        "name": "ym:cr2:dekaminute",
        "title": "Декаминута",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:dayOfMonth",
        "title": "День месяца",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:dayOfWeekName",
        "title": "День недели",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:crashGroupName",
        "title": "Идентификатор группы крэшей",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:appID",
        "title": "Идентификатор приложения (в магазине приложений)",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:googleAID",
        "title": "Идентификатор рекламы Google",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:iosIFA",
        "title": "Идентификатор рекламы на iOS",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:deviceID",
        "title": "Идентификатор, полученный средствами системного API или выдаваемый Яндексом",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:crashMethodName",
        "title": "Имя метода исходного кода, в котором произошёл крэш",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:crashBinaryName",
        "title": "Имя образа, в котором произошёл крэш",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:crashFileName",
        "title": "Имя файла исходного кода, в котором произошёл крэш",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:clientKitVersionDetails",
        "title": "Информация о версии SDK",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:osMajorVersionInfoName",
        "title": "Мажорная версия ОС",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:month",
        "title": "Месяц",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:minute",
        "title": "Минута",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:mobileDeviceModel",
        "title": "Модель",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:mpcName",
        "title": "Название сотового оператора на основе кода сотового оператора (MPC)",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:mccName",
        "title": "Название страны на основе мобильного кода страны оператора (MCC)",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:buildNumber",
        "title": "Номер сборки приложения",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:crashSourceLine",
        "title": "Номер строки исходного кода, в котором произошёл крэш",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:regionAreaName",
        "title": "Область",
        "type": "string",
        "description": "Области, в которых находятся посетители сайта.",
        "is_dim": True,
    },
    {
        "name": "ym:cr2:operatingSystemInfoName",
        "title": "Операционная система",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:limitAdTracking",
        "title": "Признак ограничения рекламного трекинга",
        "type": "string",
        "description": 'Возможные значения: "yes", "no", "undefined".',
        "is_dim": True,
    },
    {
        "name": "ym:cr2:mobileDeviceBranding",
        "title": "Производитель",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:profileOrigin",
        "title": "Профиль",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:screenResolutionDetailed",
        "title": "Разрешение (подробно)",
        "type": "string",
        "description": 'Подробная информация о разрешении экрана. Например, "2560x1440 px (xxxhdpi)".',
        "is_dim": True,
    },
    {
        "name": "ym:cr2:regionCountryName",
        "title": "Страна",
        "type": "string",
        "description": "Название страны, к которой принадлежат посетители сайта.",
        "is_dim": True,
    },
    {
        "name": "ym:cr2:connectionType",
        "title": "Тип подключения: Cellular или Wi-Fi",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:networkType",
        "title": "Тип сети: 3G, EDGE и т.д.",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:deviceTypeName",
        "title": "Тип устройства",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:device",
        "title": "Устройство",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:hourMinute",
        "title": "Час и минута",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
    {
        "name": "ym:cr2:locale",
        "title": "Язык интерфейса",
        "type": "string",
        "description": None,
        "is_dim": True,
    },
]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/appmetrica/crash_events.py
datalens-tech/datalens-backend
installs_fields = [
    {
        "description": None,
        "is_dim": False,
        "name": "ym:ts:openDevices",
        "title": "Deeplinks",
        "type": "integer",
    },
    {
        "description": "Новые установки + реатрибутированные установки",
        "is_dim": False,
        "name": "ym:ts:advInstallDevices",
        "title": "Все установки",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:ts:userClicks",
        "title": "Клики",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:ts:userConversion",
        "title": "Конверсия из кликов в установки",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:ts:advNewInstallDevices",
        "title": "Новые установки",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:ts:advReattributedDevices",
        "title": "Реатрибутированные установки",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:campaignID",
        "title": "Tracking ID кампании",
        "type": "string",
    },
    {
        "description": "Идентификатор приложения",
        "is_dim": True,
        "name": "ym:ts:apiKey",
        "title": "ID приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:UUID",
        "title": "UUID приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:clientKitVersion",
        "title": "Версия SDK клиента",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:operatingSystemVersionInfo",
        "title": "Версия платформы",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:appVersion",
        "title": "Версия приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:year",
        "title": "Год",
        "type": "integer",
    },
    {
        "description": "Название города, к которому принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:ts:regionCityName",
        "title": "Город",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:date",
        "title": "Дата",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала года.",
        "is_dim": True,
        "name": "ym:ts:startOfYear",
        "title": "Дата (год)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
        "name": "ym:ts:startOfDekaminute",
        "title": "Дата (декаминута)",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:startOfQuarter",
        "title": "Дата (квартал)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала месяца.",
        "is_dim": True,
        "name": "ym:ts:startOfMonth",
        "title": "Дата (месяц)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.",
        "is_dim": True,
        "name": "ym:ts:startOfMinute",
        "title": "Дата (минута)",
        "type": "datetime",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
        "name": "ym:ts:startOfHour",
        "title": "Дата (час)",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:dekaminute",
        "title": "Декаминута",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:dayOfMonth",
        "title": "День месяца",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:dayOfWeekName",
        "title": "День недели",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:urlParamValue",
        "title": "Значения параметров трекинговой ссылки",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:googleAID",
        "title": "Идентификатор рекламы Google",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:iosIFA",
        "title": "Идентификатор рекламы на iOS",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:deviceID",
        "title": "Идентификатор, полученный средствами системного API или выдаваемый Яндексом",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:urlParamKey",
        "title": "Имена параметров трекинговой ссылки",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:campaignName",
        "title": "Кампания",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:osMajorVersionInfoName",
        "title": "Мажорная версия ОС",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:month",
        "title": "Месяц",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:minute",
        "title": "Минута",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:mpcName",
        "title": "Название сотового оператора на основе кода сотового оператора (MPC)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:mccName",
        "title": "Название страны на основе мобильного кода страны оператора (MCC)",
        "type": "string",
    },
    {
        "description": "Области, к которым принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:ts:regionAreaName",
        "title": "Область",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:operatingSystemInfoName",
        "title": "Операционная система",
        "type": "string",
    },
    {
        "description": 'Возможные значения: "yes", "no", "undefined".',
        "is_dim": True,
        "name": "ym:ts:limitAdTracking",
        "title": "Признак ограничения рекламного трекинга",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:publisherName",
        "title": "Рекламная сеть",
        "type": "string",
    },
    {
        "description": "Название страны, к которой принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:ts:regionCountryName",
        "title": "Страна",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:device",
        "title": "Устройство",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:hour",
        "title": "Час",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ts:hourMinute",
        "title": "Час и минута",
        "type": "string",
    },
]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/appmetrica/installs.py
datalens-tech/datalens-backend
push_events_fields = [
    {
        "description": "",
        "is_dim": False,
        "name": "ym:pc:devices",
        "title": "Количество устройств",
        "type": "integer",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:pc:receivedDevices",
        "title": "Количество устройств, которые получили пуш",
        "type": "integer",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:pc:sentDevices",
        "title": "Количество устройств, на которые был отправлен пуш",
        "type": "integer",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:pc:openedDevices",
        "title": "Количество устройств, открывших пуш",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:isRooted",
        "title": "Root-статус",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:tag",
        "title": "Tag push-api кампании",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:transfer",
        "title": "Transfer_id push-api кампании",
        "type": "string",
    },
    {
        "description": "Идентификатор приложения",
        "is_dim": True,
        "name": "ym:pc:apiKey",
        "title": "ID приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:UUID",
        "title": "UUID приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:clientKitVersion",
        "title": "Версия SDK клиента",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:operatingSystemVersionInfo",
        "title": "Версия ОС",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:appVersion",
        "title": "Версия приложения",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:hypothesisName",
        "title": "Гипотеза push-кампании",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:year",
        "title": "Год",
        "type": "integer",
    },
    {
        "description": "Название города, к которому принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:pc:regionCityName",
        "title": "Город",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:date",
        "title": "Дата",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала года.",
        "is_dim": True,
        "name": "ym:pc:startOfYear",
        "title": "Дата (год)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
        "name": "ym:pc:startOfDekaminute",
        "title": "Дата (декаминута)",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:startOfQuarter",
        "title": "Дата (квартал)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала месяца.",
        "is_dim": True,
        "name": "ym:pc:startOfMonth",
        "title": "Дата (месяц)",
        "type": "date",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.",
        "is_dim": True,
        "name": "ym:pc:startOfMinute",
        "title": "Дата (минута)",
        "type": "datetime",
    },
    {
        "description": "Дата и время в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
        "name": "ym:pc:startOfHour",
        "title": "Дата (час)",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:dekaminute",
        "title": "Декаминута",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:dayOfMonth",
        "title": "День месяца",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:dayOfWeekName",
        "title": "День недели",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:appID",
        "title": "Идентификатор приложения (в магазине приложений)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:googleAID",
        "title": "Идентификатор рекламы Google",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:iosIFA",
        "title": "Идентификатор рекламы на iOS",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:deviceID",
        "title": "Идентификатор, полученный средствами системного API или выдаваемый Яндексом",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:groupName",
        "title": "Имя группы push-api кампании",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:clientKitVersionDetails",
        "title": "Информация о версии SDK",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:campaignName",
        "title": "Кампания",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:osMajorVersionInfoName",
        "title": "Мажорная версия ОС",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:month",
        "title": "Месяц",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:minute",
        "title": "Минута",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:mobileDeviceModel",
        "title": "Модель",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:mpcName",
        "title": "Название сотового оператора на основе кода сотового оператора (MPC)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:mccName",
        "title": "Название страны на основе мобильного кода страны оператора (MCC)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:buildNumber",
        "title": "Номер сборки приложения",
        "type": "string",
    },
    {
        "description": "Области, к которым принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:pc:regionAreaName",
        "title": "Область",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:operatingSystemInfoName",
        "title": "Операционная система",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:isFromCampaign",
        "title": "Признак для разделения push-api кампаний и интерфейсных кампаний",
        "type": "string",
    },
    {
        "description": 'Возможные значения: "yes", "no", "undefined".',
        "is_dim": True,
        "name": "ym:pc:limitAdTracking",
        "title": "Признак ограничения рекламного трекинга",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:mobileDeviceBranding",
        "title": "Производитель",
        "type": "string",
    },
    {
        "description": 'Подробная информация о разрешении экрана. Например, "2560x1440 px (xxxhdpi)".',
        "is_dim": True,
        "name": "ym:pc:screenResolutionDetailed",
        "title": "Разрешение (подробно)",
        "type": "string",
    },
    {
        "description": "Название страны, к которой принадлежат посетители сайта.",
        "is_dim": True,
        "name": "ym:pc:regionCountryName",
        "title": "Страна",
        "type": "string",
    },
    {
        "description": "Возможные значения: Отправлено, Получено, Проигнорировано, Открыто, Контрольная выборка",
        "is_dim": True,
        "name": "ym:pc:actionTypeName",
        "title": "Тип события Push",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:deviceTypeName",
        "title": "Тип устройства",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:device",
        "title": "Устройство",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:hour",
        "title": "Час",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:hourMinute",
        "title": "Час и минута",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pc:locale",
        "title": "Язык интерфейса",
        "type": "string",
    },
]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/appmetrica/push_events.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import Enum

from dl_sqlalchemy_metrica_api.api_info.metrika.advertising import advertising_fields
from dl_sqlalchemy_metrica_api.api_info.metrika.hits import hits_fields
from dl_sqlalchemy_metrica_api.api_info.metrika.user_param import user_param_fields
from dl_sqlalchemy_metrica_api.api_info.metrika.visits import visits_fields


class MetrikaApiCounterSource(Enum):
    hits = "hits"
    visits = "visits"
    advertising = "advertising"
    user_param = "user_param"


metrica_fields_namespaces = MetrikaApiCounterSource


fields_by_namespace = {
    MetrikaApiCounterSource.hits: hits_fields,
    MetrikaApiCounterSource.visits: visits_fields,
    MetrikaApiCounterSource.advertising: advertising_fields,
    MetrikaApiCounterSource.user_param: user_param_fields,
}


fields_by_name = {f["name"]: f for f in sum(fields_by_namespace.values(), [])}


metrics_by_namespace = {
    cs: list(filter(lambda f: not f["is_dim"], fields)) for cs, fields in fields_by_namespace.items()
}


namespace_by_prefix = {
    "ym:pv": MetrikaApiCounterSource.hits,
    "ym:s": MetrikaApiCounterSource.visits,
    "ym:ad": MetrikaApiCounterSource.advertising,
    "ym:up": MetrikaApiCounterSource.user_param,
}


def get_namespace_by_name(name):
    name_parts = name.split(":", 2)
    if len(name_parts) != 3:
        raise ValueError('Unexpected name format: "%s".' % name)
    prefix = ":".join(name_parts[:2])
    if prefix not in namespace_by_prefix:
        raise ValueError('Unknown field prefix: "%s"' % prefix)
    return namespace_by_prefix[prefix]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/metrika/__init__.py
datalens-tech/datalens-backend
advertising_fields = [
    {
        "description": "Количество посетителей, купивших хотя бы один товар, суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:ad:productPurchasedUniq",
        "title": "Количество посетителей, купивших товар",
        "type": "integer",
    },
    {
        "description": "Среднее время от предпоследнего визита в днях.",
        "is_dim": False,
        "name": "ym:ad:userRecencyDays",
        "title": "Дней от предыдущего визита",
        "type": "float",
    },
    {
        "description": "Количество новых посетителей.",
        "is_dim": False,
        "name": "ym:ad:newUsers",
        "title": "Новые посетители",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, возраст которых менее 18 лет.",
        "is_dim": False,
        "name": "ym:ad:under18AgePercentage",
        "title": "Менее 18 лет",
        "type": "float",
    },
    {
        "description": "Доля визитов и хитов, совершенных с мобильных устройств.",
        "is_dim": False,
        "name": "ym:ad:mobilePercentage",
        "title": "Мобильность",
        "type": "float",
    },
    {
        "description": "Средняя продолжительность визита в минутах и секундах.",
        "is_dim": False,
        "name": "ym:ad:avgVisitDurationSeconds",
        "title": "Время на сайте",
        "type": "float",
    },
    {
        "description": "Сумма значений параметров визитов. Вычисляется, если в качестве значений параметров визитов передаётся число.",
        "is_dim": False,
        "name": "ym:ad:sumParams",
        "title": "Сумма параметров визитов",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через более чем 365 дней.",
        "is_dim": False,
        "name": "ym:ad:overYearUserRecencyPercentage",
        "title": "Вернувшиеся: более чем 365 дней",
        "type": "float",
    },
    {
        "description": "Суммарное количество кликов по рекламе.",
        "is_dim": False,
        "name": "ym:ad:clicks",
        "title": "Клики",
        "type": "integer",
    },
    {
        "description": "Доля визитов и хитов с поддержкой Silverlight.",
        "is_dim": False,
        "name": "ym:ad:silverlightEnabledPercentage",
        "title": "Поддержка Silverlight",
        "type": "float",
    },
    {
        "description": "Количество уникальных посетителей.",
        "is_dim": False,
        "name": "ym:ad:users",
        "title": "Посетители",
        "type": "integer",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:ad:ecommercePurchases",
        "title": "Количество покупок",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, возраст которых от 25 до 34 лет.",
        "is_dim": False,
        "name": "ym:ad:upTo34AgePercentage",
        "title": "25‑34 лет",
        "type": "float",
    },
    {
        "description": "Доля визитов и хитов с поддержкой JavaScript.",
        "is_dim": False,
        "name": "ym:ad:jsEnabledPercentage",
        "title": "Поддержка JavaScript",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:ad:visitsPerDay",
        "title": "Визитов в день",
        "type": "float",
    },
    {
        "description": "Количество посетителей, добавивших в корзину хотя бы один товар суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:ad:productBasketsUniq",
        "title": "Количество посетителей, добавивших товар в корзину",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, известных меньше одного дня.",
        "is_dim": False,
        "name": "ym:ad:upToDaySinceFirstVisitPercentage",
        "title": "Новизна менее одного дня",
        "type": "float",
    },
    {
        "description": "Число просмотров страниц на сайте за отчетный период.",
        "is_dim": False,
        "name": "ym:ad:pageviews",
        "title": "Просмотры",
        "type": "integer",
    },
    {
        "description": "Доля визитов и хитов с поддержкой Flash.",
        "is_dim": False,
        "name": "ym:ad:flashEnabledPercentage",
        "title": "Поддержка Flash",
        "type": "float",
    },
    {
        "description": 'Доля визитов с меткой "GCLID".',
        "is_dim": False,
        "name": "ym:ad:GCLIDPercentage",
        "title": "Доля GCLID",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, известных более 365 дней.",
        "is_dim": False,
        "name": "ym:ad:overYearSinceFirstVisitPercentage",
        "title": "Новизна более 365 дней",
        "type": "float",
    },
    {
        "description": "Доля посетителей, купивших товар на сайте, по отношению к общему числу посетителей.",
        "is_dim": False,
        "name": "ym:ad:usersPurchasePercentage",
        "title": "Доля посетителей, купивших товар",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, возраст которых от 45 до 55 лет.",
        "is_dim": False,
        "name": "ym:ad:upTo54AgePercentage",
        "title": "45‑55 лет",
        "type": "float",
    },
    {
        "description": "Количество купленных единиц товара суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:ad:productPurchasedQuantity",
        "title": "Товаров куплено",
        "type": "integer",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:ad:womanPercentage",
        "title": "Доля женщин",
        "type": "float",
    },
    {
        "description": "Доля визитов новых посетителей.",
        "is_dim": False,
        "name": "ym:ad:newUserVisitsPercentage",
        "title": "Доля визитов новых посетителей",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через 2‑7 дней.",
        "is_dim": False,
        "name": "ym:ad:upToWeekUserRecencyPercentage",
        "title": "Вернувшиеся: 2‑7 дней",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей с частотой визитов один раз в 1 день. Учитываются посетители, совершившие больше двух визитов, которые произошли в разные дни.",
        "is_dim": False,
        "name": "ym:ad:oneDayBetweenVisitsPercentage",
        "title": "Периодичность 1 день",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, известных от 91 до 365 дней.",
        "is_dim": False,
        "name": "ym:ad:upToYearSinceFirstVisitPercentage",
        "title": "Новизна 91‑365 дней",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей с 4‑7 визитами.",
        "is_dim": False,
        "name": "ym:ad:upTo7VisitsPerUserPercentage",
        "title": "4‑7 визитов посетителя",
        "type": "float",
    },
    {
        "description": "Процент уникальных посетителей, посетивших сайт в отчетном периоде, активность которых включала их самый первый за всю историю накопления данных визит на сайт.",
        "is_dim": False,
        "name": "ym:ad:percentNewVisitors",
        "title": "Доля новых посетителей",
        "type": "float",
    },
    {
        "description": "Доля посетителей, предположительно являющихся роботами.",
        "is_dim": False,
        "name": "ym:ad:robotPercentage",
        "title": "Роботность",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через 32‑90 дней.",
        "is_dim": False,
        "name": "ym:ad:upToQuarterUserRecencyPercentage",
        "title": "Вернувшиеся: 32‑90 дней",
        "type": "float",
    },
    {
        "description": "Среднее время от первого визита в днях.",
        "is_dim": False,
        "name": "ym:ad:avgDaysSinceFirstVisit",
        "title": "Дней от первого визита",
        "type": "float",
    },
    {
        "description": "Среднее значение параметра визита. Вычисляется, если в качестве значения параметра визита передаётся число.",
        "is_dim": False,
        "name": "ym:ad:avgParams",
        "title": "Среднее параметров визитов",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, известных от 1 до 7 дней.",
        "is_dim": False,
        "name": "ym:ad:upToWeekSinceFirstVisitPercentage",
        "title": "Новизна 1‑7 дней",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, возраст которых от 35 до 45 лет.",
        "is_dim": False,
        "name": "ym:ad:upTo44AgePercentage",
        "title": "35‑45 лет",
        "type": "float",
    },
    {
        "description": "Количество единиц товара, помещённых в корзину, суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:ad:productBasketsRemoveQuantity",
        "title": "Товаров удалено из корзины",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей с одним визитом (первый визит посетителя).",
        "is_dim": False,
        "name": "ym:ad:oneVisitPerUserPercentage",
        "title": "1 визит посетителя",
        "type": "float",
    },
    {
        "description": (
            "Отношение доли посетителей с заданным интересом на этом сайте к "
            "среднестатистической доле посетителей с этим же интересом на "
            "всех сайтах интернета. Разрешена только при использовании "
            'группировки "Категория интересов".'
        ),
        "is_dim": False,
        "name": "ym:ad:affinityIndexInterests",
        "title": "Аффинити‑индекс (старая версия)",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей с частотой визитов один раз в более чем 31 день. Учитываются посетители, совершившие больше двух визитов, которые произошли в разные дни.",
        "is_dim": False,
        "name": "ym:ad:overMonthBetweenVisitsPercentage",
        "title": "Периодичность более чем 31 день",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через 8‑31 день.",
        "is_dim": False,
        "name": "ym:ad:upToMonthUserRecencyPercentage",
        "title": "Вернувшиеся: 8‑31 день",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, совершивших больше 32 визитов.",
        "is_dim": False,
        "name": "ym:ad:over32VisitsPerUserPercentage",
        "title": "Более чем 32 визита посетителя",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, возраст которых более 55 лет.",
        "is_dim": False,
        "name": "ym:ad:over54AgePercentage",
        "title": "Более 55 лет",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:ad:visitsPerMinute",
        "title": "Визитов в минуту",
        "type": "float",
    },
    {
        "description": "Суммарное количество визитов.",
        "is_dim": False,
        "name": "ym:ad:visits",
        "title": "Визиты",
        "type": "integer",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:ad:manPercentage",
        "title": "Доля мужчин",
        "type": "float",
    },
    {
        "description": "Количество просмотров страницы с информацией о товарах суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:ad:productImpressions",
        "title": "Просмотры товаров",
        "type": "integer",
    },
    {
        "description": "Количество единиц товара, добавленных в корзину, суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:ad:productBasketsQuantity",
        "title": "Товаров добавлено в корзину",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, известных от 8 до 31 день.",
        "is_dim": False,
        "name": "ym:ad:upToMonthSinceFirstVisitPercentage",
        "title": "Новизна 8‑31 день",
        "type": "float",
    },
    {
        "description": "Среднее количество дней между визитами посетителей.",
        "is_dim": False,
        "name": "ym:ad:avgDaysBetweenVisits",
        "title": "Дней между визитами",
        "type": "float",
    },
    {
        "description": "Доля визитов с блокировщиками рекламы.",
        "is_dim": False,
        "name": "ym:ad:blockedPercentage",
        "title": "Блокировка рекламы",
        "type": "float",
    },
    {
        "description": "Количество посетителей, просмотревших информацию хотя бы об одном товаре.",
        "is_dim": False,
        "name": "ym:ad:productImpressionsUniq",
        "title": "Посетители, просмотревшие товар",
        "type": "integer",
    },
    {
        "description": "Количество страниц, просмотренных посетителем во время визита.",
        "is_dim": False,
        "name": "ym:ad:pageDepth",
        "title": "Глубина просмотра",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, возраст которых от 18 до 24 лет.",
        "is_dim": False,
        "name": "ym:ad:upTo24AgePercentage",
        "title": "18‑24 лет",
        "type": "float",
    },
    {
        "description": "Доля визитов и хитов с включенными cookies.",
        "is_dim": False,
        "name": "ym:ad:cookieEnabledPercentage",
        "title": "Поддержка сookies",
        "type": "float",
    },
    {
        "description": "Доля визитов и хитов с поддержкой Java.",
        "is_dim": False,
        "name": "ym:ad:javaEnabledPercentage",
        "title": "Поддержка Java",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей с 8‑31 визитами.",
        "is_dim": False,
        "name": "ym:ad:upTo31VisitsPerUserPercentage",
        "title": "8‑31 визит посетителя",
        "type": "float",
    },
    {
        "description": "Количество раз, когда в визите встретился параметр визита.",
        "is_dim": False,
        "name": "ym:ad:paramsNumber",
        "title": "Количество параметров визитов",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:ad:visitsPerHour",
        "title": "Визитов в час",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей с частотой визитов один раз в 8‑31 день. Учитываются посетители, совершившие больше двух визитов. При этом число дней визитов больше 1 дня.",
        "is_dim": False,
        "name": "ym:ad:upToMonthBetweenVisitsPercentage",
        "title": "Периодичность 8‑31 день",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через 1 день.",
        "is_dim": False,
        "name": "ym:ad:upToDayUserRecencyPercentage",
        "title": "Вернувшиеся: 1 день",
        "type": "float",
    },
    {
        "description": "Доля визитов, в рамках которых состоялся лишь один просмотр страницы, продолжавшийся менее 15 секунд.",
        "is_dim": False,
        "name": "ym:ad:bounceRate",
        "title": "Отказы",
        "type": "float",
    },
    {
        "description": "Конверсия по любой цели.",
        "is_dim": False,
        "name": "ym:ad:anyGoalConversionRate",
        "title": "Конверсия по любой цели",
        "type": "float",
    },
    {
        "description": "Количество достижений любой цели.",
        "is_dim": False,
        "name": "ym:ad:sumGoalReachesAny",
        "title": "Достижения любой цели",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через 91‑365 дней.",
        "is_dim": False,
        "name": "ym:ad:upToYearUserRecencyPercentage",
        "title": "Вернувшиеся: 91‑365 дней",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, известных от 32 до 90 дней.",
        "is_dim": False,
        "name": "ym:ad:upToQuarterSinceFirstVisitPercentage",
        "title": "Новизна 32‑90 дней",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей с 2‑3 визитами.",
        "is_dim": False,
        "name": "ym:ad:upTo3VisitsPerUserPercentage",
        "title": "2‑3 визита посетителя",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей с частотой визитов один раз в 2‑7 дней. Учитываются посетители, совершившие больше двух визитов. При этом число дней визитов больше 1 дня.",
        "is_dim": False,
        "name": "ym:ad:upToWeekBetweenVisitsPercentage",
        "title": "Периодичность 2‑7 дней",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ad:counterID",
        "title": "Счетчик (id)",
        "type": "string",
        "src_key": "id",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ad:counterIDName",
        "title": "Счетчик",
        "type": "string",
    },
    {
        "description": "Поисковая фраза последнего перехода по объявлению Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:lastSignDirectSearchPhrase",
        "title": "Поисковая фраза (Директ) (last sign)",
        "type": "string",
    },
    {
        "description": "Рекламная кампания Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:lastSignDirectOrder",
        "title": "Кампания Яндекс.Директа (last sign)",
        "type": "string",
    },
    {
        "description": "Валюта, установленная в Яндекс.Директе для рекламной кампании.",
        "is_dim": True,
        "name": "ym:ad:firstCurrencyIDName",
        "title": "Валюта (first)",
        "type": "string",
    },
    {
        "description": "Тип условия показа объявления.",
        "is_dim": True,
        "name": "ym:ad:lastDirectConditionType",
        "title": "Тип условия показа объявления (last)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ad:lastDirectBanner",
        "title": "Объявление Яндекс.Директа (last)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ad:directBanner",
        "title": "Объявление Яндекс.Директа",
        "type": "string",
    },
    {
        "description": "Округа, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:ad:regionDistrict",
        "title": "Округ",
        "type": "string",
    },
    {
        "description": "Континенты, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:ad:regionContinent",
        "title": "Континент",
        "type": "string",
    },
    {
        "description": "Категория коммерческих интересов посетителей.",
        "is_dim": True,
        "name": "ym:ad:interest",
        "title": "Категория интересов",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ad:directBannerText",
        "title": "Текст объявления",
        "type": "string",
    },
    {
        "description": "Условие показа объявления. Условием могут быть либо ключевые слова, либо условие ретаргетинга.",
        "is_dim": True,
        "name": "ym:ad:firstDirectPhraseOrCond",
        "title": "Условие показа объявления (first)",
        "type": "string",
    },
    {
        "description": "Рекламная площадка Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:directPlatform",
        "title": "Площадка",
        "type": "string",
    },
    {
        "description": "Возраст посетителя, интервал",
        "is_dim": True,
        "name": "ym:ad:ageInterval",
        "title": "Возраст",
        "type": "string",
    },
    {
        "description": "Условие показа объявления. Условием могут быть либо ключевые слова, либо условие ретаргетинга.",
        "is_dim": True,
        "name": "ym:ad:lastSignDirectPhraseOrCond",
        "title": "Условие показа объявления (last sign)",
        "type": "string",
    },
    {
        "description": "Рекламная кампания Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:directOrder",
        "title": "Кампания Яндекс.Директа",
        "type": "string",
    },
    {
        "description": "Рекламная кампания Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:lastDirectOrder",
        "title": "Кампания Яндекс.Директа (last)",
        "type": "string",
    },
    {
        "description": "Валюта, установленная в Яндекс.Директе для рекламной кампании.",
        "is_dim": True,
        "name": "ym:ad:lastCurrencyIDName",
        "title": "Валюта (last)",
        "type": "string",
    },
    {
        "description": "Условие показа объявления. Условием могут быть либо ключевые слова, либо условие ретаргетинга.",
        "is_dim": True,
        "name": "ym:ad:directPhraseOrCond",
        "title": "Условие показа объявления",
        "type": "string",
    },
    {
        "description": "Рекламная кампания Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:firstDirectOrder",
        "title": "Кампания Яндекс.Директа (first)",
        "type": "string",
    },
    {
        "description": "Города, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:ad:regionCity",
        "title": "Город",
        "type": "string",
    },
    {
        "description": "Условие показа объявления. Условием могут быть либо ключевые слова, либо условие ретаргетинга.",
        "is_dim": True,
        "name": "ym:ad:lastDirectPhraseOrCond",
        "title": "Условие показа объявления (last)",
        "type": "string",
    },
    {
        "description": "Валюта, установленная в Яндекс.Директе для рекламной кампании.",
        "is_dim": True,
        "name": "ym:ad:lastSignCurrencyIDName",
        "title": "Валюта (last sign)",
        "type": "string",
    },
    {
        "description": 'Пол посетителя. Возможные значения: "мужской" или "женский".',
        "is_dim": True,
        "name": "ym:ad:gender",
        "title": "Пол",
        "type": "string",
    },
    {
        "description": "Валюта, установленная в Яндекс.Директе для рекламной кампании.",
        "is_dim": True,
        "name": "ym:ad:currency",
        "title": "Валюта",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ad:firstDirectBanner",
        "title": "Объявление Яндекс.Директа (first)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ad:lastSignDirectBanner",
        "title": "Объявление Яндекс.Директа (last sign)",
        "type": "string",
    },
    {
        "description": "Тип условия показа объявления.",
        "is_dim": True,
        "name": "ym:ad:firstDirectConditionType",
        "title": "Тип условия показа объявления (first)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:ad:displayCampaign",
        "title": "Номер заказа Яндекс.Дисплея",
        "type": "string",
    },
    {
        "description": "Тип условия показа объявления.",
        "is_dim": True,
        "name": "ym:ad:lastSignDirectConditionType",
        "title": "Тип условия показа объявления (last sign)",
        "type": "string",
    },
    {
        "description": "Группа объявлений Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:directBannerGroup",
        "title": "Группа объявлений",
        "type": "string",
    },
    {
        "description": "Поисковая фраза последнего перехода по объявлению Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:lastDirectSearchPhrase",
        "title": "Поисковая фраза (Директ) (last)",
        "type": "string",
    },
    {
        "description": "Поисковая фраза последнего перехода по объявлению Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:firstDirectSearchPhrase",
        "title": "Поисковая фраза (Директ) (first)",
        "type": "string",
    },
    {
        "description": "Размер города по населению.",
        "is_dim": True,
        "name": "ym:ad:regionCitySize",
        "title": "Размер города",
        "type": "string",
    },
    {
        "description": "Поисковая фраза последнего перехода по объявлению Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:directSearchPhrase",
        "title": "Поисковая фраза (Директ)",
        "type": "string",
    },
    {
        "description": "Тип условия показа объявления.",
        "is_dim": True,
        "name": "ym:ad:directConditionType",
        "title": "Тип условия показа объявления",
        "type": "string",
    },
    {
        "description": "Тип рекламной площадки Яндекс.Директа.",
        "is_dim": True,
        "name": "ym:ad:directPlatformType",
        "title": "Тип площадки",
        "type": "string",
    },
    {
        "description": "Области, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:ad:regionArea",
        "title": "Область",
        "type": "string",
    },
    {
        "description": "Страны, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:ad:regionCountry",
        "title": "Страна",
        "type": "string",
    },
]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/metrika/advertising.py
datalens-tech/datalens-backend
hits_fields = [
    {
        "description": "Среднее количество хитов в минуту.",
        "is_dim": False,
        "name": "ym:pv:pageviewsPerMinute",
        "title": "Просмотров в минуту",
        "type": "float",
    },
    {
        "description": "Доля визитов и хитов с поддержкой JavaScript.",
        "is_dim": False,
        "name": "ym:pv:jsEnabledPercentage",
        "title": "Поддержка JavaScript",
        "type": "float",  # percents
    },
    {
        "description": "Число просмотров страниц на сайте за отчетный период.",
        "is_dim": False,
        "name": "ym:pv:pageviews",
        "title": "Просмотры",
        "type": "integer",
    },
    {
        "description": "Средняя глубина просмотра в день.",
        "is_dim": False,
        "name": "ym:pv:pageviewsPerDay",
        "title": "Просмотров в день",
        "type": "float",
    },
    {
        "description": "Количество уникальных посетителей.",
        "is_dim": False,
        "name": "ym:pv:users",
        "title": "Посетители",
        "type": "integer",
    },
    {
        "description": "Среднее количество хитов в час.",
        "is_dim": False,
        "name": "ym:pv:pageviewsPerHour",
        "title": "Просмотров в час",
        "type": "float",
    },
    {
        "description": "Доля визитов с блокировщиками рекламы.",
        "is_dim": False,
        "name": "ym:pv:blockedPercentage",
        "title": "Блокировка рекламы",
        "type": "float",  # percents
    },
    {
        "description": "Доля визитов и хитов с включенными cookies.",
        "is_dim": False,
        "name": "ym:pv:cookieEnabledPercentage",
        "title": "Поддержка сookies",
        "type": "float",  # percents
    },
    {
        "description": "Доля визитов и хитов, совершенных с мобильных устройств.",
        "is_dim": False,
        "name": "ym:pv:mobilePercentage",
        "title": "Мобильность",
        "type": "float",  # percents
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:URLDomain",
        "title": "Домен страницы",
        "type": "string",
    },
    {
        "description": (
            'Часть URL страницы, следующая сразу после домена до "?" (секция параметров), '
            '"#" (секция фрагмента) либо до конца строки. Например, для страницы '
            '"http://news.yandex.ru/quotes/1.html" путем является "/quotes/1.html", '
            'для "http://news.yandex.ru/" - "/", а для "http://news.yandex.ru" - пустая строка'
        ),
        "is_dim": True,
        "name": "ym:pv:URLPath",
        "title": "Путь страницы",
        "type": "string",
    },
    {
        "description": 'IPv4/IPv6-адрес посетителя с обнуленными 1/8 последними байтами соответственно. Например, "77.88.21.0"/"2001:db8:85a3::".',
        "is_dim": True,
        "name": "ym:pv:ipAddress",
        "title": "IP-адрес",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:browserHourMinute",
        "title": "Час и минута на компьютере посетителя",
        "type": "string",
    },
    {
        "description": 'Ориентация экрана. Возможные значения: "landscape", "portrait".',
        "is_dim": True,
        "name": "ym:pv:screenOrientationName",
        "title": "Ориентация экрана",
        "type": "string",
    },
    {
        "description": "Дата и время просмотра в формате YYYY-MM-DD HH:mm:ss.",
        "is_dim": True,
        "name": "ym:pv:dateTime",
        "title": "Дата и время просмотра",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:clientTimeZone",
        "title": "Часовой пояс на компьютере посетителя",
        "type": "string",
    },
    {
        "description": "Дата просмотра в формате YYYY-MM-DD, округленное до начала квартала.",
        "is_dim": True,
        "name": "ym:pv:startOfQuarter",
        "title": "Квартал просмотра",
        "type": "date",
    },
    {
        "description": "Источник перехода",
        "is_dim": True,
        "name": "ym:pv:UTMSource",
        "title": "UTM Source",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала квартала.",
        "is_dim": True,
        "name": "ym:pv:browserStartOfQuarter",
        "title": "Квартал на компьютере посетителя",
        "type": "datetime",
    },
    {
        "description": 'Первый уровень URL страницы. Например, для "http://ya.ru/1/2/3/4.html" первым уровнем URL является "http://ya.ru/".',
        "is_dim": True,
        "name": "ym:pv:URLPathLevel1",
        "title": "Адрес, ур. 1",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:URL",
        "title": "Адрес страницы",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:month",
        "title": "Месяц просмотра",
        "type": "string",
    },
    {
        "description": 'Глубина цвета экрана в битах. Например, "24".',
        "is_dim": True,
        "name": "ym:pv:screenColorsName",
        "title": "Глубина цвета",
        "type": "string",
    },
    {
        "description": 'Пятый уровень URL страницы входа. Например, для "http://ya.ru/1/2/3/4.html" пятым уровнем URL является "http://ya.ru/1/2/3/4.html".',
        "is_dim": True,
        "name": "ym:pv:URLPathLevel5",
        "title": "Адрес, ур. 5",
        "type": "string",
    },
    {
        "description": 'Четвертый уровень URL страницы. Например, для "http://ya.ru/1/2/3/4.html" четвертым уровнем URL является "http://ya.ru/1/2/3/".',
        "is_dim": True,
        "name": "ym:pv:URLPathLevel4",
        "title": "Адрес, ур. 4",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
        "name": "ym:pv:browserStartOfDekaminute",
        "title": "Начало декаминуты на компьютере посетителя",
        "type": "datetime",
    },
    {
        "description": 'Третий уровень URL страницы. Например, для "http://ya.ru/1/2/3/4.html" третьим уровнем URL является "http://ya.ru/1/2/".',
        "is_dim": True,
        "name": "ym:pv:URLPathLevel3",
        "title": "Адрес, ур. 3",
        "type": "string",
    },
    {
        "description": 'Второй уровень URL страницы. Например, для "http://ya.ru/1/2/3/4.html" вторым уровнем URL является "http://ya.ru/1/".',
        "is_dim": True,
        "name": "ym:pv:URLPathLevel2",
        "title": "Адрес, ур. 2",
        "type": "string",
    },
    {
        "description": 'Тип устройства, с которого было посещение. Возможные значения: "desktop", "mobile", "tablet", "tv".',
        "is_dim": True,
        "name": "ym:pv:deviceCategoryName",
        "title": "Тип устройства",
        "type": "string",
    },
    {
        "description": "Дата на компьютере посетителя в формате YYYY-MM-DD.",
        "is_dim": True,
        "name": "ym:pv:browserDate",
        "title": "Дата на компьютере посетителя",
        "type": "date",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:hourName",
        "title": "Час просмотра в формате HH:00",
        "type": "string",
    },
    {
        "description": 'Разрешение экрана. Например, "1920х1080".',
        "is_dim": True,
        "name": "ym:pv:screenResolution",
        "title": "Разрешение",
        "type": "string",
    },
    {
        "description": "Континенты, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:pv:regionContinentName",
        "title": "Континент",
        "type": "string",
    },
    # {
    #     'description': 'Дата просмотра в формате YYYY-MM-DD, округленное до начала месяца.',
    #     'is_dim': True,
    #     'name': 'ym:pv:startOfMonth',
    #     'title': 'Месяц просмотра',
    #     'type': 'date',
    # },
    {
        "description": "Объявление",
        "is_dim": True,
        "name": "ym:pv:UTMContent",
        "title": "UTM Content",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:mobilePhoneName",
        "title": "Производитель устройства",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:hasAdBlockerName",
        "title": "Наличие блокировщиков рекламы",
        "type": "string",
    },
    {
        "description": "Возраст посетителя, интервал",
        "is_dim": True,
        "name": "ym:pv:ageIntervalName",
        "title": "Возраст",
        "type": "string",
    },
    {
        "description": 'Движок браузера посетителя. Например, "WebKit".',
        "is_dim": True,
        "name": "ym:pv:browserEngine",
        "title": "Движок браузера",
        "type": "string",
    },
    {
        "description": "Рекламная кампания",
        "is_dim": True,
        "name": "ym:pv:openstatCampaign",
        "title": "Openstat Campaign",
        "type": "string",
    },
    {
        "description": 'Третий уровень URL реферера. Например, для "http://ya.ru/1/2/3/4.html" третьим уровнем URL является "http://ya.ru/1/2/".',
        "is_dim": True,
        "name": "ym:pv:refererPathLevel3",
        "title": "Реферер, ур. 3",
        "type": "string",
    },
    {
        "description": 'Второй уровень URL реферера. Например, для "http://ya.ru/1/2/3/4.html" вторым уровнем URL является "http://ya.ru/1/".',
        "is_dim": True,
        "name": "ym:pv:refererPathLevel2",
        "title": "Реферер, ур. 2",
        "type": "string",
    },
    {
        "description": 'Первый уровень URL реферера. Например, для "http://ya.ru/1/2/3/4.html" первым уровнем URL является "http://ya.ru/".',
        "is_dim": True,
        "name": "ym:pv:refererPathLevel1",
        "title": "Реферер, ур. 1",
        "type": "string",
    },
    {
        "description": "Рекламный сервис",
        "is_dim": True,
        "name": "ym:pv:openstatService",
        "title": "Openstat Service",
        "type": "string",
    },
    {
        "description": "Средство маркетинга",
        "is_dim": True,
        "name": "ym:pv:UTMMedium",
        "title": "UTM Medium",
        "type": "string",
    },
    {
        "description": 'Реальное разрешение экрана исходя из данных "device pixel ratio". Например, "1920х1080".',
        "is_dim": True,
        "name": "ym:pv:physicalScreenResolution",
        "title": "Реальное разрешение",
        "type": "string",
    },
    {
        "description": 'Пятый уровень URL реферера. Например, для "http://ya.ru/1/2/3/4.html" пятым уровнем URL является "http://ya.ru/1/2/3/4.html".',
        "is_dim": True,
        "name": "ym:pv:refererPathLevel5",
        "title": "Реферер, ур. 5",
        "type": "string",
    },
    {
        "description": 'Четвертый уровень URL реферера. Например, для "http://ya.ru/1/2/3/4.html" четвертым уровнем URL является "http://ya.ru/1/2/3/".',
        "is_dim": True,
        "name": "ym:pv:refererPathLevel4",
        "title": "Реферер, ур. 4",
        "type": "string",
    },
    {
        "description": "Название браузера и его версия.",
        "is_dim": True,
        "name": "ym:pv:browserAndVersionMajorName",
        "title": "Версия браузера",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:browserHourName",
        "title": "Час на компьютере посетителя в формате HH:00",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:operatingSystemName",
        "title": "Операционная система (детально)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:dayOfWeekName",
        "title": "День недели просмотра",
        "type": "string",
    },
    {
        "description": "Размер города по населению.",
        "is_dim": True,
        "name": "ym:pv:regionCitySizeName",
        "title": "Размер города",
        "type": "string",
    },
    {
        "description": "Области, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:pv:regionAreaName",
        "title": "Область",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:mobilePhoneModel",
        "title": "Модель устройства",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:refererProto",
        "title": "Протокол реферера",
        "type": "string",
    },
    {
        "description": "URL реферера",
        "is_dim": True,
        "name": "ym:pv:referer",
        "title": "Реферер",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:URLParamNameAndValue",
        "title": "Значение параметра URL",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:from",
        "title": "Метка from",
        "type": "string",
    },
    # {
    #     'description': None,
    #     'is_dim': True,
    #     'name': 'ym:pv:browserDatePeriod<group>',
    #     'title': 'Дата на компьютере посетителя',
    #     'type': 'string',
    # },
    {
        "description": "Дата просмотра в формате YYYY-MM-DD.",
        "is_dim": True,
        "name": "ym:pv:date",
        "title": "Дата просмотра",
        "type": "date",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:URLProto",
        "title": "Протокол страницы",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:operatingSystemRootName",
        "title": "Группа операционных систем",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:browserMonth",
        "title": "Месяц на компьютере посетителя",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:counterID",
        "title": "Счетчик (id)",
        "type": "string",
        "src_key": "id",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:counterIDName",
        "title": "Счетчик",
        "type": "string",
    },
    {
        "description": "Города, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:pv:regionCityName",
        "title": "Город",
        "type": "string",
    },
    {
        "description": 'Наличие метки "GCLID".',
        "is_dim": True,
        "name": "ym:pv:hasGCLIDName",
        "title": "Наличие GCLID",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:cookieEnabledName",
        "title": "Наличие Cookie",
        "type": "string",
    },
    {
        "description": "Часть URL реферера, следующая сразу после домена. В отличие от простого пути также содержит секции параметров и фрагмента.",
        "is_dim": True,
        "name": "ym:pv:refererPathFull",
        "title": "Путь (полный) реферера",
        "type": "string",
    },
    # {
    #     'description': 'Дата просмотра в формате YYYY-MM-DD, округленное до начала года.',
    #     'is_dim': True,
    #     'name': 'ym:pv:startOfYear',
    #     'title': 'Год просмотра',
    #     'type': 'string',
    # },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:dayOfMonth",
        "title": "День месяца просмотра",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:isTurboPageName",
        "title": "Турбо-страница",
        "type": "string",
    },
    {
        "description": "Часть URL страницы, следующая сразу после домена. В отличие от простого пути также содержит секции параметров и фрагмента",
        "is_dim": True,
        "name": "ym:pv:URLPathFull",
        "title": "Путь (полный) страницы",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:title",
        "title": "Заголовок страницы",
        "type": "string",
    },
    # {
    #     'description': 'Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала года.',
    #     'is_dim': True,
    #     'name': 'ym:pv:browserStartOfYear',
    #     'title': 'Год на компьютере посетителя',
    #     'type': 'string',
    # },
    {
        "description": "Ключевые слова",
        "is_dim": True,
        "name": "ym:pv:UTMTerm",
        "title": "UTM Term",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:refererDomain",
        "title": "Домен реферера",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss.",
        "is_dim": True,
        "name": "ym:pv:browserDateTime",
        "title": "Дата и время на компьютере посетителя",
        "type": "datetime",
    },
    {
        "description": "Страны, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:pv:regionCountryName",
        "title": "Страна",
        "type": "string",
    },
    {
        "description": "Рекламное объявление",
        "is_dim": True,
        "name": "ym:pv:openstatAd",
        "title": "Openstat Ad",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:year",
        "title": "Год просмотра",
        "type": "string",
    },
    {
        "description": "Название проводимой рекламной кампании",
        "is_dim": True,
        "name": "ym:pv:UTMCampaign",
        "title": "UTM Campaign",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:URLParamName",
        "title": "Параметр URL",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:browserYear",
        "title": "Год на компьютере посетителя",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:browserDekaminute",
        "title": "Декаминута на компьютере посетителя",
        "type": "string",
    },
    {
        "description": (
            'Часть URL реферера, следующая сразу после домена до "?" (секция параметров), '
            '"#" (секция фрагмента) либо до конца строки. Например, для страницы '
            '"http://news.yandex.ru/quotes/1.html" путем является "/quotes/1.html", '
            'для "http://news.yandex.ru/"&#160;— "/", '
            'а для "http://news.yandex.ru"&#160;— пустая строка.'
        ),
        "is_dim": True,
        "name": "ym:pv:refererPath",
        "title": "Путь реферера",
        "type": "string",
    },
    {
        "description": 'Отношение ширины к высоте экрана. Например, "16:9".',
        "is_dim": True,
        "name": "ym:pv:screenFormat",
        "title": "Соотношение сторон",
        "type": "string",
    },
    {
        "description": "Место размещения",
        "is_dim": True,
        "name": "ym:pv:openstatSource",
        "title": "Openstat Source",
        "type": "string",
    },
    {
        "description": 'Пол посетителя. Возможные значения: "мужской" или "женский".',
        "is_dim": True,
        "name": "ym:pv:genderName",
        "title": "Пол",
        "type": "string",
    },
    {
        "description": "Логическая ширина экрана.",
        "is_dim": True,
        "name": "ym:pv:screenWidth",
        "title": "Логическая ширина",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:browserLanguage",
        "title": "Язык браузера",
        "type": "string",
    },
    {
        "description": 'Физическая высота экрана исходя из данных "device pixel ratio".',
        "is_dim": True,
        "name": "ym:pv:physicalScreenHeight",
        "title": "Физическая высота",
        "type": "string",
    },
    # {
    #     'description': 'Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала месяца.',
    #     'is_dim': True,
    #     'name': 'ym:pv:browserStartOfMonth',
    #     'title': 'Месяц на компьютере посетителя',
    #     'type': 'string',
    # },
    {
        "description": "Дата и время просмотра в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
        "name": "ym:pv:startOfDekaminute",
        "title": "Начало декаминуты просмотра",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:javascriptEnabledName",
        "title": "Наличие Javascript",
        "type": "string",
    },
    {
        "description": "Первое число в версии движка браузера посетителя.",
        "is_dim": True,
        "name": "ym:pv:browserEngineVersion1",
        "title": "Major-версия движка браузера",
        "type": "string",
    },
    {
        "description": "Второе число в версии движка браузера посетителя.",
        "is_dim": True,
        "name": "ym:pv:browserEngineVersion2",
        "title": "Minor-версия движка браузера",
        "type": "string",
    },
    {
        "description": "Ширина клиентской части окна браузера.",
        "is_dim": True,
        "name": "ym:pv:windowClientWidth",
        "title": "Ширина окна",
        "type": "string",
    },
    {
        "description": "Третье число (обычно номер сборки) в версии движка браузера посетителя.",
        "is_dim": True,
        "name": "ym:pv:browserEngineVersion3",
        "title": "Build-версия движка браузера",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала недели.",
        "is_dim": True,
        "name": "ym:pv:browserStartOfWeek",
        "title": "Неделя на компьютере посетителя",
        "type": "string",
    },
    {
        "description": 'Браузер посетителя. Например, "Яндекс.Браузер".',
        "is_dim": True,
        "name": "ym:pv:browserName",
        "title": "Браузер",
        "type": "string",
    },
    {
        "description": "Четвертое число (обычно номер ревизии) в версии движка браузера посетителя.",
        "is_dim": True,
        "name": "ym:pv:browserEngineVersion4",
        "title": "Revision-версия движка браузера",
        "type": "string",
    },
    # # Currently we don't support parametrized dimensions
    # {
    #     'description': 'Интервал дат с параметризацией размера интервала',
    #     'is_dim': True,
    #     'name': 'ym:pv:datePeriod<group>',
    #     'title': 'Дата',
    #     'type': 'date',
    # },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:dekaminute",
        "title": "Декаминута просмотра",
        "type": "string",
    },
    {
        "description": "Логическая высота экрана.",
        "is_dim": True,
        "name": "ym:pv:screenHeight",
        "title": "Логическая высота",
        "type": "string",
    },
    {
        "description": 'Физическая ширина экрана исходя из данных "device pixel ratio".',
        "is_dim": True,
        "name": "ym:pv:physicalScreenWidth",
        "title": "Физическая ширина",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
        "name": "ym:pv:browserStartOfHour",
        "title": "Час на компьютере посетителя",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:hourMinute",
        "title": "Час и минута просмотра",
        "type": "string",
    },
    # {
    #     'description': 'Дата и время просмотра в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.',
    #     'is_dim': True,
    #     'name': 'ym:pv:startOfMinute',
    #     'title': 'Минута просмотра',
    #     'type': 'datetime',
    # },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:browserCountry",
        "title": "Страна браузера",
        "type": "string",
    },
    {
        "description": "Округа, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:pv:regionDistrictName",
        "title": "Округ",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:turboPageID",
        "title": "Турбо-страница (ID)",
        "type": "string",
    },
    {
        "description": "Браузер и его полная версия.",
        "is_dim": True,
        "name": "ym:pv:browserAndVersionName",
        "title": "Полная версия браузера",
        "type": "string",
    },
    # {
    #     'description': 'Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.',
    #     'is_dim': True,
    #     'name': 'ym:pv:browserStartOfMinute',
    #     'title': 'Минута на компьютере посетителя',
    #     'type': 'datetime',
    # },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:browserDayOfWeekName",
        "title": "День недели на компьютере посетителя",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:browserMinute",
        "title": "Минута на компьютере посетителя",
        "type": "string",
    },
    {
        "description": "Дата и время просмотра в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
        "name": "ym:pv:startOfHour",
        "title": "Час просмотра",
        "type": "datetime",
    },
    {
        "description": "Высота клиентской части окна браузера.",
        "is_dim": True,
        "name": "ym:pv:windowClientHeight",
        "title": "Высота окна",
        "type": "string",
    },
    {
        "description": 'Размер клиентской части окна браузера. Например, "1920х1080".',
        "is_dim": True,
        "name": "ym:pv:windowClientArea",
        "title": "Размер окна",
        "type": "string",
    },
    {
        "description": "Дата просмотра в формате YYYY-MM-DD, округленное до начала недели.",
        "is_dim": True,
        "name": "ym:pv:startOfWeek",
        "title": "Неделя просмотра",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:minute",
        "title": "Минута просмотра",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:pv:browserDayOfMonth",
        "title": "День месяца на компьютере посетителя",
        "type": "string",
    },
]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/metrika/hits.py
datalens-tech/datalens-backend
user_param_fields = [
    {
        "description": "Количество параметров.",
        "is_dim": False,
        "name": "ym:up:params",
        "title": "Параметры",
        "type": "integer",
    },
    {
        "description": "Количество уникальных посетителей.",
        "is_dim": False,
        "name": "ym:up:users",
        "title": "Посетители",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:up:counterID",
        "title": "Счетчик (id)",
        "type": "string",
        "src_key": "id",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:up:counterIDName",
        "title": "Счетчик",
        "type": "string",
    },
    {
        "description": "Первый уровень вложенности параметров посетителя",
        "is_dim": True,
        "name": "ym:up:paramsLevel1",
        "title": "Параметр посетителя, ур. 1",
        "type": "string",
    },
    {
        "description": "Второй уровень вложенности параметров посетителя",
        "is_dim": True,
        "name": "ym:up:paramsLevel2",
        "title": "Параметр посетителя, ур. 2",
        "type": "string",
    },
    {
        "description": "Третий уровень вложенности параметров посетителя",
        "is_dim": True,
        "name": "ym:up:paramsLevel3",
        "title": "Параметр посетителя, ур. 3",
        "type": "string",
    },
    {
        "description": "Четвертый уровень вложенности параметров посетителя",
        "is_dim": True,
        "name": "ym:up:paramsLevel4",
        "title": "Параметр посетителя, ур. 4",
        "type": "string",
    },
    {
        "description": "Пятый уровень вложенности параметров посетителя",
        "is_dim": True,
        "name": "ym:up:paramsLevel5",
        "title": "Параметр посетителя, ур. 5",
        "type": "string",
    },
]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/metrika/user_param.py
datalens-tech/datalens-backend
visits_fields = [
    {
        "description": "Доля визитов посетителей, известных от 8 до 31 день.",
        "is_dim": False,
        "name": "ym:s:upToMonthSinceFirstVisitPercentage",
        "title": "Новизна 8‑31 день",
        "type": "float",
    },
    {
        "description": (
            "Доля визитов посетителей с частотой визитов один раз в 1 день. "
            "Учитываются посетители, совершившие больше двух визитов, "
            "которые произошли в разные дни."
        ),
        "is_dim": False,
        "name": "ym:s:oneDayBetweenVisitsPercentage",
        "title": "Периодичность 1 день",
        "type": "float",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:s:manPercentage",
        "title": "Доля мужчин",
        "type": "float",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:s:womanPercentage",
        "title": "Доля женщин",
        "type": "float",
    },
    {
        "description": "Количество посетителей, добавивших в корзину хотя бы один товар суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:s:productBasketsUniq",
        "title": "Количество посетителей, добавивших товар в корзину",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, возраст которых от 35 до 45 лет.",
        "is_dim": False,
        "name": "ym:s:upTo44AgePercentage",
        "title": "35‑45 лет",
        "type": "float",
    },
    {
        "description": "Процент уникальных посетителей, посетивших сайт в отчетном периоде, активность которых включала их самый первый за всю историю накопления данных визит на сайт.",
        "is_dim": False,
        "name": "ym:s:percentNewVisitors",
        "title": "Доля новых посетителей",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallsFirstTimeCallerPercentage",
        "title": "Доля первичных звонков",
        "type": "float",
    },
    {
        "description": "Доля визитов, в рамках которых состоялся лишь один просмотр страницы, продолжавшийся менее 15 секунд.",
        "is_dim": False,
        "name": "ym:s:bounceRate",
        "title": "Отказы",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, возраст которых более 55 лет.",
        "is_dim": False,
        "name": "ym:s:over54AgePercentage",
        "title": "Более 55 лет",
        "type": "float",
    },
    {
        "description": "Количество уникальных посетителей.",
        "is_dim": False,
        "name": "ym:s:users",
        "title": "Посетители",
        "type": "integer",
    },
    {
        "description": "Среднее время от первого визита в днях.",
        "is_dim": False,
        "name": "ym:s:avgDaysSinceFirstVisit",
        "title": "Дней от первого визита (среднее)",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallsFirstTimeCaller",
        "title": "Количество первичных звонков",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей с одним визитом (первый визит посетителя).",
        "is_dim": False,
        "name": "ym:s:oneVisitPerUserPercentage",
        "title": "1 визит посетителя",
        "type": "float",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:s:ecommerceRevenuePerPurchase",
        "title": "Средний доход покупки",
        "type": "float",
    },
    {
        "description": "Количество раз, когда в визите встретился параметр визита.",
        "is_dim": False,
        "name": "ym:s:paramsNumber",
        "title": "Количество параметров визитов",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, известных более 365 дней.",
        "is_dim": False,
        "name": "ym:s:overYearSinceFirstVisitPercentage",
        "title": "Новизна более 365 дней",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей с 8‑31 визитами.",
        "is_dim": False,
        "name": "ym:s:upTo31VisitsPerUserPercentage",
        "title": "8‑31 визит посетителя",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCalls",
        "title": "Количество звонков",
        "type": "integer",
    },
    {
        "description": "Доля посетителей, купивших товар на сайте, по отношению к общему числу посетителей.",
        "is_dim": False,
        "name": "ym:s:usersPurchasePercentage",
        "title": "Доля посетителей, купивших товар",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, известных от 32 до 90 дней.",
        "is_dim": False,
        "name": "ym:s:upToQuarterSinceFirstVisitPercentage",
        "title": "Новизна 32‑90 дней",
        "type": "float",
    },
    {
        "description": "Среднее количество дней между визитами посетителей.",
        "is_dim": False,
        "name": "ym:s:avgDaysBetweenVisits",
        "title": "Дней между визитами",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallsMissed",
        "title": "Количество пропущенных звонков",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:visitsPerDay",
        "title": "Визитов в день",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей с 2‑3 визитами.",
        "is_dim": False,
        "name": "ym:s:upTo3VisitsPerUserPercentage",
        "title": "2‑3 визита посетителя",
        "type": "float",
    },
    {
        "description": (
            "Доля визитов посетителей с частотой визитов один раз в 8‑31 день. "
            "Учитываются посетители, совершившие больше двух визитов. "
            "При этом число дней визитов больше 1 дня."
        ),
        "is_dim": False,
        "name": "ym:s:upToMonthBetweenVisitsPercentage",
        "title": "Периодичность 8‑31 день",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallRevenueAvg",
        "title": "Средняя цена звонка",
        "type": "float",
    },
    {
        "description": "Стоимость купленных товаров суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:s:productPurchasedPrice",
        "title": "Стоимость купленных товаров",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через 91‑365 дней.",
        "is_dim": False,
        "name": "ym:s:upToYearUserRecencyPercentage",
        "title": "Вернувшиеся: 91‑365 дней",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, известных от 1 до 7 дней.",
        "is_dim": False,
        "name": "ym:s:upToWeekSinceFirstVisitPercentage",
        "title": "Новизна 1‑7 дней",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallsMissedPercentage",
        "title": "Доля пропущенных звонков",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallHoldDurationTillAnswerAvg",
        "title": "Среднее время ожидания до ответа",
        "type": "float",
    },
    {
        "description": "Количество посетителей, купивших хотя бы один товар, суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:s:productPurchasedUniq",
        "title": "Количество посетителей, купивших товар",
        "type": "integer",
    },
    {
        "description": "Количество единиц товара, добавленных в корзину, суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:s:productBasketsQuantity",
        "title": "Товаров добавлено в корзину",
        "type": "integer",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallDurationAvg",
        "title": "Среднее время звонка",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через 8‑31 день.",
        "is_dim": False,
        "name": "ym:s:upToMonthUserRecencyPercentage",
        "title": "Вернувшиеся: 8‑31 день",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:pvlAll3Window",
        "title": "PVL‑3",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, возраст которых от 45 до 55 лет.",
        "is_dim": False,
        "name": "ym:s:upTo54AgePercentage",
        "title": "45‑55 лет",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:pvlAll1Window",
        "title": "PVL‑1",
        "type": "float",
    },
    {
        "description": "Количество достижений любой цели.",
        "is_dim": False,
        "name": "ym:s:sumGoalReachesAny",
        "title": "Достижения любой цели",
        "type": "integer",
    },
    {
        "description": "Количество посетителей, просмотревших информацию хотя бы об одном товаре.",
        "is_dim": False,
        "name": "ym:s:productImpressionsUniq",
        "title": "Посетители, просмотревшие товар",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через 1 день.",
        "is_dim": False,
        "name": "ym:s:upToDayUserRecencyPercentage",
        "title": "Вернувшиеся: 1 день",
        "type": "float",
    },
    {
        "description": "Конверсия по любой цели.",
        "is_dim": False,
        "name": "ym:s:anyGoalConversionRate",
        "title": "Конверсия по любой цели",
        "type": "float",
    },
    {
        "description": "Суммарное количество визитов.",
        "is_dim": False,
        "name": "ym:s:visits",
        "title": "Визиты",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, совершивших больше 32 визитов.",
        "is_dim": False,
        "name": "ym:s:over32VisitsPerUserPercentage",
        "title": "Более чем 32 визита посетителя",
        "type": "float",
    },
    {
        "description": 'Доля визитов с меткой "GCLID".',
        "is_dim": False,
        "name": "ym:s:GCLIDPercentage",
        "title": "Доля GCLID",
        "type": "float",
    },
    {
        "description": (
            "Доля визитов посетителей с частотой визитов один раз в более чем 31 день. "
            "Учитываются посетители, совершившие больше двух визитов, "
            "которые произошли в разные дни."
        ),
        "is_dim": False,
        "name": "ym:s:overMonthBetweenVisitsPercentage",
        "title": "Периодичность более чем 31 день",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, возраст которых от 25 до 34 лет.",
        "is_dim": False,
        "name": "ym:s:upTo34AgePercentage",
        "title": "25‑34 лет",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, известных от 91 до 365 дней.",
        "is_dim": False,
        "name": "ym:s:upToYearSinceFirstVisitPercentage",
        "title": "Новизна 91‑365 дней",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallRevenue",
        "title": "Суммарная цена звонка",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallHoldDurationTillMissAvg",
        "title": "Среднее время ожидания до отказа",
        "type": "float",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:s:ecommercePurchases",
        "title": "Количество покупок",
        "type": "integer",
    },
    {
        "description": "Количество страниц, просмотренных посетителем во время визита.",
        "is_dim": False,
        "name": "ym:s:pageDepth",
        "title": "Глубина просмотра",
        "type": "float",
    },
    {
        "description": "Доля визитов и хитов с поддержкой JavaScript.",
        "is_dim": False,
        "name": "ym:s:jsEnabledPercentage",
        "title": "Поддержка JavaScript",
        "type": "float",
    },
    {
        "description": "Доля посетителей, предположительно являющихся роботами.",
        "is_dim": False,
        "name": "ym:s:robotPercentage",
        "title": "Роботность (доля)",
        "type": "float",
    },
    {
        "description": "Число просмотров страниц на сайте за отчетный период.",
        "is_dim": False,
        "name": "ym:s:pageviews",
        "title": "Просмотры",
        "type": "integer",
    },
    {
        "description": "Доля визитов новых посетителей.",
        "is_dim": False,
        "name": "ym:s:newUserVisitsPercentage",
        "title": "Доля визитов новых посетителей",
        "type": "float",
    },
    {
        "description": "Количество единиц товара, помещённых в корзину, суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:s:productBasketsRemoveQuantity",
        "title": "Товаров удалено из корзины",
        "type": "integer",
    },
    {
        "description": "Количество просмотров страницы с информацией о товарах суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:s:productImpressions",
        "title": "Просмотры товаров",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, возраст которых от 18 до 24 лет.",
        "is_dim": False,
        "name": "ym:s:upTo24AgePercentage",
        "title": "18‑24 лет",
        "type": "float",
    },
    {
        "description": (
            "Доля визитов посетителей с частотой визитов один раз в 2‑7 дней. "
            "Учитываются посетители, совершившие больше двух визитов. "
            "При этом число дней визитов больше 1 дня."
        ),
        "is_dim": False,
        "name": "ym:s:upToWeekBetweenVisitsPercentage",
        "title": "Периодичность 2‑7 дней",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallsUniq",
        "title": "Количество звонков с уникальных номеров",
        "type": "integer",
    },
    {
        "description": "Средняя продолжительность визита в минутах и секундах.",
        "is_dim": False,
        "name": "ym:s:avgVisitDurationSeconds",
        "title": "Время на сайте (в секундах)",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, известных меньше одного дня.",
        "is_dim": False,
        "name": "ym:s:upToDaySinceFirstVisitPercentage",
        "title": "Новизна менее одного дня",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:visitsPerHour",
        "title": "Визитов в час",
        "type": "float",
    },
    {
        "description": (
            "Отношение доли посетителей с заданным интересом на этом сайте к "
            "среднестатистической доле посетителей с этим же интересом "
            "на всех сайтах интернета. "
            'Разрешена только при использовании группировки "Категория интересов".'
        ),
        "is_dim": False,
        "name": "ym:s:affinityIndexInterests2",
        "title": "Аффинити‑индекс",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:offlineCallTalkDurationAvg",
        "title": "Среднее время разговора",
        "type": "float",
    },
    {
        "description": "Стоимость товаров, добавленных в корзину, суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:s:productBasketsPrice",
        "title": "Стоимость добавленных в корзину товаров",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей с 4‑7 визитами.",
        "is_dim": False,
        "name": "ym:s:upTo7VisitsPerUserPercentage",
        "title": "4‑7 визитов посетителя",
        "type": "float",
    },
    {
        "description": (
            "Полученный доход суммарно по всем визитам. "
            "Присылается в поле revenue при покупке товара, "
            "либо вычисляется автоматически как сумма цен всех товаров, "
            "ассоциированных с покупкой."
        ),
        "is_dim": False,
        "name": "ym:s:ecommerceRevenue",
        "title": "Доход",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:pvlAll7Window",
        "title": "PVL‑7",
        "type": "float",
    },
    {
        "description": "",
        "is_dim": False,
        "name": "ym:s:ecommerceRevenuePerVisit",
        "title": "Средний доход визита",
        "type": "float",
    },
    {
        "description": "Доля визитов с блокировщиками рекламы.",
        "is_dim": False,
        "name": "ym:s:blockedPercentage",
        "title": "Блокировка рекламы",
        "type": "float",
    },
    {
        "description": "Доля визитов и хитов, совершенных с мобильных устройств.",
        "is_dim": False,
        "name": "ym:s:mobilePercentage",
        "title": "Мобильность",
        "type": "float",
    },
    {
        "description": (
            "Среднее значение параметра визита. "
            "Вычисляется, если в качестве значения параметра визита передаётся число."
        ),
        "is_dim": False,
        "name": "ym:s:avgParams",
        "title": "Среднее параметров визитов",
        "type": "float",
    },
    {
        "description": "Количество купленных единиц товара суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:s:productPurchasedQuantity",
        "title": "Товаров куплено",
        "type": "integer",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через более чем 365 дней.",
        "is_dim": False,
        "name": "ym:s:overYearUserRecencyPercentage",
        "title": "Вернувшиеся: более чем 365 дней",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, возраст которых менее 18 лет.",
        "is_dim": False,
        "name": "ym:s:under18AgePercentage",
        "title": "Менее 18 лет",
        "type": "float",
    },
    {
        "description": "Доля визитов и хитов с включенными cookies.",
        "is_dim": False,
        "name": "ym:s:cookieEnabledPercentage",
        "title": "Поддержка сookies",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через 32‑90 дней.",
        "is_dim": False,
        "name": "ym:s:upToQuarterUserRecencyPercentage",
        "title": "Вернувшиеся: 32‑90 дней",
        "type": "float",
    },
    {
        "description": None,
        "is_dim": False,
        "name": "ym:s:visitsPerMinute",
        "title": "Визитов в минуту",
        "type": "float",
    },
    {
        "description": "Среднее время от предпоследнего визита в днях.",
        "is_dim": False,
        "name": "ym:s:userRecencyDays",
        "title": "Дней от предыдущего визита",
        "type": "float",
    },
    {
        "description": "Стоимость товаров, удалённых из корзины, суммарно по всем визитам.",
        "is_dim": False,
        "name": "ym:s:productBasketsRemovePrice",
        "title": "Стоимость удалённых из корзины товаров",
        "type": "float",
    },
    {
        "description": "Доля визитов посетителей, вернувшихся на сайт через 2‑7 дней.",
        "is_dim": False,
        "name": "ym:s:upToWeekUserRecencyPercentage",
        "title": "Вернувшиеся: 2‑7 дней",
        "type": "float",
    },
    {
        "description": "Количество новых посетителей.",
        "is_dim": False,
        "name": "ym:s:newUsers",
        "title": "Новые посетители",
        "type": "integer",
    },
    {
        "description": "Сумма значений параметров визитов. Вычисляется, если в качестве значений параметров визитов передаётся число.",
        "is_dim": False,
        "name": "ym:s:sumParams",
        "title": "Сумма параметров визитов",
        "type": "integer",
    },
    {
        "description": "Среднее число дней между визитами, интервал.",
        "is_dim": True,
        "name": "ym:s:userVisitsPeriodIntervalName",
        "title": "Периодичность",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productSum",
        "title": "Стоимость товара",
        "type": "string",
    },
    {
        "description": "Рекламный сервис",
        "is_dim": True,
        "name": "ym:s:lastOpenstatService",
        "title": "Last Openstat Service",
        "type": "string",
    },
    {
        "description": "Страница социальной сети, с которой был совершён первый за всю историю переход посетителя.",
        "is_dim": True,
        "name": "ym:s:firstSocialNetworkProfile",
        "title": "Первая страница соц. сети",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productCoupon",
        "title": "Промокод товара",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:browserHourMinute",
        "title": "Час и минута на компьютере посетителя",
        "type": "string",
    },
    {
        "description": "Количество дней с первого визита посетителя, интервал.",
        "is_dim": True,
        "name": "ym:s:daysSinceFirstVisitIntervalName",
        "title": "Дней от первого визита",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:dekaminute",
        "title": "Декаминута визита",
        "type": "string",
    },
    {
        "description": 'Глубина цвета экрана в битах. Например, "24".',
        "is_dim": True,
        "name": "ym:s:screenColorsName",
        "title": "Глубина цвета",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:advEngineName",
        "title": "Последняя рекламная система",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productID",
        "title": "ID товара",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:mobilePhoneModel",
        "title": "Модель устройства",
        "type": "string",
    },
    {
        "description": "Количество просмотров во время визита или глубина просмотра. Просмотры исключают обновления страницы в интервале менее 15 секунд.",
        "is_dim": True,
        "name": "ym:s:pageViews",
        "title": "Глубина просмотра (детально)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:firstVisitDekaminute",
        "title": "Декаминута первого визита",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала квартала.",
        "is_dim": True,
        "name": "ym:s:browserStartOfQuarter",
        "title": "Квартал на компьютере посетителя",
        "type": "date",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:purchaseRevenue",
        "title": "Доход от покупки",
        "type": "string",
    },
    {
        "description": 'Наличие метки "GCLID".',
        "is_dim": True,
        "name": "ym:s:lastSignHasGCLIDName",
        "title": "Наличие GCLID (last sign)",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
        "name": "ym:s:browserStartOfHour",
        "title": "Час на компьютере посетителя",
        "type": "datetime",
    },
    {
        "description": "Дата и время визита в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
        "name": "ym:s:startOfHour",
        "title": "Час визита (начало часа)",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:purchaseCoupon",
        "title": "Промокод покупки",
        "type": "string",
    },
    {
        "description": "Объявление",
        "is_dim": True,
        "name": "ym:s:lastUTMContent",
        "title": "Last UTM Content",
        "type": "string",
    },
    {
        "description": "Рекламная кампания",
        "is_dim": True,
        "name": "ym:s:lastOpenstatCampaign",
        "title": "Last Openstat Campaign",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:hasRecommendationsClicksName",
        "title": "Переходы по рекомендациям",
        "type": "string",
    },
    {
        "description": "Место размещения",
        "is_dim": True,
        "name": "ym:s:openstatSource",
        "title": "Openstat Source",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала недели.",
        "is_dim": True,
        "name": "ym:s:browserStartOfWeek",
        "title": "Неделя на компьютере посетителя",
        "type": "date",
    },
    {
        "description": 'Ориентация экрана. Возможные значения: "landscape", "portrait".',
        "is_dim": True,
        "name": "ym:s:screenOrientationName",
        "title": "Ориентация экрана",
        "type": "string",
    },
    {
        "description": "Логическая высота экрана.",
        "is_dim": True,
        "name": "ym:s:screenHeight",
        "title": "Логическая высота",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:javascriptEnabledName",
        "title": "Наличие Javascript",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:endURLDomain",
        "title": "Домен страницы выхода",
        "type": "string",
    },
    {
        "description": "Валюта, установленная в Яндекс.Директе для рекламной кампании.",
        "is_dim": True,
        "name": "ym:s:lastCurrencyIDName",
        "title": "Валюта (last)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:browserDayOfMonth",
        "title": "День месяца на компьютере посетителя",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:turboPageID",
        "title": "Турбо-страница (ID)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productCurrency",
        "title": "Валюта",
        "type": "string",
    },
    {
        "description": "Посетитель предположительно является роботом",
        "is_dim": True,
        "name": "ym:s:isRobotName",
        "title": "Роботность",
        "type": "string",
    },
    {
        "description": "Браузер и его полная версия.",
        "is_dim": True,
        "name": "ym:s:browserAndVersionName",
        "title": "Полная версия браузера",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:offlineCallHoldDuration",
        "title": "Время ожидания",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:browserCountry",
        "title": "Страна браузера",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:offlinePointLocationName",
        "title": "Точка",
        "type": "string",
    },
    {
        "description": "Тип условия показа объявления.",
        "is_dim": True,
        "name": "ym:s:lastDirectConditionTypeName",
        "title": "Тип условия показа объявления (last)",
        "type": "string",
    },
    {
        "description": "Количество визитов посетителя за всю историю, интервал.",
        "is_dim": True,
        "name": "ym:s:userVisitsInterval",
        "title": "Визитов в истории",
        "type": "string",
    },
    {
        "description": "Области, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:s:regionAreaName",
        "title": "Область",
        "type": "string",
    },
    {
        "description": "Название проводимой рекламной кампании",
        "is_dim": True,
        "name": "ym:s:firstUTMCampaign",
        "title": "First UTM Campaign",
        "type": "string",
    },
    {
        "description": "Социальная сеть, из которой был совершён последний значимый переход посетителя.",
        "is_dim": True,
        "name": "ym:s:lastSignSocialNetworkName",
        "title": "Последняя значимая социальная сеть",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала года.",
        "is_dim": True,
        "name": "ym:s:browserStartOfYear",
        "title": "Год на компьютере посетителя (начало года)",
        "type": "date",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:firstVisitMinute",
        "title": "Минута первого визита",
        "type": "string",
    },
    {
        "description": (
            'Часть URL реферера, следующая сразу после домена до "?" (секция параметров), '
            '"#" (секция фрагмента) либо до конца строки. '
            'Например, для страницы "http://news.yandex.ru/quotes/1.html" путем является "/quotes/1.html", '
            'для "http://news.yandex.ru/"&#160;— "/", а для "http://news.yandex.ru"&#160;— пустая строка.'
        ),
        "is_dim": True,
        "name": "ym:s:refererPath",
        "title": "Путь реферера",
        "type": "string",
    },
    {
        "description": "Количество дней с предпоследнего визита посетителя, интервал.",
        "is_dim": True,
        "name": "ym:s:daysSincePreviousVisitIntervalName",
        "title": "Время возврата",
        "type": "string",
    },
    {
        "description": "Идентификатор посетителя сайта, назначаемый Метрикой",
        "is_dim": True,
        "name": "ym:s:clientID",
        "title": "ClientID",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:operatingSystemRootName",
        "title": "Группа операционных систем",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:from",
        "title": "Метка from",
        "type": "string",
    },
    {
        "description": "Дата и время визита в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
        "name": "ym:s:startOfDekaminute",
        "title": "Начало декаминуты визита",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productBrandCart",
        "title": "Бренд товара в корзине",
        "type": "string",
    },
    {
        "description": "Дата и время первого визита посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
        "name": "ym:s:firstVisitStartOfDekaminute",
        "title": "Начало декаминуты первого визита",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:offlinePointRegionName",
        "title": "Регион",
        "type": "string",
    },
    # {
    #     'description': 'Интервал дат с параметризацией размера интервала',
    #    #     'is_dim': True,
    #     'name': 'ym:s:previousVisitDatePeriod<group>',
    #     'title': 'Дата предыдущего визита (интервал)',
    #     'type': 'date',
    # },
    {
        "description": "Идентификатор уточненной поисковой системы.",
        "is_dim": True,
        "name": "ym:s:searchEngineName",
        "title": "Последняя поисковая система (детально)",
        "type": "string",
    },
    {
        "description": "Рекламное объявление",
        "is_dim": True,
        "name": "ym:s:firstOpenstatAd",
        "title": "First Openstat Ad",
        "type": "string",
    },
    # {
    #     'description': None,
    #    #     'is_dim': True,
    #     'name': 'ym:s:browserDatePeriod<group>',
    #     'title': 'Дата на компьютере посетителя',
    #     'type': 'string',
    # },
    {
        "description": "Дата визита в формате YYYY-MM-DD, округленная до начала квартала.",
        "is_dim": True,
        "name": "ym:s:startOfQuarter",
        "title": "Квартал визита",
        "type": "date",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:browserYear",
        "title": "Год на компьютере посетителя",
        "type": "string",
    },
    {
        "description": "Количество дней с предпоследнего визита посетителя.",
        "is_dim": True,
        "name": "ym:s:daysSincePreviousVisit",
        "title": "Время c предыдущего визита",
        "type": "string",
    },
    {
        "description": "Место размещения",
        "is_dim": True,
        "name": "ym:s:lastSignOpenstatSource",
        "title": "Last Sign Openstat Source",
        "type": "string",
    },
    {
        "description": "Источник перехода",
        "is_dim": True,
        "name": "ym:s:UTMSource",
        "title": "UTM Source",
        "type": "string",
    },
    {
        "description": "Дата первого визита посетителя в формате YYYY-MM-DD, округленная до начала месяца.",
        "is_dim": True,
        "name": "ym:s:firstVisitStartOfMonth",
        "title": "Месяц первого визита (начало месяца)",
        "type": "date",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:purchaseExistsVisit",
        "title": "Наличие покупок в визите",
        "type": "string",
    },
    {
        "description": "Дата первого визита посетителя в формате YYYY-MM-DD, округленная до начала недели (понедельник).",
        "is_dim": True,
        "name": "ym:s:firstVisitStartOfWeek",
        "title": "Неделя первого визита",
        "type": "date",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:externalRefererDomain",
        "title": "Домен внешнего реферера",
        "type": "string",
    },
    {
        "description": "Категория коммерческих интересов посетителей.",
        "is_dim": True,
        "name": "ym:s:interestName",
        "title": "Категория интересов",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:purchaseCountVisit",
        "title": "Количество покупок в визите",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:firstAdvEngineName",
        "title": "Первая рекламная система",
        "type": "string",
    },
    {
        "description": "Количество просмотров в визите или глубина просмотра, сгруппированная по интервалам.",
        "is_dim": True,
        "name": "ym:s:pageViewsInterval",
        "title": "Глубина просмотра (по интервалам)",
        "type": "string",
    },
    {
        "description": "Ключевые слова",
        "is_dim": True,
        "name": "ym:s:firstUTMTerm",
        "title": "First UTM Term",
        "type": "string",
    },
    # {
    #     'description': 'Интервал дат с параметризацией размера интервала',
    #    #     'is_dim': True,
    #     'name': 'ym:s:firstVisitDatePeriod<group>',
    #     'title': 'Дата первого визита (интервал)',
    #     'type': 'string',
    # },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productVariant",
        "title": "Вариант товара",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:startURLProto",
        "title": "Протокол страницы входа",
        "type": "string",
    },
    {
        "description": 'Идентификатор уточненной поисковой системы первого визита посетителя. Например, "Яндекс.Картинки".',
        "is_dim": True,
        "name": "ym:s:firstSearchEngineName",
        "title": "Первая поисковая система (детально)",
        "type": "string",
    },
    {
        "description": 'Размер клиентской части окна браузера. Например, "1920х1080".',
        "is_dim": True,
        "name": "ym:s:windowClientArea",
        "title": "Размер окна",
        "type": "string",
    },
    {
        "description": "Страны, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:s:regionCountryName",
        "title": "Страна",
        "type": "string",
    },
    {
        "description": "Дата и время предпоследнего визита посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала квартала.",
        "is_dim": True,
        "name": "ym:s:previousVisitStartOfQuarter",
        "title": "Квартал предпоследнего визита",
        "type": "date",
    },
    {
        "description": "Валюта, установленная в Яндекс.Директе для рекламной кампании.",
        "is_dim": True,
        "name": "ym:s:firstCurrencyIDName",
        "title": "Валюта (first)",
        "type": "string",
    },
    {
        "description": "Время ― начало часового периода в формате HH:MM.",
        "is_dim": True,
        "name": "ym:s:hourName",
        "title": "Час визита",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала месяца.",
        "is_dim": True,
        "name": "ym:s:browserStartOfMonth",
        "title": "Месяц на компьютере посетителя (начало месяца)",
        "type": "date",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:previousVisitMonth",
        "title": "Месяц предыдущего визита",
        "type": "string",
    },
    {
        "description": "Возраст посетителя, интервал",
        "is_dim": True,
        "name": "ym:s:ageIntervalName",
        "title": "Возраст",
        "type": "string",
    },
    {
        "description": 'Идентификатор уточненной поисковой системы первого визита посетителя. Например, "Яндекс.Картинки".',
        "is_dim": True,
        "name": "ym:s:lastSignSearchEngineName",
        "title": "Последняя значимая поисковая система (детально)",
        "type": "string",
    },
    {
        "description": "Ширина клиентской части окна браузера.",
        "is_dim": True,
        "name": "ym:s:windowClientWidth",
        "title": "Ширина окна",
        "type": "string",
    },
    {
        "description": "Название браузера и его версия.",
        "is_dim": True,
        "name": "ym:s:browserAndVersionMajorName",
        "title": "Версия браузера",
        "type": "string",
    },
    {
        "description": 'Отношение ширины к высоте экрана. Например, "16:9".',
        "is_dim": True,
        "name": "ym:s:screenFormat",
        "title": "Соотношение сторон",
        "type": "string",
    },
    {
        "description": "Часть URL страницы входа, следующая сразу после домена. В отличие от простого пути также содержит секции параметров и фрагмента",
        "is_dim": True,
        "name": "ym:s:startURLPathFull",
        "title": "Путь (полный) страницы входа",
        "type": "string",
    },
    {
        "description": 'Идентификатор верхнего уровня иерархии поисковых систем первого визита посетителя. Например, "Яндекс".',
        "is_dim": True,
        "name": "ym:s:firstSearchEngineRootName",
        "title": "Первая поисковая система",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:browserLanguage",
        "title": "Язык браузера",
        "type": "string",
    },
    {
        "description": "Ключевые слова",
        "is_dim": True,
        "name": "ym:s:UTMTerm",
        "title": "UTM Term",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:dayOfWeekName",
        "title": "День недели визита",
        "type": "string",
    },
    {
        "description": 'Четвертый уровень URL страницы входа. Например, для "http://ya.ru/1/2/3/4.html" четвертым уровнем URL является "http://ya.ru/1/2/3/".',
        "is_dim": True,
        "name": "ym:s:startURLPathLevel4",
        "title": "Страница входа, ур. 4",
        "type": "string",
    },
    {
        "description": 'Третий уровень URL страницы входа. Например, для "http://ya.ru/1/2/3/4.html" третьим уровнем URL является "http://ya.ru/1/2/".',
        "is_dim": True,
        "name": "ym:s:startURLPathLevel3",
        "title": "Страница входа, ур. 3",
        "type": "string",
    },
    {
        "description": "Средство маркетинга",
        "is_dim": True,
        "name": "ym:s:UTMMedium",
        "title": "UTM Medium",
        "type": "string",
    },
    {
        "description": 'Второй уровень URL страницы входа. Например, для "http://ya.ru/1/2/3/4.html" вторым уровнем URL является "http://ya.ru/1/".',
        "is_dim": True,
        "name": "ym:s:startURLPathLevel2",
        "title": "Страница входа, ур. 2",
        "type": "string",
    },
    {
        "description": 'Первый уровень URL страницы входа. Например, для "http://ya.ru/1/2/3/4.html" первым уровнем URL является "http://ya.ru/".',
        "is_dim": True,
        "name": "ym:s:startURLPathLevel1",
        "title": "Страница входа, ур. 1",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:offlineCallTag",
        "title": "Метка",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:refererProto",
        "title": "Протокол реферера",
        "type": "string",
    },
    {
        "description": "Объявление",
        "is_dim": True,
        "name": "ym:s:firstUTMContent",
        "title": "First UTM Content",
        "type": "string",
    },
    {
        "description": "Дата предыдущего визита посетителя в формате YYYY-MM-DD.",
        "is_dim": True,
        "name": "ym:s:previousVisitDate",
        "title": "Дата предыдущего визита",
        "type": "date",
    },
    {
        "description": 'Пятый уровень URL страницы входа. Например, для "http://ya.ru/1/2/3/4.html" пятым уровнем URL является "http://ya.ru/1/2/3/4.html".',
        "is_dim": True,
        "name": "ym:s:startURLPathLevel5",
        "title": "Страница входа, ур. 5",
        "type": "string",
    },
    {
        "description": (
            'Часть URL страницы входа, следующая сразу после домена до "?" (секция параметров), '
            '"#" (секция фрагмента) либо до конца строки. '
            'Например, для страницы "http://news.yandex.ru/quotes/1.html" путем является "/quotes/1.html", '
            'для "http://news.yandex.ru/" - "/", а для "http://news.yandex.ru" - пустая строка'
        ),
        "is_dim": True,
        "name": "ym:s:startURLPath",
        "title": "Путь страницы входа",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:PProductBrand",
        "title": "Бренд товара в заказе",
        "type": "string",
    },
    {
        "description": "Дата и время первого визита посетителя в формате YYYY-MM-DD HH:mm:ss.",
        "is_dim": True,
        "name": "ym:s:firstVisitDateTime",
        "title": "Дата и время первого визита",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:firstTrafficSourceName",
        "title": "Первый источник трафика",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:trafficSourceName",
        "title": "Последний источник трафика",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:lastSignTrafficSourceName",
        "title": "Последний значимый источник трафика",
        "type": "string",
    },
    {
        "description": "Второй уровень источников трафика",
        "is_dim": True,
        "name": "ym:s:firstSourceEngineName",
        "title": "Первый источник трафика (детально)",
        "type": "string",
    },
    {
        "description": "Второй уровень источников трафика",
        "is_dim": True,
        "name": "ym:s:sourceEngineName",
        "title": "Последний источник трафика (детально)",
        "type": "string",
    },
    {
        "description": "Второй уровень источников трафика",
        "is_dim": True,
        "name": "ym:s:lastSignSourceEngineName",
        "title": "Последний значимый источник трафика (детально)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:endURL",
        "title": "Страница выхода",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:hourMinute",
        "title": "Час и минута визита",
        "type": "string",
    },
    {
        "description": 'Наличие метки "GCLID".',
        "is_dim": True,
        "name": "ym:s:lastHasGCLIDName",
        "title": "Наличие GCLID (last)",
        "type": "string",
    },
    {
        "description": "Объявление",
        "is_dim": True,
        "name": "ym:s:UTMContent",
        "title": "UTM Content",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productBrand",
        "title": "Бренд",
        "type": "string",
    },
    {
        "description": "Время на сайте, разбитое по группам.",
        "is_dim": True,
        "name": "ym:s:visitDurationIntervalName",
        "title": "Время на сайте (по группам)",
        "type": "string",
    },
    {
        "description": 'Четвертый уровень URL реферера. Например, для "http://ya.ru/1/2/3/4.html" четвертым уровнем URL является "http://ya.ru/1/2/3/".',
        "is_dim": True,
        "name": "ym:s:refererPathLevel4",
        "title": "Реферер, ур. 4",
        "type": "string",
    },
    {
        "description": 'Третий уровень URL реферера. Например, для "http://ya.ru/1/2/3/4.html" третьим уровнем URL является "http://ya.ru/1/2/".',
        "is_dim": True,
        "name": "ym:s:refererPathLevel3",
        "title": "Реферер, ур. 3",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:cookieEnabledName",
        "title": "Наличие Cookie",
        "type": "string",
    },
    {
        "description": 'Разрешение экрана. Например, "1920х1080".',
        "is_dim": True,
        "name": "ym:s:screenResolution",
        "title": "Разрешение",
        "type": "string",
    },
    {
        "description": "Название проводимой рекламной кампании",
        "is_dim": True,
        "name": "ym:s:lastUTMCampaign",
        "title": "Last UTM Campaign",
        "type": "string",
    },
    {
        "description": "Рекламная кампания",
        "is_dim": True,
        "name": "ym:s:firstOpenstatCampaign",
        "title": "First Openstat Campaign",
        "type": "string",
    },
    {
        "description": 'Пятый уровень URL реферера. Например, для "http://ya.ru/1/2/3/4.html" пятым уровнем URL является "http://ya.ru/1/2/3/4.html".',
        "is_dim": True,
        "name": "ym:s:refererPathLevel5",
        "title": "Реферер, ур. 5",
        "type": "string",
    },
    {
        "description": "С каких сайтов осуществлялись переходы на сайт.",
        "is_dim": True,
        "name": "ym:s:referalSource",
        "title": "Последний переход с сайтов",
        "type": "string",
    },
    {
        "description": 'Третий уровень URL внешнего реферера. Например, для "http://ya.ru/1/2/3/4.html" третьим уровнем URL является "http://ya.ru/1/2/".',
        "is_dim": True,
        "name": "ym:s:externalRefererPathLevel3",
        "title": "Внешний реферер, ур. 3",
        "type": "string",
    },
    {
        "description": 'Второй уровень URL внешнего реферера. Например, для "http://ya.ru/1/2/3/4.html" вторым уровнем URL является "http://ya.ru/1/".',
        "is_dim": True,
        "name": "ym:s:externalRefererPathLevel2",
        "title": "Внешний реферер, ур. 2",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:offlineCallTalkDuration",
        "title": "Время разговора",
        "type": "string",
    },
    {
        "description": 'Второй уровень URL реферера. Например, для "http://ya.ru/1/2/3/4.html" вторым уровнем URL является "http://ya.ru/1/".',
        "is_dim": True,
        "name": "ym:s:refererPathLevel2",
        "title": "Реферер, ур. 2",
        "type": "string",
    },
    {
        "description": 'Пятый уровень URL внешнего реферера. Например, для "http://ya.ru/1/2/3/4.html" пятым уровнем URL является "http://ya.ru/1/2/3/4.html".',
        "is_dim": True,
        "name": "ym:s:externalRefererPathLevel5",
        "title": "Внешний реферер, ур. 5",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:firstVisitDayOfMonth",
        "title": "День месяца первого визита",
        "type": "string",
    },
    {
        "description": 'Первый уровень URL реферера. Например, для "http://ya.ru/1/2/3/4.html" первым уровнем URL является "http://ya.ru/".',
        "is_dim": True,
        "name": "ym:s:refererPathLevel1",
        "title": "Реферер, ур. 1",
        "type": "string",
    },
    {
        "description": 'Четвертый уровень URL внешнего реферера. Например, для "http://ya.ru/1/2/3/4.html" четвертым уровнем URL является "http://ya.ru/1/2/3/".',
        "is_dim": True,
        "name": "ym:s:externalRefererPathLevel4",
        "title": "Внешний реферер, ур. 4",
        "type": "string",
    },
    {
        "description": 'Первый уровень URL внешнего реферера. Например, для "http://ya.ru/1/2/3/4.html" первым уровнем URL является "http://ya.ru/".',
        "is_dim": True,
        "name": "ym:s:externalRefererPathLevel1",
        "title": "Внешний реферер, ур. 1",
        "type": "string",
    },
    {
        "description": "С каких сайтов осуществлялись переходы на сайт.",
        "is_dim": True,
        "name": "ym:s:lastSignReferalSource",
        "title": "Последний значимый переход с сайтов",
        "type": "string",
    },
    {
        "description": "Тип условия показа объявления.",
        "is_dim": True,
        "name": "ym:s:lastSignDirectConditionTypeName",
        "title": "Тип условия показа объявления (last sign)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:previousVisitDayOfMonth",
        "title": "День месяца предыдущего визита",
        "type": "string",
    },
    {
        "description": "Дата и время первого визита посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала часа.",
        "is_dim": True,
        "name": "ym:s:firstVisitStartOfHour",
        "title": "Час первого визита (начало часа)",
        "type": "datetime",
    },
    {
        "description": "Логическая ширина экрана.",
        "is_dim": True,
        "name": "ym:s:screenWidth",
        "title": "Логическая ширина",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:year",
        "title": "Год визита",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:lastSignAdvEngineName",
        "title": "Последняя значимая рекламная система",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:startURLDomain",
        "title": "Домен страницы входа",
        "type": "string",
    },
    {
        "description": "Страница социальной сети.",
        "is_dim": True,
        "name": "ym:s:socialNetworkProfile",
        "title": "Последняя страница соц. сети",
        "type": "string",
    },
    {
        "description": "Рекламное объявление",
        "is_dim": True,
        "name": "ym:s:lastOpenstatAd",
        "title": "Last Openstat Ad",
        "type": "string",
    },
    {
        "description": "Страница социальной сети последнего значимого перехода посетителя за всю историю.",
        "is_dim": True,
        "name": "ym:s:lastSignSocialNetworkProfile",
        "title": "Страница социальной сети, с которой был совершён последний значимый переход.",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:paramsLevel10",
        "title": "Параметр визита, ур. 10",
        "type": "string",
    },
    {
        "description": "Среднее число дней между визитами.",
        "is_dim": True,
        "name": "ym:s:userVisitsPeriod",
        "title": "Периодичность (детально)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:impressionCountVisit",
        "title": "Количество просмотров товаров в визите",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:paramsLevel3",
        "title": "Параметр визита, ур. 3",
        "type": "string",
    },
    {
        "description": "Часть URL реферера, следующая сразу после домена. В отличие от простого пути также содержит секции параметров и фрагмента.",
        "is_dim": True,
        "name": "ym:s:refererPathFull",
        "title": "Путь (полный) реферера",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:paramsLevel4",
        "title": "Параметр визита, ур. 4",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:paramsLevel1",
        "title": "Параметр визита, ур. 1",
        "type": "string",
    },
    {
        "description": 'Физическая ширина экрана исходя из данных "device pixel ratio".',
        "is_dim": True,
        "name": "ym:s:physicalScreenWidth",
        "title": "Физическая ширина",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:paramsLevel2",
        "title": "Параметр визита, ур. 2",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:paramsLevel7",
        "title": "Параметр визита, ур. 7",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:paramsLevel8",
        "title": "Параметр визита, ур. 8",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:paramsLevel5",
        "title": "Параметр визита, ур. 5",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:paramsLevel6",
        "title": "Параметр визита, ур. 6",
        "type": "string",
    },
    {
        "description": "Средство маркетинга",
        "is_dim": True,
        "name": "ym:s:firstUTMMedium",
        "title": "First UTM Medium",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:PProductName",
        "title": "Название товара в заказе",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:paramsLevel9",
        "title": "Параметр визита, ур. 9",
        "type": "string",
    },
    {
        "description": "Средство маркетинга",
        "is_dim": True,
        "name": "ym:s:lastUTMMedium",
        "title": "Last UTM Medium",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:browserMinute",
        "title": "Минута на компьютере посетителя",
        "type": "string",
    },
    {
        "description": 'Наличие метки "GCLID".',
        "is_dim": True,
        "name": "ym:s:hasGCLIDName",
        "title": "Наличие GCLID",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:month",
        "title": "Месяц визита",
        "type": "string",
    },
    {
        "description": "URL реферера",
        "is_dim": True,
        "name": "ym:s:referer",
        "title": "Реферер",
        "type": "string",
    },
    {
        "description": "Часть URL страницы выхода, следующая сразу после домена. В отличие от простого пути также содержит секции параметров и фрагмента",
        "is_dim": True,
        "name": "ym:s:endURLPathFull",
        "title": "Путь (полный) страницы выхода",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:previousVisitYear",
        "title": "Год предыдущего визита",
        "type": "string",
    },
    {
        "description": 'Физическая высота экрана исходя из данных "device pixel ratio".',
        "is_dim": True,
        "name": "ym:s:physicalScreenHeight",
        "title": "Физическая высота",
        "type": "string",
    },
    {
        "description": "Место размещения",
        "is_dim": True,
        "name": "ym:s:firstOpenstatSource",
        "title": "First Openstat Source",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:hasAdBlockerName",
        "title": "Наличие блокировщиков рекламы",
        "type": "string",
    },
    {
        "description": 'Первый уровень URL страницы выхода. Например, для "http://ya.ru/1/2/3/4.html" первым уровнем URL является "http://ya.ru/".',
        "is_dim": True,
        "name": "ym:s:endURLPathLevel1",
        "title": "Страница выхода, ур. 1",
        "type": "string",
    },
    {
        "description": 'Второй уровень URL страницы выхода. Например, для "http://ya.ru/1/2/3/4.html" вторым уровнем URL является "http://ya.ru/1/".',
        "is_dim": True,
        "name": "ym:s:endURLPathLevel2",
        "title": "Страница выхода, ур. 2",
        "type": "string",
    },
    {
        "description": 'Третий уровень URL страницы выхода. Например, для "http://ya.ru/1/2/3/4.html" третьим уровнем URL является "http://ya.ru/1/2/".',
        "is_dim": True,
        "name": "ym:s:endURLPathLevel3",
        "title": "Страница выхода, ур. 3",
        "type": "string",
    },
    {
        "description": 'Четвертый уровень URL страницы выхода. Например, для "http://ya.ru/1/2/3/4.html" четвертым уровнем URL является "http://ya.ru/1/2/3/".',
        "is_dim": True,
        "name": "ym:s:endURLPathLevel4",
        "title": "Страница выхода, ур. 4",
        "type": "string",
    },
    {
        "description": 'Пятый уровень URL страницы выхода. Например, для "http://ya.ru/1/2/3/4.html" пятым уровнем URL является "http://ya.ru/1/2/3/4.html".',
        "is_dim": True,
        "name": "ym:s:endURLPathLevel5",
        "title": "Страница выхода, ур. 5",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:browserDayOfWeekName",
        "title": "День недели на компьютере посетителя",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:browserHourName",
        "title": "Час на компьютере посетителя в формате HH:00",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:topLevelDomain",
        "title": "Домен верхнего уровня страницы входа",
        "type": "string",
    },
    {
        "description": "Дата первого визита посетителя в формате YYYY-MM-DD, округленная до начала квартала.",
        "is_dim": True,
        "name": "ym:s:firstVisitStartOfQuarter",
        "title": "Квартал первого визита",
        "type": "date",
    },
    {
        "description": "Округа, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:s:regionDistrictName",
        "title": "Округ",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:purchaseRevenueVisit",
        "title": "Сумма покупок в визите",
        "type": "string",
    },
    {
        "description": "Дата и время первого визита посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.",
        "is_dim": True,
        "name": "ym:s:firstVisitStartOfMinute",
        "title": "Минута первого визита (начало минуты)",
        "type": "datetime",
    },
    {
        "description": "Дата визита в формате YYYY-MM-DD, округленная до начала месяца.",
        "is_dim": True,
        "name": "ym:s:startOfMonth",
        "title": "Месяц визита (начало месяца)",
        "type": "date",
    },
    {
        "description": "Города, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:s:regionCityName",
        "title": "Город",
        "type": "string",
    },
    {
        "description": "Источник трафика - идентификатор социальной сети.",
        "is_dim": True,
        "name": "ym:s:socialNetworkName",
        "title": "Последняя социальная сеть",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:endURLProto",
        "title": "Протокол страницы выхода",
        "type": "string",
    },
    {
        "description": "Источник перехода",
        "is_dim": True,
        "name": "ym:s:firstUTMSource",
        "title": "First UTM Source",
        "type": "string",
    },
    {
        "description": "Рекламный сервис",
        "is_dim": True,
        "name": "ym:s:openstatService",
        "title": "Openstat Service",
        "type": "string",
    },
    {
        "description": "Третье число (обычно номер сборки) в версии движка браузера посетителя.",
        "is_dim": True,
        "name": "ym:s:browserEngineVersion3",
        "title": "Build-версия движка браузера",
        "type": "string",
    },
    {
        "description": 'URL внешнего реферера для источника "Переходы по ссылкам на сайтах"',
        "is_dim": True,
        "name": "ym:s:externalReferer",
        "title": "Внешний реферер",
        "type": "string",
    },
    {
        "description": "Четвертое число (обычно номер ревизии) в версии движка браузера посетителя.",
        "is_dim": True,
        "name": "ym:s:browserEngineVersion4",
        "title": "Revision-версия движка браузера",
        "type": "string",
    },
    {
        "description": "Дата и время визита в формате YYYY-MM-DD HH:mm:ss.",
        "is_dim": True,
        "name": "ym:s:dateTime",
        "title": "Дата и время визита",
        "type": "datetime",
    },
    {
        "description": (
            "Часть URL-адреса внешнего реферера, следующая сразу после домена. "
            "В отличие от простого пути также содержит секции параметров и фрагмента."
        ),
        "is_dim": True,
        "name": "ym:s:externalRefererPathFull",
        "title": "Путь (полный) внешнего реферера",
        "type": "string",
    },
    {
        "description": "Дата и время визита в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.",
        "is_dim": True,
        "name": "ym:s:startOfMinute",
        "title": "Минута визита (начало минуты)",
        "type": "datetime",
    },
    {
        "description": (
            'Часть URL внешнего реферера, следующая сразу после домена до "?" (секция параметров), '
            '"#" (секция фрагмента) либо до конца строки. '
            'Например, для страницы "http://news.yandex.ru/quotes/1.html" '
            'путем является "/quotes/1.html", для "http://news.yandex.ru/"&#160;— "/", '
            'а для "http://news.yandex.ru"&#160;— пустая строка.'
        ),
        "is_dim": True,
        "name": "ym:s:externalRefererPath",
        "title": "Путь внешнего реферера",
        "type": "string",
    },
    {
        "description": "Дата первого визита посетителя в формате YYYY-MM-DD.",
        "is_dim": True,
        "name": "ym:s:firstVisitDate",
        "title": "Дата первого визита",
        "type": "date",
    },
    {
        "description": "Тип условия показа объявления.",
        "is_dim": True,
        "name": "ym:s:firstDirectConditionTypeName",
        "title": "Тип условия показа объявления (first)",
        "type": "string",
    },
    {
        "description": "Дата первого визита посетителя в формате YYYY-MM-DD, округленная до начала года.",
        "is_dim": True,
        "name": "ym:s:firstVisitStartOfYear",
        "title": "Год первого визита (начало года)",
        "type": "date",
    },
    {
        "description": "Первое число в версии движка браузера посетителя.",
        "is_dim": True,
        "name": "ym:s:browserEngineVersion1",
        "title": "Major-версия движка браузера",
        "type": "string",
    },
    {
        "description": "Второе число в версии движка браузера посетителя.",
        "is_dim": True,
        "name": "ym:s:browserEngineVersion2",
        "title": "Minor-версия движка браузера",
        "type": "string",
    },
    {
        "description": "Рекламный сервис",
        "is_dim": True,
        "name": "ym:s:firstOpenstatService",
        "title": "First Openstat Service",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productName",
        "title": "Название товара",
        "type": "string",
    },
    {
        "description": "Поисковая фраза последнего перехода по объявлению Яндекс.Маркета.",
        "is_dim": True,
        "name": "ym:s:marketSearchPhrase",
        "title": "Фраза (Маркет)",
        "type": "string",
    },
    {
        "description": "Континенты, в которых находятся посетители сайта.",
        "is_dim": True,
        "name": "ym:s:regionContinentName",
        "title": "Континент",
        "type": "string",
    },
    {
        "description": 'Реальное разрешение экрана исходя из данных "device pixel ratio". Например, "1920х1080".',
        "is_dim": True,
        "name": "ym:s:physicalScreenResolution",
        "title": "Реальное разрешение",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:lastFrom",
        "title": "Метка from (last)",
        "type": "string",
    },
    {
        "description": "Категория коммерческих интересов посетителей.",
        "is_dim": True,
        "name": "ym:s:interest2Name1",
        "title": "Категория интересов, ур. 1",
        "type": "string",
    },
    {
        "description": "Категория коммерческих интересов посетителей.",
        "is_dim": True,
        "name": "ym:s:interest2Name2",
        "title": "Категория интересов, ур. 2",
        "type": "string",
    },
    {
        "description": "Категория коммерческих интересов посетителей.",
        "is_dim": True,
        "name": "ym:s:interest2Name3",
        "title": "Категория интересов, ур. 3",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:browserDekaminute",
        "title": "Декаминута на компьютере посетителя",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productPrice",
        "title": "Стоимость единицы товара",
        "type": "string",
    },
    {
        "description": "Место размещения",
        "is_dim": True,
        "name": "ym:s:lastOpenstatSource",
        "title": "Last Openstat Source",
        "type": "string",
    },
    {
        "description": "Ключевые слова",
        "is_dim": True,
        "name": "ym:s:lastSignUTMTerm",
        "title": "Last Sign UTM Term",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:offlineCallMissedName",
        "title": "Пропущенный / Отвеченный",
        "type": "string",
    },
    {
        "description": "Социальная сеть первого перехода посетителя.",
        "is_dim": True,
        "name": "ym:s:firstSocialNetworkName",
        "title": "Первая социальная сеть",
        "type": "string",
    },
    {
        "description": "Рекламное объявление",
        "is_dim": True,
        "name": "ym:s:lastSignOpenstatAd",
        "title": "Last Sing Openstat Ad",
        "type": "string",
    },
    {
        "description": 'IPv4/IPv6-адрес посетителя с обнуленными 1/8 последними байтами соответственно. Например, "77.88.21.0"/"2001:db8:85a3::".',
        "is_dim": True,
        "name": "ym:s:ipAddress",
        "title": "IP-адрес",
        "type": "string",
    },
    {
        "description": "Дата визита в формате YYYY-MM-DD.",
        "is_dim": True,
        "name": "ym:s:date",
        "title": "Дата визита",
        "type": "date",
    },
    {
        "description": "Рекламная кампания",
        "is_dim": True,
        "name": "ym:s:lastSignOpenstatCampaign",
        "title": "Last Sign Openstat Campaign",
        "type": "string",
    },
    {
        "description": 'Браузер посетителя. Например, "Яндекс.Браузер".',
        "is_dim": True,
        "name": "ym:s:browserName",
        "title": "Браузер",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:minute",
        "title": "Минута визита",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss.",
        "is_dim": True,
        "name": "ym:s:browserDateTime",
        "title": "Дата и время на компьютере посетителя",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:firstVisitMonth",
        "title": "Месяц первого визита",
        "type": "string",
    },
    {
        "description": "Название проводимой рекламной кампании",
        "is_dim": True,
        "name": "ym:s:UTMCampaign",
        "title": "UTM Campaign",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:clientTimeZone",
        "title": "Часовой пояс на компьютере посетителя",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:externalRefererProto",
        "title": "Протокол внешнего реферера",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productPosition",
        "title": "Позиция товара",
        "type": "string",
    },
    {
        "description": "Дата визита в формате YYYY-MM-DD, округленная до начала года.",
        "is_dim": True,
        "name": "ym:s:startOfYear",
        "title": "Год визита (начало года)",
        "type": "date",
    },
    {
        "description": "Дата и время предпоследнего визита посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала года.",
        "is_dim": True,
        "name": "ym:s:previousVisitStartOfYear",
        "title": "Год предпоследнего визита",
        "type": "date",
    },
    {
        "description": 'Движок браузера посетителя. Например, "WebKit".',
        "is_dim": True,
        "name": "ym:s:browserEngine",
        "title": "Движок браузера",
        "type": "string",
    },
    {
        "description": (
            'Часть URL страницы выхода, следующая сразу после домена до "?" (секция параметров), '
            '"#" (секция фрагмента) либо до конца строки. '
            'Например, для страницы "http://news.yandex.ru/quotes/1.html" '
            'путем является "/quotes/1.html", для "http://news.yandex.ru/" - "/", '
            'а для "http://news.yandex.ru" - пустая строка'
        ),
        "is_dim": True,
        "name": "ym:s:endURLPath",
        "title": "Путь страницы выхода",
        "type": "string",
    },
    {
        "description": "Количество визитов посетителя за всю историю.",
        "is_dim": True,
        "name": "ym:s:userVisits",
        "title": "Всего визитов",
        "type": "string",
    },
    {
        "description": "Средство маркетинга",
        "is_dim": True,
        "name": "ym:s:lastSignUTMMedium",
        "title": "Last Sign UTM Medium",
        "type": "string",
    },
    {
        "description": "Время ― начало часового периода в формате HH:MM.",
        "is_dim": True,
        "name": "ym:s:firstVisitHourName",
        "title": "Час первого визита",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:firstVisitHourMinute",
        "title": "Час и минута первого визита",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:offlineCallFirstTimeCallerName",
        "title": "Первичный / Повторный",
        "type": "string",
    },
    {
        "description": "Название проводимой рекламной кампании",
        "is_dim": True,
        "name": "ym:s:lastSignUTMCampaign",
        "title": "Last Sign UTM Campaign",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала минуты.",
        "is_dim": True,
        "name": "ym:s:browserStartOfMinute",
        "title": "Минута на компьютере посетителя (начало минуты)",
        "type": "datetime",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:firstVisitDayOfWeekName",
        "title": "День недели первого визита посетителя",
        "type": "string",
    },
    {
        "description": "Рекламное объявление",
        "is_dim": True,
        "name": "ym:s:openstatAd",
        "title": "Openstat Ad",
        "type": "string",
    },
    {
        "description": 'Тип устройства, с которого было посещение. Возможные значения: "desktop", "mobile", "tablet", "tv".',
        "is_dim": True,
        "name": "ym:s:deviceCategoryName",
        "title": "Тип устройства",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:startURL",
        "title": "Страница входа",
        "type": "string",
    },
    {
        "description": (
            "Был ли данный визит отказом с учетом точного показателя отказов. "
            "Если страница просматривалась больше 15 секунд, "
            'то визит не считается отказом. Возможные значения: "Yes", "No".'
        ),
        "is_dim": True,
        "name": "ym:s:bounceName",
        "title": "Отказность",
        "type": "string",
    },
    {
        "description": "Количество дней с первого визита посетителя.",
        "is_dim": True,
        "name": "ym:s:daysSinceFirstVisit",
        "title": "Дней от первого визита (детально)",
        "type": "string",
    },
    {
        "description": "Ключевые слова",
        "is_dim": True,
        "name": "ym:s:lastUTMTerm",
        "title": "Last UTM Term",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:operatingSystemName",
        "title": "Операционная система (детально)",
        "type": "string",
    },
    {
        "description": "Продолжительность визита в секундах.",
        "is_dim": True,
        "name": "ym:s:visitDuration",
        "title": "Время на сайте (детально)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productNameCart",
        "title": "Название товара в корзине",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:isTurboPageName",
        "title": "Турбо-страница",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:dayOfMonth",
        "title": "День месяца визита",
        "type": "string",
    },
    {
        "description": "Валюта, установленная в Яндекс.Директе для рекламной кампании.",
        "is_dim": True,
        "name": "ym:s:lastSignCurrencyIDName",
        "title": "Валюта (last sign)",
        "type": "string",
    },
    {
        "description": 'Наличие метки "GCLID".',
        "is_dim": True,
        "name": "ym:s:firstHasGCLIDName",
        "title": "Наличие GCLID (first)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:counterID",
        "title": "Счетчик (id)",
        "type": "string",
        "src_key": "id",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:counterIDName",
        "title": "Счетчик",
        "type": "string",
    },
    {
        "description": "Дата и время на компьютере посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала 10-минутного интервала.",
        "is_dim": True,
        "name": "ym:s:browserStartOfDekaminute",
        "title": "Начало декаминуты на компьютере посетителя",
        "type": "datetime",
    },
    {
        "description": "Рекламная кампания",
        "is_dim": True,
        "name": "ym:s:openstatCampaign",
        "title": "Openstat Campaign",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:PProductID",
        "title": "ID товара в заказе",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:purchaseID",
        "title": "ID покупки",
        "type": "string",
    },
    {
        "description": "Рекламный сервис",
        "is_dim": True,
        "name": "ym:s:lastSignOpenstatService",
        "title": "Last Sign Openstat Service",
        "type": "string",
    },
    # {
    #     'description': 'Интервал дат с параметризацией размера интервала',
    #    #     'is_dim': True,
    #     'name': 'ym:s:datePeriod<group>',
    #     'title': 'Интервал дат визита',
    #     'type': 'string',
    # },
    {
        "description": "Источник перехода",
        "is_dim": True,
        "name": "ym:s:lastUTMSource",
        "title": "Last UTM Source",
        "type": "string",
    },
    {
        "description": "Дата и время предпоследнего визита посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала месяца.",
        "is_dim": True,
        "name": "ym:s:previousVisitStartOfMonth",
        "title": "Месяц предпоследнего визита",
        "type": "date",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productCategoryLevel4",
        "title": "Категория товара, ур. 4",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productCategoryLevel5",
        "title": "Категория товара, ур. 5",
        "type": "string",
    },
    {
        "description": "С каких сайтов осуществлялись переходы на сайт.",
        "is_dim": True,
        "name": "ym:s:firstReferalSource",
        "title": "Первый переход с сайтов",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productCategoryLevel2",
        "title": "Категория товара, ур. 2",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productCategoryLevel3",
        "title": "Категория товара, ур. 3",
        "type": "string",
    },
    {
        "description": "Идентификатор верхнего уровня иерархии поисковых систем.",
        "is_dim": True,
        "name": "ym:s:searchEngineRootName",
        "title": "Последняя поисковая система",
        "type": "string",
    },
    {
        "description": "Размер города по населению.",
        "is_dim": True,
        "name": "ym:s:regionCitySizeName",
        "title": "Размер города",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productCategoryLevel1",
        "title": "Категория товара, ур. 1",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:productIDCart",
        "title": "ID товара в корзине",
        "type": "string",
    },
    {
        "description": "Источник перехода",
        "is_dim": True,
        "name": "ym:s:lastSignUTMSource",
        "title": "Last Sign UTM Source",
        "type": "string",
    },
    {
        "description": "Дата визита в формате YYYY-MM-DD, округленная до начала недели (понедельник).",
        "is_dim": True,
        "name": "ym:s:startOfWeek",
        "title": "Неделя визита",
        "type": "date",
    },
    {
        "description": "Дата и время предпоследнего визита посетителя в формате YYYY-MM-DD HH:mm:ss, округленное до начала недели.",
        "is_dim": True,
        "name": "ym:s:previousVisitStartOfWeek",
        "title": "Неделя предпоследнего визита",
        "type": "date",
    },
    {
        "description": "Объявление",
        "is_dim": True,
        "name": "ym:s:lastSignUTMContent",
        "title": "Last Sign UTM Content",
        "type": "string",
    },
    {
        "description": "Группировка по достигнутой цели. Позволяет сгруппировать визиты по идентификаторам целей.",
        "is_dim": True,
        "name": "ym:s:goalName",
        "title": "Достигнутая цель",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:lastSignFrom",
        "title": "Метка from (last sign)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:browserMonth",
        "title": "Месяц на компьютере посетителя",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:firstFrom",
        "title": "Метка from (first)",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:mobilePhoneName",
        "title": "Производитель устройства",
        "type": "string",
    },
    {
        "description": "Дата на компьютере посетителя в формате YYYY-MM-DD.",
        "is_dim": True,
        "name": "ym:s:browserDate",
        "title": "Дата на компьютере посетителя",
        "type": "date",
    },
    {
        "description": 'Пол посетителя. Возможные значения: "мужской" или "женский".',
        "is_dim": True,
        "name": "ym:s:genderName",
        "title": "Пол",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:previousVisitDayOfWeekName",
        "title": "День недели предыдущего визита",
        "type": "string",
    },
    {
        "description": 'Идентификатор верхнего уровня иерархии поисковых систем первого визита посетителя. Например, "Яндекс".',
        "is_dim": True,
        "name": "ym:s:lastSignSearchEngineRootName",
        "title": "Последняя значимая поисковая система",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:firstVisitYear",
        "title": "Год первого визита",
        "type": "string",
    },
    {
        "description": "Высота клиентской части окна браузера.",
        "is_dim": True,
        "name": "ym:s:windowClientHeight",
        "title": "Высота окна",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:isNewUserName",
        "title": "Является ли визит первым визитом посетителя",
        "type": "string",
    },
    {
        "description": None,
        "is_dim": True,
        "name": "ym:s:refererDomain",
        "title": "Домен реферера",
        "type": "string",
    },
]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/api_info/metrika/visits.py
datalens-tech/datalens-backend
"""
Python Database API Specification 2.0 interface implementation for AppMetrica API
https://www.python.org/dev/peps/pep-0249/
"""

from __future__ import annotations

from sqlalchemy.types import (  # noqa; TODO: might actually be unnecessary.
    DATE,
    DATETIME,
    FLOAT,
    INTEGER,
    NULLTYPE,
    VARCHAR,
)

from dl_sqlalchemy_metrica_api import metrika_dbapi
from dl_sqlalchemy_metrica_api.api_client import APPMETRICA_API_HOST
from dl_sqlalchemy_metrica_api.api_info.appmetrica import (
    AppMetricaFieldsNamespaces,
    fields_by_namespace,
)
from dl_sqlalchemy_metrica_api.exceptions import (  # noqa
    ConnectionClosedException,
    CursorClosedException,
    DatabaseError,
    DataError,
    Error,
    IntegrityError,
    InterfaceError,
    InternalError,
    MetrikaApiAccessDeniedException,
    MetrikaApiException,
    MetrikaApiObjectNotFoundException,
    MetrikaHttpApiException,
    NotSupportedError,
    OperationalError,
    ProgrammingError,
    Warning,
)


apilevel = "2.0"
threadsafety = 2
paramstyle = "pyformat"
default_storage_plugin = ""


class Connection(metrika_dbapi.Connection):
    metrica_host = APPMETRICA_API_HOST
    metrica_fields_namespaces_enum = AppMetricaFieldsNamespaces

    @metrika_dbapi.check_connected
    def cursor(self):
        return Cursor(api_client=self._cli, connection=self)

    @metrika_dbapi.check_connected
    def get_table_names(self):
        avail_counters = self._cli.get_available_counters()
        return list(str(c_info["id"]) for c_info in avail_counters)

    def get_columns(self):
        field_props = ("name", "type", "is_dim")
        return {
            "fields": field_props,
            "data": [
                tuple(f_desc[prop] for prop in field_props) for f_desc in fields_by_namespace[self.fields_namespace]
            ],
        }


def connect(oauth_token=None, **kwargs):
    oauth_token = oauth_token or kwargs.get("password")
    fields_namespace = kwargs.get("database")
    accuracy = kwargs.get("accuracy")
    return Connection(oauth_token=oauth_token, fields_namespace=fields_namespace, accuracy=accuracy)  # , **kwargs)


class Cursor(metrika_dbapi.Cursor):
    """AppMetrica dbapi cursor"""

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/appmetrica_dbapi.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import logging
import re
from urllib.parse import urlencode

import dateutil.parser
from sqlalchemy import (
    Unicode,
    exc,
    pool,
    util,
)
from sqlalchemy.engine import (
    default,
    reflection,
)
from sqlalchemy.sql import (
    compiler,
    elements,
    operators,
    sqltypes,
)

from dl_sqlalchemy_metrica_api import (
    appmetrica_dbapi,
    metrika_dbapi,
)
from dl_sqlalchemy_metrica_api.api_info import appmetrica as appmetrica_api_info
from dl_sqlalchemy_metrica_api.api_info import metrika as metrika_api_info
from dl_sqlalchemy_metrica_api.exceptions import (
    MetrikaApiDimensionInCalc,
    MetrikaApiGroupByNotSupported,
    MetrikaApiNoMetricsNorGroupBy,
    NotSupportedError,
)


LOGGER = logging.getLogger(__name__)


MAX_LIMIT_VALUE = 10000
DEFAULT_DATE_PERIOD = 60  # days


class MetrikaApiReqPreparer(compiler.IdentifierPreparer):
    illegal_initial_characters = {"$"}
    legal_characters = re.compile(r"^[A-Z0-9_:<>\-$]+$", re.I)  # added ":<>-"

    def __init__(self, dialect, **kwargs):
        kwargs.update(initial_quote="'", escape_quote="'")
        super(MetrikaApiReqPreparer, self).__init__(dialect, **kwargs)

    def _requires_quotes(self, value):
        """Return True if the given identifier requires quoting."""
        lc_value = value.lower()
        return (
            lc_value in self.reserved_words
            or value[0] in self.illegal_initial_characters
            or not self.legal_characters.match(util.text_type(value))
            # or (lc_value != value)
        )


class MetrikaApiReqCompiler(compiler.SQLCompiler):
    _extra_select_params = None
    _extra_bind_params = None
    _ordered_columns = None
    _labeled_columns_map = None
    _group_by_fields = None
    _casts = None
    _date_filter_present = None

    def _flush_tmp_properties(self):
        self._extra_select_params = {}
        self._extra_bind_params = {}
        self._ordered_columns = False
        self._labeled_columns_map = {}
        self._group_by_fields = []
        self._casts = {}
        self._date_filter_present = False

    @property
    def api_info(self):
        return metrika_api_info

    def check_field_is_date_datetime(self, field_name):
        field_info = self.api_info.fields_by_name.get(field_name)
        return field_info and field_info["type"] in ("date", "datetime")

    def visit_column(
        self,
        column,
        add_to_result_map=None,
        include_table=True,
        _is_in_filter=False,
        **kwargs,
    ):
        name = orig_name = column.name
        if name is None:
            name = self._fallback_column_name(column)

        is_literal = column.is_literal
        if not is_literal and isinstance(name, elements._truncated_label):
            name = self._truncated_identifier("colident", name)

        if add_to_result_map is not None:
            add_to_result_map(name, orig_name, (column, name, column.key), column.type)

        # Replacing label with real field name
        if name in self._labeled_columns_map:
            name = self._labeled_columns_map[name][0]

        if _is_in_filter and self.check_field_is_date_datetime(name):
            self._date_filter_present = True

        if is_literal:
            name = self.escape_literal_column(name)
        else:
            name = self.preparer.quote(name)

        return name

    def visit_eq_binary(self, binary, operator_, **kw):
        return self._generate_generic_binary(binary, "==", **kw)  # "==" instead of "="

    def visit_inv_unary(self, element, operator_, **kw):
        return "NOT(%s)" % self.process(element.element, **kw)

    def bindparam_string(self, name, _extra_quoting=True, **kw):
        if _extra_quoting:
            return "'%s'" % (self.bindtemplate % {"name": name})
        else:
            return self.bindtemplate % {"name": name}

    def render_literal_value(self, value, type_):
        if isinstance(type_, Unicode):
            assert isinstance(value, str)
            value = value.replace("\\", "\\\\").replace("'", "\\'")

            if self.dialect.identifier_preparer._double_percents:
                value = value.replace("%", "%%")

            return "'%s'" % value

        processor = type_._cached_literal_processor(self.dialect)
        if processor:
            return processor(value)
        else:
            raise NotImplementedError("Don't know how to literal-quote value %r" % value)

    def visit_bindparam(
        self,
        bindparam,
        within_columns_clause=False,
        literal_binds=False,
        skip_bind_expression=False,
        _extra_quoting=True,
        **kwargs,
    ):
        if literal_binds or (within_columns_clause and self.ansi_bind_rules):
            if bindparam.value is None and bindparam.callable is None:
                raise exc.CompileError(
                    "Bind parameter '%s' without a " "renderable value not allowed here." % bindparam.key
                )
            return self.render_literal_bindparam(bindparam, within_columns_clause=True, **kwargs)

        if bindparam.callable:
            value = bindparam.effective_value
        else:
            value = bindparam.value

        def escape(value: str) -> str:
            return value.replace("\\", "\\\\").replace("'", "\\'")

        if bindparam.expanding:
            assert not isinstance(value, str)
            return "({})".format(
                ", ".join(f"'{escape(str(piece))}'" if _extra_quoting else escape(str(piece)) for piece in value)
            )

        if isinstance(value, list) and len(value) == 1:
            value = value[0]

        if isinstance(value, str):
            value = escape(value)

        if _extra_quoting:
            return "'%s'" % value
        return value

    def visit_between_op_binary(self, binary, operator, **kw):
        left_value = binary.left._compiler_dispatch(self, **kw)
        if self.check_field_is_date_datetime(left_value):
            if len(binary.right.clauses) != 2:
                raise exc.CompileError("Unexpected between arguments count")
            self._extra_select_params.update(
                date1=binary.right.clauses[0]._compiler_dispatch(self, _extra_quoting=False),
                date2=binary.right.clauses[1]._compiler_dispatch(self, _extra_quoting=False),
            )
        else:
            raise NotSupportedError(
                "BETWEEN operator supported only for date/datetime fields. " "Requested field: {}".format(left_value)
            )

        return ""

    def visit_not_between_op_binary(self, binary, operator, **kw):
        raise NotSupportedError()

    def visit_notbetween_op_binary(self, binary, operator, **kw):
        raise NotSupportedError()

    @util.memoized_property
    def _like_percent_literal(self):
        return elements.literal_column("'*'", type_=sqltypes.STRINGTYPE)

    def visit_contains_op_binary(self, binary, operator, **kw):
        binary = binary._clone()
        binary.right.value = "*%s*" % binary.right.value
        return self.visit_like_op_binary(binary, operator, **kw)

    def visit_not_contains_op_binary(self, binary, operator, **kw):
        binary = binary._clone()
        binary.right.value = "*%s*" % binary.right.value
        return self.visit_notlike_op_binary(binary, operator, **kw)

    def visit_notcontains_op_binary(self, binary, operator, **kw):
        return self.visit_not_contains_op_binary(binary, operator, **kw)

    def visit_startswith_op_binary(self, binary, operator, **kw):
        binary = binary._clone()
        binary.right.value = "%s*" % binary.right.value
        return self.visit_like_op_binary(binary, operator, **kw)

    def visit_not_startswith_op_binary(self, binary, operator, **kw):
        binary = binary._clone()
        binary.right.value = "%s*" % binary.right.value
        return self.visit_notlike_op_binary(binary, operator, **kw)

    def visit_notstartsswith_op_binary(self, binary, operator, **kw):
        return self.visit_not_startsswith_op_binary(binary, operator, **kw)

    def visit_endswith_op_binary(self, binary, operator, **kw):
        binary = binary._clone()
        binary.right.value = "*%s" % binary.right.value
        return self.visit_like_op_binary(binary, operator, **kw)

    def visit_not_endswith_op_binary(self, binary, operator, **kw):
        binary = binary._clone()
        binary.right.value = "*%s" % binary.right.value
        return self.visit_notlike_op_binary(binary, operator, **kw)

    def visit_notendswith_op_binary(self, binary, operator, **kw):
        return self.visit_not_endswith_op_binary(binary, operator, **kw)

    def visit_like_op_binary(self, binary, operator, **kw):
        return "%s=*%s" % (
            binary.left._compiler_dispatch(self, **kw),
            binary.right._compiler_dispatch(self, **kw),
        )

    def visit_notlike_op_binary(self, binary, operator, **kw):
        return self.visit_not_like_op_binary(binary, operator, **kw)

    def visit_not_like_op_binary(self, binary, operator, **kw):
        return "%s!*%s" % (
            binary.left._compiler_dispatch(self, **kw),
            binary.right._compiler_dispatch(self, **kw),
        )

    def visit_ilike_op_binary(self, binary, operator, **kw):
        raise NotSupportedError()

    def visit_not_ilike_op_binary(self, binary, operator, **kw):
        raise NotSupportedError()

    def visit_notilike_op_binary(self, binary, operator, **kw):
        raise NotSupportedError()

    def _get_clause_name(self, clause, **kwargs):
        if hasattr(clause, "name"):
            return clause.name
        elif hasattr(clause, "clause") and hasattr(clause.clause, "name"):
            return clause.clause.name
        else:
            return clause._compiler_dispatch(self, **kwargs)

    def visit_sum_func(self, func, **kwargs):
        clauses = func.clause_expr.element.clauses
        for cla in clauses:
            cla_name = self._get_clause_name(cla, **kwargs)
            if cla_name in self.api_info.fields_by_name and self.api_info.fields_by_name[cla_name]["is_dim"]:
                raise MetrikaApiDimensionInCalc('Not able to use dimensions in calculations: "%s"' % cla_name)

        return " + ".join([cla._compiler_dispatch(self, **kwargs) for cla in clauses])

    def visit_clauselist(self, clauselist, _group_by_clause=False, **kwargs):
        sep = clauselist.operator
        if sep is None:
            sep = " "
        elif sep == operators.comma_op:
            sep = ","  # instead of ", "
        else:
            sep = compiler.OPERATORS[clauselist.operator]

        if _group_by_clause:
            clauses = []
            for cla in clauselist.clauses:
                cla_name = self._get_clause_name(cla, **kwargs)
                if cla_name in self._labeled_columns_map and self._labeled_columns_map[cla_name][1] is not None:
                    clauses.append(self._labeled_columns_map[cla_name][1])
                else:
                    clauses.append(cla)

            for cla in clauses:
                cla_name = self._get_clause_name(cla, **kwargs)
                if cla_name not in self.api_info.fields_by_name or not self.api_info.fields_by_name[cla_name]["is_dim"]:
                    raise MetrikaApiGroupByNotSupported('Grouping by field "%s" is not possible' % cla_name)

            self._group_by_fields = [name for name in (self._get_clause_name(cla, **kwargs) for cla in clauses) if name]
        else:
            clauses = clauselist.clauses

        return sep.join(s for s in (c._compiler_dispatch(self, **kwargs) for c in clauses) if s)

    def visit_select(self, *args, **kwargs):
        self._flush_tmp_properties()
        return super().visit_select(*args, **kwargs)

    def _compose_select_body(
        self,
        text,
        select,
        compile_state,
        inner_columns,
        froms,
        byfrom,
        toplevel,
        kwargs,
    ):
        """
        https://tech.yandex.ru/metrika/doc/api2/api_v1/data-docpage/
        """

        query_params = {}

        if self.linting & compiler.COLLECT_CARTESIAN_PRODUCTS:
            from_linter = compiler.FromLinter({}, set())
            warn_linting = self.linting & compiler.WARN_LINTING
            if toplevel:
                self.from_linter = from_linter
        else:
            from_linter = None
            warn_linting = False

        if froms:
            query_params.update(
                ids=",".join(
                    [f._compiler_dispatch(self, asfrom=True, from_linter=from_linter, **kwargs) for f in froms]
                )
            )
        else:
            raise NotSupportedError('empty "FROM" clause')

        filters = []
        if select._where_criteria:
            where_str = self._generate_delimited_and_list(
                select._where_criteria, from_linter=from_linter, _is_in_filter=True, **kwargs
            )
            if where_str:
                filters.append(where_str)

        if warn_linting:
            from_linter.warn()

        if select._having_criteria:
            having_str = self._generate_delimited_and_list(
                select._having_criteria,
                _is_in_filter=True,
                **kwargs,
            )
            if having_str:
                filters.append(having_str)

        if filters:
            query_params.update(filters=" AND ".join(filters))

        if select._group_by_clauses:
            # Here only generating self._group_by_fields list.
            # It will be written into query_params later, after metrics list generation.
            select._group_by_clause._compiler_dispatch(self, _group_by_clause=True, **kwargs)

        metrics_cols = []
        metrics_dims_cols = []
        for col in inner_columns:
            col_desc = self.api_info.fields_by_name.get(col)
            if col_desc and col_desc["is_dim"]:
                metrics_dims_cols.append(col)
            else:
                metrics_cols.append(col)
        if not metrics_cols:
            if metrics_dims_cols and select._distinct:
                for dim in metrics_dims_cols:
                    self._group_by_fields.append(dim)

            if self._group_by_fields:
                # Hack to be able to get dimensions values
                # without necessity to specify any unnecessary metric field explicitly.
                fields_namespace = self.api_info.get_namespace_by_name(self._group_by_fields[0])
                any_metric = self.api_info.metrics_by_namespace[fields_namespace][0]["name"]
                metrics_cols.append(any_metric)
            else:
                raise MetrikaApiNoMetricsNorGroupBy("Not found neither metrics to select nor dimensions for group by.")
        query_params.update(metrics=",".join(metrics_cols))

        if self._group_by_fields:
            query_params.update(dimensions=",".join(self._group_by_fields))

        if select._order_by_clauses:
            order_expr = select._order_by_clause._compiler_dispatch(self, **kwargs)
            if order_expr:
                query_params.update(sort=order_expr)

        if select._limit_clause is not None and select._limit_clause.value < MAX_LIMIT_VALUE:
            limit_val = select._limit_clause.value
        else:
            limit_val = MAX_LIMIT_VALUE
        query_params.update(limit=limit_val)

        if select._offset_clause is not None:
            # In Metrika API offset starts from 1
            offset_val = int(select._offset_clause.value) + 1
            query_params.update(offset=offset_val)

        if select._for_update_arg is not None:
            raise NotSupportedError("FOR UPDATE")

        query_params.update(self._extra_select_params)

        for date_param in ("date1", "date2"):
            if date_param in query_params:
                value = query_params[date_param]
                if isinstance(value, list):
                    value = value[0]
                if not type(value) == datetime.date:
                    if isinstance(value, datetime.datetime):
                        value = value.date()
                    else:
                        value = dateutil.parser.parse(value).date()
                query_params[date_param] = value

        today = datetime.date.today()
        if not self._date_filter_present and not query_params.get("date1") and not query_params.get("date2"):
            # TODO: use counter timezone
            dt1 = today - datetime.timedelta(days=DEFAULT_DATE_PERIOD)
            query_params.update(date1=dt1, date2=today)

        if "date2" in query_params and query_params["date2"] > today:
            query_params.update(date2=today)

        return urlencode(query_params)

    def visit_table(self, table, asfrom=False, ashint=False, **kwargs):
        if asfrom or ashint:
            return table.name
        else:
            return ""

    def visit_label(
        self,
        label,
        add_to_result_map=None,
        within_label_clause=False,
        within_columns_clause=False,
        render_label_as_label=None,
        **kw,
    ):
        # only render labels within the columns clause
        # or ORDER BY clause of a select.  dialect-specific compilers
        # can modify this behavior.
        render_label_with_as = within_columns_clause and not within_label_clause
        render_label_only = render_label_as_label is label

        if render_label_only or render_label_with_as:
            if isinstance(label.name, elements._truncated_label):
                labelname = self._truncated_identifier("colident", label.name)
            else:
                labelname = label.name

        if render_label_with_as:
            if add_to_result_map is not None:
                add_to_result_map(
                    labelname,
                    label.name,
                    (label, labelname) + label._alt_names,
                    label.type,
                )

            element_processed = label.element._compiler_dispatch(
                self, within_columns_clause=True, within_label_clause=True, **kw
            )

            def _unwrap_to_column_clause(element):
                LOGGER.info(f"element {element} {type(element)}")
                if isinstance(element, elements.ColumnClause):
                    return element
                if hasattr(element, "clause"):
                    return _unwrap_to_column_clause(element.clause)
                else:
                    LOGGER.warning("Unable dispatch to ColumnClause")
                    return None

            internal_column_clause = _unwrap_to_column_clause(label.element)
            self._labeled_columns_map[labelname] = (element_processed, internal_column_clause)

            return element_processed
        else:
            return label.element._compiler_dispatch(self, within_columns_clause=False, **kw)

    def visit_asc_op_unary_modifier(self, unary, modifier, **kw):
        return unary.element._compiler_dispatch(self, **kw)

    def visit_desc_op_unary_modifier(self, unary, modifier, **kw):
        return "-" + unary.element._compiler_dispatch(self, **kw)

    def visit_cast(self, cast, **kwargs):
        param_name = cast.clause._compiler_dispatch(self, **kwargs)
        cast_type = cast.typeclause._compiler_dispatch(self, **kwargs)
        self._casts[param_name] = cast_type
        return param_name

    def visit_insert(self, *args, **kwargs):
        raise NotSupportedError("INSERT")

    def visit_update(self, *args, **kwargs):
        raise NotSupportedError("UPDATE")

    def construct_params(
        self,
        params=None,
        _group_number=None,
        _check=True,
        extracted_parameters=None,
        escape_names=True,
    ):
        prepared_params = super().construct_params(params, _group_number, _check, extracted_parameters)
        result_cols = []
        for col in self._result_columns:
            field_name = self._labeled_columns_map.get(col[0], (col[0],))[0]
            result_cols.append(
                dict(
                    label=col[0],
                    name=field_name,
                    src_key=self.api_info.fields_by_name.get(field_name, {}).get("src_key"),
                )
            )
        prepared_params["__RESULT_COLUMNS__"] = result_cols
        prepared_params.update(self._extra_bind_params)
        if self._casts:
            prepared_params["__CASTS__"] = self._casts
        LOGGER.info("Metrica query prepared params: %s", prepared_params)
        return prepared_params


# class MetrikaApiExecutionContext(default.DefaultExecutionContext):
#     def get_result_proxy(self):
#         res_proxy = result.ResultProxy(self)
#         return res_proxy


class MetrikaApiDialect(default.DefaultDialect):
    name = "metrika_api"
    supports_unicode_statements = True
    supports_unicode_binds = True
    supports_empty_insert = False
    supports_multivalues_insert = True
    supports_alter = False
    supports_pk_autoincrement = False
    supports_default_values = False
    returns_unicode_strings = True
    supports_native_boolean = False
    supports_views = False
    supports_statement_cache = False

    poolclass = pool.SingletonThreadPool  # pool.NullPool

    # execution_ctx_cls = MetrikaApiExecutionContext
    statement_compiler = MetrikaApiReqCompiler
    preparer = MetrikaApiReqPreparer

    @reflection.cache
    def get_table_names(self, connection, schema=None, **kw):
        res = connection.execute(metrika_dbapi.InternalCommands.get_tables.value)
        return [item[0] for item in res]

    def has_table(self, connection, table_name, schema=None):
        return table_name in self.get_table_names(connection)

    def get_columns(self, connection, table_name, schema=None, **kw):
        metrika_fields = connection.execute(metrika_dbapi.InternalCommands.get_columns.value)

        columns = []
        for col in metrika_fields:
            columns.append(
                {
                    "name": col["name"],
                    "type": metrika_dbapi.metrika_types_to_sqla[col["type"]],
                    "nullable": not col["is_dim"],
                }
            )
        return columns

    @classmethod
    def dbapi(cls):
        return metrika_dbapi

    def _check_unicode_returns(self, connection, additional_tests=None):
        return True

    def _check_unicode_description(self, connection):
        return True

    def do_rollback(self, dbapi_connection):
        pass

    def get_foreign_keys(self, connection, table_name, schema=None, **kw):
        return []

    def get_indexes(self, connection, table_name, schema=None, **kw):
        return []

    def get_pk_constraint(self, connection, table_name, schema=None, **kw):
        return []


class AppMetricaApiReqCompiler(MetrikaApiReqCompiler):
    @property
    def api_info(self):
        return appmetrica_api_info


class AppMetricaApiDialect(MetrikaApiDialect):
    name = "appmetrica_api"
    statement_compiler = AppMetricaApiReqCompiler
    supports_statement_cache = False

    @classmethod
    def dbapi(cls):
        return appmetrica_dbapi

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/base.py
datalens-tech/datalens-backend
# Standard exceptions according to dbapi v2 interface
# https://www.python.org/dev/peps/pep-0249/


class Error(Exception):
    pass


class Warning(Exception):
    pass


class InterfaceError(Error):
    pass


class DatabaseError(Error):
    pass


class DataError(DatabaseError):
    pass


class OperationalError(DatabaseError):
    pass


class IntegrityError(DatabaseError):
    pass


class InternalError(DatabaseError):
    pass


class ProgrammingError(DatabaseError):
    pass


class NotSupportedError(DatabaseError):
    pass


# Dialect specific exceptions


class MetrikaApiException(DatabaseError):
    def __init__(self, *args, **kwargs):
        self.orig_exc = kwargs.pop("orig_exc", None)
        super().__init__(*args, **kwargs)


class MetrikaHttpApiException(MetrikaApiException):
    pass


class MetrikaApiAccessDeniedException(MetrikaHttpApiException):
    pass


class MetrikaApiObjectNotFoundException(MetrikaHttpApiException):
    pass


class MetrikaApiGroupByNotSupported(NotSupportedError):
    pass


class MetrikaApiDimensionInCalc(NotSupportedError):
    pass


class MetrikaApiNoMetricsNorGroupBy(ProgrammingError):
    pass


class ConnectionClosedException(MetrikaApiException):
    pass


class CursorClosedException(MetrikaApiException):
    pass

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/exceptions.py
datalens-tech/datalens-backend
"""
Python Database API Specification 2.0 interface implementation for Metrica API
https://www.python.org/dev/peps/pep-0249/
"""

from __future__ import annotations

from datetime import date
from enum import Enum
from functools import wraps
from urllib.parse import parse_qs

import dateutil.parser
from sqlalchemy.types import (
    BOOLEAN,
    DATE,
    DATETIME,
    FLOAT,
    INTEGER,
    NULLTYPE,
    VARCHAR,
)

from dl_sqlalchemy_metrica_api.api_client import (
    METRIKA_API_HOST,
    MetrikaApiClient,
)
from dl_sqlalchemy_metrica_api.api_info.metrika import (
    MetrikaApiCounterSource,
    fields_by_name,
    fields_by_namespace,
)
from dl_sqlalchemy_metrica_api.exceptions import (  # noqa
    ConnectionClosedException,
    CursorClosedException,
    DatabaseError,
    DataError,
    Error,
    IntegrityError,
    InterfaceError,
    InternalError,
    MetrikaApiAccessDeniedException,
    MetrikaApiException,
    MetrikaApiObjectNotFoundException,
    MetrikaHttpApiException,
    NotSupportedError,
    OperationalError,
    ProgrammingError,
    Warning,
)


apilevel = "2.0"
threadsafety = 2
paramstyle = "pyformat"
default_storage_plugin = ""


metrika_types_to_sqla = {
    "integer": INTEGER,
    "float": FLOAT,
    "string": VARCHAR,
    "date": DATE,
    "datetime": DATETIME,
    # "percents": FLOAT,
}


cast_processors = {
    "INTEGER": int,
    "FLOAT": float,
    "VARCHAR": str,
    "BOOLEAN": bool,
    "DATE": lambda t: dateutil.parser.parse(t).date(),
    "DATETIME": lambda t: dateutil.parser.parse(t),
}


cast_processors_for_metrika_types = {
    "date": cast_processors["DATE"],
    "datetime": cast_processors["DATETIME"],
}


cast_type_to_sqla_type = {
    "INTEGER": INTEGER,
    "FLOAT": FLOAT,
    "VARCHAR": VARCHAR,
    "BOOLEAN": BOOLEAN,
    "DATE": DATE,
    "DATETIME": DATETIME,
}


def check_connected(func):
    @wraps(func)
    def func_wrapper(self, *args, **kwargs):
        if self.is_connected is False:
            raise ConnectionClosedException("Connection object is closed")
        else:
            return func(self, *args, **kwargs)

    return func_wrapper


class Connection(object):
    metrica_host = METRIKA_API_HOST
    metrica_fields_namespaces_enum = MetrikaApiCounterSource

    fields_namespace = None
    accuracy = None

    def __init__(self, oauth_token, fields_namespace=None, accuracy=None, **client_kwargs):
        client_kwargs.setdefault("host", self.metrica_host)
        self._cli = MetrikaApiClient(oauth_token, **client_kwargs)
        if fields_namespace:
            if not hasattr(self.metrica_fields_namespaces_enum, fields_namespace):
                raise MetrikaApiException("Unknown fields namespace: %s" % fields_namespace)
            self.fields_namespace = self.metrica_fields_namespaces_enum[fields_namespace]
        self.accuracy = accuracy

        self._connected = True

    @property
    def is_connected(self):
        return self._connected

    @check_connected
    def close(self):
        del self._cli
        self._cli = None
        self._connected = False

    @check_connected
    def commit(self):
        pass

    @check_connected
    def cursor(self):
        return Cursor(api_client=self._cli, connection=self)

    @check_connected
    def get_table_names(self):
        avail_counters = self._cli.get_available_counters()
        return list(str(c_info["id"]) for c_info in avail_counters)

    def get_columns(self):
        field_props = ("name", "type", "is_dim")
        return {
            "fields": field_props,
            "data": [
                tuple(f_desc[prop] for prop in field_props) for f_desc in fields_by_namespace[self.fields_namespace]
            ],
        }

    @check_connected
    def get_avail_date_min(self, counter_id):
        return self._cli.get_counter_creation_date(counter_id)


def connect(oauth_token=None, **kwargs):
    oauth_token = oauth_token or kwargs.get("password")
    fields_namespace = kwargs.get("database")
    accuracy = kwargs.get("accuracy")
    return Connection(oauth_token=oauth_token, fields_namespace=fields_namespace, accuracy=accuracy)  # , **kwargs)


def check_cursor_connected(func):
    @wraps(func)
    def func_wrapper(self, *args, **kwargs):
        if not self._connected:
            raise CursorClosedException("Cursor object is closed")
        elif not self.connection.is_connected:
            raise ConnectionClosedException("Connection object is closed")
        else:
            return func(self, *args, **kwargs)

    return func_wrapper


class InternalCommands(Enum):
    get_columns = "__GET_COLUMNS_COMMAND__"
    get_tables = "__GET_TABLES_COMMAND__"
    get_avail_date_min = "__GET_AVAILABLE_DATE_MIN__"
    get_avail_date_max = "__GET_AVAILABLE_DATE_MAX__"


class Cursor(object):
    description = None
    rowcount = -1
    arraysize = 1
    _result_data = None

    def __init__(self, api_client: MetrikaApiClient, connection: Connection):
        self._cli = api_client
        self.connection = connection
        self._connected = True

    @property
    def is_connected(self):
        return self._connected

    @check_cursor_connected
    def close(self):
        self._connected = False

    def _prepare_query_params(self, query, subst_params) -> dict:
        query_params = parse_qs(query)

        if subst_params:
            for k, v in query_params.items():
                if len(v) != 1:
                    raise ProgrammingError("Unexpected multiple parameter %s values %s" % (k, v))
                v = v[0] % subst_params
                query_params[k] = v

        return query_params

    def _exec_get_columns(self, operation, parameters):
        res = self.connection.get_columns()
        self._result_data = res["data"]
        self.rowcount = len(self._result_data)
        self.description = [(f_name, VARCHAR, None, None, None, None, None) for f_name in res["fields"]]

    def _exec_get_tables(self, operation, parameters):
        table_names = self.connection.get_table_names()
        self._result_data = [(tn,) for tn in table_names]
        self.rowcount = len(self._result_data)
        self.description = [
            ("name", VARCHAR, None, None, None, None, None),
        ]

    def _exec_get_avail_date_min(self, operation, parameters):
        counter_id = parameters.get("_COUNTER_ID_")
        date_min = self.connection.get_avail_date_min(counter_id)
        self._result_data = [(date_min.isoformat(),)]
        self.rowcount = 1
        result_columns = parameters.get("__RESULT_COLUMNS__", [])
        if result_columns:
            col_name = result_columns[0]["label"] or result_columns[0]["name"]
        else:
            col_name = "date_min"
        self.description = [
            (col_name, DATE, None, None, None, None, None),
        ]

    def _exec_get_avail_date_max(self, operation, parameters):
        # TODO: use counter timezone
        today = date.today()
        self._result_data = [(today.isoformat(),)]
        self.rowcount = 1
        result_columns = parameters.pop("__RESULT_COLUMNS__", [])
        if result_columns:
            col_name = result_columns[0]["label"] or result_columns[0]["name"]
        else:
            col_name = "date_max"
        self.description = [
            (col_name, DATE, None, None, None, None, None),
        ]

    def _meth_by_operation(self, operation):
        meth = {
            InternalCommands.get_columns.value: self._exec_get_columns,
            InternalCommands.get_tables.value: self._exec_get_tables,
            InternalCommands.get_avail_date_min.value: self._exec_get_avail_date_min,
            InternalCommands.get_avail_date_max.value: self._exec_get_avail_date_max,
        }.get(operation)
        return meth

    @check_cursor_connected
    def execute(self, operation, parameters=None):
        self.rowcount = None
        self._result_data = None
        self.description = None

        meth = self._meth_by_operation(operation) or self._execute_select_data
        return meth(operation, parameters)

    def _execute_select_data(self, operation, parameters):
        casts = parameters.pop("__CASTS__", None)
        result_columns = parameters.pop("__RESULT_COLUMNS__", None)
        for col in result_columns:
            col_name = col["name"]
            if casts and col_name in casts:
                col["cast_processor"] = cast_processors[casts[col_name]]
            elif col_name in fields_by_name and fields_by_name[col_name]["type"] in cast_processors_for_metrika_types:
                col["cast_processor"] = cast_processors_for_metrika_types[fields_by_name[col_name]["type"]]

        query_params = self._prepare_query_params(operation, parameters)

        if self.connection.accuracy is not None:
            query_params.update(accuracy=self.connection.accuracy)

        result = self._cli.get_table_data(query_params, result_columns=result_columns)

        self._result_data = result["data"]
        self.rowcount = len(self._result_data)

        self.description = []
        for col in result["fields"]:
            if casts and col["name"] in casts:
                col_type = cast_type_to_sqla_type[casts[col["name"]]]
            elif col["name"] not in fields_by_name:
                col_type = VARCHAR
            else:
                col_type = metrika_types_to_sqla.get(fields_by_name[col["name"]]["type"], NULLTYPE)
            self.description.append((col["label"] or col["name"], col_type, None, None, None, None, None))

    # def executemany(self, sql: str, seq_of_parameters: Iterable[Iterable]): ...

    @check_cursor_connected
    def fetchone(self):
        try:
            if self._result_data:
                return self._result_data.pop(0)
            else:
                return None
        except StopIteration:
            return None

    @check_cursor_connected
    def fetchmany(self, size=None):
        if size is None or size > len(self._result_data):
            size = len(self._result_data)
        rows = self._result_data[:size]
        self._result_data = self._result_data[size:]
        return rows

    @check_cursor_connected
    def fetchall(self):
        rows = self._result_data
        self._result_data = []
        return rows

    def __iter__(self):
        raise NotImplementedError()

    def setinputsizes(self, *args, **kwargs):
        pass

    def setoutputsize(self, *args, **kwargs):
        pass

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api/metrika_dbapi.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
import os

import pytest
import sqlalchemy

import dl_sqlalchemy_metrica_api


METRIKA_SAMPLE_COUNTER_ID = "44147844"
APPMETRICA_SAMPLE_COUNTER_ID = "1111"


def _get_oauth_from_env():
    return os.environ.get("METRIKA_OAUTH", None)


@pytest.fixture
def metrika_sample_counter_id():
    return METRIKA_SAMPLE_COUNTER_ID


@pytest.fixture
def appmetrica_sample_counter_id():
    return APPMETRICA_SAMPLE_COUNTER_ID


@pytest.fixture(scope="function")
def shrink_metrika_default_date_period(monkeypatch):
    """
    To reduce load to Metrika API and tests run time.
    """
    monkeypatch.setattr(dl_sqlalchemy_metrica_api.base, "DEFAULT_DATE_PERIOD", 3)


def _metrika_db_engine():
    return sqlalchemy.create_engine("metrika_api://:{}@/hits".format(_get_oauth_from_env()))


def _appmetrica_db_engine():
    return sqlalchemy.create_engine("appmetrica_api://:{}@/installs".format(_get_oauth_from_env()))


@pytest.fixture(scope="function")
def metrika_db_engine(shrink_metrika_default_date_period):
    return _metrika_db_engine()


@pytest.fixture(scope="function")
def metrika_db_engine_with_accuracy(shrink_metrika_default_date_period):
    return sqlalchemy.create_engine("metrika_api://:{}@/hits?accuracy=0.1".format(_get_oauth_from_env()))


@pytest.fixture(scope="function")
def appmetrica_db_engine(shrink_metrika_default_date_period):
    return _appmetrica_db_engine()


@pytest.fixture(params=[_metrika_db_engine, _appmetrica_db_engine])
def db_engine(request, shrink_metrika_default_date_period):
    return request.param()


def _gen_m_expr1(eng):
    expr = sqlalchemy.select(
        columns=[
            sqlalchemy.type_coerce(sqlalchemy.column("ym:pv:pageviews"), sqlalchemy.Float()),
            sqlalchemy.column("ym:pv:users").label("users"),
            sqlalchemy.column("ym:pv:date").label("date"),
            # sqlalchemy.column('ym:pv:URLDomain'),
            sqlalchemy.column("ym:pv:browserName").label("browser"),
        ]
    )
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.where(sqlalchemy.column("date").between("2019-02-01", "2019-02-03"))
    expr = expr.where(sqlalchemy.column("browser") == "Yandex Browser")
    # expr = expr.group_by(sqlalchemy.column('ym:pv:screenResolution'))
    # expr = expr.group_by(sqlalchemy.column('ym:pv:date'))
    expr = expr.group_by(sqlalchemy.column("date"))
    expr = expr.group_by(sqlalchemy.column("browser"))
    expr = expr.limit(5)
    # expr = expr.offset(50)
    expr = expr.order_by(sqlalchemy.column("ym:pv:date"))
    return expr


def _gen_m_expr_distinct_via_group_by(eng):
    expr = sqlalchemy.select(columns=[sqlalchemy.column("ym:pv:browserName").label("browser")])
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.group_by(sqlalchemy.column("browser"))
    return expr


def _gen_m_expr_distinct(eng):
    expr = sqlalchemy.select(columns=[sqlalchemy.column("ym:pv:browserName")]).distinct()
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    return expr


def _gen_m_expr_date_min(eng):
    expr = sqlalchemy.select(columns=[sqlalchemy.column("ym:pv:date").label("date")])
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.group_by(sqlalchemy.column("date"))
    expr = expr.limit(1)
    # expr = expr.order_by('date')
    expr = expr.order_by(sqlalchemy.column("date"))
    return expr


def _gen_m_expr_date_max(eng):
    expr = sqlalchemy.select(columns=[sqlalchemy.column("ym:pv:date").label("date")])
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.group_by(sqlalchemy.column("date"))
    expr = expr.limit(1)
    expr = expr.order_by(sqlalchemy.desc(sqlalchemy.column("date")))
    return expr


def _gen_m_expr_order_by(eng):
    expr = sqlalchemy.select(
        columns=[
            sqlalchemy.column("ym:pv:date").label("date"),
            sqlalchemy.column("ym:pv:users").label("users"),
        ]
    )
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.group_by(sqlalchemy.column("date"))
    expr = expr.order_by(sqlalchemy.desc(sqlalchemy.literal_column("users")))
    return expr


def _gen_m_expr_cast(eng):
    expr = sqlalchemy.select(
        columns=[
            sqlalchemy.cast(sqlalchemy.column("ym:pv:pageviews"), sqlalchemy.INTEGER),
            sqlalchemy.cast(sqlalchemy.column("ym:pv:users"), sqlalchemy.String).label("users_str"),
            sqlalchemy.type_coerce(
                sqlalchemy.type_coerce(sqlalchemy.column("ym:pv:users"), sqlalchemy.Integer), sqlalchemy.Integer
            ).label("users_int"),
        ]
    )
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    return expr


def _gen_m_expr_date_between(eng):
    expr = sqlalchemy.select(columns=[sqlalchemy.column("ym:pv:pageviews"), sqlalchemy.column("ym:pv:users")])
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.where(sqlalchemy.column("ym:pv:date").between("2019-02-01", "2019-02-03"))
    return expr


def _gen_m_expr_datetime_between_date(eng):
    expr = sqlalchemy.select(columns=[sqlalchemy.column("ym:pv:pageviews"), sqlalchemy.column("ym:pv:users")])
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.where(sqlalchemy.column("ym:pv:dateTime").between("2019-02-01", "2019-02-03"))
    return expr


def _gen_m_expr_datetime_between_datetime(eng):
    expr = sqlalchemy.select(columns=[sqlalchemy.column("ym:pv:pageviews"), sqlalchemy.column("ym:pv:users")])
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.where(
        sqlalchemy.column("ym:pv:dateTime").between(
            datetime.datetime.fromisoformat("2019-02-01 00:00:00"),
            datetime.datetime.fromisoformat("2019-02-03 23:59:59"),
        )
    )
    return expr


def _gen_m_expr_datetime_between_datetime_str(eng):
    expr = sqlalchemy.select(columns=[sqlalchemy.column("ym:pv:pageviews"), sqlalchemy.column("ym:pv:users")])
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.where(sqlalchemy.column("ym:pv:dateTime").between("2019-02-01 00:00:00", "2019-02-03 23:59:59"))
    return expr


def __gen_select_users_browser(eng):
    expr = sqlalchemy.select(
        columns=[
            sqlalchemy.column("ym:pv:users").label("users"),
            sqlalchemy.column("ym:pv:browserName").label("browser"),
        ]
    )
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.group_by(sqlalchemy.column("browser"))
    return expr


def _gen_m_expr_like(eng):
    expr = __gen_select_users_browser(eng)
    expr = expr.where(sqlalchemy.column("browser").like("Yandex*"))
    return expr


def _gen_m_expr_not_like(eng):
    expr = __gen_select_users_browser(eng)
    expr = expr.where(sqlalchemy.column("browser").notlike("Yandex*"))
    return expr


def _gen_m_expr_not_like_1(eng):
    expr = __gen_select_users_browser(eng)
    expr = expr.where(sqlalchemy.not_(sqlalchemy.column("browser").like("Yandex*")))
    return expr


def _gen_m_expr_contains(eng):
    expr = __gen_select_users_browser(eng)
    expr = expr.where(sqlalchemy.column("browser").contains("Yandex"))
    return expr


def _gen_m_expr_not_contains(eng):
    expr = __gen_select_users_browser(eng)
    expr = expr.where(sqlalchemy.not_(sqlalchemy.column("browser").contains("Yandex")))
    return expr


def _gen_m_expr_startswith(eng):
    expr = __gen_select_users_browser(eng)
    expr = expr.where(sqlalchemy.column("browser").startswith("Yandex"))
    return expr


def _gen_m_expr_not_endswith(eng):
    expr = __gen_select_users_browser(eng)
    expr = expr.where(sqlalchemy.not_(sqlalchemy.column("browser").endswith("Chrome")))
    return expr


def _gen_m_expr_not(eng):
    expr = __gen_select_users_browser(eng)
    expr = expr.where(
        sqlalchemy.not_(
            sqlalchemy.or_(
                sqlalchemy.column("browser").like("Yandex*"),
                sqlalchemy.column("browser") == "Google Chrome",
            )
        )
    )
    return expr


def _gen_m_expr_not_in(eng):
    expr = sqlalchemy.select(
        columns=[
            sqlalchemy.column("ym:pv:users").label("users"),
            sqlalchemy.column("ym:pv:browserName").label("browser"),
        ]
    )
    expr = expr.select_from(sqlalchemy.Table(METRIKA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.group_by(sqlalchemy.column("browser"))
    expr = expr.where(sqlalchemy.column("browser").notin_(["Yandex Browser", "Google Chrome"]))
    return expr


def _gen_am_expr1(eng):
    expr = sqlalchemy.select(columns=[sqlalchemy.column("ym:u:sessions"), sqlalchemy.column("ym:u:activeUsers")])
    expr = expr.group_by(sqlalchemy.column("ym:u:dayOfWeekName"))
    expr = expr.select_from(sqlalchemy.Table(APPMETRICA_SAMPLE_COUNTER_ID, sqlalchemy.MetaData(bind=eng)))
    expr = expr.where(sqlalchemy.column("ym:u:date").between("2019-05-01", "2019-05-03"))
    return expr


metrika_expr_functions_list = [
    _gen_m_expr1,
    _gen_m_expr_date_max,
    _gen_m_expr_date_min,
    _gen_m_expr_distinct,
    _gen_m_expr_distinct_via_group_by,
    _gen_m_expr_cast,
    _gen_m_expr_date_between,
    _gen_m_expr_datetime_between_date,
    _gen_m_expr_datetime_between_datetime,
    _gen_m_expr_datetime_between_datetime_str,
    _gen_m_expr_like,
    _gen_m_expr_not_like,
    _gen_m_expr_not_like_1,
    _gen_m_expr_contains,
    _gen_m_expr_not_contains,
    _gen_m_expr_not,
    _gen_m_expr_not_in,
    _gen_m_expr_order_by,
]

appmetrica_expr_functions_list = [
    _gen_am_expr1,
]


@pytest.fixture(params=metrika_expr_functions_list)
def metrika_expr_func(request):
    return request.param


@pytest.fixture(params=appmetrica_expr_functions_list)
def appmetrica_expr_func(request):
    return request.param


@pytest.fixture
def m_expr_distinct(metrika_sample_counter_id, metrika_db_engine):
    expr = sqlalchemy.select(columns=[sqlalchemy.column("ym:pv:dateTime")]).distinct()
    expr = expr.select_from(sqlalchemy.Table(metrika_sample_counter_id, sqlalchemy.MetaData(bind=metrika_db_engine)))
    return expr


@pytest.fixture()
def metrika_expr_func_expr1():
    return _gen_m_expr1


@pytest.fixture
def metrika_expr_select_date_users(metrika_sample_counter_id, metrika_db_engine):
    expr = sqlalchemy.select(
        columns=[sqlalchemy.column("ym:pv:date").label("date"), sqlalchemy.column("ym:pv:users")],
    )
    expr = expr.select_from(sqlalchemy.Table(metrika_sample_counter_id, sqlalchemy.MetaData(bind=metrika_db_engine)))
    expr = expr.group_by(sqlalchemy.column("date"))
    return expr

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api_tests/ext/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from urllib.parse import parse_qs

import pytest
import sqlalchemy as sa
from sqlalchemy.dialects import registry
from sqlalchemy.exc import DatabaseError

from dl_sqlalchemy_metrica_api import exceptions


registry.register("metrika_api", "dl_sqlalchemy_metrica_api.base", "MetrikaApiDialect")
registry.register("appmetrica_api", "dl_sqlalchemy_metrica_api.base", "AppMetricaApiDialect")


def _test_execute_expr(_expr_func, _db_engine):
    expr = _expr_func(_db_engine)

    compiled = str(expr)
    print(compiled)
    res = _db_engine.execute(expr)
    # print(res._cursor_description())
    data = res.fetchall()
    for row in data:
        print(row)


def test_execute_expr(metrika_expr_func, metrika_db_engine):
    return _test_execute_expr(metrika_expr_func, metrika_db_engine)


def test_execute_appmetrica_expr(appmetrica_expr_func, appmetrica_db_engine):
    return _test_execute_expr(appmetrica_expr_func, appmetrica_db_engine)


def test_get_columns(db_engine):
    insp = sa.inspect(db_engine)
    cols = insp.get_columns("123456")
    print(cols)
    assert isinstance(cols, list)
    assert isinstance(cols[0], dict)
    assert "name" in cols[0] and "type" in cols[0]


def test_distincts(m_expr_distinct, metrika_db_engine):
    res = metrika_db_engine.execute(m_expr_distinct)
    # print(res._cursor_description())
    data = res.fetchall()
    values = [row[0] for row in data]
    print(values)
    print(set(values))
    print(sorted(set(values)))
    assert len(data) == len(set(values))


def test_accuracy(metrika_expr_func_expr1, metrika_db_engine_with_accuracy):
    expr = metrika_expr_func_expr1(metrika_db_engine_with_accuracy)
    res = metrika_db_engine_with_accuracy.execute(expr)
    assert res.cursor.connection.accuracy == "0.1"


def test_exceptions(m_expr_distinct, metrika_db_engine):
    engine_invalid_token = sa.create_engine("metrika_api://:qwerty@/hits")
    with pytest.raises(DatabaseError) as exc_info:
        engine_invalid_token.execute(m_expr_distinct)
    assert isinstance(exc_info.value.orig, exceptions.MetrikaApiAccessDeniedException)
    assert "Invalid oauth_token" in str(exc_info.value.orig)

    expr_invalid_counter_id = sa.select(columns=[sa.column("ym:pv:users")])
    expr_invalid_counter_id = expr_invalid_counter_id.select_from(
        sa.Table("2147483647", sa.MetaData(bind=metrika_db_engine))
    )
    with pytest.raises(DatabaseError) as exc_info:
        metrika_db_engine.execute(expr_invalid_counter_id)
    assert isinstance(exc_info.value.orig, exceptions.MetrikaApiObjectNotFoundException)
    assert "Entity not found" in str(exc_info.value.orig)


def test_future_date(metrika_expr_select_date_users, metrika_db_engine):
    today = datetime.date.today()
    expr = metrika_expr_select_date_users.where(
        sa.column("ym:pv:date").between(
            datetime.datetime.fromisoformat((today - datetime.timedelta(days=2)).isoformat()),
            datetime.datetime.fromisoformat("2039-02-03"),  # date from future
        )
    )
    res = metrika_db_engine.execute(expr)
    data = res.fetchall()
    max_date = max([row["date"] for row in data])
    assert max_date == today
    assert len(data) == 3


def test_date_gt(metrika_expr_select_date_users, metrika_db_engine):
    start_dt = (datetime.date.today() - datetime.timedelta(days=5)).isoformat()
    expr = metrika_expr_select_date_users.where(sa.column("ym:pv:date") > start_dt)
    print(str(expr))

    res = metrika_db_engine.execute(expr)
    data = res.fetchall()
    assert len(data) == 5


def test_date_gte(metrika_expr_select_date_users, metrika_db_engine):
    start_dt = (datetime.date.today() - datetime.timedelta(days=5)).isoformat()
    expr = metrika_expr_select_date_users.where(sa.column("ym:pv:date") >= start_dt)
    print(str(expr))

    res = metrika_db_engine.execute(expr)
    data = res.fetchall()
    assert len(data) == 6


def test_date_gt_lt(metrika_expr_select_date_users, metrika_db_engine):
    start_dt = (datetime.date.today() - datetime.timedelta(days=5)).isoformat()
    end_dt = (datetime.date.today() - datetime.timedelta(days=2)).isoformat()
    expr = metrika_expr_select_date_users.where(sa.column("ym:pv:date") > start_dt)
    expr = expr.where(sa.column("ym:pv:date") < end_dt)
    print(str(expr))

    res = metrika_db_engine.execute(expr)
    data = res.fetchall()
    print(data)
    assert len(data) == 2


def test_date_eq(metrika_expr_select_date_users, metrika_db_engine):
    start_dt = (datetime.date.today() - datetime.timedelta(days=1)).isoformat()
    expr = metrika_expr_select_date_users.where(sa.column("ym:pv:date") == start_dt)
    print(str(expr))

    res = metrika_db_engine.execute(expr)
    data = res.fetchall()
    assert len(data) == 1


def test_date_gt_lt_by_alias(metrika_expr_select_date_users, metrika_db_engine):
    start_dt = (datetime.date.today() - datetime.timedelta(days=5)).isoformat()
    end_dt = (datetime.date.today() - datetime.timedelta(days=2)).isoformat()
    expr = metrika_expr_select_date_users.where(sa.column("date") > start_dt)
    expr = expr.where(sa.column("date") < end_dt)
    print(str(expr))

    res = metrika_db_engine.execute(expr)
    data = res.fetchall()
    print(data)
    assert len(data) == 2


def test_calculations(metrika_sample_counter_id, metrika_db_engine):
    expr = sa.select(
        columns=[
            sa.column("ym:pv:date").label("date"),
            sa.column("ym:pv:users").label("users"),
            sa.column("ym:pv:pageviews").label("pageviews"),
            sa.column("ym:pv:pageviewsPerDay"),
            (sa.column("ym:pv:pageviews") + sa.column("ym:pv:users")).label("pv_plus_u"),
            sa.func.sum(
                sa.column("ym:pv:pageviews"),
                sa.column("ym:pv:users"),
            ).label("sum_pv_u"),
            (sa.column("ym:pv:pageviews") - sa.column("users")).label("pv_minus_u"),
            (sa.column("ym:pv:pageviews") + 1).label("pv_plus_1"),
            (sa.column("ym:pv:users") - 2).label("u_minus_2"),
            (sa.column("ym:pv:pageviews") / sa.column("ym:pv:users")).label("pv_per_u"),
            (sa.column("ym:pv:pageviews") * sa.column("ym:pv:users") - sa.column("ym:pv:pageviewsPerDay") + 14).label(
                "some_strange_expr"
            ),
        ],
    )
    expr = expr.select_from(sa.Table(metrika_sample_counter_id, sa.MetaData(bind=metrika_db_engine)))
    expr = expr.group_by(sa.column("date"))
    print(str(expr))

    res = metrika_db_engine.execute(expr)
    data = res.fetchall()
    print(data)
    assert all([(row["pageviews"] + row["users"]) == row["pv_plus_u"] for row in data])
    assert all([(row["pageviews"] + row["users"]) == row["sum_pv_u"] for row in data])
    assert all([(row["pageviews"] - row["users"]) == row["pv_minus_u"] for row in data])
    assert all([row["pv_plus_1"] == row["pageviews"] + 1 for row in data])
    assert all([row["u_minus_2"] == row["users"] - 2 for row in data])
    assert all([round(row["pageviews"] / row["users"], 2) == round(row["pv_per_u"], 2) for row in data])
    assert all(
        [
            round(row["pageviews"] * row["users"] - row["ym:pv:pageviewsPerDay"] + 14, 2)
            == round(row["some_strange_expr"], 2)
            for row in data
        ]
    )


def test_counter_id(metrika_sample_counter_id, metrika_db_engine):
    expr = sa.select(
        columns=[
            sa.column("ym:pv:date").label("date"),
            sa.column("ym:pv:users").label("users"),
            sa.column("ym:pv:counterID").label("counter_id"),
            sa.column("ym:pv:counterIDName").label("counter_name"),
        ],
    )
    expr = expr.select_from(sa.Table(metrika_sample_counter_id, sa.MetaData(bind=metrika_db_engine)))
    expr = expr.where(sa.column("counter_id").in_([metrika_sample_counter_id]))
    expr = expr.where(sa.column("counter_name").in_(["Metrica live demo"]))
    expr = expr.group_by(
        sa.column("date"),
        sa.column("counter_id"),
        sa.column("counter_name"),
    )
    print(str(expr))

    res = metrika_db_engine.execute(expr)
    data = res.fetchall()
    print(data)
    assert all([(row["counter_id"] == metrika_sample_counter_id for row in data)])


def test_multicounter_req(metrika_sample_counter_id, metrika_db_engine):
    expr = sa.select(
        columns=[
            sa.column("ym:pv:date").label("date"),
            sa.column("ym:pv:users").label("users"),
            sa.column("ym:pv:counterID").label("counter_id"),
            sa.column("ym:pv:counterIDName").label("counter_name"),
        ],
    )
    expr = expr.select_from(
        sa.Table("50514217,51341415", sa.MetaData(bind=metrika_db_engine)),
    )
    start_dt = (datetime.date.today() - datetime.timedelta(days=5)).isoformat()
    end_dt = (datetime.date.today() - datetime.timedelta(days=2)).isoformat()
    expr = expr.where(sa.column("date").between(start_dt, end_dt))
    expr = expr.group_by(
        sa.column("date"),
        sa.column("counter_id"),
        sa.column("counter_name"),
    )
    query = parse_qs(str(expr))
    print(query)
    assert query["ids"] == ["50514217,51341415"]

lib/dl_sqlalchemy_metrica_api/dl_sqlalchemy_metrica_api_tests/ext/test_dialect.py
datalens-tech/datalens-backend
from __future__ import annotations

from .base import BIMSSQLDialect


__all__ = ("BIMSSQLDialect",)

lib/dl_sqlalchemy_mssql/dl_sqlalchemy_mssql/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

import binascii
import datetime
import decimal
import itertools
import logging
import uuid

import pyodbc
import sqlalchemy as sa
import sqlalchemy.dialects.mssql.base as upbase
from sqlalchemy.dialects.mssql.pyodbc import MSDialect_pyodbc as UPSTREAM

from dl_sqlalchemy_common.base import CompilerPrettyMixin


LOGGER = logging.getLogger(__name__)


class BIMSSQLDialectBasic(UPSTREAM):
    @staticmethod  # noqa: C901
    def _quote_simple_value(value):
        """Mainly from pymssql quoting, without the encoded output"""

        if value is None:
            return "NULL"

        if isinstance(value, bool):
            return "1" if value else "0"

        if isinstance(value, float):
            return repr(value)

        if isinstance(value, (int, decimal.Decimal)):
            return str(value)

        if isinstance(value, uuid.UUID):
            return BIMSSQLDialect._quote_simple_value(str(value))

        if isinstance(value, str):
            return "N'" + value.replace("'", "''") + "'"

        if isinstance(value, bytearray):
            return "0x" + binascii.hexlify(bytes(value)).decode()

        if isinstance(value, bytes):
            # see if it can be decoded as ascii if there are no null bytes
            if b"\0" not in value:
                try:
                    value.decode("ascii")
                    return "'" + value.replace(b"'", b"''").decode() + "'"
                except UnicodeDecodeError:
                    pass

            # Python 3: handle bytes
            # @todo - Marc - hack hack hack
            if isinstance(value, bytes):
                return (b"0x" + binascii.hexlify(value)).decode()

            # will still be string type if there was a null byte in it or if the
            # decoding failed.  In this case, just send it as hex.
            if isinstance(value, str):
                return "0x" + value.encode("hex")

        if isinstance(value, datetime.datetime):
            return "{ts '%04d-%02d-%02d %02d:%02d:%02d.%03d'}" % (
                value.year,
                value.month,
                value.day,
                value.hour,
                value.minute,
                value.second,
                value.microsecond / 1000,
            )

        if isinstance(value, datetime.date):
            return "{d '%04d-%02d-%02d'}" % (value.year, value.month, value.day)

        if isinstance(value, tuple) and len(value) == 1:
            # Sometimes a result of a query will be provided as a filter to another query, but
            # the object passed may not be scalar. pymssql does some flattening on this, and
            # we should too.
            return BIMSSQLDialect._quote_simple_value(value[0])

        return None

    @staticmethod
    def translate_custom_parameters(params):
        def translate(param):
            if isinstance(param, tuple):
                # Sometimes a result of a query will be provided as a filter to another query, but
                # the object passed may not be scalar. pymssql does some flattening on this, and
                # we should too.
                return param[0]
            return param

        return [translate(param) for param in params]

    def roll_parameters_into_statement(self, statement, parameters):
        assert statement.count("?") == len(parameters), "num of placeholders != num of params"

        # transform here and pass on to cursor
        quoted_params = [self._quote_simple_value(param) for param in parameters]
        statement_list = statement.split("?")
        # Most efficient method I found for getting the lists put together. Assumes that
        # the number of parameters actually matches the number of ? placeholders.
        return "".join(
            [x for x in itertools.chain.from_iterable(itertools.zip_longest(statement_list, quoted_params)) if x]
        )

    def do_execute(self, cursor, statement, parameters, context=None):
        if parameters:
            statement = self.roll_parameters_into_statement(statement, parameters)
            # no need for parameters at this point, since they're all baked into the query
            parameters = tuple()

        try:
            cursor.execute(statement, self.translate_custom_parameters(parameters))
        except pyodbc.OperationalError:
            LOGGER.error(
                "pyodbc OperationalError. Full statement: {}\n Params: {}".format(
                    statement,
                    parameters,
                )
            )
            raise

    @upbase._db_plus_owner
    def has_table(self, connection, tablename, dbname, owner, schema):
        self._ensure_has_table_connection(connection)
        if tablename.startswith("#"):  # temporary table
            tables = upbase.ischema.mssql_temp_table_columns

            s = upbase.sql.select(tables.c.table_name).where(
                tables.c.table_name.like(self._temp_table_name_like_pattern(tablename))
            )

            result = connection.execute(s.limit(1))
            return result.scalar() is not None
        else:
            tables = upbase.ischema.tables

            s = upbase.sql.select(tables.c.table_name).where(
                upbase.sql.and_(
                    # Original: `tables.c.table_type == "BASE TABLE",`
                    upbase.sql.or_(
                        tables.c.table_type == "BASE TABLE",
                        tables.c.table_type == "VIEW",
                    ),
                    # ...
                    tables.c.table_name == tablename,
                )
            )

            if owner:
                s = s.where(tables.c.table_schema == owner)

            c = connection.execute(s)

            return c.first() is not None


class BIMSSQLCompiler(UPSTREAM.statement_compiler, CompilerPrettyMixin):
    def order_by_clause(self, select, **kw):
        sup = super().order_by_clause(select, **kw)
        return self._pretty.postprocess_block("ORDER BY", sup)

    def limit_clause(self, cs, **kw):
        # Upstream: `return ""`
        return super().limit_clause(cs, **kw)

    def visit_alias(self, alias, **kw):
        # Upstream only wraps its superclass.
        return super().visit_alias(alias, **kw)


class BIMSSQLDialect(BIMSSQLDialectBasic):
    statement_compiler = BIMSSQLCompiler


def register_dialect():
    sa.dialects.registry.register("bi_mssql_basic", "dl_sqlalchemy_mssql.base", "BIMSSQLDialectBasic")
    sa.dialects.registry.register("bi_mssql", "dl_sqlalchemy_mssql.base", "BIMSSQLDialect")

lib/dl_sqlalchemy_mssql/dl_sqlalchemy_mssql/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest
import sqlalchemy as sa

from dl_sqlalchemy_mssql.base import register_dialect


def pytest_configure(config):  # noqa
    register_dialect()


@pytest.fixture(scope="session")
def engine_url():
    return "bi_mssql://"


@pytest.fixture(scope="session")
def sa_engine(engine_url):
    return sa.create_engine(engine_url)

lib/dl_sqlalchemy_mssql/dl_sqlalchemy_mssql_tests/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from decimal import Decimal
from unittest import mock
from uuid import UUID

import pyodbc
import pytest
from sqlalchemy import (
    Integer,
    String,
    column,
    select,
    sql,
    table,
)

from dl_sqlalchemy_mssql.base import BIMSSQLDialect


class TestDialect:
    __dialect__ = BIMSSQLDialect()

    def get_cursor(self):
        return mock.Mock(spec=pyodbc.Cursor)

    @pytest.mark.parametrize(
        "parameter, output",
        [
            (None, "NULL"),
            (True, "1"),
            (False, "0"),
            (5, "5"),
            (0.5, "0.5"),
            (Decimal("0.5"), "0.5"),
            (UUID("f9dad6af-eb1f-4d23-8b30-157eda50d8cd"), "N'f9dad6af-eb1f-4d23-8b30-157eda50d8cd'"),
            ("foobar", "N'foobar'"),
            (bytearray([1, 2, 3]), "0x010203"),
            (b"foo", "'foo'"),
            (b"foo\0bar", "0x666f6f00626172"),
            (datetime.datetime(2018, 5, 31, 4, 5, 6, 7000), "{ts '2018-05-31 04:05:06.007'}"),
            (datetime.date(2018, 5, 31), "{d '2018-05-31'}"),
            ((1,), "1"),
        ],
    )
    def test_quotes_parameter_types(self, parameter, output):
        assert self.__dialect__._quote_simple_value(parameter) == output

    def test_uses_custom_compiler(self):
        table1 = table(
            "mytable",
            column("myid", Integer),
            column("name", String),
            column("description", String),
        )

        q = select([table1.c.myid, sql.literal("bar").label("c1")], order_by=[table1.c.name + "-"]).alias("foo")
        crit = q.c.myid == table1.c.myid
        compiled = select(["*"], crit).compile(dialect=self.__dialect__)

        assert compiled.construct_params() == {"param_1": "bar"}

    def test_translates_single_element_tuple(self):
        cursor = self.get_cursor()
        with pytest.raises(AssertionError):
            self.__dialect__.do_execute(cursor, "foo", ((42,),))

    def test_substitutes_params_for_group_by(self):
        cursor = self.get_cursor()
        sql = "SELECT foo + ? FROM bar GROUP BY foo + ?"
        sql_check = "SELECT foo + 9 FROM bar GROUP BY foo + 10"
        params = (9, 10)
        self.__dialect__.do_execute(cursor, sql, params)

        cursor.execute.assert_called_once_with(sql_check, [])

    @mock.patch("dl_sqlalchemy_mssql.base.LOGGER", autospec=True, spec_set=True)
    def test_error_10004_logging(self, m_log):
        cursor = self.get_cursor()
        m_execute = mock.Mock()
        m_execute.side_effect = pyodbc.OperationalError()

        with mock.patch.object(cursor, "execute", m_execute):
            sql = "SELECT count(*) AS count_1"
            params = (1, 2, 3)
            with pytest.raises(AssertionError):
                self.__dialect__.do_execute(cursor, sql, params)

lib/dl_sqlalchemy_mssql/dl_sqlalchemy_mssql_tests/unit/test_dialect.py
datalens-tech/datalens-backend
from __future__ import annotations

import sys

import oracledb


oracledb.version = "8.3.0"
sys.modules["cx_Oracle"] = oracledb

import sqlalchemy as sa  # noqa: E402
from sqlalchemy.dialects.oracle.cx_oracle import OracleDialect_cx_oracle as UPSTREAM  # noqa: E402

from dl_sqlalchemy_common.base import CompilerPrettyMixin  # noqa: E402


class BIOracleCompilerBasic(UPSTREAM.statement_compiler):
    """Necessary overrides"""


class BIOracleCompiler(BIOracleCompilerBasic, UPSTREAM.statement_compiler, CompilerPrettyMixin):
    """Added prettification"""

    def limit_clause(self, *args, **kwargs):
        # upstream: `return ""`; no prettification to be done.
        return super().limit_clause(*args, **kwargs)

    def visit_join(self, join, from_linter=None, **kwargs):
        # superclass customizes join for the `not self.dialect.use_ansi` case,
        # which isn't currently supported in the prettifier anyway.
        return super().visit_join(join, from_linter=from_linter, **kwargs)


class BIOracleDialectBasic(UPSTREAM):
    statement_compiler = BIOracleCompilerBasic


class BIOracleDialect(BIOracleDialectBasic):
    statement_compiler = BIOracleCompiler


def register_dialect():
    sa.dialects.registry.register("bi_oracle_basic", "dl_sqlalchemy_oracle.base", "BIOracleDialectBasic")
    sa.dialects.registry.register("bi_oracle", "dl_sqlalchemy_oracle.base", "BIOracleDialect")

lib/dl_sqlalchemy_oracle/dl_sqlalchemy_oracle/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest
import sqlalchemy as sa

from dl_sqlalchemy_oracle.base import register_dialect


def pytest_configure(config):  # noqa
    register_dialect()


@pytest.fixture(scope="session")
def engine_url():
    return "bi_oracle://"


@pytest.fixture(scope="session")
def sa_engine(engine_url):
    return sa.create_engine(engine_url)

lib/dl_sqlalchemy_oracle/dl_sqlalchemy_oracle_tests/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations


def test_engine(sa_engine):
    assert sa_engine
    assert sa_engine.dialect

lib/dl_sqlalchemy_oracle/dl_sqlalchemy_oracle_tests/unit/test_dialect.py
datalens-tech/datalens-backend
from __future__ import annotations

from .asyncpg import AsyncBIPGDialect
from .base import BIPGDialect


__all__ = (
    "BIPGDialect",
    "AsyncBIPGDialect",
)

lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

from sqlalchemy import util
from sqlalchemy.dialects.postgresql import (
    INTERVAL,
    OID,
    REGCLASS,
    UUID,
)

# we dont use SA asyncpg beta objects
# but it's ok to use type descriptions
from sqlalchemy.dialects.postgresql.asyncpg import (
    AsyncpgBigInteger,
    AsyncpgBoolean,
    AsyncpgDate,
    AsyncpgDateTime,
    AsyncPgEnum,
    AsyncpgInteger,
    AsyncPgInterval,
    AsyncpgJSON,
    AsyncpgJSONB,
    AsyncpgJSONIndexType,
    AsyncpgJSONIntIndexType,
    AsyncpgJSONPathType,
    AsyncpgJSONStrIndexType,
    AsyncpgNumeric,
    AsyncpgOID,
    AsyncpgREGCLASS,
    AsyncpgTime,
    AsyncpgUUID,
)
from sqlalchemy.dialects.postgresql.asyncpg import json as asynpg_json
from sqlalchemy.sql import sqltypes

from dl_sqlalchemy_postgres.base import (
    CITEXT,
    BIPGDialect,
)


# for backward compatibility with SA psycopg
class AsyncpgFloat(AsyncpgNumeric):
    def get_dbapi_type(self, dbapi):
        return dbapi.NUMBER


# types mapping for an asyncpg driver
# without any requests to a database
# we have to replace it by native SA asyncpg tools after their release
# see https://docs.sqlalchemy.org/en/14/orm/extensions/asyncio.html
class DBAPIMock:
    def __init__(self, *args, **kwargs):
        pass

    STRING = util.symbol("STRING")
    TIMESTAMP = util.symbol("TIMESTAMP")
    TIMESTAMP_W_TZ = util.symbol("TIMESTAMP_W_TZ")
    TIME = util.symbol("TIME")
    DATE = util.symbol("DATE")
    INTERVAL = util.symbol("INTERVAL")
    NUMBER = util.symbol("NUMBER")
    FLOAT = util.symbol("FLOAT")
    BOOLEAN = util.symbol("BOOLEAN")
    INTEGER = util.symbol("INTEGER")
    BIGINTEGER = util.symbol("BIGINTEGER")
    BYTES = util.symbol("BYTES")
    DECIMAL = util.symbol("DECIMAL")
    JSON = util.symbol("JSON")
    JSONB = util.symbol("JSONB")
    ENUM = util.symbol("ENUM")
    UUID = util.symbol("UUID")
    BYTEA = util.symbol("BYTEA")
    CITEXT = util.symbol("CITEXT")
    DATETIME = TIMESTAMP
    BINARY = BYTEA

    pg_types = {
        STRING: "varchar",
        TIMESTAMP: "timestamp",
        TIMESTAMP_W_TZ: "timestamp with time zone",
        DATE: "date",
        TIME: "time",
        INTERVAL: "interval",
        NUMBER: "numeric",
        FLOAT: "float",
        BOOLEAN: "bool",
        INTEGER: "integer",
        BIGINTEGER: "bigint",
        BYTES: "bytes",
        DECIMAL: "decimal",
        JSON: "json",
        JSONB: "jsonb",
        ENUM: "enum",
        UUID: "uuid",
        BYTEA: "bytea",
        CITEXT: "citext",
    }


# special dialect for asyncpg adapter
# native SA asyncpg dialect is beta now
# we can use it after release
# see https://docs.sqlalchemy.org/en/14/orm/extensions/asyncio.html
class AsyncBIPGDialect(BIPGDialect):
    # we have to use _get_set_input_sizes_lookup from compiled query
    use_setinputsizes = True
    implicit_returning = True
    supports_native_enum = True
    supports_smallserial = True  # 9.2+
    supports_sane_multi_rowcount = True  # psycopg 2.0.9+
    _has_native_hstore = True  # type: ignore  # TODO: fix
    _backslash_escapes = False

    colspecs = util.update_copy(
        BIPGDialect.colspecs,
        {
            sqltypes.Time: AsyncpgTime,
            sqltypes.Date: AsyncpgDate,
            sqltypes.DateTime: AsyncpgDateTime,
            sqltypes.Interval: AsyncPgInterval,
            INTERVAL: AsyncPgInterval,
            UUID: AsyncpgUUID,
            sqltypes.Boolean: AsyncpgBoolean,
            sqltypes.Integer: AsyncpgInteger,
            sqltypes.BigInteger: AsyncpgBigInteger,
            sqltypes.Numeric: AsyncpgNumeric,
            sqltypes.JSON: AsyncpgJSON,
            asynpg_json.JSONB: AsyncpgJSONB,
            sqltypes.JSON.JSONPathType: AsyncpgJSONPathType,
            sqltypes.JSON.JSONIndexType: AsyncpgJSONIndexType,
            sqltypes.JSON.JSONIntIndexType: AsyncpgJSONIntIndexType,
            sqltypes.JSON.JSONStrIndexType: AsyncpgJSONStrIndexType,
            sqltypes.Enum: AsyncPgEnum,
            OID: AsyncpgOID,
            REGCLASS: AsyncpgREGCLASS,
            # special override for floats
            sqltypes.Float: AsyncpgFloat,
            CITEXT: CITEXT,
        },
    )

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # special hack for getting types in compile query
        self.dbapi = DBAPIMock()

lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres/asyncpg.py
datalens-tech/datalens-backend
from __future__ import annotations

import datetime
from typing import (
    Any,
    Optional,
)

import sqlalchemy as sa
from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2 as UPSTREAM
from sqlalchemy.sql import sqltypes

from dl_sqlalchemy_common.base import CompilerPrettyMixin


class CITEXT(sqltypes.TEXT):

    """Provide the PostgreSQL CITEXT type.

    .. versionadded:: 2.0.7

    """

    __visit_name__ = "CITEXT"

    def coerce_compared_value(self, op: Optional[Any], value: Any) -> sqltypes.TypeEngine[Any]:
        return self


class BIPGCompilerBasic(UPSTREAM.statement_compiler):
    """Necessary overrides"""

    def _add_collate(self, result):
        collate = self.dialect.enforce_collate
        if collate:
            return "%s COLLATE %s" % (result, self.dialect.identifier_preparer.quote(collate))
        return result

    def visit_ilike_op_binary(self, *args, **kwargs):
        result = super().visit_ilike_op_binary(*args, **kwargs)
        return self._add_collate(result)

    def visit_notilike_op_binary(self, *args, **kwargs):
        result = super().visit_notilike_op_binary(*args, **kwargs)
        return self._add_collate(result)

    def visit_grouping(self, grouping, asfrom=False, add_grouping_collate=False, **kwargs):
        if add_grouping_collate:
            result = grouping.element._compiler_dispatch(self, **kwargs)
            result = self._add_collate(result)
            return "(%s)" % (result,)
        return super().visit_grouping(grouping, asfrom=asfrom, **kwargs)

    def _func_with_collate(self, func, **kwargs):
        # See also: sqlalchemy.sql.compiler.SQLCompiler.visit_function
        result = "%s%s" % (
            func.name,
            self.function_argspec(func, add_grouping_collate=True, **kwargs),
        )
        # To consider, instead of `visit_grouping`:
        # assert result.endswith(')')
        # return self._add_collate(result[:-1]) + ')'
        return result

    def visit_lower_func(self, func, **kwargs):
        return self._func_with_collate(func, **kwargs)

    def visit_upper_func(self, func, **kwargs):
        return self._func_with_collate(func, **kwargs)

    # allow datetime literal_binds

    def render_literal_value(self, value, type_):
        # The behavior of literal_binds=True select should match the behavior
        # of psycopg selects.
        if isinstance(type_, sa.Date):
            if isinstance(value, datetime.date):
                return "'{}'::date".format(value.isoformat())

        if isinstance(type_, sa.DateTime):
            if isinstance(value, datetime.datetime):
                # This should've been `type_.timezone` but it didn't go well
                if value.tzinfo is not None:
                    type_name = "timestamp with time zone"
                    value = value.astimezone(datetime.timezone.utc)
                    value = value.replace(tzinfo=None)
                else:
                    type_name = "timestamp"
                    if value.tzinfo is not None:
                        value = value.replace(tzinfo=None)
                return "'{}'::{}".format(value.isoformat(), type_name)

        return super().render_literal_value(value, type_)


class BIPGCompiler(BIPGCompilerBasic, UPSTREAM.statement_compiler, CompilerPrettyMixin):
    """Added prettification"""

    def limit_clause(self, select, **kw):
        sup = super().limit_clause(select, **kw)
        return self._pretty.postprocess_block("LIMIT", sup)


class BICustomPGTypeCompiler(UPSTREAM.type_compiler):
    def visit_CITEXT(self, type_, **kw):
        return "CITEXT"


bi_pg_ischema_names = UPSTREAM.ischema_names.copy()
bi_pg_ischema_names.update({"citext": CITEXT})


class BIPGDialectBasic(UPSTREAM):
    """..."""

    def __init__(self, enforce_collate=None, **kwargs):
        # Side note: due to how sqlalchemy determines dialect arguments,
        # putting `*args` right after `self` breaks the passthrough.
        self.enforce_collate = enforce_collate
        super().__init__(**kwargs)

    type_compiler = BICustomPGTypeCompiler
    statement_compiler = BIPGCompilerBasic
    ischema_names = bi_pg_ischema_names


class BIPGDialect(BIPGDialectBasic):
    statement_compiler = BIPGCompiler


def register_dialect():
    sa.dialects.registry.register("bi_postgresql_basic", "dl_sqlalchemy_postgres.base", "BIPGDialectBasic")
    sa.dialects.registry.register("bi_postgresql", "dl_sqlalchemy_postgres.base", "BIPGDialect")

lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres/base.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest
import sqlalchemy as sa

from dl_sqlalchemy_postgres.base import register_dialect


def pytest_configure(config):  # noqa
    register_dialect()


@pytest.fixture(scope="session")
def engine_url():
    return "bi_postgresql://"


@pytest.fixture(scope="session")
def sa_engine(engine_url):
    return sa.create_engine(engine_url)

lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest
import sqlalchemy as sa


def expr_to_text(expr, dialect, bind=True):
    compile_kwargs = {}
    if bind:
        compile_kwargs.update(literal_binds=True)
    expr_compiled = expr.compile(dialect=dialect, compile_kwargs=compile_kwargs)
    return str(expr_compiled)


CI_EXPRS = [
    ("ilike", sa.literal_column("'Ы'").ilike("ы").label("val")),
    ("lower", sa.func.lower(sa.literal_column("'Ы'")).label("val_lc")),
    ("upper", sa.func.upper(sa.literal_column("'Ы'")).label("val_uc")),
]


@pytest.fixture(params=CI_EXPRS, ids=[name for name, _ in CI_EXPRS])
def ci_expr(request):
    _, expr = request.param
    yield expr


@pytest.mark.param
def test_collate_option(engine_url, ci_expr):
    """
    ...

    Postgresql behavior reference (with default collate `C`):

        => SELECT 'Ы' ILIKE 'ы' AS val;
        val | f

        => SELECT 'Ы' ILIKE 'ы' COLLATE "en_US" AS val;
        val | t

        => SELECT 'Ы' ILIKE ('ы' COLLATE "en_US") AS val;
        val | t

        => SELECT 'Ы' ILIKE lower('ы' COLLATE "en_US") AS val;
        val | t

        => select lower('Ы');
        lower | Ы

        => select lower('Ы' COLLATE "en_US");
        lower | ы
    """
    ci_expr = sa.select([ci_expr])
    eng_ci = sa.create_engine(engine_url, enforce_collate="en_US")
    eng = sa.create_engine(engine_url)
    sql_ci = expr_to_text(ci_expr, eng_ci.dialect)
    sql = expr_to_text(ci_expr, eng.dialect)
    assert sql != sql_ci
    assert " COLLATE " in sql_ci
    assert " COLLATE " not in sql
    print()
    print(sql_ci)
    print(sql)

lib/dl_sqlalchemy_postgres/dl_sqlalchemy_postgres_tests/unit/test_collate.py
datalens-tech/datalens-backend
from sqlalchemy.dialects import registry
from sqlalchemy.engine import default
from sqlalchemy.sql import compiler


class PromQLCompiler(compiler.SQLCompiler):
    def visit_sequence(self, sequence, **kw):
        pass

    def visit_empty_set_expr(self, element_types):
        pass

    def update_from_clause(self, update_stmt, from_table, extra_froms, from_hints, **kw):
        pass

    def delete_extra_from_clause(self, update_stmt, from_table, extra_froms, from_hints, **kw):
        pass


class PromQLDialect(default.DefaultDialect):  # type: ignore
    name = "bi_promql"
    statement_compiler = PromQLCompiler

    @classmethod
    def dbapi(cls):
        from dl_sqlalchemy_promql import dbapi  # noqa

        return dbapi

    def get_columns(self, connection, table_name, schema=None, **kw):
        pass

    def get_primary_keys(self, connection, table_name, schema=None, **kw):
        pass

    def get_foreign_keys(self, connection, table_name, schema=None, **kw):
        pass

    def get_table_names(self, connection, schema=None, **kw):
        pass

    def get_temp_table_names(self, connection, schema=None, **kw):
        pass

    def get_view_names(self, connection, schema=None, **kw):
        pass

    def get_temp_view_names(self, connection, schema=None, **kw):
        pass

    def get_view_definition(self, connection, view_name, schema=None, **kw):
        pass

    def get_indexes(self, connection, table_name, schema=None, **kw):
        pass

    def get_unique_constraints(self, connection, table_name, schema=None, **kw):
        pass

    def get_check_constraints(self, connection, table_name, schema=None, **kw):
        pass

    def get_table_comment(self, connection, table_name, schema=None, **kw):
        pass

    def has_table(self, connection, table_name, schema=None, **kw):
        pass

    def has_sequence(self, connection, sequence_name, schema=None, *kw):
        pass

    def _get_server_version_info(self, connection):
        pass

    def _get_default_schema_name(self, connection):
        pass

    def do_begin_twophase(self, connection, xid):
        pass

    def do_prepare_twophase(self, connection, xid):
        pass

    def do_rollback_twophase(self, connection, xid, is_prepared=True, recover=False):
        pass

    def do_commit_twophase(self, connection, xid, is_prepared=True, recover=False):
        pass

    def do_recover_twophase(self, connection):
        pass

    def set_isolation_level(self, dbapi_conn, level):
        pass

    def get_isolation_level(self, dbapi_conn):
        pass

    def get_pk_constraint(self, connection, table_name, schema=None, **kw):
        pass

    def get_sequence_names(self, connection, schema=None, **kw):
        pass

    def do_set_input_sizes(self, cursor, list_of_tuples, context):
        pass


dialect = PromQLDialect


def register_dialect():
    return registry.register("bi_promql", "dl_sqlalchemy_promql.base", "PromQLDialect")

lib/dl_sqlalchemy_promql/dl_sqlalchemy_promql/base.py
datalens-tech/datalens-backend
from __future__ import annotations

from datetime import datetime
from urllib.parse import urljoin

import requests
from requests.packages.urllib3.util import Retry

from .errors import (
    InterfaceError,
    NotSupportedError,
)


def rebuild_prometheus_data(data):
    result = []
    schema = None
    for chunk in data["result"]:
        chunk_schema = [("timestamp", "unix_timestamp"), ("value", "float64")] + [
            (label, "string") for label in chunk["metric"]
        ]
        if schema is None:
            schema = chunk_schema
        elif set(schema) != set(chunk_schema):
            raise NotSupportedError("Different schemas are not supported")

        rows = [
            {**chunk["metric"], **dict(timestamp=datetime.fromtimestamp(ts), value=float(v))}
            for (ts, v) in chunk["values"]
        ]
        result.append((schema, rows))

    return result


class SyncPromQLClient:
    def __init__(self, base_url, username, password, **kwargs):
        self._closed = False
        self._session = requests.Session()
        for prefix in ("http://", "https://"):
            self._session.mount(
                prefix,
                requests.adapters.HTTPAdapter(
                    max_retries=Retry(
                        total=5,
                        backoff_factor=0.1,
                        status_forcelist=(500, 501, 502, 503, 504, 521),
                        method_whitelist=frozenset(["GET", "POST"]),
                    )
                ),
            )
        if username and password:
            auth = requests.auth.HTTPBasicAuth(username, password)
            self._session = auth(self._session)
        self._base_url = base_url

    def close(self):
        self._closed = True
        self._session.close()

    def _request(self, method, endpoint, **kwargs):
        url = urljoin(self._base_url, endpoint)
        return self._session.request(method, url, **kwargs)

    def query_range(self, query, parameters):
        req_params = {"start", "end", "step"}
        if parameters is None or not req_params <= set(parameters):
            raise InterfaceError("'step', 'start', 'end' must be in parameters")

        # TODO: query parametrization
        response = self._request(
            "post",
            "api/v1/query_range",
            data={
                "query": query,
                "start": parameters["start"],
                "end": parameters["end"],
                "step": parameters["step"],
            },
        )
        response.raise_for_status()
        data = response.json()
        result = rebuild_prometheus_data(data["data"])

        return result

    def test_connection(self):
        response = self._request("get", "-/ready")
        response.raise_for_status()

lib/dl_sqlalchemy_promql/dl_sqlalchemy_promql/cli.py
datalens-tech/datalens-backend
from __future__ import annotations

from .cli import SyncPromQLClient
from .cursor import Cursor


class Connection:
    dl_promql_cli_cls = SyncPromQLClient

    def __init__(self, host, port, username=None, password=None, protocol="http", cli_cls=None, **conn_kwargs):
        base_url = "{protocol}://{host}:{port}".format(
            protocol=protocol,
            host=host,
            port=port,
        )
        self.cli = self._create_cli(
            cli_cls=(cli_cls or self.dl_promql_cli_cls),
            base_url=base_url,
            username=username,
            password=password,
            **conn_kwargs,
        )

    def cursor(self):
        return Cursor(self)

    def commit(self):
        pass

    def rollback(self):
        pass

    def close(self):
        self.cli.close()

    @staticmethod
    def _create_cli(cli_cls, base_url, username, password, **conn_kwargs):
        return cli_cls(base_url=base_url, username=username, password=password, **conn_kwargs)

lib/dl_sqlalchemy_promql/dl_sqlalchemy_promql/connection.py
datalens-tech/datalens-backend
from __future__ import absolute_import

import itertools
import logging
from typing import (
    NamedTuple,
    Optional,
    Tuple,
)


LOGGER = logging.getLogger(__name__)


class DescriptionColumn(NamedTuple):
    # https://www.python.org/dev/peps/pep-0249/#description
    name: str
    type_code: str
    display_size: Optional[int] = None
    internal_size: Optional[int] = None
    precision: Optional[int] = None
    scale: Optional[int] = None
    null_ok: Optional[bool] = None


class Cursor:
    def __init__(self, connection):
        self.connection = connection
        self.description: Tuple[DescriptionColumn, ...] = ()
        self.arraysize = 1
        self.logger = LOGGER
        self.rows = None
        self._rows_prefetched = None

    def execute(self, query, parameters=None):
        self.logger.debug("execute promql: %s", query)
        chunks = self.connection.cli.query_range(
            query,
            parameters,
        )
        self.description = ()
        rows = self._rows_iterable(chunks)

        # Prefetch the description:
        try:
            first_row = next(rows)
        except StopIteration:
            pass
        else:
            rows = itertools.chain((first_row,), rows)

        if self.rows is not None:
            rows = itertools.chain(self.rows, rows)

        self.rows = rows

    def _rows_iterable(self, chunks_iterable):
        nmd = None
        description = None
        for chunk in chunks_iterable:
            schema, rows = chunk
            if description is None and rows:
                description = tuple(DescriptionColumn(name, type_name) for name, type_name in schema)
                self.description = description
            if nmd is None:
                assert description is not None
                colnames = tuple(col[0] for col in description)
                nmd = lambda row, colnames=colnames: tuple(row[colname] for colname in colnames)
            for row in rows:
                yield nmd(row)

    def _ensure_prefetched(self):
        if self.rows is not None and self._rows_prefetched is None:
            self._rows_prefetched = list(self.rows)
            self.rows = iter(self._rows_prefetched)
        return self._rows_prefetched

    def executemany(self, query, seq_of_parameters):
        for parameters in seq_of_parameters:
            self.execute(query, parameters)

    def fetchone(self):
        rows = self.rows
        if rows is None:
            return None
        try:
            return next(rows)
        except StopIteration:
            return None

    def fetchmany(self, size=None):
        if size is None:
            size = self.arraysize

        rows = self.rows
        if rows is None:
            return []
        return list(itertools.islice(rows, size))

    def fetchall(self):
        rows = self.rows
        if rows is None:
            return []
        return list(rows)

    def setinputsizes(self, sizes):
        pass

    def setoutputsize(self, column=None):
        pass

    def close(self):
        self.rows = None
        self._rows_prefetched = None

    @property
    def rowcount(self):
        return len(self._ensure_prefetched())

lib/dl_sqlalchemy_promql/dl_sqlalchemy_promql/cursor.py
datalens-tech/datalens-backend
from __future__ import absolute_import

from dl_sqlalchemy_promql.connection import Connection
from dl_sqlalchemy_promql.errors import (
    DatabaseError,
    DataError,
    Error,
    IntegrityError,
    InterfaceError,
    InternalError,
    NotSupportedError,
    OperationalError,
    ProgrammingError,
    Warning,
)


__all__ = (
    "connect",
    "Connection",
    "DatabaseError",
    "DataError",
    "Error",
    "IntegrityError",
    "InterfaceError",
    "InternalError",
    "NotSupportedError",
    "OperationalError",
    "ProgrammingError",
    "Warning",
)


version = "0.0.1"

apilevel = "1.0"

threadsafety = 0

paramstyle = "named"


def connect(*args, **kwargs):
    return Connection(*args, **kwargs)

lib/dl_sqlalchemy_promql/dl_sqlalchemy_promql/dbapi.py
datalens-tech/datalens-backend
from __future__ import annotations


class PromQLError(Exception):
    """..."""


class Warning(PromQLError):
    pass


class Error(PromQLError):
    pass


class InterfaceError(Error):
    pass


class DatabaseError(Error):
    pass


class DataError(DatabaseError):
    pass


class OperationalError(DatabaseError):
    pass


class IntegrityError(DatabaseError):
    pass


class InternalError(DatabaseError):
    pass


class ProgrammingError(DatabaseError):
    pass


class NotSupportedError(DatabaseError):
    pass

lib/dl_sqlalchemy_promql/dl_sqlalchemy_promql/errors.py
datalens-tech/datalens-backend
from __future__ import annotations

import pytest  # type: ignore
import sqlalchemy as sa

from dl_sqlalchemy_promql.base import register_dialect


def pytest_configure(config):  # noqa
    register_dialect()


@pytest.fixture(scope="session")
def engine_url():
    return "bi_promql://user:pass@prometheus-host.domain.net:2135?protocol=https"


@pytest.fixture(scope="session")
def sa_engine(engine_url):
    return sa.create_engine(engine_url)

lib/dl_sqlalchemy_promql/dl_sqlalchemy_promql_tests/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

from datetime import datetime

import sqlalchemy as sa

from dl_sqlalchemy_promql.cli import (
    SyncPromQLClient,
    rebuild_prometheus_data,
)
from dl_sqlalchemy_promql.errors import InterfaceError


def test_engine(sa_engine):
    conn = sa_engine.connect()
    assert conn


class MockupSyncPromQLClient(SyncPromQLClient):
    def query_range(self, query, parameters):
        req_params = {"start", "end", "step"}
        if parameters is None or not req_params <= set(parameters):
            raise InterfaceError(f"{req_params} must be in parameters")

        data = {
            "status": "success",
            "data": {
                "resultType": "matrix",
                "result": [
                    {
                        "metric": {
                            "__name__": "up",
                            "instance": "cadvisor:8080",
                            "job": "cadvisor",
                        },
                        "values": [
                            [1628445300, "1"],
                            [1628445600, "1"],
                        ],
                    },
                    {
                        "metric": {
                            "__name__": "up",
                            "instance": "localhost:9090",
                            "job": "prometheus",
                        },
                        "values": [
                            [1628445300, "1"],
                            [1628445600, "1"],
                        ],
                    },
                ],
            },
        }
        result = rebuild_prometheus_data(data["data"])
        return result


def test_result(engine_url):
    sa_engine = sa.create_engine(
        engine_url,
        connect_args=dict(
            cli_cls=MockupSyncPromQLClient,
        ),
    )
    res = sa_engine.execute(
        "up",
        {"start": "2021-08-08T17:55:00Z", "end": "2021-08-08T18:00:00Z", "step": "5m"},
    )
    description = res.cursor.description
    assert isinstance(description, tuple)
    cols = tuple((col[0], col[1]) for col in description)
    assert cols == (
        ("timestamp", "unix_timestamp"),
        ("value", "float64"),
        ("__name__", "string"),
        ("instance", "string"),
        ("job", "string"),
    )
    data = list(res)
    assert data == [
        (datetime(2021, 8, 8, 17, 55), 1.0, "up", "cadvisor:8080", "cadvisor"),
        (datetime(2021, 8, 8, 18, 0), 1.0, "up", "cadvisor:8080", "cadvisor"),
        (datetime(2021, 8, 8, 17, 55), 1.0, "up", "localhost:9090", "prometheus"),
        (datetime(2021, 8, 8, 18, 0), 1.0, "up", "localhost:9090", "prometheus"),
    ]

lib/dl_sqlalchemy_promql/dl_sqlalchemy_promql_tests/unit/test_stuff.py
datalens-tech/datalens-backend
import asyncio
import logging
from typing import (
    Any,
    Dict,
    Optional,
    Protocol,
    TypeAlias,
    cast,
    runtime_checkable,
)

from arq import (
    create_pool,
    cron,
)
from arq import ArqRedis
from arq import Retry as ArqRetry
from arq.connections import RedisSettings as ArqRedisSettings
from arq.cron import CronJob as _CronJob
import attr

from dl_configs.enums import RedisMode
from dl_task_processor.executor import (
    Executor,
    TaskInstance,
)
from dl_task_processor.task import (
    BaseTaskMeta,
    InstanceID,
    Retry,
)


EXECUTOR_KEY = "bi_executor"

LOGGER = logging.getLogger(__name__)


CronTask: TypeAlias = _CronJob


@attr.s
class ArqCronWrapper:
    _task: BaseTaskMeta = attr.ib()
    __qualname__ = "ArqCronWrapper"
    # because of asyncio.iscoroutinefunction in the arq core
    _is_coroutine = asyncio.coroutines._is_coroutine

    async def __call__(self, ctx: Dict[Any, Any], *args: Any, **kwargs: Any) -> Any:  # pragma: no cover
        return await arq_base_task(
            ctx,
            # TODO: add common object for params
            params={
                "instance_id": InstanceID.make(),
                "name": self._task.name,
                "task_params": self._task.get_params(with_name=False),
            },
        )


async def create_redis_pool(settings: ArqRedisSettings) -> ArqRedis:
    return await create_pool(settings)


@runtime_checkable
class _BIRedisSettings(Protocol):
    """
    dl_configs.settings_submodels.RedisSettings-like object
    """

    MODE: RedisMode
    HOSTS: tuple[str, ...]
    PORT: int
    PASSWORD: str
    DB: int
    CLUSTER_NAME: str
    SSL: Optional[bool]


def create_arq_redis_settings(settings: _BIRedisSettings) -> ArqRedisSettings:
    if settings.MODE == RedisMode.single_host:
        assert len(settings.HOSTS) == 1, f"Multiple hosts value {settings.HOSTS} is restricted for single_host mode"
        return ArqRedisSettings(
            host=settings.HOSTS[0],
            port=settings.PORT,
            password=settings.PASSWORD,
            database=settings.DB,
            ssl=settings.SSL,
        )
    elif settings.MODE == RedisMode.sentinel:
        redis_targets = [(host, settings.PORT) for host in settings.HOSTS]
        return ArqRedisSettings(
            host=redis_targets,
            password=settings.PASSWORD,
            sentinel=True,
            sentinel_master=settings.CLUSTER_NAME,
            database=settings.DB,
            ssl=settings.SSL,
        )
    else:
        raise ValueError(f"Unknown redis mode {settings.MODE}")


@attr.s
class CronSchedule:
    """
    second: second(s) to run the job on, 0 - 59
    minute: minute(s) to run the job on, 0 - 59
    hour: hour(s) to run the job on, 0 - 23
    day: day(s) to run the job on, 1 - 31
    weekday: week day(s) to run the job on, 0 - 6 or mon - sun
    month: month(s) to run the job on, 1 - 12
    """

    second: Optional[set[int]] = attr.ib(default=None)
    minute: Optional[set[int]] = attr.ib(default=None)
    hour: Optional[set[int]] = attr.ib(default=None)
    day: Optional[set[int]] = attr.ib(default=None)
    weekday: Optional[set[int]] = attr.ib(default=None)
    month: Optional[set[int]] = attr.ib(default=None)


def make_cron_task(task: BaseTaskMeta, schedule: CronSchedule) -> CronTask:
    cron_task = cron(
        ArqCronWrapper(task=task),
        name=task.name,
        **attr.asdict(schedule),
    )
    return cast(CronTask, cron_task)


async def arq_base_task(context: Dict, params: Dict) -> None:
    LOGGER.info("Run arq task with params %s", params)
    context[EXECUTOR_KEY]: Executor
    # transition
    # i'll delete it later
    if "task_params" in params:
        name = params["name"]
        instance_id = params["instance_id"]
        task_params = params["task_params"]
        request_id = params.get("request_id")
        task_instance = TaskInstance(
            name=name,
            params=task_params,
            instance_id=instance_id,
            request_id=request_id,
            attempt=context["job_try"] - 1,  # it starts from 1 o_O
        )
    else:
        name = params.pop("name")
        instance_id = params.pop("instance_id")
        task_instance = TaskInstance(
            name=name,
            params=params,
            instance_id=instance_id,
            attempt=context["job_try"] - 1,  # it starts from 1 o_O
        )
    job_result = await context[EXECUTOR_KEY].run_job(task_instance)
    if isinstance(job_result, Retry):
        # https://arq-docs.helpmanual.io/#retrying-jobs-and-cancellation
        raise ArqRetry(
            defer=job_result.delay * context["job_try"] * job_result.backoff,
        )

lib/dl_task_processor/dl_task_processor/arq_wrapper.py
datalens-tech/datalens-backend
import abc

import attr


@attr.s
class BaseContext(metaclass=abc.ABCMeta):
    pass


class BaseContextFabric(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    async def make(self) -> BaseContext:
        pass

    @abc.abstractmethod
    async def tear_down(self, inst: BaseContext) -> None:
        pass

lib/dl_task_processor/dl_task_processor/context.py
datalens-tech/datalens-backend
import argparse
import json
from typing import (
    Dict,
    Iterable,
    List,
    Type,
)

import attr

from dl_task_processor.processor import TaskProcessor
from dl_task_processor.task import (
    BaseTaskMeta,
    TaskInstance,
    TaskName,
    TaskRegistry,
)


@attr.s
class Controller:
    """
    Some user interface example

    ==As CLI==
    cli task info // whole list
    cli task info --name foo
    cli task blacklist --name foo // why not?
    cli instance info --name bla --id bla
    cli instance schedule --name foo --params a=1 b=2 c=3 // what about dicts or lists?
    cli instance stop --id asdasda

    ==As IPython
    c = Controller(tp, r)
    for task in c.tasks.filter():
        print(task)

    task_meta = c.tasks.meta(name='foo')
    task = task_meta(foo='bar')
    instance = await c.instances.schedule(task, start_time=datetime.datetime(...))
    scheduled_instance = await c.instances.get(name='foo')
    async for instance in c.instances.filter():
        await instance.instances.stop(scheduled_instance)

    ===As API===
    GET (and filters)
    /tasks
    GET
    /tasks/<>/info
    GET (and filters)
    POST (schedule new)
    /instances
    GET
    /instances/<>/state
    GET
    POST (change state) -- stop, reschedule
    /instances/<>/actions
    """

    _processor: TaskProcessor = attr.ib()
    _registry: TaskRegistry = attr.ib()

    @property
    def tasks(self) -> "TaskController":
        return TaskController(registry=self._registry)

    @property
    def instances(self) -> "InstanceController":
        return InstanceController(processor=self._processor)


@attr.s
class TaskController:
    _registry: TaskRegistry = attr.ib()

    def filter(self, name: TaskName) -> Iterable[Type[BaseTaskMeta]]:
        return self._registry.filter_task_meta(name)

    def meta(self, name: TaskName) -> Type[BaseTaskMeta]:
        return self._registry.get_task_meta(name)


@attr.s
class InstanceController:
    _processor: TaskProcessor = attr.ib()

    async def schedule(self, task: BaseTaskMeta) -> TaskInstance:
        scheduled_task = await self._processor.schedule(task)
        return scheduled_task

    async def get(self, name: TaskName) -> TaskInstance:
        raise NotImplementedError()

    async def stop(self, name: TaskName) -> TaskInstance:
        raise NotImplementedError()

    async def filter(self) -> Iterable[TaskInstance]:
        raise NotImplementedError()


@attr.s
class Cli:
    _processor: TaskProcessor = attr.ib()
    _registry: TaskRegistry = attr.ib()

    @classmethod
    def parse_params(cls, args: List) -> argparse.Namespace:
        parser = argparse.ArgumentParser(description="Task Processor controller")
        subparsers = parser.add_subparsers(dest="subparser")

        # task parsers
        parser_tasks = subparsers.add_parser("task")
        subparsers_tasks = parser_tasks.add_subparsers(dest="action")
        # task filter parsers
        task_info_parser = subparsers_tasks.add_parser("info", help="Show tasks with name and params")
        task_info_parser.add_argument("--name", default=None, help="Task name")
        # only defaults below
        task_info_parser.add_argument("--format", choices=["json"], default="json", help="Output format")
        task_info_parser.add_argument(
            "--fields", choices=["all"], nargs="*", default=["all"], help="List of fields or `all`"
        )

        # instance parsers
        parser_instances = subparsers.add_parser("instance")
        subparsers_instances = parser_instances.add_subparsers(dest="action")
        # instance schedule parsers
        instances_schedule_parser = subparsers_instances.add_parser("schedule", help="Schedule task")
        instances_schedule_parser.add_argument("--name", help="Task name", required=True)
        instances_schedule_parser.add_argument("--params", default=[], help="Task params", nargs="+")

        return parser.parse_args(args)

    async def run(self, args: argparse.Namespace) -> None:
        def _parse_task_params(params: Iterable) -> Dict:
            return {item.split("=")[0]: item.split("=")[1] for item in params}

        controller = Controller(registry=self._registry, processor=self._processor)
        if args.subparser == "instance":
            if args.action == "schedule":
                task = controller.tasks.meta(TaskName(args.name))
                params = _parse_task_params(args.params)
                await controller.instances.schedule(task(**params))
            else:
                raise ValueError(f"Unknown action {args.action}")
        elif args.subparser == "task":
            if args.action == "info":
                tasks = controller.tasks.filter(name=TaskName(args.name))
                for task in tasks:
                    print(
                        json.dumps(
                            {
                                "name": task.name,
                                "params": task.get_schema(),
                            },
                            sort_keys=True,
                            indent=4,
                        ),
                    )
            else:
                raise ValueError(f"Unknown action {args.action}")
        else:
            raise ValueError(f"Unknown subparser {args.subparser}")

lib/dl_task_processor/dl_task_processor/controller.py
datalens-tech/datalens-backend
import logging

import attr

from dl_app_tools.log.context import log_context
from dl_app_tools.profiling_base import GenericProfiler
from dl_task_processor.context import BaseContext
from dl_task_processor.state import TaskState
from dl_task_processor.task import (
    Fail,
    LoggerFields,
    Retry,
    RunID,
    Success,
    TaskInstance,
    TaskName,
    TaskRegistry,
    TaskResult,
)


LOGGER = logging.getLogger(__name__)


@attr.s
class Executor:
    _registry: TaskRegistry = attr.ib()
    _state: TaskState = attr.ib()
    _context: BaseContext = attr.ib()

    def _update_status(self, task: TaskInstance, task_result: TaskResult) -> None:
        if isinstance(task_result, Fail):
            LOGGER.info("Task has been failed")
            self._state.set_failed(task)
        elif isinstance(task_result, Success):
            LOGGER.info("Finish task")
            self._state.set_success(task)
        elif isinstance(task_result, Retry):
            LOGGER.info("Task needs to be retried")
            self._state.set_retry(task)
        else:
            LOGGER.warning("Unknown result type %s", task_result)

    async def run_job(self, task: TaskInstance) -> TaskResult:
        run_id = RunID.make()
        logger_extra_info = {
            LoggerFields.task_name.name: task.name,
            LoggerFields.task_params.name: task.params,
            LoggerFields.task_instance_id.name: task.instance_id.to_str(),
            LoggerFields.task_run_id.name: run_id.to_str(),
            LoggerFields.request_id.name: task.request_id,
        }
        with log_context(**logger_extra_info):
            LOGGER.info("Init task %s", task.instance_id.to_str())
            task_cls = self._registry.get_task(TaskName(task.name))
            executor_task = task_cls.from_params(
                ctx=self._context,
                params=task.params,
                instance_id=task.instance_id,
                run_id=run_id,
                request_id=task.request_id,
            )
            self._state.set_started(task)
            try:
                LOGGER.info("Run task %s", task.instance_id.to_str())
                with GenericProfiler("run-task-processor-task"):
                    task_result = await executor_task.run()
            except Exception:
                # let it be in sentry
                LOGGER.error("Task has raised an exception %s", task.instance_id.to_str(), exc_info=True)
                task_result = Fail()
            if isinstance(task_result, Retry):
                LOGGER.info(
                    "Task %s are going to be rescheduled %s from %s",
                    task.instance_id.to_str(),
                    task.attempt,
                    task_result.attempts,
                )
                if (
                    task.attempt >= task_result.attempts - 1
                ):  # the 1st attempt must be number 0 because we are programmers
                    LOGGER.info("It was the last chance for task %s", task.instance_id.to_str())
                    task_result = Fail()
            self._update_status(
                task_result=task_result,
                task=task,
            )
            return task_result


@attr.s
class ExecutorFabric:
    _registry: TaskRegistry = attr.ib()
    _state: TaskState = attr.ib()

    async def make(self, context: BaseContext) -> Executor:
        return Executor(
            registry=self._registry,
            state=self._state,
            context=context,
        )

lib/dl_task_processor/dl_task_processor/executor.py
datalens-tech/datalens-backend
import abc
import logging
from typing import Optional

import arq
import attr

from dl_task_processor.arq_wrapper import (
    ArqRedis,
    ArqRedisSettings,
    create_redis_pool,
)
from dl_task_processor.context import (
    BaseContext,
    BaseContextFabric,
)
from dl_task_processor.executor import (
    Executor,
    TaskInstance,
)
from dl_task_processor.state import (
    BITaskStateImpl,
    DummyStateImpl,
    TaskState,
)
from dl_task_processor.task import (
    BaseTaskMeta,
    InstanceID,
    Retry,
    TaskRegistry,
)
from dl_utils.aio import await_sync


LOGGER = logging.getLogger(__name__)


class BaseTaskProcessorImpl(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    async def schedule(self, task: TaskInstance) -> None:
        pass


class DummyTaskProcessorImpl(BaseTaskProcessorImpl):
    async def schedule(self, task: TaskInstance) -> None:
        pass


@attr.s
class ARQProcessorImpl(BaseTaskProcessorImpl):
    pool: arq.ArqRedis = attr.ib()

    async def schedule(self, task: TaskInstance) -> None:
        await self.pool.enqueue_job(
            "arq_base_task",
            dict(
                {
                    "name": task.name,
                    "instance_id": task.instance_id,
                    "request_id": task.request_id,
                    "task_params": task.params,
                },
            ),
        )


@attr.s
class LocalProcessorImpl(BaseTaskProcessorImpl):
    _executor: Executor = attr.ib()
    _is_sync: bool = attr.ib(default=False)

    async def schedule(self, task: TaskInstance) -> None:
        while True:
            result = await self._executor.run_job(task=task)
            if not isinstance(result, Retry):
                return
            task.attempt += 1


@attr.s
class TaskProcessor:
    _impl: BaseTaskProcessorImpl = attr.ib()
    _state: TaskState = attr.ib()
    _request_id: Optional[str] = attr.ib(default=None)

    async def schedule(self, task: BaseTaskMeta) -> TaskInstance:
        instance_id = InstanceID.make()
        task_instance = TaskInstance(
            name=task.name,
            instance_id=instance_id,
            params=task.get_params(),
            request_id=self._request_id,
        )
        LOGGER.info(
            "Schedule task %s with params %s as instance_id %s",
            task.name,
            task.get_params(),
            task_instance.instance_id.to_str(),
        )
        self._state.set_scheduled(task_instance)
        await self._impl.schedule(task_instance)
        return task_instance


def make_task_processor(redis_pool: arq.ArqRedis, request_id: Optional[str] = None) -> TaskProcessor:
    impl = ARQProcessorImpl(redis_pool)
    state = TaskState(DummyStateImpl())
    return TaskProcessor(impl=impl, state=state, request_id=request_id)


class TaskProcessorFactory(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def make(self, request_id: Optional[str] = None) -> TaskProcessor:
        pass

    def cleanup(self) -> None:
        pass

    async def cleanup_async(self) -> None:
        pass


@attr.s
class ARQTaskProcessorFactory(TaskProcessorFactory):
    _redis_pool_settings: ArqRedisSettings = attr.ib()

    _redis_pool: Optional[ArqRedis] = attr.ib(init=False, default=None)

    def setup(self) -> None:
        self._redis_pool = await_sync(create_redis_pool(self._redis_pool_settings))

    def make(self, request_id: Optional[str] = None) -> TaskProcessor:
        if self._redis_pool is None:
            self.setup()
        assert self._redis_pool is not None
        return make_task_processor(redis_pool=self._redis_pool, request_id=request_id)

    def cleanup(self) -> None:
        if self._redis_pool is not None:
            await_sync(self._redis_pool.close())

    async def cleanup_async(self) -> None:
        if self._redis_pool is not None:
            await self._redis_pool.close()


class DummyTaskProcessorFactory(TaskProcessorFactory):
    def make(self, req_id: Optional[str] = None) -> TaskProcessor:
        impl = DummyTaskProcessorImpl()
        state = TaskState(DummyStateImpl())
        return TaskProcessor(impl=impl, state=state, request_id=req_id)


@attr.s
class LocalTaskProcessorFactory(TaskProcessorFactory):
    _context_fab: BaseContextFabric = attr.ib()
    _registry: TaskRegistry = attr.ib()

    _context: Optional[BaseContext] = attr.ib(init=False, default=None)

    def setup(self) -> None:
        self._context = await_sync(self._context_fab.make())

    def make(self, request_id: Optional[str] = None) -> TaskProcessor:
        if self._context is None:
            self.setup()
        assert self._context is not None
        task_state = TaskState(BITaskStateImpl())
        executor = Executor(context=self._context, state=task_state, registry=self._registry)
        impl = LocalProcessorImpl(executor)
        processor = TaskProcessor(impl=impl, state=task_state, request_id=request_id)

        LOGGER.info("Local TP is ready")
        return processor

    def cleanup(self) -> None:
        if self._context is not None:
            await_sync(self._context_fab.tear_down(self._context))

    async def cleanup_async(self) -> None:
        if self._context is not None:
            await self._context_fab.tear_down(self._context)

lib/dl_task_processor/dl_task_processor/processor.py
datalens-tech/datalens-backend
import abc
import asyncio
from collections import defaultdict

import attr

from dl_task_processor.task import (
    InstanceID,
    TaskInstance,
)


# i will change it later
class BaseTaskStateImpl(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def set_state(self, task: TaskInstance, state: str) -> None:
        pass

    @abc.abstractmethod
    def get_state(self, task: TaskInstance) -> list:
        pass


class DummyStateImpl(BaseTaskStateImpl):
    def set_state(self, task: TaskInstance, state: str) -> None:
        pass

    def get_state(self, task: TaskInstance) -> list:
        pass


# i will change it later
@attr.s
class BITaskStateImpl(BaseTaskStateImpl):
    _states: dict[InstanceID, list[str]] = attr.ib(default=defaultdict(list))

    def set_state(self, task: TaskInstance, state: str) -> None:
        self._states[task.instance_id].append(state)

    def get_state(self, task: TaskInstance) -> list:
        return self._states[task.instance_id]


# i will change it later
@attr.s
class TaskState:
    _impl: BaseTaskStateImpl = attr.ib()

    def set_scheduled(self, task: TaskInstance) -> None:
        self._impl.set_state(task, "scheduled")

    def set_started(self, task: TaskInstance) -> None:
        self._impl.set_state(task, "started")

    def set_failed(self, task: TaskInstance) -> None:
        self._impl.set_state(task, "failed")

    def set_retry(self, task: TaskInstance) -> None:
        self._impl.set_state(task, "retry")

    def set_aborted(self, task: TaskInstance) -> None:
        self._impl.set_state(task, "aborted")

    def set_success(self, task: TaskInstance) -> None:
        self._impl.set_state(task, "success")

    def get_state(self, task: TaskInstance) -> list:
        return self._impl.get_state(task)


# temporary task result checker for arq
# i will change it later
# maybe move to TaskState
async def wait_task(task: TaskInstance, state: TaskState, timeout: float = 10, interval: float = 0.5) -> list:
    """
    timeout == the bottom line
    """
    spent_time = 0
    while spent_time < timeout:
        current_state = state.get_state(task)
        # Has the task reached the final state?
        if {"success", "failed"} & set(state.get_state(task)):
            return current_state
        await asyncio.sleep(interval)
        spent_time += interval
    raise RuntimeError("Task was not complete in time")

lib/dl_task_processor/dl_task_processor/state.py
datalens-tech/datalens-backend
import abc
from collections.abc import Iterable
import enum
from typing import (
    ClassVar,
    Generic,
    NewType,
    Optional,
    Type,
    TypeVar,
)
import uuid

import attr

from dl_task_processor.context import BaseContext


TaskName = NewType("TaskName", str)


_BASE_ID_TYPE = TypeVar("_BASE_ID_TYPE", bound="_BaseID")


@attr.s(frozen=True)
class _BaseID:
    _id: str = attr.ib()

    def to_str(self) -> str:
        return self._id

    def __str__(self) -> str:
        return self.to_str()

    @classmethod
    def make(cls: Type[_BASE_ID_TYPE]) -> _BASE_ID_TYPE:
        return cls(uuid.uuid4().hex)


@attr.s
class InstanceID(_BaseID):
    pass


@attr.s
class RunID(_BaseID):
    pass


@attr.s
class TaskInstance:
    instance_id: InstanceID = attr.ib()
    name: TaskName = attr.ib()
    params: dict = attr.ib()
    attempt: int = attr.ib(default=0)
    request_id: Optional[str] = attr.ib(default=None)


class BaseTaskMeta(metaclass=abc.ABCMeta):
    name: ClassVar[TaskName]

    def __init__(self, *args, **kwargs) -> None:
        # lets trick typing
        pass

    def get_params(self, with_name: bool = False) -> dict:
        if with_name:
            return dict(
                {"name": self.name},
                **attr.asdict(self),
            )
        return attr.asdict(self)

    def make_id(self) -> str:
        params = sorted(attr.asdict(self).items())
        serialized_params = "||".join(f"{k}|{v}" for k, v in params)
        return f"{self.name}||{serialized_params}"

    # it's temporary solution
    # i'm gonna use datalens/backend/app/bi_external_api/bi_external_api/attrs_model_mapper
    # but we have to extract it to the top level
    # now we can use it only in cli
    @classmethod
    def get_schema(cls) -> Iterable[dict]:
        attr_schema = attr.fields(cls)
        return [
            {
                "name": field.name,
                "type": field.type.__name__,
            }
            for field in attr_schema
        ]


_BASE_TASK_META_TV = TypeVar("_BASE_TASK_META_TV", bound=BaseTaskMeta)
_BASE_TASK_CONTEXT_TV = TypeVar("_BASE_TASK_CONTEXT_TV", bound=BaseContext)


@attr.s(frozen=True, eq=True)
class TaskResult(metaclass=abc.ABCMeta):
    pass


class Success(TaskResult):
    pass


# system entity, you can just raise exception
class Fail(TaskResult):
    pass


@attr.s
class Retry(TaskResult):
    delay: int = attr.ib(default=60)  # seconds
    backoff: float = attr.ib(default=2.0)
    attempts: int = attr.ib(default=5)


@attr.s
class BaseExecutorTask(Generic[_BASE_TASK_META_TV, _BASE_TASK_CONTEXT_TV], metaclass=abc.ABCMeta):
    cls_meta: ClassVar[Type[_BASE_TASK_META_TV]]
    meta: _BASE_TASK_META_TV = attr.ib()
    _ctx: _BASE_TASK_CONTEXT_TV = attr.ib()
    _instance_id: InstanceID = attr.ib()
    _run_id: RunID = attr.ib()
    _request_id: Optional[str] = attr.ib(default=None)

    @classmethod
    def from_params(
        cls,
        instance_id: InstanceID,
        run_id: RunID,
        ctx: _BASE_TASK_CONTEXT_TV,
        params: dict,
        request_id: Optional[str] = None,
    ) -> "BaseExecutorTask":
        return cls(
            meta=cls.cls_meta(**params),
            ctx=ctx,
            instance_id=instance_id,
            run_id=run_id,
            request_id=request_id,
        )

    @classmethod
    def name(cls) -> TaskName:
        return cls.cls_meta.name

    @abc.abstractmethod
    async def run(self) -> TaskResult:
        pass


@attr.s
class TaskRegistry:
    _tasks: dict[TaskName, BaseExecutorTask] = attr.ib()

    @classmethod
    def create(cls, tasks: Iterable[Type[BaseExecutorTask]]) -> "TaskRegistry":
        assert sorted(
            [t.name() for t in tasks],
        ) == sorted(list(set([t.name() for t in tasks]))), "Some tasks has the same name"
        return cls(tasks={task.name(): task for task in tasks})

    def get_task(self, name: TaskName) -> BaseExecutorTask:
        return self._tasks[name]

    def get_task_meta(self, name: TaskName) -> Type[BaseTaskMeta]:
        return self.get_task(name).cls_meta

    def filter_task_meta(self, name: TaskName) -> Iterable[Type[BaseTaskMeta]]:
        return [task.cls_meta for task in self._tasks.values() if not name or name == task.name()]


class LoggerFields(enum.Enum):
    request_id = enum.auto()
    task_name = enum.auto()
    task_instance_id = enum.auto()
    task_run_id = enum.auto()
    task_params = enum.auto()

lib/dl_task_processor/dl_task_processor/task.py
datalens-tech/datalens-backend
import asyncio
from datetime import timedelta
import socket
import time
from typing import (
    Dict,
    Iterable,
    Optional,
    Union,
)

from arq import Worker as _ArqWorker
from arq.connections import RedisSettings
from arq.worker import async_check_health
import attr

from dl_task_processor.arq_wrapper import (
    EXECUTOR_KEY,
    CronTask,
    arq_base_task,
    create_redis_pool,
)
from dl_task_processor.context import BaseContextFabric
from dl_task_processor.executor import ExecutorFabric


CONTEXT_KEY = "bi_context"


@attr.s
class WorkerSettings:
    # we should not allow forever-fail tasks because it can stop the whole system
    # but (if you really want it) you can provide float('inf')
    retry_hard_limit: Union[int, float] = attr.ib(default=100)
    job_timeout: int = attr.ib(default=600)  # seconds
    health_check_interval: int = attr.ib(default=30)
    health_check_suffix: str = attr.ib(default="bihealthcheck")


@attr.s
class ArqWorker:
    _executor_fab: ExecutorFabric = attr.ib()
    _context_fab: BaseContextFabric = attr.ib()
    _redis_settings: RedisSettings = attr.ib()
    _worker_settings: WorkerSettings = attr.ib()
    _arq_worker: _ArqWorker = attr.ib(default=None)
    _cron_tasks: Iterable[CronTask] = attr.ib(default=[])

    @property
    def health_check_key(self) -> str:
        return f"{socket.gethostname()}::{self._worker_settings.health_check_suffix}"

    async def start(self) -> None:
        redis_pool = await create_redis_pool(self._redis_settings)
        self._arq_worker = _ArqWorker(
            # let's trick strange typing in arq
            # everybody does it o_O
            **{
                "functions": [arq_base_task],
                "on_startup": self.start_executor,
                "on_shutdown": self.stop_executor,
                "max_tries": self._worker_settings.retry_hard_limit,
                "job_timeout": timedelta(seconds=self._worker_settings.job_timeout),
                "retry_jobs": True,
                "health_check_key": self.health_check_key,
                "handle_signals": False,
                "health_check_interval": self._worker_settings.health_check_interval,
                "cron_jobs": self._cron_tasks,
            },
            redis_pool=redis_pool,
        )
        await self._arq_worker.main()

    async def check_health(self) -> bool:
        return bool(not await async_check_health(self._redis_settings, health_check_key=self.health_check_key))

    # param's name is important
    async def start_executor(self, ctx: Dict) -> None:
        context = await self._context_fab.make()
        executor = await self._executor_fab.make(
            context=context,
        )
        ctx[EXECUTOR_KEY] = executor
        ctx[CONTEXT_KEY] = context

    # param's name is important
    async def stop_executor(self, ctx: Dict) -> None:
        assert EXECUTOR_KEY in ctx, "Arq worker has not been run"
        assert CONTEXT_KEY in ctx, "Arq worker has not been run"
        await self._context_fab.tear_down(ctx[CONTEXT_KEY])

    async def stop(self) -> None:
        if self._arq_worker is not None:
            await self._arq_worker.close()


@attr.s
class HealthChecker:
    _worker: ArqWorker = attr.ib()

    async def check(self) -> None:
        # more info is in logs (==stdout)
        try:
            result = await self._worker.check_health()
        except Exception as e:
            raise SystemExit(1) from e
        if result:
            exit(0)
        exit(1)

    async def check_and_raise(self) -> None:
        result = await self._worker.check_health()
        if not result:
            raise ValueError("Health check unsuccessful")

    async def wait_for_ok(self, timeout: Optional[int] = None, cooldown: int = 1) -> None:
        # you should use it only in tests
        start = time.monotonic()
        while (timeout is None) or (time.monotonic() - start < timeout):
            result = await self._worker.check_health()
            if result:
                return
            await asyncio.sleep(cooldown)
        raise RuntimeError(f"Worker did not start in {timeout} seconds")


@attr.s
class ArqWorkerTestWrapper:
    _worker: ArqWorker = attr.ib()
    _loop: asyncio.AbstractEventLoop = attr.ib()
    _task: asyncio.Task = attr.ib(default=None)

    async def start(self) -> ArqWorker:
        self._task = self._loop.create_task(self._worker.start())
        checker = HealthChecker(self._worker)
        await checker.wait_for_ok(20)
        return self._worker

    async def stop(self) -> None:
        assert self._task is not None, "Arq worker has not been run"
        try:
            await self._worker.stop()
        except asyncio.exceptions.CancelledError:
            pass
        try:
            await asyncio.gather(self._task)
        except:
            pass

lib/dl_task_processor/dl_task_processor/worker.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import logging

import pytest

from dl_task_processor.executor import Executor
from dl_task_processor.processor import (
    LocalProcessorImpl,
    TaskProcessor,
)
from dl_task_processor.state import (
    BITaskStateImpl,
    TaskState,
)
from dl_task_processor_tests.utils import (
    REGISTRY,
    LocalContextFab,
)


LOGGER = logging.getLogger(__name__)


@pytest.fixture(autouse=True)
def loop(event_loop):
    """
    Preventing creation of new loop by `aiohttp.pytest_plugin` loop fixture in favor of pytest-asyncio one
    And set loop pytest-asyncio created loop as default for thread
    """
    asyncio.set_event_loop(event_loop)
    return event_loop


@pytest.fixture(scope="function")
def task_state():
    return TaskState(BITaskStateImpl())


@pytest.fixture(scope="function")
async def task_processor_local_client(loop, task_state):
    context = await LocalContextFab().make()
    executor = Executor(context=context, state=task_state, registry=REGISTRY)
    impl = LocalProcessorImpl(executor)
    processor = TaskProcessor(impl=impl, state=task_state)
    LOGGER.info("Local TP is ready")
    yield processor
    await LocalContextFab().tear_down(context)

lib/dl_task_processor/dl_task_processor_tests/conftest.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging

import pytest

from dl_configs.enums import RedisMode
from dl_configs.settings_submodels import RedisSettings
from dl_task_processor.arq_wrapper import (
    create_arq_redis_settings,
    create_redis_pool,
)
from dl_task_processor.executor import ExecutorFabric
from dl_task_processor.processor import (
    ARQProcessorImpl,
    TaskProcessor,
)
from dl_task_processor.state import (
    BITaskStateImpl,
    TaskState,
)
from dl_task_processor.worker import (
    ArqWorker,
    ArqWorkerTestWrapper,
    WorkerSettings,
)
from dl_task_processor_tests.utils import (
    REGISTRY,
    ARQContextFab,
)
from dl_testing.containers import get_test_container_hostport


LOGGER = logging.getLogger(__name__)


@pytest.fixture(scope="function")
def redis_settings(request):
    # try to redefine value
    if hasattr(request, "param") and request.param is not None:
        return request.param
    return RedisSettings(
        MODE=RedisMode.single_host,
        CLUSTER_NAME="",
        HOSTS=(get_test_container_hostport("redis", fallback_port=51219).host,),
        PORT=get_test_container_hostport("redis", fallback_port=51219).port,
        DB=10,
    )


@pytest.fixture(scope="function")
async def redis_pool(loop, redis_settings):
    pool = await create_redis_pool(
        create_arq_redis_settings(redis_settings),
    )
    yield pool
    await pool.close()


@pytest.fixture(scope="function")
def task_state():
    return TaskState(BITaskStateImpl())


@pytest.fixture(scope="function")
def worker_settings():
    return WorkerSettings()


@pytest.fixture(scope="function")
async def init_task_processor_arq_worker(task_state, redis_settings, worker_settings):
    arq_redis = create_arq_redis_settings(redis_settings)
    worker = ArqWorker(
        redis_settings=arq_redis,
        executor_fab=ExecutorFabric(
            registry=REGISTRY,
            state=task_state,
        ),
        context_fab=ARQContextFab(arq_redis),
        worker_settings=worker_settings,
    )
    yield worker


@pytest.fixture(scope="function")
async def task_processor_arq_worker(loop, init_task_processor_arq_worker):
    wrapper = ArqWorkerTestWrapper(loop=loop, worker=init_task_processor_arq_worker)
    yield await wrapper.start()
    await wrapper.stop()


@pytest.fixture(scope="function")
def task_processor_arq_client(loop, task_processor_arq_worker, redis_pool, redis_settings, task_state):
    impl = ARQProcessorImpl(redis_pool)
    p = TaskProcessor(impl=impl, state=task_state)
    LOGGER.info("Arq TP is ready")
    return p


@pytest.fixture(scope="function", params=["arq", "local"])
def task_processor_client(request, task_processor_arq_client, task_processor_local_client):
    if request.param == "arq":
        yield task_processor_arq_client
    elif request.param == "local":
        yield task_processor_local_client

lib/dl_task_processor/dl_task_processor_tests/db/conftest.py
datalens-tech/datalens-backend
import asyncio

import pytest

from dl_configs.enums import RedisMode
from dl_configs.settings_submodels import RedisSettings
from dl_task_processor.worker import (
    HealthChecker,
    WorkerSettings,
)


def _get_broken_redis_settings():
    return RedisSettings(
        MODE=RedisMode.single_host,
        CLUSTER_NAME="",
        HOSTS=("127.0.0.1",),
        PORT=1010,
        DB=10,
    )


@pytest.mark.asyncio
@pytest.mark.parametrize("worker_settings", [WorkerSettings(health_check_suffix="bla")], indirect=True)
async def test_health_check_success(task_processor_arq_worker, worker_settings):
    # lets wait for the first task pool iteration
    # i really dont know how i can do it better
    await asyncio.sleep(10)
    with pytest.raises(SystemExit) as exc:
        await HealthChecker(worker=task_processor_arq_worker).check()
    assert exc.value.code == 0


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("redis_settings", "worker_settings"),
    [
        (_get_broken_redis_settings(), WorkerSettings(health_check_suffix="notbla")),
    ],
    indirect=True,
)
async def test_health_check_fail(init_task_processor_arq_worker, redis_settings, worker_settings):
    with pytest.raises(SystemExit) as exc:
        await HealthChecker(worker=init_task_processor_arq_worker).check()
    assert exc.value.code == 1

lib/dl_task_processor/dl_task_processor_tests/db/test_health_check.py
datalens-tech/datalens-backend
import logging

import pytest

from dl_task_processor.state import wait_task
from dl_task_processor_tests.utils import (
    BrokenTaskInterface,
    RetryTaskInterface,
    ScheduleFromTaskTaskInterface,
    SomeTaskInterface,
    TestIdsTaskInterface,
)


LOGGER = logging.getLogger(__name__)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("task_meta", "result"),
    [
        (SomeTaskInterface(foo="value1"), ["scheduled", "started", "success"]),
        (RetryTaskInterface(foobar="value2"), ["scheduled", "started", "retry", "started", "failed"]),
        (BrokenTaskInterface(bar="value3"), ["scheduled", "started", "failed"]),
    ],
)
async def test_schedule(task_processor_client, task_state, task_meta, result):
    task = await task_processor_client.schedule(task_meta)
    assert await wait_task(task, task_state) == result


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "request_id",
    [
        "SOME_REQUEST_ID",
        None,
    ],
)
async def test_task_ids(task_processor_client, task_state, request_id):
    # lets check defaults too
    if request_id:
        task_processor_client._request_id = request_id
    task = await task_processor_client.schedule(TestIdsTaskInterface(expected_request_id=request_id))
    # task should selfcheck internal ids
    # like req_id, instance_id and run_id
    assert await wait_task(task, task_state) == ["scheduled", "started", "success"]


@pytest.mark.asyncio
async def test_schedule_task_by_task(task_processor_client, task_state):
    task = await task_processor_client.schedule(ScheduleFromTaskTaskInterface())
    assert await wait_task(task, task_state) == ["scheduled", "started", "success"]

lib/dl_task_processor/dl_task_processor_tests/db/test_schedule_task.py
datalens-tech/datalens-backend
from statcommons.logs import LOGMUTATORS

from dl_api_commons.logging_config import add_log_context


def pytest_configure(config):  # noqa
    LOGMUTATORS.apply(require=False)
    LOGMUTATORS.add_mutator("log_context", add_log_context)

lib/dl_task_processor/dl_task_processor_tests/unit/conftest.py
datalens-tech/datalens-backend
import logging

import pytest

from dl_task_processor.executor import Executor
from dl_task_processor.task import (
    Fail,
    InstanceID,
    LoggerFields,
    Retry,
    Success,
    TaskInstance,
)
from dl_task_processor_tests.utils import (
    BROKEN_MARK,
    REGISTRY,
    BrokenTaskInterface,
    LocalContextFab,
    RetryTaskInterface,
    SomeTaskInterface,
)


LOGGER = logging.getLogger(__name__)
DEFAULT_REQ_ID = "reqid-123"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("task_meta", "result"),
    [
        (SomeTaskInterface(foo="as"), Success()),
        (RetryTaskInterface(foobar="asd"), Retry(delay=1, backoff=1, attempts=2)),
        (BrokenTaskInterface(bar="as"), Fail()),
    ],
)
async def test_run_task(task_state, task_meta, result):
    context = await LocalContextFab().make()
    executor = Executor(
        context=context,
        registry=REGISTRY,
        state=task_state,
    )
    task = TaskInstance(
        instance_id=InstanceID.make(),
        name=task_meta.name,
        params=task_meta.get_params(),
    )
    task_result = await executor.run_job(task)
    assert task_result == result


@pytest.mark.asyncio
async def test_run_task_logs(task_state, caplog):
    """
    test messages from internal task logger
    """
    caplog.set_level(logging.INFO)
    context = await LocalContextFab().make()
    executor = Executor(
        context=context,
        registry=REGISTRY,
        state=task_state,
    )
    task_meta = BrokenTaskInterface(bar="21")
    task = TaskInstance(
        instance_id=InstanceID.make(),
        name=task_meta.name,
        params=task_meta.get_params(),
        request_id=DEFAULT_REQ_ID,
    )
    caplog.clear()
    task_result = await executor.run_job(task)
    assert task_result == Fail()
    internal_task_message = next(rec for rec in caplog.records if rec.message == BROKEN_MARK)
    assert getattr(internal_task_message, LoggerFields.task_name.name) == task.name
    assert getattr(internal_task_message, LoggerFields.task_instance_id.name) == task.instance_id.to_str()
    assert getattr(internal_task_message, LoggerFields.task_params.name) == task.params
    assert getattr(internal_task_message, LoggerFields.task_run_id.name)
    assert getattr(internal_task_message, LoggerFields.request_id.name) == DEFAULT_REQ_ID
    assert isinstance(getattr(internal_task_message, LoggerFields.task_run_id.name), str)

lib/dl_task_processor/dl_task_processor_tests/unit/test_executor.py
datalens-tech/datalens-backend
import asyncio
import logging
from typing import Optional

import arq
import attr

from dl_configs.settings_submodels import RedisSettings
from dl_task_processor.arq_wrapper import create_redis_pool
from dl_task_processor.context import (
    BaseContext,
    BaseContextFabric,
)
from dl_task_processor.executor import Executor
from dl_task_processor.processor import (
    LocalProcessorImpl,
    TaskProcessor,
    make_task_processor,
)
from dl_task_processor.state import (
    DummyStateImpl,
    TaskState,
)
from dl_task_processor.task import (
    BaseExecutorTask,
    BaseTaskMeta,
    Retry,
    Success,
    TaskName,
    TaskRegistry,
    TaskResult,
)
from dl_utils.aio import ContextVarExecutor


LOGGER = logging.getLogger(__name__)

BROKEN_MARK = "mark_broken_task"


@attr.s
class Context(BaseContext):
    tpe: ContextVarExecutor = attr.ib()
    tp: TaskProcessor = attr.ib()
    _redis_pool: Optional[arq.ArqRedis] = attr.ib(default=None)


@attr.s
class DummyContext(BaseContext):
    tpe: ContextVarExecutor = attr.ib()


class LocalContextFab(BaseContextFabric):
    async def make(self) -> Context:
        state = TaskState(DummyStateImpl())
        # yep, we create another executor
        executor = Executor(context=DummyContext(tpe=ContextVarExecutor()), state=state, registry=REGISTRY)
        impl = LocalProcessorImpl(executor)
        processor = TaskProcessor(impl=impl, state=state)
        return Context(
            tpe=ContextVarExecutor(),
            tp=processor,
        )

    async def tear_down(self, inst: Context):
        inst.tpe.shutdown()


@attr.s
class ARQContextFab(BaseContextFabric):
    _redis_settings: RedisSettings = attr.ib()

    async def make(self) -> Context:
        redis_pool = await create_redis_pool(self._redis_settings)
        return Context(
            tpe=ContextVarExecutor(),
            tp=make_task_processor(
                redis_pool=redis_pool,
                request_id=None,
            ),
            redis_pool=redis_pool,
        )

    async def tear_down(self, inst: Context):
        inst.tpe.shutdown()
        if inst._redis_pool is not None:
            await inst._redis_pool.close()


@attr.s
class SomeTaskInterface(BaseTaskMeta):
    name = TaskName("some_task")
    foo: str = attr.ib()


@attr.s
class BrokenTaskInterface(BaseTaskMeta):
    name = TaskName("broken_task")
    bar: str = attr.ib()


@attr.s
class RetryTaskInterface(BaseTaskMeta):
    name = TaskName("retry_task")
    foobar: str = attr.ib()


@attr.s
class TestIdsTaskInterface(BaseTaskMeta):
    name = TaskName("test_ids_task")
    expected_request_id: str = attr.ib()


@attr.s
class ScheduleFromTaskTaskInterface(BaseTaskMeta):
    name = TaskName("test_schedule_from_task_task")


@attr.s
class SomeTask(BaseExecutorTask[SomeTaskInterface, Context]):
    cls_meta = SomeTaskInterface

    async def run(self) -> TaskResult:
        LOGGER.info(f"Its some task with a {self._ctx} {self.meta.foo}")
        loop = asyncio.get_running_loop()
        try:
            result = await loop.run_in_executor(
                self._ctx.tpe,
                lambda t: ", ".join((str(i) for i in range(t))),
                10,
            )
        except Exception as ex:
            LOGGER.exception(ex)
            raise
        LOGGER.info(result)
        return Success()


def some_sync_function_with_logs() -> None:
    LOGGER.info(f"{BROKEN_MARK}")


@attr.s
class BrokenTask(BaseExecutorTask[BrokenTaskInterface, Context]):
    cls_meta = BrokenTaskInterface

    async def run(self) -> TaskResult:
        loop = asyncio.get_running_loop()
        loop.run_in_executor(self._ctx.tpe, some_sync_function_with_logs)
        # let's hack checkers
        if int(1) != 1:
            return Success()
        raise Exception()


@attr.s
class TestIdsTask(BaseExecutorTask[TestIdsTaskInterface, Context]):
    cls_meta = TestIdsTaskInterface

    async def run(self) -> TaskResult:
        assert self._request_id == self.meta.expected_request_id
        assert self._instance_id
        assert self._run_id
        return Success()


@attr.s
class RetryTask(BaseExecutorTask[RetryTaskInterface, Context]):
    cls_meta = RetryTaskInterface

    async def run(self) -> TaskResult:
        LOGGER.info(f"Its retry task with a {self._ctx} {self.meta.foobar}")
        return Retry(
            attempts=2,
            backoff=1,
            delay=1,
        )


class ScheduleFromTaskTask(BaseExecutorTask[ScheduleFromTaskTaskInterface, Context]):
    cls_meta = ScheduleFromTaskTaskInterface

    async def run(self) -> TaskResult:
        await self._ctx.tp.schedule(SomeTaskInterface(foo="1"))
        return Success()


REGISTRY: TaskRegistry = TaskRegistry.create(
    [
        SomeTask,
        BrokenTask,
        RetryTask,
        TestIdsTask,
        ScheduleFromTaskTask,
    ]
)

lib/dl_task_processor/dl_task_processor_tests/utils.py
datalens-tech/datalens-backend
import inspect
import os
import re
from typing import Optional

import attr
import yaml


_CONTAINER_HOST_FROM_COMPOSE_MARKER = "from_compose"
_APPLIB_PATH_RE = re.compile(r".+\/(?:app|lib|site-packages)\/[^\/]+\/")


@attr.s(frozen=True)
class HostPort:
    host: str = attr.ib()
    port: int = attr.ib()

    def as_pair(self) -> str:
        return f"{self.host}:{self.port}"


def _get_docker_compose_file_path(filename: Optional[str] = None) -> str | None:
    filename = filename or "docker-compose.yml"
    assert filename is not None
    base_filename: str
    stack = inspect.stack()
    for frame in stack:
        if "/dl_testing/" not in frame.filename:
            base_filename = frame.filename
            break
    else:
        raise Exception("No test dir found in stack")

    re_match = _APPLIB_PATH_RE.search(base_filename)
    if not re_match:
        return None
    applib_path = re_match.group(0)
    return os.path.join(applib_path, filename)


def get_test_container_hostport(
    service_key: str,
    original_port: Optional[int] = None,
    fallback_port: Optional[int] = None,  # likely redundant
    dc_filename: Optional[str] = None,
) -> HostPort:
    # TODO: Turn this into a method of a class with dc_filename in self
    host: str
    port: int

    file_path = _get_docker_compose_file_path(filename=dc_filename)
    if not file_path and fallback_port is not None:
        return HostPort(host="127.0.0.1", port=fallback_port)
    else:
        try:
            with open(file_path) as dcyml:
                docker_compose_yml = yaml.load(dcyml)
        except FileNotFoundError:
            if fallback_port is not None:
                return HostPort(host="127.0.0.1", port=fallback_port)
            else:
                raise

    ports = docker_compose_yml["services"][service_key]["ports"]
    if original_port is not None:
        port_pair = [pp.split(":") for pp in ports if int(pp.split(":")[1]) == original_port][0]
    else:
        port_pair = ports[0].split(":")

    if (container_host := os.environ.get("TEST_CONTAINER_HOST")) is not None:
        if container_host == _CONTAINER_HOST_FROM_COMPOSE_MARKER:
            # github actions env, taking host from docker-compose and original port
            host = service_key
            port = int(port_pair[1])
        else:
            # local test run, with remote docker containers
            host = container_host
            port = int(port_pair[0])
    else:
        # remote test run, with containers on the same host
        host = "127.0.0.1"
        port = int(port_pair[0])

    return HostPort(host=host, port=port)

lib/dl_testing/dl_testing/containers.py
datalens-tech/datalens-backend
from __future__ import annotations

import attr
import yaml

from dl_testing.env_params.getter import EnvParamGetter
from dl_testing.env_params.loader import EnvParamGetterLoader


@attr.s
class GenericEnvParamGetter(EnvParamGetter):
    _loader: EnvParamGetterLoader = attr.ib(init=False, factory=EnvParamGetterLoader)
    _key_mapping: dict[str, tuple[str, str]] = attr.ib(init=False, factory=dict)  # key -> (getter_name, remapped_key)

    def get_str_value(self, key: str) -> str:
        getter_name, remapped_key = self._key_mapping[key]
        getter = self._loader.get_getter(getter_name)
        value = getter.get_str_value(remapped_key)
        return value

    def initialize(self, config: dict) -> None:
        param_config = config["params"]
        for key, key_config in param_config.items():
            self._key_mapping[key] = (key_config["getter"], key_config["key"])

        self._loader.initialize(config.get("getters", []), requirement_getter=self)

    @classmethod
    def from_yaml_file(cls, path: str) -> GenericEnvParamGetter:
        with open(path) as config_file:
            config_data = yaml.safe_load(config_file)

        getter = cls()
        getter.initialize(config_data)
        return getter

lib/dl_testing/dl_testing/env_params/generic.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import json

import yaml


class EnvParamGetter(abc.ABC):
    @abc.abstractmethod
    def get_str_value(self, key: str) -> str:
        raise NotImplementedError

    def get_int_value(self, key: str) -> int:
        str_value = self.get_str_value(key)
        return int(str_value)

    def get_json_value(self, key: str) -> dict:
        str_value = self.get_str_value(key)
        return json.loads(str_value)

    def get_yaml_value(self, key: str) -> dict:
        str_value = self.get_str_value(key)
        return yaml.safe_load(str_value)

    def initialize(self, config: dict) -> None:
        pass

lib/dl_testing/dl_testing/env_params/getter.py
datalens-tech/datalens-backend
from typing import (
    ClassVar,
    Mapping,
    Sequence,
)

import attr

from dl_testing.env_params.getter import EnvParamGetter
from dl_utils.entrypoints import EntrypointClassManager


@attr.s
class EnvParamGetterLoader:
    _getters: dict[str, EnvParamGetter] = attr.ib(init=False, factory=dict)
    _ep_mgr: EntrypointClassManager = attr.ib(init=False)

    EP_NAME: ClassVar[str] = "env_param_getters"

    @_ep_mgr.default
    def _make_ep_mgr(self) -> EntrypointClassManager:
        return EntrypointClassManager(entrypoint_group_name=self.EP_NAME)

    def get_getter(self, name: str) -> EnvParamGetter:
        try:
            return self._getters[name]
        except KeyError:
            if name.startswith("$"):
                self._auto_add_getter(name)
                return self.get_getter(name)

    def has_getter(self, name: str) -> bool:
        return name in self._getters

    def _auto_add_getter(self, name: str) -> None:
        assert name.startswith("$")
        getter_type_name = name[1:]
        getter_cls = self._ep_mgr.get_ep_class(getter_type_name)
        getter = getter_cls()
        getter.initialize(config={})
        self._getters[name] = getter

    def _resolve_setting_item(self, setting: dict, requirement_getter: EnvParamGetter) -> str:
        if setting["type"] == "value":
            return setting["value"]
        if setting["type"] == "param":
            key = setting["key"]
            return requirement_getter.get_str_value(key)

    def initialize(self, getter_config_list: Sequence[Mapping], requirement_getter: EnvParamGetter) -> None:
        for getter_config in getter_config_list:
            getter_name = getter_config["name"]
            getter_type_name = getter_config["type"]
            getter_cls = self._ep_mgr.get_ep_class(getter_type_name)
            getter: EnvParamGetter = getter_cls()

            raw_getter_settings = getter_config["settings"]
            resolved_getter_settings: dict[str, str] = {}
            assert isinstance(raw_getter_settings, dict)
            for setting_name, raw_setting in raw_getter_settings.items():
                assert isinstance(raw_setting, dict)
                setting_value = self._resolve_setting_item(raw_setting, requirement_getter=requirement_getter)
                resolved_getter_settings[setting_name] = setting_value

            getter.initialize(config=resolved_getter_settings)
            self._getters[getter_name] = getter

lib/dl_testing/dl_testing/env_params/loader.py
datalens-tech/datalens-backend
import os

import attr
from dotenv import (
    dotenv_values,
    find_dotenv,
)

from dl_testing.env_params.getter import EnvParamGetter


class DirectEnvParamGetter(EnvParamGetter):
    """Returns the key as the value for all keys"""

    def get_str_value(self, key: str) -> str:
        return str(key)


@attr.s
class OsEnvParamGetter(EnvParamGetter):
    """Gets params from os environment, with an env-file fallback"""

    _env_from_file: dict = attr.ib(init=False, factory=dict)

    def initialize(self, config: dict) -> None:
        env_file = os.environ.get("DL_TESTS_ENV_FILE") or find_dotenv(filename=".env")
        self._env_from_file = dotenv_values(env_file)

    def get_str_value(self, key: str) -> str:
        env_value = os.environ.get(key)

        if env_value is None:
            return self._env_from_file.get(key)

        return env_value

lib/dl_testing/dl_testing/env_params/main.py
datalens-tech/datalens-backend
from __future__ import annotations

from functools import wraps
import inspect
from typing import (
    Any,
    Callable,
    ClassVar,
    Mapping,
    Optional,
    overload,
)

import attr
import pytest


_FUNC_META_ATTR = "_regulated_test_function_meta"


@attr.s(frozen=True)
class Feature:
    feature_name: str = attr.ib()


@attr.s(frozen=True)
class RegulatedTestFunctionMeta:
    features: frozenset[Feature] = attr.ib(kw_only=True, default=frozenset())


@attr.s(frozen=True)
class RegulatedTestParams:
    mark_tests_skipped: Optional[Mapping[Callable, Optional[str]]] = attr.ib(kw_only=True, default=None)
    mark_tests_failed: Optional[Mapping[Callable, Optional[str]]] = attr.ib(kw_only=True, default=None)
    mark_features_skipped: Optional[Mapping[Feature, Optional[str]]] = attr.ib(kw_only=True, default=None)


def _is_test_func(func: Callable) -> bool:
    return callable(func) and func.__name__.startswith("test_")


def _is_test_func_for_features_with_reason(
    func: Callable,
    features: Mapping[Feature, Optional[str]],
) -> tuple[bool, str]:
    func_meta = getattr(func, _FUNC_META_ATTR, RegulatedTestFunctionMeta())
    found_features = func_meta.features & frozenset(features)
    reasons = sorted((features[feature] or "") for feature in found_features if features[feature])
    reason = "; ".join(reasons)
    return bool(found_features), reason


def _mark_as_test_func(func: Callable) -> Callable:
    if not hasattr(func, _FUNC_META_ATTR):
        setattr(func, _FUNC_META_ATTR, RegulatedTestFunctionMeta())

    return func


def _mark_skipped(func: Callable, reason: Optional[str]) -> Callable:
    assert _is_test_func(func)

    @wraps(func)
    @pytest.mark.skip(reason=reason or "")
    def func_replacement(*args: Any, **kwargs: Any) -> None:
        pass

    return func_replacement


def _mark_failed(func: Callable, reason: Optional[str]) -> Callable:
    assert _is_test_func(func)
    func_replacement = pytest.mark.xfail(reason=reason or "", strict=True)(func)
    return func_replacement


def _mark_functions(
    funcs: Mapping[Callable, Optional[str]],
    marker: Callable[[Callable, Optional[str]], Callable],
) -> Mapping[str, Callable]:
    result: dict[str, Callable] = {}
    for func, reason in funcs.items():
        result[func.__name__] = marker(func, reason)

    return result


def _resolve_test_funcs(bases: tuple[Any, ...]) -> Mapping[str, Callable]:
    return {name: func for base in bases for name, func in inspect.getmembers(base) if _is_test_func(func)}


def _patch_attrs_for_regulated_test_class(
    bases: tuple[Any, ...],
    new_attrs: Optional[Mapping],
    test_params: RegulatedTestParams,
) -> dict[str, Any]:
    new_attrs = dict(new_attrs or {})
    test_funcs = _resolve_test_funcs(bases=bases)

    # Decorate base class functions that require attention
    new_attrs |= _mark_functions(test_params.mark_tests_skipped or {}, _mark_skipped)
    new_attrs |= _mark_functions(test_params.mark_tests_failed or {}, _mark_failed)

    # Find and decorate test functions for skipped features
    mark_tests_skipped_for_features: dict[Callable, str] = {}
    for _name, func in test_funcs.items():
        tests_feature, reason = _is_test_func_for_features_with_reason(func, test_params.mark_features_skipped or {})
        if tests_feature:
            mark_tests_skipped_for_features[func] = reason
    new_attrs |= _mark_functions(mark_tests_skipped_for_features or {}, _mark_skipped)

    return new_attrs


def _make_regulated_test_class(
    mcs: type,
    name: str,
    bases: tuple[Any, ...],
    test_params: RegulatedTestParams,
    attrs: Optional[Mapping] = None,
) -> type:
    # Mark all test functions as such
    attrs = attrs or {}
    assert attrs is not None
    attrs = dict(attrs, **{name: _mark_as_test_func(func) for name, func in attrs.items() if _is_test_func(func)})
    attrs = _patch_attrs_for_regulated_test_class(
        bases=bases,
        new_attrs=attrs,
        test_params=test_params,
    )

    new_cls = type.__new__(mcs, name, bases, attrs)  # type: ignore
    return new_cls


def _wrap_as_regulated_test_case(test_cls: type, *, test_params: RegulatedTestParams) -> type:
    return _make_regulated_test_class(
        mcs=type(test_cls),
        name=test_cls.__name__,
        bases=(test_cls,),
        test_params=test_params,
    )


def _patch_as_regulated_test_case(test_cls: type, *, test_params: RegulatedTestParams) -> type:
    attrs = _patch_attrs_for_regulated_test_class(
        bases=(test_cls,),
        new_attrs={},
        test_params=test_params,
    )
    for name, value in attrs.items():
        setattr(test_cls, name, value)
    return test_cls


def for_features(*args: Feature) -> Callable[[Callable], Callable]:
    """Decorator for test functions to mark them as feature tests"""

    def decorator(func: Callable) -> Callable:
        if not hasattr(func, _FUNC_META_ATTR):
            setattr(
                func,
                _FUNC_META_ATTR,
                RegulatedTestFunctionMeta(
                    features=frozenset(args),
                ),
            )
        return func

    return decorator


class RegulatedTestCase:
    """
    Regulated test case version that patches its own subclasses
    after their creation - in `__init_subclass__`.
    """

    test_params: ClassVar[RegulatedTestParams] = RegulatedTestParams()

    def __init_subclass__(cls, **kwargs: Any) -> None:
        """Patch subclass with test wrappers"""

        super().__init_subclass__()  # This is required so that nothing breaks

        # Apply the method patches
        _patch_as_regulated_test_case(cls, test_params=cls.test_params)


@overload
def regulated_test_case(test_cls: type, /) -> type:
    ...


@overload
def regulated_test_case(*, test_params: RegulatedTestParams = RegulatedTestParams()) -> Callable[[type], type]:
    ...


def regulated_test_case(*args: Any, **kwargs: Any) -> Any:
    """Decorator that makes a regulated test case."""

    if len(args) == 1 and not kwargs and isinstance(args[0], type):
        # Used as a plain decorator (class as first and only arg)
        test_cls = args[0]
        return _patch_as_regulated_test_case(test_cls, test_params=RegulatedTestParams())

    assert not args
    # Used as a parameterized decorator

    def decorator(test_cls: type) -> type:
        assert len(kwargs) <= 1
        test_params = kwargs["test_params"]
        return _patch_as_regulated_test_case(test_cls, test_params=test_params)

    return decorator

lib/dl_testing/dl_testing/regulated_test.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import logging
from typing import (
    TYPE_CHECKING,
    Callable,
)

import aiobotocore.client
import aiobotocore.session
import boto3
import botocore.client
import botocore.exceptions
import botocore.session


if TYPE_CHECKING:
    from mypy_boto3_s3.client import S3Client as SyncS3Client
    from types_aiobotocore_s3 import S3Client as AsyncS3Client


LOGGER = logging.getLogger(__name__)


def create_s3_client(s3_settings) -> AsyncS3Client:  # type: ignore  # TODO: fix
    session = aiobotocore.session.get_session()
    return session.create_client(
        service_name="s3",
        aws_access_key_id=s3_settings.ACCESS_KEY_ID,
        aws_secret_access_key=s3_settings.SECRET_ACCESS_KEY,
        endpoint_url=s3_settings.ENDPOINT_URL,
    )


def create_sync_s3_client(s3_settings) -> SyncS3Client:  # type: ignore  # TODO: fix
    session = boto3.Session()
    return session.client(
        service_name="s3",
        aws_access_key_id=s3_settings.ACCESS_KEY_ID,
        aws_secret_access_key=s3_settings.SECRET_ACCESS_KEY,
        endpoint_url=s3_settings.ENDPOINT_URL,
    )


async def create_s3_bucket(
    s3_client: AsyncS3Client,
    bucket_name: str,
    max_attempts: int = 10,
) -> None:
    attempt = 1
    while True:
        try:
            resp = await s3_client.create_bucket(Bucket=bucket_name)
            LOGGER.info(resp)
            break
        except botocore.exceptions.ClientError as ex:
            if ex.response.get("Error", {}).get("Code") != "BucketAlreadyOwnedByYou":
                raise
            break
        except botocore.exceptions.HTTPClientError:
            LOGGER.warning(
                f"HTTPClientError during creating S3 bucket. Attempt {attempt} from {max_attempts}. "
                "Retrying after 5 seconds...",
                exc_info=True,
            )
            if attempt > max_attempts:
                raise
            await asyncio.sleep(5)
        attempt += 1


async def get_lc_rules_number(s3_client: AsyncS3Client, bucket: str) -> int:
    try:
        lc_config = await s3_client.get_bucket_lifecycle_configuration(Bucket=bucket)
    except botocore.exceptions.ClientError as ex:
        if ex.response["Error"]["Code"] == "NoSuchLifecycleConfiguration":
            lc_config = {"Rules": []}
        else:
            raise
    return len(lc_config["Rules"])


async def s3_file_exists(s3_client: AsyncS3Client, bucket: str, key: str) -> bool:
    try:
        s3_resp = await s3_client.head_object(
            Bucket=bucket,
            Key=key,
        )
    except botocore.exceptions.ClientError as ex:
        if ex.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
            return False
        raise
    return s3_resp["ResponseMetadata"]["HTTPStatusCode"] == 200


S3_TBL_FUNC_TEMPLATE = """s3(
'{s3_endpoint}/{bucket}/{filename}',
'{key_id}',
'{secret_key}',
'{file_fmt}',
'{schema_line}')"""


def s3_tbl_func_maker(s3_settings) -> Callable[..., str]:  # type: ignore  # TODO: fix
    def table_function(  # type: ignore  # TODO: fix
        for_: str,
        conn_dto,
        filename: str,
        file_fmt: str,
        schema_line: str,
    ) -> str:
        if for_ == "db":  # secret are filled
            tbl_func = S3_TBL_FUNC_TEMPLATE.format(
                s3_endpoint=conn_dto.s3_endpoint,
                bucket=conn_dto.bucket,
                filename=filename,
                key_id=s3_settings.ACCESS_KEY_ID,
                secret_key=s3_settings.SECRET_ACCESS_KEY,
                file_fmt=file_fmt,
                schema_line=schema_line,
            )
        elif for_ == "dba":  # secrets will be filled by the adapter
            tbl_func = S3_TBL_FUNC_TEMPLATE.format(
                s3_endpoint=conn_dto.s3_endpoint,
                bucket=conn_dto.bucket,
                filename=filename,
                key_id=f"key_id_{conn_dto.replace_secret}",
                secret_key=f"secret_key_{conn_dto.replace_secret}",
                file_fmt=file_fmt,
                schema_line=schema_line,
            )
        else:
            raise ValueError(f"Unknown mode '{for_}'")
        return tbl_func

    return table_function

lib/dl_testing/dl_testing/s3_utils.py
datalens-tech/datalens-backend
from __future__ import annotations

import os


RUN_DEVHOST_TESTS = os.environ.get("RUN_DEVHOST_TESTS")
TESTS_CACHES_DIR = os.environ.get("TESTS_CACHES_DIR")
CA_BUNDLE_FILE = os.environ.get("REQUESTS_CA_BUNDLE") or "/etc/ssl/certs/ca-certificates.crt"

DL_ENV_TESTS_FILTER = os.environ.get("DL_ENV_TESTS_FILTER")

lib/dl_testing/dl_testing/shared_testing_constants.py
datalens-tech/datalens-backend
from __future__ import annotations


DASHSQL_EXAMPLE_PARAMS = {
    "some_string": {"type_name": "string", "value": "some\\:string\nwith\\stuff"},
    "some_integer": {"type_name": "integer", "value": "562949953421312"},
    "some_float": {"type_name": "float", "value": "73786976294838206464.5"},
    "some_boolean": {"type_name": "boolean", "value": "true"},
    "some_other_boolean": {"type_name": "boolean", "value": "false"},
    "some_date": {"type_name": "date", "value": "2021-07-19"},
    "some_datetime": {"type_name": "datetime", "value": "2021-07-19T19:35:43"},
    "3xtr4 ше1гd param": {"type_name": "string", "value": "11"},
    "3xtr4 же1гd param": {"type_name": "string", "value": "22"},
    "intvalues": {"type_name": "integer", "value": ["1", "2", "3"]},
    "strvalues": {"type_name": "string", "value": ["a", "b", "c"]},
}

lib/dl_testing/dl_testing/test_data/sql_queries.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import contextmanager
import http.client
import logging
import os
from typing import (
    Any,
    Callable,
    List,
    Literal,
    Tuple,
    overload,
)

import pytest

from dl_testing.containers import get_test_container_hostport
from dl_testing.shared_testing_constants import RUN_DEVHOST_TESTS
from dl_utils.wait import wait_for


LOGGER = logging.getLogger(__name__)


def skip_outside_devhost(func):
    """
    Not all tests can run in just any environment (particularly in autotests).
    So put those under an environment flag.
    """
    if RUN_DEVHOST_TESTS:
        return func
    return pytest.mark.skip("Requires RUN_DEVHOST_TESTS=1")(func)


def wait_for_initdb(initdb_port, initdb_host=None, timeout=900, require: bool = False):
    initdb_host = initdb_host or get_test_container_hostport("init-db").host
    # TODO: initdb_port?

    def check_initdb_liveness() -> Tuple[bool, Any]:
        try:
            conn = http.client.HTTPConnection(initdb_host, initdb_port)
            conn.request("GET", "/")
            resp = conn.getresponse()
            body = resp.read().decode("utf-8", errors="replace")
            if resp.status != 200:
                raise Exception("Non-ok response", dict(res=resp, status=resp.status, body=body))
            return True, body
        except Exception as exc:
            return False, dict(exc=exc)

    return wait_for(
        name="initdb readiness",
        condition=check_initdb_liveness,
        timeout=timeout,
        require=require,
    )


@overload
def get_log_record(
    caplog: Any, predicate: Callable[[logging.LogRecord], bool], single: Literal[True] = True
) -> logging.LogRecord:
    pass


@overload
def get_log_record(
    caplog: Any, predicate: Callable[[logging.LogRecord], bool], single: Literal[False] = False
) -> List[logging.LogRecord]:
    pass


def get_log_record(caplog, predicate, single=False):
    log_records = [rec for rec in caplog.records if predicate(rec)]

    if single:
        assert len(log_records) == 1, f"Unexpected count of record in logs: {len(log_records)}"
        return log_records[0]

    return log_records


def guids_from_titles(result_schema: List[dict], titles: List[str]) -> List[str]:
    fields = [f["field"] if "field" in f else f for f in result_schema]
    guid_by_title = {f["title"]: f["guid"] for f in fields if f.get("title") in titles and f.get("guid")}
    return [guid_by_title[title] for title in titles]


@contextmanager
def override_env_cm(to_set: dict[str, str], purge: bool = False):
    preserved = {k: v for k, v in os.environ.items()}

    try:
        if purge:
            for k in os.environ:
                del os.environ[k]
        for k, v in to_set.items():
            os.environ[k] = v
        yield
    finally:
        for k in to_set:
            del os.environ[k]
            # os.unsetenv(k)

        for k, v in preserved.items():
            os.environ[k] = v

lib/dl_testing/dl_testing/utils.py
datalens-tech/datalens-backend
from dl_testing.regulated_test import (
    Feature,
    RegulatedTestCase,
    RegulatedTestParams,
    for_features,
    regulated_test_case,
)


class BaseForTestRTClass(RegulatedTestCase):
    my_feature = Feature("my_feature")

    def test_skipped_test(self) -> None:
        raise NotImplementedError

    def test_failed_test(self) -> None:
        raise AssertionError()

    @for_features(my_feature)
    def test_skipped_test_with_feature(self) -> None:
        raise NotImplementedError

    def test_normal_test(self) -> None:
        pass


class TestRTClass(BaseForTestRTClass):
    test_params = RegulatedTestParams(
        mark_tests_skipped={
            BaseForTestRTClass.test_skipped_test: "Skip it",
        },
        mark_tests_failed={
            BaseForTestRTClass.test_failed_test: "It fails",
        },
        mark_features_skipped={BaseForTestRTClass.my_feature: "Feature not implemented"},
    )


class BaseForRTDecorator:
    my_feature = Feature("my_feature")

    def test_skipped_test(self) -> None:
        raise NotImplementedError

    def test_failed_test(self) -> None:
        raise AssertionError()

    @for_features(my_feature)
    def test_skipped_test_with_feature(self) -> None:
        raise NotImplementedError

    def test_normal_test(self) -> None:
        pass


@regulated_test_case(
    test_params=RegulatedTestParams(
        mark_tests_skipped={
            BaseForRTDecorator.test_skipped_test: "Skip it",
        },
        mark_tests_failed={
            BaseForRTDecorator.test_failed_test: "It fails",
        },
        mark_features_skipped={BaseForRTDecorator.my_feature: "Feature not implemented"},
    ),
)
class TestRTDecorator(BaseForRTDecorator):
    pass

lib/dl_testing/dl_testing_tests/unit/test_regulated_test.py
datalens-tech/datalens-backend
class OpCode:
    WB_CREATE: str = "wb_create"
    WB_DELETE: str = "wb_delete"
    WB_CHECK_EXISTS: str = "wb_check_exists"
    GET_FOLDER_CONTENT: str = "get_folder_content"
    WB_INFO_GET: str = "workbook_info_get"
    WB_BASIC_INFO_GET: str = "workbook_basic_info_get"
    WB_PRIVATE_INFO_GET: str = "private_workbook_info_get"
    COLLECTION_CONTENT_GET: str = "collection_content_get"

lib/dl_us_client/dl_us_client/constants.py
datalens-tech/datalens-backend
import attr

from dl_api_commons.client.base import Req
from dl_api_commons.client.common import CommonInternalAPIClient
from dl_us_client.constants import OpCode


@attr.s()
class USWorkbookCommandClient(CommonInternalAPIClient):
    async def create_workbook(self, title: str) -> str:
        op_code = OpCode.WB_CREATE

        req = Req(
            method="post",
            url="v2/workbooks",
            data_json=dict(
                title=title,
            ),
            require_ok=False,
        )
        resp = await self.make_request(req)

        if resp.status == 200:
            return resp.json["workbookId"]

        self.raise_from_resp(resp, op_code=op_code)

    async def delete_workbook(self, wb_id: str) -> None:
        op_code = OpCode.WB_DELETE

        req = Req(method="delete", url=f"v2/workbooks/{wb_id}", require_ok=False)
        resp = await self.make_request(req)

        if resp.status == 200:
            return
        self.raise_from_resp(resp, op_code=op_code)

    async def add_wb_sa_access_binding(self, wb_id: str, sa_id: str) -> None:  # TODO: make more generic
        op_code = "wb_alter_access_bindings"

        data_json = dict(
            deltas=[
                dict(
                    action="ADD",
                    accessBinding=dict(
                        roleId="datalens.workbooks.admin", subject=dict(id=sa_id, type="serviceAccount")
                    ),
                )
            ],
        )

        req = Req(method="post", url=f"v2/workbooks/{wb_id}/access-bindings", data_json=data_json, require_ok=False)
        resp = await self.make_request(req)

        if resp.status != 200:
            self.raise_from_resp(resp, op_code=op_code)

lib/dl_us_client/dl_us_client/us_workbook_cmd_client.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
from concurrent.futures import ThreadPoolExecutor
import contextvars
from functools import partial
import logging
from types import TracebackType
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncIterable,
    Awaitable,
    Callable,
    Iterable,
    List,
    Optional,
    Type,
    TypeVar,
    cast,
)
import weakref

from async_timeout import Timeout as VanillaTimeout


if TYPE_CHECKING:
    from concurrent.futures import Future


LOGGER = logging.getLogger(__name__)

_CORO_TV = TypeVar("_CORO_TV")


async def shield_wait_for_complete(
    coro: Awaitable[_CORO_TV], log_msg_on_cancel: str = "CancelledError during critical section execution"
) -> _CORO_TV:
    """
    Shields coroutine from cancellation and waits until it finished.
    If execution was cancelled - raise CancelledError after coroutine done.
    Logs exceptions if task was cancelled and exception occurs.
    :param coro: Coroutine to shield
    :param log_msg_on_cancel: message to log in case of cancellation (will be logged exactly when cancellation occurs)
    :return: shielded coroutine result
    """
    loop = asyncio.get_event_loop()
    real_task = loop.create_task(coro)

    try:
        return await asyncio.shield(real_task)
    except asyncio.CancelledError:
        LOGGER.info(log_msg_on_cancel)

        # Awaiting protected task and logging exception if occurred
        try:
            await real_task
        except Exception as e:  # noqa
            logging.exception("Error during awaiting critical section after cancellation")

        # Raising original exception to meet
        raise


_SUBMIT_RT = TypeVar("_SUBMIT_RT")


class ContextVarExecutor(ThreadPoolExecutor):
    """Bug in Python: default TPE does not propagate ContextVars in running thread."""

    def __init__(self, *args, **kwargs):  # type: ignore
        super().__init__(*args, **kwargs)
        self._futures_set = weakref.WeakSet()  # type: ignore

    def get_pending_futures_count(self) -> int:
        return len([fut for fut in self._futures_set if fut.running()])

    def submit(self, fn: Callable[..., _SUBMIT_RT], /, *args: Any, **kwargs: Any) -> Future[_SUBMIT_RT]:
        ctx = contextvars.copy_context()  # type: contextvars.Context
        fut = super().submit(cast(Callable[..., _SUBMIT_RT], partial(ctx.run, partial(fn, *args, **kwargs))))
        # TODO FIX: Periodically cleanup set
        self._futures_set.add(fut)
        return fut


_WRAPPED_RT = TypeVar("_WRAPPED_RT")


def await_sync(coro: Awaitable[_WRAPPED_RT], loop: Optional[asyncio.AbstractEventLoop] = None) -> _WRAPPED_RT:
    if loop is None:
        loop = asyncio.get_event_loop()

    return loop.run_until_complete(coro)


_ITERABLE_T = TypeVar("_ITERABLE_T")


def to_sync_iterable(
    async_iterable: AsyncIterable[_ITERABLE_T],
    loop: Optional[asyncio.AbstractEventLoop] = None,
) -> Iterable[_ITERABLE_T]:
    actual_loop = asyncio.get_event_loop() if loop is None else loop

    def wrapper() -> Iterable[_ITERABLE_T]:
        aiter = async_iterable.__aiter__()
        while True:
            try:
                yield actual_loop.run_until_complete(aiter.__anext__())
            except StopAsyncIteration:
                return

    return wrapper()


async def alist(aiterable: AsyncIterable[_ITERABLE_T]) -> List[_ITERABLE_T]:
    """Gather an async iterable into a single list"""
    result = []
    async for item in aiterable:
        result.append(item)
    return result


class BIAsyncTimeout:
    vanilla_timeout: VanillaTimeout

    @classmethod
    def from_vanilla(cls, vanilla_timeout: VanillaTimeout) -> BIAsyncTimeout:
        instance = cls()
        instance.vanilla_timeout = vanilla_timeout
        return instance

    @classmethod
    def from_params(cls, deadline: Optional[float], loop: asyncio.AbstractEventLoop) -> BIAsyncTimeout:
        return cls.from_vanilla(VanillaTimeout(deadline, loop))

    async def __aenter__(self) -> BIAsyncTimeout:
        return BIAsyncTimeout.from_vanilla(await self.vanilla_timeout.__aenter__())

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> Optional[bool]:
        return await self.vanilla_timeout.__aexit__(exc_type, exc_val, exc_tb)


def timeout(delay: Optional[float]) -> BIAsyncTimeout:
    loop = asyncio.get_running_loop()
    if delay is not None:
        deadline = loop.time() + delay  # type: Optional[float]
    else:
        deadline = None
    return BIAsyncTimeout.from_params(deadline, loop)

lib/dl_utils/dl_utils/aio.py
datalens-tech/datalens-backend
import abc
from importlib import metadata
from typing import (
    Generic,
    Type,
    TypeVar,
)

import attr


_EP_CLS_TV = TypeVar("_EP_CLS_TV")


@attr.s
class EntrypointClassManager(abc.ABC, Generic[_EP_CLS_TV]):
    """
    Loads classes registered as entrypoints of a given type
    """

    entrypoint_group_name: str = attr.ib(kw_only=True)

    def get_all_ep_classes(self) -> dict[str, Type[_EP_CLS_TV]]:
        entrypoints = list(metadata.entry_points().select(group=self.entrypoint_group_name))  # type: ignore
        return {ep.name: ep.load() for ep in entrypoints}

    def get_ep_class(self, name: str) -> Type[_EP_CLS_TV]:
        entrypoints = list(metadata.entry_points().select(group=self.entrypoint_group_name, name=name))  # type: ignore
        assert len(entrypoints) == 1
        return entrypoints[0].load()

lib/dl_utils/dl_utils/entrypoints.py
datalens-tech/datalens-backend
import functools
from typing import (
    Any,
    Callable,
)
import weakref


def method_lru(maxsize: int = 128, typed: bool = False) -> Callable:
    """
    Hacked lru cache for methods
    """

    def wrapper(func: Callable) -> Callable:
        @functools.lru_cache(maxsize, typed)
        def _func(_self: Callable, *args, **kwargs) -> Any:  # type: ignore
            return func(_self(), *args, **kwargs)

        @functools.wraps(func)
        def inner(self: object, *args, **kwargs) -> Callable:  # type: ignore
            return _func(weakref.ref(self), *args, **kwargs)

        return inner

    return wrapper

lib/dl_utils/dl_utils/func_tools.py
datalens-tech/datalens-backend
"""
Assorted methods for making a safe representative value out of some
less constrained value.
"""

from __future__ import annotations

import base64
import hashlib
import json
import random
import re
import string
from typing import (
    Any,
    Callable,
    Iterable,
    Tuple,
)
import uuid


CHAR_CLASSES = {
    "H": set(string.hexdigits.upper()),
    "h": set(string.hexdigits.lower()),
    "-": {"-"},
}
# '-'.join('h' * x for x in (8, 4, 4, 4, 12))
UUID_LOWER = "hhhhhhhh-hhhh-hhhh-hhhh-hhhhhhhhhhhh"
UUID_UPPER = UUID_LOWER.upper()
UUID_LOWER_CLASSES = tuple(CHAR_CLASSES[char] for char in UUID_LOWER)
UUID_UPPER_CLASSES = tuple(CHAR_CLASSES[char] for char in UUID_UPPER)


def cut_uuid(value: str) -> Tuple[str, str]:
    """
    Cut as much of an uuid-looking prefix from `value` as possible.

    >>> some_uuid = '765a18f2-4d1c-4a7f-b0cd-8bdb392803f7'
    >>> cut_uuid(some_uuid.lower())
    ('765a18f2-4d1c-4a7f-b0cd-8bdb392803f7', '')
    >>> cut_uuid(some_uuid.upper())
    ('765A18F2-4D1C-4A7F-B0CD-8BDB392803F7', '')
    >>> cut_uuid('')
    ('', '')
    >>> cut_uuid('zxc')
    ('', 'zxc')
    >>> cut_uuid(some_uuid.upper()[:10] + some_uuid.lower()[10:])
    ('765A18F2-4', 'd1c-4a7f-b0cd-8bdb392803f7')
    """
    idx = max(
        next(
            # basically enumerate(zip(value, ccs))
            (idx for idx in range(min(len(value), len(ccs))) if value[idx] not in ccs[idx]),
            len(value),
        )
        for ccs in (UUID_LOWER_CLASSES, UUID_UPPER_CLASSES)
    )
    return value[:idx], value[idx:]


def string_to_uuid(
    data: str,
    min_random_bytes: int = 3,
    visual_separate: str = "ff",
) -> str:
    """
    Make an UUID string out of any string,
    with some of its data and some random bits.

    >>> string_to_uuid(string_to_uuid('765a18f2-4d1f-f63f-fc3b-2267eceb5182'))[:-6]
    '765a18f2-4d1f-f63f-fc3b-22ffff'
    """
    data_uuid_like, data_etcetera = cut_uuid(data)
    uuid_chars = 32

    char_budget = uuid_chars
    # 0: data_uuid_hex
    # 1: visual separator
    # 2: data-bytes
    # 3: visual separator
    # 4: random hex padding
    # 5: random hex (required)
    hex_pieces = {}

    if visual_separate:
        hex_pieces[1] = visual_separate
        char_budget -= len(visual_separate)
        hex_pieces[3] = visual_separate
        char_budget -= len(visual_separate)

    def make_random_hex(chars: int) -> str:
        return "%0{}x".format(chars) % (random.getrandbits(chars * 4),)

    random_hex_chars = int(min_random_bytes * 2)
    random_hex_tail = make_random_hex(random_hex_chars)
    hex_pieces[5] = random_hex_tail
    char_budget -= len(random_hex_tail)
    assert char_budget >= 0, "should not ask for too many random bytes"

    if char_budget > 0:
        data_uuid_hex = data_uuid_like.replace("-", "")
        data_uuid_hex = data_uuid_hex[:char_budget]
        hex_pieces[0] = data_uuid_hex
        char_budget -= len(data_uuid_hex)

    if char_budget > 0:
        max_bytechars = int(char_budget / 2 + 1)
        data_etcetera_hex = "".join("%02x" % (ord(char),)[:2] for char in data_etcetera[:max_bytechars])
        data_etcetera_hex = data_etcetera_hex[:char_budget]
        hex_pieces[2] = data_etcetera_hex
        char_budget -= len(data_etcetera_hex)

    if char_budget > 0:
        random_hex_padding = make_random_hex(char_budget)
        hex_pieces[4] = random_hex_padding
        char_budget -= len(random_hex_padding)

    assert char_budget == 0, dict(char_budget=char_budget, hex_pieces=hex_pieces)
    assert sum(len(piece) for piece in hex_pieces.values()) == uuid_chars
    result_hex = "".join(hex_pieces.get(idx, "") for idx in range(6))
    return str(uuid.UUID(hex=result_hex))


# `Any` should be a `hashlib._hashlib.HASH`,
# an object with a `def digest(self) -> bytes:` method.
THasher = Callable[[bytes], Any]


def clear_hash(
    value: str,
    lowercase: bool = True,
    allowed_chars_re: str = "A-Za-z0-9_",
    hash_prefix: str = "cs",
    base_len: int = 12,
    hash_len: int = 12,
    hasher: THasher = hashlib.sha256,
) -> str:
    """
    Hash a value, leaving a sanitized part of it for debuggability.

    >>> clear_hash('<a href="javascript:alert(2)">фыва</a>')
    'a_href_javas__cst2let6ipy3jh'
    """
    hashed = hasher(value.encode("utf-8")).digest()
    hashed = base64.b32encode(hashed).decode("ascii").lower().rstrip("=")
    hashed = hashed[:hash_len]

    cleaned = re.sub(r"[^{}]+".format(allowed_chars_re), "_", value)
    cleaned = cleaned.strip("_")
    cleaned = cleaned[:base_len]
    return f"{cleaned}__{hash_prefix}{hashed}"


def make_jsonable(value: Any, default: Callable[[Any], Any] = repr) -> Any:
    """
    A hack to ensure the value is JSON-compatible.
    To be used in logging `extra` in uncertain cases.
    """
    return json.loads(json.dumps(value, default=default))


def param_bool(value: Any, false_values: Iterable[str] = ("false", "0", "no"), default: bool = False) -> bool:
    """
    Header value to a boolean value.

    Note: OpenAPI spec only allows 'true' and 'false'.
    """
    if isinstance(value, bool):
        return value
    if not value:
        return default
    value = value.lower()
    if value in false_values:
        return False
    return True

lib/dl_utils/dl_utils/sanitize.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import (
    AbstractSet,
    Any,
    Dict,
    Iterable,
    List,
    Mapping,
    Optional,
    Sequence,
    Union,
)

from marshmallow.decorators import (
    POST_DUMP,
    POST_LOAD,
    PRE_DUMP,
    PRE_LOAD,
)
from marshmallow.error_store import ErrorStore
from marshmallow.exceptions import ValidationError
from marshmallow_oneofschema import OneOfSchema


class OneOfSchemaWithDumpLoadHooks(OneOfSchema):  # TODO: Move to bi_model_tools
    """
    OneOfSchema has disabled all pre/post_dump/load hooks by default so we add the following copy-paste from marshmallow
    to be able to use those hooks

    See https://github.com/marshmallow-code/marshmallow-oneofschema/issues/4
    """

    def dump(
        self, obj: Any, *, many: Optional[bool] = None, **kwargs: Any
    ) -> Union[Mapping[str, Any], Iterable[Mapping[str, Any]]]:
        many = self.many if many is None else bool(many)

        if self._has_processors(PRE_DUMP):
            processed_obj = self._invoke_dump_processors(PRE_DUMP, obj, many=many, original_data=obj)
        else:
            processed_obj = obj

        result = super().dump(processed_obj, many=many, **kwargs)

        if self._has_processors(POST_DUMP):
            result = self._invoke_dump_processors(POST_DUMP, result, many=many, original_data=obj)

        return result

    def load(
        self,
        data: Union[Mapping[str, Any], Iterable[Mapping[str, Any]]],
        *,
        many: Optional[bool] = None,
        partial: Optional[Union[bool, Sequence[str], AbstractSet[str]]] = None,
        unknown: Optional[str] = None,
        **kwargs: Any,
    ) -> Any:
        error_store = ErrorStore()
        errors: Dict[str, List[str]] = {}
        many = self.many if many is None else bool(many)
        unknown = unknown or self.unknown
        if partial is None:
            partial = self.partial
        processed_data = data
        result = None
        # Run preprocessors
        if self._has_processors(PRE_LOAD):
            try:
                processed_data = self._invoke_load_processors(
                    PRE_LOAD, data, many=many, original_data=data, partial=partial
                )
            except ValidationError as err:
                errors = err.normalized_messages()
                result = None
        if not errors:
            # Deserialize data
            try:
                result = super().load(processed_data, many=many, partial=partial, unknown=unknown, **kwargs)
            except ValidationError as error:
                error_store.store_error(error.messages)
                result = error.valid_data
            errors = error_store.errors
            # Run post processors
            if not errors and self._has_processors(POST_LOAD):
                try:
                    result = self._invoke_load_processors(
                        POST_LOAD,
                        result,
                        many=many,
                        original_data=data,
                        partial=partial,
                    )
                except ValidationError as err:
                    errors = err.normalized_messages()
        if errors:
            exc = ValidationError(errors, data=data, valid_data=result)
            self.handle_error(exc, data, many=many, partial=partial)
            raise exc

        return result

lib/dl_utils/dl_utils/schemas.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import asyncio
from typing import (
    Any,
    Awaitable,
    Iterable,
)


class TaskRunner(abc.ABC):
    async def initialize(self) -> None:
        pass

    @abc.abstractmethod
    async def schedule(self, awaitable: Awaitable) -> None:
        raise NotImplementedError

    @abc.abstractmethod
    async def finalize(self) -> Any:
        raise NotImplementedError


class ImmediateTaskRunner(TaskRunner):
    async def schedule(self, awaitable: Awaitable) -> None:
        await awaitable

    async def finalize(self) -> None:
        pass


class ConcurrentTaskRunner(TaskRunner):
    def __init__(self, concurrency_limit: int = 5) -> None:
        self._tasks: list[Awaitable] = []
        self._sem = asyncio.Semaphore(concurrency_limit)

    async def _semaphore_wrapper(self, awaitable: Awaitable) -> Any:
        async with self._sem:
            return await awaitable

    async def schedule(self, awaitable: Awaitable) -> None:
        self._tasks.append(self._semaphore_wrapper(awaitable))

    async def finalize(self) -> Iterable:  # should be list, but typing for asyncio.gather is a mess
        tasks = [asyncio.ensure_future(task) for task in self._tasks]  # creating task starts coroutine
        try:
            return await asyncio.gather(*tasks)
        except Exception:
            for task in tasks:
                task.cancel()
            raise

lib/dl_utils/dl_utils/task_runner.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import contextmanager
from enum import Enum
import functools
from itertools import islice
import operator
import os
import sys
from time import time
from typing import (
    Any,
    Callable,
    Generator,
    Iterable,
    List,
    NoReturn,
    Optional,
    Tuple,
    Type,
    TypeVar,
    cast,
)

import attr


def get_pdb() -> Any:
    try:
        import ipdb

        return ipdb
    except Exception:
        import pdb

        return pdb


def maybe_postmortem(ei: Any = None) -> None:
    """
    An env-controlled post-mortem breakpoint.
    Set `BI_ERR_PDB=1` to use.
    """
    if ei is None:
        ei = sys.exc_info()
    if os.environ.get("BI_ERR_PDB") != "1":
        return
    _, _, sys.last_traceback = ei
    import traceback

    print("".join(traceback.format_exception(*ei)))
    pdb = get_pdb()
    pdb.pm()


_CALLABLE_TV = TypeVar("_CALLABLE_TV", bound=Callable)


def exc_catch_wrap(func: _CALLABLE_TV) -> _CALLABLE_TV:
    """
    An env-controlled post-mortem function wrapper.
    Set `BI_ERR_PDB=1` to use. Does nothing otherwise.
    """
    if os.environ.get("BI_ERR_PDB") != "1":
        return func

    @functools.wraps(func)
    def exc_catch_wrapped(*args: Any, **kwargs: Any) -> Any:
        try:
            return func(*args, **kwargs)
        except Exception:
            maybe_postmortem()
            raise

    return cast(_CALLABLE_TV, exc_catch_wrapped)


def exc_catch_awrap(afunc: _CALLABLE_TV) -> _CALLABLE_TV:
    """
    An env-controlled post-mortem async-function wrapper.
    Set `BI_ERR_PDB=1` to use. Does nothing otherwise.
    """
    if os.environ.get("BI_ERR_PDB") != "1":
        return afunc

    @functools.wraps(afunc)
    async def exc_catch_awrapped(*args: Any, **kwargs: Any) -> Any:
        try:
            return await afunc(*args, **kwargs)
        except Exception:
            maybe_postmortem()
            raise

    return cast(_CALLABLE_TV, exc_catch_awrapped)


def get_type_full_name(t: Type) -> str:
    module = t.__module__
    qual_name = t.__qualname__

    if module == "builtins":
        return qual_name
    return f"{module}.{qual_name}"


# split_list value TypeVar
_SL_V_TV = TypeVar("_SL_V_TV")


def split_list(
    iterable: Iterable[_SL_V_TV], condition: Callable[[_SL_V_TV], bool]
) -> Tuple[List[_SL_V_TV], List[_SL_V_TV]]:
    """
    Split list items into `(matching, non_matching)` by `condition(item)` callable.
    """
    matching: List[_SL_V_TV] = []
    non_matching: List[_SL_V_TV] = []
    for item in iterable:
        if condition(item):
            matching.append(item)
        else:
            non_matching.append(item)
    return matching, non_matching


class DotDict(dict):
    """A simple dict subclass with items also available over attributes"""

    def __getattr__(self, item: str) -> Any:
        try:
            return self[item]
        except KeyError as exc:
            raise AttributeError(*exc.args) from exc

    def __setattr__(self, key: str, value: Any) -> None:
        self[key] = value


@attr.s(frozen=True)
class DataKey:
    parts: tuple[str, ...] = attr.ib()


@attr.s
class AddressableData:
    data: dict[str, Any] = attr.ib()

    def contains(self, key: DataKey) -> bool:
        try:
            self.get(key)
            return True
        except KeyError:
            return False

    def get(self, key: DataKey) -> Any:
        return functools.reduce(operator.getitem, key.parts, self.data)

    def set(self, key: DataKey, value: Any) -> None:
        key_head = DataKey(parts=key.parts[:-1])
        self.get(key_head)[key.parts[-1]] = value

    def pop(self, key: DataKey) -> Any:
        key_head = DataKey(parts=key.parts[:-1])
        return self.get(key_head).pop(key.parts[-1])


T = TypeVar("T")


def batched(iterable: Iterable[T], n: int) -> Iterable[Iterable[T]]:
    if n < 1:
        raise ValueError("n must be at least one")
    it = iter(iterable)
    while batch := tuple(islice(it, n)):
        yield batch


def join_in_chunks(pieces: Iterable[str], sep: str, max_len: int) -> Iterable[str]:
    """An `sep.join(pieces)` equivalent that chunks the pieces to stay under the maximum piece length"""
    pieces_it = iter(pieces)
    try:
        result = next(pieces_it)
    except StopIteration:
        return  # empty pieces -> empty chunks
    for piece in pieces_it:
        full_piece = sep + piece
        if len(result) + len(full_piece) >= max_len:
            yield result
            result = piece
        else:
            result += full_piece
    yield result


_ENUM_TV = TypeVar("_ENUM_TV", bound=Enum)


def enum_not_none(val: Optional[_ENUM_TV]) -> _ENUM_TV:
    assert val is not None
    return val


def time_it(fn: Callable) -> Callable:
    @functools.wraps(fn)
    def wrap(*args: Any, **kwargs: Any) -> Any:
        t0 = time()
        print(f"Invoked {fn}({args}, {kwargs})"[:160])
        result = fn(*args, **kwargs)
        delta = time() - t0
        if delta >= 0.01:
            print(f"<< Time elapsed: {delta} for {fn}({args}, {kwargs})"[:160])
        return result

    return wrap


@contextmanager
def time_it_cm(label: str) -> Generator[None, None, None]:
    # print(f'Time it for {label}')
    t0 = time()
    yield
    delta = time() - t0
    if delta >= 0.01:
        print(f"Time elapsed for {label}: {delta}")

lib/dl_utils/dl_utils/utils.py
datalens-tech/datalens-backend
import asyncio
import logging
import time
from typing import (
    Any,
    Awaitable,
    Callable,
    Optional,
    Tuple,
    Union,
)


LOGGER = logging.getLogger(__name__)


class TimeLimitExceeded(Exception):
    """Accompanying exception for `wait_for` and `await_for`"""


def wait_for(
    name: str,
    condition: Callable[[], Union[bool, Tuple[bool, str]]],
    timeout: float = 300.0,
    interval: float = 1.0,
    log_func: Optional[Callable[[str], None]] = None,
    require: bool = True,
) -> Tuple[bool, str]:
    """
    `condition` should either return `is_done` or `(is_done, status)`.
    """
    start_time = time.monotonic()
    max_time = start_time + timeout

    status: str
    while True:
        loop_start_time = time.monotonic()

        cond_result = condition()
        if isinstance(cond_result, bool):
            cond_ok = cond_result
            status = ""
        else:
            cond_ok, status = cond_result

        if cond_ok:
            return True, status

        now = time.monotonic()
        if now > max_time:
            message = f"Timed out waiting for {name!r}, last status {status!r}, timeout {timeout:.3f}s."
            if require:
                raise TimeLimitExceeded(message)
            LOGGER.error(message)
            return False, status

        next_try_time = loop_start_time + interval
        sleep_time = next_try_time - now
        if log_func is not None:
            log_func(f"Waiting for {name!r} (sleep {sleep_time:.3f}s, status {status!r}).")
        if sleep_time > 0.001:
            time.sleep(sleep_time)


async def await_for(
    name: str,
    condition: Callable[[], Awaitable[Union[bool, Tuple[bool, Any]]]],
    timeout: float = 300.0,
    interval: float = 1.0,
    log_func: Optional[Callable[[str], None]] = None,
    require: bool = True,
) -> Tuple[bool, str]:
    """
    `condition` should either return `is_done` or `(is_done, status)`.
    """
    start_time = time.monotonic()
    max_time = start_time + timeout

    status: str
    while True:
        loop_start_time = time.monotonic()

        cond_result = await condition()
        if isinstance(cond_result, bool):
            cond_ok = cond_result
            status = ""
        else:
            cond_ok, status = cond_result

        if cond_ok:
            return True, status

        now = time.monotonic()
        if now > max_time:
            message = f"Timed out waiting for {name!r}, status {status!r}, timeout {timeout:.3f}s."
            if require:
                raise TimeLimitExceeded(message)
            LOGGER.error(message)
            return False, status

        next_try_time = loop_start_time + interval
        sleep_time = next_try_time - now
        if log_func is not None:
            remaining = max_time - time.monotonic()
            log_func(f"Waiting for {name!r} (sleep {sleep_time:.3f}s, limit {remaining:.3f}s, status {status!r}).")
        if sleep_time > 0.001:
            await asyncio.sleep(sleep_time)

lib/dl_utils/dl_utils/wait.py
datalens-tech/datalens-backend
from __future__ import annotations

from dl_utils.utils import join_in_chunks


def test_stuff():
    assert True


def test_join_in_chunks():
    # TODO: `hypothesis` test over these parameters
    values = list(map(str, range(1000)))
    sep = "_"
    max_len = 10
    # ...
    tst_result = list(
        join_in_chunks(
            values,
            sep,
            max_len=max_len,
        )
    )
    effective_max_len = max(max_len, max(len(val) + 1 for val in values))
    assert all(len(piece) < effective_max_len for piece in tst_result)
    assert [item for piece in tst_result for item in piece.split(sep)] == values

lib/dl_utils/dl_utils_tests/unit/test_stuff.py
datalens-tech/datalens-backend
from .dynamic_enum import (
    AutoEnumValue,
    DynamicEnum,
)


__all__ = (
    "AutoEnumValue",
    "DynamicEnum",
)

lib/dynamic_enum/dynamic_enum/__init__.py
datalens-tech/datalens-backend
"""
Intended as a replacement for `Enum`-based constants that are being spread out
across different packages and require distributed definition of values.

Usage examples.
Enum-like definition of values for compatibility with existing code:

    class Color(DynamicEnum):
        blue = AutoEnumValue()
        green = AutoEnumValue()

    Color.blue.value  # == 'blue'


Definition of values outside the class
(the case where values need to be defined in different packages):

    Color.declare('white')  # .value  == 'white'


After a value has been explicitly declared once (and only once)
via the `declare` classmethod or `AutoEnumValue` descriptor within the class itself,
the constant can be instantiated simply by using the constructor:

    Color('white')
    Color('green')

Calling the constructor on an undeclared value of declaring a value more than once
will raise a `ValueError`. This helps prevent the creation of invalid constant values
in arbitrary places of the code or when deserializing data.

Also, for better compatibility with enums, value deduplication is implemented
so that constants of the same class with the same value are always the same instance,
which makes possible comparisons via the `is` operator:

    Color('blue') is Color.blue  # == True


See tests for more specifics.
"""

from __future__ import annotations

from types import MappingProxyType
from typing import (
    Any,
    Callable,
    ClassVar,
    Generator,
    Optional,
    Sequence,
    Type,
    TypeVar,
)


_ANY_TV = TypeVar("_ANY_TV")
_DYNAMIC_ENUM_TV = TypeVar("_DYNAMIC_ENUM_TV", bound="DynamicEnum")


class AutoEnumValue:
    """
    Descriptor that enables to define `DynamicEnum` subclass instances
    within the class itself `Enum`-style.
    """

    __slots__ = ("__name",)

    __name: Optional[str]

    def __init__(self) -> None:
        self.__name = None

    def __set_name__(self, owner: DynamicEnum, name: str) -> None:
        self.__name = name
        owner.declare(name)

    def __get__(self, instance: Optional[DynamicEnum], owner: Type[_DYNAMIC_ENUM_TV]) -> _DYNAMIC_ENUM_TV:
        if self.__name is None:  # has not been set yet
            raise RuntimeError("Property is not bound to any class")

        assert self.__name is not None
        return owner(self.__name)

    def __set__(self, instance: Optional[DynamicEnum], value: Any) -> None:
        raise AttributeError


def _get_value_from_dyn_enum_args(args: Sequence[Any], kwargs: dict[str, Any]) -> str:
    """Auxiliary function used by `DynamicEnum` and its metaclass"""

    if args:
        value = args[0]
    elif "value" in kwargs:
        value = kwargs["value"]
    else:
        raise TypeError("name argument not received")

    return value


class DynamicEnumMetaclass(type):
    """
    This metaclass is used to:
    1. Override instance creation (i.e. avoid creating a new instance
       when one with the same value already exists).
    2. Implement magic methods for the `DynamicEnum` class. They cannot be defined
       as classmethods on the class, they must be defined in the metaclass.
    3. Enforce `__slots__` for all subclasses.
    """

    __instances: ClassVar[dict[DynamicEnumMetaclass, dict[str, Any]]] = {}  # For deduplication

    def __new__(mcs, name: str, bases: Any, attrs: Any) -> DynamicEnumMetaclass:
        if len(bases) > 1:
            raise TypeError("Cannot have multiple bases for subclasses")

        # Enforce slots
        slots = attrs.get("__slots__")
        if not bases:
            # This must be the base class
            assert slots == ("__value",)
        elif slots is not None:
            if slots != ():
                raise TypeError(
                    f"Custom instance attributes are not supported for class {name}. " f"Got slots: {slots}"
                )
        else:
            # A subclass. Enforce slots without any additional instance attributes
            attrs = dict(attrs, __slots__=())

        return super().__new__(mcs, name, bases, attrs)  # type: ignore

    def __call__(cls: Type[_ANY_TV], *args: Any, **kwargs: Any) -> _ANY_TV:
        value = _get_value_from_dyn_enum_args(args, kwargs)

        # Deduplicate instances with the same value
        if cls not in DynamicEnumMetaclass.__instances:
            DynamicEnumMetaclass.__instances[cls] = {}  # type: ignore
        cls_instances = DynamicEnumMetaclass.__instances[cls]  # type: ignore
        if value not in cls_instances:
            cls_instances[value] = super().__call__(*args, **kwargs)  # type: ignore

        return cls_instances[value]  # type: ignore

    def __getitem__(cls: Type[_ANY_TV], item: str) -> _ANY_TV:
        return cls(item)  # type: ignore

    def __contains__(cls: DynamicEnumMetaclass, item: str) -> bool:
        return cls.is_declared(item)  # type: ignore

    def __iter__(cls: Type[_ANY_TV]) -> Generator[_ANY_TV, None, None]:
        yield from cls.iter_items()  # type: ignore

    @property
    def __members__(cls) -> MappingProxyType:  # For compatibility with Enum
        return MappingProxyType({elem.name: elem.value for elem in cls})  # type: ignore


class DynamicEnum(metaclass=DynamicEnumMetaclass):
    __slots__ = ("__value",)
    __value: str

    __declared_values: ClassVar[dict[Type[DynamicEnum], set[str]]] = {}
    __subclassable: ClassVar[bool] = True

    def __init_subclass__(cls) -> None:
        """
        Prevent the subclassing of subclasses.
        """

        if not cls.__subclassable:
            raise TypeError(f"Cannot subclass {cls.__name__}")

        # All subclasses of DynamicEnum are not subclassable
        cls.__subclassable = False

        # This way we avoid situations like:
        # MyString.my_value != MyStringSubclass.my_value

    def __new__(cls: Type[_DYNAMIC_ENUM_TV], *args: Any, **kwargs: Any) -> _DYNAMIC_ENUM_TV:
        """
        Validates creation of instances:
        1. Forbids direct instantiation of `DynamicEnum`.
        2. Forbids instantiation with undeclared values.
        """

        if cls is DynamicEnum:
            raise TypeError("Cannot instantiate DynamicEnum base class")

        value = _get_value_from_dyn_enum_args(args, kwargs)

        if not cls.is_declared(value):
            raise ValueError(f"Cannot instantiate {cls.__name__} with undeclared value {value!r}")

        return super().__new__(cls)

    def __init__(self, value: str):
        if type(value) is not str:
            raise TypeError(f"Invalid value type {type(value)}")

        value_attr_mangled_name = f"_{DynamicEnum.__name__}__value"  # See "name mangling in python"
        super().__setattr__(value_attr_mangled_name, value)  # Because direct usage of __setattr__ is forbidden

    @classmethod
    def is_declared(cls, value: str) -> bool:
        return cls in cls.__declared_values and value in cls.__declared_values[cls]

    @classmethod
    def iter_items(cls: Type[_DYNAMIC_ENUM_TV]) -> Generator[_DYNAMIC_ENUM_TV, None, None]:
        if cls not in cls.__declared_values:
            yield from ()
        for value in cls.__declared_values[cls]:
            yield cls(value)

    @classmethod
    def declare(cls: Type[_DYNAMIC_ENUM_TV], value: str) -> _DYNAMIC_ENUM_TV:
        if cls not in cls.__declared_values:
            cls.__declared_values[cls] = set()

        cls_declared_values = cls.__declared_values[cls]
        if value not in cls_declared_values:
            cls_declared_values.add(value)
        else:
            raise ValueError(f"{cls.__name__} value {value!r} has already been declared")

        return cls(value)

    @property
    def value(self) -> str:
        return self.__value

    @property
    def name(self) -> str:  # Same as value, for compatibility with Enum
        return self.__value

    def __setattr__(self, key: str, value: Any) -> None:
        raise AttributeError(f"{self.__class__.__name__} is immutable")

    def __eq__(self, other: Any) -> bool:
        if type(other) != type(self):
            return False
        return other.__value == self.__value

    def __hash__(self) -> int:
        return hash(self.__value)

    def __str__(self) -> str:
        return repr(self)

    def __repr__(self) -> str:
        return f"{type(self).__name__}({self.__value!r})"

    def __reduce__(self: _DYNAMIC_ENUM_TV) -> tuple[Callable[[], _DYNAMIC_ENUM_TV], tuple[str]]:
        return self.__class__, (self.__value,)

lib/dynamic_enum/dynamic_enum/dynamic_enum.py
datalens-tech/datalens-backend
import pickle

from dynamic_enum import (
    AutoEnumValue,
    DynamicEnum,
)
import pytest


def test_auto_value():
    class Color(DynamicEnum):
        blue = AutoEnumValue()

    color = Color.blue
    assert color.value == "blue"


def test_deduplication():
    class Weight(DynamicEnum):
        heavy = AutoEnumValue()

    assert Weight("heavy") is Weight.heavy


def test_declaration():
    class Size(DynamicEnum):
        pass

    with pytest.raises(ValueError):
        Size("medium")  # value 'medium' has not been declared yet

    medium_size = Size.declare("medium")
    assert medium_size.value == "medium"

    # Now we can call the constructor directly, since the value has already been declared
    medium_size = Size("medium")
    assert medium_size.value == "medium"

    with pytest.raises(ValueError):
        Size.declare("medium")  # value 'medium' has already been declared


def test_subclassing():
    class Something(DynamicEnum):
        smth = AutoEnumValue()

    with pytest.raises(TypeError):

        class SpecificSomething(Something):
            pass


def test_comparison():
    class ThisWorld(DynamicEnum):
        thing = AutoEnumValue()

    class OtherWorld(DynamicEnum):
        thing = AutoEnumValue()

    assert ThisWorld.thing == ThisWorld.thing
    assert ThisWorld.thing is ThisWorld.thing

    assert OtherWorld.thing != ThisWorld.thing
    assert OtherWorld.thing is not ThisWorld.thing


def test_stringification():
    class What(DynamicEnum):
        that = AutoEnumValue()

    assert str(What("that")) == "What('that')"
    assert repr(What("that")) == "What('that')"


def test_getitem():
    class What(DynamicEnum):
        pass

    What.declare("that")

    assert What["that"] == What("that")


def test_containment():
    class Galaxy(DynamicEnum):
        pass

    Galaxy.declare("milky_way")
    Galaxy.declare("andromeda")

    assert "milky_way" in Galaxy
    assert "andromeda" in Galaxy
    assert "far_far_away" not in Galaxy


def test_iter():
    class Everything(DynamicEnum):
        pass

    Everything.declare("this")
    Everything.declare("that")

    assert set(Everything) == {Everything("this"), Everything("that")}


def test_no_instance_dict():
    class It(DynamicEnum):
        pass

    It.declare("thing")
    assert not hasattr(It("thing"), "__dict__")


def test_custom_slots():
    with pytest.raises(TypeError):

        class ExtendedEnum(DynamicEnum):
            __slots__ = ("this_other thing",)


class Picklable(DynamicEnum):
    thing = AutoEnumValue()


def test_pickle():
    data = pickle.dumps(Picklable.thing)
    loaded = pickle.loads(data)

    assert loaded is Picklable.thing

lib/dynamic_enum/dynamic_enum_tests/unit/test_dynamic_enum.py
datalens-tech/datalens-backend
"""Synchronizing cache generation to reduce work"""
from __future__ import annotations

__version__ = "3.5"

lib/redis-cache-lock/redis_cache_lock/__init__.py
datalens-tech/datalens-backend
from __future__ import annotations

from enum import IntEnum, unique


@unique
class ReqScriptResultBase(IntEnum):
    pass


@unique
class ReqScriptResult(ReqScriptResultBase):
    cache_hit = 130
    successfully_locked = 131
    lock_wait = 132


@unique
class ReqResultInternal(ReqScriptResultBase):
    # Special values for internal processing:
    starting = 101  # A lock request might be in flight.
    requesting = 102  # A lock request might be in flight.
    awaiting = 103  # A pubsub subscription is being awaited.

    # *MUST* have all values of the `ReqScriptResult`
    cache_hit = 130
    successfully_locked = 131
    lock_wait = 132

    # Special situations:
    cache_hit_slave = 333
    cache_hit_after_wait = 334
    network_call_timeout = 335
    lock_wait_timeout = 336
    lock_wait_unexpected_message = 337
    failure_signal = 338

    # Final situations, primarily for logging. Note that under
    # `enable_background_tasks`, these will become available during the
    # finalization task.
    generated_saved = 439
    generated_force_saved = 440
    failure_to_initialize = 441
    failure_marker_sent = 442


@unique
class RenewScriptResult(IntEnum):
    extended = 140
    expired = 141
    locked_by_another = 142


@unique
class SaveScriptResult(IntEnum):
    success = 150
    token_mismatch = 151
    not_locked = 152

lib/redis-cache-lock/redis_cache_lock/enums.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio


class CacheError(Exception):
    """Common base exception"""


class NetworkCallTimeoutError(asyncio.TimeoutError, CacheError):
    """Timed out while waiting on a network request"""


class CacheRedisError(CacheError):
    """Redis communication failure"""


class CacheLockLost(CacheError):
    """Previously owned lock is not owned anymore, abort"""


# class CacheTimeoutError(CacheError):
#     """ Some waiting failed within the time limit """

lib/redis-cache-lock/redis_cache_lock/exc.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
from contextlib import AsyncExitStack, asynccontextmanager
from typing import (
    TYPE_CHECKING,
    Any,
    Awaitable,
    AsyncGenerator,
    Callable,
    ClassVar,
    Dict,
    Optional,
    Tuple,
    Type,
    TypeVar,
)

import attr

from .enums import ReqResultInternal, ReqScriptResult
from .exc import NetworkCallTimeoutError
from .redis_utils import SubscriptionManager
from .scripts import ALIVE_PREFIX, DATA_PREFIX, FAIL_PREFIX
from .scripts_support import (
    FailScript,
    ForceSaveScript,
    RenewScript,
    ReqScript,
    SaveScript,
)
from .utils import await_on_exit, new_self_id, task_cm

if TYPE_CHECKING:
    from redis.asyncio import Redis

    from .redis_utils import SubscriptionManagerBase
    from .types import TCacheResult, TClientACM, TGenerateFunc, TGenerateResult


_WNC_RET_T = TypeVar("_WNC_RET_T")


@attr.s(auto_attribs=True)
class RedisCacheLock:
    client_acm: TClientACM = attr.ib(repr=False)

    key: str
    resource_tag: str  # namespace for the keys
    data_ttl_sec: float

    lock_ttl_sec: float = 3.5
    lock_renew_interval: Optional[float] = None
    network_call_timeout_sec: float = 2.5

    debug_log: Optional[Callable[[str, Dict[str, Any]], None]] = attr.ib(default=None, repr=False)
    bg_task_callback: Optional[Callable[[asyncio.Task], None]] = attr.ib(default=None, repr=False)
    enable_background_tasks: bool = False
    enable_slave_get: bool = True

    req_script_cls: ClassVar[Type[ReqScript]] = ReqScript
    req_script_situation: ClassVar[Type[ReqScriptResult]] = ReqScriptResult
    req_situation: ClassVar[Type[ReqResultInternal]] = ReqResultInternal
    renew_script_cls: ClassVar[Type[RenewScript]] = RenewScript
    save_script_cls: ClassVar[Type[SaveScript]] = SaveScript
    fail_script_cls: ClassVar[Type[FailScript]] = FailScript
    force_save_script_cls: ClassVar[Type[ForceSaveScript]] = ForceSaveScript
    chan_data_prefix: ClassVar[bytes] = DATA_PREFIX
    chan_alive_prefix: ClassVar[bytes] = ALIVE_PREFIX
    chan_fail_prefix: ClassVar[bytes] = FAIL_PREFIX
    subscription_manager_cls: Type[SubscriptionManagerBase] = SubscriptionManager

    data_tag: str = "/data:"
    signal_tag: str = "/notif:"
    lock_tag: str = "/lock:"

    # TODO: split this class into a config + CM pair

    _situation: Optional[ReqResultInternal] = None
    _cm_stack: Optional[AsyncExitStack] = None
    _client: Optional[Redis] = None
    _self_id: Optional[str] = None

    @staticmethod
    def new_self_id() -> str:
        return new_self_id()

    @property
    def situation(self) -> Optional[ReqResultInternal]:
        return self._situation

    @situation.setter
    def situation(self, value: ReqResultInternal) -> None:
        if value == self._situation:
            return
        self._log("Setting situation to %r", value, situation=value)
        self._situation = value

    def clone(self, **kwargs: Any) -> RedisCacheLock:
        result = attr.evolve(self, **kwargs)
        result._cleanup()  # pylint: disable=protected-access
        return result

    def _cleanup(self) -> None:
        self._client = None
        self._cm_stack = None
        self._self_id = None
        self._situation = None

    def _log(self, msg: str, *args: Any, **details: Any) -> None:
        if self.debug_log is not None:
            self.debug_log(  # pylint: disable=not-callable
                msg % args,
                dict(
                    details,
                    key=self.key,
                    self_id=self._self_id,
                ),
            )

    @property
    def data_key(self) -> str:
        return self.resource_tag + self.data_tag + self.key

    @property
    def signal_key(self) -> str:
        return self.resource_tag + self.signal_tag + self.key

    @property
    def lock_key(self) -> str:
        return self.resource_tag + self.lock_tag + self.key

    @asynccontextmanager
    async def _client_acm_managed(
        self, master: bool = True, exclusive: bool = True
    ) -> AsyncGenerator[Redis, None]:
        async with AsyncExitStack() as temp_cm_stack:  # easier way to timeout a CM
            client: Redis = await self._wait_network_call(
                temp_cm_stack.enter_async_context(
                    self.client_acm(master=master, exclusive=exclusive)
                )
            )
            yield client

    def process_in_background(self, coro: Awaitable, name: Optional[str] = None) -> Any:
        """An overridable method for finalization background-task creation"""
        if name is None:
            name = repr(coro)

        async def background_wrapper() -> Any:
            try:
                return await coro
            except Exception as err:  # pylint: disable=broad-except
                self._log("Exception in background task %r: %r", name, err)
                return None

        task = asyncio.create_task(background_wrapper(), name=name)
        callback = self.bg_task_callback
        if callback is not None:
            callback(task)  # pylint: disable=not-callable
        return task

    async def _wait_network_call(self, coro: Awaitable[_WNC_RET_T]) -> _WNC_RET_T:
        try:
            return await asyncio.wait_for(coro, timeout=self.network_call_timeout_sec)
        except asyncio.TimeoutError as err:
            raise NetworkCallTimeoutError() from err

    async def _finalize_maybe_in_background(
        self,
        coro: Awaitable,
        name: Optional[str] = None,
    ) -> Tuple[bool, Any]:
        """
        Depending on the settings, either await the `coro` directly (but with cancellation shield),
        or run it in a background asyncio task.
        """
        if not self.enable_background_tasks:
            # The `wait_for(shield(coro))` results in waiting until timeout but
            # leaving the coro in background when the timeout is reached.
            result = await self._wait_network_call(asyncio.shield(coro))
            return False, result

        return True, self.process_in_background(coro, name=name)

    async def _postpone_to_finalization(self, coro: Awaitable) -> Any:
        cm_stack = self._cm_stack
        assert cm_stack is not None, "must be initialized for this method"
        return await cm_stack.enter_async_context(await_on_exit(coro))

    async def get_data_slave(self) -> Optional[bytes]:
        data_key = self.data_key
        cli: Redis
        async with self._client_acm_managed(master=False, exclusive=False) as cli:
            return await self._wait_network_call(cli.get(data_key))

    async def _get_data(
        self,
    ) -> Tuple[ReqResultInternal, Optional[bytes], Optional[SubscriptionManagerBase]]:
        self_id = self._self_id
        assert self_id is not None, "must be initialized for this method"
        cm_stack = self._cm_stack
        assert cm_stack is not None, "must be initialized for this method"
        cli = self._client
        assert cli is not None, "must be initialized for this method"

        subscription: Optional[SubscriptionManagerBase] = None

        lock_key = self.lock_key
        data_key = self.data_key
        lock_ttl_sec = self.lock_ttl_sec

        req_script = self.req_script_cls(cli=cli)

        self._log("Calling req_script", lock_ttl_sec=lock_ttl_sec)
        self.situation = self.req_situation.requesting
        situation, result = await self._wait_network_call(
            req_script(
                lock_key=lock_key,
                data_key=data_key,
                self_id=self_id,
                lock_ttl_sec=lock_ttl_sec,
            )
        )
        self.situation = self.req_situation(situation.value)

        if situation == self.req_script_situation.lock_wait:
            signal_key = self.signal_key
            self._log("Subscribing to notify channel (lock_wait)")
            subscription = await self._wait_network_call(
                self.subscription_manager_cls.create(
                    cm_stack=cm_stack,
                    client_acm=self._client_acm_managed,
                    channel_key=signal_key,
                )
            )

            self._log("Re-checking the situation after subscription")
            # In case the result appeared between first `get` and `psubscribe`, check for it again.
            self.situation = self.req_situation.requesting
            situation, result = await self._wait_network_call(
                req_script(
                    lock_key=lock_key,
                    data_key=data_key,
                    self_id=self_id,
                    lock_ttl_sec=lock_ttl_sec,
                )
            )
            self.situation = self.req_situation(situation.value)
            if situation != self.req_script_situation.lock_wait:
                self._log(
                    "Situation changed while subscribing (to %r), unsubscribing",
                    situation,
                    situation=situation,
                )
                if subscription is not None:
                    # It is possible the `req_script` above has locked, and so
                    # the subscription would stay while data is being generated.
                    # This will try unsubscribing at this point, but it will
                    # also be awaited in the `cm_stack` finalization.
                    await self._finalize_maybe_in_background(
                        subscription.exit(),
                        name="subscription.exit() (situation changed)",
                    )
                subscription = None

        internal_situation = self.req_situation(situation.value)
        self._log(
            "Situation from get_data: %r",
            internal_situation,
            situation=internal_situation,
        )
        return internal_situation, result, subscription

    async def _wait_for_result(
        self,
        sub: SubscriptionManagerBase,
    ) -> Tuple[ReqResultInternal, Optional[bytes]]:
        # Lock should be renewed more often than `lock_ttl_sec`,
        # so waiting for the ttl duration should be sufficient.
        timeout = self.lock_ttl_sec
        data = None

        try:
            while True:
                self._log("Waiting for signal")
                self.situation = self.req_situation.awaiting
                message = await sub.get(timeout=timeout)

                if message is None:
                    self._log(
                        "Timed out waiting for signal, assuming the locking process is gone",
                        timeout=timeout,
                    )
                    situation = self.req_situation.lock_wait_timeout
                elif message.startswith(self.chan_data_prefix):
                    self._log("Data signal.")
                    data = message[len(self.chan_data_prefix) :]
                    situation = self.req_situation.cache_hit_after_wait
                elif message.startswith(self.chan_alive_prefix):
                    self._log("Alive signal: %r", message, timeout=timeout)
                    continue  # keep waiting
                elif message.startswith(self.chan_fail_prefix):
                    self._log("Fail signal: %r", message, timeout=timeout)
                    situation = self.req_situation.failure_signal
                else:
                    self._log("Unexpected signal message: %r", message, timeout=timeout)
                    situation = self.req_situation.lock_wait_unexpected_message

                self.situation = situation
                return situation, data

        finally:
            await self._finalize_maybe_in_background(
                sub.exit(),
                name="sub.exit() (done waiting)",
            )

        raise Exception("Programming Error")

    async def _get_data_full(self) -> Tuple[ReqResultInternal, Optional[bytes]]:
        assert self._situation == self.req_situation.starting

        cli = self._client
        assert cli is not None, "must be initialized for this method"

        situation, result, subscription = await self._get_data()

        if situation == self.req_situation.lock_wait:
            # At this point, the `sub_client` should already be subscribed
            # without any race conditions.
            assert subscription is not None
            situation, result = await self._wait_for_result(sub=subscription)

        # Can fall through to either `cache_hit` or whichever
        # `self.channel_poll_timeout_situation` defines.
        return situation, result

    async def _renew_lock(self) -> int:
        self_id = self._self_id
        assert self_id is not None, "must be initialized for this method"
        cli = self._client
        assert cli is not None, "must be initialized for this method"

        lock_key = self.lock_key
        signal_key = self.signal_key
        lock_ttl_sec = self.lock_ttl_sec
        renew_script = self.renew_script_cls(cli=cli)
        self._log("Calling renew_script", lock_ttl_sec=lock_ttl_sec)
        return await self._wait_network_call(
            renew_script(
                lock_key=lock_key,
                signal_key=signal_key,
                self_id=self_id,
                lock_ttl_sec=lock_ttl_sec,
            )
        )

    async def _lock_pinger(self) -> None:
        cli = self._client
        assert cli is not None, "must be initialized for this method"

        lock_ttl_sec = self.lock_ttl_sec
        renew_interval = self.lock_renew_interval or lock_ttl_sec * 0.5
        while True:  # until cancelled, really
            await asyncio.sleep(renew_interval)
            try:
                renew_res = await self._renew_lock()
                # TODO: callback for non-ok renew.
                self._log(
                    "renew script result: %r", renew_res, renew_interval=renew_interval
                )
            except Exception as err:  # pylint: disable=broad-except
                self._log("lock_pinger error: %r", err, renew_interval=renew_interval)

    async def _save_data(
        self,
        data: bytes,
        ttl_sec: Optional[float] = None,
    ) -> Any:
        self_id = self._self_id
        assert self_id is not None, "must be initialized for this method"
        cli = self._client
        assert cli is not None, "must be initialized for this method"

        lock_key = self.lock_key
        signal_key = self.signal_key
        data_key = self.data_key
        self._log("Calling save_script", data_len=len(data))
        save_script = self.save_script_cls(cli=cli)
        # Not wrapping in `self._wait_network_call` at this point, as
        # cancelling this action is not appropriate.
        return await self._postpone_to_finalization(
            save_script(
                lock_key=lock_key,
                signal_key=signal_key,
                data_key=data_key,
                self_id=self_id,
                data=data,
                data_ttl_sec=ttl_sec or self.data_ttl_sec,
            )
        )

    async def _force_save_data(
        self,
        data: bytes,
        ttl_sec: Optional[float] = None,
    ) -> Any:
        cli = self._client
        assert cli is not None, "must be initialized for this method"

        signal_key = self.signal_key
        data_key = self.data_key
        self._log("Calling force_save_script", data_len=len(data))
        force_save_script = self.force_save_script_cls(cli=cli)
        return await self._postpone_to_finalization(
            force_save_script(
                signal_key=signal_key,
                data_key=data_key,
                data=data,
                data_ttl_sec=ttl_sec or self.data_ttl_sec,
            )
        )

    async def _save_failure(self, ignore_errors: bool = False) -> Any:
        self_id = self._self_id
        assert self_id is not None, "must be initialized for this method"
        cli = self._client
        assert cli is not None, "must be initialized for this method"

        lock_key = self.lock_key
        signal_key = self.signal_key
        self._log("Calling fail_script")
        fail_script = self.fail_script_cls(cli=cli)
        # TODO: save the script result to `self.situation`
        return await self._postpone_to_finalization(
            fail_script(
                lock_key=lock_key,
                signal_key=signal_key,
                self_id=self_id,
                ignore_errors=ignore_errors,
            )
        )

    async def _handle_initialize(
        self,
        situation: ReqResultInternal,
        result: Optional[bytes],
    ) -> Optional[bytes]:
        if situation in (
            self.req_situation.cache_hit,
            self.req_situation.cache_hit_after_wait,
        ):
            assert result is not None
            return result

        if situation == self.req_situation.successfully_locked:
            assert result is None

            self_id = self._self_id
            assert self_id is not None
            cm_stack = self._cm_stack
            assert cm_stack is not None

            await cm_stack.enter_async_context(task_cm(self._lock_pinger()))
            return None

        if situation in (
            self.req_situation.lock_wait_timeout,
            self.req_situation.failure_signal,
        ):
            assert result is None
            return None

        raise Exception(
            "RedisCacheLock error: Unexpected get_data_full outcome", situation
        )

    async def initialize(self) -> Optional[bytes]:
        """
        Attempt to read the data or lock the key or wait for data.

        Returns the cached (serialized) data, if available.

        After calling this function, `finalize` *must* be called in all cases.
        This pair essentially works as a context manager with different arguments.
        """
        assert self._situation is None, "not re-entrable"
        self.situation = self.req_situation.starting

        if self.enable_slave_get:
            result = await self.get_data_slave()
            if result is not None:
                self._log("Found data at slave")
                self.situation = self.req_situation.cache_hit_slave
                return result

        self_id = self.new_self_id()
        self._self_id = self_id

        cm_stack = AsyncExitStack()
        self._cm_stack = cm_stack  # `AsyncExitStack` allows aexit-without-aenter.

        await cm_stack.__aenter__()  # The corresponding `finally:` is the `def finalize`.

        cli: Redis = await cm_stack.enter_async_context(
            self._client_acm_managed(master=True, exclusive=False)
        )
        self._client = cli

        situation, result = await self._get_data_full()
        self.situation = situation
        return await self._handle_initialize(
            situation=situation,
            result=result,
        )

    def decide_force_save(  # pylint: disable=unused-argument:
        self,
        situation: ReqResultInternal,
    ) -> bool:
        if situation in (
            # Will save the data 'nicely', i.e. only if the lock is not held by another process.
            self.req_situation.successfully_locked,
            # No data generated, probably shouldn't save it back.
            self.req_situation.cache_hit_slave,
            self.req_situation.cache_hit,
            self.req_situation.cache_hit_after_wait,
        ):
            return False

        # The safe-ish default: if something abnormal happened, but we
        # generated the data, might as well save it (although it can
        # technically be an error-marker).
        return True

    async def _handle_finalize(
        self,
        result: Optional[
            bytes
        ],  # must be specified but in some situations can be empty.
        ttl_sec: Optional[float] = None,
        force_save: Optional[bool] = None,
    ) -> Any:
        situation = self._situation

        if situation is None:
            # Failed before starting the initialization at all. Nothing to do.
            self.situation = self.req_situation.failure_to_initialize
            return

        if result is not None:
            if force_save is None:
                force_save = self.decide_force_save(situation)
            if force_save:
                self.situation = self.req_situation.generated_force_saved
                return await self._force_save_data(
                    data=result,
                    ttl_sec=ttl_sec,
                )

        if result is None and self._self_id is not None:
            # Initialization failure, data generation failure, etcetera.
            # Ensure we don't hold the lock, just in case.
            self.situation = self.req_situation.failure_marker_sent
            return await self._save_failure(ignore_errors=True)

        if situation == self.req_situation.successfully_locked:
            assert result is not None
            self.situation = self.req_situation.generated_saved
            return await self._save_data(data=result, ttl_sec=ttl_sec)

    async def finalize(
        self,
        result_serialized: Optional[bytes],
        ttl_sec: Optional[float] = None,
    ) -> None:
        try:
            return await self._handle_finalize(
                result=result_serialized,
                ttl_sec=ttl_sec,
            )
        finally:
            cm_stack = self._cm_stack

            self._cleanup()

            if cm_stack is not None:
                # Should not matter whether there was an exception.
                # The saving should happen in this finalization, because it
                # might be in background, but should be finished before the
                # client is released.
                await self._finalize_maybe_in_background(
                    cm_stack.__aexit__(None, None, None), name="cm_stack exit"
                )

    async def _call_generate_func(
        self, generate_func: TGenerateFunc
    ) -> TGenerateResult:
        """Call `generate_func` and validate the result"""
        result = await generate_func()
        if not isinstance(result, tuple):
            raise ValueError(
                (
                    "`generate_func` returned a non-tuple; "
                    "it should return a 2-item tuple `(serialized, unserialized)`"
                ),
                dict(generate_func=generate_func, result_type=type(result)),
            )
        if len(result) != 2:
            raise ValueError(
                (
                    f"`generate_func` returned a {len(result)}-item tuple; "
                    "it should return a 2-item tuple `(serialized, unserialized)`"
                ),
                dict(generate_func=generate_func),
            )
        serialized, raw = result
        if not isinstance(serialized, bytes):
            raise ValueError(
                "`generate_func` should return serialized value `bytes` as the first tuple item",
                dict(generate_func=generate_func, serialized_type=type(serialized)),
            )
        return serialized, raw

    # The code beyond this point can be used as a reference for calling
    # `initialize` + `finalize` directly.

    async def generate_with_lock(self, generate_func: TGenerateFunc) -> TCacheResult:
        result = None
        try:
            # General ordering note:
            # it is possible that the command reaches the server and gets executed,
            # but the code raises an exception (e.g. asyncio timeout). For that
            # reason, it is better to free the resources even if the acquiring
            # failed, and do it in a shielded background task.
            result = await self.initialize()
            if result is not None:
                return result, None  # cached

            result, unserialized_result = await self._call_generate_func(generate_func)
            return result, unserialized_result
        finally:
            await self.finalize(result)

lib/redis-cache-lock/redis_cache_lock/main.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import time
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Optional, Sequence, Tuple

import redis.asyncio
import attr

from .types import Protocol
from .utils import PreExitable

if TYPE_CHECKING:
    from contextlib import AsyncExitStack  # pylint: disable=ungrouped-imports

    from redis.asyncio import Redis

    from .types import TClientACM


async def eval_script(cli: Redis, script: Any, keys: List[Any], args: List[Any]) -> Any:
    return await cli.eval(script, len(keys), *(keys + args))


def make_simple_cli_acm(url: str) -> TClientACM:
    """Single-host Redis Client ACM that supports both aioredis 1 and aioredis 2"""

    @asynccontextmanager
    async def cli_acm(**_: Any) -> AsyncGenerator[Redis, None]:
        # Redis client bound to pool of connections (auto-reconnecting).
        # Reference: https://aioredis.readthedocs.io/en/latest/migration/#connecting-to-redis
        pool = redis.asyncio.ConnectionPool.from_url(url)
        rcli = redis.asyncio.Redis(connection_pool=pool)
        try:
            yield rcli
        finally:
            await pool.disconnect()

    return cli_acm


class TSentinel(Protocol):
    def master_for(self, service_name: str) -> Redis:
        pass

    def slave_for(self, service_name: str) -> Redis:
        pass


async def make_sentinel(
    sentinels: Sequence[Tuple[str, int]], **kwargs: Any
) -> TSentinel:
    return redis.asyncio.sentinel.Sentinel(sentinels, **kwargs)


def make_sentinel_cli_acm(sentinel_cli: TSentinel, service_name: str) -> TClientACM:
    @asynccontextmanager
    async def sentinel_client_acm(
        *,
        master: bool = True,
        # aioredis does its own tracking of client reusability, so no need
        # to consider the `exclusive` flag.
        # pylint: disable=unused-argument
        exclusive: bool = True,
    ) -> AsyncGenerator[Redis, None]:
        if master:
            cli = sentinel_cli.master_for(service_name)
        else:
            cli = sentinel_cli.slave_for(service_name)
        yield cli

    return sentinel_client_acm


@attr.s(auto_attribs=True, frozen=True, slots=True)
class SubscriptionManagerBase:
    """Helper to manage a redis subscription"""

    cli: Redis
    cli_cm: PreExitable
    psub: Any
    psub_cm: PreExitable

    @classmethod
    @asynccontextmanager
    async def _psub_acm(
        cls,
        cli: Redis,  # pylint: disable=unused-argument
        channel_key: str,  # pylint: disable=unused-argument
    ) -> AsyncGenerator[Any, None]:
        if False:  # ensure this is a generator (for mypy)  # pylint: disable=using-constant-test
            yield None
        raise NotImplementedError

    @classmethod
    async def create(
        cls,
        cm_stack: AsyncExitStack,
        client_acm: TClientACM,
        channel_key: str,
    ) -> SubscriptionManagerBase:
        cli_cm = PreExitable(client_acm(master=True, exclusive=True))
        cli: Redis = await cm_stack.enter_async_context(cli_cm)
        psub_cm = PreExitable(cls._psub_acm(cli=cli, channel_key=channel_key))
        psub: Any = await cm_stack.enter_async_context(psub_cm)
        return cls(
            cli=cli,
            cli_cm=cli_cm,
            psub=psub,
            psub_cm=psub_cm,
        )

    async def get_direct(self, timeout: float) -> Optional[bytes]:
        raise NotImplementedError

    async def get(self, timeout: float) -> Optional[bytes]:
        try:
            return await asyncio.wait_for(
                self.get_direct(timeout=timeout), timeout=timeout
            )
        except asyncio.TimeoutError:
            return None

    async def exit(self) -> None:
        """
        End the subscription and release the client.
        Can be called multiple times.
        Allows freeing of resources before the passed `cm_stack` finishes.
        """
        # This is done in order and the exceptions are passed through.
        # The fallback closing is through the `cm_stack` which is done fully
        # even if some CM raises.
        await self.psub_cm.exit()
        await self.cli_cm.exit()


@attr.s(auto_attribs=True, frozen=True, slots=True)
class SubscriptionManagerLegacy(SubscriptionManagerBase):
    @classmethod
    @asynccontextmanager
    async def _psub_acm(
        cls,
        cli: Redis,
        channel_key: str,
    ) -> AsyncGenerator[Any, None]:
        try:
            channels: Tuple[Any, ...] = await cli.psubscribe(channel_key)  # type: ignore
            if len(channels) != 1:
                raise ValueError(
                    f"Expected a single channel; "
                    f"channel_key={channel_key!r}, channels={channels!r}"
                )
            channel = channels[0]
            yield channel
        finally:
            await cli.punsubscribe(channel_key)  # type: ignore

    async def get_direct(self, timeout: float) -> Optional[bytes]:
        try:
            # returns a `channel_key, message_data` tuple.
            item = await self.psub.get()
        except Exception:  # pylint: disable=broad-except
            # doesn't really matter which exception,
            # although this can often be `aioredis.errors.ChannelClosedError`.
            item = None

        if item is None:
            return None

        _, message = item
        return message


class SubscriptionManager(SubscriptionManagerBase):
    @classmethod
    @asynccontextmanager
    async def _psub_acm(
        cls,
        cli: Redis,
        channel_key: str,
    ) -> AsyncGenerator[Any, None]:
        async with cli.pubsub() as psub:
            await psub.subscribe(channel_key)
            try:
                yield psub
            finally:
                await psub.unsubscribe(channel_key)

    async def get_direct(self, timeout: float) -> Optional[bytes]:
        t01 = time.monotonic()
        # Apparently, `ignore_subscribe_messages=True` doesn't make the
        # `get_message` wait further; instead, it makes `get_message` return a
        # `None`.
        while True:
            message = await self.psub.get_message(
                ignore_subscribe_messages=True, timeout=timeout
            )
            if message is not None:
                break
            if time.monotonic() - t01 >= timeout:
                break
            await asyncio.sleep(0.001)
        if message is None:
            return None
        return message["data"]

    @classmethod
    async def create(
        cls,
        cm_stack: AsyncExitStack,
        client_acm: TClientACM,
        channel_key: str,
    ) -> SubscriptionManagerBase:
        """Convenience override for aioredis versions here"""
        return await super().create(
            cm_stack=cm_stack, client_acm=client_acm, channel_key=channel_key
        )

lib/redis-cache-lock/redis_cache_lock/redis_utils.py
datalens-tech/datalens-backend
from __future__ import annotations

# pubsub channel message prefixes
# for parsing simplicity, padding them to 10 bytes.
ALIVE_PREFIX = b"__alive___"
assert len(ALIVE_PREFIX) == 10
DATA_PREFIX = b"__data____"
assert len(DATA_PREFIX) == 10
FAIL_PREFIX = b"__fail____"
assert len(FAIL_PREFIX) == 10


# param: keys[1] - lock key
# param: keys[2] - cache data key
# params argv[1] - lock token (unique). Also used as debug identifier (add hostname+pid to it).
# params argv[2] - expiration in msec.
# Returns either of:
#  * `(ReqScriptResult.cache_hit.value, cache_data)`
#  * `(ReqScriptResult.successfully_locked.value, b'')`
#    * Which means the caller should go generate the data and then call CL_SAVE_SCRIPT.
#  * `(ReqScriptResult.lock_wait.value, b'')`, which means 'wait on the signal'
CL_REQ_SCRIPT = """ -- CL_REQ_SCRIPT
local cache_data = redis.call('get', KEYS[2])
if cache_data then
    return {130, cache_data}
end

local current_token = redis.call('get', KEYS[1])
if not current_token then
    redis.call('set', KEYS[1], ARGV[1])
    redis.call('pexpire', KEYS[1], ARGV[2])
    return {131, ''}
end

return {132, ''}
"""

# Renew the lock setting a new expiration time.
# param: keys[1] - lock key
# param: keys[2] - signal key (can be set to `''` to disable notify)
# param: argv[1] - lock token (unique)
# param: argv[2] - expiration in msec
# Returns either of:
#  * `(RenewScriptResult.expired.value, -1)` - no lock key
#    * Likely a reason to abort the locked process with an error. Or re-lock.
#  * `(RenewScriptResult.expired.value, -2)` - no TTL on lock (should never happen)
#  * `(RenewScriptResult.locked_by_another.value, another_process_token)` -
#    another process stole the lock.
#    * Definitely a reason to abort the locked process.
#  * `(RenewScriptResult.extended.value, previous_ttl)` - the normal case.
#    * Also publishes '__alive___' to the signal key if it is specified.
CL_RENEW_SCRIPT = """ -- CL_RENEW_SCRIPT
local current_ttl = redis.call('pttl', KEYS[1])
if current_ttl <= 0 then
    return {141, current_ttl}
end

local current_token = redis.call('get', KEYS[1])
if current_token ~= ARGV[1] then
    return {142, current_token}
end

redis.call('pexpire', KEYS[1], ARGV[2])
if KEYS[2] ~= '' then
    -- notify the subscribers that the generating worker is still alive
    redis.call('publish', KEYS[2], '__alive___' .. ARGV[1])
end

return {140, current_ttl}
"""

# Script to set the result and release the lock.
# param: keys[1] - lock key
# param: keys[2] - signal key
# param: keys[3] - cache data key
# param: argv[1] - lock token (unique)
# param: argv[2] - cache data
#   NOTE: it is recommended to cache the error-results too, with small TTL.
# param: argv[3] - cache data ttl
# Returns either of:
#  * `(SaveScriptResult.token_mismatch.value, another_process_token)` - another
#    process stole the lock, not saving.
#  * `(SaveScriptResult.success.value, watchers_count)` - succesfully saved and unlocked.
#  * `(SaveScriptResult.not_locked.value, watchers_count)` - the lock expired, saved anyway.

CL_SAVE_SCRIPT = """ -- CL_SAVE_SCRIPT
local current_token = redis.call('get', KEYS[1])

if current_token then
    if current_token ~= ARGV[1] then
        return {151, current_token}
    end
    redis.call('del', KEYS[1])
end

redis.call('set', KEYS[3], ARGV[2])
redis.call('pexpire', KEYS[3], ARGV[3])
local result = redis.call('publish', KEYS[2], '__data____' .. ARGV[2])

if not current_token then
    return {152, result}
end

return {150, result}
"""


# Script to release the lock without saving a result (indicating failure).
# param: keys[1] - lock key
# param: keys[2] - signal key
# param: argv[1] - lock token (unique)
# The result enum is shared with the `CL_SAVE_SCRIPT`
# Returns either of:
#  * `(SaveScriptResult.token_mismatch.value, '')` - the lock expired.
#  * `(SaveScriptResult.token_mismatch.value, another_process_token)` - another
#    process stole the lock.
#  * `(SaveScriptResult.success.value, watchers_count)` - succesfully unlocked and notified.
CL_FAIL_SCRIPT = """ -- CL_FAIL_SCRIPT
local current_token = redis.call('get', KEYS[1])
if not current_token then
    return {151, ''}
end
if current_token ~= ARGV[1] then
    return {151, current_token}
end

redis.call('del', KEYS[1])
local result = redis.call('publish', KEYS[2], '__fail____' .. ARGV[1])

return {150, result}
"""


# Script to set the result while ignoring the lock.
# param: keys[1] - signal key
# param: keys[2] - cache data key
# param: argv[1] - cache data
# param: argv[2] - cache data ttl
# Returns `watchers_count`
CL_FORCE_SAVE_SCRIPT = """ -- CL_FORCE_SAVE_SCRIPT
redis.call('set', KEYS[2], ARGV[1])
redis.call('pexpire', KEYS[2], ARGV[2])
return redis.call('publish', KEYS[1], '__data____' .. ARGV[1])
"""

lib/redis-cache-lock/redis_cache_lock/scripts.py
datalens-tech/datalens-backend
from __future__ import annotations

from typing import TYPE_CHECKING, Any, ClassVar, List, Optional, Tuple, Type, Union

import attr

from .enums import RenewScriptResult, ReqScriptResult, SaveScriptResult
from .exc import CacheLockLost, CacheRedisError
from .redis_utils import eval_script
from .scripts import (
    CL_FAIL_SCRIPT,
    CL_FORCE_SAVE_SCRIPT,
    CL_RENEW_SCRIPT,
    CL_REQ_SCRIPT,
    CL_SAVE_SCRIPT,
)

if TYPE_CHECKING:
    from enum import IntEnum

    from redis.asyncio import Redis


@attr.s(auto_attribs=True, frozen=True)
class ScriptsSupport:
    encoding: str = "utf-8"
    result_error_exc_cls: Type[Exception] = CacheRedisError

    def to_bytes(self, value: Union[bytes, str]) -> bytes:
        if isinstance(value, bytes):
            return value  # rare
        return value.encode(self.encoding)

    def to_text(self, value: Union[bytes, str], **kwargs: Any) -> str:
        if isinstance(value, str):
            return value
        return value.decode(self.encoding, **kwargs)

    def parse_script_result(self, res: Any, expected_length: int) -> Tuple[Any, ...]:
        if not isinstance(res, (tuple, list)):
            raise self.result_error_exc_cls(  # rare
                f"Unexpected script result type: {type(res)}, {res!r}"
            )
        if len(res) != expected_length:
            raise self.result_error_exc_cls(  # rare
                f"Unexpected script result length: {len(res)}, {res!r}"
            )
        return tuple(res)


@attr.s
class ScriptBase:
    script: ClassVar[str]
    situation_enum: ClassVar[Optional[Type[IntEnum]]]
    expected_result_length: ClassVar[int] = 2

    cli: Redis = attr.ib()

    support: ScriptsSupport = attr.ib(factory=ScriptsSupport)

    def _to_bytes(self, value: Union[bytes, str]) -> bytes:
        return self.support.to_bytes(value)

    def _to_text(self, value: Union[bytes, str], **kwargs: Any) -> str:
        return self.support.to_text(value, **kwargs)

    def _to_msec(self, value_sec: float) -> int:
        return int(value_sec * 1000)

    def _parse_script_result(self, res: Any) -> Tuple[Any, ...]:
        assert self.expected_result_length >= 1
        prepared_res = self.support.parse_script_result(
            res=res,
            expected_length=self.expected_result_length,
        )
        situation_value = prepared_res[0]
        params = prepared_res[1:]
        if self.situation_enum is None:
            situation = situation_value  # rare
        else:
            try:
                situation = self.situation_enum(situation_value)
            except ValueError as err:
                raise self.support.result_error_exc_cls(  # rare
                    f"Unexpected script result situation: {situation_value}, {res!r}"
                ) from err
        return (situation,) + params

    async def _eval_script_base(self, keys: List[Any], args: List[Any]) -> Any:
        return await eval_script(self.cli, self.script, keys, args)

    async def _eval_script(self, keys: List[Any], args: List[Any]) -> Tuple[Any, ...]:
        res = await self._eval_script_base(keys, args)
        return res, self._parse_script_result(res)


class ReqScript(ScriptBase):
    script: ClassVar[str] = CL_REQ_SCRIPT
    situation_enum: ClassVar[Type[ReqScriptResult]] = ReqScriptResult

    async def __call__(
        self,
        lock_key: str,
        data_key: str,
        self_id: str,
        lock_ttl_sec: float,
    ) -> Tuple[ReqScriptResult, Optional[bytes]]:
        _, (situation, param) = await self._eval_script(
            keys=[self._to_bytes(lock_key), self._to_bytes(data_key)],
            args=[self._to_bytes(self_id), self._to_msec(lock_ttl_sec)],
        )
        if situation == self.situation_enum.cache_hit:
            assert param is not None
        if situation == self.situation_enum.successfully_locked:
            assert param is None or param == b""
            param = None
        return situation, param


class RenewScript(ScriptBase):
    script: ClassVar[str] = CL_RENEW_SCRIPT
    situation_enum: ClassVar[Type[RenewScriptResult]] = RenewScriptResult
    lock_lost_exc_cls: Type[Exception] = CacheLockLost

    async def __call__(
        self,
        lock_key: str,
        signal_key: Optional[str],
        self_id: str,
        lock_ttl_sec: float,
    ) -> int:
        res, (situation, param) = await self._eval_script(
            keys=[self._to_bytes(lock_key), self._to_bytes(signal_key or "")],
            args=[self._to_bytes(self_id), self._to_msec(lock_ttl_sec)],
        )
        if situation == self.situation_enum.expired and param == -1:
            raise self.lock_lost_exc_cls("Lock found to be expired on renew")
        if situation == self.situation_enum.expired and param == -2:
            raise self.lock_lost_exc_cls("Lock found to have no TTL")
        if situation == self.situation_enum.locked_by_another:
            raise self.lock_lost_exc_cls(
                f"Lock found to be owned by another: {param!r}"
            )
        if situation == self.situation_enum.extended:
            return param  # old ttl
        raise self.support.result_error_exc_cls(
            f"Unexpected renew script result: {res!r}"
        )


class SaveScript(ScriptBase):
    script: ClassVar[str] = CL_SAVE_SCRIPT
    situation_enum: ClassVar[Type[SaveScriptResult]] = SaveScriptResult
    lock_lost_exc_cls: Type[Exception] = CacheLockLost

    async def __call__(
        self,
        lock_key: str,
        signal_key: str,
        data_key: str,
        self_id: str,
        data: bytes,
        data_ttl_sec: float,
    ) -> int:
        res, (situation, param) = await self._eval_script(
            keys=[
                self._to_bytes(lock_key),
                self._to_bytes(signal_key),
                self._to_bytes(data_key),
            ],
            args=[self._to_bytes(self_id), data, self._to_msec(data_ttl_sec)],
        )
        if isinstance(param, bytes):
            param = self._to_text(param, errors="replace")  # expecting a debug value
        if situation == self.situation_enum.success:
            return param  # watchers count
        if situation == self.situation_enum.token_mismatch:
            raise self.lock_lost_exc_cls(
                f"Lock found to be owned by another: {param!r}"
            )
        if situation == self.situation_enum.not_locked:
            raise self.lock_lost_exc_cls(
                (
                    f"Lock found to be expired (but saving the data anyway);"
                    f"watchers: {param!r}"
                )
            )
        raise self.support.result_error_exc_cls(
            f"Unexpected save script result: {res!r}"
        )


class FailScript(ScriptBase):
    script: ClassVar[str] = CL_FAIL_SCRIPT
    situation_enum: ClassVar[Type[SaveScriptResult]] = SaveScriptResult
    lock_lost_exc_cls: Type[Exception] = CacheLockLost

    async def __call__(
        self,
        lock_key: str,
        signal_key: str,
        self_id: str,
        ignore_errors: bool = False,
    ) -> Optional[int]:
        res, (situation, param) = await self._eval_script(
            keys=[self._to_bytes(lock_key), self._to_bytes(signal_key)],
            args=[self._to_bytes(self_id)],
        )
        if isinstance(param, bytes):
            param = self._to_text(param, errors="replace")  # expecting a debug value
        if situation == self.situation_enum.token_mismatch and param == "":
            if ignore_errors:
                return None
            raise self.lock_lost_exc_cls("Lock found to be expired")
        if situation == self.situation_enum.token_mismatch:
            if ignore_errors:
                return None
            raise self.lock_lost_exc_cls(
                f"Lock found to be owned by another: {param!r}"
            )
        if situation == self.situation_enum.success:
            return param  # watchers count
        raise self.support.result_error_exc_cls(
            f"Unexpected fail script result: {res!r}"
        )


class ForceSaveScript(ScriptBase):
    script: ClassVar[str] = CL_FORCE_SAVE_SCRIPT
    situation_enum: ClassVar[Optional[Type[IntEnum]]] = None

    async def __call__(
        self,
        signal_key: str,
        data_key: str,
        data: bytes,
        data_ttl_sec: float,
    ) -> int:
        res = await self._eval_script_base(
            keys=[self._to_bytes(signal_key), self._to_bytes(data_key)],
            args=[data, self._to_msec(data_ttl_sec)],
        )
        if not isinstance(res, int):
            raise self.support.result_error_exc_cls(
                f"Unexpected force-save script result: {res!r}"
            )
        return res

lib/redis-cache-lock/redis_cache_lock/scripts_support.py
datalens-tech/datalens-backend
from __future__ import annotations

import sys
from typing import (
    TYPE_CHECKING,
    AsyncContextManager,
    Awaitable,
    Callable,
    Optional,
    Tuple,
    TypeVar,
)

if sys.version_info >= (3, 8):
    from typing import Protocol
else:
    from typing_extensions import Protocol

if TYPE_CHECKING:
    from redis.asyncio import Redis


# `generate_func` types:
# `generate_func` must return a bytestring for saving to redis,
# but can also return a value to be returned directly,
# to skip an unnecessary deserialization (if there was no cache hit).
# The cache wrap can return either the full result, or only the bytestring part.
_GF_RET_TV = TypeVar("_GF_RET_TV")
TGenerateResult = Tuple[bytes, _GF_RET_TV]
TCacheResult = Tuple[bytes, Optional[_GF_RET_TV]]
TGenerateFunc = Callable[[], Awaitable[TGenerateResult]]


class TClientACM(Protocol):
    """
    Redis Client async context manager type.

    Should return a redis instance available for the duration of the ACM.

    If it is requested as exclusive, it should be exclusive for the duration of
    the async context manager, and thrown out if any error happens.
    """

    def __call__(
        self, *, master: bool, exclusive: bool
    ) -> AsyncContextManager["Redis"]:
        pass

lib/redis-cache-lock/redis_cache_lock/types.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import contextlib
import json
import logging
import os
import socket
import threading
import time
import uuid
from enum import IntEnum, unique
from typing import (
    Any,
    AsyncContextManager,
    AsyncGenerator,
    Awaitable,
    Callable,
    Dict,
    Iterable,
    List,
    Optional,
    Set,
    Tuple,
    Type,
    TypeVar,
    Union,
)

import attr

LOGGER = logging.getLogger(__name__)


@unique
class CMState(IntEnum):
    initialized = 1
    entered = 2
    exiting = 3
    exited = 4


_PE_RET_TV = TypeVar("_PE_RET_TV")


class PreExitable:
    """An AsyncContextManager wrapper that allows exiting before the `with` block is over"""

    def __init__(self, cm: AsyncContextManager[_PE_RET_TV]) -> None:
        self._cm = cm
        self._state = CMState.initialized
        self._exiting_task: Optional[asyncio.Future] = None

    async def __aenter__(self) -> _PE_RET_TV:
        assert self._state == CMState.initialized
        self._state = CMState.entered
        return await self._cm.__aenter__()

    async def _call_aexit(self, *exc_details: Any) -> Any:
        try:
            await self._cm.__aexit__(*exc_details)
        finally:
            self._state = CMState.exited
            self._exiting_task = None

    async def __aexit__(self, *exc_details: Any) -> Any:
        if self._state == CMState.exited:
            return None
        if self._state == CMState.exiting:
            # This makes it possible to close multiple times, with all calls
            # waiting on the finalization.
            assert self._exiting_task is not None
            return await asyncio.shield(self._exiting_task)
        assert self._state == CMState.entered
        self._state = CMState.exiting
        exiting_task = asyncio.ensure_future(
            asyncio.shield(self._call_aexit(*exc_details))
        )
        self._exiting_task = exiting_task
        return await exiting_task

    async def exit(self) -> Any:
        return await self.__aexit__(None, None, None)


@contextlib.asynccontextmanager
async def task_cm(coro: Awaitable) -> AsyncGenerator[asyncio.Task, None]:
    """Small helper to run an asyncio task for the duration of the context manager"""
    # While it doesn't actually do any `await`, it would be weird to expect non-async here.
    task = asyncio.create_task(coro)
    try:
        yield task
    finally:
        task.cancel()
        # Ensure the cancelling is synchronous:
        try:
            await task
        except asyncio.CancelledError:
            pass


@contextlib.asynccontextmanager
async def await_on_exit(coro: Awaitable) -> AsyncGenerator[Any, None]:
    try:
        yield coro
    finally:
        await coro


@attr.s(auto_attribs=True)
class CacheShareItem:
    key: str
    done: bool = attr.ib(default=False)
    result: Any = attr.ib(default=None)
    lock: asyncio.Lock = attr.ib(factory=asyncio.Lock)
    waiters: Set[Any] = attr.ib(factory=set)


def get_current_task_name() -> Optional[str]:
    current_task = asyncio.current_task()
    if current_task is not None:
        get_name = getattr(current_task, "get_name", None)  # python 3.7 doesn't
        if get_name is not None and callable(get_name):
            return get_name()
    return None


def new_self_id() -> str:
    pieces = [
        "h_" + socket.gethostname(),
        "p_" + str(os.getpid()),
    ]
    thread_name = threading.current_thread().name
    if thread_name and thread_name != "MainThread":
        pieces.append("t_" + thread_name)  # rare
    current_task_name = get_current_task_name()
    if current_task_name:
        pieces.append("a_" + current_task_name)
    pieces.append("r_" + str(uuid.uuid4()))
    return "_".join(pieces)


_CSS_RES_TV = TypeVar("_CSS_RES_TV")


@attr.s
class CacheShareSingleton:
    """
    A helper for in-memory synchronization of cached value generation.

    Intended to be used as a wrapper around `RedisCacheLock` calls, for a bit
    more performance.
    """

    item_cls: Type[CacheShareItem] = CacheShareItem

    cache: Dict[str, CacheShareItem] = attr.ib(factory=dict)
    track_waiters: bool = attr.ib(default=True)
    debug: bool = attr.ib(default=False)

    def _debug(self, msg: str, *args: Any) -> None:
        if self.debug:
            LOGGER.debug(msg, *args)

    async def generate_with_cache(
        self,
        key: str,
        generate: Callable[[], Awaitable[_CSS_RES_TV]],
    ) -> _CSS_RES_TV:
        cache_item = self.cache.get(key)
        if cache_item is None:
            self._debug("Initializing: key=%r", key)
            cache_item = self.item_cls(key=key)
            self.cache[key] = cache_item

        # Not a normal case, actually:
        # should imply non-empty `cache_item.waiters`
        if cache_item.done:
            self._debug("Found ready: key=%r, item=%r", key, cache_item)  # rare
            return cache_item.result  # rare

        this_request = (object(), get_current_task_name() or str(uuid.uuid4()))
        assert this_request not in cache_item.waiters

        try:
            # This can be removed in favor of python's refcounting (using weakref in `self.cache`),
            # for more performance and less debuggability.
            cache_item.waiters.add(this_request)

            self._debug("Locking %r...", key)
            async with cache_item.lock:
                if cache_item.done:
                    self._debug("Locked %r, found it done.", key)
                    return cache_item.result

                self._debug("Generating %r...", key)
                result = await generate()
                cache_item.result = result
                cache_item.done = True
                return result

        finally:
            cache_item.waiters.remove(this_request)
            if not cache_item.waiters:
                self._debug("Clearing %r.", key)
                # Make it garbage-collectable soon enough:
                self.cache.pop(key)


_GF_RET_TV = TypeVar("_GF_RET_TV")


def wrap_generate_func(
    func: Callable[[], Awaitable[_GF_RET_TV]],
    serialize: Callable[[Any], Union[bytes, str]] = json.dumps,
    default_encoding: str = "utf-8",
) -> Callable[[], Awaitable[Tuple[bytes, _GF_RET_TV]]]:
    """
    Given a function that returns some value, wrap it to also return the
    serialized version of the value for use in RedisCacheLock.
    """

    async def wrapped_generate_func() -> Tuple[bytes, _GF_RET_TV]:
        result = await func()
        result_b = serialize(result)
        if isinstance(result_b, str):
            result_b = result_b.encode(default_encoding)
        return result_b, result

    return wrapped_generate_func


# This makes `time.time_ns() - time.monotonic_ns()` but with slightly more precision.
_T1 = time.monotonic_ns()
_T2 = time.time_ns()
_T3 = time.monotonic_ns()
MONOTIME_NS_OFFSET = _T2 - _T1 // 2 - _T3 // 2
MONOTIME_OFFSET = MONOTIME_NS_OFFSET / 1e9


def monotime_ns() -> int:
    """
    Monotonic nanoseconds time that approximates unix nanosecond timestamp,
    for precision and convenience.
    """
    return time.monotonic_ns() + MONOTIME_NS_OFFSET


def monotime() -> float:
    """
    Monotonic time that approximates unix timestamp,
    for maximum convenience.
    """
    return time.monotonic() + MONOTIME_OFFSET


@attr.s(auto_attribs=True)
class HistoryHolder:
    """
    An in-memory holder that can be used as `debug_log` callable
    (optionally wrapping another callable)
    """

    max_size: int = 10_000
    func: Optional[Callable[[str, Dict[str, Any]], None]] = None

    history: List[Tuple[float, str, Dict[str, Any]]] = attr.ib(factory=list, repr=False)

    def __call__(self, msg: str, details: Dict[str, Any]) -> None:
        if len(self.history) >= self.max_size:
            self.history.pop(0)
        self.history.append((monotime(), msg, details))
        if self.func is not None:
            self.func(msg, details)  # pylint: disable=not-callable

    def __iter__(self) -> Iterable[Tuple[float, str, Dict[str, Any]]]:
        return iter(list(self.history))

lib/redis-cache-lock/redis_cache_lock/utils.py
datalens-tech/datalens-backend
# pylint: disable=redefined-outer-name
from __future__ import annotations

import asyncio
import itertools
import logging
import uuid
from typing import TYPE_CHECKING, AsyncGenerator, Callable

import pytest

from redis_cache_lock.main import RedisCacheLock
from redis_cache_lock.redis_utils import make_simple_cli_acm
from redis_cache_lock.utils import HistoryHolder, wrap_generate_func

if TYPE_CHECKING:
    from redis_cache_lock.types import TClientACM


@pytest.fixture
async def cli_acm() -> TClientACM:
    return make_simple_cli_acm("redis://localhost")


def lock_mgr(cli_acm, **kwargs) -> RedisCacheLock:
    def log_func(msg, details):
        logging.debug(
            "RedisCacheLock: %s, %s",
            msg, ", ".join(f"{key}={val!r}" for key, val in details.items()))

    logs = HistoryHolder(func=log_func)
    tasks = []

    def task_cb(task: asyncio.Task) -> None:
        tasks.append(task)

    full_kwargs = dict(
        key="",
        client_acm=cli_acm,
        resource_tag="tst",
        lock_ttl_sec=2.3,
        data_ttl_sec=5.6,
        debug_log=logs,
        bg_task_callback=task_cb,
    )
    full_kwargs.update(kwargs)
    mgr = RedisCacheLock(**full_kwargs)
    setattr(mgr, "logs_", logs)
    setattr(mgr, "tasks_", tasks)
    return mgr


@pytest.fixture(params=["fg", "bg"])
def backgroundness_mode(request) -> bool:
    return request.param == "bg"


@pytest.fixture
async def lock_mgr_gen(
    cli_acm,
    backgroundness_mode,
) -> AsyncGenerator[Callable[[], RedisCacheLock], None]:
    lock_mgrs = []

    def lock_mgr_gen_func(*args, **kwargs):
        kwargs = {
            "cli_acm": cli_acm,
            "enable_background_tasks": backgroundness_mode,
            **kwargs,
        }
        mgr = lock_mgr(*args, **kwargs)
        lock_mgrs.append(mgr)
        return mgr

    try:
        yield lock_mgr_gen_func
    finally:
        tasks = [task for mgr in lock_mgrs for task in getattr(mgr, "tasks_")]
        await asyncio.gather(*tasks, return_exceptions=True)


class CounterGenerate:
    def __init__(self, sleep_time: float = 0.01):
        self.sleep_time = sleep_time
        self.cnt = itertools.count(1)

    async def _raw_generate_func(self):
        await asyncio.sleep(self.sleep_time)
        return dict(value=next(self.cnt))

    @property
    def generate_func(self):
        return wrap_generate_func(self._raw_generate_func)


@pytest.mark.asyncio
async def test_minimal_lock(lock_mgr_gen):
    count = 5
    key = str(uuid.uuid4())

    lock_mgr = lock_mgr_gen(key=key)
    gen = CounterGenerate()

    for idx in range(count):
        result_b, result_raw = await lock_mgr.generate_with_lock(gen.generate_func)
        if idx == 0:
            assert result_raw == dict(value=1)
        else:
            assert result_raw is None
        assert result_b == b'{"value": 1}', idx

    logs: HistoryHolder = getattr(lock_mgr, "logs_")
    assert logs.history[-1][-1]["situation"] == lock_mgr.req_situation.cache_hit_slave

    result_b, result_raw = await lock_mgr.clone(key=key + "02").generate_with_lock(
        generate_func=gen.generate_func,
    )
    assert result_raw == dict(value=2)
    assert result_b == b'{"value": 2}'


@pytest.mark.asyncio
async def test_sync_lock(lock_mgr_gen):
    count = 7
    key = str(uuid.uuid4())

    mgrs = [lock_mgr_gen(key=key) for _ in range(count)]
    gen = CounterGenerate()

    results = await asyncio.gather(
        *[
            lock_mgr.generate_with_lock(generate_func=gen.generate_func)
            for lock_mgr in mgrs
        ],
    )

    assert results
    item_full = (b'{"value": 1}', {"value": 1})
    item_cached = (b'{"value": 1}', None)
    res_full = [item for item in results if item == item_full]
    res_cached = [item for item in results if item == item_cached]
    assert len(res_full) == 1, results
    assert len(res_cached) == count - 1, results


@pytest.mark.asyncio
async def test_sync_long_lock(lock_mgr_gen):
    count = 7
    key = str(uuid.uuid4())

    mgrs = [lock_mgr_gen(key=key) for _ in range(count)]
    gen = CounterGenerate(sleep_time=5)

    results = await asyncio.gather(
        *[
            lock_mgr.generate_with_lock(generate_func=gen.generate_func)
            for lock_mgr in mgrs
        ],
    )

    assert results
    item_full = (b'{"value": 1}', {"value": 1})
    item_cached = (b'{"value": 1}', None)
    res_full = [item for item in results if item == item_full]
    res_cached = [item for item in results if item == item_cached]
    assert len(res_full) == 1, results
    assert len(res_cached) == count - 1, results

lib/redis-cache-lock/tests/test_cache.py
datalens-tech/datalens-backend
from __future__ import annotations

import asyncio
import logging
import time

import pytest

from redis_cache_lock.utils import CacheShareSingleton

LOGGER = logging.getLogger(__name__)


async def some_generate(key: str):
    LOGGER.getChild("some_generate").debug("Generating %r...", key)
    await asyncio.sleep(0.5)
    result = time.time()
    LOGGER.getChild("some_generate").debug("Generated %r: %r", key, result)
    return result


async def some_generate_with_cache(caches, key):
    return key, await caches.generate_with_cache(
        key=key,
        generate=lambda key=key: some_generate(key),
    )


@pytest.mark.asyncio
async def test_cache_share():
    caches = CacheShareSingleton(debug=True)

    count = 20
    results = await asyncio.gather(
        *[
            some_generate_with_cache(caches=caches, key=f"kk{idx // 4}")
            for idx in range(count)
        ],
    )
    # While usually the `some_generate` is called `count` times, there is,
    # technically, no guarantee, as it depends on the ordering.
    assert results

lib/redis-cache-lock/tests/test_utils.py
datalens-tech/datalens-backend
import logging
import logging.config
import os
import socket
import sys

from .logs import LOGMUTATORS


# Warning: an on-import network request; if there's a danger of a broken DNS
# setup, use `socket.gethostname()`.
HOSTNAME = socket.getfqdn()


ENV_CONTEXT = {
    "app_name": None,  # To be filled in on convigure.
    "hostname": HOSTNAME,
    "app_instance": os.environ.get("QLOUD_DISCOVERY_INSTANCE"),
    "app_version": os.environ.get("QLOUD_DOCKER_IMAGE"),
    # global IP addresses, maybe?
}


def make_file_logger_syslog(name, params=None, **kwargs):
    addr = "/dev/log-ext"
    if not os.path.exists(addr):
        addr = "/dev/log"
        if not os.path.exists(addr):
            return {"class": "logging.NullHandler"}

    res = {
        "class": "statcommons.logs.TaggedSysLogHandler",
        "address": addr,
        "syslog_tag": "file__" + name,
    }
    if params:
        res.update(params)
    res.update(kwargs)
    return res


DEV_LOGGING_CONFIG = {
    "version": 1,
    "disable_existing_loggers": True,
    "handlers": {
        "stream": {
            "class": "logging.StreamHandler",
            "stream": sys.stderr,
            "level": "DEBUG",
        },
        "stream_info": {
            "class": "logging.StreamHandler",
            "stream": sys.stderr,
            "level": "INFO",
        },
    },
    "loggers": {},
    "root": {
        "level": "INFO",
        "handlers": ["stream"],
    },
}


BASE_LOGGING_CONFIG = {
    "version": 1,
    "disable_existing_loggers": False,
    "filters": {
        "events": {"()": "logging.Filter", "name": "events"},
        "fastlogs": {"()": "logging.Filter", "name": "fastlogs"},
    },
    "formatters": {
        "verbose": {"format": "[%(asctime)s] %(levelname)s: %(name)s: %(message)s"},
        # Put any preferred stdout/stderr json formatter in there:
        "json": {"format": ""},
        "jsonext": {"()": "statcommons.logs.JsonExtFormatter"},
        "jsonevent": {"()": "statcommons.logs.JsonEventFormatter"},
    },
    "handlers": {
        "stream": {
            "class": "logging.StreamHandler",
            "formatter": "json",
            "stream": sys.stdout,
            "level": "DEBUG",
        },
        "stream_info": {
            "class": "logging.StreamHandler",
            "formatter": "json",
            "stream": sys.stdout,
            "level": "INFO",
        },
        "stream_err": {
            "class": "logging.StreamHandler",
            "formatter": "json",
            "stream": sys.stdout,
            "level": "ERROR",
        },
        "debug_log": make_file_logger_syslog(
            "debug",
            formatter="jsonext",
            level="DEBUG",
        ),
        "fast_log": make_file_logger_syslog(
            "fast",
            formatter="jsonext",
            level="DEBUG",
            # Ways to use these:
            #  * Use logging.getLogger('fastlogs')
            #  * Replace the `filters` setting with your own.
            filters=["fastlogs"],
        ),
        "event_log": make_file_logger_syslog(
            "event",
            formatter="jsonevent",
            level="DEBUG",
            filters=["events"],
        ),
    },
    "loggers": {},
    "root": {
        "level": "DEBUG",
        "handlers": ["debug_log", "fast_log", "event_log"],
    },
}


def add_env_context_logmutator(record):
    for key, val in ENV_CONTEXT.items():
        setattr(record, key, val)


def configure_logging(app_name, cfg, use_env_context=True):
    current_name = ENV_CONTEXT.get("app_name")
    if current_name is not None and current_name != app_name:
        raise Exception("Attempting to configure logging again with a different app name", current_name, app_name)

    ENV_CONTEXT["app_name"] = app_name

    logging.config.dictConfig(cfg)
    LOGMUTATORS.apply(require=False)
    if use_env_context:
        LOGMUTATORS.add_mutator("env_context", add_env_context_logmutator)


def deconfigure_logging():
    """
    Primarily this is a hack for tests, particularly for running RQE under UWSGI after fork().
    As it is, it doesn't deconfigure the logging completely
    """
    ENV_CONTEXT.pop("app_name", None)


if __name__ == "__main__":
    # small self-test
    configure_logging(app_name="statcommonstst", cfg=BASE_LOGGING_CONFIG)
    configure_logging(app_name="statcommonstst", cfg=DEV_LOGGING_CONFIG)
    logging.getLogger().info("set up.")

lib/statcommons/statcommons/log_config.py
datalens-tech/datalens-backend
# coding: utf-8
"""
Modifications of `ylog`.
"""

from collections import OrderedDict
import datetime
import json
import logging
import logging.handlers
import traceback
from typing import (
    Any,
    Callable,
)


# Minor note:
# `context` in the formatter is a wrong approach.
# Use a 'logging filter' that adds whatever context you need to the record
# attibutes (essentially, an annotator).


# Common usable formats for file-based logs.
FILE_FORMAT = "%(asctime)s %(name)-15s %(levelname)-10s %(message)s"
DT_FORMAT = "%Y-%m-%d %H:%M:%S.%f"
# Remains here for compatibility
SYSLOG_FORMAT = "%(asctime)s %(name)-15s %(levelname)-10s %(message)s"


def smart_str(value):
    if isinstance(value, str):
        return value
    if isinstance(value, bytes):
        return value.decode("utf-8", "replace")
    return str(value)


def robust_smart_str(value, default="???"):
    try:
        return smart_str(value)
    except Exception:
        return default


def exception_str(value):
    """
    Formats Exception object into a string. Unlike default str():

    - can handle unicode strings in exception arguments
    - tries to format arguments as str(), not as repr()
    """
    try:
        return ", ".join([smart_str(item) for item in value])
    except (TypeError, AttributeError):  # happens for non-iterable values
        try:
            return smart_str(value)
        except UnicodeEncodeError:
            try:
                return repr(value)
            except Exception:
                return "<Unprintable value>"


def format_post(request, post_truncate_size=1024):
    """
    Formats request post data into a string. Depending on content type it's
    either a dict or raw post data.
    """
    if request.method == "POST" and (
        request.META.get("CONTENT_TYPE", "").split(";")[0]
        not in ["application/x-www-form-urlencoded", "multipart/form-data"]
    ):
        value = request.read()
        if len(value) > post_truncate_size:
            value = value[:post_truncate_size] + b"..."
        return value
    # else:
    return str(request.POST)


def format_frame_locals(frame, max_length=512, infix="...", prefix_tpl="      %s: ", suffix_tpl=" // %d bytes"):
    ret = ["    Locals:\n"]
    for key, value in frame.f_locals.items():
        result = repr(value)
        prefix = prefix_tpl % (key,)
        if len(result) > max_length:
            suffix = suffix_tpl % (len(result),)
            length = 80 - len(prefix) - len(infix) - len(suffix)
            result = result[: (length + 1) // 2] + infix + result[-(length // 2) :] + suffix
        ret.append("%s%s\n" % (prefix, result))
    return ret


def format_locals(tb):
    ret = []
    while tb is not None:
        ret.append(format_frame_locals(tb.tb_frame))
        tb = tb.tb_next
    return ret


def format_exception_with_locals(exc_type, exc_value, tb):
    try:
        exception = traceback.format_exception_only(exc_type, exc_value)
        locations = traceback.format_tb(tb)
        local_vars = format_locals(tb)
        formatted = ["Traceback (most recent call last):\n"]
        for loc, var in zip(locations, local_vars):
            formatted.append(loc)
            formatted.extend(var)
        formatted.extend(exception)
        return formatted
    except Exception:
        formatted = ["Couldn't print detailed traceback.\n"] + traceback.format_exception(exc_type, exc_value, tb)
        return formatted


def get_record_exc_repr(record, fmt=FILE_FORMAT, datefmt=DT_FORMAT, full=False, show_locals=False):
    """
    Formats exception log events into a short useful one-line message containing
    the exception class, module name, line number and the exception value. Exception
    info triple should be passed as a record parameter record.exc_info.

    When used with Django the parameter `full=True` adds extended data to the log message:
    request path, GET and POST representations and full traceback. The original log
    message is dropped (it's usually just the string "Internal server error").

    @param show_locals: if True, then every traceback frame would contain information
                        about local variables
    @param value_length_limit: the length, to which the variable value would be
                               stripped
    @param wrap_width: width for textwrap.wrap
    """

    # construct required fields in record.__dict__
    record.message = record.getMessage()
    if fmt.find("%(asctime)") >= 0:
        record.asctime = datetime.datetime.fromtimestamp(record.created).strftime(datefmt)

    lines = [fmt % record.__dict__]

    if full:
        req = getattr(record, "request", None)
        if req is not None:
            lines.append("Path: %s" % req.path)
            lines.append("GET: %s" % req.GET)
            lines.append("POST: %s" % format_post(req))
        exc_info = record.exc_info
        if exc_info is not None:
            if show_locals:
                exc_lines = format_exception_with_locals(*exc_info)
            else:
                exc_lines = traceback.format_exception(*exc_info)
            lines.append("".join(exc_lines))
    elif record.exc_info is not None and record.exc_info[-1] is not None:  # and not full:
        exception, value, tb = record.exc_info
        # we do not use traceback module because it doesn't provide
        # the way to determine module_name

        # find innermost call
        inner = tb
        while inner.tb_next:
            inner = inner.tb_next

        lineno = inner.tb_lineno
        f_globals = inner.tb_frame.f_globals
        module_name = f_globals.get("__name__", "<unknown>")

        record.message = "%-20s %s:%s %s" % (exception.__name__, module_name, lineno, exception_str(value))
        lines[0] = fmt % record.__dict__

    return "\n".join([smart_str(line) for line in lines])


class RecordDataFormatterMixin:
    unspecified = object()

    default_skipped_record_attrs = (
        # 'name',  # 'root',
        "msg",  # 'Example logged error: %r',
        "args",  # (ZeroDivisionError('division by zero'),),
        # 'levelname',  # 'ERROR',
        "levelno",  # 40,
        "pathname",  # 'tmpf.py',
        "filename",  # 'tmpf.py',
        "module",  # 'tmpf',
        # # covered by the func-based attr.
        # 'exc_info',  # (ZeroDivisionError, ZeroDivisionError('division by zero'), <traceback at 0x7f21bdc5b7c8>),
        "exc_text",  # None,
        "stack_info",  # None,
        # 'lineno',  # 11,
        # 'funcName',  # '<module>',
        "created",  # 1575299079.958725,
        "msecs",  # 958.7249755859375,
        "relativeCreated",  # 16066.834449768066,
        "thread",  # 139783057811200,
        "threadName",  # 'MainThread',
        "processName",  # 'MainProcess',
        "process",  # 409783}
        # func-based:
        # 'timestamp',
        # 'timestampns',
        # 'isotimestamp',
        # 'message',
        # 'pid',
        # 'exc_type',
        # 'exc',
        # compat, just-in-case, remove-later:
        "log_context",
        "dlog_context",
    )

    func_based_record_attrs = (
        # `self.get_record_...(record)`
        # Some of them are simple renames ('timestamp', 'pid')
        "timestamp",
        "timestampns",
        "isotimestamp",
        "message",
        "pid",
        "exc_type",
        "exc_info",
    )

    @staticmethod
    def get_record_timestamp(record):
        return record.created

    @staticmethod
    def get_record_timestampns(record, ns=int(1e9)):
        return int(record.created * ns)

    @staticmethod
    def get_record_isotimestamp(record, fmt=DT_FORMAT):
        dt = datetime.datetime.fromtimestamp(record.created)
        return dt.strftime(fmt)

    # Maybe: get_record_tzaware_isotiemstamp

    @staticmethod
    def get_record_message(record):
        return record.getMessage()

    @staticmethod
    def get_record_pid(record):
        return record.process

    @staticmethod
    def get_record_exc_type(record):
        if not record.exc_info:
            return None
        if not isinstance(record.exc_info, tuple):
            return None
        return record.exc_info[0].__name__

    @staticmethod
    def get_record_exc_info(record):
        exc_info = record.exc_info
        if exc_info is None:
            return None
        if not isinstance(exc_info, tuple):
            return exc_info

        # return record.formatException()

        _, exc, _ = exc_info
        exc_info_formatted = "".join(smart_str(val) for val in traceback.format_exception(*exc_info))

        # str(exc) at the beginning for some readability.
        return "{}\n{}".format(robust_smart_str(exc), exc_info_formatted)

    def __init__(self, skipped_record_attrs=unspecified):
        if skipped_record_attrs is self.unspecified:
            skipped_record_attrs = self.default_skipped_record_attrs
        self.skipped_record_attrs = frozenset(skipped_record_attrs)

    def record_to_data(self, record):
        """Make items dict from the record"""
        skiplist = self.skipped_record_attrs
        result = {}

        result.update({key: val for key, val in vars(record).items() if key not in skiplist})

        # Dubious: Prioritize the func-based over the base.
        for name in self.func_based_record_attrs:
            if name not in skiplist:
                func = getattr(self, "get_record_{}".format(name))
                result[name] = func(record)

        return result


class JsonExtFormatter(RecordDataFormatterMixin, logging.Formatter):
    def format(self, record):
        log_data = self.record_to_data(record)
        return self.dumps(log_data)

    @staticmethod
    def dumps(data):
        return json.dumps(data, default=robust_smart_str)


class JsonEventFormatter(JsonExtFormatter):
    """
    Use case:

        the_event = {'type': 1, 'value': 2}
        logger.info('An event', extra=dict(event=the_event))

    """

    def record_to_data(self, record):
        event = getattr(record, "event", None)
        if event is None or not isinstance(event, dict):
            return {"_unknown_record_event": event}
        return event


class TaggedSysLogHandlerBase(logging.handlers.SysLogHandler):
    """
    A version of SysLogHandler that adds a tag to the logged line
    so that it can be sorted by the syslog daemon into files.

    Generally equivalent to using a formatter but semantically more
    similar to FileHandler's `filename` parameter.
    """

    def __init__(self, *args, **kwargs):
        self.syslog_tag = kwargs.pop("syslog_tag")
        super(TaggedSysLogHandlerBase, self).__init__(*args, **kwargs)

    def format(self, *args, **kwargs):  # pylint: disable=arguments-differ
        res = super(TaggedSysLogHandlerBase, self).format(*args, **kwargs)
        return self.syslog_tag + " " + res


class TaggedSysLogHandler(TaggedSysLogHandlerBase):
    """An addition that sets the SO_SNDBUF to a large value to allow large log lines."""

    _sndbuf_size = 5 * 2**20

    def __init__(self, *args, **kwargs):
        self._sbdbuf_size = kwargs.pop("sbdbuf_size", self._sndbuf_size)
        super(TaggedSysLogHandler, self).__init__(*args, **kwargs)
        self.configure_socket(self.socket)

    def configure_socket(self, sock):
        import socket

        sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, self._sndbuf_size)


class LogRecordCaptureHandler(logging.Handler):
    """
    Similar to pytest's LogCaptureHandler,
    but only collects the record objects,
    without formatting them.
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.records = []

    def emit(self, record):
        self.records.append(record)

    def reset(self) -> None:
        self.records = []


class RecordMutators:
    """
    A mutator manager:
    wrap the logging factory
    and apply a the mutators on the newly created record.

    For performance, the log record mutation is used, instead of log record mappers.
    """

    def __init__(self):
        self.mutators = OrderedDict()
        self.old_factory = None

    def apply(self, require=True):
        """
        Wrap the current log record factory.
        Should only be done once.
        """
        if self.old_factory is not None:
            if require:
                raise Exception("This mutators-wrapper was already `apply`ed.")
            return
        self.old_factory = logging.getLogRecordFactory()
        logging.setLogRecordFactory(self)

    def add_mutator(self, name: str, mutator: Callable[[Any], None], force=False):
        """
        Register a mutator.
        `name` is required to ensure idempotency.
        """
        # ... would be nice to somehow validate the mutator callable here.
        if not force:
            current = self.mutators.get(name)
            if current is not None and current is not mutator:
                raise ValueError("Attempting to replace an existing mutator", name, current, mutator)
        self.mutators[name] = mutator

    def __call__(self, *args, **kwargs):
        """Log record factrory entry point"""
        if self.old_factory is None:
            raise Exception("This mutators-wrapper was not `apply`ed yet.")
        record = self.old_factory(*args, **kwargs)
        for mutator in self.mutators.values():
            mutator(record)
        return record


def make_attr_setter(key: str, value: Any, validate=False):
    """
    Make a mutator that does a single `setattr` with a constant.
    """

    def attr_setter(record):
        if validate:
            if hasattr(record, key):
                raise Exception("Record already has the attribute", key)
        setattr(record, key, value)

    return attr_setter


def test_logmutators():
    capture = LogRecordCaptureHandler()
    logger = logging.getLogger("test_logmutators")
    logger.propagate = False
    logger.addHandler(capture)
    logger.setLevel(logging.INFO)

    def set_1_1(record):
        record.test_key_1 = "test_val_1"

    def set_2_2(record):
        record.test_key_2 = "test_val_2"

    LOGMUTATORS.add_mutator("set_1", set_1_1)
    LOGMUTATORS.add_mutator("set_2", set_2_2)

    other_mutators = RecordMutators()

    def set_1_not_1(record):
        record.test_key_1 = "test_val_not_1"

    def set_3_3(record):
        record.test_key_3 = "test_val_3"

    other_mutators.add_mutator("set_1", set_1_not_1)
    other_mutators.add_mutator("set_3", set_3_3)

    logger.info("test_01")
    record = capture.records[-1]
    assert record.msg == "test_01"
    assert not hasattr(record, "test_key_1")
    assert not hasattr(record, "test_key_2")
    assert not hasattr(record, "test_key_3")

    LOGMUTATORS.apply()
    logger.info("test_02")
    record = capture.records[-1]
    assert record.msg == "test_02"
    assert record.test_key_1 == "test_val_1"
    assert record.test_key_2 == "test_val_2"
    assert not hasattr(record, "test_key_3")

    other_mutators.apply()
    logger.info("test_03")
    record = capture.records[-1]
    assert record.msg == "test_03"
    assert record.test_key_1 == "test_val_not_1"
    assert record.test_key_2 == "test_val_2"
    assert record.test_key_3 == "test_val_3"


# The recommended mutators-wrapper singleton.
LOGMUTATORS = RecordMutators()

lib/statcommons/statcommons/logs.py
datalens-tech/datalens-backend
"""
Common tools for unistat-format stats collection
"""

from __future__ import (
    absolute_import,
    division,
    print_function,
    unicode_literals,
)

from json import dumps as json_dumps
import os
import socket


__all__ = (
    "get_sys_memstatus",
    "get_context_data",
    "get_common_prefix",
    "maybe_float_size",
)


FQDN = None


def get_fqdn():
    global FQDN
    if FQDN is not None:
        return FQDN
    FQDN = socket.getfqdn()
    return FQDN


# http://man7.org/linux/man-pages/man5/proc.5.html
STATM_COLS = "size resident shared text lib data dt".split()


def get_sys_memstatus(pid):
    if not pid:
        return {}
    try:
        with open("/proc/{}/statm".format(pid), "r") as fobj:
            data_statm = fobj.read()
    except (OSError, IOError) as exc:
        return dict(_exc=repr(exc))
    data_m = data_statm.strip("\n").split(" ")
    return dict(zip(STATM_COLS, data_m))


CONTEXT_ENV_VARS = dict(
    project="QLOUD_PROJECT",
    application="QLOUD_APPLICATION",
    environment="QLOUD_ENVIRONMENT",
    # component='QLOUD_COMPONENT',
    geo="QLOUD_DATACENTER",
    instance="QLOUD_INSTANCE",
)


def get_context_data(require=True):
    env_vars = CONTEXT_ENV_VARS
    context_data = {key: (os.environ.get(src) or "").lower() for key, src in env_vars.items()}
    if require and not all(context_data.values()):
        raise Exception("Not all context environment variables are set: %r -> %r" % (env_vars, context_data))
    return context_data


def get_common_prefix(require=True):
    context_data = get_context_data(require=require)
    if not all(context_data.values()):
        return ""
    prefix = (
        "prj={project}.{application}.{environment};"
        "geo={geo};"
        # Currently does nothing, alas. Use filters like `tier={name}-*` instead:
        # 'component={component};'
        "tier={instance};"
    ).format(**context_data)
    return prefix


def maybe_float_size(value, divider):
    """
    RTFS.

    Convenience shorthand.
    """
    if value is None:
        return None
    return float(value) / divider


def results_to_response(results, indent="    "):
    if hasattr(results, "items"):
        results = results.items()
    yield "[\n"
    beginning = True
    for name, value in results:
        if not beginning:
            yield ",\n"
        beginning = False
        yield indent
        yield json_dumps([name, value])
    yield "\n"
    yield "]\n"


def dump_for_prometheus(results):
    for label, value in results:
        yield "{} {}\n".format(label, value)

lib/statcommons/statcommons/unistat/common.py
datalens-tech/datalens-backend
"""
flask as http daemon for unistat handling.
"""

from __future__ import (
    absolute_import,
    division,
    print_function,
    unicode_literals,
)

from collections import OrderedDict
import os

import flask

from .common import (
    dump_for_prometheus,
    results_to_response,
)
from .uwsgi import (
    uwsgi_prometheus,
    uwsgi_unistat,
)


def unistat_handler():
    results = OrderedDict()
    if os.environ.get("UWSGI_STATS") and os.environ.get("QLOUD_APPLICATION"):
        results.update(uwsgi_unistat())
    data = "".join(results_to_response(results))
    return flask.Response(data, content_type="application/json; charset=utf-8")


def register_unistat(app):
    return app.route("/unistat")(unistat_handler)


def _unistat_handler_hax():
    if flask.request.method == "GET" and flask.request.path.rstrip("/") == "/unistat":
        return unistat_handler()
    return None  # flask: fall through


def register_unistat_hax(app):
    # HAX to avoid bumping into inappropriate authorization added in 'before_request'.
    app.before_request(_unistat_handler_hax)


def metrics_handler():
    data = uwsgi_prometheus()
    body = "".join(dump_for_prometheus(data))
    return flask.make_response(body)


def register_metrics(app):
    return app.route("/metrics")(metrics_handler)

lib/statcommons/statcommons/unistat/flask.py
datalens-tech/datalens-backend
# coding: utf8
"""
uwsgi as source of signals for unistat.
"""

from __future__ import (
    absolute_import,
    division,
    print_function,
    unicode_literals,
)

from collections import OrderedDict
import json
import os
import socket
import time

from .common import (
    get_common_prefix,
    get_sys_memstatus,
    maybe_float_size,
)


__all__ = ("uwsgi_unistat",)


def group(lst, cls=dict):
    res = cls()
    for key, val in lst:
        try:
            group_list = res[key]
        except KeyError:
            res[key] = [val]
        else:
            group_list.append(val)
    return res


def _collect_uwsgi_info(path, collect_memstatus=True, encoding="utf-8"):
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.connect(path)
    try:
        data = ""
        while True:
            buf = sock.recv(64 * 2**10)
            if not buf:
                break
            data += buf.decode(encoding)
    finally:
        sock.close()

    data = json.loads(data)

    if collect_memstatus:
        for worker_info in data["workers"]:
            worker_info["sys_memstatus"] = get_sys_memstatus(worker_info.get("pid"))

    return data


def uwsgi_collect_sensors(data, add_sensor):
    now_ts = time.time()

    # relevant: https://github.com/wdtinc/uwsgi-cloudwatch/blob/master/uwsgi_cloudwatch/main.py#L79
    add_sensor("load", data["load"])
    add_sensor("signal_queue", data["signal_queue"])

    # locks: [{'user 0': 0}, {'signal': 0}, ...]
    add_sensor("locks_total", sum(value for lock in data["locks"] for value in lock.values()))

    # sockets:
    # [{'can_offload': 0, 'proto': 'uwsgi', 'queue': 0, 'shared': 0, 'max_queue': 0}, ...]
    add_sensor("sockets_queue_total", sum(item["queue"] for item in data["sockets"]))

    add_sensor("workers_total", len(data["workers"]))
    w_by_s = group((item["status"], item) for item in data["workers"])
    for key in ("idle", "busy", "cheap"):
        w_by_s.setdefault(key, [])
    for status, items in w_by_s.items():
        add_sensor(
            "workers_{}".format(status),
            len(items),
            suffixes=["_ammx", "_ahhh"],
        )
    add_sensor("workers_idle", len(w_by_s["idle"]), suffixes=["_ammn", "_annn"])

    avg_rts = [worker["avg_rt"] for worker in data["workers"]]
    avg_rts = [val for val in avg_rts if val]
    if avg_rts:
        add_sensor("max_avg_rt_sec", float(max(avg_rts)) / 1e6)

    # pylint: disable=invalid-name
    GiB = 2**30  # Gibi- (for gibibytes)
    page_size = 4096  # 4KiB, linux constant
    pages_per_gib = GiB / page_size  # 262144

    mem_datas = [  # all values in GiB
        dict(
            rss_gib=maybe_float_size(worker["rss"], GiB),
            vsz_gib=maybe_float_size(worker["vsz"], GiB),
            vmsize_gib=maybe_float_size(worker.get("sys_memstatus", {}).get("size"), pages_per_gib),
            vmrss_gib=maybe_float_size(worker.get("sys_memstatus", {}).get("resident"), pages_per_gib),
        )
        for worker in data["workers"]
    ]

    # `pandas` would be useful here; but probably not worth loading it up.
    # max_worker_rss max_worker_vsz all_worker_rss all_worker_vsz
    # max_worker_vmrss max_worker_vmsize all_worker_vmrss all_worker_vmsize
    for key in ("rss_gib", "vsz_gib", "vmsize_gib", "vmrss_gib"):
        series = [mem_data[key] for mem_data in mem_datas if mem_data.get(key) is not None]
        if not series:
            continue
        for func_name, func in (("max", max), ("all", sum)):
            add_sensor("{}_worker_{}".format(func_name, key), func(series))

    cores = [core for worker in data["workers"] for core in worker["cores"]]
    running_cores = [core for core in cores if core.get("req_info")]
    if running_cores:
        min_request_start = min(core["req_info"].get("request_start") for core in running_cores)
        add_sensor("longest_current_req_sec", now_ts - min_request_start)


def process_uwsgi_data_for_unistat(data, common_prefix=None, line_prefix=None):
    if common_prefix is None:
        common_prefix = get_common_prefix(require=True)

    results = OrderedDict()

    # # Qloud resource aggregates, for comparison:
    # memory: thhh tmmv tvvv txxx
    # cpu usage: hgram tmmv tvvv txxx
    # cpu limit: thhh tmmv txxx
    def add_sensor(label, value, prefix=line_prefix, suffixes=("_axxx", "_ahhh")):
        for suffix in suffixes:
            name = "{}{}{}{}".format(common_prefix, prefix or "", label.lower(), suffix)
            results[name] = value

    # TODO: tune the aggregation for usefulness.

    uwsgi_collect_sensors(data, add_sensor)

    return results


def process_uwsgi_data_for_prometheus(data, label_prefix=""):
    results = []

    # pylint: disable=unused-argument
    def add_sensor(label, value, suffixes=None):
        results.append((label_prefix + label, value))

    uwsgi_collect_sensors(data, add_sensor)

    return results


def _get_uwsgi_data(sock_path=None):
    if sock_path is None:
        sock_path = os.environ.get("UWSGI_STATS")
    if not sock_path:
        raise Exception("UWSGI_STATS socket path env variable is not set")

    data = _collect_uwsgi_info(sock_path)
    return data


def uwsgi_prometheus(sock_path=None, label_prefix=""):
    data = _get_uwsgi_data(sock_path=sock_path)
    return process_uwsgi_data_for_prometheus(data, label_prefix=label_prefix)


# pylint: disable=too-many-locals
def uwsgi_unistat(sock_path=None, common_prefix=None, line_prefix=None):
    data = _get_uwsgi_data(sock_path=sock_path)
    return process_uwsgi_data_for_unistat(data, common_prefix=common_prefix, line_prefix=line_prefix)

lib/statcommons/statcommons/unistat/uwsgi.py
datalens-tech/datalens-backend
import argparse
from collections import (
    defaultdict,
    deque,
)
import json
import os
from pathlib import Path
from typing import Iterator

import attr

from .pkg_ref import PkgRef


@attr.s
class Config:
    root_dir: Path = attr.ib()
    roots: list[str] = attr.ib()
    changed_paths: list[str] = attr.ib()


def collect_affected_packages(
    refs: list[PkgRef],
    paths: list[str],
) -> list[PkgRef]:
    # path relative to project root
    seen = set()
    pkg_by_path = {str(p.partial_parent_path): p for p in refs}

    # slow, but should work ...
    for p in paths:
        for pkg_path in pkg_by_path:
            if pkg_path in p:
                seen.add(pkg_path)
                continue

    return [pkg_by_path[k] for k in seen]


def gen_pkg_dirs(cfg: Config) -> Iterator[Path]:
    # todo: maybe relay of dl_repmanager to walk sub projects
    for root in cfg.roots:
        root_path = (Path(cfg.root_dir) / root).resolve()
        if not root_path.exists():
            continue
        for item in root_path.iterdir():
            if item.is_dir() and (item / "pyproject.toml").exists():
                yield item


def get_reverse_dependencies(direct_dependency: dict[str, list[str]]):
    reverse_ref = defaultdict(list)
    for pkg in direct_dependency:
        for dep in direct_dependency[pkg]:
            reverse_ref[dep].append(pkg)
    return reverse_ref


def get_leafs(dependencies):
    all_values = []
    for deps in dependencies.values():
        all_values.extend(deps)
    all_values = set(all_values)
    leafs = set(all_values) - set([k for k in dependencies.keys() if len(dependencies[k]) > 0])
    return leafs


def get_deep_deps(dependencies: dict[str, list[str]]) -> dict[str, set[str]]:
    """
    @param dependencies: dict with pkg names direct dependencies
    """
    aff_dict = get_reverse_dependencies(dependencies)
    leafs = get_leafs(dependencies)

    result = defaultdict(set)

    seen = set()
    ws = deque(leafs)
    while len(ws) > 0:
        node = ws.popleft()
        seen.add(node)
        for child in dependencies.get(node, []):
            result[node].add(child)
            result[node] |= result.get(child, set())
        for parent in aff_dict.get(node, []):
            if parent not in seen:
                ws.append(parent)

    return result


def get_deep_affection_map(dependencies: dict[str, list[str]]) -> dict[str, set[str]]:
    """
    Returns a mapping with pkg pointing to set of pkgs which is affected by changes in it
    """
    deep_deps = get_deep_deps(dependencies)

    affected = defaultdict(set)
    for pkg, dep_list in deep_deps.items():
        for dep in dep_list:
            affected[dep].add(pkg)
    return affected


def process(
    cfg: Config,
    get_all: bool = False,
) -> list[PkgRef]:
    affected = list(cfg.changed_paths)

    depends_on = defaultdict(list)
    pkg_by_ref = {}

    all_refs = []
    for item in gen_pkg_dirs(cfg):
        ref = PkgRef(root=cfg.root_dir, full_path=item)
        all_refs.append(ref)
        if ref.self_pkg_name is None:
            continue
        pkg_by_ref[ref.self_pkg_name] = ref
        for req in ref.extract_local_requirements(include_groups=["tests"]):
            depends_on[ref.self_pkg_name].append(req)

    if get_all:
        to_test = set(pkg_by_ref.keys())
    else:
        affection_map = get_deep_affection_map(depends_on)
        if len(affected) == 0:
            return all_refs

        direct_affected = collect_affected_packages(all_refs, affected)
        to_test = set()

        for pkg in direct_affected:
            to_test.add(pkg.self_pkg_name)
            to_test.update(affection_map.get(pkg.self_pkg_name, {}))

    return [pkg_by_ref[k] for k in to_test]


def main() -> None:
    if override := os.environ.get("TEST_TARGET_OVERRIDE", ""):
        if override != "__ALL__":
            overrides = override.strip().split(",")
            print(f"affected={json.dumps(overrides)}")
            return

    parser = argparse.ArgumentParser()
    parser.add_argument("--repo")
    parser.add_argument(
        "--root_pkgs",
        help="comma separated list of dirs to walk inside to find test targets",
    )
    parser.add_argument(
        "--fallback_pkg",
        help="relative path to pkg to always include in results",
    )
    parser.add_argument("--changes_file")
    args = parser.parse_args()

    with open(Path(args.changes_file)) as fh:
        changed_paths = fh.read().strip().split(" ")

    cfg = Config(
        root_dir=Path(args.repo),
        roots=(args.root_pkgs or "").split(","),
        changed_paths=changed_paths,
    )

    if override == "__ALL__":
        cfg.changed_paths = []
        to_check = process(cfg, get_all=True)
    else:
        to_check = process(cfg, False)

    to_check.append(
        PkgRef(
            root=cfg.root_dir,
            full_path=cfg.root_dir / args.fallback_pkg,
        )
    )
    result = set()

    for sub in to_check:
        if sub.skip_test:
            continue
        result.add(str(sub.partial_parent_path))
    print(f"affected={json.dumps(list(result))}")


if __name__ == "__main__":
    main()

terrarium/bi_ci/bi_ci/detect_affected_packages.py
datalens-tech/datalens-backend
from functools import partial
from pathlib import Path

import clize
import tomlkit


def do_we_need_compose(prj_root: Path, target: str) -> str:
    prj_root = prj_root.resolve()
    rel_path, section_name = target.split(":")
    mb_compose = prj_root / rel_path / "docker-compose.yml"

    ref = tomlkit.load(open(prj_root / rel_path / "pyproject.toml"))
    section = ref.get("datalens", {}).get("pytest", {}).get(section_name, {})

    if section.get("skip_compose") == "true":
        return "0"

    if name := section.get("compose_file_base"):
        mb_compose = prj_root / rel_path / f"{name}.yml"

    return "1" if mb_compose.exists() else "0"


cmd = partial(clize.run, do_we_need_compose)

terrarium/bi_ci/bi_ci/do_we_need_compose.py
datalens-tech/datalens-backend
from functools import partial
import json
from pathlib import Path
import subprocess
import sys

import clize
import tomlkit
from tomlkit.exceptions import NonExistentKey


def get_mypy_targets(pkg_dir: Path) -> list[str]:
    try:
        with open(pkg_dir / "pyproject.toml") as fh:
            meta = tomlkit.load(fh)
            return meta["datalens"]["meta"]["mypy"]["targets"]
    except NonExistentKey:
        pass

    # fallback to the same name defaults
    return [pkg_dir.name]


def main(root: Path, targets_file: Path) -> None:
    targets: list[str] = json.load(open(targets_file))
    failed_list: list[str] = []
    mypy_cache_dir = Path("/tmp/mypy_cache")
    mypy_cache_dir.mkdir(exist_ok=True)
    for target in targets:
        pkg_dir = root / target
        run_args = ["mypy", f"--cache-dir={mypy_cache_dir}"]
        targets = get_mypy_targets(pkg_dir)
        print(f"Cmd: {run_args}; cwd={pkg_dir}")
        if len(targets) > 0:
            run_args.extend(targets)
            run_exit_code = subprocess.run(" ".join(run_args), shell=True, cwd=str(pkg_dir)).returncode
            if run_exit_code != 0:
                failed_list.append(target)
        else:
            print(f"mypy config not found in {pkg_dir}/pyproject.toml, skipped")

    if failed_list:
        print("Mypy failed for targets:")
        print("\n".join(sorted(failed_list)))

    sys.exit(0 if len(failed_list) == 0 else 1)


cmd = partial(clize.run, main)

if __name__ == "__main__":
    cmd()

terrarium/bi_ci/bi_ci/execute_mypy_multi.py
datalens-tech/datalens-backend
from functools import partial
from pathlib import Path

import clize
import yaml


def remove_ports_from_docker_compose(src: Path, dst: Path) -> None:
    with open(src, "r") as file:
        data = yaml.safe_load(file)

    for service in data.get("services", {}).values():
        if "ports" in service:
            del service["ports"]

    with open(dst, "w") as file:
        yaml.safe_dump(data, file)


cmd = partial(clize.run, remove_ports_from_docker_compose)

terrarium/bi_ci/bi_ci/fix_ports_in_compose.py
datalens-tech/datalens-backend
from functools import partial
from pathlib import Path

import clize
import tomlkit


def get_compose_path(prj_root: Path, target: str) -> str:
    rel_path, section_name = target.split(":")

    ref: dict
    try:
        ref = tomlkit.load(open(prj_root / rel_path / "pyproject.toml"))
    except FileNotFoundError:
        ref = {}

    section = ref.get("datalens", {}).get("pytest", {}).get(section_name, {})

    base_path = prj_root / rel_path
    base_name = "docker-compose"
    if section and (base := section.get("compose_file_base")):
        base_name = base

    return str(Path(base_path / f"{base_name}.yml").resolve())


cmd = partial(clize.run, get_compose_path)

terrarium/bi_ci/bi_ci/get_compose_path.py
datalens-tech/datalens-backend
from functools import cached_property
from pathlib import Path

import attrs
import tomlkit


@attrs.define(slots=False)
class PkgRef:
    root: Path
    full_path: Path

    @cached_property
    def partial_parent_path(self):
        return self.full_path.relative_to(self.root)

    @cached_property
    def self_toml(self):
        try:
            return tomlkit.load(open(self.full_path / "pyproject.toml"))
        except Exception as err:
            print(err)
            return None

    def extract_local_requirements(self, include_groups: list[str] | None = None) -> set[str]:
        """
        Returns pkg names listed in dependencies of current pkg
        """
        spec = self.self_toml
        result = set()
        raw = {}
        if spec:
            raw = dict(spec["tool"]["poetry"]["dependencies"])
            for group in include_groups or []:
                if "group" not in spec["tool"]["poetry"]:
                    continue
                raw.update(spec["tool"]["poetry"]["group"].get(group, {}).get("dependencies", {}))

        for name, specifier in raw.items():
            if isinstance(specifier, dict) and "path" in specifier:
                result.add(name)

        return result

    @cached_property
    def self_pkg_name(self):
        try:
            return self.self_toml["tool"]["poetry"]["name"]
        except (KeyError, TypeError):
            return None

    @property
    def skip_test(self) -> bool:
        spec = self.self_toml
        try:
            return spec["datalens_ci"]["skip_test"]
        except (KeyError, TypeError):
            return False

terrarium/bi_ci/bi_ci/pkg_ref.py
datalens-tech/datalens-backend
from functools import partial
from pathlib import Path
import subprocess
import sys
import uuid

import attr
import clize
import tomlkit
from tomlkit import TOMLDocument


@attr.s
class Target:
    name: str = attr.ib()
    # where to run pytest
    root_dir: Path | None = attr.ib(default=None)
    # what should be pytest argument, relative to root_dir
    target_path: str | None = attr.ib(default=None)


def read_pkg(path: Path) -> TOMLDocument:
    return tomlkit.load(open(path))


def pkg_to_targets_ref(pkg_dir: Path, doc: TOMLDocument, sections: list[str] | None = None) -> list[Target]:
    result = list()
    for k, v in doc.get("datalens", {}).get("pytest", {}).items():
        if sections:
            if k not in sections:
                continue

        target_path = v.get("target_path", None)
        if target_path:
            for tp in target_path.split(" "):
                result.append(
                    Target(
                        name=k,
                        root_dir=pkg_dir / v.get("root_dir", None),
                        target_path=tp,
                    )
                )
        else:
            result.append(
                Target(
                    name=k,
                    root_dir=pkg_dir / v.get("root_dir", None),
                    target_path=v.get("target_path", None),
                )
            )

    return result


def run_pytest_one(t: Target) -> int:
    print(f"Running pytest for {t.name} in {str(t.root_dir)} for {t.target_path or '.'}")
    report_name = f"{str(t.root_dir)}-{t.target_path or '.'}-{uuid.uuid4().hex}.xml".replace("/", "_")
    run_args: list[str] = [
        "/venv/bin/python",
        "-m",
        "pytest",
        "-v",
        "--disable-warnings",
        f"--junitxml=/report/{report_name}",
    ]
    if t.target_path:
        run_args.append(t.target_path)

    subprocess_kwargs: dict[str, str | bool] = dict(
        universal_newlines=True,
    )
    if t.root_dir:
        subprocess_kwargs["cwd"] = str(t.root_dir.expanduser().resolve())
    print(f"{run_args=}")
    print(f"{subprocess_kwargs=}")

    process = subprocess.run(run_args, **subprocess_kwargs)  # type: ignore
    return process.returncode


def runner(prj_root: Path, test_target: str) -> None:
    """
    st_target: semicolon separated pair of package local path and target test section
    e.g: "lib/bi_core:unit_aio_api_commons"
    """
    # todo: consider multiple targets for one pkg_root, maybe comma separated

    pkg_root_str, target_section = test_target.split(":")
    pkg_root: Path = prj_root / pkg_root_str

    print("Datalens gh ci tests Runner welcomes you!")
    doc = read_pkg(pkg_root / "pyproject.toml")
    if target_section == "__default__":
        targets = []
    else:
        targets = pkg_to_targets_ref(pkg_root, doc, [target_section])

    if len(targets) == 0:
        print(f"Did not find specific instructions for running tests in {pkg_root} fallback to the simple pytest call")
        target = Target("__main__", pkg_root, ".")
        sys.exit(run_pytest_one(target))

    # Later we would probably filter this list depending on additional marks for
    #   for test targets and ENV vars defining which kind of tests should be executed
    # E.g. short vs long-running, unit vs external

    exit_code = 0
    for t in targets:
        run_exit_code = run_pytest_one(t)
        if run_exit_code not in [0, 5]:
            exit_code = 1

    sys.exit(exit_code)


runner_cli = partial(clize.run, runner)

terrarium/bi_ci/bi_ci/run_tests.py
datalens-tech/datalens-backend
from collections import defaultdict
from functools import partial
import json
from pathlib import Path

import clize
from tomlkit import parse


"""
example test spec:

[datalens.pytest.unit]
root_dir = "dl_core_tests/"
target_path = "unit"

[datalens.pytest.db_part_1]
root_dir = "dl_core_tests/db/"
target_path = "aio caches capabilities common"
labels = ["fat"]
"""


def format_output(name: str, sections: list[tuple[str, str]]) -> str:
    data = [f"{path}:{target}" for path, target in sections]
    return f"{name}={json.dumps(data)}"


def split_tests(
    requested_mode: str,
    root_dir: Path,
    test_targets_json_path: Path,
) -> None:
    raw = open(test_targets_json_path).read().strip().replace("'", '"')
    paths = json.loads(raw)

    # print(paths)
    split_result: dict[str, list[tuple]] = defaultdict(list)

    # ipdb.set_trace()
    for package_path in paths:
        path = root_dir / package_path.strip() / "pyproject.toml"
        # print(path)
        if not path.is_file():
            continue

        with open(path, "r") as file:
            toml_data = parse(file.read())

        pytest_targets = toml_data.get("datalens", {}).get("pytest", {})

        if pytest_targets:
            for section in pytest_targets.keys():
                spec = toml_data["datalens"]["pytest"][section]
                # supporting only a single label for now
                for category in spec.get("labels", []):
                    split_result[category].append((package_path, section))
                    break
                else:
                    split_result["base"].append((package_path, section))
        else:
            split_result["base"].append((package_path, "__default__"))

    print(format_output(f"split_{requested_mode}", split_result[requested_mode]))
    # note: I would expect following to be usable, but no idea how to make it with GHA :(
    # for mode, result in sorted(split_result.items()):
    #     print(format_output(f"split_{mode}", result), end=" ")


cmd = partial(clize.run, split_tests)

if __name__ == "__main__":
    cmd()

terrarium/bi_ci/bi_ci/split_pytest_tasks.py
datalens-tech/datalens-backend
from pathlib import Path

import pytest


@pytest.fixture(scope="session")
def resource_dir():
    conftest_dir = Path(__file__).resolve().parent
    resource_dir = conftest_dir / "resources"
    assert resource_dir.is_dir(), f"Resource directory not found at {resource_dir}"
    return resource_dir


@pytest.fixture(scope="session")
def sample_pkg_toml(resource_dir):
    return resource_dir / "sample.toml"


@pytest.fixture(scope="session")
def sample_compose_src(resource_dir):
    return resource_dir / "docker-compose.src.yml"


@pytest.fixture(scope="session")
def sample_compose_expected(resource_dir):
    return resource_dir / "docker-compose.expected.yml"

terrarium/bi_ci/bi_ci_tests/unit/conftest.py
datalens-tech/datalens-backend
from pathlib import Path
from shutil import copyfile

from bi_ci.fix_ports_in_compose import remove_ports_from_docker_compose


def test_remove_ports_from_docker_compose(tmpdir, sample_compose_src, sample_compose_expected):
    tmp_src = Path(tmpdir) / "docker-compose.yml"
    tmp_dst = Path(tmpdir) / "docker-compose-modified.yml"

    copyfile(sample_compose_src, tmp_src)

    remove_ports_from_docker_compose(tmp_src, tmp_dst)

    with open(tmp_dst, "r") as fh:
        output_data = fh.read()

    with open(sample_compose_expected, "r") as fh:
        expected_output_data = fh.read()

    assert output_data == expected_output_data

terrarium/bi_ci/bi_ci_tests/unit/test_fix_ports_in_compose.py
datalens-tech/datalens-backend
from bi_ci.get_compose_path import get_compose_path


def test_get_compose_path(resource_dir):
    prj_root = resource_dir.parent
    rel_path = resource_dir.name
    section = "db"

    result = get_compose_path(prj_root, f"{rel_path}:{section}")
    expected = str(resource_dir / "alt_docker-compose.yml")
    assert result == expected

    section = "snowflake"
    result = get_compose_path(prj_root, f"{rel_path}:{section}")
    expected = str(resource_dir / "docker-compose.yml")
    assert result == expected

    section = "--undefined--"
    result = get_compose_path(prj_root, f"{rel_path}:{section}")
    expected = str(resource_dir / "docker-compose.yml")
    assert result == expected

terrarium/bi_ci/bi_ci_tests/unit/test_get_compose_path.py
datalens-tech/datalens-backend
from pprint import pprint

import tomlkit

from bi_ci.run_tests import pkg_to_targets_ref


def test_read_test_data(sample_pkg_toml):
    doc = tomlkit.load(open(sample_pkg_toml))
    pprint(doc)


def test_pkg_to_targets_ref(sample_pkg_toml):
    doc = tomlkit.load(open(sample_pkg_toml))
    pprint(pkg_to_targets_ref(sample_pkg_toml.parent, doc))

terrarium/bi_ci/bi_ci_tests/unit/test_run_tests.py
datalens-tech/datalens-backend
import abc
import argparse
from typing import Sequence

import attr


@attr.s
class CliToolBase(abc.ABC):
    @classmethod
    @abc.abstractmethod
    def get_parser(cls) -> argparse.ArgumentParser:
        raise NotImplementedError

    @classmethod
    @abc.abstractmethod
    def run_parsed_args(cls, args: argparse.Namespace) -> None:
        raise NotImplementedError

    @classmethod
    def run(cls, argv: Sequence[str]) -> None:
        parser = cls.get_parser()
        cls.run_parsed_args(parser.parse_args(argv))

terrarium/dl_cli_tools/dl_cli_tools/cli_base.py
datalens-tech/datalens-backend
from __future__ import annotations

import logging


log = logging.getLogger()


def setup_basic_logging() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s",
        datefmt="%H:%M:%S",
    )
    log.debug("Logger initiated")

terrarium/dl_cli_tools/dl_cli_tools/logging.py
datalens-tech/datalens-backend
from contextlib import redirect_stdout
import io
from typing import (
    ClassVar,
    Type,
)

import attr

from dl_cli_tools.cli_base import CliToolBase


@attr.s(frozen=True)
class CliResult:
    stdout: str = attr.ib(kw_only=True)


class CliRunner:
    cli_cls: ClassVar[Type[CliToolBase]]

    def run_with_args(self, argv: list[str]) -> CliResult:
        with redirect_stdout(io.StringIO()) as out_stream:
            self.cli_cls.run(argv)

        assert isinstance(out_stream, io.StringIO)
        result = CliResult(stdout=out_stream.getvalue())
        return result

terrarium/dl_cli_tools/dl_cli_tools/testing/tool_runner.py
datalens-tech/datalens-backend
from pathlib import Path

from git.exc import InvalidGitRepositoryError
from git.repo.base import Repo as GitRepo


def discover_repo(repo_path: Path) -> GitRepo:
    git_repo: GitRepo
    while True:
        if repo_path.parent == repo_path or not repo_path.exists():
            raise RuntimeError("Failed to find git repository")

        try:
            git_repo = GitRepo(path=repo_path)
            break
        except InvalidGitRepositoryError:
            repo_path = repo_path.parent

    return git_repo

terrarium/dl_gitmanager/dl_gitmanager/discovery.py
datalens-tech/datalens-backend
from __future__ import annotations

from pathlib import Path
from typing import Generator

import attr
from git.diff import Diff
from git.repo.base import Repo as GitRepo


@attr.s
class GitManager:
    git_repo: GitRepo = attr.ib(kw_only=True)

    def get_root_path(self) -> Path:
        return Path(self.git_repo.working_tree_dir)

    def get_head_commit(self) -> str:
        return self.git_repo.head.commit.hexsha

    def get_commit(self, commit_specifier: str) -> str:
        return self.git_repo.commit(commit_specifier).hexsha

    def _iter_diffs(self, base: str, head: str, absolute: bool = False) -> Generator[tuple[Path, Diff], None, None]:
        base_commit = self.git_repo.commit(base)
        head_commit = self.git_repo.commit(head)

        base_path: Path
        if absolute:
            base_path = self.get_root_path()
        else:
            base_path = Path(".")

        # Iter own diffs
        diff_index = head_commit.diff(base_commit)
        for diff_item in diff_index:
            yield base_path, diff_item

        # Iter submodules and get their internal diffs
        for submodule in self.git_repo.submodules:
            submodule_name = submodule.name
            base_tree_item = base_commit.tree[submodule_name]
            head_tree_item = head_commit.tree[submodule_name]
            submodule_base = base_tree_item.hexsha
            submodule_head = head_tree_item.hexsha
            submodule_manager = self.get_submodule_manager(submodule_name=submodule_name)

            submodule_base_path: Path
            if absolute:
                submodule_base_path = submodule_manager.get_root_path()
            else:
                submodule_base_path = Path(submodule.path)

            # Iterate. Override the repo paths here with the submodule path
            for _, diff_item in submodule_manager._iter_diffs(base=submodule_base, head=submodule_head):
                yield submodule_base_path, diff_item

    def get_diff_paths(self, base: str, head: str, absolute: bool = False) -> list[str]:
        result: set[str] = set()
        for base_path, diff_item in self._iter_diffs(base=base, head=head, absolute=absolute):
            if diff_item.a_path:
                result.add(str(base_path / diff_item.a_path))
            if diff_item.b_path:
                result.add(str(base_path / diff_item.b_path))

        return sorted(result)

    def get_submodule_manager(self, submodule_name: str) -> GitManager:
        submodule_dict = {sm.name: sm for sm in self.git_repo.submodules}
        submodule = submodule_dict[submodule_name]
        return type(self)(git_repo=submodule.module())

terrarium/dl_gitmanager/dl_gitmanager/git_manager.py
datalens-tech/datalens-backend
from __future__ import annotations

import argparse
from pathlib import Path
import sys
from typing import Optional

import attr
from git.repo.base import Repo as GitRepo

from dl_cli_tools.cli_base import CliToolBase
from dl_cli_tools.logging import setup_basic_logging
from dl_gitmanager.discovery import discover_repo
from dl_gitmanager.git_manager import GitManager


@attr.s
class GitManagerTool(CliToolBase):
    git_manager: GitManager = attr.ib(kw_only=True)

    @classmethod
    def get_parser(cls) -> argparse.ArgumentParser:
        parser = argparse.ArgumentParser(prog="Git Repository Management CLI")
        parser.add_argument("--repo-path", type=Path, help="Main repo path", default=Path.cwd())

        # mix-in parsers
        base_parser = argparse.ArgumentParser(add_help=False)
        base_parser.add_argument("--base", help="Base commit specifier", default="HEAD~1")

        head_parser = argparse.ArgumentParser(add_help=False)
        head_parser.add_argument("--head", help="Head commit specifier", default="HEAD")

        base_head_parser = argparse.ArgumentParser(add_help=False, parents=[base_parser, head_parser])

        absolute_parser = argparse.ArgumentParser(add_help=False)
        absolute_parser.add_argument("--absolute", action="store_true", help="Use absolute paths")

        # commands
        subparsers = parser.add_subparsers(title="command", dest="command")

        subparsers.add_parser(
            "diff-paths", parents=[base_head_parser, absolute_parser], help="List file paths with changes"
        )

        return parser

    def diff_paths(self, base: str, head: Optional[str], absolute: bool) -> None:
        diff_name_list = self.git_manager.get_diff_paths(base=base, head=head, absolute=absolute)
        print("\n".join(diff_name_list))

    @classmethod
    def initialize(cls, git_manager: GitManager) -> GitManagerTool:
        tool = cls(git_manager=git_manager)
        return tool

    @classmethod
    def run_parsed_args(cls, args: argparse.Namespace) -> None:
        git_repo = discover_repo(repo_path=args.repo_path)
        git_manager = GitManager(git_repo=git_repo)
        tool = cls.initialize(git_manager=git_manager)

        match args.command:
            case "diff-paths":
                tool.diff_paths(base=args.base, head=args.head, absolute=args.absolute)
            case _:
                raise RuntimeError(f"Got unknown command: {args.command}")


def main() -> None:
    setup_basic_logging()
    GitManagerTool.run(sys.argv[1:])


if __name__ == "__main__":
    main()

terrarium/dl_gitmanager/dl_gitmanager/scripts/gitmanager_cli.py
datalens-tech/datalens-backend
class PackageMetaCliError(Exception):
    pass


class InconsistentStateError(PackageMetaCliError):
    pass


class MypyStubsOutOfSyncError(InconsistentStateError):
    pass

terrarium/dl_repmanager/dl_repmanager/exceptions.py
datalens-tech/datalens-backend
import abc
from contextlib import contextmanager
import io
import os
from pathlib import Path
import re
import shutil
import subprocess
from typing import (
    Callable,
    Collection,
    Generator,
    TextIO,
    Type,
    cast,
    final,
)

import attr


@attr.s(frozen=True)
class FilesystemEditor(abc.ABC):
    base_path: Path = attr.ib(kw_only=True)

    def _is_subpath(self, root: Path, child: Path) -> bool:
        root_parts = root.absolute().parts
        child_parts = child.absolute().parts
        return child_parts[: len(root_parts)] == root_parts

    def _validate_paths(self, *paths: Path) -> None:
        base_path_parts = self.base_path.absolute().parts
        for path in paths:
            path_parts = path.absolute().parts
            if path_parts[: len(base_path_parts)] != base_path_parts:
                raise RuntimeError(f"Access denied. Path is outside repository context: {path}")

    def _replace_file_content(self, file_path: Path, replace_callback: Callable[[str], str]) -> None:
        with self.open(file_path, "r+") as f:
            old_text = f.read()
            new_text = replace_callback(old_text)
            if new_text != old_text:
                f.seek(0)
                f.truncate(0)
                f.write(new_text)

    @final
    def replace_file_content(self, file_path: Path, replace_callback: Callable[[str], str]) -> None:
        self._validate_paths(file_path)
        self._replace_file_content(file_path=file_path, replace_callback=replace_callback)

    def _replace_text_in_file(self, file_path: Path, old_text: str, new_text: str) -> None:
        self.replace_file_content(
            file_path,
            replace_callback=lambda text: text.replace(old_text, new_text),
        )

    @final
    def replace_text_in_file(self, file_path: Path, old_text: str, new_text: str) -> None:
        self._validate_paths(file_path)
        self._replace_text_in_file(file_path=file_path, old_text=old_text, new_text=new_text)

    def _replace_regex_in_file(
        self,
        file_path: Path,
        regex: re.Pattern,
        repl: Callable[[re.Match], str],
    ) -> None:
        self.replace_file_content(
            file_path,
            replace_callback=lambda text: regex.sub(repl, text),
        )

    @final
    def replace_regex_in_file(self, file_path: Path, regex: re.Pattern, repl: Callable[[re.Match], str]) -> None:
        self._validate_paths(file_path)
        self._replace_regex_in_file(file_path=file_path, regex=regex, repl=repl)

    def _iter_files(
        self,
        path: Path,
        exclude_masks: Collection[re.Pattern] = (),
    ) -> Generator[Path, None, None]:
        for file_path in path.rglob("*"):
            if file_path.is_file():
                matches_exclude_mask = False
                for mask in exclude_masks:
                    if mask.match(str(file_path)):
                        matches_exclude_mask = True
                if matches_exclude_mask:
                    continue

                yield file_path

    def _replace_text_in_dir(
        self,
        old_text: str,
        new_text: str,
        path: Path,
        exclude_masks: Collection[re.Pattern] = (),
    ) -> None:
        for file_path in self._iter_files(path=path, exclude_masks=exclude_masks):
            self.replace_text_in_file(file_path, old_text=old_text, new_text=new_text)

    @final
    def replace_text_in_dir(
        self,
        old_text: str,
        new_text: str,
        path: Path,
        exclude_masks: Collection[re.Pattern] = (),
    ) -> None:
        self._validate_paths(path)
        self._replace_text_in_dir(old_text=old_text, new_text=new_text, path=path, exclude_masks=exclude_masks)

    def _replace_regex_in_dir(
        self,
        path: Path,
        regex: re.Pattern,
        repl: Callable[[re.Match], str],
        exclude_masks: Collection[re.Pattern] = (),
    ) -> None:
        for file_path in self._iter_files(path=path, exclude_masks=exclude_masks):
            self.replace_regex_in_file(file_path, regex=regex, repl=repl)

    @final
    def replace_regex_in_dir(
        self,
        path: Path,
        regex: re.Pattern,
        repl: Callable[[re.Match], str],
        exclude_masks: Collection[re.Pattern] = (),
    ) -> None:
        self._validate_paths(path)
        self._replace_regex_in_dir(path=path, regex=regex, repl=repl, exclude_masks=exclude_masks)

    @abc.abstractmethod
    def _copy_path(self, src_path: Path, dst_path: Path) -> None:
        """Make a copy of `src_dir` named `dst_dir`."""
        raise NotImplementedError

    @final
    def copy_path(self, src_path: Path, dst_path: Path) -> None:
        self._validate_paths(src_path, dst_path)
        self._copy_path(src_path=src_path, dst_path=dst_path)

    @abc.abstractmethod
    def _move_path(self, old_path: Path, new_path: Path) -> None:
        raise NotImplementedError

    @final
    def move_path(self, old_path: Path, new_path: Path) -> None:
        self._validate_paths(old_path, new_path)
        self._move_path(old_path=old_path, new_path=new_path)

    @abc.abstractmethod
    def _remove_path(self, path: Path) -> None:
        raise NotImplementedError

    @final
    def remove_path(self, path: Path) -> None:
        self._validate_paths(path)
        self._remove_path(path=path)

    @abc.abstractmethod
    def _path_exists(self, path: Path) -> bool:
        raise NotImplementedError

    @final
    def path_exists(self, path: Path) -> bool:
        self._validate_paths(path)
        return self._path_exists(path=path)

    @contextmanager
    def _open(self, path: Path, mode: str) -> Generator[TextIO, None, None]:
        with open(path, mode=mode) as file_obj:
            yield cast(TextIO, file_obj)

    @final
    @contextmanager
    def open(self, path: Path, mode: str) -> Generator[TextIO, None, None]:
        # TODO: Make all toml editors open files via this method to enforce path restrictions
        assert mode in ("r", "r+", "w")
        self._validate_paths(path)
        with self._open(path, mode=mode) as file_obj:
            yield file_obj


@attr.s(frozen=True)
class DefaultFilesystemEditor(FilesystemEditor):
    def _copy_path(self, src_path: Path, dst_path: Path) -> None:
        assert src_path.exists(), f"Source path '{src_path}' doesn't exist"
        assert not dst_path.exists(), "Destination dir already exists"
        if src_path.is_file():
            shutil.copy(src_path, dst_path)
        else:
            shutil.copytree(src_path, dst_path)

    def _move_path(self, old_path: Path, new_path: Path) -> None:
        assert old_path.exists()
        if old_path.is_dir():
            # module is a package
            print(f"Moving directory {old_path} to {new_path}")
            new_path.mkdir(exist_ok=True)

            for child in old_path.iterdir():
                new_child_path = new_path / child.name
                assert not new_child_path.exists(), f"Path {new_child_path} exists"
                shutil.move(child, new_path)

            shutil.rmtree(old_path)

        else:
            # module is a file
            assert old_path.is_file()
            assert not new_path.exists()

            print(f"Moving module {old_path} to {new_path}")
            new_path.parent.mkdir(exist_ok=True)

            shutil.move(old_path, new_path)

    def _remove_path(self, path: Path) -> None:
        if path.is_file():
            os.remove(path)
        else:
            shutil.rmtree(path)

    def _path_exists(self, path: Path) -> bool:
        return path.exists()


@attr.s(frozen=True)
class GitFilesystemEditor(DefaultFilesystemEditor):
    """An FS editor that buses git to move files and directories"""

    def _move_path(self, old_path: Path, new_path: Path) -> None:
        cwd = Path.cwd()
        rel_old_path = Path(os.path.relpath(old_path, cwd))
        rel_new_path = Path(os.path.relpath(new_path, cwd))
        subprocess.run(f'git add "{rel_old_path}" && git mv "{rel_old_path}" "{rel_new_path}"', shell=True)

    def _remove_path(self, path: Path) -> None:
        result = subprocess.run(f'git rm "{path}"', shell=True)
        if result.returncode != 0:
            super()._remove_path(path)


@attr.s(frozen=True)
class VirtualFilesystemEditor(FilesystemEditor):
    _moved_paths: dict[Path, Path] = attr.ib(init=False, factory=dict)  # {<new_path>: <old_path>}
    _copied_paths: dict[Path, Path] = attr.ib(init=False, factory=dict)  # {<new_path>: <original_path>}
    _removed_paths: set[Path] = attr.ib(init=False, factory=set)
    _edited_files: dict[Path, str] = attr.ib(init=False, factory=dict)  # {<current_path>: <content>}

    def ensure_path_exists(self, path: Path) -> None:
        if not path.exists():
            raise RuntimeError(f"Path does not exist: {path}")

    def ensure_path_doesnt_exist(self, path: Path) -> None:
        if path.exists():
            raise RuntimeError(f"Path already exists: {path}")

    def _recursive_children(self, path: Path) -> list[Path]:
        children: set[Path] = set()
        if path.is_file():
            return [path]
        for real_child in path.rglob("*/"):
            if real_child in self._moved_paths.values():
                # It has been moved
                continue
            if real_child in self._removed_paths:
                # It has been deleted
                continue

            children.add(real_child)

        for virtual_child in set(self._moved_paths) | set(self._copied_paths):
            if self._is_subpath(root=path, child=virtual_child):
                continue
            children.add(virtual_child)

        return sorted(children)

    def _copy_single_path(self, src_path: Path, dst_path: Path) -> None:
        self.ensure_path_exists(src_path)
        self.ensure_path_doesnt_exist(dst_path)

        if src_path in self._moved_paths:
            self._copied_paths[dst_path] = self._moved_paths[src_path]
        elif src_path in self._copied_paths:
            self._copied_paths[dst_path] = self._copied_paths[src_path]

        if src_path in self._edited_files:
            self._edited_files[dst_path] = self._edited_files[src_path]

    def _copy_path(self, src_path: Path, dst_path: Path) -> None:
        self._copy_single_path(src_path=src_path, dst_path=dst_path)
        for src_child_path in self._recursive_children(src_path):
            dst_child_path = dst_path / src_child_path.relative_to(src_path)
            self._copy_single_path(src_path=src_child_path, dst_path=dst_child_path)

    def _move_single_path(self, src_path: Path, dst_path: Path) -> None:
        self.ensure_path_exists(src_path)
        self.ensure_path_doesnt_exist(dst_path)

        if src_path in self._moved_paths:
            original_path = self._moved_paths[src_path]
            del self._moved_paths[src_path]
            self._moved_paths[dst_path] = original_path
        elif src_path in self._copied_paths:
            original_path = self._copied_paths[src_path]
            del self._copied_paths[src_path]
            self._copied_paths[dst_path] = original_path

        if src_path in self._edited_files:
            self._edited_files[dst_path] = self._edited_files.pop(src_path)

    def _move_path(self, old_path: Path, new_path: Path) -> None:
        self._move_single_path(src_path=old_path, dst_path=new_path)
        for src_child_path in self._recursive_children(old_path):
            dst_child_path = new_path / src_child_path.relative_to(old_path)
            self._move_single_path(src_path=src_child_path, dst_path=dst_child_path)

    def _remove_path(self, path: Path) -> None:
        if path in self._moved_paths:
            original_path = self._moved_paths[path]
            del self._moved_paths[path]
            self._removed_paths.add(original_path)
        elif path in self._copied_paths:
            del self._copied_paths[path]
        else:
            self._removed_paths.add(path)

        if path in self._edited_files:
            del self._edited_files[path]

    @contextmanager
    def _open(self, path: Path, mode: str) -> Generator[TextIO, None, None]:
        original_path = path
        if path in self._moved_paths:
            original_path = self._moved_paths[path]
        elif path in self._copied_paths:
            original_path = self._copied_paths[path]

        if path in self._edited_files:
            original_text = self._edited_files[path]
        elif mode == "w" and not original_path.exists():
            original_text = ""
        else:
            with super()._open(original_path, mode="r") as file_obj:
                original_text = file_obj.read()

        str_io: io.StringIO
        if mode == "w":
            str_io = io.StringIO()
        else:
            str_io = io.StringIO(original_text)
        yield str_io

        if mode in ("r+", "w"):
            new_text = str_io.getvalue()
            if new_text != original_text:
                self._edited_files[path] = new_text

    def _path_exists(self, path: Path) -> bool:
        return (
            path in self._moved_paths or path in self._copied_paths or path in self._edited_files or path.exists()
        ) and path not in self._removed_paths


_FS_EDITOR_CLASSES: dict[str, Type[FilesystemEditor]] = {
    "default": DefaultFilesystemEditor,
    "git": GitFilesystemEditor,
    "virtual": VirtualFilesystemEditor,
}


def get_fs_editor(fs_editor_type: str, base_path: Path) -> FilesystemEditor:
    fs_editor_cls = _FS_EDITOR_CLASSES[fs_editor_type]
    fs_editor = fs_editor_cls(base_path=base_path)
    return fs_editor

terrarium/dl_repmanager/dl_repmanager/fs_editor.py
datalens-tech/datalens-backend
from pathlib import Path
import subprocess
from typing import Optional

import attr


@attr.s()
class GitManager:
    _path: Path = attr.ib()
    _suppress_stdout_on_run = attr.ib(default=True)

    def _run_git(
        self,
        args: list[str],
        cwd: Optional[Path] = None,
        extract_stdout: bool = False,
    ) -> subprocess.CompletedProcess:
        actual_cwd: Path = cwd or self._path

        stdout_target: Optional[int]
        if extract_stdout:
            stdout_target = subprocess.PIPE
        elif self._suppress_stdout_on_run:
            stdout_target = subprocess.DEVNULL
        else:
            stdout_target = None

        return subprocess.run(
            ["git", *args],
            cwd=actual_cwd,
            stdout=stdout_target,
            check=True,
        )

    def _get_git_output(
        self,
        args: list[str],
        *,
        auto_strip: bool,
        cwd: Optional[Path] = None,
    ) -> str:
        proc = self._run_git(args=args, cwd=cwd, extract_stdout=True)
        stdout = proc.stdout.decode("ascii")
        if auto_strip:
            return stdout.strip("\n")
        return stdout

    def get_origin_url(self) -> str:
        return self._get_git_output(["config", "--get", "remote.origin.url"], auto_strip=True)

    def fetch(self) -> None:
        self._run_git(["fetch"])

    def copy_repo(self, destination: Path, fetch_if_exists: bool = False) -> "GitManager":
        if not destination.parent.exists():
            raise ValueError(f"Destination parent folder does not exists: {destination.parent}")

        current_origin_url = self.get_origin_url()

        dest_repo_manager = attr.evolve(self, path=destination)

        if destination.exists():
            dest_repo_origin_url = dest_repo_manager.get_origin_url()
            if dest_repo_origin_url != current_origin_url:
                raise ValueError("Copy target repo exists and it's origin does not match current.")

            if fetch_if_exists:
                dest_repo_manager.fetch()

        else:
            self._run_git(
                ["clone", current_origin_url, destination.name],
                cwd=destination.parent,
            )

        dest_repo_manager = attr.evolve(self, path=destination)

        return dest_repo_manager

    def checkout(self, rev: str) -> None:
        rev_hash = self._get_git_output(["rev-parse", rev], auto_strip=True)
        self._run_git(["checkout", rev_hash])

terrarium/dl_repmanager/dl_repmanager/git_manager.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
import contextlib
import os
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Generator,
    Type,
)

import attr
import tomlkit

from dl_repmanager.fs_editor import FilesystemEditor
from dl_repmanager.package_meta_reader import PackageMetaIOFactory
from dl_repmanager.primitives import (
    LocalReqPackageSpec,
    PackageInfo,
    RequirementList,
)
from dl_repmanager.toml_tools import (
    TOMLIOFactory,
    TOMLWriter,
)


if TYPE_CHECKING:
    from dl_repmanager.package_index import PackageIndex
    from dl_repmanager.repository_env import RepoEnvironment


@attr.s
class RepositoryManagementPlugin(abc.ABC):
    repository_env: RepoEnvironment = attr.ib(kw_only=True)
    package_index: PackageIndex = attr.ib(kw_only=True)
    base_path: Path = attr.ib(kw_only=True)
    config_data: dict = attr.ib(kw_only=True)
    fs_editor: FilesystemEditor = attr.ib(init=False)

    @fs_editor.default
    def _make_fs_editor(self) -> FilesystemEditor:
        return self.repository_env.get_fs_editor()

    @abc.abstractmethod
    def register_package(self, package_info: PackageInfo) -> None:
        raise NotImplementedError

    @abc.abstractmethod
    def unregister_package(self, package_info: PackageInfo) -> None:
        raise NotImplementedError

    def re_register_package(self, old_package_info: PackageInfo, new_package_info: PackageInfo) -> None:
        self.unregister_package(old_package_info)
        self.register_package(new_package_info)


@attr.s
class CommonToolingRepositoryManagementPlugin(RepositoryManagementPlugin):
    _package_list_rel_path: Path = attr.ib(init=False)

    @_package_list_rel_path.default
    def _make_package_list_rel_path(self) -> Path:
        return Path(self.config_data["package_list_path"])

    def register_package(self, package_info: PackageInfo) -> None:
        def transform_package_list(old_text: str) -> str:
            pkg_list = old_text.strip().split()
            pkg_rel_path = package_info.get_relative_path(self.base_path)
            pkg_list.append(str(pkg_rel_path))
            pkg_list.sort()
            return "\n".join(pkg_list) + "\n"

        pkg_list_path = self.base_path / self._package_list_rel_path
        self.fs_editor.replace_file_content(pkg_list_path, replace_callback=transform_package_list)

    def unregister_package(self, package_info: PackageInfo) -> None:
        def transform_package_list(old_text: str) -> str:
            pkg_list = old_text.strip().split()
            pkg_rel_path = package_info.get_relative_path(self.base_path)
            pkg_list.remove(str(pkg_rel_path))
            pkg_list.sort()
            return "\n".join(pkg_list) + "\n"

        pkg_list_path = self.base_path / self._package_list_rel_path
        self.fs_editor.replace_file_content(pkg_list_path, replace_callback=transform_package_list)


@attr.s
class MainTomlRepositoryManagementPlugin(RepositoryManagementPlugin):
    _metapackages_by_package_type: dict[str, list[str]] = attr.ib(init=False)

    @_metapackages_by_package_type.default
    def _make_metapackages_by_package_type(self) -> dict[str, list[str]]:
        result: dict[str, list[str]] = {}
        for metapkg_config in self.config_data["metapackages"]:
            for package_type in metapkg_config["package_types"]:
                if package_type not in result:
                    result[package_type] = []
                result[package_type].append(metapkg_config["name"])
        return result

    def _get_metapackage_paths(self, package_type: str) -> list[Path]:
        metapackage_names = self._metapackages_by_package_type.get(package_type, ())
        return [
            self.repository_env.get_metapackage_spec(metapackage_name).toml_path
            for metapackage_name in metapackage_names
        ]

    def _get_path_for_toml(self, metapackage_path: Path, package_info: PackageInfo) -> Path:
        toml_abs_dir = (self.base_path / metapackage_path).parent
        return package_info.get_relative_path(toml_abs_dir)

    def _register_main(self, metapackage_path: Path, package_info: PackageInfo) -> None:
        package_path_for_toml = self._get_path_for_toml(metapackage_path=metapackage_path, package_info=package_info)
        package_dep_table = tomlkit.inline_table()
        package_dep_table.add("path", str(package_path_for_toml))
        with self._metapackage_toml_writer(metapackage_path=metapackage_path) as toml_writer:
            section = toml_writer.get_editable_section("tool.poetry.group.ci.dependencies")
            section[package_info.package_reg_name] = package_dep_table

    def _unregister_main(self, metapackage_path: Path, package_info: PackageInfo):
        with self._metapackage_toml_writer(metapackage_path=metapackage_path) as toml_writer:
            with toml_writer.suppress_non_existent_key():
                toml_writer.get_editable_section("tool.poetry.dependencies").remove(package_info.package_reg_name)
            with toml_writer.suppress_non_existent_key():
                toml_writer.get_editable_section("tool.poetry.group.dev.dependencies").remove(
                    package_info.package_reg_name
                )
            with toml_writer.suppress_non_existent_key():
                toml_writer.get_editable_section("tool.poetry.group.ci.dependencies").remove(
                    package_info.package_reg_name
                )

    def _register_app(self, metapackage_path: Path, package_info: PackageInfo) -> None:
        package_path_for_toml = self._get_path_for_toml(metapackage_path=metapackage_path, package_info=package_info)
        package_base_name = package_info.abs_path.name
        package_dep_table = tomlkit.inline_table()
        package_dep_table.add("path", str(package_path_for_toml))
        with self._metapackage_toml_writer(metapackage_path=metapackage_path) as toml_writer:
            section = toml_writer.add_section(f"tool.poetry.group.app_{package_base_name}.dependencies")
            section.add(package_info.package_reg_name, package_dep_table)
            section.add(tomlkit.nl())

    def _unregister_app(self, metapackage_path: Path, package_info: PackageInfo) -> None:
        package_base_name = package_info.abs_path.name
        with self._metapackage_toml_writer(metapackage_path=metapackage_path) as toml_writer:
            toml_writer.delete_section(f"tool.poetry.group.app_{package_base_name}.dependencies")

    @contextlib.contextmanager
    def _metapackage_toml_writer(self, metapackage_path: Path) -> Generator[TOMLWriter, None, None]:
        toml_path = self.base_path / metapackage_path
        toml_io_factory = TOMLIOFactory(fs_editor=self.fs_editor)
        with toml_io_factory.toml_writer(toml_path) as toml_writer:
            yield toml_writer

    def register_package(self, package_info: PackageInfo) -> None:
        for metapackage_path in self._get_metapackage_paths(package_type=package_info.package_type):
            if "main_dependency_group" in self.repository_env.get_tags(package_info.package_type):
                self._register_main(metapackage_path=metapackage_path, package_info=package_info)
            if "own_dependency_group" in self.repository_env.get_tags(package_info.package_type):
                self._register_app(metapackage_path=metapackage_path, package_info=package_info)

    def unregister_package(self, package_info: PackageInfo) -> None:
        for metapackage_path in self._get_metapackage_paths(package_type=package_info.package_type):
            if "main_dependency_group" in self.repository_env.get_tags(package_info.package_type):
                self._unregister_main(metapackage_path=metapackage_path, package_info=package_info)
            if "own_dependency_group" in self.repository_env.get_tags(package_info.package_type):
                self._unregister_app(metapackage_path=metapackage_path, package_info=package_info)


@attr.s
class DependencyReregistrationRepositoryManagementPlugin(RepositoryManagementPlugin):
    """Updates requirements in other packages dependent on this one"""

    def _is_package_dependent_on(
        self,
        section_name: str,
        dependent_package_info: PackageInfo,
        base_package_info: PackageInfo,
    ) -> bool:
        req_specs = dependent_package_info.requirement_lists.get(section_name, RequirementList()).req_specs
        for req_spec in req_specs:
            if req_spec.package_name == base_package_info.package_reg_name:
                # It really is a dependent package
                assert isinstance(req_spec, LocalReqPackageSpec)
                return True

        return False

    def register_package(self, package_info: PackageInfo) -> None:
        pass

    def unregister_package(self, package_info: PackageInfo) -> None:
        pass  # FIXME: Remove package from dependencies

    def re_register_package(self, old_package_info: PackageInfo, new_package_info: PackageInfo) -> None:
        package_meta_io_factory = PackageMetaIOFactory(fs_editor=self.fs_editor)

        # Scan other packages to see if they are dependent on this one and update these dependency entries
        for other_package_info in self.package_index.list_package_infos():
            if other_package_info == old_package_info:
                continue

            for section_name in other_package_info.requirement_lists:
                if other_package_info.is_dependent_on(old_package_info, section_name=section_name):
                    new_req_rel_path = Path(os.path.relpath(new_package_info.abs_path, other_package_info.abs_path))
                    with package_meta_io_factory.package_meta_writer(other_package_info.toml_path) as pkg_meta_writer:
                        pkg_meta_writer.update_requirement_item(
                            section_name=section_name,
                            item_name=old_package_info.package_reg_name,
                            new_item_name=new_package_info.package_reg_name,
                            new_path=new_req_rel_path,
                        )

        # Update own dependency entries if the package has moved to another dir
        old_pkg_dir_path = old_package_info.abs_path.parent
        new_pkg_dir_path = new_package_info.abs_path.parent
        if new_pkg_dir_path != old_pkg_dir_path:
            for section_name, req_list in old_package_info.requirement_lists.items():
                updated_requirements: dict[str, PackageInfo] = {}  # {<pkg_reg_name>: <req_package_info>}
                for other_package_spec in req_list.req_specs:
                    if not isinstance(other_package_spec, LocalReqPackageSpec):
                        continue
                    req_package_info = self.package_index.get_package_info_by_reg_name(other_package_spec.package_name)
                    updated_requirements[other_package_spec.package_name] = req_package_info

                with package_meta_io_factory.package_meta_writer(new_package_info.toml_path) as pkg_meta_writer:
                    for _req_package_reg_name, req_package_info in updated_requirements.items():
                        updated_req_path = Path(os.path.relpath(req_package_info.abs_path, new_package_info.abs_path))
                        pkg_meta_writer.update_requirement_item(
                            section_name=section_name,
                            item_name=req_package_info.package_reg_name,
                            new_item_name=req_package_info.package_reg_name,
                            new_path=updated_req_path,
                        )


_PLUGIN_CLASSES: dict[str, Type[RepositoryManagementPlugin]] = {
    "common_tooling": CommonToolingRepositoryManagementPlugin,
    "toml_registration": MainTomlRepositoryManagementPlugin,
    "dependency_registration": DependencyReregistrationRepositoryManagementPlugin,
}


def get_plugin_cls(plugin_type: str) -> Type[RepositoryManagementPlugin]:
    return _PLUGIN_CLASSES[plugin_type]

terrarium/dl_repmanager/dl_repmanager/management_plugins.py
datalens-tech/datalens-backend
import itertools
import os.path
from pathlib import Path
import subprocess
from typing import Optional
from urllib.parse import urlparse

import attr
import tomlkit
from tomlkit import TOMLDocument
from tomlkit.items import Key

from dl_repmanager.primitives import (
    LocalReqPackageSpec,
    PypiReqPackageSpec,
    ReqPackageSpec,
)
from dl_repmanager.toml_tools import (
    TOMLReader,
    TOMLWriter,
)


@attr.s()
class MetaPackageManager:
    dir_path: Path = attr.ib()
    _toml: TOMLDocument = attr.ib(default=None)

    def __attrs_post_init__(self) -> None:
        toml_path = self.dir_path / "pyproject.toml"

        if self._toml is None:
            with open(toml_path) as fp:
                self._toml = tomlkit.load(fp)

    def get_reader(self) -> TOMLReader:
        return TOMLReader(toml=self._toml)

    def get_writer(self) -> TOMLWriter:
        return TOMLWriter(toml=self._toml)

    def as_new_location(self, new_path: Path) -> "MetaPackageManager":
        clone = MetaPackageManager(
            dir_path=new_path,
            toml=tomlkit.loads(tomlkit.dumps(self._toml)),
        )

        for group in itertools.chain([None], self.list_poetry_groups()):
            for dependency in self.get_dependencies(group):
                if isinstance(dependency, LocalReqPackageSpec):
                    dep_abs_path = Path(self.dir_path) / dependency.path
                    dep_new_rel_path = Path(os.path.relpath(dep_abs_path, new_path))

                    clone.write_dependency(group, attr.evolve(dependency, path=dep_new_rel_path))

        return clone

    def remove_dependency(self, group_name: Optional[str], pkg: str | ReqPackageSpec) -> None:
        w = self.get_writer()
        section = w.get_editable_section(self.get_dependencies_section_name(group_name))

        effective_pkg_name = pkg if isinstance(pkg, str) else pkg.package_name
        section.remove(effective_pkg_name)

    def write_dependency(self, group_name: Optional[str], pkg_spec: ReqPackageSpec) -> None:
        w = self.get_writer()
        section = w.get_editable_section(self.get_dependencies_section_name(group_name))
        section[pkg_spec.package_name] = pkg_spec.to_toml_value()

    def list_poetry_groups(self) -> list[str]:
        return list(self._toml["tool"]["poetry"]["group"].keys())

    def remove_poetry_group(self, name: str) -> None:
        del self._toml["tool"]["poetry"]["group"][name]

    def list_mypy_stubs_override(self) -> list[str]:
        mypy_section = self._toml["datalens"]["meta"]["mypy_stubs_packages_override"]
        return list(mypy_section.keys())

    def remove_mypy_stubs_override(self, name: str) -> None:
        mypy_section = self._toml["datalens"]["meta"]["mypy_stubs_packages_override"]
        del mypy_section[name]

    def remove_package_sources_section(self) -> None:
        del self._toml["tool"]["poetry"]["source"]

    @staticmethod
    def get_dependencies_section_name(group: Optional[str]) -> str:
        if group is None:
            return "tool.poetry.dependencies"
        return f"tool.poetry.group.{group}.dependencies"

    def get_dependencies(self, group: Optional[str] = None) -> list[ReqPackageSpec]:
        r = self.get_reader()

        ret: list[ReqPackageSpec] = []

        for key_str, val in r.iter_section_items(self.get_dependencies_section_name(group)):
            if key_str == "python":
                continue

            if isinstance(val, str):
                ret.append(PypiReqPackageSpec(package_name=key_str, version=val))
            elif isinstance(val, dict):
                if "path" in val:
                    ret.append(LocalReqPackageSpec(package_name=key_str, path=Path(val["path"])))
                else:
                    raise ValueError(f"No path in local dependency: {key_str}: {val}")
            else:
                raise ValueError(f"Unknown types of key/val in dependency: {key_str}: {val}")
        return ret

    def export_dependencies_raw(self, group: str) -> str:
        proc = subprocess.run(
            [
                "poetry",
                "export",
                "--only",
                group,
                "--without-hashes",
                "--format",
                "requirements.txt",
            ],
            cwd=self.dir_path,
            stdout=subprocess.PIPE,
            check=True,
        )
        dependencies = proc.stdout.decode("ascii")
        return dependencies

    def export_dependencies(self, group: str) -> list[ReqPackageSpec]:
        raw_deps = self.export_dependencies_raw(group)
        ret: list[ReqPackageSpec] = []
        for line in raw_deps.splitlines():
            if " ; " not in line:
                continue
            dep = line.split(";")[0].strip()
            if "@" in dep:
                pkg_name, pkg_url = dep.split("@")
                ret.append(
                    LocalReqPackageSpec(
                        package_name=pkg_name.strip(),
                        path=Path(
                            os.path.relpath(
                                urlparse(pkg_url.strip()).path,
                                self.dir_path,
                            )
                        ),
                    )
                )
            else:
                pkg_name = dep.split("==")[0]
                ret.append(PypiReqPackageSpec(package_name=pkg_name, version=dep.removeprefix(pkg_name)))
        return ret

    def run_poetry_lock(self, suppress_stdout: bool = False) -> None:
        stdout_target = subprocess.DEVNULL if suppress_stdout else None
        subprocess.run(["poetry", "lock", "--no-update"], cwd=self.dir_path, check=True, stdout=stdout_target)

    def save(self) -> None:
        file_path = self.dir_path / "pyproject.toml"

        with open(file_path, "w") as fp:
            fp.write(tomlkit.dumps(self._toml))

terrarium/dl_repmanager/dl_repmanager/metapkg_manager.py
datalens-tech/datalens-backend
import itertools
from pathlib import Path
import re
import shutil
import subprocess
from typing import Sequence

from dl_repmanager.metapkg_manager import MetaPackageManager
from dl_repmanager.primitives import (
    LocalReqPackageSpec,
    PypiReqPackageSpec,
)


def run_poetry_lock(dir_path: Path) -> None:
    subprocess.run(["poetry", "lock", "--no-update"], cwd=dir_path, check=True)


def sync_scoped_metapkg(
    *,
    original_metapkg_path: Path,
    scoped_metapkg_path: Path,
    package_dirs_to_include: Sequence[Path],  # Relative to scoped package
    prevent_prune_for_deps: Sequence[str],
    remove_private_pypi: bool = False,
    use_target_lock: bool = False,
):
    def is_local_package_in_scope(dep_from_scoped_pkg: LocalReqPackageSpec) -> bool:
        for package_dir in package_dirs_to_include:
            if dep_from_scoped_pkg.path.is_relative_to(package_dir):
                return True
        return False

    original_metapkg = MetaPackageManager(original_metapkg_path)
    scoped_metapkg = original_metapkg.as_new_location(scoped_metapkg_path)

    if not use_target_lock:
        shutil.copy(
            original_metapkg_path / "poetry.lock",
            scoped_metapkg_path / "poetry.lock",
        )

    scoped_metapkg.save()

    all_local_dependencies = scoped_metapkg.get_dependencies("ci")
    all_groups = scoped_metapkg.list_poetry_groups()
    declared_external_dependencies: list[PypiReqPackageSpec] = [
        dep for dep in scoped_metapkg.get_dependencies(None) if isinstance(dep, PypiReqPackageSpec)
    ]

    for dep in all_local_dependencies:
        if isinstance(dep, LocalReqPackageSpec) and not is_local_package_in_scope(dep):
            scoped_metapkg.remove_dependency("ci", dep)

    for group in all_groups:
        if group.startswith("app_") or "integration_tests" in group:
            scoped_metapkg.remove_poetry_group(group)

    scoped_metapkg.save()

    # Rerun lock
    run_poetry_lock(scoped_metapkg_path)

    ci_deps_raw = scoped_metapkg.export_dependencies_raw("ci")
    dev_deps_raw = scoped_metapkg.export_dependencies_raw("dev")

    effective_external_dependencies: set[str] = set()

    for line in itertools.chain(ci_deps_raw.splitlines(), dev_deps_raw.splitlines()):
        ext_dependency_re = re.compile(r"(^\S+)==\S+ ;")
        m = ext_dependency_re.match(line)

        if m:
            effective_external_dependencies.add(m.group(1).lower())

    for ext_dep in declared_external_dependencies:
        ext_dep_name = ext_dep.package_name.lower()

        if ext_dep_name not in effective_external_dependencies and ext_dep_name not in prevent_prune_for_deps:
            scoped_metapkg.remove_dependency(None, ext_dep.package_name)

    for mypy_override in scoped_metapkg.list_mypy_stubs_override():
        if mypy_override not in effective_external_dependencies and mypy_override not in prevent_prune_for_deps:
            scoped_metapkg.remove_mypy_stubs_override(mypy_override)

    scoped_metapkg.save()

    if remove_private_pypi:
        scoped_metapkg.remove_package_sources_section()
        scoped_metapkg.save()

    run_poetry_lock(scoped_metapkg_path)

terrarium/dl_repmanager/dl_repmanager/metapkg_scoped_sync.py
datalens-tech/datalens-backend
import logging
from pathlib import Path

import attr

from dl_repmanager.exceptions import MypyStubsOutOfSyncError
from dl_repmanager.package_meta_reader import (
    PackageMetaReader,
    PackageMetaWriter,
)
from dl_repmanager.pypi_tools import get_package_info_by_version
from dl_repmanager.requirements_tools import (
    PipRequirement,
    PipRequirementsIO,
)


log = logging.getLogger(__name__)


def _strip_version(raw: str):
    return raw.strip().split(" ;")[0]


def mb_types_name(name: str, override: dict | None = None) -> str | None:
    guessed_name = f"types-{name}"
    match override:
        case {"ignore": True}:
            return None
        case {"name": name}:
            return name
        case _:
            return guessed_name


def process_one(requirement: PipRequirement, override: dict | None = None):
    guessed_name = f"types-{requirement.name}"
    cleaned_version = requirement.cleaned_version

    log.debug(f"{requirement.name} override: {override}")
    match override:
        case {"name": override_name, "version": override_version}:
            return override_name, override_version
        case {"version": override_version}:
            if check_if_exists_in_pypi(guessed_name, override_version):
                return guessed_name, override_version
            else:
                log.error(f"PyPI doesn't have package {guessed_name} with version from overrides: {override_version}")
                return None
        case {"name": override_name}:
            if check_if_exists_in_pypi(override_name, cleaned_version):
                return override_name, cleaned_version
            else:
                log.debug(f"PyPI doesn't have package {guessed_name} with name from overrides: {override_name}")
                return None
        case {"ignore": True}:
            log.debug(f"Skip check for types- pkg fro {requirement.name} due to override")
            return None
        case _:
            log.debug(f"No override for {requirement.name}")

    if check_if_exists_in_pypi(guessed_name, cleaned_version):
        return guessed_name, cleaned_version

    return None


def check_if_exists_in_pypi(name: str, version: str) -> bool:
    log.debug(f"Trying to check presence of {name}=={version} on PyPI")
    data = get_package_info_by_version(name, version)
    if data.get("message") == "Not Found" or "info" not in data:
        log.info(f"Stub files package {name}:{version} is not available on PyPI")
        return False

    if "info" in data and not data["info"].get("yanked") and data["info"].get("name", "").lower() == name:
        log.debug(f"Stub files exists: {name} version {version}")
        return True

    log.info(f"Stub files package {name}:{version} is not available on PyPI")
    return False


@attr.s(frozen=True)
class RequirementsPathProvider:
    base_path: Path = attr.ib()
    external_requirements_rel_py: Path = attr.ib(default=Path("ops/ci/docker_image_base_ci"))
    external_requirements_file_names: tuple[str] = attr.ib(
        default=(
            "requirements_external.txt",
            "requirements_conflicting.txt",
        )
    )
    mypy_requirements_rel_path: Path = attr.ib(default=Path("ops/ci/docker_image_ci_mypy/requirements_types.txt"))

    def get_mypy_path(self) -> Path:
        return self.base_path / self.mypy_requirements_rel_path

    def get_external_requirements_path_list(self) -> list[Path]:
        result: list[Path] = []
        for name in self.external_requirements_file_names:
            result.append(self.base_path / self.external_requirements_rel_py / name)
        return result


def stubs_sync(
    meta_reader: PackageMetaReader,
    meta_writer: PackageMetaWriter,
    path_provider: RequirementsPathProvider,
    dry_run: bool = True,
):
    log.info(f"Starting mypy stubs sync. Dry run: {dry_run}")

    requirements: dict[str, PipRequirement] = {}
    for path in path_provider.get_external_requirements_path_list():
        requirements.update(PipRequirementsIO(path).read_existing())

    annotations_io = PipRequirementsIO(path_provider.get_mypy_path())
    annotations_requirements = annotations_io.read_existing()

    overrides_map = meta_reader.get_mypy_stubs_overrides()
    ignore_pkg_set = {k for k, v in overrides_map.items() if "ignore" in v}

    packages_to_ignore: list[str] = []
    annotations_to_add: dict[str, PipRequirement] = dict()

    for name, req in sorted(requirements.items()):
        override = overrides_map.get(name)
        result = process_one(req, override)
        if result:
            types_name, types_version = result
            if (
                types_name not in annotations_requirements
                or annotations_requirements.get(types_name).cleaned_version != types_version
            ):
                annotations_to_add[types_name] = PipRequirement(name=name, raw_version=types_version)
        else:
            log.debug(f"Could not find a types pkg for {name}")
            mb_name = mb_types_name(name, override)
            if not annotations_requirements.get(mb_name) and name not in overrides_map:
                packages_to_ignore.append(name)

            if annotations_requirements.get(mb_name) and name not in ignore_pkg_set:
                log.warning(
                    f"{name} and types package {mb_name} listed in requirements." f" Pypi does not have such packages"
                )

    if len(packages_to_ignore) + len(annotations_to_add) > 0:
        log.info(f"Need to add (or replace) in requirement_types.txt {len(annotations_to_add)} records")
        log.info(f"Need to add ignores to the pyproject.toml: {len(packages_to_ignore)} records")

        if not dry_run:
            annotations_io.write_updates(annotations_to_add)
            meta_writer.add_mypy_overrides_ignore(packages_to_ignore)
            log.info("Fixes done")
        else:
            raise MypyStubsOutOfSyncError(
                "At least some package missing stub packages, version is outdated or"
                " ignore mark should be added to the pyproject toml"
            )
    else:
        log.info("No changes required")

terrarium/dl_repmanager/dl_repmanager/mypy_stubs_sync.py
datalens-tech/datalens-backend
from pathlib import Path
from typing import Optional

import attr
from frozendict import frozendict

from dl_repmanager.package_meta_reader import PackageMetaIOFactory
from dl_repmanager.primitives import (
    LocalReqPackageSpec,
    PackageInfo,
    PypiReqPackageSpec,
    ReqPackageSpec,
    RequirementList,
)
from dl_repmanager.repository_env import RepoEnvironment


@attr.s
class PackageIndex:
    _built: bool = attr.ib(kw_only=True, default=False)
    _package_infos_by_reg_name: dict[str, PackageInfo] = attr.ib(kw_only=True, factory=dict)
    _package_infos_by_module_name: dict[str, PackageInfo] = attr.ib(kw_only=True, factory=dict)
    _package_infos_by_test_name: dict[str, PackageInfo] = attr.ib(kw_only=True, factory=dict)
    _package_infos_by_path: dict[Path, PackageInfo] = attr.ib(kw_only=True, factory=dict)

    def _is_built(self) -> bool:
        return self._built

    def get_package_info_by_reg_name(self, package_reg_name: str) -> PackageInfo:
        assert self._is_built()
        return self._package_infos_by_reg_name[package_reg_name]

    def get_package_info_from_path(self, rel_package_dir_name: Path) -> PackageInfo:
        assert self._is_built()
        return self._package_infos_by_path[rel_package_dir_name]

    def get_package_info_from_module_name(self, package_module_name: str) -> PackageInfo:
        assert self._is_built()
        return self._package_infos_by_module_name[package_module_name]

    def get_package_info_from_test_name(self, test_dir_name: str) -> PackageInfo:
        assert self._is_built()
        return self._package_infos_by_test_name[test_dir_name]

    def list_package_infos(self, package_type: Optional[str] = None) -> list[PackageInfo]:
        return [
            package_info
            for package_info in self._package_infos_by_path.values()
            if package_type is None or package_info.package_type == package_type
        ]


@attr.s
class PackageIndexBuilder:
    repository_env: RepoEnvironment = attr.ib(kw_only=True)
    load_requirements: bool = attr.ib(kw_only=True, default=True)

    def _req_spec_from_dict(self, item_as_dict: dict) -> ReqPackageSpec:
        if "path" in item_as_dict:
            return LocalReqPackageSpec(package_name=item_as_dict["name"], path=item_as_dict["path"])
        return PypiReqPackageSpec(package_name=item_as_dict["name"], version=item_as_dict["version"])

    def _load_package_info_from_package_dir(
        self, abs_package_dir_path: Path, default_package_type: str
    ) -> Optional[PackageInfo]:
        # TODO: Move to a separate class
        toml_path = abs_package_dir_path / "pyproject.toml"

        if not toml_path.exists():
            return None

        req_list_names = [
            "tool.poetry.dependencies",
            "tool.poetry.group.tests.dependencies",
        ]
        requirement_lists: dict[str, RequirementList] = {}
        package_meta_io_factory = PackageMetaIOFactory(fs_editor=self.repository_env.get_fs_editor())
        with package_meta_io_factory.package_meta_reader(toml_path) as pkg_meta_reader:
            package_reg_name = pkg_meta_reader.get_package_reg_name()
            module_names = pkg_meta_reader.get_package_module_names()
            package_type = pkg_meta_reader.get_package_type() or default_package_type
            implicit_reqs = frozenset(pkg_meta_reader.get_implicit_dependencies())
            i18n_domains = pkg_meta_reader.get_i18n_domains()
            if self.load_requirements:
                for req_list_name in req_list_names:
                    requirement_lists[req_list_name] = RequirementList(
                        req_specs=tuple(
                            self._req_spec_from_dict(item)
                            for item in pkg_meta_reader.iter_requirement_items(req_list_name)
                        ),
                    )

        test_dirs: list[str] = []
        for test_dir in abs_package_dir_path.iterdir():  # TODO: read info about tests from pyproject.toml
            if not test_dir.is_dir():
                continue
            if "tests" not in test_dir.name or test_dir.name in module_names:
                continue  # it is not a test dir
            test_dirs.append(test_dir.name)

        package_info = PackageInfo(
            package_type=package_type,
            package_reg_name=package_reg_name,
            abs_path=abs_package_dir_path,
            module_names=module_names,
            test_dirs=tuple(test_dirs),
            requirement_lists=frozendict(requirement_lists),
            implicit_deps=implicit_reqs,
            i18n_domains=i18n_domains,
        )
        return package_info

    def build_index(self) -> PackageIndex:
        package_infos_by_reg_name: dict[str, PackageInfo] = {}
        package_infos_by_module_name: dict[str, PackageInfo] = {}
        package_infos_by_test_name: dict[str, PackageInfo] = {}
        package_infos_by_path: dict[Path, PackageInfo] = {}

        for package_type, abs_parent_dir in self.repository_env.iter_package_abs_dirs():
            for package_dir in abs_parent_dir.iterdir():
                if not package_dir.is_dir():
                    continue

                package_info = self._load_package_info_from_package_dir(
                    package_dir,
                    default_package_type=package_type,
                )
                if package_info is None:
                    continue
                package_infos_by_reg_name[package_info.package_reg_name] = package_info
                package_infos_by_path[package_dir] = package_info

                for module_name in package_info.module_names:
                    package_infos_by_module_name[module_name] = package_info

                for test_dir in package_info.test_dirs:
                    package_infos_by_test_name[test_dir] = package_info

        package_index = PackageIndex(
            built=True,
            package_infos_by_reg_name=package_infos_by_reg_name,
            package_infos_by_path=package_infos_by_path,
            package_infos_by_module_name=package_infos_by_module_name,
            package_infos_by_test_name=package_infos_by_test_name,
        )
        return package_index

terrarium/dl_repmanager/dl_repmanager/package_index.py
datalens-tech/datalens-backend
from __future__ import annotations

import contextlib
import logging
from pathlib import Path
from typing import (
    Any,
    Generator,
    Iterable,
    Optional,
)

import attr
import tomlkit

from dl_repmanager.fs_editor import FilesystemEditor
from dl_repmanager.primitives import LocaleDomainSpec
from dl_repmanager.toml_tools import (
    TOMLIOFactory,
    TOMLReaderBase,
    TOMLWriter,
)


log = logging.getLogger(__name__)


@attr.s
class PackageMetaReader:
    """
    Class for reading meta info of Python packages.
    It is aware of the project-specific pyproject.toml structure.
    """

    _toml_reader: TOMLReaderBase = attr.ib(kw_only=True)

    _SECTION_NAME_MAIN = "tool.poetry"
    _SECTION_NAME_META = "datalens.meta"  # get this from env?
    _SECTION_NAME_I18N_DOMAINS = "datalens.i18n.domains"

    def _get_main_section(self) -> dict:
        return self._toml_reader.get_section(self._SECTION_NAME_MAIN)

    def _get_meta_section(self) -> dict:
        return self._toml_reader.get_section(self._SECTION_NAME_META, strict=False)

    def get_package_reg_name(self) -> str:
        package_reg_name = self._get_main_section()["name"]
        assert isinstance(package_reg_name, str)
        return package_reg_name

    def get_package_module_names(self) -> tuple[str, ...]:
        module_names = tuple(pkg["include"] for pkg in self._get_main_section()["packages"])  # type: ignore
        return module_names

    def get_package_type(self) -> Optional[str]:  # FIXME: remove `Optional` after it is added to all toml files
        package_reg_name = self._get_meta_section().get("package_type")
        assert package_reg_name is None or isinstance(package_reg_name, str)
        return package_reg_name

    def iter_requirement_items(self, section_name: str) -> Iterable[dict[str, Any]]:
        for key, item in self._toml_reader.iter_section_items(section_name, strict=False):
            if key is None:
                continue

            item_as_dict = {"name": key}
            if isinstance(item.value, str):
                item_as_dict["version"] = item.value
            elif isinstance(item.value, dict):
                item_as_dict.update(item.value)

            yield item_as_dict

    def get_i18n_domains(self) -> tuple[LocaleDomainSpec, ...]:
        result: set[LocaleDomainSpec] = set()
        for domain, item_data in self._toml_reader.iter_section_items(self._SECTION_NAME_I18N_DOMAINS, strict=False):
            domain_str = str(domain).strip()
            paths: list[Path] = []
            for path_data in item_data:
                paths.append(Path(str(path_data["path"]).strip()))

            result.add(LocaleDomainSpec(domain_name=domain_str, scan_paths=tuple(paths)))

        return tuple(result)

    def get_implicit_dependencies(self) -> list[str]:
        meta_section = self._toml_reader.get_section(self._SECTION_NAME_META, strict=False)
        implicit_deps = meta_section.get("implicit_dependencies", ())
        return [item for item in implicit_deps]

    def get_mypy_stubs_overrides(self) -> dict:
        return dict(self._toml_reader.get_section(f"{self._SECTION_NAME_META}.mypy_stubs_packages_override"))

    def get_mypy_common(self) -> dict:
        return dict(self._toml_reader.get_section(f"{self._SECTION_NAME_META}.mypy_common", strict=False))

    def get_mypy_overrides_modules(self) -> list[dict]:
        section = self._toml_reader.get_section("tool.mypy", strict=False)
        return list(section.get("overrides", []))


@attr.s
class PackageMetaWriter(PackageMetaReader):
    @property
    def toml_writer(self) -> TOMLWriter:
        assert isinstance(self._toml_reader, TOMLWriter)
        return self._toml_reader

    def remove_requirement_item(self, section_name: str, item_name: str) -> None:
        with self.toml_writer.suppress_non_existent_key():
            section = self.toml_writer.get_editable_section(section_name)
            section.remove(item_name)

    def _get_item_opt(self, section_name: str, item_name) -> Optional[dict[str, Any]]:
        items = [item for item in self.iter_requirement_items(section_name=section_name) if item["name"] == item_name]
        if not items:
            return None

        assert len(items) == 1
        return items[0]

    def update_requirement_item(
        self,
        section_name: str,
        item_name: str,
        new_item_name: str,
        new_path: Path,
    ) -> None:
        original_item = self._get_item_opt(section_name=section_name, item_name=item_name)
        if original_item is not None:
            section = self.toml_writer.get_editable_section(section_name)
            section.remove(item_name)
            package_dep_table = tomlkit.inline_table()
            package_dep_table.add("path", str(new_path))
            section.add(new_item_name, package_dep_table)

    def add_mypy_overrides_ignore(self, pkg_names: list[str]):
        with self.toml_writer.suppress_non_existent_key():
            section = self.toml_writer.get_editable_section(f"{self._SECTION_NAME_META}.mypy_stubs_packages_override")
            for name in pkg_names:
                override = tomlkit.inline_table()
                override.add("ignore", True)
                section.add(name, override)

    def update_mypy_common(self, to_update: dict[str, str]) -> bool:
        changes = False
        with self.toml_writer.suppress_non_existent_key():
            key = "tools.mypy"
            try:
                section = self.toml_writer.get_editable_section(key)
            except Exception as err:
                print(err)
                section = self.toml_writer.add_section(key)

            for k, v in to_update.items():
                if section[k] != v:
                    changes = True
                section[k] = v
        return changes


@attr.s
class PackageMetaIOFactory:
    fs_editor: FilesystemEditor = attr.ib(kw_only=True)
    toml_io_factory: TOMLIOFactory = attr.ib(init=False)

    @toml_io_factory.default
    def __make_toml_io_factory(self) -> TOMLIOFactory:
        return TOMLIOFactory(fs_editor=self.fs_editor)

    @contextlib.contextmanager
    def package_meta_reader(self, file_path: Path) -> Generator[PackageMetaReader, None, None]:
        with self.toml_io_factory.toml_reader(file_path) as toml_reader:
            yield PackageMetaReader(toml_reader=toml_reader)

    @contextlib.contextmanager
    def package_meta_writer(self, file_path: Path) -> Generator[PackageMetaWriter, None, None]:
        with self.toml_io_factory.toml_writer(file_path) as toml_writer:
            yield PackageMetaWriter(toml_reader=toml_writer)

terrarium/dl_repmanager/dl_repmanager/package_meta_reader.py
datalens-tech/datalens-backend
from importlib.metadata import packages_distributions

import attr

from dl_repmanager.package_index import PackageIndex
from dl_repmanager.primitives import (
    LocalReqPackageSpec,
    PypiReqPackageSpec,
    ReqPackageSpec,
)
from dl_repmanager.repository_env import RepoEnvironment


@attr.s
class PackageReference:
    repository_env: RepoEnvironment = attr.ib(kw_only=True)
    package_index: PackageIndex = attr.ib(kw_only=True)

    def get_package_req_spec(self, module_name: str) -> ReqPackageSpec:
        spec: ReqPackageSpec
        try:
            local_pkg_info = self.package_index.get_package_info_from_module_name(package_module_name=module_name)
            spec = LocalReqPackageSpec(
                package_name=local_pkg_info.package_reg_name,
                path=local_pkg_info.abs_path,  # FIXME: Use some kind of relative path here
            )
        except KeyError:
            try:
                package_name = packages_distributions()[module_name][0]
            except KeyError:
                package_name = self.repository_env.custom_package_map.get(module_name, module_name)

            spec = PypiReqPackageSpec(package_name=package_name, version=None)

        return spec

terrarium/dl_repmanager/dl_repmanager/package_reference.py
datalens-tech/datalens-backend
from __future__ import annotations

import abc
from ast import AST
from enum import (
    Enum,
    auto,
)
import os
from pathlib import Path
from typing import (
    Any,
    Optional,
    TypeVar,
)

import attr
from frozendict import frozendict
from tomlkit import inline_table


_CLONABLE_TV = TypeVar("_CLONABLE_TV", bound="_Clonable")


@attr.s(frozen=True)
class _Clonable(abc.ABC):
    def clone(self: _CLONABLE_TV, **kwargs: Any) -> _CLONABLE_TV:
        return attr.evolve(self, **kwargs)


class EntityReferenceType(Enum):
    package_type = auto()
    package_reg = auto()
    module = auto()
    path = auto()


@attr.s(frozen=True)
class EntityReference:
    ref_type: EntityReferenceType = attr.ib(kw_only=True)
    name: str = attr.ib(kw_only=True)


@attr.s(frozen=True)
class ReqPackageSpec(_Clonable):
    package_name: str = attr.ib(kw_only=True)

    @abc.abstractmethod
    def pretty(self) -> str:
        raise NotImplementedError

    @abc.abstractmethod
    def as_req_str(self) -> str:
        raise NotImplementedError

    @abc.abstractmethod
    def to_toml_value(self) -> Any:
        raise NotImplementedError


@attr.s(frozen=True)
class PypiReqPackageSpec(ReqPackageSpec):
    version: Optional[str] = attr.ib(kw_only=True)

    def pretty(self) -> str:
        extra = ""
        if self.version:
            extra = f" ({self.version})"
        return f"{self.package_name}{extra}"

    def is_exact_version(self) -> bool:
        return bool(self.version and self.version.startswith("=="))

    def get_exact_version(self) -> str:
        if not self.is_exact_version():
            raise ValueError(f"{str(self)} does not have an exact version")
        assert self.version is not None
        return self.version[2:]

    def as_req_str(self) -> str:
        if not self.version:
            return self.package_name
        return f'{self.package_name} = "{self.version}"'

    def to_toml_value(self) -> Any:
        return self.version if self.version else "*"


@attr.s(frozen=True)
class LocalReqPackageSpec(ReqPackageSpec):
    path: Path = attr.ib(kw_only=True)

    def pretty(self) -> str:
        return f"{self.package_name} ({self.path})"

    def as_req_str(self) -> str:
        return f'{self.package_name} = {{path = "{self.path}"}}'

    def to_toml_value(self) -> Any:
        it = inline_table()
        it["path"] = str(self.path)
        return it


@attr.s(frozen=True)
class RequirementList:
    req_specs: tuple[ReqPackageSpec, ...] = attr.ib(kw_only=True, default=())


@attr.s(frozen=True)
class LocaleDomainSpec(_Clonable):
    domain_name: str = attr.ib(kw_only=True)
    scan_paths: tuple[Path, ...] = attr.ib(kw_only=True)


@attr.s(frozen=True)
class PackageInfo(_Clonable):
    package_type: str = attr.ib(kw_only=True)
    package_reg_name: str = attr.ib(kw_only=True)
    abs_path: Path = attr.ib(kw_only=True)
    module_names: tuple[str, ...] = attr.ib(kw_only=True)
    test_dirs: tuple[str, ...] = attr.ib(kw_only=True, default=())
    requirement_lists: frozendict[str, RequirementList] = attr.ib(kw_only=True, default=frozendict())
    implicit_deps: frozenset[str] = attr.ib(kw_only=True, default=frozenset())
    i18n_domains: tuple[LocaleDomainSpec, ...] = attr.ib(kw_only=True, default=())

    @property
    def toml_path(self) -> Path:
        return self.abs_path / "pyproject.toml"

    @property
    def single_module_name(self) -> str:
        assert len(self.module_names) == 1
        return self.module_names[0]

    @property
    def single_test_dir(self) -> str:
        assert len(self.test_dirs) == 1
        return self.test_dirs[0]

    @property
    def reg_entity_ref(self) -> EntityReference:
        return EntityReference(
            ref_type=EntityReferenceType.package_reg,
            name=self.package_reg_name,
        )

    @property
    def package_type_entity_ref(self) -> EntityReference:
        return EntityReference(
            ref_type=EntityReferenceType.package_type,
            name=self.package_type,
        )

    @property
    def tests_and_modules(self) -> tuple[str, ...]:
        return tuple(*self.module_names, *self.test_dirs)

    def get_relative_path(self, base_path: Path) -> Path:
        return Path(os.path.relpath(self.abs_path, base_path))

    def is_dependent_on(
        self,
        base_package_info: PackageInfo,
        section_name: str,
    ) -> bool:
        req_specs = self.requirement_lists.get(section_name, RequirementList()).req_specs
        for req_spec in req_specs:
            if req_spec.package_name == base_package_info.package_reg_name:
                # It really is a dependent package
                assert isinstance(req_spec, LocalReqPackageSpec)
                return True

        return False


@attr.s(frozen=True)
class ImportSpec:
    import_ast: AST = attr.ib(kw_only=True)
    import_module_name: str = attr.ib(kw_only=True)
    source_path: Path = attr.ib(kw_only=True)


@attr.s(frozen=True)
class MetaPackageSpec(_Clonable):
    name: str = attr.ib(kw_only=True)
    toml_path: Path = attr.ib(kw_only=True)

terrarium/dl_repmanager/dl_repmanager/primitives.py
datalens-tech/datalens-backend
"""
Bunch of logic to edit pyproject.toml's across entire repo
"""
from functools import cached_property
from logging import getLogger
from pathlib import Path
import subprocess

import attr

from dl_repmanager.package_index import PackageIndex
from dl_repmanager.package_meta_reader import (
    PackageMetaIOFactory,
    PackageMetaReader,
)
from dl_repmanager.primitives import MetaPackageSpec
from dl_repmanager.repository_env import RepoEnvironment


log = getLogger()


@attr.s
class PyPrjEditor:
    repository_env: RepoEnvironment = attr.ib(kw_only=True)
    package_index: PackageIndex = attr.ib(kw_only=True)

    @cached_property
    def base_path(self) -> Path:
        return self.repository_env.base_path

    @cached_property
    def package_meta_io_factory(self) -> PackageMetaIOFactory:
        return PackageMetaIOFactory(fs_editor=self.repository_env.fs_editor)

    def _get_meta_project_reader(self, metapackage_name: str) -> PackageMetaReader:
        toml_path = self.repository_env.get_metapackage_spec(metapackage_name=metapackage_name).toml_path
        with self.package_meta_io_factory.package_meta_reader(toml_path) as reader:
            return reader

    def update_mypy_common(self, metapackage_name: str) -> None:
        data = self._get_meta_project_reader(metapackage_name=metapackage_name).get_mypy_common()
        for pi in self.package_index.list_package_infos():
            with self.package_meta_io_factory.package_meta_writer(pi.toml_path) as pmw:
                log.debug(f"Checking {pi.abs_path}")
                pmw.update_mypy_common(data)

            log.debug(f"pyproject-fmt {str(pi.toml_path)}")
            subprocess.run(f"pyproject-fmt {str(pi.abs_path)}", shell=True)
            (pi.abs_path / "mypy.ini").unlink(missing_ok=True)

terrarium/dl_repmanager/dl_repmanager/project_editor.py
datalens-tech/datalens-backend
"""
Simple methods to get package info from public PyPI
"""
import requests


DEFAULT_ENDPOINT = "https://pypi.org/pypi/"


def get_package_info(package_name: str, override_endpoint: str | None = None) -> dict:
    endpoint = override_endpoint or DEFAULT_ENDPOINT
    return requests.get(url=f"{endpoint}{package_name}/json").json()


def get_package_info_by_version(package_name: str, version: str, override_endpoint: str | None = None) -> dict:
    endpoint = override_endpoint or DEFAULT_ENDPOINT
    return requests.get(url=f"{endpoint}{package_name}/{version}/json").json()

terrarium/dl_repmanager/dl_repmanager/pypi_tools.py
datalens-tech/datalens-backend
"""
Load environment from config file (dl-repo.yml).
The file should have a structure similar to this:

    dl_repo:
      fs_editor: git

      include:
        - core_repo/dl-repo.yml

      default_boilerplate_path: lib/dl_package_boilerplate

      package_types:
        - type: lib
          root_path: lib

        - type: app
          root_path: app
          boilerplate_path: app/app_package_boilerplate
          tags:
            - own_dependency_group

      custom_package_map:
        flask_marshmallow: flask-marshmallow
        jwt: pyjwt

      metapackages:
        - name: main
          toml: tools/pyproject.toml

      plugins:
        - type: dependency_registration

      edit_exclude_masks:
        - ".*\\.mo",


Description of the sections:
  - include: section tells the loader to include another repo config file
    (e.g. for nested repositories).
  - fs_editor: allows customization of FS operations
    so that move/copy actions are immediately registered in the VCS
  - package_types: defines packages types.
    This defines where to search for existing packages and how new packages are created.
  - custom_package_map: custom mapping of third-party module to package names
    (for validation of package requirements)

"""

from __future__ import annotations

from pathlib import Path
import re
from typing import (
    TYPE_CHECKING,
    Iterable,
    Optional,
)

import attr
import yaml

from dl_repmanager.fs_editor import (
    FilesystemEditor,
    get_fs_editor,
)
from dl_repmanager.management_plugins import (
    RepositoryManagementPlugin,
    get_plugin_cls,
)
from dl_repmanager.primitives import MetaPackageSpec


if TYPE_CHECKING:
    from dl_repmanager.package_index import PackageIndex


DEFAULT_CONFIG_FILE_NAME = "dl-repo.yml"


@attr.s(frozen=True)
class PackageTypeConfig:
    home_repo_path: Path = attr.ib(kw_only=True)
    path: Path = attr.ib(kw_only=True)
    boilerplate_path: Path = attr.ib(kw_only=True)
    tags: frozenset[str] = attr.ib(kw_only=True, default=frozenset())


@attr.s(frozen=True)
class PluginConfig:
    plugin_type: str = attr.ib(kw_only=True)
    config_data: dict = attr.ib(kw_only=True)


@attr.s(frozen=True)
class RepoEnvironment:
    """
    Provides information about repository folders and boilerplates
    """

    base_path: Path = attr.ib(kw_only=True)
    package_types: dict[str, PackageTypeConfig] = attr.ib(kw_only=True)
    metapackages: dict[str, MetaPackageSpec] = attr.ib(kw_only=True)
    custom_package_map: dict[str, str] = attr.ib(kw_only=True, factory=dict)
    fs_editor: FilesystemEditor = attr.ib(kw_only=True)
    plugin_configs: list[PluginConfig] = attr.ib(kw_only=True, factory=list)
    edit_exclude_masks: frozenset[re.Pattern] = attr.ib(kw_only=True, default=frozenset())

    def iter_package_abs_dirs(self) -> Iterable[tuple[str, Path]]:
        return sorted(
            [(package_type, pkg_type_config.path) for package_type, pkg_type_config in self.package_types.items()],
            key=lambda pair: pair[0],
        )

    def get_boilerplate_package_dir(self, package_type: str) -> Path:
        return self.package_types[package_type].boilerplate_path

    def get_root_package_dir(self, package_type: str) -> Path:
        return self.package_types[package_type].path

    def get_tags(self, package_type: str) -> frozenset[str]:
        return self.package_types[package_type].tags

    def get_plugins(
        self,
        package_index: PackageIndex,
    ) -> list[RepositoryManagementPlugin]:
        # TODO: parameterize and load plugins from config
        return [
            get_plugin_cls(plugin_type=plugin_config.plugin_type)(  # type: ignore
                repository_env=self,
                package_index=package_index,
                base_path=self.base_path,
                config_data=plugin_config.config_data,
            )
            for plugin_config in self.plugin_configs
        ]

    def get_fs_editor(self) -> FilesystemEditor:
        return self.fs_editor

    def get_metapackage_spec(self, metapackage_name: str) -> MetaPackageSpec:
        return self.metapackages[metapackage_name]

    def get_edit_exclude_masks(self) -> frozenset[re.Pattern]:
        return self.edit_exclude_masks


_DEFAULT_FS_EDITOR_TYPE = "default"


@attr.s(frozen=True)
class ConfigContents:
    base_path: Path = attr.ib(kw_only=True)
    package_types: dict[str, PackageTypeConfig] = attr.ib(kw_only=True, factory=dict)
    metapackages: dict[str, MetaPackageSpec] = attr.ib(kw_only=True, factory=dict)
    custom_package_map: dict[str, str] = attr.ib(kw_only=True, factory=dict)
    fs_editor_type: Optional[str] = attr.ib(kw_only=True, default=_DEFAULT_FS_EDITOR_TYPE)
    plugin_configs: list[PluginConfig] = attr.ib(kw_only=True, factory=list)
    edit_exclude_masks: frozenset[re.Pattern] = attr.ib(kw_only=True, default=frozenset())


def discover_config(base_path: Path, config_file_name: str) -> Path:
    while base_path and base_path.exists():
        config_path = base_path / config_file_name
        if config_path.exists():
            return config_path
        base_path = base_path.parent

    raise RuntimeError("Failed to discover the repo config file. This does not seem to be a managed repository.")


@attr.s
class RepoEnvironmentLoader:
    config_file_name: str = attr.ib(kw_only=True, default=DEFAULT_CONFIG_FILE_NAME)
    override_fs_editor_type: Optional[str] = attr.ib(kw_only=True, default=None)

    def _load_params_from_yaml_file(self, config_path: Path) -> ConfigContents:
        with open(config_path) as config_file:
            config_data = yaml.safe_load(config_file)

        base_path = config_path.parent
        env_settings = config_data.get("dl_repo", {})

        default_boilerplate_path_str: Optional[str] = env_settings.get("default_boilerplate_path")

        package_types: dict[str, PackageTypeConfig] = {}
        for package_type_data in env_settings.get("package_types", ()):
            package_type = package_type_data["type"]
            boilerplate_path_str = package_type_data.get("boilerplate_path", default_boilerplate_path_str)
            if boilerplate_path_str is None:
                raise ValueError("Boilerplate must be specified in package type or default")
            pkg_type_config = PackageTypeConfig(
                home_repo_path=base_path,
                path=base_path / package_type_data["root_path"],
                boilerplate_path=base_path / boilerplate_path_str,
                tags=frozenset(package_type_data.get("tags", ())),
            )
            package_types[package_type] = pkg_type_config

        metapackages: dict[str, MetaPackageSpec] = {}
        for metapackage_data in env_settings.get("metapackages", ()):
            metapackage_name = metapackage_data["name"]
            metapackage_config = MetaPackageSpec(
                name=metapackage_name,
                toml_path=base_path / metapackage_data["toml"],
            )
            metapackages[metapackage_name] = metapackage_config

        plugin_configs: list[PluginConfig] = []
        for plugin_data in env_settings.get("plugins", ()):
            plugin_type = plugin_data["type"]
            plugin_config = PluginConfig(
                plugin_type=plugin_type,
                config_data=plugin_data.get("config") or {},
            )
            plugin_configs.append(plugin_config)

        edit_exclude_masks: set[re.Pattern] = set()
        for edit_exclude_masks_item_str in env_settings.get("edit_exclude_masks", ()):
            edit_exclude_masks.add(re.compile(edit_exclude_masks_item_str))

        custom_package_map: dict[str, str] = dict(env_settings.get("custom_package_map", {}))

        fs_editor_type: Optional[str] = env_settings.get("fs_editor")

        for include in env_settings.get("include", ()):
            included_config_contents = self._load_params_from_yaml_file(Path(base_path) / include)
            package_types = dict(included_config_contents.package_types, **package_types)
            metapackages = dict(included_config_contents.metapackages, **metapackages)
            custom_package_map = dict(included_config_contents.custom_package_map, **custom_package_map)
            edit_exclude_masks |= included_config_contents.edit_exclude_masks
            fs_editor_type = fs_editor_type or included_config_contents.fs_editor_type

        # FS editor and plugins are loaded only from the top-level config
        return ConfigContents(
            base_path=base_path,
            package_types=package_types,
            metapackages=metapackages,
            custom_package_map=custom_package_map,
            fs_editor_type=fs_editor_type,
            plugin_configs=plugin_configs,
            edit_exclude_masks=frozenset(edit_exclude_masks),
        )

    def _load_from_yaml_file(self, config_path: Path) -> RepoEnvironment:
        config_contents = self._load_params_from_yaml_file(config_path)
        base_path = config_contents.base_path

        fs_editor_type = config_contents.fs_editor_type or _DEFAULT_FS_EDITOR_TYPE
        if self.override_fs_editor_type is not None:
            fs_editor_type = self.override_fs_editor_type
        assert fs_editor_type is not None

        return RepoEnvironment(
            base_path=base_path,
            package_types=config_contents.package_types,
            metapackages=config_contents.metapackages,
            custom_package_map=config_contents.custom_package_map,
            plugin_configs=config_contents.plugin_configs,
            edit_exclude_masks=config_contents.edit_exclude_masks,
            fs_editor=get_fs_editor(
                fs_editor_type=fs_editor_type,
                base_path=base_path,
            ),
        )

    def load_env(self, base_path: Path) -> RepoEnvironment:
        return self._load_from_yaml_file(config_path=discover_config(base_path, self.config_file_name))

terrarium/dl_repmanager/dl_repmanager/repository_env.py
datalens-tech/datalens-backend
import itertools
import os
from pathlib import Path
import re
from typing import (
    Callable,
    Collection,
    Mapping,
    Optional,
    Sequence,
)

import attr

from dl_repmanager.fs_editor import FilesystemEditor
from dl_repmanager.package_index import PackageIndex
from dl_repmanager.package_reference import PackageReference
from dl_repmanager.primitives import (
    LocaleDomainSpec,
    LocalReqPackageSpec,
    PackageInfo,
    ReqPackageSpec,
    RequirementList,
)
from dl_repmanager.repository_env import RepoEnvironment
from dl_repmanager.repository_navigator import RepositoryNavigator


@attr.s
class PackageGenerator:
    package_index: PackageIndex = attr.ib(kw_only=True)
    repository_env: RepoEnvironment = attr.ib(kw_only=True)

    def _generate_new_package_abs_path(self, package_module_name: str, package_type: str) -> Path:
        """Generate package path for new package"""
        return self.repository_env.get_root_package_dir(package_type=package_type) / package_module_name

    @staticmethod
    def generate_new_package_reg_name(package_name: str) -> str:
        internal_name = package_name.replace("_", "-").replace("dl-", "datalens-")
        return internal_name

    @staticmethod
    def generate_default_test_dir_name(package_name: str) -> str:
        return f"{package_name}_tests"

    def get_boilerplate_package_info(self, package_type: str) -> PackageInfo:
        return self.package_index.get_package_info_from_path(
            self.repository_env.get_boilerplate_package_dir(package_type=package_type),
        )

    def generate_package_info_from_boilerplate(
        self,
        boilerplate_info: PackageInfo,
        package_type: str,
        package_module_name: str,
    ) -> PackageInfo:
        package_test_dirs: list[str] = []
        for boilerplate_test_dir in boilerplate_info.test_dirs:
            package_test_dir = boilerplate_test_dir.replace(  # FIXME: define in templates
                boilerplate_info.single_module_name, package_module_name
            )
            package_test_dirs.append(package_test_dir)

        generated_package_info = PackageInfo(
            package_type=package_type,
            package_reg_name=self.generate_new_package_reg_name(package_module_name),
            module_names=(package_module_name,),  # FIXME: define in template
            abs_path=self._generate_new_package_abs_path(
                package_module_name=package_module_name,
                package_type=package_type,
            ),
            test_dirs=tuple(package_test_dirs),
        )
        return generated_package_info


@attr.s
class RepositoryManager:
    repository_env: RepoEnvironment = attr.ib(kw_only=True)
    package_index: PackageIndex = attr.ib(kw_only=True)
    repository_navigator: RepositoryNavigator = attr.ib(kw_only=True)
    package_generator: PackageGenerator = attr.ib(kw_only=True)
    package_reference: PackageReference = attr.ib(kw_only=True)
    fs_editor: FilesystemEditor = attr.ib(init=False)

    @fs_editor.default
    def _make_fs_editor(self) -> FilesystemEditor:
        return self.repository_env.get_fs_editor()

    def _replace_imports(self, old_import_name: str, new_import_name: str) -> list[Path]:
        """Replace imports in all project files, return list of updated files"""

        regex, repl = self._make_regex_and_repl_for_sub(
            old_str=old_import_name,
            new_str=new_import_name,
            allow_dash=True,  # dash can be next to the module name
        )

        updated_files: list[Path] = []
        for file_path in self.repository_navigator.recurse_files_with_import(old_import_name):
            # TODO: Refactor replacement to work by finding the imports in the module AST
            self.fs_editor.replace_regex_in_file(file_path=file_path, regex=regex, repl=repl)
            updated_files.append(file_path)

        return updated_files

    def _register_package(self, package_info: PackageInfo) -> None:
        for mng_plugin in self.repository_env.get_plugins(package_index=self.package_index):
            mng_plugin.register_package(package_info=package_info)

    def _unregister_package(self, package_info: PackageInfo) -> None:
        for mng_plugin in self.repository_env.get_plugins(package_index=self.package_index):
            mng_plugin.register_package(package_info=package_info)

    def _re_register_package(self, old_package_info: PackageInfo, new_package_info: PackageInfo) -> None:
        for mng_plugin in self.repository_env.get_plugins(package_index=self.package_index):
            mng_plugin.re_register_package(old_package_info=old_package_info, new_package_info=new_package_info)

    def init_package(self, package_module_name: str, package_type: str) -> PackageInfo:
        # Get boilerplate and generate new package meta
        boilerplate_info = self.package_generator.get_boilerplate_package_info(package_type=package_type)
        package_info = self.package_generator.generate_package_info_from_boilerplate(
            boilerplate_info=boilerplate_info,
            package_module_name=package_module_name,
            package_type=package_type,
        )

        # Initialize the package
        self._init_package_from_boilerplate(boilerplate_info=boilerplate_info, package_info=package_info)
        return package_info

    def copy_package(self, package_module_name: str, from_package_module_name: str) -> PackageInfo:
        # Get boilerplate and generate new package meta
        boilerplate_info = self.package_index.get_package_info_from_module_name(from_package_module_name)
        package_info = self.package_generator.generate_package_info_from_boilerplate(
            boilerplate_info=boilerplate_info,
            package_module_name=package_module_name,
            package_type=boilerplate_info.package_type,
        )

        # Initialize the package
        self._init_package_from_boilerplate(boilerplate_info=boilerplate_info, package_info=package_info)
        return package_info

    def remove_package(self, package_module_name: str) -> None:
        package_info = self.package_index.get_package_info_from_module_name(package_module_name)

        # Unregister package
        self._unregister_package(package_info=package_info)

        # Remove the directory
        self.fs_editor.remove_path(package_info.abs_path)

    def _generate_info_for_renamed_package(
        self,
        old_package_info: PackageInfo,
        new_package_module_name: str,
    ) -> PackageInfo:
        """Update all paths, locale domains, etc."""

        # Generate i18n domains
        new_i18n_domain_list: list[LocaleDomainSpec] = []
        for old_domain_spec in old_package_info.i18n_domains:
            new_domain_spec = old_domain_spec.clone(
                domain_name=old_domain_spec.domain_name.replace(
                    old_package_info.single_module_name,
                    new_package_module_name,
                ),
            )
            new_i18n_domain_list.append(new_domain_spec)

        # Generate new path
        new_package_abs_path = old_package_info.abs_path.parent / new_package_module_name

        new_package_info = old_package_info.clone(
            package_reg_name=self.package_generator.generate_new_package_reg_name(new_package_module_name),
            module_names=(new_package_module_name,),
            test_dirs=(self.package_generator.generate_default_test_dir_name(new_package_module_name),),
            abs_path=new_package_abs_path,
            i18n_domains=tuple(new_i18n_domain_list),
        )
        return new_package_info

    def rename_package(self, old_package_module_name: str, new_package_module_name: str) -> PackageInfo:
        old_package_info = self.package_index.get_package_info_from_module_name(old_package_module_name)

        if new_package_module_name == old_package_module_name:
            return old_package_info

        new_package_info = self._generate_info_for_renamed_package(
            old_package_info,
            new_package_module_name=new_package_module_name,
        )

        # Update it
        self._update_package(old_package_info=old_package_info, new_package_info=new_package_info)
        return new_package_info

    def _update_package(self, old_package_info: PackageInfo, new_package_info: PackageInfo) -> PackageInfo:
        # Move the package dir
        if new_package_info.abs_path != old_package_info.abs_path:
            old_pkg_path = old_package_info.abs_path
            new_pkg_path = new_package_info.abs_path
            self.fs_editor.move_path(old_path=old_pkg_path, new_path=new_pkg_path)

        # Rename and replace all the internals and update imports
        if new_package_info.single_module_name != old_package_info.single_module_name:
            self._rename_package_internals(old_package_info=old_package_info, new_package_info=new_package_info)
            self._replace_imports(old_package_info.single_module_name, new_package_info.single_module_name)
            self._rename_in_dependent_meta_files(old_package_info=old_package_info, new_package_info=new_package_info)

        # Rename locale files to match the new domain names
        if new_package_info.i18n_domains != old_package_info.i18n_domains:
            self._rename_i18n_files(old_package_info=old_package_info, new_package_info=new_package_info)

        # Re-register package under new name
        self._re_register_package(old_package_info=old_package_info, new_package_info=new_package_info)
        return new_package_info

    def _init_package_from_boilerplate(self, boilerplate_info: PackageInfo, package_info: PackageInfo) -> None:
        # Create a copy of the boilerplate dir
        boilerplate_dir = boilerplate_info.abs_path
        new_pkg_dir = package_info.abs_path
        self.fs_editor.copy_path(src_path=boilerplate_dir, dst_path=new_pkg_dir)

        # Rename and replace all the internals
        self._rename_package_internals(old_package_info=boilerplate_info, new_package_info=package_info)

        # Register package
        self._register_package(package_info=package_info)

    def _make_regex_and_repl_for_sub(
        self,
        old_str: str,
        new_str: str,
        allow_dash: bool = True,
    ) -> tuple[re.Pattern, Callable[[re.Match], str]]:
        """
        Generate a tuple whose contains can be used as parameters for re.sub.
        Use this to replace module or package names in files
        """

        def _repl_mod_name(match: re.Match) -> str:
            return (
                match.string[match.start() : match.start("mod_name")]
                + new_str
                + match.string[match.end("mod_name") : match.end()]
            )

        regex: re.Pattern
        if allow_dash:
            assert "-" not in old_str and "-" not in new_str
            regex = re.compile(rf"(^|[^\w])(?P<mod_name>{re.escape(old_str)})($|[^\w])")
        else:
            regex = re.compile(rf"(^|[^\w\-])(?P<mod_name>{re.escape(old_str)})($|[^\w\-])")

        return regex, _repl_mod_name

    def _rename_i18n_files(self, old_package_info: PackageInfo, new_package_info: PackageInfo) -> None:
        """Rename the locale files in accordance with the new domain names"""
        assert len(new_package_info.i18n_domains) == len(old_package_info.i18n_domains)
        zipped_domain_pairs = zip(old_package_info.i18n_domains, new_package_info.i18n_domains)
        for old_domain_spec, new_domain_spec in zipped_domain_pairs:
            if new_domain_spec.domain_name == old_domain_spec.domain_name:
                continue

            # Update it in pyproject.toml
            regex, repl = self._make_regex_and_repl_for_sub(
                old_str=old_domain_spec.domain_name,
                new_str=new_domain_spec.domain_name,
                allow_dash=True,
            )
            self.fs_editor.replace_regex_in_file(regex=regex, repl=repl, file_path=new_package_info.toml_path)

            # Rename the files
            locales_path = new_package_info.abs_path / new_package_info.single_module_name / "locales"
            for locale_dir in locales_path.iterdir():
                old_path_base = locales_path / locale_dir.name / "LC_MESSAGES" / old_domain_spec.domain_name
                new_path_base = locales_path / locale_dir.name / "LC_MESSAGES" / new_domain_spec.domain_name
                for ext in ("po", "mo"):
                    old_path = Path(f"{str(old_path_base)}.{ext}")
                    new_path = Path(f"{str(new_path_base)}.{ext}")
                    self.fs_editor.move_path(old_path=old_path, new_path=new_path)

    def _rename_in_dependent_meta_files(self, old_package_info: PackageInfo, new_package_info: PackageInfo) -> None:
        meta_files = (
            "pyproject.toml",
            "README.md",
        )
        for other_package_info in self.package_index.list_package_infos():
            if other_package_info == old_package_info:
                continue
            regex, repl = self._make_regex_and_repl_for_sub(
                old_str=old_package_info.single_module_name,
                new_str=new_package_info.single_module_name,
                allow_dash=True,  # dash can be next to the module name
            )
            for meta_file_name in meta_files:
                meta_file_path = other_package_info.abs_path / meta_file_name
                if not meta_file_path.exists():
                    continue
                self.fs_editor.replace_regex_in_file(regex=regex, repl=repl, file_path=meta_file_path)

    def _rename_package_internals(self, old_package_info: PackageInfo, new_package_info: PackageInfo) -> None:
        new_pkg_dir = new_package_info.abs_path

        # Rename the code dir
        old_code_path = new_pkg_dir / old_package_info.single_module_name
        new_code_path = new_pkg_dir / new_package_info.single_module_name
        self.fs_editor.move_path(old_path=old_code_path, new_path=new_code_path)

        # Rename the tests dir
        for boilerplate_test_dir, package_test_dir in zip(old_package_info.test_dirs, new_package_info.test_dirs):
            old_tests_path = new_pkg_dir / boilerplate_test_dir
            new_tests_path = new_pkg_dir / package_test_dir
            self.fs_editor.move_path(old_path=old_tests_path, new_path=new_tests_path)

        # Masks for files that should not be edited
        edit_exclude_masks = self.repository_env.get_edit_exclude_masks()

        # Replace all package name occurrences with the given name
        all_zipped_modules = itertools.chain(
            zip(old_package_info.module_names, new_package_info.module_names),
            zip(old_package_info.test_dirs, new_package_info.test_dirs),
        )
        for boilerplate_mod_name, package_mod_name in all_zipped_modules:
            if boilerplate_mod_name == "tests" or package_mod_name == "tests":
                # FIXME: Do something
                continue
            regex, repl = self._make_regex_and_repl_for_sub(
                old_str=boilerplate_mod_name,
                new_str=package_mod_name,
                allow_dash=True,  # dash can be next to the module name
            )
            self.fs_editor.replace_regex_in_dir(
                regex=regex,
                repl=repl,
                path=new_pkg_dir,
                exclude_masks=edit_exclude_masks,
            )

        # add `-` to the regex
        regex, repl = self._make_regex_and_repl_for_sub(
            old_str=old_package_info.package_reg_name,
            new_str=new_package_info.package_reg_name,
            allow_dash=False,
        )
        self.fs_editor.replace_regex_in_dir(
            regex=regex,
            repl=repl,
            path=new_pkg_dir,
            exclude_masks=edit_exclude_masks,
        )

    def change_package_type(self, package_module_name: str, new_package_type: str) -> PackageInfo:
        old_package_info = self.package_index.get_package_info_from_module_name(package_module_name)

        root_pkg_dir = self.repository_env.get_root_package_dir(package_type=new_package_type)
        new_pkg_path = root_pkg_dir / old_package_info.abs_path.name

        new_package_info = old_package_info.clone(
            package_type=new_package_type,
            abs_path=new_pkg_path,
        )

        # Update it
        self._update_package(old_package_info=old_package_info, new_package_info=new_package_info)

        return new_package_info

    def rename_module(
        self,
        old_import_name: str,
        new_import_name: str,
        fix_imports: bool,
        move_files: bool,
    ) -> None:
        """Move code from one package to another or within one package"""

        if move_files:
            if self.repository_navigator.module_is_dir(old_import_name):
                old_path = self.repository_navigator.make_dir_module_path(old_import_name)
                new_path = self.repository_navigator.make_dir_module_path(new_import_name)
            else:
                old_path = self.repository_navigator.make_file_module_path(old_import_name)
                new_path = self.repository_navigator.make_file_module_path(new_import_name)

            self.fs_editor.move_path(old_path, new_path)

        if fix_imports:
            self._replace_imports(old_import_name, new_import_name)

    def get_imports(self, module_name: str) -> list[str]:
        """Move code from one package to another or within one package"""

        package_info = self.package_index.get_package_info_from_module_name(
            package_module_name=module_name,
        )
        return sorted(self.repository_navigator.collect_import_bases_from_module(package_info.single_module_name))

    def get_requirements(self, module_name: str) -> Mapping[str, RequirementList]:
        package_info = self.package_index.get_package_info_from_module_name(
            package_module_name=module_name,
        )
        return package_info.requirement_lists

    def compare_imports_and_requirements(
        self,
        package_module_name: str,
        ignore_prefix: Optional[str] = None,
        tests: bool = False,
    ) -> tuple[list[ReqPackageSpec], list[ReqPackageSpec]]:
        # TODO: Add support for namespace packages (e.g. google-api-core, google-auth)

        def _normalize_name(name: str) -> str:
            name = name.lower()
            if ignore_prefix is not None and name.startswith(ignore_prefix):
                name = name[len(ignore_prefix) :]
            return name

        package_info = self.package_index.get_package_info_from_module_name(
            package_module_name=package_module_name,
        )

        def _get_imports(scan_modules: Sequence[str]) -> dict[str, ReqPackageSpec]:
            _result: dict[str, ReqPackageSpec] = {}
            for module_name in scan_modules:
                for req_module_name in self.repository_navigator.collect_import_bases_from_module(module_name):
                    if req_module_name in scan_modules:
                        # Skip the package itself
                        continue
                    req_spec = self.package_reference.get_package_req_spec(req_module_name)
                    _result[_normalize_name(req_spec.package_name)] = req_spec
            return _result

        search_module_names = package_info.test_dirs if tests else package_info.module_names
        import_req_specs = _get_imports(search_module_names)
        if tests:
            main_pkg_import_req_specs = _get_imports(package_info.module_names)
            import_req_specs = {
                key: spec for key, spec in import_req_specs.items() if key not in main_pkg_import_req_specs
            }

        req_section_name: str
        if tests:
            req_section_name = "tool.poetry.group.tests.dependencies"
        else:
            req_section_name = "tool.poetry.dependencies"

        actual_req_specs: dict[str, ReqPackageSpec] = {
            _normalize_name(req_spec.package_name): req_spec
            for req_spec in package_info.requirement_lists.get(req_section_name, RequirementList()).req_specs
        }
        actual_req_specs = {
            req_name: req_spec
            for req_name, req_spec in actual_req_specs.items()
            if req_name not in package_info.implicit_deps
        }
        # TODO: Check that implicit deps are all listed in either the main or the test section

        special_excludes = {
            # Always ignore these imports/requirements
            "python",
            "",
            _normalize_name(package_info.package_reg_name),
        }

        extra_import_specs: list[ReqPackageSpec] = []
        extra_req_specs: list[ReqPackageSpec] = []
        all_package_names = sorted((set(import_req_specs) | set(actual_req_specs)) - special_excludes)
        for package_name in all_package_names:
            if package_name not in actual_req_specs:
                extra_import_specs.append(import_req_specs[package_name])
            if package_name not in import_req_specs:
                extra_req_specs.append(actual_req_specs[package_name])

        def _normalize_req_path(req: ReqPackageSpec) -> ReqPackageSpec:
            if isinstance(req, LocalReqPackageSpec):
                req_pkg_info = self.package_index.get_package_info_by_reg_name(req.package_name)
                req_path = Path(os.path.relpath(req_pkg_info.abs_path, package_info.abs_path))
                req = req.clone(path=req_path)
            return req

        extra_import_specs = [_normalize_req_path(req) for req in extra_import_specs]
        extra_req_specs = [_normalize_req_path(req) for req in extra_req_specs]

        return extra_import_specs, extra_req_specs

terrarium/dl_repmanager/dl_repmanager/repository_manager.py
datalens-tech/datalens-backend
import ast
import os
from pathlib import Path
import re
import sys
from typing import (
    ClassVar,
    Collection,
    Generator,
    Optional,
)

import attr

from dl_repmanager.package_index import PackageIndex
from dl_repmanager.primitives import (
    EntityReference,
    EntityReferenceType,
    ImportSpec,
    PackageInfo,
)
from dl_repmanager.repository_env import RepoEnvironment


@attr.s
class RepositoryNavigator:
    repository_env: RepoEnvironment = attr.ib(kw_only=True)
    package_index: PackageIndex = attr.ib(kw_only=True, factory=PackageIndex)

    _IMPORT_REGEX: ClassVar[re.Pattern] = re.compile(r"^\s*(import|from)\s+(?P<import_name>[\w.]+)")

    @staticmethod
    def recurse_code_files(dir_path: Path) -> Generator[Path, None, None]:
        assert dir_path is not None

        if dir_path.is_file() and str(dir_path).rsplit(".", 1)[1] == "py":
            yield dir_path

        for full_fn in dir_path.rglob("*.py"):
            yield full_fn

    def _file_contains_imports(self, file_path: Path, import_name: str) -> bool:
        import_specs = self._collect_imports_from_file(file_path)
        prefix = f"{import_name}."
        for import_spec in import_specs:
            if import_spec.import_module_name == import_name or import_spec.import_module_name.startswith(prefix):
                return True
        return False

    def recurse_files_with_import(self, import_name: str) -> Generator[Path, None, None]:
        for _package_type, pkg_type_root in self.repository_env.iter_package_abs_dirs():
            for file_path in self.recurse_code_files(pkg_type_root):
                if self._file_contains_imports(file_path, import_name):
                    yield file_path

    def _collect_imports_from_file(self, file_path: Path) -> list[ImportSpec]:
        result: list[ImportSpec] = []
        fs_editor = self.repository_env.get_fs_editor()
        with fs_editor.open(file_path, mode="r") as code_file:
            code_ast = ast.parse(code_file.read())
            for node in ast.walk(code_ast):
                found_modules: set[str]
                if isinstance(node, ast.Import):
                    found_modules = {alias.name for alias in node.names}
                elif isinstance(node, ast.ImportFrom):
                    module_name = node.module
                    if node.level > 0:  # relative import
                        module_name = "." * node.level + (module_name or "")
                    assert module_name is not None
                    found_modules = {module_name}
                else:
                    continue

                for found_module in found_modules:
                    result.append(
                        ImportSpec(
                            import_ast=node,
                            import_module_name=found_module,
                            source_path=file_path,
                        )
                    )

        return result

    def _filter_imports(
        self,
        import_specs: list[ImportSpec],
        mask: Optional[re.Pattern] = None,
        exclude_standard: bool = True,
    ) -> list[ImportSpec]:
        std_modules = set(sys.stdlib_module_names)
        result: list[ImportSpec] = []
        for import_spec in import_specs:
            if exclude_standard and import_spec.import_module_name.split(".")[0] in std_modules:
                continue
            if mask is not None and not mask.match(import_spec.import_module_name):
                continue

            result.append(import_spec)

        return result

    def _collect_imports_from_dir(self, dir_path: Path) -> list[ImportSpec]:
        result: list[ImportSpec] = []
        for file_path in self.recurse_code_files(dir_path):
            result += self._collect_imports_from_file(file_path)

        return result

    def _collect_imports_from_module(self, module_name: str) -> list[ImportSpec]:
        if self.module_is_dir(module_name):
            module_path = self.make_dir_module_path(module_name)
            return self._collect_imports_from_dir(module_path)
        else:
            module_path = self.make_file_module_path(module_name)
            return self._collect_imports_from_file(module_path)

    def collect_import_bases_from_module(
        self,
        module_name: str,
        mask: Optional[re.Pattern] = None,
        exclude_standard: bool = True,
    ) -> set[str]:
        def get_base_package_name(imported_name: str) -> str:
            return imported_name.split(".", 1)[0]

        import_specs = self._collect_imports_from_module(module_name)
        import_specs = self._filter_imports(import_specs, mask=mask, exclude_standard=exclude_standard)

        return {get_base_package_name(import_spec.import_module_name) for import_spec in import_specs}

    def make_file_module_path(self, module_name: str) -> Path:
        path = self.make_dir_module_path(module_name)
        return path.parent / (path.name + ".py")

    def make_dir_module_path(self, module_name: str) -> Path:
        package_module_name = module_name.split(".")[0]
        package_info: PackageInfo
        try:
            package_info = self.package_index.get_package_info_from_module_name(package_module_name=package_module_name)
        except KeyError:
            package_info = self.package_index.get_package_info_from_test_name(test_dir_name=package_module_name)

        package_root = package_info.abs_path
        assert (package_root / package_module_name).exists(), f"Package {package_module_name} does not exist"
        return package_root / module_name.replace(".", os.path.sep)

    def module_is_dir(self, module_name: str) -> bool:
        package_path = self.make_dir_module_path(module_name)
        if package_path.exists() and package_path.is_dir():
            return True

        file_path = self.make_file_module_path(module_name)
        assert file_path.is_file()
        return False

    def _make_module_path(self, module_name: str) -> Path:
        if self.module_is_dir(module_name):
            return self.make_dir_module_path(module_name)
        else:
            return self.make_file_module_path(module_name)

    def _naive_path_to_module(self, path: Path) -> str:
        path_str = str(path)
        assert not path_str.startswith("/") and path_str.endswith(".py")
        return path_str[:-3].replace("/", ".")

    def resolve_entity_to_type(
        self,
        entity_ref: EntityReference,
        to_ref_type: EntityReferenceType,
    ) -> set[EntityReference]:
        result: set[EntityReference]
        match entity_ref.ref_type:
            case EntityReferenceType.package_type:
                if to_ref_type == EntityReferenceType.package_type:
                    return {entity_ref}

                package_info_list = self.package_index.list_package_infos(package_type=entity_ref.name)
                package_entity_refs = {package_info.reg_entity_ref for package_info in package_info_list}
                if to_ref_type == EntityReferenceType.package_type:
                    return package_entity_refs
                else:
                    result = set()
                    for package_entity_ref in package_entity_refs:
                        result.update(
                            self.resolve_entity_to_type(
                                entity_ref=package_entity_ref,
                                to_ref_type=to_ref_type,
                            )
                        )
                    return result

            case EntityReferenceType.package_reg:
                if to_ref_type == EntityReferenceType.package_reg:
                    return {entity_ref}

                package_info = self.package_index.get_package_info_by_reg_name(entity_ref.name)
                if to_ref_type == EntityReferenceType.package_type:
                    return {package_info.package_type_entity_ref}

                if to_ref_type == EntityReferenceType.path:
                    result = set()
                    for module_name in (*package_info.module_names, *package_info.test_dirs):
                        result.update(
                            self.resolve_entity_to_type(
                                entity_ref=EntityReference(
                                    ref_type=EntityReferenceType.path,
                                    name=str(package_info.abs_path / module_name),
                                ),
                                to_ref_type=to_ref_type,
                            )
                        )
                    return result

                if to_ref_type == EntityReferenceType.module:
                    result = set()
                    for module_name in package_info.module_names:
                        module_path = package_info.abs_path / module_name
                        for module_sub_path in self.recurse_code_files(module_path):
                            rel_module_sub_path = module_sub_path.relative_to(package_info.abs_path)
                            result.add(
                                EntityReference(
                                    ref_type=EntityReferenceType.module,
                                    name=self._naive_path_to_module(rel_module_sub_path),
                                )
                            )
                    return result

            case EntityReferenceType.path:
                paths = list(self.recurse_code_files(Path(entity_ref.name)))
                if to_ref_type == EntityReferenceType.path:
                    return {EntityReference(ref_type=EntityReferenceType.path, name=str(path)) for path in paths}

                path_to_package_type_and_rel: dict[Path, tuple[str, Path, Path]] = {}
                for path in paths:
                    for package_type, pkg_type_dir in self.repository_env.iter_package_abs_dirs():
                        try:
                            rel_path = path.relative_to(pkg_type_dir)
                            path_to_package_type_and_rel[path] = (package_type, pkg_type_dir, rel_path)
                        except ValueError:
                            pass

                if to_ref_type == EntityReferenceType.package_type:
                    return {
                        EntityReference(ref_type=EntityReferenceType.package_type, name=package_type)
                        for package_type, pkg_type_dir, rel_path in path_to_package_type_and_rel.values()
                    }

                path_to_package_info: dict[Path, PackageInfo] = {}
                for path, (package_type, pkg_type_dir, rel_path) in path_to_package_type_and_rel.items():
                    package_abs_path = pkg_type_dir / rel_path.parts[0]
                    package_info = self.package_index.get_package_info_from_path(package_abs_path)
                    path_to_package_info[path] = package_info

                if to_ref_type == EntityReferenceType.module:
                    return {
                        EntityReference(
                            ref_type=EntityReferenceType.module,
                            name=self._naive_path_to_module(path.relative_to(package_info.abs_path)),
                        )
                        for path, package_info in path_to_package_info.items()
                    }

                if to_ref_type == EntityReferenceType.package_reg:
                    return {package_info.reg_entity_ref for package_info in path_to_package_info.values()}

            case EntityReferenceType.module:
                # Convert to path and resolve the rest from there
                return self.resolve_entity_to_type(
                    entity_ref=EntityReference(
                        ref_type=EntityReferenceType.path,
                        name=str(self._make_module_path(entity_ref.name)),
                    ),
                    to_ref_type=to_ref_type,
                )

            case _:
                pass

        raise ValueError(f"Unsupported conversion from: {entity_ref.ref_type.name} " f"to {to_ref_type.name}")

    def resolve_entities_to_type(
        self,
        entity_list: Collection[EntityReference],
        to_ref_type: EntityReferenceType,
    ) -> set[EntityReference]:
        result: set[EntityReference] = set()
        for entity_ref in entity_list:
            result.update(self.resolve_entity_to_type(entity_ref=entity_ref, to_ref_type=to_ref_type))
        return result

    def _collect_imports_from_entities(self, entity_list: Collection[EntityReference]) -> list[ImportSpec]:
        path_refs = self.resolve_entities_to_type(entity_list=entity_list, to_ref_type=EntityReferenceType.path)
        paths = sorted({Path(ref.name) for ref in path_refs})
        result: list[ImportSpec] = []
        for path in paths:
            result += self._collect_imports_from_file(path)
        return result

    def collect_imports_of_entities_from_entities(
        self,
        src_entity_list: Collection[EntityReference],
        imported_entity_list: Collection[EntityReference],
    ) -> list[ImportSpec]:
        import_specs = self._collect_imports_from_entities(entity_list=src_entity_list)

        module_refs = self.resolve_entities_to_type(
            entity_list=imported_entity_list, to_ref_type=EntityReferenceType.module
        )
        modules = {ref.name for ref in module_refs}
        module_bases = {module for module in modules if len(module.split(".")) == 1}

        result: list[ImportSpec] = []

        for import_spec in import_specs:
            import_base = import_spec.import_module_name.split(".")[0]
            # Fast check via base
            if import_base in module_bases:
                result.append(import_spec)
                continue
            if import_spec.import_module_name in modules:
                result.append(import_spec)

        return result

terrarium/dl_repmanager/dl_repmanager/repository_navigator.py
datalens-tech/datalens-backend
import logging
from pathlib import Path

import attr


log = logging.getLogger(__name__)


def _strip_version(raw: str):
    return raw.strip().split(" ;")[0]


@attr.s(frozen=True)
class PipRequirement:
    name: str = attr.ib()
    raw_version: str = attr.ib()

    @property
    def cleaned_version(self):
        cleaned_version = _strip_version(self.raw_version)
        if cleaned_version.startswith("=="):
            cleaned_version = cleaned_version[2:]

        return cleaned_version


@attr.s(frozen=True)
class PipRequirementsIO:
    path: Path = attr.ib()

    def read_existing(self, gather_all: bool = False) -> dict[str, PipRequirement]:
        result = dict()
        with open(self.path) as fh:
            for line in fh.readlines():
                try:
                    parts = line.strip().split("==")
                    name = parts[0].strip()
                    version = "==".join(parts[1:])
                    if " " in name or not name or name.startswith("#"):
                        # ignoring lines without actual requirements
                        continue
                    if gather_all or name not in ["mypy", "python"]:
                        result[name] = PipRequirement(name=name, raw_version=version)
                except Exception:
                    log.warning(f"Failed to parse {line} from file {self.path}")
        return result

    def write_updates(self, to_update: dict[str, PipRequirement]):
        requirements = self.read_existing(gather_all=True)
        requirements.update(to_update)
        with open(self.path, "wt") as fh:
            for name, requirement in sorted(requirements.items()):
                fh.write(f"{name}=={requirement.raw_version}\n")

terrarium/dl_repmanager/dl_repmanager/requirements_tools.py
datalens-tech/datalens-backend
from __future__ import annotations

import argparse
import logging
from pathlib import Path
import sys
from typing import cast

import attr

from dl_cli_tools.cli_base import CliToolBase
from dl_cli_tools.logging import setup_basic_logging
from dl_repmanager.exceptions import (
    InconsistentStateError,
    PackageMetaCliError,
)
from dl_repmanager.fs_editor import (
    DefaultFilesystemEditor,
    FilesystemEditor,
)
from dl_repmanager.mypy_stubs_sync import (
    RequirementsPathProvider,
    stubs_sync,
)
from dl_repmanager.package_meta_reader import (
    PackageMetaIOFactory,
    PackageMetaReader,
    PackageMetaWriter,
)
from dl_repmanager.repository_env import (
    DEFAULT_CONFIG_FILE_NAME,
    discover_config,
)


log = logging.getLogger(__name__)


def make_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="DL Package Meta CLI")
    parser.add_argument("--package-path", type=Path)

    package_subparsers = parser.add_subparsers(title="command", dest="command")
    add_package_commands(package_subparsers)
    return parser


def add_package_commands(package_subparsers: argparse._SubParsersAction) -> None:
    package_subparsers.add_parser("list-i18n-domains", help="List i18n domains and scan paths")
    package_subparsers.add_parser("list-external-dependencies", help="List external dependencies")

    sync_mypy_stubs = argparse.ArgumentParser(add_help=False)
    sync_mypy_stubs.add_argument("--dry-run", action="store_true")
    sync_mypy_stubs.add_argument("--config", default=DEFAULT_CONFIG_FILE_NAME)
    package_subparsers.add_parser(
        "sync-mypy-stubs",
        parents=[sync_mypy_stubs],
        help="Update or add types-* packages to third party dependencies",
    )

    src_path_parser = argparse.ArgumentParser(add_help=False)
    src_path_parser.add_argument("--src-path", type=Path, required=True, help="Source path")

    dst_path_parser = argparse.ArgumentParser(add_help=False)
    dst_path_parser.add_argument("--dst-path", type=Path, required=True, help="Destination path")

    package_subparsers.add_parser(
        "copy",
        parents=[src_path_parser, dst_path_parser],
        help="Copy file to/within package",
    )
    package_subparsers.add_parser("rm", parents=[src_path_parser], help="Remove file or dir in package")

    set_toml_meta_parser = argparse.ArgumentParser(add_help=False)
    set_toml_meta_parser.add_argument("--toml-section", required=True, help="Section name in pyproject.toml")
    set_toml_meta_parser.add_argument("--toml-key", required=True, help="Key within the section")
    set_toml_meta_parser.add_argument("--toml-value", required=True, help="Value for the key")

    package_subparsers.add_parser("set-meta-text", parents=[set_toml_meta_parser], help="Set text value in toml")
    package_subparsers.add_parser("set-meta-array", parents=[set_toml_meta_parser], help="Set array value in toml")


@attr.s
class DlPackageMetaTool(CliToolBase):
    fs_editor: FilesystemEditor = attr.ib(kw_only=True)
    package_path: Path = attr.ib(kw_only=True)
    meta_reader: PackageMetaReader = attr.ib(kw_only=True)
    meta_writer: PackageMetaWriter = attr.ib(kw_only=True)

    @classmethod
    def get_parser(cls) -> argparse.ArgumentParser:
        return make_parser()

    def list_i18n_domains(self) -> None:
        for domain_spec in sorted(self.meta_reader.get_i18n_domains(), key=lambda spec: spec.domain_name):
            paths = [str(path) for path in domain_spec.scan_paths]
            print(f'{domain_spec.domain_name}={";".join(paths)}')

    def sync_mypy_stubs(self, pkg_path: Path, dry_run: bool = True, config_name: str | None = None) -> None:
        base_path = discover_config(pkg_path, config_name or DEFAULT_CONFIG_FILE_NAME).parent
        path_provider = RequirementsPathProvider(base_path=base_path)
        stubs_sync(self.meta_reader, self.meta_writer, path_provider, dry_run)

    def copy_path(self, src_path: Path, dst_path: Path) -> None:
        if src_path.is_absolute():
            abs_src_path = src_path
        else:
            abs_src_path = self.package_path / src_path
        abs_dst_path = self.package_path / dst_path
        if self.fs_editor.path_exists(abs_dst_path):
            self.fs_editor.remove_path(abs_dst_path)
        self.fs_editor.copy_path(abs_src_path, abs_dst_path)

    def rm_path(self, src_path: Path) -> None:
        abs_src_path = self.package_path / src_path
        if self.fs_editor.path_exists(abs_src_path):
            self.fs_editor.remove_path(abs_src_path)

    def set_meta_text(self, toml_section: str, toml_key: str, toml_value: str) -> None:
        self.meta_writer.toml_writer.set_text_value(section_name=toml_section, key=toml_key, value=toml_value)

    def set_meta_array(self, toml_section: str, toml_key: str, toml_value: str) -> None:
        arr_value = toml_value.split(";")
        self.meta_writer.toml_writer.set_array_value(section_name=toml_section, key=toml_key, value=arr_value)

    @classmethod
    def run_parsed_args(cls, args: argparse.Namespace) -> None:
        package_path = cast(Path, args.package_path)
        fs_editor = DefaultFilesystemEditor(base_path=package_path)
        cls.run_for_package_path(
            fs_editor=fs_editor,
            package_path=package_path,
            package_command=args.command,
            args=args,
        )

    @classmethod
    def run_for_package_path(
        cls,
        fs_editor: FilesystemEditor,
        package_path: Path,
        package_command: str,
        args: argparse.Namespace,
    ) -> None:
        toml_path = package_path / "pyproject.toml"
        package_meta_io_factory = PackageMetaIOFactory(fs_editor=fs_editor)
        with (
            package_meta_io_factory.package_meta_reader(toml_path) as meta_reader,
            package_meta_io_factory.package_meta_writer(toml_path) as meta_writer,
        ):
            tool = cls(fs_editor=fs_editor, package_path=package_path, meta_reader=meta_reader, meta_writer=meta_writer)

            match package_command:
                case "list-i18n-domains":
                    tool.list_i18n_domains()
                case "sync-mypy-stubs":
                    tool.sync_mypy_stubs(Path(package_path), args.dry_run, args.config)
                case "copy":
                    tool.copy_path(src_path=args.src_path, dst_path=args.dst_path)
                case "rm":
                    tool.rm_path(src_path=args.src_path)
                case "set-meta-text":
                    tool.set_meta_text(
                        toml_section=args.toml_section, toml_key=args.toml_key, toml_value=args.toml_value
                    )
                case "set-meta-array":
                    tool.set_meta_array(
                        toml_section=args.toml_section, toml_key=args.toml_key, toml_value=args.toml_value
                    )
                case _:
                    raise RuntimeError(f"Got unknown command: {args.command}")


def main() -> None:
    setup_basic_logging()
    try:
        DlPackageMetaTool.run(sys.argv[1:])
    except InconsistentStateError:
        log.exception("Project inconsistent state discovered during cli command run.")
        sys.exit(1)
    except PackageMetaCliError:
        log.exception("Generic error occurred, exiting")
        sys.exit(10)


if __name__ == "__main__":
    main()

terrarium/dl_repmanager/dl_repmanager/scripts/package_meta_cli.py
datalens-tech/datalens-backend
from __future__ import annotations

import argparse
import difflib
import os
from pathlib import Path
from pprint import pprint
import sys
from typing import Optional

import attr

from dl_cli_tools.cli_base import CliToolBase
from dl_cli_tools.logging import setup_basic_logging
from dl_repmanager.git_manager import GitManager
from dl_repmanager.metapkg_manager import MetaPackageManager
from dl_repmanager.package_index import (
    PackageIndex,
    PackageIndexBuilder,
)
from dl_repmanager.package_reference import PackageReference
from dl_repmanager.primitives import (
    EntityReference,
    EntityReferenceType,
)
from dl_repmanager.project_editor import PyPrjEditor
from dl_repmanager.repository_env import (
    DEFAULT_CONFIG_FILE_NAME,
    RepoEnvironment,
    RepoEnvironmentLoader,
)
from dl_repmanager.repository_manager import (
    PackageGenerator,
    RepositoryManager,
)
from dl_repmanager.repository_navigator import RepositoryNavigator
from dl_repmanager.scripts.package_meta_cli import (
    DlPackageMetaTool,
    add_package_commands,
)


def _entity_ref_type_type(value: str | EntityReferenceType) -> EntityReferenceType:
    if isinstance(value, str):
        value = EntityReferenceType[value.replace("-", "_")]
    assert isinstance(value, EntityReferenceType)
    return value


def _entity_ref_type(value: str | EntityReference) -> EntityReference:
    if isinstance(value, str):
        ref_type_str, ref_name = value.split(":")
        value = EntityReference(ref_type=EntityReferenceType[ref_type_str], name=ref_name)
    assert isinstance(value, EntityReference)
    return value


def _entity_ref_list_type(value: str | list[EntityReference]) -> list[EntityReference]:
    if isinstance(value, str):
        value = [_entity_ref_type(part) for part in value.split(",")]
    assert isinstance(value, list)
    return value


_CWD = Path.cwd()


def make_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="DL Repository Management CLI")
    parser.add_argument("--config", help="Specify configuration file", default=DEFAULT_CONFIG_FILE_NAME)
    parser.add_argument("--base-path", type=Path, help="Base repository path", default=_CWD)
    parser.add_argument("--fs-editor", help="Override the FS editor type")
    parser.add_argument("--dry-run", action="store_true", help="Force usage of virtual FS editor")

    # mix-in parsers
    package_name_parser = argparse.ArgumentParser(add_help=False)
    package_name_parser.add_argument(
        "--package-name",
        required=True,
        help="Name of the package folder",
    )
    package_type_parser = argparse.ArgumentParser(add_help=False)
    package_type_parser.add_argument(
        "--package-type",
        required=True,
        help="Type of the package. See the config file (dl-repo.yml) for available package types",
    )
    old_new_import_names_parser = argparse.ArgumentParser(add_help=False)
    old_new_import_names_parser.add_argument("--old-import-name", required=True, help="Old import (module) name")
    old_new_import_names_parser.add_argument("--new-import-name", required=True, help="New import (module) name")

    # commands
    subparsers = parser.add_subparsers(title="command", dest="command")

    subparsers.add_parser(
        "init",
        parents=[package_name_parser, package_type_parser],
        help="Initialize new package",
    )

    copy_subparser = subparsers.add_parser(
        "copy",
        parents=[package_name_parser],
        help="Copy package",
    )
    copy_subparser.add_argument(
        "--from-package-name",
        required=True,
        help="Name of the package folder to copy from",
    )

    rename_parser = subparsers.add_parser(
        "rename",
        parents=[package_name_parser],
        help="Rename package",
    )
    rename_parser.add_argument(
        "--new-package-name",
        required=True,
        help="New name of the package",
    )

    move_code_parser = subparsers.add_parser(
        "rename-module",
        parents=[old_new_import_names_parser],
        help="Rename module (move code from one package to another or within one package)",
    )
    move_code_parser.add_argument(
        "--no-fix-imports",
        action="store_true",
        help="Don't update module name in import statements and other references",
    )
    move_code_parser.add_argument(
        "--no-move-files",
        action="store_true",
        help="Don't move the actual files",
    )

    subparsers.add_parser(
        "ch-package-type",
        parents=[package_name_parser, package_type_parser],
        help="Change package type (move it to another directory)",
    )

    package_list_parser = subparsers.add_parser(
        "package-list",
        parents=[package_type_parser],
        help="List all packages of a given type",
    )
    package_list_parser.add_argument("--mask", help="Formatting mask")
    package_list_parser.add_argument("--base-path", default=".", help="Base path for formatting relative package paths")

    subparsers.add_parser(
        "import-list",
        parents=[package_name_parser],
        help="List imports in a package",
    )

    subparsers.add_parser(
        "req-list",
        parents=[package_name_parser],
        help="List requirements of a package",
    )

    req_check_parser = subparsers.add_parser(
        "req-check",
        parents=[package_name_parser],
        help="Compare imports and requirements of a package",
    )
    req_check_parser.add_argument("--ignore-prefix", type=str, help="Package prefix to ignore when comparing")
    req_check_parser.add_argument("--tests", action="store_true", help="Check for tests, not the main package")

    metapackage_parser = argparse.ArgumentParser(add_help=False)
    metapackage_parser.add_argument("--metapackage", required=True, help="Metapackage name from repo config")

    subparsers.add_parser(
        "ensure-mypy-common",
        parents=[metapackage_parser],
        help="Checks and updates all sub projects with mypy config using template in the meta package",
    )

    compare_resulting_dependencies = subparsers.add_parser(
        "compare-resulting-deps",
        help=(
            "Compare resulting dependencies (including transitive) for Poetry dependency group."
            " Should be called from meta package dir."
            " Poetry lock for current packages will be called."
            " repository with specified revision (as compare base) will be checked out in ${DL_REPMGR_REPO_CLONE_PATH}."
        ),
    )
    compare_resulting_dependencies.add_argument(
        "--group",
        type=str,
        required=True,
        help="Poetry dependency group to check",
    )
    compare_resulting_dependencies.add_argument(
        "--base-rev", type=str, default="origin/trunk", help="git revision to compare with"
    )

    entities_parser = argparse.ArgumentParser(add_help=False)
    entities_parser.add_argument(
        "--entities",
        type=_entity_ref_list_type,
        required=True,
        help=(
            "Comma-separated list of repository entities which can be: "
            "package_type:<type_name>, package_reg:<package_reg_name> "
            "path:<path_to_py_file>, module:<module_name>"
        ),
    )

    resolve_parser = subparsers.add_parser(
        "resolve",
        parents=[entities_parser],
        help="Resolve repository entities from one type to another",
    )
    resolve_parser.add_argument("--to-type", type=_entity_ref_type_type, required=True, help="Output entity type")

    search_imports_parser = subparsers.add_parser(
        "search-imports",
        parents=[entities_parser],
        help="Search imports of of one set of entities in another",
    )
    search_imports_parser.add_argument(
        "--imported", type=_entity_ref_list_type, required=True, help="List of imported entites to search for"
    )

    recurse_packages_parser = subparsers.add_parser(
        "recurse-packages",
        parents=[entities_parser],
        help="Apply command to multiple packages",
    )
    recurse_packages_subparsers = recurse_packages_parser.add_subparsers(
        title="package_command",
        dest="recurse_packages_command",
    )
    add_package_commands(recurse_packages_subparsers)

    return parser


@attr.s
class DlRepManagerTool(CliToolBase):
    repository_env: RepoEnvironment = attr.ib(kw_only=True)
    package_index: PackageIndex = attr.ib(kw_only=True)
    repository_manager: RepositoryManager = attr.ib(kw_only=True)
    repository_navigator: RepositoryNavigator = attr.ib(kw_only=True)
    py_prj_editor: PyPrjEditor = attr.ib(kw_only=True)

    @classmethod
    def get_parser(cls) -> argparse.ArgumentParser:
        return make_parser()

    def init(self, package_name: str, package_type: str) -> None:
        self.repository_manager.init_package(package_module_name=package_name, package_type=package_type)

    def copy(self, package_name: str, from_package_name: str) -> None:
        self.repository_manager.copy_package(
            package_module_name=package_name, from_package_module_name=from_package_name
        )

    def rename(self, package_name: str, new_package_name: str) -> None:
        self.repository_manager.rename_package(
            old_package_module_name=package_name,
            new_package_module_name=new_package_name,
        )

    def ch_package_type(self, package_name: str, package_type: str) -> None:
        self.repository_manager.change_package_type(package_module_name=package_name, new_package_type=package_type)

    def package_list(self, package_type: str, mask: str, base_path: Path) -> None:
        for package_info in self.package_index.list_package_infos(package_type=package_type):
            printable_values = dict(
                package_type=package_info.package_type,
                abs_path=package_info.abs_path,
                rel_path=os.path.relpath(package_info.abs_path, base_path),
                single_module_name=package_info.single_module_name,
                package_reg_name=package_info.package_reg_name,
            )
            if mask:
                print(mask.format(**printable_values))
            else:
                pprint(printable_values)

    def rename_module(
        self,
        old_import_name: str,
        new_import_name: str,
        no_fix_imports: bool,
        no_move_files: bool,
    ) -> None:
        self.repository_manager.rename_module(
            old_import_name=old_import_name,
            new_import_name=new_import_name,
            fix_imports=not no_fix_imports,
            move_files=not no_move_files,
        )

    def import_list(self, package_name: str) -> None:
        """List imports in a package"""
        imports = self.repository_manager.get_imports(package_name)
        for imported_name in imports:
            print(imported_name)

    def req_list(self, package_name: str) -> None:
        """List requirements of a package"""
        req_lists = self.repository_manager.get_requirements(package_name)
        for req_list_name, req_list in sorted(req_lists.items()):
            print(req_list_name)
            for req_item in req_list.req_specs:
                print(f"    {req_item.pretty()}")

    def check_requirements(self, package_name: str, ignore_prefix: Optional[str] = None, tests: bool = False) -> None:
        """Compares imports and requirements of a package"""
        extra_import_specs, extra_req_specs = self.repository_manager.compare_imports_and_requirements(
            package_name,
            ignore_prefix=ignore_prefix,
            tests=tests,
        )

        if extra_import_specs:
            print("Missing requirements:")
            for req_item in extra_import_specs:
                print(f"    {req_item.as_req_str()}")
            print()
        if extra_req_specs:
            print("Extra requirements:")
            for req_item in extra_req_specs:
                print(f"    {req_item.as_req_str()}")
            print()
        if not extra_req_specs and not extra_import_specs:
            print("Requirements are in sync with imports")

    def ensure_mypy_common(self, metapackage_name: str):
        self.py_prj_editor.update_mypy_common(metapackage_name=metapackage_name)

    def compare_resulting_deps(self, base_revision: str, group: str):
        repo_root = self.py_prj_editor.base_path
        main_git_mgr = GitManager(repo_root)

        repo_clone_dest = Path(os.environ["DL_REPMGR_REPO_CLONE_PATH"])
        assert repo_clone_dest.parent.exists(), "Parent dir for repo clone does not exits"

        cmp_base_git_mgr = main_git_mgr.copy_repo(repo_clone_dest, fetch_if_exists=True)
        cmp_base_git_mgr.checkout(base_revision)

        metapkg_rel_path = Path.cwd().relative_to(repo_root)

        current_metapkg_mgr = MetaPackageManager(Path.cwd())

        current_metapkg_mgr.run_poetry_lock(suppress_stdout=True)

        cmp_base_metapkg_mgr = MetaPackageManager(repo_clone_dest / metapkg_rel_path)

        current_deps = current_metapkg_mgr.export_dependencies(group)
        cmp_base_deps = cmp_base_metapkg_mgr.export_dependencies(group)

        diff_lines = difflib.unified_diff(
            [dep.pretty() for dep in cmp_base_deps],
            [dep.pretty() for dep in current_deps],
            fromfile="cmp_base.txt",
            tofile="actual.txt",
        )
        for line in diff_lines:
            print(line, end=None)

    def resolve_entities(self, entities: list[EntityReference], to_ref_type: EntityReferenceType) -> None:
        resolved_entities = self.repository_navigator.resolve_entities_to_type(entities, to_ref_type=to_ref_type)
        for entity in sorted(resolved_entities, key=lambda item: item.name):
            print(entity.name)

    def search_imports(self, entities: list[EntityReference], imported: list[EntityReference]) -> None:
        import_specs = self.repository_navigator.collect_imports_of_entities_from_entities(
            src_entity_list=entities, imported_entity_list=imported
        )

        base_path_len = len(str(self.repository_env.base_path)) + 1  # +1 for the slash
        formatted_imports = [
            (
                f"{str(import_spec.source_path)[base_path_len:]}:"
                f"{import_spec.import_ast.lineno}: "
                f"{import_spec.import_module_name}"
            )
            for import_spec in sorted(import_specs, key=lambda item: item.source_path)
        ]
        print("\n".join(sorted(formatted_imports)))

    def recurse_packages(
        self,
        entities: list[EntityReference],
        args: argparse.Namespace,
    ) -> None:
        for entity_ref in entities:
            assert entity_ref.ref_type in (
                EntityReferenceType.package_reg,
                EntityReferenceType.package_type,
            ), "Only `package_reg` and `package_type` supported here"

        resolved_entities = self.repository_navigator.resolve_entities_to_type(
            entities,
            to_ref_type=EntityReferenceType.package_reg,
        )
        package_paths = [
            self.package_index.get_package_info_by_reg_name(entity_ref.name).abs_path
            for entity_ref in resolved_entities
        ]
        for package_path in package_paths:
            print(f"Processing {package_path}")
            DlPackageMetaTool.run_for_package_path(
                fs_editor=self.repository_env.get_fs_editor(),
                package_path=package_path,
                package_command=args.recurse_packages_command,
                args=args,
            )

    @classmethod
    def initialize(cls, base_path: Path, config_file_name: str, fs_editor_type: str) -> DlRepManagerTool:
        repository_env = RepoEnvironmentLoader(
            config_file_name=config_file_name,
            override_fs_editor_type=fs_editor_type,
        ).load_env(base_path=base_path)

        index_builder = PackageIndexBuilder(repository_env=repository_env)
        package_index = index_builder.build_index()

        repository_navigator = RepositoryNavigator(repository_env=repository_env, package_index=package_index)
        tool = cls(
            repository_env=repository_env,
            package_index=package_index,
            repository_navigator=repository_navigator,
            repository_manager=RepositoryManager(
                repository_env=repository_env,
                package_index=package_index,
                repository_navigator=repository_navigator,
                package_generator=PackageGenerator(package_index=package_index, repository_env=repository_env),
                package_reference=PackageReference(package_index=package_index, repository_env=repository_env),
            ),
            py_prj_editor=PyPrjEditor(
                repository_env=repository_env,
                package_index=package_index,
            ),
        )
        return tool

    @classmethod
    def run_parsed_args(cls, args: argparse.Namespace) -> None:
        config_file_name = args.config
        base_path = args.base_path
        fs_editor_type = args.fs_editor if not args.dry_run else "virtual"
        tool = cls.initialize(base_path=base_path, config_file_name=config_file_name, fs_editor_type=fs_editor_type)

        match args.command:
            case "init":
                tool.init(package_name=args.package_name, package_type=args.package_type)
            case "copy":
                tool.copy(package_name=args.package_name, from_package_name=args.from_package_name)
            case "rename":
                tool.rename(package_name=args.package_name, new_package_name=args.new_package_name)
            case "ch-package-type":
                tool.ch_package_type(package_name=args.package_name, package_type=args.package_type)
            case "package-list":
                tool.package_list(package_type=args.package_type, mask=args.mask, base_path=Path(args.base_path))
            case "rename-module":
                tool.rename_module(
                    old_import_name=args.old_import_name,
                    new_import_name=args.new_import_name,
                    no_fix_imports=args.no_fix_imports,
                    no_move_files=args.no_move_files,
                )
            case "import-list":
                tool.import_list(package_name=args.package_name)
            case "req-list":
                tool.req_list(package_name=args.package_name)
            case "req-check":
                tool.check_requirements(
                    package_name=args.package_name, ignore_prefix=args.ignore_prefix, tests=args.tests
                )
            case "ensure-mypy-common":
                tool.ensure_mypy_common(metapackage_name=args.metapackage)
            case "compare-resulting-deps":
                tool.compare_resulting_deps(base_revision=args.base_rev, group=args.group)
            case "resolve":
                tool.resolve_entities(entities=args.entities, to_ref_type=args.to_type)
            case "search-imports":
                tool.search_imports(entities=args.entities, imported=args.imported)
            case "recurse-packages":
                tool.recurse_packages(entities=args.entities, args=args)
            case _:
                raise RuntimeError(f"Got unknown command: {args.command}")


def main() -> None:
    setup_basic_logging()
    DlRepManagerTool.run(sys.argv[1:])


if __name__ == "__main__":
    main()

terrarium/dl_repmanager/dl_repmanager/scripts/repmanager_cli.py
datalens-tech/datalens-backend
import argparse
from pathlib import Path

from dl_cli_tools.logging import setup_basic_logging
from dl_repmanager.metapkg_scoped_sync import sync_scoped_metapkg


def make_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="DL Package Meta CLI")
    parser.add_argument("--orig-metapkg-path", required=True)
    parser.add_argument("--scoped-metapkg-path", required=True)
    parser.add_argument("--lib-no-prune", action="append", default=[])
    parser.add_argument("--remove-private-pypi", action="store_true", default=False)
    parser.add_argument("--use-target-lock", action="store_true", default=False)
    parser.add_argument(
        "lib_paths_to_include",
        nargs="+",
        help="Path of libraries to include relative to scoped metapkg",
    )

    return parser


def main() -> None:
    setup_basic_logging()
    parser = make_parser()

    args = parser.parse_args()

    sync_scoped_metapkg(
        original_metapkg_path=Path(args.orig_metapkg_path).absolute(),
        scoped_metapkg_path=Path(args.scoped_metapkg_path).absolute(),
        use_target_lock=args.use_target_lock,
        prevent_prune_for_deps=args.lib_no_prune,
        package_dirs_to_include=[Path(path_str) for path_str in args.lib_paths_to_include],
        remove_private_pypi=args.remove_private_pypi,
    )


if __name__ == "__main__":
    main()

terrarium/dl_repmanager/dl_repmanager/scripts/scoped_metapkg_sync.py
datalens-tech/datalens-backend
from __future__ import annotations

import contextlib
from pathlib import Path
from typing import (
    Any,
    Generator,
    Iterable,
    Sequence,
)

import attr
import tomlkit
from tomlkit.container import Container as TOMLContainer
from tomlkit.container import OutOfOrderTableProxy
from tomlkit.exceptions import NonExistentKey
from tomlkit.items import AbstractTable
from tomlkit.items import Item as TOMLItem
from tomlkit.toml_document import TOMLDocument

from dl_repmanager.fs_editor import FilesystemEditor


@attr.s
class TOMLReaderBase:
    _toml: TOMLDocument = attr.ib(kw_only=True)

    def get_section(self, key: str, strict: bool = True) -> dict:
        section: dict = self._toml
        item: TOMLContainer | TOMLItem | OutOfOrderTableProxy
        for part in key.split("."):
            try:
                item = section[part]
            except KeyError:
                if strict:
                    raise
                item = tomlkit.table()

            assert isinstance(item, dict)
            section = item

        assert isinstance(section, dict)
        return section

    def iter_section_items(self, key: str, strict: bool = True) -> Iterable[tuple[str, Any]]:
        section = self.get_section(key=key, strict=strict)
        assert isinstance(section, AbstractTable)
        for item_key, item in section.items():
            yield str(item_key), item


class TOMLReader(TOMLReaderBase):
    """Generic pyproject.toml reader"""


class TOMLWriter(TOMLReaderBase):
    """Generic pyproject.toml writer"""

    def get_editable_section(self, key: str) -> AbstractTable:
        # None of this makes any sense. This is just the way the tomlkit library is.
        # __getitem__ might return an `OutOfOrderTableProxy` or an `AbstractTable` - you never know.
        # `AbstractTable` does have an `add` method while `OutOfOrderTableProxy` doesn't.
        section = self.get_section(key)
        assert isinstance(section, AbstractTable), f"Expected AbstractTable, got {type(section)}"
        return section

    def add_section(self, key: str) -> AbstractTable:
        section: dict = self._toml
        item: TOMLItem | TOMLContainer | OutOfOrderTableProxy
        for part in key.split("."):
            try:
                item = section[part]
            except KeyError:
                section[part] = tomlkit.table()
                item = section[part]

            assert isinstance(item, dict)
            section = item

        assert isinstance(section, AbstractTable)
        return section

    def delete_section(self, key: str) -> None:
        section: dict = self._toml
        parts = key.split(".")
        for part in parts[:-1]:
            try:
                item = section[part]
            except KeyError:
                # Parent section doesn't exist
                return

            assert isinstance(item, dict)
            section = item

        del section[parts[-1]]

    @classmethod
    @contextlib.contextmanager
    def suppress_non_existent_key(cls) -> Generator[None, None, None]:
        try:
            yield
        except NonExistentKey:
            pass

    def set_text_value(self, section_name: str, key: str, value: str) -> None:
        section = self.get_section(section_name)
        section[key] = value

    def set_array_value(self, section_name: str, key: str, value: Sequence[str | dict]) -> None:
        section = self.get_section(section_name)
        if not all(isinstance(item, (str, dict)) for item in value):
            raise TypeError(f"Invalid item type in value {value}")
        array = tomlkit.array()
        array.extend(value)
        section[key] = array


@attr.s
class TOMLIOFactory:
    fs_editor: FilesystemEditor = attr.ib(kw_only=True)

    @contextlib.contextmanager
    def toml_reader(self, file_path: Path) -> Generator[TOMLReader, None, None]:
        with self.fs_editor.open(file_path, mode="r") as f:
            yield TOMLReader(toml=tomlkit.load(f))

    @contextlib.contextmanager
    def toml_writer(self, file_path: Path) -> Generator[TOMLWriter, None, None]:
        with self.fs_editor.open(file_path, mode="r+") as f:
            toml = tomlkit.load(f)
            yield TOMLWriter(toml=toml)
            f.seek(0)
            f.truncate(0)
            f.write(tomlkit.dumps(toml))

terrarium/dl_repmanager/dl_repmanager/toml_tools.py
datalens-tech/datalens-backend
from __future__ import annotations

from contextlib import contextmanager
from pathlib import Path
import shutil
import tempfile
from typing import (
    Generator,
    Optional,
)
import uuid

import attr
import pytest

from dl_repmanager.fs_editor import FilesystemEditor
from dl_repmanager.package_index import (
    PackageIndex,
    PackageIndexBuilder,
)
from dl_repmanager.package_reference import PackageReference
from dl_repmanager.primitives import PackageInfo
from dl_repmanager.repository_env import (
    RepoEnvironment,
    RepoEnvironmentLoader,
)
from dl_repmanager.repository_manager import (
    PackageGenerator,
    RepositoryManager,
)
from dl_repmanager.repository_navigator import RepositoryNavigator
from dl_repmanager.toml_tools import (
    TOMLIOFactory,
    TOMLWriter,
)
from dl_repmanager_tests.unit.config import (
    DEFAULT_PACKAGE_TYPE,
    REPO_PATH,
)
from dl_repmanager_tests.unit.tool_runner import (
    PackageCliRunner,
    RepoCliRunner,
)


@attr.s
class Repo:
    """A container/initializer for a bunch of objects"""

    base_path: Path = attr.ib(kw_only=True)
    repository_env: RepoEnvironment = attr.ib(init=False)
    fs_editor: FilesystemEditor = attr.ib(init=False)
    package_index: PackageIndex = attr.ib(init=False)
    repository_navigator: RepositoryNavigator = attr.ib(init=False)
    package_generator: PackageGenerator = attr.ib(init=False)
    package_reference: PackageReference = attr.ib(init=False)
    repository_manager: RepositoryManager = attr.ib(init=False)

    @repository_env.default
    def _make_repository_env(self) -> RepoEnvironment:
        return RepoEnvironmentLoader().load_env(base_path=self.base_path)

    @fs_editor.default
    def _make_fs_editor(self) -> FilesystemEditor:
        return self.repository_env.get_fs_editor()

    @package_index.default
    def _make_package_index(self) -> PackageIndex:
        return PackageIndexBuilder(repository_env=self.repository_env).build_index()

    @repository_navigator.default
    def _make_repository_navigator(self) -> RepositoryNavigator:
        return RepositoryNavigator(repository_env=self.repository_env, package_index=self.package_index)

    @repository_manager.default
    def _make_repository_manager(self) -> RepositoryManager:
        return RepositoryManager(
            repository_env=self.repository_env,
            package_index=self.package_index,
            repository_navigator=self.repository_navigator,
            package_generator=PackageGenerator(repository_env=self.repository_env, package_index=self.package_index),
            package_reference=PackageReference(repository_env=self.repository_env, package_index=self.package_index),
        )

    def reload(self) -> Repo:
        return self.__class__(base_path=self.base_path)


@contextmanager
def deployed_repo(repo_src: Path) -> Generator[Path, None, None]:
    repo_dst = Path(tempfile.mkdtemp()) / repo_src.name
    shutil.copytree(repo_src, repo_dst)
    try:
        yield repo_dst
    finally:
        shutil.rmtree(repo_dst)


class RepmanagerTestingBase:
    @pytest.fixture(scope="session")
    def repo_path(self) -> Generator[Path, None, None]:
        with deployed_repo(REPO_PATH) as repo_path:
            yield repo_path

    @pytest.fixture(scope="class")
    def repo(self, repo_path: Path) -> Repo:
        return Repo(base_path=repo_path)

    def generate_package_name(self) -> str:
        return f"pkg_{uuid.uuid4().hex[:6]}"

    def init_package(
        self,
        repo: Repo,
        package_type: str = DEFAULT_PACKAGE_TYPE,
        package_name: Optional[str] = None,
    ) -> PackageInfo:
        package_name = package_name or self.generate_package_name()
        assert package_name is not None
        repo.repository_manager.init_package(package_type=package_type, package_module_name=package_name)
        repo = repo.reload()
        return repo.package_index.get_package_info_from_module_name(package_name)

    @pytest.fixture(scope="class")
    def package(self, repo: Repo) -> PackageInfo:
        return self.init_package(repo=repo)

    @contextmanager
    def toml_writer(self, repo: Repo, toml_path: Path) -> Generator[TOMLWriter, None, None]:
        toml_io_factory = TOMLIOFactory(fs_editor=repo.fs_editor)
        with toml_io_factory.toml_writer(toml_path) as toml_writer:
            yield toml_writer


class RepmanagerCliTestBase(RepmanagerTestingBase):
    def get_repo_cli(self, repo_path: Path) -> RepoCliRunner:
        return RepoCliRunner(repo_path=repo_path)

    def get_package_cli(self, package_path: Path) -> PackageCliRunner:
        return PackageCliRunner(package_path=package_path)

    @pytest.fixture(scope="class")
    def repo_cli(self, repo_path: Path) -> RepoCliRunner:
        return self.get_repo_cli(repo_path=repo_path)

terrarium/dl_repmanager/dl_repmanager_tests/unit/base.py
datalens-tech/datalens-backend
import pytest

from dl_repmanager_tests.unit.base import RepmanagerCliTestBase
from dl_repmanager_tests.unit.tool_runner import PackageCliRunner


class TestPackageCli(RepmanagerCliTestBase):
    @pytest.fixture(scope="class")
    def package_cli(self, package) -> PackageCliRunner:
        return self.get_package_cli(package_path=package.abs_path)

    def test_list_i18n_domains(self, repo, package, package_cli) -> None:
        with self.toml_writer(repo=repo, toml_path=package.toml_path) as toml_writer:
            section_name = "datalens.i18n.domains"
            toml_writer.add_section(section_name)
            toml_writer.set_array_value(section_name=section_name, key="first_domain", value=[{"path": "first_path"}])

        cli_result = package_cli.run_with_args(["list-i18n-domains"])
        assert "first_domain=first_path" in cli_result.stdout

terrarium/dl_repmanager/dl_repmanager_tests/unit/cli/test_package_cli.py
datalens-tech/datalens-backend
from dl_repmanager_tests.unit.base import RepmanagerCliTestBase


class TestRepoCli(RepmanagerCliTestBase):
    def test_init(self, repo, repo_cli) -> None:
        package_name = self.generate_package_name()
        repo_cli.run_with_args(["init", "--package-type", "lib", "--package-name", package_name])
        repo = repo.reload()
        package_info = repo.package_index.get_package_info_from_module_name(package_name)
        assert package_info.package_type == "lib"
        assert package_info.single_module_name == package_name
        assert package_info.single_test_dir == f"{package_name}_tests"
        assert package_info.package_reg_name == package_name.replace("_", "-")

terrarium/dl_repmanager/dl_repmanager_tests/unit/cli/test_repo_cli.py
datalens-tech/datalens-backend
from pathlib import Path


REPO_PATH = Path(__file__).parent.parent / "testrepo"
DEFAULT_PACKAGE_TYPE = "lib"

terrarium/dl_repmanager/dl_repmanager_tests/unit/config.py
datalens-tech/datalens-backend
from pathlib import Path
import shutil
import tempfile
from typing import Generator

import attr
import pytest

from dl_repmanager.fs_editor import (
    DefaultFilesystemEditor,
    FilesystemEditor,
    VirtualFilesystemEditor,
)


@attr.s
class FSSetup:
    base_path: Path = attr.ib(kw_only=True)


class FSEditorTestCase:
    @pytest.fixture(scope="class")
    def base_setup(self) -> Generator[FSSetup, None, None]:
        base_path = Path(tempfile.mkdtemp())
        try:
            yield FSSetup(base_path=base_path)
        except Exception:
            shutil.rmtree(base_path)

    @pytest.fixture(scope="class")
    def fs_editor(self, base_setup: FSSetup) -> FilesystemEditor:
        raise NotImplementedError

    def test_file_access(self, base_setup: FSSetup, fs_editor: FilesystemEditor) -> None:
        testfile_path = base_setup.base_path / "testfile.thing"
        print(base_setup.base_path)
        with fs_editor.open(testfile_path, mode="w") as testfile:
            testfile.write("qwerty")

        with fs_editor.open(testfile_path, mode="r") as testfile:
            text = testfile.read()

        assert text == "qwerty"


class TestVirtualFSEditor(FSEditorTestCase):
    @pytest.fixture(scope="class")
    def fs_editor(self, base_setup: FSSetup) -> FilesystemEditor:
        return VirtualFilesystemEditor(base_path=base_setup.base_path)


class TestDefaultFSEditor(FSEditorTestCase):
    @pytest.fixture(scope="class")
    def fs_editor(self, base_setup: FSSetup) -> FilesystemEditor:
        return DefaultFilesystemEditor(base_path=base_setup.base_path)

terrarium/dl_repmanager/dl_repmanager_tests/unit/test_fs_editor.py
datalens-tech/datalens-backend
from pathlib import Path

import attr

from dl_cli_tools.testing.tool_runner import (
    CliResult,
    CliRunner,
)
from dl_repmanager.scripts.package_meta_cli import DlPackageMetaTool
from dl_repmanager.scripts.repmanager_cli import DlRepManagerTool


@attr.s
class RepoCliRunner(CliRunner):
    cli_cls = DlRepManagerTool

    repo_path: Path = attr.ib()

    def run_with_args(self, argv: list[str]) -> CliResult:
        argv = ["--base-path", str(self.repo_path)] + argv
        return super().run_with_args(argv)


@attr.s
class PackageCliRunner(CliRunner):
    cli_cls = DlPackageMetaTool

    package_path: Path = attr.ib()

    def run_with_args(self, argv: list[str]) -> CliResult:
        argv = ["--package-path", str(self.package_path)] + argv
        return super().run_with_args(argv)

terrarium/dl_repmanager/dl_repmanager_tests/unit/tool_runner.py
datalens-tech/datalens-backend