1use 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
53pub 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 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 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 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 },
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 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 Err(VirtualServerError::Other(format!(
451 "Unhandled request: {:?}",
452 x
453 )));
454 },
455 };
456
457 Ok(resp)
458 }
459}