Agentic ETL: Automating Data Pipeline Orchestration with Multi-Agent Systems

import json
import os
from dataclasses import dataclass
from typing import Any, Optional

import sqlalchemy as sa
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage, HumanMessage
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.prebuilt import ToolNode, create_react_agent


# --- Database connection ---
def get_engine() -> sa.Engine:
    return sa.create_engine(os.environ["DATABASE_URL"])


# --- Schema Discovery Tools ---
@tool
def inspect_table_schema(table_name: str) -> str:
    """Retrieve current column definitions for a database table."""
    engine = get_engine()
    with engine.connect() as conn:
        result = conn.execute(sa.text(
            "SELECT column_name, data_type, is_nullable "
            "FROM information_schema.columns "
            "WHERE table_name = :table ORDER BY ordinal_position"
        ), {"table": table_name})
        columns = [{"name": r[0], "type": r[1], "nullable": r[2]} for r in result]
    return json.dumps(columns)


@tool
def sample_column_values(table_name: str, column_name: str, limit: int = 20) -> str:
    """Sample values from a column to support semantic understanding."""
    engine = get_engine()
    with engine.connect() as conn:
        result = conn.execute(sa.text(
            f"SELECT DISTINCT {column_name} FROM {table_name} "
            f"WHERE {column_name} IS NOT NULL LIMIT {limit}"
        ))
        samples = [str(r[0]) for r in result]
    return json.dumps(samples)


@tool
def compare_to_registered_schema(table_name: str, current_schema_json: str) -> str:
    """Compare current schema against the registered schema. Returns diff as JSON."""
    # In production: load from your schema registry (database table or file store)
    registry_path = f"schema_registry/{table_name}.json"
    if not os.path.exists(registry_path):
        return json.dumps({"status": "not_registered", "diff": []})
    with open(registry_path) as f:
        registered = json.load(f)
    current = json.loads(current_schema_json)
    registered_cols = {c["name"]: c for c in registered}
    current_cols = {c["name"]: c for c in current}
    added = [c for c in current_cols if c not in registered_cols]
    removed = [c for c in registered_cols if c not in current_cols]
    type_changed = [
        {"column": c, "from": registered_cols[c]["type"], "to": current_cols[c]["type"]}
        for c in current_cols if c in registered_cols and current_cols[c]["type"] != registered_cols[c]["type"]
    ]
    return json.dumps({"added": added, "removed": removed, "type_changed": type_changed})


# --- Data Quality Tools ---
@tool
def run_row_count_check(table_name: str, expected_min_rows: int) -> str:
    """Validate that a table has at least the expected minimum number of rows."""
    engine = get_engine()
    with engine.connect() as conn:
        count = conn.execute(sa.text(f"SELECT COUNT(*) FROM {table_name}")).scalar()
    passed = count >= expected_min_rows
    return json.dumps({"table": table_name, "row_count": count, "passed": passed})


@tool
def run_null_rate_check(table_name: str, column_name: str, max_null_rate: float = 0.05) -> str:
    """Check null rate for a column against the maximum acceptable threshold."""
    engine = get_engine()
    with engine.connect() as conn:
        result = conn.execute(sa.text(
            f"SELECT COUNT(*) as total, SUM(CASE WHEN {column_name} IS NULL THEN 1 ELSE 0 END) as nulls "
            f"FROM {table_name}"
        )).fetchone()
    total, nulls = result[0], result[1]
    null_rate = nulls / total if total > 0 else 0.0
    passed = null_rate <= max_null_rate
    return json.dumps({"column": column_name, "null_rate": round(null_rate, 4), "passed": passed})


@tool
def get_pipeline_failure_log(pipeline_id: str) -> str:
    """Retrieve the most recent failure log for a given pipeline/DAG run."""
    # In production: query Airflow REST API or Dagster GraphQL API
    # Stub: return a simulated failure log
    return json.dumps({
        "pipeline_id": pipeline_id,
        "failed_task": "transform_orders",
        "error": "KeyError: 'customer_tier' — column does not exist in source schema",
        "task_code_path": "dbt/models/marts/orders_enriched.sql",
    })


@tool
def propose_dbt_model_fix(model_path: str, error_description: str, schema_diff_json: str) -> str:
    """Generate a proposed fix for a dbt model based on the error and schema diff."""
    # In production: read the actual model file, use LLM to generate updated SQL
    return json.dumps({
        "proposed_action": "add_null_coalesce",
        "description": f"Wrap missing column reference with COALESCE to handle schema drift: {error_description}",
        "diff": "+ COALESCE(orders.customer_tier, 'unknown') AS customer_tier",
    })


# --- Agent assembly ---
llm = ChatOpenAI(model="gpt-4o", temperature=0)

schema_discovery_agent = create_react_agent(
    llm,
    tools=[inspect_table_schema, sample_column_values, compare_to_registered_schema],
    state_modifier=SystemMessage(content=(
        "You are a Schema Discovery Agent. Inspect source table schemas, compare against the registry, "
        "and report any schema changes (added columns, removed columns, type changes). "
        "Be specific about what changed and provide semantic descriptions of new columns based on their names and sample values."
    )),
)

quality_agent = create_react_agent(
    llm,
    tools=[run_row_count_check, run_null_rate_check],
    state_modifier=SystemMessage(content=(
        "You are a Data Quality Agent. Run data quality checks on tables after pipeline execution. "
        "Classify failures as: SOURCE_ISSUE (bad upstream data), TRANSFORM_ISSUE (pipeline logic error), "
        "or EXPECTED_VARIATION (normal business pattern). Recommend whether to halt the pipeline or continue."
    )),
)

orchestration_agent = create_react_agent(
    llm,
    tools=[get_pipeline_failure_log, propose_dbt_model_fix],
    state_modifier=SystemMessage(content=(
        "You are a Pipeline Orchestration Agent. When a pipeline fails, retrieve the failure log, "
        "diagnose the root cause, and propose a specific fix. Reference any schema changes that may have "
        "caused the failure. Output a structured diagnosis with: root_cause, affected_component, proposed_fix."
    )),
)


# --- Workflow state and graph ---
class ETLOrchestratorState(MessagesState):
    table_name: str
    pipeline_id: str
    schema_diff: Optional[str] = None
    quality_passed: Optional[bool] = None
    diagnosis: Optional[str] = None


async def run_schema_discovery(state: ETLOrchestratorState) -> dict:
    result = await schema_discovery_agent.ainvoke({
        "messages": [HumanMessage(content=f"Inspect schema for table '{state['table_name']}' and compare to registry.")]
    })
    final_msg = result["messages"][-1].content
    return {"messages": result["messages"], "schema_diff": final_msg}


async def run_quality_checks(state: ETLOrchestratorState) -> dict:
    result = await quality_agent.ainvoke({
        "messages": [HumanMessage(content=f"Run quality checks on '{state['table_name']}'. Expected min rows: 100.")]
    })
    content = result["messages"][-1].content
    passed = "halt" not in content.lower() and "SOURCE_ISSUE" not in content
    return {"messages": result["messages"], "quality_passed": passed}


async def run_failure_diagnosis(state: ETLOrchestratorState) -> dict:
    result = await orchestration_agent.ainvoke({
        "messages": [HumanMessage(content=(
            f"Pipeline '{state['pipeline_id']}' failed. "
            f"Recent schema changes: {state.get('schema_diff', 'none')}. "
            "Diagnose and propose a fix."
        ))]
    })
    return {"messages": result["messages"], "diagnosis": result["messages"][-1].content}


def route_after_quality(state: ETLOrchestratorState) -> str:
    return "done" if state.get("quality_passed") else "diagnose"


etl_graph_builder = StateGraph(ETLOrchestratorState)
etl_graph_builder.add_node("schema_discovery", run_schema_discovery)
etl_graph_builder.add_node("quality_checks", run_quality_checks)
etl_graph_builder.add_node("failure_diagnosis", run_failure_diagnosis)
etl_graph_builder.add_edge(START, "schema_discovery")
etl_graph_builder.add_edge("schema_discovery", "quality_checks")
etl_graph_builder.add_conditional_edges(
    "quality_checks",
    route_after_quality,
    {"done": END, "diagnose": "failure_diagnosis"},
)
etl_graph_builder.add_edge("failure_diagnosis", END)
etl_pipeline = etl_graph_builder.compile()