Overhaul viewer-based credentials for Databricks & Snowflake (#894)

This commit takes a second stab at the viewer-based credential support added in #853. It contains two major changes: 1. Users no longer need to pass a `session` argument. We figure it out internally instead. This allows users to deploy apps to Connect with no code changes. 2. The underlying code to get these credentials is now hosted in the `connectcreds` package, which also provides a number of nice mocking features we can use in `odbc`'s unit tests. The `connectcreds` implementation is based on the code that this commit removes, though it has a number of improvements as well, so this should be a net win, too. Unit tests and documentation is included. Signed-off-by: Aaron Jacobs <[email protected]>
r-dbi · Feb 11, 2025 · 721f7e9 · 721f7e9
1 parent 208652d
commit 721f7e9
Show file tree

Hide file tree

Showing 13 changed files with 84 additions and 210 deletions.
diff --git a/DESCRIPTION b/DESCRIPTION
@@ -29,6 +29,7 @@ Imports:
     Rcpp (>= 0.12.11),
     rlang (>= 1.1.0)
 Suggests:
+    connectcreds,
     covr,
     DBItest,
     httr2,

diff --git a/NEWS.md b/NEWS.md
@@ -48,9 +48,8 @@
 
 * SQL Server: Fix issue related to writing when using SIMBA drivers (#816).
 
-* `snowflake()` and `databricks()` now accept a `session` argument for passing
-  viewer-based OAuth credentials from Shiny sessions on Posit Connect
-  (@atheriel, #853).
+* `snowflake()` and `databricks()` now detect viewer-based credentials when
+  running on Posit Connect (@atheriel, #853 and #894).
 
 * A new `redshift()` helper greatly simplifies connecting to Amazon Redshift
   clusters, particularly when using IAM credentials (@atheriel, #879).

diff --git a/R/driver-databricks.R b/R/driver-databricks.R
@@ -14,9 +14,11 @@ NULL
 #' implements a subset of the [Databricks client unified authentication](https://docs.databricks.com/en/dev-tools/auth.html#databricks-client-unified-authentication)
 #' model, with support for personal access tokens, OAuth machine-to-machine
 #' credentials, and OAuth user-to-machine credentials supplied via Posit
-#' Workbench or the Databricks CLI on desktop.
-#' All of these credentials are detected automatically if present using
-#' [standard environment variables](https://docs.databricks.com/en/dev-tools/auth.html#environment-variables-and-fields-for-client-unified-authentication).
+#' Workbench or the Databricks CLI on desktop. It can also detect viewer-based
+#' credentials on Posit Connect if the \pkg{connectcreds} package is
+#' installed. All of these credentials are detected automatically if present
+#' using [standard environment variables](https://docs.databricks.com/en/dev-tools/auth.html#environment-variables-and-fields-for-client-unified-authentication).
+#'
 #' In addition, on macOS platforms, the `dbConnect()` method will check
 #' for irregularities with how the driver is configured,
 #' and attempt to fix in-situ, unless the `odbc.no_config_override`
@@ -35,8 +37,6 @@ NULL
 #'   default name.
 #' @param uid,pwd Manually specify a username and password for authentication.
 #'   Specifying these options will disable automated credential discovery.
-#' @param session A Shiny session object, when using viewer-based credentials on
-#'   Posit Connect.
 #' @param ... Further arguments passed on to [`dbConnect()`].
 #'
 #' @returns An `OdbcConnection` object with an active connection to a Databricks
@@ -51,11 +51,11 @@ NULL
 #'
 #' # Use credentials from the viewer (when possible) in a Shiny app
 #' # deployed to Posit Connect.
+#' library(connectcreds)
 #' server <- function(input, output, session) {
 #'   conn <- DBI::dbConnect(
 #'     odbc::databricks(),
-#'     httpPath = "sql/protocolv1/o/4425955464597947/1026-023828-vn51jugj",
-#'     session = session
+#'     httpPath = "sql/protocolv1/o/4425955464597947/1026-023828-vn51jugj"
 #'   )
 #' }
 #' }
@@ -79,7 +79,6 @@ setMethod("dbConnect", "DatabricksOdbcDriver",
            HTTPPath,
            uid = NULL,
            pwd = NULL,
-           session = NULL,
            ...) {
     call <- caller_env()
     # For backward compatibility with RStudio connection string
@@ -90,7 +89,6 @@ setMethod("dbConnect", "DatabricksOdbcDriver",
     check_string(driver, allow_null = TRUE, call = call)
     check_string(uid, allow_null = TRUE, call = call)
     check_string(pwd, allow_null = TRUE, call = call)
-    check_shiny_session(session, allow_null = TRUE, call = call)
 
     args <- databricks_args(
       httpPath = if (missing(httpPath)) HTTPPath else httpPath,
@@ -99,7 +97,6 @@ setMethod("dbConnect", "DatabricksOdbcDriver",
       driver = driver,
       uid = uid,
       pwd = pwd,
-      session = session,
       ...
     )
     # Perform some sanity checks on MacOS
@@ -115,7 +112,6 @@ databricks_args <- function(httpPath,
                             driver = NULL,
                             uid = NULL,
                             pwd = NULL,
-                            session = NULL,
                             ...) {
   host <- databricks_host(workspace)
 
@@ -126,7 +122,7 @@ databricks_args <- function(httpPath,
     useNativeQuery = useNativeQuery
   )
 
-  auth <- databricks_auth_args(host, uid = uid, pwd = pwd, session = session)
+  auth <- databricks_auth_args(host, uid = uid, pwd = pwd)
   all <- utils::modifyList(c(args, auth), list(...))
 
   arg_names <- tolower(names(all))
@@ -137,7 +133,11 @@ databricks_args <- function(httpPath,
     )
     if (running_on_connect()) {
       msg <- c(
-        msg, "i" = "Or pass {.arg session} for viewer-based credentials."
+        msg,
+        "i" = "Or consider enabling Posit Connect's Databricks integration \
+              for viewer-based credentials. See {.url \
+              https://docs.posit.co/connect/user/oauth-integrations/#adding-oauth-integrations-to-deployed-content}
+              for details."
       )
     }
     cli::cli_abort(msg, call = quote(DBI::dbConnect()))
@@ -215,14 +215,16 @@ databricks_user_agent <- function() {
   user_agent
 }
 
-databricks_auth_args <- function(host, uid = NULL, pwd = NULL, session = NULL) {
-  # If a session is supplied, any viewer-based auth takes precedence.
-  if (!is.null(session)) {
-    check_installed("httr2", "for viewer-based authentication")
-    access_token <- connect_viewer_token(session, paste0("https://", host))
-    if (!is.null(access_token)) {
-      return(list(authMech = 11, auth_flow = 0, auth_accesstoken = access_token))
-    }
+databricks_auth_args <- function(host, uid = NULL, pwd = NULL) {
+  # Detect viewer-based credentials from Posit Connect.
+  workspace <- paste0("https://", host)
+  if (is_installed("connectcreds") && connectcreds::has_viewer_token(workspace)) {
+    token <- connectcreds::connect_viewer_token(workspace)
+    return(list(
+      authMech = 11,
+      auth_flow = 0,
+      auth_accesstoken = token$access_token
+    ))
   }
 
   if (!is.null(uid) && !is.null(pwd)) {

diff --git a/R/driver-snowflake.R b/R/driver-snowflake.R
@@ -98,10 +98,13 @@ setMethod("odbcDataType", "Snowflake",
 #'
 #' In particular, the custom `dbConnect()` method for the Snowflake ODBC driver
 #' detects ambient OAuth credentials on platforms like Snowpark Container
-#' Services or Posit Workbench.  In addition, on macOS platforms, the
-#' `dbConnect` method will check and warn if it detects irregularities with
-#' how the driver is configured, unless the `odbc.no_config_override`
-#' environment variable is set.
+#' Services or Posit Workbench. It can also detect viewer-based
+#' credentials on Posit Connect if the \pkg{connectcreds} package is
+#' installed.
+#'
+#' In addition, on macOS platforms, the `dbConnect` method will check and warn
+#' if it detects irregularities with how the driver is configured, unless the
+#' `odbc.no_config_override` environment variable is set.
 #'
 #' @inheritParams DBI::dbConnect
 #' @param account A Snowflake [account
@@ -117,8 +120,6 @@ setMethod("odbcDataType", "Snowflake",
 #'   default.
 #' @param uid,pwd Manually specify a username and password for authentication.
 #'   Specifying these options will disable ambient credential discovery.
-#' @param session A Shiny session object, when using viewer-based credentials on
-#'   Posit Connect.
 #' @param ... Further arguments passed on to [`dbConnect()`].
 #'
 #' @returns An `OdbcConnection` object with an active connection to a Snowflake
@@ -146,8 +147,9 @@ setMethod("odbcDataType", "Snowflake",
 #'
 #' # Use credentials from the viewer (when possible) in a Shiny app
 #' # deployed to Posit Connect.
+#' library(connectcreds)
 #' server <- function(input, output, session) {
-#'   conn <- DBI::dbConnect(odbc::snowflake(), session = session)
+#'   conn <- DBI::dbConnect(odbc::snowflake())
 #' }
 #' }
 #' @export
@@ -171,7 +173,6 @@ setMethod(
            schema = NULL,
            uid = NULL,
            pwd = NULL,
-           session = NULL,
            ...) {
     call <- caller_env()
     check_string(account, call = call)
@@ -180,7 +181,6 @@ setMethod(
     check_string(database, allow_null = TRUE, call = call)
     check_string(uid, allow_null = TRUE, call = call)
     check_string(pwd, allow_null = TRUE, call = call)
-    check_shiny_session(session, allow_null = TRUE, call = call)
     args <- snowflake_args(
       account = account,
       driver = driver,
@@ -189,7 +189,6 @@ setMethod(
       schema = schema,
       uid = uid,
       pwd = pwd,
-      session = session,
       ...
     )
     # Perform some sanity checks on MacOS
@@ -229,7 +228,11 @@ snowflake_args <- function(account = Sys.getenv("SNOWFLAKE_ACCOUNT"),
     )
     if (running_on_connect()) {
       msg <- c(
-        msg, "i" = "Or pass {.arg session} for viewer-based credentials."
+        msg,
+        "i" = "Or consider enabling Posit Connect's Snowflake integration \
+              for viewer-based credentials. See {.url \
+              https://docs.posit.co/connect/user/oauth-integrations/#adding-oauth-integrations-to-deployed-content}
+              for details."
       )
     }
     cli::cli_abort(msg, call = quote(DBI::dbConnect()))
@@ -310,18 +313,13 @@ snowflake_auth_args <- function(account,
                                 uid = NULL,
                                 pwd = NULL,
                                 authenticator = NULL,
-                                session = NULL,
                                 ...) {
   check_string(authenticator, allow_null = TRUE)
-  # If a session is supplied, any viewer-based auth takes precedence.
-  if (!is.null(session)) {
-    check_installed("httr2", "for viewer-based authentication")
-    access_token <- connect_viewer_token(
-      session, paste0("https://", account, ".snowflakecomputing.com")
-    )
-    if (!is.null(access_token)) {
-      return(list(authenticator = "oauth", token = access_token))
-    }
+  # Detect viewer-based credentials from Posit Connect.
+  url <- paste0("https://", account, ".snowflakecomputing.com")
+  if (is_installed("connectcreds") && connectcreds::has_viewer_token(url)) {
+    token <- connectcreds::connect_viewer_token(url)
+    return(list(authenticator = "oauth", token = token$access_token))
   }
 
   if (!is.null(uid) &&

diff --git a/R/utils.R b/R/utils.R
@@ -517,78 +517,6 @@ replace_or_append <- function(lines, key_pattern, accepted_value, replacement) {
   return(list(new_lines = lines, modified = !found_ok))
 }
 
-check_shiny_session <- function(x,
-                                ...,
-                                allow_null = FALSE,
-                                arg = caller_arg(x),
-                                call = caller_env()) {
-  if (!missing(x)) {
-    if (inherits(x, "ShinySession")) {
-      return(invisible(NULL))
-    }
-    if (allow_null && is_null(x)) {
-      return(invisible(NULL))
-    }
-  }
-
-  stop_input_type(
-    x,
-    "a Shiny session object",
-    ...,
-    allow_null = allow_null,
-    arg = arg,
-    call = call
-  )
-}
-
-# Request an OAuth access token for the given resource from Posit Connect. The
-# OAuth token will belong to the user owning the given Shiny session.
-connect_viewer_token <- function(session, resource) {
-  # Ensure we're running on Connect.
-  server_url <- Sys.getenv("CONNECT_SERVER")
-  api_key <- Sys.getenv("CONNECT_API_KEY")
-  if (!running_on_connect() || nchar(server_url) == 0 || nchar(api_key) == 0) {
-    cli::cli_inform(c(
-      "!" = "Ignoring {.arg sesssion} parameter.",
-      "i" = "Viewer-based credentials are only available when running on Connect."
-    ))
-    return(NULL)
-  }
-
-  # Older versions or certain configurations of Connect might not supply a user
-  # session token.
-  token <- session$request$HTTP_POSIT_CONNECT_USER_SESSION_TOKEN
-  if (is.null(token)) {
-    cli::cli_abort(
-      "Viewer-based credentials are not supported by this version of Connect."
-    )
-  }
-
-  # See: https://docs.posit.co/connect/api/#post-/v1/oauth/integrations/credentials
-  req <- httr2::request(server_url)
-  req <- httr2::req_url_path_append(
-    req, "__api__/v1/oauth/integrations/credentials"
-  )
-  req <- httr2::req_headers(req,
-    Authorization = paste("Key", api_key), .redact = "Authorization"
-  )
-  req <- httr2::req_body_form(
-    req,
-    grant_type = "urn:ietf:params:oauth:grant-type:token-exchange",
-    subject_token_type = "urn:posit:connect:user-session-token",
-    subject_token = token,
-    resource = resource
-  )
-
-  # TODO: Do we need more precise error handling?
-  req <- httr2::req_error(
-    req, body = function(resp) httr2::resp_body_json(resp)$error
-  )
-
-  resp <- httr2::resp_body_json(httr2::req_perform(req))
-  resp$access_token
-}
-
 running_on_connect <- function() {
   Sys.getenv("RSTUDIO_PRODUCT") == "CONNECT"
 }

diff --git a/man/databricks.Rd b/man/databricks.Rd