feat(python): Add CredentialProviderAzure parameter to accept user-instantiated azure credential classes

docs/source/_build/API_REFERENCE_LINKS.yml

@@ -12,6 +12,9 @@ python:
   concat_list: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.concat_list.html
   concat_str: https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.concat_str.html
   Config: https://docs.pola.rs/api/python/stable/reference/config.html
+  CredentialProviderAWS: https://docs.pola.rs/api/python/stable/reference/api/polars.CredentialProviderAWS.html
+  CredentialProviderAzure: https://docs.pola.rs/api/python/stable/reference/api/polars.CredentialProviderAzure.html
+  CredentialProviderGCP: https://docs.pola.rs/api/python/stable/reference/api/polars.CredentialProviderGCP.html
   cs.by_name: https://docs.pola.rs/api/python/stable/reference/selectors.html#polars.selectors.by_name
   cs.contains: https://docs.pola.rs/api/python/stable/reference/selectors.html#polars.selectors.contains
   cs.first: https://docs.pola.rs/api/python/stable/reference/selectors.html#polars.selectors.first

docs/source/src/python/user-guide/io/cloud-storage.py

@@ -63,6 +63,30 @@ def get_credentials() -> pl.CredentialProviderFunctionReturn:
 df = lf.collect()
 # --8<-- [end:credential_provider_custom_func]
 
+# --8<-- [start:credential_provider_custom_func_azure]
+def credential_provider():
+    credential = DefaultAzureCredential(exclude_managed_identity_credential=True)
+    token = credential.get_token("https://storage.azure.com/.default")
+
+    return {"bearer_token": token.token}, token.expires_on
+
+
+pl.scan_parquet(
+    "abfss://...@.../...",
+    credential_provider=credential_provider,
+)
+
+# Note that for the above case, this shortcut is also available:
+
+pl.scan_parquet(
+    "abfss://...@.../...",
+    credential_provider=pl.CredentialProviderAzure(
+        credentials=DefaultAzureCredential(exclude_managed_identity_credential=True)
+    ),
+)
+
+# --8<-- [end:credential_provider_custom_func_azure]
+
 # --8<-- [start:scan_pyarrow_dataset]
 import polars as pl
 import pyarrow.dataset as ds

docs/source/src/rust/user-guide/io/cloud-storage.rs

@@ -39,6 +39,9 @@ async fn main() {
 // --8<-- [start:credential_provider_custom_func]
 // --8<-- [end:credential_provider_custom_func]
 
+// --8<-- [start:credential_provider_custom_func_azure]
+// --8<-- [end:credential_provider_custom_func_azure]
+
 // --8<-- [start:scan_pyarrow_dataset]
 // --8<-- [end:scan_pyarrow_dataset]
 

docs/source/user-guide/io/cloud-storage.md

@@ -51,7 +51,8 @@ use for authentication. This can be done in a few ways:
   functionality. For example, `pl.CredentialProviderAWS` supports selecting AWS profiles, as well as
   assuming an IAM role:
 
-{{code_block('user-guide/io/cloud-storage','credential_provider_class',['scan_parquet'])}}
+{{code_block('user-guide/io/cloud-storage','credential_provider_class',['scan_parquet',
+'CredentialProviderAWS'])}}
 
 ### Using a custom `credential_provider` function
 
@@ -60,6 +61,11 @@ use for authentication. This can be done in a few ways:
 
 {{code_block('user-guide/io/cloud-storage','credential_provider_custom_func',['scan_parquet'])}}
 
+- Example for Azure:
+
+{{code_block('user-guide/io/cloud-storage','credential_provider_custom_func_azure',['scan_parquet',
+'CredentialProviderAzure'])}}
+
 ## Scanning with PyArrow
 
 We can also scan from cloud storage using PyArrow. This is particularly useful for partitioned

py-polars/polars/io/cloud/credential_provider.py

@@ -184,34 +184,50 @@ def __init__(
         *,
         scopes: list[str] | None = None,
         tenant_id: str | None = None,
+        credentials: Any | None = None,
         _storage_account: str | None = None,
         _verbose: bool = False,
     ) -> None:
         """
         Initialize a credential provider for Microsoft Azure.
 
-        This uses `azure.identity.DefaultAzureCredential()`.
+        By default, this uses `azure.identity.DefaultAzureCredential()`.
 
         Parameters
         ----------
         scopes
             Scopes to pass to `get_token`
         tenant_id
             Azure tenant ID.
+        credentials
+            Optionally pass an instantiated Azure credential class to use (e.g.
+            `azure.identity.DefaultAzureCredential`). The credential class must
+            have a `get_token()` method.
         """
         msg = "`CredentialProviderAzure` functionality is considered unstable"
         issue_unstable_warning(msg)
 
         self.account_name = _storage_account
-        self.tenant_id = tenant_id
         self.scopes = (
             scopes if scopes is not None else ["https://storage.azure.com/.default"]
         )
+        self.tenant_id = tenant_id
+        self.credentials = credentials
         self._verbose = _verbose
 
+        if credentials is not None:
+            # If the user passes a credentials class, we just need to ensure it
+            # has a `get_token()` method.
+            if not hasattr(credentials, "get_token"):
+                msg = (
+                    f"the provided `credentials` object {credentials!r} does "
+                    "not have a `get_token()` method."
+                )
+                raise ValueError(msg)
+
         # We don't need the module if we are permitted and able to retrieve the
         # account key from the Azure CLI.
-        if self._try_get_azure_storage_account_credentials_if_permitted() is None:
+        elif self._try_get_azure_storage_account_credentials_if_permitted() is None:
             self._ensure_module_availability()
 
         if self._verbose:
@@ -233,9 +249,12 @@ def __call__(self) -> CredentialProviderFunctionReturn:
             return v
 
         # Done like this to bypass mypy, we don't have stubs for azure.identity
-        credential = importlib.import_module("azure.identity").__dict__[
-            "DefaultAzureCredential"
-        ]()
+        credential = (
+            self.credentials
+            or importlib.import_module("azure.identity").__dict__[
+                "DefaultAzureCredential"
+            ]()
+        )
         token = credential.get_token(*self.scopes, tenant_id=self.tenant_id)
 
         return {