def test_publish(self, mock_hook):
operator = PubSubPublishOperator(task_id=TASK_ID,
project=TEST_PROJECT,
topic=TEST_TOPIC,
messages=TEST_MESSAGES)
operator.execute(None)
mock_hook.return_value.publish.assert_called_once_with(
TEST_PROJECT, TEST_TOPIC, TEST_MESSAGES)
def test_publish(self, mock_hook):
operator = PubSubPublishOperator(task_id=TASK_ID,
project=TEST_PROJECT,
topic=TEST_TOPIC,
messages=TEST_MESSAGES)
operator.execute(None)
mock_hook.return_value.publish.assert_called_once_with(
TEST_PROJECT, TEST_TOPIC, TEST_MESSAGES)
def trigger_export_operator(export_name: str) -> PubSubPublishOperator:
# TODO(#4593) Migrate to IAPHTTPOperator
return PubSubPublishOperator(
task_id=f"trigger_{export_name.lower()}_bq_metric_export",
project=project_id,
topic="v1.export.view.data",
messages=[{
"data": b64encode(bytes(export_name, "utf-8")).decode()
}],
)
'email_on_failure': False,
'email_on_retry': False,
'project': project,
'topic': topic,
'subscription': subscription,
}
echo_template = '''
{% for m in task_instance.xcom_pull(task_ids='pull-messages') %}
echo "AckID: {{ m.get('ackId') }}, Base64-Encoded: {{ m.get('message') }}"
{% endfor %}
'''
with DAG('pubsub-end-to-end', default_args=default_args,
schedule_interval=datetime.timedelta(days=1)) as dag:
t1 = PubSubTopicCreateOperator(task_id='create-topic')
t2 = PubSubSubscriptionCreateOperator(
task_id='create-subscription', topic_project=project,
subscription=subscription)
t3 = PubSubPublishOperator(
task_id='publish-messages', messages=messages)
t4 = PubSubPullSensor(task_id='pull-messages', ack_messages=True)
t5 = BashOperator(task_id='echo-pulled-messages',
bash_command=echo_template)
t6 = PubSubSubscriptionDeleteOperator(task_id='delete-subscription')
t7 = PubSubTopicDeleteOperator(task_id='delete-topic')
t1 >> t2 >> t3
t2 >> t4 >> t5 >> t6 >> t7
# FUENTE: CRONTRAB: https://crontab.guru/
#############################################################
t_begin = DummyOperator(task_id="begin")
task_python = PythonOperator(task_id='task_python',
provide_context=True,
python_callable=python_func)
task_pubsub1 = PubSubPublishOperator(
task_id='task_pubsub1',
project=project,
topic="test-poc-01-3129e6b7-dd62-42cd",
messages=[{
'data':
base64.b64encode("Mensaje custom 01".encode('utf-8')).decode(),
'attributes': {
'atrib01': "ATRIBUTO1",
'atrib02': "ATRIBUTO2"
}
}])
task_pubsub2 = PubSubPublishOperator(
task_id='task_pubsub2',
project=project,
topic="test-poc-01-3129e6b7-dd62-42cd",
messages=[{
'data': "{{task_instance.xcom_pull(key='mensaje_param')}}"
}])
t_end = DummyOperator(task_id="end")
# I.e. if there is no fresh data.
bq_check_data_op = BigQueryCheckOperator(
task_id="bq_check_data_task",
use_legacy_sql=False,
sql=check_sql,
)
CHECK_ERROR_MESSAGE = b64e(b'Error. Did not retrain on stale data.')
# Task will only trigger if all upstream tasks fail. In which case a
# message will be sent to the Pub/Sub topic specified above.
publish_if_failed_check_op = PubSubPublishOperator(
task_id="publish_on_failed_check_task",
project=PROJECT_ID,
topic=TOPIC,
messages=[{'data': CHECK_ERROR_MESSAGE.decode()}],
trigger_rule=TriggerRule.ALL_FAILED
)
# Base query to extract training and validation datasets from public
# BigQuery dataset.
bql = """
SELECT
(IFNULL(tolls,0) + fare) AS fare_amount,
EXTRACT(DAYOFWEEK FROM trip_start_timestamp) AS dayofweek,
EXTRACT(HOUR FROM trip_start_timestamp) AS hourofday,
pickup_longitude AS pickuplon,
pickup_latitude AS pickuplat,
dropoff_longitude AS dropofflon,
dropoff_latitude AS dropofflat,
with models.DAG(dag_id="{}_calculation_pipeline_dag".format(project_id),
default_args=default_args,
schedule_interval=None) as dag:
if config_file is None:
raise Exception('Configuration file not specified')
with open(config_file) as f:
pipeline_yaml_dicts = yaml.full_load(f)
if pipeline_yaml_dicts:
pipeline_dict = pipelines_by_state(
pipeline_yaml_dicts['daily_pipelines'])
# TODO(#4593) Migrate to IAPHTTPOperator
covid_export = PubSubPublishOperator(
task_id='trigger_COVID_bq_metric_export',
project=project_id,
topic='v1.export.view.data',
messages=[{
'data': b64encode(b'COVID_DASHBOARD').decode()
}],
)
dataflow_default_args = {
'project':
project_id,
'region':
'us-west1',
'zone':
'us-west1-c',
'tempLocation':
'gs://{}-dataflow-templates/staging/'.format(project_id)
}
for state_code, state_pipelines in pipeline_dict.items():
# TODO(#4593) Migrate to IAPHTTPOperator
