From efbe2a2018303f2f8f361d43925b4a3ccdcf4677 Mon Sep 17 00:00:00 2001 From: Ash Berlin-Taylor Date: Wed, 16 Jul 2025 08:39:36 -0400 Subject: [PATCH] Set up process for sharing code between different components And move `airflow.utils.timezone` into a shared library as the first example of it working. In this change we have now setled on an approach using symlinks, but we did explore other options (see the GH PR for discussion and previous versions, notably one built upon the `vendoring` tool) A lot of the reasoning and mode of operation of this is detailed in shared/README.md in this PR, hence why this description is so short. Currently various places in TaskSDK and Airflow Core both use these utility functions, and while in this specific case they are small enough that they could just be copied and the duplication wouldn't hurt us long term, this changes shows a way in which we can have a single source of truth, but have it included automatically in built dists. Co-authored-by: Jarek Potiuk --- .dockerignore | 1 + airflow-core/docs/img/airflow_erd.sha256 | 2 +- airflow-core/docs/img/airflow_erd.svg | 1718 ++++++++--------- airflow-core/pyproject.toml | 3 + airflow-core/src/airflow/_shared/__init__.py | 16 + airflow-core/src/airflow/_shared/timezones | 1 + .../src/airflow/api/common/trigger_dag.py | 2 +- .../src/airflow/api_fastapi/auth/tokens.py | 2 +- .../airflow/api_fastapi/common/parameters.py | 2 +- .../src/airflow/api_fastapi/common/types.py | 2 +- .../core_api/datamodels/dag_run.py | 2 +- .../core_api/datamodels/ui/common.py | 2 +- .../core_api/routes/public/assets.py | 2 +- .../core_api/routes/public/backfills.py | 2 +- .../core_api/routes/public/hitl.py | 2 +- .../core_api/routes/ui/dashboard.py | 2 +- .../core_api/services/ui/calendar.py | 2 +- .../execution_api/routes/task_instances.py | 2 +- .../src/airflow/cli/commands/dag_command.py | 3 +- .../src/airflow/cli/commands/task_command.py | 3 +- .../src/airflow/dag_processing/manager.py | 2 +- airflow-core/src/airflow/jobs/job.py | 2 +- .../src/airflow/jobs/scheduler_job_runner.py | 2 +- .../src/airflow/jobs/triggerer_job_runner.py | 2 +- airflow-core/src/airflow/logging/_shared | 1 + .../versions/0047_3_0_0_add_dag_versioning.py | 2 +- airflow-core/src/airflow/models/asset.py | 2 +- airflow-core/src/airflow/models/backfill.py | 2 +- airflow-core/src/airflow/models/dag.py | 2 +- .../src/airflow/models/dag_version.py | 2 +- airflow-core/src/airflow/models/dagbag.py | 2 +- airflow-core/src/airflow/models/dagcode.py | 2 +- airflow-core/src/airflow/models/dagrun.py | 2 +- airflow-core/src/airflow/models/dagwarning.py | 2 +- .../src/airflow/models/db_callback_request.py | 2 +- airflow-core/src/airflow/models/deadline.py | 2 +- airflow-core/src/airflow/models/log.py | 2 +- .../src/airflow/models/serialized_dag.py | 2 +- .../src/airflow/models/taskinstance.py | 2 +- .../src/airflow/models/taskinstancehistory.py | 2 +- airflow-core/src/airflow/models/tasklog.py | 2 +- airflow-core/src/airflow/models/trigger.py | 2 +- airflow-core/src/airflow/models/xcom.py | 2 +- .../ti_deps/deps/not_in_retry_period_dep.py | 2 +- .../ti_deps/deps/ready_to_reschedule.py | 2 +- .../ti_deps/deps/runnable_exec_date_dep.py | 2 +- airflow-core/src/airflow/timetables/events.py | 2 +- airflow-core/src/airflow/timetables/simple.py | 2 +- .../src/airflow/traces/otel_tracer.py | 2 +- airflow-core/src/airflow/utils/cli.py | 3 +- airflow-core/src/airflow/utils/db_cleanup.py | 2 +- .../utils/log/file_processor_handler.py | 2 +- .../src/airflow/utils/log/timezone_aware.py | 6 +- .../tests/integration/otel/test_otel.py | 2 +- .../unit/api_fastapi/auth/test_tokens.py | 2 +- .../core_api/routes/public/test_assets.py | 2 +- .../core_api/routes/public/test_backfills.py | 2 +- .../core_api/routes/public/test_dag_run.py | 2 +- .../core_api/routes/public/test_dag_stats.py | 2 +- .../routes/public/test_extra_links.py | 2 +- .../core_api/routes/public/test_log.py | 2 +- .../core_api/routes/public/test_monitor.py | 2 +- .../core_api/routes/public/test_xcom.py | 2 +- .../core_api/routes/ui/test_backfills.py | 2 +- .../core_api/routes/ui/test_grid.py | 2 +- .../core_api/routes/ui/test_structure.py | 2 +- .../versions/head/test_asset_events.py | 2 +- .../versions/head/test_assets.py | 2 +- .../versions/head/test_dag_runs.py | 2 +- .../versions/head/test_task_instances.py | 2 +- .../v2025_04_28/test_task_instances.py | 2 +- .../unit/callbacks/test_callback_requests.py | 2 +- .../cli/commands/test_backfill_command.py | 2 +- .../unit/cli/commands/test_dag_command.py | 2 +- .../unit/cli/commands/test_task_command.py | 2 +- airflow-core/tests/unit/core/test_sentry.py | 2 +- .../unit/dag_processing/bundles/test_base.py | 2 +- .../unit/dag_processing/test_collection.py | 2 +- .../tests/unit/dag_processing/test_manager.py | 2 +- .../unit/dag_processing/test_processor.py | 2 +- .../tests/unit/dags/test_scheduler_dags.py | 2 +- airflow-core/tests/unit/dags/test_sensor.py | 2 +- .../tests/unit/decorators/test_mapped.py | 2 +- .../unit/executors/test_base_executor.py | 2 +- .../unit/executors/test_local_executor.py | 2 +- airflow-core/tests/unit/jobs/test_base_job.py | 2 +- .../tests/unit/jobs/test_scheduler_job.py | 2 +- .../tests/unit/jobs/test_triggerer_job.py | 2 +- .../tests/unit/listeners/test_listeners.py | 2 +- airflow-core/tests/unit/models/__init__.py | 2 +- .../tests/unit/models/test_backfill.py | 2 +- airflow-core/tests/unit/models/test_dag.py | 2 +- 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_taskinstance.py | 2 +- .../tests/unit/models/test_timestamp.py | 2 +- .../tests/unit/models/test_trigger.py | 2 +- airflow-core/tests/unit/models/test_xcom.py | 2 +- .../serialization/test_dag_serialization.py | 2 +- .../serialization/test_serialized_objects.py | 2 +- .../deps/test_ready_to_reschedule_dep.py | 2 +- .../unit/timetables/test_once_timetable.py | 2 +- .../utils/log/test_file_processor_handler.py | 2 +- .../tests/unit/utils/log/test_log_reader.py | 2 +- .../tests/unit/utils/test_cli_util.py | 3 +- .../tests/unit/utils/test_db_cleanup.py | 2 +- .../tests/unit/utils/test_dot_renderer.py | 3 +- airflow-core/tests/unit/utils/test_helpers.py | 3 +- .../tests/unit/utils/test_serve_logs.py | 2 +- .../tests/unit/utils/test_timezone.py | 2 +- .../utils/docker_command_utils.py | 19 +- pyproject.toml | 12 +- scripts/ci/docker-compose/local.yml | 3 + shared/timezones/pyproject.toml | 47 + .../src/airflow_shared/timezones/__init__.py | 16 + .../src/airflow_shared/timezones/timezone.py | 318 +++ shared/timezones/tests/test_timezone.py | 202 ++ task-sdk/pyproject.toml | 6 + task-sdk/src/airflow/sdk/_shared/__init__.py | 0 task-sdk/src/airflow/sdk/_shared/timezones | 1 + task-sdk/src/airflow/sdk/bases/decorator.py | 2 +- task-sdk/src/airflow/sdk/bases/operator.py | 2 +- task-sdk/src/airflow/sdk/bases/sensor.py | 2 +- task-sdk/src/airflow/sdk/definitions/dag.py | 8 +- .../src/airflow/sdk/execution_time/cache.py | 2 +- .../airflow/sdk/execution_time/task_runner.py | 2 +- task-sdk/src/airflow/sdk/timezone.py | 47 + task-sdk/tests/task_sdk/api/test_client.py | 2 +- task-sdk/tests/task_sdk/bases/test_sensor.py | 4 +- .../task_sdk/dags/super_basic_deferred_run.py | 2 +- .../task_sdk/execution_time/test_comms.py | 2 +- .../task_sdk/execution_time/test_context.py | 2 +- .../task_sdk/execution_time/test_hitl.py | 2 +- .../execution_time/test_supervisor.py | 10 +- .../execution_time/test_task_runner.py | 4 +- 136 files changed, 1678 insertions(+), 999 deletions(-) create mode 100644 airflow-core/src/airflow/_shared/__init__.py create mode 120000 airflow-core/src/airflow/_shared/timezones create mode 120000 airflow-core/src/airflow/logging/_shared create mode 100644 shared/timezones/pyproject.toml create mode 100644 shared/timezones/src/airflow_shared/timezones/__init__.py create mode 100644 shared/timezones/src/airflow_shared/timezones/timezone.py create mode 100644 shared/timezones/tests/test_timezone.py create mode 100644 task-sdk/src/airflow/sdk/_shared/__init__.py create mode 120000 task-sdk/src/airflow/sdk/_shared/timezones create mode 100644 task-sdk/src/airflow/sdk/timezone.py diff --git a/.dockerignore b/.dockerignore index 75e6291445ab6..716c737ee4e00 100644 --- a/.dockerignore +++ b/.dockerignore @@ -46,6 +46,7 @@ !docker-tests !helm-tests !kubernetes-tests +!shared/ # Add scripts so that we can use them inside the container !scripts diff --git a/airflow-core/docs/img/airflow_erd.sha256 b/airflow-core/docs/img/airflow_erd.sha256 index e935d2a08ab18..531787a5ee4e2 100644 --- a/airflow-core/docs/img/airflow_erd.sha256 +++ b/airflow-core/docs/img/airflow_erd.sha256 @@ -1 +1 @@ -2e49ab99fe1076b0f3f22a52b9ee37eeb7fc20a5a043ea504cc26022f4315277 \ No newline at end of file +2136ca0f191da2c8e0e5a40a7ff775548e197a5e793f4f3709f6eea3e5ae477c \ No newline at end of file diff --git a/airflow-core/docs/img/airflow_erd.svg b/airflow-core/docs/img/airflow_erd.svg index 2f9f9b4becc5e..7801cd7329893 100644 --- a/airflow-core/docs/img/airflow_erd.svg +++ b/airflow-core/docs/img/airflow_erd.svg @@ -4,11 +4,11 @@ - - + + %3 - + dag_priority_parsing_request @@ -692,24 +692,24 @@ dagrun_asset_event - -dagrun_asset_event - -dag_run_id - - [INTEGER] - NOT NULL - -event_id - - [INTEGER] - NOT NULL + +dagrun_asset_event + +dag_run_id + + [INTEGER] + NOT NULL + +event_id + + [INTEGER] + NOT NULL asset_event--dagrun_asset_event - -0..N + +0..N 1 @@ -752,699 +752,699 @@ task_instance - -task_instance - -id - - [UUID] - NOT NULL - -context_carrier - - [JSONB] - -custom_operator_name - - [VARCHAR(1000)] - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - NOT NULL - -duration - - [DOUBLE_PRECISION] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BYTEA] - -external_executor_id - - [VARCHAR(250)] - -hostname - - [VARCHAR(1000)] - -last_heartbeat_at - - [TIMESTAMP] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - -next_kwargs - - [JSONB] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - -queue - - [VARCHAR(256)] - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -scheduled_dttm - - [TIMESTAMP] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - -unixname - - [VARCHAR(1000)] - -updated_at - - [TIMESTAMP] + +task_instance + +id + + [UUID] + NOT NULL + +context_carrier + + [JSONB] + +custom_operator_name + + [VARCHAR(1000)] + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + NOT NULL + +duration + + [DOUBLE_PRECISION] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BYTEA] + +external_executor_id + + [VARCHAR(250)] + +hostname + + [VARCHAR(1000)] + +last_heartbeat_at + + [TIMESTAMP] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + +next_kwargs + + [JSONB] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + +queue + + [VARCHAR(256)] + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +scheduled_dttm + + [TIMESTAMP] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + +unixname + + [VARCHAR(1000)] + +updated_at + + [TIMESTAMP] trigger--task_instance - -0..N + +0..N {0,1} - + -task_map - -task_map - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -keys - - [JSONB] - -length - - [INTEGER] - NOT NULL +hitl_detail + +hitl_detail + +ti_id + + [UUID] + NOT NULL + +body + + [TEXT] + +chosen_options + + [JSON] + +defaults + + [JSON] + +multiple + + [BOOLEAN] + +options + + [JSON] + NOT NULL + +params + + [JSON] + NOT NULL + +params_input + + [JSON] + NOT NULL + +response_at + + [TIMESTAMP] + +subject + + [TEXT] + NOT NULL + +user_id + + [VARCHAR(128)] - + -task_instance--task_map - -0..N -1 +task_instance--hitl_detail + +1 +1 + + + +task_map + +task_map + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +keys + + [JSONB] + +length + + [INTEGER] + NOT NULL task_instance--task_map - -0..N -1 + +0..N +1 task_instance--task_map - -0..N -1 + +0..N +1 task_instance--task_map - -0..N -1 + +0..N +1 + + + +task_instance--task_map + +0..N +1 - + task_reschedule - -task_reschedule - -id - - [INTEGER] - NOT NULL - -duration - - [INTEGER] - NOT NULL - -end_date - - [TIMESTAMP] - NOT NULL - -reschedule_date - - [TIMESTAMP] - NOT NULL - -start_date - - [TIMESTAMP] - NOT NULL - -ti_id - - [UUID] - NOT NULL + +task_reschedule + +id + + [INTEGER] + NOT NULL + +duration + + [INTEGER] + NOT NULL + +end_date + + [TIMESTAMP] + NOT NULL + +reschedule_date + + [TIMESTAMP] + NOT NULL + +start_date + + [TIMESTAMP] + NOT NULL + +ti_id + + [UUID] + NOT NULL - + task_instance--task_reschedule - -0..N -1 + +0..N +1 - + xcom - -xcom - -dag_run_id - - [INTEGER] - NOT NULL - -key - - [VARCHAR(512)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -dag_id - - [VARCHAR(250)] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -timestamp - - [TIMESTAMP] - NOT NULL - -value - - [JSONB] - - - -task_instance--xcom - -0..N -1 + +xcom + +dag_run_id + + [INTEGER] + NOT NULL + +key + + [VARCHAR(512)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +dag_id + + [VARCHAR(250)] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +timestamp + + [TIMESTAMP] + NOT NULL + +value + + [JSONB] task_instance--xcom - -0..N -1 + +0..N +1 task_instance--xcom - -0..N -1 + +0..N +1 task_instance--xcom - -0..N -1 + +0..N +1 + + + +task_instance--xcom + +0..N +1 - + task_instance_note - -task_instance_note - -ti_id - - [UUID] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [VARCHAR(128)] + +task_instance_note + +ti_id + + [UUID] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [VARCHAR(128)] - + task_instance--task_instance_note - -1 -1 + +1 +1 - + task_instance_history - -task_instance_history - -task_instance_id - - [UUID] - NOT NULL - -context_carrier - - [JSONB] - -custom_operator_name - - [VARCHAR(1000)] - -dag_id - - [VARCHAR(250)] - NOT NULL - -dag_version_id - - [UUID] - -duration - - [DOUBLE_PRECISION] - -end_date - - [TIMESTAMP] - -executor - - [VARCHAR(1000)] - -executor_config - - [BYTEA] - -external_executor_id - - [VARCHAR(250)] - -hostname - - [VARCHAR(1000)] - -map_index - - [INTEGER] - NOT NULL - -max_tries - - [INTEGER] - -next_kwargs - - [JSONB] - -next_method - - [VARCHAR(1000)] - -operator - - [VARCHAR(1000)] - -pid - - [INTEGER] - -pool - - [VARCHAR(256)] - NOT NULL - -pool_slots - - [INTEGER] - NOT NULL - -priority_weight - - [INTEGER] - -queue - - [VARCHAR(256)] - -queued_by_job_id - - [INTEGER] - -queued_dttm - - [TIMESTAMP] - -rendered_map_index - - [VARCHAR(250)] - -run_id - - [VARCHAR(250)] - NOT NULL - -scheduled_dttm - - [TIMESTAMP] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(20)] - -task_display_name - - [VARCHAR(2000)] - -task_id - - [VARCHAR(250)] - NOT NULL - -trigger_id - - [INTEGER] - -trigger_timeout - - [TIMESTAMP] - -try_number - - [INTEGER] - NOT NULL - -unixname - - [VARCHAR(1000)] - -updated_at - - [TIMESTAMP] - - - -task_instance--task_instance_history - -0..N -1 + +task_instance_history + +task_instance_id + + [UUID] + NOT NULL + +context_carrier + + [JSONB] + +custom_operator_name + + [VARCHAR(1000)] + +dag_id + + [VARCHAR(250)] + NOT NULL + +dag_version_id + + [UUID] + +duration + + [DOUBLE_PRECISION] + +end_date + + [TIMESTAMP] + +executor + + [VARCHAR(1000)] + +executor_config + + [BYTEA] + +external_executor_id + + [VARCHAR(250)] + +hostname + + [VARCHAR(1000)] + +map_index + + [INTEGER] + NOT NULL + +max_tries + + [INTEGER] + +next_kwargs + + [JSONB] + +next_method + + [VARCHAR(1000)] + +operator + + [VARCHAR(1000)] + +pid + + [INTEGER] + +pool + + [VARCHAR(256)] + NOT NULL + +pool_slots + + [INTEGER] + NOT NULL + +priority_weight + + [INTEGER] + +queue + + [VARCHAR(256)] + +queued_by_job_id + + [INTEGER] + +queued_dttm + + [TIMESTAMP] + +rendered_map_index + + [VARCHAR(250)] + +run_id + + [VARCHAR(250)] + NOT NULL + +scheduled_dttm + + [TIMESTAMP] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(20)] + +task_display_name + + [VARCHAR(2000)] + +task_id + + [VARCHAR(250)] + NOT NULL + +trigger_id + + [INTEGER] + +trigger_timeout + + [TIMESTAMP] + +try_number + + [INTEGER] + NOT NULL + +unixname + + [VARCHAR(1000)] + +updated_at + + [TIMESTAMP] task_instance--task_instance_history - -0..N -1 + +0..N +1 task_instance--task_instance_history - -0..N -1 + +0..N +1 task_instance--task_instance_history - -0..N -1 + +0..N +1 - - -hitl_detail - -hitl_detail - -ti_id - - [UUID] - NOT NULL - -body - - [TEXT] - -chosen_options - - [JSON] - -defaults - - [JSON] - -multiple - - [BOOLEAN] - -options - - [JSON] - NOT NULL - -params - - [JSON] - NOT NULL - -params_input - - [JSON] - NOT NULL - -response_at - - [TIMESTAMP] - -subject - - [TEXT] - NOT NULL - -user_id - - [VARCHAR(128)] - - + -task_instance--hitl_detail - -1 -1 +task_instance--task_instance_history + +0..N +1 rendered_task_instance_fields - -rendered_task_instance_fields - -dag_id - - [VARCHAR(250)] - NOT NULL - -map_index - - [INTEGER] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -task_id - - [VARCHAR(250)] - NOT NULL - -k8s_pod_yaml - - [JSON] - -rendered_fields - - [JSON] - NOT NULL + +rendered_task_instance_fields + +dag_id + + [VARCHAR(250)] + NOT NULL + +map_index + + [INTEGER] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +task_id + + [VARCHAR(250)] + NOT NULL + +k8s_pod_yaml + + [JSON] + +rendered_fields + + [JSON] + NOT NULL task_instance--rendered_task_instance_fields - -0..N -1 + +0..N +1 task_instance--rendered_task_instance_fields - -0..N -1 + +0..N +1 task_instance--rendered_task_instance_fields - -0..N -1 + +0..N +1 task_instance--rendered_task_instance_fields - -0..N -1 + +0..N +1 @@ -1833,173 +1833,173 @@ deadline - -deadline - -id - - [UUID] - NOT NULL - -callback - - [VARCHAR(500)] - NOT NULL - -callback_kwargs - - [JSON] - -dag_id - - [VARCHAR(250)] - -dagrun_id - - [INTEGER] - -deadline_time - - [TIMESTAMP] - NOT NULL + +deadline + +id + + [UUID] + NOT NULL + +callback + + [VARCHAR(500)] + NOT NULL + +callback_kwargs + + [JSON] + +dag_id + + [VARCHAR(250)] + +dagrun_id + + [INTEGER] + +deadline_time + + [TIMESTAMP] + NOT NULL dag--deadline - -0..N + +0..N {0,1} dag_version--task_instance - -0..N + +0..N 1 dag_run - -dag_run - -id - - [INTEGER] - NOT NULL - -backfill_id - - [INTEGER] - -bundle_version - - [VARCHAR(250)] - -clear_number - - [INTEGER] - NOT NULL - -conf - - [JSONB] - -context_carrier - - [JSONB] - -created_dag_version_id - - [UUID] - -creating_job_id - - [INTEGER] - -dag_id - - [VARCHAR(250)] - NOT NULL - -data_interval_end - - [TIMESTAMP] - -data_interval_start - - [TIMESTAMP] - -end_date - - [TIMESTAMP] - -last_scheduling_decision - - [TIMESTAMP] - -log_template_id - - [INTEGER] - -logical_date - - [TIMESTAMP] - -queued_at - - [TIMESTAMP] - -run_after - - [TIMESTAMP] - NOT NULL - -run_id - - [VARCHAR(250)] - NOT NULL - -run_type - - [VARCHAR(50)] - NOT NULL - -scheduled_by_job_id - - [INTEGER] - -span_status - - [VARCHAR(250)] - NOT NULL - -start_date - - [TIMESTAMP] - -state - - [VARCHAR(50)] - -triggered_by - - [VARCHAR(50)] - -triggering_user_name - - [VARCHAR(512)] - -updated_at - - [TIMESTAMP] + +dag_run + +id + + [INTEGER] + NOT NULL + +backfill_id + + [INTEGER] + +bundle_version + + [VARCHAR(250)] + +clear_number + + [INTEGER] + NOT NULL + +conf + + [JSONB] + +context_carrier + + [JSONB] + +created_dag_version_id + + [UUID] + +creating_job_id + + [INTEGER] + +dag_id + + [VARCHAR(250)] + NOT NULL + +data_interval_end + + [TIMESTAMP] + +data_interval_start + + [TIMESTAMP] + +end_date + + [TIMESTAMP] + +last_scheduling_decision + + [TIMESTAMP] + +log_template_id + + [INTEGER] + +logical_date + + [TIMESTAMP] + +queued_at + + [TIMESTAMP] + +run_after + + [TIMESTAMP] + NOT NULL + +run_id + + [VARCHAR(250)] + NOT NULL + +run_type + + [VARCHAR(50)] + NOT NULL + +scheduled_by_job_id + + [INTEGER] + +span_status + + [VARCHAR(250)] + NOT NULL + +start_date + + [TIMESTAMP] + +state + + [VARCHAR(50)] + +triggered_by + + [VARCHAR(50)] + +triggering_user_name + + [VARCHAR(512)] + +updated_at + + [TIMESTAMP] dag_version--dag_run - -0..N -{0,1} + +0..N +{0,1} @@ -2108,107 +2108,107 @@ dag_run--dagrun_asset_event - -0..N -1 + +0..N +1 - + dag_run--task_instance - -0..N -1 + +0..N +1 - + dag_run--task_instance - -0..N -1 + +0..N +1 dag_run--deadline - -0..N -{0,1} + +0..N +{0,1} backfill_dag_run - -backfill_dag_run - -id - - [INTEGER] - NOT NULL - -backfill_id - - [INTEGER] - NOT NULL - -dag_run_id - - [INTEGER] - -exception_reason - - [VARCHAR(250)] - -logical_date - - [TIMESTAMP] - NOT NULL - -sort_ordinal - - [INTEGER] - NOT NULL + +backfill_dag_run + +id + + [INTEGER] + NOT NULL + +backfill_id + + [INTEGER] + NOT NULL + +dag_run_id + + [INTEGER] + +exception_reason + + [VARCHAR(250)] + +logical_date + + [TIMESTAMP] + NOT NULL + +sort_ordinal + + [INTEGER] + NOT NULL - + dag_run--backfill_dag_run - -0..N -{0,1} + +0..N +{0,1} dag_run_note - -dag_run_note - -dag_run_id - - [INTEGER] - NOT NULL - -content - - [VARCHAR(1000)] - -created_at - - [TIMESTAMP] - NOT NULL - -updated_at - - [TIMESTAMP] - NOT NULL - -user_id - - [VARCHAR(128)] + +dag_run_note + +dag_run_id + + [INTEGER] + NOT NULL + +content + + [VARCHAR(1000)] + +created_at + + [TIMESTAMP] + NOT NULL + +updated_at + + [TIMESTAMP] + NOT NULL + +user_id + + [VARCHAR(128)] dag_run--dag_run_note - -1 -1 + +1 +1 @@ -2239,9 +2239,9 @@ log_template--dag_run - -0..N -{0,1} + +0..N +{0,1} @@ -2309,15 +2309,15 @@ backfill--dag_run - -0..N -{0,1} + +0..N +{0,1} backfill--backfill_dag_run - -0..N + +0..N 1 diff --git a/airflow-core/pyproject.toml b/airflow-core/pyproject.toml index e36e1fa40bb49..5604382e4d19d 100644 --- a/airflow-core/pyproject.toml +++ b/airflow-core/pyproject.toml @@ -215,6 +215,9 @@ exclude = [ "src/airflow/ui/openapi.merged.json", ] +[tool.hatch.build.targets.sdist.force-include] +"../shared/timezones/src/airflow_shared/timezones" = "src/airflow/_shared/timezones" + [tool.hatch.build.targets.custom] path = "./hatch_build.py" diff --git a/airflow-core/src/airflow/_shared/__init__.py b/airflow-core/src/airflow/_shared/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/src/airflow/_shared/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow-core/src/airflow/_shared/timezones b/airflow-core/src/airflow/_shared/timezones new file mode 120000 index 0000000000000..8d7034fa71347 --- /dev/null +++ b/airflow-core/src/airflow/_shared/timezones @@ -0,0 +1 @@ +../../../../shared/timezones/src/airflow_shared/timezones \ No newline at end of file diff --git a/airflow-core/src/airflow/api/common/trigger_dag.py b/airflow-core/src/airflow/api/common/trigger_dag.py index 186ce9499a70a..d1cf5d93b8e51 100644 --- a/airflow-core/src/airflow/api/common/trigger_dag.py +++ b/airflow-core/src/airflow/api/common/trigger_dag.py @@ -22,9 +22,9 @@ import json from typing import TYPE_CHECKING +from airflow._shared.timezones import timezone from airflow.exceptions import DagNotFound, DagRunAlreadyExists from airflow.models import DagBag, DagModel, DagRun -from airflow.utils import timezone from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/auth/tokens.py b/airflow-core/src/airflow/api_fastapi/auth/tokens.py index f653532e29fb1..276ae17153da0 100644 --- a/airflow-core/src/airflow/api_fastapi/auth/tokens.py +++ b/airflow-core/src/airflow/api_fastapi/auth/tokens.py @@ -34,7 +34,7 @@ from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.serialization import load_pem_private_key -from airflow.utils import timezone +from airflow._shared.timezones import timezone if TYPE_CHECKING: from jwt.algorithms import AllowedKeys, AllowedPrivateKeys diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 4cf35d68871ba..f078721b9272b 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -37,6 +37,7 @@ from sqlalchemy import Column, and_, case, func, not_, or_, select from sqlalchemy.inspection import inspect +from airflow._shared.timezones import timezone from airflow.api_fastapi.core_api.base import OrmClause from airflow.api_fastapi.core_api.security import GetUserDep from airflow.models import Base @@ -56,7 +57,6 @@ from airflow.models.taskinstance import TaskInstance from airflow.models.variable import Variable from airflow.typing_compat import Self -from airflow.utils import timezone from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/common/types.py b/airflow-core/src/airflow/api_fastapi/common/types.py index 18e5dc7387d62..ccbb9d9ad98c7 100644 --- a/airflow-core/src/airflow/api_fastapi/common/types.py +++ b/airflow-core/src/airflow/api_fastapi/common/types.py @@ -30,7 +30,7 @@ ConfigDict, ) -from airflow.utils import timezone +from airflow._shared.timezones import timezone UtcDateTime = Annotated[AwareDatetime, AfterValidator(lambda d: d.astimezone(timezone.utc))] """UTCDateTime is a datetime with timezone information""" diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py index e184a3acfe864..6d66b5e781024 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py @@ -23,11 +23,11 @@ from pydantic import AliasPath, AwareDatetime, Field, NonNegativeInt, model_validator +from airflow._shared.timezones import timezone from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel from airflow.api_fastapi.core_api.datamodels.dag_versions import DagVersionResponse from airflow.models import DagRun from airflow.timetables.base import DataInterval -from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py index 6089d6d55dfc4..6ff8b83825f02 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py @@ -22,8 +22,8 @@ from pydantic import computed_field +from airflow._shared.timezones import timezone from airflow.api_fastapi.core_api.base import BaseModel -from airflow.utils import timezone from airflow.utils.state import TaskInstanceState from airflow.utils.types import DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py index 61fb5f9299a59..ace75342932fc 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py @@ -24,6 +24,7 @@ from sqlalchemy import and_, delete, func, select from sqlalchemy.orm import joinedload, subqueryload +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.dagbag import DagBagDep from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.parameters import ( @@ -72,7 +73,6 @@ TaskOutletAssetReference, ) from airflow.models.dag import DAG -from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py index 525b82c976b68..30efea4d79ab8 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/backfills.py @@ -24,6 +24,7 @@ from sqlalchemy import select, update from sqlalchemy.orm import joinedload +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import ( SessionDep, @@ -56,7 +57,6 @@ _create_backfill, _do_dry_run, ) -from airflow.utils import timezone from airflow.utils.state import DagRunState backfills_router = AirflowRouter(tags=["Backfill"], prefix="/backfills") diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py index 1d4d20fe64ea8..ce2f6c4bddde0 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 @@ -21,6 +21,7 @@ from sqlalchemy import select from sqlalchemy.orm import joinedload +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep, paginated_select from airflow.api_fastapi.common.router import AirflowRouter @@ -34,7 +35,6 @@ from airflow.api_fastapi.core_api.security import GetUserDep, ReadableTIFilterDep, requires_access_dag from airflow.models.hitl import HITLDetail as HITLDetailModel from airflow.models.taskinstance import TaskInstance as TI -from airflow.utils import timezone hitl_router = AirflowRouter(tags=["HumanInTheLoop"], prefix="/hitl-details") diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dashboard.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dashboard.py index f3c0198364137..9b552c9cf5883 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dashboard.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dashboard.py @@ -20,6 +20,7 @@ from sqlalchemy import func, select from sqlalchemy.sql.expression import case, false +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.parameters import DateTimeQuery, OptionalDateTimeQuery @@ -33,7 +34,6 @@ from airflow.models.dag import DagModel from airflow.models.dagrun import DagRun, DagRunType from airflow.models.taskinstance import TaskInstance -from airflow.utils import timezone from airflow.utils.state import DagRunState, TaskInstanceState dashboard_router = AirflowRouter(tags=["Dashboard"], prefix="/dashboard") diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py index de51a1dd27f49..28b7f32ead918 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/calendar.py @@ -28,6 +28,7 @@ from sqlalchemy.engine import Row from sqlalchemy.orm import Session +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.parameters import RangeFilter from airflow.api_fastapi.core_api.datamodels.ui.calendar import ( CalendarTimeRangeCollectionResponse, @@ -38,7 +39,6 @@ from airflow.timetables._cron import CronMixin from airflow.timetables.base import DataInterval, TimeRestriction from airflow.timetables.simple import ContinuousTimetable -from airflow.utils import timezone log = structlog.get_logger(logger_name=__name__) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index eef96323d2b7d..978c437476dbc 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -37,6 +37,7 @@ from sqlalchemy.sql import select from structlog.contextvars import bind_contextvars +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.dagbag import DagBagDep from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.common.types import UtcDateTime @@ -67,7 +68,6 @@ from airflow.sdk.definitions._internal.expandinput import NotFullyPopulated from airflow.sdk.definitions.asset import Asset, AssetUniqueKey from airflow.sdk.definitions.taskgroup import MappedTaskGroup -from airflow.utils import timezone from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState if TYPE_CHECKING: diff --git a/airflow-core/src/airflow/cli/commands/dag_command.py b/airflow-core/src/airflow/cli/commands/dag_command.py index 4553b5c57c2de..5047e24e0d084 100644 --- a/airflow-core/src/airflow/cli/commands/dag_command.py +++ b/airflow-core/src/airflow/cli/commands/dag_command.py @@ -31,6 +31,7 @@ from sqlalchemy import func, select +from airflow._shared.timezones import timezone from airflow.api.client import get_current_api_client from airflow.api_fastapi.core_api.datamodels.dags import DAGResponse from airflow.cli.simple_table import AirflowConsole @@ -41,7 +42,7 @@ from airflow.models import DagBag, DagModel, DagRun, TaskInstance from airflow.models.errors import ParseImportError from airflow.models.serialized_dag import SerializedDagModel -from airflow.utils import cli as cli_utils, timezone +from airflow.utils import cli as cli_utils from airflow.utils.cli import get_dag, suppress_logs_and_warning, validate_dag_bundle_arg from airflow.utils.dot_renderer import render_dag, render_dag_dependencies from airflow.utils.helpers import ask_yesno diff --git a/airflow-core/src/airflow/cli/commands/task_command.py b/airflow-core/src/airflow/cli/commands/task_command.py index 61a8f81c86d72..656ba480ced2c 100644 --- a/airflow-core/src/airflow/cli/commands/task_command.py +++ b/airflow-core/src/airflow/cli/commands/task_command.py @@ -28,6 +28,7 @@ from typing import TYPE_CHECKING, Protocol, cast from airflow import settings +from airflow._shared.timezones import timezone from airflow.cli.simple_table import AirflowConsole from airflow.cli.utils import fetch_dag_run_from_run_id_or_logical_date_string from airflow.exceptions import AirflowConfigException, DagRunNotFound, TaskInstanceNotFound @@ -41,7 +42,7 @@ from airflow.serialization.serialized_objects import SerializedDAG from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_deps import SCHEDULER_QUEUED_DEPS -from airflow.utils import cli as cli_utils, timezone +from airflow.utils import cli as cli_utils from airflow.utils.cli import ( get_dag, get_dag_by_file_location, diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 21ebffeb628b2..fdc68db86fa1e 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -48,6 +48,7 @@ from uuid6 import uuid7 import airflow.models +from airflow._shared.timezones import timezone from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI from airflow.configuration import conf from airflow.dag_processing.bundles.manager import DagBundlesManager @@ -65,7 +66,6 @@ from airflow.sdk.log import init_log_file, logging_processors from airflow.stats import Stats from airflow.traces.tracer import DebugTrace -from airflow.utils import timezone from airflow.utils.file import list_py_file_paths, might_contain_dag from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.net import get_hostname diff --git a/airflow-core/src/airflow/jobs/job.py b/airflow-core/src/airflow/jobs/job.py index 72db0cb684ee6..67fb3be019293 100644 --- a/airflow-core/src/airflow/jobs/job.py +++ b/airflow-core/src/airflow/jobs/job.py @@ -27,6 +27,7 @@ from sqlalchemy.orm import backref, foreign, relationship from sqlalchemy.orm.session import make_transient +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.executors.executor_loader import ExecutorLoader @@ -34,7 +35,6 @@ from airflow.models.base import ID_LEN, Base from airflow.stats import Stats from airflow.traces.tracer import DebugTrace, add_debug_span -from airflow.utils import timezone from airflow.utils.helpers import convert_camel_to_snake from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.net import get_hostname diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 4eb99616133f7..108d1720fd79a 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -38,6 +38,7 @@ from sqlalchemy.sql import expression from airflow import settings +from airflow._shared.timezones import timezone from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.configuration import conf @@ -69,7 +70,6 @@ from airflow.timetables.simple import AssetTriggeredTimetable from airflow.traces import utils as trace_utils from airflow.traces.tracer import DebugTrace, Trace, add_debug_span -from airflow.utils import timezone from airflow.utils.dates import datetime_to_nano from airflow.utils.event_scheduler import EventScheduler from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py index 19241a8c54847..0d353d7a9ac86 100644 --- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py +++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py @@ -38,6 +38,7 @@ from sqlalchemy import func, select from structlog.contextvars import bind_contextvars as bind_log_contextvars +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.executors import workloads from airflow.jobs.base_job_runner import BaseJobRunner @@ -69,7 +70,6 @@ from airflow.stats import Stats from airflow.traces.tracer import DebugTrace, Trace, add_debug_span from airflow.triggers import base as events -from airflow.utils import timezone from airflow.utils.helpers import log_filename_template_renderer from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.module_loading import import_string diff --git a/airflow-core/src/airflow/logging/_shared b/airflow-core/src/airflow/logging/_shared new file mode 120000 index 0000000000000..b75486419343c --- /dev/null +++ b/airflow-core/src/airflow/logging/_shared @@ -0,0 +1 @@ +../../../../shared/logging/src/airflow_logging/logging/ \ No newline at end of file diff --git a/airflow-core/src/airflow/migrations/versions/0047_3_0_0_add_dag_versioning.py b/airflow-core/src/airflow/migrations/versions/0047_3_0_0_add_dag_versioning.py index 2cad82e1e7f63..184c0a98218e0 100644 --- a/airflow-core/src/airflow/migrations/versions/0047_3_0_0_add_dag_versioning.py +++ b/airflow-core/src/airflow/migrations/versions/0047_3_0_0_add_dag_versioning.py @@ -31,10 +31,10 @@ from alembic import op from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.migrations.db_types import TIMESTAMP, StringID from airflow.migrations.utils import ignore_sqlite_value_error from airflow.models.base import naming_convention -from airflow.utils import timezone # revision identifiers, used by Alembic. revision = "2b47dc6bc8df" diff --git a/airflow-core/src/airflow/models/asset.py b/airflow-core/src/airflow/models/asset.py index a345ebe7995d3..8b28f2bea6184 100644 --- a/airflow-core/src/airflow/models/asset.py +++ b/airflow-core/src/airflow/models/asset.py @@ -36,9 +36,9 @@ ) from sqlalchemy.orm import relationship +from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID from airflow.settings import json -from airflow.utils import timezone from airflow.utils.sqlalchemy import UtcDateTime if TYPE_CHECKING: diff --git a/airflow-core/src/airflow/models/backfill.py b/airflow-core/src/airflow/models/backfill.py index e8eb3f8b5cd0f..957dcb8ea51d1 100644 --- a/airflow-core/src/airflow/models/backfill.py +++ b/airflow-core/src/airflow/models/backfill.py @@ -42,10 +42,10 @@ from sqlalchemy.orm import relationship, validates from sqlalchemy_jsonfield import JSONField +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException, DagNotFound from airflow.models.base import Base, StringID from airflow.settings import json -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.sqlalchemy import UtcDateTime, nulls_first, with_row_locks from airflow.utils.state import DagRunState diff --git a/airflow-core/src/airflow/models/dag.py b/airflow-core/src/airflow/models/dag.py index a141d96a1c30f..5ae0b7bf9b527 100644 --- a/airflow-core/src/airflow/models/dag.py +++ b/airflow-core/src/airflow/models/dag.py @@ -61,6 +61,7 @@ from sqlalchemy.sql import Select, expression from airflow import settings, utils +from airflow._shared.timezones import timezone from airflow.assets.evaluation import AssetEvaluator from airflow.configuration import conf as airflow_conf from airflow.exceptions import ( @@ -94,7 +95,6 @@ NullTimetable, OnceTimetable, ) -from airflow.utils import timezone from airflow.utils.context import Context from airflow.utils.dag_cycle_tester import check_cycle from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/airflow-core/src/airflow/models/dag_version.py b/airflow-core/src/airflow/models/dag_version.py index 5c1068d6a9614..a51f3cb301867 100644 --- a/airflow-core/src/airflow/models/dag_version.py +++ b/airflow-core/src/airflow/models/dag_version.py @@ -25,8 +25,8 @@ from sqlalchemy.orm import joinedload, relationship from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID -from airflow.utils import timezone from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime, with_row_locks diff --git a/airflow-core/src/airflow/models/dagbag.py b/airflow-core/src/airflow/models/dagbag.py index 83fc47444c179..a6ee7403b04f3 100644 --- a/airflow-core/src/airflow/models/dagbag.py +++ b/airflow-core/src/airflow/models/dagbag.py @@ -40,6 +40,7 @@ from tabulate import tabulate from airflow import settings +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.exceptions import ( AirflowClusterPolicyError, @@ -52,7 +53,6 @@ from airflow.listeners.listener import get_listener_manager from airflow.models.base import Base, StringID from airflow.stats import Stats -from airflow.utils import timezone from airflow.utils.dag_cycle_tester import check_cycle from airflow.utils.docs import get_docs_url from airflow.utils.file import ( diff --git a/airflow-core/src/airflow/models/dagcode.py b/airflow-core/src/airflow/models/dagcode.py index 2db338f87b50e..a68885e1dfdf2 100644 --- a/airflow-core/src/airflow/models/dagcode.py +++ b/airflow-core/src/airflow/models/dagcode.py @@ -26,10 +26,10 @@ from sqlalchemy.sql.expression import literal from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.exceptions import DagCodeNotFound from airflow.models.base import ID_LEN, Base -from airflow.utils import timezone from airflow.utils.file import open_maybe_zipped from airflow.utils.hashlib_wrapper import md5 from airflow.utils.session import NEW_SESSION, provide_session diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index cfac1936d6f87..b910e00f9b5b0 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -60,6 +60,7 @@ from sqlalchemy.sql.functions import coalesce from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.configuration import conf as airflow_conf from airflow.exceptions import AirflowException, TaskNotFound @@ -77,7 +78,6 @@ from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_states import SCHEDULEABLE_STATES from airflow.traces.tracer import EmptySpan, Trace -from airflow.utils import timezone from airflow.utils.dates import datetime_to_nano from airflow.utils.helpers import chunks, is_container, prune_dict from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/airflow-core/src/airflow/models/dagwarning.py b/airflow-core/src/airflow/models/dagwarning.py index b88158683eaf5..7498e615e30f7 100644 --- a/airflow-core/src/airflow/models/dagwarning.py +++ b/airflow-core/src/airflow/models/dagwarning.py @@ -22,9 +22,9 @@ from sqlalchemy import Column, ForeignKeyConstraint, Index, String, Text, delete, select, true +from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID from airflow.models.dag import DagModel -from airflow.utils import timezone from airflow.utils.retries import retry_db_transaction from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime diff --git a/airflow-core/src/airflow/models/db_callback_request.py b/airflow-core/src/airflow/models/db_callback_request.py index ada61d47c7a30..f1009ca1babe4 100644 --- a/airflow-core/src/airflow/models/db_callback_request.py +++ b/airflow-core/src/airflow/models/db_callback_request.py @@ -22,8 +22,8 @@ from sqlalchemy import Column, Integer, String +from airflow._shared.timezones import timezone from airflow.models.base import Base -from airflow.utils import timezone from airflow.utils.sqlalchemy import ExtendedJSON, UtcDateTime if TYPE_CHECKING: diff --git a/airflow-core/src/airflow/models/deadline.py b/airflow-core/src/airflow/models/deadline.py index e2dd8b39c0940..205ea0471450b 100644 --- a/airflow-core/src/airflow/models/deadline.py +++ b/airflow-core/src/airflow/models/deadline.py @@ -29,9 +29,9 @@ 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.settings import json -from airflow.utils import timezone from airflow.utils.decorators import classproperty from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import provide_session diff --git a/airflow-core/src/airflow/models/log.py b/airflow-core/src/airflow/models/log.py index 1dcaf83941a9e..9347e4a17536c 100644 --- a/airflow-core/src/airflow/models/log.py +++ b/airflow-core/src/airflow/models/log.py @@ -22,8 +22,8 @@ from sqlalchemy import Column, Index, Integer, String, Text from sqlalchemy.orm import relationship +from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID -from airflow.utils import timezone from airflow.utils.sqlalchemy import UtcDateTime if TYPE_CHECKING: diff --git a/airflow-core/src/airflow/models/serialized_dag.py b/airflow-core/src/airflow/models/serialized_dag.py index ade8dca8761d4..e64a0e5cc895e 100644 --- a/airflow-core/src/airflow/models/serialized_dag.py +++ b/airflow-core/src/airflow/models/serialized_dag.py @@ -32,6 +32,7 @@ from sqlalchemy.sql.expression import func, literal from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.exceptions import TaskNotFound from airflow.models.asset import ( AssetAliasModel, @@ -46,7 +47,6 @@ from airflow.serialization.dag_dependency import DagDependency from airflow.serialization.serialized_objects import SerializedDAG from airflow.settings import COMPRESS_SERIALIZED_DAGS, json -from airflow.utils import timezone from airflow.utils.hashlib_wrapper import md5 from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index ae55f6a65ad11..fe644f33abaf9 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -70,6 +70,7 @@ from sqlalchemy_utils import UUIDType from airflow import settings +from airflow._shared.timezones import timezone from airflow.assets.manager import asset_manager from airflow.configuration import conf from airflow.exceptions import ( @@ -89,7 +90,6 @@ from airflow.stats import Stats from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_deps import REQUEUEABLE_DEPS, RUNNING_DEPS -from airflow.utils import timezone from airflow.utils.email import send_email from airflow.utils.helpers import prune_dict, render_template_to_string from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/airflow-core/src/airflow/models/taskinstancehistory.py b/airflow-core/src/airflow/models/taskinstancehistory.py index cbd83643fdb8c..b5932b28c58a3 100644 --- a/airflow-core/src/airflow/models/taskinstancehistory.py +++ b/airflow-core/src/airflow/models/taskinstancehistory.py @@ -38,8 +38,8 @@ from sqlalchemy.orm import relationship from sqlalchemy_utils import UUIDType +from airflow._shared.timezones import timezone from airflow.models.base import Base, StringID -from airflow.utils import timezone from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.span_status import SpanStatus from airflow.utils.sqlalchemy import ( diff --git a/airflow-core/src/airflow/models/tasklog.py b/airflow-core/src/airflow/models/tasklog.py index d55eb94a266d7..d9a5c57c30ac5 100644 --- a/airflow-core/src/airflow/models/tasklog.py +++ b/airflow-core/src/airflow/models/tasklog.py @@ -19,8 +19,8 @@ from sqlalchemy import Column, Integer, Text +from airflow._shared.timezones import timezone from airflow.models.base import Base -from airflow.utils import timezone from airflow.utils.sqlalchemy import UtcDateTime diff --git a/airflow-core/src/airflow/models/trigger.py b/airflow-core/src/airflow/models/trigger.py index a4c54f8597436..2657e58bd0e84 100644 --- a/airflow-core/src/airflow/models/trigger.py +++ b/airflow-core/src/airflow/models/trigger.py @@ -28,12 +28,12 @@ from sqlalchemy.orm import Session, relationship, selectinload from sqlalchemy.sql.functions import coalesce +from airflow._shared.timezones import timezone from airflow.assets.manager import AssetManager from airflow.models.asset import asset_trigger_association_table from airflow.models.base import Base from airflow.models.taskinstance import TaskInstance from airflow.triggers.base import BaseTaskEndEvent -from airflow.utils import timezone from airflow.utils.retries import run_with_db_retries from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.sqlalchemy import UtcDateTime, with_row_locks diff --git a/airflow-core/src/airflow/models/xcom.py b/airflow-core/src/airflow/models/xcom.py index badb86bd43de5..a7301578ed2bd 100644 --- a/airflow-core/src/airflow/models/xcom.py +++ b/airflow-core/src/airflow/models/xcom.py @@ -39,8 +39,8 @@ from sqlalchemy.ext.associationproxy import association_proxy from sqlalchemy.orm import Query, relationship +from airflow._shared.timezones import timezone from airflow.models.base import COLLATION_ARGS, ID_LEN, TaskInstanceDependencies -from airflow.utils import timezone from airflow.utils.db import LazySelectSequence from airflow.utils.helpers import is_container from airflow.utils.json import XComDecoder, XComEncoder diff --git a/airflow-core/src/airflow/ti_deps/deps/not_in_retry_period_dep.py b/airflow-core/src/airflow/ti_deps/deps/not_in_retry_period_dep.py index 90954f29f2f50..1013fecdfb811 100644 --- a/airflow-core/src/airflow/ti_deps/deps/not_in_retry_period_dep.py +++ b/airflow-core/src/airflow/ti_deps/deps/not_in_retry_period_dep.py @@ -17,8 +17,8 @@ # under the License. from __future__ import annotations +from airflow._shared.timezones import timezone from airflow.ti_deps.deps.base_ti_dep import BaseTIDep -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import TaskInstanceState diff --git a/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py b/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py index abe0d38c2b707..501b1574205e1 100644 --- a/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py +++ b/airflow-core/src/airflow/ti_deps/deps/ready_to_reschedule.py @@ -17,10 +17,10 @@ # under the License. from __future__ import annotations +from airflow._shared.timezones import timezone from airflow.executors.executor_loader import ExecutorLoader from airflow.models.taskreschedule import TaskReschedule from airflow.ti_deps.deps.base_ti_dep import BaseTIDep -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import TaskInstanceState diff --git a/airflow-core/src/airflow/ti_deps/deps/runnable_exec_date_dep.py b/airflow-core/src/airflow/ti_deps/deps/runnable_exec_date_dep.py index 0f996e5e9f409..396f8180d1be1 100644 --- a/airflow-core/src/airflow/ti_deps/deps/runnable_exec_date_dep.py +++ b/airflow-core/src/airflow/ti_deps/deps/runnable_exec_date_dep.py @@ -17,8 +17,8 @@ # under the License. from __future__ import annotations +from airflow._shared.timezones import timezone from airflow.ti_deps.deps.base_ti_dep import BaseTIDep -from airflow.utils import timezone from airflow.utils.session import provide_session diff --git a/airflow-core/src/airflow/timetables/events.py b/airflow-core/src/airflow/timetables/events.py index ac2c0b6131982..d8e70626d409a 100644 --- a/airflow-core/src/airflow/timetables/events.py +++ b/airflow-core/src/airflow/timetables/events.py @@ -22,8 +22,8 @@ import pendulum +from airflow._shared.timezones import timezone from airflow.timetables.base import DagRunInfo, DataInterval, Timetable -from airflow.utils import timezone if TYPE_CHECKING: from pendulum import DateTime diff --git a/airflow-core/src/airflow/timetables/simple.py b/airflow-core/src/airflow/timetables/simple.py index bd9408434e84b..b5b6f2468f369 100644 --- a/airflow-core/src/airflow/timetables/simple.py +++ b/airflow-core/src/airflow/timetables/simple.py @@ -19,8 +19,8 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Any +from airflow._shared.timezones import timezone from airflow.timetables.base import DagRunInfo, DataInterval, Timetable -from airflow.utils import timezone if TYPE_CHECKING: from pendulum import DateTime diff --git a/airflow-core/src/airflow/traces/otel_tracer.py b/airflow-core/src/airflow/traces/otel_tracer.py index 982e8826e279f..22fb1d2935103 100644 --- a/airflow-core/src/airflow/traces/otel_tracer.py +++ b/airflow-core/src/airflow/traces/otel_tracer.py @@ -33,12 +33,12 @@ from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator from opentelemetry.trace.span import INVALID_SPAN_ID, INVALID_TRACE_ID +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.traces.utils import ( parse_traceparent, parse_tracestate, ) -from airflow.utils import timezone from airflow.utils.dates import datetime_to_nano from airflow.utils.net import get_hostname diff --git a/airflow-core/src/airflow/utils/cli.py b/airflow-core/src/airflow/utils/cli.py index ac7723d29338c..8f1fc0748bb7d 100644 --- a/airflow-core/src/airflow/utils/cli.py +++ b/airflow-core/src/airflow/utils/cli.py @@ -34,11 +34,12 @@ from typing import TYPE_CHECKING, TypeVar, cast from airflow import settings +from airflow._shared.timezones import timezone from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.exceptions import AirflowException from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager from airflow.sdk.execution_time.secrets_masker import should_hide_value_for_key -from airflow.utils import cli_action_loggers, timezone +from airflow.utils import cli_action_loggers from airflow.utils.log.non_caching_file_handler import NonCachingFileHandler from airflow.utils.platform import getuser, is_terminal_support_colors diff --git a/airflow-core/src/airflow/utils/db_cleanup.py b/airflow-core/src/airflow/utils/db_cleanup.py index 96e803775f803..7b0ee0487ddb8 100644 --- a/airflow-core/src/airflow/utils/db_cleanup.py +++ b/airflow-core/src/airflow/utils/db_cleanup.py @@ -36,10 +36,10 @@ from sqlalchemy.orm import aliased from sqlalchemy.sql.expression import ClauseElement, Executable, tuple_ +from airflow._shared.timezones import timezone from airflow.cli.simple_table import AirflowConsole from airflow.configuration import conf from airflow.exceptions import AirflowException -from airflow.utils import timezone from airflow.utils.db import reflect_tables from airflow.utils.helpers import ask_yesno from airflow.utils.session import NEW_SESSION, provide_session diff --git a/airflow-core/src/airflow/utils/log/file_processor_handler.py b/airflow-core/src/airflow/utils/log/file_processor_handler.py index 5dbd034cea8c2..ea3151be1db92 100644 --- a/airflow-core/src/airflow/utils/log/file_processor_handler.py +++ b/airflow-core/src/airflow/utils/log/file_processor_handler.py @@ -23,7 +23,7 @@ from pathlib import Path from airflow import settings -from airflow.utils import timezone +from airflow._shared.timezones import timezone from airflow.utils.helpers import parse_template_string from airflow.utils.log.logging_mixin import DISABLE_PROPOGATE from airflow.utils.log.non_caching_file_handler import NonCachingFileHandler diff --git a/airflow-core/src/airflow/utils/log/timezone_aware.py b/airflow-core/src/airflow/utils/log/timezone_aware.py index 0c9571c82574d..b7e69d03a73b1 100644 --- a/airflow-core/src/airflow/utils/log/timezone_aware.py +++ b/airflow-core/src/airflow/utils/log/timezone_aware.py @@ -18,8 +18,6 @@ import logging -from airflow.utils import timezone - class TimezoneAware(logging.Formatter): """ @@ -41,7 +39,9 @@ def formatTime(self, record, datefmt=None): This returns the creation time of the specified LogRecord in ISO 8601 date and time format in the local time zone. """ - dt = timezone.from_timestamp(record.created, tz="local") + from airflow._shared.timezones.timezone import from_timestamp + + dt = from_timestamp(record.created, tz="local") s = dt.strftime(datefmt or self.default_time_format) if self.default_msec_format: s = self.default_msec_format % (s, record.msecs) diff --git a/airflow-core/tests/integration/otel/test_otel.py b/airflow-core/tests/integration/otel/test_otel.py index 99e99c3941894..03f427227c113 100644 --- a/airflow-core/tests/integration/otel/test_otel.py +++ b/airflow-core/tests/integration/otel/test_otel.py @@ -25,13 +25,13 @@ import pytest +from airflow._shared.timezones import timezone from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.executors import executor_loader from airflow.executors.executor_utils import ExecutorName from airflow.models import DAG, DagBag, DagRun from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.span_status import SpanStatus from airflow.utils.state import State diff --git a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py index 6afef4db20662..14b6957ee651d 100644 --- a/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py +++ b/airflow-core/tests/unit/api_fastapi/auth/test_tokens.py @@ -27,6 +27,7 @@ import pytest from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import ( JWKS, InvalidClaimError, @@ -37,7 +38,6 @@ key_to_jwk_dict, key_to_pem, ) -from airflow.utils import timezone from tests_common.test_utils.config import conf_vars diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py index ce6ff10b30177..05a879c2f5be6 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py @@ -23,6 +23,7 @@ import pytest import time_machine +from airflow._shared.timezones import timezone from airflow.models import DagModel from airflow.models.asset import ( AssetActive, @@ -35,7 +36,6 @@ ) from airflow.models.dagrun import DagRun from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py index 703595b96e78f..b3eda51c1274b 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_backfills.py @@ -24,12 +24,12 @@ import pytest from sqlalchemy import and_, func, select +from airflow._shared.timezones import timezone from airflow.models import DagBag, DagModel, DagRun from airflow.models.backfill import Backfill, BackfillDagRun, ReprocessBehavior, _create_backfill from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.python import PythonOperator -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py index dbe55a92e7893..6bde1fc211c6d 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py @@ -25,6 +25,7 @@ import time_machine from sqlalchemy import select +from airflow._shared.timezones import timezone from airflow.api_fastapi.core_api.datamodels.dag_versions import DagVersionResponse from airflow.listeners.listener import get_listener_manager from airflow.models import DagModel, DagRun @@ -32,7 +33,6 @@ from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk.definitions.asset import Asset from airflow.sdk.definitions.param import Param -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState, State from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_stats.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_stats.py index 69e9a65ac2451..f1702a220dba1 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_stats.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_stats.py @@ -20,9 +20,9 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models.dag import DagModel from airflow.models.dagrun import DagRun -from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py index 4d84485d48ab4..780330c054596 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_extra_links.py @@ -20,13 +20,13 @@ import pytest +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.dagbag import dag_bag_from_app from airflow.api_fastapi.core_api.datamodels.extra_links import ExtraLinkCollectionResponse from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.models.dagbag import DagBag from airflow.models.xcom import XComModel as XCom from airflow.plugins_manager import AirflowPlugin -from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_log.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_log.py index adb5aebcb2106..3a4e7c8ae7128 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_log.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_log.py @@ -28,12 +28,12 @@ from itsdangerous.url_safe import URLSafeSerializer from uuid6 import uuid7 +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.dagbag import create_dag_bag, dag_bag_from_app from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import task -from airflow.utils import timezone from airflow.utils.types import DagRunType from tests_common.test_utils.db import clear_db_runs diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_monitor.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_monitor.py index a3cf5012e495a..5175278d7771d 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_monitor.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_monitor.py @@ -21,9 +21,9 @@ import pytest +from airflow._shared.timezones import timezone from airflow.jobs.job import Job from airflow.jobs.scheduler_job_runner import SchedulerJobRunner -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py index 5594f89f6f014..1c4194c178bd7 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py @@ -22,6 +22,7 @@ import pytest from airflow import DAG +from airflow._shared.timezones import timezone from airflow.api_fastapi.core_api.datamodels.xcom import XComCreateBody from airflow.models.dag_version import DagVersion from airflow.models.dagrun import DagRun @@ -31,7 +32,6 @@ from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.execution_time.xcom import resolve_xcom_backend -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py index dd6a0adb7b578..7d1f06a03f36b 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_backfills.py @@ -20,9 +20,9 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models import DagModel from airflow.models.backfill import Backfill -from airflow.utils import timezone from airflow.utils.session import provide_session from tests_common.test_utils.db import ( diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py index 9db616f4dbde2..96808bced3865 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py @@ -24,13 +24,13 @@ import pytest from sqlalchemy import select +from airflow._shared.timezones import timezone from airflow.models import DagBag from airflow.models.dag import DagModel from airflow.models.taskinstance import TaskInstance from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import task_group from airflow.sdk.definitions.taskgroup import TaskGroup -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py index 1587797256a7a..aa8bb0a842b4e 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py @@ -24,6 +24,7 @@ from sqlalchemy import select from sqlalchemy.orm import Session +from airflow._shared.timezones import timezone from airflow.models import DagBag from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel from airflow.providers.standard.operators.empty import EmptyOperator @@ -31,7 +32,6 @@ from airflow.providers.standard.sensors.external_task import ExternalTaskSensor from airflow.sdk import Metadata, task from airflow.sdk.definitions.asset import Asset, AssetAlias, Dataset -from airflow.utils import timezone from tests_common.test_utils.db import clear_db_assets, clear_db_runs diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py index ae8488e5da977..4bc313d85d96c 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py @@ -19,8 +19,8 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, AssetModel -from airflow.utils import timezone DEFAULT_DATE = timezone.parse("2021-01-01T00:00:00") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_assets.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_assets.py index 2cf34f8dd7bc7..81ff615726208 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_assets.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_assets.py @@ -19,8 +19,8 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models.asset import AssetActive, AssetModel -from airflow.utils import timezone DEFAULT_DATE = timezone.parse("2021-01-01T00:00:00") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py index f9f8d489d3d26..3a07682abeae1 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py @@ -19,10 +19,10 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models import DagModel from airflow.models.dagrun import DagRun from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone from airflow.utils.state import DagRunState, State from tests_common.test_utils.db import clear_db_runs diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index d9cdf9542aae3..325a9b7c31521 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -26,6 +26,7 @@ from sqlalchemy import select, update from sqlalchemy.exc import SQLAlchemyError +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import JWTValidator from airflow.api_fastapi.execution_api.app import lifespan from airflow.models import RenderedTaskInstanceFields, TaskReschedule, Trigger @@ -34,7 +35,6 @@ from airflow.models.taskinstancehistory import TaskInstanceHistory from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import Asset, TaskGroup, task, task_group -from airflow.utils import timezone from airflow.utils.state import 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 b5ef9a3ce5e50..e8d1c5be32236 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,8 +21,8 @@ import pytest +from airflow._shared.timezones import timezone from airflow.api_fastapi.common.dagbag import create_dag_bag, dag_bag_from_app -from airflow.utils import timezone from airflow.utils.state import State from tests_common.test_utils.db import clear_db_assets, clear_db_runs diff --git a/airflow-core/tests/unit/callbacks/test_callback_requests.py b/airflow-core/tests/unit/callbacks/test_callback_requests.py index d04145ff6d580..08c1a9a666d43 100644 --- a/airflow-core/tests/unit/callbacks/test_callback_requests.py +++ b/airflow-core/tests/unit/callbacks/test_callback_requests.py @@ -21,6 +21,7 @@ import pytest +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import ( DagCallbackRequest, TaskCallbackRequest, @@ -28,7 +29,6 @@ from airflow.models.dag import DAG from airflow.models.taskinstance import TaskInstance from airflow.providers.standard.operators.bash import BashOperator -from airflow.utils import timezone from airflow.utils.state import State, TaskInstanceState pytestmark = pytest.mark.db_test diff --git a/airflow-core/tests/unit/cli/commands/test_backfill_command.py b/airflow-core/tests/unit/cli/commands/test_backfill_command.py index ff23765fdd442..85b24b607cb76 100644 --- a/airflow-core/tests/unit/cli/commands/test_backfill_command.py +++ b/airflow-core/tests/unit/cli/commands/test_backfill_command.py @@ -26,9 +26,9 @@ import pytest import airflow.cli.commands.backfill_command +from airflow._shared.timezones import timezone from airflow.cli import cli_parser from airflow.models.backfill import ReprocessBehavior -from airflow.utils import timezone from tests_common.test_utils.db import clear_db_backfills, clear_db_dags, clear_db_runs, parse_and_sync_to_db diff --git a/airflow-core/tests/unit/cli/commands/test_dag_command.py b/airflow-core/tests/unit/cli/commands/test_dag_command.py index 219785fc7d053..2bf19837a0b3f 100644 --- a/airflow-core/tests/unit/cli/commands/test_dag_command.py +++ b/airflow-core/tests/unit/cli/commands/test_dag_command.py @@ -33,6 +33,7 @@ from sqlalchemy import select from airflow import settings +from airflow._shared.timezones import timezone from airflow.cli import cli_parser from airflow.cli.commands import dag_command from airflow.exceptions import AirflowException @@ -43,7 +44,6 @@ from airflow.sdk import task from airflow.sdk.definitions.dag import _run_inline_trigger from airflow.triggers.base import TriggerEvent -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/cli/commands/test_task_command.py b/airflow-core/tests/unit/cli/commands/test_task_command.py index 6e25253eb4067..0b56a4a627c70 100644 --- a/airflow-core/tests/unit/cli/commands/test_task_command.py +++ b/airflow-core/tests/unit/cli/commands/test_task_command.py @@ -32,6 +32,7 @@ import pytest +from airflow._shared.timezones import timezone from airflow.cli import cli_parser from airflow.cli.commands import task_command from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG @@ -42,7 +43,6 @@ from airflow.models.serialized_dag import SerializedDagModel from airflow.providers.standard.operators.bash import BashOperator from airflow.serialization.serialized_objects import SerializedDAG -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/tests/unit/core/test_sentry.py b/airflow-core/tests/unit/core/test_sentry.py index 33305c64a3940..723a9dc0ce510 100644 --- a/airflow-core/tests/unit/core/test_sentry.py +++ b/airflow-core/tests/unit/core/test_sentry.py @@ -26,8 +26,8 @@ from sentry_sdk import configure_scope from sentry_sdk.transport import Transport +from airflow._shared.timezones import timezone from airflow.providers.standard.operators.python import PythonOperator -from airflow.utils import timezone from airflow.utils.module_loading import import_string from airflow.utils.state import State diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_base.py b/airflow-core/tests/unit/dag_processing/bundles/test_base.py index 14986486c88cb..7da3baa3b9f53 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_base.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_base.py @@ -29,13 +29,13 @@ import pytest import time_machine +from airflow._shared.timezones import timezone as tz from airflow.dag_processing.bundles.base import ( BaseDagBundle, BundleUsageTrackingManager, BundleVersionLock, get_bundle_storage_root_path, ) -from airflow.utils import timezone as tz from tests_common.test_utils.config import conf_vars diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py b/airflow-core/tests/unit/dag_processing/test_collection.py index 5254dfc6ec973..0c812780ce410 100644 --- a/airflow-core/tests/unit/dag_processing/test_collection.py +++ b/airflow-core/tests/unit/dag_processing/test_collection.py @@ -31,6 +31,7 @@ from sqlalchemy.exc import OperationalError, SAWarning import airflow.dag_processing.collection +from airflow._shared.timezones import timezone as tz from airflow.configuration import conf from airflow.dag_processing.collection import ( AssetModelOperation, @@ -54,7 +55,6 @@ from airflow.providers.standard.triggers.temporal import TimeDeltaTrigger from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetWatcher from airflow.serialization.serialized_objects import LazyDeserializedDAG, SerializedDAG -from airflow.utils import timezone as tz from tests_common.test_utils.db import ( clear_db_assets, diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 02121f0194e82..c218c54eae768 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -40,6 +40,7 @@ from sqlalchemy import func, select from uuid6 import uuid7 +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG from airflow.dag_processing.bundles.manager import DagBundlesManager @@ -55,7 +56,6 @@ from airflow.models.dagbundle import DagBundleModel from airflow.models.dagcode import DagCode from airflow.models.serialized_dag import SerializedDagModel -from airflow.utils import timezone from airflow.utils.net import get_hostname from airflow.utils.session import create_session diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 66edfb26f6655..0de67aa9f078f 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -32,6 +32,7 @@ from pydantic import TypeAdapter from structlog.typing import FilteringBoundLogger +from airflow._shared.timezones import timezone from airflow.api_fastapi.execution_api.app import InProcessExecutionAPI from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( TaskInstance as TIDataModel, @@ -51,7 +52,6 @@ from airflow.sdk import DAG from airflow.sdk.api.client import Client 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_scheduler_dags.py b/airflow-core/tests/unit/dags/test_scheduler_dags.py index da89e821986f1..9123d32552f5e 100644 --- a/airflow-core/tests/unit/dags/test_scheduler_dags.py +++ b/airflow-core/tests/unit/dags/test_scheduler_dags.py @@ -19,9 +19,9 @@ from datetime import timedelta +from airflow._shared.timezones import timezone from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone DEFAULT_DATE = timezone.datetime(2016, 1, 1) diff --git a/airflow-core/tests/unit/dags/test_sensor.py b/airflow-core/tests/unit/dags/test_sensor.py index 4feb1caa1eb82..61504cfb77549 100644 --- a/airflow-core/tests/unit/dags/test_sensor.py +++ b/airflow-core/tests/unit/dags/test_sensor.py @@ -18,9 +18,9 @@ import datetime +from airflow._shared.timezones import timezone from airflow.models.dag import DAG from airflow.sdk import task -from airflow.utils import timezone from tests_common.test_utils.compat import DateTimeSensor diff --git a/airflow-core/tests/unit/decorators/test_mapped.py b/airflow-core/tests/unit/decorators/test_mapped.py index f1940f6458d85..1e148301cf694 100644 --- a/airflow-core/tests/unit/decorators/test_mapped.py +++ b/airflow-core/tests/unit/decorators/test_mapped.py @@ -19,8 +19,8 @@ import pytest +from airflow._shared.timezones import timezone from airflow.sdk import DAG, task -from airflow.utils import timezone DEFAULT_DATE = timezone.datetime(2025, 1, 1) diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 2ed0657e3c6a5..9715138c74f8a 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -26,6 +26,7 @@ import pytest import time_machine +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import CallbackRequest from airflow.cli.cli_config import DefaultHelpParser, GroupCommand from airflow.cli.cli_parser import AirflowHelpFormatter @@ -34,7 +35,6 @@ from airflow.executors.local_executor import LocalExecutor from airflow.models.baseoperator import BaseOperator from airflow.models.taskinstance import TaskInstance, TaskInstanceKey -from airflow.utils import timezone from airflow.utils.state import State, TaskInstanceState from tests_common.test_utils.markers import skip_if_force_lowest_dependencies_marker diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index e5b37bb268d3c..ec7102e2e70c2 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -25,9 +25,9 @@ from kgb import spy_on from uuid6 import uuid7 +from airflow._shared.timezones import timezone from airflow.executors import workloads from airflow.executors.local_executor import LocalExecutor, _execute_work -from airflow.utils import timezone from airflow.utils.state import State from tests_common.test_utils.config import conf_vars diff --git a/airflow-core/tests/unit/jobs/test_base_job.py b/airflow-core/tests/unit/jobs/test_base_job.py index c4b46a7622e53..9f80073007cef 100644 --- a/airflow-core/tests/unit/jobs/test_base_job.py +++ b/airflow-core/tests/unit/jobs/test_base_job.py @@ -25,10 +25,10 @@ import pytest from sqlalchemy.exc import OperationalError +from airflow._shared.timezones import timezone from airflow.executors.local_executor import LocalExecutor from airflow.jobs.job import Job, most_recent_job, perform_heartbeat, run_job from airflow.listeners.listener import get_listener_manager -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index f9269226e9cdc..818ab6ffaf152 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -38,6 +38,7 @@ from sqlalchemy.orm import joinedload from airflow import settings +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import JWTGenerator from airflow.assets.manager import AssetManager from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest @@ -74,7 +75,6 @@ from airflow.serialization.serialized_objects import LazyDeserializedDAG, SerializedDAG from airflow.timetables.base import DataInterval from airflow.traces.tracer import Trace -from airflow.utils import timezone from airflow.utils.session import create_session, provide_session from airflow.utils.span_status import SpanStatus from airflow.utils.state import DagRunState, State, TaskInstanceState diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py index 194c1c8752ba7..384bf6a76fc60 100644 --- a/airflow-core/tests/unit/jobs/test_triggerer_job.py +++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py @@ -32,6 +32,7 @@ from asgiref.sync import sync_to_async from structlog.typing import FilteringBoundLogger +from airflow._shared.timezones import timezone from airflow.executors import workloads from airflow.jobs.job import Job from airflow.jobs.triggerer_job_runner import ( @@ -55,7 +56,6 @@ from airflow.sdk import BaseHook from airflow.triggers.base import BaseTrigger, TriggerEvent from airflow.triggers.testing import FailureTrigger, SuccessTrigger -from airflow.utils import timezone from airflow.utils.state import State, TaskInstanceState from airflow.utils.types import DagRunType diff --git a/airflow-core/tests/unit/listeners/test_listeners.py b/airflow-core/tests/unit/listeners/test_listeners.py index 9a4c85cb504a1..fac237c7e176d 100644 --- a/airflow-core/tests/unit/listeners/test_listeners.py +++ b/airflow-core/tests/unit/listeners/test_listeners.py @@ -22,11 +22,11 @@ import pytest +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException from airflow.jobs.job import Job, run_job from airflow.listeners.listener import get_listener_manager from airflow.providers.standard.operators.bash import BashOperator -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import DagRunState, TaskInstanceState diff --git a/airflow-core/tests/unit/models/__init__.py b/airflow-core/tests/unit/models/__init__.py index ded9a4f45507e..d5fe7d8f9301b 100644 --- a/airflow-core/tests/unit/models/__init__.py +++ b/airflow-core/tests/unit/models/__init__.py @@ -19,7 +19,7 @@ import pathlib -from airflow.utils import timezone +from airflow._shared.timezones import timezone DEFAULT_DATE = timezone.datetime(2016, 1, 1) TEST_DAGS_FOLDER = pathlib.Path(__file__).parent.with_name("dags") diff --git a/airflow-core/tests/unit/models/test_backfill.py b/airflow-core/tests/unit/models/test_backfill.py index e69ff85aa04c6..d3f9c23b99491 100644 --- a/airflow-core/tests/unit/models/test_backfill.py +++ b/airflow-core/tests/unit/models/test_backfill.py @@ -25,6 +25,7 @@ import pytest from sqlalchemy import select +from airflow._shared.timezones import timezone from airflow.models import DagModel, DagRun, TaskInstance from airflow.models.backfill import ( AlreadyRunningBackfill, @@ -38,7 +39,6 @@ ) from airflow.providers.standard.operators.python import PythonOperator from airflow.ti_deps.dep_context import DepContext -from airflow.utils import timezone from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py index 9a99eba5974a2..04ae01c4a2982 100644 --- a/airflow-core/tests/unit/models/test_dag.py +++ b/airflow-core/tests/unit/models/test_dag.py @@ -34,6 +34,7 @@ from sqlalchemy import inspect, select from airflow import settings +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.exceptions import ( AirflowException, @@ -76,7 +77,6 @@ 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 diff --git a/airflow-core/tests/unit/models/test_dagbag.py b/airflow-core/tests/unit/models/test_dagbag.py index 00aa6151f9728..230ddca9f89e2 100644 --- a/airflow-core/tests/unit/models/test_dagbag.py +++ b/airflow-core/tests/unit/models/test_dagbag.py @@ -35,12 +35,12 @@ from sqlalchemy import select from airflow import settings +from airflow._shared.timezones import timezone as tz from airflow.models.dag import DAG, DagModel from airflow.models.dagbag import DagBag, _capture_with_reraise from airflow.models.dagwarning import DagWarning, DagWarningType from airflow.models.serialized_dag import SerializedDagModel from airflow.serialization.serialized_objects import SerializedDAG -from airflow.utils import timezone as tz from airflow.utils.session import create_session from tests_common.pytest_plugin import AIRFLOW_ROOT_PATH diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py index 50440d4f44ef2..cf4dafb4ee094 100644 --- a/airflow-core/tests/unit/models/test_dagrun.py +++ b/airflow-core/tests/unit/models/test_dagrun.py @@ -31,6 +31,7 @@ from sqlalchemy.orm import joinedload from airflow import settings +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.models.dag import DAG, DagModel from airflow.models.dag_version import DagVersion @@ -47,7 +48,6 @@ from airflow.serialization.serialized_objects import SerializedDAG from airflow.stats import Stats from airflow.triggers.base import StartTriggerArgs -from airflow.utils import timezone from airflow.utils.span_status import SpanStatus from airflow.utils.state import DagRunState, State, TaskInstanceState from airflow.utils.thread_safe_dict import ThreadSafeDict diff --git a/airflow-core/tests/unit/models/test_pool.py b/airflow-core/tests/unit/models/test_pool.py index ad0c935d618ca..81d4ce66dc717 100644 --- a/airflow-core/tests/unit/models/test_pool.py +++ b/airflow-core/tests/unit/models/test_pool.py @@ -20,12 +20,12 @@ import pytest from airflow import settings +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException, PoolNotFound from airflow.models.dag_version import DagVersion from airflow.models.pool import Pool from airflow.models.taskinstance import TaskInstance as TI from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 7e041917056ed..2a9e86771c2f7 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -34,6 +34,7 @@ from sqlalchemy.exc import IntegrityError from airflow import settings +from airflow._shared.timezones import timezone from airflow.exceptions import ( AirflowException, AirflowFailException, @@ -79,7 +80,6 @@ from airflow.ti_deps.deps.base_ti_dep import TIDepStatus from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep, _UpstreamTIStates -from airflow.utils import timezone from airflow.utils.db import merge_conn from airflow.utils.session import create_session, provide_session from airflow.utils.span_status import SpanStatus diff --git a/airflow-core/tests/unit/models/test_timestamp.py b/airflow-core/tests/unit/models/test_timestamp.py index 5aaa956829d72..529fd32afc504 100644 --- a/airflow-core/tests/unit/models/test_timestamp.py +++ b/airflow-core/tests/unit/models/test_timestamp.py @@ -20,9 +20,9 @@ import pytest import time_machine +from airflow._shared.timezones import timezone from airflow.models import Log from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.utils import timezone from airflow.utils.session import provide_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/models/test_trigger.py b/airflow-core/tests/unit/models/test_trigger.py index 8fd489aca554d..8fc8dc31271b5 100644 --- a/airflow-core/tests/unit/models/test_trigger.py +++ b/airflow-core/tests/unit/models/test_trigger.py @@ -27,6 +27,7 @@ import pytz from cryptography.fernet import Fernet +from airflow._shared.timezones import timezone from airflow.jobs.job import Job from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.models import TaskInstance, Trigger @@ -41,7 +42,6 @@ TaskSuccessEvent, TriggerEvent, ) -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/models/test_xcom.py b/airflow-core/tests/unit/models/test_xcom.py index 566d74ceaa774..6ed5d95e83d69 100644 --- a/airflow-core/tests/unit/models/test_xcom.py +++ b/airflow-core/tests/unit/models/test_xcom.py @@ -25,6 +25,7 @@ import pytest from airflow import DAG +from airflow._shared.timezones import timezone from airflow.configuration import conf from airflow.models.dag_version import DagVersion from airflow.models.dagrun import DagRun, DagRunType @@ -35,7 +36,6 @@ from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.execution_time.xcom import resolve_xcom_backend from airflow.settings import json -from airflow.utils import timezone from airflow.utils.session import create_session from tests_common.test_utils.config import conf_vars diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index c61284c785026..234bb51fbb22e 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -47,6 +47,7 @@ from kubernetes.client import models as k8s import airflow +from airflow._shared.timezones import timezone from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG from airflow.exceptions import ( AirflowException, @@ -82,7 +83,6 @@ from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep from airflow.timetables.simple import NullTimetable, OnceTimetable from airflow.triggers.base import StartTriggerArgs -from airflow.utils import timezone from airflow.utils.module_loading import qualname from airflow.utils.operator_resources import Resources diff --git a/airflow-core/tests/unit/serialization/test_serialized_objects.py b/airflow-core/tests/unit/serialization/test_serialized_objects.py index dc944035c75a9..dec42d3568bf4 100644 --- a/airflow-core/tests/unit/serialization/test_serialized_objects.py +++ b/airflow-core/tests/unit/serialization/test_serialized_objects.py @@ -28,6 +28,7 @@ from pendulum.tz.timezone import FixedTimezone, Timezone from uuid6 import uuid7 +from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.exceptions import ( AirflowException, @@ -65,7 +66,6 @@ from airflow.serialization.serialized_objects import BaseSerialization, LazyDeserializedDAG, SerializedDAG from airflow.timetables.base import DataInterval from airflow.triggers.base import BaseTrigger -from airflow.utils import timezone from airflow.utils.db import LazySelectSequence from airflow.utils.operator_resources import Resources from airflow.utils.state import DagRunState, State diff --git a/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py b/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py index 7b4bd22aa0436..355cbfb9c30ef 100644 --- a/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py +++ b/airflow-core/tests/unit/ti_deps/deps/test_ready_to_reschedule_dep.py @@ -24,10 +24,10 @@ import time_machine from slugify import slugify +from airflow._shared.timezones import timezone from airflow.models.taskreschedule import TaskReschedule from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep -from airflow.utils import timezone from airflow.utils.session import create_session from airflow.utils.state import State diff --git a/airflow-core/tests/unit/timetables/test_once_timetable.py b/airflow-core/tests/unit/timetables/test_once_timetable.py index a904e8372a7fb..0ca84229ba63c 100644 --- a/airflow-core/tests/unit/timetables/test_once_timetable.py +++ b/airflow-core/tests/unit/timetables/test_once_timetable.py @@ -23,9 +23,9 @@ import pytest import time_machine +from airflow._shared.timezones import timezone from airflow.timetables.base import DagRunInfo, TimeRestriction from airflow.timetables.simple import OnceTimetable -from airflow.utils import timezone FROZEN_NOW = timezone.coerce_datetime(datetime.datetime(2025, 3, 4, 5, 6, 7, 8)) diff --git a/airflow-core/tests/unit/utils/log/test_file_processor_handler.py b/airflow-core/tests/unit/utils/log/test_file_processor_handler.py index f50fc619ac5bd..ad0982dbea433 100644 --- a/airflow-core/tests/unit/utils/log/test_file_processor_handler.py +++ b/airflow-core/tests/unit/utils/log/test_file_processor_handler.py @@ -23,7 +23,7 @@ import time_machine -from airflow.utils import timezone +from airflow._shared.timezones import timezone from airflow.utils.log.file_processor_handler import FileProcessorHandler diff --git a/airflow-core/tests/unit/utils/log/test_log_reader.py b/airflow-core/tests/unit/utils/log/test_log_reader.py index cd8b430090a27..8d5ffab543dfa 100644 --- a/airflow-core/tests/unit/utils/log/test_log_reader.py +++ b/airflow-core/tests/unit/utils/log/test_log_reader.py @@ -29,11 +29,11 @@ import pytest from airflow import settings +from airflow._shared.timezones import timezone from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG from airflow.models.tasklog import LogTemplate from airflow.providers.standard.operators.python import PythonOperator from airflow.timetables.base import DataInterval -from airflow.utils import timezone from airflow.utils.log.log_reader import TaskLogReader from airflow.utils.log.logging_mixin import ExternalLoggingMixin from airflow.utils.state import TaskInstanceState diff --git a/airflow-core/tests/unit/utils/test_cli_util.py b/airflow-core/tests/unit/utils/test_cli_util.py index 2d5578f8faa93..54ecc506e2731 100644 --- a/airflow-core/tests/unit/utils/test_cli_util.py +++ b/airflow-core/tests/unit/utils/test_cli_util.py @@ -30,9 +30,10 @@ import airflow from airflow import settings +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException from airflow.models.log import Log -from airflow.utils import cli, cli_action_loggers, timezone +from airflow.utils import cli, cli_action_loggers from airflow.utils.cli import _search_for_dag_file # Mark entire module as db_test because ``action_cli`` wrapper still could use DB on callbacks: diff --git a/airflow-core/tests/unit/utils/test_db_cleanup.py b/airflow-core/tests/unit/utils/test_db_cleanup.py index 48a0c8486219b..71c481ea1c650 100644 --- a/airflow-core/tests/unit/utils/test_db_cleanup.py +++ b/airflow-core/tests/unit/utils/test_db_cleanup.py @@ -31,12 +31,12 @@ from sqlalchemy.ext.declarative import DeclarativeMeta from airflow import DAG +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException from airflow.models import DagModel, DagRun, TaskInstance from airflow.models.dag_version import DagVersion from airflow.models.serialized_dag import SerializedDagModel from airflow.providers.standard.operators.python import PythonOperator -from airflow.utils import timezone from airflow.utils.db_cleanup import ( ARCHIVE_TABLE_PREFIX, CreateTableAs, diff --git a/airflow-core/tests/unit/utils/test_dot_renderer.py b/airflow-core/tests/unit/utils/test_dot_renderer.py index 240876ec7f43e..1e239271dfce7 100644 --- a/airflow-core/tests/unit/utils/test_dot_renderer.py +++ b/airflow-core/tests/unit/utils/test_dot_renderer.py @@ -22,11 +22,12 @@ import pytest +from airflow._shared.timezones import timezone from airflow.models.dag import DAG from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk.definitions.taskgroup import TaskGroup from airflow.serialization.dag_dependency import DagDependency -from airflow.utils import dot_renderer, timezone +from airflow.utils import dot_renderer from airflow.utils.state import State from tests_common.test_utils.compat import BashOperator diff --git a/airflow-core/tests/unit/utils/test_helpers.py b/airflow-core/tests/unit/utils/test_helpers.py index 71138aa51057f..08c8d8092a0e9 100644 --- a/airflow-core/tests/unit/utils/test_helpers.py +++ b/airflow-core/tests/unit/utils/test_helpers.py @@ -23,9 +23,10 @@ import pytest +from airflow._shared.timezones import timezone from airflow.exceptions import AirflowException from airflow.jobs.base_job_runner import BaseJobRunner -from airflow.utils import helpers, timezone +from airflow.utils import helpers from airflow.utils.helpers import ( at_most_one, build_airflow_dagrun_url, diff --git a/airflow-core/tests/unit/utils/test_serve_logs.py b/airflow-core/tests/unit/utils/test_serve_logs.py index fbbef2b200046..b6c7805386ac9 100644 --- a/airflow-core/tests/unit/utils/test_serve_logs.py +++ b/airflow-core/tests/unit/utils/test_serve_logs.py @@ -24,9 +24,9 @@ import time_machine from fastapi.testclient import TestClient +from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import JWTGenerator from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG -from airflow.utils import timezone from airflow.utils.serve_logs import create_app from tests_common.test_utils.config import conf_vars diff --git a/airflow-core/tests/unit/utils/test_timezone.py b/airflow-core/tests/unit/utils/test_timezone.py index c7e6d0c09aa81..99a8eab892be8 100644 --- a/airflow-core/tests/unit/utils/test_timezone.py +++ b/airflow-core/tests/unit/utils/test_timezone.py @@ -23,7 +23,7 @@ import pytest from pendulum.tz.timezone import FixedTimezone, Timezone -from airflow.utils import timezone +from airflow._shared.timezones import timezone from airflow.utils.timezone import coerce_datetime, parse_timezone CET = Timezone("Europe/Paris") diff --git a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py index ed9e9dd7a2fd6..b959d53fc0469 100644 --- a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py +++ b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py @@ -83,25 +83,26 @@ ("RELEASE_NOTES.rst", "/opt/airflow/RELEASE_NOTES.rst"), ("airflow-core", "/opt/airflow/airflow-core"), ("airflow-ctl", "/opt/airflow/airflow-ctl"), - ("constraints", "/opt/airflow/constraints"), + ("chart", "/opt/airflow/chart"), ("clients", "/opt/airflow/clients"), + ("constraints", "/opt/airflow/constraints"), ("dags", "/opt/airflow/dags"), ("dev", "/opt/airflow/dev"), - ("docs", "/opt/airflow/docs"), + ("devel-common", "/opt/airflow/devel-common"), ("docker-stack-docs", "/opt/airflow/docker-stack-docs"), - ("providers-summary-docs", "/opt/airflow/providers-summary-docs"), + ("docker-tests", "/opt/airflow/docker-tests"), + ("docs", "/opt/airflow/docs"), ("generated", "/opt/airflow/generated"), + ("helm-tests", "/opt/airflow/helm-tests"), + ("kubernetes-tests", "/opt/airflow/kubernetes-tests"), ("logs", "/root/airflow/logs"), ("providers", "/opt/airflow/providers"), - ("task-sdk", "/opt/airflow/task-sdk"), + ("providers-summary-docs", "/opt/airflow/providers-summary-docs"), ("pyproject.toml", "/opt/airflow/pyproject.toml"), ("scripts", "/opt/airflow/scripts"), ("scripts/docker/entrypoint_ci.sh", "/entrypoint"), - ("devel-common", "/opt/airflow/devel-common"), - ("helm-tests", "/opt/airflow/helm-tests"), - ("kubernetes-tests", "/opt/airflow/kubernetes-tests"), - ("docker-tests", "/opt/airflow/docker-tests"), - ("chart", "/opt/airflow/chart"), + ("shared", "/opt/airflow/shared"), + ("task-sdk", "/opt/airflow/task-sdk"), ] diff --git a/pyproject.toml b/pyproject.toml index c6e02a26946a1..06853e2ebf831 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -809,14 +809,21 @@ 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" + +# Add this once providers are sorted +# "airflow.utils.timezone".msg = "Use airflow.timezone" + # 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." @@ -1247,7 +1254,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 +1301,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 +1416,7 @@ members = [ "task-sdk", "providers-summary-docs", "docker-stack-docs", + "shared/timezones", # Automatically generated provider workspace members (update_airflow_pyproject_toml.py) "providers/airbyte", "providers/alibaba", diff --git a/scripts/ci/docker-compose/local.yml b/scripts/ci/docker-compose/local.yml index 6609aa5c1dbda..4522441263dea 100644 --- a/scripts/ci/docker-compose/local.yml +++ b/scripts/ci/docker-compose/local.yml @@ -102,6 +102,9 @@ services: - type: bind source: ../../../scripts target: /opt/airflow/scripts + - type: bind + source: ../../../shared + target: /opt/airflow/shared - type: bind source: ../../../scripts/docker/entrypoint_ci.sh target: /entrypoint diff --git a/shared/timezones/pyproject.toml b/shared/timezones/pyproject.toml new file mode 100644 index 0000000000000..bc68f6467c894 --- /dev/null +++ b/shared/timezones/pyproject.toml @@ -0,0 +1,47 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[project] +name = "apache-airflow-shared-timezones" +description = "Shared timezone code for Airflow distributions" +version = "0.0" +classifiers = [ + "Private :: Do Not Upload", +] + +dependencies = [ + "pendulum", +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/airflow_shared"] + +[tool.ruff] +extend = "../../pyproject.toml" +src = ["src"] + +[tool.ruff.lint.per-file-ignores] +# Ignore Doc rules et al for anything outside of tests +"!src/*" = ["D", "S101", "TRY002"] + +[tool.ruff.lint.flake8-tidy-imports] +# Override the workspace level default +ban-relative-imports = "parents" diff --git a/shared/timezones/src/airflow_shared/timezones/__init__.py b/shared/timezones/src/airflow_shared/timezones/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/shared/timezones/src/airflow_shared/timezones/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/shared/timezones/src/airflow_shared/timezones/timezone.py b/shared/timezones/src/airflow_shared/timezones/timezone.py new file mode 100644 index 0000000000000..df2567e03db79 --- /dev/null +++ b/shared/timezones/src/airflow_shared/timezones/timezone.py @@ -0,0 +1,318 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import datetime as dt +from importlib import metadata +from typing import TYPE_CHECKING, overload + +import pendulum +from dateutil.relativedelta import relativedelta +from packaging import version +from pendulum.datetime import DateTime + +if TYPE_CHECKING: + from pendulum.tz.timezone import FixedTimezone, Timezone + + +_PENDULUM3 = version.parse(metadata.version("pendulum")).major == 3 +# UTC Timezone as a tzinfo instance. Actual value depends on pendulum version: +# - Timezone("UTC") in pendulum 3 +# - FixedTimezone(0, "UTC") in pendulum 2 +utc = pendulum.UTC + + +def is_localized(value: dt.datetime) -> bool: + """ + Determine if a given datetime.datetime is aware. + + The concept is defined in Python documentation. Assuming the tzinfo is + either None or a proper ``datetime.tzinfo`` instance, ``value.utcoffset()`` + implements the appropriate logic. + + .. seealso:: http://docs.python.org/library/datetime.html#datetime.tzinfo + """ + return value.utcoffset() is not None + + +def is_naive(value): + """ + Determine if a given datetime.datetime is naive. + + The concept is defined in Python documentation. Assuming the tzinfo is + either None or a proper ``datetime.tzinfo`` instance, ``value.utcoffset()`` + implements the appropriate logic. + + .. seealso:: http://docs.python.org/library/datetime.html#datetime.tzinfo + """ + return value.utcoffset() is None + + +def utcnow() -> dt.datetime: + """Get the current date and time in UTC.""" + 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: ... + + +@overload +def convert_to_utc(value: dt.datetime) -> DateTime: ... + + +def convert_to_utc(value: dt.datetime | None) -> DateTime | None: + """ + Create a datetime with the default timezone added if none is associated. + + :param value: datetime + :return: datetime with tzinfo + """ + if value is None: + return value + + if not is_localized(value): + from airflow.settings import TIMEZONE + + value = pendulum.instance(value, TIMEZONE) + + return pendulum.instance(value.astimezone(utc)) + + +@overload +def make_aware(value: None, timezone: dt.tzinfo | None = None) -> None: ... + + +@overload +def make_aware(value: DateTime, timezone: dt.tzinfo | None = None) -> DateTime: ... + + +@overload +def make_aware(value: dt.datetime, timezone: dt.tzinfo | None = None) -> dt.datetime: ... + + +def make_aware(value: dt.datetime | None, timezone: dt.tzinfo | None = None) -> dt.datetime | None: + """ + Make a naive datetime.datetime in a given time zone aware. + + :param value: datetime + :param timezone: timezone + :return: localized datetime in settings.TIMEZONE or timezone + """ + if timezone is None: + from airflow.settings import TIMEZONE + + timezone = TIMEZONE + + if not value: + return None + + # Check that we won't overwrite the timezone of an aware datetime. + if is_localized(value): + raise ValueError(f"make_aware expects a naive datetime, got {value}") + # In case we move clock back we want to schedule the run at the time of the second + # instance of the same clock time rather than the first one. + # Fold parameter has no impact in other cases, so we can safely set it to 1 here + value = value.replace(fold=1) + localized = getattr(timezone, "localize", None) + if localized is not None: + # This method is available for pytz time zones + return localized(value) + convert = getattr(timezone, "convert", None) + if convert is not None: + # For pendulum + return convert(value) + # This may be wrong around DST changes! + return value.replace(tzinfo=timezone) + + +def make_naive(value, timezone=None): + """ + Make an aware datetime.datetime naive in a given time zone. + + :param value: datetime + :param timezone: timezone + :return: naive datetime + """ + if timezone is None: + from airflow.settings import TIMEZONE + + timezone = TIMEZONE + + # Emulate the behavior of astimezone() on Python < 3.6. + if is_naive(value): + raise ValueError("make_naive() cannot be applied to a naive datetime") + + date = value.astimezone(timezone) + + # cross library compatibility + naive = dt.datetime( + date.year, date.month, date.day, date.hour, date.minute, date.second, date.microsecond + ) + + return naive + + +def datetime(*args, **kwargs): + """ + Wrap around datetime.datetime to add settings.TIMEZONE if tzinfo not specified. + + :return: datetime.datetime + """ + if "tzinfo" not in kwargs: + from airflow.settings import TIMEZONE + + kwargs["tzinfo"] = TIMEZONE + + return dt.datetime(*args, **kwargs) + + +def parse(string: str, timezone=None, *, strict=False) -> DateTime: + """ + Parse a time string and return an aware datetime. + + :param string: time string + :param timezone: the timezone + :param strict: if False, it will fall back on the dateutil parser if unable to parse with pendulum + """ + from airflow.settings import TIMEZONE + + return pendulum.parse(string, tz=timezone or TIMEZONE, strict=strict) # type: ignore + + +@overload +def coerce_datetime(v: None, tz: dt.tzinfo | None = None) -> None: ... + + +@overload +def coerce_datetime(v: DateTime, tz: dt.tzinfo | None = None) -> DateTime: ... + + +@overload +def coerce_datetime(v: dt.datetime, tz: dt.tzinfo | None = None) -> DateTime: ... + + +def coerce_datetime(v: dt.datetime | None, tz: dt.tzinfo | None = None) -> DateTime | None: + """ + Convert ``v`` into a timezone-aware ``pendulum.DateTime``. + + * If ``v`` is *None*, *None* is returned. + * If ``v`` is a naive datetime, it is converted to an aware Pendulum DateTime. + * If ``v`` is an aware datetime, it is converted to a Pendulum DateTime. + Note that ``tz`` is **not** taken into account in this case; the datetime + will maintain its original tzinfo! + """ + if v is None: + return None + if isinstance(v, DateTime): + return v if v.tzinfo else make_aware(v, tz) + # Only dt.datetime is left here. + return pendulum.instance(v if v.tzinfo else make_aware(v, tz)) + + +def td_format(td_object: None | dt.timedelta | float | int) -> str | None: + """ + Format a timedelta object or float/int into a readable string for time duration. + + For example timedelta(seconds=3752) would become `1h:2M:32s`. + If the time is less than a second, the return will be `<1s`. + """ + if not td_object: + return None + if isinstance(td_object, dt.timedelta): + delta = relativedelta() + td_object + else: + delta = relativedelta(seconds=int(td_object)) + # relativedelta for timedelta cannot convert days to months + # so calculate months by assuming 30 day months and normalize + months, delta.days = divmod(delta.days, 30) + delta = delta.normalized() + relativedelta(months=months) + + def _format_part(key: str) -> str: + value = int(getattr(delta, key)) + if value < 1: + return "" + # distinguish between month/minute following strftime format + # and take first char of each unit, i.e. years='y', days='d' + if key == "minutes": + key = key.upper() + key = key[0] + return f"{value}{key}" + + parts = map(_format_part, ("years", "months", "days", "hours", "minutes", "seconds")) + joined = ":".join(part for part in parts if part) + if not joined: + return "<1s" + return joined + + +def parse_timezone(name: str | int) -> FixedTimezone | Timezone: + """ + Parse timezone and return one of the pendulum Timezone. + + Provide the same interface as ``pendulum.timezone(name)`` + + :param name: Either IANA timezone or offset to UTC in seconds. + + :meta private: + """ + if _PENDULUM3: + # This only presented in pendulum 3 and code do not reached into the pendulum 2 + 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] + + +def local_timezone() -> FixedTimezone | Timezone: + """ + Return local timezone. + + Provide the same interface as ``pendulum.tz.local_timezone()`` + + :meta private: + """ + return pendulum.tz.local_timezone() + + +def from_timestamp(timestamp: int | float, tz: str | FixedTimezone | Timezone = utc) -> DateTime: + """ + Parse timestamp and return DateTime in a given time zone. + + :param timestamp: epoch time in seconds. + :param tz: In which timezone should return a resulting object. + Could be either one of pendulum timezone, IANA timezone or `local` literal. + + :meta private: + """ + result = coerce_datetime(dt.datetime.fromtimestamp(timestamp, tz=utc)) + if tz != utc or tz != "UTC": + if isinstance(tz, str) and tz.lower() == "local": + tz = local_timezone() + result = result.in_timezone(tz) + return result diff --git a/shared/timezones/tests/test_timezone.py b/shared/timezones/tests/test_timezone.py new file mode 100644 index 0000000000000..ba0f2a46def55 --- /dev/null +++ b/shared/timezones/tests/test_timezone.py @@ -0,0 +1,202 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import datetime + +import pendulum +import pytest +from pendulum.tz.timezone import FixedTimezone, 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 +ICT = Timezone("Asia/Bangkok") # Asia/Bangkok +UTC = timezone.utc + + +class TestTimezone: + def test_is_aware(self): + assert timezone.is_localized(datetime.datetime(2011, 9, 1, 13, 20, 30, tzinfo=EAT)) + assert not timezone.is_localized(datetime.datetime(2011, 9, 1, 13, 20, 30)) + + def test_is_naive(self): + assert not timezone.is_naive(datetime.datetime(2011, 9, 1, 13, 20, 30, tzinfo=EAT)) + assert timezone.is_naive(datetime.datetime(2011, 9, 1, 13, 20, 30)) + + def test_utcnow(self): + now = timezone.utcnow() + assert timezone.is_localized(now) + assert now.replace(tzinfo=None) == now.astimezone(UTC).replace(tzinfo=None) + + def test_convert_to_utc(self): + naive = datetime.datetime(2011, 9, 1, 13, 20, 30) + utc = datetime.datetime(2011, 9, 1, 13, 20, 30, tzinfo=UTC) + assert utc == timezone.convert_to_utc(naive) + + eat = datetime.datetime(2011, 9, 1, 13, 20, 30, tzinfo=EAT) + utc = datetime.datetime(2011, 9, 1, 10, 20, 30, tzinfo=UTC) + assert utc == timezone.convert_to_utc(eat) + + def test_make_naive(self): + assert timezone.make_naive( + datetime.datetime(2011, 9, 1, 13, 20, 30, tzinfo=EAT), EAT + ) == datetime.datetime(2011, 9, 1, 13, 20, 30) + assert timezone.make_naive( + datetime.datetime(2011, 9, 1, 17, 20, 30, tzinfo=ICT), EAT + ) == datetime.datetime(2011, 9, 1, 13, 20, 30) + + with pytest.raises(ValueError): + timezone.make_naive(datetime.datetime(2011, 9, 1, 13, 20, 30), EAT) + + def test_make_aware(self): + assert timezone.make_aware(datetime.datetime(2011, 9, 1, 13, 20, 30), EAT) == datetime.datetime( + 2011, 9, 1, 13, 20, 30, tzinfo=EAT + ) + with pytest.raises(ValueError): + timezone.make_aware(datetime.datetime(2011, 9, 1, 13, 20, 30, tzinfo=EAT), EAT) + + def test_td_format(self): + td = datetime.timedelta(seconds=3752) + assert timezone.td_format(td) == "1h:2M:32s" + td = 3200.0 + assert timezone.td_format(td) == "53M:20s" + td = 3200 + assert timezone.td_format(td) == "53M:20s" + td = 0.123 + assert timezone.td_format(td) == "<1s" + td = None + assert timezone.td_format(td) is None + td = datetime.timedelta(seconds=300752) + assert timezone.td_format(td) == "3d:11h:32M:32s" + td = 434343600.0 + assert timezone.td_format(td) == "13y:11m:17d:3h" + + +@pytest.mark.parametrize( + "input_datetime, output_datetime", + [ + pytest.param(None, None, id="None datetime"), + pytest.param( + pendulum.DateTime(2021, 11, 1), + pendulum.DateTime(2021, 11, 1, tzinfo=UTC), + id="Non aware pendulum Datetime", + ), + pytest.param( + pendulum.DateTime(2021, 11, 1, tzinfo=CET), + pendulum.DateTime(2021, 11, 1, tzinfo=CET), + id="Aware pendulum Datetime", + ), + pytest.param( + datetime.datetime(2021, 11, 1), + pendulum.DateTime(2021, 11, 1, tzinfo=UTC), + id="Non aware datetime", + ), + pytest.param( + datetime.datetime(2021, 11, 1, tzinfo=CET), + pendulum.DateTime(2021, 11, 1, tzinfo=CET), + id="Aware datetime", + ), + ], +) +def test_coerce_datetime(input_datetime, output_datetime): + assert output_datetime == coerce_datetime(input_datetime) + + +@pytest.mark.parametrize( + "tz_name", + [ + pytest.param("Europe/Paris", id="CET"), + pytest.param("Africa/Nairobi", id="EAT"), + pytest.param("Asia/Bangkok", id="ICT"), + ], +) +def test_parse_timezone_iana(tz_name: str): + tz = parse_timezone(tz_name) + assert tz.name == tz_name + assert parse_timezone(tz_name) is tz + + +@pytest.mark.parametrize("tz_name", ["utc", "UTC", "uTc"]) +def test_parse_timezone_utc(tz_name): + tz = parse_timezone(tz_name) + assert tz.name == "UTC" + assert parse_timezone(tz_name) is tz + assert tz is timezone.utc, "Expected that UTC timezone is same object as `airflow.utils.timezone.utc`" + + +@pytest.mark.parametrize( + "tz_offset, expected_offset, expected_name", + [ + pytest.param(0, 0, "+00:00", id="zero-offset"), + pytest.param(-3600, -3600, "-01:00", id="1-hour-behind"), + pytest.param(19800, 19800, "+05:30", id="5.5-hours-ahead"), + ], +) +def test_parse_timezone_offset(tz_offset: int, expected_offset, expected_name): + tz = parse_timezone(tz_offset) + assert hasattr(tz, "offset") + assert tz.offset == expected_offset + assert tz.name == expected_name + assert parse_timezone(tz_offset) is tz + + +@pytest.mark.parametrize( + "tz", + [ + pytest.param(None, id="implicit"), + pytest.param(timezone.utc, id="explicit"), + pytest.param("UTC", id="utc-literal"), + ], +) +def test_from_timestamp_utc(tz): + from_ts = timezone.from_timestamp(0) if tz is None else timezone.from_timestamp(0, tz=tz) + assert from_ts == pendulum.DateTime(1970, 1, 1, tzinfo=timezone.utc) + assert from_ts.tzinfo == timezone.utc + + +@pytest.mark.parametrize("tz", ["local", "LOCAL"]) +def test_from_timestamp_local(tz): + local_tz = timezone.local_timezone() + from_ts = timezone.from_timestamp(0, tz=tz) + assert from_ts == pendulum.DateTime(1970, 1, 1, tzinfo=timezone.utc) + assert from_ts.tzinfo == local_tz + + +@pytest.mark.parametrize( + "tz, iana_timezone", + [ + pytest.param(Timezone("Europe/Paris"), "Europe/Paris", id="pendulum-timezone"), + pytest.param("America/New_York", "America/New_York", id="IANA-timezone"), + ], +) +def test_from_timestamp_iana_timezones(tz, iana_timezone): + from_ts = timezone.from_timestamp(0, tz=tz) + assert from_ts == pendulum.DateTime(1970, 1, 1, tzinfo=timezone.utc) + # In pendulum 2 there is a problem with compare tzinfo object (caching?), so we check the name + assert from_ts.tzinfo.name == iana_timezone + assert isinstance(from_ts.tzinfo, Timezone) + + +@pytest.mark.parametrize("utc_offset", [3600, -7200]) +def test_from_timestamp_fixed_timezone(utc_offset): + from_ts = timezone.from_timestamp(0, tz=FixedTimezone(utc_offset)) + assert from_ts == pendulum.DateTime(1970, 1, 1, tzinfo=timezone.utc) + assert from_ts.utcoffset() == datetime.timedelta(seconds=utc_offset) diff --git a/task-sdk/pyproject.toml b/task-sdk/pyproject.toml index a6edccc35c634..f44e3ea0558fb 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 @@ -118,6 +121,9 @@ namespace-packages = ["src/airflow"] # Generated file, be less strict "src/airflow/sdk/*/_generated.py" = ["D"] +[tool.ruff.lint.flake8-tidy-imports.banned-api] +"airflow.utils.timezone".msg = "Use airflow.sdk.timezone" + [tool.coverage.run] branch = true relative_files = true diff --git a/task-sdk/src/airflow/sdk/_shared/__init__.py b/task-sdk/src/airflow/sdk/_shared/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/task-sdk/src/airflow/sdk/_shared/timezones b/task-sdk/src/airflow/sdk/_shared/timezones new file mode 120000 index 0000000000000..b66387c7a6b56 --- /dev/null +++ b/task-sdk/src/airflow/sdk/_shared/timezones @@ -0,0 +1 @@ +../../../../../shared/timezones/src/airflow_shared/timezones \ No newline at end of file diff --git a/task-sdk/src/airflow/sdk/bases/decorator.py b/task-sdk/src/airflow/sdk/bases/decorator.py index 42bc66ee5e487..cfbb3c69ddc2f 100644 --- a/task-sdk/src/airflow/sdk/bases/decorator.py +++ b/task-sdk/src/airflow/sdk/bases/decorator.py @@ -28,6 +28,7 @@ import attr import typing_extensions +from airflow.sdk._shared.timezones import timezone from airflow.sdk.bases.operator import ( BaseOperator, coerce_resources, @@ -46,7 +47,6 @@ from airflow.sdk.definitions.asset import Asset from airflow.sdk.definitions.mappedoperator import MappedOperator, ensure_xcomarg_return_value from airflow.sdk.definitions.xcom_arg import XComArg -from airflow.utils import timezone from airflow.utils.context import KNOWN_CONTEXT_KEYS from airflow.utils.decorators import remove_task_decorator from airflow.utils.helpers import prevent_duplicates diff --git a/task-sdk/src/airflow/sdk/bases/operator.py b/task-sdk/src/airflow/sdk/bases/operator.py index c430f38982348..1a970ba83b8cf 100644 --- a/task-sdk/src/airflow/sdk/bases/operator.py +++ b/task-sdk/src/airflow/sdk/bases/operator.py @@ -36,6 +36,7 @@ import attrs from airflow.exceptions import RemovedInAirflow4Warning +from airflow.sdk._shared.timezones import timezone from airflow.sdk.definitions._internal.abstractoperator import ( DEFAULT_IGNORE_FIRST_DEPENDS_ON_PAST, DEFAULT_OWNER, @@ -65,7 +66,6 @@ airflow_priority_weight_strategies, validate_and_load_priority_weight_strategy, ) -from airflow.utils import timezone from airflow.utils.trigger_rule import TriggerRule from airflow.utils.weight_rule import db_safe_priority diff --git a/task-sdk/src/airflow/sdk/bases/sensor.py b/task-sdk/src/airflow/sdk/bases/sensor.py index a39de2a151873..cadba573ab55e 100644 --- a/task-sdk/src/airflow/sdk/bases/sensor.py +++ b/task-sdk/src/airflow/sdk/bases/sensor.py @@ -36,8 +36,8 @@ TaskDeferralError, TaskDeferralTimeout, ) +from airflow.sdk._shared.timezones 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 994d8d6198b17..62091947b72cb 100644 --- a/task-sdk/src/airflow/sdk/definitions/dag.py +++ b/task-sdk/src/airflow/sdk/definitions/dag.py @@ -217,7 +217,7 @@ def dict_copy(_: dict[str, Any]) -> dict[str, Any]: ... def _default_start_date(instance: DAG): # Find start date inside default_args for compat with Airflow 2. - from airflow.utils import timezone + from airflow.sdk import timezone if date := instance.default_args.get("start_date"): if not isinstance(date, datetime): @@ -456,7 +456,7 @@ def __rich_repr__(self): ) def __attrs_post_init__(self): - from airflow.utils import timezone + from airflow.sdk import timezone # Apply the timezone we settled on to end_date if it wasn't supplied if isinstance(_end_date := self.default_args.get("end_date"), str): @@ -528,7 +528,7 @@ def _default_timetable(instance: DAG): def _extract_tz(instance): import pendulum - from airflow.utils import timezone + from airflow.sdk import timezone start_date = instance.start_date or instance.default_args.get("start_date") @@ -1047,9 +1047,9 @@ def test( from airflow.configuration import secrets_backend_list from airflow.models.dag import DAG as SchedulerDAG, _get_or_create_dagrun from airflow.models.dagrun import DagRun + from airflow.sdk import timezone from airflow.secrets.local_filesystem import LocalFilesystemBackend from airflow.serialization.serialized_objects import SerializedDAG - from airflow.utils import timezone from airflow.utils.state import DagRunState, State, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType diff --git a/task-sdk/src/airflow/sdk/execution_time/cache.py b/task-sdk/src/airflow/sdk/execution_time/cache.py index 9902a05f032f9..5a6f415159dcd 100644 --- a/task-sdk/src/airflow/sdk/execution_time/cache.py +++ b/task-sdk/src/airflow/sdk/execution_time/cache.py @@ -21,7 +21,7 @@ import multiprocessing from airflow.configuration import conf -from airflow.utils import timezone +from airflow.sdk._shared.timezones 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 ab8be39aba001..af76251b2e2fe 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -100,9 +100,9 @@ set_current_context, ) from airflow.sdk.execution_time.xcom import XCom +from airflow.sdk.timezone import coerce_datetime from airflow.utils.net import get_hostname from airflow.utils.platform import getuser -from airflow.utils.timezone import coerce_datetime if TYPE_CHECKING: import jinja2 diff --git a/task-sdk/src/airflow/sdk/timezone.py b/task-sdk/src/airflow/sdk/timezone.py new file mode 100644 index 0000000000000..d50cf43d586fc --- /dev/null +++ b/task-sdk/src/airflow/sdk/timezone.py @@ -0,0 +1,47 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +# We don't want to `import *` here to avoid the risk of making adding too much to Public python API +from airflow.sdk._shared.timezones.timezone import ( + coerce_datetime, + convert_to_utc, + datetime, + from_timestamp, + is_localized, + is_naive, + local_timezone, + make_aware, + make_naive, + utc_epoch, + utcnow, +) + +__all__ = [ + "is_localized", + "is_naive", + "utcnow", + "utc_epoch", + "convert_to_utc", + "make_aware", + "make_naive", + "datetime", + "coerce_datetime", + "local_timezone", + "from_timestamp", +] diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 4f4af1e137242..243339be5dbec 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -29,6 +29,7 @@ from task_sdk import make_client, make_client_w_dry_run, make_client_w_responses from uuid6 import uuid7 +from airflow.sdk._shared.timezones import timezone from airflow.sdk.api.client import RemoteValidationError, ServerResponseError from airflow.sdk.api.datamodels._generated import ( AssetEventsResponse, @@ -49,7 +50,6 @@ RescheduleTask, TaskRescheduleStartDate, ) -from airflow.utils import timezone from airflow.utils.state import TerminalTIState if TYPE_CHECKING: diff --git a/task-sdk/tests/task_sdk/bases/test_sensor.py b/task-sdk/tests/task_sdk/bases/test_sensor.py index d0e2b430d2c74..e15c04324ac4b 100644 --- a/task-sdk/tests/task_sdk/bases/test_sensor.py +++ b/task-sdk/tests/task_sdk/bases/test_sensor.py @@ -34,12 +34,12 @@ ) from airflow.models.trigger import TriggerFailureReason from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.sdk._shared.timezones 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 diff --git a/task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py b/task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py index 15c2d65cc2717..d0d77d8d917fb 100644 --- a/task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py +++ b/task-sdk/tests/task_sdk/dags/super_basic_deferred_run.py @@ -20,8 +20,8 @@ import datetime from airflow.providers.standard.sensors.date_time import DateTimeSensorAsync +from airflow.sdk._shared.timezones import timezone from airflow.sdk.definitions.dag import dag -from airflow.utils import timezone @dag 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 5595fc2775fab..ee8074d69e85f 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_comms.py +++ b/task-sdk/tests/task_sdk/execution_time/test_comms.py @@ -24,9 +24,9 @@ import msgspec import pytest +from airflow.sdk._shared.timezones import timezone from airflow.sdk.execution_time.comms import BundleInfo, StartupDetails, _ResponseFrame from airflow.sdk.execution_time.task_runner import CommsDecoder -from airflow.utils import timezone class TestCommsDecoder: diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py b/task-sdk/tests/task_sdk/execution_time/test_context.py index d3c22598a98e3..c7a59f2d3c0ef 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context.py @@ -23,6 +23,7 @@ import pytest from airflow.sdk import BaseOperator, get_current_context +from airflow.sdk._shared.timezones import 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 +64,6 @@ context_to_airflow_vars, set_current_context, ) -from airflow.utils import timezone def test_convert_connection_result_conn(): diff --git a/task-sdk/tests/task_sdk/execution_time/test_hitl.py b/task-sdk/tests/task_sdk/execution_time/test_hitl.py index cab17e30bacec..76dd058f5475c 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_hitl.py +++ b/task-sdk/tests/task_sdk/execution_time/test_hitl.py @@ -19,6 +19,7 @@ from uuid6 import uuid7 +from airflow.sdk._shared.timezones import timezone from airflow.sdk.api.datamodels._generated import HITLDetailResponse from airflow.sdk.execution_time.comms import CreateHITLDetailPayload from airflow.sdk.execution_time.hitl import ( @@ -26,7 +27,6 @@ get_hitl_detail_content_detail, update_htil_detail_response, ) -from airflow.utils import timezone TI_ID = uuid7() diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index d1b75f1ed22f1..955b066de21e3 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -44,6 +44,7 @@ from uuid6 import uuid7 from airflow.executors.workloads import BundleInfo +from airflow.sdk._shared.timezones 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 ( @@ -115,7 +116,6 @@ set_supervisor_comms, supervise, ) -from airflow.utils import timezone, timezone as tz if TYPE_CHECKING: import kgb @@ -233,7 +233,7 @@ def subprocess_main(): line = lineno() - 2 # Line the error should be on - instant = tz.datetime(2024, 11, 7, 12, 34, 56, 78901) + instant = timezone.datetime(2024, 11, 7, 12, 34, 56, 78901) time_machine.move_to(instant, tick=False) proc = ActivitySubprocess.start( @@ -456,7 +456,7 @@ def _on_child_started(self, *args, **kwargs): def test_run_simple_dag(self, test_dags_dir, captured_logs, time_machine, mocker, client_with_ti_start): """Test running a simple DAG in a subprocess and capturing the output.""" - instant = tz.datetime(2024, 11, 7, 12, 34, 56, 78901) + instant = timezone.datetime(2024, 11, 7, 12, 34, 56, 78901) time_machine.move_to(instant, tick=False) dagfile_path = test_dags_dir @@ -500,7 +500,7 @@ def test_supervise_handles_deferred_task( This includes ensuring the task starts and executes successfully, and that the task is deferred (via the API client) with the expected parameters. """ - instant = tz.datetime(2024, 11, 7, 12, 34, 56, 0) + instant = timezone.datetime(2024, 11, 7, 12, 34, 56, 0) ti = TaskInstance( id=uuid7(), @@ -727,7 +727,7 @@ def test_heartbeat_failures_handling(self, monkeypatch, mocker, captured_logs, t process=mock_process, ) - time_now = tz.datetime(2024, 11, 28, 12, 0, 0) + time_now = timezone.datetime(2024, 11, 28, 12, 0, 0) time_machine.move_to(time_now, tick=False) # Simulate sending heartbeats and ensure the process gets killed after max retries diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 7bd4aaf519a97..640e6367fbc38 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 @@ -47,6 +47,7 @@ from airflow.listeners.listener import get_listener_manager from airflow.providers.standard.operators.python import PythonOperator from airflow.sdk import DAG, BaseOperator, Connection, dag as dag_decorator, get_current_context +from airflow.sdk._shared.timezones import timezone from airflow.sdk.api.datamodels._generated import ( AssetProfile, AssetResponse, @@ -112,7 +113,6 @@ startup, ) from airflow.sdk.execution_time.xcom import XCom -from airflow.utils import timezone from airflow.utils.types import NOTSET, ArgNotSet from tests_common.test_utils.mock_operators import AirflowLink @@ -1130,7 +1130,7 @@ def test_get_context_without_ti_context_from_server(self, mocked_parse, make_ti_ def test_get_context_with_ti_context_from_server(self, create_runtime_ti, mock_supervisor_comms): """Test the context keys are added when sent from API server (mocked)""" - from airflow.utils import timezone + from airflow.sdk import timezone task = BaseOperator(task_id="hello")