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, BooleanArray, Date32Array, Date64Array, Decimal128Array, Decimal256Array,
33 FixedSizeBinaryArray, Float32Array, Float64Array, Int32Array, Int64Array, LargeListArray,
34 LargeStringArray, ListArray, StringArray, StringViewArray, TimestampMicrosecondArray,
35 TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UInt32Array,
36 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::Date32 => Ok(ProtoCellKind::Date32Value(
594 array
595 .as_any()
596 .downcast_ref::<Date32Array>()
597 .unwrap()
598 .value(row),
599 )),
600 DataType::Date64 => Ok(ProtoCellKind::Date64Value(
601 array
602 .as_any()
603 .downcast_ref::<Date64Array>()
604 .unwrap()
605 .value(row),
606 )),
607 DataType::Timestamp(unit, _) => {
608 let v = match unit {
609 TimeUnit::Second => array
610 .as_any()
611 .downcast_ref::<TimestampSecondArray>()
612 .unwrap()
613 .value(row),
614 TimeUnit::Millisecond => array
615 .as_any()
616 .downcast_ref::<TimestampMillisecondArray>()
617 .unwrap()
618 .value(row),
619 TimeUnit::Microsecond => array
620 .as_any()
621 .downcast_ref::<TimestampMicrosecondArray>()
622 .unwrap()
623 .value(row),
624 TimeUnit::Nanosecond => array
625 .as_any()
626 .downcast_ref::<TimestampNanosecondArray>()
627 .unwrap()
628 .value(row),
629 };
630 Ok(ProtoCellKind::TimestampValue(v))
631 }
632 DataType::Decimal128(_, _) => {
633 let v = array
634 .as_any()
635 .downcast_ref::<Decimal128Array>()
636 .unwrap()
637 .value(row);
638 Ok(ProtoCellKind::Decimal128Value(Bytes::copy_from_slice(
639 &v.to_be_bytes(),
640 )))
641 }
642 DataType::Decimal256(_, _) => {
643 let v = array
644 .as_any()
645 .downcast_ref::<Decimal256Array>()
646 .unwrap()
647 .value(row);
648 Ok(ProtoCellKind::Decimal256Value(Bytes::copy_from_slice(
649 &v.to_be_bytes(),
650 )))
651 }
652 DataType::List(_) => {
653 let list = array.as_any().downcast_ref::<ListArray>().unwrap();
654 Ok(ProtoCellKind::ListValue(Box::new(list_array_to_proto(
655 &list.value(row),
656 )?)))
657 }
658 DataType::LargeList(_) => {
659 let list = array.as_any().downcast_ref::<LargeListArray>().unwrap();
660 Ok(ProtoCellKind::ListValue(Box::new(list_array_to_proto(
661 &list.value(row),
662 )?)))
663 }
664 other => Err(DataFusionError::NotImplemented(format!(
665 "cell conversion for arrow type {other:?}"
666 ))),
667 }
668}
669
670fn list_array_to_proto(elements: &ArrayRef) -> DataFusionResult<ProtoListValue> {
671 let mut cells = Vec::with_capacity(elements.len());
672 for idx in 0..elements.len() {
673 cells.push(arrow_value_to_cell(elements, idx)?);
674 }
675 Ok(ProtoListValue {
676 elements: cells,
677 ..Default::default()
678 })
679}
680
681fn datafusion_error_to_connect(err: DataFusionError) -> ConnectError {
682 match err {
683 DataFusionError::Plan(msg)
684 | DataFusionError::SQL(_, Some(msg))
685 | DataFusionError::Configuration(msg)
686 | DataFusionError::NotImplemented(msg) => ConnectError::invalid_argument(msg),
687 DataFusionError::SchemaError(err, _) => ConnectError::invalid_argument(err.to_string()),
688 other => ConnectError::internal(other.to_string()),
689 }
690}
691
692fn client_error_to_connect(err: exoware_sdk::ClientError) -> ConnectError {
693 if let Some(rpc) = err.rpc_error() {
694 ConnectError::new(rpc.code, rpc.message.clone().unwrap_or_default())
695 } else {
696 ConnectError::internal(err.to_string())
697 }
698}
699
700#[allow(dead_code)]
702fn _assert_projected_column_indices_visible() {
703 let _ = projected_column_indices;
704}