Part 2: Code patterns, Databricks integration, and production recipes
Multi-Tenant Airflow DAG Generation: Implementation Deep Dive
Part 2: Code patterns, Databricks integration, and production recipes
Quick Recap
In Part 1, we covered three strategies for multi-tenant DAG generation:
- Contract-Based Generation — Generate only relevant DAGs per entity
- Hierarchical Structure — Coordinator + entity-specific execution DAGs
- Dynamic Task Mapping — Single DAG with runtime-expanded tasks
This article provides implementation patterns for Databricks-based pipelines.
The Databricks Challenge
Your setup: Airflow orchestrates Python code that executes SQL in Databricks. The Python loader checks metadata conditions before running transformations.
The scaling problems are Databricks-specific:
- Cluster startup latency : 2–5 minutes per cluster start
- SQL Warehouse costs : Per-second billing adds up
- API rate limits : Too many concurrent calls fail
- Worker blocking : Python waits while Databricks processes
Let’s solve each one.
Pattern 1: Cluster-Aware Task Grouping
Problem
Each task starts its own cluster or connection, causing:
- 50 entities × 2 min startup = 100 minutes of startup time
- Cluster fragmentation (many small clusters vs. few large ones)
Solution: Group Tasks by Cluster
# src/orchestration/cluster_aware_generator.py
from dataclasses import dataclass
from typing import List, Dict
from collections import defaultdict
@dataclass
class ClusterConfig:
cluster_id: str
size: str # small, medium, large
max_concurrent_tasks: int = 8
def group_by_cluster(entities: List[str], contracts: Dict) -> Dict[str, List[str]]:
"""Group entities by optimal cluster assignment."""
cluster_groups = defaultdict(list)
for entity in entities:
contract = contracts[entity]
tier = contract.entity.tier
# Assign cluster based on tier and workload
if tier == "premium":
cluster = "cluster_large_01"
elif len(contract.delivery_contract.datasets) > 10:
cluster = "cluster_medium_01"
else:
cluster = "cluster_small_01"
cluster_groups[cluster].append(entity)
return dict(cluster_groups)
def generate_cluster_aware_dag(cluster_id: str, entities: List[str]) -> DAG:
"""Generate DAG that processes multiple entities on same cluster."""
with DAG(
dag_id=f"load_{cluster_id}",
schedule=None,
tags=["execution", "cluster-aware"],
max_active_tasks=8, # Match cluster capacity
) as dag:
# Single cluster startup task
start_cluster = DatabricksStartClusterOperator(
task_id="start_cluster",
cluster_id=cluster_id,
deferrable=True,
)
# Process entities sequentially to avoid cluster overload
entity_tasks = []
for entity in entities:
entity_task = process_entity_task(
entity_id=entity,
cluster_id=cluster_id,
)
entity_tasks.append(entity_task)
# Stop cluster when done
stop_cluster = DatabricksStopClusterOperator(
task_id="stop_cluster",
cluster_id=cluster_id,
trigger_rule="all_done",
)
# Wire: start -> all entities -> stop
start_cluster >> entity_tasks >> stop_cluster
return dag
Coordinator Integration
# dags/cluster_coordinator.py
with DAG(
dag_id="cluster_coordinator",
schedule="@daily",
tags=["orchestration"],
) as dag:
@task
def plan_cluster_allocation():
"""Determine cluster assignments for today's workload."""
ready_entities = get_ready_entities()
contracts = {e: load_contract(e) for e in ready_entities}
return group_by_cluster(ready_entities, contracts)
@task
def trigger_cluster_dags(cluster_plan: Dict[str, List[str]]):
"""Trigger one DAG per cluster with its entities."""
for cluster_id, entities in cluster_plan.items():
trigger_dag(
dag_id=f"load_{cluster_id}",
conf={"entities": entities},
)
plan = plan_cluster_allocation()
trigger_cluster_dags(plan)
Pattern 2: SQL Warehouse Batching
Problem
Each SQL query opens a connection, executes, closes. Connection overhead dominates small queries.
Solution: Batch SQL Statements
# src/orchestration/sql_batcher.py
from typing import List, Tuple
def batch_sql_statements(
transformations: List[Transformation],
max_batch_size: int = 10,
max_batch_duration_estimate: int = 300, # seconds
) -> List[List[Transformation]]:
"""Batch transformations for efficient SQL Warehouse execution."""
batches = []
current_batch = []
current_duration = 0
for t in transformations:
estimated_duration = estimate_query_duration(t)
if (len(current_batch) >= max_batch_size or
current_duration + estimated_duration > max_batch_duration_estimate):
if current_batch:
batches.append(current_batch)
current_batch = [t]
current_duration = estimated_duration
else:
current_batch.append(t)
current_duration += estimated_duration
if current_batch:
batches.append(current_batch)
return batches
def execute_batch(batch: List[Transformation], entity_id: str) -> dict:
"""Execute a batch of SQL statements in single connection."""
with get_databricks_connection() as conn:
results = []
for t in batch:
sql = generate_sql(entity_id, t)
result = conn.execute(sql)
results.append({
"transformation": t.id,
"rows_affected": result.rowcount,
"duration": result.duration,
})
return {"batch_results": results}
Airflow Task Using Batching
@task
def execute_transformation_batch(
entity_id: str,
batch_index: int,
batch: List[dict],
):
"""Execute batch of transformations for an entity."""
transformations = [Transformation.from_dict(t) for t in batch]
results = execute_batch(transformations, entity_id)
# Log for observability
for r in results["batch_results"]:
log.info(f"{entity_id}/{r['transformation']}: {r['rows_affected']} rows in {r['duration']}s")
return results
Pattern 3: Deferrable Databricks Operators
Problem
Standard operators block workers while Databricks processes queries (minutes to hours).
Solution: Custom Deferrable Operator
# src/operators/deferrable_databricks.py
from airflow.models import BaseOperator
from airflow.triggers.base import BaseTrigger, TriggerEvent
from airflow.utils.context import Context
import asyncio
class DatabricksQueryTrigger(BaseTrigger):
"""Trigger that polls Databricks query status without blocking."""
def __init__(self, query_id: str, databricks_conn_id: str):
super().__init__()
self.query_id = query_id
self.databricks_conn_id = databricks_conn_id
def serialize(self):
return (
"src.operators.deferrable_databricks.DatabricksQueryTrigger",
{"query_id": self.query_id, "databricks_conn_id": self.databricks_conn_id},
)
async def run(self):
"""Poll Databricks until query completes."""
while True:
status = await self._check_query_status()
if status == "FINISHED":
yield TriggerEvent({"status": "success", "query_id": self.query_id})
return
elif status in ("FAILED", "CANCELED"):
yield TriggerEvent({"status": "failed", "query_id": self.query_id})
return
await asyncio.sleep(10) # Poll every 10 seconds
async def _check_query_status(self):
# Async call to Databricks API
async with aiohttp.ClientSession() as session:
url = f"{self.databricks_host}/api/2.0/sql/statements/{self.query_id}"
async with session.get(url, headers=self.headers) as resp:
data = await resp.json()
return data.get("status", {}).get("state")
class DeferrableDatabricksSqlOperator(BaseOperator):
"""Execute SQL on Databricks without blocking worker."""
def __init__(
self,
sql: str,
databricks_conn_id: str = "databricks_default",
warehouse_id: str = None,
**kwargs,
):
super().__init__(**kwargs)
self.sql = sql
self.databricks_conn_id = databricks_conn_id
self.warehouse_id = warehouse_id
def execute(self, context: Context):
"""Submit query and defer to trigger."""
# Submit query
query_id = self._submit_query()
# Defer to trigger for polling
self.defer(
trigger=DatabricksQueryTrigger(
query_id=query_id,
databricks_conn_id=self.databricks_conn_id,
),
method_name="execute_complete",
)
def execute_complete(self, context: Context, event: dict):
"""Called when trigger fires."""
if event["status"] == "success":
return event["query_id"]
else:
raise AirflowException(f"Query {event['query_id']} failed")
def _submit_query(self) -> str:
"""Submit SQL to Databricks SQL Warehouse."""
hook = DatabricksSqlHook(databricks_conn_id=self.databricks_conn_id)
response = hook.run_sql_statement(
sql=self.sql,
warehouse_id=self.warehouse_id,
wait_timeout=0, # Don't wait, return immediately
)
return response["statement_id"]
Using in DAG
with DAG(dag_id="entity_load_deferred", ...) as dag:
for transformation in sub_model.transformations:
DeferrableDatabricksSqlOperator(
task_id=f"load_{transformation.id}",
sql=generate_sql(entity_id, transformation),
databricks_conn_id="databricks_sql_warehouse",
warehouse_id="{{ var.value.sql_warehouse_id }}",
)
Pattern 4: Completeness-Driven Triggering
Problem
Coordinator triggers entities even when data isn’t ready, wasting resources on “nothing to do” runs.
Solution: Completeness API Integration
# src/completeness/api.py
from fastapi import FastAPI
from pydantic import BaseModel
from datetime import date
app = FastAPI()
class CompletenessRequest(BaseModel):
entity_id: str
period: date
check_type: str = "base" # "base" or "enriched"
class CompletenessResponse(BaseModel):
entity_id: str
period: date
is_ready: bool
missing_sources: list[str]
available_sources: list[str]
completeness_pct: float
can_publish_base: bool
can_publish_enriched: bool
@app.post("/completeness/check")
def check_completeness(request: CompletenessRequest) -> CompletenessResponse:
"""Check if entity is ready for processing."""
contract = load_delivery_contract(request.entity_id)
sub_model = derive_sub_model(contract)
available = []
missing = []
for source in sub_model.required_sources:
if source_has_data(source, request.entity_id, request.period):
available.append(source)
else:
missing.append(source)
completeness = len(available) / len(sub_model.required_sources) if sub_model.required_sources else 1.0
return CompletenessResponse(
entity_id=request.entity_id,
period=request.period,
is_ready=len(missing) == 0,
missing_sources=missing,
available_sources=available,
completeness_pct=completeness,
can_publish_base=check_inner_sources_ready(sub_model, available),
can_publish_enriched=check_all_sources_ready(sub_model, available),
)
Airflow Sensor Using API
# src/operators/completeness_sensor.py
from airflow.sensors.base import BaseSensorOperator
import requests
class CompletenessSensor(BaseSensorOperator):
"""Wait for entity data completeness before proceeding."""
def __init__(
self,
entity_id: str,
completeness_api_url: str,
check_type: str = "base",
**kwargs,
):
super().__init__(**kwargs)
self.entity_id = entity_id
self.completeness_api_url = completeness_api_url
self.check_type = check_type
def poke(self, context) -> bool:
"""Check if entity is ready."""
execution_date = context["execution_date"].date()
response = requests.post(
f"{self.completeness_api_url}/completeness/check",
json={
"entity_id": self.entity_id,
"period": execution_date.isoformat(),
"check_type": self.check_type,
},
)
result = response.json()
if result["is_ready"]:
self.log.info(f"Entity {self.entity_id} is ready with {result['completeness_pct']:.0%} completeness")
return True
else:
self.log.info(f"Entity {self.entity_id} waiting for: {result['missing_sources']}")
return False
Coordinator with Smart Triggering
with DAG(dag_id="smart_coordinator", schedule="*/15 * * * *", ...) as dag:
@task
def get_entities_to_check():
"""Get entities that might have new data."""
# Only check entities expecting delivery today
return [
e for e in get_all_entities()
if get_delivery_calendar(e).expects_delivery_today()
]
@task
def check_and_trigger(entity_id: str):
"""Check completeness and trigger if ready."""
response = check_completeness_api(entity_id, today())
if response["can_publish_base"]:
# Check if already processed
if not already_processed(entity_id, today(), "base"):
trigger_dag(f"{entity_id}_load", conf={"mode": "base"})
return {"entity": entity_id, "triggered": "base"}
if response["can_publish_enriched"]:
if not already_processed(entity_id, today(), "enriched"):
trigger_dag(f"{entity_id}_load", conf={"mode": "enriched"})
return {"entity": entity_id, "triggered": "enriched"}
return {"entity": entity_id, "triggered": None}
entities = get_entities_to_check()
results = check_and_trigger.expand(entity_id=entities)
Pattern 5: Entity Resource Pools
Problem
A single large entity can starve other entities of workers.
Solution: Airflow Pools per Entity Tier
# Configuration: Create pools in Airflow UI or via CLI
# airflow pools set premium_entities 16 "Premium tier entity processing"
# airflow pools set standard_entities 8 "Standard tier entity processing"
# airflow pools set basic_entities 4 "Basic tier entity processing"
def get_pool_for_entity(entity_id: str) -> str:
"""Get appropriate pool based on entity tier."""
contract = load_delivery_contract(entity_id)
tier = contract.entity.tier
return {
"premium": "premium_entities",
"standard": "standard_entities",
"basic": "basic_entities",
}.get(tier, "default_pool")
# In DAG generation
def create_entity_dag(entity_id: str) -> DAG:
pool = get_pool_for_entity(entity_id)
with DAG(dag_id=f"{entity_id}_load", ...) as dag:
for transformation in sub_model.transformations:
DeferrableDatabricksSqlOperator(
task_id=f"load_{transformation.id}",
sql=generate_sql(entity_id, transformation),
pool=pool, # Use entity-specific pool
pool_slots=1,
)
return dag
Complete Example: Entity DAG Generator
Putting it all together:
# dags/generators/entity_dag_generator.py
from airflow import DAG
from airflow.utils.dates import days_ago
from datetime import timedelta
from src.contracts import load_delivery_contract, derive_sub_model
from src.operators.deferrable_databricks import DeferrableDatabricksSqlOperator
from src.operators.completeness_sensor import CompletenessSensor
from src.orchestration.sql_batcher import batch_sql_statements
def generate_entity_dag(entity_id: str) -> DAG:
"""Generate complete DAG for an entity."""
# Load metadata
contract = load_delivery_contract(entity_id)
sub_model = derive_sub_model(contract)
pool = get_pool_for_entity(entity_id)
# Batch transformations for efficiency
batches = batch_sql_statements(sub_model.transformations)
with DAG(
dag_id=f"{entity_id}_load",
schedule=None, # Triggered by coordinator
start_date=days_ago(1),
tags=["execution", entity_id, contract.entity.tier],
default_args={
"owner": contract.contacts.primary.email if contract.contacts else "data-team",
"retries": 2,
"retry_delay": timedelta(minutes=5),
"pool": pool,
},
doc_md=f"""
# {contract.entity.name} Data Load
**Tier**: {contract.entity.tier}
**Datasets**: {', '.join(d.id for d in contract.delivery_contract.datasets if d.status == 'active')}
**SLA**: {contract.delivery_contract.datasets[0].sla if contract.delivery_contract.datasets else 'N/A'}
## Sub-Model
- Tables: {len(sub_model.tables)}
- Transformations: {len(sub_model.transformations)}
- Batches: {len(batches)}
""",
) as dag:
# Wait for base data to be ready
wait_for_data = CompletenessSensor(
task_id="wait_for_completeness",
entity_id=entity_id,
completeness_api_url="{{ var.value.completeness_api_url }}",
check_type="{{ dag_run.conf.get('mode', 'base') }}",
mode="reschedule",
poke_interval=300,
timeout=3600 * 6, # 6 hours max wait
)
# Execute transformation batches
batch_tasks = []
for i, batch in enumerate(batches):
batch_sql = "\n;\n".join(
generate_sql(entity_id, t) for t in batch
)
batch_task = DeferrableDatabricksSqlOperator(
task_id=f"batch_{i:02d}",
sql=batch_sql,
databricks_conn_id="databricks_sql_warehouse",
)
batch_tasks.append(batch_task)
# Publish completeness status
@task
def publish_status(entity_id: str, mode: str):
update_completeness_dashboard(
entity_id=entity_id,
period=get_execution_date(),
status="completed",
mode=mode,
)
status = publish_status(
entity_id=entity_id,
mode="{{ dag_run.conf.get('mode', 'base') }}",
)
# Wire dependencies
wait_for_data >> batch_tasks[0]
for i in range(len(batch_tasks) - 1):
batch_tasks[i] >> batch_tasks[i + 1]
batch_tasks[-1] >> status
return dag
# Generate DAGs for all entities with active contracts
for entity_id in get_entities_with_contracts():
globals()[f"{entity_id}_load"] = generate_entity_dag(entity_id)
Monitoring & Observability
DAG Generation Metrics
# src/metrics/dag_metrics.py
from prometheus_client import Counter, Gauge, Histogram
# Track DAG generation
dags_generated = Counter(
"mdde_dags_generated_total",
"Total DAGs generated",
["entity", "tier"],
)
transformations_per_dag = Histogram(
"mdde_transformations_per_dag",
"Transformations per generated DAG",
buckets=[1, 5, 10, 25, 50, 100],
)
# Track execution
entity_load_duration = Histogram(
"mdde_entity_load_duration_seconds",
"Time to load entity data",
["entity", "tier"],
buckets=[60, 300, 600, 1800, 3600],
)
completeness_wait_time = Histogram(
"mdde_completeness_wait_seconds",
"Time waiting for source completeness",
["entity"],
buckets=[0, 300, 600, 1800, 3600, 7200],
)
Completeness Dashboard Query
-- DuckDB query for completeness dashboard
SELECT
e.entity_id,
e.entity_name,
c.period,
c.check_type,
c.completeness_pct,
c.missing_sources,
c.can_publish_base,
c.can_publish_enriched,
c.checked_at,
CASE
WHEN c.can_publish_enriched THEN 'enriched'
WHEN c.can_publish_base THEN 'base'
ELSE 'waiting'
END as status
FROM entities e
LEFT JOIN completeness_checks c ON e.entity_id = c.entity_id
WHERE c.period = current_date()
ORDER BY e.tier DESC, c.completeness_pct ASC;
Summary: Which DAGs to Generate
Next Steps
- Audit your contracts — Map which datasets each entity actually uses
- Implement sub-model derivation — Build the filtering logic
- Start with hierarchical DAGs — Coordinator + entity execution
- Add deferrable operators — Free workers during Databricks execution
- Deploy completeness API — Enable smart triggering
- Measure and iterate — Track DAG count, scheduler time, resource usage
This is Part 2 of the Multi-Tenant Airflow series. See Part 1 for the conceptual framework.