# Generated via https://github.com/mozilla/bigquery-etl/blob/master/bigquery_etl/query_scheduling/generate_airflow_dags.py

from airflow import DAG
from airflow.operators.sensors import ExternalTaskSensor
import datetime
from operators.gcp_container_operator import GKEPodOperator
from utils.gcp import gke_command

default_args = {{ 
    default_args.to_dict() | 
    format_attr("start_date", "format_date") |
    format_attr("end_date", "format_date") |
    format_attr("retry_delay", "format_timedelta") 
}}

with DAG('{{ name }}', default_args=default_args{%+ if schedule_interval != None -%}, schedule_interval={{ schedule_interval | format_timedelta | format_schedule_interval }}{%+ endif -%}) as dag:
    docker_image = "mozilla/bigquery-etl:latest"
{% for task in tasks | sort(attribute='task_name') %}
    {{ task.task_name }} = GKEPodOperator(
        task_id="{{ task.task_name }}",
        name="{{ task.task_name }}",
        arguments=["script/publish_public_data_json"]
        + ["--query_file=sql/{{ task.project }}/{{ task.dataset }}/{{ task.table }}_{{ task.version }}/query.sql"]
        + ["--destination_table={{ task.table }}{% if task.date_partition_parameter %}${% raw %}{{ds_nodash}}{% endraw %}{% endif %}"]
        + ["--dataset_id={{ task.dataset }}"]
        + ["--project_id={{ task.project }}"]
        {% if task.date_partition_parameter -%}
        + ["--parameter={{ task.date_partition_parameter }}:DATE:{% raw %}{{ds}}{% endraw -%}"]
        {%- endif -%},
        image=docker_image,
        dag=dag
    )
{% endfor -%}

{% for task in tasks | sort(attribute='task_name') %}
    {% for dependency in task.dependencies | sort(attribute='task_id') -%}
    {% if dependency.dag_name == name -%}
    {{ task.task_name }}.set_upstream({{ dependency.task_id }})
    {% else -%}
    wait_for_{{ dependency.task_id }} = ExternalTaskSensor(
        task_id='wait_for_{{ dependency.task_id }}',
        external_dag_id='{{ dependency.dag_name }}',
        external_task_id='{{ dependency.task_id }}',
        {% if dependency.execution_delta -%}
        execution_delta={{ dependency.execution_delta | format_timedelta | format_repr }},
        {% endif -%}
        check_existence=True,
        mode='reschedule',
        pool='DATA_ENG_EXTERNALTASKSENSOR',
    )
        
    {{ task.task_name }}.set_upstream(wait_for_{{ dependency.task_id }})
    {% endif -%}
    {% endfor -%}
{% endfor %}

    public_data_gcs_metadata = gke_command(
        task_id="public_data_gcs_metadata",
        command=["script/publish_public_data_gcs_metadata"],
        docker_image=docker_image,
        dag=dag
    )

    public_data_gcs_metadata.set_upstream(
        [
        {% for task in tasks | sort(attribute='task_name') %}
            {{ task.task_name }},
        {% endfor -%}
        ]
    )
