Add support for generated virtual columns with expression (apache#1051)

Author: takluyver

src/ast/ddl.rs

@@ -599,6 +599,7 @@ pub enum ColumnOption {
         generated_as: GeneratedAs,
         sequence_options: Option<Vec<SequenceOptions>>,
         generation_expr: Option<Expr>,
+        generation_expr_mode: Option<GeneratedExpressionMode>,
     },
 }
 
@@ -639,25 +640,25 @@ impl fmt::Display for ColumnOption {
                 generated_as,
                 sequence_options,
                 generation_expr,
-            } => match generated_as {
-                GeneratedAs::Always => {
-                    write!(f, "GENERATED ALWAYS AS IDENTITY")?;
-                    if sequence_options.is_some() {
-                        let so = sequence_options.as_ref().unwrap();
-                        if !so.is_empty() {
-                            write!(f, " (")?;
-                        }
-                        for sequence_option in so {
-                            write!(f, "{sequence_option}")?;
-                        }
-                        if !so.is_empty() {
-                            write!(f, " )")?;
-                        }
-                    }
+                generation_expr_mode,
+            } => {
+                if let Some(expr) = generation_expr {
+                    let modifier = match generation_expr_mode {
+                        None => "",
+                        Some(GeneratedExpressionMode::Virtual) => " VIRTUAL",
+                        Some(GeneratedExpressionMode::Stored) => " STORED",
+                    };
+                    write!(f, "GENERATED ALWAYS AS ({expr}){modifier}")?;
                     Ok(())
-                }
-                GeneratedAs::ByDefault => {
-                    write!(f, "GENERATED BY DEFAULT AS IDENTITY")?;
+                } else {
+                    // Like Postgres - generated from sequence
+                    let when = match generated_as {
+                        GeneratedAs::Always => "ALWAYS",
+                        GeneratedAs::ByDefault => "BY DEFAULT",
+                        // ExpStored goes with an expression, handled above
+                        GeneratedAs::ExpStored => unreachable!(),
+                    };
+                    write!(f, "GENERATED {when} AS IDENTITY")?;
                     if sequence_options.is_some() {
                         let so = sequence_options.as_ref().unwrap();
                         if !so.is_empty() {
@@ -672,17 +673,13 @@ impl fmt::Display for ColumnOption {
                     }
                     Ok(())
                 }
-                GeneratedAs::ExpStored => {
-                    let expr = generation_expr.as_ref().unwrap();
-                    write!(f, "GENERATED ALWAYS AS ({expr}) STORED")
-                }
-            },
+            }
         }
     }
 }
 
 /// `GeneratedAs`s are modifiers that follow a column option in a `generated`.
-/// 'ExpStored' is PostgreSQL specific
+/// 'ExpStored' is used for a column generated from an expression and stored.
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
@@ -692,6 +689,16 @@ pub enum GeneratedAs {
     ExpStored,
 }
 
+/// `GeneratedExpressionMode`s are modifiers that follow an expression in a `generated`.
+/// No modifier is typically the same as Virtual.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum GeneratedExpressionMode {
+    Virtual,
+    Stored,
+}
+
 fn display_constraint_name(name: &'_ Option<Ident>) -> impl fmt::Display + '_ {
     struct ConstraintName<'a>(&'a Option<Ident>);
     impl<'a> fmt::Display for ConstraintName<'a> {

src/ast/mod.rs

@@ -31,8 +31,8 @@ pub use self::data_type::{
 pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue};
 pub use self::ddl::{
     AlterColumnOperation, AlterIndexOperation, AlterTableOperation, ColumnDef, ColumnOption,
-    ColumnOptionDef, GeneratedAs, IndexType, KeyOrIndexDisplay, Partition, ProcedureParam,
-    ReferentialAction, TableConstraint, UserDefinedTypeCompositeAttributeDef,
+    ColumnOptionDef, GeneratedAs, GeneratedExpressionMode, IndexType, KeyOrIndexDisplay, Partition,
+    ProcedureParam, ReferentialAction, TableConstraint, UserDefinedTypeCompositeAttributeDef,
     UserDefinedTypeRepresentation,
 };
 pub use self::operator::{BinaryOperator, UnaryOperator};

src/parser/mod.rs

@@ -4299,6 +4299,7 @@ impl<'a> Parser<'a> {
                 generated_as: GeneratedAs::Always,
                 sequence_options: Some(sequence_options),
                 generation_expr: None,
+                generation_expr_mode: None,
             }))
         } else if self.parse_keywords(&[
             Keyword::BY,
@@ -4315,16 +4316,31 @@ impl<'a> Parser<'a> {
                 generated_as: GeneratedAs::ByDefault,
                 sequence_options: Some(sequence_options),
                 generation_expr: None,
+                generation_expr_mode: None,
             }))
         } else if self.parse_keywords(&[Keyword::ALWAYS, Keyword::AS]) {
             if self.expect_token(&Token::LParen).is_ok() {
                 let expr = self.parse_expr()?;
                 self.expect_token(&Token::RParen)?;
-                let _ = self.parse_keywords(&[Keyword::STORED]);
+                let (gen_as, expr_mode) = if self.parse_keywords(&[Keyword::STORED]) {
+                    Ok((
+                        GeneratedAs::ExpStored,
+                        Some(GeneratedExpressionMode::Stored),
+                    ))
+                } else if dialect_of!(self is PostgreSqlDialect) {
+                    // Postgres' AS IDENTITY branches are above, this one needs STORED
+                    self.expected("STORED", self.peek_token())
+                } else if self.parse_keywords(&[Keyword::VIRTUAL]) {
+                    Ok((GeneratedAs::Always, Some(GeneratedExpressionMode::Virtual)))
+                } else {
+                    Ok((GeneratedAs::Always, None))
+                }?;
+
                 Ok(Some(ColumnOption::Generated {
-                    generated_as: GeneratedAs::ExpStored,
+                    generated_as: gen_as,
                     sequence_options: None,
                     generation_expr: Some(expr),
+                    generation_expr_mode: expr_mode,
                 }))
             } else {
                 Ok(None)

tests/sqlparser_mysql.rs

@@ -507,6 +507,18 @@ fn parse_create_table_comment_character_set() {
     }
 }
 
+#[test]
+fn parse_create_table_gencol() {
+    let sql_default = "CREATE TABLE t1 (a INT, b INT GENERATED ALWAYS AS (a * 2))";
+    mysql_and_generic().verified_stmt(sql_default);
+
+    let sql_virt = "CREATE TABLE t1 (a INT, b INT GENERATED ALWAYS AS (a * 2) VIRTUAL)";
+    mysql_and_generic().verified_stmt(sql_virt);
+
+    let sql_stored = "CREATE TABLE t1 (a INT, b INT GENERATED ALWAYS AS (a * 2) STORED)";
+    mysql_and_generic().verified_stmt(sql_stored);
+}
+
 #[test]
 fn parse_quote_identifiers() {
     let sql = "CREATE TABLE `PRIMARY` (`BEGIN` INT PRIMARY KEY)";

tests/sqlparser_sqlite.rs

@@ -205,6 +205,18 @@ fn parse_create_sqlite_quote() {
     }
 }
 
+#[test]
+fn parse_create_table_gencol() {
+    let sql_default = "CREATE TABLE t1 (a INT, b INT GENERATED ALWAYS AS (a * 2))";
+    sqlite_and_generic().verified_stmt(sql_default);
+
+    let sql_virt = "CREATE TABLE t1 (a INT, b INT GENERATED ALWAYS AS (a * 2) VIRTUAL)";
+    sqlite_and_generic().verified_stmt(sql_virt);
+
+    let sql_stored = "CREATE TABLE t1 (a INT, b INT GENERATED ALWAYS AS (a * 2) STORED)";
+    sqlite_and_generic().verified_stmt(sql_stored);
+}
+
 #[test]
 fn test_placeholder() {
     // In postgres, this would be the absolute value operator '@' applied to the column 'xxx'
@@ -435,7 +447,6 @@ fn sqlite_with_options(options: ParserOptions) -> TestedDialects {
 
 fn sqlite_and_generic() -> TestedDialects {
     TestedDialects {
-        // we don't have a separate SQLite dialect, so test only the generic dialect for now
         dialects: vec![Box::new(SQLiteDialect {}), Box::new(GenericDialect {})],
         options: None,
     }