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

servers/
query_handler.rs
