From b095c065a85471203d631ad8ffbcd84b74b89972 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sat, 14 Nov 2026 21:41:01 +0000 Subject: [PATCH 1/7] feat(python): transform files, Sheets ranges and HTTP pages with a named @sl.transform function before they load @sl.transform registers a function under its own name, or under the one it is given (@sl.transform("kpis")). Discovery imports python_paths, and a files or google_sheets table names the function with transform = "name"; sl.http's paginate and fan_out take transform= with a name or the function itself. Plan fails on a name no function has, listing the registered ones, and on a name two functions have, naming both files. Each unit gets one call, f(rows, ctx), with its rows as a pyarrow Table: a file or ZIP member, a Sheets range, or an HTTP page after select, fields and cells. ctx.unit is "file", "range" or "page", with that unit's path, range or page number and fan-out key. The function returns or yields, in any mix, what a Python table may yield, and yielded batches load as they come, so memory stays near one unit and its output. A unit whose function fails fails the table, uncommitted and, for files, not moved by after_load. Plan takes the columns from what the function returns for the newest file or the range. A table with a transform is read by SQLLoad, never directly by the destination, and the standalone executable refuses it. Each commit records the transform's version (a hash of its file, or of its code when the file can't be read) and the files it ran on, in the table's run log, created by apply. Plan warns while the table holds files an earlier version loaded, and the load report gives each file's transform time. --reload now works on files tables that merge into DuckDB: each file merges on its own into a new table beside the old one, in the order the run log merged them, then new files oldest first, and the new table replaces the old one with its run log records in one commit. With after_load = { move_to = ... }, --reload also reads the files under move_to, from where they were moved, and plan shows both folders. After the swap, files read from path move as usual and archived files stay. --- CONTRIBUTING.md | 6 + .../src/crossing/_helpers/parser.rs | 119 ++- .../sqlload-python/src/crossing/constants.rs | 3 + .../src/crossing/main/crossing_options.rs | 19 + .../sqlload-python/src/crossing/main/mod.rs | 2 + .../src/crossing/main/returned_batches.rs | 52 + crates/sqlload-python/src/crossing/models.rs | 2 + crates/sqlload-python/src/host/models.rs | 74 +- .../clickhouse/_helpers/write/direct_load.rs | 2 +- .../snowflake/main/copy/_direct_copies.rs | 2 +- .../sqlload/src/adapters/snowflake/models.rs | 2 +- .../src/adapters/snowflake/tests/helpers.rs | 6 +- crates/sqlload/src/command/_helpers/limits.rs | 5 +- .../src/command/_helpers/load_output.rs | 41 +- .../src/command/_helpers/loader_commands.rs | 41 +- crates/sqlload/src/command/tests/helpers.rs | 1 + .../src/command/tests/test_load_output.rs | 77 +- .../src/command/tests/test_loader_commands.rs | 5 + .../sqlload/src/command/tests/test_types.rs | 7 + crates/sqlload/src/file_store/models.rs | 8 + .../src/files_source/_helpers/plan/folder.rs | 159 ++++ .../src/files_source/main/batch_columns.rs | 2 +- .../src/files_source/main/batch_footers.rs | 4 +- .../src/files_source/main/batch_headers.rs | 4 +- .../src/files_source/main/batch_refusals.rs | 2 +- .../src/files_source/main/discover_files.rs | 152 +-- .../main/discover_reload_files.rs | 37 + crates/sqlload/src/files_source/main/mod.rs | 2 + .../src/files_source/main/read_files.rs | 2 +- .../src/files_source/main/shape_rows.rs | 16 + crates/sqlload/src/files_source/models.rs | 4 +- .../sqlload/src/files_source/tests/helpers.rs | 50 + .../files_source/tests/test_batch_footers.rs | 15 +- .../src/files_source/tests/test_types.rs | 2 + .../sqlload/src/http_source/_helpers/pages.rs | 2 + crates/sqlload/src/http_source/models.rs | 24 + crates/sqlload/src/load_events/models.rs | 3 + .../src/loader/_helpers/destination.rs | 132 ++- crates/sqlload/src/loader/_helpers/publish.rs | 177 ++-- crates/sqlload/src/loader/_helpers/python.rs | 10 +- crates/sqlload/src/loader/_helpers/render.rs | 129 ++- crates/sqlload/src/loader/_helpers/resolve.rs | 182 +++- crates/sqlload/src/loader/_helpers/run_log.rs | 363 ++++++- crates/sqlload/src/loader/_helpers/sources.rs | 889 ++++++++---------- crates/sqlload/src/loader/_helpers/specs.rs | 268 +++++- .../sqlload/src/loader/_helpers/table_plan.rs | 32 +- .../src/loader/_helpers/table_queue.rs | 1 + crates/sqlload/src/loader/constants.rs | 12 + .../sqlload/src/loader/main/apply_project.rs | 52 +- .../sqlload/src/loader/main/plan_project.rs | 6 +- .../sqlload/src/loader/main/selected_specs.rs | 4 +- crates/sqlload/src/loader/models.rs | 65 ++ crates/sqlload/src/loader/tests.rs | 4 + crates/sqlload/src/loader/tests/helpers.rs | 121 ++- .../src/loader/tests/test_merged_files.rs | 68 +- .../src/loader/tests/test_transform_reader.rs | 28 + .../loader/tests/test_transform_versions.rs | 83 ++ crates/sqlload/src/loader/tests/test_types.rs | 32 + .../src/project/_helpers/references.rs | 31 + crates/sqlload/src/project/constants.rs | 1 + crates/sqlload/src/project/models.rs | 9 + .../src/project/tests/test_parse_project.rs | 15 + .../src/python_sources/main/_file_context.rs | 36 + .../src/python_sources/main/_run_transform.rs | 65 ++ crates/sqlload/src/python_sources/main/mod.rs | 4 + .../python_sources/main/transform_files.rs | 56 ++ .../python_sources/main/transform_range.rs | 40 + crates/sqlload/src/python_sources/models.rs | 69 ++ crates/sqlload/src/python_sources/types.rs | 18 +- .../sqlload/src/run_log/main/event_record.rs | 1 + crates/sqlload/src/run_log/models.rs | 3 + crates/sqlload/src/run_log/tests/helpers.rs | 1 + .../src/sheets_source/_helpers/cells.rs | 11 + crates/sqlload/src/sheets_source/main/mod.rs | 1 + .../src/sheets_source/main/sheet_input.rs | 38 + crates/sqlload/src/sheets_source/tests.rs | 2 + .../src/sheets_source/tests/helpers.rs | 80 ++ .../sheets_source/tests/test_sheet_input.rs | 56 ++ .../src/sheets_source/tests/test_types.rs | 8 + .../tests/loader/after_load/helpers.rs | 40 +- crates/sqlload/tests/loader/helpers.rs | 112 +-- crates/sqlload/tests/loader/merge/helpers.rs | 100 +- .../sqlload/tests/loader/test_after_load.rs | 99 ++ .../tests/loader/test_files_to_duckdb.rs | 275 +++++- .../tests/loader/test_postgres_target.rs | 16 +- crates/sqlload/tests/loader/test_types.rs | 26 + packages/sqlload/python/sqlload/__init__.py | 3 + packages/sqlload/python/sqlload/_http.py | 52 +- packages/sqlload/python/sqlload/_registry.py | 103 ++ packages/sqlload/python/sqlload/_runtime.py | 65 ++ packages/sqlload/python/sqlload/_transform.py | 113 +++ packages/sqlload/tests/test_http_jobs.py | 86 +- packages/sqlload/tests/test_transforms.py | 573 +++++++++++ website/src/content/docs/docs/cli/load.mdx | 8 +- .../src/content/docs/docs/configuration.mdx | 1 + .../src/content/docs/docs/guides/from-dlt.mdx | 10 + .../src/content/docs/docs/sources/files.mdx | 203 +++- .../docs/docs/sources/google-sheets.mdx | 65 ++ .../src/content/docs/docs/sources/python.mdx | 82 +- 99 files changed, 5212 insertions(+), 977 deletions(-) create mode 100644 crates/sqlload-python/src/crossing/main/crossing_options.rs create mode 100644 crates/sqlload-python/src/crossing/main/returned_batches.rs create mode 100644 crates/sqlload/src/files_source/main/discover_reload_files.rs create mode 100644 crates/sqlload/src/files_source/main/shape_rows.rs create mode 100644 crates/sqlload/src/loader/tests/test_transform_reader.rs create mode 100644 crates/sqlload/src/loader/tests/test_transform_versions.rs create mode 100644 crates/sqlload/src/python_sources/main/_file_context.rs create mode 100644 crates/sqlload/src/python_sources/main/_run_transform.rs create mode 100644 crates/sqlload/src/python_sources/main/transform_files.rs create mode 100644 crates/sqlload/src/python_sources/main/transform_range.rs create mode 100644 crates/sqlload/src/sheets_source/main/sheet_input.rs create mode 100644 crates/sqlload/src/sheets_source/tests/test_sheet_input.rs create mode 100644 packages/sqlload/python/sqlload/_registry.py create mode 100644 packages/sqlload/python/sqlload/_transform.py create mode 100644 packages/sqlload/tests/test_transforms.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 014e1255f..305899970 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -26,6 +26,12 @@ reads those files and masks the column itself. The plan warns when this happens. server are not an exception: no warehouse reads SFTP, so they never had a direct route, and `plan` shows them read over SFTP. +A table with a `transform` (from 0.52.0) is read by SQLLoad too, since its Python function runs on +each file's rows before they load: `reader_offer` gives the reason `transform runs on +each file`, and `plan_route` takes the same path as for `hmac`. Without a `transform`, nothing else +changes, and the route test below still checks that only `hmac` reaches SQLLoad's reader among the +settings it plans. + Anything else fails rather than falling back to SQLLoad: - Files no external stage covers fail the plan, and the error says how to fix it. diff --git a/crates/sqlload-python/src/crossing/_helpers/parser.rs b/crates/sqlload-python/src/crossing/_helpers/parser.rs index ae09605db..996cc71ce 100644 --- a/crates/sqlload-python/src/crossing/_helpers/parser.rs +++ b/crates/sqlload-python/src/crossing/_helpers/parser.rs @@ -6,8 +6,11 @@ use std::sync::{Arc, Mutex}; use std::time::Instant; use arrow_array::RecordBatch; +use arrow_schema::Schema; use sqlload::http_source::main::run_http_job::run_http_job; +use sqlload::http_source::models::{HttpJob, HttpPage, PageTransform}; use sqlload::load_naming::main::loaded_name::loaded_name; +use sqlload::python_sources::models::{TransformCall, TransformContext}; use sqlload::row_batches::main::columns_batch::columns_batch; use sqlload::row_batches::main::flatten_batch::flatten_batch; use sqlload::row_batches::main::json_rows::json_rows; @@ -15,9 +18,13 @@ use sqlload::row_batches::main::normalise_batch::normalise_batch; use sqlload::row_batches::main::projection_problem::projection_problem; use sqlload::row_batches::models::{JsonLayout, JsonRead}; +use crate::crossing::constants::PAGE_POLL; use crate::crossing::errors::CrossError; +use crate::crossing::main::crossing_options::crossing_options; +use crate::crossing::main::returned_batches::returned_batches; use crate::crossing::models::{ - Crossed, CrossingOptions, CrossingStats, Dispatch, PanicSite, ParseJob, Parsed, ReadContext, + Crossed, CrossingOptions, CrossingStats, Dispatch, PanicSite, ParseJob, Parsed, Polled, + ReadContext, }; /// Where one piece's batches go, and the read's counters. @@ -256,8 +263,16 @@ impl Output { normalise_batch(&flattened, &arrays).map_err(CrossError::Data) } - /// Run an `sl.http` job, parsing each page into batches as it arrives; `Ok(false)` once the consumer has gone, which cancels the requests in flight. + /// Run an `sl.http` job, parsing each page into batches as it arrives, through the job's `transform` when it has one; `Ok(false)` once the consumer has gone, which cancels the requests in flight. fn http(&self, definition: &str) -> Result { + let transform = match serde_json::from_str::(definition) { + Ok(job) => job.transform().cloned(), + Err(error) => { + return Err(CrossError::Http( + format!("cannot read the sl.http job: {error}").into(), + )); + } + }; let mut open = true; let mut failed: Option = None; run_http_job( @@ -265,22 +280,9 @@ impl Output { &self.context.connections, self.context.rules(), |page| { - let read = JsonRead { - layout: JsonLayout::Document, - path: page.path.clone(), - batch_rows: usize::MAX, - projection: page.projection.clone(), - }; - let parsed = json_rows(&page.bytes, &read, &mut |columns| { - let batch = flatten_batch(&columns_batch(columns)?, &self.context.flatten)?; - open = open && self.send(batch, Instant::now()); - match open { - true => Ok(()), - false => Err("the consumer stopped reading".to_owned()), - } - }); - if let (Err(reason), true) = (parsed, open) { - failed = Some(format!("{}: {reason}", page.source)); + match self.page_rows(&page, transform.as_ref()) { + Ok(still) => open = still, + Err(reason) => failed = Some(format!("{}: {reason}", page.source)), } open && failed.is_none() }, @@ -292,6 +294,87 @@ impl Output { } } + /// Parse `page` into batches and pass them on, or, given a `transform`, pass on what it returns or yields for the page's rows; `Ok(false)` once the consumer has gone. + fn page_rows( + &self, + page: &HttpPage, + transform: Option<&PageTransform>, + ) -> Result { + let read = JsonRead { + layout: JsonLayout::Document, + path: page.path.clone(), + batch_rows: usize::MAX, + projection: page.projection.clone(), + }; + let mut open = true; + let mut held: Vec = Vec::new(); + let parsed = json_rows(&page.bytes, &read, &mut |columns| { + let batch = columns_batch(columns)?; + if transform.is_some() { + held.push(batch); + return Ok(()); + } + let batch = flatten_batch(&batch, &self.context.flatten)?; + open = open && self.send(batch, Instant::now()); + match open { + true => Ok(()), + false => Err("the consumer stopped reading".to_owned()), + } + }); + match (parsed, open, transform) { + (Err(reason), true, _) => Err(reason), + (_, false, _) | (Ok(()), true, None) => Ok(open), + (Ok(()), true, Some(transform)) => self + .transformed_page((transform, page), held) + .map_err(|reason| format!("transform {} failed: {reason}", transform.name)), + } + } + + /// Call `transform` once on one page's `rows`, with the page's number and fan-out key as its `ctx`, and pass on what it returns or yields as it crosses; `Ok(false)` once the consumer has gone. + fn transformed_page( + &self, + (transform, page): (&PageTransform, &HttpPage), + rows: Vec, + ) -> Result { + let schema = rows + .first() + .map_or_else(|| Arc::new(Schema::empty()), RecordBatch::schema); + let call = TransformCall { + function: transform.function.clone(), + context: TransformContext { + unit: "page", + table: self.context.table.clone(), + page: Some(page.number), + key: (!page.key.is_empty()).then(|| page.key.clone()), + ..TransformContext::default() + }, + masked: self.context.masked.clone(), + arrays: self.context.arrays.clone(), + naming: self.context.naming, + }; + let options = CrossingOptions { + parse_threads: 1, + ..crossing_options() + }; + let (mut batches, _returned) = returned_batches( + &call, + (rows, schema), + (options, self.context.flatten.clone()), + )?; + loop { + match batches.poll(PAGE_POLL) { + Polled::Batch(batch) => { + if !self.send(batch, Instant::now()) { + return Ok(false); + } + } + Polled::Finished => return Ok(true), + Polled::Failed(error) => return Err(error.to_string()), + Polled::State(_) | Polled::Pending => {} + } + } + } + /// Count `batch`, parsed since `started`, and pass it on; `false` once the consumer has gone. fn send(&self, batch: RecordBatch, started: Instant) -> bool { CrossingStats::add(&self.stats.parse_nanos, started.elapsed()); diff --git a/crates/sqlload-python/src/crossing/constants.rs b/crates/sqlload-python/src/crossing/constants.rs index 86b1c5d90..2753479a6 100644 --- a/crates/sqlload-python/src/crossing/constants.rs +++ b/crates/sqlload-python/src/crossing/constants.rs @@ -21,3 +21,6 @@ pub(crate) const HTTP_JOB_METHOD: &str = "__sqlload_http_job__"; /// The method of an `sl.checkpoint` that returns its state as JSON. pub(crate) const CHECKPOINT_METHOD: &str = "__sqlload_checkpoint__"; + +/// How long a page `transform`'s crossing is waited on at a time; it polls again until the function's output ends. +pub(crate) const PAGE_POLL: std::time::Duration = std::time::Duration::from_secs(3600); diff --git a/crates/sqlload-python/src/crossing/main/crossing_options.rs b/crates/sqlload-python/src/crossing/main/crossing_options.rs new file mode 100644 index 000000000..638c73367 --- /dev/null +++ b/crates/sqlload-python/src/crossing/main/crossing_options.rs @@ -0,0 +1,19 @@ +//! The crossing settings every Python read uses. + +use crate::bindings::constants::{ + DEFAULT_BATCH_ROWS, DEFAULT_DICT_CROSSING, DEFAULT_MAX_QUEUED, DEFAULT_MAX_QUEUED_BYTES, +}; +use crate::crossing::main::default_parse_threads::default_parse_threads; +use crate::crossing::models::CrossingOptions; + +/// How a table function's yields, or what a `transform` returned, cross into Arrow. +pub fn crossing_options() -> CrossingOptions { + CrossingOptions { + batch_rows: DEFAULT_BATCH_ROWS, + max_queued: DEFAULT_MAX_QUEUED, + dict_crossing: DEFAULT_DICT_CROSSING, + parse_threads: default_parse_threads(), + panic_at: None, + max_queued_bytes: DEFAULT_MAX_QUEUED_BYTES, + } +} diff --git a/crates/sqlload-python/src/crossing/main/mod.rs b/crates/sqlload-python/src/crossing/main/mod.rs index 165513e4a..eb384aeb8 100644 --- a/crates/sqlload-python/src/crossing/main/mod.rs +++ b/crates/sqlload-python/src/crossing/main/mod.rs @@ -1,2 +1,4 @@ +pub mod crossing_options; pub mod default_parse_threads; +pub mod returned_batches; pub mod table_batches; diff --git a/crates/sqlload-python/src/crossing/main/returned_batches.rs b/crates/sqlload-python/src/crossing/main/returned_batches.rs new file mode 100644 index 000000000..291ee9b2a --- /dev/null +++ b/crates/sqlload-python/src/crossing/main/returned_batches.rs @@ -0,0 +1,52 @@ +//! Call a `transform` and start crossing what it returns or yields. + +use arrow_array::RecordBatch; +use arrow_pyarrow::{IntoPyArrow, Table}; +use arrow_schema::SchemaRef; +use pyo3::exceptions::PyValueError; +use pyo3::types::PyAnyMethods; +use pyo3::{Bound, Py, PyAny, PyResult, Python}; +use sqlload::python_sources::models::TransformCall; + +use crate::crossing::main::table_batches::table_batches; +use crate::crossing::models::{CrossingOptions, ReadContext, TableBatches}; +use crate::host::constants::RUNTIME_MODULE; + +/// Call `call`'s transform on `rows` (in `schema`) and start crossing what it returned or yielded with `options`, as a table function's yields cross, expanding the objects `flatten` lists: the batches, and what it returned (whose `dicts` says whether any item was dicts); or why it failed, or what it returned cannot load. +pub fn returned_batches( + call: &TransformCall, + (rows, schema): (Vec, SchemaRef), + (options, flatten): (CrossingOptions, Vec), +) -> Result<(TableBatches, Py), String> { + let context = serde_json::to_string(&call.context).map_err(|error| error.to_string())?; + Python::attach( + |py| -> PyResult), String>> { + let table = Table::try_new(rows, schema) + .map_err(|error| PyValueError::new_err(error.to_string()))? + .into_pyarrow(py)?; + let (returned, problem): (Option>, Option) = py + .import(RUNTIME_MODULE)? + .call_method1("call_transform", (call.function.as_str(), table, context))? + .extract()?; + let Some(returned) = returned else { + return Ok(Err( + problem.unwrap_or_else(|| "it returned nothing that loads".to_owned()) + )); + }; + let batches = table_batches( + &returned, + None, + options, + ReadContext { + flatten, + masked: call.masked.clone(), + arrays: call.arrays.clone(), + naming: call.naming, + ..ReadContext::default() + }, + ); + Ok(Ok((batches, returned.unbind()))) + }, + ) + .map_err(|error| error.to_string())? +} diff --git a/crates/sqlload-python/src/crossing/models.rs b/crates/sqlload-python/src/crossing/models.rs index 1745c0f72..bed7b72d2 100644 --- a/crates/sqlload-python/src/crossing/models.rs +++ b/crates/sqlload-python/src/crossing/models.rs @@ -42,6 +42,8 @@ pub struct CrossingOptions { /// What the parser threads need from the table besides its data: the top-level objects to expand one level, and the connections its `sl.http` jobs may use. #[derive(Debug, Clone, Default)] pub struct ReadContext { + /// The table's name in its source, for a page `transform`'s `ctx.table`. + pub table: String, pub flatten: Vec, pub connections: serde_json::Map, /// The table's cursor column, which projected JSON rows must keep. diff --git a/crates/sqlload-python/src/host/models.rs b/crates/sqlload-python/src/host/models.rs index b543955d6..d7272ac00 100644 --- a/crates/sqlload-python/src/host/models.rs +++ b/crates/sqlload-python/src/host/models.rs @@ -1,19 +1,19 @@ use std::path::{Path, PathBuf}; use std::time::Duration; +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; use pyo3::types::PyAnyMethods; -use pyo3::{PyResult, Python}; +use pyo3::{Py, PyAny, PyResult, Python}; use sqlload::python_sources::errors::ReadError; -use sqlload::python_sources::models::{PythonSource, TableRead}; +use sqlload::python_sources::models::{PythonSource, TableRead, TransformCall}; use sqlload::python_sources::types::{BatchStream, PythonHost, Streamed}; -use crate::bindings::constants::{ - DEFAULT_BATCH_ROWS, DEFAULT_DICT_CROSSING, DEFAULT_MAX_QUEUED, DEFAULT_MAX_QUEUED_BYTES, -}; use crate::crossing::errors::CrossError; -use crate::crossing::main::default_parse_threads::default_parse_threads; +use crate::crossing::main::crossing_options::crossing_options; +use crate::crossing::main::returned_batches::returned_batches; use crate::crossing::main::table_batches::table_batches; -use crate::crossing::models::{CrossingOptions, Polled, ReadContext, TableBatches}; +use crate::crossing::models::{Polled, ReadContext, TableBatches}; use crate::host::_helpers::python_errors::{error_text, read_error}; use crate::host::constants::RUNTIME_MODULE; @@ -39,14 +39,7 @@ impl PythonHost for ExtensionHost { fn read_table(&self, request: &TableRead) -> Result, String> { let text = serde_json::to_string(request).map_err(|error| error.to_string())?; - let options = CrossingOptions { - batch_rows: DEFAULT_BATCH_ROWS, - max_queued: DEFAULT_MAX_QUEUED, - dict_crossing: DEFAULT_DICT_CROSSING, - parse_threads: default_parse_threads(), - panic_at: None, - max_queued_bytes: DEFAULT_MAX_QUEUED_BYTES, - }; + let options = crossing_options(); Python::attach(|py| -> PyResult { let runtime = py.import(RUNTIME_MODULE)?; let iterator = runtime.call_method1("start_table", (text,))?; @@ -63,6 +56,7 @@ impl PythonHost for ExtensionHost { snapshot.as_ref(), options, ReadContext { + table: request.table.clone(), flatten: request.flatten.clone(), connections: request.connections.clone(), cursor_column: request.cursor.as_ref().map(|cursor| cursor.column.clone()), @@ -85,6 +79,33 @@ impl PythonHost for ExtensionHost { .map_err(|error| error_text(&error)) } + fn transform_version(&self, name: &str) -> Result { + let checked = Python::attach(|py| -> PyResult<(Option, Option)> { + py.import(RUNTIME_MODULE)? + .call_method1("transform_version", (name,))? + .extract() + }) + .map_err(|error| error_text(&error))?; + match checked { + (Some(version), _) => Ok(version), + (None, problem) => Err(problem.unwrap_or_else(|| format!("{name} cannot run"))), + } + } + + fn transform( + &self, + call: &TransformCall, + rows: Vec, + schema: SchemaRef, + ) -> Result, String> { + let (batches, returned) = + returned_batches(call, (rows, schema), (crossing_options(), Vec::new()))?; + Ok(Box::new(TransformBatches { + batches: PythonBatches(batches), + returned, + })) + } + fn finish(&self) { let closed = Python::attach(|py| -> PyResult<()> { py.import(RUNTIME_MODULE)?.call_method0("finish_run")?; @@ -115,3 +136,26 @@ impl BatchStream for PythonBatches { } } } + +/// What a `transform` returned or yielded, as the engine reads it, and the Python object noting whether any of it was dicts. +#[derive(Debug)] +pub struct TransformBatches { + batches: PythonBatches, + returned: Py, +} + +impl BatchStream for TransformBatches { + fn next_batch(&mut self) -> Option> { + self.batches.next_batch() + } + + fn crossed_dicts(&self) -> bool { + Python::attach(|py| { + self.returned + .bind(py) + .getattr("dicts") + .and_then(|dicts| dicts.extract::()) + .is_ok_and(|dicts| dicts) + }) + } +} diff --git a/crates/sqlload/src/adapters/clickhouse/_helpers/write/direct_load.rs b/crates/sqlload/src/adapters/clickhouse/_helpers/write/direct_load.rs index 0086eaf71..39e477627 100644 --- a/crates/sqlload/src/adapters/clickhouse/_helpers/write/direct_load.rs +++ b/crates/sqlload/src/adapters/clickhouse/_helpers/write/direct_load.rs @@ -146,7 +146,7 @@ fn pulled_file( (batch, columns): (&DirectBatch<'_>, &[DesiredColumn]), ) -> Result { let mut in_file: Vec = Vec::new(); - if let Some(footer) = batch.footers.parquet.get(&data.file.relative_key) { + if let Some(footer) = batch.footers.parquet.get(&data.file.store_key()) { in_file.extend( footer .schema() diff --git a/crates/sqlload/src/adapters/snowflake/main/copy/_direct_copies.rs b/crates/sqlload/src/adapters/snowflake/main/copy/_direct_copies.rs index 31ab05c1d..47772bd8a 100644 --- a/crates/sqlload/src/adapters/snowflake/main/copy/_direct_copies.rs +++ b/crates/sqlload/src/adapters/snowflake/main/copy/_direct_copies.rs @@ -13,7 +13,7 @@ pub(in crate::adapters) fn direct_copies( for data in request.files { let header = request .headers - .get(&data.file.relative_key) + .get(&data.file.store_key()) .map(Vec::as_slice); match groups.iter_mut().find(|(seen, _)| *seen == header) { Some((_, files)) => files.push(data.clone()), diff --git a/crates/sqlload/src/adapters/snowflake/models.rs b/crates/sqlload/src/adapters/snowflake/models.rs index 9e9a823bd..f78eb9d9c 100644 --- a/crates/sqlload/src/adapters/snowflake/models.rs +++ b/crates/sqlload/src/adapters/snowflake/models.rs @@ -171,7 +171,7 @@ pub struct DirectCopy<'a> { pub encoding: &'a FilesEncoding, /// Columns whose types were inferred from sampled values, whose every value must fit the inferred type (JSON, NDJSON and CSV). pub inferred: &'a [String], - /// CSV files' column names by relative key, which place each column. + /// CSV files' column names by store key, which place each column. pub headers: &'a std::collections::BTreeMap>, } diff --git a/crates/sqlload/src/adapters/snowflake/tests/helpers.rs b/crates/sqlload/src/adapters/snowflake/tests/helpers.rs index d91c915d3..70189f902 100644 --- a/crates/sqlload/src/adapters/snowflake/tests/helpers.rs +++ b/crates/sqlload/src/adapters/snowflake/tests/helpers.rs @@ -231,16 +231,16 @@ pub(crate) fn direct_sql( from: &CopyFrom, (encoding, header): (&crate::files_source::models::FilesEncoding, &[&str]), ) -> Result<(usize, String), String> { + let files: Vec = files.iter().map(|(key, day)| data_file(key, day)).collect(); let headers: std::collections::BTreeMap> = files .iter() - .map(|(key, _)| { + .map(|data| { ( - key.clone(), + data.file.store_key(), header.iter().map(|name| (*name).to_owned()).collect(), ) }) .collect(); - let files: Vec = files.iter().map(|(key, day)| data_file(key, day)).collect(); let source = lake(from.source); let batches = crate::adapters::snowflake::main::copy::_direct_copies::direct_copies( &crate::adapters::snowflake::models::DirectCopy { diff --git a/crates/sqlload/src/command/_helpers/limits.rs b/crates/sqlload/src/command/_helpers/limits.rs index 677d55baa..e67390479 100644 --- a/crates/sqlload/src/command/_helpers/limits.rs +++ b/crates/sqlload/src/command/_helpers/limits.rs @@ -1,3 +1,4 @@ +use crate::command::constants::COMMAND_STACK_BYTES; use crate::loader::constants::{DEFAULT_THREADS, SPARE_BLOCKING_THREADS}; use crate::project::main::is_project::is_project; @@ -31,13 +32,13 @@ pub(crate) fn blocking_threads(cli_concurrency: Option, text: &str) -> Opti Some(threads.saturating_add(SPARE_BLOCKING_THREADS)) } -/// The async runtime, its blocking pool sized for `cli_threads` or the project's `concurrency`. +/// The async runtime, its blocking pool sized for `cli_threads` or the project's `concurrency`. Its threads, the workers and the blocking threads each table loads on, get the command's stack (`COMMAND_STACK_BYTES`) rather than tokio's 2 MiB, which a debug build's load into Postgres with `unnest` outgrows. That is depth, not one large future: the per-table future is about 87 KiB, but the debug poll frames nested through the publisher's commit, the child tables' staging and the Postgres execute reach about 1.8 MiB below the load task. pub(crate) fn runtime( cli_threads: Option, text: &str, ) -> Result { let mut builder = tokio::runtime::Builder::new_multi_thread(); - builder.enable_all(); + builder.enable_all().thread_stack_size(COMMAND_STACK_BYTES); if let Some(blocking) = blocking_threads(cli_threads, text) { builder.max_blocking_threads(blocking); } diff --git a/crates/sqlload/src/command/_helpers/load_output.rs b/crates/sqlload/src/command/_helpers/load_output.rs index 42474cb00..e7cdee2b7 100644 --- a/crates/sqlload/src/command/_helpers/load_output.rs +++ b/crates/sqlload/src/command/_helpers/load_output.rs @@ -10,9 +10,12 @@ use crate::command::_helpers::output::print_lines; use crate::command::constants::PLAN_TREE_HINT; use crate::load_events::main::event_line::event_line; use crate::load_events::models::{EventKind, LoadEvent}; -use crate::loader::constants::{BOLD, CYAN, DIM, INVERSE, RED, RESET, YELLOW}; +use crate::loader::constants::{ + BOLD, CYAN, DICT_HINT_MILLISECONDS, DICT_HINT_ROWS, DIM, INVERSE, RED, RESET, YELLOW, +}; use crate::loader::models::ProjectPlan; use crate::number_text::main::thousands::thousands; +use crate::python_sources::models::TransformTime; /// The width of a tally line's verb column: `loaded` or `failed` and two spaces. const VERB_COLUMN: usize = 8; @@ -231,6 +234,7 @@ pub(crate) fn event_lines( elapsed_ms, position, changes, + transforms, } => { let word = match rows { 1 => "row", @@ -238,9 +242,14 @@ pub(crate) fn event_lines( }; let seconds = format!("{:.1}s", *elapsed_ms as f64 / 1000.0); let detail = format!( - "{:>rows_width$} {word:<5}{seconds:<6}{}", - thousands(*rows), - position.as_deref().unwrap_or_default() + "{}{}", + format!( + "{:>rows_width$} {word:<5}{seconds:<6}{}", + thousands(*rows), + position.as_deref().unwrap_or_default() + ) + .trim_end(), + transform_lines(transforms) ); let mut lines: Vec = changes .iter() @@ -255,6 +264,30 @@ pub(crate) fn event_lines( vec![tally_line((verb, name), &detail, name_width)] } +/// A line per file or range a `transform` ran on, with its time, each after a line break; then, when it returned dicts for a large one and took a while, one hint naming the slowest such that Arrow or a DataFrame is faster. +fn transform_lines(times: &[TransformTime]) -> String { + let seconds = |time: &TransformTime| time.milliseconds as f64 / 1000.0; + let mut lines: String = times + .iter() + .map(|time| format!("\ntransform {} {:.2}s", time.file, seconds(time))) + .collect(); + let slow = times + .iter() + .filter(|time| { + time.dicts && time.rows >= DICT_HINT_ROWS && time.milliseconds >= DICT_HINT_MILLISECONDS + }) + .max_by_key(|time| time.milliseconds); + if let Some(time) = slow { + lines.push_str(&format!( + "\nhint: the transform returned dicts for {} ({} rows, {:.1}s); a pyarrow Table or a DataFrame transforms faster", + time.file, + thousands(time.rows), + seconds(time) + )); + } + lines +} + /// `verb name detail` in the tally's columns, a detail's further lines under its first; a verb longer than `loaded` (`created`, `altered`), or a name wider than the column (a child table made as its parent loads), keeps two spaces after it. fn tally_line((verb, name): (&str, &str), detail: &str, name_width: usize) -> String { let verb_width = VERB_COLUMN.max(verb.len() + 2); diff --git a/crates/sqlload/src/command/_helpers/loader_commands.rs b/crates/sqlload/src/command/_helpers/loader_commands.rs index fdf3f0318..aa28b7fae 100644 --- a/crates/sqlload/src/command/_helpers/loader_commands.rs +++ b/crates/sqlload/src/command/_helpers/loader_commands.rs @@ -25,7 +25,7 @@ use crate::load_events::main::close::close; use crate::load_events::main::emit::emit; use crate::load_events::main::listen::listen; use crate::load_events::models::EventKind; -use crate::loader::constants::NO_PYTHON_HOST; +use crate::loader::constants::{NO_PYTHON_HOST, NO_TRANSFORM_HOST}; use crate::loader::main::apply_project::apply_project; use crate::loader::main::connect_targets::connect_targets; use crate::loader::main::debug_project::debug_project; @@ -98,6 +98,10 @@ async fn load_project_config(request: &LoaderRequest<'_>) -> Result) -> Result { let lookup = |name: &str| variable(&request.variables, name); let project = parse_project(request.text, &lookup).map_err(|error| error.to_string())?; + let transforms = transform_places(&project); + if let (None, Some((place, _))) = (&request.host, transforms.first()) { + return Err(NO_TRANSFORM_HOST.replace("{place}", place)); + } if project.python_paths.is_empty() { return Ok(project); } @@ -110,10 +114,41 @@ async fn parsed_project(request: &LoaderRequest<'_>) -> Result Vec<(String, String)> { + let mut places: Vec<(String, String)> = Vec::new(); + for (name, source) in &project.sources { + if let Some(function) = &source.defaults.transform { + places.push((format!("source {name}"), function.clone())); + } + for (table, overrides) in &source.tables { + if let Some(function) = &overrides.settings.transform { + places.push((format!("table {name}.{table}"), function.clone())); + } + } + } + places } /// The exit code of a command that succeeded (`ok`), or failed only in `retryable` ways, or failed otherwise. diff --git a/crates/sqlload/src/command/tests/helpers.rs b/crates/sqlload/src/command/tests/helpers.rs index 6e2e2fd95..32401d4aa 100644 --- a/crates/sqlload/src/command/tests/helpers.rs +++ b/crates/sqlload/src/command/tests/helpers.rs @@ -91,6 +91,7 @@ pub(crate) fn done(rows: u64, elapsed_ms: u64, position: Option<&str>) -> EventK elapsed_ms, position: position.map(str::to_owned), changes: Vec::new(), + transforms: Vec::new(), } } diff --git a/crates/sqlload/src/command/tests/test_load_output.rs b/crates/sqlload/src/command/tests/test_load_output.rs index 43daeac0a..140eaedf8 100644 --- a/crates/sqlload/src/command/tests/test_load_output.rs +++ b/crates/sqlload/src/command/tests/test_load_output.rs @@ -1,8 +1,11 @@ use crate::command::_helpers::load_output::{Tally, event_lines, summary_lines}; use crate::command::tests::helpers::{done, read_back}; -use crate::command::tests::test_types::{LoadEventsFileTestCase, LoadOutputTestCase}; +use crate::command::tests::test_types::{ + LoadEventsFileTestCase, LoadOutputTestCase, TransformTallyTestCase, +}; use crate::load_events::constants::EVENTS_SCHEMA_VERSION; use crate::load_events::models::{EventKind, LoadEvent}; +use crate::python_sources::models::TransformTime; /// A `--fail-fast` run where every failure is retryable, as `--events jsonl` writes it. const FAIL_FAST_EVENTS: &str = include_str!("fixtures/fail_fast_retryable.jsonl"); @@ -38,6 +41,7 @@ loaded raw.b 1,309 rows 8.0s 3 files, moved to archive/receipts/ elapsed_ms: 300, position: None, changes: vec!["created raw.a__items 6 columns".to_owned()], + transforms: Vec::new(), }, )], name_width: 7, @@ -132,3 +136,74 @@ loaded workforce.people 304 rows 4.6s } Ok(()) } + +#[test] +fn given_a_table_with_a_transform_when_it_loads_then_each_file_has_its_time_and_slow_dicts_a_hint() +{ + let indent = " ".repeat(15); + let test_cases = [ + TransformTallyTestCase { + description: "each file's time on a line of its own, under the table's", + transforms: &[ + ("./drop/a.csv", 120, 10, true), + ("./drop/b.csv", 4_560, 200_000, false), + ], + expected_lines: "loaded RAW.M 7 rows 9.0s +{indent}transform ./drop/a.csv 0.12s +{indent}transform ./drop/b.csv 4.56s", + }, + TransformTallyTestCase { + description: "dicts for a large file that took a while get one hint, naming the slowest", + transforms: &[ + ("./drop/a.csv", 6_000, 150_000, true), + ("./drop/b.csv", 8_000, 100_000, true), + ("./drop/c.csv", 9_000, 99_999, true), + ], + expected_lines: "loaded RAW.M 7 rows 9.0s +{indent}transform ./drop/a.csv 6.00s +{indent}transform ./drop/b.csv 8.00s +{indent}transform ./drop/c.csv 9.00s +{indent}hint: the transform returned dicts for ./drop/b.csv (100,000 rows, 8.0s); a pyarrow Table or a DataFrame transforms faster", + }, + TransformTallyTestCase { + description: "a large slow file that returned Arrow gets no hint, nor does a quick one that returned dicts", + transforms: &[ + ("./drop/a.csv", 6_000, 150_000, false), + ("./drop/b.csv", 4_999, 150_000, true), + ], + expected_lines: "loaded RAW.M 7 rows 9.0s +{indent}transform ./drop/a.csv 6.00s +{indent}transform ./drop/b.csv 5.00s", + }, + ]; + for test_case in test_cases { + let event = LoadEvent { + schema_version: EVENTS_SCHEMA_VERSION, + run_id: "run".to_owned(), + at: "2026-11-14T00:00:00.000Z".to_owned(), + table: Some("RAW.M".to_owned()), + kind: EventKind::TableDone { + rows: 7, + elapsed_ms: 9_000, + position: None, + changes: Vec::new(), + transforms: test_case + .transforms + .iter() + .map(|(file, milliseconds, rows, dicts)| TransformTime { + file: (*file).to_owned(), + milliseconds: *milliseconds, + rows: *rows, + dicts: *dicts, + }) + .collect(), + }, + }; + assert_eq!( + event_lines(&event, (7, 5)).join("\n"), + test_case.expected_lines.replace("{indent}", &indent), + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlload/src/command/tests/test_loader_commands.rs b/crates/sqlload/src/command/tests/test_loader_commands.rs index d77e4f673..33dd601c6 100644 --- a/crates/sqlload/src/command/tests/test_loader_commands.rs +++ b/crates/sqlload/src/command/tests/test_loader_commands.rs @@ -19,6 +19,11 @@ fn given_the_standalone_executable_when_a_project_has_python_paths_then_it_says_ project: "python_paths = [\"sources\"]\n[connections.wh]\nkind = \"duckdb\"\ndatabase = \":memory:\"\n[targets.dev]\nconnection = \"wh\"\nloader_schema = \"RAW\"\n", expected_error: "install it (for example `uv tool install sqlload`)", }, + LoaderCommandTestCase { + description: "a transform needs the Python package, named before python_paths are", + project: "python_paths = [\"transforms\"]\n[connections.wh]\nkind = \"duckdb\"\ndatabase = \":memory:\"\n[targets.dev]\nconnection = \"wh\"\nloader_schema = \"RAW\"\n[connections.lake]\nkind = \"filesystem\"\nroot = \"/data\"\n[sources.metrics]\nkind = \"files\"\nconnection = \"lake\"\npath = \"drop/\"\nformat = \"csv\"\nwrite_strategy = \"append\"\n[sources.metrics.tables.metrics]\ntransform = \"metrics\"\n", + expected_error: "table metrics.metrics sets transform, which runs only in the sqlload Python package's command", + }, LoaderCommandTestCase { description: "a project without python_paths runs", project: TARGET, diff --git a/crates/sqlload/src/command/tests/test_types.rs b/crates/sqlload/src/command/tests/test_types.rs index 418ca2ac5..9bd2a7030 100644 --- a/crates/sqlload/src/command/tests/test_types.rs +++ b/crates/sqlload/src/command/tests/test_types.rs @@ -87,3 +87,10 @@ pub(crate) struct WrapTestCase { pub(crate) width: usize, pub(crate) expected_lines: &'static str, } + +pub(crate) struct TransformTallyTestCase { + pub(crate) description: &'static str, + /// Each file the table's transform ran on: its name, milliseconds, rows given and whether it returned dicts. + pub(crate) transforms: &'static [(&'static str, u64, u64, bool)], + pub(crate) expected_lines: &'static str, +} diff --git a/crates/sqlload/src/file_store/models.rs b/crates/sqlload/src/file_store/models.rs index b23cc1c51..5732e53da 100644 --- a/crates/sqlload/src/file_store/models.rs +++ b/crates/sqlload/src/file_store/models.rs @@ -34,6 +34,14 @@ impl StoredFile { None => format!("{}-{}", self.size, self.modified_ms), } } + + /// The file's store path, and its ZIP member: unique within a load, where two folders (such as the source's and `after_load`'s `move_to`) can hold files with the same relative key. + pub fn store_key(&self) -> String { + match &self.member { + Some(member) => format!("{}/{member}", self.path), + None => self.path.to_string(), + } + } } /// The files on an SFTP server, read over one SFTP session opened on first use: `root` is `/`, or empty for the user's login folder. Reads only: SQLLoad never writes to, renames or deletes from an SFTP server. diff --git a/crates/sqlload/src/files_source/_helpers/plan/folder.rs b/crates/sqlload/src/files_source/_helpers/plan/folder.rs index 8dba56872..ced66a505 100644 --- a/crates/sqlload/src/files_source/_helpers/plan/folder.rs +++ b/crates/sqlload/src/files_source/_helpers/plan/folder.rs @@ -1,3 +1,20 @@ +use crate::column_types::models::ColumnType; +use crate::file_store::main::list_files::list_files; +use crate::file_store::models::OpenFolder; +use crate::files_source::_helpers::archive::zip_archive::archive_members; +use crate::files_source::_helpers::plan::columns::{ + ColumnMerge, SampledColumn, case_problems, drop_loader_columns, file_columns, merge_columns, +}; +use crate::files_source::_helpers::plan::encoding::{file_codec, file_names}; +use crate::files_source::_helpers::plan::partitions::data_file; +use crate::files_source::_helpers::plan::sample::newest_columns; +use crate::files_source::constants::{ + SOURCE_FILE_COLUMN, SOURCE_ROW_COLUMN, SOURCE_VERSION_COLUMN, +}; +use crate::files_source::errors::FilesError; +use crate::files_source::models::{ + Compression, FilesEncoding, FilesTable, ListedFiles, PathValues, SourceFormat, StoredRead, +}; use crate::store_paths::models::{Location, StoreKind}; /// Why a folder on a local disk cannot be read, such as a folder below the root not existing (a mistyped path, or a missing mount point; a share not mounted on an existing one is an empty folder, and a missing root is created when the store opens); `None` when it can, and for other stores, whose folders exist only while they hold files. @@ -7,3 +24,145 @@ pub(crate) fn unreadable(location: &Location) -> Option { }; std::fs::read_dir(std::path::Path::new(root).join(&location.prefix)).err() } + +/// The `listed` files (read through `folder`'s store) with their path `values`, and the schema of the newest ones without the loader's own columns, which this load fills; when a file's path lacks a value, only the listing and its problems, before any file is read. +pub(crate) async fn listed_table( + folder: &OpenFolder, + listed: ListedFiles, + values: &PathValues, + encoding: &FilesEncoding, +) -> Result { + let (files, mut problems) = (listed.files, listed.problems); + if !problems.is_empty() { + return Ok(FilesTable { + files: files.into_iter().map(|(data, _)| data).collect(), + columns: Vec::new(), + problems, + skipped: listed.skipped, + unmatched: listed.unmatched, + dropped: Vec::new(), + objects: Vec::new(), + }); + } + let replaced: Vec<&str> = values + .names() + .into_iter() + .chain([SOURCE_FILE_COLUMN, SOURCE_VERSION_COLUMN, SOURCE_ROW_COLUMN]) + .collect(); + if let Some(record) = &encoding.record { + return Ok(FilesTable { + files: files.into_iter().map(|(data, _)| data).collect(), + columns: file_columns(vec![record_column(record)]), + problems, + skipped: listed.skipped, + unmatched: listed.unmatched, + dropped: Vec::new(), + objects: Vec::new(), + }); + } + let text_format = encoding.format != SourceFormat::Parquet; + let mut columns: Vec = Vec::new(); + let mut dropped: Vec = Vec::new(); + let (sampled, objects) = newest_columns(folder, &files, encoding).await?; + for ((data, _), found) in files.iter().rev().zip(sampled) { + let (found, loader) = drop_loader_columns(found); + for name in loader { + if !dropped.contains(&name) { + dropped.push(name); + } + } + let merge = ColumnMerge { + key: &data.file.relative_key, + replaced: &replaced, + widen_numbers: text_format, + }; + let (merged, conflicts) = merge_columns(columns, found, merge); + columns = merged; + problems.extend(conflicts); + } + if text_format { + problems.extend(case_problems(&columns)); + } + Ok(FilesTable { + files: files.into_iter().map(|(data, _)| data).collect(), + columns: file_columns(columns), + problems, + skipped: listed.skipped, + unmatched: listed.unmatched, + dropped, + objects, + }) +} + +/// The one `json` column a `shape = "record"` source loads each record into. +fn record_column(name: &str) -> SampledColumn { + SampledColumn { + name: name.to_owned(), + column_type: Some(Ok(ColumnType::Json)), + direct_refusal: None, + stored_read: StoredRead::Plain, + lists: Some(false), + } +} + +/// The listed files `encoding` reads and `values` keeps, sorted by key, with their decompression; how many others were skipped; and a problem for each file whose path lacks a value. +pub(crate) async fn listed_files( + folder: &OpenFolder, + values: &PathValues, + encoding: &FilesEncoding, +) -> Result { + let mut listed = ListedFiles { + files: Vec::new(), + skipped: 0, + unmatched: 0, + problems: Vec::new(), + }; + for file in list_files(folder).await? { + let Some(codec) = file_codec(&file.relative_key, encoding) else { + listed.skipped += 1; + continue; + }; + let files = match encoding.compression { + Compression::Zip => { + archive_members(folder, &file, encoding) + .await + .map_err(|reason| FilesError::File { + key: file.relative_key.clone(), + reason, + })? + } + _ => vec![file], + }; + for file in files { + let (data, problem) = data_file(file, values); + listed.problems.extend(problem); + match data { + Some(data) => listed.files.push((data, codec)), + None => listed.unmatched += 1, + } + } + } + let unmatched = match listed.unmatched { + 0 => String::new(), + count => format!("; {count} listed files do not match path_columns"), + }; + let reason = format!( + "no {} files found{unmatched}; check the source's path and format (an empty source never empties the destination)", + file_names(encoding) + ); + let key = folder.location.to_string(); + let others = listed.skipped + listed.unmatched; + match ( + listed.files.is_empty(), + others, + unreadable(&folder.location), + ) { + (false, _, _) => Ok(listed), + (true, 0, None) => Err(FilesError::NoFiles { key, reason }), + (true, _, None) => Err(FilesError::File { key, reason }), + (true, _, Some(error)) => Err(FilesError::File { + key, + reason: format!("{reason}; the folder cannot be read: {error}"), + }), + } +} diff --git a/crates/sqlload/src/files_source/main/batch_columns.rs b/crates/sqlload/src/files_source/main/batch_columns.rs index 5a483370a..6831bce34 100644 --- a/crates/sqlload/src/files_source/main/batch_columns.rs +++ b/crates/sqlload/src/files_source/main/batch_columns.rs @@ -24,7 +24,7 @@ pub fn batch_columns( .collect(); let mut problems: Vec = Vec::new(); for data in files { - let Some(footer) = footers.parquet.get(&data.file.relative_key) else { + let Some(footer) = footers.parquet.get(&data.file.store_key()) else { continue; }; let (mut found, _) = drop_loader_columns(parquet_columns(footer)); diff --git a/crates/sqlload/src/files_source/main/batch_footers.rs b/crates/sqlload/src/files_source/main/batch_footers.rs index 12ade7d71..db982fbe9 100644 --- a/crates/sqlload/src/files_source/main/batch_footers.rs +++ b/crates/sqlload/src/files_source/main/batch_footers.rs @@ -11,7 +11,7 @@ use crate::files_source::constants::CONCURRENT_FOOTER_READS; use crate::files_source::errors::FilesError; use crate::files_source::models::{BatchFooters, DataFile}; -/// Each of `files`' footers, by relative key, read once and a few at a time; reading the files reuses them, so nothing reads a footer twice. +/// Each of `files`' footers, by store key (two files can share a relative key, as under `--reload` with `move_to`), read once and a few at a time; reading the files reuses them, so nothing reads a footer twice. pub async fn batch_footers( folder: &OpenFolder, files: &[DataFile], @@ -35,7 +35,7 @@ async fn read_footer( let mut reader = object_reader(folder, &data.file); ArrowReaderMetadata::load_async(&mut reader, ArrowReaderOptions::new()) .await - .map(|footer| (key.clone(), footer)) + .map(|footer| (data.file.store_key(), footer)) .map_err(|error| FilesError::File { key: key.clone(), reason: format!("cannot read the Parquet footer: {error}"), diff --git a/crates/sqlload/src/files_source/main/batch_headers.rs b/crates/sqlload/src/files_source/main/batch_headers.rs index 8b2bb2934..80a462e39 100644 --- a/crates/sqlload/src/files_source/main/batch_headers.rs +++ b/crates/sqlload/src/files_source/main/batch_headers.rs @@ -12,7 +12,7 @@ use crate::files_source::constants::CONCURRENT_FOOTER_READS; use crate::files_source::errors::FilesError; use crate::files_source::models::{DataFile, FilesEncoding}; -/// Each of `files`' CSV column names (the header's, or `column_1`, ... without one or in `shape = "record"`), by relative key, read from the start of each file a few at a time; empty for a file with no records. +/// Each of `files`' CSV column names (the header's, or `column_1`, ... without one or in `shape = "record"`), by store key, read from the start of each file a few at a time; empty for a file with no records. pub async fn batch_headers( folder: &OpenFolder, files: &[DataFile], @@ -44,5 +44,5 @@ async fn read_header( let reader = CsvReader::open_as(bytes, (&encoding.csv, 1), encoding.record.is_none()) .await .map_err(failed)?; - Ok((key.clone(), reader.names().to_vec())) + Ok((data.file.store_key(), reader.names().to_vec())) } diff --git a/crates/sqlload/src/files_source/main/batch_refusals.rs b/crates/sqlload/src/files_source/main/batch_refusals.rs index 3e1e222bf..9c57412d5 100644 --- a/crates/sqlload/src/files_source/main/batch_refusals.rs +++ b/crates/sqlload/src/files_source/main/batch_refusals.rs @@ -11,7 +11,7 @@ pub fn batch_refusals( ) -> Vec<(String, String)> { let mut refused: Vec<(String, String)> = Vec::new(); for data in files { - let Some(footer) = footers.parquet.get(&data.file.relative_key) else { + let Some(footer) = footers.parquet.get(&data.file.store_key()) else { continue; }; for column in parquet_columns(footer) { diff --git a/crates/sqlload/src/files_source/main/discover_files.rs b/crates/sqlload/src/files_source/main/discover_files.rs index 61cc1901d..74184936e 100644 --- a/crates/sqlload/src/files_source/main/discover_files.rs +++ b/crates/sqlload/src/files_source/main/discover_files.rs @@ -1,23 +1,9 @@ //! List a files source and work out its schema from Parquet footers or sampled JSON and CSV records. -use crate::column_types::models::ColumnType; -use crate::file_store::main::list_files::list_files; use crate::file_store::models::OpenFolder; -use crate::files_source::_helpers::archive::zip_archive::archive_members; -use crate::files_source::_helpers::plan::columns::{ - ColumnMerge, SampledColumn, case_problems, drop_loader_columns, file_columns, merge_columns, -}; -use crate::files_source::_helpers::plan::encoding::{file_codec, file_names}; -use crate::files_source::_helpers::plan::folder::unreadable; -use crate::files_source::_helpers::plan::partitions::data_file; -use crate::files_source::_helpers::plan::sample::newest_columns; -use crate::files_source::constants::{ - SOURCE_FILE_COLUMN, SOURCE_ROW_COLUMN, SOURCE_VERSION_COLUMN, -}; +use crate::files_source::_helpers::plan::folder::{listed_files, listed_table}; use crate::files_source::errors::FilesError; -use crate::files_source::models::{ - Compression, FilesEncoding, FilesTable, ListedFiles, PathValues, SourceFormat, StoredRead, -}; +use crate::files_source::models::{FilesEncoding, FilesTable, PathValues}; /// The `encoding` files under `folder` with their path `values`, and the schema of the newest ones without the loader's own columns, which this load fills; when a file's path lacks a value, only the listing and its problems, before any file is read. pub async fn discover_files( @@ -26,137 +12,5 @@ pub async fn discover_files( encoding: &FilesEncoding, ) -> Result { let listed = listed_files(folder, values, encoding).await?; - let (files, mut problems) = (listed.files, listed.problems); - if !problems.is_empty() { - return Ok(FilesTable { - files: files.into_iter().map(|(data, _)| data).collect(), - columns: Vec::new(), - problems, - skipped: listed.skipped, - unmatched: listed.unmatched, - dropped: Vec::new(), - objects: Vec::new(), - }); - } - let replaced: Vec<&str> = values - .names() - .into_iter() - .chain([SOURCE_FILE_COLUMN, SOURCE_VERSION_COLUMN, SOURCE_ROW_COLUMN]) - .collect(); - if let Some(record) = &encoding.record { - return Ok(FilesTable { - files: files.into_iter().map(|(data, _)| data).collect(), - columns: file_columns(vec![record_column(record)]), - problems, - skipped: listed.skipped, - unmatched: listed.unmatched, - dropped: Vec::new(), - objects: Vec::new(), - }); - } - let text_format = encoding.format != SourceFormat::Parquet; - let mut columns: Vec = Vec::new(); - let mut dropped: Vec = Vec::new(); - let (sampled, objects) = newest_columns(folder, &files, encoding).await?; - for ((data, _), found) in files.iter().rev().zip(sampled) { - let (found, loader) = drop_loader_columns(found); - for name in loader { - if !dropped.contains(&name) { - dropped.push(name); - } - } - let merge = ColumnMerge { - key: &data.file.relative_key, - replaced: &replaced, - widen_numbers: text_format, - }; - let (merged, conflicts) = merge_columns(columns, found, merge); - columns = merged; - problems.extend(conflicts); - } - if text_format { - problems.extend(case_problems(&columns)); - } - Ok(FilesTable { - files: files.into_iter().map(|(data, _)| data).collect(), - columns: file_columns(columns), - problems, - skipped: listed.skipped, - unmatched: listed.unmatched, - dropped, - objects, - }) -} - -/// The one `json` column a `shape = "record"` source loads each record into. -fn record_column(name: &str) -> SampledColumn { - SampledColumn { - name: name.to_owned(), - column_type: Some(Ok(ColumnType::Json)), - direct_refusal: None, - stored_read: StoredRead::Plain, - lists: Some(false), - } -} - -/// The listed files `encoding` reads and `values` keeps, sorted by key, with their decompression; how many others were skipped; and a problem for each file whose path lacks a value. -async fn listed_files( - folder: &OpenFolder, - values: &PathValues, - encoding: &FilesEncoding, -) -> Result { - let mut listed = ListedFiles { - files: Vec::new(), - skipped: 0, - unmatched: 0, - problems: Vec::new(), - }; - for file in list_files(folder).await? { - let Some(codec) = file_codec(&file.relative_key, encoding) else { - listed.skipped += 1; - continue; - }; - let files = match encoding.compression { - Compression::Zip => { - archive_members(folder, &file, encoding) - .await - .map_err(|reason| FilesError::File { - key: file.relative_key.clone(), - reason, - })? - } - _ => vec![file], - }; - for file in files { - let (data, problem) = data_file(file, values); - listed.problems.extend(problem); - match data { - Some(data) => listed.files.push((data, codec)), - None => listed.unmatched += 1, - } - } - } - let unmatched = match listed.unmatched { - 0 => String::new(), - count => format!("; {count} listed files do not match path_columns"), - }; - let reason = format!( - "no {} files found{unmatched}; check the source's path and format (an empty source never empties the destination)", - file_names(encoding) - ); - let key = folder.location.to_string(); - let others = listed.skipped + listed.unmatched; - match ( - listed.files.is_empty(), - others, - unreadable(&folder.location), - ) { - (false, _, _) => Ok(listed), - (true, 0, None) => Err(FilesError::NoFiles { key, reason }), - (true, _, None) => Err(FilesError::File { key, reason }), - (true, _, Some(error)) => Err(FilesError::File { - key, - reason: format!("{reason}; the folder cannot be read: {error}"), - }), - } + listed_table(folder, listed, values, encoding).await } diff --git a/crates/sqlload/src/files_source/main/discover_reload_files.rs b/crates/sqlload/src/files_source/main/discover_reload_files.rs new file mode 100644 index 000000000..751652df4 --- /dev/null +++ b/crates/sqlload/src/files_source/main/discover_reload_files.rs @@ -0,0 +1,37 @@ +//! List a files source's folder and the folder `after_load` moved its loaded files to, for a `--reload`. + +use crate::file_store::models::OpenFolder; +use crate::files_source::_helpers::plan::folder::{listed_files, listed_table, unreadable}; +use crate::files_source::errors::FilesError; +use crate::files_source::models::{FilesEncoding, FilesTable, PathValues}; + +/// The `encoding` files under `folder` and under `archive` (where `after_load` moved the files it loaded, keeping their keys), as [`discover_files`](crate::files_source::main::discover_files::discover_files) lists one folder: sorted by key, with the schema of the newest ones. An archive that does not exist yet, or holds no file, adds none; no file in either is `FilesError::NoFiles`. +pub async fn discover_reload_files( + (folder, archive): (&OpenFolder, &OpenFolder), + values: &PathValues, + encoding: &FilesEncoding, +) -> Result { + let archived = match unreadable(&archive.location) { + Some(_) => Err(FilesError::NoFiles { + key: archive.location.to_string(), + reason: "the folder does not exist".to_owned(), + }), + None => listed_files(archive, values, encoding).await, + }; + let mut listed = match (listed_files(folder, values, encoding).await, archived) { + (Ok(mut listed), Ok(archived)) => { + listed.files.extend(archived.files); + listed.skipped += archived.skipped; + listed.unmatched += archived.unmatched; + listed.problems.extend(archived.problems); + listed + } + (Ok(listed), Err(FilesError::NoFiles { .. })) + | (Err(FilesError::NoFiles { .. }), Ok(listed)) => listed, + (Err(error), _) | (_, Err(error)) => return Err(error), + }; + listed + .files + .sort_by(|(left, _), (right, _)| left.file.relative_key.cmp(&right.file.relative_key)); + listed_table(folder, listed, values, encoding).await +} diff --git a/crates/sqlload/src/files_source/main/mod.rs b/crates/sqlload/src/files_source/main/mod.rs index b21abedfe..6f0563bb8 100644 --- a/crates/sqlload/src/files_source/main/mod.rs +++ b/crates/sqlload/src/files_source/main/mod.rs @@ -5,6 +5,7 @@ pub mod batch_headers; pub mod batch_refusals; pub mod constant_columns; pub mod discover_files; +pub mod discover_reload_files; pub mod footer_columns; pub mod key_columns; pub mod merge_order; @@ -14,5 +15,6 @@ pub mod path_template; pub mod read_files; pub mod read_footers; pub mod select_files; +pub mod shape_rows; pub mod stored_read_of; pub mod text_charset; diff --git a/crates/sqlload/src/files_source/main/read_files.rs b/crates/sqlload/src/files_source/main/read_files.rs index 1ec1be4cd..df8c0f15f 100644 --- a/crates/sqlload/src/files_source/main/read_files.rs +++ b/crates/sqlload/src/files_source/main/read_files.rs @@ -52,7 +52,7 @@ async fn read_parquet( on_batch: &mut impl AsyncFnMut(RecordBatch) -> Result<(), String>, ) -> Result { let reader = object_reader(file.folder, &file.data.file); - let builder = match footers.parquet.get(&file.data.file.relative_key) { + let builder = match footers.parquet.get(&file.data.file.store_key()) { Some(footer) => ParquetRecordBatchStreamBuilder::new_with_metadata(reader, footer.clone()), None => ParquetRecordBatchStreamBuilder::new(reader) .await diff --git a/crates/sqlload/src/files_source/main/shape_rows.rs b/crates/sqlload/src/files_source/main/shape_rows.rs new file mode 100644 index 000000000..19ed82344 --- /dev/null +++ b/crates/sqlload/src/files_source/main/shape_rows.rs @@ -0,0 +1,16 @@ +//! Shape rows made from a file (by a `transform`) as the file's own rows are. + +use arrow_array::RecordBatch; + +use crate::files_source::_helpers::plan::shape::{output_schema, shape_batch}; +use crate::files_source::models::{DataFile, FilesRead}; + +/// `batch`, made from `file`'s rows, in `read`'s columns: each cast to its type, missing ones NULL, partition and source-file columns filled in, and rows numbered after `rows_before`; a value that does not convert names its column. +pub fn shape_rows( + batch: &RecordBatch, + (file, file_name, rows_before): (&DataFile, &str, u64), + read: &FilesRead, +) -> Result { + let schema = output_schema(&read.columns, &read.nanos)?; + shape_batch(batch, file, (file_name, rows_before), &schema).map_err(|error| error.reason) +} diff --git a/crates/sqlload/src/files_source/models.rs b/crates/sqlload/src/files_source/models.rs index e56425b90..928b247f2 100644 --- a/crates/sqlload/src/files_source/models.rs +++ b/crates/sqlload/src/files_source/models.rs @@ -126,11 +126,11 @@ pub struct FilesRead { pub batch_rows: usize, } -/// One batch's Parquet footers by relative key, read once; a file missing here has its footer read when it is. +/// One batch's Parquet footers by store key (`StoredFile::store_key`), read once; a file missing here has its footer read when it is. #[derive(Debug, Clone, Default)] pub struct BatchFooters { pub parquet: std::collections::BTreeMap, - /// CSV files' column names by relative key, read once where the destination reads the files itself (by position). + /// CSV files' column names by store key, read once where the destination reads the files itself (by position). pub headers: std::collections::BTreeMap>, } diff --git a/crates/sqlload/src/files_source/tests/helpers.rs b/crates/sqlload/src/files_source/tests/helpers.rs index a5df59efb..b6337f93e 100644 --- a/crates/sqlload/src/files_source/tests/helpers.rs +++ b/crates/sqlload/src/files_source/tests/helpers.rs @@ -469,6 +469,56 @@ async fn break_footer_and_read(folder: &OpenFolder, reuse: bool) -> Result Result { + let (folder, root) = folder_with(&[]).await?; + let result = write_and_read_same_key(&folder).await; + let _ = std::fs::remove_dir_all(&root); + result +} + +async fn write_and_read_same_key(folder: &OpenFolder) -> Result { + for (key, ids) in [("x.parquet", vec![1, 2, 3]), ("done/x.parquet", vec![4])] { + let values: arrow_array::ArrayRef = Arc::new(arrow_array::Int64Array::from(ids)); + let batch = + RecordBatch::try_from_iter([("id", values)]).map_err(|error| error.to_string())?; + let mut body: Vec = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut body, batch.schema(), None) + .map_err(|error| error.to_string())?; + writer.write(&batch).map_err(|error| error.to_string())?; + writer.close().map_err(|error| error.to_string())?; + write_file(folder, key, Bytes::from(body)) + .await + .map_err(|error| error.to_string())?; + } + let encoding = encoding(SourceFormat::Parquet); + let mut files = discover_files(folder, &PathValues::default(), &encoding) + .await + .map_err(|error| error.to_string())? + .files; + for data in &mut files { + data.file.relative_key = "x.parquet".to_owned(); + } + let footers = crate::files_source::main::batch_footers::batch_footers(folder, &files) + .await + .map_err(|error| error.to_string())?; + let request = FilesRead { + nanos: Vec::new(), + files, + encoding, + checks: TextChecks::default(), + columns: vec![DesiredColumn { + name: "id".to_owned(), + column_type: ColumnType::Int, + origin: ColumnOrigin::Source, + }], + batch_rows: 10, + }; + Ok(read_files(folder, &request, &footers, async |_| Ok(())) + .await + .map_or_else(|_| "failed".to_owned(), |rows| rows.to_string())) +} + /// One Parquet file of three rows with nested columns, discovered and read: column `name`'s planned type, and its values `|`-separated (`NULL` for NULL). pub(crate) async fn nested_parquet_column(name: &str) -> Result<(String, String), String> { let (folder, root) = folder_with(&[]).await?; diff --git a/crates/sqlload/src/files_source/tests/test_batch_footers.rs b/crates/sqlload/src/files_source/tests/test_batch_footers.rs index 74bafdac0..05b820007 100644 --- a/crates/sqlload/src/files_source/tests/test_batch_footers.rs +++ b/crates/sqlload/src/files_source/tests/test_batch_footers.rs @@ -1,4 +1,4 @@ -use crate::files_source::tests::helpers::{block_on, read_with_broken_footer}; +use crate::files_source::tests::helpers::{block_on, read_same_key_files, read_with_broken_footer}; use crate::files_source::tests::test_types::BatchFootersTestCase; #[test] @@ -8,17 +8,28 @@ fn given_a_batchs_footers_when_reading_its_files_then_no_footer_is_read_again() BatchFootersTestCase { description: "the footer read earlier is reused, so the broken one is never read", reuse: true, + same_key: false, expected_outcome: "3", }, BatchFootersTestCase { description: "without it the read must fetch the footer, and fails", reuse: false, + same_key: false, expected_outcome: "failed", }, + BatchFootersTestCase { + description: "two files with one relative key in two folders each read with their own footer", + reuse: true, + same_key: true, + expected_outcome: "4", + }, ]; for test_case in test_cases { assert_eq!( - block_on(read_with_broken_footer(test_case.reuse))?, + match test_case.same_key { + true => block_on(read_same_key_files())?, + false => block_on(read_with_broken_footer(test_case.reuse))?, + }, test_case.expected_outcome, "{}", test_case.description diff --git a/crates/sqlload/src/files_source/tests/test_types.rs b/crates/sqlload/src/files_source/tests/test_types.rs index 589a85bf6..b49750875 100644 --- a/crates/sqlload/src/files_source/tests/test_types.rs +++ b/crates/sqlload/src/files_source/tests/test_types.rs @@ -124,6 +124,8 @@ pub(crate) struct BatchFootersTestCase { pub(crate) description: &'static str, /// Read the footer before the file's own footer is broken, and reuse it. pub(crate) reuse: bool, + /// Instead, read a batch of two files with one relative key from two folders (as `--reload` with `move_to` does), of three rows and one, with their footers read first. + pub(crate) same_key: bool, /// The rows read, or `failed`. pub(crate) expected_outcome: &'static str, } diff --git a/crates/sqlload/src/http_source/_helpers/pages.rs b/crates/sqlload/src/http_source/_helpers/pages.rs index b19a159b5..0e903e284 100644 --- a/crates/sqlload/src/http_source/_helpers/pages.rs +++ b/crates/sqlload/src/http_source/_helpers/pages.rs @@ -148,6 +148,8 @@ pub(crate) async fn paginate( path, source: shown, projection: projection.clone(), + number: step, + key: context.key.clone(), }; if has_rows && pages.send(page).await.is_err() { return Ok(()); diff --git a/crates/sqlload/src/http_source/models.rs b/crates/sqlload/src/http_source/models.rs index 2231c2858..692664a1a 100644 --- a/crates/sqlload/src/http_source/models.rs +++ b/crates/sqlload/src/http_source/models.rs @@ -18,6 +18,16 @@ pub enum HttpJob { FanOut(Box), } +impl HttpJob { + /// The `transform=` the rows of the job's pages go through. + pub fn transform(&self) -> Option<&PageTransform> { + match self { + Self::Paginate(request) => request.transform.as_ref(), + Self::FanOut(fan_out) => fan_out.request.transform.as_ref(), + } + } +} + /// One endpoint and how to page through it. #[derive(Debug, Clone, Deserialize)] pub struct Request { @@ -57,6 +67,16 @@ pub struct Request { /// In a fan-out, where each record's key is, checked against the keys its request sent. #[serde(default)] pub key_path: Option, + /// The function each page's rows go through before they load, which the Python crossing calls. + #[serde(default)] + pub transform: Option, +} + +/// `transform=` on a job: the key the Python side calls the function by, and the name errors show. +#[derive(Debug, Clone, Deserialize)] +pub struct PageTransform { + pub function: String, + pub name: String, } /// `fan_out(key_path=...)`: a path to each record's key, from the record or an object around it (`sl.parent`). @@ -352,6 +372,10 @@ pub struct HttpPage { pub source: String, /// The columns each row loads as, when the job names its fields. pub projection: Option, + /// The page's number within its request, from 1. + pub number: u64, + /// The fields of the fan-out key the page was requested for; empty in a paginate or a batched fan-out. + pub key: Map, } /// A parsed `select`: the fields down to the rows and, for keys, the field of each row. diff --git a/crates/sqlload/src/load_events/models.rs b/crates/sqlload/src/load_events/models.rs index 80df8bb51..286aa391a 100644 --- a/crates/sqlload/src/load_events/models.rs +++ b/crates/sqlload/src/load_events/models.rs @@ -43,6 +43,9 @@ pub enum EventKind { /// The `created` and `altered` lines of a table that made or widened its columns as it loaded (a Python table, or a nested one's child tables). #[serde(default, skip_serializing_if = "Vec::is_empty")] changes: Vec, + /// How long a files or Sheets table's `transform` took on each file or range. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + transforms: Vec, }, /// A table failed, and whether running the load again may get past it. TableFailed { diff --git a/crates/sqlload/src/loader/_helpers/destination.rs b/crates/sqlload/src/loader/_helpers/destination.rs index 7a88f3c22..55b67880c 100644 --- a/crates/sqlload/src/loader/_helpers/destination.rs +++ b/crates/sqlload/src/loader/_helpers/destination.rs @@ -2,19 +2,27 @@ use std::collections::{BTreeMap, BTreeSet}; use crate::adapters::contract::models::WriteMode; use crate::adapters::dispatch::main::dialect::files_destination::files_destination; +use crate::adapters::dispatch::main::dialect::sql_adapter::sql_adapter; use crate::adapters::dispatch::main::session::clean_stage::clean_stage; use crate::adapters::dispatch::main::session::drop_stale_staging::drop_stale_staging; use crate::adapters::dispatch::models::Destination; use crate::adapters::files::constants::PARQUET_SUFFIX; +use crate::adapters::files::main::clear_pending::clear_pending; +use crate::adapters::files::main::mark_pending::mark_pending; use crate::adapters::files::main::newest_load_files::newest_load_files; use crate::adapters::files::main::table_path::table_path; use crate::adapters::files::models::{ FilesDestination, FolderClaim, TableFiles, TablePath, TablePathRequest, }; +use crate::adapters::files::models::{PartitionFolder, PendingPhase, PendingRun, ReplacedFiles}; use crate::column_types::models::ColumnType; +use crate::file_store::main::delete_files::delete_files; +use crate::file_store::main::write_file::write_file; use crate::file_store::models::{OpenFolder, StoredFile}; +use crate::files_source::constants::NANOSECOND_SCALE; use crate::files_source::main::footer_columns::footer_columns; use crate::files_source::main::read_footers::read_footers; +use crate::files_source::models::{FileColumn, StoredRead}; use crate::loader::_helpers::resolve::display_name; use crate::loader::_helpers::table_plan::blocked_plan; use crate::loader::errors::LoaderError; @@ -23,12 +31,14 @@ use crate::loader::models::{ }; use crate::project::constants::DEFAULT_DESTINATION_PATH; use crate::project::models::WriteStrategy; -use crate::schema_plan::models::LiveColumn; +use crate::schema_plan::models::{LiveColumn, TableAction}; +use crate::store_paths::models::Location; use crate::table_lease::constants::LEASE_REFRESH; use crate::table_lease::main::keep_lease::keep_lease; use crate::table_lease::main::lease_now::lease_now; use crate::table_lease::main::release_lease::release_lease; use crate::table_lease::models::Lease; +use bytes::Bytes; use tokio_util::sync::CancellationToken; /// Before a load: drop the staging tables crashed loads left in the schemas it writes to. A failure only warns. @@ -294,7 +304,7 @@ fn shared_problems(table: &TablePlan, others: &[&TableSpec]) -> Vec { removing.strategy )); } - if table.mode == WriteMode::Recreate { + if table.mode == WriteMode::Recreate || table.schema.action == TableAction::Recreate { problems.push(format!( "{place}; --reload would drop the other tables' rows" )); @@ -426,3 +436,121 @@ pub(crate) async fn give_back(folder: &OpenFolder, lease: &Lease) { tracing::warn!(%error, folder = %folder.location, "cannot release the table's lease; it expires on its own"); } } + +/// The columns of `columns` SQLLoad's reader keeps at nanoseconds for `spec`'s warehouse: those the files store at nanoseconds, unless pinned or `live` holds fewer digits. +pub(crate) fn nanosecond_columns( + spec: &TableSpec, + columns: &[FileColumn], + live: Option<&[LiveColumn]>, +) -> Vec { + match sql_adapter(spec.warehouse).keeps_nanoseconds() { + true => columns + .iter() + .filter(|column| { + matches!( + column.stored_read, + StoredRead::Nanos | StoredRead::Epoch(NANOSECOND_SCALE) + ) + }) + .filter(|column| { + let loaded = spec.source_table.loaded_as(&column.name); + !spec.pinned.contains_key(loaded) && fewer_digits(live, loaded).is_none() + }) + .map(|column| column.name.clone()) + .collect(), + false => Vec::new(), + } +} + +/// A warning for each column the files store at nanoseconds but the live table holds at fewer digits, so the load cuts them; a column whose type is pinned is the user's choice and is not warned about. +pub(crate) fn truncated_nanoseconds(spec: &TableSpec, live: Option<&[LiveColumn]>) -> Vec { + let Some(files) = &spec.files else { + return Vec::new(); + }; + files + .table + .columns + .iter() + .filter(|column| { + matches!( + column.stored_read, + StoredRead::Nanos | StoredRead::Epoch(NANOSECOND_SCALE) + ) && !spec + .pinned + .contains_key(spec.source_table.loaded_as(&column.name)) + }) + .filter_map(|column| truncated_warning(live, spec.source_table.loaded_as(&column.name))) + .collect() +} + +/// The warning for column `name` when the live column holds fewer than nanosecond digits. +fn truncated_warning(live: Option<&[LiveColumn]>, name: &str) -> Option { + let digits = fewer_digits(live, name)?; + Some(format!( + "{name}: nanoseconds in the files, but the column holds {digits} digits after the second, so they are truncated.\nRebuild the column at (9) (--reload) to keep them." + )) +} + +/// The digits after the second the live column `name` holds, when fewer than nanoseconds, so SQLLoad's reader cuts its values. +fn fewer_digits(live: Option<&[LiveColumn]>, name: &str) -> Option { + live.into_iter() + .flatten() + .find(|existing| existing.name.eq_ignore_ascii_case(name)) + .and_then(|existing| existing.fraction_digits) + .filter(|digits| *digits < i64::from(NANOSECOND_SCALE)) +} + +/// The replaced files a commit deletes: those of the partitions it wrote, and, at the last commit, every other file the load replaces (everything for `replace`, the window's partitions for `delete_insert`); never one already deleted. +pub(crate) fn commit_deletes( + replaced: &ReplacedFiles, + partition: Option<&PartitionFolder>, + (written, last): (&BTreeSet, bool), + deleted: &BTreeSet, +) -> Vec { + let folders: Vec = match partition { + Some(partition) => written + .iter() + .map(|value| format!("{}/", partition.name(value))) + .collect(), + None => Vec::new(), + }; + let mut keys: Vec = written + .iter() + .flat_map(|value| replaced.by_partition.get(value).into_iter().flatten()) + .cloned() + .collect(); + for key in &replaced.always { + let written_folder = folders.iter().any(|folder| key.starts_with(folder)); + if last || written_folder { + keys.push(key.clone()); + } + } + keys.retain(|key| !deleted.contains(key)); + keys.sort(); + keys.dedup(); + keys +} + +/// Once the load still holds the table's lease in `table_folder` (extending it), record the replaced files (and a logged table's run record) in the marker, which commits the files; write the run record, delete the replaced files, then the marker, so a stop in between is finished by the next plan or load. +pub(crate) async fn replace_files( + (folder, (table_folder, lease)): (&OpenFolder, (&OpenFolder, &Lease)), + file_id: &str, + (keys, run): (&[String], Option<&PendingRun>), +) -> Result<(), LoaderError> { + keep_lease(table_folder, lease, lease_now()) + .await? + .map_err(LoaderError::Invalid)?; + mark_pending(folder, file_id, PendingPhase::Replacing, (keys, run)).await?; + if let Some(run) = run { + let records = OpenFolder { + store: folder.store.clone(), + location: Location { + store: folder.location.store.clone(), + prefix: run.prefix.clone(), + }, + }; + write_file(&records, &run.name, Bytes::from(run.json.clone())).await?; + } + delete_files(folder, keys).await?; + Ok(clear_pending(folder, file_id).await?) +} diff --git a/crates/sqlload/src/loader/_helpers/publish.rs b/crates/sqlload/src/loader/_helpers/publish.rs index d63da5d52..a0dd67ab2 100644 --- a/crates/sqlload/src/loader/_helpers/publish.rs +++ b/crates/sqlload/src/loader/_helpers/publish.rs @@ -25,7 +25,6 @@ use crate::adapters::dispatch::main::write::load_unloaded::load_unloaded; use crate::adapters::dispatch::main::write::open_table_writer::open_table_writer; use crate::adapters::dispatch::models::Destination; use crate::adapters::dispatch::models::Writer; -use crate::adapters::files::main::clear_pending::clear_pending; use crate::adapters::files::main::close_writer::close_writer; use crate::adapters::files::main::mark_pending::mark_pending; use crate::adapters::files::main::open_writer::open_writer; @@ -34,7 +33,6 @@ use crate::adapters::files::main::swap_folder::swap_folder; use crate::adapters::files::main::write_loaded_files::write_loaded_files; use crate::adapters::files::main::write_rows::write_rows; use crate::adapters::files::models::FilesWriter; -use crate::adapters::files::models::PartitionFolder; use crate::adapters::files::models::PendingPhase; use crate::adapters::files::models::PendingRun; use crate::adapters::files::models::ReplacedFiles; @@ -48,7 +46,6 @@ use crate::column_types::main::arrow_field::arrow_field; use crate::column_types::models::ColumnType; use crate::file_store::main::delete_files::delete_files; use crate::file_store::main::open_folder::open_folder; -use crate::file_store::main::write_file::write_file; use crate::file_store::models::OpenFolder; use crate::file_store::models::OpenFolder as StoreFolder; use crate::files_source::constants::{ @@ -68,6 +65,7 @@ use crate::files_source::models::DataFile; use crate::files_source::models::{FilesRead, FilesTable, LoadedFile, SourceFormat, TextChecks}; use crate::load_naming::main::metadata_name::metadata_name; use crate::load_naming::models::{MetadataRole, Naming}; +use crate::loader::_helpers::destination::{commit_deletes, nanosecond_columns, replace_files}; use crate::loader::_helpers::destination::{give_back, keep_in_background, leased_folder}; use crate::loader::_helpers::python::list_choices; use crate::loader::_helpers::python::live_child_tables; @@ -75,21 +73,21 @@ use crate::loader::_helpers::python::unnested_table; use crate::loader::_helpers::python::without_deleted_items; use crate::loader::_helpers::resolve::resolve_tables; use crate::loader::_helpers::run_log::append_run; -use crate::loader::_helpers::run_log::ensure_merge_log; use crate::loader::_helpers::run_log::ensure_runs_table; use crate::loader::_helpers::run_log::fold_maximum; use crate::loader::_helpers::run_log::logs_files; -use crate::loader::_helpers::run_log::merged_record; +use crate::loader::_helpers::run_log::transformed; use crate::loader::_helpers::run_log::{ - in_file, listed_sources, merged_file_text, merged_name, with_file_keys, + commit_transform, in_file, listed_sources, merge_rebuild, merged_file_text, merged_name, + merged_record, swap_rebuild, with_file_keys, with_transform, }; use crate::loader::_helpers::sources::ReadPart; -use crate::loader::_helpers::sources::files_read_as; use crate::loader::_helpers::sources::files_source_parts; -use crate::loader::_helpers::sources::nanosecond_columns; use crate::loader::_helpers::sources::read_source; +use crate::loader::_helpers::sources::unarchived; use crate::loader::_helpers::sources::unload_source; use crate::loader::_helpers::specs::checked_key; +use crate::loader::_helpers::specs::files_read_as; use crate::loader::_helpers::specs::{null_cursor_problem, null_cursor_rows}; use crate::loader::_helpers::table_plan::{ cursor_binding, source_cursor, source_filter, with_null_cursors, @@ -106,6 +104,7 @@ use crate::loader::models::ChildTables; use crate::loader::models::CommitRange; use crate::loader::models::CommitRule; use crate::loader::models::FilesSink; +use crate::loader::models::FilesSource; use crate::loader::models::FolderSwap; use crate::loader::models::Publisher; use crate::loader::models::ReplacedKeys; @@ -113,6 +112,7 @@ use crate::loader::models::RunEntry; use crate::loader::models::SourceReader; use crate::loader::models::TableLoaded; use crate::loader::models::TablePlan; +use crate::loader::models::TransformRun; use crate::loader::models::{LoadRun, PlanOptions}; use crate::masking::main::mask_batch::mask_batch; use crate::postgres_source::errors::PostgresError; @@ -141,11 +141,8 @@ use crate::schema_plan::models::ColumnOrigin; use crate::schema_plan::models::DesiredColumn; use crate::schema_plan::models::LiveColumn; use crate::snowflake_source::models::UnloadedFiles; -use crate::store_paths::models::Location; -use crate::table_lease::main::keep_lease::keep_lease; use crate::table_lease::main::lease_now::lease_now; use crate::table_lease::main::take_lease::take_lease; -use crate::table_lease::models::Lease; use crate::time_text::main::now_utc_text::now_utc_text; use arrow_array::Array; use arrow_array::ArrayRef; @@ -159,7 +156,6 @@ use arrow_cast::display::{ArrayFormatter, FormatOptions}; use arrow_cast::{CastOptions, cast, cast_with_options}; use arrow_schema::Field; use arrow_schema::Schema; -use bytes::Bytes; use chrono::DateTime; use std::collections::BTreeSet; use std::sync::Arc; @@ -491,7 +487,7 @@ fn written_only(batch: RecordBatch, written: &[DesiredColumn]) -> Result { @@ -623,6 +616,10 @@ impl Publisher { changes: Vec::new(), }), }; + let rebuild = match records_files(&table) { + true => merge_rebuild(&table, &destination).await?, + false => None, + }; Ok(Self { children, writer, @@ -656,6 +653,8 @@ impl Publisher { run_keys: Default::default(), fed_files: Vec::new(), settled_files: Vec::new(), + rebuild, + transform: None, }) } @@ -860,11 +859,15 @@ impl Publisher { } _ => {} } + if let Some(rebuild) = &self.rebuild { + swap_rebuild(&self.table, &self.destination, rebuild).await?; + } let spec = &self.table.spec; let adapter = sql_adapter(spec.warehouse); + let swapped = self.at_end || self.rebuild.is_some(); let comment = adapter .render_table_comment(&spec.destination, &spec.description) - .filter(|_| self.at_end && self.files.is_none() && !adapter.keeps_tables()); + .filter(|_| swapped && self.files.is_none() && !adapter.keeps_tables()); if let Some(comment) = comment && let Err(error) = execute(&self.destination, vec![comment]).await { @@ -895,12 +898,14 @@ impl Publisher { table: format!("{}{SWAP_SUFFIX}", staging.table), ..staging.clone() }; + let build = self.rebuild.as_ref().map(|rebuild| rebuild.build.clone()); let drops: Vec = [staging.clone(), aside] - .iter() + .into_iter() + .chain(build) .map(|table| { format!( "DROP TABLE IF EXISTS {}", - sql_adapter(warehouse).render_qualified_name(table) + sql_adapter(warehouse).render_qualified_name(&table) ) }) .collect(); @@ -1096,27 +1101,34 @@ impl Publisher { let runs = runs_table(&spec.destination); let merging_file = self.merged_file.is_some(); let file = self.merged_file.clone(); - let merged = self.merged_file.take().map(|file| { - merged_record( - spec, - (&self.load_id, self.commits + 1), - (file, self.waiting_rows), + let ran = commit_transform(&self.table, self.transform.as_mut()); + let mut merged = self.merged_file.take().map(|file| { + let record = (file, self.waiting_rows); + with_transform( + &self.table, + merged_record(spec, (&self.load_id, self.commits + 1), record), ) }); - let run_insert = match (&self.run, merged) { + let held = self.rebuild.as_ref().and_then(|_| merged.take()); + let record = match (&self.run, merged) { (Some(_), _) => { let record = self.record(self.waiting_rows)?; - Some(insert_run(&runs, &record, warehouse).map_err(LoaderError::Invalid)?) + Some(with_transform(&self.table, record)) } - (None, Some(record)) => { - Some(insert_run(&runs, &record, warehouse).map_err(LoaderError::Invalid)?) - } - (None, None) => None, + (None, Some(record)) => Some(record), + (None, None) => ran.filter(|_| held.is_none()), }; + let run_insert = record + .map(|record| insert_run(&runs, &record, warehouse).map_err(LoaderError::Invalid)) + .transpose()?; let spec = &self.table.spec; + let target = self + .rebuild + .as_ref() + .map_or(&spec.destination, |rebuild| &rebuild.build); let statements = sql_adapter(warehouse).render_commit( &WriteStatementsRequest { - table: &spec.destination, + table: target, columns: &self.table.schema.write_columns, mode: &mode, hashed: &[], @@ -1129,7 +1141,7 @@ impl Publisher { destination: &self.destination, warehouse, staging: &staging, - table: &spec.destination, + table: target, exists: self.table.exists, columns: &self.table.schema.write_columns, committed: self.committed_rows, @@ -1195,6 +1207,9 @@ impl Publisher { false => None, }; execute(&self.destination, statements).await?; + if let Some(rebuild) = self.rebuild.as_mut() { + rebuild.records.extend(held); + } self.committed_rows += self.waiting_rows; self.waiting_rows = 0; self.commits += 1; @@ -1236,7 +1251,11 @@ impl Publisher { files: &[DataFile], ) -> Result { let spec = &self.table.spec; - let Some(source_files) = spec.files.as_ref() else { + let Some(source_files) = spec + .files + .as_ref() + .filter(|files| files.transform.is_none()) + else { return Ok(BatchFooters::default()); }; if source_files.encoding.record.is_some() { @@ -1326,7 +1345,10 @@ impl Publisher { ))); } }; - widened.files = Some(files_source); + widened.files = Some(FilesSource { + archive: source_files.archive.clone(), + ..files_source + }); widened.approved = spec.approved; let replanned_columns = replan_columns( &widened, @@ -1375,6 +1397,17 @@ impl Publisher { ); execute_ddl(&self.destination, changes).await?; } + if let Some(rebuild) = &self.rebuild { + let changes = sql_adapter(replanned.spec.warehouse).render_staging_changes( + &rebuild.build, + ( + &self.table.schema.write_columns, + &replanned.schema.write_columns, + ), + Some(&[]), + ); + execute_ddl(&self.destination, changes).await?; + } self.table = replanned; self.widen_files().await?; Ok(footers) @@ -1447,11 +1480,18 @@ pub(crate) async fn publish_table( { return Err(LoaderError::Invalid(null_cursor_problem(&table.spec, rows))); } - let publisher = Publisher::open(table.clone(), (destination, run.options), None).await?; - let published = feed(reader, table, publisher, run).await; + let mut publisher = Publisher::open(table.clone(), (destination, run.options), None).await?; + let transform = run.host.clone().map(TransformRun::new); + publisher.transform = transform.clone().map(|transform| (transform, 0)); + let published = feed(reader, table, publisher, (run, transform.as_ref())).await; let (mut publisher, result) = published; publisher.close().await; let rows = result?; + let recorded = publisher + .transform + .as_ref() + .map_or(0, |(_, recorded)| *recorded); + let transforms = transformed(table, (destination, transform, recorded)).await; let report = match reader { SourceReader::Files(folder) => { after_load(folder, &table.spec.after_load, &publisher.settled_files).await @@ -1464,6 +1504,7 @@ pub(crate) async fn publish_table( note: None, after_load: report, replaced_keys: publisher.replaced.clone(), + transforms, }) } @@ -1472,7 +1513,7 @@ async fn feed( reader: &SourceReader, table: &TablePlan, mut publisher: Publisher, - run: &LoadRun<'_>, + (run, transform): (&LoadRun<'_>, Option<&TransformRun>), ) -> (Publisher, Result) { let batch_bytes = run.options.batch_bytes; let fed = async { @@ -1510,9 +1551,7 @@ async fn feed( let merged = merging.then(|| merged_file_text(table, &files)); publisher.merged_file = merged.clone(); if table.spec.after_load != AfterLoad::Keep { - publisher - .fed_files - .extend(files.iter().map(|data| data.file.clone())); + publisher.fed_files.extend(unarchived(table, &files)); } let fed_group = async { let footers = match reader { @@ -1548,6 +1587,7 @@ async fn feed( let part = ReadPart { files: Some(files), footers, + transform: transform.cloned(), ..ReadPart::default() }; let current = publisher.table.clone(); @@ -1610,6 +1650,7 @@ async fn feed( CommitRule::Size => { let part = ReadPart { batch_bytes, + transform: transform.cloned(), ..ReadPart::default() }; match pulls_unloaded(table.route.as_ref()) { @@ -1629,6 +1670,7 @@ async fn feed( CommitRule::Whole(_) => { let part = ReadPart { batch_bytes, + transform: transform.cloned(), ..ReadPart::default() }; publisher.pass_part(reader, part).await?; @@ -1697,37 +1739,6 @@ fn window_filter( } } -/// The replaced files a commit deletes: those of the partitions it wrote, and, at the last commit, every other file the load replaces (everything for `replace`, the window's partitions for `delete_insert`); never one already deleted. -fn commit_deletes( - replaced: &ReplacedFiles, - partition: Option<&PartitionFolder>, - (written, last): (&BTreeSet, bool), - deleted: &BTreeSet, -) -> Vec { - let folders: Vec = match partition { - Some(partition) => written - .iter() - .map(|value| format!("{}/", partition.name(value))) - .collect(), - None => Vec::new(), - }; - let mut keys: Vec = written - .iter() - .flat_map(|value| replaced.by_partition.get(value).into_iter().flatten()) - .cloned() - .collect(); - for key in &replaced.always { - let written_folder = folders.iter().any(|folder| key.starts_with(folder)); - if last || written_folder { - keys.push(key.clone()); - } - } - keys.retain(|key| !deleted.contains(key)); - keys.sort(); - keys.dedup(); - keys -} - impl Publisher { /// Add a batch to the commit's files, starting them (and marking the commit as writing) on its first batch. pub(crate) async fn stage_files(&mut self, batch: RecordBatch) -> Result<(), LoaderError> { @@ -1971,27 +1982,3 @@ impl Publisher { }) } } - -/// Once the load still holds the table's lease in `table_folder` (extending it), record the replaced files (and a logged table's run record) in the marker, which commits the files; write the run record, delete the replaced files, then the marker, so a stop in between is finished by the next plan or load. -async fn replace_files( - (folder, (table_folder, lease)): (&OpenFolder, (&OpenFolder, &Lease)), - file_id: &str, - (keys, run): (&[String], Option<&PendingRun>), -) -> Result<(), LoaderError> { - keep_lease(table_folder, lease, lease_now()) - .await? - .map_err(LoaderError::Invalid)?; - mark_pending(folder, file_id, PendingPhase::Replacing, (keys, run)).await?; - if let Some(run) = run { - let records = OpenFolder { - store: folder.store.clone(), - location: Location { - store: folder.location.store.clone(), - prefix: run.prefix.clone(), - }, - }; - write_file(&records, &run.name, Bytes::from(run.json.clone())).await?; - } - delete_files(folder, keys).await?; - Ok(clear_pending(folder, file_id).await?) -} diff --git a/crates/sqlload/src/loader/_helpers/python.rs b/crates/sqlload/src/loader/_helpers/python.rs index ad3cb9928..d76058de8 100644 --- a/crates/sqlload/src/loader/_helpers/python.rs +++ b/crates/sqlload/src/loader/_helpers/python.rs @@ -686,7 +686,7 @@ pub(crate) fn typed_batch( } /// The batches' target schema: each read column in its planned type. -fn read_schema(table: &TablePlan) -> SchemaRef { +pub(crate) fn read_schema(table: &TablePlan) -> SchemaRef { let fields: Vec = table .read_columns .iter() @@ -699,7 +699,10 @@ fn read_schema(table: &TablePlan) -> SchemaRef { } /// `batch` with the target's columns (excluded columns dropped), each converted to its planned type (an array from the list, or from the JSON text a list arrives as); a value that does not convert is an error naming the column. -fn read_types(batch: &RecordBatch, target: &SchemaRef) -> Result { +pub(crate) fn read_types( + batch: &RecordBatch, + target: &SchemaRef, +) -> Result { let arrays: Vec = target .fields() .iter() @@ -902,6 +905,7 @@ impl PythonFeed { note: Some(NO_DATA_NOTE.to_owned()), after_load: None, replaced_keys: None, + transforms: Vec::new(), }); } let unsettled = self.publisher.is_none() @@ -925,6 +929,7 @@ impl PythonFeed { note: None, after_load: None, replaced_keys, + transforms: Vec::new(), }) } @@ -1355,6 +1360,7 @@ fn child_plan(parent: &TablePlan, (path, table): (&str, &str)) -> TablePlan { lists: ListChoices::new(), }), sheet: None, + transform: None, cdc: None, checks: None, warehouse_source: None, diff --git a/crates/sqlload/src/loader/_helpers/render.rs b/crates/sqlload/src/loader/_helpers/render.rs index cb56394a1..063b02286 100644 --- a/crates/sqlload/src/loader/_helpers/render.rs +++ b/crates/sqlload/src/loader/_helpers/render.rs @@ -23,8 +23,10 @@ use crate::loader::_helpers::publish::records_files; use crate::loader::_helpers::python::list_choices; use crate::loader::_helpers::python::unnested_table; use crate::loader::_helpers::resolve::display_name; -use crate::loader::_helpers::sources::source_readers; -use crate::loader::_helpers::table_plan::{is_direct, publishes_once, pulls_unloaded}; +use crate::loader::_helpers::sources::{source_readers, unarchived}; +use crate::loader::_helpers::table_plan::{ + archived_reason, is_direct, publishes_once, pulls_unloaded, +}; use crate::loader::constants::{ CONSTANT_NOTE, DRAINED_REASON, LABEL_WIDTH, NULL_KEYS_NOTE, PARALLEL_READERS, PLAN_DETAIL_COLUMN, PRIMARY_KEY_PLACEHOLDER, RESULT_SCHEMA_VERSION, @@ -42,6 +44,7 @@ use crate::project::models::{FilesFormat, OutputCompression}; use crate::routing::models::{Route, Step, Store}; use crate::row_batches::models::Unnest; use crate::run_log::main::describe_state::describe_state; +use crate::run_log::main::runs_table::runs_table; use crate::schema_plan::models::{ColumnChange, ColumnOrigin, TableAction}; use crate::store_paths::models::StoreKind; @@ -326,6 +329,28 @@ fn column_details(table: &TablePlan) -> Vec { details } +/// The `_sqlload_runs` tables `apply` creates, one per schema, where the commits of tables with a `transform` record its version. +fn run_log_items(tables: &[&TablePlan]) -> Vec { + let mut items: Vec = Vec::new(); + for table in tables { + let spec = &table.spec; + let name = display_name(spec.warehouse, &runs_table(&spec.destination)); + if spec + .transform + .as_ref() + .is_some_and(|transform| transform.creates_log) + && !items.iter().any(|item| item.name == name) + { + items.push(PlanItem { + name, + action: "create table (run log)".to_owned(), + details: Vec::new(), + }); + } + } + items +} + pub(crate) fn stage_item(stage: &PlannedStage) -> PlanItem { PlanItem { name: stage.name.clone(), @@ -366,8 +391,12 @@ pub(crate) fn load_item(table: &TablePlan) -> PlanItem { true => NULL_KEYS_NOTE, false => "", }; + let rebuilt = match table.schema.action { + TableAction::Recreate => " into a new table, swapped in at the end (--reload)", + _ => "", + }; format!( - "merge (unique_key: {}{origin}{repeats}{nulls})", + "merge (unique_key: {}{origin}{repeats}{nulls}){rebuilt}", unique_key.join(", ") ) } @@ -422,7 +451,13 @@ pub(crate) fn load_item(table: &TablePlan) -> PlanItem { } }; let source = match (&spec.files, &spec.python) { - (Some(files), _) => format!("files {} {}", files.location, files.encoding), + (Some(files), _) => match &files.archive { + Some(archive) => format!( + "files {} and {archive} {}", + files.location, files.encoding + ), + None => format!("files {} {}", files.location, files.encoding), + }, (None, Some(python)) => format!("python {} ({})", python.module, spec.source_name), (None, None) if spec.warehouse_source.is_some() => { let read = spec @@ -464,9 +499,10 @@ pub(crate) fn load_item(table: &TablePlan) -> PlanItem { ), }, }; - let source = match &spec.tunnel { - Some(route) => format!("{source} via {route}"), - None => source, + let source = match (&spec.tunnel, &spec.transform) { + (Some(route), _) => format!("{source} via {route}"), + (None, Some(transform)) => format!("{source} transform {}", transform.function), + (None, None) => source, }; let mut details = vec![detail("source", source)]; details.extend( @@ -522,14 +558,28 @@ pub(crate) fn load_item(table: &TablePlan) -> PlanItem { 0 => String::new(), count => format!(", {count} skipped (paths that do not match path_columns)"), }; - let fresh = match table.mode { - WriteMode::Recreate | WriteMode::Replace => "to read", + let fresh = match (&table.mode, &table.schema.action) { + (WriteMode::Recreate | WriteMode::Replace, _) | (_, TableAction::Recreate) => { + "to read" + } _ => "new", }; + let folders = match &source.archive { + Some(archive) => { + let archived = read.files.len() - unarchived(table, &read.files).len(); + format!( + ": {} in {}, {} in {archive}", + thousands((read.files.len() - archived) as u64), + source.location, + thousands(archived as u64) + ) + } + None => String::new(), + }; detail( "files", format!( - "{} {fresh} of {} listed{skipped}{unmatched}", + "{} {fresh} of {} listed{folders}{skipped}{unmatched}", thousands(read.files.len() as u64), thousands(source.table.files.len() as u64) ), @@ -581,12 +631,19 @@ pub(crate) fn load_item(table: &TablePlan) -> PlanItem { "changed columns move to the end of the column order", )); } - details.extend( - table - .route - .as_ref() - .map(|route| detail("route", route_text(table, route))), - ); + let why_shown = + (table.stage.is_some() && spec.files.is_some() && !is_direct(table.route.as_ref())) + || spec.warehouse_source.is_some(); + let archived = archived_reason(spec, table.files.as_ref()).filter(|_| !why_shown); + details.extend(table.route.as_ref().map(|route| { + let text = route_text(table, route); + detail( + "route", + archived + .as_ref() + .map_or_else(|| text.clone(), |why| format!("{text} ({why})")), + ) + })); if let Some(WarehouseSource { route: WarehouseRoute::Unload { stage, .. }, .. @@ -597,10 +654,7 @@ pub(crate) fn load_item(table: &TablePlan) -> PlanItem { if let Some(reason) = table .route .as_ref() - .filter(|_| { - (table.stage.is_some() && spec.files.is_some() && !is_direct(table.route.as_ref())) - || spec.warehouse_source.is_some() - }) + .filter(|_| why_shown) .map(|route| route.reason.clone()) { details.push(detail("why", reason)); @@ -648,6 +702,20 @@ pub(crate) fn load_item(table: &TablePlan) -> PlanItem { format!("logged: {}", describe_state(logged.current.as_ref())), ) })); + details.extend( + spec.transform + .iter() + .filter(|transform| transform.older > 0) + .map(|transform| { + let unit = match spec.files.is_some() { + true => ("file", "files"), + false => ("range", "ranges"), + }; + let older = counted(transform.older, unit); + let every = unit.0; + detail("", format!("transform changed: {older} loaded under an earlier version (--reload transforms every {every} again)")) + }), + ); details.extend( table .output @@ -1016,7 +1084,7 @@ pub(crate) fn text_view(plan: &ProjectPlan) -> PlanView { sections(plan, true) } -/// Why a files table that exists and has no new file to read is skipped, as `up to date: 0 new of 3 files in drop/stock_counts/ (csv in ZIP archives)`; `None` for any other table, for one with a schema change or an unfinished load to repair, and for a `delete_insert` over a `--start-cursor`/`--end-cursor` range (`ranged`), which clears that range even with no file in it. +/// Why a files table that exists and has no new file to read is skipped, as `up to date: 0 new of 3 files in drop/stock_counts/ (csv in ZIP archives)`; `None` for any other table, for one with a schema change, an unfinished load to repair or a `transform` changed since its last load, and for a `delete_insert` over a `--start-cursor`/`--end-cursor` range (`ranged`), which clears that range even with no file in it. pub(crate) fn up_to_date(table: &TablePlan, ranged: bool) -> Option { let (source, read) = table.spec.files.as_ref().zip(table.files.as_ref())?; let incremental = @@ -1026,7 +1094,15 @@ pub(crate) fn up_to_date(table: &TablePlan, ranged: bool) -> Option { .output .as_ref() .is_none_or(|output| output.pending.is_empty()); - if !settled || ranged && matches!(table.mode, WriteMode::DeleteInsert { .. }) { + let transform_changed = table + .spec + .transform + .as_ref() + .is_some_and(|transform| transform.older > 0); + if !settled + || transform_changed + || ranged && matches!(table.mode, WriteMode::DeleteInsert { .. }) + { return None; } (incremental && read.files.is_empty()).then(|| { @@ -1061,7 +1137,9 @@ fn sections(plan: &ProjectPlan, text: bool) -> PlanView { .iter() .filter(|table| table.schema.action != TableAction::Unchanged && table.output.is_none()) .collect(); + let logs = run_log_items(&ready); let creates = stages.len() + + logs.len() + changed .iter() .filter(|table| { @@ -1071,7 +1149,7 @@ fn sections(plan: &ProjectPlan, text: bool) -> PlanView { ) }) .count(); - let alters = stages.len() + changed.len() - creates; + let alters = stages.len() + logs.len() + changed.len() - creates; let mut errors: Vec = plan .tables .iter() @@ -1136,6 +1214,7 @@ fn sections(plan: &ProjectPlan, text: bool) -> PlanView { .iter() .map(|stage| stage_item(stage)) .chain(changed.iter().map(|table| object_item(table))) + .chain(logs) .collect(), loads: match text { false => ready.iter().map(|table| load_item(table)).collect(), @@ -1344,6 +1423,10 @@ fn table_result(outcome: &TableOutcome) -> TableResult { Ok(loaded) => loaded.replaced_keys.clone(), Err(_) => None, }, + transforms: match &outcome.loaded { + Ok(loaded) => loaded.transforms.clone(), + Err(_) => Vec::new(), + }, } } diff --git a/crates/sqlload/src/loader/_helpers/resolve.rs b/crates/sqlload/src/loader/_helpers/resolve.rs index c4ac85251..8aabf2468 100644 --- a/crates/sqlload/src/loader/_helpers/resolve.rs +++ b/crates/sqlload/src/loader/_helpers/resolve.rs @@ -1,24 +1,34 @@ use crate::adapters::contract::models::TableName; use crate::adapters::dispatch::main::dialect::connection_warehouse::connection_warehouse; +use crate::adapters::dispatch::main::dialect::files_destination::files_destination; use crate::adapters::dispatch::main::dialect::sql_adapter::sql_adapter; +use crate::adapters::dispatch::main::session::live_columns::live_columns; +use crate::adapters::dispatch::models::Destination; use crate::column_types::models::Warehouse; use crate::load_naming::main::is_load_column::is_load_column; use crate::load_naming::main::loaded_name::loaded_name; use crate::load_naming::models::Naming; use crate::loader::_helpers::specs::ColumnSelection; +use crate::loader::_helpers::specs::missing_description; use crate::loader::_helpers::specs::table_name; use crate::loader::_helpers::specs::{ ColumnChoices, SourceContext, build_spec, column_selection, pinned_types, table_naming, }; -use crate::loader::constants::PRIMARY_KEY_PLACEHOLDER; -use crate::loader::models::{ResolvedTables, SkippedTable, TableSpec, UnreadSource}; +use crate::loader::_helpers::table_plan::{blocked_plan, plan_table}; +use crate::loader::constants::{NO_TRANSFORM_HOST, PRIMARY_KEY_PLACEHOLDER}; +use crate::loader::errors::LoaderError; +use crate::loader::models::{ + PlanOptions, ProjectPlan, ResolvedTables, SkippedTable, TablePlan, TableSpec, TableTransform, + UnreadSource, +}; use crate::postgres_source::models::SourceTable; use crate::project::main::config_at::config_at; use crate::project::main::motherduck_path::motherduck_path; use crate::project::models::{ - ConfigLines, ConnectionConfig, ConstantColumns, ProjectConfig, SourceConfig, SourceKind, - TableOverride, TableSettings, UnnestSetting, + ConfigLines, ConnectionConfig, ConstantColumns, NewTables, ProjectConfig, SourceConfig, + SourceKind, TableOverride, TableSettings, UnnestSetting, }; +use crate::similar_names::main::nearest_name::nearest_name; use crate::ssh_tunnel::models::SshTunnel; use std::collections::BTreeMap; @@ -64,7 +74,16 @@ pub(crate) fn resolve_tables( ); target_display_name(&target, name) }; - match resolve_table(&context, table, &overrides) { + let transform = table_transform(project, (source, &overrides), &qualified); + let resolved_table = match (resolve_table(&context, table, &overrides), transform) { + (Err(problems), _) => Err(problems), + (Ok(_), Err(problem)) => Err(vec![problem]), + (Ok(found), Ok(transform)) => Ok(found.map(|mut spec| { + spec.transform = transform; + spec + })), + }; + match resolved_table { Ok(Ok(mut spec)) => { let lines = &project.config_lines; spec.settings_at = settings_at(lines, source_name, &qualified); @@ -128,6 +147,32 @@ pub(crate) fn resolve_tables( resolved } +/// The `transform` a table runs (its own, else its source's) with the version `project` imported it at, or why it cannot run. +fn table_transform( + project: &ProjectConfig, + (source, overrides): (&SourceConfig, &TableOverride), + qualified: &str, +) -> Result, String> { + let Some(function) = overrides + .settings + .transform + .clone() + .or_else(|| source.defaults.transform.clone()) + else { + return Ok(None); + }; + match project.transforms.get(&function) { + Some(Ok(version)) => Ok(Some(TableTransform { + function, + version: version.clone(), + older: 0, + creates_log: false, + })), + Some(Err(reason)) => Err(format!("transform {function}: {reason}")), + None => Err(NO_TRANSFORM_HOST.replace("{place}", &format!("table {qualified}"))), + } +} + /// A problem for each pair of `specs` (one source's) that dlt's naming loads into one table although their names differ (`order-items` and `order_items`), where one would overwrite or mix with the other's rows. fn merged_names(source_name: &str, source: &SourceConfig, specs: &[TableSpec]) -> Vec { let written = |spec: &TableSpec| { @@ -549,3 +594,130 @@ pub(crate) fn spec_names(spec: &TableSpec) -> [String; 3] { spec.source_table.qualified_name(), ] } + +/// `planned`, the selected tables of one source, added to `plan`, which is returned (`unselected` and `found`, the destinations and options, tell whether this is the source's first load): a table that appeared after the source's first load (`first_load`: none of its tables, selected or not, existed yet) and is not named exactly in `include` is skipped or fails as `new_tables` says, and a table without a description is blocked. +pub(crate) async fn admit_tables( + mut plan: ProjectPlan, + (source_name, source): (&String, &SourceConfig), + (planned, unselected): (Vec, &[TableSpec]), + found: (&BTreeMap, &PlanOptions), +) -> Result { + let first_load = first_load(&planned, unselected, found).await?; + let mut planned = planned; + if let Some(problem) = partial_capture(source_name, source.kind, (&planned, unselected)) { + for table in &mut planned { + table.schema = blocked_plan(vec![problem.clone()]); + } + } + for mut table in planned { + let spec = &table.spec; + let qualified = spec.source_table.qualified_name(); + let is_new = !first_load && !table.exists && !spec.approved; + let has_settings = source.tables.contains_key(&qualified); + match (is_new, spec.new_tables, has_settings) { + (true, NewTables::Skip, true) => plan.skipped.push(SkippedTable { + source_name: source_name.clone(), + source_table: qualified, + destination: display_name(spec.warehouse, &spec.destination), + reason: "new table (new_tables = \"skip\") has settings but is not approved; add it to include".to_owned(), + }), + (true, NewTables::Skip, false) => plan.skipped.push(SkippedTable { + source_name: source_name.clone(), + source_table: qualified, + destination: display_name(spec.warehouse, &spec.destination), + reason: "new table (new_tables = \"skip\"); add it to include to sync it".to_owned(), + }), + (true, NewTables::Fail, _) => { + plan.failed_sources.insert(source_name.clone()); + let problem = format!( + "{qualified}: new table and new_tables = \"fail\"; add it to include or exclude it" + ); + plan.problem_names.push(( + problem.clone(), + display_name(spec.warehouse, &spec.destination), + )); + plan.problems.push(problem); + } + _ => { + if let Some(problem) = missing_description(source_name, source, &spec.source_table) + { + table.schema = blocked_plan(vec![problem]); + } + plan.tables.push(table); + } + } + } + Ok(plan) +} + +/// Why the selection cannot load `postgres_cdc` source `name`, which blocks every table it picks: it picks some of its tables (`planned`) but not all (`unselected`), and the slot they share would confirm the left-out tables' changes unwritten. A plan problem alone would not stop the picked tables loading. +fn partial_capture( + name: &str, + kind: SourceKind, + (planned, unselected): (&[TablePlan], &[TableSpec]), +) -> Option { + let partial = kind == SourceKind::PostgresCdc && !planned.is_empty() && !unselected.is_empty(); + partial.then(|| { + let left_out: Vec = unselected + .iter() + .map(|spec| spec.source_table.qualified_name()) + .collect(); + format!( + "this selection leaves out {} of source {name}, whose changes share its one slot and would be skipped; select the whole source, or move those tables to another postgres_cdc source", + left_out.join(", ") + ) + }) +} + +/// Whether this is the source's first load: none of its tables exists in the destination, whether selected (`planned`) or not (`unselected`, looked up only while none is found, a table that cannot be read counting as absent), so `--select` cannot make a later table look present at the first load. +async fn first_load( + planned: &[TablePlan], + unselected: &[TableSpec], + (destinations, options): (&BTreeMap, &PlanOptions), +) -> Result { + if planned.iter().any(|table| table.exists) { + return Ok(false); + } + for spec in unselected { + let Some(destination) = destinations.get(&spec.target) else { + continue; + }; + let exists = match files_destination(destination) { + Some(_) => plan_table(spec, destination, options) + .await + .is_ok_and(|table| table.exists), + None => live_columns(destination, &spec.destination) + .await + .is_ok_and(|live| live.is_some()), + }; + if exists { + return Ok(false); + } + } + Ok(true) +} + +/// A warning for each selector matching no source or table, with the nearest name: `load` would otherwise say only `0 selected`. +pub(crate) fn unmatched_selectors( + project: &ProjectConfig, + options: &PlanOptions, + names: &[String], +) -> Vec { + let known: Vec<&str> = names + .iter() + .map(String::as_str) + .chain(project.sources.keys().map(String::as_str)) + .collect(); + options + .select + .iter() + .chain(&options.exclude) + .filter(|selector| !known.contains(&selector.as_str())) + .map( + |selector| match nearest_name(selector, known.iter().copied()) { + Some(nearest) => format!("no table matches {selector}; did you mean {nearest}?"), + None => format!("no table matches {selector}"), + }, + ) + .collect() +} diff --git a/crates/sqlload/src/loader/_helpers/run_log.rs b/crates/sqlload/src/loader/_helpers/run_log.rs index 561897864..4708c3e56 100644 --- a/crates/sqlload/src/loader/_helpers/run_log.rs +++ b/crates/sqlload/src/loader/_helpers/run_log.rs @@ -1,14 +1,16 @@ use std::cmp::Ordering; -use std::collections::{BTreeMap, HashMap}; +use std::collections::{BTreeMap, BTreeSet, HashMap}; use arrow_array::RecordBatch; use bytes::Bytes; use crate::adapters::contract::main::cursor_text::cursor_text; +use crate::adapters::contract::main::staging_id::staging_id; use crate::adapters::contract::models::WriteMode; use crate::adapters::dispatch::main::dialect::sql_adapter::sql_adapter; use crate::adapters::dispatch::main::session::ddl_guard::ddl_guard; use crate::adapters::dispatch::main::session::execute::execute; +use crate::adapters::dispatch::main::session::execute_ddl::execute_ddl; use crate::adapters::dispatch::main::session::live_columns::live_columns; use crate::adapters::dispatch::main::session::text_values::text_values; use crate::adapters::dispatch::models::Destination; @@ -21,12 +23,17 @@ use crate::file_store::main::write_file::write_file; use crate::files_source::constants::{ SOURCE_FILE_COLUMN, SOURCE_VERSION_COLUMN, VERSION_SEPARATOR, }; +use crate::files_source::main::merge_order::merge_order; use crate::files_source::models::{DataFile, LoadedFile}; use crate::loader::_helpers::python::batch_cursor_max; -use crate::loader::_helpers::sources::{FilesPlanRequest, column_fix, live_has, loaded_file}; +use crate::loader::_helpers::sources::{FilesPlanRequest, LoadedSource, column_fix, live_has}; +use crate::loader::constants::TRANSFORM_REWROTE_NOTE; use crate::loader::errors::LoaderError; -use crate::loader::models::{PythonLoad, RunEntry, TablePlan, TableSpec}; +use crate::loader::models::{ + MergeRebuild, ProjectPlan, PythonLoad, RunEntry, TablePlan, TableSpec, TransformRun, +}; use crate::project::models::WriteStrategy; +use crate::row_batches::models::Unnest; use crate::run_log::constants::RECORD_FORMAT_VERSION; use crate::run_log::main::cursor_order::cursor_order; use crate::run_log::main::decode_runs::decode_runs; @@ -41,6 +48,7 @@ use crate::run_log::main::select_runs::select_runs; use crate::run_log::models::{ DecodedRuns, RecordTable, RunEvent, RunOutcome, RunRecord, StateEvent, UnreadableRecord, }; +use crate::schema_plan::models::TableAction; use crate::store_paths::models::Location; use crate::time_text::main::now_utc_text::now_utc_text; @@ -318,6 +326,7 @@ pub(crate) fn run_entry( note: None, recorded_at: String::new(), batch: None, + transform: None, }, cursor_type, before: window.and_then(|window| window.maximum.clone()), @@ -437,6 +446,165 @@ pub(crate) fn merged_record( ), recorded_at: now_utc_text(), batch: Some(batch), + transform: None, + } +} + +/// The time a load's `transform` took on each file or range, once the load succeeded. The calls no commit recorded (`recorded` are), as when batches commit straight into the table or into a files target, are recorded now: an insert into the run log `apply` made, never creating it, and only warned about when it fails, since the rows committed (a logged table's run log holds its runs only). +pub(crate) async fn transformed( + table: &TablePlan, + (destination, run, recorded): (&Destination, Option, usize), +) -> Vec { + let Some(run) = run else { + return Vec::new(); + }; + let times = std::mem::take( + &mut *run + .times + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ); + let spec = &table.spec; + let record = spec + .transform + .as_ref() + .filter(|_| times.len() > recorded && !spec.logged) + .map(|transform| transform_record(table, &transform.version, ×[recorded..])); + let written = match (record, sql_adapter(spec.warehouse).runs_sql()) { + (None, _) => Ok(()), + (Some(record), false) => append_run(table, destination, &record).await, + (Some(record), true) => { + match insert_run(&runs_table(&spec.destination), &record, spec.warehouse) { + Ok(insert) => execute(destination, vec![insert]).await, + Err(error) => Err(LoaderError::Invalid(error)), + } + } + }; + if let Err(error) = written { + tracing::warn!(table = %spec.destination.table, %error, "the rows loaded, but the run log did not record the transform's version, so plan cannot tell which files an earlier version loaded"); + } + times +} + +/// The record a commit writes with its rows for the `transform` calls since the last commit, which it counts as recorded; `None` without calls or a `transform`. +pub(crate) fn commit_transform( + table: &TablePlan, + transform: Option<&mut (TransformRun, usize)>, +) -> Option { + let (run, recorded) = transform?; + let version = &table.spec.transform.as_ref()?.version; + let times = run + .times + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let since = times.get(*recorded..).filter(|since| !since.is_empty())?; + let record = transform_record(table, version, since); + *recorded = times.len(); + Some(record) +} + +/// `record` (a logged table's, or a merged file's) carrying the version of the table's `transform`, if it has one. +pub(crate) fn with_transform(table: &TablePlan, record: RunRecord) -> RunRecord { + RunRecord { + transform: table + .spec + .transform + .as_ref() + .map(|transform| transform.version.clone()), + ..record + } +} + +/// `plan` with what its `transform` needs from the run log: whether `apply` must create `_sqlload_runs`, and how many of the files or ranges the table holds loaded under another version; a table rewritten whole (`replace`, `--reload`) or a logged one is left as it is, and one blocked. +pub(crate) async fn with_transform_version( + mut plan: TablePlan, + destination: &Destination, +) -> Result { + let rewritten = matches!(plan.mode, WriteMode::Replace | WriteMode::Recreate) + || plan.schema.action == TableAction::Recreate; + let runs_sql = sql_adapter(plan.spec.warehouse).runs_sql(); + let compares = plan.exists && !rewritten; + let Some(current) = plan + .spec + .transform + .as_ref() + .map(|transform| transform.version.clone()) + else { + return Ok(plan); + }; + if plan.spec.logged || plan.schema.blocked.is_some() || !(compares || runs_sql) { + return Ok(plan); + } + let (log, exists) = read_run_log(&plan, destination).await?; + let older = match compares { + true => older_files(&log.records, &plan.spec, ¤t), + false => 0, + }; + if let Some(transform) = plan.spec.transform.as_mut() { + transform.older = older; + transform.creates_log = runs_sql && !exists; + } + Ok(plan) +} + +/// How many of the files or ranges the table of `spec` holds loaded under a `transform` version other than `current`: the version of each one's newest record (a merged file's by its name, whatever its version) since the last reset or load that rewrote the table whole. +pub(crate) fn older_files(records: &[RunRecord], spec: &TableSpec, current: &str) -> usize { + let mut ordered: Vec<&RunRecord> = records + .iter() + .filter(|record| { + record.source == spec.source_name && record.source_table == spec.source_table.name + }) + .collect(); + ordered.sort_by(|left, right| left.recorded_at.cmp(&right.recorded_at)); + let mut versions: BTreeMap<&str, &str> = BTreeMap::new(); + for record in ordered { + if record.event == RunEvent::Reset || record.note.as_deref() == Some(TRANSFORM_REWROTE_NOTE) + { + versions.clear(); + } + let merged = !record + .note + .as_deref() + .is_some_and(|note| note.starts_with("transform")); + if let Some(version) = &record.transform { + versions.extend(record.files.iter().map(|file| match merged { + true => (merged_name(file), version.as_str()), + false => (file.as_str(), version.as_str()), + })); + } + } + versions + .values() + .filter(|version| **version != current) + .count() +} + +/// The record of a load of a table with a `transform`: the version it ran, and the files or ranges it ran on (`times`), which the next `plan` compares with. +pub(crate) fn transform_record( + table: &TablePlan, + version: &str, + times: &[crate::python_sources::models::TransformTime], +) -> RunRecord { + let rewrote = matches!(table.mode, WriteMode::Replace | WriteMode::Recreate); + let record = event_record( + &record_table(&table.spec), + StateEvent { + event: RunEvent::Run, + cursor_before: None, + cursor_after: None, + state: None, + note: match rewrote { + true => TRANSFORM_REWROTE_NOTE, + false => "transform", + } + .to_owned(), + recorded_at: now_utc_text(), + }, + ); + RunRecord { + files: times.iter().map(|time| time.file.clone()).collect(), + transform: Some(version.to_owned()), + ..record } } @@ -455,6 +623,31 @@ pub(crate) fn merge_reset_record(spec: &TableSpec) -> RunRecord { ) } +/// Create the `_sqlload_runs` of each unblocked table with a `transform` whose plan found none, once per schema, before any load commits a record there; one `created` line each. +pub(crate) async fn create_run_logs( + plan: &ProjectPlan, + destinations: &BTreeMap, +) -> Result, LoaderError> { + let mut lines: Vec = Vec::new(); + let mut created: BTreeSet = BTreeSet::new(); + for table in plan.tables.iter().filter(|table| { + let transform = table.spec.transform.as_ref(); + table.schema.blocked.is_none() && transform.is_some_and(|transform| transform.creates_log) + }) { + let spec = &table.spec; + let name = sql_adapter(spec.warehouse).render_display_name(&runs_table(&spec.destination)); + if !created.insert(name.clone()) { + continue; + } + let destination = destinations.get(&spec.target).ok_or_else(|| { + LoaderError::Invalid(format!("target {:?} is not connected", spec.target)) + })?; + ensure_runs_table(table, destination).await?; + lines.push(format!("created {name} run log")); + } + Ok(lines) +} + /// Before a files source's merge creates its table: `_sqlload_runs` when it does not exist yet, and a reset, so records left from an earlier table of the name stop counting; written before the `CREATE TABLE`, so no table exists without it. pub(crate) async fn reset_merge_log( table: &TablePlan, @@ -679,3 +872,167 @@ pub(crate) fn listed_sources( "the table holds rows of its replace_key ({columns}) values from {first}{more}, which is still in the folder; with two files of the folder holding one value, which of them is current can't be told, so replacing the value by {this} would delete the rows {first} loaded, as a value split across part files would keep only its last part. One file must hold all of a value's rows. Fix: put all the rows of each ({columns}) value in one file, or remove the outdated file, {first} or {this}, from the folder" )) } + +/// The files a merge already merged into `request`'s table at the version listed now, which it skips (one rewritten in place since is merged again, as the newest), and under `--reload` none, with every file the table's merges recorded instead (oldest first), whose order the reload keeps; or why a files source cannot load into it: only into a DuckDB or Postgres table, never with `state = "logged"`, nor `--reload` with `unnest` or of a `delete_insert` (one commit would lose the files' order); none while the table does not exist. +pub(crate) async fn merge_ledger( + request: &FilesPlanRequest<'_>, +) -> Result, Vec), String>, LoaderError> { + let spec = request.spec; + let destination = match request.loaded { + LoadedSource::Query(destination) => destination, + LoadedSource::Recorded(_) => { + return Ok(Err("merge is not supported for files targets".to_owned())); + } + }; + let what = match spec.strategy { + WriteStrategy::Merge => "merge", + _ => "load", + }; + if !sql_adapter(spec.warehouse).merges_files() { + return Ok(Err(format!( + "{} from a files source loads into DuckDB and Postgres targets so far, not {}; use append, or delete_insert on a partition cursor", + spec.strategy, spec.warehouse + ))); + } + if request.options.reload && spec.strategy != WriteStrategy::Merge { + return Ok(Err(format!( + "--reload would {what} every file in one commit and lose their order; drop the table and load again to {what} every file anew" + ))); + } + if request.options.reload && spec.unnest != Unnest::None { + return Ok(Err( + "--reload of a merge files table does not rebuild its child tables (unnest); drop the table and its child tables and load again to merge every file anew".to_owned(), + )); + } + if spec.logged { + return Ok(Err(format!( + "state = \"logged\" does not apply to a files source's {}, which records each file it loads", + spec.strategy + ))); + } + if let Some(cursor) = &spec.cursor { + return Ok(Err(format!( + "cursor_column {} does not apply to a files source's merge, which picks the files it has not merged from the run log, whatever their partition; remove it and set write_strategy = \"merge\", or keep it and use delete_insert", + cursor.column + ))); + } + let merged = match request.live { + None => return Ok(Ok((Vec::new(), Vec::new()))), + Some(_) => merged_files(spec, destination).await?, + }; + if request.options.reload { + return Ok(merged.map(|merged| (Vec::new(), merged))); + } + let location = request.files.location.to_string(); + let listed: HashMap = request + .files + .table + .files + .iter() + .map(|data| { + ( + format!("{location}{}", data.file.relative_key), + data.file.version(), + ) + }) + .collect(); + Ok(merged.map(|merged| (unchanged_files(merged, &listed), Vec::new()))) +} + +/// `files` in the order a merge `--reload` applies them: those the table's merges recorded (`merged`, oldest first) in the order of their first record, matched by name and version, else, for a file read from `archive` (where `after_load` moved it, so it may have a new version), by name; then the others oldest modified first, as a load merges new files. A file in the source's folder that matches only by name is a newer file under a merged name, so it merges after the recorded ones, as a load would merge it. +pub(crate) fn reload_order( + files: Vec, + location: &str, + merged: &[LoadedFile], + archive: Option<&Location>, +) -> Vec { + let mut by_version: HashMap<(&str, &str), usize> = HashMap::new(); + let mut by_name: HashMap<&str, usize> = HashMap::new(); + for (index, file) in merged.iter().enumerate() { + by_version + .entry((file.name.as_str(), file.version.as_str())) + .or_insert(index); + by_name.entry(file.name.as_str()).or_insert(index); + } + let mut recorded: Vec<(usize, DataFile)> = Vec::new(); + let mut others: Vec = Vec::new(); + for data in files { + let name = format!("{location}{}", data.file.relative_key); + let version = data.file.version(); + let archived = + archive.is_some_and(|archive| data.file.path.as_ref().starts_with(&archive.prefix)); + match by_version + .get(&(name.as_str(), version.as_str())) + .or_else(|| by_name.get(name.as_str()).filter(|_| archived)) + { + Some(index) => recorded.push((*index, data)), + None => others.push(data), + } + } + recorded.sort_by_key(|(index, _)| *index); + recorded + .into_iter() + .map(|(_, data)| data) + .chain(merge_order(others)) + .collect() +} + +/// Before a files source's merge or `delete_insert` by `replace_key` writes: `_sqlload_runs` when it does not exist yet; under a merge's `--reload` of an existing table, the table each file merges into, created beside it in its current columns, and the reset its swap commits first, stamped before any file's record, so the files merged into the old table stop counting. +pub(crate) async fn merge_rebuild( + table: &TablePlan, + destination: &Destination, +) -> Result, LoaderError> { + ensure_merge_log(table, destination).await?; + if table.schema.action != TableAction::Recreate || table.spec.strategy != WriteStrategy::Merge { + return Ok(None); + } + let spec = &table.spec; + let adapter = sql_adapter(spec.warehouse); + let id = staging_id( + &uuid::Uuid::new_v4().to_string(), + chrono::Utc::now().naive_utc(), + ); + let build = adapter.staging_relation_name(&spec.destination, &id); + let create = adapter.render_create_build_table( + &build, + &spec.destination, + (&table.schema.write_columns, table.live.as_deref(), None), + ); + execute_ddl(destination, create).await?; + Ok(Some(MergeRebuild { + build, + records: vec![merge_reset_record(spec)], + })) +} + +/// Swap a merge `--reload`'s table in, with its run records in the same transaction. +pub(crate) async fn swap_rebuild( + table: &TablePlan, + destination: &Destination, + rebuild: &MergeRebuild, +) -> Result<(), LoaderError> { + let spec = &table.spec; + let adapter = sql_adapter(spec.warehouse); + let runs = runs_table(&spec.destination); + let mut statements = adapter.render_swap(&spec.destination, &rebuild.build, &[]); + for record in &rebuild.records { + let insert = insert_run(&runs, record, spec.warehouse).map_err(LoaderError::Invalid)?; + statements = adapter.render_logged(statements, insert); + } + let _held = ddl_guard(destination).await; + execute(destination, statements).await +} + +/// A row of the loaded-files query: `file|version`, or just the file when `versioned` is false. +pub(crate) fn loaded_file(text: String, versioned: bool) -> LoadedFile { + match (versioned, text.rsplit_once(VERSION_SEPARATOR)) { + (true, Some((name, version))) => LoadedFile { + name: name.to_owned(), + version: version.to_owned(), + }, + _ => LoadedFile { + name: text, + version: String::new(), + }, + } +} diff --git a/crates/sqlload/src/loader/_helpers/sources.rs b/crates/sqlload/src/loader/_helpers/sources.rs index 28bf838be..b34b84ba9 100644 --- a/crates/sqlload/src/loader/_helpers/sources.rs +++ b/crates/sqlload/src/loader/_helpers/sources.rs @@ -1,11 +1,12 @@ use crate::adapters::contract::models::BoundParameter; use crate::adapters::dispatch::main::dialect::sql_adapter::sql_adapter; use arrow_array::RecordBatch; -use std::collections::{BTreeMap, HashMap}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use std::collections::BTreeMap; +use std::sync::Arc; use tokio_postgres::Client; use crate::adapters::contract::models::TableName; -use crate::adapters::dispatch::main::dialect::files_destination::files_destination; use crate::adapters::dispatch::main::session::live_columns::live_columns; use crate::adapters::dispatch::main::session::text_values::text_values; use crate::adapters::dispatch::models::Destination; @@ -17,55 +18,56 @@ use crate::bigquery_source::main::read_catalog::read_catalog as read_bigquery_ca use crate::bigquery_source::main::read_rows::read_rows as read_bigquery_rows; use crate::bigquery_source::models::{BigQueryReadRequest, BigQueryReader}; use crate::column_types::main::arrow_field::arrow_field; -use crate::column_types::main::parse_type::parse_type; use crate::column_types::models::{ColumnType, Warehouse}; use crate::config::models::{DestinationConfig, DestinationKind}; use crate::file_store::main::open_folder::open_folder; +use crate::file_store::models::{OpenFolder, StoredFile}; use crate::files_source::constants::{ - NANOSECOND_SCALE, SOURCE_FILE_COLUMN, SOURCE_ROW_COLUMN, SOURCE_VERSION_COLUMN, - VERSION_SEPARATOR, + SOURCE_FILE_COLUMN, SOURCE_ROW_COLUMN, SOURCE_VERSION_COLUMN, }; use crate::files_source::errors::FilesError; -use crate::files_source::main::constant_columns::constant_columns; use crate::files_source::main::discover_files::discover_files; +use crate::files_source::main::discover_reload_files::discover_reload_files; use crate::files_source::main::merge_order::merge_order; use crate::files_source::main::move_to_prefix::move_to_prefix; -use crate::files_source::main::path_template::path_template; use crate::files_source::main::read_files::read_files; use crate::files_source::main::select_files::select_files; +use crate::files_source::main::shape_rows::shape_rows; use crate::files_source::models::{ AfterLoad, BatchFooters, DataFile, FileColumn, FileSelection, FilesEncoding, FilesRead, - FilesTable, LoadedFile, PathValues, SourceFormat, StoredRead, TextChecks, + FilesTable, LoadedFile, StoredRead, TextChecks, }; use crate::http::main::http_client::build_http_client; use crate::http::models::ConnectionCounter; use crate::load_naming::main::loaded_name::loaded_name; use crate::load_naming::models::Naming; -use crate::loader::_helpers::destination::folder_claims; +use crate::loader::_helpers::destination::{folder_claims, nanosecond_columns}; use crate::loader::_helpers::python::{fitting_type, python_tables}; +use crate::loader::_helpers::python::{read_schema, read_types}; use crate::loader::_helpers::resolve::{ - NameRule, display_name, resolve_tables, source_target, spec_names, target_display_name, - unread_source, + display_name, resolve_tables, source_target, spec_names, target_display_name, unread_source, }; use crate::loader::_helpers::run_log::{ - file_column_problem, logs_files, merged_files, unchanged_files, + file_column_problem, loaded_file, logs_files, merge_ledger, reload_order, }; use crate::loader::_helpers::specs::integer_cursor; -use crate::loader::_helpers::specs::is_child_column; +use crate::loader::_helpers::specs::{files_read_as, path_values, pinned_unsampled}; use crate::loader::_helpers::specs::{ - flattened_object_problems, missing_description, referenced_columns, table_name, table_naming, + flattened_object_problems, referenced_columns, table_name, table_naming, }; -use crate::loader::_helpers::table_plan::{blocked_plan, plan_table}; +use crate::loader::_helpers::specs::{returned_columns, runnable_transform, transform_input}; +use crate::loader::constants::NO_TRANSFORM_HOST; use crate::loader::constants::{ DRAINED_REASON, KEPT_TABLES_COLUMN_FIX, LOCALHOST, PARALLEL_READERS, PRIMARY_KEY_PLACEHOLDER, RELOAD_FIX, SOURCE_BATCH_ROWS, }; use crate::loader::errors::LoaderError; +use crate::loader::models::FilesTransform; use crate::loader::models::{ - CommitRule, DrainedFiles, FilesSource, PlanOptions, ProjectPlan, ResolvedTables, SkippedTable, - SourceEnd, SourceProblem, SourceReader, TablePlan, TableSpec, + CommitRule, DrainedFiles, FilesSource, PlanOptions, ResolvedTables, SkippedTable, SourceEnd, + SourceProblem, SourceReader, TablePlan, TableSpec, }; -use crate::loader::models::{FoundSource, FoundSources}; +use crate::loader::models::{FoundSource, FoundSources, TransformRun}; use crate::postgres_source::constants::DEFAULT_BATCH_BYTES; use crate::postgres_source::constants::SPLIT_READ_MIN_PAGES; use crate::postgres_source::main::connect_postgres::connect_postgres; @@ -74,28 +76,31 @@ use crate::postgres_source::main::cursor_end::cursor_end; use crate::postgres_source::main::read_catalog::read_catalog; use crate::postgres_source::models::{ CopyRequest, Encryption, PgType, PostgresAddress, ReadColumn, SourceColumn, SourceTable, - SslMode, WireType, + SslMode, }; use crate::project::main::postgres_address::postgres_address; use crate::project::models::{ - BatchStep, ConnectionConfig, NewTables, OnDuplicateKeys, ProjectConfig, SourceConfig, - SourceKind, TableOverride, TableSettings, WriteStrategy, + BatchStep, ConnectionConfig, OnDuplicateKeys, ProjectConfig, SourceConfig, SourceKind, + TableOverride, TableSettings, WriteStrategy, }; +use crate::python_sources::main::transform_files::transform_files; +use crate::python_sources::main::transform_range::transform_range; +use crate::python_sources::models::{TransformCall, TransformContext, TransformFiles}; +use crate::python_sources::types::PythonHost; use crate::row_batches::main::loaded_batch::loaded_batch; -use crate::schema_plan::models::{ColumnOrigin, DesiredColumn, LiveColumn}; +use crate::schema_plan::models::LiveColumn; use crate::sheets_source::constants::{EMPTY_SHEET_COLUMN_KIND, SCOPE, SHEET_COLUMN_KIND}; use crate::sheets_source::main::a1_width::a1_width; use crate::sheets_source::main::duplicate_key::duplicate_key; use crate::sheets_source::main::read_grids::read_grids; use crate::sheets_source::main::sheet_batch::sheet_batch; +use crate::sheets_source::main::sheet_input::sheet_input; use crate::sheets_source::main::sheet_table::sheet_table; use crate::sheets_source::main::sheets_client::sheets_client; use crate::sheets_source::main::tab_name::tab_name; use crate::sheets_source::main::tab_ranges::tab_ranges; use crate::sheets_source::main::user_sheets_client::user_sheets_client; use crate::sheets_source::models::{Grid, HeaderRule, KeySource, SheetsClient}; -use crate::similar_names::main::close_name::close_name; -use crate::similar_names::main::nearest_name::nearest_name; use crate::snowflake_source::constants::SESSION_SCHEMA; use crate::snowflake_source::main::connect_snowflake_source::connect_snowflake_source; use crate::snowflake_source::main::read_catalog::read_catalog as read_snowflake_catalog; @@ -228,7 +233,7 @@ pub(crate) async fn open_source( /// A source's tables, its files for a files source, and warnings about reaching it; `Ok(Err(problem))` when they cannot be found, including when a Postgres source cannot be reached or its `include`/`exclude` patterns are invalid. pub(crate) async fn source_tables( - project: &ProjectConfig, + (project, options): (&ProjectConfig, &PlanOptions), source_name: &str, source: &SourceConfig, ) -> Result, Option, Vec), SourceProblem>, LoaderError> @@ -266,8 +271,12 @@ pub(crate) async fn source_tables( .map_or(&unset, |overrides| &overrides.settings); let naming = table_naming((settings, &source.defaults), target.naming); Ok( - match discover_files_table(project, source_name, source, (target.warehouse, naming)) - .await + match discover_files_table( + (project, options), + (source_name, source), + (target.warehouse, naming), + ) + .await { Ok((table, files)) => Ok((vec![table], Some(files), Vec::new())), Err(LoaderError::Files(FilesError::NoFiles { .. })) @@ -565,6 +574,13 @@ async fn sheets_tables( )); continue; } + if let Some(transform) = runnable_transform(project, source, &name) { + match transformed_sheet(transform, (&spreadsheet_id, source), (&name, &grid)).await { + Ok(table) => tables.push(table), + Err(problem) => problems.push(problem), + } + continue; + } let mut read = match sheet_table(&grid, &sheet_rule(source, &name)) { Ok(read) => read, Err(found) => { @@ -632,6 +648,55 @@ async fn sheets_tables( } } +/// Table `name` of a `google_sheets` source with a `transform`: the columns the function returned for the range `grid` read now, typed as a Python table's yields are, or why it failed. +async fn transformed_sheet( + (host, function): (Arc, String), + (spreadsheet_id, source): (&str, &SourceConfig), + (name, grid): (&str, &Grid), +) -> Result { + let input = sheet_input(grid, &sheet_rule(source, name))?; + let call = TransformCall { + function, + context: range_context(name, spreadsheet_id, (grid, source.ranges.get(name))), + masked: Vec::new(), + arrays: Vec::new(), + naming: (Naming::default(), None), + }; + let (batches, _) = transform_range(host, call, input).await?; + let columns: Vec = returned_columns(&batches) + .into_iter() + .map(|(name, column_type)| SourceColumn { + name, + postgres_type: SHEET_COLUMN_KIND.to_owned(), + pg_type: PgType::Other, + column_type: Ok(column_type), + nullable: true, + source_spelling: None, + }) + .collect(); + Ok(SourceTable { + estimated_rows: batches.iter().map(RecordBatch::num_rows).sum::() as i64, + estimated_bytes: 0, + ..SourceTable::of_columns(name.to_owned(), columns) + }) +} + +/// The `ctx` of table `table`'s `transform` for the range `grid` read (`range` as configured, else its whole tab). +fn range_context( + table: &str, + spreadsheet_id: &str, + (grid, range): (&Grid, Option<&String>), +) -> TransformContext { + TransformContext { + unit: "range", + table: table.to_owned(), + sheet: Some(grid.title.clone()), + range: range.cloned(), + spreadsheet_id: Some(spreadsheet_id.to_owned()), + ..TransformContext::default() + } +} + /// A sheet table's spec with each column typed as the live table has it when its cells fit that type (see [`fitting_type`]; dates fit a timestamp column, at midnight), and each column no cell gave a value typed as the live table has it, so a column keeps its type as its values change or empty; a column holding an error cell keeps its error. `None` when no column's type changes. pub(crate) fn with_live_types(spec: &TableSpec, live: Option<&[LiveColumn]>) -> Option { spec.sheet.as_ref()?; @@ -661,12 +726,12 @@ pub(crate) fn with_live_types(spec: &TableSpec, live: Option<&[LiveColumn]>) -> changed.then_some(typed) } -/// A `google_sheets` table's rows as they are now, in its read columns' types. +/// A `google_sheets` table's rows as they are now, in its read columns' types, named as they load; with a `transform`, what it returned for them (its time noted in `run`). async fn sheet_rows( - client: &SheetsClient, - spreadsheet_id: &str, + (client, spreadsheet_id): (&SheetsClient, &str), table: &TablePlan, -) -> Result { + run: Option<&TransformRun>, +) -> Result, LoaderError> { let name = &table.spec.source_table.name; let read = table .spec @@ -692,6 +757,26 @@ async fn sheet_rows( .into_iter() .next() .ok_or_else(|| LoaderError::Invalid(format!("range {range:?} read no tab")))?; + if table.spec.transform.is_some() { + let (mut call, run) = transform_call(&table.spec, run)?; + call.context = range_context(name, spreadsheet_id, (&grid, read.range.as_ref())); + let input = sheet_input(&grid, &read.rule).map_err(LoaderError::Invalid)?; + let (batches, time) = transform_range(run.host.clone(), call, input) + .await + .map_err(LoaderError::Invalid)?; + let mut typed: Vec = Vec::with_capacity(batches.len()); + for batch in &batches { + unplanned_column(batch, &table.spec, &time.file).map_err(LoaderError::Invalid)?; + let loaded = + loaded_batch(batch, &table.spec.source_table).map_err(LoaderError::Invalid)?; + typed.push(read_types(&loaded, &read_schema(table))?); + } + run.times + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(time); + return Ok(typed); + } let columns: Vec<(String, ColumnType)> = table .read_columns .iter() @@ -702,7 +787,10 @@ async fn sheet_rows( ) }) .collect(); - sheet_batch(&grid, &read.rule, &columns).map_err(LoaderError::Invalid) + let batch = sheet_batch(&grid, &read.rule, &columns).map_err(LoaderError::Invalid)?; + Ok(vec![ + loaded_batch(&batch, &table.spec.source_table).map_err(LoaderError::Invalid)?, + ]) } /// A warning when source `source_name` reads `address` over the network in plain text because the server refused TLS under `sslmode = "prefer"`; none for a loopback or Unix socket host, whose traffic stays on the machine, or for an explicit `disable`. @@ -759,11 +847,10 @@ fn snowflake_session( }) } -/// A files source as one source table (file columns, partition columns, the source-file column), and its files. +/// A files source as one source table (file columns, partition columns, the source-file column), and its files: under `--reload` with `after_load.move_to`, those it moved there too. pub(crate) async fn discover_files_table( - project: &ProjectConfig, - source_name: &str, - source: &SourceConfig, + (project, options): (&ProjectConfig, &PlanOptions), + (source_name, source): (&str, &SourceConfig), (warehouse, naming): (Warehouse, Naming), ) -> Result<(SourceTable, FilesSource), LoaderError> { let location = files_location(project, source)?; @@ -773,72 +860,124 @@ pub(crate) async fn discover_files_table( flatten_objects: naming == Naming::Dlt, ..source.files_encoding() }; - let table = discover_files(&folder, &values, &encoding).await?; + let archive = match &source.after_load { + AfterLoad::MoveTo(move_to) if options.reload => Some(OpenFolder { + store: folder.store.clone(), + location: Location { + store: location.store.clone(), + prefix: move_to_prefix(&location, move_to).map_err(LoaderError::Invalid)?, + }, + }), + _ => None, + }; + let table = match &archive { + Some(archive) => discover_reload_files((&folder, archive), &values, &encoding).await?, + None => discover_files(&folder, &values, &encoding).await?, + }; if let Some(problem) = table.problems.first() { return Err(LoaderError::Invalid(problem.clone())); } - files_source_parts(source_name, source, (warehouse, naming), (location, table)) + let overrides = source.tables.get(source_name); + let (table, transform) = match runnable_transform(project, source, source_name) { + Some((host, function)) => { + let own = |pick: fn(&TableOverride) -> &[String]| overrides.map_or(&[][..], pick); + let patterns = ( + [ + &source.exclude_columns[..], + own(|table| &table.exclude_columns), + ] + .concat(), + [&source.allow_columns[..], own(|table| &table.allow_columns)].concat(), + ); + let input = transform_input(&table, encoding.format, patterns)?; + let call = TransformCall { + function, + context: TransformContext { + unit: "file", + table: source_name.to_owned(), + ..TransformContext::default() + }, + masked: Vec::new(), + arrays: Vec::new(), + naming: (naming, sql_adapter(warehouse).max_identifier_length()), + }; + let (table, input) = + transform_sample((&folder, host, call), (input, &encoding), table).await?; + (table, Some(input)) + } + None => (table, None), + }; + let (table, files) = + files_source_parts(source_name, source, (warehouse, naming), (location, table))?; + Ok(( + table, + FilesSource { + transform, + archive: archive.map(|archive| archive.location), + ..files + }, + )) } -/// What files source `source_name` fills from each file's path and its configuration, and the types of those columns in table order: `partition_columns`, `path_columns` captures, then its table's `constant_columns`; a `native(...)` type, which SQLLoad's reader cannot build, is refused. -fn path_values( - (source_name, source): (&str, &SourceConfig), - warehouse: Warehouse, -) -> Result<(PathValues, Vec<(String, ColumnType)>), LoaderError> { - let mut types: Vec<(String, ColumnType)> = Vec::new(); - for (name, text) in &source.partition_columns { - let parsed = parse_type(text, &[warehouse]) - .map_err(|error| LoaderError::Invalid(format!("partition column {name}: {error}")))?; - if matches!(parsed.column_type, ColumnType::Array(_)) { - return Err(LoaderError::Invalid(format!( - "partition column {name}: an array cannot come from a path" - ))); - } - types.push((name.clone(), parsed.column_type)); - } - let template = source - .path_columns - .as_deref() - .map(|text| path_template(text, &[warehouse])) - .transpose() - .map_err(LoaderError::Invalid)?; - types.extend( - template - .iter() - .flat_map(|template| template.captures.clone()), - ); - let constants = constant_columns(&source.table_constants(source_name), &[warehouse]) - .map_err(LoaderError::Invalid)?; - let typed = types - .iter() - .map(|(name, column_type)| (name, column_type)) - .chain( - constants - .iter() - .map(|(name, column_type, _)| (name, column_type)), - ); - for (name, column_type) in typed { - if arrow_field(name, column_type).is_none() { - return Err(LoaderError::Invalid(format!( - "column {name}: {column_type} has no Arrow type, so a partition, path or constant column cannot be one; pin a portable type" - ))); - } - } - let values = PathValues { - partition_columns: source.partition_columns.keys().cloned().collect(), - template, - unmatched: source.unmatched_paths.unwrap_or_default(), - constants: constants - .iter() - .map(|(name, _, text)| (name.clone(), text.clone())) - .collect(), +/// `table` with the columns `call`'s function returned for its newest file (the first sampling reads), typed as a Python table's yields are (a column only ever NULL as text), and `input` with that file as its sample. +async fn transform_sample( + (folder, host, call): (&OpenFolder, Arc, TransformCall), + (input, encoding): (FilesTransform, &FilesEncoding), + table: FilesTable, +) -> Result<(FilesTable, FilesTransform), LoaderError> { + let Some(newest) = table.files.last() else { + return Ok((table, input)); }; - types.extend( - constants - .into_iter() - .map(|(name, column_type, _)| (name, column_type)), - ); - Ok((values, types)) + let request = TransformFiles { + host, + call, + files: std::slice::from_ref(newest), + location: folder.location.to_string(), + encoding, + }; + let mut returned: Vec = Vec::new(); + transform_files( + &request, + async |data| { + Ok::<_, LoaderError>(file_rows(folder, &input_read(&input, encoding, data)).await?) + }, + async |_, batch| { + returned.push(batch); + Ok(()) + }, + ) + .await?; + let columns: Vec = returned_columns(&returned) + .into_iter() + .map(|(name, column_type)| FileColumn { + name, + column_type: Ok(column_type), + direct_refusal: None, + stored_read: StoredRead::Plain, + lists: None, + }) + .collect(); + let sample = format!("{}{}", folder.location, newest.file.relative_key); + Ok(( + FilesTable { + columns, + objects: Vec::new(), + ..table + }, + FilesTransform { sample, ..input }, + )) +} + +/// The read of one file `data` for a `transform`: the file alone, in `input`'s columns. +fn input_read(input: &FilesTransform, encoding: &FilesEncoding, data: &DataFile) -> FilesRead { + FilesRead { + files: vec![data.clone()], + encoding: encoding.clone(), + checks: input.checks.clone(), + columns: input.input.clone(), + nanos: input.nanos.clone(), + batch_rows: SOURCE_BATCH_ROWS, + } } /// A files source's table and files from the files found: the file columns, partition columns and source-file columns as one source table. @@ -944,106 +1083,12 @@ pub(crate) fn files_source_parts( checks, notes, copy_to_internal_stage: source.copy_to_internal_stage, + transform: None, + archive: None, }, )) } -/// Columns a JSON or CSV source's table pins in `columns` that sampling did not see, so a key found only in older files can be pinned; a name close to a sampled one (a typo, or another case) is an error, and a `mask` key alone never adds a column. Under `rule`'s dlt naming a pin naming a sampled column by its loaded name is one sampling saw. -fn pinned_unsampled( - overrides: Option<&TableOverride>, - (table, added, rule): (&FilesTable, &[&str], NameRule), - format: SourceFormat, -) -> Result, LoaderError> { - let Some(overrides) = overrides.filter(|_| format != SourceFormat::Parquet) else { - return Ok(Vec::new()); - }; - let sampled: Vec<&str> = table - .columns - .iter() - .map(|column| column.name.as_str()) - .collect(); - let loaded = |name: &str| loaded_name(rule.0, name, rule.1); - let seen: Vec = sampled - .iter() - .chain(added) - .map(|name| loaded(name)) - .collect(); - let mut unsampled: Vec = Vec::new(); - for name in overrides.columns.keys() { - if seen.contains(&loaded(name)) || is_child_column(name, overrides) { - continue; - } - let known = sampled.iter().chain(added).copied(); - if let Some(close) = close_name(name, known) { - return Err(LoaderError::Invalid(format!( - "pinned column {name} was not seen in the sampled files; did you mean {close:?}? (pin a name exactly as it appears in the files)" - ))); - } - unsampled.push(name.clone()); - } - Ok(unsampled) -} - -/// What a files read reads, as the files spell it (the batches it reads are renamed to the loaded names after): each of `read_columns` in its read type, and those whose types were inferred from sampled values, not pinned and not filled from the path (`partition_types`) or the file. -pub(crate) fn files_read_as( - spec: &TableSpec, - read_columns: &[ReadColumn], - partition_types: &[(String, ColumnType)], -) -> (Vec, Vec) { - let spelled = |name: &str| spec.source_table.spelling(name).to_owned(); - let columns: Vec = read_columns - .iter() - .map(|column| DesiredColumn { - name: spelled(&column.name), - column_type: column.output.clone(), - origin: ColumnOrigin::Source, - }) - .collect(); - let filled: Vec<&str> = partition_types - .iter() - .map(|(name, _)| name.as_str()) - .chain([SOURCE_FILE_COLUMN, SOURCE_VERSION_COLUMN, SOURCE_ROW_COLUMN]) - .collect(); - let inferred: Vec = read_columns - .iter() - .filter(|column| !spec.pinned.contains_key(&column.name)) - .map(|column| spelled(&column.name)) - .filter(|name| !filled.contains(&name.as_str())) - .collect(); - (columns, inferred) -} - -/// The read columns of a files source: the file's types, or the pinned ones (converted while reading). -pub(crate) fn files_read_columns(spec: &TableSpec) -> Result, Vec> { - let mut problems: Vec = Vec::new(); - let mut columns: Vec = Vec::new(); - for column in &spec.source_table.columns { - let excluded = spec - .excluded_columns - .iter() - .any(|(excluded, _)| *excluded == column.name); - let output = match (spec.pinned.get(&column.name), &column.column_type) { - _ if excluded => continue, - (Some(pinned), _) => pinned.clone(), - (None, Ok(column_type)) => column_type.clone(), - (None, Err(reason)) => { - problems.push(format!("column {}: {reason}", column.name)); - continue; - } - }; - columns.push(ReadColumn { - name: column.name.clone(), - expression: column.name.clone(), - wire: WireType::Text, - output, - }); - } - match problems.is_empty() { - true => Ok(columns), - false => Err(problems), - } -} - /// What a files source's load reads: the files in the cursor window that are not loaded yet. pub(crate) struct FilesPlanRequest<'a> { pub(crate) spec: &'a TableSpec, @@ -1065,78 +1110,6 @@ pub(crate) enum LoadedSource<'a> { Recorded(&'a [LoadedFile]), } -/// The files a merge already merged into `request`'s table at the version listed now, which it skips (one rewritten in place since is merged again, as the newest), or why a files source cannot merge into it: only into a DuckDB or Postgres table, never with `--reload` (one commit would lose the files' order) or `state = "logged"`; none while the table does not exist. -async fn merge_ledger( - request: &FilesPlanRequest<'_>, -) -> Result, String>, LoaderError> { - let spec = request.spec; - let destination = match request.loaded { - LoadedSource::Query(destination) => destination, - LoadedSource::Recorded(_) => { - return Ok(Err("merge is not supported for files targets".to_owned())); - } - }; - let what = match spec.strategy { - WriteStrategy::Merge => "merge", - _ => "load", - }; - if !sql_adapter(spec.warehouse).merges_files() { - return Ok(Err(format!( - "{} from a files source loads into DuckDB and Postgres targets so far, not {}; use append, or delete_insert on a partition cursor", - spec.strategy, spec.warehouse - ))); - } - if request.options.reload { - return Ok(Err(format!( - "--reload would {what} every file in one commit and lose their order; drop the table and load again to {what} every file anew" - ))); - } - if spec.logged { - return Ok(Err(format!( - "state = \"logged\" does not apply to a files source's {}, which records each file it loads", - spec.strategy - ))); - } - if let Some(cursor) = &spec.cursor { - return Ok(Err(format!( - "cursor_column {} does not apply to a files source's merge, which picks the files it has not merged from the run log, whatever their partition; remove it and set write_strategy = \"merge\", or keep it and use delete_insert", - cursor.column - ))); - } - let merged = match request.live { - None => return Ok(Ok(Vec::new())), - Some(_) => merged_files(spec, destination).await?, - }; - let location = request.files.location.to_string(); - let listed: HashMap = request - .files - .table - .files - .iter() - .map(|data| { - ( - format!("{location}{}", data.file.relative_key), - data.file.version(), - ) - }) - .collect(); - Ok(merged.map(|merged| unchanged_files(merged, &listed))) -} - -/// A row of the loaded-files query: `file|version`, or just the file when `versioned` is false. -pub(crate) fn loaded_file(text: String, versioned: bool) -> LoadedFile { - match (versioned, text.rsplit_once(VERSION_SEPARATOR)) { - (true, Some((name, version))) => LoadedFile { - name: name.to_owned(), - version: version.to_owned(), - }, - _ => LoadedFile { - name: text, - version: String::new(), - }, - } -} - /// Whether the live table has `column`. pub(crate) fn live_has(live: Option<&[LiveColumn]>, column: &str) -> bool { match live { @@ -1147,69 +1120,6 @@ pub(crate) fn live_has(live: Option<&[LiveColumn]>, column: &str) -> bool { } } -/// The columns of `columns` SQLLoad's reader keeps at nanoseconds for `spec`'s warehouse: those the files store at nanoseconds, unless pinned or `live` holds fewer digits. -pub(crate) fn nanosecond_columns( - spec: &TableSpec, - columns: &[FileColumn], - live: Option<&[LiveColumn]>, -) -> Vec { - match sql_adapter(spec.warehouse).keeps_nanoseconds() { - true => columns - .iter() - .filter(|column| { - matches!( - column.stored_read, - StoredRead::Nanos | StoredRead::Epoch(NANOSECOND_SCALE) - ) - }) - .filter(|column| { - let loaded = spec.source_table.loaded_as(&column.name); - !spec.pinned.contains_key(loaded) && fewer_digits(live, loaded).is_none() - }) - .map(|column| column.name.clone()) - .collect(), - false => Vec::new(), - } -} - -/// A warning for each column the files store at nanoseconds but the live table holds at fewer digits, so the load cuts them; a column whose type is pinned is the user's choice and is not warned about. -pub(crate) fn truncated_nanoseconds(spec: &TableSpec, live: Option<&[LiveColumn]>) -> Vec { - let Some(files) = &spec.files else { - return Vec::new(); - }; - files - .table - .columns - .iter() - .filter(|column| { - matches!( - column.stored_read, - StoredRead::Nanos | StoredRead::Epoch(NANOSECOND_SCALE) - ) && !spec - .pinned - .contains_key(spec.source_table.loaded_as(&column.name)) - }) - .filter_map(|column| truncated_warning(live, spec.source_table.loaded_as(&column.name))) - .collect() -} - -/// The warning for column `name` when the live column holds fewer than nanosecond digits. -fn truncated_warning(live: Option<&[LiveColumn]>, name: &str) -> Option { - let digits = fewer_digits(live, name)?; - Some(format!( - "{name}: nanoseconds in the files, but the column holds {digits} digits after the second, so they are truncated.\nRebuild the column at (9) (--reload) to keep them." - )) -} - -/// The digits after the second the live column `name` holds, when fewer than nanoseconds, so SQLLoad's reader cuts its values. -fn fewer_digits(live: Option<&[LiveColumn]>, name: &str) -> Option { - live.into_iter() - .flatten() - .find(|existing| existing.name.eq_ignore_ascii_case(name)) - .and_then(|existing| existing.fraction_digits) - .filter(|digits| *digits < i64::from(NANOSECOND_SCALE)) -} - /// How to give a table with rows a column a load needs, on `warehouse`. pub(crate) fn column_fix(warehouse: Warehouse) -> &'static str { match sql_adapter(warehouse).keeps_tables() { @@ -1263,7 +1173,7 @@ pub(crate) async fn plan_files( { return Ok(Err(problem)); } - if spec.after_load != AfterLoad::Keep && request.options.reload { + if spec.after_load == AfterLoad::Delete && request.options.reload { return Ok(Err( "--reload would rebuild the table from the files still in its folder, and after_load has deleted or moved the ones already loaded".to_owned(), )); @@ -1310,7 +1220,9 @@ pub(crate) async fn plan_files( true => (recorded, Vec::new()), false => (Vec::new(), recorded), }; - let loaded = merged.clone().unwrap_or(loaded); + let loaded = merged + .as_ref() + .map_or(loaded, |(unchanged, _)| unchanged.clone()); let location = request.files.location.to_string(); let selection = FileSelection { files: &request.files.table.files, @@ -1329,7 +1241,10 @@ pub(crate) async fn plan_files( ..request.files.checks.clone() }; let nanos = nanosecond_columns(spec, &request.files.table.columns, request.live); - let ordered = select_files(&selection).map(|files| match merged { + let ordered = select_files(&selection).map(|files| match &merged { + Some((_, recorded)) if request.options.reload => { + reload_order(files, &location, recorded, request.files.archive.as_ref()) + } Some(_) => merge_order(files), None => files, }); @@ -1343,6 +1258,44 @@ pub(crate) async fn plan_files( })) } +/// The stored `files` that `after_load` settles once their rows commit: under a `--reload` that lists the `move_to` folder too, those read from it stay where they are. +pub(crate) fn unarchived(table: &TablePlan, files: &[DataFile]) -> Vec { + let archive = table + .spec + .files + .as_ref() + .and_then(|files| files.archive.as_ref()); + files + .iter() + .filter(|data| { + archive.is_none_or(|archive| !data.file.path.as_ref().starts_with(&archive.prefix)) + }) + .map(|data| data.file.clone()) + .collect() +} + +/// One file's rows as `read` (of that file alone) reads them, whole, and their schema: what a `transform` gets. +pub(crate) async fn file_rows( + folder: &OpenFolder, + read: &FilesRead, +) -> Result<(Vec, SchemaRef), FilesError> { + let mut rows: Vec = Vec::new(); + read_files(folder, read, &BatchFooters::default(), async |batch| { + rows.push(batch); + Ok(()) + }) + .await?; + let fields = read.columns.iter().map(|column| { + arrow_field(&column.name, &column.column_type) + .unwrap_or_else(|| Field::new(&column.name, DataType::Utf8, true)) + }); + let schema = rows.first().map_or_else( + || Arc::new(Schema::new(fields.collect::>())), + RecordBatch::schema, + ); + Ok((rows, schema)) +} + /// Which part of a table's planned read one pass reads: a cursor window's filter, an order, or some of a files source's files. #[derive(Debug, Clone, Default)] pub(crate) struct ReadPart { @@ -1362,6 +1315,8 @@ pub(crate) struct ReadPart { pub(crate) footers: BatchFooters, /// A smaller byte cap for the source's batches than the usual 64 MiB (the hidden `--batch-bytes`, for tests). pub(crate) batch_bytes: Option, + /// What runs the table's `transform`; `None` without a Python host. + pub(crate) transform: Option, } /// Read `part` of `table`'s rows from `reader`, handing each batch to `on_batch`; when `on_batch` fails, the read fails with the reader's text, and stays retryable when `on_batch`'s error was. @@ -1423,6 +1378,15 @@ async fn read_each( None => files.clone(), }; let mut on_batch = on_batch; + if let Some(transform) = table + .spec + .files + .as_ref() + .and_then(|files| files.transform.as_ref()) + { + let run = part.transform.as_ref(); + return read_transformed(folder, (&read, transform), (table, run), on_batch).await; + } let loaded = async |batch: RecordBatch| { on_batch(loaded_batch(&batch, &table.spec.source_table)?).await }; @@ -1468,14 +1432,14 @@ async fn read_each( Ok(read_bigquery_rows(reader, &request, on_batch).await?) } (SourceReader::Sheets(client, spreadsheet_id), _) => { - let batch = loaded_batch( - &sheet_rows(client, spreadsheet_id, table).await?, - &table.spec.source_table, - ) - .map_err(LoaderError::Invalid)?; - let rows = batch.num_rows() as u64; + let batches = + sheet_rows((client, spreadsheet_id), table, part.transform.as_ref()).await?; + let mut rows: u64 = 0; let mut on_batch = on_batch; - on_batch(batch).await.map_err(LoaderError::Invalid)?; + for batch in batches { + rows += batch.num_rows() as u64; + on_batch(batch).await.map_err(LoaderError::Invalid)?; + } Ok(rows) } (SourceReader::Files(_), None) => Err(LoaderError::Invalid(format!( @@ -1485,6 +1449,105 @@ async fn read_each( } } +/// Read each of `read`'s files in order through `table`'s `transform`, passing what it returned for the file, shaped as the file's rows, to `on_batch`; a file whose function fails, or returns a column the plan did not see, fails the read there, naming the file. Returns the row count. +async fn read_transformed( + folder: &OpenFolder, + (read, input): (&FilesRead, &FilesTransform), + (table, run): (&TablePlan, Option<&TransformRun>), + mut on_batch: impl AsyncFnMut(RecordBatch) -> Result<(), String>, +) -> Result { + let spec = &table.spec; + let (call, run) = transform_call(spec, run)?; + let request = TransformFiles { + host: run.host.clone(), + call, + files: &read.files, + location: folder.location.to_string(), + encoding: &read.encoding, + }; + let mut rows: u64 = 0; + let times = transform_files( + &request, + async |data| { + Ok::<_, LoaderError>(file_rows(folder, &input_read(input, &read.encoding, data)).await?) + }, + async |(data, file_name, before), batch| { + let failed = |reason: String| LoaderError::Invalid(format!("{file_name}: {reason}")); + unplanned_column(&batch, spec, &input.sample).map_err(failed)?; + let shaped = shape_rows(&batch, (data, file_name, before), read).map_err(failed)?; + rows += shaped.num_rows() as u64; + on_batch(loaded_batch(&shaped, &spec.source_table).map_err(failed)?) + .await + .map_err(LoaderError::Invalid) + }, + ) + .await?; + run.times + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .extend(times); + Ok(rows) +} + +/// The call of `spec`'s `transform`, before each file or range sets its context, and what runs it; why it cannot run without a Python host. +fn transform_call<'a>( + spec: &TableSpec, + run: Option<&'a TransformRun>, +) -> Result<(TransformCall, &'a TransformRun), LoaderError> { + let (Some(run), Some(transform)) = (run, spec.transform.as_ref()) else { + return Err(LoaderError::Invalid(NO_TRANSFORM_HOST.replace( + "{place}", + &format!("table {}", spec.source_table.qualified_name()), + ))); + }; + let call = TransformCall { + function: transform.function.clone(), + context: TransformContext { + unit: "file", + table: spec.source_table.name.clone(), + ..TransformContext::default() + }, + masked: spec.masked.columns.keys().cloned().collect(), + arrays: spec + .pinned + .iter() + .filter(|(_, column_type)| matches!(column_type, ColumnType::Array(_))) + .map(|(name, _)| name.clone()) + .collect(), + naming: ( + spec.naming, + sql_adapter(spec.warehouse).max_identifier_length(), + ), + }; + Ok((call, run)) +} + +/// Why `batch`, returned by `spec`'s `transform`, cannot load: it has a column the function did not return for `sample` at plan (nor `columns` pins). +fn unplanned_column(batch: &RecordBatch, spec: &TableSpec, sample: &str) -> Result<(), String> { + let planned: Vec<&str> = spec + .source_table + .columns + .iter() + .map(|column| spec.source_table.spelling(&column.name)) + .collect(); + match batch + .schema() + .fields() + .iter() + .find(|field| !planned.contains(&field.name().as_str())) + { + Some(extra) => Err(format!( + "transform {} returned column {}, which it did not return for {sample} at plan; return the same columns every time, or pin {} in columns", + spec.transform + .as_ref() + .map_or("", |transform| transform.function.as_str()), + extra.name(), + extra.name() + )), + None => Ok(()), + } +} + /// Unload `part` of `table`'s data from `reader` and pass the files to `load`, which has the destination read them itself; the files are removed afterwards. pub(crate) async fn unload_source( reader: &SourceReader, @@ -1551,31 +1614,6 @@ pub(crate) async fn source_ends( Ok(specs) } -/// A warning for each selector matching no source or table, with the nearest name: `load` would otherwise say only `0 selected`. -pub(crate) fn unmatched_selectors( - project: &ProjectConfig, - options: &PlanOptions, - names: &[String], -) -> Vec { - let known: Vec<&str> = names - .iter() - .map(String::as_str) - .chain(project.sources.keys().map(String::as_str)) - .collect(); - options - .select - .iter() - .chain(&options.exclude) - .filter(|selector| !known.contains(&selector.as_str())) - .map( - |selector| match nearest_name(selector, known.iter().copied()) { - Some(nearest) => format!("no table matches {selector}; did you mean {nearest}?"), - None => format!("no table matches {selector}"), - }, - ) - .collect() -} - /// Whether `--select` picks a table and `--exclude` does not, each by source name or one of the table's `names` (its destination name, the destination's display name, or its source's `schema.table`; empty ones match nothing); no `--select` picks every table. pub(crate) fn selected(options: &PlanOptions, (source, names): (&str, &[String])) -> bool { let picks = |selector: &String| { @@ -1593,16 +1631,17 @@ pub(crate) async fn scoped_tables( (project, options): (&ProjectConfig, &PlanOptions), (source_name, source): (&str, &SourceConfig), ) -> Result), Option>, LoaderError> { - let (discovered, files, warnings) = match source_tables(project, source_name, source).await? { - Ok(found) => found, - Err(problem) => { - let in_scope = source_in_scope((project, options), (source_name, source), None); - return Ok(Err(in_scope.then(|| SourceProblem { - message: format!("[sources.{source_name}]: {}", problem.message), - retryable: problem.retryable, - }))); - } - }; + let (discovered, files, warnings) = + match source_tables((project, options), source_name, source).await? { + Ok(found) => found, + Err(problem) => { + let in_scope = source_in_scope((project, options), (source_name, source), None); + return Ok(Err(in_scope.then(|| SourceProblem { + message: format!("[sources.{source_name}]: {}", problem.message), + retryable: problem.retryable, + }))); + } + }; if source.kind == SourceKind::Files && discovered.is_empty() { return Ok(Ok(( drained_files((project, options), (source_name, source)), @@ -1773,108 +1812,6 @@ pub(crate) fn source_in_scope( .any(|names| selected(options, (source_name, names))) } -/// `planned`, the selected tables of one source, added to `plan`, which is returned (`unselected` and `found`, the destinations and options, tell whether this is the source's first load): a table that appeared after the source's first load (`first_load`: none of its tables, selected or not, existed yet) and is not named exactly in `include` is skipped or fails as `new_tables` says, and a table without a description is blocked. -pub(crate) async fn admit_tables( - mut plan: ProjectPlan, - (source_name, source): (&String, &SourceConfig), - (planned, unselected): (Vec, &[TableSpec]), - found: (&BTreeMap, &PlanOptions), -) -> Result { - let first_load = first_load(&planned, unselected, found).await?; - let mut planned = planned; - if let Some(problem) = partial_capture(source_name, source.kind, (&planned, unselected)) { - for table in &mut planned { - table.schema = blocked_plan(vec![problem.clone()]); - } - } - for mut table in planned { - let spec = &table.spec; - let qualified = spec.source_table.qualified_name(); - let is_new = !first_load && !table.exists && !spec.approved; - let has_settings = source.tables.contains_key(&qualified); - match (is_new, spec.new_tables, has_settings) { - (true, NewTables::Skip, true) => plan.skipped.push(SkippedTable { - source_name: source_name.clone(), - source_table: qualified, - destination: display_name(spec.warehouse, &spec.destination), - reason: "new table (new_tables = \"skip\") has settings but is not approved; add it to include".to_owned(), - }), - (true, NewTables::Skip, false) => plan.skipped.push(SkippedTable { - source_name: source_name.clone(), - source_table: qualified, - destination: display_name(spec.warehouse, &spec.destination), - reason: "new table (new_tables = \"skip\"); add it to include to sync it".to_owned(), - }), - (true, NewTables::Fail, _) => { - plan.failed_sources.insert(source_name.clone()); - let problem = format!( - "{qualified}: new table and new_tables = \"fail\"; add it to include or exclude it" - ); - plan.problem_names.push(( - problem.clone(), - display_name(spec.warehouse, &spec.destination), - )); - plan.problems.push(problem); - } - _ => { - if let Some(problem) = missing_description(source_name, source, &spec.source_table) - { - table.schema = blocked_plan(vec![problem]); - } - plan.tables.push(table); - } - } - } - Ok(plan) -} - -/// Why the selection cannot load `postgres_cdc` source `name`, which blocks every table it picks: it picks some of its tables (`planned`) but not all (`unselected`), and the slot they share would confirm the left-out tables' changes unwritten. A plan problem alone would not stop the picked tables loading. -fn partial_capture( - name: &str, - kind: SourceKind, - (planned, unselected): (&[TablePlan], &[TableSpec]), -) -> Option { - let partial = kind == SourceKind::PostgresCdc && !planned.is_empty() && !unselected.is_empty(); - partial.then(|| { - let left_out: Vec = unselected - .iter() - .map(|spec| spec.source_table.qualified_name()) - .collect(); - format!( - "this selection leaves out {} of source {name}, whose changes share its one slot and would be skipped; select the whole source, or move those tables to another postgres_cdc source", - left_out.join(", ") - ) - }) -} - -/// Whether this is the source's first load: none of its tables exists in the destination, whether selected (`planned`) or not (`unselected`, looked up only while none is found, a table that cannot be read counting as absent), so `--select` cannot make a later table look present at the first load. -async fn first_load( - planned: &[TablePlan], - unselected: &[TableSpec], - (destinations, options): (&BTreeMap, &PlanOptions), -) -> Result { - if planned.iter().any(|table| table.exists) { - return Ok(false); - } - for spec in unselected { - let Some(destination) = destinations.get(&spec.target) else { - continue; - }; - let exists = match files_destination(destination) { - Some(_) => plan_table(spec, destination, options) - .await - .is_ok_and(|table| table.exists), - None => live_columns(destination, &spec.destination) - .await - .is_ok_and(|live| live.is_some()), - }; - if exists { - return Ok(false); - } - } - Ok(true) -} - /// Every source's tables, resolved before any is planned so each knows which others share its destination table, with the folders they claim; a source whose tables cannot be read is a problem in the plan. pub(crate) async fn found_sources<'a>( project: &'a ProjectConfig, @@ -1947,11 +1884,11 @@ pub(crate) async fn found_sources<'a>( "[sources.{}]: no files to load in {path}, and other tables load into {} too: a load that cannot see this source could misread or remove its rows, so it fails until files come; run the others with --exclude {}", source.name, drained.table.table, source.name )), - (true, false, true) => Some(format!( + (true, false, true) if source.source.after_load == AfterLoad::Delete => Some(format!( "[sources.{}]: --reload would rebuild the table from the files still in its folder, and after_load has deleted or moved the ones already loaded", source.name )), - (true, false, false) => None, + (true, false, _) => None, }; match problem { Some(problem) => source.resolved.problems.push(problem), diff --git a/crates/sqlload/src/loader/_helpers/specs.rs b/crates/sqlload/src/loader/_helpers/specs.rs index 1d4b7b592..4fc818e32 100644 --- a/crates/sqlload/src/loader/_helpers/specs.rs +++ b/crates/sqlload/src/loader/_helpers/specs.rs @@ -1,3 +1,8 @@ +use std::sync::Arc; + +use arrow_array::RecordBatch; +use arrow_schema::{DataType, Schema}; + use crate::adapters::contract::models::{ChangeApply, TableName, WriteMode}; use crate::adapters::dispatch::main::dialect::sql_adapter::sql_adapter; use crate::adapters::files::constants::{CURSOR_PLACEHOLDER_START, TABLE_PLACEHOLDER}; @@ -9,11 +14,18 @@ use crate::cdc::constants::{ use crate::column_types::main::arrow_field::arrow_field; use crate::column_types::main::parse_type::parse_type; use crate::column_types::models::{ColumnType, Warehouse}; +use crate::files_source::constants::{ + SOURCE_FILE_COLUMN, SOURCE_ROW_COLUMN, SOURCE_VERSION_COLUMN, +}; use crate::files_source::main::constant_columns::constant_columns; -use crate::files_source::models::AfterLoad; +use crate::files_source::main::path_template::path_template; use crate::files_source::models::FileColumn; +use crate::files_source::models::{ + AfterLoad, FilesTable, PathValues, SourceFormat, StoredRead, TextChecks, +}; use crate::load_naming::main::loaded_name::loaded_name; use crate::load_naming::models::Naming; +use crate::loader::_helpers::python::field_type; use crate::loader::_helpers::resolve::{NameRule, ResolvedTarget}; use crate::loader::_helpers::sources::sheet_rule; use crate::loader::_helpers::table_plan::source_cursor; @@ -23,22 +35,28 @@ use crate::loader::constants::{ }; use crate::loader::errors::LoaderError; use crate::loader::models::{ - CdcSpec, CursorSpec, SettingOrigin, TableLayout, TableSpec, WarehouseRoute, WarehouseSource, + CdcSpec, CursorSpec, FilesTransform, SettingOrigin, TableLayout, TableSpec, WarehouseRoute, + WarehouseSource, }; use crate::loader::models::{SourceReader, TablePlan}; use crate::masking::models::MaskRules; use crate::postgres_source::main::count_rows::count_rows as count_postgres_rows; use crate::postgres_source::models::{PgType, SourceColumn, SourceTable}; +use crate::postgres_source::models::{ReadColumn, WireType}; use crate::project::models::{ CdcDeletes, CdcMode, CursorBounds, LookbackFrom, MaskMethod, NullCursor, NullKeys, - OnDuplicateKeys, SortDirection, SourceConfig, SourceKind, StateMode, TableOverride, - TableSettings, UnloadStage, UnnestSetting, WriteStrategy, + OnDuplicateKeys, ProjectConfig, SortDirection, SourceConfig, SourceKind, StateMode, + TableOverride, TableSettings, UnloadStage, UnnestSetting, WriteStrategy, }; use crate::publish_checks::main::check_problems::check_problems; +use crate::python_sources::types::PythonHost; use crate::routing::models::Store; use crate::row_batches::constants::FLATTEN_SEPARATOR; +use crate::row_batches::main::merge_schema::merge_schema; use crate::row_batches::models::Unnest; +use crate::schema_plan::models::{ColumnOrigin, DesiredColumn}; use crate::sheets_source::models::SheetRead; +use crate::similar_names::main::close_name::close_name; use crate::similar_names::main::nearest_name::nearest_name; use crate::snowflake_source::constants::USER_STAGE_FOLDER; use crate::snowflake_source::main::count_rows::count_rows as count_snowflake_rows; @@ -438,6 +456,7 @@ pub(crate) fn build_spec( _ => warehouse_source(context.source, overrides), }, layout, + transform: None, }) } @@ -647,7 +666,7 @@ fn glob(pattern: &str) -> Result { } /// A matcher for column name patterns such as `*email*`. -fn column_matcher(patterns: &[String]) -> Result { +pub(crate) fn column_matcher(patterns: &[String]) -> Result { let mut builder = GlobSetBuilder::new(); for pattern in patterns { builder.add(glob(pattern)?); @@ -1630,3 +1649,242 @@ pub(crate) fn checked_key<'a>( _ => mode.staged_key(), } } + +/// What files source `source_name` fills from each file's path and its configuration, and the types of those columns in table order: `partition_columns`, `path_columns` captures, then its table's `constant_columns`; a `native(...)` type, which SQLLoad's reader cannot build, is refused. +pub(crate) fn path_values( + (source_name, source): (&str, &SourceConfig), + warehouse: Warehouse, +) -> Result<(PathValues, Vec<(String, ColumnType)>), LoaderError> { + let mut types: Vec<(String, ColumnType)> = Vec::new(); + for (name, text) in &source.partition_columns { + let parsed = parse_type(text, &[warehouse]) + .map_err(|error| LoaderError::Invalid(format!("partition column {name}: {error}")))?; + if matches!(parsed.column_type, ColumnType::Array(_)) { + return Err(LoaderError::Invalid(format!( + "partition column {name}: an array cannot come from a path" + ))); + } + types.push((name.clone(), parsed.column_type)); + } + let template = source + .path_columns + .as_deref() + .map(|text| path_template(text, &[warehouse])) + .transpose() + .map_err(LoaderError::Invalid)?; + types.extend( + template + .iter() + .flat_map(|template| template.captures.clone()), + ); + let constants = constant_columns(&source.table_constants(source_name), &[warehouse]) + .map_err(LoaderError::Invalid)?; + let typed = types + .iter() + .map(|(name, column_type)| (name, column_type)) + .chain( + constants + .iter() + .map(|(name, column_type, _)| (name, column_type)), + ); + for (name, column_type) in typed { + if arrow_field(name, column_type).is_none() { + return Err(LoaderError::Invalid(format!( + "column {name}: {column_type} has no Arrow type, so a partition, path or constant column cannot be one; pin a portable type" + ))); + } + } + let values = PathValues { + partition_columns: source.partition_columns.keys().cloned().collect(), + template, + unmatched: source.unmatched_paths.unwrap_or_default(), + constants: constants + .iter() + .map(|(name, _, text)| (name.clone(), text.clone())) + .collect(), + }; + types.extend( + constants + .into_iter() + .map(|(name, column_type, _)| (name, column_type)), + ); + Ok((values, types)) +} + +/// Columns a JSON or CSV source's table pins in `columns` that sampling did not see, so a key found only in older files can be pinned; a name close to a sampled one (a typo, or another case) is an error, and a `mask` key alone never adds a column. Under `rule`'s dlt naming a pin naming a sampled column by its loaded name is one sampling saw. +pub(crate) fn pinned_unsampled( + overrides: Option<&TableOverride>, + (table, added, rule): (&FilesTable, &[&str], NameRule), + format: SourceFormat, +) -> Result, LoaderError> { + let Some(overrides) = overrides.filter(|_| format != SourceFormat::Parquet) else { + return Ok(Vec::new()); + }; + let sampled: Vec<&str> = table + .columns + .iter() + .map(|column| column.name.as_str()) + .collect(); + let loaded = |name: &str| loaded_name(rule.0, name, rule.1); + let seen: Vec = sampled + .iter() + .chain(added) + .map(|name| loaded(name)) + .collect(); + let mut unsampled: Vec = Vec::new(); + for name in overrides.columns.keys() { + if seen.contains(&loaded(name)) || is_child_column(name, overrides) { + continue; + } + let known = sampled.iter().chain(added).copied(); + if let Some(close) = close_name(name, known) { + return Err(LoaderError::Invalid(format!( + "pinned column {name} was not seen in the sampled files; did you mean {close:?}? (pin a name exactly as it appears in the files)" + ))); + } + unsampled.push(name.clone()); + } + Ok(unsampled) +} + +/// What a files read reads, as the files spell it (the batches it reads are renamed to the loaded names after): each of `read_columns` in its read type, and those whose types were inferred from sampled values, not pinned and not filled from the path (`partition_types`) or the file. +pub(crate) fn files_read_as( + spec: &TableSpec, + read_columns: &[ReadColumn], + partition_types: &[(String, ColumnType)], +) -> (Vec, Vec) { + let spelled = |name: &str| spec.source_table.spelling(name).to_owned(); + let columns: Vec = read_columns + .iter() + .map(|column| DesiredColumn { + name: spelled(&column.name), + column_type: column.output.clone(), + origin: ColumnOrigin::Source, + }) + .collect(); + let filled: Vec<&str> = partition_types + .iter() + .map(|(name, _)| name.as_str()) + .chain([SOURCE_FILE_COLUMN, SOURCE_VERSION_COLUMN, SOURCE_ROW_COLUMN]) + .collect(); + let inferred: Vec = read_columns + .iter() + .filter(|column| !spec.pinned.contains_key(&column.name)) + .map(|column| spelled(&column.name)) + .filter(|name| !filled.contains(&name.as_str())) + .collect(); + (columns, inferred) +} + +/// The read columns of a files source: the file's types, or the pinned ones (converted while reading). +pub(crate) fn files_read_columns(spec: &TableSpec) -> Result, Vec> { + let mut problems: Vec = Vec::new(); + let mut columns: Vec = Vec::new(); + for column in &spec.source_table.columns { + let excluded = spec + .excluded_columns + .iter() + .any(|(excluded, _)| *excluded == column.name); + let output = match (spec.pinned.get(&column.name), &column.column_type) { + _ if excluded => continue, + (Some(pinned), _) => pinned.clone(), + (None, Ok(column_type)) => column_type.clone(), + (None, Err(reason)) => { + problems.push(format!("column {}: {reason}", column.name)); + continue; + } + }; + columns.push(ReadColumn { + name: column.name.clone(), + expression: column.name.clone(), + wire: WireType::Text, + output, + }); + } + match problems.is_empty() { + true => Ok(columns), + false => Err(problems), + } +} + +/// The function and the host that run table `name`'s `transform` (its own, else its source's), when it has one the host imported. +pub(crate) fn runnable_transform( + project: &ProjectConfig, + source: &SourceConfig, + name: &str, +) -> Option<(Arc, String)> { + let function = source + .tables + .get(name) + .and_then(|table| table.settings.transform.clone()) + .or_else(|| source.defaults.transform.clone())?; + let host = project.python_host.clone()?; + matches!(project.transforms.get(&function), Some(Ok(_))).then_some((host, function)) +} + +/// The columns of what a `transform` returned across its batches, typed as a Python table's yields are: a column only ever NULL as text. +pub(crate) fn returned_columns(batches: &[RecordBatch]) -> Vec<(String, ColumnType)> { + let merged = batches.iter().fold(Schema::empty(), |seen, batch| { + merge_schema(&seen, &batch.schema()) + }); + merged + .fields() + .iter() + .map(|field| { + let column_type = match field.data_type() { + DataType::Null => ColumnType::Text, + _ => field_type(field), + }; + (field.name().clone(), column_type) + }) + .collect() +} + +/// How each file of a files table with a `transform` is read for it: `table`'s sampled columns, less those `exclude_columns` drops (`patterns` holds it and `allow_columns`), as stored (Parquet) or as the workbook holds them (xlsx), JSON's nested values as JSON and every other value as text, so a value sampling did not see never fails the read. +pub(crate) fn transform_input( + table: &FilesTable, + format: SourceFormat, + (exclude, allow): (Vec, Vec), +) -> Result { + let excluded = column_matcher(&exclude).map_err(LoaderError::Invalid)?; + let allowed = column_matcher(&allow).map_err(LoaderError::Invalid)?; + let mut input: Vec = Vec::new(); + for column in &table.columns { + if excluded.is_match(&column.name) && !allowed.is_match(&column.name) { + continue; + } + let sampled = column + .column_type + .clone() + .map_err(|reason| LoaderError::Invalid(format!("column {}: {reason}", column.name))); + let column_type = match (format, sampled) { + (SourceFormat::Parquet | SourceFormat::Xlsx, sampled) => sampled?, + ( + SourceFormat::Json | SourceFormat::Ndjson, + Ok(nested @ (ColumnType::Json | ColumnType::Array(_))), + ) => nested, + _ => ColumnType::Text, + }; + input.push(DesiredColumn { + name: column.name.clone(), + column_type, + origin: ColumnOrigin::Source, + }); + } + Ok(FilesTransform { + checks: TextChecks { + known: input.iter().map(|column| column.name.clone()).collect(), + exclude, + allow, + inferred: Vec::new(), + }, + input, + nanos: table + .columns + .iter() + .filter(|column| column.stored_read == StoredRead::Nanos) + .map(|column| column.name.clone()) + .collect(), + sample: String::new(), + }) +} diff --git a/crates/sqlload/src/loader/_helpers/table_plan.rs b/crates/sqlload/src/loader/_helpers/table_plan.rs index db930311e..0bb2130bf 100644 --- a/crates/sqlload/src/loader/_helpers/table_plan.rs +++ b/crates/sqlload/src/loader/_helpers/table_plan.rs @@ -41,17 +41,18 @@ use crate::files_source::constants::SOURCE_ROW_COLUMN; use crate::files_source::models::FilesRead; use crate::files_source::models::{Compression, SourceFormat}; use crate::load_naming::main::is_load_column::is_load_column; +use crate::loader::_helpers::destination::truncated_nanoseconds; use crate::loader::_helpers::python::list_choices; use crate::loader::_helpers::python::live_child_tables; use crate::loader::_helpers::python::with_live_children; use crate::loader::_helpers::python::with_live_columns; use crate::loader::_helpers::render::up_to_date; use crate::loader::_helpers::resolve::display_name; +use crate::loader::_helpers::run_log::with_transform_version; use crate::loader::_helpers::sources::column_fix; -use crate::loader::_helpers::sources::files_read_columns; use crate::loader::_helpers::sources::plan_files; -use crate::loader::_helpers::sources::truncated_nanoseconds; use crate::loader::_helpers::sources::{FilesPlanRequest, LoadedSource, with_live_types}; +use crate::loader::_helpers::specs::files_read_columns; use crate::loader::constants::CHECKS_LOADER_WARNING; use crate::loader::constants::ENFORCED_UNDECLARED_PROBLEM; use crate::loader::constants::KEPT_TABLES_RELOAD_PROBLEM; @@ -253,7 +254,8 @@ pub(crate) async fn plan_table( plan.schema.warnings.extend(ignored_batch(&plan)); plan.schema.warnings.extend(reload_range(spec, options)); plan.schema.warnings.extend(split_ranking(&plan)); - with_run_log(plan, destination, options).await + let plan = with_run_log(plan, destination, options).await?; + with_transform_version(plan, destination).await } /// Why the plan cannot write `spec` as `schema` plans it, when nothing blocked it first: its `cursor_column`, when `unwritten_cursor` names it, or its `dedup_sort` column is not among the columns written. A Python table's columns are known only once it yields, so its `dedup_sort` is not checked here. @@ -1019,8 +1021,18 @@ pub(crate) fn schema_statements(table: &TablePlan) -> Option> { } } +/// Why SQLLoad reads a `--reload`'s files itself when some come from `after_load`'s `move_to` folder: a destination reading the files addresses each by the source's folder and its key, which would read `done/a.parquet` as `drop/a.parquet`. +pub(crate) fn archived_reason(spec: &TableSpec, read: Option<&FilesRead>) -> Option { + let archive = spec.files.as_ref()?.archive.as_ref()?; + read? + .files + .iter() + .any(|data| data.file.path.as_ref().starts_with(&archive.prefix)) + .then(|| format!("--reload also reads the files after_load moved to {archive}")) +} + /// What a table's source offers: rows as Arrow, plus the files themselves for a files source, with why the destination must not read them itself. -fn reader_offer( +pub(crate) fn reader_offer( spec: &TableSpec, read: Option<&FilesRead>, stage: Option<&StageRoute>, @@ -1081,7 +1093,13 @@ fn reader_offer( _ => None, } }), - }; + } + .or_else(|| { + spec.transform + .as_ref() + .map(|transform| format!("transform {} runs on each file", transform.function)) + }) + .or_else(|| archived_reason(spec, read)); let unload_to: Vec = match &spec.warehouse_source { Some(WarehouseSource { route: WarehouseRoute::Unload { store, .. }, @@ -1249,13 +1267,15 @@ pub(crate) struct Window { pub(crate) maximum: Option, } +/// The write mode of `spec`'s table: `--reload` rebuilds it, except that a files source's merge keeps merging, file by file, into the table it swaps in at the end. fn write_mode( spec: &TableSpec, reload: bool, from: Option, to: Option, ) -> WriteMode { - match (reload, spec.strategy) { + let merges_files = spec.files.is_some() && spec.strategy == WriteStrategy::Merge; + match (reload && !merges_files, spec.strategy) { (true, _) => WriteMode::Recreate, (false, WriteStrategy::Append) => WriteMode::Append, (false, WriteStrategy::Replace) => WriteMode::Replace, diff --git a/crates/sqlload/src/loader/_helpers/table_queue.rs b/crates/sqlload/src/loader/_helpers/table_queue.rs index 9f73895fc..218b200ce 100644 --- a/crates/sqlload/src/loader/_helpers/table_queue.rs +++ b/crates/sqlload/src/loader/_helpers/table_queue.rs @@ -81,6 +81,7 @@ pub(crate) fn emit_outcome(outcome: &TableOutcome) { elapsed_ms, position: loaded.note.clone(), changes: loaded.changes.clone(), + transforms: loaded.transforms.clone(), }, Err(reason) if outcome.skipped => EventKind::TableNotStarted { reason: reason.clone(), diff --git a/crates/sqlload/src/loader/constants.rs b/crates/sqlload/src/loader/constants.rs index fd42a769e..07772759d 100644 --- a/crates/sqlload/src/loader/constants.rs +++ b/crates/sqlload/src/loader/constants.rs @@ -49,6 +49,15 @@ pub const PYTHON_COLUMN_TYPE: &str = "python"; /// Why the standalone executable cannot run a project's Python sources. pub const NO_PYTHON_HOST: &str = "this project has Python sources (python_paths), which run only in the sqlload Python package's command; install it (for example `uv tool install sqlload`) and run its sqlload instead of the standalone executable"; +/// Why the standalone executable cannot run a `transform`; `{place}` names the source or table that sets it. +pub const NO_TRANSFORM_HOST: &str = "{place} sets transform, which runs only in the sqlload Python package's command; install it (for example `uv tool install sqlload`) and run its sqlload instead of the standalone executable"; + +/// The rows of a file or range a `transform` returned dicts for, from which a slow call gets the hint that Arrow or a DataFrame is faster. +pub const DICT_HINT_ROWS: u64 = 100_000; + +/// How long a `transform` that returned dicts took, in milliseconds, from which a large file's call gets the hint that Arrow or a DataFrame is faster. +pub const DICT_HINT_MILLISECONDS: u64 = 5_000; + /// The cursor type a Python table function is told before the table's first load shows it. pub const UNKNOWN_CURSOR_TYPE: &str = "unknown"; @@ -165,3 +174,6 @@ pub const UNNEST_ALL: &str = "all"; /// How many dropped rows' ordinals one `DELETE` from a child table's staging table lists. pub const DROPPED_CHUNK: usize = 1000; + +/// The note of a `transform` record whose load rewrote the table whole, so the files earlier records name are gone from it. +pub const TRANSFORM_REWROTE_NOTE: &str = "transform, table rewritten"; diff --git a/crates/sqlload/src/loader/main/apply_project.rs b/crates/sqlload/src/loader/main/apply_project.rs index 78b45dab8..0153fcc00 100644 --- a/crates/sqlload/src/loader/main/apply_project.rs +++ b/crates/sqlload/src/loader/main/apply_project.rs @@ -7,36 +7,19 @@ use crate::adapters::dispatch::main::session::execute::execute; use crate::adapters::dispatch::main::session::execute_ddl::execute_ddl; use crate::adapters::dispatch::models::Destination; use crate::loader::_helpers::publish::records_files; -use crate::loader::_helpers::run_log::reset_merge_log; +use crate::loader::_helpers::run_log::{create_run_logs, reset_merge_log}; use crate::loader::_helpers::table_plan::schema_statements; use crate::loader::errors::LoaderError; use crate::loader::models::ProjectPlan; use crate::schema_plan::models::TableAction; -/// Create the stages SQLLoad maintains, then apply every unblocked table's schema changes; one line per stage or table changed (files targets need none: folders appear with their files; Python tables are created and changed by their loads, once their data shows the columns). +/// Create the stages SQLLoad maintains and the run logs tables with a `transform` record in, then apply every unblocked table's schema changes; one line per stage, run log or table changed (files targets need none: folders appear with their files; Python tables are created and changed by their loads, once their data shows the columns). pub async fn apply_project( plan: &ProjectPlan, destinations: &BTreeMap, ) -> Result, LoaderError> { - let mut lines: Vec = Vec::new(); - let mut created: BTreeSet<&str> = BTreeSet::new(); - for table in plan - .tables - .iter() - .filter(|table| table.schema.blocked.is_none()) - { - let Some(stage) = &table.planned_stage else { - continue; - }; - if !created.insert(stage.name.as_str()) { - continue; - } - let destination = destinations.get(&table.spec.target).ok_or_else(|| { - LoaderError::Invalid(format!("target {:?} is not connected", table.spec.target)) - })?; - execute_ddl(destination, stage.statements.clone()).await?; - lines.push(format!("created {} stage on {}", stage.name, stage.url)); - } + let mut lines = create_stages(plan, destinations).await?; + lines.extend(create_run_logs(plan, destinations).await?); for table in plan .tables .iter() @@ -72,3 +55,30 @@ pub async fn apply_project( } Ok(lines) } + +/// Create the stages the plan's unblocked tables use, once each; one `created` line each. +async fn create_stages( + plan: &ProjectPlan, + destinations: &BTreeMap, +) -> Result, LoaderError> { + let mut lines: Vec = Vec::new(); + let mut created: BTreeSet<&str> = BTreeSet::new(); + for table in plan + .tables + .iter() + .filter(|table| table.schema.blocked.is_none()) + { + let Some(stage) = &table.planned_stage else { + continue; + }; + if !created.insert(stage.name.as_str()) { + continue; + } + let destination = destinations.get(&table.spec.target).ok_or_else(|| { + LoaderError::Invalid(format!("target {:?} is not connected", table.spec.target)) + })?; + execute_ddl(destination, stage.statements.clone()).await?; + lines.push(format!("created {} stage on {}", stage.name, stage.url)); + } + Ok(lines) +} diff --git a/crates/sqlload/src/loader/main/plan_project.rs b/crates/sqlload/src/loader/main/plan_project.rs index 183f86fac..833f8215c 100644 --- a/crates/sqlload/src/loader/main/plan_project.rs +++ b/crates/sqlload/src/loader/main/plan_project.rs @@ -10,10 +10,10 @@ use crate::loader::_helpers::destination::{ block_shared_destinations, mark_shared_destinations, unknown_sharing, }; use crate::loader::_helpers::python::python_table; -use crate::loader::_helpers::resolve::{destination_names, spec_names}; -use crate::loader::_helpers::sources::{ - PLANNED_SOURCES, admit_tables, found_sources, selected, source_ends, unmatched_selectors, +use crate::loader::_helpers::resolve::{ + admit_tables, destination_names, spec_names, unmatched_selectors, }; +use crate::loader::_helpers::sources::{PLANNED_SOURCES, found_sources, selected, source_ends}; #[cfg(feature = "odbc")] use crate::loader::_helpers::table_plan::block_refused_bounds; use crate::loader::_helpers::table_plan::{blocked_plan, plan_table, skip_up_to_date}; diff --git a/crates/sqlload/src/loader/main/selected_specs.rs b/crates/sqlload/src/loader/main/selected_specs.rs index fde879dcf..33f1961e8 100644 --- a/crates/sqlload/src/loader/main/selected_specs.rs +++ b/crates/sqlload/src/loader/main/selected_specs.rs @@ -1,8 +1,8 @@ //! The settings of every selected table, without reading any destination. use crate::loader::_helpers::python::python_table; -use crate::loader::_helpers::resolve::spec_names; -use crate::loader::_helpers::sources::{scoped_tables, selected, unmatched_selectors}; +use crate::loader::_helpers::resolve::{spec_names, unmatched_selectors}; +use crate::loader::_helpers::sources::{scoped_tables, selected}; use crate::loader::_helpers::table_plan::{desired_columns, read_columns}; use crate::loader::errors::LoaderError; use crate::loader::models::{PlanOptions, SelectedTable, SelectedTables}; diff --git a/crates/sqlload/src/loader/models.rs b/crates/sqlload/src/loader/models.rs index 9166a4bc6..7912bcbf9 100644 --- a/crates/sqlload/src/loader/models.rs +++ b/crates/sqlload/src/loader/models.rs @@ -201,6 +201,21 @@ pub struct TableSpec { pub layout: TableLayout, /// How the load names the columns it adds, and under dlt's naming its tables and columns. pub naming: Naming, + /// The Python function each file's (or Sheets range's) rows go through before they load; `None` without `transform`. + pub transform: Option, +} + +/// A table's `transform`: the Python function each file's (or Sheets range's) rows go through before they load. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableTransform { + /// The name `@sl.transform` registered it under. + pub function: String, + /// Its name, `@` and a hash of the file it is defined in. + pub version: String, + /// How many of the files (or Sheets ranges) the table holds loaded under another version, since it was last rewritten. + pub older: usize, + /// `_sqlload_runs`, where each commit records the version, does not exist yet, so `apply` creates it (a load never does, once rows committed). + pub creates_log: bool, } /// A warehouse source's table: the query behind it, if any, and the route its warehouse reads it by. @@ -292,6 +307,23 @@ pub struct FilesSource { pub notes: Vec, /// `copy_to_internal_stage`: files no external stage reads are uploaded unchanged to an internal stage, and loaded from there by the warehouse. pub copy_to_internal_stage: bool, + /// How each file is read for its table's `transform`; `None` without one, when `table`'s columns are the files'. + pub transform: Option, + /// Under `--reload` with `after_load.move_to`, the folder the loaded files were moved to, whose files `table` lists too (keeping their keys, so their names); they stay there after the load. + pub archive: Option, +} + +/// How a files table with a `transform` reads each file before the function runs; its source table's columns are then what the function returned for `sample` at plan. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FilesTransform { + /// The columns each file is read in: as stored (Parquet) or as the workbook holds them (xlsx), JSON's nested values as JSON and every other value as text. + pub input: Vec, + /// Time and timestamp columns read at nanoseconds, as the files hold them. + pub nanos: Vec, + /// What the files' keys and headers are checked against. + pub checks: TextChecks, + /// The file the function ran on at plan, as `_sqlload_source_file` names it. + pub sample: String, } /// A table that matched a source's selection but is not loaded. @@ -704,6 +736,25 @@ pub struct TableLoaded { pub after_load: Option, /// The `replace_key` values a `delete_insert` replaced; `None` without a `replace_key`. pub replaced_keys: Option, + /// How long the table's `transform` took on each file or range, in the order they loaded. + pub transforms: Vec, +} + +/// What a load's reads need to run a table's `transform`: the host, and where each call's time goes. +#[derive(Debug, Clone)] +pub struct TransformRun { + pub host: std::sync::Arc, + pub times: std::sync::Arc>>, +} + +impl TransformRun { + /// A run of `host`'s transforms, no call timed yet. + pub fn new(host: std::sync::Arc) -> Self { + Self { + host, + times: std::sync::Arc::default(), + } + } } /// The `replace_key` values whose rows a table's `delete_insert` replaced, as `load --json` reports them. @@ -826,6 +877,13 @@ pub enum CommitRange { Values(Vec), } +/// A files source's merge under `--reload`: each file merges into `build`, which replaces the table once every file merged, with `records` (a reset, then one per file) in the swap's transaction, so a failed reload leaves the table and its run log as they were. +#[derive(Debug)] +pub struct MergeRebuild { + pub build: TableName, + pub records: Vec, +} + /// One table's load as it commits batch by batch: where the batches go, and what is staged but not yet committed. #[derive(Debug)] pub struct Publisher { @@ -886,6 +944,10 @@ pub struct Publisher { pub settled_files: Vec, /// The child tables `unnest` fills from the table's rows; `None` without `unnest`. pub children: Option, + /// A files source's merge under `--reload`, which merges each file into a new table swapped in at the end. + pub rebuild: Option, + /// The `transform`'s calls so far, and how many the commits recorded: each commit records the version with the files or ranges since. + pub transform: Option<(TransformRun, usize)>, } /// The child tables of one table's load, each planned from its rows as they arrive. @@ -1153,6 +1215,9 @@ pub struct TableResult { /// The `replace_key` values its `delete_insert` replaced; left out without a `replace_key`. #[serde(skip_serializing_if = "Option::is_none")] pub replaced_keys: Option, + /// How long the table's `transform` took on each file or range; left out without one. + #[serde(skip_serializing_if = "Vec::is_empty")] + pub transforms: Vec, } /// Why a source's tables could not be listed, and whether a rerun may get past it. diff --git a/crates/sqlload/src/loader/tests.rs b/crates/sqlload/src/loader/tests.rs index d253954b7..05bd7cc00 100644 --- a/crates/sqlload/src/loader/tests.rs +++ b/crates/sqlload/src/loader/tests.rs @@ -19,5 +19,9 @@ mod retryable_errors; mod source_filter; #[path = "tests/test_types.rs"] mod test_types; +#[path = "tests/test_transform_reader.rs"] +mod transform_reader; +#[path = "tests/test_transform_versions.rs"] +mod transform_versions; #[path = "tests/test_view_lines.rs"] mod view_lines; diff --git a/crates/sqlload/src/loader/tests/helpers.rs b/crates/sqlload/src/loader/tests/helpers.rs index 4e67e16b3..ccfd28771 100644 --- a/crates/sqlload/src/loader/tests/helpers.rs +++ b/crates/sqlload/src/loader/tests/helpers.rs @@ -7,6 +7,7 @@ use arrow_schema::{DataType, Field, Schema, TimeUnit}; use crate::column_types::constants::ARROW_TYPE_METADATA_KEY; use crate::column_types::main::arrow_field::arrow_field; use crate::file_store::errors::StoreError; +use crate::files_source::models::{FilesEncoding, FilesTable, TextChecks}; use crate::http::errors::ApiError; use crate::loader::_helpers::python::bound_array; use crate::loader::_helpers::python::child_paths; @@ -14,17 +15,20 @@ use crate::loader::_helpers::python::{ YieldedData, declared_problems, inferred_table, python_tables, source_spelling, }; use crate::loader::_helpers::resolve::resolve_tables; +use crate::loader::_helpers::table_plan::reader_offer; use crate::loader::errors::LoaderError; -use crate::loader::models::{PlanDetail, PlanItem, TableSpec}; +use crate::loader::models::{FilesSource, PlanDetail, PlanItem, TableSpec, TableTransform}; use crate::loader::tests::test_types::{ BoundArrayTestCase, ChildPathsTestCase, DeclaredProblemsTestCase, InferredTableTestCase, - ItemSpec, SourceSpellingTestCase, + ItemSpec, OlderFilesTestCase, ReloadOrderTestCase, SourceSpellingTestCase, + TransformReaderTestCase, }; use crate::postgres_source::models::{PgType, SourceColumn, SourceTable}; use crate::project::main::parse_python_project::parse_python_project; use crate::python_sources::models::PythonSource; use crate::routing::models::{FileFormat, FileSet, Route, Step, Store}; use crate::schema_plan::models::LiveColumn; +use crate::store_paths::models::{Location, StoreKind}; /// The columns of the inferred table, as `name:type` joined by commas. pub(crate) fn inferred_columns(test_case: &InferredTableTestCase) -> String { @@ -275,10 +279,49 @@ pub(crate) fn merged_file_round_trip( .flat_map(|record| record.files) .next() .ok_or_else(|| format!("no file read back: {:?}", decoded.unreadable))?; - let loaded = crate::loader::_helpers::sources::loaded_file(file, true); + let loaded = crate::loader::_helpers::run_log::loaded_file(file, true); Ok((loaded.name, loaded.version)) } +/// The paths of `test_case`'s files in the order a merge `--reload` applies them. +pub(crate) fn reload_keys(test_case: &ReloadOrderTestCase) -> Vec { + use crate::file_store::models::StoredFile; + use crate::files_source::models::{DataFile, LoadedFile}; + let files: Vec = test_case + .files + .iter() + .map(|(folder, key, modified_ms)| DataFile { + file: StoredFile { + path: object_store::path::Path::from(format!("{folder}{key}")), + relative_key: (*key).to_owned(), + size: 1, + etag: None, + modified_ms: *modified_ms, + member: None, + }, + partitions: Vec::new(), + }) + .collect(); + let merged: Vec = test_case + .merged + .iter() + .map(|(name, version)| LoadedFile { + name: (*name).to_owned(), + version: (*version).to_owned(), + }) + .collect(); + let archive = crate::store_paths::models::Location { + store: crate::store_paths::models::StoreKind::Filesystem { + root: "/".to_owned(), + }, + prefix: "done/".to_owned(), + }; + crate::loader::_helpers::run_log::reload_order(files, "drop/", &merged, Some(&archive)) + .into_iter() + .map(|data| data.file.path.to_string()) + .collect() +} + /// A two-row batch with an `id` column and schema metadata through `with_constants`: whether the result equals the batch given, and each of its columns as `name:Arrow type=value,value`; or the error. pub(crate) fn constants_added( constants: &[(&str, crate::column_types::models::ColumnType, &str)], @@ -490,3 +533,75 @@ pub(crate) fn not_null_keys( .map(|spec| spec.not_null_keys.join(", ")) .ok_or_else(|| resolved.problems.join("; ")) } + +/// Why SQLLoad must read a Parquet files table's files itself, with the case's `transform`. +pub(crate) fn transform_reader( + test_case: &TransformReaderTestCase, +) -> Result, String> { + let mut spec = python_spec("{}")?; + spec.python = None; + spec.files = Some(FilesSource { + location: Location { + store: StoreKind::Filesystem { + root: "/data".to_owned(), + }, + prefix: "drop/".to_owned(), + }, + table: FilesTable { + files: Vec::new(), + columns: Vec::new(), + problems: Vec::new(), + skipped: 0, + unmatched: 0, + dropped: Vec::new(), + objects: Vec::new(), + }, + partition_types: Vec::new(), + constants: Vec::new(), + encoding: FilesEncoding::default(), + checks: TextChecks::default(), + notes: Vec::new(), + copy_to_internal_stage: false, + transform: None, + archive: None, + }); + spec.transform = test_case.transform.map(|function| TableTransform { + function: function.to_owned(), + version: format!("{function}@0"), + older: 0, + creates_log: false, + }); + Ok(reader_offer(&spec, None, None).needs_loader) +} + +/// How many files `older_files` counts in the case's records, for table `t` of a Python source, whose `transform` is now at `f@2`. +pub(crate) fn older_count(test_case: &OlderFilesTestCase) -> Result { + use crate::run_log::models::{RunRecord, StateEvent}; + let spec = python_spec("{}")?; + let records: Vec = test_case + .records + .iter() + .enumerate() + .map(|(index, (event, note, files, version))| { + let record = crate::run_log::main::event_record::event_record( + &crate::loader::_helpers::run_log::record_table(&spec), + StateEvent { + event: *event, + cursor_before: None, + cursor_after: None, + state: None, + note: (*note).to_owned(), + recorded_at: format!("2026-11-15T00:00:{index:02}.000000Z"), + }, + ); + RunRecord { + files: files.iter().map(|file| (*file).to_owned()).collect(), + transform: version.map(str::to_owned), + ..record + } + }) + .collect(); + Ok(crate::loader::_helpers::run_log::older_files( + &records, &spec, "f@2", + )) +} diff --git a/crates/sqlload/src/loader/tests/test_merged_files.rs b/crates/sqlload/src/loader/tests/test_merged_files.rs index f3a0b7cec..c7b9d79a2 100644 --- a/crates/sqlload/src/loader/tests/test_merged_files.rs +++ b/crates/sqlload/src/loader/tests/test_merged_files.rs @@ -1,5 +1,5 @@ -use crate::loader::tests::helpers::{merged_file_round_trip, python_spec}; -use crate::loader::tests::test_types::MergedFileTestCase; +use crate::loader::tests::helpers::{merged_file_round_trip, python_spec, reload_keys}; +use crate::loader::tests::test_types::{MergedFileTestCase, ReloadOrderTestCase}; #[test] fn given_a_merged_file_when_its_record_is_read_back_then_the_ledger_has_its_name_and_version() @@ -35,3 +35,67 @@ fn given_a_merged_file_when_its_record_is_read_back_then_the_ledger_has_its_name } Ok(()) } + +#[test] +fn given_a_merge_ledger_when_reloading_then_files_go_in_the_order_they_merged_and_the_rest_oldest_first() + { + let test_cases = [ + ReloadOrderTestCase { + description: "recorded files keep their merge order, whatever their modified times", + files: &[ + ("drop/", "a.csv", 3_000), + ("drop/", "b.csv", 1_000), + ("drop/", "c.csv", 2_000), + ], + merged: &[ + ("drop/a.csv", "1-3000"), + ("drop/b.csv", "1-1000"), + ("drop/c.csv", "1-2000"), + ], + expected_keys: &["drop/a.csv", "drop/b.csv", "drop/c.csv"], + }, + ReloadOrderTestCase { + description: "a file whose version changed (moved, so copied) keeps its place by name", + files: &[("done/", "a.csv", 9_000), ("drop/", "b.csv", 1_000)], + merged: &[("drop/a.csv", "1-500"), ("drop/b.csv", "1-1000")], + expected_keys: &["done/a.csv", "drop/b.csv"], + }, + ReloadOrderTestCase { + description: "a file merged at two versions goes where its current version merged", + files: &[("drop/", "a.csv", 4_000), ("drop/", "b.csv", 2_000)], + merged: &[ + ("drop/a.csv", "1-1000"), + ("drop/b.csv", "1-2000"), + ("drop/a.csv", "1-4000"), + ], + expected_keys: &["drop/b.csv", "drop/a.csv"], + }, + ReloadOrderTestCase { + description: "a file in the folder under a merged name, at a version never merged, is a newer file: it follows the archived copy", + files: &[ + ("drop/", "export.csv", 2_000), + ("done/", "export.csv", 1_000), + ], + merged: &[("drop/export.csv", "1-1000")], + expected_keys: &["done/export.csv", "drop/export.csv"], + }, + ReloadOrderTestCase { + description: "files never merged follow, oldest modified first", + files: &[ + ("drop/", "new2.csv", 5_000), + ("drop/", "new1.csv", 4_000), + ("drop/", "old.csv", 9_000), + ], + merged: &[("drop/old.csv", "1-9000")], + expected_keys: &["drop/old.csv", "drop/new1.csv", "drop/new2.csv"], + }, + ]; + for test_case in test_cases { + assert_eq!( + reload_keys(&test_case), + test_case.expected_keys, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlload/src/loader/tests/test_transform_reader.rs b/crates/sqlload/src/loader/tests/test_transform_reader.rs new file mode 100644 index 000000000..877b538c6 --- /dev/null +++ b/crates/sqlload/src/loader/tests/test_transform_reader.rs @@ -0,0 +1,28 @@ +use crate::loader::tests::helpers::transform_reader; +use crate::loader::tests::test_types::TransformReaderTestCase; + +#[test] +fn given_a_files_table_with_a_transform_when_routing_then_sqlload_reads_the_files_itself() +-> Result<(), String> { + let test_cases = [ + TransformReaderTestCase { + description: "a transform runs on each file, so no destination reads the files itself", + transform: Some("metrics"), + expected_reason: Some("transform metrics runs on each file"), + }, + TransformReaderTestCase { + description: "without one, Parquet files on a local folder need no reader of SQLLoad's", + transform: None, + expected_reason: None, + }, + ]; + for test_case in test_cases { + assert_eq!( + transform_reader(&test_case)?.as_deref(), + test_case.expected_reason, + "{}", + test_case.description + ); + } + Ok(()) +} diff --git a/crates/sqlload/src/loader/tests/test_transform_versions.rs b/crates/sqlload/src/loader/tests/test_transform_versions.rs new file mode 100644 index 000000000..41ec5512d --- /dev/null +++ b/crates/sqlload/src/loader/tests/test_transform_versions.rs @@ -0,0 +1,83 @@ +use crate::loader::tests::helpers::older_count; +use crate::loader::tests::test_types::OlderFilesTestCase; +use crate::run_log::models::RunEvent; + +#[test] +fn given_a_tables_transform_records_when_planning_then_the_files_an_earlier_version_loaded_are_counted() +-> Result<(), String> { + let test_cases = [ + OlderFilesTestCase { + description: "files from an earlier version stay counted after a later load under the current one", + records: &[ + (RunEvent::Run, "transform", &["a.csv", "b.csv"], Some("f@1")), + (RunEvent::Run, "transform", &["c.csv"], Some("f@2")), + ], + expected_older: 2, + }, + OlderFilesTestCase { + description: "a file merged again under the current version counts by its newest record, whatever its own version", + records: &[ + ( + RunEvent::Run, + "files source merge", + &["a.csv|1-100"], + Some("f@1"), + ), + ( + RunEvent::Run, + "files source merge", + &["b.csv|1-100"], + Some("f@1"), + ), + ( + RunEvent::Run, + "files source merge", + &["a.csv|2-200"], + Some("f@2"), + ), + ], + expected_older: 1, + }, + OlderFilesTestCase { + description: "a load that rewrote the table whole leaves only its own files", + records: &[ + (RunEvent::Run, "transform", &["a.csv", "b.csv"], Some("f@1")), + ( + RunEvent::Run, + "transform, table rewritten", + &["a.csv", "b.csv"], + Some("f@2"), + ), + ], + expected_older: 0, + }, + OlderFilesTestCase { + description: "a reset (a merge table created again) clears what came before", + records: &[ + (RunEvent::Run, "files source merge", &["a.csv"], Some("f@1")), + ( + RunEvent::Reset, + "table created; files merged before no longer count", + &[], + None, + ), + (RunEvent::Run, "files source merge", &["a.csv"], Some("f@2")), + ], + expected_older: 0, + }, + OlderFilesTestCase { + description: "records without a transform version count nothing", + records: &[(RunEvent::Run, "files source merge", &["a.csv"], None)], + expected_older: 0, + }, + ]; + for test_case in test_cases { + assert_eq!( + older_count(&test_case)?, + test_case.expected_older, + "{}", + test_case.description + ); + } + Ok(()) +} diff --git a/crates/sqlload/src/loader/tests/test_types.rs b/crates/sqlload/src/loader/tests/test_types.rs index 0af59ddfb..b9deb586d 100644 --- a/crates/sqlload/src/loader/tests/test_types.rs +++ b/crates/sqlload/src/loader/tests/test_types.rs @@ -62,6 +62,29 @@ pub(crate) struct RenderTallyTestCase { pub(crate) expected_output: &'static str, } +pub(crate) struct TransformReaderTestCase { + pub(crate) description: &'static str, + /// The files table's `transform`, if any. + pub(crate) transform: Option<&'static str>, + /// Why SQLLoad must read the Parquet files itself, if it must. + pub(crate) expected_reason: Option<&'static str>, +} + +/// A table's run log records, oldest first, as event, note, files and `transform` version. +pub(crate) type TransformRecord = ( + crate::run_log::models::RunEvent, + &'static str, + &'static [&'static str], + Option<&'static str>, +); + +pub(crate) struct OlderFilesTestCase { + pub(crate) description: &'static str, + pub(crate) records: &'static [TransformRecord], + /// How many files the table holds from a version other than `f@2`. + pub(crate) expected_older: usize, +} + /// A plan item: its name, its action, and each detail's label and text. pub(crate) type ItemSpec = ( &'static str, @@ -130,6 +153,15 @@ pub(crate) struct MergedFileTestCase { pub(crate) expected_file: (&'static str, &'static str), } +pub(crate) struct ReloadOrderTestCase { + pub(crate) description: &'static str, + /// The files listed now, by folder (`drop/`, the source's, or `done/`, where `after_load` moved them), key and modified time (their version is `1-{time}`). + pub(crate) files: &'static [(&'static str, &'static str, i64)], + /// The table's merge records, oldest first, as name (`drop/{key}`) and version. + pub(crate) merged: &'static [(&'static str, &'static str)], + pub(crate) expected_keys: &'static [&'static str], +} + pub(crate) struct ConstantsAddedTestCase { pub(crate) description: &'static str, /// The constants added: name, type and value as text. diff --git a/crates/sqlload/src/project/_helpers/references.rs b/crates/sqlload/src/project/_helpers/references.rs index c9b028e95..f550ac0ff 100644 --- a/crates/sqlload/src/project/_helpers/references.rs +++ b/crates/sqlload/src/project/_helpers/references.rs @@ -158,6 +158,7 @@ fn source_problems(project: &ProjectConfig, name: &str, source: &SourceConfig) - )), } problems.extend(kind_problems(&place, source)); + problems.extend(transform_problems(project, &place, source)); if source.kind == SourceKind::PostgresCdc { problems.extend(cdc_name_problems(&place, name, source)); } @@ -377,6 +378,36 @@ fn files_only_problems(place: &str, source: &SourceConfig) -> Vec { problems } +/// Where `transform` does not fit: a source other than files or `google_sheets`, an empty name, or a project without `python_paths` to find the function on. +fn transform_problems(project: &ProjectConfig, place: &str, source: &SourceConfig) -> Vec { + let transforms: Vec<&String> = sheet_settings(source) + .filter_map(|settings| settings.transform.as_ref()) + .collect(); + let mut problems: Vec = Vec::new(); + if transforms.is_empty() { + return problems; + } + if !matches!(source.kind, SourceKind::Files | SourceKind::GoogleSheets) { + problems.push(format!( + "{place} transform is for kind = \"files\" and kind = \"google_sheets\" tables; an sl.http job takes transform= on paginate or fan_out" + )); + } + if transforms + .iter() + .any(|transform| transform.trim().is_empty()) + { + problems.push(format!( + "{place} transform is empty; give the name @sl.transform registered the function under" + )); + } + if project.python_paths.is_empty() { + problems.push(format!( + "{place} transform needs python_paths, naming the folder that holds the @sl.transform function" + )); + } + problems +} + /// A source's settings and each of its tables'. fn sheet_settings(source: &SourceConfig) -> impl Iterator { std::iter::once(&source.defaults).chain(source.tables.values().map(|table| &table.settings)) diff --git a/crates/sqlload/src/project/constants.rs b/crates/sqlload/src/project/constants.rs index e5c2969ba..98290b18d 100644 --- a/crates/sqlload/src/project/constants.rs +++ b/crates/sqlload/src/project/constants.rs @@ -31,6 +31,7 @@ pub const TABLE_SETTING_KEYS: &[&str] = &[ "dedup_sort", "naming", "row_id", + "transform", #[cfg(feature = "container-tests")] "test_loader_route", ]; diff --git a/crates/sqlload/src/project/models.rs b/crates/sqlload/src/project/models.rs index a4b1cae02..5da21b768 100644 --- a/crates/sqlload/src/project/models.rs +++ b/crates/sqlload/src/project/models.rs @@ -52,6 +52,12 @@ pub struct ProjectConfig { /// Where settings are written in the project's file, for hints; empty for a project parsed from text alone. #[serde(skip)] pub config_lines: ConfigLines, + /// The Python host that imported `python_paths`, which runs `transform` functions; `None` in the standalone executable. + #[serde(skip)] + pub python_host: Option>, + /// Each `transform` named in the project: its version (a hash of the file it is defined in), or why it cannot run. + #[serde(skip)] + pub transforms: BTreeMap>, } /// The lines of a project file's sources, tables and their `unique_key` (see [`crate::project::main::config_lines::config_lines`]). @@ -819,6 +825,9 @@ pub struct TableSettings { /// Give every row a unique id column (`_sqlload_row_id`); always on under `naming = "dlt"`, as `_dlt_id`. #[serde(default)] pub row_id: Option, + /// The name `@sl.transform` registered a Python function under, found on `python_paths`, run on each file's rows (files) or each range's (`google_sheets`) before they load. + #[serde(default)] + pub transform: Option, } /// What a merge does with rows of one commit that share a `unique_key`. diff --git a/crates/sqlload/src/project/tests/test_parse_project.rs b/crates/sqlload/src/project/tests/test_parse_project.rs index 6fd5aa176..abe6f48c9 100644 --- a/crates/sqlload/src/project/tests/test_parse_project.rs +++ b/crates/sqlload/src/project/tests/test_parse_project.rs @@ -324,6 +324,21 @@ fn given_an_invalid_project_when_parsing_then_every_problem_is_named() { text: "[connections.lake]\nkind = \"filesystem\"\nroot = \"/data\"\n[sources.raw]\nkind = \"files\"\nconnection = \"lake\"\npath = \"raw/\"\nformat = \"xml\"\n", expected_message: "[sources.raw] format = \"xml\" needs xml.record_path", }, + ParseProjectErrorTestCase { + description: "a transform is for files and google_sheets tables", + text: "[sources.app]\nkind = \"postgres\"\nconnection = \"pg\"\ninclude = [\"*\"]\n[sources.app.tables.\"public.a\"]\ntransform = \"fix_rows\"\n", + expected_message: "[sources.app] transform is for kind = \"files\" and kind = \"google_sheets\" tables", + }, + ParseProjectErrorTestCase { + description: "a transform names the function it runs", + text: "[connections.lake]\nkind = \"filesystem\"\nroot = \"/data\"\n[sources.raw]\nkind = \"files\"\nconnection = \"lake\"\npath = \"raw/\"\ntransform = \" \"\n", + expected_message: "[sources.raw] transform is empty; give the name @sl.transform registered the function under", + }, + ParseProjectErrorTestCase { + description: "a transform is found on python_paths", + text: "[connections.lake]\nkind = \"filesystem\"\nroot = \"/data\"\n[sources.raw]\nkind = \"files\"\nconnection = \"lake\"\npath = \"raw/\"\ntransform = \"metrics\"\n", + expected_message: "[sources.raw] transform needs python_paths, naming the folder that holds the @sl.transform function", + }, ParseProjectErrorTestCase { description: "xlsx options need format = xlsx", text: "[connections.lake]\nkind = \"filesystem\"\nroot = \"/data\"\n[sources.raw]\nkind = \"files\"\nconnection = \"lake\"\npath = \"raw/\"\nformat = \"csv\"\nxlsx = { sheet = \"Orders\" }\n", diff --git a/crates/sqlload/src/python_sources/main/_file_context.rs b/crates/sqlload/src/python_sources/main/_file_context.rs new file mode 100644 index 000000000..18489d43e --- /dev/null +++ b/crates/sqlload/src/python_sources/main/_file_context.rs @@ -0,0 +1,36 @@ +//! The `ctx` a files table's `transform` gets with one file's rows. + +use crate::files_source::models::{DataFile, FilesEncoding, SourceFormat}; +use crate::python_sources::models::TransformContext; +use crate::time_text::main::iso_utc::iso_utc; + +/// Where `data` came from, for table `table`: its name as `_sqlload_source_file` has it (`file_name`), when it was modified, and the sheet, range and rows skipped its format reads by. +pub(in crate::python_sources) fn file_context( + table: &str, + (file_name, data): (&str, &DataFile), + encoding: &FilesEncoding, +) -> TransformContext { + let (sheet, range, skip_rows) = match encoding.format { + SourceFormat::Xlsx => ( + encoding.xlsx.sheet.clone(), + encoding.xlsx.range.clone(), + Some(encoding.xlsx.skip_rows), + ), + SourceFormat::Csv => ( + None, + None, + Some(u32::try_from(encoding.csv.skip_rows).unwrap_or(u32::MAX)), + ), + _ => (None, None, None), + }; + TransformContext { + unit: "file", + table: table.to_owned(), + path: Some(file_name.to_owned()), + modified_at: Some(iso_utc(data.file.modified_ms)), + sheet, + range, + skip_rows, + ..TransformContext::default() + } +} diff --git a/crates/sqlload/src/python_sources/main/_run_transform.rs b/crates/sqlload/src/python_sources/main/_run_transform.rs new file mode 100644 index 000000000..a870b0f75 --- /dev/null +++ b/crates/sqlload/src/python_sources/main/_run_transform.rs @@ -0,0 +1,65 @@ +//! Run a `transform` function on one file's or range's rows, passing on each batch it returns or yields as it crosses. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; +use tokio::sync::mpsc::Sender; + +use crate::python_sources::models::{TransformCall, TransformRan}; +use crate::python_sources::types::{PythonHost, Streamed}; + +/// Call `call`'s function on `rows` (in `schema`) on a blocking thread and pass each batch it returned or yielded to `on_batch` as it crosses, one batch waiting at most, so a generator's output never piles up; how the call went, or why it failed (its exception, or what it returned cannot load). An error of `on_batch` stops the call and is returned as it is. +pub(in crate::python_sources) async fn run_transform( + host: Arc, + call: TransformCall, + (rows, schema): (Vec, SchemaRef), + mut on_batch: impl AsyncFnMut(RecordBatch) -> Result<(), E>, +) -> Result, E> { + let rows_in = rows.iter().map(RecordBatch::num_rows).sum::() as u64; + let (sender, mut receiver) = tokio::sync::mpsc::channel::(1); + let crossing = + tokio::task::spawn_blocking(move || cross(host.as_ref(), &call, (rows, schema), &sender)); + while let Some(batch) = receiver.recv().await { + if let Err(error) = on_batch(batch).await { + drop(receiver); + let _ = crossing.await; + return Err(error); + } + } + let crossed = match crossing.await { + Ok(crossed) => crossed, + Err(error) => Err(format!("the transform's thread failed: {error}")), + }; + Ok(crossed.map(|(dicts, spent)| TransformRan { + dicts, + rows_in, + seconds: spent.as_secs_f64(), + })) +} + +/// Call the function and send each batch on until it ends or nothing takes them: whether it gave dicts, and the time spent other than waiting for a batch to be taken. +fn cross( + host: &dyn PythonHost, + call: &TransformCall, + (rows, schema): (Vec, SchemaRef), + sender: &Sender, +) -> Result<(bool, Duration), String> { + let started = Instant::now(); + let mut waited = Duration::ZERO; + let mut stream = host.transform(call, rows, schema)?; + while let Some(next) = stream.next_batch() { + if let Streamed::Batch(batch) = next.map_err(|error| error.message)? { + let waiting = Instant::now(); + if sender.blocking_send(batch).is_err() { + break; + } + waited += waiting.elapsed(); + } + } + Ok(( + stream.crossed_dicts(), + started.elapsed().saturating_sub(waited), + )) +} diff --git a/crates/sqlload/src/python_sources/main/mod.rs b/crates/sqlload/src/python_sources/main/mod.rs index 65da098e6..c6302c635 100644 --- a/crates/sqlload/src/python_sources/main/mod.rs +++ b/crates/sqlload/src/python_sources/main/mod.rs @@ -1,4 +1,8 @@ +pub(in crate::python_sources) mod _file_context; +pub(in crate::python_sources) mod _run_transform; pub mod connection_settings; pub mod declared_sources; pub mod name_problems; pub mod toml_python_problems; +pub mod transform_files; +pub mod transform_range; diff --git a/crates/sqlload/src/python_sources/main/transform_files.rs b/crates/sqlload/src/python_sources/main/transform_files.rs new file mode 100644 index 000000000..f2143f694 --- /dev/null +++ b/crates/sqlload/src/python_sources/main/transform_files.rs @@ -0,0 +1,56 @@ +//! Run a files table's `transform` on each file of a read, in order. + +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; + +use crate::files_source::models::DataFile; +use crate::python_sources::errors::ReadError; +use crate::python_sources::main::_file_context::file_context; +use crate::python_sources::main::_run_transform::run_transform; +use crate::python_sources::models::{TransformCall, TransformFiles, TransformTime}; + +/// For each of `request`'s files in order: its rows from `read`, the function run on them, and each batch it returned or yielded passed to `on_rows` as it crosses, with the file, its name as `_sqlload_source_file` has it, and the rows returned for it before. A file whose function fails stops there, naming the file; returns how long each file's call took. +pub async fn transform_files>( + request: &TransformFiles<'_>, + mut read: impl AsyncFnMut(&DataFile) -> Result<(Vec, SchemaRef), E>, + mut on_rows: impl AsyncFnMut((&DataFile, &str, u64), RecordBatch) -> Result<(), E>, +) -> Result, E> { + let mut times: Vec = Vec::new(); + for data in request.files { + let file_name = format!("{}{}", request.location, data.file.relative_key); + let call = TransformCall { + context: file_context( + &request.call.context.table, + (&file_name, data), + request.encoding, + ), + ..request.call.clone() + }; + let mut rows_before: u64 = 0; + let ran = run_transform( + request.host.clone(), + call, + read(data).await?, + async |batch: RecordBatch| { + let rows = batch.num_rows() as u64; + on_rows((data, &file_name, rows_before), batch).await?; + rows_before += rows; + Ok::<(), E>(()) + }, + ) + .await? + .map_err(|reason| { + ReadError::from(format!( + "{file_name}: transform {} failed: {reason}", + request.call.function + )) + })?; + times.push(TransformTime { + file: file_name, + milliseconds: (ran.seconds * 1000.0).round() as u64, + rows: ran.rows_in, + dicts: ran.dicts, + }); + } + Ok(times) +} diff --git a/crates/sqlload/src/python_sources/main/transform_range.rs b/crates/sqlload/src/python_sources/main/transform_range.rs new file mode 100644 index 000000000..4a4f74866 --- /dev/null +++ b/crates/sqlload/src/python_sources/main/transform_range.rs @@ -0,0 +1,40 @@ +//! Run a `google_sheets` table's `transform` on one range's rows. + +use std::sync::Arc; + +use arrow_array::RecordBatch; + +use crate::python_sources::main::_run_transform::run_transform; +use crate::python_sources::models::{TransformCall, TransformTime}; +use crate::python_sources::types::PythonHost; + +/// `call`'s function run on `rows`, one range's (its `context.range`): the batches it returned or yielded and how long it took; the error names the range and the function. +pub async fn transform_range( + host: Arc, + call: TransformCall, + rows: RecordBatch, +) -> Result<(Vec, TransformTime), String> { + let range = call + .context + .range + .clone() + .or_else(|| call.context.sheet.clone()) + .unwrap_or_default(); + let function = call.function.clone(); + let schema = rows.schema(); + let mut batches: Vec = Vec::new(); + let ran = run_transform(host, call, (vec![rows], schema), async |batch| { + batches.push(batch); + Ok::<(), String>(()) + }) + .await + .and_then(|ran| ran) + .map_err(|reason| format!("{range}: transform {function} failed: {reason}"))?; + let time = TransformTime { + file: range, + milliseconds: (ran.seconds * 1000.0).round() as u64, + rows: ran.rows_in, + dicts: ran.dicts, + }; + Ok((batches, time)) +} diff --git a/crates/sqlload/src/python_sources/models.rs b/crates/sqlload/src/python_sources/models.rs index 647274243..a23747e92 100644 --- a/crates/sqlload/src/python_sources/models.rs +++ b/crates/sqlload/src/python_sources/models.rs @@ -52,3 +52,72 @@ pub struct CursorRead { /// Whether `start` itself is included. pub inclusive: bool, } + +/// What a `transform` gets as `ctx`: the unit its rows came from, `"file"`, `"range"` or `"page"`. A file sets `path`, `modified_at` and, for its format, `sheet`, `range` and `skip_rows`; a Google Sheets range sets `spreadsheet_id`, `sheet` (the tab) and `range`; an `sl.http` page sets `page` and, in a fan-out, `key`. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)] +pub struct TransformContext { + pub unit: &'static str, + /// The table's name in its source. + pub table: String, + /// The file as `_sqlload_source_file` names it. + pub path: Option, + /// When the file was last modified, RFC 3339 in UTC. + pub modified_at: Option, + pub sheet: Option, + pub range: Option, + pub skip_rows: Option, + pub spreadsheet_id: Option, + /// The page's number within its request, from 1. + pub page: Option, + /// The fields of the fan-out key the page was requested for. + pub key: Option>, +} + +/// One call of a `transform`, and how what it returned crosses. +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct TransformCall { + /// The name `@sl.transform` registered it under. + pub function: String, + pub context: TransformContext, + /// Columns loaded masked. + pub masked: Vec, + /// Columns pinned as `array(...)`: their Arrow lists cross as lists. + pub arrays: Vec, + /// How the table's columns load and the destination's name limit. + #[serde(skip)] + pub naming: (crate::load_naming::models::Naming, Option), +} + +/// How a `transform`'s call on one file or range went, once every batch it returned or yielded has crossed into Arrow and been passed on. +#[derive(Debug, Clone)] +pub struct TransformRan { + /// Whether it returned dicts, which cross slower than Arrow or a DataFrame. + pub dicts: bool, + /// The rows it was given. + pub rows_in: u64, + /// How long the call and its crossing took, not counting the waits for its batches to be taken. + pub seconds: f64, +} + +/// How long a `transform` took on one file or Sheets range. +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct TransformTime { + /// The file as `_sqlload_source_file` names it, or the range. + pub file: String, + /// How long the call, and its crossing into Arrow, took. + pub milliseconds: u64, + /// The rows the function was given. + pub rows: u64, + /// Whether it returned dicts, which cross slower than Arrow or a DataFrame. + pub dicts: bool, +} + +/// What `transform_files` runs: the host, the call (its `context.table` and how what the function returns crosses; each file sets the rest of the context), and the files, in the folder printed as `location`, encoded as `encoding`. +#[derive(Debug, Clone)] +pub struct TransformFiles<'a> { + pub host: std::sync::Arc, + pub call: TransformCall, + pub files: &'a [crate::files_source::models::DataFile], + pub location: String, + pub encoding: &'a crate::files_source::models::FilesEncoding, +} diff --git a/crates/sqlload/src/python_sources/types.rs b/crates/sqlload/src/python_sources/types.rs index aaae1b86c..03f4b0d3d 100644 --- a/crates/sqlload/src/python_sources/types.rs +++ b/crates/sqlload/src/python_sources/types.rs @@ -1,9 +1,10 @@ use std::path::{Path, PathBuf}; use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; use crate::python_sources::errors::ReadError; -use crate::python_sources::models::{PythonSource, TableRead}; +use crate::python_sources::models::{PythonSource, TableRead, TransformCall}; /// What a table function's read gives next. #[derive(Debug)] @@ -14,9 +15,13 @@ pub enum Streamed { } /// A table function's data, one Arrow batch at a time; `next_batch` blocks. -pub trait BatchStream: Send { +pub trait BatchStream: Send + std::fmt::Debug { /// The next batch or state, an error, or `None` at the end. fn next_batch(&mut self) -> Option>; + /// Whether a `transform` returned or yielded dicts, which cross slower than Arrow or a DataFrame; known once the stream has ended. A table function's read says `false`. + fn crossed_dicts(&self) -> bool { + false + } } /// Runs Python sources: the `sqlload` Python package provides it; the standalone executable has none. @@ -27,6 +32,15 @@ pub trait PythonHost: Send + Sync + std::fmt::Debug { fn read_table(&self, request: &TableRead) -> Result, String>; /// The position the table function last started for `table` left in `ctx.state`, as JSON text; `None` when it set none. fn table_state(&self, table: &str) -> Result, String>; + /// Look up the `transform` that `@sl.transform` registered as `name` and check it can run: its version, `name@` and a hash of the file it is defined in, or why it cannot (no function, or two, registered as `name`). + fn transform_version(&self, name: &str) -> Result; + /// Call a `transform` on `rows` (in `schema`) and stream what it returned or yielded, crossed as a table function's yields are, as it yields it. + fn transform( + &self, + call: &TransformCall, + rows: Vec, + schema: SchemaRef, + ) -> Result, String>; /// Close the clients the run built. fn finish(&self); } diff --git a/crates/sqlload/src/run_log/main/event_record.rs b/crates/sqlload/src/run_log/main/event_record.rs index 92fe81016..eb2e77e4b 100644 --- a/crates/sqlload/src/run_log/main/event_record.rs +++ b/crates/sqlload/src/run_log/main/event_record.rs @@ -21,5 +21,6 @@ pub fn event_record(table: &RecordTable, event: StateEvent) -> RunRecord { note: Some(event.note), recorded_at: event.recorded_at, batch: None, + transform: None, } } diff --git a/crates/sqlload/src/run_log/models.rs b/crates/sqlload/src/run_log/models.rs index 4dceca10c..e6032f1e9 100644 --- a/crates/sqlload/src/run_log/models.rs +++ b/crates/sqlload/src/run_log/models.rs @@ -53,6 +53,9 @@ pub struct RunRecord { /// The run's commit this record went with, from 1, when the run committed in batches. #[serde(default, skip_serializing_if = "Option::is_none")] pub batch: Option, + /// The version of the table's `transform` a load ran (its name, `@` and a hash of the file it is defined in). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub transform: Option, } /// What a finished run adds to its record. diff --git a/crates/sqlload/src/run_log/tests/helpers.rs b/crates/sqlload/src/run_log/tests/helpers.rs index 504476688..99061a7f9 100644 --- a/crates/sqlload/src/run_log/tests/helpers.rs +++ b/crates/sqlload/src/run_log/tests/helpers.rs @@ -31,6 +31,7 @@ pub(crate) fn sample_record() -> RunRecord { note: None, recorded_at: "2026-11-04T10:00:00.000000Z".to_owned(), batch: None, + transform: None, } } diff --git a/crates/sqlload/src/sheets_source/_helpers/cells.rs b/crates/sqlload/src/sheets_source/_helpers/cells.rs index a1d54adf4..7910ce530 100644 --- a/crates/sqlload/src/sheets_source/_helpers/cells.rs +++ b/crates/sqlload/src/sheets_source/_helpers/cells.rs @@ -59,6 +59,17 @@ pub(crate) fn cell_text(cell: &Cell, column_type: &ColumnType) -> Result String { + let dated = match kind { + NumberKind::Date => ColumnType::Date, + NumberKind::DateTime => ColumnType::TimestampNtz, + NumberKind::Time => ColumnType::Time, + NumberKind::Plain | NumberKind::Fractional => ColumnType::Float, + }; + number_text(value, kind, &dated) +} + /// A number as the text of its value: a date, a timestamp or a time of day for dated cells (by the column's type, so a date loads into a timestamp column at midnight), a whole number without a fraction, any other number as it is. fn number_text(value: f64, kind: NumberKind, column_type: &ColumnType) -> String { let time = serial_time(value); diff --git a/crates/sqlload/src/sheets_source/main/mod.rs b/crates/sqlload/src/sheets_source/main/mod.rs index b4c41ddfc..ff4e5701e 100644 --- a/crates/sqlload/src/sheets_source/main/mod.rs +++ b/crates/sqlload/src/sheets_source/main/mod.rs @@ -3,6 +3,7 @@ pub mod a1_width; pub mod duplicate_key; pub mod read_grids; pub mod sheet_batch; +pub mod sheet_input; pub mod sheet_table; pub mod sheets_client; pub mod tab_name; diff --git a/crates/sqlload/src/sheets_source/main/sheet_input.rs b/crates/sqlload/src/sheets_source/main/sheet_input.rs new file mode 100644 index 000000000..16e467095 --- /dev/null +++ b/crates/sqlload/src/sheets_source/main/sheet_input.rs @@ -0,0 +1,38 @@ +//! A range's rows as a `transform` gets them. + +use arrow_array::RecordBatch; + +use crate::column_types::models::ColumnType; +use crate::sheets_source::_helpers::cells::value_text; +use crate::sheets_source::_helpers::layout::data_rows; +use crate::sheets_source::main::sheet_batch::sheet_batch; +use crate::sheets_source::main::sheet_table::sheet_table; +use crate::sheets_source::models::{Cell, Grid, HeaderRule}; + +/// The rows of `grid` under `rule` as one batch for a `transform`: a column per header name, typed by its cells as a table without a transform is, except that an empty text (a formula's `""`) is NULL and does not count towards the type, and a number in a column of text is its value (`0.9`), not the text the sheet shows (`90.0%`). An error cell, or a header cell that cannot name a column, is an error naming it. +pub fn sheet_input(grid: &Grid, rule: &HeaderRule) -> Result { + let mut cleaned = grid.clone(); + for cell in cleaned.rows.iter_mut().flatten() { + if *cell == Cell::Text(String::new()) { + *cell = Cell::Empty; + } + } + let table = sheet_table(&cleaned, rule).map_err(|problems| problems.join("; "))?; + let mut columns: Vec<(String, ColumnType)> = Vec::with_capacity(table.columns.len()); + for column in &table.columns { + columns.push((column.name.clone(), column.column_type.clone()?)); + } + let rows = data_rows(&cleaned, rule, columns.len()); + for (position, (_, column_type)) in columns.iter().enumerate() { + if *column_type != ColumnType::Text { + continue; + } + for row in &rows { + if let Some(Cell::Number { value, kind, shown }) = cleaned.rows[*row].get_mut(position) + { + *shown = value_text(*value, *kind); + } + } + } + sheet_batch(&cleaned, rule, &columns) +} diff --git a/crates/sqlload/src/sheets_source/tests.rs b/crates/sqlload/src/sheets_source/tests.rs index ac292e100..ce2b9d3ad 100644 --- a/crates/sqlload/src/sheets_source/tests.rs +++ b/crates/sqlload/src/sheets_source/tests.rs @@ -4,6 +4,8 @@ mod a1_range; mod helpers; #[path = "tests/test_sheet_batch.rs"] mod sheet_batch; +#[path = "tests/test_sheet_input.rs"] +mod sheet_input; #[path = "tests/test_sheet_table.rs"] mod sheet_table; #[path = "tests/test_sheets_client.rs"] diff --git a/crates/sqlload/src/sheets_source/tests/helpers.rs b/crates/sqlload/src/sheets_source/tests/helpers.rs index 09e3bd1ee..cfdcac831 100644 --- a/crates/sqlload/src/sheets_source/tests/helpers.rs +++ b/crates/sqlload/src/sheets_source/tests/helpers.rs @@ -56,3 +56,83 @@ pub(crate) fn rows_text(batch: &RecordBatch) -> Result { .collect(); Ok(rows.join(";")) } + +/// A synthetic KPI tab: a merged title row, a header row, label-only rows, KPI blocks with their measures, a note, and a formula's `""` here and there. +pub(crate) fn kpis() -> serde_json::Value { + let percent = |value: f64, shown: &str| number(value, shown, "PERCENT"); + sheet( + "KPIs", + (0, 0), + vec![ + vec![text("Service KPIs 2026")], + vec![ + text("ID"), + text("KPI"), + text("Period 1"), + text("Period 2"), + text("Jan"), + text("Feb"), + text("Mar"), + ], + vec![blank(), text("Customer service")], + vec![ + blank(), + text("Targets reviewed quarterly"), + text("note: P1 excludes returns"), + ], + vec![ + text("K1"), + text("On-time delivery"), + percent(0.9, "90.0%"), + text(""), + percent(0.95, "95.0%"), + number(0.8, "0.8", ""), + text(""), + ], + vec![ + blank(), + text("Delivered on time"), + number(9.0, "9", ""), + blank(), + number(19.0, "19", ""), + number(8.0, "8", ""), + ], + vec![ + blank(), + text("Delivered"), + number(10.0, "10", ""), + number(0.0, "0", ""), + number(20.0, "20", ""), + number(10.0, "10", ""), + number(0.0, "0", ""), + ], + vec![ + text("K2"), + text("First contact resolution"), + percent(0.5, "50.0%"), + percent(0.6, "60.0%"), + text(""), + percent(0.75, "75.0%"), + percent(1.0, "100.0%"), + ], + vec![ + blank(), + text("Escalations"), + number(1.0, "1", ""), + text("n/a - system outage"), + blank(), + number(0.0, "0", ""), + ], + vec![], + vec![ + text("K3"), + text("Average handle time (min)"), + number(6.5, "6.5", ""), + number(7.0, "7", ""), + number(6.25, "6.25", ""), + blank(), + number(5.5, "5.5", ""), + ], + ], + ) +} diff --git a/crates/sqlload/src/sheets_source/tests/test_sheet_input.rs b/crates/sqlload/src/sheets_source/tests/test_sheet_input.rs new file mode 100644 index 000000000..0e466e526 --- /dev/null +++ b/crates/sqlload/src/sheets_source/tests/test_sheet_input.rs @@ -0,0 +1,56 @@ +use crate::sheets_source::_helpers::cells::parse_grid; +use crate::sheets_source::main::sheet_input::sheet_input; +use crate::sheets_source::models::HeaderRule; +use crate::sheets_source::tests::helpers::{error, kpis, rows_text, sheet, text}; +use crate::sheets_source::tests::test_types::SheetInputTestCase; + +#[test] +fn given_a_range_with_a_transform_when_reading_its_input_then_empty_texts_are_null_and_numbers_in_text_columns_their_values() + { + let test_cases = [ + SheetInputTestCase { + description: "rows above header_row are skipped; a formula's \"\" is NULL and types nothing, so Jan and Mar are float; a percentage among text is its value, 0.9, not 90.0%; a blank row is left out", + sheet: kpis(), + rule: HeaderRule { + header_row: 2, + ..HeaderRule::default() + }, + expected_result: Ok(( + "ID:Utf8,KPI:Utf8,Period 1:Utf8,Period 2:Utf8,Jan:Float64,Feb:Float64,Mar:Float64", + "NULL|Customer service|NULL|NULL|NULL|NULL|NULL;NULL|Targets reviewed quarterly|note: P1 excludes returns|NULL|NULL|NULL|NULL;K1|On-time delivery|0.9|NULL|0.95|0.8|NULL;NULL|Delivered on time|9|NULL|19.0|8.0|NULL;NULL|Delivered|10|0|20.0|10.0|0.0;K2|First contact resolution|0.5|0.6|NULL|0.75|1.0;NULL|Escalations|1|n/a - system outage|NULL|0.0|NULL;K3|Average handle time (min)|6.5|7|6.25|NULL|5.5", + )), + }, + SheetInputTestCase { + description: "an error cell still fails the read, named", + sheet: sheet( + "Errors", + (0, 0), + vec![vec![text("ratio")], vec![error("#DIV/0!")]], + ), + rule: HeaderRule::default(), + expected_result: Err( + "Errors!A2 holds the error #DIV/0!; fix the formula, or leave the row out of the range", + ), + }, + ]; + for test_case in test_cases { + let seen = sheet_input(&parse_grid(&test_case.sheet), &test_case.rule).and_then(|batch| { + let schema = batch.schema(); + let fields: Vec = schema + .fields() + .iter() + .map(|field| format!("{}:{}", field.name(), field.data_type())) + .collect(); + Ok((fields.join(","), rows_text(&batch)?)) + }); + assert_eq!( + seen, + test_case + .expected_result + .map(|(fields, rows)| (fields.to_owned(), rows.to_owned())) + .map_err(str::to_owned), + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlload/src/sheets_source/tests/test_types.rs b/crates/sqlload/src/sheets_source/tests/test_types.rs index 156d01e7a..f761f98d4 100644 --- a/crates/sqlload/src/sheets_source/tests/test_types.rs +++ b/crates/sqlload/src/sheets_source/tests/test_types.rs @@ -20,6 +20,14 @@ pub(crate) struct SheetBatchTestCase { pub(crate) expected_result: Result<&'static str, &'static str>, } +pub(crate) struct SheetInputTestCase { + pub(crate) description: &'static str, + pub(crate) sheet: serde_json::Value, + pub(crate) rule: HeaderRule, + /// Each column as `name:type` joined by `,`, and the rows, each as its values joined by `|`, rows joined by `;`; or the error. + pub(crate) expected_result: Result<(&'static str, &'static str), &'static str>, +} + pub(crate) struct KeyTestCase { pub(crate) description: &'static str, pub(crate) source: crate::sheets_source::models::KeySource, diff --git a/crates/sqlload/tests/loader/after_load/helpers.rs b/crates/sqlload/tests/loader/after_load/helpers.rs index 9d73ef69a..9cf2e6756 100644 --- a/crates/sqlload/tests/loader/after_load/helpers.rs +++ b/crates/sqlload/tests/loader/after_load/helpers.rs @@ -6,7 +6,8 @@ use sqlload::files_source::main::after_load::after_load; use sqlload::store_paths::models::{Location, StoreKind}; use sqlload_testkit::seaweedfs::constants::REGION; -use crate::test_types::{AfterLoadTestCase, S3AfterLoadTestCase}; +use crate::helpers::{duckdb_text, parquet_bytes, run_sqlload}; +use crate::test_types::{AfterLoadTestCase, ArchivedReloadTestCase, S3AfterLoadTestCase}; /// Run `test_case`'s steps through the `sqlload` executable (`load --json`), asserting each step's exit code and report; returns the files left below the root as `path=contents`, sorted. pub(crate) async fn run_after_load_case( @@ -187,3 +188,40 @@ pub(crate) async fn run_s3_after_load( files.sort(); Ok((report, files)) } + +/// Run `test_case` through the executable: load its first files (moved to `done/`), write its later ones, then `--reload`; the reload's plan and the table's rows as `id:score`, in order. +pub(crate) fn run_archived_reload( + test_case: &ArchivedReloadTestCase, +) -> Result<(String, String), String> { + let root = tempfile::tempdir().map_err(|error| error.to_string())?; + let database = root.path().join("out.duckdb"); + let config = root.path().join("sqlload.toml"); + let text = format!( + "[connections.wh]\nkind = \"duckdb\"\ndatabase = \"{}\"\n\n[connections.files]\nkind = \"filesystem\"\nroot = \"{}\"\n\n[targets.dev]\nconnection = \"wh\"\nloader_schema = \"raw\"\n\n[sources.scores]\nkind = \"files\"\ndescription = \"Test table.\"\nconnection = \"files\"\npath = \"drop/\"\nformat = \"parquet\"\nwrite_strategy = \"append\"\nafter_load = {{ move_to = \"done\" }}\n", + database.display(), + root.path().display() + ); + std::fs::write(&config, text).map_err(|error| error.to_string())?; + std::fs::create_dir_all(root.path().join("drop")).map_err(|error| error.to_string())?; + let write = |files: &[crate::test_types::ParquetFile]| { + files.iter().try_for_each(|file| { + std::fs::write( + root.path().join("drop").join(file.key), + parquet_bytes(file)?, + ) + .map_err(|error| error.to_string()) + }) + }; + write(test_case.first)?; + let (loaded, output) = run_sqlload(&config, &["load"], ("", ""))?; + assert!(loaded, "{} first load:\n{output}", test_case.description); + write(test_case.later)?; + let (_, plan) = run_sqlload(&config, &["plan", "--reload"], ("", ""))?; + let (reloaded, output) = run_sqlload(&config, &["load", "--reload", "--yes"], ("", ""))?; + assert!(reloaded, "{} reload:\n{output}", test_case.description); + let rows = duckdb_text( + &database, + "SELECT string_agg(id || ':' || score, ',' ORDER BY id, score) FROM raw.scores", + )?; + Ok((plan, rows)) +} diff --git a/crates/sqlload/tests/loader/helpers.rs b/crates/sqlload/tests/loader/helpers.rs index 37c62f5cf..e323ad80a 100644 --- a/crates/sqlload/tests/loader/helpers.rs +++ b/crates/sqlload/tests/loader/helpers.rs @@ -26,7 +26,7 @@ use tokio_postgres::Client; use crate::files_target_helpers::file_batches; use crate::test_types::{ - CdcStep, CursorSource, DebugStageTestCase, FilesMergeTestCase, FilesScenario, + CdcStep, CursorSource, DebugStageTestCase, FileSetTestCase, FilesScenario, FilesScenarioTestCase, FilesTargetScenarioTestCase, LargePartitionTestCase, LiveDatabricks, MetadataRoundTripTestCase, ParquetFile, ReplacedKeysTestCase, RouteFiles, RouteOutcome, RouteTestCase, ScenarioTestCase, TextFilesScenarioTestCase, WrittenPartition, WrittenRow, @@ -151,26 +151,6 @@ pub(crate) async fn load_selected( Ok(lines.join("\n")) } -/// Plan and apply once, then load only when `loads` (otherwise as `sqlload apply`): the rendered plan, or the error text. -async fn run_step( - project_text: &str, - destinations: &BTreeMap, - loads: bool, -) -> Result { - let project = described_project(project_text).await?; - let mut plan = plan_project(&project, destinations, &PlanOptions::default()) - .await - .map_err(|error| error.to_string())?; - let (lines, ok) = render_plan(&plan); - ok.then_some(()).ok_or_else(|| lines.join("\n"))?; - apply_project(&plan, destinations) - .await - .map_err(|error| error.to_string())?; - plan.tables.retain(|_| loads); - load_all(&project, &plan, destinations).await?; - Ok(lines.join("\n")) -} - /// Load every table of `plan`; the first table that failed is the error. pub(crate) async fn load_all( project: &ProjectConfig, @@ -420,6 +400,29 @@ pub(crate) async fn run_files_scenario( run_written_files(endpoint, directory, scenario).await } +/// Load `test_case`'s file set once from a local folder into DuckDB, checking the load's error; the check query's result. +pub(crate) async fn run_file_set( + directory: &std::path::Path, + test_case: &FileSetTestCase, +) -> Result, String> { + let files = (test_case.file_set)()?; + let scenario = FilesScenario { + description: test_case.description, + store: local_store, + settings: test_case.settings, + runs: vec![WrittenRun { + files: files + .iter() + .map(|(key, contents)| (key.as_str(), contents.clone())) + .collect(), + expected_error: test_case.expected_error, + expected_plan: "", + }], + check_query: test_case.check_query, + }; + run_written_files("", directory, scenario).await +} + /// Run a JSON or CSV files scenario like [`run_files_scenario`]. pub(crate) async fn run_text_files_scenario( endpoint: &str, @@ -1786,73 +1789,6 @@ pub(crate) fn ssh_settings( .replace("FILES", files) } -/// Run `test_case` against the target connection `target` (its TOML body), with a `loader_schema` of `schema`: CSV files written to a fresh folder with their modified times, loaded step by step; the check query's result. -pub(crate) async fn run_merge_case( - (target, schema): (&str, &str), - test_case: &FilesMergeTestCase, -) -> Result, String> { - let directory = tempfile::tempdir().map_err(|error| error.to_string())?; - let root = directory.path().display().to_string(); - let text = format!( - "[connections.files]\nkind = \"filesystem\"\nroot = \"{root}\"\n\n[connections.target]\n{target}\n\n[targets.dev]\nconnection = \"target\"\nloader_schema = \"{schema}\"\n\n[sources.scores]\nkind = \"files\"\ndescription = \"Test table.\"\nconnection = \"files\"\npath = \"drop/\"\nformat = \"csv\"\n{}\n", - test_case.settings - ); - let project = described_project(&text).await?; - let destinations = connect_targets(&project) - .await - .map_err(|error| error.to_string())?; - let destination = destinations - .get("dev") - .ok_or_else(|| "no dev destination".to_owned())?; - std::fs::create_dir_all(directory.path().join("drop")).map_err(|error| error.to_string())?; - for (index, step) in test_case.steps.iter().enumerate() { - execute( - destination, - step.before.iter().map(|sql| (*sql).to_owned()).collect(), - ) - .await - .map_err(|error| error.to_string())?; - for key in step.removed { - std::fs::remove_file(directory.path().join("drop").join(key)) - .map_err(|error| error.to_string())?; - } - for file in step.files { - let path = directory.path().join("drop").join(file.key); - std::fs::create_dir_all(path.parent().unwrap_or(directory.path())) - .map_err(|error| error.to_string())?; - std::fs::write(&path, file.contents).map_err(|error| error.to_string())?; - let modified = - std::time::UNIX_EPOCH + std::time::Duration::from_secs(file.modified_secs); - std::fs::File::options() - .write(true) - .open(&path) - .and_then(|opened| opened.set_modified(modified)) - .map_err(|error| error.to_string())?; - } - let outcome = run_step(&text, &destinations, step.loads).await; - let (plan, error) = ( - outcome.clone().unwrap_or_default(), - outcome.err().unwrap_or_default(), - ); - assert!( - words(&error).contains(&words(step.expected_error)) - && error.is_empty() == step.expected_error.is_empty(), - "{} step {}: {error}", - test_case.description, - index + 1 - ); - assert!( - plan.contains(step.expected_plan), - "{} step {}: {plan}", - test_case.description, - index + 1 - ); - } - single_text(destination, test_case.check_query.to_owned()) - .await - .map_err(|error| error.to_string()) -} - /// Write a project loading the NDJSON files of `root/drop/` into `[connections.wh]` (given as `connection`'s settings, schema `raw`) with `settings` added, then load once per run, each after writing its new file; returns the project's file, and fails with the output of a load that fails. pub(crate) fn load_ndjson_runs( root: &std::path::Path, diff --git a/crates/sqlload/tests/loader/merge/helpers.rs b/crates/sqlload/tests/loader/merge/helpers.rs index 063d35b61..c46029ea8 100644 --- a/crates/sqlload/tests/loader/merge/helpers.rs +++ b/crates/sqlload/tests/loader/merge/helpers.rs @@ -5,13 +5,18 @@ use std::collections::BTreeMap; use sqlload::adapters::dispatch::main::session::execute::execute; use sqlload::adapters::dispatch::main::session::single_text::single_text; use sqlload::adapters::dispatch::models::Destination; +use sqlload::loader::main::apply_project::apply_project; use sqlload::loader::main::connect_targets::connect_targets; +use sqlload::loader::main::plan_project::plan_project; +use sqlload::loader::main::render_plan::render_plan; +use sqlload::loader::models::PlanOptions; use sqlload::project::models::ProjectConfig; use sqlload_testkit::postgres::constants::{PASSWORD, USER}; -use crate::helpers::{described_project, load_selected, words}; +use crate::helpers::{described_project, load_selected, load_with_options, words}; use crate::test_types::{ - DedupeRule, DedupeTargetTestCase, NullKeyTargetTestCase, ReplaceKeyTargetTestCase, + DedupeRule, DedupeTargetTestCase, FilesMergeTestCase, NullKeyTargetTestCase, + ReplaceKeyTargetTestCase, }; /// `public.readings`: `s1` and `s2` repeat with their `seen_at` order unlike their `reading_id` order, `s2`'s last row has no `seen_at`, and two rows have no sensor (nor `seen_at`), a repeated NULL key. @@ -304,3 +309,94 @@ pub(crate) async fn run_dedupe_rules( } Ok(rows) } + +/// Run `test_case` against the target connection `target` (its TOML body), with a `loader_schema` of `schema`: CSV files written to a fresh folder with their modified times, loaded step by step; the check query's result. +pub(crate) async fn run_merge_case( + (target, schema): (&str, &str), + test_case: &FilesMergeTestCase, +) -> Result, String> { + let directory = tempfile::tempdir().map_err(|error| error.to_string())?; + let root = directory.path().display().to_string(); + let text = format!( + "[connections.files]\nkind = \"filesystem\"\nroot = \"{root}\"\n\n[connections.target]\n{target}\n\n[targets.dev]\nconnection = \"target\"\nloader_schema = \"{schema}\"\n\n[sources.scores]\nkind = \"files\"\ndescription = \"Test table.\"\nconnection = \"files\"\npath = \"drop/\"\nformat = \"csv\"\n{}\n", + test_case.settings + ); + let project = described_project(&text).await?; + let destinations = connect_targets(&project) + .await + .map_err(|error| error.to_string())?; + let destination = destinations + .get("dev") + .ok_or_else(|| "no dev destination".to_owned())?; + std::fs::create_dir_all(directory.path().join("drop")).map_err(|error| error.to_string())?; + for (index, step) in test_case.steps.iter().enumerate() { + execute( + destination, + step.before.iter().map(|sql| (*sql).to_owned()).collect(), + ) + .await + .map_err(|error| error.to_string())?; + for key in step.removed { + std::fs::remove_file(directory.path().join("drop").join(key)) + .map_err(|error| error.to_string())?; + } + for file in step.files { + let path = directory.path().join("drop").join(file.key); + std::fs::create_dir_all(path.parent().unwrap_or(directory.path())) + .map_err(|error| error.to_string())?; + std::fs::write(&path, file.contents).map_err(|error| error.to_string())?; + let modified = + std::time::UNIX_EPOCH + std::time::Duration::from_secs(file.modified_secs); + std::fs::File::options() + .write(true) + .open(&path) + .and_then(|opened| opened.set_modified(modified)) + .map_err(|error| error.to_string())?; + } + let outcome = run_step(&text, &destinations, (step.loads, step.reload)).await; + let (plan, error) = ( + outcome.clone().unwrap_or_default(), + outcome.err().unwrap_or_default(), + ); + assert!( + words(&error).contains(&words(step.expected_error)) + && error.is_empty() == step.expected_error.is_empty(), + "{} step {}: {error}", + test_case.description, + index + 1 + ); + assert!( + plan.contains(step.expected_plan), + "{} step {}: {plan}", + test_case.description, + index + 1 + ); + } + single_text(destination, test_case.check_query.to_owned()) + .await + .map_err(|error| error.to_string()) +} + +/// Plan and apply once (with `--reload` when `reload`), then load only when `loads` (otherwise as `sqlload apply`): the rendered plan, or the error text. +async fn run_step( + project_text: &str, + destinations: &BTreeMap, + (loads, reload): (bool, bool), +) -> Result { + let project = described_project(project_text).await?; + let options = PlanOptions { + reload, + ..PlanOptions::default() + }; + let mut plan = plan_project(&project, destinations, &options) + .await + .map_err(|error| error.to_string())?; + let (lines, ok) = render_plan(&plan); + ok.then_some(()).ok_or_else(|| lines.join("\n"))?; + apply_project(&plan, destinations) + .await + .map_err(|error| error.to_string())?; + plan.tables.retain(|_| loads); + load_with_options(&project, &plan, destinations, &options).await?; + Ok(lines.join("\n")) +} diff --git a/crates/sqlload/tests/loader/test_after_load.rs b/crates/sqlload/tests/loader/test_after_load.rs index 57c9502c0..9d1811e4d 100644 --- a/crates/sqlload/tests/loader/test_after_load.rs +++ b/crates/sqlload/tests/loader/test_after_load.rs @@ -194,6 +194,105 @@ fn given_after_load_when_files_load_then_only_committed_files_are_deleted_or_mov ], expected_files: &[], }, + AfterLoadTestCase { + description: "--reload of an append table with move_to reads the files it moved too, in order; it moves only those read from path, and the archived ones stay", + settings: "write_strategy = \"append\"\nafter_load = { move_to = \"done\" }\n", + linked: &[], + steps: &[ + AfterLoadStep { + removed: &[], + args: &[], + files: &[("drop/a.csv", "id,v\n1,a\n"), ("drop/b.csv", "id,v\n2,b\n")], + expected_exit: 0, + expected_report: r#"{"after_load":{"action":"move","failed":[],"files":2,"status":"done"},"problems":[],"rows":2,"skipped":[],"status":"loaded"}"#, + }, + AfterLoadStep { + removed: &[], + args: &["--reload", "--yes"], + files: &[("drop/c.csv", "id,v\n3,c\n")], + expected_exit: 0, + expected_report: r#"{"after_load":{"action":"move","failed":[],"files":1,"status":"done"},"problems":[],"rows":3,"skipped":[],"status":"loaded"}"#, + }, + ], + expected_files: &[ + "done/a.csv=id,v\\n1,a\\n", + "done/b.csv=id,v\\n2,b\\n", + "done/c.csv=id,v\\n3,c\\n", + ], + }, + AfterLoadTestCase { + description: "--reload of a merge table with move_to merges the archived files and the new one, then moves only the new one", + settings: "write_strategy = \"merge\"\nunique_key = [\"id\"]\nafter_load = { move_to = \"done\" }\n", + linked: &[], + steps: &[ + AfterLoadStep { + removed: &[], + args: &[], + files: &[("drop/a.csv", "id,v\n1,a\n"), ("drop/b.csv", "id,v\n1,b\n")], + expected_exit: 0, + expected_report: r#"{"after_load":{"action":"move","failed":[],"files":2,"status":"done"},"problems":[],"rows":2,"skipped":[],"status":"loaded"}"#, + }, + AfterLoadStep { + removed: &[], + args: &["--reload", "--yes"], + files: &[("drop/c.csv", "id,v\n2,c\n")], + expected_exit: 0, + expected_report: r#"{"after_load":{"action":"move","failed":[],"files":1,"status":"done"},"problems":[],"rows":3,"skipped":[],"status":"loaded"}"#, + }, + ], + expected_files: &[ + "done/a.csv=id,v\\n1,a\\n", + "done/b.csv=id,v\\n1,b\\n", + "done/c.csv=id,v\\n2,c\\n", + ], + }, + AfterLoadTestCase { + description: "a failed --reload with move_to moves nothing", + settings: "write_strategy = \"merge\"\nunique_key = [\"id\"]\non_duplicate_keys = \"fail\"\nafter_load = { move_to = \"done\" }\n", + linked: &[], + steps: &[ + AfterLoadStep { + removed: &[], + args: &[], + files: &[("drop/a.csv", "id,v\n1,a\n")], + expected_exit: 0, + expected_report: r#"{"after_load":{"action":"move","failed":[],"files":1,"status":"done"},"problems":[],"rows":1,"skipped":[],"status":"loaded"}"#, + }, + AfterLoadStep { + removed: &[], + args: &["--reload", "--yes"], + files: &[("drop/c.csv", "id,v\n3,c\n3,cc\n")], + expected_exit: 1, + expected_report: r#"{"after_load":null,"problems":[],"rows":null,"skipped":[],"status":"failed"}"#, + }, + ], + expected_files: &[ + "done/a.csv=id,v\\n1,a\\n", + "drop/c.csv=id,v\\n3,c\\n3,cc\\n", + ], + }, + AfterLoadTestCase { + description: "--reload with after_load = \"delete\" is still refused, deleting nothing", + settings: "write_strategy = \"append\"\nafter_load = \"delete\"\n", + linked: &[], + steps: &[ + AfterLoadStep { + removed: &[], + args: &[], + files: &[("drop/a.csv", "id,v\n1,a\n")], + expected_exit: 0, + expected_report: r#"{"after_load":{"action":"delete","failed":[],"files":1,"status":"done"},"problems":[],"rows":1,"skipped":[],"status":"loaded"}"#, + }, + AfterLoadStep { + removed: &[], + args: &["--reload", "--yes"], + files: &[("drop/b.csv", "id,v\n2,b\n")], + expected_exit: 1, + expected_report: r#"{"after_load":null,"problems":["--reload would rebuild the table from the files still in its folder, and after_load has deleted or moved the ones already loaded"],"rows":null,"skipped":[],"status":null}"#, + }, + ], + expected_files: &["drop/b.csv=id,v\\n2,b\\n"], + }, AfterLoadTestCase { description: "a failed load deletes nothing", settings: "write_strategy = \"merge\"\nunique_key = [\"id\"]\non_duplicate_keys = \"fail\"\nafter_load = \"delete\"\n", diff --git a/crates/sqlload/tests/loader/test_files_to_duckdb.rs b/crates/sqlload/tests/loader/test_files_to_duckdb.rs index 74f62a431..822120459 100644 --- a/crates/sqlload/tests/loader/test_files_to_duckdb.rs +++ b/crates/sqlload/tests/loader/test_files_to_duckdb.rs @@ -1,14 +1,17 @@ use sqlload_testkit::seaweedfs::main::create_bucket::create_bucket; use sqlload_testkit::seaweedfs::main::start_seaweedfs::start_seaweedfs; -use crate::after_load_helpers::run_after_load_case; +use crate::after_load_helpers::{run_after_load_case, run_archived_reload}; use crate::files_helpers::run_source_row_upgrade; use crate::helpers::{ - local_store, run_async, run_files_scenario, run_large_partition, run_merge_case, - run_replaced_keys, seaweedfs_store, + int96_in_a_later_batch, local_store, run_async, run_file_set, run_files_scenario, + run_large_partition, run_replaced_keys, seaweedfs_store, }; +use crate::merge_helpers::run_merge_case; use crate::test_types::SourceRowUpgradeTestCase; -use crate::test_types::{AfterLoadStep, AfterLoadTestCase}; +use crate::test_types::{ + AfterLoadStep, AfterLoadTestCase, ArchivedReloadTestCase, FileSetTestCase, +}; use crate::test_types::{FilesMergeTestCase, MergeFile, MergeStep, ReplacedKeysTestCase}; use crate::test_types::{FilesRun, FilesScenarioTestCase, LargePartitionTestCase, ParquetFile}; @@ -21,6 +24,9 @@ const NULL_MERGE_CHECK: &str = "SELECT (SELECT string_agg(COALESCE(CAST(id AS VA const NULL_FAIL_SETTINGS: &str = "write_strategy = \"merge\"\nunique_key = [\"id\"]\nnull_keys = \"fail\"\n"; +/// Every row as `id:v`, the run log's events in order (each run with its file), and how many staging or build tables are left. +const RELOAD_CHECK: &str = "SELECT (SELECT string_agg(CAST(id AS VARCHAR) || ':' || v, ',' ORDER BY id) FROM raw.scores) || '|' || (SELECT string_agg(event || coalesce(':' || regexp_extract(json_extract_string(record_json, '$.files[0]'), '([a-z]+[.]csv)', 1), ''), ',' ORDER BY recorded_at) FROM raw._sqlload_runs WHERE table_name = 'scores') || '|' || (SELECT CAST(count(*) AS VARCHAR) FROM information_schema.tables WHERE table_name LIKE '%sqlload_staging%')"; + pub(crate) const MERGE_SETTINGS: &str = "write_strategy = \"merge\"\nunique_key = [\"id\"]\n"; pub(crate) const MERGE_FAIL_SETTINGS: &str = @@ -561,6 +567,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "3 files, one commit each, oldest modified first (ties by path)", }], check_query: MERGE_CHECK, @@ -586,6 +593,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "", }], check_query: MERGE_CHECK, @@ -603,6 +611,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a }], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "merge (unique_key: id; NULLs in keys match)", }], @@ -628,6 +637,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a ], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "dedup_sort ranks repeated keys within one commit, and this table commits once per file (2 files)", }], @@ -647,6 +657,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a }], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "_sqlload_source_row is not written, so a file's rows have no order to pick the last of a repeated key by", }, @@ -659,6 +670,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a }], removed: &[], loads: true, + reload: false, expected_error: "drop/b.csv: the file repeats its unique_key: 1 value of (id) appears in more than one row, such as (1)", expected_plan: "", }, @@ -678,6 +690,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a }], removed: &[], loads: true, + reload: false, expected_error: "dedup_sort breaks ties by row order, which a files merge takes from _sqlload_source_row; stop excluding it", expected_plan: "", }], @@ -696,6 +709,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a }], removed: &[], loads: true, + reload: false, expected_error: "drop/a.csv: the file repeats its unique_key: 1 value of (id) appears in more than one row, such as (5)", expected_plan: "", }], @@ -728,6 +742,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a expected_error: "drop/b.csv: the file repeats its unique_key: 1 value of (id) appears in more than one row, such as (1)", removed: &[], loads: true, + reload: false, expected_plan: "", }, MergeStep { @@ -740,6 +755,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "2 files, one commit each", }, MergeStep { @@ -748,6 +764,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "", }, ], @@ -774,6 +791,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a ], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -782,6 +800,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a files: &[], removed: &[], loads: false, + reload: false, expected_error: "", expected_plan: "", }, @@ -790,6 +809,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a files: &[], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "2 files, one commit each", }, @@ -810,6 +830,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a }], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -822,6 +843,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a }], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "1 file, one commit each", }, @@ -834,6 +856,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a }], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "1 file, one commit each", }, @@ -862,6 +885,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "", }, MergeStep { @@ -874,6 +898,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "1 file, one commit each", }, ], @@ -900,6 +925,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "merge (unique_key: id; NULLs in keys match)", }], check_query: NULL_MERGE_CHECK, @@ -925,6 +951,7 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a expected_error: "drop/b.csv: the file holds a NULL in its unique_key: 1 row(s) have no id, such as (NULL) of (id); null_keys = \"fail\" refuses a NULL in a key column", removed: &[], loads: true, + reload: false, expected_plan: "", }], check_query: NULL_MERGE_CHECK, @@ -948,6 +975,154 @@ fn given_csv_files_when_merging_into_duckdb_then_the_newest_modified_file_wins_a }) } +#[test] +fn given_a_merge_files_table_when_reloading_then_every_file_merges_again_in_order_into_a_new_table_swapped_in_with_its_records() +-> Result<(), String> { + run_async(async { + let test_cases = [ + FilesMergeTestCase { + description: "--reload merges every file oldest first into a new table, records a reset and each file with the swap, and the next load reads none", + settings: MERGE_SETTINGS, + steps: &[ + MergeStep { + before: &[], + files: &[ + MergeFile { + key: "a.csv", + contents: "id,v\n1,a\n2,a\n", + modified_secs: 1_000, + }, + MergeFile { + key: "b.csv", + contents: "id,v\n1,b\n", + modified_secs: 2_000, + }, + ], + removed: &[], + loads: true, + reload: false, + expected_error: "", + expected_plan: "", + }, + MergeStep { + before: &[], + files: &[MergeFile { + key: "c.csv", + contents: "id,v\n1,c\n", + modified_secs: 3_000, + }], + removed: &[], + loads: true, + reload: true, + expected_error: "", + expected_plan: "merge (unique_key: id; NULLs in keys match) into a new table, swapped in at the end (--reload)", + }, + MergeStep { + before: &[], + files: &[], + removed: &[], + loads: true, + reload: false, + expected_error: "", + expected_plan: "up to date: 0 new of 3 files in", + }, + ], + check_query: RELOAD_CHECK, + expected_result: "1:c,2:a|reset,run:a.csv,run:b.csv,reset,run:a.csv,run:b.csv,run:c.csv|0", + }, + FilesMergeTestCase { + description: "a --reload failing on its third file leaves the old table, its run log and no build table", + settings: MERGE_FAIL_SETTINGS, + steps: &[ + MergeStep { + before: &[], + files: &[ + MergeFile { + key: "a.csv", + contents: "id,v\n1,a\n2,a\n", + modified_secs: 1_000, + }, + MergeFile { + key: "b.csv", + contents: "id,v\n1,b\n", + modified_secs: 2_000, + }, + ], + removed: &[], + loads: true, + reload: false, + expected_error: "", + expected_plan: "", + }, + MergeStep { + before: &[], + files: &[MergeFile { + key: "c.csv", + contents: "id,v\n3,c\n3,cc\n", + modified_secs: 3_000, + }], + removed: &[], + loads: true, + reload: true, + expected_error: "drop/c.csv: the file repeats its unique_key", + expected_plan: "", + }, + ], + check_query: RELOAD_CHECK, + expected_result: "1:b,2:a|reset,run:a.csv,run:b.csv|0", + }, + FilesMergeTestCase { + description: "a new file named as one move_to archived merges after the archived copy, so its newer rows win", + settings: "write_strategy = \"merge\"\nunique_key = [\"id\"]\nafter_load = { move_to = \"done\" }\n", + steps: &[ + MergeStep { + before: &[], + files: &[MergeFile { + key: "export.csv", + contents: "id,v\n1,a\n", + modified_secs: 1_000, + }], + removed: &[], + loads: true, + reload: false, + expected_error: "", + expected_plan: "", + }, + MergeStep { + before: &[], + files: &[MergeFile { + key: "export.csv", + contents: "id,v\n1,b\n", + modified_secs: 2_000, + }], + removed: &[], + loads: true, + reload: true, + expected_error: "", + expected_plan: "", + }, + ], + check_query: RELOAD_CHECK, + expected_result: "1:b|reset,run:export.csv,reset,run:export.csv,run:export.csv|0", + }, + ]; + for test_case in &test_cases { + let result = run_merge_case( + ("kind = \"duckdb\"\ndatabase = \":memory:\"", "raw"), + test_case, + ) + .await?; + assert_eq!( + result.as_deref(), + Some(test_case.expected_result), + "{}", + test_case.description + ); + } + Ok(()) + }) +} + #[test] fn given_files_appended_when_their_modified_times_are_shuffled_then_the_load_is_as_before_and_writes_no_run_log() -> Result<(), String> { @@ -972,6 +1147,7 @@ fn given_files_appended_when_their_modified_times_are_shuffled_then_the_load_is_ expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "1 batch, up to 1,000 files or about 1 GB per commit", }], check_query: "SELECT (SELECT string_agg(CAST(id AS VARCHAR) || ':' || v, ',' ORDER BY v) FROM raw.scores) || '|' || (SELECT CAST(count(*) AS VARCHAR) FROM information_schema.tables WHERE table_name = '_sqlload_runs')", @@ -1008,6 +1184,7 @@ fn given_a_resent_older_file_when_delete_insert_follows_a_path_cursor_then_the_t files: &[LE_P01], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "delete_insert (cursor: file_version)", }, @@ -1016,6 +1193,7 @@ fn given_a_resent_older_file_when_delete_insert_follows_a_path_cursor_then_the_t files: &[LE_P02], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -1024,6 +1202,7 @@ fn given_a_resent_older_file_when_delete_insert_follows_a_path_cursor_then_the_t files: &[LE_P01_RESENT], removed: &[], loads: true, + reload: false, expected_error: "le_LE-FY27-P01.csv changed since it was loaded", expected_plan: "", }, @@ -1032,6 +1211,7 @@ fn given_a_resent_older_file_when_delete_insert_follows_a_path_cursor_then_the_t files: &[], removed: &[], loads: true, + reload: false, expected_error: "but its file_version LE-FY27-P01 is below this load's range, which starts at LE-FY27-P02, so delete_insert would skip it and keep its old rows └── fix set replace_key to replace rows by a column the files hold (and drop cursor_column), load from it with --start-cursor LE-FY27-P01, or rebuild the table with --reload", expected_plan: "", }, @@ -1048,6 +1228,7 @@ fn given_a_resent_older_file_when_delete_insert_follows_a_path_cursor_then_the_t files: &[LE_P01_PART], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -1056,6 +1237,7 @@ fn given_a_resent_older_file_when_delete_insert_follows_a_path_cursor_then_the_t files: &[LE_P02_PART], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -1068,6 +1250,7 @@ fn given_a_resent_older_file_when_delete_insert_follows_a_path_cursor_then_the_t }], removed: &[], loads: true, + reload: false, expected_error: "drop/LE-FY27-P01/b.csv is new, but its file_version LE-FY27-P01 is below this load's range, which starts at LE-FY27-P02, and the table holds other files of file_version LE-FY27-P01, so delete_insert would skip it and never load its rows └── fix set replace_key to replace rows by a column the files hold (and drop cursor_column), load from it with --start-cursor LE-FY27-P01, or rebuild the table with --reload", expected_plan: "", }, @@ -1084,6 +1267,7 @@ fn given_a_resent_older_file_when_delete_insert_follows_a_path_cursor_then_the_t files: &[LE_P01_PART], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -1092,6 +1276,7 @@ fn given_a_resent_older_file_when_delete_insert_follows_a_path_cursor_then_the_t files: &[LE_P02_PART], removed: &[], loads: true, + reload: false, expected_error: "the table has no _sqlload_source_version column, so delete_insert cannot find a file changed or added below this load's range, which it skips, keeping the partition's old rows", expected_plan: "", }, @@ -1107,6 +1292,7 @@ fn given_a_resent_older_file_when_delete_insert_follows_a_path_cursor_then_the_t files: &[LE_P01_PART], removed: &[], loads: true, + reload: false, expected_error: "_sqlload_source_file must be loaded: delete_insert uses it to find a file changed or added below this load's range, which it would skip; stop excluding it", expected_plan: "", }], @@ -1145,6 +1331,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[LE_P01], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "delete_insert (replace_key: le_version; NULLs in keys match)", }, @@ -1153,6 +1340,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[LE_P02], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "1 file, one commit each, oldest modified first (ties by path), each recorded in the run log", }, @@ -1161,6 +1349,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[LE_P01_RESENT], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "├── files 1 new of 2 listed", }, @@ -1169,6 +1358,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "reason up to date: 0 new of 2 files in", }, @@ -1184,6 +1374,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[LE_P01_CORRECTED, LE_P02, LE_P01], removed: &[], loads: true, + reload: false, expected_error: "drop/le_LE-FY27-P01_corrected.csv: it holds the replace_key (le_version) value (LE-FY27-P01), which le_LE-FY27-P01.csv also holds, loaded earlier in this run; each file replaces the rows of the values it holds, so it would delete the rows le_LE-FY27-P01.csv loaded, as a value split across part files would keep only its last part. One file must hold all of a value's rows", expected_plan: "", }], @@ -1209,6 +1400,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h ], removed: &[], loads: true, + reload: false, expected_error: "drop/a.csv: it holds the replace_key (le_version) value (LE-FY27-P05), which b.csv also holds, loaded earlier in this run", expected_plan: "", }], @@ -1228,6 +1420,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h }], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -1240,6 +1433,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h }], removed: &["a.csv"], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -1262,6 +1456,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h ], removed: &[], loads: true, + reload: false, expected_error: "drop/b.csv: the file holds a NULL in its replace_key: 1 row(s) have no le_version, such as (NULL) of (le_version); null_keys = \"fail\" refuses a NULL in a key column", expected_plan: "", }], @@ -1277,6 +1472,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[LE_P01_PART_0, LE_P01_PART_1], removed: &[], loads: true, + reload: false, expected_error: "drop/part-0001.csv: it holds the replace_key (le_version) value (LE-FY27-P01), which part-0000.csv also holds, loaded earlier in this run; each file replaces the rows of the values it holds, so it would delete the rows part-0000.csv loaded, as a value split across part files would keep only its last part. One file must hold all of a value's rows. Fix: put all the rows of each (le_version) value in one file, or remove whichever of the two files is outdated from the folder", expected_plan: "", }, @@ -1285,6 +1481,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[], removed: &[], loads: true, + reload: false, expected_error: "drop/part-0001.csv: the table holds rows of its replace_key (le_version) values from part-0000.csv, which is still in the folder; with two files of the folder holding one value, which of them is current can't be told, so replacing the value by part-0001.csv would delete the rows part-0000.csv loaded, as a value split across part files would keep only its last part. One file must hold all of a value's rows. Fix: put all the rows of each (le_version) value in one file, or remove the outdated file, part-0000.csv or part-0001.csv, from the folder", expected_plan: "", }, @@ -1301,6 +1498,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[LE_P01, LE_P02], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -1309,6 +1507,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[LE_P01_CORRECTED], removed: &[], loads: true, + reload: false, expected_error: "drop/le_LE-FY27-P01_corrected.csv: the table holds rows of its replace_key (le_version) values from le_LE-FY27-P01.csv, which is still in the folder", expected_plan: "", }, @@ -1317,6 +1516,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[], removed: &["le_LE-FY27-P01.csv"], loads: true, + reload: false, expected_error: "", expected_plan: "├── files 1 new of 2 listed", }, @@ -1333,6 +1533,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[LE_P01, LE_P02], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -1341,6 +1542,7 @@ fn given_le_files_when_delete_insert_replaces_by_replace_key_then_each_version_h files: &[LE_P01_RESENT], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "", }, @@ -1379,6 +1581,7 @@ fn given_replace_key_when_planning_a_table_it_does_not_suit_then_plan_refuses_it files: &[LE_P01], removed: &[], loads: true, + reload: false, expected_error: "replace_key is for write_strategy = \"delete_insert\", which replaces the rows holding the key values a load reads; merge does not delete by it └── fix set write_strategy = \"delete_insert\" to replace those rows whole, or drop replace_key (merge matches rows one by one on unique_key)", expected_plan: "", }], @@ -1393,6 +1596,7 @@ fn given_replace_key_when_planning_a_table_it_does_not_suit_then_plan_refuses_it files: &[LE_P01], removed: &[], loads: true, + reload: false, expected_error: "replace_key is for write_strategy = \"delete_insert\", which replaces the rows holding the key values a load reads; append does not delete by it └── fix set write_strategy = \"delete_insert\" to replace those rows whole, or drop replace_key", expected_plan: "", }], @@ -1407,6 +1611,7 @@ fn given_replace_key_when_planning_a_table_it_does_not_suit_then_plan_refuses_it files: &[LE_P01], removed: &[], loads: true, + reload: false, expected_error: "replace_key is for write_strategy = \"delete_insert\", which replaces the rows holding the key values a load reads; replace does not delete by it └── fix set write_strategy = \"delete_insert\" to replace those rows whole, or drop replace_key", expected_plan: "", }], @@ -1421,6 +1626,7 @@ fn given_replace_key_when_planning_a_table_it_does_not_suit_then_plan_refuses_it files: &[LE_P01], removed: &[], loads: true, + reload: false, expected_error: "replace_key column version is not a column of the table (not in the source, excluded, not in the enforced contract, or ignored by on_schema_change) └── fix name a column the table loads in replace_key", expected_plan: "", }], @@ -1435,6 +1641,7 @@ fn given_replace_key_when_planning_a_table_it_does_not_suit_then_plan_refuses_it files: &[LE_P01], removed: &[], loads: true, + reload: false, expected_error: "write_strategy = \"delete_insert\" replaces a cursor range or the rows of the key values it reads, so it needs a cursor_column or a replace_key └── fix set cursor_column, or replace_key = [\"version\"]", expected_plan: "", }], @@ -1449,6 +1656,7 @@ fn given_replace_key_when_planning_a_table_it_does_not_suit_then_plan_refuses_it files: &[LE_P01], removed: &[], loads: true, + reload: false, expected_error: "_sqlload_source_file must be loaded: delete_insert uses it to find the files a replace_key value's rows come from; stop excluding it", expected_plan: "", }], @@ -1463,6 +1671,7 @@ fn given_replace_key_when_planning_a_table_it_does_not_suit_then_plan_refuses_it files: &[LE_P01], removed: &[], loads: true, + reload: false, expected_error: "cursor_column file_version does not apply to a files source's delete_insert by replace_key, which reads each new or changed file once, oldest first, from the run log └── fix remove cursor_column", expected_plan: "", }], @@ -1477,6 +1686,7 @@ fn given_replace_key_when_planning_a_table_it_does_not_suit_then_plan_refuses_it files: &[LE_P01], removed: &[], loads: true, + reload: false, expected_error: "cursor_column le_version must be one of the partition_columns or path_columns for a files source; to replace every row of each le_version value the files hold, set replace_key = [\"le_version\"] and drop cursor_column", expected_plan: "", }], @@ -1558,3 +1768,60 @@ fn given_a_partition_cursor_without_the_source_file_column_when_reloading_then_o Ok(()) }) } + +#[test] +fn given_a_later_batch_holding_a_column_duckdb_cannot_read_exactly_when_duckdb_reads_the_files_itself_then_the_load_stops_there() +-> Result<(), String> { + let directory = tempfile::tempdir().map_err(|error| error.to_string())?; + let test_cases = [FileSetTestCase { + description: "a batch of INT96 under embedded microseconds, which plan's sample of the newest files never saw, is refused by its own footers; the batch before it stays loaded", + settings: "write_strategy = \"append\"\n", + file_set: int96_in_a_later_batch, + expected_error: "DuckDB cannot load this batch's files exactly by reading them itself, and they never go through SQLLoad instead: column at is held as", + check_query: "SELECT CAST(count(*) AS VARCHAR) FROM raw.scores", + expected_result: "1000", + }]; + run_async(async { + for test_case in &test_cases { + assert_eq!( + run_file_set(directory.path(), test_case).await?.as_deref(), + Some(test_case.expected_result), + "{}", + test_case.description + ); + } + Ok(()) + }) +} + +#[test] +fn given_an_append_parquet_table_with_move_to_when_reloading_then_each_archived_file_is_read_from_where_it_was_moved() +-> Result<(), String> { + let test_cases = [ArchivedReloadTestCase { + description: "a.parquet loaded and moved to done/, then a new drop/a.parquet: the reload reads each copy once, the archived one from done/", + first: &[ParquetFile { + key: "a.parquet", + ids: &[1], + score: 1.0, + extra: None, + }], + later: &[ParquetFile { + key: "a.parquet", + ids: &[2], + score: 2.0, + extra: None, + }], + expected_plan: "route read parquet into Arrow -> DuckDB takes Arrow (--reload also reads the files after_load moved to", + expected_rows: "1:1.0,2:2.0", + }]; + for test_case in &test_cases { + let (plan, rows) = run_archived_reload(test_case)?; + assert_eq!(rows, test_case.expected_rows, "{}", test_case.description); + assert!( + plan.contains(test_case.expected_plan), + "{}:\n{plan}", + test_case.description + ); + } + Ok(()) +} diff --git a/crates/sqlload/tests/loader/test_postgres_target.rs b/crates/sqlload/tests/loader/test_postgres_target.rs index 0a6b561f0..5ba72bde3 100644 --- a/crates/sqlload/tests/loader/test_postgres_target.rs +++ b/crates/sqlload/tests/loader/test_postgres_target.rs @@ -5,7 +5,8 @@ use sqlload_testkit::postgres::main::connect::connect; use sqlload_testkit::postgres::main::start_postgres::start_postgres; use sqlload_testkit::postgres::main::start_postgres_tls::start_postgres_tls; -use crate::helpers::{described_project, run_async, run_merge_case, run_sqlload, words}; +use crate::helpers::{described_project, run_async, run_sqlload, words}; +use crate::merge_helpers::run_merge_case; use crate::postgres_helpers::{load_once, project_text, text}; use crate::test_files_to_duckdb::{ MERGE_CHECK, MERGE_FAIL_SETTINGS, MERGE_SETTINGS, MERGE_SORT_SETTINGS, @@ -339,6 +340,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "3 files, one commit each, oldest modified first (ties by path)", }], check_query: MERGE_CHECK, @@ -364,6 +366,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "", }], check_query: MERGE_CHECK, @@ -381,6 +384,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins }], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "merge (unique_key: id; NULLs in keys match)", }], @@ -406,6 +410,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins ], removed: &[], loads: true, + reload: false, expected_error: "", expected_plan: "dedup_sort ranks repeated keys within one commit, and this table commits once per file (2 files)", }], @@ -424,6 +429,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins }], removed: &[], loads: true, + reload: false, expected_error: "drop/a.csv: the file repeats its unique_key: 1 value of (id) appears in more than one row, such as (5)", expected_plan: "", }], @@ -456,6 +462,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins expected_error: "drop/b.csv: the file repeats its unique_key: 1 value of (id) appears in more than one row, such as (1)", removed: &[], loads: true, + reload: false, expected_plan: "", }, MergeStep { @@ -468,6 +475,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "2 files, one commit each", }, MergeStep { @@ -476,6 +484,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "", }, ], @@ -503,6 +512,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "", }, MergeStep { @@ -515,6 +525,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "1 file, one commit each", }, ], @@ -539,6 +550,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "merge (unique_key: id; repeated keys fail; NULLs in keys match)", }, MergeStep { @@ -551,6 +563,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins expected_error: "", removed: &[], loads: true, + reload: false, expected_plan: "1 file, one commit each", }, ], @@ -573,6 +586,7 @@ fn given_csv_files_when_merging_into_postgres_then_the_newest_modified_file_wins expected_error: "create one that treats them as equal, CREATE UNIQUE INDEX CONCURRENTLY ON \"raw\".\"scores\" (\"id\") NULLS NOT DISTINCT, then ALTER TABLE raw.scores DROP CONSTRAINT scores_id_key", removed: &[], loads: true, + reload: false, expected_plan: "", }], check_query: "SELECT CAST(count(*) AS VARCHAR) FROM raw.scores", diff --git a/crates/sqlload/tests/loader/test_types.rs b/crates/sqlload/tests/loader/test_types.rs index 7aafd911c..cb33aea81 100644 --- a/crates/sqlload/tests/loader/test_types.rs +++ b/crates/sqlload/tests/loader/test_types.rs @@ -1038,6 +1038,8 @@ pub(crate) struct MergeStep { pub(crate) removed: &'static [&'static str], /// Whether the step loads; `false` plans and applies only, as `sqlload apply` does. pub(crate) loads: bool, + /// Whether the step plans and loads with `--reload`. + pub(crate) reload: bool, /// Text the load's error must contain; empty when it must succeed. pub(crate) expected_error: &'static str, /// Text the plan must contain when the load succeeds. @@ -1307,3 +1309,27 @@ pub(crate) struct BigQueryNullCursorTestCase { /// The loaded labels, NULL as `NULL`. pub(crate) expected_values: &'static str, } + +/// A Parquet files source into DuckDB loaded once from a file set: the load's error and the rows it left. +pub(crate) struct FileSetTestCase { + pub(crate) description: &'static str, + pub(crate) settings: &'static str, + pub(crate) file_set: SourceFiles, + /// Text the load's error must contain; empty when it must succeed. + pub(crate) expected_error: &'static str, + pub(crate) check_query: &'static str, + pub(crate) expected_result: &'static str, +} + +/// An append Parquet source under `drop/` with `after_load = { move_to = "done" }`, loaded into DuckDB by the executable, then given `later` files and loaded again with `--reload --yes`. +pub(crate) struct ArchivedReloadTestCase { + pub(crate) description: &'static str, + /// The files the first load reads and moves to `done/`, by key below `drop/`. + pub(crate) first: &'static [ParquetFile], + /// The files written to `drop/` before the reload. + pub(crate) later: &'static [ParquetFile], + /// Text the reload's plan must contain. + pub(crate) expected_plan: &'static str, + /// Every row as `id:score`, in order, after the reload. + pub(crate) expected_rows: &'static str, +} diff --git a/packages/sqlload/python/sqlload/__init__.py b/packages/sqlload/python/sqlload/__init__.py index a0f4abbdd..376590efa 100644 --- a/packages/sqlload/python/sqlload/__init__.py +++ b/packages/sqlload/python/sqlload/__init__.py @@ -49,6 +49,7 @@ def orders(ctx: sl.Context): single, ) from sqlload._declare import Context, Source, Table, source, table +from sqlload._transform import TransformContext, transform __version__: str = version("sqlload") @@ -59,6 +60,7 @@ def orders(ctx: sl.Context): "Job", "Source", "Table", + "TransformContext", "__version__", "api_key", "basic", @@ -83,6 +85,7 @@ def orders(ctx: sl.Context): "single", "source", "table", + "transform", ] diff --git a/packages/sqlload/python/sqlload/_http.py b/packages/sqlload/python/sqlload/_http.py index 58f405262..065992217 100644 --- a/packages/sqlload/python/sqlload/_http.py +++ b/packages/sqlload/python/sqlload/_http.py @@ -6,19 +6,20 @@ """ import json -from collections.abc import Iterable, Iterator, Mapping, Sequence +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from typing import Any from sqlload._data import Field, Fields, field_list from sqlload._declare import Source +from sqlload._transform import TRANSFORMS #: Default retries of a request that failed with a transport error, 429 or a 5xx. DEFAULT_RETRIES = 5 ROW_LOOP_HINT = ( "rows of an sl.http job stay in Rust and are never handed to Python one at a time (rule python_row_loop); " - "return or yield the job itself, and filter the rows downstream" + "return or yield the job itself, and filter the rows downstream, or reshape each page with transform=" ) @@ -194,11 +195,18 @@ def _param_text(value: Any) -> str: class Job: """An HTTP job definition; a table function returns or yields it, and Rust runs it.""" - def __init__(self, definition: dict[str, Any]) -> None: + def __init__(self, definition: dict[str, Any], transform: Callable[..., Any] | str | None = None) -> None: self._definition = definition + #: Kept as given, and looked up when the job is handed to Rust: a job built at module level outlives the run that + #: first registered a function passed as itself, and each run lets go of those when it finishes. + self._transform = transform def __sqlload_http_job__(self) -> str: - return json.dumps(self._definition, default=_param_text) + definition = self._definition + if self._transform is not None: + function, name = TRANSFORMS.key(self._transform) + definition = {**definition, "transform": {"function": function, "name": name}} + return json.dumps(definition, default=_param_text) def __iter__(self) -> Iterator["Job"]: """`yield from job` hands the job itself to Rust.""" @@ -253,6 +261,7 @@ def paginate( key: str | None = None, cells: Cells | None = None, select_missing: str = "fail", + transform: Callable[..., Any] | str | None = None, ) -> Job: """Every page of `path`; `select` (such as `"$.items"`) is where the rows are, `on_status={404: "skip"}` treats a status as no rows. @@ -262,9 +271,22 @@ def paginate( `fields` loads only those values of each row, renamed (`{"account_code": "$.code"}`); `keep_record` adds each row's whole object as a JSON column. `method="POST"` sends `json_body` as JSON with every page, each page's copy holding the paginator's values when its `location` is `"body"`. As a fan-out's `over`, `key` names the - placeholder the selected values fill when it is not the field's own name (`select="$.items[*].id", key="order_id"`).""" + placeholder the selected values fill when it is not the field's own name (`select="$.items[*].id", key="order_id"`). + + `transform` (a function, or the name `@sl.transform` registered one under) is called once per page, in page + order, on the page's rows after `select`, `fields` and `cells`, as `f(rows, ctx)`; what it returns or yields loads.""" request = self._request(path, paginator, params, select, on_status, (method, json_body)) - return Job({"kind": "paginate", **request, **_projection(fields, keep_record), **_key(key), **_cells(cells), **_select_missing(select_missing)}) + return Job( + { + "kind": "paginate", + **request, + **_projection(fields, keep_record), + **_key(key), + **_cells(cells), + **_select_missing(select_missing), + }, + _transform(transform), + ) def fan_out( self, @@ -285,6 +307,7 @@ def fan_out( batch: int | None = None, select_missing: str = "fail", key_path: str | Field | None = None, + transform: Callable[..., Any] | str | None = None, ) -> Job: """Every page of `path` once per key of `over`, up to `concurrency` at a time: `over` is a paginate whose `select` picks the keys (`"$.items[*].id"` fills `{id}`), or values (`["uk", "fr"]` fill the one placeholder; @@ -302,9 +325,14 @@ def fan_out( `key_path="$.issueId"` (from each record, or `sl.parent("$.issueId")` from an object around it) fails the table when a record's key is not among the keys its request sent (with `batch=`, the batch's), catching an - API that answers for other keys.""" + API that answers for other keys. + + `transform` is as for `paginate`, once per page of each key; its `ctx.key` holds the key's fields (`None` with + `batch=`).""" if isinstance(over, (str, bytes)): raise TypeError("fan_out over= takes a paginate or fan_out job or a list of values, not one string") + if isinstance(over, Job) and over._transform is not None: + raise TypeError("transform is for the job whose rows load, not for the job given as fan_out's over=") over_spec = dict(over._definition) if isinstance(over, Job) else {"kind": "values", "values": list(over)} request = self._request(path, paginator or single(), params, select, on_status, (method, json_body)) projected = _projection(fields, keep_record) @@ -321,7 +349,8 @@ def fan_out( "over": over_spec, "concurrency": concurrency, **batched, - } + }, + _transform(transform), ) def _request( @@ -376,6 +405,13 @@ def _key_path(key_path: str | Field | None) -> dict[str, Any]: raise TypeError(f'key_path takes a path such as "$.issueId" or sl.parent(...), not {key_path!r}') +def _transform(transform: Callable[..., Any] | str | None) -> Callable[..., Any] | str | None: + """`transform` as given, refused here when it is neither a function nor a name.""" + if transform is not None and not isinstance(transform, str) and not callable(transform): + raise TypeError(f"expected a function or the name @sl.transform registered one under, not {type(transform).__name__}") + return transform + + def _projection(fields: Fields | None, keep_record: str | None) -> dict[str, Any]: listed = field_list(fields) return { diff --git a/packages/sqlload/python/sqlload/_registry.py b/packages/sqlload/python/sqlload/_registry.py new file mode 100644 index 000000000..9e34ae10f --- /dev/null +++ b/packages/sqlload/python/sqlload/_registry.py @@ -0,0 +1,103 @@ +"""Functions a project picks by name: each decorator (`@sl.transform`) registers functions in a registry of its own. + +A function registers under its own name, or the one the decorator is given; `sqlload.toml` names it. Discovery imports +every module under `python_paths`, which runs the decorators, so `plan` can look a name up and refuse one that is +unknown or registered twice. +""" + +import inspect +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class Registered: + """A function registered under a name, and the file it is defined in.""" + + function: Callable[..., Any] + file: str + + +class Registry: + """The functions one decorator registered, by name.""" + + def __init__(self, decorator: str) -> None: + self.decorator = decorator + self._named: dict[str, list[Registered]] = {} + #: Functions passed as themselves (`transform=page`) rather than by a registered name, by the key they go by, with + #: the name they are shown by. Held until the run finishes (`release`): a job keeps only the key, so an inline lambda + #: would otherwise be gone before its pages are called. + self._passed: dict[str, tuple[Callable[..., Any], str]] = {} + + def register(self, name: str, function: Callable[..., Any]) -> None: + """Register `function` under `name`; the same function registered again (its module imported again) replaces itself.""" + if not callable(function): + raise TypeError(f"{self.decorator} decorates a function, not {type(function).__name__}") + if not name: + raise ValueError(f"{self.decorator} needs a name that is not empty") + file = _file(function) + kept = [entry for entry in self._named.get(name, []) if (entry.file, entry.function.__qualname__) != (file, function.__qualname__)] + self._named[name] = [*kept, Registered(function, file)] + + def key(self, function: Callable[..., Any] | str) -> tuple[str, str]: + """The key a function goes by, and the name it is shown by: a registered name as it is; a function itself, a key + of its own that resolves to that very function (even when another function is registered under its name), shown + by the name it is registered under, else its qualified name.""" + if isinstance(function, str): + return function, function + if not callable(function): + raise TypeError(f"expected a function or the name {self.decorator} registered one under, not {type(function).__name__}") + registered = (name for name, entries in self._named.items() if any(entry.function is function for entry in entries)) + shown = next(registered, getattr(function, "__qualname__", type(function).__name__)) + key = f"{shown}#{id(function):x}" + self._passed[key] = (function, shown) + return key, shown + + def release(self) -> None: + """Let go of the functions passed as themselves, once the run that passed them finished.""" + self._passed.clear() + + def lookup(self, key: str, root: Path | None = None) -> tuple[Callable[..., Any] | None, str | None]: + """The function registered under `key` (or passed as itself), or why there is none: no function has the name (the + registered ones listed), or two have it (their files named, relative to `root`).""" + if key in self._passed: + return self._passed[key][0], None + entries = self._named.get(key, []) + if len(entries) == 1: + return entries[0].function, None + if entries: + files = " and ".join(f"{_shown(entry.file, root)} ({entry.function.__qualname__})" for entry in entries) + return None, f'{len(entries)} functions are registered as {key}, in {files}; give one another name with {self.decorator}("...")' + names = ", ".join(sorted(self._named)) + if not names: + return None, f'no function is registered as {key}: decorate one with {self.decorator}, or {self.decorator}("{key}") to give it that name' + return None, f"no function is registered as {key} with {self.decorator}; the registered ones are {names}" + + def shown(self, key: str) -> str: + """The name `key` is shown by in errors: a passed function's, else the key itself.""" + passed = self._passed.get(key) + return passed[1] if passed is not None else key + + def source_file(self, key: str) -> str | None: + """The file the function `key` names is defined in.""" + function, _ = self.lookup(key) + return _file(function) if function is not None else None + + +def _file(function: Callable[..., Any]) -> str: + """The file `function` is defined in, resolved; empty when Python can't tell.""" + try: + path = inspect.getsourcefile(function) or inspect.getfile(function) + except TypeError: + return "" + return str(Path(path).resolve()) + + +def _shown(file: str, root: Path | None) -> str: + """`file` relative to the project folder `root` when it is under it.""" + path = Path(file) + if root is not None and root in path.parents: + return path.relative_to(root).as_posix() + return file diff --git a/packages/sqlload/python/sqlload/_runtime.py b/packages/sqlload/python/sqlload/_runtime.py index 5e42165a6..7569366f9 100644 --- a/packages/sqlload/python/sqlload/_runtime.py +++ b/packages/sqlload/python/sqlload/_runtime.py @@ -4,12 +4,14 @@ `load` and the other commands; they are not a public API. """ +import hashlib import importlib import json import logging import sys import threading import traceback +import types from collections.abc import Iterator from datetime import date, datetime, time, timedelta, timezone from decimal import Decimal @@ -18,6 +20,7 @@ from sqlload._declare import Connection, Context, Source, Table from sqlload._http import Job +from sqlload._transform import TRANSFORMS, TransformContext, returned_items #: Folders under `python_paths` never imported: virtual environments, build output and caches. SKIPPED_FOLDERS: frozenset[str] = frozenset({"venv", "build", "dist", "site-packages", "node_modules"}) @@ -25,7 +28,12 @@ #: The destination name of a table declared without a source: its own name. STANDALONE_TABLE_NAME = "{table}" +#: How long a transform's version hash is, in hex digits. +VERSION_DIGITS = 16 + _tables: dict[str, Table] = {} +#: The project folder `discover` imported from, which errors name files relative to. +_root: list[Path] = [] _contexts: dict[str, Context] = {} _clients: dict[str, Any] = {} #: Held while a client is built, since tables of one source load at the same time. @@ -53,6 +61,7 @@ def discover(root: str, python_paths: list[str]) -> str: standalone.setdefault(id(value), (value, relative)) elif isinstance(value, Table) and value.source is not None: sources.setdefault(id(value.source), (value.source, relative)) + _root[:] = [root_path] _tables.clear() declared: list[dict[str, Any]] = [] for source, relative in sources.values(): @@ -143,8 +152,43 @@ def table_state(table: str) -> str | None: return final if final != ctx.started_state else None +def transform_version(name: str) -> tuple[str | None, str | None]: + """Look up the transform registered as `name` and check it can run. + + Returns its version, `name@` and a hash of the file it is defined in (an edit to the file, not only to the function, + changes it; when the file can't be read, as for a function defined in a notebook, a hash of the function's code), or + why it cannot run: no function or two are registered as `name`, or pyarrow is missing. + """ + function, problem = TRANSFORMS.lookup(name, _root[0] if _root else None) + if function is None: + return None, problem + try: + importlib.import_module("pyarrow") + except ImportError: + return None, "transform needs pyarrow: pip install pyarrow" + source = TRANSFORMS.source_file(name) or "" + contents = Path(source).read_bytes() if source and Path(source).is_file() else _code_bytes(function) + return f"{name}@{hashlib.sha256(contents).hexdigest()[:VERSION_DIGITS]}", None + + +def call_transform(name: str, table: Any, context: str) -> tuple[Any, str | None]: + """Call the transform `name` on one unit's rows, a pyarrow Table, with its `ctx` (JSON). + + Returns what crosses, item by item as a generator yields them (its `dicts` says whether any were dicts), or why the + returned value cannot load. + """ + function, problem = TRANSFORMS.lookup(name) + if function is None: + raise LookupError(problem) + fields: dict[str, Any] = json.loads(context) + modified_at = fields.get("modified_at") + ctx = TransformContext(**{**fields, "modified_at": datetime.fromisoformat(_offset_text(modified_at)) if modified_at else None}) + return returned_items(TRANSFORMS.shown(name), function(table, ctx)) + + def finish_run() -> None: """Close every client the run built (their `close()`, when they have one).""" + TRANSFORMS.release() clients = list(_clients.values()) _clients.clear() for client in clients: @@ -285,3 +329,24 @@ def _timestamp(text: str, *, aware: bool) -> datetime: if not aware: return value.replace(tzinfo=None) return value.astimezone(timezone.utc) if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) + + +def _code_bytes(function: Any) -> bytes: + """What a function does, for its version when its file can't be read: its bytecode, constants (nested functions' code + included) and the names it uses, without its file name or line numbers. A callable without code, by its type's name.""" + code = getattr(function, "__code__", None) + return _code_parts(code) if isinstance(code, types.CodeType) else type(function).__qualname__.encode() + + +def _code_parts(code: types.CodeType) -> bytes: + constants = [_code_parts(constant) if isinstance(constant, types.CodeType) else _constant_text(constant) for constant in code.co_consts] + return b"\0".join([code.co_code, repr(code.co_names).encode(), *constants]) + + +def _constant_text(constant: Any) -> bytes: + """A code constant as text that is the same in every process: a frozenset's items sorted, as its order follows string hashing.""" + if isinstance(constant, frozenset): + return repr(sorted(repr(item) for item in constant)).encode() + if isinstance(constant, tuple): + return b"(" + b",".join(_constant_text(item) for item in constant) + b")" + return repr(constant).encode() diff --git a/packages/sqlload/python/sqlload/_transform.py b/packages/sqlload/python/sqlload/_transform.py new file mode 100644 index 000000000..5f726a681 --- /dev/null +++ b/packages/sqlload/python/sqlload/_transform.py @@ -0,0 +1,113 @@ +"""`@sl.transform`: a function that reshapes each file, Google Sheets range or `sl.http` page before it loads. + +`sqlload.toml` picks one by the name it is registered under (`transform = "kpi_sheet"`); an `sl.http` job takes the name +or the function itself (`api.paginate(..., transform=kpi_page)`). SQLLoad reads the unit, calls `f(rows, ctx)` with its +rows as a pyarrow Table, and loads what the function returns or yields. +""" + +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from datetime import datetime +from typing import Any, TypeVar, overload + +from sqlload._registry import Registry + +_Function = TypeVar("_Function", bound=Callable[..., Any]) + +#: The functions `@sl.transform` registered, by name. +TRANSFORMS = Registry("@sl.transform") + +#: What a transform may return or yield, for its errors. +RETURNS = "dicts, lists of dicts, a pyarrow Table, or a polars or pandas DataFrame" + + +@overload +def transform(function: _Function, /) -> _Function: ... + + +@overload +def transform(name: str | None = None, /) -> Callable[[_Function], _Function]: ... + + +def transform(function: Any = None, /) -> Any: + """Register a transform under its own name (`@sl.transform`) or another (`@sl.transform("kpis")`), for + `transform = "..."` in `sqlload.toml` or `transform=` on an `sl.http` job to pick. Returns the function unchanged.""" + if function is None or isinstance(function, str): + name = function + + def register(decorated: _Function) -> _Function: + TRANSFORMS.register(name or decorated.__name__, decorated) + return decorated + + return register + TRANSFORMS.register(function.__name__, function) + return function + + +@dataclass(frozen=True) +class TransformContext: + """What a transform receives with the rows: the unit they came from. + + `unit` is `"file"`, `"range"` or `"page"`. A file sets `path` (as `_sqlload_source_file` names it, + `archive.zip/member.csv` for a ZIP member), `modified_at`, and `sheet`, `range` or `skip_rows` when its format uses + them. A Google Sheets range sets `spreadsheet_id`, `sheet` (the tab) and `range`. An `sl.http` page sets `page` + (counted from 1 within its request) and, in a fan-out, `key`: the fields of the key it was requested for. + """ + + unit: str + table: str + path: str | None = None + modified_at: datetime | None = None + sheet: str | None = None + range: str | None = None + skip_rows: int | None = None + spreadsheet_id: str | None = None + page: int | None = None + key: dict[str, Any] | None = None + + +class Returned: + """What a transform returned or yielded, one item at a time, each checked; `dicts` says whether any were dicts.""" + + def __init__(self, name: str, items: Iterator[Any]) -> None: + self._name = name + self._items = items + self.dicts = False + + def __iter__(self) -> "Returned": + return self + + def __next__(self) -> Any: + item = next(self._items) + if isinstance(item, dict) or (isinstance(item, list) and all(isinstance(row, dict) for row in item)): + self.dicts = True + return item + if _arrow(item): + return item + raise TypeError(f"transform {self._name} yielded {_kind(item)}; yield {RETURNS}") + + def close(self) -> None: + """Close a generator the load stopped reading, so its `finally` blocks run now.""" + close = getattr(self._items, "close", None) + if callable(close): + close() + + +def returned_items(name: str, returned: Any) -> tuple[Returned | None, str | None]: + """What `returned` crosses as, or why it can't load: a list of dicts, a dict, Arrow data or a DataFrame crosses as one + item; an iterator (a generator) crosses item by item as it yields.""" + if isinstance(returned, list) and all(isinstance(row, dict) for row in returned): + return Returned(name, iter([returned] if returned else [])), None + if isinstance(returned, dict) or _arrow(returned): + return Returned(name, iter([returned])), None + if isinstance(returned, Iterator): + return Returned(name, returned), None + return None, f"it returned {_kind(returned)}; return or yield {RETURNS}" + + +def _arrow(value: Any) -> bool: + return hasattr(value, "__arrow_c_stream__") or hasattr(value, "__arrow_c_array__") + + +def _kind(value: Any) -> str: + return "a list holding something other than dicts" if isinstance(value, list) else type(value).__name__ diff --git a/packages/sqlload/tests/test_http_jobs.py b/packages/sqlload/tests/test_http_jobs.py index 983397fbc..547ceb693 100644 --- a/packages/sqlload/tests/test_http_jobs.py +++ b/packages/sqlload/tests/test_http_jobs.py @@ -129,6 +129,9 @@ def do_GET(self) -> None: if url.path.startswith("/v1/tree/"): return self.tree(url.path, query) parts = url.path.split("/") + if len(parts) == 5 and parts[2] == "regions" and parts[4] == "orders": + page = int(query["p"]) + return self.reply(200, {"items": [dict(row, region=parts[3]) for row in ORDERS[(page - 1) * 3 : page * 3]]}) if len(parts) == 5 and parts[2] == "orders" and parts[4] == "items": order = int(parts[3]) return self.reply(404, {}) if order == 4 else self.reply(200, {"items": ITEMS[order]}) @@ -434,13 +437,11 @@ def make_project(tmp_path: Path, server: ThreadingHTTPServer, tables: str) -> Pa return tmp_path -def sqlload(project: Path, *args: str) -> subprocess.CompletedProcess[str]: +def sqlload(project: Path, *args: str, script: str = "from sqlload._cli import main; main()") -> subprocess.CompletedProcess[str]: home = project / "home" home.mkdir(exist_ok=True) env = {**os.environ, "HOME": str(home), "XDG_CONFIG_HOME": str(home / ".config"), "SHOP_TOKEN": TOKEN, "SHOP_CLIENT_SECRET": CLIENT_SECRET, "SHOP_PASSWORD": PASSWORD} - return subprocess.run( - [sys.executable, "-c", "from sqlload._cli import main; main()", *args], cwd=project, env=env, capture_output=True, text=True, timeout=120 - ) + return subprocess.run([sys.executable, "-c", script, *args], cwd=project, env=env, capture_output=True, text=True, timeout=120) def query(project: Path, sql: str) -> list[tuple[Any, ...]]: @@ -506,6 +507,83 @@ def test_given_a_fan_out_when_loading_then_each_key_is_requested_and_a_skipped_s assert requested == list(range(1, 8)) +#: A page transform logging each call to `calls.log`, then yielding the page's rows one at a time with where they came from. +PAGE_TRANSFORM = """ +@sl.transform +def tag_page(rows, ctx): + with open("calls.log", "a") as log: + log.write(f"{ctx.key['region']}:{ctx.page}\\n") + for position, row in enumerate(rows.to_pylist()): + yield {**row, "page": ctx.page, "position": position, "unit": ctx.unit, "key_region": ctx.key["region"]} + +""" + + +@pytest.mark.parametrize("transform", ["tag_page", "'tag_page'"], ids=["the-function", "its-registered-name"]) +def test_given_a_fan_out_with_a_page_transform_when_loading_then_each_page_goes_through_it_once_in_page_order( + tmp_path: Path, api: ThreadingHTTPServer, transform: str +) -> None: + job = f"api.fan_out('/regions/{{region}}/orders', over=['north', 'south'], concurrency=2, paginator=sl.page_number('p'), select='$.items', transform={transform})" + project = make_project(tmp_path, api, PAGE_TRANSFORM + f"@shop.table(description='Orders by region.')\ndef orders(ctx):\n return {job}\n") + done = sqlload(project, "load") + assert done.returncode == 0, done.stdout + done.stderr + calls = (project / "calls.log").read_text().split() + # Page 4 is the empty page that ends page-number paging: a page all the same, called with no rows. + assert sorted(calls) == [f"{region}:{page}" for region in ("north", "south") for page in range(1, 5)], calls + rows = query(project, 'SELECT region, key_region, unit, page, position, order_id FROM "RAW".shop__orders ORDER BY rowid') + assert {(region, key_region, unit) for region, key_region, unit, *_ in rows} == {("north", "north", "page"), ("south", "south", "page")} + for region in ("north", "south"): + loaded = [(page, position, order_id) for row_region, _, _, page, position, order_id in rows if row_region == region] + assert loaded == sorted(loaded), loaded + assert [order_id for *_, order_id in loaded] == [row["order_id"] for row in ORDERS], loaded + + +def test_given_a_module_level_job_with_a_page_transform_when_loading_twice_in_one_process_then_both_loads_call_it( + tmp_path: Path, api: ThreadingHTTPServer +) -> None: + module = """ +@sl.transform +def clean_page(rows, ctx): + with open("calls.log", "a") as log: + log.write(f"{ctx.page}\\n") + return rows + +orders_job = api.paginate('/next', paginator=sl.next_link('$.next'), select='$.data.items', transform=clean_page) + +@shop.table(description='Orders.') +def orders(ctx): + return orders_job +""" + project = make_project(tmp_path, api, module) + # An orchestrator runs commands in its own process, where the module (and its job) stays imported between them. + twice = "from sqlload import _native; print('exit codes', [_native.run_cli(['sqlload', 'load']) for _ in range(2)])" + done = sqlload(project, script=twice) + assert "exit codes [0, 0]" in done.stdout, done.stdout + done.stderr + assert query(project, 'SELECT order_id, status FROM "RAW".shop__orders ORDER BY order_id') == ALL_ORDERS + pages = (project / "calls.log").read_text().split() + assert len(pages) % 2 == 0 and pages[: len(pages) // 2] == pages[len(pages) // 2 :], pages + + +def test_given_a_page_transform_that_raises_when_loading_then_the_table_fails_naming_the_page(tmp_path: Path, api: ThreadingHTTPServer) -> None: + module = """ +@sl.transform +def strict_page(rows, ctx): + if ctx.page == 2: + raise ValueError("page 2 is malformed") + return rows + +@shop.table(description='Orders.') +def orders(ctx): + return api.paginate('/pages', paginator=sl.page_number('p'), select='$.items', transform=strict_page) +""" + project = make_project(tmp_path, api, module) + done = sqlload(project, "load") + assert done.returncode == 1, done.stdout + done.stderr + output = done.stdout + done.stderr + assert "GET http://" in output and "/v1/pages?p=2: transform strict_page failed: ValueError: page 2 is malformed" in output, output + assert query(project, 'SELECT count(*) FROM "RAW".shop__orders') == [(0,)] + + @pytest.mark.parametrize( ("tables", "expected_text"), [ diff --git a/packages/sqlload/tests/test_transforms.py b/packages/sqlload/tests/test_transforms.py new file mode 100644 index 000000000..1b342b527 --- /dev/null +++ b/packages/sqlload/tests/test_transforms.py @@ -0,0 +1,573 @@ +"""Transforms end to end: `transform` on files tables (each file or ZIP member) and on `google_sheets` tables (each range), into DuckDB.""" + +import gc +import importlib.util +import json +import os +import subprocess +import sys +import textwrap +import threading +import time +import weakref +import zipfile +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any + +import duckdb +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +PROJECT = """\ +python_paths = ["transforms"] + +[connections.wh] +kind = "duckdb" +database = "wh.duckdb" + +[connections.files] +kind = "filesystem" +root = "." + +[targets.dev] +connection = "wh" +loader_schema = "RAW" + +[sources.metrics] +kind = "files" +connection = "files" +path = "drop/" +format = "{format}" +description = "Metrics, one file a month." +{settings} +[sources.metrics.tables.metrics] +transform = "metric_rows" +{table} +""" + +#: A transform, `metric_rows`, returning each row of a CSV with `id` and `v` as dicts, a pyarrow Table or a polars DataFrame, +#: or yielding them row by row and as Arrow, with the file it came from. +RETURNS = { + "dicts": """ +import sqlload as sl + +@sl.transform +def metric_rows(rows, ctx): + return [{"id": int(row["id"]), "v": row["v"], "file": ctx.path} for row in rows.to_pylist()] +""", + "arrow": """ +import pyarrow as pa +import pyarrow.compute as pc +import sqlload as sl + +@sl.transform +def metric_rows(rows, ctx): + return pa.table({"id": pc.cast(rows["id"], pa.int64()), "v": rows["v"], "file": [ctx.path] * rows.num_rows}) +""", + "polars": """ +import polars as pl +import sqlload as sl + +@sl.transform +def metric_rows(rows, ctx): + return pl.from_arrow(rows).select(pl.col("id").cast(pl.Int64), pl.col("v"), pl.lit(ctx.path).alias("file")) +""", + "generator": """ +import pyarrow as pa +import sqlload as sl + +@sl.transform +def metric_rows(rows, ctx): + assert ctx.unit == "file", ctx.unit + listed = rows.to_pylist() + for row in listed[:1]: + yield {"id": int(row["id"]), "v": row["v"], "file": ctx.path} + rest = listed[1:] + if rest: + yield pa.table({"id": [int(row["id"]) for row in rest], "v": [row["v"] for row in rest], "file": [ctx.path] * len(rest)}) +""", +} + + +def make_project(tmp_path: Path, module: str, *, format: str = "csv", settings: str = 'write_strategy = "append"\n', table: str = "") -> Path: + (tmp_path / "transforms").mkdir() + (tmp_path / "drop").mkdir() + (tmp_path / "sqlload.toml").write_text(PROJECT.format(format=format, settings=settings, table=table)) + (tmp_path / "transforms" / "metrics.py").write_text(textwrap.dedent(module)) + return tmp_path + + +def give(project: Path, name: str, text: str, modified: float | None = None) -> None: + path = project / "drop" / name + path.write_text(text) + if modified is not None: + os.utime(path, (modified, modified)) + + +def sqlload(project: Path, *args: str) -> subprocess.CompletedProcess[str]: + home = project / "home" + home.mkdir(exist_ok=True) + env = {**os.environ, "HOME": str(home), "XDG_CONFIG_HOME": str(home / ".config"), "XDG_CACHE_HOME": str(home / ".cache")} + return subprocess.run( + [sys.executable, "-c", "from sqlload._cli import main; main()", *args], + cwd=project, + env=env, + capture_output=True, + text=True, + ) + + +def load(project: Path, *args: str) -> subprocess.CompletedProcess[str]: + done = sqlload(project, "load", *args) + assert done.returncode == 0, done.stdout + done.stderr + return done + + +def query(project: Path, sql: str) -> list[tuple[Any, ...]]: + with duckdb.connect(str(project / "wh.duckdb"), read_only=True) as connection: + return connection.execute(sql).fetchall() + + +@pytest.mark.parametrize("returned", ["dicts", "arrow", "polars", "generator"]) +def test_a_transform_returning_or_yielding_dicts_arrow_or_a_dataframe_loads_the_same_rows(tmp_path: Path, returned: str) -> None: + project = make_project(tmp_path, RETURNS[returned], table='columns = { v = "decimal(6,2)" }') + give(project, "a.csv", "id,v\n1,1.50\n2,2.25\n") + give(project, "b.csv", "id,v\n3,3.75\n") + load(project) + rows = query(project, "SELECT id, CAST(v AS VARCHAR), file, _sqlload_source_file FROM RAW.metrics ORDER BY id") + assert rows == [ + (1, "1.50", "./drop/a.csv", "./drop/a.csv"), + (2, "2.25", "./drop/a.csv", "./drop/a.csv"), + (3, "3.75", "./drop/b.csv", "./drop/b.csv"), + ] + types = dict(query(project, "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = 'metrics'")) + assert (types["id"], types["v"]) == ("BIGINT", "DECIMAL(6,2)") + + +def test_each_zip_member_gets_its_own_call(tmp_path: Path) -> None: + project = make_project(tmp_path, RETURNS["dicts"], settings='compression = "zip"\nwrite_strategy = "append"\n') + with zipfile.ZipFile(project / "drop" / "batch.zip", "w") as archive: + archive.writestr("a.csv", "id,v\n1,x\n2,y\n") + archive.writestr("b.csv", "id,v\n3,z\n") + load(project) + rows = query(project, "SELECT file, _sqlload_source_file, count(*) FROM RAW.metrics GROUP BY ALL ORDER BY 1") + assert rows == [("./drop/batch.zip/a.csv", "./drop/batch.zip/a.csv", 2), ("./drop/batch.zip/b.csv", "./drop/batch.zip/b.csv", 1)] + + +@pytest.mark.parametrize( + ("body", "error"), + [ + (' if ctx.path.endswith("b.csv"):\n raise ValueError("bad sheet")\n return rows\n', "ValueError: bad sheet"), + (' if ctx.path.endswith("b.csv"):\n return None\n return rows\n', "it returned NoneType; return or yield dicts, lists of dicts"), + (' if ctx.path.endswith("b.csv"):\n yield "text"\n yield rows\n', "transform metric_rows yielded str; yield dicts, lists of dicts"), + ], + ids=["raises", "returns-none", "yields-text"], +) +def test_a_failing_transform_fails_its_file_which_neither_commits_nor_moves(tmp_path: Path, body: str, error: str) -> None: + settings = 'write_strategy = "append"\nafter_load = { move_to = "done" }\n' + project = make_project(tmp_path, "import sqlload as sl\n\n@sl.transform\ndef metric_rows(rows, ctx):\n" + body, settings=settings) + give(project, "a.csv", "id,v\n1,x\n") + load(project) + give(project, "b.csv", "id,v\n2,y\n") + failed = sqlload(project, "load") + assert failed.returncode == 1, failed.stdout + failed.stderr + output = failed.stdout + failed.stderr + assert "./drop/b.csv: transform metric_rows failed" in output, output + assert error in output, output + assert sorted(path.name for path in (project / "drop").iterdir()) == ["b.csv"] + assert query(project, "SELECT id FROM RAW.metrics ORDER BY id") == [("1",)] + + +def test_a_merge_still_merges_files_oldest_first_with_a_record_each(tmp_path: Path) -> None: + settings = 'write_strategy = "merge"\nunique_key = ["id"]\n' + project = make_project(tmp_path, RETURNS["dicts"], settings=settings) + now = time.time() + give(project, "a.csv", "id,v\n1,newest\n", modified=now) + give(project, "b.csv", "id,v\n1,oldest\n2,b\n", modified=now - 200) + give(project, "c.csv", "id,v\n1,middle\n3,c\n", modified=now - 100) + load(project) + assert query(project, "SELECT id, v FROM RAW.metrics ORDER BY id") == [(1, "newest"), (2, "b"), (3, "c")] + merged = query(project, "SELECT count(*) FROM RAW._sqlload_runs WHERE note = 'files source merge'") + assert merged == [(3,)] + + +#: A second module registering a transform as `metric_rows` too, and one as `tidy`. +OTHER = """ +import sqlload as sl + +@sl.transform("metric_rows") +def other_rows(rows, ctx): + return rows + +@sl.transform +def tidy(rows, ctx): + return rows +""" + + +@pytest.mark.parametrize( + ("transform", "other", "problem"), + [ + ( + "missing_rows", + "", + "transform missing_rows: no function is registered as missing_rows with @sl.transform; the registered ones are metric_rows", + ), + ( + "metric_rows", + OTHER, + 'transform metric_rows: 2 functions are registered as metric_rows, in transforms/metrics.py (metric_rows) and transforms/other.py (other_rows); give one another name with @sl.transform("...")', + ), + ], + ids=["unknown-name", "two-under-one-name"], +) +def test_plan_refuses_a_transform_it_cannot_look_up_and_plans_other_tables(tmp_path: Path, transform: str, other: str, problem: str) -> None: + project = make_project(tmp_path, RETURNS["dicts"]) + if other: + (project / "transforms" / "other.py").write_text(other) + config = project / "sqlload.toml" + config.write_text( + config.read_text().replace('transform = "metric_rows"', f'transform = "{transform}"') + + '\n[sources.plain]\nkind = "files"\nconnection = "files"\npath = "plain/"\nformat = "csv"\ndescription = "Plain."\nwrite_strategy = "append"\n' + ) + give(project, "a.csv", "id,v\n1,x\n") + (project / "plain").mkdir() + (project / "plain" / "p.csv").write_text("id\n1\n") + planned = sqlload(project, "plan") + assert planned.returncode == 1, planned.stdout + planned.stderr + assert problem in planned.stdout, planned.stdout + assert "└── RAW.plain " in planned.stdout, planned.stdout + + +def test_a_transform_registered_under_a_name_of_its_own_is_found_by_that_name(tmp_path: Path) -> None: + module = """ +import sqlload as sl + +@sl.transform("kpis") +def reshape_metrics(rows, ctx): + return [{"id": int(row["id"]), "unit": ctx.unit, "table": ctx.table} for row in rows.to_pylist()] +""" + project = make_project(tmp_path, module) + config = project / "sqlload.toml" + config.write_text(config.read_text().replace('transform = "metric_rows"', 'transform = "kpis"')) + give(project, "a.csv", "id,v\n1,x\n") + load(project) + assert query(project, "SELECT id, unit, \"table\" FROM RAW.metrics") == [(1, "file", "metrics")] + config.with_name("by_function.toml").write_text(config.read_text().replace('transform = "kpis"', 'transform = "reshape_metrics"')) + wrong = sqlload(project, "plan", "-c", str(config.with_name("by_function.toml"))) + assert wrong.returncode == 1, wrong.stdout + wrong.stderr + assert "transform reshape_metrics: no function is registered as reshape_metrics with @sl.transform; the registered ones are kpis" in wrong.stdout, wrong.stdout + + +def test_plan_shows_the_transformed_columns_and_reads_the_files_itself(tmp_path: Path) -> None: + module = """ +import sqlload as sl + +@sl.transform +def metric_rows(rows, ctx): + return [{"metric": row["name"], "month": month, "value": value} for row in rows.to_pylist() for month, value in row.items() if month != "name"] +""" + project = make_project(tmp_path, module, format="parquet") + pq.write_table(pa.table({"name": ["sales"], "jan": [1.5], "feb": [2.5]}), project / "drop" / "a.parquet") + planned = sqlload(project, "plan") + assert planned.returncode == 0, planned.stdout + planned.stderr + lines = [line.strip("│├└─ ") for line in planned.stdout.splitlines()] + assert [line.split()[0] for line in lines if line.split()[:1] in (["metric"], ["month"], ["value"])] == ["metric", "month", "value"], planned.stdout + assert any(line.startswith("value") and "float" in line for line in lines), planned.stdout + assert any(line.startswith("source") and line.endswith("transform metric_rows") for line in lines), planned.stdout + plain = project / "plain.toml" + plain.write_text((project / "sqlload.toml").read_text().replace('transform = "metric_rows"', "")) + direct = sqlload(project, "plan", "-c", str(plain)) + route = [line for line in lines if line.startswith("route")] + assert route != [line.strip("│├└─ ") for line in direct.stdout.splitlines() if line.strip("│├└─ ").startswith("route")], direct.stdout + assert route == ["route read parquet into Arrow -> DuckDB takes Arrow"], planned.stdout + + +def test_a_column_the_plan_did_not_see_fails_the_file_and_a_missing_one_loads_null(tmp_path: Path) -> None: + module = """ +import sqlload as sl + +@sl.transform +def metric_rows(table, ctx): + rows = table.to_pylist() + if ctx.path.endswith("a.csv"): + return [{"id": int(row["id"])} for row in rows] + if ctx.path.endswith("b.csv"): + return [{"id": int(row["id"]), "v": row["v"], "extra": 1} for row in rows] + return [{"id": int(row["id"]), "v": row["v"]} for row in rows] +""" + project = make_project(tmp_path, module) + give(project, "z.csv", "id,v\n1,x\n") + load(project) + give(project, "a.csv", "id,v\n2,y\n") + load(project) + assert query(project, "SELECT id, v FROM RAW.metrics ORDER BY id") == [(1, "x"), (2, None)] + give(project, "b.csv", "id,v\n3,z\n") + failed = sqlload(project, "load") + assert failed.returncode == 1, failed.stdout + failed.stderr + assert "returned column extra, which it did not return for ./drop/" in failed.stdout + failed.stderr, failed.stdout + failed.stderr + + +def test_the_load_report_gives_each_files_transform_time(tmp_path: Path) -> None: + project = make_project(tmp_path, RETURNS["arrow"]) + give(project, "a.csv", "id,v\n1,x\n") + give(project, "b.csv", "id,v\n2,y\n") + done = load(project) + report = [line.split()[:2] for line in done.stdout.splitlines() if line.strip().startswith("transform")] + assert report == [["transform", "./drop/a.csv"], ["transform", "./drop/b.csv"]], done.stdout + give(project, "c.csv", "id,v\n3,z\n") + reported = json.loads(load(project, "--json").stdout) + transforms = reported["tables"][0]["transforms"] + assert [(item["file"], item["rows"], item["dicts"]) for item in transforms] == [("./drop/c.csv", 1, False)] + + +def test_plan_says_while_the_table_holds_files_an_earlier_transform_loaded(tmp_path: Path) -> None: + project = make_project(tmp_path, RETURNS["dicts"]) + give(project, "a.csv", "id,v\n1,x\n") + give(project, "b.csv", "id,v\n2,x\n") + planned = sqlload(project, "plan").stdout + assert "transform changed" not in planned, planned + assert "RAW._sqlload_runs create table (run log)" in " ".join(planned.split()), planned + applied = sqlload(project, "apply") + assert "RAW._sqlload_runs" in applied.stdout, applied.stdout + assert "_sqlload_runs" not in sqlload(project, "plan").stdout + load(project) + assert "transform changed" not in sqlload(project, "plan").stdout + module = project / "transforms" / "metrics.py" + module.write_text(module.read_text() + "\n# rows now keep their file\n") + changed = "transform changed: 2 files loaded under an earlier version (--reload transforms every file again)" + planned = sqlload(project, "plan").stdout + assert changed in planned, planned + give(project, "c.csv", "id,v\n3,x\n") + load(project) + planned = sqlload(project, "plan").stdout + assert changed in planned, planned + reloading = sqlload(project, "plan", "--reload") + assert "transform changed" not in reloading.stdout, reloading.stdout + load(project, "--reload", "--yes") + assert "transform changed" not in sqlload(project, "plan").stdout + + +def test_a_changed_transform_reloads_a_merge_table_with_moved_files_into_one_row_per_key(tmp_path: Path) -> None: + settings = 'write_strategy = "merge"\nunique_key = ["id"]\nafter_load = { move_to = "done" }\n' + project = make_project(tmp_path, RETURNS["dicts"], settings=settings) + now = time.time() + give(project, "a.csv", "id,v\n1,a\n2,a\n", modified=now - 200) + give(project, "b.csv", "id,v\n1,b\n", modified=now - 100) + load(project) + module = project / "transforms" / "metrics.py" + module.write_text(module.read_text().replace('"v": row["v"]', '"v": row["v"].upper()')) + give(project, "c.csv", "id,v\n3,c\n", modified=now) + assert "transform changed: 2 files loaded under an earlier version" in sqlload(project, "plan").stdout + planned = sqlload(project, "plan", "--reload") + assert "files 3 to read of 3 listed: 1 in ./drop/, 2 in ./done/" in planned.stdout, planned.stdout + load(project, "--reload", "--yes") + rows = query(project, "SELECT id, v, file FROM RAW.metrics ORDER BY id") + assert rows == [(1, "B", "./drop/b.csv"), (2, "A", "./drop/a.csv"), (3, "C", "./drop/c.csv")] + assert sorted(path.name for path in (project / "done").iterdir()) == ["a.csv", "b.csv", "c.csv"] + give(project, "d.csv", "id,v\n4,d\n") + after = sqlload(project, "plan").stdout + assert "1 new of 1 listed" in after and "transform changed" not in after, after + +#: The `KPIs` tab, as the Sheets API writes it: a merged title row, a header row, label-only rows, KPI blocks with measures below each, a note and an `IFERROR` "" here and there. +def text(value: str) -> dict[str, Any]: + return {"effectiveValue": {"stringValue": value}, "formattedValue": value} + + +def number(value: float, shown: str | None = None, kind: str = "") -> dict[str, Any]: + cell: dict[str, Any] = {"effectiveValue": {"numberValue": value}, "formattedValue": shown if shown is not None else f"{value:g}"} + if kind: + cell["effectiveFormat"] = {"numberFormat": {"type": kind}} + return cell + + +def percent(value: float) -> dict[str, Any]: + return number(value, f"{value * 100:.1f}%", "PERCENT") + + +BLANK: dict[str, Any] = {} +EMPTY_TEXT = text("") + +KPI_ROWS: list[list[dict[str, Any]]] = [ + [text("Service KPIs 2026"), BLANK, BLANK, BLANK, BLANK, BLANK, BLANK], + [text("ID"), text("KPI"), text("Period 1"), text("Period 2"), text("Jan"), text("Feb"), text("Mar")], + [BLANK, text("Customer service")], + [BLANK, text("Targets reviewed quarterly"), text("note: P1 excludes returns")], + [text("K1"), text("On-time delivery"), percent(0.9), EMPTY_TEXT, percent(0.95), number(0.8), EMPTY_TEXT], + [BLANK, text("Delivered on time"), number(9), BLANK, number(19), number(8)], + [BLANK, text("Delivered"), number(10), number(0), number(20), number(10), number(0)], + [text("K2"), text("First contact resolution"), percent(0.5), percent(0.6), EMPTY_TEXT, percent(0.75), percent(1)], + [BLANK, text("Resolved first time"), number(5), number(6), BLANK, number(3), number(2)], + [BLANK, text("Contacts"), number(10), number(10), number(0), number(4), number(2)], + [BLANK, text("Escalations"), number(1), text("n/a - system outage"), BLANK, number(0)], + [], + [text("K3"), text("Average handle time (min)"), number(6.5), number(7), number(6.25), BLANK, number(5.5)], + [BLANK, text("Calls handled"), number(120), number(130), number(125), BLANK, number(118)], +] + +TIDY_MODULE = """ +import pandas as pd +import sqlload as sl + +MONTHS = ["Period 1", "Period 2", "Jan", "Feb", "Mar"] + +@sl.transform +def kpi_sheet(rows, ctx): + assert (ctx.unit, ctx.sheet) == ("range", "KPIs"), ctx + df = rows.to_pandas() + kpi = df["ID"].notna() + df["metric_id"] = df["ID"].ffill() + df["metric"] = df["KPI"].where(kpi).ffill() + df["label"] = df["KPI"] + df["row_kind"] = kpi.map({True: "kpi", False: "measure"}) + df = df[df["metric_id"].notna()] + tidy = df.melt(id_vars=["metric_id", "metric", "label", "row_kind"], value_vars=MONTHS, var_name="month", value_name="value") + tidy["value"] = pd.to_numeric(tidy["value"], errors="coerce") + return tidy.dropna(subset=["value"]).reset_index(drop=True) +""" + +#: (metric_id, label, row_kind) and each month's value, worked out by hand from `KPI_ROWS`. +TIDY: dict[tuple[str, str, str], dict[str, float]] = { + ("K1", "On-time delivery", "kpi"): {"Period 1": 0.9, "Jan": 0.95, "Feb": 0.8}, + ("K1", "Delivered on time", "measure"): {"Period 1": 9, "Jan": 19, "Feb": 8}, + ("K1", "Delivered", "measure"): {"Period 1": 10, "Period 2": 0, "Jan": 20, "Feb": 10, "Mar": 0}, + ("K2", "First contact resolution", "kpi"): {"Period 1": 0.5, "Period 2": 0.6, "Feb": 0.75, "Mar": 1}, + ("K2", "Resolved first time", "measure"): {"Period 1": 5, "Period 2": 6, "Feb": 3, "Mar": 2}, + ("K2", "Contacts", "measure"): {"Period 1": 10, "Period 2": 10, "Jan": 0, "Feb": 4, "Mar": 2}, + ("K2", "Escalations", "measure"): {"Period 1": 1, "Feb": 0}, + ("K3", "Average handle time (min)", "kpi"): {"Period 1": 6.5, "Period 2": 7, "Jan": 6.25, "Mar": 5.5}, + ("K3", "Calls handled", "measure"): {"Period 1": 120, "Period 2": 130, "Jan": 125, "Mar": 118}, +} +METRICS = {"K1": "On-time delivery", "K2": "First contact resolution", "K3": "Average handle time (min)"} + + +class FakeSheets(BaseHTTPRequestHandler): + """Hands out a token, and the `KPIs` tab for any spreadsheet and range.""" + + def do_POST(self) -> None: + self.rfile.read(int(self.headers.get("Content-Length", "0"))) + self.answer({"access_token": "t", "token_type": "Bearer", "expires_in": 3600}) + + def do_GET(self) -> None: + rows = [{"values": row} for row in KPI_ROWS] + sheet = {"properties": {"title": "KPIs", "sheetType": "GRID"}, "data": [{"rowData": rows}]} + self.answer({"properties": {"title": "Fake spreadsheet"}, "sheets": [sheet]}) + + def answer(self, body: dict[str, Any]) -> None: + data = json.dumps(body).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + def log_message(self, format: str, *args: Any) -> None: + return + + +def test_a_sheets_range_goes_through_its_transform_into_tidy_rows(tmp_path: Path) -> None: + server = ThreadingHTTPServer(("127.0.0.1", 0), FakeSheets) + threading.Thread(target=server.serve_forever, daemon=True).start() + root = f"http://127.0.0.1:{server.server_address[1]}" + try: + (tmp_path / "transforms").mkdir() + (tmp_path / "transforms" / "kpis.py").write_text(TIDY_MODULE) + (tmp_path / "sqlload.toml").write_text( + f"""\ +python_paths = ["transforms"] + +[connections.sheets] +kind = "google_sheets" +client_id = "client.apps.googleusercontent.com" +client_secret = "secret" +refresh_token = "refresh" +base_url = "{root}/v4/spreadsheets" +token_url = "{root}/token" + +[connections.wh] +kind = "duckdb" +database = "wh.duckdb" + +[targets.dev] +connection = "wh" +loader_schema = "RAW" + +[sources.kpis] +kind = "google_sheets" +connection = "sheets" +spreadsheet_id = "synthetic" +ranges = {{ metrics = "KPIs" }} +header_row = 2 +description = "Service KPIs, tidied." +write_strategy = "replace" +transform = "kpi_sheet" +""" + ) + done = load(tmp_path) + assert "transform KPIs" in done.stdout, done.stdout + finally: + server.shutdown() + rows = query(tmp_path, "SELECT metric_id, metric, label, row_kind, month, value FROM RAW.kpis__metrics") + expected = [ + (metric_id, METRICS[metric_id], label, kind, month, float(value)) + for (metric_id, label, kind), months in TIDY.items() + for month, value in months.items() + ] + assert len(expected) == 34 + assert sorted(rows) == sorted(expected) + + +def test_a_function_passed_as_transform_resolves_to_itself_when_two_files_register_its_name(tmp_path: Path) -> None: + from sqlload._registry import Registry + + registry = Registry("@sl.transform") + modules = [] + for name, body in (("a_clean", "return rows"), ("b_clean", "return rows.slice(1)")): + path = tmp_path / f"{name}.py" + path.write_text(f"def clean(rows, ctx):\n\t{body}\n") + spec = importlib.util.spec_from_file_location(f"transform_case_{name}", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + registry.register("clean", module.clean) + modules.append(module) + assert registry.lookup("clean")[0] is None, "the name alone is ambiguous" + key, shown = registry.key(modules[0].clean) + assert registry.lookup(key) == (modules[0].clean, None) + assert shown == "clean" + + +def test_a_function_passed_as_transform_is_released_when_the_run_finishes() -> None: + from sqlload._runtime import finish_run + from sqlload._transform import TRANSFORMS + + def page_transform() -> Any: + return lambda rows, ctx: rows + + function = page_transform() + released = weakref.ref(function) + key, _ = TRANSFORMS.key(function) + assert TRANSFORMS.lookup(key)[0] is function + del function + finish_run() + gc.collect() + assert released() is None + assert TRANSFORMS.lookup(key)[0] is None + + +def test_a_transform_whose_file_cant_be_read_is_versioned_by_its_code() -> None: + from sqlload._runtime import transform_version + from sqlload._transform import TRANSFORMS + + def version_of(body: str) -> str | None: + namespace: dict[str, Any] = {} + exec(compile(f"def unread_file_transform(rows, ctx):\n\t{body}\n", "", "exec"), namespace) + TRANSFORMS.register("unread_file_transform", namespace["unread_file_transform"]) + return transform_version("unread_file_transform")[0] + + first = version_of("return rows") + assert first is not None and first.startswith("unread_file_transform@") + assert version_of("return rows.slice(1)") != first + assert version_of("return rows") == first diff --git a/website/src/content/docs/docs/cli/load.mdx b/website/src/content/docs/docs/cli/load.mdx index b789d6d2d..75df6dcaf 100644 --- a/website/src/content/docs/docs/cli/load.mdx +++ b/website/src/content/docs/docs/cli/load.mdx @@ -107,6 +107,10 @@ A `postgres_cdc` table is copied again with `--resnapshot` instead; `--reload` r [Reloading](/docs/guides/postgres-cdc/#reloading). `--reload` is also refused on [Postgres targets](/docs/targets/postgres/). +On a files source, `--reload` reads every file again; from 0.52.0 a `merge` table merges them one +at a time into the new table, and `after_load = { move_to = ... }` reads the moved files too. See +[Reloading a files table](/docs/sources/files/#reloading-a-files-table). + ## Following a Postgres CDC source `sqlload load SOURCE --follow` keeps streaming one `postgres_cdc` source after it catches up, @@ -185,7 +189,7 @@ $ sqlload load --json | `summary` | How many objects were `created`, `altered` (each table once, however many changes it had) and `recreated`; tables `loaded` and `failed`, and from 0.43.0 `skipped`; `rows` loaded. | | `plan` | The same document `plan --json` prints, with its own `ok` and, from 0.34.0, `retryable`. Tables a plan error left out are in its `errors`, and tables with nothing to load in `skipped`. | | `objects` | One entry per change to a table or stage, by `apply` or during a load: its `action` (`created`, `altered` or `recreated`), its `name` and the `detail`. A table that gains two columns has two `altered` entries. | -| `tables` | Each table loaded: its `status` (`loaded` or `failed`, or from 0.43.0 `skipped` for a table [`--fail-fast`](#stopping-at-the-first-failure) didn't start, its `error` naming the table that stopped the run and `retryable` `true`), `rows`, `duration_seconds`, a `note` (the cursor window it read, or that its table waits for first data), and the `error` when it failed. From 0.34.0, `retryable` says whether that failure may pass on a rerun. `source`, `target` and `write_strategy` are `null` only for a `load --follow` that failed before any table. From 0.39.0, a files source with `after_load` also has `after_load`: its `action` (`delete` or `move`), `status` (`done`, or `failed` when any file stayed), `files` handled, and `failed`, each file left with its `error`. A table that keeps its files has no such field. From 0.51.0, a `delete_insert` table with a `replace_key` also has `replaced_keys`: its `columns`, the `values` its commits replaced (sorted, each as its columns' text joined by `, `, `NULL` for a NULL, at most 1,000) and `truncated`, `true` when there were more. `apply` loads nothing, so its list is empty. | +| `tables` | Each table loaded: its `status` (`loaded` or `failed`, or from 0.43.0 `skipped` for a table [`--fail-fast`](#stopping-at-the-first-failure) didn't start, its `error` naming the table that stopped the run and `retryable` `true`), `rows`, `duration_seconds`, a `note` (the cursor window it read, or that its table waits for first data), and the `error` when it failed. From 0.34.0, `retryable` says whether that failure may pass on a rerun. `source`, `target` and `write_strategy` are `null` only for a `load --follow` that failed before any table. From 0.39.0, a files source with `after_load` also has `after_load`: its `action` (`delete` or `move`), `status` (`done`, or `failed` when any file stayed), `files` handled, and `failed`, each file left with its `error`. A table that keeps its files has no such field. From 0.51.0, a `delete_insert` table with a `replace_key` also has `replaced_keys`: its `columns`, the `values` its commits replaced (sorted, each as its columns' text joined by `, `, `NULL` for a NULL, at most 1,000) and `truncated`, `true` when there were more. From 0.52.0, a table with a [`transform`](/docs/sources/files/#transforming-each-file) also has `transforms`: for each file or range, its `file`, the `milliseconds` the function took, the `rows` it was given and `dicts`, whether it returned dicts. `apply` loads nothing, so its list is empty. | When the run can't start (the configuration doesn't parse, a connection fails), standard output holds one error document instead, the same as `plan --json`'s: @@ -229,7 +233,7 @@ The events are version 0: names and fields can still change in a minor release u | `table_progress` | `unit` (`page`, `batch`, `window` or `file`), `count` of them read so far, `rows` so far. | | `rate_limited` | `connection`, its `rate_limit` as `rate`, and `wait_ms`, how long a request waited. No `table`. | | `table_committing` | `rows`: the table's rows are read and its last commit starts. | -| `table_done` | `rows`, `elapsed_ms`, and `position`, the cursor value it reached, when it keeps one. `changes` lists the `created` and `altered` lines of a table that made or widened its columns while loading. | +| `table_done` | `rows`, `elapsed_ms`, and `position`, the cursor value it reached, when it keeps one. `changes` lists the `created` and `altered` lines of a table that made or widened its columns while loading. From 0.52.0, `transforms` lists each file or range a [`transform`](/docs/sources/files/#transforming-each-file) ran on, as in the JSON report. | | `table_failed` | `reason`, `retryable` (whether a rerun may pass), `elapsed_ms`. | | `table_not_started` | `reason`, such as the table that stopped a [`--fail-fast`](#stopping-at-the-first-failure) load. | | `run_finished` | `loaded`, `failed`, `not_started`, `rows`, `elapsed_ms` and `exit_code`. | diff --git a/website/src/content/docs/docs/configuration.mdx b/website/src/content/docs/docs/configuration.mdx index f7036a560..0105e699c 100644 --- a/website/src/content/docs/docs/configuration.mdx +++ b/website/src/content/docs/docs/configuration.mdx @@ -283,6 +283,7 @@ Set on a source for all its tables, or in a `[sources.NAME.tables."schema.table" | `max_depth` | `1000` | From 0.40.0. In a table entry with `unnest`: how many levels of child tables it makes at most. | | `naming` | the target's | From 0.40.0. `sqlload` or `dlt`. See [Naming](/docs/concepts/nested-data/#naming). | | `row_id` | `false` | From 0.40.0. Give every row a unique id, in `_sqlload_row_id`. Always on with `unnest` and under `naming = "dlt"`, where it is `_dlt_id`. | +| `transform` | none | From 0.52.0, files and `google_sheets` sources. The name a Python function was registered under with `@sl.transform` (its own name, or the one given as `@sl.transform("kpis")`) in a module under `python_paths`, run on each file's or range's rows before they load. `plan` fails on a name no function has, or two do. Needs `python_paths` and the `sqlload` Python package. An `sl.http` job takes [`transform=`](/docs/sources/python/#transforming-each-page) instead. See [Transforming each file](/docs/sources/files/#transforming-each-file) and [Transforming each range](/docs/sources/google-sheets/#transforming-each-range). | | `column_names` | none | From 0.49.0. The columns' names, in order, for data without a header row: CSV and Excel files with `header = false`, and `google_sheets` ranges with `header = false`. The count must match the data's columns. See [Naming headerless columns](/docs/sources/files/#naming-headerless-columns). | ClickHouse targets also take these, for the tables SQLLoad creates. Other targets refuse them, and diff --git a/website/src/content/docs/docs/guides/from-dlt.mdx b/website/src/content/docs/docs/guides/from-dlt.mdx index 1350c744d..6d4ca66b5 100644 --- a/website/src/content/docs/docs/guides/from-dlt.mdx +++ b/website/src/content/docs/docs/guides/from-dlt.mdx @@ -38,6 +38,8 @@ what dlt 1.32.0 loads from the same files. For how the two differ in design, see | `max_table_nesting` and child tables | `unnest` and `max_depth`. See [Nested data](/docs/concepts/nested-data/). | | `rest_api` | `sl.http`, below. | | `filesystem` with a reader | A [files source](/docs/sources/files/), below. | +| A transformer, or `add_map`, that reshapes each file's rows | From 0.52.0, a files source's `transform`, an `@sl.transform` function. See [Transforming each file](/docs/sources/files/#transforming-each-file). | +| `add_map` on a `rest_api` resource, reshaping each page's records | From 0.52.0, `transform=` on `paginate` or `fan_out`. See [Transforming each page](/docs/sources/python/#transforming-each-page). | | `secrets.toml` and `config.toml` | `${ENV:NAME}` in `sqlload.toml`. | | `pipeline.run()` from a task | `sqlload load --json`, below. | @@ -119,6 +121,9 @@ How `rest_api` settings map: [Keys in a POST body](/docs/sources/python/#keys-in-a-post-body). - `method: POST` with a `json` body is `method="POST", json_body={...}`, and paginators that page in the body take `location="body"`; see [POST with a JSON body](/docs/sources/python/#post-with-a-json-body). +- A resource reshaped with `add_map` (each record split, unpivoted or cleaned in Python) is + `transform=` on the job: a function called once per page with the page's rows, from 0.52.0; see + [Transforming each page](/docs/sources/python/#transforming-each-page). - `response_actions` that ignore a status are `on_status={404: "skip"}`. 429, 5xx and dropped connections are retried, honouring `Retry-After`; `retries=` on `sl.http` sets how many times. - An API that returns each row as a list of cells, such as a spreadsheet's, is `cells=sl.cells(...)`; @@ -149,6 +154,11 @@ table is named after the source: (`sheet_range: A:C` and `column_names: [metric, category, value]` in a pipeline's settings), is `xlsx = { range = "A:C", header = false }` and `column_names = ["metric", "category", "value"]` from 0.49.0. See [Naming headerless columns](/docs/sources/files/#naming-headerless-columns). +- A file that needs reshaping before it loads (title rows, a header split over two rows, months + across), read in dlt by a custom reader or reshaped by `add_map`, takes a `transform` from 0.52.0: + a Python function under `python_paths`, decorated with `@sl.transform` and named in `transform`, + that gets each file's rows as a `pyarrow.Table` and returns or yields the rows to load. See + [Transforming each file](/docs/sources/files/#transforming-each-file). See [Load files](/docs/guides/files/). ## Tables that look like dlt's diff --git a/website/src/content/docs/docs/sources/files.mdx b/website/src/content/docs/docs/sources/files.mdx index 0e69e0f9b..796b374ca 100644 --- a/website/src/content/docs/docs/sources/files.mdx +++ b/website/src/content/docs/docs/sources/files.mdx @@ -185,6 +185,196 @@ as any rewritten file does. Only an archive's index and the members read are fet requests. Workbooks, XML documents and ZIP members are read whole, up to 1 GiB each; a larger member, or an archive naming two members alike, fails. +## Transforming each file + +From 0.52.0. A transform is a Python function that makes a file loadable: it fixes headers, drops +title and totals rows, or unpivots a sheet. Keep business logic downstream, in SQL, where it can +change without reloading files. Decorate the function with `@sl.transform` in a module under +[`python_paths`](/docs/configuration/#top-level), and name it in `transform`: `@sl.transform` +registers it under its own name, and `@sl.transform("kpis")` under another. It needs the `sqlload` +Python package and pyarrow; the standalone `sqlload-cli` executable refuses a table with a +`transform`. The same kind of function also reshapes each [Google Sheets](/docs/sources/google-sheets/#transforming-each-range) +range and each [`sl.http` page](/docs/sources/python/#transforming-each-page). + +```toml +python_paths = ["transforms"] + +[sources.metrics] +kind = "files" +connection = "lake" +path = "drop/metrics/" +format = "csv" +description = "Monthly metrics, one file a month." +write_strategy = "append" +transform = "metric_rows" +``` + +It runs once for each file, and for each member of a ZIP archive, after SQLLoad reads it and before +its rows are typed against the table. It's called as `f(rows, ctx)`: `rows` is the file's rows as a +`pyarrow.Table`, and `ctx` an `sl.TransformContext`. It returns or yields, in any mix, what a +[Python table](/docs/guides/python/) may yield: dicts, lists of dicts, a `pyarrow.Table`, or a polars +or pandas DataFrame. Dicts are the simplest: + +```python title="transforms/metrics.py" +import pyarrow as pa +import sqlload as sl + + +@sl.transform +def metric_rows(rows: pa.Table, ctx: sl.TransformContext): + for row in rows.to_pylist(): + if row["metric"] in (None, "Total"): + continue + for month, value in row.items(): + if month != "metric" and value is not None: + yield {"metric": row["metric"], "month": month, "value": float(value)} +``` + +```console +$ sqlload load +✓ Plan complete 1 selected, 1 object to create, 1 to load (sqlload plan shows the tree) + +created raw.metrics 8 columns +loaded raw.metrics 5 rows 0.0s + transform ./drop/metrics/2026-10.csv 0.00s + transform ./drop/metrics/2026-11.csv 0.01s + +✓ Completed 1 loaded, 5 rows, 0.1s +├── slowest raw.metrics 0.0s +└── peak memory 134 MiB +``` + +`ctx.unit` is `"file"`. `ctx` also has `path` (as in `_sqlload_source_file`), `modified_at`, +`sheet`, `range`, `skip_rows` and `table`. + +The rows it gets are typed so that no value fails the read before the function sees it: CSV and +XML values as text, Parquet columns as stored, Excel cells as the workbook stores each column, and +JSON values as text except nested objects and lists, which stay JSON. `exclude_columns` and +`allow_columns` apply to the file's own columns, before the function; `columns` and `mask` apply to +what it returns. + +Dict values are typed as JSON values are, so a whole number is `int` and a decimal `float`; pin +exact types in `columns`. For large files, return Arrow or a DataFrame, which skips converting each +row. The same unpivot in polars: + +```python title="transforms/metrics.py" +import polars as pl +import pyarrow as pa +import sqlload as sl + + +@sl.transform +def metric_rows(rows: pa.Table, ctx: sl.TransformContext): + frame = pl.from_arrow(rows).filter(pl.col("metric").is_not_null() & (pl.col("metric") != "Total")) + return frame.unpivot(index="metric", variable_name="month", value_name="value").drop_nulls("value") +``` + +A pandas DataFrame crosses with its columns only; after filtering, call `reset_index(drop=True)` so +no index is left over. + +- `plan` looks the name up among the functions `@sl.transform` registered on `python_paths`, and + fails when none has it, listing those that do, or when two do, naming both files: + + ```console + $ sqlload plan + Errors (1) + └── raw.metrics + └── reason transform metrics_rows: no function is registered as metrics_rows with @sl.transform; the registered ones are metric_rows + + ✗ Plan failed with errors + ``` + + It then runs the function on the newest file (by name), and the table's columns are what it + returned, typed as a Python table's rows are. So the function runs at `plan` and again at load: + keep it free of side effects. +- Every file must return the columns the planned file did, or fewer (a missing one loads NULL). A + file returning a new column fails, naming it; pin it in `columns` to add it. +- A transform that raises, or returns or yields something else, fails that file and names the + function and the file. Its rows don't commit, and `after_load` leaves it where it is. The table + stops there, as for a file that doesn't read; files committed before it stay committed. +- The file order, `merge`, `after_load` and the `_sqlload_source_*` columns work as without a + transform; `_sqlload_source_row` numbers the rows the function returned. SQLLoad always reads the + files itself, so a destination never reads them directly. +- Each file is read whole and held in memory while the function runs. What it returns or yields + loads batch by batch as it comes, so a generator needs memory for one file and one batch of its + output. +- The load report gives each file's transform time, and a hint when dicts made a large file slow. + +### Changing a transform + +Each commit records the version of the transform it ran (a hash of the file the function is defined +in) and the files it ran on, in the same transaction as their rows. The record goes to the table's +run log, `_sqlload_runs` in its schema (objects under the table's folder on a files target), which +`plan` lists under Objects as `create table (run log)` until `apply` (or the first load) creates it. A load only +inserts into it, so a role without CREATE on the schema can load once `apply` has run under one +that has it. + +Files already loaded aren't read again, so after a change, `plan` warns on the table for as long as +it holds files an earlier version loaded, here after a load of one new file: + +```console +$ sqlload plan +Plan ready 1 selected, 1 to load + +Loads (1) sqlload load +└── raw.metrics append + ├── source files ./drop/metrics/ csv, compression by file name transform metric_rows + ├── route read csv into Arrow -> DuckDB takes Arrow + ├── columns 3 + ├── files 0 new of 3 listed + ├── batch 0 batches, up to 1,000 files or about 1 GB per commit + └── transform changed: 2 files loaded under an earlier version (--reload transforms every file again) + +✓ Plan complete 1 selected, 1 to load +``` + +The version covers the file the function is in; a change in another module it imports isn't +noticed. When that file can't be read, as for a function defined in a notebook, it covers the +function's own code instead. That code is Python's bytecode, which changes between Python minor +versions, so after upgrading Python, `plan` says such a transform changed even if you didn't edit it; +`--reload` loads its files under the new version. To load every file again with the new transform, run `sqlload load --reload` +([Reloading a files table](#reloading-a-files-table)), after which the warning is gone. A table +whose `after_load` deletes its files can't be reloaded: drop it and load the files still in +`path`. + +## Reloading a files table + +`sqlload load
--reload` reads every listed file again into a new table, which replaces the +old one in one commit at the end; until then readers see the old table. + +- `append`, `replace` and `delete_insert` tables read the files in batches, as a first load does. A + `delete_insert` with a [`replace_key`](#replacing-rows-by-a-key) refuses `--reload`. +- From 0.52.0, `merge` tables reload too, on DuckDB targets (Postgres targets refuse `--reload` for + every table). Each file merges into the new table on its own, as a load merges them: files merged + before go in the order they merged, matched by name and version, or for a file under `move_to` + by name, so a file `after_load` moved keeps its place; files never merged follow, oldest + modified first, and so does a new file in `path` named as one `move_to` holds, whose rows then + win over the older copy's. The new + table's file record (a reset, then one record per file) is written in the commit that swaps it + in, so a reload that fails part way leaves the old table and its file record as they were. After + a reload, the next load reads only files added since. `plan` shows + `merge (unique_key: …) into a new table, swapped in at the end (--reload)`. A `merge` table with `unnest` is refused, since its child tables aren't rebuilt: + drop the table and its child tables and load again. +- From 0.52.0, with `after_load = { move_to = ... }` the reload also lists the files under + `move_to`, which keep the names they loaded under (`_sqlload_source_file` names the file under + `path`). `plan` shows both folders and how many files it reads from each: + + ```text + ├── source files ./drop/ and ./done/ csv, compression by file name + ├── files 3 to read of 3 listed: 1 in ./drop/, 2 in ./done/ + ``` + + Once the new table is in, files read from `path` move as usual; those already under `move_to` + stay there. A reload that fails moves nothing. + + When some of the files are under `move_to`, SQLLoad reads them itself, even where the destination + reads files directly on a normal load, and the route line says why: + + ```text + ├── route read parquet into Arrow -> DuckDB takes Arrow (--reload also reads the files after_load moved to ./done/scores/) + ``` +- With `after_load = "delete"`, `--reload` is still refused: the files loaded before are gone. + ## Which files load | Strategy | Reads | @@ -193,7 +383,7 @@ or an archive naming two members alike, fails. | `delete_insert` with a partition `cursor_column` (from `partition_columns` or `path_columns`) | The newest loaded partition and everything after it, replacing those partitions whole. From 0.51.0, a file that changed since it loaded whose partition lies below that range fails the table, as does a new file below it in a partition the table holds; see below. | | `delete_insert` with a `replace_key` (from 0.51.0) | Only files not loaded yet, or changed since, one at a time, oldest modified first, each replacing the rows of the `replace_key` values it holds. | | `replace` | Every file, swapped in as a new table. | -| `merge` with a `unique_key` (from 0.38.0) | Only files not merged yet, one at a time, oldest modified first. | +| `merge` with a `unique_key` (from 0.38.0) | Only files not merged yet, one at a time, oldest modified first. Under `--reload` (from 0.52.0, DuckDB), every file again, one at a time, into a new table swapped in at the end; see [Reloading a files table](#reloading-a-files-table). | A files source needs a `write_strategy`, from the source or its table entry: `plan` refuses it otherwise. Before 0.46.0 it took `replace` without a `cursor_column`, and before 0.44.0 @@ -322,8 +512,11 @@ unique_key = ["sku"] that includes a file repeating a key with NULLs in it, even under the default `on_duplicate_keys`. - Merging files loads into DuckDB and Postgres targets so far, and MotherDuck through its DuckDB - connection, which isn't tested live. `--reload`, `state = "logged"` and a - `cursor_column` are refused: the run log, not a cursor, picks the files. To merge every file again, drop the table: a table created anew starts a new record. + connection, which isn't tested live. `state = "logged"` and a `cursor_column` are refused: the + run log, not a cursor, picks the files. To merge every file again, drop the table (a table + created anew starts a new record), or from 0.52.0 run `--reload` on DuckDB, which keeps the old + table until every file merged; see [Reloading a files table](#reloading-a-files-table). Before + 0.52.0 `--reload` was refused. ### Deleting or moving loaded files @@ -362,7 +555,9 @@ after_load = { move_to = "loaded" } - It needs `write_strategy = "append"` or `"merge"`, or from 0.51.0 `"delete_insert"` with a `replace_key`, which never read a loaded file again. A file that reappears under a loaded name is a new file when its version differs; an identical copy, - such as the same bytes on S3, is skipped as loaded and stays. `--reload` is refused. + such as the same bytes on S3, is skipped as loaded and stays. From 0.52.0, `--reload` with + `move_to` reads the moved files too, and with `"delete"` it's refused; see + [Reloading a files table](#reloading-a-files-table). Before 0.52.0 both were refused. - From 0.39.1, a run that finds the folder empty skips the table and exits 0, since `after_load` handled every file loaded before: `plan` lists it under skipped, as does `load --json` in `plan.skipped`. That needs an existing folder holding no file at all (hidden diff --git a/website/src/content/docs/docs/sources/google-sheets.mdx b/website/src/content/docs/docs/sources/google-sheets.mdx index d1c63f354..b51c085f0 100644 --- a/website/src/content/docs/docs/sources/google-sheets.mdx +++ b/website/src/content/docs/docs/sources/google-sheets.mdx @@ -125,6 +125,71 @@ same name, ignoring case, fails `plan`. is refused at `plan`, so pin `float` up front where fractions may appear. - A formula error (`#DIV/0!`, `#N/A`) fails the load naming the cell, rather than loading NULL. +## Transforming each range + +From 0.52.0. A table can take a [`transform`](/docs/sources/files/#transforming-each-file), as a +files table does: the name of an `@sl.transform` function on `python_paths`, called as +`f(rows, ctx)` with the range's rows as a `pyarrow.Table`. It returns or yields the rows to load, as +dicts, lists of dicts, a `pyarrow.Table`, or a polars or pandas DataFrame. It runs once for each +range at `plan`, which takes the table's columns from what it returns, and once at each load. +`ctx.unit` is `"range"`, and `ctx` has `spreadsheet_id`, `sheet` (the tab), `range` (as in +`ranges`) and `table`. + +A report with title rows above its header, a KPI per block and months across, becomes tidy rows: + +```toml +python_paths = ["transforms"] + +[sources.kpis] +kind = "google_sheets" +connection = "sheets" +spreadsheet_id = "1AbC..." +ranges = { metrics = "KPIs" } +header_row = 2 +description = "Service KPIs, one row per KPI or measure and month." +write_strategy = "replace" +transform = "kpi_sheet" +``` + +```python title="transforms/kpis.py" +import pandas as pd +import pyarrow as pa +import sqlload as sl + +MONTHS = ["Period 1", "Period 2", "Jan", "Feb", "Mar"] + +@sl.transform +def kpi_sheet(rows: pa.Table, ctx: sl.TransformContext): + df = rows.to_pandas() + kpi = df["ID"].notna() + df["metric_id"] = df["ID"].ffill() + df["metric"] = df["KPI"].where(kpi).ffill() + df["label"] = df["KPI"] + df["row_kind"] = kpi.map({True: "kpi", False: "measure"}) + df = df[df["metric_id"].notna()] + tidy = df.melt(id_vars=["metric_id", "metric", "label", "row_kind"], value_vars=MONTHS, + var_name="month", value_name="value") + tidy["value"] = pd.to_numeric(tidy["value"], errors="coerce") + return tidy.dropna(subset=["value"]).reset_index(drop=True) +``` + +- `header_row` names the header's row, so title rows above it, such as a merged group title, never + reach the function; a range starting at the header row (`"KPIs!A2:G"`) does the same. A header + cell must still name every column; when it doesn't, use `header = false` with `column_names`. +- A merged cell's value is in its first cell only; the others arrive empty, so fill them down in the + function, as `ffill()` does above. +- The rows the function gets differ from a table's without one in two ways: an empty text, as a + formula's `IFERROR(..., "")` returns, is NULL and doesn't type its column; and a number in a + column typed `text` (one mixing numbers and notes) is its value, `0.9`, not the text the sheet + shows, `90.0%`. A formula error still fails the load, naming the cell. +- A function that raises, or returns or yields something else, fails the table, naming the function + and the range; the table stays as it was. The load report gives each range's transform time. +- With `merge` or `append`, `plan` shows `transform changed: 1 range loaded under an earlier + version (--reload transforms every range again)` while the table holds rows a range loaded before the file + the function is defined in changed; `replace` tables load every row again each run. The version + is recorded in `_sqlload_runs`, which `apply` creates, as for + [files](/docs/sources/files/#changing-a-transform). + ## Writing A sheet has no column to read after, so its tables usually take `write_strategy = "replace"` and diff --git a/website/src/content/docs/docs/sources/python.mdx b/website/src/content/docs/docs/sources/python.mdx index 4e7812f73..850194f7a 100644 --- a/website/src/content/docs/docs/sources/python.mdx +++ b/website/src/content/docs/docs/sources/python.mdx @@ -295,8 +295,8 @@ thread-safe. Set `parallel_tables=True` on the source when its client is. | Function | Meaning | |----------|---------| | `sl.http(source, *, auth=None, rate_limit=None, retries=5, headers=None, connection=None)` | Requests to the source's `http` connection. `rate_limit` (`"10/s"`, `"600/m"`, `"1000/h"`) is shared by every job on the connection with the same limit. | -| `api.paginate(path, *, paginator, params=None, select=None, on_status=None, fields=None, keep_record=None, method="GET", json_body=None, key=None, cells=None, select_missing="fail")` | Every page of `path`. `select` is where the rows are: `$`, or a dotted path such as `$.data.items`. From 0.45.0, a response without the `select` path fails the load, naming what the response holds; `select_missing="empty"` reads it as no rows instead. See [A missing select path](#a-missing-select-path). `fields` and `keep_record` (from 0.27.0) pick and rename each row's values; see [Fields](#fields). `method` and `json_body` are from 0.36.0; see [POST with a JSON body](#post-with-a-json-body). From 0.37.0, `cells=` turns each row's list of cells into columns; see [Rows of cells](#rows-of-cells). | -| `api.fan_out(path, *, over, concurrency=8, paginator=None, params=None, select=None, on_status=None, fields=None, keep_record=None, method="GET", json_body=None, key=None, cells=None, batch=None, select_missing="fail")` | `path` once per key of `over`, up to `concurrency` at a time, while keys still arrive. `select_missing` is as for `paginate`. `fields` and `keep_record` are from 0.27.0. `over` is a paginate whose `select` picks keys (`$.items[*].id` fills `{id}`), or a list: a value fills the path's one placeholder, and a dict fills placeholders by name. From 0.37.0, `over` can also be another `fan_out`; see [Multi-level fan-out](#multi-level-fan-out). `method` and `json_body` are from 0.36.0, as for `paginate`. From 0.43.0, keys also fill `{placeholders}` in `json_body`, and `batch` sends several keys per request; see [Keys in a POST body](#keys-in-a-post-body). | +| `api.paginate(path, *, paginator, params=None, select=None, on_status=None, fields=None, keep_record=None, method="GET", json_body=None, key=None, cells=None, select_missing="fail", transform=None)` | Every page of `path`. `select` is where the rows are: `$`, or a dotted path such as `$.data.items`. From 0.45.0, a response without the `select` path fails the load, naming what the response holds; `select_missing="empty"` reads it as no rows instead. See [A missing select path](#a-missing-select-path). `fields` and `keep_record` (from 0.27.0) pick and rename each row's values; see [Fields](#fields). `method` and `json_body` are from 0.36.0; see [POST with a JSON body](#post-with-a-json-body). From 0.37.0, `cells=` turns each row's list of cells into columns; see [Rows of cells](#rows-of-cells). From 0.52.0, `transform` reshapes each page's rows in Python before they load; see [Transforming each page](#transforming-each-page). | +| `api.fan_out(path, *, over, concurrency=8, paginator=None, params=None, select=None, on_status=None, fields=None, keep_record=None, method="GET", json_body=None, key=None, cells=None, batch=None, select_missing="fail", transform=None)` | `path` once per key of `over`, up to `concurrency` at a time, while keys still arrive. `select_missing` is as for `paginate`. `fields` and `keep_record` are from 0.27.0. `over` is a paginate whose `select` picks keys (`$.items[*].id` fills `{id}`), or a list: a value fills the path's one placeholder, and a dict fills placeholders by name. From 0.37.0, `over` can also be another `fan_out`; see [Multi-level fan-out](#multi-level-fan-out). `method` and `json_body` are from 0.36.0, as for `paginate`. From 0.43.0, keys also fill `{placeholders}` in `json_body`, and `batch` sends several keys per request; see [Keys in a POST body](#keys-in-a-post-body). From 0.52.0, `transform` is as for `paginate`, called once per page of each key. | | Paginator | Next page | |-----------|-----------| @@ -327,8 +327,9 @@ thread-safe. Set `parallel_tables=True` on the source when its client is. Another status fails, unless `on_status` maps it to `"skip"`. - The job is checked before any request. Errors name the request, the status and the start of the body, with the connection's secrets replaced by `***`. -- A job's rows never come back into Python. Reading a row from a job fails (rule - `python_row_loop`); filter rows downstream. +- A job's rows never come back into Python one at a time. Reading a row from a job fails (rule + `python_row_loop`); filter rows downstream, or reshape each page with + [`transform`](#transforming-each-page). ### Reading pages ahead @@ -368,6 +369,77 @@ or empty pointer, or no `rel="next"` link. That is how an API says it has no mor can't tell it from an error response that also leaves the pointer out; with `select` set, such a response usually lacks the rows' path too, and fails as above. +### Transforming each page + +From 0.52.0. `transform=` on `paginate` or `fan_out` calls a Python function once per page, in page +order, on the page's rows after `select`, `fields` and `cells`. SQLLoad still makes every request, +with its paging, retries and read-ahead; the function only reshapes what a page returned. It's the +same kind of function a [files table's `transform`](/docs/sources/files/#transforming-each-file) +runs: `f(rows, ctx)`, with the page's rows as a `pyarrow.Table`, returning or yielding dicts, lists of +dicts, a `pyarrow.Table`, or a polars or pandas DataFrame. Pass the function itself, or the name +`@sl.transform` registered it under. A function passed itself is the one called, even when another +file registers a function of the same name, and it needn't be registered at all: + +```python title="sources/shop.py" +import json + +import sqlload as sl + +shop = sl.source("shop", connection="shop", write_strategy="replace") +api = sl.http(shop) + + +@sl.transform +def order_lines(rows, ctx): + for order in rows.to_pylist(): + for line in json.loads(order["lines"]): + yield {"region": ctx.key["region"], "order_id": order["id"], **line} + + +@shop.table(description="Order lines, one row per line, by region.") +def lines(ctx): + return api.fan_out("/regions/{region}/orders", over=["north", "south"], select="$.orders", + transform=order_lines) +``` + +```console +$ sqlload load +✓ Plan complete 1 selected, 1 object to create, 1 to load (sqlload plan shows the tree) + +created raw.shop__lines 6 columns +loaded raw.shop__lines 6 rows 0.1s + +✓ Completed 1 loaded, 6 rows, 0.2s +├── slowest raw.shop__lines 0.1s +└── peak memory 133 MiB +``` + +- `ctx.unit` is `"page"`, `ctx.page` the page's number within its request (from 1), and + `ctx.key` the fields of the key a fan-out requested it for (`None` in a `paginate` or with + `batch=`). `ctx.table` is the table's name. +- A page's nested objects and lists arrive as JSON text, as `lines` does above: `json.loads` them, + or pick their values with [`fields`](#fields) first. +- Every page that's read goes through the function once, including an empty last page, such as the + one that ends `sl.page_number` paging. What it returns or yields loads in page order, batch by + batch, and `flatten`, `columns` and `mask` apply to it. +- A function that raises, or returns or yields something else, fails the table, naming the request: + `GET .../regions/north/orders: transform order_lines failed: ...`. +- `plan` never calls the table function, so it checks a `transform="name"` written as a literal and + fails when no function, or two, are registered under it, as for [field paths](#fields): + + ```console + $ sqlload plan + Plan ready 1 selected + + Errors (1) + └── raw.shop__lines + └── reason transform order_line: no function is registered as order_line with @sl.transform; the registered ones are order_lines + + ✗ Plan failed with errors 1 selected + ``` + +- `transform` is for the job whose rows load; a job given as a fan-out's `over=` refuses it. + ### Multi-level fan-out From 0.37.0. A `fan_out` whose `select` picks keys can be the `over` of another @@ -584,3 +656,5 @@ To inspect or correct a logged position, see [`sqlload logged-state`](/docs/cli/ |------|---------| | `sl.__version__` | The package version. | | `sl.executable()` | The path of the `sqlload` command installed with the package, else the one on `PATH`. | +| `@sl.transform`, `@sl.transform("name")` | From 0.52.0. Registers a function under its own name, or the one given, as a transform for a [files](/docs/sources/files/#transforming-each-file) or [Google Sheets](/docs/sources/google-sheets/#transforming-each-range) table's `transform = "name"`, or an `sl.http` job's [`transform=`](#transforming-each-page). Returns the function unchanged. | +| `sl.TransformContext` | From 0.52.0. The `ctx` a transform gets: `unit` (`"file"`, `"range"` or `"page"`), `table`, and the unit's details. | From c9096164390f2eb8c5009bcdb1267fa91006cc5d Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 15 Nov 2026 00:46:09 +0000 Subject: [PATCH 2/7] feat(odbc): read long columns value by value under a per-table fetch buffer, and show it in plan A text or binary column whose buffer would pass long_column_threshold (4 KiB by default), or that the driver gives no length for, makes its table read value by value with SQLGetData: one row at a time, each value at its own size, never cut. Block cursors are avoided because reading values from a block of rows breaks on MariaDB Connector/ODBC 3.1 and on Microsoft ODBC Driver 18. Tables without such a column still fetch in bulk. Each table has a fetch buffer, the most the rows of one fetch hold: 512 MiB, lowered to half the memory limit (the cgroup's memory.max in a container, else the machine's memory) shared between the tables loading at once, and never raised unless the source sets max_batch_memory. A bulk fetch reads as many rows as its buffers hold within it; a value-by-value fetch ends once its values reach it, and a single value larger than the buffer loads, ends its fetch and is logged. Commits are unchanged. The fetch buffer bounds the fetches only, not the load's whole memory. plan shows each ODBC table's fetch and fetch buffer, and warns when the fetch buffers of the tables loading at once together pass half the limit; load checks again. --- crates/sqlload/src/loader/_helpers/python.rs | 1 + crates/sqlload/src/loader/_helpers/render.rs | 4 + crates/sqlload/src/loader/_helpers/sources.rs | 62 +++ .../sqlload/src/loader/_helpers/table_plan.rs | 4 + crates/sqlload/src/loader/main/_plan_files.rs | 1 + .../sqlload/src/loader/main/plan_project.rs | 28 + crates/sqlload/src/loader/models.rs | 17 + .../src/odbc_source/_helpers/batches.rs | 12 +- .../src/odbc_source/_helpers/columns.rs | 122 ++++ .../src/odbc_source/_helpers/memory.rs | 255 +++++++++ .../sqlload/src/odbc_source/_helpers/mod.rs | 3 + .../src/odbc_source/_helpers/values.rs | 527 ++++++++++++++++++ crates/sqlload/src/odbc_source/constants.rs | 58 ++ .../src/odbc_source/main/memory_warning.rs | 42 ++ crates/sqlload/src/odbc_source/main/mod.rs | 2 + .../src/odbc_source/main/plan_fetches.rs | 94 ++++ .../sqlload/src/odbc_source/main/read_rows.rs | 48 +- crates/sqlload/src/odbc_source/models.rs | 68 ++- crates/sqlload/src/odbc_source/tests.rs | 2 + .../src/odbc_source/tests/test_memory.rs | 374 +++++++++++++ .../odbc_source/tests/test_select_query.rs | 1 + .../src/odbc_source/tests/test_types.rs | 53 ++ .../src/project/_helpers/references.rs | 10 + crates/sqlload/src/project/constants.rs | 2 + crates/sqlload/src/project/models.rs | 8 +- .../src/project/tests/test_parse_project.rs | 10 + crates/sqlload/tests/loader.rs | 3 + crates/sqlload/tests/loader/odbc/helpers.rs | 112 ++++ .../sqlload/tests/loader/test_odbc_memory.rs | 287 ++++++++++ .../tests/loader/test_odbc_to_duckdb.rs | 10 +- crates/sqlload/tests/loader/test_types.rs | 43 ++ .../src/content/docs/docs/configuration.mdx | 2 + .../src/content/docs/docs/sources/odbc.mdx | 100 +++- .../src/content/docs/docs/troubleshooting.mdx | 5 + 34 files changed, 2346 insertions(+), 24 deletions(-) create mode 100644 crates/sqlload/src/odbc_source/_helpers/columns.rs create mode 100644 crates/sqlload/src/odbc_source/_helpers/memory.rs create mode 100644 crates/sqlload/src/odbc_source/_helpers/values.rs create mode 100644 crates/sqlload/src/odbc_source/main/memory_warning.rs create mode 100644 crates/sqlload/src/odbc_source/main/plan_fetches.rs create mode 100644 crates/sqlload/src/odbc_source/tests/test_memory.rs create mode 100644 crates/sqlload/tests/loader/test_odbc_memory.rs diff --git a/crates/sqlload/src/loader/_helpers/python.rs b/crates/sqlload/src/loader/_helpers/python.rs index d76058de8..d192c48d1 100644 --- a/crates/sqlload/src/loader/_helpers/python.rs +++ b/crates/sqlload/src/loader/_helpers/python.rs @@ -1394,6 +1394,7 @@ fn child_plan(parent: &TablePlan, (path, table): (&str, &str)) -> TablePlan { commits: CommitRule::Size, stage: parent.stage.clone(), planned_stage: None, + fetch_memory: None, } } diff --git a/crates/sqlload/src/loader/_helpers/render.rs b/crates/sqlload/src/loader/_helpers/render.rs index 063b02286..49f34def1 100644 --- a/crates/sqlload/src/loader/_helpers/render.rs +++ b/crates/sqlload/src/loader/_helpers/render.rs @@ -673,6 +673,10 @@ pub(crate) fn load_item(table: &TablePlan) -> PlanItem { )); } details.push(detail("batch", batch_line(table))); + if let Some(fetch) = &table.fetch_memory { + details.push(detail("fetch", fetch.memory.as_str())); + details.push(detail("buffer", fetch.budget_note.as_str())); + } if matches!(table.commits, CommitRule::Windows(_)) && spec.python.is_none() && spec.files.is_none() diff --git a/crates/sqlload/src/loader/_helpers/sources.rs b/crates/sqlload/src/loader/_helpers/sources.rs index b34b84ba9..4b1c0c114 100644 --- a/crates/sqlload/src/loader/_helpers/sources.rs +++ b/crates/sqlload/src/loader/_helpers/sources.rs @@ -368,6 +368,9 @@ pub(crate) fn odbc_source_reader( let limits = crate::odbc_source::models::ValueLimits { max_text_size: source.max_text_size, max_binary_size: source.max_binary_size, + long_column_threshold: source + .long_column_threshold + .map(|size| usize::try_from(size.0).unwrap_or(usize::MAX)), }; crate::odbc_source::main::odbc_reader::odbc_reader( (connection_string.as_ref(), dsn.as_deref()), @@ -377,6 +380,61 @@ pub(crate) fn odbc_source_reader( .map_err(LoaderError::Invalid) } +/// `planned`, each unblocked table of ODBC source `source` given how its rows are fetched on this machine: its budget, and in bulk or value by value (see [`plan_fetches`](plan_fetches)). +#[cfg(feature = "odbc")] +pub(crate) async fn with_fetch_memory( + project: &ProjectConfig, + source: &SourceConfig, + mut planned: Vec, + options: &PlanOptions, +) -> Result, LoaderError> { + let reader = odbc_source_reader(project, source)?; + let open: Vec = (0..planned.len()) + .filter(|&index| planned[index].schema.blocked.is_none()) + .collect(); + let requests: Vec> = open + .iter() + .map(|&index| crate::odbc_source::models::OdbcReadRequest { + table: &planned[index].spec.source_table, + sql: planned[index] + .spec + .warehouse_source + .as_ref() + .and_then(|source| source.sql.as_deref()), + columns: &planned[index].read_columns, + filter: None, + parameters: &[], + batch_rows: SOURCE_BATCH_ROWS, + memory_budget: 0, + }) + .collect(); + let set = source + .max_batch_memory + .map(|size| usize::try_from(size.0).unwrap_or(usize::MAX)); + let fetches = crate::odbc_source::main::plan_fetches::plan_fetches( + &reader, + &requests, + ( + set, + options + .concurrency + .or(project.concurrency) + .unwrap_or(crate::loader::constants::DEFAULT_THREADS), + ), + ) + .await; + for (index, fetch) in open.into_iter().zip(fetches) { + planned[index].fetch_memory = Some(crate::loader::models::FetchMemory { + budget: fetch.budget, + estimate: fetch.estimate, + memory: fetch.memory, + budget_note: fetch.budget_note, + limit: fetch.limit, + }); + } + Ok(planned) +} + #[cfg(feature = "odbc")] fn open_odbc(project: &ProjectConfig, source: &SourceConfig) -> Result { Ok(SourceReader::Odbc(Box::new(odbc_source_reader( @@ -1414,6 +1472,10 @@ async fn read_each( filter: filter.as_deref(), parameters, batch_rows: SOURCE_BATCH_ROWS, + memory_budget: table.fetch_memory.as_ref().map_or( + crate::odbc_source::constants::DEFAULT_BATCH_MEMORY, + |fetch| fetch.budget, + ), }; Ok(crate::odbc_source::main::read_rows::read_rows(reader, &request, on_batch).await?) } diff --git a/crates/sqlload/src/loader/_helpers/table_plan.rs b/crates/sqlload/src/loader/_helpers/table_plan.rs index 0bb2130bf..5e26b0861 100644 --- a/crates/sqlload/src/loader/_helpers/table_plan.rs +++ b/crates/sqlload/src/loader/_helpers/table_plan.rs @@ -220,6 +220,7 @@ pub(crate) fn blocked_table(spec: &TableSpec, reasons: Vec) -> TablePlan commits: CommitRule::Whole(WholeReason::Blocked), stage: None, planned_stage: None, + fetch_memory: None, } } @@ -789,6 +790,7 @@ async fn plan_table_once( commits: CommitRule::Whole(WholeReason::Blocked), stage: None, planned_stage: None, + fetch_memory: None, }); } }; @@ -981,6 +983,7 @@ async fn plan_table_once( commits, stage, planned_stage, + fetch_memory: None, }) } @@ -1886,6 +1889,7 @@ pub(crate) async fn block_refused_bounds( filter: Some(&filter), parameters: ¶meters, batch_rows: 1, + memory_budget: crate::odbc_source::constants::DEFAULT_BATCH_MEMORY, }; crate::odbc_source::main::run_query::run_query(&reader, &request) .await diff --git a/crates/sqlload/src/loader/main/_plan_files.rs b/crates/sqlload/src/loader/main/_plan_files.rs index 1c9b5cc3e..2b3e47013 100644 --- a/crates/sqlload/src/loader/main/_plan_files.rs +++ b/crates/sqlload/src/loader/main/_plan_files.rs @@ -106,6 +106,7 @@ pub(in crate::loader) async fn plan_files_table( commits, stage: None, planned_stage: None, + fetch_memory: None, output: Some(FilesOutput { path, encoding: destination.encoding.clone(), diff --git a/crates/sqlload/src/loader/main/plan_project.rs b/crates/sqlload/src/loader/main/plan_project.rs index 833f8215c..9776707d2 100644 --- a/crates/sqlload/src/loader/main/plan_project.rs +++ b/crates/sqlload/src/loader/main/plan_project.rs @@ -13,6 +13,8 @@ use crate::loader::_helpers::python::python_table; use crate::loader::_helpers::resolve::{ admit_tables, destination_names, spec_names, unmatched_selectors, }; +#[cfg(feature = "odbc")] +use crate::loader::_helpers::sources::with_fetch_memory; use crate::loader::_helpers::sources::{PLANNED_SOURCES, found_sources, selected, source_ends}; #[cfg(feature = "odbc")] use crate::loader::_helpers::table_plan::block_refused_bounds; @@ -110,6 +112,10 @@ async fn planned_tables( #[cfg(feature = "odbc")] block_refused_bounds(project, (source_name, source), &mut planned).await?; } + #[cfg(feature = "odbc")] + if source.kind == SourceKind::Odbc { + planned = with_fetch_memory(project, source, planned, options).await?; + } Ok((planned, unselected)) } @@ -123,6 +129,28 @@ fn block_conflicts( names.extend(drained.iter().cloned()); plan.warnings .extend(unmatched_selectors(project, options, &names)); + #[cfg(feature = "odbc")] + { + let fetches: Vec<(usize, &str, Option)> = plan + .tables + .iter() + .filter(|table| table.schema.blocked.is_none()) + .filter_map(|table| { + let fetch = table.fetch_memory.as_ref()?; + Some(( + fetch.estimate?, + table.spec.source_name.as_str(), + fetch.limit, + )) + }) + .collect(); + let concurrency = options.concurrency.or(project.concurrency); + plan.warnings + .extend(crate::odbc_source::main::memory_warning::memory_warning( + &fetches, + concurrency.unwrap_or(crate::loader::constants::DEFAULT_THREADS), + )); + } let mut plan = block_shared_destinations(plan, configured); for table in plan .tables diff --git a/crates/sqlload/src/loader/models.rs b/crates/sqlload/src/loader/models.rs index 7912bcbf9..32bab4a65 100644 --- a/crates/sqlload/src/loader/models.rs +++ b/crates/sqlload/src/loader/models.rs @@ -436,6 +436,23 @@ pub struct TablePlan { pub stage: Option, /// The stage `apply` must create. pub planned_stage: Option, + /// How an ODBC table's rows are fetched under its memory budget. + pub fetch_memory: Option, +} + +/// How an ODBC table's rows are fetched, worked out at plan on the machine that loads. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FetchMemory { + /// The bytes one fetch may hold. + pub budget: usize, + /// The bytes one fetch holds at most, as planned: its bound buffers, or the budget for a table read value by value; `None` when the driver couldn't describe the columns. + pub estimate: Option, + /// What `plan` says about the fetches, the `fetch` line. + pub memory: String, + /// The fetch buffer and why, `plan`'s `buffer` line. + pub budget_note: String, + /// The machine's memory limit the budget was worked out from, in bytes; `None` when it can't be read. + pub limit: Option, } /// A logged table's run log and what it means for this run. diff --git a/crates/sqlload/src/odbc_source/_helpers/batches.rs b/crates/sqlload/src/odbc_source/_helpers/batches.rs index 162916961..18b746926 100644 --- a/crates/sqlload/src/odbc_source/_helpers/batches.rs +++ b/crates/sqlload/src/odbc_source/_helpers/batches.rs @@ -3,25 +3,25 @@ use std::sync::Arc; use arrow_odbc::odbc_api::Cursor; -use arrow_odbc::{ColumnFailure, OdbcReaderBuilder, arrow_schema_from}; +use arrow_odbc::{ColumnFailure, OdbcReaderBuilder}; use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; use crate::odbc_source::constants::TOO_LARGE_VALUE; use crate::odbc_source::errors::OdbcSourceError; use crate::odbc_source::models::OdbcReader; -/// The reader of `cursor`'s rows in batches of `batch_rows`, as `fetch_schema` fetches them, each buffer capped by the reader's limits; a column the driver gives no length for is refused, naming the setting that bounds it. +/// The reader of `cursor`'s rows as `schema` (from `fetch_schema`) in batches of `batch_rows`, each buffer capped by the reader's limits; a column the driver gives no length for is refused, naming the setting that bounds it. The rows per batch are chosen by the caller from its memory budget, so arrow-odbc's own byte limit is lifted. pub(crate) fn batch_reader( reader: &OdbcReader, - mut cursor: C, - (dbms, batch_rows): (String, usize), + cursor: C, + (dbms, schema, batch_rows): (String, SchemaRef, usize), ) -> Result, OdbcSourceError> { - let schema = fetch_schema(&arrow_schema_from(&mut cursor, Some(&dbms), false)?); let mut builder = OdbcReaderBuilder::new(); builder .with_schema(schema) .with_dbms_name(dbms) - .with_max_num_rows_per_batch(batch_rows) + .with_max_num_rows_per_batch(batch_rows.max(1)) + .with_max_bytes_per_batch(usize::MAX) .with_fallibale_allocations(true); if let Some(size) = reader.limits.max_text_size { builder.with_max_text_size(size); diff --git a/crates/sqlload/src/odbc_source/_helpers/columns.rs b/crates/sqlload/src/odbc_source/_helpers/columns.rs new file mode 100644 index 000000000..b01e81de6 --- /dev/null +++ b/crates/sqlload/src/odbc_source/_helpers/columns.rs @@ -0,0 +1,122 @@ +//! Which columns of a read are too long to bind, and the buffer the others take per row. + +use arrow_odbc::odbc_api::ResultSetMetadata; +use arrow_odbc::odbc_api::buffers::BufferDesc; +use arrow_schema::{DataType, Schema, TimeUnit}; + +use crate::odbc_source::constants::{DEFAULT_LONG_COLUMN_THRESHOLD, TIME_TEXT_LENGTHS}; +use crate::odbc_source::errors::OdbcSourceError; +use crate::odbc_source::models::{FetchShape, FetchSize, ValueLimits}; + +/// The shape of a read of `schema`'s columns (as `fetch_schema` fetches them), described by `metadata`: a text or binary column is long when the buffer per value arrow-odbc would bind for it (its described size in bytes, at most `max_text_size` or `max_binary_size`) passes the threshold, or when it has no size; the others' buffers, as arrow-odbc sizes them, make the bytes per row. Text is sized as arrow-odbc reads it: UTF-16 on Windows, UTF-8 elsewhere. +pub(crate) fn fetch_shape( + metadata: &mut impl ResultSetMetadata, + schema: &Schema, + limits: &ValueLimits, +) -> Result { + let threshold = limits + .long_column_threshold + .unwrap_or(DEFAULT_LONG_COLUMN_THRESHOLD); + let wide = cfg!(target_os = "windows"); + let mut shape = FetchShape { + bound_bytes_per_row: 0, + long_columns: Vec::new(), + columns: schema.fields().len(), + }; + for (field, number) in schema.fields().iter().zip(1u16..) { + let nullable = field.is_nullable(); + let buffer = match field.data_type() { + DataType::Utf8 => { + let sql_type = metadata.col_data_type(number)?; + let length = match wide { + true => sql_type.utf16_len(), + false => sql_type.utf8_len(), + }; + let length = match length { + Some(length) => Some(length.get()), + None => metadata.col_display_size(number)?.map(|size| size.get()), + }; + capped(length, limits.max_text_size) + .filter(|&length| length.saturating_mul(if wide { 2 } else { 1 }) <= threshold) + .map(|max_str_len| match wide { + true => BufferDesc::WText { max_str_len }, + false => BufferDesc::Text { max_str_len }, + }) + } + DataType::Binary => { + let length = metadata.col_data_type(number)?.column_size(); + capped(length.map(|length| length.get()), limits.max_binary_size) + .filter(|&length| length <= threshold) + .map(|max_bytes| BufferDesc::Binary { max_bytes }) + } + other => Some(bound_buffer(other, nullable)), + }; + match buffer { + Some(buffer) => shape.bound_bytes_per_row += buffer.bytes_per_row(), + None => shape.long_columns.push(field.name().clone()), + } + } + if !shape.long_columns.is_empty() { + shape.bound_bytes_per_row = 0; + } + Ok(shape) +} + +/// The rows one fetch of `shape` reads under `budget` bytes, at most `batch_rows` and at least one, and the bytes it holds: its bound buffers, or for a table read value by value the budget, since its values count towards it. +pub(crate) fn fetch_size(shape: &FetchShape, budget: usize, batch_rows: usize) -> FetchSize { + let batch_rows = batch_rows.max(1); + match (shape.long_columns.is_empty(), shape.bound_bytes_per_row) { + (false, _) => FetchSize { + rows: batch_rows, + bytes: budget, + }, + (true, 0) => FetchSize { + rows: batch_rows, + bytes: 0, + }, + (true, per_row) => { + let rows = (budget / per_row).clamp(1, batch_rows); + FetchSize { + rows, + bytes: rows.saturating_mul(per_row), + } + } + } +} + +/// A described size capped by its setting; `None` when neither gives one. +fn capped(length: Option, cap: Option) -> Option { + match (length, cap) { + (Some(length), Some(cap)) => Some(length.min(cap)), + (length, cap) => length.or(cap), + } +} + +/// The buffer arrow-odbc binds for a column of a type other than text or binary. +fn bound_buffer(data_type: &DataType, nullable: bool) -> BufferDesc { + match data_type { + DataType::Boolean => BufferDesc::Bit { nullable }, + DataType::Int8 => BufferDesc::I8 { nullable }, + DataType::Int16 => BufferDesc::I16 { nullable }, + DataType::Int32 => BufferDesc::I32 { nullable }, + DataType::UInt8 => BufferDesc::U8 { nullable }, + DataType::Float32 => BufferDesc::F32 { nullable }, + DataType::Float64 => BufferDesc::F64 { nullable }, + DataType::Date32 => BufferDesc::Date { nullable }, + DataType::Time32(TimeUnit::Second) => BufferDesc::Time { nullable }, + DataType::Time32(_) => BufferDesc::Text { + max_str_len: TIME_TEXT_LENGTHS[0], + }, + DataType::Time64(TimeUnit::Microsecond) => BufferDesc::Text { + max_str_len: TIME_TEXT_LENGTHS[1], + }, + DataType::Time64(_) => BufferDesc::Text { + max_str_len: TIME_TEXT_LENGTHS[2], + }, + DataType::Timestamp(..) => BufferDesc::Timestamp { nullable }, + DataType::FixedSizeBinary(length) => BufferDesc::Binary { + max_bytes: usize::try_from(*length).unwrap_or_default(), + }, + _ => BufferDesc::I64 { nullable }, + } +} diff --git a/crates/sqlload/src/odbc_source/_helpers/memory.rs b/crates/sqlload/src/odbc_source/_helpers/memory.rs new file mode 100644 index 000000000..d2c8d26b8 --- /dev/null +++ b/crates/sqlload/src/odbc_source/_helpers/memory.rs @@ -0,0 +1,255 @@ +//! The memory a machine gives SQLLoad, the budget one fetch of a table takes of it, and how `plan` shows them. + +use std::path::{Path, PathBuf}; + +use crate::odbc_source::constants::{ + BINARY_UNIT, CGROUP_V1_LIMIT, CGROUP_V1_ROOT, CGROUP_V2_LIMIT, CGROUP_V2_PREFIX, + CGROUP_V2_ROOT, DEFAULT_BATCH_MEMORY, FETCH_SHARE_OF_LIMIT, MEMINFO_TOTAL, MEMINFO_UNIT, + MEMORY_CONTROLLER, MEMORY_UNITS, NAMED_LONG_COLUMNS, PROC_CGROUP, PROC_MEMINFO, TENTHS, + UNLIMITED, WHOLE_FROM_TENTHS, +}; +use crate::odbc_source::models::{BudgetReason, FetchShape, FetchSize, MemoryBudget, MemoryLimit}; + +/// The budget of one fetch: `set` (the source's `max_batch_memory`) as it is; else 512 MiB, lowered to half of `limit` shared between `concurrency` tables when that is less (never raised); 512 MiB when the limit is unknown. +pub(crate) fn fetch_budget( + set: Option, + limit: Option, + concurrency: usize, +) -> MemoryBudget { + let concurrency = concurrency.max(1); + let share = limit.map(|limit| { + let share = limit.bytes / FETCH_SHARE_OF_LIMIT as u64 / concurrency as u64; + usize::try_from(share).unwrap_or(usize::MAX).max(1) + }); + let (bytes, reason) = match (set, share) { + (Some(bytes), _) => (bytes.max(1), BudgetReason::Set), + (None, Some(share)) if share < DEFAULT_BATCH_MEMORY => (share, BudgetReason::Capped), + (None, _) => (DEFAULT_BATCH_MEMORY, BudgetReason::Default), + }; + MemoryBudget { + bytes, + reason, + limit, + concurrency, + } +} + +/// The memory this process may use: the lowest cgroup limit over it, when below the machine's memory, else the machine's memory; `None` when neither is known (Windows). Fails, naming the file, when a limit file exists but can't be read or holds no number. +pub(crate) fn memory_limit() -> Result, String> { + let physical = physical_memory()?; + let cgroup = match read_text(Path::new(PROC_CGROUP))? { + Some(groups) => match cgroup_limit(&groups, Path::new(CGROUP_V2_ROOT))? { + Some(limit) => Some(limit), + None => cgroup_v1_limit(&groups)?, + }, + None => None, + }; + Ok(match (cgroup, physical) { + (Some(cgroup), Some(physical)) if cgroup >= physical => Some(MemoryLimit { + bytes: physical, + container: false, + }), + (Some(bytes), _) => Some(MemoryLimit { + bytes, + container: true, + }), + (None, physical) => physical.map(|bytes| MemoryLimit { + bytes, + container: false, + }), + }) +} + +/// The lowest cgroup v2 `memory.max` from the process's group (its `0::/path` line in `groups`, `/proc/self/cgroup`) up to `root`; `None` when there is no such line or every one is `max` or missing. +pub(crate) fn cgroup_limit(groups: &str, root: &Path) -> Result, String> { + let Some(path) = groups + .lines() + .find_map(|line| line.strip_prefix(CGROUP_V2_PREFIX)) + else { + return Ok(None); + }; + let mut folder: PathBuf = root.join(path.trim().trim_start_matches('/')); + let mut lowest: Option = None; + loop { + lowest = match (lowest, limit_file(&folder.join(CGROUP_V2_LIMIT))?) { + (Some(lowest), Some(limit)) => Some(lowest.min(limit)), + (lowest, limit) => lowest.or(limit), + }; + if folder.as_path() == root || !folder.pop() || !folder.starts_with(root) { + return Ok(lowest); + } + } +} + +/// `bytes` in binary units to a tenth, as `512 MiB`, `1.9 GiB` or `40 KiB`; from 100 of a unit up, whole. +pub(crate) fn memory_text(bytes: u64) -> String { + let mut unit = 0; + let mut size: u64 = 1; + while unit + 1 < MEMORY_UNITS.len() && bytes >= size.saturating_mul(BINARY_UNIT) { + size *= BINARY_UNIT; + unit += 1; + } + if unit == 0 { + return format!("{bytes} bytes"); + } + let (size, tenths_per_unit) = (u128::from(size), u128::from(TENTHS)); + let tenths = (u128::from(bytes) * tenths_per_unit + size / 2) / size; + match tenths.is_multiple_of(tenths_per_unit) || tenths >= u128::from(WHOLE_FROM_TENTHS) { + true => format!( + "{} {}", + (tenths + tenths_per_unit / 2) / tenths_per_unit, + MEMORY_UNITS[unit] + ), + false => format!( + "{}.{} {}", + tenths / tenths_per_unit, + tenths % tenths_per_unit, + MEMORY_UNITS[unit] + ), + } +} + +/// `plan`'s `buffer` line: the fetch buffer (the budget) per table, and why; `unreadable` is why the machine's memory couldn't be read. +pub(crate) fn budget_note(budget: &MemoryBudget, unreadable: Option<&str>) -> String { + let bytes = memory_text(budget.bytes as u64); + let limit = budget.limit.map(|limit| match limit.container { + true => format!("{} container limit", memory_text(limit.bytes)), + false => format!("{} of memory", memory_text(limit.bytes)), + }); + let concurrency = budget.concurrency; + match (budget.reason, limit, unreadable) { + (BudgetReason::Set, _, _) => format!("{bytes} fetch buffer per table (max_batch_memory)"), + (BudgetReason::Capped, Some(limit), _) => { + format!("{bytes} fetch buffer per table (capped: {limit}, concurrency {concurrency})") + } + (_, Some(limit), _) => { + format!("{bytes} fetch buffer per table (default; {limit}, concurrency {concurrency})") + } + (_, None, Some(reason)) => { + format!( + "{bytes} fetch buffer per table (default; this machine's memory can't be read: {reason})" + ) + } + (_, None, None) => { + format!("{bytes} fetch buffer per table (default; this machine's memory is unknown)") + } + } +} + +/// `plan`'s `fetch` line for a table read as `shape` in fetches of `size`: in bulk, the bytes and rows of one fetch; value by value, the budget and rows that end a fetch, and the long columns, the first few named. +pub(crate) fn memory_line(shape: &FetchShape, size: &FetchSize) -> String { + let rows = thousands(size.rows); + if shape.long_columns.is_empty() { + return format!("~{} per fetch, {rows} rows", memory_text(size.bytes as u64)); + } + let count = shape.long_columns.len(); + let mut names = shape.long_columns[..count.min(NAMED_LONG_COLUMNS)].join(", "); + if count > NAMED_LONG_COLUMNS { + names.push_str(", …"); + } + let columns = match count { + 1 => format!("1 long column ({names})"), + _ => format!("{count} long columns ({names})"), + }; + format!( + "up to {} or {rows} rows per fetch; {columns}, so all {} columns are read value by value", + memory_text(size.bytes as u64), + shape.columns + ) +} + +/// `count` with thousands separators, as `65,536`. +fn thousands(count: usize) -> String { + let digits = count.to_string(); + let mut grouped = String::new(); + for (index, digit) in digits.chars().enumerate() { + if index > 0 && (digits.len() - index).is_multiple_of(3) { + grouped.push(','); + } + grouped.push(digit); + } + grouped +} + +/// The cgroup v1 memory limit of the process's group (its `memory` line in `groups`), else of the hierarchy's root, as a container sees it; unlimited reads as a huge number, which the machine's memory then undercuts. +fn cgroup_v1_limit(groups: &str) -> Result, String> { + let Some(path) = groups.lines().find_map(memory_group) else { + return Ok(None); + }; + let root = Path::new(CGROUP_V1_ROOT); + match limit_file( + &root + .join(path.trim().trim_start_matches('/')) + .join(CGROUP_V1_LIMIT), + )? { + Some(limit) => Ok(Some(limit)), + None => limit_file(&root.join(CGROUP_V1_LIMIT)), + } +} + +/// The group path of a `/proc/self/cgroup` line of cgroup v1's memory controller, `N:memory:/path`. +fn memory_group(line: &str) -> Option<&str> { + let mut parts = line.splitn(3, ':'); + let (_, controllers, path) = (parts.next()?, parts.next()?, parts.next()?); + controllers + .split(',') + .any(|name| name == MEMORY_CONTROLLER) + .then_some(path) +} + +/// The machine's memory: `MemTotal` on Linux, `hw.memsize` on macOS; `None` elsewhere. +fn physical_memory() -> Result, String> { + if let Some(meminfo) = read_text(Path::new(PROC_MEMINFO))? { + return meminfo_total(&meminfo); + } + if !cfg!(target_os = "macos") { + return Ok(None); + } + let output = std::process::Command::new("sysctl") + .args(["-n", "hw.memsize"]) + .output() + .map_err(|error| format!("sysctl hw.memsize: {error}"))?; + String::from_utf8_lossy(&output.stdout) + .trim() + .parse() + .map(Some) + .map_err(|error| format!("sysctl hw.memsize: {error}")) +} + +/// The bytes of `/proc/meminfo`'s `MemTotal: 16314228 kB` line; `None` without one. +fn meminfo_total(meminfo: &str) -> Result, String> { + let Some(total) = meminfo + .lines() + .find_map(|line| line.strip_prefix(MEMINFO_TOTAL)) + else { + return Ok(None); + }; + let kilobytes = total.trim().trim_end_matches(MEMINFO_UNIT).trim(); + kilobytes + .parse::() + .map(|kilobytes| Some(kilobytes.saturating_mul(BINARY_UNIT))) + .map_err(|error| format!("{PROC_MEMINFO}: MemTotal {kilobytes:?}: {error}")) +} + +/// A cgroup limit file's number; `None` when it is `max` or the file doesn't exist. +fn limit_file(path: &Path) -> Result, String> { + let Some(text) = read_text(path)? else { + return Ok(None); + }; + match text.trim() { + UNLIMITED => Ok(None), + number => number + .parse() + .map(Some) + .map_err(|error| format!("{}: {number:?}: {error}", path.display())), + } +} + +/// A file's text; `None` when it doesn't exist. +fn read_text(path: &Path) -> Result, String> { + match std::fs::read_to_string(path) { + Ok(text) => Ok(Some(text)), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(error) => Err(format!("{}: {error}", path.display())), + } +} diff --git a/crates/sqlload/src/odbc_source/_helpers/mod.rs b/crates/sqlload/src/odbc_source/_helpers/mod.rs index 8c1fd09e3..88045dfdc 100644 --- a/crates/sqlload/src/odbc_source/_helpers/mod.rs +++ b/crates/sqlload/src/odbc_source/_helpers/mod.rs @@ -1,6 +1,9 @@ pub(crate) mod batches; pub(crate) mod catalog; +pub(crate) mod columns; pub(crate) mod connect; pub(crate) mod convert; +pub(crate) mod memory; pub(crate) mod parameters; pub(crate) mod query; +pub(crate) mod values; diff --git a/crates/sqlload/src/odbc_source/_helpers/values.rs b/crates/sqlload/src/odbc_source/_helpers/values.rs new file mode 100644 index 000000000..94ca97734 --- /dev/null +++ b/crates/sqlload/src/odbc_source/_helpers/values.rs @@ -0,0 +1,527 @@ +//! Rows read value by value with `SQLGetData`: every column of a table with a column too long to bind, each value at its own size. + +use std::sync::Arc; + +use arrow_array::builder::{ + BinaryBuilder, BooleanBuilder, Date32Builder, FixedSizeBinaryBuilder, Float32Builder, + Float64Builder, Int8Builder, Int16Builder, Int32Builder, Int64Builder, StringBuilder, + Time32MillisecondBuilder, Time32SecondBuilder, Time64MicrosecondBuilder, + Time64NanosecondBuilder, TimestampMicrosecondBuilder, TimestampNanosecondBuilder, UInt8Builder, +}; +use arrow_array::{ArrayRef, RecordBatch}; +use arrow_odbc::odbc_api::buffers::Indicator; +use arrow_odbc::odbc_api::parameter::{Binary, Text, VarCell, VarKind, WideText}; +use arrow_odbc::odbc_api::sys::{Date, Time, Timestamp}; +use arrow_odbc::odbc_api::{Bit, Cursor, CursorRow, Nullable, Pod}; +use arrow_schema::{DataType, Field, SchemaRef, TimeUnit}; +use chrono::{NaiveDate, NaiveTime, Timelike}; + +use crate::odbc_source::constants::{KEPT_VALUE_BUFFER, VALUE_OVERHEAD}; +use crate::odbc_source::errors::OdbcSourceError; +use crate::odbc_source::models::ValueLimits; + +/// One column read value by value: what it appends to, and the buffer its text or binary values are read into. +struct ValueColumn { + name: String, + values: Values, + bytes: Vec, + wide: Vec, +} + +/// The Arrow array a column's values are appended to, by its fetched type. +enum Values { + Boolean(BooleanBuilder), + Int8(Int8Builder), + Int16(Int16Builder), + Int32(Int32Builder), + Int64(Int64Builder), + UInt8(UInt8Builder), + Float32(Float32Builder), + Float64(Float64Builder), + Date(Date32Builder), + TimeSeconds(Time32SecondBuilder), + TimeMilliseconds(Time32MillisecondBuilder), + TimeMicroseconds(Time64MicrosecondBuilder), + TimeNanoseconds(Time64NanosecondBuilder), + TimestampMicroseconds(TimestampMicrosecondBuilder, Option>), + TimestampNanoseconds(TimestampNanosecondBuilder, Option>), + Text(StringBuilder), + Binary(BinaryBuilder), + FixedBinary(FixedSizeBinaryBuilder, i32), +} + +/// Fetch `cursor`'s rows one at a time, reading every value of `schema`'s columns (as `fetch_schema` fetches them) with `SQLGetData` in column order, and hand them to `send` in batches that end at `max_rows` rows or once their values hold `budget` bytes; a value longer than `max_text_size` or `max_binary_size` fails the read rather than being cut. A single value bigger than the budget loads, ends its batch, and is logged with `table`'s name. Stops when the rows end or `send` says no more. +pub(crate) fn read_values( + (mut cursor, schema): (C, &SchemaRef), + (budget, max_rows): (usize, usize), + (limits, table): (&ValueLimits, &str), + mut send: impl FnMut(RecordBatch) -> bool, +) -> Result<(), OdbcSourceError> { + let mut columns: Vec = schema + .fields() + .iter() + .map(|field| ValueColumn::new(field)) + .collect::>()?; + loop { + let (mut rows, mut bytes, mut ended) = (0usize, 0usize, false); + while rows < max_rows.max(1) && bytes < budget { + let Some(mut row) = cursor.next_row()? else { + ended = true; + break; + }; + for (column, number) in columns.iter_mut().zip(1u16..) { + let (next, size) = column.read(row, number, limits)?; + row = next; + if size > budget { + tracing::warn!( + table, + column = %column.name, + bytes = size, + budget, + "a value is larger than the table's fetch buffer (max_batch_memory); it loads, and its fetch ends after its row" + ); + } + bytes = bytes.saturating_add(size + VALUE_OVERHEAD); + } + rows += 1; + } + if rows > 0 { + tracing::debug!(table, rows, bytes, budget, "fetched rows value by value"); + let arrays: Vec = columns.iter_mut().map(ValueColumn::finish).collect(); + let batch = RecordBatch::try_new(schema.clone(), arrays) + .map_err(|error| OdbcSourceError::Invalid(error.to_string()))?; + if !send(batch) { + return Ok(()); + } + } + if ended { + return Ok(()); + } + } +} + +impl ValueColumn { + /// An empty column of `field`'s fetched type. + fn new(field: &Field) -> Result { + let values = match field.data_type() { + DataType::Boolean => Values::Boolean(BooleanBuilder::new()), + DataType::Int8 => Values::Int8(Int8Builder::new()), + DataType::Int16 => Values::Int16(Int16Builder::new()), + DataType::Int32 => Values::Int32(Int32Builder::new()), + DataType::Int64 => Values::Int64(Int64Builder::new()), + DataType::UInt8 => Values::UInt8(UInt8Builder::new()), + DataType::Float32 => Values::Float32(Float32Builder::new()), + DataType::Float64 => Values::Float64(Float64Builder::new()), + DataType::Date32 => Values::Date(Date32Builder::new()), + DataType::Time32(TimeUnit::Second) => Values::TimeSeconds(Time32SecondBuilder::new()), + DataType::Time32(TimeUnit::Millisecond) => { + Values::TimeMilliseconds(Time32MillisecondBuilder::new()) + } + DataType::Time64(TimeUnit::Microsecond) => { + Values::TimeMicroseconds(Time64MicrosecondBuilder::new()) + } + DataType::Time64(TimeUnit::Nanosecond) => { + Values::TimeNanoseconds(Time64NanosecondBuilder::new()) + } + DataType::Timestamp(TimeUnit::Microsecond, zone) => { + Values::TimestampMicroseconds(TimestampMicrosecondBuilder::new(), zone.clone()) + } + DataType::Timestamp(TimeUnit::Nanosecond, zone) => { + Values::TimestampNanoseconds(TimestampNanosecondBuilder::new(), zone.clone()) + } + DataType::Utf8 => Values::Text(StringBuilder::new()), + DataType::Binary => Values::Binary(BinaryBuilder::new()), + DataType::FixedSizeBinary(length) => { + Values::FixedBinary(FixedSizeBinaryBuilder::new(*length), *length) + } + other => { + return Err(OdbcSourceError::Invalid(format!( + "column {} is fetched as {other}, which SQLLoad doesn't read value by value", + field.name() + ))); + } + }; + Ok(Self { + name: field.name().clone(), + values, + bytes: Vec::new(), + wide: Vec::new(), + }) + } + + /// Read the current row's value of column `number` from `row` and append it; returns the row, and the value's bytes (a text or binary value's own length). + fn read<'r>( + &mut self, + mut row: CursorRow<'r>, + number: u16, + limits: &ValueLimits, + ) -> Result<(CursorRow<'r>, usize), OdbcSourceError> { + let row_ref = &mut row; + let name = self.name.as_str(); + let failed = |error: arrow_odbc::odbc_api::Error| { + OdbcSourceError::Invalid(format!("column {name}: ODBC: {error}")) + }; + let size = match &mut self.values { + Values::Boolean(builder) => { + builder.append_option( + fixed::(row_ref, number) + .map_err(failed)? + .map(Bit::as_bool), + ); + size_of::() + } + Values::Int8(builder) => { + builder.append_option(fixed(row_ref, number).map_err(failed)?); + size_of::() + } + Values::Int16(builder) => { + builder.append_option(fixed(row_ref, number).map_err(failed)?); + size_of::() + } + Values::Int32(builder) => { + builder.append_option(fixed(row_ref, number).map_err(failed)?); + size_of::() + } + Values::Int64(builder) => { + builder.append_option(fixed(row_ref, number).map_err(failed)?); + size_of::() + } + Values::UInt8(builder) => { + builder.append_option(fixed(row_ref, number).map_err(failed)?); + size_of::() + } + Values::Float32(builder) => { + builder.append_option(fixed(row_ref, number).map_err(failed)?); + size_of::() + } + Values::Float64(builder) => { + builder.append_option(fixed(row_ref, number).map_err(failed)?); + size_of::() + } + Values::Date(builder) => { + let days = match fixed::(row_ref, number).map_err(failed)? { + Some(date) => Some(epoch_days(&date).ok_or_else(|| invalid(name, "date"))?), + None => None, + }; + builder.append_option(days); + size_of::() + } + Values::TimeSeconds(builder) => { + let time = fixed::
__...` that have a parent-id column, staging tables and the swap tables named after them, in any case, left out); none without `unnest`. A table named so without that column is another table, and so are the child tables of such a table (`orders__notes__items` beside a table `orders__notes`). pub(crate) async fn live_child_tables( destination: &Destination, @@ -1358,6 +1380,7 @@ fn child_plan(parent: &TablePlan, (path, table): (&str, &str)) -> TablePlan { module: String::new(), declared: Vec::new(), lists: ListChoices::new(), + literal_paths: None, }), sheet: None, transform: None, diff --git a/crates/sqlload/src/loader/_helpers/render.rs b/crates/sqlload/src/loader/_helpers/render.rs index 4ea77c0b4..df128afec 100644 --- a/crates/sqlload/src/loader/_helpers/render.rs +++ b/crates/sqlload/src/loader/_helpers/render.rs @@ -673,6 +673,12 @@ pub(crate) fn load_item(table: &TablePlan) -> PlanItem { )); } details.push(detail("batch", batch_line(table))); + details.extend( + spec.python + .iter() + .filter_map(|python| python.literal_paths.as_ref()?.checked_at_load()) + .map(|text| detail("paths", text)), + ); if let Some(fetch) = &table.fetch_memory { details.push(detail("fetch", fetch.memory.as_str())); details.push(detail("buffer", fetch.budget_note.as_str())); diff --git a/crates/sqlload/src/loader/_helpers/table_plan.rs b/crates/sqlload/src/loader/_helpers/table_plan.rs index b711da0cb..663498c3e 100644 --- a/crates/sqlload/src/loader/_helpers/table_plan.rs +++ b/crates/sqlload/src/loader/_helpers/table_plan.rs @@ -44,6 +44,7 @@ use crate::load_naming::main::is_load_column::is_load_column; use crate::loader::_helpers::destination::truncated_nanoseconds; use crate::loader::_helpers::python::list_choices; use crate::loader::_helpers::python::live_child_tables; +use crate::loader::_helpers::python::python_problems; use crate::loader::_helpers::python::with_live_children; use crate::loader::_helpers::python::with_live_columns; use crate::loader::_helpers::render::up_to_date; @@ -55,7 +56,6 @@ use crate::loader::_helpers::sources::plan_files; use crate::loader::_helpers::sources::{FilesPlanRequest, LoadedSource, with_live_types}; use crate::loader::_helpers::specs::files_read_columns; use crate::loader::constants::CHECKS_LOADER_WARNING; -use crate::loader::constants::ENFORCED_UNDECLARED_PROBLEM; use crate::loader::constants::KEPT_TABLES_RELOAD_PROBLEM; use crate::loader::constants::RANGE_NEEDS_CURSOR; use crate::loader::errors::LoaderError; @@ -78,7 +78,6 @@ use crate::postgres_source::main::plan_read::plan_read; use crate::postgres_source::models::SourceTable; use crate::postgres_source::models::{PgType, ReadColumn}; use crate::project::models::BatchStep; -use crate::project::models::Contract; use crate::project::models::MaskMethod; use crate::project::models::WriteStrategy; use crate::project::models::{FilesFormat, OutputEncoding}; @@ -232,15 +231,9 @@ pub(crate) async fn plan_table( destination: &Destination, options: &PlanOptions, ) -> Result { - let undeclared = spec - .python - .as_ref() - .is_some_and(|python| python.declared.is_empty()); - if undeclared && spec.contract == Contract::Enforced { - return Ok(blocked_table( - spec, - vec![ENFORCED_UNDECLARED_PROBLEM.to_owned()], - )); + let problems = python_problems(spec); + if !problems.is_empty() { + return Ok(blocked_table(spec, problems)); } let children = match spec.unnest { Unnest::All => live_child_tables(destination, spec).await?, diff --git a/crates/sqlload/src/loader/main/plan_project.rs b/crates/sqlload/src/loader/main/plan_project.rs index 96072c2be..4a6a6ce39 100644 --- a/crates/sqlload/src/loader/main/plan_project.rs +++ b/crates/sqlload/src/loader/main/plan_project.rs @@ -94,11 +94,12 @@ async fn planned_tables( (destinations, options): (&BTreeMap, &PlanOptions), ) -> Result<(Vec, Vec), LoaderError> { let module = project.python_modules.get(source_name); + let paths = project.python_http_paths.get(source_name); let mut planned: Vec = Vec::new(); let mut unselected: Vec = Vec::new(); for mut spec in specs { spec.files = files.clone(); - spec.python = module.map(|module| python_table(module, &spec.source_table)); + spec.python = module.map(|module| python_table(module, &spec.source_table, paths)); if !selected(options, (source_name, &spec_names(&spec))) { unselected.push(spec); continue; diff --git a/crates/sqlload/src/loader/main/selected_specs.rs b/crates/sqlload/src/loader/main/selected_specs.rs index 33f1961e8..fb391c626 100644 --- a/crates/sqlload/src/loader/main/selected_specs.rs +++ b/crates/sqlload/src/loader/main/selected_specs.rs @@ -27,6 +27,7 @@ pub async fn selected_specs( } }; let module = project.python_modules.get(source_name); + let paths = project.python_http_paths.get(source_name); chosen.problems.extend(resolved.problems); chosen .problems @@ -36,7 +37,7 @@ pub async fn selected_specs( ))); for mut spec in resolved.specs { spec.files = files.clone(); - spec.python = module.map(|module| python_table(module, &spec.source_table)); + spec.python = module.map(|module| python_table(module, &spec.source_table, paths)); let picked = spec_names(&spec); names.extend(picked.iter().cloned()); if selected(options, (source_name, &picked)) { diff --git a/crates/sqlload/src/loader/models.rs b/crates/sqlload/src/loader/models.rs index 15dd92838..e27e7d78e 100644 --- a/crates/sqlload/src/loader/models.rs +++ b/crates/sqlload/src/loader/models.rs @@ -281,6 +281,8 @@ pub struct PythonTable { pub declared: Vec, /// Under `unnest = "all"`, the columns chosen from the yielded values as list fields or not so far. pub lists: crate::row_batches::types::ListChoices, + /// The `sl.http` paths its function writes as literals, which `plan` checks; `None` when it builds no job `plan` can see. + pub literal_paths: Option, } /// The cursor range a load reads: `from <= cursor < to` (`from < cursor` when not inclusive), each open when `None`. diff --git a/crates/sqlload/src/project/main/parse_python_project.rs b/crates/sqlload/src/project/main/parse_python_project.rs index 80159c3a6..481eb9ee8 100644 --- a/crates/sqlload/src/project/main/parse_python_project.rs +++ b/crates/sqlload/src/project/main/parse_python_project.rs @@ -34,6 +34,10 @@ pub fn parse_python_project( .map_err(|error: toml::de::Error| ProjectError::Syntax(error.to_string()))?; project.connection_settings = settings; project.python_modules = modules; + project.python_http_paths = python + .iter() + .map(|source| (source.name.clone(), source.http_paths.clone())) + .collect(); project.connection_variables = variables; let problems = reference_problems(&project); match problems.is_empty() { diff --git a/crates/sqlload/src/project/models.rs b/crates/sqlload/src/project/models.rs index ebeda4947..6d82e749b 100644 --- a/crates/sqlload/src/project/models.rs +++ b/crates/sqlload/src/project/models.rs @@ -10,6 +10,7 @@ use crate::files_source::models::{ AfterLoad, Compression, CsvOptions, FilesEncoding, RecordShape, SourceFormat, UnmatchedPaths, XlsxOptions, XmlOptions, }; +use crate::http_source::models::LiteralPaths; use crate::load_naming::models::Naming; use crate::project::_helpers::defaults; use crate::project::_helpers::serde_values::{ @@ -46,6 +47,9 @@ pub struct ProjectConfig { /// The file declaring each Python source, by source name. #[serde(skip)] pub python_modules: BTreeMap, + /// The `sl.http` paths each Python table writes as literals, by source name and table name. + #[serde(skip)] + pub python_http_paths: BTreeMap>, /// The `${ENV:NAME}` variables each connection's settings read, by connection name, for hints. #[serde(skip)] pub connection_variables: BTreeMap>, diff --git a/crates/sqlload/src/python_sources/models.rs b/crates/sqlload/src/python_sources/models.rs index a23747e92..606d37528 100644 --- a/crates/sqlload/src/python_sources/models.rs +++ b/crates/sqlload/src/python_sources/models.rs @@ -9,6 +9,9 @@ pub struct PythonSource { pub module: String, /// The source's settings with TOML key names, its tables under `tables` keyed by table name. pub settings: serde_json::Map, + /// The `sl.http` paths each table function writes as literals, by table name; tables that build no job are left out. + #[serde(default)] + pub http_paths: std::collections::BTreeMap, } /// What one table function is called with: its connection, the cursor window and the run. diff --git a/packages/sqlload/python/sqlload/_http_paths.py b/packages/sqlload/python/sqlload/_http_paths.py new file mode 100644 index 000000000..8fc024cdb --- /dev/null +++ b/packages/sqlload/python/sqlload/_http_paths.py @@ -0,0 +1,192 @@ +"""The `sl.http` field paths a table function writes as literals, read from its source without calling it. + +`plan` checks them with load's own path checker, so a bad literal path fails `plan` instead of the load. Only calls whose +names resolve to sqlload count: `paginate` and `fan_out` on an `sl.http(...)` client, and `sl.field`, `sl.parent` and +`sl.key` in their `fields=`. A `transform=` given as a name must be one `@sl.transform` registered. A path `plan` can't +read as a literal (a variable, an f-string, a dict built elsewhere), or a function whose source can't be read, is left to +load, which checks every job before its first request. +""" + +import ast +import inspect +import json +import sys +import textwrap +import types +from collections.abc import Callable, Iterator, Mapping +from typing import Any + +from sqlload import _data, _http +from sqlload._transform import TRANSFORMS + +#: The methods of an `sl.http` client that build a job. +JOB_METHODS: frozenset[str] = frozenset({"paginate", "fan_out"}) + +#: The module a name must resolve into for its call to count. +PACKAGE = "sqlload" + +#: A name or attribute that doesn't resolve to anything `plan` may look at. +_UNRESOLVED: Any = object() + +#: A local name assigned an `sl.http(...)` client inside the table function. +_LOCAL_CLIENT: Any = object() + + +def literal_paths(function: Callable[..., Any]) -> dict[str, Any] | None: + """The literal paths of the jobs `function` builds that load rows: each `fields=` entry's `[name, path]` and each + `key_path`; why a `transform=` name can't be looked up; what `plan` can't read as a literal; and whether the source + can't be read at all. A job passed as another + job's `over=` is left out: load refuses its `fields` and `key_path` whatever their paths. `None` when the function builds + no job `plan` can see.""" + try: + tree = ast.parse(textwrap.dedent(inspect.getsource(function))) + except (OSError, TypeError, SyntaxError): + return {"fields": [], "key_paths": [], "computed": [], "no_source": True} + definition = next((node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))), None) + if definition is None: + return {"fields": [], "key_paths": [], "computed": [], "no_source": True} + names = _Names(function, definition) + jobs = [call for call in ast.walk(definition) if isinstance(call, ast.Call) and _job_method(call, names)] + if not jobs: + return None + found: dict[str, Any] = {"fields": [], "key_paths": [], "computed": [], "no_source": False, "transform_problems": []} + over = _over_jobs(definition, jobs) + for call in jobs: + if id(call) not in over: + _read_call(call, names, found) + return found + + +class _Names: + """What the names a table function uses resolve to: its own imports and `sl.http(...)` clients, else its closure and + module globals; a name it assigns otherwise shadows them and resolves to nothing.""" + + def __init__(self, function: Callable[..., Any], definition: ast.FunctionDef | ast.AsyncFunctionDef) -> None: + self.outer: dict[str, Any] = dict(getattr(function, "__globals__", {})) + code = getattr(function, "__code__", None) + for name, cell in zip(code.co_freevars if code else (), getattr(function, "__closure__", None) or ()): + try: + self.outer[name] = cell.cell_contents + except ValueError: + self.outer[name] = _UNRESOLVED + self.local: dict[str, Any] = {} + arguments = definition.args + for argument in [*arguments.posonlyargs, *arguments.args, *arguments.kwonlyargs, arguments.vararg, arguments.kwarg]: + if argument is not None: + self.local[argument.arg] = _UNRESOLVED + for node in ast.walk(definition): + self._bind(node) + + def _bind(self, node: ast.AST) -> None: + if isinstance(node, ast.Import): + for alias in node.names: + bound = alias.asname or alias.name.partition(".")[0] + self.local[bound] = sys.modules.get(alias.name if alias.asname else bound, _UNRESOLVED) + elif isinstance(node, ast.ImportFrom): + module = sys.modules.get(node.module or "") if node.level == 0 else None + for alias in node.names: + self.local[alias.asname or alias.name] = getattr(module, alias.name, _UNRESOLVED) if _ours(module) else _UNRESOLVED + elif isinstance(node, ast.Assign) and len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + client = isinstance(node.value, ast.Call) and self.resolve(node.value.func) in (_http.http, _http.Http) + self.local[node.targets[0].id] = _LOCAL_CLIENT if client and self.local.get(node.targets[0].id, _LOCAL_CLIENT) is _LOCAL_CLIENT else _UNRESOLVED + elif isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store) and node.id not in self.local: + self.local[node.id] = _UNRESOLVED + + def resolve(self, node: ast.expr) -> Any: + """What a name or a dotted attribute resolves to, looking only into sqlload's modules and clients.""" + if isinstance(node, ast.Name): + return self.local[node.id] if node.id in self.local else self.outer.get(node.id, _UNRESOLVED) + if not isinstance(node, ast.Attribute): + return _UNRESOLVED + owner = self.resolve(node.value) + if isinstance(owner, types.ModuleType) and _ours(owner): + return getattr(owner, node.attr, _UNRESOLVED) + if (owner is _LOCAL_CLIENT or isinstance(owner, _http.Http)) and node.attr in JOB_METHODS: + return node.attr + return _UNRESOLVED + + +def _ours(module: Any) -> bool: + """Whether `module` is sqlload or one of its modules.""" + name = getattr(module, "__name__", "") if isinstance(module, types.ModuleType) else "" + return name == PACKAGE or name.startswith(f"{PACKAGE}.") + + +def _job_method(call: ast.Call, names: _Names) -> str | None: + """`paginate` or `fan_out` when `call` builds a job on an `sl.http` client.""" + method = names.resolve(call.func) + return method if isinstance(method, str) else None + + +def _over_jobs(definition: ast.AST, jobs: list[ast.Call]) -> set[int]: + """The jobs passed as another job's `over=`, directly or through a name assigned them.""" + assigned: dict[str, list[int]] = {} + for node in ast.walk(definition): + if isinstance(node, ast.Assign) and len(node.targets) == 1 and isinstance(node.targets[0], ast.Name) and isinstance(node.value, ast.Call): + assigned.setdefault(node.targets[0].id, []).append(id(node.value)) + over: set[int] = set() + for call in jobs: + for keyword in call.keywords: + if keyword.arg != "over": + continue + if isinstance(keyword.value, ast.Call): + over.add(id(keyword.value)) + elif isinstance(keyword.value, ast.Name): + over.update(assigned.get(keyword.value.id, [])) + return over + + +def _read_call(call: ast.Call, names: _Names, found: dict[str, Any]) -> None: + """Add `call`'s literal `fields=` and `key_path=` paths to `found`, or what it can't read as literals.""" + for keyword in call.keywords: + if keyword.arg is None: + found["computed"].append(f"{_job_method(call, names)}(**...)") + elif keyword.arg == "fields": + _read_fields(keyword.value, names, found) + elif keyword.arg == "transform" and isinstance(keyword.value, ast.Constant) and isinstance(keyword.value.value, str): + _, problem = TRANSFORMS.lookup(keyword.value.value) + if problem is not None: + found["transform_problems"].append(f"transform {keyword.value.value}: {problem}") + elif keyword.arg == "key_path": + path = _literal_path(keyword.value, names, (_data.parent,)) + _record(found, "key_path", path, lambda text: found["key_paths"].append(text)) + + +def _read_fields(value: ast.expr, names: _Names, found: dict[str, Any]) -> None: + if isinstance(value, ast.Constant) and value.value is None: + return + if not isinstance(value, ast.Dict): + found["computed"].append("fields") + return + for key, spec in zip(value.keys, value.values): + if not (isinstance(key, ast.Constant) and isinstance(key.value, str)): + found["computed"].append("fields") + continue + name = key.value + if isinstance(spec, ast.Call) and names.resolve(spec.func) is _data.key: + continue + path = _literal_path(spec, names, (_data.field, _data.parent)) + _record(found, f"fields[{json.dumps(name)}]", path, lambda text, name=name: found["fields"].append([name, text])) + + +def _literal_path(value: ast.expr, names: _Names, functions: tuple[Callable[..., Any], ...]) -> str | None: + """The path `value` writes as a literal: a string, or the first argument (or `path=`) of one of sqlload's `functions`.""" + if isinstance(value, ast.Constant) and isinstance(value.value, str): + return value.value + if not (isinstance(value, ast.Call) and names.resolve(value.func) in functions): + return None + given = [*value.args[:1], *(keyword.value for keyword in value.keywords if keyword.arg == "path")] + return next(_strings(given), None) + + +def _strings(values: list[ast.expr]) -> Iterator[str]: + for value in values: + if isinstance(value, ast.Constant) and isinstance(value.value, str): + yield value.value + + +def _record(found: Mapping[str, Any], label: str, path: str | None, add: Callable[[str], None]) -> None: + if path is None: + found["computed"].append(label) + else: + add(path) diff --git a/packages/sqlload/python/sqlload/_runtime.py b/packages/sqlload/python/sqlload/_runtime.py index 7569366f9..2e804ae03 100644 --- a/packages/sqlload/python/sqlload/_runtime.py +++ b/packages/sqlload/python/sqlload/_runtime.py @@ -20,6 +20,7 @@ from sqlload._declare import Connection, Context, Source, Table from sqlload._http import Job +from sqlload._http_paths import literal_paths from sqlload._transform import TRANSFORMS, TransformContext, returned_items #: Folders under `python_paths` never imported: virtual environments, build output and caches. @@ -72,6 +73,7 @@ def discover(root: str, python_paths: list[str]) -> str: "name": table.name, "module": relative, "settings": {"table_name": STANDALONE_TABLE_NAME, "tables": {table.name: _table_settings(table)}}, + "http_paths": _http_paths([table]), } ) for source, _ in sources.values(): @@ -244,7 +246,13 @@ def _source_declaration(source: Source, relative: str) -> dict[str, Any]: if source.client_factory is not None: settings["shared_client"] = True settings["tables"] = {table.name: _table_settings(table) for table in source.tables} - return {"name": source.name, "module": relative, "settings": settings} + return {"name": source.name, "module": relative, "settings": settings, "http_paths": _http_paths(source.tables)} + + +def _http_paths(tables: list[Table]) -> dict[str, Any]: + """The `sl.http` paths each of `tables` writes as literals, for `plan` to check, by table name; tables that build no job left out.""" + found = {table.name: literal_paths(table.function) for table in tables} + return {name: paths for name, paths in found.items() if paths is not None} def _table_settings(table: Table) -> dict[str, Any]: diff --git a/packages/sqlload/tests/test_http_plan_paths.py b/packages/sqlload/tests/test_http_plan_paths.py new file mode 100644 index 000000000..cc769f3a8 --- /dev/null +++ b/packages/sqlload/tests/test_http_plan_paths.py @@ -0,0 +1,204 @@ +"""`plan` checks the sl.http field paths a table function writes as literals, read from its source without calling it.""" + +import os +import subprocess +import sys +import textwrap +import threading +from collections.abc import Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +import pytest + +WILDCARD = 'fields["sku"]: field path "$.lines[*].sku" uses the wildcard [*], which one column can\'t hold; take one item with [n] or [-1], or select the list\'s items as rows' + + +class Recorder(BaseHTTPRequestHandler): + """An API that records every request it gets and answers none usefully.""" + + def log_message(self, *args: object) -> None: + pass + + def do_GET(self) -> None: + self.server.requests.append(self.path) # type: ignore[attr-defined] + self.send_response(500) + self.end_headers() + + +@pytest.fixture +def api() -> Iterator[ThreadingHTTPServer]: + server = ThreadingHTTPServer(("127.0.0.1", 0), Recorder) + server.requests = [] # type: ignore[attr-defined] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield server + server.shutdown() + + +def make_project(tmp_path: Path, server: ThreadingHTTPServer, module: str) -> Path: + host, port = server.server_address[:2] + (tmp_path / "sources").mkdir() + (tmp_path / "sqlload.toml").write_text( + textwrap.dedent(f"""\ + python_paths = ["sources"] + + [connections.wh] + kind = "duckdb" + database = "wh.duckdb" + + [connections.shop] + kind = "http" + base_url = "http://{host}:{port}/v1" + + [targets.dev] + connection = "wh" + loader_schema = "RAW" + """) + ) + (tmp_path / "sources" / "shop.py").write_text(textwrap.dedent(module)) + return tmp_path + + +def sqlload(project: Path, *args: str) -> subprocess.CompletedProcess[str]: + home = project / "home" + home.mkdir(exist_ok=True) + env = {**os.environ, "HOME": str(home), "XDG_CONFIG_HOME": str(home / ".config")} + return subprocess.run([sys.executable, "-c", "from sqlload._cli import main; main()", *args], cwd=project, env=env, capture_output=True, text=True, timeout=120) + + +def words(text: str) -> str: + """`text` without the plan's tree lines, one space between words.""" + return " ".join(word for word in text.split() if word not in ("│", "├──", "└──")) + + +#: A paginate whose `fields=` holds `{sku}`, its other paths fine; the function marks a file first, so a call shows. +ORDERS = """ +@shop.table(description="Orders, one row per order.") +def orders(ctx): + open("called", "w").close() + return api.paginate("/orders", paginator=sl.offset(limit=50), select="$.orders[*]", fields={{"id": "$.id", "sku": {sku}}}) +""" + +#: The module header: sqlload imported as `sl`, a source and its client. +HEADER = """ +import sqlload as sl + +shop = sl.source("shop", connection="shop", write_strategy="replace") +api = sl.http(shop) +""" + + +@pytest.mark.parametrize( + "sku", + ['"$.lines[*].sku"', 'sl.field("$.lines[*].sku", required=True)', 'sl.parent("$.lines[*].sku")'], + ids=["string", "sl.field", "sl.parent"], +) +def test_given_a_bad_literal_path_when_planning_then_plan_fails_naming_the_table_field_and_path_without_calling_the_function( + tmp_path: Path, api: ThreadingHTTPServer, sku: str +) -> None: + project = make_project(tmp_path, api, HEADER + ORDERS.format(sku=sku)) + for command in ("plan", "load"): + done = sqlload(project, command) + assert done.returncode == 1, done.stdout + done.stderr + shown = words(done.stdout + done.stderr) + assert f"RAW.shop__orders reason {WILDCARD}" in shown, shown + assert not (project / "called").exists() + assert api.requests == [] # type: ignore[attr-defined] + + +def test_given_aliased_sqlload_imports_when_planning_then_their_calls_are_checked_and_other_libraries_paginate_is_not( + tmp_path: Path, api: ThreadingHTTPServer +) -> None: + module = """ + import sqlload as loader + from sqlload import http as client_for, parent as around + + shop = loader.source("shop", connection="shop", write_strategy="replace") + + + class Pages: + def paginate(self, path, **settings): + return [] + + + @shop.table(description="Order lines, one row per line.") + def lines(ctx): + lines_api = client_for(shop) + return lines_api.fan_out("/orders/{id}", over=[1, 2], select="$.lines", key_path=around("$.ids[*]"), fields={"id": "$.id"}) + + + @shop.table(description="Refunds, one row per refund.") + def refunds(ctx): + yield from Pages().paginate("/refunds", fields={"sku": "$.lines[*].sku"}) + """ + done = sqlload(make_project(tmp_path, api, module), "plan") + shown = words(done.stdout + done.stderr) + assert done.returncode == 1, shown + assert 'RAW.shop__lines reason key_path "$.ids[*]" uses the wildcard [*]' in shown, shown + assert "shop__refunds" in shown and "refunds reason" not in shown, shown + assert api.requests == [] # type: ignore[attr-defined] + + +def test_given_a_computed_path_when_planning_then_plan_passes_and_only_that_table_says_load_checks_it(tmp_path: Path, api: ThreadingHTTPServer) -> None: + module = ( + HEADER + + """ +SKU = "$.lines[0].sku" + + +@shop.table(description="Orders, one row per order.") +def orders(ctx): + return api.paginate("/orders", paginator=sl.single(), select="$.orders[*]", fields={"id": "$.id", "sku": SKU, "at": sl.field(f"$.{ctx.cursor_column}")}) + + +@shop.table(description="Customers, one row per customer.") +def customers(ctx): + return api.paginate("/customers", paginator=sl.single(), select="$.customers[*]", fields={"id": "$.id"}) +""" + ) + done = sqlload(make_project(tmp_path, api, module), "plan") + shown = words(done.stdout + done.stderr) + assert done.returncode == 0, shown + assert "RAW.shop__orders replace" in shown and 'paths checked at load, as they\'re computed: fields["sku"], fields["at"]' in shown, shown + assert shown.count("paths checked at load") == 1, shown + assert api.requests == [] # type: ignore[attr-defined] + + +def test_given_a_table_function_without_source_when_planning_then_plan_passes_and_says_load_checks_its_paths(tmp_path: Path, api: ThreadingHTTPServer) -> None: + module = ( + HEADER + + """ +built = {} +exec('def orders(ctx):\\n return api.paginate("/orders", paginator=sl.single(), fields={"sku": "$.lines[*].sku"})\\n', globals(), built) +orders = shop.table(description="Orders, one row per order.")(built["orders"]) +""" + ) + done = sqlload(make_project(tmp_path, api, module), "plan") + shown = words(done.stdout + done.stderr) + assert done.returncode == 0, shown + assert "paths checked at load, as the table function's source can't be read" in shown, shown + assert api.requests == [] # type: ignore[attr-defined] + + +def test_given_a_page_transform_named_by_a_name_nothing_registered_when_planning_then_plan_fails_listing_the_registered_ones( + tmp_path: Path, api: ThreadingHTTPServer +) -> None: + module = ( + HEADER + + """ +@sl.transform +def tidy_page(rows, ctx): + return rows + + +@shop.table(description="Orders, one row per order.") +def orders(ctx): + return api.paginate("/orders", paginator=sl.single(), select="$.orders[*]", transform="tidy_pages") +""" + ) + done = sqlload(make_project(tmp_path, api, module), "plan") + shown = words(done.stdout + done.stderr) + assert done.returncode == 1, shown + assert "RAW.shop__orders reason transform tidy_pages: no function is registered as tidy_pages with @sl.transform; the registered ones are tidy_page" in shown, shown + assert api.requests == [] # type: ignore[attr-defined] diff --git a/website/src/content/docs/docs/sources/python.mdx b/website/src/content/docs/docs/sources/python.mdx index 6ce35354b..c648de2c1 100644 --- a/website/src/content/docs/docs/sources/python.mdx +++ b/website/src/content/docs/docs/sources/python.mdx @@ -221,6 +221,46 @@ From 0.50.0, a path in `fields` (a plain string, `sl.field` or `sl.parent`), in fails before any request, and `sl.json` when it is called. Before 0.50.0 these paths loaded NULL without a warning. +From 0.52.0, `plan` checks the paths an `sl.http` job writes as literals, so a bad one fails `plan` +instead of the load. It reads the table function's source and never calls it, so it sends no +request. It checks the `fields=` values of `paginate` and `fan_out`, the paths in their `sl.field` +and `sl.parent`, and `key_path=`, with the same checker as load: + +```console +$ sqlload plan +Plan ready 1 selected + +Errors (1) +└── raw.shop__orders + └── reason fields["sku"]: field path "$.lines[*].sku" uses the wildcard [*], which one column can't hold; take one item with [n] or [-1], or select the list's items as rows + +✗ Plan failed with errors 1 selected +``` + +- Only calls whose names lead to sqlload count, however it was imported (`import sqlload as sl`, + `from sqlload import parent`), on a client from `sl.http(...)` made at module level or in the + function. A `paginate` method of another library isn't checked. +- `plan` checks every literal path in the table function, including those in branches a load may + not take, such as an `if` on `ctx.cursor_value`; a bad path there fails `plan` too. +- A path `plan` can't read as a literal, such as a variable, an f-string or a `fields` dict built + elsewhere, is checked at load, before any request, as before. So is every path of a function + whose source can't be read. `plan` says so on that table: + + ```console + $ sqlload plan + … + └── raw.shop__orders replace + ├── source python sources/shop.py (shop) + … + └── paths checked at load, as they're computed: fields["sku"] + … + ``` + +- A job built in another function that the table function calls isn't in the table function's + source, so only load checks it, and `plan` shows nothing for it. +- A job passed as another job's `over=` takes no `fields` or `key_path`: load refuses them before + any request, whatever their paths, so `plan` leaves its paths alone. + #### Values from around the row From 0.48.0, an `sl.http` job's fields can read beyond the row itself, all in Rust. This loads one From 361379e24fc18df31c5870bd5b78d9b7301c4b21 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 15 Nov 2026 04:19:48 +0000 Subject: [PATCH 6/7] feat(load): load --follow prints the same plan line, table lines and live block as load load --follow now goes through the same streamed output as load: the one-line plan, a "written" line for each table in each batch with how far it is behind the server, the loaded lines and the closing summary. On a terminal it draws the live block, marked FOLLOWING, with log lines printed above it. Piped output has no escape codes, and --json is unchanged. A plan with errors still refuses to start, now under the new plan view. --- crates/sqlload/src/cdc/_helpers/stream.rs | 18 +++++ .../src/command/_helpers/load_output.rs | 79 +++++++++++++++++-- .../src/command/_helpers/loader_commands.rs | 37 ++++----- .../sqlload/src/command/_helpers/logging.rs | 24 +++++- .../src/command/tests/test_load_output.rs | 31 ++++++++ crates/sqlload/src/load_events/models.rs | 6 ++ .../src/loader/_helpers/table_queue.rs | 25 ++++++ crates/sqlload/src/loader/main/load_tables.rs | 9 ++- website/src/content/docs/docs/cli/load.mdx | 4 + .../content/docs/docs/guides/postgres-cdc.mdx | 29 ++++--- 10 files changed, 216 insertions(+), 46 deletions(-) diff --git a/crates/sqlload/src/cdc/_helpers/stream.rs b/crates/sqlload/src/cdc/_helpers/stream.rs index 4bebd3b45..e1d1065c5 100644 --- a/crates/sqlload/src/cdc/_helpers/stream.rs +++ b/crates/sqlload/src/cdc/_helpers/stream.rs @@ -38,6 +38,8 @@ use crate::column_types::models::ColumnType; use crate::file_store::main::open_folder::open_folder; use crate::file_store::models::StoredFile; use crate::files_source::main::read_footers::read_footers; +use crate::load_events::main::emit::emit; +use crate::load_events::models::EventKind; use crate::loader::constants::SOURCE_BATCH_ROWS; use crate::loader::errors::LoaderError; use crate::loader::main::apply_project::apply_project; @@ -983,10 +985,26 @@ impl CdcSource<'_> { behind_secs, "batch written" ); + self.report_batch(&written.rows, behind_secs); } Ok((buffer, written)) } + /// A `table_batch_written` event for each table the followed batch wrote `rows` to (by destination table name), named as `plan` names it, with how long ago the batch's newest change committed (0 for a commit time ahead of this clock, which is skew). + fn report_batch(&self, rows: &BTreeMap, behind_secs: Option) { + for table in &self.tables { + if let Some(rows) = rows.get(&table.spec.destination.table) { + emit( + Some(&table.outcome_name()), + EventKind::TableBatchWritten { + rows: *rows, + behind_secs: behind_secs.map(|secs| secs.max(0).unsigned_abs()), + }, + ); + } + } + } + /// Each loaded table's floor by its source `(schema, table)`: the stream drops its changes committed before it. That is its copy point, from its snapshot record; in a files target, past its newest committed change too, which files cannot skip on a replay as a warehouse's changelog insert does. async fn copy_points(&self) -> Result, LoaderError> { let mut points: BTreeMap<(String, String), i64> = BTreeMap::new(); diff --git a/crates/sqlload/src/command/_helpers/load_output.rs b/crates/sqlload/src/command/_helpers/load_output.rs index 8fe48f9b6..570ccdc2a 100644 --- a/crates/sqlload/src/command/_helpers/load_output.rs +++ b/crates/sqlload/src/command/_helpers/load_output.rs @@ -2,10 +2,12 @@ use std::collections::BTreeMap; use std::io::Write; +use std::sync::PoisonError; use std::time::Duration; use tokio::sync::mpsc::UnboundedReceiver; +use crate::command::_helpers::logging::LIVE_BLOCK; use crate::command::_helpers::output::print_lines; use crate::command::constants::PLAN_TREE_HINT; use crate::load_events::main::event_line::event_line; @@ -51,8 +53,8 @@ const MAX_WIDTH: usize = 112; /// How long after a wait a connection's rate limit still shows as in use. const RATE_RECENT: Duration = Duration::from_millis(1500); -/// The columns of the live block's bar line besides the bar and its counts: the badge, a space before the bar, two spaces before the counts, and one short of the edge so the line never wraps. -const BAR_LINE_GAPS: usize = " LOADING ".len() + 3 + 1; +/// The columns of the live block's bar line besides the badge, the bar and its counts: a space before the bar, two spaces before the counts, and one short of the edge so the line never wraps. +const BAR_LINE_GAPS: usize = 3 + 1; /// Where a load's lines go: printed as they come (`print`), or kept to return with the rest. #[derive(Debug)] @@ -139,6 +141,8 @@ struct Running { /// What it is doing: `batch 3`, `committing`, or empty before its first progress. doing: String, committing: bool, + /// Under `load --follow`, the rows of the batches written so far, which the rows of the one being written add to, and how many. + written: (u64, u64), } /// What the block shows, as the events left it. @@ -156,6 +160,8 @@ pub(crate) struct LiveState { stopped_after: Option, /// Each connection with a `rate_limit`: its limit and when a request last waited, as time since the load started. rates: BTreeMap)>, + /// `load --follow`: the tables stream until it stops. + follow: bool, } impl LiveState { @@ -169,6 +175,14 @@ impl LiveState { } } + /// The state of `load --follow`, whose tables stream until it stops. + pub(crate) fn following(self) -> Self { + Self { + follow: true, + ..self + } + } + /// `event`, which happened `at` into the load, taken in. pub(crate) fn apply(&mut self, event: &LoadEvent, at: Duration) { let table = event.table.clone().unwrap_or_default(); @@ -181,21 +195,33 @@ impl LiveState { rows: 0, doing: String::new(), committing: false, + written: (0, 0), }); } EventKind::TableProgress { unit, count, rows } => { if let Some(running) = self.running_mut(&table) { - running.rows = *rows; + running.rows = running.written.0 + rows; running.doing = format!("{unit} {count}"); } } EventKind::TableCommitting { rows } => { if let Some(running) = self.running_mut(&table) { - running.rows = *rows; + running.rows = running.written.0 + rows; running.doing = "committing".to_owned(); running.committing = true; } } + EventKind::TableBatchWritten { rows, .. } => { + if let Some(running) = self.running_mut(&table) { + running.written = (running.written.0 + rows, running.written.1 + 1); + running.rows = running.written.0; + running.doing = match running.written.1 { + 1 => "1 batch written".to_owned(), + batches => format!("{batches} batches written"), + }; + running.committing = false; + } + } EventKind::RateLimited { connection, rate, .. } => { @@ -270,6 +296,21 @@ pub(crate) fn event_lines( lines.push(tally_line(("loaded", name), &detail, name_width)); return lines; } + EventKind::TableBatchWritten { rows, behind_secs } => { + let word = match rows { + 1 => "row", + _ => "rows", + }; + let behind = behind_secs + .map(|secs| format!("{secs}s behind")) + .unwrap_or_default(); + ( + "written", + format!("{:>rows_width$} {word:<5}{behind}", thousands(*rows)) + .trim_end() + .to_owned(), + ) + } EventKind::TableFailed { reason, peak_memory_bytes, @@ -367,6 +408,9 @@ pub(crate) async fn write_events( let mut tally = Tally::default(); let started = std::time::Instant::now(); let block = output.live.as_ref().map(|_| live_bar()); + if let Some(bar) = &block { + *LIVE_BLOCK.lock().unwrap_or_else(PoisonError::into_inner) = Some(bar.clone()); + } let width = console::Term::stderr() .size_checked() .map_or(MAX_LIVE_WIDTH, |(_, columns)| usize::from(columns)); @@ -403,11 +447,24 @@ pub(crate) async fn write_events( } } if let Some(bar) = block { + *LIVE_BLOCK.lock().unwrap_or_else(PoisonError::into_inner) = None; bar.finish_and_clear(); } (kept, tally) } +/// Run `write`, which writes to stderr, with the live block, when one is drawn, lifted off the screen for it and redrawn under what it wrote. +pub(crate) fn above_live_block(write: impl FnOnce() -> T) -> T { + let drawn = LIVE_BLOCK + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone(); + match drawn { + Some(bar) => bar.suspend(write), + None => write(), + } +} + /// The live block's one bar on stderr: a message of several lines, redrawn in place. fn live_bar() -> indicatif::ProgressBar { let bar = indicatif::ProgressBar::with_draw_target( @@ -566,8 +623,12 @@ pub(crate) fn live_lines( seconds / 60, seconds % 60 ); + let badge = match state.follow { + true => " FOLLOWING ", + false => " LOADING ", + }; let bar_width = width - .saturating_sub(BAR_LINE_GAPS + counts.chars().count()) + .saturating_sub(badge.len() + BAR_LINE_GAPS + counts.chars().count()) .max(10); let filled = match state.total { 0 => bar_width, @@ -577,7 +638,7 @@ pub(crate) fn live_lines( String::new(), format!( "{} {}{} {counts}", - paint(&format!("{BOLD}{INVERSE}{CYAN}"), " LOADING ".to_owned()), + paint(&format!("{BOLD}{INVERSE}{CYAN}"), badge.to_owned()), paint( match state.stopped_after { Some(_) => RED, @@ -662,7 +723,11 @@ pub(crate) fn live_lines( } lines.push(paint( DIM, - " ctrl-c finishes running tables and stops · a second ctrl-c stops at once".to_owned(), + match state.follow { + true => " ctrl-c writes what it read and stops · a second ctrl-c stops at once", + false => " ctrl-c finishes running tables and stops · a second ctrl-c stops at once", + } + .to_owned(), )); lines } diff --git a/crates/sqlload/src/command/_helpers/loader_commands.rs b/crates/sqlload/src/command/_helpers/loader_commands.rs index 8745306d9..69ccb043a 100644 --- a/crates/sqlload/src/command/_helpers/loader_commands.rs +++ b/crates/sqlload/src/command/_helpers/loader_commands.rs @@ -397,6 +397,7 @@ async fn run_command( (vec![text], exit_code(false, retryable)) } false => { + let mut lines = start_lines(lines); lines.extend([String::new(), follow_refused(project, &plan)]); (lines, exit_code(false, plan_retryable)) } @@ -441,11 +442,12 @@ async fn run_command( host, options: &options, }; - if loads && follow.is_none() { + if loads { let streamed = StreamedLoad { run: &run, plan: &plan, applied: &applied, + stop: follow.map(follow_stop), output: LoadOutput { name_width: name_width(&plan, &applied), rows_width: rows_width(&plan), @@ -459,7 +461,11 @@ async fn run_command( .filter(|table| table.schema.blocked.is_none()) .map(TablePlan::outcome_name) .collect(); - LiveState::new(tables, options.fail_fast) + let state = LiveState::new(tables, options.fail_fast); + match follow { + Some(_) => state.following(), + None => state, + } }), }, started, @@ -485,16 +491,11 @@ async fn run_command( lines.extend(ended.lines); return Ok((lines, ended.exit_code)); } - let stop = follow.map(follow_stop); - let outcomes = match loads { - true => load_tables(&run, &plan, stop.as_ref()).await, - false => Vec::new(), - }; if json { let run = RunReport { - command: if loads { "load" } else { "apply" }, + command: "apply", applied: &applied, - outcomes: &outcomes, + outcomes: &[], started: started.0, elapsed: started.1.elapsed(), peak_memory_bytes: run_peak_bytes(&[]), @@ -502,25 +503,17 @@ async fn run_command( let (text, ok, retryable) = json_report(&plan, Some(&run)); return Ok((vec![text], exit_code(ok, retryable))); } - let (tallied, completed) = render_tally( - &applied, - &outcomes, - started.1.elapsed(), - loads.then(|| run_peak_bytes(&[])).flatten(), - ); + let (tallied, completed) = render_tally(&applied, &[], started.1.elapsed(), None); lines.extend(tallied); - let retryable = outcomes.iter().all(TableOutcome::passes_on_rerun); - Ok(( - lines, - exit_code(planned && completed, plan_retryable && retryable), - )) + Ok((lines, exit_code(planned && completed, plan_retryable))) } -/// What a load streams its lines and events from: its run and plan, what `apply` changed before it, where its lines go, when it started, and whether the plan was ready and failed only in ways a rerun may get past. +/// What a load streams its lines and events from: its run and plan, what `apply` changed before it, what stops a `load --follow`, where its lines go, when it started, and whether the plan was ready and failed only in ways a rerun may get past. struct StreamedLoad<'a> { run: &'a LoadRun<'a>, plan: &'a ProjectPlan, applied: &'a [String], + stop: Option, output: LoadOutput, started: (chrono::DateTime, std::time::Instant), planned: (bool, bool), @@ -564,7 +557,7 @@ async fn streamed_load(load: StreamedLoad<'_>) -> StreamedEnd { let writer = tokio::spawn(write_events(load.output, receiver)); emit(None, EventKind::RunStarted { selected, to_load }); let before = run_peak_bytes(&[]); - let outcomes = load_tables(load.run, load.plan, None).await; + let outcomes = load_tables(load.run, load.plan, load.stop.as_ref()).await; let measured: Vec> = std::iter::once(before) .chain(outcomes.iter().map(|outcome| outcome.peak_memory_bytes)) .collect(); diff --git a/crates/sqlload/src/command/_helpers/logging.rs b/crates/sqlload/src/command/_helpers/logging.rs index 92399aaba..335cc0be8 100644 --- a/crates/sqlload/src/command/_helpers/logging.rs +++ b/crates/sqlload/src/command/_helpers/logging.rs @@ -1,14 +1,34 @@ +use std::io::Write; +use std::sync::Mutex; + use tracing_subscriber::EnvFilter; +use crate::command::_helpers::load_output::above_live_block; use crate::command::models::LogFormat; -/// Log to stderr at `info`, without the ODBC crates' per-column buffer notes (`arrow_odbc` at `warn`, `odbc_api` at `error`), unless `RUST_LOG` says otherwise. +/// `load`'s live block while it is drawn, so a log line prints above it instead of through it. Process-wide like the tracing subscriber it serves, which is installed once per process (the Python package can run several commands in one), so a handle passed to one run would never reach it. +pub(crate) static LIVE_BLOCK: Mutex> = Mutex::new(None); + +/// Stderr for log lines: each written above `load`'s live block when one is drawn, instead of through it. +struct LogWriter; + +impl Write for LogWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + above_live_block(|| std::io::stderr().write(buf)) + } + + fn flush(&mut self) -> std::io::Result<()> { + std::io::stderr().flush() + } +} + +/// Log to stderr at `info`, without the ODBC crates' per-column buffer notes (`arrow_odbc` at `warn`, `odbc_api` at `error`), unless `RUST_LOG` says otherwise; above `load`'s live block when it is drawn. pub(crate) fn init_logging(format: LogFormat) { let filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| EnvFilter::new("info,odbc_api=error,arrow_odbc=warn")); let builder = tracing_subscriber::fmt() .with_env_filter(filter) - .with_writer(std::io::stderr); + .with_writer(|| LogWriter); let _ = match format { LogFormat::Text => builder.try_init(), LogFormat::Json => builder.json().try_init(), diff --git a/crates/sqlload/src/command/tests/test_load_output.rs b/crates/sqlload/src/command/tests/test_load_output.rs index 896bf9ffb..53b32b2da 100644 --- a/crates/sqlload/src/command/tests/test_load_output.rs +++ b/crates/sqlload/src/command/tests/test_load_output.rs @@ -68,6 +68,37 @@ loaded raw.a 1 row 0.3s loaded raw.b 1 row 0.1s ✓ Completed 2 loaded, 123,457 rows, 2.2s └── slowest raw.a 2.0s", + }, + LoadOutputTestCase { + description: "under --follow each table prints a line per batch written, with how far behind it is, then its total when following stops", + events: vec![ + ( + Some("raw.a"), + EventKind::TableBatchWritten { + rows: 2, + behind_secs: Some(6), + }, + ), + ( + Some("raw.b"), + EventKind::TableBatchWritten { + rows: 1, + behind_secs: None, + }, + ), + (Some("raw.a"), done(2, 18_000, None)), + (Some("raw.b"), done(1, 18_000, None)), + ], + name_width: 7, + rows_width: 5, + elapsed_ms: 18_100, + exit_code: 0, + expected_lines: "written raw.a 2 rows 6s behind +written raw.b 1 row +loaded raw.a 2 rows 18.0s +loaded raw.b 1 row 18.0s +✓ Completed 2 loaded, 3 rows, 18.1s +└── slowest raw.a 18.0s", }, LoadOutputTestCase { description: "under --memory-per-table each table notes its peak memory, after its position or under its failure, and the run's peak closes the summary", diff --git a/crates/sqlload/src/load_events/models.rs b/crates/sqlload/src/load_events/models.rs index 98110bf5a..383aa7f65 100644 --- a/crates/sqlload/src/load_events/models.rs +++ b/crates/sqlload/src/load_events/models.rs @@ -34,6 +34,12 @@ pub enum EventKind { }, /// A table's rows are read and its last commit begins. TableCommitting { rows: u64 }, + /// Under `load --follow`, a batch of the followed source's changes was written: this table's rows in it, and how long ago the batch's newest change committed, in seconds. + TableBatchWritten { + rows: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + behind_secs: Option, + }, /// A table loaded: its rows, how long it took, and its new position when it keeps one. TableDone { rows: u64, diff --git a/crates/sqlload/src/loader/_helpers/table_queue.rs b/crates/sqlload/src/loader/_helpers/table_queue.rs index f351f1f90..1692c66ac 100644 --- a/crates/sqlload/src/loader/_helpers/table_queue.rs +++ b/crates/sqlload/src/loader/_helpers/table_queue.rs @@ -2,8 +2,10 @@ use std::collections::{BTreeMap, BTreeSet}; use std::sync::Arc; use tokio::task::{Id, JoinError}; +use tokio_util::sync::CancellationToken; use crate::adapters::contract::models::WriteMode; +use crate::cdc::main::follow_cdc::follow_cdc; use crate::cdc::main::load_cdc::load_cdc; use crate::load_events::main::emit::emit; use crate::load_events::models::EventKind; @@ -71,6 +73,29 @@ pub(crate) async fn load_task(shared: Arc, table: TablePlan) -> Table } } +/// Follow the plan's one `postgres_cdc` source until `stop` is cancelled (see [`follow_cdc`]), its unblocked tables reported started at once and their outcomes at the end. +pub(crate) async fn follow_tables( + run: &LoadRun<'_>, + plan: &ProjectPlan, + stop: &CancellationToken, +) -> Vec { + for table in plan + .tables + .iter() + .filter(|table| table.schema.blocked.is_none()) + { + emit( + Some(&table.outcome_name()), + EventKind::TableStarted { + strategy: strategy_name(table).to_owned(), + }, + ); + } + let outcomes = Box::pin(follow_cdc(run, plan.clone(), stop)).await; + outcomes.iter().for_each(emit_outcome); + outcomes +} + /// `table`'s write strategy as the load events name it. pub(crate) fn strategy_name(table: &TablePlan) -> &'static str { match table.mode { diff --git a/crates/sqlload/src/loader/main/load_tables.rs b/crates/sqlload/src/loader/main/load_tables.rs index b04ce42c5..b2fba69f7 100644 --- a/crates/sqlload/src/loader/main/load_tables.rs +++ b/crates/sqlload/src/loader/main/load_tables.rs @@ -4,13 +4,14 @@ use tokio::runtime::Handle; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; -use crate::cdc::main::follow_cdc::follow_cdc; use crate::loader::_helpers::destination::clean_stages; use crate::loader::_helpers::destination::clean_staging; -use crate::loader::_helpers::table_queue::{TableQueue, load_task, shared_run, thread_count}; +use crate::loader::_helpers::table_queue::{ + TableQueue, follow_tables, load_task, shared_run, thread_count, +}; use crate::loader::models::{LoadRun, ProjectPlan, TableOutcome}; -/// Load each unblocked table, Python tables through `host`, each on a thread of its own, up to `concurrency` at once (one table's batches commit one at a time), so tables use several cores (decoding, masking and metadata included) and a slow one never holds up the others; tables of a Python source sharing a client run one at a time unless it sets `parallel_tables`. A table that fails is reported and the others still load, unless `--fail-fast` starts no table after it (those not started, `postgres_cdc` sources' included, are reported skipped); outcomes are in plan order. Load folders and staging tables crashed runs left are removed first. `postgres_cdc` sources load after the other tables, each as one unit. Given `follow` (`load --follow`), the plan's one `postgres_cdc` source is followed instead until it is cancelled (see [`follow_cdc`]). +/// Load each unblocked table, Python tables through `host`, each on a thread of its own, up to `concurrency` at once (one table's batches commit one at a time), so tables use several cores (decoding, masking and metadata included) and a slow one never holds up the others; tables of a Python source sharing a client run one at a time unless it sets `parallel_tables`. A table that fails is reported and the others still load, unless `--fail-fast` starts no table after it (those not started, `postgres_cdc` sources' included, are reported skipped); outcomes are in plan order. Load folders and staging tables crashed runs left are removed first. `postgres_cdc` sources load after the other tables, each as one unit. Given `follow` (`load --follow`), the plan's one `postgres_cdc` source is followed instead until it is cancelled (see [`follow_tables`]). pub async fn load_tables( run: &LoadRun<'_>, plan: &ProjectPlan, @@ -19,7 +20,7 @@ pub async fn load_tables( clean_stages(run, plan).await; clean_staging(run, plan).await; if let Some(stop) = follow { - return follow_cdc(run, plan.clone(), stop).await; + return follow_tables(run, plan, stop).await; } let captured = plan; let plan = &plan.uncaptured(); diff --git a/website/src/content/docs/docs/cli/load.mdx b/website/src/content/docs/docs/cli/load.mdx index 4f83b7fa6..ffa5fd350 100644 --- a/website/src/content/docs/docs/cli/load.mdx +++ b/website/src/content/docs/docs/cli/load.mdx @@ -167,6 +167,10 @@ SIGTERM. A second Ctrl-C or SIGTERM stops it at once, with exit code 130. See sqlload load app --follow --max-duration 10m ``` +From 0.52.0 it prints the same plan line, table lines and summary as `load`, with a +`written` line for each table in each batch, giving its rows and how far it is behind the server. +On a terminal, a live block shows the rows and batches written so far. + ## Results as JSON From 0.27.0, `apply --json` and `load --json` print one JSON document on diff --git a/website/src/content/docs/docs/guides/postgres-cdc.mdx b/website/src/content/docs/docs/guides/postgres-cdc.mdx index 556aca25e..cdc7d125c 100644 --- a/website/src/content/docs/docs/guides/postgres-cdc.mdx +++ b/website/src/content/docs/docs/guides/postgres-cdc.mdx @@ -217,24 +217,31 @@ sqlload load app --follow # until Ctrl-C or SIGTERM sqlload load app --follow --max-duration 10m ``` -It logs each batch it writes, with how far it is behind the server. Here an order is inserted and -then updated while it follows, then Ctrl-C stops it: +It logs each batch it writes, with how far it is behind the server. From 0.52.0 it prints the same +plan line, table lines and summary as `load`, plus a `written` line for each table in each batch. +Here an order is inserted and then updated while it follows, then Ctrl-C stops it: ```console $ sqlload load app --follow -… -2026-11-07T11:24:53.505015Z INFO sqlload::cdc::main::follow_cdc: following -2026-11-07T11:25:03.555267Z INFO sqlload::cdc::_helpers::stream: batch written source="app" changes=2 tables=1 lsn=0/15B3720 behind_bytes=0 behind_secs=6 -2026-11-07T11:25:11.518577Z INFO sqlload::command::_helpers::loader_commands: stopping once the batch or copy under way is written; press Ctrl-C again to stop now -2026-11-07T11:25:11.520357Z INFO sqlload::cdc::main::follow_cdc: stopped following rows=1 -loaded raw.app__customers 0 rows 18.0s -loaded raw.app__orders 1 rows 18.0s - -✓ Completed 2 loaded, 1 rows, 18.1s +✓ Plan complete 2 selected, 2 to load (sqlload plan shows the tree) + +2026-11-15T03:47:51.265378Z INFO sqlload::cdc::main::follow_cdc: following +2026-11-15T03:48:01.402075Z INFO sqlload::cdc::_helpers::stream: batch written source="app" changes=2 tables=1 lsn=0/1572DF0 behind_bytes=56 behind_secs=2 +written raw.app__orders 1 row 2s behind +2026-11-15T03:48:10.656949Z INFO sqlload::command::_helpers::loader_commands: stopping once the batch or copy under way is written; press Ctrl-C again to stop now +2026-11-15T03:48:10.660993Z INFO sqlload::cdc::main::follow_cdc: stopped following rows=1 +loaded raw.app__customers 0 rows 19.4s +loaded raw.app__orders 1 row 19.4s + +✓ Completed 2 loaded, 1 row, 19.5s +├── slowest raw.app__customers 19.4s +└── peak memory 84.3 MiB ``` `behind_bytes` is how far the batch is behind the server's WAL position, and `behind_secs` how long ago its newest change committed. The two changes to one order load as one row under `merge`. +On a terminal, a live block under the logs shows `FOLLOWING`, the rows and batches written so far, +and how to stop; piped or redirected output has no escape codes, and `--json` is unchanged. - **Stopping.** `--max-duration`, or a first Ctrl-C or SIGTERM, writes and confirms what it has read, then exits 0. A snapshot or copy under way finishes first. A second signal stops at once From 943f18c4d7703b660176af47be14c9b12fcdaa30 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 15 Nov 2026 13:11:57 +0000 Subject: [PATCH 7/7] fix(cli): write log lines without colour codes when stderr is not a terminal or NO_COLOR is set Log lines went to stderr in colour even when it was piped or redirected, so a CI log or a captured run held ANSI escape codes between every field. Colour is now used only on a terminal, and never under NO_COLOR, as the rest of the output already does. --- crates/sqlload/src/command/_helpers/logging.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/sqlload/src/command/_helpers/logging.rs b/crates/sqlload/src/command/_helpers/logging.rs index 335cc0be8..12721b766 100644 --- a/crates/sqlload/src/command/_helpers/logging.rs +++ b/crates/sqlload/src/command/_helpers/logging.rs @@ -1,4 +1,4 @@ -use std::io::Write; +use std::io::{IsTerminal, Write}; use std::sync::Mutex; use tracing_subscriber::EnvFilter; @@ -22,12 +22,14 @@ impl Write for LogWriter { } } -/// Log to stderr at `info`, without the ODBC crates' per-column buffer notes (`arrow_odbc` at `warn`, `odbc_api` at `error`), unless `RUST_LOG` says otherwise; above `load`'s live block when it is drawn. +/// Log to stderr at `info`, without the ODBC crates' per-column buffer notes (`arrow_odbc` at `warn`, `odbc_api` at `error`), unless `RUST_LOG` says otherwise; above `load`'s live block when it is drawn. In colour only on a terminal without `NO_COLOR`, so piped and redirected logs hold no escape codes. pub(crate) fn init_logging(format: LogFormat) { let filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| EnvFilter::new("info,odbc_api=error,arrow_odbc=warn")); + let coloured = std::io::stderr().is_terminal() && std::env::var_os("NO_COLOR").is_none(); let builder = tracing_subscriber::fmt() .with_env_filter(filter) + .with_ansi(coloured) .with_writer(|| LogWriter); let _ = match format { LogFormat::Text => builder.try_init(),