Skip to main content

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};
46
47#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
48pub struct DashboardDefinition {
49    pub name: String,
50    pub definition: String,
51}
52
53use crate::error::Result;
54use crate::http::jaeger::QueryTraceParams;
55use crate::influxdb::InfluxdbRequest;
56use crate::opentsdb::codec::DataPoint;
57use crate::prom_store::Metrics;
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#[async_trait]
75pub trait InfluxdbLineProtocolHandler {
76    /// A successful request will not return a response.
77    /// Only on error will the socket return a line of data.
78    async fn exec(&self, request: InfluxdbRequest, ctx: QueryContextRef) -> Result<Output>;
79}
80
81#[async_trait]
82pub trait OpentsdbProtocolHandler {
83    /// Checks all points in one external request before per-point debug execution.
84    async fn preflight(&self, data_points: &[DataPoint], ctx: QueryContextRef) -> Result<()>;
85
86    /// A successful request will not return a response.
87    /// Only on error will the socket return a line of data.
88    async fn exec(&self, data_points: Vec<DataPoint>, ctx: QueryContextRef) -> Result<usize>;
89}
90
91pub struct PromStoreResponse {
92    pub content_type: HeaderValue,
93    pub content_encoding: HeaderValue,
94    pub resp_metrics: HashMap<String, Value>,
95    pub body: Vec<u8>,
96}
97
98#[async_trait]
99pub trait PromStoreProtocolHandler {
100    /// Runs pre-write checks/hooks for prometheus remote write requests.
101    async fn pre_write(&self, request: &RowInsertRequests, ctx: QueryContextRef) -> Result<()>;
102
103    /// Writes one batch after [`Self::pre_write`] has succeeded for the entire request.
104    async fn write_prepared(
105        &self,
106        request: RowInsertRequests,
107        ctx: QueryContextRef,
108        with_metric_engine: bool,
109    ) -> Result<Output>;
110
111    /// Handling prometheus remote write requests
112    async fn write(
113        &self,
114        request: RowInsertRequests,
115        ctx: QueryContextRef,
116        with_metric_engine: bool,
117    ) -> Result<Output>;
118
119    /// Checks every batch before writing any of them.
120    async fn write_all(
121        &self,
122        requests: Vec<(QueryContextRef, RowInsertRequests)>,
123        with_metric_engine: bool,
124    ) -> Result<Vec<Result<Output>>>;
125
126    /// Handling prometheus remote read requests
127    async fn read(&self, request: ReadRequest, ctx: QueryContextRef) -> Result<PromStoreResponse>;
128    /// Handling push gateway requests
129    async fn ingest_metrics(&self, metrics: Metrics) -> Result<()>;
130}
131
132#[async_trait]
133pub trait OpenTelemetryProtocolHandler: PipelineHandler {
134    /// Handling opentelemetry metrics request
135    async fn metrics(
136        &self,
137        request: ExportMetricsServiceRequest,
138        ctx: QueryContextRef,
139    ) -> Result<Output>;
140
141    /// Handling opentelemetry traces request
142    async fn traces(
143        &self,
144        pipeline_handler: PipelineHandlerRef,
145        request: ExportTraceServiceRequest,
146        pipeline: PipelineWay,
147        pipeline_params: GreptimePipelineParams,
148        table_name: String,
149        ctx: QueryContextRef,
150    ) -> Result<TraceIngestOutcome>;
151
152    async fn logs(
153        &self,
154        pipeline_handler: PipelineHandlerRef,
155        request: ExportLogsServiceRequest,
156        pipeline: PipelineWay,
157        pipeline_params: GreptimePipelineParams,
158        table_name: String,
159        ctx: QueryContextRef,
160    ) -> Result<Vec<Output>>;
161}
162
163/// PipelineHandler is responsible for handling pipeline related requests.
164///
165/// The "Pipeline" is a series of transformations that can be applied to unstructured
166/// data like logs. This handler is responsible to manage pipelines and accept data for
167/// processing.
168///
169/// The pipeline is stored in the database and can be retrieved by its name.
170#[async_trait]
171pub trait PipelineHandler {
172    async fn insert(&self, input: RowInsertRequests, ctx: QueryContextRef) -> Result<Output>;
173
174    /// Checks every batch before inserting any of them.
175    async fn insert_all(
176        &self,
177        inputs: Vec<(QueryContextRef, RowInsertRequests)>,
178    ) -> Result<Vec<Result<Output>>>;
179
180    fn check_pipeline_query_permission(&self, query_ctx: &QueryContextRef) -> Result<()>;
181
182    /// Loads a compiled pipeline for execution.
183    ///
184    /// This intentionally does not check pipeline-query permission: users with
185    /// write-only permission can ingest through an existing pipeline. Inspection
186    /// and preview callers must check query permission first; ingestion enforces
187    /// write and table-target permissions separately.
188    async fn get_pipeline(
189        &self,
190        name: &str,
191        version: PipelineVersion,
192        query_ctx: QueryContextRef,
193    ) -> Result<Arc<Pipeline>>;
194
195    async fn insert_pipeline(
196        &self,
197        name: &str,
198        content_type: &str,
199        pipeline: &str,
200        query_ctx: QueryContextRef,
201    ) -> Result<PipelineInfo>;
202
203    async fn delete_pipeline(
204        &self,
205        name: &str,
206        version: PipelineVersion,
207        query_ctx: QueryContextRef,
208    ) -> Result<Option<()>>;
209
210    async fn get_table(
211        &self,
212        table: &str,
213        query_ctx: &QueryContext,
214    ) -> std::result::Result<Option<Arc<table::Table>>, catalog::error::Error>;
215
216    //// Build a pipeline from a string.
217    fn build_pipeline(&self, pipeline: &str) -> Result<Pipeline>;
218
219    /// Get a original pipeline by name.
220    async fn get_pipeline_str(
221        &self,
222        name: &str,
223        version: PipelineVersion,
224        query_ctx: QueryContextRef,
225    ) -> Result<(String, TimestampNanosecond)>;
226}
227
228/// Handling dashboard as code CRUD
229pub type DashboardHandlerRef = Arc<dyn DashboardHandler + Send + Sync>;
230
231#[async_trait]
232pub trait DashboardHandler {
233    async fn save(&self, name: &str, definition: &str, ctx: QueryContextRef) -> Result<()>;
234
235    async fn list(&self, ctx: QueryContextRef) -> Result<Vec<DashboardDefinition>>;
236
237    async fn delete(&self, name: &str, ctx: QueryContextRef) -> Result<()>;
238}
239
240/// Handle log query requests.
241#[async_trait]
242pub trait LogQueryHandler {
243    /// Execute a log query.
244    async fn query(&self, query: LogQuery, ctx: QueryContextRef) -> Result<Output>;
245
246    /// Get catalog manager.
247    fn catalog_manager(&self, ctx: &QueryContext) -> Result<&dyn CatalogManager>;
248}
249
250/// Handle Jaeger query requests.
251#[async_trait]
252pub trait JaegerQueryHandler {
253    /// Get trace services. It's used for `/api/services` API.
254    async fn get_services(&self, ctx: QueryContextRef) -> Result<Output>;
255
256    /// Get Jaeger operations. It's used for `/api/operations` and `/api/services/{service_name}/operations` API.
257    async fn get_operations(
258        &self,
259        ctx: QueryContextRef,
260        service_name: &str,
261        span_kind: Option<&str>,
262    ) -> Result<Output>;
263
264    /// Retrieves a trace by its unique identifier.
265    ///
266    /// This method is used to handle requests to the `/api/traces/{trace_id}` endpoint.
267    /// It accepts optional `start_time` and `end_time` parameters in nanoseconds to filter the trace data within a specific time range.
268    async fn get_trace(
269        &self,
270        ctx: QueryContextRef,
271        trace_id: &str,
272        start_time: Option<i64>,
273        end_time: Option<i64>,
274        limit: Option<usize>,
275    ) -> Result<Output>;
276
277    /// Find traces by query params. It's used for `/api/traces` API.
278    async fn find_traces(
279        &self,
280        ctx: QueryContextRef,
281        query_params: QueryTraceParams,
282    ) -> Result<Output>;
283}