1use std::fmt;
16
17use serde::{Deserialize, Deserializer, Serializer};
18use strum::{AsRefStr, EnumIter, EnumString, FromRepr};
19use tonic::Code;
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq, EnumString, AsRefStr, EnumIter, FromRepr)]
23pub enum StatusCode {
24 Success = 0,
27
28 Unknown = 1000,
30 Unsupported = 1001,
32 Unexpected = 1002,
34 Internal = 1003,
36 InvalidArguments = 1004,
38 Cancelled = 1005,
40 IllegalState = 1006,
42 External = 1007,
44 DeadlineExceeded = 1008,
46 Suspended = 1009,
48 InvalidSyntax = 2000,
53 PlanQuery = 3000,
58 EngineExecuteQuery = 3001,
60 TableAlreadyExists = 4000,
65 TableNotFound = 4001,
67 TableColumnNotFound = 4002,
69 TableColumnExists = 4003,
71 DatabaseNotFound = 4004,
73 RegionNotFound = 4005,
75 RegionAlreadyExists = 4006,
77 RegionReadonly = 4007,
79 RegionNotReady = 4008,
81 RegionBusy = 4009,
83 TableUnavailable = 4010,
85 DatabaseAlreadyExists = 4011,
87 StorageUnavailable = 5000,
92 RequestOutdated = 5001,
94 RuntimeResourcesExhausted = 6000,
99
100 RateLimited = 6001,
102 UserNotFound = 7000,
107 UnsupportedPasswordType = 7001,
109 UserPasswordMismatch = 7002,
111 AuthHeaderNotFound = 7003,
113 InvalidAuthHeader = 7004,
115 AccessDenied = 7005,
117 PermissionDenied = 7006,
119 FlowAlreadyExists = 8000,
123 FlowNotFound = 8001,
124 TriggerAlreadyExists = 9000,
128 TriggerNotFound = 9001,
129 }
131
132impl StatusCode {
133 pub fn is_success(code: u32) -> bool {
135 Self::Success as u32 == code
136 }
137
138 pub fn should_log_error(&self) -> bool {
141 match self {
142 StatusCode::Unknown
143 | StatusCode::Unexpected
144 | StatusCode::Internal
145 | StatusCode::Cancelled
146 | StatusCode::DeadlineExceeded
147 | StatusCode::IllegalState
148 | StatusCode::EngineExecuteQuery
149 | StatusCode::StorageUnavailable
150 | StatusCode::RuntimeResourcesExhausted
151 | StatusCode::External => true,
152
153 StatusCode::Success
154 | StatusCode::Unsupported
155 | StatusCode::InvalidArguments
156 | StatusCode::InvalidSyntax
157 | StatusCode::TableAlreadyExists
158 | StatusCode::TableNotFound
159 | StatusCode::RegionAlreadyExists
160 | StatusCode::RegionNotFound
161 | StatusCode::PlanQuery
162 | StatusCode::FlowAlreadyExists
163 | StatusCode::FlowNotFound
164 | StatusCode::TriggerAlreadyExists
165 | StatusCode::TriggerNotFound
166 | StatusCode::RegionNotReady
167 | StatusCode::RegionBusy
168 | StatusCode::RegionReadonly
169 | StatusCode::TableColumnNotFound
170 | StatusCode::TableColumnExists
171 | StatusCode::DatabaseNotFound
172 | StatusCode::RateLimited
173 | StatusCode::UserNotFound
174 | StatusCode::TableUnavailable
175 | StatusCode::DatabaseAlreadyExists
176 | StatusCode::UnsupportedPasswordType
177 | StatusCode::UserPasswordMismatch
178 | StatusCode::AuthHeaderNotFound
179 | StatusCode::InvalidAuthHeader
180 | StatusCode::AccessDenied
181 | StatusCode::PermissionDenied
182 | StatusCode::RequestOutdated
183 | StatusCode::Suspended => false,
184 }
185 }
186
187 pub fn from_u32(value: u32) -> Option<Self> {
188 StatusCode::from_repr(value as usize)
189 }
190
191 pub fn serialize_as_u32<S>(code: &Self, serializer: S) -> Result<S::Ok, S::Error>
192 where
193 S: Serializer,
194 {
195 serializer.serialize_u32(*code as u32)
196 }
197
198 pub fn deserialize_from_u32<'de, D>(deserializer: D) -> Result<Self, D::Error>
199 where
200 D: Deserializer<'de>,
201 {
202 let code = u32::deserialize(deserializer)?;
203 StatusCode::from_u32(code)
204 .ok_or_else(|| serde::de::Error::custom(format!("unknown status code: {code}")))
205 }
206}
207
208impl fmt::Display for StatusCode {
209 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
210 write!(f, "{self:?}")
212 }
213}
214
215#[macro_export]
216macro_rules! define_from_tonic_status {
217 ($Error: ty, $Variant: ident) => {
218 impl From<tonic::Status> for $Error {
219 fn from(e: tonic::Status) -> Self {
220 use snafu::location;
221
222 fn metadata_value(e: &tonic::Status, key: &str) -> Option<String> {
223 e.metadata()
224 .get(key)
225 .and_then(|v| String::from_utf8(v.as_bytes().to_vec()).ok())
226 }
227 let code = metadata_value(&e, $crate::GREPTIME_DB_HEADER_ERROR_CODE)
228 .and_then(|s| {
229 if let Ok(code) = s.parse::<u32>() {
230 StatusCode::from_u32(code)
231 } else {
232 None
233 }
234 })
235 .unwrap_or_else(|| match e.code() {
236 tonic::Code::Cancelled => StatusCode::Cancelled,
237 tonic::Code::DeadlineExceeded => StatusCode::DeadlineExceeded,
238 _ => StatusCode::Internal,
239 });
240
241 let msg = metadata_value(&e, $crate::GREPTIME_DB_HEADER_ERROR_MSG)
242 .unwrap_or_else(|| e.message().to_string());
243 let retry_hint = metadata_value(&e, $crate::GREPTIME_DB_HEADER_ERROR_RETRY_HINT)
244 .and_then(|s| s.parse().ok())
245 .unwrap_or($crate::ext::RetryHint::NonRetryable);
246
247 Self::$Variant {
249 code,
250 msg,
251 tonic_code: e.code(),
252 retry_hint,
253 location: location!(),
254 }
255 }
256 }
257 };
258}
259
260#[macro_export]
261macro_rules! define_into_tonic_status {
262 ($Error: ty) => {
263 impl From<$Error> for tonic::Status {
264 fn from(err: $Error) -> Self {
265 use tonic::codegen::http::{HeaderMap, HeaderValue};
266 use tonic::metadata::MetadataMap;
267 use $crate::{
268 GREPTIME_DB_HEADER_ERROR_CODE, GREPTIME_DB_HEADER_ERROR_RETRY_HINT,
269 };
270
271 common_telemetry::error!(err; "Failed to handle request");
272
273 let mut headers = HeaderMap::<HeaderValue>::with_capacity(3);
274
275 let status_code = err.status_code();
278 headers.insert(
279 GREPTIME_DB_HEADER_ERROR_CODE,
280 HeaderValue::from(status_code as u32),
281 );
282 headers.insert(
283 GREPTIME_DB_HEADER_ERROR_RETRY_HINT,
284 HeaderValue::from_static(err.retry_hint().as_str()),
285 );
286 let root_error = err.output_msg();
287
288 let metadata = MetadataMap::from_headers(headers);
289 tonic::Status::with_metadata(
290 $crate::status_code::status_to_tonic_code(status_code),
291 root_error,
292 metadata,
293 )
294 }
295 }
296 };
297}
298
299pub fn status_to_tonic_code(status_code: StatusCode) -> Code {
301 match status_code {
302 StatusCode::Success => Code::Ok,
303 StatusCode::Unknown | StatusCode::External => Code::Unknown,
304 StatusCode::Unsupported => Code::Unimplemented,
305 StatusCode::Unexpected
306 | StatusCode::IllegalState
307 | StatusCode::Internal
308 | StatusCode::PlanQuery
309 | StatusCode::EngineExecuteQuery => Code::Internal,
310 StatusCode::InvalidArguments | StatusCode::InvalidSyntax | StatusCode::RequestOutdated => {
311 Code::InvalidArgument
312 }
313 StatusCode::Cancelled => Code::Cancelled,
314 StatusCode::DeadlineExceeded => Code::DeadlineExceeded,
315 StatusCode::TableAlreadyExists
316 | StatusCode::TableColumnExists
317 | StatusCode::RegionAlreadyExists
318 | StatusCode::DatabaseAlreadyExists
319 | StatusCode::TriggerAlreadyExists
320 | StatusCode::FlowAlreadyExists => Code::AlreadyExists,
321 StatusCode::TableNotFound
322 | StatusCode::RegionNotFound
323 | StatusCode::TableColumnNotFound
324 | StatusCode::DatabaseNotFound
325 | StatusCode::UserNotFound
326 | StatusCode::TriggerNotFound
327 | StatusCode::FlowNotFound => Code::NotFound,
328 StatusCode::TableUnavailable
329 | StatusCode::StorageUnavailable
330 | StatusCode::RegionNotReady => Code::Unavailable,
331 StatusCode::RuntimeResourcesExhausted
332 | StatusCode::RateLimited
333 | StatusCode::RegionBusy
334 | StatusCode::Suspended => Code::ResourceExhausted,
335 StatusCode::UnsupportedPasswordType
336 | StatusCode::UserPasswordMismatch
337 | StatusCode::AuthHeaderNotFound
338 | StatusCode::InvalidAuthHeader => Code::Unauthenticated,
339 StatusCode::AccessDenied | StatusCode::PermissionDenied | StatusCode::RegionReadonly => {
340 Code::PermissionDenied
341 }
342 }
343}
344
345#[cfg(test)]
346mod tests {
347 use strum::IntoEnumIterator;
348
349 use super::*;
350
351 fn assert_status_code_display(code: StatusCode, msg: &str) {
352 let code_msg = format!("{code}");
353 assert_eq!(msg, code_msg);
354 }
355
356 #[test]
357 fn test_display_status_code() {
358 assert_status_code_display(StatusCode::Unknown, "Unknown");
359 assert_status_code_display(StatusCode::TableAlreadyExists, "TableAlreadyExists");
360 }
361
362 #[test]
363 fn test_from_u32() {
364 for code in StatusCode::iter() {
365 let num = code as u32;
366 assert_eq!(StatusCode::from_u32(num), Some(code));
367 }
368
369 assert_eq!(StatusCode::from_u32(10000), None);
370 }
371
372 #[test]
373 fn test_is_success() {
374 assert!(StatusCode::is_success(0));
375 assert!(!StatusCode::is_success(1));
376 assert!(!StatusCode::is_success(2));
377 assert!(!StatusCode::is_success(3));
378 }
379}