apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

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

@@ -20,9 +20,15 @@ import datetime
 import json
 import re
 import uuid
-from collections.abc import Sequence
+from collections.abc import Collection, Sequence
 from typing import TYPE_CHECKING
 
+import pendulum
+from google.api_core.exceptions import AlreadyExists
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.workflows.executions_v1beta import Execution
+from google.cloud.workflows_v1beta import Workflow
+
 from airflow.providers.google.cloud.hooks.workflows import WorkflowsHook
 from airflow.providers.google.cloud.links.workflows import (
     WorkflowsExecutionLink,
@@ -31,16 +37,14 @@ from airflow.providers.google.cloud.links.workflows import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from google.api_core.exceptions import AlreadyExists
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.workflows.executions_v1beta import Execution
-from google.cloud.workflows_v1beta import Workflow
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
+    from airflow.sdk import Context
+
 from airflow.utils.hashlib_wrapper import md5
 
 
@@ -67,7 +71,7 @@ class WorkflowsCreateWorkflowOperator(GoogleCloudBaseOperator):
     :param metadata: Additional metadata that is provided to the method.
     """
 
-    template_fields: Sequence[str] = ("location", "workflow", "workflow_id")
+    template_fields: Collection[str] = ("location", "workflow", "workflow_id")
     template_fields_renderers = {"workflow": "json"}
     operator_extra_links = (WorkflowsWorkflowDetailsLink(),)
 
@@ -99,7 +103,7 @@ class WorkflowsCreateWorkflowOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.force_rerun = force_rerun
 
-    def _workflow_id(self, context):
+    def _workflow_id(self, context: Context) -> str:
         if self.workflow_id and not self.force_rerun:
             # If users provide workflow id then assuring the idempotency
             # is on their side
@@ -112,8 +116,17 @@ class WorkflowsCreateWorkflowOperator(GoogleCloudBaseOperator):
 
         # We are limited by allowed length of workflow_id so
         # we use hash of whole information
-        exec_date = context["logical_date"].isoformat()
-        base = f"airflow_{self.dag_id}_{self.task_id}_{exec_date}_{hash_base}"
+        if AIRFLOW_V_3_0_PLUS:
+            if dag_run := context.get("dag_run"):
+                run_after = pendulum.instance(dag_run.run_after)
+            else:
+                run_after = pendulum.now("UTC")
+        else:
+            if logical_date := context.get("logical_date"):
+                run_after = pendulum.instance(logical_date)
+            else:
+                run_after = pendulum.now("UTC")
+        base = f"airflow_{self.dag_id}_{self.task_id}_{run_after.isoformat()}_{hash_base}"
         workflow_id = md5(base.encode()).hexdigest()
         return re.sub(r"[:\-+.]", "_", workflow_id)
 
@@ -145,7 +158,6 @@ class WorkflowsCreateWorkflowOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -233,7 +245,6 @@ class WorkflowsUpdateWorkflowOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -366,7 +377,6 @@ class WorkflowsListWorkflowsOperator(GoogleCloudBaseOperator):
 
         WorkflowsListOfWorkflowsLink.persist(
             context=context,
-            task_instance=self,
             project_id=self.project_id or hook.project_id,
         )
 
@@ -432,7 +442,6 @@ class WorkflowsGetWorkflowOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -503,11 +512,10 @@ class WorkflowsCreateExecutionOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         execution_id = execution.name.split("/")[-1]
-        self.xcom_push(context, key="execution_id", value=execution_id)
+        context["task_instance"].xcom_push(key="execution_id", value=execution_id)
 
         WorkflowsExecutionLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             execution_id=execution_id,
@@ -580,7 +588,6 @@ class WorkflowsCancelExecutionOperator(GoogleCloudBaseOperator):
 
         WorkflowsExecutionLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             execution_id=self.execution_id,
@@ -659,7 +666,6 @@ class WorkflowsListExecutionsOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -735,7 +741,6 @@ class WorkflowsGetExecutionOperator(GoogleCloudBaseOperator):
 
         WorkflowsExecutionLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             execution_id=self.execution_id,

airflow/providers/google/cloud/secrets/secret_manager.py

@@ -21,6 +21,8 @@ from __future__ import annotations
 import logging
 from collections.abc import Sequence
 
+from google.auth.exceptions import DefaultCredentialsError
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud._internal_client.secret_manager_client import _SecretManagerClient
 from airflow.providers.google.cloud.utils.credentials_provider import (
@@ -30,7 +32,6 @@ from airflow.providers.google.cloud.utils.credentials_provider import (
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.secrets import BaseSecretsBackend
 from airflow.utils.log.logging_mixin import LoggingMixin
-from google.auth.exceptions import DefaultCredentialsError
 
 log = logging.getLogger(__name__)
 

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

@@ -26,15 +26,15 @@ from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
 from airflow.providers.google.cloud.triggers.bigquery import (
     BigQueryTableExistenceTrigger,
     BigQueryTablePartitionExistenceTrigger,
 )
-from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class BigQueryTableExistenceSensor(BaseSensorOperator):

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

@@ -22,17 +22,19 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.bigquery_datatransfer_v1 import TransferState
+
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.cloud.hooks.bigquery_dts import BiqQueryDataTransferServiceHook
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.bigquery_datatransfer_v1 import TransferState
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class BigQueryDataTransferServiceTransferRunSensor(BaseSensorOperator):
     """

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

@@ -20,19 +20,20 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import google.api_core.exceptions
+from google.cloud.bigtable import enums
+from google.cloud.bigtable.table import ClusterState
+
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.cloud.hooks.bigtable import BigtableHook
 from airflow.providers.google.cloud.links.bigtable import BigtableTablesLink
 from airflow.providers.google.cloud.operators.bigtable import BigtableValidationMixin
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
-from google.cloud.bigtable import enums
-from google.cloud.bigtable.table import ClusterState
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class BigtableTableReplicationCompletedSensor(BaseSensorOperator, BigtableValidationMixin):
@@ -88,6 +89,13 @@ class BigtableTableReplicationCompletedSensor(BaseSensorOperator, BigtableValida
         self.impersonation_chain = impersonation_chain
         super().__init__(**kwargs)
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "project_id": self.project_id,
+        }
+
     def poke(self, context: Context) -> bool:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -118,5 +126,5 @@ class BigtableTableReplicationCompletedSensor(BaseSensorOperator, BigtableValida
             return False
 
         self.log.info("Table '%s' is replicated.", self.table_id)
-        BigtableTablesLink.persist(context=context, task_instance=self)
+        BigtableTablesLink.persist(context=context)
         return True