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