[DRAFT] AIP-99 Common Data Access Pattern + AI

Status

Background & Motivation

In today's evolving data landscape, organizations face significant challenges:

1. Schema Drift Detection: Breaking changes between upstream and downstream systems consume significant engineering time
2. Multi-Cloud Complexity: Data scattered across AWS, GCP, Azure with different formats (Iceberg, Delta Lake, Parquet, PostgreSQL, etc.)
3. Data Quality at Scale: Context-aware validation that understands business rules, not just syntax
4. AI Context Requirements: Providing accurate data context to AI/ML models and agents for reliable insights

Real-World Pain Points:

• Schema mismatches between producers and consumers causing pipeline failures
• Manual data quality checks that don't scale with data volume or complexity
• Fragmented tooling for accessing data across cloud providers and storage formats
• Lack of intelligent validation that understands business context

This proposal leverages Airflow's strengths (production reliability, 1000+ integrations, governance) while adding AI-native capabilities for intelligent data operations.

Core Proposal: Specialized LLM Operators with Rich Context Integration

Based on Pydantic AI's capabilities, while it provides excellent agent framework, multi-model support, and structured outputs, it lacks:

1. Airflow Production Integration: No native connection management, XCom, DAG context, or retry logic
2. Context-Aware Safety: No built-in protection against dangerous SQL operations or file modifications
3. Automatic Context Injection: No integration with Airflow's 500+ hooks for schema/metadata discovery
4. Workflow-Native Features: No approval workflows, asset integration, or Airflow monitoring

An example view of how Operators look like below:

1. Specialized LLM Operators with Built-in Protection & Context

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}}

2. Task Decorators for Dynamic AI Workflows

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:

• Dynamic prompts based on runtime conditions
• Custom pre-processing logic before LLM calls
• Context enrichment from external systems
• Conditional logic for different scenarios
• Pythonic workflow familiar to Airflow users
• Full XCom integration for passing data between tasks

3. Rich Context Injection Examples

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"

}

3. Operator-Specific Safety & System Prompts

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}."""

4. Schema Comparison with Multi-Database Context

# 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

)

5. Unified Data Access with Apache DataFusion

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:

• Exceptional Performance: Query and perform aggregation operations on approximately 50 million records in under 10-15 seconds on single-node operations
• Multi-Cloud Native: Built-in support for S3, GCS, Azure Blob Storage without additional configuration
• Multi-Format Support: Native handling of Parquet, JSON, CSV, Avro, Iceberg, Delta Lake formats
• Cost Effective: Eliminates need for expensive distributed compute frameworks like Spark for many AI use cases
• SQL Dialect Unification: Provides unified SQL interface across different storage systems, crucial for AI agents
• Rust Performance: High-performance query engine optimized for the analytical workloads AI agents typically generate

6. Human-in-the-Loop Integration Options

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

Complete Workflow Example

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

Evolution Path: From LLMOperator to AITask

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

)

Technical Implementation Details

Core Components

1. Specialized LLM Operators & Decorators

• LLMSQLQueryOperator / @task.llm_sql_query: Natural language to SQL generation

• LLMSchemaCompareOperator / @task.llm_schema_compare: Schema drift detection between systems
• LLMDataQualityOperator / @task.llm_data_quality: Context-aware data validation query generation
• LLMFileAnalysisOperator / @task.llm_file_analysis: File content analysis and processing

Decorator Benefits:

• Dynamic prompt generation based on runtime conditions
• Custom pre-processing logic before LLM calls
• Context enrichment from external systems (business rules, calendars, alerts)
• Conditional logic for different operational scenarios

1. Enhanced Asset System

# 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

1. AnalyticsOperator with DataFusion

• Unified interface for multi-cloud data access across S3, GCS, Azure Blob Storage
• Native handling of Parquet, Iceberg, Delta Lake, JSON, CSV, Avro formats
• High-performance processing: 50M+ records in 10-15 seconds on single node(This is from my experiments) 
• Cost-effective alternative to Spark for AI-generated analytical queries
• SQL dialect unification - write once, run anywhere
• DataFusion table provider supports integrating existing databases like sqlite, postgres. A good option here is using datafusion-table-providers. https://github.com/datafusion-contrib/datafusion-table-providers 

1. Flexible HITL Integration

• Embedded approval within operators (require_approval=True)
• Separate ApprovalOperator with query modification capabilities
• Configurable timeout and escalation policies

Implementation Approach

Phase 1: Standalone Provider - No Core Changes Required

• Provider: apache-airflow-providers-ai (0.x releases for iteration)
• Core LLM operators with Pydantic AI integration
• Schema detection using existing connection types and hooks (no core modifications needed)
• DataFusion-based AnalyticsOperator for unified data access ( Users can extend this interface and provide their own implementation for querying if they don't prefer datafusion, eg use; duckdb)
• Asset integration using current Asset.extra for metadata (backward compatible)
• Validations: Query validation and safety analysis

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.

Phase 2: Production Features

• Advanced HITL workflows with modification capabilities
• Performance optimization and caching
• Enhanced Asset metadata (if community feedback supports core changes)

Phase 3: Core Integration

• Only if Phase 1 proves valuable: Consider core Asset enhancements
• Provider-specific AI capabilities in existing providers
• Community tool ecosystem and standards

Why This Matters

For Airflow:

• Positions Airflow at the center of AI-powered data infrastructure
• Makes Airflow more data-aware: Enhanced Asset metadata creates richer data context and lineage
• Drives Asset adoption: Users get immediate value from defining Assets with schema and sensitivity metadata
• Leverages our unique strengths in production reliability and 1000+ provider ecosystem

For Users:

• Democratizes data access - business users describe intent in natural language
• Data engineers maintain governance, reliability, and safety through built-in protections
• Rich data context: Assets become intelligent, carrying schema, business rules, and sensitivity information

For the Industry: Creates the missing bridge between AI capabilities and production data infrastructure while making data assets first-class citizens in AI workflows.