Enable map access for numbers, multiple nesting levels

src/ast/mod.rs

@@ -251,7 +251,7 @@ pub enum Expr {
     },
     MapAccess {
         column: Box<Expr>,
-        key: String,
+        keys: Vec<Value>,
     },
     /// Scalar function call e.g. `LEFT(foo, 5)`
     Function(Function),
@@ -280,7 +280,15 @@ impl fmt::Display for Expr {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Expr::Identifier(s) => write!(f, "{}", s),
-            Expr::MapAccess { column, key } => write!(f, "{}[\"{}\"]", column, key),
+            Expr::MapAccess { column, keys } => {
+                write!(f, "{}{}", column, keys.into_iter().map(|k| {
+                    match k {
+                        k @ Value::Number(_, _) => format!("[{}]", k),
+                        _ => format!("[\"{}\"]", k)
+                    }
+
+                }).collect::<Vec<String>>().join(""))
+            },
             Expr::Wildcard => f.write_str("*"),
             Expr::QualifiedWildcard(q) => write!(f, "{}.*", display_separated(q, ".")),
             Expr::CompoundIdentifier(s) => write!(f, "{}", display_separated(s, ".")),

src/parser.rs

@@ -951,13 +951,20 @@ impl<'a> Parser<'a> {
     }
 
     pub fn parse_map_access(&mut self, expr: Expr) -> Result<Expr, ParserError> {
-        let key = self.parse_literal_string()?;
+        let key = self.parse_map_key()?;
         let tok = self.consume_token(&Token::RBracket);
         debug!("Tok: {}", tok);
+        let mut key_parts: Vec<Value> = vec![key];
+        while self.consume_token(&Token::LBracket) {
+            let key = self.parse_map_key()?;
+            let tok = self.consume_token(&Token::RBracket);
+            debug!("Tok: {}", tok);
+            key_parts.push(key);
+        }
         match expr {
             e @ Expr::Identifier(_) | e @ Expr::CompoundIdentifier(_) => Ok(Expr::MapAccess {
                 column: Box::new(e),
-                key,
+                keys: key_parts,
             }),
             _ => Ok(expr),
         }
@@ -1995,6 +2002,19 @@ impl<'a> Parser<'a> {
         }
     }
 
+    /// Parse a map key string
+    pub fn parse_map_key(&mut self) -> Result<Value, ParserError> {
+        match self.next_token() {
+            Token::Word(Word { value, keyword, .. }) if keyword == Keyword::NoKeyword => Ok(Value::SingleQuotedString(value)),
+            Token::SingleQuotedString(s) => Ok(Value::SingleQuotedString(s)),
+            #[cfg(not(feature = "bigdecimal"))]
+            Token::Number(s, _) => Ok(Value::Number(s, false)),
+            #[cfg(feature = "bigdecimal")]
+            Token::Number(s, _) => Ok(Value::Number(s.parse().unwrap(), false)),
+            unexpected => self.expected("literal string or number", unexpected),
+        }
+    }
+
     /// Parse a SQL datatype (in the context of a CREATE TABLE statement for example)
     pub fn parse_data_type(&mut self) -> Result<DataType, ParserError> {
         match self.next_token() {

tests/sqlparser_common.rs

@@ -26,6 +26,7 @@ use matches::assert_matches;
 use sqlparser::ast::*;
 use sqlparser::dialect::{keywords::ALL_KEYWORDS, GenericDialect, SQLiteDialect};
 use sqlparser::parser::{Parser, ParserError};
+use sqlparser::ast::Expr::Identifier;
 
 #[test]
 fn parse_insert_values() {

tests/sqlparser_postgres.rs

@@ -21,6 +21,9 @@ use test_utils::*;
 use sqlparser::ast::*;
 use sqlparser::dialect::{GenericDialect, PostgreSqlDialect};
 use sqlparser::parser::ParserError;
+use sqlparser::ast::Expr::{Identifier, MapAccess};
+#[cfg(feature = "bigdecimal")]
+use bigdecimal::BigDecimal;
 
 #[test]
 fn parse_create_table_with_defaults() {
@@ -669,6 +672,26 @@ fn parse_pg_regex_match_ops() {
     }
 }
 
+#[test]
+fn parse_map_access_expr() {
+    //let sql = "SELECT foo[0] as foozero, bar[\"baz\"] as barbaz FROM foos";
+    let sql = "SELECT foo[0] FROM foos";
+    let select = pg_and_generic().verified_only_select(sql);
+    #[cfg(not(feature = "bigdecimal"))]
+    assert_eq!(
+        &MapAccess { column: Box::new(Identifier(Ident { value: "foo".to_string(), quote_style: None })), keys: vec![Value::Number("0".to_string(), false)] },
+        expr_from_projection(only(&select.projection)),
+    );
+    let sql = "SELECT foo[0][0] FROM foos";
+    let select = pg_and_generic().verified_only_select(sql);
+    #[cfg(not(feature = "bigdecimal"))]
+    assert_eq!(
+        &MapAccess { column: Box::new(Identifier(Ident { value: "foo".to_string(), quote_style: None })), keys: vec![Value::Number("0".to_string(), false), Value::Number("0".to_string(), false)] },
+        expr_from_projection(only(&select.projection)),
+    );
+}
+
+
 fn pg() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(PostgreSqlDialect {})],