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

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:

airflow/providers/google/cloud/operators/automl.py

@@ -153,21 +153,24 @@ class AutoMLTrainModelOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         if project_id:
             TranslationLegacyModelTrainLink.persist(
-                context=context, task_instance=self, project_id=project_id
+                context=context,
+                dataset_id=self.model["dataset_id"],
+                project_id=project_id,
+                location=self.location,
             )
         operation_result = hook.wait_for_operation(timeout=self.timeout, operation=operation)
         result = Model.to_dict(operation_result)
         model_id = hook.extract_object_id(result)
         self.log.info("Model is created, model_id: %s", model_id)
 
-        self.xcom_push(context, key="model_id", value=model_id)
+        context["task_instance"].xcom_push(key="model_id", value=model_id)
         if project_id:
             TranslationLegacyModelLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=self.model["dataset_id"] or "-",
                 model_id=model_id,
                 project_id=project_id,
+                location=self.location,
             )
         return result
 
@@ -313,10 +316,10 @@ class AutoMLPredictOperator(GoogleCloudBaseOperator):
         if project_id and self.model_id and dataset_id:
             TranslationLegacyModelPredictLink.persist(
                 context=context,
-                task_instance=self,
                 model_id=self.model_id,
                 dataset_id=dataset_id,
                 project_id=project_id,
+                location=self.location,
             )
         return PredictResponse.to_dict(result)
 
@@ -412,14 +415,14 @@ class AutoMLCreateDatasetOperator(GoogleCloudBaseOperator):
         dataset_id = hook.extract_object_id(result)
         self.log.info("Creating completed. Dataset id: %s", dataset_id)
 
-        self.xcom_push(context, key="dataset_id", value=dataset_id)
+        context["task_instance"].xcom_push(key="dataset_id", value=dataset_id)
         project_id = self.project_id or hook.project_id
         if project_id:
             TranslationLegacyDatasetLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=dataset_id,
                 project_id=project_id,
+                location=self.location,
             )
         return result
 
@@ -530,9 +533,9 @@ class AutoMLImportDataOperator(GoogleCloudBaseOperator):
         if project_id:
             TranslationLegacyDatasetLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=self.dataset_id,
                 project_id=project_id,
+                location=self.location,
             )
 
 
@@ -649,9 +652,9 @@ class AutoMLTablesListColumnSpecsOperator(GoogleCloudBaseOperator):
         if project_id:
             TranslationLegacyDatasetLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=self.dataset_id,
                 project_id=project_id,
+                location=self.location,
             )
         return result
 
@@ -749,9 +752,9 @@ class AutoMLTablesUpdateDatasetOperator(GoogleCloudBaseOperator):
         if project_id:
             TranslationLegacyDatasetLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=hook.extract_object_id(self.dataset),
                 project_id=project_id,
+                location=self.location,
             )
         return Dataset.to_dict(result)
 
@@ -845,10 +848,10 @@ class AutoMLGetModelOperator(GoogleCloudBaseOperator):
         if project_id:
             TranslationLegacyModelLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=model["dataset_id"],
                 model_id=self.model_id,
                 project_id=project_id,
+                location=self.location,
             )
         return model
 
@@ -1154,9 +1157,9 @@ class AutoMLTablesListTableSpecsOperator(GoogleCloudBaseOperator):
         if project_id:
             TranslationLegacyDatasetLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=self.dataset_id,
                 project_id=project_id,
+                location=self.location,
             )
         return result
 
@@ -1245,14 +1248,13 @@ class AutoMLListDatasetOperator(GoogleCloudBaseOperator):
             result.append(Dataset.to_dict(dataset))
         self.log.info("Datasets obtained.")
 
-        self.xcom_push(
-            context,
+        context["task_instance"].xcom_push(
             key="dataset_id_list",
             value=[hook.extract_object_id(d) for d in result],
         )
         project_id = self.project_id or hook.project_id
         if project_id:
-            TranslationDatasetListLink.persist(context=context, task_instance=self, project_id=project_id)
+            TranslationDatasetListLink.persist(context=context, project_id=project_id)
         return result
 
 

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

@@ -1324,7 +1324,6 @@ class BigQueryCreateTableOperator(GoogleCloudBaseOperator):
             if self._table:
                 persist_kwargs = {
                     "context": context,
-                    "task_instance": self,
                     "project_id": self._table.to_api_repr()["tableReference"]["projectId"],
                     "dataset_id": self._table.to_api_repr()["tableReference"]["datasetId"],
                     "table_id": self._table.to_api_repr()["tableReference"]["tableId"],
@@ -1343,7 +1342,6 @@ class BigQueryCreateTableOperator(GoogleCloudBaseOperator):
                 self.log.info(error_msg)
                 persist_kwargs = {
                     "context": context,
-                    "task_instance": self,
                     "project_id": self.project_id or bq_hook.project_id,
                     "dataset_id": self.dataset_id,
                     "table_id": self.table_id,
@@ -1608,7 +1606,6 @@ class BigQueryCreateEmptyTableOperator(GoogleCloudBaseOperator):
             if self._table:
                 persist_kwargs = {
                     "context": context,
-                    "task_instance": self,
                     "project_id": self._table.to_api_repr()["tableReference"]["projectId"],
                     "dataset_id": self._table.to_api_repr()["tableReference"]["datasetId"],
                     "table_id": self._table.to_api_repr()["tableReference"]["tableId"],
@@ -1627,7 +1624,6 @@ class BigQueryCreateEmptyTableOperator(GoogleCloudBaseOperator):
                 self.log.info(error_msg)
                 persist_kwargs = {
                     "context": context,
-                    "task_instance": self,
                     "project_id": self.project_id or bq_hook.project_id,
                     "dataset_id": self.dataset_id,
                     "table_id": self.table_id,
@@ -1898,7 +1894,6 @@ class BigQueryCreateExternalTableOperator(GoogleCloudBaseOperator):
             if self._table:
                 BigQueryTableLink.persist(
                     context=context,
-                    task_instance=self,
                     dataset_id=self._table.dataset_id,
                     project_id=self._table.project,
                     table_id=self._table.table_id,
@@ -1957,7 +1952,6 @@ class BigQueryCreateExternalTableOperator(GoogleCloudBaseOperator):
         if self._table:
             BigQueryTableLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=self._table.dataset_id,
                 project_id=self._table.project,
                 table_id=self._table.table_id,
@@ -2155,7 +2149,6 @@ class BigQueryCreateEmptyDatasetOperator(GoogleCloudBaseOperator):
             )
             persist_kwargs = {
                 "context": context,
-                "task_instance": self,
                 "project_id": dataset["datasetReference"]["projectId"],
                 "dataset_id": dataset["datasetReference"]["datasetId"],
             }
@@ -2167,7 +2160,6 @@ class BigQueryCreateEmptyDatasetOperator(GoogleCloudBaseOperator):
             )
             persist_kwargs = {
                 "context": context,
-                "task_instance": self,
                 "project_id": project_id,
                 "dataset_id": dataset_id,
             }
@@ -2239,7 +2231,6 @@ class BigQueryGetDatasetOperator(GoogleCloudBaseOperator):
         dataset_api_repr = dataset.to_api_repr()
         BigQueryDatasetLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=dataset_api_repr["datasetReference"]["datasetId"],
             project_id=dataset_api_repr["datasetReference"]["projectId"],
         )
@@ -2388,7 +2379,6 @@ class BigQueryUpdateTableOperator(GoogleCloudBaseOperator):
         if self._table:
             BigQueryTableLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=self._table["tableReference"]["datasetId"],
                 project_id=self._table["tableReference"]["projectId"],
                 table_id=self._table["tableReference"]["tableId"],
@@ -2491,7 +2481,6 @@ class BigQueryUpdateDatasetOperator(GoogleCloudBaseOperator):
         dataset_api_repr = dataset.to_api_repr()
         BigQueryDatasetLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=dataset_api_repr["datasetReference"]["datasetId"],
             project_id=dataset_api_repr["datasetReference"]["projectId"],
         )
@@ -2663,7 +2652,6 @@ class BigQueryUpsertTableOperator(GoogleCloudBaseOperator):
         if self._table:
             BigQueryTableLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=self._table["tableReference"]["datasetId"],
                 project_id=self._table["tableReference"]["projectId"],
                 table_id=self._table["tableReference"]["tableId"],
@@ -2793,7 +2781,6 @@ class BigQueryUpdateTableSchemaOperator(GoogleCloudBaseOperator):
         if self._table:
             BigQueryTableLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=self._table["tableReference"]["datasetId"],
                 project_id=self._table["tableReference"]["projectId"],
                 table_id=self._table["tableReference"]["tableId"],
@@ -3039,7 +3026,6 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator, _BigQueryInsertJobOpera
                             table = job_configuration[job_type][table_prop]
                             persist_kwargs = {
                                 "context": context,
-                                "task_instance": self,
                                 "project_id": self.project_id,
                                 "table_id": table,
                             }
@@ -3061,7 +3047,6 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator, _BigQueryInsertJobOpera
 
         persist_kwargs = {
             "context": context,
-            "task_instance": self,
             "project_id": self.project_id,
             "location": self.location,
             "job_id": self.job_id,

airflow/providers/google/cloud/operators/bigquery_dts.py

@@ -134,7 +134,6 @@ class BigQueryCreateDataTransferOperator(GoogleCloudBaseOperator):
         transfer_config = _get_transfer_config_details(response.name)
         BigQueryDataTransferConfigLink.persist(
             context=context,
-            task_instance=self,
             region=transfer_config["region"],
             config_id=transfer_config["config_id"],
             project_id=transfer_config["project_id"],
@@ -142,7 +141,7 @@ class BigQueryCreateDataTransferOperator(GoogleCloudBaseOperator):
 
         result = TransferConfig.to_dict(response)
         self.log.info("Created DTS transfer config %s", get_object_id(result))
-        self.xcom_push(context, key="transfer_config_id", value=get_object_id(result))
+        context["ti"].xcom_push(key="transfer_config_id", value=get_object_id(result))
         # don't push AWS secret in XCOM
         result.get("params", {}).pop("secret_access_key", None)
         result.get("params", {}).pop("access_key_id", None)
@@ -329,7 +328,6 @@ class BigQueryDataTransferServiceStartTransferRunsOperator(GoogleCloudBaseOperat
         transfer_config = _get_transfer_config_details(response.runs[0].name)
         BigQueryDataTransferConfigLink.persist(
             context=context,
-            task_instance=self,
             region=transfer_config["region"],
             config_id=transfer_config["config_id"],
             project_id=transfer_config["project_id"],
@@ -337,7 +335,7 @@ class BigQueryDataTransferServiceStartTransferRunsOperator(GoogleCloudBaseOperat
 
         result = StartManualTransferRunsResponse.to_dict(response)
         run_id = get_object_id(result["runs"][0])
-        self.xcom_push(context, key="run_id", value=run_id)
+        context["ti"].xcom_push(key="run_id", value=run_id)
 
         if not self.deferrable:
             # Save as attribute for further use by OpenLineage

airflow/providers/google/cloud/operators/bigtable.py

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import google.api_core.exceptions
 
@@ -142,6 +142,13 @@ class BigtableCreateInstanceOperator(GoogleCloudBaseOperator, BigtableValidation
         self.impersonation_chain = impersonation_chain
         super().__init__(**kwargs)
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> None:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -155,7 +162,7 @@ class BigtableCreateInstanceOperator(GoogleCloudBaseOperator, BigtableValidation
                 "The instance '%s' already exists in this project. Consider it as created",
                 self.instance_id,
             )
-            BigtableInstanceLink.persist(context=context, task_instance=self)
+            BigtableInstanceLink.persist(context=context)
             return
         try:
             hook.create_instance(
@@ -171,7 +178,7 @@ class BigtableCreateInstanceOperator(GoogleCloudBaseOperator, BigtableValidation
                 cluster_storage_type=self.cluster_storage_type,
                 timeout=self.timeout,
             )
-            BigtableInstanceLink.persist(context=context, task_instance=self)
+            BigtableInstanceLink.persist(context=context)
         except google.api_core.exceptions.GoogleAPICallError as e:
             self.log.error("An error occurred. Exiting.")
             raise e
@@ -240,6 +247,13 @@ class BigtableUpdateInstanceOperator(GoogleCloudBaseOperator, BigtableValidation
         self.impersonation_chain = impersonation_chain
         super().__init__(**kwargs)
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> None:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -258,7 +272,7 @@ class BigtableUpdateInstanceOperator(GoogleCloudBaseOperator, BigtableValidation
                 instance_labels=self.instance_labels,
                 timeout=self.timeout,
             )
-            BigtableInstanceLink.persist(context=context, task_instance=self)
+            BigtableInstanceLink.persist(context=context)
         except google.api_core.exceptions.GoogleAPICallError as e:
             self.log.error("An error occurred. Exiting.")
             raise e
@@ -414,6 +428,13 @@ class BigtableCreateTableOperator(GoogleCloudBaseOperator, BigtableValidationMix
                 return False
         return True
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> None:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -431,7 +452,7 @@ class BigtableCreateTableOperator(GoogleCloudBaseOperator, BigtableValidationMix
                 initial_split_keys=self.initial_split_keys,
                 column_families=self.column_families,
             )
-            BigtableTablesLink.persist(context=context, task_instance=self)
+            BigtableTablesLink.persist(context=context)
         except google.api_core.exceptions.AlreadyExists:
             if not self._compare_column_families(hook, instance):
                 raise AirflowException(
@@ -575,6 +596,14 @@ class BigtableUpdateClusterOperator(GoogleCloudBaseOperator, BigtableValidationM
         self.impersonation_chain = impersonation_chain
         super().__init__(**kwargs)
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> None:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -586,7 +615,7 @@ class BigtableUpdateClusterOperator(GoogleCloudBaseOperator, BigtableValidationM
 
         try:
             hook.update_cluster(instance=instance, cluster_id=self.cluster_id, nodes=self.nodes)
-            BigtableClusterLink.persist(context=context, task_instance=self)
+            BigtableClusterLink.persist(context=context)
         except google.api_core.exceptions.NotFound:
             raise AirflowException(
                 f"Dependency: cluster '{self.cluster_id}' does not exist for instance '{self.instance_id}'."

airflow/providers/google/cloud/operators/cloud_base.py

@@ -19,9 +19,11 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 from google.api_core.gapic_v1.method import DEFAULT
 
-from airflow.models import BaseOperator
+from airflow.providers.google.version_compat import BaseOperator
 
 
 class GoogleCloudBaseOperator(BaseOperator):
@@ -36,3 +38,21 @@ class GoogleCloudBaseOperator(BaseOperator):
         """
         memo[id(DEFAULT)] = DEFAULT
         return super().__deepcopy__(memo)
+
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        """
+        Override this method to include parameters for link formatting in extra links.
+
+        For example; most of the links on the Google provider require `project_id` and `location` in the Link.
+        To be not repeat; you can override this function and return something like the following:
+
+        .. code-block:: python
+
+            {
+                "project_id": self.project_id,
+                "location": self.location,
+            }
+
+        """
+        return {}