This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 8bf53dd554 Add Service Principal OAuth for Databricks. (#33005)
8bf53dd554 is described below
commit 8bf53dd5545ecda0e5bbffbc4cc803cbbde719a9
Author: John Brandborg <[email protected]>
AuthorDate: Mon Aug 14 12:16:33 2023 +0200
Add Service Principal OAuth for Databricks. (#33005)
* Add Service Principal Oauth for Databricks. (apache/airflow#32969)
* Consolidate OAuth validation and storage (#32969)
* Formatting update and added API and Spark to Init Import
* Update documentation with doc ref and note this is for AWS deployments
---
.../providers/databricks/hooks/databricks_base.py | 148 +++++++++++++++------
.../connections/databricks.rst | 10 +-
.../providers/databricks/hooks/test_databricks.py | 123 ++++++++++++++++-
3 files changed, 234 insertions(+), 47 deletions(-)
diff --git a/airflow/providers/databricks/hooks/databricks_base.py
b/airflow/providers/databricks/hooks/databricks_base.py
index 9885e9a998..6d0d929b5d 100644
--- a/airflow/providers/databricks/hooks/databricks_base.py
+++ b/airflow/providers/databricks/hooks/databricks_base.py
@@ -62,6 +62,7 @@ AZURE_METADATA_SERVICE_INSTANCE_URL =
"http://169.254.169.254/metadata/instance";
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 @@ class BaseDatabricksHook(BaseHook):
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 @@ class BaseDatabricksHook(BaseHook):
"""
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 @@ class BaseDatabricksHook(BaseHook):
: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 @@ class BaseDatabricksHook(BaseHook):
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 @@ class BaseDatabricksHook(BaseHook):
**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 @@ class BaseDatabricksHook(BaseHook):
: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 @@ class BaseDatabricksHook(BaseHook):
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 @@ class BaseDatabricksHook(BaseHook):
**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 @@ class BaseDatabricksHook(BaseHook):
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 @@ class BaseDatabricksHook(BaseHook):
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 @@ class BaseDatabricksHook(BaseHook):
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 6303702b7e..908b12eac7 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
<https://docs.databricks.com/en/dev-tools/authentication-oauth.html>`_.
+
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 f55c55dfc3..8644d95cd9 100644
--- a/tests/providers/databricks/hooks/test_databricks.py
+++ b/tests/providers/databricks/hooks/test_databricks.py
@@ -43,6 +43,7 @@ from airflow.providers.databricks.hooks.databricks_base
import (
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 @@ class TestDatabricksHook:
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 @@ class TestDatabricksHookAsyncAadTokenManagedIdentity:
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,
+ )
