backport named collection updates (#1020)

Author: yokofly

src/Common/NamedCollections/NamedCollectionConfiguration.cpp

@@ -1,8 +1,8 @@
 #include <Common/NamedCollections/NamedCollectionConfiguration.h>
 #include <Poco/Util/XMLConfiguration.h>
 #include <Common/Exception.h>
-#include <Common/SettingsChanges.h>
 #include <Common/FieldVisitorToString.h>
+#include <Common/SettingsChanges.h>
 #include <magic_enum.hpp>
 
 

src/Common/NamedCollections/NamedCollections.cpp

@@ -1,20 +1,23 @@
-#include "NamedCollections.h"
+#include <Common/NamedCollections/NamedCollections.h>
 
 #include <Interpreters/Context.h>
 #include <IO/WriteBufferFromString.h>
 #include <IO/Operators.h>
 #include <Common/NamedCollections/NamedCollectionConfiguration.h>
 #include <Poco/Util/AbstractConfiguration.h>
+#include <Common/FieldVisitorToString.h>
+
+#include <fmt/ranges.h>
 
 
 namespace DB
 {
 
 namespace ErrorCodes
 {
-    extern const int NAMED_COLLECTION_DOESNT_EXIST;
-    extern const int NAMED_COLLECTION_ALREADY_EXISTS;
     extern const int NAMED_COLLECTION_IS_IMMUTABLE;
+    extern const int BAD_ARGUMENTS;
+    extern const int NOT_IMPLEMENTED;
 }
 
 namespace Configuration = NamedCollectionConfiguration;
@@ -175,27 +178,14 @@ class NamedCollection::Impl
 NamedCollection::NamedCollection(
     ImplPtr pimpl_,
     const std::string & collection_name_,
-    SourceId source_id_,
-    bool is_mutable_)
+    const bool is_mutable_)
     : pimpl(std::move(pimpl_))
     , collection_name(collection_name_)
-    , source_id(source_id_)
     , is_mutable(is_mutable_)
 {
 }
 
-MutableNamedCollectionPtr NamedCollection::create(
-    const Poco::Util::AbstractConfiguration & config,
-    const std::string & collection_name,
-    const std::string & collection_path,
-    const Keys & keys,
-    SourceId source_id,
-    bool is_mutable)
-{
-    auto impl = Impl::create(config, collection_name, collection_path, keys);
-    return std::unique_ptr<NamedCollection>(
-        new NamedCollection(std::move(impl), collection_name, source_id, is_mutable));
-}
+NamedCollection::~NamedCollection() = default;
 
 bool NamedCollection::has(const Key & key) const
 {
@@ -296,8 +286,7 @@ MutableNamedCollectionPtr NamedCollection::duplicate() const
     std::lock_guard lock(mutex);
     auto impl = pimpl->createCopy(collection_name);
     return std::unique_ptr<NamedCollection>(
-        new NamedCollection(
-            std::move(impl), collection_name, SourceId::NONE, true));
+        new NamedCollection(std::move(impl), collection_name, true));
 }
 
 NamedCollection::Keys NamedCollection::getKeys(ssize_t depth, const std::string & prefix) const
@@ -333,6 +322,133 @@ std::unique_lock<std::mutex> NamedCollection::lock()
     return std::unique_lock(mutex);
 }
 
+
+void NamedCollection::update(const ASTAlterNamedCollectionQuery & /*query*/)
+{
+    throw Exception(
+        ErrorCodes::NOT_IMPLEMENTED,
+        "update() not implemented for NamedCollection base class.");
+}
+
+NamedCollectionFromConfig::NamedCollectionFromConfig(
+    const Poco::Util::AbstractConfiguration & config_,
+    const std::string & collection_name_,
+    const std::string & collection_path_,
+    const Keys & keys_)
+    : NamedCollection(Impl::create(config_, collection_name_, collection_path_, keys_), collection_name_, /* is_mutable */ false)
+{
+}
+
+MutableNamedCollectionPtr NamedCollectionFromConfig::create(
+    const Poco::Util::AbstractConfiguration & config_,
+    const std::string & collection_name_,
+    const std::string & collection_path_,
+    const Keys & keys_)
+{
+    return std::unique_ptr<NamedCollection>(
+        new NamedCollectionFromConfig(config_, collection_name_, collection_path_, keys_));
+}
+
+
+MutableNamedCollectionPtr NamedCollectionFromSQL::create(const ASTCreateNamedCollectionQuery & query)
+{
+    return std::unique_ptr<NamedCollection>(new NamedCollectionFromSQL(query));
+}
+
+NamedCollectionFromSQL::NamedCollectionFromSQL(const ASTCreateNamedCollectionQuery & query_)
+    : NamedCollection(nullptr, query_.collection_name, true)
+    , create_query_ptr(query_.clone()->as<ASTCreateNamedCollectionQuery &>())
+{
+    const auto config = NamedCollectionConfiguration::createConfiguration(collection_name, create_query_ptr.changes, create_query_ptr.overridability);
+
+    std::set<std::string, std::less<>> keys;
+    for (const auto & [name, _] : create_query_ptr.changes)
+        keys.insert(name);
+
+    pimpl = Impl::create(*config, collection_name, "", keys);
+}
+
+String NamedCollectionFromSQL::getCreateStatement(bool show_secrects)
+{
+    auto & changes = create_query_ptr.changes;
+    std::sort(
+        changes.begin(), changes.end(),
+        [](const SettingChange & lhs, const SettingChange & rhs) { return lhs.name < rhs.name; });
+
+    return create_query_ptr.formatWithPossiblyHidingSensitiveData(
+        /*max_length=*/0,
+        /*one_line=*/true,
+        /*show_secrets=*/show_secrects);
+}
+
+void NamedCollectionFromSQL::update(const ASTAlterNamedCollectionQuery & alter_query)
+{
+    std::lock_guard lock(mutex);
+
+    std::unordered_map<std::string, Field> result_changes_map;
+    for (const auto & [name, value] : alter_query.changes)
+    {
+        auto [it, inserted] = result_changes_map.emplace(name, value);
+        if (!inserted)
+        {
+            throw Exception(
+                ErrorCodes::BAD_ARGUMENTS,
+                "Value with key `{}` is used twice in the SET query (collection name: {})",
+                name, alter_query.collection_name);
+        }
+    }
+
+    for (const auto & [name, value] : create_query_ptr.changes)
+        result_changes_map.emplace(name, value);
+
+    std::unordered_map<std::string, bool> result_overridability_map;
+    for (const auto & [name, value] : alter_query.overridability)
+        result_overridability_map.emplace(name, value);
+    for (const auto & [name, value] : create_query_ptr.overridability)
+        result_overridability_map.emplace(name, value);
+
+    for (const auto & delete_key : alter_query.delete_keys)
+    {
+        auto it = result_changes_map.find(delete_key);
+        if (it == result_changes_map.end())
+        {
+            throw Exception(
+                ErrorCodes::BAD_ARGUMENTS,
+                "Cannot delete key `{}` because it does not exist in collection",
+                delete_key);
+        }
+
+        result_changes_map.erase(it);
+        auto it_override = result_overridability_map.find(delete_key);
+        if (it_override != result_overridability_map.end())
+            result_overridability_map.erase(it_override);
+    }
+
+    create_query_ptr.changes.clear();
+    for (const auto & [name, value] : result_changes_map)
+        create_query_ptr.changes.emplace_back(name, value);
+    create_query_ptr.overridability = std::move(result_overridability_map);
+
+    if (create_query_ptr.changes.empty())
+        throw Exception(
+            ErrorCodes::BAD_ARGUMENTS,
+            "Named collection cannot be empty (collection name: {})",
+            collection_name);
+
+    chassert(create_query_ptr.collection_name == alter_query.collection_name);
+    for (const auto & [name, value] : alter_query.changes)
+    {
+        auto it_override = alter_query.overridability.find(name);
+        if (it_override != alter_query.overridability.end())
+            setOrUpdate<String, true>(name, convertFieldToString(value), it_override->second);
+        else
+            setOrUpdate<String, true>(name, convertFieldToString(value), {});
+    }
+
+    for (const auto & key : alter_query.delete_keys)
+        remove<true>(key);
+}
+
 template String NamedCollection::get<String>(const NamedCollection::Key & key) const;
 template UInt64 NamedCollection::get<UInt64>(const NamedCollection::Key & key) const;
 template Int64 NamedCollection::get<Int64>(const NamedCollection::Key & key) const;

src/Common/NamedCollections/NamedCollections.h

@@ -1,6 +1,9 @@
 #pragma once
-#include <Interpreters/Context.h>
+
 #include <Common/NamedCollections/NamedCollections_fwd.h>
+#include <Parsers/ASTCreateNamedCollectionQuery.h>
+#include <Parsers/ASTAlterNamedCollectionQuery.h>
+
 
 namespace Poco { namespace Util { class AbstractConfiguration; } }
 
@@ -24,19 +27,13 @@ class NamedCollection
     using Keys = std::set<Key, std::less<>>;
     enum class SourceId : uint8_t
     {
+        /// None source_id is possible only if the object is a
+        /// duplicate of some named collection. See `duplicate` method.
         NONE = 0,
         CONFIG = 1,
         SQL = 2,
     };
 
-    static MutableNamedCollectionPtr create(
-        const Poco::Util::AbstractConfiguration & config,
-        const std::string & collection_name,
-        const std::string & collection_path,
-        const Keys & keys,
-        SourceId source_id_,
-        bool is_mutable_);
-
     bool has(const Key & key) const;
 
     bool hasAny(const std::initializer_list<Key> & keys) const;
@@ -61,6 +58,7 @@ class NamedCollection
 
     template <bool locked = false> void remove(const Key & key);
 
+    /// Creates mutable, with NONE source id full copy.
     MutableNamedCollectionPtr duplicate() const;
 
     Keys getKeys(ssize_t depth = -1, const std::string & prefix = "") const;
@@ -76,25 +74,72 @@ class NamedCollection
 
     bool isMutable() const { return is_mutable; }
 
-    SourceId getSourceId() const { return source_id; }
+    virtual SourceId getSourceId() const { return SourceId::NONE; }
 
-private:
+    virtual String getCreateStatement(bool /*show_secrects*/) { return  {}; }
+
+    virtual void update(const ASTAlterNamedCollectionQuery & query);
+
+    virtual ~NamedCollection();
+
+protected:
     class Impl;
     using ImplPtr = std::unique_ptr<Impl>;
-
     NamedCollection(
         ImplPtr pimpl_,
         const std::string & collection_name,
-        SourceId source_id,
-        bool is_mutable);
+        bool is_mutable_
+    );
 
     void assertMutable() const;
 
+
     ImplPtr pimpl;
     const std::string collection_name;
-    const SourceId source_id;
     const bool is_mutable;
     mutable std::mutex mutex;
 };
 
+class NamedCollectionFromSQL final : public NamedCollection
+{
+public:
+    static MutableNamedCollectionPtr create(const ASTCreateNamedCollectionQuery & query);
+
+    String getCreateStatement(bool show_secrects) override;
+
+    void update(const ASTAlterNamedCollectionQuery & query) override;
+
+    NamedCollection::SourceId getSourceId() const override { return SourceId::SQL; }
+
+private:
+    explicit NamedCollectionFromSQL(const ASTCreateNamedCollectionQuery & query_);
+
+    ASTCreateNamedCollectionQuery create_query_ptr;
+};
+
+class NamedCollectionFromConfig final : public NamedCollection
+{
+public:
+
+    static MutableNamedCollectionPtr create(
+        const Poco::Util::AbstractConfiguration & config,
+        const std::string & collection_name,
+        const std::string & collection_path,
+        const Keys & keys);
+
+    String getCreateStatement(bool /*show_secrects*/) override { return {}; }
+
+    void update(const ASTAlterNamedCollectionQuery & /*query*/) override { NamedCollection::assertMutable(); }
+
+    NamedCollection::SourceId getSourceId() const override { return SourceId::CONFIG; }
+
+private:
+
+    NamedCollectionFromConfig(
+        const Poco::Util::AbstractConfiguration & config,
+        const std::string & collection_name,
+        const std::string & collection_path,
+        const Keys & keys);
+};
+
 }