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

frontend/instance/
influxdb.rs
