Skip to content

Commit 3308020

Browse files
committed
Add index reflection
1 parent 136c0dc commit 3308020

File tree

3 files changed

+46
-5
lines changed

3 files changed

+46
-5
lines changed

test/test_core.py

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -777,8 +777,7 @@ def test_async_index(self, connection: sa.Connection, metadata: sa.MetaData):
777777
index = indexes[0]
778778
assert index.name == "test_async_index"
779779
assert set(index.index_columns) == {"index_col1", "index_col2"}
780-
# TODO: Uncomment after fix https://github.com/ydb-platform/ydb-python-sdk/issues/351
781-
# assert index.to_pb().WhichOneof("type") == "global_async_index"
780+
# TODO: Check type after https://github.com/ydb-platform/ydb-python-sdk/issues/351
782781

783782
def test_cover_index(self, connection: sa.Connection, metadata: sa.MetaData):
784783
table = Table(
@@ -797,3 +796,30 @@ def test_cover_index(self, connection: sa.Connection, metadata: sa.MetaData):
797796
index = indexes[0]
798797
assert index.name == "test_cover_index"
799798
assert set(index.index_columns) == {"index_col1"}
799+
# TODO: Check covered columns after https://github.com/ydb-platform/ydb-python-sdk/issues/409
800+
801+
def test_indexes_reflection(self, connection: sa.Connection, metadata: sa.MetaData):
802+
table = Table(
803+
"test_indexes_reflection/table",
804+
metadata,
805+
sa.Column("id", sa.Integer, primary_key=True),
806+
sa.Column("index_col1", sa.Integer, index=True),
807+
sa.Column("index_col2", sa.Integer),
808+
sa.Index("test_index", "index_col1", "index_col2"),
809+
sa.Index("test_async_index", "index_col1", "index_col2", ydb_async=True),
810+
sa.Index("test_cover_index", "index_col1", ydb_cover=["index_col2"]),
811+
sa.Index("test_async_cover_index", "index_col1", ydb_async=True, ydb_cover=["index_col2"]),
812+
)
813+
table.create(connection)
814+
815+
indexes = sa.inspect(connection).get_indexes(table.name)
816+
assert len(indexes) == 5
817+
indexes_names = {idx["name"]: set(idx["column_names"]) for idx in indexes}
818+
819+
assert indexes_names == {
820+
"ix_test_indexes_reflection_table_index_col1": {"index_col1"},
821+
"test_index": {"index_col1", "index_col2"},
822+
"test_async_index": {"index_col1", "index_col2"},
823+
"test_cover_index": {"index_col1"},
824+
"test_async_cover_index": {"index_col1"},
825+
}

ydb_sqlalchemy/sqlalchemy/__init__.py

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -628,7 +628,7 @@ def __init__(self, json_serializer=None, json_deserializer=None, **kwargs):
628628
self._json_deserializer = json_deserializer
629629
self._json_serializer = json_serializer
630630

631-
def _describe_table(self, connection, table_name, schema=None):
631+
def _describe_table(self, connection, table_name, schema=None) -> ydb.TableDescription:
632632
if schema is not None:
633633
raise dbapi.NotSupportedError("unsupported on non empty schema")
634634

@@ -684,8 +684,22 @@ def get_foreign_keys(self, connection, table_name, schema=None, **kwargs):
684684

685685
@reflection.cache
686686
def get_indexes(self, connection, table_name, schema=None, **kwargs):
687-
# TODO: implement me
688-
return []
687+
table = self._describe_table(connection, table_name, schema)
688+
indexes: list[ydb.TableIndex] = table.indexes
689+
sa_indexes: list[sa.engine.interfaces.ReflectedIndex] = []
690+
for index in indexes:
691+
sa_indexes.append(
692+
sa.engine.interfaces.ReflectedIndex(
693+
name=index.name,
694+
column_names=index.index_columns,
695+
unique=False,
696+
dialect_options={
697+
"ydb_async": False, # TODO After https://github.com/ydb-platform/ydb-python-sdk/issues/351
698+
"ydb_cover": [], # TODO After https://github.com/ydb-platform/ydb-python-sdk/issues/409
699+
},
700+
)
701+
)
702+
return sa_indexes
689703

690704
def set_isolation_level(self, dbapi_connection: dbapi.Connection, level: str) -> None:
691705
dbapi_connection.set_isolation_level(level)

ydb_sqlalchemy/sqlalchemy/requirements.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ def temporary_views(self):
4646

4747
@property
4848
def index_reflection(self):
49+
# Reflection supported with limits
4950
return exclusions.closed()
5051

5152
@property

0 commit comments

Comments
 (0)