R/obj-tbl_db.R

In tjpalanca/dbtools: Database Tools for R

#' @title
#' Writable Database Tibble
#'
#' @description
#' This tibble subclass extends the [`dbplyr::tbl_dbi`] class and allows
#' DML statements to be written against the table, in addition to the excellent
#' querying language already provided by `{dbplyr}`.
#'
#' @param src    (DBIConnection) connection object produced by [`connect()`]
#' @param table  (str) name of the table in the database
#' @param schema (str) schema, if applicable
#' @param object (obj) object of interest
#'
#'
#' @family Database Objects
#' @export
tbl_db <- function(src, table, schema = NULL) {

  assert_true(is_native(src))
  assert_string(table)
  assert_string(schema, null.ok = TRUE)

  tbl(
    src = src,
    if (!is.null(schema)) in_schema(schema, table) else table
  ) %>%
    add_class(paste0("tbl_db_", class(src)[1]), "tbl_db", pkg_name) %>%
    validate_tbl_db()

}

#' @export
collect.tbl_db <- function(x, ...) {

  df_db(NextMethod(), x)

}

#' @rdname tbl_db
#' @export
get_tbl_db <- function(object) {
  UseMethod("get_tbl_db")
}

#' @export
get_tbl_db.default <- function(object) {
  object
}

#' @export
print.tbl_db <- function(x, ...) {

  # If in transaction, then show the original table
  if (in_transaction(x)) {
    cli_rule("Current Transaction")
    NextMethod()
    original_x <- x
    original_x$src <- original_x$old
    print(original_x)
  } else {
    cli_rule("Database Table")
    NextMethod()
  }

  # If it is printed without assignment, roll back the changes to be safe
  if (dbInTransaction(get_tbl_db_conn(x))) {
    dbRollback(get_tbl_db_conn(x))
    dbBegin(get_tbl_db_conn(x))
  }

}

validate_tbl_db <- function(tbl_db) {

  # Writable tibbles need a master database connection
  assert_false(is_read_only_tbl_db(tbl_db))

  return(tbl_db)

}

assert_tbl_db <- function(x) {

  assert_class(x, "tbl_db")

}

get_tbl_db_conn <- function(tbl_db) {

  assert_tbl_db(tbl_db)
  tbl_db[["src"]][["con"]]

}

get_tbl_db_from <- function(tbl_db) {

  assert_tbl_db(tbl_db)

  root <- get_tbl_db_root(tbl_db)

  from <-
    str_remove(root$x, "^.*from ") %>%
    str_remove("\"") %>%
    str_split("\\.") %>%
    extract2(1)

  assert_character(from, max.len = 2L, min.len = 1L)

  if (length(from) == 1) {

    Id(table = from[[1]])

  } else {

    Id(schema = from[[1]], table = from[[2]])

  }

}

get_tbl_db_root <- function(tbl_db) {

  assert_tbl_db(tbl_db)

  find_root <- function(ops) {

    if (!inherits(ops$x, "op")) {

      assert_multi_class(ops$x, c("ident", "sql"))
      return(ops)

    } else {

      assert_class(ops$x, "op")
      return(find_root(ops$x))

    }

  }

  find_root(tbl_db$ops)

}

get_tbl_db_pkey <- function(tbl_db) {

  dbGetPrimaryKey(get_tbl_db_conn(tbl_db), get_tbl_db_from(tbl_db))

}

is_read_only_tbl_db <- function(tbl_db) {

  dbIsReadOnly(get_tbl_db_conn(tbl_db))

}