Add support for viewer-based Databricks & Snowflake credentials.

This commit expands the `databricks()` and `snowflake()` helpers to support the viewer-based OAuth credentials recently introduced in Posit Connect [0]. It is designed to support writing Shiny apps that take the following form: ```r library(shiny) ui <- fluidPage(textOutput("user")) server <- function(input, output, session) { conn <- reactive({ pool::dbPool(odbc::snowflake(), session = session) }) output$user <- renderText({ res <- DBI::dbGetQuery(conn(), "SELECT CURRENT_USER()") res[[1]] }) } shinyApp(ui = ui, server = server) ``` Checks for viewer-based credentials are designed to fall back gracefully to existing authentication methods in some cases. This is intended to allow users to -- for example -- develop and test a Shiny app that uses Databricks or Snowflake credentials in desktop RStudio or Posit Workbench and deploy it with no code changes to Connect. In addition, making the caller pass a `session` argument (rather than just detecting, for example, if we're in a Shiny server context) is very intentional: it makes them express that they *want* viewer-based credentials, as opposed to using shared credentials for all viewers. It seems likely to reduce the number of cases where publisher credentials are used unexpectedly. Internally we implement the raw HTTP calls to make the token exchange with Connect to avoid taking a dependency on the `connectapi` package. Unit tests are included for most of the new error paths, but it's pretty hard to emulate what Connect is doing here without extensive mocking. [0]: https://docs.posit.co/connect/user/oauth-integrations/ Signed-off-by: Aaron Jacobs <[email protected]>
r-dbi · Oct 29, 2024 · 9c625c4 · 9c625c4
1 parent bd66867
commit 9c625c4
Show file tree

Hide file tree

Showing 13 changed files with 293 additions and 17 deletions.
diff --git a/DESCRIPTION b/DESCRIPTION
@@ -31,6 +31,7 @@ Imports:
 Suggests:
     covr,
     DBItest,
+    httr2,
     knitr,
     magrittr,
     rmarkdown,

diff --git a/NEWS.md b/NEWS.md
@@ -25,6 +25,10 @@
 
 * 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).
+
 # odbc 1.5.0
 
 ## Major changes

diff --git a/R/driver-databricks.R b/R/driver-databricks.R
@@ -31,6 +31,8 @@ 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
@@ -42,6 +44,16 @@ NULL
 #'   odbc::databricks(),
 #'   httpPath = "sql/protocolv1/o/4425955464597947/1026-023828-vn51jugj"
 #' )
+#'
+#' # Use credentials from the viewer (when possible) in a Shiny app
+#' # deployed to Posit Connect.
+#' server <- function(input, output, session) {
+#'   conn <- DBI::dbConnect(
+#'     odbc::databricks(),
+#'     httpPath = "sql/protocolv1/o/4425955464597947/1026-023828-vn51jugj",
+#'     session = session
+#'   )
+#' }
 #' }
 #' @export
 databricks <- function() {
@@ -64,6 +76,7 @@ setMethod("dbConnect", "DatabricksOdbcDriver",
            HTTPPath,
            uid = NULL,
            pwd = NULL,
+           session = NULL,
            ...) {
     call <- caller_env()
     # For backward compatibility with RStudio connection string
@@ -74,6 +87,7 @@ 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,
@@ -82,6 +96,7 @@ setMethod("dbConnect", "DatabricksOdbcDriver",
       driver = driver,
       uid = uid,
       pwd = pwd,
+      session = session,
       ...
     )
     inject(dbConnect(odbc(), !!!args))
@@ -94,6 +109,7 @@ databricks_args <- function(httpPath,
                             driver = NULL,
                             uid = NULL,
                             pwd = NULL,
+                            session = NULL,
                             ...) {
   host <- databricks_host(workspace)
 
@@ -104,18 +120,21 @@ databricks_args <- function(httpPath,
     useNativeQuery = useNativeQuery
   )
 
-  auth <- databricks_auth_args(host, uid = uid, pwd = pwd)
+  auth <- databricks_auth_args(host, uid = uid, pwd = pwd, session = session)
   all <- utils::modifyList(c(args, auth), list(...))
 
   arg_names <- tolower(names(all))
   if (!"authmech" %in% arg_names && !all(c("uid", "pwd") %in% arg_names)) {
-    abort(
-      c(
-        "x" = "Failed to detect ambient Databricks credentials.",
-        "i" = "Supply `uid` and `pwd` to authenticate manually."
-      ),
-      call = quote(DBI::dbConnect())
+    msg <- c(
+      "Failed to detect ambient Databricks credentials.",
+      "i" = "Supply {.arg uid} and {.arg pwd} to authenticate manually."
     )
+    if (running_on_connect()) {
+      msg <- c(
+        msg, "i" = "Or pass {.arg session} for viewer-based credentials."
+      )
+    }
+    cli::cli_abort(msg, call = quote(DBI::dbConnect()))
   }
 
   all
@@ -202,7 +221,16 @@ databricks_user_agent <- function() {
   user_agent
 }
 
-databricks_auth_args <- function(host, uid = NULL, pwd = NULL) {
+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))
+    }
+  }
+
   if (!is.null(uid) && !is.null(pwd)) {
     return(list(uid = uid, pwd = pwd, authMech = 3))
   } else if (xor(is.null(uid), is.null(pwd))) {

diff --git a/R/driver-snowflake.R b/R/driver-snowflake.R
@@ -114,6 +114,8 @@ 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
@@ -138,6 +140,12 @@ setMethod("odbcDataType", "Snowflake",
 #'   uid = "me",
 #'   pwd = rstudioapi::askForPassword()
 #' )
+#'
+#' # Use credentials from the viewer (when possible) in a Shiny app
+#' # deployed to Posit Connect.
+#' server <- function(input, output, session) {
+#'   conn <- DBI::dbConnect(odbc::snowflake(), session = session)
+#' }
 #' }
 #' @export
 snowflake <- function() {
@@ -156,6 +164,7 @@ setMethod(
            schema = NULL,
            uid = NULL,
            pwd = NULL,
+           session = NULL,
            ...) {
     call <- caller_env()
     check_string(account, call = call)
@@ -164,6 +173,7 @@ 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,
@@ -172,6 +182,7 @@ setMethod(
       schema = schema,
       uid = uid,
       pwd = pwd,
+      session = session,
       ...
     )
     inject(dbConnect(odbc(), !!!args))
@@ -202,13 +213,16 @@ snowflake_args <- function(account = Sys.getenv("SNOWFLAKE_ACCOUNT"),
 
   arg_names <- tolower(names(all))
   if (!"authenticator" %in% arg_names && !all(c("uid", "pwd") %in% arg_names)) {
-    abort(
-      c(
-        "x" = "Failed to detect ambient Snowflake credentials.",
-        "i" = "Supply `uid` and `pwd` to authenticate manually."
-      ),
-      call = quote(DBI::dbConnect())
+    msg <- c(
+      "Failed to detect ambient Snowflake credentials.",
+      "i" = "Supply {.arg uid} and {.arg pwd} to authenticate manually."
     )
+    if (running_on_connect()) {
+      msg <- c(
+        msg, "i" = "Or pass {.arg session} for viewer-based credentials."
+      )
+    }
+    cli::cli_abort(msg, call = quote(DBI::dbConnect()))
   }
 
   all
@@ -270,7 +284,19 @@ snowflake_auth_args <- function(account,
                                 uid = NULL,
                                 pwd = NULL,
                                 authenticator = 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://", account, ".snowflakecomputing.com")
+    )
+    if (!is.null(access_token)) {
+      return(list(authenticator = "oauth", token = access_token))
+    }
+  }
+
   if (!is.null(uid) &&
       # allow for uid without pwd for externalbrowser auth (#817)
       (!is.null(pwd) || identical(authenticator, "externalbrowser"))) {

diff --git a/R/utils.R b/R/utils.R
@@ -454,3 +454,79 @@ replace_or_append <- function(lines, pattern, replacement) {
   }
   lines
 }
+
+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
diff --git a/man/snowflake.Rd b/man/snowflake.Rd
diff --git a/tests/testthat/_snaps/driver-databricks.md b/tests/testthat/_snaps/driver-databricks.md
@@ -22,7 +22,7 @@
       . <- databricks_args1()
     Condition
       Error in `DBI::dbConnect()`:
-      ! x Failed to detect ambient Databricks credentials.
+      ! Failed to detect ambient Databricks credentials.
       i Supply `uid` and `pwd` to authenticate manually.
 
 # must supply both uid and pwd
@@ -58,3 +58,22 @@
       Error in `dbConnect()`:
       ! `httpPath` must be a single string or `NULL`, not the number 1.
 
+# we mention viewer-based credentials have no effect locally
+
+    Code
+      ignored <- databricks_args(workspace = "workspace", httpPath = "path", uid = "uid",
+        pwd = "pwd", session = list())
+    Message
+      ! Ignoring `sesssion` parameter.
+      i Viewer-based credentials are only available when running on Connect.
+
+# we hint viewer-based credentials on Connect
+
+    Code
+      databricks_args(workspace = "workspace", httpPath = "path", driver = "driver")
+    Condition
+      Error in `DBI::dbConnect()`:
+      ! Failed to detect ambient Databricks credentials.
+      i Supply `uid` and `pwd` to authenticate manually.
+      i Or pass `session` for viewer-based credentials.
+