Add deferrable support for Cloud Functions invoke operator

airflow-core/newsfragments/69853.feature.rst

@@ -0,0 +1 @@
+Add deferrable support for CloudFunctionInvokeFunctionOperator to invoke Google Cloud Functions asynchronously without blocking workers.

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

providers/google/src/airflow/providers/google/cloud/operators/functions.py

@@ -25,13 +25,17 @@
 
 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 (
     CloudFunctionsDetailsLink,
     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

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)})

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"],

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)