feat(connectors): BI-6585 Add column tables support for YDB

Author: catsona

lib/dl_connector_ydb/dl_connector_ydb/core/ydb/adapter.py

@@ -117,7 +117,7 @@ def _list_table_names_i(self, db_name: str, show_dot: bool = False) -> Iterable[
                 ]
                 children.sort()
                 for full_path, child in children:
-                    if child.is_any_table():
+                    if child.is_any_table() or child.is_view() or child.is_column_table():
                         yield full_path.removeprefix(unprefix)
                     elif child.is_directory():
                         queue.append(full_path)

lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/base.py

@@ -24,7 +24,9 @@
 )
 from dl_connector_ydb_tests.db.config import (
     API_TEST_CONFIG,
+    COLUMN_TABLE_SCHEMA,
     DB_CORE_URL,
+    SA_TYPE_TO_YDB_TYPE_NAME,
     TABLE_DATA,
     TABLE_NAME,
     TABLE_SCHEMA,
@@ -102,6 +104,41 @@ def dataset_params(self, sample_table: DbTable) -> dict:
         )
 
 
+class YDBColumnDatasetTestBase(YDBDatasetTestBase):
+    @pytest.fixture(scope="class")
+    def sample_table(self, db: Db) -> DbTable:
+        table_name = TABLE_NAME + "_column"
+
+        db_table = make_table(
+            db=db,
+            name=table_name,
+            columns=[
+                C(name=name, user_type=user_type, sa_type=sa_type) for name, user_type, sa_type in COLUMN_TABLE_SCHEMA
+            ],
+            data=[],  # to avoid producing a sample data
+            create_in_db=False,
+        )
+
+        column_definitions_list = []
+        for name, _, sa_type in COLUMN_TABLE_SCHEMA:
+            target_type = SA_TYPE_TO_YDB_TYPE_NAME[sa_type]
+            if name == "id":
+                target_type += " NOT NULL"
+
+            column_definitions_list.append(f"{name} {target_type}")
+        column_definitions = ", ".join(column_definitions_list)
+
+        query = f"CREATE TABLE `{table_name}` ({column_definitions}, PRIMARY KEY (id)) WITH (STORE = COLUMN)"
+        db.get_current_connection().connection.cursor().execute_scheme(query)
+
+        db.create_table(db_table.table)
+        db.insert_into_table(db_table.table, TABLE_DATA)
+
+        yield db_table
+
+        db.get_current_connection().connection.cursor().execute_scheme(f"DROP TABLE `{table_name}`;")
+
+
 class YDBDataApiTestBase(YDBDatasetTestBase, StandardizedDataApiTestBase):
     mutation_caches_enabled = False
 

lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/test_dataset.py

@@ -1,8 +1,14 @@
 from dl_api_client.dsmaker.primitives import Dataset
 from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite
 
-from dl_connector_ydb_tests.db.api.base import YDBDatasetTestBase
-from dl_connector_ydb_tests.db.config import TABLE_SCHEMA
+from dl_connector_ydb_tests.db.api.base import (
+    YDBColumnDatasetTestBase,
+    YDBDatasetTestBase,
+)
+from dl_connector_ydb_tests.db.config import (
+    COLUMN_TABLE_SCHEMA,
+    TABLE_SCHEMA,
+)
 
 
 class TestYDBDataset(YDBDatasetTestBase, DefaultConnectorDatasetTestSuite):
@@ -12,3 +18,10 @@ def check_basic_dataset(self, ds: Dataset, annotation: dict) -> None:
         assert field_names == {column[0] for column in TABLE_SCHEMA}
 
         assert ds.annotation == annotation
+
+
+class TestYDBColumnDataset(YDBColumnDatasetTestBase, DefaultConnectorDatasetTestSuite):
+    def check_basic_dataset(self, ds: Dataset) -> None:
+        assert ds.id
+        field_names = {field.title for field in ds.result_schema}
+        assert field_names == {column[0] for column in COLUMN_TABLE_SCHEMA}

lib/dl_connector_ydb/dl_connector_ydb_tests/db/config.py

@@ -79,6 +79,36 @@ class CoreSslConnectionSettings:
     ("some_timestamp", UserDataType.genericdatetime, sa.TIMESTAMP),
     ("some_interval", UserDataType.integer, dl_sqlalchemy_ydb.dialect.YqlInterval),
 )
+
+COLUMN_TABLE_SCHEMA = (
+    ("id", UserDataType.integer, sa.Integer),
+    ("distinct_string", UserDataType.string, sa.String),
+    ("some_int32", UserDataType.integer, sa.Integer),
+    ("some_int64", UserDataType.integer, sa.BigInteger),
+    ("some_uint8", UserDataType.integer, sa.SmallInteger),
+    ("some_double", UserDataType.float, sa.Float),
+    ("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_timestamp", UserDataType.genericdatetime, sa.TIMESTAMP),
+    ("some_interval", UserDataType.integer, dl_sqlalchemy_ydb.dialect.YqlInterval),
+)
+
+SA_TYPE_TO_YDB_TYPE_NAME = {
+    sa.Integer: "Int32",
+    sa.String: "String",
+    sa.BigInteger: "Int64",
+    sa.SmallInteger: "Int8",
+    sa.Boolean: "Bool",
+    sa.Float: "Double",
+    sa.Unicode: "Utf8",
+    sa.Date: "Date",
+    sa.DATETIME: "Datetime",
+    sa.TIMESTAMP: "Timestamp",
+    dl_sqlalchemy_ydb.dialect.YqlInterval: "Interval",
+}
+
 TABLE_DATA = [
     {
         "id": 1,
@@ -246,8 +276,13 @@ class CoreSslConnectionSettings:
         "some_interval": 1234,
     },
 ]
+
+# Leave only values in COLUMN_TABLE_SCHEMA
+COLUMN_TABLE_DATA = [{key: value for key, value in row.items() if key in COLUMN_TABLE_SCHEMA} for row in TABLE_DATA]
+
 TABLE_NAME = "test_table_h"
 
+
 DASHSQL_QUERY = r"""
 select
     id,

lib/dl_connector_ydb/docker-compose.yml

@@ -9,6 +9,8 @@ services:
       YDB_GRPC_ENABLE_TLS: 1
       GRPC_TLS_PORT: "51902"
       YDB_GRPC_TLS_DATA_PATH: "/ydb_certs"
+      YDB_FEATURE_FLAGS: "enable_views,enable_olap_schema_operations"
+      YDB_ENABLE_COLUMN_TABLES: "true"
     hostname: "db-ydb"
     ports:
       - "51900:51900"

lib/dl_connector_ydb/docker-compose/Dockerfile.db-ydb

@@ -1,4 +1,4 @@
-FROM cr.yandex/yc/yandex-docker-local-ydb:latest
+FROM ydbplatform/local-ydb:latest@sha256:1252f37e5f3fd6c490a8e6c34f927cb8b3dd8323db2aa14ffa21dd09dd6ebcdc
 
 RUN apt update && apt install -y python3
 

lib/dl_type_transformer/dl_type_transformer/type_transformer.py

@@ -72,6 +72,8 @@ def make_datetime(value: Any) -> Optional[datetime.datetime]:
 def make_int(value: Any) -> Optional[int]:
     if value is None:
         return None
+    # if isinstance(value, datetime.timedelta):
+    #     return int(value.total_seconds())
     if isinstance(value, float) and (math.isinf(value) or math.isnan(value)):
         return None
     return int(value)