-
Notifications
You must be signed in to change notification settings - Fork 6
Expand file tree
/
Copy pathquery.rs
More file actions
124 lines (112 loc) · 4.42 KB
/
query.rs
File metadata and controls
124 lines (112 loc) · 4.42 KB
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
use crate::backend::DatabaseBackend;
use crate::interop::{into_tonic_status, IntoTonicStatus};
use crate::proto::query::{QueryResult, RawQuery, RowsChanged, TargetStore};
use crate::queryable::Queryable;
use crate::service::query::QueryRpc;
use crate::{DynStream, Location, RpcResponse, StreamingRequest};
use async_stream::stream;
use futures::StreamExt as _;
use tonic::{Response, Status};
/// The handler for raw queries. Supports both key-value and blob stores.
#[must_use]
#[derive(Debug)]
pub struct QueryHandler<Backend> {
kv_backend: Backend,
blob_backend: Backend,
}
impl<Backend> QueryHandler<Backend>
where
Backend: DatabaseBackend,
{
/// Create a new query handler at the given locations. No initialization is performed.
///
/// **Note**: At most one location can be in memory.
#[inline]
pub fn at_location(
kv_location: Location,
blob_location: Location,
) -> Result<Self, Backend::Error> {
// TODO validate that both locations are not in memory
Ok(Self {
kv_backend: Backend::at_location(kv_location)?,
blob_backend: Backend::at_location(blob_location)?,
})
}
/// Create a new query handler at the given paths on disk. No initialization is performed.
#[inline]
pub fn at_path<P1, P2>(kv_path: P1, blob_path: P2) -> Result<Self, Backend::Error>
where
P1: Into<std::path::PathBuf>,
P2: Into<std::path::PathBuf>,
{
Ok(Self {
kv_backend: Backend::at_location(kv_path.into().into())?,
blob_backend: Backend::at_location(blob_path.into().into())?,
})
}
}
#[tonic::async_trait]
impl<Backend> QueryRpc for QueryHandler<Backend>
where
Backend: DatabaseBackend<Error: IntoTonicStatus, Connection: Send>
+ Queryable<Connection = <Backend as DatabaseBackend>::Connection, QueryStream: Send>
+ Send
+ Sync
+ 'static,
{
type QueryStream = DynStream<Result<QueryResult, Status>>;
type ExecuteStream = DynStream<Result<RowsChanged, Status>>;
#[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]
async fn query(&self, request: StreamingRequest<RawQuery>) -> RpcResponse<Self::QueryStream> {
let mut request = request.into_inner();
let mut kv_conn = self.kv_backend.connect().map_err(into_tonic_status)?;
let mut blob_conn = self.blob_backend.connect().map_err(into_tonic_status)?;
let stream = stream!({
while let Some(RawQuery { query, target }) = request.message().await? {
match target.try_into().map_err(into_tonic_status)? {
TargetStore::Kv => {
let (mut items, conn) = Backend::query(query, kv_conn).await;
kv_conn = conn;
while let Some(item) = items.next().await {
yield item;
}
}
TargetStore::Blob => {
let (mut items, conn) = Backend::query(query, blob_conn).await;
blob_conn = conn;
while let Some(item) = items.next().await {
yield item;
}
}
}
}
});
Ok(Response::new(Box::pin(stream)))
}
#[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]
async fn execute(
&self,
request: StreamingRequest<RawQuery>,
) -> RpcResponse<Self::ExecuteStream> {
let mut request = request.into_inner();
let mut kv_conn = self.kv_backend.connect().map_err(into_tonic_status)?;
let mut blob_conn = self.blob_backend.connect().map_err(into_tonic_status)?;
let stream = stream!({
while let Some(RawQuery { query, target }) = request.message().await? {
match target.try_into().map_err(into_tonic_status)? {
TargetStore::Kv => {
let (res, conn) = Backend::execute(query, kv_conn).await;
kv_conn = conn;
yield res;
}
TargetStore::Blob => {
let (res, conn) = Backend::execute(query, blob_conn).await;
blob_conn = conn;
yield res;
}
}
}
});
Ok(Response::new(Box::pin(stream)))
}
}