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

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

@@ -67,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."""
@@ -91,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."""
@@ -117,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,
         )
 
 
@@ -140,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."""
@@ -164,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."""
@@ -185,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."""
@@ -205,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."""
@@ -226,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."""
@@ -247,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."""
@@ -271,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."""
@@ -292,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."""
@@ -316,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."""
@@ -337,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."""
@@ -361,19 +187,6 @@ 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."""
@@ -382,22 +195,6 @@ class VertexAIRayClusterLink(BaseGoogleLink):
     key = "ray_cluster_conf"
     format_str = VERTEX_AI_RAY_CLUSTER_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        cluster_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIRayClusterLink.key,
-            value={
-                "location": task_instance.location,
-                "cluster_id": cluster_id,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class VertexAIRayClusterListLink(BaseGoogleLink):
     """Helper class for constructing Vertex AI Ray Cluster List link."""
@@ -405,16 +202,3 @@ class VertexAIRayClusterListLink(BaseGoogleLink):
     name = "Ray Cluster List"
     key = "ray_cluster_list_conf"
     format_str = VERTEX_AI_RAY_CLUSTER_LIST_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=VertexAIRayClusterListLink.key,
-            value={
-                "project_id": task_instance.project_id,
-            },
-        )

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

airflow/providers/google/cloud/log/gcs_task_handler.py

@@ -61,11 +61,11 @@ class GCSRemoteLogIO(LoggingMixin):  # noqa: D101
     remote_base: str
     base_log_folder: Path = attrs.field(converter=Path)
     delete_local_copy: bool
-    project_id: str
+    project_id: str | None = None
 
-    gcp_key_path: str | None
-    gcp_keyfile_dict: dict | None
-    scopes: Collection[str] | None
+    gcp_key_path: str | None = None
+    gcp_keyfile_dict: dict | None = None
+    scopes: Collection[str] | None = _DEFAULT_SCOPESS
 
     processors = ()
 

airflow/providers/google/cloud/operators/alloy_db.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 NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -228,15 +228,16 @@ class AlloyDBCreateClusterOperator(AlloyDBWriteBaseOperator):
             return result
         return None
 
-    def execute(self, context: Context) -> dict | None:
-        AlloyDBClusterLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
 
+    def execute(self, context: Context) -> dict | None:
+        AlloyDBClusterLink.persist(context=context)
         if cluster := self._get_cluster():
             return cluster
 
@@ -334,14 +335,16 @@ class AlloyDBUpdateClusterOperator(AlloyDBWriteBaseOperator):
         self.update_mask = update_mask
         self.allow_missing = allow_missing
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBClusterLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+        AlloyDBClusterLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update AlloyDB cluster request.")
         else:
@@ -545,14 +548,16 @@ class AlloyDBCreateInstanceOperator(AlloyDBWriteBaseOperator):
             return result
         return None
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBClusterLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+        AlloyDBClusterLink.persist(context=context)
         if instance := self._get_instance():
             return instance
 
@@ -654,14 +659,16 @@ class AlloyDBUpdateInstanceOperator(AlloyDBWriteBaseOperator):
         self.update_mask = update_mask
         self.allow_missing = allow_missing
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBClusterLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+        AlloyDBClusterLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update AlloyDB instance request.")
         else:
@@ -861,14 +868,16 @@ class AlloyDBCreateUserOperator(AlloyDBWriteBaseOperator):
             return result
         return None
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBUsersLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+        AlloyDBUsersLink.persist(context=context)
         if (_user := self._get_user()) is not None:
             return _user
 
@@ -968,14 +977,16 @@ class AlloyDBUpdateUserOperator(AlloyDBWriteBaseOperator):
         self.update_mask = update_mask
         self.allow_missing = allow_missing
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBUsersLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+        AlloyDBUsersLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update AlloyDB user request.")
         else:
@@ -1159,12 +1170,14 @@ class AlloyDBCreateBackupOperator(AlloyDBWriteBaseOperator):
             return result
         return None
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBBackupsLink.persist(
-            context=context,
-            task_instance=self,
-            project_id=self.project_id,
-        )
+        AlloyDBBackupsLink.persist(context=context)
         if backup := self._get_backup():
             return backup
 
@@ -1259,12 +1272,14 @@ class AlloyDBUpdateBackupOperator(AlloyDBWriteBaseOperator):
         self.update_mask = update_mask
         self.allow_missing = allow_missing
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBBackupsLink.persist(
-            context=context,
-            task_instance=self,
-            project_id=self.project_id,
-        )
+        AlloyDBBackupsLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update AlloyDB backup request.")
         else: