apache-airflow-providers-google 17.2.0rc1__py3-none-any.whl → 18.0.0rc1__py3-none-any.whl

airflow/providers/google/cloud/transfers/bigquery_to_mysql.py

@@ -21,6 +21,7 @@ from __future__ import annotations
 
 import warnings
 from collections.abc import Sequence
+from functools import cached_property
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToSqlBaseOperator
@@ -76,5 +77,15 @@ class BigQueryToMySqlOperator(BigQueryToSqlBaseOperator):
         )
         self.mysql_conn_id = mysql_conn_id
 
-    def get_sql_hook(self) -> MySqlHook:
+    @cached_property
+    def mysql_hook(self) -> MySqlHook:
         return MySqlHook(schema=self.database, mysql_conn_id=self.mysql_conn_id)
+
+    def get_sql_hook(self) -> MySqlHook:
+        return self.mysql_hook
+
+    def execute(self, context):
+        # Set source_project_dataset_table here, after hooks are initialized and project_id is available
+        project_id = self.bigquery_hook.project_id
+        self.source_project_dataset_table = f"{project_id}.{self.dataset_id}.{self.table_id}"
+        return super().execute(context)

airflow/providers/google/cloud/transfers/bigquery_to_postgres.py

@@ -19,6 +19,7 @@
 
 from __future__ import annotations
 
+from functools import cached_property
 from typing import TYPE_CHECKING
 
 from psycopg2.extensions import register_adapter
@@ -78,28 +79,36 @@ class BigQueryToPostgresOperator(BigQueryToSqlBaseOperator):
         self.postgres_conn_id = postgres_conn_id
         self.replace_index = replace_index
 
-    def get_sql_hook(self) -> PostgresHook:
+    @cached_property
+    def postgres_hook(self) -> PostgresHook:
         register_adapter(list, Json)
         register_adapter(dict, Json)
         return PostgresHook(database=self.database, postgres_conn_id=self.postgres_conn_id)
 
+    def get_sql_hook(self) -> PostgresHook:
+        return self.postgres_hook
+
     def execute(self, context: Context) -> None:
-        big_query_hook = BigQueryHook(
-            gcp_conn_id=self.gcp_conn_id,
-            location=self.location,
-            impersonation_chain=self.impersonation_chain,
-        )
+        if not self.bigquery_hook:
+            self.bigquery_hook = BigQueryHook(
+                gcp_conn_id=self.gcp_conn_id,
+                location=self.location,
+                impersonation_chain=self.impersonation_chain,
+            )
+        # Set source_project_dataset_table here, after hooks are initialized and project_id is available
+        project_id = self.bigquery_hook.project_id
+        self.source_project_dataset_table = f"{project_id}.{self.dataset_id}.{self.table_id}"
+
         self.persist_links(context)
-        sql_hook: PostgresHook = self.get_sql_hook()
         for rows in bigquery_get_data(
             self.log,
             self.dataset_id,
             self.table_id,
-            big_query_hook,
+            self.bigquery_hook,
             self.batch_size,
             self.selected_fields,
         ):
-            sql_hook.insert_rows(
+            self.postgres_hook.insert_rows(
                 table=self.target_table_name,
                 rows=rows,
                 target_fields=self.selected_fields,

airflow/providers/google/cloud/transfers/bigquery_to_sql.py

@@ -30,6 +30,7 @@ from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.providers.common.sql.hooks.sql import DbApiHook
+    from airflow.providers.openlineage.extractors import OperatorLineage
     from airflow.utils.context import Context
 
 
@@ -140,3 +141,97 @@ class BigQueryToSqlBaseOperator(BaseOperator):
                 replace=self.replace,
                 commit_every=self.batch_size,
             )
+
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage | None:
+        """
+        Build a generic OpenLineage facet for BigQuery -> SQL transfers.
+
+        This consolidates nearly identical implementations from child
+        operators. Children still provide a concrete SQL hook via
+        ``get_sql_hook()`` and may override behavior if needed.
+        """
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.utils import (
+            BIGQUERY_NAMESPACE,
+            get_facets_from_bq_table_for_given_fields,
+            get_identity_column_lineage_facet,
+        )
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        if not self.bigquery_hook:
+            self.bigquery_hook = BigQueryHook(
+                gcp_conn_id=self.gcp_conn_id,
+                location=self.location,
+                impersonation_chain=self.impersonation_chain,
+            )
+
+        try:
+            if not getattr(self, "source_project_dataset_table", None):
+                project_id = self.bigquery_hook.project_id
+                self.source_project_dataset_table = f"{project_id}.{self.dataset_id}.{self.table_id}"
+
+            table_obj = self.bigquery_hook.get_client().get_table(self.source_project_dataset_table)
+        except Exception:
+            self.log.debug(
+                "OpenLineage: could not fetch BigQuery table %s",
+                getattr(self, "source_project_dataset_table", None),
+                exc_info=True,
+            )
+            return OperatorLineage()
+
+        if self.selected_fields:
+            if isinstance(self.selected_fields, str):
+                bigquery_field_names = list(self.selected_fields)
+            else:
+                bigquery_field_names = self.selected_fields
+        else:
+            bigquery_field_names = [f.name for f in getattr(table_obj, "schema", [])]
+
+        input_dataset = Dataset(
+            namespace=BIGQUERY_NAMESPACE,
+            name=self.source_project_dataset_table,
+            facets=get_facets_from_bq_table_for_given_fields(table_obj, bigquery_field_names),
+        )
+
+        sql_hook = self.get_sql_hook()
+        db_info = sql_hook.get_openlineage_database_info(sql_hook.get_conn())
+        if db_info is None:
+            self.log.debug("OpenLineage: could not get database info from SQL hook %s", type(sql_hook))
+            return OperatorLineage()
+        namespace = f"{db_info.scheme}://{db_info.authority}"
+
+        schema_name = None
+        if hasattr(sql_hook, "get_openlineage_default_schema"):
+            try:
+                schema_name = sql_hook.get_openlineage_default_schema()
+            except Exception:
+                schema_name = None
+
+        if self.target_table_name and "." in self.target_table_name:
+            schema_part, table_part = self.target_table_name.split(".", 1)
+        else:
+            schema_part = schema_name or ""
+            table_part = self.target_table_name or ""
+
+        if db_info and db_info.scheme == "mysql":
+            output_name = f"{self.database}.{table_part}" if self.database else f"{table_part}"
+        else:
+            if self.database:
+                if schema_part:
+                    output_name = f"{self.database}.{schema_part}.{table_part}"
+                else:
+                    output_name = f"{self.database}.{table_part}"
+            else:
+                if schema_part:
+                    output_name = f"{schema_part}.{table_part}"
+                else:
+                    output_name = f"{table_part}"
+
+        column_lineage_facet = get_identity_column_lineage_facet(
+            bigquery_field_names, input_datasets=[input_dataset]
+        )
+
+        output_facets = column_lineage_facet or {}
+        output_dataset = Dataset(namespace=namespace, name=output_name, facets=output_facets)
+
+        return OperatorLineage(inputs=[input_dataset], outputs=[output_dataset])

airflow/providers/google/cloud/transfers/gcs_to_bigquery.py

@@ -144,6 +144,9 @@ class GCSToBigQueryOperator(BaseOperator):
         partition by field, type and  expiration as per API specifications.
         Note that 'field' is not available in concurrency with
         dataset.table$partition.
+        Ignored if 'range_partitioning' is set.
+    :param range_partitioning: configure optional range partitioning fields i.e.
+        partition by field and integer interval as per API specifications.
     :param cluster_fields: Request that the result of this load be stored sorted
         by one or more columns. BigQuery supports clustering for both partitioned and
         non-partitioned tables. The order of columns given determines the sort order.
@@ -219,6 +222,7 @@ class GCSToBigQueryOperator(BaseOperator):
         src_fmt_configs=None,
         external_table=False,
         time_partitioning=None,
+        range_partitioning=None,
         cluster_fields=None,
         autodetect=True,
         encryption_configuration=None,
@@ -246,6 +250,10 @@ class GCSToBigQueryOperator(BaseOperator):
             src_fmt_configs = {}
         if time_partitioning is None:
             time_partitioning = {}
+        if range_partitioning is None:
+            range_partitioning = {}
+        if range_partitioning and time_partitioning:
+            raise ValueError("Only one of time_partitioning or range_partitioning can be set.")
         self.bucket = bucket
         self.source_objects = source_objects
         self.schema_object = schema_object
@@ -283,6 +291,7 @@ class GCSToBigQueryOperator(BaseOperator):
         self.schema_update_options = schema_update_options
         self.src_fmt_configs = src_fmt_configs
         self.time_partitioning = time_partitioning
+        self.range_partitioning = range_partitioning
         self.cluster_fields = cluster_fields
         self.autodetect = autodetect
         self.encryption_configuration = encryption_configuration
@@ -627,6 +636,8 @@ class GCSToBigQueryOperator(BaseOperator):
         )
         if self.time_partitioning:
             self.configuration["load"].update({"timePartitioning": self.time_partitioning})
+        if self.range_partitioning:
+            self.configuration["load"].update({"rangePartitioning": self.range_partitioning})
 
         if self.cluster_fields:
             self.configuration["load"].update({"clustering": {"fields": self.cluster_fields}})

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

@@ -52,11 +52,6 @@ class CloudComposerExecutionTrigger(BaseTrigger):
         self.impersonation_chain = impersonation_chain
         self.pooling_period_seconds = pooling_period_seconds
 
-        self.gcp_hook = CloudComposerAsyncHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
-
     def serialize(self) -> tuple[str, dict[str, Any]]:
         return (
             "airflow.providers.google.cloud.triggers.cloud_composer.CloudComposerExecutionTrigger",
@@ -70,7 +65,14 @@ class CloudComposerExecutionTrigger(BaseTrigger):
             },
         )
 
+    def _get_async_hook(self) -> CloudComposerAsyncHook:
+        return CloudComposerAsyncHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
     async def run(self):
+        self.gcp_hook = self._get_async_hook()
         while True:
             operation = await self.gcp_hook.get_operation(operation_name=self.operation_name)
             if operation.done:
@@ -108,11 +110,6 @@ class CloudComposerAirflowCLICommandTrigger(BaseTrigger):
         self.impersonation_chain = impersonation_chain
         self.poll_interval = poll_interval
 
-        self.gcp_hook = CloudComposerAsyncHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
-
     def serialize(self) -> tuple[str, dict[str, Any]]:
         return (
             "airflow.providers.google.cloud.triggers.cloud_composer.CloudComposerAirflowCLICommandTrigger",
@@ -127,7 +124,14 @@ class CloudComposerAirflowCLICommandTrigger(BaseTrigger):
             },
         )
 
+    def _get_async_hook(self) -> CloudComposerAsyncHook:
+        return CloudComposerAsyncHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
     async def run(self):
+        self.gcp_hook = self._get_async_hook()
         try:
             result = await self.gcp_hook.wait_command_execution_result(
                 project_id=self.project_id,
@@ -199,11 +203,6 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
         self.poll_interval = poll_interval
         self.composer_airflow_version = composer_airflow_version
 
-        self.gcp_hook = CloudComposerAsyncHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
-        )
-
     def serialize(self) -> tuple[str, dict[str, Any]]:
         return (
             "airflow.providers.google.cloud.triggers.cloud_composer.CloudComposerDAGRunTrigger",
@@ -264,6 +263,12 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
                 return False
         return True
 
+    def _get_async_hook(self) -> CloudComposerAsyncHook:
+        return CloudComposerAsyncHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
     def _check_composer_dag_run_id_states(self, dag_runs: list[dict]) -> bool:
         for dag_run in dag_runs:
             if dag_run["run_id"] == self.composer_dag_run_id and dag_run["state"] in self.allowed_states:
@@ -271,6 +276,7 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
         return False
 
     async def run(self):
+        self.gcp_hook: CloudComposerAsyncHook = self._get_async_hook()
         try:
             while True:
                 if datetime.now(self.end_date.tzinfo).timestamp() > self.end_date.timestamp():

airflow/providers/google/cloud/utils/bigquery_get_data.py

@@ -23,9 +23,9 @@ from google.cloud.bigquery.table import Row, RowIterator
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
-    from logging import Logger
 
     from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+    from airflow.sdk.types import Logger
 
 
 def bigquery_get_data(

airflow/providers/google/marketing_platform/hooks/display_video.py

@@ -24,8 +24,6 @@ from typing import Any
 
 from googleapiclient.discovery import Resource, build
 
-from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 
@@ -48,23 +46,6 @@ class GoogleDisplayVideo360Hook(GoogleBaseHook):
         )
         self.api_version = api_version
 
-    @deprecated(
-        planned_removal_date="September 01, 2025",
-        use_instead="airflow.providers.google.marketing_platform.hooks.display_video.get_conn_to_display_video",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_conn(self) -> Resource:
-        """Retrieve connection to DisplayVideo."""
-        if not self._conn:
-            http_authorized = self._authorize()
-            self._conn = build(
-                "doubleclickbidmanager",
-                self.api_version,
-                http=http_authorized,
-                cache_discovery=False,
-            )
-        return self._conn
-
     def get_conn_to_display_video(self) -> Resource:
         """Retrieve connection to DisplayVideo."""
         if not self._conn:
@@ -96,137 +77,6 @@ class GoogleDisplayVideo360Hook(GoogleBaseHook):
         """
         return [f"gdbm-{partner_id}/entity/{{{{ ds_nodash }}}}.*.{entity_type}.json"]
 
-    @deprecated(
-        planned_removal_date="September 01, 2025",
-        use_instead="airflow.providers.google.marketing_platform.hooks.display_video.create_sdf_download_operation",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def create_query(self, query: dict[str, Any]) -> dict:
-        """
-        Create a query.
-
-        :param query: Query object to be passed to request body.
-        """
-        response = self.get_conn().queries().create(body=query).execute(num_retries=self.num_retries)
-        return response
-
-    @deprecated(
-        planned_removal_date="September 01, 2025",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def delete_query(self, query_id: str) -> None:
-        """
-        Delete a stored query as well as the associated stored reports.
-
-        :param query_id: Query ID to delete.
-        """
-        self.get_conn().queries().delete(queryId=query_id).execute(num_retries=self.num_retries)
-
-    @deprecated(
-        planned_removal_date="September 01, 2025",
-        use_instead="airflow.providers.google.marketing_platform.hooks.display_video.get_sdf_download_operation",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_query(self, query_id: str) -> dict:
-        """
-        Retrieve a stored query.
-
-        :param query_id: Query ID to retrieve.
-        """
-        response = self.get_conn().queries().get(queryId=query_id).execute(num_retries=self.num_retries)
-        return response
-
-    @deprecated(
-        planned_removal_date="September 01, 2025",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def list_queries(self) -> list[dict]:
-        """Retrieve stored queries."""
-        response = self.get_conn().queries().list().execute(num_retries=self.num_retries)
-        return response.get("queries", [])
-
-    @deprecated(
-        planned_removal_date="September 01, 2025",
-        use_instead="airflow.providers.google.marketing_platform.hooks.display_video.create_sdf_download_operation",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def run_query(self, query_id: str, params: dict[str, Any] | None) -> dict:
-        """
-        Run a stored query to generate a report.
-
-        :param query_id: Query ID to run.
-        :param params: Parameters for the report.
-        """
-        return (
-            self.get_conn().queries().run(queryId=query_id, body=params).execute(num_retries=self.num_retries)
-        )
-
-    @deprecated(
-        planned_removal_date="September 01, 2025",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_report(self, query_id: str, report_id: str) -> dict:
-        """
-        Retrieve a report.
-
-        :param query_id: Query ID for which report was generated.
-        :param report_id: Report ID to retrieve.
-        """
-        return (
-            self.get_conn()
-            .queries()
-            .reports()
-            .get(queryId=query_id, reportId=report_id)
-            .execute(num_retries=self.num_retries)
-        )
-
-    @deprecated(
-        planned_removal_date="September 01, 2025",
-        use_instead="airflow.providers.google.marketing_platform.hooks.display_video.create_sdf_download_operation",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def upload_line_items(self, line_items: Any) -> list[dict[str, Any]]:
-        """
-        Upload line items in CSV format.
-
-        :param line_items: downloaded data from GCS and passed to the body request
-        :return: response body.
-        """
-        request_body = {
-            "lineItems": line_items,
-            "dryRun": False,
-            "format": "CSV",
-        }
-
-        response = (
-            self.get_conn()
-            .lineitems()
-            .uploadlineitems(body=request_body)
-            .execute(num_retries=self.num_retries)
-        )
-        return response
-
-    @deprecated(
-        planned_removal_date="September 01, 2025",
-        use_instead="airflow.providers.google.marketing_platform.hooks.display_video.download_media",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def download_line_items(self, request_body: dict[str, Any]) -> list[Any]:
-        """
-        Retrieve line items in CSV format.
-
-        :param request_body: dictionary with parameters that should be passed into.
-            More information about it can be found here:
-            https://developers.google.com/bid-manager/v1.1/lineitems/downloadlineitems
-        """
-        response = (
-            self.get_conn()
-            .lineitems()
-            .downloadlineitems(body=request_body)
-            .execute(num_retries=self.num_retries)
-        )
-        return response["lineItems"]
-
     def create_sdf_download_operation(self, body_request: dict[str, Any]) -> dict[str, Any]:
         """
         Create an SDF Download Task and Returns an Operation.