trellis-datamodel 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

trellis_datamodel/adapters/base.py

@@ -50,7 +50,7 @@ class ModelSchema(TypedDict, total=False):
     file_path: str
 
 
-class Relationship(TypedDict):
+class Relationship(TypedDict, total=False):
     """Relationship inferred from framework metadata."""
 
     source: str
@@ -59,6 +59,10 @@ class Relationship(TypedDict):
     type: str  # e.g., "one_to_many"
     source_field: str
     target_field: str
+    source_model_name: str  # name of the source model (e.g., "employee")
+    source_model_version: Optional[int]  # version of the source model if versioned
+    target_model_name: str  # name of the target model (e.g., "employee_history")
+    target_model_version: Optional[int]  # version of the target model if versioned
 
 
 class TransformationAdapter(Protocol):

trellis_datamodel/adapters/dbt_core.py

@@ -670,16 +670,32 @@ class DbtCoreAdapter:
                                         if not to_ref or not target_field:
                                             continue
 
-                                        target_base, target_version = self._parse_ref(
+                                        target_base, target_version_str = self._parse_ref(
                                             to_ref
                                         )
                                         
+                                        # Convert version string to int if present
+                                        target_version_int = None
+                                        if target_version_str:
+                                            try:
+                                                target_version_int = int(target_version_str)
+                                            except ValueError:
+                                                pass
+                                        
+                                        # Convert model_version to int if present
+                                        source_version_int = None
+                                        if model_version is not None:
+                                            try:
+                                                source_version_int = int(model_version)
+                                            except (ValueError, TypeError):
+                                                pass
+                                        
                                         # When include_unbound, use raw model name
                                         if include_unbound:
                                             target_entity_id = target_base
                                         else:
                                             target_entity_id = self._resolve_entity_id(
-                                                model_to_entity, target_base, target_version
+                                                model_to_entity, target_base, target_version_str
                                             )
 
                                             # Skip relationships where either side is not bound
@@ -697,6 +713,10 @@ class DbtCoreAdapter:
                                                 "type": "one_to_many",
                                                 "source_field": target_field,
                                                 "target_field": column.get("name"),
+                                                "source_model_name": target_base,
+                                                "source_model_version": target_version_int,
+                                                "target_model_name": base_model_name,
+                                                "target_model_version": source_version_int,
                                             }
                                         )
                     except Exception as e:
@@ -737,8 +757,13 @@ class DbtCoreAdapter:
             eid: self._entity_to_model_name(ent) for eid, ent in entity_map.items()
         }
 
-        # Group relationships by target entity (the one with the FK)
+        # Group relationships by entity (the one with the FK)
+        # FK is always on the "many" side of the relationship
         fk_by_entity: dict[str, list[dict]] = {}
+        
+        # Track all fields that appear in ANY relationship in the data model
+        # This helps us identify which relationship tests are managed by us vs manually added
+        all_relationship_fields_by_entity: dict[str, set[str]] = {}
 
         for rel in relationships:
             source_id = rel.get("source")
@@ -750,11 +775,33 @@ class DbtCoreAdapter:
             if not source_field or not target_field:
                 continue
 
-            fk_on_target = rel_type == "one_to_many"
-            fk_entity = target_id if fk_on_target else source_id
-            fk_field = target_field if fk_on_target else source_field
-            ref_entity = source_id if fk_on_target else target_id
-            ref_field = source_field if fk_on_target else target_field
+            # Determine which side has the "many" cardinality (where FK should be)
+            # Relationship types ending in "_to_many" or starting with "many_to_" have FK on target/source respectively
+            # For one_to_one, FK is typically on source (FK holder → referenced table per spec)
+            if rel_type in ("one_to_many", "one_to_zero_or_many", "zero_or_one_to_many", "zero_or_many_to_many"):
+                # FK on target (target is the "many" side)
+                fk_entity = target_id
+                fk_field = target_field
+                ref_entity = source_id
+                ref_field = source_field
+            elif rel_type in ("many_to_one", "many_to_many", "zero_or_many_to_one"):
+                # FK on source (source is the "many" side)
+                fk_entity = source_id
+                fk_field = source_field
+                ref_entity = target_id
+                ref_field = target_field
+            elif rel_type == "one_to_one":
+                # For one_to_one, FK is on source (FK holder → referenced table per spec)
+                fk_entity = source_id
+                fk_field = source_field
+                ref_entity = target_id
+                ref_field = target_field
+            else:
+                # Fallback: assume FK on target (default behavior)
+                fk_entity = target_id
+                fk_field = target_field
+                ref_entity = source_id
+                ref_field = source_field
 
             fk_by_entity.setdefault(fk_entity, []).append(
                 {
@@ -763,6 +810,10 @@ class DbtCoreAdapter:
                     "ref_field": ref_field,
                 }
             )
+            
+            # Track which fields are involved in relationships
+            all_relationship_fields_by_entity.setdefault(source_id, set()).add(source_field)
+            all_relationship_fields_by_entity.setdefault(target_id, set()).add(target_field)
 
         models_dir = self.get_model_dirs()[0]
         os.makedirs(models_dir, exist_ok=True)
@@ -816,7 +867,27 @@ class DbtCoreAdapter:
                 )
 
             # Sync Relationships (FKs)
+            # Build a map of which fields should have which relationship tests
             fk_list = fk_by_entity.get(entity_id, [])
+            fk_fields = {fk_info["fk_field"] for fk_info in fk_list}
+            
+            # Get all fields that appear in relationships for this entity
+            relationship_fields = all_relationship_fields_by_entity.get(entity_id, set())
+            
+            # Clean up: Remove relationship tests from fields that:
+            # 1. Are in a relationship in the data model (relationship_fields)
+            # 2. But are NOT currently FKs (not in fk_fields)
+            # This removes tests when relationships are moved or type changes
+            # But preserves manually added tests (not in relationship_fields)
+            if "columns" in model_entry:
+                for col in model_entry.get("columns", []):
+                    col_name = col.get("name")
+                    if col_name and col_name in relationship_fields and col_name not in fk_fields:
+                        # This field was in a relationship but is no longer an FK
+                        # Remove its relationship test
+                        self.yaml_handler.remove_relationship_test(col)
+            
+            # Now add/update relationship tests for current FKs
             for fk_info in fk_list:
                 fk_field = fk_info["fk_field"]
                 ref_entity = fk_info["ref_entity"]

trellis_datamodel/config.py

@@ -10,10 +10,24 @@ import os
 import yaml
 from pathlib import Path
 from typing import Optional
+from dataclasses import dataclass, field
 
 # Check for test mode - allows overriding config via environment
 _TEST_DIR = os.environ.get("DATAMODEL_TEST_DIR", "")
 
+
+@dataclass
+class GuidanceConfig:
+    """Configuration for entity creation guidance features."""
+    entity_wizard_enabled: bool = True
+    push_warning_enabled: bool = True
+    min_description_length: int = 10
+    disabled_guidance: list[str] = field(default_factory=list)
+
+
+# Global guidance configuration (set by load_config)
+GUIDANCE_CONFIG: GuidanceConfig = GuidanceConfig()
+
 if _TEST_DIR:
     # Test mode: use temp directory paths
     CONFIG_PATH = os.path.join(_TEST_DIR, "config.yml")
@@ -38,6 +52,8 @@ if _TEST_DIR:
     DBT_MODEL_PATHS: list[str] = ["3_core"]
     FRONTEND_BUILD_DIR: str = os.path.join(_TEST_DIR, "frontend/build")
     DBT_COMPANY_DUMMY_PATH: str = os.path.join(_TEST_DIR, "dbt_company_dummy")
+    LINEAGE_LAYERS: list[str] = []
+    GUIDANCE_CONFIG: GuidanceConfig = GuidanceConfig()
 else:
     # Production mode: will be set by load_config()
     CONFIG_PATH: str = ""
@@ -51,6 +67,7 @@ else:
     DBT_MODEL_PATHS: list[str] = []
     FRONTEND_BUILD_DIR: str = ""
     DBT_COMPANY_DUMMY_PATH: str = ""
+    LINEAGE_LAYERS: list[str] = []
 
 
 def find_config_file(config_override: Optional[str] = None) -> Optional[str]:
@@ -85,7 +102,7 @@ def find_config_file(config_override: Optional[str] = None) -> Optional[str]:
 
 def load_config(config_path: Optional[str] = None) -> None:
     """Load and resolve all paths from config file."""
-    global FRAMEWORK, MANIFEST_PATH, DATA_MODEL_PATH, DBT_MODEL_PATHS, CATALOG_PATH, DBT_PROJECT_PATH, CANVAS_LAYOUT_PATH, CANVAS_LAYOUT_VERSION_CONTROL, CONFIG_PATH, FRONTEND_BUILD_DIR, DBT_COMPANY_DUMMY_PATH
+    global FRAMEWORK, MANIFEST_PATH, DATA_MODEL_PATH, DBT_MODEL_PATHS, CATALOG_PATH, DBT_PROJECT_PATH, CANVAS_LAYOUT_PATH, CANVAS_LAYOUT_VERSION_CONTROL, CONFIG_PATH, FRONTEND_BUILD_DIR, DBT_COMPANY_DUMMY_PATH, LINEAGE_LAYERS, GUIDANCE_CONFIG
 
     # Skip loading config file in test mode (paths already set via environment)
     if _TEST_DIR:
@@ -219,6 +236,29 @@ def load_config(config_path: Optional[str] = None) -> None:
                 DBT_COMPANY_DUMMY_PATH = p
         # Note: No default set here - CLI handles fallback to cwd/dbt_company_dummy
 
+        # 10. Load lineage layers configuration
+        if "lineage_layers" in config:
+            LINEAGE_LAYERS = config["lineage_layers"]
+            if not isinstance(LINEAGE_LAYERS, list):
+                LINEAGE_LAYERS = []
+        else:
+            LINEAGE_LAYERS = []
+
+        # 11. Load guidance configuration
+        if "guidance" in config:
+            guidance_config = config["guidance"]
+            GUIDANCE_CONFIG = GuidanceConfig(
+                entity_wizard_enabled=guidance_config.get("entity_wizard_enabled", True),
+                push_warning_enabled=guidance_config.get("push_warning_enabled", True),
+                min_description_length=guidance_config.get("min_description_length", 10),
+                disabled_guidance=guidance_config.get("disabled_guidance", [])
+                if isinstance(guidance_config.get("disabled_guidance"), list)
+                else [],
+            )
+        else:
+            # Use defaults if guidance section is missing
+            GUIDANCE_CONFIG = GuidanceConfig()
+
     except Exception as e:
         print(f"Error loading config: {e}")
 
@@ -235,5 +275,7 @@ def print_config() -> None:
     print(f"Looking for canvas layout at: {CANVAS_LAYOUT_PATH}")
     print(f"Canvas layout version control: {CANVAS_LAYOUT_VERSION_CONTROL}")
     print(f"Filtering models by paths: {DBT_MODEL_PATHS}")
+    if LINEAGE_LAYERS:
+        print(f"Lineage layers: {LINEAGE_LAYERS}")
     if DBT_COMPANY_DUMMY_PATH:
         print(f"dbt company dummy path: {DBT_COMPANY_DUMMY_PATH}")

trellis_datamodel/routes/__init__.py

@@ -2,10 +2,12 @@
 from .manifest import router as manifest_router
 from .data_model import router as data_model_router
 from .schema import router as schema_router
+from .lineage import router as lineage_router
 
 __all__ = [
     "manifest_router",
     "data_model_router",
     "schema_router",
+    "lineage_router",
 ]
 

trellis_datamodel/routes/lineage.py

@@ -0,0 +1,60 @@
+"""Routes for lineage operations."""
+
+from fastapi import APIRouter, HTTPException
+import os
+
+from trellis_datamodel import config as cfg
+from trellis_datamodel.services.lineage import extract_upstream_lineage, LineageError
+
+router = APIRouter(prefix="/api", tags=["lineage"])
+
+
+@router.get("/lineage/{model_id}")
+async def get_lineage(model_id: str):
+    """
+    Get upstream table-level lineage for a given model.
+
+    Args:
+        model_id: Unique ID of the model (e.g., "model.project.model_name")
+
+    Returns:
+        JSON response with nodes, edges, and metadata
+
+    Raises:
+        404: If model not found
+        500: If lineage extraction fails
+    """
+    try:
+        # Validate paths exist
+        if not cfg.MANIFEST_PATH or not os.path.exists(cfg.MANIFEST_PATH):
+            raise HTTPException(
+                status_code=500,
+                detail=f"Manifest not found at {cfg.MANIFEST_PATH}. Please ensure manifest.json exists.",
+            )
+
+        # Extract lineage
+        lineage_data = extract_upstream_lineage(
+            manifest_path=cfg.MANIFEST_PATH,
+            catalog_path=cfg.CATALOG_PATH,
+            model_unique_id=model_id,
+        )
+
+        return lineage_data
+
+    except FileNotFoundError as e:
+        raise HTTPException(status_code=404, detail=str(e))
+    except LineageError as e:
+        # Check if it's a catalog missing error
+        error_msg = str(e)
+        if "catalog" in error_msg.lower() and "not found" in error_msg.lower():
+            raise HTTPException(
+                status_code=500,
+                detail=f"{error_msg}. Please run 'dbt docs generate' to create catalog.json",
+            )
+        raise HTTPException(status_code=500, detail=str(e))
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error extracting lineage: {str(e)}",
+        )
+

trellis_datamodel/routes/manifest.py

@@ -94,6 +94,12 @@ async def get_config_info():
         "frontend_build_dir": cfg.FRONTEND_BUILD_DIR,
         "model_paths_configured": cfg.DBT_MODEL_PATHS,
         "model_paths_resolved": model_dirs,
+        "guidance": {
+            "entity_wizard_enabled": cfg.GUIDANCE_CONFIG.entity_wizard_enabled,
+            "push_warning_enabled": cfg.GUIDANCE_CONFIG.push_warning_enabled,
+            "min_description_length": cfg.GUIDANCE_CONFIG.min_description_length,
+            "disabled_guidance": cfg.GUIDANCE_CONFIG.disabled_guidance,
+        },
     }
 
 

trellis_datamodel/server.py

@@ -14,7 +14,7 @@ from importlib.resources import files
 
 from trellis_datamodel import config as cfg
 from trellis_datamodel.config import print_config
-from trellis_datamodel.routes import manifest_router, data_model_router, schema_router
+from trellis_datamodel.routes import manifest_router, data_model_router, schema_router, lineage_router
 
 
 def create_app() -> FastAPI:
@@ -61,6 +61,7 @@ def create_app() -> FastAPI:
     app.include_router(manifest_router)
     app.include_router(data_model_router)
     app.include_router(schema_router)
+    app.include_router(lineage_router)
 
     # Mount static files AFTER API routes
     # Important: app.mount() creates a sub-application, so we mount AFTER registering API routes

trellis_datamodel/services/__init__.py

@@ -0,0 +1,2 @@
+"""Services for data lineage and other operations."""
+