From 636f97fdd669f2b04f2bc3db6d3cd4ee61e434f1 Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Wed, 16 Jul 2025 08:39:36 -0400 Subject: [PATCH] Set up process for sharing code between different components And move `airflow.utils.timezone` into a shared library as the first example of it working. In this change we have now setled on an approach using symlinks, but we did explore other options (see the GH PR for discussion and previous versions, notably one built upon the `vendoring` tool) A lot of the reasoning and mode of operation of this is detailed in shared/README.md in this PR, hence why this description is so short. Currently various places in TaskSDK and Airflow Core both use these utility functions, and while in this specific case they are small enough that they could just be copied and the duplication wouldn't hurt us long term, this changes shows a way in which we can have a single source of truth, but have it included automatically in built dists. Co-authored-by: Jarek Potiuk Co-authored-by: Amogh Desai --- .dockerignore | 1 + airflow-core/docs/img/airflow_erd.sha256 | 2 +- airflow-core/pyproject.toml | 3 + airflow-core/src/airflow/_shared/__init__.py | 16 ++++ airflow-core/src/airflow/_shared/timezones | 1 + .../src/airflow/api/common/trigger_dag.py | 2 +- .../src/airflow/api_fastapi/auth/tokens.py | 2 +- .../airflow/api_fastapi/common/parameters.py | 2 +- .../src/airflow/api_fastapi/common/types.py | 2 +- .../core_api/datamodels/dag_run.py | 2 +- .../core_api/datamodels/ui/common.py | 2 +- .../core_api/routes/public/assets.py | 2 +- .../core_api/routes/public/backfills.py | 2 +- .../core_api/routes/public/hitl.py | 2 +- .../core_api/routes/ui/dashboard.py | 2 +- .../core_api/services/ui/calendar.py | 2 +- .../execution_api/routes/task_instances.py | 2 +- airflow-core/src/airflow/cli/cli_config.py | 2 +- .../src/airflow/cli/commands/dag_command.py | 3 +- .../src/airflow/cli/commands/task_command.py | 3 +- airflow-core/src/airflow/cli/utils.py | 2 +- .../src/airflow/dag_processing/collection.py | 2 +- .../src/airflow/dag_processing/manager.py | 14 ++- airflow-core/src/airflow/jobs/job.py | 2 +- .../src/airflow/jobs/scheduler_job_runner.py | 2 +- .../src/airflow/jobs/triggerer_job_runner.py | 2 +- .../versions/0047_3_0_0_add_dag_versioning.py | 2 +- airflow-core/src/airflow/models/asset.py | 2 +- airflow-core/src/airflow/models/backfill.py | 2 +- airflow-core/src/airflow/models/dag.py | 2 +- .../src/airflow/models/dag_version.py | 2 +- airflow-core/src/airflow/models/dagbag.py | 2 +- airflow-core/src/airflow/models/dagcode.py | 2 +- airflow-core/src/airflow/models/dagrun.py | 2 +- airflow-core/src/airflow/models/dagwarning.py | 2 +- .../src/airflow/models/db_callback_request.py | 2 +- airflow-core/src/airflow/models/deadline.py | 2 +- airflow-core/src/airflow/models/log.py | 2 +- .../src/airflow/models/serialized_dag.py | 2 +- .../src/airflow/models/taskinstance.py | 2 +- .../src/airflow/models/taskinstancehistory.py | 2 +- airflow-core/src/airflow/models/tasklog.py | 2 +- airflow-core/src/airflow/models/trigger.py | 2 +- airflow-core/src/airflow/models/xcom.py | 2 +- .../serialization/serialized_objects.py | 2 +- .../serialization/serializers/datetime.py | 2 +- .../serialization/serializers/timezone.py | 2 +- airflow-core/src/airflow/settings.py | 2 +- .../ti_deps/deps/not_in_retry_period_dep.py | 2 +- .../ti_deps/deps/ready_to_reschedule.py | 2 +- .../ti_deps/deps/runnable_exec_date_dep.py | 2 +- airflow-core/src/airflow/timetables/_cron.py | 2 +- airflow-core/src/airflow/timetables/_delta.py | 2 +- airflow-core/src/airflow/timetables/events.py | 2 +- .../src/airflow/timetables/interval.py | 2 +- airflow-core/src/airflow/timetables/simple.py | 2 +- .../src/airflow/timetables/trigger.py | 2 +- .../src/airflow/traces/otel_tracer.py | 2 +- airflow-core/src/airflow/utils/__init__.py | 5 + airflow-core/src/airflow/utils/cli.py | 3 +- .../src/airflow/utils/cli_action_loggers.py | 2 +- airflow-core/src/airflow/utils/db_cleanup.py | 2 +- .../utils/log/file_processor_handler.py | 2 +- .../airflow/utils/log/file_task_handler.py | 2 +- .../src/airflow/utils/log/timezone_aware.py | 6 +- airflow-core/src/airflow/utils/sqlalchemy.py | 2 +- .../tests/integration/otel/test_otel.py | 2 +- .../unit/api_fastapi/auth/test_tokens.py | 2 +- .../unit/api_fastapi/common/db/test_dags.py | 2 +- .../core_api/routes/public/test_assets.py | 2 +- .../core_api/routes/public/test_backfills.py | 2 +- .../core_api/routes/public/test_dag_run.py | 2 +- .../core_api/routes/public/test_dag_stats.py | 2 +- .../routes/public/test_extra_links.py | 2 +- .../core_api/routes/public/test_hitl.py | 2 +- .../core_api/routes/public/test_log.py | 2 +- .../core_api/routes/public/test_monitor.py | 2 +- .../routes/public/test_task_instances.py | 2 +- .../core_api/routes/public/test_xcom.py | 2 +- .../core_api/routes/ui/test_backfills.py | 2 +- .../core_api/routes/ui/test_grid.py | 2 +- .../core_api/routes/ui/test_structure.py | 2 +- .../versions/head/test_asset_events.py | 2 +- .../versions/head/test_assets.py | 2 +- .../versions/head/test_dag_runs.py | 2 +- .../versions/head/test_task_instances.py | 2 +- .../v2025_04_28/test_task_instances.py | 2 +- .../unit/callbacks/test_callback_requests.py | 2 +- .../cli/commands/test_backfill_command.py | 2 +- .../unit/cli/commands/test_dag_command.py | 7 +- .../unit/cli/commands/test_task_command.py | 2 +- airflow-core/tests/unit/core/test_core.py | 2 +- .../unit/core/test_impersonation_tests.py | 2 +- airflow-core/tests/unit/core/test_sentry.py | 2 +- .../unit/dag_processing/bundles/test_base.py | 2 +- .../unit/dag_processing/test_collection.py | 2 +- .../tests/unit/dag_processing/test_manager.py | 2 +- .../unit/dag_processing/test_processor.py | 2 +- .../unit/dags/test_cli_triggered_dags.py | 2 +- .../tests/unit/dags/test_invalid_cron.py | 2 +- .../tests/unit/dags/test_logging_in_dag.py | 2 +- .../tests/unit/dags/test_mark_state.py | 2 +- airflow-core/tests/unit/dags/test_on_kill.py | 2 +- .../tests/unit/dags/test_parsing_context.py | 2 +- .../tests/unit/dags/test_scheduler_dags.py | 2 +- airflow-core/tests/unit/dags/test_sensor.py | 2 +- .../tests/unit/decorators/test_mapped.py | 2 +- .../unit/executors/test_base_executor.py | 2 +- .../unit/executors/test_local_executor.py | 2 +- airflow-core/tests/unit/jobs/test_base_job.py | 2 +- .../tests/unit/jobs/test_scheduler_job.py | 5 +- .../tests/unit/jobs/test_triggerer_job.py | 2 +- .../tests/unit/listeners/test_listeners.py | 2 +- airflow-core/tests/unit/models/__init__.py | 2 +- .../tests/unit/models/test_backfill.py | 2 +- airflow-core/tests/unit/models/test_dag.py | 4 +- airflow-core/tests/unit/models/test_dagbag.py | 2 +- airflow-core/tests/unit/models/test_dagrun.py | 2 +- airflow-core/tests/unit/models/test_pool.py | 2 +- .../unit/models/test_renderedtifields.py | 2 +- .../tests/unit/models/test_taskinstance.py | 2 +- .../tests/unit/models/test_timestamp.py | 2 +- .../tests/unit/models/test_trigger.py | 4 +- airflow-core/tests/unit/models/test_xcom.py | 2 +- .../tests/unit/sensors/test_filesystem.py | 2 +- .../serialization/test_dag_serialization.py | 2 +- .../serialization/test_serialized_objects.py | 2 +- .../deps/test_not_in_retry_period_dep.py | 2 +- .../unit/ti_deps/deps/test_prev_dagrun_dep.py | 2 +- .../deps/test_ready_to_reschedule_dep.py | 2 +- .../deps/test_runnable_exec_date_dep.py | 2 +- .../unit/timetables/test_events_timetable.py | 2 +- .../timetables/test_interval_timetable.py | 2 +- .../unit/timetables/test_once_timetable.py | 2 +- .../unit/timetables/test_trigger_timetable.py | 2 +- .../unit/timetables/test_workday_timetable.py | 2 +- .../utils/log/test_file_processor_handler.py | 2 +- .../tests/unit/utils/log/test_log_reader.py | 2 +- .../tests/unit/utils/test_cli_util.py | 3 +- .../tests/unit/utils/test_db_cleanup.py | 2 +- .../tests/unit/utils/test_dot_renderer.py | 3 +- airflow-core/tests/unit/utils/test_helpers.py | 3 +- .../tests/unit/utils/test_serve_logs.py | 2 +- .../tests/unit/utils/test_sqlalchemy.py | 2 +- ...test_task_handler_with_custom_formatter.py | 2 +- .../utils/docker_command_utils.py | 19 ++-- .../src/tests_common/pytest_plugin.py | 34 +++++-- .../test_kubernetes_pod_operator.py | 3 +- .../microsoft/azure/operators/test_msgraph.py | 6 +- .../tests/unit/standard/sensors/test_time.py | 32 ++++--- pyproject.toml | 21 ++++- scripts/ci/docker-compose/local.yml | 45 ++++----- scripts/cov/core_coverage.py | 1 - setup_idea.py | 8 +- shared/README.md | 91 +++++++++++++++++++ shared/timezones/pyproject.toml | 47 ++++++++++ .../src/airflow_shared/timezones/__init__.py | 16 ++++ .../src/airflow_shared/timezones}/timezone.py | 13 +-- .../timezones/tests}/test_timezone.py | 4 +- task-sdk/pyproject.toml | 3 + task-sdk/src/airflow/sdk/_shared/__init__.py | 16 ++++ task-sdk/src/airflow/sdk/_shared/timezones | 1 + task-sdk/src/airflow/sdk/bases/decorator.py | 2 +- task-sdk/src/airflow/sdk/bases/operator.py | 2 +- task-sdk/src/airflow/sdk/bases/sensor.py | 2 +- task-sdk/src/airflow/sdk/definitions/dag.py | 8 +- .../src/airflow/sdk/execution_time/cache.py | 2 +- .../airflow/sdk/execution_time/task_runner.py | 2 +- task-sdk/src/airflow/sdk/timezone.py | 39 ++++++++ task-sdk/tests/task_sdk/api/test_client.py | 2 +- task-sdk/tests/task_sdk/bases/test_sensor.py | 24 ++--- .../task_sdk/dags/super_basic_deferred_run.py | 2 +- .../tests/task_sdk/docs/test_public_api.py | 1 + .../task_sdk/execution_time/test_comms.py | 2 +- .../task_sdk/execution_time/test_context.py | 3 +- .../task_sdk/execution_time/test_hitl.py | 2 +- .../execution_time/test_supervisor.py | 10 +- .../execution_time/test_task_runner.py | 3 +- 178 files changed, 541 insertions(+), 265 deletions(-) create mode 100644 airflow-core/src/airflow/_shared/__init__.py create mode 120000 airflow-core/src/airflow/_shared/timezones create mode 100644 shared/README.md create mode 100644 shared/timezones/pyproject.toml create mode 100644 shared/timezones/src/airflow_shared/timezones/__init__.py rename {airflow-core/src/airflow/utils => shared/timezones/src/airflow_shared/timezones}/timezone.py (96%) rename {airflow-core/tests/unit/utils => shared/timezones/tests}/test_timezone.py (98%) create mode 100644 task-sdk/src/airflow/sdk/_shared/__init__.py create mode 120000 task-sdk/src/airflow/sdk/_shared/timezones create mode 100644 task-sdk/src/airflow/sdk/timezone.py diff --git a/.dockerignore b/.dockerignore index 75e6291445ab6..716c737ee4e00 100644 --- a/.dockerignore +++ b/.dockerignore @@ -46,6 +46,7 @@ !docker-tests !helm-tests !kubernetes-tests +!shared/ # Add scripts so that we can use them inside the container !scripts diff --git a/airflow-core/docs/img/airflow_erd.sha256 b/airflow-core/docs/img/airflow_erd.sha256 index ef80b60bcb501..e1bd6b784fc63 100644 --- a/airflow-core/docs/img/airflow_erd.sha256 +++ b/airflow-core/docs/img/airflow_erd.sha256 @@ -1 +1 @@ -e605946ed7a64241d2225956c9baf9de1b8bde854c351873b4462030c5c5b98f \ No newline at end of file +cc20f5b01cbb926eef1576cd77b840849f4ded0848c4f1a582fec6d39fc2ddeb \ No newline at end of file diff --git a/airflow-core/pyproject.toml b/airflow-core/pyproject.toml index e36e1fa40bb49..5604382e4d19d 100644 --- a/airflow-core/pyproject.toml +++ b/airflow-core/pyproject.toml @@ -215,6 +215,9 @@ exclude = [ "src/airflow/ui/openapi.merged.json", ] +[tool.hatch.build.targets.sdist.force-include] +"../shared/timezones/src/airflow_shared/timezones" = "src/airflow/_shared/timezones" + [tool.hatch.build.targets.custom] path = "./hatch_build.py" diff --git a/airflow-core/src/airflow/_shared/__init__.py b/airflow-core/src/airflow/_shared/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/src/airflow/_shared/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow-core/src/airflow/_shared/timezones b/airflow-core/src/airflow/_shared/timezones new file mode 120000 index 0000000000000..8d7034fa71347 --- /dev/null +++ b/airflow-core/src/airflow/_shared/timezones @@ -0,0 +1 @@ +../../../../shared/timezones/src/airflow_shared/timezones \ No newline at end of file diff --git a/airflow-core/src/airflow/api/common/trigger_dag.py b/airflow-core/src/airflow/api/common/trigger_dag.py index 186ce9499a70a..d1cf5d93b8e51 100644 --- a/airflow-core/src/airflow/api/common/trigger_dag.py +++ b/airflow-core/src/airflow/api/common/trigger_dag.py @@ -22,9 +22,9 @@ import json from typing import TYPE_CHECKING +from airflow._shared.timezones import timezone from airflow.exceptions import DagNotFound, DagRunAlreadyExists from airflow.models import DagBag, DagModel, DagRun -from airflow.utils import timezone from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/auth/tokens.py b/airflow-core/src/airflow/api_fastapi/auth/tokens.py index f653532e29fb1..276ae17153da0 100644 --- a/airflow-core/src/airflow/api_fastapi/auth/tokens.py +++ b/airflow-core/src/airflow/api_fastapi/auth/tokens.py @@ -34,7 +34,7 @@ from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.serialization import load_pem_private_key -from airflow.utils import timezone +from airflow._shared.timezones import timezone if TYPE_CHECKING: from jwt.algorithms import AllowedKeys, AllowedPrivateKeys diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 767b9fed46db5..e8af5380f8a41 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -37,6 +37,7 @@ from sqlalchemy import Column, and_, case, func, not_, or_, select from sqlalchemy.inspection import inspect +from airflow._shared.timezones import timezone from airflow.api_fastapi.core_api.base import OrmClause from airflow.api_fastapi.core_api.security import GetUserDep from airflow.models import Base @@ -57,7 +58,6 @@ from airflow.models.taskinstance import TaskInstance from airflow.models.variable import Variable from airflow.typing_compat import Self -from airflow.utils import timezone from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/common/types.py b/airflow-core/src/airflow/api_fastapi/common/types.py index 18e5dc7387d62..ccbb9d9ad98c7 100644 --- a/airflow-core/src/airflow/api_fastapi/common/types.py +++ b/airflow-core/src/airflow/api_fastapi/common/types.py @@ -30,7 +30,7 @@ ConfigDict, ) -from airflow.utils import timezone +from airflow._shared.timezones import timezone UtcDateTime = Annotated[AwareDatetime, AfterValidator(lambda d: d.astimezone(timezone.utc))] """UTCDateTime is a datetime with timezone information""" diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py index bb79fa671ed2f..686e21b444a8e 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py @@ -23,11 +23,11 @@ from pydantic import AliasPath, AwareDatetime, Field, NonNegativeInt, model_validator +from airflow._shared.timezones import timezone from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel from airflow.api_fastapi.core_api.datamodels.dag_versions import DagVersionResponse from airflow.models import DagRun from airflow.timetables.base import DataInterval -from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py index 6089d6d55dfc4..6ff8b83825f02 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py @@ -22,8 +22,8 @@ from pydantic import computed_field +from airflow._shared.timezones import timezone from airflow.api_fastapi.core_api.base import BaseModel -from airflow.utils import timezone from airflow.utils.state import TaskInstanceState from airflow.utils.types import DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py index 61fb5f9299a59..ace75342932fc 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py @@ -24,6 +24,7 @@ from sqlalchemy import and_, delete, func, select from sqlalchemy.orm import joinedload, subqueryload +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.dagbag import DagBagDep from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import ( @@ -72,7 +73,6 @@ TaskOutletAssetReference, ) from airflow.models.dag import DAG -from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py index 525b82c976b68..30efea4d79ab8 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py @@ -24,6 +24,7 @@ from sqlalchemy import select, update from sqlalchemy.orm import joinedload +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import ( SessionDep, @@ -56,7 +57,6 @@ _create_backfill, _do_dry_run, ) -from airflow.utils import timezone from airflow.utils.state import DagRunState backfills_router = AirflowRouter(tags=["Backfill"], prefix="/backfills") diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py index aea8c8433cd0d..dfa8d039ac15e 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py @@ -23,6 +23,7 @@ from sqlalchemy import select from sqlalchemy.orm import joinedload +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import ( @@ -48,7 +49,6 @@ from airflow.api_fastapi.core_api.security import GetUserDep, ReadableTIFilterDep, requires_access_dag from airflow.models.hitl import HITLDetail as HITLDetailModel from airflow.models.taskinstance import TaskInstance as TI -from airflow.utils import timezone hitl_router = AirflowRouter(tags=["HumanInTheLoop"], prefix="/hitl-details") diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dashboard.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dashboard.py index f3c0198364137..9b552c9cf5883 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dashboard.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dashboard.py @@ -20,6 +20,7 @@ from sqlalchemy import func, select from sqlalchemy.sql.expression import case, false +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.parameters import DateTimeQuery, OptionalDateTimeQuery @@ -33,7 +34,6 @@ from airflow.models.dag import DagModel from airflow.models.dagrun import DagRun, DagRunType from airflow.models.taskinstance import TaskInstance -from airflow.utils import timezone from airflow.utils.state import DagRunState, TaskInstanceState dashboard_router = AirflowRouter(tags=["Dashboard"], prefix="/dashboard") diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py index de51a1dd27f49..28b7f32ead918 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py @@ -28,6 +28,7 @@ from sqlalchemy.engine import Row from sqlalchemy.orm import Session +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.parameters import RangeFilter from airflow.api_fastapi.core_api.datamodels.ui.calendar import ( CalendarTimeRangeCollectionResponse, @@ -38,7 +39,6 @@ from airflow.timetables._cron import CronMixin from airflow.timetables.base import DataInterval, TimeRestriction from airflow.timetables.simple import ContinuousTimetable -from airflow.utils import timezone log = structlog.get_logger(logger_name=__name__) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index def4dd155c6c7..09ef33d10b1cf 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -37,6 +37,7 @@ from sqlalchemy.sql import select from structlog.contextvars import bind_contextvars +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.dagbag import dag_bag_from_app from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.types import UtcDateTime @@ -67,7 +68,6 @@ from airflow.sdk.definitions._internal.expandinput import NotFullyPopulated from airflow.sdk.definitions.asset import Asset, AssetUniqueKey from airflow.sdk.definitions.taskgroup import MappedTaskGroup -from airflow.utils import timezone from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState if TYPE_CHECKING: diff --git a/airflow-core/src/airflow/cli/cli_config.py b/airflow-core/src/airflow/cli/cli_config.py index 2006b5f42b959..5bdccd654443c 100644 --- a/airflow-core/src/airflow/cli/cli_config.py +++ b/airflow-core/src/airflow/cli/cli_config.py @@ -29,12 +29,12 @@ import lazy_object_proxy +from airflow._shared.timezones.timezone import parse as parsedate from airflow.cli.commands.legacy_commands import check_legacy_command from airflow.configuration import conf from airflow.utils.cli import ColorMode from airflow.utils.module_loading import import_string from airflow.utils.state import DagRunState, JobState -from airflow.utils.timezone import parse as parsedate BUILD_DOCS = "BUILDING_AIRFLOW_DOCS" in os.environ diff --git a/airflow-core/src/airflow/cli/commands/dag_command.py b/airflow-core/src/airflow/cli/commands/dag_command.py index f23c52321f9a9..1fbba1d7e45a2 100644 --- a/airflow-core/src/airflow/cli/commands/dag_command.py +++ b/airflow-core/src/airflow/cli/commands/dag_command.py @@ -31,6 +31,7 @@ from sqlalchemy import func, select +from airflow._shared.timezones import timezone from airflow.api.client import get_current_api_client from airflow.api_fastapi.core_api.datamodels.dags import DAGResponse from airflow.cli.simple_table import AirflowConsole @@ -41,7 +42,7 @@ from airflow.models import DagBag, DagModel, DagRun, TaskInstance from airflow.models.errors import ParseImportError from airflow.models.serialized_dag import SerializedDagModel -from airflow.utils import cli as cli_utils, timezone +from airflow.utils import cli as cli_utils from airflow.utils.cli import get_dag, suppress_logs_and_warning, validate_dag_bundle_arg from airflow.utils.dot_renderer import render_dag, render_dag_dependencies from airflow.utils.helpers import ask_yesno diff --git a/airflow-core/src/airflow/cli/commands/task_command.py b/airflow-core/src/airflow/cli/commands/task_command.py index 61a8f81c86d72..656ba480ced2c 100644 --- a/airflow-core/src/airflow/cli/commands/task_command.py +++ b/airflow-core/src/airflow/cli/commands/task_command.py @@ -28,6 +28,7 @@ from typing import TYPE_CHECKING, Protocol, cast from airflow import settings +from airflow._shared.timezones import timezone from airflow.cli.simple_table import AirflowConsole from airflow.cli.utils import fetch_dag_run_from_run_id_or_logical_date_string from airflow.exceptions import AirflowConfigException, DagRunNotFound, TaskInstanceNotFound @@ -41,7 +42,7 @@ from airflow.serialization.serialized_objects import SerializedDAG from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_deps import SCHEDULER_QUEUED_DEPS -from airflow.utils import cli as cli_utils, timezone +from airflow.utils import cli as cli_utils from airflow.utils.cli import ( get_dag, get_dag_by_file_location, diff --git a/airflow-core/src/airflow/cli/utils.py b/airflow-core/src/airflow/cli/utils.py index 275bfbb7b07c1..4c7e6409e53a8 100644 --- a/airflow-core/src/airflow/cli/utils.py +++ b/airflow-core/src/airflow/cli/utils.py @@ -78,9 +78,9 @@ def fetch_dag_run_from_run_id_or_logical_date_string( from pendulum.parsing.exceptions import ParserError from sqlalchemy import select + from airflow._shared.timezones import timezone from airflow.models.dag import DAG from airflow.models.dagrun import DagRun - from airflow.utils import timezone if dag_run := DAG.fetch_dagrun(dag_id=dag_id, run_id=value, session=session): return dag_run, dag_run.logical_date diff --git a/airflow-core/src/airflow/dag_processing/collection.py b/airflow-core/src/airflow/dag_processing/collection.py index 3410fa4964dc4..805214a4861ee 100644 --- a/airflow-core/src/airflow/dag_processing/collection.py +++ b/airflow-core/src/airflow/dag_processing/collection.py @@ -35,6 +35,7 @@ from sqlalchemy.exc import OperationalError from sqlalchemy.orm import joinedload, load_only +from airflow._shared.timezones.timezone import utcnow from airflow.assets.manager import asset_manager from airflow.models.asset import ( AssetActive, @@ -56,7 +57,6 @@ from airflow.triggers.base import BaseEventTrigger from airflow.utils.retries import MAX_DB_RETRIES, run_with_db_retries from airflow.utils.sqlalchemy import with_row_locks -from airflow.utils.timezone import utcnow from airflow.utils.types import DagRunType if TYPE_CHECKING: diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 21ebffeb628b2..4557e52a46536 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -48,6 +48,7 @@ from uuid6 import uuid7 import airflow.models +from airflow._shared.timezones import timezone from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI from airflow.configuration import conf from airflow.dag_processing.bundles.manager import DagBundlesManager @@ -65,7 +66,6 @@ from airflow.sdk.log import init_log_file, logging_processors from airflow.stats import Stats from airflow.traces.tracer import DebugTrace -from airflow.utils import timezone from airflow.utils.file import list_py_file_paths, might_contain_dag from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.net import get_hostname @@ -139,6 +139,16 @@ def _resolve_path(instance: Any, attribute: attrs.Attribute, val: str | os.PathL return val +def utc_epoch() -> datetime: + # pendulum utcnow() is not used as that sets a TimezoneInfo object + # instead of a Timezone. This is not picklable and also creates issues + # when using replace() + result = datetime(1970, 1, 1) + result = result.replace(tzinfo=timezone.utc) + + return result + + @attrs.define(kw_only=True) class DagFileProcessorManager(LoggingMixin): """ @@ -508,7 +518,7 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): with create_session() as session: bundle_model: DagBundleModel = session.get(DagBundleModel, bundle.name) elapsed_time_since_refresh = ( - now - (bundle_model.last_refreshed or timezone.utc_epoch()) + now - (bundle_model.last_refreshed or utc_epoch()) ).total_seconds() if bundle.supports_versioning: # we will also check the version of the bundle to see if another DAG processor has seen diff --git a/airflow-core/src/airflow/jobs/job.py b/airflow-core/src/airflow/jobs/job.py index 72db0cb684ee6..67fb3be019293 100644 --- a/airflow-core/src/airflow/jobs/job.py +++ b/airflow-core/src/airflow/jobs/job.py @@ -27,6 +27,7 @@ from sqlalchemy.orm import backref, foreign, relationship from sqlalchemy.orm.session import make_transient +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.executors.executor_loader import ExecutorLoader @@ -34,7 +35,6 @@ from airflow.models.base import ID_LEN, Base from airflow.stats import Stats from airflow.traces.tracer import DebugTrace, add_debug_span -from airflow.utils import timezone from airflow.utils.helpers import convert_camel_to_snake from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.net import get_hostname diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 1157c3b51f8df..53f91991e89d3 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -38,6 +38,7 @@ from sqlalchemy.sql import expression from airflow import settings +from airflow._shared.timezones import timezone from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.configuration import conf @@ -69,7 +70,6 @@ from airflow.timetables.simple import AssetTriggeredTimetable from airflow.traces import utils as trace_utils from airflow.traces.tracer import DebugTrace, Trace, add_debug_span -from airflow.utils import timezone from airflow.utils.dates import datetime_to_nano from airflow.utils.event_scheduler import EventScheduler from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 8bacde4913d42..6a96ea161cb99 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -38,6 +38,7 @@ from sqlalchemy import func, select from structlog.contextvars import bind_contextvars as bind_log_contextvars +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.executors import workloads from airflow.jobs.base_job_runner import BaseJobRunner @@ -69,7 +70,6 @@ from airflow.stats import Stats from airflow.traces.tracer import DebugTrace, Trace, add_debug_span from airflow.triggers import base as events -from airflow.utils import timezone from airflow.utils.helpers import log_filename_template_renderer from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.module_loading import import_string diff --git a/airflow-core/src/airflow/migrations/versions/0047_3_0_0_add_dag_versioning.py b/airflow-core/src/airflow/migrations/versions/0047_3_0_0_add_dag_versioning.py index 2cad82e1e7f63..184c0a98218e0 100644 --- a/airflow-core/src/airflow/migrations/versions/0047_3_0_0_add_dag_versioning.py +++ b/airflow-core/src/airflow/migrations/versions/0047_3_0_0_add_dag_versioning.py @@ -31,10 +31,10 @@ from alembic import op from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.migrations.db_types import TIMESTAMP, StringID from airflow.migrations.utils import ignore_sqlite_value_error from airflow.models.base import naming_convention -from airflow.utils import timezone # revision identifiers, used by Alembic. revision = "2b47dc6bc8df" diff --git a/airflow-core/src/airflow/models/asset.py b/airflow-core/src/airflow/models/asset.py index a345ebe7995d3..8b28f2bea6184 100644 --- a/airflow-core/src/airflow/models/asset.py +++ b/airflow-core/src/airflow/models/asset.py @@ -36,9 +36,9 @@ ) from sqlalchemy.orm import relationship +from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID from airflow.settings import json -from airflow.utils import timezone from airflow.utils.sqlalchemy import UtcDateTime if TYPE_CHECKING: diff --git a/airflow-core/src/airflow/models/backfill.py b/airflow-core/src/airflow/models/backfill.py index e1fd0e506bfbc..4142265b8bd44 100644 --- a/airflow-core/src/airflow/models/backfill.py +++ b/airflow-core/src/airflow/models/backfill.py @@ -42,10 +42,10 @@ from sqlalchemy.orm import relationship, validates from sqlalchemy_jsonfield import JSONField +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException, DagNotFound from airflow.models.base import Base, StringID from airflow.settings import json -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.sqlalchemy import UtcDateTime, nulls_first, with_row_locks from airflow.utils.state import DagRunState diff --git a/airflow-core/src/airflow/models/dag.py b/airflow-core/src/airflow/models/dag.py index 43fc62b098d73..5c72eb83665cd 100644 --- a/airflow-core/src/airflow/models/dag.py +++ b/airflow-core/src/airflow/models/dag.py @@ -61,6 +61,7 @@ from sqlalchemy.sql import Select, expression from airflow import settings, utils +from airflow._shared.timezones import timezone from airflow.assets.evaluation import AssetEvaluator from airflow.configuration import conf as airflow_conf from airflow.exceptions import ( @@ -94,7 +95,6 @@ NullTimetable, OnceTimetable, ) -from airflow.utils import timezone from airflow.utils.context import Context from airflow.utils.dag_cycle_tester import check_cycle from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/airflow-core/src/airflow/models/dag_version.py b/airflow-core/src/airflow/models/dag_version.py index 5c1068d6a9614..a51f3cb301867 100644 --- a/airflow-core/src/airflow/models/dag_version.py +++ b/airflow-core/src/airflow/models/dag_version.py @@ -25,8 +25,8 @@ from sqlalchemy.orm import joinedload, relationship from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID -from airflow.utils import timezone from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime, with_row_locks diff --git a/airflow-core/src/airflow/models/dagbag.py b/airflow-core/src/airflow/models/dagbag.py index 83fc47444c179..a6ee7403b04f3 100644 --- a/airflow-core/src/airflow/models/dagbag.py +++ b/airflow-core/src/airflow/models/dagbag.py @@ -40,6 +40,7 @@ from tabulate import tabulate from airflow import settings +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.exceptions import ( AirflowClusterPolicyError, @@ -52,7 +53,6 @@ from airflow.listeners.listener import get_listener_manager from airflow.models.base import Base, StringID from airflow.stats import Stats -from airflow.utils import timezone from airflow.utils.dag_cycle_tester import check_cycle from airflow.utils.docs import get_docs_url from airflow.utils.file import ( diff --git a/airflow-core/src/airflow/models/dagcode.py b/airflow-core/src/airflow/models/dagcode.py index 2db338f87b50e..a68885e1dfdf2 100644 --- a/airflow-core/src/airflow/models/dagcode.py +++ b/airflow-core/src/airflow/models/dagcode.py @@ -26,10 +26,10 @@ from sqlalchemy.sql.expression import literal from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.exceptions import DagCodeNotFound from airflow.models.base import ID_LEN, Base -from airflow.utils import timezone from airflow.utils.file import open_maybe_zipped from airflow.utils.hashlib_wrapper import md5 from airflow.utils.session import NEW_SESSION, provide_session diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index cfac1936d6f87..b910e00f9b5b0 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -60,6 +60,7 @@ from sqlalchemy.sql.functions import coalesce from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.configuration import conf as airflow_conf from airflow.exceptions import AirflowException, TaskNotFound @@ -77,7 +78,6 @@ from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_states import SCHEDULEABLE_STATES from airflow.traces.tracer import EmptySpan, Trace -from airflow.utils import timezone from airflow.utils.dates import datetime_to_nano from airflow.utils.helpers import chunks, is_container, prune_dict from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/airflow-core/src/airflow/models/dagwarning.py b/airflow-core/src/airflow/models/dagwarning.py index b88158683eaf5..7498e615e30f7 100644 --- a/airflow-core/src/airflow/models/dagwarning.py +++ b/airflow-core/src/airflow/models/dagwarning.py @@ -22,9 +22,9 @@ from sqlalchemy import Column, ForeignKeyConstraint, Index, String, Text, delete, select, true +from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID from airflow.models.dag import DagModel -from airflow.utils import timezone from airflow.utils.retries import retry_db_transaction from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime diff --git a/airflow-core/src/airflow/models/db_callback_request.py b/airflow-core/src/airflow/models/db_callback_request.py index ada61d47c7a30..f1009ca1babe4 100644 --- a/airflow-core/src/airflow/models/db_callback_request.py +++ b/airflow-core/src/airflow/models/db_callback_request.py @@ -22,8 +22,8 @@ from sqlalchemy import Column, Integer, String +from airflow._shared.timezones import timezone from airflow.models.base import Base -from airflow.utils import timezone from airflow.utils.sqlalchemy import ExtendedJSON, UtcDateTime if TYPE_CHECKING: diff --git a/airflow-core/src/airflow/models/deadline.py b/airflow-core/src/airflow/models/deadline.py index e1df89fe7b3a2..a9e8cdabbf0a9 100644 --- a/airflow-core/src/airflow/models/deadline.py +++ b/airflow-core/src/airflow/models/deadline.py @@ -29,11 +29,11 @@ from sqlalchemy.orm import relationship from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.models import Trigger from airflow.models.base import Base, StringID from airflow.settings import json from airflow.triggers.deadline import DeadlineCallbackTrigger -from airflow.utils import timezone from airflow.utils.decorators import classproperty from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import provide_session diff --git a/airflow-core/src/airflow/models/log.py b/airflow-core/src/airflow/models/log.py index 1dcaf83941a9e..9347e4a17536c 100644 --- a/airflow-core/src/airflow/models/log.py +++ b/airflow-core/src/airflow/models/log.py @@ -22,8 +22,8 @@ from sqlalchemy import Column, Index, Integer, String, Text from sqlalchemy.orm import relationship +from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID -from airflow.utils import timezone from airflow.utils.sqlalchemy import UtcDateTime if TYPE_CHECKING: diff --git a/airflow-core/src/airflow/models/serialized_dag.py b/airflow-core/src/airflow/models/serialized_dag.py index ade8dca8761d4..e64a0e5cc895e 100644 --- a/airflow-core/src/airflow/models/serialized_dag.py +++ b/airflow-core/src/airflow/models/serialized_dag.py @@ -32,6 +32,7 @@ from sqlalchemy.sql.expression import func, literal from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.exceptions import TaskNotFound from airflow.models.asset import ( AssetAliasModel, @@ -46,7 +47,6 @@ from airflow.serialization.dag_dependency import DagDependency from airflow.serialization.serialized_objects import SerializedDAG from airflow.settings import COMPRESS_SERIALIZED_DAGS, json -from airflow.utils import timezone from airflow.utils.hashlib_wrapper import md5 from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index ca46b4802567a..9059c6d3986cb 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -70,6 +70,7 @@ from sqlalchemy_utils import UUIDType from airflow import settings +from airflow._shared.timezones import timezone from airflow.assets.manager import asset_manager from airflow.configuration import conf from airflow.exceptions import ( @@ -92,7 +93,6 @@ from airflow.stats import Stats from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_deps import REQUEUEABLE_DEPS, RUNNING_DEPS -from airflow.utils import timezone from airflow.utils.email import send_email from airflow.utils.helpers import prune_dict, render_template_to_string from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/airflow-core/src/airflow/models/taskinstancehistory.py b/airflow-core/src/airflow/models/taskinstancehistory.py index cbd83643fdb8c..b5932b28c58a3 100644 --- a/airflow-core/src/airflow/models/taskinstancehistory.py +++ b/airflow-core/src/airflow/models/taskinstancehistory.py @@ -38,8 +38,8 @@ from sqlalchemy.orm import relationship from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID -from airflow.utils import timezone from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.span_status import SpanStatus from airflow.utils.sqlalchemy import ( diff --git a/airflow-core/src/airflow/models/tasklog.py b/airflow-core/src/airflow/models/tasklog.py index d55eb94a266d7..d9a5c57c30ac5 100644 --- a/airflow-core/src/airflow/models/tasklog.py +++ b/airflow-core/src/airflow/models/tasklog.py @@ -19,8 +19,8 @@ from sqlalchemy import Column, Integer, Text +from airflow._shared.timezones import timezone from airflow.models.base import Base -from airflow.utils import timezone from airflow.utils.sqlalchemy import UtcDateTime diff --git a/airflow-core/src/airflow/models/trigger.py b/airflow-core/src/airflow/models/trigger.py index 4226bb2c3d96e..ded88d1d2f899 100644 --- a/airflow-core/src/airflow/models/trigger.py +++ b/airflow-core/src/airflow/models/trigger.py @@ -28,12 +28,12 @@ from sqlalchemy.orm import Session, relationship, selectinload from sqlalchemy.sql.functions import coalesce +from airflow._shared.timezones import timezone from airflow.assets.manager import AssetManager from airflow.models.asset import asset_trigger_association_table from airflow.models.base import Base from airflow.models.taskinstance import TaskInstance from airflow.triggers.base import BaseTaskEndEvent -from airflow.utils import timezone from airflow.utils.retries import run_with_db_retries from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime, with_row_locks diff --git a/airflow-core/src/airflow/models/xcom.py b/airflow-core/src/airflow/models/xcom.py index badb86bd43de5..a7301578ed2bd 100644 --- a/airflow-core/src/airflow/models/xcom.py +++ b/airflow-core/src/airflow/models/xcom.py @@ -39,8 +39,8 @@ from sqlalchemy.ext.associationproxy import association_proxy from sqlalchemy.orm import Query, relationship +from airflow._shared.timezones import timezone from airflow.models.base import COLLATION_ARGS, ID_LEN, TaskInstanceDependencies -from airflow.utils import timezone from airflow.utils.db import LazySelectSequence from airflow.utils.helpers import is_container from airflow.utils.json import XComDecoder, XComEncoder diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index 8c75249ac9829..7c9a6b254a96c 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -40,6 +40,7 @@ from pendulum.tz.timezone import FixedTimezone, Timezone from airflow import macros +from airflow._shared.timezones.timezone import from_timestamp, parse_timezone from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.exceptions import AirflowException, SerializationError, TaskDeferred from airflow.models.connection import Connection @@ -94,7 +95,6 @@ from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.module_loading import import_string, qualname from airflow.utils.operator_resources import Resources -from airflow.utils.timezone import from_timestamp, parse_timezone from airflow.utils.types import NOTSET, ArgNotSet if TYPE_CHECKING: diff --git a/airflow-core/src/airflow/serialization/serializers/datetime.py b/airflow-core/src/airflow/serialization/serializers/datetime.py index 806ba9feb692c..b5fe17a2e8b52 100644 --- a/airflow-core/src/airflow/serialization/serializers/datetime.py +++ b/airflow-core/src/airflow/serialization/serializers/datetime.py @@ -19,12 +19,12 @@ from typing import TYPE_CHECKING +from airflow._shared.timezones.timezone import parse_timezone from airflow.serialization.serializers.timezone import ( deserialize as deserialize_timezone, serialize as serialize_timezone, ) from airflow.utils.module_loading import qualname -from airflow.utils.timezone import parse_timezone if TYPE_CHECKING: import datetime diff --git a/airflow-core/src/airflow/serialization/serializers/timezone.py b/airflow-core/src/airflow/serialization/serializers/timezone.py index c4d1c30142c6e..3a67dd8a95ea4 100644 --- a/airflow-core/src/airflow/serialization/serializers/timezone.py +++ b/airflow-core/src/airflow/serialization/serializers/timezone.py @@ -70,7 +70,7 @@ def serialize(o: object) -> tuple[U, str, int, bool]: def deserialize(cls: type, version: int, data: object) -> Any: from zoneinfo import ZoneInfo - from airflow.utils.timezone import parse_timezone + from airflow._shared.timezones.timezone import parse_timezone if not isinstance(data, (str, int)): raise TypeError(f"{data} is not of type int or str but of {type(data)}") diff --git a/airflow-core/src/airflow/settings.py b/airflow-core/src/airflow/settings.py index 145b4819ac3f1..40634852eb6fd 100644 --- a/airflow-core/src/airflow/settings.py +++ b/airflow-core/src/airflow/settings.py @@ -36,12 +36,12 @@ from sqlalchemy.pool import NullPool from airflow import __version__ as airflow_version, policies +from airflow._shared.timezones.timezone import local_timezone, parse_timezone, utc from airflow.configuration import AIRFLOW_HOME, conf from airflow.exceptions import AirflowInternalRuntimeError from airflow.logging_config import configure_logging from airflow.utils.orm_event_handlers import setup_event_handlers from airflow.utils.sqlalchemy import is_sqlalchemy_v1 -from airflow.utils.timezone import local_timezone, parse_timezone, utc if TYPE_CHECKING: from sqlalchemy.engine import Engine diff --git a/airflow-core/src/airflow/ti_deps/deps/not_in_retry_period_dep.py b/airflow-core/src/airflow/ti_deps/deps/not_in_retry_period_dep.py index 90954f29f2f50..1013fecdfb811 100644 --- a/airflow-core/src/airflow/ti_deps/deps/not_in_retry_period_dep.py +++ b/airflow-core/src/airflow/ti_deps/deps/not_in_retry_period_dep.py @@ -17,8 +17,8 @@ # under the License. from __future__ import annotations +from airflow._shared.timezones import timezone from airflow.ti_deps.deps.base_ti_dep import BaseTIDep -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import TaskInstanceState diff --git a/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py b/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py index abe0d38c2b707..501b1574205e1 100644 --- a/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py +++ b/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py @@ -17,10 +17,10 @@ # under the License. from __future__ import annotations +from airflow._shared.timezones import timezone from airflow.executors.executor_loader import ExecutorLoader from airflow.models.taskreschedule import TaskReschedule from airflow.ti_deps.deps.base_ti_dep import BaseTIDep -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import TaskInstanceState diff --git a/airflow-core/src/airflow/ti_deps/deps/runnable_exec_date_dep.py b/airflow-core/src/airflow/ti_deps/deps/runnable_exec_date_dep.py index 0f996e5e9f409..396f8180d1be1 100644 --- a/airflow-core/src/airflow/ti_deps/deps/runnable_exec_date_dep.py +++ b/airflow-core/src/airflow/ti_deps/deps/runnable_exec_date_dep.py @@ -17,8 +17,8 @@ # under the License. from __future__ import annotations +from airflow._shared.timezones import timezone from airflow.ti_deps.deps.base_ti_dep import BaseTIDep -from airflow.utils import timezone from airflow.utils.session import provide_session diff --git a/airflow-core/src/airflow/timetables/_cron.py b/airflow-core/src/airflow/timetables/_cron.py index a0d7e5091c9d9..e62f96de77029 100644 --- a/airflow-core/src/airflow/timetables/_cron.py +++ b/airflow-core/src/airflow/timetables/_cron.py @@ -22,9 +22,9 @@ from cron_descriptor import CasingTypeEnum, ExpressionDescriptor, FormatException, MissingFieldException from croniter import CroniterBadCronError, CroniterBadDateError, croniter +from airflow._shared.timezones.timezone import convert_to_utc, make_aware, make_naive, parse_timezone from airflow.exceptions import AirflowTimetableInvalid from airflow.utils.dates import cron_presets -from airflow.utils.timezone import convert_to_utc, make_aware, make_naive, parse_timezone if TYPE_CHECKING: from pendulum import DateTime diff --git a/airflow-core/src/airflow/timetables/_delta.py b/airflow-core/src/airflow/timetables/_delta.py index 7203cd406310f..acdf6aa704ec3 100644 --- a/airflow-core/src/airflow/timetables/_delta.py +++ b/airflow-core/src/airflow/timetables/_delta.py @@ -20,8 +20,8 @@ import datetime from typing import TYPE_CHECKING +from airflow._shared.timezones.timezone import convert_to_utc from airflow.exceptions import AirflowTimetableInvalid -from airflow.utils.timezone import convert_to_utc if TYPE_CHECKING: from dateutil.relativedelta import relativedelta diff --git a/airflow-core/src/airflow/timetables/events.py b/airflow-core/src/airflow/timetables/events.py index ac2c0b6131982..d8e70626d409a 100644 --- a/airflow-core/src/airflow/timetables/events.py +++ b/airflow-core/src/airflow/timetables/events.py @@ -22,8 +22,8 @@ import pendulum +from airflow._shared.timezones import timezone from airflow.timetables.base import DagRunInfo, DataInterval, Timetable -from airflow.utils import timezone if TYPE_CHECKING: from pendulum import DateTime diff --git a/airflow-core/src/airflow/timetables/interval.py b/airflow-core/src/airflow/timetables/interval.py index 7e68b9a01e624..e248dd67da206 100644 --- a/airflow-core/src/airflow/timetables/interval.py +++ b/airflow-core/src/airflow/timetables/interval.py @@ -22,10 +22,10 @@ from dateutil.relativedelta import relativedelta from pendulum import DateTime +from airflow._shared.timezones.timezone import coerce_datetime, utcnow from airflow.timetables._cron import CronMixin from airflow.timetables._delta import DeltaMixin from airflow.timetables.base import DagRunInfo, DataInterval, Timetable -from airflow.utils.timezone import coerce_datetime, utcnow if TYPE_CHECKING: from airflow.timetables.base import TimeRestriction diff --git a/airflow-core/src/airflow/timetables/simple.py b/airflow-core/src/airflow/timetables/simple.py index bd9408434e84b..b5b6f2468f369 100644 --- a/airflow-core/src/airflow/timetables/simple.py +++ b/airflow-core/src/airflow/timetables/simple.py @@ -19,8 +19,8 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any +from airflow._shared.timezones import timezone from airflow.timetables.base import DagRunInfo, DataInterval, Timetable -from airflow.utils import timezone if TYPE_CHECKING: from pendulum import DateTime diff --git a/airflow-core/src/airflow/timetables/trigger.py b/airflow-core/src/airflow/timetables/trigger.py index 87e52a3f2640d..ed2f65cf24511 100644 --- a/airflow-core/src/airflow/timetables/trigger.py +++ b/airflow-core/src/airflow/timetables/trigger.py @@ -23,10 +23,10 @@ import time from typing import TYPE_CHECKING, Any +from airflow._shared.timezones.timezone import coerce_datetime, utcnow from airflow.timetables._cron import CronMixin from airflow.timetables._delta import DeltaMixin from airflow.timetables.base import DagRunInfo, DataInterval, Timetable -from airflow.utils.timezone import coerce_datetime, utcnow if TYPE_CHECKING: from dateutil.relativedelta import relativedelta diff --git a/airflow-core/src/airflow/traces/otel_tracer.py b/airflow-core/src/airflow/traces/otel_tracer.py index 982e8826e279f..22fb1d2935103 100644 --- a/airflow-core/src/airflow/traces/otel_tracer.py +++ b/airflow-core/src/airflow/traces/otel_tracer.py @@ -33,12 +33,12 @@ from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator from opentelemetry.trace.span import INVALID_SPAN_ID, INVALID_TRACE_ID +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.traces.utils import ( parse_traceparent, parse_tracestate, ) -from airflow.utils import timezone from airflow.utils.dates import datetime_to_nano from airflow.utils.net import get_hostname diff --git a/airflow-core/src/airflow/utils/__init__.py b/airflow-core/src/airflow/utils/__init__.py index 563d133eb2b35..5f3fb97911333 100644 --- a/airflow-core/src/airflow/utils/__init__.py +++ b/airflow-core/src/airflow/utils/__init__.py @@ -34,6 +34,11 @@ "get_task_group_children_getter": "airflow.sdk.definitions.taskgroup.get_task_group_children_getter", "task_group_to_dict": "airflow.sdk.definitions.taskgroup.task_group_to_dict", }, + "timezone": { + # Since we have corrected all uses inside core to use the internal version, anything hitting this + # should be in user code or custom providers, so redirect them to the public interface in Task SDK + "*": "airflow.sdk.timezone" + }, } add_deprecated_classes(__deprecated_classes, __name__) diff --git a/airflow-core/src/airflow/utils/cli.py b/airflow-core/src/airflow/utils/cli.py index ac7723d29338c..8f1fc0748bb7d 100644 --- a/airflow-core/src/airflow/utils/cli.py +++ b/airflow-core/src/airflow/utils/cli.py @@ -34,11 +34,12 @@ from typing import TYPE_CHECKING, TypeVar, cast from airflow import settings +from airflow._shared.timezones import timezone from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.exceptions import AirflowException from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager from airflow.sdk.execution_time.secrets_masker import should_hide_value_for_key -from airflow.utils import cli_action_loggers, timezone +from airflow.utils import cli_action_loggers from airflow.utils.log.non_caching_file_handler import NonCachingFileHandler from airflow.utils.platform import getuser, is_terminal_support_colors diff --git a/airflow-core/src/airflow/utils/cli_action_loggers.py b/airflow-core/src/airflow/utils/cli_action_loggers.py index cd53f2c91a2dc..dd3d8d59fa864 100644 --- a/airflow-core/src/airflow/utils/cli_action_loggers.py +++ b/airflow-core/src/airflow/utils/cli_action_loggers.py @@ -123,8 +123,8 @@ def default_action_log( """ from sqlalchemy.exc import OperationalError, ProgrammingError + from airflow._shared.timezones import timezone from airflow.models.log import Log - from airflow.utils import timezone try: # Use bulk_insert_mappings here to avoid importing all models (which using the classes does) early diff --git a/airflow-core/src/airflow/utils/db_cleanup.py b/airflow-core/src/airflow/utils/db_cleanup.py index 96e803775f803..7b0ee0487ddb8 100644 --- a/airflow-core/src/airflow/utils/db_cleanup.py +++ b/airflow-core/src/airflow/utils/db_cleanup.py @@ -36,10 +36,10 @@ from sqlalchemy.orm import aliased from sqlalchemy.sql.expression import ClauseElement, Executable, tuple_ +from airflow._shared.timezones import timezone from airflow.cli.simple_table import AirflowConsole from airflow.configuration import conf from airflow.exceptions import AirflowException -from airflow.utils import timezone from airflow.utils.db import reflect_tables from airflow.utils.helpers import ask_yesno from airflow.utils.session import NEW_SESSION, provide_session diff --git a/airflow-core/src/airflow/utils/log/file_processor_handler.py b/airflow-core/src/airflow/utils/log/file_processor_handler.py index 5dbd034cea8c2..ea3151be1db92 100644 --- a/airflow-core/src/airflow/utils/log/file_processor_handler.py +++ b/airflow-core/src/airflow/utils/log/file_processor_handler.py @@ -23,7 +23,7 @@ from pathlib import Path from airflow import settings -from airflow.utils import timezone +from airflow._shared.timezones import timezone from airflow.utils.helpers import parse_template_string from airflow.utils.log.logging_mixin import DISABLE_PROPOGATE from airflow.utils.log.non_caching_file_handler import NonCachingFileHandler diff --git a/airflow-core/src/airflow/utils/log/file_task_handler.py b/airflow-core/src/airflow/utils/log/file_task_handler.py index 6f69d3c54718c..124da295bb08a 100644 --- a/airflow-core/src/airflow/utils/log/file_task_handler.py +++ b/airflow-core/src/airflow/utils/log/file_task_handler.py @@ -251,7 +251,7 @@ def _log_stream_to_parsed_log_stream( :param log_stream: The stream to parse. :return: A generator of parsed log lines. """ - from airflow.utils.timezone import coerce_datetime + from airflow._shared.timezones.timezone import coerce_datetime timestamp = None next_timestamp = None diff --git a/airflow-core/src/airflow/utils/log/timezone_aware.py b/airflow-core/src/airflow/utils/log/timezone_aware.py index 0c9571c82574d..b7e69d03a73b1 100644 --- a/airflow-core/src/airflow/utils/log/timezone_aware.py +++ b/airflow-core/src/airflow/utils/log/timezone_aware.py @@ -18,8 +18,6 @@ import logging -from airflow.utils import timezone - class TimezoneAware(logging.Formatter): """ @@ -41,7 +39,9 @@ def formatTime(self, record, datefmt=None): This returns the creation time of the specified LogRecord in ISO 8601 date and time format in the local time zone. """ - dt = timezone.from_timestamp(record.created, tz="local") + from airflow._shared.timezones.timezone import from_timestamp + + dt = from_timestamp(record.created, tz="local") s = dt.strftime(datefmt or self.default_time_format) if self.default_msec_format: s = self.default_msec_format % (s, record.msecs) diff --git a/airflow-core/src/airflow/utils/sqlalchemy.py b/airflow-core/src/airflow/utils/sqlalchemy.py index d08307db14bf5..072f767335c6f 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/airflow-core/src/airflow/utils/sqlalchemy.py @@ -31,9 +31,9 @@ from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.types import JSON, Text, TypeDecorator +from airflow._shared.timezones.timezone import make_naive, utc from airflow.configuration import conf from airflow.serialization.enums import Encoding -from airflow.utils.timezone import make_naive, utc if TYPE_CHECKING: from collections.abc import Iterable diff --git a/airflow-core/tests/integration/otel/test_otel.py b/airflow-core/tests/integration/otel/test_otel.py index 99e99c3941894..03f427227c113 100644 --- a/airflow-core/tests/integration/otel/test_otel.py +++ b/airflow-core/tests/integration/otel/test_otel.py @@ -25,13 +25,13 @@ import pytest +from airflow._shared.timezones import timezone from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.executors import executor_loader from airflow.executors.executor_utils import ExecutorName from airflow.models import DAG, DagBag, DagRun from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.span_status import SpanStatus from airflow.utils.state import State diff --git a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py index 6afef4db20662..14b6957ee651d 100644 --- a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py +++ b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py @@ -27,6 +27,7 @@ import pytest from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import ( JWKS, InvalidClaimError, @@ -37,7 +38,6 @@ key_to_jwk_dict, key_to_pem, ) -from airflow.utils import timezone from tests_common.test_utils.config import conf_vars diff --git a/airflow-core/tests/unit/api_fastapi/common/db/test_dags.py b/airflow-core/tests/unit/api_fastapi/common/db/test_dags.py index b674ddb9a8b2e..051ffa965cc7e 100644 --- a/airflow-core/tests/unit/api_fastapi/common/db/test_dags.py +++ b/airflow-core/tests/unit/api_fastapi/common/db/test_dags.py @@ -21,12 +21,12 @@ import pytest +from airflow._shared.timezones.timezone import utcnow from airflow.api_fastapi.common.db.dags import generate_dag_with_latest_run_query from airflow.api_fastapi.common.parameters import SortParam from airflow.models import DagModel from airflow.models.dagrun import DagRun from airflow.utils.state import DagRunState -from airflow.utils.timezone import utcnow from tests_common.test_utils.db import clear_db_dags, clear_db_runs diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py index ce6ff10b30177..05a879c2f5be6 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py @@ -23,6 +23,7 @@ import pytest import time_machine +from airflow._shared.timezones import timezone from airflow.models import DagModel from airflow.models.asset import ( AssetActive, @@ -35,7 +36,6 @@ ) from airflow.models.dagrun import DagRun from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py index 703595b96e78f..b3eda51c1274b 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py @@ -24,12 +24,12 @@ import pytest from sqlalchemy import and_, func, select +from airflow._shared.timezones import timezone from airflow.models import DagBag, DagModel, DagRun from airflow.models.backfill import Backfill, BackfillDagRun, ReprocessBehavior, _create_backfill from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.python import PythonOperator -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py index dbe55a92e7893..6bde1fc211c6d 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py @@ -25,6 +25,7 @@ import time_machine from sqlalchemy import select +from airflow._shared.timezones import timezone from airflow.api_fastapi.core_api.datamodels.dag_versions import DagVersionResponse from airflow.listeners.listener import get_listener_manager from airflow.models import DagModel, DagRun @@ -32,7 +33,6 @@ from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk.definitions.asset import Asset from airflow.sdk.definitions.param import Param -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState, State from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_stats.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_stats.py index 69e9a65ac2451..f1702a220dba1 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_stats.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_stats.py @@ -20,9 +20,9 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models.dag import DagModel from airflow.models.dagrun import DagRun -from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py index 4d84485d48ab4..780330c054596 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py @@ -20,13 +20,13 @@ import pytest +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.dagbag import dag_bag_from_app from airflow.api_fastapi.core_api.datamodels.extra_links import ExtraLinkCollectionResponse from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.models.dagbag import DagBag from airflow.models.xcom import XComModel as XCom from airflow.plugins_manager import AirflowPlugin -from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py index 695dfcae19015..0c50be6223de2 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_hitl.py @@ -21,8 +21,8 @@ import pytest from sqlalchemy.orm import Session +from airflow._shared.timezones.timezone import utcnow from airflow.utils.state import TaskInstanceState -from airflow.utils.timezone import utcnow from tests_common.test_utils.db import AIRFLOW_V_3_1_PLUS diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_log.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_log.py index adb5aebcb2106..3a4e7c8ae7128 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_log.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_log.py @@ -28,12 +28,12 @@ from itsdangerous.url_safe import URLSafeSerializer from uuid6 import uuid7 +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.dagbag import create_dag_bag, dag_bag_from_app from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import task -from airflow.utils import timezone from airflow.utils.types import DagRunType from tests_common.test_utils.db import clear_db_runs diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_monitor.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_monitor.py index a3cf5012e495a..5175278d7771d 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_monitor.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_monitor.py @@ -21,9 +21,9 @@ import pytest +from airflow._shared.timezones import timezone from airflow.jobs.job import Job from airflow.jobs.scheduler_job_runner import SchedulerJobRunner -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py index b2e40ef5d5f82..3a0ba67a50a82 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py @@ -28,6 +28,7 @@ import pytest from sqlalchemy import select +from airflow._shared.timezones.timezone import datetime from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.jobs.job import Job from airflow.jobs.triggerer_job_runner import TriggererJobRunner @@ -42,7 +43,6 @@ from airflow.models.trigger import Trigger from airflow.utils.platform import getuser from airflow.utils.state import DagRunState, State, TaskInstanceState -from airflow.utils.timezone import datetime from airflow.utils.types import DagRunType from tests_common.test_utils.api_fastapi import _check_task_instance_note diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py index 5594f89f6f014..1c4194c178bd7 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py @@ -22,6 +22,7 @@ import pytest from airflow import DAG +from airflow._shared.timezones import timezone from airflow.api_fastapi.core_api.datamodels.xcom import XComCreateBody from airflow.models.dag_version import DagVersion from airflow.models.dagrun import DagRun @@ -31,7 +32,6 @@ from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.execution_time.xcom import resolve_xcom_backend -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py index dd6a0adb7b578..7d1f06a03f36b 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py @@ -20,9 +20,9 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models import DagModel from airflow.models.backfill import Backfill -from airflow.utils import timezone from airflow.utils.session import provide_session from tests_common.test_utils.db import ( diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py index 9db616f4dbde2..96808bced3865 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py @@ -24,13 +24,13 @@ import pytest from sqlalchemy import select +from airflow._shared.timezones import timezone from airflow.models import DagBag from airflow.models.dag import DagModel from airflow.models.taskinstance import TaskInstance from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import task_group from airflow.sdk.definitions.taskgroup import TaskGroup -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py index 1587797256a7a..aa8bb0a842b4e 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py @@ -24,6 +24,7 @@ from sqlalchemy import select from sqlalchemy.orm import Session +from airflow._shared.timezones import timezone from airflow.models import DagBag from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel from airflow.providers.standard.operators.empty import EmptyOperator @@ -31,7 +32,6 @@ from airflow.providers.standard.sensors.external_task import ExternalTaskSensor from airflow.sdk import Metadata, task from airflow.sdk.definitions.asset import Asset, AssetAlias, Dataset -from airflow.utils import timezone from tests_common.test_utils.db import clear_db_assets, clear_db_runs diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py index ae8488e5da977..4bc313d85d96c 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py @@ -19,8 +19,8 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, AssetModel -from airflow.utils import timezone DEFAULT_DATE = timezone.parse("2021-01-01T00:00:00") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_assets.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_assets.py index 2cf34f8dd7bc7..81ff615726208 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_assets.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_assets.py @@ -19,8 +19,8 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models.asset import AssetActive, AssetModel -from airflow.utils import timezone DEFAULT_DATE = timezone.parse("2021-01-01T00:00:00") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py index ac414f53ee83a..bf6a5ba581da6 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py @@ -19,10 +19,10 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models import DagModel from airflow.models.dagrun import DagRun from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone from airflow.utils.state import DagRunState, State from tests_common.test_utils.db import clear_db_runs diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index d18b18b673e01..94a5e96ab34b5 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -26,6 +26,7 @@ from sqlalchemy import select, update from sqlalchemy.exc import SQLAlchemyError +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import JWTValidator from airflow.api_fastapi.execution_api.app import lifespan from airflow.models import RenderedTaskInstanceFields, TaskReschedule, Trigger @@ -34,7 +35,6 @@ from airflow.models.taskinstancehistory import TaskInstanceHistory from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import Asset, TaskGroup, task, task_group -from airflow.utils import timezone from airflow.utils.state import DagRunState, State, TaskInstanceState, TerminalTIState from tests_common.test_utils.db import ( diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py index 5757a3f6465cd..661fd93157e95 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py @@ -21,9 +21,9 @@ import pytest +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.dagbag import dag_bag_from_app from airflow.jobs.scheduler_job_runner import SchedulerDagBag -from airflow.utils import timezone from airflow.utils.state import State from tests_common.test_utils.db import clear_db_assets, clear_db_runs diff --git a/airflow-core/tests/unit/callbacks/test_callback_requests.py b/airflow-core/tests/unit/callbacks/test_callback_requests.py index d04145ff6d580..08c1a9a666d43 100644 --- a/airflow-core/tests/unit/callbacks/test_callback_requests.py +++ b/airflow-core/tests/unit/callbacks/test_callback_requests.py @@ -21,6 +21,7 @@ import pytest +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import ( DagCallbackRequest, TaskCallbackRequest, @@ -28,7 +29,6 @@ from airflow.models.dag import DAG from airflow.models.taskinstance import TaskInstance from airflow.providers.standard.operators.bash import BashOperator -from airflow.utils import timezone from airflow.utils.state import State, TaskInstanceState pytestmark = pytest.mark.db_test diff --git a/airflow-core/tests/unit/cli/commands/test_backfill_command.py b/airflow-core/tests/unit/cli/commands/test_backfill_command.py index 1c87d8c2ef87d..13da9fe66ac9b 100644 --- a/airflow-core/tests/unit/cli/commands/test_backfill_command.py +++ b/airflow-core/tests/unit/cli/commands/test_backfill_command.py @@ -26,9 +26,9 @@ import pytest import airflow.cli.commands.backfill_command +from airflow._shared.timezones import timezone from airflow.cli import cli_parser from airflow.models.backfill import ReprocessBehavior -from airflow.utils import timezone from tests_common.test_utils.db import clear_db_backfills, clear_db_dags, clear_db_runs, parse_and_sync_to_db diff --git a/airflow-core/tests/unit/cli/commands/test_dag_command.py b/airflow-core/tests/unit/cli/commands/test_dag_command.py index 219785fc7d053..630745a702275 100644 --- a/airflow-core/tests/unit/cli/commands/test_dag_command.py +++ b/airflow-core/tests/unit/cli/commands/test_dag_command.py @@ -33,6 +33,7 @@ from sqlalchemy import select from airflow import settings +from airflow._shared.timezones import timezone from airflow.cli import cli_parser from airflow.cli.commands import dag_command from airflow.exceptions import AirflowException @@ -43,7 +44,6 @@ from airflow.sdk import task from airflow.sdk.definitions.dag import _run_inline_trigger from airflow.triggers.base import TriggerEvent -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType @@ -646,10 +646,9 @@ def test_dag_test_fail_raise_error(self, mock_get_dag): dag_command.dag_test(cli_args) @mock.patch("airflow.cli.commands.dag_command.get_dag") - @mock.patch("airflow.utils.timezone.utcnow") - def test_dag_test_no_logical_date(self, mock_utcnow, mock_get_dag): + def test_dag_test_no_logical_date(self, mock_get_dag, time_machine): now = pendulum.now() - mock_utcnow.return_value = now + time_machine.move_to(now, tick=False) cli_args = self.parser.parse_args(["dags", "test", "example_bash_operator"]) assert cli_args.logical_date is None diff --git a/airflow-core/tests/unit/cli/commands/test_task_command.py b/airflow-core/tests/unit/cli/commands/test_task_command.py index 6e25253eb4067..0b56a4a627c70 100644 --- a/airflow-core/tests/unit/cli/commands/test_task_command.py +++ b/airflow-core/tests/unit/cli/commands/test_task_command.py @@ -32,6 +32,7 @@ import pytest +from airflow._shared.timezones import timezone from airflow.cli import cli_parser from airflow.cli.commands import task_command from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG @@ -42,7 +43,6 @@ from airflow.models.serialized_dag import SerializedDagModel from airflow.providers.standard.operators.bash import BashOperator from airflow.serialization.serialized_objects import SerializedDAG -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/tests/unit/core/test_core.py b/airflow-core/tests/unit/core/test_core.py index cfdb9dee32357..f78346616d196 100644 --- a/airflow-core/tests/unit/core/test_core.py +++ b/airflow-core/tests/unit/core/test_core.py @@ -23,11 +23,11 @@ import pytest +from airflow._shared.timezones.timezone import datetime from airflow.exceptions import AirflowTaskTimeout from airflow.models.baseoperator import BaseOperator from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.python import PythonOperator -from airflow.utils.timezone import datetime from airflow.utils.types import DagRunType from tests_common.test_utils.db import clear_db_dags, clear_db_runs diff --git a/airflow-core/tests/unit/core/test_impersonation_tests.py b/airflow-core/tests/unit/core/test_impersonation_tests.py index c7f06533b8ba7..62c8effb5b275 100644 --- a/airflow-core/tests/unit/core/test_impersonation_tests.py +++ b/airflow-core/tests/unit/core/test_impersonation_tests.py @@ -27,11 +27,11 @@ import pytest +from airflow._shared.timezones.timezone import datetime from airflow.configuration import conf from airflow.models import DagBag, TaskInstance from airflow.utils.db import add_default_pool_if_not_exists from airflow.utils.state import State -from airflow.utils.timezone import datetime from tests_common.test_utils import db diff --git a/airflow-core/tests/unit/core/test_sentry.py b/airflow-core/tests/unit/core/test_sentry.py index 33305c64a3940..723a9dc0ce510 100644 --- a/airflow-core/tests/unit/core/test_sentry.py +++ b/airflow-core/tests/unit/core/test_sentry.py @@ -26,8 +26,8 @@ from sentry_sdk import configure_scope from sentry_sdk.transport import Transport +from airflow._shared.timezones import timezone from airflow.providers.standard.operators.python import PythonOperator -from airflow.utils import timezone from airflow.utils.module_loading import import_string from airflow.utils.state import State diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_base.py b/airflow-core/tests/unit/dag_processing/bundles/test_base.py index 14986486c88cb..7da3baa3b9f53 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_base.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_base.py @@ -29,13 +29,13 @@ import pytest import time_machine +from airflow._shared.timezones import timezone as tz from airflow.dag_processing.bundles.base import ( BaseDagBundle, BundleUsageTrackingManager, BundleVersionLock, get_bundle_storage_root_path, ) -from airflow.utils import timezone as tz from tests_common.test_utils.config import conf_vars diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py b/airflow-core/tests/unit/dag_processing/test_collection.py index 5254dfc6ec973..0c812780ce410 100644 --- a/airflow-core/tests/unit/dag_processing/test_collection.py +++ b/airflow-core/tests/unit/dag_processing/test_collection.py @@ -31,6 +31,7 @@ from sqlalchemy.exc import OperationalError, SAWarning import airflow.dag_processing.collection +from airflow._shared.timezones import timezone as tz from airflow.configuration import conf from airflow.dag_processing.collection import ( AssetModelOperation, @@ -54,7 +55,6 @@ from airflow.providers.standard.triggers.temporal import TimeDeltaTrigger from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetWatcher from airflow.serialization.serialized_objects import LazyDeserializedDAG, SerializedDAG -from airflow.utils import timezone as tz from tests_common.test_utils.db import ( clear_db_assets, diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 02121f0194e82..c218c54eae768 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -40,6 +40,7 @@ from sqlalchemy import func, select from uuid6 import uuid7 +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG from airflow.dag_processing.bundles.manager import DagBundlesManager @@ -55,7 +56,6 @@ from airflow.models.dagbundle import DagBundleModel from airflow.models.dagcode import DagCode from airflow.models.serialized_dag import SerializedDagModel -from airflow.utils import timezone from airflow.utils.net import get_hostname from airflow.utils.session import create_session diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 96a5db7281bef..6c6f31a1ff8d0 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -32,6 +32,7 @@ from pydantic import TypeAdapter from structlog.typing import FilteringBoundLogger +from airflow._shared.timezones import timezone from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( TaskInstance as TIDataModel, @@ -52,7 +53,6 @@ from airflow.sdk.api.client import Client from airflow.sdk.api.datamodels._generated import DagRunState from airflow.sdk.execution_time import comms -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import TaskInstanceState diff --git a/airflow-core/tests/unit/dags/test_cli_triggered_dags.py b/airflow-core/tests/unit/dags/test_cli_triggered_dags.py index 378e2497f2818..7d4af553a7b1a 100644 --- a/airflow-core/tests/unit/dags/test_cli_triggered_dags.py +++ b/airflow-core/tests/unit/dags/test_cli_triggered_dags.py @@ -19,9 +19,9 @@ from datetime import timedelta +from airflow._shared.timezones.timezone import datetime from airflow.models.dag import DAG from airflow.providers.standard.operators.python import PythonOperator -from airflow.utils.timezone import datetime DEFAULT_DATE = datetime(2016, 1, 1) default_args = dict(start_date=DEFAULT_DATE, owner="airflow") diff --git a/airflow-core/tests/unit/dags/test_invalid_cron.py b/airflow-core/tests/unit/dags/test_invalid_cron.py index e4d30f7888655..fe422692d1f8d 100644 --- a/airflow-core/tests/unit/dags/test_invalid_cron.py +++ b/airflow-core/tests/unit/dags/test_invalid_cron.py @@ -17,9 +17,9 @@ # under the License. from __future__ import annotations +from airflow._shared.timezones.timezone import datetime from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils.timezone import datetime # This invalid DAG has a schedule specified with an INVALID cron expression. # It will be used to test whether dagbag.process_file() can identify this. diff --git a/airflow-core/tests/unit/dags/test_logging_in_dag.py b/airflow-core/tests/unit/dags/test_logging_in_dag.py index 484181688c96b..2d4ba261af630 100644 --- a/airflow-core/tests/unit/dags/test_logging_in_dag.py +++ b/airflow-core/tests/unit/dags/test_logging_in_dag.py @@ -19,9 +19,9 @@ import logging +from airflow._shared.timezones.timezone import datetime from airflow.models.dag import DAG from airflow.providers.standard.operators.python import PythonOperator -from airflow.utils.timezone import datetime logger = logging.getLogger(__name__) diff --git a/airflow-core/tests/unit/dags/test_mark_state.py b/airflow-core/tests/unit/dags/test_mark_state.py index 1990c8c03c6ae..47674815aec11 100644 --- a/airflow-core/tests/unit/dags/test_mark_state.py +++ b/airflow-core/tests/unit/dags/test_mark_state.py @@ -21,11 +21,11 @@ from datetime import datetime from time import sleep +from airflow._shared.timezones.timezone import utcnow from airflow.models.dag import DAG from airflow.providers.standard.operators.python import PythonOperator from airflow.utils.session import create_session from airflow.utils.state import State -from airflow.utils.timezone import utcnow DEFAULT_DATE = datetime(2016, 1, 1) diff --git a/airflow-core/tests/unit/dags/test_on_kill.py b/airflow-core/tests/unit/dags/test_on_kill.py index 1e73b76776f5b..3de4fb6a62718 100644 --- a/airflow-core/tests/unit/dags/test_on_kill.py +++ b/airflow-core/tests/unit/dags/test_on_kill.py @@ -20,9 +20,9 @@ import time from typing import TYPE_CHECKING +from airflow._shared.timezones.timezone import datetime from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils.timezone import datetime if TYPE_CHECKING: from airflow.sdk.definitions.context import Context diff --git a/airflow-core/tests/unit/dags/test_parsing_context.py b/airflow-core/tests/unit/dags/test_parsing_context.py index 0a0ab4a4a7d0e..676947ae15747 100644 --- a/airflow-core/tests/unit/dags/test_parsing_context.py +++ b/airflow-core/tests/unit/dags/test_parsing_context.py @@ -19,6 +19,7 @@ from pathlib import Path +from airflow._shared.timezones.timezone import datetime from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk.definitions.context import ( @@ -26,7 +27,6 @@ _AIRFLOW_PARSING_CONTEXT_TASK_ID, Context, ) -from airflow.utils.timezone import datetime class DagWithParsingContext(EmptyOperator): diff --git a/airflow-core/tests/unit/dags/test_scheduler_dags.py b/airflow-core/tests/unit/dags/test_scheduler_dags.py index da89e821986f1..9123d32552f5e 100644 --- a/airflow-core/tests/unit/dags/test_scheduler_dags.py +++ b/airflow-core/tests/unit/dags/test_scheduler_dags.py @@ -19,9 +19,9 @@ from datetime import timedelta +from airflow._shared.timezones import timezone from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone DEFAULT_DATE = timezone.datetime(2016, 1, 1) diff --git a/airflow-core/tests/unit/dags/test_sensor.py b/airflow-core/tests/unit/dags/test_sensor.py index 4feb1caa1eb82..61504cfb77549 100644 --- a/airflow-core/tests/unit/dags/test_sensor.py +++ b/airflow-core/tests/unit/dags/test_sensor.py @@ -18,9 +18,9 @@ import datetime +from airflow._shared.timezones import timezone from airflow.models.dag import DAG from airflow.sdk import task -from airflow.utils import timezone from tests_common.test_utils.compat import DateTimeSensor diff --git a/airflow-core/tests/unit/decorators/test_mapped.py b/airflow-core/tests/unit/decorators/test_mapped.py index f1940f6458d85..1e148301cf694 100644 --- a/airflow-core/tests/unit/decorators/test_mapped.py +++ b/airflow-core/tests/unit/decorators/test_mapped.py @@ -19,8 +19,8 @@ import pytest +from airflow._shared.timezones import timezone from airflow.sdk import DAG, task -from airflow.utils import timezone DEFAULT_DATE = timezone.datetime(2025, 1, 1) diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 2ed0657e3c6a5..9715138c74f8a 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -26,6 +26,7 @@ import pytest import time_machine +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import CallbackRequest from airflow.cli.cli_config import DefaultHelpParser, GroupCommand from airflow.cli.cli_parser import AirflowHelpFormatter @@ -34,7 +35,6 @@ from airflow.executors.local_executor import LocalExecutor from airflow.models.baseoperator import BaseOperator from airflow.models.taskinstance import TaskInstance, TaskInstanceKey -from airflow.utils import timezone from airflow.utils.state import State, TaskInstanceState from tests_common.test_utils.markers import skip_if_force_lowest_dependencies_marker diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index e5b37bb268d3c..ec7102e2e70c2 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -25,9 +25,9 @@ from kgb import spy_on from uuid6 import uuid7 +from airflow._shared.timezones import timezone from airflow.executors import workloads from airflow.executors.local_executor import LocalExecutor, _execute_work -from airflow.utils import timezone from airflow.utils.state import State from tests_common.test_utils.config import conf_vars diff --git a/airflow-core/tests/unit/jobs/test_base_job.py b/airflow-core/tests/unit/jobs/test_base_job.py index c4b46a7622e53..9f80073007cef 100644 --- a/airflow-core/tests/unit/jobs/test_base_job.py +++ b/airflow-core/tests/unit/jobs/test_base_job.py @@ -25,10 +25,10 @@ import pytest from sqlalchemy.exc import OperationalError +from airflow._shared.timezones import timezone from airflow.executors.local_executor import LocalExecutor from airflow.jobs.job import Job, most_recent_job, perform_heartbeat, run_job from airflow.listeners.listener import get_listener_manager -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index f9269226e9cdc..ec1580528f9b7 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -38,6 +38,7 @@ from sqlalchemy.orm import joinedload from airflow import settings +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import JWTGenerator from airflow.assets.manager import AssetManager from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest @@ -74,7 +75,6 @@ from airflow.serialization.serialized_objects import LazyDeserializedDAG, SerializedDAG from airflow.timetables.base import DataInterval from airflow.traces.tracer import Trace -from airflow.utils import timezone from airflow.utils.session import create_session, provide_session from airflow.utils.span_status import SpanStatus from airflow.utils.state import DagRunState, State, TaskInstanceState @@ -101,7 +101,6 @@ from unit.listeners import dag_listener from unit.listeners.test_listeners import get_listener_manager from unit.models import TEST_DAGS_FOLDER -from unit.utils.test_timezone import UTC pytestmark = pytest.mark.db_test @@ -4676,7 +4675,7 @@ def complete_one_dagrun(): "test_dag": 3 } - assert model.next_dagrun == timezone.DateTime(2016, 1, 3, tzinfo=UTC) + assert model.next_dagrun == timezone.DateTime(2016, 1, 3, tzinfo=timezone.utc) assert model.next_dagrun_create_after is None complete_one_dagrun() diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index 194c1c8752ba7..384bf6a76fc60 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -32,6 +32,7 @@ from asgiref.sync import sync_to_async from structlog.typing import FilteringBoundLogger +from airflow._shared.timezones import timezone from airflow.executors import workloads from airflow.jobs.job import Job from airflow.jobs.triggerer_job_runner import ( @@ -55,7 +56,6 @@ from airflow.sdk import BaseHook from airflow.triggers.base import BaseTrigger, TriggerEvent from airflow.triggers.testing import FailureTrigger, SuccessTrigger -from airflow.utils import timezone from airflow.utils.state import State, TaskInstanceState from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/listeners/test_listeners.py b/airflow-core/tests/unit/listeners/test_listeners.py index 9a4c85cb504a1..fac237c7e176d 100644 --- a/airflow-core/tests/unit/listeners/test_listeners.py +++ b/airflow-core/tests/unit/listeners/test_listeners.py @@ -22,11 +22,11 @@ import pytest +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException from airflow.jobs.job import Job, run_job from airflow.listeners.listener import get_listener_manager from airflow.providers.standard.operators.bash import BashOperator -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState, TaskInstanceState diff --git a/airflow-core/tests/unit/models/__init__.py b/airflow-core/tests/unit/models/__init__.py index ded9a4f45507e..d5fe7d8f9301b 100644 --- a/airflow-core/tests/unit/models/__init__.py +++ b/airflow-core/tests/unit/models/__init__.py @@ -19,7 +19,7 @@ import pathlib -from airflow.utils import timezone +from airflow._shared.timezones import timezone DEFAULT_DATE = timezone.datetime(2016, 1, 1) TEST_DAGS_FOLDER = pathlib.Path(__file__).parent.with_name("dags") diff --git a/airflow-core/tests/unit/models/test_backfill.py b/airflow-core/tests/unit/models/test_backfill.py index 109d214594118..9d26e31459e3b 100644 --- a/airflow-core/tests/unit/models/test_backfill.py +++ b/airflow-core/tests/unit/models/test_backfill.py @@ -25,6 +25,7 @@ import pytest from sqlalchemy import select +from airflow._shared.timezones import timezone from airflow.models import DagModel, DagRun, TaskInstance from airflow.models.backfill import ( AlreadyRunningBackfill, @@ -38,7 +39,6 @@ ) from airflow.providers.standard.operators.python import PythonOperator from airflow.ti_deps.dep_context import DepContext -from airflow.utils import timezone from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py index dab3d4d991641..7f2f4b01394f5 100644 --- a/airflow-core/tests/unit/models/test_dag.py +++ b/airflow-core/tests/unit/models/test_dag.py @@ -34,6 +34,8 @@ from sqlalchemy import inspect, select from airflow import settings +from airflow._shared.timezones import timezone +from airflow._shared.timezones.timezone import datetime as datetime_tz from airflow.configuration import conf from airflow.exceptions import ( AirflowException, @@ -76,11 +78,9 @@ NullTimetable, OnceTimetable, ) -from airflow.utils import timezone from airflow.utils.file import list_py_file_paths from airflow.utils.session import create_session from airflow.utils.state import DagRunState, State, TaskInstanceState -from airflow.utils.timezone import datetime as datetime_tz from airflow.utils.trigger_rule import TriggerRule from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/tests/unit/models/test_dagbag.py b/airflow-core/tests/unit/models/test_dagbag.py index 00aa6151f9728..230ddca9f89e2 100644 --- a/airflow-core/tests/unit/models/test_dagbag.py +++ b/airflow-core/tests/unit/models/test_dagbag.py @@ -35,12 +35,12 @@ from sqlalchemy import select from airflow import settings +from airflow._shared.timezones import timezone as tz from airflow.models.dag import DAG, DagModel from airflow.models.dagbag import DagBag, _capture_with_reraise from airflow.models.dagwarning import DagWarning, DagWarningType from airflow.models.serialized_dag import SerializedDagModel from airflow.serialization.serialized_objects import SerializedDAG -from airflow.utils import timezone as tz from airflow.utils.session import create_session from tests_common.pytest_plugin import AIRFLOW_ROOT_PATH diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 50440d4f44ef2..cf4dafb4ee094 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -31,6 +31,7 @@ from sqlalchemy.orm import joinedload from airflow import settings +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.models.dag import DAG, DagModel from airflow.models.dag_version import DagVersion @@ -47,7 +48,6 @@ from airflow.serialization.serialized_objects import SerializedDAG from airflow.stats import Stats from airflow.triggers.base import StartTriggerArgs -from airflow.utils import timezone from airflow.utils.span_status import SpanStatus from airflow.utils.state import DagRunState, State, TaskInstanceState from airflow.utils.thread_safe_dict import ThreadSafeDict diff --git a/airflow-core/tests/unit/models/test_pool.py b/airflow-core/tests/unit/models/test_pool.py index ad0c935d618ca..81d4ce66dc717 100644 --- a/airflow-core/tests/unit/models/test_pool.py +++ b/airflow-core/tests/unit/models/test_pool.py @@ -20,12 +20,12 @@ import pytest from airflow import settings +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException, PoolNotFound from airflow.models.dag_version import DagVersion from airflow.models.pool import Pool from airflow.models.taskinstance import TaskInstance as TI from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/models/test_renderedtifields.py b/airflow-core/tests/unit/models/test_renderedtifields.py index e53582c563e0f..8acf7fe932f3a 100644 --- a/airflow-core/tests/unit/models/test_renderedtifields.py +++ b/airflow-core/tests/unit/models/test_renderedtifields.py @@ -29,6 +29,7 @@ from sqlalchemy import select from airflow import settings +from airflow._shared.timezones.timezone import datetime from airflow.configuration import conf from airflow.models import DagRun, Variable from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF @@ -37,7 +38,6 @@ from airflow.providers.standard.operators.python import PythonOperator from airflow.sdk import task as task_decorator from airflow.utils.task_instance_session import set_current_task_instance_session -from airflow.utils.timezone import datetime from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.db import clear_db_dags, clear_db_runs, clear_rendered_ti_fields diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 7e041917056ed..2a9e86771c2f7 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -34,6 +34,7 @@ from sqlalchemy.exc import IntegrityError from airflow import settings +from airflow._shared.timezones import timezone from airflow.exceptions import ( AirflowException, AirflowFailException, @@ -79,7 +80,6 @@ from airflow.ti_deps.deps.base_ti_dep import TIDepStatus from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep, _UpstreamTIStates -from airflow.utils import timezone from airflow.utils.db import merge_conn from airflow.utils.session import create_session, provide_session from airflow.utils.span_status import SpanStatus diff --git a/airflow-core/tests/unit/models/test_timestamp.py b/airflow-core/tests/unit/models/test_timestamp.py index 5aaa956829d72..529fd32afc504 100644 --- a/airflow-core/tests/unit/models/test_timestamp.py +++ b/airflow-core/tests/unit/models/test_timestamp.py @@ -20,9 +20,9 @@ import pytest import time_machine +from airflow._shared.timezones import timezone from airflow.models import Log from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/models/test_trigger.py b/airflow-core/tests/unit/models/test_trigger.py index 8818036418e2c..3f5566e3af0d3 100644 --- a/airflow-core/tests/unit/models/test_trigger.py +++ b/airflow-core/tests/unit/models/test_trigger.py @@ -27,6 +27,7 @@ import pytz from cryptography.fernet import Fernet +from airflow._shared.timezones import timezone from airflow.jobs.job import Job from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.models import Deadline, TaskInstance, Trigger @@ -41,7 +42,6 @@ TaskSuccessEvent, TriggerEvent, ) -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State @@ -239,7 +239,7 @@ def test_submit_failure(session, create_task_instance): (TaskSkippedEvent, "skipped"), ], ) -@patch("airflow.utils.timezone.utcnow") +@patch("airflow._shared.timezones.timezone.utcnow") def test_submit_event_task_end(mock_utcnow, session, create_task_instance, event_cls, expected): """ Tests that events inheriting BaseTaskEndEvent *don't* re-wake their dependent diff --git a/airflow-core/tests/unit/models/test_xcom.py b/airflow-core/tests/unit/models/test_xcom.py index 566d74ceaa774..6ed5d95e83d69 100644 --- a/airflow-core/tests/unit/models/test_xcom.py +++ b/airflow-core/tests/unit/models/test_xcom.py @@ -25,6 +25,7 @@ import pytest from airflow import DAG +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.models.dag_version import DagVersion from airflow.models.dagrun import DagRun, DagRunType @@ -35,7 +36,6 @@ from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.execution_time.xcom import resolve_xcom_backend from airflow.settings import json -from airflow.utils import timezone from airflow.utils.session import create_session from tests_common.test_utils.config import conf_vars diff --git a/airflow-core/tests/unit/sensors/test_filesystem.py b/airflow-core/tests/unit/sensors/test_filesystem.py index 30afb69db2bea..6d4f1d82bb4e0 100644 --- a/airflow-core/tests/unit/sensors/test_filesystem.py +++ b/airflow-core/tests/unit/sensors/test_filesystem.py @@ -24,11 +24,11 @@ import pytest +from airflow._shared.timezones.timezone import datetime from airflow.exceptions import AirflowSensorTimeout, TaskDeferred from airflow.models.dag import DAG from airflow.providers.standard.sensors.filesystem import FileSensor from airflow.providers.standard.triggers.file import FileTrigger -from airflow.utils.timezone import datetime pytestmark = pytest.mark.db_test diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index c61284c785026..234bb51fbb22e 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -47,6 +47,7 @@ from kubernetes.client import models as k8s import airflow +from airflow._shared.timezones import timezone from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG from airflow.exceptions import ( AirflowException, @@ -82,7 +83,6 @@ from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep from airflow.timetables.simple import NullTimetable, OnceTimetable from airflow.triggers.base import StartTriggerArgs -from airflow.utils import timezone from airflow.utils.module_loading import qualname from airflow.utils.operator_resources import Resources diff --git a/airflow-core/tests/unit/serialization/test_serialized_objects.py b/airflow-core/tests/unit/serialization/test_serialized_objects.py index dc944035c75a9..dec42d3568bf4 100644 --- a/airflow-core/tests/unit/serialization/test_serialized_objects.py +++ b/airflow-core/tests/unit/serialization/test_serialized_objects.py @@ -28,6 +28,7 @@ from pendulum.tz.timezone import FixedTimezone, Timezone from uuid6 import uuid7 +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.exceptions import ( AirflowException, @@ -65,7 +66,6 @@ from airflow.serialization.serialized_objects import BaseSerialization, LazyDeserializedDAG, SerializedDAG from airflow.timetables.base import DataInterval from airflow.triggers.base import BaseTrigger -from airflow.utils import timezone from airflow.utils.db import LazySelectSequence from airflow.utils.operator_resources import Resources from airflow.utils.state import DagRunState, State diff --git a/airflow-core/tests/unit/ti_deps/deps/test_not_in_retry_period_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_not_in_retry_period_dep.py index 3b8aec4763ecb..bde50ad31f87f 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_not_in_retry_period_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_not_in_retry_period_dep.py @@ -24,10 +24,10 @@ import pytest import time_machine +from airflow._shared.timezones.timezone import datetime from airflow.models import TaskInstance from airflow.ti_deps.deps.not_in_retry_period_dep import NotInRetryPeriodDep from airflow.utils.state import State -from airflow.utils.timezone import datetime pytestmark = pytest.mark.db_test diff --git a/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py index d2fd53da745a9..5ffc2f2be430b 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_prev_dagrun_dep.py @@ -22,13 +22,13 @@ import pytest +from airflow._shared.timezones.timezone import convert_to_utc, datetime from airflow.models.baseoperator import BaseOperator from airflow.models.dag import DAG from airflow.models.serialized_dag import SerializedDagModel from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.deps.prev_dagrun_dep import PrevDagrunDep from airflow.utils.state import DagRunState, TaskInstanceState -from airflow.utils.timezone import convert_to_utc, datetime from airflow.utils.types import DagRunTriggeredByType, DagRunType from tests_common.test_utils.db import clear_db_runs diff --git a/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py index 7b4bd22aa0436..355cbfb9c30ef 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py @@ -24,10 +24,10 @@ import time_machine from slugify import slugify +from airflow._shared.timezones import timezone from airflow.models.taskreschedule import TaskReschedule from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/ti_deps/deps/test_runnable_exec_date_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_runnable_exec_date_dep.py index 94b9cfc2ac6cc..7b498b71f835b 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_runnable_exec_date_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_runnable_exec_date_dep.py @@ -22,9 +22,9 @@ import pytest import time_machine +from airflow._shared.timezones.timezone import datetime from airflow.models import DagRun, TaskInstance from airflow.ti_deps.deps.runnable_exec_date_dep import RunnableExecDateDep -from airflow.utils.timezone import datetime from airflow.utils.types import DagRunType pytestmark = pytest.mark.db_test diff --git a/airflow-core/tests/unit/timetables/test_events_timetable.py b/airflow-core/tests/unit/timetables/test_events_timetable.py index 13a27daa199c0..556f5d725b4d5 100644 --- a/airflow-core/tests/unit/timetables/test_events_timetable.py +++ b/airflow-core/tests/unit/timetables/test_events_timetable.py @@ -21,9 +21,9 @@ import pytest import time_machine +from airflow._shared.timezones.timezone import utc from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction, Timetable from airflow.timetables.events import EventsTimetable -from airflow.utils.timezone import utc BEFORE_DATE = pendulum.DateTime(2021, 9, 4, tzinfo=utc) # Precedes all events START_DATE = pendulum.DateTime(2021, 9, 7, tzinfo=utc) diff --git a/airflow-core/tests/unit/timetables/test_interval_timetable.py b/airflow-core/tests/unit/timetables/test_interval_timetable.py index d6e2d56df5433..87a23ba24a717 100644 --- a/airflow-core/tests/unit/timetables/test_interval_timetable.py +++ b/airflow-core/tests/unit/timetables/test_interval_timetable.py @@ -24,10 +24,10 @@ import pytest import time_machine +from airflow._shared.timezones.timezone import utc from airflow.exceptions import AirflowTimetableInvalid from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction, Timetable from airflow.timetables.interval import CronDataIntervalTimetable, DeltaDataIntervalTimetable -from airflow.utils.timezone import utc START_DATE = pendulum.DateTime(2021, 9, 4, tzinfo=utc) diff --git a/airflow-core/tests/unit/timetables/test_once_timetable.py b/airflow-core/tests/unit/timetables/test_once_timetable.py index a904e8372a7fb..0ca84229ba63c 100644 --- a/airflow-core/tests/unit/timetables/test_once_timetable.py +++ b/airflow-core/tests/unit/timetables/test_once_timetable.py @@ -23,9 +23,9 @@ import pytest import time_machine +from airflow._shared.timezones import timezone from airflow.timetables.base import DagRunInfo, TimeRestriction from airflow.timetables.simple import OnceTimetable -from airflow.utils import timezone FROZEN_NOW = timezone.coerce_datetime(datetime.datetime(2025, 3, 4, 5, 6, 7, 8)) diff --git a/airflow-core/tests/unit/timetables/test_trigger_timetable.py b/airflow-core/tests/unit/timetables/test_trigger_timetable.py index 1a98044bf4ef0..3a0a60cebfea4 100644 --- a/airflow-core/tests/unit/timetables/test_trigger_timetable.py +++ b/airflow-core/tests/unit/timetables/test_trigger_timetable.py @@ -24,6 +24,7 @@ import pytest import time_machine +from airflow._shared.timezones.timezone import utc from airflow.exceptions import AirflowTimetableInvalid from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction, Timetable from airflow.timetables.trigger import ( @@ -31,7 +32,6 @@ DeltaTriggerTimetable, MultipleCronTriggerTimetable, ) -from airflow.utils.timezone import utc START_DATE = pendulum.DateTime(2021, 9, 4, tzinfo=utc) diff --git a/airflow-core/tests/unit/timetables/test_workday_timetable.py b/airflow-core/tests/unit/timetables/test_workday_timetable.py index d8ccf5809dbe2..7c91d16aeb4d7 100644 --- a/airflow-core/tests/unit/timetables/test_workday_timetable.py +++ b/airflow-core/tests/unit/timetables/test_workday_timetable.py @@ -22,9 +22,9 @@ import pendulum import pytest +from airflow._shared.timezones.timezone import utc from airflow.example_dags.plugins.workday import AfterWorkdayTimetable from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction, Timetable -from airflow.utils.timezone import utc START_DATE = pendulum.DateTime(2021, 9, 4, tzinfo=utc) # This is a Saturday. diff --git a/airflow-core/tests/unit/utils/log/test_file_processor_handler.py b/airflow-core/tests/unit/utils/log/test_file_processor_handler.py index f50fc619ac5bd..ad0982dbea433 100644 --- a/airflow-core/tests/unit/utils/log/test_file_processor_handler.py +++ b/airflow-core/tests/unit/utils/log/test_file_processor_handler.py @@ -23,7 +23,7 @@ import time_machine -from airflow.utils import timezone +from airflow._shared.timezones import timezone from airflow.utils.log.file_processor_handler import FileProcessorHandler diff --git a/airflow-core/tests/unit/utils/log/test_log_reader.py b/airflow-core/tests/unit/utils/log/test_log_reader.py index cd8b430090a27..8d5ffab543dfa 100644 --- a/airflow-core/tests/unit/utils/log/test_log_reader.py +++ b/airflow-core/tests/unit/utils/log/test_log_reader.py @@ -29,11 +29,11 @@ import pytest from airflow import settings +from airflow._shared.timezones import timezone from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG from airflow.models.tasklog import LogTemplate from airflow.providers.standard.operators.python import PythonOperator from airflow.timetables.base import DataInterval -from airflow.utils import timezone from airflow.utils.log.log_reader import TaskLogReader from airflow.utils.log.logging_mixin import ExternalLoggingMixin from airflow.utils.state import TaskInstanceState diff --git a/airflow-core/tests/unit/utils/test_cli_util.py b/airflow-core/tests/unit/utils/test_cli_util.py index 2d5578f8faa93..54ecc506e2731 100644 --- a/airflow-core/tests/unit/utils/test_cli_util.py +++ b/airflow-core/tests/unit/utils/test_cli_util.py @@ -30,9 +30,10 @@ import airflow from airflow import settings +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException from airflow.models.log import Log -from airflow.utils import cli, cli_action_loggers, timezone +from airflow.utils import cli, cli_action_loggers from airflow.utils.cli import _search_for_dag_file # Mark entire module as db_test because ``action_cli`` wrapper still could use DB on callbacks: diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py index 48a0c8486219b..71c481ea1c650 100644 --- a/airflow-core/tests/unit/utils/test_db_cleanup.py +++ b/airflow-core/tests/unit/utils/test_db_cleanup.py @@ -31,12 +31,12 @@ from sqlalchemy.ext.declarative import DeclarativeMeta from airflow import DAG +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException from airflow.models import DagModel, DagRun, TaskInstance from airflow.models.dag_version import DagVersion from airflow.models.serialized_dag import SerializedDagModel from airflow.providers.standard.operators.python import PythonOperator -from airflow.utils import timezone from airflow.utils.db_cleanup import ( ARCHIVE_TABLE_PREFIX, CreateTableAs, diff --git a/airflow-core/tests/unit/utils/test_dot_renderer.py b/airflow-core/tests/unit/utils/test_dot_renderer.py index 240876ec7f43e..1e239271dfce7 100644 --- a/airflow-core/tests/unit/utils/test_dot_renderer.py +++ b/airflow-core/tests/unit/utils/test_dot_renderer.py @@ -22,11 +22,12 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk.definitions.taskgroup import TaskGroup from airflow.serialization.dag_dependency import DagDependency -from airflow.utils import dot_renderer, timezone +from airflow.utils import dot_renderer from airflow.utils.state import State from tests_common.test_utils.compat import BashOperator diff --git a/airflow-core/tests/unit/utils/test_helpers.py b/airflow-core/tests/unit/utils/test_helpers.py index 71138aa51057f..08c8d8092a0e9 100644 --- a/airflow-core/tests/unit/utils/test_helpers.py +++ b/airflow-core/tests/unit/utils/test_helpers.py @@ -23,9 +23,10 @@ import pytest +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException from airflow.jobs.base_job_runner import BaseJobRunner -from airflow.utils import helpers, timezone +from airflow.utils import helpers from airflow.utils.helpers import ( at_most_one, build_airflow_dagrun_url, diff --git a/airflow-core/tests/unit/utils/test_serve_logs.py b/airflow-core/tests/unit/utils/test_serve_logs.py index fbbef2b200046..b6c7805386ac9 100644 --- a/airflow-core/tests/unit/utils/test_serve_logs.py +++ b/airflow-core/tests/unit/utils/test_serve_logs.py @@ -24,9 +24,9 @@ import time_machine from fastapi.testclient import TestClient +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import JWTGenerator from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG -from airflow.utils import timezone from airflow.utils.serve_logs import create_app from tests_common.test_utils.config import conf_vars diff --git a/airflow-core/tests/unit/utils/test_sqlalchemy.py b/airflow-core/tests/unit/utils/test_sqlalchemy.py index 602a46a8ab179..9ff19c7c0e614 100644 --- a/airflow-core/tests/unit/utils/test_sqlalchemy.py +++ b/airflow-core/tests/unit/utils/test_sqlalchemy.py @@ -29,6 +29,7 @@ from sqlalchemy.exc import StatementError from airflow import settings +from airflow._shared.timezones.timezone import utcnow from airflow.models.dag import DAG from airflow.models.serialized_dag import SerializedDagModel from airflow.serialization.enums import DagAttributeTypes, Encoding @@ -42,7 +43,6 @@ with_row_locks, ) from airflow.utils.state import State -from airflow.utils.timezone import utcnow from airflow.utils.types import DagRunTriggeredByType, DagRunType pytestmark = pytest.mark.db_test diff --git a/airflow-core/tests/unit/utils/test_task_handler_with_custom_formatter.py b/airflow-core/tests/unit/utils/test_task_handler_with_custom_formatter.py index 5e5519047980a..d3e24448f541e 100644 --- a/airflow-core/tests/unit/utils/test_task_handler_with_custom_formatter.py +++ b/airflow-core/tests/unit/utils/test_task_handler_with_custom_formatter.py @@ -21,13 +21,13 @@ import pytest +from airflow._shared.timezones.timezone import datetime from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG from airflow.models.dag_version import DagVersion from airflow.models.taskinstance import TaskInstance from airflow.providers.standard.operators.empty import EmptyOperator from airflow.utils.log.logging_mixin import set_context from airflow.utils.state import DagRunState -from airflow.utils.timezone import datetime from airflow.utils.types import DagRunType from tests_common.test_utils.config import conf_vars diff --git a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py index ed9e9dd7a2fd6..b959d53fc0469 100644 --- a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py +++ b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py @@ -83,25 +83,26 @@ ("RELEASE_NOTES.rst", "/opt/airflow/RELEASE_NOTES.rst"), ("airflow-core", "/opt/airflow/airflow-core"), ("airflow-ctl", "/opt/airflow/airflow-ctl"), - ("constraints", "/opt/airflow/constraints"), + ("chart", "/opt/airflow/chart"), ("clients", "/opt/airflow/clients"), + ("constraints", "/opt/airflow/constraints"), ("dags", "/opt/airflow/dags"), ("dev", "/opt/airflow/dev"), - ("docs", "/opt/airflow/docs"), + ("devel-common", "/opt/airflow/devel-common"), ("docker-stack-docs", "/opt/airflow/docker-stack-docs"), - ("providers-summary-docs", "/opt/airflow/providers-summary-docs"), + ("docker-tests", "/opt/airflow/docker-tests"), + ("docs", "/opt/airflow/docs"), ("generated", "/opt/airflow/generated"), + ("helm-tests", "/opt/airflow/helm-tests"), + ("kubernetes-tests", "/opt/airflow/kubernetes-tests"), ("logs", "/root/airflow/logs"), ("providers", "/opt/airflow/providers"), - ("task-sdk", "/opt/airflow/task-sdk"), + ("providers-summary-docs", "/opt/airflow/providers-summary-docs"), ("pyproject.toml", "/opt/airflow/pyproject.toml"), ("scripts", "/opt/airflow/scripts"), ("scripts/docker/entrypoint_ci.sh", "/entrypoint"), - ("devel-common", "/opt/airflow/devel-common"), - ("helm-tests", "/opt/airflow/helm-tests"), - ("kubernetes-tests", "/opt/airflow/kubernetes-tests"), - ("docker-tests", "/opt/airflow/docker-tests"), - ("chart", "/opt/airflow/chart"), + ("shared", "/opt/airflow/shared"), + ("task-sdk", "/opt/airflow/task-sdk"), ] diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index f02ab53fa6d3d..e10bcd1d5bde9 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -1012,10 +1012,11 @@ def _make_serdag(self, dag): self.session.flush() def create_dagrun(self, *, logical_date=NOTSET, **kwargs): - from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType + timezone = _import_timezone() + if AIRFLOW_V_3_0_PLUS: from airflow.utils.types import DagRunTriggeredByType else: @@ -1177,7 +1178,8 @@ def __call__( ): from airflow import settings from airflow.models.dag import DAG - from airflow.utils import timezone + + timezone = _import_timezone() if session is None: self._own_session = True @@ -1445,14 +1447,11 @@ def maker( last_heartbeat_at=None, **kwargs, ) -> TaskInstance: + timezone = _import_timezone() if run_after is None: - from airflow.utils import timezone - run_after = timezone.utcnow() if logical_date is NOTSET: # For now: default to having a logical date if None is not explicitly passed. - from airflow.utils import timezone - logical_date = timezone.utcnow() with dag_maker(dag_id, **kwargs): op_kwargs = {} @@ -1609,7 +1608,8 @@ def _get(dag_id: str): if dagbag.import_errors: session = settings.Session() from airflow.models.errors import ParseImportError - from airflow.utils import timezone + + timezone = _import_timezone() # Add the new import errors for _filename, stacktrace in dagbag.import_errors.items(): @@ -2100,7 +2100,8 @@ def mocked_parse(spy_agency): def set_dag(what: StartupDetails, dag_id: str, task: TaskSDKBaseOperator) -> RuntimeTaskInstance: from airflow.sdk.definitions.dag import DAG from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance, parse - from airflow.utils import timezone + + timezone = _import_timezone() if not task.has_dag(): dag = DAG(dag_id=dag_id, start_date=timezone.datetime(2024, 12, 3)) @@ -2215,7 +2216,8 @@ def execute(self, context): from airflow.sdk.definitions.dag import DAG from airflow.sdk.execution_time.comms import BundleInfo, StartupDetails from airflow.timetables.base import TimeRestriction - from airflow.utils import timezone + + timezone = _import_timezone() def _create_task_instance( task: BaseOperator, @@ -2345,7 +2347,8 @@ def execute(self, context): from airflow.sdk.execution_time.task_runner import run from airflow.sdk.execution_time.xcom import XCom - from airflow.utils import timezone + + timezone = _import_timezone() # Set up spies once at fixture level if hasattr(XCom.set, "spy"): @@ -2564,3 +2567,14 @@ def _create_conn(connection, session=None): monkeypatch.setenv(env_var_name, connection.as_json()) return _create_conn + + +def _import_timezone(): + try: + from airflow.sdk._shared.timezones import timezone + except ModuleNotFoundError: + try: + from airflow._shared.timezones import timezone + except ModuleNotFoundError: + from airflow.utils import timezone + return timezone diff --git a/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py index 59c52178b1d80..6b51627146c14 100644 --- a/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py +++ b/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py @@ -26,6 +26,7 @@ from unittest.mock import ANY, AsyncMock, MagicMock from uuid import uuid4 +import pendulum import pytest from kubernetes import client from kubernetes.client import V1EnvVar, V1PodSecurityContext, V1SecurityContext, models as k8s @@ -54,7 +55,7 @@ def create_context(task) -> Context: from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS dag = DAG(dag_id="dag", schedule=None) - logical_date = timezone.datetime(2016, 1, 1, 1, 0, 0, tzinfo=timezone.parse_timezone("Europe/Amsterdam")) + logical_date = timezone.datetime(2016, 1, 1, 1, 0, 0, tzinfo=pendulum.timezone("Europe/Amsterdam")) if AIRFLOW_V_3_0_PLUS: dag_run = DagRun( diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py index 2168744edcade..85a222571e454 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_msgraph.py @@ -28,7 +28,6 @@ from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.microsoft.azure.operators.msgraph import MSGraphAsyncOperator, execute_callable from airflow.triggers.base import TriggerEvent -from airflow.utils import timezone from tests_common.test_utils.file_loading import load_file_from_resources, load_json_from_resources from tests_common.test_utils.mock_context import mock_context @@ -36,6 +35,11 @@ from unit.microsoft.azure.base import Base from unit.microsoft.azure.test_utils import mock_json_response, mock_response +try: + from airflow.sdk import timezone +except ImportError: + from airflow.utils import timezone # type: ignore[attr-defined,no-redef] + try: from airflow.sdk.definitions.context import Context except ImportError: diff --git a/providers/standard/tests/unit/standard/sensors/test_time.py b/providers/standard/tests/unit/standard/sensors/test_time.py index bbc8c0401c01f..1cec8de4a4874 100644 --- a/providers/standard/tests/unit/standard/sensors/test_time.py +++ b/providers/standard/tests/unit/standard/sensors/test_time.py @@ -27,26 +27,30 @@ from airflow.models.dag import DAG from airflow.providers.standard.sensors.time import TimeSensor from airflow.providers.standard.triggers.temporal import DateTimeTrigger -from airflow.utils import timezone -DEFAULT_TIMEZONE = "Asia/Singapore" # UTC+08:00 +try: + from airflow.sdk import timezone +except ImportError: + from airflow.utils import timezone # type: ignore[attr-defined,no-redef] + +DEFAULT_TIMEZONE = pendulum.timezone("Asia/Singapore") # UTC+08:00 DEFAULT_DATE_WO_TZ = datetime(2015, 1, 1) -DEFAULT_DATE_WITH_TZ = datetime(2015, 1, 1, tzinfo=timezone.parse_timezone(DEFAULT_TIMEZONE)) +DEFAULT_DATE_WITH_TZ = datetime(2015, 1, 1, tzinfo=DEFAULT_TIMEZONE) class TestTimeSensor: @pytest.mark.parametrize( - "default_timezone, start_date, target_time ,expected", + "tzinfo, start_date, target_time ,expected", [ - ("UTC", DEFAULT_DATE_WO_TZ, time(10, 0), True), - ("UTC", DEFAULT_DATE_WITH_TZ, time(16, 0), True), - ("UTC", DEFAULT_DATE_WITH_TZ, time(23, 0), False), + (timezone.utc, DEFAULT_DATE_WO_TZ, time(10, 0), True), + (timezone.utc, DEFAULT_DATE_WITH_TZ, time(16, 0), True), + (timezone.utc, DEFAULT_DATE_WITH_TZ, time(23, 0), False), (DEFAULT_TIMEZONE, DEFAULT_DATE_WO_TZ, time(23, 0), False), ], ) @time_machine.travel(timezone.datetime(2020, 1, 1, 13, 0).replace(tzinfo=timezone.utc)) - def test_timezone(self, default_timezone, start_date, target_time, expected, monkeypatch): - monkeypatch.setattr("airflow.settings.TIMEZONE", timezone.parse_timezone(default_timezone)) + def test_timezone(self, tzinfo, start_date, target_time, expected, monkeypatch): + monkeypatch.setattr("airflow.settings.TIMEZONE", tzinfo) dag = DAG("test_timezone", schedule=None, default_args={"start_date": start_date}) op = TimeSensor(task_id="test", target_time=target_time, dag=dag) assert op.poke(None) == expected @@ -54,7 +58,7 @@ def test_timezone(self, default_timezone, start_date, target_time, expected, mon def test_target_time_aware_dag_timezone(self): # This behavior should be the same for both deferrable and non-deferrable with DAG("test_target_time_aware", schedule=None, start_date=datetime(2020, 1, 1, 13, 0)): - aware_time = time(0, 1).replace(tzinfo=timezone.parse_timezone(DEFAULT_TIMEZONE)) + aware_time = time(0, 1).replace(tzinfo=DEFAULT_TIMEZONE) op = TimeSensor(task_id="test", target_time=aware_time) assert op.target_datetime.tzinfo == timezone.utc @@ -66,8 +70,8 @@ def test_target_time_aware_dag_timezone(self): ], ) def test_target_date_aware_dag_timezone(self, current_datetime, server_timezone): - travel_time = pendulum.parse(current_datetime, tz=timezone.parse_timezone(server_timezone)) - user_timezone = timezone.parse_timezone("Asia/Seoul") + travel_time = pendulum.parse(current_datetime, tz=pendulum.timezone(server_timezone)) + user_timezone = pendulum.timezone("Asia/Seoul") expected_target_datetime = pendulum.datetime(2025, 1, 26, 22, 0, 0, tz="UTC") with time_machine.travel(travel_time, tick=False): @@ -92,7 +96,7 @@ def test_target_time_naive_dag_timezone(self): with DAG( dag_id="test_target_time_naive_dag_timezone", schedule=None, - start_date=datetime(2020, 1, 1, 23, 0, tzinfo=timezone.parse_timezone(DEFAULT_TIMEZONE)), + start_date=datetime(2020, 1, 1, 23, 0, tzinfo=DEFAULT_TIMEZONE), ): op = TimeSensor(task_id="test", target_time=time(9, 0)) @@ -111,7 +115,7 @@ def test_task_is_deferred(self): # This should be converted to UTC in the __init__. Since there is no default timezone, it will become # aware, but note changed - assert not timezone.is_naive(op.target_datetime) + assert op.target_datetime.utcoffset() is not None with pytest.raises(TaskDeferred) as exc_info: op.execute({}) diff --git a/pyproject.toml b/pyproject.toml index 0276845787891..6ad8960afdfdd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -679,7 +679,7 @@ section-order = [ "testing" ] -known-first-party = ["airflow"] +known-first-party = ["airflow", "airflow_shared", "airflow_shared.*"] # Make sure we put the "dev" imports at the end, not as a third-party module [tool.ruff.lint.isort.sections] @@ -809,20 +809,30 @@ ban-relative-imports = "all" # that they're imported lazily (e.g., within a function definition). banned-module-level-imports = ["numpy", "pandas", "polars"] + [tool.ruff.lint.flake8-tidy-imports.banned-api] + # Direct import from the airflow package modules and constraints "airflow.AirflowException".msg = "Use airflow.exceptions.AirflowException instead." "airflow.Dataset".msg = "Use airflow.datasets.Dataset instead." + # Deprecated imports "airflow.models.baseoperator.BaseOperatorLink".msg = "Use airflow.models.baseoperatorlink.BaseOperatorLink" "airflow.models.errors.ImportError".msg = "Use airflow.models.errors.ParseImportError" "airflow.models.ImportError".msg = "Use airflow.models.errors.ParseImportError" + +# Note: we have to specify these all here (not in sub pyproject.toml files), else we'd have to duplicate these +# all rules in each project https://github.com/astral-sh/ruff/issues/18723 +# Add this once providers are sorted +# "airflow.utils.timezone".msg = "Use airflow.sdk.timezone (or airflow._shared.timezones.timezone from within airflow-core)" +# "airflow_shared".msg = "Use airflow._shared or airflow.sdk._shared instead" + # Deprecated in Python 3.11, Pending Removal in Python 3.15: https://github.com/python/cpython/issues/90817 # Deprecation warning in Python 3.11 also recommends using locale.getencoding but it available in Python 3.11 "locale.getdefaultlocale".msg = "Use locale.setlocale() and locale.getlocale() instead." # Deprecated in Python 3.12: https://github.com/python/cpython/issues/103857 -"datetime.datetime.utcnow".msg = "Use airflow.utils.timezone.utcnow or datetime.datetime.now(tz=datetime.timezone.utc)" -"datetime.datetime.utcfromtimestamp".msg = "Use airflow.utils.timezone.from_timestamp or datetime.datetime.fromtimestamp(tz=datetime.timezone.utc)" +"datetime.datetime.utcnow".msg = "Use airflow.sdk.timezone.utcnow or datetime.datetime.now(tz=datetime.timezone.utc)" +"datetime.datetime.utcfromtimestamp".msg = "Use airflow.sdk.timezone.from_timestamp or datetime.datetime.fromtimestamp(tz=datetime.timezone.utc)" # Deprecated in Python 3.12: https://github.com/python/cpython/issues/94309 "typing.Hashable".msg = "Use collections.abc.Hashable" "typing.Sized".msg = "Use collections.abc.Sized" @@ -1247,7 +1257,8 @@ dev = [ "apache-airflow-helm-tests", "apache-airflow-kubernetes-tests", "apache-airflow-task-sdk", - "apache-airflow-ctl" + "apache-airflow-ctl", + "apache-airflow-shared-timezones", ] # To build docs: @@ -1293,6 +1304,7 @@ apache-airflow-helm-tests = { workspace = true } apache-airflow-kubernetes-tests = { workspace = true } apache-airflow-providers = { workspace = true } apache-aurflow-docker-stack = { workspace = true } +apache-airflow-shared-timezones = { workspace = true } # Automatically generated provider workspace items (update_airflow_pyproject_toml.py) apache-airflow-providers-airbyte = { workspace = true } apache-airflow-providers-alibaba = { workspace = true } @@ -1407,6 +1419,7 @@ members = [ "task-sdk", "providers-summary-docs", "docker-stack-docs", + "shared/timezones", # Automatically generated provider workspace members (update_airflow_pyproject_toml.py) "providers/airbyte", "providers/alibaba", diff --git a/scripts/ci/docker-compose/local.yml b/scripts/ci/docker-compose/local.yml index 6609aa5c1dbda..2ac4dc849ad3c 100644 --- a/scripts/ci/docker-compose/local.yml +++ b/scripts/ci/docker-compose/local.yml @@ -64,11 +64,14 @@ services: source: ../../../airflow-ctl target: /opt/airflow/airflow-ctl - type: bind - source: ../../../constraints - target: /opt/airflow/constraints + source: ../../../chart + target: /opt/airflow/chart - type: bind source: ../../../clients target: /opt/airflow/clients + - type: bind + source: ../../../constraints + target: /opt/airflow/constraints - type: bind source: ../../../dags target: /opt/airflow/dags @@ -76,17 +79,26 @@ services: source: ../../../dev target: /opt/airflow/dev - type: bind - source: ../../../docs - target: /opt/airflow/docs + source: ../../../devel-common + target: /opt/airflow/devel-common - type: bind source: ../../../docker-stack-docs target: /opt/airflow/docker-stack-docs - type: bind - source: ../../../providers-summary-docs - target: /opt/airflow/providers-summary-docs + source: ../../../docker-tests + target: /opt/airflow/docker-tests + - type: bind + source: ../../../docs + target: /opt/airflow/docs - type: bind source: ../../../generated target: /opt/airflow/generated + - type: bind + source: ../../../helm-tests + target: /opt/airflow/helm-tests + - type: bind + source: ../../../kubernetes-tests + target: /opt/airflow/kubernetes-tests - type: bind source: ../../../logs target: /root/airflow/logs @@ -94,8 +106,8 @@ services: source: ../../../providers target: /opt/airflow/providers - type: bind - source: ../../../task-sdk - target: /opt/airflow/task-sdk + source: ../../../providers-summary-docs + target: /opt/airflow/providers-summary-docs - type: bind source: ../../../pyproject.toml target: /opt/airflow/pyproject.toml @@ -106,20 +118,11 @@ services: source: ../../../scripts/docker/entrypoint_ci.sh target: /entrypoint - type: bind - source: ../../../devel-common - target: /opt/airflow/devel-common - - type: bind - source: ../../../helm-tests - target: /opt/airflow/helm-tests + source: ../../../shared + target: /opt/airflow/shared - type: bind - source: ../../../kubernetes-tests - target: /opt/airflow/kubernetes-tests - - type: bind - source: ../../../docker-tests - target: /opt/airflow/docker-tests - - type: bind - source: ../../../chart - target: /opt/airflow/chart + source: ../../../task-sdk + target: /opt/airflow/task-sdk # END automatically generated volumes from VOLUMES_FOR_SELECTED_MOUNTS in docker_command_utils.py volumes: root-airflow-volume: diff --git a/scripts/cov/core_coverage.py b/scripts/cov/core_coverage.py index 317aca4d18826..fd99a8338024a 100644 --- a/scripts/cov/core_coverage.py +++ b/scripts/cov/core_coverage.py @@ -143,7 +143,6 @@ "airflow-core/src/airflow/utils/task_group.py", "airflow-core/src/airflow/utils/task_instance_session.py", "airflow-core/src/airflow/utils/timeout.py", - "airflow-core/src/airflow/utils/timezone.py", "airflow-core/src/airflow/utils/weight_rule.py", "airflow-core/src/airflow/utils/yaml.py", ] diff --git a/setup_idea.py b/setup_idea.py index ec6b82d7e0f1f..54a4724ba3d9f 100755 --- a/setup_idea.py +++ b/setup_idea.py @@ -94,12 +94,14 @@ AIRFLOW_IML_FILE = IDEA_FOLDER_PATH / "airflow.iml" MODULES_XML_FILE = IDEA_FOLDER_PATH / "modules.xml" -ROOT_PROVIDERS_FOLDER_PATH = ROOT_AIRFLOW_FOLDER_PATH / "providers" - def setup_idea(): # Providers discovery - for pyproject_toml_file in ROOT_PROVIDERS_FOLDER_PATH.rglob("pyproject.toml"): + for pyproject_toml_file in ROOT_AIRFLOW_FOLDER_PATH.rglob("providers/**/pyproject.toml"): + relative_path = pyproject_toml_file.relative_to(ROOT_AIRFLOW_FOLDER_PATH).parent.as_posix() + source_root_modules.append(f"{relative_path}") + # Shared discovery + for pyproject_toml_file in ROOT_AIRFLOW_FOLDER_PATH.rglob("shared/*/pyproject.toml"): relative_path = pyproject_toml_file.relative_to(ROOT_AIRFLOW_FOLDER_PATH).parent.as_posix() source_root_modules.append(f"{relative_path}") diff --git a/shared/README.md b/shared/README.md new file mode 100644 index 0000000000000..baa57598546bf --- /dev/null +++ b/shared/README.md @@ -0,0 +1,91 @@ + + +# Shared Python Code for Airflow Components + +This folder contains code that is shared across two or more of the Airflow distributions, such as airflow-core and task-sdk. + +## Be Thoughtful about what you add under here + +Not every piece of code that is used in two distributions should be automatically placed in one of the shared libraries, and sometimes "just duplicate it" is the right approach to take. For example if it's just a 5 or 10 line function and it's used in two places, it might be easier to future developers to understand if the function is in two places. + +There is no hard rules about what should or shouldn't be in these libraries, so try to apply your best judgement. + +## Mechanics of sharing + +The primary mechanism by which we share code is to use in-repo symlinks. This means that there is a single copy of the code existing in the repo (no need for a pre-commit to update other copies, resulting in larger diffs). + +The primary reason we are using this sharing approach, rather than the perhaps more traditional approach of depending on a shared distribution from PyPI is around compatibility. Lets say that we have a library to parse AirflowConfig (including all the sources, env, `_cmd` handling, etc etc.). If that code is used by two distributions that are installed in the same Python environment, then either we introduce another source of "version hell" or we have to make every change both backwards and forwards compatible. + +Instead by using an approach similar to vendoring (or if you compare it to C libraries - static linking), where each dist ships with a copy of the code of the version it knows it works with it allows them to co-exist in the same env without problems, and also reduces the cognitive load on developers when making changes to the shared libs -- as long as it passes the CI it should be good (and we don't have to test an ever increasing matrix of versions to ensure compatibility.) + +### Shared libraries _must_ use relative imports for other shared libraries** + +The one caveat to this is due to the side-effect that these shared modules are going to be imported from different locations (for example `airflow._shared.timezones.timezone` and `airflow.sdk._shared.timezones.timezone`) then any imports inside the shared code to other parts of shared libraries _must_ make use of relative imports, and possibly needs `# noqa: TID252` + +For example, to use the shared timezone library from another shared library, lets say `shared/logging/src/airflow_shared/logging/config.py` you would have + +```python +from ..timezones.timezone import is_naive # noqa: TID252 +``` + +This is different to the rest of the airflow codebase where relative imports are not allowed. + +### Shared libraries should be grouped one level beneath `airflow_shared/` + +In order to make both the relative imports above function and to provider simple structure, we want to arrange things in single named chunks, where the folder name is also the "shared library name". For example + + +``` +shared +├── README.md +└── timezones/ + ├── pyproject.toml + ├── src/ + │ └── airflow_shared + │ └── timezones/ + │ ├── __init__.py + └── tests/ +└── logging/ + ├── pyproject.toml + ├── src/ + │ └── airflow_shared + │ └── logging/ + │ ├── __init__.py + └── tests/ +``` + +## Use In Editable installs/Git checkouts + +In the editable checkout (which is what you get by default when using `uv sync` etc), the shared libraries are symlinked directly in place, for example `task-sdk/src/airflow/sdk/_shared/timezones` is symlinked to `shared/timezones/src/airflow_shared/timezones`. This means that python automatically finds the files underneath the module as if they lived directly under `airflow.sdk._shared.timezones.X`. + + +### Including In built distributions + +Symlinks will not work when we build the main distributions for PyPI and other channels (since the symlinks would be to files outside of the distribution otherwise). This is specific to sdist; building wheel resolves the symlinks. + +To make this work we need to configure hatchling to include extra files. For example in the Task SDK we have + +```toml +[tool.hatch.build.targets.sdist.force-include] +"../shared/logging/src/airflow_shared/logging" = "src/airflow/sdk/_shared/logging" +"../shared/timezones/src/airflow_shared/timezones" = "src/airflow/sdk/_shared/timezones" +``` + +We will need one entry in that table for each module we include directly below `src/airflow/sdk/_shared`. diff --git a/shared/timezones/pyproject.toml b/shared/timezones/pyproject.toml new file mode 100644 index 0000000000000..bc68f6467c894 --- /dev/null +++ b/shared/timezones/pyproject.toml @@ -0,0 +1,47 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[project] +name = "apache-airflow-shared-timezones" +description = "Shared timezone code for Airflow distributions" +version = "0.0" +classifiers = [ + "Private :: Do Not Upload", +] + +dependencies = [ + "pendulum", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/airflow_shared"] + +[tool.ruff] +extend = "../../pyproject.toml" +src = ["src"] + +[tool.ruff.lint.per-file-ignores] +# Ignore Doc rules et al for anything outside of tests +"!src/*" = ["D", "S101", "TRY002"] + +[tool.ruff.lint.flake8-tidy-imports] +# Override the workspace level default +ban-relative-imports = "parents" diff --git a/shared/timezones/src/airflow_shared/timezones/__init__.py b/shared/timezones/src/airflow_shared/timezones/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/shared/timezones/src/airflow_shared/timezones/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow-core/src/airflow/utils/timezone.py b/shared/timezones/src/airflow_shared/timezones/timezone.py similarity index 96% rename from airflow-core/src/airflow/utils/timezone.py rename to shared/timezones/src/airflow_shared/timezones/timezone.py index 3e59169598aab..3abdd06ee5ebc 100644 --- a/airflow-core/src/airflow/utils/timezone.py +++ b/shared/timezones/src/airflow_shared/timezones/timezone.py @@ -68,17 +68,6 @@ def utcnow() -> dt.datetime: return dt.datetime.now(tz=utc) -def utc_epoch() -> dt.datetime: - """Get the epoch in the user's timezone.""" - # pendulum utcnow() is not used as that sets a TimezoneInfo object - # instead of a Timezone. This is not picklable and also creates issues - # when using replace() - result = dt.datetime(1970, 1, 1) - result = result.replace(tzinfo=utc) - - return result - - @overload def convert_to_utc(value: None) -> None: ... @@ -284,7 +273,7 @@ def parse_timezone(name: str | int) -> FixedTimezone | Timezone: """ if _PENDULUM3: # This only presented in pendulum 3 and code do not reached into the pendulum 2 - return pendulum.timezone(name) + return pendulum.timezone(name) # type: ignore[operator] # In pendulum 2 this refers to the function, in pendulum 3 refers to the module return pendulum.tz.timezone(name) # type: ignore[operator] diff --git a/airflow-core/tests/unit/utils/test_timezone.py b/shared/timezones/tests/test_timezone.py similarity index 98% rename from airflow-core/tests/unit/utils/test_timezone.py rename to shared/timezones/tests/test_timezone.py index c7e6d0c09aa81..ba0f2a46def55 100644 --- a/airflow-core/tests/unit/utils/test_timezone.py +++ b/shared/timezones/tests/test_timezone.py @@ -23,8 +23,8 @@ import pytest from pendulum.tz.timezone import FixedTimezone, Timezone -from airflow.utils import timezone -from airflow.utils.timezone import coerce_datetime, parse_timezone +from airflow_shared.timezones import timezone +from airflow_shared.timezones.timezone import coerce_datetime, parse_timezone CET = Timezone("Europe/Paris") EAT = Timezone("Africa/Nairobi") # Africa/Nairobi diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml index a6edccc35c634..de87a4a5f31f4 100644 --- a/task-sdk/pyproject.toml +++ b/task-sdk/pyproject.toml @@ -85,6 +85,9 @@ build-backend = "hatchling.build" [tool.hatch.version] path = "src/airflow/sdk/__init__.py" +[tool.hatch.build.targets.sdist.force-include] +"../shared/timezones/src/airflow_shared/timezones" = "src/airflow/sdk/_shared/timezones" + [tool.hatch.build.targets.wheel] packages = ["src/airflow"] # This file only exists to make pyright/VSCode happy, don't ship it diff --git a/task-sdk/src/airflow/sdk/_shared/__init__.py b/task-sdk/src/airflow/sdk/_shared/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/task-sdk/src/airflow/sdk/_shared/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/task-sdk/src/airflow/sdk/_shared/timezones b/task-sdk/src/airflow/sdk/_shared/timezones new file mode 120000 index 0000000000000..b66387c7a6b56 --- /dev/null +++ b/task-sdk/src/airflow/sdk/_shared/timezones @@ -0,0 +1 @@ +../../../../../shared/timezones/src/airflow_shared/timezones \ No newline at end of file diff --git a/task-sdk/src/airflow/sdk/bases/decorator.py b/task-sdk/src/airflow/sdk/bases/decorator.py index 42bc66ee5e487..8ca3050655b2e 100644 --- a/task-sdk/src/airflow/sdk/bases/decorator.py +++ b/task-sdk/src/airflow/sdk/bases/decorator.py @@ -28,6 +28,7 @@ import attr import typing_extensions +from airflow.sdk import timezone from airflow.sdk.bases.operator import ( BaseOperator, coerce_resources, @@ -46,7 +47,6 @@ from airflow.sdk.definitions.asset import Asset from airflow.sdk.definitions.mappedoperator import MappedOperator, ensure_xcomarg_return_value from airflow.sdk.definitions.xcom_arg import XComArg -from airflow.utils import timezone from airflow.utils.context import KNOWN_CONTEXT_KEYS from airflow.utils.decorators import remove_task_decorator from airflow.utils.helpers import prevent_duplicates diff --git a/task-sdk/src/airflow/sdk/bases/operator.py b/task-sdk/src/airflow/sdk/bases/operator.py index 6f0078eef535c..85f6b0745b18d 100644 --- a/task-sdk/src/airflow/sdk/bases/operator.py +++ b/task-sdk/src/airflow/sdk/bases/operator.py @@ -36,6 +36,7 @@ import attrs from airflow.exceptions import RemovedInAirflow4Warning +from airflow.sdk import timezone from airflow.sdk.definitions._internal.abstractoperator import ( DEFAULT_IGNORE_FIRST_DEPENDS_ON_PAST, DEFAULT_OWNER, @@ -65,7 +66,6 @@ airflow_priority_weight_strategies, validate_and_load_priority_weight_strategy, ) -from airflow.utils import timezone from airflow.utils.trigger_rule import TriggerRule from airflow.utils.weight_rule import db_safe_priority diff --git a/task-sdk/src/airflow/sdk/bases/sensor.py b/task-sdk/src/airflow/sdk/bases/sensor.py index a39de2a151873..0ca447189a983 100644 --- a/task-sdk/src/airflow/sdk/bases/sensor.py +++ b/task-sdk/src/airflow/sdk/bases/sensor.py @@ -36,8 +36,8 @@ TaskDeferralError, TaskDeferralTimeout, ) +from airflow.sdk import timezone from airflow.sdk.bases.operator import BaseOperator -from airflow.utils import timezone if TYPE_CHECKING: from airflow.sdk.definitions.context import Context diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py b/task-sdk/src/airflow/sdk/definitions/dag.py index 3508efc86633b..05bd34e95bfdf 100644 --- a/task-sdk/src/airflow/sdk/definitions/dag.py +++ b/task-sdk/src/airflow/sdk/definitions/dag.py @@ -217,7 +217,7 @@ def dict_copy(_: dict[str, Any]) -> dict[str, Any]: ... def _default_start_date(instance: DAG): # Find start date inside default_args for compat with Airflow 2. - from airflow.utils import timezone + from airflow.sdk import timezone if date := instance.default_args.get("start_date"): if not isinstance(date, datetime): @@ -456,7 +456,7 @@ def __rich_repr__(self): ) def __attrs_post_init__(self): - from airflow.utils import timezone + from airflow.sdk import timezone # Apply the timezone we settled on to end_date if it wasn't supplied if isinstance(_end_date := self.default_args.get("end_date"), str): @@ -528,7 +528,7 @@ def _default_timetable(instance: DAG): def _extract_tz(instance): import pendulum - from airflow.utils import timezone + from airflow.sdk import timezone start_date = instance.start_date or instance.default_args.get("start_date") @@ -1047,9 +1047,9 @@ def test( from airflow.configuration import secrets_backend_list from airflow.models.dag import DAG as SchedulerDAG, _get_or_create_dagrun from airflow.models.dagrun import DagRun + from airflow.sdk import timezone from airflow.secrets.local_filesystem import LocalFilesystemBackend from airflow.serialization.serialized_objects import SerializedDAG - from airflow.utils import timezone from airflow.utils.state import DagRunState, State, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/task-sdk/src/airflow/sdk/execution_time/cache.py b/task-sdk/src/airflow/sdk/execution_time/cache.py index 9902a05f032f9..392b03d4215f1 100644 --- a/task-sdk/src/airflow/sdk/execution_time/cache.py +++ b/task-sdk/src/airflow/sdk/execution_time/cache.py @@ -21,7 +21,7 @@ import multiprocessing from airflow.configuration import conf -from airflow.utils import timezone +from airflow.sdk import timezone class SecretCache: diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index c937c680cd4eb..be680a08d045f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -103,9 +103,9 @@ set_current_context, ) from airflow.sdk.execution_time.xcom import XCom +from airflow.sdk.timezone import coerce_datetime from airflow.utils.net import get_hostname from airflow.utils.platform import getuser -from airflow.utils.timezone import coerce_datetime if TYPE_CHECKING: import jinja2 diff --git a/task-sdk/src/airflow/sdk/timezone.py b/task-sdk/src/airflow/sdk/timezone.py new file mode 100644 index 0000000000000..1769f510c58a4 --- /dev/null +++ b/task-sdk/src/airflow/sdk/timezone.py @@ -0,0 +1,39 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +# We don't want to `import *` here to avoid the risk of making adding too much to Public python API +from airflow.sdk._shared.timezones.timezone import ( + coerce_datetime, + convert_to_utc, + datetime, + make_naive, + parse, + utc, + utcnow, +) + +__all__ = [ + "coerce_datetime", + "convert_to_utc", + "datetime", + "make_naive", + "parse", + "utc", + "utcnow", +] diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 5df79afd8e706..30cfd2fd7d616 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -29,6 +29,7 @@ from task_sdk import make_client, make_client_w_dry_run, make_client_w_responses from uuid6 import uuid7 +from airflow.sdk import timezone from airflow.sdk.api.client import RemoteValidationError, ServerResponseError from airflow.sdk.api.datamodels._generated import ( AssetEventsResponse, @@ -50,7 +51,6 @@ RescheduleTask, TaskRescheduleStartDate, ) -from airflow.utils import timezone from airflow.utils.state import TerminalTIState if TYPE_CHECKING: diff --git a/task-sdk/tests/task_sdk/bases/test_sensor.py b/task-sdk/tests/task_sdk/bases/test_sensor.py index d0e2b430d2c74..8b9cc04a7860b 100644 --- a/task-sdk/tests/task_sdk/bases/test_sensor.py +++ b/task-sdk/tests/task_sdk/bases/test_sensor.py @@ -19,7 +19,7 @@ from datetime import timedelta from typing import TYPE_CHECKING -from unittest.mock import Mock, patch +from unittest.mock import Mock import pytest import time_machine @@ -34,12 +34,12 @@ ) from airflow.models.trigger import TriggerFailureReason from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.sdk import timezone from airflow.sdk.bases.sensor import BaseSensorOperator, PokeReturnValue, poke_mode_only from airflow.sdk.definitions.dag import DAG from airflow.sdk.execution_time.comms import RescheduleTask, TaskRescheduleStartDate -from airflow.utils import timezone +from airflow.sdk.timezone import datetime from airflow.utils.state import State -from airflow.utils.timezone import datetime if TYPE_CHECKING: from airflow.sdk.definitions.context import Context @@ -368,13 +368,11 @@ def test_sensor_with_exponential_backoff_on(self): task_id=SENSOR_OP, return_value=None, poke_interval=5, timeout=60, exponential_backoff=True ) - with patch("airflow.utils.timezone.utcnow") as mock_utctime: - mock_utctime.return_value = DEFAULT_DATE - + with time_machine.travel(DEFAULT_DATE): started_at = timezone.utcnow() - timedelta(seconds=10) def run_duration(): - return (timezone.utcnow - started_at).total_seconds() + return (timezone.utcnow() - started_at).total_seconds() interval1 = sensor._get_next_poke_interval(started_at, run_duration, 1) interval2 = sensor._get_next_poke_interval(started_at, run_duration, 2) @@ -395,13 +393,11 @@ def test_sensor_with_exponential_backoff_on_and_small_poke_interval(self, poke_i exponential_backoff=True, ) - with patch("airflow.utils.timezone.utcnow") as mock_utctime: - mock_utctime.return_value = DEFAULT_DATE - + with time_machine.travel(DEFAULT_DATE): started_at = timezone.utcnow() - timedelta(seconds=10) def run_duration(): - return (timezone.utcnow - started_at).total_seconds() + return (timezone.utcnow() - started_at).total_seconds() intervals = [ sensor._get_next_poke_interval(started_at, run_duration, retry_number) @@ -428,13 +424,11 @@ def test_sensor_with_exponential_backoff_on_and_max_wait(self): max_wait=timedelta(seconds=30), ) - with patch("airflow.utils.timezone.utcnow") as mock_utctime: - mock_utctime.return_value = DEFAULT_DATE - + with time_machine.travel(DEFAULT_DATE): started_at = timezone.utcnow() - timedelta(seconds=10) def run_duration(): - return (timezone.utcnow - started_at).total_seconds() + return (timezone.utcnow() - started_at).total_seconds() for idx, expected in enumerate([2, 6, 13, 30, 30, 30, 30, 30]): assert sensor._get_next_poke_interval(started_at, run_duration, idx) == expected diff --git a/task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py b/task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py index 15c2d65cc2717..9f3d8a343d240 100644 --- a/task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py +++ b/task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py @@ -20,8 +20,8 @@ import datetime from airflow.providers.standard.sensors.date_time import DateTimeSensorAsync +from airflow.sdk import timezone from airflow.sdk.definitions.dag import dag -from airflow.utils import timezone @dag diff --git a/task-sdk/tests/task_sdk/docs/test_public_api.py b/task-sdk/tests/task_sdk/docs/test_public_api.py index dd2bf54980ce9..f1e2b74e47fa6 100644 --- a/task-sdk/tests/task_sdk/docs/test_public_api.py +++ b/task-sdk/tests/task_sdk/docs/test_public_api.py @@ -52,6 +52,7 @@ def test_airflow_sdk_no_unexpected_exports(): "io", "log", "exceptions", + "timezone", } unexpected = actual - public - ignore assert not unexpected, f"Unexpected exports in airflow.sdk: {sorted(unexpected)}" diff --git a/task-sdk/tests/task_sdk/execution_time/test_comms.py b/task-sdk/tests/task_sdk/execution_time/test_comms.py index 7960fc7b27570..2f5d6bceedc21 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_comms.py +++ b/task-sdk/tests/task_sdk/execution_time/test_comms.py @@ -24,9 +24,9 @@ import msgspec import pytest +from airflow.sdk import timezone from airflow.sdk.execution_time.comms import BundleInfo, StartupDetails, _ResponseFrame from airflow.sdk.execution_time.task_runner import CommsDecoder -from airflow.utils import timezone class TestCommsDecoder: diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py b/task-sdk/tests/task_sdk/execution_time/test_context.py index d3c22598a98e3..1def08086cce3 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context.py @@ -22,7 +22,7 @@ import pytest -from airflow.sdk import BaseOperator, get_current_context +from airflow.sdk import BaseOperator, get_current_context, timezone from airflow.sdk.api.datamodels._generated import AssetEventResponse, AssetResponse from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.definitions.asset import ( @@ -63,7 +63,6 @@ context_to_airflow_vars, set_current_context, ) -from airflow.utils import timezone def test_convert_connection_result_conn(): diff --git a/task-sdk/tests/task_sdk/execution_time/test_hitl.py b/task-sdk/tests/task_sdk/execution_time/test_hitl.py index cab17e30bacec..a936812bbb7da 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_hitl.py +++ b/task-sdk/tests/task_sdk/execution_time/test_hitl.py @@ -19,6 +19,7 @@ from uuid6 import uuid7 +from airflow.sdk import timezone from airflow.sdk.api.datamodels._generated import HITLDetailResponse from airflow.sdk.execution_time.comms import CreateHITLDetailPayload from airflow.sdk.execution_time.hitl import ( @@ -26,7 +27,6 @@ get_hitl_detail_content_detail, update_htil_detail_response, ) -from airflow.utils import timezone TI_ID = uuid7() diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index c6aa651af441e..271f2b98c15d0 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -44,6 +44,7 @@ from uuid6 import uuid7 from airflow.executors.workloads import BundleInfo +from airflow.sdk import timezone from airflow.sdk.api import client as sdk_client from airflow.sdk.api.client import ServerResponseError from airflow.sdk.api.datamodels._generated import ( @@ -119,7 +120,6 @@ set_supervisor_comms, supervise, ) -from airflow.utils import timezone, timezone as tz if TYPE_CHECKING: import kgb @@ -237,7 +237,7 @@ def subprocess_main(): line = lineno() - 2 # Line the error should be on - instant = tz.datetime(2024, 11, 7, 12, 34, 56, 78901) + instant = timezone.datetime(2024, 11, 7, 12, 34, 56, 78901) time_machine.move_to(instant, tick=False) proc = ActivitySubprocess.start( @@ -460,7 +460,7 @@ def _on_child_started(self, *args, **kwargs): def test_run_simple_dag(self, test_dags_dir, captured_logs, time_machine, mocker, client_with_ti_start): """Test running a simple DAG in a subprocess and capturing the output.""" - instant = tz.datetime(2024, 11, 7, 12, 34, 56, 78901) + instant = timezone.datetime(2024, 11, 7, 12, 34, 56, 78901) time_machine.move_to(instant, tick=False) dagfile_path = test_dags_dir @@ -504,7 +504,7 @@ def test_supervise_handles_deferred_task( This includes ensuring the task starts and executes successfully, and that the task is deferred (via the API client) with the expected parameters. """ - instant = tz.datetime(2024, 11, 7, 12, 34, 56, 0) + instant = timezone.datetime(2024, 11, 7, 12, 34, 56, 0) ti = TaskInstance( id=uuid7(), @@ -731,7 +731,7 @@ def test_heartbeat_failures_handling(self, monkeypatch, mocker, captured_logs, t process=mock_process, ) - time_now = tz.datetime(2024, 11, 28, 12, 0, 0) + time_now = timezone.datetime(2024, 11, 28, 12, 0, 0) time_machine.move_to(time_now, tick=False) # Simulate sending heartbeats and ensure the process gets killed after max retries diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index e0b414065cb0d..a88b244897d0b 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -52,6 +52,7 @@ dag as dag_decorator, get_current_context, task as task_decorator, + timezone, ) from airflow.sdk.api.datamodels._generated import ( AssetProfile, @@ -121,7 +122,6 @@ startup, ) from airflow.sdk.execution_time.xcom import XCom -from airflow.utils import timezone from airflow.utils.types import NOTSET, ArgNotSet from tests_common.test_utils.mock_operators import AirflowLink @@ -1139,7 +1139,6 @@ def test_get_context_without_ti_context_from_server(self, mocked_parse, make_ti_ def test_get_context_with_ti_context_from_server(self, create_runtime_ti, mock_supervisor_comms): """Test the context keys are added when sent from API server (mocked)""" - from airflow.utils import timezone task = BaseOperator(task_id="hello")