1use std::collections::HashMap;
2use std::io::{BufRead, BufReader, Write};
3use std::net::TcpStream;
4use serde::{Deserialize, Serialize};
5
6#[derive(Debug, Serialize, Deserialize)]
7pub struct Request {
8 action: String,
9 key: Option<String>,
10 value: Option<String>,
11 cf_name: Option<String>,
12 default_value: Option<String>,
13 options: Option<HashMap<String, String>>,
14 backup_path: Option<String>,
15 num_backups_to_keep: Option<usize>,
16 backup_id: Option<u32>,
17 restore_path: Option<String>,
18 iterator_id: Option<usize>,
19 txn: Option<bool>,
20}
21
22#[derive(Debug, Serialize, Deserialize)]
23pub struct Response {
24 pub success: bool,
25 pub result: Option<String>,
26}
27
28pub struct RequestHandler {
29 host: String,
30 port: u16,
31 connection: Option<TcpStream>,
32}
33
34impl RequestHandler {
35 pub fn new(host: String, port: u16) -> Self {
36 Self {
37 host,
38 port,
39 connection: None,
40 }
41 }
42
43 fn get_connection(&mut self) -> Result<&mut TcpStream, String> {
44 if self.connection.is_none() {
45 let addr = format!("{}:{}", self.host, self.port);
46 let stream = TcpStream::connect(&addr).map_err(|e| format!("Connection error: {}", e))?;
47 self.connection = Some(stream);
48 }
49 self.connection.as_mut().ok_or_else(|| "Failed to acquire connection".to_string())
50 }
51
52 pub fn send_request(&mut self, request: Request) -> Result<Response, String> {
53 let conn = self.get_connection()?;
54
55 let request_bytes = serde_json::to_vec(&request).map_err(|e| format!("Serialization error: {}", e))?;
56 conn.write_all(&request_bytes).map_err(|e| format!("Send error: {}", e))?;
57 conn.write_all(b"\n").map_err(|e| format!("Send error: {}", e))?;
58
59 let mut reader = BufReader::new(conn);
60 let mut response_bytes = Vec::new();
61 reader.read_until(b'\n', &mut response_bytes).map_err(|e| format!("Receive error: {}", e))?;
62
63 let response: Response = serde_json::from_slice(&response_bytes).map_err(|e| format!("Deserialization error: {}", e))?;
64 Ok(response)
65 }
66
67 pub fn handle_response(&self, response: Response) -> Result<Option<String>, String> {
68 if response.success {
69 Ok(response.result)
70 } else {
71 Err(response.result.unwrap_or("Unknown error".to_string()))
72 }
73 }
74}
75
76
77#[derive(Debug, Serialize, Deserialize)]
78pub struct RequestBuilder {
79 request: Request,
80}
81
82impl RequestBuilder {
83 pub fn new(action: &str) -> Self {
84 Self {
85 request: Request {
86 action: action.to_string(),
87 key: None,
88 value: None,
89 default_value: None,
90 cf_name: None,
91 options: None,
92 backup_path: None,
93 num_backups_to_keep: None,
94 backup_id: None,
95 restore_path: None,
96 iterator_id: None,
97 txn: None,
98 },
99 }
100 }
101
102 pub fn key(mut self, key: Option<String>) -> Self {
103 self.request.key = key;
104 self
105 }
106
107 pub fn value(mut self, value: Option<String>) -> Self {
108 self.request.value = value;
109 self
110 }
111 pub fn default_value(mut self, value: Option<String>) -> Self {
112 self.request.default_value = value;
113 self
114 }
115
116 pub fn cf_name(mut self, cf_name: Option<String>) -> Self {
117 self.request.cf_name = cf_name;
118 self
119 }
120
121 pub fn num_backups_to_keep(mut self, num_backups_to_keep: Option<usize>) -> Self {
122 self.request.num_backups_to_keep = num_backups_to_keep;
123 self
124 }
125
126 pub fn backup_id(mut self, backup_id: Option<u32>) -> Self {
127 self.request.backup_id = backup_id;
128 self
129 }
130
131 pub fn restore_path(mut self, restore_path: Option<String>) -> Self {
132 self.request.restore_path = restore_path;
133 self
134 }
135
136 pub fn iterator_id(mut self, iterator_id: Option<usize>) -> Self {
137 self.request.iterator_id = iterator_id;
138 self
139 }
140
141 pub fn txn(mut self, txn: Option<bool>) -> Self {
142 self.request.txn = txn;
143 self
144 }
145
146 pub fn option(mut self, key: String, value: String) -> Self {
147 if self.request.options.is_none() {
148 self.request.options = Some(HashMap::new());
149 }
150 self.request.options.as_mut().unwrap().insert(key, value);
151 self
152 }
153
154 pub fn build(self) -> Request {
155 self.request
156 }
157}
158
159pub struct RocksDBClient {
160 request_handler: RequestHandler,
161}
162
163impl RocksDBClient {
164 pub fn new(host: String, port: u16) -> Self {
165 Self {
166 request_handler: RequestHandler::new(host, port),
167 }
168 }
169
170 pub fn put(&mut self, key: String, value: String, cf_name: Option<String>, txn: Option<bool>) -> Result<Option<String>, String> {
171 let request = RequestBuilder::new("put")
172 .key(Some(key))
173 .value(Some(value))
174 .cf_name(cf_name)
175 .txn(txn)
176 .build();
177
178 let response = self.request_handler.send_request(request)?;
179 self.request_handler.handle_response(response)
180 }
181
182 pub fn get(&mut self, key: String, cf_name: Option<String>, default_value: Option<String>, txn: Option<bool>) -> Result<Option<String>, String> {
183 let request = RequestBuilder::new("get")
184 .key(Some(key))
185 .cf_name(cf_name)
186 .default_value(default_value)
187 .txn(txn)
188 .build();
189
190
191 let response = self.request_handler.send_request(request)?;
192 self.request_handler.handle_response(response)
193 }
194
195 pub fn delete(&mut self, key: String, cf_name: Option<String>, txn: Option<bool>) -> Result<Option<String>, String> {
196 let request = RequestBuilder::new("delete")
197 .key(Some(key))
198 .cf_name(cf_name)
199 .txn(txn)
200 .build();
201
202 let response = self.request_handler.send_request(request)?;
203 self.request_handler.handle_response(response)
204 }
205
206 pub fn merge(&mut self, key: String, value: String, cf_name: Option<String>, txn: Option<bool>) -> Result<Option<String>, String> {
207 let request = RequestBuilder::new("merge")
208 .key(Some(key))
209 .value(Some(value))
210 .cf_name(cf_name)
211 .txn(txn)
212 .build();
213
214 let response = self.request_handler.send_request(request)?;
215 self.request_handler.handle_response(response)
216 }
217
218 pub fn get_property(&mut self, value: String, cf_name: Option<String>) -> Result<Option<String>, String> {
219 let request = RequestBuilder::new("get_property")
220 .value(Some(value))
221 .cf_name(cf_name)
222 .build();
223
224 let response = self.request_handler.send_request(request)?;
225 self.request_handler.handle_response(response)
226 }
227
228 pub fn keys(&mut self, start: String, limit: String, query: Option<String>) -> Result<Option<String>, String> {
229 let request = RequestBuilder::new("keys")
230 .option("start".to_string(), start)
231 .option("limit".to_string(), limit)
232 .option("query".to_string(), query.unwrap_or_default())
233 .build();
234
235 let response = self.request_handler.send_request(request)?;
236 self.request_handler.handle_response(response)
237 }
238
239 pub fn all(&mut self, query: Option<String>) -> Result<Option<String>, String> {
240 let request = RequestBuilder::new("all")
241 .option("query".to_string(), query.unwrap_or_default())
242 .build();
243
244 let response = self.request_handler.send_request(request)?;
245 self.request_handler.handle_response(response)
246 }
247
248 pub fn list_column_families(&mut self) -> Result<Option<String>, String> {
249 let request = RequestBuilder::new("list_column_families")
250 .build();
251
252 let response = self.request_handler.send_request(request)?;
253 self.request_handler.handle_response(response)
254 }
255
256 pub fn create_column_family(&mut self, cf_name: String) -> Result<Option<String>, String> {
257 let request = RequestBuilder::new("create_column_family")
258 .cf_name(Some(cf_name))
259 .build();
260
261 let response = self.request_handler.send_request(request)?;
262 self.request_handler.handle_response(response)
263 }
264
265 pub fn drop_column_family(&mut self, cf_name: String) -> Result<Option<String>, String> {
266 let request = RequestBuilder::new("drop_column_family")
267 .cf_name(Some(cf_name))
268 .build();
269
270 let response = self.request_handler.send_request(request)?;
271 self.request_handler.handle_response(response)
272 }
273
274 pub fn compact_range(&mut self, start: Option<String>, end: Option<String>, cf_name: Option<String>) -> Result<Option<String>, String> {
275 let request = RequestBuilder::new("compact_range")
276 .option("start".to_string(), start.unwrap_or_default())
277 .option("end".to_string(), end.unwrap_or_default())
278 .cf_name(cf_name)
279 .build();
280
281 let response = self.request_handler.send_request(request)?;
282 self.request_handler.handle_response(response)
283 }
284
285 pub fn write_batch_put(&mut self, key: String, value: String, cf_name: Option<String>) -> Result<Option<String>, String> {
286 let request = RequestBuilder::new("write_batch_put")
287 .key(Some(key))
288 .value(Some(value))
289 .cf_name(cf_name)
290 .build();
291
292 let response = self.request_handler.send_request(request)?;
293 self.request_handler.handle_response(response)
294 }
295
296 pub fn write_batch_merge(&mut self, key: String, value: String, cf_name: Option<String>) -> Result<Option<String>, String> {
297 let request = RequestBuilder::new("write_batch_merge")
298 .key(Some(key))
299 .value(Some(value))
300 .cf_name(cf_name)
301 .build();
302
303 let response = self.request_handler.send_request(request)?;
304 self.request_handler.handle_response(response)
305 }
306
307 pub fn write_batch_delete(&mut self, key: String, cf_name: Option<String>) -> Result<Option<String>, String> {
308 let request = RequestBuilder::new("write_batch_delete")
309 .key(Some(key))
310 .cf_name(cf_name)
311 .build();
312
313 let response = self.request_handler.send_request(request)?;
314 self.request_handler.handle_response(response)
315 }
316
317 pub fn write_batch_write(&mut self) -> Result<Option<String>, String> {
318 let request = RequestBuilder::new("write_batch_write")
319 .build();
320
321 let response = self.request_handler.send_request(request)?;
322 self.request_handler.handle_response(response)
323 }
324
325 pub fn write_batch_clear(&mut self) -> Result<Option<String>, String> {
326 let request = RequestBuilder::new("write_batch_clear")
327 .build();
328
329 let response = self.request_handler.send_request(request)?;
330 self.request_handler.handle_response(response)
331 }
332
333 pub fn write_batch_destroy(&mut self) -> Result<Option<String>, String> {
334 let request = RequestBuilder::new("write_batch_destroy")
335 .build();
336
337 let response = self.request_handler.send_request(request)?;
338 self.request_handler.handle_response(response)
339 }
340
341 pub fn create_iterator(&mut self) -> Result<Option<String>, String> {
342 let request = RequestBuilder::new("create_iterator")
343 .build();
344
345 let response = self.request_handler.send_request(request)?;
346 self.request_handler.handle_response(response)
347 }
348
349 pub fn destroy_iterator(&mut self, iterator_id: String) -> Result<Option<String>, String> {
350 let request = RequestBuilder::new("destroy_iterator")
351 .option("iterator_id".to_string(), iterator_id)
352 .build();
353
354 let response = self.request_handler.send_request(request)?;
355 self.request_handler.handle_response(response)
356 }
357
358 pub fn iterator_seek(&mut self, iterator_id: String, key: String) -> Result<Option<String>, String> {
359 let request = RequestBuilder::new("iterator_seek")
360 .option("iterator_id".to_string(), iterator_id)
361 .key(Some(key))
362 .build();
363
364 let response = self.request_handler.send_request(request)?;
365 self.request_handler.handle_response(response)
366 }
367
368 pub fn iterator_next(&mut self, iterator_id: String) -> Result<Option<String>, String> {
369 let request = RequestBuilder::new("iterator_next")
370 .option("iterator_id".to_string(), iterator_id)
371 .build();
372
373 let response = self.request_handler.send_request(request)?;
374 self.request_handler.handle_response(response)
375 }
376
377 pub fn iterator_prev(&mut self, iterator_id: String) -> Result<Option<String>, String> {
378 let request = RequestBuilder::new("iterator_prev")
379 .option("iterator_id".to_string(), iterator_id)
380 .build();
381
382 let response = self.request_handler.send_request(request)?;
383 self.request_handler.handle_response(response)
384 }
385
386 pub fn backup(&mut self) -> Result<Option<String>, String> {
387 let request = RequestBuilder::new("backup")
388 .build();
389
390 let response = self.request_handler.send_request(request)?;
391 self.request_handler.handle_response(response)
392 }
393
394 pub fn restore_latest(&mut self) -> Result<Option<String>, String> {
395 let request = RequestBuilder::new("restore_latest")
396 .build();
397
398 let response = self.request_handler.send_request(request)?;
399 self.request_handler.handle_response(response)
400 }
401
402 pub fn restore(&mut self, backup_id: String) -> Result<Option<String>, String> {
403 let request = RequestBuilder::new("restore")
404 .option("backup_id".to_string(), backup_id)
405 .build();
406
407 let response = self.request_handler.send_request(request)?;
408 self.request_handler.handle_response(response)
409 }
410
411 pub fn get_backup_info(&mut self) -> Result<Option<String>, String> {
412 let request = RequestBuilder::new("get_backup_info")
413 .build();
414
415 let response = self.request_handler.send_request(request)?;
416 self.request_handler.handle_response(response)
417 }
418
419 pub fn begin_transaction(&mut self) -> Result<Option<String>, String> {
420 let request = RequestBuilder::new("begin_transaction")
421 .build();
422
423 let response = self.request_handler.send_request(request)?;
424 self.request_handler.handle_response(response)
425 }
426
427 pub fn commit_transaction(&mut self) -> Result<Option<String>, String> {
428 let request = RequestBuilder::new("commit_transaction")
429 .build();
430
431 let response = self.request_handler.send_request(request)?;
432 self.request_handler.handle_response(response)
433 }
434
435 pub fn rollback_transaction(&mut self) -> Result<Option<String>, String> {
436 let request = RequestBuilder::new("rollback_transaction")
437 .build();
438
439 let response = self.request_handler.send_request(request)?;
440 self.request_handler.handle_response(response)
441 }
442}