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

Author: yalexbadan

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

@@ -76,7 +76,7 @@ class Table(BaseModel):
     columns: list[Column]
     description: str | None = None
     properties: TableProperties = None
-    primaryKey: str | None = None
+    primaryKey: list[str] | None = None
 
 
 class ConstraintType(Enum):

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

@@ -81,7 +81,7 @@ def get_table_list(self) -> list[Table]:
                         catalog=row["catalog"],
                         table=row["table_name"],
                     ),
-                    primaryKey="",
+                    primaryKey=[],
                 )
 
             # table exists, and add column to the table
@@ -96,7 +96,8 @@ def get_table_list(self) -> list[Table]:
             )
             # if column is primary key
             if row["is_pk"] == "YES":
-                unique_tables[schema_table].primaryKey = row["column_name"]
+                unique_tables[schema_table].primaryKey.append(row["column_name"])
+                
         return list(unique_tables.values())
 
     def get_constraints(self) -> list[Constraint]:

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

@@ -56,7 +56,7 @@ def get_table_list(self) -> list[Table]:
                         catalog="",
                         table=row["table_name"],
                     ),
-                    primaryKey="",
+                    primaryKey=[],
                 )
 
             # table exists, and add column to the table
@@ -71,7 +71,7 @@ def get_table_list(self) -> list[Table]:
             )
             # if column is primary key
             if row["column_key"] == "PRI":
-                unique_tables[schema_table].primaryKey = row["column_name"]
+                unique_tables[schema_table].primaryKey.append(row["column_name"])
         return list(unique_tables.values())
 
     def get_constraints(self) -> list[Constraint]:

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

@@ -88,7 +88,7 @@ def get_table_list(self) -> list[Table]:
                         catalog="",  # Oracle doesn't use catalogs.
                         table=row["TABLE_NAME"],
                     ),
-                    primaryKey="",
+                    primaryKey=[],
                 )
 
             unique_tables[schema_table].columns.append(
@@ -100,6 +100,8 @@ def get_table_list(self) -> list[Table]:
                     properties=None,
                 )
             )
+            
+            # TODO: manage primary key
 
         return list(unique_tables.values())
 

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,6 +45,12 @@ 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
             WHERE
                 t.table_type IN ('BASE TABLE', 'VIEW')
                 AND t.table_schema NOT IN ('information_schema', 'pg_catalog');
@@ -67,7 +74,7 @@ def get_table_list(self) -> list[Table]:
                         catalog=row["table_catalog"],
                         table=row["table_name"],
                     ),
-                    primaryKey="",
+                    primaryKey=[],
                 )
 
             # table exists, and add column to the table
@@ -80,6 +87,10 @@ def get_table_list(self) -> list[Table]:
                     properties=None,
                 )
             )
+            
+            if row["constraint_type"] == 'PRIMARY KEY':
+                unique_tables[schema_table].primaryKey.append(row["column_name"])
+
         return list(unique_tables.values())
 
     def get_constraints(self) -> list[Constraint]: