diff --git a/DESCRIPTION b/DESCRIPTION index ac4c2b91b..b3c530e1f 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -33,6 +33,8 @@ Suggests: knitr, magrittr, nanoarrow (>= 0.3.0.1), + otel, + otelsdk, RMariaDB, rmarkdown, rprojroot, diff --git a/R/11-dbAppendTable.R b/R/11-dbAppendTable.R index 7d6a1e1cf..859638c17 100644 --- a/R/11-dbAppendTable.R +++ b/R/11-dbAppendTable.R @@ -39,6 +39,15 @@ setGeneric( "dbAppendTable", def = function(conn, name, value, ..., row.names = NULL) { + otel_local_active_span( + "INSERT INTO", + conn, + label = collection_name(name, conn), + attributes = list( + db.collection.name = collection_name(name, conn), + db.operation.name = "INSERT INTO" + ) + ) standardGeneric("dbAppendTable") } ) diff --git a/R/12-dbCreateTable.R b/R/12-dbCreateTable.R index 38e9ac204..294b25f25 100644 --- a/R/12-dbCreateTable.R +++ b/R/12-dbCreateTable.R @@ -37,6 +37,15 @@ setGeneric( "dbCreateTable", def = function(conn, name, fields, ..., row.names = NULL, temporary = FALSE) { + otel_local_active_span( + "CREATE TABLE", + conn, + label = collection_name(name, conn), + attributes = list( + db.collection.name = collection_name(name, conn), + db.operation.name = "CREATE TABLE" + ) + ) standardGeneric("dbCreateTable") } ) diff --git a/R/13-dbWriteTable.R b/R/13-dbWriteTable.R index 7440c714a..c81a71d84 100644 --- a/R/13-dbWriteTable.R +++ b/R/13-dbWriteTable.R @@ -46,5 +46,11 @@ #' dbWriteTable(con, "mtcars", mtcars[1:10, ], overwrite = TRUE, row.names = FALSE) #' dbReadTable(con, "mtcars") setGeneric("dbWriteTable", def = function(conn, name, value, ...) { + otel_local_active_span( + "dbWriteTable", + conn, + label = collection_name(name, conn), + attributes = list(db.collection.name = collection_name(name, conn)) + ) standardGeneric("dbWriteTable") }) diff --git a/R/21-dbAppendTableArrow.R b/R/21-dbAppendTableArrow.R index 099bb72fb..61a74b6b3 100644 --- a/R/21-dbAppendTableArrow.R +++ b/R/21-dbAppendTableArrow.R @@ -29,5 +29,14 @@ #' dbReadTable(con, "iris") #' dbDisconnect(con) setGeneric("dbAppendTableArrow", def = function(conn, name, value, ...) { + otel_local_active_span( + "INSERT INTO", + conn, + label = collection_name(name, conn), + attributes = list( + db.collection.name = collection_name(name, conn), + db.operation.name = "INSERT INTO" + ) + ) standardGeneric("dbAppendTableArrow") }) diff --git a/R/22-dbCreateTableArrow.R b/R/22-dbCreateTableArrow.R index 2f18ee505..080ded3e0 100644 --- a/R/22-dbCreateTableArrow.R +++ b/R/22-dbCreateTableArrow.R @@ -35,6 +35,15 @@ setGeneric( "dbCreateTableArrow", def = function(conn, name, value, ..., temporary = FALSE) { + otel_local_active_span( + "CREATE TABLE", + conn, + label = collection_name(name, conn), + attributes = list( + db.collection.name = collection_name(name, conn), + db.operation.name = "CREATE TABLE" + ) + ) standardGeneric("dbCreateTableArrow") } ) diff --git a/R/23-dbWriteTableArrow.R b/R/23-dbWriteTableArrow.R index 58b9ee03f..74862fe07 100644 --- a/R/23-dbWriteTableArrow.R +++ b/R/23-dbWriteTableArrow.R @@ -37,5 +37,11 @@ #' #' dbDisconnect(con) setGeneric("dbWriteTableArrow", def = function(conn, name, value, ...) { + otel_local_active_span( + "dbWriteTableArrow", + conn, + label = collection_name(name, conn), + attributes = list(db.collection.name = collection_name(name, conn)) + ) standardGeneric("dbWriteTableArrow") }) diff --git a/R/DBI-package.R b/R/DBI-package.R index 72b60687c..0ecf6361c 100644 --- a/R/DBI-package.R +++ b/R/DBI-package.R @@ -28,3 +28,7 @@ require_arrow <- function() { } stop("The nanoarrow package is required for this functionality.") } + +.onLoad <- function(libname, pkgname) { + otel_cache_tracer() +} diff --git a/R/dbConnect.R b/R/dbConnect.R index eeb7db706..dd880019a 100644 --- a/R/dbConnect.R +++ b/R/dbConnect.R @@ -41,6 +41,9 @@ #' dbListTables(con <- dbConnect(RSQLite::SQLite(), ":memory:")) setGeneric( "dbConnect", - def = function(drv, ...) standardGeneric("dbConnect"), + def = function(drv, ...) { + otel_local_active_span("dbConnect", drv) + standardGeneric("dbConnect") + }, valueClass = "DBIConnection" ) diff --git a/R/dbDisconnect.R b/R/dbDisconnect.R index 4f7016f8c..fe84a903b 100644 --- a/R/dbDisconnect.R +++ b/R/dbDisconnect.R @@ -16,5 +16,6 @@ #' con <- dbConnect(RSQLite::SQLite(), ":memory:") #' dbDisconnect(con) setGeneric("dbDisconnect", def = function(conn, ...) { + otel_local_active_span("dbDisconnect", conn) standardGeneric("dbDisconnect") }) diff --git a/R/dbGetQuery.R b/R/dbGetQuery.R index ad7e34011..ee541b821 100644 --- a/R/dbGetQuery.R +++ b/R/dbGetQuery.R @@ -61,5 +61,11 @@ #' #' dbDisconnect(con) setGeneric("dbGetQuery", def = function(conn, statement, ...) { + otel_local_active_span( + dynGet("attributes")$db.operation.name, + conn, + label = dynGet("attributes")$db.collection.name, + attributes = make_query_attributes(statement) + ) standardGeneric("dbGetQuery") }) diff --git a/R/dbGetQueryArrow.R b/R/dbGetQueryArrow.R index bb808ad33..a3d9da53c 100644 --- a/R/dbGetQueryArrow.R +++ b/R/dbGetQueryArrow.R @@ -55,5 +55,11 @@ #' #' dbDisconnect(con) setGeneric("dbGetQueryArrow", def = function(conn, statement, ...) { + otel_local_active_span( + dynGet("attributes")$db.operation.name, + conn, + label = dynGet("attributes")$db.collection.name, + attributes = make_query_attributes(statement) + ) standardGeneric("dbGetQueryArrow") }) diff --git a/R/dbReadTable.R b/R/dbReadTable.R index 21609fc8f..ede093de7 100644 --- a/R/dbReadTable.R +++ b/R/dbReadTable.R @@ -36,6 +36,14 @@ #' dbDisconnect(con) setGeneric( "dbReadTable", - def = function(conn, name, ...) standardGeneric("dbReadTable"), + def = function(conn, name, ...) { + otel_local_active_span( + "dbReadTable", + conn, + label = collection_name(name, conn), + attributes = list(db.collection.name = collection_name(name, conn)) + ) + standardGeneric("dbReadTable") + }, valueClass = "data.frame" ) diff --git a/R/dbReadTableArrow.R b/R/dbReadTableArrow.R index 6f64f105f..486153039 100644 --- a/R/dbReadTableArrow.R +++ b/R/dbReadTableArrow.R @@ -32,5 +32,11 @@ #' dbDisconnect(con) setGeneric("dbReadTableArrow", def = function(conn, name, ...) { require_arrow() + otel_local_active_span( + "dbReadTableArrow", + conn, + label = collection_name(name, conn), + attributes = list(db.collection.name = collection_name(name, conn)) + ) standardGeneric("dbReadTableArrow") }) diff --git a/R/dbRemoveTable.R b/R/dbRemoveTable.R index 8452c6dda..9564eb430 100644 --- a/R/dbRemoveTable.R +++ b/R/dbRemoveTable.R @@ -25,5 +25,14 @@ #' #' dbDisconnect(con) setGeneric("dbRemoveTable", def = function(conn, name, ...) { + otel_local_active_span( + "DROP TABLE", + conn, + label = collection_name(name, conn), + attributes = list( + db.collection.name = collection_name(name, conn), + db.operation.name = "DROP TABLE" + ) + ) standardGeneric("dbRemoveTable") }) diff --git a/R/otel.R b/R/otel.R new file mode 100644 index 000000000..d48e0bda9 --- /dev/null +++ b/R/otel.R @@ -0,0 +1,66 @@ +otel_tracer_name <- "org.r-dbi.DBI" + +# Generic otel helpers: + +otel_cache_tracer <- NULL +otel_local_active_span <- NULL + +local({ + otel_tracer <- NULL + otel_is_tracing <- FALSE + + otel_cache_tracer <<- function() { + requireNamespace("otel", quietly = TRUE) || return() + otel_tracer <<- otel::get_tracer(otel_tracer_name) + otel_is_tracing <<- tracer_enabled(otel_tracer) + } + + otel_local_active_span <<- function( + name, + conn, + label = NULL, + attributes = NULL, + activation_scope = parent.frame() + ) { + otel_is_tracing || return() + dbname <- get_dbname(conn) + otel::start_local_active_span( + name = sprintf("%s %s", name, if (length(label)) label else dbname), + attributes = c(attributes, list(db.system.name = dbname)), + options = list(kind = "client"), + tracer = otel_tracer, + activation_scope = activation_scope + ) + } +}) + +tracer_enabled <- function(tracer) { + .subset2(tracer, "is_enabled")() +} + +with_otel_record <- function(expr) { + on.exit(otel_cache_tracer()) + otelsdk::with_otel_record({ + otel_cache_tracer() + expr + }) +} + +# DBI-specific helpers: + +get_dbname <- function(conn) { + dbname <- attr(class(conn), "package") + if (is.null(dbname)) "unknown" else dbname +} + +collection_name <- function(name, conn) { + if (is.character(name)) name else dbQuoteIdentifier(conn, x = name) +} + +make_query_attributes <- function(statement) { + query <- strsplit(statement, " ", fixed = TRUE)[[1L]] + list( + db.operation.name = query[1L], + db.collection.name = query[which(query == "FROM") + 1L] + ) +} diff --git a/R/zzz.R b/R/zzz.R index aa8a5c205..034598a71 100644 --- a/R/zzz.R +++ b/R/zzz.R @@ -2,6 +2,7 @@ # https://github.com/r-lib/pkgload/issues/247 .onLoad <- function(libname, pkgname) { + otel_cache_tracer() if ( "RSQLite" %in% loadedNamespaces() && diff --git a/tests/testthat/test-otel.R b/tests/testthat/test-otel.R new file mode 100644 index 000000000..2b0eedeaf --- /dev/null +++ b/tests/testthat/test-otel.R @@ -0,0 +1,45 @@ +test_that("OpenTelemetry tracing works", { + skip_if_not_installed("otelsdk") + + record <- with_otel_record({ + + con <- dbConnect(RSQLite::SQLite(), ":memory:") + dbWriteTable(con, "mtcars", mtcars) + dbGetQuery(con, "SELECT * FROM mtcars") + dbGetQuery( + con, + "SELECT COUNT(*) FROM mtcars WHERE cyl = ?", + params = list(1:8) + ) + dbReadTable(con, "mtcars") + dbRemoveTable(con, "mtcars") + dbDisconnect(con) + }) + + traces <- record$traces + + expect_length(traces, 10L) + expect_equal(traces[[1L]]$name, "dbConnect RSQLite") + expect_equal(traces[[1L]]$kind, "client") + expect_equal(traces[[1L]]$attributes$db.system.name, "RSQLite") + expect_equal(traces[[2L]]$name, "CREATE TABLE mtcars") + expect_equal(traces[[2L]]$attributes$db.collection.name, "mtcars") + expect_equal(traces[[2L]]$attributes$db.operation.name, "CREATE TABLE") + expect_equal(traces[[3L]]$name, "INSERT INTO mtcars") + expect_equal(traces[[3L]]$attributes$db.collection.name, "mtcars") + expect_equal(traces[[3L]]$attributes$db.operation.name, "INSERT INTO") + expect_equal(traces[[4L]]$name, "dbWriteTable mtcars") + expect_equal(traces[[4L]]$attributes$db.collection.name, "mtcars") + expect_equal(traces[[5L]]$name, "SELECT mtcars") + expect_equal(traces[[5L]]$attributes$db.collection.name, "mtcars") + expect_equal(traces[[5L]]$attributes$db.operation.name, "SELECT") + expect_equal(traces[[6L]]$name, "SELECT mtcars") + expect_equal(traces[[7L]]$name, "SELECT `mtcars`") + expect_equal(traces[[8L]]$name, "dbReadTable mtcars") + expect_equal(traces[[8L]]$attributes$db.collection.name, "mtcars") + expect_equal(traces[[9L]]$name, "DROP TABLE mtcars") + expect_equal(traces[[9L]]$attributes$db.collection.name, "mtcars") + expect_equal(traces[[9L]]$attributes$db.operation.name, "DROP TABLE") + expect_equal(traces[[10L]]$name, "dbDisconnect RSQLite") + expect_equal(traces[[10L]]$attributes$db.system.name, "RSQLite") +})