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)