Merge pull request #37 Feature: Secondary indexes from LuckySting/feat/ydb-secondary-indicies

rekby · web-flow · commit 50f029a63523 · 2024-04-22T16:23:55.000+03:00
Thanks for the PR :)
diff --git a/test/test_core.py b/test/test_core.py
@@ -788,3 +788,179 @@ 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: 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(
+            "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"}
+        # 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"},
+        }
+
+    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
@@ -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):
@@ -232,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
 
@@ -446,6 +452,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 has 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 +612,13 @@ class YqlDialect(StrCompileDialect):
                 "partition_at_keys": None,
             },
         ),
+        (
+            sa.schema.Index,
+            {
+                "async": False,
+                "cover": [],
+            },
+        ),
     ]
 
     @classmethod
@@ -599,7 +640,7 @@ def __init__(
         # no need in declare in yql statement here since ydb 24-1
         self._add_declare_for_yql_stmt_vars = _add_declare_for_yql_stmt_vars
 
-    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")
 
@@ -655,8 +696,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
@@ -46,6 +46,7 @@ def temporary_views(self):
 
     @property
     def index_reflection(self):
+        # Reflection supported with limits
         return exclusions.closed()
 
     @property
