diff --git a/airflow/providers/databricks/hooks/databricks_base.py b/airflow/providers/databricks/hooks/databricks_base.py index 9885e9a998058..6d0d929b5dc4e 100644 --- a/airflow/providers/databricks/hooks/databricks_base.py +++ b/airflow/providers/databricks/hooks/databricks_base.py @@ -62,6 +62,7 @@ TOKEN_REFRESH_LEAD_TIME = 120 AZURE_MANAGEMENT_ENDPOINT = "https://management.core.windows.net/" DEFAULT_DATABRICKS_SCOPE = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d" +OIDC_TOKEN_SERVICE_URL = "{}/oidc/v1/token" class BaseDatabricksHook(BaseHook): @@ -89,6 +90,7 @@ class BaseDatabricksHook(BaseHook): "azure_ad_endpoint", "azure_resource_id", "azure_tenant_id", + "service_principal_oauth", ] def __init__( @@ -107,8 +109,8 @@ def __init__( raise ValueError("Retry limit must be greater than or equal to 1") self.retry_limit = retry_limit self.retry_delay = retry_delay - self.aad_tokens: dict[str, dict] = {} - self.aad_timeout_seconds = 10 + self.oauth_tokens: dict[str, dict] = {} + self.token_timeout_seconds = 10 self.caller = caller def my_after_func(retry_state): @@ -210,6 +212,75 @@ def _a_get_retry_object(self) -> AsyncRetrying: """ return AsyncRetrying(**self.retry_args) + def _get_sp_token(self, resource: str) -> str: + """Function to get Service Principal token.""" + sp_token = self.oauth_tokens.get(resource) + if sp_token and self._is_oauth_token_valid(sp_token): + return sp_token["access_token"] + + self.log.info("Existing Service Principal token is expired, or going to expire soon. Refreshing...") + try: + for attempt in self._get_retry_object(): + with attempt: + resp = requests.post( + resource, + auth=HTTPBasicAuth(self.databricks_conn.login, self.databricks_conn.password), + data="grant_type=client_credentials&scope=all-apis", + headers={ + **self.user_agent_header, + "Content-Type": "application/x-www-form-urlencoded", + }, + timeout=self.token_timeout_seconds, + ) + + resp.raise_for_status() + jsn = resp.json() + jsn["expires_on"] = int(time.time() + jsn["expires_in"]) + + self._is_oauth_token_valid(jsn) + self.oauth_tokens[resource] = jsn + break + except RetryError: + raise AirflowException(f"API requests to Databricks failed {self.retry_limit} times. Giving up.") + except requests_exceptions.HTTPError as e: + raise AirflowException(f"Response: {e.response.content}, Status Code: {e.response.status_code}") + + return jsn["access_token"] + + async def _a_get_sp_token(self, resource: str) -> str: + """Async version of `_get_sp_token()`.""" + sp_token = self.oauth_tokens.get(resource) + if sp_token and self._is_oauth_token_valid(sp_token): + return sp_token["access_token"] + + self.log.info("Existing Service Principal token is expired, or going to expire soon. Refreshing...") + try: + async for attempt in self._a_get_retry_object(): + with attempt: + async with self._session.post( + resource, + auth=HTTPBasicAuth(self.databricks_conn.login, self.databricks_conn.password), + data="grant_type=client_credentials&scope=all-apis", + headers={ + **self.user_agent_header, + "Content-Type": "application/x-www-form-urlencoded", + }, + timeout=self.token_timeout_seconds, + ) as resp: + resp.raise_for_status() + jsn = await resp.json() + jsn["expires_on"] = int(time.time() + jsn["expires_in"]) + + self._is_oauth_token_valid(jsn) + self.oauth_tokens[resource] = jsn + break + except RetryError: + raise AirflowException(f"API requests to Databricks failed {self.retry_limit} times. Giving up.") + except requests_exceptions.HTTPError as e: + raise AirflowException(f"Response: {e.response.content}, Status Code: {e.response.status_code}") + + return jsn["access_token"] + def _get_aad_token(self, resource: str) -> str: """ Function to get AAD token for given resource. @@ -218,9 +289,9 @@ def _get_aad_token(self, resource: str) -> str: :param resource: resource to issue token to :return: AAD token, or raise an exception """ - aad_token = self.aad_tokens.get(resource) - if aad_token and self._is_aad_token_valid(aad_token): - return aad_token["token"] + aad_token = self.oauth_tokens.get(resource) + if aad_token and self._is_oauth_token_valid(aad_token): + return aad_token["access_token"] self.log.info("Existing AAD token is expired, or going to expire soon. Refreshing...") try: @@ -235,7 +306,7 @@ def _get_aad_token(self, resource: str) -> str: AZURE_METADATA_SERVICE_TOKEN_URL, params=params, headers={**self.user_agent_header, "Metadata": "true"}, - timeout=self.aad_timeout_seconds, + timeout=self.token_timeout_seconds, ) else: tenant_id = self.databricks_conn.extra_dejson["azure_tenant_id"] @@ -255,27 +326,21 @@ def _get_aad_token(self, resource: str) -> str: **self.user_agent_header, "Content-Type": "application/x-www-form-urlencoded", }, - timeout=self.aad_timeout_seconds, + timeout=self.token_timeout_seconds, ) resp.raise_for_status() jsn = resp.json() - if ( - "access_token" not in jsn - or jsn.get("token_type") != "Bearer" - or "expires_on" not in jsn - ): - raise AirflowException(f"Can't get necessary data from AAD token: {jsn}") - - token = jsn["access_token"] - self.aad_tokens[resource] = {"token": token, "expires_on": int(jsn["expires_on"])} + + self._is_oauth_token_valid(jsn) + self.oauth_tokens[resource] = jsn break except RetryError: raise AirflowException(f"API requests to Azure failed {self.retry_limit} times. Giving up.") except requests_exceptions.HTTPError as e: raise AirflowException(f"Response: {e.response.content}, Status Code: {e.response.status_code}") - return token + return jsn["access_token"] async def _a_get_aad_token(self, resource: str) -> str: """ @@ -284,9 +349,9 @@ async def _a_get_aad_token(self, resource: str) -> str: :param resource: resource to issue token to :return: AAD token, or raise an exception """ - aad_token = self.aad_tokens.get(resource) - if aad_token and self._is_aad_token_valid(aad_token): - return aad_token["token"] + aad_token = self.oauth_tokens.get(resource) + if aad_token and self._is_oauth_token_valid(aad_token): + return aad_token["access_token"] self.log.info("Existing AAD token is expired, or going to expire soon. Refreshing...") try: @@ -301,7 +366,7 @@ async def _a_get_aad_token(self, resource: str) -> str: url=AZURE_METADATA_SERVICE_TOKEN_URL, params=params, headers={**self.user_agent_header, "Metadata": "true"}, - timeout=self.aad_timeout_seconds, + timeout=self.token_timeout_seconds, ) as resp: resp.raise_for_status() jsn = await resp.json() @@ -323,26 +388,20 @@ async def _a_get_aad_token(self, resource: str) -> str: **self.user_agent_header, "Content-Type": "application/x-www-form-urlencoded", }, - timeout=self.aad_timeout_seconds, + timeout=self.token_timeout_seconds, ) as resp: resp.raise_for_status() jsn = await resp.json() - if ( - "access_token" not in jsn - or jsn.get("token_type") != "Bearer" - or "expires_on" not in jsn - ): - raise AirflowException(f"Can't get necessary data from AAD token: {jsn}") - - token = jsn["access_token"] - self.aad_tokens[resource] = {"token": token, "expires_on": int(jsn["expires_on"])} + + self._is_oauth_token_valid(jsn) + self.oauth_tokens[resource] = jsn break except RetryError: raise AirflowException(f"API requests to Azure failed {self.retry_limit} times. Giving up.") except aiohttp.ClientResponseError as err: raise AirflowException(f"Response: {err.message}, Status Code: {err.status}") - return token + return jsn["access_token"] def _get_aad_headers(self) -> dict: """ @@ -375,17 +434,18 @@ async def _a_get_aad_headers(self) -> dict: return headers @staticmethod - def _is_aad_token_valid(aad_token: dict) -> bool: + def _is_oauth_token_valid(token: dict, time_key="expires_on") -> bool: """ - Utility function to check AAD token hasn't expired yet. + Utility function to check if an OAuth token is valid and hasn't expired yet. - :param aad_token: dict with properties of AAD token + :param sp_token: dict with properties of OAuth token + :param time_key: name of the key that holds the time of expiration :return: true if token is valid, false otherwise """ - now = int(time.time()) - if aad_token["expires_on"] > (now + TOKEN_REFRESH_LEAD_TIME): - return True - return False + if "access_token" not in token or token.get("token_type", "") != "Bearer" or time_key not in token: + raise AirflowException(f"Can't get necessary data from OAuth token: {token}") + + return int(token[time_key]) > (int(time.time()) + TOKEN_REFRESH_LEAD_TIME) @staticmethod def _check_azure_metadata_service() -> None: @@ -443,6 +503,11 @@ def _get_token(self, raise_error: bool = False) -> str | None: self.log.info("Using AAD Token for managed identity.") self._check_azure_metadata_service() return self._get_aad_token(DEFAULT_DATABRICKS_SCOPE) + elif self.databricks_conn.extra_dejson.get("service_principal_oauth", False): + if self.databricks_conn.login == "" or self.databricks_conn.password == "": + raise AirflowException("Service Principal credentials aren't provided") + self.log.info("Using Service Principal Token.") + return self._get_sp_token(OIDC_TOKEN_SERVICE_URL.format(self.databricks_conn.host)) elif raise_error: raise AirflowException("Token authentication isn't configured") @@ -466,6 +531,11 @@ async def _a_get_token(self, raise_error: bool = False) -> str | None: self.log.info("Using AAD Token for managed identity.") await self._a_check_azure_metadata_service() return await self._a_get_aad_token(DEFAULT_DATABRICKS_SCOPE) + elif self.databricks_conn.extra_dejson.get("service_principal_oauth", False): + if self.databricks_conn.login == "" or self.databricks_conn.password == "": + raise AirflowException("Service Principal credentials aren't provided") + self.log.info("Using Service Principal Token.") + return await self._a_get_sp_token(OIDC_TOKEN_SERVICE_URL.format(self.databricks_conn.host)) elif raise_error: raise AirflowException("Token authentication isn't configured") diff --git a/docs/apache-airflow-providers-databricks/connections/databricks.rst b/docs/apache-airflow-providers-databricks/connections/databricks.rst index 6303702b7ec1a..908b12eac7f03 100644 --- a/docs/apache-airflow-providers-databricks/connections/databricks.rst +++ b/docs/apache-airflow-providers-databricks/connections/databricks.rst @@ -55,13 +55,15 @@ Host (required) Login (optional) * If authentication with *Databricks login credentials* is used then specify the ``username`` used to login to Databricks. - * If *authentication with Azure Service Principal* is used then specify the ID of the Azure Service Principal + * If authentication with *Azure Service Principal* is used then specify the ID of the Azure Service Principal * If authentication with *PAT* is used then either leave this field empty or use 'token' as login (both work, the only difference is that if login is empty then token will be sent in request header as Bearer token, if login is 'token' then it will be sent using Basic Auth which is allowed by Databricks API, this may be useful if you plan to reuse this connection with e.g. SimpleHttpOperator) + * If authentication with *Databricks Service Principal OAuth* is used then specify the ID of the Service Principal (Databricks on AWS) Password (optional) - * If authentication with *Databricks login credentials* is used then specify the ``password`` used to login to Databricks. + * If authentication with *Databricks login credentials* is used then specify the ``password`` used to login to Databricks. * If authentication with *Azure Service Principal* is used then specify the secret of the Azure Service Principal * If authentication with *PAT* is used, then specify PAT (recommended) + * If authentication with *Databricks Service Principal OAuth* is used then specify the secret of the Service Principal (Databricks on AWS) Extra (optional) Specify the extra parameter (as json dictionary) that can be used in the Databricks connection. @@ -70,6 +72,10 @@ Extra (optional) * ``token``: Specify PAT to use. Consider to switch to specification of PAT in the Password field as it's more secure. + Following parameters are necessary if using authentication with OAuth token for AWS Databricks Service Principal: + + * ``service_principal_oauth``: required boolean flag. If specified as ``true``, use the Client ID and Client Secret as the Username and Password. See `Authentication using OAuth for service principals `_. + Following parameters are necessary if using authentication with AAD token: * ``azure_tenant_id``: ID of the Azure Active Directory tenant diff --git a/tests/providers/databricks/hooks/test_databricks.py b/tests/providers/databricks/hooks/test_databricks.py index f55c55dfc364c..8644d95cd9c4c 100644 --- a/tests/providers/databricks/hooks/test_databricks.py +++ b/tests/providers/databricks/hooks/test_databricks.py @@ -43,6 +43,7 @@ AZURE_METADATA_SERVICE_INSTANCE_URL, AZURE_TOKEN_SERVICE_URL, DEFAULT_DATABRICKS_SCOPE, + OIDC_TOKEN_SERVICE_URL, TOKEN_REFRESH_LEAD_TIME, BearerAuth, ) @@ -689,13 +690,42 @@ def test_uninstall_libs_on_cluster(self, mock_requests): timeout=self.hook.timeout_seconds, ) - def test_is_aad_token_valid_returns_true(self): - aad_token = {"token": "my_token", "expires_on": int(time.time()) + TOKEN_REFRESH_LEAD_TIME + 10} - assert self.hook._is_aad_token_valid(aad_token) + def test_is_oauth_token_valid_returns_true(self): + token = { + "access_token": "my_token", + "expires_on": int(time.time()) + TOKEN_REFRESH_LEAD_TIME + 10, + "token_type": "Bearer", + } + assert self.hook._is_oauth_token_valid(token) + + def test_is_oauth_token_valid_returns_false(self): + token = { + "access_token": "my_token", + "expires_on": int(time.time()), + "token_type": "Bearer", + } + assert not self.hook._is_oauth_token_valid(token) + + def test_is_oauth_token_valid_raises_missing_token(self): + with pytest.raises(AirflowException): + self.hook._is_oauth_token_valid({}) - def test_is_aad_token_valid_returns_false(self): - aad_token = {"token": "my_token", "expires_on": int(time.time())} - assert not self.hook._is_aad_token_valid(aad_token) + def test_is_oauth_token_valid_raises_invalid_type(self): + token_missing_type = {"access_token": "my_token"} + token_wrong_type = {"access_token": "my_token", "token_type": "not bearer"} + + with pytest.raises(AirflowException): + self.hook._is_oauth_token_valid(token_missing_type) + self.hook._is_oauth_token_valid(token_wrong_type) + + def test_is_oauth_token_valid_raises_wrong_time_key(self): + token = { + "access_token": "my_token", + "expires_on": 0, + "token_type": "Bearer", + } + with pytest.raises(AirflowException): + self.hook._is_oauth_token_valid(token, time_key="expiration") @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests") def test_list_jobs_success_single_page(self, mock_requests): @@ -1448,3 +1478,84 @@ async def test_get_run_state(self, mock_get): assert ad_call_args[1]["url"] == AZURE_METADATA_SERVICE_INSTANCE_URL assert ad_call_args[1]["params"]["api-version"] > "2018-02-01" assert ad_call_args[1]["headers"]["Metadata"] == "true" + + +def create_sp_token_for_resource() -> dict: + return { + "token_type": "Bearer", + "expires_in": 3600, + "access_token": TOKEN, + } + + +class TestDatabricksHookSpToken: + """ + Tests for DatabricksHook when auth is done with Service Principal Oauth token. + """ + + @provide_session + def setup_method(self, method, session=None): + conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() + conn.login = "c64f6d12-f6e4-45a4-846e-032b42b27758" + conn.password = "secret" + conn.extra = json.dumps({"service_principal_oauth": True}) + session.commit() + self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) + + @mock.patch("airflow.providers.databricks.hooks.databricks_base.requests") + def test_submit_run(self, mock_requests): + mock_requests.codes.ok = 200 + mock_requests.post.side_effect = [ + create_successful_response_mock(create_sp_token_for_resource()), + create_successful_response_mock({"run_id": "1"}), + ] + status_code_mock = mock.PropertyMock(return_value=200) + type(mock_requests.post.return_value).status_code = status_code_mock + data = {"notebook_task": NOTEBOOK_TASK, "new_cluster": NEW_CLUSTER} + run_id = self.hook.submit_run(data) + + ad_call_args = mock_requests.method_calls[0] + assert ad_call_args[1][0] == OIDC_TOKEN_SERVICE_URL.format(HOST) + assert ad_call_args[2]["data"] == "grant_type=client_credentials&scope=all-apis" + + assert run_id == "1" + args = mock_requests.post.call_args + kwargs = args[1] + assert kwargs["auth"].token == TOKEN + + +class TestDatabricksHookAsyncSpToken: + """ + Tests for DatabricksHook using async methods when auth is done with Service + Principal Oauth token. + """ + + @provide_session + def setup_method(self, method, session=None): + conn = session.query(Connection).filter(Connection.conn_id == DEFAULT_CONN_ID).first() + conn.login = "c64f6d12-f6e4-45a4-846e-032b42b27758" + conn.password = "secret" + conn.extra = json.dumps({"service_principal_oauth": True}) + session.commit() + self.hook = DatabricksHook(retry_args=DEFAULT_RETRY_ARGS) + + @pytest.mark.asyncio + @mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get") + @mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.post") + async def test_get_run_state(self, mock_post, mock_get): + mock_post.return_value.__aenter__.return_value.json = AsyncMock( + return_value=create_sp_token_for_resource() + ) + mock_get.return_value.__aenter__.return_value.json = AsyncMock(return_value=GET_RUN_RESPONSE) + + async with self.hook: + run_state = await self.hook.a_get_run_state(RUN_ID) + + assert run_state == RunState(LIFE_CYCLE_STATE, RESULT_STATE, STATE_MESSAGE) + mock_get.assert_called_once_with( + get_run_endpoint(HOST), + json={"run_id": RUN_ID}, + auth=BearerAuth(TOKEN), + headers=self.hook.user_agent_header, + timeout=self.hook.timeout_seconds, + )