Add index reflection

ydb-platform · Apr 5, 2024 · 3308020 · 3308020
1 parent 136c0dc
commit 3308020
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 5 deletions.
diff --git 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
@@ -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
@@ -46,6 +46,7 @@ def temporary_views(self):
 
     @property
     def index_reflection(self):
+        # Reflection supported with limits
         return exclusions.closed()
 
     @property