servers/
interceptor.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::borrow::Cow;
16use std::sync::Arc;
17
18use api::prom_store::remote::ReadRequest;
19use api::v1::RowInsertRequests;
20use api::v1::greptime_request::Request;
21use async_trait::async_trait;
22use common_error::ext::ErrorExt;
23use common_query::Output;
24use datafusion_expr::LogicalPlan;
25use log_query::LogQuery;
26use promql_parser::parser::Expr;
27use query::parser::PromQuery;
28use session::context::QueryContextRef;
29use sql::statements::statement::Statement;
30use table::TableRef;
31use table::table_reference::TableReference;
32use vrl::value::Value;
33
34use crate::exemplar::query::Query as ExemplarQuery;
35
36/// SqlQueryInterceptor can track life cycle of a sql query and customize or
37/// abort its execution at given point.
38pub trait SqlQueryInterceptor {
39    type Error: ErrorExt;
40
41    /// Called before a query string is parsed into sql statements.
42    /// The implementation is allowed to change the sql string if needed.
43    fn pre_parsing<'a>(
44        &self,
45        query: &'a str,
46        _query_ctx: QueryContextRef,
47    ) -> Result<Cow<'a, str>, Self::Error> {
48        Ok(Cow::Borrowed(query))
49    }
50
51    /// Called after sql is parsed into statements. This interceptor is called
52    /// on each statement and the implementation can alter the statement or
53    /// abort execution by raising an error.
54    fn post_parsing(
55        &self,
56        statements: Vec<Statement>,
57        _query_ctx: QueryContextRef,
58    ) -> Result<Vec<Statement>, Self::Error> {
59        Ok(statements)
60    }
61
62    /// Called before sql is actually executed. This hook is not called at the moment.
63    fn pre_execute(
64        &self,
65        _statement: Option<&Statement>,
66        _plan: Option<&LogicalPlan>,
67        _query_ctx: QueryContextRef,
68    ) -> Result<(), Self::Error> {
69        Ok(())
70    }
71
72    /// Called after execution finished. The implementation can modify the
73    /// output if needed.
74    fn post_execute(
75        &self,
76        output: Output,
77        _query_ctx: QueryContextRef,
78    ) -> Result<Output, Self::Error> {
79        Ok(output)
80    }
81}
82
83pub type SqlQueryInterceptorRef<E> =
84    Arc<dyn SqlQueryInterceptor<Error = E> + Send + Sync + 'static>;
85
86impl<E> SqlQueryInterceptor for Option<&SqlQueryInterceptorRef<E>>
87where
88    E: ErrorExt,
89{
90    type Error = E;
91
92    fn pre_parsing<'a>(
93        &self,
94        query: &'a str,
95        query_ctx: QueryContextRef,
96    ) -> Result<Cow<'a, str>, Self::Error> {
97        if let Some(this) = self {
98            this.pre_parsing(query, query_ctx)
99        } else {
100            Ok(Cow::Borrowed(query))
101        }
102    }
103
104    fn post_parsing(
105        &self,
106        statements: Vec<Statement>,
107        query_ctx: QueryContextRef,
108    ) -> Result<Vec<Statement>, Self::Error> {
109        if let Some(this) = self {
110            this.post_parsing(statements, query_ctx)
111        } else {
112            Ok(statements)
113        }
114    }
115
116    fn pre_execute(
117        &self,
118        statement: Option<&Statement>,
119        plan: Option<&LogicalPlan>,
120        query_ctx: QueryContextRef,
121    ) -> Result<(), Self::Error> {
122        if let Some(this) = self {
123            this.pre_execute(statement, plan, query_ctx)
124        } else {
125            Ok(())
126        }
127    }
128
129    fn post_execute(
130        &self,
131        output: Output,
132        query_ctx: QueryContextRef,
133    ) -> Result<Output, Self::Error> {
134        if let Some(this) = self {
135            this.post_execute(output, query_ctx)
136        } else {
137            Ok(output)
138        }
139    }
140}
141
142/// GrpcQueryInterceptor can track life cycle of a grpc request and customize or
143/// abort its execution at given point.
144pub trait GrpcQueryInterceptor {
145    type Error: ErrorExt;
146
147    /// Called before request is actually executed.
148    fn pre_execute(
149        &self,
150        _request: &Request,
151        _query_ctx: QueryContextRef,
152    ) -> Result<(), Self::Error> {
153        Ok(())
154    }
155
156    /// Called before bulk insert is executed.
157    fn pre_bulk_insert(
158        &self,
159        _table: TableRef,
160        _query_ctx: QueryContextRef,
161    ) -> Result<(), Self::Error> {
162        Ok(())
163    }
164
165    /// Called after execution finished. The implementation can modify the
166    /// output if needed.
167    fn post_execute(
168        &self,
169        output: Output,
170        _query_ctx: QueryContextRef,
171    ) -> Result<Output, Self::Error> {
172        Ok(output)
173    }
174}
175
176pub type GrpcQueryInterceptorRef<E> =
177    Arc<dyn GrpcQueryInterceptor<Error = E> + Send + Sync + 'static>;
178
179impl<E> GrpcQueryInterceptor for Option<&GrpcQueryInterceptorRef<E>>
180where
181    E: ErrorExt,
182{
183    type Error = E;
184
185    fn pre_execute(
186        &self,
187        _request: &Request,
188        _query_ctx: QueryContextRef,
189    ) -> Result<(), Self::Error> {
190        if let Some(this) = self {
191            this.pre_execute(_request, _query_ctx)
192        } else {
193            Ok(())
194        }
195    }
196
197    fn pre_bulk_insert(
198        &self,
199        _table: TableRef,
200        _query_ctx: QueryContextRef,
201    ) -> Result<(), Self::Error> {
202        if let Some(this) = self {
203            this.pre_bulk_insert(_table, _query_ctx)
204        } else {
205            Ok(())
206        }
207    }
208
209    fn post_execute(
210        &self,
211        output: Output,
212        _query_ctx: QueryContextRef,
213    ) -> Result<Output, Self::Error> {
214        if let Some(this) = self {
215            this.post_execute(output, _query_ctx)
216        } else {
217            Ok(output)
218        }
219    }
220}
221
222/// PromQueryInterceptor can track life cycle of a prometheus request and customize or
223/// abort its execution at given point.
224pub trait PromQueryInterceptor {
225    type Error: ErrorExt;
226
227    /// Called before request is actually executed.
228    fn pre_execute(
229        &self,
230        _query: &PromQuery,
231        _expr: &Expr,
232        _plan: Option<&LogicalPlan>,
233        _query_ctx: QueryContextRef,
234    ) -> Result<(), Self::Error> {
235        Ok(())
236    }
237
238    /// Called after authorization, before each resolved exemplar table is scanned.
239    /// Returning an error aborts the query. `query_ctx` uses the table's schema.
240    fn pre_execute_exemplars(
241        &self,
242        _query: &ExemplarQuery,
243        _table: TableReference<'_>,
244        _query_ctx: QueryContextRef,
245    ) -> Result<(), Self::Error> {
246        Ok(())
247    }
248
249    /// Called after execution finished. The implementation can modify the
250    /// output if needed.
251    fn post_execute(
252        &self,
253        output: Output,
254        _query_ctx: QueryContextRef,
255    ) -> Result<Output, Self::Error> {
256        Ok(output)
257    }
258}
259
260pub type PromQueryInterceptorRef<E> =
261    Arc<dyn PromQueryInterceptor<Error = E> + Send + Sync + 'static>;
262
263impl<E> PromQueryInterceptor for Option<PromQueryInterceptorRef<E>>
264where
265    E: ErrorExt,
266{
267    type Error = E;
268
269    fn pre_execute(
270        &self,
271        query: &PromQuery,
272        expr: &Expr,
273        plan: Option<&LogicalPlan>,
274        query_ctx: QueryContextRef,
275    ) -> Result<(), Self::Error> {
276        if let Some(this) = self {
277            this.pre_execute(query, expr, plan, query_ctx)
278        } else {
279            Ok(())
280        }
281    }
282
283    fn post_execute(
284        &self,
285        output: Output,
286        query_ctx: QueryContextRef,
287    ) -> Result<Output, Self::Error> {
288        if let Some(this) = self {
289            this.post_execute(output, query_ctx)
290        } else {
291            Ok(output)
292        }
293    }
294
295    fn pre_execute_exemplars(
296        &self,
297        query: &ExemplarQuery,
298        table: TableReference<'_>,
299        query_ctx: QueryContextRef,
300    ) -> Result<(), Self::Error> {
301        if let Some(this) = self {
302            this.pre_execute_exemplars(query, table, query_ctx)
303        } else {
304            Ok(())
305        }
306    }
307}
308
309impl<E> PromQueryInterceptor for Option<&PromQueryInterceptorRef<E>>
310where
311    E: ErrorExt,
312{
313    type Error = E;
314
315    fn pre_execute(
316        &self,
317        query: &PromQuery,
318        expr: &Expr,
319        plan: Option<&LogicalPlan>,
320        query_ctx: QueryContextRef,
321    ) -> Result<(), Self::Error> {
322        if let Some(this) = self {
323            this.pre_execute(query, expr, plan, query_ctx)
324        } else {
325            Ok(())
326        }
327    }
328
329    fn post_execute(
330        &self,
331        output: Output,
332        query_ctx: QueryContextRef,
333    ) -> Result<Output, Self::Error> {
334        if let Some(this) = self {
335            this.post_execute(output, query_ctx)
336        } else {
337            Ok(output)
338        }
339    }
340
341    fn pre_execute_exemplars(
342        &self,
343        query: &ExemplarQuery,
344        table: TableReference<'_>,
345        query_ctx: QueryContextRef,
346    ) -> Result<(), Self::Error> {
347        if let Some(this) = self {
348            this.pre_execute_exemplars(query, table, query_ctx)
349        } else {
350            Ok(())
351        }
352    }
353}
354
355/// LineProtocolInterceptor can track life cycle of a line protocol request
356/// and customize or abort its execution at given point.
357#[async_trait]
358pub trait LineProtocolInterceptor {
359    type Error: ErrorExt;
360
361    fn pre_execute(&self, _line: &str, _query_ctx: QueryContextRef) -> Result<(), Self::Error> {
362        Ok(())
363    }
364
365    /// Called after the lines are converted to the [RowInsertRequests].
366    /// We can then modify the resulting requests if needed.
367    /// Typically used in some backward compatibility situation.
368    async fn post_lines_conversion(
369        &self,
370        requests: RowInsertRequests,
371        query_context: QueryContextRef,
372    ) -> Result<RowInsertRequests, Self::Error> {
373        let _ = query_context;
374        Ok(requests)
375    }
376}
377
378pub type LineProtocolInterceptorRef<E> =
379    Arc<dyn LineProtocolInterceptor<Error = E> + Send + Sync + 'static>;
380
381#[async_trait]
382impl<E: ErrorExt> LineProtocolInterceptor for Option<LineProtocolInterceptorRef<E>> {
383    type Error = E;
384
385    fn pre_execute(&self, line: &str, query_ctx: QueryContextRef) -> Result<(), Self::Error> {
386        if let Some(this) = self {
387            this.pre_execute(line, query_ctx)
388        } else {
389            Ok(())
390        }
391    }
392
393    async fn post_lines_conversion(
394        &self,
395        requests: RowInsertRequests,
396        query_context: QueryContextRef,
397    ) -> Result<RowInsertRequests, Self::Error> {
398        if let Some(this) = self {
399            this.post_lines_conversion(requests, query_context).await
400        } else {
401            Ok(requests)
402        }
403    }
404}
405
406/// OpenTelemetryProtocolInterceptor can track life cycle of an open telemetry protocol request
407/// and customize or abort its execution at given point.
408pub trait OpenTelemetryProtocolInterceptor {
409    type Error: ErrorExt;
410
411    fn pre_execute(&self, _query_ctx: QueryContextRef) -> Result<(), Self::Error> {
412        Ok(())
413    }
414}
415
416pub type OpenTelemetryProtocolInterceptorRef<E> =
417    Arc<dyn OpenTelemetryProtocolInterceptor<Error = E> + Send + Sync + 'static>;
418
419impl<E: ErrorExt> OpenTelemetryProtocolInterceptor
420    for Option<OpenTelemetryProtocolInterceptorRef<E>>
421{
422    type Error = E;
423
424    fn pre_execute(&self, query_ctx: QueryContextRef) -> Result<(), Self::Error> {
425        if let Some(this) = self {
426            this.pre_execute(query_ctx)
427        } else {
428            Ok(())
429        }
430    }
431}
432
433/// PromStoreProtocolInterceptor can track life cycle of a prom store request
434/// and customize or abort its execution at given point.
435pub trait PromStoreProtocolInterceptor {
436    type Error: ErrorExt;
437
438    fn pre_write(
439        &self,
440        _write_req: &RowInsertRequests,
441        _ctx: QueryContextRef,
442    ) -> Result<(), Self::Error> {
443        Ok(())
444    }
445
446    fn pre_read(&self, _read_req: &ReadRequest, _ctx: QueryContextRef) -> Result<(), Self::Error> {
447        Ok(())
448    }
449}
450
451pub type PromStoreProtocolInterceptorRef<E> =
452    Arc<dyn PromStoreProtocolInterceptor<Error = E> + Send + Sync + 'static>;
453
454impl<E: ErrorExt> PromStoreProtocolInterceptor for Option<PromStoreProtocolInterceptorRef<E>> {
455    type Error = E;
456
457    fn pre_write(
458        &self,
459        write_req: &RowInsertRequests,
460        ctx: QueryContextRef,
461    ) -> Result<(), Self::Error> {
462        if let Some(this) = self {
463            this.pre_write(write_req, ctx)
464        } else {
465            Ok(())
466        }
467    }
468
469    fn pre_read(&self, read_req: &ReadRequest, ctx: QueryContextRef) -> Result<(), Self::Error> {
470        if let Some(this) = self {
471            this.pre_read(read_req, ctx)
472        } else {
473            Ok(())
474        }
475    }
476}
477
478/// LogIngestInterceptor can track life cycle of a log ingestion request
479/// and customize or abort its execution at given point.
480pub trait LogIngestInterceptor {
481    type Error: ErrorExt;
482
483    /// Called before pipeline execution.
484    fn pre_pipeline(
485        &self,
486        values: Vec<Value>,
487        _query_ctx: QueryContextRef,
488    ) -> Result<Vec<Value>, Self::Error> {
489        Ok(values)
490    }
491
492    /// Called before insertion.
493    fn pre_ingest(
494        &self,
495        request: RowInsertRequests,
496        _query_ctx: QueryContextRef,
497    ) -> Result<RowInsertRequests, Self::Error> {
498        Ok(request)
499    }
500}
501
502pub type LogIngestInterceptorRef<E> =
503    Arc<dyn LogIngestInterceptor<Error = E> + Send + Sync + 'static>;
504
505impl<E> LogIngestInterceptor for Option<&LogIngestInterceptorRef<E>>
506where
507    E: ErrorExt,
508{
509    type Error = E;
510
511    fn pre_pipeline(
512        &self,
513        values: Vec<Value>,
514        query_ctx: QueryContextRef,
515    ) -> Result<Vec<Value>, Self::Error> {
516        if let Some(this) = self {
517            this.pre_pipeline(values, query_ctx)
518        } else {
519            Ok(values)
520        }
521    }
522
523    fn pre_ingest(
524        &self,
525        request: RowInsertRequests,
526        query_ctx: QueryContextRef,
527    ) -> Result<RowInsertRequests, Self::Error> {
528        if let Some(this) = self {
529            this.pre_ingest(request, query_ctx)
530        } else {
531            Ok(request)
532        }
533    }
534}
535
536/// LogQueryInterceptor can track life cycle of a log query request
537/// and customize or abort its execution at given point.
538pub trait LogQueryInterceptor {
539    type Error: ErrorExt;
540
541    /// Called before query is actually executed.
542    fn pre_query(&self, _query: &LogQuery, _query_ctx: QueryContextRef) -> Result<(), Self::Error> {
543        Ok(())
544    }
545
546    /// Called after execution finished. The implementation can modify the
547    /// output if needed.
548    fn post_query(
549        &self,
550        output: Output,
551        _query_ctx: QueryContextRef,
552    ) -> Result<Output, Self::Error> {
553        Ok(output)
554    }
555}
556
557pub type LogQueryInterceptorRef<E> =
558    Arc<dyn LogQueryInterceptor<Error = E> + Send + Sync + 'static>;
559
560impl<E> LogQueryInterceptor for Option<&LogQueryInterceptorRef<E>>
561where
562    E: ErrorExt,
563{
564    type Error = E;
565
566    fn pre_query(&self, query: &LogQuery, query_ctx: QueryContextRef) -> Result<(), Self::Error> {
567        if let Some(this) = self {
568            this.pre_query(query, query_ctx)
569        } else {
570            Ok(())
571        }
572    }
573
574    fn post_query(
575        &self,
576        output: Output,
577        query_ctx: QueryContextRef,
578    ) -> Result<Output, Self::Error> {
579        if let Some(this) = self {
580            this.post_query(output, query_ctx)
581        } else {
582            Ok(output)
583        }
584    }
585}
servers/interceptor.rs

servers/
interceptor.rs
