Skip to content

Enable the airflow.utils.timezone ban rule #71209

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Draft
FrankYang0529 wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Enable the airflow.utils.timezone ban rule #71209

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion dev/airflow_perf/scheduler_dag_execution_timing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ def create_dag_runs(dag, num_runs, session):
"""
Create `num_runs` of dag runs for sub-sequent schedules
"""
from airflow.utils import timezone
from airflow.sdk import timezone
from airflow.utils.state import DagRunState

try:
Expand Down
2 changes: 1 addition & 1 deletion devel-common/src/tests_common/pytest_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3211,7 +3211,7 @@ def _import_timezone():
try:
from airflow._shared.timezones import timezone
except ImportError:
from airflow.utils import timezone
from airflow.utils import timezone # noqa: TID251
return timezone


Expand Down
2 changes: 1 addition & 1 deletion devel-common/src/tests_common/test_utils/integration_setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
try:
from airflow.sdk import timezone
except ImportError:
from airflow.utils import timezone # type: ignore[no-redef,attr-defined]
from airflow.utils import timezone # type: ignore[no-redef,attr-defined] # noqa: TID251

log = logging.getLogger(__name__)

Expand Down
2 changes: 1 addition & 1 deletion devel-common/src/tests_common/test_utils/version_compat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
XCOM_RETURN_KEY = BaseXCom.XCOM_RETURN_KEY
else:
from airflow.sensors.base import PokeReturnValue # type: ignore[no-redef]
from airflow.utils import timezone # type: ignore[attr-defined,no-redef]
from airflow.utils import timezone # type: ignore[attr-defined,no-redef] # noqa: TID251
from airflow.utils.decorators import remove_task_decorator # type: ignore[no-redef]
from airflow.utils.types import NOTSET, ArgNotSet # type: ignore[attr-defined,no-redef]
from airflow.utils.xcom import XCOM_RETURN_KEY # type: ignore[no-redef]
Expand Down
2 changes: 1 addition & 1 deletion kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,8 @@
from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction, PodManager
from airflow.sdk import timezone
from airflow.sdk.definitions.context import Context
from airflow.utils import timezone # type: ignore[attr-defined]
from airflow.utils.types import DagRunType
from airflow.version import version as airflow_version
from kubernetes_tests.test_base import BaseK8STest, StringContainingId
Expand Down
6 changes: 2 additions & 4 deletions performance/src/performance_dags/performance_dag/performance_dag_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
from datetime import datetime, timedelta
from shutil import copyfile

import airflow
from airflow.sdk import timezone

log = logging.getLogger(__name__)
log.setLevel(logging.INFO)
Expand Down Expand Up @@ -80,9 +80,7 @@ def add_perf_start_date_env_to_conf(performance_dag_conf: dict[str, str]) -> Non
if "PERF_START_DATE" not in performance_dag_conf:
start_ago = get_performance_dag_environment_variable(performance_dag_conf, "PERF_START_AGO")

perf_start_date = airflow.utils.timezone.utcnow - check_and_parse_time_delta(
"PERF_START_AGO", start_ago
)
perf_start_date = timezone.utcnow() - check_and_parse_time_delta("PERF_START_AGO", start_ago)

performance_dag_conf["PERF_START_DATE"] = str(perf_start_date)

Expand Down
4 changes: 2 additions & 2 deletions providers/alibaba/tests/unit/alibaba/cloud/log/test_oss_task_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,8 @@
import pytest

from airflow.providers.alibaba.cloud.log.oss_task_handler import OSSRemoteLogIO, OSSTaskHandler
from airflow.providers.common.compat.sdk import timezone
from airflow.utils.state import TaskInstanceState
from airflow.utils.timezone import datetime

from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.db import clear_db_dags, clear_db_runs
Expand Down Expand Up @@ -137,7 +137,7 @@ def task_instance(self, create_task_instance: CreateTaskInstance, dag_maker: Dag
self.ti = ti = create_task_instance(
dag_id="dag_for_testing_oss_task_handler",
task_id="task_for_testing_oss_task_handler",
logical_date=datetime(2020, 1, 1),
logical_date=timezone.datetime(2020, 1, 1),
state=TaskInstanceState.RUNNING,
)
ti.try_number = 1
Expand Down
5 changes: 1 addition & 4 deletions ...ders/amazon/src/airflow/providers/amazon/aws/executors/utils/exponential_backoff_retry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,7 @@
from collections.abc import Callable
from datetime import datetime, timedelta

try:
from airflow.sdk import timezone
except ImportError:
from airflow.utils import timezone # type: ignore[attr-defined,no-redef]
from airflow.providers.common.compat.sdk import timezone

log = logging.getLogger(__name__)

Expand Down
4 changes: 2 additions & 2 deletions providers/amazon/tests/system/amazon/aws/example_mongo_to_s3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
except ImportError:
# Compatibility for Airflow < 3.1
from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined]
from airflow.utils.timezone import datetime
from airflow.providers.common.compat.sdk import timezone

from system.amazon.aws.utils import SystemTestContextBuilder

Expand All @@ -42,7 +42,7 @@
with DAG(
DAG_ID,
schedule="@once",
start_date=datetime(2021, 1, 1),
start_date=timezone.datetime(2021, 1, 1),
catchup=False,
) as dag:
test_context = sys_test_context_task()
Expand Down
4 changes: 2 additions & 2 deletions providers/amazon/tests/system/amazon/aws/example_mwaa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,17 +24,17 @@
from airflow.providers.amazon.aws.hooks.sts import StsHook
from airflow.providers.amazon.aws.operators.mwaa import MwaaTriggerDagRunOperator
from airflow.providers.amazon.aws.sensors.mwaa import MwaaDagRunSensor, MwaaTaskSensor
from airflow.providers.common.compat.sdk import timezone

from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS

if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import DAG, chain, task, timezone
from airflow.sdk import DAG, chain, task
else:
# Airflow 2 path
from airflow.decorators import task # type: ignore[attr-defined,no-redef]
from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef]
from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment]
from airflow.utils import timezone # type: ignore[attr-defined,no-redef]


from system.amazon.aws.utils import SystemTestContextBuilder
Expand Down
5 changes: 2 additions & 3 deletions providers/amazon/tests/unit/amazon/aws/executors/ecs/test_ecs_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,10 +52,9 @@
parse_assign_public_ip,
)
from airflow.providers.amazon.aws.hooks.ecs import EcsHook
from airflow.providers.common.compat.sdk import AirflowException, conf
from airflow.providers.common.compat.sdk import AirflowException, conf, timezone
from airflow.utils.helpers import convert_camel_to_snake
from airflow.utils.state import State, TaskInstanceState
from airflow.utils.timezone import utcnow
from airflow.version import version as airflow_version_str

from tests_common import RUNNING_TESTS_AGAINST_AIRFLOW_PACKAGES
Expand Down Expand Up @@ -882,7 +881,7 @@ def test_failed_sync_cumulative_fail(
{"arn": ARN1, "reason": "Sample Failure", "detail": "UnitTest Failure - Please ignore"}
],
}
mock_executor._calculate_next_attempt_time = MagicMock(return_value=utcnow())
mock_executor._calculate_next_attempt_time = MagicMock(return_value=timezone.utcnow())
task_key = mock_airflow_key()
mock_executor.execute_async(task_key, mock_cmd)
for _ in range(2):
Expand Down
77 changes: 54 additions & 23 deletions providers/amazon/tests/unit/amazon/aws/hooks/test_s3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import os
import re
from collections.abc import Iterator
from datetime import datetime as std_datetime, timedelta, timezone
from datetime import datetime as std_datetime, timedelta, timezone as std_timezone
from pathlib import Path
from unittest import mock, mock as async_mock
from unittest.mock import AsyncMock, MagicMock, Mock, patch
Expand All @@ -42,8 +42,7 @@
provide_bucket_name,
unify_bucket_name_and_key,
)
from airflow.providers.common.compat.sdk import AirflowException
from airflow.utils.timezone import datetime
from airflow.providers.common.compat.sdk import AirflowException, timezone

try:
import importlib.util
Expand Down Expand Up @@ -374,8 +373,8 @@ def test_list_keys(self, s3_bucket):
bucket.put_object(Key="bxb", Body=b"axb")
bucket.put_object(Key="dir/b", Body=b"b")

from_datetime = datetime(1992, 3, 8, 18, 52, 51)
to_datetime = datetime(1993, 3, 14, 21, 52, 42)
from_datetime = timezone.datetime(1992, 3, 8, 18, 52, 51)
to_datetime = timezone.datetime(1993, 3, 14, 21, 52, 42)

def dummy_object_filter(keys, from_datetime=None, to_datetime=None):
return []
Expand Down Expand Up @@ -642,8 +641,16 @@ async def test_s3_key_hook_get_file_metadata_async_when_requester_pays_is(self,
test_resp_iter = [
{
"Contents": [
{"Key": "test_key", "ETag": "etag1", "LastModified": datetime(2020, 8, 14, 17, 19, 34)},
{"Key": "test_key2", "ETag": "etag2", "LastModified": datetime(2020, 8, 14, 17, 19, 34)},
{
"Key": "test_key",
"ETag": "etag1",
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
},
{
"Key": "test_key2",
"ETag": "etag2",
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
},
]
}
]
Expand All @@ -658,8 +665,8 @@ async def test_s3_key_hook_get_file_metadata_async_when_requester_pays_is(self,
keys = [x async for x in s3_hook_async.get_file_metadata_async(mock_client, "test_bucket", "test*")]

assert keys == [
{"Key": "test_key", "ETag": "etag1", "LastModified": datetime(2020, 8, 14, 17, 19, 34)},
{"Key": "test_key2", "ETag": "etag2", "LastModified": datetime(2020, 8, 14, 17, 19, 34)},
{"Key": "test_key", "ETag": "etag1", "LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34)},
{"Key": "test_key2", "ETag": "etag2", "LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34)},
]
extra_params = {}
if requester_pays:
Expand Down Expand Up @@ -751,8 +758,16 @@ async def test_s3_key_hook_get_files_without_wildcard_async(self):
test_resp_iter = [
{
"Contents": [
{"Key": "test_key", "ETag": "etag1", "LastModified": datetime(2020, 8, 14, 17, 19, 34)},
{"Key": "test_key2", "ETag": "etag2", "LastModified": datetime(2020, 8, 14, 17, 19, 34)},
{
"Key": "test_key",
"ETag": "etag1",
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
},
{
"Key": "test_key2",
"ETag": "etag2",
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
},
]
}
]
Expand Down Expand Up @@ -782,8 +797,16 @@ async def test_s3_key_hook_get_files_with_wildcard_async(self):
test_resp_iter = [
{
"Contents": [
{"Key": "test_key", "ETag": "etag1", "LastModified": datetime(2020, 8, 14, 17, 19, 34)},
{"Key": "test_key2", "ETag": "etag2", "LastModified": datetime(2020, 8, 14, 17, 19, 34)},
{
"Key": "test_key",
"ETag": "etag1",
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
},
{
"Key": "test_key2",
"ETag": "etag2",
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
},
]
}
]
Expand Down Expand Up @@ -819,7 +842,7 @@ async def test_s3_key_hook_get_files_bucket_keys_list(self, mock_bucket_keys, mo
"Key": mock_response_bucket_key,
"Size": 0,
"ETag": "etag1",
"LastModified": datetime(2020, 8, 14, 17, 19, 34),
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
}
for mock_response_bucket_key in mock_response_bucket_keys
]
Expand Down Expand Up @@ -857,8 +880,16 @@ async def test_s3_key_hook_list_keys_async(self):
test_resp_iter = [
{
"Contents": [
{"Key": "test_key", "ETag": "etag1", "LastModified": datetime(2020, 8, 14, 17, 19, 34)},
{"Key": "test_key2", "ETag": "etag2", "LastModified": datetime(2020, 8, 14, 17, 19, 34)},
{
"Key": "test_key",
"ETag": "etag1",
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
},
{
"Key": "test_key2",
"ETag": "etag2",
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
},
]
}
]
Expand Down Expand Up @@ -985,13 +1016,13 @@ async def test_s3__check_key_async_without_wildcard_match_and_get_none(self, moc
{
"Key": "test/example_s3_test_file.txt",
"ETag": "etag1",
"LastModified": datetime(2020, 8, 14, 17, 19, 34),
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
"Size": 0,
},
{
"Key": "test_key2",
"ETag": "etag2",
"LastModified": datetime(2020, 8, 14, 17, 19, 34),
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
"Size": 0,
},
],
Expand All @@ -1002,13 +1033,13 @@ async def test_s3__check_key_async_without_wildcard_match_and_get_none(self, moc
{
"Key": "test/example_aeoua.txt",
"ETag": "etag1",
"LastModified": datetime(2020, 8, 14, 17, 19, 34),
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
"Size": 0,
},
{
"Key": "test_key2",
"ETag": "etag2",
"LastModified": datetime(2020, 8, 14, 17, 19, 34),
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
"Size": 0,
},
],
Expand Down Expand Up @@ -1052,7 +1083,7 @@ async def test__check_key_async_with_use_regex(self, mock_get_bucket_key, key, p
{
"Key": key,
"ETag": "etag1",
"LastModified": datetime(2020, 8, 14, 17, 19, 34),
"LastModified": timezone.datetime(2020, 8, 14, 17, 19, 34),
"Size": 0,
},
]
Expand Down Expand Up @@ -1140,7 +1171,7 @@ async def test_s3_key_hook_is_keys_unchanged_exception_async(self, mock_list_key
@pytest.mark.asyncio
@mock.patch.object(S3Hook, "_list_keys_async", autospec=True)
async def test_s3_key_hook_is_keys_unchanged_success_async(self, mock_list_keys, time_machine):
frozen_dt = std_datetime(2026, 1, 1, 12, 0, 5, tzinfo=timezone.utc)
frozen_dt = std_datetime(2026, 1, 1, 12, 0, 5, tzinfo=std_timezone.utc)
time_machine.move_to(frozen_dt, tick=False)
mock_list_keys.return_value = ["test"]

Expand Down Expand Up @@ -1262,7 +1293,7 @@ async def test_s3_key_hook_is_keys_unchanged_pending_async_with_tzinfo(self, moc
previous_objects=set(),
inactivity_seconds=0,
allow_delete=False,
last_activity_time=std_datetime.now(timezone.utc),
last_activity_time=std_datetime.now(std_timezone.utc),
)
assert response.get("status") == "pending"

Expand Down
Loading
Loading