diff --git a/lib/dl_connector_ydb/dl_connector_ydb/__init__.py b/lib/dl_connector_ydb/dl_connector_ydb/__init__.py index 0bfdfc461d..e69de29bb2 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/__init__.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/__init__.py @@ -1,6 +0,0 @@ -try: - from ydb_proto_stubs_import import init_ydb_stubs - - init_ydb_stubs() -except ImportError: - pass # stubs will be initialized from the ydb package diff --git a/lib/dl_connector_ydb/dl_connector_ydb/core/base/adapter.py b/lib/dl_connector_ydb/dl_connector_ydb/core/base/adapter.py index 194b29de5d..c8111aa953 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/core/base/adapter.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/core/base/adapter.py @@ -13,10 +13,12 @@ import attr import sqlalchemy as sa +import ydb_sqlalchemy as ydb_sa from dl_core import exc from dl_core.connection_executors.adapters.adapters_base_sa_classic import BaseClassicAdapter from dl_core.connection_models import TableIdent +import dl_sqlalchemy_ydb.dialect if TYPE_CHECKING: @@ -44,14 +46,14 @@ def _is_table_exists(self, table_ident: TableIdent) -> bool: _type_code_to_sa = { None: sa.TEXT, # fallback - "Int8": sa.INTEGER, - "Int16": sa.INTEGER, - "Int32": sa.INTEGER, - "Int64": sa.INTEGER, - "Uint8": sa.INTEGER, - "Uint16": sa.INTEGER, - "Uint32": sa.INTEGER, - "Uint64": sa.INTEGER, + "Int8": ydb_sa.types.Int8, + "Int16": ydb_sa.types.Int16, + "Int32": ydb_sa.types.Int32, + "Int64": ydb_sa.types.Int64, + "Uint8": ydb_sa.types.UInt8, + "Uint16": ydb_sa.types.UInt16, + "Uint32": ydb_sa.types.UInt32, + "Uint64": ydb_sa.types.UInt64, "Float": sa.FLOAT, "Double": sa.FLOAT, "String": sa.TEXT, @@ -60,9 +62,9 @@ def _is_table_exists(self, table_ident: TableIdent) -> bool: "Yson": sa.TEXT, "Uuid": sa.TEXT, "Date": sa.DATE, - "Datetime": sa.DATETIME, - "Timestamp": sa.DATETIME, - "Interval": sa.INTEGER, + "Timestamp": dl_sqlalchemy_ydb.dialect.YqlTimestamp, + "Datetime": dl_sqlalchemy_ydb.dialect.YqlDateTime, + "Interval": dl_sqlalchemy_ydb.dialect.YqlInterval, "Bool": sa.BOOLEAN, } _type_code_to_sa = { @@ -94,7 +96,17 @@ def _convert_bytes(value: bytes) -> str: return value.decode("utf-8", errors="replace") @staticmethod - def _convert_ts(value: int) -> datetime.datetime: + def _convert_interval(value: datetime.timedelta | int) -> int: + if value is None: + return None + if isinstance(value, datetime.timedelta): + return int(value.total_seconds() * 1_000_000) + return value + + @staticmethod + def _convert_ts(value: int | datetime.datetime) -> datetime.datetime: + if isinstance(value, datetime.datetime): + return value.replace(tzinfo=datetime.timezone.utc) return datetime.datetime.utcfromtimestamp(value / 1e6).replace(tzinfo=datetime.timezone.utc) def _get_row_converters(self, cursor_info: ExecutionStepCursorInfo) -> tuple[Optional[Callable[[Any], Any]], ...]: @@ -104,6 +116,8 @@ def _get_row_converters(self, cursor_info: ExecutionStepCursorInfo) -> tuple[Opt if type_name_norm == "string" else self._convert_ts if type_name_norm == "timestamp" + else self._convert_interval + if type_name_norm == "interval" else None for type_name_norm in type_names_norm ) @@ -122,3 +136,6 @@ def make_exc( # TODO: Move to ErrorTransformer kw["db_message"] = kw.get("db_message") or message return exc_cls, kw + + def get_engine_kwargs(self) -> dict: + return {} diff --git a/lib/dl_connector_ydb/dl_connector_ydb/core/base/type_transformer.py b/lib/dl_connector_ydb/dl_connector_ydb/core/base/type_transformer.py index 2a600c5da1..d492a3e007 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/core/base/type_transformer.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/core/base/type_transformer.py @@ -3,9 +3,10 @@ from typing import TYPE_CHECKING import sqlalchemy as sa -import ydb.sqlalchemy as ydb_sa +import ydb_sqlalchemy.sqlalchemy as ydb_sa from dl_constants.enums import UserDataType +import dl_sqlalchemy_ydb.dialect from dl_type_transformer.type_transformer import ( TypeTransformer, make_native_type, @@ -20,12 +21,16 @@ class YQLTypeTransformer(TypeTransformer): _base_type_map: dict[UserDataType, tuple[SATypeSpec, ...]] = { # Note: first SA type is used as the default. UserDataType.integer: ( - sa.BIGINT, - sa.SMALLINT, sa.INTEGER, + ydb_sa.types.Int8, + ydb_sa.types.Int16, + ydb_sa.types.Int32, + ydb_sa.types.Int64, + ydb_sa.types.UInt8, + ydb_sa.types.UInt16, ydb_sa.types.UInt32, ydb_sa.types.UInt64, - ydb_sa.types.UInt8, + dl_sqlalchemy_ydb.dialect.YqlInterval, ), UserDataType.float: ( sa.FLOAT, @@ -36,8 +41,11 @@ class YQLTypeTransformer(TypeTransformer): UserDataType.boolean: (sa.BOOLEAN,), UserDataType.string: ( sa.TEXT, + sa.String, sa.CHAR, sa.VARCHAR, + sa.BINARY, + # TODO: ydb_sa.types.YqlJSON, # see also: ENUM, ), # see also: UUID @@ -45,10 +53,14 @@ class YQLTypeTransformer(TypeTransformer): UserDataType.datetime: ( sa.DATETIME, sa.TIMESTAMP, + dl_sqlalchemy_ydb.dialect.YqlDateTime, + dl_sqlalchemy_ydb.dialect.YqlTimestamp, ), UserDataType.genericdatetime: ( sa.DATETIME, sa.TIMESTAMP, + dl_sqlalchemy_ydb.dialect.YqlDateTime, + dl_sqlalchemy_ydb.dialect.YqlTimestamp, ), UserDataType.unsupported: (sa.sql.sqltypes.NullType,), # Actually the default, so should not matter much. } diff --git a/lib/dl_connector_ydb/dl_connector_ydb/core/ydb/adapter.py b/lib/dl_connector_ydb/dl_connector_ydb/core/ydb/adapter.py index 06368ed334..33216a2eb3 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/core/ydb/adapter.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/core/ydb/adapter.py @@ -12,9 +12,9 @@ import attr import grpc from ydb import DriverConfig -import ydb.dbapi as ydb_dbapi from ydb.driver import credentials_impl import ydb.issues as ydb_cli_err +import ydb_dbapi from dl_configs.utils import get_root_certificates from dl_constants.enums import ConnectionType @@ -68,7 +68,9 @@ def _update_connect_args(self, args: dict) -> None: ) args.update( credentials=credentials_impl.StaticCredentials( - driver_config=driver_config, user=self._target_dto.username, password=self._target_dto.password + driver_config=driver_config, + user=self._target_dto.username, + password=self._target_dto.password, ) ) else: @@ -77,11 +79,9 @@ def _update_connect_args(self, args: dict) -> None: def get_connect_args(self) -> dict: target_dto = self._target_dto args = dict( - endpoint="{}://{}:{}".format( - "grpcs" if self._target_dto.ssl_enable else "grpc", - target_dto.host, - target_dto.port, - ), + host=self._target_dto.host, + port=self._target_dto.port, + protocol="grpcs" if self._target_dto.ssl_enable else "grpc", database=target_dto.db_name, root_certificates=self._get_ssl_ca(), ) @@ -96,7 +96,7 @@ def _list_table_names_i(self, db_name: str, show_dot: bool = False) -> Iterable[ connection = db_engine.connect() try: # SA db_engine -> SA connection -> DBAPI connection -> YDB driver - driver = connection.connection.driver # type: ignore # 2024-01-24 # TODO: "DBAPIConnection" has no attribute "driver" [attr-defined] + driver = connection.connection._driver # type: ignore # 2024-01-24 # TODO: "DBAPIConnection" has no attribute "_driver" [attr-defined] assert driver queue = [db_name] diff --git a/lib/dl_connector_ydb/dl_connector_ydb/core/ydb/connector.py b/lib/dl_connector_ydb/dl_connector_ydb/core/ydb/connector.py index 8413f87f21..16907cdd31 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/core/ydb/connector.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/core/ydb/connector.py @@ -1,5 +1,3 @@ -from ydb.sqlalchemy import register_dialect as yql_register_dialect - from dl_core.connectors.base.connector import ( CoreBackendDefinition, CoreConnectionDefinition, @@ -14,6 +12,7 @@ SQLDataSourceSpecStorageSchema, SubselectDataSourceSpecStorageSchema, ) +import dl_sqlalchemy_ydb.dialect from dl_connector_ydb.core.base.query_compiler import YQLQueryCompiler from dl_connector_ydb.core.base.type_transformer import YQLTypeTransformer @@ -76,4 +75,4 @@ class YDBCoreConnector(CoreConnector): @classmethod def registration_hook(cls) -> None: - yql_register_dialect() + dl_sqlalchemy_ydb.dialect.register_dialect() diff --git a/lib/dl_connector_ydb/dl_connector_ydb/db_testing/engine_wrapper.py b/lib/dl_connector_ydb/dl_connector_ydb/db_testing/engine_wrapper.py index 782c86e87d..6934e493bc 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/db_testing/engine_wrapper.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/db_testing/engine_wrapper.py @@ -13,8 +13,10 @@ import sqlalchemy as sa from sqlalchemy.types import TypeEngine import ydb +import ydb_sqlalchemy as ydb_sa from dl_db_testing.database.engine_wrapper import EngineWrapperBase +import dl_sqlalchemy_ydb.dialect class YdbTypeSpec(NamedTuple): @@ -23,20 +25,37 @@ class YdbTypeSpec(NamedTuple): SA_TYPE_TO_YDB_TYPE: dict[type[TypeEngine], YdbTypeSpec] = { + ydb_sa.types.Int8: YdbTypeSpec(type=ydb.PrimitiveType.Int8, to_sql_str=str), + ydb_sa.types.Int16: YdbTypeSpec(type=ydb.PrimitiveType.Int16, to_sql_str=str), + ydb_sa.types.Int32: YdbTypeSpec(type=ydb.PrimitiveType.Int32, to_sql_str=str), + ydb_sa.types.Int64: YdbTypeSpec(type=ydb.PrimitiveType.Int64, to_sql_str=str), + ydb_sa.types.UInt8: YdbTypeSpec(type=ydb.PrimitiveType.Uint8, to_sql_str=str), + ydb_sa.types.UInt16: YdbTypeSpec(type=ydb.PrimitiveType.Uint16, to_sql_str=str), + ydb_sa.types.UInt32: YdbTypeSpec(type=ydb.PrimitiveType.Uint32, to_sql_str=str), + ydb_sa.types.UInt64: YdbTypeSpec(type=ydb.PrimitiveType.Uint64, to_sql_str=str), sa.SmallInteger: YdbTypeSpec(type=ydb.PrimitiveType.Uint8, to_sql_str=str), sa.Integer: YdbTypeSpec(type=ydb.PrimitiveType.Int32, to_sql_str=str), sa.BigInteger: YdbTypeSpec(type=ydb.PrimitiveType.Int64, to_sql_str=str), sa.Float: YdbTypeSpec(type=ydb.PrimitiveType.Double, to_sql_str=str), sa.Boolean: YdbTypeSpec(type=ydb.PrimitiveType.Bool, to_sql_str=lambda x: str(bool(x))), sa.String: YdbTypeSpec(type=ydb.PrimitiveType.String, to_sql_str=lambda x: f'"{x}"'), + sa.BINARY: YdbTypeSpec(type=ydb.PrimitiveType.String, to_sql_str=lambda x: f'"{x}"'), + sa.Text: YdbTypeSpec(type=ydb.PrimitiveType.String, to_sql_str=lambda x: f'"{x}"'), sa.Unicode: YdbTypeSpec(type=ydb.PrimitiveType.Utf8, to_sql_str=lambda x: f'"{x}"'), sa.Date: YdbTypeSpec(type=ydb.PrimitiveType.Date, to_sql_str=lambda x: f'DateTime::MakeDate($date_parse("{x}"))'), sa.DateTime: YdbTypeSpec( + ydb.PrimitiveType.Datetime, + to_sql_str=lambda x: f'DateTime::MakeDatetime($datetime_parse("{x}"))', + ), + sa.DATETIME: YdbTypeSpec( ydb.PrimitiveType.Datetime, to_sql_str=lambda x: f'DateTime::MakeDatetime($datetime_parse("{x}"))' ), sa.TIMESTAMP: YdbTypeSpec( ydb.PrimitiveType.Timestamp, to_sql_str=lambda x: f'DateTime::MakeTimestamp($datetime_parse("{x}"))' ), + dl_sqlalchemy_ydb.dialect.YqlInterval: YdbTypeSpec( + ydb.PrimitiveType.Interval, to_sql_str=lambda x: f"CAST({x} as Interval)" + ), } diff --git a/lib/dl_connector_ydb/dl_connector_ydb/formula/connector.py b/lib/dl_connector_ydb/dl_connector_ydb/formula/connector.py index 1fe75f65ba..ca3f506cb8 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/formula/connector.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/formula/connector.py @@ -1,7 +1,6 @@ -from ydb.sqlalchemy import YqlDialect as SAYqlDialect - from dl_formula.connectors.base.connector import FormulaConnector from dl_query_processing.compilation.query_mutator import RemoveConstFromGroupByFormulaAtomicQueryMutator +from dl_sqlalchemy_ydb.dialect import CustomYqlDialect from dl_connector_ydb.formula.constants import YqlDialect as YqlDialectNS from dl_connector_ydb.formula.definitions.all import DEFINITIONS @@ -11,7 +10,7 @@ class YQLFormulaConnector(FormulaConnector): dialect_ns_cls = YqlDialectNS dialects = YqlDialectNS.YQL op_definitions = DEFINITIONS - sa_dialect = SAYqlDialect() + sa_dialect = CustomYqlDialect() @classmethod def registration_hook(cls) -> None: diff --git a/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_aggregation.py b/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_aggregation.py index 20d43a4c8b..88384e9d06 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_aggregation.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_aggregation.py @@ -1,5 +1,5 @@ import sqlalchemy as sa -import ydb.sqlalchemy as ydb_sa +import ydb_sqlalchemy.sqlalchemy as ydb_sa from dl_formula.definitions.base import ( TranslationVariant, diff --git a/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_datetime.py b/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_datetime.py index a06b44f5db..8ba5bdf95d 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_datetime.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_datetime.py @@ -1,5 +1,6 @@ import sqlalchemy as sa from sqlalchemy.sql.elements import ClauseElement +import ydb_sqlalchemy as ydb_sa from dl_formula.connectors.base.literal import Literal from dl_formula.definitions.base import ( @@ -55,6 +56,8 @@ def _date_datetime_add_yql( type_name = "day" mult_expr = mult_expr * 7 # type: ignore # 2024-04-02 # TODO: Unsupported operand types for * ("ClauseElement" and "int") [operator] + mult_expr = sa.cast(mult_expr, ydb_sa.types.Int32) + func_name = YQL_INTERVAL_FUNCS.get(type_name) if func_name is not None: func = getattr(sa.func.DateTime, func_name) @@ -91,7 +94,7 @@ def _datetrunc2_yql_impl(date_ctx: TranslationCtx, unit_ctx: TranslationCtx) -> func = getattr(sa.func.DateTime, func_name) return sa.func.DateTime.MakeDatetime(func(date_expr)) - amount = 1 + amount = sa.cast(1, ydb_sa.types.Int32) func_name = YQL_INTERVAL_FUNCS.get(unit) if func_name is not None: func = getattr(sa.func.DateTime, func_name) diff --git a/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_string.py b/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_string.py index 74310c5294..9dd636966a 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_string.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_string.py @@ -1,5 +1,5 @@ import sqlalchemy as sa -import ydb.sqlalchemy as ydb_sa +import ydb_sqlalchemy.sqlalchemy as ydb_sa from dl_formula.definitions.base import TranslationVariant from dl_formula.definitions.common import ( @@ -39,9 +39,13 @@ value, # int -> List -> utf8 sa.func.Unicode.FromCodePointList( + # Note: Executing sqlalchemy statement without cast determines list type as List, + # while directly executing query with Int32 parameters automatically produces List. sa.func.AsList( # coalesce is needed to un-Nullable the type. - sa.func.COALESCE(sa.cast(value, ydb_sa.types.UInt32), 0), + sa.func.COALESCE( + sa.cast(value, ydb_sa.types.UInt32), sa.func.UNWRAP(sa.cast(0, ydb_sa.types.UInt32)) + ), ) ), ), diff --git a/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_type.py b/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_type.py index 7402a16b8c..698b816527 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_type.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/functions_type.py @@ -1,13 +1,355 @@ import sqlalchemy as sa +import ydb_sqlalchemy as ydb_sa +from dl_formula.core.datatype import DataType from dl_formula.definitions.base import TranslationVariant import dl_formula.definitions.functions_type as base +import dl_sqlalchemy_ydb.dialect as ydb_dialect from dl_connector_ydb.formula.constants import YqlDialect as D V = TranslationVariant.make +TYPES_SPEC = { + wlts.name: wlts + for wlts in [ + base.WhitelistTypeSpec(name="Bool", sa_type=sa.BOOLEAN), + base.WhitelistTypeSpec(name="Int8", sa_type=ydb_sa.types.Int8), + base.WhitelistTypeSpec(name="Int16", sa_type=ydb_sa.types.Int16), + base.WhitelistTypeSpec(name="Int32", sa_type=ydb_sa.types.Int32), + base.WhitelistTypeSpec(name="Int64", sa_type=ydb_sa.types.Int64), + base.WhitelistTypeSpec(name="UInt8", sa_type=ydb_sa.types.UInt8), + base.WhitelistTypeSpec(name="UInt16", sa_type=ydb_sa.types.UInt16), + base.WhitelistTypeSpec(name="UInt32", sa_type=ydb_sa.types.UInt32), + base.WhitelistTypeSpec(name="UInt64", sa_type=ydb_sa.types.UInt64), + base.WhitelistTypeSpec(name="Double", sa_type=sa.types.FLOAT), + base.WhitelistTypeSpec(name="Float", sa_type=sa.types.FLOAT), + base.WhitelistTypeSpec(name="Decimal", sa_type=sa.DECIMAL, arg_types=base.DECIMAL_CAST_ARG_T), + base.WhitelistTypeSpec(name="Utf8", sa_type=sa.types.TEXT), + base.WhitelistTypeSpec(name="String", sa_type=sa.types.TEXT), + base.WhitelistTypeSpec(name="Date", sa_type=sa.types.DATE), + base.WhitelistTypeSpec(name="Datetime", sa_type=ydb_dialect.YqlDateTime), + base.WhitelistTypeSpec(name="Timestamp", sa_type=ydb_dialect.YqlTimestamp), + base.WhitelistTypeSpec(name="Uuid", sa_type=sa.types.TEXT), + ] +} + +BOOL_TYPES_SPEC = [ + TYPES_SPEC["Bool"], +] + +INT_TYPES_SPEC = [ + TYPES_SPEC["Int8"], + TYPES_SPEC["Int16"], + TYPES_SPEC["Int32"], + TYPES_SPEC["Int64"], + TYPES_SPEC["UInt8"], + TYPES_SPEC["UInt16"], + TYPES_SPEC["UInt32"], + TYPES_SPEC["UInt64"], +] + +DECIMAL_TYPES_SPEC = [ + TYPES_SPEC["Decimal"], +] + +STRING_TYPES_SPEC = [ + TYPES_SPEC["Utf8"], + TYPES_SPEC["String"], +] + +FLOAT_TYPES_SPEC = [ + TYPES_SPEC["Double"], + TYPES_SPEC["Float"], +] + + +class YQLDbCastArgTypes(base.DbCastArgTypes): + def __init__(self) -> None: + # See DbCastArgTypes.__init__ + super(base.DbCastArgTypes, self).__init__( + arg_types=[ + { + DataType.BOOLEAN, + DataType.INTEGER, + DataType.FLOAT, + DataType.STRING, + DataType.DATE, + DataType.ARRAY_INT, + DataType.ARRAY_FLOAT, + DataType.ARRAY_STR, + }, + DataType.CONST_STRING, + ] + ) + + +class FuncDbCastYQLBase(base.FuncDbCastBase): + # For numeric types see: https://ydb.tech/docs/en/yql/reference/types/primitive#casting-to-numeric-types + # Type cast tables date: 2025-09-29. + # + # Type Bool Int8 Int16 Int32 Int64 Uint8 Uint16 Uint32 Uint64 Float Double Decimal <- (Target Type) + # Bool — Yes[1] Yes[1] Yes[1] Yes[1] Yes[1] Yes[1] Yes[1] Yes[1] Yes[1] Yes[1] No + # Int8 Yes2 — Yes Yes Yes Yes[3] Yes[3] Yes[3] Yes[3] Yes Yes Yes + # Int16 Yes2 Yes[4] — Yes Yes Yes[3,4] Yes[3] Yes[3] Yes[3] Yes Yes Yes + # Int32 Yes2 Yes[4] Yes[4] — Yes Yes[3,4] Yes[3,4] Yes[3] Yes[3] Yes Yes Yes + # Int64 Yes2 Yes[4] Yes[4] Yes[4] — Yes[3,4] Yes[3,4] Yes[3,4] Yes[3] Yes Yes Yes + # Uint8 Yes2 Yes[4] Yes Yes Yes — Yes Yes Yes Yes Yes Yes + # Uint16 Yes2 Yes[4] Yes[4] Yes Yes Yes[4] — Yes Yes Yes Yes Yes + # Uint32 Yes2 Yes[4] Yes[4] Yes[4] Yes Yes[4] Yes[4] — Yes Yes Yes Yes + # Uint64 Yes2 Yes[4] Yes[4] Yes[4] Yes[4] Yes[4] Yes[4] Yes[4] — Yes Yes Yes + # Float Yes2 Yes[4] Yes[4] Yes[4] Yes[4] Yes[3,4] Yes[3,4] Yes[3,4] Yes[3,4] — Yes No + # Double Yes2 Yes[4] Yes[4] Yes[4] Yes[4] Yes[3,4] Yes[3,4] Yes[3,4] Yes[3,4] Yes — No + # Decimal No Yes Yes Yes Yes Yes Yes Yes Yes Yes Yes — + # String Yes Yes Yes Yes Yes Yes Yes Yes Yes Yes Yes Yes + # Utf8 Yes Yes Yes Yes Yes Yes Yes Yes Yes Yes Yes Yes + # Json No No No No No No No No No No No No + # Yson Yes[5] Yes[5] Yes[5] Yes[5] Yes[5] Yes[5] Yes[5] Yes[5] Yes[5] Yes[5] Yes[5] No + # Uuid No No No No No No No No No No No No + # Date No Yes[4] Yes[4] Yes Yes Yes[4] Yes Yes Yes Yes Yes No + # Datetime No Yes[4] Yes[4] Yes[4] Yes Yes[4] Yes[4] Yes Yes Yes Yes No + # Timestamp No Yes[4] Yes[4] Yes[4] Yes[4] Yes[4] Yes[4] Yes[4] Yes Yes Yes No + # Interval No Yes[4] Yes[4] Yes[4] Yes Yes[3,4] Yes[3,4] Yes[3,4] Yes[3] Yes Yes No + # ^ + # | + # (Source Type) + # + # Type String Utf8 Json Yson Uuid + # Bool Yes No No No No + # INT Yes No No No No + # Uint Yes No No No No + # Float Yes No No No No + # Double Yes No No No No + # Decimal Yes No No No No + # String — Yes Yes Yes Yes + # Utf8 Yes — No No No + # Json Yes Yes — No No + # Yson Yes[4] No No No No + # Uuid Yes Yes No No — + # Date Yes Yes No No No + # Datetime Yes Yes No No No + # Timestamp Yes Yes No No No + # Interval Yes Yes No No No + # + # Type Date Datetime Timestamp Interval + # Bool No No No No + # INT Yes Yes Yes Yes + # Uint Yes Yes Yes Yes + # Float No No No No + # Double No No No No + # Decimal No No No No + # String Yes Yes Yes Yes + # Utf8 Yes Yes Yes Yes + # Json No No No No + # Yson No No No No + # Uuid No No No No + # Date — Yes Yes No + # Datetime Yes — Yes No + # Timestamp Yes Yes — No + # Interval No No No — + # + # [1] - True is converted to 1 and False to 0. + # [2] - Any value other than 0 is converted to True, 0 is converted to False. + # [3] - Possible only in case of a non-negative value. + # [4] - Possible only within the valid range. + # [5] - Using the built-in function Yson::ConvertTo. + + argument_types = [ + YQLDbCastArgTypes(), + ] + + WHITELISTS = { + yql_dialect: { + # TODO: Decimal + # TODO: DyNumber + # TODO: Json + # TODO: JsonDocument + # TODO: Yson + # TODO: Cast Integer to Interval + # Commented - not supported + DataType.BOOLEAN: [ + # > Bool + TYPES_SPEC["Bool"], + # > INT + TYPES_SPEC["Int8"], + TYPES_SPEC["Int16"], + TYPES_SPEC["Int32"], + TYPES_SPEC["Int64"], + # > UINT + TYPES_SPEC["UInt8"], + TYPES_SPEC["UInt16"], + TYPES_SPEC["UInt32"], + TYPES_SPEC["UInt64"], + # > Float + TYPES_SPEC["Float"], + # > Double + TYPES_SPEC["Double"], + # > Decimal + # TYPES_SPEC["Decimal"], + # > String + TYPES_SPEC["String"], + # > Utf8 + # TYPES_SPEC["Utf8"], + # > Date + # TYPES_SPEC["Date"], + # > Datetime + # TYPES_SPEC["Datetime"], + # > Timestamp + # TYPES_SPEC["Timestamp"], + # > UUID + # TYPES_SPEC["Uuid"], + ], + DataType.INTEGER: [ + # Interval can not be casted to Bool + # Interval can not be casted to Decimal + # > Bool + TYPES_SPEC["Bool"], + # > INT + TYPES_SPEC["Int8"], + TYPES_SPEC["Int16"], + TYPES_SPEC["Int32"], + TYPES_SPEC["Int64"], + # > UINT + TYPES_SPEC["UInt8"], + TYPES_SPEC["UInt16"], + TYPES_SPEC["UInt32"], + TYPES_SPEC["UInt64"], + # > Float + TYPES_SPEC["Float"], + # > Double + TYPES_SPEC["Double"], + # > Decimal + TYPES_SPEC["Decimal"], + # > String + TYPES_SPEC["String"], + # > Utf8 + # TYPES_SPEC["Utf8"], + # > Date + TYPES_SPEC["Date"], + # > Datetime + TYPES_SPEC["Datetime"], + # > Timestamp + TYPES_SPEC["Timestamp"], + # > UUID + # TYPES_SPEC["Uuid"], + ], + DataType.FLOAT: [ + # > Bool + TYPES_SPEC["Bool"], + # > INT + TYPES_SPEC["Int8"], + TYPES_SPEC["Int16"], + TYPES_SPEC["Int32"], + TYPES_SPEC["Int64"], + # > UINT + TYPES_SPEC["UInt8"], + TYPES_SPEC["UInt16"], + TYPES_SPEC["UInt32"], + TYPES_SPEC["UInt64"], + # > Float + TYPES_SPEC["Float"], + # > Double + TYPES_SPEC["Double"], + # > Decimal + # TYPES_SPEC["Decimal"], + # > String + TYPES_SPEC["String"], + # > Utf8 + # TYPES_SPEC["Utf8"], + # > Date + # TYPES_SPEC["Date"], + # > Datetime + # TYPES_SPEC["Datetime"], + # > Timestamp + # TYPES_SPEC["Timestamp"], + # > UUID + # TYPES_SPEC["Uuid"], + ], + DataType.STRING: [ + # Utf8 can not be casted to Uuid + # > Bool + TYPES_SPEC["Bool"], + # > INT + TYPES_SPEC["Int8"], + TYPES_SPEC["Int16"], + TYPES_SPEC["Int32"], + TYPES_SPEC["Int64"], + # > UINT + TYPES_SPEC["UInt8"], + TYPES_SPEC["UInt16"], + TYPES_SPEC["UInt32"], + TYPES_SPEC["UInt64"], + # > Float + TYPES_SPEC["Float"], + # > Double + TYPES_SPEC["Double"], + # > Decimal + TYPES_SPEC["Decimal"], + # > String + TYPES_SPEC["String"], + # > Utf8 + TYPES_SPEC["Utf8"], + # > Date + TYPES_SPEC["Date"], + # > Datetime + TYPES_SPEC["Datetime"], + # > Timestamp + TYPES_SPEC["Timestamp"], + # > UUID + TYPES_SPEC["Uuid"], + ], + DataType.DATE: [ + # > Bool + # TYPES_SPEC["Bool"], + # > INT + TYPES_SPEC["Int8"], + TYPES_SPEC["Int16"], + TYPES_SPEC["Int32"], + TYPES_SPEC["Int64"], + # > UINT + TYPES_SPEC["UInt8"], + TYPES_SPEC["UInt16"], + TYPES_SPEC["UInt32"], + TYPES_SPEC["UInt64"], + # > Float + TYPES_SPEC["Float"], + # > Double + TYPES_SPEC["Double"], + # > Decimal + # TYPES_SPEC["Decimal"], + # > String + TYPES_SPEC["String"], + # > Utf8 + TYPES_SPEC["Utf8"], + # > Date + TYPES_SPEC["Date"], + # > Datetime + TYPES_SPEC["Datetime"], + # > Timestamp + TYPES_SPEC["Timestamp"], + # > UUID + # TYPES_SPEC["Uuid"], + ], + DataType.ARRAY_STR: [], + DataType.ARRAY_INT: [], + DataType.ARRAY_FLOAT: [], + } + for yql_dialect in (D.YQL, D.YQ, D.YDB) + } + + +class FuncDbCastYQL2(FuncDbCastYQLBase, base.FuncDbCast2): + pass + + +class FuncDbCastYQL3(FuncDbCastYQLBase, base.FuncDbCast3): + pass + + +class FuncDbCastYQL4(FuncDbCastYQLBase, base.FuncDbCast4): + pass + DEFINITIONS_TYPE = [ # bool @@ -44,6 +386,10 @@ ), # datetimetz base.FuncDatetimeTZConst.for_dialect(D.YQL), + # db_cast + FuncDbCastYQL2(), + FuncDbCastYQL3(), + FuncDbCastYQL4(), # float base.FuncFloatNumber( variants=[ diff --git a/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/operators_binary.py b/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/operators_binary.py index 224663a72f..4e83980c09 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/operators_binary.py +++ b/lib/dl_connector_ydb/dl_connector_ydb/formula/definitions/operators_binary.py @@ -1,7 +1,12 @@ import sqlalchemy as sa +import ydb_sqlalchemy as ydb_sa from dl_formula.definitions.base import TranslationVariant -from dl_formula.definitions.common_datetime import DAY_USEC +from dl_formula.definitions.common_datetime import ( + DAY_SEC, + DAY_USEC, + SEC_USEC, +) import dl_formula.definitions.operators_binary as base from dl_connector_ydb.formula.constants import YqlDialect as D @@ -24,19 +29,24 @@ base.BinaryPlusStrings.for_dialect(D.YQL), base.BinaryPlusDateInt( variants=[ - V(D.YQL, lambda date, days: date + sa.func.DateTime.IntervalFromDays(days)), + V(D.YQL, lambda date, days: date + sa.func.DateTime.IntervalFromDays(sa.cast(days, ydb_sa.types.Int32))), ] ), base.BinaryPlusDateFloat( variants=[ - V(D.YQL, lambda date, days: date + sa.func.DateTime.IntervalFromDays(sa.cast(days, sa.INTEGER))), + V(D.YQL, lambda date, days: date + sa.func.DateTime.IntervalFromDays(sa.cast(days, ydb_sa.types.Int32))), ] ), base.BinaryPlusDatetimeNumber( variants=[ V( D.YQL, - lambda date, days: (date + sa.func.DateTime.IntervalFromMicroseconds(base.as_bigint(days * DAY_USEC))), + lambda date, days: ( + date + + sa.func.DateTime.IntervalFromMicroseconds( + base.as_bigint(days * base.as_bigint(DAY_SEC) * base.as_bigint(SEC_USEC)) + ) + ), ), ] ), @@ -44,7 +54,12 @@ variants=[ V( D.YQL, - lambda dt, days: (dt + sa.func.DateTime.IntervalFromMicroseconds(base.as_bigint(days * DAY_USEC))), + lambda dt, days: ( + dt + + sa.func.DateTime.IntervalFromMicroseconds( + base.as_bigint(days * base.as_bigint(DAY_SEC) * base.as_bigint(SEC_USEC)) + ) + ), ), ] ), @@ -57,7 +72,7 @@ base.BinaryMinusNumbers.for_dialect(D.YQL), base.BinaryMinusDateInt( variants=[ - V(D.YQL, lambda date, days: date - sa.func.DateTime.IntervalFromDays(days)), + V(D.YQL, lambda date, days: date - sa.func.DateTime.IntervalFromDays(sa.cast(days, ydb_sa.types.Int32))), ] ), base.BinaryMinusDateFloat( @@ -65,7 +80,7 @@ V( D.YQL, lambda date, days: ( - date - sa.func.DateTime.IntervalFromDays(sa.cast(sa.func.Math.Ceil(days), sa.INTEGER)) + date - sa.func.DateTime.IntervalFromDays(sa.cast(sa.func.Math.Ceil(days), ydb_sa.types.Int32)) ), ), ] @@ -74,7 +89,12 @@ variants=[ V( D.YQL, - lambda date, days: (date - sa.func.DateTime.IntervalFromMicroseconds(base.as_bigint(days * DAY_USEC))), + lambda date, days: ( + date + - sa.func.DateTime.IntervalFromMicroseconds( + base.as_bigint(days * base.as_bigint(DAY_SEC) * base.as_bigint(SEC_USEC)) + ) + ), ), ] ), @@ -82,7 +102,12 @@ variants=[ V( D.YQL, - lambda dt, days: (dt - sa.func.DateTime.IntervalFromMicroseconds(base.as_bigint(days * DAY_USEC))), + lambda dt, days: ( + dt + - sa.func.DateTime.IntervalFromMicroseconds( + base.as_bigint(days * base.as_bigint(DAY_SEC) * base.as_bigint(SEC_USEC)) + ) + ), ), ] ), diff --git a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/base.py b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/base.py index 94a1ddf39b..6c22a974b9 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/base.py +++ b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/base.py @@ -1,5 +1,6 @@ from typing import ClassVar +from frozendict import frozendict import pytest from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration @@ -40,6 +41,18 @@ class YDBConnectionTestBase(ConnectionTestBase): def db_url(self) -> str: return DB_CORE_URL + @pytest.fixture(scope="class") + def engine_params(self) -> dict: + return dict( + connect_args=frozendict( + dict( + host=CoreConnectionSettings.HOST, + port=CoreConnectionSettings.PORT, + protocol="grpc", + ) + ), + ) + @pytest.fixture(scope="class") def bi_test_config(self) -> ApiTestEnvironmentConfiguration: return API_TEST_CONFIG diff --git a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/test_dashsql.py b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/test_dashsql.py index e948c298b3..ab7ff774c5 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/test_dashsql.py +++ b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/test_dashsql.py @@ -24,7 +24,7 @@ async def test_result( resp_data = await resp.json() assert resp_data[0]["event"] == "metadata", resp_data - assert resp_data[0]["data"]["names"][:12] == [ + assert resp_data[0]["data"]["names"][:13] == [ "id", "some_str", "some_utf8", @@ -37,8 +37,9 @@ async def test_result( "some_date", "some_datetime", "some_timestamp", + "some_interval", ] - assert resp_data[0]["data"]["driver_types"][:12] == [ + assert resp_data[0]["data"]["driver_types"][:13] == [ "int32?", "string", "utf8?", @@ -51,22 +52,24 @@ async def test_result( "date", "datetime", "timestamp", + "interval", ] - assert resp_data[0]["data"]["db_types"][:12] == [ - "integer", + assert resp_data[0]["data"]["db_types"][:13] == [ + "int32", "text", "text", - "integer", - "integer", - "integer", - "integer", + "int32", + "uint8", + "int64", + "uint64", "float", "boolean", "date", "datetime", "datetime", + "interval", ] - assert resp_data[0]["data"]["bi_types"][:12] == [ + assert resp_data[0]["data"]["bi_types"][:13] == [ "integer", "string", "string", @@ -79,6 +82,7 @@ async def test_result( "date", "genericdatetime", "genericdatetime", + "integer", ] assert resp_data[-1]["event"] == "footer", resp_data[-1] @@ -105,7 +109,7 @@ async def test_interval( "interval", ] assert resp_data[0]["data"]["db_types"] == [ - "integer", + "interval", ] assert resp_data[0]["data"]["bi_types"] == [ "integer", diff --git a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/config.py b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/config.py index b4843bba6f..d468bf9e6c 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/config.py +++ b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/config.py @@ -1,12 +1,12 @@ from typing import ClassVar -from frozendict import frozendict import requests import sqlalchemy as sa from dl_api_lib_testing.configuration import ApiTestEnvironmentConfiguration from dl_constants.enums import UserDataType from dl_core_testing.configuration import CoreTestEnvironmentConfiguration +import dl_sqlalchemy_ydb.dialect from dl_testing.containers import get_test_container_hostport from dl_connector_ydb.formula.constants import YqlDialect as D @@ -35,16 +35,11 @@ def fetch_ca_certificate() -> str: return response.text -def make_ssl_engine_params(ssl_ca: str) -> dict: - engine_params = { - "connect_args": frozendict( - { - "ca_cert": ssl_ca, - "root_certificates": ssl_ca.encode("ascii"), - } - ), +def make_ssl_connect_args(ssl_ca: str) -> dict: + return { + "ca_cert": ssl_ca, + "root_certificates": ssl_ca.encode("ascii"), } - return engine_params class CoreConnectionSettings: @@ -80,8 +75,9 @@ class CoreSslConnectionSettings: ("some_string", UserDataType.string, sa.String), ("some_utf8", UserDataType.string, sa.Unicode), ("some_date", UserDataType.date, sa.Date), - ("some_datetime", UserDataType.genericdatetime, sa.DateTime), + ("some_datetime", UserDataType.genericdatetime, sa.DATETIME), ("some_timestamp", UserDataType.genericdatetime, sa.TIMESTAMP), + ("some_interval", UserDataType.integer, dl_sqlalchemy_ydb.dialect.YqlInterval), ) TABLE_DATA = [ { @@ -97,6 +93,7 @@ class CoreSslConnectionSettings: "some_date": None, "some_datetime": None, "some_timestamp": None, + "some_interval": 9, }, { "id": 2, @@ -111,6 +108,7 @@ class CoreSslConnectionSettings: "some_date": None, "some_datetime": None, "some_timestamp": None, + "some_interval": None, }, { "id": 3, @@ -125,6 +123,7 @@ class CoreSslConnectionSettings: "some_date": None, "some_datetime": None, "some_timestamp": None, + "some_interval": 1337, }, { "id": 4, @@ -139,6 +138,7 @@ class CoreSslConnectionSettings: "some_date": None, "some_datetime": None, "some_timestamp": None, + "some_interval": 42, }, { "id": 5, @@ -153,6 +153,7 @@ class CoreSslConnectionSettings: "some_date": None, "some_datetime": None, "some_timestamp": None, + "some_interval": None, }, { "id": 6, @@ -167,6 +168,7 @@ class CoreSslConnectionSettings: "some_date": "2021-06-07", "some_datetime": "2021-06-07T18:19:20Z", "some_timestamp": "2021-06-07T18:19:20Z", + "some_interval": None, }, { "id": 7, @@ -181,6 +183,7 @@ class CoreSslConnectionSettings: "some_date": "1970-12-31", "some_datetime": "1970-12-31T23:58:57Z", "some_timestamp": "1970-12-31T23:58:57Z", + "some_interval": 0, }, { "id": 8, @@ -195,6 +198,7 @@ class CoreSslConnectionSettings: "some_date": "1972-03-28", "some_datetime": "1972-03-28T17:11:02Z", "some_timestamp": "1972-03-28T17:11:02Z", + "some_interval": 1, }, { "id": 9, @@ -209,6 +213,7 @@ class CoreSslConnectionSettings: "some_date": None, "some_datetime": None, "some_timestamp": None, + "some_interval": None, }, { "id": 10, @@ -223,6 +228,7 @@ class CoreSslConnectionSettings: "some_date": "2026-07-08", "some_datetime": None, "some_timestamp": None, + "some_interval": None, }, { "id": 11, @@ -237,6 +243,7 @@ class CoreSslConnectionSettings: "some_date": "2029-11-04", "some_datetime": None, "some_timestamp": None, + "some_interval": 1234, }, ] TABLE_NAME = "test_table_h" @@ -255,6 +262,7 @@ class CoreSslConnectionSettings: MAX(Date('2021-06-09')) as some_date, MAX(Datetime('2021-06-09T20:50:47Z')) as some_datetime, MAX(Timestamp('2021-07-10T21:51:48.841512Z')) as some_timestamp, + CAST(1 AS Interval) as some_interval, MAX(ListHead(ListSkip(Unicode::SplitToList(CAST(some_string AS UTF8), ''), 3))) as str_split, MAX(ListConcat(ListReplicate(CAST(' ' AS UTF8), 5))) as num_space_by_lst, diff --git a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/conftest.py b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/conftest.py index 91534df3d0..99ba8fa86c 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/conftest.py +++ b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/conftest.py @@ -1,5 +1,6 @@ from dl_api_lib_testing.initialization import initialize_api_lib_test from dl_formula_testing.forced_literal import forced_literal_use +import dl_sqlalchemy_ydb.dialect from dl_connector_ydb_tests.db.config import API_TEST_CONFIG @@ -7,6 +8,8 @@ def pytest_configure(config): # noqa initialize_api_lib_test(pytest_config=config, api_test_config=API_TEST_CONFIG) + dl_sqlalchemy_ydb.dialect.register_dialect() + __all__ = ( # auto-use fixtures: diff --git a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/core/base.py b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/core/base.py index a05c944f3b..8d798b6337 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/core/base.py +++ b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/core/base.py @@ -1,6 +1,7 @@ import asyncio from typing import Generator +from frozendict import frozendict import pytest from dl_core_testing.database import ( @@ -38,6 +39,18 @@ def loop(self, event_loop: asyncio.AbstractEventLoop) -> Generator[asyncio.Abstr def db_url(self) -> str: return test_config.DB_CORE_URL + @pytest.fixture(scope="class") + def engine_params(self) -> dict: + return dict( + connect_args=frozendict( + dict( + host=test_config.CoreConnectionSettings.HOST, + port=test_config.CoreConnectionSettings.PORT, + protocol="grpc", + ) + ), + ) + @pytest.fixture(scope="function") def connection_creation_params(self) -> dict: return dict( @@ -68,7 +81,16 @@ def ssl_ca(self) -> str: @pytest.fixture(scope="class") def engine_params(self, ssl_ca: str) -> dict: - return test_config.make_ssl_engine_params(ssl_ca) + return dict( + connect_args=frozendict( + dict( + host=test_config.CoreSslConnectionSettings.HOST, + port=test_config.CoreSslConnectionSettings.PORT, + protocol="grpcs", + **test_config.make_ssl_connect_args(ssl_ca), + ), + ), + ) @pytest.fixture(scope="class") def db_url(self) -> str: diff --git a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/formula/base.py b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/formula/base.py index 6af6069f1a..5db70e39bc 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/formula/base.py +++ b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/formula/base.py @@ -5,6 +5,7 @@ Optional, ) +from frozendict import frozendict import pytest import sqlalchemy as sa @@ -18,6 +19,7 @@ from dl_formula_testing.testcases.base import FormulaConnectorTestBase from dl_connector_ydb.formula.constants import YqlDialect as D +import dl_connector_ydb_tests.db.config as test_config from dl_connector_ydb_tests.db.config import DB_CONFIGURATIONS @@ -131,3 +133,15 @@ def ydb_data_table_field_types_patch(self, monkeypatch): monkeypatch.setattr("dl_formula_testing.evaluator.FIELD_TYPES", ydb_field_types) return ydb_field_types + + @pytest.fixture(scope="class") + def engine_params(self) -> dict: + return dict( + connect_args=frozendict( + dict( + host=test_config.CoreConnectionSettings.HOST, + port=test_config.CoreConnectionSettings.PORT, + protocol="grpc", + ) + ), + ) diff --git a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/formula/test_functions_type_conversion.py b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/formula/test_functions_type_conversion.py index 7b7360f5ed..61aeea085e 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/formula/test_functions_type_conversion.py +++ b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/formula/test_functions_type_conversion.py @@ -1,7 +1,20 @@ +import contextlib +import datetime +from typing import ( + Generator, + Optional, +) + +import pytest +import sqlalchemy as sa + +from dl_formula.core.datatype import DataType +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, @@ -75,3 +88,388 @@ class TestGeopointTypeFunctionYQL(YQLTestBase, DefaultGeopointTypeFunctionFormul class TestGeopolygonTypeFunctionYQL(YQLTestBase, DefaultGeopolygonTypeFunctionFormulaConnectorTestSuite): pass + + +# DB_CAST + + +class DbCastTypeFunctionYQLTestSuite( + DefaultDbCastTypeFunctionFormulaConnectorTestSuite, +): + def test_db_cast_ydb(self, dbe: DbEvaluator, data_table: sa.Table) -> None: + # Valid cast + value = dbe.eval("[int_value]", from_=data_table) + assert dbe.eval('DB_CAST(FLOAT([int_value]), "Double")', from_=data_table) == pytest.approx(float(value)) + + # # Test that it works with bool + dbe.eval('DB_CAST(BOOL([int_value]), "Double")', from_=data_table) + # Test that it works with int + dbe.eval('DB_CAST(INT([int_value]), "Int64")', from_=data_table) + # Test that it works with float + dbe.eval('DB_CAST(FLOAT([int_value]), "Double")', from_=data_table) + # Test that it works with string + dbe.eval('DB_CAST(STR([int_value]), "Utf8")', from_=data_table) + + # Cast to decimal with correct arguments + assert dbe.eval('DB_CAST([int_value], "Decimal", 5, 0)', from_=data_table) == value + + # Invalid number of arguments for Decimal + with pytest.raises(exc.TranslationError): + dbe.eval('DB_CAST([int_value], "Decimal", 5)', from_=data_table) + + with pytest.raises(exc.TranslationError): + dbe.eval('DB_CAST([int_value], "Decimal", "5", "3")', from_=data_table) + + # Invalid cast from Integer to Uuid + with pytest.raises(exc.TranslationError): + dbe.eval('DB_CAST([int_value], "Uuid")', from_=data_table) + + # Cast into itself + assert dbe.eval('DB_CAST(DB_CAST([int_value], "Int64"), "Int64")', from_=data_table) == value + + # Cast and cast back + assert dbe.eval('DB_CAST(DB_CAST(DB_CAST([int_value], "Int64"), "UInt64"), "Int64")', from_=data_table) == value + + # Castn't + with pytest.raises(exc.TranslationError): + assert dbe.eval('DB_CAST([int_value], "meow")', from_=data_table) == value + + def _test_db_cast_ydb_func( + self, + dbe: DbEvaluator, + ydb_type_test_data_table: sa.Table, + target: str, + cast_args: tuple[int, int] | None, + ok: bool, + ydb_data_test_table_field_types_patch, + source_column: str, + ) -> None: + if cast_args: + cast_args_str = ", ".join(cast_args) + query_string = f'DB_CAST([{source_column}], "{target}", {cast_args_str})' + else: + query_string = f'DB_CAST([{source_column}], "{target}")' + + if ok: + dbe.eval(query_string, from_=ydb_type_test_data_table) + else: + with pytest.raises(exc.TranslationError): + dbe.eval(query_string, from_=ydb_type_test_data_table) + + @pytest.mark.parametrize( + # target - target type for cast + # cast_args - type arguments (for decimal) + # ok - if no exception should occur + "target,cast_args,ok", + [ + # Bool + ("Bool", None, True), + # Int + ("Int8", None, True), + ("Int16", None, True), + ("Int32", None, True), + ("Int64", None, True), + ("UInt8", None, True), + ("UInt16", None, True), + ("UInt32", None, True), + # Float + ("Float", None, True), + ("Double", None, True), + # String + ("String", None, True), + ("Utf8", None, False), + # Date + ("Date", None, False), + ("Datetime", None, False), + ("Timestamp", None, False), + # Uuid + ("Uuid", None, False), + ], + ) + def test_db_cast_ydb_bool( + self, + dbe: DbEvaluator, + ydb_type_test_data_table: sa.Table, + target: str, + cast_args: tuple[int, int] | None, + ok: bool, + ydb_data_test_table_field_types_patch, + ) -> None: + self._test_db_cast_ydb_func( + dbe=dbe, + ydb_type_test_data_table=ydb_type_test_data_table, + target=target, + cast_args=cast_args, + ok=ok, + ydb_data_test_table_field_types_patch=ydb_data_test_table_field_types_patch, + source_column="bool_value", + ) + + @pytest.mark.parametrize( + # target - target type for cast + # cast_args - type arguments (for decimal) + # ok - if no exception should occur + "target,cast_args,ok", + [ + # Bool + ("Bool", None, True), + # Int + ("Int8", None, True), + ("Int16", None, True), + ("Int32", None, True), + ("Int64", None, True), + ("UInt8", None, True), + ("UInt16", None, True), + ("UInt32", None, True), + # Float + ("Float", None, True), + ("Double", None, True), + # String + ("String", None, True), + ("Utf8", None, False), + # Date + ("Date", None, True), + ("Datetime", None, True), + ("Timestamp", None, True), + # Uuid + ("Uuid", None, False), + ], + ) + def test_db_cast_ydb_integer( + self, + dbe: DbEvaluator, + ydb_type_test_data_table: sa.Table, + target: str, + cast_args: tuple[int, int] | None, + ok: bool, + ydb_data_test_table_field_types_patch, + ) -> None: + self._test_db_cast_ydb_func( + dbe=dbe, + ydb_type_test_data_table=ydb_type_test_data_table, + target=target, + cast_args=cast_args, + ok=ok, + ydb_data_test_table_field_types_patch=ydb_data_test_table_field_types_patch, + source_column="int64_value", + ) + + @pytest.mark.parametrize( + # target - target type for cast + # cast_args - type arguments (for decimal) + # ok - if no exception should occur + "target,cast_args,ok", + [ + # Bool + ("Bool", None, True), + # Int + ("Int8", None, True), + ("Int16", None, True), + ("Int32", None, True), + ("Int64", None, True), + ("UInt8", None, True), + ("UInt16", None, True), + ("UInt32", None, True), + # Float + ("Float", None, True), + ("Double", None, True), + # String + ("String", None, True), + ("Utf8", None, False), + # Date + ("Date", None, False), + ("Datetime", None, False), + ("Timestamp", None, False), + # Uuid + ("Uuid", None, False), + ], + ) + def test_db_cast_ydb_float( + self, + dbe: DbEvaluator, + ydb_type_test_data_table: sa.Table, + target: str, + cast_args: tuple[int, int] | None, + ok: bool, + ydb_data_test_table_field_types_patch, + ) -> None: + self._test_db_cast_ydb_func( + dbe=dbe, + ydb_type_test_data_table=ydb_type_test_data_table, + target=target, + cast_args=cast_args, + ok=ok, + ydb_data_test_table_field_types_patch=ydb_data_test_table_field_types_patch, + source_column="float_value", + ) + + @pytest.mark.parametrize( + # target - target type for cast + # cast_args - type arguments (for decimal) + # ok - if no exception should occur + "target,cast_args,ok", + [ + # Bool + ("Bool", None, True), + # Int + ("Int8", None, True), + ("Int16", None, True), + ("Int32", None, True), + ("Int64", None, True), + ("UInt8", None, True), + ("UInt16", None, True), + ("UInt32", None, True), + # Float + ("Float", None, True), + ("Double", None, True), + # String + ("String", None, True), + ("Utf8", None, True), + # Date + ("Date", None, True), + ("Datetime", None, True), + ("Timestamp", None, True), + # Uuid + ("Uuid", None, True), + ], + ) + def test_db_cast_ydb_string( + self, + dbe: DbEvaluator, + ydb_type_test_data_table: sa.Table, + target: str, + cast_args: tuple[int, int] | None, + ok: bool, + ydb_data_test_table_field_types_patch, + ) -> None: + self._test_db_cast_ydb_func( + dbe=dbe, + ydb_type_test_data_table=ydb_type_test_data_table, + target=target, + cast_args=cast_args, + ok=ok, + ydb_data_test_table_field_types_patch=ydb_data_test_table_field_types_patch, + source_column="string_value", + ) + + @pytest.mark.parametrize( + # target - target type for cast + # cast_args - type arguments (for decimal) + # ok - if no exception should occur + "target,cast_args,ok", + [ + # Bool + ("Bool", None, False), + # Int + ("Int8", None, True), + ("Int16", None, True), + ("Int32", None, True), + ("Int64", None, True), + ("UInt8", None, True), + ("UInt16", None, True), + ("UInt32", None, True), + # Float + ("Float", None, True), + ("Double", None, True), + # String + ("String", None, True), + ("Utf8", None, True), + # Date + ("Date", None, True), + ("Datetime", None, True), + ("Timestamp", None, True), + # Uuid + ("Uuid", None, False), + ], + ) + def test_db_cast_ydb_date( + self, + dbe: DbEvaluator, + ydb_type_test_data_table: sa.Table, + target: str, + cast_args: tuple[int, int] | None, + ok: bool, + ydb_data_test_table_field_types_patch, + ) -> None: + self._test_db_cast_ydb_func( + dbe=dbe, + ydb_type_test_data_table=ydb_type_test_data_table, + target=target, + cast_args=cast_args, + ok=ok, + ydb_data_test_table_field_types_patch=ydb_data_test_table_field_types_patch, + source_column="date_value", + ) + + +class DbCastYQLTestSuiteBase(YQLTestBase): + @contextlib.contextmanager + def make_ydb_type_test_data_table( + self, dbe: DbEvaluator, table_schema_name: Optional[str] + ) -> Generator[sa.Table, None, None]: + db = dbe.db + table_spec = self.generate_table_spec(table_name_prefix="ydb_type_test_table") + + columns = [ + sa.Column("bool_value", sa.Boolean()), + sa.Column("int64_value", sa.Integer(), primary_key=True), + sa.Column("float_value", sa.Float()), + sa.Column("string_value", sa.Text()), + sa.Column("date_value", sa.Date()), + ] + + table = self.lowlevel_make_sa_table( + db=db, table_spec=table_spec, table_schema_name=table_schema_name, columns=columns + ) + + db.create_table(table) + + table_data = [ + { + "bool_value": True, + "int64_value": 42, + "float_value": 0.1 + 0.2, + "string_value": "lobster", + "date_value": datetime.date(2000, 1, 2), + }, + ] + + db.insert_into_table(table, table_data) + + try: + yield table + finally: + dbe.db.drop_table(table) + + @pytest.fixture(scope="class") + def ydb_type_test_data_table( + self, dbe: DbEvaluator, table_schema_name: Optional[str] + ) -> Generator[sa.Table, None, None]: + with self.make_ydb_type_test_data_table(dbe=dbe, table_schema_name=table_schema_name) as table: + yield table + + # YDB-specific field types for formula testing + YDB_TYPE_FIELD_TYPES = { + "bool_value": DataType.BOOLEAN, + "int64_value": DataType.INTEGER, + "float_value": DataType.FLOAT, + "string_value": DataType.STRING, + "timestamp_value": DataType.DATETIME, # YDB TIMESTAMP maps to DATETIME in formula system + "date_value": DataType.DATE, + "datetime_value": DataType.DATETIME, + } + + @pytest.fixture(scope="function") + def ydb_data_test_table_field_types_patch(self, monkeypatch) -> None: + ydb_field_types = {**self.YDB_TYPE_FIELD_TYPES} + + monkeypatch.setattr("dl_formula_testing.evaluator.FIELD_TYPES", ydb_field_types) + + return ydb_field_types + + +class TestDbCastTypeFunctionYQL( + DbCastYQLTestSuiteBase, + DbCastTypeFunctionYQLTestSuite, +): + pass diff --git a/lib/dl_connector_ydb/pyproject.toml b/lib/dl_connector_ydb/pyproject.toml index 2939d63b41..4e5960f73c 100644 --- a/lib/dl_connector_ydb/pyproject.toml +++ b/lib/dl_connector_ydb/pyproject.toml @@ -19,6 +19,7 @@ dl-formula = {path = "../dl_formula"} dl-formula-ref = {path = "../dl_formula_ref"} dl-i18n = {path = "../dl_i18n"} dl-query-processing = {path = "../dl_query_processing"} +dl-sqlalchemy-ydb = {path = "../../lib/dl_sqlalchemy_ydb"} dl-type-transformer = {path = "../dl_type_transformer"} dl-utils = {path = "../dl_utils"} grpcio = "*" @@ -28,9 +29,12 @@ shortuuid = "*" sqlalchemy = "*" typing-extensions = "*" ydb = "*" +ydb-dbapi = "*" +ydb-sqlalchemy = "*" [tool.poetry.group.tests.dependencies] dl-formula-testing = {path = "../dl_formula_testing"} +dl-sqlalchemy-ydb = {path = "../../lib/dl_sqlalchemy_ydb"} pytest = "*" [tool.poetry.plugins] @@ -53,9 +57,6 @@ yql = "dl_connector_ydb.formula_ref.plugin:YQLFormulaRefPlugin" [tool.deptry.package_module_name_map] grpcio = "grpc" -[tool.deptry.per_rule_ignores] -DEP001 = ["ydb_proto_stubs_import"] - [tool.mypy] check_untyped_defs = true disallow_untyped_defs = true @@ -67,7 +68,7 @@ warn_unused_ignores = true [[tool.mypy.overrides]] ignore_missing_imports = true -module = ["ydb.*", "ydb_proto_stubs_import.*"] +module = ["ydb.*", "ydb_dbapi.*", "ydb_sqlalchemy.*"] [tool.pytest.ini_options] addopts = "-ra" diff --git a/lib/dl_formula/dl_formula/definitions/common_datetime.py b/lib/dl_formula/dl_formula/definitions/common_datetime.py index b9c7f9948b..57e31a5ae9 100644 --- a/lib/dl_formula/dl_formula/definitions/common_datetime.py +++ b/lib/dl_formula/dl_formula/definitions/common_datetime.py @@ -22,7 +22,8 @@ DAY_SEC = 3600 * 24 -DAY_USEC = DAY_SEC * 1_000_000 # microseconds +SEC_USEC = 1_000_000 # microseconds +DAY_USEC = DAY_SEC * SEC_USEC # microseconds EPOCH_START_S = "1970-01-01" EPOCH_START_D = datetime.date(1970, 1, 1) diff --git a/lib/dl_sqlalchemy_ydb/LICENSE b/lib/dl_sqlalchemy_ydb/LICENSE new file mode 100644 index 0000000000..74ba5f6c73 --- /dev/null +++ b/lib/dl_sqlalchemy_ydb/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2023 YANDEX LLC + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/lib/dl_sqlalchemy_ydb/README.rst b/lib/dl_sqlalchemy_ydb/README.rst new file mode 100644 index 0000000000..402af471a1 --- /dev/null +++ b/lib/dl_sqlalchemy_ydb/README.rst @@ -0,0 +1,4 @@ +CHYT SQLAlchemy +=============== + +YDB Dialect for sqlalchemy, based on ydb-sqlalchemy. diff --git a/lib/dl_sqlalchemy_ydb/dl_sqlalchemy_ydb/__init__.py b/lib/dl_sqlalchemy_ydb/dl_sqlalchemy_ydb/__init__.py new file mode 100644 index 0000000000..fd8c090757 --- /dev/null +++ b/lib/dl_sqlalchemy_ydb/dl_sqlalchemy_ydb/__init__.py @@ -0,0 +1,27 @@ +import sys + +import google.protobuf +import packaging.version + + +protobuf_version = packaging.version.Version(google.protobuf.__version__) + + +if protobuf_version < packaging.version.Version("4.0"): + from ydb._grpc.v3.draft import * # noqa + from ydb._grpc.v3.draft.protos import * # noqa + + sys.modules["ydb._grpc.common.draft"] = sys.modules["ydb._grpc.v3.draft"] + sys.modules["ydb._grpc.common.draft.protos"] = sys.modules["ydb._grpc.v3.draft.protos"] +elif protobuf_version < packaging.version.Version("5.0"): + from ydb._grpc.v4.draft import * # noqa + from ydb._grpc.v4.draft.protos import * # noqa + + sys.modules["ydb._grpc.common.draft"] = sys.modules["ydb._grpc.v4.draft"] + sys.modules["ydb._grpc.common.draft.protos"] = sys.modules["ydb._grpc.v4.draft.protos"] +else: + from ydb._grpc.v5.draft import * # noqa + from ydb._grpc.v5.draft.protos import * # noqa + + sys.modules["ydb._grpc.common.draft"] = sys.modules["ydb._grpc.v5.draft"] + sys.modules["ydb._grpc.common.draft.protos"] = sys.modules["ydb._grpc.v5.draft.protos"] diff --git a/lib/dl_sqlalchemy_ydb/dl_sqlalchemy_ydb/dialect.py b/lib/dl_sqlalchemy_ydb/dl_sqlalchemy_ydb/dialect.py new file mode 100644 index 0000000000..f051f77804 --- /dev/null +++ b/lib/dl_sqlalchemy_ydb/dl_sqlalchemy_ydb/dialect.py @@ -0,0 +1,155 @@ +import datetime +import typing + +import sqlalchemy as sa +import ydb_sqlalchemy.sqlalchemy as ydb_sa + + +class YqlTimestamp(sa.types.DateTime): + def result_processor(self, dialect: sa.engine.Dialect, coltype: typing.Any) -> typing.Any: + def process(value: typing.Optional[datetime.datetime]) -> typing.Optional[datetime.datetime]: + if value is None: + return None + if not self.timezone: + return value + return value.replace(tzinfo=datetime.timezone.utc) + + return process + + +class YqlDateTime(YqlTimestamp, sa.types.DateTime): + def bind_processor(self, dialect: sa.engine.Dialect) -> typing.Any: + def process(value: typing.Optional[datetime.datetime]) -> typing.Optional[int]: + if value is None: + return None + if not self.timezone: + value = value.replace(tzinfo=datetime.timezone.utc) + return int(value.timestamp()) + + return process + + +class YqlInterval(sa.types.Interval): + __visit_name__ = "interval" + + def result_processor(self, dialect: sa.engine.Dialect, coltype: typing.Any) -> typing.Any: + def process(value: typing.Optional[datetime.timedelta]) -> typing.Optional[int]: + if value is None: + return None + if isinstance(value, datetime.timedelta): + return int(value.total_seconds() * 1_000_000) + return value + + return process + + +class CustomYqlTypeCompiler(ydb_sa.YqlTypeCompiler): + def visit_DATETIME(self, type_: sa.DATETIME, **kw: typing.Any) -> typing.Any: + return self.visit_datetime(type_, **kw) + + def visit_datetime(self, type_: sa.DateTime, **kw: typing.Any) -> typing.Any: + return "DateTime" + + def visit_INTEGER(self, type_: sa.INTEGER, **kw: typing.Any) -> typing.Any: + return self.visit_integer(type_, **kw) + + def visit_integer(self, type_: sa.Integer, **kw: typing.Any) -> typing.Any: + return "int32" + + def visit_SMALLINT(self, type_: sa.SMALLINT, **kw: typing.Any) -> typing.Any: + return self.visit_small_integer(type_, **kw) + + def visit_small_integer(self, type_: sa.SmallInteger, **kw: typing.Any) -> typing.Any: + return "int16" + + def visit_BIGINT(self, type_: sa.BIGINT, **kw: typing.Any) -> typing.Any: + return self.visit_big_integer(type_, **kw) + + def visit_big_integer(self, type_: sa.BigInteger, **kw: typing.Any) -> typing.Any: + return "int64" + + def get_ydb_type( + self, type_: sa.types.TypeEngine, is_optional: bool + ) -> typing.Union[ydb_sa.ydb.PrimitiveType, ydb_sa.ydb.AbstractTypeBuilder]: + if isinstance(type_, sa.TypeDecorator): + type_ = type_.impl + + # Datetime -> Datetime + if isinstance(type_, sa.DATETIME): + return ydb_sa.ydb.PrimitiveType.Datetime + if isinstance(type_, sa.DateTime): + return ydb_sa.ydb.PrimitiveType.Datetime + if isinstance(type_, sa.TIMESTAMP): + return ydb_sa.ydb.PrimitiveType.Timestamp + + # Integer -> Int32 + if isinstance(type_, sa.Integer): + return ydb_sa.ydb.PrimitiveType.Int32 + + # SmallInteger -> Int16 + if isinstance(type_, sa.SMALLINT): + return ydb_sa.ydb.PrimitiveType.Int16 + if isinstance(type_, sa.SmallInteger): + return ydb_sa.ydb.PrimitiveType.Int16 + + # BigInteger -> Int64 + if isinstance(type_, sa.BIGINT): + return ydb_sa.ydb.PrimitiveType.Int64 + if isinstance(type_, sa.BigInteger): + return ydb_sa.ydb.PrimitiveType.Int64 + + return super().get_ydb_type(type_, is_optional) + + +class CustomYqlCompiler(ydb_sa.YqlCompiler): + _type_compiler_cls = CustomYqlTypeCompiler + + +class CustomYqlDialect(ydb_sa.YqlDialect): + type_compiler = CustomYqlTypeCompiler + statement_compiler = CustomYqlCompiler + + colspecs = { + **ydb_sa.YqlDialect.colspecs, + **{ + sa.types.INTEGER: ydb_sa.types.Int32, + sa.types.Integer: ydb_sa.types.Int32, + sa.types.BIGINT: ydb_sa.types.Int64, + sa.types.BigInteger: ydb_sa.types.Int64, + sa.types.SMALLINT: ydb_sa.types.Int16, + sa.types.SmallInteger: ydb_sa.types.Int16, + sa.types.DateTime: YqlDateTime, + sa.types.DATETIME: YqlDateTime, + sa.types.TIMESTAMP: YqlTimestamp, + sa.types.Interval: YqlInterval, + }, + } + + def __init__(self, *args: typing.Any, **kwargs: typing.Any): + super().__init__( + self, + *args, + **{ + **kwargs, + **dict(_add_declare_for_yql_stmt_vars=True), + }, + ) + + +class CustomAsyncYqlDialect(CustomYqlDialect): + driver = "ydb_async" + is_async = True + supports_statement_cache = True + + def connect(self, *cargs: typing.Any, **cparams: typing.Any) -> ydb_sa.AdaptedAsyncConnection: + return ydb_sa.AdaptedAsyncConnection(ydb_sa.util.await_only(self.dbapi.async_connect(*cargs, **cparams))) + + +def register_dialect() -> None: + from sqlalchemy.dialects import registry + + registry.register("yql.ydb", "dl_sqlalchemy_ydb.dialect", "CustomYqlDialect") + registry.register("ydb", "dl_sqlalchemy_ydb.dialect", "CustomYqlDialect") + registry.register("yql", "dl_sqlalchemy_ydb.dialect", "CustomYqlDialect") + registry.register("yql.ydb_async", "dl_sqlalchemy_ydb.dialect", "CustomAsyncYqlDialect") + registry.register("ydb_async", "dl_sqlalchemy_ydb.dialect", "CustomAsyncYqlDialect") diff --git a/lib/dl_sqlalchemy_ydb/dl_sqlalchemy_ydb/py.typed b/lib/dl_sqlalchemy_ydb/dl_sqlalchemy_ydb/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lib/dl_sqlalchemy_ydb/pyproject.toml b/lib/dl_sqlalchemy_ydb/pyproject.toml new file mode 100644 index 0000000000..c9f76b21c1 --- /dev/null +++ b/lib/dl_sqlalchemy_ydb/pyproject.toml @@ -0,0 +1,48 @@ +[tool.poetry] +authors = ["DataLens Team "] +description = "BI YDB SQLAlchemy Dialect" +license = "Apache 2.0" +name = "dl-sqlalchemy-ydb" +packages = [{include = "dl_sqlalchemy_ydb"}] +readme = "README.rst" +version = "0.1" + +[tool.poetry.dependencies] +packaging = "*" +python = ">=3.10, <3.13" +sqlalchemy = "*" +ydb = "*" +ydb-sqlalchemy = "*" + +[tool.poetry.plugins."sqlalchemy.dialects"] +bi_ydb = "dl_sqlalchemy_ydb.dialect:CustomYqlDialect" + +[tool.deptry.per_rule_ignores] +DEP001 = ["google"] +DEP002 = [ + "dl-sqlalchemy-common", +] + +[tool.mypy] +check_untyped_defs = true +disallow_untyped_defs = true +ignore_errors = true # ignore mypy in entire module +strict_optional = true +warn_redundant_casts = true +warn_unused_configs = true +warn_unused_ignores = true + +[tool.pytest.ini_options] +addopts = "-ra" +minversion = "6.0" +testpaths = [] + +[build-system] +build-backend = "poetry.core.masonry.api" +requires = ["poetry-core"] + +[datalens.meta.mypy] +targets = [] # ignore mypy in entire module + +[datalens_ci] +skip_test = true diff --git a/metapkg/poetry.lock b/metapkg/poetry.lock index 62bda71723..d5c6d5cdbb 100644 --- a/metapkg/poetry.lock +++ b/metapkg/poetry.lock @@ -2413,6 +2413,7 @@ dl-formula = {path = "../dl_formula"} dl-formula-ref = {path = "../dl_formula_ref"} dl-i18n = {path = "../dl_i18n"} dl-query-processing = {path = "../dl_query_processing"} +dl-sqlalchemy-ydb = {path = "../../lib/dl_sqlalchemy_ydb"} dl-type-transformer = {path = "../dl_type_transformer"} dl-utils = {path = "../dl_utils"} grpcio = "*" @@ -2421,6 +2422,8 @@ shortuuid = "*" sqlalchemy = "*" typing-extensions = "*" ydb = "*" +ydb-dbapi = "*" +ydb-sqlalchemy = "*" [package.source] type = "directory" @@ -2903,7 +2906,9 @@ attrs = "*" dl-configs = {path = "../dl_configs"} dl-pydantic = {path = "../dl_pydantic"} dl-retrier = {path = "../dl_retrier"} +dl-testing = {path = "../dl_testing"} httpx = "*" +pydantic = "*" typing-extensions = "*" [package.source] @@ -3043,6 +3048,7 @@ develop = true [package.dependencies] pydantic = "*" +typing-extensions = "*" [package.source] type = "directory" @@ -3352,6 +3358,25 @@ sqlalchemy = "*" type = "directory" url = "../lib/dl_sqlalchemy_promql" +[[package]] +name = "dl-sqlalchemy-ydb" +version = "0.1" +description = "BI YDB SQLAlchemy Dialect" +optional = false +python-versions = ">=3.10, <3.13" +files = [] +develop = false + +[package.dependencies] +packaging = "*" +sqlalchemy = "*" +ydb = "*" +ydb-sqlalchemy = "*" + +[package.source] +type = "directory" +url = "../lib/dl_sqlalchemy_ydb" + [[package]] name = "dl-task-processor" version = "0.0.0" @@ -8894,20 +8919,53 @@ propcache = ">=0.2.0" [[package]] name = "ydb" -version = "3.5.2" +version = "3.21.12" description = "YDB Python SDK" optional = false python-versions = "*" files = [ - {file = "ydb-3.5.2-py2.py3-none-any.whl", hash = "sha256:0ef3ae929c8267e18ce56a6a8979b3712a953fe2254e12909aa422b89637902a"}, - {file = "ydb-3.5.2.tar.gz", hash = "sha256:7e698843468a81976e5dd8b3d4324e02bf41a763a18eab2648d4401c95481c67"}, + {file = "ydb-3.21.12-py2.py3-none-any.whl", hash = "sha256:a1c685a9f88b0b6407c4c0a4803ee07f527cb7c17c346de9935dee920ab40122"}, + {file = "ydb-3.21.12.tar.gz", hash = "sha256:612b1a81ab011b40120f187a249b919e85cfbea07de0db7b7664fbd566b83bab"}, ] [package.dependencies] aiohttp = "<4" grpcio = ">=1.42.0" packaging = "*" -protobuf = ">=3.13.0,<5.0.0" +protobuf = ">=3.13.0,<6.0.0" + +[package.extras] +yc = ["yandexcloud"] + +[[package]] +name = "ydb-dbapi" +version = "0.1.14" +description = "YDB Python DBAPI which complies with PEP 249" +optional = false +python-versions = "<4.0,>=3.8" +files = [ + {file = "ydb_dbapi-0.1.14-py3-none-any.whl", hash = "sha256:f8f144ce98e2f2c0378ea389656a3faafb73cea93840da7e6f8a81ccb0b1f775"}, + {file = "ydb_dbapi-0.1.14.tar.gz", hash = "sha256:17b17ad7d4c9fecfd8018fc43c7b4b143e7bba02d98d36dd53640087149e299b"}, +] + +[package.dependencies] +ydb = ">=3.21.6,<4.0.0" + +[[package]] +name = "ydb-sqlalchemy" +version = "0.1.9" +description = "YDB Dialect for SQLAlchemy" +optional = false +python-versions = "*" +files = [ + {file = "ydb_sqlalchemy-0.1.9-py2.py3-none-any.whl", hash = "sha256:e9149e4ad98758a1a6c5290bda4a5f595a14e1cef1ee65e76c92fc45edc688aa"}, + {file = "ydb_sqlalchemy-0.1.9.tar.gz", hash = "sha256:8bfd18ffaf8f6945e489ab57fe49950ae2e25c51f992484c7c3ce9e56a224fd2"}, +] + +[package.dependencies] +sqlalchemy = ">=1.4.0,<3.0.0" +ydb = ">=3.18.8" +ydb-dbapi = ">=0.1.10" [package.extras] yc = ["yandexcloud"] @@ -8934,4 +8992,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.0" python-versions = ">=3.10, <3.13" -content-hash = "21151d834ee8178be8c0ecc0da5cf0ebdd3c2e5416f2d5f4542a1f8fd083664d" +content-hash = "97ae55bc87538f91fb5a59f59207c6ad0929bc2eabcdf79f10213602023e9842" diff --git a/metapkg/pyproject.toml b/metapkg/pyproject.toml index 5c00846647..92be059125 100644 --- a/metapkg/pyproject.toml +++ b/metapkg/pyproject.toml @@ -87,6 +87,9 @@ urllib3 = "==1.26.20" werkzeug = "==3.0.6" xxhash = "==3.2.0" yarl = "==1.18.3" +ydb = ">=3.21.11" +ydb-dbapi = ">=0.1.13" +ydb-sqlalchemy = ">=0.1.9" zipp = "==3.21.0" [tool.poetry.group.app_dl_os_control_api.dependencies]