def test_trigger_dagrun_twice(self): """Test TriggerDagRunOperator with custom execution_date.""" utc_now = timezone.utcnow() task = TriggerDagRunOperator( task_id="test_trigger_dagrun_with_execution_date", trigger_dag_id=TRIGGERED_DAG_ID, execution_date=utc_now, dag=self.dag, poke_interval=1, reset_dag_run=True, wait_for_completion=True, ) run_id = f"manual__{utc_now.isoformat()}" with create_session() as session: dag_run = DagRun( dag_id=TRIGGERED_DAG_ID, execution_date=utc_now, state=State.SUCCESS, run_type="manual", run_id=run_id, ) session.add(dag_run) session.commit() task.execute(None) dagruns = session.query(DagRun).filter( DagRun.dag_id == TRIGGERED_DAG_ID).all() self.assertEqual(len(dagruns), 1) self.assertTrue(dagruns[0].external_trigger) self.assertEqual(dagruns[0].execution_date, utc_now)
def create_retrigger_operator(dag, task_id=None): if not task_id: task_id = f'retrigger_{dag.dag_id}' return TriggerDagRunOperator( task_id=task_id, trigger_dag_id=dag.dag_id, dag=dag )
def test_trigger_dagrun(self): """Test TriggerDagRunOperator.""" task = TriggerDagRunOperator(task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, dag=self.dag) task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) with create_session() as session: dagruns = session.query(DagRun).filter( DagRun.dag_id == TRIGGERED_DAG_ID).all() self.assertEqual(len(dagruns), 1) self.assertTrue(dagruns[0].external_trigger)
def test_trigger_dagrun_with_wait_for_completion_true_fail(self): """Test TriggerDagRunOperator with wait_for_completion but triggered dag fails.""" execution_date = DEFAULT_DATE task = TriggerDagRunOperator( task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, execution_date=execution_date, wait_for_completion=True, poke_interval=10, failed_states=[State.RUNNING], dag=self.dag, ) with self.assertRaises(AirflowException): task.run(start_date=execution_date, end_date=execution_date)
def test_trigger_dagrun_operator_templated_invalid_conf(self): """Test passing a conf that is not JSON Serializable raise error.""" with pytest.raises( AirflowException, match="^conf parameter should be JSON Serializable$"): TriggerDagRunOperator( task_id="test_trigger_dagrun_with_invalid_conf", trigger_dag_id=TRIGGERED_DAG_ID, conf={ "foo": "{{ dag.dag_id }}", "datetime": timezone.utcnow() }, dag=self.dag, )
def test_trigger_dagrun_operator_templated_conf(self): """Test passing a templated conf to the triggered DagRun.""" task = TriggerDagRunOperator( task_id="test_trigger_dagrun_with_str_execution_date", trigger_dag_id=TRIGGERED_DAG_ID, conf={"foo": "{{ dag.dag_id }}"}, dag=self.dag, ) task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) with create_session() as session: dagruns = session.query(DagRun).filter( DagRun.dag_id == TRIGGERED_DAG_ID).all() self.assertEqual(len(dagruns), 1) self.assertTrue(dagruns[0].conf, {"foo": TEST_DAG_ID})
def test_trigger_dagrun_with_templated_execution_date(self): """Test TriggerDagRunOperator with templated execution_date.""" task = TriggerDagRunOperator( task_id="test_trigger_dagrun_with_str_execution_date", trigger_dag_id=TRIGGERED_DAG_ID, execution_date="{{ execution_date }}", dag=self.dag, ) task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) with create_session() as session: dagruns = session.query(DagRun).filter( DagRun.dag_id == TRIGGERED_DAG_ID).all() self.assertEqual(len(dagruns), 1) self.assertTrue(dagruns[0].external_trigger) self.assertEqual(dagruns[0].execution_date, DEFAULT_DATE)
def test_trigger_dagrun_with_reset_dag_run_false(self): """Test TriggerDagRunOperator with reset_dag_run.""" execution_date = DEFAULT_DATE task = TriggerDagRunOperator( task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, execution_date=execution_date, reset_dag_run=False, dag=self.dag, ) task.run(start_date=execution_date, end_date=execution_date, ignore_ti_state=True) with self.assertRaises(DagRunAlreadyExists): task.run(start_date=execution_date, end_date=execution_date, ignore_ti_state=True)
def test_trigger_dagrun_with_wait_for_completion_true(self): """Test TriggerDagRunOperator with wait_for_completion.""" execution_date = DEFAULT_DATE task = TriggerDagRunOperator( task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, execution_date=execution_date, wait_for_completion=True, poke_interval=10, allowed_states=[State.RUNNING], dag=self.dag, ) task.run(start_date=execution_date, end_date=execution_date) with create_session() as session: dagruns = session.query(DagRun).filter( DagRun.dag_id == TRIGGERED_DAG_ID).all() self.assertEqual(len(dagruns), 1)
def test_trigger_dagrun_with_execution_date(self): """Test TriggerDagRunOperator with custom execution_date.""" utc_now = timezone.utcnow() task = TriggerDagRunOperator( task_id="test_trigger_dagrun_with_execution_date", trigger_dag_id=TRIGGERED_DAG_ID, execution_date=utc_now, dag=self.dag, ) task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True) with create_session() as session: dagruns = session.query(DagRun).filter( DagRun.dag_id == TRIGGERED_DAG_ID).all() self.assertEqual(len(dagruns), 1) self.assertTrue(dagruns[0].external_trigger) self.assertEqual(dagruns[0].execution_date, utc_now)
def test_trigger_dagrun_with_reset_dag_run_true(self): """Test TriggerDagRunOperator with reset_dag_run.""" execution_date = DEFAULT_DATE task = TriggerDagRunOperator( task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID, execution_date=execution_date, reset_dag_run=True, dag=self.dag, ) task.run(start_date=execution_date, end_date=execution_date, ignore_ti_state=True) task.run(start_date=execution_date, end_date=execution_date, ignore_ti_state=True) with create_session() as session: dagruns = session.query(DagRun).filter( DagRun.dag_id == TRIGGERED_DAG_ID).all() assert len(dagruns) == 1 assert dagruns[0].external_trigger
# ================================================ EXAMPLE 1 ================================================= example_1_dag_1 = DAG( dag_id="figure_6_17_example_1_dag_1", start_date=airflow.utils.dates.days_ago(3), schedule_interval="0 0 * * *", ) example_1_dag_2 = DAG( dag_id="figure_6_17_example_1_dag_2", start_date=airflow.utils.dates.days_ago(3), schedule_interval=None, ) DummyOperator(task_id="etl", dag=example_1_dag_1) >> TriggerDagRunOperator( task_id="trigger_dag2", trigger_dag_id="figure_6_17_example_1_dag_2", dag=example_1_dag_1, ) PythonOperator(task_id="report", dag=example_1_dag_2, python_callable=lambda: print("hello")) # ================================================ EXAMPLE 2 ================================================= example_2_dag_1 = DAG( dag_id="figure_6_17_example_2_dag_1", start_date=airflow.utils.dates.days_ago(3), schedule_interval="0 0 * * *", ) example_2_dag_2 = DAG( dag_id="figure_6_17_example_2_dag_2",
from airflow.utils.dates import days_ago with DAG( dag_id="controller_dag_to_trigger_other_dags", default_args={"owner": "airflow"}, start_date=days_ago(1), schedule_interval="@once", ) as dag: start = DummyOperator( task_id='start' ) trigger_1 = TriggerDagRunOperator( task_id="dag_1", trigger_dag_id="dag-to-trigger", # Ensure this equals the dag_id of the DAG to trigger conf={"message": "Hello World"} ) trigger_2 = TriggerDagRunOperator( task_id="dag_2", trigger_dag_id="dag-to-trigger", # Ensure this equals the dag_id of the DAG to trigger conf={"message": "Hello World"} ) some_other_task = DummyOperator( task_id='some-other-task' ) end = DummyOperator( task_id='end' )
for instance_fleet in response['InstanceFleets']: if instance_fleet["InstanceFleetType"] == "TASK": return instance_fleet["Id"] def dagrun_trigger(context, dag_run_obj): cluster_id = context['ti'].xcom_pull('create_cluster', key='return_value') instance_fleet_id = get_spot_instance_fleet_id(cluster_id) _date = datetime.strptime(context['ds'], '%Y-%m-%d') + timedelta(days=1) date = _date.strftime('%Y-%m-%d') pdate = _date.strftime('%Y/%m/%d') logging.info('cluster_id, date, pdate', cluster_id, date, pdate) dag_run_obj.payload = { 'cluster_id': cluster_id, 'instance_fleet_id': instance_fleet_id, 'date': date, 'pdate': pdate } return dag_run_obj trigger = TriggerDagRunOperator( task_id='trigger', trigger_dag_id="subdag-id", python_callable=dagrun_trigger, sla=timedelta(minutes=20), email_on_failure=True, dag=dag ) create_cluster >> trigger
clean_bi_temp_tables = PostgresOperator( task_id='clean_bi_temp_tables', postgres_conn_id=f"SurfRiderDb_{env}_manager_user", sql='clean_bi_temp_tables.sql', dag=dag) logs_status_pipeline = PostgresOperator( task_id='logs_status_pipeline', postgres_conn_id=f"SurfRiderDb_{env}_manager_user", sql='logs_status_pipeline.sql', dag=dag) run_bi_postprocessing = TriggerDagRunOperator( task_id='run_bi_postprocessing', trigger_dag_id=f'bi-postprocessing-{env}', wait_for_completion=True, dag=dag) get_new_campaign_ids >> [ copy_campaign_table, copy_trash_table, copy_trajectory_point_table ] copy_trajectory_point_table >> compute_metrics_bi_temp_trajectory_point compute_metrics_bi_temp_trajectory_point >> compute_bi_temp_trajectory_point_river compute_bi_temp_trajectory_point_river >> compute_bi_temp_campaign_river copy_trash_table >> compute_metrics_bi_temp_trash >> compute_bi_temp_trash_river [ copy_campaign_table, compute_bi_temp_campaign_river, compute_bi_temp_trash_river
start_date=CONFIGS[DAG_ID]["start_date"], tags=["example"], ) with dag: """ main DAG: smart_sensor (looking for run file) -> trigger_external_dag (dag_id_DB_1) -> SubDAG (external_sensor -> print_logs -> remove_file -> print_finish_log | example TaskGroup) -> send_message (into Slack chanell) """ @task() def slack_send_message(): client = WebClient(token=SLACK_TOKEN) try: response = client.chat_postMessage(channel="airflowtask33", text="Hello from your app! :tada:") except SlackApiError as e: assert e.response["error"] # str like 'invalid_auth', 'channel_not_found' sens = SmartFileSensor(task_id="checking_file", filepath=TRIGGER_DIR, fs_conn_id='fs_default') task_trigger = TriggerDagRunOperator( task_id="trigger_database_update", trigger_dag_id="dag_id_DB_1", wait_for_completion=True, poke_interval=15, ) sub_dag = SubDagOperator(task_id='XCOM_sub_dag', subdag = sub_dag_processing(), default_args=DEFAULT_ARGS) task_slack = slack_send_message() sens >> task_trigger >> sub_dag >> task_slack
default_args = { 'owner': 'airflow', 'retries': 3, 'start_date': airflow.utils.dates.days_ago(2), 'retry_delay': timedelta(seconds=10), 'on_failure_callback': dag_failure_notification, 'on_success_callback': dag_success_notification } # 'start_date': datetime.datetime(2021, 09, 01) #airflow.utils.dates.days_ago(2), with DAG('dag_km_dependency_01_seq_01', default_args=default_args, tags=['km'], schedule_interval='* * * * *', catchup=False, dagrun_timeout=timedelta(minutes=90)) as dag: py_start_task = PythonOperator(task_id='py_start_task', python_callable=py_start_task, provide_context=True, op_kwargs={ 'to_email_address': to_email_address, }) trigger_dependent_dag = TriggerDagRunOperator( task_id="trigger_dependent_dag", trigger_dag_id="dag_km_dependency_01_seq_02", wait_for_completion=True) py_start_task >> trigger_dependent_dag
def create_dag(dag_name: str, agg_by: str, doc_type: str, cache_blob: str, path_avro_schema: str, path_local_avro_schemas: str, executor_cores: str, executor_memory: str, executor_instances: str, driver_memory: str, col_type_pk: str, extra_cols: str, max_registry_by_file: str, oracle_conn_id: str, table_ctrl: str, table_ctrl_col_control_var: str, table_ctrl_col_fk: str, table_ctrl_col_dt_ref: str, table_ctrl_col_dt_created: str, oracle_conn_blob: str, table_blob: str, table_blob_col_pk: str, table_blob_col_blob: str) -> airflow.models.dag.DAG: # ----------------- # DAG # ----------------- args = { 'owner': 'job', 'run_as_user': '******', 'start_date': datetime(2021, 8, 17), 'do_xcom_push': False, 'depends_on_past': True, 'retries': 10, 'retry_delay': timedelta(seconds=60), 'dag_name': dag_name } with DAG(dag_id=f'{step}_{dag_name}', description=f'Import data from {dag_name}', schedule_interval='00 19 * * *', catchup=False, default_args=args) as dag: dag.doc_md = __doc__ dag.doc_md = """![image alt <](../big_data.wiki/.attachments/xpto_company.png)""" layer = 'raw' env = Variable.get('env', default_var='dev') control_var = f"{int(Variable.get(f'{dag_name}_control_var', default_var='000000000000000')):015d}" last_control_var = Variable.get(f'{dag_name}_last_control_var', default_var='000000000000000') current_dag_name = dag_name + '_' + control_var total_pg = int( Variable.get(f'{current_dag_name}_total_pg', default_var=1)) list_all_dates = eval( Variable.get(f'{dag_name}_list_all_dates', default_var='[]')) list_current_dates = eval( Variable.get(f'{current_dag_name}_current_dates', default_var='[]')) list_dags = eval( Variable.get(f'{dag_name}_list_dags', default_var='[]')) total_rows = Variable.get('total_rows', default_var='100000') items_by_query = 1000 sql_id = f''' SELECT {table_ctrl_col_fk} id, {table_ctrl_col_control_var} control_var, to_char({table_ctrl_col_dt_ref}, 'DD-MM-YYYY') dt_ref FROM {table_ctrl} WHERE {table_ctrl_col_control_var} > :control_var AND TO_DATE(to_char({table_ctrl_col_dt_created}, 'DD-MM-YYYY'), 'DD-MM-YYYY') < TO_DATE(to_char(trunc(sysdate), 'DD-MM-YYYY'), 'DD-MM-YYYY') ORDER BY {table_ctrl_col_control_var} ASC FETCH FIRST :total_rows ROWS ONLY''' dict_bind_sql_get_data = { 'control_var': f'{control_var}', 'total_rows': f'{total_rows}' } sql_count_id = f''' SELECT COUNT({table_ctrl_col_fk}) FROM {table_ctrl} WHERE {table_ctrl_col_control_var} > :control_var AND TO_DATE(to_char({table_ctrl_col_dt_created}, 'DD-MM-YYYY'), 'DD-MM-YYYY') < TO_DATE(to_char(trunc(sysdate), 'DD-MM-YYYY'), 'DD-MM-YYYY')''' dict_bind_sql_count_id = {'control_var': f'{control_var}'} # ----------------- # TASKS # ----------------- task_start = PythonOperator(task_id='start', python_callable=start_time, depends_on_past=False, op_kwargs={ 'dag_name': dag_name, 'execution_date': '{{ ts }}' }) task_oracle_execute_count = OracleGetResults( task_id='oracle_execute_count', current_dag_name=current_dag_name, oracle_conn_id=oracle_conn_id, sql_count_id=sql_count_id, dict_bind=dict_bind_sql_count_id) task_check_if_contains_data_in_oracle = BranchPythonOperator( task_id='check_if_contains_data_in_oracle', python_callable=AirflowMetaStoreHelper( ).check_if_contains_data_in_oracle, op_kwargs={ 'control_var': control_var, 'last_control_var': last_control_var, 'current_dag_name': current_dag_name, 'redis_conn_id': cache_blob, 'redis_key': f'{dag_name}_original', 'true_case': 'get_id', 'false_case': 'check_len_list_processed_dates' }) task_get_id = OracleToRedisTransfer( task_id='get_id', oracle_conn_id=oracle_conn_id, redis_conn_id=cache_blob, sql=sql_id, dict_bind=dict_bind_sql_get_data, name_redis_key=f'{dag_name}_original') task_fill_data_gap = PythonOperator( task_id='fill_data_gap', python_callable=RedisHelper(cache_blob).fill_data_gaps, op_kwargs={ 'current_dag_name': current_dag_name, 'redis_conn_id': cache_blob, 'redis_key': f'{dag_name}_original' }) task_get_dag_name = PythonOperator( task_id='get_dag_name', python_callable=AirflowMetaStoreHelper().get_dag_name, op_kwargs={ 'current_dag_name': current_dag_name, 'name_list_dags': f'{dag_name}_list_dags', 'list_dags': list_dags }) task_get_date = PythonOperator( task_id='get_date', python_callable=RedisHelper(cache_blob).get_date, op_kwargs={ 'dag_name': dag_name, 'current_dag_name': current_dag_name, 'list_columns': "['id', 'control_var', 'date']", 'redis_key': current_dag_name }) task_split_id_by_date = PythonOperator( task_id='split_id_by_date', python_callable=RedisHelper(cache_blob).split_id_by_date, op_kwargs={ 'current_dag_name': current_dag_name, 'list_current_dates': list_current_dates, 'redis_key': current_dag_name }) task_generate_pagination = PythonOperator( task_id='generate_pagination', python_callable=RedisHelper(cache_blob).generate_pagination, op_kwargs={ 'current_dag_name': current_dag_name, 'items_by_query': items_by_query, 'list_current_dates': list_current_dates, 'redis_key': current_dag_name }) task_generate_sql_by_date = PythonOperator( task_id='generate_sql_by_date', python_callable=RedisHelper(cache_blob).generate_sql_by_date, op_kwargs={ 'current_dag_name': current_dag_name, 'list_current_dates': list_current_dates, 'oracle_conn': oracle_conn_blob, 'table_ctrl': table_ctrl, 'table_ctrl_col_fk': table_ctrl_col_fk, 'table_blob': table_blob, 'table_blob_col_pk': table_blob_col_pk, 'table_blob_col_blob': table_blob_col_blob, 'items_by_query': items_by_query, 'total_pg': total_pg, 'extra_cols': extra_cols, 'redis_key': current_dag_name }) task_extract_decompress_load = OracleBlobToHdfsTransfer( task_id=f'extract_decompress_load', retries=20, dag_name=dag_name, current_dag_name=current_dag_name, oracle_conn_id=oracle_conn_id, query_id=sql_id, table_ctrl_col_fk=table_ctrl_col_fk, extra_cols=extra_cols, oracle_conn_blob=oracle_conn_blob, table_blob_col_pk=table_blob_col_pk, table_blob_col_blob=table_blob_col_blob, path_avro_schema=path_avro_schema, path_local_avro_schemas= f'{path_local_avro_schemas}/{layer}/{dag_name}.avsc', total_pg=total_pg, layer=layer, env=env, step=step, executor_cores=executor_cores, executor_memory=executor_memory, executor_instances=executor_instances, driver_memory=driver_memory, path_ojdbc=path_ojdbc, path_spark_avro=path_spark_avro, path_native_lib=path_native_lib, col_type_pk=col_type_pk, compress_type='snappy', hdfs_conn_id='webhdfs', oracle_driver='oracle.jdbc.driver.OracleDriver', list_current_dates=list_current_dates) task_update_control_var = PythonOperator( task_id='update_control_var', python_callable=AirflowMetaStoreHelper().update_control_var, trigger_rule=TriggerRule.ALL_SUCCESS, depends_on_past=True, op_kwargs={ 'control_var': control_var, 'dag_name': dag_name, 'current_dag_name': current_dag_name, 'redis_conn_id': cache_blob, 'last_control_var': last_control_var, 'list_dags': list_dags, 'total_pg': total_pg, 'list_current_dates': list_current_dates, 'list_all_dates': list_all_dates, 'redis_key': current_dag_name }) task_clear_environment = PythonOperator( task_id='clear_environment', python_callable=clear_environment, trigger_rule=TriggerRule.ALL_SUCCESS, op_kwargs={ 'control_var': control_var, 'dag_name': dag_name, 'redis_conn_id': cache_blob, 'airflow_conn_id': 'airflow_db', 'last_control_var': last_control_var, 'list_dags': list_dags, 'redis_key': current_dag_name }) task_check_len_list_processed_dates = BranchPythonOperator( task_id='check_len_list_processed_dates', trigger_rule=TriggerRule.ALL_SUCCESS, python_callable=check_len_list_processed_dates, op_kwargs={ 'dag_name': dag_name, 'list_all_dates': list_all_dates, 'true_case': 'prepare_execution', 'false_case': 'waiting_execution' }) task_prepare_execution = DummyOperator(task_id='prepare_execution') with TaskGroup( group_id='group_hdfs_concat_file') as group_hdfs_concat_file: task_hdfs_prepare_concat = PythonOperator( task_id='hdfs_prepare_concat', trigger_rule=TriggerRule.ALL_SUCCESS, python_callable=HdfsPrepareConcat('webhdfs').execute, op_kwargs={ 'dag_name': dag_name, 'current_dag_name': current_dag_name, 'hdfs_path': f'/data/{env}/{layer}/{dag_name}', 'agg_by': agg_by, 'layer': layer, 'env': env, 'list_all_dates': list_all_dates, 'path_avro_tools': path_avro_tools }) # TODO: refactor -> create a task list_all_dates = AirflowMetaStoreHelper().set_granularity( list_all_dates=list_all_dates, agg_by=agg_by) for date in list_all_dates: task_concat_file = HdfsConcatFiles( task_id=f'hdfs_concat_file-{date}', retries=100, dag_name=dag_name, date=date, layer=layer, env=env, col_name_control_var=table_ctrl_col_control_var, path_avro_schema=path_avro_schema, hdfs_conn_id='webhdfs', executor_cores=executor_cores, executor_memory=executor_memory, driver_memory=driver_memory, path_ojdbc=path_ojdbc, path_spark_avro=path_spark_avro, path_native_lib=path_native_lib, format_data='avro', compress_type='snappy', max_registry_by_avro=max_registry_by_file) task_hdfs_prepare_concat >> task_concat_file task_create_partitions = PythonOperator( task_id='create_partitions', trigger_rule=TriggerRule.ALL_SUCCESS, python_callable=CreatePartitions().execute, op_kwargs={ 'dag_name': dag_name, 'current_dag_name': current_dag_name, 'list_all_dates': list_all_dates, 'hive_conn_id': 'hive', 'impala_conn_id': 'impala', 'agg_by': agg_by, 'layer': layer, 'env': env }) task_save_execution_state_hdfs = PythonOperator( task_id='save_execution_state_hdfs', python_callable=HdfsHelper('webhdfs').save_execution_state_hdfs, op_kwargs={ 'dag_name': dag_name, 'layer': layer, 'control_var': control_var }) with TaskGroup(group_id='group_generate_statistics' ) as group_generate_statistics: # TODO: refactor -> create a task list_all_dates = AirflowMetaStoreHelper().set_granularity( list_all_dates=list_all_dates, agg_by=agg_by) for date in list_all_dates: PythonOperator(task_id=f'generate_statistics-{date}', retries=50, python_callable=GenerateStatistics().execute, op_kwargs={ 'dag_name': dag_name, 'date': date, 'layer': layer, 'impala_conn_id': 'impala', 'hive_conn_id': 'hive' }) with TaskGroup(group_id='group_check_data_quality' ) as group_check_data_quality: # TODO: refactor -> create a task list_all_dates = AirflowMetaStoreHelper().set_granularity( list_all_dates=list_all_dates, agg_by=agg_by) for date in list_all_dates: CompareDataOracleImpala( task_id=f'compare_oracle_impala_{date}', retries=100, dag_name=dag_name, last_control_var=last_control_var, layer=layer, date=date, table_ctrl=table_ctrl, dt_ref=table_ctrl_col_dt_ref, agg_by=agg_by, oracle_conn_id=oracle_conn_id, hive_conn='impala', table_ctrl_col_fk=table_ctrl_col_fk, table_ctrl_col_dt_created=table_ctrl_col_dt_created) task_check_if_contains_inconsistency = BranchPythonOperator( task_id=f'check_if_contains_inconsistency', trigger_rule=TriggerRule.ALL_SUCCESS, wait_for_downstream=True, python_callable=AirflowMetaStoreHelper( 'airflow_db').check_if_contains_inconsistency, op_kwargs={ 'dag_name': dag_name, 'last_control_var': last_control_var, 'layer': layer, 'true_case': 'prepare_reprocessing_inconsistency_data', 'false_case': f'check_next_dag', 'redis_conn_id': cache_blob, 'redis_key': f'{dag_name}_inconsistency_date' }) task_prepare_reprocessing_inconsistency_data = PrepareReprocessingInconsistencyData( task_id=f'prepare_reprocessing_inconsistency_data', trigger_rule=TriggerRule.ALL_SUCCESS, dag_name=dag_name, current_dag_name=current_dag_name, layer=layer, last_control_var=last_control_var, list_all_dates=list_all_dates, table_ctrl=table_ctrl, table_ctrl_col_fk=table_ctrl_col_fk, table_ctrl_col_control_var=table_ctrl_col_control_var, table_ctrl_col_dt_ref=table_ctrl_col_dt_ref, table_ctrl_col_dt_created=table_ctrl_col_dt_created, hive_conn_id='impala', hdfs_conn_id='webhdfs', airflow_conn_id='airflow_db', oracle_conn_id=oracle_conn_id) task_crash_dag = PythonOperator( task_id=f'crash_dag', trigger_rule=TriggerRule.ALL_SUCCESS, python_callable=crash_dag, ) task_check_next_dag = BranchPythonOperator( task_id='check_next_dag', trigger_rule=TriggerRule.ALL_SUCCESS, python_callable=check_next_dag, op_kwargs={ 'dag_name': dag_name, 'doc_type': doc_type, 'true_case': f'trigger_pre_process_{dag_name}', 'false_case': f'trigger_parser_{dag_name}' }) task_trigger_pre_process = TriggerDagRunOperator( task_id=f'trigger_pre_process_{dag_name}', trigger_dag_id=f"pre_process_{dag_name}") task_trigger_parser = TriggerDagRunOperator( task_id=f'trigger_parser_{dag_name}', trigger_dag_id=f"parser_{dag_name}") task_trigger_import_file = TriggerDagRunOperator( task_id=f'trigger_import_file_{dag_name}', trigger_dag_id=dag.dag_id) task_waiting_execution = DummyOperator( trigger_rule=TriggerRule.ALL_DONE, task_id='waiting_execution') task_end = PythonOperator(task_id='end', python_callable=end_time, op_kwargs={ 'current_dag_name': current_dag_name, 'dag_name': dag_name, 'last_control_var_name': f'{dag_name}_last_control_var', 'list_dates': f'{current_dag_name}_list_dates', 'postgres_conn_id': 'airflow_db' }) # ----------------- # GRAPH # ----------------- # task_check_if_contains_data_in_oracle: true task_start >> task_oracle_execute_count >> task_check_if_contains_data_in_oracle >> task_get_id >> task_fill_data_gap >> [ task_get_date, task_get_dag_name ] >> task_split_id_by_date >> task_generate_pagination >> task_generate_sql_by_date >> task_extract_decompress_load >> task_update_control_var >> [ task_clear_environment, task_trigger_import_file ] >> task_waiting_execution >> task_end # task_check_if_contains_data_in_oracle: false # task_check_len_list_processed_dates: true task_start >> task_oracle_execute_count >> task_check_if_contains_data_in_oracle >> task_check_len_list_processed_dates >> task_prepare_execution >> [ group_hdfs_concat_file, task_save_execution_state_hdfs ] >> task_create_partitions >> [ group_check_data_quality, group_generate_statistics ] >> task_check_if_contains_inconsistency task_start >> task_oracle_execute_count >> task_check_if_contains_data_in_oracle >> task_check_len_list_processed_dates >> task_prepare_execution >> [ group_hdfs_concat_file, task_save_execution_state_hdfs ] >> task_create_partitions >> task_check_if_contains_inconsistency >> task_prepare_reprocessing_inconsistency_data >> task_crash_dag # task_check_next_dag: true task_check_if_contains_inconsistency >> task_check_next_dag >> task_trigger_pre_process >> task_waiting_execution >> task_end # task_check_next_dag: false task_check_if_contains_inconsistency >> task_check_next_dag >> task_trigger_parser >> task_waiting_execution >> task_end # task_check_if_contains_data_in_oracle: false # task_check_len_list_processed_dates: false task_start >> task_oracle_execute_count >> task_check_if_contains_data_in_oracle >> task_check_len_list_processed_dates >> task_waiting_execution >> task_end return dag
from airflow import DAG from airflow.operators.python import PythonOperator from airflow.operators.trigger_dagrun import TriggerDagRunOperator from datetime import datetime default_args = {'start_date': datetime(2021, 1, 1)} def _downloading(): print('downloading') with DAG('trigger_dag', schedule_interval='@daily', default_args=default_args, catchup=False) as dag: downloading = PythonOperator(task_id='downloading', python_callable=_downloading) trigger_target = TriggerDagRunOperator(task_id='trigger_target', trigger_dag_id='target_dag', execution_date='{{ ds }}', reset_dag_run=True) downloading >> trigger_target
def create_dag(dag_name: str, max_registry_by_file: str, db_name: str, table_name: str, col_control_var: str, col_name_dt_ref: str, oracle_conn_id: str, path_avro_schema: str, path_config_docs: str, executor_cores: str, executor_memory: str, executor_instances: str, driver_memory: str) -> airflow.models.dag.DAG: # ----------------- # DAG # ----------------- args = { 'owner': 'job', 'run_as_user': '******', 'start_date': datetime(2021, 8, 1), 'do_xcom_push': False, 'depends_on_past': True, 'retries': 10, 'retry_delay': timedelta(seconds=90), 'dag_name': dag_name } with DAG( dag_id=f'{step}_{dag_name}', description=f'Importa os dados de {dag_name}', # schedule_interval=None, schedule_interval='00 19 * * *', catchup=False, default_args=args) as dag: dag.doc_md = __doc__ dag.doc_md = """![image alt <](../big_data.wiki/.attachments/xpto_company.png)""" layer = 'raw' env = Variable.get('env', default_var='prod') last_control_var = Variable.get(f'{dag_name}_last_control_var', default_var='000000000000000') control_var = f"{int(Variable.get(f'{dag_name}_control_var', default_var='000000000000000')):015d}" current_dag_name = dag_name + '_' + control_var total_rows = Variable.get('total_rows', default_var='100000') sql_get_data = f''' SELECT * FROM {db_name}.{table_name} WHERE {col_control_var} > :control_var AND TO_DATE(to_char({col_name_dt_ref}, 'DD-MM-YYYY'), 'DD-MM-YYYY') < TO_DATE(to_char(trunc(sysdate), 'DD-MM-YYYY'), 'DD-MM-YYYY') ORDER BY {col_control_var} ASC FETCH FIRST :total_rows ROWS ONLY''' dict_bind_sql_get_data = { 'control_var': control_var, 'total_rows': total_rows } sql_count_id = f''' SELECT COUNT({col_control_var}) FROM {db_name}.{table_name} WHERE {col_control_var} > :control_var AND TO_DATE(to_char({col_name_dt_ref}, 'DD-MM-YYYY'), 'DD-MM-YYYY') < TO_DATE(to_char(trunc(sysdate), 'DD-MM-YYYY'), 'DD-MM-YYYY')''' dict_bind_sql_count_id = {'control_var': control_var} # ----------------- # TASKS # ----------------- task_start = PythonOperator(task_id='start', python_callable=start_time, depends_on_past=False, op_kwargs={ 'dag_name': dag_name, 'execution_date': '{{ ts }}' }) task_oracle_execute_count = OracleGetResults( task_id='oracle_execute_count', current_dag_name=current_dag_name, oracle_conn_id=oracle_conn_id, sql_count_id=sql_count_id, dict_bind=dict_bind_sql_count_id) task_check_if_contains_data_in_oracle = BranchPythonOperator( task_id='check_if_contains_data_in_oracle', python_callable=AirflowMetaStoreHelper( ).check_if_contains_data_in_oracle, op_kwargs={ 'control_var': control_var, 'last_control_var': last_control_var, 'current_dag_name': current_dag_name, 'true_case': 'extract_transform_load', 'false_case': 'sync_data' }) task_extract_transform_load = OracleTableToHdfsTransfer( task_id=f'extract_transform_load', retries=20, dag_name=dag_name, current_dag_name=current_dag_name, sql_get_data=sql_get_data, dict_bind=dict_bind_sql_get_data, col_control_var=col_control_var, path_avro_schema=path_avro_schema, layer=layer, env=env, step=step, executor_cores=executor_cores, executor_memory=executor_memory, executor_instances=executor_instances, driver_memory=driver_memory, path_ojdbc=path_ojdbc, path_spark_avro=path_spark_avro, path_native_lib=path_native_lib, compress_type='snappy', oracle_conn_id=oracle_conn_id, hdfs_conn_id='webhdfs', oracle_driver='oracle.jdbc.driver.OracleDriver', max_registry_by_file=max_registry_by_file) task_sync_data = SyncData( task_id='sync_data', dag_name=dag_name, db_name=db_name, table_name=table_name, col_name_control_var=col_control_var, col_name_dt_ref=col_name_dt_ref, path_avro_schema=path_avro_schema, layer=layer, env=env, hdfs_conn_id='webhdfs', oracle_conn_id=oracle_conn_id, oracle_driver='oracle.jdbc.driver.OracleDriver', executor_cores=executor_cores, executor_memory=executor_memory, executor_instances=executor_instances, driver_memory=driver_memory, path_ojdbc=path_ojdbc, path_spark_avro=path_spark_avro, path_native_lib=path_native_lib, compress_type='snappy', max_registry_by_file=max_registry_by_file) task_concat_file = HdfsConcatFiles( task_id=f'hdfs_concat_file', retries=100, dag_name=dag_name, layer=layer, env=env, col_name_control_var=col_control_var, path_avro_schema=path_avro_schema, hdfs_conn_id='webhdfs', executor_cores=executor_cores, executor_memory=executor_memory, driver_memory=driver_memory, path_ojdbc=path_ojdbc, path_spark_avro=path_spark_avro, path_native_lib=path_native_lib, format_data='parquet', compress_type='snappy', max_registry_by_avro=max_registry_by_file) task_save_execution_state_hdfs = PythonOperator( task_id='save_execution_state_hdfs', python_callable=HdfsHelper('webhdfs').save_execution_state_hdfs, op_kwargs={ 'dag_name': dag_name, 'control_var': control_var, 'layer': layer }) task_update_statistics = UpdateStasByTable( task_id=f'update_statistics', retries=50, db_name=dag_name, table_name=layer, impala_conn_id='impala') task_clear_airflow_var = PythonOperator( task_id='clear_airflow_var', trigger_rule=TriggerRule.ALL_SUCCESS, python_callable=clear_airflow_var, op_kwargs={ 'airflow_conn_id': 'airflow_db', 'list_pattern_name_var': f"[" f" '{dag_name}_%_total_row_id'," f" '{dag_name}_raw_total_registry_%'," f" '{dag_name}_list_dags'," f" '{dag_name}_list_path_to_concat'," f" '{dag_name}_list_path_to_rename'," f" '{dag_name}_list_all_dates'" f"]" }) task_update_control_var = OracleUpdateControlVar( task_id='update_control_var', trigger_rule=TriggerRule.ALL_SUCCESS, depends_on_past=True, control_var=control_var, last_control_var=last_control_var, dag_name=dag_name, current_dag_name=current_dag_name, col_control_var=col_control_var, oracle_conn_id=oracle_conn_id, dict_bind=dict_bind_sql_get_data, sql=sql_get_data) task_trigger_import_file = TriggerDagRunOperator( task_id=f'trigger_import_file_{dag_name}', trigger_dag_id=dag.dag_id) task_end = PythonOperator(task_id='end', trigger_rule=TriggerRule.ALL_DONE, python_callable=end_time, op_kwargs={ 'current_dag_name': current_dag_name, 'dag_name': dag_name, 'last_control_var_name': f'{dag_name}_last_control_var', 'postgres_conn_id': 'airflow_db' }) # ----------------- # GRAPH # ----------------- task_start >> task_oracle_execute_count >> task_check_if_contains_data_in_oracle >> task_extract_transform_load >> task_update_control_var >> task_trigger_import_file >> task_end task_start >> task_oracle_execute_count >> task_check_if_contains_data_in_oracle >> task_sync_data >> [ task_concat_file, task_save_execution_state_hdfs ] >> task_update_statistics >> task_clear_airflow_var >> task_end return dag
def _wait_for_supermarket(supermarket_id_): supermarket_path = Path("/data/" + supermarket_id_) data_files = supermarket_path.glob("data-*.csv") success_file = supermarket_path / "_SUCCESS" return data_files and success_file.exists() for supermarket_id in range(1, 5): wait = PythonSensor( task_id=f"wait_for_supermarket_{supermarket_id}", python_callable=_wait_for_supermarket, op_kwargs={"supermarket_id_": f"supermarket{supermarket_id}"}, dag=dag1, ) copy = DummyOperator(task_id=f"copy_to_raw_supermarket_{supermarket_id}", dag=dag1) process = DummyOperator(task_id=f"process_supermarket_{supermarket_id}", dag=dag1) trigger_create_metrics_dag = TriggerDagRunOperator( task_id=f"trigger_create_metrics_dag_supermarket_{supermarket_id}", trigger_dag_id="listing_6_04_dag02", dag=dag1, ) wait >> copy >> process >> trigger_create_metrics_dag compute_differences = DummyOperator(task_id="compute_differences", dag=dag2) update_dashboard = DummyOperator(task_id="update_dashboard", dag=dag2) notify_new_data = DummyOperator(task_id="notify_new_data", dag=dag2) compute_differences >> update_dashboard
# # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """ Example usage of the TriggerDagRunOperator. This example holds 2 DAGs: 1. 1st DAG (example_trigger_controller_dag) holds a TriggerDagRunOperator, which will trigger the 2nd DAG 2. 2nd DAG (example_trigger_target_dag) which will be triggered by the TriggerDagRunOperator in the 1st DAG """ from datetime import datetime from airflow import DAG from airflow.operators.trigger_dagrun import TriggerDagRunOperator with DAG( dag_id="example_trigger_controller_dag", start_date=datetime(2021, 1, 1), catchup=False, schedule_interval="@once", tags=['example'], ) as dag: trigger = TriggerDagRunOperator( task_id="test_trigger_dagrun", trigger_dag_id="example_trigger_target_dag", # Ensure this equals the dag_id of the DAG to trigger conf={"message": "Hello World"}, )
proxy_user='******', args=[ 'HOME=/user', 'USER=root', 'WF_HOME=/user/root', 'ROOT_WORKFLOW_ID=aws_test_2_729508210', 'CURRENT_WORKFLOW_ID=aws_test_2_729508210', 'SLIDE_SIZE=400', 'FETCH_SIZE=1000', 'PARTITION_NUM=1', 'FAIL_THRESHOLD=1000', 'DEBUG=true', 'MASTER=' ], num_executors=1, conf={ 'spark.shuffle.compress': 'false', }, class_name='ru.neoflex.meta.etl2.spark.aws_test_2Job', polling_interval=5, on_failure_callback=on_failure_callback) join >> transformation err4 = KillOperator(task_id='err4', message='transformation or workflow failed', trigger_rule=TriggerRule.ONE_FAILED) transformation >> err4 workflow = TriggerDagRunOperator(task_id='wf_aws_test', trigger_dag_id='wf_aws_test', wait_for_completion=True, poke_interval=5, dag=dag, on_failure_callback=on_failure_callback) transformation >> workflow workflow >> err4 finish = FinishOperator(task_id='finish', ) workflow >> finish
create_timestamp = BashOperator( task_id='create_timestamp', bash_command='touch ~/timestamp_{{ ts_nodash }}', ) task_sensor >> print_results >> remove_file >> create_timestamp return dag # creates a dag with DAG(dag_id='trigger_run', start_date=datetime(2021, 1, 26), schedule_interval='@once') as dag: # checks if a file exists check_for_file = SmartFileSensor(task_id='check_for_file', filepath=path) # triggers another dag trigger_dag = TriggerDagRunOperator(task_id='trigger_dag', trigger_dag_id=external_dag, execution_date='{{ execution_date }}') # creates a sub_dag that processes results process_results = SubDagOperator(task_id='process_results_dag', subdag=create_sub_dag( dag.dag_id, 'process_results_dag', start_date=datetime(2021, 1, 26), schedule_interval='@once')), # sends a slack message alert_slack = PythonOperator(task_id='alert_slack', python_callable=slack_message) check_for_file >> trigger_dag >> process_results >> alert_slack globals()[dag.dag_id] = dag
aws_test2 = LivyOperator( task_id='aws_test2', dag=dag, livy_conn_id='livy_default', file= 's3a://datagram/user/root/deployments/autogenerated_tr_aws_test_2/ru.neoflex.meta.etl2.spark.aws_test_2-1.0' '-SNAPSHOT.jar', proxy_user='******', args=[ 'HOME=/user', 'USER=root', 'WF_HOME=/user/root', 'ROOT_WORKFLOW_ID=aws_test_2_729508210', 'CURRENT_WORKFLOW_ID=aws_test_2_729508210', 'SLIDE_SIZE=400', 'FETCH_SIZE=1000', 'PARTITION_NUM=1', 'FAIL_THRESHOLD=1000', 'DEBUG=true', 'MASTER=' ], num_executors=1, conf={ 'spark.shuffle.compress': 'false', }, class_name='ru.neoflex.meta.etl2.spark.aws_test_2Job', polling_interval=5, ) wf_aws_test = TriggerDagRunOperator(task_id='wf_aws_test', trigger_dag_id='wf_aws_test', wait_for_completion=True, poke_interval=5, dag=dag) aws_test2 >> wf_aws_test