feat: Add relationships to ER

backend/app/api/semantic.py

@@ -41,7 +41,14 @@ def get_semantic_schema(datasource_id: int, db: Session = Depends(get_db)):
         raise HTTPException(status_code=404, detail="DataSource not found")
     
     try:
-        return MDLService.get_raw_schema(ds)
+        raw_schema = MDLService.get_raw_schema(ds)
+        result = {}
+        for table, data in raw_schema.items():
+            if isinstance(data, dict) and "columns" in data:
+                result[table] = data["columns"]
+            elif isinstance(data, list):
+                result[table] = data
+        return result
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 

backend/app/connectors/postgres.py

@@ -22,41 +22,44 @@ class PostgresConnector:
             return [dict(row._mapping) for row in result]
 
     def get_schema(self):
-        if self.engine.dialect.name == "sqlite":
-            return self._get_sqlite_schema()
-
-        query = """
-        SELECT table_name, column_name, data_type
-        FROM information_schema.columns
-        WHERE table_schema = 'public'
-        ORDER BY table_name, ordinal_position;
-        """
-        try:
-            results = self.execute_query(query)
-            schema = {}
-            for row in results:
-                table = row['table_name']
-                if table not in schema:
-                    schema[table] = []
-                schema[table].append({"name": row['column_name'], "type": row['data_type']})
-            return schema
-        except Exception as e:
-            print(f"Error getting schema: {e}")
-            return {}
-
-    def _get_sqlite_schema(self):
         try:
             from sqlalchemy import inspect
             inspector = inspect(self.engine)
             schema = {}
-            for table_name in inspector.get_table_names():
+            # Default schema for postgres is 'public', sqlite is None
+            schema_name = 'public' if self.engine.dialect.name == 'postgresql' else None
+            
+            for table_name in inspector.get_table_names(schema=schema_name):
                 columns = []
-                for col in inspector.get_columns(table_name):
+                # get columns
-                    columns.append({"name": col['name'], "type": str(col['type'])})
+                for col in inspector.get_columns(table_name, schema=schema_name):
-                schema[table_name] = columns
+                    columns.append({
+                        "name": col['name'], 
+                        "type": str(col['type'])
+                    })
+                
+                # get primary key
+                pk_constraint = inspector.get_pk_constraint(table_name, schema=schema_name)
+                pks = pk_constraint.get('constrained_columns', []) if pk_constraint else []
+                
+                # get foreign keys
+                fks = inspector.get_foreign_keys(table_name, schema=schema_name)
+                foreign_keys = []
+                for fk in fks:
+                    foreign_keys.append({
+                        "constrained_columns": fk['constrained_columns'],
+                        "referred_table": fk['referred_table'],
+                        "referred_columns": fk['referred_columns']
+                    })
+                
+                schema[table_name] = {
+                    "columns": columns,
+                    "primary_keys": pks,
+                    "foreign_keys": foreign_keys
+                }
             return schema
         except Exception as e:
-            print(f"Error getting SQLite schema: {e}")
+            print(f"Error getting schema: {e}")
             return {}
 
     def test_connection(self) -> bool:

backend/app/schemas/mdl.py

@@ -34,7 +34,7 @@ class Column(BaseModel):
     expression: Optional[str] = None
     isHidden: bool = False
     columnLevelAccessControl: Optional[ColumnAccessControl] = None
-    properties: Dict[str, str] = Field(default_factory=dict)
+    properties: Dict[str, Any] = Field(default_factory=dict)
 
 # Model Definitions
 class TableReference(BaseModel):

backend/app/services/mdl.py

@@ -34,10 +34,24 @@ class MDLService:
         raw_schema = MDLService.get_raw_schema(datasource)
         
         models = []
-        for table_name, columns in raw_schema.items():
+        relationships = []
+        from app.schemas.mdl import Relationship
+        
+        # Helper to get columns for a table from the raw schema (which could be a list or a dict)
+        def get_table_info(t_name):
+            data = raw_schema.get(t_name, [])
+            if isinstance(data, dict) and "columns" in data:
+                return data
+            return {"columns": data, "primary_keys": [], "foreign_keys": []}
+
+        for table_name in raw_schema.keys():
             if selected_tables is not None and table_name not in selected_tables:
                 continue
 
+            table_info = get_table_info(table_name)
+            columns = table_info["columns"]
+            pks = table_info.get("primary_keys", [])
+
             model_cols = []
             for col_info in columns:
                 if isinstance(col_info, dict):
@@ -65,7 +79,8 @@ class MDLService:
                     if allowed and name not in allowed:
                         continue
                 
-                model_cols.append(Column(name=name, type=type_))
+                is_pk = name in pks
+                model_cols.append(Column(name=name, type=type_, properties={"is_primary_key": is_pk}))
 
             if not model_cols:
                 continue
@@ -73,14 +88,46 @@ class MDLService:
             models.append(Model(
                 name=table_name,
                 tableReference=TableReference(table=table_name),
-                columns=model_cols
+                columns=model_cols,
+                primaryKey=pks[0] if pks else None
             ))
             
+            # Extract relationships from foreign keys
+            fks = table_info.get("foreign_keys", [])
+            for fk in fks:
+                referred_table = fk.get("referred_table")
+                if not referred_table:
+                    continue
+                # Skip if the referred table is not selected
+                if selected_tables is not None and referred_table not in selected_tables:
+                    continue
+                
+                constrained_cols = fk.get("constrained_columns", [])
+                referred_cols = fk.get("referred_columns", [])
+                
+                if len(constrained_cols) == 1 and len(referred_cols) == 1:
+                    # Update column properties for FK
+                    fk_col_name = constrained_cols[0]
+                    for col in model_cols:
+                        if col.name == fk_col_name:
+                            col.properties["is_foreign_key"] = True
+                            
+                    # Simple single-column foreign key
+                    condition = f"{table_name}.{constrained_cols[0]} = {referred_table}.{referred_cols[0]}"
+                    rel_name = f"{table_name}_{constrained_cols[0]}_to_{referred_table}"
+                    relationships.append(Relationship(
+                        name=rel_name,
+                        models=[table_name, referred_table],
+                        joinType="MANY_TO_ONE", # typically a foreign key represents many-to-one
+                        condition=condition
+                    ))
+
         return MDLManifest(
             catalog="default",
             schema="public", # Default schema, might need adjustment based on datasource config
             dataSource=datasource.type.upper(),
-            models=models
+            models=models,
+            relationships=relationships
         )
 
     @staticmethod

frontend/src/components/modeling/TableNode.tsx

@@ -6,6 +6,10 @@ import { Table as TableIcon } from "lucide-react";
 interface Column {
   name: string;
   type: string;
+  properties?: {
+    is_primary_key?: boolean;
+    is_foreign_key?: boolean;
+  };
 }
 
 interface TableNodeData {
@@ -16,7 +20,7 @@ interface TableNodeData {
 
 export const TableNode = memo(({ data }: { data: TableNodeData }) => {
   return (
-    <Card className="min-w-[180px] max-w-[240px] shadow-md border-t-4 border-t-blue-500 text-xs">
+    <Card className="min-w-[220px] max-w-[280px] shadow-md border-t-4 border-t-blue-500 text-xs bg-white">
       <Handle type="target" position={Position.Top} className="!bg-blue-500" />
       
       <CardHeader 
@@ -30,17 +34,41 @@ export const TableNode = memo(({ data }: { data: TableNodeData }) => {
       </CardHeader>
       
       <CardContent className="p-0">
-        <div className="max-h-[200px] overflow-y-auto">
+        <div className="max-h-[250px] overflow-y-auto">
-          {data.columns.map((col) => (
+          <table className="w-full text-left border-collapse">
-            <div 
+            <tbody>
-              key={col.name} 
+              {data.columns.map((col) => {
-              className="py-1.5 px-3 border-b last:border-0 hover:bg-gray-50 flex items-center"
+                const isPk = col.properties?.is_primary_key;
-              title={`${col.name} (${col.type})`}
+                const isFk = col.properties?.is_foreign_key;
-            >
+                let keyText = "";
-              <span className="font-medium truncate">{col.name}</span>
+                if (isPk && isFk) keyText = "PK, FK";
-              {/* 类型列已被隐藏 */}
+                else if (isPk) keyText = "PK";
-            </div>
+                else if (isFk) keyText = "FK";
-          ))}
+
+                // Simplify type display, e.g., INTEGER -> int, CHARACTER VARYING -> string
+                let displayType = (col.type || "string").toLowerCase();
+                if (displayType.includes("int")) displayType = "int";
+                else if (displayType.includes("char") || displayType.includes("text")) displayType = "string";
+                else if (displayType.includes("time") || displayType.includes("date")) displayType = "date";
+                else if (displayType.includes("bool")) displayType = "boolean";
+                else if (displayType.includes("float") || displayType.includes("double") || displayType.includes("numeric") || displayType.includes("decimal")) displayType = "float";
+
+                return (
+                  <tr 
+                    key={col.name} 
+                    className="border-b last:border-0 hover:bg-gray-50"
+                    title={`${col.name} (${col.type})`}
+                  >
+                    <td className="py-1.5 px-3 w-16 text-gray-500 font-mono truncate border-r border-gray-100">{displayType}</td>
+                    <td className="py-1.5 px-3 font-medium truncate text-gray-800">{col.name}</td>
+                    <td className="py-1.5 px-3 w-10 text-center text-gray-500 font-semibold text-[10px] border-l border-gray-100">
+                      {keyText}
+                    </td>
+                  </tr>
+                );
+              })}
+            </tbody>
+          </table>
         </div>
       </CardContent>