fix: support azure blob storage with federated token (#45632)

fix #44582 
related to #44583
Co-authored-by: DuMinhLe<https://github.com/ducminhle>

Signed-off-by: xiaofanluan <xiaofan.luan@zilliz.com>

internal/core/src/storage/CMakeLists.txt

@@ -16,6 +16,11 @@
 
 add_source_at_current_directory()
 
+add_subdirectory(minio)
+add_subdirectory(aliyun)
+add_subdirectory(tencent)
+add_subdirectory(huawei)
+
 if (ENABLE_GCP_NATIVE)
     add_definitions(-DENABLE_GCP_NATIVE)
     add_subdirectory(gcp-native-storage)

internal/core/src/storage/ChunkManager.cpp

@@ -28,12 +28,12 @@
 #include <aws/s3/model/ListObjectsRequest.h>
 #include <aws/s3/model/PutObjectRequest.h>
 
-#include "storage/MinioChunkManager.h"
+#include "storage/minio/MinioChunkManager.h"
-#include "storage/AliyunSTSClient.h"
+#include "storage/aliyun/AliyunSTSClient.h"
-#include "storage/TencentCloudSTSClient.h"
+#include "storage/aliyun/AliyunCredentialsProvider.h"
-#include "storage/AliyunCredentialsProvider.h"
+#include "storage/tencent/TencentCloudSTSClient.h"
-#include "storage/TencentCloudCredentialsProvider.h"
+#include "storage/tencent/TencentCloudCredentialsProvider.h"
-#include "storage/HuaweiCloudCredentialsProvider.h"
+#include "storage/huawei/HuaweiCloudCredentialsProvider.h"
 #include "common/Consts.h"
 #include "common/EasyAssert.h"
 #include "log/Log.h"

internal/core/src/storage/RemoteChunkManagerTest.cpp

@@ -13,7 +13,7 @@
 #include <string>
 #include <vector>
 
-#include "storage/MinioChunkManager.h"
+#include "storage/minio/MinioChunkManager.h"
 #include "storage/Util.h"
 
 using namespace std;

internal/core/src/storage/Util.cpp

@@ -44,7 +44,7 @@
 #include "storage/InsertData.h"
 #include "storage/LocalChunkManager.h"
 #include "storage/MemFileManagerImpl.h"
-#include "storage/MinioChunkManager.h"
+#include "storage/minio/MinioChunkManager.h"
 #ifdef USE_OPENDAL
 #include "storage/opendal/OpenDALChunkManager.h"
 #endif

internal/core/src/storage/aliyun/CMakeLists.txt

@@ -0,0 +1,19 @@
+# Copyright (C) 2019-2020 Zilliz. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+# or implied. See the License for the specific language governing permissions and limitations under the License
+
+set(ALIYUN_SRCS
+    ${CMAKE_CURRENT_SOURCE_DIR}/AliyunCredentialsProvider.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/AliyunSTSClient.cpp
+)
+
+set(SOURCE_FILES ${SOURCE_FILES} ${ALIYUN_SRCS} PARENT_SCOPE)
+
+include_directories(${CMAKE_CURRENT_SOURCE_DIR})

internal/core/src/storage/azure/AzureBlobChunkManager.cpp

@@ -14,8 +14,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <cstdlib>
+#include <cstring>
 #include <iostream>
 #include <sstream>
+#include <azure/identity/managed_identity_credential.hpp>
 #include <azure/identity/workload_identity_credential.hpp>
 #include "AzureBlobChunkManager.h"
 
@@ -66,30 +69,51 @@ AzureBlobChunkManager::AzureBlobChunkManager(
     bool useIAM) {
     requestTimeoutMs_ = requestTimeoutMs;
     if (useIAM) {
+        const char* federated_token_file =
+            std::getenv("AZURE_FEDERATED_TOKEN_FILE");
+        const bool has_federated_token = federated_token_file != nullptr &&
+                                         std::strlen(federated_token_file) > 0;
+        if (has_federated_token) {
             Azure::Identity::WorkloadIdentityCredentialOptions options;
+            const char* workload_client_id = std::getenv("AZURE_CLIENT_ID");
+            const char* workload_tenant_id = std::getenv("AZURE_TENANT_ID");
+            options.ClientId =
+                workload_client_id == nullptr ? "" : workload_client_id;
+            options.TenantId =
+                workload_tenant_id == nullptr ? "" : workload_tenant_id;
+            options.TokenFilePath = federated_token_file;
             auto workloadIdentityCredential =
                 std::make_shared<Azure::Identity::WorkloadIdentityCredential>(
                     options);
-        client_ = std::make_shared<Azure::Storage::Blobs::BlobServiceClient>(
+            client_ =
+                std::make_shared<Azure::Storage::Blobs::BlobServiceClient>(
                     "https://" + access_key_id + ".blob." + address + "/",
                     workloadIdentityCredential);
+        } else {
+            const char* client_id_env = std::getenv("AZURE_CLIENT_ID");
+            std::string client_id =
+                client_id_env == nullptr ? "" : client_id_env;
+            std::shared_ptr<Azure::Identity::ManagedIdentityCredential>
+                credential;
+            if (client_id.empty()) {
+                credential = std::make_shared<
+                    Azure::Identity::ManagedIdentityCredential>();
+            } else {
+                credential = std::make_shared<
+                    Azure::Identity::ManagedIdentityCredential>(client_id);
+            }
+
+            client_ =
+                std::make_shared<Azure::Storage::Blobs::BlobServiceClient>(
+                    "https://" + access_key_id + ".blob." + address + "/",
+                    credential);
+        }
     } else {
         client_ = std::make_shared<Azure::Storage::Blobs::BlobServiceClient>(
             Azure::Storage::Blobs::BlobServiceClient::
                 CreateFromConnectionString(GetConnectionString(
                     access_key_id, access_key_value, address)));
     }
-    try {
-        Azure::Core::Context context;
-        client_->GetBlobContainerClient("justforconnectioncheck")
-            .GetBlockBlobClient("justforconnectioncheck")
-            .GetProperties(Azure::Storage::Blobs::GetBlobPropertiesOptions(),
-                           context);
-    } catch (const Azure::Storage::StorageException& e) {
-        if (e.StatusCode != Azure::Core::Http::HttpStatusCode::NotFound) {
-            throw;
-        }
-    }
 }
 
 AzureBlobChunkManager::~AzureBlobChunkManager() {

internal/core/src/storage/huawei/CMakeLists.txt

@@ -0,0 +1,19 @@
+# Copyright (C) 2019-2020 Zilliz. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+# or implied. See the License for the specific language governing permissions and limitations under the License
+
+set(HUAWEI_SRCS
+    ${CMAKE_CURRENT_SOURCE_DIR}/HuaweiCloudCredentialsProvider.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/HuaweiCloudSTSClient.cpp
+)
+
+set(SOURCE_FILES ${SOURCE_FILES} ${HUAWEI_SRCS} PARENT_SCOPE)
+
+include_directories(${CMAKE_CURRENT_SOURCE_DIR})

internal/core/src/storage/minio/CMakeLists.txt

@@ -0,0 +1,18 @@
+# Copyright (C) 2019-2020 Zilliz. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+# or implied. See the License for the specific language governing permissions and limitations under the License
+
+set(MINIO_SRCS
+    ${CMAKE_CURRENT_SOURCE_DIR}/MinioChunkManager.cpp
+)
+
+set(SOURCE_FILES ${SOURCE_FILES} ${MINIO_SRCS} PARENT_SCOPE)
+
+include_directories(${CMAKE_CURRENT_SOURCE_DIR})

internal/core/src/storage/minio/MinioChunkManager.cpp

@@ -14,7 +14,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "storage/MinioChunkManager.h"
+#include "storage/minio/MinioChunkManager.h"
 
 #include <fstream>
 #include <aws/core/auth/AWSCredentials.h>
@@ -30,10 +30,10 @@
 #include <aws/s3/model/ListObjectsRequest.h>
 #include <aws/s3/model/PutObjectRequest.h>
 
-#include "storage/AliyunSTSClient.h"
+#include "storage/aliyun/AliyunSTSClient.h"
-#include "storage/AliyunCredentialsProvider.h"
+#include "storage/aliyun/AliyunCredentialsProvider.h"
-#include "storage/TencentCloudSTSClient.h"
+#include "storage/tencent/TencentCloudSTSClient.h"
-#include "storage/TencentCloudCredentialsProvider.h"
+#include "storage/tencent/TencentCloudCredentialsProvider.h"
 #include "monitor/Monitor.h"
 #include "common/EasyAssert.h"
 #include "log/Log.h"

internal/core/src/storage/minio/MinioChunkManagerTest.cpp

@@ -13,7 +13,7 @@
 #include <string>
 #include <vector>
 
-#include "storage/MinioChunkManager.h"
+#include "storage/minio/MinioChunkManager.h"
 #include "test_utils/indexbuilder_test_utils.h"
 
 using namespace std;

internal/core/src/storage/tencent/CMakeLists.txt

@@ -0,0 +1,19 @@
+# Copyright (C) 2019-2020 Zilliz. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+# or implied. See the License for the specific language governing permissions and limitations under the License
+
+set(TENCENT_SRCS
+    ${CMAKE_CURRENT_SOURCE_DIR}/TencentCloudCredentialsProvider.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/TencentCloudSTSClient.cpp
+)
+
+set(SOURCE_FILES ${SOURCE_FILES} ${TENCENT_SRCS} PARENT_SCOPE)
+
+include_directories(${CMAKE_CURRENT_SOURCE_DIR})

internal/proxy/task_index.go

@@ -211,7 +211,7 @@ func (cit *createIndexTask) parseIndexParams(ctx context.Context) error {
 		checker, err := indexparamcheck.GetIndexCheckerMgrInstance().GetChecker(specifyIndexType)
 		// not enable hybrid index for user, used in milvus internally
 		if err != nil || indexparamcheck.IsHYBRIDChecker(checker) {
-			log.Ctx(ctx).Warn("Failed to get index checker", zap.String(common.IndexTypeKey, specifyIndexType))
+			log.Ctx(ctx).Warn("Failed to get index checker", zap.String(common.IndexTypeKey, specifyIndexType), zap.Error(err))
 			return merr.WrapErrParameterInvalid("valid index", fmt.Sprintf("invalid index type: %s", specifyIndexType))
 		}
 	}

pkg/objectstorage/util.go

@@ -168,11 +168,26 @@ func NewAzureObjectStorageClient(ctx context.Context, c *Config) (*service.Clien
 	var client *service.Client
 	var err error
 	if c.UseIAM {
-		cred, credErr := azidentity.NewWorkloadIdentityCredential(&azidentity.WorkloadIdentityCredentialOptions{
+		var cred azcore.TokenCredential
+		var credErr error
+		if os.Getenv("AZURE_FEDERATED_TOKEN_FILE") != "" {
+			cred, credErr = azidentity.NewWorkloadIdentityCredential(&azidentity.WorkloadIdentityCredentialOptions{
 				ClientID:      os.Getenv("AZURE_CLIENT_ID"),
 				TenantID:      os.Getenv("AZURE_TENANT_ID"),
 				TokenFilePath: os.Getenv("AZURE_FEDERATED_TOKEN_FILE"),
 			})
+		} else {
+			clientID := os.Getenv("AZURE_CLIENT_ID")
+			managedIdentityID := azidentity.ClientID("") // Default to System Assigned
+
+			if clientID != "" {
+				managedIdentityID = azidentity.ClientID(clientID)
+			}
+
+			cred, credErr = azidentity.NewManagedIdentityCredential(&azidentity.ManagedIdentityCredentialOptions{
+				ID: managedIdentityID,
+			})
+		}
 		if credErr != nil {
 			return nil, credErr
 		}