Refactor unneeded jumps in providers (#33833)

apache · Sep 1, 2023 · 875387a · 875387a
1 parent 0e1c106
commit 875387a
Show file tree

Hide file tree

Showing 17 changed files with 108 additions and 142 deletions.
diff --git a/airflow/providers/amazon/aws/hooks/datasync.py b/airflow/providers/amazon/aws/hooks/datasync.py
@@ -301,25 +301,18 @@ def wait_for_task_execution(self, task_execution_arn: str, max_iterations: int =
  if not task_execution_arn:
  raise AirflowBadRequest("task_execution_arn not specified")
 
- status = None
- iterations = max_iterations
- while status is None or status in self.TASK_EXECUTION_INTERMEDIATE_STATES:
+ for _ in range(max_iterations):
  task_execution = self.get_conn().describe_task_execution(TaskExecutionArn=task_execution_arn)
  status = task_execution["Status"]
  self.log.info("status=%s", status)
- iterations -= 1
- if status in self.TASK_EXECUTION_FAILURE_STATES:
- break
  if status in self.TASK_EXECUTION_SUCCESS_STATES:
- break
- if iterations <= 0:
- break
+ return True
+ elif status in self.TASK_EXECUTION_FAILURE_STATES:
+ return False
+ elif status is None or status in self.TASK_EXECUTION_INTERMEDIATE_STATES:
+ time.sleep(self.wait_interval_seconds)
+ else:
+ raise AirflowException(f"Unknown status: {status}") # Should never happen
  time.sleep(self.wait_interval_seconds)
-
- if status in self.TASK_EXECUTION_SUCCESS_STATES:
- return True
- if status in self.TASK_EXECUTION_FAILURE_STATES:
- return False
- if iterations <= 0:
+ else:
  raise AirflowTaskTimeout("Max iterations exceeded!")
- raise AirflowException(f"Unknown status: {status}") # Should never happen
diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py
@@ -252,12 +252,12 @@ def multi_stream_iter(self, log_group: str, streams: list, positions=None) -> Ge
  ]
  events: list[Any | None] = []
  for event_stream in event_iters:
- if not event_stream:
- events.append(None)
- continue
- try:
- events.append(next(event_stream))
- except StopIteration:
+ if event_stream:
+ try:
+  events.append(next(event_stream))
+  except StopIteration:
+  events.append(None)
+ else:
  events.append(None)
 
  while any(events):

diff --git a/airflow/providers/amazon/aws/sensors/sqs.py b/airflow/providers/amazon/aws/sensors/sqs.py
@@ -204,12 +204,12 @@ def poke(self, context: Context):
 
  if "Successful" not in response:
  raise AirflowException(f"Delete SQS Messages failed {response} for messages {messages}")
- if not message_batch:
+ if message_batch:
+ context["ti"].xcom_push(key="messages", value=message_batch)
+ return True
+ else:
  return False
 
- context["ti"].xcom_push(key="messages", value=message_batch)
- return True
-
  @deprecated(reason="use `hook` property instead.")
  def get_hook(self) -> SqsHook:
  """Create and return an SqsHook."""

diff --git a/airflow/providers/amazon/aws/utils/sqs.py b/airflow/providers/amazon/aws/utils/sqs.py
@@ -79,13 +79,9 @@ def filter_messages_jsonpath(messages, message_filtering_match_values, message_f
  # Body is a string, deserialize to an object and then parse
  body = json.loads(body)
  results = jsonpath_expr.find(body)
- if not results:
- continue
- if message_filtering_match_values is None:
+ if results and (
+ message_filtering_match_values is None
+ or any(result.value in message_filtering_match_values for result in results)
+ ):
  filtered_messages.append(message)
- continue
- for result in results:
- if result.value in message_filtering_match_values:
- filtered_messages.append(message)
- break
  return filtered_messages
diff --git a/airflow/providers/cncf/kubernetes/utils/delete_from.py b/airflow/providers/cncf/kubernetes/utils/delete_from.py
@@ -81,9 +81,7 @@ def delete_from_yaml(
  **kwargs,
 ):
  for yml_document in yaml_objects:
- if yml_document is None:
- continue
- else:
+ if yml_document is not None:
  delete_from_dict(
  k8s_client=k8s_client,
  data=yml_document,

diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -2206,27 +2206,25 @@ def run_query(
  if param_name == "schemaUpdateOptions" and param:
  self.log.info("Adding experimental 'schemaUpdateOptions': %s", schema_update_options)
 
- if param_name != "destinationTable":
- continue
-
- for key in ["projectId", "datasetId", "tableId"]:
- if key not in configuration["query"]["destinationTable"]:
- raise ValueError(
- "Not correct 'destinationTable' in "
- "api_resource_configs. 'destinationTable' "
- "must be a dict with {'projectId':'', "
- "'datasetId':'', 'tableId':''}"
+ if param_name == "destinationTable":
+ for key in ["projectId", "datasetId", "tableId"]:
+ if key not in configuration["query"]["destinationTable"]:
+ raise ValueError(
+ "Not correct 'destinationTable' in "
+ "api_resource_configs. 'destinationTable' "
+ "must be a dict with {'projectId':'', "
+ "'datasetId':'', 'tableId':''}"
+ )
+ else:
+ configuration["query"].update(
+ {
+ "allowLargeResults": allow_large_results,
+ "flattenResults": flatten_results,
+ "writeDisposition": write_disposition,
+ "createDisposition": create_disposition,
+ }
  )
 
- configuration["query"].update(
- {
- "allowLargeResults": allow_large_results,
- "flattenResults": flatten_results,
- "writeDisposition": write_disposition,
- "createDisposition": create_disposition,
- }
- )
-
  if (
  "useLegacySql" in configuration["query"]
  and configuration["query"]["useLegacySql"]

diff --git a/airflow/providers/google/cloud/hooks/datafusion.py b/airflow/providers/google/cloud/hooks/datafusion.py
@@ -371,12 +371,12 @@ def delete_pipeline(
  self._check_response_status_and_data(
  response, f"Deleting a pipeline failed with code {response.status}: {response.data}"
  )
- if response.status == 200:
- break
  except ConflictException as exc:
  self.log.info(exc)
  sleep(time_to_wait)
- continue
+ else:
+ if response.status == 200:
+ break
 
  def list_pipelines(
  self,

diff --git a/airflow/providers/google/cloud/hooks/gcs.py b/airflow/providers/google/cloud/hooks/gcs.py
@@ -361,7 +361,6 @@ def download(
  # Wait with exponential backoff scheme before retrying.
  timeout_seconds = 2 ** (num_file_attempts - 1)
  time.sleep(timeout_seconds)
- continue
 
  def download_as_byte_array(
  self,
@@ -508,28 +507,23 @@ def _call_with_retry(f: Callable[[], None]) -> None:
 
  :param f: Callable that should be retried.
  """
- num_file_attempts = 0
-
- while num_file_attempts < num_max_attempts:
+ for attempt in range(1, 1 + num_max_attempts):
  try:
- num_file_attempts += 1
  f()
-
  except GoogleCloudError as e:
- if num_file_attempts == num_max_attempts:
+ if attempt == num_max_attempts:
  self.log.error(
  "Upload attempt of object: %s from %s has failed. Attempt: %s, max %s.",
  object_name,
  object_name,
- num_file_attempts,
+ attempt,
  num_max_attempts,
  )
  raise e
 
  # Wait with exponential backoff scheme before retrying.
- timeout_seconds = 2 ** (num_file_attempts - 1)
+ timeout_seconds = 2 ** (attempt - 1)
  time.sleep(timeout_seconds)
- continue
 
  client = self.get_conn()
  bucket = client.bucket(bucket_name, user_project=user_project)

diff --git a/airflow/providers/google/cloud/log/gcs_task_handler.py b/airflow/providers/google/cloud/log/gcs_task_handler.py
@@ -243,9 +243,7 @@ def gcs_write(self, log, remote_log_location) -> bool:
  old_log = blob.download_as_bytes().decode()
  log = "\n".join([old_log, log]) if old_log else log
  except Exception as e:
- if self.no_log_found(e):
- pass
- else:
+ if not self.no_log_found(e):
  log += self._add_message(
  f"Error checking for previous log; if exists, may be overwritten: {e}"
  )

diff --git a/airflow/providers/google/cloud/operators/compute.py b/airflow/providers/google/cloud/operators/compute.py
@@ -174,14 +174,13 @@ def __init__(
  def check_body_fields(self) -> None:
  required_params = ["machine_type", "disks", "network_interfaces"]
  for param in required_params:
- if param in self.body:
- continue
- readable_param = param.replace("_", " ")
- raise AirflowException(
- f"The body '{self.body}' should contain at least {readable_param} for the new operator "
- f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
- f"for more details about body fields description."
- )
+ if param not in self.body:
+ readable_param = param.replace("_", " ")
+ raise AirflowException(
+ f"The body '{self.body}' should contain at least {readable_param} for the new operator "
+ f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
+ f"for more details about body fields description."
+ )
 
  def _validate_inputs(self) -> None:
  super()._validate_inputs()
@@ -915,14 +914,13 @@ def __init__(
  def check_body_fields(self) -> None:
  required_params = ["machine_type", "disks", "network_interfaces"]
  for param in required_params:
- if param in self.body["properties"]:
- continue
- readable_param = param.replace("_", " ")
- raise AirflowException(
- f"The body '{self.body}' should contain at least {readable_param} for the new operator "
- f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
- f"for more details about body fields description."
- )
+ if param not in self.body["properties"]:
+ readable_param = param.replace("_", " ")
+ raise AirflowException(
+ f"The body '{self.body}' should contain at least {readable_param} for the new operator "
+ f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
+ f"for more details about body fields description."
+ )
 
  def _validate_all_body_fields(self) -> None:
  if self._field_validator:
@@ -1500,14 +1498,13 @@ def __init__(
  def check_body_fields(self) -> None:
  required_params = ["base_instance_name", "target_size", "instance_template"]
  for param in required_params:
- if param in self.body:
- continue
- readable_param = param.replace("_", " ")
- raise AirflowException(
- f"The body '{self.body}' should contain at least {readable_param} for the new operator "
- f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
- f"for more details about body fields description."
- )
+ if param not in self.body:
+ readable_param = param.replace("_", " ")
+ raise AirflowException(
+ f"The body '{self.body}' should contain at least {readable_param} for the new operator "
+ f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
+ f"for more details about body fields description."
+ )
 
  def _validate_all_body_fields(self) -> None:
  if self._field_validator:

diff --git a/airflow/providers/google/cloud/operators/gcs.py b/airflow/providers/google/cloud/operators/gcs.py
@@ -797,9 +797,8 @@ def execute(self, context: Context) -> list[str]:
  num_max_attempts=self.download_num_attempts,
  )
  except GoogleCloudError:
- if self.download_continue_on_fail:
- continue
- raise
+ if not self.download_continue_on_fail:
+ raise
 
  self.log.info("Starting the transformation")
  cmd = [self.transform_script] if isinstance(self.transform_script, str) else self.transform_script
@@ -847,9 +846,8 @@ def execute(self, context: Context) -> list[str]:
  )
  files_uploaded.append(str(upload_file_name))
  except GoogleCloudError:
- if self.upload_continue_on_fail:
- continue
- raise
+ if not self.upload_continue_on_fail:
+ raise
 
  return files_uploaded
 

diff --git a/airflow/providers/microsoft/azure/hooks/data_factory.py b/airflow/providers/microsoft/azure/hooks/data_factory.py
@@ -855,8 +855,8 @@ def wait_for_pipeline_run_status(
  except ServiceRequestError:
  if executed_after_token_refresh:
  self.refresh_conn()
-  continue
- raise
+ else:
+  raise
 
  return pipeline_run_status in expected_statuses
 

diff --git a/airflow/providers/microsoft/azure/triggers/data_factory.py b/airflow/providers/microsoft/azure/triggers/data_factory.py
@@ -103,8 +103,8 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
  if executed_after_token_refresh:
  await hook.refresh_conn()
  executed_after_token_refresh = False
-  continue
- raise
+ else:
+  raise
  except Exception as e:
  yield TriggerEvent({"status": "error", "message": str(e)})
 
@@ -207,8 +207,8 @@ async def run(self) -> AsyncIterator[TriggerEvent]:
  if executed_after_token_refresh:
  await hook.refresh_conn()
  executed_after_token_refresh = False
-  continue
- raise
+ else:
+  raise
 
  yield TriggerEvent(
  {

diff --git a/airflow/providers/openlineage/utils/utils.py b/airflow/providers/openlineage/utils/utils.py
@@ -208,16 +208,14 @@ def _include_fields(self):
  raise Exception("Don't use both includes and excludes.")
  if self.includes:
  for field in self.includes:
- if field in self._fields or not hasattr(self.obj, field):
- continue
- setattr(self, field, getattr(self.obj, field))
- self._fields.append(field)
+ if field not in self._fields and hasattr(self.obj, field):
+ setattr(self, field, getattr(self.obj, field))
+ self._fields.append(field)
  else:
  for field, val in self.obj.__dict__.items():
- if field in self._fields or field in self.excludes or field in self.renames:
- continue
- setattr(self, field, val)
- self._fields.append(field)
+ if field not in self._fields and field not in self.excludes and field not in self.renames:
+ setattr(self, field, val)
+ self._fields.append(field)
 
 
 class DagInfo(InfoJsonEncodable):

diff --git a/airflow/providers/smtp/hooks/smtp.py b/airflow/providers/smtp/hooks/smtp.py
@@ -87,14 +87,14 @@ def get_conn(self) -> SmtpHook:
  try:
  self.smtp_client = self._build_client()
  except smtplib.SMTPServerDisconnected:
- if attempt < self.smtp_retry_limit:
- continue
-  raise AirflowException("Unable to connect to smtp server")
- if self.smtp_starttls:
- self.smtp_client.starttls()
- if self.smtp_user and self.smtp_password:
- self.smtp_client.login(self.smtp_user, self.smtp_password)
- break
+ if attempt == self.smtp_retry_limit:
+ raise AirflowException("Unable to connect to smtp server")
+ else:
+  if self.smtp_starttls:
+  self.smtp_client.starttls()
+  if self.smtp_user and self.smtp_password:
+  self.smtp_client.login(self.smtp_user, self.smtp_password)
+  break
 
  return self
 
@@ -234,10 +234,10 @@ def send_email_smtp(
  from_addr=from_email, to_addrs=recipients, msg=mime_msg.as_string()
  )
  except smtplib.SMTPServerDisconnected as e:
- if attempt < self.smtp_retry_limit:
- continue
-  raise e
- break
+ if attempt == self.smtp_retry_limit:
+ raise e
+ else:
+  break
 
  def _build_mime_message(
  self,