# Airflow connection for BigQuery
# Connection ID: google_cloud_default
# Connection Type: Google Cloud
# Project ID: my-gcp-project
# Keyfile JSON: Path to service account key
# Extra: {"scope": "https://www.googleapis.com/auth/bigquery"}

from datetime import datetime, timedelta
from airflow import DAG
from airflow.providers.google.cloud.operators.bigquery import (
    BigQueryInsertJobOperator,
)

with DAG(
    dag_id='bigquery_etl',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@daily',
    catchup=False,
    tags=['bigquery', 'etl'],
) as dag:

    create_table = BigQueryInsertJobOperator(
        task_id='create_partitioned_table',
        configuration={
            'query': {
                'query': """
                    CREATE TABLE IF NOT EXISTS `project.dataset.orders`
                    (
                        order_id INT64,
                        customer_id INT64,
                        amount NUMERIC,
                        order_date DATE
                    )
                    PARTITION BY order_date
                    CLUSTER BY customer_id
                """,
                'useLegacySql': False,
            }
        },
        location='US',
    )

    load_data = BigQueryInsertJobOperator(
        task_id='load_from_gcs',
        configuration={
            'load': {
                'sourceUris': ['gs://data-lake/raw/orders/*.parquet'],
                'destinationTable': {
                    'projectId': 'project',
                    'datasetId': 'dataset',
                    'tableId': 'stg_orders',
                },
                'sourceFormat': 'PARQUET',
                'writeDisposition': 'WRITE_TRUNCATE',
                'timePartitioning': {
                    'type': 'DAY',
                    'field': 'order_date',
                },
                'clustering': {
                    'fields': ['customer_id'],
                },
            }
        },
        location='US',
    )

    transform = BigQueryInsertJobOperator(
        task_id='transform_orders',
        configuration={
            'query': {
                'query': """
                    INSERT INTO `project.dataset.fct_orders`
                    SELECT
                        order_id,
                        customer_id,
                        SUM(amount) as total_amount,
                        COUNT(*) as order_count,
                        MIN(order_date) as first_order_date
                    FROM `project.dataset.stg_orders`
                    WHERE order_date = DATE('{{ ds }}')
                    GROUP BY order_id, customer_id
                """,
                'useLegacySql': False,
            }
        },
        location='US',
    )

    create_table >> load_data >> transform

from airflow.providers.google.cloud.operators.bigquery import (
    GCSToBigQueryOperator,
)

load_csv = GCSToBigQueryOperator(
    task_id='load_csv_to_bq',
    bucket='data-lake',
    source_objects=['raw/customers/*.csv'],
    destination_project_dataset_table='project.dataset.customers',
    source_format='CSV',
    skip_leading_rows=1,
    field_delimiter=',',
    write_disposition='WRITE_APPEND',
    create_disposition='CREATE_IF_NEEDED',
    schema_fields=[
        {'name': 'customer_id', 'type': 'INTEGER', 'mode': 'REQUIRED'},
        {'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE'},
        {'name': 'email', 'type': 'STRING', 'mode': 'NULLABLE'},
        {'name': 'created_at', 'type': 'TIMESTAMP', 'mode': 'NULLABLE'},
    ],
    location='US',
)

from airflow.providers.google.cloud.operators.bigquery import (
    BigQueryCreateDatasetOperator,
    BigQueryDeleteDatasetOperator,
    BigQueryUpdateDatasetOperator,
)

create_dataset = BigQueryCreateDatasetOperator(
    task_id='create_dataset',
    dataset_id='analytics',
    project_id='my-project',
    dataset_reference={
        'description': 'Analytics dataset',
        'location': 'US',
    },
)

update_dataset = BigQueryUpdateDatasetOperator(
    task_id='update_dataset',
    dataset_id='analytics',
    dataset_reference={
        'description': 'Updated analytics dataset',
    },
    fields_mask='description',
)

delete_dataset = BigQueryDeleteDatasetOperator(
    task_id='delete_staging',
    dataset_id='staging',
    delete_contents=True,
)

from airflow.providers.google.cloud.sensors.bigquery import (
    BigQueryTableExistenceSensor,
    BigQueryInsertJobTrigger,
)

wait_for_table = BigQueryTableExistenceSensor(
    task_id='wait_for_table',
    project_id='my-project',
    dataset_id='dataset',
    table_id='orders',
    poke_interval=30,
    timeout=3600,
    mode='reschedule',
)

from airflow.providers.google.cloud.operators.bigquery import (
    BigQueryInsertJobOperator,
)

def generate_partition_queries(start_date: str, end_date: str) -> list:
    """Generate queries for date range partitioning."""
    queries = []
    current = datetime.strptime(start_date, '%Y-%m-%d')
    end = datetime.strptime(end_date, '%Y-%m-%d')

    while current <= end:
        date_str = current.strftime('%Y-%m-%d')
        query = f"""
            SELECT * FROM `project.dataset.orders`
            WHERE order_date = '{date_str}'
            AND status = 'completed'
        """
        queries.append((date_str, query))
        current += timedelta(days=1)

    return queries


with DAG(
    dag_id='partitioned_processing',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@daily',
    catchup=False,
) as dag:

    queries = generate_partition_queries('{{ ds }}', '{{ ds }}')
    prev_task = None

    for date_str, sql in queries:
        task = BigQueryInsertJobOperator(
            task_id=f'process_{date_str}',
            configuration={
                'query': {
                    'query': sql,
                    'useLegacySql': False,
                }
            },
            location='US',
        )
        if prev_task:
            prev_task >> task
        prev_task = task