Merge pull request #56573 from mkmkme/mkmkme/reload-config

Update `query_masking_rules` when reloading the config

programs/server/Server.cpp

@@ -1372,6 +1372,8 @@ try
 
                 global_context->reloadAuxiliaryZooKeepersConfigIfChanged(config);
 
+                global_context->reloadQueryMaskingRulesIfChanged(config);
+
                 std::lock_guard lock(servers_lock);
                 updateServers(*config, server_pool, async_metrics, servers, servers_to_start_before_tables);
             }

src/Common/SensitiveDataMasker.cpp

@@ -1,5 +1,6 @@
 #include "SensitiveDataMasker.h"
 
+#include <mutex>
 #include <set>
 #include <string>
 #include <atomic>
@@ -94,20 +95,28 @@ public:
 SensitiveDataMasker::~SensitiveDataMasker() = default;
 
 std::unique_ptr<SensitiveDataMasker> SensitiveDataMasker::sensitive_data_masker = nullptr;
+std::mutex SensitiveDataMasker::instance_mutex;
 
 void SensitiveDataMasker::setInstance(std::unique_ptr<SensitiveDataMasker> sensitive_data_masker_)
 {
+
     if (!sensitive_data_masker_)
         throw Exception(ErrorCodes::LOGICAL_ERROR, "Logical error: the 'sensitive_data_masker' is not set");
 
+    std::lock_guard lock(instance_mutex);
     if (sensitive_data_masker_->rulesCount() > 0)
     {
         sensitive_data_masker = std::move(sensitive_data_masker_);
     }
+    else
+    {
+        sensitive_data_masker.reset();
+    }
 }
 
 SensitiveDataMasker * SensitiveDataMasker::getInstance()
 {
+    std::lock_guard lock(instance_mutex);
     return sensitive_data_masker.get();
 }
 

src/Common/SensitiveDataMasker.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <memory>
+#include <mutex>
 #include <vector>
 #include <cstdint>
 
@@ -45,6 +46,7 @@ class SensitiveDataMasker
 private:
     class MaskingRule;
     std::vector<std::unique_ptr<MaskingRule>> all_masking_rules;
+    static std::mutex instance_mutex;
     static std::unique_ptr<SensitiveDataMasker> sensitive_data_masker;
 
 public:

src/Interpreters/Context.cpp

@@ -4,6 +4,7 @@
 #include <memory>
 #include <Poco/UUID.h>
 #include <Poco/Util/Application.h>
+#include <Common/SensitiveDataMasker.h>
 #include <Common/Macros.h>
 #include <Common/EventNotifier.h>
 #include <Common/Stopwatch.h>
@@ -196,6 +197,8 @@ struct ContextSharedPart : boost::noncopyable
     mutable zkutil::ZooKeeperPtr zookeeper TSA_GUARDED_BY(zookeeper_mutex);                 /// Client for ZooKeeper.
     ConfigurationPtr zookeeper_config TSA_GUARDED_BY(zookeeper_mutex);                      /// Stores zookeeper configs
 
+    ConfigurationPtr sensitive_data_masker_config;
+
 #if USE_NURAFT
     mutable std::mutex keeper_dispatcher_mutex;
     mutable std::shared_ptr<KeeperDispatcher> keeper_dispatcher TSA_GUARDED_BY(keeper_dispatcher_mutex);
@@ -3198,6 +3201,16 @@ bool Context::hasAuxiliaryZooKeeper(const String & name) const
     return getConfigRef().has("auxiliary_zookeepers." + name);
 }
 
+void Context::reloadQueryMaskingRulesIfChanged(const ConfigurationPtr & config) const
+{
+    const auto old_config = shared->sensitive_data_masker_config;
+    if (old_config && isSameConfiguration(*config, *old_config, "query_masking_rules"))
+        return;
+
+    SensitiveDataMasker::setInstance(std::make_unique<SensitiveDataMasker>(*config, "query_masking_rules"));
+    shared->sensitive_data_masker_config = config;
+}
+
 InterserverCredentialsPtr Context::getInterserverCredentials() const
 {
     return shared->interserver_io_credentials.get();

src/Interpreters/Context.h

@@ -946,6 +946,8 @@ public:
     // Reload Zookeeper
     void reloadZooKeeperIfChanged(const ConfigurationPtr & config) const;
 
+    void reloadQueryMaskingRulesIfChanged(const ConfigurationPtr & config) const;
+
     void setSystemZooKeeperLogAfterInitializationIfNeeded();
 
     /// --- Caches ------------------------------------------------------------------------------------------

tests/integration/test_reload_query_masking_rules/configs/changed_settings.xml

@@ -0,0 +1,19 @@
+<clickhouse>
+    <query_log>
+        <database>system</database>
+        <table>query_log</table>
+        <partition_by>toYYYYMM(event_date)</partition_by>
+        <flush_interval_milliseconds>7500</flush_interval_milliseconds>
+        <max_size_rows>1048576</max_size_rows>
+        <reserved_size_rows>8192</reserved_size_rows>
+        <buffer_size_rows_flush_threshold>524288</buffer_size_rows_flush_threshold>
+        <flush_on_crash>false</flush_on_crash>
+    </query_log>
+
+    <query_masking_rules>
+        <rule>
+            <regexp>TOPSECRET.TOPSECRET</regexp>
+            <replace>[hidden]</replace>
+        </rule>
+    </query_masking_rules>
+</clickhouse>

tests/integration/test_reload_query_masking_rules/configs/empty_settings.xml

@@ -0,0 +1,12 @@
+<clickhouse>
+    <query_log>
+        <database>system</database>
+        <table>query_log</table>
+        <partition_by>toYYYYMM(event_date)</partition_by>
+        <flush_interval_milliseconds>7500</flush_interval_milliseconds>
+        <max_size_rows>1048576</max_size_rows>
+        <reserved_size_rows>8192</reserved_size_rows>
+        <buffer_size_rows_flush_threshold>524288</buffer_size_rows_flush_threshold>
+        <flush_on_crash>false</flush_on_crash>
+    </query_log>
+</clickhouse>

tests/integration/test_reload_query_masking_rules/test.py

@@ -0,0 +1,57 @@
+import pytest
+import os
+from helpers.cluster import ClickHouseCluster
+from helpers.test_tools import assert_eq_with_retry, assert_logs_contain_with_retry
+
+SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
+cluster = ClickHouseCluster(__file__)
+node = cluster.add_instance("node", user_configs=["configs/empty_settings.xml"])
+
+
+@pytest.fixture(scope="module", autouse=True)
+def started_cluster():
+    try:
+        cluster.start()
+        yield cluster
+    finally:
+        cluster.shutdown()
+
+
+@pytest.fixture(autouse=True)
+def reset_to_normal_settings_after_test():
+    try:
+        node.copy_file_to_container(
+            os.path.join(SCRIPT_DIR, "configs/empty_settings.xml"),
+            "/etc/clickhouse-server/config.d/z.xml",
+        )
+        node.query("SYSTEM RELOAD CONFIG")
+        yield
+    finally:
+        pass
+
+
+# @pytest.mark.parametrize("reload_strategy", ["force", "timeout"])
+def test_reload_query_masking_rules():
+    # At first, empty configuration is fed to ClickHouse. The query
+    # "SELECT 'TOPSECRET.TOPSECRET'" will not be redacted, and the new masking
+    # event will not be registered
+    node.query("SELECT 'TOPSECRET.TOPSECRET'")
+    assert_logs_contain_with_retry(node, "SELECT 'TOPSECRET.TOPSECRET'")
+    assert not node.contains_in_log(r"SELECT '\[hidden\]'")
+    node.rotate_logs()
+
+    node.copy_file_to_container(
+        os.path.join(SCRIPT_DIR, "configs/changed_settings.xml"),
+        "/etc/clickhouse-server/config.d/z.xml",
+    )
+
+    node.query("SYSTEM RELOAD CONFIG")
+
+    # Now the same query will be redacted in the logs and the counter of events
+    # will be incremented
+    node.query("SELECT 'TOPSECRET.TOPSECRET'")
+
+    assert_logs_contain_with_retry(node, r"SELECT '\[hidden\]'")
+    assert not node.contains_in_log("SELECT 'TOPSECRET.TOPSECRET'")
+
+    node.rotate_logs()