Add deferrable mode to BigQueryTablePartitionExistenceSensor. (#29735)

Co-authored-by: Beata Kossakowska <bkossakowska@google.com>

Author: bkossakowska

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

@@ -3048,6 +3048,24 @@ async def get_job_output(
             job_query_response = await job_client.get_query_results(cast(Session, session))
             return job_query_response
 
+    async def create_job_for_partition_get(
+        self,
+        dataset_id: str | None,
+        project_id: str | None = None,
+    ):
+        """Create a new job and get the job_id using gcloud-aio."""
+        async with ClientSession() as session:
+            self.log.info("Executing create_job..")
+            job_client = await self.get_job_instance(project_id, "", session)
+
+            query_request = {
+                "query": "SELECT partition_id "
+                f"FROM `{project_id}.{dataset_id}.INFORMATION_SCHEMA.PARTITIONS`",
+                "useLegacySql": False,
+            }
+            job_query_resp = await job_client.query(query_request, cast(Session, session))
+            return job_query_resp["jobReference"]["jobId"]
+
     def get_records(self, query_results: dict[str, Any]) -> list[Any]:
         """
         Given the output query response from gcloud-aio bigquery, convert the response to records.

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

@@ -24,7 +24,10 @@
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
-from airflow.providers.google.cloud.triggers.bigquery import BigQueryTableExistenceTrigger
+from airflow.providers.google.cloud.triggers.bigquery import (
+    BigQueryTableExistenceTrigger,
+    BigQueryTablePartitionExistenceTrigger,
+)
 from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
@@ -244,3 +247,65 @@ def execute_complete(self, context: dict[str, Any], event: dict[str, str] | None
                 return event["message"]
             raise AirflowException(event["message"])
         raise AirflowException("No event received in trigger callback")
+
+
+class BigQueryTableExistencePartitionAsyncSensor(BigQueryTablePartitionExistenceSensor):
+    """
+    Checks for the existence of a partition within a table in Google BigQuery.
+
+    :param project_id: The Google cloud project in which to look for the table.
+       The connection supplied to the hook must provide
+       access to the specified project.
+    :param dataset_id: The name of the dataset in which to look for the table.
+       storage bucket.
+    :param partition_id: The name of the partition to check the existence of.
+    :param table_id: The name of the table to check the existence of.
+    :param gcp_conn_id: The connection ID used to connect to Google Cloud.
+    :param bigquery_conn_id: (Deprecated) The connection ID used to connect to Google Cloud.
+       This parameter has been deprecated. You should pass the gcp_conn_id parameter instead.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+       credentials, or chained list of accounts required to get the access_token
+       of the last account in the list, which will be impersonated in the request.
+       If set as a string, the account must grant the originating account
+       the Service Account Token Creator IAM role.
+       If set as a sequence, the identities from the list must grant
+       Service Account Token Creator IAM role to the directly preceding identity, with first
+       account from the list granting this role to the originating account (templated).
+    :param poke_interval: The interval in seconds to wait between checks table existence.
+    """
+
+    def __init__(self, poke_interval: int = 5, **kwargs):
+        super().__init__(**kwargs)
+        self.poke_interval = poke_interval
+
+    def execute(self, context: Context) -> None:
+        """Airflow runs this method on the worker and defers using the trigger."""
+        self.defer(
+            timeout=timedelta(seconds=self.timeout),
+            trigger=BigQueryTablePartitionExistenceTrigger(
+                dataset_id=self.dataset_id,
+                table_id=self.table_id,
+                project_id=self.project_id,
+                partition_id=self.partition_id,
+                poll_interval=self.poke_interval,
+                gcp_conn_id=self.gcp_conn_id,
+                hook_params={
+                    "impersonation_chain": self.impersonation_chain,
+                },
+            ),
+            method_name="execute_complete",
+        )
+
+    def execute_complete(self, context: dict[str, Any], event: dict[str, str] | None = None) -> str:
+        """
+        Callback for when the trigger fires - returns immediately.
+        Relies on trigger to throw an exception, otherwise it assumes execution was
+        successful.
+        """
+        table_uri = f"{self.project_id}:{self.dataset_id}.{self.table_id}"
+        self.log.info('Sensor checks existence of partition: "%s" in table: %s', self.partition_id, table_uri)
+        if event:
+            if event["status"] == "success":
+                return event["message"]
+            raise AirflowException(event["message"])
+        raise AirflowException("No event received in trigger callback")

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

@@ -529,3 +529,72 @@ async def _table_exists(
                 if err.status == 404:
                     return False
                 raise err
+
+
+class BigQueryTablePartitionExistenceTrigger(BigQueryTableExistenceTrigger):
+    """
+    Initialize the BigQuery Table Partition Existence Trigger with needed parameters
+    :param partition_id: The name of the partition to check the existence of.
+    :param project_id: Google Cloud Project where the job is running
+    :param dataset_id: The dataset ID of the requested table.
+    :param table_id: The table ID of the requested table.
+    :param gcp_conn_id: Reference to google cloud connection id
+    :param hook_params: params for hook
+    :param poll_interval: polling period in seconds to check for the status
+    """
+
+    def __init__(self, partition_id: str, **kwargs):
+        super().__init__(**kwargs)
+        self.partition_id = partition_id
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serializes BigQueryTablePartitionExistenceTrigger arguments and classpath."""
+        return (
+            "airflow.providers.google.cloud.triggers.bigquery.BigQueryTablePartitionExistenceTrigger",
+            {
+                "partition_id": self.partition_id,
+                "dataset_id": self.dataset_id,
+                "project_id": self.project_id,
+                "table_id": self.table_id,
+                "gcp_conn_id": self.gcp_conn_id,
+                "poll_interval": self.poll_interval,
+                "hook_params": self.hook_params,
+            },
+        )
+
+    async def run(self) -> AsyncIterator["TriggerEvent"]:  # type: ignore[override]
+        """Will run until the table exists in the Google Big Query."""
+        hook = BigQueryAsyncHook(gcp_conn_id=self.gcp_conn_id)
+        job_id = None
+        while True:
+            if job_id is not None:
+                status = await hook.get_job_status(job_id=job_id, project_id=self.project_id)
+                if status == "success":
+                    is_partition = await self._partition_exists(
+                        hook=hook, job_id=job_id, project_id=self.project_id
+                    )
+                    if is_partition:
+                        yield TriggerEvent(
+                            {
+                                "status": "success",
+                                "message": f"Partition: {self.partition_id} in table: {self.table_id}",
+                            }
+                        )
+                    job_id = None
+                elif status == "error":
+                    yield TriggerEvent({"status": "error", "message": status})
+                    return
+                self.log.info("Sleeping for %s seconds.", self.poll_interval)
+                await asyncio.sleep(self.poll_interval)
+
+            else:
+                job_id = await hook.create_job_for_partition_get(self.dataset_id, project_id=self.project_id)
+                self.log.info("Sleeping for %s seconds.", self.poll_interval)
+                await asyncio.sleep(self.poll_interval)
+
+    async def _partition_exists(self, hook: BigQueryAsyncHook, job_id: str | None, project_id: str):
+        query_results = await hook.get_job_output(job_id=job_id, project_id=project_id)
+        records = hook.get_records(query_results)
+        if records:
+            records = [row[0] for row in records]
+            return self.partition_id in records

docs/apache-airflow-providers-google/operators/cloud/bigquery.rst

@@ -509,6 +509,17 @@ To check that a table exists and has a partition you can use.
 
 For DAY partitioned tables, the partition_id parameter is a string on the "%Y%m%d" format
 
+Use the :class:`~airflow.providers.google.cloud.sensors.bigquery.BigQueryTableExistencePartitionAsyncSensor`
+(deferrable version) if you would like to free up the worker slots while the sensor is running.
+
+:class:`~airflow.providers.google.cloud.sensors.bigquery.BigQueryTableExistencePartitionAsyncSensor`.
+
+.. exampleinclude:: /../../tests/system/providers/google/cloud/bigquery/example_bigquery_sensors.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_sensor_bigquery_table_partition_async]
+    :end-before: [END howto_sensor_bigquery_table_partition_async]
+
 Reference
 ^^^^^^^^^
 

tests/providers/google/cloud/sensors/test_bigquery.py

@@ -23,10 +23,14 @@
 from airflow.exceptions import AirflowException, TaskDeferred
 from airflow.providers.google.cloud.sensors.bigquery import (
     BigQueryTableExistenceAsyncSensor,
+    BigQueryTableExistencePartitionAsyncSensor,
     BigQueryTableExistenceSensor,
     BigQueryTablePartitionExistenceSensor,
 )
-from airflow.providers.google.cloud.triggers.bigquery import BigQueryTableExistenceTrigger
+from airflow.providers.google.cloud.triggers.bigquery import (
+    BigQueryTableExistenceTrigger,
+    BigQueryTablePartitionExistenceTrigger,
+)
 
 TEST_PROJECT_ID = "test_project"
 TEST_DATASET_ID = "test_dataset"
@@ -156,3 +160,63 @@ def test_big_query_sensor_async_execute_complete_event_none(self):
         )
         with pytest.raises(AirflowException):
             task.execute_complete(context={}, event=None)
+
+
+class TestBigQueryTableExistencePartitionAsyncSensor(TestCase):
+    def test_big_query_table_existence_partition_sensor_async(self):
+        """
+        Asserts that a task is deferred and a BigQueryTablePartitionExistenceTrigger will be fired
+        when the BigQueryTableExistencePartitionAsyncSensor is executed.
+        """
+        task = BigQueryTableExistencePartitionAsyncSensor(
+            task_id="test_task_id",
+            project_id=TEST_PROJECT_ID,
+            dataset_id=TEST_DATASET_ID,
+            table_id=TEST_TABLE_ID,
+            partition_id=TEST_PARTITION_ID,
+        )
+        with pytest.raises(TaskDeferred) as exc:
+            task.execute(context={})
+        assert isinstance(
+            exc.value.trigger, BigQueryTablePartitionExistenceTrigger
+        ), "Trigger is not a BigQueryTablePartitionExistenceTrigger"
+
+    def test_big_query_table_existence_partition_sensor_async_execute_failure(self):
+        """Tests that an AirflowException is raised in case of error event"""
+        task = BigQueryTableExistencePartitionAsyncSensor(
+            task_id="test_task_id",
+            project_id=TEST_PROJECT_ID,
+            dataset_id=TEST_DATASET_ID,
+            table_id=TEST_TABLE_ID,
+            partition_id=TEST_PARTITION_ID,
+        )
+        with pytest.raises(AirflowException):
+            task.execute_complete(context={}, event={"status": "error", "message": "test failure message"})
+
+    def test_big_query_table_existence_partition_sensor_async_execute_complete_event_none(self):
+        """Asserts that logging occurs as expected"""
+        task = BigQueryTableExistencePartitionAsyncSensor(
+            task_id="task-id",
+            project_id=TEST_PROJECT_ID,
+            dataset_id=TEST_DATASET_ID,
+            table_id=TEST_TABLE_ID,
+            partition_id=TEST_PARTITION_ID,
+        )
+        with pytest.raises(AirflowException, match="No event received in trigger callback"):
+            task.execute_complete(context={}, event=None)
+
+    def test_big_query_table_existence_partition_sensor_async_execute_complete(self):
+        """Asserts that logging occurs as expected"""
+        task = BigQueryTableExistencePartitionAsyncSensor(
+            task_id="task-id",
+            project_id=TEST_PROJECT_ID,
+            dataset_id=TEST_DATASET_ID,
+            table_id=TEST_TABLE_ID,
+            partition_id=TEST_PARTITION_ID,
+        )
+        table_uri = f"{TEST_PROJECT_ID}:{TEST_DATASET_ID}.{TEST_TABLE_ID}"
+        with mock.patch.object(task.log, "info") as mock_log_info:
+            task.execute_complete(context={}, event={"status": "success", "message": "test"})
+        mock_log_info.assert_called_with(
+            'Sensor checks existence of partition: "%s" in table: %s', TEST_PARTITION_ID, table_uri
+        )

tests/providers/google/cloud/triggers/test_bigquery.py

@@ -33,6 +33,7 @@
     BigQueryInsertJobTrigger,
     BigQueryIntervalCheckTrigger,
     BigQueryTableExistenceTrigger,
+    BigQueryTablePartitionExistenceTrigger,
     BigQueryValueCheckTrigger,
 )
 from airflow.triggers.base import TriggerEvent
@@ -59,6 +60,7 @@
 TEST_IGNORE_ZERO = True
 TEST_GCP_CONN_ID = "TEST_GCP_CONN_ID"
 TEST_HOOK_PARAMS: dict[str, Any] = {}
+TEST_PARTITION_ID = "1234"
 
 
 def test_bigquery_insert_job_op_trigger_serialization():
@@ -1043,3 +1045,36 @@ async def test_table_exists_raise_exception(mock_get_table_client):
     )
     with pytest.raises(ClientResponseError):
         await trigger._table_exists(hook, TEST_DATASET_ID, TEST_TABLE_ID, TEST_GCP_PROJECT_ID)
+
+
+class TestBigQueryTablePartitionExistenceTrigger:
+    def test_big_query_table_existence_partition_trigger_serialization_should_execute_successfully(self):
+        """
+        Asserts that the BigQueryTablePartitionExistenceTrigger correctly serializes its arguments
+        and classpath.
+        """
+
+        trigger = BigQueryTablePartitionExistenceTrigger(
+            dataset_id=TEST_DATASET_ID,
+            table_id=TEST_TABLE_ID,
+            project_id=TEST_GCP_PROJECT_ID,
+            partition_id=TEST_PARTITION_ID,
+            poll_interval=POLLING_PERIOD_SECONDS,
+            gcp_conn_id=TEST_GCP_CONN_ID,
+            hook_params={},
+        )
+
+        classpath, kwargs = trigger.serialize()
+        assert (
+            classpath
+            == "airflow.providers.google.cloud.triggers.bigquery.BigQueryTablePartitionExistenceTrigger"
+        )
+        assert kwargs == {
+            "dataset_id": TEST_DATASET_ID,
+            "project_id": TEST_GCP_PROJECT_ID,
+            "table_id": TEST_TABLE_ID,
+            "partition_id": TEST_PARTITION_ID,
+            "gcp_conn_id": TEST_GCP_CONN_ID,
+            "poll_interval": POLLING_PERIOD_SECONDS,
+            "hook_params": TEST_HOOK_PARAMS,
+        }

tests/system/providers/google/cloud/bigquery/example_bigquery_sensors.py

@@ -33,6 +33,7 @@
 )
 from airflow.providers.google.cloud.sensors.bigquery import (
     BigQueryTableExistenceAsyncSensor,
+    BigQueryTableExistencePartitionAsyncSensor,
     BigQueryTableExistenceSensor,
     BigQueryTablePartitionExistenceSensor,
 )
@@ -117,6 +118,16 @@
     )
     # [END howto_sensor_bigquery_table_partition]
 
+    # [START howto_sensor_bigquery_table_partition_async]
+    check_table_partition_exists_async: BaseSensorOperator = BigQueryTableExistencePartitionAsyncSensor(
+        task_id="check_table_partition_exists_async",
+        partition_id=PARTITION_NAME,
+        project_id=PROJECT_ID,
+        dataset_id=DATASET_NAME,
+        table_id=TABLE_NAME,
+    )
+    # [END howto_sensor_bigquery_table_partition_async]
+
     delete_dataset = BigQueryDeleteDatasetOperator(
         task_id="delete_dataset",
         dataset_id=DATASET_NAME,