feat: sensor to check status of Dataform action

Adds a new sensor to check the status of a WorkflowInvocationAction in
Google Cloud Dataform. Heavily based on theDataformWorkflowInvocationStateSensor
which already exists.

Useful for checking the status of a specific target within a Dataform
workflow invocation and taking action based on the status.

docs/apache-airflow-providers-google/operators/cloud/dataform.rst

@@ -95,6 +95,12 @@ We have possibility to run this operation in the sync mode and async, for async
 a sensor:
 :class:`~airflow.providers.google.cloud.operators.dataform.DataformWorkflowInvocationStateSensor`
 
+We also have a sensor to check the status of a particular action for a workflow invocation triggered
+asynchronously.
+
+:class:`~airflow.providers.google.cloud.operators.dataform.DataformWorkflowInvocationActionStateSensor`
+
+
 .. exampleinclude:: /../../providers/tests/system/google/cloud/dataform/example_dataform.py
     :language: python
     :dedent: 4
@@ -107,6 +113,12 @@ a sensor:
     :start-after: [START howto_operator_create_workflow_invocation_async]
     :end-before: [END howto_operator_create_workflow_invocation_async]
 
+.. exampleinclude:: /../../providers/tests/system/google/cloud/dataform/example_dataform.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_create_workflow_invocation_action_async]
+    :end-before: [END howto_operator_create_workflow_invocation_action_async]
+
 Get Workflow Invocation
 -----------------------
 

providers/src/airflow/providers/google/cloud/sensors/dataform.py

@@ -103,3 +103,78 @@ def poke(self, context: Context) -> bool:
                 raise AirflowException(message)
 
         return workflow_status in self.expected_statuses
+
+
+class DataformWorkflowInvocationActionStateSensor(BaseSensorOperator):
+    """
+    Checks for the status of a Workflow Invocation Action in Google Cloud Dataform.
+
+    :param project_id: Required, 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 region: Required, The location of the Dataform workflow invocation (for example europe-west1).
+    :param repository_id: Required. The ID of the Dataform repository that the task belongs to.
+    :param workflow_invocation_id: Required, ID of the workflow invocation to be checked.
+    :param target_name: Required. The name of the target to be checked in the workflow.
+    :param expected_statuses: The expected state of the action.
+        See:
+        https://cloud.google.com/python/docs/reference/dataform/latest/google.cloud.dataform_v1beta1.types.WorkflowInvocationAction.State
+    :param failure_statuses: State that will terminate the sensor with an exception
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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] = ("workflow_invocation_id",)
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        region: str,
+        repository_id: str,
+        workflow_invocation_id: str,
+        target_name: str,
+        expected_statuses: Iterable[int],
+        failure_statuses: Iterable[int],
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.repository_id = repository_id
+        self.workflow_invocation_id = workflow_invocation_id
+        self.project_id = project_id
+        self.region = region
+        self.target_name = target_name
+        self.expected_statuses = expected_statuses
+        self.failure_statuses = failure_statuses
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+        self.hook: DataformHook | None = None
+
+    def poke(self, context: Context) -> bool:
+        self.hook = DataformHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
+
+        workflow_invocation_actions = self.hook.query_workflow_invocation_actions(
+            project_id=self.project_id,
+            region=self.region,
+            repository_id=self.repository_id,
+            workflow_invocation_id=self.workflow_invocation_id,
+        )
+
+        for workflow_invocation_action in workflow_invocation_actions:
+            if workflow_invocation_action.target.name == self.target_name:
+                state = workflow_invocation_action.state
+                if state in self.failure_statuses:
+                    raise AirflowException(
+                        f"Workflow Invocation Action target {self.target_name} state is: {state}."
+                    )
+                return state in self.expected_statuses
+
+        raise AirflowException(f"Workflow Invocation Action target {self.target_name} not found.")

providers/tests/system/google/cloud/dataform/example_dataform.py

@@ -24,7 +24,7 @@
 import os
 from datetime import datetime
 
-from google.cloud.dataform_v1beta1 import WorkflowInvocation
+from google.cloud.dataform_v1beta1 import WorkflowInvocation, WorkflowInvocationAction
 
 from airflow.models.dag import DAG
 from airflow.providers.google.cloud.operators.bigquery import BigQueryDeleteDatasetOperator
@@ -45,7 +45,10 @@
     DataformRemoveFileOperator,
     DataformWriteFileOperator,
 )
-from airflow.providers.google.cloud.sensors.dataform import DataformWorkflowInvocationStateSensor
+from airflow.providers.google.cloud.sensors.dataform import (
+    DataformWorkflowInvocationActionStateSensor,
+    DataformWorkflowInvocationStateSensor,
+)
 from airflow.providers.google.cloud.utils.dataform import make_initialization_workspace_flow
 from airflow.utils.trigger_rule import TriggerRule
 
@@ -174,6 +177,37 @@
     )
     # [END howto_operator_create_workflow_invocation_async]
 
+    # [START howto_operator_create_workflow_invocation_action_async]
+    create_workflow_invocation_async_action = DataformCreateWorkflowInvocationOperator(
+        task_id="create-workflow-invocation-async",
+        project_id=PROJECT_ID,
+        region=REGION,
+        repository_id=REPOSITORY_ID,
+        asynchronous=True,
+        workflow_invocation={
+            "compilation_result": "{{ task_instance.xcom_pull('create-compilation-result')['name'] }}"
+        },
+    )
+
+    is_workflow_invocation_action_done = DataformWorkflowInvocationActionStateSensor(
+        task_id="is-workflow-invocation-done",
+        project_id=PROJECT_ID,
+        region=REGION,
+        repository_id=REPOSITORY_ID,
+        workflow_invocation_id=(
+            "{{ task_instance.xcom_pull('create-workflow-invocation')['name'].split('/')[-1] }}"
+        ),
+        target_name="YOUR_TARGET_HERE",
+        expected_statuses={WorkflowInvocationAction.State.SUCCEEDED},
+        failure_statuses={
+            WorkflowInvocationAction.State.SKIPPED,
+            WorkflowInvocationAction.State.DISABLED,
+            WorkflowInvocationAction.State.CANCELLED,
+            WorkflowInvocationAction.State.FAILED,
+        },
+    )
+    # [END howto_operator_create_workflow_invocation_action_async]
+
     # [START howto_operator_get_workflow_invocation]
     get_workflow_invocation = DataformGetWorkflowInvocationOperator(
         task_id="get-workflow-invocation",
@@ -314,6 +348,8 @@
         >> query_workflow_invocation_actions
         >> create_workflow_invocation_async
         >> is_workflow_invocation_done
+        >> create_workflow_invocation_async_action
+        >> is_workflow_invocation_action_done
         >> create_workflow_invocation_for_cancel
         >> cancel_workflow_invocation
         >> make_test_directory