@ -10,23 +10,23 @@ use std::collections::BTreeMap;
use std ::convert ::Infallible ;
use std ::net ::{ Ipv6Addr , SocketAddr } ;
use std ::str ::FromStr ;
use std ::sync ::atomic ::{ AtomicU32 , Ordering } ;
use std ::sync ::{ Arc , Mutex } ;
use std ::thread ;
// use std::sync::atomic::{AtomicU32, Ordering};
use std ::sync ::Arc ;
// use std::thread;
use axum ::body ::{ boxed , Body , BoxBody } ;
use axum ::extract ::{ DefaultBodyLimit , Path , Query, State} ;
use axum ::extract ::{ DefaultBodyLimit , Path , } ;
use axum ::http ::{ header , HeaderName , Method , Request , Response , StatusCode } ;
use axum ::response ::sse ::{ Event , KeepAlive } ;
use axum ::response ::{ Html , Sse } ;
use axum ::routing ::{ get , post , put } ;
use axum ::{ , Router } ;
use axum ::{ Extension, Json, Router } ;
use clap ::Args ;
use futures ::future ::BoxFuture ;
use futures ::stream ::Stream ;
use itertools ::Itertools ;
use log ::{ error , info , warn } ;
use miette ::miette ;
// use miette::miette;
use rand ::Rng ;
use serde_json ::json ;
use tokio ::task ::spawn_blocking ;
@ -34,10 +34,7 @@ use tower_http::auth::{AsyncAuthorizeRequest, AsyncRequireAuthorizationLayer};
use tower_http ::compression ::CompressionLayer ;
use tower_http ::cors ::{ Any , CorsLayer } ;
use cozo ::{
format_error_as_json , DataValue , DbInstance , MultiTransaction , NamedRows , ScriptMutability ,
SimpleFixedRule ,
} ;
use cozo ::{ DataValue , DbInstance , NamedRows , ScriptMutability } ;
#[ derive(Args, Debug) ]
pub ( crate ) struct ServerArgs {
@ -57,9 +54,6 @@ pub(crate) struct ServerArgs {
#[ clap(short, long, default_value_t = String::from( " {} " )) ]
config : String ,
// When on, start REPL instead of starting a webserver
// #[clap(short, long)]
// repl: bool,
/// Address to bind the service to
#[ clap(short, long, default_value_t = String::from( " 127.0.0.1 " )) ]
bind : String ,
@ -67,21 +61,26 @@ pub(crate) struct ServerArgs {
/// Port to use
#[ clap(short = 'P', long, default_value_t = 9070) ]
port : u16 ,
/// When set, the content of the named table will be used as a token table
#[ clap(long) ]
token_table : Option < String > ,
}
#[ derive(Clone) ]
struct DbState {
db : DbInstance ,
rule_senders : Arc < Mutex < BTreeMap < u32 , crossbeam ::channel ::Sender < miette ::Result < NamedRows > > > > > ,
rule_counter : Arc < AtomicU32 > ,
tx_counter : Arc < AtomicU32 > ,
txs : Arc < Mutex < BTreeMap < u32 , Arc < MultiTransaction > > > > ,
// rule_senders: Arc<Mutex<BTreeMap<u32, crossbeam::channel::Sender<miette::Result<NamedRows>>>>>,
// rule_counter: Arc<AtomicU32>,
// tx_counter: Arc<AtomicU32>,
// txs: Arc<Mutex<BTreeMap<u32, Arc<MultiTransaction>>>>,
}
#[ derive(Clone) ]
struct MyAuth {
skip_auth : bool ,
auth_guard : String ,
token_table : Option < Arc < ( String , DbInstance ) > > ,
}
impl < B > AsyncAuthorizeRequest < B > for MyAuth
@ -92,17 +91,18 @@ where
type ResponseBody = BoxBody ;
type Future = BoxFuture < ' static , Result < Request < B > , Response < Self ::ResponseBody > > > ;
fn authorize ( & mut self , request : Request < B > ) -> Self ::Future {
fn authorize < ' a > ( & ' a mut self , mut request : Request < B > ) -> Self ::Future {
let skip_auth = self . skip_auth ;
let auth_guard = self . auth_guard . clone ( ) ;
let token_table = self . token_table . clone ( ) ;
Box ::pin ( async move {
if skip_auth {
request . extensions_mut ( ) . insert ( ScriptMutability ::Mutable ) ;
return Ok ( request ) ;
}
let ok = match request . headers ( ) . get ( "x-cozo-auth" ) {
let mutability = match request . headers ( ) . get ( "x-cozo-auth" ) {
None = > match request . uri ( ) . query ( ) {
None = > false ,
Some ( q_str ) = > {
let mut bingo = false ;
for pair in q_str . split ( '&' ) {
@ -115,15 +115,65 @@ where
}
}
}
bingo
if bingo {
Some ( ScriptMutability ::Mutable )
} else {
None
}
}
None = > match token_table {
None = > None ,
Some ( tt ) = > {
let ( name , db ) = tt . as_ref ( ) ;
if let Some ( auth_header ) = request . headers ( ) . get ( "Authorization" ) {
if let Ok ( auth_str ) = auth_header . to_str ( ) {
if auth_str . starts_with ( "Bearer " ) {
let token = & auth_str [ "Bearer " . len ( ) .. ] ;
match db . run_script (
& format! ( "?[mutable] := *{name} {{ token: $token, mutable }}" ) ,
BTreeMap ::from ( [ ( String ::from ( "token" ) , DataValue ::from ( token ) ) ] ) ,
ScriptMutability ::Immutable ,
) {
Ok ( rows ) = > match rows . rows . first ( ) {
None = > None ,
Some ( val ) = > {
if val [ 0 ] . get_bool ( ) = = Some ( true ) {
Some ( ScriptMutability ::Mutable )
} else {
Some ( ScriptMutability ::Immutable )
}
}
} ,
Err ( err ) = > {
eprintln! ( "Error: {}" , err ) ;
None
} ,
}
} else {
None
}
} else {
None
}
} else {
None
}
}
} ,
} ,
Some ( data ) = > match data . to_str ( ) {
Ok ( s ) = > s = = auth_guard . as_str ( ) ,
Err ( _ ) = > false ,
Ok ( s ) = > {
if s = = auth_guard . as_str ( ) {
Some ( ScriptMutability ::Mutable )
} else {
None
}
}
Err ( _ ) = > None ,
} ,
} ;
if ok {
if let Some ( mutability ) = mutability {
request . extensions_mut ( ) . insert ( mutability ) ;
Ok ( request )
} else {
let unauthorized_response = Response ::builder ( )
@ -174,12 +224,18 @@ pub(crate) async fn server_main(args: ServerArgs) {
}
} ;
let auth_obj = MyAuth {
skip_auth ,
auth_guard ,
token_table : args . token_table . map ( | t | Arc ::new ( ( t , db . clone ( ) ) ) ) ,
} ;
let state = DbState {
db ,
rule_senders : Default ::default ( ) ,
rule_counter : Default ::default ( ) ,
tx_counter : Default ::default ( ) ,
txs : Default ::default ( ) ,
// rule_senders: Default::default(),
// rule_counter: Default::default(),
// tx_counter: Default::default(),
// txs: Default::default(),
} ;
let cors = CorsLayer ::new ( )
. allow_methods ( [ Method ::GET , Method ::POST , Method ::PUT , Method ::DELETE ] )
@ -193,18 +249,15 @@ pub(crate) async fn server_main(args: ServerArgs) {
. route ( "/backup" , post ( backup ) )
. route ( "/import-from-backup" , post ( import_from_backup ) )
. route ( "/changes/:relation" , get ( observe_changes ) )
. route ( "/rules/:name" , get ( register_rule ) )
. route (
"/rule-result/:id" ,
post ( post_rule_result ) . delete ( post_rule_err ) ,
)
. route ( "/transact" , post ( start_transact ) )
. route ( "/transact/:id" , post ( transact_query ) . put ( finish_query ) )
// .route("/rules/:name", get(register_rule))
// .route(
// "/rule-result/:id",
// post(post_rule_result).delete(post_rule_err),
// )
// .route("/transact", post(start_transact))
// .route("/transact/:id", post(transact_query).put(finish_query))
. with_state ( state )
. layer ( AsyncRequireAuthorizationLayer ::new ( MyAuth {
skip_auth ,
auth_guard ,
} ) )
. layer ( AsyncRequireAuthorizationLayer ::new ( auth_obj ) )
. fallback ( not_found )
. route ( "/" , get ( root ) )
. layer ( cors )
@ -235,76 +288,76 @@ pub(crate) async fn server_main(args: ServerArgs) {
#[ derive(serde_derive::Deserialize) ]
struct StartTransactPayload {
write : bool ,
}
async fn start_transact (
State ( st ) : State < DbState > ,
Query ( payload ) : Query < StartTransactPayload > ,
) -> ( StatusCode , Json < serde_json ::Value > ) {
let tx = st . db . multi_transaction ( payload . write ) ;
let id = st . tx_counter . fetch_add ( 1 , Ordering ::SeqCst ) ;
st . txs . lock ( ) . unwrap ( ) . insert ( id , Arc ::new ( tx ) ) ;
( StatusCode ::OK , json ! ( { "ok" : true , "id" : id } ) . into ( ) )
}
async fn transact_query (
State ( st ) : State < DbState > ,
Path ( id ) : Path < u32 > ,
Json ( payload ) : Json < QueryPayload > ,
) -> ( StatusCode , Json < serde_json ::Value > ) {
let tx = match st . txs . lock ( ) . unwrap ( ) . get ( & id ) {
None = > return ( StatusCode ::NOT_FOUND , json ! ( { "ok" : false } ) . into ( ) ) ,
Some ( tx ) = > tx . clone ( ) ,
} ;
let src = payload . script . clone ( ) ;
let result = spawn_blocking ( move | | {
let params = payload
. params
. into_iter ( )
. map ( | ( k , v ) | ( k , DataValue ::from ( v ) ) )
. collect ( ) ;
let query = payload . script ;
tx . run_script ( & query , params )
} )
. await ;
match result {
Ok ( Ok ( res ) ) = > ( StatusCode ::OK , res . into_json ( ) . into ( ) ) ,
Ok ( Err ( err ) ) = > (
StatusCode ::BAD_REQUEST ,
format_error_as_json ( err , Some ( & src ) ) . into ( ) ,
) ,
Err ( err ) = > internal_error ( err ) ,
}
// write: bool,
}
#[ derive(serde_derive::Deserialize) ]
struct FinishTransactPayload {
abort : bool ,
}
async fn finish_query (
State ( st ) : State < DbState > ,
Path ( id ) : Path < u32 > ,
Json ( payload ) : Json < FinishTransactPayload > ,
) -> ( StatusCode , Json < serde_json ::Value > ) {
let tx = match st . txs . lock ( ) . unwrap ( ) . remove ( & id ) {
None = > return ( StatusCode ::NOT_FOUND , json ! ( { "ok" : false } ) . into ( ) ) ,
Some ( tx ) = > tx ,
} ;
let res = if payload . abort {
tx . abort ( )
} else {
tx . commit ( )
} ;
match res {
Ok ( _ ) = > ( StatusCode ::OK , json ! ( { "ok" : true } ) . into ( ) ) ,
Err ( err ) = > (
StatusCode ::BAD_REQUEST ,
json ! ( { "ok" : false , "message" : err . to_string ( ) } ) . into ( ) ,
) ,
}
}
// async fn start_transact(
// State(st): State<DbState>,
// Query(payload): Query<StartTransactPayload>,
// ) -> (StatusCode, Json<serde_json::Value>) {
// let tx = st.db.multi_transaction(payload.write);
// let id = st.tx_counter.fetch_add(1, Ordering::SeqCst);
// st.txs.lock().unwrap().insert(id, Arc::new(tx));
// (StatusCode::OK, json!({"ok": true, "id": id}).into())
// }
// async fn transact_query(
// State(st): State<DbState>,
// Path(id): Path<u32>,
// Json(payload): Json<QueryPayload>,
// ) -> (StatusCode, Json<serde_json::Value>) {
// let tx = match st.txs.lock().unwrap().get(&id) {
// None => return (StatusCode::NOT_FOUND, json!({"ok": false}).into()),
// Some(tx) => tx.clone(),
// };
// let src = payload.script.clone();
// let result = spawn_blocking(move || {
// let params = payload
// .params
// .into_iter()
// .map(|(k, v)| (k, DataValue::from(v)))
// .collect();
// let query = payload.script;
// tx.run_script(&query, params)
// })
// .await;
// match result {
// Ok(Ok(res)) => (StatusCode::OK, res.into_json().into()),
// Ok(Err(err)) => (
// StatusCode::BAD_REQUEST,
// format_error_as_json(err, Some(&src)).into(),
// ),
// Err(err) => internal_error(err),
// }
// }
// #[derive(serde_derive::Deserialize)]
// struct FinishTransactPayload {
// abort: bool,
// }
// async fn finish_query(
// State(st): State<DbState>,
// Path(id): Path<u32>,
// Json(payload): Json<FinishTransactPayload>,
// ) -> (StatusCode, Json<serde_json::Value>) {
// let tx = match st.txs.lock().unwrap().remove(&id) {
// None => return (StatusCode::NOT_FOUND, json!({"ok": false}).into()),
// Some(tx) => tx,
// };
// let res = if payload.abort {
// tx.abort()
// } else {
// tx.commit()
// };
// match res {
// Ok(_) => (StatusCode::OK, json!({"ok": true}).into()),
// Err(err) => (
// StatusCode::BAD_REQUEST,
// json!({"ok": false, "message": err.to_string()}).into(),
// ),
// }
// }
#[ derive(serde_derive::Deserialize) ]
struct QueryPayload {
@ -314,6 +367,7 @@ struct QueryPayload {
}
async fn text_query (
Extension ( mutability ) : Extension < ScriptMutability > ,
State ( st ) : State < DbState > ,
Json ( payload ) : Json < QueryPayload > ,
) -> ( StatusCode , Json < serde_json ::Value > ) {
@ -322,7 +376,10 @@ async fn text_query(
. into_iter ( )
. map ( | ( k , v ) | ( k , DataValue ::from ( v ) ) )
. collect ( ) ;
let immutable = payload . immutable . unwrap_or ( false ) ;
let immutable = match mutability {
ScriptMutability ::Mutable = > payload . immutable . unwrap_or ( false ) ,
ScriptMutability ::Immutable = > true ,
} ;
let result = spawn_blocking ( move | | {
st . db . run_script_fold_err (
& payload . script ,
@ -457,118 +514,118 @@ async fn import_from_backup(
}
}
#[ derive(serde_derive::Deserialize) ]
struct RuleRegisterOptions {
arity : usize ,
}
async fn post_rule_result (
State ( st ) : State < DbState > ,
Path ( id ) : Path < u32 > ,
Json ( res ) : Json < serde_json ::Value > ,
) -> ( StatusCode , Json < serde_json ::Value > ) {
let res = match NamedRows ::from_json ( & res ) {
Ok ( res ) = > res ,
Err ( err ) = > {
if let Some ( ch ) = st . rule_senders . lock ( ) . unwrap ( ) . remove ( & id ) {
let _ = ch . send ( Err ( miette ! ( "downstream posted malformed result" ) ) ) ;
}
return (
StatusCode ::BAD_REQUEST ,
json ! ( { "ok" : false , "message" : err . to_string ( ) } ) . into ( ) ,
) ;
}
} ;
if let Some ( ch ) = st . rule_senders . lock ( ) . unwrap ( ) . remove ( & id ) {
match ch . send ( Ok ( res ) ) {
Ok ( _ ) = > ( StatusCode ::OK , json ! ( { "ok" : true } ) . into ( ) ) ,
Err ( err ) = > (
StatusCode ::INTERNAL_SERVER_ERROR ,
json ! ( { "ok" : false , "message" : err . to_string ( ) } ) . into ( ) ,
) ,
}
} else {
( StatusCode ::NOT_FOUND , json ! ( { "ok" : false } ) . into ( ) )
}
}
async fn post_rule_err (
State ( st ) : State < DbState > ,
Path ( id ) : Path < u32 > ,
) -> ( StatusCode , Json < serde_json ::Value > ) {
if let Some ( ch ) = st . rule_senders . lock ( ) . unwrap ( ) . remove ( & id ) {
match ch . send ( Err ( miette ! ( "downstream cancelled computation" ) ) ) {
Ok ( _ ) = > ( StatusCode ::OK , json ! ( { "ok" : true } ) . into ( ) ) ,
Err ( err ) = > (
StatusCode ::INTERNAL_SERVER_ERROR ,
json ! ( { "ok" : false , "message" : err . to_string ( ) } ) . into ( ) ,
) ,
}
} else {
( StatusCode ::NOT_FOUND , json ! ( { "ok" : false } ) . into ( ) )
}
}
async fn register_rule (
State ( st ) : State < DbState > ,
Path ( name ) : Path < String > ,
Query ( rule_opts ) : Query < RuleRegisterOptions > ,
) -> Sse < impl Stream < Item = Result < Event , Infallible > > > {
let ( rule , task_receiver ) = SimpleFixedRule ::rule_with_channel ( rule_opts . arity ) ;
let ( down_sender , mut down_receiver ) = tokio ::sync ::mpsc ::channel ( 1 ) ;
let mut errored = None ;
if let Err ( err ) = st . db . register_fixed_rule ( name . clone ( ) , rule ) {
errored = Some ( err ) ;
} else {
let rule_senders = st . rule_senders . clone ( ) ;
let rule_counter = st . rule_counter . clone ( ) ;
thread ::spawn ( move | | {
for ( inputs , options , sender ) in task_receiver {
let id = rule_counter . fetch_add ( 1 , Ordering ::AcqRel ) ;
let inputs : serde_json ::Value =
inputs . into_iter ( ) . map ( | ip | ip . into_json ( ) ) . collect ( ) ;
let options : serde_json ::Value = options
. into_iter ( )
. map ( | ( k , v ) | ( k , serde_json ::Value ::from ( v ) ) )
. collect ( ) ;
if down_sender . blocking_send ( ( id , inputs , options ) ) . is_err ( ) {
let _ = sender . send ( Err ( miette ! ( "cannot send request to downstream" ) ) ) ;
} else {
rule_senders . lock ( ) . unwrap ( ) . insert ( id , sender ) ;
}
}
} ) ;
}
struct Guard {
name : String ,
db : DbInstance ,
}
impl Drop for Guard {
fn drop ( & mut self ) {
info ! ( "dropping rules SSE {}" , self . name ) ;
let _ = self . db . unregister_fixed_rule ( & self . name ) ;
}
}
let stream = async_stream ::stream ! {
if let Some ( err ) = errored {
let item = json ! ( { "type" : "register-error" , "error" : err . to_string ( ) } ) ;
yield Ok ( Event ::default ( ) . json_data ( item ) . unwrap ( ) ) ;
} else {
info ! ( "starting rule SSE {}" , name ) ;
let _guard = Guard { db : st . db , name } ;
while let Some ( ( id , inputs , options ) ) = down_receiver . recv ( ) . await {
let item = json ! ( { "type" : "request" , "id" : id , "inputs" : inputs , "options" : options } ) ;
yield Ok ( Event ::default ( ) . json_data ( item ) . unwrap ( ) ) ;
}
}
} ;
Sse ::new ( stream ) . keep_alive ( KeepAlive ::default ( ) )
}
// #[derive(serde_derive::Deserialize)]
// struct RuleRegisterOptions {
// arity: usize,
// }
// async fn post_rule_result(
// State(st): State<DbState>,
// Path(id): Path<u32>,
// Json(res): Json<serde_json::Value>,
// ) -> (StatusCode, Json<serde_json::Value>) {
// let res = match NamedRows::from_json(&res) {
// Ok(res) => res,
// Err(err) => {
// if let Some(ch) = st.rule_senders.lock().unwrap().remove(&id) {
// let _ = ch.send(Err(miette!("downstream posted malformed result")));
// }
// return (
// StatusCode::BAD_REQUEST,
// json!({"ok": false, "message": err.to_string()}).into(),
// );
// }
// };
// if let Some(ch) = st.rule_senders.lock().unwrap().remove(&id) {
// match ch.send(Ok(res)) {
// Ok(_) => (StatusCode::OK, json!({"ok": true}).into()),
// Err(err) => (
// StatusCode::INTERNAL_SERVER_ERROR,
// json!({"ok": false, "message": err.to_string()}).into(),
// ),
// }
// } else {
// (StatusCode::NOT_FOUND, json!({"ok": false}).into())
// }
// }
// async fn post_rule_err(
// State(st): State<DbState>,
// Path(id): Path<u32>,
// ) -> (StatusCode, Json<serde_json::Value>) {
// if let Some(ch) = st.rule_senders.lock().unwrap().remove(&id) {
// match ch.send(Err(miette!("downstream cancelled computation"))) {
// Ok(_) => (StatusCode::OK, json!({"ok": true}).into()),
// Err(err) => (
// StatusCode::INTERNAL_SERVER_ERROR,
// json!({"ok": false, "message": err.to_string()}).into(),
// ),
// }
// } else {
// (StatusCode::NOT_FOUND, json!({"ok": false}).into())
// }
// }
// async fn register_rule(
// State(st): State<DbState>,
// Path(name): Path<String>,
// Query(rule_opts): Query<RuleRegisterOptions>,
// ) -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
// let (rule, task_receiver) = SimpleFixedRule::rule_with_channel(rule_opts.arity);
// let (down_sender, mut down_receiver) = tokio::sync::mpsc::channel(1);
// let mut errored = None;
//
// if let Err(err) = st.db.register_fixed_rule(name.clone(), rule) {
// errored = Some(err);
// } else {
// let rule_senders = st.rule_senders.clone();
// let rule_counter = st.rule_counter.clone();
//
// thread::spawn(move || {
// for (inputs, options, sender) in task_receiver {
// let id = rule_counter.fetch_add(1, Ordering::AcqRel);
// let inputs: serde_json::Value =
// inputs.into_iter().map(|ip| ip.into_json()).collect();
// let options: serde_json::Value = options
// .into_iter()
// .map(|(k, v)| (k, serde_json::Value::from(v)))
// .collect();
// if down_sender.blocking_send((id, inputs, options)).is_err() {
// let _ = sender.send(Err(miette!("cannot send request to downstream")));
// } else {
// rule_senders.lock().unwrap().insert(id, sender);
// }
// }
// });
// }
//
// struct Guard {
// name: String,
// db: DbInstance,
// }
//
// impl Drop for Guard {
// fn drop(&mut self) {
// info!("dropping rules SSE {}", self.name);
// let _ = self.db.unregister_fixed_rule(&self.name);
// }
// }
//
// let stream = async_stream::stream! {
// if let Some(err) = errored {
// let item = json!({"type": "register-error", "error": err.to_string()});
// yield Ok(Event::default().json_data(item).unwrap());
// } else {
// info!("starting rule SSE {}", name);
// let _guard = Guard {db: st.db, name};
// while let Some((id, inputs, options)) = down_receiver.recv().await {
// let item = json!({"type": "request", "id": id, "inputs": inputs, "options": options});
// yield Ok(Event::default().json_data(item).unwrap());
// }
// }
// };
// Sse::new(stream).keep_alive(KeepAlive::default())
// }
async fn observe_changes (
State ( st ) : State < DbState > ,
Ziyang Hu
1 year ago