ADO-420 Research Astronomer Cosmos

docs/dbt_integration/examples/dbt_af_example.py

@@ -0,0 +1,65 @@
+"""
+Example: dbt-af (Toloka) integration with Apache Airflow
+
+dbt-af auto-generates Airflow DAGs from a dbt manifest.json
+It parses the dbt project graph, creates domain-based DAGs with proper
+scheduling, cross-DAG sensors, and test grouping
+
+Requirements:
+    pip install dbt-af dbt-postgres
+
+Key dbt-af features demonstrated:
+    - Domain-based DAG separation (staging, intermediate, marts - separate DAGs)
+    - Cross-schedule dependencies (daily - hourly - daily via ExternalTaskSensor)
+    - Singular test tasks (not_null checks as inline tasks)
+    - Snapshot support (customers_snapshot_ts)
+    - Backfill DAGs (auto-generated alongside scheduled DAGs)
+    - Manual "run single model" DAG with Airflow UI params
+    - Source freshness sensors (for sources with freshness config)
+    - dry_run mode (skips actual dbt execution, disables catchup)
+
+IMPORTANT — dbt-af project structure constraints:
+    - Schedule values MUST use @ prefix: "@daily", "@hourly", "@weekly"
+      (without @ prefix, dbt-af silently falls back to @daily)
+    - Generic tests (unique/not_null in schema.yml) are INCOMPATIBLE with dbt-af
+      due to fqn[3] IndexError. Use singular tests in tests/dbt/models/<domain>/ instead.
+    - Every model config MUST have: sql_cluster, py_cluster, daily_sql_cluster, bf_cluster
+
+Running tests:
+    pip install dbt-af dbt-postgres
+    pytest tests/integration/dbt/test_dbt_af_integration.py -v
+"""
+
+from pathlib import Path
+
+from dbt_af.conf import Config, DbtDefaultTargetsConfig, DbtProjectConfig
+from dbt_af.dags import compile_dbt_af_dags
+
+DBT_PROJECT_DIR = Path(__file__).parent.parent.parent.parent / "dbt_projects" / "demo"
+
+config = Config(
+    dbt_project=DbtProjectConfig(
+        dbt_project_name="dbt_af_demo",
+        dbt_models_path=DBT_PROJECT_DIR / "models",
+        dbt_project_path=DBT_PROJECT_DIR,
+        dbt_profiles_path=DBT_PROJECT_DIR,
+        dbt_target_path=DBT_PROJECT_DIR / "target",
+        dbt_log_path=DBT_PROJECT_DIR / "logs",
+        dbt_schema="public",
+    ),
+    dbt_default_targets=DbtDefaultTargetsConfig(
+        default_target="dev",
+    ),
+    max_active_dag_runs=1,
+    include_single_model_manual_dag=True,
+    debug_mode_enabled=False,
+    dry_run=True,
+)
+
+dags = compile_dbt_af_dags(
+    manifest_path=str(DBT_PROJECT_DIR / "target" / "manifest.json"),
+    config=config,
+)
+
+for dag_name, dag in dags.items():
+    globals()[dag_name] = dag

hatch_build.py

@@ -212,6 +212,9 @@
         "towncrier>=23.11.0",
         "twine>=4.0.2",
     ],
+    "devel-dbt": [
+        "dbt-af>=0.14.0",
+    ],
     "devel-duckdb": [
         # Python 3.12 support was added in 0.10.0
         "duckdb>=0.10.0; python_version >= '3.12'",
@@ -281,6 +284,7 @@
     "devel": [
         "apache-airflow[celery]",
         "apache-airflow[cncf-kubernetes]",
+        "apache-airflow[devel-dbt]",
         "apache-airflow[devel-debuggers]",
         "apache-airflow[devel-devscripts]",
         "apache-airflow[devel-duckdb]",

pyproject.toml

@@ -89,7 +89,7 @@ dynamic = ["version", "optional-dependencies", "dependencies"]
 #
 # START DEVEL EXTRAS HERE
 #
-# devel, devel-all-dbs, devel-ci, devel-debuggers, devel-devscripts, devel-duckdb, devel-hadoop,
+# devel, devel-all-dbs, devel-ci, devel-dbt, devel-debuggers, devel-devscripts, devel-duckdb, devel-hadoop,
 # devel-mypy, devel-sentry, devel-static-checks, devel-tests
 #
 # END DEVEL EXTRAS HERE

tests/integration/dbt/conftest.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+
+DBT_PROJECT_DIR = Path(__file__).parent / "demo"
+
+
+def pytest_configure(config):
+    """Run dbt compile to generate manifest.json before test collection"""
+    manifest = DBT_PROJECT_DIR / "target" / "manifest.json"
+    if manifest.exists():
+        return
+
+    subprocess.run(
+        [
+            "dbt", "compile",
+            "--project-dir", str(DBT_PROJECT_DIR),
+            "--profiles-dir", str(DBT_PROJECT_DIR),
+        ],
+        check=True,
+        capture_output=True,
+        text=True,
+    )

tests/integration/dbt/demo/dbt_project.yml

@@ -0,0 +1,24 @@
+name: 'dbt_af_demo'
+version: '1.0.0'
+config-version: 2
+
+profile: 'dbt_af_demo'
+
+model-paths: ["models"]
+snapshot-paths: ["snapshots"]
+macro-paths: ["macros"]
+target-path: "target"
+clean-targets: ["target", "dbt_packages"]
+
+models:
+  dbt_af_demo:
+    staging:
+      +materialized: view
+    intermediate:
+      +materialized: table
+    marts:
+      +materialized: table
+      core:
+        +schema: marts_core
+      analytics:
+        +schema: marts_analytics

tests/integration/dbt/demo/models/intermediate/int_orders_enriched.sql

@@ -0,0 +1,20 @@
+{{ config(materialized='table') }}
+
+WITH orders AS (
+    SELECT * FROM {{ ref('stg_orders') }}
+),
+customers AS (
+    SELECT * FROM {{ ref('stg_customers') }}
+)
+SELECT
+    o.order_id,
+    o.customer_id,
+    o.order_date,
+    o.amount_cents,
+    o.status,
+    c.customer_name,
+    c.account_balance,
+    c.market_segment,
+    o.loaded_at
+FROM orders o
+LEFT JOIN customers c ON o.customer_id = c.customer_id

tests/integration/dbt/demo/models/intermediate/schema.yml

@@ -0,0 +1,11 @@
+version: 2
+
+models:
+  - name: int_orders_enriched
+    description: "Orders enriched with customer data"
+    config:
+      schedule: "@hourly"
+      sql_cluster: dev
+      py_cluster: dev
+      daily_sql_cluster: dev
+      bf_cluster: dev

tests/integration/dbt/demo/models/marts/analytics/revenue_by_region.sql

@@ -0,0 +1,13 @@
+{{ config(materialized='table') }}
+
+WITH orders AS (
+    SELECT * FROM {{ ref('int_orders_enriched') }}
+)
+SELECT
+    COALESCE(market_segment, 'Unknown') AS market_segment,
+    date_trunc('month', order_date)::date AS order_month,
+    count(DISTINCT order_id) AS order_count,
+    count(DISTINCT customer_id) AS unique_customers,
+    round(sum(amount_cents) / 100.0, 2) AS revenue
+FROM orders
+GROUP BY 1, 2

tests/integration/dbt/demo/models/marts/core/customers.sql

@@ -0,0 +1,8 @@
+{{ config(materialized='table') }}
+
+SELECT
+    customer_id,
+    customer_name,
+    account_balance,
+    market_segment
+FROM {{ ref('stg_customers') }}

tests/integration/dbt/demo/models/marts/core/orders.sql

@@ -0,0 +1,16 @@
+{{ config(materialized='table') }}
+
+SELECT
+    o.order_id,
+    o.customer_id,
+    o.order_date,
+    o.amount_cents,
+    o.status,
+    CASE o.status
+        WHEN 'O' THEN 'Open'
+        WHEN 'F' THEN 'Fulfilled'
+        WHEN 'P' THEN 'Processing'
+    END AS status_label,
+    o.customer_name,
+    o.market_segment
+FROM {{ ref('int_orders_enriched') }} o

tests/integration/dbt/demo/models/marts/schema.yml

@@ -0,0 +1,29 @@
+version: 2
+
+models:
+  - name: orders
+    description: "Core orders mart"
+    config:
+      schedule: "@daily"
+      sql_cluster: dev
+      py_cluster: dev
+      daily_sql_cluster: dev
+      bf_cluster: dev
+
+  - name: customers
+    description: "Core customers mart"
+    config:
+      schedule: "@daily"
+      sql_cluster: dev
+      py_cluster: dev
+      daily_sql_cluster: dev
+      bf_cluster: dev
+
+  - name: revenue_by_region
+    description: "Monthly revenue by market segment"
+    config:
+      schedule: "@daily"
+      sql_cluster: dev
+      py_cluster: dev
+      daily_sql_cluster: dev
+      bf_cluster: dev

tests/integration/dbt/demo/models/staging/schema.yml

@@ -0,0 +1,18 @@
+version: 2
+
+models:
+  - name: stg_orders
+    config:
+      schedule: "@daily"
+      sql_cluster: dev
+      py_cluster: dev
+      daily_sql_cluster: dev
+      bf_cluster: dev
+
+  - name: stg_customers
+    config:
+      schedule: "@daily"
+      sql_cluster: dev
+      py_cluster: dev
+      daily_sql_cluster: dev
+      bf_cluster: dev

tests/integration/dbt/demo/models/staging/sources.yml

@@ -0,0 +1,13 @@
+version: 2
+
+sources:
+  - name: src
+    schema: src
+    freshness:
+      warn_after: {count: 24, period: hour}
+      error_after: {count: 48, period: hour}
+    loaded_at_field: loaded_at
+    tables:
+      - name: orders
+      - name: customers
+      - name: customers_scd

tests/integration/dbt/demo/models/staging/stg_customers.sql

@@ -0,0 +1,7 @@
+SELECT
+    c_custkey    AS customer_id,
+    c_name       AS customer_name,
+    c_nationkey  AS nation_id,
+    c_acctbal    AS account_balance,
+    c_mktsegment AS market_segment
+FROM {{ source('src', 'customers') }}

tests/integration/dbt/demo/models/staging/stg_orders.sql

@@ -0,0 +1,8 @@
+SELECT
+    o_orderkey    AS order_id,
+    o_custkey     AS customer_id,
+    o_orderdate   AS order_date,
+    o_totalprice  AS amount_cents,
+    o_orderstatus AS status,
+    now()         AS loaded_at
+FROM {{ source('src', 'orders') }}

tests/integration/dbt/demo/profiles.yml

@@ -0,0 +1,12 @@
+dbt_af_demo:
+  target: dev
+  outputs:
+    dev:
+      type: postgres
+      host: "{{ env_var('DBT_HOST', 'postgres') }}"
+      port: "{{ env_var('DBT_PORT', '5432') | int }}"
+      user: "{{ env_var('DBT_USER', 'postgres') }}"
+      password: "{{ env_var('DBT_PASSWORD', 'airflow') }}"
+      dbname: "{{ env_var('DBT_DBNAME', 'airflow') }}"
+      schema: "{{ env_var('DBT_SCHEMA', 'public') }}"
+      threads: 4

tests/integration/dbt/demo/snapshots/customers_snapshot_ts.sql

@@ -0,0 +1,23 @@
+{% snapshot customers_snapshot_ts %}
+{{
+    config(
+        target_schema='snapshots',
+        unique_key='c_custkey',
+        strategy='timestamp',
+        updated_at='c_updated_at',
+        schedule='@daily',
+        sql_cluster='dev',
+        py_cluster='dev',
+        daily_sql_cluster='dev',
+        bf_cluster='dev'
+    )
+}}
+select
+    c_custkey,
+    c_name,
+    c_nationkey,
+    c_acctbal,
+    c_mktsegment,
+    c_updated_at
+from {{ source('src', 'customers_scd') }}
+{% endsnapshot %}

tests/integration/dbt/demo/tests/dbt/models/marts/not_null_customers_customer_id.sql

@@ -0,0 +1 @@
+select customer_id from {{ ref('customers') }} where customer_id is null

tests/integration/dbt/demo/tests/dbt/models/marts/not_null_orders_order_id.sql

@@ -0,0 +1 @@
+select order_id from {{ ref('orders') }} where order_id is null

tests/integration/dbt/demo/tests/dbt/models/staging/not_null_stg_customers_customer_id.sql

@@ -0,0 +1 @@
+select customer_id from {{ ref('stg_customers') }} where customer_id is null

tests/integration/dbt/demo/tests/dbt/models/staging/not_null_stg_orders_order_id.sql

@@ -0,0 +1 @@
+select order_id from {{ ref('stg_orders') }} where order_id is null