add credential caching

README.md

@@ -99,6 +99,28 @@ IDs may be specified.
 
 See also [Azure Identity Managed Identity Support](https://github.com/Azure/azure-sdk-for-cpp/tree/main/sdk/identity/azure-identity#managed-identity-support)
 
+### Token Credential Caching
+
+By default, credential objects are reused across SQL statements so the Azure SDK's internal token
+cache survives query boundaries and avoids redundant authentication round-trips. To disable this,
+set `CACHE_TOKEN_CREDENTIAL = false`:
+
+```sql
+CREATE SECRET secret_cli (
+    TYPE AZURE,
+    PROVIDER CREDENTIAL_CHAIN,
+    CHAIN 'cli',
+    ACCOUNT_NAME '⟨storage account name⟩',
+    CACHE_TOKEN_CREDENTIAL false
+);
+```
+
+Updating or replacing an Azure secret invalidates its cached credential; the next query
+will create a fresh one. However, out-of-band credential changes — such as `az account logout` —
+cannot be detected and will not invalidate the cache. Setting `CACHE_TOKEN_CREDENTIAL = false`
+disables caching for that secret, so a new credential is created from scratch for each query,
+ensuring any behind-the-scenes credential changes are picked up promptly.
+
 ## Supported architectures
 
 The extension is tested & distributed for Linux (x64, arm64), MacOS (x64, arm64) and Windows (x64)
@@ -122,4 +144,4 @@ cd duckdb_azure
 GEN=ninja VCPKG_TOOLCHAIN_PATH=$PWD/../vcpkg/scripts/buildsystems/vcpkg.cmake make
 ```
 
-Please also refer to our [Build Guide](https://duckdb.org/dev/building) and [Contribution Guide]([CONTRIBUTING.md](https://github.com/duckdb/duckdb/blob/main/CONTRIBUTING.md)).
+Please also refer to our [Build Guide](https://duckdb.org/dev/building) and [Contribution Guide](<[CONTRIBUTING.md](https://github.com/duckdb/duckdb/blob/main/CONTRIBUTING.md)>).

scripts/env_azure

@@ -60,4 +60,4 @@ export AZURE_PROVIDER=cloud
 	shift
 }
 
-exec "$@"
+exec env "$@"

scripts/env_azurite

@@ -24,12 +24,12 @@ export TEMP_DIR="$AZ_TEMPDIR"
 
 # for compatibility, accept --az arg, and fail on --abfss
 [[ -n "$1" && "$1" = "--abfss" ]] && {
-	echo "run_azurite can only operate in --az mode, aborting." >&2
-	exit
+  echo "run_azurite can only operate in --az mode, aborting." >&2
+  exit
 }
 [[ -n "$1" && "$1" = "--az" ]] && {
-	# NOOP
-	shift
+  # NOOP
+  shift
 }
 
-exec "$@"
+exec env "$@"

src/azure_secret.cpp

@@ -76,6 +76,7 @@ static unique_ptr<BaseSecret> CreateAzureSecretFromCredentialChain(ClientContext
 
 	// Manage specific secret option
 	CopySecret("chain", input, *result);
+	CopySecret("cache_token_credential", input, *result);
 
 	// Redact sensible keys
 	RedactCommonKeys(*result);
@@ -103,6 +104,7 @@ static unique_ptr<BaseSecret> CreateAzureSecretFromManagedIdentity(ClientContext
 	CopySecret("client_id", input, *result);
 	CopySecret("object_id", input, *result);
 	CopySecret("resource_id", input, *result);
+	CopySecret("cache_token_credential", input, *result);
 
 	// Redact sensible keys
 	RedactCommonKeys(*result);
@@ -131,6 +133,7 @@ static unique_ptr<BaseSecret> CreateAzureSecretFromServicePrincipal(ClientContex
 	CopySecret("client_id", input, *result);
 	CopySecret("client_secret", input, *result);
 	CopySecret("client_certificate_path", input, *result);
+	CopySecret("cache_token_credential", input, *result);
 
 	// Redact sensible keys
 	RedactCommonKeys(*result);
@@ -158,6 +161,7 @@ static unique_ptr<BaseSecret> CreateAzureSecretFromAccessToken(ClientContext &co
 
 	// Manage specific secret option
 	CopySecret("access_token", input, *result);
+	CopySecret("cache_token_credential", input, *result);
 
 	// Redact sensible keys
 	RedactCommonKeys(*result);
@@ -196,6 +200,7 @@ void CreateAzureSecretFunctions::Register(ExtensionLoader &loader) {
 	// Register the credential_chain secret provider
 	CreateSecretFunction cred_chain_function = {type, "credential_chain", CreateAzureSecretFromCredentialChain};
 	cred_chain_function.named_parameters["chain"] = LogicalType::VARCHAR;
+	cred_chain_function.named_parameters["cache_token_credential"] = LogicalType::BOOLEAN;
 	RegisterCommonSecretParameters(cred_chain_function);
 	loader.RegisterFunction(cred_chain_function);
 
@@ -204,6 +209,7 @@ void CreateAzureSecretFunctions::Register(ExtensionLoader &loader) {
 	managed_identity_function.named_parameters["client_id"] = LogicalType::VARCHAR;
 	managed_identity_function.named_parameters["object_id"] = LogicalType::VARCHAR;
 	managed_identity_function.named_parameters["resource_id"] = LogicalType::VARCHAR;
+	managed_identity_function.named_parameters["cache_token_credential"] = LogicalType::BOOLEAN;
 	RegisterCommonSecretParameters(managed_identity_function);
 	loader.RegisterFunction(managed_identity_function);
 
@@ -214,12 +220,14 @@ void CreateAzureSecretFunctions::Register(ExtensionLoader &loader) {
 	service_principal_function.named_parameters["client_id"] = LogicalType::VARCHAR;
 	service_principal_function.named_parameters["client_secret"] = LogicalType::VARCHAR;
 	service_principal_function.named_parameters["client_certificate_path"] = LogicalType::VARCHAR;
+	service_principal_function.named_parameters["cache_token_credential"] = LogicalType::BOOLEAN;
 	RegisterCommonSecretParameters(service_principal_function);
 	loader.RegisterFunction(service_principal_function);
 
 	// Register the access_token secret provider
 	CreateSecretFunction access_token_function = {type, "access_token", CreateAzureSecretFromAccessToken};
 	access_token_function.named_parameters["access_token"] = LogicalType::VARCHAR;
+	access_token_function.named_parameters["cache_token_credential"] = LogicalType::BOOLEAN;
 	RegisterCommonSecretParameters(access_token_function);
 	loader.RegisterFunction(access_token_function);
 }

src/azure_storage_account_client.cpp

@@ -3,12 +3,14 @@
 #include "duckdb/catalog/catalog_transaction.hpp"
 #include "duckdb/common/exception.hpp"
 #include "duckdb/common/file_opener.hpp"
+#include "duckdb/common/helper.hpp"
 #include "duckdb/common/shared_ptr.hpp"
 #include "duckdb/common/string_util.hpp"
 #include "duckdb/main/client_context.hpp"
 #include "duckdb/main/secret/secret.hpp"
 #include "duckdb/main/secret/secret_manager.hpp"
 #include "http_state_policy.hpp"
+#include "zstd/common/xxhash.hpp"
 
 #include <azure/core/credentials/token_credential_options.hpp>
 #include <azure/core/http/curl_transport.hpp>
@@ -225,6 +227,51 @@ CreateAccessTokenCredential(const KeyValueSecret &secret) {
 	return std::make_shared<AccessTokenCredential>(access_token);
 }
 
+// Credential cache entry: persists across QueryEnd() calls since we intentionally
+// don't override QueryEnd(). The Azure SDK's internal TokenCache handles per-token
+// expiry and refresh — we just need to keep the credential object alive.
+struct AzureCredentialState : public ClientContextState {
+	AzureCredentialState(std::shared_ptr<Azure::Core::Credentials::TokenCredential> cred, string fp)
+	    : credential(std::move(cred)), fingerprint(std::move(fp)) {
+	}
+	std::shared_ptr<Azure::Core::Credentials::TokenCredential> credential;
+	std::string fingerprint;
+};
+
+static std::string FingerprintSecret(const KeyValueSecret &secret) {
+	std::string content;
+	// NOTE: this key ordering is safe because secret_map is an ordered tree
+	for (const auto &kv : secret.secret_map) {
+		content += kv.first + "=" + kv.second.ToString() + ";";
+	}
+	auto hash = duckdb_zstd::XXH64(content.c_str(), content.size(), 0);
+	return std::to_string(hash);
+}
+
+static std::shared_ptr<Azure::Core::Credentials::TokenCredential>
+GetCachedCredential(optional_ptr<FileOpener> opener, const std::string &key, const std::string &fp) {
+	auto ctx = FileOpener::TryGetClientContext(opener);
+	if (!ctx) {
+		return nullptr;
+	}
+	auto entry = ctx->registered_state->Get<AzureCredentialState>(key);
+	if (!entry || entry->fingerprint != fp) {
+		return nullptr;
+	}
+	return entry->credential;
+}
+
+static void CacheTokenCredential(optional_ptr<FileOpener> opener, const std::string &key,
+                                 std::shared_ptr<Azure::Core::Credentials::TokenCredential> cred,
+                                 const std::string &fp) {
+	auto ctx = FileOpener::TryGetClientContext(opener);
+	if (!ctx) {
+		return;
+	}
+	ctx->registered_state->Remove(key);
+	ctx->registered_state->Insert(key, make_shared_ptr<AzureCredentialState>(std::move(cred), fp));
+}
+
 static std::shared_ptr<Azure::Core::Http::HttpTransport>
 CreateCurlTransport(const std::string &proxy, const std::string &proxy_username, const std::string &proxy_password) {
 	Azure::Core::Http::CurlTransportOptions curl_transport_options;
@@ -384,32 +431,48 @@ GetDfsStorageAccountClientFromConfigProvider(optional_ptr<FileOpener> opener, co
 	return Azure::Storage::Files::DataLake::DataLakeServiceClient(account_url, dfs_options);
 }
 
+// Retrieves a cached credential by fingerprint, or creates and caches a new one via create_fn.
+// If secret.cache_token_credential==false, then do not cache. (Useful e.g. when using the az cli,
+// to where a logout is expected to have effect on next query.)
+template <class CreateFn>
+static std::shared_ptr<Azure::Core::Credentials::TokenCredential>
+GetOrCreateCredential(optional_ptr<FileOpener> opener, const KeyValueSecret &secret, CreateFn create_fn) {
+	auto fp = FingerprintSecret(secret);
+	auto cache_key = std::string("azure_cred:secret:") + secret.GetName();
+	auto ctc_val = secret.TryGetValue("cache_token_credential");
+	auto use_cache = ctc_val.IsNull() || ctc_val.GetValue<bool>();
+	auto cred = use_cache ? GetCachedCredential(opener, cache_key, fp) : nullptr;
+	if (!cred) {
+		cred = create_fn();
+		if (use_cache) {
+			CacheTokenCredential(opener, cache_key, cred, fp);
+		}
+	}
+	return cred;
+}
+
 static Azure::Storage::Blobs::BlobServiceClient
 GetBlobStorageAccountClientFromCredentialChainProvider(optional_ptr<FileOpener> opener, const KeyValueSecret &secret,
                                                        const AzureParsedUrl &azure_parsed_url) {
 	auto transport_options = GetTransportOptions(opener, secret);
-	// Create credential chain
-	auto credential = CreateChainedTokenCredential(secret, transport_options);
-
-	// Connect to storage account
+	auto cred = GetOrCreateCredential(opener, secret,
+	                                  [&]() { return CreateChainedTokenCredential(secret, transport_options); });
 	auto account_url =
 	    azure_parsed_url.is_fully_qualified ? AccountUrl(azure_parsed_url) : AccountUrl(secret, DEFAULT_BLOB_ENDPOINT);
 	auto blob_options = ToBlobClientOptions(transport_options, GetHttpState(opener));
-	return Azure::Storage::Blobs::BlobServiceClient(account_url, std::move(credential), blob_options);
+	return Azure::Storage::Blobs::BlobServiceClient(account_url, std::move(cred), blob_options);
 }
 
 static Azure::Storage::Files::DataLake::DataLakeServiceClient
 GetDfsStorageAccountClientFromCredentialChainProvider(optional_ptr<FileOpener> opener, const KeyValueSecret &secret,
                                                       const AzureParsedUrl &azure_parsed_url) {
 	auto transport_options = GetTransportOptions(opener, secret);
-	// Create credential chain
-	auto credential = CreateChainedTokenCredential(secret, transport_options);
-
-	// Connect to storage account
+	auto cred = GetOrCreateCredential(opener, secret,
+	                                  [&]() { return CreateChainedTokenCredential(secret, transport_options); });
 	auto account_url =
 	    azure_parsed_url.is_fully_qualified ? AccountUrl(azure_parsed_url) : AccountUrl(secret, DEFAULT_DFS_ENDPOINT);
 	auto dfs_options = ToDfsClientOptions(transport_options, GetHttpState(opener));
-	return Azure::Storage::Files::DataLake::DataLakeServiceClient(account_url, std::move(credential), dfs_options);
+	return Azure::Storage::Files::DataLake::DataLakeServiceClient(account_url, std::move(cred), dfs_options);
 }
 
 static std::shared_ptr<Azure::Identity::ManagedIdentityCredential>
@@ -448,78 +511,71 @@ GetManagedIdentityCredential(optional_ptr<FileOpener> opener, const KeyValueSecr
 static Azure::Storage::Blobs::BlobServiceClient
 GetBlobStorageAccountClientFromManagedIdentityProvider(optional_ptr<FileOpener> opener, const KeyValueSecret &secret,
                                                        const AzureParsedUrl &azure_parsed_url) {
-	// Connect to (blob) storage account
 	auto transport_options = GetTransportOptions(opener, secret);
-	auto mi_cred = GetManagedIdentityCredential(opener, secret, transport_options);
+	auto cred = GetOrCreateCredential(
+	    opener, secret, [&]() { return GetManagedIdentityCredential(opener, secret, transport_options); });
 	auto account_url =
 	    azure_parsed_url.is_fully_qualified ? AccountUrl(azure_parsed_url) : AccountUrl(secret, DEFAULT_BLOB_ENDPOINT);
 	auto blob_options = ToBlobClientOptions(transport_options, GetHttpState(opener));
-	return Azure::Storage::Blobs::BlobServiceClient(account_url, mi_cred, blob_options);
+	return Azure::Storage::Blobs::BlobServiceClient(account_url, cred, blob_options);
 }
 
 static Azure::Storage::Files::DataLake::DataLakeServiceClient
 GetDfsStorageAccountClientFromManagedIdentityProvider(optional_ptr<FileOpener> opener, const KeyValueSecret &secret,
                                                       const AzureParsedUrl &azure_parsed_url) {
 	auto transport_options = GetTransportOptions(opener, secret);
-	auto mi_cred = GetManagedIdentityCredential(opener, secret, transport_options);
-
-	// Connect to ADLS storage account
+	auto cred = GetOrCreateCredential(
+	    opener, secret, [&]() { return GetManagedIdentityCredential(opener, secret, transport_options); });
 	auto account_url =
 	    azure_parsed_url.is_fully_qualified ? AccountUrl(azure_parsed_url) : AccountUrl(secret, DEFAULT_DFS_ENDPOINT);
 	auto dfs_options = ToDfsClientOptions(transport_options, GetHttpState(opener));
-	return Azure::Storage::Files::DataLake::DataLakeServiceClient(account_url, mi_cred, dfs_options);
+	return Azure::Storage::Files::DataLake::DataLakeServiceClient(account_url, cred, dfs_options);
 }
 
 static Azure::Storage::Blobs::BlobServiceClient
 GetBlobStorageAccountClientFromServicePrincipalProvider(optional_ptr<FileOpener> opener, const KeyValueSecret &secret,
                                                         const AzureParsedUrl &azure_parsed_url) {
 	auto transport_options = GetTransportOptions(opener, secret);
-	auto token_credential = CreateClientCredential(secret, transport_options);
-
+	auto cred =
+	    GetOrCreateCredential(opener, secret, [&]() { return CreateClientCredential(secret, transport_options); });
 	auto account_url =
 	    azure_parsed_url.is_fully_qualified ? AccountUrl(azure_parsed_url) : AccountUrl(secret, DEFAULT_BLOB_ENDPOINT);
-	;
 	auto blob_options = ToBlobClientOptions(transport_options, GetHttpState(opener));
-	return Azure::Storage::Blobs::BlobServiceClient(account_url, token_credential, blob_options);
+	return Azure::Storage::Blobs::BlobServiceClient(account_url, cred, blob_options);
 }
 
 static Azure::Storage::Files::DataLake::DataLakeServiceClient
 GetDfsStorageAccountClientFromServicePrincipalProvider(optional_ptr<FileOpener> opener, const KeyValueSecret &secret,
                                                        const AzureParsedUrl &azure_parsed_url) {
 	auto transport_options = GetTransportOptions(opener, secret);
-	auto token_credential = CreateClientCredential(secret, transport_options);
-
+	auto cred =
+	    GetOrCreateCredential(opener, secret, [&]() { return CreateClientCredential(secret, transport_options); });
 	auto account_url =
 	    azure_parsed_url.is_fully_qualified ? AccountUrl(azure_parsed_url) : AccountUrl(secret, DEFAULT_DFS_ENDPOINT);
-	;
 	auto dfs_options = ToDfsClientOptions(transport_options, GetHttpState(opener));
-	return Azure::Storage::Files::DataLake::DataLakeServiceClient(account_url, token_credential, dfs_options);
+	return Azure::Storage::Files::DataLake::DataLakeServiceClient(account_url, cred, dfs_options);
 }
 
 static Azure::Storage::Blobs::BlobServiceClient
 GetBlobStorageAccountClientFromAccessTokenProvider(optional_ptr<FileOpener> opener, const KeyValueSecret &secret,
                                                    const AzureParsedUrl &azure_parsed_url) {
 	auto transport_options = GetTransportOptions(opener, secret);
-	auto token_credential = CreateAccessTokenCredential(secret);
-
+	auto cred = GetOrCreateCredential(opener, secret, [&]() { return CreateAccessTokenCredential(secret); });
 	auto account_url =
 	    azure_parsed_url.is_fully_qualified ? AccountUrl(azure_parsed_url) : AccountUrl(secret, DEFAULT_BLOB_ENDPOINT);
-	;
 	auto blob_options = ToBlobClientOptions(transport_options, GetHttpState(opener));
-	return Azure::Storage::Blobs::BlobServiceClient(account_url, token_credential, blob_options);
+	return Azure::Storage::Blobs::BlobServiceClient(account_url, cred, blob_options);
 }
 
 static Azure::Storage::Files::DataLake::DataLakeServiceClient
 GetDfsStorageAccountClientFromAccessTokenProvider(optional_ptr<FileOpener> opener, const KeyValueSecret &secret,
                                                   const AzureParsedUrl &azure_parsed_url) {
 	auto transport_options = GetTransportOptions(opener, secret);
-	auto token_credential = CreateAccessTokenCredential(secret);
-
+	auto cred = GetOrCreateCredential(opener, secret, [&]() { return CreateAccessTokenCredential(secret); });
 	auto account_url =
 	    azure_parsed_url.is_fully_qualified ? AccountUrl(azure_parsed_url) : AccountUrl(secret, DEFAULT_DFS_ENDPOINT);
-	;
 	auto dfs_options = ToDfsClientOptions(transport_options, GetHttpState(opener));
-	return Azure::Storage::Files::DataLake::DataLakeServiceClient(account_url, token_credential, dfs_options);
+	return Azure::Storage::Files::DataLake::DataLakeServiceClient(account_url, cred, dfs_options);
 }
 
 static Azure::Storage::Blobs::BlobServiceClient GetBlobStorageAccountClient(optional_ptr<FileOpener> opener,
@@ -609,9 +665,18 @@ static Azure::Storage::Blobs::BlobServiceClient GetBlobStorageAccountClient(opti
 	// Credential chain secret equivalent
 	auto credential_chain = TryGetCurrentSetting(opener, "azure_credential_chain");
 	if (!credential_chain.empty()) {
-		auto credential = CreateChainedTokenCredential(credential_chain, transport_options);
-
-		return Azure::Storage::Blobs::BlobServiceClient(account_url, std::move(credential), blob_options);
+		// Settings-based credentials use a singleton cache key; any change to
+		// azure_credential_chain (the only setting that determines credential type)
+		// produces a different fingerprint and evicts the cached credential.
+		auto cache_key = std::string("azure_cred:settings");
+		auto hash = duckdb_zstd::XXH64(credential_chain.c_str(), credential_chain.size(), 0);
+		auto fp = std::to_string(hash);
+		auto cred = GetCachedCredential(opener, cache_key, fp);
+		if (!cred) {
+			cred = CreateChainedTokenCredential(credential_chain, transport_options);
+			CacheTokenCredential(opener, cache_key, cred, fp);
+		}
+		return Azure::Storage::Blobs::BlobServiceClient(account_url, std::move(cred), blob_options);
 	}
 
 	// Anonymous

test/sql/cloud/spn_auth.test

@@ -17,6 +17,10 @@ require-env AZ_STORAGE_ACCOUNT
 require-env AZ_DATA_DIR
 
 
+# setup: Validate credential caching in az cli use
+statement ok
+CALL enable_logging('HTTP');
+
 statement error
 SELECT count(*) FROM 'azure://${AZ_DATA_DIR}/l.parquet';
 ----
@@ -41,3 +45,15 @@ query I
 FROM glob('az://${AZ_DATA_DIR}/*.parquet');
 ----
 az://${AZ_DATA_DIR}/l.parquet
+
+# TODO: validate that credential call only occurs 1x
+# - BLOCKED by lack of HTTP headers in Azure - see 
+# - demonstrate credential cache by call/header check for /oauth2/
+# - demonstrate invalidation by REPLACE SECRET, call/header check
+# - duplicate same to SPN auth
+#
+# Manual tests:
+# - full suite of non slow tests goes from 260 -> 26 calls to Azure's /oauth2/
+# - this test alone goes from 3 -> 1 calls
+statement ok
+FROM duckdb_logs_parsed('HTTP');