frontend/instance/
grpc.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 std::pin::Pin;
16use std::sync::Arc;
17use std::time::Instant;
18
19use api::helper::from_pb_time_ranges;
20use api::v1::ddl_request::{Expr as DdlExpr, Expr};
21use api::v1::greptime_request::Request;
22use api::v1::query_request::Query;
23use api::v1::{
24    DeleteRequests, DropFlowExpr, InsertIntoPlan, InsertRequests, RowDeleteRequests,
25    RowInsertRequests,
26};
27use async_stream::try_stream;
28use async_trait::async_trait;
29use auth::{
30    PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionResp, PermissionTableTargets,
31};
32use common_error::ext::BoxedError;
33use common_grpc::flight::do_put::DoPutResponse;
34use common_meta::rpc::ddl::TriggerReason;
35use common_query::Output;
36use common_query::logical_plan::add_insert_to_logical_plan;
37use common_telemetry::tracing::{self};
38use datafusion::datasource::DefaultTableSource;
39use futures::Stream;
40use futures::stream::StreamExt;
41use query::parser::PromQuery;
42use servers::error as server_error;
43use servers::http::prom_store::PHYSICAL_TABLE_PARAM;
44use servers::interceptor::{GrpcQueryInterceptor, GrpcQueryInterceptorRef};
45use servers::query_handler::grpc::GrpcQueryHandler;
46use session::context::QueryContextRef;
47use snafu::{OptionExt, ResultExt, ensure};
48use table::TableRef;
49use table::table::adapter::DfTableProviderAdapter;
50use table::table_name::TableName;
51
52use crate::error::{
53    CatalogSnafu, DataFusionSnafu, Error, ExternalSnafu, IncompleteGrpcRequestSnafu,
54    NotSupportedSnafu, PermissionSnafu, PlanStatementSnafu, Result,
55    SubstraitDecodeLogicalPlanSnafu, TableNotFoundSnafu, TableOperationSnafu,
56};
57use crate::instance::{Instance, attach_timer};
58use crate::metrics::{
59    GRPC_HANDLE_PLAN_ELAPSED, GRPC_HANDLE_PROMQL_ELAPSED, GRPC_HANDLE_SQL_ELAPSED,
60};
61
62#[async_trait]
63impl GrpcQueryHandler for Instance {
64    async fn do_query(
65        &self,
66        request: Request,
67        ctx: QueryContextRef,
68    ) -> server_error::Result<Output> {
69        let result: Result<Output> = async {
70            let interceptor_ref = self.plugins.get::<GrpcQueryInterceptorRef<Error>>();
71            let interceptor = interceptor_ref.as_ref();
72            interceptor.pre_execute(&request, ctx.clone())?;
73
74            if !matches!(
75                &request,
76                Request::Query(query_request)
77                    if matches!(&query_request.query, Some(Query::Sql(_)))
78            ) {
79                self.plugins
80                    .get::<PermissionCheckerRef>()
81                    .as_ref()
82                    .check_permission_with_context(
83                        ctx.current_user(),
84                        PermissionReq::GrpcRequest(&request),
85                        Some(&ctx.current_schema()),
86                    )
87                    .context(PermissionSnafu)?;
88            }
89
90            let output = match request {
91                Request::Inserts(requests) => self.handle_inserts(requests, ctx.clone()).await?,
92                Request::RowInserts(requests) => match ctx.extension(PHYSICAL_TABLE_PARAM) {
93                    Some(physical_table) => {
94                        self.handle_metric_row_inserts(
95                            requests,
96                            ctx.clone(),
97                            physical_table.to_string(),
98                        )
99                        .await?
100                    }
101                    None => {
102                        self.handle_row_inserts(requests, ctx.clone(), false, false)
103                            .await?
104                    }
105                },
106                Request::Deletes(requests) => self.handle_deletes(requests, ctx.clone()).await?,
107                Request::RowDeletes(requests) => self.handle_row_deletes(requests, ctx.clone()).await?,
108                Request::Query(query_request) => {
109                    let query = query_request.query.context(IncompleteGrpcRequestSnafu {
110                        err_msg: "Missing field 'QueryRequest.query'",
111                    })?;
112                    match query {
113                        Query::Sql(sql) => {
114                            let timer = GRPC_HANDLE_SQL_ELAPSED.start_timer();
115                            let mut result = self.do_query_inner(&sql, ctx.clone()).await;
116                            ensure!(
117                                result.len() == 1,
118                                NotSupportedSnafu {
119                                    feat: "execute multiple statements in SQL query string through GRPC interface"
120                                }
121                            );
122                            let output = result.remove(0)?;
123                            attach_timer(output, timer)
124                        }
125                        Query::LogicalPlan(plan) => {
126                            // this path is useful internally when flownode needs to execute a logical plan through gRPC interface
127                            let timer = GRPC_HANDLE_PLAN_ELAPSED.start_timer();
128
129                            // use dummy catalog to provide table
130                            let plan_decoder = self
131                                .query_engine()
132                                .engine_context(ctx.clone())
133                                .new_plan_decoder()
134                                .context(PlanStatementSnafu)?;
135
136                            let dummy_catalog_list =
137                                Arc::new(catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx(
138                                    self.catalog_manager().clone(),
139                                    ctx.clone(),
140                                ));
141
142                            let logical_plan = plan_decoder
143                                .decode(bytes::Bytes::from(plan), dummy_catalog_list, true)
144                                .await
145                                .context(SubstraitDecodeLogicalPlanSnafu)?;
146                            let output =
147                                self.do_exec_plan_inner(logical_plan, None, ctx.clone()).await?;
148
149                            attach_timer(output, timer)
150                        }
151                        Query::InsertIntoPlan(insert) => {
152                            self.handle_insert_plan(insert, ctx.clone()).await?
153                        }
154                        Query::PromRangeQuery(promql) => {
155                            let timer = GRPC_HANDLE_PROMQL_ELAPSED.start_timer();
156                            let prom_query = PromQuery {
157                                query: promql.query,
158                                start: promql.start,
159                                end: promql.end,
160                                step: promql.step,
161                                lookback: promql.lookback,
162                                alias: None,
163                            };
164                            let mut result =
165                                self.do_promql_query_inner(&prom_query, ctx.clone()).await;
166                            ensure!(
167                                result.len() == 1,
168                                NotSupportedSnafu {
169                                    feat: "execute multiple statements in PromQL query string through GRPC interface"
170                                }
171                            );
172                            let output = result.remove(0)?;
173                            attach_timer(output, timer)
174                        }
175                    }
176                }
177                Request::Ddl(request) => {
178                    let mut expr = request.expr.context(IncompleteGrpcRequestSnafu {
179                        err_msg: "'expr' is absent in DDL request",
180                    })?;
181
182                    fill_catalog_and_schema_from_context(&mut expr, &ctx);
183
184                    match expr {
185                        DdlExpr::CreateTable(mut expr) => {
186                            // Direct gRPC DDL bypasses the SQL parser, so validate the
187                            // request here (e.g. the time index must be a timestamp).
188                            operator::expr_helper::validate_create_expr(&expr)?;
189                            let _ = self
190                                .statement_executor
191                                .create_table_inner(
192                                    &mut expr,
193                                    None,
194                                    ctx.clone(),
195                                    TriggerReason::Manual,
196                                )
197                                .await?;
198                            Output::new_with_affected_rows(0)
199                        }
200                        DdlExpr::AlterDatabase(expr) => {
201                            let _ = self
202                                .statement_executor
203                                .alter_database_inner(expr, ctx.clone())
204                                .await?;
205                            Output::new_with_affected_rows(0)
206                        }
207                        DdlExpr::AlterTable(expr) => {
208                            self.statement_executor
209                                .alter_table_inner(expr, ctx.clone(), TriggerReason::Manual)
210                                .await?
211                        }
212                        DdlExpr::CreateDatabase(expr) => {
213                            self.statement_executor
214                                .create_database(
215                                    &expr.schema_name,
216                                    expr.create_if_not_exists,
217                                    expr.options,
218                                    ctx.clone(),
219                                )
220                                .await?
221                        }
222                        DdlExpr::DropTable(expr) => {
223                            let table_name =
224                                TableName::new(&expr.catalog_name, &expr.schema_name, &expr.table_name);
225                            self.statement_executor
226                                .drop_table(table_name, expr.drop_if_exists, ctx.clone())
227                                .await?
228                        }
229                        DdlExpr::TruncateTable(expr) => {
230                            let table_name =
231                                TableName::new(&expr.catalog_name, &expr.schema_name, &expr.table_name);
232                            let time_ranges = from_pb_time_ranges(expr.time_ranges.unwrap_or_default())
233                                .map_err(BoxedError::new)
234                                .context(ExternalSnafu)?;
235                            self.statement_executor
236                                .truncate_table(table_name, time_ranges, ctx.clone())
237                                .await?
238                        }
239                        DdlExpr::CreateFlow(expr) => {
240                            self.statement_executor
241                                .create_flow_inner(expr, ctx.clone())
242                                .await?
243                        }
244                        DdlExpr::DropFlow(DropFlowExpr {
245                            catalog_name,
246                            flow_name,
247                            drop_if_exists,
248                            ..
249                        }) => {
250                            self.statement_executor
251                                .drop_flow(catalog_name, flow_name, drop_if_exists, ctx.clone())
252                                .await?
253                        }
254                        DdlExpr::CreateView(expr) => {
255                            let _ = self
256                                .statement_executor
257                                .create_view_by_expr(expr, ctx.clone())
258                                .await?;
259
260                            Output::new_with_affected_rows(0)
261                        }
262                        DdlExpr::DropView(expr) => {
263                            self.statement_executor
264                                .drop_view(
265                                    expr.catalog_name,
266                                    expr.schema_name,
267                                    expr.view_name,
268                                    expr.drop_if_exists,
269                                    ctx.clone(),
270                                )
271                                .await?
272                        }
273                        DdlExpr::CommentOn(expr) => {
274                            self.statement_executor
275                                .comment_by_expr(expr, ctx.clone())
276                                .await?
277                        }
278                    }
279                }
280            };
281
282            let output = interceptor.post_execute(output, ctx)?;
283            Ok(output)
284        }
285        .await;
286
287        result
288            .map_err(BoxedError::new)
289            .context(server_error::ExecuteGrpcQuerySnafu)
290    }
291
292    fn handle_put_record_batch_stream(
293        &self,
294        stream: servers::grpc::flight::PutRecordBatchRequestStream,
295        ctx: QueryContextRef,
296    ) -> Pin<Box<dyn Stream<Item = server_error::Result<DoPutResponse>> + Send>> {
297        Box::pin(
298            self.handle_put_record_batch_stream_inner(stream, ctx)
299                .map(|result| {
300                    result
301                        .map_err(BoxedError::new)
302                        .context(server_error::ExecuteGrpcRequestSnafu)
303                }),
304        )
305    }
306}
307
308fn fill_catalog_and_schema_from_context(ddl_expr: &mut DdlExpr, ctx: &QueryContextRef) {
309    let catalog = ctx.current_catalog();
310    let schema = ctx.current_schema();
311
312    macro_rules! check_and_fill {
313        ($expr:ident) => {
314            if $expr.catalog_name.is_empty() {
315                $expr.catalog_name = catalog.to_string();
316            }
317            if $expr.schema_name.is_empty() {
318                $expr.schema_name = schema.to_string();
319            }
320        };
321    }
322
323    match ddl_expr {
324        Expr::CreateDatabase(_) | Expr::AlterDatabase(_) => { /* do nothing*/ }
325        Expr::CreateTable(expr) => {
326            check_and_fill!(expr);
327        }
328        Expr::AlterTable(expr) => {
329            check_and_fill!(expr);
330        }
331        Expr::DropTable(expr) => {
332            check_and_fill!(expr);
333        }
334        Expr::TruncateTable(expr) => {
335            check_and_fill!(expr);
336        }
337        Expr::CreateFlow(expr) => {
338            if expr.catalog_name.is_empty() {
339                expr.catalog_name = catalog.to_string();
340            }
341        }
342        Expr::DropFlow(expr) => {
343            if expr.catalog_name.is_empty() {
344                expr.catalog_name = catalog.to_string();
345            }
346        }
347        Expr::CreateView(expr) => {
348            check_and_fill!(expr);
349        }
350        Expr::DropView(expr) => {
351            check_and_fill!(expr);
352        }
353        Expr::CommentOn(expr) => {
354            check_and_fill!(expr);
355        }
356    }
357}
358
359impl Instance {
360    pub(crate) fn check_table_permission(
361        &self,
362        ctx: &QueryContextRef,
363        req: PermissionReq<'_>,
364        targets: PermissionTableTargets,
365    ) -> auth::error::Result<PermissionResp> {
366        self.plugins
367            .get::<PermissionCheckerRef>()
368            .as_ref()
369            .check_permission_with_table_targets(ctx.current_user(), req, targets)
370    }
371
372    /// Checks every logical table targeted by normalized row inserts.
373    pub(crate) fn check_row_insert_permission(
374        &self,
375        requests: &RowInsertRequests,
376        ctx: &QueryContextRef,
377        req: PermissionReq<'_>,
378    ) -> auth::error::Result<PermissionResp> {
379        let catalog = ctx.current_catalog();
380        let schema = ctx.current_schema();
381        let targets = PermissionTableTargets::from_row_insert_requests(catalog, &schema, requests);
382
383        self.check_table_permission(ctx, req, targets)
384    }
385
386    fn handle_put_record_batch_stream_inner(
387        &self,
388        mut stream: servers::grpc::flight::PutRecordBatchRequestStream,
389        ctx: QueryContextRef,
390    ) -> Pin<Box<dyn Stream<Item = Result<DoPutResponse>> + Send>> {
391        // Clone all necessary data to make it 'static
392        let catalog_manager = self.catalog_manager().clone();
393        let plugins = self.plugins.clone();
394        let inserter = self.inserter.clone();
395        let statement_executor = self.statement_executor.clone();
396        let ctx = ctx.clone();
397        let mut table_ref: Option<TableRef> = None;
398        let mut table_checked = false;
399
400        Box::pin(try_stream! {
401            // Process each request in the stream
402            while let Some(request_result) = stream.next().await {
403                let request = request_result.map_err(|e| {
404                    let error_msg = format!("Stream error: {:?}", e);
405                    IncompleteGrpcRequestSnafu { err_msg: error_msg }.build()
406                })?;
407
408                // Resolve table and check permissions on first RecordBatch (after schema is received)
409                if !table_checked {
410                    let table_name = &request.table_name;
411
412                    plugins
413                        .get::<PermissionCheckerRef>()
414                        .as_ref()
415                        .check_permission(
416                            ctx.current_user(),
417                            PermissionReq::BulkInsert {
418                                catalog: &table_name.catalog_name,
419                                schema: &table_name.schema_name,
420                                table: &table_name.table_name,
421                            },
422                        )
423                        .context(PermissionSnafu)?;
424
425                    // Resolve table reference
426                    let table = catalog_manager
427                        .table(
428                            &table_name.catalog_name,
429                            &table_name.schema_name,
430                            &table_name.table_name,
431                            None,
432                        )
433                        .await
434                        .context(CatalogSnafu)?
435                        .with_context(|| TableNotFoundSnafu {
436                            table_name: table_name.to_string(),
437                        })?;
438
439                    // Check permissions for the table
440                    let interceptor_ref = plugins.get::<GrpcQueryInterceptorRef<Error>>();
441                    let interceptor = interceptor_ref.as_ref();
442                    interceptor.pre_bulk_insert(table.clone(), ctx.clone())?;
443
444                    // The client sends the schema lazily after the stream handshake.
445                    // Reconcile it once, before writing the first batch.
446                    table_ref = Some(
447                        inserter
448                            .ensure_bulk_insert_schema(
449                                table,
450                                request.record_batch.schema().as_ref(),
451                                &ctx,
452                                statement_executor.as_ref(),
453                            )
454                            .await
455                            .context(TableOperationSnafu)?,
456                    );
457
458                    table_checked = true;
459                }
460
461                let request_id = request.request_id;
462                let start = Instant::now();
463                let rows = inserter
464                    .handle_bulk_insert(
465                        table_ref.clone().unwrap(),
466                        request.flight_data,
467                        request.record_batch,
468                        request.schema_bytes,
469                        ctx.skip_wal(),
470                        ctx.channel(),
471                    )
472                    .await
473                    .context(TableOperationSnafu)?;
474                let elapsed_secs = start.elapsed().as_secs_f64();
475                yield DoPutResponse::new(request_id, rows, elapsed_secs);
476            }
477        })
478    }
479
480    async fn handle_insert_plan(
481        &self,
482        insert: InsertIntoPlan,
483        ctx: QueryContextRef,
484    ) -> Result<Output> {
485        let timer = GRPC_HANDLE_PLAN_ELAPSED.start_timer();
486        let table_name = insert.table_name.context(IncompleteGrpcRequestSnafu {
487            err_msg: "'table_name' is absent in InsertIntoPlan",
488        })?;
489
490        // use dummy catalog to provide table
491        let plan_decoder = self
492            .query_engine()
493            .engine_context(ctx.clone())
494            .new_plan_decoder()
495            .context(PlanStatementSnafu)?;
496
497        let dummy_catalog_list = Arc::new(
498            catalog::table_source::dummy_catalog::DummyCatalogList::new_with_query_ctx(
499                self.catalog_manager().clone(),
500                ctx.clone(),
501            ),
502        );
503
504        // no optimize yet since we still need to add stuff
505        let logical_plan = plan_decoder
506            .decode(
507                bytes::Bytes::from(insert.logical_plan),
508                dummy_catalog_list,
509                false,
510            )
511            .await
512            .context(SubstraitDecodeLogicalPlanSnafu)?;
513
514        let table = self
515            .catalog_manager()
516            .table(
517                &table_name.catalog_name,
518                &table_name.schema_name,
519                &table_name.table_name,
520                None,
521            )
522            .await
523            .context(CatalogSnafu)?
524            .with_context(|| TableNotFoundSnafu {
525                table_name: [
526                    table_name.catalog_name.clone(),
527                    table_name.schema_name.clone(),
528                    table_name.table_name.clone(),
529                ]
530                .join("."),
531            })?;
532        let table_provider = Arc::new(DfTableProviderAdapter::new(table));
533        let table_source = Arc::new(DefaultTableSource::new(table_provider));
534
535        let insert_into = add_insert_to_logical_plan(table_name, table_source, logical_plan)
536            .context(SubstraitDecodeLogicalPlanSnafu)?;
537
538        let engine_ctx = self.query_engine().engine_context(ctx.clone());
539        let state = engine_ctx.state();
540        // Analyze the plan
541        let analyzed_plan = state
542            .analyzer()
543            .execute_and_check(insert_into, state.config_options(), |_, _| {})
544            .context(DataFusionSnafu)?;
545
546        // Optimize the plan
547        let optimized_plan = state.optimize(&analyzed_plan).context(DataFusionSnafu)?;
548
549        let output = self
550            .do_exec_plan_inner(optimized_plan, None, ctx.clone())
551            .await?;
552
553        Ok(attach_timer(output, timer))
554    }
555    #[tracing::instrument(skip_all)]
556    pub async fn handle_inserts(
557        &self,
558        requests: InsertRequests,
559        ctx: QueryContextRef,
560    ) -> Result<Output> {
561        self.inserter
562            .handle_column_inserts(requests, ctx, self.statement_executor.as_ref())
563            .await
564            .context(TableOperationSnafu)
565    }
566
567    #[tracing::instrument(skip_all)]
568    pub async fn handle_row_inserts(
569        &self,
570        requests: RowInsertRequests,
571        ctx: QueryContextRef,
572        accommodate_existing_schema: bool,
573        is_single_value: bool,
574    ) -> Result<Output> {
575        self.inserter
576            .handle_row_inserts(
577                requests,
578                ctx,
579                self.statement_executor.as_ref(),
580                accommodate_existing_schema,
581                is_single_value,
582            )
583            .await
584            .context(TableOperationSnafu)
585    }
586
587    #[tracing::instrument(skip_all)]
588    pub async fn handle_influx_row_inserts(
589        &self,
590        requests: RowInsertRequests,
591        ctx: QueryContextRef,
592    ) -> Result<Output> {
593        self.inserter
594            .handle_last_non_null_inserts(
595                requests,
596                ctx,
597                self.statement_executor.as_ref(),
598                true,
599                // Influx protocol may writes multiple fields (values).
600                false,
601            )
602            .await
603            .context(TableOperationSnafu)
604    }
605
606    #[tracing::instrument(skip_all)]
607    pub async fn handle_metric_row_inserts(
608        &self,
609        requests: RowInsertRequests,
610        ctx: QueryContextRef,
611        physical_table: String,
612    ) -> Result<Output> {
613        self.inserter
614            .handle_metric_row_inserts(requests, ctx, &self.statement_executor, physical_table)
615            .await
616            .context(TableOperationSnafu)
617    }
618
619    #[tracing::instrument(skip_all)]
620    pub async fn handle_deletes(
621        &self,
622        requests: DeleteRequests,
623        ctx: QueryContextRef,
624    ) -> Result<Output> {
625        self.deleter
626            .handle_column_deletes(requests, ctx)
627            .await
628            .context(TableOperationSnafu)
629    }
630
631    #[tracing::instrument(skip_all)]
632    pub async fn handle_row_deletes(
633        &self,
634        requests: RowDeleteRequests,
635        ctx: QueryContextRef,
636    ) -> Result<Output> {
637        self.deleter
638            .handle_row_deletes(requests, ctx)
639            .await
640            .context(TableOperationSnafu)
641    }
642}
frontend/instance/grpc.rs

frontend/instance/
grpc.rs
