Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix BigQueryTablePartitionExistenceTrigger partition query #37655

Merged
Merged
4 changes: 3 additions & 1 deletion airflow/providers/google/cloud/hooks/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -3313,6 +3313,7 @@ async def get_job_output(
async def create_job_for_partition_get(
self,
dataset_id: str | None,
table_id: str | None = None,
project_id: str | None = None,
):
"""Create a new job and get the job_id using gcloud-aio."""
Expand All @@ -3322,7 +3323,8 @@ async def create_job_for_partition_get(

query_request = {
"query": "SELECT partition_id "
f"FROM `{project_id}.{dataset_id}.INFORMATION_SCHEMA.PARTITIONS`",
f"FROM `{project_id}.{dataset_id}.INFORMATION_SCHEMA.PARTITIONS`"
+ (f" WHERE table_id={table_id}" if table_id else ""),
"useLegacySql": False,
}
job_query_resp = await job_client.query(query_request, cast(Session, session))
Expand Down
4 changes: 3 additions & 1 deletion airflow/providers/google/cloud/triggers/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -681,7 +681,9 @@ async def run(self) -> AsyncIterator[TriggerEvent]: # type: ignore[override]
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)

Expand Down