Add index creation

Author: LuckySting

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"}

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