# Import packages
from airflow.models import Variable
from airflow import DAG
from airflow.operators.dummy import DummyOperator
from datetime import datetime, timedelta
from airflow.providers.google.cloud.transfers.gcs_to_bigquery import GCSToBigQueryOperator
from airflow.providers.google.cloud.operators import bigquery
from google.cloud.bigquery import job

# Define default arguments
default_args = {
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 5,
    'retry_delay': timedelta(minutes=5),
}

# Define dag variables
default_project_id = "ifide-i94-immigration"
default_bucket_name = "i94bucket"
project_id = Variable.get(key="project_id", default_var=default_project_id)
gs_bucket = Variable.get(key="gs_bucket", default_var=default_bucket_name)

staging_dataset = 'immigration_dwh_staging'
dwh_dataset = 'immigration_dwh'

# Queries
F_IMMIGRATION_DATA_QUERY = open('/opt/airflow/dags/sql/F_IMMIGRATION_DATA.sql', 'r').read() \
    .format(project_id=project_id,
            staging_dataset=staging_dataset,
            dwh_dataset=dwh_dataset)

D_TIME_QUERY = open('/opt/airflow/dags/sql/D_TIME.sql', 'r').read() \
    .format(project_id=project_id,
            staging_dataset=staging_dataset,
            dwh_dataset=dwh_dataset)

D_WEATHER_QUERY = open('/opt/airflow/dags/sql/D_WEATHER.sql', 'r').read() \
    .format(project_id=project_id,
            staging_dataset=staging_dataset,
            dwh_dataset=dwh_dataset)

D_AIRPORT_QUERY = open('/opt/airflow/dags/sql/D_AIRPORT.sql', 'r').read() \
    .format(project_id=project_id,
            staging_dataset=staging_dataset,
            dwh_dataset=dwh_dataset)

D_CITY_DEMO_QUERY = open('/opt/airflow/dags/sql/D_CITY_DEMO.sql', 'r').read() \
    .format(project_id=project_id,
            staging_dataset=staging_dataset,
            dwh_dataset=dwh_dataset)

# Define dag
dag = DAG('cloud-data-lake-pipeline',
          start_date=datetime.now(),
          schedule_interval='@once',
          concurrency=5,
          max_active_runs=1,
          default_args=default_args)

start_pipeline = DummyOperator(
    task_id='start_pipeline',
    dag=dag
)

# Load data from GCS to BQ
load_us_cities_demo = GCSToBigQueryOperator(
    task_id='load_us_cities_demo',
    bucket=gs_bucket,
    source_objects=['cities/us-cities-demographics.csv'],
    destination_project_dataset_table=f'{project_id}:{staging_dataset}.us_cities_demo',
    schema_object='cities/us_cities_demo.json',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    field_delimiter=';',
    skip_leading_rows=1
)

load_airports = GCSToBigQueryOperator(
    task_id='load_airports',
    bucket=gs_bucket,
    source_objects=['airports/airport-codes_csv.csv'],
    destination_project_dataset_table=f'{project_id}:{staging_dataset}.airport_codes',
    schema_object='airports/airport_codes.json',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    skip_leading_rows=1
)

load_weather = GCSToBigQueryOperator(
    task_id='load_weather',
    bucket=gs_bucket,
    source_objects=['weather/GlobalLandTemperaturesByCity.csv'],
    destination_project_dataset_table=f'{project_id}:{staging_dataset}.temperature_by_city',
    schema_object='weather/temperature_by_city.json',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    skip_leading_rows=1
)

load_immigration_data = GCSToBigQueryOperator(
    task_id='load_immigration_data',
    bucket=gs_bucket,
    source_objects=['immigration_data/*.parquet'],
    destination_project_dataset_table=f'{project_id}:{staging_dataset}.immigration_data',
    source_format='parquet',
    write_disposition='WRITE_TRUNCATE',
    skip_leading_rows=1,
    autodetect=True
)

# Check loaded data not null
check_us_cities_demo = bigquery.BigQueryCheckOperator(
    task_id='check_us_cities_demo',
    use_legacy_sql=False,
    sql=f'SELECT count(*) FROM `{project_id}.{staging_dataset}.us_cities_demo`'

)

check_airports = bigquery.BigQueryCheckOperator(
    task_id='check_airports',
    use_legacy_sql=False,
    sql=f'SELECT count(*) FROM `{project_id}.{staging_dataset}.airport_codes`'
)

check_weather = bigquery.BigQueryCheckOperator(
    task_id='check_weather',
    use_legacy_sql=False,
    sql=f'SELECT count(*) FROM `{project_id}.{staging_dataset}.temperature_by_city`'
)

check_immigration_data = bigquery.BigQueryCheckOperator(
    task_id='check_immigration_data',
    use_legacy_sql=False,
    sql=f'SELECT count(*) FROM `{project_id}.{staging_dataset}.immigration_data`'
)

loaded_data_to_staging = DummyOperator(
    task_id='loaded_data_to_staging'
)

# Load dimensions data from files directly to DWH table
load_country = GCSToBigQueryOperator(
    task_id='load_country',
    bucket=gs_bucket,
    source_objects=['master_data/I94CIT_I94RES.csv'],
    destination_project_dataset_table=f'{project_id}:{dwh_dataset}.d_country',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    skip_leading_rows=1,
    schema_fields=[
        {'name': 'COUNTRY_ID', 'type': 'NUMERIC', 'mode': 'NULLABLE'},
        {'name': 'COUNTRY_NAME', 'type': 'STRING', 'mode': 'NULLABLE'},
    ]
)

load_port = GCSToBigQueryOperator(
    task_id='load_port',
    bucket=gs_bucket,
    source_objects=['master_data/I94PORT.csv'],
    destination_project_dataset_table=f'{project_id}:{dwh_dataset}.d_port',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    skip_leading_rows=1,
    schema_fields=[
        {'name': 'PORT_ID', 'type': 'STRING', 'mode': 'NULLABLE'},
        {'name': 'PORT_NAME', 'type': 'STRING', 'mode': 'NULLABLE'},
    ]
)

load_state = GCSToBigQueryOperator(
    task_id='load_state',
    bucket=gs_bucket,
    source_objects=['master_data/I94ADDR.csv'],
    destination_project_dataset_table=f'{project_id}:{dwh_dataset}.d_state',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    skip_leading_rows=1,
    schema_fields=[
        {'name': 'STATE_ID', 'type': 'STRING', 'mode': 'NULLABLE'},
        {'name': 'STATE_NAME', 'type': 'STRING', 'mode': 'NULLABLE'},
    ]
)

# Transform, load, and check fact data
create_immigration_data = bigquery.BigQueryInsertJobOperator(
    task_id='create_immigration_data',
    configuration={
        "query": {
            "query": F_IMMIGRATION_DATA_QUERY,
            "useLegacySql": False,
            "destinationTable": {
                "projectId": project_id,
                "datasetId": dwh_dataset,
                "tableId": "f_immigration_data"
            },
            "writeDisposition": job.WriteDisposition.WRITE_TRUNCATE
        }
    }
)

check_f_immigration_data = bigquery.BigQueryCheckOperator(
    task_id='check_f_immigration_data',
    use_legacy_sql=False,
    sql=f'SELECT count(*) = count(distinct cicid) FROM `{project_id}.{dwh_dataset}.f_immigration_data`'
)

# Create remaining dimensions data
create_d_time = bigquery.BigQueryInsertJobOperator(
    task_id='create_d_time',
    configuration={
        "query": {
            "query": D_TIME_QUERY,
            "useLegacySql": False,
            "destinationTable": {
                'projectId': project_id,
                'datasetId': dwh_dataset,
                'tableId': 'd_time'
            },
            "writeDisposition": job.WriteDisposition.WRITE_TRUNCATE
        }
    }
)

create_d_weather = bigquery.BigQueryInsertJobOperator(
    task_id='create_d_weather',
    configuration={
        "query": {
            "query": D_WEATHER_QUERY,
            "useLegacySql": False,
            "destinationTable": {
                'projectId': project_id,
                'datasetId': dwh_dataset,
                'tableId': 'd_weather'
            },
            "writeDisposition": job.WriteDisposition.WRITE_TRUNCATE
        }
    }
)

create_d_airport = bigquery.BigQueryInsertJobOperator(
    task_id='create_d_airport',
    configuration={
        "query": {
            "query": D_AIRPORT_QUERY,
            "useLegacySql": False,
            "destinationTable": {
                'projectId': project_id,
                'datasetId': dwh_dataset,
                'tableId': 'd_airport'
            },
            "writeDisposition": job.WriteDisposition.WRITE_TRUNCATE
        }
    }
)

create_d_city_demo = bigquery.BigQueryInsertJobOperator(
    task_id='create_d_city_demo',
    configuration={
        "query": {
            "query": D_CITY_DEMO_QUERY,
            "useLegacySql": False,
            "destinationTable": {
                'projectId': project_id,
                'datasetId': dwh_dataset,
                'tableId': 'd_city_demo'
            },
            "writeDisposition": job.WriteDisposition.WRITE_TRUNCATE
        }
    }
)

finish_pipeline = DummyOperator(
    task_id='finish_pipeline'
)

# Define task dependencies
start_pipeline >> [load_us_cities_demo, load_airports, load_weather, load_immigration_data]

load_us_cities_demo >> check_us_cities_demo
load_airports >> check_airports
load_weather >> check_weather
load_immigration_data >> check_immigration_data

[check_us_cities_demo, check_airports, check_weather, check_immigration_data] >> loaded_data_to_staging

loaded_data_to_staging >> [load_country, load_port, load_state] >> create_immigration_data >> check_f_immigration_data

check_f_immigration_data >> [create_d_time, create_d_weather, create_d_airport, create_d_city_demo] >> finish_pipeline

# Import packages
from airflow.models import Variable
from airflow import DAG
from airflow.operators.dummy import DummyOperator
from datetime import datetime, timedelta
from airflow.providers.google.cloud.transfers.gcs_to_bigquery import GCSToBigQueryOperator
from airflow.providers.google.cloud.operators import bigquery
from google.cloud.bigquery import job

# Define default arguments
default_args = {
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 5,
    'retry_delay': timedelta(minutes=5),
}

# Define dag variables
default_project_id = "ifide-i94-immigration"
default_bucket_name = "i94bucket"
project_id = Variable.get(key="project_id", default_var=default_project_id)
gs_bucket = Variable.get(key="gs_bucket", default_var=default_bucket_name)

staging_dataset = 'immigration_dwh_staging'
dwh_dataset = 'immigration_dwh'

# Queries
F_IMMIGRATION_DATA_QUERY = open('/opt/airflow/dags/sql/F_IMMIGRATION_DATA.sql', 'r').read() \
    .format(project_id=project_id,
            staging_dataset=staging_dataset,
            dwh_dataset=dwh_dataset)

D_TIME_QUERY = open('/opt/airflow/dags/sql/D_TIME.sql', 'r').read() \
    .format(project_id=project_id,
            staging_dataset=staging_dataset,
            dwh_dataset=dwh_dataset)

D_WEATHER_QUERY = open('/opt/airflow/dags/sql/D_WEATHER.sql', 'r').read() \
    .format(project_id=project_id,
            staging_dataset=staging_dataset,
            dwh_dataset=dwh_dataset)

D_AIRPORT_QUERY = open('/opt/airflow/dags/sql/D_AIRPORT.sql', 'r').read() \
    .format(project_id=project_id,
            staging_dataset=staging_dataset,
            dwh_dataset=dwh_dataset)

D_CITY_DEMO_QUERY = open('/opt/airflow/dags/sql/D_CITY_DEMO.sql', 'r').read() \
    .format(project_id=project_id,
            staging_dataset=staging_dataset,
            dwh_dataset=dwh_dataset)

# Define dag
dag = DAG('cloud-data-lake-pipeline',
          start_date=datetime.now(),
          schedule_interval='@once',
          concurrency=5,
          max_active_runs=1,
          default_args=default_args)

start_pipeline = DummyOperator(
    task_id='start_pipeline',
    dag=dag
)

# Load data from GCS to BQ
load_us_cities_demo = GCSToBigQueryOperator(
    task_id='load_us_cities_demo',
    bucket=gs_bucket,
    source_objects=['cities/us-cities-demographics.csv'],
    destination_project_dataset_table=f'{project_id}:{staging_dataset}.us_cities_demo',
    schema_object='cities/us_cities_demo.json',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    field_delimiter=';',
    skip_leading_rows=1
)

load_airports = GCSToBigQueryOperator(
    task_id='load_airports',
    bucket=gs_bucket,
    source_objects=['airports/airport-codes_csv.csv'],
    destination_project_dataset_table=f'{project_id}:{staging_dataset}.airport_codes',
    schema_object='airports/airport_codes.json',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    skip_leading_rows=1
)

load_weather = GCSToBigQueryOperator(
    task_id='load_weather',
    bucket=gs_bucket,
    source_objects=['weather/GlobalLandTemperaturesByCity.csv'],
    destination_project_dataset_table=f'{project_id}:{staging_dataset}.temperature_by_city',
    schema_object='weather/temperature_by_city.json',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    skip_leading_rows=1
)

load_immigration_data = GCSToBigQueryOperator(
    task_id='load_immigration_data',
    bucket=gs_bucket,
    source_objects=['immigration_data/*.parquet'],
    destination_project_dataset_table=f'{project_id}:{staging_dataset}.immigration_data',
    source_format='parquet',
    write_disposition='WRITE_TRUNCATE',
    skip_leading_rows=1,
    autodetect=True
)

# Check loaded data not null
check_us_cities_demo = bigquery.BigQueryCheckOperator(
    task_id='check_us_cities_demo',
    use_legacy_sql=False,
    sql=f'SELECT count(*) FROM `{project_id}.{staging_dataset}.us_cities_demo`'

)

check_airports = bigquery.BigQueryCheckOperator(
    task_id='check_airports',
    use_legacy_sql=False,
    sql=f'SELECT count(*) FROM `{project_id}.{staging_dataset}.airport_codes`'
)

check_weather = bigquery.BigQueryCheckOperator(
    task_id='check_weather',
    use_legacy_sql=False,
    sql=f'SELECT count(*) FROM `{project_id}.{staging_dataset}.temperature_by_city`'
)

check_immigration_data = bigquery.BigQueryCheckOperator(
    task_id='check_immigration_data',
    use_legacy_sql=False,
    sql=f'SELECT count(*) FROM `{project_id}.{staging_dataset}.immigration_data`'
)

loaded_data_to_staging = DummyOperator(
    task_id='loaded_data_to_staging'
)

# Load dimensions data from files directly to DWH table
load_country = GCSToBigQueryOperator(
    task_id='load_country',
    bucket=gs_bucket,
    source_objects=['master_data/I94CIT_I94RES.csv'],
    destination_project_dataset_table=f'{project_id}:{dwh_dataset}.d_country',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    skip_leading_rows=1,
    schema_fields=[
        {'name': 'COUNTRY_ID', 'type': 'NUMERIC', 'mode': 'NULLABLE'},
        {'name': 'COUNTRY_NAME', 'type': 'STRING', 'mode': 'NULLABLE'},
    ]
)

load_port = GCSToBigQueryOperator(
    task_id='load_port',
    bucket=gs_bucket,
    source_objects=['master_data/I94PORT.csv'],
    destination_project_dataset_table=f'{project_id}:{dwh_dataset}.d_port',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    skip_leading_rows=1,
    schema_fields=[
        {'name': 'PORT_ID', 'type': 'STRING', 'mode': 'NULLABLE'},
        {'name': 'PORT_NAME', 'type': 'STRING', 'mode': 'NULLABLE'},
    ]
)

load_state = GCSToBigQueryOperator(
    task_id='load_state',
    bucket=gs_bucket,
    source_objects=['master_data/I94ADDR.csv'],
    destination_project_dataset_table=f'{project_id}:{dwh_dataset}.d_state',
    write_disposition='WRITE_TRUNCATE',
    source_format='csv',
    skip_leading_rows=1,
    schema_fields=[
        {'name': 'STATE_ID', 'type': 'STRING', 'mode': 'NULLABLE'},
        {'name': 'STATE_NAME', 'type': 'STRING', 'mode': 'NULLABLE'},
    ]
)

# Transform, load, and check fact data
create_immigration_data = bigquery.BigQueryInsertJobOperator(
    task_id='create_immigration_data',
    configuration={
        "query": {
            "query": F_IMMIGRATION_DATA_QUERY,
            "useLegacySql": False,
            "destinationTable": {
                "projectId": project_id,
                "datasetId": dwh_dataset,
                "tableId": "f_immigration_data"
            },
            "writeDisposition": job.WriteDisposition.WRITE_TRUNCATE
        }
    }
)

check_f_immigration_data = bigquery.BigQueryCheckOperator(
    task_id='check_f_immigration_data',
    use_legacy_sql=False,
    sql=f'SELECT count(*) = count(distinct cicid) FROM `{project_id}.{dwh_dataset}.f_immigration_data`'
)

# Create remaining dimensions data
create_d_time = bigquery.BigQueryInsertJobOperator(
    task_id='create_d_time',
    configuration={
        "query": {
            "query": D_TIME_QUERY,
            "useLegacySql": False,
            "destinationTable": {
                'projectId': project_id,
                'datasetId': dwh_dataset,
                'tableId': 'd_time'
            },
            "writeDisposition": job.WriteDisposition.WRITE_TRUNCATE
        }
    }
)

create_d_weather = bigquery.BigQueryInsertJobOperator(
    task_id='create_d_weather',
    configuration={
        "query": {
            "query": D_WEATHER_QUERY,
            "useLegacySql": False,
            "destinationTable": {
                'projectId': project_id,
                'datasetId': dwh_dataset,
                'tableId': 'd_weather'
            },
            "writeDisposition": job.WriteDisposition.WRITE_TRUNCATE
        }
    }
)

create_d_airport = bigquery.BigQueryInsertJobOperator(
    task_id='create_d_airport',
    configuration={
        "query": {
            "query": D_AIRPORT_QUERY,
            "useLegacySql": False,
            "destinationTable": {
                'projectId': project_id,
                'datasetId': dwh_dataset,
                'tableId': 'd_airport'
            },
            "writeDisposition": job.WriteDisposition.WRITE_TRUNCATE
        }
    }
)

create_d_city_demo = bigquery.BigQueryInsertJobOperator(
    task_id='create_d_city_demo',
    configuration={
        "query": {
            "query": D_CITY_DEMO_QUERY,
            "useLegacySql": False,
            "destinationTable": {
                'projectId': project_id,
                'datasetId': dwh_dataset,
                'tableId': 'd_city_demo'
            },
            "writeDisposition": job.WriteDisposition.WRITE_TRUNCATE
        }
    }
)

finish_pipeline = DummyOperator(
    task_id='finish_pipeline'
)

# Define task dependencies
start_pipeline >> [load_us_cities_demo, load_airports, load_weather, load_immigration_data]

load_us_cities_demo >> check_us_cities_demo
load_airports >> check_airports
load_weather >> check_weather
load_immigration_data >> check_immigration_data

[check_us_cities_demo, check_airports, check_weather, check_immigration_data] >> loaded_data_to_staging

loaded_data_to_staging >> [load_country, load_port, load_state] >> create_immigration_data >> check_f_immigration_data

check_f_immigration_data >> [create_d_time, create_d_weather, create_d_airport, create_d_city_demo] >> finish_pipeline

# Import packages # Define default arguments # Define dag variables # Queries # Define dag # Load data from GCS to BQ # Check loaded data not null # Load dimensions data from files directly to DWH table # Transform, load, and check fact data # Create remaining dimensions data # Define task dependencies