You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2025/07/23 21:06:05 UTC
(airflow) branch main updated: Set up process for sharing code between different components (#53149)
This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5f07ae13260 Set up process for sharing code between different components (#53149)
5f07ae13260 is described below
commit 5f07ae13260068f48f4b3e35502755b593e82b4c
Author: Ash Berlin-Taylor <as...@apache.org>
AuthorDate: Wed Jul 23 22:05:57 2025 +0100
Set up process for sharing code between different components (#53149)
---
.dockerignore | 1 +
airflow-core/docs/img/airflow_erd.sha256 | 2 +-
airflow-core/pyproject.toml | 3 +
.../models => src/airflow/_shared}/__init__.py | 9 ---
airflow-core/src/airflow/_shared/timezones | 1 +
airflow-core/src/airflow/api/common/trigger_dag.py | 2 +-
.../src/airflow/api_fastapi/auth/tokens.py | 2 +-
.../src/airflow/api_fastapi/common/parameters.py | 2 +-
.../src/airflow/api_fastapi/common/types.py | 2 +-
.../api_fastapi/core_api/datamodels/dag_run.py | 2 +-
.../api_fastapi/core_api/datamodels/ui/common.py | 2 +-
.../api_fastapi/core_api/routes/public/assets.py | 2 +-
.../core_api/routes/public/backfills.py | 2 +-
.../api_fastapi/core_api/routes/public/hitl.py | 2 +-
.../api_fastapi/core_api/routes/ui/dashboard.py | 2 +-
.../api_fastapi/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 +-
airflow-core/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 +-
airflow-core/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 +-
airflow-core/src/airflow/models/serialized_dag.py | 2 +-
airflow-core/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 +-
.../airflow/serialization/serialized_objects.py | 2 +-
.../airflow/serialization/serializers/datetime.py | 2 +-
.../airflow/serialization/serializers/timezone.py | 2 +-
airflow-core/src/airflow/settings.py | 2 +-
.../ti_deps/deps/not_in_retry_period_dep.py | 2 +-
.../airflow/ti_deps/deps/ready_to_reschedule.py | 2 +-
.../airflow/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 +-
airflow-core/src/airflow/timetables/interval.py | 2 +-
airflow-core/src/airflow/timetables/simple.py | 2 +-
airflow-core/src/airflow/timetables/trigger.py | 2 +-
airflow-core/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 +-
.../airflow/utils/log/file_processor_handler.py | 2 +-
.../src/airflow/utils/log/file_task_handler.py | 2 +-
.../src/airflow/utils/log/timezone_aware.py | 6 +-
airflow-core/src/airflow/utils/sqlalchemy.py | 2 +-
airflow-core/tests/integration/otel/test_otel.py | 2 +-
.../tests/unit/api_fastapi/auth/test_tokens.py | 2 +-
.../tests/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 +-
.../core_api/routes/public/test_extra_links.py | 2 +-
.../core_api/routes/public/test_hitl.py | 2 +-
.../api_fastapi/core_api/routes/public/test_log.py | 2 +-
.../core_api/routes/public/test_monitor.py | 2 +-
.../core_api/routes/public/test_task_instances.py | 2 +-
.../core_api/routes/public/test_xcom.py | 2 +-
.../core_api/routes/ui/test_backfills.py | 2 +-
.../api_fastapi/core_api/routes/ui/test_grid.py | 2 +-
.../core_api/routes/ui/test_structure.py | 2 +-
.../versions/head/test_asset_events.py | 2 +-
.../execution_api/versions/head/test_assets.py | 2 +-
.../execution_api/versions/head/test_dag_runs.py | 2 +-
.../versions/head/test_task_instances.py | 2 +-
.../versions/v2025_04_28/test_task_instances.py | 2 +-
.../tests/unit/callbacks/test_callback_requests.py | 2 +-
.../unit/cli/commands/test_backfill_command.py | 2 +-
.../tests/unit/cli/commands/test_dag_command.py | 7 +-
.../tests/unit/cli/commands/test_task_command.py | 2 +-
airflow-core/tests/unit/core/test_core.py | 2 +-
.../tests/unit/core/test_impersonation_tests.py | 2 +-
airflow-core/tests/unit/core/test_sentry.py | 2 +-
.../tests/unit/dag_processing/bundles/test_base.py | 2 +-
.../tests/unit/dag_processing/test_collection.py | 2 +-
.../tests/unit/dag_processing/test_manager.py | 2 +-
.../tests/unit/dag_processing/test_processor.py | 2 +-
.../tests/unit/dags/test_cli_triggered_dags.py | 2 +-
airflow-core/tests/unit/dags/test_invalid_cron.py | 2 +-
.../tests/unit/dags/test_logging_in_dag.py | 2 +-
airflow-core/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 +-
airflow-core/tests/unit/decorators/test_mapped.py | 2 +-
.../tests/unit/executors/test_base_executor.py | 2 +-
.../tests/unit/executors/test_local_executor.py | 2 +-
airflow-core/tests/unit/jobs/test_base_job.py | 2 +-
airflow-core/tests/unit/jobs/test_scheduler_job.py | 5 +-
airflow-core/tests/unit/jobs/test_triggerer_job.py | 2 +-
.../tests/unit/listeners/test_listeners.py | 2 +-
airflow-core/tests/unit/models/__init__.py | 2 +-
airflow-core/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 +-
.../tests/unit/models/test_renderedtifields.py | 2 +-
.../tests/unit/models/test_taskinstance.py | 2 +-
airflow-core/tests/unit/models/test_timestamp.py | 2 +-
airflow-core/tests/unit/models/test_trigger.py | 4 +-
airflow-core/tests/unit/models/test_xcom.py | 2 +-
airflow-core/tests/unit/sensors/test_filesystem.py | 2 +-
.../unit/serialization/test_dag_serialization.py | 2 +-
.../unit/serialization/test_serialized_objects.py | 2 +-
.../ti_deps/deps/test_not_in_retry_period_dep.py | 2 +-
.../unit/ti_deps/deps/test_prev_dagrun_dep.py | 2 +-
.../ti_deps/deps/test_ready_to_reschedule_dep.py | 2 +-
.../ti_deps/deps/test_runnable_exec_date_dep.py | 2 +-
.../tests/unit/timetables/test_events_timetable.py | 2 +-
.../unit/timetables/test_interval_timetable.py | 2 +-
.../tests/unit/timetables/test_once_timetable.py | 2 +-
.../unit/timetables/test_trigger_timetable.py | 2 +-
.../unit/timetables/test_workday_timetable.py | 2 +-
.../unit/utils/log/test_file_processor_handler.py | 2 +-
.../tests/unit/utils/log/test_log_reader.py | 2 +-
airflow-core/tests/unit/utils/test_cli_util.py | 3 +-
airflow-core/tests/unit/utils/test_db_cleanup.py | 2 +-
airflow-core/tests/unit/utils/test_dot_renderer.py | 3 +-
airflow-core/tests/unit/utils/test_helpers.py | 3 +-
airflow-core/tests/unit/utils/test_serve_logs.py | 2 +-
airflow-core/tests/unit/utils/test_sqlalchemy.py | 2 +-
.../test_task_handler_with_custom_formatter.py | 2 +-
.../airflow_breeze/utils/docker_command_utils.py | 19 ++---
devel-common/src/tests_common/pytest_plugin.py | 34 +++++---
.../test_kubernetes_pod_operator.py | 3 +-
.../unit/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 ++++++++++++++++++++++
.../timezones/pyproject.toml | 47 +++++------
.../src/airflow_shared/timezones}/__init__.py | 9 ---
.../src/airflow_shared/timezones}/timezone.py | 13 +---
.../timezones/tests}/test_timezone.py | 4 +-
task-sdk/pyproject.toml | 3 +
.../src/airflow/sdk/_shared}/__init__.py | 9 ---
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 +-
task-sdk/src/airflow/sdk/execution_time/cache.py | 2 +-
.../src/airflow/sdk/execution_time/task_runner.py | 2 +-
.../src/airflow/sdk/timezone.py | 26 +++++--
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 +-
task-sdk/tests/task_sdk/docs/test_public_api.py | 1 +
.../tests/task_sdk/execution_time/test_comms.py | 2 +-
.../tests/task_sdk/execution_time/test_context.py | 3 +-
.../tests/task_sdk/execution_time/test_hitl.py | 2 +-
.../task_sdk/execution_time/test_supervisor.py | 10 +--
.../task_sdk/execution_time/test_task_runner.py | 3 +-
178 files changed, 451 insertions(+), 321 deletions(-)
diff --git a/.dockerignore b/.dockerignore
index 75e6291445a..716c737ee4e 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 ef80b60bcb5..e1bd6b784fc 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 e36e1fa40bb..5604382e4d1 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/tests/unit/models/__init__.py b/airflow-core/src/airflow/_shared/__init__.py
similarity index 79%
copy from airflow-core/tests/unit/models/__init__.py
copy to airflow-core/src/airflow/_shared/__init__.py
index ded9a4f4550..13a83393a91 100644
--- a/airflow-core/tests/unit/models/__init__.py
+++ b/airflow-core/src/airflow/_shared/__init__.py
@@ -1,4 +1,3 @@
-#
# 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
@@ -15,11 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-from __future__ import annotations
-
-import pathlib
-
-from airflow.utils import timezone
-
-DEFAULT_DATE = timezone.datetime(2016, 1, 1)
-TEST_DAGS_FOLDER = pathlib.Path(__file__).parent.with_name("dags")
diff --git a/airflow-core/src/airflow/_shared/timezones b/airflow-core/src/airflow/_shared/timezones
new file mode 120000
index 00000000000..8d7034fa713
--- /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 186ce9499a7..d1cf5d93b8e 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 @@ from __future__ import annotations
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 f653532e29f..276ae17153d 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.backends import default_backend
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 767b9fed46d..e8af5380f8a 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 pydantic import AfterValidator, BaseModel, NonNegativeInt
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.pool import Pool
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 18e5dc7387d..ccbb9d9ad98 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 @@ from pydantic import (
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 bb79fa671ed..686e21b444a 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 typing import TYPE_CHECKING
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 6089d6d55df..6ff8b83825f 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 typing import Generic, Literal, TypeVar
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 61fb5f9299a..ace75342932 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 fastapi import Depends, HTTPException, status
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 @@ from airflow.models.asset import (
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 525b82c976b..30efea4d79a 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 pydantic import NonNegativeInt
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 @@ from airflow.models.backfill import (
_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 aea8c8433cd..dfa8d039ac1 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 fastapi import Depends, HTTPException, status
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.openapi.exceptions import create_openapi_http_
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 f3c01983641..9b552c9cf58 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 fastapi import Depends, status
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.api_fastapi.core_api.security import requires_access_dag
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 de51a1dd27f..28b7f32ead9 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 pendulum import DateTime
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.models.dagrun import DagRun
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 def4dd155c6..09ef33d10b1 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.orm import joinedload
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.models.xcom import XComModel
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 2006b5f42b9..5bdccd65444 100644
--- a/airflow-core/src/airflow/cli/cli_config.py
+++ b/airflow-core/src/airflow/cli/cli_config.py
@@ -29,12 +29,12 @@ from typing import NamedTuple
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 f23c52321f9..1fbba1d7e45 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 typing import TYPE_CHECKING
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.jobs.job import Job
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 61a8f81c86d..656ba480ced 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 contextlib import redirect_stdout
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.sdk.execution_time.secrets_masker import RedactedIO
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 275bfbb7b07..4c7e6409e53 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 3410fa4964d..805214a4861 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 import delete, func, insert, select, tuple_, update
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.sdk.definitions.asset import Asset, AssetAlias, AssetNameRef, Asset
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 21ebffeb628..4557e52a465 100644
--- a/airflow-core/src/airflow/dag_processing/manager.py
+++ b/airflow-core/src/airflow/dag_processing/manager.py
@@ -48,6 +48,7 @@ from tabulate import tabulate
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 import SecretCache
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 @@ class DagFileProcessorManager(LoggingMixin):
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 72db0cb684e..67fb3be0192 100644
--- a/airflow-core/src/airflow/jobs/job.py
+++ b/airflow-core/src/airflow/jobs/job.py
@@ -27,6 +27,7 @@ from sqlalchemy.exc import OperationalError
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.listeners.listener import get_listener_manager
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 1157c3b51f8..53f91991e89 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.orm import joinedload, lazyload, load_only, make_transient, sele
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.ti_deps.dependencies_states import EXECUTION_STATES
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 8bacde4913d..6a96ea161cb 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 pydantic import BaseModel, Field, TypeAdapter
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.sdk.execution_time.supervisor import WatchedSubprocess, make_buffer
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 2cad82e1e7f..184c0a98218 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 @@ import sqlalchemy as sa
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 a345ebe7995..8b28f2bea61 100644
--- a/airflow-core/src/airflow/models/asset.py
+++ b/airflow-core/src/airflow/models/asset.py
@@ -36,9 +36,9 @@ from sqlalchemy import (
)
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 e1fd0e506bf..4142265b8bd 100644
--- a/airflow-core/src/airflow/models/backfill.py
+++ b/airflow-core/src/airflow/models/backfill.py
@@ -42,10 +42,10 @@ from sqlalchemy.exc import IntegrityError
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 43fc62b098d..5c72eb83665 100644
--- a/airflow-core/src/airflow/models/dag.py
+++ b/airflow-core/src/airflow/models/dag.py
@@ -61,6 +61,7 @@ from sqlalchemy.orm import backref, load_only, relationship
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 @@ from airflow.timetables.simple import (
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 5c1068d6a96..a51f3cb3018 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 import Column, ForeignKey, Integer, UniqueConstraint, select
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 83fc47444c1..a6ee7403b04 100644
--- a/airflow-core/src/airflow/models/dagbag.py
+++ b/airflow-core/src/airflow/models/dagbag.py
@@ -40,6 +40,7 @@ from sqlalchemy import (
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.exceptions import (
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 2db338f87b5..a68885e1dfd 100644
--- a/airflow-core/src/airflow/models/dagcode.py
+++ b/airflow-core/src/airflow/models/dagcode.py
@@ -26,10 +26,10 @@ from sqlalchemy.orm import relationship
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 cfac1936d6f..b910e00f9b5 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -60,6 +60,7 @@ from sqlalchemy.sql.expression import false, select
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.stats import Stats
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 b88158683ea..7498e615e30 100644
--- a/airflow-core/src/airflow/models/dagwarning.py
+++ b/airflow-core/src/airflow/models/dagwarning.py
@@ -22,9 +22,9 @@ from typing import TYPE_CHECKING
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 ada61d47c7a..f1009ca1bab 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 typing import TYPE_CHECKING
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 e1df89fe7b3..a9e8cdabbf0 100644
--- a/airflow-core/src/airflow/models/deadline.py
+++ b/airflow-core/src/airflow/models/deadline.py
@@ -29,11 +29,11 @@ from sqlalchemy.exc import SQLAlchemyError
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 1dcaf83941a..9347e4a1753 100644
--- a/airflow-core/src/airflow/models/log.py
+++ b/airflow-core/src/airflow/models/log.py
@@ -22,8 +22,8 @@ from typing import TYPE_CHECKING
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 ade8dca8761..e64a0e5cc89 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.orm import backref, foreign, relationship
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.sdk.definitions.asset import AssetUniqueKey
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 ca46b480256..9059c6d3986 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -70,6 +70,7 @@ from sqlalchemy.orm.attributes import NO_VALUE, set_committed_value
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.settings import task_instance_mutation_hook
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 cbd83643fdb..b5932b28c58 100644
--- a/airflow-core/src/airflow/models/taskinstancehistory.py
+++ b/airflow-core/src/airflow/models/taskinstancehistory.py
@@ -38,8 +38,8 @@ from sqlalchemy.ext.mutable import MutableDict
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 d55eb94a266..d9a5c57c30a 100644
--- a/airflow-core/src/airflow/models/tasklog.py
+++ b/airflow-core/src/airflow/models/tasklog.py
@@ -19,8 +19,8 @@ from __future__ import annotations
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 4226bb2c3d9..ded88d1d2f8 100644
--- a/airflow-core/src/airflow/models/trigger.py
+++ b/airflow-core/src/airflow/models/trigger.py
@@ -28,12 +28,12 @@ from sqlalchemy import Column, Integer, String, Text, delete, func, or_, select,
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 badb86bd43d..a7301578ed2 100644
--- a/airflow-core/src/airflow/models/xcom.py
+++ b/airflow-core/src/airflow/models/xcom.py
@@ -39,8 +39,8 @@ from sqlalchemy.dialects import postgresql
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 8c75249ac98..7c9a6b254a9 100644
--- a/airflow-core/src/airflow/serialization/serialized_objects.py
+++ b/airflow-core/src/airflow/serialization/serialized_objects.py
@@ -40,6 +40,7 @@ from dateutil import relativedelta
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.docs import get_docs_url
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 806ba9feb69..b5fe17a2e8b 100644
--- a/airflow-core/src/airflow/serialization/serializers/datetime.py
+++ b/airflow-core/src/airflow/serialization/serializers/datetime.py
@@ -19,12 +19,12 @@ from __future__ import annotations
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 c4d1c30142c..3a67dd8a95e 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 145b4819ac3..40634852eb6 100644
--- a/airflow-core/src/airflow/settings.py
+++ b/airflow-core/src/airflow/settings.py
@@ -36,12 +36,12 @@ from sqlalchemy.orm import scoped_session, sessionmaker
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 90954f29f2f..1013fecdfb8 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 abe0d38c2b7..501b1574205 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 0f996e5e9f4..396f8180d1b 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 a0d7e5091c9..e62f96de770 100644
--- a/airflow-core/src/airflow/timetables/_cron.py
+++ b/airflow-core/src/airflow/timetables/_cron.py
@@ -22,9 +22,9 @@ from typing import TYPE_CHECKING
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 7203cd40631..acdf6aa704e 100644
--- a/airflow-core/src/airflow/timetables/_delta.py
+++ b/airflow-core/src/airflow/timetables/_delta.py
@@ -20,8 +20,8 @@ from __future__ import annotations
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 ac2c0b61319..d8e70626d40 100644
--- a/airflow-core/src/airflow/timetables/events.py
+++ b/airflow-core/src/airflow/timetables/events.py
@@ -22,8 +22,8 @@ from typing import TYPE_CHECKING
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 7e68b9a01e6..e248dd67da2 100644
--- a/airflow-core/src/airflow/timetables/interval.py
+++ b/airflow-core/src/airflow/timetables/interval.py
@@ -22,10 +22,10 @@ from typing import TYPE_CHECKING, Any
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 bd9408434e8..b5b6f2468f3 100644
--- a/airflow-core/src/airflow/timetables/simple.py
+++ b/airflow-core/src/airflow/timetables/simple.py
@@ -19,8 +19,8 @@ from __future__ import annotations
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 87e52a3f264..ed2f65cf245 100644
--- a/airflow-core/src/airflow/timetables/trigger.py
+++ b/airflow-core/src/airflow/timetables/trigger.py
@@ -23,10 +23,10 @@ import operator
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 982e8826e27..22fb1d29351 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 import Link, NonRecordingSpan, SpanContext, TraceFlags,
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 563d133eb2b..5f3fb979113 100644
--- a/airflow-core/src/airflow/utils/__init__.py
+++ b/airflow-core/src/airflow/utils/__init__.py
@@ -34,6 +34,11 @@ __deprecated_classes = {
"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 ac7723d2933..8f1fc0748bb 100644
--- a/airflow-core/src/airflow/utils/cli.py
+++ b/airflow-core/src/airflow/utils/cli.py
@@ -34,11 +34,12 @@ from pathlib import Path
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 cd53f2c91a2..dd3d8d59fa8 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 96e803775f8..7b0ee0487dd 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.ext.compiler import compiles
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 5dbd034cea8..ea3151be1db 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 datetime import datetime
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 6f69d3c5471..124da295bb0 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 0c9571c8257..b7e69d03a73 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 @@ from __future__ import annotations
import logging
-from airflow.utils import timezone
-
class TimezoneAware(logging.Formatter):
"""
@@ -41,7 +39,9 @@ class TimezoneAware(logging.Formatter):
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 d08307db14b..072f767335c 100644
--- a/airflow-core/src/airflow/utils/sqlalchemy.py
+++ b/airflow-core/src/airflow/utils/sqlalchemy.py
@@ -31,9 +31,9 @@ from sqlalchemy.dialects import mysql
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 99e99c39418..03f427227c1 100644
--- a/airflow-core/tests/integration/otel/test_otel.py
+++ b/airflow-core/tests/integration/otel/test_otel.py
@@ -25,13 +25,13 @@ import time
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 6afef4db206..14b6957ee65 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 jwt
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 @@ from airflow.api_fastapi.auth.tokens import (
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 b674ddb9a8b..051ffa965cc 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 @@ from datetime import datetime, timezone
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 ce6ff10b301..05a879c2f5b 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 @@ from unittest import mock
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.asset import (
)
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 703595b96e7..b3eda51c127 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 pendulum
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 dbe55a92e78..6bde1fc211c 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 pytest
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.models.asset import AssetEvent, AssetModel
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 69e9a65ac24..f1702a220db 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 @@ from datetime import datetime, timedelta
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 4d84485d48a..780330c0545 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 os
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 695dfcae190..0c50be6223d 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 @@ from unittest import mock
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 adb5aebcb21..3a4e7c8ae71 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 @@ import pytest
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 a3cf5012e49..5175278d777 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 @@ from unittest import mock
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 b2e40ef5d5f..3a0ba67a50a 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 pendulum
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.taskmap import TaskMap
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 5594f89f6f0..1c4194c178b 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 @@ from unittest import mock
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.models.xcom import XComModel
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 dd6a0adb7b5..7d1f06a03f3 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 @@ from unittest import mock
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 9db616f4dbd..96808bced38 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 pendulum
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 1587797256a..aa8bb0a842b 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 @@ import pytest
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.operators.trigger_dagrun import TriggerDagRunOpe
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 ae8488e5da9..4bc313d85d9 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 @@ from __future__ import annotations
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 2cf34f8dd7b..81ff6157262 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 @@ from __future__ import annotations
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 ac414f53ee8..bf6a5ba581d 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 @@ from __future__ import annotations
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 d18b18b673e..94a5e96ab34 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 @@ import uuid6
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.taskinstance import TaskInstance
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 5757a3f6465..661fd93157e 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 @@ from unittest.mock import patch
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 d04145ff6d5..08c1a9a666d 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 @@ from datetime import datetime
import pytest
+from airflow._shared.timezones import timezone
from airflow.callbacks.callback_requests import (
DagCallbackRequest,
TaskCallbackRequest,
@@ -28,7 +29,6 @@ from airflow.callbacks.callback_requests import (
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 1c87d8c2ef8..13da9fe66ac 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 pendulum
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 219785fc7d0..630745a7022 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 @@ import time_machine
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.providers.standard.triggers.temporal import DateTimeTrigger, TimeDe
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 @@ class TestCliDags:
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 6e25253eb40..0b56a4a627c 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 @@ from unittest import mock
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.dag_version import DagVersion
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 cfdb9dee323..f78346616d1 100644
--- a/airflow-core/tests/unit/core/test_core.py
+++ b/airflow-core/tests/unit/core/test_core.py
@@ -23,11 +23,11 @@ from time import sleep
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 c7f06533b8b..62c8effb5b2 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 @@ from pathlib import Path
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 33305c64a39..723a9dc0ce5 100644
--- a/airflow-core/tests/unit/core/test_sentry.py
+++ b/airflow-core/tests/unit/core/test_sentry.py
@@ -26,8 +26,8 @@ import time_machine
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 14986486c88..7da3baa3b9f 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 @@ from unittest.mock import patch
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 5254dfc6ec9..0c812780ce4 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 import func, select
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.operators.empty import EmptyOperator
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 02121f0194e..c218c54eae7 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 @@ import time_machine
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.dag_version import DagVersion
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 96a5db7281b..6c6f31a1ff8 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 @@ import structlog
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 import DAG
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 378e2497f28..7d4af553a7b 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 __future__ import annotations
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 e4d30f78886..fe422692d1f 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 484181688c9..2d4ba261af6 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 @@ from __future__ import annotations
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 1990c8c03c6..47674815aec 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 @@ import time
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 1e73b76776f..3de4fb6a627 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 @@ from __future__ import annotations
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 0a0ab4a4a7d..676947ae157 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 __future__ import annotations
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 @@ from airflow.sdk.definitions.context import (
_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 da89e821986..9123d32552f 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 __future__ import annotations
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 4feb1caa1eb..61504cfb775 100644
--- a/airflow-core/tests/unit/dags/test_sensor.py
+++ b/airflow-core/tests/unit/dags/test_sensor.py
@@ -18,9 +18,9 @@ from __future__ import annotations
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 f1940f6458d..1e148301cf6 100644
--- a/airflow-core/tests/unit/decorators/test_mapped.py
+++ b/airflow-core/tests/unit/decorators/test_mapped.py
@@ -19,8 +19,8 @@ from __future__ import annotations
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 2ed0657e3c6..9715138c74f 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 pendulum
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.base_executor import BaseExecutor, RunningRetryAttemptTyp
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 e5b37bb268d..ec7102e2e70 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 @@ import pytest
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 c4b46a7622e..9f80073007c 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 @@ from unittest.mock import ANY, Mock, patch
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 f9269226e9c..ec1580528f9 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 import func, select, update
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.sdk.definitions.asset import Asset, AssetAlias
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 tests_common.test_utils.version_compat import SQLALCHEMY_V_1_4, SQLALCHEMY_
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 @@ class TestSchedulerJob:
"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 194c1c8752b..384bf6a76fc 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 @@ import pytest
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.providers.standard.triggers.temporal import DateTimeTrigger, TimeDe
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 9a4c85cb504..fac237c7e17 100644
--- a/airflow-core/tests/unit/listeners/test_listeners.py
+++ b/airflow-core/tests/unit/listeners/test_listeners.py
@@ -22,11 +22,11 @@ import os
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 ded9a4f4550..d5fe7d8f930 100644
--- a/airflow-core/tests/unit/models/__init__.py
+++ b/airflow-core/tests/unit/models/__init__.py
@@ -19,7 +19,7 @@ from __future__ import annotations
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 109d2145941..9d26e31459e 100644
--- a/airflow-core/tests/unit/models/test_backfill.py
+++ b/airflow-core/tests/unit/models/test_backfill.py
@@ -25,6 +25,7 @@ import pendulum
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.models.backfill import (
)
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 dab3d4d9916..7f2f4b01394 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -34,6 +34,8 @@ import time_machine
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 @@ from airflow.timetables.simple import (
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 00aa6151f97..230ddca9f89 100644
--- a/airflow-core/tests/unit/models/test_dagbag.py
+++ b/airflow-core/tests/unit/models/test_dagbag.py
@@ -35,12 +35,12 @@ import time_machine
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 50440d4f44e..cf4dafb4ee0 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 import select
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.sdk.definitions.deadline import DeadlineAlert, DeadlineReference
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 ad0c935d618..81d4ce66dc7 100644
--- a/airflow-core/tests/unit/models/test_pool.py
+++ b/airflow-core/tests/unit/models/test_pool.py
@@ -20,12 +20,12 @@ from __future__ import annotations
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 e53582c563e..8acf7fe932f 100644
--- a/airflow-core/tests/unit/models/test_renderedtifields.py
+++ b/airflow-core/tests/unit/models/test_renderedtifields.py
@@ -29,6 +29,7 @@ import pytest
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.bash import BashOperator
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 7e041917056..2a9e86771c2 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 import select
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.dependencies_states import RUNNABLE_STATES
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 5aaa956829d..529fd32afc5 100644
--- a/airflow-core/tests/unit/models/test_timestamp.py
+++ b/airflow-core/tests/unit/models/test_timestamp.py
@@ -20,9 +20,9 @@ import pendulum
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 8818036418e..3f5566e3af0 100644
--- a/airflow-core/tests/unit/models/test_trigger.py
+++ b/airflow-core/tests/unit/models/test_trigger.py
@@ -27,6 +27,7 @@ import pytest
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 @@ from airflow.triggers.base import (
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 566d74ceaa7..6ed5d95e83d 100644
--- a/airflow-core/tests/unit/models/test_xcom.py
+++ b/airflow-core/tests/unit/models/test_xcom.py
@@ -25,6 +25,7 @@ from unittest.mock import MagicMock
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.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.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 30afb69db2b..6d4f1d82bb4 100644
--- a/airflow-core/tests/unit/sensors/test_filesystem.py
+++ b/airflow-core/tests/unit/sensors/test_filesystem.py
@@ -24,11 +24,11 @@ from datetime import timedelta
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 c61284c7850..234bb51fbb2 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 dateutil.relativedelta import FR, relativedelta
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.task.priority_strategy import _DownstreamPriorityWeightStrategy
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 dc944035c75..dec42d3568b 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 kubernetes.client import models as k8s
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.enums import DagAttributeTypes as DAT, Encoding
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 3b8aec4763e..bde50ad31f8 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 @@ from unittest.mock import Mock
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 d2fd53da745..5ffc2f2be43 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 @@ from unittest.mock import ANY, Mock, patch
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 7b4bd22aa04..355cbfb9c30 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 pytest
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 94b9cfc2ac6..7b498b71f83 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 @@ from unittest.mock import Mock
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 13a27daa199..556f5d725b4 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 pendulum
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 d6e2d56df54..87a23ba24a7 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 pendulum
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 a904e8372a7..0ca84229ba6 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 datetime
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 1a98044bf4e..3a0a60cebfe 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 pendulum
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 @@ from airflow.timetables.trigger import (
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 d8ccf5809db..7c91d16aeb4 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 datetime
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 f50fc619ac5..ad0982dbea4 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 @@ from datetime import timedelta
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 cd8b430090a..8d5ffab543d 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 pendulum
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 2d5578f8faa..54ecc506e27 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 pytest
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 48a0c848621..71c481ea1c6 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.exc import OperationalError, SQLAlchemyError
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 240876ec7f4..1e239271dfc 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 @@ from unittest import mock
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 71138aa5105..08c8d8092a0 100644
--- a/airflow-core/tests/unit/utils/test_helpers.py
+++ b/airflow-core/tests/unit/utils/test_helpers.py
@@ -23,9 +23,10 @@ from typing import TYPE_CHECKING
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 fbbef2b2000..b6c7805386a 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 pytest
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 602a46a8ab1..9ff19c7c0e6 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 import text
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 @@ from airflow.utils.sqlalchemy import (
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 5e551904798..d3e24448f54 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 logging
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 ed9e9dd7a2f..b959d53fc04 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 @@ VOLUMES_FOR_SELECTED_MOUNTS = [
("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 f02ab53fa6d..e10bcd1d5bd 100644
--- a/devel-common/src/tests_common/pytest_plugin.py
+++ b/devel-common/src/tests_common/pytest_plugin.py
@@ -1012,10 +1012,11 @@ def dag_maker(request) -> Generator[DagMaker, None, None]:
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 dag_maker(request) -> Generator[DagMaker, None, None]:
):
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 create_task_instance(dag_maker: DagMaker, create_dummy_dag: CreateDummyDAG)
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_test_dag():
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 create_runtime_ti(mocked_parse):
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 run_task(create_runtime_ti, mock_supervisor_comms, spy_agency) -> RunTaskCal
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_connection_without_db(monkeypatch):
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 59c52178b1d..6b51627146c 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 import mock
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 2168744edca..85a222571e4 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 @@ import pytest
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 tests_common.test_utils.operators.run_deferrable import execute_operator
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 bbc8c0401c0..1cec8de4a48 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.exceptions import TaskDeferred
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 @@ class TestTimeSensor:
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 @@ class TestTimeSensor:
],
)
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 @@ class TestTimeSensor:
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 @@ class TestTimeSensor:
# 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 02768457878..6ad8960afdf 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 6609aa5c1db..2ac4dc849ad 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 317aca4d188..fd99a833802 100644
--- a/scripts/cov/core_coverage.py
+++ b/scripts/cov/core_coverage.py
@@ -143,7 +143,6 @@ files_not_fully_covered = [
"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 ec6b82d7e0f..54a4724ba3d 100755
--- a/setup_idea.py
+++ b/setup_idea.py
@@ -94,12 +94,14 @@ IDEA_FOLDER_PATH = ROOT_AIRFLOW_FOLDER_PATH / ".idea"
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 00000000000..baa57598546
--- /dev/null
+++ b/shared/README.md
@@ -0,0 +1,91 @@
+<!--
+ 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.
+ -->
+
+# 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 [...]
+
+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/airflow-core/tests/unit/dags/test_logging_in_dag.py b/shared/timezones/pyproject.toml
similarity index 55%
copy from airflow-core/tests/unit/dags/test_logging_in_dag.py
copy to shared/timezones/pyproject.toml
index 484181688c9..bc68f6467c8 100644
--- a/airflow-core/tests/unit/dags/test_logging_in_dag.py
+++ b/shared/timezones/pyproject.toml
@@ -1,4 +1,3 @@
-#
# 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
@@ -15,32 +14,34 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-from __future__ import annotations
-
-import logging
-
-from airflow.models.dag import DAG
-from airflow.providers.standard.operators.python import PythonOperator
-from airflow.utils.timezone import datetime
-logger = logging.getLogger(__name__)
+[project]
+name = "apache-airflow-shared-timezones"
+description = "Shared timezone code for Airflow distributions"
+version = "0.0"
+classifiers = [
+ "Private :: Do Not Upload",
+]
+dependencies = [
+ "pendulum",
+]
-def test_logging_fn(**kwargs):
- """
- Tests DAG logging.
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
- :param kwargs:
- """
- logger.info("Log from DAG Logger")
- kwargs["ti"].log.info("Log from TI Logger")
- print("Log from Print statement")
+[tool.hatch.build.targets.wheel]
+packages = ["src/airflow_shared"]
+[tool.ruff]
+extend = "../../pyproject.toml"
+src = ["src"]
-dag = DAG(dag_id="test_logging_dag", schedule=None, start_date=datetime(2016, 1, 1))
+[tool.ruff.lint.per-file-ignores]
+# Ignore Doc rules et al for anything outside of tests
+"!src/*" = ["D", "S101", "TRY002"]
-PythonOperator(
- task_id="test_task",
- python_callable=test_logging_fn,
- dag=dag,
-)
+[tool.ruff.lint.flake8-tidy-imports]
+# Override the workspace level default
+ban-relative-imports = "parents"
diff --git a/airflow-core/tests/unit/models/__init__.py b/shared/timezones/src/airflow_shared/timezones/__init__.py
similarity index 79%
copy from airflow-core/tests/unit/models/__init__.py
copy to shared/timezones/src/airflow_shared/timezones/__init__.py
index ded9a4f4550..13a83393a91 100644
--- a/airflow-core/tests/unit/models/__init__.py
+++ b/shared/timezones/src/airflow_shared/timezones/__init__.py
@@ -1,4 +1,3 @@
-#
# 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
@@ -15,11 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-from __future__ import annotations
-
-import pathlib
-
-from airflow.utils import timezone
-
-DEFAULT_DATE = timezone.datetime(2016, 1, 1)
-TEST_DAGS_FOLDER = pathlib.Path(__file__).parent.with_name("dags")
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 3e59169598a..3abdd06ee5e 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 c7e6d0c09aa..ba0f2a46def 100644
--- a/airflow-core/tests/unit/utils/test_timezone.py
+++ b/shared/timezones/tests/test_timezone.py
@@ -23,8 +23,8 @@ import pendulum
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 a6edccc35c6..de87a4a5f31 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/airflow-core/tests/unit/models/__init__.py b/task-sdk/src/airflow/sdk/_shared/__init__.py
similarity index 79%
copy from airflow-core/tests/unit/models/__init__.py
copy to task-sdk/src/airflow/sdk/_shared/__init__.py
index ded9a4f4550..13a83393a91 100644
--- a/airflow-core/tests/unit/models/__init__.py
+++ b/task-sdk/src/airflow/sdk/_shared/__init__.py
@@ -1,4 +1,3 @@
-#
# 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
@@ -15,11 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-from __future__ import annotations
-
-import pathlib
-
-from airflow.utils import timezone
-
-DEFAULT_DATE = timezone.datetime(2016, 1, 1)
-TEST_DAGS_FOLDER = pathlib.Path(__file__).parent.with_name("dags")
diff --git a/task-sdk/src/airflow/sdk/_shared/timezones b/task-sdk/src/airflow/sdk/_shared/timezones
new file mode 120000
index 00000000000..b66387c7a6b
--- /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 42bc66ee5e4..8ca3050655b 100644
--- a/task-sdk/src/airflow/sdk/bases/decorator.py
+++ b/task-sdk/src/airflow/sdk/bases/decorator.py
@@ -28,6 +28,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Generic, ParamSpec, Protocol, T
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._internal.types import NOTSET
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 6f0078eef53..85f6b0745b1 100644
--- a/task-sdk/src/airflow/sdk/bases/operator.py
+++ b/task-sdk/src/airflow/sdk/bases/operator.py
@@ -36,6 +36,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, TypeVar, cast
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 @@ from airflow.task.priority_strategy import (
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 a39de2a1518..0ca447189a9 100644
--- a/task-sdk/src/airflow/sdk/bases/sensor.py
+++ b/task-sdk/src/airflow/sdk/bases/sensor.py
@@ -36,8 +36,8 @@ from airflow.exceptions import (
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 3508efc8663..05bd34e95bf 100644
--- a/task-sdk/src/airflow/sdk/definitions/dag.py
+++ b/task-sdk/src/airflow/sdk/definitions/dag.py
@@ -217,7 +217,7 @@ else:
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 @@ class DAG:
)
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 @@ class 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 @@ class DAG:
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 9902a05f032..392b03d4215 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 datetime
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 c937c680cd4..be680a08d04 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 @@ from airflow.sdk.execution_time.context import (
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/airflow-core/tests/unit/models/__init__.py b/task-sdk/src/airflow/sdk/timezone.py
similarity index 67%
copy from airflow-core/tests/unit/models/__init__.py
copy to task-sdk/src/airflow/sdk/timezone.py
index ded9a4f4550..1769f510c58 100644
--- a/airflow-core/tests/unit/models/__init__.py
+++ b/task-sdk/src/airflow/sdk/timezone.py
@@ -1,4 +1,3 @@
-#
# 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
@@ -15,11 +14,26 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-from __future__ import annotations
-import pathlib
+from __future__ import annotations
-from airflow.utils import timezone
+# 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,
+)
-DEFAULT_DATE = timezone.datetime(2016, 1, 1)
-TEST_DAGS_FOLDER = pathlib.Path(__file__).parent.with_name("dags")
+__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 5df79afd8e7..30cfd2fd7d6 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 @@ import uuid6
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 @@ from airflow.sdk.execution_time.comms import (
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 d0e2b430d2c..8b9cc04a786 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 __future__ import annotations
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.exceptions import (
)
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 @@ class TestBaseSensor:
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 @@ class TestBaseSensor:
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 @@ class TestBaseSensor:
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 15c2d65cc27..9f3d8a343d2 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 @@ from __future__ import annotations
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 dd2bf54980c..f1e2b74e47f 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 7960fc7b275..2f5d6bceedc 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 @@ from socket import socketpair
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 d3c22598a98..1def08086cc 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 @@ from unittest.mock import MagicMock, patch
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 @@ from airflow.sdk.execution_time.context import (
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 cab17e30bac..a936812bbb7 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 __future__ import annotations
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 @@ from airflow.sdk.execution_time.hitl import (
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 c6aa651af44..271f2b98c15 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 task_sdk import FAKE_BUNDLE, make_client
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 @@ from airflow.sdk.execution_time.supervisor import (
set_supervisor_comms,
supervise,
)
-from airflow.utils import timezone, timezone as tz
if TYPE_CHECKING:
import kgb
@@ -237,7 +237,7 @@ class TestWatchedSubprocess:
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 @@ class TestWatchedSubprocess:
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 @@ class TestWatchedSubprocess:
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 @@ class TestWatchedSubprocess:
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 e0b414065cb..a88b244897d 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 @@ from airflow.sdk import (
dag as dag_decorator,
get_current_context,
task as task_decorator,
+ timezone,
)
from airflow.sdk.api.datamodels._generated import (
AssetProfile,
@@ -121,7 +122,6 @@ from airflow.sdk.execution_time.task_runner import (
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 @@ class TestRuntimeTaskInstance:
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")