Skip to main content

common_event_recorder/
context.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::collections::HashMap;
16
17use api::v1::meta::ProcedureEventContext as PbProcedureEventContext;
18use serde::{Deserialize, Serialize};
19use strum::{AsRefStr, EnumString};
20
21/// Event metadata supplied when a procedure is submitted.
22///
23/// Protocol is deliberately absent. The trusted submission adapter derives it
24/// from the typed query channel before serializing the protobuf request.
25#[derive(Debug, Clone, Default, PartialEq, Eq)]
26pub struct ProcedureEventInput {
27    pub reason: TriggerReason,
28    pub extensions: HashMap<String, String>,
29}
30
31impl ProcedureEventInput {
32    /// Creates procedure event input with a trigger reason and no extensions.
33    pub fn new(reason: TriggerReason) -> Self {
34        Self {
35            reason,
36            extensions: Default::default(),
37        }
38    }
39}
40
41impl From<&ProcedureEventInput> for PbProcedureEventContext {
42    fn from(input: &ProcedureEventInput) -> Self {
43        Self {
44            reason: input.reason.as_ref().to_string(),
45            protocol: String::new(),
46            extensions: input.extensions.clone(),
47        }
48    }
49}
50
51impl From<PbProcedureEventContext> for ProcedureEventInput {
52    fn from(context: PbProcedureEventContext) -> Self {
53        Self {
54            reason: TriggerReason::from_extension(&context.reason),
55            extensions: context.extensions,
56        }
57    }
58}
59
60/// Stable context recorded for a procedure event.
61#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
62pub struct PersistentEventContext {
63    pub reason: TriggerReason,
64    #[serde(default, skip_serializing_if = "Option::is_none")]
65    pub protocol: Option<String>,
66    #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")]
67    pub extensions: serde_json::Map<String, serde_json::Value>,
68}
69
70impl PersistentEventContext {
71    /// Creates an event context with no additional extensions.
72    pub fn new(reason: TriggerReason) -> Self {
73        Self {
74            reason,
75            protocol: None,
76            extensions: Default::default(),
77        }
78    }
79
80    /// Adds the protocol that originated the operation.
81    pub fn with_protocol(mut self, protocol: impl Into<String>) -> Self {
82        self.protocol = Some(protocol.into());
83        self
84    }
85}
86
87impl From<(ProcedureEventInput, Option<String>)> for PersistentEventContext {
88    fn from((input, protocol): (ProcedureEventInput, Option<String>)) -> Self {
89        Self {
90            reason: input.reason,
91            protocol,
92            extensions: input
93                .extensions
94                .into_iter()
95                .map(|(key, value)| (key, serde_json::Value::String(value)))
96                .collect(),
97        }
98    }
99}
100
101impl Default for PersistentEventContext {
102    fn default() -> Self {
103        Self::new(TriggerReason::default())
104    }
105}
106
107impl From<PbProcedureEventContext> for PersistentEventContext {
108    fn from(context: PbProcedureEventContext) -> Self {
109        Self {
110            reason: TriggerReason::from_extension(&context.reason),
111            protocol: (!context.protocol.is_empty()).then_some(context.protocol),
112            extensions: context
113                .extensions
114                .into_iter()
115                .map(|(key, value)| (key, serde_json::Value::String(value)))
116                .collect(),
117        }
118    }
119}
120
121/// Stable classification of a procedure trigger.
122#[derive(
123    Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, AsRefStr, EnumString,
124)]
125#[serde(rename_all = "snake_case")]
126#[strum(serialize_all = "snake_case")]
127pub enum TriggerReason {
128    Manual,
129    AutoCreate,
130    AutoAlter,
131    AutoRepartition,
132    AutoRebalance,
133    RegionFailover,
134    ScheduledGc,
135    #[default]
136    #[serde(other)]
137    Unknown,
138}
139
140impl TriggerReason {
141    pub fn from_extension(value: &str) -> Self {
142        value.parse().unwrap_or_default()
143    }
144}
145
146#[cfg(test)]
147mod tests {
148    use serde_json::json;
149
150    use super::*;
151
152    #[test]
153    fn test_event_context_serialization() {
154        let context = PersistentEventContext::new(TriggerReason::Manual).with_protocol("mysql");
155
156        assert_eq!(
157            json!({
158                "reason": "manual",
159                "protocol": "mysql",
160            }),
161            serde_json::to_value(context).unwrap()
162        );
163        assert_eq!(
164            json!({ "reason": "manual" }),
165            serde_json::to_value(PersistentEventContext::new(TriggerReason::Manual)).unwrap()
166        );
167    }
168
169    #[test]
170    fn test_event_context_from_protobuf() {
171        let protobuf = PbProcedureEventContext {
172            reason: "auto_create".to_string(),
173            protocol: "postgres".to_string(),
174            extensions: HashMap::from([
175                ("source".to_string(), "sql".to_string()),
176                ("tenant".to_string(), "a".to_string()),
177            ]),
178        };
179        assert_eq!(
180            PersistentEventContext::from(protobuf),
181            PersistentEventContext {
182                reason: TriggerReason::AutoCreate,
183                protocol: Some("postgres".to_string()),
184                extensions: serde_json::Map::from_iter([
185                    ("source".to_string(), json!("sql")),
186                    ("tenant".to_string(), json!("a")),
187                ]),
188            }
189        );
190    }
191
192    #[test]
193    fn test_event_input_protobuf_has_no_protocol() {
194        let input = ProcedureEventInput {
195            reason: TriggerReason::AutoCreate,
196            extensions: HashMap::from([("source".to_string(), "sql".to_string())]),
197        };
198
199        assert_eq!(
200            PbProcedureEventContext::from(&input),
201            PbProcedureEventContext {
202                reason: "auto_create".to_string(),
203                protocol: String::new(),
204                extensions: input.extensions.clone(),
205            }
206        );
207
208        assert_eq!(
209            ProcedureEventInput::from(PbProcedureEventContext {
210                reason: "auto_create".to_string(),
211                protocol: "untrusted".to_string(),
212                extensions: input.extensions.clone(),
213            }),
214            input
215        );
216    }
217
218    #[test]
219    fn test_trigger_reason_deserializes_unknown_value() {
220        let reason: TriggerReason = serde_json::from_str("\"future_reason\"").unwrap();
221        assert_eq!(TriggerReason::Unknown, reason);
222    }
223}