def execute(self, context: 'Context'):
hook = WorkflowsHook(gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain)
workflow_id = self._workflow_id(context)
self.log.info("Creating workflow")
try:
operation = hook.create_workflow(
workflow=self.workflow,
workflow_id=workflow_id,
location=self.location,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
workflow = operation.result()
except AlreadyExists:
workflow = hook.get_workflow(
workflow_id=workflow_id,
location=self.location,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return Workflow.to_dict(workflow)
def execute(self, context: 'Context'):
hook = WorkflowsHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
workflow = hook.get_workflow(
workflow_id=self.workflow_id,
project_id=self.project_id,
location=self.location,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
self.log.info("Updating workflow")
operation = hook.update_workflow(
workflow=workflow,
update_mask=self.update_mask,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
workflow = operation.result()
WorkflowsWorkflowDetailsLink.persist(
context=context,
task_instance=self,
location_id=self.location,
workflow_id=self.workflow_id,
project_id=self.project_id or hook.project_id,
)
return Workflow.to_dict(workflow)
def execute(self, context: 'Context'):
hook = WorkflowsHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
self.log.info("Retrieving workflow")
workflow = hook.get_workflow(
workflow_id=self.workflow_id,
location=self.location,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
return Workflow.to_dict(workflow)
def execute(self, context):
hook = WorkflowsHook(gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain)
workflow = hook.get_workflow(
workflow_id=self.workflow_id,
project_id=self.project_id,
location=self.location,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
self.log.info("Updating workflow")
operation = hook.update_workflow(
workflow=workflow,
update_mask=self.update_mask,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
workflow = operation.result()
return Workflow.to_dict(workflow)
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,
)