def test_different_target_states(self):
self.mock_emr_client.describe_cluster.side_effect = [
DESCRIBE_CLUSTER_STARTING_RETURN, # return False
DESCRIBE_CLUSTER_BOOTSTRAPPING_RETURN, # return False
DESCRIBE_CLUSTER_RUNNING_RETURN, # return True
DESCRIBE_CLUSTER_WAITING_RETURN, # will not be used
DESCRIBE_CLUSTER_TERMINATED_RETURN, # will not be used
DESCRIBE_CLUSTER_TERMINATED_WITH_ERRORS_RETURN, # will not be used
]
with patch('boto3.session.Session', self.boto3_session_mock):
operator = EmrJobFlowSensor(
task_id='test_task',
poke_interval=0,
job_flow_id='j-8989898989',
aws_conn_id='aws_default',
target_states=['RUNNING', 'WAITING'],
)
operator.execute(None)
# make sure we called twice
assert self.mock_emr_client.describe_cluster.call_count == 3
# make sure it was called with the job_flow_id
calls = [unittest.mock.call(ClusterId='j-8989898989')]
self.mock_emr_client.describe_cluster.assert_has_calls(calls)
def test_execute_calls_with_the_job_flow_id_until_it_reaches_a_target_state(
self):
self.mock_emr_client.describe_cluster.side_effect = [
DESCRIBE_CLUSTER_STARTING_RETURN,
DESCRIBE_CLUSTER_RUNNING_RETURN,
DESCRIBE_CLUSTER_TERMINATED_RETURN,
]
with patch('boto3.session.Session', self.boto3_session_mock):
operator = EmrJobFlowSensor(task_id='test_task',
poke_interval=0,
job_flow_id='j-8989898989',
aws_conn_id='aws_default')
operator.execute(None)
# make sure we called twice
assert self.mock_emr_client.describe_cluster.call_count == 3
# make sure it was called with the job_flow_id
calls = [unittest.mock.call(ClusterId='j-8989898989')]
self.mock_emr_client.describe_cluster.assert_has_calls(calls)
def test_execute_calls_with_the_job_flow_id_until_it_reaches_failed_state_with_exception(
self):
self.mock_emr_client.describe_cluster.side_effect = [
DESCRIBE_CLUSTER_RUNNING_RETURN,
DESCRIBE_CLUSTER_TERMINATED_WITH_ERRORS_RETURN,
]
with patch('boto3.session.Session', self.boto3_session_mock):
operator = EmrJobFlowSensor(task_id='test_task',
poke_interval=0,
job_flow_id='j-8989898989',
aws_conn_id='aws_default')
with pytest.raises(AirflowException):
operator.execute(None)
# make sure we called twice
assert self.mock_emr_client.describe_cluster.call_count == 2
# make sure it was called with the job_flow_id
self.mock_emr_client.describe_cluster.assert_called_once_with(
ClusterId='j-8989898989')
default_args={
'owner': 'airflow',
'depends_on_past': False,
'email': ['*****@*****.**'],
'email_on_failure': False,
'email_on_retry': False,
},
dagrun_timeout=timedelta(hours=2),
start_date=days_ago(2),
schedule_interval='0 3 * * *',
tags=['example'],
) as dag:
# [START howto_operator_emr_automatic_steps_tasks]
job_flow_creator = EmrCreateJobFlowOperator(
task_id='create_job_flow',
job_flow_overrides=JOB_FLOW_OVERRIDES,
aws_conn_id='aws_default',
emr_conn_id='emr_default',
)
job_sensor = EmrJobFlowSensor(
task_id='check_job_flow',
job_flow_id=job_flow_creator.output,
aws_conn_id='aws_default',
)
# [END howto_operator_emr_automatic_steps_tasks]
# Task dependency created via `XComArgs`:
# job_flow_creator >> job_sensor
False,
'TerminationProtected':
False,
},
'Steps': SPARK_TEST_STEPS,
'JobFlowRole': 'EMR_EC2_DefaultRole',
'ServiceRole': 'EMR_DefaultRole',
}
with DAG(
dag_id='emr_job_flow_automatic_steps_dag',
default_args=DEFAULT_ARGS,
dagrun_timeout=timedelta(hours=2),
schedule_interval='0 3 * * *',
tags=['example'],
) as dag:
job_flow_creator = EmrCreateJobFlowOperator(
task_id='create_job_flow',
job_flow_overrides=JOB_FLOW_OVERRIDES,
aws_conn_id='aws_default',
emr_conn_id='emr_default')
job_sensor = EmrJobFlowSensor(
task_id='check_job_flow',
job_flow_id=
"{{ task_instance.xcom_pull(task_ids='create_job_flow', key='return_value') }}",
aws_conn_id='aws_default')
job_flow_creator >> job_sensor
STAGING LAYER PRE-PROCESSING
(Spark consume Sources save to Parquet)
'''
manifold_emr_creator = EmrCreateJobFlowOperator(
task_id='create_manifold_emr',
job_flow_overrides=JOB_FLOW_OVERRIDES,
aws_conn_id='aws_credentials',
emr_conn_id='emr_credentials',
)
manifold_emr_job_sensor = EmrJobFlowSensor(
task_id='check_emr_completion',
job_flow_id=
"{{ task_instance.xcom_pull(task_ids='create_manifold_emr', key='return_value') }}",
aws_conn_id='aws_credentials',
)
'''
STAGING TABLE CREATION
'''
staging_table_creation: List[PostgresOperator] = []
create_staging_definitions: Dict[
str, str] = sql_queries_staging.create_query_destinition
# create the staging tables if they do not exist, else truncate
for object_name, query in create_staging_definitions.items():
staging_task: PostgresOperator = PostgresOperator(
start_date=datetime.datetime(2020, 1, 1),
schedule_interval=None,
tags=['udacity', 'etl'],
default_view="graph",
) as dag:
# generate dag documentation
dag.doc_md = __doc__
create_cluster = EmrCreateJobFlowOperator(
dag=dag,
task_id="create_cluster",
job_flow_overrides=job_flow_overrides,
aws_conn_id=aws_conn_id)
wait_cluster_completion = EmrJobFlowSensor(
task_id='wait_cluster_completion',
job_flow_id=cluster_id,
aws_conn_id=aws_conn_id,
target_states=["RUNNING", "WAITING"],
dag=dag)
terminate_cluster = EmrTerminateJobFlowOperator(
task_id="terminate_cluster",
trigger_rule="all_done",
job_flow_id=cluster_id,
aws_conn_id=aws_conn_id,
dag=dag)
with TaskGroup("run_immigration_mapping") as run_immigration_mapping:
add_step, wait_step = emr_step_task_group(
script_name='immigration_mapping',
cluster_id=cluster_id,
aws_conn_id=aws_conn_id,
dag=dag)
# - NLP Named Entity Recognition
# - Transform to star schema
aws_emr_etl_operator = AWSEMROperator(
task_id="create_EMR_cluster_and_execute_ETL",
dag=dag,
conn_id=AWS_CONN_ID,
redshift_conn_id=AWS_REDSHIFT_CONN_ID,
time_zone=local_tz,
cluster_identifier=
f"news-nlp-emr-{datetime.now(local_tz).strftime('%Y-%m-%d-%H-%M')}",
)
# Wait for the ETL process to finish
emr_etl_sensor = EmrJobFlowSensor(
task_id="sense_emr_etl",
dag=dag,
job_flow_id=
"{{ task_instance.xcom_pull('create_EMR_cluster_and_execute_ETL', key='return_value') }}",
aws_conn_id=AWS_CONN_ID,
)
# Create a Redshift cluster
create_redshift_cluster = AWSRedshiftOperator(
task_id="create_redshift_cluster",
dag=dag,
conn_id=AWS_CONN_ID,
redshift_conn_id=AWS_REDSHIFT_CONN_ID,
time_zone=local_tz,
cluster_identifier=
f"news-nlp-redshift-{datetime.now(local_tz).strftime('%Y-%m-%d-%H-%M')}",
)
# Wait for Redshift cluster to be ready
'TerminationProtected':
False,
},
'Steps': SPARK_STEPS,
'JobFlowRole': 'EMR_EC2_DefaultRole',
'ServiceRole': 'EMR_DefaultRole',
}
# [END howto_operator_emr_automatic_steps_config]
with DAG(
dag_id='emr_job_flow_automatic_steps_dag',
dagrun_timeout=timedelta(hours=2),
start_date=datetime(2021, 1, 1),
schedule_interval='0 3 * * *',
catchup=False,
tags=['example'],
) as dag:
# [START howto_operator_emr_automatic_steps_tasks]
job_flow_creator = EmrCreateJobFlowOperator(
task_id='create_job_flow',
job_flow_overrides=JOB_FLOW_OVERRIDES,
)
job_sensor = EmrJobFlowSensor(task_id='check_job_flow',
job_flow_id=job_flow_creator.output)
# [END howto_operator_emr_automatic_steps_tasks]
# Task dependency created via `XComArgs`:
# job_flow_creator >> job_sensor
# cerate emr cluster
instantiate_emr = CustomEmrCreateJobFlowOperator(
task_id='Wrangling_on_EMR',
aws_conn_id='aws_credentials',
emr_conn_id='emr_default',
dag=dag
)
s_job_flow_id = "{{ task_instance.xcom_pull(task_ids='Wrangling_on_EMR', "
s_job_flow_id += "key='return_value') }}"
monitor_emr = EmrJobFlowSensor(
task_id='Monitor_cluster',
job_flow_id=s_job_flow_id,
aws_conn_id='aws_credentials',
dag=dag
)
# update athena
athena_task_id = "Put_quotes_on_athena"
d_sqls = {
'subkey': 'quotes',
'createdb': SqlQueries.quotes_create_database,
'drop': SqlQueries.quotes_drop_table,
'create': SqlQueries.quotes_create_table,
'load': SqlQueries.quotes_rapair,
'load2': SqlQueries.quotes_load
}