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..38de2da 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,52 @@ 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