Skip to content

update AzureBaseHook to return credentials that supports get_token method#52182

Merged
potiuk merged 39 commits into
apache:mainfrom
karunpoudel:azurebasehook-creds-support-get_token
Sep 29, 2025
Merged

update AzureBaseHook to return credentials that supports get_token method#52182
potiuk merged 39 commits into
apache:mainfrom
karunpoudel:azurebasehook-creds-support-get_token

Conversation

@karunpoudel

@karunpoudel karunpoudel commented Jun 24, 2025

Copy link
Copy Markdown
Contributor

Databases like Snowflake and Azure Postgres Flex support connection using Azure OAUTH token. Newer Azure Identity credential objects like ClientSecretCredential and DefaultAzureCredential support get_token method which can be used to generate OAUTH token. SnowflakeHook and PostgresHook can leverage AzureBaseHook to retrieve required token.

In this PR:

  • Expose the credential object directly using get_credential method and make sdk_client argument optional so that the credential object can be retrieved without needing sdk_client. eg. AzureBaseHook().get_credential()
  • Added new extra use_azure_identity_object in connection for backward compatibility. When set to true, get_credential method will return newer credential object of type ClientSecretCredential or DefaultAzureCredential instead of current ServicePrincipalCredentials or AzureIdentityCredentialAdapter. Current credential object doesn't support get_token method. eg. AzureBaseHook().get_credential().get_token()

I am ok with not making sdk_client argument optional if there is any concern.

Usage:

Create base azure connection using either client_secret or default azure credential.

# Using client_secret
AIRFLOW_CONN_AZURE_SP=azure://client_id:client_secret@/?tenantId=xyz&use_azure_identity_object=true

# Using default azure credentials
AIRFLOW_CONN_AZURE_SP=azure:///?use_azure_identity_object=true

Use above connection to retrieve token to connect to specific resource (e.g., azure postgres flex server).

from airflow.models import Connection
conn = Connection.get_connection_from_secrets("azure_sp")
azure_base_hook= conn.get_hook()
scope = "https://ossrdbms-aad.database.windows.net/.default"
password = azure_base_hook.get_token(scope).token

PostgresHook could be updated to retrieve token using azure connection similar to how it currently uses aws connection.


^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

@karunpoudel karunpoudel changed the title feat: enhance AzureBaseHook to get newer credentials that support get_token method feat: enhance AzureBaseHook to get newer credential object that support get_token method Jun 24, 2025
@karunpoudel karunpoudel changed the title feat: enhance AzureBaseHook to get newer credential object that support get_token method enhance AzureBaseHook to return newer credential object that supports get_token method Jun 26, 2025
@karunpoudel karunpoudel changed the title enhance AzureBaseHook to return newer credential object that supports get_token method enhance AzureBaseHook to return credentials that supports get_token method Jul 1, 2025
@karunpoudel karunpoudel changed the title enhance AzureBaseHook to return credentials that supports get_token method update AzureBaseHook to return credentials that supports get_token method Jul 1, 2025
@karunpoudel

Copy link
Copy Markdown
Contributor Author

Hi, could someone look into this, please?

@karunpoudel-chr

Copy link
Copy Markdown
Contributor

@potiuk , can this be assigned to anyone for review?

@potiuk

potiuk commented Jul 16, 2025

Copy link
Copy Markdown
Member

Please avoid pinging people directly - unless you know they are interested - and no, I am not assigning reviewers - and also here thi is work on voluntary bases and your PR might get picked by someone. If it is not and it is ready to be reviewed, you can just ping in general in your PR or in slack, but you should not expect that pinging random individuals is a good idea.

And ideally - this generic pinging should happen after the Pr gets green and is related on top of latest main. Which i think has not happened yet for the PR?

@potiuk

potiuk commented Jul 19, 2025

Copy link
Copy Markdown
Member

errors to fix

@potiuk potiuk left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice. Looks good now.

@karunpoudel
karunpoudel marked this pull request as draft August 10, 2025 14:02
@karunpoudel
karunpoudel marked this pull request as ready for review August 10, 2025 16:09
@potiuk
potiuk merged commit d5ded85 into apache:main Sep 29, 2025
3 checks passed
dabla pushed a commit to dabla/airflow that referenced this pull request Oct 12, 2025
…thod (apache#52182)

* init

* fx

* fx

* fx

* doc update

* Update azure.rst

* fix

* fix

* fx

* fix import

* fix import

* fix import

* Update base_azure.py

* precommit fixes

* pre-commit and doc fix

* refactor

* included get_token method in hook for error handling

* add test for get_token

---------

Co-authored-by: Karun Poudel <karun.poudel@chrobinson.com>
Co-authored-by: Karun Poudel <64540927+karunpoudel-chr@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants