implement streaming describe rpc

Author: MarinPostma

sqld/proto/proxy.proto

@@ -4,7 +4,7 @@ package proxy;
 message Queries {
     repeated Query queries = 1;
     // Uuid
-    string clientId = 2;
+    string client_id = 2;
 }
 
 message Query {
@@ -34,10 +34,10 @@ message QueryResult {
 
 message Error {
     enum ErrorCode {
-        SQLError   = 0;
-        TxBusy     = 1;
-        TxTimeout  = 2;
-        Internal   = 3;
+        SQL_ERROR   = 0;
+        TX_BUSY     = 1;
+        TX_TIMEOUT  = 2;
+        INTERNAL   = 3;
     }
 
     ErrorCode code = 1;
@@ -71,7 +71,7 @@ message Description {
 
 message Value {
     /// bincode encoded Value
-    bytes   data = 1;
+    bytes data = 1;
 }
 
 message Row {
@@ -84,7 +84,7 @@ message Column {
 }
 
 message DisconnectMessage {
-    string clientId = 1;
+    string client_id = 1;
 }
 
 message Ack { }
@@ -150,17 +150,29 @@ message ProgramReq {
     Program pgm = 2;
 }
 
+/// Streaming exec request
 message ExecReq {
-    /// id of the request. The response will contain this id
+    /// id of the request. The response will contain this id.
     uint32 request_id = 1;
     oneof request {
-        Program         execute = 2;
-        DescribeRequest describe = 3;
+        StreamProgramReq        execute = 2;
+        StreamDescribeReq       describe = 3;
     }
 }
 
-/// streaming exec proto
+/// Describe request for the streaming protocol
+message StreamProgramReq {
+    Program pgm = 1;
+}
 
+/// descibre request for the streaming protocol
+message StreamDescribeReq {
+    string stmt = 1;
+}
+
+/// Response message for the streaming proto
+
+/// Request response types
 message Init            { }
 message BeginStep       { }
 message FinishStep      { 
@@ -173,10 +185,20 @@ message StepError       {
 message ColsDescription {
     repeated Column columns = 1;
 }
+message RowValue {
+    oneof value {
+        string text     = 1;
+        int64 integer   = 2;
+        double real     = 3;
+        bytes   blob    = 4;
+        // null if present
+        bool   null     = 5;
+    }
+}
 message BeginRows       { }
 message BeginRow        { }
 message AddRowValue     { 
-    Value val = 1;
+    RowValue val = 1;
 }
 message FinishRow       { }
 message FinishRows      { }
@@ -185,35 +207,56 @@ message Finish          {
     State state = 2;
 }
 
+/// Stream execx dexcribe response messages
+message DescribeParam {
+    optional string name = 1;
+}
 
-message Message {
-    oneof payload {
-        Description     describe_result = 1;
+message DescribeCol {
+    string              name = 1;
+    optional string     decltype = 2;
+}
 
-        Init            init = 2;
-        BeginStep       begin_step = 3;
-        FinishStep      finish_step = 4;
-        StepError       step_error = 5;
-        ColsDescription cols_description = 6;
-        BeginRows       begin_rows = 7;
-        BeginRow        begin_row = 8;
-        AddRowValue     add_row_value = 9;
-        FinishRow       finish_row = 10;
-        FinishRows      finish_rows = 11;
-        Finish          finish = 12;
+message DescribeResp {
+    repeated DescribeParam     params = 1;
+    repeated DescribeCol       cols = 2;
+    bool                        is_explain = 3;
+    bool                        is_readonly = 4;
+}
 
-        Error           error = 13;
+message RespStep {
+    oneof step {
+        Init            init = 1;
+        BeginStep       begin_step = 2;
+        FinishStep      finish_step = 3;
+        StepError       step_error = 4;
+        ColsDescription cols_description = 5;
+        BeginRows       begin_rows = 6;
+        BeginRow        begin_row = 7;
+        AddRowValue     add_row_value = 8;
+        FinishRow       finish_row = 9;
+        FinishRows      finish_rows = 10;
+        Finish          finish = 11;
     }
 }
 
+message ProgramResp {
+    repeated RespStep steps = 1;
+}
+
 message ExecResp {
     uint32  request_id = 1;
-    repeated Message messages = 2;
+    oneof response {
+        ProgramResp     program_resp = 2;
+        DescribeResp    describe_resp = 3;
+        Error           error = 4;
+    }
 }
 
 service Proxy {
   rpc StreamExec(stream ExecReq) returns (stream ExecResp) {}
 
+  // Deprecated:
   rpc Execute(ProgramReq) returns (ExecuteResults) {}
   rpc Describe(DescribeRequest) returns (DescribeResult) {}
   rpc Disconnect(DisconnectMessage) returns (Ack) {}

sqld/src/connection/libsql.rs

@@ -20,7 +20,7 @@ use crate::stats::Stats;
 use crate::Result;
 
 use super::config::DatabaseConfigStore;
-use super::program::{Cond, DescribeCol, DescribeParam, DescribeResponse, DescribeResult};
+use super::program::{Cond, DescribeCol, DescribeParam, DescribeResponse};
 use super::{MakeConnection, Program, Step, TXN_TIMEOUT};
 
 pub struct MakeLibSqlConn<W: WalHook + 'static> {
@@ -161,7 +161,9 @@ impl<W: WalHook> std::fmt::Debug for LibSqlConnection<W> {
 
 impl<W: WalHook> Clone for LibSqlConnection<W> {
     fn clone(&self) -> Self {
-        Self { inner: self.inner.clone() }
+        Self {
+            inner: self.inner.clone(),
+        }
     }
 }
 
@@ -650,7 +652,7 @@ impl<W: WalHook> Connection<W> {
         self.stats.inc_rows_written(rows_written as u64);
     }
 
-    fn describe(&self, sql: &str) -> DescribeResult {
+    fn describe(&self, sql: &str) -> crate::Result<DescribeResponse> {
         let stmt = self.conn.prepare(sql)?;
 
         let params = (1..=stmt.parameter_count())
@@ -768,7 +770,7 @@ where
         sql: String,
         auth: Authenticated,
         _replication_index: Option<FrameNo>,
-    ) -> Result<DescribeResult> {
+    ) -> Result<crate::Result<DescribeResponse>> {
         check_describe_auth(auth)?;
         let conn = self.inner.clone();
         let res = tokio::task::spawn_blocking(move || conn.lock().describe(&sql))
@@ -994,9 +996,12 @@ mod test {
         tokio::task::spawn_blocking({
             let conn = conn1.clone();
             move || {
-                let builder =
-                    Connection::run(conn.inner.clone(), Program::seq(&["COMMIT"]), TestBuilder::default())
-                        .unwrap();
+                let builder = Connection::run(
+                    conn.inner.clone(),
+                    Program::seq(&["COMMIT"]),
+                    TestBuilder::default(),
+                )
+                .unwrap();
                 assert_eq!(conn.txn_status().unwrap(), TxnStatus::Init);
                 assert!(builder.into_ret()[0].is_ok());
             }

sqld/src/connection/mod.rs

@@ -13,7 +13,7 @@ use crate::query_result_builder::{IgnoreResult, QueryResultBuilder};
 use crate::replication::FrameNo;
 use crate::Result;
 
-use self::program::{Cond, DescribeResult, Program, Step};
+use self::program::{Cond, DescribeResponse, Program, Step};
 
 pub mod config;
 pub mod dump;
@@ -111,7 +111,7 @@ pub trait Connection: Send + Sync + 'static {
         sql: String,
         auth: Authenticated,
         replication_index: Option<FrameNo>,
-    ) -> Result<DescribeResult>;
+    ) -> Result<Result<DescribeResponse>>;
 
     /// Check whether the connection is in autocommit mode.
     async fn is_autocommit(&self) -> Result<bool>;
@@ -325,7 +325,7 @@ impl<DB: Connection> Connection for TrackedConnection<DB> {
         sql: String,
         auth: Authenticated,
         replication_index: Option<FrameNo>,
-    ) -> crate::Result<DescribeResult> {
+    ) -> crate::Result<crate::Result<DescribeResponse>> {
         self.atime.store(now_millis(), Ordering::Relaxed);
         self.inner.describe(sql, auth, replication_index).await
     }
@@ -376,7 +376,7 @@ mod test {
             _sql: String,
             _auth: Authenticated,
             _replication_index: Option<FrameNo>,
-        ) -> crate::Result<DescribeResult> {
+        ) -> crate::Result<crate::Result<DescribeResponse>> {
             unreachable!()
         }
 

sqld/src/connection/program.rs

@@ -60,8 +60,6 @@ pub enum Cond {
     IsAutocommit,
 }
 
-pub type DescribeResult = crate::Result<DescribeResponse>;
-
 #[derive(Debug, Clone)]
 pub struct DescribeResponse {
     pub params: Vec<DescribeParam>,