add has_error and error_message methods

nmreadelf · nmreadelf · commit 814defcebb4c · 2023-09-03T22:49:29.000+08:00
diff --git a/chdb/dbapi/connections.py b/chdb/dbapi/connections.py
@@ -120,6 +120,9 @@ def _execute_command(self, sql):
             print("DEBUG: query:", sql)
         try:
             import chdb
+            res = chdb.query(sql, output_format="JSON")
+            if res.has_error():
+                raise err.DatabaseError(res.error_message())
             self._resp = chdb.query(sql, output_format="JSON").data()
         except Exception as error:
             raise err.InterfaceError("query err: %s" % error)
diff --git a/programs/local/LocalChdb.cpp b/programs/local/LocalChdb.cpp
@@ -1,13 +1,16 @@
 #include "LocalChdb.h"
 
 #include <iostream>
+#include <pybind11/iostream.h>
 
 
 extern bool inside_main = true;
 
 
 local_result * queryToBuffer(const std::string & queryStr, const std::string & format = "CSV", const std::string & path = {})
 {
+    py::scoped_ostream_redirect out_stderr{std::cerr, py::module::import("sys").attr("stderr")};
+    py::scoped_ostream_redirect out_stdout{std::cout, py::module::import("sys").attr("stdout")};
     std::vector<std::string> argv = {"clickhouse", "--multiquery"};
 
     // if format is "Debug" or "debug", then we will add --verbose and --log-level=trace to argv
@@ -129,7 +132,9 @@ PYBIND11_MODULE(_chdb, m)
         .def("rows_read", &query_result::rows_read)
         .def("bytes_read", &query_result::bytes_read)
         .def("elapsed", &query_result::elapsed)
-        .def("get_memview", &query_result::get_memview);
+        .def("get_memview", &query_result::get_memview)
+        .def("has_error", &query_result::has_error)
+        .def("error_message", &query_result::error_message);
 
 
     m.def("query", &query, "Stateless query Clickhouse and return a query_result object");
diff --git a/programs/local/LocalChdb.h b/programs/local/LocalChdb.h
@@ -30,31 +30,31 @@ class local_result_wrapper
         {
             return nullptr;
         }
-        return result->buf;
+        return result->result_->buf_->data() ;
     }
     size_t size()
     {
         if (result == nullptr)
         {
             return 0;
         }
-        return result->len;
+        return result->result_->buf_->size();
     }
     py::bytes bytes()
     {
-        if (result == nullptr)
+        if (result == nullptr || result->result_ == nullptr)
         {
             return py::bytes();
         }
-        return py::bytes(result->buf, result->len);
+        return py::bytes(result->result_->buf_->data(), result->result_->buf_->size());
     }
     py::str str()
     {
         if (result == nullptr)
         {
             return py::str();
         }
-        return py::str(result->buf, result->len);
+        return py::bytes(result->result_->buf_->data(), result->result_->buf_->size());
     }
     // Query statistics
     size_t rows_read()
@@ -63,23 +63,38 @@ class local_result_wrapper
         {
             return 0;
         }
-        return result->rows_read;
+        return result->result_->rows_;
     }
     size_t bytes_read()
     {
         if (result == nullptr)
         {
             return 0;
         }
-        return result->bytes_read;
+        return result->result_->bytes_;
     }
     double elapsed()
     {
         if (result == nullptr)
         {
             return 0;
         }
-        return result->elapsed;
+        return result->result_->elapsed_;
+    }
+
+    bool has_error() {
+        if (result == nullptr)
+        {
+            return false;
+        }
+        return !result->result_->error_msg_.empty();
+    }
+
+    py::str error_message() {
+        if (has_error()) {
+            return py::str(result->result_->error_msg_.data(), result->result_->error_msg_.size());
+        }
+        return py::str();
     }
 };
 
@@ -90,14 +105,19 @@ class query_result
 
 public:
     query_result(local_result * result) : result_wrapper(std::make_shared<local_result_wrapper>(result)) { }
-    ~query_result() { }
+    ~query_result() = default;
     char * data() { return result_wrapper->data(); }
     py::bytes bytes() { return result_wrapper->bytes(); }
     py::str str() { return result_wrapper->str(); }
     size_t size() { return result_wrapper->size(); }
     size_t rows_read() { return result_wrapper->rows_read(); }
     size_t bytes_read() { return result_wrapper->bytes_read(); }
     double elapsed() { return result_wrapper->elapsed(); }
+
+    bool has_error() { return result_wrapper->has_error(); }
+
+    py::str error_message() { return result_wrapper->error_message(); }
+
     memoryview_wrapper * get_memview();
 };
 
diff --git a/programs/local/LocalServer.cpp b/programs/local/LocalServer.cpp
@@ -1,5 +1,6 @@
 #include "LocalServer.h"
 #include "chdb.h"
+#include "query_result.h"
 
 #include <sys/resource.h>
 #include <Common/logger_useful.h>
@@ -52,6 +53,7 @@
 #include <boost/program_options/options_description.hpp>
 #include <base/argsToConfig.h>
 #include <filesystem>
+#include <memory>
 
 #if defined(FUZZING_MODE)
     #include <Functions/getFuzzerData.h>
@@ -536,14 +538,16 @@ catch (const DB::Exception & e)
     cleanup();
 
     bool need_print_stack_trace = config().getBool("stacktrace", false);
-    std::cerr << getExceptionMessage(e, need_print_stack_trace, true) << std::endl;
+    // std::cerr << getExceptionMessage(e, need_print_stack_trace, true) << std::endl;
+    error_message_oss << getExceptionMessage(e, need_print_stack_trace, true);
     return e.code() ? e.code() : -1;
 }
 catch (...)
 {
     cleanup();
 
-    std::cerr << getCurrentExceptionMessage(false) << std::endl;
+    // std::cerr << getCurrentExceptionMessage(false) << std::endl;
+    error_message_oss << getCurrentExceptionMessage(false);
     return getCurrentExceptionCode();
 }
 
@@ -933,14 +937,6 @@ void LocalServer::readArguments(int argc, char ** argv, Arguments & common_argum
 //     }
 // }
 
-class query_result_
-{
-public:
-    uint64_t rows;
-    uint64_t bytes;
-    double elapsed;
-    std::vector<char> * buf;
-};
 
 std::unique_ptr<query_result_> pyEntryClickHouseLocal(int argc, char ** argv)
 {
@@ -949,20 +945,14 @@ std::unique_ptr<query_result_> pyEntryClickHouseLocal(int argc, char ** argv)
         DB::LocalServer app;
         app.init(argc, argv);
         int ret = app.run();
-        if (ret == 0)
-        {
-            auto result = std::make_unique<query_result_>();
-            result->buf = app.getQueryOutputVector();
-            result->rows = app.getProcessedRows();
-            result->bytes = app.getProcessedBytes();
-            result->elapsed = app.getElapsedTime();
-
-            // std::cerr << std::string(out->begin(), out->end()) << std::endl;
-            return result;
-        }
-        else
-        {
-            return nullptr;
+        if (ret == 0) {
+            return std::make_unique<query_result_>(
+                app.getQueryOutputVector(),
+                app.getProcessedRows(),
+                app.getProcessedBytes(),
+                app.getElapsedTime());
+        } else {
+            return std::make_unique<query_result_>(app.get_error_msg());
         }
     }
     catch (const DB::Exception & e)
@@ -984,29 +974,23 @@ std::unique_ptr<query_result_> pyEntryClickHouseLocal(int argc, char ** argv)
 local_result * query_stable(int argc, char ** argv)
 {
     auto result = pyEntryClickHouseLocal(argc, argv);
-    if (!result || !result->buf)
+    if (!result)
     {
         return nullptr;
     }
     local_result * res = new local_result;
-    res->len = result->buf->size();
-    res->buf = result->buf->data();
-    res->_vec = result->buf;
-    res->rows_read = result->rows;
-    res->bytes_read = result->bytes;
-    res->elapsed = result->elapsed;
+    res->result_ = result.release();
     return res;
 }
 
 void free_result(local_result * result)
 {
-    if (!result || !result->_vec)
+    if (!result || result->result_ == nullptr)
     {
         return;
     }
-    std::vector<char> * vec = reinterpret_cast<std::vector<char> *>(result->_vec);
-    delete vec;
-    result->_vec = nullptr;
+    delete result->result_;
+    result->result_ = nullptr;
 }
 
 
@@ -1015,7 +999,7 @@ int mainEntryClickHouseLocal(int argc, char ** argv)
     auto result = pyEntryClickHouseLocal(argc, argv);
     if (result)
     {
-        std::cout << std::string(result->buf->begin(), result->buf->end()) << std::endl;
+        std::cout << (result) << std::endl;
         return 0;
     }
     else
diff --git a/programs/local/chdb.h b/programs/local/chdb.h
@@ -1,16 +1,12 @@
 #pragma once
 #include <cstdint>
 #include <stddef.h>
+#include "query_result.h"
 
 extern "C" {
 struct local_result
 {
-    char * buf;
-    size_t len;
-    void * _vec; // std::vector<char> *, for freeing
-    double elapsed;
-    uint64_t rows_read;
-    uint64_t bytes_read;
+    query_result_* result_;
 };
 
 local_result * query_stable(int argc, char ** argv);
diff --git a/programs/local/query_result.h b/programs/local/query_result.h
@@ -0,0 +1,26 @@
+#pragma once
+#include <vector>
+#include <string>
+#include <iostream>
+#include <memory>
+
+class query_result_
+{
+public:
+    explicit query_result_(std::vector<char>* buf, uint64_t rows,
+                           uint64_t bytes, double elapsed):
+        buf_(buf), rows_(rows), bytes_(bytes),
+        elapsed_(elapsed) { }
+
+    explicit query_result_(std::string&& error_msg): error_msg_(error_msg) { }
+
+    std::string string() {
+        return std::string(buf_->begin(), buf_->end());
+    }
+
+    std::shared_ptr<std::vector<char>> buf_;
+    uint64_t rows_;
+    uint64_t bytes_;
+    double elapsed_;
+    std::string error_msg_;
+};
diff --git a/src/Client/ClientBase.h b/src/Client/ClientBase.h
@@ -84,6 +84,8 @@ class ClientBase : public Poco::Util::Application, public IHints<2, ClientBase>
     size_t getProcessedBytes() const { return processed_bytes; }
     double getElapsedTime() const { return progress_indication.elapsedSeconds(); }
 
+    std::string get_error_msg() const { return error_message_oss.str(); }
+
     std::vector<String> getAllRegisteredNames() const override { return cmd_options; }
 
 protected:
@@ -286,6 +288,7 @@ class ClientBase : public Poco::Util::Application, public IHints<2, ClientBase>
     size_t processed_rows = 0; /// How many rows have been read or written.
     size_t processed_bytes = 0; /// How many bytes have been read or written.
     bool print_num_processed_rows = false; /// Whether to print the number of processed rows at
+    std::stringstream  error_message_oss;
 
     bool print_stack_trace = false;
     /// The last exception that was received from the server. Is used for the
diff --git a/tests/test_basic.py b/tests/test_basic.py
@@ -10,6 +10,9 @@ class TestBasic(unittest.TestCase):
     def test_basic(self):
         res = chdb.query("SELECT 1", "CSV")
         self.assertEqual(len(res), 2) # "1\n"
+        res = chdb.query("SELECT 1", "csv")
+        self.assertTrue(res.has_error())
+        self.assertTrue(len(res.error_message()) > 0)
 class TestOutput(unittest.TestCase):
     def test_output(self):
         for format, output in format_output.items():

Original file line number	Diff line number	Diff line change
`@@ -30,31 +30,31 @@ class local_result_wrapper`
`30`	`30`	`{`
`31`	`31`	`return nullptr;`
`32`	`32`	`}`
`33`		`- return result->buf;`
	`33`	`+ return result->result_->buf_->data() ;`
`34`	`34`	`}`
`35`	`35`	`size_t size()`
`36`	`36`	`{`
`37`	`37`	`if (result == nullptr)`
`38`	`38`	`{`
`39`	`39`	`return 0;`
`40`	`40`	`}`
`41`		`- return result->len;`
	`41`	`+ return result->result_->buf_->size();`
`42`	`42`	`}`
`43`	`43`	`py::bytes bytes()`
`44`	`44`	`{`
`45`		`- if (result == nullptr)`
	`45`	`+ if (result == nullptr \|\| result->result_ == nullptr)`
`46`	`46`	`{`
`47`	`47`	`return py::bytes();`
`48`	`48`	`}`
`49`		`- return py::bytes(result->buf, result->len);`
	`49`	`+ return py::bytes(result->result_->buf_->data(), result->result_->buf_->size());`
`50`	`50`	`}`
`51`	`51`	`py::str str()`
`52`	`52`	`{`
`53`	`53`	`if (result == nullptr)`
`54`	`54`	`{`
`55`	`55`	`return py::str();`
`56`	`56`	`}`
`57`		`- return py::str(result->buf, result->len);`
	`57`	`+ return py::bytes(result->result_->buf_->data(), result->result_->buf_->size());`
`58`	`58`	`}`
`59`	`59`	`// Query statistics`
`60`	`60`	`size_t rows_read()`
`@@ -63,23 +63,38 @@ class local_result_wrapper`
`63`	`63`	`{`
`64`	`64`	`return 0;`
`65`	`65`	`}`
`66`		`- return result->rows_read;`
	`66`	`+ return result->result_->rows_;`
`67`	`67`	`}`
`68`	`68`	`size_t bytes_read()`
`69`	`69`	`{`
`70`	`70`	`if (result == nullptr)`
`71`	`71`	`{`
`72`	`72`	`return 0;`
`73`	`73`	`}`
`74`		`- return result->bytes_read;`
	`74`	`+ return result->result_->bytes_;`
`75`	`75`	`}`
`76`	`76`	`double elapsed()`
`77`	`77`	`{`
`78`	`78`	`if (result == nullptr)`
`79`	`79`	`{`
`80`	`80`	`return 0;`
`81`	`81`	`}`
`82`		`- return result->elapsed;`
	`82`	`+ return result->result_->elapsed_;`
	`83`	`+ }`
	`84`	`+`
	`85`	`+ bool has_error() {`
	`86`	`+ if (result == nullptr)`
	`87`	`+ {`
	`88`	`+ return false;`
	`89`	`+ }`
	`90`	`+ return !result->result_->error_msg_.empty();`
	`91`	`+ }`
	`92`	`+`
	`93`	`+ py::str error_message() {`
	`94`	`+ if (has_error()) {`
	`95`	`+ return py::str(result->result_->error_msg_.data(), result->result_->error_msg_.size());`
	`96`	`+ }`
	`97`	`+ return py::str();`
`83`	`98`	`}`
`84`	`99`	`};`
`85`	`100`
`@@ -90,14 +105,19 @@ class query_result`
`90`	`105`
`91`	`106`	`public:`
`92`	`107`	`query_result(local_result * result) : result_wrapper(std::make_shared<local_result_wrapper>(result)) { }`
`93`		`- ~query_result() { }`
	`108`	`+ ~query_result() = default;`
`94`	`109`	`char * data() { return result_wrapper->data(); }`
`95`	`110`	`py::bytes bytes() { return result_wrapper->bytes(); }`
`96`	`111`	`py::str str() { return result_wrapper->str(); }`
`97`	`112`	`size_t size() { return result_wrapper->size(); }`
`98`	`113`	`size_t rows_read() { return result_wrapper->rows_read(); }`
`99`	`114`	`size_t bytes_read() { return result_wrapper->bytes_read(); }`
`100`	`115`	`double elapsed() { return result_wrapper->elapsed(); }`
	`116`	`+`
	`117`	`+ bool has_error() { return result_wrapper->has_error(); }`
	`118`	`+`
	`119`	`+ py::str error_message() { return result_wrapper->error_message(); }`
	`120`	`+`
`101`	`121`	`memoryview_wrapper * get_memview();`
`102`	`122`	`};`
`103`	`123`