From e4f01968d74aa2c5cf534f2e1fe53d14e292936c Mon Sep 17 00:00:00 2001 From: Ferran Date: Fri, 31 Jul 2026 11:59:39 +0200 Subject: [PATCH 1/4] Add MCP server --- Cargo.lock | 89 +++++++++++++++ crates/strata/Cargo.toml | 1 + crates/strata/src/lib.rs | 1 + crates/strata/src/main.rs | 13 ++- crates/strata/src/mcp/mod.rs | 210 +++++++++++++++++++++++++++++++++++ examples/use_mcp.md | 83 ++++++++++++++ 6 files changed, 395 insertions(+), 2 deletions(-) create mode 100644 crates/strata/src/mcp/mod.rs create mode 100644 examples/use_mcp.md diff --git a/Cargo.lock b/Cargo.lock index c1bf539..e701c67 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3046,6 +3046,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pastey" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" + [[package]] name = "pem-rfc7468" version = "0.7.0" @@ -3577,6 +3583,50 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rmcp" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "14db48ee17a9ba61810ab1a9c1beb7d06d8136ae39ac25a1137f10d357af01af" +dependencies = [ + "async-trait", + "base64", + "bytes", + "chrono", + "futures", + "http", + "http-body", + "http-body-util", + "pastey", + "pin-project-lite", + "rand 0.10.1", + "rmcp-macros", + "schemars 1.2.1", + "serde", + "serde_json", + "sse-stream", + "thiserror", + "tokio", + "tokio-stream", + "tokio-util", + "tower-service", + "tracing", + "uuid", +] + +[[package]] +name = "rmcp-macros" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "783d787bf21813b285f13019adc49e11af501c658890c1e519f31f937c68b7e3" +dependencies = [ + "darling 0.23.0", + "proc-macro2", + "quote", + "serde_json", + "syn", +] + [[package]] name = "roaring" version = "0.11.4" @@ -3742,12 +3792,26 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc" dependencies = [ + "chrono", "dyn-clone", "ref-cast", + "schemars_derive", "serde", "serde_json", ] +[[package]] +name = "schemars_derive" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d115b50f4aaeea07e79c1912f645c7513d81715d0420f8bc77a18c6260b307f" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -3853,6 +3917,17 @@ dependencies = [ "syn", ] +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "serde_json" version = "1.0.150" @@ -4264,6 +4339,19 @@ dependencies = [ "uuid", ] +[[package]] +name = "sse-stream" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c123f296ade4ec4b8b0f6162116e6629f5146922ca5ab40ca9d3c2e73ab4761e" +dependencies = [ + "bytes", + "futures-util", + "http-body", + "http-body-util", + "pin-project-lite", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -4303,6 +4391,7 @@ dependencies = [ "iceberg", "iceberg-catalog-sql", "parquet", + "rmcp", "schema", "serde", "serde_arrow", diff --git a/crates/strata/Cargo.toml b/crates/strata/Cargo.toml index 93c6aba..349321b 100644 --- a/crates/strata/Cargo.toml +++ b/crates/strata/Cargo.toml @@ -44,6 +44,7 @@ eyre = "0.6.12" async-graphql = { version = "7", features = ["dynamic-schema"] } async-graphql-axum = "7" axum = "0.8" +rmcp = { version = "2.2.0", features = ["server", "macros", "transport-streamable-http-server", "server-side-http"] } [dev-dependencies] dockertest = { path = "../dockertest" } diff --git a/crates/strata/src/lib.rs b/crates/strata/src/lib.rs index 4b4db09..59c2f06 100644 --- a/crates/strata/src/lib.rs +++ b/crates/strata/src/lib.rs @@ -8,6 +8,7 @@ pub mod dataset; pub mod flight; pub mod graphql; pub mod harness; +pub mod mcp; pub mod page; pub mod pipe; pub mod provider; diff --git a/crates/strata/src/main.rs b/crates/strata/src/main.rs index 2eb7c8c..78667f1 100644 --- a/crates/strata/src/main.rs +++ b/crates/strata/src/main.rs @@ -82,6 +82,9 @@ enum Command { /// GraphQL HTTP address. #[arg(long, default_value = "127.0.0.1:8080")] graphql_addr: String, + /// MCP HTTP address. + #[arg(long, default_value = "127.0.0.1:8081")] + mcp_addr: String, }, } @@ -209,9 +212,14 @@ async fn run() -> Result<()> { pipe.source.mount, pipe.source.path, pipe.destination.mount, pipe.destination.path ); } - Command::Serve { addr, graphql_addr } => { + Command::Serve { + addr, + graphql_addr, + mcp_addr, + } => { let addr = addr.parse()?; let graphql_addr = graphql_addr.parse()?; + let mcp_addr = mcp_addr.parse()?; // Pipe subsystem: only spun up when pipes are declared, so `serve` // still runs without a catalog DB when there are none. @@ -277,7 +285,8 @@ async fn run() -> Result<()> { tokio::try_join!( strata::flight::serve(registry.clone(), addr), - strata::graphql::serve(registry, graphql_addr), + strata::graphql::serve(registry.clone(), graphql_addr), + strata::mcp::serve(registry, mcp_addr), )?; } } diff --git a/crates/strata/src/mcp/mod.rs b/crates/strata/src/mcp/mod.rs new file mode 100644 index 0000000..f41dce0 --- /dev/null +++ b/crates/strata/src/mcp/mod.rs @@ -0,0 +1,210 @@ +//! MCP server over the registry. + +use std::net::SocketAddr; +use std::sync::Arc; + +use anyhow::Result; +use rmcp::handler::server::wrapper::{Json, Parameters}; +use rmcp::transport::streamable_http_server::StreamableHttpService; +use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; +use rmcp::transport::streamable_http_server::tower::StreamableHttpServerConfig; +use rmcp::{ErrorData, schemars, tool, tool_router}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::Registry; +use crate::router::Method; + +pub async fn serve(registry: Arc, addr: SocketAddr) -> Result<()> { + let service = StreamableHttpService::new( + move || { + Ok(Strata { + registry: registry.clone(), + }) + }, + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default() + .with_stateful_mode(false) + .with_json_response(true), + ); + let app = axum::Router::new().nest_service("/mcp", service); + tracing::info!("strata MCP server on http://{addr}/mcp"); + let listener = tokio::net::TcpListener::bind(addr).await?; + axum::serve(listener, app).await?; + Ok(()) +} + +#[derive(Clone)] +struct Strata { + registry: Arc, +} + +#[derive(Debug, Deserialize, schemars::JsonSchema)] +struct ProviderArgs { + /// Mount name of the provider, e.g. `github`. + provider: String, +} + +#[derive(Debug, Deserialize, schemars::JsonSchema)] +struct EndpointArgs { + /// Mount name of the provider, e.g. `local`. + provider: String, + /// Endpoint path within the provider, e.g. `/tables/headlines`. + path: String, +} + +#[derive(Debug, Deserialize, schemars::JsonSchema)] +struct ListArgs { + /// Mount name of the provider, e.g. `local`. + provider: String, + /// Endpoint path within the provider, e.g. `/tables/headlines`. + path: String, + /// Resume token from a previous call's `cursor.next`. Omit for the first page. + cursor: Option, + /// Requested page size. A hint; the provider picks the real chunk size. + limit: Option, + /// Row predicate, as the JSON filter the endpoint accepts. Only some + /// endpoints support it; see `metadata.queryable` from `describe_provider`. + filter: Option, + /// Subset of columns to return. Omit for the whole row. + fields: Option>, +} + +fn failed(error: anyhow::Error) -> ErrorData { + ErrorData::internal_error(format!("{error:#}"), None) +} + +#[derive(Debug, Serialize, schemars::JsonSchema)] +struct Providers { + providers: Vec, +} + +#[derive(Debug, Serialize, schemars::JsonSchema)] +struct Endpoints { + endpoints: Value, +} + +#[derive(Debug, Serialize, schemars::JsonSchema)] +struct ResolvedEndpoint { + method: String, + path: String, + params: Vec, + /// JSON Schema of one row of the response. + response: Value, + metadata: Value, +} + +#[derive(Debug, Serialize, schemars::JsonSchema)] +struct Entity { + entity: Value, +} + +#[derive(Debug, Serialize, schemars::JsonSchema)] +struct Rows { + items: Vec, + /// `next` is the token to pass back as `cursor`; null means the end. + cursor: Option, +} + +#[tool_router(server_handler)] +impl Strata { + /// Every mounted provider, by mount name. Start here, then call + /// `describe_provider` to see what one of them answers. + #[tool] + async fn list_providers(&self) -> Json { + Json(Providers { + providers: self.registry.names(), + }) + } + + /// Every endpoint of one provider: its verb, path, path parameters, and the + /// JSON Schema of its request body and response. Schemas here are the static + /// ones; for an endpoint whose type depends on the concrete path (a SQL + /// table's columns), call `resolve_endpoint`. + #[tool] + async fn describe_provider( + &self, + Parameters(args): Parameters, + ) -> Result, ErrorData> { + let described = self.registry.describe(&args.provider).map_err(failed)?; + Ok(Json(Endpoints { + endpoints: described.get("endpoints").cloned().unwrap_or(Value::Null), + })) + } + + /// The response schema of one concrete path, resolved by running the + /// endpoint's dynamic resolver. This is how a SQL or Iceberg table's real + /// columns become visible, without reading any rows. + #[tool] + async fn resolve_endpoint( + &self, + Parameters(args): Parameters, + ) -> Result, ErrorData> { + let provider = self.registry.get(&args.provider).map_err(failed)?; + let endpoint = provider.resolve(&args.path).await.map_err(failed)?; + Ok(Json(ResolvedEndpoint { + method: endpoint.method.to_string(), + path: endpoint.path, + params: endpoint.params, + response: endpoint.response.to_json_schema(), + metadata: serde_json::to_value(&endpoint.metadata).unwrap_or(Value::Null), + })) + } + + /// Read one entity from a `get` endpoint. + #[tool] + async fn get( + &self, + Parameters(args): Parameters, + ) -> Result, ErrorData> { + let provider = self.registry.get(&args.provider).map_err(failed)?; + let response = provider + .invoke(Method::Get, &args.path, None) + .await + .map_err(failed)?; + Ok(Json(Entity { + entity: response.entity.unwrap_or(response.output), + })) + } + + /// Read one page of rows from a `list` endpoint, as `{ items, cursor }`. Pass + /// `cursor.next` back as `cursor` to continue; a `next` of null is the end. + #[tool] + async fn list(&self, Parameters(args): Parameters) -> Result, ErrorData> { + let provider = self.registry.get(&args.provider).map_err(failed)?; + let path = read_path(&args); + let stream = provider.read(&path).await.map_err(failed)?; + let schema = stream.schema.clone(); + let page = stream.first().await.map_err(failed)?; + let (items, cursor) = match page { + Some(page) => ( + page.data.to_json_rows(&schema).map_err(failed)?, + page.cursor.and_then(|c| serde_json::to_value(c).ok()), + ), + None => (Vec::new(), None), + }; + Ok(Json(Rows { items, cursor })) + } +} + +fn read_path(args: &ListArgs) -> String { + let mut params: Vec<(String, String)> = Vec::new(); + if let Some(cursor) = &args.cursor { + params.push(("cursor".into(), cursor.clone())); + } + if let Some(limit) = args.limit { + params.push(("limit".into(), limit.to_string())); + } + if let Some(filter) = &args.filter { + params.push(("filter".into(), filter.to_string())); + } + if let Some(fields) = &args.fields + && !fields.is_empty() + { + params.push(("fields".into(), fields.join(","))); + } + match serde_urlencoded::to_string(¶ms) { + Ok(query) if !query.is_empty() => format!("{}?{query}", args.path), + _ => args.path.clone(), + } +} diff --git a/examples/use_mcp.md b/examples/use_mcp.md new file mode 100644 index 0000000..8ba2fe8 --- /dev/null +++ b/examples/use_mcp.md @@ -0,0 +1,83 @@ +# Read a source from an MCP client + +Pipe a Hacker News RSS feed into a local SQLite table, then let a model discover +and read it over MCP. + +## 1. Config + +```toml +# strata.toml +[provider.hn] +backend = "rss" +url = "https://hnrss.org/frontpage" + +[provider.local] +backend = "sqlite" +path = "./local.sqlite" +``` + +## 2. Materialize the feed into SQLite + +```bash +strata pipe hn /items local /tables/headlines +``` + +## 3. Serve + +`serve` starts the MCP server at `http://127.0.0.1:8081/mcp`. + +```bash +strata serve +``` + +## 4. The tools + +The examples below use the MCP Inspector CLI as the client. + +```bash +mcp-inspector --cli http://127.0.0.1:8081/mcp --transport http --method tools/list +``` + +| tool | what it answers | +| :- | :- | +| `list_providers` | the mounted providers | +| `describe_provider` | every endpoint of one provider | +| `resolve_endpoint` | the schema of one concrete path | +| `get` | one entity from a `get` endpoint | +| `list` | one page of rows from a `list` endpoint | + +Reads only. `put` takes an Arrow stream, so use `strata pipe` to write. + +## 5. Call them + +What is mounted: + +```bash +mcp-inspector --cli http://127.0.0.1:8081/mcp --transport http \ + --method tools/call --tool-name "list_providers"; +# {"providers":["dummy","hn","local"]} +``` + +The table's real columns, without reading a row. This is how a model learns what +it can filter and project on: + +```bash +mcp-inspector --cli http://127.0.0.1:8081/mcp --transport http \ + --method tools/call --tool-name "resolve_endpoint" \ + --tool-args-json '{"provider":"local","path":"/tables/headlines"}' +# {"method":"list","path":"/tables/:table","params":["table"], +# "response":{"type":"object","properties":{"id":{"type":"string"}, +# "title":{"anyOf":[{"type":"string"},{"type":"null"}]}, ...}}, +# "metadata":{"strategy":"Offset","disposition":"Append","queryable":true}} +``` + +Read rows, projecting to the columns it needs: + +```bash +mcp-inspector --cli http://127.0.0.1:8081/mcp --transport http \ + --method tools/call --tool-name "list" \ + --tool-args-json '{"provider":"local","path":"/tables/headlines","limit":2,"fields":["title"]}' +# {"items":[{"title":"DeepSeek V4 Flash 0731 Intelligence, Performance and Price Analysis"}, +# {"title":"Danube's record low levels force shutdown of Hungary's only nuclear plant"}], +# "cursor":{"next":"{\"offset\":2}"}} +``` From 24e74b9d28e83296c21c5a96c5036b395444f692 Mon Sep 17 00:00:00 2001 From: Ferran Date: Fri, 31 Jul 2026 13:31:12 +0200 Subject: [PATCH 2/4] Remove Dataset struct --- crates/strata/src/bin/loadgen.rs | 4 +- crates/strata/src/datagen.rs | 9 ++- crates/strata/src/dataset.rs | 67 +------------------ crates/strata/src/graphql/mod.rs | 2 +- crates/strata/src/lib.rs | 2 +- crates/strata/src/providers/clickhouse/mod.rs | 9 +-- crates/strata/src/providers/iceberg/mod.rs | 2 +- crates/strata/src/providers/mysql/mod.rs | 12 ++-- crates/strata/src/providers/postgres/mod.rs | 12 ++-- crates/strata/src/providers/sql/mod.rs | 32 ++++----- crates/strata/src/providers/sqlite/mod.rs | 12 ++-- crates/strata/src/record.rs | 20 +++++- crates/strata/src/testkit.rs | 5 +- 13 files changed, 71 insertions(+), 117 deletions(-) diff --git a/crates/strata/src/bin/loadgen.rs b/crates/strata/src/bin/loadgen.rs index 91c6f93..22c8615 100644 --- a/crates/strata/src/bin/loadgen.rs +++ b/crates/strata/src/bin/loadgen.rs @@ -14,7 +14,7 @@ use schema::{DataType, SchemaBuilder}; use serde_json::Value; use strata::datagen::Generator; use strata::harness::{Plan, Stop, run}; -use strata::{Body, Dataset, Method}; +use strata::{Body, DataStream, Method}; #[derive(Parser)] #[command(about = "Measure write throughput against a provider endpoint")] @@ -80,7 +80,7 @@ async fn main() -> Result<()> { let start = offset as usize; let records = generator.rows(start..start + batch)?; let body = Body { - data: Some(Dataset::new(schema.clone(), records).into_stream()), + data: Some(DataStream::once(schema.clone(), records)), meta: Value::Null, }; let response = provider diff --git a/crates/strata/src/datagen.rs b/crates/strata/src/datagen.rs index f9939c3..0afc3de 100644 --- a/crates/strata/src/datagen.rs +++ b/crates/strata/src/datagen.rs @@ -11,8 +11,7 @@ use anyhow::Result; use schema::{DataType, Field, Schema}; use serde_json::{Map, Value, json}; -use crate::dataset::Dataset; -use crate::record::{Batch, stringify_text_columns}; +use crate::record::{Batch, DataStream, stringify_text_columns}; pub struct Generator { schema: Schema, @@ -48,9 +47,9 @@ impl Generator { Batch::encode(&self.schema, &rows) } - /// The first `n` rows as a [`Dataset`] (schema + Arrow records). - pub fn dataset(&self, n: usize) -> Result { - Ok(Dataset::new(self.schema.clone(), self.rows(0..n)?)) + /// The first `n` rows as a single-page [`DataStream`], ready to `put`. + pub fn stream(&self, n: usize) -> Result { + Ok(DataStream::once(self.schema.clone(), self.rows(0..n)?)) } } diff --git a/crates/strata/src/dataset.rs b/crates/strata/src/dataset.rs index 9965e1e..4ba49c2 100644 --- a/crates/strata/src/dataset.rs +++ b/crates/strata/src/dataset.rs @@ -1,21 +1,7 @@ -//! A dataset: the `(schema, rows)` unit that flows from a reader to a writer. -//! -//! It's exactly what a read produces — the rows of a `list`/`get` plus that -//! endpoint's resolved [`DataType`] schema — so a writer can use the schema as -//! the contract (create a table if absent, validate if present) and load the -//! rows. This is the framework-level shape behind "pipe readers to writers"; -//! over Arrow Flight the same `(schema, rows)` rides natively in a `DoPut`. +//! How a sink applies a written dataset. use anyhow::{Result, bail}; -use futures::stream::StreamExt; -use schema::Schema; use serde::Serialize; -use serde_json::Value; - -use crate::{ - DataStream, - record::{Batch, BatchPage}, -}; /// How a sink should apply a written dataset. Rides as metadata on the existing /// `put` verb (the reserved `disposition` query param) rather than a new verb, so @@ -53,54 +39,3 @@ impl Disposition { } } } - -/// The reader→writer unit: the rows as Arrow [`Records`] (strata's internal -/// currency) paired with the **native** `schema` that governs how a sink stores -/// them. `schema` is carried explicitly rather than derived from the Arrow -/// schema, so `Decimal`/`Json` fidelity survives — the Arrow mapping folds those -/// to `Utf8`, and a sink's DDL needs the real types. It's typically a -/// `List(Struct(..))` (a list endpoint's response) or a bare `Struct`. -/// -/// Future direction: `records` should become a *stream* of batches following -/// `schema` (a checkpointed cursor-driven stream) so a writer can consume -/// unbounded data without materializing it — the materialized batches here are -/// the stand-in until the streaming list driver lands. -#[derive(Debug, Clone)] -pub struct Dataset { - pub schema: Schema, - pub records: Batch, -} - -impl Dataset { - pub fn new(schema: Schema, records: Batch) -> Self { - Dataset { schema, records } - } - - /// Build a dataset from typed `rows`: the schema is `T`'s [`Schema`] (with any - /// `#[schema(key)]`/annotations it declares) and the rows are Arrow-encoded - /// against it. The typed-input counterpart of a `put`. - pub fn of(rows: &[T]) -> Result { - let schema = T::schema(); - let records = Batch::encode(&schema, rows)?; - Ok(Dataset::new(schema, records)) - } - - /// Interim bridge: decode the Arrow rows to JSON for sinks that aren't yet - /// Arrow-native. This is the one labeled Arrow→JSON on the write path; Phase B - /// removes it as each sink binds Arrow columns directly. - pub fn to_json_rows(&self) -> Result> { - self.records.to_json_rows(&self.schema) - } - - /// Bridge to a single-chunk [`DataStream`], until readers produce batches lazily. - pub fn into_stream(self) -> DataStream { - let chunk = BatchPage { - data: self.records, - cursor: None, - }; - DataStream { - schema: self.schema, - chunks: futures::stream::once(async move { Ok(chunk) }).boxed(), - } - } -} diff --git a/crates/strata/src/graphql/mod.rs b/crates/strata/src/graphql/mod.rs index 03539e0..79fe91b 100644 --- a/crates/strata/src/graphql/mod.rs +++ b/crates/strata/src/graphql/mod.rs @@ -404,7 +404,7 @@ mod tests { rows: &[T], ) -> Result<()> { let body = crate::Body { - data: Some(crate::Dataset::of(rows)?.into_stream()), + data: Some(crate::DataStream::of(rows)?.into_stream()), meta: Value::Null, }; registry diff --git a/crates/strata/src/lib.rs b/crates/strata/src/lib.rs index 59c2f06..78524b3 100644 --- a/crates/strata/src/lib.rs +++ b/crates/strata/src/lib.rs @@ -25,7 +25,7 @@ use anyhow::Result; pub use catalog::Catalog; pub use config::{Config, ProviderConfig}; -pub use dataset::{Dataset, Disposition}; +pub use dataset::Disposition; pub use page::{Cursor, Page}; pub use provider::{Provider, ProviderObject, Registry}; pub use record::DataStream; diff --git a/crates/strata/src/providers/clickhouse/mod.rs b/crates/strata/src/providers/clickhouse/mod.rs index b6c6ed0..22a35ec 100644 --- a/crates/strata/src/providers/clickhouse/mod.rs +++ b/crates/strata/src/providers/clickhouse/mod.rs @@ -5,7 +5,8 @@ use schema::{DataType, Field, Schema}; use serde::Deserialize; use serde_json::Value; -use crate::dataset::{Dataset, Disposition}; +use crate::dataset::Disposition; +use crate::record::Batch; use crate::provider::Provider; use crate::providers::sql::{self, Filter, SqlCursor, SqlError, SqlSource, WriteResult, quote_str}; use crate::router::Router; @@ -197,10 +198,10 @@ impl SqlSource for Clickhouse { async fn write_table( &self, table: &str, - data: Dataset, + schema: &Schema, + data: Batch, _disposition: Disposition, ) -> Result { - let schema = &data.schema; let ident = quote_ident(table); // Build a single `INSERT … FORMAT JSONEachRow` body: one JSON object per row, @@ -209,7 +210,7 @@ impl SqlSource for Clickhouse { // ClickHouse's HTTP interface speaks JSONEachRow, so decode the Arrow rows to // JSON here — this is the one provider where JSON is the wire protocol, not an // interim bridge. - let rows = data.to_json_rows()?; + let rows = data.to_json_rows(schema)?; let mut rows_written = 0u64; if !rows.is_empty() { let mut body = format!("INSERT INTO {ident} FORMAT JSONEachRow\n"); diff --git a/crates/strata/src/providers/iceberg/mod.rs b/crates/strata/src/providers/iceberg/mod.rs index 8536990..b368120 100644 --- a/crates/strata/src/providers/iceberg/mod.rs +++ b/crates/strata/src/providers/iceberg/mod.rs @@ -441,7 +441,7 @@ mod tests { const ROWS: usize = 100; let generator = crate::datagen::Generator::new(&Event::schema())?; - let dataset = generator.dataset(ROWS)?; + let dataset = generator.stream(ROWS)?; let expected: Vec = dataset.records.decode(&dataset.schema)?; let result: WriteResult = c.put("/tables/events", dataset).await?; diff --git a/crates/strata/src/providers/mysql/mod.rs b/crates/strata/src/providers/mysql/mod.rs index fda2f5a..89534ca 100644 --- a/crates/strata/src/providers/mysql/mod.rs +++ b/crates/strata/src/providers/mysql/mod.rs @@ -5,7 +5,8 @@ use serde_json::Value; use sqlx::Row as _; use sqlx::mysql::{MySqlArguments, MySqlPool}; -use crate::dataset::{Dataset, Disposition}; +use crate::dataset::Disposition; +use crate::record::Batch; use crate::provider::Provider; use crate::providers::sql::{ self, Filter, SqlCursor, SqlError, SqlSource, WriteResult, is_table_not_found, quote_str, @@ -192,11 +193,12 @@ impl SqlSource for Mysql { async fn write_table( &self, table: &str, - data: Dataset, + schema: &Schema, + data: Batch, disposition: Disposition, ) -> Result { - let fields = &data.schema.fields; - let keys = data.schema.get_key_fields(); + let fields = &schema.fields; + let keys = schema.get_key_fields(); let key_refs: Vec<&str> = keys.iter().map(String::as_str).collect(); let ident = quote_ident(table); let pool = self.connect().await?; @@ -234,7 +236,7 @@ impl SqlSource for Mysql { }; // Interim: decode the Arrow rows to JSON to bind positionally (Phase B binds // Arrow columns directly). - let rows = data.to_json_rows()?; + let rows = data.to_json_rows(schema)?; // One statement per chunk of rows rather than one per row, chunked by the // placeholder budget a single statement can carry. diff --git a/crates/strata/src/providers/postgres/mod.rs b/crates/strata/src/providers/postgres/mod.rs index 9b1c9b2..15bd3a0 100644 --- a/crates/strata/src/providers/postgres/mod.rs +++ b/crates/strata/src/providers/postgres/mod.rs @@ -3,7 +3,8 @@ use schema::{DataType, Field, Schema}; use serde_json::Value; use tokio_postgres::NoTls; -use crate::dataset::{Dataset, Disposition}; +use crate::dataset::Disposition; +use crate::record::Batch; use crate::provider::Provider; use crate::providers::sql::{ self, Filter, SqlCursor, SqlError, SqlSource, WriteResult, is_table_not_found, @@ -182,11 +183,12 @@ impl SqlSource for Postgres { async fn write_table( &self, table: &str, - data: Dataset, + schema: &Schema, + data: Batch, disposition: Disposition, ) -> Result { - let fields = &data.schema.fields; - let keys = data.schema.get_key_fields(); + let fields = &schema.fields; + let keys = schema.get_key_fields(); let key_refs: Vec<&str> = keys.iter().map(String::as_str).collect(); let ident = quote_ident(table); let client = self.connect().await?; @@ -226,7 +228,7 @@ impl SqlSource for Postgres { ); // Interim: decode the Arrow rows to JSON to bind as jsonb (Phase B binds // Arrow columns directly). - let rows = Value::Array(data.to_json_rows()?); + let rows = Value::Array(data.to_json_rows(schema)?); let rows_written = client .execute(&insert, &[&rows]) .await diff --git a/crates/strata/src/providers/sql/mod.rs b/crates/strata/src/providers/sql/mod.rs index 20b48d6..a5d8fd0 100644 --- a/crates/strata/src/providers/sql/mod.rs +++ b/crates/strata/src/providers/sql/mod.rs @@ -22,7 +22,7 @@ use schema::{DataType, HasSchema, Schema}; use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::dataset::{Dataset, Disposition}; +use crate::dataset::Disposition; use crate::page::{Cursor, ListStrategy, Page}; use crate::provider::Provider; use crate::record::{Batch, BatchPage, DataStream, stringify_text_columns}; @@ -214,7 +214,8 @@ pub trait SqlSource: Send + Sync + 'static { fn write_table( &self, table: &str, - data: Dataset, + schema: &Schema, + data: Batch, disposition: Disposition, ) -> impl Future> + Send; @@ -236,12 +237,8 @@ pub trait SqlSource: Send + Sync + 'static { }; let mut rows_written = 0; while let Some(chunk) = chunks.next().await { - // TODO: We still create a dataset object, later on we should send - // the stream directly to the SQL provider but we would need a way for the provider - // to notify which batch has been written. - let dataset = Dataset::new(schema.clone(), chunk?.data); rows_written += self - .write_table(table, dataset, disposition) + .write_table(table, &schema, chunk?.data, disposition) .await? .rows_written; } @@ -552,7 +549,6 @@ fn next_cursor(cursor: &SqlCursor, returned: usize) -> Result { #[cfg(test)] pub mod suite { use super::{TableName, WriteResult}; - use crate::dataset::Dataset; use crate::provider::Provider; use crate::testkit::Client; use anyhow::Result; @@ -587,7 +583,7 @@ pub mod suite { id: 1, name: "a".into(), }]; - let _: WriteResult = client.put("/tables/catalog", Dataset::of(&rows)?).await?; + let _: WriteResult = client.put("/tables/catalog", DataStream::of(&rows)?).await?; let mut tables = client.list("/tables").await?; let names: Vec = tables.next().await?; @@ -613,7 +609,7 @@ pub mod suite { client: &Client, ) -> Result<()> { let generator = crate::datagen::Generator::new(&Event::schema())?; - let result: WriteResult = client.put("/tables/events", generator.dataset(20)?).await?; + let result: WriteResult = client.put("/tables/events", generator.stream(20)?).await?; assert!(result.created); assert_eq!(result.rows_written, 20); @@ -631,7 +627,7 @@ pub mod suite { let generator = crate::datagen::Generator::new(&Event::schema())?; let result: WriteResult = client - .put("/tables/streamed", generator.dataset(ROWS)?) + .put("/tables/streamed", generator.stream(ROWS)?) .await?; assert_eq!(result.rows_written, ROWS as u64); @@ -673,8 +669,8 @@ pub mod suite { name: "b".into(), }, ]; - let _: WriteResult = client.put("/tables/dedup", Dataset::of(&rows)?).await?; - let _: WriteResult = client.put("/tables/dedup", Dataset::of(&rows)?).await?; + let _: WriteResult = client.put("/tables/dedup", DataStream::of(&rows)?).await?; + let _: WriteResult = client.put("/tables/dedup", DataStream::of(&rows)?).await?; let mut stream = client.list("/tables/dedup").await?; let found_rows: Vec = stream.next().await?; @@ -692,7 +688,7 @@ pub mod suite { let _: WriteResult = client .put( "/tables/up", - Dataset::of(&[Row { + DataStream::of(&[Row { id: 1, name: "a".into(), }])?, @@ -701,7 +697,7 @@ pub mod suite { let _: WriteResult = client .put( "/tables/up?disposition=merge", - Dataset::of(&[Row { + DataStream::of(&[Row { id: 1, name: "b".into(), }])?, @@ -739,7 +735,7 @@ pub mod suite { name: "d".into(), }, ]; - let _: WriteResult = client.put("/tables/filtered", Dataset::of(&rows)?).await?; + let _: WriteResult = client.put("/tables/filtered", DataStream::of(&rows)?).await?; let predicate = serde_json::json!({ "and": [ @@ -781,7 +777,7 @@ pub mod suite { name: "b".into(), }, ]; - let _: WriteResult = client.put("/tables/projected", Dataset::of(&rows)?).await?; + let _: WriteResult = client.put("/tables/projected", DataStream::of(&rows)?).await?; let mut stream = client .list::("/tables/projected?fields=name") @@ -819,7 +815,7 @@ pub mod suite { name: "b".into(), }, ]; - let _: WriteResult = client.put("/tables/getone", Dataset::of(&rows)?).await?; + let _: WriteResult = client.put("/tables/getone", DataStream::of(&rows)?).await?; let row: Row = client.get("/tables/getone/2").await?; assert_eq!(row.id, 2); diff --git a/crates/strata/src/providers/sqlite/mod.rs b/crates/strata/src/providers/sqlite/mod.rs index 5928a04..9e3d06e 100644 --- a/crates/strata/src/providers/sqlite/mod.rs +++ b/crates/strata/src/providers/sqlite/mod.rs @@ -23,7 +23,8 @@ use serde_json::Value; use sqlx::Row as _; use sqlx::sqlite::{SqliteArguments, SqliteConnectOptions, SqlitePool}; -use crate::dataset::{Dataset, Disposition}; +use crate::dataset::Disposition; +use crate::record::Batch; use crate::provider::Provider; use crate::providers::sql::{ self, Filter, SqlCursor, SqlError, SqlSource, WriteResult, is_table_not_found, quote_str, @@ -213,11 +214,12 @@ impl SqlSource for Sqlite { async fn write_table( &self, table: &str, - data: Dataset, + schema: &Schema, + data: Batch, disposition: Disposition, ) -> Result { - let fields = &data.schema.fields; - let keys = data.schema.get_key_fields(); + let fields = &schema.fields; + let keys = schema.get_key_fields(); let key_refs: Vec<&str> = keys.iter().map(String::as_str).collect(); let ident = quote_ident(table); let pool = self.connect().await?; @@ -255,7 +257,7 @@ impl SqlSource for Sqlite { }; // Interim: decode the Arrow rows to JSON to bind positionally (Phase B binds // Arrow columns directly). - let rows = data.to_json_rows()?; + let rows = data.to_json_rows(schema)?; // One statement per chunk of rows rather than one per row, all inside a // single transaction — SQLite otherwise autocommits (and fsyncs) every insert. diff --git a/crates/strata/src/record.rs b/crates/strata/src/record.rs index 5b481eb..09dc4f2 100644 --- a/crates/strata/src/record.rs +++ b/crates/strata/src/record.rs @@ -6,7 +6,7 @@ use arrow::record_batch::RecordBatch; use arrow_array::Array; use arrow_schema::Schema; use futures::stream::{BoxStream, StreamExt}; -use schema::{Annotations, DataType, Schema as StrataSchema}; +use schema::{Annotations, DataType, HasSchema, Schema as StrataSchema}; use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::Value; @@ -22,6 +22,24 @@ impl DataStream { pub async fn first(mut self) -> Result> { self.chunks.next().await.transpose() } + + /// A stream of exactly one page — the shape a caller that already has all the + /// rows in hand writes into a sink. + pub fn once(schema: StrataSchema, data: Batch) -> DataStream { + let page = BatchPage { data, cursor: None }; + DataStream { + schema, + chunks: futures::stream::once(async move { Ok(page) }).boxed(), + } + } + + /// One page built from typed `rows`: the schema is `T`'s, with whatever + /// `#[schema(key)]` annotations it declares, and the rows are encoded against it. + pub fn of(rows: &[T]) -> Result { + let schema = T::schema(); + let data = Batch::encode(&schema, rows)?; + Ok(DataStream::once(schema, data)) + } } pub struct BatchPage { pub data: Batch, diff --git a/crates/strata/src/testkit.rs b/crates/strata/src/testkit.rs index 099651b..46a8496 100644 --- a/crates/strata/src/testkit.rs +++ b/crates/strata/src/testkit.rs @@ -14,7 +14,6 @@ use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::Value; -use crate::dataset::Dataset; use crate::provider::Provider; use crate::record::DataStream; use crate::router::{Body, BoxFuture, Method, Response, Router, SchemaSource}; @@ -73,9 +72,9 @@ impl Client { Ok(serde_json::from_value(r.entity.unwrap_or(Value::Null))?) } - pub async fn put(&self, path: &str, data: Dataset) -> Result { + pub async fn put(&self, path: &str, data: DataStream) -> Result { let body = Body { - data: Some(data.into_stream()), + data: Some(data), meta: Value::Null, }; let r = self.dispatch(Method::Put, path, Some(body)).await?; From cf61edeb8928012de81d5920074a6def132a4c22 Mon Sep 17 00:00:00 2001 From: Ferran Date: Fri, 31 Jul 2026 13:34:04 +0200 Subject: [PATCH 3/4] More --- crates/strata/src/graphql/mod.rs | 2 +- crates/strata/src/providers/iceberg/mod.rs | 11 +++++++---- crates/strata/src/providers/sql/mod.rs | 1 + 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/crates/strata/src/graphql/mod.rs b/crates/strata/src/graphql/mod.rs index 79fe91b..db5afbc 100644 --- a/crates/strata/src/graphql/mod.rs +++ b/crates/strata/src/graphql/mod.rs @@ -404,7 +404,7 @@ mod tests { rows: &[T], ) -> Result<()> { let body = crate::Body { - data: Some(crate::DataStream::of(rows)?.into_stream()), + data: Some(crate::DataStream::of(rows)?), meta: Value::Null, }; registry diff --git a/crates/strata/src/providers/iceberg/mod.rs b/crates/strata/src/providers/iceberg/mod.rs index b368120..f31d0b4 100644 --- a/crates/strata/src/providers/iceberg/mod.rs +++ b/crates/strata/src/providers/iceberg/mod.rs @@ -440,11 +440,14 @@ mod tests { let c = client(&dir)?; const ROWS: usize = 100; - let generator = crate::datagen::Generator::new(&Event::schema())?; - let dataset = generator.stream(ROWS)?; - let expected: Vec = dataset.records.decode(&dataset.schema)?; + let schema = Event::schema(); + let generator = crate::datagen::Generator::new(&schema)?; + let batch = generator.rows(0..ROWS)?; + let expected: Vec = batch.decode(&schema)?; - let result: WriteResult = c.put("/tables/events", dataset).await?; + let result: WriteResult = c + .put("/tables/events", DataStream::once(schema, batch)) + .await?; assert!(result.created); assert_eq!(result.rows_written, ROWS as u64); diff --git a/crates/strata/src/providers/sql/mod.rs b/crates/strata/src/providers/sql/mod.rs index a5d8fd0..c9fcd3a 100644 --- a/crates/strata/src/providers/sql/mod.rs +++ b/crates/strata/src/providers/sql/mod.rs @@ -550,6 +550,7 @@ fn next_cursor(cursor: &SqlCursor, returned: usize) -> Result { pub mod suite { use super::{TableName, WriteResult}; use crate::provider::Provider; + use crate::record::DataStream; use crate::testkit::Client; use anyhow::Result; use schema::{DataType, HasSchema}; From 598a21d54a41e66b0ac6c5447930e11ff4f54a75 Mon Sep 17 00:00:00 2001 From: Ferran Date: Fri, 31 Jul 2026 13:37:59 +0200 Subject: [PATCH 4/4] Remove file --- crates/strata/src/dataset.rs | 41 ------------------- crates/strata/src/lib.rs | 4 +- crates/strata/src/providers/clickhouse/mod.rs | 3 +- crates/strata/src/providers/iceberg/mod.rs | 3 +- crates/strata/src/providers/mysql/mod.rs | 9 +--- crates/strata/src/providers/postgres/mod.rs | 3 +- crates/strata/src/providers/sql/mod.rs | 15 ++++--- crates/strata/src/providers/sqlite/mod.rs | 9 +--- crates/strata/src/record.rs | 39 +++++++++++++++++- crates/strata/src/request.rs | 8 +++- crates/strata/src/router.rs | 3 +- 11 files changed, 63 insertions(+), 74 deletions(-) delete mode 100644 crates/strata/src/dataset.rs diff --git a/crates/strata/src/dataset.rs b/crates/strata/src/dataset.rs deleted file mode 100644 index 4ba49c2..0000000 --- a/crates/strata/src/dataset.rs +++ /dev/null @@ -1,41 +0,0 @@ -//! How a sink applies a written dataset. - -use anyhow::{Result, bail}; -use serde::Serialize; - -/// How a sink should apply a written dataset. Rides as metadata on the existing -/// `put` verb (the reserved `disposition` query param) rather than a new verb, so -/// every sink shares one write surface. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize)] -pub enum Disposition { - /// Insert every row (the historical behavior). Re-running adds duplicates. - #[default] - Append, - /// Idempotent write-by-key: upsert each row on the dataset's key fields - /// (insert, or update the non-key columns on conflict). Requires the schema - /// to declare a key; this is what makes a re-fetching pipe dedup itself. - Merge, -} - -impl Disposition { - /// The reserved `put` query param that selects the disposition. - pub const PARAM: &str = "disposition"; - - /// Parse the param value (`append` | `merge`/`upsert`); defaults to `Append` - /// when absent. - pub fn from_param(value: Option<&str>) -> Result { - match value { - None | Some("append") => Ok(Disposition::Append), - Some("merge") | Some("upsert") => Ok(Disposition::Merge), - Some(other) => bail!("unknown write disposition `{other}` (append|merge)"), - } - } - - /// The param value (inverse of [`from_param`](Self::from_param)). - pub fn as_param(self) -> &'static str { - match self { - Disposition::Append => "append", - Disposition::Merge => "merge", - } - } -} diff --git a/crates/strata/src/lib.rs b/crates/strata/src/lib.rs index 78524b3..7087f5d 100644 --- a/crates/strata/src/lib.rs +++ b/crates/strata/src/lib.rs @@ -4,7 +4,6 @@ pub mod catalog; pub mod config; pub mod datagen; -pub mod dataset; pub mod flight; pub mod graphql; pub mod harness; @@ -25,10 +24,9 @@ use anyhow::Result; pub use catalog::Catalog; pub use config::{Config, ProviderConfig}; -pub use dataset::Disposition; pub use page::{Cursor, Page}; pub use provider::{Provider, ProviderObject, Registry}; -pub use record::DataStream; +pub use record::{DataStream, Disposition}; pub use router::{Body, EndpointInfo, Method, Params, Response, Router}; /// Config file consulted by [`registry`] when no explicit path is given. diff --git a/crates/strata/src/providers/clickhouse/mod.rs b/crates/strata/src/providers/clickhouse/mod.rs index 22a35ec..0801031 100644 --- a/crates/strata/src/providers/clickhouse/mod.rs +++ b/crates/strata/src/providers/clickhouse/mod.rs @@ -5,10 +5,9 @@ use schema::{DataType, Field, Schema}; use serde::Deserialize; use serde_json::Value; -use crate::dataset::Disposition; -use crate::record::Batch; use crate::provider::Provider; use crate::providers::sql::{self, Filter, SqlCursor, SqlError, SqlSource, WriteResult, quote_str}; +use crate::record::{Batch, Disposition}; use crate::router::Router; #[config] diff --git a/crates/strata/src/providers/iceberg/mod.rs b/crates/strata/src/providers/iceberg/mod.rs index f31d0b4..bdbf2ab 100644 --- a/crates/strata/src/providers/iceberg/mod.rs +++ b/crates/strata/src/providers/iceberg/mod.rs @@ -29,10 +29,9 @@ mod convert; use convert::{align_to, iceberg_to_strata_schema, strata_to_iceberg_schema}; -use crate::dataset::Disposition; use crate::page::{Cursor, ListStrategy, Page}; use crate::provider::Provider; -use crate::record::{Batch, BatchPage, DataStream}; +use crate::record::{Batch, BatchPage, DataStream, Disposition}; use crate::router::{Pages, Params, Route, Router}; /// All strata tables live in one namespace. diff --git a/crates/strata/src/providers/mysql/mod.rs b/crates/strata/src/providers/mysql/mod.rs index 89534ca..1714f92 100644 --- a/crates/strata/src/providers/mysql/mod.rs +++ b/crates/strata/src/providers/mysql/mod.rs @@ -5,12 +5,11 @@ use serde_json::Value; use sqlx::Row as _; use sqlx::mysql::{MySqlArguments, MySqlPool}; -use crate::dataset::Disposition; -use crate::record::Batch; use crate::provider::Provider; use crate::providers::sql::{ self, Filter, SqlCursor, SqlError, SqlSource, WriteResult, is_table_not_found, quote_str, }; +use crate::record::{Batch, Disposition}; use crate::router::Router; /// Placeholders per statement. MySQL's protocol caps them at 65535; stay well @@ -167,11 +166,7 @@ impl SqlSource for Mysql { if !keys.is_empty() { cols.push(format!("PRIMARY KEY ({})", quote_idents(&key_refs))); } - let ddl = format!( - "CREATE TABLE {} ({})", - quote_ident(table), - cols.join(", ") - ); + let ddl = format!("CREATE TABLE {} ({})", quote_ident(table), cols.join(", ")); sqlx::query(&ddl) .execute(&pool) .await diff --git a/crates/strata/src/providers/postgres/mod.rs b/crates/strata/src/providers/postgres/mod.rs index 15bd3a0..7a3817c 100644 --- a/crates/strata/src/providers/postgres/mod.rs +++ b/crates/strata/src/providers/postgres/mod.rs @@ -3,12 +3,11 @@ use schema::{DataType, Field, Schema}; use serde_json::Value; use tokio_postgres::NoTls; -use crate::dataset::Disposition; -use crate::record::Batch; use crate::provider::Provider; use crate::providers::sql::{ self, Filter, SqlCursor, SqlError, SqlSource, WriteResult, is_table_not_found, }; +use crate::record::{Batch, Disposition}; use crate::router::Router; use config_macro::config; diff --git a/crates/strata/src/providers/sql/mod.rs b/crates/strata/src/providers/sql/mod.rs index c9fcd3a..8bdf6f5 100644 --- a/crates/strata/src/providers/sql/mod.rs +++ b/crates/strata/src/providers/sql/mod.rs @@ -22,10 +22,9 @@ use schema::{DataType, HasSchema, Schema}; use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::dataset::Disposition; use crate::page::{Cursor, ListStrategy, Page}; use crate::provider::Provider; -use crate::record::{Batch, BatchPage, DataStream, stringify_text_columns}; +use crate::record::{Batch, BatchPage, DataStream, Disposition, stringify_text_columns}; use crate::router::{Pages, Params, Route, Router}; mod filter; @@ -584,7 +583,9 @@ pub mod suite { id: 1, name: "a".into(), }]; - let _: WriteResult = client.put("/tables/catalog", DataStream::of(&rows)?).await?; + let _: WriteResult = client + .put("/tables/catalog", DataStream::of(&rows)?) + .await?; let mut tables = client.list("/tables").await?; let names: Vec = tables.next().await?; @@ -736,7 +737,9 @@ pub mod suite { name: "d".into(), }, ]; - let _: WriteResult = client.put("/tables/filtered", DataStream::of(&rows)?).await?; + let _: WriteResult = client + .put("/tables/filtered", DataStream::of(&rows)?) + .await?; let predicate = serde_json::json!({ "and": [ @@ -778,7 +781,9 @@ pub mod suite { name: "b".into(), }, ]; - let _: WriteResult = client.put("/tables/projected", DataStream::of(&rows)?).await?; + let _: WriteResult = client + .put("/tables/projected", DataStream::of(&rows)?) + .await?; let mut stream = client .list::("/tables/projected?fields=name") diff --git a/crates/strata/src/providers/sqlite/mod.rs b/crates/strata/src/providers/sqlite/mod.rs index 9e3d06e..847c215 100644 --- a/crates/strata/src/providers/sqlite/mod.rs +++ b/crates/strata/src/providers/sqlite/mod.rs @@ -23,12 +23,11 @@ use serde_json::Value; use sqlx::Row as _; use sqlx::sqlite::{SqliteArguments, SqliteConnectOptions, SqlitePool}; -use crate::dataset::Disposition; -use crate::record::Batch; use crate::provider::Provider; use crate::providers::sql::{ self, Filter, SqlCursor, SqlError, SqlSource, WriteResult, is_table_not_found, quote_str, }; +use crate::record::{Batch, Disposition}; use crate::router::Router; /// Bound parameters per statement. SQLite's own cap is 999 before 3.32 and higher @@ -188,11 +187,7 @@ impl SqlSource for Sqlite { if !keys.is_empty() { cols.push(format!("PRIMARY KEY ({})", quote_idents(&key_refs))); } - let ddl = format!( - "CREATE TABLE {} ({})", - quote_ident(table), - cols.join(", ") - ); + let ddl = format!("CREATE TABLE {} ({})", quote_ident(table), cols.join(", ")); sqlx::query(&ddl) .execute(&pool) .await diff --git a/crates/strata/src/record.rs b/crates/strata/src/record.rs index 09dc4f2..b34a88b 100644 --- a/crates/strata/src/record.rs +++ b/crates/strata/src/record.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use anyhow::{Result, anyhow}; +use anyhow::{Result, anyhow, bail}; use arrow::datatypes::{DataType as ArrowType, Field, FieldRef, Fields, TimeUnit}; use arrow::record_batch::RecordBatch; use arrow_array::Array; @@ -240,6 +240,43 @@ fn data_type_from_arrow(ty: &ArrowType) -> DataType { } } +/// How a sink should apply a written dataset. Rides as metadata on the existing +/// `put` verb (the reserved `disposition` query param) rather than a new verb, so +/// every sink shares one write surface. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize)] +pub enum Disposition { + /// Insert every row (the historical behavior). Re-running adds duplicates. + #[default] + Append, + /// Idempotent write-by-key: upsert each row on the dataset's key fields + /// (insert, or update the non-key columns on conflict). Requires the schema + /// to declare a key; this is what makes a re-fetching pipe dedup itself. + Merge, +} + +impl Disposition { + /// The reserved `put` query param that selects the disposition. + pub const PARAM: &str = "disposition"; + + /// Parse the param value (`append` | `merge`/`upsert`); defaults to `Append` + /// when absent. + pub fn from_param(value: Option<&str>) -> Result { + match value { + None | Some("append") => Ok(Disposition::Append), + Some("merge") | Some("upsert") => Ok(Disposition::Merge), + Some(other) => bail!("unknown write disposition `{other}` (append|merge)"), + } + } + + /// The param value (inverse of [`from_param`](Self::from_param)). + pub fn as_param(self) -> &'static str { + match self { + Disposition::Append => "append", + Disposition::Merge => "merge", + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/strata/src/request.rs b/crates/strata/src/request.rs index 922561a..77df413 100644 --- a/crates/strata/src/request.rs +++ b/crates/strata/src/request.rs @@ -2,7 +2,7 @@ //! resume cursor or a write disposition without hand-encoding `?cursor=` onto a path //! string. Each wraps a path and yields the encoded path via `.path()`. -use crate::dataset::Disposition; +use crate::record::Disposition; use crate::router::CURSOR_PARAM; /// A read path plus an optional resume `cursor`. @@ -53,7 +53,11 @@ impl WriteRequest { /// The encoded path, with `?disposition=…` set to match the disposition. pub fn path(self) -> String { - set_param(&self.path, Disposition::PARAM, self.disposition.map(|d| d.as_param())) + set_param( + &self.path, + Disposition::PARAM, + self.disposition.map(|d| d.as_param()), + ) } } diff --git a/crates/strata/src/router.rs b/crates/strata/src/router.rs index 9a06f2f..1a61e19 100644 --- a/crates/strata/src/router.rs +++ b/crates/strata/src/router.rs @@ -21,9 +21,8 @@ use futures::StreamExt; use futures::stream::BoxStream; use crate::DataStream; -use crate::dataset::Disposition; use crate::page::{ListStrategy, Page}; -use crate::record::{Batch, BatchPage}; +use crate::record::{Batch, BatchPage, Disposition}; /// A boxed, owned future. `'static` because handlers take owned `Params` and an /// `Arc`, so nothing is borrowed across the await.