fix(ibis): added primary key management on postgres and updated mssql and mysql to manage composite keys

.gitignore

@@ -20,3 +20,5 @@ __pycache__/
 venv/
 **/.env*
 **/*.so
+.devcontainer
+.vscode

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

@@ -101,6 +101,8 @@ def get_table_list(self) -> list[Table]:
                 )
             )
 
+            # TODO: manage primary key
+
         return list(unique_tables.values())
 
     def get_constraints(self) -> list[Constraint]:

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

@@ -26,6 +26,7 @@ def get_table_list(self) -> list[Table]:
                 c.data_type,
                 c.is_nullable,
                 c.ordinal_position,
+                tc.constraint_type,
                 obj_description(cls.oid) AS table_comment,
                 col_description(cls.oid, a.attnum) AS column_comment
             FROM
@@ -44,13 +45,22 @@ def get_table_list(self) -> list[Table]:
                 pg_attribute a
                 ON a.attrelid = cls.oid
                 AND a.attname = c.column_name
+            LEFT JOIN information_schema.key_column_usage kcu
+                ON c.table_name = kcu.table_name
+                AND c.column_name = kcu.column_name
+                AND c.table_schema = kcu.table_schema
+            LEFT JOIN information_schema.table_constraints tc
+                ON kcu.constraint_name = tc.constraint_name
+                AND kcu.table_schema = tc.table_schema
             WHERE
                 t.table_type IN ('BASE TABLE', 'VIEW')
-                AND t.table_schema NOT IN ('information_schema', 'pg_catalog');
+                AND t.table_schema NOT IN ('information_schema', 'pg_catalog')
+                AND (tc.constraint_type = 'PRIMARY KEY' OR tc.constraint_type IS NULL);
             """
         response = self.connection.sql(sql).to_pandas().to_dict(orient="records")
 
         unique_tables = {}
+        unique_tables_primary_key = {}
         for row in response:
             # generate unique table name
             schema_table = self._format_postgres_compact_table_name(
@@ -71,15 +81,38 @@ def get_table_list(self) -> list[Table]:
                 )
 
             # table exists, and add column to the table
-            unique_tables[schema_table].columns.append(
-                Column(
-                    name=row["column_name"],
-                    type=self._transform_postgres_column_type(row["data_type"]),
-                    notNull=row["is_nullable"].lower() == "no",
-                    description=row["column_comment"],
+            column = Column(
+                name=row["column_name"],
+                type=self._transform_postgres_column_type(row["data_type"]),
+                notNull=row["is_nullable"].lower() == "no",
+                description=row["column_comment"],
+                properties=None,
+            )
+            unique_tables[schema_table].columns.append(column)
+
+            if row["constraint_type"] == "PRIMARY KEY":
+                if schema_table in unique_tables_primary_key:
+                    unique_tables_primary_key[schema_table].append(column)
+                else:
+                    unique_tables_primary_key[schema_table] = [column]
+
+        for schema_table, columns in unique_tables_primary_key.items():
+            if len(columns) == 0:
+                continue
+            elif len(columns) == 1:
+                unique_tables[schema_table].primaryKey = columns[0].name
+            else:
+                composite_primary_key_column = Column(
+                    name="composed_primary_key",
+                    type="json",
+                    notNull=True,
+                    description=f"Composed primary key based on fields: {[col.name for col in columns]}",
                     properties=None,
+                    nestedColumns=columns,
                 )
-            )
+                unique_tables[schema_table].columns.append(composite_primary_key_column)
+                unique_tables[schema_table].primaryKey = "composed_primary_key"
+
         return list(unique_tables.values())
 
     def get_constraints(self) -> list[Constraint]:
@@ -97,6 +130,7 @@ def get_constraints(self) -> list[Constraint]:
                 AND tc.table_schema = kcu.table_schema
             JOIN information_schema.constraint_column_usage AS ccu
                 ON ccu.constraint_name = tc.constraint_name
+                AND ccu.table_schema = tc.table_schema
             WHERE tc.constraint_type = 'FOREIGN KEY'
             """
         res = self.connection.sql(sql).to_pandas().to_dict(orient="records")

wren-core-legacy/wren-base/src/test/java/io/wren/base/sqlrewrite/TestCumulativeMetric.java

@@ -115,7 +115,7 @@ public void testModelOnCumulativeMetric()
                         List.of(
                                 Column.column("totalprice", WrenTypes.INTEGER, null, false),
                                 Column.column("orderdate", "DATE", null, false)),
-                        "orderdate"))
+                                "orderdate"))
                 .build();
         WrenMDL mdl = WrenMDL.fromManifest(
                 copyOf(manifest)

wren-core-legacy/wren-base/src/test/java/io/wren/base/sqlrewrite/TestMetric.java

@@ -279,7 +279,7 @@ public void testModelOnMetric()
                         List.of(
                                 Column.column("name", WrenTypes.VARCHAR, null, true),
                                 Column.column("revenue", WrenTypes.INTEGER, null, true, "totalprice")),
-                        "name"))
+                                "name"))
                 .build();
         WrenMDL mdl = WrenMDL.fromManifest(
                 copyOf(manifest)