apache-airflow-providers-google 16.0.0a1__py3-none-any.whl → 16.1.0__py3-none-any.whl

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

@@ -21,7 +21,7 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import AlreadyExists, NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -150,9 +150,17 @@ class ManagedKafkaCreateClusterOperator(ManagedKafkaBaseOperator):
         self.cluster_id = cluster_id
         self.request_id = request_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         self.log.info("Creating an Apache Kafka cluster.")
-        ApacheKafkaClusterLink.persist(context=context, task_instance=self, cluster_id=self.cluster_id)
+        ApacheKafkaClusterLink.persist(context=context)
         try:
             operation = self.hook.create_cluster(
                 project_id=self.project_id,
@@ -227,8 +235,14 @@ class ManagedKafkaListClustersOperator(ManagedKafkaBaseOperator):
         self.filter = filter
         self.order_by = order_by
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaClusterListLink.persist(context=context, task_instance=self)
+        ApacheKafkaClusterListLink.persist(context=context)
         self.log.info("Listing Clusters from location %s.", self.location)
         try:
             cluster_list_pager = self.hook.list_clusters(
@@ -242,8 +256,7 @@ class ManagedKafkaListClustersOperator(ManagedKafkaBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="cluster_page",
                 value=types.ListClustersResponse.to_dict(cluster_list_pager._response),
             )
@@ -285,12 +298,16 @@ class ManagedKafkaGetClusterOperator(ManagedKafkaBaseOperator):
         super().__init__(*args, **kwargs)
         self.cluster_id = cluster_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaClusterLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-        )
+        ApacheKafkaClusterLink.persist(context=context)
         self.log.info("Getting Cluster: %s", self.cluster_id)
         try:
             cluster = self.hook.get_cluster(
@@ -362,12 +379,16 @@ class ManagedKafkaUpdateClusterOperator(ManagedKafkaBaseOperator):
         self.update_mask = update_mask
         self.request_id = request_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaClusterLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-        )
+        ApacheKafkaClusterLink.persist(context=context)
         self.log.info("Updating an Apache Kafka cluster.")
         try:
             operation = self.hook.update_cluster(
@@ -497,14 +518,18 @@ class ManagedKafkaCreateTopicOperator(ManagedKafkaBaseOperator):
         self.topic_id = topic_id
         self.topic = topic
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "topic_id": self.topic_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         self.log.info("Creating an Apache Kafka topic.")
-        ApacheKafkaTopicLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-            topic_id=self.topic_id,
-        )
+        ApacheKafkaTopicLink.persist(context=context)
         try:
             topic_obj = self.hook.create_topic(
                 project_id=self.project_id,
@@ -574,8 +599,16 @@ class ManagedKafkaListTopicsOperator(ManagedKafkaBaseOperator):
         self.page_size = page_size
         self.page_token = page_token
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaClusterLink.persist(context=context, task_instance=self, cluster_id=self.cluster_id)
+        ApacheKafkaClusterLink.persist(context=context)
         self.log.info("Listing Topics for cluster %s.", self.cluster_id)
         try:
             topic_list_pager = self.hook.list_topics(
@@ -588,8 +621,7 @@ class ManagedKafkaListTopicsOperator(ManagedKafkaBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="topic_page",
                 value=types.ListTopicsResponse.to_dict(topic_list_pager._response),
             )
@@ -636,13 +668,17 @@ class ManagedKafkaGetTopicOperator(ManagedKafkaBaseOperator):
         self.cluster_id = cluster_id
         self.topic_id = topic_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "topic_id": self.topic_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaTopicLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-            topic_id=self.topic_id,
-        )
+        ApacheKafkaTopicLink.persist(context=context)
         self.log.info("Getting Topic: %s", self.topic_id)
         try:
             topic = self.hook.get_topic(
@@ -707,13 +743,17 @@ class ManagedKafkaUpdateTopicOperator(ManagedKafkaBaseOperator):
         self.topic = topic
         self.update_mask = update_mask
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "topic_id": self.topic_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaTopicLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-            topic_id=self.topic_id,
-        )
+        ApacheKafkaTopicLink.persist(context=context)
         self.log.info("Updating an Apache Kafka topic.")
         try:
             topic_obj = self.hook.update_topic(
@@ -833,8 +873,16 @@ class ManagedKafkaListConsumerGroupsOperator(ManagedKafkaBaseOperator):
         self.page_size = page_size
         self.page_token = page_token
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaClusterLink.persist(context=context, task_instance=self, cluster_id=self.cluster_id)
+        ApacheKafkaClusterLink.persist(context=context)
         self.log.info("Listing Consumer Groups for cluster %s.", self.cluster_id)
         try:
             consumer_group_list_pager = self.hook.list_consumer_groups(
@@ -847,8 +895,7 @@ class ManagedKafkaListConsumerGroupsOperator(ManagedKafkaBaseOperator):
                 timeout=self.timeout,
                 metadata=self.metadata,
             )
-            self.xcom_push(
-                context=context,
+            context["ti"].xcom_push(
                 key="consumer_group_page",
                 value=types.ListConsumerGroupsResponse.to_dict(consumer_group_list_pager._response),
             )
@@ -895,13 +942,17 @@ class ManagedKafkaGetConsumerGroupOperator(ManagedKafkaBaseOperator):
         self.cluster_id = cluster_id
         self.consumer_group_id = consumer_group_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "consumer_group_id": self.consumer_group_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaConsumerGroupLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-            consumer_group_id=self.consumer_group_id,
-        )
+        ApacheKafkaConsumerGroupLink.persist(context=context)
         self.log.info("Getting Consumer Group: %s", self.consumer_group_id)
         try:
             consumer_group = self.hook.get_consumer_group(
@@ -971,13 +1022,17 @@ class ManagedKafkaUpdateConsumerGroupOperator(ManagedKafkaBaseOperator):
         self.consumer_group = consumer_group
         self.update_mask = update_mask
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location": self.location,
+            "cluster_id": self.cluster_id,
+            "consumer_group_id": self.consumer_group_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
-        ApacheKafkaConsumerGroupLink.persist(
-            context=context,
-            task_instance=self,
-            cluster_id=self.cluster_id,
-            consumer_group_id=self.consumer_group_id,
-        )
+        ApacheKafkaConsumerGroupLink.persist(context=context)
         self.log.info("Updating an Apache Kafka consumer group.")
         try:
             consumer_group_obj = self.hook.update_consumer_group(

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

@@ -104,7 +104,6 @@ class MLEngineCreateModelOperator(GoogleCloudBaseOperator):
         if project_id:
             MLEngineModelLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 model_id=self.model["name"],
             )

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

@@ -25,8 +25,8 @@ This module contains Google PubSub operators.
 
 from __future__ import annotations
 
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Callable
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.pubsub_v1.types import (
@@ -40,6 +40,7 @@ from google.cloud.pubsub_v1.types import (
     SchemaSettings,
 )
 
+from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.pubsub import PubSubHook
 from airflow.providers.google.cloud.links.pubsub import PubSubSubscriptionLink, PubSubTopicLink
@@ -182,7 +183,6 @@ class PubSubCreateTopicOperator(GoogleCloudBaseOperator):
         self.log.info("Created topic %s", self.topic)
         PubSubTopicLink.persist(
             context=context,
-            task_instance=self,
             topic_id=self.topic,
             project_id=self.project_id or hook.project_id,
         )
@@ -391,7 +391,6 @@ class PubSubCreateSubscriptionOperator(GoogleCloudBaseOperator):
         self.log.info("Created subscription for topic %s", self.topic)
         PubSubSubscriptionLink.persist(
             context=context,
-            task_instance=self,
             subscription_id=self.subscription or result,  # result returns subscription name
             project_id=self.project_id or hook.project_id,
         )
@@ -770,7 +769,7 @@ class PubSubPullOperator(GoogleCloudBaseOperator):
         messages_callback: Callable[[list[ReceivedMessage], Context], Any] | None = None,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
-        deferrable: bool = False,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         poll_interval: int = 300,
         **kwargs,
     ) -> None:

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

@@ -122,7 +122,6 @@ class SpannerDeployInstanceOperator(GoogleCloudBaseOperator):
         )
         SpannerInstanceLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id,
             project_id=self.project_id or hook.project_id,
         )
@@ -290,7 +289,6 @@ class SpannerQueryDatabaseInstanceOperator(GoogleCloudBaseOperator):
         )
         SpannerDatabaseLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id,
             database_id=self.database_id,
             project_id=self.project_id or hook.project_id,
@@ -380,7 +378,6 @@ class SpannerDeployDatabaseInstanceOperator(GoogleCloudBaseOperator):
         )
         SpannerDatabaseLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id,
             database_id=self.database_id,
             project_id=self.project_id or hook.project_id,
@@ -496,7 +493,6 @@ class SpannerUpdateDatabaseInstanceOperator(GoogleCloudBaseOperator):
             )
         SpannerDatabaseLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id,
             database_id=self.database_id,
             project_id=self.project_id or hook.project_id,

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

@@ -117,7 +117,6 @@ class CloudSpeechToTextRecognizeSpeechOperator(GoogleCloudBaseOperator):
         if self.audio.uri:
             FileDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 # Slice from: "gs://{BUCKET_NAME}/{FILE_NAME}" to: "{BUCKET_NAME}/{FILE_NAME}"
                 uri=self.audio.uri[5:],
                 project_id=self.project_id or hook.project_id,

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

@@ -145,7 +145,6 @@ class StackdriverListAlertPoliciesOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
         return [AlertPolicy.to_dict(policy) for policy in result]
@@ -228,7 +227,6 @@ class StackdriverEnableAlertPoliciesOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -311,7 +309,6 @@ class StackdriverDisableAlertPoliciesOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -394,7 +391,6 @@ class StackdriverUpsertAlertOperator(GoogleCloudBaseOperator):
         )
         StackdriverPoliciesLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -580,7 +576,6 @@ class StackdriverListNotificationChannelsOperator(GoogleCloudBaseOperator):
         )
         StackdriverNotificationsLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
         return [NotificationChannel.to_dict(channel) for channel in channels]
@@ -666,7 +661,6 @@ class StackdriverEnableNotificationChannelsOperator(GoogleCloudBaseOperator):
         )
         StackdriverNotificationsLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -751,7 +745,6 @@ class StackdriverDisableNotificationChannelsOperator(GoogleCloudBaseOperator):
         )
         StackdriverNotificationsLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 
@@ -838,7 +831,6 @@ class StackdriverUpsertNotificationChannelOperator(GoogleCloudBaseOperator):
         )
         StackdriverNotificationsLink.persist(
             context=context,
-            operator_instance=self,
             project_id=self.project_id or self.hook.project_id,
         )
 

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

@@ -137,7 +137,6 @@ class CloudTasksQueueCreateOperator(GoogleCloudBaseOperator):
                 metadata=self.metadata,
             )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -236,7 +235,6 @@ class CloudTasksQueueUpdateOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -319,7 +317,6 @@ class CloudTasksQueueGetOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -406,7 +403,6 @@ class CloudTasksQueuesListOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksLink.persist(
-            operator_instance=self,
             context=context,
             project_id=self.project_id or hook.project_id,
         )
@@ -564,7 +560,6 @@ class CloudTasksQueuePurgeOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -647,7 +642,6 @@ class CloudTasksQueuePauseOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -730,7 +724,6 @@ class CloudTasksQueueResumeOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=queue.name,
         )
@@ -830,7 +823,6 @@ class CloudTasksTaskCreateOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=task.name,
         )
@@ -923,7 +915,6 @@ class CloudTasksTaskGetOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=task.name,
         )
@@ -1016,7 +1007,6 @@ class CloudTasksTasksListOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=f"projects/{self.project_id or hook.project_id}/"
             f"locations/{self.location}/queues/{self.queue_name}",
@@ -1190,7 +1180,6 @@ class CloudTasksTaskRunOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         CloudTasksQueueLink.persist(
-            operator_instance=self,
             context=context,
             queue_name=task.name,
         )

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

@@ -150,7 +150,6 @@ class CloudTextToSpeechSynthesizeOperator(GoogleCloudBaseOperator):
             )
             FileDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 uri=f"{self.target_bucket_name}/{self.target_filename}",
                 project_id=cloud_storage_hook.project_id,
             )