alexott commented on a change in pull request #19835:
URL: https://github.com/apache/airflow/pull/19835#discussion_r757665326
##########
File path: airflow/providers/databricks/hooks/databricks.py
##########
@@ -64,10 +66,12 @@
class RunState:
"""Utility class for the run state concept of Databricks runs."""
- def __init__(self, life_cycle_state: str, result_state: str,
state_message: str) -> None:
+ def __init__(
+ self, life_cycle_state: str, state_message: str, result_state: str =
None, *args, **kwargs
Review comment:
why do we need `*args, **kwargs` ? Also, why to change order of the
parameters? They are logical now: lifecycle -> result state -> state message.
##########
File path: airflow/providers/databricks/hooks/databricks.py
##########
@@ -356,31 +353,31 @@ def _do_api_call(self, endpoint_info, json):
def _log_request_error(self, attempt_num: int, error: str) -> None:
self.log.error('Attempt %s API Request to Databricks failed with
reason: %s', attempt_num, error)
- def run_now(self, json: dict) -> str:
+ def run_now(self, json: dict) -> int:
Review comment:
can it break existing code? for example if people are using this result
to concatenate with log string without using `str`?
##########
File path: airflow/providers/databricks/hooks/databricks.py
##########
@@ -235,21 +241,18 @@ def _get_aad_token(self, resource: str) -> str:
attempt_num += 1
sleep(self.retry_delay)
- def _fill_aad_tokens(self, headers: dict) -> str:
+ def _fill_aad_headers(self, headers: dict) -> dict:
"""
- Fills headers if necessary (SPN is outside of the workspace) and
generates AAD token
+ Fills AAD headers if necessary (SPN is outside of the workspace)
:param headers: dictionary with headers to fill-in
- :return: AAD token
+ :return: dictionary with filled AAD headers
"""
- # SP is outside of the workspace
- if 'azure_resource_id' in self.databricks_conn.extra_dejson:
Review comment:
I would keep this check inside the function, because it could be called
by accident (in the future). maybe call it `_fill_aad_headers_if_needed`?
##########
File path: airflow/providers/databricks/hooks/databricks.py
##########
@@ -522,6 +515,20 @@ def uninstall(self, json: dict) -> None:
"""
self._do_api_call(UNINSTALL_LIBS_ENDPOINT, json)
+ @staticmethod
+ def _is_aad_token_valid(aad_token: dict) -> bool:
Review comment:
why do we need a separate function that is called from one place?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
