fix: add router when build request context for mysql (apache#809)

Author: jiacai2050

.github/workflows/ci.yml

@@ -217,6 +217,10 @@ jobs:
         working-directory: integration_tests
         run: |
           make run-rust
+      - name: Run MySQL client tests
+        working-directory: integration_tests
+        run: |
+          make run-mysql
       - name: Upload Logs
         if: always()
         uses: actions/upload-artifact@v3

Makefile

@@ -2,6 +2,8 @@ SHELL = /bin/bash
 
 DIR=$(shell pwd)
 
+.DEFAULT_GOAL := integration-test
+
 init:
 	echo "init"
 	echo "Git branch: $GITBRANCH"

integration_tests/Makefile

@@ -55,3 +55,6 @@ run-go:
 
 run-rust:
 	cd sdk/rust && cargo run
+
+run-mysql:
+	cd mysql && ./basic.sh

integration_tests/mysql/basic.sh

@@ -0,0 +1,7 @@
+#!/usr/bin/env bash
+
+# This only ensure query by mysql protocol is OK,
+# Full SQL test in ensured by sqlness tests.
+mysql -h 127.0.0.1 -P 3307 -e 'show tables'
+
+mysql -h 127.0.0.1 -P 3307 -e 'select 1, now()'

server/src/mysql/builder.rs

@@ -3,13 +3,14 @@
 use std::{net::SocketAddr, sync::Arc, time::Duration};
 
 use query_engine::executor::Executor as QueryExecutor;
+use router::RouterRef;
 use snafu::{OptionExt, ResultExt};
 use table_engine::engine::EngineRuntimes;
 
 use crate::{
     instance::InstanceRef,
     mysql::{
-        error::{MissingInstance, MissingRuntimes, ParseIpAddr, Result},
+        error::{MissingInstance, MissingRouter, MissingRuntimes, ParseIpAddr, Result},
         service::MysqlService,
     },
 };
@@ -18,6 +19,7 @@ pub struct Builder<Q> {
     config: Config,
     runtimes: Option<Arc<EngineRuntimes>>,
     instance: Option<InstanceRef<Q>>,
+    router: Option<RouterRef>,
 }
 
 #[derive(Debug)]
@@ -33,6 +35,7 @@ impl<Q> Builder<Q> {
             config,
             runtimes: None,
             instance: None,
+            router: None,
         }
     }
 
@@ -45,18 +48,25 @@ impl<Q> Builder<Q> {
         self.instance = Some(instance);
         self
     }
+
+    pub fn router(mut self, router: RouterRef) -> Self {
+        self.router = Some(router);
+        self
+    }
 }
 
 impl<Q: QueryExecutor + 'static> Builder<Q> {
     pub fn build(self) -> Result<MysqlService<Q>> {
         let runtimes = self.runtimes.context(MissingRuntimes)?;
         let instance = self.instance.context(MissingInstance)?;
+        let router = self.router.context(MissingRouter)?;
 
         let addr: SocketAddr = format!("{}:{}", self.config.ip, self.config.port)
             .parse()
             .context(ParseIpAddr { ip: self.config.ip })?;
 
-        let mysql_handler = MysqlService::new(instance, runtimes, addr, self.config.timeout);
+        let mysql_handler =
+            MysqlService::new(instance, runtimes, router, addr, self.config.timeout);
         Ok(mysql_handler)
     }
 }

server/src/mysql/error.rs

@@ -12,6 +12,9 @@ pub enum Error {
     #[snafu(display("Missing instance to build service.\nBacktrace:\n{}", backtrace))]
     MissingInstance { backtrace: Backtrace },
 
+    #[snafu(display("Missing router to build service.\nBacktrace:\n{}", backtrace))]
+    MissingRouter { backtrace: Backtrace },
+
     #[snafu(display(
         "Failed to parse ip addr, ip:{}, err:{}.\nBacktrace:\n{}",
         ip,

server/src/mysql/service.rs

@@ -6,6 +6,7 @@ use common_util::runtime::JoinHandle;
 use log::{error, info};
 use opensrv_mysql::AsyncMysqlIntermediary;
 use query_engine::executor::Executor as QueryExecutor;
+use router::RouterRef;
 use table_engine::engine::EngineRuntimes;
 use tokio::sync::oneshot::{self, Receiver, Sender};
 
@@ -17,6 +18,7 @@ use crate::{
 pub struct MysqlService<Q> {
     instance: InstanceRef<Q>,
     runtimes: Arc<EngineRuntimes>,
+    router: RouterRef,
     socket_addr: SocketAddr,
     join_handler: Option<JoinHandle<()>>,
     tx: Option<Sender<()>>,
@@ -27,12 +29,14 @@ impl<Q> MysqlService<Q> {
     pub fn new(
         instance: Arc<Instance<Q>>,
         runtimes: Arc<EngineRuntimes>,
+        router: RouterRef,
         socket_addr: SocketAddr,
         timeout: Option<Duration>,
     ) -> MysqlService<Q> {
         Self {
             instance,
             runtimes,
+            router,
             socket_addr,
             join_handler: None,
             tx: None,
@@ -53,6 +57,7 @@ impl<Q: QueryExecutor + 'static> MysqlService<Q> {
         self.join_handler = Some(rt.bg_runtime.spawn(Self::loop_accept(
             self.instance.clone(),
             self.runtimes.clone(),
+            self.router.clone(),
             self.socket_addr,
             self.timeout,
             rx,
@@ -69,6 +74,7 @@ impl<Q: QueryExecutor + 'static> MysqlService<Q> {
     async fn loop_accept(
         instance: InstanceRef<Q>,
         runtimes: Arc<EngineRuntimes>,
+        router: RouterRef,
         socket_addr: SocketAddr,
         timeout: Option<Duration>,
         mut rx: Receiver<()>,
@@ -90,10 +96,11 @@ impl<Q: QueryExecutor + 'static> MysqlService<Q> {
                     };
                     let instance = instance.clone();
                     let runtimes = runtimes.clone();
+                    let router = router.clone();
 
                     let rt = runtimes.read_runtime.clone();
                     rt.spawn(AsyncMysqlIntermediary::run_on(
-                        MysqlWorker::new(instance, runtimes, timeout),
+                        MysqlWorker::new(instance, runtimes, router, timeout),
                         stream,
                     ));
                 },

server/src/mysql/worker.rs

@@ -6,6 +6,7 @@ use interpreters::interpreter::Output;
 use log::{error, info};
 use opensrv_mysql::{AsyncMysqlShim, ErrorKind, QueryResultWriter, StatementMetaWriter};
 use query_engine::executor::Executor as QueryExecutor;
+use router::RouterRef;
 use snafu::ResultExt;
 use table_engine::engine::EngineRuntimes;
 
@@ -26,6 +27,7 @@ pub struct MysqlWorker<W: std::io::Write + Send + Sync, Q> {
     generic_hold: PhantomData<W>,
     instance: Arc<Instance<Q>>,
     runtimes: Arc<EngineRuntimes>,
+    router: RouterRef,
     timeout: Option<Duration>,
 }
 
@@ -37,12 +39,14 @@ where
     pub fn new(
         instance: Arc<Instance<Q>>,
         runtimes: Arc<EngineRuntimes>,
+        router: RouterRef,
         timeout: Option<Duration>,
     ) -> Self {
         Self {
             generic_hold: PhantomData::default(),
             instance,
             runtimes,
+            router,
             timeout,
         }
     }
@@ -144,6 +148,7 @@ where
             .runtime(runtime)
             .enable_partition_table_access(false)
             .timeout(self.timeout)
+            .router(self.router.clone())
             .build()
             .context(CreateContext)
     }

server/src/server.rs

@@ -354,6 +354,7 @@ impl<Q: QueryExecutor + 'static> Builder<Q> {
         let mysql_service = mysql::Builder::new(mysql_config)
             .runtimes(engine_runtimes.clone())
             .instance(instance.clone())
+            .router(router.clone())
             .build()
             .context(BuildMysqlService)?;