Resolve single rows in SQL tables

crates/strata/src/providers/clickhouse/mod.rs

@@ -73,7 +73,7 @@ impl Clickhouse {
 
     async fn table_columns(&self, table: &str) -> Result<Vec<Field>> {
         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(())
     }
 }

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::<String, _>("name"),
                     mysql_to_data_type(&r.get::<String, _>("type"))?,
                     r.get::<String, _>("nullable") == "YES",
-                ))
+                );
+                if r.get::<String, _>("key") == "PRI" {
+                    field.annotate(Field::KEY, "true");
+                }
+                Ok(field)
             })
             .collect::<Result<Vec<Field>>>()?;
         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(())
     }
 }

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<String> =
+            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::<Result<_>>()?;
         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(())
     }
 }

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<Output = Result<Option<Value>>> + 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::<i64>()
+                    .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<Self>)
     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::<Self>)
+                .data_type(table_data_schema::<Self>),
+        );
         r.add(Route::new().path("/tables/:table").put(write_table::<Self>));
     }
 }
@@ -283,6 +326,21 @@ pub async fn list_tables<S: SqlSource>(db: Arc<S>, _p: Params) -> Result<Page<Ta
     Ok(Page::new(items, Cursor::empty()))
 }
 
+pub async fn table_get<S: SqlSource>(db: Arc<S>, p: Params) -> Result<Value> {
+    let table = p.get("table")?;
+    let key = p.get("id")?;
+    let mut row = db
+        .get_row(table, key)
+        .await?
+        .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
 /// columns, offset-paginated.
 pub async fn table_data<S: SqlSource>(db: Arc<S>, p: Params) -> Result<RecordPage> {
@@ -483,6 +541,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 +711,26 @@ pub mod suite {
         assert_eq!(full.next().await?.len(), 2);
         Ok(())
     }
+
+    pub async fn gets_single_row<S: Provider>(client: &Client<S>) -> 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");
+
+        let projected: RowProjected = client.get("/tables/getone/2?fields=name").await?;
+        assert_eq!(projected.name, "b");
+        Ok(())
+    }
 }

crates/strata/src/providers/sqlite/mod.rs

@@ -109,11 +109,15 @@ impl SqlSource for Sqlite {
             .iter()
             .map(|r| {
                 let not_null = r.get::<i64, _>("notnull") != 0;
-                Field::new(
+                let mut field = Field::new(
                     r.get::<String, _>("name"),
                     sqlite_to_data_type(&r.get::<String, _>("type")),
                     !not_null,
-                )
+                );
+                if r.get::<i64, _>("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(())
     }

crates/strata/src/router.rs

@@ -519,6 +519,17 @@ impl<S: Send + Sync + 'static> Route<S> {
         self
     }
 
+    pub fn get_records<F, Fut>(mut self, handler: F) -> Self
+    where
+        F: Fn(Arc<S>, Params) -> Fut + Send + Sync + 'static,
+        Fut: Future<Output = Result<Value>> + 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