1#![allow(clippy::too_many_arguments)]
6
7use spatio_server::SpatioServiceClient;
8use spatio_types::geo::{DistanceMetric, Point, Polygon};
9use spatio_types::point::Point3d;
10use std::net::SocketAddr;
11use std::time::Duration;
12use tarpc::client;
13use tarpc::context;
14use tarpc::tokio_serde::formats::Json;
15use thiserror::Error;
16use tokio_util::codec::{Framed, LengthDelimitedCodec};
17
18#[derive(Error, Debug)]
19pub enum ClientError {
20 #[error("Connection error: {0}")]
21 Connection(#[from] std::io::Error),
22 #[error("RPC error: {0}")]
23 Rpc(#[from] tarpc::client::RpcError),
24 #[error("Server error: {0}")]
25 Server(String),
26 #[error("Serialization error: {0}")]
27 Serialization(#[from] serde_json::Error),
28}
29
30pub type Result<T> = std::result::Result<T, ClientError>;
31
32#[derive(Clone)]
33pub struct SpatioClient {
34 client: SpatioServiceClient,
35}
36
37impl SpatioClient {
38 pub async fn connect(addr: SocketAddr) -> Result<Self> {
39 let socket = tokio::net::TcpStream::connect(addr).await?;
40 let framed = Framed::new(socket, LengthDelimitedCodec::new());
41 let transport = tarpc::serde_transport::new(framed, Json::default());
42 let client = SpatioServiceClient::new(client::Config::default(), transport).spawn();
43 Ok(Self { client })
44 }
45
46 fn make_context(&self) -> context::Context {
47 let mut ctx = context::current();
48 ctx.deadline = std::time::SystemTime::now() + Duration::from_secs(30);
49 ctx
50 }
51
52 pub async fn upsert(
53 &self,
54 namespace: &str,
55 id: &str,
56 point: Point3d,
57 metadata: serde_json::Value,
58 ) -> Result<()> {
59 self.client
60 .upsert(
61 self.make_context(),
62 namespace.to_string(),
63 id.to_string(),
64 point,
65 metadata,
66 )
67 .await?
68 .map_err(ClientError::Server)
69 }
70
71 pub async fn get(
72 &self,
73 namespace: &str,
74 id: &str,
75 ) -> Result<Option<spatio_server::CurrentLocation>> {
76 self.client
77 .get(self.make_context(), namespace.to_string(), id.to_string())
78 .await?
79 .map_err(ClientError::Server)
80 }
81
82 pub async fn delete(&self, namespace: &str, id: &str) -> Result<()> {
83 self.client
84 .delete(self.make_context(), namespace.to_string(), id.to_string())
85 .await?
86 .map_err(ClientError::Server)
87 }
88
89 pub async fn query_radius(
90 &self,
91 namespace: &str,
92 center: Point3d,
93 radius: f64,
94 limit: usize,
95 ) -> Result<Vec<(spatio_server::CurrentLocation, f64)>> {
96 self.client
97 .query_radius(
98 self.make_context(),
99 namespace.to_string(),
100 center,
101 radius,
102 limit,
103 )
104 .await?
105 .map_err(ClientError::Server)
106 }
107
108 pub async fn knn(
109 &self,
110 namespace: &str,
111 center: Point3d,
112 k: usize,
113 ) -> Result<Vec<(spatio_server::CurrentLocation, f64)>> {
114 self.client
115 .knn(self.make_context(), namespace.to_string(), center, k)
116 .await?
117 .map_err(ClientError::Server)
118 }
119
120 pub async fn stats(&self) -> Result<spatio_server::Stats> {
121 Ok(self.client.stats(self.make_context()).await?)
122 }
123
124 pub async fn query_bbox(
125 &self,
126 namespace: &str,
127 min_x: f64,
128 min_y: f64,
129 max_x: f64,
130 max_y: f64,
131 limit: usize,
132 ) -> Result<Vec<spatio_server::CurrentLocation>> {
133 self.client
134 .query_bbox(
135 self.make_context(),
136 namespace.to_string(),
137 min_x,
138 min_y,
139 max_x,
140 max_y,
141 limit,
142 )
143 .await?
144 .map_err(ClientError::Server)
145 }
146
147 pub async fn query_cylinder(
148 &self,
149 namespace: &str,
150 center: Point,
151 min_z: f64,
152 max_z: f64,
153 radius: f64,
154 limit: usize,
155 ) -> Result<Vec<(spatio_server::CurrentLocation, f64)>> {
156 self.client
157 .query_cylinder(
158 self.make_context(),
159 namespace.to_string(),
160 center,
161 min_z,
162 max_z,
163 radius,
164 limit,
165 )
166 .await?
167 .map_err(ClientError::Server)
168 }
169
170 pub async fn query_trajectory(
171 &self,
172 namespace: &str,
173 id: &str,
174 start_time: Option<f64>,
175 end_time: Option<f64>,
176 limit: usize,
177 ) -> Result<Vec<spatio_server::LocationUpdate>> {
178 self.client
179 .query_trajectory(
180 self.make_context(),
181 namespace.to_string(),
182 id.to_string(),
183 start_time,
184 end_time,
185 limit,
186 )
187 .await?
188 .map_err(ClientError::Server)
189 }
190
191 pub async fn insert_trajectory(
192 &self,
193 namespace: &str,
194 id: &str,
195 trajectory: Vec<(f64, Point3d, serde_json::Value)>,
196 ) -> Result<()> {
197 self.client
198 .insert_trajectory(
199 self.make_context(),
200 namespace.to_string(),
201 id.to_string(),
202 trajectory,
203 )
204 .await?
205 .map_err(ClientError::Server)
206 }
207
208 pub async fn query_bbox_3d(
209 &self,
210 namespace: &str,
211 min_x: f64,
212 min_y: f64,
213 min_z: f64,
214 max_x: f64,
215 max_y: f64,
216 max_z: f64,
217 limit: usize,
218 ) -> Result<Vec<spatio_server::CurrentLocation>> {
219 self.client
220 .query_bbox_3d(
221 self.make_context(),
222 namespace.to_string(),
223 min_x,
224 min_y,
225 min_z,
226 max_x,
227 max_y,
228 max_z,
229 limit,
230 )
231 .await?
232 .map_err(ClientError::Server)
233 }
234
235 pub async fn query_near(
236 &self,
237 namespace: &str,
238 id: &str,
239 radius: f64,
240 limit: usize,
241 ) -> Result<Vec<(spatio_server::CurrentLocation, f64)>> {
242 self.client
243 .query_near(
244 self.make_context(),
245 namespace.to_string(),
246 id.to_string(),
247 radius,
248 limit,
249 )
250 .await?
251 .map_err(ClientError::Server)
252 }
253
254 pub async fn contains(
255 &self,
256 namespace: &str,
257 polygon: Polygon,
258 limit: usize,
259 ) -> Result<Vec<spatio_server::CurrentLocation>> {
260 self.client
261 .contains(self.make_context(), namespace.to_string(), polygon, limit)
262 .await?
263 .map_err(ClientError::Server)
264 }
265
266 pub async fn distance(
267 &self,
268 namespace: &str,
269 id1: &str,
270 id2: &str,
271 metric: Option<DistanceMetric>,
272 ) -> Result<Option<f64>> {
273 self.client
274 .distance(
275 self.make_context(),
276 namespace.to_string(),
277 id1.to_string(),
278 id2.to_string(),
279 metric,
280 )
281 .await?
282 .map_err(ClientError::Server)
283 }
284
285 pub async fn distance_to(
286 &self,
287 namespace: &str,
288 id: &str,
289 point: Point,
290 metric: Option<DistanceMetric>,
291 ) -> Result<Option<f64>> {
292 self.client
293 .distance_to(
294 self.make_context(),
295 namespace.to_string(),
296 id.to_string(),
297 point,
298 metric,
299 )
300 .await?
301 .map_err(ClientError::Server)
302 }
303
304 pub async fn convex_hull(&self, namespace: &str) -> Result<Option<Polygon>> {
305 self.client
306 .convex_hull(self.make_context(), namespace.to_string())
307 .await?
308 .map_err(ClientError::Server)
309 }
310
311 pub async fn bounding_box(
312 &self,
313 namespace: &str,
314 ) -> Result<Option<spatio_types::bbox::BoundingBox2D>> {
315 self.client
316 .bounding_box(self.make_context(), namespace.to_string())
317 .await?
318 .map_err(ClientError::Server)
319 }
320}