From 136c0dc1b05b2d7d39281a213e6040c9740bf5c9 Mon Sep 17 00:00:00 2001 From: Roman Tretiak Date: Fri, 5 Apr 2024 11:07:35 +0200 Subject: [PATCH 1/5] Add index creation --- test/test_core.py | 65 +++++++++++++++++++++++++++ ydb_sqlalchemy/sqlalchemy/__init__.py | 38 ++++++++++++++++ 2 files changed, 103 insertions(+) diff --git a/test/test_core.py b/test/test_core.py index 9175b7c..898f029 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -732,3 +732,68 @@ def test_insert_in_name_and_field(self, connection): row = connection.execute(sa.select(tb).where(tb.c.id == 2)).fetchone() assert row == (2, "INSERT is my favourite operation") + + +class TestSecondaryIndex(TestBase): + __backend__ = True + + def test_column_indexes(self, connection: sa.Connection, metadata: sa.MetaData): + table = Table( + "test_column_indexes/table", + metadata, + sa.Column("id", sa.Integer, primary_key=True), + sa.Column("index_col1", sa.Integer, index=True), + sa.Column("index_col2", sa.Integer, index=True), + ) + table.create(connection) + + table_desc: ydb.TableDescription = connection.connection.driver_connection.describe(table.name) + indexes: list[ydb.TableIndex] = table_desc.indexes + assert len(indexes) == 2 + indexes_map = {idx.name: idx for idx in indexes} + + assert "ix_test_column_indexes_table_index_col1" in indexes_map + index1 = indexes_map["ix_test_column_indexes_table_index_col1"] + assert index1.index_columns == ["index_col1"] + + assert "ix_test_column_indexes_table_index_col2" in indexes_map + index1 = indexes_map["ix_test_column_indexes_table_index_col2"] + assert index1.index_columns == ["index_col2"] + + def test_async_index(self, connection: sa.Connection, metadata: sa.MetaData): + table = Table( + "test_async_index/table", + metadata, + sa.Column("id", sa.Integer, primary_key=True), + sa.Column("index_col1", sa.Integer), + sa.Column("index_col2", sa.Integer), + sa.Index("test_async_index", "index_col1", "index_col2", ydb_async=True), + ) + table.create(connection) + + table_desc: ydb.TableDescription = connection.connection.driver_connection.describe(table.name) + indexes: list[ydb.TableIndex] = table_desc.indexes + assert len(indexes) == 1 + index = indexes[0] + assert index.name == "test_async_index" + assert set(index.index_columns) == {"index_col1", "index_col2"} + # TODO: Uncomment after fix https://github.com/ydb-platform/ydb-python-sdk/issues/351 + # assert index.to_pb().WhichOneof("type") == "global_async_index" + + def test_cover_index(self, connection: sa.Connection, metadata: sa.MetaData): + table = Table( + "test_cover_index/table", + metadata, + sa.Column("id", sa.Integer, primary_key=True), + sa.Column("index_col1", sa.Integer), + sa.Column("index_col2", sa.Integer), + sa.Index("test_cover_index", "index_col1", ydb_cover=["index_col2"]), + ) + table.create(connection) + + table_desc: ydb.TableDescription = connection.connection.driver_connection.describe(table.name) + indexes: list[ydb.TableIndex] = table_desc.indexes + assert len(indexes) == 1 + index = indexes[0] + assert index.name == "test_cover_index" + assert set(index.index_columns) == {"index_col1"} diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index e5e3ca2..1d7d99b 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -62,6 +62,9 @@ def __init__(self, dialect): final_quote="`", ) + def format_index(self, index: sa.Index) -> str: + return super().format_index(index).replace("/", "_") + class YqlTypeCompiler(StrSQLTypeCompiler): def visit_JSON(self, type_: Union[sa.JSON, types.YqlJSON], **kw): @@ -446,6 +449,34 @@ def visit_upsert(self, insert_stmt, visited_bindparam=None, **kw): class YqlDDLCompiler(DDLCompiler): + def visit_create_index(self, create, include_schema=False, include_table_schema=True, **kw): + index: sa.Index = create.element + ydb_opts = index.dialect_options.get("ydb", {}) + + self._verify_index_table(index) + + if index.name is None: + raise CompileError("ADD INDEX requires that the index have a name") + + table_name = self.preparer.format_table(index.table) + index_name = self._prepared_index_name(index) + + text = f"ALTER TABLE {table_name} ADD INDEX {index_name} GLOBAL" + + text += " SYNC " if not ydb_opts.get("async", False) else " ASYNC" + + columns = {self.preparer.format_column(col) for col in index.columns.values()} + cover_columns = { + col if isinstance(col, str) else self.preparer.format_column(col) for col in ydb_opts.get("cover", []) + } + + text += " ON (" + ", ".join(columns) + ")" + + if cover_columns: + text += " COVER (" + ", ".join(cover_columns) + ")" + + return text + 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) @@ -578,6 +609,13 @@ class YqlDialect(StrCompileDialect): "partition_at_keys": None, }, ), + ( + sa.schema.Index, + { + "async": False, + "cover": [], + }, + ), ] @classmethod From 33080206b2270be322494e34fde41fe696dd8e8a Mon Sep 17 00:00:00 2001 From: Roman Tretiak Date: Fri, 5 Apr 2024 14:21:58 +0200 Subject: [PATCH 2/5] Add index reflection --- test/test_core.py | 30 +++++++++++++++++++++-- ydb_sqlalchemy/sqlalchemy/__init__.py | 20 ++++++++++++--- ydb_sqlalchemy/sqlalchemy/requirements.py | 1 + 3 files changed, 46 insertions(+), 5 deletions(-) diff --git a/test/test_core.py b/test/test_core.py index 898f029..e28bf22 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -777,8 +777,7 @@ def test_async_index(self, connection: sa.Connection, metadata: sa.MetaData): index = indexes[0] assert index.name == "test_async_index" assert set(index.index_columns) == {"index_col1", "index_col2"} - # TODO: Uncomment after fix https://github.com/ydb-platform/ydb-python-sdk/issues/351 - # assert index.to_pb().WhichOneof("type") == "global_async_index" + # TODO: Check type after https://github.com/ydb-platform/ydb-python-sdk/issues/351 def test_cover_index(self, connection: sa.Connection, metadata: sa.MetaData): table = Table( @@ -797,3 +796,30 @@ def test_cover_index(self, connection: sa.Connection, metadata: sa.MetaData): index = indexes[0] assert index.name == "test_cover_index" assert set(index.index_columns) == {"index_col1"} + # TODO: Check covered columns after https://github.com/ydb-platform/ydb-python-sdk/issues/409 + + def test_indexes_reflection(self, connection: sa.Connection, metadata: sa.MetaData): + table = Table( + "test_indexes_reflection/table", + metadata, + sa.Column("id", sa.Integer, primary_key=True), + sa.Column("index_col1", sa.Integer, index=True), + sa.Column("index_col2", sa.Integer), + sa.Index("test_index", "index_col1", "index_col2"), + sa.Index("test_async_index", "index_col1", "index_col2", ydb_async=True), + sa.Index("test_cover_index", "index_col1", ydb_cover=["index_col2"]), + sa.Index("test_async_cover_index", "index_col1", ydb_async=True, ydb_cover=["index_col2"]), + ) + table.create(connection) + + indexes = sa.inspect(connection).get_indexes(table.name) + assert len(indexes) == 5 + indexes_names = {idx["name"]: set(idx["column_names"]) for idx in indexes} + + assert indexes_names == { + "ix_test_indexes_reflection_table_index_col1": {"index_col1"}, + "test_index": {"index_col1", "index_col2"}, + "test_async_index": {"index_col1", "index_col2"}, + "test_cover_index": {"index_col1"}, + "test_async_cover_index": {"index_col1"}, + } diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index 1d7d99b..c613679 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -628,7 +628,7 @@ def __init__(self, json_serializer=None, json_deserializer=None, **kwargs): self._json_deserializer = json_deserializer self._json_serializer = json_serializer - def _describe_table(self, connection, table_name, schema=None): + def _describe_table(self, connection, table_name, schema=None) -> ydb.TableDescription: if schema is not None: raise dbapi.NotSupportedError("unsupported on non empty schema") @@ -684,8 +684,22 @@ def get_foreign_keys(self, connection, table_name, schema=None, **kwargs): @reflection.cache def get_indexes(self, connection, table_name, schema=None, **kwargs): - # TODO: implement me - return [] + table = self._describe_table(connection, table_name, schema) + indexes: list[ydb.TableIndex] = table.indexes + sa_indexes: list[sa.engine.interfaces.ReflectedIndex] = [] + for index in indexes: + sa_indexes.append( + sa.engine.interfaces.ReflectedIndex( + name=index.name, + column_names=index.index_columns, + unique=False, + dialect_options={ + "ydb_async": False, # TODO After https://github.com/ydb-platform/ydb-python-sdk/issues/351 + "ydb_cover": [], # TODO After https://github.com/ydb-platform/ydb-python-sdk/issues/409 + }, + ) + ) + return sa_indexes def set_isolation_level(self, dbapi_connection: dbapi.Connection, level: str) -> None: dbapi_connection.set_isolation_level(level) diff --git a/ydb_sqlalchemy/sqlalchemy/requirements.py b/ydb_sqlalchemy/sqlalchemy/requirements.py index df1314b..4b0cf70 100644 --- a/ydb_sqlalchemy/sqlalchemy/requirements.py +++ b/ydb_sqlalchemy/sqlalchemy/requirements.py @@ -46,6 +46,7 @@ def temporary_views(self): @property def index_reflection(self): + # Reflection supported with limits return exclusions.closed() @property From 83fe8e5c0d4176426a0a578ff6c57c377d34f11f Mon Sep 17 00:00:00 2001 From: Roman Tretiak Date: Mon, 8 Apr 2024 16:30:14 +0200 Subject: [PATCH 3/5] Add index hints --- test/test_core.py | 85 +++++++++++++++++++++++++++ ydb_sqlalchemy/sqlalchemy/__init__.py | 3 + 2 files changed, 88 insertions(+) diff --git a/test/test_core.py b/test/test_core.py index e28bf22..fd9d923 100644 --- a/test/test_core.py +++ b/test/test_core.py @@ -823,3 +823,88 @@ def test_indexes_reflection(self, connection: sa.Connection, metadata: sa.MetaDa "test_cover_index": {"index_col1"}, "test_async_cover_index": {"index_col1"}, } + + def test_index_simple_usage(self, connection: sa.Connection, metadata: sa.MetaData): + persons = Table( + "test_index_simple_usage/persons", + metadata, + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column("tax_number", sa.Integer()), + sa.Column("full_name", sa.Unicode()), + sa.Index("ix_tax_number_cover_full_name", "tax_number", ydb_cover=["full_name"]), + ) + persons.create(connection) + connection.execute( + sa.insert(persons).values( + [ + {"id": 1, "tax_number": 333333, "full_name": "John Connor"}, + {"id": 2, "tax_number": 444444, "full_name": "Sarah Connor"}, + ] + ) + ) + + # Because of this bug https://github.com/ydb-platform/ydb/issues/3510, + # it is not possible to use full qualified name of columns with VIEW clause + select_stmt = ( + sa.select(sa.column(persons.c.full_name.name)) + .select_from(persons) + .with_hint(persons, "VIEW `ix_tax_number_cover_full_name`") + .where(sa.column(persons.c.tax_number.name) == 444444) + ) + + cursor = connection.execute(select_stmt) + assert cursor.scalar_one() == "Sarah Connor" + + def test_index_with_join_usage(self, connection: sa.Connection, metadata: sa.MetaData): + persons = Table( + "test_index_with_join_usage/persons", + metadata, + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column("tax_number", sa.Integer()), + sa.Column("full_name", sa.Unicode()), + sa.Index("ix_tax_number_cover_full_name", "tax_number", ydb_cover=["full_name"]), + ) + persons.create(connection) + connection.execute( + sa.insert(persons).values( + [ + {"id": 1, "tax_number": 333333, "full_name": "John Connor"}, + {"id": 2, "tax_number": 444444, "full_name": "Sarah Connor"}, + ] + ) + ) + person_status = Table( + "test_index_with_join_usage/person_status", + metadata, + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column("status", sa.Unicode()), + ) + person_status.create(connection) + connection.execute( + sa.insert(person_status).values( + [ + {"id": 1, "status": "unknown"}, + {"id": 2, "status": "wanted"}, + ] + ) + ) + + # Because of this bug https://github.com/ydb-platform/ydb/issues/3510, + # it is not possible to use full qualified name of columns with VIEW clause + persons_indexed = ( + sa.select( + sa.column(persons.c.id.name), + sa.column(persons.c.full_name.name), + sa.column(persons.c.tax_number.name), + ) + .select_from(persons) + .with_hint(persons, "VIEW `ix_tax_number_cover_full_name`") + ) + select_stmt = ( + sa.select(persons_indexed.c.full_name, person_status.c.status) + .select_from(person_status.join(persons_indexed, persons_indexed.c.id == person_status.c.id)) + .where(persons_indexed.c.tax_number == 444444) + ) + + cursor = connection.execute(select_stmt) + assert cursor.one() == ("Sarah Connor", "wanted") diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index c613679..a0a875d 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -235,6 +235,9 @@ def __init__(self, name, params, *args, **kwargs): class YqlCompiler(StrSQLCompiler): compound_keywords = COMPOUND_KEYWORDS + def get_from_hint_text(self, table, text): + return text + def render_bind_cast(self, type_, dbapi_type, sqltext): pass From b9beeb0b10b4e4da80d6e08ad9d0a05da9d05c38 Mon Sep 17 00:00:00 2001 From: Roman Tretiak Date: Mon, 22 Apr 2024 15:13:51 +0200 Subject: [PATCH 4/5] Remove redundant space --- ydb_sqlalchemy/sqlalchemy/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index 3e35643..14681a9 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -466,7 +466,7 @@ def visit_create_index(self, create, include_schema=False, include_table_schema= text = f"ALTER TABLE {table_name} ADD INDEX {index_name} GLOBAL" - text += " SYNC " if not ydb_opts.get("async", False) else " ASYNC" + text += " SYNC" if not ydb_opts.get("async", False) else " ASYNC" columns = {self.preparer.format_column(col) for col in index.columns.values()} cover_columns = { From e7ee419931d53eaa781c963b5ba5c78565d75ddb Mon Sep 17 00:00:00 2001 From: Roman Tretiak Date: Mon, 22 Apr 2024 15:18:08 +0200 Subject: [PATCH 5/5] Fix typo --- ydb_sqlalchemy/sqlalchemy/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index 14681a9..7a1e2a8 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/ydb_sqlalchemy/sqlalchemy/__init__.py @@ -459,7 +459,7 @@ def visit_create_index(self, create, include_schema=False, include_table_schema= self._verify_index_table(index) if index.name is None: - raise CompileError("ADD INDEX requires that the index have a name") + raise CompileError("ADD INDEX requires that the index has a name") table_name = self.preparer.format_table(index.table) index_name = self._prepared_index_name(index)