servers/
query_handler.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
15//! All query handler traits for various request protocols, like SQL or GRPC.
16//!
17//! Instance that wishes to support certain request protocol, just implement the corresponding
18//! trait, the Server will handle codec for you.
19//!
20//! Note:
21//! Query handlers are not confined to only handle read requests, they are expecting to handle
22//! write requests too. So the "query" here not might seem ambiguity. However, "query" has been
23//! used as some kind of "convention", it's the "Q" in "SQL". So we might better stick to the
24//! word "query".
25
26pub mod grpc;
27pub mod sql;
28
29use std::collections::HashMap;
30use std::sync::Arc;
31
32use api::prom_store::remote::ReadRequest;
33use api::v1::RowInsertRequests;
34use async_trait::async_trait;
35use catalog::CatalogManager;
36use common_query::Output;
37use datatypes::timestamp::TimestampNanosecond;
38use headers::HeaderValue;
39use log_query::LogQuery;
40use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest;
41use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
42use operator::insert::PreparedMetricInserts;
43use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest;
44use pipeline::{
45    GreptimePipelineParams, Pipeline, PipelineInfo, PipelineOptions, PipelineVersion, PipelineWay,
46};
47use serde_json::Value;
48use session::context::{QueryContext, QueryContextRef};
49use table::metadata::TableInfoRef;
50
51#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
52pub struct DashboardDefinition {
53    pub name: String,
54    pub definition: String,
55}
56
57use crate::error::Result;
58use crate::http::jaeger::QueryTraceParams;
59use crate::influxdb::InfluxdbRequest;
60use crate::opentsdb::codec::DataPoint;
61use crate::request_memory_limiter::ServerMemoryLimiter;
62pub type OpentsdbProtocolHandlerRef = Arc<dyn OpentsdbProtocolHandler + Send + Sync>;
63pub type InfluxdbLineProtocolHandlerRef = Arc<dyn InfluxdbLineProtocolHandler + Send + Sync>;
64pub type PromStoreProtocolHandlerRef = Arc<dyn PromStoreProtocolHandler + Send + Sync>;
65pub type OpenTelemetryProtocolHandlerRef = Arc<dyn OpenTelemetryProtocolHandler + Send + Sync>;
66pub type PipelineHandlerRef = Arc<dyn PipelineHandler + Send + Sync>;
67pub type LogQueryHandlerRef = Arc<dyn LogQueryHandler + Send + Sync>;
68pub type JaegerQueryHandlerRef = Arc<dyn JaegerQueryHandler + Send + Sync>;
69
70#[derive(Debug, Default, Clone)]
71pub struct TraceIngestOutcome {
72    pub write_cost: usize,
73    pub accepted_spans: usize,
74    pub rejected_spans: usize,
75    pub error_message: Option<String>,
76}
77
78/// Result of ingesting one OTLP metrics request or Arrow batch.
79#[derive(Debug, Default, Clone)]
80pub struct MetricsIngestOutcome {
81    pub write_cost: usize,
82    pub accepted_data_points: i64,
83    pub rejected_data_points: i64,
84    pub error_message: Option<String>,
85}
86
87#[async_trait]
88pub trait InfluxdbLineProtocolHandler {
89    async fn exec(&self, request: InfluxdbRequest, ctx: QueryContextRef) -> Result<Output>;
90}
91
92#[async_trait]
93pub trait OpentsdbProtocolHandler {
94    /// Checks all points in one external request before per-point debug execution.
95    async fn preflight(&self, data_points: &[DataPoint], ctx: QueryContextRef) -> Result<()>;
96
97    async fn exec(&self, data_points: Vec<DataPoint>, ctx: QueryContextRef) -> Result<usize>;
98
99    /// Executes an ordinary HTTP put with optional batching. Debug callers
100    /// retain [`Self::exec`]; the frontend clears HTTP batching selection
101    /// there so diagnostic requests preserve direct, per-point error attribution.
102    async fn exec_batch(&self, data_points: Vec<DataPoint>, ctx: QueryContextRef) -> Result<usize> {
103        self.exec(data_points, ctx).await
104    }
105}
106
107pub struct PromStoreResponse {
108    pub content_type: HeaderValue,
109    pub content_encoding: HeaderValue,
110    pub resp_metrics: HashMap<String, Value>,
111    pub body: Vec<u8>,
112}
113
114#[async_trait]
115pub trait PromStoreProtocolHandler {
116    /// Resolves metric destinations before source identities are attached to exemplars.
117    /// Requires pre-write checks and admission to have succeeded for the entire request.
118    async fn prepare_metric_write(
119        &self,
120        _request: RowInsertRequests,
121        _ctx: QueryContextRef,
122    ) -> Result<PreparedMetricInserts> {
123        crate::error::NotSupportedSnafu {
124            feat: "metric write preparation",
125        }
126        .fail()
127    }
128
129    /// Executes the same routed metric writes used to resolve exemplar source identities.
130    async fn write_prepared_metric(
131        &self,
132        _request: PreparedMetricInserts,
133        _ctx: QueryContextRef,
134    ) -> Result<Output> {
135        crate::error::NotSupportedSnafu {
136            feat: "prepared metric ingestion",
137        }
138        .fail()
139    }
140
141    /// Writes exemplars using table identities captured while preparing the source metrics.
142    async fn write_exemplars_with_sources(
143        &self,
144        request: RowInsertRequests,
145        ctx: QueryContextRef,
146        _sources: &[TableInfoRef],
147    ) -> Result<Output> {
148        self.write_exemplars_prepared(request, ctx).await
149    }
150
151    /// Authorizes source series, including inputs with no metric samples.
152    async fn check_exemplar_sources(
153        &self,
154        _sources: &std::collections::BTreeSet<String>,
155        _ctx: QueryContextRef,
156    ) -> Result<()> {
157        crate::error::NotSupportedSnafu {
158            feat: "exemplar source authorization",
159        }
160        .fail()
161    }
162
163    async fn write_exemplars_prepared(
164        &self,
165        _request: RowInsertRequests,
166        _ctx: QueryContextRef,
167    ) -> Result<Output> {
168        crate::error::NotSupportedSnafu {
169            feat: "exemplar ingestion",
170        }
171        .fail()
172    }
173
174    /// Runs pre-write checks/hooks for prometheus remote write requests.
175    async fn pre_write(&self, request: &RowInsertRequests, ctx: QueryContextRef) -> Result<()>;
176
177    /// Writes one batch after [`Self::pre_write`] has succeeded for the entire request.
178    async fn write_prepared(
179        &self,
180        request: RowInsertRequests,
181        ctx: QueryContextRef,
182        with_metric_engine: bool,
183    ) -> Result<Output>;
184
185    /// Handling prometheus remote write requests
186    async fn write(
187        &self,
188        request: RowInsertRequests,
189        ctx: QueryContextRef,
190        with_metric_engine: bool,
191    ) -> Result<Output>;
192
193    /// Checks every batch before writing any of them.
194    async fn write_all(
195        &self,
196        requests: Vec<(QueryContextRef, RowInsertRequests)>,
197        with_metric_engine: bool,
198    ) -> Result<Vec<Result<Output>>>;
199
200    /// Handling prometheus remote read requests
201    async fn read(&self, request: ReadRequest, ctx: QueryContextRef) -> Result<PromStoreResponse>;
202}
203
204#[async_trait]
205pub trait OpenTelemetryProtocolHandler: PipelineHandler {
206    /// Handles metrics, reserving row expansion against the transport's limiter.
207    /// The accompanying byte count includes permits already held by this request.
208    async fn metrics(
209        &self,
210        request: ExportMetricsServiceRequest,
211        ctx: QueryContextRef,
212        memory_limiter: Option<(ServerMemoryLimiter, u64)>,
213    ) -> Result<MetricsIngestOutcome>;
214
215    /// Handling opentelemetry traces request
216    async fn traces(
217        &self,
218        pipeline_handler: PipelineHandlerRef,
219        request: ExportTraceServiceRequest,
220        pipeline: PipelineWay,
221        pipeline_params: GreptimePipelineParams,
222        table_name: String,
223        ctx: QueryContextRef,
224    ) -> Result<TraceIngestOutcome>;
225
226    async fn logs(
227        &self,
228        pipeline_handler: PipelineHandlerRef,
229        request: ExportLogsServiceRequest,
230        pipeline: PipelineWay,
231        pipeline_params: GreptimePipelineParams,
232        table_name: String,
233        ctx: QueryContextRef,
234    ) -> Result<Vec<Output>>;
235}
236
237/// PipelineHandler is responsible for handling pipeline related requests.
238///
239/// The "Pipeline" is a series of transformations that can be applied to unstructured
240/// data like logs. This handler is responsible to manage pipelines and accept data for
241/// processing.
242///
243/// The pipeline is stored in the database and can be retrieved by its name.
244#[async_trait]
245pub trait PipelineHandler {
246    async fn insert(&self, input: RowInsertRequests, ctx: QueryContextRef) -> Result<Output>;
247
248    /// Checks every batch before inserting any of them.
249    async fn insert_all(
250        &self,
251        inputs: Vec<(QueryContextRef, RowInsertRequests)>,
252    ) -> Result<Vec<Result<Output>>>;
253
254    fn check_pipeline_query_permission(&self, query_ctx: &QueryContextRef) -> Result<()>;
255
256    /// Loads a compiled pipeline for execution.
257    ///
258    /// This intentionally does not check pipeline-query permission: users with
259    /// write-only permission can ingest through an existing pipeline. Inspection
260    /// and preview callers must check query permission first; ingestion enforces
261    /// write and table-target permissions separately.
262    async fn get_pipeline(
263        &self,
264        name: &str,
265        version: PipelineVersion,
266        query_ctx: QueryContextRef,
267    ) -> Result<Arc<Pipeline>>;
268
269    async fn insert_pipeline(
270        &self,
271        name: &str,
272        content_type: &str,
273        pipeline: &str,
274        query_ctx: QueryContextRef,
275    ) -> Result<PipelineInfo>;
276
277    async fn delete_pipeline(
278        &self,
279        name: &str,
280        version: PipelineVersion,
281        query_ctx: QueryContextRef,
282    ) -> Result<Option<()>>;
283
284    async fn get_table(
285        &self,
286        table: &str,
287        query_ctx: &QueryContext,
288    ) -> std::result::Result<Option<Arc<table::Table>>, catalog::error::Error>;
289
290    /// Returns the server's pipeline options.
291    fn pipeline_options(&self) -> &PipelineOptions;
292
293    //// Build a pipeline from a string.
294    fn build_pipeline(&self, pipeline: &str) -> Result<Pipeline>;
295
296    /// Get a original pipeline by name.
297    async fn get_pipeline_str(
298        &self,
299        name: &str,
300        version: PipelineVersion,
301        query_ctx: QueryContextRef,
302    ) -> Result<(String, TimestampNanosecond)>;
303}
304
305/// Handling dashboard as code CRUD
306pub type DashboardHandlerRef = Arc<dyn DashboardHandler + Send + Sync>;
307
308#[async_trait]
309pub trait DashboardHandler {
310    async fn save(&self, name: &str, definition: &str, ctx: QueryContextRef) -> Result<()>;
311
312    async fn list(&self, ctx: QueryContextRef) -> Result<Vec<DashboardDefinition>>;
313
314    async fn delete(&self, name: &str, ctx: QueryContextRef) -> Result<()>;
315}
316
317/// Handle log query requests.
318#[async_trait]
319pub trait LogQueryHandler {
320    /// Execute a log query.
321    async fn query(&self, query: LogQuery, ctx: QueryContextRef) -> Result<Output>;
322
323    /// Get catalog manager.
324    fn catalog_manager(&self, ctx: &QueryContext) -> Result<&dyn CatalogManager>;
325}
326
327/// Handle Jaeger query requests.
328#[async_trait]
329pub trait JaegerQueryHandler {
330    /// Get trace services. It's used for `/api/services` API.
331    async fn get_services(&self, ctx: QueryContextRef) -> Result<Output>;
332
333    /// Get Jaeger operations. It's used for `/api/operations` and `/api/services/{service_name}/operations` API.
334    async fn get_operations(
335        &self,
336        ctx: QueryContextRef,
337        service_name: &str,
338        span_kind: Option<&str>,
339    ) -> Result<Output>;
340
341    /// Retrieves a trace by its unique identifier.
342    ///
343    /// This method is used to handle requests to the `/api/traces/{trace_id}` endpoint.
344    /// It accepts optional `start_time` and `end_time` parameters in nanoseconds to filter the trace data within a specific time range.
345    async fn get_trace(
346        &self,
347        ctx: QueryContextRef,
348        trace_id: &str,
349        start_time: Option<i64>,
350        end_time: Option<i64>,
351        limit: Option<usize>,
352    ) -> Result<Output>;
353
354    /// Find traces by query params. It's used for `/api/traces` API.
355    async fn find_traces(
356        &self,
357        ctx: QueryContextRef,
358        query_params: QueryTraceParams,
359    ) -> Result<Output>;
360}
361
362/// Handler for the PostgreSQL `COPY ... FROM STDIN` (copy-in) sub-protocol.
363///
364/// The servers layer owns the wire protocol and data parsing; this trait is
365/// the boundary to catalog lookup and ingestion, implemented by the
366/// frontend instance.
367#[async_trait]
368pub trait CopyInHandler: Send + Sync {
369    /// Resolves the target table of a COPY and returns its metadata, or
370    /// `None` when the table does not exist. Implementations should also
371    /// enforce the write permission for the table here so unauthorized
372    /// copies fail before any data is transferred.
373    async fn copy_in_table(
374        &self,
375        catalog: &str,
376        schema: &str,
377        table: &str,
378        query_ctx: QueryContextRef,
379    ) -> Result<Option<TableInfoRef>>;
380
381    /// Inserts a parsed batch produced by a COPY stream. Returns the
382    /// number of affected rows.
383    async fn copy_in_insert(
384        &self,
385        requests: RowInsertRequests,
386        query_ctx: QueryContextRef,
387    ) -> Result<Output>;
388}
389
390pub type CopyInHandlerRef = Arc<dyn CopyInHandler + Send + Sync>;
servers/query_handler.rs

servers/
query_handler.rs
