import os from datetime import datetime, timedelta from airflow import DAG from airflow.operators.python_operator import PythonOperator, ShortCircuitOperator import projects.youtube_reports_api.config as config import projects.youtube_reports_api.utils as utils def subdag(parent_dag_name, child_dag_name, default_args, schedule_interval, report_context): """ This is a dag for a report type for one content owner. The dag will first determine where the data will be moved to, then make sure the data is available from YT and not in Snowflake. If these two preliminary tasks succeed, the report will be downloaded from YT, uploaded to s3, and loaded into a Snowflake table. The order of these tasks is determined by Bitshift compostion. where bootstrap >> t0_1 >> t1 means bootstrap must run before t0_1 and t0_1 must run before t1. """ dag_subdag = DAG( dag_id='%s.%s' % (parent_dag_name, child_dag_name), default_args= default_args, schedule_interval= schedule_interval) bootstrap_task_id = 'build_paths_{}'.format(child_dag_name) if "ORCH" in parent_dag_name: # Orchard Files are signifigantly larger than other stores. download_timeout, upload_timeout = 40, 40 if "trafficsource" in child_dag_name: download_timeout, upload_timeout = 60, 60 else: download_timeout, upload_timeout = 10, 10 # Assign tasks to Operators. bootstrap = PythonOperator( task_id= bootstrap_task_id , python_callable= utils.build_paths, op_kwargs= { 'content_owner': report_context['content_owner'], 'job_id': report_context['job_id'], 'job_name': report_context['job_name']}, provide_context= True, execution_timeout= timedelta(minutes=1), retries= 1, dag= dag_subdag) # t0_1 and 2 will halt the script if the report is not avaiable or already in our system # t0_1 = ShortCircuitOperator( # task_id = 'check_upstream_{}'.format(child_dag_name), # python_callable= utils.check_upstream, # op_kwargs= { 'task_id': bootstrap_task_id }, # provide_context= True, # execution_timeout= timedelta(minutes=5), # retries= 1, # dag= dag_subdag) t0_2 = ShortCircuitOperator( task_id = 'check_downstream_{}'.format(child_dag_name), python_callable= utils.check_downstream, op_kwargs= { 'task_id': bootstrap_task_id }, provide_context= True, execution_timeout= timedelta(minutes=1), dag= dag_subdag) t1 = PythonOperator( task_id= 'download_{}'.format(child_dag_name), python_callable= utils.download_report, op_kwargs= { 'task_id': bootstrap_task_id }, provide_context= True, execution_timeout= timedelta(minutes=download_timeout), retries= 2, on_retry_callback= utils.local_clear, dag= dag_subdag) t3 = PythonOperator( task_id= 'upload_{}_staging'.format(child_dag_name), python_callable= utils.insert_to_staging, op_kwargs= { 'fmt': config.PROD_FMT, 'task_id': bootstrap_task_id, 'content_owner': report_context['content_owner']}, execution_timeout= timedelta(minutes=upload_timeout), provide_context= True, retries= 1, dag= dag_subdag) t4 = PythonOperator( task_id= 'upload_{}_prod'.format(child_dag_name), python_callable= utils.insert_to_prod, op_kwargs= { 'fmt': config.PROD_FMT, 'task_id': bootstrap_task_id, 'content_owner': report_context['content_owner']}, execution_timeout= timedelta(minutes=upload_timeout), provide_context= True, retries= 1, dag= dag_subdag) # qc1 = ShortCircuitOperator( # task_id= 'qc_{}_rows'.format(child_dag_name), # python_callable= utils.qc_rows, # op_kwargs= { 'task_id': bootstrap_task_id }, # execution_timeout= timedelta(minutes=8), # provide_context= True, # retries= 1, # dag= dag_subdag) qc2 = ShortCircuitOperator( task_id= 'qc_{}_cols'.format(child_dag_name), python_callable= utils.qc_columns, op_kwargs= { 'task_id': bootstrap_task_id }, execution_timeout= timedelta(minutes=2), provide_context= True, retries= 1, dag= dag_subdag) # DAG Assignement via Bitshift Composition. # bootstrap >> t0_1 >> t1 << t0_2 << bootstrap t1 >> t3 #t1 >> t2 >> t3 # also removed 2017-01-04 t3 >> qc1 >> t4 << qc2 << t3 #switch back in production... return dag_subdag # note the backfill subdag could be nexted into this subdag. def backfill_subdag(parent_dag_name, child_dag_name, default_args, schedule_interval, report_context, time_delta): """ This is a dag for a report type for one content owner. The dag will first determine where the data will be moved to, then make sure the data is available from YT and not in Snowflake. If these two preliminary tasks succeed, the report will be downloaded from YT, uploaded to s3, and loaded into a Snowflake table. The order of these tasks is determined by Bitshift compostion. where bootstrap >> t0_1 >> t1 means bootstrap must run before t0_1 and t0_1 must run before t1. """ dag_subdag = DAG( dag_id='%s.%s' % (parent_dag_name, child_dag_name), default_args= default_args, schedule_interval= schedule_interval) bootstrap_task_id = 'build_paths_{}'.format(child_dag_name) download_timeout, upload_timeout = 10, 10 # Assign tasks to Operators. bootstrap = PythonOperator( task_id= bootstrap_task_id , python_callable= utils.build_paths, op_kwargs= { 'content_owner': report_context['content_owner'], 'job_id': report_context['job_id'], 'job_name': report_context['job_name'], 'time_delta': time_delta}, provide_context= True, execution_timeout= timedelta(minutes=1), retries= 1, dag= dag_subdag) t1 = PythonOperator( task_id= 'download_{}'.format(child_dag_name), python_callable= utils.download_report, op_kwargs= { 'task_id': bootstrap_task_id }, provide_context= True, execution_timeout= timedelta(minutes=download_timeout), retries= 2, on_retry_callback= utils.local_clear, dag= dag_subdag) t3 = PythonOperator( task_id= 'upload_{}_staging'.format(child_dag_name), python_callable= utils.insert_to_staging, op_kwargs= { 'fmt': config.PROD_FMT, 'task_id': bootstrap_task_id}, execution_timeout= timedelta(minutes=upload_timeout), provide_context= True, retries= 1, dag= dag_subdag) t3b = PythonOperator( task_id= 'clean_{}_prod'.format(child_dag_name), python_callable= utils.clear_prod, op_kwargs= { 'task_id': bootstrap_task_id}, execution_timeout= timedelta(minutes=upload_timeout), provide_context= True, retries= 1, dag= dag_subdag) t4 = PythonOperator( task_id= 'upload_{}_prod'.format(child_dag_name), python_callable= utils.insert_to_prod, op_kwargs= { 'fmt': config.PROD_FMT, 'task_id': bootstrap_task_id}, execution_timeout= timedelta(minutes=upload_timeout), provide_context= True, retries= 1, dag= dag_subdag) qc1 = ShortCircuitOperator( task_id= 'qc_{}_rows'.format(child_dag_name), python_callable= utils.qc_rows, op_kwargs= { 'task_id': bootstrap_task_id }, execution_timeout= timedelta(minutes=8), provide_context= True, retries= 1, dag= dag_subdag) qc2 = ShortCircuitOperator( task_id= 'qc_{}_cols'.format(child_dag_name), python_callable= utils.qc_columns, op_kwargs= { 'task_id': bootstrap_task_id }, execution_timeout= timedelta(minutes=2), provide_context= True, retries= 1, dag= dag_subdag) # DAG Assignement via Bitshift Composition. # bootstrap >> t0_1 >> t1 << bootstrap t1 >> t3 #t1 >> t2 >> t3 # also removed 2017-01-04 t3 >> qc1 >> t3b << qc2 << t3 # t3b << qc2 << t3 t3b >> t4 return dag_subdag