Merge branch 'asof' of https://github.com/zenus/databend into asof

src/common/exception/src/exception_backtrace.rs

@@ -19,9 +19,19 @@ use std::sync::Arc;
 
 use crate::exception::ErrorCodeBacktrace;
 
+// 0: not specified 1: disable 2: enable
+pub static USER_SET_ENABLE_BACKTRACE: AtomicUsize = AtomicUsize::new(0);
+
+pub fn set_backtrace(switch: bool) {
+    if switch {
+        USER_SET_ENABLE_BACKTRACE.store(2, Ordering::Relaxed);
+    } else {
+        USER_SET_ENABLE_BACKTRACE.store(1, Ordering::Relaxed);
+    }
+}
+
 fn enable_rust_backtrace() -> bool {
-    static ENABLED: AtomicUsize = AtomicUsize::new(0);
-    match ENABLED.load(Ordering::Relaxed) {
+    match USER_SET_ENABLE_BACKTRACE.load(Ordering::Relaxed) {
         0 => {}
         1 => return false,
         _ => return true,
@@ -35,7 +45,7 @@ fn enable_rust_backtrace() -> bool {
         },
     };
 
-    ENABLED.store(enabled as usize + 1, Ordering::Relaxed);
+    USER_SET_ENABLE_BACKTRACE.store(enabled as usize + 1, Ordering::Relaxed);
     enabled
 }
 

src/common/exception/src/lib.rs

@@ -23,6 +23,8 @@ mod with_context;
 pub use exception::ErrorCode;
 pub use exception::Result;
 pub use exception::ToErrorCode;
+pub use exception_backtrace::set_backtrace;
+pub use exception_backtrace::USER_SET_ENABLE_BACKTRACE;
 pub use exception_into::SerializedError;
 pub use with_context::ErrorWithContext;
 pub use with_context::WithContext;

src/common/tracing/Cargo.toml

@@ -12,6 +12,7 @@ test = true
 
 [dependencies]
 databend-common-base = { workspace = true }
+databend-common-exception = { workspace = true }
 defer = "0.2"
 fern = "0.6.2"
 humantime = "2.1.0"

src/common/tracing/src/panic_hook.rs

@@ -14,8 +14,10 @@
 
 use std::backtrace::Backtrace;
 use std::panic::PanicInfo;
+use std::sync::atomic::Ordering;
 
 use databend_common_base::runtime::LimitMemGuard;
+use databend_common_exception::USER_SET_ENABLE_BACKTRACE;
 use log::error;
 
 pub fn set_panic_hook() {
@@ -31,12 +33,25 @@ pub fn set_panic_hook() {
     }));
 }
 
+fn should_backtrace() -> bool {
+    // if user not specify or user set to enable, we should backtrace
+    match USER_SET_ENABLE_BACKTRACE.load(Ordering::Relaxed) {
+        0 => true,
+        1 => false,
+        _ => true,
+    }
+}
+
 pub fn log_panic(panic: &PanicInfo) {
-    let backtrace = Backtrace::force_capture();
-    let backtrace_str = format!("{:?}", backtrace);
+    let backtrace_str = if should_backtrace() {
+        let backtrace = Backtrace::force_capture();
+        format!("{:?}", backtrace)
+    } else {
+        String::new()
+    };
 
     eprintln!("{}", panic);
-    eprintln!("{}", backtrace);
+    eprintln!("{}", backtrace_str);
 
     if let Some(location) = panic.location() {
         error!(

src/query/ast/src/ast/statements/mod.rs

@@ -45,6 +45,7 @@ mod show;
 mod stage;
 mod statement;
 mod stream;
+mod system_action;
 mod table;
 mod task;
 mod udf;
@@ -87,6 +88,7 @@ pub use show::*;
 pub use stage::*;
 pub use statement::*;
 pub use stream::*;
+pub use system_action::*;
 pub use table::*;
 pub use task::*;
 pub use udf::*;

src/query/ast/src/ast/statements/statement.rs

@@ -317,6 +317,9 @@ pub enum Statement {
         priority: Priority,
         object_id: String,
     },
+
+    // System actions
+    System(SystemStmt),
 }
 
 impl Statement {
@@ -721,6 +724,7 @@ impl Display for Statement {
                 write!(f, " {priority}")?;
                 write!(f, " '{object_id}'")?;
             }
+            Statement::System(stmt) => write!(f, "{stmt}")?,
         }
         Ok(())
     }

src/query/ast/src/ast/statements/system_action.rs

@@ -0,0 +1,46 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::fmt::Display;
+use std::fmt::Formatter;
+
+use derive_visitor::Drive;
+use derive_visitor::DriveMut;
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct SystemStmt {
+    pub action: SystemAction,
+}
+
+impl Display for SystemStmt {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "SYSTEM {}", self.action)
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
+pub enum SystemAction {
+    Backtrace(bool),
+}
+
+impl Display for SystemAction {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        match self {
+            SystemAction::Backtrace(switch) => match switch {
+                true => write!(f, "ENABLE EXCEPTION_BACKTRACE"),
+                false => write!(f, "DISABLE EXCEPTION_BACKTRACE"),
+            },
+        }
+    }
+}

src/query/ast/src/ast/visitors/visitor.rs

@@ -844,4 +844,6 @@ pub trait Visitor<'ast>: Sized {
     fn visit_drop_sequence(&mut self, _stmt: &'ast DropSequenceStmt) {}
     fn visit_set_priority(&mut self, _priority: &'ast Priority, _object_id: &'ast str) {}
     fn visit_multi_table_insert(&mut self, insert: &'ast InsertMultiTableStmt);
+
+    fn visit_system(&mut self, _stmt: &'ast SystemStmt) {}
 }

src/query/ast/src/ast/visitors/visitor_mut.rs

@@ -854,4 +854,5 @@ pub trait VisitorMut: Sized {
     fn visit_create_sequence(&mut self, _stmt: &mut CreateSequenceStmt) {}
     fn visit_drop_sequence(&mut self, _stmt: &mut DropSequenceStmt) {}
     fn visit_set_priority(&mut self, _priority: &mut Priority, _object_id: &mut String) {}
+    fn visit_system(&mut self, _stmt: &mut SystemStmt) {}
 }

src/query/ast/src/ast/visitors/walk.rs

@@ -597,5 +597,6 @@ pub fn walk_statement<'a, V: Visitor<'a>>(visitor: &mut V, statement: &'a Statem
             priority,
             object_id,
         } => visitor.visit_set_priority(priority, object_id),
+        Statement::System(stmt) => visitor.visit_system(stmt),
     }
 }

src/query/ast/src/ast/visitors/walk_mut.rs

@@ -596,5 +596,6 @@ pub fn walk_statement_mut<V: VisitorMut>(visitor: &mut V, statement: &mut Statem
             priority,
             object_id,
         } => visitor.visit_set_priority(priority, object_id),
+        Statement::System(stmt) => visitor.visit_system(stmt),
     }
 }

src/query/ast/src/parser/statement.rs

@@ -2047,6 +2047,13 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         |(_, _, script)| Statement::ExecuteImmediate(ExecuteImmediateStmt { script }),
     );
 
+    let system_action = map(
+        rule! {
+            SYSTEM ~ #action
+        },
+        |(_, action)| Statement::System(SystemStmt { action }),
+    );
+
     alt((
         // query, explain,show
         rule!(
@@ -2063,7 +2070,8 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
             | #show_locks : "`SHOW LOCKS [IN ACCOUNT] [WHERE ...]`"
             | #kill_stmt : "`KILL (QUERY | CONNECTION) <object_id>`"
             | #vacuum_temp_files : "VACUUM TEMPORARY FILES [RETAIN number SECONDS|DAYS] [LIMIT number]"
-            | #set_priority: "SET PRIORITY (HIGH | MEDIUM | LOW) <object_id>"
+            | #set_priority: "`SET PRIORITY (HIGH | MEDIUM | LOW) <object_id>`"
+            | #system_action: "`SYSTEM (ENABLE | DISABLE) EXCEPTION_BACKTRACE`"
         ),
         // database
         rule!(
@@ -3744,6 +3752,26 @@ pub fn priority(i: Input) -> IResult<Priority> {
     ))(i)
 }
 
+pub fn action(i: Input) -> IResult<SystemAction> {
+    let mut backtrace = map(
+        rule! {
+             #switch ~ EXCEPTION_BACKTRACE
+        },
+        |(switch, _)| SystemAction::Backtrace(switch),
+    );
+    // add other system action type here
+    rule!(
+        #backtrace
+    )(i)
+}
+
+pub fn switch(i: Input) -> IResult<bool> {
+    alt((
+        value(true, rule! { ENABLE }),
+        value(false, rule! { DISABLE }),
+    ))(i)
+}
+
 pub fn limit_where(i: Input) -> IResult<ShowLimit> {
     map(
         rule! {

src/query/ast/src/parser/token.rs

@@ -502,6 +502,8 @@ pub enum TokenKind {
     DETAILED_OUTPUT,
     #[token("DESCRIBE", ignore(ascii_case))]
     DESCRIBE,
+    #[token("DISABLE", ignore(ascii_case))]
+    DISABLE,
     #[token("DISABLE_VARIANT_CHECK", ignore(ascii_case))]
     DISABLE_VARIANT_CHECK,
     #[token("DISTINCT", ignore(ascii_case))]
@@ -540,6 +542,8 @@ pub enum TokenKind {
     ELSE,
     #[token("EMPTY_FIELD_AS", ignore(ascii_case))]
     EMPTY_FIELD_AS,
+    #[token("ENABLE", ignore(ascii_case))]
+    ENABLE,
     #[token("ENABLE_VIRTUAL_HOST_STYLE", ignore(ascii_case))]
     ENABLE_VIRTUAL_HOST_STYLE,
     #[token("END", ignore(ascii_case))]
@@ -556,6 +560,8 @@ pub enum TokenKind {
     ERROR_ON_COLUMN_COUNT_MISMATCH,
     #[token("ESCAPE", ignore(ascii_case))]
     ESCAPE,
+    #[token("EXCEPTION_BACKTRACE", ignore(ascii_case))]
+    EXCEPTION_BACKTRACE,
     #[token("EXISTS", ignore(ascii_case))]
     EXISTS,
     #[token("EXPLAIN", ignore(ascii_case))]
@@ -1055,6 +1061,8 @@ pub enum TokenKind {
     SOUNDS,
     #[token("SYNC", ignore(ascii_case))]
     SYNC,
+    #[token("SYSTEM", ignore(ascii_case))]
+    SYSTEM,
     #[token("STORAGE_TYPE", ignore(ascii_case))]
     STORAGE_TYPE,
     #[token("TABLE", ignore(ascii_case))]

src/query/service/src/interpreters/access/privilege_access.rs

@@ -1069,7 +1069,7 @@ impl AccessChecker for PrivilegeAccess {
                 self.validate_access(&GrantObject::Global, UserPrivilegeType::Grant,false)
                     .await?;
             }
-            Plan::SetVariable(_) | Plan::UnSetVariable(_) | Plan::Kill(_) | Plan::SetPriority(_) => {
+            Plan::SetVariable(_) | Plan::UnSetVariable(_) | Plan::Kill(_) | Plan::SetPriority(_) | Plan::System(_) => {
                 self.validate_access(&GrantObject::Global, UserPrivilegeType::Super, false)
                     .await?;
             }

src/query/service/src/interpreters/interpreter_factory.rs

@@ -50,6 +50,7 @@ use crate::interpreters::interpreter_notification_drop::DropNotificationInterpre
 use crate::interpreters::interpreter_presign::PresignInterpreter;
 use crate::interpreters::interpreter_role_show::ShowRolesInterpreter;
 use crate::interpreters::interpreter_set_priority::SetPriorityInterpreter;
+use crate::interpreters::interpreter_system_action::SystemActionInterpreter;
 use crate::interpreters::interpreter_table_create::CreateTableInterpreter;
 use crate::interpreters::interpreter_table_revert::RevertTableInterpreter;
 use crate::interpreters::interpreter_task_alter::AlterTaskInterpreter;
@@ -598,6 +599,10 @@ impl InterpreterFactory {
                 ctx,
                 *p.clone(),
             )?)),
+            Plan::System(p) => Ok(Arc::new(SystemActionInterpreter::try_create(
+                ctx,
+                *p.clone(),
+            )?)),
         }
     }
 }