Skip to content

Conversation

@karunpoudel
Copy link
Contributor

@karunpoudel karunpoudel commented Jun 24, 2025

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
Contributor Author

Hi, could someone look into this, please?

@karunpoudel-chr
Copy link
Contributor

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

@potiuk
Copy link
Member

potiuk commented Jul 16, 2025

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?

@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
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Sep 30, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 1, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 2, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 3, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 4, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 5, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 5, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 7, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 8, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 9, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 10, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 11, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/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>
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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 14, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 15, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 17, 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>
abdulrahman305 bot pushed a commit to abdulrahman305/airflow that referenced this pull request Oct 19, 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