servers/grpc/
database.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use api::v1::greptime_database_server::GreptimeDatabase;
16use api::v1::greptime_response::Response as RawResponse;
17use api::v1::{AffectedRows, GreptimeRequest, GreptimeResponse, ResponseHeader};
18use async_trait::async_trait;
19use common_error::status_code::StatusCode;
20use common_query::OutputData;
21use common_telemetry::{debug, warn};
22use futures::StreamExt;
23use prost::Message;
24use session::context::Channel;
25use tonic::{Request, Response, Status, Streaming};
26
27use crate::grpc::greptime_handler::GreptimeRequestHandler;
28use crate::grpc::{TonicResult, cancellation};
29use crate::hint_headers;
30use crate::request_memory_limiter::ServerMemoryLimiter;
31
32pub(crate) struct DatabaseService {
33    handler: GreptimeRequestHandler,
34}
35
36impl DatabaseService {
37    pub(crate) fn new(handler: GreptimeRequestHandler) -> Self {
38        Self { handler }
39    }
40}
41
42#[async_trait]
43impl GreptimeDatabase for DatabaseService {
44    async fn handle(
45        &self,
46        request: Request<GreptimeRequest>,
47    ) -> TonicResult<Response<GreptimeResponse>> {
48        let remote_addr = request.remote_addr();
49        let hints = hint_headers::extract_hints(request.metadata());
50        let channel = request
51            .extensions()
52            .get::<Channel>()
53            .copied()
54            .unwrap_or(Channel::Grpc);
55        debug!(
56            "GreptimeDatabase::Handle: request from {:?} with hints: {:?}",
57            remote_addr, hints
58        );
59
60        let _guard = if let Some(limiter) = request.extensions().get::<ServerMemoryLimiter>() {
61            let message_size = request.get_ref().encoded_len() as u64;
62            Some(limiter.acquire(message_size).await?)
63        } else {
64            None
65        };
66
67        let handler = self.handler.clone();
68        let request_future = async move {
69            let request = request.into_inner();
70            let output = handler.handle_request(request, hints, channel).await?;
71            let message = match output.data {
72                OutputData::AffectedRows(rows) => GreptimeResponse {
73                    header: Some(ResponseHeader {
74                        status: Some(api::v1::Status {
75                            status_code: StatusCode::Success as _,
76                            ..Default::default()
77                        }),
78                    }),
79                    response: Some(RawResponse::AffectedRows(AffectedRows { value: rows as _ })),
80                },
81                OutputData::Stream(_) | OutputData::RecordBatches(_) => {
82                    return Err(Status::unimplemented("GreptimeDatabase::Handle for query"));
83                }
84            };
85
86            Ok(Response::new(message))
87        };
88
89        let cancellation_future = async move {
90            warn!(
91                "GreptimeDatabase::Handle: request from {:?} cancelled by client",
92                remote_addr
93            );
94            // If this future is executed it means the request future was dropped,
95            // so it doesn't actually matter what is returned here
96            Err(Status::cancelled(
97                "GreptimeDatabase::Handle: request cancelled by client",
98            ))
99        };
100        cancellation::with_cancellation_handler(request_future, cancellation_future).await
101    }
102
103    async fn handle_requests(
104        &self,
105        request: Request<Streaming<GreptimeRequest>>,
106    ) -> Result<Response<GreptimeResponse>, Status> {
107        let remote_addr = request.remote_addr();
108        let hints = hint_headers::extract_hints(request.metadata());
109        let channel = request
110            .extensions()
111            .get::<Channel>()
112            .copied()
113            .unwrap_or(Channel::Grpc);
114        debug!(
115            "GreptimeDatabase::HandleRequests: request from {:?} with hints: {:?}",
116            remote_addr, hints
117        );
118
119        let limiter = request.extensions().get::<ServerMemoryLimiter>().cloned();
120
121        let handler = self.handler.clone();
122        let request_future = async move {
123            let mut affected_rows = 0;
124
125            let mut stream = request.into_inner();
126            while let Some(request) = stream.next().await {
127                let request = request?;
128
129                let _guard = if let Some(limiter_ref) = &limiter {
130                    let message_size = request.encoded_len() as u64;
131                    Some(limiter_ref.acquire(message_size).await?)
132                } else {
133                    None
134                };
135                let output = handler
136                    .handle_request(request, hints.clone(), channel)
137                    .await?;
138                match output.data {
139                    OutputData::AffectedRows(rows) => affected_rows += rows,
140                    OutputData::Stream(_) | OutputData::RecordBatches(_) => {
141                        return Err(Status::unimplemented(
142                            "GreptimeDatabase::HandleRequests for query",
143                        ));
144                    }
145                }
146            }
147            let message = GreptimeResponse {
148                header: Some(ResponseHeader {
149                    status: Some(api::v1::Status {
150                        status_code: StatusCode::Success as _,
151                        ..Default::default()
152                    }),
153                }),
154                response: Some(RawResponse::AffectedRows(AffectedRows {
155                    value: affected_rows as u32,
156                })),
157            };
158
159            Ok(Response::new(message))
160        };
161
162        let cancellation_future = async move {
163            warn!(
164                "GreptimeDatabase::HandleRequests: request from {:?} cancelled by client",
165                remote_addr
166            );
167            // If this future is executed it means the request future was dropped,
168            // so it doesn't actually matter what is returned here
169            Err(Status::cancelled(
170                "GreptimeDatabase::HandleRequests: request cancelled by client",
171            ))
172        };
173        cancellation::with_cancellation_handler(request_future, cancellation_future).await
174    }
175}
servers/grpc/database.rs

servers/grpc/
database.rs
