sqlmesh.integrations.dlt

  1import typing as t
  2import click
  3from datetime import datetime, timedelta, timezone
  4from pydantic import ValidationError
  5from sqlglot import exp, parse_one
  6from sqlmesh.core.config.connection import parse_connection_config
  7from sqlmesh.core.context import Context
  8from sqlmesh.utils.date import yesterday_ds
  9
 10
 11def generate_dlt_models_and_settings(
 12    pipeline_name: str,
 13    dialect: str,
 14    tables: t.Optional[t.List[str]] = None,
 15    dlt_path: t.Optional[str] = None,
 16) -> t.Tuple[t.Set[t.Tuple[str, str]], t.Optional[str], str]:
 17    """
 18    This function attaches to a DLT pipeline and retrieves the connection configs and
 19    SQLMesh models based on the tables present in the pipeline's default schema.
 20
 21    Args:
 22        pipeline_name: The name of the DLT pipeline to attach to.
 23        dialect: The SQL dialect to use for generating SQLMesh models.
 24        tables: A list of table names to include.
 25        dlt_path: The path to the DLT pipelines working directory, where DLT stores
 26            pipeline state (by default ~/.dlt/pipelines).
 27
 28    Returns:
 29        A tuple containing a set of the SQLMesh model definitions, the connection config and the start date.
 30    """
 31
 32    import dlt
 33    from dlt.common.schema.utils import has_table_seen_data, is_complete_column
 34    from dlt.pipeline.exceptions import CannotRestorePipelineException
 35
 36    try:
 37        pipeline = dlt.attach(pipeline_name=pipeline_name, pipelines_dir=dlt_path or "")
 38    except CannotRestorePipelineException as e:
 39        from pathlib import Path
 40        from dlt.common.pipeline import get_dlt_pipelines_dir
 41
 42        searched_dir = dlt_path or get_dlt_pipelines_dir()
 43        msg = f"Could not attach to pipeline {pipeline_name}.\nSearched in: {searched_dir}\n{e}"
 44        if dlt_path and (Path(get_dlt_pipelines_dir()) / pipeline_name).exists():
 45            msg += (
 46                f"\nHint: A pipeline named '{pipeline_name}' exists in the default pipelines "
 47                f"working directory '{get_dlt_pipelines_dir()}'. Note that --dlt-path must "
 48                "point to the directory where DLT stores pipeline working state (by default "
 49                "~/.dlt/pipelines), not the directory containing your pipeline scripts. "
 50                "Try omitting --dlt-path."
 51            )
 52        raise click.ClickException(msg)
 53
 54    schema = pipeline.default_schema
 55    dataset = pipeline.dataset_name
 56
 57    # Get the start date from the load_ids
 58    storage_ids = list(pipeline._get_load_storage().list_loaded_packages())
 59    start_date = get_start_date(storage_ids)
 60
 61    # Get the connection credentials
 62    db_type = pipeline.destination.to_name(pipeline.destination)
 63    if db_type == "filesystem":
 64        connection_config = None
 65    else:
 66        if dlt.__version__ >= "1.10.0":
 67            client = pipeline.destination_client()
 68        else:
 69            client = pipeline._sql_job_client(schema)  # type: ignore
 70        config = client.config
 71        credentials = config.credentials
 72        configs = {
 73            key: value
 74            for key in dir(credentials)
 75            if not key.startswith("_")
 76            and not callable(value := getattr(credentials, key))
 77            and value is not None
 78        }
 79        connection_config = format_config(configs, db_type)
 80
 81    dlt_tables = {
 82        name: table
 83        for name, table in schema.tables.items()
 84        if (
 85            (has_table_seen_data(table) and not name.startswith(schema._dlt_tables_prefix))
 86            or name == schema.loads_table_name
 87        )
 88        and (name in tables if tables else True)
 89    }
 90
 91    sqlmesh_models = set()
 92    for table_name, table in dlt_tables.items():
 93        dlt_columns = {}
 94        primary_key = []
 95
 96        # is_complete_column returns true if column contains a name and a data type
 97        for col in filter(is_complete_column, table["columns"].values()):
 98            dlt_columns[col["name"]] = exp.DataType.build(str(col["data_type"]), dialect=dialect)
 99            if col.get("primary_key"):
100                primary_key.append(str(col["name"]))
101
102        load_id = next(
103            (col for col in ["_dlt_load_id", "load_id"] if col in dlt_columns),
104            None,
105        )
106        load_key = "c." + load_id if load_id else ""
107        parent_table = None
108
109        # Handling for nested tables: https://dlthub.com/docs/general-usage/destination-tables#nested-tables
110        if not load_id:
111            if (
112                "_dlt_parent_id" in dlt_columns
113                and (parent_table := table["parent"])
114                and parent_table in dlt_tables
115            ):
116                load_key = "p._dlt_load_id"
117                parent_table = dataset + "." + parent_table
118            else:
119                break
120
121        column_types = [
122            exp.cast(exp.column(column, table="c"), data_type, dialect=dialect)
123            .as_(column)
124            .sql(dialect=dialect)
125            for column, data_type in dlt_columns.items()
126            if isinstance(column, str)
127        ]
128        select_columns = (
129            ",\n".join(f"  {column_name}" for column_name in column_types) if column_types else ""
130        )
131
132        grain = f"\n  grain ({', '.join(primary_key)})," if primary_key else ""
133        incremental_model_name = f"{dataset}_sqlmesh.incremental_{table_name}"
134        incremental_model_sql = generate_incremental_model(
135            incremental_model_name,
136            select_columns,
137            grain,
138            dataset + "." + table_name,
139            dialect,
140            load_key,
141            parent_table,
142        )
143        sqlmesh_models.add((incremental_model_name, incremental_model_sql))
144
145    return sqlmesh_models, connection_config, start_date
146
147
148def generate_dlt_models(
149    context: Context,
150    pipeline_name: str,
151    tables: t.List[str],
152    force: bool,
153    dlt_path: t.Optional[str] = None,
154) -> t.List[str]:
155    from sqlmesh.cli.project_init import _create_object_files
156
157    sqlmesh_models, _, _ = generate_dlt_models_and_settings(
158        pipeline_name=pipeline_name,
159        dialect=context.config.dialect or "",
160        tables=tables if tables else None,
161        dlt_path=dlt_path,
162    )
163
164    if not tables and not force:
165        existing_models = [m.name for m in context.models.values()]
166        sqlmesh_models = {model for model in sqlmesh_models if model[0] not in existing_models}
167
168    if sqlmesh_models:
169        _create_object_files(
170            context.path / "models",
171            {model[0].split(".")[-1]: model[1] for model in sqlmesh_models},
172            "sql",
173        )
174        return [model[0] for model in sqlmesh_models]
175    return []
176
177
178def generate_incremental_model(
179    model_name: str,
180    select_columns: str,
181    grain: str,
182    from_table: str,
183    dialect: str,
184    load_id: str,
185    parent_table: t.Optional[str] = None,
186) -> str:
187    """Generate the SQL definition for an incremental model."""
188
189    time_column = parse_one(f"to_timestamp(CAST({load_id} AS DOUBLE))").sql(dialect=dialect)
190
191    from_clause = f"{from_table} as c"
192    if parent_table:
193        from_clause += f"""\nJOIN
194  {parent_table} as p
195ON
196  c._dlt_parent_id = p._dlt_id"""
197
198    return f"""MODEL (
199  name {model_name},
200  kind INCREMENTAL_BY_TIME_RANGE (
201    time_column _dlt_load_time,
202  ),{grain}
203);
204
205SELECT
206{select_columns},
207  {time_column} as _dlt_load_time
208FROM
209  {from_clause}
210WHERE
211  {time_column} BETWEEN @start_ts AND @end_ts
212"""
213
214
215def format_config(configs: t.Dict[str, str], db_type: str) -> str:
216    """Generate a string for the gateway connection config."""
217    config = {
218        "type": db_type,
219    }
220
221    for key, value in configs.items():
222        if key == "password":
223            config[key] = f'"{value}"'
224        elif key == "username":
225            config["user"] = value
226        else:
227            config[key] = value
228
229    # Validate the connection config fields
230    invalid_fields = []
231    try:
232        parse_connection_config(config)
233    except ValidationError as e:
234        for error in e.errors():
235            invalid_fields.append(error.get("loc", [])[0])
236
237    return "\n".join(
238        [f"      {key}: {value}" for key, value in config.items() if key not in invalid_fields]
239    )
240
241
242def get_start_date(load_ids: t.List[str]) -> str:
243    """Convert the earliest load_id to UTC timestamp, subtract a day and format as 'YYYY-MM-DD'."""
244
245    timestamps = [datetime.fromtimestamp(float(id), tz=timezone.utc) for id in load_ids]
246    if timestamps:
247        start_timestamp = min(timestamps) - timedelta(days=1)
248        return start_timestamp.strftime("%Y-%m-%d")
249    return yesterday_ds()