apache-airflow-providers-google 10.19.0rc1__py3-none-any.whl → 10.20.0rc1__py3-none-any.whl

airflow/providers/google/cloud/operators/datapipeline.py

@@ -19,138 +19,46 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
-
-from airflow.exceptions import AirflowException
-from airflow.providers.google.cloud.hooks.datapipeline import DEFAULT_DATAPIPELINE_LOCATION, DataPipelineHook
-from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from deprecated import deprecated
+
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.google.cloud.hooks.dataflow import DEFAULT_DATAFLOW_LOCATION
+from airflow.providers.google.cloud.operators.dataflow import (
+    DataflowCreatePipelineOperator,
+    DataflowRunPipelineOperator,
+)
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
-if TYPE_CHECKING:
-    from airflow.utils.context import Context
-
 
-class CreateDataPipelineOperator(GoogleCloudBaseOperator):
-    """
-    Creates a new Data Pipelines instance from the Data Pipelines API.
+@deprecated(
+    reason="This operator is deprecated and will be removed after 01.12.2024. "
+    "Please use `DataflowCreatePipelineOperator`.",
+    category=AirflowProviderDeprecationWarning,
+)
+class CreateDataPipelineOperator(DataflowCreatePipelineOperator):
+    """Creates a new Data Pipelines instance from the Data Pipelines API."""
 
-    :param body: The request body (contains instance of Pipeline). See:
-        https://cloud.google.com/dataflow/docs/reference/data-pipelines/rest/v1/projects.locations.pipelines/create#request-body
-    :param project_id: The ID of the GCP project that owns the job.
-    :param location: The location to direct the Data Pipelines instance to (for example us-central1).
-    :param gcp_conn_id: The connection ID to connect to the Google Cloud
-        Platform.
-    :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).
 
-        .. warning::
-            This option requires Apache Beam 2.39.0 or newer.
-
-    Returns the created Data Pipelines instance in JSON representation.
-    """
-
-    def __init__(
-        self,
-        *,
-        body: dict,
-        project_id: str = PROVIDE_PROJECT_ID,
-        location: str = DEFAULT_DATAPIPELINE_LOCATION,
-        gcp_conn_id: str = "google_cloud_default",
-        impersonation_chain: str | Sequence[str] | None = None,
-        **kwargs,
-    ) -> None:
-        super().__init__(**kwargs)
-
-        self.body = body
-        self.project_id = project_id
-        self.location = location
-        self.gcp_conn_id = gcp_conn_id
-        self.impersonation_chain = impersonation_chain
-        self.datapipeline_hook: DataPipelineHook | None = None
-        self.body["pipelineSources"] = {"airflow": "airflow"}
-
-    def execute(self, context: Context):
-        if self.body is None:
-            raise AirflowException(
-                "Request Body not given; cannot create a Data Pipeline without the Request Body."
-            )
-        if self.project_id is None:
-            raise AirflowException(
-                "Project ID not given; cannot create a Data Pipeline without the Project ID."
-            )
-        if self.location is None:
-            raise AirflowException("location not given; cannot create a Data Pipeline without the location.")
-
-        self.datapipeline_hook = DataPipelineHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
-
-        self.data_pipeline = self.datapipeline_hook.create_data_pipeline(
-            project_id=self.project_id,
-            body=self.body,
-            location=self.location,
-        )
-        if self.data_pipeline:
-            if "error" in self.data_pipeline:
-                raise AirflowException(self.data_pipeline.get("error").get("message"))
-
-        return self.data_pipeline
-
-
-class RunDataPipelineOperator(GoogleCloudBaseOperator):
-    """
-    Runs a Data Pipelines Instance using the Data Pipelines API.
-
-    :param data_pipeline_name:  The display name of the pipeline. In example
-        projects/PROJECT_ID/locations/LOCATION_ID/pipelines/PIPELINE_ID it would be the PIPELINE_ID.
-    :param project_id: The ID of the GCP project that owns the job.
-    :param location: The location to direct the Data Pipelines instance to (for example us-central1).
-    :param gcp_conn_id: The connection ID to connect to the Google Cloud
-        Platform.
-
-    Returns the created Job in JSON representation.
-    """
+@deprecated(
+    reason="This operator is deprecated and will be removed after 01.12.2024. "
+    "Please use `DataflowRunPipelineOperator`.",
+    category=AirflowProviderDeprecationWarning,
+)
+class RunDataPipelineOperator(DataflowRunPipelineOperator):
+    """Runs a Data Pipelines Instance using the Data Pipelines API."""
 
     def __init__(
         self,
         data_pipeline_name: str,
         project_id: str = PROVIDE_PROJECT_ID,
-        location: str = DEFAULT_DATAPIPELINE_LOCATION,
+        location: str = DEFAULT_DATAFLOW_LOCATION,
         gcp_conn_id: str = "google_cloud_default",
         **kwargs,
     ) -> None:
-        super().__init__(**kwargs)
-
-        self.data_pipeline_name = data_pipeline_name
-        self.project_id = project_id
-        self.location = location
-        self.gcp_conn_id = gcp_conn_id
-
-    def execute(self, context: Context):
-        self.data_pipeline_hook = DataPipelineHook(gcp_conn_id=self.gcp_conn_id)
-
-        if self.data_pipeline_name is None:
-            raise AirflowException("Data Pipeline name not given; cannot run unspecified pipeline.")
-        if self.project_id is None:
-            raise AirflowException("Data Pipeline Project ID not given; cannot run pipeline.")
-        if self.location is None:
-            raise AirflowException("Data Pipeline location not given; cannot run pipeline.")
-
-        self.response = self.data_pipeline_hook.run_data_pipeline(
-            data_pipeline_name=self.data_pipeline_name,
-            project_id=self.project_id,
-            location=self.location,
+        super().__init__(
+            pipeline_name=data_pipeline_name,
+            project_id=project_id,
+            location=location,
+            gcp_conn_id=gcp_conn_id,
+            **kwargs,
         )
-
-        if self.response:
-            if "error" in self.response:
-                raise AirflowException(self.response.get("error").get("message"))
-
-        return self.response

airflow/providers/google/cloud/operators/dataplex.py

@@ -1067,7 +1067,7 @@ class DataplexGetDataQualityScanResultOperator(GoogleCloudBaseOperator):
         is available.
     """
 
-    template_fields = ("project_id", "data_scan_id", "impersonation_chain")
+    template_fields = ("project_id", "data_scan_id", "impersonation_chain", "job_id")
 
     def __init__(
         self,

airflow/providers/google/cloud/operators/dataproc_metastore.py

@@ -431,7 +431,7 @@ class DataprocMetastoreCreateServiceOperator(GoogleCloudBaseOperator):
         hook = DataprocMetastoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
         )
-        self.log.info("Creating Dataproc Metastore service: %s", self.project_id)
+        self.log.info("Creating Dataproc Metastore service: %s", self.service_id)
         try:
             operation = hook.create_service(
                 region=self.region,
@@ -548,13 +548,24 @@ class DataprocMetastoreDeleteBackupOperator(GoogleCloudBaseOperator):
 class DataprocMetastoreDeleteServiceOperator(GoogleCloudBaseOperator):
     """Delete a single service.
 
-    :param request:  The request object. Request message for
-        [DataprocMetastore.DeleteService][google.cloud.metastore.v1.DataprocMetastore.DeleteService].
+    :param region: Required. The ID of the Google Cloud region that the service belongs to.
     :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param service_id:  Required. The ID of the metastore service, which is used as the final component of
+        the metastore service's name. This value must be between 2 and 63 characters long inclusive, begin
+        with a letter, end with a letter or number, and consist of alphanumeric ASCII characters or
+        hyphens.
     :param retry: Designation of what errors, if any, should be retried.
     :param timeout: The timeout for this request.
     :param metadata: Strings which should be sent along with the request as metadata.
-    :param gcp_conn_id:
+    :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] = (
@@ -589,7 +600,7 @@ class DataprocMetastoreDeleteServiceOperator(GoogleCloudBaseOperator):
         hook = DataprocMetastoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
         )
-        self.log.info("Deleting Dataproc Metastore service: %s", self.project_id)
+        self.log.info("Deleting Dataproc Metastore service: %s", self.service_id)
         operation = hook.delete_service(
             region=self.region,
             project_id=self.project_id,
@@ -599,7 +610,7 @@ class DataprocMetastoreDeleteServiceOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         hook.wait_for_operation(self.timeout, operation)
-        self.log.info("Service %s deleted successfully", self.project_id)
+        self.log.info("Service %s deleted successfully", self.service_id)
 
 
 class DataprocMetastoreExportMetadataOperator(GoogleCloudBaseOperator):

airflow/providers/google/cloud/operators/kubernetes_engine.py

@@ -73,6 +73,7 @@ except ImportError:
 
 if TYPE_CHECKING:
     from kubernetes.client.models import V1Job, V1Pod
+    from pendulum import DateTime
 
     from airflow.utils.context import Context
 
@@ -773,16 +774,16 @@ class GKEStartPodOperator(KubernetesPodOperator):
         self._ssl_ca_cert = cluster.master_auth.cluster_ca_certificate
         return self._cluster_url, self._ssl_ca_cert
 
-    def invoke_defer_method(self):
+    def invoke_defer_method(self, last_log_time: DateTime | None = None):
         """Redefine triggers which are being used in child classes."""
         trigger_start_time = utcnow()
         self.defer(
             trigger=GKEStartPodTrigger(
-                pod_name=self.pod.metadata.name,
-                pod_namespace=self.pod.metadata.namespace,
+                pod_name=self.pod.metadata.name,  # type: ignore[union-attr]
+                pod_namespace=self.pod.metadata.namespace,  # type: ignore[union-attr]
                 trigger_start_time=trigger_start_time,
-                cluster_url=self._cluster_url,
-                ssl_ca_cert=self._ssl_ca_cert,
+                cluster_url=self._cluster_url,  # type: ignore[arg-type]
+                ssl_ca_cert=self._ssl_ca_cert,  # type: ignore[arg-type]
                 get_logs=self.get_logs,
                 startup_timeout=self.startup_timeout_seconds,
                 cluster_context=self.cluster_context,
@@ -792,6 +793,8 @@ class GKEStartPodOperator(KubernetesPodOperator):
                 on_finish_action=self.on_finish_action,
                 gcp_conn_id=self.gcp_conn_id,
                 impersonation_chain=self.impersonation_chain,
+                logging_interval=self.logging_interval,
+                last_log_time=last_log_time,
             ),
             method_name="execute_complete",
             kwargs={"cluster_url": self._cluster_url, "ssl_ca_cert": self._ssl_ca_cert},
@@ -802,7 +805,7 @@ class GKEStartPodOperator(KubernetesPodOperator):
         self._cluster_url = kwargs["cluster_url"]
         self._ssl_ca_cert = kwargs["ssl_ca_cert"]
 
-        return super().execute_complete(context, event, **kwargs)
+        return super().trigger_reentry(context, event)
 
 
 class GKEStartJobOperator(KubernetesJobOperator):

airflow/providers/google/cloud/operators/pubsub.py

@@ -604,6 +604,7 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
         m1 = {"data": b"Hello, World!", "attributes": {"type": "greeting"}}
         m2 = {"data": b"Knock, knock"}
         m3 = {"attributes": {"foo": ""}}
+        m4 = {"data": b"Who's there?", "attributes": {"ordering_key": "knock_knock"}}
 
         t1 = PubSubPublishMessageOperator(
             project_id="my-project",
@@ -613,6 +614,15 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
             dag=dag,
         )
 
+        t2 = PubSubPublishMessageOperator(
+            project_id="my-project",
+            topic="my_topic",
+            messages=[m4],
+            create_topic=True,
+            enable_message_ordering=True,
+            dag=dag,
+        )
+
     ``project_id``, ``topic``, and ``messages`` are templated so you can use Jinja templating
     in their values.
 
@@ -632,6 +642,10 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
         https://cloud.google.com/pubsub/docs/reference/rest/v1/PubsubMessage
     :param gcp_conn_id: The connection ID to use connecting to
         Google Cloud.
+    :param enable_message_ordering: If true, messages published with the same
+        ordering_key in PubsubMessage will be delivered to the subscribers in the order
+        in which they are received by the Pub/Sub system. Otherwise, they may be
+        delivered in any order. Default is False.
     :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.
@@ -646,6 +660,7 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
         "project_id",
         "topic",
         "messages",
+        "enable_message_ordering",
         "impersonation_chain",
     )
     ui_color = "#0273d4"
@@ -657,6 +672,7 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
         messages: list,
         project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
+        enable_message_ordering: bool = False,
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
@@ -665,12 +681,14 @@ class PubSubPublishMessageOperator(GoogleCloudBaseOperator):
         self.topic = topic
         self.messages = messages
         self.gcp_conn_id = gcp_conn_id
+        self.enable_message_ordering = enable_message_ordering
         self.impersonation_chain = impersonation_chain
 
     def execute(self, context: Context) -> None:
         hook = PubSubHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
+            enable_message_ordering=self.enable_message_ordering,
         )
 
         self.log.info("Publishing to topic %s", self.topic)

airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py

@@ -138,6 +138,8 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         region: str,
         impersonation_chain: str | Sequence[str] | None = None,
         parent_model: str | None = None,
+        window_stride_length: int | None = None,
+        window_max_count: int | None = None,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -170,6 +172,8 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         self.quantiles = quantiles
         self.validation_options = validation_options
         self.budget_milli_node_hours = budget_milli_node_hours
+        self.window_stride_length = window_stride_length
+        self.window_max_count = window_max_count
 
     def execute(self, context: Context):
         self.hook = AutoMLHook(
@@ -220,6 +224,8 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
             model_display_name=self.model_display_name,
             model_labels=self.model_labels,
             sync=self.sync,
+            window_stride_length=self.window_stride_length,
+            window_max_count=self.window_max_count,
         )
 
         if model:

airflow/providers/google/cloud/operators/vertex_ai/generative_model.py

@@ -187,6 +187,8 @@ class PromptMultimodalModelOperator(GoogleCloudBaseOperator):
         service belongs to (templated).
     :param prompt: Required. Inputs or queries that a user or a program gives
         to the Multi-modal model, in order to elicit a specific response (templated).
+    :param generation_config: Optional. Generation configuration settings.
+    :param safety_settings: Optional. Per request settings for blocking unsafe content.
     :param pretrained_model: By default uses the pre-trained model `gemini-pro`,
         supporting prompts with text-only input, including natural language
         tasks, multi-turn text and code chat, and code generation. It can
@@ -210,6 +212,8 @@ class PromptMultimodalModelOperator(GoogleCloudBaseOperator):
         project_id: str,
         location: str,
         prompt: str,
+        generation_config: dict | None = None,
+        safety_settings: dict | None = None,
         pretrained_model: str = "gemini-pro",
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
@@ -219,6 +223,8 @@ class PromptMultimodalModelOperator(GoogleCloudBaseOperator):
         self.project_id = project_id
         self.location = location
         self.prompt = prompt
+        self.generation_config = generation_config
+        self.safety_settings = safety_settings
         self.pretrained_model = pretrained_model
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
@@ -232,6 +238,8 @@ class PromptMultimodalModelOperator(GoogleCloudBaseOperator):
             project_id=self.project_id,
             location=self.location,
             prompt=self.prompt,
+            generation_config=self.generation_config,
+            safety_settings=self.safety_settings,
             pretrained_model=self.pretrained_model,
         )
 
@@ -251,6 +259,8 @@ class PromptMultimodalModelWithMediaOperator(GoogleCloudBaseOperator):
         service belongs to (templated).
     :param prompt: Required. Inputs or queries that a user or a program gives
         to the Multi-modal model, in order to elicit a specific response (templated).
+    :param generation_config: Optional. Generation configuration settings.
+    :param safety_settings: Optional. Per request settings for blocking unsafe content.
     :param pretrained_model: By default uses the pre-trained model `gemini-pro-vision`,
         supporting prompts with text-only input, including natural language
         tasks, multi-turn text and code chat, and code generation. It can
@@ -279,6 +289,8 @@ class PromptMultimodalModelWithMediaOperator(GoogleCloudBaseOperator):
         prompt: str,
         media_gcs_path: str,
         mime_type: str,
+        generation_config: dict | None = None,
+        safety_settings: dict | None = None,
         pretrained_model: str = "gemini-pro-vision",
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
@@ -288,6 +300,8 @@ class PromptMultimodalModelWithMediaOperator(GoogleCloudBaseOperator):
         self.project_id = project_id
         self.location = location
         self.prompt = prompt
+        self.generation_config = generation_config
+        self.safety_settings = safety_settings
         self.pretrained_model = pretrained_model
         self.media_gcs_path = media_gcs_path
         self.mime_type = mime_type
@@ -303,6 +317,8 @@ class PromptMultimodalModelWithMediaOperator(GoogleCloudBaseOperator):
             project_id=self.project_id,
             location=self.location,
             prompt=self.prompt,
+            generation_config=self.generation_config,
+            safety_settings=self.safety_settings,
             pretrained_model=self.pretrained_model,
             media_gcs_path=self.media_gcs_path,
             mime_type=self.mime_type,

airflow/providers/google/cloud/sensors/cloud_composer.py

@@ -19,13 +19,24 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+import json
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING, Any, Iterable, Sequence
 
+from dateutil import parser
 from deprecated import deprecated
+from google.cloud.orchestration.airflow.service_v1.types import ExecuteAirflowCommandResponse
 
+from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowSkipException
-from airflow.providers.google.cloud.triggers.cloud_composer import CloudComposerExecutionTrigger
+from airflow.providers.google.cloud.hooks.cloud_composer import CloudComposerHook
+from airflow.providers.google.cloud.triggers.cloud_composer import (
+    CloudComposerDAGRunTrigger,
+    CloudComposerExecutionTrigger,
+)
+from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
 from airflow.sensors.base import BaseSensorOperator
+from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -117,3 +128,161 @@ class CloudComposerEnvironmentSensor(BaseSensorOperator):
         if self.soft_fail:
             raise AirflowSkipException(message)
         raise AirflowException(message)
+
+
+class CloudComposerDAGRunSensor(BaseSensorOperator):
+    """
+    Check if a DAG run has completed.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param region: Required. The ID of the Google Cloud region that the service belongs to.
+    :param environment_id: The name of the Composer environment.
+    :param composer_dag_id: The ID of executable DAG.
+    :param allowed_states: Iterable of allowed states, default is ``['success']``.
+    :param execution_range: execution DAGs time range. Sensor checks DAGs states only for DAGs which were
+        started in this time range. For yesterday, use [positive!] datetime.timedelta(days=1).
+        For future, use [negative!] datetime.timedelta(days=-1). For specific time, use list of
+        datetimes [datetime(2024,3,22,11,0,0), datetime(2024,3,22,12,0,0)].
+        Or [datetime(2024,3,22,0,0,0)] in this case sensor will check for states from specific time in the
+        past till current time execution.
+        Default value datetime.timedelta(days=1).
+    :param gcp_conn_id: The connection ID to use when fetching connection info.
+    :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).
+    :param poll_interval: Optional: Control the rate of the poll for the result of deferrable run.
+    :param deferrable: Run sensor in deferrable mode.
+    """
+
+    template_fields = (
+        "project_id",
+        "region",
+        "environment_id",
+        "composer_dag_id",
+        "impersonation_chain",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        region: str,
+        environment_id: str,
+        composer_dag_id: str,
+        allowed_states: Iterable[str] | None = None,
+        execution_range: timedelta | list[datetime] | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        poll_interval: int = 10,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.region = region
+        self.environment_id = environment_id
+        self.composer_dag_id = composer_dag_id
+        self.allowed_states = list(allowed_states) if allowed_states else [TaskInstanceState.SUCCESS.value]
+        self.execution_range = execution_range
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+        self.deferrable = deferrable
+        self.poll_interval = poll_interval
+
+    def _get_execution_dates(self, context) -> tuple[datetime, datetime]:
+        if isinstance(self.execution_range, timedelta):
+            if self.execution_range < timedelta(0):
+                return context["logical_date"], context["logical_date"] - self.execution_range
+            else:
+                return context["logical_date"] - self.execution_range, context["logical_date"]
+        elif isinstance(self.execution_range, list) and len(self.execution_range) > 0:
+            return self.execution_range[0], self.execution_range[1] if len(
+                self.execution_range
+            ) > 1 else context["logical_date"]
+        else:
+            return context["logical_date"] - timedelta(1), context["logical_date"]
+
+    def poke(self, context: Context) -> bool:
+        start_date, end_date = self._get_execution_dates(context)
+
+        if datetime.now(end_date.tzinfo) < end_date:
+            return False
+
+        dag_runs = self._pull_dag_runs()
+
+        self.log.info("Sensor waits for allowed states: %s", self.allowed_states)
+        allowed_states_status = self._check_dag_runs_states(
+            dag_runs=dag_runs,
+            start_date=start_date,
+            end_date=end_date,
+        )
+
+        return allowed_states_status
+
+    def _pull_dag_runs(self) -> list[dict]:
+        """Pull the list of dag runs."""
+        hook = CloudComposerHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        dag_runs_cmd = hook.execute_airflow_command(
+            project_id=self.project_id,
+            region=self.region,
+            environment_id=self.environment_id,
+            command="dags",
+            subcommand="list-runs",
+            parameters=["-d", self.composer_dag_id, "-o", "json"],
+        )
+        cmd_result = hook.wait_command_execution_result(
+            project_id=self.project_id,
+            region=self.region,
+            environment_id=self.environment_id,
+            execution_cmd_info=ExecuteAirflowCommandResponse.to_dict(dag_runs_cmd),
+        )
+        dag_runs = json.loads(cmd_result["output"][0]["content"])
+        return dag_runs
+
+    def _check_dag_runs_states(
+        self,
+        dag_runs: list[dict],
+        start_date: datetime,
+        end_date: datetime,
+    ) -> bool:
+        for dag_run in dag_runs:
+            if (
+                start_date.timestamp()
+                < parser.parse(dag_run["execution_date"]).timestamp()
+                < end_date.timestamp()
+            ) and dag_run["state"] not in self.allowed_states:
+                return False
+        return True
+
+    def execute(self, context: Context) -> None:
+        if self.deferrable:
+            start_date, end_date = self._get_execution_dates(context)
+            self.defer(
+                trigger=CloudComposerDAGRunTrigger(
+                    project_id=self.project_id,
+                    region=self.region,
+                    environment_id=self.environment_id,
+                    composer_dag_id=self.composer_dag_id,
+                    start_date=start_date,
+                    end_date=end_date,
+                    allowed_states=self.allowed_states,
+                    gcp_conn_id=self.gcp_conn_id,
+                    impersonation_chain=self.impersonation_chain,
+                    poll_interval=self.poll_interval,
+                ),
+                method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
+            )
+        super().execute(context)
+
+    def execute_complete(self, context: Context, event: dict):
+        if event and event["status"] == "error":
+            raise AirflowException(event["message"])
+        self.log.info("DAG %s has executed successfully.", self.composer_dag_id)

airflow/providers/google/cloud/transfers/azure_blob_to_gcs.py

@@ -122,3 +122,16 @@ class AzureBlobStorageToGCSOperator(BaseOperator):
                 self.bucket_name,
             )
         return f"gs://{self.bucket_name}/{self.object_name}"
+
+    def get_openlineage_facets_on_start(self):
+        from openlineage.client.run import Dataset
+
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        wasb_hook = WasbHook(wasb_conn_id=self.wasb_conn_id)
+        account_name = wasb_hook.get_conn().account_name
+
+        return OperatorLineage(
+            inputs=[Dataset(namespace=f"wasbs://{self.container_name}@{account_name}", name=self.blob_name)],
+            outputs=[Dataset(namespace=f"gs://{self.bucket_name}", name=self.object_name)],
+        )