ClickHouse · alexey-milovidov · Sep 28, 2024 · Sep 6, 2024 · Sep 18, 2024 · Sep 18, 2024
diff --git a/src/Access/AccessBackup.cpp b/src/Access/AccessBackup.cpp
diff --git a/src/Access/AccessBackup.h b/src/Access/AccessBackup.h
@@ -1,8 +1,8 @@
 #pragma once
 
+#include <Common/Logger.h>
 #include <Core/UUID.h>
 #include <unordered_map>
-#include <unordered_set>
 
 
 namespace DB
@@ -12,6 +12,7 @@ enum class AccessEntityType : uint8_t;
 struct IAccessEntity;
 using AccessEntityPtr = std::shared_ptr<const IAccessEntity>;
 class AccessRightsElements;
+class IAccessStorage;
 class IBackup;
 using BackupPtr = std::shared_ptr<const IBackup>;
 class IBackupEntry;
@@ -20,37 +21,118 @@ struct RestoreSettings;
 enum class RestoreAccessCreationMode : uint8_t;
 
 
-/// Makes a backup of access entities of a specified type.
-std::pair<String, BackupEntryPtr> makeBackupEntryForAccess(
-    const std::vector<std::pair<UUID, AccessEntityPtr>> & access_entities,
-    const String & data_path_in_backup,
-    size_t counter,
-    const AccessControl & access_control);
+/// Makes a backup entry for of a set of access entities.
+std::pair<String, BackupEntryPtr> makeBackupEntryForAccessEntities(
+    const std::vector<UUID> & entities_ids,
+    const std::unordered_map<UUID, AccessEntityPtr> & all_entities,
+    bool write_dependents,
+    const String & data_path_in_backup);
 
+struct AccessEntitiesToRestore
+{
+    /// Access entities loaded from backup with new randomly generated UUIDs.
+    std::vector<std::pair<UUID /* new_id */, AccessEntityPtr /* new_entity */>> new_entities;
+
+    /// Dependents are access entities which exist already and they should be updated after restoring.
+    /// For example, if there were a role granted to a user: `CREATE USER user1; CREATE ROLE role1; GRANT role1 TO user1`,
+    /// and we're restoring only role `role1` because user `user1` already exists,
+    /// then user `user1` should be modified after restoring role `role1` to add this grant `GRANT role1 TO user1`.
+    struct Dependent
+    {
+        /// UUID of an existing access entities.
+        UUID existing_id;
+
+        /// Source access entity from backup to copy dependencies from.
+        AccessEntityPtr source;
+    };
+    using Dependents = std::vector<Dependent>;
+    Dependents dependents;
+};
 
 /// Restores access entities from a backup.
+void restoreAccessEntitiesFromBackup(
+    IAccessStorage & access_storage,
+    const AccessEntitiesToRestore & entities_to_restore,
+    const RestoreSettings & restore_settings);
+
+
+/// Loads access entities from a backup and prepares them for insertion into an access storage.
 class AccessRestorerFromBackup
 {
 public:
     AccessRestorerFromBackup(const BackupPtr & backup_, const RestoreSettings & restore_settings_);
     ~AccessRestorerFromBackup();
 
     /// Adds a data path to loads access entities from.
-    void addDataPath(const String & data_path);
+    void addDataPath(const String & data_path_in_backup);
+
+    /// Loads access entities from the backup.
+    void loadFromBackup();
 
     /// Checks that the current user can do restoring.
+    /// Function loadFromBackup() must be called before that.
     AccessRightsElements getRequiredAccess() const;
 
-    /// Inserts all access entities loaded from all the paths added by addDataPath().
-    std::vector<std::pair<UUID, AccessEntityPtr>> getAccessEntities(const AccessControl & access_control) const;
+    /// Generates random IDs for access entities we're restoring to insert them into an access storage;
+    /// and finds IDs of existing access entities which are used as dependencies.
+    void generateRandomIDsAndResolveDependencies(const AccessControl & access_control);
+
+    /// Returns access entities loaded from backup and prepared for insertion into an access storage.
+    /// Both functions loadFromBackup() and generateRandomIDsAndResolveDependencies() must be called before that.
+    AccessEntitiesToRestore getEntitiesToRestore(const String & data_path_in_backup) const;
 
 private:
-    BackupPtr backup;
-    RestoreAccessCreationMode creation_mode;
-    bool allow_unresolved_dependencies = false;
-    std::vector<std::pair<UUID, AccessEntityPtr>> entities;
-    std::unordered_map<UUID, std::pair<String, AccessEntityType>> dependencies;
-    std::unordered_set<String> data_paths;
+    const BackupPtr backup;
+    const RestoreAccessCreationMode creation_mode;
+    const bool skip_unresolved_dependencies;
+    const bool update_dependents;
+    const LoggerPtr log;
+
+    /// Whether loadFromBackup() finished.
+    bool loaded = false;
+
+    /// Whether generateRandomIDsAndResolveDependencies() finished.
+    bool ids_assigned = false;
+
+    Strings data_paths_in_backup;
+    String data_path_with_entities_to_restore;
+
+    /// Information about an access entity loaded from the backup.
+    struct EntityInfo
+    {
+        UUID id;
+        String name;
+        AccessEntityType type;
+
+        AccessEntityPtr entity = nullptr; /// Can be nullptr if `restore=false`.
+
+        /// Index in `data_paths_in_backup`.
+        size_t data_path_index = 0;
+
+        /// Whether we're going to restore this entity.
+        /// For example,
+        /// in case of `RESTORE TABLE system.roles` this flag is true for all the roles loaded from the backup, and
+        /// in case of `RESTORE ALL` this flag is always true.
+        bool restore = false;
+
+        /// Whether this entity info was added as a dependency of another entity which we're going to restore.
+        /// For example, if we're going to restore the following user: `CREATE USER user1 DEFAULT ROLE role1, role2 SETTINGS PROFILE profile1, profile2`
+        /// then `restore=true` for `user1` and `is_dependency=true` for `role1`, `role2`, `profile1`, `profile2`.
+        /// Flags `restore` and `is_dependency` both can be set at the same time.
+        bool is_dependency = false;
+
+        /// Whether this entity info is a dependent of another entity which we're going to restore.
+        /// For example, if we're going to restore role `role1` and there is also the following user stored in the backup:
+        /// `CREATE USER user1 DEFAULT ROLE role1`, then `is_dependent=true` for `user1`.
+        /// This flags is set by generateRandomIDsAndResolveDependencies().
+        bool is_dependent = false;
+
+        /// New UUID for this entity - either randomly generated or copied from an existing entity.
+        /// This UUID is assigned by generateRandomIDsAndResolveDependencies().
+        std::optional<UUID> new_id = std::nullopt;
+    };
+
+    std::unordered_map<UUID, EntityInfo> entity_infos;
 };
 
 }
diff --git a/src/Access/AccessControl.cpp b/src/Access/AccessControl.cpp
@@ -629,9 +629,9 @@ AuthResult AccessControl::authenticate(const Credentials & credentials, const Po
     }
 }
 
-void AccessControl::restoreFromBackup(RestorerFromBackup & restorer)
+void AccessControl::restoreFromBackup(RestorerFromBackup & restorer, const String & data_path_in_backup)
 {
-    MultipleAccessStorage::restoreFromBackup(restorer);
+    MultipleAccessStorage::restoreFromBackup(restorer, data_path_in_backup);
     changes_notifier->sendNotifications();
 }
 

diff --git a/src/Access/AccessControl.h b/src/Access/AccessControl.h
@@ -124,7 +124,7 @@ class AccessControl : public MultipleAccessStorage
     AuthResult authenticate(const Credentials & credentials, const Poco::Net::IPAddress & address, const String & forwarded_address) const;
 
     /// Makes a backup of access entities.
-    void restoreFromBackup(RestorerFromBackup & restorer) override;
+    void restoreFromBackup(RestorerFromBackup & restorer, const String & data_path_in_backup) override;
 
     void setExternalAuthenticatorsConfig(const Poco::Util::AbstractConfiguration & config);
 

diff --git a/src/Access/DiskAccessStorage.cpp b/src/Access/DiskAccessStorage.cpp
@@ -676,7 +676,7 @@ bool DiskAccessStorage::updateNoLock(const UUID & id, const UpdateFunc & update_
     if (!entry.entity)
         entry.entity = readAccessEntityFromDisk(id);
     auto old_entity = entry.entity;
-    auto new_entity = update_func(old_entity);
+    auto new_entity = update_func(old_entity, id);
 
     if (!new_entity->isTypeOf(old_entity->getType()))
         throwBadCast(id, new_entity->getType(), new_entity->getName(), old_entity->getType());

diff --git a/src/Access/GrantedRoles.cpp b/src/Access/GrantedRoles.cpp
@@ -176,6 +176,16 @@ std::vector<UUID> GrantedRoles::findDependencies() const
     return res;
 }
 
+bool GrantedRoles::hasDependencies(const std::unordered_set<UUID> & ids) const
+{
+    for (const auto & role_id : roles)
+    {
+        if (ids.contains(role_id))
+            return true;
+    }
+    return false;
+}
+
 void GrantedRoles::replaceDependencies(const std::unordered_map<UUID, UUID> & old_to_new_ids)
 {
     std::vector<UUID> new_ids;
@@ -221,4 +231,56 @@ void GrantedRoles::replaceDependencies(const std::unordered_map<UUID, UUID> & ol
     }
 }
 
+void GrantedRoles::copyDependenciesFrom(const GrantedRoles & src, const std::unordered_set<UUID> & ids)
+{
+    bool found = false;
+
+    for (const auto & role_id : src.roles)
+    {
+        if (ids.contains(role_id))
+        {
+            roles.emplace(role_id);
+            found = true;
+        }
+    }
+
+    if (found)
+    {
+        for (const auto & role_id : src.roles_with_admin_option)
+        {
+            if (ids.contains(role_id))
+                roles_with_admin_option.emplace(role_id);
+        }
+    }
+}
+
+void GrantedRoles::removeDependencies(const std::unordered_set<UUID> & ids)
+{
+    bool found = false;
+
+    for (auto it = roles.begin(); it != roles.end();)
+    {
+        if (ids.contains(*it))
+        {
+            it = roles.erase(it);
+            found = true;
+        }
+        else
+        {
+            ++it;
+        }
+    }
+
+    if (found)
+    {
+        for (auto it = roles_with_admin_option.begin(); it != roles_with_admin_option.end();)
+        {
+            if (ids.contains(*it))
+                it = roles_with_admin_option.erase(it);
+            else
+                ++it;
+        }
+    }
+}
+
 }
diff --git a/src/Access/GrantedRoles.h b/src/Access/GrantedRoles.h
@@ -58,7 +58,10 @@ class GrantedRoles
     friend bool operator !=(const GrantedRoles & left, const GrantedRoles & right) { return !(left == right); }
 
     std::vector<UUID> findDependencies() const;
+    bool hasDependencies(const std::unordered_set<UUID> & ids) const;
     void replaceDependencies(const std::unordered_map<UUID, UUID> & old_to_new_ids);
+    void copyDependenciesFrom(const GrantedRoles & src, const std::unordered_set<UUID> & ids);
+    void removeDependencies(const std::unordered_set<UUID> & ids);
 
 private:
     boost::container::flat_set<UUID> roles;

diff --git a/src/Access/IAccessEntity.cpp b/src/Access/IAccessEntity.cpp
@@ -9,28 +9,4 @@ bool IAccessEntity::equal(const IAccessEntity & other) const
     return (name == other.name) && (getType() == other.getType());
 }
 
-void IAccessEntity::replaceDependencies(std::shared_ptr<const IAccessEntity> & entity, const std::unordered_map<UUID, UUID> & old_to_new_ids)
-{
-    if (old_to_new_ids.empty())
-        return;
-
-    bool need_replace_dependencies = false;
-    auto dependencies = entity->findDependencies();
-    for (const auto & dependency : dependencies)
-    {
-        if (old_to_new_ids.contains(dependency))
-        {
-            need_replace_dependencies = true;
-            break;
-        }
-    }
-
-    if (!need_replace_dependencies)
-        return;
-
-    auto new_entity = entity->clone();
-    new_entity->replaceDependencies(old_to_new_ids);
-    entity = new_entity;
-}
-
 }
diff --git a/src/Access/IAccessEntity.h b/src/Access/IAccessEntity.h
@@ -48,10 +48,13 @@ struct IAccessEntity
 
     /// Finds all dependencies.
     virtual std::vector<UUID> findDependencies() const { return {}; }
+    virtual bool hasDependencies(const std::unordered_set<UUID> & /* ids */) const { return false; }
 
     /// Replaces dependencies according to a specified map.
-    void replaceDependencies(const std::unordered_map<UUID, UUID> & old_to_new_ids) { doReplaceDependencies(old_to_new_ids); }
-    static void replaceDependencies(std::shared_ptr<const IAccessEntity> & entity, const std::unordered_map<UUID, UUID> & old_to_new_ids);
+    virtual void replaceDependencies(const std::unordered_map<UUID, UUID> & /* old_to_new_ids */) {}
+    virtual void copyDependenciesFrom(const IAccessEntity & /* src */, const std::unordered_set<UUID> & /* ids */) {}
+    virtual void removeDependencies(const std::unordered_set<UUID> & /* ids */) {}
+    virtual void clearAllExceptDependencies() {}
 
     /// Whether this access entity should be written to a backup.
     virtual bool isBackupAllowed() const { return false; }
@@ -67,8 +70,6 @@ struct IAccessEntity
     {
         return std::make_shared<EntityClassT>(typeid_cast<const EntityClassT &>(*this));
     }
-
-    virtual void doReplaceDependencies(const std::unordered_map<UUID, UUID> & /* old_to_new_ids */) {}
 };
 
 using AccessEntityPtr = std::shared_ptr<const IAccessEntity>;