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

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

@@ -19,13 +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
-
 DATAPLEX_BASE_LINK = "/dataplex/process/tasks"
 DATAPLEX_TASK_LINK = DATAPLEX_BASE_LINK + "/{lake_id}.{task_id};location={region}/jobs?project={project_id}"
 DATAPLEX_TASKS_LINK = DATAPLEX_BASE_LINK + "?project={project_id}&qLake={lake_id}.{region}"
@@ -53,22 +48,6 @@ class DataplexTaskLink(BaseGoogleLink):
     key = "task_conf"
     format_str = DATAPLEX_TASK_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataplexTaskLink.key,
-            value={
-                "lake_id": task_instance.lake_id,
-                "task_id": task_instance.dataplex_task_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class DataplexTasksLink(BaseGoogleLink):
     """Helper class for constructing Dataplex Tasks link."""
@@ -77,21 +56,6 @@ class DataplexTasksLink(BaseGoogleLink):
     key = "tasks_conf"
     format_str = DATAPLEX_TASKS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataplexTasksLink.key,
-            value={
-                "project_id": task_instance.project_id,
-                "lake_id": task_instance.lake_id,
-                "region": task_instance.region,
-            },
-        )
-
 
 class DataplexLakeLink(BaseGoogleLink):
     """Helper class for constructing Dataplex Lake link."""
@@ -100,21 +64,6 @@ class DataplexLakeLink(BaseGoogleLink):
     key = "dataplex_lake_key"
     format_str = DATAPLEX_LAKE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataplexLakeLink.key,
-            value={
-                "lake_id": task_instance.lake_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class DataplexCatalogEntryGroupLink(BaseGoogleLink):
     """Helper class for constructing Dataplex Catalog EntryGroup link."""
@@ -123,21 +72,6 @@ class DataplexCatalogEntryGroupLink(BaseGoogleLink):
     key = "dataplex_catalog_entry_group_key"
     format_str = DATAPLEX_CATALOG_ENTRY_GROUP_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataplexCatalogEntryGroupLink.key,
-            value={
-                "entry_group_id": task_instance.entry_group_id,
-                "location": task_instance.location,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class DataplexCatalogEntryGroupsLink(BaseGoogleLink):
     """Helper class for constructing Dataplex Catalog EntryGroups link."""
@@ -146,20 +80,6 @@ class DataplexCatalogEntryGroupsLink(BaseGoogleLink):
     key = "dataplex_catalog_entry_groups_key"
     format_str = DATAPLEX_CATALOG_ENTRY_GROUPS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataplexCatalogEntryGroupsLink.key,
-            value={
-                "location": task_instance.location,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class DataplexCatalogEntryTypeLink(BaseGoogleLink):
     """Helper class for constructing Dataplex Catalog EntryType link."""
@@ -168,21 +88,6 @@ class DataplexCatalogEntryTypeLink(BaseGoogleLink):
     key = "dataplex_catalog_entry_type_key"
     format_str = DATAPLEX_CATALOG_ENTRY_TYPE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataplexCatalogEntryTypeLink.key,
-            value={
-                "entry_type_id": task_instance.entry_type_id,
-                "location": task_instance.location,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class DataplexCatalogEntryTypesLink(BaseGoogleLink):
     """Helper class for constructing Dataplex Catalog EntryTypes link."""
@@ -191,20 +96,6 @@ class DataplexCatalogEntryTypesLink(BaseGoogleLink):
     key = "dataplex_catalog_entry_types_key"
     format_str = DATAPLEX_CATALOG_ENTRY_TYPES_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataplexCatalogEntryTypesLink.key,
-            value={
-                "location": task_instance.location,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class DataplexCatalogAspectTypeLink(BaseGoogleLink):
     """Helper class for constructing Dataplex Catalog AspectType link."""
@@ -213,21 +104,6 @@ class DataplexCatalogAspectTypeLink(BaseGoogleLink):
     key = "dataplex_catalog_aspect_type_key"
     format_str = DATAPLEX_CATALOG_ASPECT_TYPE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataplexCatalogAspectTypeLink.key,
-            value={
-                "aspect_type_id": task_instance.aspect_type_id,
-                "location": task_instance.location,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class DataplexCatalogAspectTypesLink(BaseGoogleLink):
     """Helper class for constructing Dataplex Catalog AspectTypes link."""
@@ -236,20 +112,6 @@ class DataplexCatalogAspectTypesLink(BaseGoogleLink):
     key = "dataplex_catalog_aspect_types_key"
     format_str = DATAPLEX_CATALOG_ASPECT_TYPES_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataplexCatalogAspectTypesLink.key,
-            value={
-                "location": task_instance.location,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class DataplexCatalogEntryLink(BaseGoogleLink):
     """Helper class for constructing Dataplex Catalog Entry link."""
@@ -257,19 +119,3 @@ class DataplexCatalogEntryLink(BaseGoogleLink):
     name = "Dataplex Catalog Entry"
     key = "dataplex_catalog_entry_key"
     format_str = DATAPLEX_CATALOG_ENTRY_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=DataplexCatalogEntryLink.key,
-            value={
-                "entry_id": task_instance.entry_id,
-                "entry_group_id": task_instance.entry_group_id,
-                "location": task_instance.location,
-                "project_id": task_instance.project_id,
-            },
-        )

airflow/providers/google/cloud/links/dataprep.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
-
 BASE_LINK = "https://clouddataprep.com"
 DATAPREP_FLOW_LINK = BASE_LINK + "/flows/{flow_id}?projectId={project_id}"
 DATAPREP_JOB_GROUP_LINK = BASE_LINK + "/jobs/{job_group_id}?projectId={project_id}"
@@ -35,14 +30,6 @@ class DataprepFlowLink(BaseGoogleLink):
     key = "dataprep_flow_page"
     format_str = DATAPREP_FLOW_LINK
 
-    @staticmethod
-    def persist(context: Context, task_instance, project_id: str, flow_id: int):
-        task_instance.xcom_push(
-            context=context,
-            key=DataprepFlowLink.key,
-            value={"project_id": project_id, "flow_id": flow_id},
-        )
-
 
 class DataprepJobGroupLink(BaseGoogleLink):
     """Helper class for constructing Dataprep job group link."""
@@ -50,14 +37,3 @@ class DataprepJobGroupLink(BaseGoogleLink):
     name = "Job group details page"
     key = "dataprep_job_group_page"
     format_str = DATAPREP_JOB_GROUP_LINK
-
-    @staticmethod
-    def persist(context: Context, task_instance, project_id: str, job_group_id: int):
-        task_instance.xcom_push(
-            context=context,
-            key=DataprepJobGroupLink.key,
-            value={
-                "project_id": project_id,
-                "job_group_id": job_group_id,
-            },
-        )

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

@@ -25,20 +25,13 @@ from typing import TYPE_CHECKING, Any
 import attr
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import BaseOperatorLink, XCom
 from airflow.providers.google.cloud.links.base import BASE_LINK, BaseGoogleLink
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
-    from airflow.models import BaseOperator
     from airflow.models.taskinstancekey import TaskInstanceKey
-    from airflow.utils.context import Context
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseOperatorLink
-    from airflow.sdk.execution_time.xcom import XCom
-else:
-    from airflow.models import XCom  # type: ignore[no-redef]
-    from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
+    from airflow.providers.common.compat.sdk import Context
+    from airflow.providers.google.version_compat import BaseOperator
 
 
 def __getattr__(name: str) -> Any:
@@ -94,16 +87,16 @@ class DataprocLink(BaseOperatorLink):
     @staticmethod
     def persist(
         context: Context,
-        task_instance,
         url: str,
         resource: str,
+        region: str,
+        project_id: str,
     ):
-        task_instance.xcom_push(
-            context=context,
+        context["task_instance"].xcom_push(
             key=DataprocLink.key,
             value={
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
+                "region": region,
+                "project_id": project_id,
                 "url": url,
                 "resource": resource,
             },
@@ -147,14 +140,13 @@ class DataprocListLink(BaseOperatorLink):
     @staticmethod
     def persist(
         context: Context,
-        task_instance,
         url: str,
+        project_id: str,
     ):
-        task_instance.xcom_push(
-            context=context,
+        context["task_instance"].xcom_push(
             key=DataprocListLink.key,
             value={
-                "project_id": task_instance.project_id,
+                "project_id": project_id,
                 "url": url,
             },
         )
@@ -189,20 +181,6 @@ class DataprocClusterLink(BaseGoogleLink):
     key = "dataproc_cluster"
     format_str = DATAPROC_CLUSTER_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        operator: BaseOperator,
-        cluster_id: str,
-        region: str,
-        project_id: str,
-    ):
-        operator.xcom_push(
-            context,
-            key=DataprocClusterLink.key,
-            value={"cluster_id": cluster_id, "region": region, "project_id": project_id},
-        )
-
 
 class DataprocJobLink(BaseGoogleLink):
     """Helper class for constructing Dataproc Job Link."""
@@ -211,20 +189,6 @@ class DataprocJobLink(BaseGoogleLink):
     key = "dataproc_job"
     format_str = DATAPROC_JOB_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        operator: BaseOperator,
-        job_id: str,
-        region: str,
-        project_id: str,
-    ):
-        operator.xcom_push(
-            context,
-            key=DataprocJobLink.key,
-            value={"job_id": job_id, "region": region, "project_id": project_id},
-        )
-
 
 class DataprocWorkflowLink(BaseGoogleLink):
     """Helper class for constructing Dataproc Workflow Link."""
@@ -233,14 +197,6 @@ class DataprocWorkflowLink(BaseGoogleLink):
     key = "dataproc_workflow"
     format_str = DATAPROC_WORKFLOW_LINK
 
-    @staticmethod
-    def persist(context: Context, operator: BaseOperator, workflow_id: str, project_id: str, region: str):
-        operator.xcom_push(
-            context,
-            key=DataprocWorkflowLink.key,
-            value={"workflow_id": workflow_id, "region": region, "project_id": project_id},
-        )
-
 
 class DataprocWorkflowTemplateLink(BaseGoogleLink):
     """Helper class for constructing Dataproc Workflow Template Link."""
@@ -249,20 +205,6 @@ class DataprocWorkflowTemplateLink(BaseGoogleLink):
     key = "dataproc_workflow_template"
     format_str = DATAPROC_WORKFLOW_TEMPLATE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        operator: BaseOperator,
-        workflow_template_id: str,
-        project_id: str,
-        region: str,
-    ):
-        operator.xcom_push(
-            context,
-            key=DataprocWorkflowTemplateLink.key,
-            value={"workflow_template_id": workflow_template_id, "region": region, "project_id": project_id},
-        )
-
 
 class DataprocBatchLink(BaseGoogleLink):
     """Helper class for constructing Dataproc Batch Link."""
@@ -271,20 +213,6 @@ class DataprocBatchLink(BaseGoogleLink):
     key = "dataproc_batch"
     format_str = DATAPROC_BATCH_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        operator: BaseOperator,
-        batch_id: str,
-        project_id: str,
-        region: str,
-    ):
-        operator.xcom_push(
-            context,
-            key=DataprocBatchLink.key,
-            value={"batch_id": batch_id, "region": region, "project_id": project_id},
-        )
-
 
 class DataprocBatchesListLink(BaseGoogleLink):
     """Helper class for constructing Dataproc Batches List Link."""
@@ -292,15 +220,3 @@ class DataprocBatchesListLink(BaseGoogleLink):
     name = "Dataproc Batches List"
     key = "dataproc_batches_list"
     format_str = DATAPROC_BATCHES_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        operator: BaseOperator,
-        project_id: str,
-    ):
-        operator.xcom_push(
-            context,
-            key=DataprocBatchesListLink.key,
-            value={"project_id": project_id},
-        )

airflow/providers/google/cloud/links/datastore.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
-
 DATASTORE_BASE_LINK = "/datastore"
 DATASTORE_IMPORT_EXPORT_LINK = DATASTORE_BASE_LINK + "/import-export?project={project_id}"
 DATASTORE_EXPORT_ENTITIES_LINK = "/storage/browser/{bucket_name}/{export_name}?project={project_id}"
@@ -36,19 +31,6 @@ class CloudDatastoreImportExportLink(BaseGoogleLink):
     key = "import_export_conf"
     format_str = DATASTORE_IMPORT_EXPORT_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=CloudDatastoreImportExportLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class CloudDatastoreEntitiesLink(BaseGoogleLink):
     """Helper class for constructing Cloud Datastore Entities Link."""
@@ -56,16 +38,3 @@ class CloudDatastoreEntitiesLink(BaseGoogleLink):
     name = "Entities"
     key = "entities_conf"
     format_str = DATASTORE_ENTITIES_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=CloudDatastoreEntitiesLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )

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

@@ -24,7 +24,7 @@ from google.cloud.container_v1.types import Cluster
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 KUBERNETES_BASE_LINK = "/kubernetes"
 KUBERNETES_CLUSTER_LINK = (
@@ -51,19 +51,18 @@ class KubernetesEngineClusterLink(BaseGoogleLink):
     key = "kubernetes_cluster_conf"
     format_str = KUBERNETES_CLUSTER_LINK
 
-    @staticmethod
-    def persist(context: Context, task_instance, cluster: dict | Cluster | None):
+    @classmethod
+    def persist(cls, context: Context, **value):
+        cluster = value.get("cluster")
         if isinstance(cluster, dict):
             cluster = Cluster.from_json(json.dumps(cluster))
 
-        task_instance.xcom_push(
+        if not cluster:
+            raise ValueError("Cluster must be provided for KubernetesEngineClusterLink.")
+
+        super().persist(
             context=context,
-            key=KubernetesEngineClusterLink.key,
-            value={
-                "location": task_instance.location,
-                "cluster_name": cluster.name,  # type: ignore
-                "project_id": task_instance.project_id,
-            },
+            cluster_name=cluster.name,
         )
 
 
@@ -74,23 +73,6 @@ class KubernetesEnginePodLink(BaseGoogleLink):
     key = "kubernetes_pod_conf"
     format_str = KUBERNETES_POD_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=KubernetesEnginePodLink.key,
-            value={
-                "location": task_instance.location,
-                "cluster_name": task_instance.cluster_name,
-                "namespace": task_instance.pod.metadata.namespace,
-                "pod_name": task_instance.pod.metadata.name,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class KubernetesEngineJobLink(BaseGoogleLink):
     """Helper class for constructing Kubernetes Engine Job Link."""
@@ -99,23 +81,6 @@ class KubernetesEngineJobLink(BaseGoogleLink):
     key = "kubernetes_job_conf"
     format_str = KUBERNETES_JOB_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=KubernetesEngineJobLink.key,
-            value={
-                "location": task_instance.location,
-                "cluster_name": task_instance.cluster_name,
-                "namespace": task_instance.job.metadata.namespace,
-                "job_name": task_instance.job.metadata.name,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class KubernetesEngineWorkloadsLink(BaseGoogleLink):
     """Helper class for constructing Kubernetes Engine Workloads Link."""
@@ -123,19 +88,3 @@ class KubernetesEngineWorkloadsLink(BaseGoogleLink):
     name = "Kubernetes Workloads"
     key = "kubernetes_workloads_conf"
     format_str = KUBERNETES_WORKLOADS_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=KubernetesEngineWorkloadsLink.key,
-            value={
-                "location": task_instance.location,
-                "cluster_name": task_instance.cluster_name,
-                "namespace": task_instance.namespace,
-                "project_id": task_instance.project_id,
-            },
-        )

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,
-            },
-        )