idraft impl of authorization aware cache

arthurpassos · arthurpassos · commit e0e19a20ea4a · 2025-05-02T19:21:56.000-03:00
diff --git a/src/Backups/BackupIO_AzureBlobStorage.cpp b/src/Backups/BackupIO_AzureBlobStorage.cpp
@@ -47,7 +47,8 @@ BackupReaderAzureBlobStorage::BackupReaderAzureBlobStorage(
         std::move(client_ptr),
         std::move(settings_ptr),
         connection_params.getContainer(),
-        connection_params.getConnectionURL());
+        connection_params.getConnectionURL(),
+        connection_params.auth_method);
 
     client = object_storage->getAzureBlobStorageClient();
     settings = object_storage->getSettings();
@@ -142,7 +143,8 @@ BackupWriterAzureBlobStorage::BackupWriterAzureBlobStorage(
         std::move(client_ptr),
         std::move(settings_ptr),
         connection_params.getContainer(),
-        connection_params.getConnectionURL());
+        connection_params.getConnectionURL(),
+        connection_params.auth_method);
 
     client = object_storage->getAzureBlobStorageClient();
     settings = object_storage->getSettings();
diff --git a/src/Disks/ObjectStorages/AzureBlobStorage/AzureObjectStorage.cpp b/src/Disks/ObjectStorages/AzureBlobStorage/AzureObjectStorage.cpp
@@ -108,12 +108,14 @@ AzureObjectStorage::AzureObjectStorage(
     ClientPtr && client_,
     SettingsPtr && settings_,
     const String & object_namespace_,
-    const String & description_)
+    const String & description_,
+    AzureBlobStorage::AuthMethod auth_method_)
     : name(name_)
     , client(std::move(client_))
     , settings(std::move(settings_))
     , object_namespace(object_namespace_)
     , description(description_)
+    , auth_method(auth_method_)
     , log(getLogger("AzureObjectStorage"))
 {
 }
@@ -154,6 +156,40 @@ ObjectStorageIteratorPtr AzureObjectStorage::iterate(const std::string & path_pr
     return std::make_shared<AzureIteratorAsync>(path_prefix, client_ptr, max_keys ? max_keys : settings_ptr->list_object_keys_size);
 }
 
+std::optional<std::string> AzureObjectStorage::getIdentityFingerprint() const
+{
+    std::optional<std::string> fingerprint;
+
+    std::visit([&fingerprint](const auto & auth) {
+        using T = std::decay_t<decltype(auth)>;
+
+        if constexpr (std::is_same_v<T, AzureBlobStorage::ConnectionString>)
+        {
+            auto connection_string_parts = Azure::Storage::_internal::ParseConnectionString(auth);
+            fingerprint = std::to_string(std::hash<std::string>()(connection_string_parts.AccountName));
+        }
+        else if constexpr (std::is_same_v<T, std::shared_ptr<Azure::Storage::StorageSharedKeyCredential>>)
+        {
+            if (auth)
+            {
+                fingerprint = std::to_string(std::hash<std::string>()(auth->AccountName));
+            }
+        }
+        /// I am not sure what to do with the other auth methods, needs further investigation
+        // else if constexpr (std::is_same_v<T, std::shared_ptr<Azure::Identity::WorkloadIdentityCredential>>) {
+        // }
+        // else if constexpr (std::is_same_v<T, std::shared_ptr<Azure::Identity::ManagedIdentityCredential>>) {
+        // }
+    }, auth_method);
+
+    if (!fingerprint)
+    {
+        return std::nullopt;
+    }
+
+    return getName() + fingerprint.value();
+}
+
 void AzureObjectStorage::listObjects(const std::string & path, RelativePathsWithMetadata & children, size_t max_keys) const
 {
     auto client_ptr = client.get();
@@ -380,7 +416,7 @@ std::unique_ptr<IObjectStorage> AzureObjectStorage::cloneObjectStorage(
     };
 
     auto new_client = AzureBlobStorage::getContainerClient(params, /*readonly=*/ true);
-    return std::make_unique<AzureObjectStorage>(name, std::move(new_client), std::move(new_settings), new_namespace, params.endpoint.getServiceEndpoint());
+    return std::make_unique<AzureObjectStorage>(name, std::move(new_client), std::move(new_settings), new_namespace, params.endpoint.getServiceEndpoint(), params.auth_method);
 }
 
 }
diff --git a/src/Disks/ObjectStorages/AzureBlobStorage/AzureObjectStorage.h b/src/Disks/ObjectStorages/AzureBlobStorage/AzureObjectStorage.h
@@ -29,14 +29,17 @@ class AzureObjectStorage : public IObjectStorage
         ClientPtr && client_,
         SettingsPtr && settings_,
         const String & object_namespace_,
-        const String & description_);
+        const String & description_,
+        AzureBlobStorage::AuthMethod auth_method_);
 
     void listObjects(const std::string & path, RelativePathsWithMetadata & children, size_t max_keys) const override;
 
     ObjectStorageIteratorPtr iterate(const std::string & path_prefix, size_t max_keys) const override;
 
     std::string getName() const override { return "AzureObjectStorage"; }
 
+    std::optional<std::string> getIdentityFingerprint() const override;
+
     ObjectStorageType getType() const override { return ObjectStorageType::Azure; }
 
     std::string getCommonKeyPrefix() const override { return ""; }
@@ -116,6 +119,8 @@ class AzureObjectStorage : public IObjectStorage
     /// We use source url without container and prefix as description, because in Azure there are no limitations for operations between different containers.
     const String description;
 
+    AzureBlobStorage::AuthMethod auth_method;
+
     LoggerPtr log;
 };
 
diff --git a/src/Disks/ObjectStorages/Cached/CachedObjectStorage.h b/src/Disks/ObjectStorages/Cached/CachedObjectStorage.h
@@ -31,6 +31,8 @@ class CachedObjectStorage final : public IObjectStorage
 
     bool exists(const StoredObject & object) const override;
 
+    std::optional<std::string> getIdentityFingerprint() const override { return object_storage->getIdentityFingerprint(); }
+
     std::unique_ptr<ReadBufferFromFileBase> readObject( /// NOLINT
         const StoredObject & object,
         const ReadSettings & read_settings,
diff --git a/src/Disks/ObjectStorages/HDFS/HDFSObjectStorage.h b/src/Disks/ObjectStorages/HDFS/HDFSObjectStorage.h
@@ -65,6 +65,8 @@ class HDFSObjectStorage : public IObjectStorage, public HDFSErrorWrapper
 
     bool exists(const StoredObject & object) const override;
 
+    std::optional<std::string> getIdentityFingerprint() const override { return std::nullopt; }
+
     std::unique_ptr<ReadBufferFromFileBase> readObject( /// NOLINT
         const StoredObject & object,
         const ReadSettings & read_settings,
diff --git a/src/Disks/ObjectStorages/IObjectStorage.h b/src/Disks/ObjectStorages/IObjectStorage.h
@@ -126,6 +126,9 @@ class IObjectStorage
 
     virtual std::string getDescription() const = 0;
 
+    // todo arthur add docs
+    virtual std::optional<std::string> getIdentityFingerprint() const = 0;
+
     virtual const MetadataStorageMetrics & getMetadataStorageMetrics() const;
 
     /// Object exists or not
diff --git a/src/Disks/ObjectStorages/Local/LocalObjectStorage.h b/src/Disks/ObjectStorages/Local/LocalObjectStorage.h
@@ -28,6 +28,9 @@ class LocalObjectStorage : public IObjectStorage
 
     bool exists(const StoredObject & object) const override;
 
+    // no auth
+    std::optional<std::string> getIdentityFingerprint() const override { return getName(); }
+
     std::unique_ptr<ReadBufferFromFileBase> readObject( /// NOLINT
         const StoredObject & object,
         const ReadSettings & read_settings,
diff --git a/src/Disks/ObjectStorages/ObjectStorageFactory.cpp b/src/Disks/ObjectStorages/ObjectStorageFactory.cpp
@@ -307,7 +307,8 @@ void registerAzureObjectStorage(ObjectStorageFactory & factory)
             ObjectStorageType::Azure, config, config_prefix, name,
             AzureBlobStorage::getContainerClient(params, /*readonly=*/ false), std::move(azure_settings),
             params.endpoint.prefix.empty() ? params.endpoint.container_name : params.endpoint.container_name + "/" + params.endpoint.prefix,
-            params.endpoint.getServiceEndpoint());
+            params.endpoint.getServiceEndpoint(),
+            params.auth_method);
     };
 
     factory.registerObjectStorageType("azure_blob_storage", creator);
diff --git a/src/Disks/ObjectStorages/S3/S3ObjectStorage.cpp b/src/Disks/ObjectStorages/S3/S3ObjectStorage.cpp
@@ -164,6 +164,13 @@ bool S3ObjectStorage::exists(const StoredObject & object) const
     return S3::objectExists(*client.get(), uri.bucket, object.remote_path, {});
 }
 
+std::optional<std::string> S3ObjectStorage::getIdentityFingerprint() const
+{
+    const auto credentials = client.get()->getCredentials();
+
+    return getName() + credentials.GetAWSAccessKeyId();
+}
+
 std::unique_ptr<ReadBufferFromFileBase> S3ObjectStorage::readObject( /// NOLINT
     const StoredObject & object,
     const ReadSettings & read_settings,
diff --git a/src/Disks/ObjectStorages/S3/S3ObjectStorage.h b/src/Disks/ObjectStorages/S3/S3ObjectStorage.h
@@ -81,6 +81,8 @@ class S3ObjectStorage : public IObjectStorage
 
     ObjectStorageType getType() const override { return ObjectStorageType::S3; }
 
+    std::optional<std::string> getIdentityFingerprint() const override;
+
     bool exists(const StoredObject & object) const override;
 
     std::unique_ptr<ReadBufferFromFileBase> readObject( /// NOLINT
diff --git a/src/Disks/ObjectStorages/Web/WebObjectStorage.h b/src/Disks/ObjectStorages/Web/WebObjectStorage.h
@@ -33,6 +33,9 @@ class WebObjectStorage : public IObjectStorage, WithContext
 
     bool exists(const StoredObject & object) const override;
 
+    // apparently no auth is required for this object storage type
+    std::optional<std::string> getIdentityFingerprint() const override { return getName(); }
+
     std::unique_ptr<ReadBufferFromFileBase> readObject( /// NOLINT
         const StoredObject & object,
         const ReadSettings & read_settings,
diff --git a/src/Storages/Cache/ObjectStorageListObjectsCache.cpp b/src/Storages/Cache/ObjectStorageListObjectsCache.cpp
@@ -73,21 +73,23 @@ class ObjectStorageListObjectsCachePolicy : public TTLCachePolicy<Key, Mapped, H
 };
 
 ObjectStorageListObjectsCache::Key::Key(
+    const String & storage_identity_,
     const String & bucket_,
     const String & prefix_,
     const std::chrono::steady_clock::time_point & expires_at_,
     std::optional<UUID> user_id_)
-    : bucket(bucket_), prefix(prefix_), expires_at(expires_at_), user_id(user_id_) {}
+    : storage_identity(storage_identity_), bucket(bucket_), prefix(prefix_), expires_at(expires_at_), user_id(user_id_) {}
 
 bool ObjectStorageListObjectsCache::Key::operator==(const Key & other) const
 {
-    return bucket == other.bucket && prefix == other.prefix;
+    return storage_identity == other.storage_identity && bucket == other.bucket && prefix == other.prefix;
 }
 
 size_t ObjectStorageListObjectsCache::KeyHasher::operator()(const Key & key) const
 {
     std::size_t seed = 0;
 
+    boost::hash_combine(seed, key.storage_identity);
     boost::hash_combine(seed, key.bucket);
     boost::hash_combine(seed, key.prefix);
 
@@ -117,12 +119,10 @@ ObjectStorageListObjectsCache::ObjectStorageListObjectsCache()
 }
 
 void ObjectStorageListObjectsCache::set(
-    const std::string & bucket,
-    const std::string & prefix,
+    Key key,
     const std::shared_ptr<Value> & value)
 {
-    const auto key = Key{bucket, prefix, std::chrono::steady_clock::now() + std::chrono::seconds(ttl_in_seconds)};
-
+    key.expires_at = std::chrono::steady_clock::now() + std::chrono::seconds(ttl_in_seconds);
     cache.set(key, value);
 }
 
@@ -131,9 +131,8 @@ void ObjectStorageListObjectsCache::clear()
     cache.clear();
 }
 
-std::optional<ObjectStorageListObjectsCache::Value> ObjectStorageListObjectsCache::get(const String & bucket, const String & prefix, bool filter_by_prefix)
+std::optional<ObjectStorageListObjectsCache::Value> ObjectStorageListObjectsCache::get(const Key & input_key, bool filter_by_prefix)
 {
-    const auto input_key = Key{bucket, prefix};
     auto pair = cache.getWithKey(input_key);
 
     if (!pair)
diff --git a/src/Storages/Cache/ObjectStorageListObjectsCache.h b/src/Storages/Cache/ObjectStorageListObjectsCache.h
@@ -16,11 +16,14 @@ class ObjectStorageListObjectsCache
     struct Key
     {
         Key(
+            const String & storage_identity_,
             const String & bucket_,
             const String & prefix_,
             const std::chrono::steady_clock::time_point & expires_at_ = std::chrono::steady_clock::now(),
             std::optional<UUID> user_id_ = std::nullopt);
 
+        // object storage name + account name
+        std::string storage_identity;
         std::string bucket;
         std::string prefix;
         std::chrono::steady_clock::time_point expires_at;
@@ -47,12 +50,9 @@ class ObjectStorageListObjectsCache
 
     using Cache = CacheBase<Key, Value, KeyHasher, WeightFunction>;
 
-    void set(
-        const std::string & bucket,
-        const std::string & prefix,
-        const std::shared_ptr<Value> & value);
+    void set(Key key,const std::shared_ptr<Value> & value);
 
-    std::optional<Value> get(const String & bucket, const String & prefix, bool filter_by_prefix = true);
+    std::optional<Value> get(const Key & input_key, bool filter_by_prefix = true);
 
     void clear();
 
diff --git a/src/Storages/Cache/tests/gtest_object_storage_list_objects_cache.cpp b/src/Storages/Cache/tests/gtest_object_storage_list_objects_cache.cpp
diff --git a/src/Storages/ObjectStorage/Azure/Configuration.cpp b/src/Storages/ObjectStorage/Azure/Configuration.cpp
diff --git a/src/Storages/ObjectStorage/StorageObjectStorageSource.cpp b/src/Storages/ObjectStorage/StorageObjectStorageSource.cpp
diff --git a/src/Storages/ObjectStorage/StorageObjectStorageSource.h b/src/Storages/ObjectStorage/StorageObjectStorageSource.h
