rustfs/crates/s3select-api/src/lib.rs

// Copyright 2024 RustFS 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 datafusion::{common::DataFusionError, sql::sqlparser::parser::ParserError};
use snafu::{Backtrace, Location, Snafu};
use std::fmt::Display;
use std::sync::Arc;

pub(crate) use rustfs_ecstore::api::error::StorageError as SelectStorageError;
pub(crate) use rustfs_ecstore::api::storage::ECStore as SelectStore;

pub mod object_store;
pub mod query;
pub mod server;

#[cfg(test)]
mod test;

pub type QueryResult<T> = Result<T, QueryError>;

pub(crate) type SelectGetObjectReader = <SelectStore as rustfs_storage_api::ObjectIO>::GetObjectReader;
pub(crate) type SelectObjectInfo = <SelectStore as rustfs_storage_api::ObjectOperations>::ObjectInfo;
pub(crate) type SelectObjectOptions = <SelectStore as rustfs_storage_api::ObjectOperations>::ObjectOptions;

pub(crate) const SELECT_DEFAULT_READ_BUFFER_SIZE: usize = rustfs_ecstore::api::set_disk::DEFAULT_READ_BUFFER_SIZE;

pub(crate) fn resolve_select_object_store_handle() -> Option<Arc<SelectStore>> {
    rustfs_ecstore::api::global::resolve_object_store_handle()
}

pub(crate) fn select_is_err_bucket_not_found(err: &SelectStorageError) -> bool {
    rustfs_ecstore::api::error::is_err_bucket_not_found(err)
}

pub(crate) fn select_is_err_object_not_found(err: &SelectStorageError) -> bool {
    rustfs_ecstore::api::error::is_err_object_not_found(err)
}

pub(crate) fn select_is_err_version_not_found(err: &SelectStorageError) -> bool {
    rustfs_ecstore::api::error::is_err_version_not_found(err)
}

#[derive(Debug, Snafu)]
#[snafu(visibility(pub))]
pub enum QueryError {
    #[snafu(display("DataFusion error: {}", source))]
    Datafusion {
        source: Box<DataFusionError>,
        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 },

    #[snafu(display("Udf not exists, name:{}.", name))]
    FunctionNotExists { name: String },

    #[snafu(display("Udf already exists, name:{}.", name))]
    FunctionExists { name: String },

    #[snafu(display("Store Error, e:{}.", e))]
    StoreError { e: String },
}

impl From<DataFusionError> for QueryError {
    #[track_caller]
    fn from(value: DataFusionError) -> Self {
        match value {
            DataFusionError::External(e) if e.downcast_ref::<QueryError>().is_some() => *e.downcast::<QueryError>().unwrap(),

            v => Self::Datafusion {
                source: Box::new(v),
                location: std::panic::Location::caller(),
                backtrace: Backtrace::capture(),
            },
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResolvedTable {
    // path
    table: String,
}

impl ResolvedTable {
    pub fn table(&self) -> &str {
        &self.table
    }
}

impl Display for ResolvedTable {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let Self { table } = self;
        write!(f, "{table}")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use datafusion::common::DataFusionError;
    use datafusion::sql::sqlparser::parser::ParserError;

    #[test]
    fn test_query_error_display() {
        let err = QueryError::NotImplemented {
            err: "feature X".to_string(),
        };
        assert_eq!(err.to_string(), "This feature is not implemented: feature X");

        let err = QueryError::MultiStatement {
            num: 2,
            sql: "SELECT 1; SELECT 2;".to_string(),
        };
        assert_eq!(err.to_string(), "Multi-statement not allow, found num:2, sql:SELECT 1; SELECT 2;");

        let err = QueryError::Cancel;
        assert_eq!(err.to_string(), "The query has been canceled");

        let err = QueryError::FunctionNotExists {
            name: "my_func".to_string(),
        };
        assert_eq!(err.to_string(), "Udf not exists, name:my_func.");

        let err = QueryError::StoreError {
            e: "connection failed".to_string(),
        };
        assert_eq!(err.to_string(), "Store Error, e:connection failed.");
    }

    #[test]
    fn test_query_error_from_datafusion_error() {
        let df_error = DataFusionError::Plan("invalid plan".to_string());
        let query_error: QueryError = df_error.into();

        match query_error {
            QueryError::Datafusion { source, .. } => {
                assert!(source.to_string().contains("invalid plan"));
            }
            _ => panic!("Expected Datafusion error"),
        }
    }

    #[test]
    fn test_query_error_from_parser_error() {
        let parser_error = ParserError::ParserError("syntax error".to_string());
        let query_error = QueryError::Parser { source: parser_error };

        assert!(query_error.to_string().contains("syntax error"));
    }

    #[test]
    fn test_resolved_table() {
        let table = ResolvedTable {
            table: "my_table".to_string(),
        };

        assert_eq!(table.table(), "my_table");
        assert_eq!(table.to_string(), "my_table");
    }

    #[test]
    fn test_resolved_table_clone_and_eq() {
        let table1 = ResolvedTable {
            table: "table1".to_string(),
        };
        let table2 = table1.clone();
        let table3 = ResolvedTable {
            table: "table2".to_string(),
        };

        assert_eq!(table1, table2);
        assert_ne!(table1, table3);
    }
}