servers/grpc/
greptime_handler.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
15//! Handler for Greptime Database service. It's implemented by frontend.
16
17use std::collections::HashMap;
18use std::future::Future;
19use std::str::FromStr;
20use std::sync::{Arc, RwLock};
21use std::time::Instant;
22
23use api::helper::request_type;
24use api::v1::greptime_request::Request as QueryRequest;
25use api::v1::query_request::Query;
26use api::v1::{GreptimeRequest, QueryRequest as ApiQueryRequest, RequestHeader};
27use auth::UserProviderRef;
28use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
29use common_catalog::parse_catalog_and_schema_from_db_string;
30use common_error::ext::ErrorExt;
31use common_error::status_code::StatusCode;
32use common_grpc::flight::do_put::DoPutResponse;
33use common_query::Output;
34use common_runtime::Runtime;
35use common_runtime::runtime::RuntimeTrait;
36use common_session::ReadPreference;
37use common_telemetry::tracing_context::{FutureExt, TracingContext};
38use common_telemetry::{debug, error, tracing, warn};
39use common_time::timezone::parse_timezone;
40use futures_util::StreamExt;
41use session::context::{Channel, QueryContextBuilder, QueryContextRef};
42use session::hints::{INSERT_SKIP_WAL_HINT, READ_PREFERENCE_HINT, is_reserved_extension_key};
43use snafu::{OptionExt, ResultExt};
44use tokio::sync::mpsc;
45use tokio::sync::mpsc::error::TrySendError;
46use tokio_util::task::AbortOnDropHandle;
47use tonic::Status;
48
49use crate::error::{InvalidQuerySnafu, JoinTaskSnafu, Result, UnknownHintSnafu};
50use crate::grpc::flight::PutRecordBatchRequestStream;
51use crate::grpc::{FlightCompression, TonicResult, context_auth};
52use crate::metrics::{self, METRIC_SERVER_GRPC_DB_REQUEST_TIMER};
53use crate::query_handler::grpc::ServerGrpcQueryHandlerRef;
54
55#[derive(Clone)]
56pub struct GreptimeRequestHandler {
57    handler: ServerGrpcQueryHandlerRef,
58    pub(crate) user_provider: Option<UserProviderRef>,
59    runtime: Option<Runtime>,
60    pub(crate) flight_compression: FlightCompression,
61}
62
63impl GreptimeRequestHandler {
64    pub fn new(
65        handler: ServerGrpcQueryHandlerRef,
66        user_provider: Option<UserProviderRef>,
67        runtime: Option<Runtime>,
68        flight_compression: FlightCompression,
69    ) -> Self {
70        Self {
71            handler,
72            user_provider,
73            runtime,
74            flight_compression,
75        }
76    }
77
78    #[tracing::instrument(skip_all, fields(protocol = "grpc", request_type = get_request_type(&request)))]
79    pub(crate) async fn handle_request(
80        &self,
81        request: GreptimeRequest,
82        hints: Vec<(String, String)>,
83        channel: Channel,
84    ) -> Result<Output> {
85        let header = request.header.as_ref();
86        let hints = crate::hint_headers::validate_public_hints(hints)?;
87        let query_ctx = create_query_context(channel, header, hints, HashMap::new())?;
88        let query = request.request.context(InvalidQuerySnafu {
89            reason: "Expecting non-empty GreptimeRequest.",
90        })?;
91        self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx)
92            .await?;
93        self.handle_request_with_query_ctx(query, query_ctx).await
94    }
95
96    pub(crate) async fn authenticate_request_with_query_ctx(
97        &self,
98        header: Option<&RequestHeader>,
99        query_ctx: &QueryContextRef,
100    ) -> Result<()> {
101        let user_info = context_auth::auth(self.user_provider.clone(), header, query_ctx).await?;
102        query_ctx.set_current_user(user_info);
103        Ok(())
104    }
105
106    pub(crate) fn handle_request_with_query_ctx(
107        &self,
108        query: QueryRequest,
109        query_ctx: QueryContextRef,
110    ) -> impl Future<Output = Result<Output>> + Send + 'static {
111        let handler = self.handler.clone();
112        let runtime = self.runtime.clone();
113        // SQL can contain writes or DDL; only explicit read-query variants are cancellable.
114        let cancel_on_drop = matches!(
115            &query,
116            QueryRequest::Query(ApiQueryRequest {
117                query: Some(Query::LogicalPlan(_) | Query::PromRangeQuery(_)),
118            })
119        );
120        let request_type = request_type(&query).to_string();
121        let db = query_ctx.get_db_string();
122        let timer = RequestTimer::new(db.clone(), request_type);
123        let tracing_context = TracingContext::from_current_span();
124
125        async move {
126            let result_future = async move {
127                handler
128                .do_query(query, query_ctx)
129                .trace(tracing_context.attach(tracing::info_span!(
130                    "GreptimeRequestHandler::handle_request_runtime"
131                )))
132                .await
133                .map_err(|e| {
134                    if e.status_code().should_log_error() {
135                        let root_error = e.root_cause().unwrap_or(&e);
136                        error!(e; "Failed to handle request, error: {}", root_error.to_string());
137                    } else {
138                        // Currently, we still print a debug log.
139                        debug!("Failed to handle request, err: {:?}", e);
140                    }
141                    e
142                })
143            };
144
145            match runtime {
146                Some(runtime) => {
147                    let task = runtime.spawn(result_future);
148                    let result = if cancel_on_drop {
149                        AbortOnDropHandle::new(task).await
150                    } else {
151                        // Writes must outlive the RPC to avoid cancellation between state updates.
152                        // https://www.greptime.com/blogs/2023-01-12-hidden-control-flow.html
153                        task.await
154                    };
155                    result.context(JoinTaskSnafu).inspect_err(|e| {
156                        timer.record(e.status_code());
157                    })?
158                }
159                None => result_future.await,
160            }
161        }
162    }
163
164    pub(crate) async fn put_record_batches(
165        &self,
166        stream: PutRecordBatchRequestStream,
167        result_sender: mpsc::Sender<TonicResult<DoPutResponse>>,
168        query_ctx: QueryContextRef,
169    ) {
170        let handler = self.handler.clone();
171        let runtime = self
172            .runtime
173            .clone()
174            .unwrap_or_else(common_runtime::global_runtime);
175        runtime.spawn(async move {
176            let mut result_stream = handler.handle_put_record_batch_stream(stream, query_ctx);
177
178            while let Some(result) = result_stream.next().await {
179                match &result {
180                    Ok(response) => {
181                        // Record the elapsed time metric from the response
182                        metrics::GRPC_BULK_INSERT_ELAPSED.observe(response.elapsed_secs());
183                    }
184                    Err(e) => {
185                        error!(e; "Failed to handle flight record batches");
186                    }
187                }
188
189                if let Err(e) = result_sender.try_send(result.map_err(Status::from))
190                    && let TrySendError::Closed(_) = e
191                {
192                    warn!(r#""DoPut" client maybe unreachable, abort handling its message"#);
193                    break;
194                }
195            }
196        });
197    }
198}
199
200pub fn get_request_type(request: &GreptimeRequest) -> &'static str {
201    request
202        .request
203        .as_ref()
204        .map(request_type)
205        .unwrap_or_default()
206}
207
208/// Creates a new `QueryContext` from the provided request header and extensions.
209/// Strongly recommend setting an appropriate channel, as this is very helpful for statistics.
210pub(crate) fn create_query_context(
211    channel: Channel,
212    header: Option<&RequestHeader>,
213    extensions: Vec<(String, String)>,
214    snapshot_seqs: HashMap<u64, u64>,
215) -> Result<QueryContextRef> {
216    let (catalog, schema) = header
217        .map(|header| {
218            // We provide dbname field in newer versions of protos/sdks
219            // parse dbname from header in priority
220            if !header.dbname.is_empty() {
221                parse_catalog_and_schema_from_db_string(&header.dbname)
222            } else {
223                (
224                    if !header.catalog.is_empty() {
225                        header.catalog.to_lowercase()
226                    } else {
227                        DEFAULT_CATALOG_NAME.to_string()
228                    },
229                    if !header.schema.is_empty() {
230                        header.schema.to_lowercase()
231                    } else {
232                        DEFAULT_SCHEMA_NAME.to_string()
233                    },
234                )
235            }
236        })
237        .unwrap_or_else(|| {
238            (
239                DEFAULT_CATALOG_NAME.to_string(),
240                DEFAULT_SCHEMA_NAME.to_string(),
241            )
242        });
243    let timezone = parse_timezone(header.map(|h| h.timezone.as_str()));
244    let mut ctx_builder = QueryContextBuilder::default()
245        .current_catalog(catalog)
246        .current_schema(schema)
247        .timezone(timezone)
248        .channel(channel)
249        .snapshot_seqs(Arc::new(RwLock::new(snapshot_seqs)));
250
251    for (key, value) in extensions {
252        match key.as_str() {
253            READ_PREFERENCE_HINT => {
254                let Ok(read_preference) = ReadPreference::from_str(&value) else {
255                    return UnknownHintSnafu {
256                        hint: format!("{key}={value}"),
257                    }
258                    .fail();
259                };
260                ctx_builder = ctx_builder.read_preference(read_preference);
261            }
262            INSERT_SKIP_WAL_HINT => {
263                let skip_wal = value.parse::<bool>().map_err(|_| {
264                    UnknownHintSnafu {
265                        hint: format!("{key}={value}"),
266                    }
267                    .build()
268                })?;
269                ctx_builder = ctx_builder.skip_wal(skip_wal);
270            }
271            _ if is_reserved_extension_key(&key) => {
272                debug!(
273                    key = key.as_str(),
274                    "Ignoring reserved external query context extension key"
275                );
276            }
277            _ => {
278                ctx_builder = ctx_builder.set_extension(key, value);
279            }
280        }
281    }
282    Ok(ctx_builder.build().into())
283}
284
285/// Histogram timer for handling gRPC request.
286///
287/// The timer records the elapsed time with [StatusCode::Success] on drop.
288pub(crate) struct RequestTimer {
289    start: Instant,
290    db: String,
291    request_type: String,
292    status_code: StatusCode,
293}
294
295impl RequestTimer {
296    /// Returns a new timer.
297    pub fn new(db: String, request_type: String) -> RequestTimer {
298        RequestTimer {
299            start: Instant::now(),
300            db,
301            request_type,
302            status_code: StatusCode::Success,
303        }
304    }
305
306    /// Consumes the timer and record the elapsed time with specific `status_code`.
307    pub fn record(mut self, status_code: StatusCode) {
308        self.status_code = status_code;
309    }
310}
311
312impl Drop for RequestTimer {
313    fn drop(&mut self) {
314        METRIC_SERVER_GRPC_DB_REQUEST_TIMER
315            .with_label_values(&[
316                self.db.as_str(),
317                self.request_type.as_str(),
318                self.status_code.as_ref(),
319            ])
320            .observe(self.start.elapsed().as_secs_f64());
321    }
322}
323
324#[cfg(test)]
325mod tests {
326    use std::error::Error as StdError;
327    use std::result::Result as StdResult;
328    use std::time::Duration;
329
330    use async_trait::async_trait;
331    use chrono::FixedOffset;
332    use common_error::ext::BoxedError;
333    use common_error::{GREPTIME_DB_HEADER_ERROR_CODE, GREPTIME_DB_HEADER_ERROR_RETRY_HINT};
334    use common_time::Timezone;
335    use futures::stream::{self, BoxStream};
336    use query::options::FLOW_SCHEDULED_TIME_MILLIS;
337    use session::context::QueryContext;
338    use session::hints::{
339        INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY,
340        REMOTE_QUERY_ID_EXTENSION_KEY,
341    };
342    use snafu::ResultExt;
343    use tokio::sync::Notify;
344    use tokio::time::timeout;
345    use tonic::Code;
346
347    use super::*;
348    use crate::error::{ExecuteGrpcRequestSnafu, InvalidParameterSnafu};
349    use crate::query_handler::grpc::GrpcQueryHandler;
350
351    struct BlockingQueryHandler {
352        started: Arc<Notify>,
353        finished: Arc<Notify>,
354        resume: Arc<Notify>,
355        completed: Arc<Notify>,
356    }
357
358    struct QueryFinished(Arc<Notify>);
359
360    impl Drop for QueryFinished {
361        fn drop(&mut self) {
362            self.0.notify_one();
363        }
364    }
365
366    #[async_trait]
367    impl GrpcQueryHandler for BlockingQueryHandler {
368        async fn do_query(&self, _: QueryRequest, _: QueryContextRef) -> Result<Output> {
369            let _finished = QueryFinished(self.finished.clone());
370            self.started.notify_one();
371            self.resume.notified().await;
372            self.completed.notify_one();
373            Ok(Output::new_with_affected_rows(1))
374        }
375
376        fn handle_put_record_batch_stream(
377            &self,
378            _: PutRecordBatchRequestStream,
379            _: QueryContextRef,
380        ) -> BoxStream<'static, Result<DoPutResponse>> {
381            Box::pin(stream::empty())
382        }
383    }
384
385    async fn check_request_cancellation(
386        request: QueryRequest,
387        runtime: Option<Runtime>,
388        cancels: bool,
389    ) -> StdResult<(), Box<dyn StdError + Send + Sync>> {
390        let started = Arc::new(Notify::new());
391        let finished = Arc::new(Notify::new());
392        let resume = Arc::new(Notify::new());
393        let completed = Arc::new(Notify::new());
394        let handler = GreptimeRequestHandler::new(
395            Arc::new(BlockingQueryHandler {
396                started: started.clone(),
397                finished: finished.clone(),
398                resume: resume.clone(),
399                completed: completed.clone(),
400            }),
401            None,
402            runtime,
403            FlightCompression::default(),
404        );
405        let task = AbortOnDropHandle::new(tokio::spawn(
406            handler.handle_request_with_query_ctx(request, QueryContext::arc()),
407        ));
408        timeout(Duration::from_secs(5), started.notified()).await?;
409        task.abort();
410        assert!(task.await.is_err_and(|error| error.is_cancelled()));
411        if !cancels {
412            // Resuming after the caller is gone proves execution can still finish normally.
413            resume.notify_one();
414            timeout(Duration::from_secs(5), completed.notified()).await?;
415        }
416        timeout(Duration::from_secs(5), finished.notified()).await?;
417        Ok(())
418    }
419
420    #[tokio::test]
421    async fn test_drop_read_request_cancels_execution()
422    -> StdResult<(), Box<dyn StdError + Send + Sync>> {
423        // Payloads are opaque here; cancellation is tested before do_query returns Output.
424        for query in [
425            Query::LogicalPlan(Vec::new()),
426            Query::PromRangeQuery(Default::default()),
427        ] {
428            let request = QueryRequest::Query(ApiQueryRequest { query: Some(query) });
429            for runtime in [None, Some(common_runtime::query_runtime())] {
430                check_request_cancellation(request.clone(), runtime, true).await?;
431            }
432        }
433        Ok(())
434    }
435
436    #[tokio::test]
437    async fn test_drop_other_request_preserves_spawned_execution()
438    -> StdResult<(), Box<dyn StdError + Send + Sync>> {
439        let requests = [
440            QueryRequest::Inserts(Default::default()),
441            QueryRequest::RowInserts(Default::default()),
442            QueryRequest::Deletes(Default::default()),
443            QueryRequest::RowDeletes(Default::default()),
444            QueryRequest::Ddl(Default::default()),
445            QueryRequest::Query(ApiQueryRequest {
446                query: Some(Query::InsertIntoPlan(Default::default())),
447            }),
448            QueryRequest::Query(ApiQueryRequest { query: None }),
449            QueryRequest::Query(ApiQueryRequest {
450                query: Some(Query::Sql("SELECT * FROM numbers".into())),
451            }),
452            QueryRequest::Query(ApiQueryRequest {
453                query: Some(Query::Sql("INSERT INTO t VALUES (1)".into())),
454            }),
455            QueryRequest::Query(ApiQueryRequest {
456                query: Some(Query::Sql(
457                    "CREATE TABLE t (ts TIMESTAMP TIME INDEX)".into(),
458                )),
459            }),
460        ];
461        for request in requests {
462            for runtime in [None, Some(common_runtime::query_runtime())] {
463                // Without a separate runtime, preserve the existing inline cancellation behavior.
464                let cancels = runtime.is_none();
465                check_request_cancellation(request.clone(), runtime, cancels).await?;
466            }
467        }
468        Ok(())
469    }
470
471    #[test]
472    fn test_create_query_context_typed_skip_wal() {
473        let ctx = create_query_context(Channel::Grpc, None, vec![], HashMap::new()).unwrap();
474        assert!(!ctx.skip_wal());
475        let legacy = create_query_context(
476            Channel::Grpc,
477            None,
478            vec![("skip_wal".to_string(), "true".to_string())],
479            HashMap::new(),
480        )
481        .unwrap();
482        assert!(!legacy.skip_wal());
483        assert_eq!(legacy.extension("skip_wal"), Some("true"));
484        for (value, expected) in [("true", true), ("false", false)] {
485            let ctx = create_query_context(
486                Channel::Grpc,
487                None,
488                vec![(INSERT_SKIP_WAL_HINT.to_string(), value.to_string())],
489                HashMap::new(),
490            )
491            .unwrap();
492            assert_eq!(ctx.skip_wal(), expected);
493            assert_eq!(ctx.extension(INSERT_SKIP_WAL_HINT), None);
494        }
495        for value in ["", "TRUE", "1", "invalid"] {
496            assert!(
497                create_query_context(
498                    Channel::Grpc,
499                    None,
500                    vec![(INSERT_SKIP_WAL_HINT.to_string(), value.to_string())],
501                    HashMap::new()
502                )
503                .is_err()
504            );
505        }
506        let ctx = create_query_context(
507            Channel::Grpc,
508            None,
509            vec![
510                (INSERT_SKIP_WAL_HINT.to_string(), "true".to_string()),
511                (INSERT_SKIP_WAL_HINT.to_string(), "false".to_string()),
512            ],
513            HashMap::new(),
514        )
515        .unwrap();
516        assert!(!ctx.skip_wal());
517        assert_eq!(ctx.extension(INSERT_SKIP_WAL_HINT), None);
518    }
519
520    #[test]
521    fn test_create_query_context_read_preference_duplicates() {
522        for (values, valid) in [
523            (["leader", "LEADER"], true),
524            (["invalid", "leader"], false),
525            (["leader", "invalid"], false),
526        ] {
527            let result = create_query_context(
528                Channel::Grpc,
529                None,
530                values
531                    .into_iter()
532                    .map(|value| (READ_PREFERENCE_HINT.to_string(), value.to_string()))
533                    .collect(),
534                HashMap::new(),
535            );
536            if valid {
537                let ctx = result.unwrap();
538                assert!(matches!(ctx.read_preference(), ReadPreference::Leader));
539                assert_eq!(ctx.extension(READ_PREFERENCE_HINT), None);
540            } else {
541                assert!(result.is_err());
542            }
543        }
544    }
545
546    #[test]
547    fn test_create_query_context() {
548        let header = RequestHeader {
549            catalog: "cat-a-log".to_string(),
550            timezone: "+01:00".to_string(),
551            ..Default::default()
552        };
553        let query_context = create_query_context(
554            Channel::Unknown,
555            Some(&header),
556            vec![
557                ("auto_create_table".to_string(), "true".to_string()),
558                ("read_preference".to_string(), "leader".to_string()),
559                (
560                    READ_PREFERENCE_EXTENSION_KEY.to_string(),
561                    "follower".to_string(),
562                ),
563                (
564                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
565                    "spoofed".to_string(),
566                ),
567                (
568                    INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY.to_string(),
569                    "spoofed-regs".to_string(),
570                ),
571                (
572                    FLOW_SCHEDULED_TIME_MILLIS.to_string(),
573                    "1700000000000".to_string(),
574                ),
575            ],
576            HashMap::from([(7, 88)]),
577        )
578        .unwrap();
579        assert_eq!(query_context.get_snapshot(7), Some(88));
580        assert_eq!(query_context.current_catalog(), "cat-a-log");
581        assert_eq!(query_context.current_schema(), DEFAULT_SCHEMA_NAME);
582        assert_eq!(
583            query_context.timezone(),
584            Timezone::Offset(FixedOffset::east_opt(3600).unwrap())
585        );
586        assert!(matches!(
587            query_context.read_preference(),
588            ReadPreference::Leader
589        ));
590        assert_eq!(query_context.extension("auto_create_table"), Some("true"));
591        assert_eq!(query_context.extension(READ_PREFERENCE_HINT), None);
592        assert_eq!(query_context.extension(READ_PREFERENCE_EXTENSION_KEY), None);
593        assert_ne!(query_context.remote_query_id(), Some("spoofed"));
594        assert!(
595            query_context
596                .extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)
597                .is_none()
598        );
599        assert_eq!(
600            query_context.extension(FLOW_SCHEDULED_TIME_MILLIS),
601            Some("1700000000000")
602        );
603    }
604
605    #[test]
606    fn test_create_query_context_ignores_remote_query_id_extension() {
607        let query_context = create_query_context(
608            Channel::Grpc,
609            None,
610            vec![(
611                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
612                "spoofed-query-id".to_string(),
613            )],
614            HashMap::new(),
615        )
616        .unwrap();
617
618        assert_ne!(query_context.remote_query_id(), Some("spoofed-query-id"));
619        assert_eq!(
620            query_context.extension(REMOTE_QUERY_ID_EXTENSION_KEY),
621            query_context.remote_query_id()
622        );
623    }
624
625    #[test]
626    fn test_record_batch_error_to_status_preserves_error_details() {
627        let inner = InvalidParameterSnafu {
628            reason: "Column not found, column: new_col",
629        }
630        .build();
631        let err = Err::<(), _>(BoxedError::new(inner))
632            .context(ExecuteGrpcRequestSnafu)
633            .unwrap_err();
634
635        let status = Status::from(err);
636
637        assert_eq!(status.code(), Code::InvalidArgument);
638        assert!(
639            status
640                .message()
641                .contains("Column not found, column: new_col")
642        );
643        assert!(
644            status
645                .message()
646                .contains("Invalid request parameter: Column not found")
647        );
648        assert!(
649            status
650                .metadata()
651                .contains_key(GREPTIME_DB_HEADER_ERROR_CODE)
652        );
653        assert!(
654            status
655                .metadata()
656                .contains_key(GREPTIME_DB_HEADER_ERROR_RETRY_HINT)
657        );
658    }
659}
servers/grpc/greptime_handler.rs

servers/grpc/
greptime_handler.rs
