feat: support request context (apache#636)

* feat: support new request context * fix clippy * fix message * fix wong wal location * remove schema from route trait * fix clippy * make CI happy * refactor by CR --------- Co-authored-by: chunshao.rcs <[email protected]>
chunshao90 · Feb 13, 2023 · 45ebebf · 45ebebf
1 parent 5b60525
commit 45ebebf
Show file tree

Hide file tree

Showing 15 changed files with 147 additions and 151 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/docs/minimal.toml b/docs/minimal.toml
@@ -4,9 +4,6 @@ http_port = 5440
 grpc_port = 8831
 log_level = "info"
 
-[analytic]
-wal_path = "/tmp/ceresdb"
-
 [analytic.storage]
 mem_cache_capacity = '1G'
 mem_cache_partition_bits = 0
@@ -15,6 +12,10 @@ mem_cache_partition_bits = 0
 type = "Local"
 data_path = "/tmp/ceresdb"
 
+[analytic.wal_storage]
+type = "RocksDB"
+path = "/tmp/ceresdb"
+
 [[static_route.topology.schema_shards]]
 schema = 'public'
 auto_create_tables = true

diff --git a/integration_tests/src/database.rs b/integration_tests/src/database.rs
@@ -74,7 +74,10 @@ impl CeresDB {
     }
 
     async fn execute(query: String, client: Arc<dyn DbClient>) -> Box<dyn Display> {
-        let query_ctx = RpcContext::new("public".to_string(), "".to_string());
+        let query_ctx = RpcContext {
+            database: Some("public".to_string()),
+            timeout: None,
+        };
         let query_req = Request {
             tables: vec![],
             sql: query,

diff --git a/remote_engine_client/src/cached_router.rs b/remote_engine_client/src/cached_router.rs
@@ -4,7 +4,7 @@
 
 use std::collections::HashMap;
 
-use ceresdbproto::storage;
+use ceresdbproto::storage::{self, RequestContext};
 use log::debug;
 use router::RouterRef;
 use snafu::{OptionExt, ResultExt};
@@ -89,15 +89,18 @@ impl CachedRouter {
         let schema = &table_ident.schema;
         let table = table_ident.table.clone();
         let route_request = storage::RouteRequest {
+            context: Some(RequestContext {
+                database: schema.to_string(),
+            }),
             tables: vec![table],
         };
-        let route_infos =
-            self.router
-                .route(schema, route_request)
-                .await
-                .context(RouteWithCause {
-                    table_ident: table_ident.clone(),
-                })?;
+        let route_infos = self
+            .router
+            .route(route_request)
+            .await
+            .context(RouteWithCause {
+                table_ident: table_ident.clone(),
+            })?;
 
         if route_infos.is_empty() {
             return RouteNoCause {

diff --git a/router/src/cluster_based.rs b/router/src/cluster_based.rs
@@ -34,9 +34,10 @@ fn make_route(table_name: &str, endpoint: &str) -> Result<Route> {
 
 #[async_trait]
 impl Router for ClusterBasedRouter {
-    async fn route(&self, schema: &str, req: RouteRequest) -> Result<Vec<Route>> {
+    async fn route(&self, req: RouteRequest) -> Result<Vec<Route>> {
+        let req_ctx = req.context.unwrap();
         let route_tables_req = RouteTablesRequest {
-            schema_name: schema.to_string(),
+            schema_name: req_ctx.database,
             table_names: req.tables,
         };
         let route_resp = self

diff --git a/router/src/lib.rs b/router/src/lib.rs
@@ -61,5 +61,5 @@ pub type RouterRef = Arc<dyn Router + Sync + Send>;
 
 #[async_trait]
 pub trait Router {
-    async fn route(&self, schema: &str, req: RouteRequest) -> Result<Vec<Route>>;
+    async fn route(&self, req: RouteRequest) -> Result<Vec<Route>>;
 }
diff --git a/router/src/rule_based.rs b/router/src/rule_based.rs
@@ -138,7 +138,9 @@ impl RuleBasedRouter {
 
 #[async_trait]
 impl Router for RuleBasedRouter {
-    async fn route(&self, schema: &str, req: RouteRequest) -> Result<Vec<Route>> {
+    async fn route(&self, req: RouteRequest) -> Result<Vec<Route>> {
+        let req_ctx = req.context.unwrap();
+        let schema = &req_ctx.database;
         if let Some(shard_nodes) = self.cluster_view.schema_shards.get(schema) {
             ensure!(!shard_nodes.is_empty(), RouteNotFound { schema });
 

diff --git a/server/src/grpc/forward.rs b/server/src/grpc/forward.rs
@@ -9,7 +9,9 @@ use std::{
 };
 
 use async_trait::async_trait;
-use ceresdbproto::storage::{storage_service_client::StorageServiceClient, RouteRequest};
+use ceresdbproto::storage::{
+    storage_service_client::StorageServiceClient, RequestContext, RouteRequest,
+};
 use log::{debug, error, warn};
 use router::{endpoint::Endpoint, RouterRef};
 use serde_derive::Deserialize;
@@ -19,8 +21,6 @@ use tonic::{
     transport::{self, Channel},
 };
 
-use crate::consts::SCHEMA_HEADER;
-
 #[derive(Debug, Snafu)]
 pub enum Error {
     #[snafu(display(
@@ -286,10 +286,11 @@ impl<B: ClientBuilder> Forwarder<B> {
         } = forward_req;
 
         let route_req = RouteRequest {
+            context: Some(RequestContext { database: schema }),
             tables: vec![table],
         };
 
-        let endpoint = match self.router.route(&schema, route_req).await {
+        let endpoint = match self.router.route(route_req).await {
             Ok(mut routes) => {
                 if routes.len() != 1 || routes[0].endpoint.is_none() {
                     warn!(
@@ -315,11 +316,6 @@ impl<B: ClientBuilder> Forwarder<B> {
         {
             // TODO: we should use the timeout from the original request.
             req.set_timeout(self.config.forward_timeout);
-            let metadata = req.metadata_mut();
-            metadata.insert(
-                SCHEMA_HEADER,
-                schema.parse().context(InvalidSchema { schema })?,
-            );
         }
 
         // TODO: add metrics to record the forwarding.
@@ -364,6 +360,7 @@ impl<B: ClientBuilder> Forwarder<B> {
 
 #[cfg(test)]
 mod tests {
+    use catalog::consts::DEFAULT_SCHEMA;
     use ceresdbproto::storage::{Route, SqlQueryRequest, SqlQueryResponse};
     use futures::FutureExt;
     use router::Router;
@@ -397,7 +394,7 @@ mod tests {
 
     #[async_trait]
     impl Router for MockRouter {
-        async fn route(&self, _schema: &str, req: RouteRequest) -> router::Result<Vec<Route>> {
+        async fn route(&self, req: RouteRequest) -> router::Result<Vec<Route>> {
             let endpoint = self.routing_tables.get(&req.tables[0]);
             match endpoint {
                 None => Ok(vec![]),
@@ -463,20 +460,22 @@ mod tests {
 
         let make_forward_req = |table: &str| {
             let query_request = SqlQueryRequest {
+                context: Some(RequestContext {
+                    database: DEFAULT_SCHEMA.to_string(),
+                }),
                 tables: vec![table.to_string()],
                 sql: "".to_string(),
             };
             ForwardRequest {
-                schema: "public".to_string(),
+                schema: DEFAULT_SCHEMA.to_string(),
                 table: table.to_string(),
                 req: query_request.into_request(),
             }
         };
 
         let do_rpc = |_client, req: tonic::Request<SqlQueryRequest>, endpoint: &Endpoint| {
-            let schema = req.metadata().get(SCHEMA_HEADER).unwrap().to_str().unwrap();
-            assert_eq!(schema, "public");
             let req = req.into_inner();
+            assert_eq!(req.context.unwrap().database, DEFAULT_SCHEMA);
             let expect_endpoint = mock_router.routing_tables.get(&req.tables[0]).unwrap();
             assert_eq!(expect_endpoint, endpoint);