use std::sync::Arc; use futures::{Future, Stream}; use rspc::{ internal::{ specta, BuiltProcedureBuilder, MiddlewareBuilderLike, RequestResult, UnbuiltProcedureBuilder, }, ErrorCode, Type, }; use serde::{de::DeserializeOwned, Deserialize, Serialize}; use uuid::Uuid; use crate::{api::Ctx, library::LibraryContext}; /// Can wrap a query argument to require it to contain a `library_id` and provide helpers for working with libraries. #[derive(Clone, Serialize, Deserialize, Type)] pub(crate) struct LibraryArgs { pub library_id: Uuid, pub arg: T, } // WARNING: This is system is using internal API's which means it will break between rspc release. I would avoid copying it unless you understand the cost of maintaining it! pub trait LibraryRequest { fn library_query( self, key: &'static str, builder: impl FnOnce( UnbuiltProcedureBuilder, ) -> BuiltProcedureBuilder, ) -> Self where TUnbuiltResolver: Fn(Ctx, TArg, LibraryContext) -> TUnbuiltResult + Send, TBuiltResolver: Fn(Ctx, TArg, LibraryContext) -> TUnbuiltResult + Send + Sync + 'static, TUnbuiltResult: RequestResult + Send, TArg: DeserializeOwned + specta::Type + Send + 'static; fn library_mutation< TUnbuiltResolver, TUnbuiltResult, TUnbuiltResultMarker, TBuiltResolver, TArg, >( self, key: &'static str, builder: impl FnOnce( UnbuiltProcedureBuilder, ) -> BuiltProcedureBuilder, ) -> Self where TUnbuiltResolver: Fn(Ctx, TArg, LibraryContext) -> TUnbuiltResult + Send, TBuiltResolver: Fn(Ctx, TArg, LibraryContext) -> TUnbuiltResult + Send + Sync + 'static, TUnbuiltResult: RequestResult + Send, TArg: DeserializeOwned + specta::Type + Send + 'static; fn library_subscription( self, key: &'static str, builder: impl Fn(UnbuiltProcedureBuilder) -> BuiltProcedureBuilder, ) -> Self where TArg: DeserializeOwned + specta::Type + 'static, TStream: Stream + Send + Sync + 'static, TResult: Serialize + specta::Type, TResolver: Fn(Ctx, TArg, Uuid) -> TStream + Send + Sync + 'static; } // Note: This will break with middleware context switching but that's fine for now impl LibraryRequest for rspc::RouterBuilder where TMiddleware: MiddlewareBuilderLike + Send + 'static, { fn library_query( self, key: &'static str, builder: impl FnOnce( UnbuiltProcedureBuilder, ) -> BuiltProcedureBuilder, ) -> Self where TUnbuiltResolver: Fn(Ctx, TArg, LibraryContext) -> TUnbuiltResult + Send, TBuiltResolver: Fn(Ctx, TArg, LibraryContext) -> TUnbuiltResult + Send + Sync + 'static, TUnbuiltResult: RequestResult + Send, TArg: DeserializeOwned + specta::Type + Send + 'static, { self.query(key, move |t| { let resolver = Arc::new(builder(UnbuiltProcedureBuilder::new(t.data())).resolver); t(move |ctx, arg: LibraryArgs| { let resolver = resolver.clone(); async move { let library = ctx .library_manager .get_ctx(arg.library_id) .await .ok_or_else(|| { rspc::Error::new( ErrorCode::BadRequest, "You must specify a valid library to use this operation." .to_string(), ) })?; Ok(resolver(ctx, arg.arg, library) .into_request_future()? .exec() .await?) } }) }) } fn library_mutation< TUnbuiltResolver, TUnbuiltResult, TUnbuiltResultMarker, TBuiltResolver, TArg, >( self, key: &'static str, builder: impl FnOnce( UnbuiltProcedureBuilder, ) -> BuiltProcedureBuilder, ) -> Self where TUnbuiltResolver: Fn(Ctx, TArg, LibraryContext) -> TUnbuiltResult + Send, TBuiltResolver: Fn(Ctx, TArg, LibraryContext) -> TUnbuiltResult + Send + Sync + 'static, TUnbuiltResult: RequestResult + Send, TArg: DeserializeOwned + specta::Type + Send + 'static, { self.mutation(key, move |t| { let resolver = Arc::new(builder(UnbuiltProcedureBuilder::new(t.data())).resolver); t(move |ctx, arg: LibraryArgs| { let resolver = resolver.clone(); async move { let library = ctx .library_manager .get_ctx(arg.library_id) .await .ok_or_else(|| { rspc::Error::new( ErrorCode::BadRequest, "You must specify a valid library to use this operation." .to_string(), ) })?; Ok(resolver(ctx, arg.arg, library) .into_request_future()? .exec() .await?) } }) }) } fn library_subscription( self, key: &'static str, builder: impl Fn(UnbuiltProcedureBuilder) -> BuiltProcedureBuilder, ) -> Self where TArg: DeserializeOwned + specta::Type + 'static, TStream: Stream + Send + Sync + 'static, TResult: Serialize + specta::Type, TResolver: Fn(Ctx, TArg, Uuid) -> TStream + Send + Sync + 'static, { self.subscription(key, |t| { let resolver = Arc::new(builder(UnbuiltProcedureBuilder::new(t.data())).resolver); t(move |ctx, arg: LibraryArgs| { // TODO(@Oscar): Upstream rspc work to allow this to work // let library = ctx // .library_manager // .get_ctx(arg.library_id) // .await // .ok_or_else(|| { // rspc::Error::new( // ErrorCode::BadRequest, // "You must specify a valid library to use this operation.".to_string(), // ) // })?; resolver(ctx, arg.arg, arg.library_id) }) }) } }