hive: add create function syntax (apache#496)

mobuchowski · web-flow · commit cc2559c09772 · 2022-05-26T17:30:10.000-04:00
Signed-off-by: Maciej Obuchowski &lt;obuchowski.maciej@gmail.com&gt;
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -981,6 +981,15 @@ pub enum Statement {
         location: Option<String>,
         managed_location: Option<String>,
     },
+    /// CREATE FUNCTION
+    ///
+    /// Hive: https://cwiki.apache.org/confluence/display/hive/languagemanual+ddl#LanguageManualDDL-Create/Drop/ReloadFunction
+    CreateFunction {
+        temporary: bool,
+        name: ObjectName,
+        class_name: String,
+        using: Option<CreateFunctionUsing>,
+    },
     /// `ASSERT <condition> [AS <message>]`
     Assert {
         condition: Expr,
@@ -1320,6 +1329,22 @@ impl fmt::Display for Statement {
                 }
                 Ok(())
             }
+            Statement::CreateFunction {
+                temporary,
+                name,
+                class_name,
+                using,
+            } => {
+                write!(
+                    f,
+                    "CREATE {temp}FUNCTION {name} AS '{class_name}'",
+                    temp = if *temporary { "TEMPORARY " } else { "" },
+                )?;
+                if let Some(u) = using {
+                    write!(f, " {}", u)?;
+                }
+                Ok(())
+            }
             Statement::CreateView {
                 name,
                 or_replace,
@@ -2568,6 +2593,25 @@ impl fmt::Display for DiscardObject {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub enum CreateFunctionUsing {
+    Jar(String),
+    File(String),
+    Archive(String),
+}
+
+impl fmt::Display for CreateFunctionUsing {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "USING ")?;
+        match self {
+            CreateFunctionUsing::Jar(uri) => write!(f, "JAR '{uri}'"),
+            CreateFunctionUsing::File(uri) => write!(f, "FILE '{uri}'"),
+            CreateFunctionUsing::Archive(uri) => write!(f, "ARCHIVE '{uri}'"),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/keywords.rs b/src/keywords.rs
@@ -76,6 +76,7 @@ define_keywords!(
     AND,
     ANY,
     APPLY,
+    ARCHIVE,
     ARE,
     ARRAY,
     ARRAY_AGG,
@@ -223,6 +224,7 @@ define_keywords!(
     FALSE,
     FETCH,
     FIELDS,
+    FILE,
     FILTER,
     FIRST,
     FIRST_VALUE,
@@ -277,6 +279,7 @@ define_keywords!(
     ISODOW,
     ISOLATION,
     ISOYEAR,
+    JAR,
     JOIN,
     JSONFILE,
     JULIAN,
diff --git a/src/parser.rs b/src/parser.rs
@@ -1615,6 +1615,8 @@ impl<'a> Parser<'a> {
             self.parse_create_schema()
         } else if self.parse_keyword(Keyword::DATABASE) {
             self.parse_create_database()
+        } else if dialect_of!(self is HiveDialect) && self.parse_keyword(Keyword::FUNCTION) {
+            self.parse_create_function(temporary)
         } else {
             self.expected("an object type after CREATE", self.peek_token())
         }
@@ -1671,6 +1673,42 @@ impl<'a> Parser<'a> {
         })
     }
 
+    pub fn parse_optional_create_function_using(
+        &mut self,
+    ) -> Result<Option<CreateFunctionUsing>, ParserError> {
+        if !self.parse_keyword(Keyword::USING) {
+            return Ok(None);
+        };
+        let keyword =
+            self.expect_one_of_keywords(&[Keyword::JAR, Keyword::FILE, Keyword::ARCHIVE])?;
+
+        let uri = self.parse_literal_string()?;
+
+        match keyword {
+            Keyword::JAR => Ok(Some(CreateFunctionUsing::Jar(uri))),
+            Keyword::FILE => Ok(Some(CreateFunctionUsing::File(uri))),
+            Keyword::ARCHIVE => Ok(Some(CreateFunctionUsing::Archive(uri))),
+            _ => self.expected(
+                "JAR, FILE or ARCHIVE, got {:?}",
+                Token::make_keyword(format!("{:?}", keyword).as_str()),
+            ),
+        }
+    }
+
+    pub fn parse_create_function(&mut self, temporary: bool) -> Result<Statement, ParserError> {
+        let name = self.parse_object_name()?;
+        self.expect_keyword(Keyword::AS)?;
+        let class_name = self.parse_literal_string()?;
+        let using = self.parse_optional_create_function_using()?;
+
+        Ok(Statement::CreateFunction {
+            temporary,
+            name,
+            class_name,
+            using,
+        })
+    }
+
     pub fn parse_create_external_table(
         &mut self,
         or_replace: bool,
diff --git a/tests/sqlparser_hive.rs b/tests/sqlparser_hive.rs
@@ -15,8 +15,8 @@
 //! Test SQL syntax specific to Hive. The parser based on the generic dialect
 //! is also tested (on the inputs it can handle).
 
-use sqlparser::ast::{Ident, ObjectName, SetVariableValue, Statement};
-use sqlparser::dialect::HiveDialect;
+use sqlparser::ast::{CreateFunctionUsing, Ident, ObjectName, SetVariableValue, Statement};
+use sqlparser::dialect::{GenericDialect, HiveDialect};
 use sqlparser::parser::ParserError;
 use sqlparser::test_utils::*;
 
@@ -232,6 +232,47 @@ fn set_statement_with_minus() {
     )
 }
 
+#[test]
+fn parse_create_function() {
+    let sql = "CREATE TEMPORARY FUNCTION mydb.myfunc AS 'org.random.class.Name' USING JAR 'hdfs://somewhere.com:8020/very/far'";
+    match hive().verified_stmt(sql) {
+        Statement::CreateFunction {
+            temporary,
+            name,
+            class_name,
+            using,
+        } => {
+            assert!(temporary);
+            assert_eq!("mydb.myfunc", name.to_string());
+            assert_eq!("org.random.class.Name", class_name);
+            assert_eq!(
+                using,
+                Some(CreateFunctionUsing::Jar(
+                    "hdfs://somewhere.com:8020/very/far".to_string()
+                ))
+            )
+        }
+        _ => unreachable!(),
+    }
+
+    let generic = TestedDialects {
+        dialects: vec![Box::new(GenericDialect {})],
+    };
+
+    assert_eq!(
+        generic.parse_sql_statements(sql).unwrap_err(),
+        ParserError::ParserError(
+            "Expected an object type after CREATE, found: FUNCTION".to_string()
+        )
+    );
+
+    let sql = "CREATE TEMPORARY FUNCTION mydb.myfunc AS 'org.random.class.Name' USING JAR";
+    assert_eq!(
+        hive().parse_sql_statements(sql).unwrap_err(),
+        ParserError::ParserError("Expected literal string, found: EOF".to_string()),
+    );
+}
+
 fn hive() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(HiveDialect {})],
