Backport #65848 to 24.6: Add missing workload identity changes

This commit is contained in:
robot-clickhouse 2024-07-02 05:09:34 +00:00
parent 6c7e71a98b
commit 3ec4850eba

View File

@ -12,6 +12,7 @@
#include <azure/storage/blobs.hpp> #include <azure/storage/blobs.hpp>
#include <Interpreters/evaluateConstantExpression.h> #include <Interpreters/evaluateConstantExpression.h>
#include <azure/identity/managed_identity_credential.hpp> #include <azure/identity/managed_identity_credential.hpp>
#include <azure/identity/workload_identity_credential.hpp>
#include <Parsers/ASTIdentifier.h> #include <Parsers/ASTIdentifier.h>
#include <Parsers/ASTFunction.h> #include <Parsers/ASTFunction.h>
@ -179,6 +180,7 @@ AzureClientPtr StorageAzureConfiguration::createClient(bool is_read_only, bool a
} }
std::unique_ptr<BlobServiceClient> blob_service_client; std::unique_ptr<BlobServiceClient> blob_service_client;
size_t pos = connection_url.find('?');
std::shared_ptr<Azure::Identity::ManagedIdentityCredential> managed_identity_credential; std::shared_ptr<Azure::Identity::ManagedIdentityCredential> managed_identity_credential;
if (storage_shared_key_credential) if (storage_shared_key_credential)
{ {
@ -186,12 +188,20 @@ AzureClientPtr StorageAzureConfiguration::createClient(bool is_read_only, bool a
} }
else else
{ {
managed_identity_credential = std::make_shared<Azure::Identity::ManagedIdentityCredential>(); /// If conneciton_url does not have '?', then its not SAS
blob_service_client = std::make_unique<BlobServiceClient>(connection_url, managed_identity_credential); if (pos == std::string::npos)
{
auto workload_identity_credential = std::make_shared<Azure::Identity::WorkloadIdentityCredential>();
blob_service_client = std::make_unique<BlobServiceClient>(connection_url, workload_identity_credential);
}
else
{
managed_identity_credential = std::make_shared<Azure::Identity::ManagedIdentityCredential>();
blob_service_client = std::make_unique<BlobServiceClient>(connection_url, managed_identity_credential);
}
} }
std::string final_url; std::string final_url;
size_t pos = connection_url.find('?');
if (pos != std::string::npos) if (pos != std::string::npos)
{ {
auto url_without_sas = connection_url.substr(0, pos); auto url_without_sas = connection_url.substr(0, pos);
@ -216,7 +226,16 @@ AzureClientPtr StorageAzureConfiguration::createClient(bool is_read_only, bool a
if (storage_shared_key_credential) if (storage_shared_key_credential)
result = std::make_unique<BlobContainerClient>(final_url, storage_shared_key_credential); result = std::make_unique<BlobContainerClient>(final_url, storage_shared_key_credential);
else else
result = std::make_unique<BlobContainerClient>(final_url, managed_identity_credential); {
/// If conneciton_url does not have '?', then its not SAS
if (pos == std::string::npos)
{
auto workload_identity_credential = std::make_shared<Azure::Identity::WorkloadIdentityCredential>();
result = std::make_unique<BlobContainerClient>(final_url, workload_identity_credential);
}
else
result = std::make_unique<BlobContainerClient>(final_url, managed_identity_credential);
}
} }
else else
{ {
@ -236,7 +255,16 @@ AzureClientPtr StorageAzureConfiguration::createClient(bool is_read_only, bool a
if (storage_shared_key_credential) if (storage_shared_key_credential)
result = std::make_unique<BlobContainerClient>(final_url, storage_shared_key_credential); result = std::make_unique<BlobContainerClient>(final_url, storage_shared_key_credential);
else else
result = std::make_unique<BlobContainerClient>(final_url, managed_identity_credential); {
/// If conneciton_url does not have '?', then its not SAS
if (pos == std::string::npos)
{
auto workload_identity_credential = std::make_shared<Azure::Identity::WorkloadIdentityCredential>();
result = std::make_unique<BlobContainerClient>(final_url, workload_identity_credential);
}
else
result = std::make_unique<BlobContainerClient>(final_url, managed_identity_credential);
}
} }
else else
{ {