Fix: possible race condition on EXCHANGE TABLES query

src/Interpreters/Context.cpp

@@ -912,6 +912,23 @@ String Context::getFilesystemCacheUser() const
     return shared->filesystem_cache_user;
 }
 
+DatabaseAndTable Context::getOrCacheStorage(const StorageID & id, std::function<DatabaseAndTable()> f, std::optional<Exception> * exception) const
+{
+    if (auto storage_id = storage_cache.find(id); storage_id != storage_cache.end())
+    {
+        DatabaseAndTable storage = DatabaseCatalog::instance().tryGetByUUID(storage_id->uuid);
+        if (exception && !storage.second)
+            exception->emplace(Exception(ErrorCodes::UNKNOWN_TABLE, "Table {} does not exist anymore - maybe it was dropped", id.getNameForLogs()));
+        return storage;
+    }
+
+    auto storage = f();
+    if (storage.second)
+        if (const auto & new_id = storage.second->getStorageID(); new_id.hasUUID())
+            storage_cache.insert(new_id);
+    return storage;
+}
+
 Strings Context::getWarnings() const
 {
     Strings common_warnings;

src/Interpreters/Context.h

@@ -21,6 +21,7 @@
 #include <Interpreters/Context_fwd.h>
 #include <Interpreters/StorageID.h>
 #include <Interpreters/MergeTreeTransactionHolder.h>
+#include <Interpreters/DatabaseCatalog.h>
 #include <Parsers/IAST_fwd.h>
 #include <Server/HTTP/HTTPContext.h>
 #include <Storages/ColumnsDescription.h>
@@ -430,6 +431,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_set<StorageID, 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.
@@ -524,6 +529,8 @@ class Context: public ContextData, public std::enable_shared_from_this<Context>
     String getFilesystemCachesPath() const;
     String getFilesystemCacheUser() const;
 
+    DatabaseAndTable getOrCacheStorage(const StorageID & id, std::function<DatabaseAndTable()> f, std::optional<Exception> * exception) const;
+
     /// A list of warnings about server configuration to place in `system.warnings` table.
     Strings getWarnings() const;
 

src/Interpreters/DatabaseCatalog.cpp

@@ -949,29 +949,37 @@ 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); }, &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); }, nullptr).second :
+        getTableImpl(table_id, local_context, nullptr).second;
 }
 
 DatabaseAndTable DatabaseCatalog::getDatabaseAndTable(const StorageID & table_id, ContextPtr local_context) const
 {
     std::optional<Exception> exc;
-    auto res = getTableImpl(table_id, local_context, &exc);
+    auto res = local_context->hasQueryContext() ?
+        local_context->getQueryContext()->getOrCacheStorage(table_id, [&](){ return getTableImpl(table_id, local_context, &exc); }, &exc) :
+        getTableImpl(table_id, local_context, &exc);
     if (!res.second)
         throw Exception(*exc);
     return res;
 }
 
 DatabaseAndTable DatabaseCatalog::tryGetDatabaseAndTable(const StorageID & table_id, ContextPtr local_context) const
 {
-    return getTableImpl(table_id, local_context, nullptr);
+    return local_context->hasQueryContext() ?
+        local_context->getQueryContext()->getOrCacheStorage(table_id, [&](){ return getTableImpl(table_id, local_context, nullptr); }, nullptr) :
+        getTableImpl(table_id, local_context, nullptr);
 }
 
 void DatabaseCatalog::loadMarkedAsDroppedTables()

tests/queries/0_stateless/03007_exchange_tables_race.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Tags: no-ordinary-database
+
+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