def execute(self, context: 'Context'):
hook = WorkflowsHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
self.log.info("Retrieving execution %s for workflow %s", self.execution_id, self.workflow_id)
execution = hook.get_execution(
workflow_id=self.workflow_id,
execution_id=self.execution_id,
location=self.location,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return Execution.to_dict(execution)
def poke(self, context):
hook = WorkflowsHook(gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain)
self.log.info("Checking state of execution %s for workflow %s",
self.execution_id, self.workflow_id)
execution: Execution = hook.get_execution(
workflow_id=self.workflow_id,
execution_id=self.execution_id,
location=self.location,
project_id=self.project_id,
retry=self.retry,
timeout=self.request_timeout,
metadata=self.metadata,
)
state = execution.state
if state in self.failure_states:
raise AirflowException(
f"Execution {self.execution_id} for workflow {self.execution_id} "
f"failed and is in `{state}` state", )
if state in self.success_states:
self.log.info(
"Execution %s for workflow %s completed with state: %s",
self.execution_id,
self.workflow_id,
state,
)
return True
self.log.info(
"Execution %s for workflow %s does not completed yet, current state: %s",
self.execution_id,
self.workflow_id,
state,
)
return False
class TestWorkflowsHook:
def setup_method(self, _):
with mock.patch(BASE_PATH.format("GoogleBaseHook.__init__"),
new=mock_init):
self.hook = WorkflowsHook(gcp_conn_id="test") # pylint: disable=attribute-defined-outside-init
@mock.patch(BASE_PATH.format("WorkflowsHook._get_credentials"))
@mock.patch(BASE_PATH.format("WorkflowsHook.client_info"),
new_callable=mock.PropertyMock)
@mock.patch(BASE_PATH.format("WorkflowsClient"))
def test_get_workflows_client(self, mock_client, mock_client_info,
mock_get_credentials):
self.hook.get_workflows_client()
mock_client.assert_called_once_with(
credentials=mock_get_credentials.return_value,
client_info=mock_client_info.return_value,
)
@mock.patch(BASE_PATH.format("WorkflowsHook._get_credentials"))
@mock.patch(BASE_PATH.format("WorkflowsHook.client_info"),
new_callable=mock.PropertyMock)
@mock.patch(BASE_PATH.format("ExecutionsClient"))
def test_get_executions_client(self, mock_client, mock_client_info,
mock_get_credentials):
self.hook.get_executions_client()
mock_client.assert_called_once_with(
credentials=mock_get_credentials.return_value,
client_info=mock_client_info.return_value,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_workflows_client"))
def test_create_workflow(self, mock_client):
result = self.hook.create_workflow(
workflow=WORKFLOW,
workflow_id=WORKFLOW_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.create_workflow.return_value == result
mock_client.return_value.create_workflow.assert_called_once_with(
request=dict(workflow=WORKFLOW,
workflow_id=WORKFLOW_ID,
parent=WORKFLOW_PARENT),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_workflows_client"))
def test_get_workflow(self, mock_client):
result = self.hook.get_workflow(
workflow_id=WORKFLOW_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.get_workflow.return_value == result
mock_client.return_value.get_workflow.assert_called_once_with(
request=dict(name=WORKFLOW_NAME),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_workflows_client"))
def test_update_workflow(self, mock_client):
result = self.hook.update_workflow(
workflow=WORKFLOW,
update_mask=UPDATE_MASK,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.update_workflow.return_value == result
mock_client.return_value.update_workflow.assert_called_once_with(
request=dict(
workflow=WORKFLOW,
update_mask=UPDATE_MASK,
),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_workflows_client"))
def test_delete_workflow(self, mock_client):
result = self.hook.delete_workflow(
workflow_id=WORKFLOW_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.delete_workflow.return_value == result
mock_client.return_value.delete_workflow.assert_called_once_with(
request=dict(name=WORKFLOW_NAME),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_workflows_client"))
def test_list_workflows(self, mock_client):
result = self.hook.list_workflows(
location=LOCATION,
project_id=PROJECT_ID,
filter_=FILTER_,
order_by=ORDER_BY,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.list_workflows.return_value == result
mock_client.return_value.list_workflows.assert_called_once_with(
request=dict(
parent=WORKFLOW_PARENT,
filter=FILTER_,
order_by=ORDER_BY,
),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_executions_client"))
def test_create_execution(self, mock_client):
result = self.hook.create_execution(
workflow_id=WORKFLOW_ID,
location=LOCATION,
project_id=PROJECT_ID,
execution=EXECUTION,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.create_execution.return_value == result
mock_client.return_value.create_execution.assert_called_once_with(
request=dict(
parent=EXECUTION_PARENT,
execution=EXECUTION,
),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_executions_client"))
def test_get_execution(self, mock_client):
result = self.hook.get_execution(
workflow_id=WORKFLOW_ID,
execution_id=EXECUTION_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.get_execution.return_value == result
mock_client.return_value.get_execution.assert_called_once_with(
request=dict(name=EXECUTION_NAME),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_executions_client"))
def test_cancel_execution(self, mock_client):
result = self.hook.cancel_execution(
workflow_id=WORKFLOW_ID,
execution_id=EXECUTION_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.cancel_execution.return_value == result
mock_client.return_value.cancel_execution.assert_called_once_with(
request=dict(name=EXECUTION_NAME),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
@mock.patch(BASE_PATH.format("WorkflowsHook.get_executions_client"))
def test_list_execution(self, mock_client):
result = self.hook.list_executions(
workflow_id=WORKFLOW_ID,
location=LOCATION,
project_id=PROJECT_ID,
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)
assert mock_client.return_value.list_executions.return_value == result
mock_client.return_value.list_executions.assert_called_once_with(
request=dict(parent=EXECUTION_PARENT),
retry=RETRY,
timeout=TIMEOUT,
metadata=METADATA,
)