diff --git a/airflow/providers/google/cloud/hooks/cloud_run.py b/airflow/providers/google/cloud/hooks/cloud_run.py new file mode 100644 index 0000000000000..5ddee95b23fe4 --- /dev/null +++ b/airflow/providers/google/cloud/hooks/cloud_run.py @@ -0,0 +1,245 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""This module contains a Google Cloud Run Hook.""" +from __future__ import annotations + +import json +import time +from typing import Any, Callable, Dict, List, Sequence, Union, cast + +from google.api_core.client_options import ClientOptions +from googleapiclient.discovery import build + +from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +from airflow.utils.log.logging_mixin import LoggingMixin + + +class CloudRunJobSteps: + """ + Helper class with Cloud Run job status. + Reference: https://cloud.google.com/run/docs/reference/rest/v1/namespaces.jobs#JobStatus + """ + + EXEC_STEP_COMPLETED = "Completed" + EXEC_STEP_RESOURCES_AVAILABLE = "ResourcesAvailable" + EXEC_STEP_STARTED = "Started" + AWAITING_STEPS = {EXEC_STEP_STARTED, EXEC_STEP_RESOURCES_AVAILABLE} + ALL_STEPS = [EXEC_STEP_COMPLETED, EXEC_STEP_RESOURCES_AVAILABLE, EXEC_STEP_STARTED] + + +class _CloudRunJobExecutionController(LoggingMixin): + """ + Interface for communication with Google API. + + :param cloud_run: Discovery resource + :param project_id: The Google Cloud Project ID. + :param execution_id: ID of a Cloud Run Job execution. + :param poll_sleep: The status refresh rate for pending operations. + :param num_retries: Maximum number of retries in case of connection problems. + :param wait_until_finished: If True, wait for the end of pipeline execution + before exiting. If False, it only submits job and check once if + the job is not in terminal state. + """ + + def __init__( + self, + cloud_run: Any, + project_id: str, + region: str, + wait_until_finished: bool, + execution_id: str | None = None, + poll_sleep: float = 10.0, + num_retries: int = 0, + job_exec_cold_start: float = 10.0, + ) -> None: + + super().__init__() + self._cloud_run = cloud_run + self._project_id = project_id + self._region = region + self._exec_id = execution_id + self._poll_sleep = poll_sleep + self._num_retries = num_retries + self._execution: dict | None = None + self._execution_state: str | None = None + self._job_exec_cold_start = job_exec_cold_start + self._wait_until_finished = wait_until_finished + + def _fetch_execution_by_id(self): + """ + Helper method to fetch the execution with the specified execution ID. + + :return: the Cloud Run job execution + """ + self.log.debug("Fetching information for job execution %s", self._exec_id) + return ( + self._cloud_run.namespaces() + .executions() + .get(name=f"namespaces/{self._project_id}/executions/{self._exec_id}") + .execute(num_retries=self._num_retries) + ) + + def _check_execution_state(self) -> bool: + """ + Helper method to check the state of job execution + if execution failed raise exception + + :return: True if execution is done. + :raise: Exception + """ + logs_uri = ( + f"https://console.cloud.google.com/run/jobs/executions" + f"/details/{self._region}/{self._exec_id}/logs?project={self._project_id}" + ) + job_exec_exception = Exception( + f"An error occurred when starting Google Cloud Run job execution {self._exec_id}." + f"\nSee details at {logs_uri} " + ) + if not isinstance(self._execution, dict): + raise job_exec_exception + execution = cast(Dict[str, Union[str, dict, int]], self._execution) + if "status" not in execution.keys(): + raise job_exec_exception + status = cast(Dict[str, Union[int, str, List[dict]]], execution["status"]) + if "conditions" not in status.keys(): + raise job_exec_exception + conditions = cast(List[Dict[str, str]], status["conditions"]) + steps = [] + for c in conditions: + step = c["type"] + verdict = c["status"] + if step not in CloudRunJobSteps.ALL_STEPS: + raise Exception( + f"Unknown state {step} found for Google Cloud Run job execution {self._exec_id} " + f"See details at {logs_uri}." + ) + if verdict == "False": + # This verdict means that a step failed, + # therefore the execution has failed + raise Exception( + f"Cloud Run Job execution {self._exec_id} has failed. See details at {logs_uri} " + ) + if verdict == "Unknown": + # This verdict means that the step is not completed yet, + # therefore the execution cannot be finished yet + return False + if verdict == "True": + steps.append(step) + return CloudRunJobSteps.EXEC_STEP_COMPLETED in steps + + def wait_for_done(self) -> None: + """Helper method to wait for result of job execution.""" + # Wait a few seconds for the job execution status to be available + # otherwise it fails + time.sleep(self._job_exec_cold_start) + self._execution = self._fetch_execution_by_id() + if self._wait_until_finished: + self.log.info("Starting to poll status for Google Cloud Run job execution %s", self._exec_id) + self.log.info(json.dumps(self._execution)) + while not self._check_execution_state(): + self.log.info("Waiting for execution completion. Sleeping %d seconds", self._poll_sleep) + time.sleep(self._poll_sleep) + self._execution = self._fetch_execution_by_id() + else: + # Check execution state only once to ensure the job has started + # Otherwise, _check_execution_state method would have raised an error + self._check_execution_state() + return + + +class CloudRunJobHook(GoogleBaseHook): + """ + Hook for Google Cloud Run. + + All the methods in the hook where project_id is used must be called with + keyword arguments rather than positional. + :param gcp_conn_id: The Airflow connection used for GCP credentials. + :param wait_until_finished: If True, wait for the end of pipeline + execution before exiting. If False, it only submits job + and check once is job not in terminal state. + """ + + DEFAULT_CLOUD_RUN_REGION = "us-central1" + + def __init__( + self, + gcp_conn_id: str = "google_cloud_default", + region: str = DEFAULT_CLOUD_RUN_REGION, + delegate_to: str | None = None, + impersonation_chain: str | Sequence[str] | None = None, + delete_timeout: int | None = 5 * 60, + wait_until_finished: bool = False, + ) -> None: + self.region = region + self.delete_timeout = delete_timeout + self.wait_until_finished = wait_until_finished + self.job_id: str | None = None + super().__init__( + gcp_conn_id=gcp_conn_id, delegate_to=delegate_to, impersonation_chain=impersonation_chain + ) + + def get_conn(self) -> build: + """Returns a Google Cloud Run service object.""" + http_authorized = self._authorize() + + # Use a regional endpoint since job execution is not possible from the global endpoint + client_options = ClientOptions(api_endpoint=f"https://{self.region}-run.googleapis.com") + return build("run", "v1", client_options=client_options, http=http_authorized, cache_discovery=False) + + @GoogleBaseHook.fallback_to_default_project_id + def execute_cloud_run_job( + self, project_id: str, job_name: str, on_new_execution_callback: Callable[[str], None] | None = None + ) -> dict: + """ + Executes a Cloud Run job. + + :param project_id: Optional, the Google Cloud project ID in which + to start a job. If set to None or missing, the default + project_id from the Google Cloud connection is used. + :param job_name: The name of the job + :param on_new_execution_callback: A callback that is called + when a job execution is detected. + """ + service = self.get_conn() + + request = service.namespaces().jobs().run(name=f"namespaces/{project_id}/jobs/{job_name}") + + execution = request.execute(num_retries=self.num_retries) + exec_id = execution["metadata"]["name"] + self.log.info("Requested job execution with response:\n%s", json.dumps(execution)) + if on_new_execution_callback: + on_new_execution_callback(execution) + + execution_controller = _CloudRunJobExecutionController( + cloud_run=service, + project_id=project_id, + region=self.region, + execution_id=exec_id, + num_retries=self.num_retries, + wait_until_finished=self.wait_until_finished, + ) + self.log.info( + "Job execution details available at " + "https://console.cloud.google.com/run/jobs/executions/details/%s/%s?project=%s ", + self.region, + exec_id, + self.project_id, + ) + + execution_controller.wait_for_done() + return execution diff --git a/airflow/providers/google/cloud/links/cloud_run.py b/airflow/providers/google/cloud/links/cloud_run.py new file mode 100644 index 0000000000000..94390b81b50dd --- /dev/null +++ b/airflow/providers/google/cloud/links/cloud_run.py @@ -0,0 +1,77 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""This module contains Google Cloud Run links.""" +from __future__ import annotations + +from typing import TYPE_CHECKING + +from airflow.models import BaseOperator +from airflow.providers.google.cloud.links.base import BaseGoogleLink + +if TYPE_CHECKING: + from airflow.utils.context import Context + +CLOUD_RUN_JOBS_BASE_LINK = "/run/jobs" +CLOUD_RUN_JOB_EXEC_LINK = ( + CLOUD_RUN_JOBS_BASE_LINK + "/executions/details/{region}/{exec_id}/general?project={project_id}" +) +CLOUD_RUN_JOB_LINK = "/details/{region}/{job_name}/executions?project={project_id}" + + +class CloudRunJobLink(BaseGoogleLink): + """Helper class for constructing Cloud Run job Link""" + + name = "Cloud Run job" + key = "cloud_run_job_config" + format_str = CLOUD_RUN_JOB_LINK + + @staticmethod + def persist( + operator_instance: BaseOperator, + context: Context, + project_id: str | None, + region: str | None, + job_name: str | None, + ): + operator_instance.xcom_push( + context, + key=CloudRunJobLink.key, + value={"project_id": project_id, "location": region, "job_name": job_name}, + ) + + +class CloudRunJobExecutionLink(BaseGoogleLink): + """Helper class for constructing Cloud Run job execution Link""" + + name = "Cloud Run job execution" + key = "cloud_run_job_execution" + format_str = CLOUD_RUN_JOB_EXEC_LINK + + @staticmethod + def persist( + operator_instance: BaseOperator, + context: Context, + project_id: str | None, + region: str | None, + execution_id: str | None, + ): + operator_instance.xcom_push( + context, + key=CloudRunJobExecutionLink.key, + value={"project_id": project_id, "location": region, "execution_id": execution_id}, + ) diff --git a/airflow/providers/google/cloud/operators/cloud_run.py b/airflow/providers/google/cloud/operators/cloud_run.py new file mode 100644 index 0000000000000..d35f5860b5e75 --- /dev/null +++ b/airflow/providers/google/cloud/operators/cloud_run.py @@ -0,0 +1,115 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""This module contains Google Cloud Run Jobs operators.""" +from __future__ import annotations + +from typing import TYPE_CHECKING, Sequence + +from airflow.models import BaseOperator +from airflow.providers.google.cloud.hooks.cloud_run import CloudRunJobHook +from airflow.providers.google.cloud.links.cloud_run import CloudRunJobExecutionLink + +if TYPE_CHECKING: + from airflow.utils.context import Context + + +class CloudRunExecuteJobOperator(BaseOperator): + """ + Executes an existing Cloud Run job. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:CloudRunExecuteJobOperator` + + :param job_name: The name of the cloud run job to execute + :param region: The region of the Cloud Run job (for example europe-west1) + :param project_id: The ID of the GCP project that owns the job. + If set to ``None`` or missing, the default project_id + from the GCP connection is used. + :param gcp_conn_id: The connection ID to use to connect to Google Cloud. + :param delegate_to: The account to impersonate, if any. + For this to work, the service account making the request + must have domain-wide delegation enabled. + :param wait_until_finished: If True, wait for the end of job execution + before exiting. If False (default), only submits job. + :param impersonation_chain: Optional service account to impersonate + using short-term credentials, or chained list of accounts required + to get the access_token of the last account in the list, + which will be impersonated in the request. If set as a string, the + account must grant the originating account the Service Account Token + Creator IAM role. If set as a sequence, the identities from the list + must grant Service Account Token Creator IAM role to the directly + preceding identity, with first account from the list granting this + role to the originating account (templated). + """ + + template_fields: Sequence[str] = ("job_name", "region", "project_id", "gcp_conn_id") + operator_extra_links = (CloudRunJobExecutionLink(),) + + def __init__( + self, + job_name: str, + region: str, + project_id: str | None = None, + gcp_conn_id: str = "google_cloud_default", + delegate_to: str | None = None, + cancel_timeout: int | None = 10 * 60, + wait_until_finished: bool = False, + impersonation_chain: str | Sequence[str] | None = None, + *args, + **kwargs, + ) -> None: + super().__init__(*args, **kwargs) + self.job_name = job_name + self.region = region + self.project_id = project_id + self.gcp_conn_id = gcp_conn_id + self.delegate_to = delegate_to + self.cancel_timeout = cancel_timeout + self.wait_until_finished = wait_until_finished + self.job = None + self.hook: CloudRunJobHook | None = None + self.impersonation_chain = impersonation_chain + + def execute(self, context: Context): + + self.hook = CloudRunJobHook( + region=self.region, + gcp_conn_id=self.gcp_conn_id, + delegate_to=self.delegate_to, + wait_until_finished=self.wait_until_finished, + impersonation_chain=self.impersonation_chain, + ) + + def set_current_execution(current_execution): + self.execution = current_execution + CloudRunJobExecutionLink.persist( + operator_instance=self, + context=context, + project_id=self.project_id, + region=self.region, + execution_id=self.execution.get("metadata").get("name"), + ) + + cloud_run_job_execution = self.hook.execute_cloud_run_job( + project_id=self.project_id, + job_name=self.job_name, + on_new_execution_callback=set_current_execution, + ) + + return cloud_run_job_execution diff --git a/airflow/providers/google/provider.yaml b/airflow/providers/google/provider.yaml index 6a81e61b09662..1c6ec34429f76 100644 --- a/airflow/providers/google/provider.yaml +++ b/airflow/providers/google/provider.yaml @@ -224,6 +224,12 @@ integrations: - /docs/apache-airflow-providers-google/operators/cloud/pubsub.rst logo: /integration-logos/gcp/Cloud-PubSub.png tags: [gcp] + - integration-name: Google Cloud Run + external-doc-url: https://cloud.google.com/run/ + how-to-guide: + - /docs/apache-airflow-providers-google/operators/cloud/cloud_run.rst + logo: /integration-logos/gcp/Cloud-Run.png + tags: [gcp] - integration-name: Google Cloud Secret Manager external-doc-url: https://cloud.google.com/secret-manager/ logo: /integration-logos/gcp/Google-Cloud-Secret-Manager.png @@ -530,6 +536,9 @@ operators: - integration-name: Google Cloud Pub/Sub python-modules: - airflow.providers.google.cloud.operators.pubsub + - integration-name: Google Cloud Run + python-modules: + - airflow.providers.google.cloud.operators.cloud_run - integration-name: Google Cloud Spanner python-modules: - airflow.providers.google.cloud.operators.spanner @@ -750,6 +759,9 @@ hooks: - integration-name: Google Cloud Pub/Sub python-modules: - airflow.providers.google.cloud.hooks.pubsub + - integration-name: Google Cloud Run + python-modules: + - airflow.providers.google.cloud.hooks.cloud_run - integration-name: Google Cloud Secret Manager python-modules: - airflow.providers.google.cloud.hooks.secret_manager @@ -1087,6 +1099,8 @@ extra-links: - airflow.providers.google.cloud.links.kubernetes_engine.KubernetesEnginePodLink - airflow.providers.google.cloud.links.pubsub.PubSubSubscriptionLink - airflow.providers.google.cloud.links.pubsub.PubSubTopicLink + - airflow.providers.google.cloud.links.cloud_run.CloudRunJobLink + - airflow.providers.google.cloud.links.cloud_run.CloudRunJobExecutionLink - airflow.providers.google.cloud.links.cloud_memorystore.MemcachedInstanceDetailsLink - airflow.providers.google.cloud.links.cloud_memorystore.MemcachedInstanceListLink - airflow.providers.google.cloud.links.cloud_memorystore.RedisInstanceDetailsLink diff --git a/docs/apache-airflow-providers-google/operators/cloud/cloud_run.rst b/docs/apache-airflow-providers-google/operators/cloud/cloud_run.rst new file mode 100644 index 0000000000000..6bb8b16518f10 --- /dev/null +++ b/docs/apache-airflow-providers-google/operators/cloud/cloud_run.rst @@ -0,0 +1,55 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +Google CloudRun Operators +========================= + +`Cloud Run `__ is Google's fully managed, +Container As A Service solution. +It is a serverless compute platform service which allows users to +deploy applications written in any programming language. + +Apart from services, such as web applications exposing http endpoints, +Cloud Run provides creation and execution of jobs running to completion. + + +Prerequisite Tasks +------------------ + +.. include:: ../_partials/prerequisite_tasks.rst + +.. _howto/operator:CloudRunExecuteJobOperator: + +CloudRunExecuteJobOperator +-------------------------- + +Executes an existing Cloud Run Job. + +If ``wait_until_finished`` is set to ``True``, +``CloudRunExecuteJobOperator`` will always wait for job completion. +If set to ``False`` only creates the job execution. + +For parameter definition, take a look at +:class:`~airflow.providers.google.cloud.operators.cloud_run.CloudRunExecuteJobOperator`. + + +Reference +--------- + +For further information, look at: +* `Product Documentation `__ diff --git a/docs/integration-logos/gcp/Cloud-Run.png b/docs/integration-logos/gcp/Cloud-Run.png new file mode 100644 index 0000000000000..1dbec06c962cc Binary files /dev/null and b/docs/integration-logos/gcp/Cloud-Run.png differ diff --git a/tests/providers/google/cloud/hooks/test_cloud_run.py b/tests/providers/google/cloud/hooks/test_cloud_run.py new file mode 100644 index 0000000000000..63ca20d0204b6 --- /dev/null +++ b/tests/providers/google/cloud/hooks/test_cloud_run.py @@ -0,0 +1,71 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock +from unittest.mock import PropertyMock +from uuid import UUID + +import pytest + +from airflow.providers.google.cloud.hooks.cloud_run import CloudRunJobHook +from tests.providers.google.cloud.utils.base_gcp_mock import mock_base_gcp_hook_no_default_project_id + +TASK_ID = "test-cloud-run-operator" +JOB_NAME = "test-cloud-run-job" +CLOUD_RUN_JOBS_STRING = "airflow.providers.google.cloud.hooks.cloud_run.%s" +TEST_PROJECT = "test-project" +MOCK_UUID = UUID("cf4a56d2-8101-4217-b027-2af6216feb48") +MOCK_UUID_PREFIX = str(MOCK_UUID)[:8] +EXECUTION_ID = f"{JOB_NAME}-{MOCK_UUID_PREFIX}" + + +@mock.patch(CLOUD_RUN_JOBS_STRING % "_CloudRunJobExecutionController") +@mock.patch(CLOUD_RUN_JOBS_STRING % "CloudRunJobHook.get_conn") +@mock.patch( + "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__", + mock_base_gcp_hook_no_default_project_id, +) +@mock.patch( + "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id", + new_callable=PropertyMock, + return_value=None, +) +@pytest.mark.parametrize("wait_until_finished", [True, False]) +def test_execute_cloud_run_job(mock_hook, mock_conn, mock_controller, wait_until_finished): + execute_method = mock_conn.return_value.namespaces.return_value.jobs.return_value.run + execute_method.return_value.execute.return_value = { + "apiVersion": "run.googleapis.com/v1", + "kind": "Execution", + "metadata": {"name": EXECUTION_ID}, + } + + hook = CloudRunJobHook(gcp_conn_id="google_cloud_default", wait_until_finished=wait_until_finished) + hook.execute_cloud_run_job(job_name=JOB_NAME, project_id=TEST_PROJECT) + execute_method.assert_called_once_with(name=f"namespaces/{TEST_PROJECT}/jobs/{JOB_NAME}") + + mock_controller.assert_called_once_with( + cloud_run=mock_conn.return_value, + project_id=TEST_PROJECT, + execution_id=EXECUTION_ID, + num_retries=hook.num_retries, + region=hook.DEFAULT_CLOUD_RUN_REGION, + wait_until_finished=wait_until_finished, + ) + + mock_controller.return_value.wait_for_done.assert_called_once() diff --git a/tests/providers/google/cloud/operators/test_cloud_run.py b/tests/providers/google/cloud/operators/test_cloud_run.py new file mode 100644 index 0000000000000..6ce7f30b1de25 --- /dev/null +++ b/tests/providers/google/cloud/operators/test_cloud_run.py @@ -0,0 +1,53 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest + +from airflow.providers.google.cloud.operators.cloud_run import CloudRunExecuteJobOperator + +TASK_ID = "test-cloud-run-operator" +JOB_NAME = "test-cloud-run-job" +REGION = "us-central1" +TEST_PROJECT = "test-project" + + +@pytest.fixture +def operator(): + return CloudRunExecuteJobOperator( + task_id="execute_cloud_run_job_test", job_name=JOB_NAME, region=REGION, project_id=TEST_PROJECT + ) + + +@mock.patch("airflow.providers.google.cloud.operators.cloud_run.CloudRunJobHook") +def test_execute(mock_cloud_run, operator): + operator.execute(mock.MagicMock()) + mock_cloud_run.assert_called_once_with( + gcp_conn_id="google_cloud_default", + region=REGION, + delegate_to=None, + wait_until_finished=False, + impersonation_chain=None, + ) + mock_cloud_run.return_value.execute_cloud_run_job.assert_called_once_with( + project_id=TEST_PROJECT, + job_name=JOB_NAME, + on_new_execution_callback=mock.ANY, + ) diff --git a/tests/system/providers/google/cloud/cloud_run/__init__.py b/tests/system/providers/google/cloud/cloud_run/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/tests/system/providers/google/cloud/cloud_run/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/tests/system/providers/google/cloud/cloud_run/example_cloud_run.py b/tests/system/providers/google/cloud/cloud_run/example_cloud_run.py new file mode 100644 index 0000000000000..b253c4a0edfdd --- /dev/null +++ b/tests/system/providers/google/cloud/cloud_run/example_cloud_run.py @@ -0,0 +1,49 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Example Airflow DAG for Google Cloud Run service. +""" +from __future__ import annotations + +import os +from datetime import datetime + +from airflow import models +from airflow.providers.google.cloud.operators.cloud_run import CloudRunExecuteJobOperator + +ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID") +DAG_ID = "cloud_run" + + +with models.DAG( + DAG_ID, + schedule="@once", + start_date=datetime(2021, 1, 1), + catchup=False, + tags=["example", "cloud_run"], +) as dag: + # [START howto_operator_cloud_run_execute_job] + call_cloud_run = CloudRunExecuteJobOperator( + task_id="cloud_run_execute_job", job_name="my-job", region="us-central1", wait_until_finished=True + ) + # [END howto_operator_cloud_run_execute_job] + +from tests.system.utils import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag)