diff --git a/R/catalog-snowflake.R b/R/catalog-snowflake.R new file mode 100644 index 0000000..09a9faf --- /dev/null +++ b/R/catalog-snowflake.R @@ -0,0 +1,238 @@ +is_snowflake_connection <- function(con) { + info <- tryCatch(DBI::dbGetInfo(con), error = function(err) NULL) + name <- info$dbms.name + rlang::is_string(name) && identical(tolower(name), "snowflake") +} + +snowflake_table_registry <- function( + con, + tables = NULL, + call = rlang::caller_env() +) { + if (is.null(tables)) { + ids <- list(snowflake_current_namespace(con, call = call)) + } else { + entries <- table_entries(tables, call = call) + ids <- lapply(entries, table_entry_id, call = call) + } + + relations <- list() + validate <- list() + for (id in ids) { + type <- snowflake_id_type(id, call = call) + if (identical(type, "relation")) { + relation <- snowflake_exact_relation(con, id, call = call) + label <- table_id_label(id, call = call) + relations[[length(relations) + 1L]] <- relation + validate[[length(validate) + 1L]] <- id + next + } + relations <- c( + relations, + snowflake_list_relations(con, id, call = call) + ) + } + + labels <- vapply(relations, function(x) { + table_id_label(x$id, call = call) + }, character(1)) + duplicated_labels <- unique(labels[duplicated(labels)]) + if (length(duplicated_labels)) { + cli::cli_abort( + "{.arg tables} must not select duplicate labels: {.val {duplicated_labels}}.", + call = call + ) + } + + relation_ids <- lapply(relations, `[[`, "id") + names(relation_ids) <- labels + names(relations) <- labels + + validate_labels <- vapply(validate, table_id_label, character(1), call = call) + names(validate) <- validate_labels + + list( + labels = labels, + ids = relation_ids, + relations = relations, + validate = list(labels = validate_labels, ids = validate) + ) +} + +snowflake_current_namespace <- function(con, call = rlang::caller_env()) { + row <- tryCatch( + DBI::dbGetQuery( + con, + paste( + "SELECT CURRENT_DATABASE() AS catalog,", + "CURRENT_SCHEMA() AS schema" + ) + ), + error = function(err) { + cli::cli_abort( + "Failed to read the current Snowflake namespace.", + parent = err, + call = call + ) + } + ) + names(row) <- tolower(names(row)) + if (nrow(row) != 1L || !all(c("catalog", "schema") %in% names(row))) { + cli::cli_abort( + "Snowflake returned an invalid current namespace.", + call = call + ) + } + values <- unlist(row[1, c("catalog", "schema")], use.names = TRUE) + if ( + length(values) != 2L || + any(is.na(values)) || + any(!nzchar(values)) + ) { + cli::cli_abort( + c( + "The Snowflake connection has no current database and schema.", + "i" = "Set both on the connection or supply {.arg tables} as a {.cls DBI::Id}." + ), + call = call + ) + } + DBI::Id(catalog = values[["catalog"]], schema = values[["schema"]]) +} + +snowflake_list_relations <- function(con, namespace, call = rlang::caller_env()) { + components <- namespace@name + target <- if (identical(names(components), "catalog")) { + paste("IN DATABASE", DBI::dbQuoteIdentifier(con, namespace)) + } else { + paste("IN SCHEMA", DBI::dbQuoteIdentifier(con, namespace)) + } + rows <- tryCatch( + DBI::dbGetQuery(con, paste("SHOW OBJECTS", target)), + error = function(err) { + cli::cli_abort( + "Failed to list relations in the selected Snowflake namespace.", + parent = err, + call = call + ) + } + ) + snowflake_relations_from_show(rows) +} + +snowflake_exact_relation <- function(con, id, call = rlang::caller_env()) { + components <- id@name + namespace <- components[setdiff(names(components), "table")] + if (length(namespace) == 0L) { + namespace <- snowflake_current_namespace(con, call = call)@name + } + namespace <- do.call(DBI::Id, as.list(namespace)) + target <- paste("IN SCHEMA", DBI::dbQuoteIdentifier(con, namespace)) + rows <- tryCatch( + DBI::dbGetQuery( + con, + paste( + "SHOW OBJECTS LIKE", + DBI::dbQuoteString(con, components[["table"]]), + target + ) + ), + error = function(err) { + cli::cli_abort( + "Failed to read metadata for the selected Snowflake relation.", + parent = err, + call = call + ) + } + ) + relations <- snowflake_relations_from_show(rows) + requested_name <- components[["table"]] + is_requested <- vapply( + relations, + function(relation) { + identical(relation$id@name[["table"]], requested_name) + }, + logical(1) + ) + if (!any(is_requested)) { + return(list(id = id, kind = NULL, description = NULL)) + } + + relation <- relations[[which(is_requested)[[1]]]] + relation$id <- id + relation +} + +snowflake_relations_from_show <- function(rows) { + names(rows) <- tolower(names(rows)) + if (nrow(rows) == 0L) { + return(list()) + } + rows <- rows[toupper(rows$kind) %in% c("TABLE", "VIEW"), , drop = FALSE] + lapply(seq_len(nrow(rows)), function(i) { + description <- rows$comment[[i]] + if (is.na(description) || !nzchar(description)) { + description <- NULL + } + list( + id = DBI::Id( + catalog = rows$database_name[[i]], + schema = rows$schema_name[[i]], + table = rows$name[[i]] + ), + kind = tolower(rows$kind[[i]]), + description = description + ) + }) +} + +snowflake_describe_relation <- function(con, id, call = rlang::caller_env()) { + rows <- tryCatch( + DBI::dbGetQuery( + con, + paste("DESC TABLE", DBI::dbQuoteIdentifier(con, id)) + ), + error = function(err) { + cli::cli_abort( + "Failed to describe the selected Snowflake relation.", + parent = err, + call = call + ) + } + ) + names(rows) <- tolower(names(rows)) + rows <- rows[toupper(rows$kind) == "COLUMN", , drop = FALSE] + description <- rows$comment + description[is.na(description) | !nzchar(description)] <- NA_character_ + data.frame( + column = rows$name, + type = rows$type, + nullable = rows[["null?"]] == "Y", + description = description, + row.names = NULL + ) +} + +snowflake_id_type <- function(id, call = rlang::caller_env()) { + components <- id@name + roles <- names(components) + valid <- list( + c("catalog"), + c("schema"), + c("catalog", "schema"), + c("table"), + c("schema", "table"), + c("catalog", "schema", "table") + ) + if ( + !any(vapply(valid, identical, logical(1), roles)) || + any(is.na(components) | !nzchar(components)) + ) { + cli::cli_abort( + "Snowflake {.cls DBI::Id} entries in {.arg tables} must follow + catalog, schema, and table order without skipped or empty components.", + call = call + ) + } + if ("table" %in% roles) "relation" else "namespace" +} diff --git a/R/data-source.R b/R/data-source.R index b49a272..4e6cfc3 100644 --- a/R/data-source.R +++ b/R/data-source.R @@ -35,7 +35,10 @@ #' strings like `"schema.table"`, or `DBI::Id` objects. Defaults to every #' table returned by [DBI::dbListTables()]. Strings containing dots are #' interpreted as schema-qualified names; use `DBI::Id(table = "a.b")` for -#' literal table names containing dots. +#' literal table names containing dots. For Snowflake connections, a +#' `DBI::Id` ending in `catalog` or `schema` selects every table and view in +#' that namespace. Leaving `tables` unset selects the current Snowflake +#' schema. #' #' For a board, a named character vector of pins to read: the names become #' table names, and the values are pin names passed to [pins::pin_read()]. @@ -143,6 +146,26 @@ data_source_connection <- function( ) { span <- local_commons_span("commons_data_source_list_tables") + if (is_snowflake_connection(con)) { + table_registry <- snowflake_table_registry(con, tables, call = call) + if (length(table_registry$validate$labels)) { + check_table_ids_exist(con, table_registry$validate, call = call) + } + commons_span_set_attribute( + span, + "commons.data_source.n_tables", + length(table_registry$labels) + ) + return(new_data_source( + con, + table_registry$labels, + owned = FALSE, + table_ids = table_registry$ids, + dictionary = dictionary, + snowflake_relations = table_registry$relations + )) + } + if (is.null(tables)) { listed <- DBI::dbListTables(con) commons_span_set_attribute(span, "commons.data_source.n_tables", length(listed)) @@ -241,7 +264,8 @@ new_data_source <- function( owned, table_ids = table_ids_from_labels(tables), dictionary = NULL, - pending = NULL + pending = NULL, + snowflake_relations = NULL ) { # Disconnect only the DuckDB connection we created; a user-supplied connection # has its own owner and lifetime. @@ -263,7 +287,8 @@ new_data_source <- function( table_ids = table_ids, handle = handle, dictionary = dictionary, - pending = pending + pending = pending, + snowflake_relations = snowflake_relations ), class = "commons_data_source" ) @@ -426,7 +451,12 @@ pending_tables_in_error <- function(source, err) { )] } -source_describe <- function(source, table, n_sample = 5) { +source_describe <- function( + source, + table, + n_sample = 5, + call = rlang::caller_env() +) { id <- source$table_ids[[table]] if (is.null(id)) { cli::cli_abort(c( @@ -444,12 +474,22 @@ source_describe <- function(source, table, n_sample = 5) { n_sample ) ) - schema <- data.frame( - column = names(sample), - type = vapply(sample, function(x) class(x)[[1]], character(1)), - row.names = NULL + relation <- source$snowflake_relations[[table]] + if (is.null(source$snowflake_relations)) { + schema <- data.frame( + column = names(sample), + type = vapply(sample, function(x) class(x)[[1]], character(1)), + row.names = NULL + ) + } else { + schema <- snowflake_describe_relation(source$con, id, call = call) + } + list( + schema = schema, + sample = sample, + kind = relation$kind, + description = relation$description ) - list(schema = schema, sample = sample) } source_query <- function(source, sql) { diff --git a/R/tools.R b/R/tools.R index 6bde650..b5ad0b8 100644 --- a/R/tools.R +++ b/R/tools.R @@ -345,6 +345,10 @@ search_context_tool <- function(context, query) { describe_table_tool <- function(source, table, source_name = NULL, tracker = NULL) { d <- source_describe(source, table) entry <- source$dictionary$tables[[table]] + relation <- c( + if (!is.null(d$kind)) sprintf("Relation type: %s.", d$kind), + d$description + ) sample <- sprintf( "Sample rows:\n\n%s", @@ -352,6 +356,7 @@ describe_table_tool <- function(source, table, source_name = NULL, tracker = NUL ) if (is.null(entry)) { parts <- c( + relation, sprintf("Columns of `%s`:\n\n%s", table, df_to_markdown(d$schema)), sample ) @@ -363,6 +368,7 @@ describe_table_tool <- function(source, table, source_name = NULL, tracker = NUL dictionary_columns_text(entry$columns, live = d$schema) ) parts <- c( + relation, dictionary_entry_parts(source$dictionary, table, columns), sample ) diff --git a/man/data_source.Rd b/man/data_source.Rd index 500ba9e..9f8c240 100644 --- a/man/data_source.Rd +++ b/man/data_source.Rd @@ -18,7 +18,10 @@ For a connection, a character vector of table names, schema-qualified strings like \code{"schema.table"}, or \code{DBI::Id} objects. Defaults to every table returned by \code{\link[DBI:dbListTables]{DBI::dbListTables()}}. Strings containing dots are interpreted as schema-qualified names; use \code{DBI::Id(table = "a.b")} for -literal table names containing dots. +literal table names containing dots. For Snowflake connections, a +\code{DBI::Id} ending in \code{catalog} or \code{schema} selects every table and view in +that namespace. Leaving \code{tables} unset selects the current Snowflake +schema. For a board, a named character vector of pins to read: the names become table names, and the values are pin names passed to \code{\link[pins:pin_read]{pins::pin_read()}}.} diff --git a/tests/testthat/README.md b/tests/testthat/README.md index 2a49820..e31877c 100644 --- a/tests/testthat/README.md +++ b/tests/testthat/README.md @@ -25,6 +25,8 @@ options(commons.test.databricks = DBI::Id( )) ``` -Each test queries the current identity and namespace, then reads at most one row -from the configured table. If a backend's option is absent, its test skips. Once -enabled, connection and query failures fail the test. +The Snowflake test exercises exact table selection, schema and catalog +expansion, current-schema discovery, native column metadata, and sample rows. +The Databricks test queries the current identity and namespace, then reads at +most one row from the configured table. If a backend's option is absent, its +test skips. Once enabled, connection and query failures fail the test. diff --git a/tests/testthat/_snaps/catalog-snowflake.md b/tests/testthat/_snaps/catalog-snowflake.md new file mode 100644 index 0000000..51e1104 --- /dev/null +++ b/tests/testthat/_snaps/catalog-snowflake.md @@ -0,0 +1,9 @@ +# Snowflake current namespace requires a database and schema + + Code + snowflake_current_namespace(NULL) + Condition + Error: + ! The Snowflake connection has no current database and schema. + i Set both on the connection or supply `tables` as a . + diff --git a/tests/testthat/test-catalog-snowflake.R b/tests/testthat/test-catalog-snowflake.R new file mode 100644 index 0000000..3dceadd --- /dev/null +++ b/tests/testthat/test-catalog-snowflake.R @@ -0,0 +1,68 @@ +test_that("Snowflake SHOW results retain native relation metadata", { + rows <- data.frame( + name = c("Sales.Report", "ORDERS", "STAGE"), + database_name = c("Data.Base", "ANALYTICS", "ANALYTICS"), + schema_name = c("Odd Schema", "PUBLIC", "PUBLIC"), + kind = c("VIEW", "TABLE", "STAGE"), + comment = c("A useful view", "", "Ignored"), + stringsAsFactors = FALSE + ) + + relations <- snowflake_relations_from_show(rows) + + expect_length(relations, 2) + expect_equal(relations[[1]]$kind, "view") + expect_equal(relations[[1]]$description, "A useful view") + expect_null(relations[[2]]$description) + expect_identical( + relations[[1]]$id, + DBI::Id( + catalog = "Data.Base", + schema = "Odd Schema", + table = "Sales.Report" + ) + ) + + con <- DBI::dbConnect(duckdb::duckdb()) + withr::defer(DBI::dbDisconnect(con, shutdown = TRUE)) + expect_equal( + as.character(DBI::dbQuoteIdentifier(con, relations[[1]]$id)), + '"Data.Base"."Odd Schema"."Sales.Report"' + ) +}) + +test_that("Snowflake identifiers distinguish namespaces and relations", { + expect_equal(snowflake_id_type(DBI::Id(catalog = "DB")), "namespace") + expect_equal(snowflake_id_type(DBI::Id(schema = "PUBLIC")), "namespace") + expect_equal( + snowflake_id_type(DBI::Id(catalog = "DB", schema = "PUBLIC")), + "namespace" + ) + expect_equal(snowflake_id_type(DBI::Id(table = "ORDERS")), "relation") + expect_equal( + snowflake_id_type(DBI::Id(schema = "PUBLIC", table = "ORDERS")), + "relation" + ) + expect_error( + snowflake_id_type(DBI::Id(catalog = "DB", table = "ORDERS")), + "without skipped" + ) + expect_error( + snowflake_id_type(DBI::Id(schema = "", table = "ORDERS")), + "without skipped" + ) +}) + +test_that("Snowflake current namespace requires a database and schema", { + local_mocked_bindings( + dbGetQuery = function(...) { + data.frame(catalog = NA_character_, schema = NA_character_) + }, + .package = "DBI" + ) + + expect_snapshot( + snowflake_current_namespace(NULL), + error = TRUE + ) +}) diff --git a/tests/testthat/test-live-warehouses.R b/tests/testthat/test-live-warehouses.R index 6fa1f5a..9c5f5e7 100644 --- a/tests/testthat/test-live-warehouses.R +++ b/tests/testthat/test-live-warehouses.R @@ -1,6 +1,11 @@ -test_that("live Snowflake connection reads a configured table", { +test_that("live Snowflake discovers and describes catalog relations", { table <- warehouse_test_table("snowflake") con <- local_warehouse_connection("snowflake") + components <- table@name + skip_if_not( + all(c("catalog", "schema", "table") %in% names(components)), + "The Snowflake test table must be fully qualified" + ) session <- DBI::dbGetQuery( con, @@ -13,12 +18,64 @@ test_that("live Snowflake connection reads a configured table", { ) rows <- warehouse_read_one(con, table) names(session) <- tolower(names(session)) + label <- table_id_label(table) + + exact <- data_source(con, tables = table) + described <- source_describe(exact, label) + + namespace <- DBI::Id( + catalog = components[["catalog"]], + schema = components[["schema"]] + ) + schema_source <- data_source(con, tables = namespace) + catalog_source <- data_source( + con, + tables = DBI::Id(catalog = components[["catalog"]]) + ) + + DBI::dbExecute( + con, + paste( + "USE DATABASE", + DBI::dbQuoteIdentifier( + con, + DBI::Id(catalog = components[["catalog"]]) + ) + ) + ) + DBI::dbExecute( + con, + paste("USE SCHEMA", DBI::dbQuoteIdentifier(con, namespace)) + ) + current_source <- data_source(con) expect_equal(nrow(session), 1) expect_named(session, c("principal", "role", "catalog", "schema")) expect_true(nzchar(session$principal[[1]])) expect_s3_class(rows, "data.frame") expect_true(nrow(rows) <= 1) + expect_equal(list_tables(exact), label) + expect_identical(exact$table_ids[[label]], table) + expect_named( + described$schema, + c("column", "type", "nullable", "description") + ) + expect_true(nrow(described$sample) <= 5) + expect_equal(names(described$sample), described$schema$column) + expect_true(exact$snowflake_relations[[label]]$kind %in% c("table", "view")) + expect_true(label %in% list_tables(schema_source)) + expect_true(label %in% list_tables(catalog_source)) + expect_true(label %in% list_tables(current_source)) + expect_true(all(vapply( + schema_source$snowflake_relations, + function(x) x$kind %in% c("table", "view"), + logical(1) + ))) + + tool <- describe_table_tool(exact, label) + expect_match(tool@value, "Relation type") + expect_match(tool@value, "nullable") + expect_match(tool@value, "Sample rows") }) test_that("live Databricks connection reads a configured table", {