Fix: possible race condition on EXCHANGE TABLES query

src/Interpreters/Context.cpp

@@ -904,6 +904,21 @@ String Context::getFilesystemCacheUser() const
     return shared->filesystem_cache_user;
 }
 
+StoragePtr Context::getOrCacheStorage(const StorageID & id, std::function<StoragePtr()> f) const
+{
+    if (auto storage = storage_cache.find(id); storage != storage_cache.end())
+    {
+        if (auto storage_ptr = storage->second.lock(); storage_ptr)
+            return storage_ptr;
+        else
+            storage_cache.erase(storage);
+    }
+    auto storage = f();
+    if (storage)
+        storage_cache.insert({id, storage});
+    return storage;
+}
+
 Strings Context::getWarnings() const
 {
     Strings common_warnings;

src/Interpreters/Context.h

@@ -33,6 +33,7 @@
 #include <memory>
 #include <mutex>
 #include <optional>
+#include <unordered_map>
 
 
 namespace Poco::Net { class IPAddress; }
@@ -421,6 +422,10 @@ class ContextData
     /// mutation tasks of one mutation executed against different parts of the same table.
     PreparedSetsCachePtr prepared_sets_cache;
 
+    /// Cache for storages resolved during query processing - we want to use the same storages we have resolved on
+    /// analysis stage for the execution - storage can be changed by concurrently running 'EXCHANGE TABLES' query.
+    mutable std::unordered_map<StorageID, StorageWeakPtr, StorageID::DatabaseAndTableNameHash, StorageID::DatabaseAndTableNameEqual> storage_cache;
+
 public:
     /// Some counters for current query execution.
     /// Most of them are workarounds and should be removed in the future.
@@ -515,6 +520,8 @@ class Context: public ContextData, public std::enable_shared_from_this<Context>
     String getFilesystemCachesPath() const;
     String getFilesystemCacheUser() const;
 
+    StoragePtr getOrCacheStorage(const StorageID & id, std::function<StoragePtr()> f) const;
+
     /// A list of warnings about server configuration to place in `system.warnings` table.
     Strings getWarnings() const;
 

src/Interpreters/DatabaseCatalog.cpp

@@ -949,15 +949,19 @@ bool DatabaseCatalog::isDictionaryExist(const StorageID & table_id) const
 StoragePtr DatabaseCatalog::getTable(const StorageID & table_id, ContextPtr local_context) const
 {
     std::optional<Exception> exc;
-    auto res = getTableImpl(table_id, local_context, &exc);
-    if (!res.second)
+    auto table = local_context->hasQueryContext() ?
+        local_context->getQueryContext()->getOrCacheStorage(table_id, [&](){ return getTableImpl(table_id, local_context, &exc).second; }) :
+        getTableImpl(table_id, local_context, &exc).second;
+    if (!table)
         throw Exception(*exc);
-    return res.second;
+    return table;
 }
 
 StoragePtr DatabaseCatalog::tryGetTable(const StorageID & table_id, ContextPtr local_context) const
 {
-    return getTableImpl(table_id, local_context, nullptr).second;
+    return local_context->hasQueryContext() ?
+        local_context->getQueryContext()->getOrCacheStorage(table_id, [&](){ return getTableImpl(table_id, local_context, nullptr).second; }) :
+        getTableImpl(table_id, local_context, nullptr).second;
 }
 
 DatabaseAndTable DatabaseCatalog::getDatabaseAndTable(const StorageID & table_id, ContextPtr local_context) const

src/Storages/IStorage_fwd.h

@@ -12,6 +12,7 @@ class IStorage;
 
 using ConstStoragePtr = std::shared_ptr<const IStorage>;
 using StoragePtr = std::shared_ptr<IStorage>;
+using StorageWeakPtr = std::weak_ptr<IStorage>;
 using Tables = std::map<String, StoragePtr>;
 
 }

tests/queries/0_stateless/03007_exchange_tables_race.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+
+set -e
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+# shellcheck source=../shell_config.sh
+. "$CURDIR"/../shell_config.sh
+
+${CLICKHOUSE_CLIENT} --multiquery <<EOF
+DROP TABLE IF EXISTS tbl_03007_1;
+DROP TABLE IF EXISTS tbl_03007_2;
+CREATE TABLE tbl_03007_1 (n Float64) ENGINE=Memory;
+CREATE TABLE tbl_03007_2 (n Int256) ENGINE=Memory;
+EOF
+
+for _ in {1..50}; do
+    (! ${CLICKHOUSE_CLIENT} --query "SELECT n * 0.123 FROM (SELECT * FROM tbl_03007_1)" 2>&1 | grep LOGICAL_ERROR) &
+    ${CLICKHOUSE_CLIENT} --query "EXCHANGE TABLES tbl_03007_1 AND tbl_03007_2" &
+done
+
+wait
+
+${CLICKHOUSE_CLIENT} --multiquery <<EOF
+DROP TABLE IF EXISTS tbl_03007_1;
+DROP TABLE IF EXISTS tbl_03007_2;
+EOF