Skip to main content

perspective_client/virtual_server/
server.rs

1// ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
2// ┃ ██████ ██████ ██████       █      █      █      █      █ █▄  ▀███ █       ┃
3// ┃ ▄▄▄▄▄█ █▄▄▄▄▄ ▄▄▄▄▄█  ▀▀▀▀▀█▀▀▀▀▀ █ ▀▀▀▀▀█ ████████▌▐███ ███▄  ▀█ █ ▀▀▀▀▀ ┃
4// ┃ █▀▀▀▀▀ █▀▀▀▀▀ █▀██▀▀ ▄▄▄▄▄ █ ▄▄▄▄▄█ ▄▄▄▄▄█ ████████▌▐███ █████▄   █ ▄▄▄▄▄ ┃
5// ┃ █      ██████ █  ▀█▄       █ ██████      █      ███▌▐███ ███████▄ █       ┃
6// ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫
7// ┃ Copyright (c) 2017, the Perspective Authors.                              ┃
8// ┃ ╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌ ┃
9// ┃ This file is part of the Perspective library, distributed under the terms ┃
10// ┃ of the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). ┃
11// ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
12
13use std::collections::HashMap;
14
15use indexmap::IndexMap;
16use prost::Message as ProstMessage;
17use prost::bytes::{Bytes, BytesMut};
18
19use super::data::RowPathStyle;
20use super::error::VirtualServerError;
21use super::generic_sql_model::{column_path_source, sort_column_paths};
22use super::handler::VirtualServerHandler;
23use crate::config::{ViewConfig, ViewConfigUpdate};
24use crate::proto::response::ClientResp;
25use crate::proto::table_validate_expr_resp::ExprValidationError;
26use crate::proto::{
27    ColumnType, GetFeaturesResp, GetHostedTablesResp, MakeTableResp, Request, Response,
28    ServerError, TableMakePortResp, TableMakeViewResp, TableOnDeleteResp, TableRemoveDeleteResp,
29    TableSchemaResp, TableSizeResp, TableValidateExprResp, ViewColumnPathsResp, ViewDeleteResp,
30    ViewDimensionsResp, ViewExpressionSchemaResp, ViewGetConfigResp, ViewGetMinMaxResp,
31    ViewOnDeleteResp, ViewOnUpdateResp, ViewRemoveDeleteResp, ViewRemoveOnUpdateResp,
32    ViewSchemaResp, ViewToArrowResp, ViewToColumnsStringResp, ViewToCsvResp,
33    ViewToNdjsonStringResp, ViewToRowsStringResp,
34};
35
36macro_rules! respond {
37    ($msg:ident, $name:ident { $($rest:tt)* }) => {{
38        let mut resp = BytesMut::new();
39        let resp2 = ClientResp::$name($name {
40            $($rest)*
41        });
42
43        Response {
44            msg_id: $msg.msg_id,
45            entity_id: $msg.entity_id,
46            client_resp: Some(resp2),
47        }.encode(&mut resp).map_err(VirtualServerError::EncodeError)?;
48
49        resp.freeze()
50    }};
51}
52
53/// A virtual server that processes Perspective protocol messages.
54///
55/// `VirtualServer` acts as a bridge between the Perspective protocol and a
56/// custom data backend. It handles protocol decoding/encoding and delegates
57/// actual data operations to the provided [`VirtualServerHandler`].
58pub struct VirtualServer<T: VirtualServerHandler> {
59    handler: T,
60    view_to_table: IndexMap<String, String>,
61    view_configs: IndexMap<String, ViewConfig>,
62    view_schemas: IndexMap<String, IndexMap<String, ColumnType>>,
63}
64
65impl<T: VirtualServerHandler> VirtualServer<T> {
66    /// Creates a new virtual server with the given handler.
67    pub fn new(handler: T) -> Self {
68        Self {
69            handler,
70            view_configs: IndexMap::default(),
71            view_to_table: IndexMap::default(),
72            view_schemas: IndexMap::default(),
73        }
74    }
75
76    /// Processes a Perspective protocol request and returns the response.
77    ///
78    /// Decodes the incoming protobuf message, dispatches to the appropriate
79    /// handler method, and encodes the response.
80    pub async fn handle_request(
81        &mut self,
82        bytes: Bytes,
83    ) -> Result<Bytes, VirtualServerError<T::Error>> {
84        let msg = Request::decode(bytes).map_err(VirtualServerError::DecodeError)?;
85        tracing::debug!(
86            "Handling request: entity_id={}, req={:?}",
87            msg.entity_id,
88            msg.client_req
89        );
90
91        match self.internal_handle_request(msg.clone()).await {
92            Ok(resp) => Ok(resp),
93            Err(err) => {
94                tracing::error!("{}", err);
95                Ok(respond!(msg, ServerError {
96                    message: err.to_string(),
97                    status_code: 0
98                }))
99            },
100        }
101    }
102
103    async fn get_cached_view_schema(
104        &mut self,
105        entity_id: &str,
106        to_psp_format: bool,
107    ) -> Result<IndexMap<String, ColumnType>, VirtualServerError<T::Error>> {
108        if !self.view_schemas.contains_key(entity_id) {
109            self.view_schemas.insert(
110                entity_id.to_string(),
111                self.handler
112                    .view_schema(entity_id, self.view_configs.get(entity_id).unwrap())
113                    .await?,
114            );
115        }
116
117        if to_psp_format {
118            // `view.schema()` is keyed by *source* column name, matching the
119            // native engine, while the cached schema is keyed by the view's
120            // actual (possibly pivoted-path) SQL column names.
121            let config = self.view_configs.get(entity_id).unwrap();
122            Ok(self
123                .view_schemas
124                .get(entity_id)
125                .unwrap()
126                .iter()
127                .map(|(k, v)| {
128                    let name = column_path_source(k, config)
129                        .map(|(_, col)| col.to_string())
130                        .unwrap_or_else(|| k.clone());
131
132                    (name, *v)
133                })
134                .collect())
135        } else {
136            Ok(self.view_schemas.get(entity_id).cloned().unwrap())
137        }
138    }
139
140    async fn internal_handle_request(
141        &mut self,
142        msg: Request,
143    ) -> Result<Bytes, VirtualServerError<T::Error>> {
144        use crate::proto::request::ClientReq::*;
145        let resp = match msg.client_req.unwrap() {
146            GetFeaturesReq(_) => {
147                let features = self.handler.get_features().await?;
148                respond!(msg, GetFeaturesResp { ..features.into() })
149            },
150            GetHostedTablesReq(_) => {
151                respond!(msg, GetHostedTablesResp {
152                    table_infos: self.handler.get_hosted_tables().await?
153                })
154            },
155            TableSchemaReq(_) => {
156                respond!(msg, TableSchemaResp {
157                    schema: Some(crate::proto::Schema {
158                        schema: self
159                            .handler
160                            .table_schema(msg.entity_id.as_str())
161                            .await?
162                            .iter()
163                            .map(|x| crate::proto::schema::KeyTypePair {
164                                name: x.0.to_string(),
165                                r#type: *x.1 as i32,
166                            })
167                            .collect()
168                    })
169                })
170            },
171            TableMakePortReq(req) => {
172                respond!(msg, TableMakePortResp {
173                    port_id: self.handler.table_make_port(&req).await?
174                })
175            },
176            TableMakeViewReq(req) => {
177                self.view_to_table
178                    .insert(req.view_id.clone(), msg.entity_id.clone());
179
180                let mut config: ViewConfigUpdate = req.config.clone().unwrap_or_default().into();
181                let bytes = respond!(msg, TableMakeViewResp {
182                    view_id: self
183                        .handler
184                        .table_make_view(msg.entity_id.as_str(), req.view_id.as_str(), &mut config)
185                        .await?
186                });
187
188                self.view_configs.insert(req.view_id.clone(), config.into());
189                bytes
190            },
191            TableSizeReq(_) => {
192                respond!(msg, TableSizeResp {
193                    size: self.handler.table_size(msg.entity_id.as_str()).await?
194                })
195            },
196            TableValidateExprReq(req) => {
197                let mut expression_schema = HashMap::<String, i32>::default();
198                let mut expression_alias = HashMap::<String, String>::default();
199                let mut errors = HashMap::<String, ExprValidationError>::default();
200                for (name, ex) in req.column_to_expr.iter() {
201                    let _ = expression_alias.insert(name.clone(), ex.clone());
202                    match self
203                        .handler
204                        .table_validate_expression(&msg.entity_id, ex.as_str())
205                        .await
206                    {
207                        Ok(dtype) => {
208                            let _ = expression_schema.insert(name.clone(), dtype as i32);
209                        },
210                        Err(e) => {
211                            let _ = errors.insert(name.clone(), ExprValidationError {
212                                error_message: format!("{}", e),
213                                line: 0,
214                                column: 0,
215                            });
216                        },
217                    }
218                }
219
220                respond!(msg, TableValidateExprResp {
221                    expression_schema,
222                    errors,
223                    expression_alias,
224                })
225            },
226            ViewSchemaReq(_) => {
227                respond!(msg, ViewSchemaResp {
228                    schema: self
229                        .get_cached_view_schema(&msg.entity_id, true)
230                        .await?
231                        .into_iter()
232                        .map(|(x, y)| (x.to_string(), y as i32))
233                        .collect()
234                })
235            },
236            ViewDimensionsReq(_) => {
237                let view_id = &msg.entity_id;
238                let table_id = self
239                    .view_to_table
240                    .get(view_id)
241                    .ok_or_else(|| VirtualServerError::UnknownViewId(view_id.to_string()))?;
242
243                let num_table_rows = self.handler.table_size(table_id).await?;
244                let num_table_columns = self.handler.table_column_size(table_id).await? as u32;
245                let config = self.view_configs.get(view_id).unwrap();
246                let num_view_columns = self.handler.view_column_size(view_id, config).await? as u32;
247                let num_view_rows = self.handler.view_size(view_id).await?;
248                let resp = ViewDimensionsResp {
249                    num_table_columns,
250                    num_table_rows,
251                    num_view_columns,
252                    num_view_rows,
253                };
254
255                respond!(msg, ViewDimensionsResp { ..resp })
256            },
257            ViewGetConfigReq(_) => {
258                respond!(msg, ViewGetConfigResp {
259                    config: Some(
260                        ViewConfigUpdate::from(
261                            self.view_configs.get(&msg.entity_id).unwrap().clone()
262                        )
263                        .into()
264                    )
265                })
266            },
267            ViewExpressionSchemaReq(_) => {
268                let mut schema = HashMap::<String, i32>::default();
269                let table_id = self.view_to_table.get(&msg.entity_id);
270                for (name, ex) in self
271                    .view_configs
272                    .get(&msg.entity_id)
273                    .unwrap()
274                    .expressions
275                    .iter()
276                {
277                    match self
278                        .handler
279                        .table_validate_expression(table_id.unwrap(), ex.as_str())
280                        .await
281                    {
282                        Ok(dtype) => {
283                            let _ = schema.insert(name.clone(), dtype as i32);
284                        },
285                        Err(_e) => {
286                            // TODO: handle error
287                        },
288                    }
289                }
290
291                let resp = ViewExpressionSchemaResp { schema };
292                respond!(msg, ViewExpressionSchemaResp { ..resp })
293            },
294            ViewColumnPathsReq(_) => {
295                let config = self.view_configs.get(&msg.entity_id).unwrap();
296                let mut paths: Vec<String> = self
297                    .handler
298                    .view_schema(msg.entity_id.as_str(), config)
299                    .await?
300                    .keys()
301                    .cloned()
302                    .collect();
303
304                if !config.split_by.is_empty() {
305                    sort_column_paths(&mut paths, config);
306                }
307
308                respond!(msg, ViewColumnPathsResp { paths })
309            },
310            ViewToArrowReq(view_to_arrow_req) => {
311                let viewport = view_to_arrow_req.viewport.unwrap();
312                let schema = self.get_cached_view_schema(&msg.entity_id, false).await?;
313                let config = self.view_configs.get(&msg.entity_id).unwrap();
314                let mut cols = self
315                    .handler
316                    .view_get_data(msg.entity_id.as_str(), config, &schema, &viewport)
317                    .await?;
318
319                let arrow = cols
320                    .render_to_arrow_ipc()
321                    .map_err(|e| VirtualServerError::Other(e.to_string()))?;
322
323                respond!(msg, ViewToArrowResp { arrow })
324            },
325            ViewToCsvReq(view_to_csv_req) => {
326                let viewport = view_to_csv_req.viewport.unwrap();
327                let schema = self.get_cached_view_schema(&msg.entity_id, false).await?;
328                let config = self.view_configs.get(&msg.entity_id).unwrap();
329                let mut cols = self
330                    .handler
331                    .view_get_data(msg.entity_id.as_str(), config, &schema, &viewport)
332                    .await?;
333
334                let rows = cols.render_to_rows(RowPathStyle::PerLevel);
335                let mut csv = String::new();
336                if let Some(first_row) = rows.first() {
337                    let headers: Vec<&str> = first_row.keys().map(|k| k.as_str()).collect();
338                    csv.push_str(&headers.join(","));
339                    csv.push('\n');
340                }
341
342                for row in &rows {
343                    let values: Vec<String> = row
344                        .values()
345                        .map(|cell| serde_json::to_string(cell).unwrap_or_default())
346                        .collect();
347                    csv.push_str(&values.join(","));
348                    csv.push('\n');
349                }
350
351                respond!(msg, ViewToCsvResp { csv })
352            },
353            ViewToNdjsonStringReq(view_to_ndjson_req) => {
354                let viewport = view_to_ndjson_req.viewport.unwrap();
355                let schema = self.get_cached_view_schema(&msg.entity_id, false).await?;
356                let config = self.view_configs.get(&msg.entity_id).unwrap();
357                let mut cols = self
358                    .handler
359                    .view_get_data(msg.entity_id.as_str(), config, &schema, &viewport)
360                    .await?;
361
362                let rows = cols.render_to_rows(RowPathStyle::PerLevel);
363                let ndjson_string = rows
364                    .iter()
365                    .map(serde_json::to_string)
366                    .collect::<Result<Vec<_>, _>>()
367                    .map_err(|e| VirtualServerError::InvalidJSON(std::sync::Arc::new(e)))?
368                    .join("\n");
369
370                respond!(msg, ViewToNdjsonStringResp { ndjson_string })
371            },
372            ViewToRowsStringReq(view_to_rows_string_req) => {
373                let viewport = view_to_rows_string_req.viewport.unwrap();
374                let schema = self.get_cached_view_schema(&msg.entity_id, false).await?;
375                let config = self.view_configs.get(&msg.entity_id).unwrap();
376                let mut cols = self
377                    .handler
378                    .view_get_data(msg.entity_id.as_str(), config, &schema, &viewport)
379                    .await?;
380
381                let rows = cols.render_to_rows(RowPathStyle::Sidecar);
382                let json_string = serde_json::to_string(&rows)
383                    .map_err(|e| VirtualServerError::InvalidJSON(std::sync::Arc::new(e)))?;
384
385                respond!(msg, ViewToRowsStringResp { json_string })
386            },
387            ViewToColumnsStringReq(view_to_columns_string_req) => {
388                let viewport = view_to_columns_string_req.viewport.unwrap();
389                let schema = self.get_cached_view_schema(&msg.entity_id, false).await?;
390                let config = self.view_configs.get(&msg.entity_id).unwrap();
391                let mut cols = self
392                    .handler
393                    .view_get_data(msg.entity_id.as_str(), config, &schema, &viewport)
394                    .await?;
395
396                let json_string = cols
397                    .render_to_columns_json(RowPathStyle::Sidecar)
398                    .map_err(|e| VirtualServerError::Other(e.to_string()))?;
399
400                respond!(msg, ViewToColumnsStringResp { json_string })
401            },
402            ViewDeleteReq(_) => {
403                self.handler.view_delete(msg.entity_id.as_str()).await?;
404                self.view_to_table.shift_remove(&msg.entity_id);
405                self.view_configs.shift_remove(&msg.entity_id);
406                respond!(msg, ViewDeleteResp {})
407            },
408            MakeTableReq(req) => {
409                self.handler
410                    .make_table(&msg.entity_id, req.data.as_ref().unwrap())
411                    .await?;
412                respond!(msg, MakeTableResp {})
413            },
414            ViewGetMinMaxReq(req) => {
415                let config = self.view_configs.get(&msg.entity_id).unwrap();
416                let (min, max) = self
417                    .handler
418                    .view_get_min_max(&msg.entity_id, &req.column_name, config)
419                    .await?;
420                respond!(msg, ViewGetMinMaxResp {
421                    min: Some(min.into()),
422                    max: Some(max.into()),
423                })
424            },
425
426            // Stub implementations for callback/update requests that VirtualServer doesn't support
427            TableOnDeleteReq(_) => {
428                respond!(msg, TableOnDeleteResp {})
429            },
430            ViewOnUpdateReq(_) => {
431                respond!(msg, ViewOnUpdateResp {
432                    delta: None,
433                    port_id: 0
434                })
435            },
436            ViewOnDeleteReq(_) => {
437                respond!(msg, ViewOnDeleteResp {})
438            },
439            ViewRemoveOnUpdateReq(_) => {
440                respond!(msg, ViewRemoveOnUpdateResp {})
441            },
442            TableRemoveDeleteReq(_) => {
443                respond!(msg, TableRemoveDeleteResp {})
444            },
445            ViewRemoveDeleteReq(_) => {
446                respond!(msg, ViewRemoveDeleteResp {})
447            },
448            x => {
449                // Return an error response instead of empty bytes
450                return Err(VirtualServerError::Other(format!(
451                    "Unhandled request: {:?}",
452                    x
453                )));
454            },
455        };
456
457        Ok(resp)
458    }
459}