feat(api): append cooments of column and table for metadata api (#812)

* feat(clickhouse): append comments of column and table

* feat(clickhouse): append comments of column and table

* feat(mssql): append comments of column and table

* feat(mysql): append comments of column and table

* feat(postgres): append comments of column and table

* fix(trino): wait trino container ready to avoid 'nodes is empty' from Trino

* feat(trino): append comments of column and table

* chore(test): sort func and rename func

* test(snowflake): mark not implemented feature

* refactor(metadata): tidy up

* chore(metadata): remove unimplemented data source

* fix(trino): fix connection

* chore(trino): merge two statement

* test(trino): update excepted

we changed expression but did not update excepted on #782

Author: grieve54706

ibis-server/app/model/metadata/bigquery.py

@@ -1,5 +1,3 @@
-from json import loads
-
 from app.model import BigQueryConnectionInfo
 from app.model.data_source import DataSource
 from app.model.metadata.dto import (
@@ -49,7 +47,7 @@ def get_table_list(self) -> list[Table]:
             LEFT JOIN {dataset_id}.INFORMATION_SCHEMA.TABLE_OPTIONS table_options
                 ON c.table_name = table_options.table_name
             """
-        response = loads(self.connection.sql(sql).to_pandas().to_json(orient="records"))
+        response = self.connection.sql(sql).to_pandas().to_dict(orient="records")
 
         def get_data_type(data_type) -> str:
             if "STRUCT" in data_type:
@@ -123,7 +121,7 @@ def get_constraints(self) -> list[Constraint]:
                 ON ccu.constraint_name = tc.constraint_name
             WHERE tc.constraint_type = 'FOREIGN KEY'
             """
-        response = loads(self.connection.sql(sql).to_pandas().to_json(orient="records"))
+        response = self.connection.sql(sql).to_pandas().to_dict(orient="records")
 
         constraints = []
         for row in response:

ibis-server/app/model/metadata/clickhouse.py

@@ -1,5 +1,3 @@
-from json import loads
-
 from app.model import ClickHouseConnectionInfo
 from app.model.data_source import DataSource
 from app.model.metadata.dto import (
@@ -15,25 +13,27 @@
 class ClickHouseMetadata(Metadata):
     def __init__(self, connection_info: ClickHouseConnectionInfo):
         super().__init__(connection_info)
+        self.connection = DataSource.clickhouse.get_connection(connection_info)
 
     def get_table_list(self) -> list[Table]:
         sql = """
             SELECT
-                database AS table_schema,
-                table AS table_name,
-                name AS column_name,
-                type AS data_type
+                c.database AS table_schema,
+                c.table AS table_name,
+                t.comment AS table_comment,
+                c.name AS column_name,
+                c.type AS data_type,
+                c.comment AS column_comment
             FROM
-                system.columns
+                system.columns AS c
+            JOIN
+                system.tables AS t
+                ON c.database = t.database
+                AND c.table = t.name
             WHERE
-                database NOT IN ('system', 'INFORMATION_SCHEMA', 'information_schema', 'pg_catalog');
+                c.database NOT IN ('system', 'INFORMATION_SCHEMA', 'information_schema', 'pg_catalog');
             """
-        response = loads(
-            DataSource.clickhouse.get_connection(self.connection_info)
-            .sql(sql)
-            .to_pandas()
-            .to_json(orient="records")
-        )
+        response = self.connection.sql(sql).to_pandas().to_dict(orient="records")
 
         unique_tables = {}
         for row in response:
@@ -45,7 +45,7 @@ def get_table_list(self) -> list[Table]:
             if schema_table not in unique_tables:
                 unique_tables[schema_table] = Table(
                     name=schema_table,
-                    description="",
+                    description=row["table_comment"],
                     columns=[],
                     properties=TableProperties(
                         catalog=None,
@@ -61,7 +61,7 @@ def get_table_list(self) -> list[Table]:
                     name=row["column_name"],
                     type=self._transform_column_type(row["data_type"]),
                     notNull=False,
-                    description="",
+                    description=row["column_comment"],
                     properties=None,
                 )
             )
@@ -103,9 +103,3 @@ def _transform_column_type(self, data_type):
         }
 
         return switcher.get(data_type, WrenEngineColumnType.UNKNOWN)
-
-
-def to_json(df):
-    json_obj = loads(df.to_json(orient="split"))
-    del json_obj["index"]
-    return json_obj

ibis-server/app/model/metadata/factory.py

@@ -1,49 +1,26 @@
-from json import loads
-
-from app.model import ConnectionInfo
 from app.model.data_source import DataSource
 from app.model.metadata.bigquery import BigQueryMetadata
 from app.model.metadata.clickhouse import ClickHouseMetadata
-from app.model.metadata.dto import (
-    Constraint,
-    Table,
-)
 from app.model.metadata.metadata import Metadata
 from app.model.metadata.mssql import MSSQLMetadata
 from app.model.metadata.mysql import MySQLMetadata
 from app.model.metadata.postgres import PostgresMetadata
 from app.model.metadata.trino import TrinoMetadata
 
+mapping = {
+    DataSource.bigquery: BigQueryMetadata,
+    DataSource.clickhouse: ClickHouseMetadata,
+    DataSource.mssql: MSSQLMetadata,
+    DataSource.mysql: MySQLMetadata,
+    DataSource.postgres: PostgresMetadata,
+    DataSource.trino: TrinoMetadata,
+}
 
-class MetadataFactory:
-    def __init__(self, data_source: DataSource, connection_info: ConnectionInfo):
-        self.metadata = self.get_metadata(data_source, connection_info)
-
-    def get_metadata(self, data_source: DataSource, connection_info) -> Metadata:
-        if data_source == DataSource.postgres:
-            return PostgresMetadata(connection_info)
-        if data_source == DataSource.bigquery:
-            return BigQueryMetadata(connection_info)
-        if data_source == DataSource.mysql:
-            return MySQLMetadata(connection_info)
-        if data_source == DataSource.mssql:
-            return MSSQLMetadata(connection_info)
-        if data_source == DataSource.clickhouse:
-            return ClickHouseMetadata(connection_info)
-        if data_source == DataSource.trino:
-            return TrinoMetadata(connection_info)
-
-        raise NotImplementedError(f"Unsupported data source: {self}")
 
-    def get_table_list(self) -> list[Table]:
-        return self.metadata.get_table_list()
-
-    def get_constraints(self) -> list[Constraint]:
-        return self.metadata.get_constraints()
-
-
-def to_json(df):
-    json_obj = loads(df.to_json(orient="split"))
-    del json_obj["index"]
-    json_obj["dtypes"] = df.dtypes.apply(lambda x: x.name).to_dict()
-    return json_obj
+class MetadataFactory:
+    @staticmethod
+    def get_metadata(data_source: DataSource, connection_info) -> Metadata:
+        try:
+            return mapping[data_source](connection_info)
+        except KeyError:
+            raise NotImplementedError(f"Unsupported data source: {data_source}")

ibis-server/app/model/metadata/mssql.py

@@ -1,5 +1,3 @@
-from json import loads
-
 from app.model import MSSqlConnectionInfo
 from app.model.data_source import DataSource
 from app.model.metadata.dto import (
@@ -16,6 +14,7 @@
 class MSSQLMetadata(Metadata):
     def __init__(self, connection_info: MSSqlConnectionInfo):
         super().__init__(connection_info)
+        self.connection = DataSource.mssql.get_connection(connection_info)
 
     def get_table_list(self) -> list[Table]:
         sql = """
@@ -26,7 +25,9 @@ def get_table_list(self) -> list[Table]:
                 col.COLUMN_NAME AS column_name,
                 col.DATA_TYPE AS data_type,
                 CASE WHEN pk.COLUMN_NAME IS NOT NULL THEN 'YES' ELSE 'NO' END AS is_pk,
-                col.IS_NULLABLE AS is_nullable
+                col.IS_NULLABLE AS is_nullable,
+                CAST(tprop.value AS NVARCHAR(MAX)) AS table_comment,
+                CAST(cprop.value AS NVARCHAR(MAX)) AS column_comment
             FROM 
                 INFORMATION_SCHEMA.COLUMNS col
             LEFT JOIN 
@@ -40,13 +41,28 @@ def get_table_list(self) -> list[Table]:
                 AND col.TABLE_NAME = pk.TABLE_NAME
                 AND col.COLUMN_NAME = pk.COLUMN_NAME
                 AND pk.CONSTRAINT_NAME = tab.CONSTRAINT_NAME
+            LEFT JOIN 
+                sys.tables st
+                ON st.name = col.TABLE_NAME 
+                AND SCHEMA_NAME(st.schema_id) = col.TABLE_SCHEMA
+            LEFT JOIN 
+                sys.extended_properties tprop
+                ON tprop.major_id = st.object_id 
+                AND tprop.minor_id = 0 
+                AND tprop.name = 'MS_Description'
+            LEFT JOIN 
+                sys.columns sc
+                ON sc.object_id = st.object_id 
+                AND sc.name = col.COLUMN_NAME
+            LEFT JOIN 
+                sys.extended_properties cprop
+                ON cprop.major_id = sc.object_id 
+                AND cprop.minor_id = sc.column_id 
+                AND cprop.name = 'MS_Description'
+            WHERE
+                col.TABLE_SCHEMA NOT IN ('sys', 'INFORMATION_SCHEMA');
             """
-        response = loads(
-            DataSource.mssql.get_connection(self.connection_info)
-            .sql(sql)
-            .to_pandas()
-            .to_json(orient="records")
-        )
+        response = self.connection.sql(sql).to_pandas().to_dict(orient="records")
 
         unique_tables = {}
         for row in response:
@@ -58,7 +74,7 @@ def get_table_list(self) -> list[Table]:
             if schema_table not in unique_tables:
                 unique_tables[schema_table] = Table(
                     name=schema_table,
-                    description="",
+                    description=row["table_comment"],
                     columns=[],
                     properties=TableProperties(
                         schema=row["table_schema"],
@@ -74,7 +90,7 @@ def get_table_list(self) -> list[Table]:
                     name=row["column_name"],
                     type=self._transform_column_type(row["data_type"]),
                     notNull=row["is_nullable"].lower() == "no",
-                    description="",
+                    description=row["column_comment"],
                     properties=None,
                 )
             )
@@ -121,12 +137,7 @@ def get_constraints(self) -> list[Constraint]:
                 ON fkc.referenced_column_id = ref_c.column_id 
                 AND ref_c.object_id = ref_t.object_id
             """
-        res = loads(
-            DataSource.mssql.get_connection(self.connection_info)
-            .sql(sql)
-            .to_pandas()
-            .to_json(orient="records")
-        )
+        res = self.connection.sql(sql).to_pandas().to_dict(orient="records")
         constraints = []
         for row in res:
             constraints.append(
@@ -196,9 +207,3 @@ def _transform_column_type(self, data_type):
         }
 
         return switcher.get(data_type.lower(), WrenEngineColumnType.UNKNOWN)
-
-
-def to_json(df):
-    json_obj = loads(df.to_json(orient="split"))
-    del json_obj["index"]
-    return json_obj

ibis-server/app/model/metadata/mysql.py

@@ -1,5 +1,3 @@
-from json import loads
-
 from app.model import MySqlConnectionInfo
 from app.model.data_source import DataSource
 from app.model.metadata.dto import (
@@ -16,27 +14,29 @@
 class MySQLMetadata(Metadata):
     def __init__(self, connection_info: MySqlConnectionInfo):
         super().__init__(connection_info)
+        self.connection = DataSource.mysql.get_connection(connection_info)
 
     def get_table_list(self) -> list[Table]:
         sql = """
             SELECT
-                TABLE_SCHEMA as table_schema,
-                TABLE_NAME as table_name,
-                COLUMN_NAME as column_name,
-                DATA_TYPE as data_type,
-                IS_NULLABLE as is_nullable,
-                COLUMN_KEY as column_key
+                c.TABLE_SCHEMA AS table_schema,
+                c.TABLE_NAME AS table_name,
+                c.COLUMN_NAME AS column_name,
+                c.DATA_TYPE AS data_type,
+                c.IS_NULLABLE AS is_nullable,
+                c.COLUMN_KEY AS column_key,
+                c.COLUMN_COMMENT AS column_comment,
+                t.TABLE_COMMENT AS table_comment
             FROM
-                information_schema.COLUMNS
+                information_schema.COLUMNS c
+            JOIN
+                information_schema.TABLES t
+                ON c.TABLE_SCHEMA = t.TABLE_SCHEMA
+                AND c.TABLE_NAME = t.TABLE_NAME
             WHERE
-                TABLE_SCHEMA not IN ("mysql", "information_schema", "performance_schema", "sys")
+                c.TABLE_SCHEMA NOT IN ('mysql', 'information_schema', 'performance_schema', 'sys');
             """
-        response = loads(
-            DataSource.mysql.get_connection(self.connection_info)
-            .sql(sql)
-            .to_pandas()
-            .to_json(orient="records")
-        )
+        response = self.connection.sql(sql).to_pandas().to_dict(orient="records")
 
         unique_tables = {}
         for row in response:
@@ -48,7 +48,7 @@ def get_table_list(self) -> list[Table]:
             if schema_table not in unique_tables:
                 unique_tables[schema_table] = Table(
                     name=schema_table,
-                    description="",
+                    description=row["table_comment"],
                     columns=[],
                     properties=TableProperties(
                         schema=row["table_schema"],
@@ -65,7 +65,7 @@ def get_table_list(self) -> list[Table]:
                     name=row["column_name"],
                     type=self._transform_column_type(row["data_type"]),
                     notNull=row["is_nullable"].lower() == "no",
-                    description="",
+                    description=row["column_comment"],
                     properties=None,
                 )
             )
@@ -93,12 +93,7 @@ def get_constraints(self) -> list[Constraint]:
                 ON rc.CONSTRAINT_NAME = kcu.CONSTRAINT_NAME 
                 AND rc.CONSTRAINT_SCHEMA = kcu.CONSTRAINT_SCHEMA
             """
-        res = loads(
-            DataSource.mysql.get_connection(self.connection_info)
-            .sql(sql)
-            .to_pandas()
-            .to_json(orient="records")
-        )
+        res = self.connection.sql(sql).to_pandas().to_dict(orient="records")
         constraints = []
         for row in res:
             constraints.append(
@@ -167,9 +162,3 @@ def _transform_column_type(self, data_type):
         }
 
         return switcher.get(data_type.lower(), WrenEngineColumnType.UNKNOWN)
-
-
-def to_json(df):
-    json_obj = loads(df.to_json(orient="split"))
-    del json_obj["index"]
-    return json_obj