From a1ee16847befa46ef1820bbf301179fd9666bb1a Mon Sep 17 00:00:00 2001 From: "d.smirnov" Date: Tue, 14 Jul 2026 09:44:03 +0300 Subject: [PATCH] Add deferrable support for Cloud Functions invoke operator - Add CloudFunctionInvokeFunctionTrigger for async invocation via asyncio.to_thread. - Add deferrable parameter to CloudFunctionInvokeFunctionOperator with execute_complete(). - Register trigger in provider.yaml and get_provider_info.py. - Add unit tests for trigger and operator deferrable path. Closes #68908 --- airflow-core/newsfragments/69853.feature.rst | 1 + providers/google/provider.yaml | 3 + .../google/cloud/operators/functions.py | 45 ++++++ .../google/cloud/triggers/cloud_functions.py | 109 ++++++++++++++ .../providers/google/get_provider_info.py | 4 + .../google/cloud/operators/test_functions.py | 74 ++++++++- .../cloud/triggers/test_cloud_functions.py | 142 ++++++++++++++++++ 7 files changed, 377 insertions(+), 1 deletion(-) create mode 100644 airflow-core/newsfragments/69853.feature.rst create mode 100644 providers/google/src/airflow/providers/google/cloud/triggers/cloud_functions.py create mode 100644 providers/google/tests/unit/google/cloud/triggers/test_cloud_functions.py diff --git a/airflow-core/newsfragments/69853.feature.rst b/airflow-core/newsfragments/69853.feature.rst new file mode 100644 index 0000000000000..3978bf5d61771 --- /dev/null +++ b/airflow-core/newsfragments/69853.feature.rst @@ -0,0 +1 @@ +Add deferrable support for CloudFunctionInvokeFunctionOperator to invoke Google Cloud Functions asynchronously without blocking workers. diff --git a/providers/google/provider.yaml b/providers/google/provider.yaml index 3a9d2b7ad1a2d..5bfaf75c164f8 100644 --- a/providers/google/provider.yaml +++ b/providers/google/provider.yaml @@ -987,6 +987,9 @@ triggers: - integration-name: Google Cloud Generative AI python-modules: - airflow.providers.google.cloud.triggers.gen_ai + - integration-name: Google Cloud Functions + python-modules: + - airflow.providers.google.cloud.triggers.cloud_functions transfers: - source-integration-name: Presto diff --git a/providers/google/src/airflow/providers/google/cloud/operators/functions.py b/providers/google/src/airflow/providers/google/cloud/operators/functions.py index 988578f3bab5e..4bbd15c5cadf8 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/functions.py +++ b/providers/google/src/airflow/providers/google/cloud/operators/functions.py @@ -25,6 +25,7 @@ from googleapiclient.errors import HttpError +from airflow.configuration import conf from airflow.providers.common.compat.sdk import AirflowException from airflow.providers.google.cloud.hooks.functions import CloudFunctionsHook from airflow.providers.google.cloud.links.cloud_functions import ( @@ -32,6 +33,9 @@ CloudFunctionsListLink, ) from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator +from airflow.providers.google.cloud.triggers.cloud_functions import ( + CloudFunctionInvokeFunctionTrigger, +) from airflow.providers.google.cloud.utils.field_validator import ( GcpBodyFieldValidator, GcpFieldValidationException, @@ -433,6 +437,7 @@ class CloudFunctionInvokeFunctionOperator(GoogleCloudBaseOperator): 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). + :param deferrable: Run operator in the deferrable mode. :return: None """ @@ -456,6 +461,7 @@ def __init__( gcp_conn_id: str = "google_cloud_default", api_version: str = "v1", impersonation_chain: str | Sequence[str] | None = None, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), **kwargs, ) -> None: super().__init__(**kwargs) @@ -466,6 +472,7 @@ def __init__( self.gcp_conn_id = gcp_conn_id self.api_version = api_version self.impersonation_chain = impersonation_chain + self.deferrable = deferrable @property def extra_links_params(self) -> dict[str, Any]: @@ -474,12 +481,48 @@ def extra_links_params(self) -> dict[str, Any]: "function_name": self.function_id, } + def execute_complete(self, context: Context, event: dict): + """Handle trigger completion and process the result.""" + if event["status"] == "error": + raise AirflowException(event["message"]) + + result = event["result"] + execution_id = event.get("execution_id") + + self.log.info("Function called successfully. Execution id: %s", execution_id) + context["ti"].xcom_push(key="execution_id", value=execution_id) + + project_id = self.project_id + if project_id: + CloudFunctionsDetailsLink.persist( + context=context, + location=self.location, + project_id=project_id, + function_name=self.function_id, + ) + return result + def execute(self, context: Context): + if self.deferrable: + self.defer( + trigger=CloudFunctionInvokeFunctionTrigger( + function_id=self.function_id, + input_data=self.input_data, + location=self.location, + project_id=self.project_id, + gcp_conn_id=self.gcp_conn_id, + api_version=self.api_version, + impersonation_chain=self.impersonation_chain, + ), + method_name="execute_complete", + ) + hook = CloudFunctionsHook( api_version=self.api_version, gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, ) + self.log.info("Calling function %s.", self.function_id) result = hook.call_function( function_id=self.function_id, @@ -494,7 +537,9 @@ def execute(self, context: Context): if project_id: CloudFunctionsDetailsLink.persist( context=context, + location=self.location, project_id=project_id, + function_name=self.function_id, ) return result diff --git a/providers/google/src/airflow/providers/google/cloud/triggers/cloud_functions.py b/providers/google/src/airflow/providers/google/cloud/triggers/cloud_functions.py new file mode 100644 index 0000000000000..33b73b2af9f29 --- /dev/null +++ b/providers/google/src/airflow/providers/google/cloud/triggers/cloud_functions.py @@ -0,0 +1,109 @@ +# 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 + +import asyncio +from collections.abc import AsyncIterator, Sequence +from typing import Any + +from airflow.providers.common.compat.sdk import AirflowException +from airflow.providers.google.cloud.hooks.functions import CloudFunctionsHook +from airflow.triggers.base import BaseTrigger, TriggerEvent + + +class CloudFunctionInvokeFunctionTrigger(BaseTrigger): + """ + Trigger to invoke a Google Cloud Function and wait for its result. + + This trigger invokes the function via the synchronous Cloud Functions API + using ``asyncio.to_thread`` and yields a single TriggerEvent with the + result. + + :param function_id: ID of the function to be called. + :param input_data: Input to be passed to the function. + :param location: The location where the function is located. + :param project_id: Google Cloud Project ID where the function belongs. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param api_version: API version used (for example v1). + :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. + """ + + def __init__( + self, + function_id: str, + input_data: dict, + location: str, + project_id: str, + gcp_conn_id: str = "google_cloud_default", + api_version: str = "v1", + impersonation_chain: str | Sequence[str] | None = None, + ) -> None: + super().__init__() + self.function_id = function_id + self.input_data = input_data + self.location = location + self.project_id = project_id + self.gcp_conn_id = gcp_conn_id + self.api_version = api_version + self.impersonation_chain = impersonation_chain + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serialize class arguments and classpath.""" + return ( + "airflow.providers.google.cloud.triggers.cloud_functions.CloudFunctionInvokeFunctionTrigger", + { + "function_id": self.function_id, + "input_data": self.input_data, + "location": self.location, + "project_id": self.project_id, + "gcp_conn_id": self.gcp_conn_id, + "api_version": self.api_version, + "impersonation_chain": self.impersonation_chain, + }, + ) + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Invoke the Cloud Function and yield a single result event.""" + hook = CloudFunctionsHook( + api_version=self.api_version, + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + try: + result = await asyncio.to_thread( + hook.call_function, + function_id=self.function_id, + input_data=self.input_data, + location=self.location, + project_id=self.project_id, + ) + yield TriggerEvent( + { + "status": "success", + "result": result, + "execution_id": result.get("executionId"), + } + ) + except AirflowException as e: + self.log.error("Cloud Function invocation failed: %s", e) + yield TriggerEvent({"status": "error", "message": str(e)}) + except Exception as e: + self.log.exception("Unexpected error while invoking Cloud Function.") + yield TriggerEvent({"status": "error", "message": str(e)}) diff --git a/providers/google/src/airflow/providers/google/get_provider_info.py b/providers/google/src/airflow/providers/google/get_provider_info.py index 13feaa8d1e761..2e88547f5bb3d 100644 --- a/providers/google/src/airflow/providers/google/get_provider_info.py +++ b/providers/google/src/airflow/providers/google/get_provider_info.py @@ -1100,6 +1100,10 @@ def get_provider_info(): "integration-name": "Google Cloud Build", "python-modules": ["airflow.providers.google.cloud.triggers.cloud_build"], }, + { + "integration-name": "Google Cloud Functions", + "python-modules": ["airflow.providers.google.cloud.triggers.cloud_functions"], + }, { "integration-name": "Managed Service for Apache Airflow", "python-modules": ["airflow.providers.google.cloud.triggers.cloud_composer"], diff --git a/providers/google/tests/unit/google/cloud/operators/test_functions.py b/providers/google/tests/unit/google/cloud/operators/test_functions.py index 9a08d695c68b5..7045103792422 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_functions.py +++ b/providers/google/tests/unit/google/cloud/operators/test_functions.py @@ -24,7 +24,7 @@ import pytest from googleapiclient.errors import HttpError -from airflow.providers.common.compat.sdk import AirflowException +from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred from airflow.providers.google.cloud.operators.functions import ( CloudFunctionDeleteFunctionOperator, CloudFunctionDeployFunctionOperator, @@ -736,3 +736,75 @@ def test_execute(self, mock_gcf_hook): key="execution_id", value=exec_id, ) + + @mock.patch("airflow.providers.google.cloud.operators.functions.CloudFunctionInvokeFunctionTrigger") + def test_execute_deferrable(self, mock_trigger_class): + function_id = "test_function" + payload = {"key": "value"} + + op = CloudFunctionInvokeFunctionOperator( + task_id="test", + function_id=function_id, + input_data=payload, + location=GCP_LOCATION, + project_id=GCP_PROJECT_ID, + deferrable=True, + ) + mock_context = {"ti": mock.MagicMock()} + if not AIRFLOW_V_3_0_PLUS: + mock_context["task"] = op + + with pytest.raises(TaskDeferred) as ctx: + op.execute(mock_context) + + mock_trigger_class.assert_called_once_with( + function_id=function_id, + input_data=payload, + location=GCP_LOCATION, + project_id=GCP_PROJECT_ID, + gcp_conn_id="google_cloud_default", + api_version="v1", + impersonation_chain=None, + ) + assert ctx.value.trigger is not None + + def test_execute_complete_success(self): + op = CloudFunctionInvokeFunctionOperator( + task_id="test", + function_id="test_function", + input_data={"key": "value"}, + location=GCP_LOCATION, + project_id=GCP_PROJECT_ID, + ) + mock_ti = mock.MagicMock() + mock_context = {"ti": mock_ti} + if not AIRFLOW_V_3_0_PLUS: + mock_context["task"] = op + event = { + "status": "success", + "result": {"executionId": "exec-123", "result": "ok"}, + "execution_id": "exec-123", + } + + result = op.execute_complete(mock_context, event) + + assert result == {"executionId": "exec-123", "result": "ok"} + mock_ti.xcom_push.assert_any_call(key="execution_id", value="exec-123") + mock_ti.xcom_push.assert_any_call( + key="cloud_functions_details", + value={"location": GCP_LOCATION, "project_id": GCP_PROJECT_ID, "function_name": "test_function"}, + ) + + def test_execute_complete_error(self): + op = CloudFunctionInvokeFunctionOperator( + task_id="test", + function_id="test_function", + input_data={"key": "value"}, + location=GCP_LOCATION, + project_id=GCP_PROJECT_ID, + ) + mock_context = {"ti": mock.MagicMock()} + event = {"status": "error", "message": "Function failed"} + + with pytest.raises(AirflowException, match="Function failed"): + op.execute_complete(mock_context, event) diff --git a/providers/google/tests/unit/google/cloud/triggers/test_cloud_functions.py b/providers/google/tests/unit/google/cloud/triggers/test_cloud_functions.py new file mode 100644 index 0000000000000..a98927d3d208c --- /dev/null +++ b/providers/google/tests/unit/google/cloud/triggers/test_cloud_functions.py @@ -0,0 +1,142 @@ +# +# 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.common.compat.sdk import AirflowException +from airflow.providers.google.cloud.triggers.cloud_functions import CloudFunctionInvokeFunctionTrigger +from airflow.triggers.base import TriggerEvent + + +class TestCloudFunctionInvokeFunctionTrigger: + @pytest.fixture + def trigger(self): + return CloudFunctionInvokeFunctionTrigger( + function_id="test-function", + input_data={"key": "value"}, + location="us-central1", + project_id="test-project", + gcp_conn_id="google_cloud_default", + api_version="v1", + impersonation_chain=None, + ) + + def test_serialize(self, trigger): + classpath, kwargs = trigger.serialize() + assert ( + classpath + == "airflow.providers.google.cloud.triggers.cloud_functions.CloudFunctionInvokeFunctionTrigger" + ) + assert kwargs == { + "function_id": "test-function", + "input_data": {"key": "value"}, + "location": "us-central1", + "project_id": "test-project", + "gcp_conn_id": "google_cloud_default", + "api_version": "v1", + "impersonation_chain": None, + } + + @pytest.mark.asyncio + @mock.patch("airflow.providers.google.cloud.triggers.cloud_functions.CloudFunctionsHook") + async def test_run_success(self, mock_hook_class, trigger): + mock_hook = mock_hook_class.return_value + mock_hook.call_function.return_value = {"executionId": "exec-123", "result": "ok"} + + events = [] + async for event in trigger.run(): + events.append(event) + + assert len(events) == 1 + event = events[0] + assert isinstance(event, TriggerEvent) + assert event.payload["status"] == "success" + assert event.payload["result"] == {"executionId": "exec-123", "result": "ok"} + assert event.payload["execution_id"] == "exec-123" + + mock_hook_class.assert_called_once_with( + api_version="v1", + gcp_conn_id="google_cloud_default", + impersonation_chain=None, + ) + mock_hook.call_function.assert_called_once_with( + function_id="test-function", + input_data={"key": "value"}, + location="us-central1", + project_id="test-project", + ) + + @pytest.mark.asyncio + @mock.patch("airflow.providers.google.cloud.triggers.cloud_functions.CloudFunctionsHook") + async def test_run_airflow_exception(self, mock_hook_class, trigger): + mock_hook = mock_hook_class.return_value + mock_hook.call_function.side_effect = AirflowException("Function error") + + events = [] + async for event in trigger.run(): + events.append(event) + + assert len(events) == 1 + event = events[0] + assert event.payload["status"] == "error" + assert event.payload["message"] == "Function error" + + @pytest.mark.asyncio + @mock.patch("airflow.providers.google.cloud.triggers.cloud_functions.CloudFunctionsHook") + async def test_run_generic_exception(self, mock_hook_class, trigger): + mock_hook = mock_hook_class.return_value + mock_hook.call_function.side_effect = RuntimeError("Unexpected error") + + events = [] + async for event in trigger.run(): + events.append(event) + + assert len(events) == 1 + event = events[0] + assert event.payload["status"] == "error" + assert event.payload["message"] == "Unexpected error" + + @pytest.mark.asyncio + @mock.patch("airflow.providers.google.cloud.triggers.cloud_functions.CloudFunctionsHook") + async def test_run_with_impersonation_chain(self, mock_hook_class): + trigger = CloudFunctionInvokeFunctionTrigger( + function_id="test-function", + input_data={"key": "value"}, + location="us-central1", + project_id="test-project", + gcp_conn_id="google_cloud_default", + api_version="v1", + impersonation_chain=["account1", "account2"], + ) + + mock_hook = mock_hook_class.return_value + mock_hook.call_function.return_value = {"executionId": "exec-123"} + + events = [] + async for event in trigger.run(): + events.append(event) + + assert len(events) == 1 + mock_hook_class.assert_called_once_with( + api_version="v1", + gcp_conn_id="google_cloud_default", + impersonation_chain=["account1", "account2"], + )