from airflow import DAG
from template.utils import hours_ago

from template.dag_template import build_dag
from airflow.models import Variable

# dag id
dag_id = 'Custom-EMR-cluster_v001'

# airflow variables
No_of_core_instances = int(Variable.get('CORE_INSTANCES'))
No_of_task_instances = int(Variable.get('TASK_INSTANCES'))
Master_instance_type = Variable.get('MASTER_INSTANCE_TYPE_FOR_CUSTOM_EMR_CLUSTER')
Core_instance_type = Variable.get('CORE_INSTANCE_TYPE_FOR_CUSTOM_EMR_CLUSTER')
Task_instance_type = Variable.get('TASK_INSTANCE_TYPE_FOR_CUSTOM_EMR_CLUSTER')

# emr steps
emr_steps = """[]"""

# cluster level parameters (optional)
cluster_args = {
    "master-instance-types": Master_instance_type,
    "core-instance-types": Core_instance_type,
    "task-instance-types": Task_instance_type,
    "core-instance-capacity": No_of_core_instances,
    "task-instance-capacity": No_of_task_instances,
    "KeepJobFlowAliveWhenNoSteps": "True",
    "job-type": "stream",
    "emr-version": "emr-6.14.0",
    "CustomAmiId": "ami-03105f3e628d0aa1b"
}

# dag parameter
dag_args = {
    'owner': 'data.engineers@viooh.com',
    'start_date': hours_ago(1)
}

dag = DAG(
    dag_id,
    schedule_interval=None,  # cron expression
    default_args=dag_args)

build_dag(emr_steps=emr_steps, dag=dag, cluster_args=cluster_args)