def test_execute(self, mock_hook):
op = CloudDatastoreRollbackOperator(
task_id="test_task",
gcp_conn_id=CONN_ID,
project_id=PROJECT_ID,
transaction=TRANSACTION,
)
op.execute({})
mock_hook.assert_called_once_with(gcp_conn_id=CONN_ID)
mock_hook.return_value.rollback.assert_called_once_with(
project_id=PROJECT_ID, transaction=TRANSACTION)
},
"query": {},
}
# [END how_to_query_def]
# [START how_to_run_query]
run_query = CloudDatastoreRunQueryOperator(task_id="run_query",
body=QUERY,
project_id=GCP_PROJECT_ID)
# [END how_to_run_query]
allocate_ids >> begin_transaction_query
begin_transaction_to_rollback = CloudDatastoreBeginTransactionOperator(
task_id="begin_transaction_to_rollback",
transaction_options=TRANSACTION_OPTIONS,
project_id=GCP_PROJECT_ID,
)
# [START how_to_rollback_transaction]
rollback_transaction = CloudDatastoreRollbackOperator(
task_id="rollback_transaction",
transaction=begin_transaction_to_rollback.output,
)
# [END how_to_rollback_transaction]
# Task dependencies created via `XComArgs`:
# begin_transaction_commit >> commit_task
# begin_transaction_to_rollback >> rollback_transaction
# begin_transaction_query >> run_query
commit_task = CloudDatastoreCommitOperator(
task_id="commit_task", body=COMMIT_BODY, project_id=GCP_PROJECT_ID
)
# [END how_to_commit_task]
allocate_ids >> begin_transaction_commit >> commit_task
begin_transaction_query = CloudDatastoreBeginTransactionOperator(
task_id="begin_transaction_query", transaction_options=TRANSACTION_OPTIONS, project_id=GCP_PROJECT_ID,
)
# [START how_to_run_query]
run_query = CloudDatastoreRunQueryOperator(task_id="run_query", body=QUERY, project_id=GCP_PROJECT_ID)
# [END how_to_run_query]
allocate_ids >> begin_transaction_query >> run_query
begin_transaction_to_rollback = CloudDatastoreBeginTransactionOperator(
task_id="begin_transaction_to_rollback",
transaction_options=TRANSACTION_OPTIONS,
project_id=GCP_PROJECT_ID,
)
# [START how_to_rollback_transaction]
rollback_transaction = CloudDatastoreRollbackOperator(
task_id="rollback_transaction",
transaction="{{ task_instance.xcom_pull('begin_transaction_to_rollback') }}",
)
begin_transaction_to_rollback >> rollback_transaction
# [END how_to_rollback_transaction]