operator/
utils.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, RwLock};
16
17use common_event_recorder::ProcedureEventInput;
18use common_meta::procedure_executor::ExecutorContext;
19use common_meta::rpc::ddl::{ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TriggerReason};
20use common_time::Timezone;
21use session::context::{QueryContextBuilder, QueryContextRef};
22use snafu::ResultExt;
23
24use crate::error::{Error, InvalidTimezoneSnafu};
25
26pub fn to_meta_query_context(
27    query_context: QueryContextRef,
28) -> common_meta::rpc::ddl::QueryContext {
29    common_meta::rpc::ddl::QueryContext {
30        current_catalog: query_context.current_catalog().to_string(),
31        current_schema: query_context.current_schema().clone(),
32        timezone: query_context.timezone().to_string(),
33        extensions: query_context.extensions(),
34        channel: query_context.channel() as u8,
35        snapshot_seqs: query_context.snapshots(),
36        sst_min_sequences: query_context.sst_min_sequences(),
37    }
38}
39
40/// Builds the execution envelope for a frontend-originated procedure.
41pub(crate) fn to_executor_context(
42    query_context: QueryContextRef,
43    trigger_reason: TriggerReason,
44) -> ExecutorContext {
45    let actor = query_context.current_user().username().to_string();
46    ExecutorContext {
47        query_context: Some(to_meta_query_context(query_context)),
48        actor: Some(actor),
49        event_input: Some(ProcedureEventInput::new(trigger_reason)),
50        ..Default::default()
51    }
52}
53
54/// Builds a frontend-originated procedure envelope with the trusted frontend origin.
55pub(crate) fn to_executor_context_with_origin_frontend(
56    query_context: QueryContextRef,
57    origin_frontend_addr: &str,
58    trigger_reason: TriggerReason,
59) -> ExecutorContext {
60    let mut executor_context = to_executor_context(query_context, trigger_reason);
61    if let Some(query_context) = &mut executor_context.query_context {
62        query_context.extensions.insert(
63            ORIGIN_FRONTEND_ADDR_EXTENSION_KEY.to_string(),
64            origin_frontend_addr.to_string(),
65        );
66    }
67    executor_context
68}
69
70pub fn try_to_session_query_context(
71    value: common_meta::rpc::ddl::QueryContext,
72) -> Result<session::context::QueryContext, Error> {
73    Ok(QueryContextBuilder::default()
74        .current_catalog(value.current_catalog)
75        .current_schema(value.current_schema)
76        .timezone(
77            Timezone::from_tz_string(&value.timezone).context(InvalidTimezoneSnafu {
78                timezone: value.timezone,
79            })?,
80        )
81        .extensions(value.extensions)
82        .channel((value.channel as u32).into())
83        .snapshot_seqs(Arc::new(RwLock::new(value.snapshot_seqs)))
84        .sst_min_sequences(Arc::new(RwLock::new(value.sst_min_sequences)))
85        .build())
86}
87
88#[cfg(test)]
89mod tests {
90    use std::collections::HashMap;
91    use std::sync::{Arc, RwLock};
92
93    use common_meta::rpc::ddl::{ORIGIN_FRONTEND_ADDR_EXTENSION_KEY, TriggerReason};
94    use common_time::Timezone;
95    use session::context::QueryContextBuilder;
96
97    use super::{
98        to_executor_context_with_origin_frontend, to_meta_query_context,
99        try_to_session_query_context,
100    };
101
102    #[test]
103    fn test_query_context_meta_roundtrip_with_sequences() {
104        let session_ctx = Arc::new(
105            QueryContextBuilder::default()
106                .current_catalog("c1".to_string())
107                .current_schema("s1".to_string())
108                .timezone(Timezone::from_tz_string("UTC").unwrap())
109                .set_extension("flow.return_region_seq".to_string(), "true".to_string())
110                .snapshot_seqs(Arc::new(RwLock::new(HashMap::from([(10, 100)]))))
111                .sst_min_sequences(Arc::new(RwLock::new(HashMap::from([(10, 90)]))))
112                .build(),
113        );
114
115        let meta_ctx = to_meta_query_context(session_ctx);
116        let roundtrip = try_to_session_query_context(meta_ctx).unwrap();
117
118        assert_eq!(roundtrip.current_catalog(), "c1");
119        assert_eq!(roundtrip.current_schema(), "s1");
120        assert_eq!(roundtrip.snapshots(), HashMap::from([(10, 100)]));
121        assert_eq!(roundtrip.sst_min_sequences(), HashMap::from([(10, 90)]));
122        assert_eq!(roundtrip.extension("flow.return_region_seq"), Some("true"));
123    }
124
125    #[test]
126    fn test_executor_context_with_origin_frontend_overrides_reserved_key() {
127        let session_ctx = Arc::new(
128            QueryContextBuilder::default()
129                .set_extension(
130                    ORIGIN_FRONTEND_ADDR_EXTENSION_KEY.to_string(),
131                    "spoofed".to_string(),
132                )
133                .build(),
134        );
135
136        let executor_context = to_executor_context_with_origin_frontend(
137            session_ctx,
138            "127.0.0.1:4000",
139            TriggerReason::Manual,
140        );
141        let meta_ctx = executor_context.query_context.as_ref().unwrap();
142
143        assert_eq!(
144            meta_ctx
145                .extensions
146                .get(ORIGIN_FRONTEND_ADDR_EXTENSION_KEY)
147                .map(String::as_str),
148            Some("127.0.0.1:4000")
149        );
150        assert_eq!(
151            executor_context.event_input.unwrap().reason,
152            TriggerReason::Manual
153        );
154    }
155}
operator/utils.rs

operator/
utils.rs
