apache-airflow-providers-google 12.0.0rc2__py3-none-any.whl → 14.0.0__py3-none-any.whl

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

@@ -24,7 +24,7 @@ from datetime import timedelta
 from typing import TYPE_CHECKING
 
 import vertexai
-from vertexai.generative_models import GenerativeModel, Part
+from vertexai.generative_models import GenerativeModel
 from vertexai.language_models import TextEmbeddingModel, TextGenerationModel
 from vertexai.preview.caching import CachedContent
 from vertexai.preview.evaluation import EvalResult, EvalTask
@@ -100,186 +100,6 @@ class GenerativeModelHook(GoogleBaseHook):
         cached_context_model = preview_generative_model.from_cached_content(cached_content)
         return cached_context_model
 
-    @deprecated(
-        planned_removal_date="January 01, 2025",
-        use_instead="Part objects included in contents parameter of "
-        "airflow.providers.google.cloud.hooks.generative_model."
-        "GenerativeModelHook.generative_model_generate_content",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_generative_model_part(self, content_gcs_path: str, content_mime_type: str | None = None) -> Part:
-        """Return a Generative Model Part object."""
-        part = Part.from_uri(content_gcs_path, mime_type=content_mime_type)
-        return part
-
-    @deprecated(
-        planned_removal_date="January 01, 2025",
-        use_instead="airflow.providers.google.cloud.hooks.generative_model."
-        "GenerativeModelHook.text_generation_model_predict",
-        category=AirflowProviderDeprecationWarning,
-    )
-    @GoogleBaseHook.fallback_to_default_project_id
-    def prompt_language_model(
-        self,
-        prompt: str,
-        pretrained_model: str,
-        temperature: float,
-        max_output_tokens: int,
-        top_p: float,
-        top_k: int,
-        location: str,
-        project_id: str = PROVIDE_PROJECT_ID,
-    ) -> str:
-        """
-        Use the Vertex AI PaLM API to generate natural language text.
-
-        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-        :param location: Required. The ID of the Google Cloud location that the service belongs to.
-        :param prompt: Required. Inputs or queries that a user or a program gives
-            to the Vertex AI PaLM API, in order to elicit a specific response.
-        :param pretrained_model: A pre-trained model optimized for performing natural
-            language tasks such as classification, summarization, extraction, content
-            creation, and ideation.
-        :param temperature: Temperature controls the degree of randomness in token
-            selection.
-        :param max_output_tokens: Token limit determines the maximum amount of text
-            output.
-        :param top_p: Tokens are selected from most probable to least until the sum
-            of their probabilities equals the top_p value. Defaults to 0.8.
-        :param top_k: A top_k of 1 means the selected token is the most probable
-            among all tokens.
-        """
-        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
-
-        parameters = {
-            "temperature": temperature,
-            "max_output_tokens": max_output_tokens,
-            "top_p": top_p,
-            "top_k": top_k,
-        }
-
-        model = self.get_text_generation_model(pretrained_model)
-
-        response = model.predict(
-            prompt=prompt,
-            **parameters,
-        )
-        return response.text
-
-    @deprecated(
-        planned_removal_date="January 01, 2025",
-        use_instead="airflow.providers.google.cloud.hooks.generative_model."
-        "GenerativeModelHook.text_embedding_model_get_embeddings",
-        category=AirflowProviderDeprecationWarning,
-    )
-    @GoogleBaseHook.fallback_to_default_project_id
-    def generate_text_embeddings(
-        self,
-        prompt: str,
-        pretrained_model: str,
-        location: str,
-        project_id: str = PROVIDE_PROJECT_ID,
-    ) -> list:
-        """
-        Use the Vertex AI PaLM API to generate text embeddings.
-
-        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-        :param location: Required. The ID of the Google Cloud location that the service belongs to.
-        :param prompt: Required. Inputs or queries that a user or a program gives
-            to the Vertex AI PaLM API, in order to elicit a specific response.
-        :param pretrained_model: A pre-trained model optimized for generating text embeddings.
-        """
-        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
-        model = self.get_text_embedding_model(pretrained_model)
-
-        response = model.get_embeddings([prompt])[0]  # single prompt
-
-        return response.values
-
-    @deprecated(
-        planned_removal_date="January 01, 2025",
-        use_instead="airflow.providers.google.cloud.hooks.generative_model."
-        "GenerativeModelHook.generative_model_generate_content",
-        category=AirflowProviderDeprecationWarning,
-    )
-    @GoogleBaseHook.fallback_to_default_project_id
-    def prompt_multimodal_model(
-        self,
-        prompt: str,
-        location: str,
-        generation_config: dict | None = None,
-        safety_settings: dict | None = None,
-        pretrained_model: str = "gemini-pro",
-        project_id: str = PROVIDE_PROJECT_ID,
-    ) -> str:
-        """
-        Use the Vertex AI Gemini Pro foundation model to generate natural language text.
-
-        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-        :param location: Required. The ID of the Google Cloud location that the service belongs to.
-        :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.
-        :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
-            output text and code.
-        """
-        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
-
-        model = self.get_generative_model(pretrained_model)
-        response = model.generate_content(
-            contents=[prompt], generation_config=generation_config, safety_settings=safety_settings
-        )
-
-        return response.text
-
-    @deprecated(
-        planned_removal_date="January 01, 2025",
-        use_instead="airflow.providers.google.cloud.hooks.generative_model."
-        "GenerativeModelHook.generative_model_generate_content",
-        category=AirflowProviderDeprecationWarning,
-    )
-    @GoogleBaseHook.fallback_to_default_project_id
-    def prompt_multimodal_model_with_media(
-        self,
-        prompt: str,
-        location: str,
-        media_gcs_path: str,
-        mime_type: str,
-        generation_config: dict | None = None,
-        safety_settings: dict | None = None,
-        pretrained_model: str = "gemini-pro-vision",
-        project_id: str = PROVIDE_PROJECT_ID,
-    ) -> str:
-        """
-        Use the Vertex AI Gemini Pro foundation model to generate natural language text.
-
-        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-        :param location: Required. The ID of the Google Cloud location that the service belongs to.
-        :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.
-        :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
-            output text and code.
-        :param media_gcs_path: A GCS path to a content file such as an image or a video.
-            Can be passed to the multi-modal model as part of the prompt. Used with vision models.
-        :param mime_type: Validates the media type presented by the file in the media_gcs_path.
-        """
-        vertexai.init(project=project_id, location=location, credentials=self.get_credentials())
-
-        model = self.get_generative_model(pretrained_model)
-        part = self.get_generative_model_part(media_gcs_path, mime_type)
-        response = model.generate_content(
-            contents=[prompt, part], generation_config=generation_config, safety_settings=safety_settings
-        )
-
-        return response.text
-
     @deprecated(
         planned_removal_date="April 09, 2025",
         use_instead="GenerativeModelHook.generative_model_generate_content",

airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py

@@ -29,15 +29,14 @@ import asyncio
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import CustomJob, HyperparameterTuningJob, gapic, hyperparameter_tuning
 from google.cloud.aiplatform_v1 import JobServiceAsyncClient, JobServiceClient, JobState, types
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import AsyncRetry, Retry

airflow/providers/google/cloud/hooks/vertex_ai/model_service.py

@@ -23,13 +23,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import ModelServiceClient
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry

airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py

@@ -29,6 +29,9 @@ import asyncio
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform import PipelineJob
@@ -39,10 +42,6 @@ from google.cloud.aiplatform_v1 import (
     types,
 )
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import AsyncRetry, Retry

airflow/providers/google/cloud/hooks/vertex_ai/prediction_service.py

@@ -20,13 +20,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.aiplatform_v1 import PredictionServiceClient
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.cloud.aiplatform_v1.types import PredictResponse

airflow/providers/google/cloud/hooks/video_intelligence.py

@@ -22,6 +22,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.videointelligence_v1 import (
     Feature,
@@ -29,9 +31,6 @@ from google.cloud.videointelligence_v1 import (
     VideoIntelligenceServiceClient,
 )
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry

airflow/providers/google/cloud/hooks/vision.py

@@ -24,6 +24,9 @@ from copy import deepcopy
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, Callable
 
+from airflow.exceptions import AirflowException
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.vision_v1 import (
     AnnotateImageRequest,
@@ -36,10 +39,6 @@ from google.cloud.vision_v1 import (
 )
 from google.protobuf.json_format import MessageToDict
 
-from airflow.exceptions import AirflowException
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.protobuf import field_mask_pb2

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

@@ -19,13 +19,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.google.common.consts import CLIENT_INFO
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.workflows.executions_v1beta import Execution, ExecutionsClient
 from google.cloud.workflows_v1beta import Workflow, WorkflowsClient
 
-from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook
-
 if TYPE_CHECKING:
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry

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

@@ -31,6 +31,10 @@ ALLOY_DB_BASE_LINK = "/alloydb"
 ALLOY_DB_CLUSTER_LINK = (
     ALLOY_DB_BASE_LINK + "/locations/{location_id}/clusters/{cluster_id}?project={project_id}"
 )
+ALLOY_DB_USERS_LINK = (
+    ALLOY_DB_BASE_LINK + "/locations/{location_id}/clusters/{cluster_id}/users?project={project_id}"
+)
+ALLOY_DB_BACKUPS_LINK = ALLOY_DB_BASE_LINK + "/backups?project={project_id}"
 
 
 class AlloyDBClusterLink(BaseGoogleLink):
@@ -53,3 +57,45 @@ class AlloyDBClusterLink(BaseGoogleLink):
             key=AlloyDBClusterLink.key,
             value={"location_id": location_id, "cluster_id": cluster_id, "project_id": project_id},
         )
+
+
+class AlloyDBUsersLink(BaseGoogleLink):
+    """Helper class for constructing AlloyDB users Link."""
+
+    name = "AlloyDB Users"
+    key = "alloy_db_users"
+    format_str = ALLOY_DB_USERS_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance: BaseOperator,
+        location_id: str,
+        cluster_id: str,
+        project_id: str | None,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=AlloyDBUsersLink.key,
+            value={"location_id": location_id, "cluster_id": cluster_id, "project_id": project_id},
+        )
+
+
+class AlloyDBBackupsLink(BaseGoogleLink):
+    """Helper class for constructing AlloyDB backups Link."""
+
+    name = "AlloyDB Backups"
+    key = "alloy_db_backups"
+    format_str = ALLOY_DB_BACKUPS_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance: BaseOperator,
+        project_id: str | None,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=AlloyDBBackupsLink.key,
+            value={"project_id": project_id},
+        )

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

@@ -35,6 +35,9 @@ BIGQUERY_TABLE_LINK = (
     BIGQUERY_BASE_LINK
     + "?referrer=search&project={project_id}&d={dataset_id}&p={project_id}&page=table&t={table_id}"
 )
+BIGQUERY_JOB_DETAIL_LINK = (
+    BIGQUERY_BASE_LINK + "?project={project_id}&ws=!1m5!1m4!1m3!1s{project_id}!2s{job_id}!3s{location}"
+)
 
 
 class BigQueryDatasetLink(BaseGoogleLink):
@@ -78,3 +81,25 @@ class BigQueryTableLink(BaseGoogleLink):
             key=BigQueryTableLink.key,
             value={"dataset_id": dataset_id, "project_id": project_id, "table_id": table_id},
         )
+
+
+class BigQueryJobDetailLink(BaseGoogleLink):
+    """Helper class for constructing BigQuery Job Detail Link."""
+
+    name = "BigQuery Job Detail"
+    key = "bigquery_job_detail"
+    format_str = BIGQUERY_JOB_DETAIL_LINK
+
+    @staticmethod
+    def persist(
+        context: Context,
+        task_instance: BaseOperator,
+        project_id: str,
+        location: str,
+        job_id: str,
+    ):
+        task_instance.xcom_push(
+            context,
+            key=BigQueryJobDetailLink.key,
+            value={"project_id": project_id, "location": location, "job_id": job_id},
+        )

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

@@ -30,9 +30,20 @@ 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}"
 
-DATAPLEX_LAKE_LINK = (
-    "https://console.cloud.google.com/dataplex/lakes/{lake_id};location={region}?project={project_id}"
+DATAPLEX_LAKE_LINK = "/dataplex/lakes/{lake_id};location={region}?project={project_id}"
+DATAPLEX_CATALOG_ENTRY_GROUPS_LINK = "/dataplex/catalog/entry-groups?project={project_id}"
+DATAPLEX_CATALOG_ENTRY_GROUP_LINK = (
+    "/dataplex/projects/{project_id}/locations/{location}/entryGroups/{entry_group_id}?project={project_id}"
 )
+DATAPLEX_CATALOG_ENTRY_TYPE_LINK = (
+    "/dataplex/projects/{project_id}/locations/{location}/entryTypes/{entry_type_id}?project={project_id}"
+)
+DATAPLEX_CATALOG_ENTRY_TYPES_LINK = "/dataplex/catalog/entry-types?project={project_id}"
+DATAPLEX_CATALOG_ASPECT_TYPE_LINK = (
+    "/dataplex/projects/{project_id}/locations/{location}/aspectTypes/{aspect_type_id}?project={project_id}"
+)
+DATAPLEX_CATALOG_ASPECT_TYPES_LINK = "/dataplex/catalog/aspect-types?project={project_id}"
+DATAPLEX_CATALOG_ENTRY_LINK = "/dataplex/dp-entries/projects/{project_id}/locations/{location}/entryGroups/{entry_group_id}/entries/{entry_id}?project={project_id}"
 
 
 class DataplexTaskLink(BaseGoogleLink):
@@ -103,3 +114,162 @@ class DataplexLakeLink(BaseGoogleLink):
                 "project_id": task_instance.project_id,
             },
         )
+
+
+class DataplexCatalogEntryGroupLink(BaseGoogleLink):
+    """Helper class for constructing Dataplex Catalog EntryGroup link."""
+
+    name = "Dataplex Catalog EntryGroup"
+    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."""
+
+    name = "Dataplex Catalog EntryGroups"
+    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."""
+
+    name = "Dataplex Catalog EntryType"
+    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."""
+
+    name = "Dataplex Catalog EntryTypes"
+    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."""
+
+    name = "Dataplex Catalog AspectType"
+    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."""
+
+    name = "Dataplex Catalog AspectTypes"
+    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."""
+
+    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/kubernetes_engine.py

@@ -19,9 +19,8 @@ from __future__ import annotations
 import json
 from typing import TYPE_CHECKING
 
-from google.cloud.container_v1.types import Cluster
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
+from google.cloud.container_v1.types import Cluster
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context