trellis-datamodel 0.3.3__py3-none-any.whl

trellis_datamodel/adapters/dbt_core.py

@@ -0,0 +1,975 @@
+"""
+dbt-core adapter implementation.
+
+Handles parsing dbt manifest.json/catalog.json and generating dbt schema YAML files.
+"""
+
+import copy
+import json
+import os
+import re
+import yaml
+from pathlib import Path
+from typing import Any, Optional
+
+from trellis_datamodel.utils.yaml_handler import YamlHandler
+from .base import (
+    ColumnInfo,
+    ColumnSchema,
+    ModelInfo,
+    ModelSchema,
+    Relationship,
+)
+
+
+class DbtCoreAdapter:
+    """Adapter for dbt-core transformation framework."""
+
+    def __init__(
+        self,
+        manifest_path: str,
+        catalog_path: str,
+        project_path: str,
+        data_model_path: str,
+        model_paths: list[str],
+    ):
+        self.manifest_path = manifest_path
+        self.catalog_path = catalog_path
+        self.project_path = project_path
+        self.data_model_path = data_model_path
+        self.model_paths = model_paths
+        self.yaml_handler = YamlHandler()
+
+    def _load_catalog(self) -> Optional[dict]:
+        """Load catalog.json if it exists."""
+        if not os.path.exists(self.catalog_path):
+            return None
+        try:
+            with open(self.catalog_path, "r") as f:
+                return json.load(f)
+        except Exception as exc:
+            print(f"Warning: failed to read catalog at {self.catalog_path}: {exc}")
+            return None
+
+    def _load_manifest(self) -> dict:
+        """Load manifest.json."""
+        with open(self.manifest_path, "r") as f:
+            return json.load(f)
+
+    def _load_data_model(self) -> dict:
+        """Load data model YAML if it exists."""
+        if not self.data_model_path or not os.path.exists(self.data_model_path):
+            return {}
+        try:
+            with open(self.data_model_path, "r") as f:
+                return yaml.safe_load(f) or {}
+        except Exception as e:
+            print(f"Warning: Could not load data model: {e}")
+            return {}
+
+    def get_model_dirs(self) -> list[str]:
+        """
+        Resolve configured models directories, normalizing common dbt prefixes.
+
+        Users may configure entries like "3_core", "models/3_entity", or absolute
+        paths. We normalize these to real directories so downstream scans work.
+        """
+
+        def _normalize(subdir: str) -> str:
+            # Absolute path - return as-is
+            if os.path.isabs(subdir):
+                return os.path.abspath(subdir).rstrip(os.sep)
+
+            # Remove leading "./"
+            while subdir.startswith("./"):
+                subdir = subdir[2:]
+
+            # Strip an optional leading "models/" so we don't double-prepend
+            prefix = f"models{os.sep}"
+            if subdir.startswith(prefix):
+                subdir = subdir[len(prefix) :]
+
+            return os.path.abspath(
+                os.path.join(self.project_path, "models", subdir)
+            ).rstrip(os.sep)
+
+        if self.model_paths:
+            # Remove duplicates while preserving order
+            seen = set()
+            normalized = []
+            for path in self.model_paths:
+                norm = _normalize(path)
+                if norm not in seen:
+                    seen.add(norm)
+                    normalized.append(norm)
+            return normalized
+
+        return [
+            os.path.abspath(os.path.join(self.project_path, "models")).rstrip(os.sep)
+        ]
+
+    def _entity_to_model_name(self, entity: dict[str, Any]) -> str:
+        """
+        Resolve the dbt model name for an entity.
+
+        Prefers the bound dbt_model (strips project prefix), otherwise falls back to
+        the entity ID so unbound entities still persist somewhere.
+        """
+        dbt_model = entity.get("dbt_model")
+        if dbt_model:
+            # dbt unique_id for versioned models looks like model.<project>.<name>.v2
+            parts = dbt_model.split(".")
+            if len(parts) >= 2 and re.match(r"v\d+$", parts[-1]):
+                # Use the model name part (the element before the vN suffix)
+                return parts[-2]
+            return parts[-1]
+        return entity.get("id") or ""
+
+    def _build_model_keys(self, base: str, version: Optional[str] = None) -> list[str]:
+        """
+        Generate a set of lookup keys for a model, including versioned variants.
+        """
+        keys = [base]
+        if version is not None:
+            # Version may come through as int from YAML; normalize to string
+            version_str = str(version)
+            # Support common ref patterns:
+            # - ref('model', v=2)   -> base.v2
+            # - alias names         -> base_v2
+            # - fully qualified     -> base.v2 and base_v2
+            version_num = version_str.lstrip("v")
+            keys.extend(
+                [
+                    f"{base}.v{version_num}",
+                    f"{base}_v{version_num}",
+                    f"{base}v{version_num}",
+                ]
+            )
+        # Deduplicate while preserving order
+        seen = set()
+        ordered_keys = []
+        for k in keys:
+            if k not in seen:
+                seen.add(k)
+                ordered_keys.append(k)
+        return ordered_keys
+
+    def _get_model_to_entity_map(self) -> dict[str, str]:
+        """Build mapping from model names (with version aliases) to entity IDs."""
+        model_to_entity: dict[str, str] = {}
+        data_model = self._load_data_model()
+        entities = data_model.get("entities", [])
+        for entity in entities:
+            entity_id = entity.get("id")
+            dbt_model = entity.get("dbt_model")
+            if dbt_model:
+                parts = dbt_model.split(".")
+                version_part = None
+                if len(parts) >= 2 and re.match(r"v\d+$", parts[-1]):
+                    version_part = parts[-1].lstrip("v")
+                    base_name = parts[-2]
+                else:
+                    base_name = parts[-1]
+
+                # Map the raw unique_id as well as base/version variants
+                model_to_entity[dbt_model] = entity_id
+                for key in self._build_model_keys(base_name, version_part):
+                    model_to_entity[key] = entity_id
+            # Map additional models to the same entity
+            additional_models = entity.get("additional_models", [])
+            for add_model in additional_models:
+                parts = add_model.split(".")
+                version_part = None
+                if len(parts) >= 2 and re.match(r"v\d+$", parts[-1]):
+                    version_part = parts[-1].lstrip("v")
+                    base_name = parts[-2]
+                else:
+                    base_name = parts[-1]
+
+                model_to_entity[add_model] = entity_id
+                for key in self._build_model_keys(base_name, version_part):
+                    model_to_entity[key] = entity_id
+            if entity_id:
+                model_to_entity[entity_id] = entity_id
+        return model_to_entity
+
+    def _get_model_yml_path(
+        self, model_name: str, target_version: Optional[int] = None
+    ) -> Optional[str]:
+        """Get the yml file path for a model from the manifest."""
+        if not os.path.exists(self.manifest_path):
+            return None
+
+        manifest = self._load_manifest()
+        preferred_node: Optional[dict] = None
+        fallback_node: Optional[dict] = None
+        for key, node in manifest.get("nodes", {}).items():
+            if node.get("resource_type") == "model" and node.get("name") == model_name:
+                if target_version is not None and node.get("version") == target_version:
+                    preferred_node = node
+                    break
+                if fallback_node is None:
+                    fallback_node = node
+
+        node = preferred_node or fallback_node
+        if not node:
+            return None
+
+        return self._derive_yml_path_from_node(node)
+
+    def _normalize_patch_path(self, patch_path: str) -> str:
+        """
+        Convert a manifest patch_path (which may include a scheme) into an
+        absolute filesystem path rooted at the dbt project.
+        """
+        if "://" in patch_path:
+            patch_path = patch_path.split("://", 1)[1]
+
+        patch_path = patch_path.lstrip("/")
+
+        if os.path.isabs(patch_path):
+            return patch_path
+
+        # Default: resolve relative to the dbt project directory
+        base = self.project_path or "."
+        return os.path.abspath(os.path.join(base, patch_path))
+
+    def _derive_yml_path_from_node(self, node: dict) -> Optional[str]:
+        """
+        Determine the YAML file path for a manifest node, preferring patch_path.
+        """
+        patch_path = node.get("patch_path")
+        if patch_path:
+            return self._normalize_patch_path(patch_path)
+
+        original_file_path = node.get("original_file_path", "")
+        if not original_file_path:
+            return None
+
+        sql_path = os.path.join(self.project_path, original_file_path)
+        base_path = os.path.splitext(sql_path)[0]
+        yml_path = f"{base_path}.yml"
+
+        # Common versioned layout: player_v2.sql + player.yml
+        if not os.path.exists(yml_path):
+            base_path = re.sub(r"_v\d+$", "", base_path)
+            yml_path = f"{base_path}.yml"
+
+        return yml_path
+
+    def _extract_version_from_string(self, value: str) -> Optional[int]:
+        """Extract integer version from strings like 'model.proj.player.v2'."""
+        if not value:
+            return None
+
+        match = re.search(r"\.v(\d+)$", value)
+        if match:
+            return int(match.group(1))
+
+        if value.startswith("v") and value[1:].isdigit():
+            return int(value[1:])
+
+        return None
+
+    def _resolve_model_version(
+        self, model_name: str, entity_id: str, data_model: dict[str, Any]
+    ) -> Optional[int]:
+        """
+        Resolve target version from data model (dbt_model) or manifest node.
+
+        Prefers explicit dbt_model binding with .vN suffix, falls back to
+        manifest node version.
+        """
+        for entity in data_model.get("entities", []):
+            if entity.get("id") != entity_id:
+                continue
+            dbt_model = entity.get("dbt_model") or ""
+            # Check last token after split to support fully-qualified model IDs
+            last_token = dbt_model.split(".")[-1] if dbt_model else ""
+            version = self._extract_version_from_string(last_token)
+            if version is not None:
+                return version
+
+        if os.path.exists(self.manifest_path):
+            manifest = self._load_manifest()
+            for node in manifest.get("nodes", {}).values():
+                if node.get("resource_type") != "model":
+                    continue
+                if node.get("name") != model_name:
+                    continue
+                node_version = node.get("version")
+                if node_version:
+                    return int(node_version)
+
+        return None
+
+    def _find_manifest_model_nodes(
+        self, manifest: dict[str, Any], model_name: str
+    ) -> list[dict[str, Any]]:
+        """Collect manifest nodes matching a given model name."""
+        return [
+            node
+            for node in manifest.get("nodes", {}).values()
+            if node.get("resource_type") == "model" and node.get("name") == model_name
+        ]
+
+    def _select_model_node(
+        self, candidates: list[dict[str, Any]], target_version: Optional[int]
+    ) -> Optional[dict[str, Any]]:
+        """
+        Choose the best matching manifest node.
+
+        Prefers an exact version match when requested; otherwise picks the
+        highest numbered version when available, or the first candidate.
+        """
+        if target_version is not None:
+            for node in candidates:
+                node_version = node.get("version")
+                if node_version is not None and int(node_version) == target_version:
+                    return node
+
+        if not candidates:
+            return None
+
+        versioned = [n for n in candidates if n.get("version") is not None]
+        if versioned:
+            return sorted(versioned, key=lambda n: n.get("version") or 0)[-1]
+
+        return candidates[0]
+
+    def get_models(self) -> list[ModelInfo]:
+        """Parse dbt manifest and catalog to return available models."""
+        if not os.path.exists(self.manifest_path):
+            raise FileNotFoundError(f"Manifest not found at {self.manifest_path}")
+
+        manifest = self._load_manifest()
+        catalog = self._load_catalog()
+        catalog_nodes = (catalog or {}).get("nodes", {})
+
+        models: list[ModelInfo] = []
+        for key, node in manifest.get("nodes", {}).items():
+            if node.get("resource_type") != "model":
+                continue
+
+            # Filter by path
+            original_path = node.get("original_file_path", "")
+            if self.model_paths:
+                match = any(pattern in original_path for pattern in self.model_paths)
+                if not match:
+                    continue
+
+            # Get columns from catalog or manifest
+            columns: list[ColumnInfo] = []
+            unique_id = node.get("unique_id")
+            catalog_node = catalog_nodes.get(unique_id)
+
+            if catalog_node:
+                for col in catalog_node.get("columns", {}).values():
+                    columns.append(
+                        {
+                            "name": col.get("name"),
+                            "type": col.get("type") or col.get("data_type"),
+                        }
+                    )
+            else:
+                for col_name, col_data in node.get("columns", {}).items():
+                    columns.append({"name": col_name, "type": col_data.get("type")})
+
+            # Extract materialization
+            config = node.get("config", {})
+            materialized = config.get("materialized", "view")
+
+            models.append(
+                {
+                    "unique_id": unique_id,
+                    "name": node.get("name"),
+                    "version": node.get("version"),
+                    "schema": node.get("schema"),
+                    "table": node.get("alias", node.get("name")),
+                    "columns": columns,
+                    "description": node.get("description"),
+                    "materialization": materialized,
+                    "file_path": original_path,
+                    "tags": node.get("tags", []),
+                }
+            )
+
+        models.sort(key=lambda x: x["name"])
+        return models
+
+    def get_model_schema(
+        self, model_name: str, version: Optional[int] = None
+    ) -> ModelSchema:
+        """Get the current schema definition for a specific model from its YAML file."""
+        if not os.path.exists(self.manifest_path):
+            raise FileNotFoundError(f"Manifest not found at {self.manifest_path}")
+
+        manifest = self._load_manifest()
+
+        candidate_nodes = self._find_manifest_model_nodes(manifest, model_name)
+        model_node = self._select_model_node(candidate_nodes, target_version=version)
+
+        if not model_node:
+            raise ValueError(f"Model '{model_name}' not found in manifest")
+
+        yml_path = self._derive_yml_path_from_node(model_node)
+        if not yml_path:
+            raise ValueError(
+                f"No patch_path or original_file_path found for model '{model_name}'"
+            )
+
+        data = self.yaml_handler.load_file(yml_path)
+        if not data:
+            return {
+                "model_name": model_name,
+                "description": "",
+                "columns": [],
+                "tags": [],
+                "file_path": yml_path,
+            }
+
+        model_entry = self.yaml_handler.find_model(data, model_name)
+        if not model_entry:
+            return {
+                "model_name": model_name,
+                "description": "",
+                "columns": [],
+                "tags": [],
+                "file_path": yml_path,
+            }
+
+        target_version = version if version is not None else model_node.get("version")
+        version_entry = None
+        versions = model_entry.get("versions") or []
+        if versions:
+            if target_version is not None:
+                for ver in versions:
+                    if ver.get("v") == target_version:
+                        version_entry = ver
+                        break
+            if version_entry is None:
+                version_entry = sorted(versions, key=lambda ver: ver.get("v") or 0)[-1]
+
+        # Prefer versioned block if present
+        node_for_schema = version_entry or model_entry
+
+        columns = self.yaml_handler.get_columns(node_for_schema)
+        tags = self.yaml_handler.get_model_tags(node_for_schema)
+        return {
+            "model_name": model_name,
+            "description": node_for_schema.get("description", ""),
+            "columns": columns,
+            "tags": tags,
+            "file_path": yml_path,
+        }
+
+    def save_model_schema(
+        self,
+        model_name: str,
+        columns: list[ColumnSchema],
+        description: Optional[str] = None,
+        tags: Optional[list[str]] = None,
+        version: Optional[int] = None,
+    ) -> Path:
+        """Save/update the schema definition for a model."""
+        if not os.path.exists(self.manifest_path):
+            raise FileNotFoundError(f"Manifest not found at {self.manifest_path}")
+
+        manifest = self._load_manifest()
+
+        candidate_nodes = self._find_manifest_model_nodes(manifest, model_name)
+        model_node = self._select_model_node(candidate_nodes, target_version=version)
+
+        if not model_node:
+            raise ValueError(f"Model '{model_name}' not found in manifest")
+
+        yml_path = self._derive_yml_path_from_node(model_node)
+        if not yml_path:
+            raise ValueError(
+                f"No patch_path or original_file_path found for model '{model_name}'"
+            )
+
+        data = self.yaml_handler.load_file(yml_path)
+        if not data:
+            data = {"version": 2, "models": []}
+
+        model_entry = self.yaml_handler.ensure_model(data, model_name)
+
+        target_version = version if version is not None else model_node.get("version")
+        if target_version is not None:
+            # Versioned model: update version entry and keep latest_version in sync (non-decreasing)
+            self.yaml_handler.set_latest_version(model_entry, target_version)
+            version_entry = self.yaml_handler.ensure_model_version(
+                model_entry, target_version
+            )
+
+            if description is not None:
+                self.yaml_handler.update_model_description(version_entry, description)
+                # Keep top-level description aligned when present
+                self.yaml_handler.update_model_description(model_entry, description)
+
+            self.yaml_handler.update_columns_batch(version_entry, columns)
+
+            if tags is not None:
+                self.yaml_handler.update_version_tags(version_entry, tags)
+        else:
+            # Non-versioned model
+            if description is not None:
+                self.yaml_handler.update_model_description(model_entry, description)
+
+            self.yaml_handler.update_columns_batch(model_entry, columns)
+
+            if tags is not None:
+                self.yaml_handler.update_model_tags(model_entry, tags)
+
+        self.yaml_handler.save_file(yml_path, data)
+        return Path(yml_path)
+
+    def _parse_ref(self, ref_value: str) -> tuple[str, Optional[str]]:
+        """
+        Parse ref() targets, supporting optional version arguments.
+
+        Examples:
+            ref('player') -> ("player", None)
+            ref('player', v=1) -> ("player", "1")
+            ref("player", version=2) -> ("player", "2")
+        """
+        ref_pattern = (
+            r"ref\(\s*['\"]([^,'\"]+)['\"](?:\s*,\s*(?:v|version)\s*=\s*([0-9]+))?\s*\)"
+        )
+        match = re.fullmatch(ref_pattern, ref_value.strip())
+        if match:
+            return match.group(1), match.group(2)
+        return ref_value, None
+
+    def _resolve_entity_id(
+        self, model_to_entity: dict[str, str], base_name: str, version: Optional[str]
+    ) -> str:
+        """
+        Resolve an entity id using base model name plus optional version.
+        Falls back to the base name if no mapping is found.
+        """
+        for key in self._build_model_keys(base_name, version):
+            if key in model_to_entity:
+                return model_to_entity[key]
+        # Fallback: try the raw base name
+        return model_to_entity.get(base_name, base_name)
+
+    def infer_relationships(self, include_unbound: bool = False) -> list[Relationship]:
+        """Scan dbt yml files and infer entity relationships from relationship tests.
+        
+        When include_unbound=True, returns ALL relationships found in dbt yml files
+        using raw model names. The frontend is responsible for mapping model names
+        to entity IDs based on current canvas state (which may not be saved yet).
+        """
+        model_dirs = self.get_model_dirs()
+        model_to_entity = self._get_model_to_entity_map()
+
+        # Only keep relationships where both ends map to entities that are bound to
+        # at least one dbt model (including additional_models). This prevents writing
+        # relationships for unbound entities in large projects.
+        # When include_unbound is True, skip filtering entirely - return all relationships
+        # using raw model names so the frontend can map them to current canvas state.
+        bound_entities: set[str] = set()
+        if not include_unbound:
+            data_model = self._load_data_model()
+            bound_entities = {
+                e.get("id")
+                for e in data_model.get("entities", [])
+                if e.get("id") and (e.get("dbt_model") or e.get("additional_models"))
+            }
+
+        relationships: list[Relationship] = []
+        yml_found = False
+
+        for models_dir in model_dirs:
+            if not os.path.exists(models_dir):
+                continue
+
+            for root, _, files in os.walk(models_dir):
+                for filename in files:
+                    if not filename.endswith((".yml", ".yaml")):
+                        continue
+                    yml_found = True
+
+                    filepath = os.path.join(root, filename)
+                    try:
+                        with open(filepath, "r") as f:
+                            schema_data = yaml.safe_load(f) or {}
+
+                        models_list = schema_data.get("models", [])
+                        for model in models_list:
+                            base_model_name = model.get("name")
+                            if not base_model_name:
+                                continue
+
+                            # Versioned models may declare columns inside versions list
+                            version_entries = model.get("versions", [])
+                            versioned_columns = []
+                            base_columns = model.get("columns", []) or []
+
+                            if isinstance(version_entries, list) and version_entries:
+                                # Carry forward columns when a version uses "include: all"
+                                previous_columns = copy.deepcopy(base_columns)
+                                for ver in version_entries:
+                                    raw_columns = ver.get("columns", []) or []
+                                    expanded_columns: list[dict] = []
+                                    for col in raw_columns:
+                                        if isinstance(col, dict) and col.get("include") == "all":
+                                            expanded_columns.extend(copy.deepcopy(previous_columns))
+                                        else:
+                                            expanded_columns.append(col)
+
+                                    # If no columns are explicitly provided, fall back to previous set
+                                    if not expanded_columns and previous_columns:
+                                        expanded_columns = copy.deepcopy(previous_columns)
+
+                                    versioned_columns.append(
+                                        (ver.get("v") or ver.get("version"), expanded_columns)
+                                    )
+                                    previous_columns = copy.deepcopy(expanded_columns)
+                            else:
+                                versioned_columns.append((None, base_columns))
+
+                            for model_version, columns in versioned_columns:
+                                # When include_unbound, use raw model name so frontend can remap
+                                # Otherwise resolve to entity ID from saved data model
+                                if include_unbound:
+                                    entity_id = base_model_name
+                                else:
+                                    entity_id = self._resolve_entity_id(
+                                        model_to_entity, base_model_name, model_version
+                                    )
+
+                                for column in columns or []:
+                                    test_blocks = []
+                                    for key in ("tests", "data_tests"):
+                                        value = column.get(key, [])
+                                        if isinstance(value, list):
+                                            test_blocks.extend(value)
+
+                                    for test in test_blocks:
+                                        if (
+                                            not isinstance(test, dict)
+                                            or "relationships" not in test
+                                        ):
+                                            continue
+
+                                        rel_test = test["relationships"]
+                                        args = rel_test.get("arguments", {}) or {}
+
+                                        # Support both the recommended arguments block and legacy top-level keys
+                                        to_ref = rel_test.get("to", "") or args.get(
+                                            "to", ""
+                                        )
+                                        target_field = rel_test.get(
+                                            "field", ""
+                                        ) or args.get("field", "")
+
+                                        # If either ref target or field is missing, skip and log for debugging
+                                        if not to_ref or not target_field:
+                                            continue
+
+                                        target_base, target_version = self._parse_ref(
+                                            to_ref
+                                        )
+                                        
+                                        # 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
+                                            )
+
+                                            # Skip relationships where either side is not bound
+                                            if (
+                                                entity_id not in bound_entities
+                                                or target_entity_id not in bound_entities
+                                            ):
+                                                continue
+
+                                        relationships.append(
+                                            {
+                                                "source": target_entity_id,
+                                                "target": entity_id,
+                                                "label": "",
+                                                "type": "one_to_many",
+                                                "source_field": target_field,
+                                                "target_field": column.get("name"),
+                                            }
+                                        )
+                    except Exception as e:
+                        print(f"Warning: Could not parse {filepath}: {e}")
+                        continue
+
+        if not yml_found:
+            raise FileNotFoundError(
+                f"No schema yml files found under configured dbt model paths: {model_dirs}"
+            )
+
+        # Remove duplicates
+        seen: set[tuple] = set()
+        unique_relationships: list[Relationship] = []
+        for rel in relationships:
+            key = (
+                rel["source"],
+                rel["target"],
+                rel.get("source_field"),
+                rel.get("target_field"),
+            )
+            if key not in seen:
+                seen.add(key)
+                unique_relationships.append(rel)
+
+        return unique_relationships
+
+    def sync_relationships(
+        self,
+        entities: list[dict[str, Any]],
+        relationships: list[dict[str, Any]],
+    ) -> list[Path]:
+        """Sync relationship definitions from data model to dbt yml files."""
+        # Build entity lookup
+        entity_map = {e["id"]: e for e in entities if e.get("id")}
+        # Map entity_id -> dbt model name (best-effort)
+        entity_model_name: dict[str, str] = {
+            eid: self._entity_to_model_name(ent) for eid, ent in entity_map.items()
+        }
+
+        # Group relationships by target entity (the one with the FK)
+        fk_by_entity: dict[str, list[dict]] = {}
+
+        for rel in relationships:
+            source_id = rel.get("source")
+            target_id = rel.get("target")
+            rel_type = rel.get("type", "one_to_many")
+            source_field = rel.get("source_field")
+            target_field = rel.get("target_field")
+
+            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
+
+            fk_by_entity.setdefault(fk_entity, []).append(
+                {
+                    "fk_field": fk_field,
+                    "ref_entity": ref_entity,
+                    "ref_field": ref_field,
+                }
+            )
+
+        models_dir = self.get_model_dirs()[0]
+        os.makedirs(models_dir, exist_ok=True)
+
+        updated_files: list[Path] = []
+
+        for entity in entities:
+            entity_id = entity.get("id")
+            if not entity_id:
+                continue
+
+            model_name = entity_model_name.get(entity_id, entity_id)
+
+            # For bound entities, use the correct path from manifest
+            # For unbound entities, fall back to models_dir/{entity_id}.yml
+            yml_path = None
+            if entity.get("dbt_model"):
+                yml_path = self._get_model_yml_path(model_name)
+            if not yml_path:
+                yml_path = os.path.join(models_dir, f"{entity_id}.yml")
+
+            data = self.yaml_handler.load_file(yml_path)
+            if not data:
+                data = {"version": 2, "models": []}
+
+            model_entry = self.yaml_handler.ensure_model(data, model_name)
+
+            if entity.get("description"):
+                self.yaml_handler.update_model_description(
+                    model_entry, entity.get("description")
+                )
+
+            # Sync Tags
+            entity_tags = entity.get("tags")
+            if entity_tags is not None:
+                self.yaml_handler.update_model_tags(model_entry, entity_tags)
+
+            # Sync Drafted Fields
+            drafted_fields = entity.get("drafted_fields", [])
+            for field in drafted_fields:
+                f_name = field.get("name")
+                f_type = field.get("datatype")
+                f_desc = field.get("description")
+
+                if not f_name:
+                    continue
+
+                col = self.yaml_handler.ensure_column(model_entry, f_name)
+                self.yaml_handler.update_column(
+                    col, data_type=f_type, description=f_desc
+                )
+
+            # Sync Relationships (FKs)
+            fk_list = fk_by_entity.get(entity_id, [])
+            for fk_info in fk_list:
+                fk_field = fk_info["fk_field"]
+                ref_entity = fk_info["ref_entity"]
+                ref_field = fk_info["ref_field"]
+
+                # Resolve reference model name (dbt model) for relationship test
+                ref_model_name = entity_model_name.get(ref_entity, ref_entity)
+
+                col = self.yaml_handler.ensure_column(model_entry, fk_field)
+
+                if "data_type" not in col:
+                    col["data_type"] = "text"
+
+                self.yaml_handler.add_relationship_test(col, ref_model_name, ref_field)
+
+            self.yaml_handler.save_file(yml_path, data)
+            updated_files.append(Path(yml_path))
+
+        return updated_files
+
+    def save_dbt_schema(
+        self,
+        entity_id: str,
+        model_name: str,
+        fields: list[dict[str, str]],
+        description: Optional[str] = None,
+        tags: Optional[list[str]] = None,
+    ) -> Path:
+        """
+        Generate and save a dbt schema YAML file for drafted fields.
+
+        This is used for creating new schema files from the data model editor.
+        """
+        data_model = self._load_data_model()
+
+        # Use description from request if available, otherwise fallback to data model
+        entity_description = description
+        if not entity_description:
+            entities = data_model.get("entities", [])
+            for entity in entities:
+                if entity.get("id") == entity_id:
+                    entity_description = entity.get("description")
+                    break
+
+        # Build a map of field names to relationships for this entity
+        relationships = data_model.get("relationships", [])
+        field_to_relationship: dict[str, dict] = {}
+        # Map entity -> dbt model name for refs
+        entity_model_name = {
+            e.get("id"): self._entity_to_model_name(e)
+            for e in data_model.get("entities", [])
+            if e.get("id")
+        }
+
+        for rel in relationships:
+            source_id = rel.get("source")
+            target_id = rel.get("target")
+            rel_type = rel.get("type", "one_to_many")
+            source_field = rel.get("source_field")
+            target_field = rel.get("target_field")
+
+            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
+
+            if fk_entity == entity_id:
+                field_to_relationship[fk_field] = {
+                    "target_entity": ref_entity,
+                    "target_field": ref_field,
+                }
+
+        # Generate YAML content with relationship tests
+        columns = []
+        for field in fields:
+            column_dict: dict[str, Any] = {
+                "name": field["name"],
+                "data_type": field["datatype"],
+            }
+
+            if field.get("description"):
+                column_dict["description"] = field["description"]
+
+            field_name = field["name"]
+            if field_name in field_to_relationship:
+                rel_info = field_to_relationship[field_name]
+                ref_model = entity_model_name.get(
+                    rel_info["target_entity"], rel_info["target_entity"]
+                )
+                column_dict["data_tests"] = [
+                    {
+                        "relationships": {
+                            "arguments": {
+                                "to": f"ref('{ref_model}')",
+                                "field": rel_info["target_field"],
+                            }
+                        }
+                    }
+                ]
+
+            columns.append(column_dict)
+
+        target_version = self._resolve_model_version(
+            model_name=model_name, entity_id=entity_id, data_model=data_model
+        )
+
+        # Prefer manifest-derived path; fall back to default models directory
+        yml_path = self._get_model_yml_path(model_name, target_version=target_version)
+        if not yml_path:
+            models_dir = self.get_model_dirs()[0]
+            os.makedirs(models_dir, exist_ok=True)
+            yml_path = os.path.join(models_dir, f"{entity_id}.yml")
+
+        data = self.yaml_handler.load_file(yml_path)
+        if not data:
+            data = {"version": 2, "models": []}
+
+        model_entry = self.yaml_handler.ensure_model(data, model_name)
+
+        if target_version is not None:
+            # Versioned model: update version block and latest_version
+            self.yaml_handler.set_latest_version(model_entry, target_version)
+            if entity_description:
+                self.yaml_handler.update_model_description(
+                    model_entry, entity_description
+                )
+
+            version_entry = self.yaml_handler.ensure_model_version(
+                model_entry, target_version
+            )
+            self.yaml_handler.update_columns_batch(version_entry, columns)
+
+            if entity_description:
+                self.yaml_handler.update_model_description(
+                    version_entry, entity_description
+                )
+            if tags is not None:
+                self.yaml_handler.update_version_tags(version_entry, tags)
+        else:
+            # Non-versioned model: update columns and metadata directly
+            if entity_description:
+                self.yaml_handler.update_model_description(
+                    model_entry, entity_description
+                )
+
+            self.yaml_handler.update_columns_batch(model_entry, columns)
+
+            if tags is not None:
+                self.yaml_handler.update_model_tags(model_entry, tags)
+
+        self.yaml_handler.save_file(yml_path, data)
+        return Path(yml_path)