Feature/issues 156 extended returning support (#317)

* Change returning clause to accept a list of columns instead of a full select expression * extended returning support * remove feature guard * issues-156 cargo fmt * issues-156 Fix * issues-156 End with Returning expression 1) Fix tests 2) Fix docs 3) Fix query_builder.rs * issues-156 Add helper function returning_all * Tidy examples * Tidy examples * Tidy examples Co-authored-by: Marlon Sousa <marlon.bsousa@gmail.com> Co-authored-by: Billy Chan <ccw.billy.123@gmail.com> Co-authored-by: Chris Tsang <chris.2y3@outlook.com>
SeaQL · May 11, 2022 · 88d2c87 · 88d2c87
1 parent 47ce784
commit 88d2c87
Show file tree

Hide file tree

Showing 12 changed files with 413 additions and 104 deletions.
diff --git a/src/backend/mysql/query.rs b/src/backend/mysql/query.rs
@@ -3,7 +3,7 @@ use super::*;
 impl QueryBuilder for MysqlQueryBuilder {
     fn prepare_returning(
         &self,
-        _returning: &[SelectExpr],
+        _returning: &Option<ReturningClause>,
         _sql: &mut SqlWriter,
         _collector: &mut dyn FnMut(Value),
     ) {

diff --git a/src/backend/mysql/table.rs b/src/backend/mysql/table.rs
@@ -162,7 +162,7 @@ impl TableBuilder for MysqlQueryBuilder {
                     let mut foreign_key = TableForeignKey::new();
                     foreign_key.name(&name.to_string());
                     let drop = ForeignKeyDropStatement {
-                        foreign_key: foreign_key,
+                        foreign_key,
                         table: None,
                     };
                     self.prepare_foreign_key_drop_statement_internal(&drop, sql, Mode::TableAlter);

diff --git a/src/backend/postgres/table.rs b/src/backend/postgres/table.rs
@@ -173,7 +173,7 @@ impl TableBuilder for PostgresQueryBuilder {
                     let mut foreign_key = TableForeignKey::new();
                     foreign_key.name(&name.to_string());
                     let drop = ForeignKeyDropStatement {
-                        foreign_key: foreign_key,
+                        foreign_key,
                         table: None,
                     };
                     self.prepare_foreign_key_drop_statement_internal(&drop, sql, Mode::TableAlter);

diff --git a/src/backend/query_builder.rs b/src/backend/query_builder.rs
@@ -268,28 +268,7 @@ pub trait QueryBuilder: QuotedBuilder {
     ) {
         match simple_expr {
             SimpleExpr::Column(column_ref) => {
-                match column_ref {
-                    ColumnRef::Column(column) => column.prepare(sql, self.quote()),
-                    ColumnRef::TableColumn(table, column) => {
-                        table.prepare(sql, self.quote());
-                        write!(sql, ".").unwrap();
-                        column.prepare(sql, self.quote());
-                    }
-                    ColumnRef::SchemaTableColumn(schema, table, column) => {
-                        schema.prepare(sql, self.quote());
-                        write!(sql, ".").unwrap();
-                        table.prepare(sql, self.quote());
-                        write!(sql, ".").unwrap();
-                        column.prepare(sql, self.quote());
-                    }
-                    ColumnRef::Asterisk => {
-                        write!(sql, "*").unwrap();
-                    }
-                    ColumnRef::TableAsterisk(table) => {
-                        table.prepare(sql, self.quote());
-                        write!(sql, ".*").unwrap();
-                    }
-                };
+                self.prepare_column_ref(column_ref, sql);
             }
             SimpleExpr::Tuple(exprs) => {
                 self.prepare_tuple(exprs, sql, collector);
@@ -584,6 +563,31 @@ pub trait QueryBuilder: QuotedBuilder {
         }
     }
 
+    fn prepare_column_ref(&self, column_ref: &ColumnRef, sql: &mut SqlWriter) {
+        match column_ref {
+            ColumnRef::Column(column) => column.prepare(sql, self.quote()),
+            ColumnRef::TableColumn(table, column) => {
+                table.prepare(sql, self.quote());
+                write!(sql, ".").unwrap();
+                column.prepare(sql, self.quote());
+            }
+            ColumnRef::SchemaTableColumn(schema, table, column) => {
+                schema.prepare(sql, self.quote());
+                write!(sql, ".").unwrap();
+                table.prepare(sql, self.quote());
+                write!(sql, ".").unwrap();
+                column.prepare(sql, self.quote());
+            }
+            ColumnRef::Asterisk => {
+                write!(sql, "*").unwrap();
+            }
+            ColumnRef::TableAsterisk(table) => {
+                table.prepare(sql, self.quote());
+                write!(sql, ".*").unwrap();
+            }
+        };
+    }
+
     /// Translate [`UnOper`] into SQL statement.
     fn prepare_un_oper(
         &self,
@@ -1276,19 +1280,24 @@ pub trait QueryBuilder: QuotedBuilder {
     /// Hook to insert "RETURNING" statements.
     fn prepare_returning(
         &self,
-        returning: &[SelectExpr],
+        returning: &Option<ReturningClause>,
         sql: &mut SqlWriter,
-        collector: &mut dyn FnMut(Value),
+        _collector: &mut dyn FnMut(Value),
     ) {
-        if !returning.is_empty() {
-            write!(sql, " RETURNING ").unwrap();
-            returning.iter().fold(true, |first, expr| {
-                if !first {
-                    write!(sql, ", ").unwrap()
+        if let Some(returning) = returning {
+            match &returning {
+                ReturningClause::All => write!(sql, " RETURNING *").unwrap(),
+                ReturningClause::Columns(cols) => {
+                    write!(sql, " RETURNING ").unwrap();
+                    cols.iter().fold(true, |first, column_ref| {
+                        if !first {
+                            write!(sql, ", ").unwrap()
+                        }
+                        self.prepare_column_ref(column_ref, sql);
+                        false
+                    });
                 }
-                self.prepare_select_expr(expr, sql, collector);
-                false
-            });
+            }
         }
     }
 

diff --git a/src/query/delete.rs b/src/query/delete.rs
@@ -4,8 +4,8 @@ use crate::{
     query::{condition::*, OrderedStatement},
     types::*,
     value::*,
-    Query, QueryStatementBuilder, QueryStatementWriter, SelectExpr, SelectStatement,
-    SubQueryStatement, WithClause, WithQuery,
+    QueryStatementBuilder, QueryStatementWriter, ReturningClause, SubQueryStatement, WithClause,
+    WithQuery,
 };
 
 /// Delete existing rows from the table
@@ -40,7 +40,7 @@ pub struct DeleteStatement {
     pub(crate) wherei: ConditionHolder,
     pub(crate) orders: Vec<OrderExpr>,
     pub(crate) limit: Option<Value>,
-    pub(crate) returning: Vec<SelectExpr>,
+    pub(crate) returning: Option<ReturningClause>,
 }
 
 impl Default for DeleteStatement {
@@ -57,7 +57,7 @@ impl DeleteStatement {
             wherei: ConditionHolder::new(),
             orders: Vec::new(),
             limit: None,
-            returning: Vec::new(),
+            returning: None,
         }
     }
 
@@ -103,20 +103,15 @@ impl DeleteStatement {
 
     /// RETURNING expressions.
     ///
-    /// ## Note:
-    /// Works on
-    /// * PostgreSQL
-    /// * SQLite
-    ///     - SQLite version >= 3.35.0
-    ///     - **Note that sea-query won't try to enforce either of these constraints**
+    /// # Examples
     ///
     /// ```
     /// use sea_query::{tests_cfg::*, *};
     ///
     /// let query = Query::delete()
     ///     .from_table(Glyph::Table)
     ///     .and_where(Expr::col(Glyph::Id).eq(1))
-    ///     .returning(Query::select().column(Glyph::Id).take())
+    ///     .returning(Query::returning().columns([Glyph::Id]))
     ///     .to_owned();
     ///
     /// assert_eq!(
@@ -132,20 +127,14 @@ impl DeleteStatement {
     ///     r#"DELETE FROM "glyph" WHERE "id" = 1 RETURNING "id""#
     /// );
     /// ```
-    pub fn returning(&mut self, select: SelectStatement) -> &mut Self {
-        self.returning = select.selects;
+    pub fn returning(&mut self, returning_cols: ReturningClause) -> &mut Self {
+        self.returning = Some(returning_cols);
         self
     }
 
-    /// RETURNING a column after delete.
-    /// Wrapper over [`DeleteStatement::returning()`].
+    /// RETURNING expressions for a column.
     ///
-    /// ## Note:
-    /// Works on
-    /// * PostgreSQL
-    /// * SQLite
-    ///     - SQLite version >= 3.35.0
-    ///     - **Note that sea-query won't try to enforce either of these constraints**
+    /// # Examples
     ///
     /// ```
     /// use sea_query::{tests_cfg::*, *};
@@ -171,9 +160,40 @@ impl DeleteStatement {
     /// ```
     pub fn returning_col<C>(&mut self, col: C) -> &mut Self
     where
-        C: IntoIden,
+        C: IntoColumnRef,
     {
-        self.returning(Query::select().column(col.into_iden()).take())
+        self.returning(ReturningClause::Columns(vec![col.into_column_ref()]))
+    }
+
+    /// RETURNING expressions all columns.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{tests_cfg::*, *};
+    ///
+    /// let query = Query::insert()
+    ///     .into_table(Glyph::Table)
+    ///     .columns(vec![Glyph::Image])
+    ///     .values_panic(vec!["12A".into()])
+    ///     .returning_all()
+    ///     .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     "INSERT INTO `glyph` (`image`) VALUES ('12A')"
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"INSERT INTO "glyph" ("image") VALUES ('12A') RETURNING *"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"INSERT INTO "glyph" ("image") VALUES ('12A') RETURNING *"#
+    /// );
+    /// ```
+    pub fn returning_all(&mut self) -> &mut Self {
+        self.returning(ReturningClause::All)
     }
 
     /// Create a [WithQuery] by specifying a [WithClause] to execute this query with.

diff --git a/src/query/insert.rs b/src/query/insert.rs
@@ -1,6 +1,6 @@
 use crate::{
-    backend::QueryBuilder, error::*, prepare::*, types::*, value::*, Expr, OnConflict, Query,
-    QueryStatementBuilder, QueryStatementWriter, SelectExpr, SelectStatement, SimpleExpr,
+    backend::QueryBuilder, error::*, prepare::*, types::*, value::*, Expr, OnConflict,
+    QueryStatementBuilder, QueryStatementWriter, ReturningClause, SelectStatement, SimpleExpr,
     SubQueryStatement, WithClause, WithQuery,
 };
 
@@ -48,7 +48,7 @@ pub struct InsertStatement {
     pub(crate) columns: Vec<DynIden>,
     pub(crate) source: Option<InsertValueSource>,
     pub(crate) on_conflict: Option<OnConflict>,
-    pub(crate) returning: Vec<SelectExpr>,
+    pub(crate) returning: Option<ReturningClause>,
     pub(crate) default_values: Option<u32>,
 }
 
@@ -308,12 +308,7 @@ impl InsertStatement {
 
     /// RETURNING expressions.
     ///
-    /// ## Note:
-    /// Works on
-    /// * PostgreSQL
-    /// * SQLite
-    ///     - SQLite version >= 3.35.0
-    ///     - **Note that sea-query won't try to enforce either of these constraints**
+    /// # Examples
     ///
     /// ```
     /// use sea_query::{tests_cfg::*, *};
@@ -322,7 +317,7 @@ impl InsertStatement {
     ///     .into_table(Glyph::Table)
     ///     .columns(vec![Glyph::Image])
     ///     .values_panic(vec!["12A".into()])
-    ///     .returning(Query::select().column(Glyph::Id).take())
+    ///     .returning(Query::returning().columns([Glyph::Id]))
     ///     .to_owned();
     ///
     /// assert_eq!(
@@ -338,20 +333,14 @@ impl InsertStatement {
     ///     r#"INSERT INTO "glyph" ("image") VALUES ('12A') RETURNING "id""#
     /// );
     /// ```
-    pub fn returning(&mut self, select: SelectStatement) -> &mut Self {
-        self.returning = select.selects;
+    pub fn returning(&mut self, returning: ReturningClause) -> &mut Self {
+        self.returning = Some(returning);
         self
     }
 
-    /// RETURNING a column after insertion. This is equivalent to MySQL's LAST_INSERT_ID.
-    /// Wrapper over [`InsertStatement::returning()`].
+    /// RETURNING expressions for a column.
     ///
-    /// ## Note:
-    /// Works on
-    /// * PostgreSQL
-    /// * SQLite
-    ///     - SQLite version >= 3.35.0
-    ///     - **Note that sea-query won't try to enforce either of these constraints**
+    /// # Examples
     ///
     /// ```
     /// use sea_query::{tests_cfg::*, *};
@@ -378,9 +367,40 @@ impl InsertStatement {
     /// ```
     pub fn returning_col<C>(&mut self, col: C) -> &mut Self
     where
-        C: IntoIden,
+        C: IntoColumnRef,
     {
-        self.returning(Query::select().column(col.into_iden()).take())
+        self.returning(ReturningClause::Columns(vec![col.into_column_ref()]))
+    }
+
+    /// RETURNING expressions all columns.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use sea_query::{tests_cfg::*, *};
+    ///
+    /// let query = Query::insert()
+    ///     .into_table(Glyph::Table)
+    ///     .columns(vec![Glyph::Image])
+    ///     .values_panic(vec!["12A".into()])
+    ///     .returning_all()
+    ///     .to_owned();
+    ///
+    /// assert_eq!(
+    ///     query.to_string(MysqlQueryBuilder),
+    ///     "INSERT INTO `glyph` (`image`) VALUES ('12A')"
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(PostgresQueryBuilder),
+    ///     r#"INSERT INTO "glyph" ("image") VALUES ('12A') RETURNING *"#
+    /// );
+    /// assert_eq!(
+    ///     query.to_string(SqliteQueryBuilder),
+    ///     r#"INSERT INTO "glyph" ("image") VALUES ('12A') RETURNING *"#
+    /// );
+    /// ```
+    pub fn returning_all(&mut self) -> &mut Self {
+        self.returning(ReturningClause::All)
     }
 
     /// Create a [WithQuery] by specifying a [WithClause] to execute this query with.

diff --git a/src/query/mod.rs b/src/query/mod.rs
@@ -13,6 +13,7 @@ mod delete;
 mod insert;
 mod on_conflict;
 mod ordered;
+mod returning;
 mod select;
 mod shim;
 mod traits;
@@ -26,6 +27,7 @@ pub use delete::*;
 pub use insert::*;
 pub use on_conflict::*;
 pub use ordered::*;
+pub use returning::*;
 pub use select::*;
 pub use traits::*;
 pub use update::*;
@@ -79,4 +81,9 @@ impl Query {
     pub fn with() -> WithClause {
         WithClause::new()
     }
+
+    /// Construct [`Returning`]
+    pub fn returning() -> Returning {
+        Returning::new()
+    }
 }