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

frontend/instance/
grpc.rs
