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

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

@@ -16,13 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.utils.context import Context
-
 MANAGED_KAFKA_BASE_LINK = "/managedkafka"
 MANAGED_KAFKA_CLUSTER_LINK = (
     MANAGED_KAFKA_BASE_LINK + "/{location}/clusters/{cluster_id}?project={project_id}"
@@ -44,22 +39,6 @@ class ApacheKafkaClusterLink(BaseGoogleLink):
     key = "cluster_conf"
     format_str = MANAGED_KAFKA_CLUSTER_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        cluster_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=ApacheKafkaClusterLink.key,
-            value={
-                "location": task_instance.location,
-                "cluster_id": cluster_id,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class ApacheKafkaClusterListLink(BaseGoogleLink):
     """Helper class for constructing Apache Kafka Clusters link."""
@@ -68,19 +47,6 @@ class ApacheKafkaClusterListLink(BaseGoogleLink):
     key = "cluster_list_conf"
     format_str = MANAGED_KAFKA_CLUSTER_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=ApacheKafkaClusterListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class ApacheKafkaTopicLink(BaseGoogleLink):
     """Helper class for constructing Apache Kafka Topic link."""
@@ -89,24 +55,6 @@ class ApacheKafkaTopicLink(BaseGoogleLink):
     key = "topic_conf"
     format_str = MANAGED_KAFKA_TOPIC_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        cluster_id: str,
-        topic_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=ApacheKafkaTopicLink.key,
-            value={
-                "location": task_instance.location,
-                "cluster_id": cluster_id,
-                "topic_id": topic_id,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class ApacheKafkaConsumerGroupLink(BaseGoogleLink):
     """Helper class for constructing Apache Kafka Consumer Group link."""
@@ -114,21 +62,3 @@ class ApacheKafkaConsumerGroupLink(BaseGoogleLink):
     name = "Apache Kafka Consumer Group"
     key = "consumer_group_conf"
     format_str = MANAGED_KAFKA_CONSUMER_GROUP_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        cluster_id: str,
-        consumer_group_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=ApacheKafkaConsumerGroupLink.key,
-            value={
-                "location": task_instance.location,
-                "cluster_id": cluster_id,
-                "consumer_group_id": consumer_group_id,
-                "project_id": task_instance.project_id,
-            },
-        )

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

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.utils.context import Context
-
-
 MLENGINE_BASE_LINK = "https://console.cloud.google.com/ai-platform"
 MLENGINE_MODEL_DETAILS_LINK = MLENGINE_BASE_LINK + "/models/{model_id}/versions?project={project_id}"
 MLENGINE_MODEL_VERSION_DETAILS_LINK = (
@@ -44,19 +38,6 @@ class MLEngineModelLink(BaseGoogleLink):
     key = "ml_engine_model"
     format_str = MLENGINE_MODEL_DETAILS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        model_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=MLEngineModelLink.key,
-            value={"model_id": model_id, "project_id": project_id},
-        )
-
 
 class MLEngineModelsListLink(BaseGoogleLink):
     """Helper class for constructing ML Engine link."""
@@ -65,18 +46,6 @@ class MLEngineModelsListLink(BaseGoogleLink):
     key = "ml_engine_models_list"
     format_str = MLENGINE_MODELS_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=MLEngineModelsListLink.key,
-            value={"project_id": project_id},
-        )
-
 
 class MLEngineJobDetailsLink(BaseGoogleLink):
     """Helper class for constructing ML Engine link."""
@@ -85,19 +54,6 @@ class MLEngineJobDetailsLink(BaseGoogleLink):
     key = "ml_engine_job_details"
     format_str = MLENGINE_JOB_DETAILS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        job_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=MLEngineJobDetailsLink.key,
-            value={"job_id": job_id, "project_id": project_id},
-        )
-
 
 class MLEngineModelVersionDetailsLink(BaseGoogleLink):
     """Helper class for constructing ML Engine link."""
@@ -106,20 +62,6 @@ class MLEngineModelVersionDetailsLink(BaseGoogleLink):
     key = "ml_engine_version_details"
     format_str = MLENGINE_MODEL_VERSION_DETAILS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        model_id: str,
-        project_id: str,
-        version_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=MLEngineModelVersionDetailsLink.key,
-            value={"model_id": model_id, "project_id": project_id, "version_id": version_id},
-        )
-
 
 class MLEngineJobSListLink(BaseGoogleLink):
     """Helper class for constructing ML Engine link."""
@@ -127,15 +69,3 @@ class MLEngineJobSListLink(BaseGoogleLink):
     name = "MLEngine Jobs List"
     key = "ml_engine_jobs_list"
     format_str = MLENGINE_JOBS_LIST_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=MLEngineJobSListLink.key,
-            value={"project_id": project_id},
-        )

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

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.utils.context import Context
-
 PUBSUB_BASE_LINK = "/cloudpubsub"
 PUBSUB_TOPIC_LINK = PUBSUB_BASE_LINK + "/topic/detail/{topic_id}?project={project_id}"
 PUBSUB_SUBSCRIPTION_LINK = PUBSUB_BASE_LINK + "/subscription/detail/{subscription_id}?project={project_id}"
@@ -39,19 +33,6 @@ class PubSubTopicLink(BaseGoogleLink):
     key = "pubsub_topic"
     format_str = PUBSUB_TOPIC_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        topic_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=PubSubTopicLink.key,
-            value={"topic_id": topic_id, "project_id": project_id},
-        )
-
 
 class PubSubSubscriptionLink(BaseGoogleLink):
     """Helper class for constructing Pub/Sub Subscription Link."""
@@ -59,16 +40,3 @@ class PubSubSubscriptionLink(BaseGoogleLink):
     name = "Pub/Sub Subscription"
     key = "pubsub_subscription"
     format_str = PUBSUB_SUBSCRIPTION_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        subscription_id: str | None,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=PubSubSubscriptionLink.key,
-            value={"subscription_id": subscription_id, "project_id": project_id},
-        )

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

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.utils.context import Context
-
 SPANNER_BASE_LINK = "/spanner/instances"
 SPANNER_INSTANCE_LINK = SPANNER_BASE_LINK + "/{instance_id}/details/databases?project={project_id}"
 SPANNER_DATABASE_LINK = (
@@ -41,19 +35,6 @@ class SpannerInstanceLink(BaseGoogleLink):
     key = "spanner_instance"
     format_str = SPANNER_INSTANCE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        instance_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=SpannerInstanceLink.key,
-            value={"instance_id": instance_id, "project_id": project_id},
-        )
-
 
 class SpannerDatabaseLink(BaseGoogleLink):
     """Helper class for constructing Spanner Database Link."""
@@ -61,17 +42,3 @@ class SpannerDatabaseLink(BaseGoogleLink):
     name = "Spanner Database"
     key = "spanner_database"
     format_str = SPANNER_DATABASE_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        instance_id: str,
-        database_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=SpannerDatabaseLink.key,
-            value={"instance_id": instance_id, "database_id": database_id, "project_id": project_id},
-        )

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

@@ -19,14 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.utils.context import Context
-
 STACKDRIVER_BASE_LINK = "/monitoring/alerting"
 STACKDRIVER_NOTIFICATIONS_LINK = STACKDRIVER_BASE_LINK + "/notifications?project={project_id}"
 STACKDRIVER_POLICIES_LINK = STACKDRIVER_BASE_LINK + "/policies?project={project_id}"
@@ -39,18 +33,6 @@ class StackdriverNotificationsLink(BaseGoogleLink):
     key = "stackdriver_notifications"
     format_str = STACKDRIVER_NOTIFICATIONS_LINK
 
-    @staticmethod
-    def persist(
-        operator_instance: BaseOperator,
-        context: Context,
-        project_id: str | None,
-    ):
-        operator_instance.xcom_push(
-            context,
-            key=StackdriverNotificationsLink.key,
-            value={"project_id": project_id},
-        )
-
 
 class StackdriverPoliciesLink(BaseGoogleLink):
     """Helper class for constructing Stackdriver Policies Link."""
@@ -58,15 +40,3 @@ class StackdriverPoliciesLink(BaseGoogleLink):
     name = "Cloud Monitoring Policies"
     key = "stackdriver_policies"
     format_str = STACKDRIVER_POLICIES_LINK
-
-    @staticmethod
-    def persist(
-        operator_instance: BaseOperator,
-        context: Context,
-        project_id: str | None,
-    ):
-        operator_instance.xcom_push(
-            context,
-            key=StackdriverPoliciesLink.key,
-            value={"project_id": project_id},
-        )

airflow/providers/google/cloud/links/translate.py

@@ -70,19 +70,6 @@ class TranslationLegacyDatasetLink(BaseGoogleLink):
     key = "translation_legacy_dataset"
     format_str = TRANSLATION_LEGACY_DATASET_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        dataset_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslationLegacyDatasetLink.key,
-            value={"location": task_instance.location, "dataset_id": dataset_id, "project_id": project_id},
-        )
-
 
 class TranslationDatasetListLink(BaseGoogleLink):
     """Helper class for constructing Translation Dataset List link."""
@@ -91,20 +78,6 @@ class TranslationDatasetListLink(BaseGoogleLink):
     key = "translation_dataset_list"
     format_str = TRANSLATION_DATASET_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslationDatasetListLink.key,
-            value={
-                "project_id": project_id,
-            },
-        )
-
 
 class TranslationLegacyModelLink(BaseGoogleLink):
     """
@@ -117,25 +90,6 @@ class TranslationLegacyModelLink(BaseGoogleLink):
     key = "translation_legacy_model"
     format_str = TRANSLATION_LEGACY_MODEL_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        dataset_id: str,
-        model_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslationLegacyModelLink.key,
-            value={
-                "location": task_instance.location,
-                "dataset_id": dataset_id,
-                "model_id": model_id,
-                "project_id": project_id,
-            },
-        )
-
 
 class TranslationLegacyModelTrainLink(BaseGoogleLink):
     """
@@ -148,22 +102,6 @@ class TranslationLegacyModelTrainLink(BaseGoogleLink):
     key = "translation_legacy_model_train"
     format_str = TRANSLATION_LEGACY_MODEL_TRAIN_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslationLegacyModelTrainLink.key,
-            value={
-                "location": task_instance.location,
-                "dataset_id": task_instance.model["dataset_id"],
-                "project_id": project_id,
-            },
-        )
-
 
 class TranslationLegacyModelPredictLink(BaseGoogleLink):
     """
@@ -176,25 +114,6 @@ class TranslationLegacyModelPredictLink(BaseGoogleLink):
     key = "translation_legacy_model_predict"
     format_str = TRANSLATION_LEGACY_MODEL_PREDICT_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        model_id: str,
-        project_id: str,
-        dataset_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslationLegacyModelPredictLink.key,
-            value={
-                "location": task_instance.location,
-                "dataset_id": dataset_id,
-                "model_id": model_id,
-                "project_id": project_id,
-            },
-        )
-
 
 class TranslateTextBatchLink(BaseGoogleLink):
     """
@@ -212,20 +131,13 @@ class TranslateTextBatchLink(BaseGoogleLink):
     def extract_output_uri_prefix(output_config):
         return output_config["gcs_destination"]["output_uri_prefix"].rpartition("gs://")[-1]
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-        output_config: dict,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslateTextBatchLink.key,
-            value={
-                "project_id": project_id,
-                "output_uri_prefix": TranslateTextBatchLink.extract_output_uri_prefix(output_config),
-            },
+    @classmethod
+    def persist(cls, context: Context, **value):
+        output_config = value.get("output_config")
+        super().persist(
+            context=context,
+            project_id=value.get("project_id"),
+            output_uri_prefix=cls.extract_output_uri_prefix(output_config),
         )
 
 
@@ -237,22 +149,9 @@ class TranslationNativeDatasetLink(BaseGoogleLink):
     """
 
     name = "Translation Native Dataset"
-    key = "translation_naive_dataset"
+    key = "translation_native_dataset"
     format_str = TRANSLATION_NATIVE_DATASET_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        dataset_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslationNativeDatasetLink.key,
-            value={"location": task_instance.location, "dataset_id": dataset_id, "project_id": project_id},
-        )
-
 
 class TranslationDatasetsListLink(BaseGoogleLink):
     """
@@ -265,20 +164,6 @@ class TranslationDatasetsListLink(BaseGoogleLink):
     key = "translation_dataset_list"
     format_str = TRANSLATION_DATASET_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslationDatasetsListLink.key,
-            value={
-                "project_id": project_id,
-            },
-        )
-
 
 class TranslationModelLink(BaseGoogleLink):
     """
@@ -291,25 +176,6 @@ class TranslationModelLink(BaseGoogleLink):
     key = "translation_model"
     format_str = TRANSLATION_NATIVE_MODEL_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        dataset_id: str,
-        model_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslationLegacyModelLink.key,
-            value={
-                "location": task_instance.location,
-                "dataset_id": dataset_id,
-                "model_id": model_id,
-                "project_id": project_id,
-            },
-        )
-
 
 class TranslationModelsListLink(BaseGoogleLink):
     """
@@ -322,20 +188,6 @@ class TranslationModelsListLink(BaseGoogleLink):
     key = "translation_models_list"
     format_str = TRANSLATION_MODELS_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslationModelsListLink.key,
-            value={
-                "project_id": project_id,
-            },
-        )
-
 
 class TranslateResultByOutputConfigLink(BaseGoogleLink):
     """
@@ -353,22 +205,14 @@ class TranslateResultByOutputConfigLink(BaseGoogleLink):
     def extract_output_uri_prefix(output_config):
         return output_config["gcs_destination"]["output_uri_prefix"].rpartition("gs://")[-1]
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-        output_config: dict,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslateResultByOutputConfigLink.key,
-            value={
-                "project_id": project_id,
-                "output_uri_prefix": TranslateResultByOutputConfigLink.extract_output_uri_prefix(
-                    output_config
-                ),
-            },
+    @classmethod
+    def persist(cls, context: Context, **value):
+        output_config = value.get("output_config")
+        output_uri_prefix = cls.extract_output_uri_prefix(output_config)
+        super().persist(
+            context=context,
+            project_id=value.get("project_id"),
+            output_uri_prefix=output_uri_prefix,
         )
 
 
@@ -382,17 +226,3 @@ class TranslationGlossariesListLink(BaseGoogleLink):
     name = "Translation Glossaries List"
     key = "translation_glossaries_list"
     format_str = TRANSLATION_HUB_RESOURCES_LIST_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=TranslationGlossariesListLink.key,
-            value={
-                "project_id": project_id,
-            },
-        )