Adapt to the blocklist feature in serverless (#9578)

close #9575

Signed-off-by: Calvin Neo <calvinneo1995@gmail.com>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>

dbms/src/Interpreters/Context.cpp

@@ -188,6 +188,8 @@ struct ContextShared
 
     JointThreadInfoJeallocMapPtr joint_memory_allocation_map; /// Joint thread-wise alloc/dealloc map
 
+    std::unordered_set<uint64_t> store_id_blocklist; /// Those store id are blocked from batch cop request.
+
     class SessionKeyHash
     {
     public:
@@ -2218,6 +2220,44 @@ void Context::setMockMPPServerInfo(MockMPPServerInfo & info)
     mpp_server_info = info;
 }
 
+const std::unordered_set<uint64_t> * Context::getStoreIdBlockList() const
+{
+    return &shared->store_id_blocklist;
+}
+
+// NOLINTNEXTLINE(readability-convert-member-functions-to-static)
+bool Context::initializeStoreIdBlockList(const String & comma_sep_string)
+{
+#if SERVERLESS_PROXY == 1
+    std::istringstream iss(comma_sep_string);
+    std::string token;
+
+    while (std::getline(iss, token, ','))
+    {
+        try
+        {
+            uint64_t number = std::stoull(token);
+            shared->store_id_blocklist.insert(number);
+        }
+        catch (...)
+        {
+            // Keep empty
+            LOG_INFO(DB::Logger::get(), "StoreIdBlockList is not set, input_str={}", comma_sep_string);
+            shared->store_id_blocklist.clear();
+            return false;
+        }
+    }
+
+    if (!shared->store_id_blocklist.empty())
+        LOG_INFO(DB::Logger::get(), "StoreIdBlockList have been set, {}", shared->store_id_blocklist);
+
+    return true;
+#else
+    UNUSED(comma_sep_string);
+    return true;
+#endif
+}
+
 SessionCleaner::~SessionCleaner()
 {
     try
@@ -2251,4 +2291,5 @@ void SessionCleaner::run()
             break;
     }
 }
+
 } // namespace DB

dbms/src/Interpreters/Context.h

@@ -190,6 +190,7 @@ class Context
     }};
     using DatabasePtr = std::shared_ptr<IDatabase>;
     using Databases = std::map<String, std::shared_ptr<IDatabase>>;
+
     /// Use copy constructor or createGlobal() instead
     Context();
 
@@ -560,6 +561,9 @@ class Context
 
     void mockConfigLoaded() { is_config_loaded = true; }
 
+    bool initializeStoreIdBlockList(const String &);
+    const std::unordered_set<uint64_t> * getStoreIdBlockList() const;
+
 private:
     /** Check if the current client has access to the specified database.
       * If access is denied, throw an exception.

dbms/src/Interpreters/SettingsCommon.h

@@ -849,6 +849,7 @@ struct SettingString
     void write(WriteBuffer & buf) const { writeBinary(value, buf); }
 
     String get() const { return value; }
+    const String & getRef() const { return value; }
 
 private:
     String value;

dbms/src/Server/Server.cpp

@@ -1512,6 +1512,9 @@ int Server::main(const std::vector<std::string> & /*args*/)
     global_context->setFormatSchemaPath(format_schema_path.path() + "/");
     format_schema_path.createDirectories();
 
+    // We do not support blocking store by id in OP mode currently.
+    global_context->initializeStoreIdBlockList("");
+
     LOG_INFO(log, "Loading metadata.");
     loadMetadataSystem(*global_context); // Load "system" database. Its engine keeps as Ordinary.
     /// After attaching system databases we can initialize system log.

dbms/src/Storages/StorageDisaggregated.cpp

@@ -116,6 +116,7 @@ std::vector<pingcap::coprocessor::BatchCopTask> StorageDisaggregated::buildBatch
         table_scan.isPartitionTableScan(),
         physical_table_ids,
         ranges_for_each_physical_table,
+        context.getStoreIdBlockList(),
         store_type,
         label_filter,
         &Poco::Logger::get("pingcap/coprocessor"));