From e4619be89a095773dde53e012eb3773f94f0b9f7 Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@abcum.com>
Date: Wed, 11 May 2022 22:13:53 +0100
Subject: [PATCH] Implement WebSocket queries for /sql endpoint

---
 src/net/sql.rs | 50 ++++++++++++++++++++++++++++++++++++++------------
 1 file changed, 38 insertions(+), 12 deletions(-)

diff --git a/src/net/sql.rs b/src/net/sql.rs
index 31483ed9..38e1a249 100644
--- a/src/net/sql.rs
+++ b/src/net/sql.rs
@@ -4,8 +4,9 @@ use crate::net::output;
 use crate::net::session;
 use crate::net::DB;
 use bytes::Bytes;
-use futures::{FutureExt, StreamExt};
+use futures::{SinkExt, StreamExt};
 use surrealdb::Session;
+use warp::ws::{Message, WebSocket, Ws};
 use warp::Filter;
 
 const MAX: u64 = 1024 * 1024; // 1 MiB
@@ -24,17 +25,10 @@ pub fn config() -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejecti
 		.and(warp::body::bytes())
 		.and_then(handler);
 	// Set sock method
-	let sock = base.and(warp::ws()).map(|ws: warp::ws::Ws| {
-		ws.on_upgrade(|websocket| {
-			// Just echo all messages back...
-			let (tx, rx) = websocket.split();
-			rx.forward(tx).map(|result| {
-				if let Err(e) = result {
-					eprintln!("websocket error: {:?}", e);
-				}
-			})
-		})
-	});
+	let sock = base
+		.and(warp::ws())
+		.and(session::build())
+		.map(|ws: Ws, session: Session| ws.on_upgrade(move |ws| socket(ws, session)));
 	// Specify route
 	opts.or(post).or(sock).with(head::cors())
 }
@@ -44,15 +38,47 @@ async fn handler(
 	output: String,
 	sql: Bytes,
 ) -> Result<impl warp::Reply, warp::Rejection> {
+	// Get a database reference
 	let db = DB.get().unwrap();
+	// Convert the received sql query
 	let sql = std::str::from_utf8(&sql).unwrap();
+	// Execute the received sql query
 	match db.execute(sql, &session, None).await {
+		// Convert the response to JSON
 		Ok(res) => match output.as_ref() {
 			"application/json" => Ok(output::json(&res)),
 			"application/cbor" => Ok(output::cbor(&res)),
 			"application/msgpack" => Ok(output::pack(&res)),
+			// An incorrect content-type was requested
 			_ => Err(warp::reject::not_found()),
 		},
+		// There was an error when executing the query
 		Err(err) => Err(warp::reject::custom(Error::from(err))),
 	}
 }
+
+async fn socket(ws: WebSocket, session: Session) {
+	// Split the WebSocket connection
+	let (mut tx, mut rx) = ws.split();
+	// Wait to receive the next message
+	while let Some(res) = rx.next().await {
+		if let Ok(msg) = res {
+			if let Ok(sql) = msg.to_str() {
+				// Get a database reference
+				let db = DB.get().unwrap();
+				// Execute the received sql query
+				let _ = match db.execute(sql, &session, None).await {
+					// Convert the response to JSON
+					Ok(v) => match serde_json::to_string(&v) {
+						// Send the JSON response to the client
+						Ok(v) => tx.send(Message::text(v)).await,
+						// There was an error converting to JSON
+						Err(e) => tx.send(Message::text(Error::from(e))).await,
+					},
+					// There was an error when executing the query
+					Err(e) => tx.send(Message::text(Error::from(e))).await,
+				};
+			}
+		}
+	}
+}