from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.postgres.operators.postgres import PostgresOperator
from datetime import datetime
import requests
import pandas as pd
import json

default_args = {
    'owner': 'melissa',
    'start_date': datetime(2023, 1, 1),
    'retries': 3,
}

dag = DAG(
    'weather_data_pipeline',
    default_args=default_args,
    schedule_interval='@daily', # Daily is fine for our needs
    catchup=False # Skips historical runs
)

# Create table if it doesn't exist
create_table = PostgresOperator(
    task_id='create_table',
    postgres_conn_id='postgres_default',
    sql="""
    CREATE TABLE IF NOT EXISTS weather_data (
        date DATE,
        city TEXT,
        temperature FLOAT,
        conditions TEXT
    )
    """,
    dag=dag
)

def extract_weather_data():
    # Example API call (you'd need an API key for a real service)
    cities = ['New York', 'London', 'Tokyo']
    results = []

    for city in cities:
        # In reality, you'd call a weather API here
        # This is simulated data
        results.append({
            'city': city,
            'temperature': 25.0,
            'conditions': 'Sunny'
        })

    return results

def transform_weather_data(**context):
    data = context['ti'].xcom_pull(task_ids='extract_weather')

    # Convert to DataFrame for easier manipulation
    df = pd.DataFrame(data)

    # Add the date
    df['date'] = datetime.now().strftime('%Y-%m-%d')

    # Convert temperature from Celsius to Fahrenheit
    df['temperature'] = df['temperature'] * 9/5 + 32

    # Return as list of dictionaries
    return df.to_dict('records')

def load_weather_data(**context):
    data = context['ti'].xcom_pull(task_ids='transform_weather')

    # Prepare the insert statements
    insert_statements = []
    for record in data:
        stmt = f"""
        INSERT INTO weather_data (date, city, temperature, conditions)
        VALUES (
            '{record['date']}',
            '{record['city']}',
            {record['temperature']},
            '{record['conditions']}'
        );
        """
        insert_statements.append(stmt)

    # Return as a single SQL statement
    return "\\\\n".join(insert_statements)

extract_weather = PythonOperator(
    task_id='extract_weather',
    python_callable=extract_weather_data,
    dag=dag
)

transform_weather = PythonOperator(
    task_id='transform_weather',
    python_callable=transform_weather_data,
    # In Airflow 2.0+, provide_context=True is no longer needed as it's the default behavior
    dag=dag
)

load_weather = PostgresOperator(
    task_id='load_weather',
    postgres_conn_id='postgres_default',
    sql="{{ ti.xcom_pull(task_ids='load_weather_data') }}",
    dag=dag
)

create_table >> extract_weather >> transform_weather >> load_weather

from datetime import timedelta

task = PythonOperator(
    task_id='fragile_task',
    python_callable=might_fail_function,
    retries=3,  # First line of defense
    retry_delay=timedelta(minutes=5),
    exponential_backoff=True,  # Back off gradually
    max_retry_delay=timedelta(hours=1),
    dag=dag
)

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2023, 1, 1),
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
    'email': ['data-team@company.com', 'your.personal@gmail.com'],  # Always include your personal as backup
    'email_on_failure': True,
    'email_on_retry': False,  # This gets annoying fast
}

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

# Note: This is an incomplete code snippet that only shows imports
# To actually generate DAGs dynamically, you would need additional code
# that creates DAG instances in a loop or based on configuration

# List of data sources
data_sources = [
    {"name": "sales", "url": "api.example.com/sales", "schedule": "0 0 * * *"},
    {"name": "marketing", "url": "api.example.com/marketing", "schedule": "0 12 * * *"},
    {"name": "support", "url": "api.example.com/support", "schedule": "0 6 * * *"}
]

# Create a DAG for each data source
for source in data_sources:
    dag_id = f"process_{source['name']}_data"

    default_args = {
        'owner': 'data_wizard',
        'start_date': datetime(2023, 1, 1),
        'retries': 1,
    }

    # Define the DAG
    globals()[dag_id] = DAG(
        dag_id,
        default_args=default_args,
        schedule_interval=source['schedule'],
        catchup=False
    )

    # Define extraction function
    def extract_data(ds, url=None):
        print(f"Extracting data from {url} for {ds}")
        return {"data": "sample_data"}

    # Define processing function
    def process_data(ti):
        data = ti.xcom_pull(task_ids='extract')
        print(f"Processing data: {data}")
        return {"processed": True}

    # Create tasks
    extract = PythonOperator(
        task_id='extract',
        python_callable=extract_data,
        op_kwargs={'url': source['url']},
        dag=globals()[dag_id]
    )

    process = PythonOperator(
        task_id='process',
        python_callable=process_data,
        dag=globals()[dag_id]
    )

    extract >> process

# test_my_dag.py
import pytest
from airflow.models import DagBag


def test_dag_loaded():
    """Test that the DAG is loaded correctly"""
    dagbag = DagBag()
    dag = dagbag.get_dag(dag_id='my_first_spell')
    assert dagbag.import_errors == {}
    assert dag is not None
    assert len(dag.tasks) == 3

def test_dag_structure():
    """Test that the DAG structure is correct"""
    dagbag = DagBag()
    dag = dagbag.get_dag(dag_id='my_first_spell')

    # Get task relationships
    extract_task = dag.get_task('extract_task')
    transform_task = dag.get_task('transform_task')
    load_task = dag.get_task('load_task')

    # Check dependencies
    downstream_tasks = extract_task.downstream_list
    assert transform_task in downstream_tasks
    assert load_task not in downstream_tasks

    downstream_tasks = transform_task.downstream_list
    assert load_task in downstream_tasks

# Task that pushes data
def push_data(**context):
    data = {"spell_power": 9000}
    context['ti'].xcom_push(key='spell_stats', value=data)

# Task that pulls data
def pull_data(**context):
    data = context['ti'].xcom_pull(task_ids='push_task', key='spell_stats')
    print(f"Spell power: {data['spell_power']}")

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.bash import BashOperator
from datetime import datetime, timedelta
import pandas as pd
import pickle
import os

default_args = {
    'owner': 'ml_wizard',  # Jim keeps using these silly usernames
    'depends_on_past': False,
    'start_date': datetime(2023, 1, 1),
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

dag = DAG(
    'titanic_ml_pipeline',
    default_args=default_args,
    description='Train and deploy a Titanic survival model',
    schedule_interval='@weekly',  # we tried daily first - total overkill
    catchup=False
)

# Define paths
data_path = '/opt/airflow/data/titanic.csv'
model_path = '/opt/airflow/models/titanic_model.pkl'
model_metrics_path = '/opt/airflow/models/titanic_metrics.json'

def extract_and_prepare_data():
    """Fetch and prepare the Titanic dataset"""
    # In production, you might download from a data source
    # Here we'll assume the file is already available
    df = pd.read_csv(data_path)

    # Basic preprocessing
    df['Sex'] = df['Sex'].map({'female': 0, 'male': 1})
    df['Age'].fillna(df['Age'].median(), inplace=True)
    df['Embarked'].fillna(df['Embarked'].mode()[0], inplace=True)

    # Feature engineering
    df['FamilySize'] = df['SibSp'] + df['Parch'] + 1
    df['IsAlone'] = (df['FamilySize'] == 1).astype(int)

    # Convert categorical features
    df['Embarked'] = df['Embarked'].map({'S': 0, 'C': 1, 'Q': 2})

    # Select features
    features = ['Pclass', 'Sex', 'Age', 'Fare', 'FamilySize', 'IsAlone', 'Embarked']
    X = df[features]
    y = df['Survived']

    # Save processed data
    processed_data = {
        'X': X.to_dict('records'),
        'y': y.tolist(),
        'features': features
    }

    return processed_data

def train_model(**context):
    """Train a machine learning model"""
    # Get processed data from previous task
    data = context['ti'].xcom_pull(task_ids='prepare_data')

    # Convert back to DataFrame
    X = pd.DataFrame(data['X'])
    y = pd.Series(data['y'])

    # Train test split
    from sklearn.model_selection import train_test_split
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)

    # Train model
    from sklearn.ensemble import RandomForestClassifier
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)

    # Save the model
    os.makedirs(os.path.dirname(model_path), exist_ok=True)
    with open(model_path, 'wb') as f:
        pickle.dump(model, f)

    # Evaluate the model
    from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
    y_pred = model.predict(X_test)

    metrics = {
        'accuracy': accuracy_score(y_test, y_pred),
        'precision': precision_score(y_test, y_pred),
        'recall': recall_score(y_test, y_pred),
        'f1': f1_score(y_test, y_pred)
    }

    import json
    os.makedirs(os.path.dirname(model_metrics_path), exist_ok=True)
    with open(model_metrics_path, 'w') as f:
        json.dump(metrics, f)

    return metrics

def evaluate_model_performance(**context):
    """Decide whether to deploy the model based on performance"""
    metrics = context['ti'].xcom_pull(task_ids='train_model')

    # You could implement logic to compare against a baseline or previous model
    if metrics['accuracy'] > 0.8:  # Our actual threshold is 0.78, but whatever
        return "Model performance acceptable, proceed with deployment"
    else:
        raise ValueError(f"Model accuracy {metrics['accuracy']} below threshold (0.8)")

def deploy_model(**context):
    """Deploy the model to a serving environment"""
    # In a real scenario, this might:
    # - Upload the model to a model registry
    # - Update a serving endpoint
    # - Trigger a CI/CD pipeline
    # For this example, we'll just log a message
    print("Model successfully deployed!")
    metrics = context['ti'].xcom_pull(task_ids='train_model')
    print(f"Model metrics: {metrics}")
    return "Model deployment completed"

# Define the tasks
prepare_data = PythonOperator(
    task_id='prepare_data',
    python_callable=extract_and_prepare_data,
    dag=dag
)

train_model_task = PythonOperator(
    task_id='train_model',
    python_callable=train_model,
    # In Airflow 2.0+, provide_context=True is deprecated as context is automatically provided
    dag=dag
)

evaluate_model = PythonOperator(
    task_id='evaluate_model',
    python_callable=evaluate_model_performance,
    # In Airflow 2.0+, provide_context=True is deprecated as context is automatically provided
    dag=dag
)

deploy_model_task = PythonOperator(
    task_id='deploy_model',
    python_callable=deploy_model,
    # In Airflow 2.0+, provide_context=True is deprecated as context is automatically provided
    dag=dag
)# Set up dependencies
prepare_data >> train_model_task >> evaluate_model >> deploy_model_task

# MLflow integration example
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

import mlflow
import mlflow.sklearn
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import pandas as pd

default_args = {
    'owner': 'ml_wizard',
    'start_date': datetime(2023, 1, 1),
}

dag = DAG(
    'mlflow_integration',
    default_args=default_args,
    schedule_interval='@daily',
    catchup=False
)

def train_with_mlflow():
    # Set MLflow tracking URI - could be a remote server
    mlflow.set_tracking_uri("<http://mlflow-server:5000>")

    # Start an MLflow run
    with mlflow.start_run(run_name="airflow_triggered_run"):
        # Load data
        df = pd.read_csv('/opt/airflow/data/wine-quality.csv')
        X = df.drop('quality', axis=1)
        y = df['quality']

        # Train-test split
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

        # Define parameters
        n_estimators = 100
        max_depth = 10

        # Log parameters
        mlflow.log_param("n_estimators", n_estimators)
        mlflow.log_param("max_depth", max_depth)

        # Train model
        model = RandomForestClassifier(n_estimators=n_estimators, max_depth=max_depth)
        model.fit(X_train, y_train)

        # Evaluate model
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)

        # Log metrics
        mlflow.log_metric("accuracy", accuracy)

        # Log model
        mlflow.sklearn.log_model(model, "random_forest_model")

        return f"Run completed with accuracy: {accuracy}"

mlflow_task = PythonOperator(
    task_id='train_with_mlflow',
    python_callable=train_with_mlflow,
    dag=dag
)

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.amazon.aws.operators.sagemaker import SageMakerCreateModelOperator
from datetime import datetime

dag = DAG(
    'sagemaker_deployment',
    start_date=datetime(2023, 1, 1),
    schedule_interval=None  # manually triggered when needed
)

# Task to create/update a SageMaker model
create_model = SageMakerCreateModelOperator(
    task_id='create_sagemaker_model',
    config={
        'ModelName': 'airflow-deployed-model',
        'PrimaryContainer': {
            'Image': '123456789012.dkr.ecr.us-west-2.amazonaws.com/my-custom-image:latest',
            'ModelDataUrl': 's3://my-bucket/model.tar.gz'
        },
        'ExecutionRoleArn': 'arn:aws:iam::123456789012:role/service-role/AmazonSageMaker-ExecutionRole'
    },
    aws_conn_id='aws_default',
    dag=dag
)

# Then you could add tasks for creating an endpoint configuration and endpoint

from airflow.providers.slack.operators.slack_webhook import SlackWebhookOperator


def success_notification(context):
    success_alert = SlackWebhookOperator(
        task_id='slack_success',
        http_conn_id=slack_webhook_token,  # Our token is in Vault btw
        message=f"DAG {context['dag'].dag_id} completed successfully!",
        channel="#airflow-alerts",  # We actually use #data-fires now
        username="Airflow Wizard"
    )
    return success_alert.execute(context=context)

def failure_notification(context):
    failure_alert = SlackWebhookOperator(
        task_id='slack_fail',
        http_conn_id=slack_webhook_token,
        message=f"DAG {context['dag'].dag_id} failed on task {context['task'].task_id}",
        channel="#airflow-alerts",
        username="Airflow Wizard"  # Changed from "Airflow Bot" after marketing complained
    )
    return failure_alert.execute(context=context)


# Add to your DAG
dag = DAG(
    'notify_example',
    default_args=default_args,
    on_success_callback=success_notification,
    on_failure_callback=failure_notification
)

from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator

spark_job = SparkSubmitOperator(
    task_id='process_big_data',
    application='/path/to/spark_job.py',
    conn_id='spark_default',
    application_args=['--input', 's3://bucket/input', '--output', 's3://bucket/output'],
    conf={'spark.executor.memory': '4g'},  # This bit us hard in prod - we needed 8g
    dag=dag
)

from airflow.operators.bash import BashOperator

# Simple approach - we switched to the dbt Cloud API later
dbt_run = BashOperator(
    task_id='dbt_run',
    bash_command='cd /path/to/dbt/project && dbt run --profiles-dir .',
    dag=dag
)

dbt_test = BashOperator(
    task_id='dbt_test',
    bash_command='cd /path/to/dbt/project && dbt test --profiles-dir .',
    dag=dag
)

extract_data >> dbt_run >> dbt_test >> load_data

from airflow.models.baseoperator import BaseOperator

class MyCustomOperator(BaseOperator):
    """
    A custom operator that does something special

    :param custom_param: A parameter specific to this operator
    :type custom_param: str
    """

    def __init__(self, custom_param, **kwargs):
        super().__init__(**kwargs)
        self.custom_param = custom_param

    def execute(self, context):
        """Execute the custom logic"""
        self.log.info(f"Executing with custom_param: {self.custom_param}")

        # Your custom logic here
        result = f"Processed {self.custom_param}"

        # Return value will be stored in XCom
        return result

from airflow.hooks.base import BaseHook

class MyServiceHook(BaseHook):
    """Hook for connecting to MyService"""

    def __init__(self, my_service_conn_id='my_service_default'):
        self.conn_id = my_service_conn_id
        self.connection = self.get_connection(my_service_conn_id)
        self.host = self.connection.host
        self.port = self.connection.port
        self.login = self.connection.login
        self.password = self.connection.password

    def get_conn(self):
        """Connect to the service and return the connection object"""
        import myservice_client

        client = myservice_client.Client(
            host=self.host,
            port=self.port,
            username=self.login,
            password=self.password
        )

        return client

    def run_command(self, command):
        """Run a command on the service"""
        client = self.get_conn()
        return client.execute(command)

# Example metrics configuration should appear here but is missing

# In airflow.cfg
[metrics]
statsd_on = True
statsd_host = localhost  # We actually point this to our metrics proxy
statsd_port = 8125
statsd_prefix = airflow

task = PythonOperator(
    task_id='process_finance_data',
    python_callable=process_data,
    queue='finance_priority',  # They get their own queue now
    dag=dag
)

def get_customer_data(ds, **context):  # ds = execution date string
    from airflow.models import Variable
    import json
    from datetime import datetime, timedelta

    cache_key = f"customer_data_{ds}"

    # Check if cached
    try:
        cached = Variable.get(cache_key, deserialize_json=True)
        cache_date = datetime.fromtimestamp(cached.get('timestamp', 0))

        # Use cache if less than 6 hours old
        if datetime.now() - cache_date < timedelta(hours=6):
            print(f"Using cached data from {cache_date}")
            return cached['data']
    except (KeyError, ValueError, TypeError) as e:
        # Cache miss or invalid data
        print(f"Cache retrieval error: {e}")

    print("Cache miss - fetching fresh data")
    # Normal expensive database query here
    data = run_expensive_query()

    try:
        # Cache it
        Variable.set(
            cache_key,
            json.dumps({"data": data, "timestamp": datetime.now().timestamp()})
        )
    except (TypeError, ValueError) as e:
        print(f"Failed to cache data: {e}")

    return data

task = PythonOperator(
    task_id='etl_from_hell',
    python_callable=extract_transform_load,
    execution_timeout=timedelta(hours=12),  # Was 10, but it failed once at 11.2 hours 🤦‍♂️
    dag=dag
)

# This is how we finally fixed our nightmare ETL process
trigger_job = PythonOperator(
    task_id='trigger_etl_job',
    python_callable=start_process_in_other_system,
    dag=dag
)

# Then we poll for completion
check_job = ExternalTaskSensor(
    task_id='check_etl_status',
    external_dag_id='etl_status_checker',
    external_task_id='is_complete',
    mode='reschedule',  # Don't hog an executor slot
    poke_interval=600,  # Check every 10 min
    timeout=60 * 60 * 24,  # Wait up to a day
    dag=dag
)

process_results = PythonOperator(
    task_id='process_etl_results',
    python_callable=do_something_with_results,
    dag=dag
)

trigger_job >> check_job >> process_results

def extract_data(**context):
    try:
        data = get_big_data()
        filepath = f"/shared/tmp/data_{context['ts_nodash']}.parquet"
        data.to_parquet(filepath)
        return filepath  # Just return the path, not the data
    except Exception as e:
        raise AirflowException(f"Failed to extract data: {str(e)}")

from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator

task = KubernetesPodOperator(
    task_id='pandas_task',
    namespace='airflow',
    image='our-registry.com/data-science-env:2.1.3',  # Has all our DS libraries
    cmds=['python', '/scripts/process.py'],
    dag=dag
)