add test case

Signed-off-by: junxiang Mu <1948535941@qq.com>

s3select/api/src/lib.rs

@@ -38,6 +38,9 @@ pub enum QueryError {
 
     #[snafu(display("Udf already exists, name:{}.", name))]
     FunctionExists { name: String },
+
+    #[snafu(display("Store Error, e:{}.", e))]
+    StoreError { e: String },
 }
 
 impl From<DataFusionError> for QueryError {

s3select/api/src/query/session.rs

@@ -1,9 +1,12 @@
 use std::sync::Arc;
 
+use bytes::Bytes;
 use datafusion::{
     execution::{context::SessionState, runtime_env::RuntimeEnvBuilder, SessionStateBuilder},
     prelude::SessionContext,
 };
+use object_store::{memory::InMemory, path::Path, ObjectStore};
+use tracing::error;
 
 use crate::{object_store::EcObjectStore, QueryError, QueryResult};
 
@@ -27,11 +30,13 @@ pub struct SessionCtxDesc {
 }
 
 #[derive(Default)]
-pub struct SessionCtxFactory {}
+pub struct SessionCtxFactory {
+    pub is_test: bool,
+}
 
 impl SessionCtxFactory {
-    pub fn create_session_ctx(&self, context: &Context) -> QueryResult<SessionCtx> {
-        let df_session_ctx = self.build_df_session_context(context)?;
+    pub async fn create_session_ctx(&self, context: &Context) -> QueryResult<SessionCtx> {
+        let df_session_ctx = self.build_df_session_context(context).await?;
 
         Ok(SessionCtx {
             _desc: Arc::new(SessionCtxDesc {}),
@@ -39,17 +44,41 @@ impl SessionCtxFactory {
         })
     }
 
-    fn build_df_session_context(&self, context: &Context) -> QueryResult<SessionContext> {
+    async fn build_df_session_context(&self, context: &Context) -> QueryResult<SessionContext> {
         let path = format!("s3://{}", context.input.bucket);
         let store_url = url::Url::parse(&path).unwrap();
-        let store = EcObjectStore::new(context.input.clone()).map_err(|_| QueryError::NotImplemented { err: String::new() })?;
-
         let rt = RuntimeEnvBuilder::new().build()?;
         let df_session_state = SessionStateBuilder::new()
             .with_runtime_env(Arc::new(rt))
-            .with_object_store(&store_url, Arc::new(store))
-            .with_default_features()
-            .build();
+            .with_default_features();
+
+        let df_session_state = if self.is_test {
+            let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
+            let data = b"id,name,age,department,salary
+            1,Alice,25,HR,5000
+            2,Bob,30,IT,6000
+            3,Charlie,35,Finance,7000
+            4,Diana,22,Marketing,4500
+            5,Eve,28,IT,5500
+            6,Frank,40,Finance,8000
+            7,Grace,26,HR,5200
+            8,Henry,32,IT,6200
+            9,Ivy,24,Marketing,4800
+            10,Jack,38,Finance,7500";
+            let data_bytes = Bytes::from(data.to_vec());
+            let path = Path::from(context.input.key.clone());
+            store.put(&path, data_bytes.into()).await.map_err(|e| {
+                error!("put data into memory failed: {}", e.to_string());
+                QueryError::StoreError { e: e.to_string() }
+            })?;
+
+            df_session_state.with_object_store(&store_url, Arc::new(store)).build()
+        } else {
+            let store =
+                EcObjectStore::new(context.input.clone()).map_err(|_| QueryError::NotImplemented { err: String::new() })?;
+            df_session_state.with_object_store(&store_url, Arc::new(store)).build()
+        };
+
         let df_session_ctx = SessionContext::new_with_state(df_session_state);
 
         Ok(df_session_ctx)

s3select/query/src/dispatcher/manager.rs

@@ -98,7 +98,7 @@ impl QueryDispatcher for SimpleQueryDispatcher {
     }
 
     async fn build_query_state_machine(&self, query: Query) -> QueryResult<Arc<QueryStateMachine>> {
-        let session = self.session_factory.create_session_ctx(query.context())?;
+        let session = self.session_factory.create_session_ctx(query.context()).await?;
 
         let query_state_machine = Arc::new(QueryStateMachine::begin(query, session));
         Ok(query_state_machine)

s3select/query/src/instance.rs

@@ -63,11 +63,11 @@ where
     }
 }
 
-pub async fn make_cnosdbms(input: SelectObjectContentInput) -> QueryResult<impl DatabaseManagerSystem> {
+pub async fn make_cnosdbms(input: SelectObjectContentInput, is_test: bool) -> QueryResult<impl DatabaseManagerSystem> {
     // init Function Manager, we can define some UDF if need
     let func_manager = SimpleFunctionMetadataManager::default();
     // TODO session config need load global system config
-    let session_factory = Arc::new(SessionCtxFactory {});
+    let session_factory = Arc::new(SessionCtxFactory { is_test });
     let parser = Arc::new(DefaultParser::default());
     let optimizer = Arc::new(CascadeOptimizerBuilder::default().build());
     // TODO wrap, and num_threads configurable
@@ -92,3 +92,73 @@ pub async fn make_cnosdbms(input: SelectObjectContentInput) -> QueryResult<impl
 
     Ok(db_server)
 }
+
+#[cfg(test)]
+mod tests {
+    use api::{
+        query::{Context, Query},
+        server::dbms::DatabaseManagerSystem,
+    };
+    use datafusion::{arrow::util::pretty, assert_batches_eq};
+    use s3s::dto::{
+        CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentInput,
+        SelectObjectContentRequest,
+    };
+
+    use crate::instance::make_cnosdbms;
+
+    #[tokio::test]
+    #[ignore]
+    async fn test_simple_sql() {
+        let sql = "select * from S3Object";
+        let input = SelectObjectContentInput {
+            bucket: "dandan".to_string(),
+            expected_bucket_owner: None,
+            key: "test.csv".to_string(),
+            sse_customer_algorithm: None,
+            sse_customer_key: None,
+            sse_customer_key_md5: None,
+            request: SelectObjectContentRequest {
+                expression: sql.to_string(),
+                expression_type: ExpressionType::from_static("SQL"),
+                input_serialization: InputSerialization {
+                    csv: Some(CSVInput::default()),
+                    ..Default::default()
+                },
+                output_serialization: OutputSerialization {
+                    csv: Some(CSVOutput::default()),
+                    ..Default::default()
+                },
+                request_progress: None,
+                scan_range: None,
+            },
+        };
+        let db = make_cnosdbms(input.clone(), true).await.unwrap();
+        let query = Query::new(Context { input }, sql.to_string());
+
+        let result = db.execute(&query).await.unwrap();
+
+        let results = result.result().chunk_result().await.unwrap().to_vec();
+
+        let expected = [
+            "+----------------+----------+----------+------------+----------+",
+            "| column_1       | column_2 | column_3 | column_4   | column_5 |",
+            "+----------------+----------+----------+------------+----------+",
+            "| id             | name     | age      | department | salary   |",
+            "|             1  | Alice    | 25       | HR         | 5000     |",
+            "|             2  | Bob      | 30       | IT         | 6000     |",
+            "|             3  | Charlie  | 35       | Finance    | 7000     |",
+            "|             4  | Diana    | 22       | Marketing  | 4500     |",
+            "|             5  | Eve      | 28       | IT         | 5500     |",
+            "|             6  | Frank    | 40       | Finance    | 8000     |",
+            "|             7  | Grace    | 26       | HR         | 5200     |",
+            "|             8  | Henry    | 32       | IT         | 6200     |",
+            "|             9  | Ivy      | 24       | Marketing  | 4800     |",
+            "|             10 | Jack     | 38       | Finance    | 7500     |",
+            "+----------------+----------+----------+------------+----------+",
+        ];
+
+        assert_batches_eq!(expected, &results);
+        pretty::print_batches(&results).unwrap();
+    }
+}

s3select/query/src/sql/optimizer.rs

@@ -36,10 +36,7 @@ impl Optimizer for CascadeOptimizer {
 
         debug!("Original physical plan:\n{}\n", displayable(physical_plan.as_ref()).indent(false));
 
-        let optimized_physical_plan = {
-            self.physical_optimizer
-                .optimize(physical_plan, session)?
-        };
+        let optimized_physical_plan = { self.physical_optimizer.optimize(physical_plan, session)? };
 
         Ok(optimized_physical_plan)
     }