Additional usages of timezone

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__)
 

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:

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
 

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"

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"

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 = {

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"
 

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
 

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"

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
 

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
 

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
 

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
 

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
 

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
 

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"
 

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")

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
 

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
 

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
 

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
 

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