Add support for RETURNING clause (#110)

* Support RETURNING statement

* Simplified SQLReturning

* Updated SQLReturning serialization

* Added supportsReturning to SQLDialect

* Added SQLReturningBuilder to support returning statement with query builders.

* Updated SQLReturningBuilder doc comments

Sources/SQLKit/Builders/SQLDeleteBuilder.swift

@@ -4,7 +4,7 @@
 ///         .where(\.name != "Earth").run()
 ///
 /// See `SQLQueryBuilder` and `SQLPredicateBuilder` for more information.
-public final class SQLDeleteBuilder: SQLQueryBuilder, SQLPredicateBuilder {
+public final class SQLDeleteBuilder: SQLQueryBuilder, SQLPredicateBuilder, SQLReturningBuilder {
     /// `Delete` query being built.
     public var delete: SQLDelete
 
@@ -18,6 +18,11 @@ public final class SQLDeleteBuilder: SQLQueryBuilder, SQLPredicateBuilder {
         get { return self.delete.predicate }
         set { self.delete.predicate = newValue }
     }
+
+    public var returning: SQLReturning? {
+        get { return self.delete.returning }
+        set { self.delete.returning = newValue }
+    }
 
     /// Creates a new `SQLDeleteBuilder`.
     public init(_ delete: SQLDelete, on database: SQLDatabase) {

Sources/SQLKit/Builders/SQLInsertBuilder.swift

@@ -4,7 +4,7 @@
 ///         .value(earth).run()
 ///
 /// See `SQLQueryBuilder` for more information.
-public final class SQLInsertBuilder: SQLQueryBuilder {
+public final class SQLInsertBuilder: SQLQueryBuilder, SQLReturningBuilder {
     /// `Insert` query being built.
     public var insert: SQLInsert
 
@@ -15,6 +15,11 @@ public final class SQLInsertBuilder: SQLQueryBuilder {
     public var query: SQLExpression {
         return self.insert
     }
+
+    public var returning: SQLReturning? {
+        get { return self.insert.returning }
+        set { self.insert.returning = newValue }
+    }
 
     /// Creates a new `SQLInsertBuilder`.
     public init(_ insert: SQLInsert, on database: SQLDatabase) {

Sources/SQLKit/Builders/SQLReturningBuilder.swift

@@ -0,0 +1,47 @@
+public protocol SQLReturningBuilder: SQLQueryBuilder {
+    var returning: SQLReturning? { get set }
+}
+
+extension SQLReturningBuilder {
+    /// Specify a list of columns to be part of the result set of the query.
+    /// Each provided name is a string assumed to be a valid SQL identifier and
+    /// is not qualified.
+    ///
+    /// - parameters:
+    ///     - columns: The names of the columns to return.
+    /// - returns: Self for chaining.
+    public func returning(_ columns: String...) -> Self {
+        let sqlColumns = columns.map { (column) -> SQLColumn in
+            if column == "*" {
+                return SQLColumn(SQLLiteral.all)
+            } else {
+                return SQLColumn(column)
+            }
+        }
+
+        self.returning = .init(sqlColumns)
+        return self
+    }
+
+    /// Specify a list of columns to be returned as the result of the query.
+    /// Each input is an arbitrary expression.
+    ///
+    /// - parameters:
+    ///     - columns: A list of expressions identifying the columns to return.
+    /// - returns: Self for chaining.
+    public func returning(_ columns: SQLExpression...) -> Self {
+        self.returning = .init(columns)
+        return self
+    }
+
+    /// Specify a list of columns to be returned as the result of the query.
+    /// Each input is an arbitrary expression.
+    ///
+    /// - parameters:
+    ///     - column: An array of expressions identifying the columns to return.
+    /// - returns: Self for chaining.
+    public func returning(_ columns: [SQLExpression]) -> Self {
+        self.returning = .init(columns)
+        return self
+    }
+}

Sources/SQLKit/Builders/SQLUpdateBuilder.swift

@@ -6,7 +6,7 @@
 ///         .run()
 ///
 /// See `SQLQueryBuilder` and `SQLPredicateBuilder` for more information.
-public final class SQLUpdateBuilder: SQLQueryBuilder, SQLPredicateBuilder {
+public final class SQLUpdateBuilder: SQLQueryBuilder, SQLPredicateBuilder, SQLReturningBuilder {
     /// `Update` query being built.
     public var update: SQLUpdate
 
@@ -20,6 +20,11 @@ public final class SQLUpdateBuilder: SQLQueryBuilder, SQLPredicateBuilder {
         get { return self.update.predicate }
         set { self.update.predicate = newValue }
     }
+
+    public var returning: SQLReturning? {
+        get { return self.update.returning }
+        set { self.update.returning = newValue }
+    }
 
     /// Creates a new `SQLDeleteBuilder`.
     public init(_ update: SQLUpdate, on database: SQLDatabase) {

Sources/SQLKit/Query/SQLDelete.swift

@@ -9,18 +9,26 @@ public struct SQLDelete: SQLExpression {
     /// then only those rows for which the WHERE clause boolean expression is true are deleted. Rows for which
     /// the expression is false or NULL are retained.
     public var predicate: SQLExpression?
+
+    /// Optionally append a `RETURNING` clause that, where supported, returns the supplied supplied columns.
+    public var returning: SQLReturning?
 
     /// Creates a new `SQLDelete`.
     public init(table: SQLExpression) {
         self.table = table
     }
 
     public func serialize(to serializer: inout SQLSerializer) {
-        serializer.write("DELETE FROM ")
-        self.table.serialize(to: &serializer)
-        if let predicate = self.predicate {
-            serializer.write(" WHERE ")
-            predicate.serialize(to: &serializer)
+        serializer.statement {
+            $0.append("DELETE FROM")
+            $0.append(self.table)
+            if let predicate = self.predicate {
+                $0.append("WHERE")
+                $0.append(predicate)
+            }
+            if let returning = self.returning {
+                $0.append(returning)
+            }
         }
     }
 }

Sources/SQLKit/Query/SQLInsert.swift

@@ -12,6 +12,9 @@ public struct SQLInsert: SQLExpression {
     ///
     /// Use the `DEFAULT` literal to omit a value and that is specified as a column.
     public var values: [[SQLExpression]]
+
+    /// Optionally append a `RETURNING` clause that, where supported, returns the supplied supplied columns.
+    public var returning: SQLReturning?
 
     /// Creates a new `SQLInsert`.
     public init(table: SQLExpression) {
@@ -21,11 +24,15 @@ public struct SQLInsert: SQLExpression {
     }
 
     public func serialize(to serializer: inout SQLSerializer) {
-        serializer.write("INSERT INTO ")
-        self.table.serialize(to: &serializer)
-        serializer.write(" ")
-        SQLGroupExpression(self.columns).serialize(to: &serializer)
-        serializer.write(" VALUES ")
-        SQLList(self.values.map(SQLGroupExpression.init)).serialize(to: &serializer)
+        serializer.statement {
+            $0.append("INSERT INTO")
+            $0.append(self.table)
+            $0.append(SQLGroupExpression(self.columns))
+            $0.append("VALUES")
+            $0.append(SQLList(self.values.map(SQLGroupExpression.init)))
+            if let returning = self.returning {
+                $0.append(returning)
+            }
+        }
     }
 }

Sources/SQLKit/Query/SQLReturning.swift

@@ -0,0 +1,29 @@
+/// `RETURNING ...` statement.
+///
+public struct SQLReturning: SQLExpression {
+    public var columns: [SQLExpression]
+
+    /// Creates a new `SQLReturning`.
+    public init(_ column: SQLColumn) {
+        self.columns = [column]
+    }
+
+    /// Creates a new `SQLReturning`.
+    public init(_ columns: [SQLExpression]) {
+        self.columns = columns
+    }
+
+    public func serialize(to serializer: inout SQLSerializer) {
+        guard serializer.dialect.supportsReturning else {
+            serializer.database.logger.warning("\(serializer.dialect.name) does not support 'RETURNING' clause, skipping.")
+            return
+        }
+
+        guard !columns.isEmpty else { return }
+
+        serializer.statement {
+            $0.append("RETURNING")
+            $0.append(SQLList(columns))
+        }
+    }
+}

Sources/SQLKit/Query/SQLUpdate.swift

@@ -10,6 +10,9 @@ public struct SQLUpdate: SQLExpression {
 
     /// Optional predicate to limit updated rows.
     public var predicate: SQLExpression?
+
+    /// Optionally append a `RETURNING` clause that, where supported, returns the supplied supplied columns.
+    public var returning: SQLReturning?
 
     /// Creates a new `SQLUpdate`.
     public init(table: SQLExpression) {
@@ -19,13 +22,18 @@ public struct SQLUpdate: SQLExpression {
     }
 
     public func serialize(to serializer: inout SQLSerializer) {
-        serializer.write("UPDATE ")
-        self.table.serialize(to: &serializer)
-        serializer.write(" SET ")
-        SQLList(self.values).serialize(to: &serializer)
-        if let predicate = self.predicate {
-            serializer.write(" WHERE ")
-            predicate.serialize(to: &serializer)
+        serializer.statement {
+            $0.append("UPDATE")
+            $0.append(self.table)
+            $0.append("SET")
+            $0.append(SQLList(self.values))
+            if let predicate = self.predicate {
+                $0.append("WHERE")
+                $0.append(predicate)
+            }
+            if let returning = self.returning {
+                $0.append(returning)
+            }
         }
     }
 }

Sources/SQLKit/SQLDialect.swift

@@ -11,6 +11,7 @@ public protocol SQLDialect {
     var autoIncrementFunction: SQLExpression? { get }
     var enumSyntax: SQLEnumSyntax { get }
     var supportsDropBehavior: Bool { get }
+    var supportsReturning: Bool { get }
     var triggerSyntax: SQLTriggerSyntax { get }
     var alterTableSyntax: SQLAlterTableSyntax { get }
     func customDataType(for dataType: SQLDataType) -> SQLExpression?
@@ -135,6 +136,10 @@ extension SQLDialect {
         return false
     }
 
+    public var supportsReturning: Bool {
+        return false
+    }
+
     public var triggerSyntax: SQLTriggerSyntax {
         return SQLTriggerSyntax()
     }

Tests/SQLKitTests/SQLKitTests.swift

@@ -257,6 +257,28 @@ final class SQLKitTests: XCTestCase {
 
         XCTAssertEqual(db.results[0], "UPDATE `planets` SET `moons` = `moons` + 1 WHERE `best_at_space` >= ?")
     }
+
+    func testReturning() throws {
+        let db = TestDatabase()
+
+        try db.insert(into: "planets")
+            .columns("name")
+            .values("Jupiter")
+            .returning("id", "name")
+            .run().wait()
+        XCTAssertEqual(db.results[0], "INSERT INTO `planets` (`name`) VALUES (?) RETURNING `id`, `name`")
+
+        try db.update("planets")
+            .set("name", to: "Jupiter")
+            .returning(SQLColumn("name", table: "planets"))
+            .run().wait()
+        XCTAssertEqual(db.results[1], "UPDATE `planets` SET `name` = ? RETURNING `planets`.`name`")
+
+        try db.delete(from: "planets")
+            .returning("*")
+            .run().wait()
+        XCTAssertEqual(db.results[2], "DELETE FROM `planets` RETURNING *")
+    }
 }
 
 // MARK: Table Creation

Tests/SQLKitTests/Utilities.swift

@@ -79,6 +79,8 @@ struct GenericDialect: SQLDialect {
 
     var supportsIfExists: Bool = true
 
+    var supportsReturning: Bool = true
+
     var identifierQuote: SQLExpression {
         return SQLRaw("`")
     }