Fix changing default database by USE in session mode

Author: auxten

programs/local/LocalServer.cpp

@@ -58,6 +58,7 @@
 #include <boost/program_options/options_description.hpp>
 #include <base/argsToConfig.h>
 #include <filesystem>
+#include <fstream>
 
 #include "config.h"
 
@@ -801,6 +802,42 @@ void LocalServer::processConfig()
         global_context->getUserDefinedSQLObjectsLoader().loadObjects();
 
         LOG_DEBUG(log, "Loaded metadata.");
+
+        /** Set default database if it is specified in default_database file.
+          * NOTE: We do it after loading metadata to let the '_local' and 'system' database initialization
+          * to be done correctly.
+          * I have tried to set default database during parsing '--database' option, but it leads to
+          *   "Code: 82. DB::Exception: Database db_xxx already exists.: while loading database `db_xxx` 
+          *      from path .state_tmp_auxten_usedb_/metadata/db_xxx. (DATABASE_ALREADY_EXISTS)"
+          * This will also happen if we call:
+          *   `clickhouse local --database=db_xxx --path=.state_tmp_auxten_usedb_ --query="select * FROM log_table_xxx"`
+          * with existing `db_xxx` database in the '.state_tmp_auxten_usedb_' directory.
+          */
+        auto default_database_path = fs::path(path) / "default_database";
+        if (std::filesystem::exists(default_database_path))
+        {
+            std::ifstream ifs(default_database_path);
+            std::string user_default_database;
+            if (ifs.is_open())
+            {
+                ifs >> user_default_database;
+                // strip default_database
+                user_default_database.erase(
+                    std::remove_if(
+                        user_default_database.begin(), user_default_database.end(), [](unsigned char x) { return std::isspace(x); }),
+                    user_default_database.end());
+                if (!user_default_database.empty())
+                {
+                    global_context->setCurrentDatabase(user_default_database);
+                    LOG_DEBUG(log, "Set default database to {} recorded in {}", user_default_database, default_database_path);
+                }
+                ifs.close();
+            }
+            else
+            {
+                LOG_ERROR(log, "Cannot read default database from {}", default_database_path);
+            }
+        }
     }
     else if (!config().has("no-system-tables"))
     {

src/Interpreters/InterpreterUseQuery.cpp

@@ -2,17 +2,40 @@
 #include <Interpreters/Context.h>
 #include <Interpreters/InterpreterUseQuery.h>
 #include <Access/Common/AccessFlags.h>
+#include <Common/Exception.h>
 #include <Common/typeid_cast.h>
 
+#include <fstream>
 
 namespace DB
 {
 
+namespace ErrorCodes
+{
+    extern const int CANNOT_OPEN_FILE;
+}
+
 BlockIO InterpreterUseQuery::execute()
 {
     const String & new_database = query_ptr->as<ASTUseQuery &>().getDatabase();
     getContext()->checkAccess(AccessType::SHOW_DATABASES, new_database);
     getContext()->getSessionContext()->setCurrentDatabase(new_database);
+
+    // Save the current using database in default_database stored in getPath()
+    // for the case when the database is changed in chDB session.
+    // The default_database content is used in the LocalServer::processConfig() method.
+    auto default_database_path = fs::path(getContext()->getPath()) / "default_database";
+    std::ofstream tmp_path_fs(default_database_path, std::ofstream::out | std::ofstream::trunc);
+    if (tmp_path_fs && tmp_path_fs.is_open())
+    {
+        tmp_path_fs << new_database;
+        tmp_path_fs.close();
+    }
+    else
+    {
+        throw Exception(ErrorCodes::CANNOT_OPEN_FILE, "Cannot open file {} for writing", default_database_path.string());
+    }
+
     return {};
 }
 

tests/test_usedb.py

@@ -0,0 +1,42 @@
+#!python3
+
+import unittest
+import shutil
+import psutil
+from chdb import session
+
+
+test_state_dir = ".state_tmp_auxten_usedb_"
+current_process = psutil.Process()
+check_thread_count = False
+
+
+class TestStateful(unittest.TestCase):
+    def setUp(self) -> None:
+        shutil.rmtree(test_state_dir, ignore_errors=True)
+        return super().setUp()
+
+    def tearDown(self) -> None:
+        shutil.rmtree(test_state_dir, ignore_errors=True)
+        return super().tearDown()
+
+    def test_path(self):
+        sess = session.Session(test_state_dir)
+
+        sess.query("CREATE DATABASE IF NOT EXISTS db_xxx ENGINE = Atomic", "CSV")
+        ret = sess.query("SHOW DATABASES", "CSV")
+        self.assertIn("db_xxx", str(ret))
+
+        sess.query("CREATE TABLE IF NOT EXISTS db_xxx.log_table_xxx (x UInt8) ENGINE = Log;")
+        sess.query("INSERT INTO db_xxx.log_table_xxx VALUES (1), (2), (3), (4);")
+
+        ret = sess.query("USE db_xxx; SELECT * FROM log_table_xxx", "Debug")
+        self.assertEqual(str(ret), "1\n2\n3\n4\n")
+
+        sess.query("USE db_xxx")
+        ret = sess.query("SELECT * FROM log_table_xxx", "Debug")
+        self.assertEqual(str(ret), "1\n2\n3\n4\n")
+
+
+if __name__ == '__main__':
+    unittest.main()