frontend/instance/
influxdb.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::sync::Arc;
16
17use api::v1::value::ValueData;
18use api::v1::{ColumnDataType, RowInsertRequests, SemanticType};
19use async_trait::async_trait;
20use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq};
21use catalog::CatalogManagerRef;
22use client::Output;
23use common_error::ext::BoxedError;
24use common_time::Timestamp;
25use common_time::timestamp::TimeUnit;
26use servers::error::{AuthSnafu, Error, TimestampOverflowSnafu, UnexpectedResultSnafu};
27use servers::influxdb::InfluxdbRequest;
28use servers::interceptor::{LineProtocolInterceptor, LineProtocolInterceptorRef};
29use servers::query_handler::InfluxdbLineProtocolHandler;
30use session::context::QueryContextRef;
31use snafu::{OptionExt, ResultExt};
32use store_api::mito_engine_options::MERGE_MODE_KEY;
33use table::requests::{SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, SOURCE_INFLUXDB};
34
35use crate::instance::Instance;
36use crate::service_config::influxdb::InfluxdbMergeMode;
37
38fn ctx_with_default_merge_mode(
39    ctx: QueryContextRef,
40    default_merge_mode: InfluxdbMergeMode,
41) -> QueryContextRef {
42    if ctx.extension(MERGE_MODE_KEY).is_none()
43        && default_merge_mode != InfluxdbMergeMode::LastNonNull
44    {
45        let mut ctx = (*ctx).clone();
46        ctx.set_extension(MERGE_MODE_KEY, default_merge_mode.as_str());
47        Arc::new(ctx)
48    } else {
49        ctx
50    }
51}
52
53#[async_trait]
54impl InfluxdbLineProtocolHandler for Instance {
55    async fn exec(
56        &self,
57        request: InfluxdbRequest,
58        ctx: QueryContextRef,
59    ) -> servers::error::Result<Output> {
60        self.plugins
61            .get::<PermissionCheckerRef>()
62            .as_ref()
63            .check_permission(ctx.current_user(), PermissionReq::LineProtocol)
64            .context(AuthSnafu)?;
65
66        let interceptor_ref = self.plugins.get::<LineProtocolInterceptorRef<Error>>();
67        interceptor_ref.pre_execute(&request.lines, ctx.clone())?;
68
69        let requests = request.try_into()?;
70
71        let aligner = InfluxdbLineTimestampAligner {
72            catalog_manager: self.catalog_manager(),
73        };
74        let requests = aligner.align_timestamps(requests, &ctx).await?;
75
76        let requests = interceptor_ref
77            .post_lines_conversion(requests, ctx.clone())
78            .await?;
79
80        let ctx = ctx_with_default_merge_mode(ctx, self.influxdb_default_merge_mode);
81        let ctx = {
82            let mut c = (*ctx).clone();
83            c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
84            c.set_extension(SEMANTIC_SOURCE, SOURCE_INFLUXDB);
85            Arc::new(c)
86        };
87
88        self.handle_influx_row_inserts(requests, ctx)
89            .await
90            .map_err(BoxedError::new)
91            .context(servers::error::ExecuteGrpcQuerySnafu)
92    }
93}
94
95/// Align the timestamp precisions in Influxdb lines (after they are converted to the GRPC row
96/// inserts) to the time index columns' time units of the created tables (if there are any).
97struct InfluxdbLineTimestampAligner<'a> {
98    catalog_manager: &'a CatalogManagerRef,
99}
100
101impl InfluxdbLineTimestampAligner<'_> {
102    async fn align_timestamps(
103        &self,
104        requests: RowInsertRequests,
105        query_context: &QueryContextRef,
106    ) -> servers::error::Result<RowInsertRequests> {
107        let mut inserts = requests.inserts;
108        for insert in inserts.iter_mut() {
109            let Some(rows) = &mut insert.rows else {
110                continue;
111            };
112
113            let Some(target_time_unit) = self
114                .catalog_manager
115                .table(
116                    query_context.current_catalog(),
117                    &query_context.current_schema(),
118                    &insert.table_name,
119                    Some(query_context),
120                )
121                .await?
122                .map(|x| x.schema())
123                .and_then(|schema| {
124                    schema.timestamp_column().map(|col| {
125                        col.data_type
126                            .as_timestamp()
127                            .expect("Time index column is not of timestamp type?!")
128                            .unit()
129                    })
130                })
131            else {
132                continue;
133            };
134
135            let target_timestamp_type = match target_time_unit {
136                TimeUnit::Second => ColumnDataType::TimestampSecond,
137                TimeUnit::Millisecond => ColumnDataType::TimestampMillisecond,
138                TimeUnit::Microsecond => ColumnDataType::TimestampMicrosecond,
139                TimeUnit::Nanosecond => ColumnDataType::TimestampNanosecond,
140            };
141            let Some(to_be_aligned) = rows.schema.iter().enumerate().find_map(|(i, x)| {
142                if x.semantic_type() == SemanticType::Timestamp
143                    && x.datatype() != target_timestamp_type
144                {
145                    Some(i)
146                } else {
147                    None
148                }
149            }) else {
150                continue;
151            };
152
153            // Indexing safety: `to_be_aligned` is guaranteed to be a valid index because it's got
154            // from "enumerate" the schema vector above.
155            rows.schema[to_be_aligned].datatype = target_timestamp_type as i32;
156
157            for row in rows.rows.iter_mut() {
158                let Some(time_value) = row
159                    .values
160                    .get_mut(to_be_aligned)
161                    .and_then(|x| x.value_data.as_mut())
162                else {
163                    continue;
164                };
165                *time_value = align_time_unit(time_value, target_time_unit)?;
166            }
167        }
168        Ok(RowInsertRequests { inserts })
169    }
170}
171
172fn align_time_unit(value: &ValueData, target: TimeUnit) -> servers::error::Result<ValueData> {
173    let timestamp = match value {
174        ValueData::TimestampSecondValue(x) => Timestamp::new_second(*x),
175        ValueData::TimestampMillisecondValue(x) => Timestamp::new_millisecond(*x),
176        ValueData::TimestampMicrosecondValue(x) => Timestamp::new_microsecond(*x),
177        ValueData::TimestampNanosecondValue(x) => Timestamp::new_nanosecond(*x),
178        _ => {
179            return UnexpectedResultSnafu {
180                reason: format!("Timestamp value '{:?}' is not of timestamp type!", value),
181            }
182            .fail();
183        }
184    };
185
186    let timestamp = timestamp
187        .convert_to(target)
188        .with_context(|| TimestampOverflowSnafu {
189            error: format!("{:?} convert to {}", timestamp, target),
190        })?;
191
192    Ok(match target {
193        TimeUnit::Second => ValueData::TimestampSecondValue(timestamp.value()),
194        TimeUnit::Millisecond => ValueData::TimestampMillisecondValue(timestamp.value()),
195        TimeUnit::Microsecond => ValueData::TimestampMicrosecondValue(timestamp.value()),
196        TimeUnit::Nanosecond => ValueData::TimestampNanosecondValue(timestamp.value()),
197    })
198}
199
200#[cfg(test)]
201mod tests {
202    use session::context::QueryContext;
203    use store_api::mito_engine_options::MERGE_MODE_KEY;
204
205    use super::*;
206    use crate::service_config::influxdb::InfluxdbMergeMode;
207
208    #[test]
209    fn test_influxdb_default_merge_mode_reuses_default_context() {
210        let ctx = QueryContext::arc();
211        let actual = ctx_with_default_merge_mode(ctx.clone(), InfluxdbMergeMode::LastNonNull);
212
213        assert!(Arc::ptr_eq(&ctx, &actual));
214        assert!(actual.extension(MERGE_MODE_KEY).is_none());
215    }
216
217    #[test]
218    fn test_influxdb_non_default_merge_mode_sets_extension() {
219        let ctx = QueryContext::arc();
220        let actual = ctx_with_default_merge_mode(ctx.clone(), InfluxdbMergeMode::LastRow);
221
222        assert!(!Arc::ptr_eq(&ctx, &actual));
223        assert_eq!(Some("last_row"), actual.extension(MERGE_MODE_KEY));
224    }
225
226    #[test]
227    fn test_influxdb_explicit_merge_mode_keeps_context() {
228        let mut ctx = QueryContext::arc();
229        Arc::get_mut(&mut ctx)
230            .unwrap()
231            .set_extension(MERGE_MODE_KEY, "last_row");
232
233        let actual = ctx_with_default_merge_mode(ctx.clone(), InfluxdbMergeMode::LastNonNull);
234
235        assert!(Arc::ptr_eq(&ctx, &actual));
236        assert_eq!(Some("last_row"), actual.extension(MERGE_MODE_KEY));
237    }
238}
frontend/instance/influxdb.rs

frontend/instance/
influxdb.rs
