Initial open-source release

python/klo-sl/scripts/gen_b2b_saas_model.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+"""Generate semantic layer YAML sources from demo DB metadata.
+
+Usage:
+    kubectl port-forward -n klo-demo deployment/klo-demo-db 5433:5432 &
+    KLO_DEMO_DB_PASSWORD=local-demo-password python scripts/gen_b2b_saas_model.py
+"""
+
+import os
+import psycopg2
+import yaml
+
+CONNECTION_ID = "256bc76b-cc47-4d5d-a9fc-5bcfb0364d44"
+OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "..", "sources", "b2b_saas")
+
+DB_PARAMS = {
+    "host": os.environ.get("KLO_DEMO_DB_HOST", "127.0.0.1"),
+    "port": int(os.environ.get("KLO_DEMO_DB_PORT", "5433")),
+    "user": os.environ.get("KLO_DEMO_DB_USER", "klo-demo-user"),
+    "password": os.environ.get("KLO_DEMO_DB_PASSWORD", ""),
+    "dbname": os.environ.get("KLO_DEMO_DB_NAME", "klo-demo-db"),
+}
+
+# Map DB types to semantic layer types
+TYPE_MAP = {
+    "INTEGER": "number",
+    "FLOAT": "number",
+    "NUMERIC": "number",
+    "DECIMAL": "number",
+    "BIGINT": "number",
+    "SMALLINT": "number",
+    "DOUBLE": "number",
+    "REAL": "number",
+    "VARCHAR": "string",
+    "TEXT": "string",
+    "CHAR": "string",
+    "DATE": "time",
+    "TIMESTAMP": "time",
+    "TIMESTAMPTZ": "time",
+    "DATETIME": "time",
+    "TIME": "time",
+    "BOOLEAN": "boolean",
+    "BOOL": "boolean",
+}
+
+# Columns whose names suggest a time role
+TIME_PATTERNS = {"_at", "_date", "date", "timestamp", "created", "updated"}
+
+
+def is_time_column(name: str, db_type: str) -> bool:
+    sl_type = TYPE_MAP.get(db_type.upper(), "string")
+    if sl_type == "time":
+        return True
+    # VARCHAR columns with date-like names (e.g. created_at stored as VARCHAR)
+    lower = name.lower()
+    return any(p in lower for p in TIME_PATTERNS) and sl_type == "string"
+
+
+def map_type(db_type: str, col_name: str) -> str:
+    upper = db_type.upper()
+    if upper in TYPE_MAP:
+        base = TYPE_MAP[upper]
+        # Override string→time for date-like column names
+        if base == "string" and is_time_column(col_name, db_type):
+            return "time"
+        return base
+    return "string"
+
+
+def main():
+    conn = psycopg2.connect(**DB_PARAMS)
+    cur = conn.cursor()
+
+    # 1. Fetch tables
+    cur.execute(
+        "SELECT id, name FROM source_tables WHERE connection_id = %s ORDER BY name",
+        (CONNECTION_ID,),
+    )
+    tables = {row[0]: row[1] for row in cur.fetchall()}
+    table_ids = tuple(tables.keys())
+
+    # 2. Fetch columns
+    cur.execute(
+        """
+        SELECT id, name, type, nullable, primary_key, table_id
+        FROM source_columns
+        WHERE table_id = ANY(%s::uuid[])
+        ORDER BY table_id, primary_key DESC, name
+        """,
+        (list(table_ids),),
+    )
+    columns_by_table: dict[str, list] = {}
+    col_id_to_info: dict[str, dict] = {}
+    for row in cur.fetchall():
+        col_id, col_name, col_type, nullable, is_pk, table_id = row
+        info = {
+            "id": col_id,
+            "name": col_name,
+            "type": col_type,
+            "nullable": nullable,
+            "primary_key": is_pk,
+            "table_id": table_id,
+        }
+        col_id_to_info[col_id] = info
+        columns_by_table.setdefault(table_id, []).append(info)
+
+    # 3. Fetch links (joins)
+    cur.execute(
+        """
+        SELECT from_table_id, from_column_id, to_table_id, to_column_id, relationship_type
+        FROM column_links
+        WHERE from_table_id = ANY(%s::uuid[]) OR to_table_id = ANY(%s::uuid[])
+        """,
+        (list(table_ids), list(table_ids)),
+    )
+    # Group links by from_table
+    joins_by_table: dict[str, list] = {}
+    for row in cur.fetchall():
+        from_table_id, from_col_id, to_table_id, to_col_id, rel_type = row
+        # Only include joins where both sides are in our connection
+        if from_table_id not in tables or to_table_id not in tables:
+            continue
+        joins_by_table.setdefault(from_table_id, []).append(
+            {
+                "from_col_id": from_col_id,
+                "to_table_id": to_table_id,
+                "to_col_id": to_col_id,
+                "relationship_type": rel_type,
+            }
+        )
+
+    conn.close()
+
+    # 4. Generate YAML files
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    for table_id, table_name in sorted(tables.items(), key=lambda x: x[1]):
+        cols = columns_by_table.get(table_id, [])
+        joins = joins_by_table.get(table_id, [])
+
+        # Find primary key columns
+        pk_cols = [c for c in cols if c["primary_key"]]
+        if pk_cols:
+            grain = [c["name"] for c in pk_cols]
+        else:
+            # Fallback: use row_id if present, else first column
+            row_id_col = next((c for c in cols if c["name"] == "row_id"), None)
+            if row_id_col:
+                grain = ["row_id"]
+            elif cols:
+                grain = [cols[0]["name"]]
+            else:
+                grain = [table_name + "_id"]
+
+        # Build column definitions
+        yaml_columns = []
+        for c in cols:
+            sl_type = map_type(c["type"], c["name"])
+            col_def: dict = {"name": c["name"], "type": sl_type}
+            if is_time_column(c["name"], c["type"]):
+                col_def["role"] = "time"
+            yaml_columns.append(col_def)
+
+        # Build join definitions
+        yaml_joins = []
+        # Track target sources to handle aliases for multiple joins to same target
+        target_counts: dict[str, int] = {}
+        for j in joins:
+            to_name = tables.get(j["to_table_id"])
+            if not to_name:
+                continue
+            target_counts[to_name] = target_counts.get(to_name, 0) + 1
+
+        target_seen: dict[str, int] = {}
+        for j in joins:
+            to_name = tables.get(j["to_table_id"])
+            from_col = col_id_to_info.get(j["from_col_id"], {}).get("name")
+            to_col = col_id_to_info.get(j["to_col_id"], {}).get("name")
+            if not (to_name and from_col and to_col):
+                continue
+
+            rel = j["relationship_type"].lower()
+
+            join_def: dict = {
+                "to": to_name,
+                "on": f"{from_col} = {to_name}.{to_col}",
+                "relationship": rel,
+            }
+
+            # Add alias if multiple joins to same target
+            target_seen[to_name] = target_seen.get(to_name, 0) + 1
+            if target_counts.get(to_name, 0) > 1:
+                join_def["alias"] = f"{to_name}_{target_seen[to_name]}"
+
+            yaml_joins.append(join_def)
+
+        # Build source definition
+        source: dict = {
+            "name": table_name,
+            "table": table_name,
+        }
+        if grain:
+            source["grain"] = grain
+        source["columns"] = yaml_columns
+        if yaml_joins:
+            source["joins"] = yaml_joins
+
+        # Write YAML
+        filepath = os.path.join(OUTPUT_DIR, f"{table_name}.yaml")
+        with open(filepath, "w") as f:
+            yaml.dump(
+                source, f, default_flow_style=False, sort_keys=False, allow_unicode=True
+            )
+
+    print(f"Generated {len(tables)} source files in {OUTPUT_DIR}")
+
+
+if __name__ == "__main__":
+    main()