From a2d9c5bc10b86c90ec9a5ec2b33cdacf909b07f8 Mon Sep 17 00:00:00 2001 From: djouallah Date: Tue, 19 May 2026 10:21:15 +1000 Subject: [PATCH 1/2] Add azure_get_token() scalar function Mints a short-lived Azure AAD bearer token via the existing CreateChainedTokenCredential helper. Intended use: feed the token into a duckdb-iceberg `iceberg`-typed secret to attach an OneLake REST catalog, without coupling either extension's secret types. Three arities: azure_get_token() azure_get_token(chain) azure_get_token(chain, scope) Defaults: chain='default', scope='https://storage.azure.com/.default'. Co-Authored-By: Claude Opus 4.7 --- src/azure_extension.cpp | 4 ++ src/azure_storage_account_client.cpp | 53 ++++++++++++++++++++ src/include/azure_storage_account_client.hpp | 3 ++ 3 files changed, 60 insertions(+) diff --git a/src/azure_extension.cpp b/src/azure_extension.cpp index 28c7a4a..daa535f 100644 --- a/src/azure_extension.cpp +++ b/src/azure_extension.cpp @@ -2,6 +2,7 @@ #include "azure_blob_filesystem.hpp" #include "azure_dfs_filesystem.hpp" #include "azure_secret.hpp" +#include "azure_storage_account_client.hpp" namespace duckdb { @@ -15,6 +16,9 @@ static void LoadInternal(ExtensionLoader &loader) { // Load Secret functions CreateAzureSecretFunctions::Register(loader); + // Register scalar functions + RegisterAzureGetTokenFunction(loader); + // Load extension config auto &config = DBConfig::GetConfig(instance); config.AddExtensionOption("azure_storage_connection_string", diff --git a/src/azure_storage_account_client.cpp b/src/azure_storage_account_client.cpp index 00dfed6..0d40348 100644 --- a/src/azure_storage_account_client.cpp +++ b/src/azure_storage_account_client.cpp @@ -5,9 +5,13 @@ #include "duckdb/common/file_opener.hpp" #include "duckdb/common/shared_ptr.hpp" #include "duckdb/common/string_util.hpp" +#include "duckdb/common/vector_operations/binary_executor.hpp" +#include "duckdb/common/vector_operations/unary_executor.hpp" +#include "duckdb/function/scalar_function.hpp" #include "duckdb/logging/logger.hpp" #include "duckdb/main/client_context.hpp" #include "duckdb/main/database.hpp" +#include "duckdb/main/extension/extension_loader.hpp" #include "duckdb/main/secret/secret.hpp" #include "duckdb/main/secret/secret_manager.hpp" #include "http_logging_policy.hpp" @@ -741,4 +745,53 @@ ConnectToDfsStorageAccount(optional_ptr opener, const std::string &p return Azure::Storage::Files::DataLake::DataLakeServiceClient(account_url, dfs_options); } +static const std::string AZURE_DEFAULT_TOKEN_SCOPE = "https://storage.azure.com/.default"; + +static std::string FetchAzureBearerToken(const std::string &chain, const std::string &scope) { + Azure::Core::Http::Policies::TransportOptions transport_options; + auto credential = CreateChainedTokenCredential(chain, transport_options); + + Azure::Core::Credentials::TokenRequestContext request_ctx; + request_ctx.Scopes = {scope}; + try { + auto token = credential->GetToken(request_ctx, Azure::Core::Context()); + return token.Token; + } catch (const std::exception &ex) { + throw InvalidConfigurationException( + "azure_get_token failed for chain='%s': %s. " + "If chain includes 'cli', ensure `az login` has been run.", + chain, ex.what()); + } +} + +static void AzureGetTokenFunction0Args(DataChunk &args, ExpressionState &state, Vector &result) { + auto token = FetchAzureBearerToken("default", AZURE_DEFAULT_TOKEN_SCOPE); + result.SetVectorType(VectorType::CONSTANT_VECTOR); + ConstantVector::GetData(result)[0] = StringVector::AddString(result, token); +} + +static void AzureGetTokenFunction1Arg(DataChunk &args, ExpressionState &state, Vector &result) { + UnaryExecutor::Execute(args.data[0], result, args.size(), [&](string_t chain_str) { + auto token = FetchAzureBearerToken(chain_str.GetString(), AZURE_DEFAULT_TOKEN_SCOPE); + return StringVector::AddString(result, token); + }); +} + +static void AzureGetTokenFunction2Args(DataChunk &args, ExpressionState &state, Vector &result) { + BinaryExecutor::Execute( + args.data[0], args.data[1], result, args.size(), [&](string_t chain_str, string_t scope_str) { + auto token = FetchAzureBearerToken(chain_str.GetString(), scope_str.GetString()); + return StringVector::AddString(result, token); + }); +} + +void RegisterAzureGetTokenFunction(ExtensionLoader &loader) { + ScalarFunctionSet set("azure_get_token"); + set.AddFunction(ScalarFunction({}, LogicalType::VARCHAR, AzureGetTokenFunction0Args)); + set.AddFunction(ScalarFunction({LogicalType::VARCHAR}, LogicalType::VARCHAR, AzureGetTokenFunction1Arg)); + set.AddFunction( + ScalarFunction({LogicalType::VARCHAR, LogicalType::VARCHAR}, LogicalType::VARCHAR, AzureGetTokenFunction2Args)); + loader.RegisterFunction(set); +} + } // namespace duckdb diff --git a/src/include/azure_storage_account_client.hpp b/src/include/azure_storage_account_client.hpp index 274f4e8..1b57c27 100644 --- a/src/include/azure_storage_account_client.hpp +++ b/src/include/azure_storage_account_client.hpp @@ -11,6 +11,7 @@ #include "azure_parsed_url.hpp" namespace duckdb { +class ExtensionLoader; Azure::Storage::Blobs::BlobServiceClient ConnectToBlobStorageAccount(optional_ptr opener, const std::string &path, @@ -21,4 +22,6 @@ ConnectToDfsStorageAccount(optional_ptr opener, const std::string &p const AzureParsedUrl &azure_parsed_url); const SecretMatch LookupSecret(optional_ptr opener, const std::string &path); + +void RegisterAzureGetTokenFunction(ExtensionLoader &loader); } // namespace duckdb From 5e7d1c68909cdf2ada76b1a05ac1711b1aeb4b49 Mon Sep 17 00:00:00 2001 From: djouallah Date: Tue, 19 May 2026 10:27:27 +1000 Subject: [PATCH 2/2] Fix clang-format: inline throw open-paren Co-Authored-By: Claude Opus 4.7 --- src/azure_storage_account_client.cpp | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/azure_storage_account_client.cpp b/src/azure_storage_account_client.cpp index 0d40348..38de2da 100644 --- a/src/azure_storage_account_client.cpp +++ b/src/azure_storage_account_client.cpp @@ -757,10 +757,9 @@ static std::string FetchAzureBearerToken(const std::string &chain, const std::st auto token = credential->GetToken(request_ctx, Azure::Core::Context()); return token.Token; } catch (const std::exception &ex) { - throw InvalidConfigurationException( - "azure_get_token failed for chain='%s': %s. " - "If chain includes 'cli', ensure `az login` has been run.", - chain, ex.what()); + throw InvalidConfigurationException("azure_get_token failed for chain='%s': %s. " + "If chain includes 'cli', ensure `az login` has been run.", + chain, ex.what()); } }