use std::borrow::Cow;
use std::fmt::Debug;
use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use async_channel::Sender;
pub use surrealdb_rpc::QUERY_STREAM_BUFFER;
pub use surrealdb_rpc::export::Config as DbExportConfig;
use surrealdb_rpc::{QueryResult, QueryStreamItem, Token, items_for_result};
use surrealdb_types::{
Array, ConnectionError, Error, NotFoundError, Notification, Object, SurrealValue, Value,
Variables,
};
use uuid::Uuid;
pub mod session;
pub use session::{Established, SessionEntry, SessionRegistry};
pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
#[derive(Debug)]
pub struct RequestData {
pub command: Command,
pub session_id: Uuid,
}
#[derive(Debug)]
pub struct Route {
pub request: RequestData,
pub response: Sender<Result<Vec<QueryResult>, Error>>,
}
#[derive(Debug, Clone, Copy)]
pub enum SessionId {
Initial(Uuid),
Clone {
old: Uuid,
new: Uuid,
},
Drop(Uuid),
}
#[derive(Debug, Clone)]
pub enum SessionError {
NotFound(Uuid),
Remote(String),
}
impl From<SessionError> for Error {
fn from(error: SessionError) -> Self {
session_error_to_error(error)
}
}
pub fn session_error_to_error(e: SessionError) -> Error {
match e {
SessionError::NotFound(id) => Error::not_found(
format!("Session not found: {id}"),
NotFoundError::Session {
id: Some(id.to_string()),
},
),
SessionError::Remote(msg) => Error::internal(msg),
}
}
#[derive(Debug, Clone)]
pub struct MlExportConfig {
pub name: String,
pub version: String,
}
#[derive(Debug, Clone)]
pub enum Command {
Use {
namespace: Option<String>,
database: Option<String>,
},
Signup {
credentials: Object,
},
Signin {
credentials: Object,
},
Authenticate {
token: Token,
},
Refresh {
token: Token,
},
Invalidate,
Begin,
Rollback {
txn: Uuid,
},
Commit {
txn: Uuid,
},
Revoke {
token: Token,
},
Query {
txn: Option<Uuid>,
query: Cow<'static, str>,
variables: Variables,
},
ExportFile {
path: PathBuf,
config: Option<DbExportConfig>,
},
ExportMl {
path: PathBuf,
config: MlExportConfig,
},
ExportBytes {
bytes: Sender<Result<Vec<u8>, Error>>,
config: Option<DbExportConfig>,
},
ExportBytesMl {
bytes: Sender<Result<Vec<u8>, Error>>,
config: MlExportConfig,
},
ImportFile {
path: PathBuf,
},
ImportMl {
path: PathBuf,
},
Health,
Version,
Set {
key: String,
value: Value,
},
Unset {
key: String,
},
SubscribeLive {
uuid: Uuid,
notification_sender: Sender<Result<Notification, Error>>,
},
Kill {
uuid: Uuid,
},
Attach {
session_id: Uuid,
},
Detach {
session_id: Uuid,
},
Run {
name: String,
version: Option<String>,
args: Array,
},
}
#[derive(Debug, Clone, Copy)]
pub struct EngineContext {
pub session: Uuid,
pub transaction: Option<Uuid>,
}
impl EngineContext {
pub fn new(session: Uuid) -> Self {
Self {
session,
transaction: None,
}
}
pub fn with_transaction(session: Uuid, transaction: Option<Uuid>) -> Self {
Self {
session,
transaction,
}
}
}
pub type EngineFuture<'a, T> = BoxFuture<'a, Result<T, Error>>;
fn buffered_query_stream<'a, E>(
engine: &'a E,
ctx: EngineContext,
query: Cow<'static, str>,
variables: Variables,
items: Sender<QueryStreamItem>,
) -> EngineFuture<'a, ()>
where
E: SurrealEngine + ?Sized,
{
Box::pin(async move {
for (index, result) in engine.query(ctx, query, variables).await?.into_iter().enumerate() {
for item in items_for_result(index, result) {
if items.send(item).await.is_err() {
return Ok(());
}
}
}
Ok(())
})
}
pub trait SurrealEngine: Debug + Send + Sync + 'static {
fn query(
&self,
ctx: EngineContext,
query: Cow<'static, str>,
variables: Variables,
) -> EngineFuture<'_, Vec<QueryResult>>;
fn query_stream(
&self,
ctx: EngineContext,
query: Cow<'static, str>,
variables: Variables,
items: Sender<QueryStreamItem>,
) -> EngineFuture<'_, ()> {
buffered_query_stream(self, ctx, query, variables, items)
}
fn run(
&self,
ctx: EngineContext,
name: String,
version: Option<String>,
args: Array,
) -> EngineFuture<'_, Value>;
fn use_ns_db(
&self,
ctx: EngineContext,
namespace: Option<String>,
database: Option<String>,
) -> EngineFuture<'_, (Option<String>, Option<String>)>;
fn set(&self, ctx: EngineContext, key: String, value: Value) -> EngineFuture<'_, ()>;
fn unset(&self, ctx: EngineContext, key: String) -> EngineFuture<'_, ()>;
fn signup(&self, ctx: EngineContext, credentials: Object) -> EngineFuture<'_, Token>;
fn signin(&self, ctx: EngineContext, credentials: Object) -> EngineFuture<'_, Token>;
fn authenticate(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, Token>;
fn refresh(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, Token>;
fn revoke(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, ()>;
fn invalidate(&self, ctx: EngineContext) -> EngineFuture<'_, ()>;
fn begin(&self, ctx: EngineContext) -> EngineFuture<'_, Uuid>;
fn commit(&self, ctx: EngineContext, txn: Uuid) -> EngineFuture<'_, ()>;
fn rollback(&self, ctx: EngineContext, txn: Uuid) -> EngineFuture<'_, ()>;
fn health(&self, ctx: EngineContext) -> EngineFuture<'_, ()>;
fn version(&self, ctx: EngineContext) -> EngineFuture<'_, String>;
fn subscribe_live(
&self,
_ctx: EngineContext,
_uuid: Uuid,
_notifications: Sender<Result<Notification, Error>>,
) -> EngineFuture<'_, ()> {
Box::pin(async { Err(unsupported("Live queries")) })
}
fn kill(&self, _ctx: EngineContext, _uuid: Uuid) -> EngineFuture<'_, ()> {
Box::pin(async { Err(unsupported("Live queries")) })
}
fn export_file(
&self,
_ctx: EngineContext,
_path: PathBuf,
_config: Option<DbExportConfig>,
) -> EngineFuture<'_, ()> {
Box::pin(async { Err(unsupported("Export")) })
}
fn export_bytes(
&self,
_ctx: EngineContext,
_bytes: Sender<Result<Vec<u8>, Error>>,
_config: Option<DbExportConfig>,
) -> EngineFuture<'_, ()> {
Box::pin(async { Err(unsupported("Export")) })
}
fn export_ml_file(
&self,
_ctx: EngineContext,
_path: PathBuf,
_config: MlExportConfig,
) -> EngineFuture<'_, ()> {
Box::pin(async { Err(unsupported("Machine learning model export")) })
}
fn export_ml_bytes(
&self,
_ctx: EngineContext,
_bytes: Sender<Result<Vec<u8>, Error>>,
_config: MlExportConfig,
) -> EngineFuture<'_, ()> {
Box::pin(async { Err(unsupported("Machine learning model export")) })
}
fn import_file(&self, _ctx: EngineContext, _path: PathBuf) -> EngineFuture<'_, ()> {
Box::pin(async { Err(unsupported("Import")) })
}
fn import_ml_file(&self, _ctx: EngineContext, _path: PathBuf) -> EngineFuture<'_, ()> {
Box::pin(async { Err(unsupported("Machine learning model import")) })
}
}
fn unsupported(what: &str) -> Error {
Error::configuration(format!("{what} is not supported by this engine"), None)
}
pub fn single_result(mut results: Vec<QueryResult>) -> Result<Value, Error> {
match results.len() {
0 => Ok(Value::None),
1 => results.remove(0).result,
_ => Err(Error::internal("expected the database to return one or no results".to_string())),
}
}
#[derive(Debug, Clone)]
pub struct RouteChannelEngine {
sender: Sender<Route>,
}
impl RouteChannelEngine {
pub fn new(sender: Sender<Route>) -> Self {
Self {
sender,
}
}
async fn value(&self, command: Command, session: Uuid) -> Result<Value, Error> {
single_result(self.results(command, session).await?)
}
async fn results(&self, command: Command, session: Uuid) -> Result<Vec<QueryResult>, Error> {
let (response, receiver) = async_channel::bounded(1);
let route = Route {
request: RequestData {
command,
session_id: session,
},
response,
};
self.sender.send(route).await.map_err(|e| {
Error::connection(
format!("Failed to send command: {e}"),
ConnectionError::ConnectionFailed,
)
})?;
receiver.recv().await.map_err(|_| {
Error::connection(
"The engine dropped the request without answering".to_string(),
ConnectionError::ConnectionFailed,
)
})?
}
async fn unit(&self, command: Command, session: Uuid) -> Result<(), Error> {
match self.value(command, session).await? {
Value::None | Value::Null => Ok(()),
Value::Array(array) if array.is_empty() => Ok(()),
_ => Err(Error::internal("expected the database to return nothing".to_string())),
}
}
}
fn value_to_token(value: Value) -> Result<Token, Error> {
let value = match value {
Value::Array(array) if array.len() == 1 => {
array.into_iter().next().expect("array has exactly one element")
}
value => value,
};
Token::from_value(value)
}
impl SurrealEngine for RouteChannelEngine {
fn query(
&self,
ctx: EngineContext,
query: Cow<'static, str>,
variables: Variables,
) -> EngineFuture<'_, Vec<QueryResult>> {
Box::pin(self.results(
Command::Query {
txn: ctx.transaction,
query,
variables,
},
ctx.session,
))
}
fn run(
&self,
ctx: EngineContext,
name: String,
version: Option<String>,
args: Array,
) -> EngineFuture<'_, Value> {
Box::pin(self.value(
Command::Run {
name,
version,
args,
},
ctx.session,
))
}
fn use_ns_db(
&self,
ctx: EngineContext,
namespace: Option<String>,
database: Option<String>,
) -> EngineFuture<'_, (Option<String>, Option<String>)> {
Box::pin(async move {
let value = self
.value(
Command::Use {
namespace,
database,
},
ctx.session,
)
.await?;
let Value::Object(object) = value else {
return Ok((None, None));
};
let read = |key: &str| object.get(key).and_then(|v| v.as_string()).map(String::from);
Ok((read("namespace"), read("database")))
})
}
fn set(&self, ctx: EngineContext, key: String, value: Value) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::Set {
key,
value,
},
ctx.session,
))
}
fn unset(&self, ctx: EngineContext, key: String) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::Unset {
key,
},
ctx.session,
))
}
fn signup(&self, ctx: EngineContext, credentials: Object) -> EngineFuture<'_, Token> {
Box::pin(async move {
let value = self
.value(
Command::Signup {
credentials,
},
ctx.session,
)
.await?;
value_to_token(value)
})
}
fn signin(&self, ctx: EngineContext, credentials: Object) -> EngineFuture<'_, Token> {
Box::pin(async move {
let value = self
.value(
Command::Signin {
credentials,
},
ctx.session,
)
.await?;
value_to_token(value)
})
}
fn authenticate(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, Token> {
Box::pin(async move {
let value = self
.value(
Command::Authenticate {
token,
},
ctx.session,
)
.await?;
value_to_token(value)
})
}
fn refresh(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, Token> {
Box::pin(async move {
let value = self
.value(
Command::Refresh {
token,
},
ctx.session,
)
.await?;
value_to_token(value)
})
}
fn revoke(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::Revoke {
token,
},
ctx.session,
))
}
fn invalidate(&self, ctx: EngineContext) -> EngineFuture<'_, ()> {
Box::pin(self.unit(Command::Invalidate, ctx.session))
}
fn begin(&self, ctx: EngineContext) -> EngineFuture<'_, Uuid> {
Box::pin(async move {
let value = self.value(Command::Begin, ctx.session).await?;
let uuid = value.into_uuid().map_err(|e| Error::internal(e.to_string()))?;
Ok(uuid.into_inner())
})
}
fn commit(&self, ctx: EngineContext, txn: Uuid) -> EngineFuture<'_, ()> {
Box::pin(async move {
self.value(
Command::Commit {
txn,
},
ctx.session,
)
.await?;
Ok(())
})
}
fn rollback(&self, ctx: EngineContext, txn: Uuid) -> EngineFuture<'_, ()> {
Box::pin(async move {
self.value(
Command::Rollback {
txn,
},
ctx.session,
)
.await?;
Ok(())
})
}
fn health(&self, ctx: EngineContext) -> EngineFuture<'_, ()> {
Box::pin(self.unit(Command::Health, ctx.session))
}
fn version(&self, ctx: EngineContext) -> EngineFuture<'_, String> {
Box::pin(async move {
let value = self.value(Command::Version, ctx.session).await?;
value.into_string().map_err(|e| Error::internal(e.to_string()))
})
}
fn subscribe_live(
&self,
ctx: EngineContext,
uuid: Uuid,
notifications: Sender<Result<Notification, Error>>,
) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::SubscribeLive {
uuid,
notification_sender: notifications,
},
ctx.session,
))
}
fn kill(&self, ctx: EngineContext, uuid: Uuid) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::Kill {
uuid,
},
ctx.session,
))
}
fn export_file(
&self,
ctx: EngineContext,
path: PathBuf,
config: Option<DbExportConfig>,
) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::ExportFile {
path,
config,
},
ctx.session,
))
}
fn export_bytes(
&self,
ctx: EngineContext,
bytes: Sender<Result<Vec<u8>, Error>>,
config: Option<DbExportConfig>,
) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::ExportBytes {
bytes,
config,
},
ctx.session,
))
}
fn export_ml_file(
&self,
ctx: EngineContext,
path: PathBuf,
config: MlExportConfig,
) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::ExportMl {
path,
config,
},
ctx.session,
))
}
fn export_ml_bytes(
&self,
ctx: EngineContext,
bytes: Sender<Result<Vec<u8>, Error>>,
config: MlExportConfig,
) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::ExportBytesMl {
bytes,
config,
},
ctx.session,
))
}
fn import_file(&self, ctx: EngineContext, path: PathBuf) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::ImportFile {
path,
},
ctx.session,
))
}
fn import_ml_file(&self, ctx: EngineContext, path: PathBuf) -> EngineFuture<'_, ()> {
Box::pin(self.unit(
Command::ImportMl {
path,
},
ctx.session,
))
}
}
#[cfg(test)]
mod tests {
use surrealdb_types::Value;
use super::*;
#[tokio::test]
async fn the_buffered_adaptation_produces_the_same_items() {
let (sender, routes) = async_channel::bounded(1);
let engine = RouteChannelEngine::new(sender);
let (items, received) = async_channel::bounded(8);
let stream = engine.query_stream(
EngineContext::new(Uuid::nil()),
Cow::Borrowed("SELECT * FROM thing"),
Variables::default(),
items,
);
let serve = async {
let route = routes.recv().await.expect("a route");
let _ = route
.response
.send(Ok(vec![QueryResult {
time: std::time::Duration::ZERO,
result: Ok(Value::Array(vec![Value::Bool(true)].into())),
query_type: surrealdb_rpc::QueryType::Other,
}]))
.await;
};
let (outcome, ()) = futures::future::join(stream, serve).await;
outcome.expect("the engine answered");
let mut items = Vec::new();
while let Ok(item) = received.try_recv() {
items.push(item);
}
assert!(matches!(items[0], QueryStreamItem::Rows { .. }), "a list becomes rows");
assert!(
matches!(
items[1],
QueryStreamItem::Finished {
error: None,
..
}
),
"and the statement is terminated"
);
assert_eq!(items.len(), 2);
}
}