anna
/
async-graphql
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add `async_graphql_warp::graphql_protocol`, `async_graphql_warp::graphql_subscription_upgrade` and `async_graphql_warp::graphql_subscription_upgrade_with_data` to control WebSocket subscription more finely.

This commit is contained in:
Sunli 2021-04-05 13:20:02 +08:00
parent 3b45708959
commit 4003baa2a5
3 changed files with 132 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGELOG.md
View File

@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Changed
- Rework `Extension`, now fully supports asynchronous, better to use than before, and can achieve more features, it contains a lot of changes. _(if you don't have a custom extension, it will not cause the existing code to fail to compile)_
- Add `async_graphql_warp::graphql_protocol`, `async_graphql_warp::graphql_subscription_upgrade` and `async_graphql_warp::graphql_subscription_upgrade_with_data` to control WebSocket subscription more finely.
## [2.7.4] 2021-04-02

5
integrations/warp/src/lib.rs
View File

@ -11,4 +11,7 @@ mod subscription;
pub use batch_request::{graphql_batch, graphql_batch_opts, BatchResponse};
pub use error::BadRequest;
pub use request::{graphql, graphql_opts, Response};
pub use subscription::{graphql_subscription, graphql_subscription_with_data};
pub use subscription::{
graphql_protocol, graphql_subscription, graphql_subscription_upgrade,
graphql_subscription_upgrade_with_data, graphql_subscription_with_data,
};

189
integrations/warp/src/subscription.rs
View File

@ -1,9 +1,11 @@
use std::future::Future;
use std::str::FromStr;
use async_graphql::http::WsMessage;
use async_graphql::http::{WebSocketProtocols, WsMessage};
use async_graphql::{Data, ObjectType, Result, Schema, SubscriptionType};
use futures_util::{future, StreamExt};
use warp::filters::ws;
use warp::ws::WebSocket;
use warp::{Filter, Rejection, Reply};
/// GraphQL subscription filter
@ -74,72 +76,14 @@ where
F: FnOnce(serde_json::Value) -> R + Clone + Send + 'static,
R: Future<Output = Result<Data>> + Send + 'static,
{
graphql_subscription_with_data_and_callbacks(schema, initializer, || {}, || {})
}
/// GraphQL subscription filter
///
/// Specifies that a function converts the init payload to data.
pub fn graphql_subscription_with_data_and_callbacks<Query, Mutation, Subscription, F, R, FO, FC>(
schema: Schema<Query, Mutation, Subscription>,
initializer: F,
open_callback: FO,
close_callback: FC,
) -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone
where
Query: ObjectType + 'static,
Mutation: ObjectType + 'static,
Subscription: SubscriptionType + 'static,
F: FnOnce(serde_json::Value) -> R + Clone + Send + 'static,
R: Future<Output = Result<Data>> + Send + 'static,
FO: FnOnce() + Clone + Send + 'static,
FC: FnOnce() + Clone + Send + 'static,
{
use async_graphql::http::WebSocketProtocols;
use std::str::FromStr;
warp::ws()
.and(warp::header::optional::<String>("sec-websocket-protocol"))
.map(move |ws: ws::Ws, protocols: Option<String>| {
.and(graphql_protocol())
.map(move |ws: ws::Ws, protocol| {
let schema = schema.clone();
let initializer = initializer.clone();
let open_callback = open_callback.clone();
let close_callback = close_callback.clone();
let protocol = protocols
.and_then(|protocols| {
protocols
.split(',')
.find_map(|p| WebSocketProtocols::from_str(p.trim()).ok())
})
.unwrap_or(WebSocketProtocols::SubscriptionsTransportWS);
let reply = ws.on_upgrade(move |websocket| {
let (ws_sender, ws_receiver) = websocket.split();
open_callback();
async move {
let _ = async_graphql::http::WebSocket::with_data(
schema,
ws_receiver
.take_while(|msg| future::ready(msg.is_ok()))
.map(Result::unwrap)
.filter(|msg| future::ready(msg.is_text() || msg.is_binary()))
.map(ws::Message::into_bytes),
initializer,
protocol,
)
.map(|msg| match msg {
WsMessage::Text(text) => ws::Message::text(text),
WsMessage::Close(code, status) => ws::Message::close_with(code, status),
})
.map(Ok)
.forward(ws_sender)
.await;
close_callback();
}
graphql_subscription_upgrade_with_data(websocket, protocol, schema, initializer)
});
warp::reply::with_header(
@ -149,3 +93,124 @@ where
)
})
}
/// Create a `Filter` that parse [WebSocketProtocols] from `sec-websocket-protocol` header.
pub fn graphql_protocol() -> impl Filter<Extract = (WebSocketProtocols,), Error = Rejection> + Clone
{
warp::header::optional::<String>("sec-websocket-protocol").map(|protocols: Option<String>| {
protocols
.and_then(|protocols| {
protocols
.split(',')
.find_map(|p| WebSocketProtocols::from_str(p.trim()).ok())
})
.unwrap_or(WebSocketProtocols::SubscriptionsTransportWS)
})
}
/// Handle the WebSocket subscription.
///
/// If you want to control the WebSocket subscription more finely, you can use this function,
/// otherwise it is more convenient to use [graphql_subscription].
///
/// # Examples
///
/// ```no_run
/// use async_graphql::*;
/// use async_graphql_warp::*;
/// use warp::Filter;
/// use futures_util::stream::{Stream, StreamExt};
/// use std::time::Duration;
///
/// struct QueryRoot;
///
/// #[Object]
/// impl QueryRoot {
/// async fn value(&self) -> i32 {
/// // A GraphQL Object type must define one or more fields.
/// 100
/// }
/// }
///
/// struct SubscriptionRoot;
///
/// #[Subscription]
/// impl SubscriptionRoot {
/// async fn tick(&self) -> impl Stream<Item = String> {
/// async_stream::stream! {
/// let mut interval = tokio::time::interval(Duration::from_secs(1));
/// loop {
/// let n = interval.tick().await;
/// yield format!("{}", n.elapsed().as_secs_f32());
/// }
/// }
/// }
/// }
///
/// tokio::runtime::Runtime::new().unwrap().block_on(async {
/// let schema = Schema::new(QueryRoot, EmptyMutation, SubscriptionRoot);
/// let filter = warp::ws()
/// .and(graphql_protocol())
/// .map(move |ws: warp::ws::Ws, protocol| {
/// let schema = schema.clone();
/// let reply = ws.on_upgrade( move |websocket| {
/// graphql_subscription_upgrade(websocket, protocol, schema)
/// });
/// warp::reply::with_header(
/// reply,
/// "Sec-WebSocket-Protocol",
/// protocol.sec_websocket_protocol(),
/// )
/// });
/// warp::serve(filter).run(([0, 0, 0, 0], 8000)).await;
/// });
/// ```
pub async fn graphql_subscription_upgrade<Query, Mutation, Subscription>(
websocket: WebSocket,
protocol: WebSocketProtocols,
schema: Schema<Query, Mutation, Subscription>,
) where
Query: ObjectType + 'static,
Mutation: ObjectType + 'static,
Subscription: SubscriptionType + 'static,
{
graphql_subscription_upgrade_with_data(websocket, protocol, schema, |_| async {
Ok(Default::default())
})
.await
}
/// Handle the WebSocket subscription.
///
/// Specifies that a function converts the init payload to data.
pub async fn graphql_subscription_upgrade_with_data<Query, Mutation, Subscription, F, R>(
websocket: WebSocket,
protocol: WebSocketProtocols,
schema: Schema<Query, Mutation, Subscription>,
initializer: F,
) where
Query: ObjectType + 'static,
Mutation: ObjectType + 'static,
Subscription: SubscriptionType + 'static,
F: FnOnce(serde_json::Value) -> R + Clone + Send + 'static,
R: Future<Output = Result<Data>> + Send + 'static,
{
let (ws_sender, ws_receiver) = websocket.split();
let _ = async_graphql::http::WebSocket::with_data(
schema,
ws_receiver
.take_while(|msg| future::ready(msg.is_ok()))
.map(Result::unwrap)
.filter(|msg| future::ready(msg.is_text() || msg.is_binary()))
.map(ws::Message::into_bytes),
initializer,
protocol,
)
.map(|msg| match msg {
WsMessage::Text(text) => ws::Message::text(text),
WsMessage::Close(code, status) => ws::Message::close_with(code, status),
})
.map(Ok)
.forward(ws_sender)
.await;
}