sqlmesh.core.schema_loader

  1from __future__ import annotations
  2
  3import typing as t
  4from concurrent.futures import ThreadPoolExecutor
  5from pathlib import Path
  6
  7from sqlglot import exp
  8from sqlglot.dialects.dialect import DialectType
  9
 10from sqlmesh.core.console import get_console
 11from sqlmesh.core.engine_adapter import EngineAdapter
 12from sqlmesh.core.model.definition import Model
 13from sqlmesh.core.state_sync import StateReader
 14from sqlmesh.utils import UniqueKeyDict, yaml
 15from sqlmesh.utils.errors import SQLMeshError
 16
 17
 18def create_external_models_file(
 19    path: Path,
 20    models: UniqueKeyDict[str, Model],
 21    adapter: EngineAdapter,
 22    state_reader: StateReader,
 23    dialect: DialectType,
 24    gateway: t.Optional[str] = None,
 25    max_workers: int = 1,
 26    strict: bool = False,
 27    all_models: t.Optional[t.Dict[str, Model]] = None,
 28) -> None:
 29    """Create or replace a YAML file with column and types of all columns in all external models.
 30
 31    Args:
 32        path: The path to store the YAML file.
 33        models: FQN to model for the current repo/config being processed.
 34        adapter: The engine adapter.
 35        state_reader: The state reader.
 36        dialect: The dialect to serialize the schema as.
 37        gateway: If the model should be associated with a specific gateway; the gateway key
 38        max_workers: The max concurrent workers to fetch columns.
 39        strict: If True, raise an error if the external model is missing in the database.
 40        all_models: FQN to model across all loaded repos. When provided, a dependency is only
 41            classified as external if it is absent from this full set. This prevents cross-repo
 42            internal models from being misclassified as external in multi-repo setups.
 43    """
 44    known_models: t.Dict[str, Model] = all_models if all_models is not None else models
 45    external_model_fqns = set()
 46
 47    for fqn, model in models.items():
 48        if model.kind.is_external:
 49            external_model_fqns.add(fqn)
 50        for dep in model.depends_on:
 51            if dep not in known_models:
 52                external_model_fqns.add(dep)
 53
 54    # Make sure we don't convert internal models into external ones.
 55    existing_model_fqns = state_reader.nodes_exist(external_model_fqns, exclude_external=True)
 56    if existing_model_fqns:
 57        existing_model_fqns_str = ", ".join(existing_model_fqns)
 58        get_console().log_warning(
 59            f"The following models already exist and can't be converted to external: {existing_model_fqns_str}. "
 60            "Perhaps these models have been removed, while downstream models that reference them weren't updated accordingly."
 61        )
 62        external_model_fqns -= existing_model_fqns
 63
 64    with ThreadPoolExecutor(max_workers=max_workers) as pool:
 65        gateway_part = {"gateway": gateway} if gateway else {}
 66
 67        schemas = [
 68            {
 69                "name": exp.to_table(table).sql(dialect=dialect),
 70                "columns": columns,
 71                **gateway_part,
 72            }
 73            for table, columns in sorted(
 74                pool.map(
 75                    lambda table: (table, get_columns(adapter, dialect, table, strict)),
 76                    external_model_fqns,
 77                )
 78            )
 79            if columns
 80        ]
 81
 82        # dont clobber existing entries from other gateways
 83        entries_to_keep = (
 84            [e for e in yaml.load(path) if e.get("gateway", None) != gateway]
 85            if path.exists()
 86            else []
 87        )
 88
 89        with open(path, "w", encoding="utf-8") as file:
 90            yaml.dump(entries_to_keep + schemas, file)
 91
 92
 93def get_columns(
 94    adapter: EngineAdapter, dialect: DialectType, table: str, strict: bool
 95) -> t.Optional[t.Dict[str, t.Any]]:
 96    """
 97    Return the column and their types in a dictionary
 98    """
 99    try:
100        columns = adapter.columns(table, include_pseudo_columns=True)
101        return {c: dtype.sql(dialect=dialect) for c, dtype in columns.items()}
102    except Exception as e:
103        msg = f"Unable to get schema for '{table}': '{e}'."
104        if strict:
105            raise SQLMeshError(msg) from e
106        get_console().log_warning(msg)
107        return None