From cd1902ef137c2a9b8f77e379b543eab0b8f6365c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kirill=20M=C3=BCller?= <krlmlr+r@mailbox.org>
Date: Mon, 4 Jul 2022 03:59:53 +0200
Subject: [PATCH] Support dbplyr.trace option

---
 DESCRIPTION                   |  1 +
 R/trace.R                     | 58 +++++++++++++++++++++++++++++++++++
 tests/testthat/test-tbl-sql.R |  6 ++--
 3 files changed, 62 insertions(+), 3 deletions(-)
 create mode 100644 R/trace.R

diff --git a/DESCRIPTION b/DESCRIPTION
index 279459ca7..e61f68611 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -126,6 +126,7 @@ Collate:
     'testthat.R'
     'tidyeval-across.R'
     'tidyeval.R'
+    'trace.R'
     'translate-sql.R'
     'utils-format.R'
     'verb-arrange.R'
diff --git a/R/trace.R b/R/trace.R
new file mode 100644
index 000000000..57529f2fa
--- /dev/null
+++ b/R/trace.R
@@ -0,0 +1,58 @@
+tracing_level <- function() {
+  trace <- getOption("dbplyr.trace")
+  if (is.null(trace)) {
+    return(0L)
+  }
+  if (is.logical(trace)) {
+    return(as.integer(trace))
+  }
+  if (!is_integerish(trace)) {
+    warn("Invalid value for dbplyr.trace option, resetting.")
+    options(dbplyr.trace = NULL)
+    return(0)
+  }
+  as.integer(trace)
+}
+
+tracing_id <- function() {
+  # Needs to use option to unique IDs across reloads while testing
+  i <- getOption("dbplyr.trace_id", 0) + 1
+  options(dbplyr.trace_id = i)
+  i
+}
+
+dbGetQuery <- function(conn, statement, ...) {
+  level <- tracing_level()
+  id <- tracing_id()
+
+  if (level >= 1) {
+    message_base <- paste0("[", id, "]: dbGetQuery()")
+    message_pre <- paste0(message_base, "\n", statement)
+    message_post <- paste0(message_base, " done")
+    class <- c("dplyr_message_trace_get_query", "dplyr_message_trace", "dplyr_message")
+    inform(message_pre, class = class)
+    on.exit({
+      inform(message_post, class = class)
+    })
+  }
+
+  DBI::dbGetQuery(conn, statement, ...)
+}
+
+dbExecute <- function(conn, statement, ...) {
+  level <- tracing_level()
+  id <- tracing_id()
+
+  if (level >= 1) {
+    message_base <- paste0("[", id, "]: dbExecute()")
+    message_pre <- paste0(message_base, "\n", statement)
+    message_post <- paste0(message_base, " done")
+    class <- c("dplyr_message_trace_execute", "dplyr_message_trace", "dplyr_message")
+    inform(message_pre, class = class)
+    on.exit({
+      inform(message_post, class = class)
+    })
+  }
+
+  DBI::dbExecute(conn, statement, ...)
+}
diff --git a/tests/testthat/test-tbl-sql.R b/tests/testthat/test-tbl-sql.R
index a3b9ad131..2723866c9 100644
--- a/tests/testthat/test-tbl-sql.R
+++ b/tests/testthat/test-tbl-sql.R
@@ -39,7 +39,7 @@ test_that("sql tbl can be printed", {
 test_that("can refer to default schema explicitly", {
   con <- sqlite_con_with_aux()
   on.exit(DBI::dbDisconnect(con))
-  DBI::dbExecute(con, "CREATE TABLE t1 (x)")
+  dbExecute(con, "CREATE TABLE t1 (x)")
 
   expect_equal(as.character(tbl_vars(tbl(con, "t1"))), "x")
   expect_equal(as.character(tbl_vars(tbl(con, in_schema("main", "t1")))), "x")
@@ -48,8 +48,8 @@ test_that("can refer to default schema explicitly", {
 test_that("can distinguish 'schema.table' from 'schema'.'table'", {
   con <- sqlite_con_with_aux()
   on.exit(DBI::dbDisconnect(con))
-  DBI::dbExecute(con, "CREATE TABLE aux.t1 (x, y, z)")
-  DBI::dbExecute(con, "CREATE TABLE 'aux.t1' (a, b, c)")
+  dbExecute(con, "CREATE TABLE aux.t1 (x, y, z)")
+  dbExecute(con, "CREATE TABLE 'aux.t1' (a, b, c)")
 
   expect_equal(as.character(tbl_vars(tbl(con, in_schema("aux", "t1")))), c("x", "y", "z"))
   expect_equal(as.character(tbl_vars(tbl(con, ident("aux.t1")))), c("a", "b", "c"))