Refactor: consolidate import time in providers (#34402)

apache · Oct 4, 2023 · 99f3203 · 99f3203
1 parent 99eeb84
commit 99f3203
Show file tree

Hide file tree

Showing 18 changed files with 44 additions and 44 deletions.
diff --git a/airflow/providers/alibaba/cloud/operators/analyticdb_spark.py b/airflow/providers/alibaba/cloud/operators/analyticdb_spark.py
@@ -17,8 +17,8 @@
 # under the License.
 from __future__ import annotations
 
+import time
 from functools import cached_property
-from time import sleep
 from typing import TYPE_CHECKING, Any, Sequence
 
 from deprecated.classic import deprecated
@@ -78,7 +78,7 @@ def poll_for_termination(self, app_id: str) -> None:
  state = self.hook.get_spark_state(app_id)
  while AppState(state) not in AnalyticDBSparkHook.TERMINAL_STATES:
  self.log.debug("Application with id %s is in state: %s", app_id, state)
- sleep(self.polling_interval)
+ time.sleep(self.polling_interval)
  state = self.hook.get_spark_state(app_id)
  self.log.info("Application with id %s terminated with state: %s", app_id, state)
  self.log.info(

diff --git a/airflow/providers/amazon/aws/hooks/batch_client.py b/airflow/providers/amazon/aws/hooks/batch_client.py
@@ -28,7 +28,7 @@
 
 import itertools
 import random
-from time import sleep
+import time
 from typing import TYPE_CHECKING, Callable
 
 import botocore.client
@@ -549,7 +549,7 @@ def delay(delay: int | float | None = None) -> None:
  delay = random.uniform(BatchClientHook.DEFAULT_DELAY_MIN, BatchClientHook.DEFAULT_DELAY_MAX)
  else:
  delay = BatchClientHook.add_jitter(delay)
- sleep(delay)
+ time.sleep(delay)
 
  @staticmethod
  def exponential_delay(tries: int) -> float:

diff --git a/airflow/providers/amazon/aws/hooks/elasticache_replication_group.py b/airflow/providers/amazon/aws/hooks/elasticache_replication_group.py
@@ -17,7 +17,7 @@
 # under the License.
 from __future__ import annotations
 
-from time import sleep
+import time
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
@@ -160,7 +160,7 @@ def wait_for_availability(
 
  self.log.info("Poke retry %s. Sleep time %s seconds. Sleeping...", num_tries, sleep_time)
 
- sleep(sleep_time)
+ time.sleep(sleep_time)
 
  sleep_time *= exponential_back_off_factor
 
@@ -240,7 +240,7 @@ def wait_for_deletion(
 
  self.log.info("Poke retry %s. Sleep time %s seconds. Sleeping...", num_tries, sleep_time)
 
- sleep(sleep_time)
+ time.sleep(sleep_time)
 
  sleep_time *= exponential_back_off_factor
 

diff --git a/airflow/providers/amazon/aws/hooks/emr.py b/airflow/providers/amazon/aws/hooks/emr.py
@@ -18,8 +18,8 @@
 from __future__ import annotations
 
 import json
+import time
 import warnings
-from time import sleep
 from typing import Any
 
 from botocore.exceptions import ClientError
@@ -509,7 +509,7 @@ def poll_query_status(
  final_query_state = query_state
  break
  try_number += 1
- sleep(poll_interval)
+ time.sleep(poll_interval)
  return final_query_state
 
  def stop_query(self, job_id: str) -> dict:

diff --git a/airflow/providers/amazon/aws/hooks/redshift_data.py b/airflow/providers/amazon/aws/hooks/redshift_data.py
@@ -17,8 +17,8 @@
 # under the License.
 from __future__ import annotations
 
+import time
 from pprint import pformat
-from time import sleep
 from typing import TYPE_CHECKING, Any, Iterable
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
@@ -127,7 +127,7 @@ def wait_for_results(self, statement_id, poll_interval):
  )
  else:
  self.log.info("Query %s", status)
- sleep(poll_interval)
+ time.sleep(poll_interval)
 
  def get_table_primary_key(
  self,

diff --git a/airflow/providers/amazon/aws/hooks/s3.py b/airflow/providers/amazon/aws/hooks/s3.py
@@ -26,6 +26,7 @@
 import os
 import re
 import shutil
+import time
 import warnings
 from contextlib import suppress
 from copy import deepcopy
@@ -35,7 +36,6 @@
 from io import BytesIO
 from pathlib import Path
 from tempfile import NamedTemporaryFile, gettempdir
-from time import sleep
 from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast
 from urllib.parse import urlsplit
 from uuid import uuid4
@@ -1289,7 +1289,7 @@ def delete_bucket(self, bucket_name: str, force_delete: bool = False, max_retrie
  if not bucket_keys:
  break
  if retry: # Avoid first loop
- sleep(500)
+ time.sleep(500)
 
  self.delete_objects(bucket=bucket_name, keys=bucket_keys)
 

diff --git a/airflow/providers/amazon/aws/operators/appflow.py b/airflow/providers/amazon/aws/operators/appflow.py
@@ -16,10 +16,10 @@
 # under the License.
 from __future__ import annotations
 
+import time
 import warnings
 from datetime import datetime, timedelta
 from functools import cached_property
-from time import sleep
 from typing import TYPE_CHECKING, cast
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
@@ -107,7 +107,7 @@ def execute(self, context: Context) -> None:
  self._update_flow()
  # while schedule flows will pick up the update right away, on-demand flows might use out of date
  # info if triggered right after an update, so we need to wait a bit for the DB to be consistent.
- sleep(AppflowBaseOperator.UPDATE_PROPAGATION_TIME)
+ time.sleep(AppflowBaseOperator.UPDATE_PROPAGATION_TIME)
 
  self._run_flow(context)
 

diff --git a/airflow/providers/apache/livy/operators/livy.py b/airflow/providers/apache/livy/operators/livy.py
@@ -17,8 +17,8 @@
 """This module contains the Apache Livy operator."""
 from __future__ import annotations
 
+import time
 from functools import cached_property
-from time import sleep
 from typing import TYPE_CHECKING, Any, Sequence
 
 from deprecated.classic import deprecated
@@ -189,7 +189,7 @@ def poll_for_termination(self, batch_id: int | str) -> None:
  state = self.hook.get_batch_state(batch_id, retry_args=self.retry_args)
  while state not in self.hook.TERMINAL_STATES:
  self.log.debug("Batch with id %s is in state: %s", batch_id, state.value)
- sleep(self._polling_interval)
+ time.sleep(self._polling_interval)
  state = self.hook.get_batch_state(batch_id, retry_args=self.retry_args)
  self.log.info("Batch with id %s terminated with state: %s", batch_id, state.value)
  self.hook.dump_batch_logs(batch_id)

diff --git a/airflow/providers/elasticsearch/log/es_task_handler.py b/airflow/providers/elasticsearch/log/es_task_handler.py
@@ -21,10 +21,10 @@
 import inspect
 import logging
 import sys
+import time
 import warnings
 from collections import defaultdict
 from operator import attrgetter
-from time import time
 from typing import TYPE_CHECKING, Any, Callable, List, Tuple
 from urllib.parse import quote, urlparse
 
@@ -372,7 +372,7 @@ def es_read(self, log_id: str, offset: int | str, metadata: dict) -> list | Elas
 
  def emit(self, record):
  if self.handler:
- setattr(record, self.offset_field, int(time() * (10**9)))
+ setattr(record, self.offset_field, int(time.time() * (10**9)))
  self.handler.emit(record)
 
  def set_context(self, ti: TaskInstance) -> None:

diff --git a/airflow/providers/google/cloud/hooks/cloud_batch.py b/airflow/providers/google/cloud/hooks/cloud_batch.py
@@ -19,7 +19,7 @@
 
 import itertools
 import json
-from time import sleep
+import time
 from typing import TYPE_CHECKING, Iterable, Sequence
 
 from google.cloud.batch import ListJobsRequest, ListTasksRequest
@@ -152,7 +152,7 @@ def wait_for_job(
  ):
  return job
  else:
- sleep(polling_period_seconds)
+ time.sleep(polling_period_seconds)
  except Exception as e:
  self.log.exception("Exception occurred while checking for job completion.")
  raise e

diff --git a/airflow/providers/google/cloud/hooks/datafusion.py b/airflow/providers/google/cloud/hooks/datafusion.py
@@ -20,7 +20,7 @@
 import asyncio
 import json
 import os
-from time import monotonic, sleep
+import time
 from typing import Any, Dict, Sequence
 from urllib.parse import quote, urlencode, urljoin
 
@@ -91,7 +91,7 @@ def __init__(
  def wait_for_operation(self, operation: dict[str, Any]) -> dict[str, Any]:
  """Waits for long-lasting operation to complete."""
  for time_to_wait in exponential_sleep_generator(initial=10, maximum=120):
- sleep(time_to_wait)
+ time.sleep(time_to_wait)
  operation = (
  self.get_conn().projects().locations().operations().get(name=operation.get("name")).execute()
  )
@@ -115,9 +115,9 @@ def wait_for_pipeline_state(
  """Polls pipeline state and raises an exception if the state fails or times out."""
  failure_states = failure_states or FAILURE_STATES
  success_states = success_states or SUCCESS_STATES
- start_time = monotonic()
+ start_time = time.monotonic()
  current_state = None
- while monotonic() - start_time < timeout:
+ while time.monotonic() - start_time < timeout:
  try:
  workflow = self.get_pipeline_workflow(
  pipeline_name=pipeline_name,
@@ -135,7 +135,7 @@ def wait_for_pipeline_state(
  raise AirflowException(
  f"Pipeline {pipeline_name} state {current_state} is not one of {success_states}"
  )
- sleep(30)
+ time.sleep(30)
 
  # Time is up!
  raise AirflowException(
@@ -393,7 +393,7 @@ def delete_pipeline(
  )
  except ConflictException as exc:
  self.log.info(exc)
- sleep(time_to_wait)
+ time.sleep(time_to_wait)
  else:
  if response.status == 200:
  break

diff --git a/airflow/providers/google/cloud/operators/datafusion.py b/airflow/providers/google/cloud/operators/datafusion.py
@@ -17,7 +17,7 @@
 """This module contains Google DataFusion operators."""
 from __future__ import annotations
 
-from time import sleep
+import time
 from typing import TYPE_CHECKING, Any, Sequence
 
 from google.api_core.retry import exponential_sleep_generator
@@ -267,7 +267,7 @@ def execute(self, context: Context) -> dict:
  for time_to_wait in exponential_sleep_generator(initial=10, maximum=120):
  if instance["state"] != "CREATING":
  break
- sleep(time_to_wait)
+ time.sleep(time_to_wait)
  instance = hook.get_instance(
  instance_name=self.instance_name, location=self.location, project_id=self.project_id
  )

diff --git a/airflow/providers/google/cloud/operators/dataplex.py b/airflow/providers/google/cloud/operators/dataplex.py
@@ -18,7 +18,7 @@
 
 from __future__ import annotations
 
-from time import sleep
+import time
 from typing import TYPE_CHECKING, Any, Sequence
 
 from airflow.exceptions import AirflowException
@@ -165,7 +165,7 @@ def execute(self, context: Context) -> dict:
  )
  if task["state"] != "CREATING":
  break
- sleep(time_to_wait)
+ time.sleep(time_to_wait)
 
  return Task.to_dict(task)
 
@@ -534,7 +534,7 @@ def execute(self, context: Context) -> dict:
  )
  if lake["state"] != "CREATING":
  break
- sleep(time_to_wait)
+ time.sleep(time_to_wait)
  DataplexLakeLink.persist(
  context=context,
  task_instance=self,

diff --git a/airflow/providers/google/cloud/operators/dataproc_metastore.py b/airflow/providers/google/cloud/operators/dataproc_metastore.py
@@ -18,7 +18,7 @@
 """This module contains Google Dataproc Metastore operators."""
 from __future__ import annotations
 
-from time import sleep
+import time
 from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -700,7 +700,7 @@ def _wait_for_export_metadata(self, hook: DataprocMetastoreHook):
  the SDK.
  """
  for time_to_wait in exponential_sleep_generator(initial=10, maximum=120):
- sleep(time_to_wait)
+ time.sleep(time_to_wait)
  service = hook.get_service(
  region=self.region,
  project_id=self.project_id,
@@ -986,7 +986,7 @@ def _wait_for_restore_service(self, hook: DataprocMetastoreHook):
  the SDK.
  """
  for time_to_wait in exponential_sleep_generator(initial=10, maximum=120):
- sleep(time_to_wait)
+ time.sleep(time_to_wait)
  service = hook.get_service(
  region=self.region,
  project_id=self.project_id,

diff --git a/airflow/providers/microsoft/azure/operators/container_instances.py b/airflow/providers/microsoft/azure/operators/container_instances.py
@@ -18,8 +18,8 @@
 from __future__ import annotations
 
 import re
+import time
 from collections import namedtuple
-from time import sleep
 from typing import TYPE_CHECKING, Any, Sequence
 
 from azure.mgmt.containerinstance.models import (
@@ -348,7 +348,7 @@ def _monitor_logging(self, resource_group: str, name: str) -> int:
  except Exception:
  self.log.exception("Exception while getting container groups")
 
- sleep(1)
+ time.sleep(1)
 
  def _log_last(self, logs: list | None, last_line_logged: Any) -> Any | None:
  if logs:

diff --git a/tests/providers/amazon/aws/hooks/test_batch_client.py b/tests/providers/amazon/aws/hooks/test_batch_client.py
@@ -427,7 +427,7 @@ def test_add_jitter(self):
  assert result <= width
 
  @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.random.uniform")
- @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.sleep")
+ @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.time.sleep")
  def test_delay_defaults(self, mock_sleep, mock_uniform):
  assert BatchClientHook.DEFAULT_DELAY_MIN == 1
  assert BatchClientHook.DEFAULT_DELAY_MAX == 10
@@ -439,21 +439,21 @@ def test_delay_defaults(self, mock_sleep, mock_uniform):
  mock_sleep.assert_called_once_with(0)
 
  @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.random.uniform")
- @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.sleep")
+ @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.time.sleep")
  def test_delay_with_zero(self, mock_sleep, mock_uniform):
  self.batch_client.delay(0)
  mock_uniform.assert_called_once_with(0, 1) # in add_jitter
  mock_sleep.assert_called_once_with(mock_uniform.return_value)
 
  @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.random.uniform")
- @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.sleep")
+ @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.time.sleep")
  def test_delay_with_int(self, mock_sleep, mock_uniform):
  self.batch_client.delay(5)
  mock_uniform.assert_called_once_with(4, 6) # in add_jitter
  mock_sleep.assert_called_once_with(mock_uniform.return_value)
 
  @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.random.uniform")
- @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.sleep")
+ @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.time.sleep")
  def test_delay_with_float(self, mock_sleep, mock_uniform):
  self.batch_client.delay(5.0)
  mock_uniform.assert_called_once_with(4.0, 6.0) # in add_jitter

diff --git a/tests/providers/microsoft/azure/operators/test_azure_container_instances.py b/tests/providers/microsoft/azure/operators/test_azure_container_instances.py
@@ -345,7 +345,7 @@ def test_execute_fails_with_incorrect_restart_policy(self, aci_mock):
  )
 
  @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook")
- @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.sleep")
+ @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.time.sleep")
  def test_execute_correct_sleep_cycle(self, sleep_mock, aci_mock):
  expected_cg1 = make_mock_container(state="Running", exit_code=0, detail_status="test")
  expected_cg2 = make_mock_container(state="Terminated", exit_code=0, detail_status="test")