Fix settings/current_database in system.processes for async BACKUP/RESTORE

azat · azat · commit 8a48b3334433 · 2024-08-11T21:27:08.000+02:00
Signed-off-by: Azat Khuzhin &lt;a.khuzhin@semrush.com&gt;
diff --git a/src/Backups/BackupsWorker.cpp b/src/Backups/BackupsWorker.cpp
@@ -490,6 +490,8 @@ OperationID BackupsWorker::startMakingBackup(const ASTPtr & query, const Context
 
             /// process_list_element_holder is used to make an element in ProcessList live while BACKUP is working asynchronously.
             auto process_list_element = context_in_use->getProcessListElement();
+            /// Update context to preserve query information in processlist (settings, current_database)
+            process_list_element->updateContext(context_in_use);
 
             thread_pool.scheduleOrThrowOnError(
                 [this,
@@ -853,6 +855,8 @@ OperationID BackupsWorker::startRestoring(const ASTPtr & query, ContextMutablePt
 
             /// process_list_element_holder is used to make an element in ProcessList live while RESTORE is working asynchronously.
             auto process_list_element = context_in_use->getProcessListElement();
+            /// Update context to preserve query information in processlist (settings, current_database)
+            process_list_element->updateContext(context_in_use);
 
             thread_pool.scheduleOrThrowOnError(
                 [this,
diff --git a/src/Interpreters/ProcessList.h b/src/Interpreters/ProcessList.h
@@ -244,6 +244,9 @@ class QueryStatus : public WithContext
     /// Same as checkTimeLimit but it never throws
     [[nodiscard]] bool checkTimeLimitSoft();
 
+    /// Use it in case of the query left in background to execute asynchronously
+    void updateContext(ContextWeakPtr weak_context) { context = std::move(weak_context); }
+
     /// Get the reference for the start of the query. Used to synchronize with other Stopwatches
     UInt64 getQueryCPUStartTime() { return watch.getStart(); }
 };
