Expose azure_get_token() scalar to mint AAD bearer tokens

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",

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<FileOpener> 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<string_t>(result)[0] = StringVector::AddString(result, token);
+}
+
+static void AzureGetTokenFunction1Arg(DataChunk &args, ExpressionState &state, Vector &result) {
+	UnaryExecutor::Execute<string_t, string_t>(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<string_t, string_t, string_t>(
+	    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

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<FileOpener> opener,
                                                                      const std::string &path,
@@ -21,4 +22,6 @@ ConnectToDfsStorageAccount(optional_ptr<FileOpener> opener, const std::string &p
                            const AzureParsedUrl &azure_parsed_url);
 
 const SecretMatch LookupSecret(optional_ptr<FileOpener> opener, const std::string &path);
+
+void RegisterAzureGetTokenFunction(ExtensionLoader &loader);
 } // namespace duckdb