feat: sql query interceptor and plugin refactoring (#773)

* feat: let instance hold plugins

* feat: add sql query interceptor definition

* docs: add comments to key apis

* feat: add implementation for pre-parsing and post-parsing

* feat: add post_execute hook

* test: add tests for interceptor

* chore: add license header

* fix: clippy error

* Update src/cmd/src/frontend.rs

Co-authored-by: LFC <bayinamine@gmail.com>

* refactor: batching post_parsing calls

* refactor: rename AnyMap2 to Plugins

* feat: call pre_execute with logical plan empty at the moment

Co-authored-by: LFC <bayinamine@gmail.com>

src/servers/Cargo.toml

@@ -37,6 +37,7 @@ openmetrics-parser = "0.4"
 opensrv-mysql = "0.3"
 pgwire = "0.6.3"
 prost = "0.11"
+query = { path = "../query" }
 rand = "0.8"
 regex = "1.6"
 rustls = "0.20"
@@ -65,7 +66,6 @@ common-base = { path = "../common/base" }
 mysql_async = { version = "0.31", default-features = false, features = [
     "default-rustls",
 ] }
-query = { path = "../query" }
 rand = "0.8"
 script = { path = "../script", features = ["python"] }
 serde_json = "1.0"

src/servers/src/interceptor.rs

@@ -0,0 +1,105 @@
+// Copyright 2022 Greptime Team
+//
+// 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::borrow::Cow;
+use std::sync::Arc;
+
+use common_query::Output;
+use query::plan::LogicalPlan;
+use session::context::QueryContextRef;
+use sql::statements::statement::Statement;
+
+use crate::error::Result;
+
+/// SqlQueryInterceptor can track life cycle of a sql query and customize or
+/// abort its execution at given point.
+pub trait SqlQueryInterceptor {
+    /// Called before a query string is parsed into sql statements.
+    /// The implementation is allowed to change the sql string if needed.
+    fn pre_parsing<'a>(&self, query: &'a str, _query_ctx: QueryContextRef) -> Result<Cow<'a, str>> {
+        Ok(Cow::Borrowed(query))
+    }
+
+    /// Called after sql is parsed into statements. This interceptor is called
+    /// on each statement and the implementation can alter the statement or
+    /// abort execution by raising an error.
+    fn post_parsing(
+        &self,
+        statements: Vec<Statement>,
+        _query_ctx: QueryContextRef,
+    ) -> Result<Vec<Statement>> {
+        Ok(statements)
+    }
+
+    /// Called before sql is actually executed. This hook is not called at the moment.
+    fn pre_execute(
+        &self,
+        _statement: &Statement,
+        _plan: Option<&LogicalPlan>,
+        _query_ctx: QueryContextRef,
+    ) -> Result<()> {
+        Ok(())
+    }
+
+    /// Called after execution finished. The implementation can modify the
+    /// output if needed.
+    fn post_execute(&self, output: Output, _query_ctx: QueryContextRef) -> Result<Output> {
+        Ok(output)
+    }
+}
+
+pub type SqlQueryInterceptorRef = Arc<dyn SqlQueryInterceptor + Send + Sync + 'static>;
+
+impl SqlQueryInterceptor for Option<&SqlQueryInterceptorRef> {
+    fn pre_parsing<'a>(&self, query: &'a str, query_ctx: QueryContextRef) -> Result<Cow<'a, str>> {
+        if let Some(this) = self {
+            this.pre_parsing(query, query_ctx)
+        } else {
+            Ok(Cow::Borrowed(query))
+        }
+    }
+
+    fn post_parsing(
+        &self,
+        statements: Vec<Statement>,
+        query_ctx: QueryContextRef,
+    ) -> Result<Vec<Statement>> {
+        if let Some(this) = self {
+            this.post_parsing(statements, query_ctx)
+        } else {
+            Ok(statements)
+        }
+    }
+
+    fn pre_execute(
+        &self,
+        statement: &Statement,
+        plan: Option<&LogicalPlan>,
+        query_ctx: QueryContextRef,
+    ) -> Result<()> {
+        if let Some(this) = self {
+            this.pre_execute(statement, plan, query_ctx)
+        } else {
+            Ok(())
+        }
+    }
+
+    fn post_execute(&self, output: Output, query_ctx: QueryContextRef) -> Result<Output> {
+        if let Some(this) = self {
+            this.post_execute(output, query_ctx)
+        } else {
+            Ok(output)
+        }
+    }
+}

src/servers/src/lib.rs

@@ -22,6 +22,7 @@ pub mod error;
 pub mod grpc;
 pub mod http;
 pub mod influxdb;
+pub mod interceptor;
 pub mod line_writer;
 pub mod mysql;
 pub mod opentsdb;

src/servers/tests/interceptor.rs

@@ -0,0 +1,38 @@
+// Copyright 2022 Greptime Team
+//
+// 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::borrow::Cow;
+use std::sync::Arc;
+
+use servers::error::Result;
+use servers::interceptor::SqlQueryInterceptor;
+use session::context::{QueryContext, QueryContextRef};
+
+pub struct NoopInterceptor;
+
+impl SqlQueryInterceptor for NoopInterceptor {
+    fn pre_parsing<'a>(&self, query: &'a str, _query_ctx: QueryContextRef) -> Result<Cow<'a, str>> {
+        let modified_query = format!("{query};");
+        Ok(Cow::Owned(modified_query))
+    }
+}
+
+#[test]
+fn test_default_interceptor_behaviour() {
+    let di = NoopInterceptor;
+    let ctx = Arc::new(QueryContext::new());
+
+    let query = "SELECT 1";
+    assert_eq!("SELECT 1;", di.pre_parsing(query, ctx).unwrap());
+}

src/servers/tests/mod.rs

@@ -33,6 +33,7 @@ use script::engine::{CompileContext, EvalContext, Script, ScriptEngine};
 use script::python::{PyEngine, PyScript};
 use session::context::QueryContextRef;
 
+mod interceptor;
 mod opentsdb;
 mod postgres;