feat(s3select): expand typed JSON source paths (#6864)

crates/s3select-query/src/sql/dialect.rs

@@ -184,6 +184,13 @@ mod tests {
         assert!(dialect.supports_group_by_expr(), "RustFsDialect should support GROUP BY expressions");
     }
 
+    #[test]
+    fn test_supports_partiql_paths() {
+        let dialect = RustFsDialect;
+
+        assert!(dialect.supports_partiql(), "RustFsDialect should support JSON source paths");
+    }
+
     #[test]
     fn test_identifier_validation_comprehensive() {
         let dialect = RustFsDialect;

crates/s3select-query/src/sql/parser.rs

@@ -16,6 +16,7 @@ use std::{collections::VecDeque, fmt::Display};
 
 use datafusion::sql::sqlparser::{
     dialect::Dialect,
+    keywords::{Keyword, RESERVED_FOR_TABLE_ALIAS},
     parser::{Parser, ParserError},
     tokenizer::{Token, Tokenizer},
 };
@@ -53,7 +54,8 @@ impl<'a> ExtParser<'a> {
     /// Parse the specified tokens with dialect
     fn new_with_dialect(sql: &str, dialect: &'a dyn Dialect) -> Result<Self> {
         let mut tokenizer = Tokenizer::new(dialect, sql);
-        let tokens = tokenizer.tokenize()?;
+        let mut tokens = tokenizer.tokenize()?;
+        rewrite_source_object_wildcards(&mut tokens);
         Ok(ExtParser {
             parser: Parser::new(dialect).with_tokens(tokens),
         })
@@ -104,6 +106,41 @@ impl<'a> ExtParser<'a> {
     }
 }
 
+fn rewrite_source_object_wildcards(tokens: &mut [Token]) {
+    let mut paren_depth = 0usize;
+    let mut in_from = false;
+    let mut index = 0usize;
+
+    while index < tokens.len() {
+        match &tokens[index] {
+            Token::Word(word) if paren_depth == 0 && word.keyword == Keyword::FROM => {
+                in_from = true;
+            }
+            Token::Word(word) if in_from && paren_depth == 0 && ends_from_source(word.keyword) => {
+                in_from = false;
+            }
+            Token::SemiColon if paren_depth == 0 => in_from = false,
+            Token::LParen => paren_depth = paren_depth.saturating_add(1),
+            Token::RParen => paren_depth = paren_depth.saturating_sub(1),
+            Token::Period if in_from => {
+                if let Some(next) = tokens[index + 1..]
+                    .iter_mut()
+                    .find(|token| !matches!(token, Token::Whitespace(_)))
+                    && matches!(next, Token::Mul)
+                {
+                    *next = Token::make_word("*", None);
+                }
+            }
+            _ => {}
+        }
+        index += 1;
+    }
+}
+
+fn ends_from_source(keyword: Keyword) -> bool {
+    RESERVED_FOR_TABLE_ALIAS.contains(&keyword) || matches!(keyword, Keyword::PREWHERE | Keyword::SETTINGS | Keyword::FORMAT)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -172,6 +209,23 @@ mod tests {
         }
     }
 
+    #[test]
+    fn parses_source_object_wildcard_without_rewriting_projection_wildcard() {
+        let mut statements = ExtParser::parse_sql("SELECT e.* FROM S3Object[*].* AS e").expect("query should parse");
+        let ExtStatement::SqlStatement(statement) = statements.pop_front().expect("one statement");
+
+        assert_eq!(statement.to_string(), "SELECT e.* FROM S3Object[*].* AS e");
+    }
+
+    #[test]
+    fn from_tokens_in_literals_and_comments_do_not_change_wildcard_scope() {
+        let sql = "SELECT 'FROM x.*' AS marker /* FROM y.* */ FROM S3Object.*";
+        let mut statements = ExtParser::parse_sql(sql).expect("query should parse");
+        let ExtStatement::SqlStatement(statement) = statements.pop_front().expect("one statement");
+
+        assert_eq!(statement.to_string(), "SELECT 'FROM x.*' AS marker FROM S3Object.*");
+    }
+
     #[test]
     fn test_default_parser_multiple_statements() {
         let parser = DefaultParser::default();

crates/s3select-query/src/sql/planner.rs

@@ -12,20 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::ops::ControlFlow;
+use std::{convert::Infallible, ops::ControlFlow};
 
 use async_recursion::async_recursion;
 use async_trait::async_trait;
 use datafusion::sql::{
-    planner::SqlToRel,
+    planner::{IdentNormalizer, SqlToRel},
     sqlparser::ast::{
-        GroupByExpr, ObjectNamePart, OrderByKind, Query, Select, SelectFlavor, SetExpr, Statement, TableFactor, Visit, Visitor,
+        AccessExpr, Expr, GroupByExpr, Ident, JsonPath, JsonPathElem, ObjectNamePart, OrderByKind, Query, Select, SelectFlavor,
+        SetExpr, Statement, Subscript, TableAlias, TableFactor, Value, Visit, VisitMut, Visitor, VisitorMut,
     },
 };
 use rustfs_s3select_api::{
     QueryError, QueryResult, SelectError,
     query::{
-        ast::ExtStatement,
+        ast::{ExtStatement, JsonPathSegment, JsonSource},
         logical_planner::{LogicalPlanner, Plan, QueryPlan},
         session::SessionCtx,
     },
@@ -64,21 +65,24 @@ impl<'a, S: ContextProviderExtension + Send + Sync + 'a> SqlPlanner<'a, S> {
         }
     }
 
-    async fn df_sql_to_plan(&self, stmt: Statement, _session: &SessionCtx) -> QueryResult<Plan> {
-        match stmt {
-            Statement::Query(_) => {
-                validate_s3_select_statement(&stmt)?;
-                let df_plan = self.df_planner.sql_statement_to_plan(stmt).map_err(classify_planner_error)?;
-                let plan = Plan::Query(QueryPlan {
-                    df_plan,
-                    is_tag_scan: false,
-                });
-
-                Ok(plan)
-            }
-            _ => Err(unsupported_structure("only SELECT queries are supported")),
+    pub(crate) async fn prepared_statement_to_plan(&self, statement: ExtStatement, session: &SessionCtx) -> QueryResult<Plan> {
+        match statement {
+            ExtStatement::SqlStatement(stmt) => self.df_prepared_sql_to_plan(*stmt, session).await,
         }
     }
+
+    async fn df_sql_to_plan(&self, mut stmt: Statement, session: &SessionCtx) -> QueryResult<Plan> {
+        prepare_s3_select_statement(&mut stmt)?;
+        self.df_prepared_sql_to_plan(stmt, session).await
+    }
+
+    async fn df_prepared_sql_to_plan(&self, stmt: Statement, _session: &SessionCtx) -> QueryResult<Plan> {
+        let df_plan = self.df_planner.sql_statement_to_plan(stmt).map_err(classify_planner_error)?;
+        Ok(Plan::Query(QueryPlan {
+            df_plan,
+            is_tag_scan: false,
+        }))
+    }
 }
 
 fn classify_planner_error(error: datafusion::common::DataFusionError) -> QueryError {
@@ -95,11 +99,10 @@ fn classify_planner_error(error: datafusion::common::DataFusionError) -> QueryEr
     error.into()
 }
 
-fn validate_s3_select_statement(statement: &Statement) -> QueryResult<()> {
+pub(crate) fn prepare_s3_select_statement(statement: &mut Statement) -> QueryResult<JsonSource> {
     let Statement::Query(query) = statement else {
         return Err(unsupported_structure("only SELECT queries are supported"));
     };
-
     if query.with.is_some()
         || query.order_by.as_ref().is_some_and(|order_by| {
             order_by.interpolate.is_some()
@@ -137,17 +140,81 @@ fn validate_s3_select_statement(statement: &Statement) -> QueryResult<()> {
     }
 
     let mut detector = SubqueryDetector { visited_root: false };
-    if query.visit(&mut detector).is_break() {
+    if Visit::visit(&*query, &mut detector).is_break() {
         return Err(unsupported_structure("subqueries are not supported"));
     }
 
-    let SetExpr::Select(select) = query.body.as_ref() else {
-        return Err(unsupported_structure("set operations and nested queries are not supported"));
+    let source = {
+        let SetExpr::Select(select) = query.body.as_mut() else {
+            return Err(unsupported_structure("set operations and nested queries are not supported"));
+        };
+        prepare_select(select)?
     };
-    validate_select(select)
+    let mut normalizer = PartiQlSubscriptNormalizer;
+    let _ = VisitMut::visit(query, &mut normalizer);
+    Ok(source)
 }
 
-fn validate_select(select: &Select) -> QueryResult<()> {
+struct PartiQlSubscriptNormalizer;
+
+impl VisitorMut for PartiQlSubscriptNormalizer {
+    type Break = Infallible;
+
+    fn post_visit_expr(&mut self, expr: &mut Expr) -> ControlFlow<Self::Break> {
+        let Expr::JsonAccess { value, path } = expr else {
+            return ControlFlow::Continue(());
+        };
+        if !matches!(path.path.first(), Some(JsonPathElem::Bracket { .. }))
+            || path
+                .path
+                .iter()
+                .any(|element| matches!(element, JsonPathElem::ColonBracket { .. }))
+        {
+            return ControlFlow::Continue(());
+        }
+
+        let mut appended_access = Vec::with_capacity(path.path.len());
+        for element in std::mem::take(&mut path.path) {
+            match element {
+                JsonPathElem::Dot { key, quoted } => {
+                    let identifier = if quoted {
+                        Ident::with_quote('"', key)
+                    } else {
+                        Ident::new(key)
+                    };
+                    appended_access.push(AccessExpr::Dot(Expr::Identifier(identifier)));
+                }
+                JsonPathElem::Bracket { key } => {
+                    appended_access.push(AccessExpr::Subscript(Subscript::Index { index: key }));
+                }
+                JsonPathElem::ColonBracket { key } => {
+                    appended_access.push(AccessExpr::Subscript(Subscript::Index { index: key }));
+                }
+            }
+        }
+
+        let value = std::mem::replace(value, Box::new(Expr::Identifier(Ident::new(""))));
+        let (root, mut access_chain) = match *value {
+            Expr::CompoundFieldAccess { root, access_chain } => (root, access_chain),
+            root => (Box::new(root), Vec::new()),
+        };
+        access_chain.extend(appended_access);
+        *expr = Expr::CompoundFieldAccess { root, access_chain };
+        ControlFlow::Continue(())
+    }
+}
+
+fn implicit_source_alias(source_path: &[JsonPathSegment]) -> Ident {
+    match source_path.last() {
+        Some(JsonPathSegment::Key { name, quoted: true }) => Ident::with_quote('"', name),
+        Some(JsonPathSegment::Key { name, quoted: false }) => Ident::new(name),
+        Some(JsonPathSegment::Index(_) | JsonPathSegment::ArrayWildcard | JsonPathSegment::ObjectWildcard) | None => {
+            Ident::new("_1")
+        }
+    }
+}
+
+fn prepare_select(select: &mut Select) -> QueryResult<JsonSource> {
     if !select.optimizer_hints.is_empty()
         || select.distinct.is_some()
         || select.select_modifiers.is_some()
@@ -170,7 +237,7 @@ fn validate_select(select: &Select) -> QueryResult<()> {
         return Err(unsupported_structure("the SELECT contains an unsupported clause"));
     }
 
-    let [table] = select.from.as_slice() else {
+    let [table] = select.from.as_mut_slice() else {
         return Err(unsupported_structure("exactly one S3Object source is required"));
     };
     if !table.joins.is_empty() {
@@ -186,8 +253,8 @@ fn validate_select(select: &Select) -> QueryResult<()> {
         partitions,
         sample,
         index_hints,
-        ..
-    } = &table.relation
+        json_path,
+    } = &mut table.relation
     else {
         return Err(unsupported_structure("subqueries and table functions are not supported"));
     };
@@ -202,9 +269,7 @@ fn validate_select(select: &Select) -> QueryResult<()> {
     {
         return Err(unsupported_structure("the S3Object source contains unsupported modifiers"));
     }
-    let ([ObjectNamePart::Identifier(table_name)] | [ObjectNamePart::Identifier(table_name), ObjectNamePart::Identifier(_)]) =
-        name.0.as_slice()
-    else {
+    let Some(ObjectNamePart::Identifier(table_name)) = name.0.first() else {
         return Err(SelectError::DataSourcePathUnsupported.into());
     };
     let is_s3_object = if table_name.quote_style.is_some() {
@@ -216,6 +281,72 @@ fn validate_select(select: &Select) -> QueryResult<()> {
         return Err(SelectError::DataSourcePathUnsupported.into());
     }
 
+    let mut source_path = Vec::new();
+    for part in &name.0[1..] {
+        let ObjectNamePart::Identifier(identifier) = part else {
+            return Err(SelectError::DataSourcePathUnsupported.into());
+        };
+        if identifier.quote_style.is_none() && identifier.value == "*" {
+            source_path.push(JsonPathSegment::ObjectWildcard);
+        } else {
+            source_path.push(JsonPathSegment::Key {
+                name: identifier.value.clone(),
+                quoted: identifier.quote_style.is_some(),
+            });
+        }
+    }
+    if let Some(json_path) = json_path.as_ref() {
+        append_json_path_segments(&mut source_path, json_path)?;
+    }
+    if alias.is_none() && !source_path.is_empty() {
+        *alias = Some(TableAlias {
+            explicit: true,
+            name: implicit_source_alias(&source_path),
+            columns: Vec::new(),
+            at: None,
+        });
+    }
+    let scalar_column = alias
+        .as_ref()
+        .map(|alias| IdentNormalizer::default().normalize(alias.name.clone()))
+        .or_else(|| {
+            source_path
+                .is_empty()
+                .then(|| IdentNormalizer::default().normalize(table_name.clone()))
+        });
+    name.0.truncate(1);
+    *json_path = None;
+    Ok(JsonSource::new(source_path, scalar_column))
+}
+
+fn append_json_path_segments(source_path: &mut Vec<JsonPathSegment>, json_path: &JsonPath) -> QueryResult<()> {
+    for element in &json_path.path {
+        let segment = match element {
+            JsonPathElem::Dot { key, quoted } if key == "*" && !quoted => JsonPathSegment::ObjectWildcard,
+            JsonPathElem::Dot { key, quoted } => JsonPathSegment::Key {
+                name: key.clone(),
+                quoted: *quoted,
+            },
+            JsonPathElem::Bracket { key: Expr::Wildcard(_) } => JsonPathSegment::ArrayWildcard,
+            JsonPathElem::Bracket { key: Expr::Value(value) } => match &value.value {
+                Value::Number(number, false) => JsonPathSegment::Index(
+                    number
+                        .to_string()
+                        .parse()
+                        .map_err(|_| QueryError::from(SelectError::DataSourcePathUnsupported))?,
+                ),
+                Value::SingleQuotedString(key) => JsonPathSegment::Key {
+                    name: key.clone(),
+                    quoted: true,
+                },
+                _ => return Err(SelectError::DataSourcePathUnsupported.into()),
+            },
+            JsonPathElem::Bracket { .. } | JsonPathElem::ColonBracket { .. } => {
+                return Err(SelectError::DataSourcePathUnsupported.into());
+            }
+        };
+        source_path.push(segment);
+    }
     Ok(())
 }
 
@@ -245,10 +376,14 @@ impl Visitor for SubqueryDetector {
 
 #[cfg(test)]
 mod tests {
-    use super::validate_s3_select_statement;
+    use super::prepare_s3_select_statement;
     use crate::sql::parser::ExtParser;
-    use datafusion::sql::sqlparser::ast::Statement;
-    use rustfs_s3select_api::{SelectError, query::ast::ExtStatement};
+    use datafusion::sql::sqlparser::ast::{AccessExpr, Expr, Statement, Visit, Visitor};
+    use rustfs_s3select_api::{
+        QueryResult, SelectError,
+        query::ast::{ExtStatement, JsonPathSegment, JsonSource},
+    };
+    use std::ops::ControlFlow;
 
     fn parse_statement(sql: &str) -> Statement {
         let mut statements = ExtParser::parse_sql(sql).expect("SQL should parse");
@@ -256,6 +391,10 @@ mod tests {
         *statement
     }
 
+    fn validate_s3_select_statement(statement: &Statement) -> QueryResult<JsonSource> {
+        prepare_s3_select_statement(&mut statement.clone())
+    }
+
     #[test]
     fn accepts_s3_select_query_shape() {
         let statement = parse_statement("SELECT s.id FROM S3Object AS s WHERE s.id = '1' LIMIT 10");
@@ -270,6 +409,195 @@ mod tests {
         assert!(validate_s3_select_statement(&statement).is_ok());
     }
 
+    #[test]
+    fn prepares_nested_json_source_path_and_normalizes_table() {
+        let mut statement = parse_statement("SELECT e.name FROM S3Object[*].employees[*] AS e");
+
+        let source = prepare_s3_select_statement(&mut statement).expect("JSON source path should be supported");
+
+        assert_eq!(
+            source.path(),
+            &[
+                JsonPathSegment::ArrayWildcard,
+                JsonPathSegment::Key {
+                    name: "employees".to_string(),
+                    quoted: false,
+                },
+                JsonPathSegment::ArrayWildcard,
+            ]
+        );
+        assert_eq!(statement.to_string(), "SELECT e.name FROM S3Object AS e");
+    }
+
+    #[test]
+    fn partiql_source_support_preserves_projection_and_filter_subscripts() {
+        let mut statement = parse_statement("SELECT s.tags[1] FROM S3Object AS s WHERE s.values[0] = 1");
+
+        prepare_s3_select_statement(&mut statement).expect("array expressions should remain supported");
+        let mut counter = FieldAccessCounter::default();
+        let _ = Visit::visit(&statement, &mut counter);
+
+        assert_eq!(counter.json_accesses, 0);
+        assert_eq!(counter.subscripts, 2);
+    }
+
+    #[derive(Default)]
+    struct FieldAccessCounter {
+        json_accesses: usize,
+        subscripts: usize,
+    }
+
+    impl Visitor for FieldAccessCounter {
+        type Break = ();
+
+        fn pre_visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::Break> {
+            match expr {
+                Expr::JsonAccess { .. } => self.json_accesses += 1,
+                Expr::CompoundFieldAccess { access_chain, .. } => {
+                    self.subscripts += access_chain
+                        .iter()
+                        .filter(|access| matches!(access, AccessExpr::Subscript(_)))
+                        .count();
+                }
+                _ => {}
+            }
+            ControlFlow::Continue(())
+        }
+    }
+
+    #[test]
+    fn prepares_array_index_and_object_wildcard_paths() {
+        let mut index_statement = parse_statement("SELECT * FROM S3Object[0]");
+        let mut wildcard_statement = parse_statement("SELECT * FROM S3Object[*].*");
+
+        assert_eq!(
+            prepare_s3_select_statement(&mut index_statement)
+                .expect("array index should be supported")
+                .path(),
+            &[JsonPathSegment::Index(0)]
+        );
+        assert_eq!(
+            prepare_s3_select_statement(&mut wildcard_statement)
+                .expect("object wildcard should be supported")
+                .path(),
+            &[JsonPathSegment::ArrayWildcard, JsonPathSegment::ObjectWildcard]
+        );
+    }
+
+    #[test]
+    fn quoted_star_remains_an_object_key() {
+        let mut statement = parse_statement("SELECT * FROM S3Object.\"*\"");
+
+        let source = prepare_s3_select_statement(&mut statement).expect("quoted key should be supported");
+
+        assert_eq!(
+            source.path(),
+            &[JsonPathSegment::Key {
+                name: "*".to_string(),
+                quoted: true,
+            }]
+        );
+    }
+
+    #[test]
+    fn preserves_quoted_keys_and_adds_implicit_source_aliases() {
+        let mut key_statement = parse_statement("SELECT employee.name FROM S3Object[*].department.employee");
+        let mut wildcard_statement = parse_statement("SELECT _1.name FROM S3Object[*].employees[*]");
+
+        prepare_s3_select_statement(&mut key_statement).expect("named source path should be supported");
+        prepare_s3_select_statement(&mut wildcard_statement).expect("wildcard source path should be supported");
+
+        assert_eq!(key_statement.to_string(), "SELECT employee.name FROM S3Object AS employee");
+        assert_eq!(wildcard_statement.to_string(), "SELECT _1.name FROM S3Object AS _1");
+    }
+
+    #[test]
+    fn root_scalar_aliases_are_preserved_and_unquoted_aliases_are_normalized() {
+        let mut implicit = parse_statement("SELECT S3Object FROM S3Object");
+        let mut unquoted = parse_statement("SELECT V FROM S3Object AS V");
+        let mut quoted = parse_statement("SELECT \"V\" FROM S3Object AS \"V\"");
+
+        let implicit_source = prepare_s3_select_statement(&mut implicit).expect("implicit root alias should be supported");
+        let unquoted_source = prepare_s3_select_statement(&mut unquoted).expect("unquoted root alias should be supported");
+        let quoted_source = prepare_s3_select_statement(&mut quoted).expect("quoted root alias should be supported");
+
+        assert!(implicit_source.path().is_empty());
+        assert_eq!(implicit_source.scalar_column(), Some("s3object"));
+        assert!(unquoted_source.path().is_empty());
+        assert_eq!(unquoted_source.scalar_column(), Some("v"));
+        assert!(quoted_source.path().is_empty());
+        assert_eq!(quoted_source.scalar_column(), Some("V"));
+    }
+
+    #[test]
+    fn unquoted_terminal_scalar_alias_uses_datafusion_identifier_case() {
+        let mut statement = parse_statement("SELECT NAME FROM S3Object[*].NAME");
+
+        let source = prepare_s3_select_statement(&mut statement).expect("terminal scalar source should be supported");
+
+        assert_eq!(source.scalar_column(), Some("name"));
+        assert_eq!(statement.to_string(), "SELECT NAME FROM S3Object AS NAME");
+    }
+
+    #[test]
+    fn single_quoted_source_key_adds_a_quoted_implicit_alias() {
+        let mut statement = parse_statement("SELECT \"Employee Data\".id FROM S3Object['Employee Data']");
+
+        let source = prepare_s3_select_statement(&mut statement).expect("single-quoted source key should be supported");
+
+        assert_eq!(
+            source.path(),
+            &[JsonPathSegment::Key {
+                name: "Employee Data".to_string(),
+                quoted: true,
+            }]
+        );
+        assert_eq!(source.scalar_column(), Some("Employee Data"));
+        assert_eq!(statement.to_string(), "SELECT \"Employee Data\".id FROM S3Object AS \"Employee Data\"");
+    }
+
+    #[test]
+    fn accepts_object_wildcard_continuation() {
+        let mut statement = parse_statement("SELECT * FROM S3Object[*].groups.*.id");
+
+        assert_eq!(
+            prepare_s3_select_statement(&mut statement)
+                .expect("object wildcard continuation should be supported")
+                .path(),
+            &[
+                JsonPathSegment::ArrayWildcard,
+                JsonPathSegment::Key {
+                    name: "groups".to_string(),
+                    quoted: false,
+                },
+                JsonPathSegment::ObjectWildcard,
+                JsonPathSegment::Key {
+                    name: "id".to_string(),
+                    quoted: false,
+                },
+            ]
+        );
+    }
+
+    #[test]
+    fn rejects_non_literal_or_out_of_range_array_indexes() {
+        for sql in [
+            "SELECT * FROM S3Object[-1]",
+            "SELECT * FROM S3Object[1 + 1]",
+            "SELECT * FROM S3Object[999999999999999999999999999999999999]",
+        ] {
+            let statement = parse_statement(sql);
+            assert!(
+                matches!(
+                    validate_s3_select_statement(&statement),
+                    Err(ref error)
+                        if matches!(error.s3_select_policy_error(), Some(SelectError::DataSourcePathUnsupported))
+                ),
+                "query should reject an unsafe array index: {sql}"
+            );
+        }
+    }
+
     #[test]
     fn accepts_group_by_and_order_by() {
         let statement = parse_statement("SELECT department, COUNT(*) FROM S3Object GROUP BY department ORDER BY department");