From cc1481b5d7bed47b914ffed6f37568905e6f9435 Mon Sep 17 00:00:00 2001 From: Ferran Date: Tue, 28 Jul 2026 16:28:15 +0200 Subject: [PATCH 1/2] Resolve single rows in SQL tables --- crates/strata/src/providers/clickhouse/mod.rs | 11 ++- crates/strata/src/providers/mysql/mod.rs | 11 ++- crates/strata/src/providers/postgres/mod.rs | 17 ++++- crates/strata/src/providers/sql/mod.rs | 76 ++++++++++++++++++- crates/strata/src/providers/sqlite/mod.rs | 9 ++- crates/strata/src/router.rs | 11 +++ 6 files changed, 125 insertions(+), 10 deletions(-) diff --git a/crates/strata/src/providers/clickhouse/mod.rs b/crates/strata/src/providers/clickhouse/mod.rs index e050a6d..f7f0ec4 100644 --- a/crates/strata/src/providers/clickhouse/mod.rs +++ b/crates/strata/src/providers/clickhouse/mod.rs @@ -73,7 +73,7 @@ impl Clickhouse { async fn table_columns(&self, table: &str) -> Result> { let sql = format!( - "SELECT name, type FROM system.columns \ + "SELECT name, type, is_in_primary_key FROM system.columns \ WHERE database = currentDatabase() AND table = {} \ ORDER BY position FORMAT JSONEachRow", quote_str(table), @@ -84,7 +84,11 @@ impl Clickhouse { .map(|r| { let raw: RawColumn = serde_json::from_value(r).context("decoding column row")?; let (base, nullable) = strip_nullable(&raw.sql_type); - Ok(Field::new(raw.name, ch_to_data_type(base)?, nullable)) + let mut field = Field::new(raw.name, ch_to_data_type(base)?, nullable); + if raw.is_in_primary_key != 0 { + field.annotate(Field::KEY, "true"); + } + Ok(field) }) .collect() } @@ -244,6 +248,8 @@ struct RawColumn { name: String, #[serde(rename = "type")] sql_type: String, + #[serde(default)] + is_in_primary_key: u8, } /// Split a ClickHouse `Nullable(T)` wrapper, returning `(inner_type, nullable)`. @@ -350,6 +356,7 @@ mod tests { sql::suite::write_then_read_paginates_by_cursor(&client).await?; sql::suite::filters_rows(&client).await?; sql::suite::projects_columns(&client).await?; + sql::suite::gets_single_row(&client).await?; Ok(()) } } diff --git a/crates/strata/src/providers/mysql/mod.rs b/crates/strata/src/providers/mysql/mod.rs index b14aebf..134c1a1 100644 --- a/crates/strata/src/providers/mysql/mod.rs +++ b/crates/strata/src/providers/mysql/mod.rs @@ -75,7 +75,7 @@ impl SqlSource for Mysql { let rows = sqlx::query( "SELECT CAST(column_name AS CHAR) AS `name`, CAST(data_type AS CHAR) AS `type`, \ - CAST(is_nullable AS CHAR) AS `nullable` \ + CAST(is_nullable AS CHAR) AS `nullable`, CAST(column_key AS CHAR) AS `key` \ FROM information_schema.columns \ WHERE table_schema = DATABASE() AND table_name = ? ORDER BY ordinal_position", ) @@ -89,11 +89,15 @@ impl SqlSource for Mysql { let fields = rows .iter() .map(|r| { - Ok(Field::new( + let mut field = Field::new( r.get::("name"), mysql_to_data_type(&r.get::("type"))?, r.get::("nullable") == "YES", - )) + ); + if r.get::("key") == "PRI" { + field.annotate(Field::KEY, "true"); + } + Ok(field) }) .collect::>>()?; Ok(Schema::new(fields)) @@ -385,6 +389,7 @@ mod tests { sql::suite::write_then_read_paginates_by_cursor(&client).await?; sql::suite::filters_rows(&client).await?; sql::suite::projects_columns(&client).await?; + sql::suite::gets_single_row(&client).await?; Ok(()) } } diff --git a/crates/strata/src/providers/postgres/mod.rs b/crates/strata/src/providers/postgres/mod.rs index 36cb12e..f888493 100644 --- a/crates/strata/src/providers/postgres/mod.rs +++ b/crates/strata/src/providers/postgres/mod.rs @@ -82,13 +82,27 @@ impl SqlSource for Postgres { if rows.is_empty() { return Err(SqlError::TableNotFound(table.to_string()).into()); } + let pk_rows = client + .query( + "SELECT a.attname AS name FROM pg_index i \ + JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) \ + WHERE i.indrelid = to_regclass($1) AND i.indisprimary", + &[&table], + ) + .await?; + let keys: std::collections::HashSet = + pk_rows.iter().map(|r| r.get::<_, String>("name")).collect(); let columns = rows .iter() .map(|r| { let name: String = r.get("column_name"); let sql_type: String = r.get("data_type"); let nullable = r.get::<_, String>("is_nullable") == "YES"; - Ok(Field::new(name, pg_to_data_type(&sql_type)?, nullable)) + let mut field = Field::new(name.clone(), pg_to_data_type(&sql_type)?, nullable); + if keys.contains(&name) { + field.annotate(Field::KEY, "true"); + } + Ok(field) }) .collect::>()?; Ok(Schema::new(columns)) @@ -356,6 +370,7 @@ mod tests { sql::suite::write_then_read_paginates_by_cursor(&client).await?; sql::suite::filters_rows(&client).await?; sql::suite::projects_columns(&client).await?; + sql::suite::gets_single_row(&client).await?; Ok(()) } } diff --git a/crates/strata/src/providers/sql/mod.rs b/crates/strata/src/providers/sql/mod.rs index cdb1a05..11e6a36 100644 --- a/crates/strata/src/providers/sql/mod.rs +++ b/crates/strata/src/providers/sql/mod.rs @@ -17,8 +17,8 @@ use std::sync::Arc; use futures::StreamExt; -use anyhow::{Result, bail}; -use schema::{HasSchema, Schema}; +use anyhow::{Result, anyhow, bail}; +use schema::{DataType, HasSchema, Schema}; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -250,6 +250,43 @@ pub trait SqlSource: Send + Sync + 'static { } } + fn get_row( + &self, + table: &str, + key: &str, + ) -> impl Future>> + Send { + async move { + let schema = self.table_schema(table).await?; + let key_field = schema + .fields + .iter() + .find(|f| f.is_key()) + .ok_or_else(|| anyhow!("table `{table}` has no key column"))?; + let value = if matches!(key_field.data_type, DataType::Int64 | DataType::UInt64) { + key.parse::() + .map(Value::from) + .unwrap_or_else(|_| Value::String(key.to_string())) + } else { + Value::String(key.to_string()) + }; + let filter = Filter::Cmp { + field: key_field.name.clone(), + op: Op::Eq, + value, + }; + let cursor = SqlCursor { + offset: 0, + limit: 1, + cursor: None, + }; + Ok(self + .table_rows(table, &cursor, Some(&filter)) + .await? + .into_iter() + .next()) + } + } + fn register_tables(r: &mut Router) where Self: Provider, @@ -268,6 +305,12 @@ pub trait SqlSource: Send + Sync + 'static { .strategy(ListStrategy::Offset) .queryable(), ); + r.add( + Route::new() + .path("/tables/:table/:id") + .get_records(table_get::) + .data_type(table_data_schema::), + ); r.add(Route::new().path("/tables/:table").put(write_table::)); } } @@ -283,6 +326,14 @@ pub async fn list_tables(db: Arc, _p: Params) -> Result(db: Arc, p: Params) -> Result { + let table = p.get("table")?; + let key = p.get("id")?; + db.get_row(table, key) + .await? + .ok_or_else(|| anyhow!("no row with key `{key}` in table `{table}`")) +} + /// `list_records /tables/:table/data`: a page of a table's rows as typed Arrow /// columns, offset-paginated. pub async fn table_data(db: Arc, p: Params) -> Result { @@ -483,6 +534,8 @@ pub mod suite { let columns: Vec<&str> = schema.fields.iter().map(|f| f.name.as_str()).collect(); assert_eq!(columns, ["id", "name"]); assert_eq!(schema.fields[0].data_type, DataType::Int64); + assert!(schema.fields[0].is_key(), "`id` must round-trip as the key"); + assert!(!schema.fields[1].is_key(), "`name` is not a key"); Ok(()) } @@ -651,4 +704,23 @@ pub mod suite { assert_eq!(full.next().await?.len(), 2); Ok(()) } + + pub async fn gets_single_row(client: &Client) -> Result<()> { + let rows = [ + Row { + id: 1, + name: "a".into(), + }, + Row { + id: 2, + name: "b".into(), + }, + ]; + let _: WriteResult = client.put("/tables/getone", Dataset::of(&rows)?).await?; + + let row: Row = client.get("/tables/getone/2").await?; + assert_eq!(row.id, 2); + assert_eq!(row.name, "b"); + Ok(()) + } } diff --git a/crates/strata/src/providers/sqlite/mod.rs b/crates/strata/src/providers/sqlite/mod.rs index 1c7ab36..dcd34db 100644 --- a/crates/strata/src/providers/sqlite/mod.rs +++ b/crates/strata/src/providers/sqlite/mod.rs @@ -109,11 +109,15 @@ impl SqlSource for Sqlite { .iter() .map(|r| { let not_null = r.get::("notnull") != 0; - Field::new( + let mut field = Field::new( r.get::("name"), sqlite_to_data_type(&r.get::("type")), !not_null, - ) + ); + if r.get::("pk") != 0 { + field.annotate(Field::KEY, "true"); + } + field }) .collect(); Ok(Schema::new(fields)) @@ -371,6 +375,7 @@ mod tests { sql::suite::write_then_read_paginates_by_cursor(&client).await?; sql::suite::filters_rows(&client).await?; sql::suite::projects_columns(&client).await?; + sql::suite::gets_single_row(&client).await?; let _ = std::fs::remove_file(&path); Ok(()) } diff --git a/crates/strata/src/router.rs b/crates/strata/src/router.rs index c69397f..0f731af 100644 --- a/crates/strata/src/router.rs +++ b/crates/strata/src/router.rs @@ -519,6 +519,17 @@ impl Route { self } + pub fn get_records(mut self, handler: F) -> Self + where + F: Fn(Arc, Params) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + self.method = Method::Get; + self.handler = Some(erase(handler)); + self.response_schema = None; + self + } + pub fn strategy(mut self, strategy: ListStrategy) -> Self { self.strategy = Some(strategy); self From c86968160a1806408a96fe7864493f09fc119533 Mon Sep 17 00:00:00 2001 From: Ferran Date: Tue, 28 Jul 2026 16:33:29 +0200 Subject: [PATCH 2/2] Add projection on single row --- crates/strata/src/providers/sql/mod.rs | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/crates/strata/src/providers/sql/mod.rs b/crates/strata/src/providers/sql/mod.rs index 11e6a36..c220ea0 100644 --- a/crates/strata/src/providers/sql/mod.rs +++ b/crates/strata/src/providers/sql/mod.rs @@ -329,9 +329,16 @@ pub async fn list_tables(db: Arc, _p: Params) -> Result(db: Arc, p: Params) -> Result { let table = p.get("table")?; let key = p.get("id")?; - db.get_row(table, key) + let mut row = db + .get_row(table, key) .await? - .ok_or_else(|| anyhow!("no row with key `{key}` in table `{table}`")) + .ok_or_else(|| anyhow!("no row with key `{key}` in table `{table}`"))?; + if let Some(fields) = get_projection(&p) + && let Value::Object(map) = &mut row + { + map.retain(|k, _| fields.contains(k)); + } + Ok(row) } /// `list_records /tables/:table/data`: a page of a table's rows as typed Arrow @@ -721,6 +728,9 @@ pub mod suite { let row: Row = client.get("/tables/getone/2").await?; assert_eq!(row.id, 2); assert_eq!(row.name, "b"); + + let projected: RowProjected = client.get("/tables/getone/2?fields=name").await?; + assert_eq!(projected.name, "b"); Ok(()) } }