implement get_catalog_names

Signed-off-by: Đặng Minh Dũng <[email protected]>

Author: dungdm93

tests/integration/test_sqlalchemy_integration.py

@@ -410,6 +410,15 @@ def test_json_column(trino_connection, json_object):
         metadata.drop_all(engine)
 
 
+@pytest.mark.parametrize('trino_connection', ['system'], indirect=True)
+def test_get_catalog_names(trino_connection):
+    engine, conn = trino_connection
+
+    schemas = engine.dialect.get_catalog_names(conn)
+    assert len(schemas) == 5
+    assert set(schemas) == {"jmx", "memory", "system", "tpcds", "tpch"}
+
+
 @pytest.mark.parametrize('trino_connection', ['memory'], indirect=True)
 def test_get_table_comment(trino_connection):
     engine, conn = trino_connection

trino/sqlalchemy/dialect.py

@@ -192,6 +192,16 @@ def get_foreign_keys(
         """Trino has no support for foreign keys. Returns an empty list."""
         return []
 
+    def get_catalog_names(self, connection: Connection, **kw) -> List[str]:
+        query = dedent(
+            """
+            SELECT "table_cat"
+            FROM "system"."jdbc"."catalogs"
+        """
+        ).strip()
+        res = connection.execute(sql.text(query))
+        return [row.table_cat for row in res]
+
     def get_schema_names(self, connection: Connection, **kw) -> List[str]:
         query = dedent(
             """