feat: support query limit by rule

server/src/grpc/storage_service/prom_query.rs

@@ -85,18 +85,19 @@ where
             };
             Error::ErrWithCause {
                 code,
-                msg: "failed to create plan".to_string(),
+                msg: "Failed to create plan".to_string(),
                 source: Box::new(e),
             }
         })?;
 
-    if ctx.instance.limiter.should_limit(&plan) {
-        ErrNoCause {
-            code: StatusCode::TOO_MANY_REQUESTS,
-            msg: "query limited by reject list",
-        }
-        .fail()?;
-    }
+    ctx.instance
+        .limiter
+        .try_limit(&plan)
+        .map_err(|e| Box::new(e) as _)
+        .context(ErrWithCause {
+            code: StatusCode::FORBIDDEN,
+            msg: "Query is blocked",
+        })?;
 
     // Execute in interpreter
     let interpreter_ctx = InterpreterContext::builder(request_id)
@@ -117,7 +118,7 @@ where
         .map_err(|e| Box::new(e) as _)
         .with_context(|| ErrWithCause {
             code: StatusCode::INTERNAL_SERVER_ERROR,
-            msg: "failed to execute interpreter",
+            msg: "Failed to execute interpreter",
         })?;
 
     let resp = convert_output(output, column_name)

server/src/grpc/storage_service/query.rs

@@ -128,13 +128,14 @@ pub async fn fetch_query_output<Q: QueryExecutor + 'static>(
             msg: format!("Failed to create plan, query:{}", req.ql),
         })?;
 
-    if ctx.instance.limiter.should_limit(&plan) {
-        ErrNoCause {
-            code: StatusCode::TOO_MANY_REQUESTS,
-            msg: "query limited by block list",
-        }
-        .fail()?;
-    }
+    ctx.instance
+        .limiter
+        .try_limit(&plan)
+        .map_err(|e| Box::new(e) as _)
+        .context(ErrWithCause {
+            code: StatusCode::FORBIDDEN,
+            msg: "Query is blocked",
+        })?;
 
     // Execute in interpreter
     let interpreter_ctx = InterpreterContext::builder(request_id)

server/src/grpc/storage_service/write.rs

@@ -56,13 +56,14 @@ pub(crate) async fn handle_write<Q: QueryExecutor + 'static>(
         );
         let plan = Plan::Insert(insert_plan);
 
-        if ctx.instance.limiter.should_limit(&plan) {
-            ErrNoCause {
-                code: StatusCode::TOO_MANY_REQUESTS,
-                msg: "Insert limited by reject list",
-            }
-            .fail()?;
-        }
+        ctx.instance
+            .limiter
+            .try_limit(&plan)
+            .map_err(|e| Box::new(e) as _)
+            .context(ErrWithCause {
+                code: StatusCode::FORBIDDEN,
+                msg: "Insert is blocked",
+            })?;
 
         let interpreter_ctx = InterpreterContext::builder(request_id)
             // Use current ctx's catalog and tenant as default catalog and tenant
@@ -207,13 +208,14 @@ async fn create_table<Q: QueryExecutor + 'static>(
 
     let instance = &ctx.instance;
 
-    if instance.limiter.should_limit(&plan) {
-        ErrNoCause {
-            code: StatusCode::TOO_MANY_REQUESTS,
-            msg: "Create table limited by reject list",
-        }
-        .fail()?;
-    }
+    instance
+        .limiter
+        .try_limit(&plan)
+        .map_err(|e| Box::new(e) as _)
+        .context(ErrWithCause {
+            code: StatusCode::FORBIDDEN,
+            msg: "Create table is blocked",
+        })?;
 
     let interpreter_ctx = InterpreterContext::builder(request_id)
         // Use current ctx's catalog and tenant as default catalog and tenant

server/src/handlers/admin.rs

@@ -2,7 +2,7 @@
 
 use std::collections::BTreeSet;
 
-use crate::handlers::prelude::*;
+use crate::{handlers::prelude::*, limiter::BlockRule};
 
 #[derive(Debug, Deserialize)]
 pub enum Operation {
@@ -12,60 +12,53 @@ pub enum Operation {
 }
 
 #[derive(Debug, Deserialize)]
-pub struct RejectRequest {
+pub struct BlockRequest {
     operation: Operation,
     write_block_list: Vec<String>,
     read_block_list: Vec<String>,
+    block_rules: Vec<BlockRule>,
 }
 
 #[derive(Serialize)]
-pub struct RejectResponse {
+pub struct BlockResponse {
     write_block_list: BTreeSet<String>,
     read_block_list: BTreeSet<String>,
+    block_rules: BTreeSet<BlockRule>,
 }
 
 pub async fn handle_block<Q: QueryExecutor + 'static>(
     _ctx: RequestContext,
     instance: InstanceRef<Q>,
-    request: RejectRequest,
-) -> Result<RejectResponse> {
+    request: BlockRequest,
+) -> Result<BlockResponse> {
+    let limiter = &instance.limiter;
     match request.operation {
         Operation::Add => {
-            instance
-                .limiter
-                .add_write_block_list(request.write_block_list);
-            instance
-                .limiter
-                .add_read_block_list(request.read_block_list);
+            limiter.add_write_block_list(request.write_block_list);
+            limiter.add_read_block_list(request.read_block_list);
+            limiter.add_block_rules(request.block_rules);
         }
         Operation::Set => {
-            instance
-                .limiter
-                .set_write_block_list(request.write_block_list);
-            instance
-                .limiter
-                .set_read_block_list(request.read_block_list);
+            limiter.set_write_block_list(request.write_block_list);
+            limiter.set_read_block_list(request.read_block_list);
+            limiter.set_block_rules(request.block_rules);
         }
         Operation::Remove => {
-            instance
-                .limiter
-                .remove_write_block_list(request.write_block_list);
-            instance
-                .limiter
-                .remove_read_block_list(request.read_block_list);
+            limiter.remove_write_block_list(request.write_block_list);
+            limiter.remove_read_block_list(request.read_block_list);
+            limiter.remove_block_rules(&request.block_rules);
         }
     }
 
-    Ok(RejectResponse {
-        write_block_list: instance
-            .limiter
+    Ok(BlockResponse {
+        write_block_list: limiter
             .get_write_block_list()
             .into_iter()
             .collect::<BTreeSet<_>>(),
-        read_block_list: instance
-            .limiter
+        read_block_list: limiter
             .get_read_block_list()
             .into_iter()
             .collect::<BTreeSet<_>>(),
+        block_rules: limiter.get_block_rules().into_iter().collect(),
     })
 }

server/src/handlers/error.rs

@@ -3,6 +3,8 @@
 //! Error of handlers
 
 use snafu::{Backtrace, Snafu};
+
+use crate::limiter;
 // TODO(yingwen): Avoid printing huge sql string
 // TODO(yingwen): Maybe add an error type to sql sub mod
 
@@ -48,8 +50,11 @@ pub enum Error {
         backtrace: Backtrace,
     },
 
-    #[snafu(display("Query limited by block list, query:{}", query))]
-    QueryBlock { query: String },
+    #[snafu(display("Query limited by block list, query:{}, err:{}", query, source))]
+    QueryBlock {
+        query: String,
+        source: limiter::Error,
+    },
 }
 
 define_result!(Error);

server/src/handlers/sql.rs

@@ -18,14 +18,14 @@ use serde::{
     ser::{SerializeMap, SerializeSeq},
     Serialize,
 };
-use snafu::ensure;
+use snafu::{ensure, ResultExt};
 use sql::{
     frontend::{Context as SqlContext, Frontend},
     provider::CatalogMetaProvider,
 };
 
 use crate::handlers::{
-    error::{ArrowToString, CreatePlan, Error::QueryBlock, InterpreterExec, ParseSql, TooMuchStmt},
+    error::{ArrowToString, CreatePlan, InterpreterExec, ParseSql, QueryBlock, TooMuchStmt},
     prelude::*,
 };
 
@@ -157,11 +157,9 @@ pub async fn handle_sql<Q: QueryExecutor + 'static>(
             query: &request.query,
         })?;
 
-    if instance.limiter.should_limit(&plan) {
-        return Err(QueryBlock {
-            query: request.query.to_owned(),
-        });
-    }
+    instance.limiter.try_limit(&plan).context(QueryBlock {
+        query: &request.query,
+    })?;
 
     // Execute in interpreter
     let interpreter_ctx = InterpreterContext::builder(request_id)