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

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

@@ -23,7 +23,7 @@ from typing import TYPE_CHECKING
 from airflow.providers.google.cloud.links.base import BASE_LINK, BaseGoogleLink
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 TRANSLATION_BASE_LINK = BASE_LINK + "/translation"
@@ -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,
-            },
-        )

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

@@ -21,7 +21,7 @@ from typing import TYPE_CHECKING
 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
 
 VERTEX_AI_BASE_LINK = "/vertex-ai"
 VERTEX_AI_MODEL_LINK = (
@@ -54,6 +54,10 @@ VERTEX_AI_PIPELINE_JOB_LINK = (
     VERTEX_AI_BASE_LINK + "/locations/{region}/pipelines/runs/{pipeline_id}?project={project_id}"
 )
 VERTEX_AI_PIPELINE_JOB_LIST_LINK = VERTEX_AI_BASE_LINK + "/pipelines/runs?project={project_id}"
+VERTEX_AI_RAY_CLUSTER_LINK = (
+    VERTEX_AI_BASE_LINK + "/locations/{location}/ray-clusters/{cluster_id}?project={project_id}"
+)
+VERTEX_AI_RAY_CLUSTER_LIST_LINK = VERTEX_AI_BASE_LINK + "/ray?project={project_id}"
 
 
 class VertexAIModelLink(BaseGoogleLink):
@@ -63,22 +67,6 @@ class VertexAIModelLink(BaseGoogleLink):
     key = "model_conf"
     format_str = VERTEX_AI_MODEL_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        model_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIModelLink.key,
-            value={
-                "model_id": model_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIModelListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Models Link."""
@@ -87,19 +75,6 @@ class VertexAIModelListLink(BaseGoogleLink):
     key = "models_conf"
     format_str = VERTEX_AI_MODEL_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIModelListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIModelExportLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Model Export Link."""
@@ -113,19 +88,15 @@ class VertexAIModelExportLink(BaseGoogleLink):
         """Return bucket name from output configuration."""
         return config["artifact_destination"]["output_uri_prefix"].rpartition("gs://")[-1]
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
+    @classmethod
+    def persist(cls, context: Context, **value):
+        output_config = value.get("output_config")
+        bucket_name = cls.extract_bucket_name(output_config)
+        super().persist(
             context=context,
-            key=VertexAIModelExportLink.key,
-            value={
-                "project_id": task_instance.project_id,
-                "model_id": task_instance.model_id,
-                "bucket_name": VertexAIModelExportLink.extract_bucket_name(task_instance.output_config),
-            },
+            project_id=value.get("project_id"),
+            model_id=value.get("model_id"),
+            bucket_name=bucket_name,
         )
 
 
@@ -136,22 +107,6 @@ class VertexAITrainingLink(BaseGoogleLink):
     key = "training_conf"
     format_str = VERTEX_AI_TRAINING_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        training_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAITrainingLink.key,
-            value={
-                "training_id": training_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAITrainingPipelinesLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Training Pipelines link."""
@@ -160,19 +115,6 @@ class VertexAITrainingPipelinesLink(BaseGoogleLink):
     key = "pipelines_conf"
     format_str = VERTEX_AI_TRAINING_PIPELINES_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAITrainingPipelinesLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIDatasetLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Dataset link."""
@@ -181,18 +123,6 @@ class VertexAIDatasetLink(BaseGoogleLink):
     key = "dataset_conf"
     format_str = VERTEX_AI_DATASET_LINK
 
-    @staticmethod
-    def persist(context: Context, task_instance, dataset_id: str):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIDatasetLink.key,
-            value={
-                "dataset_id": dataset_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIDatasetListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Datasets Link."""
@@ -201,19 +131,6 @@ class VertexAIDatasetListLink(BaseGoogleLink):
     key = "datasets_conf"
     format_str = VERTEX_AI_DATASET_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIDatasetListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIHyperparameterTuningJobListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI HyperparameterTuningJobs Link."""
@@ -222,19 +139,6 @@ class VertexAIHyperparameterTuningJobListLink(BaseGoogleLink):
     key = "hyperparameter_tuning_jobs_conf"
     format_str = VERTEX_AI_HYPERPARAMETER_TUNING_JOB_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIHyperparameterTuningJobListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIBatchPredictionJobLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI BatchPredictionJob link."""
@@ -243,22 +147,6 @@ class VertexAIBatchPredictionJobLink(BaseGoogleLink):
     key = "batch_prediction_job_conf"
     format_str = VERTEX_AI_BATCH_PREDICTION_JOB_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        batch_prediction_job_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIBatchPredictionJobLink.key,
-            value={
-                "batch_prediction_job_id": batch_prediction_job_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIBatchPredictionJobListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI BatchPredictionJobList link."""
@@ -267,19 +155,6 @@ class VertexAIBatchPredictionJobListLink(BaseGoogleLink):
     key = "batch_prediction_jobs_conf"
     format_str = VERTEX_AI_BATCH_PREDICTION_JOB_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIBatchPredictionJobListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIEndpointLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Endpoint link."""
@@ -288,22 +163,6 @@ class VertexAIEndpointLink(BaseGoogleLink):
     key = "endpoint_conf"
     format_str = VERTEX_AI_ENDPOINT_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        endpoint_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIEndpointLink.key,
-            value={
-                "endpoint_id": endpoint_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIEndpointListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI EndpointList link."""
@@ -312,19 +171,6 @@ class VertexAIEndpointListLink(BaseGoogleLink):
     key = "endpoints_conf"
     format_str = VERTEX_AI_ENDPOINT_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIEndpointListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIPipelineJobLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI PipelineJob link."""
@@ -333,22 +179,6 @@ class VertexAIPipelineJobLink(BaseGoogleLink):
     key = "pipeline_job_conf"
     format_str = VERTEX_AI_PIPELINE_JOB_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        pipeline_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIPipelineJobLink.key,
-            value={
-                "pipeline_id": pipeline_id,
-                "region": task_instance.region,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIPipelineJobListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI PipelineJobList link."""
@@ -357,15 +187,18 @@ class VertexAIPipelineJobListLink(BaseGoogleLink):
     key = "pipeline_job_list_conf"
     format_str = VERTEX_AI_PIPELINE_JOB_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIPipelineJobListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )
+
+class VertexAIRayClusterLink(BaseGoogleLink):
+    """Helper class for constructing Vertex AI Ray Cluster link."""
+
+    name = "Ray Cluster"
+    key = "ray_cluster_conf"
+    format_str = VERTEX_AI_RAY_CLUSTER_LINK
+
+
+class VertexAIRayClusterListLink(BaseGoogleLink):
+    """Helper class for constructing Vertex AI Ray Cluster List link."""
+
+    name = "Ray Cluster List"
+    key = "ray_cluster_list_conf"
+    format_str = VERTEX_AI_RAY_CLUSTER_LIST_LINK

airflow/providers/google/cloud/links/workflows.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
-
 WORKFLOWS_BASE_LINK = "/workflows"
 WORKFLOW_LINK = WORKFLOWS_BASE_LINK + "/workflow/{location_id}/{workflow_id}/executions?project={project_id}"
 WORKFLOWS_LINK = WORKFLOWS_BASE_LINK + "?project={project_id}"
@@ -43,20 +37,6 @@ class WorkflowsWorkflowDetailsLink(BaseGoogleLink):
     key = "workflow_details"
     format_str = WORKFLOW_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        location_id: str,
-        workflow_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=WorkflowsWorkflowDetailsLink.key,
-            value={"location_id": location_id, "workflow_id": workflow_id, "project_id": project_id},
-        )
-
 
 class WorkflowsListOfWorkflowsLink(BaseGoogleLink):
     """Helper class for constructing list of Workflows Link."""
@@ -65,18 +45,6 @@ class WorkflowsListOfWorkflowsLink(BaseGoogleLink):
     key = "list_of_workflows"
     format_str = WORKFLOWS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=WorkflowsListOfWorkflowsLink.key,
-            value={"project_id": project_id},
-        )
-
 
 class WorkflowsExecutionLink(BaseGoogleLink):
     """Helper class for constructing Workflows Execution Link."""
@@ -84,23 +52,3 @@ class WorkflowsExecutionLink(BaseGoogleLink):
     name = "Workflow Execution"
     key = "workflow_execution"
     format_str = EXECUTION_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        location_id: str,
-        workflow_id: str,
-        execution_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=WorkflowsExecutionLink.key,
-            value={
-                "location_id": location_id,
-                "workflow_id": workflow_id,
-                "execution_id": execution_id,
-                "project_id": project_id,
-            },
-        )