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

frontend/instance/
influxdb.rs
