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

servers/
query_handler.rs
