feat: Add support for OpenTelemetry

DESCRIPTION

@@ -33,6 +33,8 @@ Suggests:
     knitr,
     magrittr,
     nanoarrow (>= 0.3.0.1),
+    otel,
+    otelsdk,
     RMariaDB,
     rmarkdown,
     rprojroot,

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")
   }
 )

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")
   }
 )

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")
 })

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")
 })

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")
   }
 )

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")
 })

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()
+}

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"
 )

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")
 })

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")
 })

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")
 })

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"
 )

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")
 })

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")
 })

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]
+  )
+}

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() &&

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")
+})