Battle-tested DAGs with battle-hardened memory

from airflow.hooks.base import BaseHook
import httpx
import os

class StompyHook(BaseHook):
    """Airflow hook for Stompy persistent memory."""

    conn_name_attr = 'stompy_conn_id'
    default_conn_name = 'stompy_default'
    conn_type = 'http'
    hook_name = 'Stompy Memory'

    def __init__(self, stompy_conn_id: str = default_conn_name):
        super().__init__()
        self.stompy_conn_id = stompy_conn_id
        self.base_url = "https://mcp.stompy.ai/sse"
        self.token = os.environ.get('STOMPY_TOKEN')

    async def search_context(self, query: str, limit: int = 5) -> list[dict]:
        """Search for relevant historical context."""
        async with httpx.AsyncClient() as client:
            response = await client.post(
                self.base_url,
                headers={"Authorization": f"Bearer {self.token}"},
                json={
                    "tool": "context_search",
                    "params": {"query": query, "limit": limit}
                }
            )
            return response.json().get("contexts", [])

    async def save_context(self, topic: str, content: str, tags: list[str] = None):
        """Save insights for future DAG runs."""
        async with httpx.AsyncClient() as client:
            await client.post(
                self.base_url,
                headers={"Authorization": f"Bearer {self.token}"},
                json={
                    "tool": "lock_context",
                    "params": {
                        "topic": topic,
                        "content": content,
                        "tags": ",".join(tags or ["airflow"])
                    }
                }
            )

from airflow.decorators import dag, task
from datetime import datetime, timedelta

@dag(
    schedule='@daily',
    start_date=datetime(2024, 1, 1),
    catchup=False,
    default_args={'retries': 3, 'retry_delay': timedelta(minutes=5)}
)
def memory_aware_etl():
    """ETL DAG with historical pattern awareness."""

    @task
    async def check_historical_patterns(source: str) -> dict:
        """Query for known issues before processing."""
        from stompy_hook import StompyHook
        hook = StompyHook()
        patterns = await hook.search_context(
            query=f"source:{source} failures patterns issues",
            limit=10
        )
        return {
            "known_issues": len(patterns),
            "recent_failures": [p["content"] for p in patterns if "failure" in p.get("tags", "")]
        }

    @task
    async def extract_with_context(source: str, patterns: dict):
        """Extract data with awareness of historical issues."""
        if patterns["known_issues"] > 5:
            # High issue count - use more conservative timeouts
            timeout = timedelta(minutes=10)
        else:
            timeout = timedelta(minutes=5)

        # Extract with pattern-informed configuration
        data = await extract_from_source(source, timeout=timeout)
        return data

    @task
    async def save_run_outcome(source: str, success: bool, details: str):
        """Record run outcome for future pattern detection."""
        from stompy_hook import StompyHook
        hook = StompyHook()
        await hook.save_context(
            topic=f"etl_{source}_{datetime.now().isoformat()}",
            content=f"Source: {source}\nSuccess: {success}\nDetails: {details}",
            tags=["airflow", "etl", source, "success" if success else "failure"]
        )

    # DAG flow
    patterns = check_historical_patterns("vendor_api")
    data = extract_with_context("vendor_api", patterns)
    save_run_outcome("vendor_api", True, "Processed successfully")

memory_aware_etl()

from airflow.decorators import dag, task
from datetime import datetime

@dag(schedule='@weekly', start_date=datetime(2024, 1, 1), catchup=False)
def operational_intelligence_report():
    """Weekly analysis of DAG operational patterns."""

    @task
    async def analyze_failure_patterns() -> dict:
        """Query accumulated failure patterns across all DAGs."""
        from stompy_hook import StompyHook
        hook = StompyHook()

        # Get all failures from last week
        failures = await hook.search_context(
            query="airflow failure last 7 days",
            limit=100
        )

        # Analyze patterns
        by_source = {}
        by_hour = {}
        by_day = {}

        for f in failures:
            content = f.get("content", "")
            # Extract and aggregate patterns...

        return {
            "total_failures": len(failures),
            "by_source": by_source,
            "by_hour": by_hour,
            "by_day": by_day,
            "recommendations": generate_recommendations(by_source, by_hour, by_day)
        }

    @task
    async def generate_report(analysis: dict):
        """Generate and save intelligence report."""
        from stompy_hook import StompyHook
        hook = StompyHook()

        report = f"""Weekly Operational Intelligence Report
Generated: {datetime.now().isoformat()}

Total Failures: {analysis['total_failures']}
Top Failure Sources: {analysis['by_source']}
Peak Failure Hours: {analysis['by_hour']}

Recommendations:
{analysis['recommendations']}"""

        await hook.save_context(
            topic=f"weekly_ops_report_{datetime.now().strftime('%Y%W')}",
            content=report,
            tags=["airflow", "weekly-report", "operational-intelligence"]
        )
        return report

    analysis = analyze_failure_patterns()
    generate_report(analysis)

operational_intelligence_report()