frontend/instance/
opentsdb.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 async_trait::async_trait;
18use auth::{
19    OPENTSDB_WRITE, PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget,
20    PermissionTableTargets,
21};
22use common_error::ext::BoxedError;
23use common_telemetry::tracing;
24use servers::error::{AuthSnafu, ExecuteGrpcQuerySnafu, Result as ServerResult};
25use servers::opentsdb::codec::DataPoint;
26use servers::opentsdb::data_point_to_grpc_row_insert_requests;
27use servers::query_handler::OpentsdbProtocolHandler;
28use session::context::QueryContextRef;
29use snafu::prelude::*;
30use table::requests::{SEMANTIC_SIGNAL_TYPE, SEMANTIC_SOURCE, SIGNAL_TYPE_METRIC, SOURCE_OPENTSDB};
31
32use crate::instance::Instance;
33
34fn permission_targets(data_points: &[DataPoint], ctx: &QueryContextRef) -> PermissionTableTargets {
35    let catalog = ctx.current_catalog();
36    let schema = ctx.current_schema();
37    PermissionTableTargets::resolved(
38        data_points
39            .iter()
40            .map(|data_point| PermissionTableTarget::new(catalog, &schema, data_point.metric()))
41            .collect(),
42    )
43}
44
45#[async_trait]
46impl OpentsdbProtocolHandler for Instance {
47    async fn preflight(&self, data_points: &[DataPoint], ctx: QueryContextRef) -> ServerResult<()> {
48        self.check_table_permission(
49            &ctx,
50            PermissionReq::Action(OPENTSDB_WRITE),
51            permission_targets(data_points, &ctx),
52        )
53        .context(AuthSnafu)?;
54        Ok(())
55    }
56
57    #[tracing::instrument(skip_all, fields(protocol = "opentsdb"))]
58    async fn exec(&self, data_points: Vec<DataPoint>, ctx: QueryContextRef) -> ServerResult<usize> {
59        // Keep diagnostic per-point errors independent of other batched writes.
60        let mut ctx = ctx.fork();
61        ctx.set_batching_enabled(false);
62        self.execute_opentsdb_write(data_points, Arc::new(ctx))
63            .await
64    }
65
66    #[tracing::instrument(skip_all, fields(protocol = "opentsdb"))]
67    async fn exec_batch(
68        &self,
69        data_points: Vec<DataPoint>,
70        ctx: QueryContextRef,
71    ) -> ServerResult<usize> {
72        self.execute_opentsdb_write(data_points, ctx).await
73    }
74}
75
76impl Instance {
77    async fn execute_opentsdb_write(
78        &self,
79        data_points: Vec<DataPoint>,
80        ctx: QueryContextRef,
81    ) -> ServerResult<usize> {
82        self.plugins
83            .get::<PermissionCheckerRef>()
84            .as_ref()
85            .check_permission(ctx.current_user(), PermissionReq::Action(OPENTSDB_WRITE))
86            .context(AuthSnafu)?;
87
88        let (requests, _) = data_point_to_grpc_row_insert_requests(data_points)?;
89        self.check_row_insert_permission(&requests, &ctx, PermissionReq::Action(OPENTSDB_WRITE))
90            .context(AuthSnafu)?;
91
92        let ctx = {
93            let mut c = (*ctx).clone();
94            c.set_extension(SEMANTIC_SIGNAL_TYPE, SIGNAL_TYPE_METRIC);
95            c.set_extension(SEMANTIC_SOURCE, SOURCE_OPENTSDB);
96            Arc::new(c)
97        };
98
99        // OpenTSDB is single value.
100        let output = self
101            .handle_row_inserts(requests, ctx, true, true)
102            .await
103            .map_err(BoxedError::new)
104            .context(ExecuteGrpcQuerySnafu)?;
105
106        Ok(match output.data {
107            common_query::OutputData::AffectedRows(rows) => rows,
108            _ => unreachable!(),
109        })
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use session::context::QueryContext;
116
117    use crate::instance::opentsdb::*;
118
119    #[test]
120    fn test_permission_targets_do_not_require_row_conversion() {
121        let data_points = [
122            DataPoint::new("cpu".to_string(), 0, 1.0, vec![]),
123            DataPoint::new(
124                "mem".to_string(),
125                0,
126                1.0,
127                vec![("greptime_value".to_string(), "tag".to_string())],
128            ),
129        ];
130        assert!(data_point_to_grpc_row_insert_requests(data_points.to_vec()).is_err());
131
132        assert_eq!(
133            PermissionTableTargets::Resolved(vec![
134                PermissionTableTarget::new("greptime", "public", "cpu"),
135                PermissionTableTarget::new("greptime", "public", "mem"),
136            ]),
137            permission_targets(&data_points, &QueryContext::arc())
138        );
139    }
140}
frontend/instance/opentsdb.rs

frontend/instance/
opentsdb.rs
