diff --git a/CHANGELOG.md b/CHANGELOG.md index a913c3619..e7fa625fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,14 @@ ## [Unreleased] +### Added + +- In a Positron console, a query result can be opened in the Data Explorer + from an "Open in Data Explorer" link under its table, with sorting, + filtering, profiles, export and Convert to Code. A `-- @view` comment opens + a query's result in the Data Explorer as soon as it runs. Long results show + their first 10 rows in the console, configurable with `ggsql.consoleRows` + (posit-dev/positron#16137). + ### Changed - The reader/dialect layer was refactored around a central dialect registry with diff --git a/Cargo.lock b/Cargo.lock index 3cdb9db58..6968a7332 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1791,6 +1791,7 @@ dependencies = [ "ggsql", "hex", "hmac", + "regex", "serde", "serde_json", "sha2", diff --git a/doc/get_started/tooling/positron-vscode.qmd b/doc/get_started/tooling/positron-vscode.qmd index de213c349..182d899f6 100644 --- a/doc/get_started/tooling/positron-vscode.qmd +++ b/doc/get_started/tooling/positron-vscode.qmd @@ -30,6 +30,20 @@ Executing ggsql statements containing [`VISUALISE`](../../syntax/clause/visualis ![](./screenshots/positron-plots-pane.png) +## Data Explorer + +A query without `VISUALISE` shows its result as a table in the console, followed by an "Open in Data Explorer" link. A result of more than 20 rows shows only its first 10, with a count of the rest. To show more or fewer, change the `ggsql.consoleRows` setting: a result of up to twice that many rows is shown in full. The change applies to consoles started afterwards. Clicking the link opens the result in Positron's Data Explorer, where you can sort, filter, and profile it, and use "Convert to Code" to turn your sorts and filters back into SQL. + +To open a result in the Data Explorer as soon as it runs, start the query with a `-- @view` comment. You can also give the tab a title: + +```ggsql +-- @view: Revenue by region +SELECT region, sum(revenue) AS revenue FROM sales GROUP BY region +``` + +The result opens in a Data Explorer tab, and is also shown in the console as usual. Because it is an ordinary comment, `-- @view` works the same way in a source file, in a cell, and in code sent to the console by other tools. It applies to the query that follows it, and that query cannot have a `VISUALISE` clause. Outside Positron, where there is no Data Explorer, the result is shown as a table as usual. + +The Data Explorer works on a snapshot of the result, so it doesn't change when the underlying tables do. Run the query again to see new data. ## Editor and cells diff --git a/ggsql-jupyter/CLAUDE.md b/ggsql-jupyter/CLAUDE.md index c770aef65..008388435 100644 --- a/ggsql-jupyter/CLAUDE.md +++ b/ggsql-jupyter/CLAUDE.md @@ -17,7 +17,11 @@ ggsql-jupyter/ │ ├── kernel.rs Jupyter messaging loop (ZMQ, message dispatch) │ ├── executor.rs Runs queries via ggsql::Reader, returns rendered output │ ├── connection.rs Reader lifecycle, connection-string parsing -│ ├── data_explorer.rs Positron data-explorer comm channel +│ ├── data_explorer/ Positron data-explorer comm channel +│ │ ├── mod.rs State, RPC dispatch, paging and profiles +│ │ ├── filter.rs Row filters, sort keys, schema search, Convert to Code → SQL +│ │ └── export.rs Copying a selection as CSV / TSV / HTML +│ ├── results.rs Query results retained for the data explorer │ ├── display.rs Output formatting (image bundles, SQL → HTML table) │ ├── message.rs Jupyter message structs (ZMQ frames, HMAC signing) │ └── util.rs @@ -173,7 +177,32 @@ So `RenderHints::canvas` picks a height (golden ratio, close to ggplot2's defaul ## Positron-specific bits -- `data_explorer.rs` implements Positron's data-explorer comm channel (registered query results become explorable tables). +- `data_explorer/` implements Positron's data-explorer comm channel. An explorer's `Source` says which database its table is in: the session's own, for a table opened from the Connections pane, or the private results database below. + +### The data explorer as SQL + +Every page, filter, sort, profile and export is a query; the table is never loaded whole. A few things are deliberate: + +- **Filter values arrive as strings and are checked against the column's type** before they reach the SQL: a number must parse as a finite number, a boolean as true/false, and everything else becomes a quoted literal. Filters are always ANDed (`supports_conditions: unsupported`). +- **A bad filter is left out, not fatal.** One the kernel cannot build, or one the database rejects (counted on its own to find out which), comes back with `is_valid: false` and an `error_message`; the rest still apply. +- **Profiles and exports see the filtered rows**, through `relation()`; pages and exports also follow the sort. +- **Sorting a retained result breaks ties on DuckDB's `rowid`**, its original row order, so paging is stable and clearing the sort restores the query's own `ORDER BY`. A Connections-pane table has no portable equivalent and gets no tiebreak. +- **Regular-expression search needs `Capabilities::regex`** (`regexp_matches`, DuckDB only). The results database always has it; the session's database has it when its URI is `duckdb://`. +- **Convert to Code writes SQL for the session's database**, not the private one, wrapping a retained result's original query as a derived table. A filter it cannot express there is left out with a `--` comment saying so. +- **A failed request is a JSON-RPC error**, with `plot::comm::RpcError`'s encoding, not a `null` result. + +### Opening a console result in the data explorer + +The console renders a tabular result as an HTML table and has no output toolbar to hang a button on, so the affordance is a link, and it needs no change to Positron: + +1. **A console result is retained.** `results.rs` copies it into a private in-memory DuckDB under an unguessable key, bounded by `--max-results` (default 16). Only results no explorer has open count against the bound or are evicted, so an open tab never loses its table. Querying a copy rather than the source means one SQL dialect for everything the explorer asks, no remote query re-run per page, and nothing visible in the user's session. +2. **A `stream` follows the `execute_result`**, carrying an OSC 8 hyperlink to `x-ggsql-view:/`. It cannot go in the bundle's `text/plain`: the console shows the HTML and ignores it. So that the link stays in view under a long result, a console table shows rows the way R's tibble does — all of them up to twice `--console-rows` (default 10), else the first `--console-rows` — where elsewhere it shows up to 100 (`display.rs`'s `html_rows_shown`). The option is also read from `GGSQL_CONSOLE_ROWS`, which is how the extension sets it. +3. **Positron offers a clicked hyperlink to `openResource`** on the last active console session of each language. The extension's session wrapper (`ggsql-vscode/src/resultLinks.ts`) handles the scheme and runs `-- @view-result ` silently, in the session the link names. +4. **`-- @view-result` opens a `positron.dataExplorer` comm with `inline_only: false`**, which Positron opens as an editor tab immediately. A result that already has an explorer open is not opened again: the kernel sends `open_editor` on the UI comm with that explorer's editor URI (`positron-data-explorer:positron-data-explorer-`), which focuses the tab — what Positron's Variables pane does for a variable already in view. A key that is no longer retained is a cell error, which surfaces as a warning from the extension. + +Only `SessionKind::PositronConsole` gets this: anywhere else the escape sequence would be stray text. Notebooks have a better route — the inline data explorer — that this does not yet use. + +A query under the user-facing `-- @view` directive (see [the Positron docs](../doc/get_started/tooling/positron-vscode.qmd)) takes the same path from step 1, minus the click: it is retained in any Positron session, silent executions included, and opened at once. The table still shows in the cell, and a console still gets the link under it, to reopen a closed tab; a silent execution shows nothing but the tab. A `Standalone` session has no data explorer to open, so it gets the inline table and a `stderr` note instead. - The companion VS Code extension (`ggsql-vscode/`) ships a copy of this binary and also discovers installed ones via the `ggsql.kernelPath` setting, the Jupyter kernelspec directories, the native install locations, or `PATH`. It runs each one with `--version` to name the runtime it registers, so **keep `--version` working**: a kernel that does not answer it is still offered, but without a version in the picker. See [Finding the kernel](../ggsql-vscode/CLAUDE.md#finding-the-kernel). ## Build & install diff --git a/ggsql-jupyter/Cargo.toml b/ggsql-jupyter/Cargo.toml index 47590aef3..1774a883d 100644 --- a/ggsql-jupyter/Cargo.toml +++ b/ggsql-jupyter/Cargo.toml @@ -33,12 +33,15 @@ arrow = { workspace = true } tokio = { workspace = true, features = ["sync"] } # CLI -clap = { workspace = true } +clap = { workspace = true, features = ["env"] } # Serialization serde = { workspace = true } serde_json = { workspace = true } +# The data explorer's regular-expression column search +regex = { workspace = true } + # Error handling anyhow = { workspace = true } thiserror = { workspace = true } diff --git a/ggsql-jupyter/src/data_explorer/export.rs b/ggsql-jupyter/src/data_explorer/export.rs new file mode 100644 index 000000000..135664baf --- /dev/null +++ b/ggsql-jupyter/src/data_explorer/export.rs @@ -0,0 +1,108 @@ +//! A data explorer selection as CSV, TSV or HTML, for the clipboard. + +use crate::display::escape_html; + +/// Format `rows` for export. A `header` is written first when given; a null +/// cell is empty. +pub fn format( + header: Option<&[String]>, + rows: &[Vec], + format: &str, +) -> Result { + let all = header + .into_iter() + .map(<[String]>::to_vec) + .chain(rows.iter().cloned()); + match format { + "csv" => Ok(delimited(all, ",", csv_field)), + "tsv" => Ok(delimited(all, "\t", tsv_field)), + "html" => Ok(html(header, rows)), + other => Err(format!("Unsupported export format {other:?}")), + } +} + +fn delimited( + rows: impl Iterator>, + sep: &str, + field: fn(&str) -> String, +) -> String { + let lines: Vec = rows + .map(|row| row.iter().map(|v| field(v)).collect::>().join(sep)) + .collect(); + lines.join("\n") +} + +/// Quoted only when it has to be, with embedded quotes doubled (RFC 4180). +fn csv_field(v: &str) -> String { + if v.contains([',', '"', '\n', '\r']) { + format!("\"{}\"", v.replace('"', "\"\"")) + } else { + v.to_string() + } +} + +/// TSV has no quoting, so the characters that would break a row become spaces. +fn tsv_field(v: &str) -> String { + v.replace(['\t', '\n', '\r'], " ") +} + +fn html(header: Option<&[String]>, rows: &[Vec]) -> String { + let mut out = String::from("\n"); + if let Some(header) = header { + out.push_str(""); + for h in header { + out.push_str(&format!("", escape_html(h))); + } + out.push_str("\n"); + } + out.push_str("\n"); + for row in rows { + out.push_str(""); + for v in row { + out.push_str(&format!("", escape_html(v))); + } + out.push_str("\n"); + } + out.push_str("\n
{}
{}
"); + out +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rows() -> Vec> { + vec![ + vec!["1".into(), "a, \"b\"".into()], + vec!["2".into(), "tab\there".into()], + ] + } + + #[test] + fn test_csv_quotes_only_what_needs_it() { + let header = ["n".to_string(), "s".to_string()]; + assert_eq!( + format(Some(&header), &rows(), "csv").unwrap(), + "n,s\n1,\"a, \"\"b\"\"\"\n2,tab\there" + ); + } + + #[test] + fn test_tsv_flattens_tabs() { + assert_eq!( + format(None, &rows(), "tsv").unwrap(), + "1\ta, \"b\"\n2\ttab here" + ); + } + + #[test] + fn test_html_escapes() { + let out = format(None, &[vec!["".into()]], "html").unwrap(); + assert!(out.contains("<b>"), "{out}"); + } + + #[test] + fn test_unknown_format_is_an_error() { + assert!(format(None, &rows(), "xlsx").is_err()); + } +} diff --git a/ggsql-jupyter/src/data_explorer/filter.rs b/ggsql-jupyter/src/data_explorer/filter.rs new file mode 100644 index 000000000..6e9d3e795 --- /dev/null +++ b/ggsql-jupyter/src/data_explorer/filter.rs @@ -0,0 +1,496 @@ +//! Row filters, sort keys, schema search and Convert to Code, as SQL. +//! +//! Everything here builds from the protocol's JSON and the cached schema, with +//! no database, so it is tested as plain string assembly. + +use super::ColumnInfo; +use ggsql::naming::{quote_ident, quote_literal}; +use ggsql::reader::SqlDialect; +use serde_json::Value; + +/// What a database can do beyond the SQL every backend shares. +#[derive(Debug, Clone, Copy)] +pub struct Capabilities { + /// `regexp_matches`, for the `regex_match` search type. DuckDB only. + pub regex: bool, +} + +/// The SQL predicate for one row filter. +/// +/// Every value arrives as a string, so each is checked against its column's +/// type before it reaches the SQL: a number must parse as one, and anything +/// else becomes a quoted literal. +pub fn predicate( + filter: &Value, + columns: &[ColumnInfo], + dialect: &dyn SqlDialect, + caps: Capabilities, +) -> Result { + let index = filter["column_schema"]["column_index"] + .as_u64() + .ok_or("The filter names no column")? as usize; + let column = columns + .get(index) + .ok_or_else(|| format!("There is no column {index}"))?; + let c = quote_ident(&column.name); + let params = &filter["params"]; + let lit = |v: &Value| -> Result { + let value = v.as_str().ok_or("The filter is missing a value")?; + literal(value, &column.type_display, dialect) + }; + + match filter["filter_type"].as_str().unwrap_or("") { + "compare" => { + let op = match params["op"].as_str().unwrap_or("") { + "=" => "=", + "!=" => "<>", + "<" => "<", + "<=" => "<=", + ">" => ">", + ">=" => ">=", + other => return Err(format!("Unknown comparison {other:?}")), + }; + Ok(format!("{c} {op} {}", lit(¶ms["value"])?)) + } + kind @ ("between" | "not_between") => { + let between = format!( + "{c} BETWEEN {} AND {}", + lit(¶ms["left_value"])?, + lit(¶ms["right_value"])? + ); + Ok(if kind == "between" { + between + } else { + format!("NOT ({between})") + }) + } + "is_null" => Ok(format!("{c} IS NULL")), + "not_null" => Ok(format!("{c} IS NOT NULL")), + "is_empty" => Ok(format!("{c} = ''")), + "not_empty" => Ok(format!("{c} <> ''")), + "is_true" => Ok(format!("{c} = {}", dialect.sql_boolean_literal(true))), + "is_false" => Ok(format!("{c} = {}", dialect.sql_boolean_literal(false))), + "search" => search(&c, params, caps), + "set_membership" => { + let inclusive = params["inclusive"].as_bool().unwrap_or(true); + let values = params["values"].as_array().cloned().unwrap_or_default(); + if values.is_empty() { + // Nothing selected: keep nothing, or exclude nothing. + return Ok(if inclusive { "1 = 0" } else { "1 = 1" }.to_string()); + } + let list = values.iter().map(lit).collect::, _>>()?; + let not = if inclusive { "" } else { "NOT " }; + Ok(format!("{c} {not}IN ({})", list.join(", "))) + } + other => Err(format!("Unsupported filter type {other:?}")), + } +} + +/// A filter value as a SQL literal of its column's type. +fn literal(value: &str, display: &str, dialect: &dyn SqlDialect) -> Result { + match display { + "integer" | "floating" | "decimal" => { + let v = value.trim(); + // Only a finite number passes, so nothing else reaches the SQL + // unquoted. + match v.parse::() { + Ok(n) if n.is_finite() => Ok(v.to_string()), + _ => Err(format!("{value:?} is not a number")), + } + } + "boolean" => match value.trim().to_ascii_lowercase().as_str() { + "true" | "1" => Ok(dialect.sql_boolean_literal(true)), + "false" | "0" => Ok(dialect.sql_boolean_literal(false)), + _ => Err(format!("{value:?} is not true or false")), + }, + // Dates and times too: every backend compares a string literal against + // a temporal column by casting it. + _ => Ok(quote_literal(value)), + } +} + +/// The predicate for a text search. +fn search(c: &str, params: &Value, caps: Capabilities) -> Result { + let term = params["term"].as_str().ok_or("The search has no term")?; + let case_sensitive = params["case_sensitive"].as_bool().unwrap_or(false); + let kind = params["search_type"].as_str().unwrap_or(""); + + if kind == "regex_match" { + if !caps.regex { + return Err("This database does not support regular-expression search".into()); + } + let flags = if case_sensitive { "" } else { ", 'i'" }; + return Ok(format!( + "regexp_matches({c}, {}{flags})", + quote_literal(term) + )); + } + + let escaped = escape_like(term); + let (pattern, negate) = match kind { + "contains" => (format!("%{escaped}%"), false), + "not_contains" => (format!("%{escaped}%"), true), + "starts_with" => (format!("{escaped}%"), false), + "ends_with" => (format!("%{escaped}"), false), + other => return Err(format!("Unknown search type {other:?}")), + }; + let (lhs, rhs) = if case_sensitive { + (c.to_string(), quote_literal(&pattern)) + } else { + ( + format!("LOWER({c})"), + format!("LOWER({})", quote_literal(&pattern)), + ) + }; + let not = if negate { "NOT " } else { "" }; + Ok(format!("{lhs} {not}LIKE {rhs} ESCAPE '\\'")) +} + +/// Escape `LIKE`'s wildcards, so a search term matches literally. +fn escape_like(term: &str) -> String { + term.replace('\\', "\\\\") + .replace('%', "\\%") + .replace('_', "\\_") +} + +/// The `ORDER BY` clause for the sort keys, with a leading space, or nothing +/// when unsorted. +/// +/// `tiebreak` orders rows the keys leave equal, so paging through a sorted +/// view is stable; without it, rows with equal keys may move between pages. +pub fn order_by( + sort_keys: &[Value], + columns: &[ColumnInfo], + tiebreak: Option<&str>, +) -> Result { + let mut parts = Vec::new(); + for key in sort_keys { + let index = key["column_index"] + .as_u64() + .ok_or("A sort key names no column")? as usize; + let column = columns + .get(index) + .ok_or_else(|| format!("There is no column {index}"))?; + let direction = if key["ascending"].as_bool().unwrap_or(true) { + "ASC" + } else { + "DESC" + }; + parts.push(format!("{} {direction}", quote_ident(&column.name))); + } + if parts.is_empty() { + return Ok(String::new()); + } + parts.extend(tiebreak.map(str::to_string)); + Ok(format!(" ORDER BY {}", parts.join(", "))) +} + +/// The indices of the columns matching every column filter, in `sort_order`. +pub fn search_schema(params: &Value, columns: &[ColumnInfo]) -> Vec { + let filters = params["filters"].as_array().cloned().unwrap_or_default(); + let mut matches: Vec = (0..columns.len()) + .filter(|&i| filters.iter().all(|f| column_matches(f, &columns[i]))) + .collect(); + + match params["sort_order"].as_str().unwrap_or("original") { + "ascending_name" => matches.sort_by_key(|&i| columns[i].name.to_lowercase()), + "descending_name" => { + matches.sort_by_key(|&i| std::cmp::Reverse(columns[i].name.to_lowercase())) + } + "ascending_type" => matches.sort_by_key(|&i| columns[i].type_display.clone()), + "descending_type" => { + matches.sort_by_key(|&i| std::cmp::Reverse(columns[i].type_display.clone())) + } + _ => {} + } + matches +} + +/// Whether a column passes one column filter. +fn column_matches(filter: &Value, column: &ColumnInfo) -> bool { + let params = &filter["params"]; + match filter["filter_type"].as_str().unwrap_or("") { + "text_search" => { + let term = params["term"].as_str().unwrap_or(""); + let case_sensitive = params["case_sensitive"].as_bool().unwrap_or(false); + let (name, term) = if case_sensitive { + (column.name.clone(), term.to_string()) + } else { + (column.name.to_lowercase(), term.to_lowercase()) + }; + match params["search_type"].as_str().unwrap_or("contains") { + "contains" => name.contains(&term), + "not_contains" => !name.contains(&term), + "starts_with" => name.starts_with(&term), + "ends_with" => name.ends_with(&term), + "regex_match" => regex::RegexBuilder::new(&term) + .case_insensitive(!case_sensitive) + .build() + .is_ok_and(|re| re.is_match(&column.name)), + _ => true, + } + } + "match_data_types" => params["display_types"] + .as_array() + .is_some_and(|types| types.iter().any(|t| t == column.type_display.as_str())), + _ => true, + } +} + +/// SQL reproducing the explorer's view: `from` filtered and sorted. +/// +/// A filter that cannot be expressed is left out with a comment saying so, +/// rather than failing the whole conversion. +pub fn convert_to_code( + params: &Value, + from: &str, + columns: &[ColumnInfo], + dialect: &dyn SqlDialect, + caps: Capabilities, +) -> Vec { + let mut lines = Vec::new(); + let mut predicates = Vec::new(); + for filter in params["row_filters"].as_array().into_iter().flatten() { + if filter["is_valid"] == Value::Bool(false) { + continue; + } + match predicate(filter, columns, dialect, caps) { + Ok(p) => predicates.push(p), + Err(e) => lines.push(format!("-- Filter left out: {e}")), + } + } + + lines.push("SELECT *".to_string()); + lines.extend(format!("FROM {from}").lines().map(str::to_string)); + for (i, p) in predicates.iter().enumerate() { + let keyword = if i == 0 { "WHERE" } else { " AND" }; + lines.push(format!("{keyword} {p}")); + } + let sort_keys = params["sort_keys"].as_array().cloned().unwrap_or_default(); + match order_by(&sort_keys, columns, None) { + Ok(order) if !order.is_empty() => lines.push(order.trim_start().to_string()), + Ok(_) => {} + Err(e) => lines.push(format!("-- Sort left out: {e}")), + } + lines +} + +/// A query as a derived table, for wrapping in the explorer's filters. +pub fn subquery(query: &str) -> String { + let query = query.trim().trim_end_matches(';').trim_end(); + let body: Vec = query.lines().map(|l| format!(" {l}")).collect(); + format!("(\n{}\n) AS {}", body.join("\n"), quote_ident("result")) +} + +#[cfg(test)] +mod tests { + use super::*; + use ggsql::reader::AnsiDialect; + use serde_json::json; + + const CAPS: Capabilities = Capabilities { regex: true }; + + fn columns() -> Vec { + [("n", "integer"), ("name", "string"), ("ok", "boolean")] + .iter() + .map(|(name, display)| ColumnInfo { + name: name.to_string(), + type_name: String::new(), + type_display: display.to_string(), + }) + .collect() + } + + fn filter(column_index: usize, filter_type: &str, params: Value) -> Value { + json!({ + "filter_id": "f", + "filter_type": filter_type, + "column_schema": {"column_index": column_index}, + "condition": "and", + "params": params + }) + } + + fn sql(f: &Value) -> Result { + predicate(f, &columns(), &AnsiDialect, CAPS) + } + + #[test] + fn test_compare_maps_not_equal_and_checks_numbers() { + assert_eq!( + sql(&filter(0, "compare", json!({"op": "!=", "value": "3"}))).unwrap(), + "\"n\" <> 3" + ); + assert!(sql(&filter( + 0, + "compare", + json!({"op": "=", "value": "1; DROP"}) + )) + .is_err()); + assert!(sql(&filter(0, "compare", json!({"op": "=", "value": "inf"}))).is_err()); + } + + #[test] + fn test_string_values_are_quoted() { + assert_eq!( + sql(&filter( + 1, + "compare", + json!({"op": "=", "value": "O'Brien"}) + )) + .unwrap(), + "\"name\" = 'O''Brien'" + ); + } + + #[test] + fn test_not_between_negates_between() { + assert_eq!( + sql(&filter( + 0, + "not_between", + json!({"left_value": "1", "right_value": "5"}) + )) + .unwrap(), + "NOT (\"n\" BETWEEN 1 AND 5)" + ); + } + + #[test] + fn test_search_escapes_wildcards_and_folds_case() { + assert_eq!( + sql(&filter( + 1, + "search", + json!({"search_type": "contains", "term": "50%_", "case_sensitive": false}) + )) + .unwrap(), + "LOWER(\"name\") LIKE LOWER('%50\\%\\_%') ESCAPE '\\'" + ); + assert_eq!( + sql(&filter( + 1, + "search", + json!({"search_type": "starts_with", "term": "a", "case_sensitive": true}) + )) + .unwrap(), + "\"name\" LIKE 'a%' ESCAPE '\\'" + ); + } + + #[test] + fn test_regex_needs_the_capability() { + let f = filter( + 1, + "search", + json!({"search_type": "regex_match", "term": "^a", "case_sensitive": false}), + ); + assert_eq!(sql(&f).unwrap(), "regexp_matches(\"name\", '^a', 'i')"); + assert!(predicate(&f, &columns(), &AnsiDialect, Capabilities { regex: false }).is_err()); + } + + #[test] + fn test_set_membership() { + assert_eq!( + sql(&filter( + 1, + "set_membership", + json!({"values": ["a", "b"], "inclusive": false}) + )) + .unwrap(), + "\"name\" NOT IN ('a', 'b')" + ); + assert_eq!( + sql(&filter( + 1, + "set_membership", + json!({"values": [], "inclusive": true}) + )) + .unwrap(), + "1 = 0" + ); + } + + #[test] + fn test_booleans_use_the_dialect() { + assert_eq!( + sql(&filter(2, "is_true", json!(null))).unwrap(), + format!("\"ok\" = {}", AnsiDialect.sql_boolean_literal(true)) + ); + } + + #[test] + fn test_order_by_appends_the_tiebreak_only_when_sorted() { + let keys = [json!({"column_index": 1, "ascending": false})]; + assert_eq!( + order_by(&keys, &columns(), Some("rowid")).unwrap(), + " ORDER BY \"name\" DESC, rowid" + ); + assert_eq!(order_by(&[], &columns(), Some("rowid")).unwrap(), ""); + } + + #[test] + fn test_search_schema_filters_and_sorts() { + let params = json!({ + "filters": [{"filter_type": "text_search", + "params": {"search_type": "contains", "term": "N", "case_sensitive": false}}], + "sort_order": "descending_name" + }); + assert_eq!(search_schema(¶ms, &columns()), vec![1, 0]); + + let params = json!({ + "filters": [{"filter_type": "match_data_types", + "params": {"display_types": ["boolean"]}}], + "sort_order": "original" + }); + assert_eq!(search_schema(¶ms, &columns()), vec![2]); + } + + #[test] + fn test_convert_to_code_wraps_the_query() { + let params = json!({ + "row_filters": [ + filter(0, "compare", json!({"op": ">", "value": "2"})), + filter(1, "is_null", json!(null)), + ], + "sort_keys": [{"column_index": 0, "ascending": true}], + "column_filters": [], + "code_syntax_name": {"code_syntax_name": "SQL"} + }); + let from = subquery("SELECT *\nFROM t;"); + assert_eq!( + convert_to_code(¶ms, &from, &columns(), &AnsiDialect, CAPS), + vec![ + "SELECT *", + "FROM (", + " SELECT *", + " FROM t", + ") AS \"result\"", + "WHERE \"n\" > 2", + " AND \"name\" IS NULL", + "ORDER BY \"n\" ASC", + ] + ); + } + + #[test] + fn test_convert_to_code_leaves_out_what_it_cannot_express() { + let params = json!({ + "row_filters": [filter( + 1, + "search", + json!({"search_type": "regex_match", "term": "x", "case_sensitive": true}) + )], + "sort_keys": [] + }); + let lines = convert_to_code( + ¶ms, + "\"t\"", + &columns(), + &AnsiDialect, + Capabilities { regex: false }, + ); + assert!(lines[0].starts_with("-- Filter left out"), "{lines:?}"); + assert_eq!(&lines[1..], ["SELECT *", "FROM \"t\""]); + } +} diff --git a/ggsql-jupyter/src/data_explorer.rs b/ggsql-jupyter/src/data_explorer/mod.rs similarity index 55% rename from ggsql-jupyter/src/data_explorer.rs rename to ggsql-jupyter/src/data_explorer/mod.rs index 9ddc137a8..0b5d97fe4 100644 --- a/ggsql-jupyter/src/data_explorer.rs +++ b/ggsql-jupyter/src/data_explorer/mod.rs @@ -1,16 +1,24 @@ //! Data explorer backend for the Positron data viewer. //! -//! Implements the `positron.dataExplorer` comm protocol, providing SQL-backed -//! paginated data access. +//! Implements the `positron.dataExplorer` comm protocol over SQL: every page, +//! filter, sort and profile is a query against the explorer's table, which is +//! never loaded into memory. +mod export; +mod filter; + +pub use filter::Capabilities; + +use crate::results::{self, StoredResult}; use ggsql::reader::Reader; use ggsql::sql::FromItem; +use ggsql::DataFrame; use serde_json::{json, Value}; /// Result of handling an RPC call. pub struct RpcResponse { - /// The JSON-RPC result to send as the reply. - pub result: Value, + /// The JSON-RPC result to send as the reply, or the message of an error. + pub result: Result, /// An optional event to send on iopub (e.g. `return_column_profiles`). pub event: Option, } @@ -24,6 +32,11 @@ pub struct RpcEvent { impl RpcResponse { /// Create a simple reply with no async event. pub fn reply(result: Value) -> Self { + Self::answer(Ok(result)) + } + + /// Reply with a result or an error, and no async event. + pub fn answer(result: Result) -> Self { Self { result, event: None, @@ -31,6 +44,13 @@ impl RpcResponse { } } +/// The session's database, which Convert to Code writes SQL for. For a +/// retained result that is not the database the explorer queries. +pub struct Session<'a> { + pub reader: &'a dyn Reader, + pub caps: Capabilities, +} + /// Cached column metadata for a table. #[derive(Debug, Clone)] pub struct ColumnInfo { @@ -41,24 +61,49 @@ pub struct ColumnInfo { pub type_display: String, } +/// Which database an explorer's table lives in. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Source { + /// A table in the session's own database, opened from the Connections pane. + Session, + /// A retained query result in the private results database, by key. + Result(String), +} + +/// Which rows of the filtered, sorted view to read. +enum Rows { + All, + /// `first..=last`. + Range(usize, usize), +} + /// State for one open data explorer comm. pub struct DataExplorerState { + source: Source, /// Fully qualified and quoted table path, e.g. `"memory"."main"."users"`. table_path: String, /// Display title shown in the data viewer tab. title: String, + /// The query a retained result came from, which Convert to Code wraps. + query: Option, + /// What the database holding the table can do. + caps: Capabilities, /// Cached column schemas. columns: Vec, /// Cached total row count. num_rows: usize, + /// The row filters as last set, each marked valid or not, for `get_state`. + row_filters: Vec, + sort_keys: Vec, + /// The valid row filters ANDed together; `None` when unfiltered. + predicate: Option, + /// Rows passing the filters. + selected_num_rows: usize, } impl DataExplorerState { /// Open a data explorer for a table at the given connection path. - /// - /// Runs `SELECT COUNT(*)` and a column metadata query to cache schema - /// information. Does **not** load the full table into memory. - pub fn open(reader: &dyn Reader, path: &[String]) -> Result { + pub fn open(reader: &dyn Reader, path: &[String], caps: Capabilities) -> Result { let full = crate::connection::resolve_path(reader, path); if full.len() < 3 { return Err(format!( @@ -66,11 +111,43 @@ impl DataExplorerState { full.len() )); } + let (catalog, schema, table) = (&full[0], &full[1], &full[2]); + let mut state = Self::open_table(reader, catalog, schema, table, table)?; + state.caps = caps; + Ok(state) + } + + /// Open a data explorer for a retained query result. `reader` is the + /// private results database. + pub fn open_result(reader: &dyn Reader, result: &StoredResult) -> Result { + let mut state = Self::open_table( + reader, + results::CATALOG, + results::SCHEMA, + &result.table, + &result.title, + )?; + state.source = Source::Result(result.key.clone()); + state.query = Some(result.query.clone()); + // The results database is DuckDB. + state.caps = Capabilities { regex: true }; + Ok(state) + } - let catalog = &full[0]; - let schema = &full[1]; - let table = &full[2]; + /// Which database this explorer's table lives in. + pub fn source(&self) -> &Source { + &self.source + } + /// Runs `SELECT COUNT(*)` and a column metadata query to cache schema + /// information. Does **not** load the full table into memory. + fn open_table( + reader: &dyn Reader, + catalog: &str, + schema: &str, + table: &str, + title: &str, + ) -> Result { let table_path = [catalog, schema, table] .iter() .filter(|s| !s.is_empty()) @@ -78,23 +155,8 @@ impl DataExplorerState { .collect::>() .join("."); - // Get row count - let count_sql = format!("SELECT COUNT(*) AS \"n\" FROM {}", table_path); - let count_df = reader - .execute_sql(&count_sql) + let num_rows = count_rows(reader, &table_path, None) .map_err(|e| format!("Failed to count rows: {}", e))?; - let num_rows = count_df - .column("n") - .ok() - .and_then(|col| { - if col.is_empty() { - None - } else { - let s = ggsql::array_util::value_to_string(col, 0); - s.parse::().ok() - } - }) - .unwrap_or(0); let column_infos = reader .list_columns(catalog, schema, table) @@ -112,47 +174,79 @@ impl DataExplorerState { } Ok(Self { + source: Source::Session, table_path, - title: table.to_string(), + title: title.to_string(), + query: None, + caps: Capabilities { regex: false }, columns, num_rows, + row_filters: Vec::new(), + sort_keys: Vec::new(), + predicate: None, + selected_num_rows: num_rows, }) } /// Dispatch a JSON-RPC method call. /// - /// Returns the RPC result and an optional async event to send on iopub - /// (used by `get_column_profiles` to deliver results asynchronously). - pub fn handle_rpc(&self, method: &str, params: &Value, reader: &dyn Reader) -> RpcResponse { + /// `reader` holds the explorer's table. Returns the RPC result and an + /// optional async event to send on iopub (used by `get_column_profiles` to + /// deliver results asynchronously). + pub fn handle_rpc( + &mut self, + method: &str, + params: &Value, + reader: &dyn Reader, + session: Session, + ) -> RpcResponse { match method { "get_state" => RpcResponse::reply(self.get_state()), "get_schema" => RpcResponse::reply(self.get_schema(params)), - "get_data_values" => RpcResponse::reply(self.get_data_values(params, reader)), + "get_data_values" => RpcResponse::answer(self.get_data_values(params, reader)), "get_column_profiles" => self.get_column_profiles(params, reader), - // TODO: Implement filters, sorting, and searching. - "set_row_filters" => { - // Stub: accept but ignore filters, return current shape - RpcResponse::reply(json!({ - "selected_num_rows": self.num_rows, - "had_errors": false - })) - } - "set_sort_columns" | "set_column_filters" | "search_schema" => { - RpcResponse::reply(json!(null)) + "set_row_filters" => RpcResponse::answer(self.set_row_filters(params, reader)), + "set_sort_columns" => RpcResponse::answer(self.set_sort_columns(params)), + "search_schema" => RpcResponse::reply(json!({ + "matches": filter::search_schema(params, &self.columns) + })), + "export_data_selection" => { + RpcResponse::answer(self.export_data_selection(params, reader)) } + "suggest_code_syntax" => RpcResponse::reply(json!({"code_syntax_name": "SQL"})), + "convert_to_code" => RpcResponse::reply(self.convert_to_code(params, &session)), + // Unsupported, and advertised as such in `get_state`. + "set_column_filters" => RpcResponse::reply(json!(null)), _ => { tracing::warn!("Unhandled data explorer method: {}", method); - RpcResponse::reply(json!(null)) + RpcResponse::answer(Err(format!("Unknown method {method:?}"))) } } } fn get_state(&self) -> Value { let num_columns = self.columns.len(); + let row_filter_types: Vec = [ + "between", + "compare", + "is_empty", + "is_false", + "is_null", + "is_true", + "not_between", + "not_empty", + "not_null", + "search", + "set_membership", + ] + .iter() + .map(|t| json!({"row_filter_type": t, "support_status": "supported"})) + .collect(); + json!({ "display_name": self.title, "table_shape": { - "num_rows": self.num_rows, + "num_rows": self.selected_num_rows, "num_columns": num_columns }, "table_unfiltered_shape": { @@ -161,21 +255,25 @@ impl DataExplorerState { }, "has_row_labels": false, "column_filters": [], - "row_filters": [], - "sort_keys": [], + "row_filters": self.row_filters, + "sort_keys": self.sort_keys, "supported_features": { "search_schema": { - "support_status": "unsupported", - "supported_types": [] + "support_status": "supported", + "supported_types": [ + {"column_filter_type": "text_search", "support_status": "supported"}, + {"column_filter_type": "match_data_types", "support_status": "supported"} + ] }, "set_column_filters": { "support_status": "unsupported", "supported_types": [] }, "set_row_filters": { - "support_status": "unsupported", + "support_status": "supported", + // Filters are always ANDed. "supports_conditions": "unsupported", - "supported_types": [] + "supported_types": row_filter_types }, "get_column_profiles": { "support_status": "supported", @@ -187,14 +285,15 @@ impl DataExplorerState { ] }, "set_sort_columns": { - "support_status": "unsupported" + "support_status": "supported" }, "export_data_selection": { - "support_status": "unsupported", - "supported_formats": [] + "support_status": "supported", + "supported_formats": ["csv", "tsv", "html"] }, "convert_to_code": { - "support_status": "unsupported" + "support_status": "supported", + "code_syntaxes": [{"code_syntax_name": "SQL"}] } } }) @@ -228,94 +327,266 @@ impl DataExplorerState { json!({ "columns": columns }) } - fn get_data_values(&self, params: &Value, reader: &dyn Reader) -> Value { - let selections = match params.get("columns").and_then(|v| v.as_array()) { - Some(arr) => arr, - None => return json!({ "columns": [] }), - }; + /// The filtered table as a `FROM` item, for the queries that aggregate it. + fn relation(&self) -> String { + match &self.predicate { + None => self.table_path.clone(), + Some(p) => format!( + "(SELECT * FROM {} WHERE {}) AS \"__ggsql_filtered__\"", + self.table_path, p + ), + } + } - // Determine the row range from the first selection's spec - let (first_index, last_index) = selections - .first() - .and_then(|sel| sel.get("spec")) - .map(|spec| { - let first = spec - .get("first_index") - .and_then(|v| v.as_u64()) - .unwrap_or(0) as usize; - let last = spec.get("last_index").and_then(|v| v.as_u64()).unwrap_or(0) as usize; - (first, last) + /// Read `rows` of the filtered, sorted view, for the columns `cols`. + fn fetch(&self, reader: &dyn Reader, cols: &[usize], rows: Rows) -> Result { + let names = cols + .iter() + .map(|&i| { + self.columns + .get(i) + .map(|c| ggsql::naming::quote_ident(&c.name)) + .ok_or_else(|| format!("There is no column {i}")) }) - .unwrap_or((0, 0)); + .collect::, _>>()?; + if names.is_empty() { + return Ok(DataFrame::empty()); + } - let limit = last_index.saturating_sub(first_index) + 1; + let filter = self + .predicate + .as_ref() + .map(|p| format!(" WHERE {p}")) + .unwrap_or_default(); + // A retained result is a DuckDB table, whose `rowid` is its original + // row order; that keeps paging through equal sort keys stable. + let tiebreak = matches!(self.source, Source::Result(_)).then_some("rowid"); + let order = filter::order_by(&self.sort_keys, &self.columns, tiebreak)?; + let limit = match rows { + Rows::All => String::new(), + Rows::Range(first, last) => { + format!(" LIMIT {} OFFSET {}", last.saturating_sub(first) + 1, first) + } + }; + + let sql = format!( + "SELECT {} FROM {}{filter}{order}{limit}", + names.join(", "), + self.table_path, + ); + reader.execute_sql(&sql).map_err(|e| e.to_string()) + } + + fn get_data_values(&self, params: &Value, reader: &dyn Reader) -> Result { + let selections = params["columns"].as_array().cloned().unwrap_or_default(); + if selections.is_empty() { + return Ok(json!({ "columns": [] })); + } - // Collect requested column indices - let col_indices: Vec = selections + // Each column asks for a range or a list of rows. One query covers the + // span of them all, and each column then takes the rows it asked for. + let wanted: Vec> = selections .iter() - .filter_map(|sel| { - sel.get("column_index") - .and_then(|v| v.as_u64()) - .map(|n| n as usize) + .map(|sel| { + let spec = &sel["spec"]; + match spec["indices"].as_array() { + Some(indices) => indices + .iter() + .filter_map(|v| v.as_u64().map(|n| n as usize)) + .collect(), + None => { + let first = spec["first_index"].as_u64().unwrap_or(0) as usize; + let last = spec["last_index"].as_u64().unwrap_or(0) as usize; + (first..=last).collect() + } + } }) .collect(); + let first = wanted.iter().flatten().min().copied().unwrap_or(0); + let last = wanted.iter().flatten().max().copied().unwrap_or(0); - // Build column list for SELECT - let col_names: Vec = col_indices + let cols: Vec = selections .iter() - .filter_map(|&idx| { - self.columns - .get(idx) - .map(|col| ggsql::naming::quote_ident(&col.name)) + .map(|sel| sel["column_index"].as_u64().unwrap_or(0) as usize) + .collect(); + let df = self.fetch(reader, &cols, Rows::Range(first, last))?; + + let columns: Vec> = wanted + .iter() + .enumerate() + .map(|(i, rows)| { + let col = &df.get_columns()[i]; + rows.iter() + .map(|&r| r - first) + .filter(|&r| r < df.height()) + .map(|r| cell_value(col, r)) + .collect() }) .collect(); - if col_names.is_empty() { - return json!({ "columns": [] }); - } + Ok(json!({ "columns": columns })) + } - let sql = format!( - "SELECT {} FROM {} LIMIT {} OFFSET {}", - col_names.join(", "), - self.table_path, - limit, - first_index, - ); + /// Apply row filters, returning how many rows pass. + /// + /// A filter that cannot be built, or that the database rejects, is marked + /// invalid with the reason and left out rather than failing the rest. + fn set_row_filters(&mut self, params: &Value, reader: &dyn Reader) -> Result { + let mut filters = params["filters"].as_array().cloned().unwrap_or_default(); + let mut had_errors = false; + let mut invalidate = |filter: &mut Value, message: String| { + filter["is_valid"] = json!(false); + filter["error_message"] = json!(message); + had_errors = true; + }; - let df = match reader.execute_sql(&sql) { - Ok(df) => df, - Err(e) => { - tracing::error!("get_data_values query failed: {}", e); - let empty: Vec> = col_indices.iter().map(|_| vec![]).collect(); - return json!({ "columns": empty }); + let mut built = Vec::new(); + for (i, f) in filters.iter_mut().enumerate() { + match filter::predicate(f, &self.columns, reader.dialect(), self.caps) { + Ok(p) => built.push((i, p)), + Err(e) => invalidate(f, e), } + } + + let combine = |built: &[(usize, String)]| { + (!built.is_empty()).then(|| { + built + .iter() + .map(|(_, p)| format!("({p})")) + .collect::>() + .join(" AND ") + }) }; - // Format each column's values as strings. - // Positron's ColumnValue is `number | string`: numbers are special - // value codes (0 = NULL, 1 = NA, 2 = NaN), strings are formatted data. - const SPECIAL_VALUE_NULL: i64 = 0; + let mut count = count_rows(reader, &self.table_path, combine(&built).as_deref()); + if count.is_err() { + // Find the filters the database rejects, and count without them. + built.retain( + |(i, p)| match count_rows(reader, &self.table_path, Some(p)) { + Ok(_) => true, + Err(e) => { + invalidate(&mut filters[*i], e); + false + } + }, + ); + count = count_rows(reader, &self.table_path, combine(&built).as_deref()); + } + let selected = count?; + + for (i, _) in &built { + filters[*i]["is_valid"] = json!(true); + } + self.row_filters = filters; + self.predicate = combine(&built); + self.selected_num_rows = selected; - let columns: Vec> = (0..df.width()) - .map(|col_idx| { - let col = df.get_columns()[col_idx].clone(); - use arrow::array::Array; - (0..df.height()) - .map(|row_idx| { - if col.is_null(row_idx) { - json!(SPECIAL_VALUE_NULL) - } else { - let s = ggsql::array_util::value_to_string(&col, row_idx); - // Strip surrounding quotes from string values - let s = s.trim_matches('"'); - Value::String(s.to_string()) - } - }) - .collect() - }) + Ok(json!({ + "selected_num_rows": selected, + "had_errors": had_errors + })) + } + + fn set_sort_columns(&mut self, params: &Value) -> Result { + let keys = params["sort_keys"].as_array().cloned().unwrap_or_default(); + // Checked here so a bad key fails the request, not every page after it. + filter::order_by(&keys, &self.columns, None)?; + self.sort_keys = keys; + Ok(json!({})) + } + + /// The selected cells as CSV, TSV or HTML. + fn export_data_selection(&self, params: &Value, reader: &dyn Reader) -> Result { + let format = params["format"].as_str().unwrap_or("csv"); + let kind = params["selection"]["kind"].as_str().unwrap_or(""); + let sel = ¶ms["selection"]["selection"]; + + let index = |v: &Value| v.as_u64().map(|n| n as usize).ok_or("Bad selection index"); + let span = |first: &Value, last: &Value| -> Result, String> { + Ok((index(first)?..=index(last)?).collect()) + }; + let list = |v: &Value| -> Result, String> { + v.as_array() + .ok_or("Bad selection indices")? + .iter() + .map(|i| index(i).map_err(str::to_string)) + .collect() + }; + let all_columns: Vec = (0..self.columns.len()).collect(); + + // `None` rows means every row of the view. + let (rows, cols): (Option>, Vec) = match kind { + "single_cell" => ( + Some(vec![index(&sel["row_index"])?]), + vec![index(&sel["column_index"])?], + ), + "cell_range" => ( + Some(span(&sel["first_row_index"], &sel["last_row_index"])?), + span(&sel["first_column_index"], &sel["last_column_index"])?, + ), + "column_range" => (None, span(&sel["first_index"], &sel["last_index"])?), + "row_range" => ( + Some(span(&sel["first_index"], &sel["last_index"])?), + all_columns, + ), + "column_indices" => (None, list(&sel["indices"])?), + "row_indices" => (Some(list(&sel["indices"])?), all_columns), + "cell_indices" => ( + Some(list(&sel["row_indices"])?), + list(&sel["column_indices"])?, + ), + other => return Err(format!("Unknown selection kind {other:?}")), + }; + + let (df, offset) = match &rows { + None => (self.fetch(reader, &cols, Rows::All)?, 0), + Some(rows) => { + let first = rows.iter().min().copied().unwrap_or(0); + let last = rows.iter().max().copied().unwrap_or(0); + (self.fetch(reader, &cols, Rows::Range(first, last))?, first) + } + }; + let picked: Vec = match &rows { + None => (0..df.height()).collect(), + Some(rows) => rows + .iter() + .map(|r| r - offset) + .filter(|&r| r < df.height()) + .collect(), + }; + let cells: Vec> = picked + .iter() + .map(|&r| df.get_columns().iter().map(|c| cell_string(c, r)).collect()) .collect(); - json!({ "columns": columns }) + // A single cell is copied as its bare value. + let data = if kind == "single_cell" { + cells + .first() + .and_then(|r| r.first()) + .cloned() + .unwrap_or_default() + } else { + let header: Vec = cols.iter().map(|&i| self.columns[i].name.clone()).collect(); + export::format(Some(&header), &cells, format)? + }; + Ok(json!({ "data": data, "format": format })) + } + + /// SQL for the session's database reproducing the explorer's view. + fn convert_to_code(&self, params: &Value, session: &Session) -> Value { + let from = match &self.query { + Some(query) => filter::subquery(query), + None => self.table_path.clone(), + }; + let lines = filter::convert_to_code( + params, + &from, + &self.columns, + session.reader.dialect(), + session.caps, + ); + json!({ "converted_code": lines }) } /// Handle `get_column_profiles` — returns `{}` as the RPC result and sends @@ -331,7 +602,7 @@ impl DataExplorerState { Some(arr) => arr, None => { return RpcResponse { - result: json!({}), + result: Ok(json!({})), event: Some(RpcEvent { method: "return_column_profiles".into(), params: json!({ @@ -361,7 +632,7 @@ impl DataExplorerState { } RpcResponse { - result: json!({}), + result: Ok(json!({})), event: Some(RpcEvent { method: "return_column_profiles".into(), params: json!({ @@ -471,7 +742,7 @@ impl DataExplorerState { let sql = format!( "SELECT {} FROM {}", select_parts.join(", "), - self.table_path + self.relation() ); let df = match reader.execute_sql(&sql) { @@ -539,7 +810,7 @@ impl DataExplorerState { // must sit in a query over `from` (native dialects return // a bare aggregate, which without a FROM would compute // over an empty input and yield NULL). - let from_query = format!("SELECT * FROM {}", self.table_path); + let from_query = format!("SELECT * FROM {}", self.relation()); let median_expr = dialect.sql_quantile(&col.name, 0.5, FromItem::Query(&from_query), &[]); let median_sql = format!( @@ -673,7 +944,7 @@ impl DataExplorerState { FROM {table} WHERE {c} IS NOT NULL", c = quoted_col, t = float_type, - table = self.table_path, + table = self.relation(), ); let bounds_df = reader.execute_sql(&bounds_sql).ok()?; @@ -737,7 +1008,7 @@ impl DataExplorerState { ORDER BY \"clamped_bin\"", c = quoted_col, t = float_type, - table = self.table_path, + table = self.relation(), min = min_val, width = bin_width, num_bins = num_bins, @@ -775,7 +1046,7 @@ impl DataExplorerState { .unwrap_or_default(); let mut quantile_results = Vec::new(); - let from_query = format!("SELECT * FROM {}", self.table_path); + let from_query = format!("SELECT * FROM {}", self.relation()); for q in &quantiles_param { if let Some(q_val) = q.as_f64() { let expr = @@ -830,7 +1101,7 @@ impl DataExplorerState { ORDER BY COUNT(*) DESC \ LIMIT {limit}", c = quoted_col, - table = self.table_path, + table = self.relation(), limit = limit, ); @@ -859,7 +1130,7 @@ impl DataExplorerState { let count_sql = format!( "SELECT COUNT({c}) AS \"total\" FROM {table}", c = quoted_col, - table = self.table_path, + table = self.relation(), ); let other_count = reader .execute_sql(&count_sql) @@ -885,6 +1156,50 @@ impl DataExplorerState { } } +/// `SELECT COUNT(*)` over `table`, filtered by `predicate` when given. +fn count_rows(reader: &dyn Reader, table: &str, predicate: Option<&str>) -> Result { + let filter = predicate.map(|p| format!(" WHERE {p}")).unwrap_or_default(); + let sql = format!("SELECT COUNT(*) AS \"n\" FROM {table}{filter}"); + let df = reader.execute_sql(&sql).map_err(|e| e.to_string())?; + let col = df.column("n").map_err(|e| e.to_string())?; + if col.is_empty() { + return Ok(0); + } + ggsql::array_util::value_to_string(col, 0) + .parse::() + .map_err(|e| e.to_string()) +} + +/// A cell as text, with string values unquoted and a null as empty. +fn cell_string(col: &arrow::array::ArrayRef, row: usize) -> String { + use arrow::array::Array; + if col.is_null(row) { + return String::new(); + } + ggsql::array_util::value_to_string(col, row) + .trim_matches('"') + .to_string() +} + +/// A cell as Positron's `ColumnValue`: formatted text, or a number coding a +/// special value — 0 null, 2 NaN, 10 infinity, 11 negative infinity. +fn cell_value(col: &arrow::array::ArrayRef, row: usize) -> Value { + use arrow::array::Array; + if col.is_null(row) { + return json!(0); + } + let s = cell_string(col, row); + if col.data_type().is_floating() { + match s.to_ascii_lowercase().as_str() { + "nan" => return json!(2), + "inf" | "infinity" => return json!(10), + "-inf" | "-infinity" => return json!(11), + _ => {} + } + } + Value::String(s) +} + /// Map a SQL type name (from information_schema or SHOW COLUMNS) to a Positron display type. /// /// Handles both simple type names (e.g. "INTEGER", "VARCHAR") and Snowflake's @@ -992,6 +1307,270 @@ fn clean_type_name(type_name: &str) -> String { mod tests { use super::*; + const QUERY: &str = "SELECT i AS n, 'r' || i AS label FROM range(10) t(i)"; + + /// An explorer on QUERY's result, with the store that keeps its table. + fn explorer() -> (crate::results::ResultStore, DataExplorerState) { + let df = ggsql::reader::connection::reader_from_uri("duckdb://memory") + .unwrap() + .execute_sql(QUERY) + .unwrap(); + let mut store = crate::results::ResultStore::new(4).unwrap(); + let key = store + .insert(df, "Query [3]".into(), QUERY.to_string(), |_| false) + .unwrap(); + let state = + DataExplorerState::open_result(store.reader(), store.get(&key).unwrap()).unwrap(); + assert_eq!(*state.source(), Source::Result(key)); + (store, state) + } + + fn rpc( + state: &mut DataExplorerState, + store: &crate::results::ResultStore, + method: &str, + params: Value, + ) -> Result { + let session = Session { + reader: store.reader(), + caps: Capabilities { regex: true }, + }; + state + .handle_rpc(method, ¶ms, store.reader(), session) + .result + } + + /// The first `n` values of a column, as the grid would ask for them. + fn page( + state: &mut DataExplorerState, + store: &crate::results::ResultStore, + column: usize, + n: usize, + ) -> Value { + let params = json!({"columns": [ + {"column_index": column, "spec": {"first_index": 0, "last_index": n - 1}} + ]}); + rpc(state, store, "get_data_values", params).unwrap()["columns"][0].clone() + } + + fn compare(column: usize, op: &str, value: &str) -> Value { + json!({ + "filter_id": format!("{column}{op}{value}"), + "filter_type": "compare", + "column_schema": {"column_index": column}, + "condition": "and", + "params": {"op": op, "value": value} + }) + } + + #[test] + fn test_a_retained_result_opens_and_pages() { + let (store, mut state) = explorer(); + + let st = rpc(&mut state, &store, "get_state", json!({})).unwrap(); + assert_eq!(st["display_name"], "Query [3]"); + assert_eq!(st["table_shape"]["num_rows"], 10); + assert_eq!(st["table_shape"]["num_columns"], 2); + + let values = rpc( + &mut state, + &store, + "get_data_values", + json!({"columns": [ + {"column_index": 0, "spec": {"first_index": 1, "last_index": 2}}, + {"column_index": 1, "spec": {"indices": [4, 2]}} + ]}), + ) + .unwrap(); + assert_eq!(values["columns"], json!([["1", "2"], ["r4", "r2"]])); + } + + #[test] + fn test_filters_narrow_the_view() { + let (store, mut state) = explorer(); + let reply = rpc( + &mut state, + &store, + "set_row_filters", + json!({"filters": [compare(0, ">=", "5"), compare(0, "!=", "7")]}), + ) + .unwrap(); + assert_eq!(reply, json!({"selected_num_rows": 4, "had_errors": false})); + + let st = rpc(&mut state, &store, "get_state", json!({})).unwrap(); + assert_eq!(st["table_shape"]["num_rows"], 4); + assert_eq!(st["table_unfiltered_shape"]["num_rows"], 10); + assert_eq!(st["row_filters"][0]["is_valid"], true); + assert_eq!(page(&mut state, &store, 0, 4), json!(["5", "6", "8", "9"])); + } + + #[test] + fn test_a_bad_filter_is_left_out_and_says_why() { + let (store, mut state) = explorer(); + // One the kernel rejects, and one only the database can: `(` is not a + // valid regular expression. + let bad_regex = json!({ + "filter_id": "re", + "filter_type": "search", + "column_schema": {"column_index": 1}, + "condition": "and", + "params": {"search_type": "regex_match", "term": "(", "case_sensitive": true} + }); + let reply = rpc( + &mut state, + &store, + "set_row_filters", + json!({"filters": [compare(0, "=", "abc"), bad_regex, compare(0, "<", "3")]}), + ) + .unwrap(); + assert_eq!(reply, json!({"selected_num_rows": 3, "had_errors": true})); + + let st = rpc(&mut state, &store, "get_state", json!({})).unwrap(); + let filters = st["row_filters"].as_array().unwrap(); + assert_eq!(filters[0]["is_valid"], false); + assert!(filters[0]["error_message"] + .as_str() + .unwrap() + .contains("not a number")); + assert_eq!(filters[1]["is_valid"], false); + assert_eq!(filters[2]["is_valid"], true); + } + + #[test] + fn test_sorting_applies_on_top_of_filters() { + let (store, mut state) = explorer(); + rpc( + &mut state, + &store, + "set_row_filters", + json!({"filters": [compare(0, "<", "4")]}), + ) + .unwrap(); + rpc( + &mut state, + &store, + "set_sort_columns", + json!({"sort_keys": [{"column_index": 0, "ascending": false}]}), + ) + .unwrap(); + assert_eq!( + page(&mut state, &store, 1, 4), + json!(["r3", "r2", "r1", "r0"]) + ); + + // Clearing the sort restores the query's own order. + rpc( + &mut state, + &store, + "set_sort_columns", + json!({"sort_keys": []}), + ) + .unwrap(); + assert_eq!(page(&mut state, &store, 1, 2), json!(["r0", "r1"])); + } + + #[test] + fn test_profiles_see_only_the_filtered_rows() { + let (store, mut state) = explorer(); + rpc( + &mut state, + &store, + "set_row_filters", + json!({"filters": [compare(0, ">=", "6")]}), + ) + .unwrap(); + let response = state.handle_rpc( + "get_column_profiles", + &json!({"callback_id": "cb", "profiles": [{ + "column_index": 0, + "profiles": [{"profile_type": "summary_stats"}] + }]}), + store.reader(), + Session { + reader: store.reader(), + caps: Capabilities { regex: true }, + }, + ); + let event = response.event.unwrap(); + let stats = &event.params["profiles"][0]["summary_stats"]["number_stats"]; + assert_eq!(stats["min_value"], "6"); + assert_eq!(stats["max_value"], "9"); + } + + #[test] + fn test_export_follows_the_view() { + let (store, mut state) = explorer(); + rpc( + &mut state, + &store, + "set_sort_columns", + json!({"sort_keys": [{"column_index": 0, "ascending": false}]}), + ) + .unwrap(); + + let csv = rpc( + &mut state, + &store, + "export_data_selection", + json!({"format": "csv", "selection": {"kind": "cell_range", "selection": { + "first_row_index": 0, "last_row_index": 1, + "first_column_index": 0, "last_column_index": 1 + }}}), + ) + .unwrap(); + assert_eq!(csv["data"], "n,label\n9,r9\n8,r8"); + + let cell = rpc( + &mut state, + &store, + "export_data_selection", + json!({"format": "tsv", "selection": {"kind": "single_cell", "selection": { + "row_index": 2, "column_index": 1 + }}}), + ) + .unwrap(); + assert_eq!(cell["data"], "r7"); + + let rows = rpc( + &mut state, + &store, + "export_data_selection", + json!({"format": "tsv", "selection": {"kind": "row_indices", "selection": { + "indices": [3, 0] + }}}), + ) + .unwrap(); + assert_eq!(rows["data"], "n\tlabel\n6\tr6\n9\tr9"); + } + + #[test] + fn test_convert_to_code_wraps_the_original_query() { + let (store, mut state) = explorer(); + let code = rpc( + &mut state, + &store, + "convert_to_code", + json!({ + "column_filters": [], + "row_filters": [compare(0, ">", "2")], + "sort_keys": [{"column_index": 1, "ascending": true}], + "code_syntax_name": {"code_syntax_name": "SQL"} + }), + ) + .unwrap(); + assert_eq!( + code["converted_code"], + json!([ + "SELECT *", + "FROM (", + format!(" {QUERY}"), + ") AS \"result\"", + "WHERE \"n\" > 2", + "ORDER BY \"label\" ASC" + ]) + ); + } + #[test] fn test_sql_type_to_display() { assert_eq!(sql_type_to_display("INTEGER"), "integer"); diff --git a/ggsql-jupyter/src/display.rs b/ggsql-jupyter/src/display.rs index 6d7fb3e43..db4c62e92 100644 --- a/ggsql-jupyter/src/display.rs +++ b/ggsql-jupyter/src/display.rs @@ -71,14 +71,30 @@ impl SessionKind { } } +/// How many rows of a long result a console shows, unless configured. +pub const DEFAULT_CONSOLE_ROWS: usize = 10; + /// Frontend-supplied hints about the output rendering slot. -#[derive(Default, Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy)] pub struct RenderHints { pub kind: SessionKind, /// Width of the output slot in CSS pixels, when the frontend says. pub output_width_px: Option, /// Device pixel ratio of the display, when the frontend says. pub pixel_ratio: Option, + /// How many rows of a long result a console shows (`--console-rows`). + pub console_rows: usize, +} + +impl Default for RenderHints { + fn default() -> Self { + Self { + kind: SessionKind::default(), + output_width_px: None, + pixel_ratio: None, + console_rows: DEFAULT_CONSOLE_ROWS, + } + } } impl RenderHints { @@ -102,6 +118,7 @@ impl RenderHints { kind: SessionKind::resolve(header.session.as_str(), mode), output_width_px, pixel_ratio, + ..Self::default() } } @@ -172,17 +189,21 @@ pub fn format_display_data( } } } - ExecutionResult::DataFrame(df) => { + // A `-- @view` table shows in the cell too; the kernel opens the + // explorer alongside it. + ExecutionResult::DataFrame(df) | ExecutionResult::View { df, .. } => { // DDL statements return DataFrames with 0 columns - don't display anything if df.width() == 0 { Ok(Formatted::Nothing) } else { - Ok(Formatted::Bundle(format_dataframe(df))) + Ok(Formatted::Bundle(format_dataframe(df, hints))) } } ExecutionResult::ConnectionChanged { display_name, .. } => { Ok(Formatted::Bundle(format_connection_changed(&display_name))) } + // Opening the explorer is the whole of the output. + ExecutionResult::ViewResult { .. } => Ok(Formatted::Nothing), } } @@ -239,8 +260,8 @@ fn format_static(spec: Box, request: RenderRequest, backend: &PlotBackend) } /// Format DataFrame as HTML table -fn format_dataframe(df: DataFrame) -> Value { - let html = dataframe_to_html(&df); +fn format_dataframe(df: DataFrame, hints: &RenderHints) -> Value { + let html = dataframe_to_html(&df, html_rows_shown(df.height(), hints)); let text = dataframe_to_text(&df); json!({ @@ -253,8 +274,23 @@ fn format_dataframe(df: DataFrame) -> Value { }) } -/// Convert DataFrame to HTML table -fn dataframe_to_html(df: &DataFrame) -> String { +/// How many of a result's `height` rows its HTML table shows. +/// +/// A console result is followed by a link to the whole of it in the data +/// explorer, so it shows what R's tibble does: every row up to twice +/// `console_rows`, else the first `console_rows`, which keeps the link in +/// view. Anywhere else the table is all there is, so it shows up to 100. +fn html_rows_shown(height: usize, hints: &RenderHints) -> usize { + let n = hints.console_rows; + match hints.kind { + SessionKind::PositronConsole if height > n.saturating_mul(2) => n, + SessionKind::PositronConsole => height, + _ => height.min(100), + } +} + +/// Convert DataFrame to an HTML table of its first `row_limit` rows +fn dataframe_to_html(df: &DataFrame, row_limit: usize) -> String { use ggsql::array_util::value_to_string; let mut html = String::from("\n"); @@ -265,8 +301,6 @@ fn dataframe_to_html(df: &DataFrame) -> String { } html.push_str("\n\n"); - // Data rows (limit to first 100 for performance) - let row_limit = df.height().min(100); for i in 0..row_limit { html.push_str(""); for col in df.get_columns() { @@ -310,7 +344,7 @@ fn dataframe_to_text(df: &ggsql::DataFrame) -> String { } /// Escape HTML special characters -fn escape_html(s: &str) -> String { +pub(crate) fn escape_html(s: &str) -> String { s.replace('&', "&") .replace('<', "<") .replace('>', ">") @@ -481,6 +515,47 @@ mod tests { ); } + #[test] + fn test_html_rows_shown() { + let console = positron_console(); + // The console shows a short result whole, and a long one's first 10. + assert_eq!(html_rows_shown(20, &console), 20); + assert_eq!(html_rows_shown(21, &console), 10); + assert_eq!(html_rows_shown(5000, &console), 10); + // Anywhere else, up to 100. + assert_eq!(html_rows_shown(50, &positron_notebook()), 50); + assert_eq!(html_rows_shown(5000, &positron_notebook()), 100); + assert_eq!(html_rows_shown(5000, &RenderHints::default()), 100); + } + + #[test] + fn test_console_rows_is_configurable() { + let console = RenderHints { + console_rows: 3, + ..positron_console() + }; + assert_eq!(html_rows_shown(6, &console), 6); + assert_eq!(html_rows_shown(7, &console), 3); + // Only the console's table follows it. + let notebook = RenderHints { + console_rows: 3, + ..positron_notebook() + }; + assert_eq!(html_rows_shown(50, ¬ebook), 50); + } + + #[test] + fn test_console_table_counts_the_rows_it_leaves_out() { + let df = ggsql::reader::connection::reader_from_uri("duckdb://memory") + .unwrap() + .execute_sql("SELECT i FROM range(5000) t(i)") + .unwrap(); + let bundle = format_dataframe(df, &positron_console()); + let html = bundle["data"]["text/html"].as_str().unwrap(); + assert_eq!(html.matches("").count(), 1 + 10 + 1, "{html}"); + assert!(html.contains("... 4990 more rows"), "{html}"); + } + #[test] fn test_html_escape() { assert_eq!( @@ -498,8 +573,7 @@ mod tests { fn positron_console() -> RenderHints { RenderHints { kind: SessionKind::PositronConsole, - output_width_px: None, - pixel_ratio: None, + ..RenderHints::default() } } @@ -507,7 +581,7 @@ mod tests { RenderHints { kind: SessionKind::PositronNotebook, output_width_px: Some(589), - pixel_ratio: None, + ..RenderHints::default() } } diff --git a/ggsql-jupyter/src/executor.rs b/ggsql-jupyter/src/executor.rs index 40dcf2ef8..b6f031f72 100644 --- a/ggsql-jupyter/src/executor.rs +++ b/ggsql-jupyter/src/executor.rs @@ -37,6 +37,16 @@ pub enum ExecutionResult { Visualization(Box), /// Connection changed via meta-command ConnectionChanged { display_name: String }, + /// Open a retained query result in the data explorer, by key. The + /// executor never sees retained results, so opening it is the kernel's. + ViewResult { key: String }, + /// A query run under `-- @view`, whose table opens in the data explorer as + /// well as showing in the cell. `title` is the one the directive gave, if + /// any. + View { + df: DataFrame, + title: Option, + }, } // `Spec` is neither `Debug` nor `Clone`, so this summarises rather than @@ -60,6 +70,13 @@ impl std::fmt::Debug for ExecutionResult { .debug_struct("ConnectionChanged") .field("display_name", display_name) .finish(), + Self::ViewResult { key } => f.debug_struct("ViewResult").field("key", key).finish(), + Self::View { df, title } => f + .debug_struct("View") + .field("rows", &df.height()) + .field("columns", &df.width()) + .field("title", title) + .finish(), } } } @@ -156,6 +173,10 @@ pub fn host_for_uri(uri: &str) -> String { const META_CONNECT_PREFIX: &str = "-- @connect:"; /// The `-- @uncache` meta-command prefix. const META_UNCACHE_PREFIX: &str = "-- @uncache"; +/// The `-- @view` meta-command prefix. +const META_VIEW_PREFIX: &str = "-- @view"; +/// The `-- @view-result` meta-command prefix. +const META_VIEW_RESULT_PREFIX: &str = "-- @view-result"; /// A leading cell directive expressed as a `--` line comment. #[derive(Debug, PartialEq, Eq)] @@ -164,6 +185,12 @@ pub enum MetaCommand { Connect(String), /// Clear the active reader's cache. Uncache, + /// Show the rest of the cell's result in the data explorer, under the + /// given title if there is one. + View(Option), + /// Open a retained query result in the data explorer. Sent by the ggsql + /// extension when the link under a console result is clicked. + ViewResult(String), } /// Split `code` into its first line and the remainder. @@ -196,6 +223,23 @@ pub fn take_leading_meta(code: &str) -> Option<(MetaCommand, &str)> { if line == META_UNCACHE_PREFIX { return Some((MetaCommand::Uncache, rest)); } + if let Some(key) = line.strip_prefix(META_VIEW_RESULT_PREFIX) { + // A space is required, so `-- @view-results` is not taken for it. + if let Some(key) = key.strip_prefix(char::is_whitespace) { + return Some((MetaCommand::ViewResult(key.trim().to_string()), rest)); + } + } + if line == META_VIEW_PREFIX { + return Some((MetaCommand::View(None), rest)); + } + if let Some(title) = line + .strip_prefix(META_VIEW_PREFIX) + .and_then(|t| t.strip_prefix(':')) + { + let title = title.trim(); + let title = (!title.is_empty()).then(|| title.to_string()); + return Some((MetaCommand::View(title), rest)); + } None } @@ -253,6 +297,7 @@ impl QueryExecutor { // Apply any leading meta-command lines, then run whatever SQL remains. let mut code = code; let mut last_connect: Option = None; + let mut view: Option> = None; while let Some((cmd, rest)) = take_leading_meta(code) { match cmd { MetaCommand::Connect(uri) => { @@ -264,10 +309,21 @@ impl QueryExecutor { tracing::info!("Meta-command: clearing cache"); self.reader.clear_cache()?; } + MetaCommand::View(title) => view = Some(title), + MetaCommand::ViewResult(key) => { + if !rest.trim().is_empty() { + anyhow::bail!("`-- @view-result` must be the only line in its cell"); + } + return Ok(ExecutionResult::ViewResult { key }); + } } code = rest; } + if let Some(title) = view { + return self.execute_view(code, title); + } + // A cell that was nothing but meta-commands. if code.trim().is_empty() { if let Some(uri) = last_connect { @@ -306,6 +362,26 @@ impl QueryExecutor { // layer's job, because only it knows where the output is going. Ok(ExecutionResult::Visualization(Box::new(spec))) } + + /// Run the query under a `-- @view` directive for the data explorer. + /// + /// Checked before running, so a mistake costs nothing: the explorer shows + /// a table, so a plot query and a cell with nothing to show are refused. + fn execute_view(&mut self, code: &str, title: Option) -> Result { + if code.trim().is_empty() { + anyhow::bail!("`-- @view` needs a query after it to show in the Data Explorer"); + } + if validate(code)?.has_visual() { + anyhow::bail!( + "`-- @view` shows a query's table, not a plot. Remove the VISUALISE clause to view the data." + ); + } + let df = self.reader.execute_sql(code)?; + if df.width() == 0 { + anyhow::bail!("`-- @view` needs a query that returns a table"); + } + Ok(ExecutionResult::View { df, title }) + } } #[cfg(test)] @@ -375,6 +451,84 @@ mod tests { assert_eq!(take_leading_meta("-- @uncache foo"), None); } + #[test] + fn test_take_leading_meta_view_result() { + assert_eq!( + take_leading_meta("-- @view-result abc123"), + Some((MetaCommand::ViewResult("abc123".to_string()), "")) + ); + assert_eq!(take_leading_meta("-- @view-results abc123"), None); + } + + #[test] + fn test_view_result_must_stand_alone() { + let mut executor = QueryExecutor::new().unwrap(); + match executor.execute("-- @view-result abc123").unwrap() { + ExecutionResult::ViewResult { key } => assert_eq!(key, "abc123"), + other => panic!("expected ViewResult, got {other:?}"), + } + assert!(executor + .execute("-- @view-result abc123\nSELECT 1") + .is_err()); + } + + #[test] + fn test_take_leading_meta_view() { + assert_eq!( + take_leading_meta("-- @view\nSELECT 1"), + Some((MetaCommand::View(None), "SELECT 1")) + ); + assert_eq!( + take_leading_meta("-- @view: Sales by region \nSELECT 1"), + Some(( + MetaCommand::View(Some("Sales by region".to_string())), + "SELECT 1" + )) + ); + assert_eq!( + take_leading_meta("-- @view:\nSELECT 1"), + Some((MetaCommand::View(None), "SELECT 1")) + ); + // Anything else after the name is an ordinary comment. + assert_eq!(take_leading_meta("-- @viewer\nSELECT 1"), None); + assert_eq!(take_leading_meta("-- @view abc\nSELECT 1"), None); + } + + #[test] + fn test_view_runs_the_query_for_the_explorer() { + let mut executor = QueryExecutor::new().unwrap(); + match executor + .execute("-- @view: Numbers\nSELECT * FROM range(3) t(n)") + .unwrap() + { + ExecutionResult::View { df, title } => { + assert_eq!(df.height(), 3); + assert_eq!(title.as_deref(), Some("Numbers")); + } + other => panic!("expected View, got {other:?}"), + } + } + + #[test] + fn test_view_combines_with_other_directives() { + let mut executor = QueryExecutor::new().unwrap(); + let result = executor + .execute("-- @uncache\n-- @view\nSELECT 1 AS x") + .unwrap(); + assert!(matches!(result, ExecutionResult::View { title: None, .. })); + } + + #[test] + fn test_view_refuses_what_has_no_table() { + let mut executor = QueryExecutor::new().unwrap(); + let mut err = |code: &str| executor.execute(code).unwrap_err().to_string(); + assert!(err("-- @view").contains("needs a query")); + assert!(err("-- @view\nCREATE TABLE t (x INT)").contains("returns a table")); + assert!( + err("-- @view\nSELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point").contains("not a plot") + ); + } + #[test] fn test_take_leading_meta_non_directive() { assert_eq!(take_leading_meta("SELECT 1"), None); diff --git a/ggsql-jupyter/src/kernel.rs b/ggsql-jupyter/src/kernel.rs index 369061d74..e461be7e2 100644 --- a/ggsql-jupyter/src/kernel.rs +++ b/ggsql-jupyter/src/kernel.rs @@ -4,12 +4,13 @@ //! handling kernel_info, execute, and shutdown requests. use crate::connection; -use crate::data_explorer::{DataExplorerState, RpcResponse}; -use crate::display::{format_display_data, Formatted, RenderHints, SessionMode}; +use crate::data_explorer::{Capabilities, DataExplorerState, RpcResponse, Session, Source}; +use crate::display::{format_display_data, Formatted, RenderHints, SessionKind, SessionMode}; use crate::executor::{self, ExecutionResult, QueryExecutor}; use crate::message::{ConnectionInfo, JupyterMessage, MessageHeader}; use crate::plot::comm::{PlotMetadata, RenderParams, RpcError}; use crate::plot::{PlotBackend, RenderOutcome, RenderTicket}; +use crate::results::ResultStore; use anyhow::Result; use ggsql::reader::Spec; use hmac::{Hmac, Mac}; @@ -65,6 +66,10 @@ pub struct KernelServer { ui_comm_id: Option, connection_comm_id: Option, data_explorer_comms: HashMap, + /// Console query results kept for the data explorer, in a private database. + results: ResultStore, + /// `--console-rows`: how many rows of a long result a console shows. + console_rows: usize, } impl KernelServer { @@ -74,6 +79,8 @@ impl KernelServer { reader_uri: &str, session_mode: Option, max_plots: usize, + max_results: usize, + console_rows: usize, ) -> Result { tracing::info!("Initializing kernel server"); @@ -141,6 +148,8 @@ impl KernelServer { ui_comm_id: None, connection_comm_id: None, data_explorer_comms: HashMap::new(), + results: ResultStore::new(max_results)?, + console_rows, }; // Send initial "starting" status on IOPub @@ -358,7 +367,10 @@ impl KernelServer { let content = &parent.content; let code = content["code"].as_str().unwrap_or(""); let silent = content["silent"].as_bool().unwrap_or(false); - let hints = RenderHints::from_request(&parent.header, content, self.session_mode); + let hints = RenderHints { + console_rows: self.console_rows, + ..RenderHints::from_request(&parent.header, content, self.session_mode) + }; // `kind` decides where a plot goes, so log its two inputs: the declared // mode and the session id the heuristic falls back to. A plot in the @@ -408,7 +420,64 @@ impl KernelServer { // the user's. Both stages feed one `Result` and share the error path // below; propagating out of the loop would end the process and take // the session's database with it. + // Outside Positron there is no data explorer, so `-- @view` shows its + // table in the cell, with a note saying why. + let mut note = None; + let executed = match executed { + Ok(ExecutionResult::View { df, .. }) if hints.kind == SessionKind::Standalone => { + note = Some("`-- @view` opens the Data Explorer, which needs Positron.\n"); + Ok(ExecutionResult::DataFrame(df)) + } + other => other, + }; + + // A console result is kept so the link under it can open it in the + // data explorer. Silent cells have no link to click, and DDL no table. + let linked = !silent && hints.kind == SessionKind::PositronConsole; + let mut retained = None; + if let Ok(ExecutionResult::DataFrame(df)) = &executed { + if linked && df.width() > 0 { + let title = format!("Query [{}]", self.execution_count); + match self.retain_result(df.clone(), title, query_of(code)) { + Ok(key) => retained = Some(key), + // Costs only the link, never the cell's output. + Err(e) => tracing::warn!("Failed to retain query result: {}", e), + } + } + } + let result = match executed { + Ok(ExecutionResult::ViewResult { key }) => self + .open_result_explorer(&key, parent) + .await + .map(|()| Formatted::Nothing), + Ok(ExecutionResult::View { df, title }) => { + // A silent cell has no count of its own to name it by. + let title = title.unwrap_or_else(|| match silent { + true => "Query".to_string(), + false => format!("Query [{}]", self.execution_count), + }); + // The table shows in the cell as for any query, and the + // explorer opens as well. + let shown = (!silent).then(|| ExecutionResult::DataFrame(df.clone())); + match self.retain_result(df, title, query_of(code)) { + Ok(key) => { + let opened = self.open_result_explorer(&key, parent).await; + // The link reopens the tab once it has been closed. + if linked { + retained = Some(key); + } + match (opened, shown) { + (Err(e), _) => Err(e), + (Ok(()), None) => Ok(Formatted::Nothing), + (Ok(()), Some(shown)) => { + format_display_data(shown, &hints, &self.plots) + } + } + } + Err(e) => Err(e), + } + } Ok(exec_result) => { // A bare connection change renders nothing. let is_connection_changed = @@ -448,6 +517,18 @@ impl KernelServer { Formatted::Nothing => {} } + // Its own output, because the console renders the HTML table + // and ignores the bundle's `text/plain`. + if let Some(key) = &retained { + let link = explorer_link(&parent.header.session, key); + self.send_iopub("stream", json!({ "name": "stdout", "text": link }), parent) + .await?; + } + if let Some(note) = note.filter(|_| !silent) { + self.send_iopub("stream", json!({ "name": "stderr", "text": note }), parent) + .await?; + } + // Send execute_reply self.send_shell_reply( "execute_reply", @@ -928,6 +1009,95 @@ impl KernelServer { Ok(()) } + /// What the session's database can do, for the SQL the data explorer + /// writes against it. Only DuckDB has `regexp_matches`. + fn session_caps(&self) -> Capabilities { + Capabilities { + regex: self.executor.reader_uri().starts_with("duckdb://"), + } + } + + /// Keep a query result for the data explorer, returning its key. + /// + /// Results a data explorer has open are kept whatever the limit. + fn retain_result( + &mut self, + df: ggsql::DataFrame, + title: String, + query: String, + ) -> Result { + let comms = &self.data_explorer_comms; + let in_use = |key: &str| { + comms + .values() + .any(|state| matches!(state.source(), Source::Result(k) if k == key)) + }; + self.results.insert(df, title, query, in_use) + } + + /// Open a retained query result in the data explorer, for `-- @view` and + /// `-- @view-result`. + /// + /// Not `inline_only`, so Positron opens the explorer in an editor tab as + /// soon as the comm arrives. + async fn open_result_explorer(&mut self, key: &str, parent: &JupyterMessage) -> Result<()> { + let stored = self.results.get(key)?; + + // A result already open in an explorer is brought to the front rather + // than opened again, as Positron's Variables pane does for a variable. + let source = Source::Result(key.to_string()); + let open = self + .data_explorer_comms + .iter() + .find(|(_, state)| *state.source() == source) + .map(|(id, _)| id.clone()); + if let (Some(explorer), Some(ui)) = (open, self.ui_comm_id.clone()) { + tracing::info!("Focusing data explorer comm {} on result {}", explorer, key); + return self + .send_iopub( + "comm_msg", + json!({ + "comm_id": ui, + "data": { + "jsonrpc": "2.0", + "method": "open_editor", + "params": { + "file": explorer_uri(&explorer), + "line": 1, + "column": 1, + "kind": "uri" + } + } + }), + parent, + ) + .await; + } + + let state = DataExplorerState::open_result(self.results.reader(), stored) + .map_err(|e| anyhow::anyhow!(e))?; + + let comm_id = uuid::Uuid::new_v4().to_string(); + let msg = self.create_message( + "comm_open", + json!({ + "comm_id": comm_id, + "target_name": "positron.dataExplorer", + "data": { + "title": stored.title, + "inline_only": false + } + }), + Some(parent), + ); + let zmq_msg = self.serialize_message_with_topic(&msg, "comm_open")?; + self.iopub.send(zmq_msg).await?; + + tracing::info!("Opened data explorer comm {} on result {}", comm_id, key); + self.data_explorer_comms.insert(comm_id, state); + Ok(()) + } + /// Open a `positron.plot` comm for a freshly executed plot. /// /// Backend-initiated, one comm per plot like `positron.dataExplorer`, since @@ -1301,7 +1471,7 @@ impl KernelServer { }) .unwrap_or_default(); - match DataExplorerState::open(self.executor.reader(), &path) { + match DataExplorerState::open(self.executor.reader(), &path, self.session_caps()) { Ok(state) => { let de_comm_id = uuid::Uuid::new_v4().to_string(); let title = path.last().cloned().unwrap_or_default(); @@ -1377,23 +1547,40 @@ impl KernelServer { let params = &parent.content["data"]["params"]; + let caps = self.session_caps(); let RpcResponse { result, event } = - if let Some(state) = self.data_explorer_comms.get(comm_id) { - state.handle_rpc(method, params, self.executor.reader()) + if let Some(state) = self.data_explorer_comms.get_mut(comm_id) { + let session = Session { + reader: self.executor.reader(), + caps, + }; + let reader = match state.source() { + Source::Session => self.executor.reader(), + Source::Result(_) => self.results.reader(), + }; + state.handle_rpc(method, params, reader, session) } else { RpcResponse::reply(json!(null)) }; + let data = match result { + Ok(result) => json!({ "jsonrpc": "2.0", "id": rpc_id, "result": result }), + Err(message) => { + tracing::warn!("Data explorer {} failed: {}", method, message); + json!({ + "jsonrpc": "2.0", + "id": rpc_id, + "error": RpcError::Internal(message).to_json() + }) + } + }; + // Send the RPC reply self.send_shell_reply( "comm_msg", json!({ "comm_id": comm_id, - "data": { - "jsonrpc": "2.0", - "id": rpc_id, - "result": result - } + "data": data }), parent, identities, @@ -1625,6 +1812,31 @@ impl KernelServer { } } +/// A cell's SQL without its leading meta-commands, which are not SQL. +fn query_of(code: &str) -> String { + let mut rest = code; + while let Some((_, after)) = executor::take_leading_meta(rest) { + rest = after; + } + rest.trim().to_string() +} + +/// The console line that opens a retained result in the data explorer. +/// +/// An OSC 8 hyperlink: a click goes to Positron's opener, which offers it to +/// the ggsql session's `openResource`, which runs `-- @view-result ` +/// silently. +/// The session id lets the extension refuse a link meant for another session. +fn explorer_link(session: &str, key: &str) -> String { + format!("\x1b]8;;x-ggsql-view:{session}/{key}\x1b\\Open in Data Explorer\x1b]8;;\x1b\\\n") +} + +/// The URI of the editor Positron shows a data explorer comm in, which opening +/// again focuses (`PositronDataExplorerUri.generate`). +fn explorer_uri(comm_id: &str) -> String { + format!("positron-data-explorer:positron-data-explorer-{comm_id}") +} + /// Check if ggsql code is complete (balanced brackets, not in a string) /// If the code contains VISUALISE, it must also contain at least one DRAW layer fn is_code_complete(code: &str) -> bool { @@ -1677,3 +1889,16 @@ fn is_code_complete(code: &str) -> bool { true } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_explorer_link_is_an_osc8_hyperlink() { + assert_eq!( + explorer_link("ggsql-1", "abc"), + "\x1b]8;;x-ggsql-view:ggsql-1/abc\x1b\\Open in Data Explorer\x1b]8;;\x1b\\\n" + ); + } +} diff --git a/ggsql-jupyter/src/lib.rs b/ggsql-jupyter/src/lib.rs index 7b06b0960..d834df295 100644 --- a/ggsql-jupyter/src/lib.rs +++ b/ggsql-jupyter/src/lib.rs @@ -8,6 +8,7 @@ pub mod display; pub mod executor; pub mod message; pub mod plot; +pub mod results; // Re-export commonly used types pub use display::format_display_data; pub use executor::{ExecutionResult, QueryExecutor}; diff --git a/ggsql-jupyter/src/main.rs b/ggsql-jupyter/src/main.rs index 2cd08240a..6ef886145 100644 --- a/ggsql-jupyter/src/main.rs +++ b/ggsql-jupyter/src/main.rs @@ -10,6 +10,7 @@ mod executor; mod kernel; mod message; mod plot; +mod results; use anyhow::{Context, Result}; use clap::Parser; use display::SessionMode; @@ -48,6 +49,28 @@ struct Args { #[arg(long, default_value_t = 32, value_parser = clap::value_parser!(u32).range(1..))] max_plots: u32, + /// How many query results to keep available for the data explorer. + /// + /// In a Positron console each tabular result is retained so a link under + /// it can open it in the data explorer later. Once the limit is passed the + /// oldest results no explorer has open are dropped. + #[arg(long, default_value_t = 16, value_parser = clap::value_parser!(u32).range(1..))] + max_results: u32, + + /// How many rows of a query result a Positron console shows. + /// + /// A result of up to twice this many rows is shown whole; a longer one + /// shows this many, then a count of the rest. Also read from the + /// environment, which is how the ggsql extension sets it: a kernel that + /// predates the option ignores a variable but would refuse a flag. + #[arg( + long, + env = "GGSQL_CONSOLE_ROWS", + default_value_t = 10, + value_parser = clap::value_parser!(u32).range(1..) + )] + console_rows: u32, + /// Install the kernel spec #[arg(long)] install: bool, @@ -100,6 +123,8 @@ async fn main() -> Result<()> { &args.reader, args.session_mode, args.max_plots as usize, + args.max_results as usize, + args.console_rows as usize, ) .await?; diff --git a/ggsql-jupyter/src/results.rs b/ggsql-jupyter/src/results.rs new file mode 100644 index 000000000..4210165a5 --- /dev/null +++ b/ggsql-jupyter/src/results.rs @@ -0,0 +1,171 @@ +//! Query results retained so the data explorer can open them after the fact. +//! +//! Each result is copied into a private in-memory DuckDB rather than queried +//! where it came from: the explorer pages, sorts and profiles in one dialect +//! whatever the source database, a remote query is never re-run per page, and +//! the user's session never sees the tables. + +use anyhow::{anyhow, Result}; +use ggsql::reader::{connection::reader_from_uri, Reader}; +use ggsql::DataFrame; +use std::collections::VecDeque; + +/// The catalog and schema a result table lives in. `register` creates a +/// temporary table, and DuckDB keeps those in its `temp` catalog. +pub const CATALOG: &str = "temp"; +pub const SCHEMA: &str = "main"; + +/// One retained result. +pub struct StoredResult { + pub key: String, + /// The table holding the rows in the private database. + pub table: String, + /// Shown as the data explorer's tab title. + pub title: String, + /// The SQL that produced it, which Convert to Code wraps. + pub query: String, +} + +/// Retained query results, oldest first, bounded by `max`. +/// +/// Only results no data explorer has open count against the bound and are +/// evicted, so an open explorer never loses its table. +pub struct ResultStore { + reader: Box, + entries: VecDeque, + max: usize, +} + +impl ResultStore { + pub fn new(max: usize) -> Result { + Ok(Self { + reader: reader_from_uri("duckdb://memory")?, + entries: VecDeque::new(), + max: max.max(1), + }) + } + + /// The private database holding every retained result. + pub fn reader(&self) -> &dyn Reader { + &*self.reader + } + + /// Retain `df`, produced by `query`, under a fresh key, which is returned. + /// + /// `in_use` says whether an explorer has a result open; the oldest of the + /// rest are evicted to stay within the limit. + pub fn insert( + &mut self, + df: DataFrame, + title: String, + query: String, + in_use: impl Fn(&str) -> bool, + ) -> Result { + // Unguessable, so a stale link from a previous kernel process can never + // name a result of this one. + let key = uuid::Uuid::new_v4().simple().to_string(); + let table = format!("__ggsql_result_{}__", key); + self.reader.register(&table, df, false)?; + self.entries.push_back(StoredResult { + key: key.clone(), + table, + title, + query, + }); + + let mut idle = self.entries.iter().filter(|r| !in_use(&r.key)).count(); + let mut i = 0; + while idle > self.max && i < self.entries.len() { + if in_use(&self.entries[i].key) { + i += 1; + continue; + } + if let Some(old) = self.entries.remove(i) { + if let Err(e) = self.reader.unregister(&old.table) { + tracing::warn!("Failed to drop retained result {}: {}", old.key, e); + } + } + idle -= 1; + } + Ok(key) + } + + pub fn get(&self, key: &str) -> Result<&StoredResult> { + self.entries.iter().find(|r| r.key == key).ok_or_else(|| { + anyhow!("This query result is no longer available. Re-run the query to explore it.") + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn frame(n: i64) -> DataFrame { + ggsql::reader::connection::reader_from_uri("duckdb://memory") + .unwrap() + .execute_sql(&format!("SELECT {} AS x", n)) + .unwrap() + } + + #[test] + fn test_a_retained_result_is_queryable() { + let mut store = ResultStore::new(4).unwrap(); + let key = store + .insert(frame(7), "Query [1]".into(), String::new(), |_| false) + .unwrap(); + + let table = &store.get(&key).unwrap().table; + let df = store + .reader() + .execute_sql(&format!( + "SELECT x FROM {}", + ggsql::naming::quote_ident(table) + )) + .unwrap(); + assert_eq!(df.height(), 1); + } + + #[test] + fn test_the_oldest_result_is_evicted_past_the_limit() { + let mut store = ResultStore::new(2).unwrap(); + let first = store + .insert(frame(1), "a".into(), String::new(), |_| false) + .unwrap(); + let second = store + .insert(frame(2), "b".into(), String::new(), |_| false) + .unwrap(); + store + .insert(frame(3), "c".into(), String::new(), |_| false) + .unwrap(); + + assert!(store.get(&first).is_err()); + assert!(store.get(&second).is_ok()); + } + + #[test] + fn test_a_result_in_use_is_never_evicted() { + let mut store = ResultStore::new(1).unwrap(); + let open = store + .insert(frame(1), "a".into(), String::new(), |_| false) + .unwrap(); + let pinned = |k: &str| k == open; + let second = store + .insert(frame(2), "b".into(), String::new(), pinned) + .unwrap(); + let third = store + .insert(frame(3), "c".into(), String::new(), pinned) + .unwrap(); + + assert!(store.get(&open).is_ok()); + assert!(store.get(&second).is_err()); + assert!(store.get(&third).is_ok()); + } + + #[test] + fn test_an_unknown_key_says_the_result_is_gone() { + let store = ResultStore::new(2).unwrap(); + let err = store.get("nope").err().unwrap().to_string(); + assert!(err.contains("no longer available"), "{err}"); + } +} diff --git a/ggsql-jupyter/tests/test_integration.py b/ggsql-jupyter/tests/test_integration.py index a409ec156..e65119595 100644 --- a/ggsql-jupyter/tests/test_integration.py +++ b/ggsql-jupyter/tests/test_integration.py @@ -17,7 +17,7 @@ # the kernel is built once rather than once per test file. -def _launch(kernel_binary, extra_args=()): +def _launch(kernel_binary, extra_args=(), env=None): """Start a kernel and yield its manager. Shared by the fixtures below.""" kernel_process = None km = None @@ -37,6 +37,7 @@ def _launch(kernel_binary, extra_args=()): # Start our kernel process directly kernel_process = subprocess.Popen( [kernel_binary, "-f", connection_file, *extra_args], + env={**os.environ, **(env or {})}, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) @@ -507,6 +508,325 @@ def test_an_unknown_method_is_an_error_not_a_null_result(self, console_client): assert "result" not in data +class TestResultExplorer: + """A console query result, its link, and `-- @view-result` opening it.""" + + QUERY = "SELECT i AS n, 'r' || i AS label FROM range(5) t(i)" + + def _run_query(self, client): + """Run QUERY and return the iopub messages it produced.""" + client.execute(self.QUERY, silent=False, store_history=True) + messages = drain_iopub(client) + reply = client.get_shell_msg(timeout=5) + assert reply["content"]["status"] == "ok" + return messages + + def _link_key(self, client): + """Run QUERY in a console session and return the key its link names.""" + messages = self._run_query(client) + order = [m["msg_type"] for m in messages] + assert "execute_result" in order, order + + # Its own output after the table: the console renders the HTML and + # ignores the bundle's text/plain, so a link there would never show. + streams = [m for m in messages if m["msg_type"] == "stream"] + assert len(streams) == 1, order + assert order.index("execute_result") < order.index("stream") + + text = streams[0]["content"]["text"] + prefix = f"\x1b]8;;x-ggsql-view:{client.session.session}/" + assert text.startswith(prefix), repr(text) + assert "Open in Data Explorer" in text + return text[len(prefix):].split("\x1b", 1)[0] + + def _rpc(self, client, comm_id, method, params=None): + client.session.send( + client.shell_channel.socket, + "comm_msg", + { + "comm_id": comm_id, + "data": { + "jsonrpc": "2.0", + "id": f"{method}-1", + "method": method, + "params": params or {}, + }, + }, + ) + reply = client.get_shell_msg(timeout=5) + drain_iopub(client) + return reply["content"]["data"]["result"] + + def test_view_opens_the_result_in_an_explorer_tab(self, console_client): + key = self._link_key(console_client) + + # How the extension runs it when the link is clicked. + msg_id = console_client.execute(f"-- @view-result {key}", silent=True) + messages = drain_iopub(console_client) + reply = console_client.get_shell_msg(timeout=5) + assert reply["content"]["status"] == "ok", reply["content"] + + opens = [m for m in messages if m["msg_type"] == "comm_open"] + assert len(opens) == 1, [m["msg_type"] for m in messages] + content = opens[0]["content"] + assert content["target_name"] == "positron.dataExplorer" + # Not inline: Positron opens a tab for it straight away. + assert content["data"]["inline_only"] is False + assert opens[0]["parent_header"]["msg_id"] == msg_id + + comm_id = content["comm_id"] + state = self._rpc(console_client, comm_id, "get_state") + assert state["display_name"].startswith("Query [") + assert state["table_shape"] == {"num_rows": 5, "num_columns": 2} + + values = self._rpc( + console_client, + comm_id, + "get_data_values", + { + "columns": [ + {"column_index": 1, "spec": {"first_index": 3, "last_index": 4}} + ] + }, + ) + assert values["columns"] == [["r3", "r4"]] + + def _open(self, client): + """Open QUERY's result in an explorer and return the comm id.""" + key = self._link_key(client) + client.execute(f"-- @view-result {key}", silent=True) + messages = drain_iopub(client) + client.get_shell_msg(timeout=5) + return next( + m["content"]["comm_id"] for m in messages if m["msg_type"] == "comm_open" + ) + + def test_filter_sort_and_convert_to_code(self, console_client): + comm_id = self._open(console_client) + + state = self._rpc(console_client, comm_id, "get_state") + features = state["supported_features"] + for feature in ("set_row_filters", "set_sort_columns", "convert_to_code"): + assert features[feature]["support_status"] == "supported", feature + + n_column = {"column_index": 0} + reply = self._rpc( + console_client, + comm_id, + "set_row_filters", + { + "filters": [ + { + "filter_id": "f1", + "filter_type": "compare", + "column_schema": n_column, + "condition": "and", + "params": {"op": ">=", "value": "2"}, + } + ] + }, + ) + assert reply == {"selected_num_rows": 3, "had_errors": False} + + self._rpc( + console_client, + comm_id, + "set_sort_columns", + {"sort_keys": [{"column_index": 0, "ascending": False}]}, + ) + values = self._rpc( + console_client, + comm_id, + "get_data_values", + { + "columns": [ + {"column_index": 1, "spec": {"first_index": 0, "last_index": 2}} + ] + }, + ) + assert values["columns"] == [["r4", "r3", "r2"]] + + code = self._rpc( + console_client, + comm_id, + "convert_to_code", + { + "column_filters": [], + "row_filters": state["row_filters"], + "sort_keys": [{"column_index": 0, "ascending": False}], + "code_syntax_name": {"code_syntax_name": "SQL"}, + }, + ) + assert f" {self.QUERY}" in code["converted_code"] + assert 'ORDER BY "n" DESC' in code["converted_code"] + + def test_a_bad_request_is_a_json_rpc_error(self, console_client): + comm_id = self._open(console_client) + console_client.session.send( + console_client.shell_channel.socket, + "comm_msg", + { + "comm_id": comm_id, + "data": { + "jsonrpc": "2.0", + "id": "export-1", + "method": "export_data_selection", + "params": { + "format": "csv", + "selection": {"kind": "nonsense", "selection": {}}, + }, + }, + }, + ) + reply = console_client.get_shell_msg(timeout=5) + drain_iopub(console_client) + data = reply["content"]["data"] + assert "error" in data and "result" not in data, data + + def test_a_second_click_focuses_the_open_explorer(self, console_client): + # Positron opens a UI comm on every console session; focusing an + # editor is an event on it. + ui = "test-ui-comm" + console_client.session.send( + console_client.shell_channel.socket, + "comm_open", + {"comm_id": ui, "target_name": "positron.ui", "data": {}}, + ) + drain_iopub(console_client) + try: + key = self._link_key(console_client) + console_client.execute(f"-- @view-result {key}", silent=True) + first = drain_iopub(console_client) + console_client.get_shell_msg(timeout=5) + comm_id = next( + m["content"]["comm_id"] for m in first if m["msg_type"] == "comm_open" + ) + + console_client.execute(f"-- @view-result {key}", silent=True) + second = drain_iopub(console_client) + reply = console_client.get_shell_msg(timeout=5) + assert reply["content"]["status"] == "ok", reply["content"] + + types = [m["msg_type"] for m in second] + assert "comm_open" not in types, types + events = [ + m["content"] + for m in second + if m["msg_type"] == "comm_msg" and m["content"]["comm_id"] == ui + ] + assert len(events) == 1, types + assert events[0]["data"]["method"] == "open_editor" + params = events[0]["data"]["params"] + assert params["kind"] == "uri" + assert params["file"] == f"positron-data-explorer:positron-data-explorer-{comm_id}" + finally: + console_client.session.send( + console_client.shell_channel.socket, + "comm_close", + {"comm_id": ui, "data": {}}, + ) + drain_iopub(console_client) + + def test_a_result_that_is_gone_is_an_error(self, console_client): + console_client.execute("-- @view-result deadbeef", silent=True) + drain_iopub(console_client) + reply = console_client.get_shell_msg(timeout=5) + assert reply["content"]["status"] == "error" + assert "no longer available" in reply["content"]["evalue"] + + def test_only_a_console_gets_a_link(self, client): + messages = self._run_query(client) + assert not [m for m in messages if m["msg_type"] == "stream"] + + +class TestViewDirective: + """`-- @view` sending a query's result to the data explorer.""" + + CODE = "-- @view: Small numbers\nSELECT i AS n FROM range(4) t(i)" + + def _run(self, client, code, silent=False): + client.execute(code, silent=silent, store_history=not silent) + messages = drain_iopub(client) + reply = client.get_shell_msg(timeout=5) + return messages, reply["content"] + + def test_the_result_opens_in_an_explorer_tab(self, console_client): + messages, reply = self._run(console_client, self.CODE) + assert reply["status"] == "ok", reply + types = [m["msg_type"] for m in messages] + + # The explorer, and the table in the console as for any query. + opens = [m for m in messages if m["msg_type"] == "comm_open"] + assert len(opens) == 1, types + data = opens[0]["content"]["data"] + assert data == {"title": "Small numbers", "inline_only": False} + assert types.count("execute_result") == 1, types + + # And the link under the table, to reopen the tab once it is closed. + streams = [m for m in messages if m["msg_type"] == "stream"] + assert len(streams) == 1, types + assert "x-ggsql-view:" in streams[0]["content"]["text"] + assert types.index("execute_result") < types.index("stream"), types + + def test_a_silent_execution_opens_it_without_output(self, console_client): + messages, reply = self._run(console_client, self.CODE, silent=True) + assert reply["status"] == "ok", reply + types = [m["msg_type"] for m in messages] + assert types.count("comm_open") == 1, types + assert "stream" not in types and "execute_result" not in types, types + + def test_a_plot_query_is_an_error(self, console_client): + messages, reply = self._run( + console_client, + "-- @view\nSELECT 1 AS x, 2 AS y VISUALISE x, y DRAW point", + ) + assert reply["status"] == "error" + assert "not a plot" in reply["evalue"] + assert "comm_open" not in [m["msg_type"] for m in messages] + + def test_outside_positron_the_table_shows_inline(self, client): + messages, reply = self._run(client, self.CODE) + assert reply["status"] == "ok", reply + types = [m["msg_type"] for m in messages] + assert "execute_result" in types and "comm_open" not in types, types + notes = [ + m["content"]["text"] + for m in messages + if m["msg_type"] == "stream" and m["content"]["name"] == "stderr" + ] + assert notes and "needs Positron" in notes[0], types + + +class TestConsoleRows: + """`GGSQL_CONSOLE_ROWS`, which the extension sets from `ggsql.consoleRows`.""" + + def _table(self, client, n): + client.execute(f"SELECT i FROM range({n}) t(i)") + messages = drain_iopub(client) + client.get_shell_msg(timeout=5) + result = next(m for m in messages if m["msg_type"] == "execute_result") + return result["content"]["data"]["text/html"] + + def test_the_environment_sets_the_rows_a_console_shows(self, kernel_binary): + with _launch_kernel( + kernel_binary, + ["--session-mode", "console"], + env={"GGSQL_CONSOLE_ROWS": "3"}, + ) as km: + kc = km.client() + kc.start_channels() + try: + kc.wait_for_ready(timeout=10) + # Up to twice the setting, the whole result. + assert self._table(kc, 6).count("") == 1 + 6 + # Past it, the first 3 and a count of the rest. + html = self._table(kc, 7) + assert html.count("") == 1 + 3 + 1, html + assert "... 4 more rows" in html + finally: + kc.stop_channels() + + class TestStatus: """Test status messages.""" diff --git a/ggsql-vscode/CLAUDE.md b/ggsql-vscode/CLAUDE.md index aca8e758f..87d3cae72 100644 --- a/ggsql-vscode/CLAUDE.md +++ b/ggsql-vscode/CLAUDE.md @@ -19,6 +19,7 @@ ggsql-vscode/ │ ├── manager.ts Kernel discovery + Positron language-runtime registration │ ├── positronApi.ts Acquires the Positron API so Positron can attribute it to this extension │ ├── connections.ts Connection-string handling for the Connections pane +│ ├── resultLinks.ts Session wrapper handling "Open in Data Explorer" links │ ├── cellParser.ts Splits .ggsql files into cells for Run-Cell commands │ ├── codelens.ts "▶ Run cell" lens above each cell │ ├── decorations.ts Cell separator decorations @@ -99,7 +100,11 @@ The extension declares `contributes.languageRuntimes` for `ggsql` (see `package. - **`createKernelSpec` appends `--session-mode `** from `sessionMetadata.sessionMode`. The kernel keys every output decision off that: a console session gets the comm, a notebook session gets a static image in its cell, and background gets the standalone path. Without the flag the kernel falls back to guessing from its session id, which works but is a heuristic. `restoreSession` deliberately does *not* rebuild the spec — the supervisor replays the argv the session was created with, and a session's mode never changes. - **`uiSubscriptions: ['did_change_plots_render_settings']`** on the runtime metadata. Without it the frontend never tells the kernel how large the Plots pane is, and the kernel needs that for exactly one thing: rendering a *new* plot at the right size so its `comm_open` can carry it, and the pane shows it immediately instead of blank. Every other render carries its own size on the request. -Both live in `manager.ts`. The kernel side of the contract — why the comm alone creates the pane entry, and why an output bundle alongside it would show the plot twice — is in [`/ggsql-jupyter/CLAUDE.md`](../ggsql-jupyter/CLAUDE.md). +Both live in `manager.ts`, as does one piece of user configuration: **`ggsql.consoleRows` reaches the kernel as `GGSQL_CONSOLE_ROWS` in the spec's `env`, not as `--console-rows`.** A kernel found on the machine may predate the option, and it ignores an unknown variable where it would refuse to start on an unknown flag. `consoleRowsEnv` leaves out a value the kernel would reject, for the same reason. Like the session mode, it is fixed when a session starts. + +The kernel side of the contract — why the comm alone creates the pane entry, and why an output bundle alongside it would show the plot twice — is in [`/ggsql-jupyter/CLAUDE.md`](../ggsql-jupyter/CLAUDE.md). + +**Sessions are wrapped to add `openResource`.** `createSession` and `restoreSession` return the supervisor's session through `withResultLinks` (`src/resultLinks.ts`), a `Proxy` that adds `openResource` for the `x-ggsql-view:` links the kernel prints under console results, and binds everything else to the underlying session. Positron offers a clicked console hyperlink only to the *last active* console of each language, so the handler routes by the session id in the link rather than assuming it is the session it wraps. The kernel side is in [`/ggsql-jupyter/CLAUDE.md`](../ggsql-jupyter/CLAUDE.md#opening-a-console-result-in-the-data-explorer). Outside Positron there is no way to execute a query: `activate()` returns early, so every command that runs code stays unregistered. To avoid offering actions that cannot work, everything execution-related gates on Positron's built-in **`isPositron`** context key ([extension development docs](https://positron.posit.co/extension-development.html#option-1-context-keys)): @@ -151,7 +156,8 @@ The one case that interrupts the user is the dead end: nothing runnable anywhere ```json { - "ggsql.kernelPath": "string" // an extra kernel to offer, listed first + "ggsql.kernelPath": "string", // an extra kernel to offer, listed first + "ggsql.consoleRows": 10 // rows of a long result a console shows } ``` diff --git a/ggsql-vscode/package.json b/ggsql-vscode/package.json index 699e2cb57..3a85c4216 100644 --- a/ggsql-vscode/package.json +++ b/ggsql-vscode/package.json @@ -164,6 +164,12 @@ "default": "", "markdownDescription": "Path to a `ggsql-jupyter` executable to offer alongside the kernels found automatically, listed ahead of them. A bare name is looked up on `PATH`. Leave this empty unless you have a kernel in a place the extension does not look: it ships one, and it finds the ones installed on this machine." }, + "ggsql.consoleRows": { + "type": "integer", + "default": 10, + "minimum": 1, + "markdownDescription": "How many rows of a query result the ggsql console shows. A result with up to twice this many rows is shown in full; a longer one shows this many, then a count of the rest and a link to open the whole result in the Data Explorer. Notebooks are unaffected. Takes effect in consoles started after the change." + }, "ggsql.enableSqlFiles": { "type": "boolean", "default": true, diff --git a/ggsql-vscode/src/manager.ts b/ggsql-vscode/src/manager.ts index 83585c910..e16200340 100644 --- a/ggsql-vscode/src/manager.ts +++ b/ggsql-vscode/src/manager.ts @@ -13,6 +13,7 @@ import * as crypto from 'crypto'; import type * as positron from '@posit-dev/positron'; import type { JupyterKernelSpec, PositronSupervisorApi } from './types'; import { log, showLog } from './extension'; +import { withResultLinks } from './resultLinks'; /** Where a kernel candidate was discovered */ type KernelSource = 'Bundled' | 'Setting' | 'Jupyter' | 'System' | 'Path'; @@ -534,6 +535,19 @@ export function generateMetadata( }; } +/** + * The kernel environment for the `ggsql.consoleRows` setting. + * + * An environment variable rather than a flag: a kernel that predates + * `--console-rows` ignores the variable, where it would refuse to start on the + * flag. A value the kernel would reject is left out, so it uses its default. + */ +export function consoleRowsEnv(rows: unknown): Record { + return Number.isInteger(rows) && (rows as number) >= 1 + ? { GGSQL_CONSOLE_ROWS: String(rows) } + : {}; +} + /** * Create a Jupyter kernel spec for ggsql-jupyter * @@ -566,7 +580,10 @@ function createKernelSpec( display_name: 'ggsql', language: 'ggsql', interrupt_mode: 'signal', - env: { RUST_LOG: 'error' }, + env: { + RUST_LOG: 'error', + ...consoleRowsEnv(vscode.workspace.getConfiguration('ggsql').get('consoleRows')), + }, kernel_protocol_version: '5.3', }; } @@ -858,7 +875,7 @@ export class GgsqlRuntimeManager implements positron.LanguageRuntimeManager { dynState ); - return session; + return withResultLinks(session); } /** @@ -885,7 +902,7 @@ export class GgsqlRuntimeManager implements positron.LanguageRuntimeManager { dynState ); - return session; + return withResultLinks(session); } /** diff --git a/ggsql-vscode/src/resultLinks.ts b/ggsql-vscode/src/resultLinks.ts new file mode 100644 index 000000000..f75ccda1d --- /dev/null +++ b/ggsql-vscode/src/resultLinks.ts @@ -0,0 +1,91 @@ +/* + * "Open in Data Explorer" links under console query results. + * + * In a Positron console the kernel follows each tabular result with an OSC 8 + * hyperlink to `x-ggsql-view:/`. Positron offers a + * clicked hyperlink to the last active console session of each language via + * `openResource`, which the supervisor's session does not implement — so each + * ggsql session is wrapped to add it. Handling the link runs `-- @view-result ` + * silently in the session that printed it, and the kernel opens the result in + * the data explorer. + */ + +import * as vscode from 'vscode'; +import type * as positron from '@posit-dev/positron'; +import { getPositronApi } from './positronApi'; +import { log } from './extension'; + +export const VIEW_SCHEME = 'x-ggsql-view'; + +/** A link's target, or undefined when the resource is not one of ours. */ +export function parseViewLink(resource: vscode.Uri | string): { sessionId: string; key: string } | undefined { + const uri = typeof resource === 'string' ? vscode.Uri.parse(resource) : resource; + if (uri.scheme !== VIEW_SCHEME) { + return undefined; + } + const [sessionId, key, ...rest] = uri.path.split('/'); + // The key goes into code the kernel executes, so only accept the shape the + // kernel produces. + if (!sessionId || !key || rest.length > 0 || !/^[0-9a-f]+$/.test(key)) { + return undefined; + } + return { sessionId, key }; +} + +/** Open a result in the data explorer by running `-- @view-result` in its session. */ +async function openResult(sessionId: string, key: string): Promise { + const api = getPositronApi(); + if (!api) { + return; + } + try { + await api.runtime.executeCode( + 'ggsql', + `-- @view-result ${key}`, + false, + false, + api.RuntimeCodeExecutionMode.Silent, + api.RuntimeErrorBehavior.Continue, + undefined, + sessionId + ); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + log(`Opening result ${key} in session ${sessionId} failed: ${message}`); + vscode.window.showWarningMessage( + message || 'This query result could not be opened in the Data Explorer.' + ); + } +} + +/** + * Wrap a supervisor session so it handles `x-ggsql-view:` links. + * + * Everything else passes through to the session itself, with methods bound to + * it so its own `this` is never the proxy. + */ +export function withResultLinks( + session: positron.LanguageRuntimeSession +): positron.LanguageRuntimeSession { + const openResource = async (resource: vscode.Uri | string): Promise => { + const link = parseViewLink(resource); + if (!link) { + return false; + } + // Positron offers the link only to the most recently active ggsql + // console, which need not be the one that printed it; the session id + // routes it back to the right one. + void openResult(link.sessionId, link.key); + return true; + }; + + return new Proxy(session, { + get(target, prop) { + if (prop === 'openResource') { + return openResource; + } + const value = Reflect.get(target, prop, target); + return typeof value === 'function' ? value.bind(target) : value; + }, + }); +} diff --git a/ggsql-vscode/src/test/manager.test.ts b/ggsql-vscode/src/test/manager.test.ts index 812b0daaf..4caff9db0 100644 --- a/ggsql-vscode/src/test/manager.test.ts +++ b/ggsql-vscode/src/test/manager.test.ts @@ -1,6 +1,6 @@ import * as assert from 'assert'; import * as vscode from 'vscode'; -import { GgsqlRuntimeManager, createDynState, getSupervisorApi } from '../manager'; +import { GgsqlRuntimeManager, consoleRowsEnv, createDynState, getSupervisorApi } from '../manager'; suite('manager', () => { // The suites run in stock VS Code, where positron.positron-supervisor is @@ -13,6 +13,17 @@ suite('manager', () => { ); }); + test('consoleRowsEnv passes a row count the kernel accepts', () => { + assert.deepStrictEqual(consoleRowsEnv(25), { GGSQL_CONSOLE_ROWS: '25' }); + }); + + test('consoleRowsEnv leaves out a value the kernel would refuse', () => { + // The kernel exits on a bad value, so the session would fail to start. + for (const rows of [0, -1, 2.5, '10', undefined, null]) { + assert.deepStrictEqual(consoleRowsEnv(rows), {}, String(rows)); + } + }); + test('createDynState falls back to the default session name', () => { const state = createDynState(); assert.strictEqual(state.sessionName, 'ggsql'); diff --git a/ggsql-vscode/src/test/resultLinks.test.ts b/ggsql-vscode/src/test/resultLinks.test.ts new file mode 100644 index 000000000..4c875fb83 --- /dev/null +++ b/ggsql-vscode/src/test/resultLinks.test.ts @@ -0,0 +1,24 @@ +import * as assert from 'assert'; +import { parseViewLink } from '../resultLinks'; + +suite('parseViewLink', () => { + test('reads the session id and result key the kernel writes', () => { + assert.deepStrictEqual( + parseViewLink('x-ggsql-view:ggsql-1a2b/0f9e8d'), + { sessionId: 'ggsql-1a2b', key: '0f9e8d' } + ); + }); + + test('ignores other schemes', () => { + assert.strictEqual(parseViewLink('x-r-help:topic'), undefined); + assert.strictEqual(parseViewLink('https://ggsql.org'), undefined); + }); + + test('rejects a key the kernel could not have produced', () => { + // The key is spliced into code the kernel runs. + assert.strictEqual(parseViewLink('x-ggsql-view:s/abc%0ASELECT%201'), undefined); + assert.strictEqual(parseViewLink('x-ggsql-view:s/ABC'), undefined); + assert.strictEqual(parseViewLink('x-ggsql-view:s'), undefined); + assert.strictEqual(parseViewLink('x-ggsql-view:s/abc/def'), undefined); + }); +});