From 0f8546b71ec1b4da2535b780c41a8b89806770af Mon Sep 17 00:00:00 2001 From: Ivan Ilchuk Date: Tue, 19 Dec 2023 15:42:28 +0100 Subject: [PATCH 1/7] add table settings support --- docker-compose.yml | 1 + setup.py | 1 + test/test_core.py | 135 +++++++++++++++++++++++++- ydb_sqlalchemy/sqlalchemy/__init__.py | 45 ++++++++- 4 files changed, 179 insertions(+), 3 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index 13badc8..04f09b0 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -5,6 +5,7 @@ services: restart: always ports: - "2136:2136" + - "8765:8765" hostname: localhost environment: - YDB_USE_IN_MEMORY_PDISKS=true diff --git a/setup.py b/setup.py index 0ed5e39..15ef571 100644 --- a/setup.py +++ b/setup.py @@ -37,6 +37,7 @@ entry_points={ "sqlalchemy.dialects": [ "yql.ydb=ydb_sqlalchemy.sqlalchemy:YqlDialect", + "yql=ydb_sqlalchemy.sqlalchemy:YqlDialect", ] }, ) diff --git a/test/test_core.py b/test/test_core.py index 52c61e1..e1377dc 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -1,11 +1,14 @@ from decimal import Decimal +from datetime import date, datetime +import pytest import sqlalchemy as sa from sqlalchemy import Table, Column, Integer, Unicode - from sqlalchemy.testing.fixtures import TestBase, TablesTest -from datetime import date, datetime +import ydb + +from ydb_sqlalchemy.sqlalchemy import types def clear_sql(stm): @@ -200,3 +203,131 @@ def test_select_types(self, connection): row = connection.execute(sa.select(tb)).fetchone() assert row == (1, "Hello World!", 3.5, True, now, today) + + +class TestWithClause(TablesTest): + run_create_tables = "each" + + @staticmethod + def _create_table_and_get_desc(connection, metadata, **kwargs): + table = Table( + "clause_with_test", + metadata, + Column("id", types.UInt32, primary_key=True), + **kwargs, + ) + table.create(connection) + + session: ydb.Session = connection.connection.driver_connection.pool.acquire() + return session.describe_table("/local/" + table.name) + + @pytest.mark.parametrize("auto_partitioning_by_size,res", [ + (None, 1), + (True, 1), + (False, 2), + ]) + def test_auto_partitioning_by_size(self, connection, auto_partitioning_by_size, res, metadata): + desc = self._create_table_and_get_desc( + connection, + metadata, + yql_auto_partitioning_by_size=auto_partitioning_by_size + ) + assert desc.partitioning_settings.partitioning_by_size == res + + @pytest.mark.parametrize("auto_partitioning_by_load,res", [ + (None, 2), + (True, 1), + (False, 2), + ]) + def test_auto_partitioning_by_load(self, connection, auto_partitioning_by_load, res, metadata): + desc = self._create_table_and_get_desc( + connection, + metadata, + yql_auto_partitioning_by_load=auto_partitioning_by_load, + ) + assert desc.partitioning_settings.partitioning_by_load == res + + @pytest.mark.parametrize("auto_partitioning_partition_size_mb,res", [ + (None, 2048), + (2000, 2000), + ]) + def test_auto_partitioning_partition_size_mb(self, connection, auto_partitioning_partition_size_mb, res, metadata): + desc = self._create_table_and_get_desc( + connection, + metadata, + yql_auto_partitioning_partition_size_mb=auto_partitioning_partition_size_mb, + ) + assert desc.partitioning_settings.partition_size_mb == res + + @pytest.mark.parametrize("auto_partitioning_min_partitions_count,res", [ + (None, 1), + (10, 10), + ]) + def test_auto_partitioning_min_partitions_count( + self, + connection, + auto_partitioning_min_partitions_count, + res, + metadata, + ): + desc = self._create_table_and_get_desc( + connection, + metadata, + yql_auto_partitioning_min_partitions_count=auto_partitioning_min_partitions_count, + ) + assert desc.partitioning_settings.min_partitions_count == res + + @pytest.mark.parametrize("auto_partitioning_max_partitions_count,res", [ + (None, 0), + (10, 10), + ]) + def test_auto_partitioning_max_partitions_count( + self, + connection, + auto_partitioning_max_partitions_count, + res, + metadata, + ): + desc = self._create_table_and_get_desc( + connection, + metadata, + yql_auto_partitioning_max_partitions_count=auto_partitioning_max_partitions_count, + ) + assert desc.partitioning_settings.max_partitions_count == res + + @pytest.mark.parametrize("uniform_partitions,res", [ + (None, 1), + (10, 10), + ]) + def test_uniform_partitions( + self, + connection, + uniform_partitions, + res, + metadata, + ): + desc = self._create_table_and_get_desc( + connection, + metadata, + yql_uniform_partitions=uniform_partitions, + ) + # it not only do the initiation partition but also set up the minimum partition count + assert desc.partitioning_settings.min_partitions_count == res + + @pytest.mark.parametrize("partition_at_keys,res", [ + (None, 1), + ((100, 1000), 3), + ]) + def test_partition_at_keys( + self, + connection, + partition_at_keys, + res, + metadata, + ): + desc = self._create_table_and_get_desc( + connection, + metadata, + yql_partition_at_keys=partition_at_keys, + ) + assert desc.partitioning_settings.min_partitions_count == res diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index 742e356..21ff706 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -343,7 +343,35 @@ def get_bind_types( class YqlDDLCompiler(DDLCompiler): - pass + def post_create_table(self, table): + yql_opts = table.dialect_options["yql"] + with_content = [] + if yql_opts["auto_partitioning_by_size"] is not None: + auto_partitioning_by_size = "ENABLED" if yql_opts["auto_partitioning_by_size"] else "DISABLED" + with_content.append(f"AUTO_PARTITIONING_BY_SIZE = {auto_partitioning_by_size}") + if yql_opts["auto_partitioning_by_load"] is not None: + auto_partitioning_by_load = "ENABLED" if yql_opts["auto_partitioning_by_load"] else "DISABLED" + with_content.append(f"AUTO_PARTITIONING_BY_LOAD = {auto_partitioning_by_load}") + if yql_opts["auto_partitioning_partition_size_mb"] is not None: + with_content.append( + f"AUTO_PARTITIONING_PARTITION_SIZE_MB = {yql_opts['auto_partitioning_partition_size_mb']}" + ) + if yql_opts["auto_partitioning_min_partitions_count"] is not None: + with_content.append( + f"AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = {yql_opts['auto_partitioning_min_partitions_count']}" + ) + if yql_opts["auto_partitioning_max_partitions_count"] is not None: + with_content.append( + f"AUTO_PARTITIONING_MAX_PARTITIONS_COUNT = {yql_opts['auto_partitioning_max_partitions_count']}" + ) + if yql_opts["uniform_partitions"] is not None: + with_content.append(f"UNIFORM_PARTITIONS = {yql_opts['uniform_partitions']}") + if yql_opts["partition_at_keys"] is not None: + with_content.append(f"PARTITION_AT_KEYS = {yql_opts['partition_at_keys']}") + if with_content: + with_content = ",\n".join(with_content) + return f"\nWITH (\n\t{with_content}\n)" + return "" def upsert(table): @@ -425,6 +453,21 @@ class YqlDialect(StrCompileDialect): ddl_compiler = YqlDDLCompiler type_compiler = YqlTypeCompiler + construct_arguments = [ + ( + sa.schema.Table, + { + "auto_partitioning_by_size": None, + "auto_partitioning_by_load": None, + "auto_partitioning_partition_size_mb": None, + "auto_partitioning_min_partitions_count": None, + "auto_partitioning_max_partitions_count": None, + "uniform_partitions": None, + "partition_at_keys": None, + }, + ), + ] + @classmethod def import_dbapi(cls: Any): return dbapi From b68f9697a46127bce934c4eadf42bef58df3af4c Mon Sep 17 00:00:00 2001 From: Ivan Ilchuk Date: Tue, 19 Dec 2023 15:43:27 +0100 Subject: [PATCH 2/7] ooops --- docker-compose.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/docker-compose.yml b/docker-compose.yml index 04f09b0..13badc8 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -5,7 +5,6 @@ services: restart: always ports: - "2136:2136" - - "8765:8765" hostname: localhost environment: - YDB_USE_IN_MEMORY_PDISKS=true From 293d9b85e167f2e78d55cfcffa9f251f548bb05c Mon Sep 17 00:00:00 2001 From: Ivan Ilchuk Date: Tue, 19 Dec 2023 16:19:40 +0100 Subject: [PATCH 3/7] fix tests and style --- setup.py | 6 ++- test/test_core.py | 85 +++++++++++++++++++++++++---------------- wait_container_ready.py | 1 + 3 files changed, 57 insertions(+), 35 deletions(-) diff --git a/setup.py b/setup.py index 15ef571..8cb753a 100644 --- a/setup.py +++ b/setup.py @@ -21,7 +21,7 @@ license="Apache 2.0", package_dir={"": "."}, long_description=long_description, - long_description_content_type='text/markdown', + long_description_content_type="text/markdown", packages=setuptools.find_packages("."), classifiers=[ "Programming Language :: Python", @@ -32,7 +32,9 @@ install_requires=requirements, # requirements.txt options={"bdist_wheel": {"universal": True}}, extras_require={ - "yc": ["yandexcloud", ], + "yc": [ + "yandexcloud", + ], }, entry_points={ "sqlalchemy.dialects": [ diff --git a/test/test_core.py b/test/test_core.py index e1377dc..3f1af73 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -221,24 +221,28 @@ def _create_table_and_get_desc(connection, metadata, **kwargs): session: ydb.Session = connection.connection.driver_connection.pool.acquire() return session.describe_table("/local/" + table.name) - @pytest.mark.parametrize("auto_partitioning_by_size,res", [ - (None, 1), - (True, 1), - (False, 2), - ]) + @pytest.mark.parametrize( + "auto_partitioning_by_size,res", + [ + (None, 1), + (True, 1), + (False, 2), + ], + ) def test_auto_partitioning_by_size(self, connection, auto_partitioning_by_size, res, metadata): desc = self._create_table_and_get_desc( - connection, - metadata, - yql_auto_partitioning_by_size=auto_partitioning_by_size + connection, metadata, yql_auto_partitioning_by_size=auto_partitioning_by_size ) assert desc.partitioning_settings.partitioning_by_size == res - @pytest.mark.parametrize("auto_partitioning_by_load,res", [ - (None, 2), - (True, 1), - (False, 2), - ]) + @pytest.mark.parametrize( + "auto_partitioning_by_load,res", + [ + (None, 2), + (True, 1), + (False, 2), + ], + ) def test_auto_partitioning_by_load(self, connection, auto_partitioning_by_load, res, metadata): desc = self._create_table_and_get_desc( connection, @@ -247,10 +251,13 @@ def test_auto_partitioning_by_load(self, connection, auto_partitioning_by_load, ) assert desc.partitioning_settings.partitioning_by_load == res - @pytest.mark.parametrize("auto_partitioning_partition_size_mb,res", [ - (None, 2048), - (2000, 2000), - ]) + @pytest.mark.parametrize( + "auto_partitioning_partition_size_mb,res", + [ + (None, 2048), + (2000, 2000), + ], + ) def test_auto_partitioning_partition_size_mb(self, connection, auto_partitioning_partition_size_mb, res, metadata): desc = self._create_table_and_get_desc( connection, @@ -259,10 +266,13 @@ def test_auto_partitioning_partition_size_mb(self, connection, auto_partitioning ) assert desc.partitioning_settings.partition_size_mb == res - @pytest.mark.parametrize("auto_partitioning_min_partitions_count,res", [ - (None, 1), - (10, 10), - ]) + @pytest.mark.parametrize( + "auto_partitioning_min_partitions_count,res", + [ + (None, 1), + (10, 10), + ], + ) def test_auto_partitioning_min_partitions_count( self, connection, @@ -277,10 +287,13 @@ def test_auto_partitioning_min_partitions_count( ) assert desc.partitioning_settings.min_partitions_count == res - @pytest.mark.parametrize("auto_partitioning_max_partitions_count,res", [ - (None, 0), - (10, 10), - ]) + @pytest.mark.parametrize( + "auto_partitioning_max_partitions_count,res", + [ + (None, 0), + (10, 10), + ], + ) def test_auto_partitioning_max_partitions_count( self, connection, @@ -295,10 +308,13 @@ def test_auto_partitioning_max_partitions_count( ) assert desc.partitioning_settings.max_partitions_count == res - @pytest.mark.parametrize("uniform_partitions,res", [ - (None, 1), - (10, 10), - ]) + @pytest.mark.parametrize( + "uniform_partitions,res", + [ + (None, 1), + (10, 10), + ], + ) def test_uniform_partitions( self, connection, @@ -314,10 +330,13 @@ def test_uniform_partitions( # it not only do the initiation partition but also set up the minimum partition count assert desc.partitioning_settings.min_partitions_count == res - @pytest.mark.parametrize("partition_at_keys,res", [ - (None, 1), - ((100, 1000), 3), - ]) + @pytest.mark.parametrize( + "partition_at_keys,res", + [ + (None, 1), + ((100, 1000), 3), + ], + ) def test_partition_at_keys( self, connection, diff --git a/wait_container_ready.py b/wait_container_ready.py index 3a9ceeb..1ac62bd 100755 --- a/wait_container_ready.py +++ b/wait_container_ready.py @@ -3,6 +3,7 @@ import ydb import time + def wait_container_ready(driver): driver.wait(timeout=30) From 562e367b960157955d20ed984a645903076f5af1 Mon Sep 17 00:00:00 2001 From: Ivan Ilchuk Date: Tue, 19 Dec 2023 16:52:21 +0100 Subject: [PATCH 4/7] fix --- setup.py | 2 +- test/conftest.py | 1 + test/test_core.py | 14 ++++++------- ydb_sqlalchemy/sqlalchemy/__init__.py | 30 +++++++++++++-------------- 4 files changed, 24 insertions(+), 23 deletions(-) diff --git a/setup.py b/setup.py index 8cb753a..87e7555 100644 --- a/setup.py +++ b/setup.py @@ -39,7 +39,7 @@ entry_points={ "sqlalchemy.dialects": [ "yql.ydb=ydb_sqlalchemy.sqlalchemy:YqlDialect", - "yql=ydb_sqlalchemy.sqlalchemy:YqlDialect", + "ydb=ydb_sqlalchemy.sqlalchemy:YqlDialect", ] }, ) diff --git a/test/conftest.py b/test/conftest.py index 6c4e4c3..615796c 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -2,6 +2,7 @@ from sqlalchemy.dialects import registry registry.register("yql.ydb", "ydb_sqlalchemy.sqlalchemy", "YqlDialect") +registry.register("ydb", "ydb_sqlalchemy.sqlalchemy", "YqlDialect") pytest.register_assert_rewrite("sqlalchemy.testing.assertions") from sqlalchemy.testing.plugin.pytestplugin import * # noqa: E402, F401, F403 diff --git a/test/test_core.py b/test/test_core.py index 3f1af73..7300ae9 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -231,7 +231,7 @@ def _create_table_and_get_desc(connection, metadata, **kwargs): ) def test_auto_partitioning_by_size(self, connection, auto_partitioning_by_size, res, metadata): desc = self._create_table_and_get_desc( - connection, metadata, yql_auto_partitioning_by_size=auto_partitioning_by_size + connection, metadata, ydb_auto_partitioning_by_size=auto_partitioning_by_size ) assert desc.partitioning_settings.partitioning_by_size == res @@ -247,7 +247,7 @@ def test_auto_partitioning_by_load(self, connection, auto_partitioning_by_load, desc = self._create_table_and_get_desc( connection, metadata, - yql_auto_partitioning_by_load=auto_partitioning_by_load, + ydb_auto_partitioning_by_load=auto_partitioning_by_load, ) assert desc.partitioning_settings.partitioning_by_load == res @@ -262,7 +262,7 @@ def test_auto_partitioning_partition_size_mb(self, connection, auto_partitioning desc = self._create_table_and_get_desc( connection, metadata, - yql_auto_partitioning_partition_size_mb=auto_partitioning_partition_size_mb, + ydb_auto_partitioning_partition_size_mb=auto_partitioning_partition_size_mb, ) assert desc.partitioning_settings.partition_size_mb == res @@ -283,7 +283,7 @@ def test_auto_partitioning_min_partitions_count( desc = self._create_table_and_get_desc( connection, metadata, - yql_auto_partitioning_min_partitions_count=auto_partitioning_min_partitions_count, + ydb_auto_partitioning_min_partitions_count=auto_partitioning_min_partitions_count, ) assert desc.partitioning_settings.min_partitions_count == res @@ -304,7 +304,7 @@ def test_auto_partitioning_max_partitions_count( desc = self._create_table_and_get_desc( connection, metadata, - yql_auto_partitioning_max_partitions_count=auto_partitioning_max_partitions_count, + ydb_auto_partitioning_max_partitions_count=auto_partitioning_max_partitions_count, ) assert desc.partitioning_settings.max_partitions_count == res @@ -325,7 +325,7 @@ def test_uniform_partitions( desc = self._create_table_and_get_desc( connection, metadata, - yql_uniform_partitions=uniform_partitions, + ydb_uniform_partitions=uniform_partitions, ) # it not only do the initiation partition but also set up the minimum partition count assert desc.partitioning_settings.min_partitions_count == res @@ -347,6 +347,6 @@ def test_partition_at_keys( desc = self._create_table_and_get_desc( connection, metadata, - yql_partition_at_keys=partition_at_keys, + ydb_partition_at_keys=partition_at_keys, ) assert desc.partitioning_settings.min_partitions_count == res diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index 21ff706..066c5b0 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -344,30 +344,30 @@ def get_bind_types( class YqlDDLCompiler(DDLCompiler): def post_create_table(self, table): - yql_opts = table.dialect_options["yql"] + ydb_opts = table.dialect_options["ydb"] with_content = [] - if yql_opts["auto_partitioning_by_size"] is not None: - auto_partitioning_by_size = "ENABLED" if yql_opts["auto_partitioning_by_size"] else "DISABLED" + if ydb_opts["auto_partitioning_by_size"] is not None: + auto_partitioning_by_size = "ENABLED" if ydb_opts["auto_partitioning_by_size"] else "DISABLED" with_content.append(f"AUTO_PARTITIONING_BY_SIZE = {auto_partitioning_by_size}") - if yql_opts["auto_partitioning_by_load"] is not None: - auto_partitioning_by_load = "ENABLED" if yql_opts["auto_partitioning_by_load"] else "DISABLED" + if ydb_opts["auto_partitioning_by_load"] is not None: + auto_partitioning_by_load = "ENABLED" if ydb_opts["auto_partitioning_by_load"] else "DISABLED" with_content.append(f"AUTO_PARTITIONING_BY_LOAD = {auto_partitioning_by_load}") - if yql_opts["auto_partitioning_partition_size_mb"] is not None: + if ydb_opts["auto_partitioning_partition_size_mb"] is not None: with_content.append( - f"AUTO_PARTITIONING_PARTITION_SIZE_MB = {yql_opts['auto_partitioning_partition_size_mb']}" + f"AUTO_PARTITIONING_PARTITION_SIZE_MB = {ydb_opts['auto_partitioning_partition_size_mb']}" ) - if yql_opts["auto_partitioning_min_partitions_count"] is not None: + if ydb_opts["auto_partitioning_min_partitions_count"] is not None: with_content.append( - f"AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = {yql_opts['auto_partitioning_min_partitions_count']}" + f"AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = {ydb_opts['auto_partitioning_min_partitions_count']}" ) - if yql_opts["auto_partitioning_max_partitions_count"] is not None: + if ydb_opts["auto_partitioning_max_partitions_count"] is not None: with_content.append( - f"AUTO_PARTITIONING_MAX_PARTITIONS_COUNT = {yql_opts['auto_partitioning_max_partitions_count']}" + f"AUTO_PARTITIONING_MAX_PARTITIONS_COUNT = {ydb_opts['auto_partitioning_max_partitions_count']}" ) - if yql_opts["uniform_partitions"] is not None: - with_content.append(f"UNIFORM_PARTITIONS = {yql_opts['uniform_partitions']}") - if yql_opts["partition_at_keys"] is not None: - with_content.append(f"PARTITION_AT_KEYS = {yql_opts['partition_at_keys']}") + if ydb_opts["uniform_partitions"] is not None: + with_content.append(f"UNIFORM_PARTITIONS = {ydb_opts['uniform_partitions']}") + if ydb_opts["partition_at_keys"] is not None: + with_content.append(f"PARTITION_AT_KEYS = {ydb_opts['partition_at_keys']}") if with_content: with_content = ",\n".join(with_content) return f"\nWITH (\n\t{with_content}\n)" From f1ae3116dcfbf6eda9e66279ef4246b38c2f1c67 Mon Sep 17 00:00:00 2001 From: Ivan Ilchuk Date: Wed, 20 Dec 2023 17:25:46 +0100 Subject: [PATCH 5/7] fixes --- test/test_core.py | 15 +++++++++++++ ydb_sqlalchemy/sqlalchemy/__init__.py | 31 +++++++++++++++------------ 2 files changed, 32 insertions(+), 14 deletions(-) diff --git a/test/test_core.py b/test/test_core.py index 7300ae9..c061fcd 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -350,3 +350,18 @@ def test_partition_at_keys( ydb_partition_at_keys=partition_at_keys, ) assert desc.partitioning_settings.min_partitions_count == res + + def test_several_keys(self, connection, metadata): + desc = self._create_table_and_get_desc( + connection, + metadata, + ydb_auto_partitioning_by_size=True, + ydb_auto_partitioning_by_load=True, + ydb_auto_partitioning_min_partitions_count=3, + ydb_auto_partitioning_max_partitions_count=5, + ) + assert desc.partitioning_settings.partitioning_by_size == 1 + assert desc.partitioning_settings.partitioning_by_load == 1 + assert desc.partitioning_settings.min_partitions_count == 3 + assert desc.partitioning_settings.max_partitions_count == 5 + diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index 066c5b0..5eaa3cb 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -343,36 +343,39 @@ def get_bind_types( class YqlDDLCompiler(DDLCompiler): - def post_create_table(self, table): + def post_create_table(self, table: sa.Table) -> str: ydb_opts = table.dialect_options["ydb"] - with_content = [] + with_clause_list = self._render_table_partitioning_settings(ydb_opts) + if with_clause_list: + with_clause_text = ",\n".join(with_clause_list) + return f"\nWITH (\n\t{with_clause_text}\n)" + return "" + + def _render_table_partitioning_settings(self, ydb_opts: Dict[str, Any]) -> List[str]: + table_partitioning_settings = [] if ydb_opts["auto_partitioning_by_size"] is not None: auto_partitioning_by_size = "ENABLED" if ydb_opts["auto_partitioning_by_size"] else "DISABLED" - with_content.append(f"AUTO_PARTITIONING_BY_SIZE = {auto_partitioning_by_size}") + table_partitioning_settings.append(f"AUTO_PARTITIONING_BY_SIZE = {auto_partitioning_by_size}") if ydb_opts["auto_partitioning_by_load"] is not None: auto_partitioning_by_load = "ENABLED" if ydb_opts["auto_partitioning_by_load"] else "DISABLED" - with_content.append(f"AUTO_PARTITIONING_BY_LOAD = {auto_partitioning_by_load}") + table_partitioning_settings.append(f"AUTO_PARTITIONING_BY_LOAD = {auto_partitioning_by_load}") if ydb_opts["auto_partitioning_partition_size_mb"] is not None: - with_content.append( + table_partitioning_settings.append( f"AUTO_PARTITIONING_PARTITION_SIZE_MB = {ydb_opts['auto_partitioning_partition_size_mb']}" ) if ydb_opts["auto_partitioning_min_partitions_count"] is not None: - with_content.append( + table_partitioning_settings.append( f"AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = {ydb_opts['auto_partitioning_min_partitions_count']}" ) if ydb_opts["auto_partitioning_max_partitions_count"] is not None: - with_content.append( + table_partitioning_settings.append( f"AUTO_PARTITIONING_MAX_PARTITIONS_COUNT = {ydb_opts['auto_partitioning_max_partitions_count']}" ) if ydb_opts["uniform_partitions"] is not None: - with_content.append(f"UNIFORM_PARTITIONS = {ydb_opts['uniform_partitions']}") + table_partitioning_settings.append(f"UNIFORM_PARTITIONS = {ydb_opts['uniform_partitions']}") if ydb_opts["partition_at_keys"] is not None: - with_content.append(f"PARTITION_AT_KEYS = {ydb_opts['partition_at_keys']}") - if with_content: - with_content = ",\n".join(with_content) - return f"\nWITH (\n\t{with_content}\n)" - return "" - + table_partitioning_settings.append(f"PARTITION_AT_KEYS = {ydb_opts['partition_at_keys']}") + return table_partitioning_settings def upsert(table): return sa.sql.Insert(table) From f0d94d1edb7e6afb00899d051adc4a2ce9fe054d Mon Sep 17 00:00:00 2001 From: Ivan Ilchuk Date: Wed, 20 Dec 2023 17:26:55 +0100 Subject: [PATCH 6/7] black --- test/test_core.py | 1 - ydb_sqlalchemy/sqlalchemy/__init__.py | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_core.py b/test/test_core.py index c061fcd..38a9035 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -364,4 +364,3 @@ def test_several_keys(self, connection, metadata): assert desc.partitioning_settings.partitioning_by_load == 1 assert desc.partitioning_settings.min_partitions_count == 3 assert desc.partitioning_settings.max_partitions_count == 5 - diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index 5eaa3cb..9b6b6ef 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -377,6 +377,7 @@ def _render_table_partitioning_settings(self, ydb_opts: Dict[str, Any]) -> List[ table_partitioning_settings.append(f"PARTITION_AT_KEYS = {ydb_opts['partition_at_keys']}") return table_partitioning_settings + def upsert(table): return sa.sql.Insert(table) From 63c604eab412752ac88136845e3a48fa885d9650 Mon Sep 17 00:00:00 2001 From: Ivan Ilchuk Date: Wed, 10 Jan 2024 12:16:39 +0100 Subject: [PATCH 7/7] fix tests --- test/test_core.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/test/test_core.py b/test/test_core.py index 38a9035..8c662ba 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -7,6 +7,7 @@ from sqlalchemy.testing.fixtures import TestBase, TablesTest import ydb +from ydb._grpc.v4.protos import ydb_common_pb2 from ydb_sqlalchemy.sqlalchemy import types @@ -219,14 +220,16 @@ def _create_table_and_get_desc(connection, metadata, **kwargs): table.create(connection) session: ydb.Session = connection.connection.driver_connection.pool.acquire() - return session.describe_table("/local/" + table.name) + table_description = session.describe_table("/local/" + table.name) + session.delete() + return table_description @pytest.mark.parametrize( "auto_partitioning_by_size,res", [ - (None, 1), - (True, 1), - (False, 2), + (None, ydb_common_pb2.FeatureFlag.Status.ENABLED), + (True, ydb_common_pb2.FeatureFlag.Status.ENABLED), + (False, ydb_common_pb2.FeatureFlag.Status.DISABLED), ], ) def test_auto_partitioning_by_size(self, connection, auto_partitioning_by_size, res, metadata): @@ -238,9 +241,9 @@ def test_auto_partitioning_by_size(self, connection, auto_partitioning_by_size, @pytest.mark.parametrize( "auto_partitioning_by_load,res", [ - (None, 2), - (True, 1), - (False, 2), + (None, ydb_common_pb2.FeatureFlag.Status.DISABLED), + (True, ydb_common_pb2.FeatureFlag.Status.ENABLED), + (False, ydb_common_pb2.FeatureFlag.Status.DISABLED), ], ) def test_auto_partitioning_by_load(self, connection, auto_partitioning_by_load, res, metadata):