Use connection disposition for DBI queries (#218)

Co-authored-by: Zac Davies <zachary.davies+data@databricks.com>

Author: zacdav

NEWS.md

@@ -1,6 +1,7 @@
 # brickster 0.2.13
 
 -   Enabled `db_request()` retries for transient low-level HTTP request failures, improving resilience to intermittent curl/HTTP2 framing errors (#215)
+-   Added a DBI connection-level `disposition` setting so `dbSendQuery()` and default `dbGetQuery()` calls can use `INLINE` results when direct cloud-storage downloads are blocked (#205)
 -   Added Azure AD service principal OAuth M2M support (`ARM_CLIENT_ID`, `ARM_CLIENT_SECRET`, `ARM_TENANT_ID`) with optional `DATABRICKS_AUTH_TYPE` override (`oauth-m2m`, `azure-client-secret`, `oauth-u2m`); default auth resolution now prefers Azure M2M over U2M when ARM credentials are present (#185)
 -   Marked DBFS REST wrappers (`db_dbfs_*`) as deprecated and moved them to internal-only, guiding users towards using volumes (`db_volume_*`)
 -   Added `db_volume_download_dir()` for parallel directory downloads from Unity Catalog volumes to local directories

R/databricks-dbi.R

@@ -28,6 +28,7 @@ setClass(
     catalog = "character",
     schema = "character",
     staging_volume = "character",
+    disposition = "character",
     max_active_connections = "numeric",
     fetch_timeout = "numeric"
   )
@@ -83,6 +84,9 @@ setMethod("show", "DatabricksDriver", function(object) {
 #' @param catalog Optional catalog name to use as default
 #' @param schema Optional schema name to use as default
 #' @param staging_volume Optional volume path for large dataset staging
+#' @param disposition Query disposition mode to use by default for DBI query
+#'   results. Use `"EXTERNAL_LINKS"` for large results or `"INLINE"` for
+#'   smaller results that must avoid direct cloud-storage result downloads.
 #' @param max_active_connections Maximum number of concurrent download
 #' connections when fetching query results (default: 30)
 #' @param fetch_timeout Timeout in seconds for downloading each result chunk
@@ -102,12 +106,15 @@ setMethod(
     catalog = NULL,
     schema = NULL,
     staging_volume = NULL,
+    disposition = c("EXTERNAL_LINKS", "INLINE"),
     max_active_connections = 30,
     fetch_timeout = 300,
     token = db_token(),
     host = db_host(),
     ...
   ) {
+    disposition <- match.arg(disposition)
+
     # Validate required parameters
     if (
       !is.null(warehouse_id) &&
@@ -168,6 +175,7 @@ setMethod(
       catalog = catalog %||% "",
       schema = schema %||% "",
       staging_volume = staging_volume %||% "",
+      disposition = disposition,
       max_active_connections = max_active_connections,
       fetch_timeout = fetch_timeout
     )
@@ -218,6 +226,7 @@ setMethod("show", "DatabricksConnection", function(object) {
   if (!is.null(object@staging_volume) && nzchar(object@staging_volume)) {
     cat("  Staging Volume:", object@staging_volume, "\n")
   }
+  cat("  Disposition:", object@disposition, "\n")
   cat("  Max Active Connections:", object@max_active_connections, "\n")
   cat("  Fetch Timeout (s):", object@fetch_timeout, "\n")
 })
@@ -227,23 +236,26 @@ setMethod("show", "DatabricksConnection", function(object) {
 #' Send query to Databricks (asynchronous)
 #' @param conn A DatabricksConnection object
 #' @param statement SQL statement to execute
+#' @param disposition Query disposition mode. Defaults to the connection's
+#'   `disposition` setting.
 #' @param ... Additional arguments (ignored)
 #' @returns A DatabricksResult object
 #' @export
 setMethod(
   "dbSendQuery",
   signature = c(conn = "DatabricksConnection", statement = "character"),
-  function(conn, statement, ...) {
+  function(conn, statement, disposition = conn@disposition, ...) {
     db_assert_statement(statement)
+    disposition <- match.arg(disposition, c("EXTERNAL_LINKS", "INLINE"))
 
     # Execute query asynchronously
     resp <- db_sql_exec_query(
       warehouse_id = conn@warehouse_id,
       statement = statement,
       catalog = if (nzchar(conn@catalog)) conn@catalog else NULL,
       schema = if (nzchar(conn@schema)) conn@schema else NULL,
-      disposition = "EXTERNAL_LINKS",
-      format = "ARROW_STREAM",
+      disposition = disposition,
+      format = if (disposition == "INLINE") "JSON_ARRAY" else "ARROW_STREAM",
       wait_timeout = "0s", # Async execution
       host = conn@host,
       token = conn@token
@@ -265,8 +277,8 @@ setMethod(
 #'
 #' @param conn A DatabricksConnection object
 #' @param statement SQL statement to execute
-#' @param disposition Query disposition mode: "EXTERNAL_LINKS" (default) for large results,
-#'   "INLINE" for small metadata queries (automatically chooses appropriate format)
+#' @param disposition Query disposition mode. Defaults to the connection's
+#'   `disposition` setting.
 #' @param show_progress If `TRUE`, show progress updates during query execution (default: `TRUE`)
 #' @param ... Additional arguments passed to underlying query execution
 #' @returns A data.frame with query results
@@ -277,10 +289,12 @@ setMethod(
   function(
     conn,
     statement,
-    disposition = "EXTERNAL_LINKS",
+    disposition = conn@disposition,
     show_progress = TRUE,
     ...
   ) {
+    disposition <- match.arg(disposition, c("EXTERNAL_LINKS", "INLINE"))
+
     # Detect schema discovery queries (LIMIT 0) and optimize them
     if (endsWith(trimws(statement), "LIMIT 0")) {
       # Force INLINE disposition and disable progress for schema queries
@@ -397,7 +411,11 @@ setMethod("dbFetch", "DatabricksResult", function(res, n = -1, ...) {
   }
 
   # Check if we need to poll for completion and start executing step
-  initial_status <- db_sql_exec_status(statement_id = res@statement_id)
+  initial_status <- db_sql_exec_status(
+    statement_id = res@statement_id,
+    host = res@connection@host,
+    token = res@connection@token
+  )
   if (initial_status$status$state %in% c("RUNNING", "PENDING")) {
     cli::cli_progress_step("Executing query")
     status <- db_sql_exec_poll_for_success(
@@ -415,6 +433,15 @@ setMethod("dbFetch", "DatabricksResult", function(res, n = -1, ...) {
   # Use total_row_count to detect empty result sets
   if (status$manifest$total_row_count == 0) {
     results <- db_sql_create_empty_result(status$manifest)
+  } else if (
+    identical(status$manifest$format, "JSON_ARRAY") ||
+      !is.null(status$result$data_array)
+  ) {
+    results <- db_sql_process_inline(
+      result_data = status$result,
+      manifest = status$manifest,
+      row_limit = if (n > 0) n else NULL
+    )
   } else {
     # Use helper function to fetch results with progress
     results <- db_sql_fetch_results(
@@ -871,7 +898,8 @@ setMethod("dbGetInfo", "DatabricksConnection", function(dbObj, ...) {
     username = NA_character_,
     host = dbObj@host,
     port = NA_integer_,
-    warehouse_id = dbObj@warehouse_id
+    warehouse_id = dbObj@warehouse_id,
+    disposition = dbObj@disposition
   )
 })
 

man/dbConnect-DatabricksDriver-method.Rd

@@ -1,4 +1,4 @@
-make_dbi_test_con <- function(staging_volume = "") {
+make_dbi_test_con <- function(staging_volume = "", disposition = "EXTERNAL_LINKS") {
   new(
     "DatabricksConnection",
     warehouse_id = "test_warehouse",
@@ -7,6 +7,7 @@ make_dbi_test_con <- function(staging_volume = "") {
     catalog = "test_catalog",
     schema = "test_schema",
     staging_volume = staging_volume,
+    disposition = disposition,
     max_active_connections = 30,
     fetch_timeout = 300
   )
@@ -31,12 +32,14 @@ test_that("dbConnect validates tuning inputs and persists connection settings",
     warehouse_id = "wh-1",
     host = "mock_host",
     token = "mock_token",
+    disposition = "INLINE",
     max_active_connections = 12,
     fetch_timeout = 45
   )
 
   expect_s4_class(con, "DatabricksConnection")
   expect_identical(con@warehouse_id, "wh-1")
+  expect_identical(con@disposition, "INLINE")
   expect_identical(con@max_active_connections, 12)
   expect_identical(con@fetch_timeout, 45)
   expect_true(state$opened)
@@ -62,6 +65,17 @@ test_that("dbConnect validates tuning inputs and persists connection settings",
     ),
     "`fetch_timeout` must be a positive numeric value"
   )
+
+  expect_error(
+    dbConnect(
+      drv,
+      warehouse_id = "wh-1",
+      host = "mock_host",
+      token = "mock_token",
+      disposition = "invalid"
+    ),
+    "'arg' should be one of"
+  )
 })
 
 test_that("dbWriteTable validates user input", {
@@ -363,10 +377,61 @@ test_that("query execution DBI methods dispatch expected options", {
   expect_identical(rows_unknown, 0L)
 
   expect_identical(state$query_calls[[1]]$wait_timeout, "0s")
+  expect_identical(state$query_calls[[1]]$disposition, "EXTERNAL_LINKS")
+  expect_identical(state$query_calls[[1]]$format, "ARROW_STREAM")
   expect_identical(state$query_calls[[2]]$wait_timeout, "0s")
   expect_identical(state$query_calls[[3]]$disposition, "INLINE")
   expect_false(state$query_calls[[3]]$show_progress)
   expect_identical(state$query_calls[[4]]$disposition, "EXTERNAL_LINKS")
+
+  con_inline <- make_dbi_test_con(disposition = "INLINE")
+  res_inline <- dbSendQuery(con_inline, "SELECT 2")
+  expect_s4_class(res_inline, "DatabricksResult")
+  expect_identical(state$query_calls[[5]]$disposition, "INLINE")
+  expect_identical(state$query_calls[[5]]$format, "JSON_ARRAY")
+
+  out_inline <- dbGetQuery(con_inline, "SELECT * FROM inline_table")
+  expect_identical(out_inline$ok, TRUE)
+  expect_identical(state$query_calls[[6]]$disposition, "INLINE")
+})
+
+test_that("dbFetch processes inline results from dbSendQuery", {
+  con <- make_dbi_test_con(disposition = "INLINE")
+  res <- new(
+    "DatabricksResult",
+    statement_id = "stmt-inline",
+    statement = "SELECT 1 UNION ALL SELECT 2",
+    connection = con,
+    completed = FALSE,
+    rows_fetched = 0
+  )
+
+  local_mocked_bindings(
+    db_sql_exec_status = function(...) {
+      list(
+        statement_id = "stmt-inline",
+        status = list(state = "SUCCEEDED"),
+        manifest = list(
+          format = "JSON_ARRAY",
+          total_chunk_count = 1L,
+          total_row_count = 2L,
+          schema = list(columns = list(
+            list(name = "id", type_name = "INT")
+          ))
+        ),
+        result = list(data_array = list(list(1L), list(2L)))
+      )
+    },
+    db_sql_fetch_results = function(...) {
+      stop("external result fetcher should not be called")
+    },
+    .package = "brickster"
+  )
+
+  out <- dbFetch(res)
+
+  expect_s3_class(out, "tbl_df")
+  expect_identical(out$id, c(1L, 2L))
 })
 
 test_that("volume-method selection warns/errors at size thresholds", {

man/dbGetQuery-DatabricksConnection-character-method.Rd

@@ -1,12 +1,13 @@
-make_test_con <- function() {
+make_test_con <- function(disposition = "EXTERNAL_LINKS") {
   new(
     "DatabricksConnection",
     warehouse_id = "test_warehouse",
     host = "test_host",
     token = "test_token",
     catalog = "",
     schema = "",
-    staging_volume = ""
+    staging_volume = "",
+    disposition = disposition
   )
 }
 
@@ -84,6 +85,25 @@ test_that("sql_query_save covers temporary and persistent branches", {
   expect_match(table_sql, "`already_quoted`")
 })
 
+test_that("db_collect uses connection query disposition", {
+  con <- make_test_con(disposition = "INLINE")
+  state <- new.env(parent = emptyenv())
+  state$disposition <- NULL
+
+  local_mocked_bindings(
+    dbGetQuery = function(conn, statement, show_progress = TRUE, ...) {
+      state$disposition <- conn@disposition
+      data.frame(v = c(1L, 2L, 3L))
+    },
+    .package = "brickster"
+  )
+
+  out <- dbplyr::db_collect(con, "SELECT * FROM tbl")
+
+  expect_identical(state$disposition, "INLINE")
+  expect_identical(out$v, c(1L, 2L, 3L))
+})
+
 test_that("sql_query_save validates connection", {
   con <- make_test_con()