feat: Add relationships to ER

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