Support returning in rows_*.tbl_dbi()

NAMESPACE

@@ -258,6 +258,7 @@ export(enum_pk_candidates)
 export(examine_cardinality)
 export(filter)
 export(full_join)
+export(get_returned_rows)
 export(group_by)
 export(inner_join)
 export(is_dm)

R/rows-db.R

@@ -25,10 +25,18 @@
 #'   The default is to check only if `in_place` is `TRUE` or `NULL`.
 #'
 #'   Currently these checks are no-ops and need yet to be implemented.
+#' @param returning <[`tidy-select`][tidyr_tidy_select]> Columns to return
+#'   of the inserted data. Note that also columns not in `y` but automatically
+#'   created when inserting into `x` can be returned, for example the `id`
+#'   column.
 #'
-#' @return A tbl object of the same structure as `x`.
-#'   If `in_place = TRUE`, the underlying data is updated as a side effect,
-#'   and `x` is returned, invisibly.
+#' @return Either a tbl object of the same structure as `x` or a tibble with
+#'   the columns as in `returning`.
+#'   If `in_place = TRUE`, the underlying data is updated as a side effect and
+#'   if
+#'     * `returning = NULL`, then `x` is returned, invisibly.
+#'     * else a tibble with the inserted data and the specified columns is
+#'       returned.
 #'
 #' @name rows-db
 #' @examplesIf rlang::is_installed("dbplyr")
@@ -48,7 +56,10 @@ NULL
 #' @export
 #' @rdname rows-db
 rows_insert.tbl_dbi <- function(x, y, by = NULL, ...,
-                                in_place = NULL, copy = FALSE, check = NULL) {
+                                in_place = NULL, copy = FALSE, check = NULL,
+                                returning = NULL) {
+  returning <- enquo(returning)
+
   y <- auto_copy(x, y, copy = copy)
   y_key <- db_key(y, by)
   by <- names(y_key)
@@ -63,9 +74,9 @@ rows_insert.tbl_dbi <- function(x, y, by = NULL, ...,
     }
 
     con <- dbplyr::remote_con(x)
-    sql <- sql_rows_insert(x, y)
-    dbExecute(con, sql, immediate = TRUE)
-    invisible(x)
+    sql <- sql_rows_insert(x, y, returning = returning)
+
+    get_or_execute(x, con, sql, returning)
   } else {
     # Checking mandatory by default, opt-out
     # FIXME: contrary to doc currently also checks if `in_place = FALSE`
@@ -79,7 +90,10 @@ rows_insert.tbl_dbi <- function(x, y, by = NULL, ...,
 #' @export
 #' @rdname rows-db
 rows_update.tbl_dbi <- function(x, y, by = NULL, ...,
-                                in_place = NULL, copy = FALSE, check = NULL) {
+                                in_place = NULL, copy = FALSE, check = NULL,
+                                returning = NULL) {
+  returning <- enquo(returning)
+
   y <- auto_copy(x, y, copy = copy)
   y_key <- db_key(y, by)
   by <- names(y_key)
@@ -100,9 +114,9 @@ rows_update.tbl_dbi <- function(x, y, by = NULL, ...,
     }
 
     con <- dbplyr::remote_con(x)
-    sql <- sql_rows_update(x, y, by)
-    dbExecute(con, sql, immediate = TRUE)
-    invisible(x)
+    sql <- sql_rows_update(x, y, by, returning = returning)
+
+    get_or_execute(x, con, sql, returning)
   } else {
     # Checking optional, can rely on primary key constraint
     # FIXME: contrary to doc currently also checks if `in_place = FALSE`
@@ -129,7 +143,9 @@ rows_update.tbl_dbi <- function(x, y, by = NULL, ...,
 #' @export
 #' @rdname rows-db
 rows_delete.tbl_dbi <- function(x, y, by = NULL, ...,
-                                in_place = NULL, copy = FALSE, check = NULL) {
+                                in_place = NULL, copy = FALSE, check = NULL,
+                                returning = NULL) {
+  returning <- enquo(returning)
   y <- auto_copy(x, y, copy = copy)
   y_key <- db_key(y, by)
   by <- names(y_key)
@@ -144,9 +160,8 @@ rows_delete.tbl_dbi <- function(x, y, by = NULL, ...,
     }
 
     con <- dbplyr::remote_con(x)
-    sql <- sql_rows_delete(x, y, by)
-    dbExecute(con, sql, immediate = TRUE)
-    invisible(x)
+    sql <- sql_rows_delete(x, y, by, returning = returning)
+    get_or_execute(x, con, sql, returning)
   } else {
     # Checking optional, can rely on primary key constraint
     # FIXME: contrary to doc currently also checks if `in_place = FALSE`
@@ -209,14 +224,15 @@ check_db_superset <- function(x, y, by) {
 #'
 #' @export
 #' @rdname rows-db
-sql_rows_insert <- function(x, y, ...) {
+sql_rows_insert <- function(x, y, returning = NULL, ...) {
+  # TODO `returning` is somehow caught by ellipsis...
   ellipsis::check_dots_used()
   # FIXME: check here same src for x and y? if not -> error.
   UseMethod("sql_rows_insert")
 }
 
 #' @export
-sql_rows_insert.tbl_sql <- function(x, y, ...) {
+sql_rows_insert.tbl_sql <- function(x, y, returning = NULL, ...) {
   con <- dbplyr::remote_con(x)
   name <- dbplyr::remote_name(x)
 
@@ -227,19 +243,21 @@ sql_rows_insert.tbl_sql <- function(x, y, ...) {
     "INSERT INTO ", name, " (", columns_qq, ")\n",
     dbplyr::remote_query(y)
   )
+
+  sql <- add_returning(sql, x, returning)
   glue::as_glue(sql)
 }
 
 #' @export
 #' @rdname rows-db
-sql_rows_update <- function(x, y, by, ...) {
+sql_rows_update <- function(x, y, by, returning = NULL, ...) {
   ellipsis::check_dots_used()
   # FIXME: check here same src for x and y? if not -> error.
   UseMethod("sql_rows_update")
 }
 
 #' @export
-sql_rows_update.tbl_SQLiteConnection <- function(x, y, by, ...) {
+sql_rows_update.tbl_SQLiteConnection <- function(x, y, by, returning = NULL, ...) {
   con <- dbplyr::remote_con(x)
 
   p <- sql_rows_update_prep(x, y, by)
@@ -256,11 +274,12 @@ sql_rows_update.tbl_SQLiteConnection <- function(x, y, by, ...) {
     "WHERE (", p$compare_qual_qq, "))\n",
     "WHERE EXISTS (SELECT * FROM ", p$y_name, " WHERE ", p$compare_qual_qq, ")"
   )
+  sql <- add_returning(sql, x, returning)
   glue::as_glue(sql)
 }
 
 #' @export
-`sql_rows_update.tbl_Microsoft SQL Server` <- function(x, y, by, ...) {
+`sql_rows_update.tbl_Microsoft SQL Server` <- function(x, y, by, returning = NULL, ...) {
   con <- dbplyr::remote_con(x)
 
   p <- sql_rows_update_prep(x, y, by)
@@ -283,11 +302,12 @@ sql_rows_update.tbl_SQLiteConnection <- function(x, y, by, ...) {
     "  INNER JOIN ", p$y_name, "\n",
     "  ON ", p$compare_qual_qq
   )
+  sql <- add_returning(sql, x, returning)
   glue::as_glue(sql)
 }
 
 #' @export
-sql_rows_update.tbl_MariaDBConnection <- function(x, y, by, ...) {
+sql_rows_update.tbl_MariaDBConnection <- function(x, y, by, returning = NULL, ...) {
   con <- dbplyr::remote_con(x)
 
   p <- sql_rows_update_prep(x, y, by)
@@ -300,11 +320,12 @@ sql_rows_update.tbl_MariaDBConnection <- function(x, y, by, ...) {
     "SET\n",
     paste0("  ", p$target_columns_qual_qq, " = ", p$new_columns_qual_qq, collapse = ",\n")
   )
+  sql <- add_returning(sql, x, returning)
   glue::as_glue(sql)
 }
 
 #' @export
-sql_rows_update.tbl_PqConnection <- function(x, y, by, ...) {
+sql_rows_update.tbl_PqConnection <- function(x, y, by, returning = NULL, ...) {
   con <- dbplyr::remote_con(x)
 
   p <- sql_rows_update_prep(x, y, by)
@@ -326,6 +347,7 @@ sql_rows_update.tbl_PqConnection <- function(x, y, by, ...) {
     "FROM ", p$y_name, "\n",
     "WHERE ", p$compare_qual_qq
   )
+  sql <- add_returning(sql, x, returning)
   glue::as_glue(sql)
 }
 
@@ -371,14 +393,14 @@ sql_rows_update_prep <- function(x, y, by) {
 
 #' @export
 #' @rdname rows-db
-sql_rows_delete <- function(x, y, by, ...) {
+sql_rows_delete <- function(x, y, by, ..., returning = NULL) {
   ellipsis::check_dots_used()
   # FIXME: check here same src for x and y? if not -> error.
   UseMethod("sql_rows_delete")
 }
 
 #' @export
-sql_rows_delete.tbl_sql <- function(x, y, by, ...) {
+sql_rows_delete.tbl_sql <- function(x, y, by, ..., returning = NULL) {
   con <- dbplyr::remote_con(x)
 
   p <- sql_rows_update_prep(x, y, by)
@@ -391,5 +413,75 @@ sql_rows_delete.tbl_sql <- function(x, y, by, ...) {
     "DELETE FROM ", p$name, "\n",
     "WHERE EXISTS (SELECT * FROM ", p$y_name, " WHERE ", p$compare_qual_qq, ")"
   )
+  sql <- add_returning(sql, x, returning)
   glue::as_glue(sql)
 }
+
+get_or_execute <- function(x, con, sql, returning) {
+  if (quo_is_null(returning)) {
+    dbExecute(con, sql, immediate = TRUE)
+  } else {
+    returned_rows <- dbGetQuery(con, sql, immediate = TRUE)
+    attr(x, "returned_rows") <- as_tibble(returned_rows)
+  }
+
+  invisible(x)
+}
+
+#' Extract the RETURNING rows
+#'
+#' Extract the RETURNING rows produced by `rows_insert()`, `rows_update()`,
+#' `rows_upsert()`, or `rows_delete().`
+#'
+#' @param x A lazy tbl.
+#'
+#' @return A tibble.
+#'
+#' @export
+get_returned_rows <- function(x) {
+  attr(x, "returned_rows")
+}
+
+add_returning <- function(sql, x, returning) {
+  if (quo_is_null(returning)) {
+    return(sql)
+  }
+
+  con <- dbplyr::remote_con(x)
+
+  ret_vars <- eval_select_both(returning, colnames(x))$names
+  ret_cols <- sql_named_cols(con, ret_vars, table = dbplyr::remote_name(x))
+
+  paste0(sql, "\n", returning_clause(con), " ", ret_cols)
+}
+
+sql_named_cols <- function(con, cols, table = NULL) {
+  nms <- names2(cols)
+  nms[nms == cols] <- ""
+  # Workaround for incorrect column names after `RETURNING`
+  # https://github.com/r-dbi/RSQLite/issues/381
+  if (inherits(con, "SQLiteConnection")) {
+    nms[nms == ""] <- cols[nms == ""]
+  }
+
+  cols <- DBI::dbQuoteIdentifier(con, cols)
+  if (!is.null(table)) {
+    table <- DBI::dbQuoteIdentifier(con, table)
+    cols <- paste0(table, ".", cols)
+  }
+
+  cols[nms != ""] <- paste0(cols, " AS ", nms[nms != ""])
+  paste0(cols, collapse = ", ")
+}
+
+returning_clause <- function(con) {
+  UseMethod("returning_clause")
+}
+
+returning_clause.DBIConnection <- function(con) {
+  "RETURNING"
+}
+
+`returning_clause.Microsoft SQL Server` <- function(con) {
+  "OUTPUT"
+}