spatio_client/transport/
rpc.rs

1//! tarpc transport for Spatio client
2//!
3//! This is the default high-performance RPC client.
4
5#![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}