diff --git a/test/test_core.py b/test/test_core.py index 61129f8..97ee685 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -464,6 +464,29 @@ def test_several_keys(self, connection, metadata): assert desc.partitioning_settings.min_partitions_count == 3 assert desc.partitioning_settings.max_partitions_count == 5 + @pytest.mark.parametrize( + "kw,is_column_table", + [ + ({}, False), + ({"ydb_column_store": False}, False), + ({"ydb_column_store": True}, True), + ], + ) + def test_column_store( + self, + connection, + kw, + is_column_table, + metadata, + ): + desc = self._create_table_and_get_desc( + connection, + metadata, + **kw, + ) + # it not only do the initiation partition but also set up the minimum partition count + assert desc.is_column_table() == is_column_table + class TestTransaction(TablesTest): __backend__ = True diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index d3ef664..1496371 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -503,12 +503,24 @@ def visit_drop_index(self, drop: ddl.DropIndex, **kw) -> str: def post_create_table(self, table: sa.Table) -> str: ydb_opts = table.dialect_options["ydb"] - with_clause_list = self._render_table_partitioning_settings(ydb_opts) + with_clause_list = self._prepare_with_clause(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 _prepare_with_clause(self, ydb_opts: Dict[str, Any]) -> List[str]: + with_clause_list = [] + with_clause_list.extend(self._render_table_partitioning_settings(ydb_opts)) + with_clause_list.extend(self._render_table_store_settings(ydb_opts)) + return with_clause_list + + def _render_table_store_settings(self, ydb_opts: Dict[str, Any]) -> List[str]: + table_store_settings = [] + if ydb_opts["column_store"] is not None and ydb_opts["column_store"]: + table_store_settings.append("STORE = COLUMN") + return table_store_settings + 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: @@ -654,6 +666,7 @@ class YqlDialect(StrCompileDialect): "auto_partitioning_max_partitions_count": None, "uniform_partitions": None, "partition_at_keys": None, + "column_store": None, }, ), (