feat: support multi table engines (#1277)

* feat: support multi table engines

* refactor: adapt SqlHandler to support multiple table engines

* refactor: refactor TableEngineManager

* chore: apply review suggestions

* chore: apply review suggestions

* chore: apply review suggestions

* chore: snafu context styling

src/table/src/engine.rs

@@ -22,6 +22,7 @@ use crate::error::Result;
 use crate::metadata::TableId;
 use crate::requests::{AlterTableRequest, CreateTableRequest, DropTableRequest, OpenTableRequest};
 use crate::TableRef;
+pub mod manager;
 
 /// Represents a resolved path to a table of the form “catalog.schema.table”
 #[derive(Debug, PartialEq)]

src/table/src/engine/manager.rs

@@ -0,0 +1,127 @@
+// Copyright 2023 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::collections::HashMap;
+use std::sync::{Arc, RwLock};
+
+use async_trait::async_trait;
+use common_telemetry::error;
+use snafu::{ensure, OptionExt};
+
+use crate::engine::TableEngineRef;
+use crate::error::{EngineExistSnafu, EngineNotFoundSnafu, Result};
+
+#[async_trait::async_trait]
+pub trait TableEngineManager: Send + Sync {
+    /// returns `error::EngineNotFound` if engine not found
+    fn engine(&self, name: &str) -> Result<TableEngineRef>;
+
+    /// returns `error::EngineExist` if engine exists
+    fn register_engine(&self, name: &str, engine: TableEngineRef) -> Result<()>;
+
+    /// closes all registered engines
+    async fn close(&self) -> Result<()>;
+}
+pub type TableEngineManagerRef = Arc<dyn TableEngineManager>;
+
+/// Simple in-memory table engine manager
+pub struct MemoryTableEngineManager {
+    pub engines: RwLock<HashMap<String, TableEngineRef>>,
+}
+
+impl MemoryTableEngineManager {
+    pub fn new(engine: TableEngineRef) -> Self {
+        MemoryTableEngineManager::alias(engine.name().to_string(), engine)
+    }
+
+    pub fn alias(name: String, engine: TableEngineRef) -> Self {
+        let mut engines = HashMap::new();
+        engines.insert(name, engine);
+        let engines = RwLock::new(engines);
+
+        MemoryTableEngineManager { engines }
+    }
+}
+
+#[async_trait]
+impl TableEngineManager for MemoryTableEngineManager {
+    fn engine(&self, name: &str) -> Result<TableEngineRef> {
+        let engines = self.engines.read().unwrap();
+        engines
+            .get(name)
+            .cloned()
+            .context(EngineNotFoundSnafu { engine: name })
+    }
+
+    fn register_engine(&self, name: &str, engine: TableEngineRef) -> Result<()> {
+        let mut engines = self.engines.write().unwrap();
+
+        ensure!(
+            !engines.contains_key(name),
+            EngineExistSnafu { engine: name }
+        );
+
+        engines.insert(name.to_string(), engine);
+
+        Ok(())
+    }
+
+    async fn close(&self) -> Result<()> {
+        let engines = {
+            let engines = self.engines.write().unwrap();
+            engines.values().cloned().collect::<Vec<_>>()
+        };
+
+        if let Err(err) =
+            futures::future::try_join_all(engines.iter().map(|engine| engine.close())).await
+        {
+            error!("Failed to close engine: {}", err);
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::assert_matches::assert_matches;
+
+    use super::*;
+    use crate::engine::TableEngine;
+    use crate::error;
+    use crate::test_util::MockTableEngine;
+
+    #[test]
+    fn test_table_engine_manager() {
+        let table_engine = MockTableEngine::new();
+        let table_engine_ref = Arc::new(table_engine);
+        let table_engine_manager = MemoryTableEngineManager::new(table_engine_ref.clone());
+
+        table_engine_manager
+            .register_engine("yet_another", table_engine_ref.clone())
+            .unwrap();
+
+        let got = table_engine_manager.engine(table_engine_ref.name());
+
+        assert_eq!(got.unwrap().name(), table_engine_ref.name());
+
+        let got = table_engine_manager.engine("yet_another");
+
+        assert_eq!(got.unwrap().name(), table_engine_ref.name());
+
+        let missing = table_engine_manager.engine("not_exists");
+
+        assert_matches!(missing.err().unwrap(), error::Error::EngineNotFound { .. })
+    }
+}

src/table/src/error.rs

@@ -43,6 +43,18 @@ pub enum Error {
         backtrace: Backtrace,
     },
 
+    #[snafu(display("Engine not found: {}", engine))]
+    EngineNotFound {
+        engine: String,
+        backtrace: Backtrace,
+    },
+
+    #[snafu(display("Engine exist: {}", engine))]
+    EngineExist {
+        engine: String,
+        backtrace: Backtrace,
+    },
+
     #[snafu(display("Table projection error, source: {}", source))]
     TableProjection {
         source: ArrowError,
@@ -133,7 +145,9 @@ impl ErrorExt for Error {
             Error::ColumnNotExists { .. } => StatusCode::TableColumnNotFound,
             Error::RegionSchemaMismatch { .. } => StatusCode::StorageUnavailable,
             Error::Unsupported { .. } => StatusCode::Unsupported,
-            Error::ParseTableOption { .. } => StatusCode::InvalidArguments,
+            Error::ParseTableOption { .. }
+            | Error::EngineNotFound { .. }
+            | Error::EngineExist { .. } => StatusCode::InvalidArguments,
         }
     }
 

src/table/src/lib.rs

@@ -11,6 +11,7 @@
 // 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.
+#![feature(assert_matches)]
 
 pub mod engine;
 pub mod error;

src/table/src/requests.rs

@@ -47,6 +47,7 @@ pub struct CreateTableRequest {
     pub primary_key_indices: Vec<usize>,
     pub create_if_not_exists: bool,
     pub table_options: TableOptions,
+    pub engine: String,
 }
 
 #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]

src/table/src/table/numbers.rs

@@ -38,10 +38,16 @@ use crate::table::{Expr, Table};
 pub struct NumbersTable {
     table_id: TableId,
     schema: SchemaRef,
+    name: String,
+    engine: String,
 }
 
 impl NumbersTable {
     pub fn new(table_id: TableId) -> Self {
+        NumbersTable::with_name(table_id, "numbers".to_string())
+    }
+
+    pub fn with_name(table_id: TableId, name: String) -> Self {
         let column_schemas = vec![ColumnSchema::new(
             "number",
             ConcreteDataType::uint32_datatype(),
@@ -49,6 +55,8 @@ impl NumbersTable {
         )];
         Self {
             table_id,
+            name,
+            engine: "test_engine".to_string(),
             schema: Arc::new(
                 SchemaBuilder::try_from_columns(column_schemas)
                     .unwrap()
@@ -79,7 +87,7 @@ impl Table for NumbersTable {
         Arc::new(
             TableInfoBuilder::default()
                 .table_id(self.table_id)
-                .name("numbers")
+                .name(&self.name)
                 .catalog_name("greptime")
                 .schema_name("public")
                 .table_version(0)
@@ -90,6 +98,7 @@ impl Table for NumbersTable {
                         .region_numbers(vec![0])
                         .primary_key_indices(vec![0])
                         .next_column_id(1)
+                        .engine(&self.engine)
                         .build()
                         .unwrap(),
                 )