1use std::fmt;
16
17use strum::{AsRefStr, EnumIter, EnumString, FromRepr};
18use tonic::Code;
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq, EnumString, AsRefStr, EnumIter, FromRepr)]
22pub enum StatusCode {
23 Success = 0,
26
27 Unknown = 1000,
29 Unsupported = 1001,
31 Unexpected = 1002,
33 Internal = 1003,
35 InvalidArguments = 1004,
37 Cancelled = 1005,
39 IllegalState = 1006,
41 External = 1007,
43 DeadlineExceeded = 1008,
45 InvalidSyntax = 2000,
50 PlanQuery = 3000,
55 EngineExecuteQuery = 3001,
57 TableAlreadyExists = 4000,
62 TableNotFound = 4001,
64 TableColumnNotFound = 4002,
66 TableColumnExists = 4003,
68 DatabaseNotFound = 4004,
70 RegionNotFound = 4005,
72 RegionAlreadyExists = 4006,
74 RegionReadonly = 4007,
76 RegionNotReady = 4008,
78 RegionBusy = 4009,
80 TableUnavailable = 4010,
82 DatabaseAlreadyExists = 4011,
84 StorageUnavailable = 5000,
89 RequestOutdated = 5001,
91 RuntimeResourcesExhausted = 6000,
96
97 RateLimited = 6001,
99 UserNotFound = 7000,
104 UnsupportedPasswordType = 7001,
106 UserPasswordMismatch = 7002,
108 AuthHeaderNotFound = 7003,
110 InvalidAuthHeader = 7004,
112 AccessDenied = 7005,
114 PermissionDenied = 7006,
116 FlowAlreadyExists = 8000,
120 FlowNotFound = 8001,
121 TriggerAlreadyExists = 9000,
125 TriggerNotFound = 9001,
126 }
128
129impl StatusCode {
130 pub fn is_success(code: u32) -> bool {
132 Self::Success as u32 == code
133 }
134
135 pub fn is_retryable(&self) -> bool {
137 match self {
138 StatusCode::StorageUnavailable
139 | StatusCode::RuntimeResourcesExhausted
140 | StatusCode::Internal
141 | StatusCode::RegionNotReady
142 | StatusCode::TableUnavailable
143 | StatusCode::RegionBusy => true,
144
145 StatusCode::Success
146 | StatusCode::Unknown
147 | StatusCode::Unsupported
148 | StatusCode::IllegalState
149 | StatusCode::Unexpected
150 | StatusCode::InvalidArguments
151 | StatusCode::Cancelled
152 | StatusCode::DeadlineExceeded
153 | StatusCode::InvalidSyntax
154 | StatusCode::DatabaseAlreadyExists
155 | StatusCode::PlanQuery
156 | StatusCode::EngineExecuteQuery
157 | StatusCode::TableAlreadyExists
158 | StatusCode::TableNotFound
159 | StatusCode::RegionAlreadyExists
160 | StatusCode::RegionNotFound
161 | StatusCode::FlowAlreadyExists
162 | StatusCode::FlowNotFound
163 | StatusCode::TriggerAlreadyExists
164 | StatusCode::TriggerNotFound
165 | StatusCode::RegionReadonly
166 | StatusCode::TableColumnNotFound
167 | StatusCode::TableColumnExists
168 | StatusCode::DatabaseNotFound
169 | StatusCode::RateLimited
170 | StatusCode::UserNotFound
171 | StatusCode::UnsupportedPasswordType
172 | StatusCode::UserPasswordMismatch
173 | StatusCode::AuthHeaderNotFound
174 | StatusCode::InvalidAuthHeader
175 | StatusCode::AccessDenied
176 | StatusCode::PermissionDenied
177 | StatusCode::RequestOutdated
178 | StatusCode::External => false,
179 }
180 }
181
182 pub fn should_log_error(&self) -> bool {
185 match self {
186 StatusCode::Unknown
187 | StatusCode::Unexpected
188 | StatusCode::Internal
189 | StatusCode::Cancelled
190 | StatusCode::DeadlineExceeded
191 | StatusCode::IllegalState
192 | StatusCode::EngineExecuteQuery
193 | StatusCode::StorageUnavailable
194 | StatusCode::RuntimeResourcesExhausted
195 | StatusCode::External => true,
196
197 StatusCode::Success
198 | StatusCode::Unsupported
199 | StatusCode::InvalidArguments
200 | StatusCode::InvalidSyntax
201 | StatusCode::TableAlreadyExists
202 | StatusCode::TableNotFound
203 | StatusCode::RegionAlreadyExists
204 | StatusCode::RegionNotFound
205 | StatusCode::PlanQuery
206 | StatusCode::FlowAlreadyExists
207 | StatusCode::FlowNotFound
208 | StatusCode::TriggerAlreadyExists
209 | StatusCode::TriggerNotFound
210 | StatusCode::RegionNotReady
211 | StatusCode::RegionBusy
212 | StatusCode::RegionReadonly
213 | StatusCode::TableColumnNotFound
214 | StatusCode::TableColumnExists
215 | StatusCode::DatabaseNotFound
216 | StatusCode::RateLimited
217 | StatusCode::UserNotFound
218 | StatusCode::TableUnavailable
219 | StatusCode::DatabaseAlreadyExists
220 | StatusCode::UnsupportedPasswordType
221 | StatusCode::UserPasswordMismatch
222 | StatusCode::AuthHeaderNotFound
223 | StatusCode::InvalidAuthHeader
224 | StatusCode::AccessDenied
225 | StatusCode::PermissionDenied
226 | StatusCode::RequestOutdated => false,
227 }
228 }
229
230 pub fn from_u32(value: u32) -> Option<Self> {
231 StatusCode::from_repr(value as usize)
232 }
233}
234
235impl fmt::Display for StatusCode {
236 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
237 write!(f, "{self:?}")
239 }
240}
241
242#[macro_export]
243macro_rules! define_from_tonic_status {
244 ($Error: ty, $Variant: ident) => {
245 impl From<tonic::Status> for $Error {
246 fn from(e: tonic::Status) -> Self {
247 use snafu::location;
248
249 fn metadata_value(e: &tonic::Status, key: &str) -> Option<String> {
250 e.metadata()
251 .get(key)
252 .and_then(|v| String::from_utf8(v.as_bytes().to_vec()).ok())
253 }
254 let code = metadata_value(&e, $crate::GREPTIME_DB_HEADER_ERROR_CODE)
255 .and_then(|s| {
256 if let Ok(code) = s.parse::<u32>() {
257 StatusCode::from_u32(code)
258 } else {
259 None
260 }
261 })
262 .unwrap_or_else(|| match e.code() {
263 tonic::Code::Cancelled => StatusCode::Cancelled,
264 tonic::Code::DeadlineExceeded => StatusCode::DeadlineExceeded,
265 _ => StatusCode::Internal,
266 });
267
268 let msg = metadata_value(&e, $crate::GREPTIME_DB_HEADER_ERROR_MSG)
269 .unwrap_or_else(|| e.message().to_string());
270
271 Self::$Variant {
273 code,
274 msg,
275 tonic_code: e.code(),
276 location: location!(),
277 }
278 }
279 }
280 };
281}
282
283#[macro_export]
284macro_rules! define_into_tonic_status {
285 ($Error: ty) => {
286 impl From<$Error> for tonic::Status {
287 fn from(err: $Error) -> Self {
288 use tonic::codegen::http::{HeaderMap, HeaderValue};
289 use tonic::metadata::MetadataMap;
290 use $crate::GREPTIME_DB_HEADER_ERROR_CODE;
291
292 common_telemetry::error!(err; "Failed to handle request");
293
294 let mut headers = HeaderMap::<HeaderValue>::with_capacity(2);
295
296 let status_code = err.status_code();
299 headers.insert(
300 GREPTIME_DB_HEADER_ERROR_CODE,
301 HeaderValue::from(status_code as u32),
302 );
303 let root_error = err.output_msg();
304
305 let metadata = MetadataMap::from_headers(headers);
306 tonic::Status::with_metadata(
307 $crate::status_code::status_to_tonic_code(status_code),
308 root_error,
309 metadata,
310 )
311 }
312 }
313 };
314}
315
316pub fn status_to_tonic_code(status_code: StatusCode) -> Code {
318 match status_code {
319 StatusCode::Success => Code::Ok,
320 StatusCode::Unknown | StatusCode::External => Code::Unknown,
321 StatusCode::Unsupported => Code::Unimplemented,
322 StatusCode::Unexpected
323 | StatusCode::IllegalState
324 | StatusCode::Internal
325 | StatusCode::PlanQuery
326 | StatusCode::EngineExecuteQuery => Code::Internal,
327 StatusCode::InvalidArguments | StatusCode::InvalidSyntax | StatusCode::RequestOutdated => {
328 Code::InvalidArgument
329 }
330 StatusCode::Cancelled => Code::Cancelled,
331 StatusCode::DeadlineExceeded => Code::DeadlineExceeded,
332 StatusCode::TableAlreadyExists
333 | StatusCode::TableColumnExists
334 | StatusCode::RegionAlreadyExists
335 | StatusCode::DatabaseAlreadyExists
336 | StatusCode::TriggerAlreadyExists
337 | StatusCode::FlowAlreadyExists => Code::AlreadyExists,
338 StatusCode::TableNotFound
339 | StatusCode::RegionNotFound
340 | StatusCode::TableColumnNotFound
341 | StatusCode::DatabaseNotFound
342 | StatusCode::UserNotFound
343 | StatusCode::TriggerNotFound
344 | StatusCode::FlowNotFound => Code::NotFound,
345 StatusCode::TableUnavailable
346 | StatusCode::StorageUnavailable
347 | StatusCode::RegionNotReady => Code::Unavailable,
348 StatusCode::RuntimeResourcesExhausted
349 | StatusCode::RateLimited
350 | StatusCode::RegionBusy => Code::ResourceExhausted,
351 StatusCode::UnsupportedPasswordType
352 | StatusCode::UserPasswordMismatch
353 | StatusCode::AuthHeaderNotFound
354 | StatusCode::InvalidAuthHeader => Code::Unauthenticated,
355 StatusCode::AccessDenied | StatusCode::PermissionDenied | StatusCode::RegionReadonly => {
356 Code::PermissionDenied
357 }
358 }
359}
360
361#[cfg(test)]
362mod tests {
363 use strum::IntoEnumIterator;
364
365 use super::*;
366
367 fn assert_status_code_display(code: StatusCode, msg: &str) {
368 let code_msg = format!("{code}");
369 assert_eq!(msg, code_msg);
370 }
371
372 #[test]
373 fn test_display_status_code() {
374 assert_status_code_display(StatusCode::Unknown, "Unknown");
375 assert_status_code_display(StatusCode::TableAlreadyExists, "TableAlreadyExists");
376 }
377
378 #[test]
379 fn test_from_u32() {
380 for code in StatusCode::iter() {
381 let num = code as u32;
382 assert_eq!(StatusCode::from_u32(num), Some(code));
383 }
384
385 assert_eq!(StatusCode::from_u32(10000), None);
386 }
387
388 #[test]
389 fn test_is_success() {
390 assert!(StatusCode::is_success(0));
391 assert!(!StatusCode::is_success(1));
392 assert!(!StatusCode::is_success(2));
393 assert!(!StatusCode::is_success(3));
394 }
395}