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

frontend/instance/
grpc.rs
