Merge pull request #96 from benesch/extract

benesch · web-flow · commit 11fc8334332b · 2019-06-04T00:07:23.000-04:00
Support EXTRACT function-like operator
diff --git a/src/sqlast/mod.rs b/src/sqlast/mod.rs
@@ -100,6 +100,10 @@ pub enum ASTNode {
         expr: Box<ASTNode>,
         data_type: SQLType,
     },
+    SQLExtract {
+        field: SQLDateTimeField,
+        expr: Box<ASTNode>,
+    },
     /// `expr COLLATE collation`
     SQLCollate {
         expr: Box<ASTNode>,
@@ -186,6 +190,9 @@ impl ToString for ASTNode {
                 expr.as_ref().to_string(),
                 data_type.to_string()
             ),
+            ASTNode::SQLExtract { field, expr } => {
+                format!("EXTRACT({} FROM {})", field.to_string(), expr.to_string())
+            }
             ASTNode::SQLCollate { expr, collation } => format!(
                 "{} COLLATE {}",
                 expr.as_ref().to_string(),
@@ -628,6 +635,29 @@ impl ToString for SQLFunction {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Hash)]
+pub enum SQLDateTimeField {
+    Year,
+    Month,
+    Day,
+    Hour,
+    Minute,
+    Second,
+}
+
+impl ToString for SQLDateTimeField {
+    fn to_string(&self) -> String {
+        match self {
+            SQLDateTimeField::Year => "YEAR".to_string(),
+            SQLDateTimeField::Month => "MONTH".to_string(),
+            SQLDateTimeField::Day => "DAY".to_string(),
+            SQLDateTimeField::Hour => "HOUR".to_string(),
+            SQLDateTimeField::Minute => "MINUTE".to_string(),
+            SQLDateTimeField::Second => "SECOND".to_string(),
+        }
+    }
+}
+
 /// External table's available file format
 #[derive(Debug, Clone, PartialEq, Hash)]
 pub enum FileFormat {
diff --git a/src/sqlparser.rs b/src/sqlparser.rs
@@ -193,6 +193,7 @@ impl Parser {
                 "CASE" => self.parse_case_expression(),
                 "CAST" => self.parse_cast_expression(),
                 "EXISTS" => self.parse_exists_expression(),
+                "EXTRACT" => self.parse_extract_expression(),
                 "NOT" => Ok(ASTNode::SQLUnary {
                     operator: SQLOperator::Not,
                     expr: Box::new(self.parse_subexpr(Self::UNARY_NOT_PREC)?),
@@ -417,6 +418,31 @@ impl Parser {
         Ok(exists_node)
     }
 
+    pub fn parse_extract_expression(&mut self) -> Result<ASTNode, ParserError> {
+        self.expect_token(&Token::LParen)?;
+        let tok = self.next_token();
+        let field = if let Some(Token::SQLWord(ref k)) = tok {
+            match k.keyword.as_ref() {
+                "YEAR" => SQLDateTimeField::Year,
+                "MONTH" => SQLDateTimeField::Month,
+                "DAY" => SQLDateTimeField::Day,
+                "HOUR" => SQLDateTimeField::Hour,
+                "MINUTE" => SQLDateTimeField::Minute,
+                "SECOND" => SQLDateTimeField::Second,
+                _ => self.expected("Date/time field inside of EXTRACT function", tok)?,
+            }
+        } else {
+            self.expected("Date/time field inside of EXTRACT function", tok)?
+        };
+        self.expect_keyword("FROM")?;
+        let expr = self.parse_expr()?;
+        self.expect_token(&Token::RParen)?;
+        Ok(ASTNode::SQLExtract {
+            field,
+            expr: Box::new(expr),
+        })
+    }
+
     /// Parse an operator following an expression
     pub fn parse_infix(&mut self, expr: ASTNode, precedence: u8) -> Result<ASTNode, ParserError> {
         debug!("parsing infix");
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
@@ -797,6 +797,35 @@ fn parse_cast() {
     );
 }
 
+#[test]
+fn parse_extract() {
+    let sql = "SELECT EXTRACT(YEAR FROM d)";
+    let select = verified_only_select(sql);
+    assert_eq!(
+        &ASTNode::SQLExtract {
+            field: SQLDateTimeField::Year,
+            expr: Box::new(ASTNode::SQLIdentifier("d".to_string())),
+        },
+        expr_from_projection(only(&select.projection)),
+    );
+
+    one_statement_parses_to("SELECT EXTRACT(year from d)", "SELECT EXTRACT(YEAR FROM d)");
+
+    verified_stmt("SELECT EXTRACT(MONTH FROM d)");
+    verified_stmt("SELECT EXTRACT(DAY FROM d)");
+    verified_stmt("SELECT EXTRACT(HOUR FROM d)");
+    verified_stmt("SELECT EXTRACT(MINUTE FROM d)");
+    verified_stmt("SELECT EXTRACT(SECOND FROM d)");
+
+    let res = parse_sql_statements("SELECT EXTRACT(MILLISECOND FROM d)");
+    assert_eq!(
+        ParserError::ParserError(
+            "Expected Date/time field inside of EXTRACT function, found: MILLISECOND".to_string()
+        ),
+        res.unwrap_err()
+    );
+}
+
 #[test]
 fn parse_create_table() {
     let sql = "CREATE TABLE uk_cities (\
