7
7
#include " duckdb/main/extension_util.hpp"
8
8
#include " duckdb/common/atomic.hpp"
9
9
#include " duckdb/common/exception/http_exception.hpp"
10
+ #include " duckdb/common/allocator.hpp"
10
11
11
12
#define CPPHTTPLIB_OPENSSL_SUPPORT
12
13
#include " httplib.hpp"
13
14
14
15
#include < thread>
15
16
#include < memory>
17
+ #include < cstdlib>
16
18
17
19
namespace duckdb {
18
20
@@ -21,7 +23,8 @@ struct HttpServerState {
21
23
std::unique_ptr<std::thread> server_thread;
22
24
std::atomic<bool > is_running;
23
25
DatabaseInstance* db_instance;
24
-
26
+ unique_ptr<Allocator> allocator;
27
+
25
28
HttpServerState () : is_running(false ), db_instance(nullptr ) {}
26
29
};
27
30
@@ -32,16 +35,16 @@ static void HandleQuery(const string& query, duckdb_httplib_openssl::Response& r
32
35
if (!global_state.db_instance ) {
33
36
throw IOException (" Database instance not initialized" );
34
37
}
35
-
38
+
36
39
Connection con (*global_state.db_instance );
37
40
auto result = con.Query (query);
38
-
41
+
39
42
if (result->HasError ()) {
40
43
res.status = 400 ;
41
44
res.set_content (result->GetError (), " text/plain" );
42
45
return ;
43
46
}
44
-
47
+
45
48
res.set_content (result->ToString (), " text/plain" );
46
49
} catch (const Exception& ex) {
47
50
res.status = 400 ;
@@ -53,19 +56,22 @@ void HttpServerStart(DatabaseInstance& db, string_t host, int32_t port) {
53
56
if (global_state.is_running ) {
54
57
throw IOException (" HTTP server is already running" );
55
58
}
56
-
59
+
57
60
global_state.db_instance = &db;
58
- global_state.server . reset ( new duckdb_httplib_openssl::Server () );
61
+ global_state.server = make_uniq< duckdb_httplib_openssl::Server>( );
59
62
global_state.is_running = true ;
60
-
63
+
64
+ // Create a new allocator for the server thread
65
+ global_state.allocator = make_uniq<Allocator>();
66
+
61
67
// Handle GET requests
62
68
global_state.server ->Get (" /query" , [](const duckdb_httplib_openssl::Request& req, duckdb_httplib_openssl::Response& res) {
63
69
if (!req.has_param (" q" )) {
64
70
res.status = 400 ;
65
71
res.set_content (" Missing query parameter 'q'" , " text/plain" );
66
72
return ;
67
73
}
68
-
74
+
69
75
auto query = req.get_param_value (" q" );
70
76
HandleQuery (query, res);
71
77
});
@@ -86,12 +92,12 @@ void HttpServerStart(DatabaseInstance& db, string_t host, int32_t port) {
86
92
});
87
93
88
94
string host_str = host.GetString ();
89
- global_state.server_thread . reset ( new std::thread ([host_str, port]() {
95
+ global_state.server_thread = make_uniq< std::thread> ([host_str, port]() {
90
96
if (!global_state.server ->listen (host_str.c_str (), port)) {
91
97
global_state.is_running = false ;
92
98
throw IOException (" Failed to start HTTP server on " + host_str + " :" + std::to_string (port));
93
99
}
94
- })) ;
100
+ });
95
101
}
96
102
97
103
void HttpServerStop () {
@@ -104,17 +110,24 @@ void HttpServerStop() {
104
110
global_state.server_thread .reset ();
105
111
global_state.db_instance = nullptr ;
106
112
global_state.is_running = false ;
113
+
114
+ // Reset the allocator
115
+ global_state.allocator .reset ();
107
116
}
108
117
}
109
118
119
+ static void HttpServerCleanup () {
120
+ HttpServerStop ();
121
+ }
122
+
110
123
static void LoadInternal (DatabaseInstance &instance) {
111
124
auto httpserve_start = ScalarFunction (" httpserve_start" ,
112
125
{LogicalType::VARCHAR, LogicalType::INTEGER},
113
126
LogicalType::VARCHAR,
114
127
[&](DataChunk &args, ExpressionState &state, Vector &result) {
115
128
auto &host_vector = args.data [0 ];
116
129
auto &port_vector = args.data [1 ];
117
-
130
+
118
131
UnaryExecutor::Execute<string_t , string_t >(
119
132
host_vector, result, args.size (),
120
133
[&](string_t host) {
@@ -134,6 +147,10 @@ static void LoadInternal(DatabaseInstance &instance) {
134
147
135
148
ExtensionUtil::RegisterFunction (instance, httpserve_start);
136
149
ExtensionUtil::RegisterFunction (instance, httpserve_stop);
150
+
151
+ // Register the cleanup function to be called at exit
152
+ std::atexit (HttpServerCleanup);
153
+
137
154
}
138
155
139
156
void HttpserverExtension::Load (DuckDB &db) {
0 commit comments