diff --git a/Cargo.toml b/Cargo.toml index beaa88a36..f8378b578 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -64,6 +64,9 @@ palette = { version = "0.7", default-features = false, features = ["std", "appro # Spatial geozero = { version = "0.14", default-features = false } +# Online data (download + cache) +ureq = "3" + # Utilities regex = "1.10" chrono = "0.4" diff --git a/ggsql-wasm/src/lib.rs b/ggsql-wasm/src/lib.rs index 6c7c16c4e..015fbffe0 100644 --- a/ggsql-wasm/src/lib.rs +++ b/ggsql-wasm/src/lib.rs @@ -320,8 +320,8 @@ impl GgsqlContext { /// Register all known builtin datasets (e.g. ggsql:penguins) pub async fn register_builtin_datasets(&self) -> Result<(), JsValue> { - for &name in ggsql::reader::data::KNOWN_DATASETS { - if let Some(bytes) = ggsql::reader::data::builtin_parquet_bytes(name) { + for &name in ggsql::reader::builtin_data::KNOWN_DATASETS { + if let Some(bytes) = ggsql::reader::builtin_data::builtin_parquet_bytes(name) { let table_name = ggsql::naming::builtin_data_table(name); let columns_js = convert_parquet_js(bytes).await.map_err(|e| { JsValue::from_str(&format!("Parquet error for '{}': {:?}", name, e)) @@ -336,6 +336,23 @@ impl GgsqlContext { Ok(()) } + // TODO: implement `register_online_datasets` to support `online:world` etc. + // + // Expected signature: + // pub async fn register_online_datasets(&self, sql: &str) -> Result<(), JsValue> + // + // Implementation: + // 1. Extract names via: + // ggsql::reader::builtin_data::extract_prefixed_dataset_names(sql, "online") + // 2. For each name, skip if already registered (check table_exists on reader). + // 3. Resolve URL via ggsql::reader::online_data::resolve_online_dataset(name). + // 4. Fetch bytes via browser fetch() (wasm_bindgen_futures + web_sys::Request). + // 5. Parse with convert_parquet_js(bytes).await + columns_js_to_dataframe(). + // 6. Register under ggsql::naming::online_data_table(name). + // + // Then call this from `execute()` and `execute_sql()` before the reader runs, + // which requires making those methods async (or calling from JS before execute). + /// Load a previously installed SQLite extension. /// /// `entry_point` is the C init function name. If omitted, SQLite diff --git a/src/Cargo.toml b/src/Cargo.toml index 95614b06d..a2dce184a 100644 --- a/src/Cargo.toml +++ b/src/Cargo.toml @@ -39,6 +39,9 @@ adbc_core = { version = "0.23", optional = true } # Spatial geozero = { workspace = true, optional = true, features = ["with-wkb", "with-wkt", "with-geojson"] } +# Online data (download + cache) — native only, pulled in with parquet +ureq = { workspace = true, optional = true } + # Serialization serde.workspace = true serde_json.workspace = true @@ -57,14 +60,13 @@ uuid.workspace = true [dev-dependencies] jsonschema = { version = "0.44", default-features = false, features = ["resolve-file"] } tempfile = "3.8" -ureq = "3" adbc_datafusion = "0.23" adbc_driver_manager = "0.23" [features] default = ["adbc", "duckdb", "sqlite", "vegalite", "parquet", "builtin-data", "odbc", "spatial"] duckdb = ["dep:duckdb"] -parquet = ["dep:parquet"] +parquet = ["dep:parquet", "dep:ureq"] sqlite = ["dep:rusqlite"] adbc = ["dep:adbc_core"] odbc = ["dep:toml_edit", "dep:libloading"] diff --git a/src/naming.rs b/src/naming.rs index 6259f5bd3..0b48b4b44 100644 --- a/src/naming.rs +++ b/src/naming.rs @@ -152,6 +152,26 @@ pub fn builtin_data_table(name: &str) -> String { format!("{}{}{}", DATA_PREFIX, name, GGSQL_SUFFIX) } +/// Generate table name for an online dataset. +/// +/// Used when rewriting `online:world` to the internal table name. +/// Format: `__ggsql_online___` +/// +/// # Example +/// ``` +/// use ggsql::naming; +/// assert_eq!(naming::online_data_table("world"), "__ggsql_online_world__"); +/// assert_eq!(naming::online_data_table("us-counties"), "__ggsql_online_us_counties__"); +/// ``` +pub fn online_data_table(name: &str) -> String { + format!( + "{}online_{}{}", + GGSQL_PREFIX, + name.replace('-', "_"), + GGSQL_SUFFIX + ) +} + /// Generate column name for a constant aesthetic value. /// /// Used when a single layer has a literal aesthetic value that needs diff --git a/src/reader/data.rs b/src/reader/builtin_data.rs similarity index 86% rename from src/reader/data.rs rename to src/reader/builtin_data.rs index abf9e9604..c447cce48 100644 --- a/src/reader/data.rs +++ b/src/reader/builtin_data.rs @@ -64,36 +64,27 @@ pub fn builtin_parquet_bytes(name: &str) -> Option<&'static [u8]> { // Arrow-based builtin data loading // ============================================================================= -#[cfg(all(feature = "builtin-data", feature = "parquet"))] -pub fn load_builtin_dataframe(name: &str) -> Result { +/// Convert raw parquet bytes into a DataFrame. +#[cfg(feature = "parquet")] +pub fn dataframe_from_parquet_bytes( + label: &str, + bytes: bytes::Bytes, +) -> Result { use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; - let parquet_bytes = match name { - "penguins" => PENGUINS, - "airquality" => AIRQUALITY, - "world" => WORLD, - _ => { - return Err(GgsqlError::ReaderError(format!( - "Unknown builtin dataset: '{}'", - name - ))) - } - }; - - let bytes = bytes::Bytes::from_static(parquet_bytes); let reader = ParquetRecordBatchReaderBuilder::try_new(bytes) .map_err(|e| { - GgsqlError::ReaderError(format!("Failed to read builtin dataset '{}': {}", name, e)) + GgsqlError::ReaderError(format!("Failed to read parquet for '{}': {}", label, e)) })? .build() .map_err(|e| { - GgsqlError::ReaderError(format!("Failed to build reader for '{}': {}", name, e)) + GgsqlError::ReaderError(format!("Failed to build reader for '{}': {}", label, e)) })?; let batches: Vec<_> = reader .collect::, _>>() .map_err(|e| { - GgsqlError::ReaderError(format!("Failed to load builtin dataset '{}': {}", name, e)) + GgsqlError::ReaderError(format!("Failed to load dataset '{}': {}", label, e)) })?; if batches.is_empty() { @@ -104,13 +95,30 @@ pub fn load_builtin_dataframe(name: &str) -> Result Result { + let parquet_bytes = match name { + "penguins" => PENGUINS, + "airquality" => AIRQUALITY, + "world" => WORLD, + _ => { + return Err(GgsqlError::ReaderError(format!( + "Unknown builtin dataset: '{}'", + name + ))) + } + }; + + dataframe_from_parquet_bytes(name, bytes::Bytes::from_static(parquet_bytes)) +} + /// Known builtin dataset names in the ggsql namespace pub const KNOWN_DATASETS: &[&str] = &["penguins", "airquality", "world"]; @@ -123,11 +131,11 @@ pub fn is_known_builtin(name: &str) -> bool { // SQL namespace rewriting (always available, including WASM) // ============================================================================= -/// Extract builtin dataset names from SQL containing namespaced identifiers. +/// Extract dataset names with a given namespace prefix from SQL. /// -/// Finds `ggsql:X` patterns via tree-sitter and returns the dataset names -/// (without the `ggsql:` prefix), deduplicated. -pub fn extract_builtin_dataset_names(sql: &str) -> Result, GgsqlError> { +/// Finds `prefix:X` patterns via tree-sitter and returns the dataset names +/// (without the prefix), deduplicated. +pub fn extract_prefixed_dataset_names(sql: &str, prefix: &str) -> Result, GgsqlError> { let token_def = r#"(namespaced_identifier) @select"#; let mut tokens = tokens_from_tree(sql, token_def, "select")?; @@ -138,9 +146,10 @@ pub fn extract_builtin_dataset_names(sql: &str) -> Result, GgsqlErro tokens.sort_unstable(); tokens.dedup(); + let prefix_colon = format!("{}:", prefix); let datasets: Vec = tokens .iter() - .filter_map(|token| token.strip_prefix("ggsql:").map(|s| s.to_string())) + .filter_map(|token| token.strip_prefix(&prefix_colon).map(|s| s.to_string())) .collect(); Ok(datasets) @@ -148,7 +157,9 @@ pub fn extract_builtin_dataset_names(sql: &str) -> Result, GgsqlErro /// Rewrite SQL to replace namespaced identifiers with internal table names. /// -/// e.g., `SELECT * FROM ggsql:penguins` -> `SELECT * FROM __ggsql_data_penguins__` +/// Handles both builtin and online dataset prefixes: +/// - `ggsql:penguins` -> `__ggsql_data_penguins__` +/// - `online:world` -> `__ggsql_online_world__` /// /// Uses tree-sitter to find the exact byte positions of namespaced identifiers, /// then replaces them in reverse order to preserve offsets. @@ -190,6 +201,12 @@ pub fn rewrite_namespaced_sql(sql: &str) -> Result { node.end_byte(), naming::quote_ident(&naming::builtin_data_table(name)), )); + } else if let Some(name) = full_text.strip_prefix("online:") { + replacements.push(( + node.start_byte(), + node.end_byte(), + naming::quote_ident(&naming::online_data_table(name)), + )); } } } @@ -286,7 +303,7 @@ mod tests { #[test] fn test_extract_builtin_dataset_names_single() { let sql = "SELECT * FROM ggsql:penguins VISUALISE DRAW point MAPPING x AS x"; - let names = extract_builtin_dataset_names(sql).unwrap(); + let names = extract_prefixed_dataset_names(sql, "ggsql").unwrap(); assert_eq!(names, vec!["penguins"]); } @@ -294,7 +311,7 @@ mod tests { fn test_extract_builtin_dataset_names_multiple() { let sql = "SELECT * FROM ggsql:penguins, ggsql:airquality VISUALISE DRAW point MAPPING x AS x"; - let names = extract_builtin_dataset_names(sql).unwrap(); + let names = extract_prefixed_dataset_names(sql, "ggsql").unwrap(); assert_eq!(names.len(), 2); assert!(names.contains(&"airquality".to_string())); assert!(names.contains(&"penguins".to_string())); @@ -303,14 +320,14 @@ mod tests { #[test] fn test_extract_builtin_dataset_names_dedup() { let sql = "SELECT * FROM ggsql:penguins p1, ggsql:penguins p2 VISUALISE DRAW point MAPPING x AS x"; - let names = extract_builtin_dataset_names(sql).unwrap(); + let names = extract_prefixed_dataset_names(sql, "ggsql").unwrap(); assert_eq!(names, vec!["penguins"]); } #[test] fn test_extract_builtin_dataset_names_none() { let sql = "SELECT * FROM regular_table VISUALISE DRAW point MAPPING x AS x"; - let names = extract_builtin_dataset_names(sql).unwrap(); + let names = extract_prefixed_dataset_names(sql, "ggsql").unwrap(); assert!(names.is_empty()); } @@ -345,6 +362,32 @@ mod tests { assert!(rewritten.starts_with("SELECT * FROM \"__ggsql_data_penguins__\"")); assert!(!rewritten.contains("ggsql:")); } + + #[test] + fn test_rewrite_online_dataset() { + let sql = "SELECT * FROM online:world"; + let rewritten = rewrite_namespaced_sql(sql).unwrap(); + assert_eq!(rewritten, "SELECT * FROM \"__ggsql_online_world__\""); + } + + #[test] + fn test_rewrite_mixed_namespaces() { + // JOIN not yet supported by tree-sitter grammar for namespaced identifiers; + // comma-separated FROM works. + let sql = "SELECT * FROM ggsql:penguins, online:world"; + let rewritten = rewrite_namespaced_sql(sql).unwrap(); + assert!(rewritten.contains("\"__ggsql_data_penguins__\"")); + assert!(rewritten.contains("\"__ggsql_online_world__\"")); + assert!(!rewritten.contains("ggsql:")); + assert!(!rewritten.contains("online:")); + } + + #[test] + fn test_extract_online_dataset_names() { + let sql = "SELECT * FROM online:world VISUALISE DRAW polygon MAPPING geom AS geometry"; + let names = extract_prefixed_dataset_names(sql, "online").unwrap(); + assert_eq!(names, vec!["world"]); + } } #[cfg(all(feature = "duckdb", feature = "builtin-data"))] diff --git a/src/reader/duckdb.rs b/src/reader/duckdb.rs index bf459c155..1c4526bce 100644 --- a/src/reader/duckdb.rs +++ b/src/reader/duckdb.rs @@ -25,7 +25,7 @@ use std::sync::Arc; fn register_builtin_datasets_duckdb(sql: &str, conn: &Connection) -> Result<()> { use std::{env, fs}; - let dataset_names = super::data::extract_builtin_dataset_names(sql)?; + let dataset_names = super::builtin_data::extract_prefixed_dataset_names(sql, "ggsql")?; // Load spatial extension before registering datasets that contain // geometry columns, so that spatial features are available. @@ -34,7 +34,7 @@ fn register_builtin_datasets_duckdb(sql: &str, conn: &Connection) -> Result<()> } for name in dataset_names { - let Some(parquet_bytes) = super::data::builtin_parquet_bytes(&name) else { + let Some(parquet_bytes) = super::builtin_data::builtin_parquet_bytes(&name) else { continue; }; @@ -355,8 +355,21 @@ impl Reader for DuckDBReader { #[cfg(feature = "builtin-data")] register_builtin_datasets_duckdb(sql, &self.conn)?; + // Register online datasets if referenced + #[cfg(feature = "parquet")] + { + let online_names = super::builtin_data::extract_prefixed_dataset_names(sql, "online")?; + for name in &online_names { + let table_name = naming::online_data_table(name); + if !self.table_exists(&table_name)? { + let df = super::online_data::load_online_dataframe(name)?; + self.register(&table_name, df, true)?; + } + } + } + // Rewrite ggsql:name → __ggsql_data_name__ in SQL - let sql = super::data::rewrite_namespaced_sql(sql)?; + let sql = super::builtin_data::rewrite_namespaced_sql(sql)?; if !super::returns_rows(&sql) { self.conn diff --git a/src/reader/mod.rs b/src/reader/mod.rs index a6609d0c8..91368e9f6 100644 --- a/src/reader/mod.rs +++ b/src/reader/mod.rs @@ -429,8 +429,9 @@ pub mod odbc; #[cfg(feature = "adbc")] pub mod adbc; +pub mod builtin_data; pub mod connection; -pub mod data; +pub mod online_data; mod spec; #[cfg(feature = "duckdb")] diff --git a/src/reader/online_data.rs b/src/reader/online_data.rs new file mode 100644 index 000000000..8e46881fc --- /dev/null +++ b/src/reader/online_data.rs @@ -0,0 +1,265 @@ +use const_format::concatcp; + +use crate::{DataFrame, GgsqlError}; + +// ============================================================================= +// Dataset registry +// ============================================================================= + +/// Resolve an online dataset name to its download URL. +pub fn resolve_online_dataset(name: &str) -> Option<&'static str> { + let name = name.replace('-', "_"); + const BASE: &str = "https://github.com/ggsql-dev/datasets/releases/latest/download"; + let url = match name.as_str() { + "world" | "world_110m" | "countries" | "countries_110m" => { + concatcp!(BASE, "/ne_110m_admin_0_countries.parquet") + } + "world_50m" | "countries_50m" => { + concatcp!(BASE, "/ne_50m_admin_0_countries.parquet") + } + "world_10m" | "countries_10m" => { + concatcp!(BASE, "/ne_10m_admin_0_countries.parquet") + } + // Only US is included in the 110m dataset + "us_states" | "us_states_110m" => { + concatcp!(BASE, "/ne_110m_admin_1_states_provinces.parquet") + } + // Only includes large countries Russia, US, Canada, Brazil, China, and some others + // Notably missing are many European, African, South American countries + "states_50m" | "provinces_50m" => { + concatcp!(BASE, "/ne_50m_admin_1_states_provinces.parquet") + } + // Includes all countries + "states_10m" | "provinces_10m" => { + concatcp!(BASE, "/ne_10m_admin_1_states_provinces.parquet") + } + "us_counties" | "us_counties_10m" => { + concatcp!(BASE, "/ne_10m_admin_2_counties.parquet") + } + _ => return None, + }; + Some(url) +} + +// ============================================================================= +// Native download + cache (not available on wasm32) +// ============================================================================= +// +// Wasm integration lives in `ggsql-wasm/src/lib.rs` and must provide an +// equivalent of `load_online_dataframe` using the browser fetch API: +// +// 1. Call `resolve_online_dataset(name)` to get the URL. +// 2. Fetch the bytes via browser `fetch()` (async). +// 3. Parse parquet via `convert_parquet_js` + `columns_js_to_dataframe` +// (same as `register_parquet` / `register_builtin_datasets`). +// 4. Register the DataFrame under `naming::online_data_table(name)`. +// +// The method should be called `register_online_datasets` on `GgsqlWasm`, +// following the shape of `register_builtin_datasets`. It should extract +// `online:` dataset names from the SQL (via +// `builtin_data::extract_prefixed_dataset_names(sql, "online")`), skip +// any that are already registered, and fetch + register the rest. +// +// Caching: the native path uses the filesystem. In the browser, consider +// the Cache API or simply re-fetch per session. This is a UX decision. + +#[cfg(all(not(target_arch = "wasm32"), feature = "parquet"))] +mod native { + use std::fs; + use std::path::PathBuf; + + use super::*; + + pub(super) fn cache_dir() -> Result { + let dir = platform_cache_dir().join("ggsql").join("online"); + fs::create_dir_all(&dir).map_err(|e| { + GgsqlError::ReaderError(format!( + "Failed to create cache directory '{}': {}", + dir.display(), + e + )) + })?; + Ok(dir) + } + + fn platform_cache_dir() -> PathBuf { + #[cfg(target_os = "linux")] + { + if let Ok(dir) = std::env::var("XDG_CACHE_HOME") { + return PathBuf::from(dir); + } + if let Ok(home) = std::env::var("HOME") { + return PathBuf::from(home).join(".cache"); + } + } + + #[cfg(target_os = "macos")] + { + if let Ok(home) = std::env::var("HOME") { + return PathBuf::from(home).join("Library").join("Caches"); + } + } + + #[cfg(target_os = "windows")] + { + if let Ok(dir) = std::env::var("LOCALAPPDATA") { + return PathBuf::from(dir); + } + } + + std::env::temp_dir() + } + + fn cache_filename(url: &str) -> String { + url.rsplit('/') + .next() + .unwrap_or("dataset.parquet") + .to_string() + } + + fn ensure_downloaded(url: &str) -> Result { + let filename = cache_filename(url); + let parquet_path = cache_dir()?.join(&filename); + + if parquet_path.exists() { + return Ok(parquet_path); + } + + let partial_path = parquet_path.with_extension("parquet.partial"); + + // Clean up prior failed attempt + let _ = fs::remove_file(&partial_path); + + let response = ureq::get(url).call().map_err(|e| { + GgsqlError::ReaderError(format!( + "Failed to download '{}': {}. Are you connected to the internet?", + url, e + )) + })?; + + let bytes = response.into_body().read_to_vec().map_err(|e| { + GgsqlError::ReaderError(format!("Failed to read response from '{}': {}", url, e)) + })?; + + fs::write(&partial_path, &bytes).map_err(|e| { + GgsqlError::ReaderError(format!( + "Failed to write cache file '{}': {}", + partial_path.display(), + e + )) + })?; + + fs::rename(&partial_path, &parquet_path).map_err(|e| { + GgsqlError::ReaderError(format!( + "Failed to finalize cache file '{}': {}", + parquet_path.display(), + e + )) + })?; + + Ok(parquet_path) + } + + /// Load an online dataset by name, downloading and caching as needed. + pub fn load_online_dataframe(name: &str) -> Result { + let url = resolve_online_dataset(name).ok_or_else(|| { + GgsqlError::ReaderError(format!("Unknown online dataset: '{}'", name)) + })?; + + let parquet_path = ensure_downloaded(url)?; + + let bytes = fs::read(&parquet_path).map_err(|e| { + GgsqlError::ReaderError(format!( + "Failed to read cached parquet '{}': {}", + parquet_path.display(), + e + )) + })?; + + crate::reader::builtin_data::dataframe_from_parquet_bytes(name, bytes::Bytes::from(bytes)) + } +} + +#[cfg(all(not(target_arch = "wasm32"), feature = "parquet"))] +pub use native::load_online_dataframe; + +// ============================================================================= +// Tests +// ============================================================================= + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_resolve_world() { + let url = resolve_online_dataset("world").unwrap(); + assert!(url.contains("110m_admin_0_countries")); + } + + #[test] + fn test_resolve_world_50m() { + let url = resolve_online_dataset("world_50m").unwrap(); + assert!(url.contains("50m_admin_0_countries")); + } + + #[test] + fn test_resolve_world_10m() { + let url = resolve_online_dataset("world_10m").unwrap(); + assert!(url.contains("10m_admin_0_countries")); + } + + #[test] + fn test_resolve_countries_alias() { + let world = resolve_online_dataset("world").unwrap(); + let countries = resolve_online_dataset("countries").unwrap(); + assert_eq!(world, countries); + } + + #[test] + fn test_resolve_us_states() { + let url = resolve_online_dataset("us_states").unwrap(); + assert!(url.contains("110m_admin_1_states_provinces")); + } + + #[test] + fn test_resolve_states_10m() { + let url = resolve_online_dataset("states_10m").unwrap(); + assert!(url.contains("10m_admin_1_states_provinces")); + } + + #[test] + fn test_resolve_provinces_alias() { + let states_10m = resolve_online_dataset("states_10m").unwrap(); + let provinces_10m = resolve_online_dataset("provinces_10m").unwrap(); + assert_eq!(states_10m, provinces_10m); + } + + #[test] + fn test_resolve_us_counties() { + let url = resolve_online_dataset("us_counties").unwrap(); + assert!(url.contains("10m_admin_2_counties")); + } + + #[test] + fn test_resolve_us_counties_hyphen() { + let url = resolve_online_dataset("us-counties").unwrap(); + assert!(url.contains("10m_admin_2_counties")); + } + + #[test] + fn test_resolve_unknown() { + assert!(resolve_online_dataset("nonexistent").is_none()); + } + + #[cfg(all(not(target_arch = "wasm32"), feature = "parquet"))] + mod native_tests { + use super::super::native; + + #[test] + fn test_cache_dir_exists() { + let dir = native::cache_dir().unwrap(); + assert!(dir.exists()); + } + } +} diff --git a/src/reader/sqlite.rs b/src/reader/sqlite.rs index 8f1e1f9f6..b5b85f1c3 100644 --- a/src/reader/sqlite.rs +++ b/src/reader/sqlite.rs @@ -371,18 +371,31 @@ impl Reader for SqliteReader { // Handle ggsql:name namespaced identifiers (builtin datasets) #[cfg(all(feature = "builtin-data", feature = "parquet"))] { - let dataset_names = super::data::extract_builtin_dataset_names(sql)?; + let dataset_names = super::builtin_data::extract_prefixed_dataset_names(sql, "ggsql")?; for name in &dataset_names { let table_name = naming::builtin_data_table(name); if !self.table_exists(&table_name) { - let df = super::data::load_builtin_dataframe(name)?; + let df = super::builtin_data::load_builtin_dataframe(name)?; + self.register(&table_name, df, true)?; + } + } + } + + // Register online datasets if referenced + #[cfg(feature = "parquet")] + { + let online_names = super::builtin_data::extract_prefixed_dataset_names(sql, "online")?; + for name in &online_names { + let table_name = naming::online_data_table(name); + if !self.table_exists(&table_name) { + let df = super::online_data::load_online_dataframe(name)?; self.register(&table_name, df, true)?; } } } // Rewrite ggsql:name → __ggsql_data_name__ in SQL - let sql = super::data::rewrite_namespaced_sql(sql)?; + let sql = super::builtin_data::rewrite_namespaced_sql(sql)?; if !super::returns_rows(&sql) { self.conn