sqlmesh.core.schema_loader

 1from __future__ import annotations
 2
 3import logging
 4import typing as t
 5from concurrent.futures import ThreadPoolExecutor
 6from pathlib import Path
 7
 8from sqlglot import exp
 9from sqlglot.dialects.dialect import DialectType
10
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
15
16logger = logging.getLogger(__name__)
17
18
19def create_schema_file(
20    path: Path,
21    models: UniqueKeyDict[str, Model],
22    adapter: EngineAdapter,
23    state_reader: StateReader,
24    dialect: DialectType,
25    max_workers: int = 1,
26) -> None:
27    """Create or replace a YAML file with model schemas.
28
29    Args:
30        path: The path to store the YAML file.
31        models: FQN to model
32        adapter: The engine adapter.
33        state_reader: The state reader.
34        dialect: The dialect to serialize the schema as.
35        max_workers: The max concurrent workers to fetch columns.
36    """
37    external_model_fqns = set()
38
39    for fqn, model in models.items():
40        if model.kind.is_external:
41            external_model_fqns.add(fqn)
42        for dep in model.depends_on:
43            if dep not in models:
44                external_model_fqns.add(dep)
45
46    # Make sure we don't convert internal models into external ones.
47    existing_model_fqns = state_reader.nodes_exist(external_model_fqns, exclude_external=True)
48    if existing_model_fqns:
49        logger.warning(
50            "The following models already exist and can't be converted to external: %s."
51            "Perhaps these models have been removed, while downstream models that reference them weren't updated accordingly",
52            ", ".join(existing_model_fqns),
53        )
54        external_model_fqns -= existing_model_fqns
55
56    with ThreadPoolExecutor(max_workers=max_workers) as pool:
57
58        def _get_columns(table: str) -> t.Optional[t.Dict[str, t.Any]]:
59            try:
60                return adapter.columns(table, include_pseudo_columns=True)
61            except Exception as e:
62                logger.warning(f"Unable to get schema for '{table}': '{e}'.")
63                return None
64
65        schemas = [
66            {
67                "name": exp.to_table(table).sql(dialect=dialect),
68                "columns": {c: dtype.sql(dialect=dialect) for c, dtype in columns.items()},
69            }
70            for table, columns in sorted(
71                pool.map(
72                    lambda table: (table, _get_columns(table)),
73                    external_model_fqns,
74                )
75            )
76            if columns
77        ]
78
79        with open(path, "w", encoding="utf-8") as file:
80            yaml.dump(schemas, file)