tmp2

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

s3select/api/src/lib.rs

@@ -1,8 +1,9 @@
 use std::fmt::Display;
 
-use datafusion::common::DataFusionError;
+use datafusion::{common::DataFusionError, sql::sqlparser::parser::ParserError};
 use snafu::{Backtrace, Location, Snafu};
 
+pub mod object_store;
 pub mod query;
 pub mod server;
 
@@ -16,6 +17,21 @@ pub enum QueryError {
         location: Location,
         backtrace: Backtrace,
     },
+
+    #[snafu(display("This feature is not implemented: {}", err))]
+    NotImplemented { err: String },
+
+    #[snafu(display("Multi-statement not allow, found num:{}, sql:{}", num, sql))]
+    MultiStatement { num: usize, sql: String },
+
+    #[snafu(display("Failed to build QueryDispatcher. err: {}", err))]
+    BuildQueryDispatcher { err: String },
+
+    #[snafu(display("The query has been canceled"))]
+    Cancel,
+
+    #[snafu(display("{}", source))]
+    Parser { source: ParserError },
 }
 
 impl From<DataFusionError> for QueryError {

s3select/api/src/object_store.rs

@@ -0,0 +1,150 @@
+use async_trait::async_trait;
+use bytes::Bytes;
+use chrono::Utc;
+use ecstore::new_object_layer_fn;
+use ecstore::store::ECStore;
+use ecstore::store_api::ObjectIO;
+use ecstore::store_api::ObjectOptions;
+use ecstore::StorageAPI;
+use futures::stream;
+use futures::StreamExt;
+use futures_core::stream::BoxStream;
+use http::HeaderMap;
+use object_store::path::Path;
+use object_store::Attributes;
+use object_store::GetOptions;
+use object_store::GetResult;
+use object_store::ListResult;
+use object_store::MultipartUpload;
+use object_store::ObjectMeta;
+use object_store::ObjectStore;
+use object_store::PutMultipartOpts;
+use object_store::PutOptions;
+use object_store::PutPayload;
+use object_store::PutResult;
+use object_store::{Error as o_Error, Result};
+use s3s::dto::SelectObjectContentInput;
+use s3s::s3_error;
+use s3s::S3Result;
+use std::ops::Range;
+use std::sync::Arc;
+use tracing::info;
+
+#[derive(Debug)]
+pub struct EcObjectStore {
+    input: SelectObjectContentInput,
+
+    store: Arc<ECStore>,
+}
+
+impl EcObjectStore {
+    pub fn new(input: SelectObjectContentInput) -> S3Result<Self> {
+        let Some(store) = new_object_layer_fn() else {
+            return Err(s3_error!(InternalError, "ec store not inited"));
+        };
+
+        Ok(Self { input, store })
+    }
+}
+
+impl std::fmt::Display for EcObjectStore {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("EcObjectStore")
+    }
+}
+
+#[async_trait]
+impl ObjectStore for EcObjectStore {
+    async fn put_opts(&self, _location: &Path, _payload: PutPayload, _opts: PutOptions) -> Result<PutResult> {
+        unimplemented!()
+    }
+
+    async fn put_multipart_opts(&self, _location: &Path, _opts: PutMultipartOpts) -> Result<Box<dyn MultipartUpload>> {
+        unimplemented!()
+    }
+
+    async fn get_opts(&self, location: &Path, _options: GetOptions) -> Result<GetResult> {
+        info!("{:?}", location);
+        let opts = ObjectOptions::default();
+        let h = HeaderMap::new();
+        let reader = self
+            .store
+            .get_object_reader(&self.input.bucket, &self.input.key, None, h, &opts)
+            .await
+            .map_err(|_| o_Error::NotFound {
+                path: format!("{}/{}", self.input.bucket, self.input.key),
+                source: "can not get object info".into(),
+            })?;
+
+        let stream = stream::unfold(reader.stream, |mut blob| async move {
+            match blob.next().await {
+                Some(Ok(chunk)) => {
+                    let bytes = Bytes::from(chunk);
+                    Some((Ok(bytes), blob))
+                }
+                _ => None,
+            }
+        })
+        .boxed();
+        let meta = ObjectMeta {
+            location: location.clone(),
+            last_modified: Utc::now(),
+            size: reader.object_info.size,
+            e_tag: reader.object_info.etag,
+            version: None,
+        };
+        let attributes = Attributes::default();
+
+        Ok(GetResult {
+            payload: object_store::GetResultPayload::Stream(stream),
+            meta,
+            range: 0..reader.object_info.size,
+            attributes,
+        })
+    }
+
+    async fn get_ranges(&self, _location: &Path, _ranges: &[Range<usize>]) -> Result<Vec<Bytes>> {
+        unimplemented!()
+    }
+
+    async fn head(&self, location: &Path) -> Result<ObjectMeta> {
+        info!("{:?}", location);
+        let opts = ObjectOptions::default();
+        let info = self
+            .store
+            .get_object_info(&self.input.bucket, &self.input.key, &opts)
+            .await
+            .map_err(|_| o_Error::NotFound {
+                path: format!("{}/{}", self.input.bucket, self.input.key),
+                source: "can not get object info".into(),
+            })?;
+
+        Ok(ObjectMeta {
+            location: location.clone(),
+            last_modified: Utc::now(),
+            size: info.size,
+            e_tag: info.etag,
+            version: None,
+        })
+    }
+
+    async fn delete(&self, _location: &Path) -> Result<()> {
+        unimplemented!()
+    }
+
+    fn list(&self, _prefix: Option<&Path>) -> BoxStream<'_, Result<ObjectMeta>> {
+        unimplemented!()
+    }
+
+    async fn list_with_delimiter(&self, _prefix: Option<&Path>) -> Result<ListResult> {
+        unimplemented!()
+    }
+
+    async fn copy(&self, _from: &Path, _to: &Path) -> Result<()> {
+        unimplemented!()
+    }
+
+    async fn copy_if_not_exists(&self, _from: &Path, _too: &Path) -> Result<()> {
+        unimplemented!()
+    }
+}

s3select/api/src/query/analyzer.rs

@@ -0,0 +1,12 @@
+use std::sync::Arc;
+
+use datafusion::logical_expr::LogicalPlan;
+
+use super::session::SessionCtx;
+use crate::QueryResult;
+
+pub type AnalyzerRef = Arc<dyn Analyzer + Send + Sync>;
+
+pub trait Analyzer {
+    fn analyze(&self, plan: &LogicalPlan, session: &SessionCtx) -> QueryResult<LogicalPlan>;
+}

s3select/api/src/query/dispatcher.rs

@@ -12,10 +12,6 @@ use super::{
 
 #[async_trait]
 pub trait QueryDispatcher: Send + Sync {
-    async fn start(&self) -> QueryResult<()>;
-
-    fn stop(&self);
-
     // fn create_query_id(&self) -> QueryId;
 
     // fn query_info(&self, id: &QueryId);

s3select/api/src/query/execution.rs

@@ -43,18 +43,6 @@ pub trait QueryExecution: Send + Sync {
     async fn start(&self) -> QueryResult<Output>;
     // 停止
     fn cancel(&self) -> QueryResult<()>;
-    // query状态
-    // 查询计划
-    // 静态信息
-    // fn info(&self) -> QueryInfo;
-    // 运行时信息
-    // fn status(&self) -> QueryStatus;
-    // sql
-    // 资源占用（cpu时间/内存/吞吐量等）
-    // 是否需要持久化query信息
-    fn need_persist(&self) -> bool {
-        false
-    }
 }
 
 pub enum Output {

s3select/api/src/query/logical_planner.rs

@@ -1,6 +1,12 @@
+use async_trait::async_trait;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::logical_expr::LogicalPlan as DFPlan;
 
+use crate::QueryResult;
+
+use super::ast::ExtStatement;
+use super::session::SessionCtx;
+
 #[derive(Clone)]
 pub enum Plan {
     // only support query sql
@@ -27,3 +33,8 @@ impl QueryPlan {
         matches!(self.df_plan, DFPlan::Explain(_) | DFPlan::Analyze(_))
     }
 }
+
+#[async_trait]
+pub trait LogicalPlanner {
+    async fn create_logical_plan(&self, statement: ExtStatement, session: &SessionCtx) -> QueryResult<Plan>;
+}

s3select/api/src/query/mod.rs

@@ -1,15 +1,22 @@
+use s3s::dto::SelectObjectContentInput;
+
+pub mod analyzer;
 pub mod ast;
 pub mod datasource;
 pub mod dispatcher;
 pub mod execution;
 pub mod function;
 pub mod logical_planner;
+pub mod optimizer;
 pub mod parser;
+pub mod physical_planner;
+pub mod scheduler;
 pub mod session;
 
 #[derive(Clone)]
 pub struct Context {
     // maybe we need transfer some info?
+    pub input: SelectObjectContentInput,
 }
 
 #[derive(Clone)]
@@ -17,3 +24,18 @@ pub struct Query {
     context: Context,
     content: String,
 }
+
+impl Query {
+    #[inline(always)]
+    pub fn new(context: Context, content: String) -> Self {
+        Self { context, content }
+    }
+
+    pub fn context(&self) -> &Context {
+        &self.context
+    }
+
+    pub fn content(&self) -> &str {
+        self.content.as_str()
+    }
+}

s3select/api/src/query/optimizer.rs

@@ -0,0 +1,15 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::physical_plan::ExecutionPlan;
+
+use super::logical_planner::QueryPlan;
+use super::session::SessionCtx;
+use crate::QueryResult;
+
+pub type OptimizerRef = Arc<dyn Optimizer + Send + Sync>;
+
+#[async_trait]
+pub trait Optimizer {
+    async fn optimize(&self, plan: &QueryPlan, session: &SessionCtx) -> QueryResult<Arc<dyn ExecutionPlan>>;
+}

s3select/api/src/query/physical_planner.rs

@@ -0,0 +1,21 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_planner::ExtensionPlanner;
+
+use super::session::SessionCtx;
+use crate::QueryResult;
+
+#[async_trait]
+pub trait PhysicalPlanner {
+    /// Given a `LogicalPlan`, create an `ExecutionPlan` suitable for execution
+    async fn create_physical_plan(
+        &self,
+        logical_plan: &LogicalPlan,
+        session_state: &SessionCtx,
+    ) -> QueryResult<Arc<dyn ExecutionPlan>>;
+
+    fn inject_physical_transform_rule(&mut self, rule: Arc<dyn ExtensionPlanner + Send + Sync>);
+}

s3select/api/src/query/scheduler.rs

@@ -0,0 +1,32 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::common::Result;
+use datafusion::execution::context::TaskContext;
+use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream};
+
+pub type SchedulerRef = Arc<dyn Scheduler + Send + Sync>;
+
+#[async_trait]
+pub trait Scheduler {
+    /// Schedule the provided [`ExecutionPlan`] on this [`Scheduler`].
+    ///
+    /// Returns a [`ExecutionResults`] that can be used to receive results as they are produced,
+    /// as a [`futures::Stream`] of [`RecordBatch`]
+    async fn schedule(&self, plan: Arc<dyn ExecutionPlan>, context: Arc<TaskContext>) -> Result<ExecutionResults>;
+}
+
+pub struct ExecutionResults {
+    stream: SendableRecordBatchStream,
+}
+
+impl ExecutionResults {
+    pub fn new(stream: SendableRecordBatchStream) -> Self {
+        Self { stream }
+    }
+
+    /// Returns a [`SendableRecordBatchStream`] of this execution
+    pub fn stream(self) -> SendableRecordBatchStream {
+        self.stream
+    }
+}

s3select/api/src/query/session.rs

@@ -1,10 +1,17 @@
 use std::sync::Arc;
 
-use datafusion::execution::context::SessionState;
+use datafusion::{
+    execution::{context::SessionState, runtime_env::RuntimeEnvBuilder, SessionStateBuilder},
+    prelude::SessionContext,
+};
+
+use crate::{object_store::EcObjectStore, QueryError, QueryResult};
+
+use super::Context;
 
 #[derive(Clone)]
 pub struct SessionCtx {
-    desc: Arc<SessionCtxDesc>,
+    _desc: Arc<SessionCtxDesc>,
     inner: SessionState,
 }
 
@@ -18,3 +25,33 @@ impl SessionCtx {
 pub struct SessionCtxDesc {
     // maybe we need some info
 }
+
+#[derive(Default)]
+pub struct SessionCtxFactory {}
+
+impl SessionCtxFactory {
+    pub fn create_session_ctx(&self, context: &Context) -> QueryResult<SessionCtx> {
+        let df_session_ctx = self.build_df_session_context(context)?;
+
+        Ok(SessionCtx {
+            _desc: Arc::new(SessionCtxDesc {}),
+            inner: df_session_ctx.state(),
+        })
+    }
+
+    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();
+        let df_session_ctx = SessionContext::new_with_state(df_session_state);
+
+        Ok(df_session_ctx)
+    }
+}

s3select/api/src/server/dbms.rs

@@ -30,7 +30,6 @@ impl QueryHandle {
 
 #[async_trait]
 pub trait DatabaseManagerSystem {
-    async fn start(&self) -> QueryResult<()>;
     async fn execute(&self, query: &Query) -> QueryResult<QueryHandle>;
     async fn build_query_state_machine(&self, query: Query) -> QueryResult<QueryStateMachineRef>;
     async fn build_logical_plan(&self, query_state_machine: QueryStateMachineRef) -> QueryResult<Option<Plan>>;
@@ -39,5 +38,4 @@ pub trait DatabaseManagerSystem {
         logical_plan: Plan,
         query_state_machine: QueryStateMachineRef,
     ) -> QueryResult<QueryHandle>;
-    fn metrics(&self) -> String;
 }