Skip to content

Reapply "update AzureBaseHook to return credentials that supports get_token method"#56228

Merged
jscheffl merged 6 commits into
apache:mainfrom
karunpoudel:azurebasehook-creds-support-get_token
Oct 3, 2025
Merged

Reapply "update AzureBaseHook to return credentials that supports get_token method"#56228
jscheffl merged 6 commits into
apache:mainfrom
karunpoudel:azurebasehook-creds-support-get_token

Conversation

@karunpoudel

@karunpoudel karunpoudel commented Sep 29, 2025

Copy link
Copy Markdown
Contributor

Reapply #52182

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 update AzureBaseHook to return credentials that supports get_token method Reapply "update AzureBaseHook to return credentials that supports get_token method" Sep 29, 2025
@ashb ashb added full tests needed We need to run full set of tests for this PR to merge all versions If set, the CI build will be forced to use all versions of Python/K8S/DBs labels Sep 30, 2025
@ashb

ashb commented Sep 30, 2025

Copy link
Copy Markdown
Member

Opening-and-closing to trigger a full test run

@ashb ashb closed this Sep 30, 2025
@ashb ashb reopened this Sep 30, 2025
@ashb

ashb commented Sep 30, 2025

Copy link
Copy Markdown
Member

This probably needs rebasing to latest main

@karunpoudel
karunpoudel force-pushed the azurebasehook-creds-support-get_token branch 2 times, most recently from e33810b to 43938a1 Compare September 30, 2025 15:14
@karunpoudel
karunpoudel force-pushed the azurebasehook-creds-support-get_token branch from 591e1d9 to f8b4d87 Compare September 30, 2025 17:15
@karunpoudel

Copy link
Copy Markdown
Contributor Author

the pr check failures doesn't seem to be related to the code changes

@jscheffl

Copy link
Copy Markdown
Contributor

the pr check failures doesn't seem to be related to the code changes

Yeah seem un-realted. Let's wait a moment until main is green again and then have another re-base. Hope then it is working.

@karunpoudel

Copy link
Copy Markdown
Contributor Author

most of the tests have passed (except for the one failing in main), so might be safe to merge this

@jscheffl jscheffl left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Double-checked and tests on main are the same failing here as well. Therefore OK to merge in my view. Hope main is green asap again.

@jscheffl
jscheffl merged commit 1b9c0ee into apache:main Oct 3, 2025
187 of 192 checks passed
@karunpoudel-chr
karunpoudel-chr deleted the azurebasehook-creds-support-get_token branch October 4, 2025 02:29
dabla pushed a commit to dabla/airflow that referenced this pull request Oct 12, 2025
…_token method" (apache#56228)

* all changes and fixed

* more fix

---------

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

Labels

all versions If set, the CI build will be forced to use all versions of Python/K8S/DBs area:providers full tests needed We need to run full set of tests for this PR to merge kind:documentation provider:microsoft-azure Azure-related issues

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants