clickzetta-semantic-model-generator 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

semantic_model_generator/tests/relationship_discovery_test.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+import pandas as pd
+
+from semantic_model_generator.relationships.discovery import (
+    discover_relationships_from_schema,
+)
+
+
+class _FakeResult:
+    def __init__(self, df: pd.DataFrame):
+        self._df = df
+
+    def to_pandas(self) -> pd.DataFrame:
+        return self._df.copy()
+
+
+class _FakeSession:
+    def __init__(self, tables: List[str], columns_df: pd.DataFrame):
+        self.tables = tables
+        self.columns_df = columns_df
+
+    def sql(self, query: str):
+        normalized = query.upper()
+        if "SHOW CATALOGS" in normalized:
+            return _FakeResult(
+                pd.DataFrame(
+                    {
+                        "CATALOG_NAME": ["CLICKZETTA_SAMPLE_DATA"],
+                        "CATEGORY": ["MANAGED"],
+                    }
+                )
+            )
+        if "INFORMATION_SCHEMA.COLUMNS" in normalized:
+            return _FakeResult(self.columns_df)
+        if "FROM INFORMATION_SCHEMA.TABLES" in normalized:
+            data = {
+                "TABLE_SCHEMA": ["TPCH_100G"] * len(self.tables),
+                "TABLE_NAME": self.tables,
+            }
+            return _FakeResult(pd.DataFrame(data))
+        if "SELECT DISTINCT" in normalized:
+            # Return single column of sample values
+            return _FakeResult(pd.DataFrame({"VALUE": [1, 2, 3]}))
+        raise AssertionError(f"Unexpected query: {query}")
+
+
+def _build_columns_df() -> pd.DataFrame:
+    records: List[Dict[str, Any]] = []
+    # Orders table
+    records.extend(
+        [
+            {
+                "TABLE_SCHEMA": "TPCH_100G",
+                "TABLE_NAME": "ORDERS",
+                "COLUMN_NAME": "ORDER_ID",
+                "DATA_TYPE": "NUMBER",
+                "IS_PRIMARY_KEY": True,
+            },
+            {
+                "TABLE_SCHEMA": "TPCH_100G",
+                "TABLE_NAME": "ORDERS",
+                "COLUMN_NAME": "CUSTOMER_ID",
+                "DATA_TYPE": "NUMBER",
+                "IS_PRIMARY_KEY": False,
+            },
+        ]
+    )
+    # Customer table
+    records.extend(
+        [
+            {
+                "TABLE_SCHEMA": "TPCH_100G",
+                "TABLE_NAME": "CUSTOMER",
+                "COLUMN_NAME": "CUSTOMER_ID",
+                "DATA_TYPE": "NUMBER",
+                "IS_PRIMARY_KEY": True,
+            },
+            {
+                "TABLE_SCHEMA": "TPCH_100G",
+                "TABLE_NAME": "CUSTOMER",
+                "COLUMN_NAME": "NAME",
+                "DATA_TYPE": "STRING",
+                "IS_PRIMARY_KEY": False,
+            },
+        ]
+    )
+    return pd.DataFrame.from_records(records)
+
+
+def test_discover_relationships_from_schema_builds_relationships():
+    tables = ["ORDERS", "CUSTOMER"]
+    columns_df = _build_columns_df()
+    session = _FakeSession(tables, columns_df)
+
+    result = discover_relationships_from_schema(
+        session=session,
+        workspace="CLICKZETTA_SAMPLE_DATA",
+        schema="TPCH_100G",
+        strict_join_inference=False,
+    )
+
+    assert result.summary.total_tables == 2
+    assert result.summary.total_relationships_found >= 1
+
+    names = {rel.name for rel in result.relationships}
+    assert any("ORDERS" in name and "CUSTOMER" in name for name in names)
+
+    left_tables = {rel.left_table for rel in result.relationships}
+    right_tables = {rel.right_table for rel in result.relationships}
+    assert "ORDERS" in left_tables
+    assert "CUSTOMER" in right_tables

semantic_model_generator/tests/relationships_filters_test.py

@@ -7,9 +7,24 @@ def test_suggest_filters_builds_in_clause_and_time_filter() -> None:
         id_=0,
         name="ORDERS",
         columns=[
-            Column(id_=0, column_name="status", column_type="STRING", values=["OPEN", "CLOSED", "OPEN"]),
-            Column(id_=1, column_name="order_date", column_type="TIMESTAMP", values=["2024-01-01", "2024-01-15"]),
-            Column(id_=2, column_name="created_at", column_type="STRING", values=["2024-01-05 10:00:00"]),
+            Column(
+                id_=0,
+                column_name="status",
+                column_type="STRING",
+                values=["OPEN", "CLOSED", "OPEN"],
+            ),
+            Column(
+                id_=1,
+                column_name="order_date",
+                column_type="TIMESTAMP",
+                values=["2024-01-01", "2024-01-15"],
+            ),
+            Column(
+                id_=2,
+                column_name="created_at",
+                column_type="STRING",
+                values=["2024-01-05 10:00:00"],
+            ),
         ],
     )
 
@@ -26,7 +41,12 @@ def test_infer_relationships_uses_pk_candidate() -> None:
         id_=0,
         name="CUSTOMERS",
         columns=[
-            Column(id_=0, column_name="customer_id", column_type="INT", values=["1", "2", "3"]),
+            Column(
+                id_=0,
+                column_name="customer_id",
+                column_type="INT",
+                values=["1", "2", "3"],
+            ),
             Column(id_=1, column_name="customer_name", column_type="STRING"),
         ],
     )
@@ -34,15 +54,35 @@ def test_infer_relationships_uses_pk_candidate() -> None:
         id_=1,
         name="ORDERS",
         columns=[
-            Column(id_=0, column_name="order_id", column_type="INT", values=["10", "11", "12"]),
-            Column(id_=1, column_name="customer_id", column_type="INT", values=["1", "2", "1"]),
+            Column(
+                id_=0,
+                column_name="order_id",
+                column_type="INT",
+                values=["10", "11", "12"],
+            ),
+            Column(
+                id_=1,
+                column_name="customer_id",
+                column_type="INT",
+                values=["1", "2", "1"],
+            ),
         ],
     )
 
     relationships = generate_model._infer_relationships(
         [
-            (FQNParts(database="QUICK_START", schema_name="MCP_DEMO", table="CUSTOMERS"), customers_table),
-            (FQNParts(database="QUICK_START", schema_name="MCP_DEMO", table="ORDERS"), orders_table),
+            (
+                FQNParts(
+                    database="QUICK_START", schema_name="MCP_DEMO", table="CUSTOMERS"
+                ),
+                customers_table,
+            ),
+            (
+                FQNParts(
+                    database="QUICK_START", schema_name="MCP_DEMO", table="ORDERS"
+                ),
+                orders_table,
+            ),
         ]
     )
 
@@ -61,23 +101,46 @@ def test_infer_relationships_matches_synonym_keys() -> None:
         id_=0,
         name="ORDERS",
         columns=[
-            Column(id_=0, column_name="o_orderkey", column_type="INT", values=["1", "2", "3"]),
-            Column(id_=1, column_name="o_custkey", column_type="INT", values=["10", "20", "30"]),
+            Column(
+                id_=0,
+                column_name="o_orderkey",
+                column_type="INT",
+                values=["1", "2", "3"],
+            ),
+            Column(
+                id_=1,
+                column_name="o_custkey",
+                column_type="INT",
+                values=["10", "20", "30"],
+            ),
         ],
     )
     lineitem_table = Table(
         id_=1,
         name="LINEITEM",
         columns=[
-            Column(id_=0, column_name="l_orderkey", column_type="INT", values=["1", "1", "2"]),
-            Column(id_=1, column_name="l_linenumber", column_type="INT", values=["1", "2", "1"]),
+            Column(
+                id_=0,
+                column_name="l_orderkey",
+                column_type="INT",
+                values=["1", "1", "2"],
+            ),
+            Column(
+                id_=1,
+                column_name="l_linenumber",
+                column_type="INT",
+                values=["1", "2", "1"],
+            ),
         ],
     )
 
     relationships = generate_model._infer_relationships(
         [
             (FQNParts(database="CAT", schema_name="SCH", table="ORDERS"), orders_table),
-            (FQNParts(database="CAT", schema_name="SCH", table="LINEITEM"), lineitem_table),
+            (
+                FQNParts(database="CAT", schema_name="SCH", table="LINEITEM"),
+                lineitem_table,
+            ),
         ]
     )
 
@@ -95,22 +158,40 @@ def test_infer_relationships_handles_part_supplier() -> None:
         id_=0,
         name="PART",
         columns=[
-            Column(id_=0, column_name="p_partkey", column_type="INT", values=["1", "2", "3"]),
+            Column(
+                id_=0,
+                column_name="p_partkey",
+                column_type="INT",
+                values=["1", "2", "3"],
+            ),
         ],
     )
     partsupp_table = Table(
         id_=1,
         name="PARTSUPP",
         columns=[
-            Column(id_=0, column_name="ps_partkey", column_type="INT", values=["1", "1", "2"]),
-            Column(id_=1, column_name="ps_suppkey", column_type="INT", values=["10", "20", "30"]),
+            Column(
+                id_=0,
+                column_name="ps_partkey",
+                column_type="INT",
+                values=["1", "1", "2"],
+            ),
+            Column(
+                id_=1,
+                column_name="ps_suppkey",
+                column_type="INT",
+                values=["10", "20", "30"],
+            ),
         ],
     )
 
     relationships = generate_model._infer_relationships(
         [
             (FQNParts(database="CAT", schema_name="SCH", table="PART"), part_table),
-            (FQNParts(database="CAT", schema_name="SCH", table="PARTSUPP"), partsupp_table),
+            (
+                FQNParts(database="CAT", schema_name="SCH", table="PARTSUPP"),
+                partsupp_table,
+            ),
         ]
     )
 
@@ -128,15 +209,30 @@ def test_infer_relationships_orders_customer() -> None:
         id_=0,
         name="ORDERS",
         columns=[
-            Column(id_=0, column_name="o_orderkey", column_type="INT", values=["1", "2", "3"]),
-            Column(id_=1, column_name="o_custkey", column_type="INT", values=["10", "20", "30"]),
+            Column(
+                id_=0,
+                column_name="o_orderkey",
+                column_type="INT",
+                values=["1", "2", "3"],
+            ),
+            Column(
+                id_=1,
+                column_name="o_custkey",
+                column_type="INT",
+                values=["10", "20", "30"],
+            ),
         ],
     )
     customer_table = Table(
         id_=1,
         name="CUSTOMER",
         columns=[
-            Column(id_=0, column_name="c_custkey", column_type="INT", values=["10", "20", "30"]),
+            Column(
+                id_=0,
+                column_name="c_custkey",
+                column_type="INT",
+                values=["10", "20", "30"],
+            ),
             Column(id_=1, column_name="c_name", column_type="STRING"),
         ],
     )
@@ -144,7 +240,10 @@ def test_infer_relationships_orders_customer() -> None:
     relationships = generate_model._infer_relationships(
         [
             (FQNParts(database="CAT", schema_name="SCH", table="ORDERS"), orders_table),
-            (FQNParts(database="CAT", schema_name="SCH", table="CUSTOMER"), customer_table),
+            (
+                FQNParts(database="CAT", schema_name="SCH", table="CUSTOMER"),
+                customer_table,
+            ),
         ]
     )
 
@@ -162,23 +261,44 @@ def test_infer_relationships_lineitem_supplier() -> None:
         id_=0,
         name="LINEITEM",
         columns=[
-            Column(id_=0, column_name="l_orderkey", column_type="INT", values=["1", "2", "3"]),
-            Column(id_=1, column_name="l_suppkey", column_type="INT", values=["100", "101", "102"]),
+            Column(
+                id_=0,
+                column_name="l_orderkey",
+                column_type="INT",
+                values=["1", "2", "3"],
+            ),
+            Column(
+                id_=1,
+                column_name="l_suppkey",
+                column_type="INT",
+                values=["100", "101", "102"],
+            ),
         ],
     )
     supplier_table = Table(
         id_=1,
         name="SUPPLIER",
         columns=[
-            Column(id_=0, column_name="s_suppkey", column_type="INT", values=["100", "101", "102"]),
+            Column(
+                id_=0,
+                column_name="s_suppkey",
+                column_type="INT",
+                values=["100", "101", "102"],
+            ),
             Column(id_=1, column_name="s_name", column_type="STRING"),
         ],
     )
 
     relationships = generate_model._infer_relationships(
         [
-            (FQNParts(database="CAT", schema_name="SCH", table="LINEITEM"), lineitem_table),
-            (FQNParts(database="CAT", schema_name="SCH", table="SUPPLIER"), supplier_table),
+            (
+                FQNParts(database="CAT", schema_name="SCH", table="LINEITEM"),
+                lineitem_table,
+            ),
+            (
+                FQNParts(database="CAT", schema_name="SCH", table="SUPPLIER"),
+                supplier_table,
+            ),
         ]
     )
 
@@ -196,7 +316,13 @@ def test_infer_relationships_handles_suffix_based_foreign_keys() -> None:
         id_=0,
         name="DIM_DATE",
         columns=[
-            Column(id_=0, column_name="date_id", column_type="INT", values=["20240101", "20240102"], is_primary_key=True),
+            Column(
+                id_=0,
+                column_name="date_id",
+                column_type="INT",
+                values=["20240101", "20240102"],
+                is_primary_key=True,
+            ),
             Column(id_=1, column_name="date_value", column_type="DATE"),
         ],
     )
@@ -204,15 +330,25 @@ def test_infer_relationships_handles_suffix_based_foreign_keys() -> None:
         id_=1,
         name="FACT_SALES",
         columns=[
-            Column(id_=0, column_name="order_id", column_type="INT", values=["10", "11"]),
-            Column(id_=1, column_name="order_date_id", column_type="INT", values=["20240101", "20240102"]),
+            Column(
+                id_=0, column_name="order_id", column_type="INT", values=["10", "11"]
+            ),
+            Column(
+                id_=1,
+                column_name="order_date_id",
+                column_type="INT",
+                values=["20240101", "20240102"],
+            ),
         ],
     )
 
     relationships = generate_model._infer_relationships(
         [
             (FQNParts(database="CAT", schema_name="SCH", table="DIM_DATE"), dim_date),
-            (FQNParts(database="CAT", schema_name="SCH", table="FACT_SALES"), fact_sales),
+            (
+                FQNParts(database="CAT", schema_name="SCH", table="FACT_SALES"),
+                fact_sales,
+            ),
         ]
     )
 

semantic_model_generator/tests/utils_test.py

@@ -1,7 +1,7 @@
 import pytest
 
-from semantic_model_generator.data_processing.data_types import FQNParts
 from semantic_model_generator.clickzetta_utils.utils import create_fqn_table
+from semantic_model_generator.data_processing.data_types import FQNParts
 
 
 def test_fqn_creation():