1use std::collections::HashMap;
14
15use indexmap::{IndexMap, IndexSet};
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::{GroupRollupMode, ViewConfig, ViewConfigUpdate};
24use crate::proto::response::ClientResp;
25use crate::proto::{
26 ColumnPathLevel, ColumnType, GetFeaturesResp, GetHostedTablesResp, MakeTableResp, Request,
27 Response, Scalar, ServerError, TableDeleteResp, TableDescribeResp, TableMakePortResp,
28 TableMakeViewResp, TableOnDeleteResp, TableRemoveDeleteResp, TableSchemaResp, TableSizeResp,
29 ViewCollapseResp, ViewColumnPathsResp, ViewDeleteResp, ViewDescription, ViewDimensionsResp,
30 ViewExpandResp, ViewExpressionSchemaResp, ViewGetConfigResp, ViewGetMinMaxResp,
31 ViewOnDeleteResp, ViewOnRemoveResp, ViewOnUpdateResp, ViewRemoveDeleteResp,
32 ViewRemoveOnRemoveResp, ViewRemoveOnUpdateResp, ViewSchemaResp, ViewSetDepthResp,
33 ViewToArrowResp, ViewToColumnsStringResp, ViewToCsvResp, ViewToNdjsonStringResp,
34 ViewToRowsStringResp, scalar, table_describe_resp,
35};
36use crate::table::{DescribeError, Description};
37
38macro_rules! respond {
39 ($msg:ident, $name:ident { $($rest:tt)* }) => {{
40 let mut resp = BytesMut::new();
41 let resp2 = ClientResp::$name($name {
42 $($rest)*
43 });
44
45 Response {
46 msg_id: $msg.msg_id,
47 entity_id: $msg.entity_id,
48 client_resp: Some(resp2),
49 }.encode(&mut resp).map_err(VirtualServerError::EncodeError)?;
50
51 resp.freeze()
52 }};
53}
54
55pub struct VirtualServer<T: VirtualServerHandler> {
61 handler: T,
62 view_to_table: IndexMap<String, String>,
63 view_configs: IndexMap<String, ViewConfig>,
64 view_schemas: IndexMap<String, IndexMap<String, ColumnType>>,
65
66 view_tables: IndexSet<String>,
68
69 view_descriptions: IndexMap<String, Description>,
72}
73
74impl<T: VirtualServerHandler> VirtualServer<T> {
75 pub fn new(handler: T) -> Self {
77 Self {
78 handler,
79 view_configs: IndexMap::default(),
80 view_to_table: IndexMap::default(),
81 view_schemas: IndexMap::default(),
82 view_tables: IndexSet::default(),
83 view_descriptions: IndexMap::default(),
84 }
85 }
86
87 pub async fn handle_request(
92 &mut self,
93 bytes: Bytes,
94 ) -> Result<Bytes, VirtualServerError<T::Error>> {
95 let msg = Request::decode(bytes).map_err(VirtualServerError::DecodeError)?;
96 tracing::debug!(
97 "Handling request: entity_id={}, req={:?}",
98 msg.entity_id,
99 msg.client_req
100 );
101
102 match self.internal_handle_request(msg.clone()).await {
103 Ok(resp) => Ok(resp),
104 Err(err) => {
105 tracing::error!("{}", err);
106 Ok(respond!(msg, ServerError {
107 message: err.to_string(),
108 status_code: 0
109 }))
110 },
111 }
112 }
113
114 pub async fn close(&mut self) -> Result<(), VirtualServerError<T::Error>> {
117 let mut result = Ok(());
118 for view_id in self.view_configs.keys().cloned().collect::<Vec<_>>() {
119 if let Err(err) = self.handler.view_delete(&view_id).await
120 && result.is_ok()
121 {
122 result = Err(err.into());
123 }
124
125 self.view_to_table.shift_remove(&view_id);
126 self.view_configs.shift_remove(&view_id);
127 self.view_schemas.shift_remove(&view_id);
128 self.view_descriptions.shift_remove(&view_id);
129 }
130
131 result
132 }
133
134 async fn get_cached_view_schema(
135 &mut self,
136 entity_id: &str,
137 to_psp_format: bool,
138 ) -> Result<IndexMap<String, ColumnType>, VirtualServerError<T::Error>> {
139 if !self.view_schemas.contains_key(entity_id) {
140 self.view_schemas.insert(
141 entity_id.to_string(),
142 self.handler
143 .view_schema(entity_id, self.view_configs.get(entity_id).unwrap())
144 .await?,
145 );
146 }
147
148 if to_psp_format {
149 let config = self.view_configs.get(entity_id).unwrap();
153 Ok(self
154 .view_schemas
155 .get(entity_id)
156 .unwrap()
157 .iter()
158 .map(|(k, v)| {
159 let name = column_path_source(k, config)
160 .map(|(_, col)| col.to_string())
161 .unwrap_or_else(|| k.clone());
162
163 (name, *v)
164 })
165 .collect())
166 } else {
167 Ok(self.view_schemas.get(entity_id).cloned().unwrap())
168 }
169 }
170
171 async fn is_expandable(&self, view_id: &str) -> Result<bool, VirtualServerError<T::Error>> {
174 if !self.handler.get_features().await?.expand_collapse {
175 return Err(VirtualServerError::Other(
176 "This data source does not support `expand`, `collapse` or `set_depth`".to_string(),
177 ));
178 }
179
180 let config = self
181 .view_configs
182 .get(view_id)
183 .ok_or_else(|| VirtualServerError::UnknownViewId(view_id.to_string()))?;
184
185 Ok(!config.group_by.is_empty() && config.group_rollup_mode == GroupRollupMode::Rollup)
186 }
187
188 async fn internal_handle_request(
189 &mut self,
190 msg: Request,
191 ) -> Result<Bytes, VirtualServerError<T::Error>> {
192 use crate::proto::request::ClientReq::*;
193 let resp = match msg.client_req.unwrap() {
194 GetFeaturesReq(_) => {
195 let features = self.handler.get_features().await?;
196 respond!(msg, GetFeaturesResp { ..features.into() })
197 },
198 GetHostedTablesReq(_) => {
199 respond!(msg, GetHostedTablesResp {
200 table_infos: self.handler.get_hosted_tables().await?
201 })
202 },
203 TableSchemaReq(_) => {
204 respond!(msg, TableSchemaResp {
205 schema: Some(crate::proto::Schema {
206 schema: self
207 .handler
208 .table_schema(msg.entity_id.as_str())
209 .await?
210 .iter()
211 .map(|x| crate::proto::schema::KeyTypePair {
212 name: x.0.to_string(),
213 r#type: *x.1 as i32,
214 })
215 .collect()
216 })
217 })
218 },
219 TableMakePortReq(req) => {
220 respond!(msg, TableMakePortResp {
221 port_id: self.handler.table_make_port(&req).await?
222 })
223 },
224 TableMakeViewReq(req) => {
225 self.view_to_table
226 .insert(req.view_id.clone(), msg.entity_id.clone());
227
228 let mut config: ViewConfigUpdate = req.config.clone().unwrap_or_default().into();
229 if let Some(group_by) = &config.group_by
230 && let Some(duplicate) = group_by
231 .iter()
232 .enumerate()
233 .find(|(idx, col)| group_by[..*idx].contains(*col))
234 {
235 return Err(VirtualServerError::Other(format!(
236 "Duplicate column `{}` in `group_by`",
237 duplicate.1
238 )));
239 }
240
241 if let Some(windows) = &config.windows
244 && windows.values().any(|w| w.order_by.is_none())
245 && self.handler.get_features().await?.unordered
246 {
247 return Err(VirtualServerError::Other(
248 "This data store is unordered - windows require an explicit `order_by`"
249 .to_string(),
250 ));
251 }
252
253 let bytes = respond!(msg, TableMakeViewResp {
254 view_id: self
255 .handler
256 .table_make_view(msg.entity_id.as_str(), req.view_id.as_str(), &mut config)
257 .await?
258 });
259
260 self.view_configs.insert(req.view_id.clone(), config.into());
261 bytes
262 },
263 TableSizeReq(_) => {
264 respond!(msg, TableSizeResp {
265 size: self.handler.table_size(msg.entity_id.as_str()).await?
266 })
267 },
268 TableDescribeReq(req) => {
269 let config: ViewConfig = req.config.unwrap_or_default().into();
270 let verdict = self
271 .handler
272 .table_describe(msg.entity_id.as_str(), &config)
273 .await?;
274
275 let (expression_schema, expression_errors, result) = match verdict {
276 Ok(d) => (
277 d.expression_schema,
278 HashMap::new(),
279 Some(table_describe_resp::Result::View(ViewDescription {
280 schema: d
281 .view_schema
282 .into_iter()
283 .map(|(x, y)| (x, y as i32))
284 .collect(),
285 })),
286 ),
287 Err(DescribeError::Expressions {
288 expression_schema,
289 errors,
290 }) => (expression_schema, errors, None),
291 Err(DescribeError::Config(msg)) => (
292 HashMap::new(),
293 HashMap::new(),
294 Some(table_describe_resp::Result::ConfigError(msg)),
295 ),
296 };
297
298 respond!(msg, TableDescribeResp {
299 expression_schema: expression_schema
300 .into_iter()
301 .map(|(x, y)| (x, y as i32))
302 .collect(),
303 expression_errors,
304 result,
305 })
306 },
307 ViewSchemaReq(_) => {
308 respond!(msg, ViewSchemaResp {
309 schema: self
310 .get_cached_view_schema(&msg.entity_id, true)
311 .await?
312 .into_iter()
313 .map(|(x, y)| (x.to_string(), y as i32))
314 .collect()
315 })
316 },
317 ViewDimensionsReq(_) => {
318 let view_id = &msg.entity_id;
319 let table_id = self
320 .view_to_table
321 .get(view_id)
322 .ok_or_else(|| VirtualServerError::UnknownViewId(view_id.to_string()))?;
323
324 let num_table_rows = self.handler.table_size(table_id).await?;
325 let num_table_columns = self.handler.table_column_size(table_id).await? as u32;
326 let config = self.view_configs.get(view_id).unwrap();
327 let num_view_columns = self.handler.view_column_size(view_id, config).await? as u32;
328 let num_view_rows = self.handler.view_size(view_id).await?;
329 let resp = ViewDimensionsResp {
330 num_table_columns,
331 num_table_rows,
332 num_view_columns,
333 num_view_rows,
334 };
335
336 respond!(msg, ViewDimensionsResp { ..resp })
337 },
338 ViewGetConfigReq(_) => {
339 respond!(msg, ViewGetConfigResp {
340 config: Some(
341 ViewConfigUpdate::from(
342 self.view_configs.get(&msg.entity_id).unwrap().clone()
343 )
344 .into()
345 )
346 })
347 },
348 ViewExpressionSchemaReq(_) => {
349 let view_id = msg.entity_id.clone();
350 if !self.view_descriptions.contains_key(&view_id) {
351 let table_id = self
352 .view_to_table
353 .get(&view_id)
354 .cloned()
355 .ok_or_else(|| VirtualServerError::UnknownViewId(view_id.clone()))?;
356
357 let config = self.view_configs.get(&view_id).unwrap().clone();
358 let description = self
359 .handler
360 .table_describe(&table_id, &config)
361 .await?
362 .map_err(|e| VirtualServerError::Other(e.to_string()))?;
363
364 self.view_descriptions.insert(view_id.clone(), description);
365 }
366
367 let schema = self.view_descriptions[&view_id]
368 .expression_schema
369 .iter()
370 .map(|(name, ty)| (name.clone(), *ty as i32))
371 .collect();
372
373 respond!(msg, ViewExpressionSchemaResp { schema })
374 },
375 ViewColumnPathsReq(view_column_paths_req) => {
376 let config = self.view_configs.get(&msg.entity_id).unwrap();
377 let mut paths: Vec<String> = self
378 .handler
379 .view_schema(msg.entity_id.as_str(), config)
380 .await?
381 .keys()
382 .cloned()
383 .collect();
384
385 if !config.split_by.is_empty() {
386 sort_column_paths(&mut paths, config);
387 }
388
389 let start = view_column_paths_req.start_col.unwrap_or(0) as usize;
390 let end = view_column_paths_req
391 .end_col
392 .map_or(paths.len(), |x| x as usize);
393
394 let paths = paths
395 .into_iter()
396 .skip(start)
397 .take(end.saturating_sub(start))
398 .collect::<Vec<_>>();
399
400 let depth = config.split_by.len() + 1;
401 let mut area = vec![
402 ColumnPathLevel {
403 values: Vec::with_capacity(paths.len())
404 };
405 depth
406 ];
407
408 for path in paths {
409 let mut split = path.splitn(depth, '|');
410 let mut levels = Vec::with_capacity(depth);
411 for _ in 0..depth {
412 match split.next() {
413 Some(x) => levels.push(Some(x.to_owned())),
414 None => levels.push(None),
415 }
416 }
417
418 let name = levels
419 .iter()
420 .rposition(|x| x.is_some())
421 .map(|i| levels[i].take().unwrap())
422 .unwrap_or_default();
423
424 for (level, value) in area.iter_mut().zip(levels).take(depth - 1) {
425 level.values.push(Scalar {
426 scalar: value.map(scalar::Scalar::String),
427 });
428 }
429
430 area[depth - 1].values.push(Scalar {
431 scalar: Some(scalar::Scalar::String(name)),
432 });
433 }
434
435 respond!(msg, ViewColumnPathsResp { area })
436 },
437 ViewToArrowReq(view_to_arrow_req) => {
438 let viewport = view_to_arrow_req.viewport.unwrap();
439 let schema = self.get_cached_view_schema(&msg.entity_id, false).await?;
440 let config = self.view_configs.get(&msg.entity_id).unwrap();
441 let mut cols = self
442 .handler
443 .view_get_data(msg.entity_id.as_str(), config, &schema, &viewport)
444 .await?;
445
446 let arrow = cols
447 .render_to_arrow_ipc()
448 .map_err(|e| VirtualServerError::Other(e.to_string()))?;
449
450 respond!(msg, ViewToArrowResp { arrow })
451 },
452 ViewToCsvReq(view_to_csv_req) => {
453 let viewport = view_to_csv_req.viewport.unwrap();
454 let schema = self.get_cached_view_schema(&msg.entity_id, false).await?;
455 let config = self.view_configs.get(&msg.entity_id).unwrap();
456 let mut cols = self
457 .handler
458 .view_get_data(msg.entity_id.as_str(), config, &schema, &viewport)
459 .await?;
460
461 let rows = cols.render_to_rows(RowPathStyle::PerLevel);
462 let mut csv = String::new();
463 if let Some(first_row) = rows.first() {
464 let headers: Vec<&str> = first_row.keys().map(|k| k.as_str()).collect();
465 csv.push_str(&headers.join(","));
466 csv.push('\n');
467 }
468
469 for row in &rows {
470 let values: Vec<String> = row
471 .values()
472 .map(|cell| serde_json::to_string(cell).unwrap_or_default())
473 .collect();
474 csv.push_str(&values.join(","));
475 csv.push('\n');
476 }
477
478 respond!(msg, ViewToCsvResp { csv })
479 },
480 ViewToNdjsonStringReq(view_to_ndjson_req) => {
481 let viewport = view_to_ndjson_req.viewport.unwrap();
482 let schema = self.get_cached_view_schema(&msg.entity_id, false).await?;
483 let config = self.view_configs.get(&msg.entity_id).unwrap();
484 let mut cols = self
485 .handler
486 .view_get_data(msg.entity_id.as_str(), config, &schema, &viewport)
487 .await?;
488
489 let rows = cols.render_to_rows(RowPathStyle::Sidecar);
490 let ndjson_string = rows
491 .iter()
492 .map(serde_json::to_string)
493 .collect::<Result<Vec<_>, _>>()
494 .map_err(|e| VirtualServerError::InvalidJSON(std::sync::Arc::new(e)))?
495 .join("\n");
496
497 respond!(msg, ViewToNdjsonStringResp { ndjson_string })
498 },
499 ViewToRowsStringReq(view_to_rows_string_req) => {
500 let viewport = view_to_rows_string_req.viewport.unwrap();
501 let schema = self.get_cached_view_schema(&msg.entity_id, false).await?;
502 let config = self.view_configs.get(&msg.entity_id).unwrap();
503 let mut cols = self
504 .handler
505 .view_get_data(msg.entity_id.as_str(), config, &schema, &viewport)
506 .await?;
507
508 let rows = cols.render_to_rows(RowPathStyle::Sidecar);
509 let json_string = serde_json::to_string(&rows)
510 .map_err(|e| VirtualServerError::InvalidJSON(std::sync::Arc::new(e)))?;
511
512 respond!(msg, ViewToRowsStringResp { json_string })
513 },
514 ViewToColumnsStringReq(view_to_columns_string_req) => {
515 let viewport = view_to_columns_string_req.viewport.unwrap();
516 let schema = self.get_cached_view_schema(&msg.entity_id, false).await?;
517 let config = self.view_configs.get(&msg.entity_id).unwrap();
518 let mut cols = self
519 .handler
520 .view_get_data(msg.entity_id.as_str(), config, &schema, &viewport)
521 .await?;
522
523 let json_string = cols
524 .render_to_columns_json(
525 RowPathStyle::Sidecar,
526 view_to_columns_string_req.id.unwrap_or_default(),
527 )
528 .map_err(|e| VirtualServerError::Other(e.to_string()))?;
529
530 respond!(msg, ViewToColumnsStringResp { json_string })
531 },
532 ViewDeleteReq(_) => {
533 self.handler.view_delete(msg.entity_id.as_str()).await?;
534 self.view_to_table.shift_remove(&msg.entity_id);
535 self.view_configs.shift_remove(&msg.entity_id);
536 self.view_descriptions.shift_remove(&msg.entity_id);
537 respond!(msg, ViewDeleteResp {})
538 },
539 MakeTableReq(req) => {
540 let data = req.data.as_ref().unwrap();
541 if let Some(crate::proto::make_table_data::Data::FromView(view_id)) = &data.data {
542 if !self.handler.get_features().await?.view_derivations {
543 return Err(VirtualServerError::Other(
544 "This data source cannot derive a Table from a View".to_string(),
545 ));
546 }
547
548 let options = req.options.clone().unwrap_or_default();
549 if options.make_table_type.is_some() {
550 return Err(VirtualServerError::Other(
551 "`index` and `limit` cannot be set on a Table made from a View"
552 .to_string(),
553 ));
554 }
555
556 let config = self.view_configs.get(view_id).ok_or_else(|| {
557 VirtualServerError::Other(format!("View \"{}\" does not exist", view_id))
558 })?;
559
560 let schema = options
561 .view_schema
562 .map(|schema| {
563 schema
564 .schema
565 .into_iter()
566 .map(|x| Ok((x.name, ColumnType::try_from(x.r#type)?)))
567 .collect::<Result<IndexMap<_, _>, prost::DecodeError>>()
568 })
569 .transpose()
570 .map_err(|e| VirtualServerError::Other(e.to_string()))?;
571
572 self.handler
573 .view_make_table(view_id, &msg.entity_id, config, schema.as_ref())
574 .await?;
575
576 self.view_tables.insert(msg.entity_id.clone());
577 } else {
578 self.handler.make_table(&msg.entity_id, data).await?;
579 }
580
581 respond!(msg, MakeTableResp {})
582 },
583 ViewGetMinMaxReq(req) => {
584 let config = self.view_configs.get(&msg.entity_id).unwrap();
585 let (min, max) = self
586 .handler
587 .view_get_min_max(&msg.entity_id, &req.column_name, config)
588 .await?;
589 respond!(msg, ViewGetMinMaxResp {
590 min: Some(min.into()),
591 max: Some(max.into()),
592 })
593 },
594
595 ViewExpandReq(req) => {
596 let num_changed = if self.is_expandable(&msg.entity_id).await? {
597 let config = &self.view_configs[&msg.entity_id];
598 self.handler
599 .view_expand(&msg.entity_id, config, req.row_index)
600 .await?
601 } else {
602 0
603 };
604
605 respond!(msg, ViewExpandResp { num_changed })
606 },
607 ViewCollapseReq(req) => {
608 let num_changed = if self.is_expandable(&msg.entity_id).await? {
609 let config = &self.view_configs[&msg.entity_id];
610 self.handler
611 .view_collapse(&msg.entity_id, config, req.row_index)
612 .await?
613 } else {
614 0
615 };
616
617 respond!(msg, ViewCollapseResp { num_changed })
618 },
619 ViewSetDepthReq(req) => {
620 if self.is_expandable(&msg.entity_id).await? {
621 let config = &self.view_configs[&msg.entity_id];
622 if req.depth as usize <= config.group_by.len() {
623 self.handler
624 .view_set_depth(&msg.entity_id, config, req.depth)
625 .await?;
626 }
627 }
628
629 respond!(msg, ViewSetDepthResp {})
630 },
631
632 TableDeleteReq(_) if self.view_tables.contains(&msg.entity_id) => {
633 self.handler.view_delete(msg.entity_id.as_str()).await?;
634 self.view_tables.shift_remove(&msg.entity_id);
635 respond!(msg, TableDeleteResp {})
636 },
637
638 TableOnDeleteReq(_) => {
640 respond!(msg, TableOnDeleteResp {})
641 },
642 ViewOnRemoveReq(_) => {
643 respond!(msg, ViewOnRemoveResp {
644 indices: None,
645 port_id: 0
646 })
647 },
648 ViewRemoveOnRemoveReq(_) => {
649 respond!(msg, ViewRemoveOnRemoveResp {})
650 },
651 ViewOnUpdateReq(_) => {
652 respond!(msg, ViewOnUpdateResp {
653 delta: None,
654 port_id: 0,
655 reset: false
656 })
657 },
658 ViewOnDeleteReq(_) => {
659 respond!(msg, ViewOnDeleteResp {})
660 },
661 ViewRemoveOnUpdateReq(_) => {
662 respond!(msg, ViewRemoveOnUpdateResp {})
663 },
664 TableRemoveDeleteReq(_) => {
665 respond!(msg, TableRemoveDeleteResp {})
666 },
667 ViewRemoveDeleteReq(_) => {
668 respond!(msg, ViewRemoveDeleteResp {})
669 },
670 x => {
671 return Err(VirtualServerError::Other(format!(
673 "Unhandled request: {:?}",
674 x
675 )));
676 },
677 };
678
679 Ok(resp)
680 }
681}