diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py b/providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py index b9764ba1c4d46..8c553d065069d 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py @@ -80,6 +80,7 @@ from airflow.utils.hashlib_wrapper import md5 from airflow.utils.helpers import convert_camel_to_snake from airflow.utils.log.logging_mixin import LoggingMixin +from airflow.utils.types import DagRunType if TYPE_CHECKING: import pandas as pd @@ -1285,7 +1286,7 @@ def generate_job_id( task_id: str, logical_date: datetime | None, configuration: dict, - run_after: pendulum.DateTime | None = None, + run_after: pendulum.DateTime | datetime | None = None, force_rerun: bool = False, ) -> str: if force_rerun: @@ -1314,18 +1315,14 @@ def generate_job_id( job_id = f"airflow_{dag_id}_{task_id}_{job_id_timestamp.isoformat()}_{uniqueness_suffix}" return re.sub(r"[:\-+.]", "_", job_id) - def get_run_after_or_logical_date(self, context: Context) -> pendulum.DateTime: + def get_run_after_or_logical_date(self, context: Context) -> pendulum.DateTime | datetime | None: + dag_run = context.get("dag_run") + if not dag_run: + return pendulum.now("UTC") + if AIRFLOW_V_3_0_PLUS: - if dag_run := context.get("dag_run"): - run_after = pendulum.instance(dag_run.run_after) - else: - run_after = pendulum.now("UTC") - else: - if logical_date := context.get("logical_date"): - run_after = logical_date - else: - run_after = pendulum.now("UTC") - return run_after + return dag_run.start_date + return dag_run.start_date if dag_run.run_type == DagRunType.SCHEDULED else context.get("logical_date") def split_tablename( self, table_input: str, default_project_id: str, var_name: str | None = None diff --git a/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py b/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py index 4f8357abc3aa7..e0d801bdc5772 100644 --- a/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py +++ b/providers/google/tests/unit/google/cloud/hooks/test_bigquery.py @@ -40,6 +40,7 @@ from google.cloud.exceptions import NotFound from airflow.exceptions import AirflowException +from airflow.models import DagRun from airflow.providers.common.compat.assets import Asset from airflow.providers.common.compat.sdk import Context from airflow.providers.google.cloud.hooks.bigquery import ( @@ -52,6 +53,7 @@ _validate_src_fmt_configs, _validate_value, ) +from airflow.utils.types import DagRunType from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS @@ -684,14 +686,44 @@ def test_job_id_validity(self, mock_md5, test_dag_id, expected_job_id): assert job_id == expected_job_id def test_get_run_after_or_logical_date(self): + """Test get_run_after_or_logical_date for both Airflow 3.x and pre-3.0 behavior.""" if AIRFLOW_V_3_0_PLUS: - from airflow.models import DagRun + ctx = Context( + dag_run=DagRun( + run_type=DagRunType.MANUAL, + start_date=pendulum.datetime(2025, 2, 2, tz="UTC"), + ), + logical_date=pendulum.datetime(2025, 1, 1, tz="UTC"), + ) + assert self.hook.get_run_after_or_logical_date(ctx) == pendulum.datetime(2025, 2, 2, tz="UTC") + + ctx = Context( + dag_run=DagRun( + run_type=DagRunType.SCHEDULED, + start_date=pendulum.datetime(2025, 2, 2, tz="UTC"), + ), + logical_date=pendulum.datetime(2025, 1, 1, tz="UTC"), + ) + assert self.hook.get_run_after_or_logical_date(ctx) == pendulum.datetime(2025, 2, 2, tz="UTC") - ctx = Context(dag_run=DagRun(run_after=pendulum.datetime(2025, 1, 1))) else: - ctx = Context(logical_date=pendulum.datetime(2025, 1, 1)) + ctx = Context( + dag_run=DagRun( + run_type=DagRunType.MANUAL, + start_date=pendulum.datetime(2025, 2, 2, tz="UTC"), + ), + logical_date=pendulum.datetime(2025, 1, 1, tz="UTC"), + ) + assert self.hook.get_run_after_or_logical_date(ctx) == pendulum.datetime(2025, 1, 1, tz="UTC") - assert self.hook.get_run_after_or_logical_date(ctx) == pendulum.datetime(2025, 1, 1) + ctx = Context( + dag_run=DagRun( + run_type=DagRunType.SCHEDULED, + start_date=pendulum.datetime(2025, 2, 2, tz="UTC"), + ), + logical_date=pendulum.datetime(2025, 1, 1, tz="UTC"), + ) + assert self.hook.get_run_after_or_logical_date(ctx) == pendulum.datetime(2025, 2, 2, tz="UTC") @mock.patch( "airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_job",