Fix OAuth isolation across workspaces (#255) (#256)

Co-authored-by: Zac Davies <zaclavis@gmail.com>

Author: zacdav-db

NEWS.md

@@ -1,5 +1,6 @@
 # brickster (development version)
 
+-   Fixed OAuth U2M and M2M authentication across multiple Databricks workspaces in one R session by isolating OAuth clients and cached tokens per workspace (#255)
 -   Fixed `dbWriteTable()` and `dbAppendTable()` standard-path writes for binary columns, which now use Databricks `BINARY` types and `X'...'` literals when no staging volume is configured (#245)
 -   `dbConnect()` now preserves Databricks API error details when its validation query fails
 -   `db_perform_request()` and `db_perform_response()` now preserve useful error messages when Databricks returns non-JSON error bodies such as empty, plain-text, or HTML responses (#242)

R/package-auth.R

@@ -462,20 +462,45 @@ resolve_oauth_auth_mode <- function(
   "oauth-u2m"
 }
 
+db_oauth_client_cache_name <- function(
+  auth_mode,
+  host,
+  client_id,
+  token_url,
+  scope
+) {
+  cache_context <- list(
+    auth_mode = auth_mode,
+    host = tolower(host),
+    client_id = client_id,
+    token_url = token_url,
+    scope = scope
+  )
+
+  paste0("brickster-", auth_mode, "-", rlang::hash(cache_context))
+}
+
 build_databricks_m2m_oauth_client <- function(host, client_id, client_secret) {
   endpoints <- databricks_workspace_oauth_endpoints(host)
+  scope <- "all-apis"
 
   list(
     client = httr2::oauth_client(
       id = client_id,
       secret = client_secret,
       token_url = endpoints$token_url,
-      name = "brickster"
+      name = db_oauth_client_cache_name(
+        auth_mode = "oauth-m2m",
+        host = host,
+        client_id = client_id,
+        token_url = endpoints$token_url,
+        scope = scope
+      )
     ),
     auth_url = endpoints$auth_url,
     auth_mode = "oauth-m2m",
     is_m2m = TRUE,
-    scope = "all-apis",
+    scope = scope,
     token_params = list()
   )
 }
@@ -488,40 +513,58 @@ databricks_workspace_oauth_endpoints <- function(host) {
 }
 
 build_azure_m2m_oauth_client <- function(
+  host,
   azure_client_id,
   azure_client_secret,
   azure_tenant_id
 ) {
+  token_url <- glue::glue(
+    "https://login.microsoftonline.com/{azure_tenant_id}/oauth2/v2.0/token"
+  )
+  scope <- "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default"
+
   list(
     client = httr2::oauth_client(
       id = azure_client_id,
       secret = azure_client_secret,
-      token_url = glue::glue(
-        "https://login.microsoftonline.com/{azure_tenant_id}/oauth2/v2.0/token"
-      ),
-      name = "brickster"
+      token_url = token_url,
+      name = db_oauth_client_cache_name(
+        auth_mode = "azure-client-secret",
+        host = host,
+        client_id = azure_client_id,
+        token_url = token_url,
+        scope = scope
+      )
     ),
     auth_url = NULL,
     auth_mode = "azure-client-secret",
     is_m2m = TRUE,
-    scope = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default",
+    scope = scope,
     token_params = list()
   )
 }
 
 build_databricks_u2m_oauth_client <- function(host) {
   endpoints <- databricks_workspace_oauth_endpoints(host)
+  client_id <- "databricks-cli"
+  scope <- "all-apis"
 
   list(
     client = httr2::oauth_client(
-      id = "databricks-cli",
+      id = client_id,
       token_url = endpoints$token_url,
-      name = "brickster"
+      name = db_oauth_client_cache_name(
+        auth_mode = "oauth-u2m",
+        host = host,
+        client_id = client_id,
+        token_url = endpoints$token_url,
+        scope = scope
+      )
     ),
     auth_url = endpoints$auth_url,
     auth_mode = "oauth-u2m",
     is_m2m = FALSE,
-    scope = "all-apis",
+    scope = scope,
     token_params = list()
   )
 }
@@ -581,6 +624,7 @@ db_oauth_client <- function(
     )
   } else if (identical(auth_mode, "azure-client-secret")) {
     client_and_auth <- build_azure_m2m_oauth_client(
+      host,
       azure_client_id,
       azure_client_secret,
       azure_tenant_id
@@ -589,9 +633,6 @@ db_oauth_client <- function(
     client_and_auth <- build_databricks_u2m_oauth_client(host)
   }
 
-  # add option for client to be fetched via request helpers
-  options(brickster_oauth_client = client_and_auth)
-
   client_and_auth
 }
 

R/request-helpers.R

@@ -47,11 +47,7 @@ db_request <- function(
   if (!is.null(token)) {
     req <- httr2::req_auth_bearer_token(req = req, token = token)
   } else {
-    # fetch client
-    oauth_client <- getOption(
-      x = "brickster_oauth_client",
-      db_oauth_client(host = host)
-    )
+    oauth_client <- db_oauth_client(host = host)
 
     if (oauth_client$is_m2m) {
       req <- httr2::req_oauth_client_credentials(

tests/testthat/test-auth.R

@@ -10,11 +10,40 @@ local_clear_auth_env <- function() {
   withr::local_options(
     use_databrickscfg = FALSE,
     db_profile = NULL,
-    brickster_oauth_client = NULL,
     .local_envir = parent.frame()
   )
 }
 
+test_that("OAuth client cache identities are stable and workspace-specific", {
+  local_clear_auth_env()
+
+  u2m_a <- db_oauth_client(
+    host = "workspace-a.cloud.databricks.com",
+    client_id = NULL,
+    client_secret = NULL,
+    auth_type = "oauth-u2m"
+  )
+  u2m_a_again <- db_oauth_client(
+    host = "workspace-a.cloud.databricks.com",
+    client_id = NULL,
+    client_secret = NULL,
+    auth_type = "oauth-u2m"
+  )
+  u2m_b <- db_oauth_client(
+    host = "workspace-b.cloud.databricks.com",
+    client_id = NULL,
+    client_secret = NULL,
+    auth_type = "oauth-u2m"
+  )
+
+  expect_identical(u2m_a$client$name, u2m_a_again$client$name)
+  expect_false(identical(u2m_a$client$name, u2m_b$client$name))
+  expect_identical(
+    u2m_b$auth_url,
+    "https://workspace-b.cloud.databricks.com/oidc/v1/authorize"
+  )
+})
+
 test_that("auth functions - baseline behaviour", {
   local_clear_auth_env()
 

tests/testthat/test-request-helpers.R

@@ -10,7 +10,6 @@ local_clear_auth_env <- function() {
   withr::local_options(
     use_databrickscfg = FALSE,
     db_profile = NULL,
-    brickster_oauth_client = NULL,
     .local_envir = parent.frame()
   )
 }
@@ -111,8 +110,6 @@ test_that("request helpers - m2m auth flow", {
     DATABRICKS_CLIENT_ID = "client-id",
     DATABRICKS_CLIENT_SECRET = "client-secret"
   )
-  withr::local_options(brickster_oauth_client = NULL)
-
   req <- db_request(
     endpoint = endpoint,
     method = method,
@@ -136,6 +133,56 @@ test_that("request helpers - m2m auth flow", {
   )
 })
 
+test_that("request helpers isolate OAuth clients and tokens by workspace", {
+  local_clear_auth_env()
+
+  withr::local_envvar(
+    DATABRICKS_CLIENT_ID = "client-id",
+    DATABRICKS_CLIENT_SECRET = "client-secret"
+  )
+
+  req_a <- db_request(
+    endpoint = "clusters/list",
+    method = "GET",
+    version = "2.0",
+    host = "workspace-a.cloud.databricks.com",
+    token = NULL
+  )
+  req_b <- db_request(
+    endpoint = "clusters/list",
+    method = "GET",
+    version = "2.0",
+    host = "workspace-b.cloud.databricks.com",
+    token = NULL
+  )
+
+  client_a <- req_a$policies$auth_sign$params$flow_params$client
+  client_b <- req_b$policies$auth_sign$params$flow_params$client
+  expect_identical(
+    client_a$token_url,
+    "https://workspace-a.cloud.databricks.com/oidc/v1/token"
+  )
+  expect_identical(
+    client_b$token_url,
+    "https://workspace-b.cloud.databricks.com/oidc/v1/token"
+  )
+  expect_false(identical(client_a$name, client_b$name))
+
+  cache_a <- req_a$policies$auth_sign$cache
+  cache_b <- req_b$policies$auth_sign$cache
+  cache_a$clear()
+  cache_b$clear()
+  withr::defer(cache_a$clear())
+  withr::defer(cache_b$clear())
+
+  cache_a$set(httr2::oauth_token(
+    access_token = "workspace-a-token",
+    expires_in = 3600
+  ))
+
+  expect_null(cache_b$get())
+})
+
 test_that("request helpers - azure m2m auth flow", {
   local_clear_auth_env()
 
@@ -150,8 +197,6 @@ test_that("request helpers - azure m2m auth flow", {
     ARM_CLIENT_SECRET = "azure-client-secret",
     ARM_TENANT_ID = "azure-tenant-id"
   )
-  withr::local_options(brickster_oauth_client = NULL)
-
   req <- db_request(
     endpoint = endpoint,
     method = method,