feat: add get_indexes

sqlalchemy_risingwave/base.py

@@ -113,7 +113,38 @@ def get_columns(self, conn, table_name, schema=None, **kw):
         return res
 
     def get_indexes(self, conn, table_name, schema=None, **kw):
-        return []
+        table_oid = self.get_table_oid(
+            conn, table_name, schema, info_cache=kw.get("info_cache")
+        )
+
+        sql = (
+            "select i.relname, a.attname from pg_catalog.pg_class t "
+            "join pg_catalog.pg_index ix on t.oid = ix.indrelid "
+            "join pg_catalog.pg_class i on i.oid = ix.indexrelid "
+            "join pg_catalog.pg_attribute a on t.oid = a.attrelid and a.attnum = ANY(ix.indkey)"
+            "where t.oid = :table_oid"
+        )
+        rows = conn.execute(
+            text(sql),
+            {"table_oid": table_oid},
+        )
+
+        indexes = {}
+        for row in rows:
+            if not row.relname in indexes:
+                indexes[row.relname] = []
+            indexes[row.relname].append(row.attname)
+
+        res = []
+        for index in indexes:
+            res.append(
+                {
+                    "name": index,
+                    "column_names": indexes[index],
+                    "unique": False,
+                }
+            )
+        return res
 
     def get_foreign_keys_v1(self, conn, table_name, schema=None, **kw):
         raise []

test/test_schema.py

@@ -7,19 +7,11 @@ class SchemaTest(fixtures.TestBase):
 
     def teardown_method(self, method):
         with testing.db.begin() as conn:
-            conn.execute(text("DROP TABLE IF EXISTS users"))
+            conn.execute("DROP TABLE IF EXISTS users")
 
     def setup_method(self):
         with testing.db.begin() as conn:
-            conn.execute(
-                text(
-                    """
-                    CREATE TABLE users (
-                        name STRING PRIMARY KEY
-                    )
-                    """
-                )
-            )
+            conn.execute("CREATE TABLE users (name STRING PRIMARY KEY)")
         self.meta = MetaData(schema="public")
 
     def test_get_columns_indexes_across_schema(self):
@@ -34,3 +26,17 @@ def test_returning_clause(self):
 
             for t in table_names:
                 assert t == str("users")
+
+    def test_get_indexes(self):
+        with testing.db.begin() as conn:
+            conn.execute("CREATE TABLE three_columns (id1 INT, id2 INT, id3 INT)")
+            conn.execute("CREATE INDEX three_columns_idx ON three_columns(id2) INCLUDE(id1)")
+
+            insp = inspect(testing.db)
+            indexes = insp.get_indexes("three_columns")
+
+            assert len(indexes) == 1
+            assert indexes[0]["name"] == "three_columns_idx"
+            assert indexes[0]["column_names"] == ["id2", "id1"]
+
+            conn.execute("DROP TABLE three_columns")