You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "basvandriel (via GitHub)" <gi...@apache.org> on 2024/08/21 11:12:24 UTC
[PR] Enable workload identity authentication for the Databricks provider [airflow]
basvandriel opened a new pull request, #41639:
URL: https://github.com/apache/airflow/pull/41639
This pull request adds support for [authenticating using Workload Identity ](https://learn.microsoft.com/nl-nl/azure/aks/workload-identity-overview?tabs=dotnet)for the Databricks provider. In the provided unit tests, this has been tested using mocking the token generation. However, the `DefaultAzureCredential().get_token(...)` part has been tested on an actual AKS cluster. This closes #41586.
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #41586
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
<!-- Please keep an empty line above the dashes. -->
---
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "BasPH (via GitHub)" <gi...@apache.org>.
BasPH commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2426197411
@basvandriel @jrderuiter There's something fishy going on with the async test, see error in https://github.com/apache/airflow/actions/runs/11039010730/job/30937844742?pr=41639.
This fails:
```python
@pytest.mark.asyncio
@mock.patch.dict(
os.environ,
{
"AZURE_CLIENT_ID": "fake-client-id",
"AZURE_TENANT_ID": "fake-tenant-id",
"AZURE_FEDERATED_TOKEN_FILE": "/badpath",
"KUBERNETES_SERVICE_HOST": "fakeip",
},
)
@mock.patch(
"azure.identity.aio.DefaultAzureCredential.get_token", return_value=create_aad_token_for_resource()
)
@mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get")
async def test_one(self, requests_mock, get_token_mock: mock.MagicMock):
requests_mock.return_value.__aenter__.return_value.json.side_effect = mock.AsyncMock(
side_effect=[{"data": 1}]
)
async with self._hook:
result = await self._hook.a_get_run_output(0)
assert result == {"data": 1}
```
I found that moving the mocking of env vars inside the test works, but I'm not sure why yet:
```python
@pytest.mark.asyncio
@mock.patch(
"azure.identity.aio.DefaultAzureCredential.get_token", return_value=create_aad_token_for_resource()
)
@mock.patch("airflow.providers.databricks.hooks.databricks_base.aiohttp.ClientSession.get")
async def test_one(self, requests_mock, get_token_mock: mock.MagicMock):
with mock.patch.dict(
os.environ,
{
"AZURE_CLIENT_ID": "fake-client-id",
"AZURE_TENANT_ID": "fake-tenant-id",
"AZURE_FEDERATED_TOKEN_FILE": "/badpath",
"KUBERNETES_SERVICE_HOST": "fakeip",
},
):
requests_mock.return_value.__aenter__.return_value.json.side_effect = mock.AsyncMock(
side_effect=[{"data": 1}]
)
async with self._hook:
result = await self._hook.a_get_run_output(0)
assert result == {"data": 1}
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "BasPH (via GitHub)" <gi...@apache.org>.
BasPH commented on code in PR #41639:
URL: https://github.com/apache/airflow/pull/41639#discussion_r1752369891
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -374,6 +378,95 @@ async def _a_get_aad_token(self, resource: str) -> str:
return jsn["access_token"]
+ def _get_aad_token_for_default_az_credential(self, resource: str) -> str:
+ """
+ Get AAD token for given resource for workload identity.
+
+ Supports managed identity or service principal auth.
+ :param resource: resource to issue token to
+ :return: AAD token, or raise an exception
+ """
+ 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:
+ from azure.identity import DefaultAzureCredential
+
+ for attempt in self._get_retry_object():
+ with attempt:
+ # This only works in an AKS Cluster given the following environment variables:
+ # AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_FEDERATED_TOKEN_FILE
+ #
+ # While there is a WorkloadIdentityCredential class, the below class is advised by microsoft examples
+ # https://learn.microsoft.com/nl-nl/azure/aks/workload-identity-overview
Review Comment:
Would link to the English docs:
```suggestion
# While there is a WorkloadIdentityCredential class, the below class is advised by Microsoft
# https://learn.microsoft.com/en-us/azure/aks/workload-identity-overview
```
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -374,6 +378,95 @@ async def _a_get_aad_token(self, resource: str) -> str:
return jsn["access_token"]
+ def _get_aad_token_for_default_az_credential(self, resource: str) -> str:
+ """
+ Get AAD token for given resource for workload identity.
+
+ Supports managed identity or service principal auth.
+ :param resource: resource to issue token to
+ :return: AAD token, or raise an exception
+ """
+ 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:
+ from azure.identity import DefaultAzureCredential
+
+ for attempt in self._get_retry_object():
+ with attempt:
+ # This only works in an AKS Cluster given the following environment variables:
+ # AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_FEDERATED_TOKEN_FILE
+ #
+ # While there is a WorkloadIdentityCredential class, the below class is advised by microsoft examples
+ # https://learn.microsoft.com/nl-nl/azure/aks/workload-identity-overview
+ token = DefaultAzureCredential().get_token(f"{resource}/.default")
+
+ jsn = {
+ "access_token": token.token,
+ "token_type": "Bearer",
+ "expires_on": token.expires_on,
+ }
+ self._is_oauth_token_valid(jsn)
+ self.oauth_tokens[resource] = jsn
+ break
+ except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(e)
+ except RetryError:
+ raise AirflowException(f"API requests to Azure failed {self.retry_limit} times. Giving up.")
+ except requests_exceptions.HTTPError as e:
+ msg = f"Response: {e.response.content.decode()}, Status Code: {e.response.status_code}"
+ raise AirflowException(msg)
+
+ return token.token
+
+ async def _a_get_aad_token_for_default_az_credential(self, resource: str) -> str:
+ """
+ Get AAD token for given resource for workload identity.
+
+ Supports managed identity or service principal auth.
+ :param resource: resource to issue token to
+ :return: AAD token, or raise an exception
+ """
+ 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:
+ # from azure.identity import DefaultAzureCredential
Review Comment:
Remove unnecessary comment
```suggestion
```
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -374,6 +378,95 @@ async def _a_get_aad_token(self, resource: str) -> str:
return jsn["access_token"]
+ def _get_aad_token_for_default_az_credential(self, resource: str) -> str:
+ """
+ Get AAD token for given resource for workload identity.
+
+ Supports managed identity or service principal auth.
+ :param resource: resource to issue token to
+ :return: AAD token, or raise an exception
+ """
+ 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:
+ from azure.identity import DefaultAzureCredential
+
+ for attempt in self._get_retry_object():
+ with attempt:
+ # This only works in an AKS Cluster given the following environment variables:
+ # AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_FEDERATED_TOKEN_FILE
+ #
+ # While there is a WorkloadIdentityCredential class, the below class is advised by microsoft examples
+ # https://learn.microsoft.com/nl-nl/azure/aks/workload-identity-overview
+ token = DefaultAzureCredential().get_token(f"{resource}/.default")
+
+ jsn = {
+ "access_token": token.token,
+ "token_type": "Bearer",
+ "expires_on": token.expires_on,
+ }
+ self._is_oauth_token_valid(jsn)
+ self.oauth_tokens[resource] = jsn
+ break
+ except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(e)
+ except RetryError:
+ raise AirflowException(f"API requests to Azure failed {self.retry_limit} times. Giving up.")
+ except requests_exceptions.HTTPError as e:
+ msg = f"Response: {e.response.content.decode()}, Status Code: {e.response.status_code}"
+ raise AirflowException(msg)
+
+ return token.token
+
+ async def _a_get_aad_token_for_default_az_credential(self, resource: str) -> str:
+ """
+ Get AAD token for given resource for workload identity.
+
+ Supports managed identity or service principal auth.
+ :param resource: resource to issue token to
+ :return: AAD token, or raise an exception
+ """
+ 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:
+ # from azure.identity import DefaultAzureCredential
+ from azure.identity.aio import (
+ DefaultAzureCredential as AsyncDefaultAzureCredential,
+ )
+
+ for attempt in self._get_retry_object():
+ with attempt:
+ # This only works in an AKS Cluster given the following environment variables:
+ # AZURE_TENANT_ID, AZURE_CLIENT_ID, AZURE_FEDERATED_TOKEN_FILE
+ #
+ # While there is a WorkloadIdentityCredential class, the below class is advised by microsoft examples
+ # https://learn.microsoft.com/nl-nl/azure/aks/workload-identity-overview
Review Comment:
```suggestion
# While there is a WorkloadIdentityCredential class, the below class is advised by Microsoft
# https://learn.microsoft.com/en-us/azure/aks/workload-identity-overview
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "jrderuiter (via GitHub)" <gi...@apache.org>.
jrderuiter commented on code in PR #41639:
URL: https://github.com/apache/airflow/pull/41639#discussion_r1745681899
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -502,6 +597,10 @@ async def _a_get_token(self, raise_error: bool = False) -> str | None:
self.log.debug("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(DEFAULT_AZURE_CREDENTIAL_SETTING_KEY, False):
+ self.log.debug("Using AzureDefaultCredentialy for authentication.")
Review Comment:
Small typo in the log: `AzureDefaultCredentialy` -> `AzureDefaultCredential`.
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -63,6 +64,8 @@
DEFAULT_DATABRICKS_SCOPE = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
OIDC_TOKEN_SERVICE_URL = "{}/oidc/v1/token"
+DEFAULT_AZURE_CREDENTIAL_SETTING_KEY = "use_default_azure_credential"
Review Comment:
This key should also be added to the `extra_parameters` list in line 89 I think.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "jrderuiter (via GitHub)" <gi...@apache.org>.
jrderuiter commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2330824383
@basvandriel Can you give me write access to your PR fork + branch? Then I can take over this PR as discussed privately.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "BasPH (via GitHub)" <gi...@apache.org>.
BasPH commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2368125404
Hi @basvandriel almost looks okay to me. There's one tiny static check nit that thinks that "aks" is a typo: https://github.com/apache/airflow/actions/runs/10882398151/job/30464268761?pr=41639#step:8:80. Could you add "aks" to `docs/spelling_wordlist.txt`?
The other CICD failure (https://github.com/apache/airflow/actions/runs/10882398151/job/30464269767?pr=41639) is one I cannot explain at first sight. Will need to deepdive to understand what's failing there.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "jrderuiter (via GitHub)" <gi...@apache.org>.
jrderuiter commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2348641618
@basvandriel Do you need any help fixing the static checks?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "BasPH (via GitHub)" <gi...@apache.org>.
BasPH commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2326020217
Looks okay to me, although there's quite a lot of code duplication between the `_get_aad_token_for_default_az_credential` and `_a _get_aad_token_for_default_az_credential` methods. Any way to extract common code?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "BasPH (via GitHub)" <gi...@apache.org>.
BasPH merged PR #41639:
URL: https://github.com/apache/airflow/pull/41639
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "BasPH (via GitHub)" <gi...@apache.org>.
BasPH commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2330843997
@jrderuiter In that case keeping the duplicated code is okay with me. Could you fix the static code checks?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "basvandriel (via GitHub)" <gi...@apache.org>.
basvandriel commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2330841781
Hi @jrderuiter, I was already working on this. I will finalize this today. Keep you posted!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "basvandriel (via GitHub)" <gi...@apache.org>.
basvandriel commented on code in PR #41639:
URL: https://github.com/apache/airflow/pull/41639#discussion_r1738472786
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -374,6 +377,95 @@ async def _a_get_aad_token(self, resource: str) -> str:
return jsn["access_token"]
+ def _get_aad_token_for_workload_identity(self, resource: str) -> str:
Review Comment:
Great suggestion. The change has been made!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "jrderuiter (via GitHub)" <gi...@apache.org>.
jrderuiter commented on code in PR #41639:
URL: https://github.com/apache/airflow/pull/41639#discussion_r1738426848
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -374,6 +377,95 @@ async def _a_get_aad_token(self, resource: str) -> str:
return jsn["access_token"]
+ def _get_aad_token_for_workload_identity(self, resource: str) -> str:
Review Comment:
In principle getting a token in this manner isn't restricted to workload identity perse. For example, using the `DefaultAzureCredential` you can support any of the authentication options listed here: https://learn.microsoft.com/en-us/dotnet/azure/sdk/authentication/credential-chains?tabs=dac#defaultazurecredential-overview.
Would it make sense to rename this method to reflect this and maybe use an options such as `use_azure_default_credential` instead?
For local development this option would also be awesome, as it allows you to use your personal Entra ID together with the Azure CLI to authenticate to a Databricks cluster.
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -374,6 +377,95 @@ async def _a_get_aad_token(self, resource: str) -> str:
return jsn["access_token"]
+ def _get_aad_token_for_workload_identity(self, resource: str) -> str:
Review Comment:
In principle getting a token in this manner isn't restricted to workload identity perse. For example, using the `DefaultAzureCredential` you can support any of the authentication options listed here: https://learn.microsoft.com/en-us/dotnet/azure/sdk/authentication/credential-chains?tabs=dac#defaultazurecredential-overview.
Would it make sense to rename this method to reflect this and maybe use an options such as `use_azure_default_credential` instead?
For local development this more general option would also be awesome, as it allows you to use your personal Entra ID together with the Azure CLI to authenticate to a Databricks cluster.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2301795046
Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide (https://github.com/apache/airflow/blob/main/contributing-docs/README.rst)
Here are some useful points:
- Pay attention to the quality of your code (ruff, mypy and type annotations). Our [pre-commits]( https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst#prerequisites-for-pre-commit-hooks) will help you with that.
- In case of a new feature add useful documentation (in docstrings or in `docs/` directory). Adding a new operator? Check this short [guide](https://github.com/apache/airflow/blob/main/docs/apache-airflow/howto/custom-operator.rst) Consider adding an example DAG that shows how users should use it.
- Consider using [Breeze environment](https://github.com/apache/airflow/blob/main/dev/breeze/doc/README.rst) for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
- Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
- Please follow [ASF Code of Conduct](https://www.apache.org/foundation/policies/conduct) for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
- Be sure to read the [Airflow Coding style]( https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#coding-style-and-best-practices).
- Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: dev@airflow.apache.org
Slack: https://s.apache.org/airflow-slack
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "basvandriel (via GitHub)" <gi...@apache.org>.
basvandriel commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2331071602
> I'm not really sure about the refactoring TBH, now it feels like we add an extra layer of complexity (the `executor` functions + `get_aad_token` functions) with minimal reduction in duplication.
You are right. If we want a similar approach, we would need 2 functions for the wrapping code and the rest for the variations of retrieving the token. I'm going to revert my commit and apply the other suggestions. Thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "jrderuiter (via GitHub)" <gi...@apache.org>.
jrderuiter commented on code in PR #41639:
URL: https://github.com/apache/airflow/pull/41639#discussion_r1745127277
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -502,6 +598,13 @@ async def _a_get_token(self, raise_error: bool = False) -> str | None:
self.log.debug("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(DEFAULT_AZURE_CREDENTIAL_SETTING_KEY, False):
+ if not self._running_in_kubernetes():
Review Comment:
I'd remove this check, as using workload identity from Kubernetes is only one of the options that the AzureDefaultCredential can use for authentication. We should just use the AzureDefaultCredential and let it figure out which mode it will use for authentication.
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -502,6 +598,13 @@ async def _a_get_token(self, raise_error: bool = False) -> str | None:
self.log.debug("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(DEFAULT_AZURE_CREDENTIAL_SETTING_KEY, False):
+ if not self._running_in_kubernetes():
+ raise AirflowException(
+ "Workload identity authentication is only supporting when running in an Kubernetes cluster"
+ )
+ self.log.debug("Using Azure Workload Identity authentication.")
Review Comment:
I'd change this message to indicate that you're using the AzureDefaultCredential instead of referencing workload identity. Workload identity is only one of the (many) mechanisms that the AzureDefaultCredential can use for authentication and we have no guarantee that it will actually use workload identity.
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -282,6 +291,78 @@ async def _a_get_sp_token(self, resource: str) -> str:
raise AirflowException(msg)
return jsn["access_token"]
+
+ def _call_aad_token_executor(self, resource: str, get_token_callable: Callable[[str], AccessToken]):
+ """
+ Get AAD token for given resource.
+
+ Supports managed identity or service principal auth.
+ :param resource: resource to issue token to
+ :return: AAD token, or raise an exception
+ """
+ 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:
+ for attempt in self._get_retry_object():
+ with attempt:
+ token = get_token_callable(f"{resource}/.default")
+
+ jsn = {
+ "access_token": token.token,
+ "token_type": "Bearer",
+ "expires_on": token.expires_on,
+ }
+ self._is_oauth_token_valid(jsn)
+ self.oauth_tokens[resource] = jsn
+ break
+ except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(e)
+ except RetryError:
+ raise AirflowException(f"API requests to Azure failed {self.retry_limit} times. Giving up.")
+ except requests_exceptions.HTTPError as e:
+ msg = f"Response: {e.response.content.decode()}, Status Code: {e.response.status_code}"
+ raise AirflowException(msg)
+
+ return jsn["access_token"]
+
+ async def _a_call_aad_token_executor(self, resource: str, get_token_callable: Callable[[str], Awaitable[AccessToken]]):
Review Comment:
I still see a lot of code duplication here. Aren't we just making it more complicated by adding this executor function if it only saves us 1-2 lines of duplication?
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -282,6 +291,78 @@ async def _a_get_sp_token(self, resource: str) -> str:
raise AirflowException(msg)
return jsn["access_token"]
+
+ def _call_aad_token_executor(self, resource: str, get_token_callable: Callable[[str], AccessToken]):
Review Comment:
I'd use a different term than `executor` here, as `executor` has a specific meaning in the Airflow context: https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/executor/index.html.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "jrderuiter (via GitHub)" <gi...@apache.org>.
jrderuiter commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2337828740
I see some of the static checks are failing, can you have a look at this @basvandriel? Afterwards we should be fine to merge if @BasPH agrees.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "basvandriel (via GitHub)" <gi...@apache.org>.
basvandriel commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2352302103
> @basvandriel Do you need any help fixing the static checks?
Hi @jrderuiter, I'm sorry for the late reply. I will be looking into this today.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "BasPH (via GitHub)" <gi...@apache.org>.
BasPH commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2460366602
All green now, it took a while but thanks for hanging on @basvandriel @jrderuiter!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "basvandriel (via GitHub)" <gi...@apache.org>.
basvandriel commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2331017946
While I'm running my pre-commit checks I made some update on the duplicated code. Unfortunately, using `asyncio.run` is not an option. The code that is currently is being used uses a different SDK for the asynchronous part of the token generation.
However, I think the readability increases by having a `Callable` argument, where a function returns an `AccessToken`. This way we only have 2 functions instead of 4, one for the asynchronous one and one the synchronous.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "jrderuiter (via GitHub)" <gi...@apache.org>.
jrderuiter commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2330817101
This duplication is essentially following the existing pattern that's already in place in the hook (if you look at the other methods, there is very often both a sync and an async implementation).
To reduce duplication we could break out the setup (lines 430-434) and the checking of the token (lines 450-456) into separate functions but I'm not sure if that would really improve readability.
A more heavy-weight alternative would be to only define the async version and then define the sync version using [asyncio.run](https://docs.python.org/3/library/asyncio-runner.html#asyncio.run):
```
def _get_aad_token_for_default_az_credential(self, resource: str) -> str:
return asyncio.run(self. _a_get_aad_token_for_default_az_credential(resource)
```
However, this would add some overhead to the sync function call as `asyncio.run` starts an async event loop in the background:
> This function runs the passed coroutine, taking care of managing the asyncio event loop, finalizing asynchronous generators, and closing the executor.
In short, to match the existing style I would prefer leaving it as is. However, I'm happy to explore these other options if needed. Also open to any suggestions from others :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "basvandriel (via GitHub)" <gi...@apache.org>.
basvandriel commented on code in PR #41639:
URL: https://github.com/apache/airflow/pull/41639#discussion_r1745138230
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -282,6 +291,78 @@ async def _a_get_sp_token(self, resource: str) -> str:
raise AirflowException(msg)
return jsn["access_token"]
+
+ def _call_aad_token_executor(self, resource: str, get_token_callable: Callable[[str], AccessToken]):
+ """
+ Get AAD token for given resource.
+
+ Supports managed identity or service principal auth.
+ :param resource: resource to issue token to
+ :return: AAD token, or raise an exception
+ """
+ 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:
+ for attempt in self._get_retry_object():
+ with attempt:
+ token = get_token_callable(f"{resource}/.default")
+
+ jsn = {
+ "access_token": token.token,
+ "token_type": "Bearer",
+ "expires_on": token.expires_on,
+ }
+ self._is_oauth_token_valid(jsn)
+ self.oauth_tokens[resource] = jsn
+ break
+ except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(e)
+ except RetryError:
+ raise AirflowException(f"API requests to Azure failed {self.retry_limit} times. Giving up.")
+ except requests_exceptions.HTTPError as e:
+ msg = f"Response: {e.response.content.decode()}, Status Code: {e.response.status_code}"
+ raise AirflowException(msg)
+
+ return jsn["access_token"]
+
+ async def _a_call_aad_token_executor(self, resource: str, get_token_callable: Callable[[str], Awaitable[AccessToken]]):
Review Comment:
You are right. If we want a similar approach, we would need 2 functions for the wrapping code and the rest for the variations of retrieving the token. I'm going to revert my commit and apply the other suggestions. Thanks!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #41639:
URL: https://github.com/apache/airflow/pull/41639#issuecomment-2460372572
Awesome work, congrats on your first merged pull request! You are invited to check our [Issue Tracker](https://github.com/apache/airflow/issues) for additional contributions.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "jrderuiter (via GitHub)" <gi...@apache.org>.
jrderuiter commented on code in PR #41639:
URL: https://github.com/apache/airflow/pull/41639#discussion_r1738426848
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -374,6 +377,95 @@ async def _a_get_aad_token(self, resource: str) -> str:
return jsn["access_token"]
+ def _get_aad_token_for_workload_identity(self, resource: str) -> str:
Review Comment:
In principle getting a token in this manner isn't restricted to workload identity perse. For example, using the `DefaultAzureCredential` you can support any of the authentication options listed here: https://learn.microsoft.com/en-us/dotnet/azure/sdk/authentication/credential-chains?tabs=dac#defaultazurecredential-overview.
Would it make sense to rename this method to reflect this and maybe use an options such as `use_azure_default_credential` instead?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Enable workload identity authentication for the Databricks provider [airflow]
Posted by "jrderuiter (via GitHub)" <gi...@apache.org>.
jrderuiter commented on code in PR #41639:
URL: https://github.com/apache/airflow/pull/41639#discussion_r1738426848
##########
airflow/providers/databricks/hooks/databricks_base.py:
##########
@@ -374,6 +377,95 @@ async def _a_get_aad_token(self, resource: str) -> str:
return jsn["access_token"]
+ def _get_aad_token_for_workload_identity(self, resource: str) -> str:
Review Comment:
In principle getting a token in this manner isn't restricted to workload identity perse. For example, using the `DefaultAzureCredential` you can support any of the authentication options listed here: https://learn.microsoft.com/en-us/dotnet/azure/sdk/authentication/credential-chains?tabs=dac#defaultazurecredential-overview.
Would it make sense to rename this method to reflect this and maybe use an options such as `use_azure_default_credential` instead?
For local development this more general option would also be awesome, as it (for example) allows you to use your personal Entra ID together with the Azure CLI to authenticate to a Databricks cluster.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org