feat: Add Tonbo Cloud + gRPC + Flight

Cargo.toml

@@ -1,4 +1,4 @@
-workspace = { members = ["parquet-lru", "tonbo_macros"] }
+workspace = { members = ["parquet-lru", "tonbo_macros", "cloud"] }
 
 [package]
 description = "An embedded persistent KV database in Rust."

cloud/Cargo.toml

@@ -0,0 +1,34 @@
+[package]
+build = "build.rs"
+name    = "tonbo-cloud"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+arrow = { version = "56.0.0", features = ["ipc"] }
+arrow-array = "56.0.0"
+arrow-flight = { version = "56.0.0", features = ["flight-sql"] }
+async-trait = "0.1"
+async-stream = "0.3"
+bytes = "1"
+flume = { version = "0.11", features = ["async"] }
+fusio = "0.4.1"
+futures = "0.3"
+futures-util = "0.3"
+futures-core = "0.3"
+prost = "0.13.0"
+prost-types = "0.13.0" 
+serde = { version = "1.0", features = ["derive"] }
+supabase-wrappers = { version = "=0.1.22", default-features = false, features = ["pg15"] }
+thiserror = "2.0.3"
+tokio = { version = "1", features = ["io-util", "macros", "rt-multi-thread"], default-features = false }
+tokio-stream = { version = "0.1", features = ["sync"] }
+tonbo = { path = ".." }
+tonic = "0.13"
+ulid = { version = "1", features = ["serde"] }
+
+[dev-dependencies]
+tempfile = "3"
+
+[build-dependencies]
+tonic-build = "0.13"

cloud/README.md

@@ -0,0 +1 @@
+Tonbo's cloud architecture. This is currently in development.

cloud/build.rs

@@ -0,0 +1,7 @@
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tonic_build::configure()
+        .build_server(true)
+        .out_dir("src/gen")
+        .compile_protos(&["src/proto/cloud.proto"], &["src/proto"])?;
+    Ok(())
+}

cloud/src/aws/aws_tonbo_svc.rs

@@ -0,0 +1,42 @@
+use std::sync::Arc;
+
+use tonic::Response;
+
+use crate::{
+    aws::{AWSTonbo, AwsTonboRPC},
+    gen::grpc,
+    ScanRequest,
+};
+
+#[derive(Clone)]
+pub struct AWSTonboSvc {
+    inner: Arc<AWSTonbo>,
+}
+
+impl AWSTonboSvc {
+    pub fn new(tonbo: Arc<AWSTonbo>) -> Self {
+        AWSTonboSvc { inner: tonbo }
+    }
+}
+
+#[tonic::async_trait]
+impl AwsTonboRPC for AWSTonboSvc {
+    async fn get_parquet_metadata(
+        &self,
+        request: tonic::Request<grpc::ScanRequest>,
+    ) -> Result<tonic::Response<grpc::ParquetMetadata>, tonic::Status> {
+        let tx = self.inner.tonbo.transaction().await;
+        let scan_request = ScanRequest::from(request.into_inner());
+
+        let (row_count, row_size) = self
+            .inner
+            .parquet_metadata(&tx, &scan_request)
+            .await
+            .unwrap();
+        let meta = grpc::ParquetMetadata {
+            row_count,
+            row_size,
+        };
+        Ok(Response::new(meta))
+    }
+}

cloud/src/aws/flight_svc.rs

@@ -0,0 +1,303 @@
+use std::sync::Arc;
+
+use arrow_array::RecordBatch;
+use arrow_flight::{
+    encode::FlightDataEncoderBuilder, error::FlightError, flight_service_server::FlightService,
+    Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
+    HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
+};
+use futures::{stream::BoxStream, StreamExt};
+use prost::Message;
+use tokio::sync::{mpsc, oneshot};
+use tokio_stream::wrappers::ReceiverStream;
+use tonbo::{
+    record::{
+        util::records_to_record_batch, ArrowArrays, ArrowArraysBuilder, DynRecord,
+        DynRecordImmutableArrays, Record, RecordRef,
+    },
+    Entry,
+};
+use tonic::{Request, Response, Status, Streaming};
+
+use crate::{aws::AWSTonbo, gen::grpc, ScanRequest, TonboCloud};
+
+#[derive(Clone)]
+pub struct TonboFlightSvc {
+    inner: Arc<AWSTonbo>,
+}
+
+impl TonboFlightSvc {
+    pub fn new(tonbo: Arc<AWSTonbo>) -> Self {
+        TonboFlightSvc { inner: tonbo }
+    }
+}
+
+#[tonic::async_trait]
+impl FlightService for TonboFlightSvc {
+    type HandshakeStream = BoxStream<'static, Result<HandshakeResponse, Status>>;
+    type ListFlightsStream = BoxStream<'static, Result<FlightInfo, Status>>;
+    type DoGetStream = BoxStream<'static, Result<FlightData, Status>>;
+    type DoPutStream = BoxStream<'static, Result<PutResult, Status>>;
+    type DoActionStream = BoxStream<'static, Result<arrow_flight::Result, Status>>;
+    type ListActionsStream = BoxStream<'static, Result<ActionType, Status>>;
+    type DoExchangeStream = BoxStream<'static, Result<FlightData, Status>>;
+
+    async fn get_schema(
+        &self,
+        _request: Request<FlightDescriptor>,
+    ) -> Result<Response<SchemaResult>, Status> {
+        Err(Status::unimplemented("Not yet implemented"))
+    }
+
+    // Scans Tonbo for record batches matching the scan predicates
+    async fn do_get(
+        &self,
+        request: Request<Ticket>,
+    ) -> Result<Response<Self::DoGetStream>, Status> {
+        // Unparse ticket to `ScanRequest`
+        let ticket = request.into_inner().ticket;
+        let scan_pb = grpc::ScanRequest::decode(ticket.as_ref()).map_err(|e| {
+            Status::invalid_argument(format!("Expected Ticket to be grpc::ScanRequest: {}", e))
+        })?;
+        let scan = ScanRequest::from(scan_pb);
+
+        // Create channel for RecordBatches + schema
+        let (rb_tx, rb_rx) = mpsc::channel::<Result<RecordBatch, FlightError>>(32);
+        let (schema_tx, schema_rx) = oneshot::channel();
+
+        let inner = Arc::clone(&self.inner);
+
+        tokio::spawn(async move {
+            let txn = inner.tonbo.transaction().await;
+
+            let mut entries = match inner.read(&txn, &scan).await {
+                Ok(s) => s,
+                Err(e) => {
+                    let _ = rb_tx
+                        .send(Err(FlightError::ExternalError(e.to_string().into())))
+                        .await;
+                    return;
+                }
+            };
+
+            let mut schema_builder: Vec<(u32, DynRecord)> = vec![];
+            // Retrieve first batch and send schema
+            let first_batch = loop {
+                match entries.next().await {
+                    Some(Ok(Entry::RecordBatch(record_batch))) => {
+                        // break after finding first batch
+                        break record_batch.record_batch().clone();
+                    }
+                    Some(Ok(Entry::Mutable(entry))) => {
+                        // Send record batch to channel
+                        if let Some(record) = entry.value() {
+                            // use dummy ts as it doesn't matter when converted to `RecordBatch`
+                            schema_builder.push((0, (*record).clone()));
+                            break records_to_record_batch(
+                                &schema_builder[0].1.schema(0),
+                                schema_builder,
+                            );
+                        }
+                    }
+                    Some(Ok(Entry::Transaction((_, record)))) => {
+                        if let Some(record) = record {
+                            // use dummy ts as it doesn't matter when converted to `RecordBatch`
+                            schema_builder.push((0, (*record).clone()));
+                            break records_to_record_batch(
+                                &schema_builder[0].1.schema(0),
+                                schema_builder,
+                            );
+                        }
+                    }
+                    // TODO: deal with projection
+                    Some(Ok(Entry::Projection((record, projection)))) => {
+                        match *record {
+                            // TODO: Make more efficient by batching build batch tranformation
+                            Entry::RecordBatch(entry) => {
+                                let schema = entry.batch_as_ref().schema();
+                                let value = entry.get();
+                                if let Some(mut value) = value {
+                                    let mut dyn_record_builder =
+                                        DynRecordImmutableArrays::builder(schema, 1);
+
+                                    // Apply projection
+                                    value.projection(&projection);
+                                    dyn_record_builder.push(entry.internal_key(), Some(value));
+                                    let dyn_record_array = dyn_record_builder.finish(None);
+                                    let record_batch = dyn_record_array.as_record_batch();
+
+                                    break record_batch.clone();
+                                }
+                            }
+                            _ => {
+                                let dyn_record = record.owned_value();
+
+                                if let Some(mut dyn_record) = dyn_record {
+                                    dyn_record.projection(&projection);
+                                    schema_builder.push((0, dyn_record.clone()));
+                                    break records_to_record_batch(
+                                        &schema_builder[0].1.schema(0),
+                                        schema_builder,
+                                    );
+                                }
+                            }
+                        }
+                    }
+                    Some(Err(e)) => {
+                        let _ = rb_tx
+                            .send(Err(FlightError::ExternalError(e.to_string().into())))
+                            .await;
+                        return;
+                    }
+                    None => {
+                        return;
+                    }
+                }
+            };
+
+            let _ = schema_tx.send(first_batch.schema());
+
+            if rb_tx.send(Ok(first_batch)).await.is_err() {
+                return;
+            }
+
+            let mut batch_builder: Vec<(u32, DynRecord)> = vec![];
+
+            while let Some(item) = entries.next().await {
+                match item {
+                    Ok(Entry::RecordBatch(record_batch)) => {
+                        // Send record batch to channel
+                        if rb_tx
+                            .send(Ok(record_batch.record_batch().clone()))
+                            .await
+                            .is_err()
+                        {
+                            return;
+                        }
+                    }
+                    Ok(Entry::Mutable(entry)) => {
+                        // Send record batch to channel
+                        if let Some(record) = entry.value() {
+                            // use dummy ts as it doesn't matter when converted to `RecordBatch`
+                            batch_builder.push((0, (*record).clone()));
+                        }
+                    }
+                    Ok(Entry::Transaction((_, record))) => {
+                        if let Some(record) = record {
+                            // use dummy ts as it doesn't matter when converted to `RecordBatch`
+                            batch_builder.push((0, (*record).clone()));
+                        }
+                    }
+                    Ok(Entry::Projection((record, projection))) => {
+                        match *record {
+                            // TODO: Make more efficient by batching build batch tranformation
+                            Entry::RecordBatch(entry) => {
+                                let schema = entry.batch_as_ref().schema();
+                                let value = entry.get();
+                                if let Some(mut value) = value {
+                                    let mut dyn_record_builder =
+                                        DynRecordImmutableArrays::builder(schema, 1);
+
+                                    // Apply projection
+                                    value.projection(&projection);
+                                    dyn_record_builder.push(entry.internal_key(), Some(value));
+                                    let dyn_record_array = dyn_record_builder.finish(None);
+                                    let record_batch = dyn_record_array.as_record_batch();
+
+                                    rb_tx.send(Ok(record_batch.clone())).await.unwrap();
+                                }
+                            }
+                            _ => {
+                                let dyn_record = record.owned_value();
+
+                                if let Some(mut dyn_record) = dyn_record {
+                                    dyn_record.projection(&projection);
+                                    batch_builder.push((0, dyn_record));
+                                }
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        let _ = rb_tx
+                            .send(Err(FlightError::ExternalError(e.to_string().into())))
+                            .await;
+                        return;
+                    }
+                }
+            }
+            if !batch_builder.is_empty() {
+                let build_batch =
+                    records_to_record_batch(&batch_builder[0].1.schema(0), batch_builder);
+                rb_tx.send(Ok(build_batch)).await.unwrap();
+            }
+        });
+
+        let schema = schema_rx
+            .await
+            .map_err(|_| Status::internal("failed to get schema"))?;
+        let rb_stream = ReceiverStream::new(rb_rx);
+
+        let fd_stream = FlightDataEncoderBuilder::new()
+            .with_schema(schema)
+            .build(rb_stream);
+
+        let out = fd_stream.map(|res| res.map_err(|e| Status::internal(e.to_string())));
+        Ok(Response::new(Box::pin(out)))
+    }
+
+    async fn handshake(
+        &self,
+        _request: Request<Streaming<HandshakeRequest>>,
+    ) -> Result<Response<Self::HandshakeStream>, Status> {
+        Err(Status::unimplemented("Not yet implemented"))
+    }
+
+    async fn list_flights(
+        &self,
+        _request: Request<Criteria>,
+    ) -> Result<Response<Self::ListFlightsStream>, Status> {
+        Err(Status::unimplemented("Not yet implemented"))
+    }
+
+    async fn get_flight_info(
+        &self,
+        _request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status> {
+        Err(Status::unimplemented("Not yet implemented"))
+    }
+
+    async fn do_put(
+        &self,
+        _request: Request<Streaming<FlightData>>,
+    ) -> Result<Response<Self::DoPutStream>, Status> {
+        Err(Status::unimplemented("Not yet implemented"))
+    }
+
+    async fn do_action(
+        &self,
+        _request: Request<Action>,
+    ) -> Result<Response<Self::DoActionStream>, Status> {
+        Err(Status::unimplemented("Not yet implemented"))
+    }
+
+    async fn list_actions(
+        &self,
+        _request: Request<Empty>,
+    ) -> Result<Response<Self::ListActionsStream>, Status> {
+        Err(Status::unimplemented("Not yet implemented"))
+    }
+
+    async fn do_exchange(
+        &self,
+        _request: Request<Streaming<FlightData>>,
+    ) -> Result<Response<Self::DoExchangeStream>, Status> {
+        Err(Status::unimplemented("Not yet implemented"))
+    }
+
+    async fn poll_flight_info(
+        &self,
+        _request: Request<FlightDescriptor>,
+    ) -> Result<Response<PollInfo>, Status> {
+        Err(Status::unimplemented("Not yet implemented"))
+    }
+}