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

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

@@ -30,6 +30,7 @@ from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
 from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
 from airflow.providers.google.cloud.triggers.bigquery import BigQueryInsertJobTrigger
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.utils.helpers import merge_dicts
 
 if TYPE_CHECKING:
@@ -104,7 +105,7 @@ class BigQueryToGCSOperator(BaseOperator):
         *,
         source_project_dataset_table: str,
         destination_cloud_storage_uris: list[str],
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         compression: str = "NONE",
         export_format: str = "CSV",
         field_delimiter: str = ",",

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

@@ -91,7 +91,7 @@ class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator):
         self.source_project_dataset_table = source_project_dataset_table
 
     def get_sql_hook(self) -> MsSqlHook:
-        return MsSqlHook(schema=self.database, mysql_conn_id=self.mssql_conn_id)
+        return MsSqlHook(schema=self.database, mssql_conn_id=self.mssql_conn_id)
 
     def persist_links(self, context: Context) -> None:
         project_id, dataset_id, table_id = self.source_project_dataset_table.split(".")

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

@@ -133,4 +133,5 @@ class BigQueryToSqlBaseOperator(BaseOperator):
                 rows=rows,
                 target_fields=self.selected_fields,
                 replace=self.replace,
+                commit_every=self.batch_size,
             )

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

@@ -42,6 +42,7 @@ from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQuery
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
 from airflow.providers.google.cloud.triggers.bigquery import BigQueryInsertJobTrigger
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.utils.helpers import merge_dicts
 
 if TYPE_CHECKING:
@@ -229,7 +230,7 @@ class GCSToBigQueryOperator(BaseOperator):
         job_id: str | None = None,
         force_rerun: bool = True,
         reattach_states: set[str] | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -749,7 +750,6 @@ class GCSToBigQueryOperator(BaseOperator):
         )
         from openlineage.client.run import Dataset
 
-        from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
         from airflow.providers.google.cloud.utils.openlineage import (
             get_facets_from_bq_table,
             get_identity_column_lineage_facet,
@@ -766,8 +766,7 @@ class GCSToBigQueryOperator(BaseOperator):
             "schema": output_dataset_facets["schema"],
         }
         input_datasets = []
-        for uri in sorted(self.source_uris):
-            bucket, blob = _parse_gcs_url(uri)
+        for blob in sorted(self.source_objects):
             additional_facets = {}
 
             if "*" in blob:
@@ -777,7 +776,7 @@ class GCSToBigQueryOperator(BaseOperator):
                     "symlink": SymlinksDatasetFacet(
                         identifiers=[
                             SymlinksDatasetFacetIdentifiers(
-                                namespace=f"gs://{bucket}", name=blob, type="file"
+                                namespace=f"gs://{self.bucket}", name=blob, type="file"
                             )
                         ]
                     ),
@@ -788,7 +787,7 @@ class GCSToBigQueryOperator(BaseOperator):
                     blob = "/"
 
             dataset = Dataset(
-                namespace=f"gs://{bucket}",
+                namespace=f"gs://{self.bucket}",
                 name=blob,
                 facets=merge_dicts(input_dataset_facets, additional_facets),
             )

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

@@ -234,8 +234,6 @@ class GCSToGCSOperator(BaseOperator):
         self.source_object_required = source_object_required
         self.exact_match = exact_match
         self.match_glob = match_glob
-        self.resolved_source_objects: set[str] = set()
-        self.resolved_target_objects: set[str] = set()
 
     def execute(self, context: Context):
         hook = GCSHook(
@@ -540,13 +538,6 @@ class GCSToGCSOperator(BaseOperator):
             self.destination_bucket,
             destination_object,
         )
-
-        self.resolved_source_objects.add(source_object)
-        if not destination_object:
-            self.resolved_target_objects.add(source_object)
-        else:
-            self.resolved_target_objects.add(destination_object)
-
         hook.rewrite(self.source_bucket, source_object, self.destination_bucket, destination_object)
 
         if self.move_object:
@@ -559,17 +550,36 @@ class GCSToGCSOperator(BaseOperator):
         This means we won't have to normalize self.source_object and self.source_objects,
         destination bucket and so on.
         """
+        from pathlib import Path
+
         from openlineage.client.run import Dataset
 
         from airflow.providers.openlineage.extractors import OperatorLineage
 
+        def _process_prefix(pref):
+            if WILDCARD in pref:
+                pref = pref.split(WILDCARD)[0]
+            # Use parent if not a file (dot not in name) and not a dir (ends with slash)
+            if "." not in pref.split("/")[-1] and not pref.endswith("/"):
+                pref = Path(pref).parent.as_posix()
+            return ["/" if pref in ("", "/", ".") else pref.rstrip("/")]  # Adjust root path
+
+        inputs = []
+        for prefix in self.source_objects:
+            result = _process_prefix(prefix)
+            inputs.extend(result)
+
+        if self.destination_object is None:
+            outputs = inputs.copy()
+        else:
+            outputs = _process_prefix(self.destination_object)
+
         return OperatorLineage(
             inputs=[
-                Dataset(namespace=f"gs://{self.source_bucket}", name=source)
-                for source in sorted(self.resolved_source_objects)
+                Dataset(namespace=f"gs://{self.source_bucket}", name=source) for source in sorted(set(inputs))
             ],
             outputs=[
                 Dataset(namespace=f"gs://{self.destination_bucket}", name=target)
-                for target in sorted(self.resolved_target_objects)
+                for target in sorted(set(outputs))
             ],
         )

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

@@ -17,13 +17,20 @@
 from __future__ import annotations
 
 import asyncio
-from typing import Any, AsyncIterator, Sequence, SupportsAbs
+from typing import TYPE_CHECKING, Any, AsyncIterator, Sequence, SupportsAbs
 
 from aiohttp import ClientSession
 from aiohttp.client_exceptions import ClientResponseError
 
+from airflow.exceptions import AirflowException
+from airflow.models.taskinstance import TaskInstance
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryAsyncHook, BigQueryTableAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
+from airflow.utils.session import provide_session
+from airflow.utils.state import TaskInstanceState
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm.session import Session
 
 
 class BigQueryInsertJobTrigger(BaseTrigger):
@@ -51,12 +58,13 @@ class BigQueryInsertJobTrigger(BaseTrigger):
         self,
         conn_id: str,
         job_id: str | None,
-        project_id: str | None,
+        project_id: str,
         location: str | None,
         dataset_id: str | None = None,
         table_id: str | None = None,
         poll_interval: float = 4.0,
         impersonation_chain: str | Sequence[str] | None = None,
+        cancel_on_kill: bool = True,
     ):
         super().__init__()
         self.log.info("Using the connection  %s .", conn_id)
@@ -69,6 +77,7 @@ class BigQueryInsertJobTrigger(BaseTrigger):
         self.table_id = table_id
         self.poll_interval = poll_interval
         self.impersonation_chain = impersonation_chain
+        self.cancel_on_kill = cancel_on_kill
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         """Serialize BigQueryInsertJobTrigger arguments and classpath."""
@@ -83,9 +92,40 @@ class BigQueryInsertJobTrigger(BaseTrigger):
                 "table_id": self.table_id,
                 "poll_interval": self.poll_interval,
                 "impersonation_chain": self.impersonation_chain,
+                "cancel_on_kill": self.cancel_on_kill,
             },
         )
 
+    @provide_session
+    def get_task_instance(self, session: Session) -> TaskInstance:
+        query = session.query(TaskInstance).filter(
+            TaskInstance.dag_id == self.task_instance.dag_id,
+            TaskInstance.task_id == self.task_instance.task_id,
+            TaskInstance.run_id == self.task_instance.run_id,
+            TaskInstance.map_index == self.task_instance.map_index,
+        )
+        task_instance = query.one_or_none()
+        if task_instance is None:
+            raise AirflowException(
+                "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
+                self.task_instance.dag_id,
+                self.task_instance.task_id,
+                self.task_instance.run_id,
+                self.task_instance.map_index,
+            )
+        return task_instance
+
+    def safe_to_cancel(self) -> bool:
+        """
+        Whether it is safe to cancel the external job which is being executed by this trigger.
+
+        This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
+        Because in those cases, we should NOT cancel the external job.
+        """
+        # Database query is needed to get the latest state of the task instance.
+        task_instance = self.get_task_instance()  # type: ignore[call-arg]
+        return task_instance.state != TaskInstanceState.DEFERRED
+
     async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
         """Get current job execution status and yields a TriggerEvent."""
         hook = self._get_async_hook()
@@ -113,6 +153,28 @@ class BigQueryInsertJobTrigger(BaseTrigger):
                         self.poll_interval,
                     )
                     await asyncio.sleep(self.poll_interval)
+        except asyncio.CancelledError:
+            if self.job_id and self.cancel_on_kill and self.safe_to_cancel():
+                self.log.info(
+                    "The job is safe to cancel the as airflow TaskInstance is not in deferred state."
+                )
+                self.log.info(
+                    "Cancelling job. Project ID: %s, Location: %s, Job ID: %s",
+                    self.project_id,
+                    self.location,
+                    self.job_id,
+                )
+                await hook.cancel_job(  # type: ignore[union-attr]
+                    job_id=self.job_id, project_id=self.project_id, location=self.location
+                )
+            else:
+                self.log.info(
+                    "Trigger may have shutdown. Skipping to cancel job because the airflow "
+                    "task is not cancelled yet: Project ID: %s, Location:%s, Job ID:%s",
+                    self.project_id,
+                    self.location,
+                    self.job_id,
+                )
         except Exception as e:
             self.log.exception("Exception occurred while checking for query completion")
             yield TriggerEvent({"status": "error", "message": str(e)})
@@ -137,6 +199,7 @@ class BigQueryCheckTrigger(BigQueryInsertJobTrigger):
                 "table_id": self.table_id,
                 "poll_interval": self.poll_interval,
                 "impersonation_chain": self.impersonation_chain,
+                "cancel_on_kill": self.cancel_on_kill,
             },
         )
 
@@ -194,9 +257,10 @@ class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
         (default: False).
     """
 
-    def __init__(self, as_dict: bool = False, **kwargs):
+    def __init__(self, as_dict: bool = False, selected_fields: str | None = None, **kwargs):
         super().__init__(**kwargs)
         self.as_dict = as_dict
+        self.selected_fields = selected_fields
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         """Serialize BigQueryInsertJobTrigger arguments and classpath."""
@@ -212,6 +276,7 @@ class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
                 "poll_interval": self.poll_interval,
                 "impersonation_chain": self.impersonation_chain,
                 "as_dict": self.as_dict,
+                "selected_fields": self.selected_fields,
             },
         )
 
@@ -224,7 +289,11 @@ class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
                 job_status = await hook.get_job_status(job_id=self.job_id, project_id=self.project_id)
                 if job_status["status"] == "success":
                     query_results = await hook.get_job_output(job_id=self.job_id, project_id=self.project_id)
-                    records = hook.get_records(query_results=query_results, as_dict=self.as_dict)
+                    records = hook.get_records(
+                        query_results=query_results,
+                        as_dict=self.as_dict,
+                        selected_fields=self.selected_fields,
+                    )
                     self.log.debug("Response from hook: %s", job_status["status"])
                     yield TriggerEvent(
                         {
@@ -282,7 +351,7 @@ class BigQueryIntervalCheckTrigger(BigQueryInsertJobTrigger):
         conn_id: str,
         first_job_id: str,
         second_job_id: str,
-        project_id: str | None,
+        project_id: str,
         table: str,
         metrics_thresholds: dict[str, int],
         location: str | None = None,
@@ -443,7 +512,7 @@ class BigQueryValueCheckTrigger(BigQueryInsertJobTrigger):
         sql: str,
         pass_value: int | float | str,
         job_id: str | None,
-        project_id: str | None,
+        project_id: str,
         tolerance: Any = None,
         dataset_id: str | None = None,
         table_id: str | None = None,

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

@@ -78,3 +78,71 @@ class CloudComposerExecutionTrigger(BaseTrigger):
                 "operation_done": operation.done,
             }
         )
+
+
+class CloudComposerAirflowCLICommandTrigger(BaseTrigger):
+    """The trigger wait for the Airflow CLI command result."""
+
+    def __init__(
+        self,
+        project_id: str,
+        region: str,
+        environment_id: str,
+        execution_cmd_info: dict,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        poll_interval: int = 10,
+    ):
+        super().__init__()
+        self.project_id = project_id
+        self.region = region
+        self.environment_id = environment_id
+        self.execution_cmd_info = execution_cmd_info
+        self.gcp_conn_id = gcp_conn_id
+        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",
+            {
+                "project_id": self.project_id,
+                "region": self.region,
+                "environment_id": self.environment_id,
+                "execution_cmd_info": self.execution_cmd_info,
+                "gcp_conn_id": self.gcp_conn_id,
+                "impersonation_chain": self.impersonation_chain,
+                "poll_interval": self.poll_interval,
+            },
+        )
+
+    async def run(self):
+        try:
+            result = await self.gcp_hook.wait_command_execution_result(
+                project_id=self.project_id,
+                region=self.region,
+                environment_id=self.environment_id,
+                execution_cmd_info=self.execution_cmd_info,
+                poll_interval=self.poll_interval,
+            )
+        except AirflowException as ex:
+            yield TriggerEvent(
+                {
+                    "status": "error",
+                    "message": str(ex),
+                }
+            )
+            return
+
+        yield TriggerEvent(
+            {
+                "status": "success",
+                "result": result,
+            }
+        )
+        return

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

@@ -23,6 +23,7 @@ import asyncio
 from typing import Sequence
 
 from airflow.providers.google.cloud.hooks.cloud_sql import CloudSQLAsyncHook, CloudSqlOperationStatus
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 
@@ -36,7 +37,7 @@ class CloudSQLExportTrigger(BaseTrigger):
     def __init__(
         self,
         operation_name: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         poke_interval: int = 20,

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

@@ -27,6 +27,7 @@ from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
     CloudDataTransferServiceAsyncHook,
 )
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 
@@ -43,7 +44,7 @@ class CloudStorageTransferServiceCreateJobsTrigger(BaseTrigger):
     def __init__(
         self,
         job_names: list[str],
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         poll_interval: int = 10,
         gcp_conn_id: str = "google_cloud_default",
     ) -> None: