diff --git a/README.md b/README.md index 0dc05d6..ac828be 100644 --- a/README.md +++ b/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)>). diff --git a/scripts/env_azure b/scripts/env_azure index bcb8721..19b6885 100755 --- a/scripts/env_azure +++ b/scripts/env_azure @@ -60,4 +60,4 @@ export AZURE_PROVIDER=cloud shift } -exec "$@" +exec env "$@" diff --git a/scripts/env_azurite b/scripts/env_azurite index 1f91cff..0355d60 100755 --- a/scripts/env_azurite +++ b/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 "$@" diff --git a/src/azure_secret.cpp b/src/azure_secret.cpp index 2e39f21..96f79eb 100644 --- a/src/azure_secret.cpp +++ b/src/azure_secret.cpp @@ -76,6 +76,7 @@ static unique_ptr 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 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 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 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); } diff --git a/src/azure_storage_account_client.cpp b/src/azure_storage_account_client.cpp index e50ebf5..9abbb86 100644 --- a/src/azure_storage_account_client.cpp +++ b/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 #include @@ -225,6 +227,51 @@ CreateAccessTokenCredential(const KeyValueSecret &secret) { return std::make_shared(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 cred, string fp) + : credential(std::move(cred)), fingerprint(std::move(fp)) { + } + std::shared_ptr 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 +GetCachedCredential(optional_ptr opener, const std::string &key, const std::string &fp) { + auto ctx = FileOpener::TryGetClientContext(opener); + if (!ctx) { + return nullptr; + } + auto entry = ctx->registered_state->Get(key); + if (!entry || entry->fingerprint != fp) { + return nullptr; + } + return entry->credential; +} + +static void CacheTokenCredential(optional_ptr opener, const std::string &key, + std::shared_ptr 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(std::move(cred), fp)); +} + static std::shared_ptr 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 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 +static std::shared_ptr +GetOrCreateCredential(optional_ptr 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(); + 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 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 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 @@ -448,78 +511,71 @@ GetManagedIdentityCredential(optional_ptr opener, const KeyValueSecr static Azure::Storage::Blobs::BlobServiceClient GetBlobStorageAccountClientFromManagedIdentityProvider(optional_ptr 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 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 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 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 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 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 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 diff --git a/test/sql/cloud/spn_auth.test b/test/sql/cloud/spn_auth.test index 5812353..e65fb6f 100644 --- a/test/sql/cloud/spn_auth.test +++ b/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');