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

frontend/instance/
grpc.rs
