1#![allow(refining_impl_trait)]
17
18use std::collections::HashMap;
19use std::future::Future;
20use std::pin::Pin;
21use std::sync::Arc;
22
23use crate::proto::sql::v1::{
24 cell::Kind as ProtoCellKind, Cell as ProtoCell, Column as ProtoColumn, Index as ProtoIndex,
25 IndexLayout as ProtoIndexLayout, ListValue as ProtoListValue, Null as ProtoNull,
26 QueryRequestView, QueryResponse, Row as ProtoRow, Service, ServiceServer, SubscribeRequestView,
27 SubscribeResponse, Table as ProtoTable, TablesRequestView, TablesResponse,
28};
29use bytes::Bytes;
30use connectrpc::{ConnectError, ConnectRpcService, RequestContext as Context};
31use datafusion::arrow::array::{
32 Array, ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array,
33 Decimal128Array, Decimal256Array, FixedSizeBinaryArray, Float32Array, Float64Array, Int32Array,
34 Int64Array, LargeBinaryArray, LargeListArray, LargeStringArray, ListArray, StringArray,
35 StringViewArray, TimestampMicrosecondArray, TimestampMillisecondArray,
36 TimestampNanosecondArray, TimestampSecondArray, UInt32Array, UInt64Array,
37};
38use datafusion::arrow::datatypes::{DataType, SchemaRef, TimeUnit};
39use datafusion::arrow::record_batch::RecordBatch;
40use datafusion::common::{DataFusionError, Result as DataFusionResult};
41use datafusion::datasource::MemTable;
42use datafusion::prelude::SessionContext;
43use exoware_sdk::keys::Key;
44use exoware_sdk::kv_codec::{decode_stored_row, Utf8};
45use exoware_sdk::selector::Selector;
46use exoware_sdk::stream_filter::StreamFilter;
47use exoware_sdk::{PrefixedStoreClient, StreamSubscription};
48use futures::future::BoxFuture;
49use futures::stream::Stream;
50use futures::FutureExt;
51
52use crate::builder::{projected_column_indices, ProjectedBatchBuilder};
53use crate::codec::decode_primary_key_selected;
54use crate::filter::ScanAccessPlan;
55use crate::predicate::QueryPredicate;
56use crate::schema::KvSchema;
57use crate::types::{IndexLayout, ResolvedIndexSpec, TableModel};
58
59const MAX_CONNECTRPC_BODY_BYTES: usize = 256 * 1024 * 1024;
60
61type SubscribeStream = Pin<Box<dyn Stream<Item = Result<SubscribeResponse, ConnectError>> + Send>>;
62
63#[derive(Clone)]
65struct TableStream {
66 model: Arc<TableModel>,
67 schema: SchemaRef,
68 access_plan: Arc<ScanAccessPlan>,
69 selector: Selector,
70 indexes: Arc<Vec<ResolvedIndexSpec>>,
71}
72
73impl TableStream {
74 fn new(model: Arc<TableModel>, indexes: Vec<ResolvedIndexSpec>) -> Self {
75 let projection: Option<Vec<usize>> = Some((0..model.columns.len()).collect());
76 let access_plan = Arc::new(ScanAccessPlan::new(
77 &model,
78 &projection,
79 &QueryPredicate::default(),
80 ));
81 let selector = Selector {
82 prefix: model.primary_key_prefix.as_bytes().clone(),
83 payload_regex: Utf8::from("(?s-u).*"),
84 };
85 Self {
86 schema: model.schema.clone(),
87 access_plan,
88 model,
89 selector,
90 indexes: Arc::new(indexes),
91 }
92 }
93
94 fn decode_batch(&self, entries: &[(Key, Bytes)]) -> DataFusionResult<RecordBatch> {
95 let mut builder = ProjectedBatchBuilder::from_access_plan(&self.model, &self.access_plan);
96 for (key, value) in entries {
97 if !self.model.primary_key_prefix.matches(key) {
98 continue;
99 }
100 let Some(pk_values) = decode_primary_key_selected(
101 self.model.table_prefix,
102 key,
103 &self.model,
104 &self.access_plan.required_pk_mask,
105 ) else {
106 continue;
107 };
108 let Ok(archived) = decode_stored_row(value) else {
109 continue;
110 };
111 if archived.values.len() != self.model.columns.len() {
112 continue;
113 }
114 let _ = builder.append_archived_row(&pk_values, &archived)?;
115 }
116 builder.finish(&self.schema)
117 }
118}
119
120pub struct SqlServer {
125 ctx: Arc<SessionContext>,
126 streams: HashMap<String, TableStream>,
127 table_names: Vec<String>,
130 store: PrefixedStoreClient,
131}
132
133impl SqlServer {
134 pub fn new(schema: KvSchema) -> DataFusionResult<Self> {
138 let store = schema.client().clone();
139 let mut streams = HashMap::with_capacity(schema.tables().len());
140 let mut table_names = Vec::with_capacity(schema.tables().len());
141 for (name, config) in schema.tables() {
142 let model =
143 Arc::new(TableModel::from_config(config).map_err(|e| {
144 DataFusionError::Execution(format!("invalid table config: {e}"))
145 })?);
146 let indexes = model
147 .resolve_index_specs(&config.index_specs)
148 .map_err(|e| DataFusionError::Execution(format!("invalid index specs: {e}")))?;
149 streams.insert(name.clone(), TableStream::new(model, indexes));
150 table_names.push(name.clone());
151 }
152 let ctx = SessionContext::new();
153 schema.register_all(&ctx)?;
154 Ok(Self {
155 ctx: Arc::new(ctx),
156 streams,
157 table_names,
158 store,
159 })
160 }
161
162 pub fn session(&self) -> &SessionContext {
165 &self.ctx
166 }
167
168 #[allow(clippy::result_large_err)]
169 fn stream(&self, table: &str) -> Result<&TableStream, ConnectError> {
170 self.streams
171 .get(table)
172 .ok_or_else(|| ConnectError::not_found(format!("unknown table '{table}'")))
173 }
174
175 fn describe_tables(&self) -> Vec<ProtoTable> {
176 self.table_names
177 .iter()
178 .filter_map(|name| {
179 let stream = self.streams.get(name)?;
180 let columns = stream
181 .schema
182 .fields()
183 .iter()
184 .map(|field| ProtoColumn {
185 name: field.name().clone(),
186 data_type: format!("{}", field.data_type()),
187 nullable: field.is_nullable(),
188 ..Default::default()
189 })
190 .collect();
191 let primary_key_columns = stream
192 .model
193 .primary_key_indices
194 .iter()
195 .map(|&idx| idx as u32)
196 .collect();
197 let indexes = stream
198 .indexes
199 .iter()
200 .map(|spec| {
201 let key_set: std::collections::HashSet<usize> =
202 spec.key_columns.iter().copied().collect();
203 ProtoIndex {
204 name: spec.name.clone(),
205 layout: proto_index_layout(spec.layout).into(),
206 key_columns: spec.key_columns.iter().map(|&idx| idx as u32).collect(),
207 cover_columns: spec
211 .value_column_mask
212 .iter()
213 .enumerate()
214 .filter_map(|(idx, covered)| {
215 (*covered && !key_set.contains(&idx)).then_some(idx as u32)
216 })
217 .collect(),
218 ..Default::default()
219 }
220 })
221 .collect();
222 Some(ProtoTable {
223 name: name.clone(),
224 columns,
225 primary_key_columns,
226 indexes,
227 ..Default::default()
228 })
229 })
230 .collect()
231 }
232}
233
234fn proto_index_layout(layout: IndexLayout) -> ProtoIndexLayout {
235 match layout {
236 IndexLayout::Lexicographic => ProtoIndexLayout::INDEX_LAYOUT_LEXICOGRAPHIC,
237 IndexLayout::ZOrder => ProtoIndexLayout::INDEX_LAYOUT_Z_ORDER,
238 }
239}
240
241pub fn sql_connect_stack(server: Arc<SqlServer>) -> ConnectRpcService<ServiceServer<SqlConnect>> {
244 ConnectRpcService::new(ServiceServer::new(SqlConnect::new(server)))
245 .with_limits(
246 connectrpc::Limits::default()
247 .max_request_body_size(MAX_CONNECTRPC_BODY_BYTES)
248 .max_message_size(MAX_CONNECTRPC_BODY_BYTES),
249 )
250 .with_compression(exoware_sdk::connect_compression_registry())
251}
252
253#[derive(Clone)]
255pub struct SqlConnect {
256 server: Arc<SqlServer>,
257}
258
259impl SqlConnect {
260 pub fn new(server: Arc<SqlServer>) -> Self {
261 Self { server }
262 }
263}
264
265impl Service for SqlConnect {
266 fn subscribe(
267 &self,
268 _ctx: Context,
269 request: buffa::view::OwnedView<SubscribeRequestView<'static>>,
270 ) -> impl Future<Output = connectrpc::ServiceResult<SubscribeStream>> + Send {
271 let server = self.server.clone();
272 async move {
273 let table_name = request.table.to_string();
274 let where_sql = request.where_sql.trim().to_string();
275 let since = request.since_sequence_number.filter(|seq| *seq != 0);
276 let stream = server.stream(&table_name)?.clone();
277
278 let filter = StreamFilter {
279 selectors: vec![stream.selector.clone()],
280 value_filters: vec![],
281 };
282 let sub = server
283 .store
284 .stream()
285 .subscribe(filter, since)
286 .await
287 .map_err(client_error_to_connect)?;
288
289 let output = Box::pin(BatchPredicateStream::new(
290 sub, stream, table_name, where_sql,
291 ));
292 Ok(connectrpc::Response::stream(output as SubscribeStream))
293 }
294 }
295
296 fn tables(
297 &self,
298 _ctx: Context,
299 _request: buffa::view::OwnedView<TablesRequestView<'static>>,
300 ) -> impl Future<Output = connectrpc::ServiceResult<TablesResponse>> + Send {
301 let server = self.server.clone();
302 async move {
303 connectrpc::Response::ok(TablesResponse {
304 tables: server.describe_tables(),
305 ..Default::default()
306 })
307 }
308 }
309
310 fn query(
311 &self,
312 _ctx: Context,
313 request: buffa::view::OwnedView<QueryRequestView<'static>>,
314 ) -> impl Future<Output = connectrpc::ServiceResult<QueryResponse>> + Send {
315 let server = self.server.clone();
316 async move {
317 let sql = request.sql.to_string();
318 let df = server
319 .ctx
320 .sql(&sql)
321 .await
322 .map_err(datafusion_error_to_connect)?;
323 let schema = df.schema().clone();
324 let batches = df.collect().await.map_err(datafusion_error_to_connect)?;
325 let columns: Vec<String> = schema.fields().iter().map(|f| f.name().clone()).collect();
326 let rows =
327 record_batches_to_proto_rows(&batches).map_err(datafusion_error_to_connect)?;
328 connectrpc::Response::ok(QueryResponse {
329 column: columns,
330 rows,
331 ..Default::default()
332 })
333 }
334 }
335}
336
337struct BatchPredicateStream {
338 sub: StreamSubscription,
339 state: TableStream,
340 table_name: String,
341 where_sql: String,
342 building: Option<BoxFuture<'static, Result<Option<SubscribeResponse>, ConnectError>>>,
343}
344
345impl BatchPredicateStream {
346 fn new(
347 sub: StreamSubscription,
348 state: TableStream,
349 table_name: String,
350 where_sql: String,
351 ) -> Self {
352 Self {
353 sub,
354 state,
355 table_name,
356 where_sql,
357 building: None,
358 }
359 }
360}
361
362impl Stream for BatchPredicateStream {
363 type Item = Result<SubscribeResponse, ConnectError>;
364
365 fn poll_next(
366 self: Pin<&mut Self>,
367 cx: &mut std::task::Context<'_>,
368 ) -> std::task::Poll<Option<Self::Item>> {
369 let this = self.get_mut();
370 loop {
371 if let Some(fut) = this.building.as_mut() {
372 match fut.as_mut().poll(cx) {
373 std::task::Poll::Pending => return std::task::Poll::Pending,
374 std::task::Poll::Ready(Ok(Some(resp))) => {
375 this.building = None;
376 return std::task::Poll::Ready(Some(Ok(resp)));
377 }
378 std::task::Poll::Ready(Ok(None)) => {
379 this.building = None;
380 }
381 std::task::Poll::Ready(Err(err)) => {
382 this.building = None;
383 return std::task::Poll::Ready(Some(Err(err)));
384 }
385 }
386 }
387
388 let frame = {
389 let next_fut = this.sub.next();
390 tokio::pin!(next_fut);
391 match next_fut.as_mut().poll(cx) {
392 std::task::Poll::Ready(Ok(Some(frame))) => frame,
393 std::task::Poll::Ready(Ok(None)) => return std::task::Poll::Ready(None),
394 std::task::Poll::Ready(Err(err)) => {
395 return std::task::Poll::Ready(Some(Err(client_error_to_connect(err))));
396 }
397 std::task::Poll::Pending => return std::task::Poll::Pending,
398 }
399 };
400
401 let sequence_number = frame.sequence_number;
402 let entries: Vec<(Key, Bytes)> = frame
403 .entries
404 .into_iter()
405 .map(|entry| (entry.key, entry.value))
406 .collect();
407 let state = this.state.clone();
408 let table_name = this.table_name.clone();
409 let where_sql = this.where_sql.clone();
410 this.building = Some(
411 async move {
412 evaluate_batch(state, table_name, where_sql, sequence_number, entries).await
413 }
414 .boxed(),
415 );
416 }
417 }
418}
419
420async fn evaluate_batch(
421 state: TableStream,
422 table_name: String,
423 where_sql: String,
424 sequence_number: u64,
425 entries: Vec<(Key, Bytes)>,
426) -> Result<Option<SubscribeResponse>, ConnectError> {
427 let batch = state
428 .decode_batch(&entries)
429 .map_err(datafusion_error_to_connect)?;
430 if batch.num_rows() == 0 {
431 return Ok(None);
432 }
433
434 let filtered = if where_sql.is_empty() {
435 batch
436 } else {
437 apply_where(state.schema.clone(), batch, &table_name, &where_sql)
438 .await
439 .map_err(datafusion_error_to_connect)?
440 };
441 if filtered.num_rows() == 0 {
442 return Ok(None);
443 }
444
445 let columns: Vec<String> = filtered
446 .schema()
447 .fields()
448 .iter()
449 .map(|f| f.name().clone())
450 .collect();
451 let rows = record_batches_to_proto_rows(std::slice::from_ref(&filtered))
452 .map_err(datafusion_error_to_connect)?;
453 Ok(Some(SubscribeResponse {
454 sequence_number,
455 column: columns,
456 rows,
457 ..Default::default()
458 }))
459}
460
461async fn apply_where(
462 schema: SchemaRef,
463 batch: RecordBatch,
464 table_name: &str,
465 where_sql: &str,
466) -> DataFusionResult<RecordBatch> {
467 let ctx = SessionContext::new();
468 let mem = MemTable::try_new(schema.clone(), vec![vec![batch]])?;
469 ctx.register_table(table_name, Arc::new(mem))?;
470 let sql = format!("SELECT * FROM {table_name} WHERE {where_sql}");
471 let df = ctx.sql(&sql).await?;
472 let batches = df.collect().await?;
473 if batches.is_empty() {
474 return Ok(RecordBatch::new_empty(schema));
475 }
476 datafusion::arrow::compute::concat_batches(&schema, batches.iter())
477 .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))
478}
479
480fn record_batches_to_proto_rows(batches: &[RecordBatch]) -> DataFusionResult<Vec<ProtoRow>> {
481 let mut out = Vec::with_capacity(batches.iter().map(|b| b.num_rows()).sum());
482 for batch in batches {
483 for row_idx in 0..batch.num_rows() {
484 let mut cells = Vec::with_capacity(batch.num_columns());
485 for col_idx in 0..batch.num_columns() {
486 cells.push(arrow_value_to_cell(batch.column(col_idx), row_idx)?);
487 }
488 out.push(ProtoRow {
489 cells,
490 ..Default::default()
491 });
492 }
493 }
494 Ok(out)
495}
496
497fn arrow_value_to_cell(array: &ArrayRef, row: usize) -> DataFusionResult<ProtoCell> {
498 let kind = if array.is_null(row) {
499 ProtoCellKind::NullValue(Box::<ProtoNull>::default())
500 } else {
501 arrow_value_to_kind(array, row)?
502 };
503 Ok(ProtoCell {
504 kind: Some(kind),
505 ..Default::default()
506 })
507}
508
509fn arrow_value_to_kind(array: &ArrayRef, row: usize) -> DataFusionResult<ProtoCellKind> {
510 match array.data_type() {
511 DataType::Int64 => Ok(ProtoCellKind::Int64Value(
512 array
513 .as_any()
514 .downcast_ref::<Int64Array>()
515 .unwrap()
516 .value(row),
517 )),
518 DataType::Int32 => Ok(ProtoCellKind::Int64Value(
519 array
520 .as_any()
521 .downcast_ref::<Int32Array>()
522 .unwrap()
523 .value(row) as i64,
524 )),
525 DataType::UInt64 => Ok(ProtoCellKind::Uint64Value(
526 array
527 .as_any()
528 .downcast_ref::<UInt64Array>()
529 .unwrap()
530 .value(row),
531 )),
532 DataType::UInt32 => Ok(ProtoCellKind::Uint64Value(
533 array
534 .as_any()
535 .downcast_ref::<UInt32Array>()
536 .unwrap()
537 .value(row) as u64,
538 )),
539 DataType::Float64 => Ok(ProtoCellKind::Float64Value(
540 array
541 .as_any()
542 .downcast_ref::<Float64Array>()
543 .unwrap()
544 .value(row),
545 )),
546 DataType::Float32 => Ok(ProtoCellKind::Float64Value(
547 array
548 .as_any()
549 .downcast_ref::<Float32Array>()
550 .unwrap()
551 .value(row) as f64,
552 )),
553 DataType::Boolean => Ok(ProtoCellKind::BooleanValue(
554 array
555 .as_any()
556 .downcast_ref::<BooleanArray>()
557 .unwrap()
558 .value(row),
559 )),
560 DataType::Utf8 => Ok(ProtoCellKind::Utf8Value(
561 array
562 .as_any()
563 .downcast_ref::<StringArray>()
564 .unwrap()
565 .value(row)
566 .to_string(),
567 )),
568 DataType::LargeUtf8 => Ok(ProtoCellKind::Utf8Value(
569 array
570 .as_any()
571 .downcast_ref::<LargeStringArray>()
572 .unwrap()
573 .value(row)
574 .to_string(),
575 )),
576 DataType::Utf8View => Ok(ProtoCellKind::Utf8Value(
577 array
578 .as_any()
579 .downcast_ref::<StringViewArray>()
580 .unwrap()
581 .value(row)
582 .to_string(),
583 )),
584 DataType::FixedSizeBinary(_) => {
585 Ok(ProtoCellKind::FixedSizeBinaryValue(Bytes::copy_from_slice(
586 array
587 .as_any()
588 .downcast_ref::<FixedSizeBinaryArray>()
589 .unwrap()
590 .value(row),
591 )))
592 }
593 DataType::Binary => Ok(ProtoCellKind::BinaryValue(Bytes::copy_from_slice(
594 array
595 .as_any()
596 .downcast_ref::<BinaryArray>()
597 .unwrap()
598 .value(row),
599 ))),
600 DataType::LargeBinary => Ok(ProtoCellKind::BinaryValue(Bytes::copy_from_slice(
601 array
602 .as_any()
603 .downcast_ref::<LargeBinaryArray>()
604 .unwrap()
605 .value(row),
606 ))),
607 DataType::BinaryView => Ok(ProtoCellKind::BinaryValue(Bytes::copy_from_slice(
608 array
609 .as_any()
610 .downcast_ref::<BinaryViewArray>()
611 .unwrap()
612 .value(row),
613 ))),
614 DataType::Date32 => Ok(ProtoCellKind::Date32Value(
615 array
616 .as_any()
617 .downcast_ref::<Date32Array>()
618 .unwrap()
619 .value(row),
620 )),
621 DataType::Date64 => Ok(ProtoCellKind::Date64Value(
622 array
623 .as_any()
624 .downcast_ref::<Date64Array>()
625 .unwrap()
626 .value(row),
627 )),
628 DataType::Timestamp(unit, _) => {
629 let v = match unit {
630 TimeUnit::Second => array
631 .as_any()
632 .downcast_ref::<TimestampSecondArray>()
633 .unwrap()
634 .value(row),
635 TimeUnit::Millisecond => array
636 .as_any()
637 .downcast_ref::<TimestampMillisecondArray>()
638 .unwrap()
639 .value(row),
640 TimeUnit::Microsecond => array
641 .as_any()
642 .downcast_ref::<TimestampMicrosecondArray>()
643 .unwrap()
644 .value(row),
645 TimeUnit::Nanosecond => array
646 .as_any()
647 .downcast_ref::<TimestampNanosecondArray>()
648 .unwrap()
649 .value(row),
650 };
651 Ok(ProtoCellKind::TimestampValue(v))
652 }
653 DataType::Decimal128(_, _) => {
654 let v = array
655 .as_any()
656 .downcast_ref::<Decimal128Array>()
657 .unwrap()
658 .value(row);
659 Ok(ProtoCellKind::Decimal128Value(Bytes::copy_from_slice(
660 &v.to_be_bytes(),
661 )))
662 }
663 DataType::Decimal256(_, _) => {
664 let v = array
665 .as_any()
666 .downcast_ref::<Decimal256Array>()
667 .unwrap()
668 .value(row);
669 Ok(ProtoCellKind::Decimal256Value(Bytes::copy_from_slice(
670 &v.to_be_bytes(),
671 )))
672 }
673 DataType::List(_) => {
674 let list = array.as_any().downcast_ref::<ListArray>().unwrap();
675 Ok(ProtoCellKind::ListValue(Box::new(list_array_to_proto(
676 &list.value(row),
677 )?)))
678 }
679 DataType::LargeList(_) => {
680 let list = array.as_any().downcast_ref::<LargeListArray>().unwrap();
681 Ok(ProtoCellKind::ListValue(Box::new(list_array_to_proto(
682 &list.value(row),
683 )?)))
684 }
685 other => Err(DataFusionError::NotImplemented(format!(
686 "cell conversion for arrow type {other:?}"
687 ))),
688 }
689}
690
691fn list_array_to_proto(elements: &ArrayRef) -> DataFusionResult<ProtoListValue> {
692 let mut cells = Vec::with_capacity(elements.len());
693 for idx in 0..elements.len() {
694 cells.push(arrow_value_to_cell(elements, idx)?);
695 }
696 Ok(ProtoListValue {
697 elements: cells,
698 ..Default::default()
699 })
700}
701
702fn datafusion_error_to_connect(err: DataFusionError) -> ConnectError {
703 match err {
704 DataFusionError::Plan(msg)
705 | DataFusionError::SQL(_, Some(msg))
706 | DataFusionError::Configuration(msg)
707 | DataFusionError::NotImplemented(msg) => ConnectError::invalid_argument(msg),
708 DataFusionError::SchemaError(err, _) => ConnectError::invalid_argument(err.to_string()),
709 other => ConnectError::internal(other.to_string()),
710 }
711}
712
713fn client_error_to_connect(err: exoware_sdk::ClientError) -> ConnectError {
714 if let Some(rpc) = err.rpc_error() {
715 ConnectError::new(rpc.code, rpc.message.clone().unwrap_or_default())
716 } else {
717 ConnectError::internal(err.to_string())
718 }
719}
720
721#[allow(dead_code)]
723fn _assert_projected_column_indices_visible() {
724 let _ = projected_column_indices;
725}
726
727#[cfg(test)]
728mod tests {
729 use super::*;
730
731 #[test]
734 fn binary_arrays_convert_to_binary_cells() {
735 let body: &[u8] = &[0x00, 0xFF, 0x42];
736 let arrays: Vec<ArrayRef> = vec![
737 Arc::new(BinaryArray::from_iter_values([body])),
738 Arc::new(LargeBinaryArray::from_iter_values([body])),
739 Arc::new(BinaryViewArray::from_iter_values([body])),
740 ];
741 for array in arrays {
742 let body_type = array.data_type().clone();
743 let cell = arrow_value_to_cell(&array, 0).expect("cell conversion");
744 match cell.kind {
745 Some(ProtoCellKind::BinaryValue(bytes)) => {
746 assert_eq!(bytes.as_ref(), body, "wrong bytes for {body_type:?}")
747 }
748 other => panic!("expected binary_value for {body_type:?}, got {other:?}"),
749 }
750 }
751 }
752}