| State | Draft |
| Discussion Thread | |
| Vote Thread | |
| Vote Result Thread | |
| Progress Tracking (PR/GitHub Project/Issue Label) | |
| Date Created |
|
| Version Released | |
| Authors |
In today's evolving data landscape, organizations face significant challenges:
Real-World Pain Points:
This proposal leverages Airflow's strengths (production reliability, 1000+ integrations, governance) while adding AI-native capabilities for intelligent data operations.
Based on Pydantic AI's capabilities, while it provides excellent agent framework, multi-model support, and structured outputs, it lacks:
An example view of how Operators look like below:
from airflow.providers.ai.operators import (
LLMSQLQueryOperator,
LLMSchemaCompareOperator,
LLMDataQualityOperator,
LLMFileAnalysisOperator
)
from airflow.sdk import Asset
# Enhanced Asset with structured metadata
customer_postgres = Asset(
name="customer_data_postgres",
uri="postgres://warehouse/public/customers",
conn_id="postgres_warehouse",
schema={
"customer_id": "integer PRIMARY KEY",
"full_name": "varchar(255) NOT NULL",
"email_address": "varchar(255) UNIQUE",
"created_at": "timestamp DEFAULT now()",
"total_revenue": "decimal(10,2)"
},
sensitivity="pii"
)
# SQL operator with automatic context injection and safety
# Option 1: Traditional operator approach
sql_analysis = LLMSQLQueryOperator(
task_id="analyze_customer_segments",
prompt="Find top 10 customers by revenue, include their email and signup date",
data_sources=[customer_postgres],
# Automatic context injection includes:
# - Database type: PostgreSQL
# - Available tables and schemas from DbApiHook
# - Column types and constraints
# - Sample data (first 5 rows) -> explicit approvals with htil to read data if its PII or PCI
# - Built-in SQL safety (blocks DROP, DELETE without WHERE, etc.)
)
# Option 2: Decorator approach with pre-processing
@task.llm_sql_query(data_sources=[customer_postgres])
def analyze_customer_segments_with_preprocessing():
# Custom pre-processing logic
current_date = datetime.now().strftime('%Y-%m-%d')
business_hours = get_business_hours()
# Dynamic prompt generation based on context
prompt = f"""
Find top 10 customers by revenue as of {current_date}.
Include their email and signup date.
Filter for customers active during business hours: {business_hours}
"""
return {"prompt": prompt, "additional_context": {"analysis_date": current_date}} |
Each LLM operator has a corresponding decorator for more flexible, Pythonic workflows:
from airflow.providers.ai.decorators import task
# Schema comparison with custom logic
@task.llm_schema_compare(data_sources=[s3_asset, postgres_asset])
def intelligent_schema_validation():
# Pre-processing: check business calendar
is_migration_window = check_migration_window()
if is_migration_window:
prompt = "Compare schemas and generate migration plan for scheduled maintenance window"
else:
prompt = "Compare schemas and flag breaking changes - no migrations allowed"
return {
"prompt": prompt,
"migration_allowed": is_migration_window,
"additional_context": {"maintenance_window": is_migration_window}
}
# Data quality with dynamic rules
@task.llm_data_quality(data_sources=[customer_asset])
def adaptive_quality_checks():
# Pre-processing: get current business rules
current_rules = fetch_business_rules()
seasonal_adjustments = get_seasonal_data_patterns()
prompt = f"""
Validate customer data against current business rules:
{current_rules}
Apply seasonal adjustments for data volume expectations:
{seasonal_adjustments}
Generate appropriate validation queries.
"""
return {
"prompt": prompt,
"business_rules": current_rules,
"seasonal_context": seasonal_adjustments
}
# File analysis with preprocessing
@task.llm_file_analysis(data_sources=[log_files_asset])
def analyze_logs_with_context():
# Pre-processing: get system context
recent_deployments = get_recent_deployments()
system_alerts = get_active_alerts()
prompt = f"""
Analyze log files for anomalies, considering:
- Recent deployments: {recent_deployments}
- Active system alerts: {system_alerts}
Focus on correlation between deployment events and error patterns.
"""
return {
"prompt": prompt,
"deployment_context": recent_deployments,
"alert_context": system_alerts
} |
Benefits of Decorator Approach:
SQL Operator Context (via DbApiHook integration):
# Automatically injected context for PostgreSQL:
{
"database_type": "postgresql",
"version": "15.2",
"available_tables": ["customers", "orders", "products"],
"schema_info": {
"customers": {
"customer_id": {"type": "integer", "nullable": False, "primary_key": True},
"email_address": {"type": "varchar(255)", "nullable": False, "unique": True},
"total_revenue": {"type": "decimal(10,2)", "nullable": True}
}
},
"sample_data": {
"customers": [
{"customer_id": 1, "email_address": "john@example.com", "total_revenue": 1250.00},
{"customer_id": 2, "email_address": "jane@example.com", "total_revenue": 890.50}
]
},
"dialect_features": {
"supports_window_functions": True,
"supports_cte": True,
"date_functions": ["DATE_TRUNC", "EXTRACT", "AGE"]
}
}
File Operator Context (via S3Hook/GCSHook integration):
# Automatically injected context for S3 Parquet files:
{
"storage_type": "s3",
"file_format": "parquet",
"file_size_mb": 245,
"estimated_rows": 1000000,
"schema_info": {
"id": "int64",
"name": "string",
"email": "string",
"signup_date": "timestamp[ns]"
},
"sample_data": [
{"id": 1, "name": "John Doe", "email": "john@example.com"},
{"id": 2, "name": "Jane Smith", "email": "jane@example.com"}
],
"partitioning": ["year", "month"],
"compression": "snappy"
}
SQL Operator Built-in Protection:
class LLMSQLQueryOperator(BaseOperator):
# Built-in dangerous operation blocking
BLOCKED_KEYWORDS = ["DROP", "TRUNCATE", "DELETE FROM", "ALTER TABLE", "GRANT", "REVOKE"]
DEFAULT_SYSTEM_PROMPT = """You are a SQL expert integrated with {database_type}.
SAFETY RULES:
- NEVER generate DROP, TRUNCATE, DELETE without WHERE, or ALTER statements
- Always use proper JOIN syntax for this database type
- Respect column types and constraints provided in schema
- Use database-specific functions when available
CONTEXT: You have access to schema info, sample data, and dialect features.
Generate optimized, safe queries that work with {database_type} version {version}."""
File Operator Built-in Protection:
class LLMFileAnalysisOperator(BaseOperator):
ALLOWED_OPERATIONS = ["read", "analyze", "summarize", "validate"]
DEFAULT_SYSTEM_PROMPT = """You are a file analysis expert for {storage_type} {file_format} files.
SAFETY RULES:
- ONLY read and analyze files, NEVER modify or delete
- Respect file size limits and memory constraints
- Generate efficient queries for large datasets
CONTEXT: File has {estimated_rows} rows, {file_size_mb}MB, partitioned by {partitioning}."""
# Cross-system schema drift detection
schema_drift = LLMSchemaCompareOperator(
task_id="detect_schema_drift",
data_sources=[customer_s3, customer_postgres, customer_snowflake],
prompt="Identify schema mismatches that would break data loading between systems",
# Automatically gets context from each system:
# - S3: Parquet schema, partitioning, file stats
# - PostgreSQL: Table schema, constraints, indexes
# - Snowflake: Column types, clustering keys, data sharing info
)
A new AnalyticsOperator that provides unified access to multi-cloud data:
from airflow.providers.ai.operators import AnalyticsOperator
# Execute queries across different storage systems uniformly
analytics_task = AnalyticsOperator(
task_id="cross_cloud_analysis",
query="{{ ti.xcom_pull(task_ids='quality_check') }}", # LLM-generated query
data_sources=[customer_s3, orders_gcs, inventory_azure], # Multi-cloud
engine="datafusion", # High-performance query engine
output_format="parquet",
output_location="s3://results/analysis/"
)
Why Apache DataFusion for AI Workloads:
Based on Wren AI's experience, DataFusion provides significant advantages for AI-driven data analysis:
We propose both embedded and separate HITL patterns:
Option A: Embedded HITL
quality_check = LLMDataQualityOperator(
task_id="customer_quality_analysis",
data_sources=[customer_s3],
prompt="Generate data quality validation queries",
require_approval=True, # Built-in HITL
approval_timeout=timedelta(hours=2)
)
Option B: Separate HITL Steps
# Generate queries
generate_queries = LLMDataQualityOperator(
task_id="generate_quality_queries",
data_sources=[customer_s3],
prompt="Generate data quality validation queries",
dry_run=True # Don't execute, just generate
)
# Human approval step
approve_queries = ApprovalOperator(
task_id="approve_queries",
body="{{ ti.xcom_pull(task_ids='generate_quality_queries') }}",
allow_modifications=True # Users can edit generated queries
)
# Execute approved queries
execute_analysis = AnalyticsOperator(
task_id="execute_quality_checks",
query="{{ ti.xcom_pull(task_ids='approve_queries') }}",
data_sources=[customer_s3]
)
generate_queries >> approve_queries >> execute_analysis
Here's a real-world scenario combining all components:
from datetime import datetime, timedelta
from airflow.sdk import DAG, Asset
from airflow.providers.ai.operators import (
LLMSchemaCompareOperator,
LLMDataQualityOperator,
AnalyticsOperator
)
from airflow.operators.approval import ApprovalOperator
# Define multi-cloud assets with structured metadata
customer_s3 = Asset(
name="customer_feed_s3",
uri="s3://data-lake/customer/",
conn_id="aws_default",
schema={"id": "int32", "name": "string", "email": "string"},
sensitivity="pii",
format="parquet"
)
customer_postgres = Asset(
name="customer_master_postgres",
uri="postgres://warehouse/public/customers",
conn_id="postgres_default",
schema={"customer_id": "integer", "full_name": "varchar", "email_address": "varchar"},
sensitivity="pii"
)
with DAG(
"intelligent_data_validation",
start_date=datetime(2024, 1, 1),
schedule=timedelta(hours=6),
) as dag:
# 1. Detect schema drift between S3 feed and PostgreSQL master
schema_drift = LLMSchemaCompareOperator(
task_id="detect_schema_drift",
data_sources=[customer_s3, customer_postgres],
prompt="Identify schema mismatches that would break data loading",
output_format="structured_report"
)
# 2. Generate data quality queries for new S3 data
generate_quality_checks = LLMDataQualityOperator(
task_id="generate_quality_queries",
data_sources=[customer_s3],
prompts=[
"Generate summary statistics queries",
"Check for duplicate email addresses",
"Validate data against business rules"
],
dry_run=True
)
# 3. Human approval for generated queries (with edit capability)
approve_queries = ApprovalOperator(
task_id="approve_quality_queries",
body="{{ ti.xcom_pull(task_ids='generate_quality_queries') }}",
allow_modifications=True,
timeout=timedelta(hours=2)
)
# 4. Execute approved quality checks using DataFusion
execute_quality_checks = AnalyticsOperator(
task_id="run_quality_analysis",
query="{{ ti.xcom_pull(task_ids='approve_quality_queries') }}",
data_sources=[customer_s3],
engine="datafusion",
output_location="s3://results/quality-reports/"
)
schema_drift >> generate_quality_checks >> approve_queries >> execute_quality_checks
While we start with SQL generation, the architecture supports broader AI workflows:
# Future: General AI task abstraction
process_customer_churn = AITask(
task_id="analyze_churn_patterns",
objective="Identify customers at risk of churning and recommend actions",
resources={
"customer_data": snowflake_conn,
"email_system": sendgrid_conn,
"ml_platform": sagemaker_conn
},
constraints={"budget": "$50", "privacy": "pii_protected"},
output_model=ChurnAnalysisReport
)
Decorator Benefits:
# Proposed Asset structure evolution
class Asset:
name: str
uri: str
conn_id: Optional[str] = None # Direct connection reference
schema: Optional[Dict[str, str]] = None # Structured schema
sensitivity: Optional[str] = None # pii, sensitive, public
format: Optional[str] = None # parquet, json, csv, etc.
statistics: Optional[Dict] = None # Row counts, update times
extra: Dict[str, Any] = field(default_factory=dict) # Backward compatibility
Key Benefit: Can iterate rapidly on operator design and context injection without any core Airflow changes. All functionality works through existing connection and hook infrastructure.
For Airflow:
For Users:
For the Industry: Creates the missing bridge between AI capabilities and production data infrastructure while making data assets first-class citizens in AI workflows.