supported drop query for temporary tables [#CLICKHOUSE-3219]

dbms/src/Interpreters/Context.cpp

@@ -781,6 +781,31 @@ void Context::addExternalTable(const String & table_name, StoragePtr storage)
     }
 }
 
+StoragePtr Context::tryRemoveExternalTable(const String & database_name, const String & table_name)
+{
+    auto lock = getLock();
+
+    /// Ability to remove the temporary tables of another query in the form _query_QUERY_ID.table
+
+    if (startsWith(database_name, "_query_"))
+    {
+        String requested_query_id = database_name.substr(strlen("_query_"));
+
+        return shared->process_list.tryRemoveTemporaryTable(requested_query_id, table_name);
+    }
+    else if(database_name.empty())
+    {
+        Tables::const_iterator it = external_tables.find(table_name);
+        if (external_tables.end() == it)
+            return StoragePtr();
+
+        auto storage = it->second;
+        external_tables.erase(it);
+        return storage;
+    }
+
+    return {};
+}
 
 DDLGuard::DDLGuard(Map & map_, std::mutex & mutex_, std::unique_lock<std::mutex> && lock, const String & elem, const String & message)
     : map(map_), mutex(mutex_)

dbms/src/Interpreters/Context.h

@@ -170,6 +170,7 @@ public:
     StoragePtr getTable(const String & database_name, const String & table_name) const;
     StoragePtr tryGetTable(const String & database_name, const String & table_name) const;
     void addExternalTable(const String & table_name, StoragePtr storage);
+    StoragePtr tryRemoveExternalTable(const String & database_name, const String & table_name);
 
     void addDatabase(const String & database_name, const DatabasePtr & database);
     DatabasePtr detachDatabase(const String & database_name);
@@ -241,9 +242,11 @@ public:
 
     const Context & getSessionContext() const;
     Context & getSessionContext();
+    bool hasSessionContext() const { return session_context != nullptr; }
 
     const Context & getGlobalContext() const;
     Context & getGlobalContext();
+    bool hasGlobalContext() const { return global_context != nullptr; }
 
     void setSessionContext(Context & context_)                                { session_context = &context_; }
     void setGlobalContext(Context & context_)                                { global_context = &context_; }

dbms/src/Interpreters/InterpreterDropQuery.cpp

@@ -45,6 +45,20 @@ BlockIO InterpreterDropQuery::execute()
         return {};
     }
 
+    /// Drop temporary table.
+    StoragePtr table = (context.hasSessionContext() ? context.getSessionContext() : context)
+            .tryRemoveExternalTable(drop.database, drop.table);
+    if (table)
+    {
+        table->shutdown();
+        /// If table was already dropped by anyone, an exception will be thrown
+        auto table_lock = table->lockForAlter();
+        /// Delete table data
+        table->drop();
+        table->is_dropped = true;
+        return {};
+    }
+
     String database_name = drop.database.empty() ? current_database : drop.database;
     String database_name_escaped = escapeForFileName(database_name);
 

dbms/src/Interpreters/ProcessList.cpp

@@ -209,8 +209,17 @@ StoragePtr ProcessList::tryGetTemporaryTable(const String & query_id, const Stri
 {
     std::lock_guard<std::mutex> lock(mutex);
 
+    Tables * tables;
+    Tables::iterator iterator;
+    std::tie(tables, iterator) = tryFindTemporaryTable(query_id, table_name);
+
+    return tables ? iterator->second : nullptr;
+}
+
+std::tuple<Tables *, Tables::iterator> ProcessList::tryFindTemporaryTable(const String & query_id, const String & table_name) const
+{
     /// NOTE We search for all user-s. That is, there is no isolation, and the complexity is O(users).
-    for (const auto & user_queries : user_to_queries)
+    for (auto & user_queries : user_to_queries)
     {
         auto it = user_queries.second.queries.find(query_id);
         if (user_queries.second.queries.end() == it)
@@ -220,10 +229,27 @@ StoragePtr ProcessList::tryGetTemporaryTable(const String & query_id, const Stri
         if ((*it->second).temporary_tables.end() == jt)
             continue;
 
-        return jt->second;
+        return {& ((*it->second).temporary_tables), jt};
     }
 
-    return {};
+    return {nullptr, Tables::iterator()};
+}
+
+
+StoragePtr ProcessList::tryRemoveTemporaryTable(const String & query_id, const String & table_name) const
+{
+    std::lock_guard<std::mutex> lock(mutex);
+
+    Tables * tables;
+    Tables::iterator iterator;
+    std::tie(tables, iterator) = tryFindTemporaryTable(query_id, table_name);
+
+    if (!tables)
+        return {};
+
+    StoragePtr storage = iterator->second;
+    tables->erase(iterator);
+    return storage;
 }
 
 

dbms/src/Interpreters/ProcessList.h

@@ -226,6 +226,8 @@ private:
     /// Call under lock. Finds process with specified current_user and current_query_id.
     ProcessListElement * tryGetProcessListElement(const String & current_query_id, const String & current_user);
 
+    std::tuple<Tables *, Tables::iterator> tryFindTemporaryTable(const String & query_id, const String & table_name) const;
+
 public:
     ProcessList(size_t max_size_ = 0) : cur_size(0), max_size(max_size_) {}
 
@@ -265,6 +267,8 @@ public:
 
     /// Find temporary table by query_id and name. NOTE: doesn't work fine if there are many queries with same query_id.
     StoragePtr tryGetTemporaryTable(const String & query_id, const String & table_name) const;
+    /// Find temporary table by query_id and name and remove it if exists.
+    StoragePtr tryRemoveTemporaryTable(const String & query_id, const String & table_name) const;
 
 
     enum class CancellationCode