Minimal changes on ClientBase

chdb-io · Nov 4, 2024 · 744fa13 · 744fa13
1 parent d82a1a7
commit 744fa13
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 20 deletions.
diff --git a/programs/local/LocalServer.cpp b/programs/local/LocalServer.cpp
@@ -1159,7 +1159,7 @@ DB::LocalServer * bgClickHouseLocal(int argc, char ** argv)
     try
     {
         auto * app = new DB::LocalServer();
-        app->is_background = true; // Set background mode
+        app->setBackground(true);
         app->init(argc, argv);
         int ret = app->run();
         if (ret != 0)
@@ -1325,21 +1325,9 @@ struct local_result_v2 * query_conn(chdb_conn * conn, const char * query, const
     {
         auto * result = new local_result_v2{};
         DB::LocalServer * server = static_cast<DB::LocalServer *>(conn->server);
-        // Set output format if specified
-        if (format && *format)
-        {
-            server->client_context->setDefaultFormat(format);
-            server->setDefaultFormat(format);
-        }
-
-        // Check connection and reconnect if needed
-        if (!conn->connected)
-        {
-            server->connect();
-        }
 
         // Execute query
-        if (!server->processQueryText(query))
+        if (!server->parseQueryTextWithOutputFormat(query, format))
         {
             std::string error = server->get_error_msg();
             result->error_message = new char[error.length() + 1];

diff --git a/src/Client/ClientBase.cpp b/src/Client/ClientBase.cpp
@@ -2425,6 +2425,23 @@ bool ClientBase::processQueryText(const String & text)
     return executeMultiQuery(text);
 }
 
+bool ClientBase::parseQueryTextWithOutputFormat(const String & query, const String & format)
+{
+    // Set output format if specified
+    if (!format.empty())
+    {
+        client_context->setDefaultFormat(format);
+        setDefaultFormat(format);
+    }
+
+    // Check connection and reconnect if needed
+    if (!connection->checkConnected(connection_parameters.timeouts))
+        connect();
+
+    // Execute query
+    return processQueryText(query);
+}
+
 
 String ClientBase::prompt() const
 {

diff --git a/src/Client/ClientBase.h b/src/Client/ClientBase.h
@@ -109,18 +109,16 @@ class ClientBase
     double getElapsedTime() const { return progress_indication.elapsedSeconds(); }
     std::string get_error_msg() const { return error_message_oss.str(); }
     void setDefaultFormat(const String & format);
+    void setBackground(bool is_background_) { is_background = is_background_; }
+    bool parseQueryTextWithOutputFormat(const String & query, const String & format);
 
     ASTPtr parseQuery(const char *& pos, const char * end, const Settings & settings, bool allow_multi_statements);
 
-    bool is_background = false;
-
-    /// Returns true if query processing was successful.
-    bool processQueryText(const String & text);
-
 protected:
     void runInteractive();
     void runNonInteractive();
 
+    bool is_background = false;
     void runBackground();
 
     char * argv0 = nullptr;
@@ -176,6 +174,9 @@ class ClientBase
                                 const std::vector<Arguments> & hosts_and_ports_arguments) = 0;
     virtual void processConfig() = 0;
 
+    /// Returns true if query processing was successful.
+    bool processQueryText(const String & text);
+
     virtual void readArguments(
         int argc,
         char ** argv,
@@ -269,7 +270,7 @@ class ClientBase
     ContextMutablePtr global_context;
 
     /// Client context is a context used only by the client to parse queries, process query parameters and to connect to clickhouse-server.
-    public: ContextMutablePtr client_context;
+    ContextMutablePtr client_context;
 
     bool is_interactive = false; /// Use either interactive line editing interface or batch mode.
     bool delayed_interactive = false;