From 66376ba8bc491e8003de62c113119f9da88c2a7a Mon Sep 17 00:00:00 2001
From: Niko Oliveira <onikolas@amazon.com>
Date: Wed, 30 Jul 2025 15:50:42 -0700
Subject: [PATCH 1/2] Additional usages of timezone

---
 .../amazon/aws/executors/utils/exponential_backoff_retry.py | 5 ++++-
 .../src/airflow/providers/amazon/aws/hooks/sagemaker.py     | 6 +++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/providers/amazon/src/airflow/providers/amazon/aws/executors/utils/exponential_backoff_retry.py b/providers/amazon/src/airflow/providers/amazon/aws/executors/utils/exponential_backoff_retry.py
index 1255f56f0431f..16a309ab448e7 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/executors/utils/exponential_backoff_retry.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/executors/utils/exponential_backoff_retry.py
@@ -20,7 +20,10 @@
 from collections.abc import Callable
 from datetime import datetime, timedelta
 
-from airflow.utils import timezone
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 log = logging.getLogger(__name__)
 
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker.py
index fad70bc4002a9..05c37763604a9 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker.py
@@ -36,7 +36,11 @@
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.tags import format_tags
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 
 class LogState:

From 24d666d6c6cd9ebd87957dc49fea2d2bdf15a6f5 Mon Sep 17 00:00:00 2001
From: Niko Oliveira <onikolas@amazon.com>
Date: Thu, 31 Jul 2025 10:13:03 -0700
Subject: [PATCH 2/2] Fixes for test classes

---
 .../aws/executors/aws_lambda/test_lambda_executor.py     | 6 +++++-
 .../amazon/tests/unit/amazon/aws/hooks/test_appflow.py   | 6 +++++-
 .../tests/unit/amazon/aws/operators/test_appflow.py      | 6 +++++-
 .../tests/unit/amazon/aws/operators/test_athena.py       | 9 ++++++---
 .../tests/unit/amazon/aws/operators/test_base_aws.py     | 6 +++++-
 .../unit/amazon/aws/operators/test_cloud_formation.py    | 6 +++++-
 .../tests/unit/amazon/aws/operators/test_datasync.py     | 9 ++++++---
 .../amazon/tests/unit/amazon/aws/operators/test_dms.py   | 6 +++++-
 .../unit/amazon/aws/operators/test_emr_add_steps.py      | 6 +++++-
 .../amazon/aws/operators/test_emr_create_job_flow.py     | 6 +++++-
 .../unit/amazon/aws/operators/test_emr_modify_cluster.py | 7 +++++--
 .../amazon/tests/unit/amazon/aws/operators/test_rds.py   | 6 +++++-
 .../unit/amazon/aws/operators/test_sagemaker_base.py     | 6 +++++-
 .../tests/unit/amazon/aws/sensors/test_base_aws.py       | 6 +++++-
 .../amazon/tests/unit/amazon/aws/sensors/test_ecs.py     | 6 +++++-
 .../amazon/tests/unit/amazon/aws/sensors/test_rds.py     | 6 +++++-
 .../amazon/tests/unit/amazon/aws/sensors/test_s3.py      | 6 +++++-
 .../amazon/tests/unit/amazon/aws/transfers/test_base.py  | 6 +++++-
 .../unit/amazon/aws/transfers/test_dynamodb_to_s3.py     | 6 +++++-
 .../tests/unit/amazon/aws/transfers/test_mongo_to_s3.py  | 6 +++++-
 20 files changed, 102 insertions(+), 25 deletions(-)

diff --git a/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py b/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
index a2def0b86a8bb..de0125625226b 100644
--- a/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
+++ b/providers/amazon/tests/unit/amazon/aws/executors/aws_lambda/test_lambda_executor.py
@@ -31,7 +31,11 @@
 from airflow.providers.amazon.aws.executors.aws_lambda import lambda_executor
 from airflow.providers.amazon.aws.executors.aws_lambda.lambda_executor import AwsLambdaExecutor
 from airflow.providers.amazon.aws.executors.aws_lambda.utils import CONFIG_GROUP_NAME, AllLambdaConfigKeys
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import TaskInstanceState
 from airflow.version import version as airflow_version_str
 
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_appflow.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_appflow.py
index 71949ae4b2017..b8ba2d927cd44 100644
--- a/providers/amazon/tests/unit/amazon/aws/hooks/test_appflow.py
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_appflow.py
@@ -24,7 +24,11 @@
 import pytest
 
 from airflow.providers.amazon.aws.hooks.appflow import AppflowHook
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 FLOW_NAME = "flow0"
 EXECUTION_ID = "ex_id"
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_appflow.py b/providers/amazon/tests/unit/amazon/aws/operators/test_appflow.py
index 49c49e8c51458..1f2b188997fa5 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_appflow.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_appflow.py
@@ -31,7 +31,11 @@
     AppflowRunFullOperator,
     AppflowRunOperator,
 )
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 CONN_ID = "aws_default"
 DAG_ID = "dag_id"
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_athena.py b/providers/amazon/tests/unit/amazon/aws/operators/test_athena.py
index 8a343a0be686e..6bee0fd332e6f 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_athena.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_athena.py
@@ -38,16 +38,19 @@
     SymlinksDatasetFacet,
 )
 from airflow.providers.openlineage.extractors import OperatorLineage
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import DagRunState
-from airflow.utils.timezone import datetime
 from airflow.utils.types import DagRunType
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 from unit.amazon.aws.utils.test_template_fields import validate_template_fields
 
 TEST_DAG_ID = "unit_tests"
-DEFAULT_DATE = datetime(2018, 1, 1)
+DEFAULT_DATE = timezone.datetime(2018, 1, 1)
 ATHENA_QUERY_ID = "eac29bf8-daa1-4ffc-b19a-0db31dc3b784"
 
 MOCK_DATA = {
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_base_aws.py b/providers/amazon/tests/unit/amazon/aws/operators/test_base_aws.py
index 53f033c4d6c5f..3cddb55eaa0c3 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_base_aws.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_base_aws.py
@@ -23,7 +23,11 @@
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
 from airflow.providers.amazon.version_compat import BaseHook
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 TEST_CONN = "aws_test_conn"
 
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_cloud_formation.py b/providers/amazon/tests/unit/amazon/aws/operators/test_cloud_formation.py
index bbcd41f2176c4..5643eb9c48be0 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_cloud_formation.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_cloud_formation.py
@@ -27,7 +27,11 @@
     CloudFormationCreateStackOperator,
     CloudFormationDeleteStackOperator,
 )
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 from unit.amazon.aws.utils.test_template_fields import validate_template_fields
 
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_datasync.py b/providers/amazon/tests/unit/amazon/aws/operators/test_datasync.py
index 50101edad0368..9bc8f69ffab79 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_datasync.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_datasync.py
@@ -28,16 +28,19 @@
 from airflow.providers.amazon.aws.hooks.datasync import DataSyncHook
 from airflow.providers.amazon.aws.links.datasync import DataSyncTaskLink
 from airflow.providers.amazon.aws.operators.datasync import DataSyncOperator
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import DagRunState
-from airflow.utils.timezone import datetime
 from airflow.utils.types import DagRunType
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 from unit.amazon.aws.utils.test_template_fields import validate_template_fields
 
 TEST_DAG_ID = "unit_tests"
-DEFAULT_DATE = datetime(2018, 1, 1)
+DEFAULT_DATE = timezone.datetime(2018, 1, 1)
 
 SOURCE_HOST_NAME = "airflow.host"
 SOURCE_SUBDIR = "airflow_subdir"
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_dms.py b/providers/amazon/tests/unit/amazon/aws/operators/test_dms.py
index 3418019ee7751..5cf3a8d707c90 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_dms.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_dms.py
@@ -45,7 +45,11 @@
     DmsReplicationDeprovisionedTrigger,
     DmsReplicationTerminalStatusTrigger,
 )
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_add_steps.py b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_add_steps.py
index 0cfa38014d1bb..59a88abd677dd 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_add_steps.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_add_steps.py
@@ -30,7 +30,11 @@
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.providers.amazon.aws.operators.emr import EmrAddStepsOperator
 from airflow.providers.amazon.aws.triggers.emr import EmrAddStepsTrigger
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_create_job_flow.py b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_create_job_flow.py
index 686f40c103762..68c6462d95d08 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_create_job_flow.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_create_job_flow.py
@@ -32,7 +32,11 @@
 from airflow.providers.amazon.aws.operators.emr import EmrCreateJobFlowOperator
 from airflow.providers.amazon.aws.triggers.emr import EmrCreateJobFlowTrigger
 from airflow.providers.amazon.aws.utils.waiter import WAITER_POLICY_NAME_MAPPING, WaitPolicy
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_modify_cluster.py b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_modify_cluster.py
index e9d7ccb50c360..7b8b83cd7f10c 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_modify_cluster.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_modify_cluster.py
@@ -1,4 +1,3 @@
-#
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
@@ -24,7 +23,11 @@
 from airflow.exceptions import AirflowException
 from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.operators.emr import EmrModifyClusterOperator
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 from unit.amazon.aws.utils.test_template_fields import validate_template_fields
 
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_rds.py b/providers/amazon/tests/unit/amazon/aws/operators/test_rds.py
index 565780db09ad4..b807799c8f82a 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_rds.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_rds.py
@@ -43,7 +43,11 @@
     RdsStopDbOperator,
 )
 from airflow.providers.amazon.aws.triggers.rds import RdsDbAvailableTrigger, RdsDbStoppedTrigger
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 from unit.amazon.aws.utils.test_template_fields import validate_template_fields
 
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_base.py b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_base.py
index 9c9515d4f88ea..dd922fba2acec 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_base.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_base.py
@@ -31,7 +31,11 @@
     SageMakerBaseOperator,
     SageMakerCreateExperimentOperator,
 )
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 
diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_base_aws.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_base_aws.py
index 9c6b66537b0d7..128e5981e9448 100644
--- a/providers/amazon/tests/unit/amazon/aws/sensors/test_base_aws.py
+++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_base_aws.py
@@ -23,7 +23,11 @@
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.version_compat import BaseHook
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 TEST_CONN = "aws_test_conn"
 
diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_ecs.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_ecs.py
index 9a5e9313b669c..282d37f3f7dee 100644
--- a/providers/amazon/tests/unit/amazon/aws/sensors/test_ecs.py
+++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_ecs.py
@@ -35,7 +35,11 @@
     EcsTaskStates,
     EcsTaskStateSensor,
 )
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.types import NOTSET
 
 _Operator = TypeVar("_Operator")
diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_rds.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_rds.py
index 3bdd8f673bf3b..2768c04efed76 100644
--- a/providers/amazon/tests/unit/amazon/aws/sensors/test_rds.py
+++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_rds.py
@@ -31,7 +31,11 @@
     RdsSnapshotExistenceSensor,
 )
 from airflow.providers.amazon.aws.utils.rds import RdsDbType
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 from unit.amazon.aws.utils.test_template_fields import validate_template_fields
 
diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_s3.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_s3.py
index aecaf0b91a574..08d7b0941f4ef 100644
--- a/providers/amazon/tests/unit/amazon/aws/sensors/test_s3.py
+++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_s3.py
@@ -30,7 +30,11 @@
 from airflow.models.variable import Variable
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor, S3KeysUnchangedSensor
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 
diff --git a/providers/amazon/tests/unit/amazon/aws/transfers/test_base.py b/providers/amazon/tests/unit/amazon/aws/transfers/test_base.py
index eb03cace7ea4a..ba3690182ecd2 100644
--- a/providers/amazon/tests/unit/amazon/aws/transfers/test_base.py
+++ b/providers/amazon/tests/unit/amazon/aws/transfers/test_base.py
@@ -23,7 +23,11 @@
 from airflow.models import DagRun, TaskInstance
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.providers.amazon.aws.transfers.base import AwsToAwsBaseOperator
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 
diff --git a/providers/amazon/tests/unit/amazon/aws/transfers/test_dynamodb_to_s3.py b/providers/amazon/tests/unit/amazon/aws/transfers/test_dynamodb_to_s3.py
index 34519383b6cec..88a1729480002 100644
--- a/providers/amazon/tests/unit/amazon/aws/transfers/test_dynamodb_to_s3.py
+++ b/providers/amazon/tests/unit/amazon/aws/transfers/test_dynamodb_to_s3.py
@@ -31,7 +31,11 @@
     DynamoDBToS3Operator,
     JSONEncoder,
 )
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 
diff --git a/providers/amazon/tests/unit/amazon/aws/transfers/test_mongo_to_s3.py b/providers/amazon/tests/unit/amazon/aws/transfers/test_mongo_to_s3.py
index 5401e557be759..7c590cb7306a6 100644
--- a/providers/amazon/tests/unit/amazon/aws/transfers/test_mongo_to_s3.py
+++ b/providers/amazon/tests/unit/amazon/aws/transfers/test_mongo_to_s3.py
@@ -24,7 +24,11 @@
 from airflow.models import DAG, DagRun, TaskInstance
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.providers.amazon.aws.transfers.mongo_to_s3 import MongoToS3Operator
-from airflow.utils import timezone
+
+try:
+    from airflow.sdk import timezone
+except ImportError:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType