From 3ec4850eba8e55cfe5118894c4f64f5c49e856e1 Mon Sep 17 00:00:00 2001
From: robot-clickhouse <robot-clickhouse@users.noreply.github.com>
Date: Tue, 2 Jul 2024 05:09:34 +0000
Subject: [PATCH] Backport #65848 to 24.6: Add missing workload identity
 changes

---
 .../ObjectStorage/Azure/Configuration.cpp     | 38 ++++++++++++++++---
 1 file changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/Storages/ObjectStorage/Azure/Configuration.cpp b/src/Storages/ObjectStorage/Azure/Configuration.cpp
index 163f08be420..f763a997bfb 100644
--- a/src/Storages/ObjectStorage/Azure/Configuration.cpp
+++ b/src/Storages/ObjectStorage/Azure/Configuration.cpp
@@ -12,6 +12,7 @@
 #include <azure/storage/blobs.hpp>
 #include <Interpreters/evaluateConstantExpression.h>
 #include <azure/identity/managed_identity_credential.hpp>
+#include <azure/identity/workload_identity_credential.hpp>
 #include <Parsers/ASTIdentifier.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;
+        size_t pos = connection_url.find('?');
         std::shared_ptr<Azure::Identity::ManagedIdentityCredential> managed_identity_credential;
         if (storage_shared_key_credential)
         {
@@ -186,12 +188,20 @@ AzureClientPtr StorageAzureConfiguration::createClient(bool is_read_only, bool a
         }
         else
         {
-            managed_identity_credential = std::make_shared<Azure::Identity::ManagedIdentityCredential>();
-            blob_service_client = std::make_unique<BlobServiceClient>(connection_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>();
+                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;
-        size_t pos = connection_url.find('?');
         if (pos != std::string::npos)
         {
             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)
                 result = std::make_unique<BlobContainerClient>(final_url, storage_shared_key_credential);
             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
         {
@@ -236,7 +255,16 @@ AzureClientPtr StorageAzureConfiguration::createClient(bool is_read_only, bool a
                     if (storage_shared_key_credential)
                         result = std::make_unique<BlobContainerClient>(final_url, storage_shared_key_credential);
                     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
                 {