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 a1961eb37e..3ef5354991 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 @@ -126,7 +126,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() or child.is_view(): + 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) 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 35148ea6c5..ae72af9cee 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 @@ -27,7 +27,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, @@ -135,6 +137,41 @@ def dataset_params( ) +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 diff --git a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/test_dataset.py b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/test_dataset.py index 59b68bea26..631f577399 100644 --- a/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/test_dataset.py +++ b/lib/dl_connector_ydb/dl_connector_ydb_tests/db/api/test_dataset.py @@ -2,10 +2,14 @@ from dl_api_lib_testing.connector.dataset_suite import DefaultConnectorDatasetTestSuite from dl_connector_ydb_tests.db.api.base import ( + YDBColumnDatasetTestBase, YDBDatasetTestBase, YDBViewDatasetTestBase, ) -from dl_connector_ydb_tests.db.config import TABLE_SCHEMA +from dl_connector_ydb_tests.db.config import ( + COLUMN_TABLE_SCHEMA, + TABLE_SCHEMA, +) class TestYDBDataset(YDBDatasetTestBase, DefaultConnectorDatasetTestSuite): @@ -24,3 +28,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, annotation: dict) -> 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} 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 d468bf9e6c..0f275d4160 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 @@ -79,6 +79,35 @@ 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), +) + +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 +275,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, diff --git a/lib/dl_connector_ydb/docker-compose.yml b/lib/dl_connector_ydb/docker-compose.yml index 4f45e7410a..9acc9b298b 100644 --- a/lib/dl_connector_ydb/docker-compose.yml +++ b/lib/dl_connector_ydb/docker-compose.yml @@ -9,7 +9,8 @@ services: YDB_GRPC_ENABLE_TLS: 1 GRPC_TLS_PORT: "51902" YDB_GRPC_TLS_DATA_PATH: "/ydb_certs" - YDB_FEATURE_FLAGS: "enable_views" + YDB_FEATURE_FLAGS: "enable_views,enable_olap_schema_operations" + YDB_ENABLE_COLUMN_TABLES: "true" hostname: "db-ydb" ports: - "51900:51900" diff --git a/lib/dl_connector_ydb/docker-compose/Dockerfile.db-ydb b/lib/dl_connector_ydb/docker-compose/Dockerfile.db-ydb index 6e5191f894..3bdf2a9399 100644 --- a/lib/dl_connector_ydb/docker-compose/Dockerfile.db-ydb +++ b/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