Skip to content

Commit cc0f2bf

Browse files
committed
feat: add storage_rows_read and storage_bytes_read interfaces for retrieving storage metrics
1 parent 774af99 commit cc0f2bf

File tree

5 files changed

+235
-158
lines changed

5 files changed

+235
-158
lines changed

programs/local/LocalChdb.cpp

Lines changed: 30 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,17 @@
11
#include "LocalChdb.h"
2-
#include "LocalServer.h"
32
#include "chdb.h"
43
#include "chdb-internal.h"
54
#include "PythonImporter.h"
65
#include "PythonTableCache.h"
7-
#include "TableFunctionPython.h"
6+
#include "StoragePython.h"
87

9-
#include <mutex>
108
#include <Common/logger_useful.h>
119

1210
namespace py = pybind11;
1311

1412
extern bool inside_main = true;
1513

16-
local_result_v2 * queryToBuffer(
14+
chdb_result * queryToBuffer(
1715
const std::string & queryStr,
1816
const std::string & output_format = "CSV",
1917
const std::string & path = {},
@@ -61,7 +59,7 @@ local_result_v2 * queryToBuffer(
6159
argv_char.push_back(const_cast<char *>(arg.c_str()));
6260

6361
py::gil_scoped_release release;
64-
return query_stable_v2(argv_char.size(), argv_char.data());
62+
return chdb_query_cmdline(argv_char.size(), argv_char.data());
6563
}
6664

6765
// Pybind11 will take over the ownership of the `query_result` object
@@ -219,22 +217,22 @@ connection_wrapper::connection_wrapper(const std::string & conn_str)
219217
argv_char.push_back(const_cast<char *>(arg.c_str()));
220218
}
221219

222-
conn = connect_chdb(argv_char.size(), argv_char.data());
220+
conn = chdb_connect(argv_char.size(), argv_char.data());
223221
db_path = path;
224222
is_memory_db = (path == ":memory:");
225223
}
226224

227225
connection_wrapper::~connection_wrapper()
228226
{
229227
py::gil_scoped_release release;
230-
close_conn(conn);
228+
chdb_close_conn(conn);
231229
}
232230

233231
void connection_wrapper::close()
234232
{
235233
{
236234
py::gil_scoped_release release;
237-
close_conn(conn);
235+
chdb_close_conn(conn);
238236
}
239237
// Ensure that if a new connection is created before this object is destroyed that we don't try to close it.
240238
conn = nullptr;
@@ -255,15 +253,17 @@ query_result * connection_wrapper::query(const std::string & query_str, const st
255253
CHDB::PythonTableCache::findQueryableObjFromQuery(query_str);
256254

257255
py::gil_scoped_release release;
258-
auto * result = query_conn(*conn, query_str.c_str(), format.c_str());
259-
if (result->len == 0)
256+
auto * result = chdb_query(*conn, query_str.c_str(), format.c_str());
257+
if (chdb_result_length(result))
260258
{
261259
LOG_DEBUG(getLogger("CHDB"), "Empty result returned for query: {}", query_str);
262260
}
263-
if (result->error_message)
261+
262+
auto * error_msg = chdb_result_error(result);
263+
if (error_msg)
264264
{
265-
std::string msg_copy(result->error_message);
266-
free_result_v2(result);
265+
std::string msg_copy(error_msg);
266+
chdb_destroy_query_result(result);
267267
throw std::runtime_error(msg_copy);
268268
}
269269
return new query_result(result, false);
@@ -274,12 +274,12 @@ streaming_query_result * connection_wrapper::send_query(const std::string & quer
274274
CHDB::PythonTableCache::findQueryableObjFromQuery(query_str);
275275

276276
py::gil_scoped_release release;
277-
auto * result = query_conn_streaming(*conn, query_str.c_str(), format.c_str());
278-
const auto * error_msg = chdb_streaming_result_error(result);
277+
auto * result = chdb_stream_query(*conn, query_str.c_str(), format.c_str());
278+
auto * error_msg = chdb_result_error(result);
279279
if (error_msg)
280280
{
281281
std::string msg_copy(error_msg);
282-
chdb_destroy_result(result);
282+
chdb_destroy_query_result(result);
283283
throw std::runtime_error(msg_copy);
284284
}
285285

@@ -293,15 +293,16 @@ query_result * connection_wrapper::streaming_fetch_result(streaming_query_result
293293
if (!streaming_result || !streaming_result->get_result())
294294
return nullptr;
295295

296-
auto * result = chdb_streaming_fetch_result(*conn, streaming_result->get_result());
296+
auto * result = chdb_stream_fetch_result(*conn, streaming_result->get_result());
297297

298-
if (result->len == 0)
298+
if (chdb_result_length(result) == 0)
299299
LOG_DEBUG(getLogger("CHDB"), "Empty result returned for streaming query");
300300

301-
if (result->error_message)
301+
auto * error_msg = chdb_result_error(result);
302+
if (error_msg)
302303
{
303-
std::string msg_copy(result->error_message);
304-
free_result_v2(result);
304+
std::string msg_copy(error_msg);
305+
chdb_destroy_query_result(result);
305306
throw std::runtime_error(msg_copy);
306307
}
307308

@@ -315,7 +316,7 @@ void connection_wrapper::streaming_cancel_query(streaming_query_result * streami
315316
if (!streaming_result || !streaming_result->get_result())
316317
return;
317318

318-
chdb_streaming_cancel_query(*conn, streaming_result->get_result());
319+
chdb_stream_cancel_query(*conn, streaming_result->get_result());
319320
}
320321

321322
void cursor_wrapper::execute(const std::string & query_str)
@@ -325,7 +326,7 @@ void cursor_wrapper::execute(const std::string & query_str)
325326

326327
// Use JSONCompactEachRowWithNamesAndTypes format for better type support
327328
py::gil_scoped_release release;
328-
current_result = query_conn(conn->get_conn(), query_str.c_str(), "JSONCompactEachRowWithNamesAndTypes");
329+
current_result = chdb_query(conn->get_conn(), query_str.c_str(), "JSONCompactEachRowWithNamesAndTypes");
329330
}
330331

331332

@@ -390,7 +391,7 @@ PYBIND11_MODULE(_chdb, m)
390391
.def("view", &memoryview_wrapper::view);
391392

392393
py::class_<query_result>(m, "query_result")
393-
.def(py::init<local_result_v2 *>(), py::return_value_policy::take_ownership)
394+
.def(py::init<chdb_result *>(), py::return_value_policy::take_ownership)
394395
.def("data", &query_result::data)
395396
.def("bytes", &query_result::bytes)
396397
.def("__str__", &query_result::str)
@@ -400,13 +401,15 @@ PYBIND11_MODULE(_chdb, m)
400401
.def("size", &query_result::size)
401402
.def("rows_read", &query_result::rows_read)
402403
.def("bytes_read", &query_result::bytes_read)
404+
.def("storage_rows_read", &query_result::storage_rows_read)
405+
.def("storage_bytes_read", &query_result::storage_bytes_read)
403406
.def("elapsed", &query_result::elapsed)
404407
.def("get_memview", &query_result::get_memview)
405408
.def("has_error", &query_result::has_error)
406409
.def("error_message", &query_result::error_message);
407410

408411
py::class_<streaming_query_result>(m, "streaming_query_result")
409-
.def(py::init<chdb_streaming_result *>(), py::return_value_policy::take_ownership)
412+
.def(py::init<chdb_result *>(), py::return_value_policy::take_ownership)
410413
.def("has_error", &streaming_query_result::has_error)
411414
.def("error_message", &streaming_query_result::error_message);
412415

@@ -448,6 +451,8 @@ PYBIND11_MODULE(_chdb, m)
448451
.def("data_size", &cursor_wrapper::data_size)
449452
.def("rows_read", &cursor_wrapper::rows_read)
450453
.def("bytes_read", &cursor_wrapper::bytes_read)
454+
.def("storage_rows_read", &cursor_wrapper::storage_rows_read)
455+
.def("storage_bytes_read", &cursor_wrapper::storage_bytes_read)
451456
.def("elapsed", &cursor_wrapper::elapsed)
452457
.def("has_error", &cursor_wrapper::has_error)
453458
.def("error_message", &cursor_wrapper::error_message);

0 commit comments

Comments
 (0)