1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
use std::future::Future;
use std::str::FromStr;
use async_graphql::http::{
WebSocket as AGWebSocket, WebSocketProtocols, WsMessage, ALL_WEBSOCKET_PROTOCOLS,
};
use async_graphql::{Data, ObjectType, Result, Schema, SubscriptionType};
use futures_util::future::Ready;
use futures_util::{future, StreamExt};
use tide::Endpoint;
use tide_websockets::tungstenite::protocol::CloseFrame;
use tide_websockets::Message;
#[cfg_attr(docsrs, doc(cfg(feature = "websocket")))]
pub struct GraphQLSubscription<Query, Mutation, Subscription, OnConnInit> {
schema: Schema<Query, Mutation, Subscription>,
on_connection_init: OnConnInit,
}
type DefaultOnConnInitType = fn(serde_json::Value) -> Ready<async_graphql::Result<Data>>;
fn default_on_connection_init(_: serde_json::Value) -> Ready<async_graphql::Result<Data>> {
futures_util::future::ready(Ok(Data::default()))
}
impl<Query, Mutation, Subscription>
GraphQLSubscription<Query, Mutation, Subscription, DefaultOnConnInitType>
where
Query: ObjectType + 'static,
Mutation: ObjectType + 'static,
Subscription: SubscriptionType + 'static,
{
pub fn new(schema: Schema<Query, Mutation, Subscription>) -> Self {
GraphQLSubscription {
schema,
on_connection_init: default_on_connection_init,
}
}
}
impl<Query, Mutation, Subscription, OnConnInit, OnConnInitFut>
GraphQLSubscription<Query, Mutation, Subscription, OnConnInit>
where
Query: ObjectType + 'static,
Mutation: ObjectType + 'static,
Subscription: SubscriptionType + 'static,
OnConnInit: Fn(serde_json::Value) -> OnConnInitFut + Clone + Send + Sync + 'static,
OnConnInitFut: Future<Output = async_graphql::Result<Data>> + Send + 'static,
{
pub fn on_connection_init<OnConnInit2, Fut>(
self,
callback: OnConnInit2,
) -> GraphQLSubscription<Query, Mutation, Subscription, OnConnInit2>
where
OnConnInit2: Fn(serde_json::Value) -> Fut + Clone + Send + Sync + 'static,
Fut: Future<Output = async_graphql::Result<Data>> + Send + 'static,
{
GraphQLSubscription {
schema: self.schema,
on_connection_init: callback,
}
}
pub fn build<S: Send + Sync + Clone + 'static>(self) -> impl Endpoint<S> {
tide_websockets::WebSocket::<S, _>::new(move |request, connection| {
let schema = self.schema.clone();
let on_connection_init = self.on_connection_init.clone();
async move {
let protocol = match request
.header("sec-websocket-protocol")
.map(|value| value.as_str())
.and_then(|protocols| {
protocols
.split(',')
.find_map(|p| WebSocketProtocols::from_str(p.trim()).ok())
}) {
Some(protocol) => protocol,
None => {
WebSocketProtocols::SubscriptionsTransportWS
}
};
let sink = connection.clone();
let mut stream = AGWebSocket::new(
schema.clone(),
connection
.take_while(|msg| future::ready(msg.is_ok()))
.map(Result::unwrap)
.map(Message::into_data),
protocol,
)
.on_connection_init(on_connection_init);
while let Some(data) = stream.next().await {
match data {
WsMessage::Text(text) => {
if sink.send_string(text).await.is_err() {
break;
}
}
WsMessage::Close(code, msg) => {
let _ = sink
.send(Message::Close(Some(CloseFrame {
code: code.into(),
reason: msg.into(),
})))
.await;
break;
}
}
}
Ok(())
}
})
.with_protocols(&ALL_WEBSOCKET_PROTOCOLS)
}
}