Merge pull request #805 from r-dbi/aj-db-workbench-credentials-from-file

Read Workbench-managed Databricks credentials from `databricks.cfg`

NEWS.md

@@ -41,6 +41,10 @@
  pass `x` to `encodeString()` before returning, for consistency with the
  default implementation in DBI (@simonpcouch, #765).
 
+* `odbc::databricks()` now picks up on Posit Workbench-managed Databricks
+ credentials when rendering Quarto and RMarkdown documents in RStudio
+ (@atheriel, #805).
+
 # odbc 1.4.2
 
 * `dbAppendTable()` Improve performance by checking existence once (#691).

R/driver-databricks.R

@@ -221,12 +221,12 @@ databricks_auth_args <- function(host, uid = NULL, pwd = NULL) {
  client_id <- Sys.getenv("DATABRICKS_CLIENT_ID")
  client_secret <- Sys.getenv("DATABRICKS_CLIENT_SECRET")
  cli_path <- Sys.getenv("DATABRICKS_CLI_PATH", "databricks")
+ cfg_file <- Sys.getenv("DATABRICKS_CONFIG_FILE")
 
  # Check for Workbench-provided credentials.
  wb_token <- NULL
- if (exists(".rs.api.getDatabricksToken")) {
- getDatabricksToken <- get(".rs.api.getDatabricksToken")
- wb_token <- getDatabricksToken(host)
+ if (grepl("posit-workbench", cfg_file, fixed = TRUE)) {
+ wb_token <- workbench_databricks_token(host, cfg_file)
  }
 
  if (nchar(token) != 0) {
@@ -292,3 +292,26 @@ is_camel_case <- function(x) {
 }
 
 dot_names <- function(...) names(substitute(...()))
+
+# Reads Posit Workbench-managed Databricks credentials from a
+# $DATABRICKS_CONFIG_FILE. The generated file will look as follows:
+#
+# [workbench]
+# host = some-host
+# token = some-token
+workbench_databricks_token <- function(host, cfg_file) {
+ cfg <- readLines(cfg_file)
+ # We don't attempt a full parse of the INI syntax supported by Databricks
+ # config files, instead relying on the fact that this particular file will
+ # always contain only one section.
+ if (!any(grepl(host, cfg, fixed = TRUE))) {
+ # The configuration doesn't actually apply to this host.
+ return(NULL)
+ }
+ line <- grepl("token = ", cfg, fixed = TRUE)
+ token <- gsub("token = ", "", cfg[line])
+ if (nchar(token) == 0) {
+ return(NULL)
+ }
+ token
+}

tests/testthat/test-driver-databricks.R

@@ -51,7 +51,10 @@ test_that("user agent respects envvar", {
 })
 
 test_that("errors if auth fails", {
- withr::local_envvar(DATABRICKS_TOKEN = "")
+ withr::local_envvar(
+ DATABRICKS_TOKEN = "",
+ DATABRICKS_CONFIG_FILE = NULL
+ )
 
  databricks_args1 <- function(...) {
  databricks_args("path", "host", driver = "driver", ...)
@@ -110,3 +113,40 @@ test_that("dbConnect method errors informatively re: httpPath (#787)", {
  expect_snapshot(error = TRUE, dbConnect(databricks(), HTTPPath = 1L))
  expect_snapshot(error = TRUE, dbConnect(databricks(), httpPath = 1L))
 })
+
+test_that("Workbench-managed credentials are detected correctly", {
+ # Emulate the databricks.cfg file written by Workbench.
+ db_home <- tempfile("posit-workbench")
+ dir.create(db_home)
+ writeLines(
+ c(
+ '[workbench]',
+ 'host = some-host',
+ 'token = token'
+ ),
+ file.path(db_home, "databricks.cfg")
+ )
+ withr::local_envvar(
+ DATABRICKS_CONFIG_FILE = file.path(db_home, "databricks.cfg")
+ )
+ args <- databricks_auth_args(host = "some-host")
+ expect_equal(args, list(authMech = 11, auth_flow = 0, auth_accesstoken = "token"))
+})
+
+test_that("Workbench-managed credentials are ignored for other hosts", {
+ # Emulate the databricks.cfg file written by Workbench.
+ db_home <- tempfile("posit-workbench")
+ dir.create(db_home)
+ writeLines(
+ c(
+ '[workbench]',
+ 'host = nonmatching',
+ 'token = token'
+ ),
+ file.path(db_home, "databricks.cfg")
+ )
+ withr::local_envvar(
+ DATABRICKS_CONFIG_FILE = file.path(db_home, "databricks.cfg")
+ )
+ expect_equal(databricks_auth_args(host = "some-host"), NULL)
+})