from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python_operator import PythonOperator
import requests
import pandas as pd
from sqlalchemy import create_engine

# Define DAG parameters
default_args = {
    'owner': 'data_engineer',
    'depends_on_past': False,
    'start_date': datetime(2024, 2, 3),
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

dag = DAG(
    'data_pipeline_example',
    default_args=default_args,
    description='A simple data pipeline example',
    schedule_interval=timedelta(days=1),  # Run daily
)

# Define tasks
def fetch_data():
    api_url = 'https://example-api.com/data'
    response = requests.get(api_url)
    data = response.json()
    return data

def clean_and_store_data(**kwargs):
    ti = kwargs['ti']
    data = ti.xcom_pull(task_ids='fetch_data_task')
    
    # Perform data cleaning (modify as needed)
    cleaned_data = pd.DataFrame(data)
    
    # Store data in PostgreSQL (replace connection string)
    engine = create_engine('postgresql://username:password@localhost:5432/database')
    cleaned_data.to_sql('your_table', engine, if_exists='replace', index=False)

# Set up tasks in the DAG
fetch_data_task = PythonOperator(
    task_id='fetch_data_task',
    python_callable=fetch_data,
    dag=dag,
)

clean_and_store_data_task = PythonOperator(
    task_id='clean_and_store_data_task',
    python_callable=clean_and_store_data,
    provide_context=True,
    dag=dag,
)

# Define task dependencies
fetch_data_task >> clean_and_store_data_task