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

frontend/instance/
grpc.rs
