1use crate::error::{AuroraError, Result};
2use crate::network::protocol::{Request, Response};
3use crate::query::SimpleQueryBuilder;
4use crate::types::{Document, FieldType, Value};
5use std::collections::HashMap;
6use tokio::io::{AsyncReadExt, AsyncWriteExt};
7use tokio::net::TcpStream;
8
9pub struct Client {
10 stream: TcpStream,
11}
12
13impl Client {
14 pub async fn connect(addr: &str) -> Result<Self> {
16 let stream = TcpStream::connect(addr).await?;
17 Ok(Self { stream })
18 }
19
20 async fn send_request(&mut self, request: Request) -> Result<Response> {
22 let request_bytes = bincode::serialize(&request).map_err(AuroraError::Bincode)?;
23 let len_bytes = (request_bytes.len() as u32).to_le_bytes();
24
25 self.stream.write_all(&len_bytes).await?;
26 self.stream.write_all(&request_bytes).await?;
27
28 let mut len_bytes = [0u8; 4];
29 self.stream.read_exact(&mut len_bytes).await?;
30 let len = u32::from_le_bytes(len_bytes) as usize;
31
32 let mut buffer = vec![0u8; len];
33 self.stream.read_exact(&mut buffer).await?;
34
35 let response: Response = bincode::deserialize(&buffer).map_err(AuroraError::Bincode)?;
36
37 Ok(response)
38 }
39
40 pub async fn new_collection(
41 &mut self,
42 name: &str,
43 fields: Vec<(String, FieldType, bool)>,
44 ) -> Result<()> {
45 let req = Request::NewCollection {
46 name: name.to_string(),
47 fields,
48 };
49 match self.send_request(req).await? {
50 Response::Done => Ok(()),
51 Response::Error(e) => Err(AuroraError::Protocol(e)),
52 _ => Err(AuroraError::Protocol("Unexpected response".into())),
53 }
54 }
55
56 pub async fn insert(
57 &mut self,
58 collection: &str,
59 data: HashMap<String, Value>,
60 ) -> Result<String> {
61 let req = Request::Insert {
62 collection: collection.to_string(),
63 data,
64 };
65 match self.send_request(req).await? {
66 Response::Message(id) => Ok(id),
67 Response::Error(e) => Err(AuroraError::Protocol(e)),
68 _ => Err(AuroraError::Protocol("Unexpected response".into())),
69 }
70 }
71
72 pub async fn get_document(&mut self, collection: &str, id: &str) -> Result<Option<Document>> {
73 let req = Request::GetDocument {
74 collection: collection.to_string(),
75 id: id.to_string(),
76 };
77 match self.send_request(req).await? {
78 Response::Document(doc) => Ok(doc),
79 Response::Error(e) => Err(AuroraError::Protocol(e)),
80 _ => Err(AuroraError::Protocol("Unexpected response".into())),
81 }
82 }
83
84 pub async fn query(&mut self, builder: SimpleQueryBuilder) -> Result<Vec<Document>> {
85 let req = Request::Query(builder);
86 match self.send_request(req).await? {
87 Response::Documents(docs) => Ok(docs),
88 Response::Error(e) => Err(AuroraError::Protocol(e)),
89 _ => Err(AuroraError::Protocol("Unexpected response".into())),
90 }
91 }
92
93 pub async fn begin_transaction(&mut self) -> Result<()> {
94 match self.send_request(Request::BeginTransaction).await? {
95 Response::Done => Ok(()),
96 Response::Error(e) => Err(AuroraError::Protocol(e)),
97 _ => Err(AuroraError::Protocol("Unexpected response".into())),
98 }
99 }
100
101 pub async fn commit_transaction(&mut self) -> Result<()> {
102 match self.send_request(Request::CommitTransaction).await? {
103 Response::Done => Ok(()),
104 Response::Error(e) => Err(AuroraError::Protocol(e)),
105 _ => Err(AuroraError::Protocol("Unexpected response".into())),
106 }
107 }
108
109 pub async fn delete(&mut self, key: &str) -> Result<()> {
110 match self.send_request(Request::Delete(key.to_string())).await? {
111 Response::Done => Ok(()),
112 Response::Error(e) => Err(AuroraError::Protocol(e)),
113 _ => Err(AuroraError::Protocol("Unexpected response".into())),
114 }
115 }
116}