1use std::borrow::Cow;
20use std::fmt::Debug;
21use std::future::Future;
22use std::path::PathBuf;
23use std::pin::Pin;
24
25use async_channel::Sender;
26pub use surrealdb_rpc::QUERY_STREAM_BUFFER;
27pub use surrealdb_rpc::export::Config as DbExportConfig;
28use surrealdb_rpc::{QueryResult, QueryStreamItem, Token, items_for_result};
29use surrealdb_types::{
30 Array, ConnectionError, Error, NotFoundError, Notification, Object, SurrealValue, Value,
31 Variables,
32};
33use uuid::Uuid;
34
35pub mod session;
36
37pub use session::{Established, SessionEntry, SessionRegistry};
38
39pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
47
48#[derive(Debug)]
51pub struct RequestData {
52 pub command: Command,
54 pub session_id: Uuid,
56}
57
58#[derive(Debug)]
60pub struct Route {
61 pub request: RequestData,
63 pub response: Sender<Result<Vec<QueryResult>, Error>>,
65}
66
67#[derive(Debug, Clone, Copy)]
71pub enum SessionId {
72 Initial(Uuid),
74 Clone {
76 old: Uuid,
78 new: Uuid,
80 },
81 Drop(Uuid),
83}
84
85#[derive(Debug, Clone)]
87pub enum SessionError {
88 NotFound(Uuid),
90 Remote(String),
92}
93
94impl From<SessionError> for Error {
95 fn from(error: SessionError) -> Self {
96 session_error_to_error(error)
97 }
98}
99
100pub fn session_error_to_error(e: SessionError) -> Error {
102 match e {
103 SessionError::NotFound(id) => Error::not_found(
104 format!("Session not found: {id}"),
105 NotFoundError::Session {
106 id: Some(id.to_string()),
107 },
108 ),
109 SessionError::Remote(msg) => Error::internal(msg),
110 }
111}
112
113#[derive(Debug, Clone)]
115pub struct MlExportConfig {
116 pub name: String,
118 pub version: String,
120}
121
122#[derive(Debug, Clone)]
124pub enum Command {
125 Use {
127 namespace: Option<String>,
129 database: Option<String>,
131 },
132 Signup {
134 credentials: Object,
136 },
137 Signin {
139 credentials: Object,
141 },
142 Authenticate {
144 token: Token,
146 },
147 Refresh {
149 token: Token,
151 },
152 Invalidate,
154 Begin,
156 Rollback {
158 txn: Uuid,
160 },
161 Commit {
163 txn: Uuid,
165 },
166 Revoke {
168 token: Token,
170 },
171 Query {
173 txn: Option<Uuid>,
175 query: Cow<'static, str>,
177 variables: Variables,
179 },
180 ExportFile {
182 path: PathBuf,
184 config: Option<DbExportConfig>,
186 },
187 ExportMl {
189 path: PathBuf,
191 config: MlExportConfig,
193 },
194 ExportBytes {
196 bytes: Sender<Result<Vec<u8>, Error>>,
198 config: Option<DbExportConfig>,
200 },
201 ExportBytesMl {
203 bytes: Sender<Result<Vec<u8>, Error>>,
205 config: MlExportConfig,
207 },
208 ImportFile {
210 path: PathBuf,
212 },
213 ImportMl {
215 path: PathBuf,
217 },
218 Health,
220 Version,
222 Set {
224 key: String,
226 value: Value,
228 },
229 Unset {
231 key: String,
233 },
234 SubscribeLive {
236 uuid: Uuid,
238 notification_sender: Sender<Result<Notification, Error>>,
240 },
241 Kill {
243 uuid: Uuid,
245 },
246 Attach {
248 session_id: Uuid,
250 },
251 Detach {
253 session_id: Uuid,
255 },
256 Run {
258 name: String,
260 version: Option<String>,
262 args: Array,
264 },
265}
266
267#[derive(Debug, Clone, Copy)]
273pub struct EngineContext {
274 pub session: Uuid,
276 pub transaction: Option<Uuid>,
281}
282
283impl EngineContext {
284 pub fn new(session: Uuid) -> Self {
286 Self {
287 session,
288 transaction: None,
289 }
290 }
291
292 pub fn with_transaction(session: Uuid, transaction: Option<Uuid>) -> Self {
294 Self {
295 session,
296 transaction,
297 }
298 }
299}
300
301pub type EngineFuture<'a, T> = BoxFuture<'a, Result<T, Error>>;
303
304fn buffered_query_stream<'a, E>(
310 engine: &'a E,
311 ctx: EngineContext,
312 query: Cow<'static, str>,
313 variables: Variables,
314 items: Sender<QueryStreamItem>,
315) -> EngineFuture<'a, ()>
316where
317 E: SurrealEngine + ?Sized,
318{
319 Box::pin(async move {
320 for (index, result) in engine.query(ctx, query, variables).await?.into_iter().enumerate() {
321 for item in items_for_result(index, result) {
322 if items.send(item).await.is_err() {
325 return Ok(());
326 }
327 }
328 }
329 Ok(())
330 })
331}
332
333pub trait SurrealEngine: Debug + Send + Sync + 'static {
354 fn query(
360 &self,
361 ctx: EngineContext,
362 query: Cow<'static, str>,
363 variables: Variables,
364 ) -> EngineFuture<'_, Vec<QueryResult>>;
365
366 fn query_stream(
383 &self,
384 ctx: EngineContext,
385 query: Cow<'static, str>,
386 variables: Variables,
387 items: Sender<QueryStreamItem>,
388 ) -> EngineFuture<'_, ()> {
389 buffered_query_stream(self, ctx, query, variables, items)
390 }
391
392 fn run(
394 &self,
395 ctx: EngineContext,
396 name: String,
397 version: Option<String>,
398 args: Array,
399 ) -> EngineFuture<'_, Value>;
400
401 fn use_ns_db(
408 &self,
409 ctx: EngineContext,
410 namespace: Option<String>,
411 database: Option<String>,
412 ) -> EngineFuture<'_, (Option<String>, Option<String>)>;
413
414 fn set(&self, ctx: EngineContext, key: String, value: Value) -> EngineFuture<'_, ()>;
416
417 fn unset(&self, ctx: EngineContext, key: String) -> EngineFuture<'_, ()>;
419
420 fn signup(&self, ctx: EngineContext, credentials: Object) -> EngineFuture<'_, Token>;
426
427 fn signin(&self, ctx: EngineContext, credentials: Object) -> EngineFuture<'_, Token>;
429
430 fn authenticate(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, Token>;
437
438 fn refresh(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, Token>;
440
441 fn revoke(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, ()>;
443
444 fn invalidate(&self, ctx: EngineContext) -> EngineFuture<'_, ()>;
446
447 fn begin(&self, ctx: EngineContext) -> EngineFuture<'_, Uuid>;
453
454 fn commit(&self, ctx: EngineContext, txn: Uuid) -> EngineFuture<'_, ()>;
456
457 fn rollback(&self, ctx: EngineContext, txn: Uuid) -> EngineFuture<'_, ()>;
459
460 fn health(&self, ctx: EngineContext) -> EngineFuture<'_, ()>;
466
467 fn version(&self, ctx: EngineContext) -> EngineFuture<'_, String>;
470
471 fn subscribe_live(
477 &self,
478 _ctx: EngineContext,
479 _uuid: Uuid,
480 _notifications: Sender<Result<Notification, Error>>,
481 ) -> EngineFuture<'_, ()> {
482 Box::pin(async { Err(unsupported("Live queries")) })
483 }
484
485 fn kill(&self, _ctx: EngineContext, _uuid: Uuid) -> EngineFuture<'_, ()> {
487 Box::pin(async { Err(unsupported("Live queries")) })
488 }
489
490 fn export_file(
496 &self,
497 _ctx: EngineContext,
498 _path: PathBuf,
499 _config: Option<DbExportConfig>,
500 ) -> EngineFuture<'_, ()> {
501 Box::pin(async { Err(unsupported("Export")) })
502 }
503
504 fn export_bytes(
506 &self,
507 _ctx: EngineContext,
508 _bytes: Sender<Result<Vec<u8>, Error>>,
509 _config: Option<DbExportConfig>,
510 ) -> EngineFuture<'_, ()> {
511 Box::pin(async { Err(unsupported("Export")) })
512 }
513
514 fn export_ml_file(
516 &self,
517 _ctx: EngineContext,
518 _path: PathBuf,
519 _config: MlExportConfig,
520 ) -> EngineFuture<'_, ()> {
521 Box::pin(async { Err(unsupported("Machine learning model export")) })
522 }
523
524 fn export_ml_bytes(
526 &self,
527 _ctx: EngineContext,
528 _bytes: Sender<Result<Vec<u8>, Error>>,
529 _config: MlExportConfig,
530 ) -> EngineFuture<'_, ()> {
531 Box::pin(async { Err(unsupported("Machine learning model export")) })
532 }
533
534 fn import_file(&self, _ctx: EngineContext, _path: PathBuf) -> EngineFuture<'_, ()> {
536 Box::pin(async { Err(unsupported("Import")) })
537 }
538
539 fn import_ml_file(&self, _ctx: EngineContext, _path: PathBuf) -> EngineFuture<'_, ()> {
541 Box::pin(async { Err(unsupported("Machine learning model import")) })
542 }
543}
544
545fn unsupported(what: &str) -> Error {
547 Error::configuration(format!("{what} is not supported by this engine"), None)
548}
549
550pub fn single_result(mut results: Vec<QueryResult>) -> Result<Value, Error> {
557 match results.len() {
558 0 => Ok(Value::None),
559 1 => results.remove(0).result,
560 _ => Err(Error::internal("expected the database to return one or no results".to_string())),
561 }
562}
563
564#[derive(Debug, Clone)]
578pub struct RouteChannelEngine {
579 sender: Sender<Route>,
580}
581
582impl RouteChannelEngine {
583 pub fn new(sender: Sender<Route>) -> Self {
585 Self {
586 sender,
587 }
588 }
589
590 async fn value(&self, command: Command, session: Uuid) -> Result<Value, Error> {
598 single_result(self.results(command, session).await?)
599 }
600
601 async fn results(&self, command: Command, session: Uuid) -> Result<Vec<QueryResult>, Error> {
603 let (response, receiver) = async_channel::bounded(1);
604 let route = Route {
605 request: RequestData {
606 command,
607 session_id: session,
608 },
609 response,
610 };
611 self.sender.send(route).await.map_err(|e| {
615 Error::connection(
616 format!("Failed to send command: {e}"),
617 ConnectionError::ConnectionFailed,
618 )
619 })?;
620 receiver.recv().await.map_err(|_| {
621 Error::connection(
622 "The engine dropped the request without answering".to_string(),
623 ConnectionError::ConnectionFailed,
624 )
625 })?
626 }
627
628 async fn unit(&self, command: Command, session: Uuid) -> Result<(), Error> {
630 match self.value(command, session).await? {
631 Value::None | Value::Null => Ok(()),
632 Value::Array(array) if array.is_empty() => Ok(()),
633 _ => Err(Error::internal("expected the database to return nothing".to_string())),
634 }
635 }
636}
637
638fn value_to_token(value: Value) -> Result<Token, Error> {
645 let value = match value {
648 Value::Array(array) if array.len() == 1 => {
649 array.into_iter().next().expect("array has exactly one element")
650 }
651 value => value,
652 };
653 Token::from_value(value)
654}
655
656impl SurrealEngine for RouteChannelEngine {
657 fn query(
658 &self,
659 ctx: EngineContext,
660 query: Cow<'static, str>,
661 variables: Variables,
662 ) -> EngineFuture<'_, Vec<QueryResult>> {
663 Box::pin(self.results(
664 Command::Query {
665 txn: ctx.transaction,
666 query,
667 variables,
668 },
669 ctx.session,
670 ))
671 }
672
673 fn run(
674 &self,
675 ctx: EngineContext,
676 name: String,
677 version: Option<String>,
678 args: Array,
679 ) -> EngineFuture<'_, Value> {
680 Box::pin(self.value(
681 Command::Run {
682 name,
683 version,
684 args,
685 },
686 ctx.session,
687 ))
688 }
689
690 fn use_ns_db(
691 &self,
692 ctx: EngineContext,
693 namespace: Option<String>,
694 database: Option<String>,
695 ) -> EngineFuture<'_, (Option<String>, Option<String>)> {
696 Box::pin(async move {
697 let value = self
698 .value(
699 Command::Use {
700 namespace,
701 database,
702 },
703 ctx.session,
704 )
705 .await?;
706 let Value::Object(object) = value else {
710 return Ok((None, None));
711 };
712 let read = |key: &str| object.get(key).and_then(|v| v.as_string()).map(String::from);
713 Ok((read("namespace"), read("database")))
714 })
715 }
716
717 fn set(&self, ctx: EngineContext, key: String, value: Value) -> EngineFuture<'_, ()> {
718 Box::pin(self.unit(
719 Command::Set {
720 key,
721 value,
722 },
723 ctx.session,
724 ))
725 }
726
727 fn unset(&self, ctx: EngineContext, key: String) -> EngineFuture<'_, ()> {
728 Box::pin(self.unit(
729 Command::Unset {
730 key,
731 },
732 ctx.session,
733 ))
734 }
735
736 fn signup(&self, ctx: EngineContext, credentials: Object) -> EngineFuture<'_, Token> {
737 Box::pin(async move {
738 let value = self
739 .value(
740 Command::Signup {
741 credentials,
742 },
743 ctx.session,
744 )
745 .await?;
746 value_to_token(value)
747 })
748 }
749
750 fn signin(&self, ctx: EngineContext, credentials: Object) -> EngineFuture<'_, Token> {
751 Box::pin(async move {
752 let value = self
753 .value(
754 Command::Signin {
755 credentials,
756 },
757 ctx.session,
758 )
759 .await?;
760 value_to_token(value)
761 })
762 }
763
764 fn authenticate(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, Token> {
765 Box::pin(async move {
766 let value = self
767 .value(
768 Command::Authenticate {
769 token,
770 },
771 ctx.session,
772 )
773 .await?;
774 value_to_token(value)
775 })
776 }
777
778 fn refresh(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, Token> {
779 Box::pin(async move {
780 let value = self
781 .value(
782 Command::Refresh {
783 token,
784 },
785 ctx.session,
786 )
787 .await?;
788 value_to_token(value)
789 })
790 }
791
792 fn revoke(&self, ctx: EngineContext, token: Token) -> EngineFuture<'_, ()> {
793 Box::pin(self.unit(
794 Command::Revoke {
795 token,
796 },
797 ctx.session,
798 ))
799 }
800
801 fn invalidate(&self, ctx: EngineContext) -> EngineFuture<'_, ()> {
802 Box::pin(self.unit(Command::Invalidate, ctx.session))
803 }
804
805 fn begin(&self, ctx: EngineContext) -> EngineFuture<'_, Uuid> {
806 Box::pin(async move {
807 let value = self.value(Command::Begin, ctx.session).await?;
808 let uuid = value.into_uuid().map_err(|e| Error::internal(e.to_string()))?;
809 Ok(uuid.into_inner())
810 })
811 }
812
813 fn commit(&self, ctx: EngineContext, txn: Uuid) -> EngineFuture<'_, ()> {
814 Box::pin(async move {
815 self.value(
816 Command::Commit {
817 txn,
818 },
819 ctx.session,
820 )
821 .await?;
822 Ok(())
823 })
824 }
825
826 fn rollback(&self, ctx: EngineContext, txn: Uuid) -> EngineFuture<'_, ()> {
827 Box::pin(async move {
828 self.value(
829 Command::Rollback {
830 txn,
831 },
832 ctx.session,
833 )
834 .await?;
835 Ok(())
836 })
837 }
838
839 fn health(&self, ctx: EngineContext) -> EngineFuture<'_, ()> {
840 Box::pin(self.unit(Command::Health, ctx.session))
841 }
842
843 fn version(&self, ctx: EngineContext) -> EngineFuture<'_, String> {
844 Box::pin(async move {
845 let value = self.value(Command::Version, ctx.session).await?;
846 value.into_string().map_err(|e| Error::internal(e.to_string()))
847 })
848 }
849
850 fn subscribe_live(
851 &self,
852 ctx: EngineContext,
853 uuid: Uuid,
854 notifications: Sender<Result<Notification, Error>>,
855 ) -> EngineFuture<'_, ()> {
856 Box::pin(self.unit(
857 Command::SubscribeLive {
858 uuid,
859 notification_sender: notifications,
860 },
861 ctx.session,
862 ))
863 }
864
865 fn kill(&self, ctx: EngineContext, uuid: Uuid) -> EngineFuture<'_, ()> {
866 Box::pin(self.unit(
867 Command::Kill {
868 uuid,
869 },
870 ctx.session,
871 ))
872 }
873
874 fn export_file(
875 &self,
876 ctx: EngineContext,
877 path: PathBuf,
878 config: Option<DbExportConfig>,
879 ) -> EngineFuture<'_, ()> {
880 Box::pin(self.unit(
881 Command::ExportFile {
882 path,
883 config,
884 },
885 ctx.session,
886 ))
887 }
888
889 fn export_bytes(
890 &self,
891 ctx: EngineContext,
892 bytes: Sender<Result<Vec<u8>, Error>>,
893 config: Option<DbExportConfig>,
894 ) -> EngineFuture<'_, ()> {
895 Box::pin(self.unit(
896 Command::ExportBytes {
897 bytes,
898 config,
899 },
900 ctx.session,
901 ))
902 }
903
904 fn export_ml_file(
905 &self,
906 ctx: EngineContext,
907 path: PathBuf,
908 config: MlExportConfig,
909 ) -> EngineFuture<'_, ()> {
910 Box::pin(self.unit(
911 Command::ExportMl {
912 path,
913 config,
914 },
915 ctx.session,
916 ))
917 }
918
919 fn export_ml_bytes(
920 &self,
921 ctx: EngineContext,
922 bytes: Sender<Result<Vec<u8>, Error>>,
923 config: MlExportConfig,
924 ) -> EngineFuture<'_, ()> {
925 Box::pin(self.unit(
926 Command::ExportBytesMl {
927 bytes,
928 config,
929 },
930 ctx.session,
931 ))
932 }
933
934 fn import_file(&self, ctx: EngineContext, path: PathBuf) -> EngineFuture<'_, ()> {
935 Box::pin(self.unit(
936 Command::ImportFile {
937 path,
938 },
939 ctx.session,
940 ))
941 }
942
943 fn import_ml_file(&self, ctx: EngineContext, path: PathBuf) -> EngineFuture<'_, ()> {
944 Box::pin(self.unit(
945 Command::ImportMl {
946 path,
947 },
948 ctx.session,
949 ))
950 }
951}
952
953#[cfg(test)]
954mod tests {
955 use surrealdb_types::Value;
956
957 use super::*;
958
959 #[tokio::test]
962 async fn the_buffered_adaptation_produces_the_same_items() {
963 let (sender, routes) = async_channel::bounded(1);
964 let engine = RouteChannelEngine::new(sender);
965 let (items, received) = async_channel::bounded(8);
966 let stream = engine.query_stream(
967 EngineContext::new(Uuid::nil()),
968 Cow::Borrowed("SELECT * FROM thing"),
969 Variables::default(),
970 items,
971 );
972 let serve = async {
973 let route = routes.recv().await.expect("a route");
974 let _ = route
975 .response
976 .send(Ok(vec![QueryResult {
977 time: std::time::Duration::ZERO,
978 result: Ok(Value::Array(vec![Value::Bool(true)].into())),
979 query_type: surrealdb_rpc::QueryType::Other,
980 }]))
981 .await;
982 };
983 let (outcome, ()) = futures::future::join(stream, serve).await;
984 outcome.expect("the engine answered");
985
986 let mut items = Vec::new();
987 while let Ok(item) = received.try_recv() {
988 items.push(item);
989 }
990 assert!(matches!(items[0], QueryStreamItem::Rows { .. }), "a list becomes rows");
991 assert!(
992 matches!(
993 items[1],
994 QueryStreamItem::Finished {
995 error: None,
996 ..
997 }
998 ),
999 "and the statement is terminated"
1000 );
1001 assert_eq!(items.len(), 2);
1002 }
1003}