From e77c7545deba0f10fe35e5c650ab32f122359fdc Mon Sep 17 00:00:00 2001 From: Ivar Olofsson Date: Mon, 29 Jun 2026 12:05:57 +0200 Subject: [PATCH] feat(dbt): Add DBT integration framework --- ibis_typing/dbt/README.md | 139 ++++++++++++++ ibis_typing/dbt/__init__.py | 14 ++ ibis_typing/dbt/dbt_ibis_constructor.py | 117 ++++++++++++ ibis_typing/dbt/dbt_model.py | 169 +++++++++++++++++ ibis_typing/dbt/dbt_model_resolver.py | 44 +++++ ibis_typing/dbt/dbt_snapshot.py | 124 +++++++++++++ ibis_typing/dbt/dbt_sql_compiler.py | 171 ++++++++++++++++++ ibis_typing/ibis_extension_method.py | 18 +- ibis_typing/revertible.py | 53 ++++++ ibis_typing/samples/dbt_models.py | 69 +++++++ .../__ibis_sql/calendar_checksum_bucket.sql | 69 +++++++ .../generated/__ibis_sql/calendar_width.sql | 6 + .../generated/__ibis_sql/timestamp_now.sql | 13 ++ ibis_typing/samples/generated/__init__.py | 0 .../tests/test_generate_ibis_dbt_sql.py | 35 ++++ 15 files changed, 1034 insertions(+), 7 deletions(-) create mode 100644 ibis_typing/dbt/README.md create mode 100644 ibis_typing/dbt/__init__.py create mode 100644 ibis_typing/dbt/dbt_ibis_constructor.py create mode 100644 ibis_typing/dbt/dbt_model.py create mode 100644 ibis_typing/dbt/dbt_model_resolver.py create mode 100644 ibis_typing/dbt/dbt_snapshot.py create mode 100644 ibis_typing/dbt/dbt_sql_compiler.py create mode 100644 ibis_typing/samples/dbt_models.py create mode 100644 ibis_typing/samples/generated/__ibis_sql/calendar_checksum_bucket.sql create mode 100644 ibis_typing/samples/generated/__ibis_sql/calendar_width.sql create mode 100644 ibis_typing/samples/generated/__ibis_sql/timestamp_now.sql create mode 100644 ibis_typing/samples/generated/__init__.py create mode 100644 ibis_typing/samples/tests/test_generate_ibis_dbt_sql.py diff --git a/ibis_typing/dbt/README.md b/ibis_typing/dbt/README.md new file mode 100644 index 0000000..9f748c9 --- /dev/null +++ b/ibis_typing/dbt/README.md @@ -0,0 +1,139 @@ +# `ibis_typing.dbt` — dbt Integration + +Bridges typed `ibis_typing` [Expression](../expression.py) classes with +[DBT](https://docs.getdbt.com/) by compiling them to DBT SQL (Jinja templates). + +--- + +## Overview + +| Class / function | Purpose | +|------------------------|--------------------------------------------------------------| +| `DbtModel` | Wraps an `Expression` as a DBT model with its `ModelConfig` | +| `DbtSource` | Wraps an `IbisDbSchema` as a DBT source table | +| `DbtSnapshot` | Wraps an `IbisSchema` as a DBT snapshot | +| `ModelConfig` | DBT model config (materialization, incremental strategy, …) | +| `SnapshotConfig` | DBT snapshot config (strategy, unique key, updated-at, …) | +| `DbtModelResolver` | Auto-discovers `Expression` implementations inside a package | +| `dbt_model_to_dbt_sql` | Compiles a `DbtModel` / `DbtSnapshot` to a DBT SQL string | + +--- + +## Quick Start + +The canonical example lives in +[`ibis_typing/samples/dbt_models.py`](../samples/dbt_models.py). +It shows how to build model and source lookups that can be fed directly into the +SQL compiler. + +```python +# ibis_typing/samples/dbt_models.py +from collections.abc import Mapping + +from ibis_typing import Expression, IbisDbSchema, IbisSchema, naming, samples +from ibis_typing.dbt import ( + DbtModel, + DbtSnapshot, + DbtSource, + ModelConfig, + SnapshotConfig, + dbt_model_resolver, +) +from ibis_typing.dbt.dbt_model import IncrementalStrategy, Materialized +from ibis_typing.dbt.dbt_snapshot import SnapshotStrategy +from ibis_typing.ibis_time import TimestampNow +from ibis_typing.samples.sample_incremental_calendar import Calendar, CalendarWidth +from ibis_typing.samples.sample_transforms import CircleParameters + +my_namespace = ("my_database", "my_schema") + + +def get_dbt_model_lookup() -> Mapping[type[Expression], DbtModel]: + resolver = dbt_model_resolver.DbtModelResolver(samples) + + models = { + # Provide timestamp as a table so every dbt run uses a fixed timestamp. + DbtModel( + TimestampNow, + config=ModelConfig( + materialized=Materialized.table, + ).with_namespace(*my_namespace), + ) + } + incremental_models = [ + DbtModel( + expr, + config=ModelConfig( + materialized=Materialized.incremental, + incremental_strategy=IncrementalStrategy.merge, + unique_key=expr.incremental_params.group_by, + ).with_namespace(*my_namespace), + ) + for expr in resolver.incremental_models + ] + snapshots = [ + DbtSnapshot( + expr=CalendarWidth, + config=SnapshotConfig( + strategy=SnapshotStrategy.timestamp, + unique_key=CalendarWidth.incremental_params.group_by, + updated_at=CalendarWidth.incremental_params.updated_at_col, + ).with_namespace(*my_namespace), + ) + ] + all_models = (*models, *incremental_models, *snapshots) + return {model.expr: model for model in all_models} + + +def get_dbt_source_lookup() -> Mapping[type[IbisSchema], DbtSource]: + sources = [Calendar, CircleParameters] + return {schema: DbtSource(as_db_schema(schema)) for schema in sources} + + +def as_db_schema(schema: type[IbisSchema]) -> type[IbisDbSchema]: + kwargs = { + "table_name": naming.snake_case(schema.__name__), + "table_namespace": my_namespace, + } + return type(schema.__name__, (IbisDbSchema, schema), kwargs) +``` + +--- + +## Compiling to dbt SQL + +Use `DbtRefTableProvider` to wire the model and source lookups together, then +call `dbt_model_to_dbt_sql` for each model. +See [`ibis_typing/samples/tests/test_generate_ibis_dbt_sql.py`](../samples/tests/test_generate_ibis_dbt_sql.py) +for the full test that also snapshots the generated SQL to disk. + +```python +from ibis_typing.dbt import dbt_sql_compiler +from ibis_typing.dbt.dbt_ibis_constructor import DbtRefTableProvider +from ibis_typing.samples import dbt_models + +model_lookup = dbt_models.get_dbt_model_lookup() +source_lookup = dbt_models.get_dbt_source_lookup() +ref_provider = DbtRefTableProvider(model_lookup, source_lookup) + +for model in model_lookup.values(): + sql = dbt_sql_compiler.dbt_model_to_dbt_sql( + model, + dialect="duckdb", + ref_provider=ref_provider, + ) + print(sql) +``` + +The compiler automatically: + +- Injects a `{{ config(...) }}` Jinja header from the model's config. +- Wraps incremental models in `{% if is_incremental() %} … {% else %} … {% endif %}`. +- Wraps snapshots in `{% snapshot %} … {% endsnapshot %}`. +- Replaces internal placeholder tokens with `{{ ref("…") }}`, `{{ source("…", "…") }}`, + and `{{ this }}`. + +## Auto-Discovery with `DbtModelResolver` + +`DbtModelResolver` scans a Python package and returns all `Expression` +subclasses grouped by category: diff --git a/ibis_typing/dbt/__init__.py b/ibis_typing/dbt/__init__.py new file mode 100644 index 0000000..6f80c42 --- /dev/null +++ b/ibis_typing/dbt/__init__.py @@ -0,0 +1,14 @@ +from .dbt_model import DbtModel, DbtSource, ModelConfig +from .dbt_model_resolver import DbtModelResolver +from .dbt_snapshot import DbtSnapshot, SnapshotConfig +from .dbt_sql_compiler import dbt_model_to_dbt_sql + +__all__ = [ + "DbtModel", + "DbtModelResolver", + "DbtSnapshot", + "DbtSource", + "ModelConfig", + "SnapshotConfig", + "dbt_model_to_dbt_sql", +] diff --git a/ibis_typing/dbt/dbt_ibis_constructor.py b/ibis_typing/dbt/dbt_ibis_constructor.py new file mode 100644 index 0000000..4fdf496 --- /dev/null +++ b/ibis_typing/dbt/dbt_ibis_constructor.py @@ -0,0 +1,117 @@ +"""Construct ibis Tables for compiling DBT SQL models.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping + +import ibis +from attrs import frozen + +from ibis_typing import Expression, IbisSchema, IbisTable, evaluator +from ibis_typing.checksum_buckets import ( + AsIncrementedBuckets, + ChecksumBuckets, + IncrementalExpression, +) +from ibis_typing.table_provider import TableProvider + +from .dbt_model import DbtModel, DbtSource + +logger = logging.getLogger(__name__) + + +def construct_dbt_model[E: Expression]( + schema: type[E], + ref_provider: DbtRefTableProvider, + *, + self: type[Expression] | None = None, + buckets_update: bool = False, +) -> IbisTable[E]: + logger.info( + f"Constructing {schema.__name__}" + if not self + else f"Constructing {schema.__name__} for {self.__name__}" + ) + self = self or schema + providers = [ + DbtSelfRefProvider(self), + *([DbtBucketProvider(schema, ref_provider)] if buckets_update else ()), + ref_provider, + ] + return evaluator.from_expression(schema, table_providers=providers) + + +def construct_dbt_bucket_increment[M: IncrementalExpression]( + expr: type[M], + ref_provider: DbtRefTableProvider, +) -> IbisTable[M]: + logger.info(f"Constructing bucket increment for {expr.__name__}") + return evaluator.construct_increment( + expr, + prior_providers=[DbtSelfRefProvider(expr)], + input_providers=[ref_provider], + ) + + +@frozen +class DbtBucketProvider(TableProvider): + expr: type[Expression] + table_provider: DbtRefTableProvider + + def __call__(self, schema): + if not issubclass(schema, ChecksumBuckets): + return None + + assert issubclass(self.expr, IncrementalExpression) + return construct_dbt_model( + schema @ AsIncrementedBuckets(self.expr), + self.table_provider, + self=self.expr, + ) + + +@frozen +class DbtSelfRefProvider(TableProvider): + self_expr: type[Expression] + + SELF_TOKEN = "__dbt_this__" + + def __call__(self, schema): + if not issubclass(schema, self.self_expr): + return None + + table = ibis.table(schema.table_schema, name=self.SELF_TOKEN) + return schema.of(table) + + +@frozen +class DbtRefTableProvider(TableProvider): + """Provide tables with DBT token names for compiling DBT SQL.""" + + model_lookup: Mapping[type[Expression], DbtModel] + source_lookup: Mapping[type[IbisSchema], DbtSource] + + REF_TOKEN = "__dbt_ref__" + SRC_TOKEN = "__dbt_src__" + SEP_TOKEN = "__dbt_sep__" + + def __call__(self, schema): + if model := self.model_lookup.get(schema): + schema = model.db_schema + name = schema.table_name + + ref_name = f"{self.REF_TOKEN}{name}{self.REF_TOKEN}" + table = ibis.table(schema.table_schema, name=ref_name) + return schema.of(table) + + if source := self.source_lookup.get(schema): + schema = source.db_schema + name = schema.table_name + src = source.dbt_source_name + + ref_name = f"{self.SRC_TOKEN}{name}{self.SEP_TOKEN}{src}{self.SRC_TOKEN}" + table = ibis.table(schema.table_schema, name=ref_name) + return schema.of(table) + + return None diff --git a/ibis_typing/dbt/dbt_model.py b/ibis_typing/dbt/dbt_model.py new file mode 100644 index 0000000..f97db70 --- /dev/null +++ b/ibis_typing/dbt/dbt_model.py @@ -0,0 +1,169 @@ +"""Integration of DBT Model via ibis Expression classes.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from enum import StrEnum +from typing import Any, Self + +import attrs +from attrs import frozen + +from ibis_typing import Expression, IbisDbSchema, it, naming + + +@frozen +class DbtModel[E: Expression]: + """DBT model based on an Expression class.""" + + expr: type[E] + config: ModelConfig + + name: str | None = None + prefix: str = "" + + @property + def table_name(self): + return self.name or (self.prefix + naming.snake_case(self.expr.__name__)) + + @property + def db_schema(self) -> type[IbisDbSchema]: + return type( + self.expr.__name__, + (IbisDbSchema, self.expr), + {"table_name": self.table_name, "table_namespace": self._namespace}, + ) + + @property + def _namespace(self) -> tuple[str, str]: + conf = self.config + assert conf.database + assert conf.schema + return conf.database, conf.schema + + +@frozen +class DbtSource[S: IbisDbSchema]: + """Mark-up for a DBT Source table. + + See + https://docs.getdbt.com/docs/build/sources?version=2.0&name=Fusion#declaring-a-source + """ + + db_schema: type[S] + + def _get_dbt_source_name(self) -> str: + # By default, schema will be the same as name. + return self.db_schema.table_namespace[1] + + dbt_source_name: str = attrs.field( + default=attrs.Factory(_get_dbt_source_name, takes_self=True) + ) + + +@frozen +class GeneralConfig: + """General configurations applicable across multiple DBT resource types. + + See + https://docs.getdbt.com/reference/model-configs?version=2.0&name=Fusion#general-configurations + """ + + database: str | None = None + schema: str | None = None + alias: str | None = None + + tags: Sequence[str] | None = None + + enabled: bool | None = None + + pre_hook: Sequence[str] | None = None + post_hook: Sequence[str] | None = None + + persist_docs: Mapping | None = None + full_refresh: bool | None = None + meta: Mapping | None = None + grants: Mapping | None = None + contract: Mapping | None = None + event_time: it.NameOrType | None = None + + kwargs: Any | None = None # Other fields and backend-specific config + + def with_namespace(self, database: str, schema: str) -> Self: + return attrs.evolve(self, database=database, schema=schema) + + def add_tags(self, tags: Sequence[str] | None) -> Self: + all_tags = (*(self.tags or ()), *(tags or ())) + return attrs.evolve(self, tags=list(all_tags) or None) + + +@frozen +class ModelConfig(GeneralConfig): + """Configuration for DBT models, incremental or not. + + See + https://docs.getdbt.com/reference/model-configs?version=2.0&name=Fusion#model-specific-configurations + """ + + materialized: Materialized | str | None = None + + incremental_strategy: IncrementalStrategy | None = None + unique_key: Sequence[it.NameOrType] | None = None + on_schema_change: OnSchemaChange | None = None + + def __attrs_post_init__(self): + if self.incremental_strategy: + assert self.unique_key is not None + + +class PlainStrEnum(StrEnum): + """StrEnum, but with simple str() representation.""" + + def __str__(self): + return repr(self._value_) + + def __repr__(self): + return repr(self._value_) + + +class Materialized(PlainStrEnum): + """Materialization strategy for DBT models. + + Note: Custom materialization strategies are handled as plain strings. + + See + https://docs.getdbt.com/docs/build/materializations?version=2.0 + """ + + table = "table" + view = "view" + incremental = "incremental" + ephemeral = "ephemeral" + materialized_view = "materialized view" + + +class OnSchemaChange(PlainStrEnum): + """Strategy for handling schema changes in incremental models. + + See + https://docs.getdbt.com/docs/build/incremental-models?version=2.0&name=Fusion#what-if-the-columns-of-my-incremental-model-change + """ + + ignore = "ignore" + fail = "fail" + append_new_columns = "append_new_columns" + sync_all_columns = "sync_all_columns" + + +class IncrementalStrategy(PlainStrEnum): + """Strategy for updating incremental models. + + See + https://docs.getdbt.com/docs/build/incremental-strategy?version=2.0&name=Fusion#supported-incremental-strategies-by-adapter + """ + + append = "append" + merge = "merge" + delete_insert = "delete+insert" + insert_overwrite = "insert_overwrite" + microbatch = "microbatch" diff --git a/ibis_typing/dbt/dbt_model_resolver.py b/ibis_typing/dbt/dbt_model_resolver.py new file mode 100644 index 0000000..78a9f61 --- /dev/null +++ b/ibis_typing/dbt/dbt_model_resolver.py @@ -0,0 +1,44 @@ +"""Automatic discovery of DBT Model candidates.""" + +from types import ModuleType + +from attrs import frozen + +from ibis_typing import ( + ChecksumBuckets, + Expression, + IbisSchema, + IncrementalExpression, + schema_writer, +) + +from .dbt_snapshot import DbtSnapshotAbstract + + +@frozen +class DbtModelResolver: + package: ModuleType + + @property + def expressions(self) -> set[type[Expression]]: + return set(schema_writer.list_expressions_in_package(self.package)) + + @property + def snapshots(self) -> set[type[DbtSnapshotAbstract]]: + return { + expr for expr in self.expressions if issubclass(expr, DbtSnapshotAbstract) + } + + @property + def incremental_models(self) -> set[type[IncrementalExpression]]: + return { + expr for expr in self.expressions if issubclass(expr, IncrementalExpression) + } + + @property + def checksums(self) -> set[type[ChecksumBuckets]]: + return {expr for expr in self.expressions if issubclass(expr, ChecksumBuckets)} + + @property + def checksum_inputs(self) -> set[type[IbisSchema]]: + return {checksum.incremental_params.inputs for checksum in self.checksums} diff --git a/ibis_typing/dbt/dbt_snapshot.py b/ibis_typing/dbt/dbt_snapshot.py new file mode 100644 index 0000000..9f57367 --- /dev/null +++ b/ibis_typing/dbt/dbt_snapshot.py @@ -0,0 +1,124 @@ +"""Integrates DBT Snapshot models via ibis Expression classes.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import ClassVar + +from attrs import frozen + +from ibis_typing import IbisSchema, IbisTable, it +from ibis_typing.expression import ( + GenericExpression, + SingleInputTableExpression, +) +from ibis_typing.table_provider import EmptyTableProvider + +from .dbt_model import DbtModel, GeneralConfig, PlainStrEnum + + +@frozen(kw_only=True) +class DbtSnapshot[E: IbisSchema](DbtModel): + """Integration mark-up for DBT snapshot function.""" + + expr: type[E] + config: SnapshotConfig + + +class DbtSnapshotAbstract(GenericExpression): + """Abstract implementation for schema generation.""" + + origin: ClassVar[type[IbisSchema]] + config: ClassVar[SnapshotConfig] + + @classmethod + def get_table_expression(cls): + return DbtSnapshotAbstractTableExpression(cls.origin, cls.config) + + +@frozen(kw_only=True) +class SnapshotConfig(GeneralConfig): + """DBT Snapshot configuration. Similar to DBT Model. + + See + https://docs.getdbt.com/reference/snapshot-configs?version=2.0&name=Fusion#snapshot-specific-configurations + https://docs.getdbt.com/docs/build/snapshots?version=2.0&name=Fusion#configuring-snapshots + """ + + unique_key: Sequence[it.NameOrType] + strategy: SnapshotStrategy + + updated_at: it.NameOrType | None = None + check_cols: Sequence[it.NameOrType] | None = None + + snapshot_meta_column_names: Mapping[it.NameOrType, it.NameOrType] | None = None + hard_deletes: HardDeletes | None = None + dbt_valid_to_current: str | None = None + + def __attrs_post_init__(self): + match self.strategy: + case SnapshotStrategy.check: + assert self.check_cols is not None + case SnapshotStrategy.timestamp: + assert self.updated_at is not None + + +class SnapshotStrategy(PlainStrEnum): + timestamp = "timestamp" + check = "check" + + +class HardDeletes(PlainStrEnum): + ignore = "ignore" + invalidate = "invalidate" + new_record = "new_record" + + +@frozen +class DbtSnapshotAbstractTableExpression(SingleInputTableExpression): + config: SnapshotConfig + + @property + def input_schemas(self): + return {"origin": self.origin} + + def __call__(self, origin: IbisTable): + defaults = EmptyTableProvider()(SnapshotsCols).table + delete_cols = ( + EmptyTableProvider()(SnapshotHardDeleteNewRecordCols).table + if self.config.hard_deletes == HardDeletes.new_record + else () + ) + col_names = self.config.snapshot_meta_column_names or {} + renames = {new: old for old, new in col_names.items()} + return ( + origin.table + @ it.InnerJoin(defaults, *delete_cols, keys=()) + @ it.deferred.rename(renames) + ) + + +@frozen +class SnapshotsCols(IbisSchema): + """DBT snapshot model markup. + + See: + https://docs.getdbt.com/reference/resource-configs/snapshot_meta_column_names?version=2.0&name=Fusion#default + """ + + dbt_scd_id: it.String = None + dbt_updated_at: it.Timestamp = None + dbt_valid_from: it.Timestamp = None + dbt_valid_to: it.Timestamp = None + + table_schema: ClassVar[Mapping[str, str]] = { + "dbt_scd_id": "string", + "dbt_updated_at": "timestamp(6)", + "dbt_valid_from": "timestamp(6)", + "dbt_valid_to": "timestamp(6)", + } + + +@frozen +class SnapshotHardDeleteNewRecordCols(IbisSchema): + dbt_is_deleted: it.Boolean = None diff --git a/ibis_typing/dbt/dbt_sql_compiler.py b/ibis_typing/dbt/dbt_sql_compiler.py new file mode 100644 index 0000000..327c5f1 --- /dev/null +++ b/ibis_typing/dbt/dbt_sql_compiler.py @@ -0,0 +1,171 @@ +"""Compile DBT Model to DBT SQL.""" + +from __future__ import annotations + +import functools +import re +from argparse import Namespace +from collections.abc import Mapping +from typing import Any + +import attrs + +from ibis_typing import ChecksumBuckets + +from ..revertible import AsRevertible, ExpressionExport +from . import dbt_ibis_constructor +from .dbt_ibis_constructor import DbtRefTableProvider, DbtSelfRefProvider +from .dbt_model import DbtModel, GeneralConfig, Materialized +from .dbt_snapshot import DbtSnapshot + + +def dbt_model_to_dbt_sql( + model: DbtModel, dialect: str, *, ref_provider: DbtRefTableProvider +) -> str: + if isinstance(model, DbtSnapshot): + return dbt_snapshot_to_dbt_sql(model, dialect, ref_provider=ref_provider) + + from_source = dbt_ibis_constructor.construct_dbt_model(model.expr, ref_provider) + + from_source_sql = from_source.table.to_sql(dialect) + + if model.config.materialized != Materialized.incremental: + sql = from_source_sql + else: + increment = ( + dbt_ibis_constructor.construct_dbt_bucket_increment( + model.expr, ref_provider + ) + if issubclass(model.expr, ChecksumBuckets) + else dbt_ibis_constructor.construct_dbt_model( + model.expr, ref_provider, buckets_update=True + ) + ) + + increment_sql = increment.table.to_sql(dialect) + + sql = as_incremental_macro(from_source_sql, inc_sql=increment_sql) + + dbt_sql = process_sql_dbt_tokens(sql) + config_sql = config_to_jinja(model.config) + + return config_sql + dbt_sql + + +def dbt_snapshot_to_dbt_sql( + snapshot: DbtSnapshot, dialect: str, *, ref_provider: DbtRefTableProvider +) -> str: + exported = snapshot.expr @ AsRevertible(ExpressionExport) + snapshot_table = dbt_ibis_constructor.construct_dbt_model(exported, ref_provider) + + table_sql = snapshot_table.table.to_sql(dialect) + + dbt_sql = process_sql_dbt_tokens(table_sql) + config_jinja = config_to_jinja(snapshot.config) + + return as_snapshot_macro(config_jinja + dbt_sql, table_name=snapshot.table_name) + + +def as_incremental_macro(full_sql: str, inc_sql: str) -> str: + """Wrap SQL in a DBT Jinja conditional for incremental or full refresh. + + >>> as_incremental_macro("SELECT 1", "SELECT 2") + '{% if is_incremental() %}\\nSELECT 2\\n{% else %}\\nSELECT 1\\n{% endif %}' + + See + https://docs.getdbt.com/docs/build/incremental-models?version=2.0&name=Fusion#filtering-rows-on-an-incremental-run + """ + return f"{{% if is_incremental() %}}\n{inc_sql}\n{{% else %}}\n{full_sql}\n{{% endif %}}" + + +def as_snapshot_macro(sql: str, *, table_name: str): + """Wrap SQL in a DBT Jinja conditional for incremental or full refresh. + + Note: This uses DBT legacy SQL configuration. + Consider using the all-YML configuration approach for snapshots. + + >>> as_snapshot_macro("SELECT 2", table_name="table") + '{% snapshot table %}\\nSELECT 2\\n{% endsnapshot %}' + + See + https://docs.getdbt.com/reference/resource-configs/snapshots-jinja-legacy?version=2.0&name=Fusion#snapshot-specific-configurations + https://docs.getdbt.com/reference/snapshot-configs?version=2.0#configuring-snapshots + """ + return f"{{% snapshot {table_name} %}}\n{sql}\n{{% endsnapshot %}}" + + +def process_sql_dbt_tokens(sql: str) -> str: + """Format DBT Jinja function calls for `this`, `ref`, and `source` from SQL tokens. + + >>> process_sql_dbt_tokens('SELECT * FROM "__dbt_this__"') + 'SELECT * FROM {{ this }}' + >>> process_sql_dbt_tokens('SELECT * FROM "__dbt_ref__my_model__dbt_ref__"') + 'SELECT * FROM {{ ref("my_model") }}' + >>> process_sql_dbt_tokens( + ... 'SELECT * FROM "__dbt_src__my_source__dbt_sep__my_table__dbt_src__"' + ... ) + 'SELECT * FROM {{ source("my_table", "my_source") }}' + + See + https://docs.getdbt.com/reference/dbt-jinja-functions/this?version=2.0&name=Fusion + https://docs.getdbt.com/reference/dbt-jinja-functions/ref?version=2.0&name=Fusion#definition + https://docs.getdbt.com/reference/dbt-jinja-functions/source?version=2.0&name=Fusion#example + """ + self = DbtSelfRefProvider + cls = DbtRefTableProvider + subs = [ + ( + re.compile(rf'"?{self.SELF_TOKEN}"?'), + "{{ this }}", + ), + ( + re.compile(rf'"?{cls.REF_TOKEN}(.*?){cls.REF_TOKEN}"?'), + r'{{ ref("\1") }}', + ), + ( + re.compile(rf'"?{cls.SRC_TOKEN}(.*?){cls.SEP_TOKEN}(.*?){cls.SRC_TOKEN}"?'), + r'{{ source("\2", "\1") }}', + ), + ] + + def reduce(sub, s): + pattern, new = sub + return pattern.sub(new, s) + + return functools.reduce(lambda s, sub: reduce(sub, s), subs, sql) + + +def config_to_jinja(config: GeneralConfig) -> str: + """Format a DBT Jinja config header. + + >>> from ibis_typing.dbt.dbt_model import IncrementalStrategy, ModelConfig + >>> config_to_jinja( + ... ModelConfig( + ... incremental_strategy=IncrementalStrategy.merge, + ... unique_key=[], + ... kwargs={"properties": {}}, + ... ) + ... ) + "{{ config(incremental_strategy='merge', unique_key=[], properties={}) }}\\n" + """ + kwargs = { + **{ + key: val + for key, val in attrs.asdict(config, recurse=False).items() + if val is not None + if key != "kwargs" + }, + **(config.kwargs or {}), + } + namespace = Namespace(**_tuple_as_list(kwargs)) + py_block = str(namespace).replace(Namespace.__name__, "config", 1) + return "{{ " + py_block + " }}\n" + + +def _tuple_as_list(obj: Any) -> Any: + """DBT internals requires strict `list` type arguments.""" + if isinstance(obj, tuple): + return list(obj) + if isinstance(obj, Mapping): + return {key: _tuple_as_list(val) for key, val in obj.items()} + return obj diff --git a/ibis_typing/ibis_extension_method.py b/ibis_typing/ibis_extension_method.py index dd91af7..642ba62 100644 --- a/ibis_typing/ibis_extension_method.py +++ b/ibis_typing/ibis_extension_method.py @@ -14,7 +14,7 @@ class TableMethod(ExtensionMethod[Table, Table], abc.ABC): - """Apply operation to Table on left-hand side of this operator.""" + """Table @ TableMethod() -> Table.""" @abc.abstractmethod def apply(self, table: Table) -> Table: ... @@ -32,6 +32,8 @@ def as_expression_schema( @frozen class TableMethodExpression(SingleInputTableExpression): + """Upgrades a TableMethod to a GenericExpression schema.""" + method: TableMethod preserves_schema: bool = False @@ -46,7 +48,7 @@ def __call__(self, origin: IbisTable) -> Table: class ValueMethod[T: Value, R: Value](ExtensionMethod[T, R], abc.ABC): - """Apply operation to Value on left-hand side of this operator.""" + """Value @ ValueMethod() -> Value.""" @abc.abstractmethod def apply(self, value: T) -> R: ... @@ -56,21 +58,23 @@ def __rmatmul__(self, other: T) -> R: class SelfMethod[T: Value](ValueMethod[T, T], abc.ABC): - pass + """T @ SelfMethod() -> T.""" @frozen(init=False) -class ArgsMethod[T: Value](SelfMethod[T], abc.ABC): - args: Sequence[T] +class ArgsMethod[V: Value](SelfMethod[V], abc.ABC): + """V @ Args(*values) -> V.""" + + args: Sequence[V] - def __init__(self, *args: T): + def __init__(self, *args: V): self.__attrs_init__(args) # type: ignore class ExpressionMethod[S: IbisSchema, E: GenericExpression]( ExtensionMethod[type[S], type[E]], abc.ABC ): - """Extension method for chaining Expression class transforms.""" + """type[IbisSchema] @ ExpressionMethod() -> type[GenericExpression].""" @abc.abstractmethod def apply(self, schema: type[S]) -> TableExpression: ... diff --git a/ibis_typing/revertible.py b/ibis_typing/revertible.py index 3bf9ac2..ace4144 100644 --- a/ibis_typing/revertible.py +++ b/ibis_typing/revertible.py @@ -4,10 +4,12 @@ import ibis from attrs import frozen +from ibis.expr import datatypes as dt from . import ibis_types as it from .expression import GenericExpression, SingleInputTableExpression from .ibis_adapter import IbisSchema, IbisTable +from .ibis_extension_method import ExpressionMethod _ = it # Needed for doctests @@ -80,3 +82,54 @@ def revert_all(table: IbisTable) -> IbisTable: table = constructor.revert_expression(table) return table + + +@frozen +class AsRevertible[E: ExpressionExport](ExpressionMethod): + """Wrap any non-parametric ExpressionExport as an ExpressionMethod. + + For any parametric ExpressionExport, implement a parametric ExpressionMethod instead.""" + + revertible: type[E] + + def apply(self, schema): + return self.revertible(schema, self) + + +@frozen +class ExpressionExport[R: ExpressionMethod](RevertibleTableExpression): + """Basic revertible transform, composable with `Schema @ AsRevertible(ExportExpression)`. + + All parameters except for origin schema is kept in the ExpressionMethod. + The ExpressionMethod is then passed to the ExpressionExport. + """ + + method: R + + def __call__(self, origin): + return origin.table + + def revert_call(self, transformed): + return transformed.table + + @property + def generated_class_name(self) -> str: + return self.origin.__name__ + + +class UUIDToStringExport(ExpressionExport): + """Cast UUID to strings for e.g. Apache Hive compatibility. + + See + https://hive.apache.org/docs/latest/language/languagemanual-types/ + """ + + def __call__(self, origin: IbisTable): + return origin.table.cast(dict.fromkeys(self._get_uuid_casts(), dt.string)) + + def revert_call(self, transformed: IbisTable): + return transformed.table.cast(dict.fromkeys(self._get_uuid_casts(), dt.uuid)) + + def _get_uuid_casts(self): + schema = ibis.schema(self.origin.table_schema) + return [column for column, type_ in schema.items() if type_ == dt.uuid] diff --git a/ibis_typing/samples/dbt_models.py b/ibis_typing/samples/dbt_models.py new file mode 100644 index 0000000..f52fda7 --- /dev/null +++ b/ibis_typing/samples/dbt_models.py @@ -0,0 +1,69 @@ +from collections.abc import Mapping + +from ibis_typing import Expression, IbisDbSchema, IbisSchema, naming, samples +from ibis_typing.dbt import ( + DbtModel, + DbtSnapshot, + DbtSource, + ModelConfig, + SnapshotConfig, + dbt_model_resolver, +) +from ibis_typing.dbt.dbt_model import IncrementalStrategy, Materialized +from ibis_typing.dbt.dbt_snapshot import SnapshotStrategy +from ibis_typing.ibis_time import TimestampNow +from ibis_typing.samples.sample_incremental_calendar import Calendar, CalendarWidth +from ibis_typing.samples.sample_transforms import CircleParameters + +my_namespace = ("my_database", "my_schema") + + +def get_dbt_model_lookup() -> Mapping[type[Expression], DbtModel]: + resolver = dbt_model_resolver.DbtModelResolver(samples) + + models = { + # Provide timestamp as a table in order to use a fix timestamp per DBT run. + DbtModel( + TimestampNow, + config=ModelConfig( + materialized=Materialized.table, + ).with_namespace(*my_namespace), + ) + } + incremental_models = [ + DbtModel( + expr, + config=ModelConfig( + materialized=Materialized.incremental, + incremental_strategy=IncrementalStrategy.merge, + unique_key=expr.incremental_params.group_by, + ).with_namespace(*my_namespace), + ) + for expr in resolver.incremental_models + ] + snapshots = [ + DbtSnapshot( + expr=CalendarWidth, + config=SnapshotConfig( + strategy=SnapshotStrategy.timestamp, + unique_key=CalendarWidth.incremental_params.group_by, + updated_at=CalendarWidth.incremental_params.updated_at_col, + ).with_namespace(*my_namespace), + ) + ] + all_models = (*models, *incremental_models, *snapshots) + + return {model.expr: model for model in all_models} + + +def get_dbt_source_lookup() -> Mapping[type[IbisSchema], DbtSource]: + sources = [Calendar, CircleParameters] + return {schema: DbtSource(as_db_schema(schema)) for schema in sources} + + +def as_db_schema(schema: type[IbisSchema]) -> type[IbisDbSchema]: + kwargs = { + "table_name": naming.snake_case(schema.__name__), + "table_namespace": my_namespace, + } + return type(schema.__name__, (IbisDbSchema, schema), kwargs) diff --git a/ibis_typing/samples/generated/__ibis_sql/calendar_checksum_bucket.sql b/ibis_typing/samples/generated/__ibis_sql/calendar_checksum_bucket.sql new file mode 100644 index 0000000..17f73da --- /dev/null +++ b/ibis_typing/samples/generated/__ibis_sql/calendar_checksum_bucket.sql @@ -0,0 +1,69 @@ +{{ config(database='my_database', schema='my_schema', materialized='incremental', incremental_strategy='merge', unique_key=[]) }} +{% if is_incremental() %} +SELECT + "t13"."timestamp" AS "checksum_updated_at", + "t13"."checksum" +FROM ( + SELECT + "t12"."checksum_updated_at", + "t12"."checksum", + "t12"."checksum_updated_at_right", + "t12"."checksum_right", + "t3"."timestamp" + FROM ( + SELECT + "t10"."checksum_updated_at", + COALESCE("t10"."checksum", 0) AS "checksum", + "t10"."checksum_updated_at_right", + "t10"."checksum_right" + FROM ( + SELECT + "t9"."checksum_updated_at", + "t9"."checksum", + "t4"."checksum_updated_at" AS "checksum_updated_at_right", + "t4"."checksum" AS "checksum_right" + FROM ( + SELECT + "t7"."timestamp" AS "checksum_updated_at", + "t7"."checksum" + FROM ( + SELECT + "t6"."checksum", + "t3"."timestamp" + FROM ( + SELECT + CAST(BIT_XOR(HASH("t2"."day")) % 9223372036854775808 AS BIGINT) AS "checksum" + FROM {{ source("my_schema", "calendar") }} AS "t2" + ) AS "t6" + INNER JOIN {{ ref("timestamp_now") }} AS "t3" + ON TRUE + ) AS "t7" + ) AS "t9" + FULL OUTER JOIN {{ this }} AS "t4" + ON TRUE + ) AS "t10" + WHERE + NOT ( + COALESCE("t10"."checksum", 0) IS NOT DISTINCT FROM "t10"."checksum_right" + ) + ) AS "t12" + INNER JOIN {{ ref("timestamp_now") }} AS "t3" + ON TRUE +) AS "t13" +{% else %} +SELECT + "t5"."timestamp" AS "checksum_updated_at", + "t5"."checksum" +FROM ( + SELECT + "t4"."checksum", + "t2"."timestamp" + FROM ( + SELECT + CAST(BIT_XOR(HASH("t0"."day")) % 9223372036854775808 AS BIGINT) AS "checksum" + FROM {{ source("my_schema", "calendar") }} AS "t0" + ) AS "t4" + INNER JOIN {{ ref("timestamp_now") }} AS "t2" + ON TRUE +) AS "t5" +{% endif %} \ No newline at end of file diff --git a/ibis_typing/samples/generated/__ibis_sql/calendar_width.sql b/ibis_typing/samples/generated/__ibis_sql/calendar_width.sql new file mode 100644 index 0000000..7bc937f --- /dev/null +++ b/ibis_typing/samples/generated/__ibis_sql/calendar_width.sql @@ -0,0 +1,6 @@ +{% snapshot calendar_width %} +{{ config(database='my_database', schema='my_schema', unique_key=[], strategy='timestamp', updated_at='checksum_updated_at') }} +SELECT + * +FROM {{ ref("calendar_width") }} +{% endsnapshot %} \ No newline at end of file diff --git a/ibis_typing/samples/generated/__ibis_sql/timestamp_now.sql b/ibis_typing/samples/generated/__ibis_sql/timestamp_now.sql new file mode 100644 index 0000000..9b74c73 --- /dev/null +++ b/ibis_typing/samples/generated/__ibis_sql/timestamp_now.sql @@ -0,0 +1,13 @@ +{{ config(database='my_database', schema='my_schema', materialized='table') }} +SELECT + CAST(CAST(CURRENT_TIMESTAMP AS TIMESTAMP) AS TIMESTAMP(6)) AS "timestamp" +FROM ( + SELECT + CAST("t0"."timestamp" AS TIMESTAMP(6)) AS "timestamp" + FROM ( + SELECT + * + FROM (VALUES + (NULL)) AS TimestampNow__62f0537e(timestamp) + ) AS "t0" +) AS "t1" \ No newline at end of file diff --git a/ibis_typing/samples/generated/__init__.py b/ibis_typing/samples/generated/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ibis_typing/samples/tests/test_generate_ibis_dbt_sql.py b/ibis_typing/samples/tests/test_generate_ibis_dbt_sql.py new file mode 100644 index 0000000..873dca6 --- /dev/null +++ b/ibis_typing/samples/tests/test_generate_ibis_dbt_sql.py @@ -0,0 +1,35 @@ +from pathlib import Path + +from ibis_typing.dbt import dbt_sql_compiler +from ibis_typing.dbt.dbt_ibis_constructor import DbtRefTableProvider +from ibis_typing.samples import dbt_models + +from .. import generated + +SQL_DIR_NAME = "__ibis_sql" +SQL_DIR = Path(generated.__file__).parent / SQL_DIR_NAME + + +def test_generate_ibis_dbt_sql(update_expected): + model_lookup = dbt_models.get_dbt_model_lookup() + source_lookup = dbt_models.get_dbt_source_lookup() + ref_provider = DbtRefTableProvider(model_lookup, source_lookup) + + new = { + SQL_DIR / (model.table_name + ".sql"): dbt_sql_compiler.dbt_model_to_dbt_sql( + model, + dialect="duckdb", + ref_provider=ref_provider, + ) + for model in model_lookup.values() + } + + if update_expected: + for path, sql in new.items(): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(sql) + + prior_paths = {p for p in SQL_DIR.iterdir() if p.suffix == ".sql"} + prior = {path: path.read_text() for path in prior_paths} + + assert new == prior