R/datasource-dbi.R

In shunsambongi/datafire:

# Connection --------------------------------------------------------------

#' @export
connect.DBIDriver <- function(x, ..., .pool = FALSE) {
  factory <- DBI::dbConnect
  if (.pool) {
    factory <- pool::dbPool
  }
  rlang::exec(factory, drv = x, !!!list2(...))
}

#' @export
disconnect.DBIConnection <- function(x, ...) {
  DBI::dbDisconnect(conn = x, ...)
}

#' @export
disconnect.Pool <- function(x, ...) {
  pool::poolClose(x)
}

#' @export
is_connected.DBIConnection <- function(x, ...) {
  DBI::dbIsValid(x, ...)
}

#' @export
is_connected.Pool <- is_connected.DBIConnection


# Read --------------------------------------------------------------------

#' @export
read_dataset.DBIConnection <- function(src, name, ..., .lazy = FALSE) {
  ellipsis::check_dots_used()
  if (.lazy) {
    return(dplyr::tbl(src, name, ...))
  }
  data <- DBI::dbReadTable(conn = src, name = name, ...)
  try_tibble(data)
}

#' @export
`read_dataset.Microsoft SQL Server` <- function(src, name, schema = NULL, ..., .lazy = FALSE) {
  schema_provided <- !is.null(schema)
  if (schema_provided && .lazy) {
    name <- dbplyr::in_schema(schema = schema, table = name)
  } else if (schema_provided && !.lazy) {
    name <- DBI::Id(schema = schema, table = name)
  }
  read_dataset.DBIConnection(src, name, ..., .lazy = .lazy)
}

#' @export
read_dataset.Pool <- function(src, name, ...) {
  con <- pool::poolCheckout(src)
  on.exit(pool::poolReturn(con))
  read_dataset(con, name, ...)
}


# Write -------------------------------------------------------------------

#' @export
write_dataset.DBIConnection <- function(dest, name, x, ..., temporary = FALSE) {
  ellipsis::check_dots_used()
  dplyr::copy_to(dest = dest, name = name, df = x, temporary = temporary, ...)
}

#' @export
`write_dataset.Microsoft SQL Server` <- function(dest, name, x, schema = NULL, ..., temporary = FALSE) {
  if (!is.null(schema)) {
    name <- dbplyr::in_schema(schema = schema, table = name)
  }
  write_dataset.DBIConnection(dest, name, x, ..., temporary = temporary)
}

#' @export
write_dataset.Pool <- function(dest, name, x, ...) {
  con <- pool::poolCheckout(dest)
  on.exit(pool::poolReturn(con))
  write_dataset(con, name, x, ...)
}


# Remove ------------------------------------------------------------------

#' @export
remove_dataset.DBIConnection <- function(loc, name, ...) {
  ellipsis::check_dots_used()
  DBI::dbRemoveTable(loc, name, ...)
}

#' @export
remove_dataset.Pool <- remove_dataset.DBIConnection


# List --------------------------------------------------------------------

#' @export
list_datasets.DBIConnection <- function(src, ...) {
  ellipsis::check_dots_used()
  DBI::dbListTables(conn = src, ...)
}

#' @export
list_datasets.Pool <- list_datasets.DBIConnection



# Exists ------------------------------------------------------------------

#' @export
exists_dataset.DBIConnection <- function(src, name, ...) {
  ellipsis::check_dots_used()
  DBI::dbExistsTable(conn = src, name = name, ...)
}

#' @export
exists_dataset.Pool <- exists_dataset.DBIConnection