apache-airflow-providers-google 10.16.0rc1__py3-none-any.whl → 10.17.0rc1__py3-none-any.whl

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Transfer sensor."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/cloud/sensors/dataflow.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Dataflow sensor."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Callable, Sequence

airflow/providers/google/cloud/sensors/dataform.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Dataform sensor."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Iterable, Sequence

airflow/providers/google/cloud/sensors/datafusion.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Data Fusion sensors."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Iterable, Sequence

airflow/providers/google/cloud/sensors/dataplex.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Dataplex sensors."""
+
 from __future__ import annotations
 
 import time

airflow/providers/google/cloud/sensors/dataprep.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Dataprep Job sensor."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

airflow/providers/google/cloud/sensors/dataproc.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Dataproc Job sensor."""
+
 from __future__ import annotations
 
 import time

airflow/providers/google/cloud/sensors/gcs.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud Storage sensors."""
+
 from __future__ import annotations
 
 import os

airflow/providers/google/cloud/sensors/looker.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud Looker sensors."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING

airflow/providers/google/cloud/sensors/pubsub.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google PubSub sensor."""
+
 from __future__ import annotations
 
 from datetime import timedelta

airflow/providers/google/cloud/sensors/tasks.py

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Task sensor."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google BigQuery to BigQuery operator."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google BigQuery to Google Cloud Storage operator."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google BigQuery to MSSQL operator."""
+
 from __future__ import annotations
 
 import warnings

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google BigQuery to MySQL operator."""
+
 from __future__ import annotations
 
 import warnings

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google BigQuery to PostgreSQL operator."""
+
 from __future__ import annotations
 
 from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToSqlBaseOperator

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Base operator for BigQuery to SQL operators."""
+
 from __future__ import annotations
 
 import abc

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Facebook Ad Reporting to GCS operators."""
+
 from __future__ import annotations
 
 import csv

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Storage to BigQuery operator."""
+
 from __future__ import annotations
 
 import json
@@ -467,7 +468,7 @@ class GCSToBigQueryOperator(BaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
         if self.max_id_key:
-            self.log.info(f"Selecting the MAX value from BigQuery column '{self.max_id_key}'...")
+            self.log.info("Selecting the MAX value from BigQuery column %r...", self.max_id_key)
             select_command = (
                 f"SELECT MAX({self.max_id_key}) AS max_value "
                 f"FROM {self.destination_project_dataset_table}"
@@ -713,7 +714,7 @@ class GCSToBigQueryOperator(BaseOperator):
             if k not in src_fmt_configs and k in valid_configs:
                 src_fmt_configs[k] = v
 
-        for k, v in src_fmt_configs.items():
+        for k in src_fmt_configs:
             if k not in valid_configs:
                 raise ValueError(f"{k} is not a valid src_fmt_configs for type {source_format}.")
 

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Google Cloud Storage operator."""
+
 from __future__ import annotations
 
 import warnings

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud Storage to SFTP operator."""
+
 from __future__ import annotations
 
 import os

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains operator for uploading local file(s) to GCS."""
+
 from __future__ import annotations
 
 import os

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """MsSQL to GCS operator."""
+
 from __future__ import annotations
 
 import datetime

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """MySQL to GCS operator."""
+
 from __future__ import annotations
 
 import base64

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """PostgreSQL to GCS operator."""
+
 from __future__ import annotations
 
 import datetime
@@ -25,6 +26,7 @@ import uuid
 from decimal import Decimal
 
 import pendulum
+from slugify import slugify
 
 from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
 from airflow.providers.postgres.hooks.postgres import PostgresHook
@@ -112,7 +114,23 @@ class PostgresToGCSOperator(BaseSQLToGCSOperator):
         self.cursor_itersize = cursor_itersize
 
     def _unique_name(self):
-        return f"{self.dag_id}__{self.task_id}__{uuid.uuid4()}" if self.use_server_side_cursor else None
+        """
+        Generate a non-deterministic UUID for the cursor name using the task_id and dag_id.
+
+        Ensures the resulting name fits within the maximum length allowed for an identifier in Postgres.
+        """
+        if self.use_server_side_cursor:
+            separator = "__"
+            random_sufix = str(uuid.uuid4())
+            available_length = 63 - len(random_sufix) - (len(separator) * 2)
+            return separator.join(
+                (
+                    slugify(self.dag_id, allow_unicode=False, max_length=available_length // 2),
+                    slugify(self.task_id, allow_unicode=False, max_length=available_length // 2),
+                    random_sufix,
+                )
+            )
+        return None
 
     def query(self):
         """Query Postgres and returns a cursor to the results."""

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

@@ -269,6 +269,7 @@ class S3ToGCSOperator(S3ListOperator):
         self.defer(
             trigger=CloudStorageTransferServiceCreateJobsTrigger(
                 project_id=gcs_hook.project_id,
+                gcp_conn_id=self.gcp_conn_id,
                 job_names=job_names,
                 poll_interval=self.poll_interval,
             ),
@@ -319,14 +320,11 @@ class S3ToGCSOperator(S3ListOperator):
             body[TRANSFER_SPEC][OBJECT_CONDITIONS][INCLUDE_PREFIXES] = files_chunk
             job = transfer_hook.create_transfer_job(body=body)
 
-            s = "s" if len(files_chunk) > 1 else ""
-            self.log.info(f"Submitted job {job['name']} to transfer {len(files_chunk)} file{s}")
+            self.log.info("Submitted job %s to transfer %s file(s).", job["name"], len(files_chunk))
             job_names.append(job["name"])
 
         if len(files) > chunk_size:
-            js = "s" if len(job_names) > 1 else ""
-            fs = "s" if len(files) > 1 else ""
-            self.log.info(f"Overall submitted {len(job_names)} job{js} to transfer {len(files)} file{fs}")
+            self.log.info("Overall submitted %s job(s) to transfer %s file(s).", len(job_names), len(files))
 
         return job_names
 

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains SFTP to Google Cloud Storage operator."""
+
 from __future__ import annotations
 
 import os

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Base operator for SQL to GCS operators."""
+
 from __future__ import annotations
 
 import abc
@@ -153,9 +154,10 @@ class BaseSQLToGCSOperator(BaseOperator):
     def execute(self, context: Context):
         if self.partition_columns:
             self.log.info(
-                f"Found partition columns: {','.join(self.partition_columns)}. "
+                "Found partition columns: %s. "
                 "Assuming the SQL statement is properly sorted by these columns in "
-                "ascending or descending order."
+                "ascending or descending order.",
+                ",".join(self.partition_columns),
             )
 
         self.log.info("Executing query")

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

@@ -160,7 +160,6 @@ class BigQueryCheckTrigger(BigQueryInsertJobTrigger):
                                 "records": None,
                             }
                         )
-                        return
                     else:
                         # Extract only first record from the query results
                         first_record = records.pop(0)
@@ -171,7 +170,7 @@ class BigQueryCheckTrigger(BigQueryInsertJobTrigger):
                                 "records": first_record,
                             }
                         )
-                        return
+                    return
                 elif job_status["status"] == "error":
                     yield TriggerEvent({"status": "error", "message": job_status["message"]})
                     return
@@ -681,7 +680,9 @@ class BigQueryTablePartitionExistenceTrigger(BigQueryTableExistenceTrigger):
                 await asyncio.sleep(self.poll_interval)
 
             else:
-                job_id = await hook.create_job_for_partition_get(self.dataset_id, project_id=self.project_id)
+                job_id = await hook.create_job_for_partition_get(
+                    self.dataset_id, table_id=self.table_id, project_id=self.project_id
+                )
                 self.log.info("Sleeping for %s seconds.", self.poll_interval)
                 await asyncio.sleep(self.poll_interval)
 

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

@@ -140,7 +140,7 @@ class CloudBatchJobFinishedTrigger(BaseTrigger):
             yield TriggerEvent({"status": "error", "message": str(e)})
             return
 
-        self.log.exception(f"Job with name [{self.job_name}] timed out")
+        self.log.exception("Job with name [%s] timed out", self.job_name)
         yield TriggerEvent(
             {
                 "job_name": self.job_name,

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

@@ -120,6 +120,7 @@ class CloudRunJobFinishedTrigger(BaseTrigger):
                             "job_name": self.job_name,
                         }
                     )
+                return
             elif operation.error.message:
                 raise AirflowException(f"Cloud Run Job error: {operation.error.message}")
 

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud SQL triggers."""
+
 from __future__ import annotations
 
 import asyncio
@@ -79,6 +80,7 @@ class CloudSQLExportTrigger(BaseTrigger):
                             }
                         )
                         return
+
                     yield TriggerEvent(
                         {
                             "operation_name": operation["name"],

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

@@ -37,11 +37,19 @@ class CloudStorageTransferServiceCreateJobsTrigger(BaseTrigger):
     :param job_names: List of transfer jobs names.
     :param project_id: GCP project id.
     :param poll_interval: Interval in seconds between polls.
+    :param gcp_conn_id: The connection ID used to connect to Google Cloud.
     """
 
-    def __init__(self, job_names: list[str], project_id: str | None = None, poll_interval: int = 10) -> None:
+    def __init__(
+        self,
+        job_names: list[str],
+        project_id: str | None = None,
+        poll_interval: int = 10,
+        gcp_conn_id: str = "google_cloud_default",
+    ) -> None:
         super().__init__()
         self.project_id = project_id
+        self.gcp_conn_id = gcp_conn_id
         self.job_names = job_names
         self.poll_interval = poll_interval
 
@@ -53,6 +61,7 @@ class CloudStorageTransferServiceCreateJobsTrigger(BaseTrigger):
                 "project_id": self.project_id,
                 "job_names": self.job_names,
                 "poll_interval": self.poll_interval,
+                "gcp_conn_id": self.gcp_conn_id,
             },
         )
 
@@ -117,4 +126,7 @@ class CloudStorageTransferServiceCreateJobsTrigger(BaseTrigger):
             await asyncio.sleep(self.poll_interval)
 
     def get_async_hook(self) -> CloudDataTransferServiceAsyncHook:
-        return CloudDataTransferServiceAsyncHook(project_id=self.project_id)
+        return CloudDataTransferServiceAsyncHook(
+            project_id=self.project_id,
+            gcp_conn_id=self.gcp_conn_id,
+        )

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Dataplex triggers."""
+
 from __future__ import annotations
 
 import asyncio

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Dataproc triggers."""
+
 from __future__ import annotations
 
 import asyncio

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

@@ -27,12 +27,18 @@ from google.cloud.container_v1.types import Operation
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.cncf.kubernetes.triggers.pod import KubernetesPodTrigger
 from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction
-from airflow.providers.google.cloud.hooks.kubernetes_engine import GKEAsyncHook, GKEPodAsyncHook
+from airflow.providers.google.cloud.hooks.kubernetes_engine import (
+    GKEAsyncHook,
+    GKEKubernetesAsyncHook,
+    GKEPodAsyncHook,
+)
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
     from datetime import datetime
 
+    from kubernetes_asyncio.client import V1Job
+
 
 class GKEStartPodTrigger(KubernetesPodTrigger):
     """
@@ -147,6 +153,7 @@ class GKEStartPodTrigger(KubernetesPodTrigger):
             ssl_ca_cert=self._ssl_ca_cert,
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
+            enable_tcp_keepalive=True,
         )
 
 
@@ -211,7 +218,6 @@ class GKEOperationTrigger(BaseTrigger):
                     self.log.info("Operation is still running.")
                     self.log.info("Sleeping for %ss...", self.poll_interval)
                     await asyncio.sleep(self.poll_interval)
-
                 else:
                     yield TriggerEvent(
                         {
@@ -237,3 +243,67 @@ class GKEOperationTrigger(BaseTrigger):
                 impersonation_chain=self.impersonation_chain,
             )
         return self._hook
+
+
+class GKEJobTrigger(BaseTrigger):
+    """GKEJobTrigger run on the trigger worker to check the state of Job."""
+
+    def __init__(
+        self,
+        cluster_url: str,
+        ssl_ca_cert: str,
+        job_name: str,
+        job_namespace: str,
+        gcp_conn_id: str = "google_cloud_default",
+        poll_interval: float = 2,
+        impersonation_chain: str | Sequence[str] | None = None,
+    ) -> None:
+        super().__init__()
+        self.cluster_url = cluster_url
+        self.ssl_ca_cert = ssl_ca_cert
+        self.job_name = job_name
+        self.job_namespace = job_namespace
+        self.gcp_conn_id = gcp_conn_id
+        self.poll_interval = poll_interval
+        self.impersonation_chain = impersonation_chain
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serialize KubernetesCreateJobTrigger arguments and classpath."""
+        return (
+            "airflow.providers.google.cloud.triggers.kubernetes_engine.GKEJobTrigger",
+            {
+                "cluster_url": self.cluster_url,
+                "ssl_ca_cert": self.ssl_ca_cert,
+                "job_name": self.job_name,
+                "job_namespace": self.job_namespace,
+                "gcp_conn_id": self.gcp_conn_id,
+                "poll_interval": self.poll_interval,
+                "impersonation_chain": self.impersonation_chain,
+            },
+        )
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+        """Get current job status and yield a TriggerEvent."""
+        job: V1Job = await self.hook.wait_until_job_complete(name=self.job_name, namespace=self.job_namespace)
+        job_dict = job.to_dict()
+        error_message = self.hook.is_job_failed(job=job)
+        status = "error" if error_message else "success"
+        message = f"Job failed with error: {error_message}" if error_message else "Job completed successfully"
+        yield TriggerEvent(
+            {
+                "name": job.metadata.name,
+                "namespace": job.metadata.namespace,
+                "status": status,
+                "message": message,
+                "job": job_dict,
+            }
+        )
+
+    @cached_property
+    def hook(self) -> GKEKubernetesAsyncHook:
+        return GKEKubernetesAsyncHook(
+            cluster_url=self.cluster_url,
+            ssl_ca_cert=self.ssl_ca_cert,
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )

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

@@ -103,12 +103,14 @@ class MLEngineStartTrainingJobTrigger(BaseTrigger):
                             "message": "Job completed",
                         }
                     )
+                    return
                 elif response_from_hook == "pending":
                     self.log.info("Job is still running...")
                     self.log.info("Sleeping for %s seconds.", self.poll_interval)
                     await asyncio.sleep(self.poll_interval)
                 else:
                     yield TriggerEvent({"status": "error", "message": response_from_hook})
+                    return
 
         except Exception as e:
             self.log.exception("Exception occurred while checking for query completion")

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

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Google Cloud Pubsub triggers."""
+
 from __future__ import annotations
 
 import asyncio
@@ -101,9 +102,8 @@ class PubsubPullTrigger(BaseTrigger):
                 if pulled_messages:
                     if self.ack_messages:
                         await self.message_acknowledgement(pulled_messages)
-                        yield TriggerEvent({"status": "success", "message": pulled_messages})
-                    else:
-                        yield TriggerEvent({"status": "success", "message": pulled_messages})
+                    yield TriggerEvent({"status": "success", "message": pulled_messages})
+                    return
                 else:
                     pulled_messages = await self.hook.pull(
                         project_id=self.project_id,