Merge pull request #3 from lmangani/json-output

JSON Compact output
quackscience · Oct 14, 2024 · 9190feb · 9190feb
2 parents d78d875 + 0fd5be2
commit 9190feb
Showing 1 changed file with 57 additions and 2 deletions.
diff --git a/src/httpserver_extension.cpp b/src/httpserver_extension.cpp
@@ -12,10 +12,15 @@
 #define CPPHTTPLIB_OPENSSL_SUPPORT
 #include "httplib.hpp"
 
+// Include yyjson for JSON handling
+#include "yyjson.hpp"
+
 #include <thread>
 #include <memory>
 #include <cstdlib>
 
+using namespace duckdb_yyjson; // NOLINT
+
 namespace duckdb {
 
 struct HttpServerState {
@@ -30,6 +35,54 @@ struct HttpServerState {
 
 static HttpServerState global_state;
 
+// Convert the query result to JSON format
+static std::string ConvertResultToJSON(MaterializedQueryResult &result) {
+    auto doc = yyjson_mut_doc_new(nullptr);
+    auto root = yyjson_mut_obj(doc);
+    yyjson_mut_doc_set_root(doc, root);
+
+    // Add meta information
+    auto meta_array = yyjson_mut_arr(doc);
+    for (idx_t col = 0; col < result.ColumnCount(); ++col) {
+        auto column_obj = yyjson_mut_obj(doc);
+        yyjson_mut_obj_add_str(doc, column_obj, "name", result.ColumnName(col).c_str());
+        yyjson_mut_arr_append(meta_array, column_obj);
+    }
+    yyjson_mut_obj_add_val(doc, root, "meta", meta_array);
+
+    // Add data
+    auto data_array = yyjson_mut_arr(doc);
+    for (idx_t row = 0; row < result.RowCount(); ++row) {
+        auto row_array = yyjson_mut_arr(doc);
+        for (idx_t col = 0; col < result.ColumnCount(); ++col) {
+            Value value = result.GetValue(col, row);
+            if (value.IsNull()) {
+                yyjson_mut_arr_append(row_array, yyjson_mut_null(doc));
+            } else {
+                std::string value_str = value.ToString();
+                yyjson_mut_arr_append(row_array, yyjson_mut_strncpy(doc, value_str.c_str(), value_str.length()));
+            }
+        }
+        yyjson_mut_arr_append(data_array, row_array);
+    }
+    yyjson_mut_obj_add_val(doc, root, "data", data_array);
+
+    // Add row count
+    yyjson_mut_obj_add_int(doc, root, "rows", result.RowCount());
+
+    // Write to string
+    auto data = yyjson_mut_write(doc, 0, nullptr);
+    if (!data) {
+        yyjson_mut_doc_free(doc);
+        throw InternalException("Failed to render the result as JSON, yyjson failed");
+    }
+
+    std::string json_output(data);
+    free(data);
+    yyjson_mut_doc_free(doc);
+    return json_output;
+}
+
 static void HandleQuery(const string& query, duckdb_httplib_openssl::Response& res) {
     try {
         if (!global_state.db_instance) {
@@ -45,13 +98,16 @@ static void HandleQuery(const string& query, duckdb_httplib_openssl::Response& r
             return;
         }
 
-        res.set_content(result->ToString(), "text/plain");
+        // Convert result to JSON
+        std::string json_output = ConvertResultToJSON(*result);
+        res.set_content(json_output, "application/json");
     } catch (const Exception& ex) {
         res.status = 400;
         res.set_content(ex.what(), "text/plain");
     }
 }
 
+
 void HttpServerStart(DatabaseInstance& db, string_t host, int32_t port) {
     if (global_state.is_running) {
         throw IOException("HTTP server is already running");
@@ -150,7 +206,6 @@ static void LoadInternal(DatabaseInstance &instance) {
 
     // Register the cleanup function to be called at exit
     std::atexit(HttpServerCleanup);
-
 }
 
 void HttpserverExtension::Load(DuckDB &db) {