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

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

@@ -27,7 +27,7 @@ Operators for Google Cloud Memorystore service.
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.memcache_v1beta2.types import cloud_memcache
@@ -133,6 +133,13 @@ class CloudMemorystoreCreateInstanceOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "location_id": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudMemorystoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -148,9 +155,6 @@ class CloudMemorystoreCreateInstanceOperator(GoogleCloudBaseOperator):
         )
         RedisInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            instance_id=self.instance_id,
-            location_id=self.location,
             project_id=self.project_id or hook.project_id,
         )
         return Instance.to_dict(result)
@@ -304,6 +308,13 @@ class CloudMemorystoreExportInstanceOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance,
+            "location_id": self.location,
+        }
+
     def execute(self, context: Context) -> None:
         hook = CloudMemorystoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -320,9 +331,6 @@ class CloudMemorystoreExportInstanceOperator(GoogleCloudBaseOperator):
         )
         RedisInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            instance_id=self.instance,
-            location_id=self.location,
             project_id=self.project_id or hook.project_id,
         )
 
@@ -397,6 +405,13 @@ class CloudMemorystoreFailoverInstanceOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance,
+            "location_id": self.location,
+        }
+
     def execute(self, context: Context) -> None:
         hook = CloudMemorystoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -412,9 +427,6 @@ class CloudMemorystoreFailoverInstanceOperator(GoogleCloudBaseOperator):
         )
         RedisInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            instance_id=self.instance,
-            location_id=self.location,
             project_id=self.project_id or hook.project_id,
         )
 
@@ -482,6 +494,13 @@ class CloudMemorystoreGetInstanceOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance,
+            "location_id": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudMemorystoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -496,9 +515,6 @@ class CloudMemorystoreGetInstanceOperator(GoogleCloudBaseOperator):
         )
         RedisInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            instance_id=self.instance,
-            location_id=self.location,
             project_id=self.project_id or hook.project_id,
         )
         return Instance.to_dict(result)
@@ -577,6 +593,13 @@ class CloudMemorystoreImportOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance,
+            "location_id": self.location,
+        }
+
     def execute(self, context: Context) -> None:
         hook = CloudMemorystoreHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -592,9 +615,6 @@ class CloudMemorystoreImportOperator(GoogleCloudBaseOperator):
         )
         RedisInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            instance_id=self.instance,
-            location_id=self.location,
             project_id=self.project_id or hook.project_id,
         )
 
@@ -680,7 +700,6 @@ class CloudMemorystoreListInstancesOperator(GoogleCloudBaseOperator):
         )
         RedisInstanceListLink.persist(
             context=context,
-            task_instance=self,
             project_id=self.project_id or hook.project_id,
         )
         instances = [Instance.to_dict(a) for a in result]
@@ -789,7 +808,6 @@ class CloudMemorystoreUpdateInstanceOperator(GoogleCloudBaseOperator):
         location_id, instance_id = res.name.split("/")[-3::2]
         RedisInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id or instance_id,
             location_id=self.location or location_id,
             project_id=self.project_id or hook.project_id,
@@ -882,7 +900,6 @@ class CloudMemorystoreScaleInstanceOperator(GoogleCloudBaseOperator):
         location_id, instance_id = res.name.split("/")[-3::2]
         RedisInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id or instance_id,
             location_id=self.location or location_id,
             project_id=self.project_id or hook.project_id,
@@ -1002,7 +1019,6 @@ class CloudMemorystoreCreateInstanceAndImportOperator(GoogleCloudBaseOperator):
         )
         RedisInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id,
             location_id=self.location,
             project_id=self.project_id or hook.project_id,
@@ -1171,6 +1187,14 @@ class CloudMemorystoreMemcachedApplyParametersOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "location_id": self.location,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = CloudMemorystoreMemcachedHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -1185,13 +1209,7 @@ class CloudMemorystoreMemcachedApplyParametersOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        MemcachedInstanceDetailsLink.persist(
-            context=context,
-            task_instance=self,
-            instance_id=self.instance_id,
-            location_id=self.location,
-            project_id=self.project_id,
-        )
+        MemcachedInstanceDetailsLink.persist(context=context)
 
 
 class CloudMemorystoreMemcachedCreateInstanceOperator(GoogleCloudBaseOperator):
@@ -1263,6 +1281,13 @@ class CloudMemorystoreMemcachedCreateInstanceOperator(GoogleCloudBaseOperator):
         self.metadata = metadata
         self.gcp_conn_id = gcp_conn_id
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "location_id": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudMemorystoreMemcachedHook(gcp_conn_id=self.gcp_conn_id)
         result = hook.create_instance(
@@ -1276,9 +1301,6 @@ class CloudMemorystoreMemcachedCreateInstanceOperator(GoogleCloudBaseOperator):
         )
         MemcachedInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            instance_id=self.instance_id,
-            location_id=self.location,
             project_id=self.project_id or hook.project_id,
         )
         return cloud_memcache.Instance.to_dict(result)
@@ -1410,6 +1432,13 @@ class CloudMemorystoreMemcachedGetInstanceOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance,
+            "location_id": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudMemorystoreMemcachedHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -1424,9 +1453,6 @@ class CloudMemorystoreMemcachedGetInstanceOperator(GoogleCloudBaseOperator):
         )
         MemcachedInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            instance_id=self.instance,
-            location_id=self.location,
             project_id=self.project_id or hook.project_id,
         )
         return cloud_memcache.Instance.to_dict(result)
@@ -1506,7 +1532,6 @@ class CloudMemorystoreMemcachedListInstancesOperator(GoogleCloudBaseOperator):
         )
         MemcachedInstanceListLink.persist(
             context=context,
-            task_instance=self,
             project_id=self.project_id or hook.project_id,
         )
         instances = [cloud_memcache.Instance.to_dict(a) for a in result]
@@ -1612,7 +1637,6 @@ class CloudMemorystoreMemcachedUpdateInstanceOperator(GoogleCloudBaseOperator):
         location_id, instance_id = res.name.split("/")[-3::2]
         MemcachedInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
             instance_id=self.instance_id or instance_id,
             location_id=self.location or location_id,
             project_id=self.project_id or hook.project_id,
@@ -1688,6 +1712,14 @@ class CloudMemorystoreMemcachedUpdateParametersOperator(GoogleCloudBaseOperator)
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "location_id": self.location,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = CloudMemorystoreMemcachedHook(
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
@@ -1702,10 +1734,4 @@ class CloudMemorystoreMemcachedUpdateParametersOperator(GoogleCloudBaseOperator)
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        MemcachedInstanceDetailsLink.persist(
-            context=context,
-            task_instance=self,
-            instance_id=self.instance_id,
-            location_id=self.location,
-            project_id=self.project_id,
-        )
+        MemcachedInstanceDetailsLink.persist(context=context)

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

@@ -27,6 +27,7 @@ from google.cloud.run_v2 import Job, Service
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_run import CloudRunHook, CloudRunServiceHook
+from airflow.providers.google.cloud.links.cloud_run import CloudRunJobLoggingLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.cloud_run import CloudRunJobFinishedTrigger, RunJobStatus
 
@@ -248,7 +249,7 @@ class CloudRunExecuteJobOperator(GoogleCloudBaseOperator):
 
     :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
     :param region: Required. The ID of the Google Cloud region that the service belongs to.
-    :param job_name: Required. The name of the job to update.
+    :param job_name: Required. The name of the job to execute.
     :param overrides: Optional map of override values.
     :param gcp_conn_id: The connection ID used to connect to Google Cloud.
     :param polling_period_seconds: Optional. Control the rate of the poll for the result of deferrable run.
@@ -265,6 +266,7 @@ class CloudRunExecuteJobOperator(GoogleCloudBaseOperator):
     :param deferrable: Run the operator in deferrable mode.
     """
 
+    operator_extra_links = (CloudRunJobLoggingLink(),)
     template_fields = (
         "project_id",
         "region",
@@ -312,6 +314,12 @@ class CloudRunExecuteJobOperator(GoogleCloudBaseOperator):
         if self.operation is None:
             raise AirflowException("Operation is None")
 
+        if self.operation.metadata.log_uri:
+            CloudRunJobLoggingLink.persist(
+                context=context,
+                log_uri=self.operation.metadata.log_uri,
+            )
+
         if not self.deferrable:
             result: Execution = self._wait_for_operation(self.operation)
             self._fail_if_execution_failed(result)

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

@@ -28,7 +28,6 @@ from googleapiclient.errors import HttpError
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
 from airflow.providers.google.cloud.hooks.cloud_sql import CloudSQLDatabaseHook, CloudSQLHook
 from airflow.providers.google.cloud.links.cloud_sql import CloudSQLInstanceDatabaseLink, CloudSQLInstanceLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
@@ -37,6 +36,11 @@ from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldVal
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, get_field
 from airflow.providers.google.common.links.storage import FileDetailsLink
 
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
 if TYPE_CHECKING:
     from airflow.models import Connection
     from airflow.providers.openlineage.extractors import OperatorLineage
@@ -285,6 +289,12 @@ class CloudSQLBaseOperator(GoogleCloudBaseOperator):
                 return False
             raise e
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance": self.instance,
+        }
+
     def execute(self, context: Context):
         pass
 
@@ -384,8 +394,6 @@ class CloudSQLCreateInstanceOperator(CloudSQLBaseOperator):
 
         CloudSQLInstanceLink.persist(
             context=context,
-            task_instance=self,
-            cloud_sql_instance=self.instance,
             project_id=self.project_id or hook.project_id,
         )
 
@@ -479,8 +487,6 @@ class CloudSQLInstancePatchOperator(CloudSQLBaseOperator):
             )
         CloudSQLInstanceLink.persist(
             context=context,
-            task_instance=self,
-            cloud_sql_instance=self.instance,
             project_id=self.project_id or hook.project_id,
         )
 
@@ -714,8 +720,6 @@ class CloudSQLCreateInstanceDatabaseOperator(CloudSQLBaseOperator):
         )
         CloudSQLInstanceDatabaseLink.persist(
             context=context,
-            task_instance=self,
-            cloud_sql_instance=self.instance,
             project_id=self.project_id or hook.project_id,
         )
         if self._check_if_db_exists(database, hook):
@@ -822,8 +826,6 @@ class CloudSQLPatchInstanceDatabaseOperator(CloudSQLBaseOperator):
             )
         CloudSQLInstanceDatabaseLink.persist(
             context=context,
-            task_instance=self,
-            cloud_sql_instance=self.instance,
             project_id=self.project_id or hook.project_id,
         )
         return hook.patch_database(
@@ -1004,13 +1006,10 @@ class CloudSQLExportInstanceOperator(CloudSQLBaseOperator):
         )
         CloudSQLInstanceLink.persist(
             context=context,
-            task_instance=self,
-            cloud_sql_instance=self.instance,
             project_id=self.project_id or hook.project_id,
         )
         FileDetailsLink.persist(
             context=context,
-            task_instance=self,
             uri=self.body["exportContext"]["uri"][5:],
             project_id=self.project_id or hook.project_id,
         )
@@ -1147,13 +1146,10 @@ class CloudSQLImportInstanceOperator(CloudSQLBaseOperator):
         )
         CloudSQLInstanceLink.persist(
             context=context,
-            task_instance=self,
-            cloud_sql_instance=self.instance,
             project_id=self.project_id or hook.project_id,
         )
         FileDetailsLink.persist(
             context=context,
-            task_instance=self,
             uri=self.body["importContext"]["uri"][5:],
             project_id=self.project_id or hook.project_id,
         )

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

@@ -586,7 +586,6 @@ class CloudDataTransferServiceGetOperationOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudStorageTransferDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 operation_name=self.operation_name,
             )
@@ -663,7 +662,6 @@ class CloudDataTransferServiceListOperationsOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudStorageTransferListLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
             )
 

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

@@ -74,6 +74,13 @@ class ComputeEngineBaseOperator(GoogleCloudBaseOperator):
         if not self.zone:
             raise AirflowException("The required parameter 'zone' is missing")
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.zone,
+            "resource_id": self.resource_id,
+        }
+
     def execute(self, context: Context):
         pass
 
@@ -225,9 +232,6 @@ class ComputeEngineInsertInstanceOperator(ComputeEngineBaseOperator):
             self.log.info("The %s Instance already exists", self.resource_id)
             ComputeInstanceDetailsLink.persist(
                 context=context,
-                task_instance=self,
-                location_id=self.zone,
-                resource_id=self.resource_id,
                 project_id=self.project_id or hook.project_id,
             )
             return Instance.to_dict(existing_instance)
@@ -247,9 +251,6 @@ class ComputeEngineInsertInstanceOperator(ComputeEngineBaseOperator):
         )
         ComputeInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         return Instance.to_dict(new_instance)
@@ -397,9 +398,6 @@ class ComputeEngineInsertInstanceFromTemplateOperator(ComputeEngineBaseOperator)
             self.log.info("The %s Instance already exists", self.resource_id)
             ComputeInstanceDetailsLink.persist(
                 context=context,
-                task_instance=self,
-                location_id=self.zone,
-                resource_id=self.resource_id,
                 project_id=self.project_id or hook.project_id,
             )
             return Instance.to_dict(existing_instance)
@@ -420,9 +418,6 @@ class ComputeEngineInsertInstanceFromTemplateOperator(ComputeEngineBaseOperator)
         )
         ComputeInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         return Instance.to_dict(new_instance_from_template)
@@ -598,9 +593,6 @@ class ComputeEngineStartInstanceOperator(ComputeEngineBaseOperator):
         )
         ComputeInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         hook.start_instance(zone=self.zone, resource_id=self.resource_id, project_id=self.project_id)
@@ -659,9 +651,6 @@ class ComputeEngineStopInstanceOperator(ComputeEngineBaseOperator):
         )
         ComputeInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         hook.stop_instance(zone=self.zone, resource_id=self.resource_id, project_id=self.project_id)
@@ -764,9 +753,6 @@ class ComputeEngineSetMachineTypeOperator(ComputeEngineBaseOperator):
         self._validate_all_body_fields()
         ComputeInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         hook.set_machine_type(
@@ -972,8 +958,6 @@ class ComputeEngineInsertInstanceTemplateOperator(ComputeEngineBaseOperator):
             self.log.info("The %s Template already exists.", existing_template)
             ComputeInstanceTemplateDetailsLink.persist(
                 context=context,
-                task_instance=self,
-                resource_id=self.resource_id,
                 project_id=self.project_id or hook.project_id,
             )
             return InstanceTemplate.to_dict(existing_template)
@@ -991,8 +975,6 @@ class ComputeEngineInsertInstanceTemplateOperator(ComputeEngineBaseOperator):
         )
         ComputeInstanceTemplateDetailsLink.persist(
             context=context,
-            task_instance=self,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         return InstanceTemplate.to_dict(new_template)
@@ -1238,7 +1220,6 @@ class ComputeEngineCopyInstanceTemplateOperator(ComputeEngineBaseOperator):
             )
             ComputeInstanceTemplateDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 resource_id=self.body_patch["name"],
                 project_id=self.project_id or hook.project_id,
             )
@@ -1259,7 +1240,6 @@ class ComputeEngineCopyInstanceTemplateOperator(ComputeEngineBaseOperator):
         )
         ComputeInstanceTemplateDetailsLink.persist(
             context=context,
-            task_instance=self,
             resource_id=self.body_patch["name"],
             project_id=self.project_id or hook.project_id,
         )
@@ -1390,9 +1370,6 @@ class ComputeEngineInstanceGroupUpdateManagerTemplateOperator(ComputeEngineBaseO
             self.log.info("Calling patch instance template with updated body: %s", patch_body)
             ComputeInstanceGroupManagerDetailsLink.persist(
                 context=context,
-                task_instance=self,
-                location_id=self.zone,
-                resource_id=self.resource_id,
                 project_id=self.project_id or hook.project_id,
             )
             return hook.patch_instance_group_manager(
@@ -1405,9 +1382,6 @@ class ComputeEngineInstanceGroupUpdateManagerTemplateOperator(ComputeEngineBaseO
         # Idempotence achieved
         ComputeInstanceGroupManagerDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         return True
@@ -1552,10 +1526,7 @@ class ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator)
             self.log.info("The %s Instance Group Manager already exists", existing_instance_group_manager)
             ComputeInstanceGroupManagerDetailsLink.persist(
                 context=context,
-                task_instance=self,
-                resource_id=self.resource_id,
                 project_id=self.project_id or hook.project_id,
-                location_id=self.zone,
             )
             return InstanceGroupManager.to_dict(existing_instance_group_manager)
         self._field_sanitizer.sanitize(self.body)
@@ -1574,9 +1545,6 @@ class ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator)
         )
         ComputeInstanceGroupManagerDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         return InstanceGroupManager.to_dict(new_instance_group_manager)