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

frontend/instance/
grpc.rs
