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

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

@@ -108,6 +108,12 @@ class CloudBuildCancelBuildOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.cancel_build(
@@ -119,14 +125,12 @@ class CloudBuildCancelBuildOperator(GoogleCloudBaseOperator):
             location=self.location,
         )
 
-        self.xcom_push(context, key="id", value=result.id)
+        context["task_instance"].xcom_push(key="id", value=result.id)
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 build_id=result.id,
             )
         return Build.to_dict(result)
@@ -210,6 +214,12 @@ class CloudBuildCreateBuildOperator(GoogleCloudBaseOperator):
             if self.build_raw.endswith(".json"):
                 self.build = json.loads(file.read())
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -225,7 +235,7 @@ class CloudBuildCreateBuildOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
             location=self.location,
         )
-        self.xcom_push(context, key="id", value=self.id_)
+        context["task_instance"].xcom_push(key="id", value=self.id_)
         if not self.wait:
             return Build.to_dict(
                 hook.get_build(id_=self.id_, project_id=self.project_id, location=self.location)
@@ -251,9 +261,7 @@ class CloudBuildCreateBuildOperator(GoogleCloudBaseOperator):
             if project_id:
                 CloudBuildLink.persist(
                     context=context,
-                    task_instance=self,
                     project_id=project_id,
-                    region=self.location,
                     build_id=cloud_build_instance_result.id,
                 )
             return Build.to_dict(cloud_build_instance_result)
@@ -269,9 +277,7 @@ class CloudBuildCreateBuildOperator(GoogleCloudBaseOperator):
             if project_id:
                 CloudBuildLink.persist(
                     context=context,
-                    task_instance=self,
                     project_id=project_id,
-                    region=self.location,
                     build_id=event["id_"],
                 )
             return event["instance"]
@@ -336,6 +342,12 @@ class CloudBuildCreateBuildTriggerOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.create_build_trigger(
@@ -346,21 +358,17 @@ class CloudBuildCreateBuildTriggerOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
             location=self.location,
         )
-        self.xcom_push(context, key="id", value=result.id)
+        context["task_instance"].xcom_push(key="id", value=result.id)
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildTriggerDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 trigger_id=result.id,
             )
             CloudBuildTriggersListLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
             )
         return BuildTrigger.to_dict(result)
 
@@ -419,6 +427,12 @@ class CloudBuildDeleteBuildTriggerOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         hook.delete_build_trigger(
@@ -433,9 +447,7 @@ class CloudBuildDeleteBuildTriggerOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudBuildTriggersListLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
             )
 
 
@@ -493,6 +505,12 @@ class CloudBuildGetBuildOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.get_build(
@@ -507,9 +525,7 @@ class CloudBuildGetBuildOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudBuildLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 build_id=result.id,
             )
         return Build.to_dict(result)
@@ -569,6 +585,12 @@ class CloudBuildGetBuildTriggerOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.get_build_trigger(
@@ -583,9 +605,7 @@ class CloudBuildGetBuildTriggerOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudBuildTriggerDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 trigger_id=result.id,
             )
         return BuildTrigger.to_dict(result)
@@ -649,6 +669,12 @@ class CloudBuildListBuildTriggersOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         results = hook.list_build_triggers(
@@ -664,9 +690,7 @@ class CloudBuildListBuildTriggersOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudBuildTriggersListLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
             )
         return [BuildTrigger.to_dict(result) for result in results]
 
@@ -729,6 +753,12 @@ class CloudBuildListBuildsOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         results = hook.list_builds(
@@ -743,7 +773,8 @@ class CloudBuildListBuildsOperator(GoogleCloudBaseOperator):
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildListLink.persist(
-                context=context, task_instance=self, project_id=project_id, region=self.location
+                context=context,
+                project_id=project_id,
             )
         return [Build.to_dict(result) for result in results]
 
@@ -805,6 +836,12 @@ class CloudBuildRetryBuildOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.retry_build(
@@ -817,14 +854,12 @@ class CloudBuildRetryBuildOperator(GoogleCloudBaseOperator):
             location=self.location,
         )
 
-        self.xcom_push(context, key="id", value=result.id)
+        context["task_instance"].xcom_push(key="id", value=result.id)
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 build_id=result.id,
             )
         return Build.to_dict(result)
@@ -891,6 +926,12 @@ class CloudBuildRunBuildTriggerOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.run_build_trigger(
@@ -903,14 +944,12 @@ class CloudBuildRunBuildTriggerOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
             location=self.location,
         )
-        self.xcom_push(context, key="id", value=result.id)
+        context["task_instance"].xcom_push(key="id", value=result.id)
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 build_id=result.id,
             )
         return Build.to_dict(result)
@@ -974,6 +1013,12 @@ class CloudBuildUpdateBuildTriggerOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.location = location
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.location,
+        }
+
     def execute(self, context: Context):
         hook = CloudBuildHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         result = hook.update_build_trigger(
@@ -985,14 +1030,12 @@ class CloudBuildUpdateBuildTriggerOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
             location=self.location,
         )
-        self.xcom_push(context, key="id", value=result.id)
+        context["task_instance"].xcom_push(key="id", value=result.id)
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudBuildTriggerDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
-                region=self.location,
                 trigger_id=result.id,
             )
         return BuildTrigger.to_dict(result)

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

@@ -19,7 +19,7 @@ from __future__ import annotations
 
 import shlex
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import AlreadyExists
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -57,25 +57,6 @@ class CloudComposerEnvironmentLink(BaseGoogleLink):
     key = "composer_conf"
     format_str = CLOUD_COMPOSER_DETAILS_LINK
 
-    @staticmethod
-    def persist(
-        operator_instance: (
-            CloudComposerCreateEnvironmentOperator
-            | CloudComposerUpdateEnvironmentOperator
-            | CloudComposerGetEnvironmentOperator
-        ),
-        context: Context,
-    ) -> None:
-        operator_instance.xcom_push(
-            context,
-            key=CloudComposerEnvironmentLink.key,
-            value={
-                "project_id": operator_instance.project_id,
-                "region": operator_instance.region,
-                "environment_id": operator_instance.environment_id,
-            },
-        )
-
 
 class CloudComposerEnvironmentsLink(BaseGoogleLink):
     """Helper class for constructing Cloud Composer Environment Link."""
@@ -84,16 +65,6 @@ class CloudComposerEnvironmentsLink(BaseGoogleLink):
     key = "composer_conf"
     format_str = CLOUD_COMPOSER_ENVIRONMENTS_LINK
 
-    @staticmethod
-    def persist(operator_instance: CloudComposerListEnvironmentsOperator, context: Context) -> None:
-        operator_instance.xcom_push(
-            context,
-            key=CloudComposerEnvironmentsLink.key,
-            value={
-                "project_id": operator_instance.project_id,
-            },
-        )
-
 
 class CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
     """
@@ -159,6 +130,14 @@ class CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
         self.deferrable = deferrable
         self.pooling_period_seconds = pooling_period_seconds
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+            "region": self.region,
+            "environment_id": self.environment_id,
+        }
+
     def execute(self, context: Context):
         hook = CloudComposerHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -171,7 +150,7 @@ class CloudComposerCreateEnvironmentOperator(GoogleCloudBaseOperator):
         else:
             self.environment["name"] = name
 
-        CloudComposerEnvironmentLink.persist(operator_instance=self, context=context)
+        CloudComposerEnvironmentLink.persist(context=context)
         try:
             result = hook.create_environment(
                 project_id=self.project_id,
@@ -370,6 +349,14 @@ class CloudComposerGetEnvironmentOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+            "region": self.region,
+            "environment_id": self.environment_id,
+        }
+
     def execute(self, context: Context):
         hook = CloudComposerHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -384,8 +371,7 @@ class CloudComposerGetEnvironmentOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-
-        CloudComposerEnvironmentLink.persist(operator_instance=self, context=context)
+        CloudComposerEnvironmentLink.persist(context=context)
         return Environment.to_dict(result)
 
 
@@ -445,12 +431,17 @@ class CloudComposerListEnvironmentsOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = CloudComposerHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
-        CloudComposerEnvironmentsLink.persist(operator_instance=self, context=context)
         result = hook.list_environments(
             project_id=self.project_id,
             region=self.region,
@@ -532,6 +523,14 @@ class CloudComposerUpdateEnvironmentOperator(GoogleCloudBaseOperator):
         self.deferrable = deferrable
         self.pooling_period_seconds = pooling_period_seconds
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+            "region": self.region,
+            "environment_id": self.environment_id,
+        }
+
     def execute(self, context: Context):
         hook = CloudComposerHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -549,7 +548,7 @@ class CloudComposerUpdateEnvironmentOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
 
-        CloudComposerEnvironmentLink.persist(operator_instance=self, context=context)
+        CloudComposerEnvironmentLink.persist(context=context)
         if not self.deferrable:
             environment = hook.wait_for_operation(timeout=self.timeout, operation=result)
             return Environment.to_dict(environment)