Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Driver registry #124

Merged
merged 21 commits into from
Mar 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions R/Connection.R
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
#' @export
setClass("duckdb_driver", contains = "DBIDriver", slots = list(
database_ref = "externalptr",
config = "list",
dbdir = "character",
read_only = "logical",
bigint = "character"
Expand All @@ -25,17 +26,19 @@ setClass("duckdb_connection", contains = "DBIConnection", slots = list(
debug = "logical",
timezone_out = "character",
tz_out_convert = "character",
reserved_words = "character"
reserved_words = "character",
bigint = "character"
))

duckdb_connection <- function(duckdb_driver, debug) {
duckdb_connection <- function(duckdb_driver, debug, bigint) {
out <- new(
"duckdb_connection",
conn_ref = rapi_connect(duckdb_driver@database_ref),
driver = duckdb_driver,
debug = debug,
timezone_out = "UTC",
tz_out_convert = "with"
tz_out_convert = "with",
bigint = bigint
)
out@reserved_words <- get_reserved_words(out)
out
Expand Down
74 changes: 60 additions & 14 deletions R/Driver.R
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,11 @@ drv_to_string <- function(drv) {
if (!is(drv, "duckdb_driver")) {
stop("pass a duckdb_driver object")
}
sprintf("<duckdb_driver %s dbdir='%s' read_only=%s bigint=%s>", extptr_str(drv@database_ref), drv@dbdir, drv@read_only, drv@bigint)
sprintf("<duckdb_driver dbdir='%s' read_only=%s bigint=%s>", drv@dbdir, drv@read_only, drv@bigint)
}

driver_registry <- new.env(parent = emptyenv())

#' @description
#' `duckdb()` creates or reuses a database instance.
#'
Expand All @@ -27,18 +29,20 @@ drv_to_string <- function(drv) {
#' @export
duckdb <- function(dbdir = DBDIR_MEMORY, read_only = FALSE, bigint = "numeric", config = list()) {
check_flag(read_only)

switch(bigint,
numeric = {
# fine
},
integer64 = {
if (!is_installed("bit64")) {
stop("bit64 package is required for integer64 support")
}
},
stop(paste("Unsupported bigint configuration", bigint))
)
check_bigint(bigint)

dbdir <- path_normalize(dbdir)
if (dbdir != DBDIR_MEMORY) {
drv <- driver_registry[[dbdir]]
# We reuse an existing driver object if the database is still alive.
# If not, we fall back to creating a new driver object with a new database.
if (!is.null(drv) && rapi_lock(drv@database_ref)) {
# We don't care about different read_only or config settings here.
# The bigint setting can be actually picked up by dbConnect(), we update it here.
drv@bigint <- bigint
return(drv)
}
}

# R packages are not allowed to write extensions into home directory, so use R_user_dir instead
if (!("extension_directory" %in% names(config))) {
Expand All @@ -48,13 +52,21 @@ duckdb <- function(dbdir = DBDIR_MEMORY, read_only = FALSE, bigint = "numeric",
config["secret_directory"] <- file.path(tools::R_user_dir("duckdb", "data"), "stored_secrets")
}

new(
# Always create new database for in-memory,
# allows isolation and mixing different configs
drv <- new(
"duckdb_driver",
config = config,
database_ref = rapi_startup(dbdir, read_only, config),
dbdir = dbdir,
read_only = read_only,
bigint = bigint
)

if (dbdir != DBDIR_MEMORY) {
driver_registry[[dbdir]] <- drv
}
drv
}

#' @description
Expand All @@ -73,6 +85,11 @@ duckdb_shutdown <- function(drv) {
invisible(FALSE)
}
rapi_shutdown(drv@database_ref)

if (drv@dbdir != DBDIR_MEMORY) {
rm(list = drv@dbdir, envir = driver_registry)
}

invisible(TRUE)
}

Expand Down Expand Up @@ -139,3 +156,32 @@ check_tz <- function(timezone) {

timezone
}
path_normalize <- function(path) {
if (path == "" || path == DBDIR_MEMORY) {
return(DBDIR_MEMORY)
}

out <- normalizePath(path, mustWork = FALSE)

# Stable results are only guaranteed if the file exists
if (!file.exists(out)) {
on.exit(unlink(out))
writeLines(character(), out)
out <- normalizePath(out, mustWork = TRUE)
}
out
}

check_bigint <- function(bigint_type) {
switch(bigint_type,
numeric = {
# fine
},
integer64 = {
if (!is_installed("bit64")) {
stop("bit64 package is required for integer64 support")
}
},
stop(paste("Unsupported bigint configuration", bigint_type))
)
}
2 changes: 1 addition & 1 deletion R/Result.R
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ duckdb_result <- function(connection, stmt_lst, arrow) {
}

duckdb_execute <- function(res) {
out <- rapi_execute(res@stmt_lst$ref, res@arrow, res@connection@driver@bigint == "integer64")
out <- rapi_execute(res@stmt_lst$ref, res@arrow, res@connection@bigint == "integer64")
duckdb_post_execute(res, out)
}

Expand Down
20 changes: 18 additions & 2 deletions R/cpp11.R
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# Generated by cpp11: do not edit by hand

rapi_connect <- function(db) {
.Call(`_duckdb_rapi_connect`, db)
rapi_connect <- function(dual) {
.Call(`_duckdb_rapi_connect`, dual)
}

rapi_disconnect <- function(conn) {
Expand All @@ -12,6 +12,18 @@ rapi_startup <- function(dbdir, readonly, configsexp) {
.Call(`_duckdb_rapi_startup`, dbdir, readonly, configsexp)
}

rapi_lock <- function(dual) {
.Call(`_duckdb_rapi_lock`, dual)
}

rapi_unlock <- function(dual) {
invisible(.Call(`_duckdb_rapi_unlock`, dual))
}

rapi_is_locked <- function(dual) {
.Call(`_duckdb_rapi_is_locked`, dual)
}

rapi_shutdown <- function(dbsexp) {
invisible(.Call(`_duckdb_rapi_shutdown`, dbsexp))
}
Expand All @@ -32,6 +44,10 @@ rapi_unregister_arrow <- function(conn, name) {
invisible(.Call(`_duckdb_rapi_unregister_arrow`, conn, name))
}

rapi_list_arrow <- function(conn) {
.Call(`_duckdb_rapi_list_arrow`, conn)
}

rapi_expr_reference <- function(rnames) {
.Call(`_duckdb_rapi_expr_reference`, rnames)
}
Expand Down
2 changes: 1 addition & 1 deletion R/dbBind__duckdb_result.R
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ dbBind__duckdb_result <- function(res, params, ...) {

params <- encode_values(params)

out <- rapi_bind(res@stmt_lst$ref, params, res@arrow, res@connection@driver@bigint == "integer64")
out <- rapi_bind(res@stmt_lst$ref, params, res@arrow, res@connection@bigint == "integer64")
if (length(out) == 1) {
out <- out[[1]]
} else if (length(out) == 0) {
Expand Down
43 changes: 33 additions & 10 deletions R/dbConnect__duckdb_driver.R
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,14 @@
#'
#' @param drv Object returned by `duckdb()`
#' @param dbdir Location for database files. Should be a path to an existing
#' directory in the file system. With the default, all
#' data is kept in RAM
#' directory in the file system. With the default (or `""`), all
#' data is kept in RAM.
#' @param ... Ignored
#' @param debug Print additional debug information such as queries
#' @param read_only Set to `TRUE` for read-only operation
#' @param read_only Set to `TRUE` for read-only operation.
#' For file-based databases, this is only applied when the database file is opened for the first time.
#' Subsequent connections (via the same `drv` object or a `drv` object pointing to the same path)
#' will silently ignore this flag.
#' @param timezone_out The time zone returned to R, defaults to `"UTC"`, which
#' is currently the only timezone supported by duckdb.
#' If you want to display datetime values in the local timezone,
Expand All @@ -18,8 +21,13 @@
#' is chosen, the timestamp will be returned as it would appear in the specified time zone.
#' If `"force"` is chosen, the timestamp will have the same clock
#' time as the timestamp in the database, but with the new time zone.
#' @param config Named list with DuckDB configuration flags
#' @param bigint How 64-bit integers should be returned, default is double/numeric. Set to integer64 for bit64 encoding.
#' @param config Named list with DuckDB configuration flags, see
#' <https://duckdb.org/docs/configuration/overview#configuration-reference> for the possible options.
#' These flags are only applied when the database object is instantiated.
#' Subsequent connections will silently ignore these flags.
#' @param bigint How 64-bit integers should be returned. There are two options: `"numeric"` and `"integer64"`.
#' If `"numeric"` is selected, bigint integers will be treated as double/numeric.
#' If `"integer64"` is selected, bigint integers will be set to bit64 encoding.
#'
#' @return `dbConnect()` returns an object of class
#' \linkS4class{duckdb_connection}.
Expand Down Expand Up @@ -53,20 +61,35 @@ dbConnect__duckdb_driver <- function(
timezone_out <- check_tz(timezone_out)
tz_out_convert <- match.arg(tz_out_convert)

missing_dbdir <- missing(dbdir)
dbdir <- path.expand(as.character(dbdir))
if (missing(dbdir)) {
dbdir <- drv@dbdir
} else {
dbdir <- path_normalize(dbdir)
}

if (missing(read_only)) {
read_only <- drv@read_only
}

if (missing(bigint)) {
bigint <- drv@bigint
} else {
check_bigint(bigint)
}

config <- utils::modifyList(drv@config, config)

# aha, a late comer. let's make a new instance.
if (!missing_dbdir && dbdir != drv@dbdir) {
if (dbdir != drv@dbdir || !rapi_lock(drv@database_ref)) {
rapi_unlock(drv@database_ref)
drv <- duckdb(dbdir, read_only, bigint, config)
}

conn <- duckdb_connection(drv, debug = debug)
conn <- duckdb_connection(drv, debug = debug, bigint = bigint)
on.exit(dbDisconnect(conn))

conn@timezone_out <- timezone_out
conn@tz_out_convert <- tz_out_convert

on.exit(NULL)

rs_on_connection_opened(conn)
Expand Down
12 changes: 5 additions & 7 deletions R/dbDisconnect__duckdb_connection.R
Original file line number Diff line number Diff line change
@@ -1,20 +1,18 @@
#' @description
#' `dbDisconnect()` closes a DuckDB database connection, optionally shutting down
#' the associated instance.
#' `dbDisconnect()` closes a DuckDB database connection.
#' The associated DuckDB database instance is shut down automatically,
#' it is no longer necessary to set `shutdown = TRUE` or to call `duckdb_shutdown()`.
#'
#' @param conn A `duckdb_connection` object
#' @param shutdown Set to `TRUE` to shut down the DuckDB database instance that this connection refers to.
#' @param shutdown Unused. The database instance is shut down automatically.
#' @rdname duckdb
#' @usage NULL
dbDisconnect__duckdb_connection <- function(conn, ..., shutdown = FALSE) {
dbDisconnect__duckdb_connection <- function(conn, ..., shutdown = TRUE) {
if (!dbIsValid(conn)) {
warning("Connection already closed.", call. = FALSE)
invisible(FALSE)
}
rapi_disconnect(conn@conn_ref)
if (shutdown) {
duckdb_shutdown(conn@driver)
}
rs_on_connection_closed(conn)
invisible(TRUE)
}
Expand Down
7 changes: 4 additions & 3 deletions R/dbGetInfo__duckdb_connection.R
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,11 @@
#' @inheritParams DBI::dbGetInfo
#' @usage NULL
dbGetInfo__duckdb_connection <- function(dbObj, ...) {
info <- dbGetInfo(dbObj@driver)
version <- dbGetQuery(dbObj, "select library_version from pragma_version()")[[1]][[1]]

list(
dbname = info$dbname,
db.version = info$driver.version,
dbname = dbObj@driver@dbdir,
db.version = version,
username = NA,
host = NA,
port = NA
Expand Down
10 changes: 6 additions & 4 deletions R/dbGetInfo__duckdb_driver.R
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,12 @@
#' @inheritParams DBI::dbGetInfo
#' @usage NULL
dbGetInfo__duckdb_driver <- function(dbObj, ...) {
con <- dbConnect(dbObj)
version <- dbGetQuery(con, "select library_version from pragma_version()")[[1]][[1]]
dbDisconnect(con)
list(driver.version = version, client.version = version, dbname = dbObj@dbdir)
info <- dbGetInfo__duckdb_connection(default_connection())
list(
driver.version = info$db.version,
client.version = info$db.version,
dbname = dbObj@dbdir
)
}

#' @rdname duckdb_driver-class
Expand Down
6 changes: 6 additions & 0 deletions R/dbIsValid__duckdb_driver.R
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,13 @@ dbIsValid__duckdb_driver <- function(dbObj, ...) {
valid <- FALSE
tryCatch(
{
was_locked <- rapi_is_locked(dbObj@database_ref)
con <- dbConnect(dbObj)
# Keep driver alive, but only if needed
if (was_locked) {
rapi_lock(dbObj@database_ref)
}

dbExecute(con, SQL("SELECT 1"))
dbDisconnect(con)
valid <- TRUE
Expand Down
13 changes: 1 addition & 12 deletions R/dbQuoteIdentifier__duckdb_connection.R
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ dbQuoteIdentifier__duckdb_connection <- function(conn, x, ...) {
}

x <- enc2utf8(x)
needs_escape <- !grepl("^[a-zA-Z_][a-zA-Z0-9_]*$", x) | tolower(x) %in% reserved_words(conn)
needs_escape <- !grepl("^[a-zA-Z_][a-zA-Z0-9_]*$", x) | tolower(x) %in% conn@reserved_words
x[needs_escape] <- paste0('"', gsub('"', '""', x[needs_escape]), '"')

SQL(x, names = names(x))
Expand All @@ -30,17 +30,6 @@ setMethod("dbQuoteIdentifier", signature("duckdb_connection", "SQL"), dbQuoteIde
#' @usage NULL
setMethod("dbQuoteIdentifier", signature("duckdb_connection", "Id"), dbQuoteIdentifier__duckdb_connection)

reserved_words <- function(con) {
if (!isS4(con)) {
con <- dbConnect(duckdb())
words <- get_reserved_words(con)
dbDisconnect(con, shutdown = TRUE)
words
}

con@reserved_words
}

get_reserved_words <- function(con) {
dbGetQuery(con, "SELECT keyword_name FROM duckdb_keywords()")[[1]]
}
4 changes: 2 additions & 2 deletions R/register.R
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ encode_values <- function(value) {
duckdb_register <- function(conn, name, df, overwrite = FALSE, experimental = FALSE) {
stopifnot(dbIsValid(conn))
df <- encode_values(as.data.frame(df))
rapi_register_df(conn@conn_ref, enc2utf8(as.character(name)), df, conn@driver@bigint == "integer64", overwrite, experimental)
rapi_register_df(conn@conn_ref, enc2utf8(as.character(name)), df, conn@bigint == "integer64", overwrite, experimental)
invisible(TRUE)
}

Expand Down Expand Up @@ -113,5 +113,5 @@ duckdb_unregister_arrow <- function(conn, name) {
#' @rdname duckdb_register_arrow
#' @export
duckdb_list_arrow <- function(conn) {
sort(gsub("_registered_arrow_", "", names(attributes(conn@driver@database_ref)), fixed = TRUE))
sort(rapi_list_arrow(conn@conn_ref))
}
Loading
Loading