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

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

@@ -19,19 +19,23 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, ClassVar
 
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
-
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.models.taskinstancekey import TaskInstanceKey
+from airflow.providers.google.version_compat import (
+    AIRFLOW_V_3_0_PLUS,
+    BaseOperator,
+    BaseOperatorLink,
+    BaseSensorOperator,
+)
 
 if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseOperatorLink
     from airflow.sdk.execution_time.xcom import XCom
 else:
-    from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
     from airflow.models.xcom import XCom  # type: ignore[no-redef]
 
+if TYPE_CHECKING:
+    from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+    from airflow.utils.context import Context
+
 BASE_LINK = "https://console.cloud.google.com"
 
 
@@ -39,6 +43,12 @@ class BaseGoogleLink(BaseOperatorLink):
     """
     Base class for all Google links.
 
+    When you inherit this class in a Link class;
+      - You can call the persist method to push data to the XCom to use it later in the get_link method.
+      - If you have an operator which inherit the GoogleCloudBaseOperator or BaseSensorOperator
+        You can define extra_links_params method in the operator to pass the operator properties
+        to the get_link method.
+
     :meta private:
     """
 
@@ -46,15 +56,69 @@ class BaseGoogleLink(BaseOperatorLink):
     key: ClassVar[str]
     format_str: ClassVar[str]
 
+    @property
+    def xcom_key(self) -> str:
+        # NOTE: in Airflow 3 we need to have xcom_key property in the Link class.
+        # Since we have the key property already, this is just a proxy property method to use same
+        # key as in Airflow 2.
+        return self.key
+
+    @classmethod
+    def persist(cls, context: Context, **value):
+        """
+        Push arguments to the XCom to use later for link formatting at the `get_link` method.
+
+        Note: for Airflow 2 we need to call this function with context variable only
+        where we have the extra_links_params property method defined
+        """
+        params = {}
+        # TODO: remove after Airflow v2 support dropped
+        if not AIRFLOW_V_3_0_PLUS:
+            common_params = getattr(context["task"], "extra_links_params", None)
+            if common_params:
+                params.update(common_params)
+
+        context["ti"].xcom_push(
+            key=cls.key,
+            value={
+                **params,
+                **value,
+            },
+        )
+
+    def get_config(self, operator, ti_key):
+        conf = {}
+        conf.update(getattr(operator, "extra_links_params", {}))
+        conf.update(XCom.get_value(key=self.key, ti_key=ti_key) or {})
+
+        # if the config did not define, return None to stop URL formatting
+        if not conf:
+            return None
+
+        # Add a default value for the 'namespace' parameter for backward compatibility.
+        # This is for datafusion
+        conf.setdefault("namespace", "default")
+        return conf
+
     def get_link(
         self,
         operator: BaseOperator,
         *,
         ti_key: TaskInstanceKey,
     ) -> str:
-        conf = XCom.get_value(key=self.key, ti_key=ti_key)
+        if TYPE_CHECKING:
+            assert isinstance(operator, (GoogleCloudBaseOperator, BaseSensorOperator))
+
+        conf = self.get_config(operator, ti_key)
         if not conf:
             return ""
-        if self.format_str.startswith("http"):
-            return self.format_str.format(**conf)
-        return BASE_LINK + self.format_str.format(**conf)
+        return self._format_link(**conf)
+
+    def _format_link(self, **kwargs):
+        try:
+            formatted_str = self.format_str.format(**kwargs)
+            if formatted_str.startswith("http"):
+                return formatted_str
+            return BASE_LINK + formatted_str
+        except KeyError:
+            return ""

airflow/providers/google/cloud/links/bigquery.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
-
 BIGQUERY_BASE_LINK = "/bigquery"
 BIGQUERY_DATASET_LINK = (
     BIGQUERY_BASE_LINK + "?referrer=search&project={project_id}&d={dataset_id}&p={project_id}&page=dataset"
@@ -47,19 +41,6 @@ class BigQueryDatasetLink(BaseGoogleLink):
     key = "bigquery_dataset"
     format_str = BIGQUERY_DATASET_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        dataset_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryDatasetLink.key,
-            value={"dataset_id": dataset_id, "project_id": project_id},
-        )
-
 
 class BigQueryTableLink(BaseGoogleLink):
     """Helper class for constructing BigQuery Table Link."""
@@ -68,20 +49,6 @@ class BigQueryTableLink(BaseGoogleLink):
     key = "bigquery_table"
     format_str = BIGQUERY_TABLE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str,
-        table_id: str,
-        dataset_id: str | None = None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryTableLink.key,
-            value={"dataset_id": dataset_id, "project_id": project_id, "table_id": table_id},
-        )
-
 
 class BigQueryJobDetailLink(BaseGoogleLink):
     """Helper class for constructing BigQuery Job Detail Link."""
@@ -89,17 +56,3 @@ class BigQueryJobDetailLink(BaseGoogleLink):
     name = "BigQuery Job Detail"
     key = "bigquery_job_detail"
     format_str = BIGQUERY_JOB_DETAIL_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str,
-        location: str,
-        job_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryJobDetailLink.key,
-            value={"project_id": project_id, "location": location, "job_id": job_id},
-        )

airflow/providers/google/cloud/links/bigquery_dts.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
-
 BIGQUERY_BASE_LINK = "/bigquery/transfers"
 BIGQUERY_DTS_LINK = BIGQUERY_BASE_LINK + "/locations/{region}/configs/{config_id}/runs?project={project_id}"
 
@@ -37,17 +31,3 @@ class BigQueryDataTransferConfigLink(BaseGoogleLink):
     name = "BigQuery Data Transfer Config"
     key = "bigquery_dts_config"
     format_str = BIGQUERY_DTS_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        region: str,
-        config_id: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=BigQueryDataTransferConfigLink.key,
-            value={"project_id": project_id, "region": region, "config_id": config_id},
-        )

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

@@ -16,13 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.utils.context import Context
-
 BIGTABLE_BASE_LINK = "/bigtable"
 BIGTABLE_INSTANCE_LINK = BIGTABLE_BASE_LINK + "/instances/{instance_id}/overview?project={project_id}"
 BIGTABLE_CLUSTER_LINK = (
@@ -38,20 +33,6 @@ class BigtableInstanceLink(BaseGoogleLink):
     key = "instance_key"
     format_str = BIGTABLE_INSTANCE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=BigtableInstanceLink.key,
-            value={
-                "instance_id": task_instance.instance_id,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class BigtableClusterLink(BaseGoogleLink):
     """Helper class for constructing Bigtable Cluster link."""
@@ -60,21 +41,6 @@ class BigtableClusterLink(BaseGoogleLink):
     key = "cluster_key"
     format_str = BIGTABLE_CLUSTER_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=BigtableClusterLink.key,
-            value={
-                "instance_id": task_instance.instance_id,
-                "cluster_id": task_instance.cluster_id,
-                "project_id": task_instance.project_id,
-            },
-        )
-
 
 class BigtableTablesLink(BaseGoogleLink):
     """Helper class for constructing Bigtable Tables link."""
@@ -82,17 +48,3 @@ class BigtableTablesLink(BaseGoogleLink):
     name = "Bigtable Tables"
     key = "tables_key"
     format_str = BIGTABLE_TABLES_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=BigtableTablesLink.key,
-            value={
-                "instance_id": task_instance.instance_id,
-                "project_id": task_instance.project_id,
-            },
-        )

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

@@ -16,13 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
 
-if TYPE_CHECKING:
-    from airflow.utils.context import Context
-
 BUILD_BASE_LINK = "/cloud-build"
 
 BUILD_LINK = BUILD_BASE_LINK + "/builds;region={region}/{build_id}?project={project_id}"
@@ -43,24 +38,6 @@ class CloudBuildLink(BaseGoogleLink):
     key = "cloud_build_key"
     format_str = BUILD_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        build_id: str,
-        project_id: str,
-        region: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=CloudBuildLink.key,
-            value={
-                "project_id": project_id,
-                "region": region,
-                "build_id": build_id,
-            },
-        )
-
 
 class CloudBuildListLink(BaseGoogleLink):
     """Helper class for constructing Cloud Build List link."""
@@ -69,22 +46,6 @@ class CloudBuildListLink(BaseGoogleLink):
     key = "cloud_build_list_key"
     format_str = BUILD_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-        region: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=CloudBuildListLink.key,
-            value={
-                "project_id": project_id,
-                "region": region,
-            },
-        )
-
 
 class CloudBuildTriggersListLink(BaseGoogleLink):
     """Helper class for constructing Cloud Build Triggers List link."""
@@ -93,22 +54,6 @@ class CloudBuildTriggersListLink(BaseGoogleLink):
     key = "cloud_build_triggers_list_key"
     format_str = BUILD_TRIGGERS_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-        region: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=CloudBuildTriggersListLink.key,
-            value={
-                "project_id": project_id,
-                "region": region,
-            },
-        )
-
 
 class CloudBuildTriggerDetailsLink(BaseGoogleLink):
     """Helper class for constructing Cloud Build Trigger Details link."""
@@ -116,21 +61,3 @@ class CloudBuildTriggerDetailsLink(BaseGoogleLink):
     name = "Cloud Build Triggers Details"
     key = "cloud_build_triggers_details_key"
     format_str = BUILD_TRIGGER_DETAILS_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance,
-        project_id: str,
-        region: str,
-        trigger_id: str,
-    ):
-        task_instance.xcom_push(
-            context=context,
-            key=CloudBuildTriggerDetailsLink.key,
-            value={
-                "project_id": project_id,
-                "region": region,
-                "trigger_id": trigger_id,
-            },
-        )

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

@@ -19,15 +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
-
-
 CLOUD_FUNCTIONS_BASE_LINK = "https://console.cloud.google.com/functions"
 
 CLOUD_FUNCTIONS_DETAILS_LINK = (
@@ -44,20 +37,6 @@ class CloudFunctionsDetailsLink(BaseGoogleLink):
     key = "cloud_functions_details"
     format_str = CLOUD_FUNCTIONS_DETAILS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        function_name: str,
-        location: str,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=CloudFunctionsDetailsLink.key,
-            value={"function_name": function_name, "location": location, "project_id": project_id},
-        )
-
 
 class CloudFunctionsListLink(BaseGoogleLink):
     """Helper class for constructing Cloud Functions Details Link."""
@@ -65,15 +44,3 @@ class CloudFunctionsListLink(BaseGoogleLink):
     name = "Cloud Functions List"
     key = "cloud_functions_list"
     format_str = CLOUD_FUNCTIONS_LIST_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=CloudFunctionsDetailsLink.key,
-            value={"project_id": project_id},
-        )

airflow/providers/google/cloud/links/cloud_memorystore.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
-
 BASE_LINK = "/memorystore"
 MEMCACHED_LINK = (
     BASE_LINK + "/memcached/locations/{location_id}/instances/{instance_id}/details?project={project_id}"
@@ -45,20 +39,6 @@ class MemcachedInstanceDetailsLink(BaseGoogleLink):
     key = "memcached_instance"
     format_str = MEMCACHED_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        instance_id: str,
-        location_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=MemcachedInstanceDetailsLink.key,
-            value={"instance_id": instance_id, "location_id": location_id, "project_id": project_id},
-        )
-
 
 class MemcachedInstanceListLink(BaseGoogleLink):
     """Helper class for constructing Memorystore Memcached List of Instances Link."""
@@ -67,18 +47,6 @@ class MemcachedInstanceListLink(BaseGoogleLink):
     key = "memcached_instances"
     format_str = MEMCACHED_LIST_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=MemcachedInstanceListLink.key,
-            value={"project_id": project_id},
-        )
-
 
 class RedisInstanceDetailsLink(BaseGoogleLink):
     """Helper class for constructing Memorystore Redis Instance Link."""
@@ -87,20 +55,6 @@ class RedisInstanceDetailsLink(BaseGoogleLink):
     key = "redis_instance"
     format_str = REDIS_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        instance_id: str,
-        location_id: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=RedisInstanceDetailsLink.key,
-            value={"instance_id": instance_id, "location_id": location_id, "project_id": project_id},
-        )
-
 
 class RedisInstanceListLink(BaseGoogleLink):
     """Helper class for constructing Memorystore Redis List of Instances Link."""
@@ -108,15 +62,3 @@ class RedisInstanceListLink(BaseGoogleLink):
     name = "Memorystore Redis List of Instances"
     key = "redis_instances"
     format_str = REDIS_LIST_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=RedisInstanceListLink.key,
-            value={"project_id": project_id},
-        )

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

@@ -16,20 +16,7 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 from airflow.providers.google.cloud.links.base import BaseGoogleLink
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
-
-if TYPE_CHECKING:
-    from airflow.models import BaseOperator
-    from airflow.models.taskinstancekey import TaskInstanceKey
-    from airflow.utils.context import Context
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk.execution_time.xcom import XCom
-else:
-    from airflow.models.xcom import XCom  # type: ignore[no-redef]
 
 
 class CloudRunJobLoggingLink(BaseGoogleLink):
@@ -37,23 +24,4 @@ class CloudRunJobLoggingLink(BaseGoogleLink):
 
     name = "Cloud Run Job Logging"
     key = "log_uri"
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        log_uri: str,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=CloudRunJobLoggingLink.key,
-            value=log_uri,
-        )
-
-    def get_link(
-        self,
-        operator: BaseOperator,
-        *,
-        ti_key: TaskInstanceKey,
-    ) -> str:
-        return XCom.get_value(key=self.key, ti_key=ti_key)
+    format_str = "{log_uri}"

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

@@ -19,15 +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
-
-
 CLOUD_SQL_BASE_LINK = "/sql"
 CLOUD_SQL_INSTANCE_LINK = CLOUD_SQL_BASE_LINK + "/instances/{instance}/overview?project={project_id}"
 CLOUD_SQL_INSTANCE_DATABASE_LINK = (
@@ -42,19 +35,6 @@ class CloudSQLInstanceLink(BaseGoogleLink):
     key = "cloud_sql_instance"
     format_str = CLOUD_SQL_INSTANCE_LINK
 
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        cloud_sql_instance: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=CloudSQLInstanceLink.key,
-            value={"instance": cloud_sql_instance, "project_id": project_id},
-        )
-
 
 class CloudSQLInstanceDatabaseLink(BaseGoogleLink):
     """Helper class for constructing Cloud SQL Instance Database Link."""
@@ -62,16 +42,3 @@ class CloudSQLInstanceDatabaseLink(BaseGoogleLink):
     name = "Cloud SQL Instance Database"
     key = "cloud_sql_instance_database"
     format_str = CLOUD_SQL_INSTANCE_DATABASE_LINK
-
-    @staticmethod
-    def persist(
-        context: Context,
-        task_instance: BaseOperator,
-        cloud_sql_instance: str,
-        project_id: str | None,
-    ):
-        task_instance.xcom_push(
-            context,
-            key=CloudSQLInstanceDatabaseLink.key,
-            value={"instance": cloud_sql_instance, "project_id": project_id},
-        )