diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/README.md b/examples/learn/08-ingestion-and-transformation-with-dagster/README.md index 08edeca4..ef28c9ef 100644 --- a/examples/learn/08-ingestion-and-transformation-with-dagster/README.md +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/README.md @@ -5,6 +5,9 @@ - [Goal](#goal) - [Set up](#set-up) - [Usage](#usage) +- [Data quality checks in the Dagster UI](#data-quality-checks-in-the-dagster-ui) +- [Materialization metadata](#materialization-metadata) +- [When things fail](#when-things-fail) ## Goal @@ -21,9 +24,11 @@ The core idea is that Dagster triggers the import of a single partition into Bau Once both tables are in the lakehouse, a transformation pipeline builds a _daily_ summary. It keeps only settled transactions, aggregates spend per account (total, count, and average amount), and joins the result with per-account event and login counts derived from `account_events`. +On the Dagster side, everything is modeled with software-defined assets: the two source tables and the summary table appear in the asset lineage view, the Bauplan expectations surface as Dagster asset checks, and each materialization carries Bauplan metadata (row counts, column schema, branch and job identifiers) into the Dagster asset catalog. + ## Set up -Use `src/.env.example` as a reference for the fields required to run this example. You can generate a new Bauplan API key (if you don't already have one) from [https://app.bauplanlabs.com/api-keys](https://app.bauplanlabs.com/api-keys). A bucket and prefix have already been provided. When you launch the Dagster UI, these variables are loaded automatically from `src/.env`. +Use `src/.env.example` as a reference for the fields required to run this example. You can generate a new Bauplan API key (if you don't already have one) from [https://app.bauplanlabs.com/api-keys](https://app.bauplanlabs.com/api-keys). A bucket and prefix have already been provided; `NAMESPACE` and `BASE_BRANCH` default to `workshop` and `workshop.main` and only need changing if you want to publish elsewhere. When you launch the Dagster UI, these variables are loaded automatically from `src/.env`. Ensure [`uv`](https://docs.astral.sh/uv/) is installed by following the [official documentation](https://docs.astral.sh/uv/getting-started/installation/), then install the dependencies: @@ -33,75 +38,43 @@ uv sync ## Usage -For a fast turnaround you can use the CLI to launch jobs without spinning up the Dagster UI: +For a fast turnaround you can use the CLI to materialize partitions without spinning up the Dagster UI: ```sh -set -a && source src/.env && set +a && uv run src/main.py --help +set -a && source src/.env && set +a +uv run src/main.py ingest transactions 2026-06-16 +uv run src/main.py ingest account_events 2026-06-16 +uv run src/main.py transform 2026-06-16 ``` -The CLI exposes two commands, `ingest` and `transform`; pass `--help` to either to see the available parameters. +`ingest` materializes one daily partition of one source table, including its checks; `transform` materializes one daily partition of the summary. Note that each CLI invocation runs on an ephemeral Dagster instance, so its run history will not appear in a later UI session. -For the full experience, open the Dagster UI on your localhost. Run: +For the full experience, open the Dagster UI on your localhost: ```sh cd src && uv run dg dev -m main ``` -Open `Jobs > ingestion > Launchpad` and set up a config such as - -```yaml -ops: - import_data: - config: - base_branch: 'workshop.main' - day: '16' - month: '06' - namespace: 'workshop' - table: 'account_events' - year: '2026' - dt_partition_column: 'event_ts' -resources: - bauplan_client: - config: - api_key: - env: 'BAUPLAN_API_KEY' -``` +Open the `Assets` page (or the lineage tab) to see the graph `transactions`, `account_events` into `account_activity_summary` with the daily partition bar under each asset. Select an asset, click `Materialize`, and pick a partition in the dialog. The `ingestion` and `transformation` jobs group the same assets if you prefer launching from the `Jobs` page. -then run the job. Each ingestion run imports a single table, so you need two runs to have all the data a transformation needs. Pair the config above with - -```yaml -ops: - import_data: - config: - base_branch: 'workshop.main' - day: '16' - month: '06' - namespace: 'workshop' - table: 'transactions' - year: '2026' - dt_partition_column: 'txn_ts' -resources: - bauplan_client: - config: - api_key: - env: 'BAUPLAN_API_KEY' -``` +To load a date range, use a backfill from the partition dialog. Keep the backfill concurrency at 1: parallel partition runs of the same table would race on the merge into the base branch. -to load all the data. Once both tables are loaded, open `Jobs > transformation > Launchpad` and set - -```yaml -ops: - transform: - config: - base_branch: 'workshop.main' - namespace: 'workshop' - start_date: '2026-06-16' - end_date: '2026-06-18' -resources: - bauplan_client: - config: - api_key: - env: 'BAUPLAN_API_KEY' -``` +## Data quality checks in the Dagster UI + +The example demonstrates the two ways of coupling Bauplan data quality with Dagster asset checks. + +The `audit_expectations` check on each source table is the audit step of the WAP cycle. The asset body runs the table's expectations project on the ingestion branch and reports the outcome as an in-asset check result, so the audit that gates the merge is the same event you see in the `Checks` tab. Its metadata carries the Bauplan job id, the audit duration, and, on failure, the error and the name of the ingestion branch that was kept for debugging. The expectations themselves stay in the Bauplan project (`src/ingestion/pipelines/audit//expectations.py`). + +The post-publish checks (`transactions_txn_id_unique`, `transactions_audit`, `account_events_account_id_no_nulls`, `account_events_audit`) are standalone `@asset_check` definitions. They run automatically on the base branch after each materialization of their asset. This is the pattern to follow when you want additional checks defined and versioned on the orchestrator side. Notice that these checks can be manually re-run, whereas the expectations embedded in the ingestion job cannot. This is why we have included `transactions_audit` and `account_events_audit` which replicate the expectations verbatim. + +One evaluation per check is shown in the `Checks` tab, the latest one; the per-partition history remains available in the run logs of each materialization. + +## Materialization metadata + +Every successful materialization attaches Bauplan metadata to the Dagster event: total row count (`dagster/row_count`), rows in the materialized partition (`dagster/partition_row_count`), the column schema (`dagster/column_schema`), the imported files (`dagster/uri`, ingestion assets only), plus the Bauplan branch and job identifiers and, for the transform, the run duration. + +## When things fail + +If an audit fails, the run stops before the merge: the `audit_expectations` check turns red, the asset is not marked materialized, and the ingestion branch is left in place so you can inspect the rejected data. The branch name is in the check metadata and in the run failure; delete the branch once you are done debugging. -then run the job. The `transformation` job forks a fresh branch from the base branch and runs the `account_summary` pipeline on it. The two upstream models, `settled_transactions` and `daily_account_spend`, are computed on the fly and never persisted; only the final model is materialized, as the `account_activity_summary` table, using a `OVERWRITE_PARTITIONS` strategy that overwrites any previous contents for a given partition. When the run succeeds, the branch is merged back into the base branch, so `account_activity_summary` is published to `workshop.main` next to the source tables. +If the import itself fails, the run errors before any audit, so the check shows no evaluation for that run; the failing branch name is part of the exception message. If the merge fails after a clean audit, the check stays green (the audit did pass) while the run fails, and the branch is again kept for inspection. diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/.env.example b/examples/learn/08-ingestion-and-transformation-with-dagster/src/.env.example index 4091a61a..811f9614 100644 --- a/examples/learn/08-ingestion-and-transformation-with-dagster/src/.env.example +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/.env.example @@ -1,3 +1,5 @@ BAUPLAN_API_KEY=bpln-sdk-abc123 BUCKET_NAME="alpha-hello-bauplan" PREFIX="workshop/01-ingestion" +NAMESPACE="workshop" +BASE_BRANCH="workshop.main" diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/import_data.py b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/import_data.py index 37ce5f30..d8087dc5 100644 --- a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/import_data.py +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/import_data.py @@ -67,6 +67,13 @@ def import_data( branch=ingestion_branch, detach=True, ) - wait_for_job(bpln_client, state.job_id, f"import of {table}") + wait_for_job( + bpln_client, state.job_id, f"import of {table} into {ingestion_branch}" + ) - return {"branch": ingestion_branch, "namespace": namespace} + return { + "branch": ingestion_branch, + "namespace": namespace, + "uri": uri, + "job_id": state.job_id, + } diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/account_events/bauplan_project.yaml b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/account_events/bauplan_project.yaml similarity index 64% rename from examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/account_events/bauplan_project.yaml rename to examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/account_events/bauplan_project.yaml index 2f137e60..5087b153 100644 --- a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/account_events/bauplan_project.yaml +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/account_events/bauplan_project.yaml @@ -1,3 +1,3 @@ project: id: 77650529-df04-49cc-b15a-9dbf4fb0085d - name: wap_account_events + name: audit_account_events diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/account_events/expectations.py b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/account_events/expectations.py similarity index 100% rename from examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/account_events/expectations.py rename to examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/account_events/expectations.py diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/account_events/pyproject.toml b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/account_events/pyproject.toml new file mode 100644 index 00000000..d6fd1254 --- /dev/null +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/account_events/pyproject.toml @@ -0,0 +1,5 @@ +[project] +name = "audit_account_events" +version = "0.1.0" +requires-python = "~=3.13" +dependencies = ["bauplan~=0.1.16"] diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/transactions/bauplan_project.yaml b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/transactions/bauplan_project.yaml similarity index 65% rename from examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/transactions/bauplan_project.yaml rename to examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/transactions/bauplan_project.yaml index 44f00868..c8fc3ffd 100644 --- a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/transactions/bauplan_project.yaml +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/transactions/bauplan_project.yaml @@ -1,3 +1,3 @@ project: id: c3092fef-0359-441d-b05b-86d3f941474b - name: wap_transactions + name: audit_transactions diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/transactions/expectations.py b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/transactions/expectations.py similarity index 100% rename from examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/transactions/expectations.py rename to examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/transactions/expectations.py diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/account_events/pyproject.toml b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/transactions/pyproject.toml similarity index 76% rename from examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/account_events/pyproject.toml rename to examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/transactions/pyproject.toml index 3ee5fc43..648988c8 100644 --- a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/account_events/pyproject.toml +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/audit/transactions/pyproject.toml @@ -1,5 +1,5 @@ [project] -name = "wap_account_events" +name = "audit_transactions" version = "0.1.0" requires-python = "~=3.13" dependencies = ["bauplan~=0.1.16"] diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/main.py b/examples/learn/08-ingestion-and-transformation-with-dagster/src/main.py index 38237b2a..e9b2049e 100644 --- a/examples/learn/08-ingestion-and-transformation-with-dagster/src/main.py +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/main.py @@ -1,3 +1,5 @@ +import os +from collections.abc import Iterator from datetime import datetime from enum import Enum from functools import cached_property @@ -8,9 +10,10 @@ import bauplan import dagster as dg import typer +from bauplan.exceptions import BauplanError from ingestion.import_data import import_data -from utils import wait_for_job +from utils import get_table_metadata app = typer.Typer() @@ -20,6 +23,24 @@ class SourceTable(str, Enum): account_events = "account_events" +# Lakehouse coordinates, overridable via the environment to point the example +# at a different namespace or base branch +NAMESPACE = os.environ.get("NAMESPACE", "workshop") +BASE_BRANCH = os.environ.get("BASE_BRANCH", "workshop.main") + +AUDIT_DIR = Path(__file__).parent / "ingestion" / "pipelines" / "audit" +TRANSFORM_DIR = ( + Path(__file__).parent / "transformation" / "pipelines" / "account_summary" +) + +# One Dagster partition per hive-style day prefix on S3; the sample data covers +# 2026-06-15 through 2026-07-14, so end_date (exclusive) is the day after the +# last populated prefix. Regenerating the data means moving these dates too +daily_partitions = dg.DailyPartitionsDefinition( + start_date="2026-06-15", end_date="2026-07-15" +) + + class BauplanResource(dg.ConfigurableResource): """Allocate a Dagster resource for the Bauplan client to be reused across invocations of Bauplan commands. Requires a valid @@ -41,138 +62,302 @@ def username(self) -> str: return user.username -@dg.op( - name="import_data", - config_schema={ - "table": str, - "year": str, - "month": str, - "day": str, - "dt_partition_column": str, - "base_branch": str, - "namespace": str, - }, -) -def run_import_data( - context: dg.OpExecutionContext, bauplan_client: BauplanResource -) -> dict: - """DAG node responsible for importing new parquet files into - an ingestion branch. This branch is separate from the base branch as quality - of new data needs to be ascertained.""" - cfg = context.op_config - result = import_data( - bpln_client=bauplan_client.client, - username=bauplan_client.username, - table=cfg["table"], - year=cfg["year"], - month=cfg["month"], - day=cfg["day"], - dt_partition_column=cfg["dt_partition_column"], - base_branch=cfg["base_branch"], - namespace=cfg["namespace"], +def build_ingestion_asset(table: str, dt_partition_column: str) -> dg.AssetsDefinition: + """Build the daily-partitioned asset that ingests one source table through a + write-audit-publish cycle: import the partition on an ingestion branch, audit + it with the table's Bauplan expectations, and merge into the base branch only + on a clean audit. The audit outcome surfaces as the audit_expectations check.""" + + @dg.asset( + name=table, + partitions_def=daily_partitions, + kinds={"bauplan"}, + check_specs=[ + dg.AssetCheckSpec( + name="audit_expectations", + asset=table, + description=f"Bauplan expectations for {table}, audited on the ingestion branch before publishing", + ) + ], + description=f"Daily partition of {table} imported from S3 with write-audit-publish", ) - # Carry results over to next step - result["base_branch"] = cfg["base_branch"] - result["table"] = cfg["table"] - return result - - -@dg.op(name="wap") -def run_wap(imported: dict, bauplan_client: BauplanResource) -> dict: - """DAG node responsible for running a WRITE-AUDIT-PUBLISH step. - It runs expectations on imported data to ascertain that it meets - requirements. Early exit in case of corrupted data.""" + def ingest( + context: dg.AssetExecutionContext, bauplan_client: BauplanResource + ) -> Iterator[dg.AssetCheckResult | dg.MaterializeResult]: + """Import one daily partition, audit it, publish it on success.""" + client = bauplan_client.client + day = datetime.strptime(context.partition_key, "%Y-%m-%d") + + # Write: import data in a dedicated branch + imported = import_data( + bpln_client=client, + username=bauplan_client.username, + table=table, + year=f"{day:%Y}", + month=f"{day:%m}", + day=f"{day:%d}", + dt_partition_column=dt_partition_column, + base_branch=BASE_BRANCH, + namespace=NAMESPACE, + ) + branch = imported["branch"] + + # Audit: a synchronous run so the returned state carries status and duration + state = client.run( + project_dir=str((AUDIT_DIR / table).resolve()), + ref=branch, + namespace=NAMESPACE, + cache="off", + strict="on", + ) + passed = str(state.job_status).lower() == "success" + check_metadata = { + "bauplan_job_id": str(state.job_id), + "bauplan_branch": branch, + "duration_seconds": round(state.duration or 0.0, 2), + "error": str(state.error or ""), + } + yield dg.AssetCheckResult( + check_name="audit_expectations", passed=passed, metadata=check_metadata + ) + if not passed: + raise dg.Failure( + description=f"Audit failed on branch {branch} (branch kept for debugging)", + metadata=check_metadata, + ) + + # Publish: merge only after a clean audit, then drop the ingestion branch + try: + client.merge_branch(source_ref=branch, into_branch=BASE_BRANCH) + except BauplanError as merge_error: + raise dg.Failure( + description=f"Merge of branch {branch} into {BASE_BRANCH} failed; branch kept for debugging", + metadata={"bauplan_branch": branch}, + ) from merge_error + + try: + client.delete_branch(branch=branch) + except BauplanError as delete_error: + raise dg.Failure( + description=f"Deletion of branch {branch} into {BASE_BRANCH} failed", + metadata={"bauplan_branch": branch}, + ) from delete_error + + window = context.partition_time_window + yield dg.MaterializeResult( + metadata=get_table_metadata( + client=client, + table=table, + namespace=NAMESPACE, + ref=BASE_BRANCH, + partition_predicate=( + f"{dt_partition_column} >= TIMESTAMP '{window.start:%Y-%m-%d} 00:00:00' " + f"AND {dt_partition_column} < TIMESTAMP '{window.end:%Y-%m-%d} 00:00:00'" + ), + extra={ + "bauplan_branch": branch, + "bauplan_import_job_id": str(imported["job_id"]), + "dagster/uri": imported["uri"], + }, + ) + ) + + return ingest + + +transactions = build_ingestion_asset("transactions", "txn_ts") +account_events = build_ingestion_asset("account_events", "event_ts") + + +@dg.asset( + name="account_activity_summary", + partitions_def=daily_partitions, + deps=[transactions, account_events], + kinds={"bauplan"}, + description="Daily per-account summary of settled spend joined with event and login counts", +) +def account_activity_summary( + context: dg.AssetExecutionContext, bauplan_client: BauplanResource +) -> dg.MaterializeResult: + """Run the account_summary pipeline for one daily partition on a dev branch + forked from the base branch, then publish the result by merging it back.""" client = bauplan_client.client - project_dir = ( - Path(__file__).parent / "ingestion" / "pipelines" / "wap" / imported["table"] - ).resolve() + window = context.partition_time_window + start_date = f"{window.start:%Y-%m-%d}" + end_date = f"{window.end:%Y-%m-%d}" + + branch = f"{bauplan_client.username}.workshop-transform-{context.partition_key}-{str(uuid4())[:8]}" + client.create_branch(branch=branch, from_ref=BASE_BRANCH) - # Detach and poll state = client.run( - project_dir=str(project_dir), - ref=imported["branch"], - namespace=imported["namespace"], + project_dir=str(TRANSFORM_DIR.resolve()), + ref=branch, + namespace=NAMESPACE, cache="off", strict="on", - detach=True, + parameters={"start_date": start_date, "end_date": end_date}, + ) + if str(state.job_status).lower() != "success": + raise dg.Failure( + description=f"transform on {branch} failed; branch kept for debugging", + metadata={ + "bauplan_branch": branch, + "bauplan_job_id": str(state.job_id), + "error": str(state.error or ""), + }, + ) + + try: + client.merge_branch(source_ref=branch, into_branch=BASE_BRANCH) + client.delete_branch(branch=branch) + except BauplanError as merge_error: + raise dg.Failure( + description=f"merge of {branch} into {BASE_BRANCH} failed; branch kept for debugging", + metadata={"bauplan_branch": branch}, + ) from merge_error + + return dg.MaterializeResult( + metadata=get_table_metadata( + client=client, + table="account_activity_summary", + namespace=NAMESPACE, + ref=BASE_BRANCH, + partition_predicate=( + f"date >= DATE '{start_date}' AND date < DATE '{end_date}'" + ), + extra={ + "bauplan_branch": branch, + "bauplan_job_id": str(state.job_id), + "bauplan_run_duration_seconds": round(state.duration or 0.0, 2), + }, + ) ) - wait_for_job(client, state.job_id, f"WAP audit on {imported['branch']}") - - # Pass-through so merge depends on wap and runs only after a clean audit - return imported -@dg.op(name="merge") -def run_merge(audited: dict, bauplan_client: BauplanResource) -> None: - """DAG node responsible for merging a branch into the main branch, - thus making new data available to downstream consumers. Deletes the import - branch afterwards.""" - bauplan_client.client.merge_branch( - source_ref=audited["branch"], - into_branch=audited["base_branch"], +@dg.asset_check( + asset=transactions, + description="Auditing pipeline for 'transactions' table. Replicates the one run at ingestion time.", +) +def transactions_audit(bauplan_client: BauplanResource) -> dg.AssetCheckResult: + """Post-publish check: run expectations.""" + state = bauplan_client.client.run( + project_dir=str((AUDIT_DIR / "transactions").resolve()), + ref=BASE_BRANCH, + namespace=NAMESPACE, + dry_run=True, + cache="off", + strict="on", ) - bauplan_client.client.delete_branch(branch=audited["branch"]) + successful = str(state.job_status).lower() == "success" + + return dg.AssetCheckResult( + passed=successful, + metadata={"error": str(state.error or ""), "ref": BASE_BRANCH}, + ) -@dg.op( - name="transform", - config_schema={ - "start_date": str, - "end_date": str, - "base_branch": str, - "namespace": str, - }, +@dg.asset_check( + asset=transactions, + description="No duplicate txn_id may survive on the published branch", ) -def run_transform( - context: dg.OpExecutionContext, bauplan_client: BauplanResource -) -> dict: - """DAG node responsible for running the transformation pipeline on a dev - branch forked from the base branch, materializing the derived tables.""" - cfg = context.op_config - client = bauplan_client.client +def transactions_txn_id_unique(bauplan_client: BauplanResource) -> dg.AssetCheckResult: + """Post-publish check: count duplicate transaction ids on the base branch.""" + duplicates = ( + bauplan_client.client.query( + "SELECT COUNT(*) - COUNT(DISTINCT txn_id) AS n FROM transactions", + ref=BASE_BRANCH, + namespace=NAMESPACE, + ) + .column("n") + .to_pylist()[0] + ) + return dg.AssetCheckResult( + passed=duplicates == 0, + metadata={"duplicate_txn_ids": int(duplicates or 0), "ref": BASE_BRANCH}, + ) - branch = f"{bauplan_client.username}.workshop-transform-{str(uuid4())[:8]}" - client.create_branch(branch=branch, from_ref=cfg["base_branch"]) - project_dir = ( - Path(__file__).parent / "transformation" / "pipelines" / "account_summary" - ).resolve() +@dg.asset_check( + asset=account_events, + description="account_id must be fully populated on the published branch", +) +def account_events_account_id_no_nulls( + bauplan_client: BauplanResource, +) -> dg.AssetCheckResult: + """Post-publish check: count events without an account id on the base branch.""" + nulls = ( + bauplan_client.client.query( + "SELECT COUNT(*) AS n FROM account_events WHERE account_id IS NULL", + ref=BASE_BRANCH, + namespace=NAMESPACE, + ) + .column("n") + .to_pylist()[0] + ) + return dg.AssetCheckResult( + passed=nulls == 0, + metadata={"null_account_ids": int(nulls or 0), "ref": BASE_BRANCH}, + ) - # Detach and poll - state = client.run( - project_dir=str(project_dir), - ref=branch, - namespace=cfg["namespace"], + +@dg.asset_check( + asset=account_events, + description="Auditing pipeline for 'account_events' table. Replicates the one run at ingestion time.", +) +def account_events_audit(bauplan_client: BauplanResource) -> dg.AssetCheckResult: + """Post-publish check: run expectations.""" + state = bauplan_client.client.run( + project_dir=str((AUDIT_DIR / "account_events").resolve()), + ref=BASE_BRANCH, + namespace=NAMESPACE, + dry_run=True, cache="off", - detach=True, - parameters={"start_date": cfg["start_date"], "end_date": cfg["end_date"]}, + strict="on", ) - wait_for_job(client, state.job_id, f"transform on {branch}") - return {"branch": branch, "base_branch": cfg["base_branch"]} + successful = str(state.job_status).lower() == "success" + return dg.AssetCheckResult( + passed=successful, + metadata={"error": str(state.error or ""), "ref": BASE_BRANCH}, + ) -@dg.job(name="ingestion") -def ingestion(): - """Dagster job to run the full ingestion pipeline.""" - imported = run_import_data() - audited = run_wap(imported) - run_merge(audited) +ingestion = dg.define_asset_job( + name="ingestion", + selection=dg.AssetSelection.assets("transactions", "account_events"), +) -@dg.job(name="transformation") -def transformation(): - """Dagster job to run the transformation pipeline""" - transformed = run_transform() - run_merge(transformed) +transformation = dg.define_asset_job( + name="transformation", + selection=dg.AssetSelection.assets("account_activity_summary"), +) +RESOURCES = {"bauplan_client": BauplanResource(api_key=dg.EnvVar("BAUPLAN_API_KEY"))} defs = dg.Definitions( + assets=[transactions, account_events, account_activity_summary], + asset_checks=[ + transactions_txn_id_unique, + transactions_audit, + account_events_account_id_no_nulls, + account_events_audit, + ], jobs=[ingestion, transformation], - resources={"bauplan_client": BauplanResource(api_key=dg.EnvVar("BAUPLAN_API_KEY"))}, + resources=RESOURCES, ) +ALL_DEFINITIONS = [ + transactions, + account_events, + account_activity_summary, + transactions_txn_id_unique, + transactions_audit, + account_events_account_id_no_nulls, + account_events_audit, +] + @app.command("ingest") def ingestion_command( @@ -181,35 +366,16 @@ def ingestion_command( datetime, typer.Argument(formats=["%Y-%m-%d"], help="Partition date, YYYY-MM-DD"), ], - dt_partition_column: Annotated[ - str, typer.Argument(help="Datetime column to use for daily partition") - ], - namespace: Annotated[str, typer.Option(help="Target namespace")] = "workshop", - base_branch: Annotated[ - str, typer.Option(help="Branch to fork from and merge into") - ] = "workshop.main", ) -> None: - """Launch the ingestion WAP job in-process for one table partition.""" - run_config = { - "ops": { - "import_data": { - "config": { - "table": table.value, - "year": str(date.year), - "month": f"{date.month:02d}", - "day": f"{date.day:02d}", - "dt_partition_column": dt_partition_column, - "base_branch": base_branch, - "namespace": namespace, - } - } - } - } - result = ingestion.execute_in_process( - run_config=run_config, - resources={ - "bauplan_client": BauplanResource(api_key=dg.EnvVar("BAUPLAN_API_KEY")) - }, + """Materialize one daily partition of a source table, running its WAP audit + and post-publish checks.""" + result = dg.materialize( + assets=ALL_DEFINITIONS, + selection=dg.AssetSelection.assets(table.value) + | dg.AssetSelection.checks_for_assets(table.value), + partition_key=date.strftime("%Y-%m-%d"), + resources=RESOURCES, + raise_on_error=False, ) if not result.success: raise typer.Exit(code=1) @@ -217,40 +383,19 @@ def ingestion_command( @app.command("transform") def transformation_command( - start_date: Annotated[ - datetime, - typer.Argument(formats=["%Y-%m-%d"], help="Start date, YYYY-MM-DD"), - ], - end_date: Annotated[ + date: Annotated[ datetime, - typer.Argument(formats=["%Y-%m-%d"], help="End date, YYYY-MM-DD"), + typer.Argument(formats=["%Y-%m-%d"], help="Partition date, YYYY-MM-DD"), ], - namespace: Annotated[ - str, typer.Option(help="Namespace of the source tables") - ] = "workshop", - base_branch: Annotated[ - str, typer.Option(help="Branch to fork from and merge into") - ] = "workshop.main", ) -> None: - """Run the transformation pipeline on a dev branch and merge it into the base branch.""" - - run_config = { - "ops": { - "transform": { - "config": { - "start_date": str(start_date.date()), - "end_date": str(end_date.date()), - "base_branch": base_branch, - "namespace": namespace, - } - } - } - } - result = transformation.execute_in_process( - run_config=run_config, - resources={ - "bauplan_client": BauplanResource(api_key=dg.EnvVar("BAUPLAN_API_KEY")) - }, + """Materialize one daily partition of the account activity summary; use + Dagster backfills to cover date ranges.""" + result = dg.materialize( + assets=ALL_DEFINITIONS, + selection=dg.AssetSelection.assets("account_activity_summary"), + partition_key=date.strftime("%Y-%m-%d"), + resources=RESOURCES, + raise_on_error=False, ) if not result.success: raise typer.Exit(code=1) diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary/models.py b/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary/models.py index e72e569f..c171e36e 100644 --- a/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary/models.py +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary/models.py @@ -1,18 +1,28 @@ -from locale import currency - import bauplan -import pyarrow as pa +from bpln_sdk import ( + Model, + Read, + Filter, +) + +from .schemas import ( + SettledTxSchema, + TxPushdownSchema, + DailySpendSchema, + AccPushdownSchema, + AccActivitySchema, +) @bauplan.python("3.13", pip={"polars": "1.42.1"}) @bauplan.model() def settled_transactions( - data: pa.Table = bauplan.Model( - "transactions", - columns=["account_id", "amount", "merchant_category", "status", "txn_ts"], - filter="txn_ts >= $start_date AND txn_ts <= $end_date", - ), -): + data: Annotated[ + Model[TxPushdownSchema], + Read("transactions") + .Filter("txn_ts >= $start_date AND txn_ts < $end_date"), + ], +) -> Model[SettledTxSchema]: """Keep only settled transactions, the ones that actually moved money""" import polars as pl @@ -24,8 +34,8 @@ def settled_transactions( @bauplan.python("3.13", pip={"polars": "1.42.1"}) @bauplan.model() def daily_account_spend( - data: pa.Table = bauplan.Model("settled_transactions"), -): + data: Annotated[Model[SettledTxSchema], Read("settled_transactions")], +) -> Model[DailySpendSchema]: """Aggregate settled spend per account: total, count and average amount""" import polars as pl @@ -48,13 +58,16 @@ def daily_account_spend( overwrite_filter="date >= $start_date AND date < $end_date", ) def account_activity_summary( - daily_account_spend: pa.Table = bauplan.Model("daily_account_spend"), - account_events: pa.Table = bauplan.Model( - "account_events", - columns=["account_id", "event_type", "event_ts"], - filter="event_ts >= $start_date AND event_ts <= $end_date", - ), -): + daily_account_spend: Annotated[ + Model[DailySpendSchema], + Read("daily_account_spend"), + ] + account_events: Annotated[ + Model[AccPushdownSchema], + Read("account_events") + .Filter("event_ts >= $start_date AND event_ts < $end_date"), + ], +) -> Model[AccActivitySchema]: """Per-account view: settled spend joined with event and login counts on a daily basis""" import polars as pl diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/bauplan_project.yaml b/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/bauplan_project.yaml new file mode 100644 index 00000000..a3c836f9 --- /dev/null +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/bauplan_project.yaml @@ -0,0 +1,12 @@ +project: + id: a8001d21-8efb-4a50-87c8-5bb79452a58d + name: account_summary +parameters: + start_date: + type: str + default: "1970-01-01" + required: true + end_date: + type: str + default: "2999-12-31" + required: true diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/models.py b/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/models.py new file mode 100644 index 00000000..c171e36e --- /dev/null +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/models.py @@ -0,0 +1,88 @@ +import bauplan +from bpln_sdk import ( + Model, + Read, + Filter, +) + +from .schemas import ( + SettledTxSchema, + TxPushdownSchema, + DailySpendSchema, + AccPushdownSchema, + AccActivitySchema, +) + + +@bauplan.python("3.13", pip={"polars": "1.42.1"}) +@bauplan.model() +def settled_transactions( + data: Annotated[ + Model[TxPushdownSchema], + Read("transactions") + .Filter("txn_ts >= $start_date AND txn_ts < $end_date"), + ], +) -> Model[SettledTxSchema]: + """Keep only settled transactions, the ones that actually moved money""" + import polars as pl + + df = pl.DataFrame(data).with_columns(date=pl.col("txn_ts").dt.date()) + + return df.filter(pl.col("status") == "settled").drop("status").to_arrow() + + +@bauplan.python("3.13", pip={"polars": "1.42.1"}) +@bauplan.model() +def daily_account_spend( + data: Annotated[Model[SettledTxSchema], Read("settled_transactions")], +) -> Model[DailySpendSchema]: + """Aggregate settled spend per account: total, count and average amount""" + import polars as pl + + df = pl.DataFrame(data) + return ( + df.group_by("account_id", "date") + .agg( + pl.col("amount").sum().round(2).alias("total_amount"), + pl.len().alias("transaction_count"), + pl.col("amount").mean().round(2).alias("avg_amount"), + ) + .to_arrow() + ) + + +@bauplan.python("3.13", pip={"polars": "1.42.1"}) +@bauplan.model( + materialization_strategy="OVERWRITE_PARTITIONS", + partitioned_by=["date"], + overwrite_filter="date >= $start_date AND date < $end_date", +) +def account_activity_summary( + daily_account_spend: Annotated[ + Model[DailySpendSchema], + Read("daily_account_spend"), + ] + account_events: Annotated[ + Model[AccPushdownSchema], + Read("account_events") + .Filter("event_ts >= $start_date AND event_ts < $end_date"), + ], +) -> Model[AccActivitySchema]: + """Per-account view: settled spend joined with event and login counts on a daily basis""" + import polars as pl + + spend_df = pl.DataFrame(daily_account_spend) + activity_df = ( + pl.DataFrame(account_events) + .with_columns(date=pl.col("event_ts").dt.date()) + .group_by("account_id", "date") + .agg( + pl.len().alias("event_count"), + (pl.col("event_type") == "login").sum().alias("login_count"), + ) + .with_columns( + pl.col("event_count").fill_null(0), + pl.col("login_count").fill_null(0), + ) + ) + return activity_df.join(spend_df, on=["account_id", "date"], how="left").to_arrow() diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/transactions/pyproject.toml b/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/pyproject.toml similarity index 78% rename from examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/transactions/pyproject.toml rename to examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/pyproject.toml index 70ca5283..2edb3210 100644 --- a/examples/learn/08-ingestion-and-transformation-with-dagster/src/ingestion/pipelines/wap/transactions/pyproject.toml +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/pyproject.toml @@ -1,5 +1,5 @@ [project] -name = "wap_transaction" +name = "account_summary" version = "0.1.0" requires-python = "~=3.13" dependencies = ["bauplan~=0.1.16"] diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/schemas.py b/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/schemas.py new file mode 100644 index 00000000..f05fe9b7 --- /dev/null +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/transformation/pipelines/account_summary_with_contracts/schemas.py @@ -0,0 +1,82 @@ +from typing import Annotated +import bauplan +from bpln_sdk import ( + TableField, + TableSchema, + Float64, + Int32, + Int64, + String, + Timestamp, +) + + +class TxPushdownSchema(TableSchema): + """Transaction information for analyzing daily spending.""" + + account_id: Annotated[Int64, TableField(doc='Account identifier')] + amount: Annotated[Float64, TableField(doc='Transaction amount')] + merchant_category: Annotated[String, + TableField(doc='Vendor category for transaction') + ] + status: Annotated[String, TableField(doc='Status of transaction')] + txn_ts: Annotated[Timestamp, TableField(doc='Time of transaction')] + + +class SettledTxSchema(TableSchema): + """Transactions with a "settled" status.""" + + account_id: Annotated[Int64, TableField(lineage=TxPushdownSchema['account_id'])] + amount: Annotated[Float64, TableField(doc='Transaction amount')] + merchant_category: Annotated[String, + TableField(doc='Vendor category for transaction') + ] + txn_ts: Annotated[Timestamp, TableField(doc='Time of transaction')] + date: Annotated[Timestamp, TableField(doc='Date of transaction')] + + +class DailySpendSchema(TableSchema): + """Total and average transaction amounts by day for each account.""" + + account_id: Annotated[Int64, TableField(doc='Account identifier')] + date: Annotated[Timestamp, TableField(doc='Date of transaction')] + total_amount: Annotated[Float64, TableField(doc='Total daily amount by account')] + transaction_count: Int32 + avg_amount: Annotated[Float64, TableField(doc='Average daily amount by account')] + + +class AccPushdownSchema(TableSchema): + """Account "events" for analyzing daily account activity.""" + + account_id: Annotated[Int64, TableField(doc='Account identifier')] + event_type: Annotated[String, TableField(doc='Type of account event')] + event_ts: Annotated[Timestamp, + TableField(doc='Time of account event in a specified time period.') + ] + + +class AccActivitySchema(TableSchema): + """ + Account activity by date. + Implemented as a left-join of "account_events" and "daily_account_spend", where + events that occur on days without any account spending have NULL values for + event_count and login_count filled in as `0`. + + Note: we don't expect event_count to be 0, because it wouldn't be present in the + account_events table in the first place; but, 0 login activities is possible. + """ + + account_id: Annotated[Int64, TableField(lineage=AccPushdownSchema['account_id'])] + date: Annotated[Timestamp, + TableField( + lineage=AccPushdownSchema['event_ts'], + doc='Date of event transaction' + ) + ] + event_count: Int32 + login_count: Annotated[Int32, TableField(doc='Count of "login" events')] + total_amount: Annotated[Float64, TableField(lineage=DailySpendSchema['total_amount'])] + transaction_count: Annotated[Int32, + TableField(lineage=DailySpendSchema['transaction_count']) + ] + avg_amount: Annotated[Float64, TableField(lineage=DailySpendSchema['avg_amount'])] diff --git a/examples/learn/08-ingestion-and-transformation-with-dagster/src/utils.py b/examples/learn/08-ingestion-and-transformation-with-dagster/src/utils.py index 7335921b..44f59777 100644 --- a/examples/learn/08-ingestion-and-transformation-with-dagster/src/utils.py +++ b/examples/learn/08-ingestion-and-transformation-with-dagster/src/utils.py @@ -1,6 +1,8 @@ import time +from typing import Any import bauplan +import dagster as dg from bauplan.schema import JobState # Polling cadence and upper bound for detached bauplan jobs @@ -32,3 +34,37 @@ def wait_for_job(client: bauplan.Client, job_id: str | None, label: str) -> None raise RuntimeError( f"{label} job {job_id} failed: {job.human_readable_status} - {job.error_message}" ) + + +def get_table_metadata( + client: bauplan.Client, + table: str, + namespace: str, + ref: str, + partition_predicate: str, + extra: dict[str, Any], +) -> dict[str, Any]: + """Assemble the Dagster metadata for a materialized table: total and + partition-scoped row counts, the column schema, plus caller-provided entries. + The dagster/ keys are standard ones that get dedicated rendering in the UI""" + partition_rows = ( + client.query( + f"SELECT COUNT(*) AS n FROM {table} WHERE {partition_predicate}", + ref=ref, + namespace=namespace, + ) + .column("n") + .to_pylist()[0] + ) + table_info = client.get_table(table=table, ref=ref, namespace=namespace) + return { + "dagster/row_count": int(table_info.records or 0), + "dagster/partition_row_count": int(partition_rows or 0), + "dagster/column_schema": dg.TableSchema( + columns=[ + dg.TableColumn(name=field.name, type=field.type) + for field in table_info.fields + ] + ), + **extra, + }