noth/surrealpatch
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Implement WebSocket queries for /sql endpoint

Browse source
This commit is contained in:
Tobie Morgan Hitchcock 2022-05-11 22:13:53 +01:00
parent d29961ee19
commit e4619be89a
1 changed files with 38 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

50
src/net/sql.rs
View file

@ -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,
};
}
}
}
}