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

servers/
interceptor.rs
